Support GroupedTensor autograd views and stream tracking

tests/pytorch/test_grouped_tensor.py

@@ -299,6 +299,76 @@ def setup_class(cls) -> None:
         torch.manual_seed(seed)
         torch.cuda.manual_seed(seed)
 
+    def test_autograd_identity_view_preserves_grouped_storage_and_gradient(self) -> None:
+        """Autograd's identity view keeps grouped metadata and passes back a dense gradient."""
+
+        class Identity(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, tensor):
+                return tensor
+
+            @staticmethod
+            def backward(ctx, grad):
+                return grad
+
+        first_dims = torch.tensor([1, 2], dtype=torch.int64, device="cuda")
+        tensor_offsets = torch.tensor([0, 4, 12], dtype=torch.int64, device="cuda")
+        source = 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,
+        )
+
+        view = Identity.apply(source)
+        assert isinstance(view, GroupedTensor)
+        assert view is not source
+        assert view._base is source
+        assert view.rowwise_data is source.rowwise_data
+        assert view.first_dims is first_dims
+        assert view.tensor_offsets is tensor_offsets
+
+        expected = torch.arange(1, 13, dtype=torch.float32, device="cuda").reshape(3, 4)
+        (grad,) = torch.autograd.grad(view, source, grad_outputs=expected)
+        assert type(grad) is torch.Tensor
+        torch.testing.assert_close(grad, expected)
+
+    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")
+        last_dims = torch.tensor([4, 4], 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"),
+            columnwise_data=torch.randn(12, device="cuda"),
+            scale_inv=torch.ones(2, device="cuda"),
+            columnwise_scale_inv=torch.ones(2, device="cuda"),
+            amax=torch.zeros(2, device="cuda"),
+            columnwise_amax=torch.zeros(2, device="cuda"),
+            scale=torch.ones(2, device="cuda"),
+            first_dims=first_dims,
+            last_dims=last_dims,
+            tensor_offsets=tensor_offsets,
+        )
+        buffers = grouped_tensor.get_data_tensors()
+        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_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

@@ -209,7 +209,7 @@ def copy_grouped_storage_metadata(dst: GroupedTensor, src: GroupedTensor) -> Non
             dst.nvfp4_e4m3_max = src.nvfp4_e4m3_max
             dst.scale_inv_dtype = src._scale_inv_dtype
 
-        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 +218,9 @@ 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,
+                # Autograd attaches the output's differentiable-view metadata after
+                # __torch_dispatch__ returns.
+                requires_grad=False,
                 device=src.device,
             )
             copy_grouped_storage_metadata(out, src)
@@ -229,9 +231,18 @@ 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__}")
+            for tensor in src.get_data_tensors():
+                if tensor is not None and tensor.is_cuda:
+                    tensor.record_stream(stream)
+            return None
 
         # Parameter construction may invoke aten.expand on tensor subclasses.
         # Handle this explicitly so grouped parameters can be created safely.
@@ -246,7 +257,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 +270,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 +281,11 @@ 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):
+                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)
@@ -358,12 +374,7 @@ def maybe_update_inplace(arg, new_arg, schema_arg):
         new_kwargs = tree_map(maybe_unwrap, kwargs)
         return super().__torch_dispatch__(func, types, new_args, new_kwargs)
 
-    @classmethod
-    def __torch_function__(cls, func, types, args=(), kwargs=None):
-        if kwargs is None:
-            kwargs = {}
-        # Do not force GroupedTensor on outputs.
-        return torch._C._disabled_torch_function_impl(func, types, args, kwargs)
+    __torch_function__ = torch._C._disabled_torch_function_impl
 
     def expand_as(self, other: torch.Tensor) -> torch.Tensor:
         # pylint: disable=missing-function-docstring