diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py index 69b5917d90..5c13d01de9 100644 --- a/tests/pytorch/test_distributed_weight.py +++ b/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 diff --git a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py index 700be4138e..381f8cece5 100644 --- a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py +++ b/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" + ) diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py index 0d491430e5..0f8f2a428a 100644 --- a/transformer_engine/pytorch/distributed_weight.py +++ b/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`. diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 828ec74d9b..6dcb0879dd 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/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)] diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index b4810f6cc8..c2aea81bd2 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/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 diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index eeb5f43335..46763a6629 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/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)]