diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py index 69b5917d90..285d4e4d62 100644 --- a/tests/pytorch/test_distributed_weight.py +++ b/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 diff --git a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py index 700be4138e..e95806192b 100644 --- a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py +++ b/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" + ) diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py index 0d491430e5..b367444047 100644 --- a/transformer_engine/pytorch/distributed_weight.py +++ b/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`. diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 32963c07ca..d1cbc55822 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/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)] diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index af3f8b5930..debb88176b 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/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 ( @@ -1495,6 +1497,7 @@ def _fuser_backward_split_quantize( final_weight_grads: list[Optional[torch.Tensor]] = ( [None] if self.single_grouped_weight else [None] * num_groups ) + wgrad_dtype = weight_grad_dtype(weights, ctx.dtype) if ctx.weight_requires_grad: weight_shape = (self.out_features, self.in_features) grouped_shape = (num_groups, *weight_shape) @@ -1511,7 +1514,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=wgrad_dtype, device=device ) grad_weights = [final_weight_grads[0][idx] for idx in range(num_groups)] else: @@ -1521,12 +1524,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 @@ -1747,6 +1745,7 @@ def _fuser_backward_grouped_tensor( # Get the right wgrad buffers for grouped gemm. # Can be a GroupedTensor or list of tensors based on single_grouped_weight. + wgrad_dtype = weight_grad_dtype(weights, dtype) if ctx.weight_requires_grad: if self.single_grouped_weight: if self._accumulate_into_main_grad: @@ -1770,7 +1769,7 @@ def _fuser_backward_grouped_tensor( shapes=[weight_shape] * num_groups, quantizer=None, device=device, - dtype=dtype, + dtype=wgrad_dtype, ) final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) wgrad_output = grouped_wgrad @@ -1782,10 +1781,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 diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 83954a9b3d..dd82d533c6 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/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: