Repository navigation
Tied weights do not accumulate with delayed wgrad #3437
Description
Describe the bug
Transformer Engine does not accumulate the gradient of a tied weight when
delay_wgrad_compute=True and fuse_wgrad_accumulation=False.
Two tied delayed te.Linear modules produce a different shared-weight gradient
than two equivalent tied normal te.Linear modules. The equivalent non-tied
control case passes.
Steps/Code to reproduce bug
import torch
import transformer_engine.pytorch as te
from torch import nn
def make_model(delay_wgrad_compute: bool) -> nn.Sequential:
model = nn.Sequential(
te.Linear(
16, 16, bias=False, params_dtype=torch.bfloat16, device="cuda",
delay_wgrad_compute=delay_wgrad_compute,
fuse_wgrad_accumulation=False,
),
te.Linear(
16, 16, bias=False, params_dtype=torch.bfloat16, device="cuda",
delay_wgrad_compute=delay_wgrad_compute,
fuse_wgrad_accumulation=False,
),
)
model[1].weight = model[0].weight
return model
torch.manual_seed(9)
delayed = make_model(True)
regular = make_model(False)
regular.load_state_dict(delayed.state_dict())
x = torch.randn(2, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True)
delayed(x).float().sum().backward()
regular(x).float().sum().backward()
delayed[0].backward_dw()
delayed[1].backward_dw()
torch.testing.assert_close(delayed[0].weight.grad, regular[0].weight.grad)Expected behavior
The shared tied-weight gradient from the delayed model should match the shared
tied-weight gradient from the equivalent normal model.
Actual behavior
AssertionError: Tensor-likes are not close!
Mismatched elements: 255 / 256 (99.6%)
Greatest absolute difference: 0.396484375 at index (3, 7) (up to 1e-05 allowed)
Greatest relative difference: 22.875 at index (5, 12) (up to 0.016 allowed)
Environment details
- OS: Linux 6.8.0
- Python: 3.12.3
- PyTorch: 2.13.0a0+8145d630e8.nv26.06
- Transformer Engine: 2.17.1+4329ff84
- CUDA: 13.3
- GPU: NVIDIA RTX A6000
- Driver: 570.211.01
Additional context
This was found while adding MFSDP support for Transformer Engine delayed weight
gradients: NVIDIA/Megatron-LM#6697
Activity
migarci2 commentedon Aug 28, 2026
Thanks for the clear reproducer. I can take this. I’ll trace how backward_dw assigns non-fused weight grads for shared parameters, then add a regression covering tied delayed weights against the normal path and microbatch accumulation.
0z5a commentedon Sep 19, 2026
@wujingyue I’d be happy to take this one.
I’ll first reproduce it on current main, then keep the initial fix scoped to te.Linear with delay_wgrad_compute=True and fuse_wgrad_accumulation=False: tied/shared parameters should accumulate delayed wgrads rather than having a later backward_dw() overwrite the existing .grad.
For regression coverage I plan to include:
I can validate the BF16 path on L20s as well. I’ll keep any analogous GroupedLinear cleanup separate unless it turns out to be required for the same correctness contract.
If nobody is already working on this, I’d like to take it.
modelpath-dev commentedon Sep 20, 2026
It looks like the issue is with the gradient mismatch when using delayed weight gradients with tied weights. I'd start by checking how the gradient accumulation is handled in the
te.Linearmodules.wujingyue commentedon Sep 21, 2026
@migarci2, @0z5a and @modelpath-dev, sure -- thoughtful and clean PRs are always welcomed and appreciated!
0z5a commentedon Sep 21, 2026
Fixed in #3552.
backward_dw()assigned instead of accumulating, so with tied weights each module's pop overwrote the previous one — 0.6015× the regular path. The tied bias was worse: the second module's bgrad is ~0 and wiped a correct contribution (1.0984 vs 33.0859). Both now initialise on the first contribution andadd_afterwards.Ships with a regression test that uses only the public
te.LinearAPI: 5 failed / 3 passed without the fix, 8 passed with it. Reproduced on L20 (sm89) and RTX 5090 (sm120).@wujingyue