From 25124ee63434799903d3bd9b89732e2d592409e9 Mon Sep 17 00:00:00 2001 From: jainapurva Date: Thu, 30 Apr 2026 23:01:27 +0000 Subject: [PATCH 1/3] fix(utils): propagate dtype in TorchAOBaseTensor._to_copy ## Problem PR #4297 added `non_blocking` propagation to TorchAOBaseTensor._to_copy, but introduced a bug: while `_get_to_kwargs` returns `device`, `dtype`, and `non_blocking`, the `_to_copy` handler only propagated `device` and `non_blocking` to inner tensors. This meant that calls like `tensor.to(dtype=torch.float16)` or `tensor.to(device='cuda', dtype=torch.bfloat16)` would change the wrapper tensor's dtype but NOT the inner tensors (qdata, scale, etc.), causing a dtype mismatch between the wrapper and its data. ## Fix - Pop `dtype` from kwargs and pass it to all inner `.to()` calls - Use explicit keyword arguments for clarity: `device=device, dtype=dtype, non_blocking=non_blocking` This ensures all three parameters are consistently propagated to inner tensors when calling `.to()` on TorchAOBaseTensor subclasses. ## Testing Added `test_to_copy_propagates_dtype_and_non_blocking` to verify: - Dtype-only changes propagate correctly - Combined device + dtype + non_blocking changes work - All existing tests continue to pass Co-Authored-By: Claude Sonnet 4.5 --- test/test_utils.py | 34 ++++++++++++++++++++++++++++++++++ torchao/utils.py | 5 +++-- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/test/test_utils.py b/test/test_utils.py index cbb9aa0191..d3fc44567c 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -89,6 +89,40 @@ def __init__(self, qdata, attr="attr", device=None): kwargs = t._get_to_kwargs(device="cpu") self.assertFalse(kwargs["non_blocking"]) + def test_to_copy_propagates_dtype_and_non_blocking(self): + """Verify that .to() propagates device, dtype, and non_blocking to inner tensors.""" + + class MyTensor(TorchAOBaseTensor): + tensor_data_names = ["qdata", "scale"] + tensor_attribute_names = ["device"] + + def __new__(cls, qdata, scale, device=None): + if device is None: + device = qdata.device + kwargs = {"device": device, "dtype": qdata.dtype} + r = torch.Tensor._make_wrapper_subclass(cls, qdata.shape, **kwargs) + r.qdata = qdata + r.scale = scale + return r + + def __init__(self, qdata, scale, device=None): + pass + + # Test dtype change propagates to inner tensors + t = MyTensor(torch.randn(4, 4, dtype=torch.float32), torch.tensor(1.0, dtype=torch.float32)) + result = t.to(dtype=torch.float16) + self.assertEqual(result.dtype, torch.float16) + self.assertEqual(result.qdata.dtype, torch.float16) + self.assertEqual(result.scale.dtype, torch.float16) + + # Test combined device and dtype change with non_blocking + t = MyTensor(torch.randn(4, 4, dtype=torch.float32), torch.tensor(1.0, dtype=torch.float32)) + result = t.to(device="cpu", dtype=torch.bfloat16, non_blocking=True) + self.assertEqual(result.dtype, torch.bfloat16) + self.assertEqual(result.qdata.dtype, torch.bfloat16) + self.assertEqual(result.scale.dtype, torch.bfloat16) + self.assertEqual(result.device.type, "cpu") + def _test_default_impls_helper(self, lp_tensor, lp_tensor_for_copy): # get `all_tensor_data_names` and `all_tensor_attribute_names` all_tensor_data_names = lp_tensor.tensor_data_names.copy() diff --git a/torchao/utils.py b/torchao/utils.py index 0f2ffa527b..98c6244e98 100644 --- a/torchao/utils.py +++ b/torchao/utils.py @@ -594,9 +594,10 @@ def _(func, types, args, kwargs): ): kwargs = self._get_to_kwargs(*args[1:], **kwargs) device = kwargs.pop("device") + dtype = kwargs.pop("dtype") non_blocking = kwargs.pop("non_blocking", False) tensors = [ - getattr(self, name).to(device, non_blocking=non_blocking) + getattr(self, name).to(device=device, dtype=dtype, non_blocking=non_blocking) for name in self.tensor_data_names ] optional_tensors = [] @@ -605,7 +606,7 @@ def _(func, types, args, kwargs): maybe_tensor = getattr(self, tensor_data_name) if maybe_tensor is not None: optional_tensors.append( - maybe_tensor.to(device, non_blocking=non_blocking) + maybe_tensor.to(device=device, dtype=dtype, non_blocking=non_blocking) ) else: optional_tensors.append(None) From 00f607141b176d23a87460e24149e7064edde02a Mon Sep 17 00:00:00 2001 From: jainapurva Date: Thu, 30 Apr 2026 23:17:42 +0000 Subject: [PATCH 2/3] Apply ruff formatting Co-Authored-By: Claude Sonnet 4.5 --- test/test_utils.py | 10 ++++++++-- torchao/utils.py | 8 ++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/test/test_utils.py b/test/test_utils.py index d3fc44567c..5e98c0a5ea 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -109,14 +109,20 @@ def __init__(self, qdata, scale, device=None): pass # Test dtype change propagates to inner tensors - t = MyTensor(torch.randn(4, 4, dtype=torch.float32), torch.tensor(1.0, dtype=torch.float32)) + t = MyTensor( + torch.randn(4, 4, dtype=torch.float32), + torch.tensor(1.0, dtype=torch.float32), + ) result = t.to(dtype=torch.float16) self.assertEqual(result.dtype, torch.float16) self.assertEqual(result.qdata.dtype, torch.float16) self.assertEqual(result.scale.dtype, torch.float16) # Test combined device and dtype change with non_blocking - t = MyTensor(torch.randn(4, 4, dtype=torch.float32), torch.tensor(1.0, dtype=torch.float32)) + t = MyTensor( + torch.randn(4, 4, dtype=torch.float32), + torch.tensor(1.0, dtype=torch.float32), + ) result = t.to(device="cpu", dtype=torch.bfloat16, non_blocking=True) self.assertEqual(result.dtype, torch.bfloat16) self.assertEqual(result.qdata.dtype, torch.bfloat16) diff --git a/torchao/utils.py b/torchao/utils.py index 98c6244e98..4c541f316f 100644 --- a/torchao/utils.py +++ b/torchao/utils.py @@ -597,7 +597,9 @@ def _(func, types, args, kwargs): dtype = kwargs.pop("dtype") non_blocking = kwargs.pop("non_blocking", False) tensors = [ - getattr(self, name).to(device=device, dtype=dtype, non_blocking=non_blocking) + getattr(self, name).to( + device=device, dtype=dtype, non_blocking=non_blocking + ) for name in self.tensor_data_names ] optional_tensors = [] @@ -606,7 +608,9 @@ def _(func, types, args, kwargs): maybe_tensor = getattr(self, tensor_data_name) if maybe_tensor is not None: optional_tensors.append( - maybe_tensor.to(device=device, dtype=dtype, non_blocking=non_blocking) + maybe_tensor.to( + device=device, dtype=dtype, non_blocking=non_blocking + ) ) else: optional_tensors.append(None) From 23df18a24f9e56d2b77270eeae70795a47d5877c Mon Sep 17 00:00:00 2001 From: jainapurva Date: Thu, 30 Apr 2026 23:48:55 +0000 Subject: [PATCH 3/3] fix(mx_formats): update float8 saturation test for PyTorch behavior change MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem PyTorch nightly now implements saturated casting to float8_e4m3fn in eager mode, matching the behavior that was previously only in compiled/triton mode. The test `test_cast_to_float8_e4m3fn_saturation_behavior` was expecting the old unsaturated behavior (out-of-range values → NaN), causing H100 CI failures. ## Fix Updated the test to verify the new saturated casting behavior: - Changed assertion from expecting NaN to expecting saturation - Added verification that out-of-range values are clamped to max_val - Updated assertions to verify eager and compiled modes produce identical results - Updated comments to reflect the completed TODO from issue #1912 ## Testing This fixes the H100 test failures on main branch where the test was asserting: ```python assert torch.all(torch.isnan(data_out_of_range_f8)) # Old behavior ``` But PyTorch now produces saturated values (448/-448) instead of NaN. Co-Authored-By: Claude Sonnet 4.5 --- test/prototype/mx_formats/test_mx_tensor.py | 22 +++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/test/prototype/mx_formats/test_mx_tensor.py b/test/prototype/mx_formats/test_mx_tensor.py index 744cbeae47..d6e58f7251 100644 --- a/test/prototype/mx_formats/test_mx_tensor.py +++ b/test/prototype/mx_formats/test_mx_tensor.py @@ -670,8 +670,9 @@ def test_index_select(): reason="float8 in triton requires CUDA capability 8.9 or greater", ) def test_cast_to_float8_e4m3fn_saturation_behavior(): - # TODO(#1912): make the saturated cast work in eager mode and remove this - # test + # PyTorch now implements saturated casting to float8_e4m3fn in both eager + # and compiled modes. This test verifies the saturation behavior. + # Relates to issue #1912. max_val = torch.finfo(torch.float8_e4m3fn).max # create example data inside the representable range @@ -694,13 +695,20 @@ def test_cast_to_float8_e4m3fn_saturation_behavior(): device="cuda", ) - # verify that in eager mode PyTorch casting to float8 is unsaturated + # verify that in eager mode PyTorch casting to float8 is saturated (no NaN) data_in_range_f8 = data_in_range_bf16.to(torch.float8_e4m3fn) data_out_of_range_f8 = data_out_of_range_bf16.to(torch.float8_e4m3fn) assert not torch.any(torch.isnan(data_in_range_f8)) - assert torch.all(torch.isnan(data_out_of_range_f8)) + assert not torch.any(torch.isnan(data_out_of_range_f8)) + # Out-of-range values should be clamped to max_val + torch.testing.assert_close( + data_out_of_range_f8.to(torch.bfloat16), + torch.tensor([max_val, -1 * max_val], dtype=torch.bfloat16, device="cuda"), + atol=0, + rtol=0, + ) - # verify that in triton, casting to float8 is saturated + # verify that in triton/compiled mode, casting to float8 is also saturated # for simplicity, use torch.compile to generate triton code def to_f8(x): x = x.to(torch.float8_e4m3fn) @@ -711,8 +719,10 @@ def to_f8(x): data_out_of_range_f8_c = to_f8_c(data_out_of_range_bf16) assert not torch.any(torch.isnan(data_in_range_f8_c)) assert not torch.any(torch.isnan(data_out_of_range_f8_c)) + # Eager and compiled should produce the same saturated results + torch.testing.assert_close(data_in_range_f8, data_in_range_f8_c, atol=0, rtol=0) torch.testing.assert_close( - data_in_range_f8_c, data_out_of_range_f8_c, atol=0, rtol=0 + data_out_of_range_f8, data_out_of_range_f8_c, atol=0, rtol=0 )