[Achyuthan-S]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Lets Transformer Engine modules built with fuse_wgrad_accumulation=True write their weight gradients directly into the experimental Megatron-FSDP (v2) reduce-scatter buffer through TE's weight.main_grad, removing the per-micro-batch gradient copy.

Linked issue: Fixes #6942

Summary

Today, after each module's backward, copy_gradients_to_partial_buffer copies every parameter's unsharded gradient into the reduce-scatter buffer. On H200 that copy is 1.2–3.9% of step time for a 4.54B-parameter Llama-shaped model (numbers in #6942). With this PR, TE writes the weight gradients of fused modules straight into the buffer.

Design

As agreed in #6942:

• Detection: a parameter group treats a 2-D weight as fused when its owning module has fuse_wgrad_accumulation=True. MFSDP follows the TE module's own setting; there is no new fully_shard argument.
• Buffer: when a module's own backward starts, its group allocates and zeroes the reduce-scatter buffer on the compute stream and sets weight.main_grad to each fused weight's view. TE accumulates into it and returns no gradient to autograd. Post-backward adds the remaining autograd gradients (biases, LayerNorm, unfused weights) into their views, removes main_grad, and hands the buffer to the reduce-scatter stream as before.
• Recompute: a checkpointed region is recomputed before its modules' backward. No buffer is prepared during recompute, so each one is allocated only when its own module's backward starts.
• Lifetime: the buffer is allocated on the compute stream, so it is released only after the compute stream's next wait on the reduce-scatter stream (the next group's reduction or the end of backward). Allocation and release stay on one stream without record_stream, and at most two fused buffers are alive at once. With delay_wgrad_compute=True, a buffer lives until backward_dw() has run and the next reduction has waited for it.
• TE contract: only the documented main_grad attribute is used, not __fsdp_param__ or is_distributed_weight.
• MXFP8: parameters with fusion raise a ValueError at fully_shard for now.

runtime_schedule.md gains a "Fused Weight Gradients" section describing this.

Scope and follow-ups

• The MCore adapter still rejects gradient_accumulation_fusion for v2 (mcore_fsdp_adapter.py), so this PR applies to fully_shard used directly with TE modules. Allowing it in the adapter, including the combined-1F1B path, is the next PR.
• overwrite_main_grad to skip the zeroing for weights with a single contribution per step.
• MXFP8 support.
• Requires PyTorch 2.8 or newer: the post-accumulate-grad hook must fire when TE returns no gradient, and 2.7's AccumulateGrad returns early on an undefined gradient.

Question for reviewers

Detection reads TE's fuse_wgrad_accumulation module attribute, a documented constructor argument that MCore also reads. The alternative is an explicit fully_shard argument. Happy to switch if you prefer that.

Testing

New tests in tests/unit_tests/distributed/mfsdp_v2/:

• test_fused_wgrad_losses_match_unfused: te.LayerNormLinear + te.Linear, fused against the copy path; losses and weights match over 5 steps. Variants: one unit per layer, one unit mixing fused and unfused weights, activation checkpointing (reentrant and non-reentrant), and HSDP; 1 and 3 micro-batches; set_to_none True and False; fp32 and bf16.
• test_fused_wgrad_main_grad_only_exists_during_backward, test_fused_wgrad_releases_buffers, test_fused_wgrad_recompute_prepares_each_buffer_at_its_backward and test_fused_wgrad_with_delayed_weight_gradient, plus unit tests for detection and the buffer.
• test_fused_wgrad_rejects_mxfp8_parameters (Blackwell only).

Results on H200 (NGC PyTorch 26.08, TE 2.18):

• 2 GPUs, test_fully_shard.py -k fused: 40 passed, 8 skipped (the HSDP variants need 4 ranks).
• 1 GPU, the whole mfsdp_v2 directory: no new failures. 7 tests fail on one rank, and the same 7 fail on main.
• Each test commit fails without the commit that follows it, and the recompute test fails if buffers are prepared during recompute.
• Not run here: the HSDP variants (4 GPUs) and the MXFP8 test (Blackwell).

Measurement

I'll post step time and peak memory before and after, on the #6942 setup with the adapter check removed locally, in a comment. The buffer still has to be zeroed each backward because TE accumulates into it, so I expect roughly half of the copy cost back.

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel, I have added or updated its bit-exact determinism test (no new kernel)
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter on my PR

cc @wujingyue

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

[Achyuthan-S]

@wujingyue small heads up,is there any issue or changes you need from this PR.? Or is this good enough.?

[wujingyue]

Sorry, this fell through. I didn't notice it was ready for review. Will take a look...

[wujingyue]

What's this for? IIRC, we have a DBuffer already for reduce scatter input. Is this the same but promoted to be a class attribute?

[Achyuthan-S]

Yes, it's the same reduce-scatter input: the same layout, Partial("avg") placements and grad_comm_dtype as allocate_partial_grad_buffer(), and it goes into the same reduce_partial_gradients(). What differs is when it exists. Today partial_grad is allocated in post-backward and filled by copying .grad, so it stays a local in _reduce_gradient_groups(). With fuse_wgrad_accumulation=True, TE writes the wgrad into weight.main_grad during the module's backward, so the buffer has to exist before that backward starts. prepare_fused_grad_buffer() allocates and zeroes it then and points each fused weight's main_grad at its view. take_fused_grad_buffer() hands it to the reduction in post-backward and resets the attribute to None. It lives on the group because it has to survive from the module's pre-backward to its post-backward. It is None outside that window and for groups without fused weights.

If it reads better, I can factor the DBuffer construction into one helper used by both paths.

[wujingyue]

Does this happen too late? IIRC, current waits for reduce-scatter at the end of the backward pass.

Also, why delaying release in the first place? We don't have to delay release today.

