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 ) diff --git a/test/test_utils.py b/test/test_utils.py index cbb9aa0191..5e98c0a5ea 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -89,6 +89,46 @@ 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..4c541f316f 100644 --- a/torchao/utils.py +++ b/torchao/utils.py @@ -594,9 +594,12 @@ 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 +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, non_blocking=non_blocking) + maybe_tensor.to( + device=device, dtype=dtype, non_blocking=non_blocking + ) ) else: optional_tensors.append(None)