diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index f1d64696ac..b2f77ecd66 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -110,6 +110,42 @@ def to_float8_CS( return quantizer(tensor) +def make_quantizer(quantization: str, device: torch.device = "cuda"): + """Construct a quantizer for the given quantization scheme.""" + if quantization in ("fp8", "fp8_delayed_scaling"): + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device=device).squeeze(), + amax=torch.zeros(1, dtype=torch.float32, device=device), + fp8_dtype=te.DType.kFloat8E4M3, + ) + if quantization == "fp8_current_scaling": + return Float8CurrentScalingQuantizer(fp8_dtype=te.DType.kFloat8E4M3, device=device) + if quantization == "fp8_blockwise": + return Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + if quantization == "mxfp8": + return MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + if quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + row_scaled_nvfp4 = quantization == "nvfp4_row_scaled" + return NVFP4Quantizer( + columnwise=not row_scaled_nvfp4, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=row_scaled_nvfp4, + with_random_sign_mask=False, + nvfp4_use_4over6=(quantization == "nvfp4_4over6"), + ) + raise ValueError(f"Unsupported quantization scheme ({quantization})") + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], @@ -139,45 +175,8 @@ def make_reference_and_test_tensors( if quantization is None: if test.data_ptr() == ref.data_ptr(): test = test.clone() - elif quantization in ("fp8", "fp8_delayed_scaling"): - quantizer = Float8Quantizer( - scale=torch.ones(1, dtype=torch.float32, device=test_device).squeeze(), - amax=torch.zeros(1, dtype=torch.float32, device=test_device), - fp8_dtype=te.DType.kFloat8E4M3, - ) - test = quantizer(test) - elif quantization == "fp8_current_scaling": - quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=te.DType.kFloat8E4M3, - device=test_device, - ) - test = quantizer(test) - elif quantization == "fp8_blockwise": - quantizer = Float8BlockQuantizer( - fp8_dtype=te.DType.kFloat8E4M3, - rowwise=True, - columnwise=True, - force_pow_2_scales=True, - amax_epsilon=0.0, - block_scaling_dim=1, - ) - test = quantizer(test) - elif quantization == "mxfp8": - test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): - row_scaled_nvfp4 = quantization == "nvfp4_row_scaled" - test = NVFP4Quantizer( - columnwise=not row_scaled_nvfp4, - with_rht=False, - with_post_rht_amax=False, - with_2d_quantization=False, - stochastic_rounding=False, - row_scaled_nvfp4=row_scaled_nvfp4, - with_random_sign_mask=False, - nvfp4_use_4over6=(quantization == "nvfp4_4over6"), - )(test) else: - raise ValueError(f"Unsupported quantization scheme ({quantization})") + test = make_quantizer(quantization, device=test_device)(test) # Make sure reference and test tensors match each other ref.copy_(test.to(dtype=ref.dtype)) @@ -818,6 +817,35 @@ def test_update_nd_tensor( assert q_x.shape == torch.Size(shape) assert_close(q_x, q_ref, rtol=0, atol=0) + @pytest.mark.parametrize("quantization", _quantization_list) + def test_quantize_dequantize_autograd( + self, + *, + quantization: str, + shape: Iterable[int] = (128, 128), + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + ) -> None: + """Autograd must survive a quantize -> dequantize round trip.""" + + quantizer = make_quantizer(quantization, device=device) + x = torch.randn(list(shape), dtype=dtype, device=device, requires_grad=True) + # Quantize with autograd enabled: a grad_fn is attached to the output. + x_q = quantizer(x) + assert isinstance(x_q, QuantizedTensor) + assert x_q.grad_fn is not None, "quantized tensor is missing its grad_fn" + # requires_grad must reflect the attached grad_fn, not a stale cache. + assert x_q.requires_grad, ( + "quantized tensor reports requires_grad=False despite having a " + "grad_fn (stale requires_grad cache)" + ) + + # Dequantize and take a loss; the gradient must reach the input. + (x_q.dequantize().float() ** 2).sum().backward() + assert ( + x.grad is not None and x.grad.norm().item() > 0 + ), "Gradient did not flow back to the input through quantize -> dequantize" + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) class TestMXFP8Tensor: diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index cfe488aae5..48fbe18590 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -469,35 +469,6 @@ def dtype(self, value: torch.dtype) -> None: """Set dtype property""" self._dtype = value - @property - def requires_grad(self) -> bool: - """ - Return whether or not the tensor requires gradient. - Attribute access of custom tensors goes through an - expensive Pyobject lookup. Since requires_grad is set during - initialization and may be updated, we cache it in a member variable. - """ - # Fallback to parent if not cached yet - if not hasattr(self, "_requires_grad"): - # pylint: disable=unnecessary-dunder-call - self._requires_grad = torch._C.TensorBase.requires_grad.__get__(self, type(self)) - return self._requires_grad - - @requires_grad.setter - def requires_grad(self, value: bool) -> None: - """Set requires_grad property so that autograd engine is aware of the change""" - # Update the cached value and call parent class method to ensure autograd engine is aware - self.requires_grad_(value) - - def requires_grad_(self, requires_grad: bool = True) -> QuantizedTensor: - """Cache requires_grad property and call parent class method""" - # pylint: disable=missing-function-docstring - # Update the cached value - self._requires_grad = requires_grad - # Call parent class method to ensure autograd engine is aware - super().requires_grad_(requires_grad) - return self - def _get_data(self) -> torch.Tensor: """Get tensor data property""" return super().data @@ -763,13 +734,7 @@ def maybe_update_inplace(arg, new_arg, schema_arg): out = super().__torch_dispatch__(func, types, args, kwargs) return out - @classmethod - def __torch_function__(cls, func, types, args=(), kwargs=None): - if kwargs is None: - kwargs = {} - - # Do not force the QuantizedTensor type on the returned tensor - return torch._C._disabled_torch_function_impl(func, types, args, kwargs) + __torch_function__ = torch._C._disabled_torch_function_impl def contiguous( self, memory_format: torch.memory_format = torch.contiguous_format diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index d2d28aecfb..6e6851877b 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -12,7 +12,10 @@ import torch import transformer_engine_torch as tex from transformer_engine.common.recipe import Float8BlockScaling, Recipe -from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage +from .storage.float8_blockwise_tensor_storage import ( + Float8BlockwiseQTensorStorage, + _FromFloat8BlockwiseFunc, +) from ..quantized_tensor import QuantizedTensor, Quantizer from ._quantization_helpers import _IdentityFunc, safe_quantized_repr from ..constants import DType @@ -312,7 +315,9 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: dequant_dtype = dtype else: dequant_dtype = self.dtype - return super().dequantize(dtype=dequant_dtype) + if torch.is_grad_enabled(): + return _FromFloat8BlockwiseFunc.apply(self, dequant_dtype) + return _FromFloat8BlockwiseFunc.forward(None, self, dequant_dtype) def detach(self) -> Float8BlockwiseQTensor: # pylint: disable=missing-function-docstring diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 993ead42ee..b24a4e9144 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -19,6 +19,96 @@ from ...utils import _empty_tensor +class _FromFloat8BlockwiseFunc(torch.autograd.Function): + """Cast from Float8 blockwise to other dtype""" + + @staticmethod + def forward( + _ctx: Optional[torch.autograd.function.FunctionCtx], # unused + tensor: Float8BlockwiseQTensorStorage, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + # pylint: disable=missing-function-docstring + if dtype is None: + dtype = tensor._dtype + + if tensor._rowwise_data is not None and tensor._rowwise_data.numel() == 0: + return torch.empty(tensor.size(), dtype=dtype, device=tensor.device) + + block_len = 128 + if not tensor._is_2D_scaled: + return tensor._dequantize_vectorwise(dtype=dtype) + + def format_scale_as_logical_shape(q_K, scales, block_len): + # The GEMM for 2D blocks required padding in the scales. + derived_scale_k_shape = math.ceil(q_K / block_len) + _, scale_K = scales.shape + if derived_scale_k_shape == scale_K: + return scales + return scales[:, :derived_scale_k_shape].contiguous() + + q_M, q_K = 1, 1 + if tensor._rowwise_data is not None: + q = tensor._rowwise_data + scale_inv = tensor._rowwise_scale_inv + transpose_output = False + if len(q.shape) >= 1: + q_K = q.shape[-1] + for i in range(len(q.shape) - 1): + q_M *= q.shape[i] + else: + assert tensor._columnwise_data is not None, "No data to dequantize" + q = tensor._columnwise_data + scale_inv = tensor._columnwise_scale_inv + transpose_output = True + if len(q.shape) >= 1: + q_M = q.shape[0] + for i in range(1, len(q.shape)): + q_K *= q.shape[i] + + orig_shape = q.shape + q = q.reshape(q_M, q_K) + formatted_scales = format_scale_as_logical_shape(q_K, scale_inv, block_len) + assert len(formatted_scales.shape) == 2 + m_tiles, k_tiles = formatted_scales.shape + unpadded_m, unpadded_k = q_M, q_K + m_block_len = block_len + k_block_len = block_len + if q_M % m_block_len != 0 or q_K % k_block_len != 0: + m_pad_amount = (m_block_len - (q_M % m_block_len)) % m_block_len + k_pad_amount = (k_block_len - (q_K % k_block_len)) % k_block_len + q = torch.nn.functional.pad( + q, (0, k_pad_amount, 0, m_pad_amount), mode="constant", value=0 + ).contiguous() + padded_M, padded_K = q.shape + q_tiled = q.reshape(m_tiles, m_block_len, k_tiles, k_block_len) + + torch_q_dtype = TE_DType_To_Torch[tensor._fp8_dtype] + + result = q_tiled.view(torch_q_dtype).to(torch.float32) * formatted_scales.view( + m_tiles, 1, k_tiles, 1 + ) + result = result.view(padded_M, padded_K).to(dtype) + if padded_M != unpadded_m or padded_K != unpadded_k: + result = result[:unpadded_m, :unpadded_k] + if len(orig_shape) == 0: + result = result.reshape([]) + else: + result = result.reshape(*orig_shape).contiguous() + if transpose_output: + return tensor._transpose_dq_columnwise_output(result) + return result + + @staticmethod + def backward( + _ctx: torch.autograd.function.FunctionCtx, # unused + grad: torch.Tensor, + ) -> Tuple[Optional[torch.Tensor], ...]: + # pylint: disable=missing-function-docstring + # Assume that we want gradients in full precision + return grad, None + + class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): """Mixin class that holds data attributes of Float8BlockwiseQTensor. @@ -220,75 +310,7 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """ Construct plain PyTorch tensor from Float8BlockwiseQTensor """ - if dtype is None: - dtype = self._dtype - - if self._rowwise_data is not None and self._rowwise_data.numel() == 0: - return torch.empty(self.size(), dtype=dtype, device=self.device) - - block_len = 128 - if not self._is_2D_scaled: - return self._dequantize_vectorwise(dtype=dtype) - - def format_scale_as_logical_shape(q_K, scales, block_len): - # The GEMM for 2D blocks required padding in the scales. - derived_scale_k_shape = math.ceil(q_K / block_len) - _, scale_K = scales.shape - if derived_scale_k_shape == scale_K: - return scales - return scales[:, :derived_scale_k_shape].contiguous() - - q_M, q_K = 1, 1 - if self._rowwise_data is not None: - q = self._rowwise_data - scale_inv = self._rowwise_scale_inv - transpose_output = False - if len(q.shape) >= 1: - q_K = q.shape[-1] - for i in range(len(q.shape) - 1): - q_M *= q.shape[i] - else: - assert self._columnwise_data is not None, "No data to dequantize" - q = self._columnwise_data - scale_inv = self._columnwise_scale_inv - transpose_output = True - if len(q.shape) >= 1: - q_M = q.shape[0] - for i in range(1, len(q.shape)): - q_K *= q.shape[i] - - orig_shape = q.shape - q = q.reshape(q_M, q_K) - formatted_scales = format_scale_as_logical_shape(q_K, scale_inv, block_len) - assert len(formatted_scales.shape) == 2 - m_tiles, k_tiles = formatted_scales.shape - unpadded_m, unpadded_k = q_M, q_K - m_block_len = block_len - k_block_len = block_len - if q_M % m_block_len != 0 or q_K % k_block_len != 0: - m_pad_amount = (m_block_len - (q_M % m_block_len)) % m_block_len - k_pad_amount = (k_block_len - (q_K % k_block_len)) % k_block_len - q = torch.nn.functional.pad( - q, (0, k_pad_amount, 0, m_pad_amount), mode="constant", value=0 - ).contiguous() - padded_M, padded_K = q.shape - q_tiled = q.reshape(m_tiles, m_block_len, k_tiles, k_block_len) - - torch_q_dtype = TE_DType_To_Torch[self._fp8_dtype] - - result = q_tiled.view(torch_q_dtype).to(torch.float32) * formatted_scales.view( - m_tiles, 1, k_tiles, 1 - ) - result = result.view(padded_M, padded_K).to(dtype) - if padded_M != unpadded_m or padded_K != unpadded_k: - result = result[:unpadded_m, :unpadded_k] - if len(orig_shape) == 0: - result = result.reshape([]) - else: - result = result.reshape(*orig_shape).contiguous() - if transpose_output: - return self._transpose_dq_columnwise_output(result) - return result + return _FromFloat8BlockwiseFunc.forward(None, self, dtype) def size(self, *args, **kwargs): # pylint: disable=missing-function-docstring