[Achyuthan-S]

Release isn't deferred to the end of backward. _reduce_gradient_groups() already makes the current stream wait for the reduce-scatter stream before packing each group; main issues that wait today. The fused path drops the previous module's buffers right after that same wait, so each buffer is freed at the next module's reduction. Only the last buffer of a backward waits for the end-of-backward wait. With one FSDP module per layer, at most two modules' buffers are alive: one being filled by TE and one being reduced. I'm merging current main into this branch and will push once the GPU tests pass. In that update the end-of-backward release moves into finish_grad_sync(), so it also happens with caller_managed_grad_sync=True.

Why delay at all: today partial_grad is allocated on the reduce-scatter stream right before the copy, so dropping it immediately is safe. The fused buffer has to exist when the module's backward starts, and TE writes it on the compute stream. If it were allocated on the reduce-scatter stream, the compute stream would have to wait for that stream before writing into it, as the copy path does before packing. At backward start, that wait would include the reduce-scatter the previous module just launched, so the module's backward would no longer overlap it. So the buffer is allocated on the compute stream. Dropping it right after launching its reduce-scatter would let the allocator hand the block to the next compute-stream allocation while the reduce-scatter is still reading it. That leaves two options: record_stream(), or holding the reference until the compute stream has waited. runtime_schedule.md avoids record_stream, so I hold the reference. This adds no wait that main doesn't already issue.

[wujingyue]

If it were allocated on the reduce-scatter stream, the compute stream would have to wait for that stream before writing into it

Yes but what's wrong with that? It sounds like a surgical change on top of the backward schedule in https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/distributed/fsdp/src/docs/runtime_schedule.md. I've yet to understand why it disables overlapping.

[Achyuthan-S]

Fair, I should have spelled that out. Going by the backward diagram:

Today AI_1 sits on the RS stream behind RS_2, but the only thing that waits on it is CI_1, and CI_1 runs after Bwd_1. So Bwd_1 overlaps RS_2.

With fusion there's no CI_1. TE writes into the buffer during Bwd_1, so Bwd_1 itself has to wait on AI_1. Allocating AI_1 when Bwd_1 starts doesn't help, because RS_2 was already launched at CI_2, so AI_1 still lands behind it on the RS stream. And it does need that wait: once DI_2 frees RS_2's input, the allocator can hand that same block to AI_1 (layers are usually the same size), and RS_2 may still be reading it. So Bwd_1 would sit there until RS_2 finishes. That's the overlap I meant.

Your idea does work if AI_1 moves one step earlier: allocate layer 1's buffer next to AI_2, before RS_2 is launched. Then the existing wait before CI_2 covers it, the buffer can be freed right after RS_1 like today, and peak memory is the same. The catch is that layer 2's post-backward has to know which module runs backward next. It would get that from the static backward order (same as prefetch), and it needs a fallback for the first module, the MTP repeat and delay_wgrad_compute.

I went with allocating at the module's own backward start so it doesn't depend on that order. Happy to switch to the one-ahead version if you'd rather keep it closer to the documented schedule. Let me know which you prefer.

[wujingyue]

Got it -- the extra copy-in (CI) kernels are critical to achieve today's overlapping. This is similar to how https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/distributed/fsdp/src/docs/runtime_schedule.md#fsdp2 achieves overlapping in forward with those copy-out kernels.

Happy to switch to the one-ahead version

Is it possible to do AI during prefetching? In addition to allgather next module's weights, pre-backward also pre-allocates its reduce-scatter input.

[Achyuthan-S]

Yes, I think that's the cleanest place for it. Module i's pre_backward already looks up the next module to prefetch its all-gather, so it can allocate that module's RS input on the RS stream right there. That's before RS_i is launched, so the wait in module i's post-backward (the one before CI_i) already covers it, and the buffer can be freed right after its own RS, like partial_grad today. That gets rid of the pending list in FsdpContext, and both paths can share the same allocation code, which also covers your first comment. I'd zero it on the RS stream too, which takes the memset off the compute stream.

Two details:

• I'd only pre-allocate for the immediate next module, even when the prefetch budget covers more, so peak memory stays at two buffers.
• So that correctness doesn't depend on the static order, I'd record an event after the allocation and have each module wait on it when its own backward starts. In the normal schedule the wait before CI_i has already covered it, so it costs nothing. It keeps things safe when the order is off and with delay_wgrad_compute. A module that didn't get a pre-allocated buffer (e.g. the first one in backward) allocates its own and waits. Anything pre-allocated for a module that never ran backward is dropped at the end of backward.

I'll push the merge with main first, then do this on top.

[wujingyue]

So that correctness doesn't depend on the static order,

We know static order can be problematic. The right fix should be something like #7998, rather than making every use case be defensive about it.

[Achyuthan-S]

Agreed, getting the order right belongs in #7998. I won't try to handle a wrong order here.

There's still one wait I need even when the order is right, though. backward_order puts a parent before its children (root first, then the last layer), and a parent's post-backward only runs after all of its children are done. So when the root pre-allocates the last layer's RS input, no post-backward wait comes between that allocation and the last layer's backward, and TE could start accumulating into the buffer before the zeroing on the RS stream has finished. delay_wgrad_compute has the same gap: the previous module's post-backward only runs after its backward_dw(), and the caller decides the order of the backward_dw() calls.

So I'd keep just a wait at the module's own backward start, on an event recorded after the allocation. That's the same thing _unshard_event already does for the prefetched all-gather. In the plain layer-to-layer case the wait before CI_i has already covered it, so it costs nothing there. I'll drop the fallback and the cleanup that were only there for a wrong order. The one exception is the root: it's first in backward_order and nothing prefetches for it, so if it owns fused weights itself it allocates its own buffer when its backward starts.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.