GroupedTensor Machinery for MXFP8 Dispatch input passed from Megatron

tests/pytorch/test_grouped_tensor.py

@@ -304,6 +304,127 @@ def setup_class(cls) -> None:
         torch.manual_seed(seed)
         torch.cuda.manual_seed(seed)
 
+    def test_autograd_identity_supports_varying_shapes(self) -> None:
+        """An identity autograd node should preserve grouped storage and metadata."""
+
+        class Identity(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, tensor):
+                return tensor
+
+            @staticmethod
+            def backward(ctx, grad_output):
+                return grad_output
+
+        first_dims = torch.tensor([1, 2], dtype=torch.int64, device="cuda")
+        tensor_offsets = torch.tensor([0, 4, 12], dtype=torch.int64, device="cuda")
+        grouped_tensor = GroupedTensor(
+            shape=(3, 4),
+            dtype=torch.float32,
+            num_tensors=2,
+            data=torch.randn(12, device="cuda"),
+            first_dims=first_dims,
+            tensor_offsets=tensor_offsets,
+            requires_grad=True,
+        )
+
+        result = Identity.apply(grouped_tensor)
+
+        assert isinstance(result, GroupedTensor)
+        assert result is not grouped_tensor
+        assert result._base is grouped_tensor
+        assert result.grad_fn is not None
+        assert result.rowwise_data is grouped_tensor.rowwise_data
+        assert result.first_dims is grouped_tensor.first_dims
+        assert result.tensor_offsets is grouped_tensor.tensor_offsets
+
+    def test_record_stream_tracks_grouped_buffers(self, monkeypatch) -> None:
+        """record_stream should reach every allocation owned by the wrapper."""
+        first_dims = torch.tensor([1, 2], dtype=torch.int64, device="cuda")
+        tensor_offsets = torch.tensor([0, 4, 12], dtype=torch.int64, device="cuda")
+        grouped_tensor = GroupedTensor(
+            shape=(3, 4),
+            dtype=torch.float32,
+            num_tensors=2,
+            data=torch.randn(12, device="cuda"),
+            scale_inv=torch.ones(2, device="cuda"),
+            first_dims=first_dims,
+            tensor_offsets=tensor_offsets,
+        )
+        buffers = [
+            tensor
+            for tensor in grouped_tensor.get_data_tensors()
+            if isinstance(tensor, torch.Tensor)
+        ]
+        recorded = []
+        original_record_stream = torch.Tensor.record_stream
+
+        def record_stream(tensor, stream):
+            recorded.append(tensor)
+            return original_record_stream(tensor, stream)
+
+        monkeypatch.setattr(torch.Tensor, "record_stream", record_stream)
+        grouped_tensor.record_stream(torch.cuda.Stream())
+
+        assert all(any(record is buffer for record in recorded) for buffer in buffers)
+
+    def test_untyped_storage_resizes_payload_storages(self) -> None:
+        """resize_(0) should resize data/scales but preserve metadata storages."""
+        first_dims = torch.tensor([1, 2], dtype=torch.int64, device="cuda")
+        last_dims = torch.tensor([4, 4], dtype=torch.int64, device="cuda")
+        tensor_offsets = torch.tensor([0, 4, 12], dtype=torch.int64, device="cuda")
+        amax = torch.ones(2, device="cuda")
+        grouped_tensor = GroupedTensor(
+            shape=(3, 4),
+            dtype=torch.float32,
+            num_tensors=2,
+            data=torch.randn(12, device="cuda"),
+            columnwise_data=torch.randn(12, device="cuda"),
+            scale_inv=torch.ones(2, device="cuda"),
+            columnwise_scale_inv=torch.ones(2, device="cuda"),
+            amax=amax,
+            first_dims=first_dims,
+            last_dims=last_dims,
+            tensor_offsets=tensor_offsets,
+            requires_grad=True,
+        )
+        detached = grouped_tensor.detach()
+        viewed = grouped_tensor.view_as(grouped_tensor)
+        aliased = torch.ops.aten.alias.default(grouped_tensor)
+        expanded = grouped_tensor.expand(grouped_tensor.shape)
+        expanded_as = grouped_tensor.expand_as(grouped_tensor)
+        data_alias = grouped_tensor.data
+        saved_data = grouped_tensor.rowwise_data
+        saved_columnwise_data = grouped_tensor.columnwise_data
+        saved_scale = grouped_tensor.scale_inv
+        saved_columnwise_scale = grouped_tensor.columnwise_scale_inv
+
+        storage = grouped_tensor.untyped_storage()
+        aliases = (grouped_tensor, detached, viewed, aliased, expanded, expanded_as, data_alias)
+        assert all(storage is alias.untyped_storage() for alias in aliases)
+        storage.resize_(0)
+
+        for alias in aliases:
+            assert alias.rowwise_data.untyped_storage().nbytes() == 0
+            assert alias.columnwise_data.untyped_storage().nbytes() == 0
+            assert alias.scale_inv.untyped_storage().nbytes() == 0
+            assert alias.columnwise_scale_inv.untyped_storage().nbytes() == 0
+        assert grouped_tensor.logical_shape == (3, 4)
+        assert grouped_tensor.num_tensors == 2
+        assert grouped_tensor.first_dims is first_dims
+        assert grouped_tensor.last_dims is last_dims
+        assert grouped_tensor.tensor_offsets is tensor_offsets
+        assert first_dims.untyped_storage().nbytes() > 0
+        assert last_dims.untyped_storage().nbytes() > 0
+        assert tensor_offsets.untyped_storage().nbytes() > 0
+        assert amax.untyped_storage().nbytes() > 0
+        assert saved_data.untyped_storage().nbytes() == 0
+        assert saved_columnwise_data.untyped_storage().nbytes() == 0
+        assert saved_scale.untyped_storage().nbytes() == 0
+        assert saved_columnwise_scale.untyped_storage().nbytes() == 0
+        with pytest.raises(RuntimeError, match=r"only supports resize_\(0\)"):
+            storage.resize_(1)
+
     def test_basic_construction_all_same_shape(self) -> None:
         """Test GroupedTensor construction with all tensors having same shape"""
         num_tensors = 4

