Overlap experimental FSDP communication with compute

[wujingyue]

Summary

This implements the “explicit prefetching” approach used in MFSDP v1 for experimental Megatron-FSDP (mfsdp_v2). The current implementation assumes that the static module registration order matches the execution order.

Each module’s pre-forward hook prefetches the next module’s parameter all-gather. During backward, the same order is traversed in reverse: each pre-backward hook prefetches the preceding module’s all-gather, which is the next one needed in backward.

An example trace from the added unit test:

Implementation

Forward prefetch

After materializing the current FSDP unit's parameters, pre_forward issues the next unit's all-gather on the shared all-gather stream. Each unit releases its own unsharded storage in stream order after its compute consumer completes, bounding the live parameter storage to the current unit and its prefetched successor.

Backward overlap

The shared FSDP context reuses its static forward_order in reverse. Each backward pre-hook:

1. waits only for that unit's prefetched all-gather;
2. prefetches the next unit needed by backward on the all-gather stream;
3. allows that all-gather to overlap the current unit'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

• Experimental FSDP v2 numerical, accumulation, mixed-precision, storage-release, symmetric-memory, and memory tests pass with two CUDA ranks: 16 passed per rank.
• The profiler test covers one forward and backward pass. For four child FSDP units, 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.
• Formatting, import ordering, lint, byte-compilation, and diff checks pass.

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

[wujingyue]

Consider a linked hashmap.

[wujingyue]

This is a fragile assumption. Fine for the prototype.