From f057413fc57754083d435dc5f90b1255649160a1 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 23 Oct 2026 09:28:08 -0700 Subject: [PATCH] Support GroupedTensor autograd views and stream tracking Preserve autograd view metadata for GroupedTensor wrappers and record streams on their backing tensors. Add focused coverage for gradient propagation through an identity view and for stream tracking. Signed-off-by: Varun Thumbe --- tests/pytorch/test_grouped_tensor.py | 70 +++++++++++++++++++ .../pytorch/tensor/grouped_tensor.py | 37 ++++++---- 2 files changed, 94 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 3272ea7d01d..a0eb48c8653 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -303,6 +303,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 diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 65b4a848228..f50061f3b5a 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/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