Skip to content

Add end-to-end distributed test suite (Phase 0: DDP smoke + env parity) - #4143

Draft
Titus-von-Koeller wants to merge 1 commit into
mainfrom
feat/distributed-e2e-tests-phase-0
Draft

Add end-to-end distributed test suite (Phase 0: DDP smoke + env parity)#4143
Titus-von-Koeller wants to merge 1 commit into
mainfrom
feat/distributed-e2e-tests-phase-0

Conversation

@Titus-von-Koeller

Copy link
Copy Markdown

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. Mirrors
the pattern introduced in huggingface/transformers#44338
and huggingface/trl#4784, adapted for
the fact that accelerate has no Trainer — the raw Accelerator loop is itself the
system 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/ with
different 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 — raw Accelerator loop over a config-constructed tiny
    LlamaForCausalLM and 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, and
    accelerator.is_main_process.
  • scripts/env_check.py — per-rank distributed environment dumper for launcher-parity
    checks.
  • accelerate_configs/ddp.yaml — 2-process DDP config.
  • accelerate_configs/fsdp2.yaml — 2-process FSDP2 config (included for the follow-up
    phase; not exercised by any test in this PR).
  • test_e2e_distributed.py — three-piece structure per backend:
    • DDPCommandsMixin — builds accelerate launch --config_file=ddp.yaml and
      torchrun command lines.
    • TestDistributedDDP (non-@slow) — asserts env-parity between the two launchers
      using env_check.py (distributed_type, device, mixed_precision,
      env_WORLD_SIZE, env_RANK, process_index, local_process_index).
    • TestDistributedDDPCommon (@slow) — smoke: launches train.py for 4 steps,
      asserts on JSON output shape and math.isfinite loss values.

Not in this PR (follow-ups)

Each will land in its own PR under the same three-piece pattern:

  • Phase 1 — L1 convergence assertion (final_loss < initial_loss * margin) and
    L2 resume-equivalence (save mid-training, load in a fresh process, assert trajectory
    match). Adds FSDP1 and FSDP2 test classes.
  • Phase 2 — DeepSpeed z2/z3 at smoke+convergence level, mixed-precision (bf16),
    and gradient-accumulation scenarios.
  • Phase 3 — L3 cross-backend parity vs DDP (loose tolerance for FSDP,
    best-effort for DeepSpeed), eval via gather_for_metrics, testing guide.

Testing

pytest -v tests/distributed/test_e2e_distributed.py                # non-slow only
RUN_SLOW=1 pytest -v tests/distributed/test_e2e_distributed.py     # full
  • Both tests pass on a 2×GPU host (~20s total).
  • Single-GPU host cleanly SKIPs both tests via @require_multi_device.
  • Three consecutive full-file runs are consistent (no flake).
  • ruff check and ruff format --check clean.
  • Integrates automatically into make test_core — no Makefile changes needed.

Notes for reviewers

  • Pytest collection: pytest picks up test_e2e_distributed.py only.
    scripts/train.py and scripts/env_check.py don't match the collection pattern and
    are only invoked as subprocesses by the harness.
  • Model instantiation is Hub-free: LlamaForCausalLM is built from a LlamaConfig
    with tiny dims (vocab=256, hidden=64, 2 layers, 2 heads). No from_pretrained,
    no cached weights, no network dependency.
  • Structural question: scripts/ currently lives under tests/distributed/
    (matching the transformers PR). The legacy convention in accelerate places test
    scripts under src/accelerate/test_utils/scripts/external_deps/. Happy to move if
    maintainers prefer the legacy location.

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>
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants