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

[vthumbe1503]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR fixes a crash during CPU offloading by replacing raise RuntimeError(...) with return torch.Tensor.size(self) in the shape property of four quantized tensor classes. When both internal data tensors are None — a valid state during CPU offloading — the fix correctly falls back to the shape stored in the PyTorch wrapper subclass metadata (set via _make_wrapper_subclass at construction time in QuantizedTensor.__new__), which always retains the original logical shape regardless of whether data storage is present.

Confidence Score: 5/5

Safe to merge — the fix is mechanically correct and all remaining findings are P2 style suggestions.

The fallback to torch.Tensor.size(self) is the standard idiom for wrapper subclasses; shape metadata is always present in the C++ tensor created by _make_wrapper_subclass regardless of whether data storage is populated. All four tensor types are fixed consistently, every type in _quantization_list has a matching isinstance branch in the new test, and no blocking issues remain.

No files require special attention; the P2 note about is_cuda/is_cpu in float8_tensor.py is a follow-up hardening suggestion, not a blocker.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/tensor/float8_tensor.py	Replaces RuntimeError with torch.Tensor.size(self) fallback in shape property when both _data and _transpose are None
transformer_engine/pytorch/tensor/float8_blockwise_tensor.py	Replaces RuntimeError with torch.Tensor.size(self) fallback in shape property when both _rowwise_data and _columnwise_data are None
transformer_engine/pytorch/tensor/mxfp8_tensor.py	Replaces RuntimeError with torch.Tensor.size(self) fallback in shape property when both rowwise and columnwise data are None
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Replaces RuntimeError with torch.Tensor.size(self) fallback in shape property when both rowwise and columnwise data are None
tests/pytorch/test_quantized_tensor.py	Adds test_shape_with_none_data covering all 4 quantization types; verifies shape is correct after all internal data tensors are set to None

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[shape property called] --> B{_rowwise_data / _data\nis not None?}
    B -- Yes --> C[Return data.shape\nor computed logical shape]
    B -- No --> D{_transpose /\n_columnwise_data\nis not None?}
    D -- Yes --> E[Return shape derived\nfrom transpose/columnwise data]
    D -- No --> F[torch.Tensor.size self\nfallback to wrapper\nsubclass metadata]
    F --> G[Returns shape set during\n_make_wrapper_subclass\nat construction time]
    style F fill:#90EE90
    style G fill:#90EE90

Loading

Comments Outside Diff (1)

1. transformer_engine/pytorch/tensor/float8_tensor.py, line 972-988 (link)

   is_cuda/is_cpu still raise when data is None

   The shape property now gracefully handles the all-None case, but is_cuda and is_cpu (and the device equivalents in MXFP8Tensor/NVFP4Tensor) still raise RuntimeError("Both data and transpose are None"). If CPU offloading code inspects these properties after clearing tensors — e.g., to decide where to restore data — the crash will resurface. Consider applying the same torch.Tensor base-class fallback pattern here for full robustness.

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[cspades]

Matches the patch I made.

    @property
    def shape(self):
        """Return the shape of the tensor. Define this to avoid expensive PyObject lookups."""
        if self._data is not None:
            return self._data.shape
        if self._transpose is not None:
            transpose_shape = self._transpose.shape
            return torch.Size(tuple(transpose_shape[1:]) + (transpose_shape[0],))
        return torch.Tensor.size(self)

I just did a quick 100-step Llama 8B test, I have loss parity as well:

[2026-04-06 16:36:16.442720] iteration      100/15258789 | consumed samples:        12800 | elapsed time per iteration (ms): 13432.5 | throughput per GPU (TFLOP/s/GPU): 1004.5 | learning rate: 4.915198E-07 | global batch size:   128 | lm loss: 1.261873E+00 | loss scale: 1.0 | grad norm: 5.902 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |

[ptrendx]

/te-ci pytorch

[timmoon10]

Overall LGTM

[timmoon10]

Nit: This is correct, but it reads unpythonic to me. The following would be more standard:

Suggested change

	return torch.Tensor.size(self)
	return super(QuantizedTensor, self).size()

[timmoon10]

We don't want spurious test passes when we add new tensor types.

Suggested change

	else:
	raise NotImplementedError(f"{type(x_test).__name__} is not supported")