Generalized Tensor Parallelism (GTP)

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

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

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"