diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index a42b6bfef1..864bcaede0 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/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( diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 654f84159f..2761743529 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/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(): diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 1b08039dda..b8c9ff43ba 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/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. diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 606ac9e74b..a1bab98476 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/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) diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 09f040ba67..169fb19347 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/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])