[Bug] Using `QuantizedTensor` on CPU causes a `RecursionError`

[janekb04]

Describe the bug

Trying to use an QuantizedTensor with certain Pytorch ops, after having moved it to cpu causes a RecursionError.

Steps/Code to reproduce bug

import torch
from transformer_engine.pytorch import NVFP4Quantizer

x = torch.randn(128, 128, device="cuda")
q = NVFP4Quantizer()
xq = q(x)
xq_cpu = xq.cpu()

xq_cpu + 0

will cause:

Traceback (most recent call last):
  File "/workspace/repro.py", line 9, in <module>
    xq_cpu + 0
    ~~~~~~~^~~
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1578, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1313, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 772, in maybe_unwrap
    return arg.dequantize()
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 444, in dequantize
    return _FromNVFP4Func.apply(self, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 625, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py", line 58, in forward
    cuda_tensor = tensor.to(device=torch.device("cuda"))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1578, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1313, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 772, in maybe_unwrap
    return arg.dequantize()
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 445, in dequantize
    return _FromNVFP4Func.forward(None, self, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py", line 58, in forward
    cuda_tensor = tensor.to(device=torch.device("cuda"))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1578, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1313, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 772, in maybe_unwrap
    return arg.dequantize()
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 445, in dequantize
    return _FromNVFP4Func.forward(None, self, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py", line 58, in forward
    cuda_tensor = tensor.to(device=torch.device("cuda"))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1578, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1313, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 772, in maybe_unwrap
    return arg.dequantize()
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 445, in dequantize
    return _FromNVFP4Func.forward(None, self, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py", line 58, in forward
    cuda_tensor = tensor.to(device=torch.device("cuda"))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1576, in tree_map
    leaves, treespec = tree_flatten(tree, is_leaf=is_leaf)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1485, in tree_flatten
    treespec = helper(tree, leaves)
               ^^^^^^^^^^^^^^^^^^^^
RecursionError: maximum recursion depth exceeded

I noticed this when trying to use torch.testing.assert_close on a tensor quantized with NVFP4Quantizer and a custom implementation of NVFP4 quantization.

Expected behavior

xq_cpu + 0 should not cause RecursionError.

Environment overview (please complete the following information)

• Environment location: [Docker]
• Method of Transformer Engine install: [build from source].

Environment details

Should be irrelevant.

Device details

Should be irrelevant.

Additional context

n/a.

[pggPL]

Maybe we should explicitly not support CPU tensors. Why do we want to have this support?

[janekb04]

@pggPL Based on a prior conversation with @timmoon10, I am under the impression that we need CPU tensor support for offloading.

[pggPL]

I removed them from offloading logic some time ago, but claude says that they are used by some checkpointing logic.

[vthumbe1503]

Yes, it is needed for supporting torch DCP checkpointing of quantized weights. We dont want to save dequantized tensor and load it since that would result in bitwise parity mismatch before and after a checkpoint load.

CPU quantized tensors as a concept was introduced in this PR and so yes we should fix this bug