diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 093f4f7cf3..cb4a2e7654 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/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 diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 65b4a84822..f50061f3b5 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