Repository navigation
[Bug] Using QuantizedTensor on CPU causes a RecursionError #3294
Description
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.
Activity
pggPL commentedon Jul 31, 2026
Maybe we should explicitly not support CPU tensors. Why do we want to have this support?
janekb04 commentedon Jul 31, 2026
@pggPL Based on a prior conversation with @timmoon10, I am under the impression that we need CPU tensor support for offloading.
pggPL commentedon Aug 3, 2026
I removed them from offloading logic some time ago, but claude says that they are used by some checkpointing logic.
vthumbe1503 commentedon Aug 10, 2026
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