Add forward all-gather overlap

[wujingyue]

Summary

Implements part of http://nv/mfsdp-schedule-design

• Add an FsdpContext that lazily owns a CUDA all-gather stream for one FSDP subtree.
• Run parameter unshard/all-gather on that context stream while compute remains on the default stream.
• Delay release of unsharded parameter storage so child-unit forward all-gathers can overlap default-stream GEMMs under a shared root context.
• Add context-sharing and memory-bound tests for nested FSDP units, plus a profiler test that asserts the forward child pipeline gives num_children - 1 all-gather/compute overlaps.

Scope

• This PR intentionally covers forward all-gather overlap only.
• Backward all-gathers are present in the profiled iteration, but they do not overlap compute in this PR because gradient reduction is still launched synchronously in post_backward before autograd reaches the next module’s pre_backward all-gather.
• The delayed gradient-reduction follow-up PR will address backward overlap.
• Reduce-scatter overlap is not included here.

Testing

• python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/megatron_fsdp/test_context.py tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py --tb=short --disable-warnings -rN
• git diff --check

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

Notes from the review meeting:

Even though the PR delayed releases, F_i still happens before AG_{i+1} on CPU. As a result, F_i's kernel(s) is often scheduled on GPU before AG_{i+1}'s ncclAllGather kernel. This may cause F_i's kernel(s) to occupy all the SMs before the ncclAllGather, preventing the overlap. In addition, the CPU overhead may delay the launch of the ncclAllGather, further reducing the chance of overlap.

I'll try to reproduce these issues by making the layers use more compute.

Potential options to alleviate this:

1. Prefetch AG_{i+1} before F_i. We'll have to capture and keep a global (per-context, more precisely) list of modules so we know what's the next module to run.
2. Make allgather SM free. This assumes certain hardware architectures that some customers might not have.
3. Green context so some SMs are dedicated for communication. But I'm told for Hopper these SMs can't be reused for compute when the communication is done.

cc @shjwudp, @Autumn1998 and @lhb8125

[wujingyue]

To follow up on #5513 (review), I prototyped and measured several approaches based on https://github.com/NVIDIA/Megatron-LM/pull/5513/changes#diff-91a39bc60a4a92f082088142637e0e73984537670406511f973df8e50355f9b8R314

Observations (Eos H100; dim=4096 fixed, M=8192 for the cited timing; M swept 4k–64k, bf16)

