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

tests/pytorch/test_distributed_weight.py

@@ -19,6 +19,8 @@
     materialize_weight_for_forward,
     materialize_weight_for_backward,
     finalize_weight_grads,
+    weight_grad_buffers,
+    weight_grad_dtype,
 )
 
 
@@ -31,9 +33,11 @@ class FakeDistributedWeight(torch.Tensor):
 
     is_distributed_weight = True
 
-    def __new__(cls, group_size=1):
+    def __new__(cls, group_size=1, marker=-1.0):
         t = torch.zeros(1).as_subclass(cls)
         t.group_size = group_size
+        # Distinguishes this member's grad_buffer from its peers'.
+        t.marker = marker
         t.calls = []
         return t
 
@@ -55,7 +59,7 @@ def finalize_group_grads(self, wgrads, **kwargs):
         return out if self.group_size > 1 else out[0]
 
     def grad_buffer(self):
-        return torch.full((2, 2), -1.0)
+        return torch.full((2, 2), self.marker)
 
 
 class FakeNonTensorWeight:
@@ -170,3 +174,58 @@ def test_finalize_grads_noop_on_plain_tensor():
     plain_w = torch.nn.Parameter(torch.zeros(2))
     g = [torch.ones(2)]
     assert finalize_weight_grads(plain_w, g) == g
+
+
+@pytest.mark.parametrize("group_size", [1, 2])
+def test_weight_grad_buffers_dispatches(group_size):
+    """EVERY member supplies its own buffer -- the leader's is not reused for the group.
+
+    Distinct markers catch a ``[weights[0].grad_buffer()] * N`` implementation, which would hand
+    the whole group one buffer and let the per-expert GEMMs clobber each other.
+    """
+    group = [
+        FakeDistributedWeight(group_size=group_size, marker=float(i)) for i in range(group_size)
+    ]
+    out = weight_grad_buffers(group, (2, 2), torch.bfloat16, "cpu")
+    assert [b[0, 0].item() for b in out] == [float(i) for i in range(group_size)]
+
+
+def test_weight_grad_buffers_rejects_mismatched_buffer():
+    """A shard-shaped buffer would let the GEMM write past the end -- fail loudly instead."""
+    w = FakeDistributedWeight()
+    with pytest.raises(RuntimeError, match="wgrad GEMM needs"):
+        weight_grad_buffers([w], (4, 5), torch.bfloat16, "cpu")
+
+
+def test_weight_grad_buffers_allocates_for_plain_weights():
+    """Plain weights have no buffer of their own, so they get fresh scratch instead."""
+    plain = [torch.nn.Parameter(torch.zeros(2)) for _ in range(3)]
+    out = weight_grad_buffers(plain, (4, 5), torch.bfloat16, "cpu")
+    assert len(out) == 3
+    assert all(tuple(b.shape) == (4, 5) and b.dtype is torch.bfloat16 for b in out)
+
+
+@pytest.mark.parametrize("group_size", [1, 2])
+def test_weight_grad_dtype_dispatches(group_size):
+    """A distributed weight types its wgrad from main_grad, not from the compute dtype."""
+    w = FakeDistributedWeight(group_size=group_size)
+    w.main_grad = torch.zeros(2, 2, dtype=torch.float32)
+    assert weight_grad_dtype(w, torch.bfloat16) == torch.float32
+    assert weight_grad_dtype([w], torch.bfloat16) == torch.float32
+
+
+def test_weight_grad_dtype_without_main_grad():
+    """No main_grad yet (pre-DDP): fall back to the compute dtype rather than raising."""
+    w = FakeDistributedWeight()
+    assert weight_grad_dtype(w, torch.bfloat16) == torch.bfloat16
+
+
+def test_weight_grad_dtype_noop_on_plain_tensor():
+    """A plain weight keeps the compute dtype even when it carries an fp32 main_grad.
+
+    Without weight sharding the wgrad is either accumulated into main_grad by the GEMM or
+    returned as an ordinary .grad, so widening here would only cost memory.
+    """
+    plain_w = torch.nn.Parameter(torch.zeros(2))
+    plain_w.main_grad = torch.zeros(2, dtype=torch.float32)
+    assert weight_grad_dtype(plain_w, torch.bfloat16) == torch.bfloat16

tests/pytorch/test_ops_grouped_linear_distributed_weight.py