transformer_engine/pytorch/tensor/grouped_tensor.py

@@ -5,6 +5,7 @@
 """Grouped tensor class for handling collections of tensors with different shapes"""
 from __future__ import annotations
 
+import weakref
 from typing import List, Optional, Tuple
 
 import torch
@@ -66,6 +67,48 @@ def backward(ctx, grad_output: torch.Tensor):
 }
 
 
+class GroupedTensorUntypedStorage(torch.UntypedStorage):
+    """Storage facade that resizes grouped data and scale storages."""
+
+    def __new__(cls, owner: "GroupedTensor"):
+        instance = super().__new__(cls, 0, device=owner.device)
+        instance._aliases = []
+        instance.register_alias(owner)
+        return instance
+
+    def register_alias(self, tensor: "GroupedTensor") -> None:
+        """Register a wrapper that shares this storage's payload allocations."""
+        self._aliases.append(weakref.ref(tensor))
+
+    def resize_(self, size: int):
+        """Resize grouped payload storages when pseudo-deallocated by a scheduler."""
+        if size != 0:
+            raise RuntimeError(
+                f"GroupedTensor untyped storage only supports resize_(0), got resize_({size})"
+            )
+        resized_storages = set()
+        live_aliases = []
+        for alias_ref in self._aliases:
+            alias = alias_ref()
+            if alias is not None:
+                live_aliases.append(alias_ref)
+                for tensor in (
+                    alias.rowwise_data,
+                    alias.columnwise_data,
+                    alias.scale_inv,
+                    alias.columnwise_scale_inv,
+                ):
+                    if not isinstance(tensor, torch.Tensor):
+                        continue
+                    storage = tensor.untyped_storage()
+                    storage_id = storage._cdata
+                    if storage_id not in resized_storages:
+                        storage.resize_(0)
+                        resized_storages.add(storage_id)
+        self._aliases = live_aliases
+        return self
+
+
 class GroupedTensor(GroupedTensorStorage, torch.Tensor):
     """Tensor wrapper class for grouped tensor storage."""
 
@@ -174,6 +217,7 @@ def __new__(
             nvfp4_e4m3_max=nvfp4_e4m3_max,
             scale_inv_dtype=scale_inv_dtype,
         )
