[PyTorch] Type the grouped wgrad from main_grad for distributed weights

[fanshiqing]

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

[greptile-apps]

Greptile Summary

The PR routes grouped distributed-weight gradients through implementer-provided buffers or allocates them using main_grad’s dtype to preserve FP32 reduction precision. It updates module, basic-op, and fused-MLP backward paths and adds dispatcher and dtype tests, but the new per-member buffer call breaks groups whose distributed protocol is carried only by the leader.

Confidence Score: 4/5

The PR should not merge until grouped backward remains compatible with distributed groups represented by a protocol-bearing leader and plain followers.

The new buffer dispatcher checks only the leader for distributed capability but calls grad_buffer on every group member, causing an immediate backward exception for an input shape already accepted by the leader-based dispatcher contract.

Files Needing Attention: transformer_engine/pytorch/distributed_weight.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed_weight.py	Adds wgrad buffer and dtype dispatchers; per-member grad_buffer invocation is incompatible with leader-only distributed groups.
transformer_engine/pytorch/module/grouped_linear.py	Types no-fusion grouped wgrad scratch from the origin distributed weight’s main_grad.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Routes discrete grouped wgrad allocations through the new dispatchers in both backward implementations.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Applies the same distributed wgrad dtype and buffer selection to fused grouped-MLP backward.
tests/pytorch/test_distributed_weight.py	Adds unit coverage for dispatcher dtype, allocation, shape, and per-member behavior but not leader-only heterogeneous groups.
tests/pytorch/test_ops_grouped_linear_distributed_weight.py	Adds BF16-compute/FP32-main-grad integration coverage using groups where every member implements the distributed protocol.

Sequence Diagram

sequenceDiagram
  participant B as Grouped backward
  participant D as weight_grad_buffers
  participant L as Distributed leader
  participant F as Plain follower
  B->>D: "weights = [leader, follower]"
  D->>L: is_distributed_weight(leader)
  L-->>D: true
  D->>L: grad_buffer()
  L-->>D: buffer
  D->>F: grad_buffer()
  F--xD: AttributeError

Loading

_{Reviews (1): Last reviewed commit: "[PyTorch] Type the grouped wgrad from ma..." | Re-trigger Greptile}

[greptile-apps]

Leader-only groups break backward

When a distributed weight group uses a protocol-bearing leader with ordinary tensor followers, weight_grad_buffers checks only the leader but calls grad_buffer() on every member, causing backward to fail with AttributeError.

Knowledge Base Used: PyTorch Distributed/Parallel Training Support

[zhongbozhu]

/te-ci pytorch L1

[vthumbe1503]

Could we move the wgrad_dtype directly in the place where this is used similar to the grouped_mlp file?

Given that wgrad_dtype is only used only for the single_grouped_weight case, it would increase code readability if we directly move it inside that branch

[vthumbe1503]

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.

[vthumbe1503]

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