Overlap FSDP communication with compute

[wujingyue]

Summary

This implements the explicit prefetching approach used in MFSDP v1 for MFSDP v2 (mfsdp_v2). The current implementation assumes that static module orders match the execution order closely enough for one-step prefetch.

Each FsdpModule's pre-forward hook prefetches the next FsdpModule's parameter all-gather from forward_order. During backward, pre-hooks follow a separately computed backward_order instead of assuming backward is strictly the reverse of forward_order, which matters for nested module trees.

An example trace from the added unit test:

Implementation

Forward prefetch

Each FsdpModule tracks materialization with _unshard_event. _unshard_parameter_groups() is idempotent: if _unshard_event is already set, the FsdpModule was already unsharded or prefetched; otherwise it enqueues the all-gather on the shared all-gather stream and records _unshard_event after materialization.

pre_forward waits on _unshard_event instead of waiting on the whole all-gather stream, so compute depends only on the current FsdpModule's materialization and not on later release work queued on that stream. It then issues the next FsdpModule's all-gather from forward_order. _reshard_parameter_groups() queues release on the same all-gather stream where unsharded storage was allocated, then clears _unshard_event so the next user enqueues a fresh all-gather.

Backward overlap

The shared FSDP context maintains separate static forward_order and backward_order sequences. Each backward pre-hook:

1. calls the same idempotent unshard path and waits only for that FsdpModule's _unshard_event;
2. prefetches the next FsdpModule needed by backward on the all-gather stream;
3. allows that all-gather to overlap the current FsdpModule's gradient computation.

Gradient reduction is split into allocation, packing, and communication stages. Once a module's gradients are packed, its reduce-scatter is launched immediately on a separate reduce-scatter stream, allowing it to overlap subsequent backward GEMMs and all-gathers. Existing symmetric-memory scaling, mixed gradient dtypes, first-backward assignment, and microbatch accumulation semantics are preserved.

An autograd-end callback registered by the root makes the caller's current stream wait for the reduce-scatter stream before backward() returns. This cannot rely on the root module's own post-backward hook: that hook tracks gradients for parameters owned by the root and may run before descendants finish, or may not run at all when the root owns no trainable parameters.

Verification

• test_fully_shard.py passes with two CUDA ranks after the _unshard_event change: 14 passed per rank.
• The nested-order unit test verifies the static forward and backward orders for a nested module tree, including the backward DFS traversal.
• The profiler test covers one forward and backward pass. For four child FsdpModules, it requires all 6 expected all-gather/GEMM overlaps and all 3 expected reduce-scatter/GEMM overlaps, rather than accepting any single overlap.
• The profiler test also verifies that all-gathers use one stream, reduce-scatters use another, and both streams are distinct from compute.
• A local Nsight Systems capture confirms backward all-gathers and reduce-scatters overlap GEMM execution on separate streams.
• Diff and byte-compilation checks pass for the latest edits.

[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.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[shjwudp]

Should we collect the states of unsharded models and release module parameter from all models except the current and next ones?

[wujingyue]

When would that be useful? Some FsdpModule's post-forward or post-backward failed to be triggered? That sounds like a bug that we shouldn't hide.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29476715799