[Pytorch] Fix RecursionError when dequantizing cpu-resident QuantizedTensor

tests/pytorch/test_quantized_tensor.py

@@ -670,6 +670,42 @@ def test_cpu_dequantize(
         assert y_cpu.shape == ref_cpu.shape
         torch.testing.assert_close(y_cpu, ref_cpu, rtol=0, atol=0)
 
+    @pytest.mark.parametrize("quantization", _quantization_list)
+    def test_cpu_torch_ops(
+        self,
+        *,
+        quantization: str,
+        shape: Iterable[int] = (128, 128),
+        dtype: torch.dtype = torch.bfloat16,
+    ) -> None:
+        """Apply plain PyTorch ops to a CPU-resident QuantizedTensor.
+
+        Ops without a quantized implementation dequantize their arguments. For
+        a CPU tensor that dequantization may need to copy the buffers to CUDA,
+        which must not dispatch back through the quantized tensor.
+
+        """
+
+        # Construct a quantized tensor on CUDA and move it to CPU
+        _, x_cuda = make_reference_and_test_tensors(
+            shape=shape,
+            quantization=quantization,
+            test_dtype=dtype,
+            requires_grad=False,
+        )
+        x_cpu = x_cuda.cpu()
+        ref_cpu = x_cuda.dequantize().to(device="cpu")
+
+        # Ops that are dispatched through the dequantizing fallback
+        torch.testing.assert_close(x_cpu + 0, ref_cpu, rtol=0, atol=0)
+        torch.testing.assert_close(x_cpu, ref_cpu, rtol=0, atol=0)
+
+        # Moving back to CUDA preserves the quantized tensor type
+        x_roundtrip = x_cpu.to(device="cuda")
+        assert isinstance(x_roundtrip, QuantizedTensor)
+        assert x_roundtrip.device.type == "cuda"
+        torch.testing.assert_close(x_roundtrip.dequantize(), x_cuda.dequantize(), rtol=0, atol=0)
+
     @pytest.mark.parametrize("quantization", _quantization_list)
     @pytest.mark.parametrize("dim", [0, 1])
     def test_chunk(

transformer_engine/pytorch/quantized_tensor.py

@@ -30,6 +30,29 @@
 _quantized_tensor_passthrough_ops: set = set()
 
 
+# Ops that copy a tensor to another device and/or dtype. The ``aten.to.*``
+# overloads are CompositeImplicitAutograd, so they normally decompose into
+# ``aten._to_copy`` before reaching __torch_dispatch__. That decomposition is
+# skipped when they are called from within another __torch_dispatch__, because
+# autograd keys are excluded there and the CompositeImplicitAutograd alias does
+# not cover the Python dispatch key. Both forms therefore need handling.
+_device_conversion_ops: set = {
+    torch.ops.aten._to_copy.default,
+    torch.ops.aten.to.device,
+    torch.ops.aten.to.dtype,
+    torch.ops.aten.to.dtype_layout,
+    torch.ops.aten.to.other,
+}
+
+
+def _bind_op_args(func, args, kwargs) -> Dict[str, Any]:
+    """Map an op's arguments to the argument names in its schema"""
+    bound_args = dict(zip((arg.name for arg in func._schema.arguments), args))
+    if kwargs:
+        bound_args.update(kwargs)
+    return bound_args
+
+
 class QuantizedTensorStorage:
     r"""Base class for all TensorStorage classes.
 
@@ -676,18 +699,20 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
                 dst.copy_(src)
             return None
 
-        # _to_copy op (used by .to(device=...), .cpu(), DCP staging).
+        # Device conversion ops (used by .to(device=...), .cpu(), DCP staging).
         # Preserve the QuantizedTensor subclass and move all internal
         # buffers (data, scales, etc.) to the requested device.
-        if func == torch.ops.aten._to_copy.default:
+        if func in _device_conversion_ops and isinstance(args[0], QuantizedTensor):
             tensor = args[0]
-            kw = dict(kwargs) if kwargs else {}
-            dtype = kw.get("dtype", None)
+            bound_args = _bind_op_args(func, args, kwargs)
+            # ``aten.to.other`` takes the target dtype and device from a tensor
+            other = bound_args.get("other", None)
+            dtype = other.dtype if other is not None else bound_args.get("dtype", None)
+            device = other.device if other is not None else bound_args.get("device", None)
             if dtype is None or dtype == tensor.dtype:
-                target_device = kw.get("device", tensor.device) or tensor.device
-                target_device = torch.device(target_device)
-                pin_memory = bool(kw.get("pin_memory", False))
-                non_blocking = bool(kw.get("non_blocking", False))
+                target_device = torch.device(device) if device is not None else tensor.device
+                pin_memory = bool(bound_args.get("pin_memory", False))
+                non_blocking = bool(bound_args.get("non_blocking", False))
                 new_metadata = {"device": target_device}
                 # Update tensor storage metadata
                 for key, value in tensor.get_metadata().items():

transformer_engine/pytorch/tensor/_quantization_helpers.py

@@ -13,7 +13,10 @@
 import torch
 
 if TYPE_CHECKING:
-    from transformer_engine.pytorch.quantized_tensor import QuantizedTensor
+    from transformer_engine.pytorch.quantized_tensor import (
+        QuantizedTensor,
+        QuantizedTensorStorage,
+    )
 
 
 class _QuantizeFunc(torch.autograd.Function):
@@ -85,6 +88,26 @@ def _stride_from_shape(shape: list[int]):
     return list(reversed(rstride))
 
 
+def cuda_storage_copy(
+    tensor: QuantizedTensorStorage,
+    storage_class: type,
+) -> QuantizedTensorStorage:
+    """Copy a quantized tensor's buffers to CUDA, as a plain tensor storage.
+
+    ``storage_class`` must be the ``QuantizedTensorStorage`` subclass that
+    matches ``tensor``. The copy is deliberately not a ``QuantizedTensor``:
+    constructing one dispatches through ``__torch_dispatch__``, which
+    dequantizes the arguments of ops it does not recognize and would therefore
+    recurse back into the dequantization that needs this copy.
+
+    """
+    metadata = {
+        key: value.to(device="cuda") if isinstance(value, torch.Tensor) else value
+        for key, value in tensor.get_metadata().items()
+    }
+    return storage_class(**metadata)
+
+
 def safe_quantized_repr(obj, cls_name, extras=None, error=None):
     """Metadata-only repr fallback for quantized tensors whose data cannot be
     materialized for any reason.

transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py

@@ -13,7 +13,7 @@
 import transformer_engine_torch as tex
 
 from ...quantized_tensor import QuantizedTensorStorage, Quantizer
-from .._quantization_helpers import safe_quantized_repr
+from .._quantization_helpers import cuda_storage_copy, safe_quantized_repr
 
 from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType
 
@@ -38,10 +38,11 @@ def forward(
         if tensor._rowwise_data is None and tensor._columnwise_data is None:
             raise ValueError("Cannot dequantize MXFP8 tensor with no data")
         te_dtype = torch_to_transformer_engine_dtype[dtype]
-        # ``tex.dequantize`` requires CUDA-resident buffers.
+        # ``tex.dequantize`` requires CUDA-resident buffers. If the tensor lives
+        # elsewhere, copy the buffers to CUDA and bring the result back.
         src_device = tensor.device
         if src_device.type != "cuda":
-            cuda_tensor = tensor.to(device=torch.device("cuda"))
+            cuda_tensor = cuda_storage_copy(tensor, MXFP8TensorStorage)
             result = tex.dequantize(cuda_tensor, te_dtype)
             return result.to(device=src_device)
         return tex.dequantize(tensor, te_dtype)

transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py

@@ -16,7 +16,7 @@
 import transformer_engine_torch as tex
 
 from ...quantized_tensor import QuantizedTensorStorage, Quantizer
-from .._quantization_helpers import safe_quantized_repr
+from .._quantization_helpers import cuda_storage_copy, safe_quantized_repr
 
 from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType
 from ...utils import _empty_tensor
@@ -52,10 +52,11 @@ def forward(
         if tensor._rowwise_data is None and tensor._columnwise_data is not None:
             raise NotImplementedError("Dequantizing column-wise NVFP4 data is not implemented yet!")
 
-        # ``tex.dequantize`` requires CUDA-resident buffers. If the tensor has
+        # ``tex.dequantize`` requires CUDA-resident buffers. If the tensor lives
+        # elsewhere, copy the buffers to CUDA and bring the result back.
         src_device = tensor.device
         if src_device.type != "cuda":
-            cuda_tensor = tensor.to(device=torch.device("cuda"))
+            cuda_tensor = cuda_storage_copy(tensor, NVFP4TensorStorage)
             result = tex.dequantize(cuda_tensor, torch_to_transformer_engine_dtype[dtype])
             return result.to(device=src_device)
         return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype])