Overlap experimental Megatron-FSDP communication with two process groups

[wujingyue]

Summary

Overlap communication in the experimental Megatron-FSDP path (mfsdp_v2) by
launching gradient reduce-scatter eagerly in post_backward on a dedicated
NCCL communicator, instead of the deferred "delayed reduction" scheduling.
All-gather stays on the mesh's default group; reduce-scatter runs on a separate
group, so the two collectives use different communicators and can run
concurrently. Delayed releases (forward prefetch) are unchanged.

What changed

• FsdpContext gains a dedicated reduce_scatter_group; each group's
  reduce-scatter is launched immediately once its gradients are packed, on
  reduce_scatter_stream.
• Removed the delayed-reduction machinery: PreparedReduction/PendingReduction,
  the prepared_reductions/pending_reductions queues, and the deferral to the
  next pre_backward. The end-of-backward callback is now a single wait_stream
  barrier so the optimizer stream orders after the reductions (net -17 SLOC in
  module.py).
• DBuffer.reduce_scatter/redistribute and
  FsdpParameterGroup.reduce_partial_gradients accept an optional reduce_group.

Why

The delayed reduction existed to overlap reduce-scatter with compute while keeping
all-gather and reduce-scatter ordered on one communicator. With a separate
reduce-scatter communicator, reduce-scatter launches eagerly without head-of-line
blocking against the next all-gather, removing the deferral/retention bookkeeping.

Validation

• Numerics: test_fully_shard_losses_match_baseline,
  test_backward_averages_across_dp_and_accumulates_across_calls, and the
  accumulation / microbatch / memory tests pass (2 ranks).
• Overlap (nsys, 2xH100, torch 2.13 / NCCL 2.30): at dim=8192, reduce-scatter
  overlaps backward GEMM (~0.45-0.66 of RS hidden under compute) and all-gather
  (~0.44); test_overlaps_communication_and_compute passes 5/5 at dim=8192. At
  dim=4096 the collectives are too small to overlap (the delayed-reduction
  baseline behaves identically) — see Increase Megatron-FSDP overlap test dim to 8192 for reliable overlap #5770 for the test-dim fix.

Known limitations

• Communicator lifecycle: the dedicated group is created via dist.new_group
  inside FsdpContext and not torn down — one leaked communicator per FSDP
  context. It should instead be created once at process-group setup and passed in
  via a ProcessGroupCollection, as mfsdp v1 does with fsdp_group_ag
  (parallel_state.create_all_gather_groups). Not addressed here.
• Convention differs from v1: v1 dedicates the all-gather group and keeps
  reduce-scatter on the default; this dedicates the reduce-scatter group.
  Equivalent for overlap.
• 1D DP mesh only in this prototype; multi-axis meshes need per-subgroup
  construction.
• End-to-end step-time impact vs. the delayed-reduction baseline is not yet
  measured.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.