• The forward GEMM is a persistent cuBLAS nvjet_sm90 kernel: grid = 2×66 = 132 CTAs = all 132 SMs, register-bound (~64.5K/65.5K regs → 1 CTA/SM), one wave. Grid is 132 regardless of M or concurrency.
• The ring all-gather (ncclDevKernel_AllGather_RING_LL) launches 24 CTAs → needs 24 SMs, contends with the GEMM.
• MFSDP v2 already overlaps implicitly: AG_{i+1} runs on the separate all-gather stream with no dependency on GEMM_i, so they co-run to the extent SMs allow — no prefetch machinery required. But limited (hidden_ag ≈ 126µs of ~490µs AG at M=8192); overlap only where an AG is resident during a GEMM (GEMM launches into SMs the AG isn't using, or the AG slips into the GEMM's tail as CTAs retire). Caveat: SM residency was inferred from grids + overlapping windows, not measured — needs --gpu-metrics-devices=all / ncu.
• Explicit prefetch prototype (branch mfsdp-forward-prefetch; issue next unit's gather in pre_forward): numerically correct (9/9 tests) but did not improve wall-clock — it front-loads the gathers; the 24-SM contention, not launch order, is the ceiling. It used static module-tree order to pick the next unit, which is fragile (assumes forward execution order == registration order; breaks for reordered/reused modules or data-dependent control flow — v1 shares this via named_parameters; FSDP2 backward instead uses runtime-recorded order).
• CTA-zero (ProcessGroupNCCL cta_policy = NCCL_CTA_POLICY_ZERO + use_symm_mem=True) makes the forward all-gather SM-free — copy-engine cudaMemcpy (NCCL log AllGather [Copy Engine]; MEMCPY-only in nsys). Only change that improved forward wall-clock (M=8192, AG ≈ 0.4× GEMM).

External corroboration — PyTorch FSDP2

• FSDP2 uses implicit forward prefetch by default (separate all-gather stream, same as MFSDP v2 — no config).
• User-specified explicit prefetch (set_modules_to_forward_prefetch) is done if needed.

Recommendations

1. Leave the implicit stream overlap as-is — free, no next-module knowledge. Launch order / explicit prefetch is not the lever (GEMM is grid=132 either way).
2. Build SM-free comm (CTA-zero symmetric memory) — the one lever that matters. Turns the implicit overlap into real, full-GEMM overlap. Fully SM-free only within one NVLink domain (a node, or a GB200 NVL72 = 72 GPUs); across domains it's hierarchical (SM-free intra-domain, network/SM-based inter-domain) — pays off most when the parameter-shard group stays inside a domain (e.g., HSDP inner group).
3. Explicit forward prefetch is an optional additive follow-up — it's mainly for CPU-bound cases and/or issuing the first all-gather earlier. MFSDP v2's implicit prefetching can prefetch for 2+ layers by tuning target_length.
4. Run realistic workload -- to be confident about or to revisit these design decisions.

cc @shjwudp, @Autumn1998 and @lhb8125

[lhb8125]

To follow up on #5513 (review), I prototyped and measured several approaches based on https://github.com/NVIDIA/Megatron-LM/pull/5513/changes#diff-91a39bc60a4a92f082088142637e0e73984537670406511f973df8e50355f9b8R314

Observations (Eos H100; dim=4096 fixed, M=8192 for the cited timing; M swept 4k–64k, bf16)

• The forward GEMM is a persistent cuBLAS nvjet_sm90 kernel: grid = 2×66 = 132 CTAs = all 132 SMs, register-bound (~64.5K/65.5K regs → 1 CTA/SM), one wave. Grid is 132 regardless of M or concurrency.
• The ring all-gather (ncclDevKernel_AllGather_RING_LL) launches 24 CTAs → needs 24 SMs, contends with the GEMM.
• MFSDP v2 already overlaps implicitly: AG_{i+1} runs on the separate all-gather stream with no dependency on GEMM_i, so they co-run to the extent SMs allow — no prefetch machinery required. But limited (hidden_ag ≈ 126µs of ~490µs AG at M=8192); overlap only where an AG is resident during a GEMM (GEMM launches into SMs the AG isn't using, or the AG slips into the GEMM's tail as CTAs retire). Caveat: SM residency was inferred from grids + overlapping windows, not measured — needs --gpu-metrics-devices=all / ncu.
• Explicit prefetch prototype (branch mfsdp-forward-prefetch; issue next unit's gather in pre_forward): numerically correct (9/9 tests) but did not improve wall-clock — it front-loads the gathers; the 24-SM contention, not launch order, is the ceiling. It used static module-tree order to pick the next unit, which is fragile (assumes forward execution order == registration order; breaks for reordered/reused modules or data-dependent control flow — v1 shares this via named_parameters; FSDP2 backward instead uses runtime-recorded order).
• CTA-zero (ProcessGroupNCCL cta_policy = NCCL_CTA_POLICY_ZERO + use_symm_mem=True) makes the forward all-gather SM-free — copy-engine cudaMemcpy (NCCL log AllGather [Copy Engine]; MEMCPY-only in nsys). Only change that improved forward wall-clock (M=8192, AG ≈ 0.4× GEMM).

External corroboration — PyTorch FSDP2

• FSDP2 uses implicit forward prefetch by default (separate all-gather stream, same as MFSDP v2 — no config).
• User-specified explicit prefetch (set_modules_to_forward_prefetch) is done if needed.

Recommendations

1. Leave the implicit stream overlap as-is — free, no next-module knowledge. Launch order / explicit prefetch is not the lever (GEMM is grid=132 either way).
2. Build SM-free comm (CTA-zero symmetric memory) — the one lever that matters. Turns the implicit overlap into real, full-GEMM overlap. Fully SM-free only within one NVLink domain (a node, or a GB200 NVL72 = 72 GPUs); across domains it's hierarchical (SM-free intra-domain, network/SM-based inter-domain) — pays off most when the parameter-shard group stays inside a domain (e.g., HSDP inner group).
3. Explicit forward prefetch is an optional additive follow-up — it's mainly for CPU-bound cases and/or issuing the first all-gather earlier. MFSDP v2's implicit prefetching can prefetch for 2+ layers by tuning target_length.
4. Run realistic workload -- to be confident about or to revisit these design decisions.

cc @shjwudp, @Autumn1998 and @lhb8125

@wujingyue I am still concerned about the launch timing in CPU cases, which are almost everywhere in real workloads. Full-iter CG and SM-free comm are not validated as universal methodologies. Could you explain a little bit about the "prefetch for 2+ layers by tuning target_length"? Does it mean prefetching the i+2 layer when computing the i layer or prefetch 2x layers? Will it increase the exposed AG or memory footprint?

[wujingyue]

/ok to test 9f5989c

[wujingyue]

I am still concerned about the launch timing in CPU cases

Yes, I'm also concerned. Therefore, I plan to go simple first (this PR) and then revisit after #5614 bringing us some realistic workload. #5719 is the draft PR that implements explicit prefetching if you'd like to preview!

Could you explain a little bit about the "prefetch for 2+ layers by tuning target_length"?

Yes. http://nv/mfsdp-schedule-design => the "prefetching" section

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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