Fuse Linear WGRAD accumulation across PP microbatches - #4386
Draft
anijain2305 wants to merge 1 commit into
Draft
Conversation
meta-cla
Bot
added
the
CLA Signed
anijain2305
commented
Aug 31, 2026
| "it mutates parameter.grad during backward" | ||
| ) | ||
| weight_NK = ( | ||
| self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight |
Contributor
Author
There was a problem hiding this comment.
dont support DTensor
Pipeline parallelism invokes backward separately for each microbatch. For BF16 Linear weights, each backward currently materializes a BF16 weight gradient. When FSDP uses an FP32 reduce dtype, it casts and accumulates these gradients into a separate FP32 buffer, requiring an additional accumulation kernel for each microbatch. Introduce a BF16 Linear implementation whose custom autograd.Function computes the first weight gradient directly in the configured reduce dtype and folds subsequent accumulation into the GEMM using addmm. This avoids the separate gradient-add kernels while preserving FSDP's reduction precision. When param_dtype is BF16 and reduce_dtype is FP32, the unsharded parameter remains BF16 while param.grad is FP32. This relies on **Tensor.grad_dtype**, which is set to the configured reduce dtype. When both dtypes are BF16, the weight gradient and its accumulation remain BF16. For subsequent backward calls, the autograd.Function reads the existing param.grad, accumulates into it, temporarily clears param.grad, and returns the same buffer. This prevents AccumulateGrad from adding the buffer to itself while still allowing autograd and FSDP hooks to observe a normal weight-gradient return value. FSDP remains responsible for the final reduce-scatter and sharded gradient.
anijain2305
force-pushed
the
fused-inplace-wgrad-accumulate
branch
from
August 31, 2026 22:12
73af371 to
1c7d7eb
Compare
anijain2305
commented
Aug 31, 2026
| weight_NK = ( | ||
| self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight | ||
| ) | ||
| weight_NK.grad_dtype = self.wgrad_accum_dtype |
Contributor
Author
There was a problem hiding this comment.
This is the KEY change.
anijain2305
commented
Aug 31, 2026
|
|
||
|
|
||
| def _get_local_wgrad(grad: torch.Tensor) -> torch.Tensor: | ||
| return grad.to_local() if isinstance(grad, DTensor) else grad |
Contributor
Author
There was a problem hiding this comment.
Do we need this?
tianyu-l
reviewed
Sep 1, 2026
| return grad_input, grad_weight_NK, None, None | ||
|
|
||
|
|
||
| class FusedWGradAccumLinear(Linear): |
Contributor
There was a problem hiding this comment.
Does / how does this compose with other override of Linear, including
- mxfp8 linear
- lora
- fp32 linear (https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/rl/models/cast_linear.py#L33)
- DistGEMM linear
I would recommend we first support mxfp8 linear and DistMoE with grad accum fusion "hardcoded" in those autograd functions, and not others.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Pipeline parallelism invokes backward separately for each microbatch. For BF16 Linear weights, each backward currently materializes a BF16 weight gradient. When FSDP uses an FP32 reduce dtype, it casts and accumulates these gradients into a separate FP32 buffer, requiring an additional accumulation kernel for each microbatch.
Introduce a BF16 Linear implementation whose custom autograd.Function computes the first weight gradient directly in the configured reduce dtype and folds subsequent accumulation into the GEMM using addmm. This avoids the separate gradient-add kernels while preserving FSDP's reduction precision.
When param_dtype is BF16 and reduce_dtype is FP32, the unsharded parameter remains BF16 while param.grad is FP32. This relies on Tensor.grad_dtype, which is set to the configured reduce dtype. When both dtypes are BF16, the weight gradient and its accumulation remain BF16.
For subsequent backward calls, the autograd.Function reads the existing param.grad, accumulates into it, temporarily clears param.grad, and returns the same buffer. This prevents AccumulateGrad from adding the buffer to itself while still allowing autograd and FSDP hooks to observe a normal weight-gradient return value.
FSDP remains responsible for the final reduce-scatter and sharded gradient.