@@ -55,22 +55,38 @@ def materialize_group_for_backward(self, **kwargs):
     def finalize_group_grads(self, wgrads, **kwargs):
         self.calls["finalize"] += 1
         wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads]
+        # A real implementer reduce-scatters these as handed over, so their dtype is the
+        # reduction dtype -- record it before the cast below papers over it.
+        self.wgrad_dtypes = [g.dtype for g in wl]
         for w, g in zip(self._group, wl):
             w.main_grad.add_(g.to(w.main_grad.dtype))  # in-place accumulate into main_grad
             w.grad_added_to_main_grad = True
         return [torch.zeros_like(g) for g in wl]  # dummy grads (real value is now in main_grad)
 
     def grad_buffer(self):
-        return self.data
+        # Mirrors a real implementer (e.g. GTP's get_wgrad_tensor): a reusable buffer typed by
+        # main_grad, not by the weight -- the wgrad GEMM writes here, then finalize reduces it.
+        if getattr(self, "_wgrad_buf", None) is None:
+            self._wgrad_buf = torch.empty_like(self.data, dtype=self.main_grad.dtype)
+        return self._wgrad_buf
 
 
-def _make_fake_dist_leader(op, num_gemms):
-    """Replace ``op.weight0`` with a fake distributed leader referencing the whole group."""
-    w0 = op.weight0
-    leader = _FakeDistWeight(w0.data)
-    leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0}
-    leader._group = [leader] + [getattr(op, f"weight{i}") for i in range(1, num_gemms)]
-    op.weight0 = leader
+def _make_fake_dist_leader(op, num_gemms, cls=_FakeDistWeight):
+    """Replace the whole weight group with fakes and return the leader, ``op.weight0``.
+
+    Every member is a distributed weight (as in a real implementer, where each shard is one), so
+    per-weight protocol calls such as ``grad_buffer`` resolve on the followers too; the dispatchers
+    only ever route through the leader.
+    """
+    group = []
+    for i in range(num_gemms):
+        w = cls(getattr(op, f"weight{i}").data)
+        w.calls = {"fwd": 0, "bwd": 0, "finalize": 0}
+        w.wgrad_dtypes = []
+        setattr(op, f"weight{i}", w)
+        group.append(w)
+    leader = group[0]
+    leader._group = group
     return leader
 
 
@@ -129,3 +145,71 @@ def test_ops_grouped_linear_distributed_weight_dispatch(num_gemms):
         # lives in main_grad). Nobody writes the real grad into op.weight.grad by design.
         assert w.grad_added_to_main_grad is True
         assert w.grad is not None, f"weight{i} should receive a dummy .grad"
+
+
+class _FakeDistWeightPlain(_FakeDistWeight):
+    """``_FakeDistWeight`` without the observability scales.
+
+    ``module.GroupedLinear`` reads ``requires_grad`` off the MATERIALIZED weight, and a tensor
+    built inside ``autograd.Function.forward`` (grad mode off) has none -- so a scaling hook makes
+    it skip the wgrad GEMM entirely. Hand back the group itself to keep both paths in play.
+    """
+
+    def materialize_group_for_forward(self):
+        self.calls["fwd"] += 1
+        return list(self._group)
+
+    def materialize_group_for_backward(self, **kwargs):
+        self.calls["bwd"] += 1
+        return list(self._group)
+
+
+@pytest.mark.parametrize("num_gemms", [2, 4])
+@pytest.mark.parametrize("fused_ops", [True, False], ids=["ops", "module"])
+def test_grouped_linear_distributed_weight_wgrad_dtype(num_gemms, fused_ops):
+    """The wgrad handed to finalize must already carry ``main_grad``'s dtype.
+
+    A distributed weight reduces its wgrad across the sharding axis BEFORE accumulating, so a
+    compute-dtype wgrad rounds on every rank -- silently defeating fp32 grad accumulation. The
+    weight cannot accumulate into a sharded ``main_grad``, so the GEMM epilogue is the only place
+    to widen. Runs the module in bf16 while ``main_grad`` stays fp32, so the two dtypes differ.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+    torch.manual_seed(0)
+    in_f, out_f, total_tokens = 32, 64, num_gemms * 8
+    dtype, device = torch.bfloat16, "cuda"
+
+    if fused_ops:
+        op = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype)
+    else:
+        # fuse_wgrad_accumulation=False is the branch with no grad_buffer to take the dtype from.
+        op = te.GroupedLinear(
+            num_gemms,
+            in_f,
+            out_f,
+            bias=False,
+            device=device,
+            params_dtype=dtype,
+            fuse_wgrad_accumulation=False,
+        )
+
+    leader = _make_fake_dist_leader(op, num_gemms, cls=_FakeDistWeightPlain)
+    for i in range(num_gemms):
+        w = getattr(op, f"weight{i}")
+        w.main_grad = torch.zeros((out_f, in_f), dtype=torch.float32, device=device)
+        w.grad_added_to_main_grad = False
+
+    m_splits = [total_tokens // num_gemms] * num_gemms
+    m_splits[-1] += total_tokens - sum(m_splits)
+    split_sizes = torch.tensor(m_splits, dtype=torch.int64, device=device)
+
+    x = torch.randn(total_tokens, in_f, dtype=dtype, device=device, requires_grad=True)
+    op(x, split_sizes).sum().backward()
+
+    assert leader.calls["finalize"] > 0, "finalize_group_grads never fired"
+    assert leader.wgrad_dtypes, "no wgrads were handed to finalize"
+    assert set(leader.wgrad_dtypes) == {torch.float32}, (
+        f"wgrad reduced in {set(leader.wgrad_dtypes)}, expected main_grad's torch.float32; "
+        "the reduce-scatter would round on every rank"
+    )

transformer_engine/pytorch/distributed_weight.py

@@ -9,7 +9,7 @@
 GroupedLinear -> N; leader is ``weights[0]``) and no-op on plain tensors.
 """
 
