diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 91d3be63d0..5d767ba4d1 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -45,6 +45,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml $TE_PATH/tests/pytorch/test_multi_tensor.py || test_fail "test_multi_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py" @@ -66,6 +67,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml # Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "test_mhc.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_linear.xml $TE_PATH/tests/pytorch/test_grouped_linear.py || test_fail "test_grouped_linear.py" +PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_linear_distributed_weight.py || test_fail "test_ops_grouped_linear_distributed_weight.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py" if [ "$RET" -ne 0 ]; then diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py new file mode 100644 index 0000000000..69b5917d90 --- /dev/null +++ b/tests/pytorch/test_distributed_weight.py @@ -0,0 +1,172 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for the GTP-agnostic DistributedWeight protocol and dispatchers. + +These tests exercise TE's weight-parallelism extension point in isolation, with +a tiny in-repo ``FakeDistributedWeight`` stub standing in for a real implementer +(e.g. Megatron's GTPShardedParam). No GPU, process group, or Megatron import is +required: the whole contract lives in TE and is verified against the fake. +""" + +import pytest +import torch + +from transformer_engine.pytorch.distributed_weight import ( + DistributedWeight, + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) + + +class FakeDistributedWeight(torch.Tensor): + """Minimal DistributedWeight implementer for dispatcher tests. + + Records how it was called and returns marker tensors so the dispatcher's + behavior (delegation, list normalization, no-op fallback) is observable. + """ + + is_distributed_weight = True + + def __new__(cls, group_size=1): + t = torch.zeros(1).as_subclass(cls) + t.group_size = group_size + t.calls = [] + return t + + def materialize_group_for_forward(self): + self.calls.append("fwd") + out = [torch.full((2, 2), float(i)) for i in range(self.group_size)] + # Match the real GTP contract: single weight returns a bare tensor. + return out if self.group_size > 1 else out[0] + + def materialize_group_for_backward(self, **kwargs): + self.calls.append(("bwd", kwargs)) + out = [torch.full((2, 2), float(10 + i)) for i in range(self.group_size)] + return out if self.group_size > 1 else out[0] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls.append(("finalize", wgrads)) + wl = wgrads if isinstance(wgrads, (list, tuple)) else [wgrads] + out = [w + 100 for w in wl] + return out if self.group_size > 1 else out[0] + + def grad_buffer(self): + return torch.full((2, 2), -1.0) + + +class FakeNonTensorWeight: + """DistributedWeight-shaped object that is NOT a torch.Tensor (contract violation).""" + + is_distributed_weight = True + + def materialize_group_for_forward(self): + return torch.zeros(2, 2) + + def materialize_group_for_backward(self, **kwargs): + return torch.zeros(2, 2) + + def finalize_group_grads(self, wgrads, **kwargs): + return wgrads + + def grad_buffer(self): + return torch.zeros(2, 2) + + +def test_protocol_runtime_checkable(): + """A conforming object passes isinstance; a plain tensor does not.""" + assert isinstance(FakeDistributedWeight(), DistributedWeight) + assert not isinstance(torch.zeros(2), DistributedWeight) + + +def test_is_distributed_weight(): + assert is_distributed_weight(FakeDistributedWeight()) + assert not is_distributed_weight(torch.zeros(2)) + assert not is_distributed_weight(torch.nn.Parameter(torch.zeros(2))) + + +def test_non_tensor_implementer_rejected(): + """Implementers must be torch.Tensor subclasses; a non-Tensor fails loudly.""" + with pytest.raises(TypeError, match="torch.Tensor subclass"): + is_distributed_weight(FakeNonTensorWeight()) + + +def test_forward_noop_on_plain_tensor(): + """Plain weights pass through unchanged — the critical non-regression.""" + w = torch.nn.Parameter(torch.randn(4, 4)) + out = materialize_weight_for_forward(w) + assert out == [w] + assert out[0] is w + + +@pytest.mark.parametrize("group_size", [1, 3]) +def test_forward_dispatches(group_size): + """Linear (N=1) and GroupedLinear (N=k): one coalesced call, full list returned.""" + w = FakeDistributedWeight(group_size=group_size) + out = materialize_weight_for_forward(w) + assert isinstance(out, list) and len(out) == group_size + # Leader is delegated to exactly once (coalesced), not once per weight. + assert w.calls == ["fwd"] + + +def test_forward_accepts_weight_list(): + """The dispatcher accepts the full per-expert list; the leader (index 0) coalesces it.""" + leader = FakeDistributedWeight(group_size=3) + followers = [torch.zeros(2, 2), torch.zeros(2, 2)] + out = materialize_weight_for_forward([leader, *followers]) + assert isinstance(out, list) and len(out) == 3 + # Leader delegated exactly once; the follower entries are not materialized separately. + assert leader.calls == ["fwd"] + + +def test_forward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + out = materialize_weight_for_forward(ws) + assert out == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_backward_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + out = materialize_weight_for_backward(w) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_accepts_weight_list(): + """Backward dispatcher also accepts the full per-expert list; the leader coalesces it.""" + leader = FakeDistributedWeight(group_size=2) + out = materialize_weight_for_backward([leader, torch.zeros(2, 2)]) + assert isinstance(out, list) and len(out) == 2 + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_noop_on_plain_tensor(): + plain = torch.zeros(2) + assert materialize_weight_for_backward(plain) == [plain] + + +def test_backward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + assert materialize_weight_for_backward(ws) == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_finalize_grads_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + wgrads = [torch.zeros(2, 2) for _ in range(group_size)] + out = finalize_weight_grads(w, wgrads) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 100.0)) + + +def test_finalize_grads_noop_on_plain_tensor(): + """No-op path leaves the grads untouched.""" + plain_w = torch.nn.Parameter(torch.zeros(2)) + g = [torch.ones(2)] + assert finalize_weight_grads(plain_w, g) == g diff --git a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py new file mode 100644 index 0000000000..700be4138e --- /dev/null +++ b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py @@ -0,0 +1,131 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DistributedWeight dispatch in the fusible ``ops.GroupedLinear`` (unfused fallback path). + +TE ships no DistributedWeight implementer (GTP etc. live in the caller), so this validates the +*dispatch wiring* with an in-repo fake. The fake applies a DISTINCT, observable scale in each +materialize hook so every plumbing point is independently checked -- a wiring bug fails a specific +assertion: + + * ``materialize_group_for_forward`` scales the weights by ``FWD_SCALE`` -> the fwd GEMM must use + the materialized weights, so ``out == FWD_SCALE * plain_out``. + * ``materialize_group_for_backward`` scales by ``BWD_SCALE`` -> dgrad must use the + re-materialized weights, so ``dgrad == BWD_SCALE * plain_dgrad``. + * ``finalize_group_grads`` reduce-scatters the wgrads into ``main_grad`` in-place and returns a + dummy -> the real wgrad lands in ``main_grad`` and the ops path returns a throwaway ``.grad`` + (it discards finalize's return; see DistributedWeight.finalize_group_grads). + +Real all-gather / reduce-scatter math is exercised by the caller's distributed tests; here the fake +is single-process and only proves the ops.GroupedLinear integration routes weights/grads through +the hooks (and does not bypass them). +""" + +import pytest +import torch + +import transformer_engine.pytorch as te + +# Distinct powers of two so each scale commutes exactly through the (possibly TF32) GEMM rounding, +# making the linear scale relations bit-exact under a tight tolerance. +FWD_SCALE = 2.0 +BWD_SCALE = 4.0 + + +class _FakeDistWeight(torch.nn.Parameter): + """Single-process fake DistributedWeight leader for dispatch testing. + + Each materialize hook applies its own scale (see module docstring) so fwd/bwd routing is + observable; ``finalize_group_grads`` models the real main-grad contract -- reduce-scatter (here + an identity accumulate) the wgrads into each shard's ``main_grad`` in-place, flag + ``grad_added_to_main_grad``, and return a dummy that the ops path discards. + """ + + is_distributed_weight = True + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return [w * FWD_SCALE for w in self._group] + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return [w * BWD_SCALE for w in self._group] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls["finalize"] += 1 + wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] + 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 + + +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 + return leader + + +@pytest.mark.parametrize("num_gemms", [2, 4]) +def test_ops_grouped_linear_distributed_weight_dispatch(num_gemms): + """Every DistributedWeight hook must be routed through the GEMM flow (and not bypassed). + + fwd/bwd use the scaled materialized weights; finalize reduce-scatters the wgrad into + ``main_grad`` in-place, and the ops path returns a throwaway dummy ``.grad``. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.manual_seed(0) + # fp32 with power-of-two scales: each scale commutes exactly through the GEMM rounding (even + # TF32), so the linear scale relations below are bit-exact under a tight tolerance. + in_f, out_f, total_tokens = 32, 64, num_gemms * 8 + dtype, device = torch.float32, "cuda" + + op = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference.load_state_dict(op.state_dict()) + + leader = _make_fake_dist_leader(op, num_gemms) + 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 # DDP initializes this on every param + + 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) + ref_x = x.detach().clone().requires_grad_(True) + + out = op(x, split_sizes) + out.sum().backward() + ref_out = reference(ref_x, split_sizes) + ref_out.sum().backward() + + # All three dispatch hooks actually fired. + assert leader.calls["fwd"] > 0 and leader.calls["bwd"] > 0 and leader.calls["finalize"] > 0 + + tols = dict(rtol=1e-5, atol=1e-5) + # fwd used the materialized (FWD_SCALE) weights. + torch.testing.assert_close(out, FWD_SCALE * ref_out, **tols) + # dgrad used the re-materialized (BWD_SCALE) weights. + torch.testing.assert_close(x.grad, BWD_SCALE * ref_x.grad, **tols) + for i in range(num_gemms): + w = getattr(op, f"weight{i}") + ref_w = getattr(reference, f"weight{i}") + # The distributed op's grad of record is main_grad (finalize reduce-scattered the identity + # wgrad there); compare it to the plain reference module's ordinary autograd .grad. + torch.testing.assert_close(w.main_grad.to(dtype), ref_w.grad, **tols) + # main_grad is flagged, and op's own .grad is a discarded dummy (not the real grad, which + # 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" diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index c050f26869..82279a6257 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -6,7 +6,7 @@ from __future__ import annotations from collections.abc import Iterable -from contextlib import contextmanager, AbstractContextManager, ContextDecorator +from contextlib import contextmanager, AbstractContextManager, ContextDecorator, nullcontext from functools import lru_cache from dataclasses import dataclass import math @@ -926,7 +926,10 @@ def fork(self, name: str = "model-parallel-rng"): def reduce_scatter_along_first_dim( - inp: torch.Tensor, tp_group: dist_group_type, async_op: bool = False + inp: torch.Tensor, + tp_group: dist_group_type, + async_op: bool = False, + output: torch.Tensor = None, ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: """Reduce-scatter the input tensor across model parallel group.""" world_size = get_distributed_world_size(tp_group) @@ -944,7 +947,8 @@ def reduce_scatter_along_first_dim( dim_size[0] = dim_size[0] // world_size - output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) + if output is None: + output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) handle = torch.distributed.reduce_scatter_tensor( output, inp.contiguous(), group=tp_group, async_op=async_op ) @@ -1311,7 +1315,8 @@ def wait(self) -> None: """Wait for the async operation to complete and post-process the tensor.""" if self._synchronized: return - self.async_handle.wait() + if self.async_handle is not None: + self.async_handle.wait() _post_process_nvfp4_gather( self.output, self.columnwise_data_interleaved, @@ -1328,6 +1333,8 @@ def _all_gather_nvfp4( async_op: bool = False, quantizer: NVFP4Quantizer, out_shape: Optional[list[int]] = None, + output_tensor=None, + external_coalescing=False, ) -> tuple[NVFP4TensorStorage, Optional[torch.distributed.Work]]: """All-gather NVFP4 tensor along first dimension.""" @@ -1404,15 +1411,23 @@ def _all_gather_nvfp4( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct NVFP4 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) # Coalesce NCCL collectives for gathering data and scale inverses. - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as gather_coalescing_manager: + if not external_coalescing: + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather NVFP4 data for row-wise usage if quantizer.rowwise_usage: @@ -1493,10 +1508,10 @@ def _all_gather_nvfp4( # Transfer amax to output. out._amax_columnwise = inp._amax_columnwise - handle = gather_coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None # Fixes interleaved data for transposed tensor/scale inv and pads scale inv if needed. - if async_op and quantizer.columnwise_usage: + if (async_op or external_coalescing) and quantizer.columnwise_usage: handle = _NVFP4AllGatherAsyncHandle( out, out_columnwise_data, out_scale_inv, world_size, handle ) @@ -1513,6 +1528,8 @@ def _all_gather_mxfp8( async_op: bool = False, quantizer: MXFP8Quantizer, out_shape: Optional[list[int]] = None, + output_tensor: torch.Tensor = None, + external_coalescing: bool = False, ) -> tuple[MXFP8TensorStorage, Optional[torch.distributed.Work]]: """All-gather MXFP8 tensor along first dimension.""" @@ -1578,15 +1595,23 @@ def _all_gather_mxfp8( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct MXFP8 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - # Coalesce NCCL collectives - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as coalescing_manager: + if not external_coalescing: + # Coalesce NCCL collectives for gathering data and scale inverses. + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather MXFP8 data for row-wise usage if quantizer.rowwise_usage: @@ -1633,7 +1658,7 @@ def _all_gather_mxfp8( group=process_group, ) - handle = coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None return out, handle @@ -1642,9 +1667,17 @@ def gather_along_first_dim( process_group: dist_group_type, async_op: bool = False, quantizer: Optional[Quantizer] = None, + output_tensor: torch.Tensor = None, + external_coalescing: bool = False, ) -> tuple[torch.Tensor, Optional[torch.distributed.Work]]: """ All-gather tensors and concatenate along first dimension. + + ``external_coalescing``: composability flag for callers that batch several gathers into + one outer ``torch.distributed._coalescing_manager``. Coalescing managers cannot nest, so + when set this call skips opening its own manager and defers any post-gather fixup + (e.g. NVFP4 columnwise de-interleave) into the returned handle; ``handle.wait()`` completes + it once the outer manager has closed. Leave ``False`` for standalone gathers. """ # Return immediately if no communication is required @@ -1732,6 +1765,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + external_coalescing=external_coalescing, ) # NVFP4 case @@ -1746,6 +1781,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + external_coalescing=external_coalescing, ) # High-precision communication for quantized tensors @@ -1775,19 +1812,20 @@ def gather_along_first_dim( inp = inp.dequantize() # Communication for plain PyTorch tensors - out = torch.empty( - out_shape, - dtype=inp.dtype, - device=inp.device, - memory_format=torch.contiguous_format, - ) + if output_tensor is None: + output_tensor = torch.empty( + out_shape, + dtype=inp.dtype, + device=inp.device, + memory_format=torch.contiguous_format, + ) handle = torch.distributed.all_gather_into_tensor( - out, + output_tensor, inp.contiguous(), group=process_group, async_op=async_op, ) - return out, handle + return output_tensor, handle # Global cache to store symmetric memory tensors diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py new file mode 100644 index 0000000000..0d491430e5 --- /dev/null +++ b/transformer_engine/pytorch/distributed_weight.py @@ -0,0 +1,117 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""GTP-agnostic weight-parallelism extension point. + +TE owns this contract but ships no implementation; the caller (e.g. Megatron GTP) implements the +protocol on the weight and injects it at construction. Dispatchers are list-shaped (Linear -> 1, +GroupedLinear -> N; leader is ``weights[0]``) and no-op on plain tensors. +""" + +from typing import Any, List, Protocol, runtime_checkable + +import torch + +__all__ = [ + "DistributedWeight", + "is_distributed_weight", + "materialize_weight_for_forward", + "materialize_weight_for_backward", + "finalize_weight_grads", +] + + +@runtime_checkable +class DistributedWeight(Protocol): + """Structural interface for a custom-weight-parallel weight (AG for the GEMM, reduce/RS the + grad, re-materialize in backward). Duck-typed ``typing.Protocol``: implementers need not + subclass it, and all state (shards, group, async handles) lives outside TE on the implementer. + + Implementers MUST be ``torch.Tensor`` subclasses (needed by ``ctx.save_for_backward``, DDP + backward hooks, and ``torch.compile``); enforced at runtime by :func:`is_distributed_weight`. + """ + + # Capability marker: True on an implementer, absent on plain tensors; TE's fwd/bwd gate on it. + is_distributed_weight: bool + + def materialize_group_for_forward(self) -> Any: + """Return the tensor(s) to feed the forward GEMM (may all-gather shards).""" + + def materialize_group_for_backward(self) -> Any: + """Re-materialize the full weight(s) for the backward GEMMs.""" + + def finalize_group_grads(self, wgrads: Any) -> Any: + """Post-process freshly computed weight grad(s) (e.g. reduce-scatter). + + May consume ``wgrads`` in-place -- reduce-scatter into ``main_grad`` and set + ``grad_added_to_main_grad`` -- returning a dummy grad (or ``None`` for an async collective) + that callers use as the parameter grad(s) or discard. + """ + + def grad_buffer(self) -> torch.Tensor: + """The gradient accumulation buffer for this weight.""" + + +def is_distributed_weight(weight: Any) -> bool: + """True if ``weight`` participates in custom weight parallelism (False on plain tensors). + + Enforces the :class:`DistributedWeight` requirement that an implementer be a ``torch.Tensor`` + subclass, failing loudly here rather than silently breaking autograd downstream. + """ + flag = bool(getattr(weight, "is_distributed_weight", False)) + if flag and not isinstance(weight, torch.Tensor): + raise TypeError( + "DistributedWeight implementers must be torch.Tensor subclasses; got " + f"{type(weight).__name__}." + ) + return flag + + +def materialize_weight_for_forward(weights: Any) -> List[Any]: + """Prepare the weight(s) fed to the forward GEMM, always returned as a list. + + Args: + weights: the module's weight(s) -- a single weight (Linear) or the full per-expert list + (GroupedLinear). A bare weight is treated as a one-element list. + + Returns: + - Distributed group: the leader ``weights[0]`` all-gathers/coalesces the whole group and + returns all N materialized weights; the follower entries ``weights[1:]`` are ignored + here (the leader already holds references to its group). + - Otherwise: the input weights, unchanged. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_forward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def materialize_weight_for_backward(weights: Any) -> List[Any]: + """Backward-GEMM mirror of :func:`materialize_weight_for_forward` (same contract).""" + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_backward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def finalize_weight_grads(weights: Any, wgrads: List[Any]) -> List[Any]: + """Finalize a weight group's grad(s), mirroring :func:`materialize_weight_for_backward`. + + Delegates to the leader's :meth:`DistributedWeight.finalize_group_grads` (which defines the + in-place / dummy / async-``None`` return contract); returns ``wgrads`` unchanged when not + distributed. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(wgrads) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 82ee7953c1..3f264900c1 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -46,6 +46,12 @@ is_fp8_activation_recompute_enabled, in_fp8_activation_recompute_phase, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, @@ -457,6 +463,11 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad + origin_weights = weights + is_dist_weight = is_distributed_weight(weights[0]) + if is_dist_weight: + weights = materialize_weight_for_forward(weights) + # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): if FP8GlobalStateManager.get_fp8_recipe().custom(): @@ -663,6 +674,10 @@ def forward( if backward_override == "high_precision" and inp.requires_grad else [None] * num_gemms ) + if is_dist_weight: + # GTP: gathered workspace is transient (re-gathered in backward), don't save it. + weights_fp8 = [None] * num_gemms + saved_weights = origin_weights tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, *weights_fp8, @@ -690,6 +705,8 @@ def forward( if hasattr(weights[0], "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + elif is_dist_weight: + ctx.main_grad_funcs = [origin_weights[i].grad_buffer for i in range(num_gemms)] else: ctx.main_grad_funcs = [ lambda j=i: weights[j].main_grad for i in range(num_gemms) @@ -943,7 +960,12 @@ def backward( # Only needed when fuse_wgrad_accumulation is enabled. origin_weights = [None] * N main_grads = [None] * N - if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + is_dist_weight = is_distributed_weight(saved_weights[0]) + if is_dist_weight: + origin_weights = saved_weights + if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] + elif ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: origin_weight_refs = ctx.origin_weight_refs ctx.origin_weight_refs = None origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] @@ -1004,13 +1026,18 @@ def backward( ctx.m_splits, ) - if ctx.is_first_microbatch is not None: + if is_dist_weight: + accumulate_wgrad_into_param_main_grad = False + elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch ) else: accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation + if is_dist_weight: + weights = materialize_weight_for_backward(origin_weights) + if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD if ctx.fp8 or ctx.debug: @@ -1079,6 +1106,9 @@ def backward( device=ctx.device, ) wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] + if is_dist_weight: + # Gathered weights are no longer needed after dgrad GEMM. + del weights if ctx.save_original_input: inp = inputmats[0] @@ -1129,7 +1159,8 @@ def backward( use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + if not is_dist_weight + and not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), ) @@ -1171,10 +1202,13 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad - wgrad_list = [ - handle_custom_ddp_from_mcore(weight, main_grad, wgrad) - for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) - ] + if is_dist_weight: + wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) + else: + wgrad_list = [ + handle_custom_ddp_from_mcore(weight, main_grad, wgrad) + for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) + ] else: wgrad_list = [None] * ctx.num_gemms diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 43799b003c..cffc5ba1e6 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -56,6 +56,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing @@ -310,6 +316,11 @@ def forward( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + origin_weight = weight + is_dist_weight = is_distributed_weight(origin_weight) + if is_dist_weight: + weight = materialize_weight_for_forward(weight)[0] + out_features = weight.shape[0] new_weight_workspace = None weightmat = weight is_weight_param_quantized = False @@ -500,10 +511,15 @@ def forward( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # Distributed weight (e.g. GTP): don't save the gathered quantized workspace; + # backward re-gathers from the saved (sharded) weight and re-quantizes. + if is_dist_weight: + wt_save = None + tensors_to_save, tensor_objects = prepare_for_saving( inputmat, wt_save, - weight, + origin_weight, bias, ln_weight, ln_out_to_save, @@ -530,6 +546,8 @@ def forward( if hasattr(weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_func = weight.get_main_grad + elif is_dist_weight: + ctx.main_grad_func = origin_weight.grad_buffer else: ctx.main_grad_func = lambda: weight.main_grad ctx.grad_input_quantizer = grad_input_quantizer @@ -623,6 +641,9 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) + is_dist_weight = is_distributed_weight(saved_weight) + if is_dist_weight: + weight = materialize_weight_for_backward(saved_weight)[0] # Restore from weakref to get original weight python object # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) # Only needed when fuse_wgrad_accumulation is enabled. @@ -640,7 +661,7 @@ def backward( ), "weight was removed while fuse_wgrad_accumulation=True" # Since main_grad can be modified inplace, it should not be a part of saved_tensors main_grad = ctx.main_grad_func() if weight is not None else None - if main_grad is not None: + if main_grad is not None and not is_dist_weight: origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -984,7 +1005,10 @@ def backward( use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if ctx.is_first_microbatch is not None: + if is_dist_weight: + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch ) @@ -1056,6 +1080,9 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) + if is_dist_weight: + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 78a4d31852..cef6d3d533 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -55,6 +55,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_gemm, ) @@ -265,6 +271,7 @@ def _linear_forward_impl( """ weight = args.weight + is_dist_weight = is_distributed_weight(args.weight) inp = args.inp bias = args.bias input_quantizer = args.input_quantizer @@ -413,6 +420,14 @@ def _linear_forward_impl( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + # Distributed weight (e.g. GTP): rebind `weight` to the all-gathered tensor; + # `args.weight` keeps the sharded-param reference for backward re-gather / grad + # finalize. No-op for a plain weight. + if is_dist_weight: + weight = materialize_weight_for_forward(args.weight)[0] + # Refresh out_features from the gathered weight (captured sharded above, pre-gather). + out_features = weight.shape[0] + new_weight_workspace = None weightmat = weight if fp8 or debug: @@ -599,6 +614,9 @@ def _linear_forward_impl( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # Distributed weight (e.g. GTP): don't save the workspace; backward re-gathers it. + if is_dist_weight: + wt_save = None # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. @@ -704,6 +722,8 @@ def _linear_setup_ctx( bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) if hasattr(weight, "__fsdp_param__"): bwd_args.main_grad_func = weight.get_main_grad + elif is_distributed_weight(weight): + bwd_args.main_grad_func = weight.grad_buffer else: bwd_args.main_grad_func = lambda: weight.main_grad @@ -749,6 +769,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. inputmat = args.inputmat weight_fp8 = args.weight_fp8 saved_weight = args.saved_weight + is_dist_weight = is_distributed_weight(saved_weight) bias = args.bias input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer @@ -793,7 +814,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" main_grad = bwd_args.main_grad_func() - origin_weight_python_object.main_grad = main_grad + if not is_dist_weight: + origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed # NOTE: weight_fp8 = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already @@ -963,6 +985,12 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. dgrad = None dgrad_work = None + + # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when + # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd. + if is_dist_weight: + weight_fp8 = materialize_weight_for_backward(saved_weight)[0] + if bwd_args.requires_dgrad: # FSDP2: Re-create workspace from all-gathered weight when @@ -977,6 +1005,16 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. elif bwd_args.weight_quantizer is not None: bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) + elif ( + is_dist_weight + and bwd_args.fp8 + and bwd_args.weight_quantizer is not None + and not isinstance(weight_fp8, QuantizedTensorStorage) + ): + # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer + # so the dgrad operand isn't cast by the delayed recipe. + bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight_fp8 = bwd_args.weight_quantizer(weight_fp8) # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): @@ -1163,7 +1201,10 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. use_split_accumulator = bwd_args.wgrad_use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if bwd_args.is_first_microbatch is not None: + if is_dist_weight: + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif bwd_args.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( bwd_args.fuse_wgrad_accumulation and not bwd_args.is_first_microbatch ) @@ -1239,6 +1280,11 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output) + # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad + # (async; overlap with the next layer's bwd via the cascade). + if is_dist_weight: + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ @@ -1287,15 +1333,19 @@ def wgrad_gemm( origin_weight_python_object, "grad_added_to_main_grad" ): origin_weight_python_object.grad_added_to_main_grad = True + # Use the param's local shape (sharded under GTP) so the dummy wgrad + # matches the saved weight shape; main_grad_func() under GTP returns + # an unsharded scratch and would otherwise mismatch. + wgrad_shape = list(origin_weight_python_object.shape) if getattr(origin_weight_python_object, "zero_out_wgrad", False): wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, zero=True, ) else: wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, ) elif bwd_args.fuse_wgrad_accumulation: diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 5c96d4658e..d01bfbc35c 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -51,6 +51,12 @@ view_main_grad_as_grouped_buffer, ) from ..op import BasicOperation, OperationContext +from ...distributed_weight import ( + finalize_weight_grads, + is_distributed_weight, + materialize_weight_for_backward, + materialize_weight_for_forward, +) from ...tensor import GroupedTensor, GroupedTensorStorage from ...triton.grouped_dbias_dscales import ( compute_grouped_dbias, @@ -898,6 +904,26 @@ def _get_weight_tensors(self) -> list[torch.nn.Parameter]: return [self.weight] return [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + def _forward_weight_list(self) -> list[torch.Tensor]: + """Per-expert forward weights, materialized (all-gathered) when distributed.""" + weights = [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + if is_distributed_weight(weights[0]): + weights = materialize_weight_for_forward(weights) + return weights + + def _backward_weight_setup(self): + """Return ``(origin_weights, is_dist_weight, dgrad_weights)``; dgrad weights are the + re-materialized (all-gathered) weights when distributed, else ``None``.""" + origin_weights = self._get_weight_tensors() + is_dist_weight = is_distributed_weight(origin_weights[0]) + dgrad_weights = materialize_weight_for_backward(origin_weights) if is_dist_weight else None + return origin_weights, is_dist_weight, dgrad_weights + + def _is_distributed_weight(self) -> bool: + """Whether this op's weights are distributed (materialized per fwd/bwd, not saved).""" + leader = self.weight if self.single_grouped_weight else self.weight0 + return is_distributed_weight(leader) + def _get_grouped_bias_for_gemm( self, dtype: torch.dtype, @@ -1143,7 +1169,7 @@ def _fuser_forward_split_quantize( if weights is None: weights = self.weight.split_into_quantized_tensors() else: - weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] + weights = self._forward_weight_list() # materialized when distributed bs = None if has_bias: bs = self._get_bias_tensors(dtype) @@ -1198,7 +1224,8 @@ def _fuser_forward_split_quantize( out_splits[i].add_(bs[i].unsqueeze(0) * scales_splits[i].unsqueeze(-1)) # Prepare weight tensors for backward pass - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so we never save the gathered weight + if not input_requires_grad or self._is_distributed_weight(): ws = [None] * num_groups elif with_quantized_compute: for w, weight_param in zip(ws, weights): @@ -1290,7 +1317,7 @@ def _fuser_forward_grouped_tensor( else: # Discrete weights grouped_weights = self._get_discrete_weights_for_gemm( - [getattr(self, f"weight{idx}") for idx in range(num_groups)], + self._forward_weight_list(), weight_quantizers, columnwise_usage=input_requires_grad, with_quantized_compute=with_quantized_compute, @@ -1333,7 +1360,8 @@ def _fuser_forward_grouped_tensor( bias_scale=bias_scale, ) - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so never save the gathered weight. + if not input_requires_grad or self._is_distributed_weight(): grouped_weights = None if self.single_grouped_weight else [None] * num_groups if not weight_requires_grad: @@ -1393,7 +1421,7 @@ def _fuser_backward_split_quantize( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device # Saved tensors from forward pass. Layout: @@ -1445,7 +1473,8 @@ def _fuser_backward_split_quantize( grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] # Initialize grad weight buffers. - accumulate_into_main_grad = self._accumulate_into_main_grad + # Distributed weights reduce their own grads (finalize); never accumulate into main_grad. + accumulate_into_main_grad = self._accumulate_into_main_grad and not is_dist_weight grad_weights = [None] * num_groups final_weight_grads: list[Optional[torch.Tensor]] = ( [None] if self.single_grouped_weight else [None] * num_groups @@ -1495,7 +1524,7 @@ def _fuser_backward_split_quantize( device=device, ) general_grouped_gemm( - ws, + dist_dgrad_weights if is_dist_weight else ws, dys, [grad_input], [None] * num_groups, # quantization_params @@ -1540,11 +1569,16 @@ def _fuser_backward_split_quantize( if not delay_wgrad: clear_tensor_data(*xs) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + finalize_weight_grads(weights, grad_weights) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups @@ -1578,7 +1612,7 @@ def _fuser_backward_grouped_tensor( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device dtype = ctx.dtype @@ -1679,7 +1713,7 @@ def _fuser_backward_grouped_tensor( tensor_offsets=base_split_offsets * self.in_features, ) general_grouped_gemm_for_grouped_tensor( - ws, + dist_dgrad_weights if is_dist_weight else ws, grouped_dy, grouped_grad_input, layout="NN", @@ -1725,7 +1759,8 @@ def _fuser_backward_grouped_tensor( final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) wgrad_output = grouped_wgrad else: - if self._accumulate_into_main_grad: + # Distributed weights finalize wgrads (below); never accumulate into main_grad. + if self._accumulate_into_main_grad and not is_dist_weight: final_weight_grads = [ get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights ] @@ -1755,11 +1790,16 @@ def _fuser_backward_grouped_tensor( else: wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + finalize_weight_grads(weights, final_weight_grads) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index e0961044ad..c91f3d46bc 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -19,6 +19,12 @@ from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor +from ...distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ...module.base import _2X_ACC_WGRAD from ...quantization import Recipe from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer @@ -539,6 +545,7 @@ def _compute_grad_params( wgrad_output = None op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" weights = fc_op._get_weight_tensors() + is_dist_weight = is_distributed_weight(weights[0]) if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: @@ -571,7 +578,9 @@ def _compute_grad_params( else: w_list = [None] * num_groups if ctx.weight_requires_grad: - if fc_op._accumulate_into_main_grad: + # Distributed weight: the GEMM produces full-sized wgrads but main_grad is sharded, + # so use a full-sized scratch buffer (the reduce-scatter below lands it in main_grad). + if fc_op._accumulate_into_main_grad and not is_dist_weight: 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: @@ -587,6 +596,10 @@ def _compute_grad_params( if ctx.weight_requires_grad: # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() + if is_dist_weight and delay_wgrad: + raise RuntimeError( + "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False." + ) if cudnn_wgrad_kernel_fn is not None: offsets = offsets if offsets.dtype == torch.int32 else offsets.to(dtype=torch.int32) gemm_fn = functools.partial( @@ -625,9 +638,14 @@ def _compute_grad_params( fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) else: gemm_fn(grouped_x, grouped_dy, wgrad_output) + # Distributed weight: reduce-scatter the wgrads into main_grad. + # Return discarded (see finalize_weight_grads); dummy wgrads returned below. + if is_dist_weight: + finalize_weight_grads(weights, w_list) # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added - if fc_op._accumulate_into_main_grad: + # (wgrad fusion, or the distributed-weight reduce-scatter above) so it doesn't double-add. + if fc_op._accumulate_into_main_grad or is_dist_weight: w_list = get_dummy_wgrads_for_params(weights) elif delay_wgrad: w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups @@ -890,6 +908,19 @@ def fuser_forward( num_groups = fc1_op.num_groups fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + + # Distributed weight: expert weights are sharded 1/N along out_features; the fused kernels + # read the full shape, so all-gather the full weight first. A plain weight is a no-op. + fc1_is_dist = is_distributed_weight(fc1_weight_param) + fc2_is_dist = is_distributed_weight(fc2_weight_param) + assert fc1_is_dist == fc2_is_dist, "FC1/FC2 must share one distributed-weight group." + if fc1_is_dist: + assert ( + not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight + ), "distributed-weight fused grouped-MLP requires single_grouped_weight=False." + assert fc1_op.weight0.is_routed_expert and fc1_op.weight0.weight_list is not None + assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None + device = fc1_weight_param.device if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") @@ -978,6 +1009,8 @@ def fuser_forward( else: quantized_fc1_weights.append(weight) grouped_fc1_weight = quantized_fc1_weights + if fc1_is_dist: + grouped_fc1_weight = materialize_weight_for_forward(grouped_fc1_weight) # Prepare FC2 grouped weight tensor for fused kernels. if fc2_op.single_grouped_weight: @@ -1016,10 +1049,6 @@ def fuser_forward( grouped_fc1_weight, "_with_gemm_swizzled_scales" ): grouped_fc1_weight._with_gemm_swizzled_scales = False - if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( - grouped_fc2_weight, "_with_gemm_swizzled_scales" - ): - grouped_fc2_weight._with_gemm_swizzled_scales = False # Group-quantize input tensor and convert dtypes if needed fc1_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) @@ -1252,6 +1281,13 @@ def fuser_forward( else: fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) + if fc2_is_dist: + grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) + if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( + grouped_fc2_weight, "_with_gemm_swizzled_scales" + ): + grouped_fc2_weight._with_gemm_swizzled_scales = False + # Unpack kernel outputs # Note: Fused kernel outputs tensors with non-contiguous # logical dims. @@ -1478,6 +1514,11 @@ def fuser_forward( fc2_weight_tensors = ( [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight ) + # Save the joint op's internal layout; distributed weights save the small shards. + if fc1_is_dist: + fc1_weight_tensors = fc1_weights + if fc2_is_dist: + fc2_weight_tensors = fc2_weights fc1_ctx.save_for_backward( split_sizes, base_split_offsets, @@ -1741,6 +1782,10 @@ def fuser_backward( glu_clamp_min=self._cudnn_glu_clamp_min, ) + fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + if is_distributed_weight(fc2_leader): + grouped_fc2_weight = materialize_weight_for_backward(fc2_leader) + if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM fc2_weight_for_gemm = grouped_fc2_weight.copy() @@ -1965,6 +2010,10 @@ def fuser_backward( if fc1_ctx.input_requires_grad: in_shape = out_shape[:-1] + [fc1_weight_shape[1]] + fc1_leader = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 + if is_distributed_weight(fc1_leader): + grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) + if use_nvfp4: grad_input = torch.empty(in_shape, dtype=dtype, device=device) if num_groups == 1: