Fix shape of new quantized tensor in make_like

[ksivaman]

Description

In functions such as split, chunk etc. in which the shape of the input and output differs, the returned tensor is correct but with the incorrect shape which leads to bugs, e.g. in FSDP2 or checkpoint loading. A small repro:

import torch
from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer
import transformer_engine_torch as tex

t = torch.randn(4, 4)
quantizer = Float8Quantizer(
    scale=torch.full([1], 1.0, dtype=torch.float32, device="cuda"),
    amax=torch.empty([1], dtype=torch.float32, device="cuda"),
    fp8_dtype=tex.DType.kFloat8E4M3,
)
x = quantizer(t.cuda())
a, b = x.chunk(2, dim=0)
print(x.shape, a.shape, b.shape)
print(x, a, b)

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

• If data

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

[ksivaman]

/te-ci pytorch

[timmoon10]

This is fine as a hack.

Rambling digression: I don't like the data kwarg. It's not generic and going forward there's no reason to expect this logic to be valid in the future (FP6 will probably require a randomly sized blob of bytes). Also, if you're providing data you should also probably provide scale_inv (and maybe column-wise data as well). I think it's fine if it's an optional kwarg in Float8Tensor and MXFP8Tensor, but we should stop exposing it in QuantizedTensor.

[ptrendx]

I agree with Tim, why not just add the shape argument to the usage of the make_like in float8tensor:

diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index 49bf4facf..0063b286a 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -402,7 +402,7 @@ class Float8Tensor(Float8TensorBase, QuantizedTensor):
                 [data] + list(args[1:]),
                 kwargs,
             )
-            return [Float8Tensor.make_like(tensor, data=split_tensor) for split_tensor in func_out]
+            return [Float8Tensor.make_like(tensor, data=split_tensor, shape=split_tensor.shape) for split_tensor in func_out]
         if func == aten.new_zeros.default:
             tensor = args[0]
             data = tensor._data
@@ -412,7 +412,7 @@ class Float8Tensor(Float8TensorBase, QuantizedTensor):
                 [data] + list(args[1:]),
                 kwargs,
             )
-            return Float8Tensor.make_like(tensor, data=func_out)
+            return Float8Tensor.make_like(tensor, data=func_out, shape=data.func_out.shape)
         if func == torch.ops.aten.as_strided.default:
             tensor = args[0]
             data = tensor._data
@@ -422,7 +422,7 @@ class Float8Tensor(Float8TensorBase, QuantizedTensor):
                 [data] + list(args[1:]),
                 kwargs,
             )
-            return Float8Tensor.make_like(tensor, data=func_out)
+            return Float8Tensor.make_like(tensor, data=func_out, shape=data.func_out.shape)
         if func == torch.ops.aten.detach.default:
             return cls.detach(args[0])
         if func == torch.ops.aten.clone.default:

I confirmed that it also resolves the given repro.
Also, could you add that repro to the sanity tests?

[ksivaman]

/te-ci pytorch

[ksivaman]

Added the shapes to make_like calls in float8_tensor, but I'm also keeping the original change to quantized_tensor since that is logically the correct shape to use in case the tensor's and data's shape do differ. We can remove the data kwarg completely from quantized_tensor which will require implementation of make_like in the Float8Tensor and MXFP8Tensor, so leaving that for a different PR.