+        instance._untyped_storage = GroupedTensorUntypedStorage(instance)
         return instance
 
     @classmethod
@@ -208,8 +252,10 @@ def copy_grouped_storage_metadata(dst: GroupedTensor, src: GroupedTensor) -> Non
             dst.nvfp4_use_4over6 = src.nvfp4_use_4over6
             dst.nvfp4_e4m3_max = src.nvfp4_e4m3_max
             dst.scale_inv_dtype = src._scale_inv_dtype
+            dst._untyped_storage = src._untyped_storage
+            dst._untyped_storage.register_alias(dst)
 
-        def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor:
+        def make_wrapper_like(src: GroupedTensor) -> GroupedTensor:
             """Create a wrapper of the same type and tensor metadata as src."""
             out = torch.Tensor._make_wrapper_subclass(
                 type(src),
@@ -218,7 +264,7 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor:
                 storage_offset=src.storage_offset(),
                 dtype=src.dtype,
                 layout=src.layout,
-                requires_grad=requires_grad,
+                requires_grad=False,
                 device=src.device,
             )
             copy_grouped_storage_metadata(out, src)
@@ -229,9 +275,24 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor:
             src = args[0]
             if not isinstance(src, GroupedTensor):
                 raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}")
-            if func == torch.ops.aten.detach.default:
-                return make_wrapper_like(src, requires_grad=False)
-            return make_wrapper_like(src, requires_grad=src.requires_grad)
+            return make_wrapper_like(src)
+
+        # The wrapper has no allocation of its own. Propagate CUDA allocator
+        # stream tracking to every tensor that backs the grouped representation.
+        if func == torch.ops.aten.record_stream.default:
+            src, stream = args
+            if not isinstance(src, GroupedTensor):
+                raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}")
+            recorded = set()
+            for tensor in src.get_data_tensors():
+                if (
+                    isinstance(tensor, torch.Tensor)
+                    and tensor.is_cuda
+                    and id(tensor) not in recorded
+                ):
+                    tensor.record_stream(stream)
+                    recorded.add(id(tensor))
+            return None
 
         # Parameter construction may invoke aten.expand on tensor subclasses.
         # Handle this explicitly so grouped parameters can be created safely.
@@ -246,7 +307,7 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor:
                     src_shape[i] if dim == -1 else dim for i, dim in enumerate(expanded_shape)
                 )
                 if normalized_shape == src_shape:
-                    return make_wrapper_like(src, requires_grad=src.requires_grad)
+                    return make_wrapper_like(src)
             return super().__torch_dispatch__(func, types, args, kwargs)
 
         # DDP and mcore use expand_as(self) to build a dummy autograd node and
@@ -259,7 +320,7 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor:
             if other is src:
                 return _GroupedIdentityFunc.apply(src)
             if tuple(other.shape) == tuple(src.shape):
-                return make_wrapper_like(src, requires_grad=src.requires_grad)
+                return make_wrapper_like(src)
             return super().__torch_dispatch__(func, types, args, kwargs)
 
         # Distributed optimizer flattens detached parameters via
@@ -270,6 +331,14 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor:
             if not isinstance(src, GroupedTensor):
                 raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}")
             target_shape = tuple(args[1])
+            # Autograd Functions that return an input unchanged create a
+            # shape-preserving view to track the output alias. Preserve the
+            # grouped layout and share its backing buffers in that case.
+            if target_shape == tuple(src.shape):
+                # Autograd attaches the differentiable-view metadata after
+                # __torch_dispatch__ returns, so the fresh wrapper must not
+                # already require gradients.
+                return make_wrapper_like(src)
             if target_shape in ((-1,), (src.numel(),)):
                 if src.rowwise_data is not None:
                     return src.rowwise_data.view(-1)
@@ -365,6 +434,10 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
         # Do not force GroupedTensor on outputs.
         return torch._C._disabled_torch_function_impl(func, types, args, kwargs)
 
+    def untyped_storage(self) -> torch.UntypedStorage:
+        """Return the shared storage facade used for scheduler pseudo-deallocation."""
+        return self._untyped_storage
+
     def expand_as(self, other: torch.Tensor) -> torch.Tensor:
         # pylint: disable=missing-function-docstring
         # Needed during parameter creation/hook registration paths.