[PyTorch] Type the grouped wgrad from main_grad for distributed weights by fanshiqing · Pull Request #3397 · NVIDIA/TransformerEngine

fanshiqing · 2026-08-18T16:05:30Z

Description

c9a1f5a routed distributed weights into the no-fusion branch, which allocates at the compute dtype -- so under fp32 grad accumulation the reduce-scatter rounded to BF16 on every rank:

  wgrad alloc fp32 ---> reduce fp32 ---> accum fp32     <- this fix
  (was compute dtype ---> compute dtype ---> fp32)     <- before fix

Two dispatchers beside finalize_weight_grads:

weight_grad_buffers() GEMM writes into the implementer's buffer (no per-backward wgrad allocation)
weight_grad_dtype() fallback for the single-grouped branchesc9a1f5a9 routed distributed weights into the no-fusion branch, which allocates at the compute dtype -- so under fp32 grad accumulation the reduce-scatter rounded to BF16 on every rank:

  wgrad alloc fp32 ---> reduce fp32 ---> accum fp32     <- this fix
  (was compute dtype ---> compute dtype ---> fp32)     <- before fix

Two dispatchers beside finalize_weight_grads:

weight_grad_buffers() GEMM writes into the implementer's buffer (no per-backward wgrad allocation)
weight_grad_dtype() fallback for the single-grouped branches

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

distributed_weight.py — new weight_grad_buffers() + weight_grad_dtype() dispatchers
ops/basic/grouped_linear.py — both backwards use them
ops/fused/grouped_mlp.py — same
module/grouped_linear.py — dtype fix for the no-fusion branch
test_distributed_weight.py — unit tests for the two dispatchers
test_ops_grouped_linear_distributed_weight.py — wgrad-dtype test, ops + module

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

c9a1f5a routed distributed weights into the no-fusion branch, which allocates at the compute dtype -- so under fp32 grad accumulation the reduce-scatter rounded to BF16 on every rank: wgrad alloc fp32 ---> reduce fp32 ---> accum fp32 <- this fix (was compute dtype ---> compute dtype ---> fp32) Two dispatchers beside finalize_weight_grads: weight_grad_buffers() GEMM writes into the implementer's buffer (no per-backward wgrad allocation) weight_grad_dtype() fallback for the single-grouped branches Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

greptile-apps · 2026-08-18T16:12:57Z

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new actionable issue remains.

Summary

This PR makes grouped weight gradients use main_grad.dtype for distributed weights, avoiding BF16 rounding before reduction.

Distributed grouped gradients use the accumulation dtype before reduction.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Original weight group] --> B{Distributed leader with main_grad?}
    B -->|Yes| C[Allocate wgrad at main_grad dtype]
    B -->|No| D[Allocate wgrad at compute dtype]
    C --> E[GEMM writes weight gradients]
    D --> E
    E --> F[Finalize distributed gradients]

_{Reviews (3) · Last reviewed commit: "[PyTorch] Drop weight_grad_buffers; type..." · Reviewed by Greptile}

zhongbozhu · 2026-08-19T04:38:28Z

/te-ci pytorch L1

vthumbe1503 · 2026-08-21T01:55:19Z

