CPU offloading fix: If Data and Transpose is None depend on super Torch tensor class for the shape (#2841)

* fix

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: vthumbe1503

tests/pytorch/test_quantized_tensor.py

@@ -18,6 +18,7 @@
     MXFP8Quantizer,
     NVFP4Quantizer,
     Float8Tensor,
+    Float8BlockwiseQTensor,
     MXFP8Tensor,
     NVFP4Tensor,
     QuantizedTensor,
@@ -657,6 +658,50 @@ def test_chunk(
             y_test = y_test.to(dtype=torch.float64, device="cpu")
             torch.testing.assert_close(y_test, y_ref, **tols)
 
+    @pytest.mark.parametrize("quantization", _quantization_list)
+    def test_shape_with_none_data(
+        self,
+        *,
+        quantization: str,
+        shape: Iterable[int] = (128, 128),
+        dtype: torch.dtype = torch.bfloat16,
+    ) -> None:
+        """Test that shape is accessible after internal data tensors are set to None.
+
+        During CPU offloading, both data and transpose tensors can be None.
+        The shape should still be available via the wrapper subclass metadata.
+        """
+
+        _, x_test = make_reference_and_test_tensors(
+            shape=shape,
+            quantization=quantization,
+            test_dtype=dtype,
+            requires_grad=False,
+        )
+
+        # Verify shape before clearing data
+        assert x_test.shape == torch.Size(shape)
+
+        # Simulate CPU offloading: None out all internal data
+        if isinstance(x_test, Float8Tensor):
+            x_test._data = None
+            x_test._transpose = None
+        elif isinstance(x_test, MXFP8Tensor):
+            x_test._rowwise_data = None
+            x_test._columnwise_data = None
+        elif isinstance(x_test, NVFP4Tensor):
+            x_test._rowwise_data = None
+            x_test._columnwise_data = None
+        elif isinstance(x_test, Float8BlockwiseQTensor):
+            x_test._rowwise_data = None
+            x_test._columnwise_data = None
+
+        # Shape must still be correct after data is cleared
+        assert x_test.shape == torch.Size(shape), (
+            f"Expected shape {shape} but got {x_test.shape} "
+            f"after setting data to None on {type(x_test).__name__}"
+        )
+
 
 @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
 class TestMXFP8Tensor:

transformer_engine/pytorch/tensor/float8_blockwise_tensor.py

@@ -598,7 +598,7 @@ def shape(self):
             return self._rowwise_data.shape
         if self._columnwise_data is not None:
             return self._columnwise_data.shape
-        raise RuntimeError("Float8BlockwiseQTensor has no data!")
+        return torch.Tensor.size(self)
 
     @property
     def is_cuda(self):

transformer_engine/pytorch/tensor/float8_tensor.py

@@ -967,7 +967,7 @@ def shape(self):
         if self._transpose is not None:
             transpose_shape = self._transpose.shape
             return torch.Size(tuple(transpose_shape[1:]) + (transpose_shape[0],))
-        raise RuntimeError("Both data and transpose are None")
+        return torch.Tensor.size(self)
 
     @property
     def is_cuda(self):

transformer_engine/pytorch/tensor/mxfp8_tensor.py

@@ -884,7 +884,7 @@ def shape(self):
             return self._rowwise_data.shape
         if self._columnwise_data is not None:
             return self._columnwise_data.shape
-        raise RuntimeError("MXFP8Tensor has no data!")
+        return torch.Tensor.size(self)
 
     @property
     def is_cuda(self):

transformer_engine/pytorch/tensor/nvfp4_tensor.py

@@ -745,7 +745,7 @@ def shape(self):
         if self._columnwise_data is not None:
             byte_shape = self._columnwise_data.shape
             return torch.Size(byte_shape[1:-1] + (byte_shape[-1] * 2, byte_shape[0]))
-        raise RuntimeError("NVFP4Tensor has no data!")
+        return torch.Tensor.size(self)
 
     @property
     def is_cuda(self):