Add end-to-end distributed test suite (Phase 0: DDP smoke + env parity) - #4143
Draft
Titus-von-Koeller wants to merge 1 commit into
Draft
Add end-to-end distributed test suite (Phase 0: DDP smoke + env parity)#4143Titus-von-Koeller wants to merge 1 commit into
Titus-von-Koeller wants to merge 1 commit into
Conversation
Introduces tests/distributed/ as a unified end-to-end layer for distributed
training tests, mirroring the pattern from transformers#44338 and trl#4784:
per-backend YAML launch configs, a shared training script that is the system
under test, and a three-piece test structure per backend.
Phase 0 lands the pattern for DDP:
- scripts/train.py: raw Accelerator loop over a config-constructed tiny
LlamaForCausalLM and hermetic seeded synthetic tokens. Emits JSON metrics
(loss history, distributed_type, num_processes, mixed_precision) that the
harness asserts on.
- scripts/env_check.py: per-rank distributed environment dumper for
launcher-parity testing.
- accelerate_configs/{ddp,fsdp2}.yaml: minimal 2-process launch configs
(fsdp2.yaml included for the follow-up phase; no tests yet).
- test_e2e_distributed.py: DDPCommandsMixin + TestDistributedDDP
(env parity between accelerate launch and torchrun, non-slow) +
TestDistributedDDPCommon (smoke, @slow).
Reproduction:
pytest -v tests/distributed/test_e2e_distributed.py
RUN_SLOW=1 pytest -v tests/distributed/test_e2e_distributed.py
Follow-up phases will add L1 convergence and L2 resume-equivalence
assertions, FSDP1/2 and DeepSpeed z2/z3, and eventual L3 cross-backend
parity checks.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
tests/distributed/, a unified end-to-end layer for distributed training tests.Backend selection is driven by per-backend YAML launch configs, and one training script
is exercised through every backend's
Accelerator.prepare()dispatch branch. Mirrorsthe pattern introduced in
huggingface/transformers#44338and
huggingface/trl#4784, adapted forthe fact that
acceleratehas noTrainer— the rawAcceleratorloop is itself thesystem under test.
Motivation
Accelerate's existing distributed tests are either unit-level (config/plugin dataclass
assertions) or backend-siloed (
tests/fsdp/,tests/deepspeed/,tests/tp/withdifferent conventions and dataset-downloading legacy scripts). There is no unified
end-to-end layer asserting on user-facing correctness properties: run-to-completion,
convergence on memorizable data, save/resume equivalence, and eventual cross-backend
parity.
What lands in this PR (Phase 0)
Foundation only. Two tests, DDP.
scripts/train.py— rawAcceleratorloop over a config-constructed tinyLlamaForCausalLMand hermetic seeded synthetic tokens. Emits a JSON summary(loss history, distributed type, num processes, mixed precision) for the harness to
assert on. Exercises
Accelerator.prepare,accelerator.accumulate,accelerator.backward,accelerator.clip_grad_norm_,accelerator.save_state/load_state,accelerator.gather_for_metrics,accelerator.wait_for_everyone, andaccelerator.is_main_process.scripts/env_check.py— per-rank distributed environment dumper for launcher-paritychecks.
accelerate_configs/ddp.yaml— 2-process DDP config.accelerate_configs/fsdp2.yaml— 2-process FSDP2 config (included for the follow-upphase; not exercised by any test in this PR).
test_e2e_distributed.py— three-piece structure per backend:DDPCommandsMixin— buildsaccelerate launch --config_file=ddp.yamlandtorchruncommand lines.TestDistributedDDP(non-@slow) — asserts env-parity between the two launchersusing
env_check.py(distributed_type,device,mixed_precision,env_WORLD_SIZE,env_RANK,process_index,local_process_index).TestDistributedDDPCommon(@slow) — smoke: launchestrain.pyfor 4 steps,asserts on JSON output shape and
math.isfiniteloss values.Not in this PR (follow-ups)
Each will land in its own PR under the same three-piece pattern:
final_loss < initial_loss * margin) andL2 resume-equivalence (save mid-training, load in a fresh process, assert trajectory
match). Adds FSDP1 and FSDP2 test classes.
and gradient-accumulation scenarios.
best-effort for DeepSpeed), eval via
gather_for_metrics, testing guide.Testing
@require_multi_device.ruff checkandruff format --checkclean.make test_core— no Makefile changes needed.Notes for reviewers
pytestpicks uptest_e2e_distributed.pyonly.scripts/train.pyandscripts/env_check.pydon't match the collection pattern andare only invoked as subprocesses by the harness.
LlamaForCausalLMis built from aLlamaConfigwith tiny dims (
vocab=256,hidden=64, 2 layers, 2 heads). Nofrom_pretrained,no cached weights, no network dependency.
scripts/currently lives undertests/distributed/(matching the
transformersPR). The legacy convention in accelerate places testscripts under
src/accelerate/test_utils/scripts/external_deps/. Happy to move ifmaintainers prefer the legacy location.