@@ -656,7 +575,7 @@ def _compute_grad_params(
                    shapes=[weight_shape] * num_groups,
                    quantizer=None,
                    device=device,
-                    dtype=dtype,
+                    dtype=weight_grad_dtype(weights, dtype),


I dont understand how this will help, since we will never enter this else condition of allocating using torch.empty, if main_grad is attached to the parameter.

Actually we will, if accumulate_into_main_grad is False, but why are we attaching main_grad to parameter in that case in Megatron?

Hi @vthumbe1503 , for you first question why this line's change is required: we need to follow main_grad's dtype instead of the compute_dtype for this scratch_buffer allocation.

Details:

The bug: for a GTP-sharded (distributed) weight whose scratch wgrad buffer isn't fused into main_grad, the buffer was allocated in the forward/backward compute dtype (e.g. bf16) — not in main_grad's dtype. That buffer is what gets reduced across the weight's GTP group before being accumulated into main_grad.

Result: whenever main_grad was fp32 (fp32 grad accumulation enabled by set --grad-reduce-in-fp32 from MLM for a best accuracy loss baseline), the reduce-scatter still ran in bf16 — every rank's contribution got rounded to bf16 before it ever reached the fp32 accumulator. fp32 grad accumulation was silently defeated for GTP-sharded weights, with no error or warning.

The fix: weight_grad_dtype(weights, dtype) returns main_grad.dtype instead of the plain compute dtype whenever the weight is distributed and main_grad is already attached, so the scratch buffer — and therefore the reduction — happens at the same precision main_grad actually uses. Falls back to compute_dtype for plain weights or when main_grad isn't attached yet.

if accumulate_into_main_grad is False, but why are we attaching main_grad to parameter in that case in Megatron

for MLM:

.main_grad attachment is unconditional.

accumulate_into_main_grad is a separate, opt-in performance flag.

Thanks for the explaination @fanshiqing! Why is this problem not present for discrete weights? Is it because in case discrete weights, we always write directly into the main_grad buffer?

Hi @vthumbe1503 , Mostly yes — what matters is whether the reduction runs before or after the wgrad lands in main_grad:

Regular weights: the GEMM writes into the fp32 main_grad (or DDP adds .grad into it), and DDP reduces later, from main_grad, so the sum is in fp32.

Distributed weights: main_grad is only the local shard, but the GEMM produces the full unsharded wgrad. It must be reduce-scattered first, from a scratch buffer.

The bug: that scratch buffer used the compute dtype, so the RS summed in BF16 even under fp32 grad accumulation.

The fix: the scratch buffer now takes main_grad's dtype, so the RS runs in fp32.

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

zhongbozhu · 2026-08-26T05:21:51Z

@@ -101,6 +108,47 @@ def materialize_weight_for_backward(weights: Any) -> List[Any]:
    return list(weights)


+def weight_grad_buffers(


can we get rid of this helper function? doesn't feel necessary

zhongbozhu · 2026-08-26T05:22:45Z

+    A distributed weight brings its own, which skips this allocation and carries ``main_grad``'s
+    dtype by construction; anything else gets fresh scratch in the compute dtype.
+    """
+    if not isinstance(weights, (list, tuple)):


this doesn't seem to support single weight? TE as a computing library supports single weight for all, so it's orthogonal to if a certain feature in Mcore doesn't consider single weight yet

It does handle a single weight: a bare tensor is wrapped, and single_grouped_weight passes [self.weight], whose main_grad dtype is used.

zhongbozhu · 2026-08-26T05:23:28Z

+    return list(packed)
+
+
+def weight_grad_dtype(weights: Any, compute_dtype: torch.dtype) -> torch.dtype:


we can also get rid of this helper?

No, I'd keep this one. It's the actual fix, and without it the same condition gets copied into 8 places

It's the whole fix now that weight_grad_buffers is gone.

Inlining it would copy the same condition into 8 call sites.

It keeps Megatron's main_grad convention in distributed_weight.py, not spread across TE's ops code.

It handles a missing main_grad by falling back to the compute dtype. Inlined, each site would need that guard or could raise AttributeError.

…_reduce_in_fp32 Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

…a weight_grad_dtype only Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

fanshiqing requested review from ksivaman and timmoon10 as code owners August 18, 2026 16:05

github-actions Bot added the org-contribution label Aug 18, 2026

fanshiqing linked an issue Aug 18, 2026 that may be closed by this pull request

GTP+TE integration NVIDIA/Megatron-LM#6619

Closed

greptile-apps Bot reviewed Aug 18, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/distributed_weight.py Outdated

vthumbe1503 reviewed Aug 21, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/basic/grouped_linear.py Outdated

vthumbe1503 reviewed Aug 21, 2026

View reviewed changes

fix comments

c73b296

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

zhongbozhu reviewed Aug 26, 2026

View reviewed changes

vthumbe1503 requested changes Sep 30, 2026

View reviewed changes

fanshiqing added 2 commits October 15, 2026 22:58

Merge remote-tracking branch 'origin/main' into shiqingf/fix_gtp_grad…

f5d8e77

…_reduce_in_fp32 Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

[PyTorch] Drop weight_grad_buffers; type distributed-weight wgrads vi…

1b2a404

…a weight_grad_dtype only Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

		@@ -101,6 +108,47 @@ def materialize_weight_for_backward(weights: Any) -> List[Any]:
		return list(weights)


		def weight_grad_buffers(

		return list(packed)


		def weight_grad_dtype(weights: Any, compute_dtype: torch.dtype) -> torch.dtype:

Conversation

fanshiqing commented Aug 18, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

zhongbozhu commented Aug 19, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fanshiqing Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

greptile-apps Bot commented Aug 18, 2026 •

edited

Loading

vthumbe1503 Aug 21, 2026 •

edited

Loading

fanshiqing Aug 24, 2026 •

edited

Loading

vthumbe1503 Sep 16, 2026 •

edited

Loading