diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index b2f77ecd66..fd8e220205 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -757,6 +757,58 @@ def test_shape_with_none_data( f"after setting data to None on {type(x_test).__name__}" ) + @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize( + "rowwise, columnwise", + [(True, True), (True, False), (False, True)], + ids=["rowwise_columnwise", "rowwise_only", "columnwise_only"], + ) + @pytest.mark.parametrize("shape", [(128, 256), (4, 128, 256)], ids=["2d", "3d"]) + def test_shape_matches_size( + self, + *, + quantization: str, + rowwise: bool, + columnwise: bool, + shape: Iterable[int], + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + ) -> None: + """shape, size() and size(dim) stay consistent for every usage combination. + + Both shape and size() are derived from whichever data buffer is present, + and classes that store columnwise data transposed have to undo that. A + columnwise-only tensor is where they can drift apart -- from each other, + and from the shape the tensor was allocated with. + """ + quantizer = make_quantizer(quantization, device=device) + # Row-scaled NVFP4 accepts set_usage(rowwise=False) but rejects the + # allocation itself, so it has to be filtered out up front. + if getattr(quantizer, "row_scaled_nvfp4", False) and not rowwise: + pytest.skip(f"{quantization} requires rowwise usage") + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + if (quantizer.rowwise_usage, quantizer.columnwise_usage) != (rowwise, columnwise): + pytest.skip(f"{quantization} does not support this usage combination") + + x = quantizer.make_empty(shape, dtype=dtype, device=device) + name = type(x).__name__ + + # shape and size() must describe the same tensor, whichever buffer they + # end up reading. + assert tuple(x.shape) == tuple(x.size()), f"{name}: {tuple(x.shape)} vs {tuple(x.size())}" + + # size(dim) must agree with the full shape, including negative indices. + # It cannot be served by forwarding dim to a transposed buffer. + for dim in range(len(x.shape)): + assert x.size(dim) == x.shape[dim], f"{name}.size({dim}) is {x.size(dim)}" + neg = dim - len(x.shape) + assert x.size(neg) == x.shape[neg], f"{name}.size({neg}) is {x.size(neg)}" + + # NVFP4 deliberately reports columnwise-only tensors flattened to 2D and + # warns about it, so only the ranks it preserves are checked here. + if not (isinstance(quantizer, NVFP4Quantizer) and not rowwise and len(shape) > 2): + assert tuple(x.shape) == tuple(shape), f"{name}.shape is {tuple(x.shape)}" + @pytest.mark.parametrize( "quantization", _quantization_list + (["nvfp4_2d"] if nvfp4_available else []), diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 09fde86f17..6699cb1ea1 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -551,7 +551,11 @@ def shape(self): if self._rowwise_data is not None: return self._rowwise_data.shape if self._columnwise_data is not None: - return self._columnwise_data.shape + # Columnwise data is stored transposed, matching size() in the storage. + dims = self._columnwise_data.shape + if len(dims) == 2: + return torch.Size((dims[1], dims[0])) + return torch.Size(tuple(dims[1:]) + (dims[0],)) return torch.Tensor.size(self) @property diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 464a7c3b23..06516d0b9e 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -317,12 +317,15 @@ def size(self, *args, **kwargs): # pylint: disable=missing-function-docstring if self._rowwise_data is not None: return self._rowwise_data.size(*args, **kwargs) - dims = list(self._columnwise_data.size(*args, **kwargs)) - reordered = [] - for i in range(1, len(dims)): - reordered.append(dims[i]) - reordered.append(dims[0]) - return torch.Size(reordered) + # Columnwise data is stored transposed, so a dim argument cannot be + # forwarded to it: rebuild the logical shape first, then index into it. + dims = self._columnwise_data.shape + if len(dims) == 2: + shape = torch.Size((dims[1], dims[0])) + else: + shape = torch.Size(tuple(dims[1:]) + (dims[0],)) + dim = args[0] if args else kwargs.get("dim") + return shape if dim is None else shape[dim] def view(self, shape): """Reshape the leading (token) dims without dequantizing. diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 374d0e1e72..429ddde97f 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -5,7 +5,6 @@ """Mixin class holding data specific for Float8Tensor""" from __future__ import annotations -import math from typing import Any, Dict, Optional, Tuple, Union import torch @@ -179,8 +178,16 @@ def size(self, *args, **kwargs): # pylint: disable=missing-function-docstring if self._data is not None: return self._data.size(*args, **kwargs) - size = self._transpose.size(*args, **kwargs) - return torch.Size([size[-1], math.prod(size[:-1])]) + # The transpose is stored as [last, *leading], so a dim argument cannot + # be forwarded to it: rebuild the logical shape first, then index into + # it. This matches the shape property on Float8Tensor. + dims = self._transpose.shape + if len(dims) == 2: + shape = torch.Size((dims[1], dims[0])) + else: + shape = torch.Size(tuple(dims[1:]) + (dims[0],)) + dim = args[0] if args else kwargs.get("dim") + return shape if dim is None else shape[dim] @property def device(self):