Generalized Tensor Parallelism (GTP)

tests/pytorch/test_distributed_weight.py

@@ -0,0 +1,124 @@
+# 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 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)
+
+
+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_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
+
+
+def test_forward_dispatches_single_weight():
+    """Linear (N=1): dispatcher delegates and returns a length-1 list."""
+    w = FakeDistributedWeight(group_size=1)
+    out = materialize_weight_for_forward(w)
+    assert isinstance(out, list) and len(out) == 1
+    assert torch.equal(out[0], torch.zeros(2, 2))
+    # The forward dispatcher owns the forward semantic; the leader is delegated to once.
+    assert w.calls == ["fwd"]
+
+
+def test_forward_dispatches_grouped_weights():
+    """GroupedLinear (N=k): one coalesced call, full list returned."""
+    w = FakeDistributedWeight(group_size=3)
+    out = materialize_weight_for_forward(w)
+    assert isinstance(out, list) and len(out) == 3
+    # Leader was called exactly once (coalesced), not once per weight.
+    assert len(w.calls) == 1
+
+
+def test_backward_dispatch_and_noop():
+    w = FakeDistributedWeight(group_size=2)
+    out = materialize_weight_for_backward(w)
+    assert len(out) == 2 and torch.equal(out[0], torch.full((2, 2), 10.0))
+
+    plain = torch.zeros(2)
+    assert materialize_weight_for_backward(plain) == [plain]
+
+
+def test_finalize_grads_dispatch_and_noop():
+    w = FakeDistributedWeight(group_size=1)
+    wgrads = [torch.zeros(2, 2)]
+    out = finalize_weight_grads(w, wgrads)
+    assert len(out) == 1 and torch.equal(out[0], torch.full((2, 2), 100.0))
+
+    # 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
+
+
+if __name__ == "__main__":
+    for name, fn in sorted(globals().items()):
+        if name.startswith("test_") and callable(fn):
+            fn()
+            print(f"{name}: PASS")

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