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

tests/pytorch/test_distributed_weight.py

@@ -19,6 +19,7 @@
     materialize_weight_for_forward,
     materialize_weight_for_backward,
     finalize_weight_grads,
+    weight_grad_dtype,
 )
 
 
@@ -170,3 +171,29 @@ 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_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,6 +55,9 @@ 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
@@ -64,11 +67,12 @@ def grad_buffer(self):
         return self.data
 
 
-def _make_fake_dist_leader(op, num_gemms):
+def _make_fake_dist_leader(op, num_gemms, cls=_FakeDistWeight):
     """Replace ``op.weight0`` with a fake distributed leader referencing the whole group."""
     w0 = op.weight0
-    leader = _FakeDistWeight(w0.data)
+    leader = cls(w0.data)
     leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0}
+    leader.wgrad_dtypes = []
     leader._group = [leader] + [getattr(op, f"weight{i}") for i in range(1, num_gemms)]
     op.weight0 = leader
     return leader
@@ -129,3 +133,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

@@ -19,6 +19,7 @@
     "materialize_weight_for_forward",
     "materialize_weight_for_backward",
     "finalize_weight_grads",
+    "weight_grad_dtype",
 ]
 
 
@@ -101,6 +102,24 @@ def materialize_weight_for_backward(weights: Any) -> List[Any]:
     return list(weights)
 
 
+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

@@ -57,6 +57,7 @@
     materialize_weight_for_forward,
     materialize_weight_for_backward,
     finalize_weight_grads,
+    weight_grad_dtype,
 )
 from ..cpp_extensions import (
     general_grouped_gemm,
@@ -655,7 +656,7 @@ def _grouped_linear_backward_impl(
                 num_gemms,
                 out_features,
                 in_features,
-                dtype=args.activation_dtype,
+                dtype=weight_grad_dtype(origin_weights, args.activation_dtype),
                 device=device,
             )
             wgrad_list = [wgrad_packed[i] for i in range(num_gemms)]
@@ -1188,13 +1189,14 @@ def _grouped_linear_graph_safe_backward(
                 wgrad_output = main_grads
                 wgrad_list = main_grads
         else:
+            wgrad_dtype = weight_grad_dtype(origin_weights, args.activation_dtype)
             if args.single_grouped_weight:
                 grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes(
                     num_tensors=N,
                     shapes=[(args.out_features, args.in_features)] * N,
                     quantizer=None,
                     device=grad_output.device,
-                    dtype=args.activation_dtype,
+                    dtype=wgrad_dtype,
                 )
                 wgrad_output = grouped_wgrad
                 wgrad_list = [
@@ -1205,7 +1207,7 @@ def _grouped_linear_graph_safe_backward(
                     N,
                     args.out_features,
                     args.in_features,
-                    dtype=args.activation_dtype,
+                    dtype=wgrad_dtype,
                     device=grad_output.device,
                 )
                 wgrad_output = [wgrad_packed[i] for i in range(N)]

transformer_engine/pytorch/ops/basic/grouped_linear.py

@@ -59,6 +59,7 @@
     is_distributed_weight,
     materialize_weight_for_backward,
     materialize_weight_for_forward,
+    weight_grad_dtype,
 )
 from ...tensor import GroupedTensor, GroupedTensorStorage
 from ...triton.grouped_dbias_dscales import (
@@ -1620,7 +1621,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:
@@ -1632,7 +1633,7 @@ def _fuser_backward_split_quantize(
                 else:
                     grad_weights_packed = torch.empty(
                         grouped_shape,
-                        dtype=ctx.dtype,
+                        dtype=weight_grad_dtype(weights, ctx.dtype),
                         device=device,
                     )
                     grad_weights = [grad_weights_packed[i] for i in range(num_groups)]
@@ -1914,7 +1915,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
@@ -1927,7 +1928,9 @@ 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)
+                        torch.empty(
+                            weight_shape, dtype=weight_grad_dtype(weights, dtype), device=device
+                        )
                         for _ in range(num_groups)
                     ]
                 wgrad_output = final_weight_grads

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -28,6 +28,7 @@
     materialize_weight_for_forward,
     materialize_weight_for_backward,
     finalize_weight_grads,
+    weight_grad_dtype,
 )
 from ...module.base import _2X_ACC_WGRAD
 from ...quantization import Recipe, get_fp8_torch_dtype
@@ -873,7 +874,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)]
@@ -889,7 +890,7 @@ def _compute_grad_params(
                 wgrad_packed = torch.empty(
                     num_groups,
                     *weight_shape,
-                    dtype=dtype,
+                    dtype=weight_grad_dtype(weights, dtype),
                     device=device,
                 )
                 w_list = [wgrad_packed[i] for i in range(num_groups)]