-from typing import Any, List, Protocol, runtime_checkable
+from typing import Any, List, Protocol, Sequence, runtime_checkable
 
 import torch
 
@@ -19,6 +19,8 @@
     "materialize_weight_for_forward",
     "materialize_weight_for_backward",
     "finalize_weight_grads",
+    "weight_grad_buffers",
+    "weight_grad_dtype",
 ]
 
 
@@ -50,7 +52,12 @@ def finalize_group_grads(self, wgrads: Any) -> Any:
         """
 
     def grad_buffer(self) -> torch.Tensor:
-        """The gradient accumulation buffer for this weight."""
+        """Where the wgrad GEMM writes this weight's gradient.
+
+        The GEMM overwrites it and :meth:`finalize_group_grads` then reduces it, so it needs the
+        full unsharded weight shape and the dtype that reduction should use. Called on every
+        member of a group, unlike the group hooks above.
+        """
 
 
 def is_distributed_weight(weight: Any) -> bool:
@@ -101,6 +108,47 @@ def materialize_weight_for_backward(weights: Any) -> List[Any]:
     return list(weights)
 
 
+def weight_grad_buffers(
+    weights: Any, weight_shape: Sequence[int], compute_dtype: torch.dtype, device: torch.device
+) -> List[torch.Tensor]:
+    """Per-weight buffers for the wgrad GEMM to write into.
+
+    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)):
+        weights = [weights]
+    if is_distributed_weight(weights[0]):
+        buffers = [w.grad_buffer() for w in weights]
+        # Spot-check the leader: a shard-shaped buffer would let the GEMM write past the end.
+        if tuple(buffers[0].shape) != tuple(weight_shape):
+            raise RuntimeError(
+                f"grad_buffer() returned shape {tuple(buffers[0].shape)}; "
+                f"the wgrad GEMM needs {tuple(weight_shape)}."
+            )
+        return buffers
+    packed = torch.empty((len(weights), *weight_shape), dtype=compute_dtype, device=device)
+    return list(packed)
+
+
+def weight_grad_dtype(weights: Any, compute_dtype: torch.dtype) -> torch.dtype:
+    """Dtype for a wgrad buffer the caller allocates itself.
+
+    A distributed weight reduces its wgrad before accumulating, so the GEMM must already emit
+    ``main_grad``'s dtype -- otherwise the reduction rounds on every rank. Falls back to
+    ``compute_dtype`` for plain weights, and for an implementer whose ``main_grad`` the framework
+    has not attached yet.
+    """
+    if not isinstance(weights, (list, tuple)):
+        weights = [weights]
+    leader = weights[0]
+    if is_distributed_weight(leader):
+        main_grad = getattr(leader, "main_grad", None)
+        if main_grad is not None:
+            return main_grad.dtype
+    return compute_dtype
+
+
 def finalize_weight_grads(weights: Any, wgrads: List[Any]) -> List[Any]:
     """Finalize a weight group's grad(s), mirroring :func:`materialize_weight_for_backward`.
 

transformer_engine/pytorch/module/grouped_linear.py

@@ -51,6 +51,7 @@
     materialize_weight_for_forward,
     materialize_weight_for_backward,
     finalize_weight_grads,
+    weight_grad_dtype,
 )
 from ..cpp_extensions import (
     general_grouped_gemm,
@@ -1185,7 +1186,7 @@ def backward(
                     wgrad_packed = torch.empty(
                         ctx.num_gemms,
                         *weights[0].size(),
-                        dtype=ctx.activation_dtype,
+                        dtype=weight_grad_dtype(origin_weights, ctx.activation_dtype),
                         device=ctx.device,
                     )
                     wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)]

transformer_engine/pytorch/ops/basic/grouped_linear.py

@@ -57,6 +57,8 @@
     is_distributed_weight,
     materialize_weight_for_backward,
     materialize_weight_for_forward,
+    weight_grad_buffers,
+    weight_grad_dtype,
 )
 from ...tensor import GroupedTensor, GroupedTensorStorage
 from ...triton.grouped_dbias_dscales import (
@@ -1511,7 +1513,7 @@ def _fuser_backward_split_quantize(
                     accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0])
                 else:
                     final_weight_grads[0] = torch.empty(
-                        grouped_shape, dtype=ctx.dtype, device=device
+                        grouped_shape, dtype=weight_grad_dtype(weights, ctx.dtype), device=device
                     )
                     grad_weights = [final_weight_grads[0][idx] for idx in range(num_groups)]
             else:
@@ -1521,12 +1523,7 @@ def _fuser_backward_split_quantize(
                     ]
                     accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0])
                 else:
-                    grad_weights_packed = torch.empty(
-                        grouped_shape,
-                        dtype=ctx.dtype,
-                        device=device,
-                    )
-                    grad_weights = [grad_weights_packed[i] for i in range(num_groups)]
+                    grad_weights = weight_grad_buffers(weights, weight_shape, ctx.dtype, device)
                 final_weight_grads = list(grad_weights)
 
         # Perform dgrad GEMMs
@@ -1770,7 +1767,7 @@ def _fuser_backward_grouped_tensor(
                         shapes=[weight_shape] * num_groups,
                         quantizer=None,
                         device=device,
-                        dtype=dtype,
+                        dtype=weight_grad_dtype(weights, dtype),
                     )
                 final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape)
                 wgrad_output = grouped_wgrad
@@ -1782,10 +1779,7 @@ def _fuser_backward_grouped_tensor(
                     ]
                     accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0])
                 else:
-                    final_weight_grads = [
-                        torch.empty(weight_shape, dtype=dtype, device=device)
-                        for _ in range(num_groups)
-                    ]
+                    final_weight_grads = weight_grad_buffers(weights, weight_shape, dtype, device)
                 wgrad_output = final_weight_grads
 
         # wgrad GEMM

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -24,6 +24,8 @@
     materialize_weight_for_forward,
     materialize_weight_for_backward,
     finalize_weight_grads,
+    weight_grad_buffers,
+    weight_grad_dtype,
 )
 from ...module.base import _2X_ACC_WGRAD
 from ...quantization import Recipe
@@ -573,7 +575,7 @@ def _compute_grad_params(
                     shapes=[weight_shape] * num_groups,
                     quantizer=None,
                     device=device,
-                    dtype=dtype,
+                    dtype=weight_grad_dtype(weights, dtype),
                 )
             wgrad_output = grouped_wgrad
             w_list = [grouped_wgrad.rowwise_data.view(num_groups, *weight_shape)]
@@ -586,13 +588,7 @@ def _compute_grad_params(
                 w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights]
                 accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0])
             else:
-                wgrad_packed = torch.empty(
-                    num_groups,
-                    *weight_shape,
-                    dtype=dtype,
-                    device=device,
-                )
-                w_list = [wgrad_packed[i] for i in range(num_groups)]
+                w_list = weight_grad_buffers(weights, weight_shape, dtype, device)
             wgrad_output = w_list
 
     if ctx.weight_requires_grad: