diff --git a/lib/python3.12/site-packages/torchao/dtypes/__init__.py b/lib/python3.12/site-packages/torchao/dtypes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eb253c11bc3f41f9eb4378e58246f30b8334feec --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/__init__.py @@ -0,0 +1,64 @@ +from . import affine_quantized_tensor_ops +from .affine_quantized_tensor import ( + AffineQuantizedTensor, + to_affine_quantized_floatx, + to_affine_quantized_floatx_static, + # experimental, will be merged into floatx in the future + to_affine_quantized_fpx, + to_affine_quantized_intx, + to_affine_quantized_intx_static, +) +from .floatx import ( + CutlassSemiSparseLayout, + Float8Layout, +) +from .nf4tensor import NF4Tensor, to_nf4 +from .uintx import ( + BlockSparseLayout, + CutlassInt4PackedLayout, + Int4CPULayout, + Int4XPULayout, + MarlinQQQLayout, + MarlinQQQTensor, + MarlinSparseLayout, + PackedLinearInt8DynamicActivationIntxWeightLayout, + QDQLayout, + SemiSparseLayout, + TensorCoreTiledLayout, + UintxLayout, + to_marlinqqq_quantized_intx, +) +from .utils import ( + Layout, + PlainLayout, +) + +__all__ = [ + "NF4Tensor", + "to_nf4", + "AffineQuantizedTensor", + "to_affine_quantized_intx", + "to_affine_quantized_intx_static", + "to_affine_quantized_fpx", + "to_affine_quantized_floatx", + "to_affine_quantized_floatx_static", + "to_marlinqqq_quantized_intx", + "Layout", + "PlainLayout", + "SemiSparseLayout", + "TensorCoreTiledLayout", + "Float8Layout", + "MarlinSparseLayout", + "affine_quantized_tensor_ops", + "BlockSparseLayout", + "UintxLayout", + "MarlinQQQTensor", + "MarlinQQQLayout", + "Int4CPULayout", + "CutlassInt4PackedLayout", + "CutlassSemiSparseLayout", + "QDQLayout", + "PackedLinearInt8DynamicActivationIntxWeightLayout", + "to_affine_quantized_packed_linear_int8_dynamic_activation_intx_weight", + "Int4XPULayout", +] diff --git a/lib/python3.12/site-packages/torchao/dtypes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5dd94f5592bfb31c6c62ce13e81d4f5032df31bb Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/__pycache__/_nf4tensor_api.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/_nf4tensor_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49e4fbaa0c1b062d331ee30afefa4f46b28fedd7 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/_nf4tensor_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/__pycache__/affine_quantized_tensor.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/affine_quantized_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7761fb7ea2db85bbfdf03e522a37e42dbcafcd87 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/affine_quantized_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/__pycache__/affine_quantized_tensor_ops.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/affine_quantized_tensor_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf5dba53f627868f676eac4f301af4ead265e4b0 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/affine_quantized_tensor_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/__pycache__/nf4tensor.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/nf4tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aebf152fa54c6056e75dabfea83096fc1cd3ed1a Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/nf4tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a1621af189f0ae513db2a11312f5441eaa46bd5 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/_nf4tensor_api.py b/lib/python3.12/site-packages/torchao/dtypes/_nf4tensor_api.py new file mode 100644 index 0000000000000000000000000000000000000000..27261856705408576504c03409f7a13ed26b26ec --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/_nf4tensor_api.py @@ -0,0 +1,39 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch + +from torchao.core.config import AOBaseConfig +from torchao.dtypes.nf4tensor import NF4Tensor +from torchao.quantization.transform_module import ( + register_quantize_module_handler, +) + + +class NF4WeightOnlyConfig(AOBaseConfig): + """ + Note: the file location of this workflow is temporary. + TODO(future PR): integrate this properly into torchao's directory structure + """ + + block_size: int = 64 + scaler_block_size: int = 256 + + +# for bc +nf4_weight_only = NF4WeightOnlyConfig + + +@register_quantize_module_handler(NF4WeightOnlyConfig) +def _nf4_weight_only_transform( + module: torch.nn.Module, + config: NF4WeightOnlyConfig, +) -> torch.nn.Module: + block_size = config.block_size + scaler_block_size = config.scaler_block_size + + new_weight = NF4Tensor.from_tensor(module.weight, block_size, scaler_block_size) + module.weight = torch.nn.Parameter(new_weight, requires_grad=False) + return module diff --git a/lib/python3.12/site-packages/torchao/dtypes/affine_quantized_tensor.py b/lib/python3.12/site-packages/torchao/dtypes/affine_quantized_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..beaac8b0e18fbac6e5de013fc80715ab41c198a7 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/affine_quantized_tensor.py @@ -0,0 +1,507 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import logging +import math +from typing import Optional, Tuple, Union + +import torch + +from torchao.dtypes.utils import ( + AQTTensorImpl, + Layout, + PlainLayout, +) +from torchao.quantization.quant_primitives import ( + FP8_TYPES, + MappingType, + ZeroPointDomain, + choose_qparams_affine, + choose_qparams_affine_floatx, + choose_qparams_and_quantize_affine_hqq, + dequantize_affine, + dequantize_affine_floatx, + quantize_affine, + quantize_affine_floatx, +) +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_5, + TorchAOBaseTensor, +) + +logger = logging.getLogger(__name__) +aten = torch.ops.aten + +__all__ = [ + "AffineQuantizedTensor", + "register_layout", + "to_affine_quantized_intx", + "to_affine_quantized_floatx", + "to_affine_quantized_intx_static", + "to_affine_quantized_floatx_static", + "to_affine_quantized_fpx", +] + + +############################## +# Tensor Subclass Definition # +############################## +class AffineQuantizedTensor(TorchAOBaseTensor): + """Affine quantized tensor subclass. Affine quantization means we quantize the floating point tensor with an affine transformation: + quantized_tensor = float_tensor / scale + zero_point + + To see what happens during choose_qparams, quantization and dequantization for affine quantization, + please checkout https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_primitives.py + and check the three quant primitive ops: choose_qparams_affine, quantize_affine qand dequantize_affine + + The shape and dtype of the tensor subclass represent how the tensor subclass looks externally, + regardless of the internal representation's type or orientation. + + fields: + - tensor_impl (AQTTensorImpl): tensor that serves as a general tensor impl storage for the quantized data, + e.g. storing plain tensors (int_data, scale, zero_point) or packed formats depending on device and operator/kernel + - block_size (Tuple[int, ...]): granularity of quantization, this means the size of the tensor elements that's sharing the same qparam + e.g. when size is the same as the input tensor dimension, we are using per tensor quantization + - shape (torch.Size): the shape for the original high precision Tensor + - quant_min (Optional[int]): minimum quantized value for the Tensor, if not specified, it will be derived from dtype of `int_data` + - quant_max (Optional[int]): maximum quantized value for the Tensor, if not specified, it will be derived from dtype of `int_data` + - zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be either integer or float + if zero_point is in integer domain, zero point is added to the quantized integer value during quantization + if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) value during quantization + default is ZeroPointDomain.INT + - dtype: dtype for original high precision tensor, e.g. torch.float32 + """ + + @staticmethod + def __new__( + cls, + tensor_impl: AQTTensorImpl, + block_size: Tuple[int, ...], + shape: torch.Size, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + dtype=None, + strides=None, + ): + if zero_point_domain is None: + raise ValueError("please use ZeroPointDomain.NONE instead of None") + kwargs = {} + kwargs["device"] = tensor_impl.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else tensor_impl.layout + ) + kwargs["dtype"] = dtype + if strides is not None: + kwargs["strides"] = strides + kwargs["requires_grad"] = False + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + tensor_impl: AQTTensorImpl, + block_size: Tuple[int, ...], + shape: torch.Size, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + dtype=None, + strides=None, + ): + self.tensor_impl = tensor_impl + self.block_size = block_size + self.quant_min = quant_min + self.quant_max = quant_max + self.zero_point_domain = zero_point_domain + + def __repr__(self): + return ( + f"{self.__class__.__name__}(tensor_impl={self.tensor_impl}, block_size={self.block_size}, " + f"shape={self.shape}, device={self.device}, dtype={self.dtype}, requires_grad={self.requires_grad})" + ) + + def _quantization_type(self): + return f"shape={self.shape}, block_size={self.block_size}, device={self.device}, _layout={self._layout}, tensor_impl_dtype={self.tensor_impl.dtype}, quant_min={self.quant_min}, quant_max={self.quant_max}" + + def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if output_dtype is None: + output_dtype = self.dtype + + from torchao.dtypes.floatx import FloatxTensorCoreLayout + + if isinstance(self._layout, FloatxTensorCoreLayout): + int_data, scale = self.tensor_impl.get_plain() + return dequantize_affine_floatx( + int_data, + scale, + self._layout.ebits, + self._layout.mbits, + output_dtype=output_dtype, + ) + else: + data, scale, zero_point = self.tensor_impl.get_plain() + dq = dequantize_affine( + data, + self.block_size, + scale, + zero_point, + data.dtype, + self.quant_min, + self.quant_max, + self.zero_point_domain, + output_dtype=output_dtype, + ) + from torchao.dtypes.uintx import TensorCoreTiledLayout + + if isinstance(self._layout, TensorCoreTiledLayout): + # need to return to original shape if tensor was padded + # in preprocessing + # TODO: we could add an API for this if there are more use cases + # (e.g. dequant_post_process) in TensorImpl or Layout + for dim, dim_size in enumerate(self.shape): + dq = dq.narrow(dim, 0, dim_size) + return dq + + def __tensor_flatten__(self): + # This is used in rumtime to unwrap AffineQuantizedTensor activations. + # AffineQuantizedTensor has __torch_function__ override: + # Each getattr will go through it, which is up to 10x slower than default attribute access. + with torch._C.DisableTorchFunctionSubclass(): + return ["tensor_impl"], [ + self.block_size, + self.shape, + self.quant_min, + self.quant_max, + self.zero_point_domain, + self.dtype, + ] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + tensor_impl = tensor_data_dict["tensor_impl"] + block_size, shape, quant_min, quant_max, zero_point_domain, dtype = ( + tensor_attributes + ) + return cls( + tensor_impl, + block_size, + shape if outer_size is None else outer_size, + quant_min, + quant_max, + zero_point_domain, + dtype=dtype, + strides=outer_stride, + ) + + @classmethod + def from_hp_to_intx( + cls, + input_float: torch.Tensor, + mapping_type: MappingType, + block_size: Tuple[int, ...], + target_dtype: torch.dtype, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + _layout: Layout = PlainLayout(), + use_hqq: bool = False, + ): + """Convert a high precision tensor to an integer affine quantized tensor.""" + original_shape = input_float.shape + input_float = _layout.pre_process(input_float) + + if use_hqq: + assert ( + zero_point_domain == ZeroPointDomain.FLOAT + and mapping_type == MappingType.ASYMMETRIC + and quant_min == 0 + ), "Invalid input parameters for HQQ quantization." + nbits = int(math.log2(quant_max + 1)) + axis = 1 if (block_size[0] == 1) else 0 + group_size = max(block_size) + compute_dtype = ( + zero_point_dtype + if (zero_point_dtype is not None) + else input_float.dtype + ) + device = input_float.device + from torchao.dtypes import Int4CPULayout + from torchao.dtypes.uintx import TensorCoreTiledLayout + + data, scale, zero_point, _ = choose_qparams_and_quantize_affine_hqq( + input_float, + nbits=nbits, + group_size=group_size, + axis=axis, + compute_dtype=compute_dtype, + device=device, + verbose=False, + raw_output=not isinstance( + _layout, (TensorCoreTiledLayout, PlainLayout, Int4CPULayout) + ), + # raw_output=False is basically the 'convert to TensorCoreTiledLayout zero_point version' option (add scale*midpoint) + # note in choose_qparams_affine, preserve_zero = False does this same thing while also controlling whether + # zero is preserved. + # TODO uncouple preserve_zero and conversion of zero_point to TensorCoreTiledLayout version + # TODO move the conversion of zero_point out of quant_primitives and into TensorCoreTiledLayout.from_plain + # TODO change PlainLayout to use raw_output. + ) + data = data.to(target_dtype) + else: + scale, zero_point = choose_qparams_affine( + input_float, + mapping_type, + block_size, + target_dtype, + quant_min, + quant_max, + eps, + scale_dtype, + zero_point_dtype, + preserve_zero, + zero_point_domain, + ) + # choose_qparams_affine is a custom op that does support returning optional Tensors. We thus set the zero_point to None if its domain is None + if zero_point_domain == ZeroPointDomain.NONE: + zero_point = None + data = quantize_affine( + input_float, + block_size, + scale, + zero_point, + target_dtype, + quant_min, + quant_max, + zero_point_domain, + ) + # Note: output will be uint8 tensor for sub byte tensors for now + + data, scale, zero_point = _layout.post_process( + data, scale, zero_point, block_size + ) + tensor_impl_ctr = get_tensor_impl_constructor(type(_layout)) + tensor_impl = tensor_impl_ctr(data, scale, zero_point, _layout) + return cls( + tensor_impl, + block_size, + original_shape, + quant_min, + quant_max, + zero_point_domain, + dtype=input_float.dtype, + ) + + @classmethod + def from_hp_to_intx_static( + cls, + input_float: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + block_size: Tuple[int, ...], + target_dtype: torch.dtype, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + _layout: Layout = PlainLayout(), + ): + """Create an integer AffineQuantizedTensor from a high precision tensor using static parameters.""" + if zero_point_domain is None: + raise ValueError("please use ZeroPointDomain.NONE instead of None") + elif zero_point_domain is ZeroPointDomain.NONE and zero_point is not None: + raise ValueError("zero_point should be None when zero_point_domain is NONE") + if target_dtype not in FP8_TYPES: + assert zero_point is not None, ( + "zero_point must be specified for non-fp8 types" + ) + original_shape = input_float.shape + input_float, scale, zero_point = _layout.pre_process_static( + input_float, scale, zero_point, block_size + ) + + int_data = quantize_affine( + input_float, + block_size, + scale, + zero_point, + target_dtype, + quant_min, + quant_max, + zero_point_domain, + ) + + int_data, scale, zero_point = _layout.post_process( + int_data, + scale, + zero_point, + block_size, + ) + + tensor_impl_ctr = get_tensor_impl_constructor(type(_layout)) + tensor_impl = tensor_impl_ctr(int_data, scale, zero_point, _layout) + return cls( + tensor_impl, + block_size, + original_shape, + quant_min, + quant_max, + zero_point_domain, + dtype=input_float.dtype, + ) + + @classmethod + def from_hp_to_floatx( + cls, + input_float: torch.Tensor, + block_size: Tuple[int, ...], + target_dtype: torch.dtype, + _layout: Layout, + scale_dtype: Optional[torch.dtype] = None, + ): + """Convert a high precision tensor to a float8 quantized tensor.""" + if target_dtype in FP8_TYPES: + return cls.from_hp_to_intx( + input_float=input_float, + mapping_type=MappingType.SYMMETRIC, + block_size=block_size, + target_dtype=target_dtype, + quant_min=math.ceil(torch.finfo(target_dtype).min), + quant_max=math.ceil(torch.finfo(target_dtype).max), + eps=torch.finfo(torch.float32).eps, + scale_dtype=scale_dtype, + zero_point_dtype=None, + preserve_zero=True, + zero_point_domain=ZeroPointDomain.NONE, + _layout=_layout, + use_hqq=False, + ) + else: + raise NotImplementedError( + f"Unsupported dtype {target_dtype} for from_hp_to_floatx" + ) + + @classmethod + def from_hp_to_floatx_static( + cls, + input_float: torch.Tensor, + scale: torch.Tensor, + block_size: Tuple[int, ...], + target_dtype: torch.dtype, + _layout: Layout, + ): + """Create a float8 AffineQuantizedTensor from a high precision tensor using static parameters.""" + if target_dtype in FP8_TYPES: + return cls.from_hp_to_intx_static( + input_float=input_float, + scale=scale, + zero_point=None, + block_size=block_size, + target_dtype=target_dtype, + quant_min=math.ceil(torch.finfo(target_dtype).min), + quant_max=math.ceil(torch.finfo(target_dtype).max), + zero_point_domain=ZeroPointDomain.NONE, + _layout=_layout, + ) + else: + raise NotImplementedError( + f"Unsupported dtype {target_dtype} for from_hp_to_floatx_static" + ) + + @classmethod + def from_hp_to_fpx( + cls, + input_float: torch.Tensor, + _layout: Layout, + ): + """Create a floatx AffineQuantizedTensor from a high precision tensor. Floatx is represented as ebits and mbits, and supports the representation of float1-float7.""" + from torchao.dtypes.floatx import FloatxTensorCoreLayout + + assert isinstance(_layout, FloatxTensorCoreLayout), ( + f"Only FloatxTensorCoreLayout is supported for floatx, got {_layout}" + ) + original_shape = input_float.shape + input_float = _layout.pre_process(input_float) + # per axis quantization, where axis = 1 + block_size = list(input_float.shape) + block_size[1] = 1 + + ebits, mbits = _layout.ebits, _layout.mbits + # Note: these ops are hardcoded to have per axis quantization (axis=1) right now + scale = choose_qparams_affine_floatx(input_float, ebits, mbits) + floatx_unpacked = quantize_affine_floatx(input_float, scale, ebits, mbits) + floatx_packed, scale, _ = _layout.post_process( + floatx_unpacked, scale, None, block_size + ) + + tensor_impl_ctr = get_tensor_impl_constructor(type(_layout)) + tensor_impl = tensor_impl_ctr(floatx_packed, scale, None, _layout) + return cls(tensor_impl, block_size, original_shape, dtype=input_float.dtype) + + @property + def _layout(self) -> Layout: + return self.tensor_impl._layout + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs.pop("device") + return self.__class__( + self.tensor_impl.to(device), + self.block_size, + self.shape, + self.quant_min, + self.quant_max, + self.zero_point_domain, + **kwargs, + ) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.tensor_impl), + self.block_size, + self.shape, + self.quant_min, + self.quant_max, + self.zero_point_domain, + dtype=self.dtype, + strides=self.stride(), + ) + + # following are the comments for __torch_function__/__torch_dispatch__, -> this is defined in affine_quantized_tensor_ops.py + # a bit later + # Note: we only added cpu path here for 8da4w, this is for executorch, in the future + # 1. we'll add cpu/cuda version (int4mm etc.) + # 2. we'll need to hide the 8da4w executorch version under things like layouts (we also have multiple impl for cpu kernel as Michael mentioned), so it will be something like + # cpu device + et laytout --> gives current 8da4w executorch representation + # cpu device + avx layout --> gives optimized kernel for 8da4w in avx cpu etc. + # cuda device + some layout --> gives cuda kernel + + # two scenarios where we currently fall back to vanilla mm: + # 1 - when tensor is on CUDA: we'll add this later, we'll also enable dispatching to optimized + # kernels in CPU as well, see the note above + # 2 - we're given non-floats - quantizing long to int8 is crazy + + +###################################################### +# Layout and TensorImpl Subclass Registration # +###################################################### +register_layout = AffineQuantizedTensor.register_layout +get_tensor_impl_constructor = AffineQuantizedTensor.get_tensor_impl_constructor + + +to_affine_quantized_intx = AffineQuantizedTensor.from_hp_to_intx +to_affine_quantized_intx_static = AffineQuantizedTensor.from_hp_to_intx_static +to_affine_quantized_floatx = AffineQuantizedTensor.from_hp_to_floatx +to_affine_quantized_floatx_static = AffineQuantizedTensor.from_hp_to_floatx_static +# experimental will be merged in to floatx +to_affine_quantized_fpx = AffineQuantizedTensor.from_hp_to_fpx + +if TORCH_VERSION_AT_LEAST_2_5: + # Allow a model with AffineQuantizedTensor weights to be loaded with `weights_only=True` + torch.serialization.add_safe_globals([AffineQuantizedTensor]) diff --git a/lib/python3.12/site-packages/torchao/dtypes/affine_quantized_tensor_ops.py b/lib/python3.12/site-packages/torchao/dtypes/affine_quantized_tensor_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..1d70f5c7f33acb0cf9d1922aa1c1fb3fe1e94859 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/affine_quantized_tensor_ops.py @@ -0,0 +1,580 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import logging + +import torch +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, +) +from torchao.dtypes.floatx.cutlass_semi_sparse_layout import ( + _linear_fp8_act_fp8_weight_sparse_cutlass_check, + _linear_fp8_act_fp8_weight_sparse_cutlass_impl, +) +from torchao.dtypes.floatx.float8_layout import ( + _linear_fp8_act_fp8_weight_check, + _linear_fp8_act_fp8_weight_impl, + _linear_fp_act_fp8_weight_check, + _linear_fp_act_fp8_weight_impl, +) +from torchao.dtypes.floatx.floatx_tensor_core_layout import ( + _linear_f16_bf16_act_floatx_weight_check, + _linear_f16_bf16_act_floatx_weight_impl, +) +from torchao.dtypes.uintx.block_sparse_layout import ( + _linear_int8_act_int8_weight_block_sparse_check, + _linear_int8_act_int8_weight_block_sparse_impl, +) +from torchao.dtypes.uintx.cutlass_int4_packed_layout import ( + _linear_int4_act_int4_weight_cutlass_check, + _linear_int4_act_int4_weight_cutlass_impl, + _linear_int8_act_int4_weight_cutlass_check, + _linear_int8_act_int4_weight_cutlass_impl, +) +from torchao.dtypes.uintx.gemlite_layout import ( + _linear_fp_act_int4_weight_gemlite_check, + _linear_fp_act_int4_weight_gemlite_impl, +) +from torchao.dtypes.uintx.int4_cpu_layout import ( + _linear_fp_act_uint4_weight_cpu_check, + _linear_fp_act_uint4_weight_cpu_impl, +) +from torchao.dtypes.uintx.int4_xpu_layout import ( + _linear_bf16_act_uint4_weight_float_zero_check, + _linear_bf16_act_uint4_weight_float_zero_impl, + _linear_bf16_act_uint4_weight_int8_zero_check, + _linear_bf16_act_uint4_weight_int8_zero_impl, +) +from torchao.dtypes.uintx.marlin_qqq_tensor import ( + _linear_int8_act_int4_weight_marlin_qqq_check, + _linear_int8_act_int4_weight_marlin_qqq_impl, +) +from torchao.dtypes.uintx.marlin_sparse_layout import ( + _linear_fp_act_int4_weight_sparse_marlin_check, + _linear_fp_act_int4_weight_sparse_marlin_impl, +) +from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import ( + _linear_check as _linear_int8_act_intx_weight_packed_check, +) +from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import ( + _linear_impl as _linear_int8_act_intx_weight_packed_impl, +) +from torchao.dtypes.uintx.plain_layout import ( + PlainAQTTensorImpl, + _linear_fp_act_int8_weight_check, + _linear_fp_act_int8_weight_impl, + _linear_int8_act_int8_weight_check, + _linear_int8_act_int8_weight_impl, +) +from torchao.dtypes.uintx.q_dq_layout import ( + _embedding_check as _embedding_q_dq_check, +) +from torchao.dtypes.uintx.q_dq_layout import ( + _embedding_impl as _embedding_q_dq_impl, +) +from torchao.dtypes.uintx.q_dq_layout import ( + _linear_check as _linear_q_dq_check, +) +from torchao.dtypes.uintx.q_dq_layout import ( + _linear_impl as _linear_q_dq_impl, +) +from torchao.dtypes.uintx.semi_sparse_layout import ( + _linear_int8_act_int8_weight_semi_structured_sparse_check, + _linear_int8_act_int8_weight_semi_structured_sparse_impl, +) +from torchao.dtypes.uintx.tensor_core_tiled_layout import ( + _linear_bf16_act_uint4_weight_check, + _linear_bf16_act_uint4_weight_impl, +) +from torchao.quantization.quant_primitives import dequantize_affine +from torchao.utils import ( + fill_defaults, +) + +logger = logging.getLogger(__name__) + + +aten = torch.ops.aten + + +_AQT_QLINEAR_DISPATCH_TABLE = {} + + +def register_aqt_quantized_linear_dispatch(dispatch_condition, impl): + """Register a dispatch for quantized linear op with dispatch_condition function and impl function + both takes three arguments: + input_tensor: dimension is (M1, M2, ..., in_features) + weight_tensor: dimension is (out_features, in_features) + bias: dimension is (out_features,) + so that these can be shared by F.linear, aten.mm, aten.addmm dispatches + + Args: + `dispatch_condition` (Callable[[torch.Tensor, torch.Tensor, torch.Tensor], bool]: the dispatch + condition for a specialized quantized linear implementation, e.g. bfloat16 activation + uint4 weight + `impl` (Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]: the specialized + quantized linear implementation + """ + _AQT_QLINEAR_DISPATCH_TABLE[dispatch_condition] = impl + + +def deregister_aqt_quantized_linear_dispatch(dispatch_condition): + if dispatch_condition in _AQT_QLINEAR_DISPATCH_TABLE: + del _AQT_QLINEAR_DISPATCH_TABLE[dispatch_condition] + else: + logger.warn( + f"Attempting to remove non-existant dispatch condition {dispatch_condition}" + ) + + +def _same_metadata(self: AffineQuantizedTensor, src: AffineQuantizedTensor): + return ( + isinstance(self, AffineQuantizedTensor) + and isinstance(src, AffineQuantizedTensor) + and all( + [ + getattr(self, attr) == getattr(src, attr) + for attr in [ + "block_size", + "shape", + "quant_min", + "quant_max", + "zero_point_domain", + "dtype", + ] + ] + ) + and isinstance(self.tensor_impl, type(src.tensor_impl)) + ) + + +class QuantizedLinearNotImplementedError(NotImplementedError): + """Thin wrapper around NotImplementedError to make it easier to catch this error in the dispatch table""" + + pass + + +# input_tensor: dimension is (M1, M2, ..., in_features) +# weight_tensor: dimension is (out_features, in_features) +# bias: dimension is (out_features,) +@staticmethod +def _quantized_linear_op(input_tensor, weight_tensor, bias): + for dispatch_condition, impl in _AQT_QLINEAR_DISPATCH_TABLE.items(): + if dispatch_condition(input_tensor, weight_tensor, bias): + return impl(input_tensor, weight_tensor, bias) + raise QuantizedLinearNotImplementedError( + "No specialized dispatch found for quantized linear op" + ) + + +# Attach the _quantized_linear_op to the AffineQuantizedTensor class +AffineQuantizedTensor._quantized_linear_op = _quantized_linear_op + + +# _register_aqt_quantized_linear_dispatches function has a list of (dispatch_condition, implementation) functions, defined in their dtype layout classes, that takes the following args: +# input_tensor: dimension is (M1, M2, ..., in_features) +# weight_tensor: dimension is (out_features, in_features) +# bias: dimension is (out_features,) +# so that these can be shared by F.linear, aten.mm, aten.addmm dispatches +def _register_aqt_quantized_linear_dispatches(): + for dispatch_condition, impl in [ + (_linear_int8_act_int8_weight_check, _linear_int8_act_int8_weight_impl), + ( + _linear_int8_act_int8_weight_semi_structured_sparse_check, + _linear_int8_act_int8_weight_semi_structured_sparse_impl, + ), + ( + _linear_int8_act_int8_weight_block_sparse_check, + _linear_int8_act_int8_weight_block_sparse_impl, + ), + (_linear_fp8_act_fp8_weight_check, _linear_fp8_act_fp8_weight_impl), + (_linear_fp_act_fp8_weight_check, _linear_fp_act_fp8_weight_impl), + (_linear_bf16_act_uint4_weight_check, _linear_bf16_act_uint4_weight_impl), + (_linear_fp_act_int8_weight_check, _linear_fp_act_int8_weight_impl), + ( + _linear_f16_bf16_act_floatx_weight_check, + _linear_f16_bf16_act_floatx_weight_impl, + ), + ( + _linear_fp_act_int4_weight_sparse_marlin_check, + _linear_fp_act_int4_weight_sparse_marlin_impl, + ), + ( + _linear_int8_act_int4_weight_marlin_qqq_check, + _linear_int8_act_int4_weight_marlin_qqq_impl, + ), + ( + _linear_fp_act_int4_weight_gemlite_check, + _linear_fp_act_int4_weight_gemlite_impl, + ), + ( + _linear_int8_act_int4_weight_cutlass_check, + _linear_int8_act_int4_weight_cutlass_impl, + ), + ( + _linear_int4_act_int4_weight_cutlass_check, + _linear_int4_act_int4_weight_cutlass_impl, + ), + ( + _linear_fp8_act_fp8_weight_sparse_cutlass_check, + _linear_fp8_act_fp8_weight_sparse_cutlass_impl, + ), + ( + _linear_fp_act_uint4_weight_cpu_check, + _linear_fp_act_uint4_weight_cpu_impl, + ), + ( + _linear_int8_act_intx_weight_packed_check, + _linear_int8_act_intx_weight_packed_impl, + ), + ( + _linear_q_dq_check, + _linear_q_dq_impl, + ), + ( + _linear_bf16_act_uint4_weight_int8_zero_check, + _linear_bf16_act_uint4_weight_int8_zero_impl, + ), + ( + _linear_bf16_act_uint4_weight_float_zero_check, + _linear_bf16_act_uint4_weight_float_zero_impl, + ), + ]: + register_aqt_quantized_linear_dispatch(dispatch_condition, impl) + + +_register_aqt_quantized_linear_dispatches() + +implements = AffineQuantizedTensor.implements + + +@implements([torch.nn.functional.linear, aten.linear.default]) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + if not input_tensor.is_floating_point(): + raise NotImplementedError( + f"{func} is not implemented for non floating point input" + ) + + # using try/except here so that we can have a general fallback when input_tensor/weight_tensor + # is not picked up by any of the dispatch paths in `_quantized_linear_op`, this allows us to + # make the branches easier to understand in `_quantized_linear_op` + try: + return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias) + except QuantizedLinearNotImplementedError as e: + # fallback path is only called when user did not specify a specfic quantized linear implementation with `_layout.quantized_linear_impl` + if ( + isinstance(weight_tensor, AffineQuantizedTensor) + and hasattr(weight_tensor._layout, "quantized_linear_impl") + and weight_tensor._layout.quantized_linear_impl is not None + ): + raise e + + if isinstance(input_tensor, AffineQuantizedTensor): + input_tensor = input_tensor.dequantize() + if isinstance(weight_tensor, AffineQuantizedTensor): + weight_tensor = weight_tensor.dequantize() + return torch.nn.functional.linear(input_tensor, weight_tensor, bias) + + +@implements(torch.nn.functional.embedding) +def _(func, types, args, kwargs): + if _embedding_q_dq_check(args, kwargs): + return _embedding_q_dq_impl(args, kwargs) + + # new_arg1 = args[1].dequantize() + # return torch.nn.embedding(args[0], new_arg1, *args[2:], **kwargs) + assert isinstance(args[1].tensor_impl, PlainAQTTensorImpl), ( + f"embedding only works with PlainAQTTensorImpl but got {type(args[1].tensor_impl)}" + ) + assert ( + kwargs["padding_idx"] is None + and kwargs["max_norm"] is None + and not kwargs["scale_grad_by_freq"] + and not kwargs["sparse"] + and kwargs["norm_type"] == 2.0 + ) + idx = args[0] + int_data, scale, zero_point = args[1].tensor_impl.get_plain() + + sliced_data, sliced_scale, sliced_zero_point = ( + int_data[idx], + scale[idx], + zero_point[idx], + ) + # Block size is expecting 2 dimensions [1, group size] but + # batchsize or other dims gets added to sliced_data, sliced_scale and sliced_zero_point so + # we need to increase block size to correct dim + new_blocks = idx.dim() - 1 + return dequantize_affine( + sliced_data, + new_blocks * [1] + list(args[1].block_size), + sliced_scale, + sliced_zero_point, + sliced_data.dtype, + args[1].quant_min, + args[1].quant_max, + args[1].zero_point_domain, + output_dtype=sliced_scale.dtype, + ) + + +@implements(aten.addmm.default) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = ( + args[1], + args[2], + args[0], + ) + if not input_tensor.is_floating_point(): + raise NotImplementedError( + f"{func} is not implemented for non floating point input" + ) + + assert input_tensor.shape[-1] == weight_tensor.shape[0], ( + f"need mat1 shape: {input_tensor.shape} final dim" + f"to match mat2 shape: {weight_tensor.shape} first dim" + ) + + # using try/except here so that we can have a general fallback when input_tensor/weight_tensor + # is not picked up by any of the dispatch paths in `_quantized_linear_op`, this allows us to + # make the branches easier to understand in `_quantized_linear_op` + try: + transposed_weight_tensor = weight_tensor.t() + return weight_tensor._quantized_linear_op( + input_tensor, transposed_weight_tensor, bias + ) + except QuantizedLinearNotImplementedError as e: + # fallback path is only called when user did not specify a specfic quantized linear implementation with `_layout.quantized_linear_impl` + if ( + isinstance(weight_tensor, AffineQuantizedTensor) + and hasattr(weight_tensor._layout, "quantized_linear_impl") + and weight_tensor._layout.quantized_linear_impl is not None + ): + raise e + + if isinstance(input_tensor, AffineQuantizedTensor): + input_tensor = input_tensor.dequantize() + if isinstance(weight_tensor, AffineQuantizedTensor): + weight_tensor = weight_tensor.dequantize() + return func(bias, input_tensor, weight_tensor) + + +@implements(aten.mm.default) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = (args[0], args[1], None) + if not input_tensor.is_floating_point(): + raise NotImplementedError( + f"{func} is not implemented for non floating point input" + ) + + assert input_tensor.shape[-1] == weight_tensor.shape[0], ( + f"need mat1 shape: {input_tensor.shape} final dim" + f"to match mat2 shape: {weight_tensor.shape} first dim" + ) + + try: + transposed_weight_tensor = weight_tensor.t() + return weight_tensor._quantized_linear_op( + input_tensor, transposed_weight_tensor, bias + ) + except QuantizedLinearNotImplementedError as e: + # fallback path is only called when user did not specify a specfic quantized linear implementation with `_layout.quantized_linear_impl` + if ( + isinstance(weight_tensor, AffineQuantizedTensor) + and hasattr(weight_tensor._layout, "quantized_linear_impl") + and weight_tensor._layout.quantized_linear_impl is not None + ): + raise e + + if isinstance(input_tensor, AffineQuantizedTensor): + input_tensor = input_tensor.dequantize() + if isinstance(weight_tensor, AffineQuantizedTensor): + weight_tensor = weight_tensor.dequantize() + return func(input_tensor, weight_tensor) + + +@implements([aten.detach.default, aten.alias.default]) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + +@implements(aten.clone.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + +@implements(aten.copy_.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[1]._apply_fn_to_data(torch.clone) + ) + + +@implements(aten._to_copy.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].to(*args[1:], **kwargs)._apply_fn_to_data(torch.clone), + ) + + +@implements(aten.copy_.default) +def _(func, types, args, kwargs): + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + + +@implements(aten.t.default) +def _(func, types, args, kwargs): + block_size = args[0].block_size + assert len(block_size) == 2 + transposed_block_size = (block_size[1], block_size[0]) + tensor = args[0] + shape = tensor.shape[::-1] + new = tensor.__class__( + tensor.tensor_impl.t(), + transposed_block_size, + shape, + tensor.quant_min, + tensor.quant_max, + tensor.zero_point_domain, + dtype=tensor.dtype, + strides=tensor.stride(), + ) + return return_and_correct_aliasing(func, args, kwargs, new) + + +@implements(aten.slice.Tensor) +def _(func, types, args, kwargs): + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + assert step == 1 + assert dim == 0 or dim == 1, f"Only dim==0 or 1 are supported, got: {dim}" + if end >= self.shape[dim]: + end = self.shape[dim] + shape = list(self.shape) + shape[dim] = end - start + block_size = self.block_size + assert len(block_size) in [ + 2, + 3, + ], f"Slice only works for 2 and 3d block_size right now, got: {block_size}" + # with slice, some shape dimension might be smaller than block_size dimension, so + # we need to make sure there is no overflow + if len(block_size) == 2: + block_size = (min(shape[0], block_size[0]), min(shape[1], block_size[1])) + + new = self.__class__( + aten.slice.Tensor(self.tensor_impl, dim, start, end, step), + block_size, + shape, + self.quant_min, + self.quant_max, + self.zero_point_domain, + dtype=self.dtype, + strides=self.stride() if len(block_size) == 2 else None, + ) + return return_and_correct_aliasing(func, args, kwargs, new) + + +@implements(aten.index.Tensor) +def _(func, types, args, kwargs): + self, indices = args + assert len(indices) == 1, ( + f"op {func} currently only implemented for single dimensional indexing but got indices: {indices}" + ) + new_tensor_impl = aten.index.Tensor(self.tensor_impl, indices) + shape = tuple([indices[0].numel(), *self.shape[1:]]) + + block_size = self.block_size + new = self.__class__( + new_tensor_impl, + block_size, + shape, + self.quant_min, + self.quant_max, + self.zero_point_domain, + dtype=self.dtype, + ) + return return_and_correct_aliasing(func, args, kwargs, new) + + +@implements(aten.select.int) +def _(func, types, args, kwargs): + self, dim, index = fill_defaults(args, 3, [0, 0]) + assert dim == 0, f"op {func} currently only implemented for dim=0 but got dim={dim}" + assert self.dim() == 3, ( + f"op {func} currently only implemented for 3 dimensional tensors but got shape={self.shape}" + ) + + new_tensor_impl = aten.select.int(self.tensor_impl, dim, index) + + shape = self.shape[1:] + block_size = self.block_size[1:] + new = self.__class__( + new_tensor_impl, + block_size, + shape, + self.quant_min, + self.quant_max, + self.zero_point_domain, + dtype=self.dtype, + ) + return return_and_correct_aliasing(func, args, kwargs, new) + + +# this is needed for DTensor.from_local() and for flattening tensor +@implements(aten.view.default) +def _(func, types, args, kwargs): + self, shape = args + + if tuple(self.shape) == tuple(shape): + return self.__class__( + self.tensor_impl, + self.block_size, + self.shape, + self.quant_min, + self.quant_max, + self.zero_point_domain, + dtype=self.dtype, + strides=self.stride(), + ) + + if len(shape) == 1 and shape[0] == -1: + assert len(self.block_size) == 2 and self.block_size[0] == 1 + block_size = (self.block_size[1],) + return self.__class__( + self.tensor_impl, + block_size, + (self.numel(),), + self.quant_min, + self.quant_max, + self.zero_point_domain, + dtype=self.dtype, + strides=self.stride(), + ) + + raise ValueError( + f"{self.__class__.__name__} only supports .view() with same shape or shape=[-1]" + ) diff --git a/lib/python3.12/site-packages/torchao/dtypes/floatx/__init__.py b/lib/python3.12/site-packages/torchao/dtypes/floatx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7e634a5211986d5b0f8ef05d798346007396cd4c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/floatx/__init__.py @@ -0,0 +1,17 @@ +from .cutlass_semi_sparse_layout import ( + CutlassSemiSparseLayout, +) +from .float8_layout import Float8Layout +from .floatx_tensor_core_layout import ( + FloatxTensorCoreLayout, + from_scaled_tc_floatx, + to_scaled_tc_floatx, +) + +__all__ = [ + "FloatxTensorCoreLayout", + "to_scaled_tc_floatx", + "from_scaled_tc_floatx", + "Float8Layout", + "CutlassSemiSparseLayout", +] diff --git a/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05493521d45d310061f390623ec2255c08c55657 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/cutlass_semi_sparse_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/cutlass_semi_sparse_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78aed361a43e8eef0a6f4bb45fd706f204962adf Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/cutlass_semi_sparse_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/float8_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/float8_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..817baea0bae904b089577b4159bd61c0b2f2847a Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/float8_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/floatx_tensor_core_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/floatx_tensor_core_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70e2bb08916d31b942d1f2492797d69f240c4e56 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/floatx/__pycache__/floatx_tensor_core_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/floatx/cutlass_semi_sparse_layout.py b/lib/python3.12/site-packages/torchao/dtypes/floatx/cutlass_semi_sparse_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..45fe4517120fde72af88f32e3d9d5b4473821f34 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/floatx/cutlass_semi_sparse_layout.py @@ -0,0 +1,212 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Optional + +import torch +from torch.utils._python_dispatch import ( + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout +from torchao.ops import ( + rowwise_scaled_linear_sparse_cutlass_f8f8, + to_sparse_semi_structured_cutlass_sm9x_f8, +) + +aten = torch.ops.aten + + +def _same_metadata( + self: "CutlassSemiSparseTensorImpl", src: "CutlassSemiSparseTensorImpl" +) -> bool: + return ( + isinstance(self, CutlassSemiSparseTensorImpl) + and isinstance(src, CutlassSemiSparseTensorImpl) + and self.shape == src.shape + and self.sparse.shape == src.sparse.shape + and self.meta.shape == src.meta.shape + and self.scale.shape == src.scale.shape + and type(self._layout) == type(src._layout) + ) + + +@dataclass(frozen=True) +class CutlassSemiSparseLayout(Layout): + """Layout class for float8 2:4 sparsity layout for affine quantized tensor, for cutlass kernel.""" + + def pre_process(self, dense: torch.Tensor) -> torch.Tensor: + # prune to 2:4 if not already + from torchao.sparsity.utils import mask_creator + + return dense * mask_creator(dense).bool() + + +@register_layout(CutlassSemiSparseLayout) +class CutlassSemiSparseTensorImpl(AQTTensorImpl): + @staticmethod + def __new__( + cls, + sparse: torch.Tensor, + meta: torch.Tensor, + scale: torch.Tensor, + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = sparse.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else sparse.layout + ) + kwargs["dtype"] = sparse.dtype + kwargs["requires_grad"] = False + shape = (sparse.shape[0], 2 * sparse.shape[-1]) + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + sparse: torch.Tensor, + meta: torch.Tensor, + scale: torch.Tensor, + _layout: Layout, + ): + self.sparse = sparse + self.meta = meta + self.scale = scale + self._layout = _layout + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + elif func is aten.copy_.default: + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + + raise NotImplementedError( + f"CutlassSemiSparseTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + def __tensor_flatten__(self): + return ["sparse", "meta", "scale"], [self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + sparse = tensor_data_dict["sparse"] + meta = tensor_data_dict["meta"] + scale = tensor_data_dict["scale"] + (_layout,) = tensor_attributes + return cls(sparse, meta, scale, _layout) + + def get_plain(self): + # No support in CUTLASS to convert back to dense from sparse + # semi-structured format, so multiplying with identity matrix, + # and using identity scale factors, for the conversion. + cols = self.shape[1] + input = torch.eye(cols, dtype=self.sparse.dtype, device=self.sparse.device) + input_scale = torch.ones( + (cols,), dtype=self.scale.dtype, device=self.sparse.device + ) + sparse_scale = torch.ones_like(self.scale) + out_dtype = torch.bfloat16 + dense = ( + rowwise_scaled_linear_sparse_cutlass_f8f8( + input, + input_scale, + self.sparse, + self.meta, + sparse_scale, + out_dtype=out_dtype, + ) + .to(self.dtype) + .t() + .contiguous() + ) + + return dense, self.scale, None + + @classmethod + def from_plain( + cls, + dense: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert zero_point is None or torch.all(zero_point == 0) + + sparse, meta = to_sparse_semi_structured_cutlass_sm9x_f8(dense) + + return cls( + sparse, + meta, + scale, + _layout, + ) + + def get_layout(self) -> Layout: + return self._layout + + def _apply_fn_to_data(self, fn): + self.sparse = fn(self.sparse) + self.meta = fn(self.meta) + self.scale = fn(self.scale) + return self + + +def _linear_fp8_act_fp8_weight_sparse_cutlass_check(input_tensor, weight_tensor, bias): + from torchao.dtypes.floatx import Float8Layout + + return ( + isinstance(input_tensor, AffineQuantizedTensor) + and isinstance(input_tensor._layout, Float8Layout) + and input_tensor.dtype in (torch.float16, torch.bfloat16) + and len(input_tensor.shape) >= 2 + and input_tensor.tensor_impl.scale.dtype == torch.float32 + and len(input_tensor.tensor_impl.scale.shape) == len(input_tensor.shape) - 1 + and isinstance(weight_tensor, AffineQuantizedTensor) + and isinstance(weight_tensor._layout, CutlassSemiSparseLayout) + and weight_tensor.dtype == input_tensor.dtype + and len(weight_tensor.shape) == 2 + and weight_tensor.tensor_impl.scale.dtype == torch.float32 + and len(weight_tensor.tensor_impl.scale.shape) == 1 + and (bias is None or bias.dtype == input_tensor.dtype) + and (bias is None or len(bias.shape) == 1) + ) + + +def _linear_fp8_act_fp8_weight_sparse_cutlass_impl(input_tensor, weight_tensor, bias): + from torchao.ops import rowwise_scaled_linear_sparse_cutlass_f8f8 + + input = input_tensor.tensor_impl.float8_data + input_scale = input_tensor.tensor_impl.scale + weight = weight_tensor.tensor_impl.sparse + weight_meta = weight_tensor.tensor_impl.meta + weight_scale = weight_tensor.tensor_impl.scale + out_dtype = input_tensor.dtype + + out = rowwise_scaled_linear_sparse_cutlass_f8f8( + input, input_scale, weight, weight_meta, weight_scale, bias, out_dtype + ) + + return out diff --git a/lib/python3.12/site-packages/torchao/dtypes/floatx/float8_layout.py b/lib/python3.12/site-packages/torchao/dtypes/floatx/float8_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..5914f00102f94fd32f308be822b8ada5ccf694f7 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/floatx/float8_layout.py @@ -0,0 +1,386 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from torch.utils._python_dispatch import ( + is_traceable_wrapper_subclass, + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout, get_out_shape +from torchao.float8.inference import ( + Float8MMConfig, + _is_rowwise_scaled, + addmm_float8_unwrapped_inference, + preprocess_data, +) +from torchao.utils import _is_float8_type, fill_defaults + +aten = torch.ops.aten + + +def _same_metadata(self: "Float8AQTTensorImpl", src: "Float8AQTTensorImpl") -> bool: + # Special handling for transposed attribute + transposed_match = (self.transposed == src.transposed) or ( + self.transposed is False and src.transposed is None + ) + + return ( + isinstance(self, Float8AQTTensorImpl) + and isinstance(src, Float8AQTTensorImpl) + and self.shape == src.shape + and self.float8_data.shape == src.float8_data.shape + and self.scale.shape == src.scale.shape + and transposed_match + and type(self._layout) == type(src._layout) + ) + + +@dataclass(frozen=True) +class Float8Layout(Layout): + """Represents the layout configuration for Float8 affine quantized tensors. + + Attributes: + mm_config (Optional[Float8MMConfig]): Configuration for matrix multiplication operations involving Float8 tensors. If None, default settings are used. + """ + + mm_config: Optional[Float8MMConfig] = None + + +_fallback_warning_shown = False + + +@register_layout(Float8Layout) +class Float8AQTTensorImpl(AQTTensorImpl): + """ + TensorImpl for float8 layout affine quantized tensor + + Note: technically we should not create a new layout for float8 we should merge this into + plain layout + """ + + float8_data: torch.Tensor + scale: torch.Tensor + transposed: bool + + def __new__( + cls, + float8_data: torch.Tensor, + scale: torch.Tensor, + transposed: bool, + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = float8_data.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else float8_data.layout + ) + kwargs["dtype"] = float8_data.dtype + kwargs["requires_grad"] = False + shape = float8_data.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + float8_data: torch.Tensor, + scale: torch.Tensor, + transposed: bool, + _layout: Layout, + ): + self.float8_data = float8_data + self.scale = scale + self.transposed = transposed + self._layout = _layout + + def _apply_fn_to_data(self, fn): + """Applys a fn to all tensor components stored on this class""" + global _fallback_warning_shown + + try: + return self.__class__( + fn(self.float8_data), + fn(self.scale), + self.transposed, + self._layout, + ) + except RuntimeError as e: + if '"index_cuda" not implemented for ' in str(e): + if not _fallback_warning_shown: + import warnings + + warnings.warn( + f"When trying to index Float8AQTTensorImpl, got known error {e}, will use slower fallback but " + + "note: You can torch.compile the model to avoid this problem.", + UserWarning, + ) + _fallback_warning_shown = True + + return self.__class__( # do indexing in bfloat16 then convert back + fn(self.float8_data.to(torch.bfloat16)).to(self.float8_data.dtype), + fn(self.scale), + self.transposed, + self._layout, + ) + else: + raise e + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + return self.__class__( + self.float8_data.to(kwargs["device"]), + self.scale.to(kwargs["device"]), + self.transposed, + self._layout, + ) + + def __tensor_flatten__(self): + return ["float8_data", "scale"], [self.transposed, self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + float8_data, scale = tensor_data_dict["float8_data"], tensor_data_dict["scale"] + ( + transposed, + _layout, + ) = tensor_attributes + return cls(float8_data, scale, transposed, _layout) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + elif func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + elif func is aten.t.default: + """we don't need to repack the weight and just rely on external + shape being changed and record the status of transpose/no-transpose + """ + args[0].transposed = not args[0].transposed + return return_and_correct_aliasing(func, args, kwargs, args[0]) + elif func is aten.copy_.default: + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + elif func in [aten.select.int, aten.index.Tensor]: + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data(lambda x: func(x, *args[1:], **kwargs)), + ) + elif func is aten.slice.Tensor: + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + if dim == 0: + # TODO: scale replecation should be dependent on block size + if self.scale.ndim == 1: + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data( + lambda x: aten.slice.Tensor(x, dim, start, end, step) + ), + ) + elif self.scale.ndim == 0: + return return_and_correct_aliasing( + func, + args, + kwargs, + Float8AQTTensorImpl( + aten.slice.Tensor(self.float8_data, dim, start, end, step), + self.scale, + None, + self._layout, + ), + ) + else: + raise NotImplementedError( + f"Float8AQTTensorImpl dispatch: attempting to run {func}, with scale ndim={dim}, that is not supported" + ) + elif dim == 1: + return return_and_correct_aliasing( + func, + args, + kwargs, + Float8AQTTensorImpl( + aten.slice.Tensor( + self.float8_data, dim, start, end, step + ).contiguous(), + self.scale, + None, + self._layout, + ), + ) + else: + raise NotImplementedError( + f"Float8AQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported" + ) + else: + raise NotImplementedError( + f"Float8AQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + __torch_function__ = torch._C._disabled_torch_function_impl + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + return self.float8_data, self.scale, None + + def get_layout(self) -> Layout: + return self._layout + + @classmethod + def from_plain( + cls, + data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + """Main entrypoint for constructing Float8TensorImpl""" + assert _is_float8_type(data.dtype), ( + f"Float8 TensorImpl must be constructed from float8 dtype but got {data.dtype}" + ) + assert isinstance(_layout, Float8Layout), ( + f"Float8 TensorImpl must be constructed from Float8Layout but got {_layout}" + ) + return cls(data, scale, False, _layout) + + def __repr__(self): + float8_data, scale, _ = self.get_plain() + _layout = self.get_layout() + return ( + f"{self.__class__.__name__}(\n" + f"float8_data={float8_data},\n" + f"scale={scale},\n" + f"transposed={self.transposed}, " + f"_layout={_layout})" + ) + + +########################## +# Float8 Dispatch Kernels +########################## + + +def _linear_fp8_act_fp8_weight_check( + input_tensor: Union[torch.Tensor, "AffineQuantizedTensor"], + weight_tensor: Union[torch.Tensor, "AffineQuantizedTensor"], + bias: Optional[torch.Tensor], +) -> bool: + def check_aqt(aqt: Union[torch.Tensor, AffineQuantizedTensor]) -> bool: + return ( + isinstance(aqt, AffineQuantizedTensor) + and isinstance(aqt._layout, Float8Layout) + and aqt.tensor_impl.dtype in [torch.float8_e4m3fn, torch.float8_e5m2] + and (aqt.shape == aqt.block_size or _is_rowwise_scaled(aqt)) + ) + + return check_aqt(input_tensor) and check_aqt(weight_tensor) + + +def preprocess_scale(input_scale: torch.Tensor, input_shape: Tuple[int]): + """Ensures input tensor is correctly formated for _scaled_mm""" + input_scale = input_scale.unsqueeze(-1) + + if input_scale.dim() > 2: + input_scale = input_scale.reshape(-1, input_scale.shape[-1]) + + return input_scale + + +def _linear_fp8_act_fp8_weight_impl( + input_tensor: "AffineQuantizedTensor", + weight_tensor: "AffineQuantizedTensor", + bias: Optional[torch.Tensor], +): + """Implements matmul between FP8 input and FP8 weight with compute using _scaled_mm""" + scaled_mm_config = weight_tensor._layout.mm_config + assert scaled_mm_config is not None + out_shape = get_out_shape(input_tensor.shape, weight_tensor.shape) + + # Weight tensor preprocessing + w_tensor_impl = weight_tensor.tensor_impl + assert not w_tensor_impl.transposed, "Weight tensor must be contiguous" + w_data = w_tensor_impl.float8_data + w_scale = w_tensor_impl.scale + + # Input tensor preprocessing + inpt_data = input_tensor.tensor_impl.float8_data + input_scale = input_tensor.tensor_impl.scale + # Handle case where input tensor is more than 2D + inpt_data = inpt_data.reshape(-1, inpt_data.shape[-1]) + + # Handle rowwise case + if _is_rowwise_scaled(weight_tensor): + assert _is_rowwise_scaled(input_tensor), ( + "Input tensor must be rowwise block size" + ) + w_scale = w_scale.unsqueeze(-1).T + input_scale = preprocess_scale(input_scale, input_tensor.shape) + + # Preprocess data + inpt_data, w_data = preprocess_data(inpt_data, w_data.T, scaled_mm_config) + + # Perform the computation + return addmm_float8_unwrapped_inference( + inpt_data, + input_scale, + w_data, + w_scale, + output_dtype=input_tensor.dtype, + bias=bias, + use_fast_accum=scaled_mm_config.use_fast_accum, + ).reshape(out_shape) + + +def _linear_fp_act_fp8_weight_check( + input_tensor: Union[torch.Tensor, "AffineQuantizedTensor"], + weight_tensor: Union[torch.Tensor, "AffineQuantizedTensor"], + bias: Optional[torch.Tensor], +) -> bool: + return ( + # input is native float tensor + not is_traceable_wrapper_subclass(input_tensor) + and input_tensor.is_floating_point() + and + # weight is float8 quantized affine quantized tensor + isinstance(weight_tensor, AffineQuantizedTensor) + and isinstance(weight_tensor._layout, Float8Layout) + and weight_tensor.tensor_impl.dtype in [torch.float8_e4m3fn, torch.float8_e5m2] + and ( + weight_tensor.shape == weight_tensor.block_size + or _is_rowwise_scaled(weight_tensor) + ) + ) + + +def _linear_fp_act_fp8_weight_impl( + input_tensor: torch.Tensor, + weight_tensor: "AffineQuantizedTensor", + bias: Optional[torch.Tensor], +): + return torch.nn.functional.linear(input_tensor, weight_tensor.dequantize(), bias) diff --git a/lib/python3.12/site-packages/torchao/dtypes/floatx/floatx_tensor_core_layout.py b/lib/python3.12/site-packages/torchao/dtypes/floatx/floatx_tensor_core_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..6871033f1a4ff2c2dab054317f12f14454f1e670 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/floatx/floatx_tensor_core_layout.py @@ -0,0 +1,666 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from functools import reduce +from typing import Optional, Tuple + +import torch +from torch import Tensor +from torch.utils._python_dispatch import ( + is_traceable_wrapper_subclass, + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.utils import ( + AQTTensorImpl, + Layout, +) +from torchao.prototype.custom_fp_utils import ( + _f32_to_floatx_unpacked, + _floatx_unpacked_to_f32, + _n_ones, +) + +aten = torch.ops.aten +_ONES_TABLE = [_n_ones(i) for i in range(8)] + + +def _pack(x: Tensor, n_bits: int) -> Tensor: + return reduce( + torch.bitwise_or, + [ + x[..., i :: (8 // n_bits)] << (8 - (i + 1) * n_bits) + for i in range(8 // n_bits) + ], + ) + + +def _unpack(x: Tensor, n_bits: int) -> Tensor: + return torch.stack( + [ + (x >> (8 - (i + 1) * n_bits)) & ((1 << n_bits) - 1) + for i in range(8 // n_bits) + ], + dim=-1, + ).flatten(-2) + + +# https://github.com/usyd-fsalab/fp6_llm/blob/5df6737cca32f604e957e3f63f03ccc2e4d1df0d/fp6_llm/csrc/utils/weight_prepacking.h#L87-L116 +def _bit_interleave(x: Tensor, n_bits: int, undo: bool = False) -> Tensor: + # the original code unpacks/packs the values from/to uint32 while we unpack/pack the values from/to uint8 + # thus, we need to reverse byte order within a uint32 word. + x = x.reshape(-1, 4).flip(1) + + x = _unpack(x, n_bits) + x = x.view(-1, 4 * (8 // n_bits)) + + if not undo: + bit_order = { + 1: [ + 1, + 5, + 9, + 13, + 17, + 21, + 25, + 29, + 3, + 7, + 11, + 15, + 19, + 23, + 27, + 31, + 0, + 4, + 8, + 12, + 16, + 20, + 24, + 28, + 2, + 6, + 10, + 14, + 18, + 22, + 26, + 30, + ], + 2: [1, 5, 9, 13, 3, 7, 11, 15, 0, 4, 8, 12, 2, 6, 10, 14], + 4: [1, 5, 3, 7, 0, 4, 2, 6], + }[n_bits] + + else: + # this is inverse of the above, obtained by running + # [v.index(i) for i in range(len(v))] + bit_order = { + 1: [ + 16, + 0, + 24, + 8, + 17, + 1, + 25, + 9, + 18, + 2, + 26, + 10, + 19, + 3, + 27, + 11, + 20, + 4, + 28, + 12, + 21, + 5, + 29, + 13, + 22, + 6, + 30, + 14, + 23, + 7, + 31, + 15, + ], + 2: [8, 0, 12, 4, 9, 1, 13, 5, 10, 2, 14, 6, 11, 3, 15, 7], + 4: [4, 0, 6, 2, 5, 1, 7, 3], + }[n_bits] + + x = x[:, bit_order] + x = _pack(x, n_bits) + + # reverse byte order within a uint32 word again. + x = x.reshape(-1, 4).flip(1) + return x.flatten() + + +# this is a literal adaptation of FP6-LLM ahead-of-time bit-level pre-packing +# https://github.com/usyd-fsalab/fp6_llm/blob/5df6737cca32f604e957e3f63f03ccc2e4d1df0d/fp6_llm/csrc/utils/weight_prepacking.h +def _pack_tc_floatx(tensor: Tensor, nbits: int) -> Tensor: + assert tensor.ndim == 2, tensor.dtype == torch.uint8 + M, N = tensor.shape + assert (M % 64 == 0) and (N % 64 == 0) + + # Pass 1 from original code + tensor = tensor.view(M // 64, 4, 2, 8, N // 16, 2, 8) + tensor = tensor.permute(0, 4, 1, 5, 2, 3, 6) + tensor = tensor.reshape(-1, 32, 2) + tensor = tensor.permute(1, 0, 2) + tensor = tensor.flatten() + + used_bits = 0 + fragments = [] + + for y in [1, 2, 4]: + if nbits & y: + mask = (1 << y) - 1 + tensor_ybit = (tensor >> (nbits - used_bits - y)) & mask + tensor_ybit = _pack(tensor_ybit, y) + + tensor_ybit = ( + tensor_ybit.view(32, -1, 4).permute(1, 0, 2).flip(2) + ) # Pass 2 from original code + tensor_ybit = _bit_interleave( + tensor_ybit.flatten(), y + ) # Pass 3 from original code + fragments.append(tensor_ybit) + used_bits += y + + return torch.cat(fragments, dim=0).view(M, -1) + + +# more optimized version of _pack_tc_floatx() for FP6 by merging ops +def _pack_tc_fp6(tensor: Tensor) -> Tensor: + assert tensor.ndim == 2, tensor.dtype == torch.uint8 + M, N = tensor.shape + assert (M % 64 == 0) and (N % 64 == 0) + + tensor = tensor.view(M // 64, 2, 2, 2, 8, N // 16, 2, 8) + tensor = tensor.flip(3) + + tensor_2bit = (tensor >> 4) & 0b11 + tensor_2bit = tensor_2bit.permute(0, 5, 1, 4, 7, 3, 2, 6) + tensor_2bit = _pack(tensor_2bit.flatten(), 2) + + tensor_4bit = tensor & 0b1111 + tensor_4bit = tensor_4bit.permute(0, 5, 1, 2, 4, 7, 3, 6) + tensor_4bit = _pack(tensor_4bit.flatten(), 4) + + return torch.cat([tensor_2bit, tensor_4bit], dim=0).view(M, -1) + + +# currently only optimize for TC-FP6 packing +def pack_tc_floatx(tensor: Tensor, nbits: int) -> Tensor: + if nbits == 6: + return _pack_tc_fp6(tensor) + return _pack_tc_floatx(tensor, nbits) + + +def to_scaled_tc_floatx( + tensor: Tensor, ebits: int, mbits: int +) -> Tuple[Tensor, Tensor]: + # _n_ones() is not compatible with torch.compile() due to << operator + # https://github.com/pytorch/pytorch/issues/119152 + # exp_bias = _n_ones(ebits - 1) + # max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2 ** mbits)) + + # workaround: global lookup table + exp_bias = _ONES_TABLE[ebits - 1] + max_normal = 2 ** (_ONES_TABLE[ebits] - exp_bias) * ( + _ONES_TABLE[mbits + 1] / (2**mbits) + ) + + dtype = tensor.dtype + tensor = tensor.float() + scale = tensor.abs().amax(1).clamp(min=1e-12) / max_normal + tensor_floatx = _f32_to_floatx_unpacked(tensor / scale.view(-1, 1), ebits, mbits) + tensor_tc_floatx = pack_tc_floatx(tensor_floatx, 1 + ebits + mbits) + return tensor_tc_floatx, scale.to(dtype) + + +# inverse of _pack_tc_floatx() +def _unpack_tc_floatx(tensor: Tensor, nbits: int) -> Tensor: + assert tensor.ndim == 2 and tensor.dtype == torch.uint8 + M = tensor.shape[0] + size = tensor.numel() + tensor = tensor.flatten() + offset = 0 + used_bits = 0 + + tensor_floatx = None + + for y in [1, 2, 4]: + if nbits & y: + size_ybit = size // nbits * y + tensor_ybit = tensor[offset : offset + size_ybit] + offset += size_ybit + + tensor_ybit = _bit_interleave(tensor_ybit, y, undo=True) # undo Pass 3 + tensor_ybit = ( + tensor_ybit.view(-1, 32, 4).flip(2).permute(1, 0, 2) + ) # undo Pass 2 + + tensor_ybit = _unpack(tensor_ybit.flatten(), y) + tensor_ybit = tensor_ybit << (nbits - used_bits - y) + used_bits += y + + if tensor_floatx is None: + tensor_floatx = tensor_ybit + else: + tensor_floatx |= tensor_ybit + + # undo Pass 1 + tensor_floatx = tensor_floatx.view(32, -1, 2).permute(1, 0, 2) + tensor_floatx = tensor_floatx.reshape(M // 64, -1, 4, 2, 2, 8, 8) + tensor_floatx = tensor_floatx.permute(0, 2, 4, 5, 1, 3, 6) + tensor_floatx = tensor_floatx.reshape(M, -1) + return tensor_floatx + + +# more optimized version of _unpack_tc_floatx() for FP6 by merging ops +# inverse of _unpack_tc_fp6() +def _unpack_tc_fp6(tensor: Tensor) -> Tensor: + assert tensor.ndim == 2 and tensor.dtype == torch.uint8 + M = tensor.shape[0] + N = tensor.shape[1] // 3 * 4 + assert (M % 64 == 0) and (N % 64 == 0) + size_2bit = M * N // 4 + size_4bit = M * N // 2 + tensor = tensor.view(-1) + assert tensor.numel() == size_2bit + size_4bit + + tensor_2bit, tensor_4bit = tensor.split([size_2bit, size_4bit]) + + tensor_2bit = _unpack(tensor_2bit, 2) + tensor_2bit = tensor_2bit.view(M // 64, N // 16, 2, 8, 8, 2, 2, 2) + tensor_2bit = tensor_2bit.permute(0, 2, 6, 5, 3, 1, 7, 4) + + tensor_4bit = _unpack(tensor_4bit, 4) + tensor_4bit = tensor_4bit.view(M // 64, N // 16, 2, 2, 8, 8, 2, 2) + tensor_4bit = tensor_4bit.permute(0, 2, 3, 6, 4, 1, 7, 5) + + tensor_fp6 = (tensor_2bit << 4) | tensor_4bit + tensor_fp6 = tensor_fp6.flip(3).reshape(M, N) + return tensor_fp6 + + +def unpack_tc_floatx(tensor: Tensor, nbits: int) -> Tensor: + if nbits == 6: + return _unpack_tc_fp6(tensor) + return _unpack_tc_floatx(tensor, nbits) + + +def from_scaled_tc_floatx(tensor: Tensor, ebits: int, mbits: int, scale=None) -> Tensor: + floatx_unpacked = unpack_tc_floatx(tensor, 1 + ebits + mbits) + tensor = _floatx_unpacked_to_f32(floatx_unpacked, ebits, mbits) + if scale is not None: + tensor = tensor * scale.float().view(-1, 1) + return tensor + + +# https://github.com/microsoft/DeepSpeed/blob/3a3a6db3332e339cc9fd94efd4982f6d60635a3d/deepspeed/inference/v2/kernels/core_ops/cuda_linear/cuda_linear.py +_SPLIT_K_MAP = [ + { # tokens: [1, 64] + 3072: 18, + 4096: 13, + 5120: 10, + 6144: 9, + 8192: 6, + 10240: 5, + 14336: 7, + 28672: 7, + 57344: 7, + }, + { # tokens: [65:128] + 3072: 9, + 4096: 6, + 5120: 5, + 6144: 9, + 8192: 3, + 10240: 5, + 14336: 7, + 28672: 7, + 57344: 6, + }, + { # tokens: [129:192] + 3072: 6, + 4096: 4, + 5120: 7, + 6144: 3, + 8192: 2, + 10240: 5, + 14336: 5, + 28672: 5, + 57344: 4, + }, + { # tokens: [193:256] + 3072: 9, + 4096: 3, + 5120: 5, + 6144: 2, + 8192: 5, + 10240: 4, + 14336: 8, + 28672: 6, + 57344: 4, + }, + { # tokens: [257:320] + 3072: 7, + 4096: 5, + 5120: 2, + 6144: 5, + 8192: 4, + 10240: 1, + 14336: 3, + 28672: 3, + 57344: 4, + }, + { # tokens: [321:384] + 3072: 3, + 4096: 2, + 5120: 5, + 6144: 3, + 8192: 1, + 10240: 8, + 14336: 3, + 28672: 4, + 57344: 3, + }, + { # tokens: [385:448] + 3072: 5, + 4096: 7, + 5120: 3, + 6144: 5, + 8192: 7, + 10240: 3, + 14336: 1, + 28672: 1, + 57344: 3, + }, + { # tokens: [449:512] + 3072: 2, + 4096: 5, + 5120: 4, + 6144: 1, + 8192: 5, + 10240: 2, + 14336: 6, + 28672: 4, + 57344: 1, + }, + { # tokens: [513:576] + 3072: 2, + 4096: 3, + 5120: 1, + 6144: 1, + 8192: 3, + 10240: 3, + 14336: 3, + 28672: 1, + 57344: 1, + }, + { # tokens: [577:640] + 3072: 5, + 4096: 4, + 5120: 1, + 6144: 4, + 8192: 2, + 10240: 1, + 14336: 1, + 28672: 1, + 57344: 1, + }, + { # tokens: [641:704] + 3072: 3, + 4096: 1, + 5120: 2, + 6144: 2, + 8192: 1, + 10240: 2, + 14336: 1, + 28672: 1, + 57344: 1, + }, + { # tokens: [705:768] + 3072: 3, + 4096: 1, + 5120: 3, + 6144: 2, + 8192: 1, + 10240: 1, + 14336: 1, + 28672: 1, + 57344: 1, + }, +] + + +# quantization api integrations +@dataclass(frozen=True) +class FloatxTensorCoreLayout(Layout): + """FloatxTensorCoreLayout is a data class that defines the layout for a tensor with a specific number of exponent bits (ebits) and mantissa bits (mbits). + This layout is used in the context of quantization and packing of tensors optimized for TensorCore operations. + """ + + ebits: int + mbits: int + + +@register_layout(FloatxTensorCoreLayout) +class FloatxTensorCoreAQTTensorImpl(AQTTensorImpl): + """FloatxTensorCoreAQTTensorImpl represents a Tensor with dtype floatx(ebits=a, mbits=b), + it has a internal tensor field of "packed_floatx_data", which is packed from the + uint8 unpacked data (the output of `quantize_affine_floatx` operator) + + The packing is optimized for TensorCore, from the fp6-llm paper: https://arxiv.org/abs/2401.14112 + github repo: https://github.com/usyd-fsalab/fp6_llm, now renamed to quant-llm + + At a high level packing is done by grouping bits into 1 bit fragments (shards), 2 bit fragments and + 4 bit fragments each fragments are packed separately and concatenated together. + For example for 6 bit dtype, we can extract the first 4 bits for all elements and pack them together + in a fragment, and extract the last 2 bits for all elements and pack them into fragment, in the end + we concatenate the fragments together. + + If original Tensor shape is (M, N), and the data is in nbit, the shape of the packed data will be + (M, N // 8 * nbit) + + FloatxTensorCoreAQTTensorImpl.from_plain takes an unpacked uint8 floatx Tensor of shape (M, N), with format of + (zero padding bits + sign bit + exponent bits + mantissa bits), e.g. 00SEEEMM for fp6_e3_m2 + it will then pack the weight and instantiate the FloatxTensorCoreAQTTensorImpl tensor + FloatxTensorCoreAQTTensorImpl.__init__() takes a packed floatx Tensor of shape (M, N // 8 * nbit) + """ + + def __new__( + cls, + packed_floatx_data: torch.Tensor, + scale: torch.Tensor, + _layout: Layout, + ): + assert packed_floatx_data.ndim == 2 + assert packed_floatx_data.dtype == torch.uint8 + shape = ( + packed_floatx_data.shape[0], + packed_floatx_data.shape[1] // (1 + _layout.ebits + _layout.mbits) * 8, + ) + kwargs = {} + kwargs["device"] = packed_floatx_data.device + kwargs["layout"] = ( + kwargs.get("layout") + if kwargs.get("layout", False) + else packed_floatx_data.layout + ) + kwargs["dtype"] = packed_floatx_data.dtype + kwargs["requires_grad"] = False + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + packed_floatx_data: torch.Tensor, + scale: torch.Tensor, + _layout: Layout, + ): + self.packed_floatx_data = packed_floatx_data + self.scale = scale + self._layout = _layout + + def __tensor_flatten__(self): + return ["packed_floatx_data", "scale"], [self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + packed_floatx_data, scale = ( + tensor_data_dict["packed_floatx_data"], + tensor_data_dict["scale"], + ) + (_layout,) = tensor_attributes + return cls(packed_floatx_data, scale, _layout) + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor]: + unpacked_floatx_data = unpack_tc_floatx( + self.packed_floatx_data, 1 + self._layout.ebits + self._layout.mbits + ) + return unpacked_floatx_data, self.scale + + @classmethod + def from_plain( + cls, + unpacked_floatx_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + """ + Format for `unpacked_floatx_data` will be: + zero padding bits | sign bit | exponent bits | mantissa bits + + For example for fp6_e3_m2, the format will be: `00SEEEMM`, where S is sign bit, E is exponent + bit, M is mantissa bit + """ + assert isinstance(_layout, FloatxTensorCoreLayout) + packed_floatx_data = pack_tc_floatx( + unpacked_floatx_data, 1 + _layout.ebits + _layout.mbits + ) + return cls(packed_floatx_data, scale, _layout) + + def __repr__(self): + unpacked_floatx_data, scale = self.get_plain() + _layout = self.get_layout() + return f"{self.__class__.__name__}(unpacked_floatx_data={unpacked_floatx_data}, scale={scale}, _layout={_layout})" + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.packed_floatx_data), + fn(self.scale), + self._layout, + ) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs.pop("device") + return self.__class__( + self.packed_floatx_data.to(device), + self.scale.to(device), + self._layout, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + elif func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + elif func is aten._to_copy.default: + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data( + lambda x: x.to(device=kwargs.pop("device", None)) + ), + ) + + raise NotImplementedError( + f"FloatxTensorCoreAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + __torch_function__ = torch._C._disabled_torch_function_impl + + def get_layout(self) -> Layout: + return self._layout + + +def _linear_f16_bf16_act_floatx_weight_check(input_tensor, weight_tensor, bias): + from torchao.dtypes.floatx import FloatxTensorCoreLayout + + return ( + # input is native float32 tensor + not is_traceable_wrapper_subclass(input_tensor) + and input_tensor.is_floating_point() + and input_tensor.dtype in (torch.float16, torch.bfloat16) + and + # weight is floatx Tensor + isinstance(weight_tensor, AffineQuantizedTensor) + and isinstance(weight_tensor._layout, FloatxTensorCoreLayout) + and ( + # weight is using fp6 quantization + (weight_tensor._layout.ebits == 3 and weight_tensor._layout.mbits == 2) + or (weight_tensor._layout.ebits == 2 and weight_tensor._layout.mbits == 3) + or + # weight is using fp5 quantization + (weight_tensor._layout.ebits == 2 and weight_tensor._layout.mbits == 2) + or (weight_tensor._layout.ebits == 3 and weight_tensor._layout.mbits == 1) + ) + ) + + +def _linear_f16_bf16_act_floatx_weight_impl(input_tensor, weight_tensor, bias): + from torchao.ops import quant_llm_linear + + act = input_tensor + weight = weight_tensor + + out_dim, in_dim = weight.shape + act_reshaped = act.view(-1, in_dim) + + # https://github.com/microsoft/DeepSpeed/blob/3a3a6db3332e339cc9fd94efd4982f6d60635a3d/deepspeed/inference/v2/kernels/core_ops/cuda_linear/cuda_linear.py + bsize = act_reshaped.shape[0] + splitK = _SPLIT_K_MAP[(bsize - 1) // 64].get(out_dim, 1) if bsize <= 768 else 1 + + out = quant_llm_linear( + weight._layout.ebits, + weight._layout.mbits, + act_reshaped, + weight.tensor_impl.packed_floatx_data, + weight.tensor_impl.scale, + splitK=splitK, + ) + + if bias is not None: + out += bias + + return out.view(*act.shape[:-1], out_dim).to(act.dtype) diff --git a/lib/python3.12/site-packages/torchao/dtypes/nf4tensor.py b/lib/python3.12/site-packages/torchao/dtypes/nf4tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..698a9391bdfe52fdf80eb3e933f1093f42003e72 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/nf4tensor.py @@ -0,0 +1,1086 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import functools +import math +import sys +from dataclasses import dataclass, replace +from enum import Enum, auto +from typing import Any, Dict, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch._prims_common import make_contiguous_strides_for +from torch.distributed.device_mesh import DeviceMesh + +from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 + +aten = torch.ops.aten + +c10d_functional = torch.ops.c10d_functional + + +NF4_OPS_TABLE: Dict[Any, Any] = {} + + +_INNER_TENSOR_NAMES_FOR_SHARDING = [ + "quantized_scalers", + "quantization_factor", + "quantized_data", +] + +# Note: Quantize in Chunks +# During quantization to NF4, one of the steps to convert from the original float number +# to the index of the nearest value in the NF4 format. This can cause a large memory spike +# Due to intermediates of the quantization process. Instead we process the original +# tensor in chunks. This is a tradeoff between memory and speed. This number seems to +# strike a good balance between memory and speed +CHUNK_SIZE = 1024**2 + + +def same_metadata(a: "NF4Tensor", b: "NF4Tensor"): + both_nf4 = isinstance(a, NF4Tensor) and isinstance(b, NF4Tensor) + return ( + both_nf4 + and a.block_size == b.block_size + and a.scaler_block_size == b.scaler_block_size + and a.n_blocks == b.n_blocks + ) + + +def implements(aten_ops): + """Use this decorator to implement a function for an aten op in __torch_dispatch__""" + + def decorator(func): + for op in aten_ops: + NF4_OPS_TABLE[op] = func + return func + + return decorator + + +def construct_nf4_args(nf4tensor: "NF4Tensor", kwargs: Optional[Dict[str, Any]] = None): + if kwargs is None: + kwargs = {} + tensor_meta = SubclassTensorArgs( + kwargs.get("size", nf4tensor.size()), + kwargs.get("stride", nf4tensor.stride()), + kwargs.get("storage_offset", nf4tensor.storage_offset()), + kwargs.get("dtype", nf4tensor.dtype), + kwargs.get("device", nf4tensor.device), + kwargs.get("requires_grad", nf4tensor.requires_grad), + ) + return ( + tensor_meta, + kwargs.get("block_size", nf4tensor.block_size), + kwargs.get("n_blocks", nf4tensor.n_blocks), + kwargs.get("scaler_block_size", nf4tensor.scaler_block_size), + kwargs.get("quantized_scalers", nf4tensor.quantized_scalers), + kwargs.get("quantization_factor", nf4tensor.quantization_factor), + kwargs.get("scaler_mean", nf4tensor.scaler_mean), + kwargs.get("quantized_data", nf4tensor.quantized_data), + kwargs.get("nf4", nf4tensor.nf4), + ) + + +# __torch_dispatch__ utils: apply aten op to inner tensors +def apply_to_inner_tensors(nf4tensor: "NF4Tensor", aten_op, args, kwargs): + attr_to_tensor = {} + for attr in _INNER_TENSOR_NAMES_FOR_SHARDING: + attr_to_tensor[attr] = aten_op(getattr(nf4tensor, attr), *args, **kwargs) + return attr_to_tensor + + +# __torch_function__ utils: call tensor ops from inner tensors +def call_from_inner_tensors(nf4tensor: "NF4Tensor", method_name: str, args, kwargs): + attr_to_tensor = {} + for attr in _INNER_TENSOR_NAMES_FOR_SHARDING: + inner_tensor = getattr(nf4tensor, attr) + func = getattr(inner_tensor, method_name) + attr_to_tensor[attr] = func(*args, **kwargs) + return attr_to_tensor + + +class CompareOp(Enum): + EQ = auto() + LT = auto() + + +def expect_num_of_args(op: CompareOp, num: int, msg: str): + def decorator(func): + @functools.wraps(func) + def wrapper(aten_op, args, kwargs=None): + if op == CompareOp.LT and not (len(args) < num): + raise NotImplementedError(msg) + return func(aten_op, args, kwargs) + + return wrapper + + return decorator + + +def expect_arg_value_at_k(k: int, op: CompareOp, value: Any, msg: str): + def decorator(func): + @functools.wraps(func) + def wrapper(aten_op, args, kwargs=None): + if op == CompareOp.EQ and not (args[k] == value): + raise NotImplementedError(msg + str(args[k])) + return func(aten_op, args, kwargs) + + return wrapper + + return decorator + + +def expect_args_len_at_k(k: int, op: CompareOp, value: Any, msg: str): + def decorator(func): + @functools.wraps(func) + def wrapper(aten_op, args, kwargs=None): + if op == CompareOp.LT and not (len(args[k]) < value): + raise NotImplementedError(msg + str(len(args[k]))) + elif op == CompareOp.EQ and not (len(args[k]) == value): + raise NotImplementedError(msg + str(len(args[k]))) + return func(aten_op, args, kwargs) + + return wrapper + + return decorator + + +@implements([torch.ops.aten.detach]) +def noop_detach(func, *args, **kwargs): + return args[0][0] + + +@implements([torch.ops.aten.clone.default]) +def clone(func, *args, **kwargs): + return to_nf4(args[0][0].get_original_weight()) + + +@implements( + [ + aten.detach.default, + ] +) +def nf4_detach(aten_op, args, kwargs=None): + nf4tensor = args[0] + updated_attrs = apply_to_inner_tensors(nf4tensor, aten_op, args[1:], kwargs) + return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs)) + + +@implements( + [ + aten.empty_like.default, + ] +) +def nf4_empty_like(aten_op, args, kwargs=None): + nf4tensor = args[0] + updated_attrs = apply_to_inner_tensors(nf4tensor, aten_op, args[1:], kwargs) + if kwargs is not None and len(kwargs): + for key, value in kwargs.items(): + updated_attrs[key] = value + return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs)) + + +@implements( + [ + aten.split.Tensor, + ] +) +def nf4_split(aten_op, args, kwargs=None): + if len(args) == 3 and args[2] != 0: + raise NotImplementedError(f"aten.split(NF4Tensor, dim={args[2]})") + nf4tensor = args[0] + num_chunks = nf4tensor.size(0) // args[1] + + attr_to_chunks = {} + for attr in _INNER_TENSOR_NAMES_FOR_SHARDING: + inner_tensor = getattr(nf4tensor, attr) + assert inner_tensor.numel() % num_chunks == 0, ( + f"{attr}.numel() not divisible by {num_chunks}" + ) + chunks = aten_op(inner_tensor, inner_tensor.numel() // num_chunks, **kwargs) + attr_to_chunks[attr] = chunks + + orig_dim = nf4tensor.dim() + if orig_dim == 1: + chunked_size = (nf4tensor.size(0) // num_chunks,) + elif orig_dim == 2: + chunked_size = (nf4tensor.size(0) // num_chunks, nf4tensor.size(1)) + else: + chunked_size = () + raise NotImplementedError( + f"aten.split(NF4Tensor) wherer NF4Tensor.dim() = {orig_dim}" + ) + + nf4_chunks = [] + for idx in range(num_chunks): + updated_attrs = {"size": chunked_size} + for attr, chunks in attr_to_chunks.items(): + updated_attrs[attr] = chunks[idx] + nf4_chunks.append(NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs))) + return nf4_chunks + + +@implements( + [ + aten.new_zeros.default, + ] +) +@expect_args_len_at_k(1, CompareOp.LT, 3, "aten.view(NF4Tensor) with len(size)=") +def nf4_new_zeros(aten_op, args, kwargs=None): + nf4tensor = args[0] + new_size = tuple(args[1]) + + if nf4tensor.numel() % math.prod(new_size) != 0: + raise NotImplementedError(f"aten.new_zeros(NF4Tensor) with new size {new_size}") + ratio = nf4tensor.numel() // math.prod(new_size) + + updated_attrs = {} + for attr in _INNER_TENSOR_NAMES_FOR_SHARDING: + inner_tensor = getattr(nf4tensor, attr) + assert inner_tensor.size(0) % ratio == 0, ( + f"{attr}.numel() must be divisible by {ratio}" + ) + inner_tensor = aten_op(inner_tensor, [inner_tensor.size(0) // ratio], **kwargs) + updated_attrs[attr] = inner_tensor + updated_attrs["size"] = new_size + + return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs)) + + +@implements( + [ + aten.slice.Tensor, + ] +) +@expect_num_of_args(CompareOp.LT, 5, "aten.slice(NF4Tensor) with customized step") +@expect_arg_value_at_k(1, CompareOp.EQ, 0, "aten.slice(NF4Tensor) with dim=") +@expect_arg_value_at_k(2, CompareOp.EQ, 0, "aten.slice(NF4Tensor) with start=") +def nf4_slice(aten_op, args, kwargs=None): + nf4tensor = args[0] + # for tensor 512 x 512, tensor[:, :512] dispatch to + # aten.slice(dim = 0, end=sys.maxsize) + if args[3] not in [nf4tensor.size(0), sys.maxsize]: + raise NotImplementedError(f"aten.slice(NF4Tensor) with end={args[3]}") + return NF4Tensor(*construct_nf4_args(nf4tensor)) + + +@implements( + [ + aten.view.default, + ] +) +@expect_args_len_at_k(1, CompareOp.EQ, 1, "aten.view(NF4Tensor) with len(size)=") +def nf4_view(aten_op, args, kwargs=None): + nf4tensor = args[0] + size = args[1] + if size[0] != -1: + raise NotImplementedError(f"aten.view(NF4Tensor) with size={size}") + updated_attrs = apply_to_inner_tensors(nf4tensor, aten_op, args[1:], kwargs) + updated_attrs.update( + { + "size": [nf4tensor.numel()], + "stride": (1,), + } + ) + return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs)) + + +@implements( + [ + aten.as_strided.default, + ] +) +@expect_args_len_at_k( + 1, CompareOp.LT, 3, "aten.as_strided(NF4Tensor) only support dim <= 2 but got dim=" +) +def nf4_as_strided(aten_op, args, kwargs=None): + nf4tensor = args[0] + size = args[1] + stride = tuple(args[2]) + storage_offset = args[3] + if math.prod(size) != nf4tensor.numel(): + raise NotImplementedError( + f"aten.as_strided(NF4Tensor) different numel={nf4tensor.numel()} and size={size}" + ) + if stride != make_contiguous_strides_for(size): + raise NotImplementedError( + f"aten.as_strided(NF4Tensor) only support continuous stride={make_contiguous_strides_for(size)} but got stride={stride}" + ) + if nf4tensor.storage_offset() != storage_offset: + raise NotImplementedError( + f"aten.as_strided(NF4Tensor) only support original storage offset {nf4tensor.storage_offset()} but got {storage_offset}" + ) + kwargs = { + "size": torch.Size(size), + "stride": stride, + "storage_offset": storage_offset, + } + return NF4Tensor(*construct_nf4_args(nf4tensor, kwargs)) + + +@implements([torch.ops.aten._to_copy.default]) +def _to_copy(func, *args, **kwargs): + if not args[0][0].is_contiguous(): + assert args[0][0].t().is_contiguous() + return func(args[0][0].t()).t() + out = args[0][0].get_original_weight().to(args[1]["dtype"]) + if "device" in args[1]: + out = out.to(args[1]["device"]) + return out + + +@implements([torch.ops.aten.to.dtype]) +def to_dtype(func, *args, **kwargs): + if not args[0][0].is_contiguous(): + assert args[0][0].t().is_contiguous() + return torch.ops.aten.to.dtype(args[0][0].t(), args[0][1]).t() + return args[0][0].get_original_weight().to(args[0][1]) + + +@implements([torch.ops.aten.t.default]) +def t_default(func, *args, **kwargs): + a = args[0][0] + tensor_meta = SubclassTensorArgs( + a.size(), + (a.stride(1), a.stride(0)), + a.storage_offset(), + a.dtype, + a.device, + a.requires_grad, + ) + b = NF4Tensor( + tensor_meta, + a.block_size, + a.n_blocks, + a.scaler_block_size, + a.quantized_scalers, + a.quantization_factor, + a.scaler_mean, + a.quantized_data, + a.nf4, + ) + return b + + +@implements([torch.ops.aten.mm.default]) +def mm_default(func, *args, **kwargs): + return linear_nf4(args[0][0], args[0][1]) + + +@implements( + [ + aten.copy_.default, + ] +) +def copy_(func, *args, **kwargs): + original: NF4Tensor = args[0][0] + copy_in: torch.Tensor = args[0][1] + + # Base Case + + if same_metadata(original, copy_in): + original_tensors = original.__tensor_flatten__()[0] + for tensor_name in original_tensors: + getattr(original, tensor_name).copy_(getattr(copy_in, tensor_name)) + return + + # Convert Non NF4Tensor into NF4 for copy in + if not isinstance(copy_in, NF4Tensor): + copy_in_nf4 = NF4Tensor.from_tensor( + copy_in.to(original.device), original.block_size, original.scaler_block_size + ) + return original.copy_(copy_in_nf4) + + # Other Tensor is not a NF4Tensor + full_precision = copy_in.get_original_weight() + same_meta_nf4 = NF4Tensor.from_tensor( + full_precision, original.block_size, original.scaler_block_size + ) + return original.copy_(same_meta_nf4) + + +@implements( + [ + aten.is_pinned.default, + ] +) +def nf4_is_pinned(aten_op, args, kwargs=None): + nf4tensor = args[0] + for attr in _INNER_TENSOR_NAMES_FOR_SHARDING: + inner_tensor = getattr(nf4tensor, attr) + if not aten_op(inner_tensor, *(args[1:]), **kwargs): + return False + return True + + +@implements( + [ + aten._pin_memory.default, + ] +) +def nf4_pin_memory(aten_op, args, kwargs=None): + nf4tensor = args[0] + updated_attrs = apply_to_inner_tensors(nf4tensor, aten_op, args[1:], kwargs) + return NF4Tensor(*construct_nf4_args(nf4tensor, updated_attrs)) + + +@implements( + [ + aten.cat.default, + ] +) +def nf4_cat(aten_op: torch._ops.OpOverload, args, kwargs=None): + tensors_to_cat = args[0] + assert all(isinstance(t, torch.Tensor) for t in tensors_to_cat) + remaining_args = args[1:] + + ts = [] + for t in tensors_to_cat: + assert isinstance(t, torch.Tensor) + + if isinstance(t, NF4Tensor): + ts.append(t.get_original_weight()) + else: + ts.append(t) + + dtype = ts[0].dtype + assert all(t.dtype == dtype for t in ts) + + if kwargs is None: + kwargs = {} + + tensors = aten_op(ts, *remaining_args, **kwargs) + return tensors + + +@dataclass(frozen=True) +class SubclassTensorArgs: + original_shape: torch.Size + + original_strides: Tuple + storage_offset: int + dtype: torch.dtype + device: torch.device + requires_grad: bool + + +def get_block_absmax(input_tensor: torch.Tensor, block_size: int) -> torch.Tensor: + """Iterate through a flattened tensor getting the absmax scalers for each block + + Args: + input_tensor: Input tensor to get scalers for + block_size: Block size for the scanning window + Returns: + torch.Tensor: Tensor of scalers for each block + """ + assert input_tensor.dim() == 1, "Input tensor must be flattened" + assert (input_tensor.numel() % block_size) == 0, ( + f"Input tensor must be divisible by block size, got {input_tensor.numel()} and {block_size}" + ) + + n_blocks = input_tensor.numel() // block_size + blocks = input_tensor.view(n_blocks, block_size) + block_scalers = blocks.abs().max(dim=1).values + return block_scalers + + +class NF4Tensor(torch.Tensor): + """NF4Tensor class for converting a weight to the QLoRA NF4 format""" + + @torch._dynamo.disable + def __new__( + cls, + # Args related for base tensor construction + tensor_meta: SubclassTensorArgs, + # Args stored on the instance + block_size: int, + n_blocks: int, + scaler_block_size: int, + quantized_scalers: torch.Tensor, + quantization_factor: torch.Tensor, + scaler_mean: torch.Tensor, + quantized_data: torch.Tensor, + nf4: torch.Tensor, + ): + """Create a new NF4Tensor object + Args: + tensor_meta: Metadata for the tensor + block_size: Size of the quantization block + n_blocks: Number of blocks to cover the full tensor + scaler_block_size: Block size for the scalar quantization + quantized_scalers: Quantized scalers data' represented a uint8 tensor + quantization_factor: Quantization factor, single scalar represented as torch.Tensor + scaler_mean: Mean of the scalers + quantized_data: Quantized data represented as uint8 tensor + nf4: NF4 tensor LUT for the quantization and dequantization + + """ + + nf4tensor = torch.Tensor._make_wrapper_subclass( + cls, + tensor_meta.original_shape, + tensor_meta.original_strides, + tensor_meta.storage_offset, + # Picked some floating dtype, but we need dtype extensibility + dtype=tensor_meta.dtype, + device=tensor_meta.device, + requires_grad=tensor_meta.requires_grad, + ) + return nf4tensor + + @torch._dynamo.disable + def __init__( + self, + tensor_meta: SubclassTensorArgs, + block_size: int, + n_blocks: int, + scaler_block_size: int, + quantized_scalers: torch.Tensor, + quantization_factor: torch.Tensor, + scaler_mean: torch.Tensor, + quantized_data: torch.Tensor, + nf4: torch.Tensor, + ): + """Initialize the NF4Tensor class""" + self.block_size = block_size + self.n_blocks = n_blocks + self.scaler_block_size = scaler_block_size + self.quantized_scalers = quantized_scalers + self.quantization_factor = quantization_factor + self.scaler_mean = scaler_mean + self.quantized_data = quantized_data + self.nf4 = nf4 + + @classmethod + @torch.no_grad() + def from_tensor( + cls, + input_tensor: torch.Tensor, + block_size: int, + scaler_block_size: int, + ): + assert input_tensor.dim() <= 2, ( + f"expect input tensor dim <= 2 but got dim = {input_tensor.dim()}" + ) + assert input_tensor.numel() % block_size == 0, ( + f"Input tensor must be divisible by block size, got {input_tensor.numel()} and {block_size}" + ) + assert input_tensor.is_contiguous, "Input tensor must be contiguous!" + # I think I want do this + # assert not input_tensor.requires_grad, "Input tensor must not require grad" + device = input_tensor.device + # Cache the tensor on the class def + nf4 = torch.tensor( + [ + -1.0000, + -0.6962, + -0.5251, + -0.3949, + -0.2844, + -0.1848, + -0.0911, + 0.0000, + 0.0796, + 0.1609, + 0.2461, + 0.3379, + 0.4407, + 0.5626, + 0.7230, + 1.0000, + ], + device=device, + dtype=input_tensor.dtype, + ) + n_blocks = input_tensor.numel() // block_size + # Double quantization + ( + quantized_scalers, + quantization_factor, + scaler_mean, + ) = cls.double_quantize_scalers( + input_tensor.flatten(), block_size, scaler_block_size + ) + quantized_data = cls.convert_to_norm_float_weight( + input_tensor, n_blocks, block_size, nf4 + ) + tensor_meta = SubclassTensorArgs( + input_tensor.size(), + input_tensor.stride(), + input_tensor.storage_offset(), + input_tensor.dtype, + input_tensor.device, + input_tensor.requires_grad, + ) + return cls( + tensor_meta, + block_size, + n_blocks, + scaler_block_size, + quantized_scalers, + quantization_factor, + scaler_mean, + quantized_data, + nf4=nf4, + ) + + @staticmethod + def double_quantize_scalers( + input_tensor: torch.Tensor, + block_size: int, + scaler_block_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Used to achieve the double quantization of the scalers + We take the input tensor first calculate the absmax quantization factors for each block. + We then find the mean of our positive absmax scalers. We subtract this mean from the scalers + And then we calculate the absmax quantization factors for each block again. We then quantize the scalers to int8. + + Args: + input_tensor: Input tensor to convert to QLoRA format, typically a weight tensor + + Returns: + torch.Tensor: Tensor of per_block quantization factors stored in int8 format + size: (n_blocks) + torch.Tensor: Tensor of per_scaler_block quantization factors stored in int16 format + size: (n_scaler_blocks) + """ + assert input_tensor.dim() == 1, "Input tensor must be flattened" + assert (input_tensor.numel() % scaler_block_size) == 0, ( + f"Input tensor must be divisible by block size, got {input_tensor.numel()} and {scaler_block_size}" + ) + + # First round of quantization + # Produces: A tensor of size (n_blocks) of input_tensor.dtype + scalers_1 = get_block_absmax(input_tensor, block_size) + scalers_1_mean = scalers_1.mean() + scalers_1 = scalers_1 - scalers_1_mean + # Second round of quantization + assert scalers_1.numel() % scaler_block_size == 0, ( + f"Number of scalers must be divisible by scaler block size, got {scalers_1.numel()} scaler_block_size {scaler_block_size} " + ) + n_scaler_blocks = scalers_1.numel() // scaler_block_size + scaler_blocks = scalers_1.view(n_scaler_blocks, scaler_block_size) + + scaler_absmax = get_block_absmax(scalers_1, scaler_block_size) + scaler_absmax = scaler_absmax.unsqueeze(-1).expand( + n_scaler_blocks, scaler_block_size + ) + + quantization_factor = 256 / (2 * scaler_absmax) + # Length equal to weight numel // block_size + quantized_scaler_blocks = scaler_blocks * quantization_factor + quantized_scaler_blocks = quantized_scaler_blocks.round() + quantized_scaler_blocks = quantized_scaler_blocks.clamp(-128, 127) + + # This is needed to make sure that quantization_factor remains a repeated view of n_scaler_blocks + # For some reason the 127/scaler_absmax realizes n_scaler entries when only n_scaler_blocks are needed + # The following will grab the first entry for the n_scaler_blocks which is the same across the scaler_block_size + + quantization_factor = quantization_factor[:, 0] + + return ( + quantized_scaler_blocks.flatten().to(torch.int8), + quantization_factor.view(n_scaler_blocks).contiguous(), + scalers_1_mean, + ) + + def dequantize_scalers( + self, + input_tensor: torch.Tensor, + quantization_factor: torch.Tensor, + scaler_block_size: int, + ) -> torch.Tensor: + """Used to unpack the double quantized scalers + + Args: + input_tensor: Input tensor to convert to QLoRA format this is the quantized scalers in int8 format + quantization_factor: Tensor of per_scaler_block quantization factors stored in inpt_weight.dtype + scaler_block_size: Scaler block size to use for double quantization. + + """ + assert input_tensor.dim() == 1, "Input tensor must be flattened" + assert (input_tensor.numel() % scaler_block_size) == 0, ( + f"Input tensor must be divisible by block size, got {input_tensor.numel()} and {scaler_block_size}" + ) + n_scaler_blocks = input_tensor.numel() // scaler_block_size + input_tensor = input_tensor.view(n_scaler_blocks, scaler_block_size) + dequantized = (input_tensor / quantization_factor.unsqueeze(-1)).flatten().to( + self.dtype + ) + self.scaler_mean + return dequantized + + @staticmethod + def convert_to_norm_float_weight( + input_tensor: torch.Tensor, n_blocks: int, block_size: int, nf4: torch.Tensor + ) -> torch.Tensor: + """Convert a tensor to the normalized float weight format""" + flattened_tensor = input_tensor.flatten() + # Since we are using uint8 we will encode 2 entries per byte + numel = input_tensor.numel() + assert numel % 2 == 0, ( + "Number of elements must be even just to not have to think about the end" + ) + # Reshape the flattened tensor into blocks of size self.block_size + blocks = flattened_tensor.view(n_blocks, block_size) + + # Scale the blocks + scalers = get_block_absmax(input_tensor.flatten(), block_size) + scales = scalers.unsqueeze(-1).expand(n_blocks, block_size) + scaled_blocks = blocks / scales + + # Returns a flattened tensor with each element quantized to nf4 index + # See Note: Quantize in Chunks + quantized_blocks = torch.empty( + numel, dtype=torch.uint8, device=input_tensor.device + ) + flattened = scaled_blocks.flatten() + for chunk_num in range(math.ceil(numel / CHUNK_SIZE)): + start = chunk_num * CHUNK_SIZE + end = min(start + CHUNK_SIZE, numel) + quantized_blocks[start:end] = NF4Tensor.quantize_tensor_nearest( + flattened[start:end], nf4 + ).to(torch.uint8) + + # Combine the quantized elements into uint8 values + # This lays out two consecutive elements in the same byte + # [a, b, c, d] -> [ab, cd] + # The size of combined blocks will be half the size of the original tensor + combined_blocks = quantized_blocks[::2] << 4 | quantized_blocks[1::2] + + return combined_blocks.to(torch.uint8) + + def get_original_weight(self) -> torch.Tensor: + """Get the original weight from the normalized float weight format""" + # Since we are using uint8 we will decode 2 entries per byte + # Shift elements down 4 and select out the bottom 4 bits + first_elements = (self.quantized_data >> 4).to(torch.long) + second_elements = (self.quantized_data & 0b1111).to(torch.long) + + # Dequantize every element + dequantized_first = self.dequantize(first_elements, self.nf4) + dequantized_second = self.dequantize(second_elements, self.nf4) + + # Build up matrix of scalers repeated for each element in the block + # Since first and second elements make up a full block + # we expand out to half the size of the full block + scalers = self.dequantize_scalers( + self.quantized_scalers, self.quantization_factor, self.scaler_block_size + ) + repeated = scalers.unsqueeze(-1).expand(scalers.size(0), self.block_size // 2) + + scaled_first = dequantized_first * repeated.flatten() + scaled_second = dequantized_second * repeated.flatten() + + # Flip them to be vertical and them stack them together horizontally + # Upon flattening this will interleave the elements + scaled_first = scaled_first.unsqueeze(-1).transpose(0, 1) + scaled_second = scaled_second.unsqueeze(-1).transpose(0, 1) + return torch.stack([scaled_first, scaled_second], dim=-1).reshape(self.shape) + + @staticmethod + def quantize_tensor_nearest(value: torch.Tensor, nf4: torch.Tensor) -> torch.Tensor: + """Quantize a float16 tensor to nf4 format to nearest and not rounded up""" + value = value.unsqueeze(-1) # (numel, 1) + # Compare the value tensor with the nf4 tensor element-wise + diff = (value - nf4).abs() + closest_nf4 = diff.min(dim=-1).indices + return closest_nf4 + + @staticmethod + def dequantize(value: torch.Tensor, nf4: torch.Tensor) -> torch.Tensor: + """Dequantize a nf4 value to bfloat16 format""" + # return nf4.index_select(0, value) + return nf4[value] + + def __repr__(self) -> str: + return f"Quantized Data: {self.quantized_data}\nScalers: {self.quantized_scalers}\n" + + def __str__(self) -> str: + return f"NF4Tensor({self.shape}, {self.block_size})" + + def __tensor_flatten__(self): + tensor_meta = SubclassTensorArgs( + self.shape, + self.stride(), + self.storage_offset(), + self.dtype, + self.device, + self.requires_grad, + ) + ctx = { + "block_size": self.block_size, + "n_blocks": self.n_blocks, + "scaler_block_size": self.scaler_block_size, + "tensor_meta": tensor_meta, + } + return [ + "quantized_data", + "scaler_mean", + "quantization_factor", + "quantized_scalers", + "nf4", + ], ctx + + @staticmethod + def __tensor_unflatten__(inner_tensors: Dict, metadata, outer_size, outer_stride): + assert len(inner_tensors) == 5, "Expected 5 inner tensors" + return NF4Tensor( + metadata["tensor_meta"], + metadata["block_size"], + metadata["n_blocks"], + metadata["scaler_block_size"], + inner_tensors["quantized_scalers"], + inner_tensors["quantization_factor"], + inner_tensors["scaler_mean"], + inner_tensors["quantized_data"], + inner_tensors["nf4"], + ) + + @classmethod + @torch._dynamo.disable + def __torch_dispatch__(cls, func, types, args, kwargs=None): + """TODO we are not supporting torch dispatch at the moment + instead we have created a Autograd.Function to handle the linear + """ + # All ops in the NF4_OPS_TABLE expect NF4 Tensors as inputs + # And don't support mixed tensor subclasses. This will trigger the handler for + # the next type in the dispatch list + + def allowed_subclasses(type): + return ( + issubclass(cls, type) + or issubclass(torch._subclasses.fake_tensor.FakeTensor, type) + or issubclass( + torch._subclasses.functional_tensor.FunctionalTensor, type + ) + ) + + if not all(allowed_subclasses(t) for t in types): + return NotImplemented("Up to the next one to handle") + + if func in NF4_OPS_TABLE: + return NF4_OPS_TABLE[func](func, args, kwargs) + raise NotImplementedError( + f"NF4Tensor dispatch: attempting to run {func}, this is not supported" + ) + + # Do not force the Float8Tensor type on the returned tensor + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + + try: + if func in NF4_TORCH_FUNCTIONS: + return NF4_TORCH_FUNCTIONS[func](*args, **kwargs) + except NotImplementedError: + pass + + with torch._C.DisableTorchFunctionSubclass(): + return func(*args, **kwargs) + + def fsdp_pre_all_gather( + self, mesh: DeviceMesh + ) -> Tuple[Tuple[torch.Tensor, ...], Any]: + return ( + self.quantized_scalers, + self.quantization_factor, + self.quantized_data, + ), ( + SubclassTensorArgs( + self.size(), + self.stride(), + self.storage_offset(), + self.dtype, + self.device, + self.requires_grad, + ), + self.block_size, + self.n_blocks, + self.scaler_block_size, + self.scaler_mean, + self.nf4, + mesh.get_group().size(), + ) + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[torch.Tensor] = None, + ) -> Union[Tuple[torch.Tensor, Tuple[torch.Tensor, ...]], None]: + (quantized_scalers, quantization_factor, quantized_data) = all_gather_outputs + ( + tensor_meta, + block_size, + n_blocks, + scaler_block_size, + scaler_mean, + nf4, + pg_size, + ) = metadata + if len(tensor_meta.original_shape) != 2: + raise NotImplementedError( + f"only support 2D shape but got dim={len(tensor_meta.original_shape)}" + ) + + new_shape = torch.Size( + (tensor_meta.original_shape[0] * pg_size, tensor_meta.original_shape[1]) + ) + new_tensor_meta = replace(tensor_meta, original_shape=new_shape) + if out is not None: + # TODO: add param dtype for mixed precision + assert isinstance(out, NF4Tensor), f"{type(out)}" + assert ( + quantized_scalers.untyped_storage().data_ptr() + == out.quantized_scalers.untyped_storage().data_ptr() + and quantization_factor.untyped_storage().data_ptr() + == out.quantization_factor.untyped_storage().data_ptr() + and quantized_data.untyped_storage().data_ptr() + == out.quantized_data.untyped_storage().data_ptr() + ), "Expects out's data to be the all-gather output" + return + + return nf4_constructor( + new_tensor_meta, + block_size, + n_blocks, + scaler_block_size, + quantized_scalers, + quantization_factor, + scaler_mean, + quantized_data, + nf4, + ), (quantized_scalers, quantization_factor, quantized_data) + + +class LinearNF4(torch.autograd.Function): + @staticmethod + def forward(ctx, input: torch.Tensor, weight: NF4Tensor): + """Save the quantized nf4 weight for backward pass""" + ctx.save_for_backward(weight) + return F.linear(input, weight.to(input.dtype)) + + @staticmethod + def backward(ctx, grad_output): + """The nf4 weight will never require grad so we can just return the grad_output @ weight.to(grad_output.dtype)""" + weight: NF4Tensor = ctx.saved_tensors[0] + return grad_output @ weight.to(grad_output.dtype), None + + +def linear_nf4(input: torch.Tensor, weight: NF4Tensor) -> torch.Tensor: + """Apply a linear operation with the NF4Tensor weight + + Args: + input: Input tensor + weight: NF4Tensor weight + """ + return LinearNF4.apply(input, weight) + + +def to_nf4(tensor, block_size: int = 64, scaler_block_size: int = 256): + """Convert a given tensor to normalized float 4-bit tensor.""" + return NF4Tensor.from_tensor(tensor, block_size, scaler_block_size) + + +NF4_TORCH_FUNCTIONS = {} + + +def implements_torch_function(torch_function): + def decorator(func): + functools.update_wrapper(func, torch_function) + NF4_TORCH_FUNCTIONS[torch_function] = func + return func + + return decorator + + +@implements_torch_function(torch.Tensor.to) +def function_to_dtype(*args, **kwargs): + tensor = args[0] + device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to( + *args[1:], **kwargs + ) + + # dtype is specified -> dequantize + if dtype is not None: + return tensor.get_original_weight().to( + device, dtype, non_blocking, memory_format=convert_to_format + ) + + # dtype is not specified -> keep NF4 + updated_attrs = dict(device=device) + tensor_attrs, _ = tensor.__tensor_flatten__() + for attr in tensor_attrs: + inner_tensor = getattr(tensor, attr) + updated_attrs[attr] = inner_tensor.to( + device, dtype, non_blocking, memory_format=convert_to_format + ) + return NF4Tensor(*construct_nf4_args(tensor, updated_attrs)) + + +@implements_torch_function(torch.Tensor.cpu) +def function_cpu(*args, **kwargs): + # Tensor.cpu(self, memory_format) + return args[0].to("cpu", *args[1:], **kwargs) + + +@implements_torch_function(torch.Tensor.cuda) +def function_cuda(*args, **kwargs): + # Tensor.cuda(self, device, non_blocking, memory_format) + tensor = args[0] + updated_attrs = dict() + tensor_attrs, _ = tensor.__tensor_flatten__() + for attr in tensor_attrs: + inner_tensor = getattr(tensor, attr) + updated_attrs[attr] = inner_tensor.cuda(*args[1:], **kwargs) + updated_attrs["device"] = updated_attrs[tensor_attrs[0]].device + return NF4Tensor(*construct_nf4_args(tensor, updated_attrs)) + + +@implements_torch_function(F.linear) +def _(*args, **kwargs): + input = args[0] + weight = args[1] + bias = args[2] if len(args) > 2 else None + out = LinearNF4.apply(input, weight) + if bias is not None: + out = out + bias + return out + + +@torch._dynamo.allow_in_graph +def nf4_constructor( + tensor_meta: SubclassTensorArgs, + block_size: int, + n_blocks: int, + scaler_block_size: int, + quantized_scalers: torch.Tensor, + quantization_factor: torch.Tensor, + scaler_mean: torch.Tensor, + quantized_data: torch.Tensor, + nf4: torch.Tensor, +): + return NF4Tensor( + tensor_meta, + block_size, + n_blocks, + scaler_block_size, + quantized_scalers, + quantization_factor, + scaler_mean, + quantized_data, + nf4, + ) + + +if TORCH_VERSION_AT_LEAST_2_5: + torch.serialization.add_safe_globals([NF4Tensor]) + torch.serialization.add_safe_globals([NF4Tensor]) diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__init__.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fee6141164ccc000240c38c2eb0db0a8d39e8613 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/__init__.py @@ -0,0 +1,51 @@ +from .block_sparse_layout import ( + BlockSparseLayout, +) +from .cutlass_int4_packed_layout import ( + CutlassInt4PackedLayout, +) +from .int4_cpu_layout import ( + Int4CPULayout, +) +from .int4_xpu_layout import ( + Int4XPULayout, +) +from .marlin_qqq_tensor import ( + MarlinQQQLayout, + MarlinQQQTensor, + to_marlinqqq_quantized_intx, +) +from .marlin_sparse_layout import ( + MarlinSparseLayout, +) +from .packed_linear_int8_dynamic_activation_intx_weight_layout import ( + PackedLinearInt8DynamicActivationIntxWeightLayout, +) +from .q_dq_layout import ( + QDQLayout, +) +from .semi_sparse_layout import ( + SemiSparseLayout, +) +from .tensor_core_tiled_layout import ( + TensorCoreTiledLayout, +) +from .uintx_layout import ( + UintxLayout, +) + +__all__ = [ + "UintxLayout", + "BlockSparseLayout", + "MarlinSparseLayout", + "SemiSparseLayout", + "TensorCoreTiledLayout", + "Int4CPULayout", + "MarlinQQQLayout", + "MarlinQQQTensor", + "to_marlinqqq_quantized_intx", + "CutlassInt4PackedLayout", + "PackedLinearInt8DynamicActivationIntxWeightLayout", + "QDQLayout", + "Int4XPULayout", +] diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dbf32fb4ef9cccd82ce5e3dd4059fee3e5b051c2 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/bitpacking.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/bitpacking.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adc8ddeba6c332ff8020f03ebe7f1211ca66bab3 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/bitpacking.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/block_sparse_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/block_sparse_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b692a76ffeedf3726593092d9f08a637ec757075 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/block_sparse_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/cutlass_int4_packed_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/cutlass_int4_packed_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b00b68e31bf2c2618a2ea00d0d9478743c917c4 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/cutlass_int4_packed_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/gemlite_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/gemlite_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45e386053f72b1d21e4cfce0084b29173f89e712 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/gemlite_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/int4_cpu_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/int4_cpu_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba748f7997ce15ac2c3360f957e1f16532a65977 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/int4_cpu_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/int4_xpu_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/int4_xpu_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63101ca7eb45508174896f24e726cdaa0e8b784a Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/int4_xpu_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/marlin_qqq_tensor.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/marlin_qqq_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a49743c34fc378125c1b37caf2e8f2020c6b27a Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/marlin_qqq_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/marlin_sparse_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/marlin_sparse_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df81a2f7915a3fafad9aa9375c0f4cfc99bfa50c Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/marlin_sparse_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/packed_linear_int8_dynamic_activation_intx_weight_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/packed_linear_int8_dynamic_activation_intx_weight_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b56db26119387e99623d32931e4a7b5dc328ff67 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/packed_linear_int8_dynamic_activation_intx_weight_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/plain_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/plain_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2da816f8958300afbbaabfa1d63cc46eaadf42a Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/plain_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/q_dq_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/q_dq_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c58711682da79dd6a6dc067de3c0774ca4eb9c7d Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/q_dq_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/semi_sparse_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/semi_sparse_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cef59c6ead9eb3228aa149b5132ac31852c64b0f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/semi_sparse_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/tensor_core_tiled_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/tensor_core_tiled_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8e55ec7a9c1df3b2d11799d5c098ec5cf24a5df Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/tensor_core_tiled_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/uint4_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/uint4_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91aef6bb336f62594c1901b6b807b03b6f115e6e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/uint4_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/uintx_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/uintx_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a764a0da4ab3299900fc9447ca2b064d0575cfff Binary files /dev/null and b/lib/python3.12/site-packages/torchao/dtypes/uintx/__pycache__/uintx_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/bitpacking.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/bitpacking.py new file mode 100644 index 0000000000000000000000000000000000000000..4c88643b0e8a2ce70b45da6e4d45523424bd3e24 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/bitpacking.py @@ -0,0 +1,265 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from functools import reduce +from typing import List, Optional + +import torch + +# for selecting the shards from 8 bits +maskbits = { + 1: (0x01,), + 2: (0x03,), + 3: (0x03, 0x04), + 4: (0x0F,), + 5: (0x0F, 0x10), + 6: (0x0F, 0x30), + 7: (0x0F, 0x30, 0x40), +} + +unpack_mask = { + 1: (0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80), + 2: (0x03, 0x0C, 0x30, 0xC0), + 4: (0x0F, 0xF0), +} + +# size of each shard +numbits = { + 1: (1,), + 2: (2,), + 3: (2, 1), + 4: (4,), + 5: (4, 1), + 6: (4, 2), + 7: (4, 2, 1), +} + +# shift amount for each shard +shifts = { + 1: (0,), + 2: (0,), + 3: (0, 2), + 4: (0,), + 5: (0, 4), + 6: (0, 4), + 7: (0, 4, 6), +} + + +# for shifting groups left but right if shift is negative +def abs_lsh(data, shift): + if shift == 0: + return data + elif shift < 0: + return data >> -shift + else: + return data << shift + + +# inverse of abs_lsh for unpacking +def abs_rsh(data, shift): + if shift == 0: + return data + elif shift < 0: + return data << -shift + else: + return data >> shift + + +def pack_cpu( + data: torch.Tensor, elem_size: int, dim: Optional[int] = -1 +) -> List[torch.Tensor]: + """ + Inputs: + data: a tensor of sub byte elements in uint8 + elem_size: the size in bits of the elements to pack + dim: the dimension to pack along + Returns: a list of packed shards + + ================================================================================================== + given an array such as [0x30,0x29,0x17,0x5,0x20,0x16,0x9,0x22] which are 8 uint6 elements + first seperate into two shards: the upper 2 bits and the lower 4 bits by using a mask (0x30 and 0x0f respectively) + 2 bit shard: + mask: 0x30 + [0x30, 0x20, 0x10, 0x00, 0x00, 0x10, 0x00, 0x20 ] + [0b00110000, 0b00100000, 0b00010000, 0b00000000, 0b00100000, 0b00010000, 0b00000000, 0b00100000] + + Group elements into subsets that will be shifted to the same position within the 8bit container + group1 >> 4, group2 >> 2, group3 >> 0, group4 << 2 + + [0b00000011, 0b00000010, 0b00000100, 0b00000000, 0b00100000, 0b00010000, 0b00000000, 0b10000000] + |------ group 1 ------| |------ group 2 ------| |------ group 3 ------| |------ group 4 ------| + + Finally bitwise-or the groups together + [0b00000011, 0b00000010, + 0b00000100, 0b00000000, + 0b00100000, 0b00010000, + 0b00000000, 0b01000000] + + [0b00100111, 0b10010010] + ================================================================================================== + Similarly for 4 bit shards: + mask: 0x0f + [0x00, 0x09, 0x07, 0x05, 0x00, 0x16, 0x9, 0x02] + [0b00000000, 0b00001001, 0b00000111, 0b00000101, 0b00000000, 0b00000110, 0b00001001, 0b00000010] + + group1 << 0, group2 << 4 + [0b00000000, 0b00001001, 0b00000111, 0b00000101, 0b00000000, 0b01100000, 0b10010000, 0b00100000] + |------------------ group 1 ------------------| |------------------ group 2 ------------------| + + bitwise-or: + [0b00000000, 0b00001001, 0b00000111, 0b00000101, + 0b00000000, 0b01100000, 0b10010000, 0b00100000] + + [0b00000000, 0b01101001, 0b10010111, 0b00100101] + ================================================================================================== + After pack, data went from 8 elements to 6: [[0, 105, 151, 37], [39, 146]] + In general this means pack reduces input tensor size from n * 8 to n * elem_size + """ + torch._assert( + data.shape[dim] % 8 == 0, + f"pack dimension size ({data.shape[dim]}) is not divisble by scale", + ) + torch._assert(data.dtype == torch.uint8, "data must be uint8") + output_shape = list(data.shape) + + output = [] + for i in range(len(numbits[elem_size])): + output_shape[dim] = data.shape[dim] * numbits[elem_size][i] // 8 + shard = torch.zeros(output_shape, dtype=torch.uint8, device=data.device) + bit_size = numbits[elem_size][i] + rel_pos = shifts[elem_size][i] + bits = data & maskbits[elem_size][i] + scale = 8 // bit_size + slice_len = bits.shape[dim] // scale + for j in range(scale): + bit_slice = bits.narrow(dim, slice_len * j, slice_len) + shard = torch.bitwise_or(shard, abs_lsh(bit_slice, j * bit_size - rel_pos)) + output.append(shard) + return output + + +def unpack_cpu( + data: List[torch.Tensor], elem_size: int, dim: Optional[int] = -1 +) -> torch.Tensor: + """ + Unpacks small dtype elements from a larger dtype. + + Inputs: + data: - a list of packed shards + elem_size: the size in bits of the elements to unpack + dim: the dimension to unpack along + + Returns: torch.Tensor - a tensor of the unpacked elements. + """ + # define the output tensor + output_shape = list(data[0].shape) + output_shape[dim] = data[0].shape[dim] * 8 // numbits[elem_size][0] + output = torch.zeros(output_shape, dtype=torch.uint8, device=data[0].device) + + for i in range(len(numbits[elem_size])): + # define variables for the current shard + bit_size = numbits[elem_size][i] + rel_pos = shifts[elem_size][i] + scale = 8 // bit_size + group_size = bit_size * output_shape[dim] // 8 + # mask and shift every group of bits to the correct position + for j in range(scale): + output_narrow = output.narrow(dim, j * group_size, group_size) + group = data[i] & unpack_mask[bit_size][j] + output_narrow.copy_( + torch.bitwise_or(output_narrow, abs_rsh(group, j * bit_size - rel_pos)) + ) + return output + + +# these are faster on the GPU + + +def _pack(data, elem_size, scale, dim): + """ + Inner for loop from above pack function + """ + packed_shape = list(data.shape) + packed_shape[dim] = packed_shape[dim] // scale + + packed = torch.zeros(packed_shape, dtype=data.dtype, device=data.device) + + for i in range(scale): + narrow_slice = data.narrow( + dim, data.shape[dim] * i // scale, data.shape[dim] // scale + ) + packed |= narrow_slice << (elem_size * i) + + return packed + + +def _unpack(data, element_size, scale, dim): + """ + Inner for loop from above unpack function + """ + unpacked_shape = list(data.shape) + unpacked_shape[dim] *= scale + + nbits = (1 << element_size) - 1 # mask for the last element_size bits + + unpacked_data = torch.zeros(unpacked_shape, dtype=data.dtype, device=data.device) + + for i in range(scale): + shift_amt = element_size * i + unpacked_data.narrow( + dim, + unpacked_data.shape[dim] * i // scale, + unpacked_data.shape[dim] // scale, + ).copy_((data >> shift_amt) & nbits) + + return unpacked_data + + +def pack( + data: torch.Tensor, elem_size: int, dim: Optional[int] = -1 +) -> List[torch.Tensor]: + """ + a less branching but more compute version so better for gpu + """ + torch._assert( + data.shape[dim] % 8 == 0, + f"pack dimension size ({data.shape[dim]}) is not divisble by scale", + ) + torch._assert(data.dtype == torch.uint8, "data must be uint8") + container_size = 8 + shards = [ + (data & maskbits[elem_size][i]) >> shifts[elem_size][i] + for i in range(len(maskbits[elem_size])) + ] + return tuple( + [ + _pack( + shards[i], + numbits[elem_size][i], + container_size // numbits[elem_size][i], + dim, + ) + for i in range(len(maskbits[elem_size])) + ] + ) + + +def unpack( + data: List[torch.Tensor], elem_size: int, dim: Optional[int] = 0 +) -> torch.Tensor: + """ + a less branching but more compute version so better for gpu + """ + container_size = 8 + # unpack each 4,2,1 bit shard and unshift them back to the correct position + data = [ + _unpack( + data[i], numbits[elem_size][i], container_size // numbits[elem_size][i], dim + ) + << shifts[elem_size][i] + for i in range(len(data)) + ] + return reduce(torch.bitwise_or, data) diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/block_sparse_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/block_sparse_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..0c6046c313bc1f1c1f71d52b129d068f0d21206a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/block_sparse_layout.py @@ -0,0 +1,233 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import logging +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from torch.utils._python_dispatch import ( + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.uintx.plain_layout import ( + PlainAQTTensorImpl, + _aqt_is_int8_reduced_range, +) +from torchao.dtypes.utils import ( + Layout, + PlainLayout, +) + +logger = logging.getLogger(__name__) + +aten = torch.ops.aten + + +@dataclass(frozen=True) +class BlockSparseLayout(Layout): + """BlockSparseLayout is a data class that represents the layout of a block sparse matrix. + + Attributes: + blocksize (int): The size of the blocks in the sparse matrix. Default is 64. + """ + + blocksize: int = 64 + + +@register_layout(BlockSparseLayout) +class BlockSparseAQTTensorImpl(PlainAQTTensorImpl): + bsr_crow_indices: Optional[torch.Tensor] + bsr_col_indices: Optional[torch.Tensor] + bsr_values: Optional[torch.Tensor] + scale: Optional[torch.Tensor] + zero_point: Optional[torch.Tensor] + + __slots__ = [ + "bsr_crow_indices", + "bsr_col_indices", + "bsr_values", + "scale", + "zero_point", + ] + + @staticmethod + def __new__( # noqa: PYI034 + cls, + shape: torch.Size, + bsr_crow_indices: Optional[torch.Tensor], + bsr_col_indices: Optional[torch.Tensor], + bsr_values: Optional[torch.Tensor], + scale: Optional[torch.Tensor], + zero_point: Optional[torch.Tensor], + _layout: Layout, + requires_grad: bool = False, + ): + if bsr_values is None: + raise ValueError("bsr values must be provided!") + else: + previous_tensor = bsr_values + + kwargs = { + "device": previous_tensor.device, + "dtype": previous_tensor.dtype, + "layout": previous_tensor.layout, + "requires_grad": requires_grad, + } + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( # noqa: PYI034 + self, + shape: torch.Size, + bsr_crow_indices: Optional[torch.Tensor], + bsr_col_indices: Optional[torch.Tensor], + bsr_values: Optional[torch.Tensor], + scale: Optional[torch.Tensor], + zero_point: Optional[torch.Tensor], + _layout: Layout, + requires_grad: bool = False, + ): + self.bsr_crow_indices = bsr_crow_indices + self.bsr_col_indices = bsr_col_indices + self.bsr_values = bsr_values + self.scale = scale + self.zero_point = zero_point + self._layout = _layout + + def __tensor_flatten__(self): + inner_tensors = list( + filter(lambda x: getattr(self, x) is not None, self.__slots__) + ) + tensor_meta = (self.shape, self._layout, self.requires_grad) + return inner_tensors, tensor_meta + + @classmethod + def __tensor_unflatten__( + cls, + inner_tensors, + tensor_meta: Tuple[torch.Size, bool], + outer_size, + outer_stride, + ) -> torch.Tensor: + shape, _layout, requires_grad = tensor_meta + return cls( + shape=shape, + bsr_crow_indices=inner_tensors.get("bsr_crow_indices", None), + bsr_col_indices=inner_tensors.get("bsr_col_indices", None), + bsr_values=inner_tensors.get("bsr_values", None), + scale=inner_tensors.get("scale", None), + zero_point=inner_tensors.get("zero_point", None), + _layout=_layout, + requires_grad=requires_grad, + ) + + @classmethod + def from_plain(cls, int_data, scale, zero_point, _layout): + bsr_tensor = int_data.to_sparse_bsr(_layout.blocksize) + return cls( + shape=int_data.shape, + bsr_crow_indices=bsr_tensor.crow_indices(), + bsr_col_indices=bsr_tensor.col_indices(), + bsr_values=bsr_tensor.values(), + scale=scale, + zero_point=zero_point, + _layout=_layout, + requires_grad=False, + ) + + def get_plain(self): + int_data_expanded = torch.ops.blocksparse.bsr_to_dense( + self.crow_indices(), + self.col_indices(), + self.values(), + self.shape[0], + self.shape[1], + ) + return int_data_expanded, self.scale, self.zero_point + + def _apply_fn_to_data(self, func): + return self.__class__( + shape=self.shape, + bsr_crow_indices=func(self.bsr_crow_indices), + bsr_col_indices=func(self.bsr_col_indices), + bsr_values=func(self.bsr_values), + scale=self.scale, + zero_point=self.zero_point, + _layout=self._layout, + requires_grad=self.requires_grad, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + if func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + # Need the following for bsr specific functions + if func is aten.crow_indices.default: + return args[0].bsr_crow_indices.detach() + + if func is aten.col_indices.default: + return args[0].bsr_col_indices.detach() + + if func is aten.values.default: + return args[0].bsr_values.detach() + + if func is aten._nnz.default: + return args[0].bsr_values.shape[0] + + raise NotImplementedError( + f"BlockSparseAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + +def _linear_int8_act_int8_weight_block_sparse_check(input_tensor, weight_tensor, bias): + return ( + isinstance(input_tensor, AffineQuantizedTensor) + and _aqt_is_int8_reduced_range(input_tensor) + and isinstance(weight_tensor, AffineQuantizedTensor) + and weight_tensor.is_cuda + and input_tensor.dtype == weight_tensor.dtype + and isinstance(input_tensor._layout, PlainLayout) + and isinstance(weight_tensor._layout, BlockSparseLayout) + ) + + +def _linear_int8_act_int8_weight_block_sparse_impl(input_tensor, weight_tensor, bias): + x_vals_int8 = input_tensor.tensor_impl.int_data + x_scales = input_tensor.tensor_impl.scale + w_vals = weight_tensor.tensor_impl + w_scales = weight_tensor.tensor_impl.scale + tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1]) + tmp_t = tmp.t() + + y = torch.ops.blocksparse.int_addmm( + w_vals.crow_indices(), + w_vals.col_indices(), + w_vals.values(), + tmp_t, + w_scales, + x_scales.reshape(-1), + ) + y_shape = (*x_vals_int8.shape[:-1], w_scales.shape[-1]) + y = y.reshape(*y_shape) + + # can downcast only at the very end + output_dtype = input_tensor.dtype + y = y.to(output_dtype) + if bias is not None: + y += bias + return y diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/cutlass_int4_packed_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/cutlass_int4_packed_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..d2fcaf90459933e6805949e29c8cfd532990fe3f --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/cutlass_int4_packed_layout.py @@ -0,0 +1,224 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Optional + +import torch +from torch.utils._python_dispatch import ( + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.uintx.plain_layout import ( + _aqt_is_int8, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout, PlainLayout + +aten = torch.ops.aten + + +def _aqt_is_int4(aqt): + """Check if an AffineQuantizedTensor is int4 quantized Tensor""" + # TODO: use torch.int4 + return ( + aqt.tensor_impl.dtype == torch.int8 + and aqt.quant_min == -8 + and aqt.quant_max == 7 + ) + + +def _same_metadata(self: "Int4PackedTensorImpl", src: "Int4PackedTensorImpl") -> bool: + return ( + isinstance(self, Int4PackedTensorImpl) + and isinstance(src, Int4PackedTensorImpl) + and self.shape == src.shape + and self.int_data.shape == src.int_data.shape + and self.scale.shape == src.scale.shape + and type(self._layout) == type(src._layout) + ) + + +@dataclass(frozen=True) +class CutlassInt4PackedLayout(Layout): + """Layout class for int4 packed layout for affine quantized tensor, for cutlass kernel.""" + + pass + + +@register_layout(CutlassInt4PackedLayout) +class Int4PackedTensorImpl(AQTTensorImpl): + """ + TensorImpl storage class for int4 packed layout for affine quantized tensor. + """ + + @staticmethod + def __new__( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = int_data.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout + ) + kwargs["dtype"] = int_data.dtype + kwargs["requires_grad"] = False + shape = int_data.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + int_data: torch.Tensor, + scale: torch.Tensor, + _layout: Layout, + ): + self.int_data = int_data + self.scale = scale + self._layout = _layout + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + elif func is aten.copy_.default: + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + + raise NotImplementedError( + f"Int4PackedTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + def __tensor_flatten__(self): + return ["int_data", "scale"], [self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + int_data = tensor_data_dict["int_data"] + scale = tensor_data_dict["scale"] + (_layout,) = tensor_attributes + return cls(int_data, scale, _layout) + + def get_plain(self): + int_data = torch.stack( + ((self.int_data << 4) >> 4, self.int_data >> 4), dim=-1 + ).view(self.int_data.shape[:-1] + (2 * self.int_data.shape[-1],)) + return int_data, self.scale, None + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert zero_point is None or torch.all(zero_point == 0) + int_data_s4 = ((int_data[..., 1::2] & 0xF) << 4) | (int_data[..., 0::2] & 0xF) + return cls( + int_data_s4, + scale, + _layout, + ) + + def get_layout(self) -> Layout: + return self._layout + + def _apply_fn_to_data(self, fn): + self.int_data = fn(self.int_data) + self.scale = fn(self.scale) + return self + + +def _linear_int8_act_int4_weight_cutlass_check(input_tensor, weight_tensor, bias): + return ( + isinstance(input_tensor, AffineQuantizedTensor) + and isinstance(input_tensor._layout, PlainLayout) + and _aqt_is_int8(input_tensor) + and input_tensor.dtype in (torch.float16, torch.bfloat16) + and len(input_tensor.shape) >= 2 + and input_tensor.tensor_impl.scale.dtype == torch.float32 + and len(input_tensor.tensor_impl.scale.shape) == len(input_tensor.shape) - 1 + and isinstance(weight_tensor, AffineQuantizedTensor) + and isinstance(weight_tensor._layout, CutlassInt4PackedLayout) + and _aqt_is_int4(weight_tensor) + and weight_tensor.dtype == input_tensor.dtype + and len(weight_tensor.shape) == 2 + and weight_tensor.tensor_impl.scale.dtype == torch.float32 + and len(weight_tensor.tensor_impl.scale.shape) == 1 + and (bias is None or bias.dtype == input_tensor.dtype) + and (bias is None or len(bias.shape) == 1) + ) + + +def _linear_int8_act_int4_weight_cutlass_impl(input_tensor, weight_tensor, bias): + from torchao.ops import rowwise_scaled_linear_cutlass_s8s4 + + weight = weight_tensor.tensor_impl.int_data + weight_scale = weight_tensor.tensor_impl.scale + input = input_tensor.tensor_impl.int_data + input_scale = input_tensor.tensor_impl.scale + out_dtype = input_tensor.dtype + + out = rowwise_scaled_linear_cutlass_s8s4( + input, input_scale, weight, weight_scale, bias, out_dtype + ) + + return out + + +def _linear_int4_act_int4_weight_cutlass_check(input_tensor, weight_tensor, bias): + return ( + isinstance(input_tensor, AffineQuantizedTensor) + and isinstance(input_tensor._layout, CutlassInt4PackedLayout) + and _aqt_is_int4(input_tensor) + and input_tensor.dtype in (torch.float16, torch.bfloat16) + and len(input_tensor.shape) >= 2 + and input_tensor.tensor_impl.scale.dtype == torch.float32 + and len(input_tensor.tensor_impl.scale.shape) == len(input_tensor.shape) - 1 + and isinstance(weight_tensor, AffineQuantizedTensor) + and isinstance(weight_tensor._layout, CutlassInt4PackedLayout) + and _aqt_is_int4(weight_tensor) + and weight_tensor.dtype == input_tensor.dtype + and len(weight_tensor.shape) == 2 + and weight_tensor.tensor_impl.scale.dtype == torch.float32 + and len(weight_tensor.tensor_impl.scale.shape) == 1 + ) + + +def _linear_int4_act_int4_weight_cutlass_impl(input_tensor, weight_tensor, bias): + from torchao.ops import rowwise_scaled_linear_cutlass_s4s4 + + weight = weight_tensor.tensor_impl.int_data + weight_scale = weight_tensor.tensor_impl.scale + input = input_tensor.tensor_impl.int_data + input_scale = input_tensor.tensor_impl.scale + out_dtype = input_tensor.dtype + + out = rowwise_scaled_linear_cutlass_s4s4( + input, input_scale, weight, weight_scale, bias, out_dtype + ) + + return out diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/gemlite_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/gemlite_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..0c124eb34313968d240eb32578dee0dbcf0ecd7b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/gemlite_layout.py @@ -0,0 +1,403 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Dict, Optional, Tuple + +import torch +from torch.utils._python_dispatch import ( + is_traceable_wrapper_subclass, + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.uintx.tensor_core_tiled_layout import TensorCoreTiledAQTTensorImpl +from torchao.dtypes.utils import Layout +from torchao.utils import fill_defaults + +try: + import gemlite + from gemlite.core import GemLiteLinearTriton +except: + gemlite = None + + +aten = torch.ops.aten + + +def _same_metadata( + self: "GemliteAQTTensorImpl", + src: "GemliteAQTTensorImpl", +) -> bool: + kwargs_match = len(self.gemlite_kwargs) == len(src.gemlite_kwargs) + for k, v in self.gemlite_kwargs.items(): + if k != "scale_activations": + kwargs_match = kwargs_match and (v == src.gemlite_kwargs[k]) + + return ( + isinstance(self, GemliteAQTTensorImpl) + and isinstance(src, GemliteAQTTensorImpl) + and self.shape == src.shape + and self.packed_weight.shape == src.packed_weight.shape + and self.scale.shape == src.scale.shape + and self.zero_point.shape == src.zero_point.shape + and kwargs_match + and type(self._layout) == type(src._layout) + ) + + +def scale_activations_no_scaling(x): + return x, None + + +def scale_activations_int8(x): + x_shape = x.shape + out_x = x.view(-1, x.shape[-1]) + scaled_x = torch.abs(out_x).amax(axis=1, keepdim=True) / 127 + out_x = torch.round(out_x / scaled_x).to(dtype=torch.int8) + return out_x.view(x_shape), scaled_x + + +def get_gemlite_quant_kwargs(bit_width, group_size, dtype): + from torchao.quantization.quant_primitives import MappingType, ZeroPointDomain + + kwargs = {} + if bit_width != 8: + kwargs["mapping_type"] = MappingType.ASYMMETRIC + kwargs["block_size"] = (1, group_size) + kwargs["target_dtype"] = torch.uint8 + kwargs["eps"] = 1e-6 + kwargs["quant_min"] = 0 + kwargs["quant_max"] = (2**bit_width) - 1 + kwargs["eps"] = 1e-6 + kwargs["zero_point_dtype"] = dtype + kwargs["zero_point_domain"] = ZeroPointDomain.FLOAT + elif bit_width == 8: + kwargs["mapping_type"] = MappingType.SYMMETRIC + kwargs["block_size"] = (1, group_size) + kwargs["target_dtype"] = torch.int8 + kwargs["quant_min"] = -128 + kwargs["quant_max"] = 127 + kwargs["eps"] = 1e-5 + kwargs["zero_point_dtype"] = None + kwargs["zero_point_domain"] = ZeroPointDomain.NONE + return kwargs + + +def get_gemlite_aqt_kwargs( + weight, + group_size=64, + bit_width=4, + packing_bitwidth=32, + contiguous=None, + use_hqq=True, +): + if gemlite is None: + raise ImportError( + "Unable to import 'gemlite'. Please ensure it is installed correctly. You can install it with: pip install gemlite" + ) + + assert bit_width in [ + 4, + 8, + ], f"gemlite only works with bit_width 4,8 but got {bit_width}" + assert packing_bitwidth in [ + 8, + 16, + 32, + None, + ], f"gemlite needs packing_bitwidth in [8, 16, 32] but got {packing_bitwidth}" + assert weight.dtype in [torch.float16, torch.bfloat16], ( + f"gemlite only works with dtype torch.float16 or torch.bfloat16 but got {weight.dtype}" + ) + assert group_size in [32, 64, 128, 256, 512, 1024, None] + assert group_size is None or bit_width != 8, ( + "gemlite only works with group_size=None for bit_width=8" + ) + + out_features, in_features = weight.shape + group_size = in_features if group_size is None else group_size + + aqt_kwargs = get_gemlite_quant_kwargs(bit_width, group_size, weight.dtype) + aqt_kwargs["_layout"] = GemlitePackedLayout( + group_size=group_size, + bit_width=bit_width, + packing_bitwidth=packing_bitwidth, + contiguous=contiguous, + ) + aqt_kwargs["use_hqq"] = use_hqq + return aqt_kwargs + + +@dataclass(frozen=True) +class GemlitePackedLayout(Layout): + group_size: Optional[int] = 64 + bit_width: int = 4 + packing_bitwidth: int = None + contiguous: bool = None + + +@register_layout(GemlitePackedLayout) +class GemliteAQTTensorImpl(TensorCoreTiledAQTTensorImpl): + def __new__( + cls, + packed_weight: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + gemlite_kwargs: Dict, + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = packed_weight.device + kwargs["layout"] = ( + kwargs.get("layout") + if kwargs.get("layout", False) + else packed_weight.layout + ) + kwargs["dtype"] = packed_weight.dtype + kwargs["requires_grad"] = False + shape = packed_weight.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + packed_weight: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + gemlite_kwargs: Dict, + _layout: Layout, + ): + self.packed_weight = packed_weight + self.scale = scale + self.zero_point = zero_point + self.gemlite_kwargs = gemlite_kwargs + self._layout = _layout + + def __tensor_flatten__(self): + return ["packed_weight", "scale", "zero_point"], [ + self._layout, + self.gemlite_kwargs, + ] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + packed_weight, scale, zero_point = ( + tensor_data_dict["packed_weight"], + tensor_data_dict["scale"], + tensor_data_dict["zero_point"], + ) + _layout, gemlite_kwargs = tensor_attributes + return cls(packed_weight, scale, zero_point, gemlite_kwargs, _layout) + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert isinstance(_layout, GemlitePackedLayout), ( + f"GemliteAQTTensorImpl only works with GemliteLinearTriton but got {_layout}" + ) + device = int_data.device + if device.type != "cuda": + int_data = ( + int_data.cuda() + ) # We need int_data on cuda device because of Triton packing + + group_size, bit_width = _layout.group_size, _layout.bit_width + out_features, in_features = int_data.shape + + gemlite_linear = gemlite.helper.A16Wn(device=int_data.device).from_weights( + int_data, scale, zero_point, bit_width, group_size, bias=None + ) + + gemlite_kwargs = { + "out_features": out_features, + "scaled_activations": gemlite_linear.scaled_activations, + "meta_args": gemlite_linear.get_meta_args(), + } + + packed_weight, scale, zero_point = gemlite_linear.get_tensor_args() + packed_weight = packed_weight.to(device) + + return cls(packed_weight, scale, zero_point, gemlite_kwargs, _layout) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs["device"] + return self.__class__( + self.packed_weight.to(device), + self.scale.to(device), + self.zero_point.to(device), + self.gemlite_kwargs, + self._layout, + ) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.packed_weight), + fn(self.scale), + fn(self.zero_point), + self.gemlite_kwargs, + self._layout, + ) + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + device = self.packed_weight.device + elements_per_sample = self._layout.packing_bitwidth // self._layout.bit_width + in_features = ( + self.packed_weight.numel() * elements_per_sample + ) // self.gemlite_kwargs["out_features"] + int_data = ( + gemlite.bitpack.unpack_over_rows( + self.packed_weight.cuda(), + W_nbits=self._layout.bit_width, + num_output_rows=in_features, + dtype=torch.uint8, + ) + .t() + .contiguous() + ).to(device) + scale = self.scale.t().contiguous() + zero_point = self.zero_point.t().contiguous() + + return int_data, scale, zero_point + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + # we don't handle transpose operations and just ignore them. In practice the only + # reason a transpsoe should occur is because the functional linear + # op can decompose into e.g. transpose + addmm so since we want + # to use the gemlite matmul kernel, which expects teh weight to be passed in as is, + # we ignore the transpose + if func is aten.detach.default or func is aten.t.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + if func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + if func is aten.slice.Tensor: + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + assert step == 1, "Only step == 1 is supported in slicing right now" + + if dim in [0, 1]: + int_data, scale, zero_point = self.get_plain() + data_len = int_data.shape[dim] + scale_len = scale.shape[dim] + ratio = data_len / scale_len + start_scale = int(start / ratio) + end_scale = int(end / ratio) + + int_data = aten.slice.Tensor(int_data, dim, start, end, step) + scale = aten.slice.Tensor(scale, dim, start_scale, end_scale, step) + if zero_point is not None and zero_point.numel() > 0: + zero_point = aten.slice.Tensor( + zero_point, dim, start_scale, end_scale, step + ) + else: + zero_point = None + # this is to handle padding + int_data, scale, zero_point = self._layout.post_process( + int_data, scale, zero_point, self.block_size + ) + + sliced = self.from_plain( + int_data, scale, zero_point, self._layout + ) # Will be transposed again + + return return_and_correct_aliasing(func, args, kwargs, sliced) + + else: + raise NotImplementedError( + f"GemliteAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported" + ) + + elif func is aten.copy_.default: + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + + raise NotImplementedError( + f"GemliteAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + __torch_function__ = torch._C._disabled_torch_function_impl + + def get_layout(self) -> Layout: + return self._layout + + @property + def block_size(self): + return (1, self._layout.group_size) + + +# logic taken from gemlite's core.py +def _matmul_type_fn(batch_size: int, bit_width: int) -> str: + if batch_size > 64: + return "GEMM" + elif batch_size > 1: + return "GEMM_SPLITK" + else: + return gemlite.core.get_default_gemv(bit_width) + + +def _linear_fp_act_int4_weight_gemlite_impl(input_tensor, weight_tensor, bias=None): + if hasattr(weight_tensor, "tensor_impl"): + weight_impl = weight_tensor.tensor_impl + else: + weight_impl = weight_tensor + + batch_size = input_tensor.view(-1, input_tensor.shape[-1]).shape[0] + matmul_type = _matmul_type_fn(batch_size, weight_impl._layout.bit_width) + + if weight_impl.gemlite_kwargs["scaled_activations"]: + scale_activations = scale_activations_int8 + else: + scale_activations = scale_activations_no_scaling + + return GemLiteLinearTriton.forward_functional( + x=input_tensor, + bias=bias, + matmul_type=matmul_type, + out_features=weight_impl.gemlite_kwargs["out_features"], + scale_activations=scale_activations, + meta_args=weight_impl.gemlite_kwargs["meta_args"], + tensor_args=( + weight_impl.packed_weight, + weight_impl.scale, + weight_impl.zero_point, + ), + ) + + +def _linear_fp_act_int4_weight_gemlite_check(input_tensor, weight_tensor, bias): + return ( + # input is native fp16 tensor + not is_traceable_wrapper_subclass(input_tensor) + # and input_tensor.dtype in [torch.float16, torch.bfloat16] + # weight is gemlite layout + and isinstance(weight_tensor, AffineQuantizedTensor) + and isinstance(weight_tensor._layout, GemlitePackedLayout) + ) diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/int4_cpu_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/int4_cpu_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..9c368fd17acd5707e4fe0b9c398d1e02b571e436 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/int4_cpu_layout.py @@ -0,0 +1,354 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from torch.utils._python_dispatch import ( + is_traceable_wrapper_subclass, + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout, is_device +from torchao.quantization.quant_primitives import ZeroPointDomain +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_5, + TORCH_VERSION_AT_LEAST_2_6, + fill_defaults, +) + +aten = torch.ops.aten + + +@dataclass(frozen=True) +class Int4CPULayout(Layout): + """Layout class for int4 CPU layout for affine quantized tensor, used by tinygemm kernels `_weight_int4pack_mm_for_cpu`. + Only for PyTorch version at least 2.6 + """ + + pass + + +@register_layout(Int4CPULayout) +class Int4CPUAQTTensorImpl(AQTTensorImpl): + """TensorImpl for int4 CPU layout for affine quantized tensor, this is for int4 only, + used by tinygemm kernels `_weight_int4pack_mm_for_cpu` + It stores the original tensor of dimension [n][k] (int32 dtype) as packed weight of 2-d tensor of + dimension: [n][k / 2] (uint8 dtype) + (unpacked Tensor shape is n * k) + Note: we also pack scale and zero point together here for tinygemm kernel + Note: technically Int4 CPU layout should be the layout for the underlying packed weight + (int Tensor) but since the scale and zero_point are also packed into the same tensor here which is not used + in plain layout, we just created a layout for AQT right now, this could be improved if we split out + int4 aqt into a separate tensor subclass + fields: + packed_weight (torch.Tensor): the 2-d packed tensor in a Int4 CPU layout + scale_and_zero (torch.Tensor): the combined scale Tensor used to map between floating point tensor to quantized tensor and zero_point Tensor + """ + + def __new__( + cls, + packed_weight: torch.Tensor, + scale_and_zero: torch.Tensor, + transposed: bool, + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = packed_weight.device + kwargs["layout"] = ( + kwargs.get("layout") + if kwargs.get("layout", False) + else packed_weight.layout + ) + kwargs["dtype"] = packed_weight.dtype + kwargs["requires_grad"] = False + shape = packed_weight.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + packed_weight: torch.Tensor, + scale_and_zero: torch.Tensor, + transposed: bool, + _layout: Layout, + ): + self.packed_weight = packed_weight + self.scale_and_zero = scale_and_zero + self.transposed = False + self._layout = _layout + + def __tensor_flatten__(self): + return ["packed_weight", "scale_and_zero"], [self.transposed, self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + packed_weight, scale_and_zero = ( + tensor_data_dict["packed_weight"], + tensor_data_dict["scale_and_zero"], + ) + ( + transposed, + _layout, + ) = tensor_attributes + return cls(packed_weight, scale_and_zero, transposed, _layout) + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert isinstance(_layout, Int4CPULayout) + + if TORCH_VERSION_AT_LEAST_2_6: + assert int_data.dtype == torch.int32, ( + "torch.ops.aten._convert_weight_to_int4pack_for_cpu expects `int32` dtype" + ) + packed_weight = torch.ops.aten._convert_weight_to_int4pack_for_cpu( + int_data, + 1, # TODO:remove + ) + elif TORCH_VERSION_AT_LEAST_2_5: + int_data = (int_data[::, ::2] << 4 | int_data[::, 1::2]).to(torch.uint8) + assert int_data.dtype == torch.uint8, ( + "torch.ops.aten._convert_weight_to_int4pack in torch 2.5 expects `uint8` dtype" + ) + packed_weight = torch.ops.aten._convert_weight_to_int4pack( + int_data, _layout.inner_k_tiles + ) + else: + assert int_data.dtype == torch.int32, ( + "torch.ops.aten._convert_weight_to_int4pack in torch 2.4 expects `int32` dtype" + ) + packed_weight = torch.ops.aten._convert_weight_to_int4pack( + int_data, _layout.inner_k_tiles + ) + + scale = scale.reshape(int_data.shape[0], -1) + zero_point = zero_point.reshape(int_data.shape[0], -1) + from torchao.quantization.utils import pack_tinygemm_scales_and_zeros + + scale_and_zero = pack_tinygemm_scales_and_zeros(scale, zero_point, scale.dtype) + return cls(packed_weight, scale_and_zero, False, _layout) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs["device"] + if not is_device(torch.device(self.device).type, device): + raise ValueError( + f"Int4CPUAQTTensorImpl does not support conversion from {self.device} to {device}" + ) + return self.__class__( + self.packed_weight.to(device), + self.scale_and_zero.to(device), + self.transposed, + self._layout, + ) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.packed_weight), + fn(self.scale_and_zero), + self.transposed, + self._layout, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + if func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + if func is aten.t.default: + """we don't need to repack the weight and just rely on external + shape being changed and record the status of transpose/no-transpose + """ + transposed = Int4CPUAQTTensorImpl( + args[0].packed_weight, + args[0].scale_and_zero, + not args[0].transposed, + args[0]._layout, + ) + return return_and_correct_aliasing(func, args, kwargs, transposed) + + if func is aten.slice.Tensor: + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + if dim in [0, 1]: + assert step == 1, "Only step == 1 is supported in slicing right now" + int_data, scale, zero_point = self.get_plain() + data_len = int_data.shape[dim] + scale_len = scale.shape[dim] + ratio = data_len / scale_len + start_scale = int(start / ratio) + end_scale = int(end / ratio) + + int_data = aten.slice.Tensor(int_data, dim, start, end, step) + scale = aten.slice.Tensor(scale, dim, start_scale, end_scale, step) + zero_point = aten.slice.Tensor( + zero_point, dim, start_scale, end_scale, step + ) + # this is to handle padding + int_data, scale, zero_point = self._layout.post_process( + int_data, scale, zero_point, self.block_size + ) + sliced = self.from_plain(int_data, scale, zero_point, self._layout) + return return_and_correct_aliasing(func, args, kwargs, sliced) + else: + raise NotImplementedError( + f"Int4CPUAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported" + ) + + raise NotImplementedError( + f"Int4CPUAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + __torch_function__ = torch._C._disabled_torch_function_impl + + @property + def block_size(self): + from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros + + scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero) + cur_shape = self.shape + assert len(cur_shape) == 4 + inner_k_tiles = cur_shape[-1] * 2 + original_shape = (cur_shape[0] * 8, cur_shape[1] * (inner_k_tiles * 16)) + groupsize = int(original_shape[1] / scale.shape[-2]) + return (1, groupsize) + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from torchao.quantization.quant_primitives import ( + ZeroPointDomain, + quantize_affine, + ) + from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros + + scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero) + + cur_shape = self.shape + assert len(cur_shape) == 2 + original_shape = (cur_shape[0], cur_shape[1] * 2) + eye_shape = original_shape[1] + groupsize = int(original_shape[1] / scale.shape[-2]) + block_size = (1, groupsize) + device = self.device + original_dtype = self.scale_and_zero.dtype + target_dtype = torch.int32 + quant_min = 0 + quant_max = 15 + zero_point_domain = ZeroPointDomain.FLOAT + assert len(block_size) == 2 and block_size[0] == 1 + dequantized = torch.ops.aten._weight_int4pack_mm_for_cpu( + torch.eye(eye_shape, device=device, dtype=original_dtype), + self.packed_weight, + groupsize, + self.scale_and_zero, + ) + dequantized = dequantized.t().contiguous() + # TODO: move this to `unpack_tinygemm_scales_and_zeros`? + scale = scale.reshape(scale.shape[:-1]).contiguous() + zero = zero.reshape(zero.shape[:-1]).contiguous() + int_data = quantize_affine( + dequantized, + block_size, + scale, + zero, + target_dtype, + quant_min, + quant_max, + zero_point_domain, + ) + return int_data, scale, zero + + def get_layout(self) -> Layout: + return self._layout + + +def _aqt_is_uint4(aqt): + """Check if an AffineQuantizedTensor is uint4 quantized Tensor""" + return ( + aqt.tensor_impl.dtype == torch.uint8 + and aqt.quant_min == 0 + and aqt.quant_max == 15 + ) + + +def _is_float(dtype): + return dtype in (torch.float, torch.half, torch.bfloat16) + + +def _linear_fp_act_uint4_weight_cpu_check(input_tensor, weight_tensor, bias): + return ( + TORCH_VERSION_AT_LEAST_2_6 + and is_device(input_tensor.device.type, "cpu") + and is_device(weight_tensor.device.type, "cpu") + and (bias is None or is_device(bias.device.type, "cpu")) + and not is_traceable_wrapper_subclass(input_tensor) + and _is_float(input_tensor.dtype) + and isinstance(weight_tensor, AffineQuantizedTensor) + and _aqt_is_uint4(weight_tensor) + and _is_float(weight_tensor.dtype) + and len(weight_tensor.shape) == 2 + and weight_tensor.zero_point_domain == ZeroPointDomain.FLOAT + and isinstance(weight_tensor._layout, Int4CPULayout) + ) + + +def _linear_fp_act_uint4_weight_cpu_impl(input_tensor, weight_tensor, bias): + assert TORCH_VERSION_AT_LEAST_2_6, ( + f"Requires PyTorch version at least 2.6, but got: {torch.__version__}" + ) + assert is_device(input_tensor.device.type, "cpu"), ( + f"For CPU device only but got: {input_tensor.device}" + ) + assert weight_tensor.block_size[0] == 1, ( + f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}" + ) + assert input_tensor.shape[-1] == weight_tensor.shape[1], ( + f"need input_tensor shape: {input_tensor.shape} final" + f"dim to match weight_tensor shape: {weight_tensor.shape} second dim " + ) + + act_mat = input_tensor + packed_weight = weight_tensor.tensor_impl.packed_weight + scale_and_zero = weight_tensor.tensor_impl.scale_and_zero + + orig_act_size = act_mat.size() + orig_dtype = act_mat.dtype + + # reshape to 2D + act_mat = act_mat.reshape(-1, act_mat.shape[-1]) + + # groupwise int4 quantization + groupsize = weight_tensor.block_size[1] + y = torch.ops.aten._weight_int4pack_mm_for_cpu( + act_mat.contiguous(), packed_weight, groupsize, scale_and_zero + ) + + # remove out_feature padding + orig_out_features = weight_tensor.shape[-2] + y = y[:, :orig_out_features] + y = y.reshape(*orig_act_size[:-1], orig_out_features) + + if bias is not None: + y += bias + return y.to(orig_dtype) diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/int4_xpu_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/int4_xpu_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..99370b084842761020a9936ca44959a083c4ff04 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/int4_xpu_layout.py @@ -0,0 +1,445 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from torch.utils._python_dispatch import ( + is_traceable_wrapper_subclass, + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout, is_device +from torchao.quantization.quant_primitives import ZeroPointDomain +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_8, + fill_defaults, +) + +aten = torch.ops.aten + + +def _aqt_is_xpu_layout_uint4(aqt): + """Check if an AffineQuantizedTensor is uint4 quantized Tensor""" + # TODO: use torch.uint4 + return ( + aqt.tensor_impl.dtype == torch.int32 + and aqt.quant_min == 0 + and aqt.quant_max == 15 + ) + + +def _linear_bf16_act_uint4_weight_float_zero_check(input_tensor, weight_tensor, bias): + return ( + # input is native bfloat16 tensor + not is_traceable_wrapper_subclass(input_tensor) + and input_tensor.dtype == torch.bfloat16 + and + # weight is uint4, group quantized tensor_core_tiled tensor impl affine quantized tensor + isinstance(weight_tensor, AffineQuantizedTensor) + and _aqt_is_xpu_layout_uint4(weight_tensor) + and weight_tensor.dtype == torch.bfloat16 + and len(weight_tensor.shape) == 2 + and weight_tensor.zero_point_domain == ZeroPointDomain.FLOAT + and weight_tensor.tensor_impl.scale_and_zero is not None + and weight_tensor.tensor_impl.scale_and_zero.dtype == torch.bfloat16 + and isinstance(weight_tensor._layout, Int4XPULayout) + ) + + +def _linear_bf16_act_uint4_weight_float_zero_impl(input_tensor, weight_tensor, bias): + assert weight_tensor.block_size[0] == 1, ( + f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}" + ) + assert input_tensor.shape[-1] == weight_tensor.shape[1], ( + f"need input_tensor shape: {input_tensor.shape} final" + f"dim to match weight_tensor shape: {weight_tensor.shape} second dim " + ) + + # TODO: check groupsize quantization + # avoid circular dep, TODO: move this to a common util.py + act_mat = input_tensor + if act_mat.is_contiguous() == False: + act_mat = act_mat.contiguous() + # weight is packed from padded (out_features, in_features) weight tensor + # (same dimension requirement as F.linear weight) + packed_weight = weight_tensor.tensor_impl.packed_weight + scales_and_zeros = weight_tensor.tensor_impl.scale_and_zero + + orig_act_size = act_mat.size() + orig_dtype = act_mat.dtype + + act_mat = act_mat.reshape(-1, act_mat.shape[-1]).to(torch.bfloat16) + + # groupwise int4 quantization + groupsize = weight_tensor.block_size[1] + y = torch.ops.aten._weight_int4pack_mm( + act_mat, packed_weight, groupsize, scales_and_zeros + ) + + # remove out_feature padding + orig_out_features = weight_tensor.shape[-2] + y = y[:, :orig_out_features] + y = y.reshape(*orig_act_size[:-1], orig_out_features) + + if bias is not None: + y += bias + return y.to(orig_dtype) + + +def _linear_bf16_act_uint4_weight_int8_zero_check(input_tensor, weight_tensor, bias): + return ( + # input is native bfloat16 tensor + not is_traceable_wrapper_subclass(input_tensor) + and input_tensor.dtype == torch.bfloat16 + and + # weight is uint4, group quantized tensor_core_tiled tensor impl affine quantized tensor + isinstance(weight_tensor, AffineQuantizedTensor) + and _aqt_is_xpu_layout_uint4(weight_tensor) + and weight_tensor.dtype == torch.bfloat16 + and len(weight_tensor.shape) == 2 + and weight_tensor.zero_point_domain == ZeroPointDomain.INT + and weight_tensor.tensor_impl.scale_and_zero is None + and weight_tensor.tensor_impl.scale.dtype == torch.bfloat16 + and weight_tensor.tensor_impl.zero.dtype == torch.int8 + and isinstance(weight_tensor._layout, Int4XPULayout) + ) + + +def _linear_bf16_act_uint4_weight_int8_zero_impl(input_tensor, weight_tensor, bias): + assert weight_tensor.block_size[0] == 1, ( + f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}" + ) + assert input_tensor.shape[-1] == weight_tensor.shape[1], ( + f"need input_tensor shape: {input_tensor.shape} final" + f"dim to match weight_tensor shape: {weight_tensor.shape} second dim " + ) + + # TODO: check groupsize quantization + # avoid circular dep, TODO: move this to a common util.py + act_mat = input_tensor + # weight is packed from padded (out_features, in_features) weight tensor + # (same dimension requirement as F.linear weight) + packed_weight = weight_tensor.tensor_impl.packed_weight + scale = weight_tensor.tensor_impl.scale + zero = weight_tensor.tensor_impl.zero + + orig_act_size = act_mat.size() + orig_dtype = act_mat.dtype + + act_mat = act_mat.reshape(-1, act_mat.shape[-1]).to(torch.bfloat16) + + # groupwise int4 quantization + groupsize = weight_tensor.block_size[1] + + y = torch.ops.aten._weight_int4pack_mm_with_scales_and_zeros( + act_mat, packed_weight, groupsize, scale, zero + ) + + # remove out_feature padding + orig_out_features = weight_tensor.shape[-2] + y = y[:, :orig_out_features] + y = y.reshape(*orig_act_size[:-1], orig_out_features) + + if bias is not None: + y += bias + return y.to(orig_dtype) + + +@dataclass(frozen=True) +class Int4XPULayout(Layout): + """Only for PyTorch version at least 2.7""" + + pass + + +@register_layout(Int4XPULayout) +class Int4XPUAQTTensorImpl(AQTTensorImpl): + """ + TensorImpl for int4 XPU layout for affine quantized tensor, this is for int4 only, + used by tinygemm kernels `_weight_int4pack_mm_xpu` and `_weight_int4pack_mm_with_zeros_and_scales` (TBD) + It stores the original tensor of dimension [n][k] (int32 dtype) as packed weight of 2-d tensor of + dimension: [n][k / 8] (int32 dtype) + (unpacked Tensor shape is n * k) + Note: we also pack scale and zero point together here for tinygemm kernel + Note: technically Int4 XPU layout should be the layout for the underlying packed weight + (int Tensor) but since the scale and zero_point are also packed into the same tensor here which is not used + in plain layout, we just created a layout for AQT right now, this could be improved if we split out + int4 aqt into a separate tensor subclass + fields: + packed_weight (torch.Tensor): the 2-d packed tensor in a Int4 XPU layout + [Optional] scale_and_zero (torch.Tensor): the combined scale Tensor used to map between floating point tensor to quantized tensor and zero_point Tensor + [Optional] scale (torch.Tensor): scale tensors, should be the same dtype of packed weight + [Optional] zeros (torch.Tensor): can be of the same dtype of packed weight or different dtype + """ + + def __new__( + cls, + packed_weight: torch.Tensor, + scale_and_zero: torch.Tensor, + transposed: bool, + _layout: Layout, + scale: torch.Tensor = None, + zero: torch.Tensor = None, + ): + kwargs = {} + kwargs["device"] = packed_weight.device + kwargs["layout"] = ( + kwargs.get("layout") + if kwargs.get("layout", False) + else packed_weight.layout + ) + kwargs["dtype"] = packed_weight.dtype + kwargs["requires_grad"] = False + shape = packed_weight.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + packed_weight: torch.Tensor, + scale_and_zero: torch.Tensor, + transposed: bool, + _layout: Layout, + scale: torch.Tensor = None, + zero: torch.Tensor = None, + ): + self.packed_weight = packed_weight + self.scale_and_zero = scale_and_zero + self.transposed = False + self._layout = _layout + self.scale = scale + self.zero = zero + + def __tensor_flatten__(self): + if self.scale_and_zero is not None: + return ["packed_weight", "scale_and_zero"], [self.transposed, self._layout] + else: + return ["packed_weight", "scale", "zero"], [self.transposed, self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + packed_weight = tensor_data_dict["packed_weight"] + scale_and_zero = ( + tensor_data_dict.get("scale_and_zero") + if "scale_and_zero" in tensor_data_dict + else None + ) + scale = tensor_data_dict.get("scale") if "scale" in tensor_data_dict else None + zero = tensor_data_dict.get("zero") if "zero" in tensor_data_dict else None + ( + transposed, + _layout, + ) = tensor_attributes + return cls(packed_weight, scale_and_zero, transposed, _layout, scale, zero) + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert isinstance(_layout, Int4XPULayout) + + from torchao.quantization.utils import convert_weight_to_int4pack_xpu + + if TORCH_VERSION_AT_LEAST_2_8: + assert int_data.dtype == torch.int32, ( + "torch.ops.aten._convert_weight_to_int4pack_for_cpu expects `int32` dtype" + ) + packed_weight = convert_weight_to_int4pack_xpu( + int_data, zero_point.dtype != scale.dtype + ) + else: + assert False, "INT4 not supported on XPU until 2.8" + + scale = scale.reshape(int_data.shape[0], -1) + zero_point = zero_point.reshape(int_data.shape[0], -1) + if zero_point.dtype == scale.dtype: + from torchao.quantization.utils import pack_tinygemm_scales_and_zeros + + scale_and_zero = pack_tinygemm_scales_and_zeros(scale, zero_point) + return cls(packed_weight, scale_and_zero, False, _layout, None, None) + else: + return cls( + packed_weight, + None, + False, + _layout, + scale.transpose(0, 1).contiguous(), + zero_point.transpose(0, 1).contiguous().to(torch.int8), + ) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs["device"] + if not is_device(torch.device(self.device).type, device): + raise ValueError( + f"Int4XPUAQTTensorImpl does not support conversion from {self.device} to {device}" + ) + return self.__class__( + self.packed_weight.to(device), + self.scale_and_zero.to(device) if self.scale_and_zero is not None else None, + self.transposed, + self._layout, + self.scale.to(device) if self.scale is not None else None, + self.zero.to(device) if self.zero is not None else None, + ) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.packed_weight), + fn(self.scale_and_zero) if self.scale_and_zero is not None else None, + self.transposed, + self._layout, + fn(self.scale) if self.scale is not None else None, + fn(self.zero) if self.zero is not None else None, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + if func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + if func is aten.t.default: + """we don't need to repack the weight and just rely on external + shape being changed and record the status of transpose/no-transpose + """ + transposed = Int4XPUAQTTensorImpl( + args[0].packed_weight, + args[0].scale_and_zero, + not args[0].transposed, + args[0]._layout, + args[0].scale, + args[0].zero, + ) + return return_and_correct_aliasing(func, args, kwargs, transposed) + + if func is torch.ops.aten.copy_.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + if func is aten.slice.Tensor: + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + if dim == 0: + int_data, scale, zero_point = self.get_plain() + int_data = aten.slice.Tensor(int_data, dim, start, end, step) + # this is to handle padding + int_data = self._layout.post_process(int_data) + sliced = self.from_plain(int_data, scale, zero_point, self._layout) + return return_and_correct_aliasing(func, args, kwargs, sliced) + elif dim == 1: + int_data, scale, zero_point = self.get_plain() + assert step == 1, "Only step == 1 is supported in slicing right now" + data_len = int_data.shape[dim] + scale_len = scale.shape[dim] + ratio = data_len / scale_len + start_scale = int(start / ratio) + end_scale = int(end / ratio) + + int_data = aten.slice.Tensor(int_data, dim, start, end, step) + # this is to handle padding + int_data = self._layout.post_process(int_data) + scale = aten.slice.Tensor(scale, dim, start_scale, end_scale, step) + zero_point = aten.slice.Tensor( + zero_point, dim, start_scale, end_scale, step + ) + sliced = self.from_plain(int_data, scale, zero_point, self._layout) + return sliced + else: + raise NotImplementedError( + f"Int4XPUAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported" + ) + + raise NotImplementedError( + f"Int4XPUAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + __torch_function__ = torch._C._disabled_torch_function_impl + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from torchao.quantization.quant_primitives import ( + ZeroPointDomain, + quantize_affine, + ) + from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros + + if self.scale_and_zero is not None: + scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero) + else: + scale = self.scale.transpose(0, 1).contiguous() + zero = self.zero.transpose(0, 1).contiguous() + + cur_shape = self.shape + assert len(cur_shape) == 2 + original_shape = (cur_shape[0], cur_shape[1] * 8) + eye_shape = original_shape[1] + groupsize = int(original_shape[1] / scale.shape[1]) + block_size = (1, groupsize) + device = self.device + original_dtype = torch.bfloat16 + target_dtype = torch.int32 + quant_min = 0 + quant_max = 15 + assert len(block_size) == 2 and block_size[0] == 1 + if self.scale_and_zero is None: + zero_point_domain = ZeroPointDomain.INT + dequantized = torch.ops.aten._weight_int4pack_mm_with_scales_and_zeros( + torch.eye(eye_shape, device=device, dtype=original_dtype), + self.packed_weight, + groupsize, + self.scale, + self.zero, + ) + dequantized = dequantized.t().contiguous() + int_data = quantize_affine( + dequantized, + block_size, + scale, + zero, + target_dtype, + quant_min, + quant_max, + zero_point_domain, + ) + else: + zero_point_domain = ZeroPointDomain.FLOAT + dequantized = torch.ops.aten._weight_int4pack_mm( + torch.eye(eye_shape, device=device, dtype=original_dtype), + self.packed_weight, + groupsize, + self.scale_and_zero, + ) + dequantized = dequantized.t().contiguous() + # TODO: move this to `unpack_tinygemm_scales_and_zeros`? + scale = scale.reshape(scale.shape[:-1]).contiguous() + zero = zero.reshape(zero.shape[:-1]).contiguous() + int_data = quantize_affine( + dequantized, + block_size, + scale, + zero, + target_dtype, + quant_min, + quant_max, + zero_point_domain, + ) + return int_data, scale, zero + + def get_layout(self) -> Layout: + return self._layout diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/marlin_qqq_tensor.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/marlin_qqq_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..3f3f4fa0756a492cc05385af145adc67e5651585 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/marlin_qqq_tensor.py @@ -0,0 +1,351 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import logging +import math +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from torch.utils._python_dispatch import ( + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + get_tensor_impl_constructor, + register_layout, +) +from torchao.dtypes.uintx.plain_layout import ( + _aqt_is_int8_reduced_range, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout +from torchao.quantization.quant_primitives import ( + ZeroPointDomain, + choose_qparams_and_quantize_affine_qqq, + dequantize_affine_qqq, +) + +logger = logging.getLogger(__name__) + +aten = torch.ops.aten + + +class MarlinQQQTensor(AffineQuantizedTensor): + """MarlinQQQ quantized tensor subclass which inherits AffineQuantizedTensor class. + + To see what happens during choose_qparams_and_quantize_affine_qqq, quantization and dequantization for marlin qqq quantization, + please checkout https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_primitives.py + and check the two quant primitive ops: choose_qparams_and_quantize_affine_qqq and dequantize_affine_qqq + """ + + def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if output_dtype is None: + output_dtype = self.dtype + + int_data, s_group, s_channel = self.tensor_impl.get_plain() + nbits = int(math.log2(self.quant_max - self.quant_min + 1)) + group_size = max(self.block_size) + return dequantize_affine_qqq( + int_data, s_group, s_channel, nbits, group_size, output_dtype + ) + + @classmethod + def from_hp_to_intx( + cls, + input_float: torch.Tensor, + block_size: Tuple[int, ...], + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + _layout: Optional[Layout] = None, + ): + """Converts a floating point tensor to a Marlin QQQ quantized tensor.""" + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + original_shape = input_float.shape + input_float = _layout.pre_process(input_float) + nbits = int(math.log2(quant_max - quant_min + 1)) + group_size = max(block_size) + data, s_group, s_channel, _ = choose_qparams_and_quantize_affine_qqq( + input_float, nbits, group_size + ) + tensor_impl_ctr = get_tensor_impl_constructor(type(_layout)) + tensor_impl = tensor_impl_ctr(data, s_group, s_channel, _layout) + return cls( + tensor_impl, + block_size, + original_shape, + quant_min, + quant_max, + zero_point_domain, + dtype=input_float.dtype, + ) + + +@dataclass(frozen=True) +class MarlinQQQLayout(Layout): + """MarlinQQQLayout is a layout class for Marlin QQQ quantization.""" + + pass + + +@register_layout(MarlinQQQLayout) +class MarlinQQQAQTTensorImpl(AQTTensorImpl): + """ + TensorImpl storage class for sparse_qqq layout for affine quantized tensor. + + Can only be used with 4 bits quantization for now. + + Original marlin documentation and information: + https://github.com/IST-DASLab/marlin/tree/master + + Marlin qqq information: + https://github.com/HandH1998/QQQ/tree/main + https://arxiv.org/pdf/2406.09904 + + fields: + original_shape (torch.Size): the original shape of the tensor. used to unpack the tensor to the original shape + group_size (int): the group size used to pack the tensor + num_bits (int): the number of bits used to quantize the tensor + """ + + @staticmethod + def __new__( + cls, + int_data: torch.Tensor, + s_group: torch.Tensor, + s_channel: torch.Tensor, + _layout: Layout, + original_shape: torch.Size, + group_size: int, + num_bits: int, + ): + kwargs = {} + kwargs["device"] = int_data.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout + ) + kwargs["dtype"] = int_data.dtype + kwargs["requires_grad"] = False + shape = int_data.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + int_data: torch.Tensor, + s_group: torch.Tensor, + s_channel: torch.Tensor, + _layout: Layout, + original_shape: torch.Size, + group_size: int, + num_bits: int, + ): + self.int_data = int_data + self.s_group = s_group + self.s_channel = s_channel + self._layout = _layout + self.original_shape = original_shape + self.group_size = group_size + self.num_bits = num_bits + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + raise NotImplementedError( + f"MarlinQQQAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + def __tensor_flatten__(self): + return ["int_data", "s_group", "s_channel"], [ + self._layout, + self.original_shape, + self.group_size, + self.num_bits, + ] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + int_data = tensor_data_dict["int_data"] + s_group = tensor_data_dict["s_group"] + s_channel = tensor_data_dict["s_channel"] + _layout, original_shape, group_size, num_bits = tensor_attributes + return cls( + int_data, s_group, s_channel, _layout, original_shape, group_size, num_bits + ) + + def get_plain(self): + from torchao.quantization.marlin_qqq import ( + unpack_from_marlin_qqq, + ) + + int_data_expanded, s_group_expanded, s_channel_expanded = ( + unpack_from_marlin_qqq( + self.int_data, + self.s_group, + self.s_channel, + self.original_shape, + self.num_bits, + self.group_size, + ) + ) + int_data_expanded_t = int_data_expanded.t() + s_group_expanded_t = s_group_expanded.t() + s_channel_expanded_t = s_channel_expanded.t() + return int_data_expanded_t, s_group_expanded_t, s_channel_expanded_t + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + s_group: torch.Tensor, + s_channel: torch.Tensor, + _layout: Layout, + ): + from torchao.quantization.marlin_qqq import ( + const, + pack_to_marlin_qqq, + ) + + assert isinstance(_layout, MarlinQQQLayout) + + # Linear layers are (in_features, out_features) but the int_data that is reaching this point + # is (out_features, in_features). We need to transpose it to match the expected shape in the marlin code. + q_w = int_data.t() + s_group_t = s_group.t() + s_channel_t = s_channel.t() + + if not torch.cuda.get_device_capability()[0] >= 8: + raise ValueError( + f"Can not use Marlin QQQ int4*int8 kernel with a device of compute capability {torch.cuda.get_device_capability()}, the minimum compute capability is 8.0 for Marlin kernel." + ) + + if q_w.dtype != torch.int32: + raise ValueError("Only `torch.int32` weights are supported.") + + in_features, out_features = q_w.shape + # (thread_k, thread_n) + thread_config = [(64, 256), (128, 128), (128, 64), (64, 128)] + if not any( + [ + in_features % thread_k == 0 and out_features % thread_n == 0 + for thread_k, thread_n in thread_config + ] + ): + raise ValueError( + "Not supported `in_features`: {} and `out_features`: {}.".format( + in_features, out_features + ) + ) + + num_bits = 4 if torch.max(q_w) - torch.min(q_w) < 16 else -1 + if num_bits not in [4]: + raise ValueError(f"Only {[4]} bits are supported, got {num_bits}.") + + if s_group.numel() == 0: + group_size = -1 + else: + group_size = in_features // s_group_t.shape[0] + assert group_size <= in_features, ( + "Group size must be less than or equal to in_features." + ) + + if group_size not in const.SUPPORTED_GROUP_SIZES: + raise ValueError( + f"Only {const.SUPPORTED_GROUP_SIZES} group sizes are supported, got {group_size}." + ) + + # Compress quantized weight to marlin format + marlin_qqq_q_w, marlin_qqq_s_group, marlin_qqq_s_channel = pack_to_marlin_qqq( + q_w, s_group_t, s_channel_t, num_bits, group_size + ) + + return cls( + marlin_qqq_q_w, + marlin_qqq_s_group, + marlin_qqq_s_channel, + _layout, + q_w.shape, + group_size, + num_bits, + ) + + def get_layout(self) -> Layout: + return self._layout + + def _apply_fn_to_data(self, fn): + self.int_data = fn(self.int_data) + self.s_group = fn(self.s_group) + self.s_channel = fn(self.s_channel) + return self + + +def _linear_int8_act_int4_weight_marlin_qqq_check(input_tensor, weight_tensor, bias): + return ( + isinstance(input_tensor, AffineQuantizedTensor) + and _aqt_is_int8_reduced_range(input_tensor) + and input_tensor.dtype == torch.float16 + and input_tensor.tensor_impl.scale.dtype == torch.float32 + and len(input_tensor.tensor_impl.scale.shape) == len(input_tensor.shape) - 1 + and isinstance(weight_tensor, AffineQuantizedTensor) + and weight_tensor.tensor_impl.dtype == torch.int32 + and len(weight_tensor.shape) == 2 + and isinstance(weight_tensor._layout, MarlinQQQLayout) + ) + + +def _linear_int8_act_int4_weight_marlin_qqq_impl(input_tensor, weight_tensor, bias): + from torchao.ops import marlin_qqq_gemm + from torchao.quantization.marlin_qqq import marlin_qqq_workspace + + assert isinstance(input_tensor, AffineQuantizedTensor) + assert isinstance(weight_tensor, AffineQuantizedTensor) + + input = input_tensor.tensor_impl.int_data + input_scale = input_tensor.tensor_impl.scale + + w_int4 = weight_tensor.tensor_impl.int_data + s_group = weight_tensor.tensor_impl.s_group + s_channel = weight_tensor.tensor_impl.s_channel + original_shape = weight_tensor.tensor_impl.original_shape + + # Folds batch dimension into the first dimension + input_2d = input.view(-1, input.shape[-1]) + input_scale = input_scale.view(1, -1) + + size_m = input_2d.shape[0] + size_n = s_channel.shape[1] + size_k = input_2d.shape[1] + workspace_qqq = marlin_qqq_workspace(original_shape[1]) + + out = marlin_qqq_gemm( + input_2d, + w_int4, + input_scale, + s_channel, + s_group, + workspace_qqq, + size_m, + size_n, + size_k, + ) + + # Unfold the batch dimension + out = out.reshape(input.shape[:-1] + (s_channel.shape[1],)) + + if bias is not None: + out += bias.to(out.dtype) + return out + + +to_marlinqqq_quantized_intx = MarlinQQQTensor.from_hp_to_intx diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/marlin_sparse_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/marlin_sparse_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..13af29d8bf56df661b99f19f0fd488232fba4c84 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/marlin_sparse_layout.py @@ -0,0 +1,310 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass + +import torch +from torch.utils._python_dispatch import ( + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.uintx.tensor_core_tiled_layout import _aqt_is_tensor_core_tile_uint4 +from torchao.dtypes.utils import AQTTensorImpl, Layout +from torchao.quantization.quant_primitives import ( + ZeroPointDomain, +) + +aten = torch.ops.aten + + +def _linear_fp_act_int4_weight_sparse_marlin_check(input_tensor, weight_tensor, bias): + return ( + isinstance(weight_tensor, AffineQuantizedTensor) + and _aqt_is_tensor_core_tile_uint4(weight_tensor) + and input_tensor.dtype == torch.float16 + and len(weight_tensor.shape) == 2 + and weight_tensor.zero_point_domain == ZeroPointDomain.INT + and isinstance(weight_tensor._layout, MarlinSparseLayout) + ) + + +def _linear_fp_act_int4_weight_sparse_marlin_impl(input_tensor, weight_tensor, bias): + from torchao.ops import marlin_24_gemm + from torchao.sparsity.marlin import marlin_24_workspace + + assert isinstance(weight_tensor, AffineQuantizedTensor) + + sparse_w_int4 = weight_tensor.tensor_impl.int_data + scale = weight_tensor.tensor_impl.scale + meta = weight_tensor.tensor_impl.meta + original_shape = weight_tensor.tensor_impl.original_shape + num_bits = weight_tensor.tensor_impl.num_bits + + # Folds batch dimension into the first dimension + input_2d = input_tensor.view(-1, input_tensor.shape[-1]) + + size_m = input_2d.shape[0] + size_n = scale.shape[1] + size_k = input_2d.shape[1] + workspace_24 = marlin_24_workspace(original_shape[1]) + + out = marlin_24_gemm( + input_2d, + sparse_w_int4, + meta, + scale, + workspace_24, + num_bits, + size_m, + size_n, + size_k, + ) + + # Unfold the batch dimension + out = out.reshape(input_tensor.shape[:-1] + (scale.shape[1],)) + + if bias is not None: + out += bias.to(out.dtype) + return out + + +@dataclass(frozen=True) +class MarlinSparseLayout(Layout): + """MarlinSparseLayout is a layout class for handling sparse tensor formats + specifically designed for the Marlin sparse kernel. This layout is used + to optimize the storage and computation of affine quantized tensors with + 2:4 sparsity patterns. + + The layout ensures that the tensor data is pre-processed and stored in a + format that is compatible with the Marlin sparse kernel operations. It + provides methods for preprocessing input tensors and managing the layout + of quantized tensors. + """ + + def pre_process(self, input: torch.Tensor) -> torch.Tensor: + """Preprocess the input tensor to be in the correct format for the Marlin sparse kernel. + - 1º: the input tensor is transposed since the linear layer keeps the weights in a transposed format + - 2º: tensor is injected with 2:4 sparsity + - 3º: transposes it again because the quantization process will compute the scales for dim=-1 + + Args: + input (torch.Tensor): the input tensor to preprocess + + Returns: + torch.Tensor: the preprocessed tensor + """ + from torchao.sparsity.marlin import inject_24 # avoid circular import + + input_t = input.t() + w_24, _ = inject_24(input_t, *input_t.shape) + return w_24.t() + + +@register_layout(MarlinSparseLayout) +class MarlinSparseAQTTensorImpl(AQTTensorImpl): + """ + TensorImpl for sparse_marlin_24 layout for affine quantized tensor. + + Can be used with 4 bits and 8 bits quantization. + + Original marlin documentation and information: + https://github.com/IST-DASLab/marlin/tree/master + + Sparse marlin documentation and information: + https://github.com/IST-DASLab/Sparse-Marlin?tab=readme-ov-file + + fields: + original_shape (torch.Size): the original shape of the tensor. used to unpack the tensor to the original shape + group_size (int): the group size used to pack the tensor + num_bits (int): the number of bits used to quantize the tensor + """ + + @staticmethod + def __new__( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + meta: torch.Tensor, + _layout: Layout, + original_shape: torch.Size, + group_size: int, + num_bits: int, + ): + kwargs = {} + kwargs["device"] = int_data.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout + ) + kwargs["dtype"] = int_data.dtype + kwargs["requires_grad"] = False + shape = int_data.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + meta: torch.Tensor, + _layout: Layout, + original_shape: torch.Size, + group_size: int, + num_bits: int, + ): + self.int_data = int_data + self.scale = scale + self.zero_point = zero_point + self.meta = meta + self._layout = _layout + self.original_shape = original_shape + self.group_size = group_size + self.num_bits = num_bits + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + raise NotImplementedError( + f"MarlinSparseAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + def __tensor_flatten__(self): + return ["int_data", "scale", "zero_point", "meta"], [ + self._layout, + self.original_shape, + self.group_size, + self.num_bits, + ] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + int_data = tensor_data_dict["int_data"] + scale = tensor_data_dict["scale"] + zero_point = tensor_data_dict["zero_point"] + meta = tensor_data_dict["meta"] + _layout, original_shape, group_size, num_bits = tensor_attributes + return cls( + int_data, + scale, + zero_point, + meta, + _layout, + original_shape, + group_size, + num_bits, + ) + + def get_plain(self): + from torchao.sparsity.marlin import ( + unpack_from_marlin_24, + ) + + int_data_expanded, scales_expanded = unpack_from_marlin_24( + self.int_data, + self.scale, + self.meta, + self.original_shape, + self.group_size, + self.num_bits, + ) + int_data_expanded_t = int_data_expanded.t() + scales_expanded_t = scales_expanded.t() + return int_data_expanded_t, scales_expanded_t, self.zero_point + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + _layout: Layout, + ): + from torchao.sparsity.marlin import ( + const, + pack_to_marlin_24, + ) + + assert isinstance(_layout, MarlinSparseLayout) + + # Linear layers are (in_features, out_features) but the int_data that is reaching this point + # is (out_features, in_features). We need to transpose it to match the expected shape in the marlin code. + q_w_24 = int_data.t() + # addressing the case when scale has dimension 1, happens when + # weight_shape[-1] == group_size == 128 + if scale.ndim == 1: + scale = scale.reshape(scale.shape[0], -1) + + scale_t = scale.t() + + if not torch.cuda.get_device_capability()[0] >= 8: + raise ValueError( + f"Can not use Sparse Marlin 2:4 int4*fp16 kernel with a device of compute capability {torch.cuda.get_device_capability()}, the minimum compute capability is 8.0 for Marlin kernel." + ) + + if q_w_24.dtype != torch.int32: + raise ValueError("Only `torch.int32` weights are supported.") + + in_features, out_features = q_w_24.shape + if in_features % 128 != 0 or out_features != 256 == 0: + raise ValueError( + "`in_features` must be divisible by 64 and `out_features` by 256." + ) + + # NOTE: The current marlin 2:4 kernel supports both 4 and 8 bits quantization but fp8 + # will require a bit more work to get our current quantization flow to work with it. + # Check the link for a reference: https://github.com/neuralmagic/nm-vllm/tree/main + num_bits = 4 if torch.max(q_w_24) < 16 else -1 + if num_bits not in [4]: + raise ValueError(f"Only {[4]} bits are supported, got {num_bits}.") + + group_size = in_features // scale_t.shape[0] + if group_size == 0: + group_size = in_features + assert group_size <= in_features, ( + "Group size must be less than or equal to in_features." + ) + + if group_size not in const.SUPPORTED_GROUP_SIZES: + raise ValueError( + f"Only {const.SUPPORTED_GROUP_SIZES} group sizes are supported, got {group_size}." + ) + + # Compress quantized weight to marlin 2:4 format + marlin_24_q_w_comp, marlin_24_s, meta = pack_to_marlin_24( + q_w_24, scale_t, num_bits, group_size + ) + + return cls( + marlin_24_q_w_comp, + marlin_24_s, + zero_point, + meta, + _layout, + q_w_24.shape, + group_size, + num_bits, + ) + + def get_layout(self) -> Layout: + return self._layout + + def _apply_fn_to_data(self, fn): + self.int_data = fn(self.int_data) + self.scale = fn(self.scale) + self.zero_point = fn(self.zero_point) + self.meta = fn(self.meta) + return self diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/packed_linear_int8_dynamic_activation_intx_weight_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/packed_linear_int8_dynamic_activation_intx_weight_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..dc7b073f32111672c786d4b573ce2fc00e584c18 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/packed_linear_int8_dynamic_activation_intx_weight_layout.py @@ -0,0 +1,467 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from enum import Enum, auto +from typing import Optional, Tuple, Union + +import torch +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.dtypes.affine_quantized_tensor import register_layout +from torchao.dtypes.utils import AQTTensorImpl, Layout +from torchao.experimental.op_lib_utils import _check_torchao_ops_loaded +from torchao.quantization.quant_primitives import ( + _DTYPE_TO_BIT_WIDTH, + _DTYPE_TO_QVALUE_BOUNDS, + ZeroPointDomain, +) +from torchao.utils import TORCH_VERSION_AT_LEAST_2_6 + +logger = logging.getLogger(__name__) +logger.setLevel(logging.WARNING) + +import sys + +handler = logging.StreamHandler(sys.stdout) +formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") +handler.setFormatter(formatter) +logger.addHandler(handler) + + +class Target(Enum): + """Enum that indicates the backend target""" + + # AUTO target will automatically select a packing format + # based on the available hardware. + AUTO = auto() + UNIVERSAL = auto() + KLEIDIAI = auto() + + # ATEN target will use the ATen operator + ATEN = auto() + + +_TARGET_AND_STR = [ + (Target.AUTO, "auto"), + (Target.ATEN, "aten"), + (Target.UNIVERSAL, "universal"), + (Target.KLEIDIAI, "kleidiai"), +] + + +def target_to_str(target: Target) -> str: + target_to_str = {t: s for t, s in _TARGET_AND_STR} + return target_to_str[target] + + +def target_from_str(target: str) -> Target: + str_to_target = {s: t for t, s in _TARGET_AND_STR} + if target.lower() in str_to_target: + return str_to_target[target.lower()] + raise ValueError(f"Invalid target: {target}") + + +class PackedLinearInt8DynamicActivationIntxWeightLayout(Layout): + bit_width: Optional[int] + group_size: Optional[int] + has_weight_zeros: Optional[bool] + has_bias: Optional[bool] + target: Optional[Target] + + def __init__( + self, + target: Union[str, Target] = "auto", + ): + if isinstance(target, str): + target = target_from_str(target) + self.target = target + + self.bit_width: Optional[int] = None + self.group_size: Optional[int] = None + self.has_weight_zeros: Optional[bool] = None + # has_bias is whether the packed weights + # have bias packed with them, not whether the + # linear operator has bias + self.has_bias: Optional[bool] = None + + def extra_repr(self): + return f"group_size={self.group_size}, bit_width={self.bit_width}, has_weight_zeros={self.has_weight_zeros}, has_bias={self.has_bias}, target={self.target}" + + def has_params_set(self) -> bool: + return ( + (self.bit_width is not None) + and (self.group_size is not None) + and (self.has_weight_zeros is not None) + and (self.has_bias is not None) + and (self.target is not None) + ) + + def set_params( + self, bit_width: int, group_size: int, has_weight_zeros: bool, has_bias: bool + ): + assert bit_width >= 1 and bit_width <= 8, "bit_width must be 1 to 8" + assert group_size >= 1, f"group_size must be positive, got {group_size}" + + self.bit_width = bit_width + self.group_size = group_size + self.has_weight_zeros = has_weight_zeros + self.has_bias = has_bias + assert self.has_params_set() + + +@register_layout(PackedLinearInt8DynamicActivationIntxWeightLayout) +class PackedLinearInt8DynamicActivationIntxWeightAQTTensorImpl(AQTTensorImpl): + def __new__( + cls, + packed_weight: torch.Tensor, + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = packed_weight.device + kwargs["dtype"] = packed_weight.dtype + assert not packed_weight.requires_grad + kwargs["requires_grad"] = False + shape = packed_weight.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + packed_weight: torch.Tensor, + _layout: Layout, + ): + assert isinstance(_layout, PackedLinearInt8DynamicActivationIntxWeightLayout) + self.packed_weight = packed_weight + self._layout = _layout + + def __repr__(self): + return f"{self.__class__.__name__}(packed_weight={str(self.packed_weight)}, layout={self.get_layout()})" + + def get_layout(self) -> Layout: + return self._layout + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + raise NotImplementedError( + "get_plain is not implemented for PackedLinearInt8DynamicActivationIntxWeightAQTTensorImpl" + ) + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + layout: Layout, + bias: Optional[torch.Tensor] = None, + *, + validate_inputs: bool = True, + ): + assert isinstance(layout, PackedLinearInt8DynamicActivationIntxWeightLayout) + assert layout.target in [t for t, _ in _TARGET_AND_STR], ( + f"Unexpected target: {layout.target}" + ) + assert layout.has_params_set(), ( + "PackedLinearInt8DynamicActivationIntxWeightLayout params must be set before calling from_plain" + ) + + if layout.target != Target.ATEN: + _check_torchao_ops_loaded() + else: + assert TORCH_VERSION_AT_LEAST_2_6, ( + "aten target is requires torch version > 2.6.0" + ) + assert torch.backends.kleidiai.is_available(), ( + "ATEN target requires torch.backends.kleidiai.is_available()" + ) + layout.bit_width == 4, "ATEN target only supports torch.int4" + assert not layout.has_weight_zeros, "ATEN target does not support zeros" + + data_dtype = getattr(torch, f"int{layout.bit_width}") + qmin, qmax = _DTYPE_TO_QVALUE_BOUNDS[data_dtype] + + int_types = [torch.int8, torch.int16, torch.int32, torch.int64] + + # Check int_data + assert int_data.device == torch.device("cpu") + assert int_data.dtype in int_types + n, k = int_data.shape + assert k % layout.group_size == 0, "k must be divisible by group_size" + if validate_inputs: + assert int_data.min().item() >= qmin + assert int_data.max().item() <= qmax + int_data = int_data.to(torch.int8) + + # Check scale + assert scale.device == torch.device("cpu") + if scale.dtype != torch.float32: + logging.info(f"scale has dtype {scale.dtype}, converting to torch.float32") + scale = scale.to(torch.float32) + n_, _ = scale.shape + assert n_ == n + assert scale.numel() * layout.group_size == int_data.numel(), ( + "must have 1 scale per group" + ) + if validate_inputs: + assert scale.min().item() > 0 + # Some targets round scales to bfloat16, give warning if scales are at higher precision + scale_is_rounded_to_bf16 = torch.allclose( + scale, scale.to(torch.bfloat16).to(torch.float32) + ) + if not scale_is_rounded_to_bf16: + if layout.target == Target.ATEN and (layout.group_size < k): + logging.warning( + "When using Target.ATEN with group_size < k, scales will be rounded to bfloat16" + ) + if layout.target in [Target.AUTO, Target.KLEIDIAI]: + logging.warning( + "When using [Target.AUTO, Target.KLEIDIAI], scales will be rounded to bfloat16" + ) + + # Check zero_point + if zero_point is None: + assert not layout.has_weight_zeros, ( + "zero_point must be provided if has_weight_zeros=True" + ) + else: + assert zero_point.device == torch.device("cpu") + assert zero_point.shape == scale.shape + assert zero_point.dtype in int_types + assert zero_point.numel() * layout.group_size == int_data.numel(), ( + "must have 1 zero_point per group" + ) + if validate_inputs: + zero_point_min = zero_point.min().item() + zero_point_max = zero_point.max().item() + assert zero_point.min().item() >= qmin + assert zero_point.max().item() <= qmax + has_weight_zeros = True + if zero_point_min == 0 and zero_point_max == 0: + has_weight_zeros = False + assert has_weight_zeros == layout.has_weight_zeros, ( + "zero_point being all zeros must be consistent with layout.has_weight_zeros" + ) + zero_point = zero_point.to(torch.int8) + + # Check bias + has_bias = bias is not None + assert has_bias == layout.has_bias, ( + "bias being None must be consistent with layout.has_bias" + ) + if has_bias: + assert bias.device == torch.device("cpu") + if bias.dtype != torch.float32: + logging.info( + f"bias has dtype {bias.dtype}, converting to torch.float32" + ) + bias = bias.to(torch.float32) + assert bias.shape == (n,) + + # Construct packed_weight + if layout.target == Target.ATEN: + int_data = int_data.add(8) + int_data = (int_data[::, 1::2] << 4 | int_data[::, ::2]).to(torch.uint8) + + # If group_size < k, convert scales to bfloat16 + # to call optimized kernel + if layout.group_size < k: + scale = scale.to(torch.bfloat16) + packed_weight = torch.ops.aten._dyn_quant_pack_4bit_weight( + int_data, scale, bias, layout.group_size, k, n + ) + return cls(packed_weight, layout) + + args = [ + int_data, + scale.reshape(-1), + zero_point.reshape(-1) if layout.has_weight_zeros else None, + layout.group_size, + bias, + target_to_str(layout.target) if layout.target != Target.AUTO else None, + ] + packed_weight = getattr( + torch.ops.torchao, + f"_pack_8bit_act_{layout.bit_width}bit_weight", + )(*args) + + return cls(packed_weight, layout) + + def _apply_fn_to_data(self, fn): + self.packed_weight = fn(self.packed_weight) + return self + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is torch.ops.aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + if func is torch.ops.aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + raise NotImplementedError( + f"PackedLinearInt8DynamicActivationIntxWeightAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + def __tensor_flatten__(self): + return ["packed_weight"], [self.get_layout()] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + packed_weight = tensor_data_dict["packed_weight"] + (layout,) = tensor_attributes + return cls(packed_weight, layout) + + +def _linear_check(input_tensor, weight_tensor, bias): + layout = weight_tensor.tensor_impl.get_layout() + return isinstance(layout, PackedLinearInt8DynamicActivationIntxWeightLayout) + + +def _linear_impl(input_tensor, weight_tensor, bias): + def _impl_2d_non_aten(input_tensor, weight_tensor): + assert input_tensor.dim() == 2 + assert weight_tensor.dim() == 2 + + m, k = input_tensor.shape + n, k_ = weight_tensor.shape + assert k_ == k + group_size = weight_tensor.tensor_impl.get_layout().group_size + + args = ( + input_tensor, + weight_tensor.tensor_impl.packed_weight, + group_size, + n, + k, + ) + + assert len(weight_tensor.block_size) == 2 + assert weight_tensor.block_size[0] == 1 + assert group_size == weight_tensor.block_size[1] + bit_width = weight_tensor.tensor_impl.get_layout().bit_width + + return getattr(torch.ops.torchao, f"_linear_8bit_act_{bit_width}bit_weight")( + *args + ) + + def _impl_2d_aten(input_tensor, weight_tensor): + assert input_tensor.dim() == 2 + assert weight_tensor.dim() == 2 + + m, k = input_tensor.shape + n, k_ = weight_tensor.shape + + assert k_ == k + group_size = weight_tensor.tensor_impl.get_layout().group_size + packed_weight = weight_tensor.tensor_impl.packed_weight + return torch.ops.aten._dyn_quant_matmul_4bit( + input_tensor, packed_weight, group_size, k, n + ) + + if input_tensor.numel() == 0: + return input_tensor + + target = weight_tensor.tensor_impl.get_layout().target + + if weight_tensor.tensor_impl.get_layout().has_bias: + assert bias is None, ( + "bias should be None because it is already packed with the weights (has_bias=True)" + ) + + if target == Target.ATEN: + assert TORCH_VERSION_AT_LEAST_2_6 == 1, "Target.ATEN requires torch >= 2.6.0" + _impl_2d = _impl_2d_aten + else: + _impl_2d = _impl_2d_non_aten + + if input_tensor.dim() == 2: + res = _impl_2d(input_tensor, weight_tensor) + else: + assert input_tensor.dim() >= 3 + lead_shape = input_tensor.shape[0:-2] + m, k = input_tensor.shape[-2], input_tensor.shape[-1] + n, k_ = weight_tensor.shape + assert k_ == k + + res = _impl_2d(input_tensor.reshape(-1, k), weight_tensor) + res = res.reshape(*lead_shape, m, n) + + if bias is not None: + res = res + bias + return res + + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout + + +def make_packed_linear_int8_dynamic_activation_intx_weight_tensor( + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + data_dtype: torch.dtype, + target: Union[str, Target] = "auto", + *, + validate_inputs: bool = True, +) -> AffineQuantizedTensor: + """ + Constructs an AffineQuantizedTensor with PackedLinearInt8DynamicActivationIntxWeightLayout + from plain data. + """ + # TORCH_VERSION_AT_LEAST_2_6 is needed for torch.intx with x < 8 + assert TORCH_VERSION_AT_LEAST_2_6, ( + "Using PackedLinearInt8DynamicActivationIntxWeightLayout requires torch version > 2.6.0" + ) + + layout = PackedLinearInt8DynamicActivationIntxWeightLayout(target=target) + + bit_width = _DTYPE_TO_BIT_WIDTH[data_dtype] + qmin, qmax = _DTYPE_TO_QVALUE_BOUNDS[data_dtype] + + n, k = int_data.shape + n_, groups_per_k = scale.shape + assert k % groups_per_k == 0 + group_size = k // groups_per_k + + has_weight_zeros = True + if zero_point is None: + has_weight_zeros = False + else: + zero_point_min = zero_point.min().item() + zero_point_max = zero_point.max().item() + if zero_point_min == 0 and zero_point_max == 0: + has_weight_zeros = False + + has_bias = bias is not None + + layout.set_params(bit_width, group_size, has_weight_zeros, has_bias) + assert layout.has_params_set() + tensor_impl = PackedLinearInt8DynamicActivationIntxWeightAQTTensorImpl.from_plain( + int_data, + scale, + zero_point, + layout, + bias, + validate_inputs=validate_inputs, + ) + + return AffineQuantizedTensor( + tensor_impl, + block_size=(1, group_size), + shape=int_data.shape, + quant_min=qmin, + quant_max=qmax, + zero_point_domain=ZeroPointDomain.INT, + ) diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/plain_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/plain_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..3551214d7ee113dac9e5117f967420c9b2212446 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/plain_layout.py @@ -0,0 +1,315 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Optional, Tuple + +import torch +from torch.utils._python_dispatch import ( + is_traceable_wrapper_subclass, + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout, PlainLayout +from torchao.kernel import ( + int_scaled_matmul, +) +from torchao.quantization.quant_primitives import ( + ZeroPointDomain, +) +from torchao.utils import fill_defaults + +aten = torch.ops.aten + + +def _same_metadata(self: "PlainAQTTensorImpl", src: "PlainAQTTensorImpl") -> bool: + return ( + isinstance(self, PlainAQTTensorImpl) + and isinstance(src, PlainAQTTensorImpl) + and self.shape == src.shape + and self.int_data.shape == src.int_data.shape + and self.scale.shape == src.scale.shape + and (self.zero_point is None and src.zero_point is None) + or ( + self.zero_point is not None + and src.zero_point is not None + and self.zero_point.shape == src.zero_point.shape + ) + and type(self._layout) == type(src._layout) + ) + + +@register_layout(PlainLayout) +class PlainAQTTensorImpl(AQTTensorImpl): + """ + TensorImpl for plain layout for affine quantized tensor, it stores int_data, scale, zero_point + tensors directly as plain tensors. + + fields: + int_data (torch.Tensor): the quantized integer data Tensor + scale (torch.Tensor): the scale Tensor used to map between floating point tensor to quantized tensor + zero_point (torch.Tensor): the zero_point Tensor used to map between floating point tensor to quantized tensor + """ + + def __new__( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = int_data.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout + ) + kwargs["dtype"] = int_data.dtype + kwargs["requires_grad"] = False + shape = int_data.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + self.int_data = int_data + self.scale = scale + self.zero_point = zero_point + self._layout = _layout + + def __tensor_flatten__(self): + if self.zero_point is None: + return ["int_data", "scale"], [self._layout] + return ["int_data", "scale", "zero_point"], [self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + int_data, scale, zero_point = ( + tensor_data_dict["int_data"], + tensor_data_dict["scale"], + tensor_data_dict.get("zero_point", None), + ) + (_layout,) = tensor_attributes + return cls(int_data, scale, zero_point, _layout) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + return self.__class__( + self.int_data.to(kwargs["device"]), + self.scale.to(kwargs["device"]), + self.zero_point.to(kwargs["device"]) + if self.zero_point is not None + else None, + self._layout, + ) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.int_data), + fn(self.scale), + fn(self.zero_point) if self.zero_point is not None else None, + self._layout, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + elif func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + elif func is aten.copy_.default: + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + + elif func is aten.t.default: + tensor = args[0] + new = tensor.__class__( + tensor.int_data.t(), tensor.scale, tensor.zero_point, tensor._layout + ) + return return_and_correct_aliasing(func, args, kwargs, new) + + elif func in [aten.select.int, aten.index.Tensor]: + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data(lambda x: func(x, *args[1:], **kwargs)), + ) + + elif func is aten.slice.Tensor: + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + if dim == 0: + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data( + lambda x: aten.slice.Tensor(x, dim, start, end, step) + ), + ) + elif dim == 1: + assert len(self.scale.shape) == 1, ( + f"slice dim==1 only works when len(scale.shape) == 1 currently, got: {self.scale.shape}" + ) + return PlainAQTTensorImpl( + aten.slice.Tensor(self.int_data, dim, start, end, step), + self.scale.view(-1), + self.zero_point.view(-1) if self.zero_point is not None else None, + self._layout, + ) + else: + raise NotImplementedError( + f"PlainAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported" + ) + + raise NotImplementedError( + f"PlainAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + __torch_function__ = torch._C._disabled_torch_function_impl + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + return self.int_data, self.scale, self.zero_point + + def get_layout(self) -> Layout: + return self._layout + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert isinstance(_layout, PlainLayout) + return cls(int_data, scale, zero_point, _layout) + + +def _aqt_is_int8(aqt): + """Check if an AffineQuantizedTensor is int8 quantized Tensor""" + return ( + aqt.tensor_impl.dtype == torch.int8 + and (aqt.quant_min is None or aqt.quant_min == -128) + and (aqt.quant_max is None or aqt.quant_max == 127) + ) + + +def _aqt_is_int8_reduced_range(aqt): + return ( + aqt.tensor_impl.dtype == torch.int8 + and aqt.quant_min == -127 + and (aqt.quant_max is None or aqt.quant_max == 127) + ) + + +def _linear_fp_act_int8_weight_check(input_tensor, weight_tensor, bias): + return ( + # input is native float tensor + not is_traceable_wrapper_subclass(input_tensor) + and input_tensor.is_floating_point() + and + # weight is int8 per channel quantized affine quantized tensor + isinstance(weight_tensor, AffineQuantizedTensor) + and _aqt_is_int8(weight_tensor) + and len(weight_tensor.shape) == 2 + and len(weight_tensor.block_size) == 2 + and weight_tensor.block_size[0] == 1 + and weight_tensor.block_size[1] == weight_tensor.shape[1] + and weight_tensor.zero_point_domain == ZeroPointDomain.INT + and isinstance(weight_tensor._layout, PlainLayout) + ) + + +def _linear_fp_act_int8_weight_impl(input_tensor, weight_tensor, bias): + # TODO: enable cpu and mps efficient path + # is_cpu and is_mps only, some issue with is_contiguous() currently + # return torch.ops.aten._weight_int8pack_mm(input_tensor.contiguous(), w_vals_int8_t, weight_tensor.tensor_impl.scale) + + # per channel int8 weight only quantizated mm + w_vals_int8_t = weight_tensor.tensor_impl.int_data.t() + scale = weight_tensor.tensor_impl.scale + m = torch.mm( + input_tensor.reshape(-1, input_tensor.shape[-1]), + w_vals_int8_t.to(input_tensor.dtype), + ) + y = m * scale.to(m.dtype) + y = y.reshape(*input_tensor.shape[:-1], y.shape[-1]) + if bias is not None: + y += bias.to(m.dtype) + return y + + +def _linear_int8_act_int8_weight_check(input_tensor, weight_tensor, bias): + return ( + isinstance(input_tensor, AffineQuantizedTensor) + and _aqt_is_int8_reduced_range(input_tensor) + and isinstance(weight_tensor, AffineQuantizedTensor) + and _aqt_is_int8(weight_tensor) + and input_tensor.dtype == weight_tensor.dtype + and isinstance(input_tensor._layout, PlainLayout) + and isinstance(weight_tensor._layout, PlainLayout) + ) + + +def _linear_int8_act_int8_weight_impl(input_tensor, weight_tensor, bias): + # + # 1. do the matrix form of dot(X_i, W_j) + # + # + # 2. rescale the output + # + # in cases with large matrices, y_dot_int32 can grow sufficiently + # large that y_dot_int32 * a float16 scale is greater than the maximum + # value of a float 16, (which results in a value of inf even if multiplying + # by the other scale would bring it within the expected range) + + x_vals_int8 = input_tensor.tensor_impl.int_data + x_scales = input_tensor.tensor_impl.scale + w_vals_int8_t = weight_tensor.tensor_impl.int_data.contiguous().t() + w_scales = weight_tensor.tensor_impl.scale + tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1]) + x_scales_dtype = x_scales.dtype + # Cast fp16 scale to float to avoid overflow in int_scaled_matmul + intermediate_dtype = torch.float if x_scales_dtype == torch.half else x_scales_dtype + y_dot_scaled = int_scaled_matmul( + tmp, w_vals_int8_t, x_scales.reshape(-1, 1).to(intermediate_dtype) + ) + y_dot_scaled = y_dot_scaled.to(x_scales_dtype) + + y = (y_dot_scaled * w_scales).reshape( + *x_vals_int8.shape[:-1], y_dot_scaled.shape[-1] + ) + + # can downcast only at the very end + output_dtype = input_tensor.dtype + y = y.to(output_dtype) + if bias is not None: + y += bias + return y diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/q_dq_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/q_dq_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..0ae1d865e8c0f6128fcc9e8449878faa3c8bc3c2 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/q_dq_layout.py @@ -0,0 +1,242 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import logging + +import torch + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) + +logger = logging.getLogger(__name__) +logger.setLevel(logging.WARNING) + +import sys + +handler = logging.StreamHandler(sys.stdout) +formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") +handler.setFormatter(formatter) +logger.addHandler(handler) + + +from dataclasses import dataclass +from typing import Optional, Tuple + +from torch.utils._python_dispatch import ( + return_and_correct_aliasing, +) + +from torchao.dtypes.utils import AQTTensorImpl, Layout +from torchao.utils import fill_defaults + +aten = torch.ops.aten + + +@dataclass(frozen=True) +class QDQLayout(Layout): + pass + + +def _same_metadata(self: "QDQTensorImpl", src: "QDQTensorImpl") -> bool: + return ( + isinstance(self, QDQTensorImpl) + and isinstance(src, QDQTensorImpl) + and self.shape == src.shape + and self.int_data.shape == src.int_data.shape + and self.scale.shape == src.scale.shape + and (self.zero_point is None and src.zero_point is None) + or ( + self.zero_point is not None + and src.zero_point is not None + and self.zero_point.shape == src.zero_point.shape + ) + and type(self._layout) == type(src._layout) + ) + + +@register_layout(QDQLayout) +class QDQTensorImpl(AQTTensorImpl): + """ + TensorImpl for QDQLayout layout for affine quantized tensor, it stores int_data, scale, zero_point + tensors directly as plain tensors. + + fields: + int_data (torch.Tensor): the quantized integer data Tensor + scale (torch.Tensor): the scale Tensor used to map between floating point tensor to quantized tensor + zero_point (torch.Tensor): the zero_point Tensor used to map between floating point tensor to quantized tensor + """ + + def __new__( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = int_data.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout + ) + kwargs["dtype"] = int_data.dtype + kwargs["requires_grad"] = False + shape = int_data.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + self.int_data = int_data + self.scale = scale + self.zero_point = zero_point + self._layout = _layout + + def __tensor_flatten__(self): + if self.zero_point is None: + return ["int_data", "scale"], [self._layout] + return ["int_data", "scale", "zero_point"], [self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + int_data, scale, zero_point = ( + tensor_data_dict["int_data"], + tensor_data_dict["scale"], + tensor_data_dict.get("zero_point", None), + ) + (_layout,) = tensor_attributes + return cls(int_data, scale, zero_point, _layout) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + return self.__class__( + self.int_data.to(kwargs["device"]), + self.scale.to(kwargs["device"]), + self.zero_point.to(kwargs["device"]) + if self.zero_point is not None + else None, + self._layout, + ) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.int_data), + fn(self.scale), + fn(self.zero_point) if self.zero_point is not None else None, + self._layout, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + elif func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + elif func is aten.copy_.default: + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + + elif func is aten.t.default: + tensor = args[0] + new = tensor.__class__( + tensor.int_data.t(), tensor.scale, tensor.zero_point, tensor._layout + ) + return return_and_correct_aliasing(func, args, kwargs, new) + + elif func is aten.slice.Tensor: + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + if dim in [0, 1]: + int_data, scale, zero_point = self.get_plain() + data_len = int_data.shape[dim] + scale_len = scale.shape[dim] + ratio = data_len / scale_len + start_scale = int(start / ratio) + end_scale = int(end / ratio) + + int_data = aten.slice.Tensor(int_data, dim, start, end, step) + scale = aten.slice.Tensor(scale, dim, start_scale, end_scale, step) + if zero_point is not None: + zero_point = aten.slice.Tensor( + zero_point, dim, start_scale, end_scale, step + ) + sliced = self.from_plain(int_data, scale, zero_point, self._layout) + return return_and_correct_aliasing(func, args, kwargs, sliced) + else: + raise NotImplementedError( + f"QDQTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported" + ) + + raise NotImplementedError( + f"QDQTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + __torch_function__ = torch._C._disabled_torch_function_impl + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + return self.int_data, self.scale, self.zero_point + + def get_layout(self) -> Layout: + return self._layout + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert isinstance(_layout, QDQLayout) + return cls(int_data, scale, zero_point, _layout) + + +def _linear_check(input_tensor, weight_tensor, bias): + layout = weight_tensor.tensor_impl.get_layout() + return isinstance(layout, QDQLayout) + + +def _linear_impl(input_tensor, weight_tensor, bias): + if isinstance(input_tensor, AffineQuantizedTensor): + input_tensor = input_tensor.dequantize() + if isinstance(weight_tensor, AffineQuantizedTensor): + weight_tensor = weight_tensor.dequantize() + return torch.nn.functional.linear(input_tensor, weight_tensor, bias) + + +def _embedding_check(args, kwargs): + _, weight_tensor = args + layout = weight_tensor.tensor_impl.get_layout() + return isinstance(layout, QDQLayout) + + +def _embedding_impl(args, kwargs): + input_tensor, weight_tensor = args + if isinstance(weight_tensor, AffineQuantizedTensor): + weight_tensor = weight_tensor.dequantize() + return torch.nn.functional.embedding(input_tensor, weight_tensor, **kwargs) diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/semi_sparse_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/semi_sparse_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..30bf4dacbd6464ce8b67b0f2e14d21f4895280f8 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/semi_sparse_layout.py @@ -0,0 +1,128 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Optional + +import torch +from torch.utils._python_dispatch import ( + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.uintx.plain_layout import ( + PlainAQTTensorImpl, + _aqt_is_int8_reduced_range, +) +from torchao.dtypes.utils import Layout, PlainLayout + +aten = torch.ops.aten + + +def _linear_int8_act_int8_weight_semi_structured_sparse_check( + input_tensor, weight_tensor, bias +): + return ( + isinstance(input_tensor, AffineQuantizedTensor) + and _aqt_is_int8_reduced_range(input_tensor) + and isinstance(weight_tensor, AffineQuantizedTensor) + and weight_tensor.is_cuda + and input_tensor.dtype == weight_tensor.dtype + and isinstance(input_tensor._layout, PlainLayout) + and isinstance(weight_tensor._layout, SemiSparseLayout) + ) + + +def _linear_int8_act_int8_weight_semi_structured_sparse_impl( + input_tensor, weight_tensor, bias +): + x_vals_int8 = input_tensor.tensor_impl.int_data + x_scales = input_tensor.tensor_impl.scale + w_vals_int8 = weight_tensor.tensor_impl.int_data + w_scales = weight_tensor.tensor_impl.scale + tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1]) + # must pad + row, col = tmp.shape + from torch.sparse import SparseSemiStructuredTensorCUSPARSELT + + tmp_padded = SparseSemiStructuredTensorCUSPARSELT._pad_dense_input(tmp) + # we fuse one of the scalar matrix multiplications (w_scales) into the sparse mm + y_dot_bf16_w_scales_fused = torch._cslt_sparse_mm( + w_vals_int8, + tmp_padded.t(), + alpha=w_scales.to(torch.float32), + out_dtype=torch.bfloat16, + ).t()[:row, :] + y = (y_dot_bf16_w_scales_fused * x_scales.reshape(-1, 1)).reshape( + *x_vals_int8.shape[:-1], y_dot_bf16_w_scales_fused.shape[-1] + ) + output_dtype = input_tensor.dtype + # TODO: waiting for jesse's test/fix + y = y.to(output_dtype).contiguous() + if bias is not None: + y += bias + return y + + +@dataclass(frozen=True) +class SemiSparseLayout(Layout): + """SemiSparseLayout is a layout class for handling semi-structured sparse + matrices in affine quantized tensors. This layout is specifically designed + to work with the 2:4 sparsity pattern, where two out of every four elements + are pruned to zero. This class provides methods for preprocessing input + tensors to conform to this sparsity pattern. + """ + + def pre_process(self, input: torch.Tensor) -> torch.Tensor: + # prune to 2:4 if not already + temp = input.detach() + pruning_inds = temp.abs().view(-1, 4).argsort(dim=1)[:, :2] + temp.view(-1, 4).scatter_(1, pruning_inds, value=0) + return temp + + +@register_layout(SemiSparseLayout) +class SemiSparseAQTTensorImpl(PlainAQTTensorImpl): + """ + TensorImpl for semi_sparse_cusparselt layout for affine quantized tensor + """ + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + raise NotImplementedError( + f"SparseAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + def get_plain(self): + # Currently we don't have cuSPARSELt expansion routines, so we matmul by + # the identity matrix to get the original dense matrix. This is slow though. + cols = self.int_data.numel() * 16 // (10 * self.scale.shape[0]) + int_data_expanded = torch._cslt_sparse_mm( + self.int_data, + torch.eye(cols, dtype=self.int_data.dtype, device=self.int_data.device).t(), + ) + return int_data_expanded, self.scale, self.zero_point + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert isinstance(_layout, SemiSparseLayout) + int_data_compressed = torch._cslt_compress(int_data) + return cls(int_data_compressed, scale, zero_point, _layout) diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/tensor_core_tiled_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/tensor_core_tiled_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..3bf9ef6b722d15b91eaa89975d44737b7a9d3462 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/tensor_core_tiled_layout.py @@ -0,0 +1,527 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from torch.utils._python_dispatch import ( + is_traceable_wrapper_subclass, + return_and_correct_aliasing, +) + +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, + register_layout, +) +from torchao.dtypes.utils import AQTTensorImpl, Layout, is_device +from torchao.quantization.quant_primitives import ZeroPointDomain, _get_reduction_params +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_5, + fill_defaults, + find_multiple, +) + +aten = torch.ops.aten + + +def _aqt_is_tensor_core_tile_uint4(aqt): + """Check if an AffineQuantizedTensor is uint4 quantized Tensor""" + # TODO: use torch.uint4 + return ( + aqt.tensor_impl.dtype == torch.int32 + and aqt.quant_min == 0 + and aqt.quant_max == 15 + ) + + +def _same_metadata( + self: "TensorCoreTiledAQTTensorImpl", src: "TensorCoreTiledAQTTensorImpl" +) -> bool: + return ( + isinstance(self, TensorCoreTiledAQTTensorImpl) + and isinstance(src, TensorCoreTiledAQTTensorImpl) + and self.shape == src.shape + and self.packed_weight.shape == src.packed_weight.shape + and self.scale_and_zero.shape == src.scale_and_zero.shape + and self.transposed == src.transposed + and type(self._layout) == type(src._layout) + ) + + +def _linear_bf16_act_uint4_weight_check(input_tensor, weight_tensor, bias): + return ( + # input is native bfloat16 tensor + not is_traceable_wrapper_subclass(input_tensor) + and input_tensor.dtype == torch.bfloat16 + and + # weight is uint4, group quantized tensor_core_tiled tensor impl affine quantized tensor + isinstance(weight_tensor, AffineQuantizedTensor) + and _aqt_is_tensor_core_tile_uint4(weight_tensor) + and weight_tensor.dtype == torch.bfloat16 + and len(weight_tensor.shape) == 2 + and weight_tensor.zero_point_domain == ZeroPointDomain.FLOAT + and isinstance(weight_tensor._layout, TensorCoreTiledLayout) + ) + + +def _linear_bf16_act_uint4_weight_impl(input_tensor, weight_tensor, bias): + assert weight_tensor.block_size[0] == 1, ( + f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}" + ) + assert input_tensor.shape[-1] == weight_tensor.shape[1], ( + f"need input_tensor shape: {input_tensor.shape} final" + f"dim to match weight_tensor shape: {weight_tensor.shape} second dim " + ) + + # TODO: check groupsize quantization + # avoid circular dep, TODO: move this to a common util.py + act_mat = input_tensor + # weight is packed from padded (out_features, in_features) weight tensor + # (same dimension requirement as F.linear weight) + packed_weight = weight_tensor.tensor_impl.packed_weight + scale_and_zero = weight_tensor.tensor_impl.scale_and_zero + + orig_act_size = act_mat.size() + orig_dtype = act_mat.dtype + + # reshape and pad activation + act_mat = act_mat.reshape(-1, act_mat.shape[-1]).to(torch.bfloat16) + pad_size = find_multiple(act_mat.shape[-1], 1024) + act_mat = torch.nn.functional.pad(act_mat, (0, pad_size - act_mat.shape[-1])) + + # groupwise int4 quantization + groupsize = weight_tensor.block_size[-1] + if act_mat.numel() == 0: # handling for empty input + y = act_mat + else: + y = torch.ops.aten._weight_int4pack_mm( + act_mat.contiguous(), packed_weight, groupsize, scale_and_zero + ) + # remove out_feature padding + orig_out_features = weight_tensor.shape[-2] + y = y[:, :orig_out_features] + y = y.reshape(*orig_act_size[:-1], orig_out_features) + + if bias is not None: + y += bias + return y.to(orig_dtype) + + +@dataclass(frozen=True) +class TensorCoreTiledLayout(Layout): + """TensorCoreTiledLayout is a layout class for handling tensor core tiled layouts in affine quantized tensors. It provides methods for pre-processing and post-processing tensors to fit the required layout for efficient computation on tensor cores. + + Attributes: + inner_k_tiles (int): An internal argument for the packing function of tensor core tiled layout that can affect the performance of the matmul kernel. Defaults to 8. + """ + + inner_k_tiles: int = 8 + + def pre_process(self, input: torch.Tensor) -> torch.Tensor: + orig_out_features, orig_in_features = input.shape[-2:] + in_features = find_multiple(orig_in_features, 1024) + out_features = find_multiple(orig_out_features, 8) + input = torch.nn.functional.pad( + input, + (0, in_features - orig_in_features, 0, out_features - orig_out_features), + ) + return input + + def pre_process_static( + self, + input: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + block_size: Tuple[int, ...], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + input = self.pre_process(input) + orig_qparam_shape = scale.shape + new_qparam_shape, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + for dim in reduction_dims: + new_qparam_shape.pop(dim) + change_in_qparam_shape = [ + new_dim_size - orig_dim_size + for new_dim_size, orig_dim_size in zip(new_qparam_shape, orig_qparam_shape) + ] + padding_changes = [] + for dim_change in change_in_qparam_shape: + padding_changes = [0, dim_change] + padding_changes + scale = torch.nn.functional.pad(scale, padding_changes) + zero_point = torch.nn.functional.pad(zero_point, padding_changes) + return input, scale, zero_point + + def post_process( + self, + input: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + block_size: Tuple[int, ...], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + orig_out_features, orig_in_features = input.shape[-2:] + in_features = find_multiple(orig_in_features, 1024) + out_features = find_multiple(orig_out_features, 8) + input = torch.nn.functional.pad( + input, + (0, in_features - orig_in_features, 0, out_features - orig_out_features), + ) + assert len(block_size) == 2 or len(block_size) == 3, ( + f"TensorCoreTiledLayout only supports len(block_size) == 2 or 3, got: {block_size}" + ) + scale_pad_dim_0 = (out_features - orig_out_features) // block_size[-2] + scale_pad_dim_1 = (in_features - orig_in_features) // block_size[-1] + scale = torch.nn.functional.pad(scale, (0, scale_pad_dim_1, 0, scale_pad_dim_0)) + zero_point = torch.nn.functional.pad( + zero_point, (0, scale_pad_dim_1, 0, scale_pad_dim_0) + ) + return input, scale, zero_point + + def extra_repr(self): + return f"inner_k_tiles={self.inner_k_tiles}" + + +@register_layout(TensorCoreTiledLayout) +class TensorCoreTiledAQTTensorImpl(AQTTensorImpl): + """TensorImpl for tensor_core_tiled layout for affine quantized tensor, this is for int4 only, + used by tinygemm kernels `_weight_int4pack_mm` + + It stores the original tensor of dimension [n][k] (int32 dtype) as packed weight of 4-d tensor of + dimension: [n / 8][k / (inner_k_tiles * 16)][32][inner_k_tiles / 2] + (unpacked Tensor shape is n * k) + where inner_k_tiles is an internal argument for packing function of tensor core tiled layout + that can affect the performance of the matmul kernel (defaults to 8) + + Note: we also pack scale and zero point together here for tinygemm kernel + + Note: technically tensor core tiled layout should be the layout for the underlying packed weight + (int Tensor) but since the scale and zero_point are also packed into the same tensor here which is not used + in plain layout, we just created a layout for AQT right now, this could be improved if we split out + int4 aqt into a separate tensor subclass + + fields: + packed_weight (torch.Tensor): the 4-d packed tensor in a tensor_core_tiled layout + scale_and_zero (torch.Tensor): the combined scale Tensor used to map between floating point tensor to quantized tensor and zero_point Tensor + """ + + def __new__( + cls, + packed_weight: torch.Tensor, + scale_and_zero: torch.Tensor, + transposed: bool, + _layout: Layout, + ): + kwargs = {} + kwargs["device"] = packed_weight.device + kwargs["layout"] = ( + kwargs.get("layout") + if kwargs.get("layout", False) + else packed_weight.layout + ) + kwargs["dtype"] = packed_weight.dtype + kwargs["requires_grad"] = False + shape = packed_weight.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + packed_weight: torch.Tensor, + scale_and_zero: torch.Tensor, + transposed: bool, + _layout: Layout, + ): + self.packed_weight = packed_weight + self.scale_and_zero = scale_and_zero + self.transposed = False + self._layout = _layout + + def __tensor_flatten__(self): + return ["packed_weight", "scale_and_zero"], [self.transposed, self._layout] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + packed_weight, scale_and_zero = ( + tensor_data_dict["packed_weight"], + tensor_data_dict["scale_and_zero"], + ) + ( + transposed, + _layout, + ) = tensor_attributes + return cls(packed_weight, scale_and_zero, transposed, _layout) + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + assert isinstance(_layout, TensorCoreTiledLayout) + assert int_data.dtype == torch.int32, ( + "torch.ops.aten._convert_weight_to_int4pack in torch 2.4 expects `int32` dtype" + ) + + def quant_2d(int_data_2d): + if TORCH_VERSION_AT_LEAST_2_5: + int_data_2d = (int_data_2d[::, ::2] << 4 | int_data_2d[::, 1::2]).to( + torch.uint8 + ) + else: + assert int_data_2d.dtype == torch.int32, ( + "torch.ops.aten._convert_weight_to_int4pack in torch 2.4 expects `int32` dtype" + ) + return torch.ops.aten._convert_weight_to_int4pack( + int_data_2d.contiguous(), _layout.inner_k_tiles + ) + + if int_data.dim() == 3: # for moe quant + num_experts = int_data.shape[0] + packed_weight_list = [] + for expert in range(num_experts): + packed_weight_list.append(quant_2d(int_data[expert]).unsqueeze(0)) + packed_weight = torch.cat(packed_weight_list, dim=0) + scale = scale.reshape(int_data.shape[0], int_data.shape[-2], -1) + zero_point = ( + zero_point.reshape(int_data.shape[0], int_data.shape[-2], -1) + if zero_point is not None + else None + ) + else: + assert int_data.dim() == 2 + packed_weight = quant_2d(int_data) + scale = scale.reshape(int_data.shape[0], -1) + zero_point = ( + zero_point.reshape(int_data.shape[0], -1) + if zero_point is not None + else None + ) + from torchao.quantization.utils import pack_tinygemm_scales_and_zeros + + scale_and_zero = pack_tinygemm_scales_and_zeros(scale, zero_point, scale.dtype) + return cls(packed_weight, scale_and_zero, False, _layout) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs["device"] + # tensor core tiled layout supports both cpu and cuda but does not support the conversion + # between these two devices, in the future we should not use the same layout for + # cpu and cuda device: https://github.com/pytorch/ao/issues/1117 + if not is_device(torch.device(self.device).type, device): + raise ValueError( + f"TensorCoreTiledAQTTensorImpl does not support conversion from {self.device} to {device}" + ) + return self.__class__( + self.packed_weight.to(device), + self.scale_and_zero.to(device), + self.transposed, + self._layout, + ) + + def _apply_fn_to_data(self, fn): + # self.packed_weight = fn(self.packed_weight) + # self.scale_and_zero = fn(self.scale_and_zero) + # return self + return self.__class__( + fn(self.packed_weight), + fn(self.scale_and_zero), + self.transposed, + self._layout, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + kwargs = {} if kwargs is None else kwargs + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + if func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + if func is aten.copy_.default: + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + + if func in [aten.select.int, aten.index.Tensor]: + assert not (func is aten.select.int and args[1] != 0), ( + "aten.select.int currently only has support for dim=0" + ) + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data(lambda x: func(x, *args[1:], **kwargs)), + ) + + if func is aten.t.default: + """we don't need to repack the weight and just rely on external + shape being changed and record the status of transpose/no-transpose + """ + transposed = TensorCoreTiledAQTTensorImpl( + args[0].packed_weight, + args[0].scale_and_zero, + not args[0].transposed, + args[0]._layout, + ) + return return_and_correct_aliasing(func, args, kwargs, transposed) + + if func is aten.slice.Tensor: + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + cur_shape = self.shape + assert len(cur_shape) == 4 + inner_k_tiles = cur_shape[-1] * 2 + original_shape = (cur_shape[0] * 8, cur_shape[1] * (inner_k_tiles * 16)) + + n_by_8, k_by_inner_tiles, _, _ = self.packed_weight.shape + sz_dim1, sz_dim0, _ = self.scale_and_zero.shape + + data_len = original_shape[dim] + assert dim in [0, 1], ( + f"TensorCoreTiledAQTTensorImpl dispatch: attempting to run {func}, with dim={dim}, that is not supported" + ) + + if dim == 0: + pw_len = n_by_8 + sz_len = sz_dim0 + else: + pw_len = k_by_inner_tiles + sz_len = sz_dim1 + + if pw_len == 0 or sz_len == 0: + return return_and_correct_aliasing( + func, + args, + kwargs, + TensorCoreTiledAQTTensorImpl( + self.packed_weight, + self.scale_and_zero, + self.transposed, + self._layout, + ), + ) + + pw_ratio = data_len / pw_len + start_pw = int(start / pw_ratio) + end_pw = int(end / pw_ratio) + + sz_ratio = data_len / sz_len + start_sz = int(start / sz_ratio) + end_sz = int(end / sz_ratio) + + packed_weight = aten.slice(self.packed_weight, dim, start_pw, end_pw, step) + scale_and_zero = aten.slice( + self.scale_and_zero, 1 - dim, start_sz, end_sz, step + ) + return return_and_correct_aliasing( + func, + args, + kwargs, + TensorCoreTiledAQTTensorImpl( + packed_weight, scale_and_zero, self.transposed, self._layout + ), + ) + + raise NotImplementedError( + f"TensorCoreTiledAQTTensorImpl dispatch: attempting to run {func}, this is not supported" + ) + + __torch_function__ = torch._C._disabled_torch_function_impl + + @property + def block_size(self): + from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros + + scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero) + cur_shape = self.shape + if len(cur_shape) == 5: + ones = [1, 1] + cur_shape = cur_shape[1:] + else: + assert len(cur_shape) == 4 + ones = [1] + inner_k_tiles = cur_shape[-1] * 2 + original_shape = (cur_shape[0] * 8, cur_shape[1] * (inner_k_tiles * 16)) + groupsize = int(original_shape[1] / scale.shape[-2]) + return tuple([*ones, groupsize]) + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from torchao.quantization.quant_primitives import ( + ZeroPointDomain, + quantize_affine, + ) + from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros + + def dequant_4d(self): + cur_shape = self.shape + scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero) + assert len(cur_shape) == 4 + inner_k_tiles = cur_shape[-1] * 2 + original_shape = (cur_shape[0] * 8, cur_shape[1] * (inner_k_tiles * 16)) + eye_shape = original_shape[1] + groupsize = int(original_shape[1] / scale.shape[-2]) + block_size = (1, groupsize) + original_dtype = torch.bfloat16 + assert len(block_size) == 2 and block_size[0] == 1 + dequantized = torch.ops.aten._weight_int4pack_mm( + torch.eye(eye_shape, device=self.device, dtype=original_dtype), + self.packed_weight, + groupsize, + self.scale_and_zero, + ) + dequantized = dequantized.t().contiguous() + return dequantized + + cur_shape = self.shape + + if len(cur_shape) == 4: + dequantized = dequant_4d(self) + else: + assert len(cur_shape) == 5 + num_experts = cur_shape[0] + dequantized_list = [] + for expert in range(num_experts): + dequantized_list.append(dequant_4d(self[expert]).unsqueeze(0)) + dequantized = torch.cat(dequantized_list, dim=0) + + scale, zero = unpack_tinygemm_scales_and_zeros(self.scale_and_zero) + # TODO: move this to `unpack_tinygemm_scales_and_zeros`? + scale = scale.reshape(scale.shape[:-1]).contiguous() + zero = zero.reshape(zero.shape[:-1]).contiguous() + + target_dtype = torch.int32 + quant_min = 0 + quant_max = 15 + zero_point_domain = ZeroPointDomain.FLOAT + int_data = quantize_affine( + dequantized, + self.block_size, + scale, + zero, + target_dtype, + quant_min, + quant_max, + zero_point_domain, + ) + return int_data, scale, zero + + def get_layout(self) -> Layout: + return self._layout diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/uint4_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/uint4_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..9ff498b314c842cd53e7690c9c4107c4cb984fe7 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/uint4_layout.py @@ -0,0 +1,321 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +import torch._prims_common as utils +import torch.utils._pytree as pytree +from torch.library import Library, impl + +from torchao.utils import fill_defaults + + +def down_size(size): + assert size[-1] % 2 == 0, f"{size} last dim not divisible by two" + return (*size[:-1], size[-1] // 2) + + +def up_size(size): + return (*size[:-1], size[-1] * 2) + + +# from +# https://github.com/drisspg/transformer_nuggets/blob/9ad3a7fc552a954eb702ade0e276b8d8e09c3db6/transformer_nuggets/quant/qlora.py#L233 + + +def unpack_uint4(uint8_data) -> torch.Tensor: + """Get the original weight from the normalized float weight format""" + # since we are using uint8 we will decode 2 entries per byte + # Shift elements down 4 and select out the bottom 4 bits + shape = uint8_data.shape + first_elements = (uint8_data >> 4).to(torch.uint8) + second_elements = (uint8_data & 0b1111).to(torch.uint8) + return torch.stack([first_elements, second_elements], dim=-1).view(up_size(shape)) + + +def pack_uint4(uint8_data) -> torch.Tensor: + # converting to uint8 for operations + shape = uint8_data.shape + assert shape[-1] % 2 == 0 + uint8_data = uint8_data.contiguous().view(-1) + return (uint8_data[::2] << 4 | uint8_data[1::2]).view(down_size(shape)) + + +qtensor_lib = Library("qtensors", "DEF") +qtensor_lib.define( + "quantize_per_tensor_uint4(Tensor input, float scale, int zero_point) -> Tensor" +) + + +@impl(qtensor_lib, "quantize_per_tensor_uint4", "CompositeExplicitAutograd") +def quantize_per_tensor_uint4( + input: torch.Tensor, + scale: float, + zero_point: int, +) -> torch.Tensor: + inv_scale = 1.0 / scale + return pack_uint4( + torch.clamp(torch.round(input * inv_scale) + zero_point, 0, 15).to(torch.uint8) + ) + + +qtensor_lib.define( + "dequantize_per_tensor_uint4(Tensor input, float scale, int zero_point) -> Tensor" +) + + +@impl(qtensor_lib, "dequantize_per_tensor_uint4", "CompositeExplicitAutograd") +def dequantize_per_tensor_uint4( + input: torch.Tensor, + scale: float, + zero_point: int, +) -> torch.Tensor: + input = unpack_uint4(input) + return (input.view(torch.uint8).to(torch.float32) - zero_point) * scale + + +class UInt4Tensor(torch.Tensor): + @staticmethod + def __new__(cls, elem, **kwargs): + assert elem.dtype is torch.uint8 + assert not kwargs.get("requires_grad", False) + kwargs["requires_grad"] = False + + return torch.Tensor._make_wrapper_subclass( + cls, up_size(elem.shape), dtype=torch.uint4, **kwargs + ) + + def __init__(self, elem, **kwargs): + self.elem = elem + + @classmethod + def from_unpacked(cls, unpacked): + return UInt4Tensor(pack_uint4(unpacked)) + + def tolist(self): + return self.to(torch.uint8).tolist() + + def __tensor_flatten__(self): + return ["elem"], None + + @staticmethod + def __tensor_unflatten__(flattened, meta, outer_size, outer_stride): + assert meta is None + elem = flattened["elem"] + return UInt4Tensor(elem) + + def __hash__(self): + return hash(self.elem) + + def __eq__(self, other): + return torch.equal(self.elem, other.elem) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if func is torch.ops.aten.view.default: + self, size = args + size = utils.infer_size(size, self.numel()) + assert not kwargs + # WARNING: views not preserved + return UInt4Tensor(self.elem.reshape(down_size(size))) + elif func is torch.ops.aten.view.dtype: + self, dtype = args + if dtype == torch.uint8: + return unpack_uint4(self.elem).view(torch.uint8) + return NotImplementedError(f"view {args}") + elif func is torch.ops.aten.to.dtype: + self, dtype = args + if dtype == torch.uint8: + return unpack_uint4(self.elem).view(torch.uint8) + return NotImplementedError(f"to {args}") + elif func is torch.ops.aten.eq.Tensor: + args = pytree.tree_map_only( + UInt4Tensor, lambda x: x.elem.view(torch.uint8), args + ) + kwargs = pytree.tree_map_only( + UInt4Tensor, lambda x: x.elem.view(torch.uint8), kwargs + ) + return torch.ops.aten.eq.Tensor(*args, **kwargs) + elif func is torch.ops.aten._to_copy.default: + (self,) = args + if kwargs == {"dtype": torch.uint8}: + return unpack_uint4(self.elem).view(self.shape) # no wrap + else: + raise NotImplementedError(f"_to_copy {kwargs}") + elif func is torch.ops.aten.unbind.int: + # This is tricky. Given torch.tensor([0, 1, 2, 3]) we want to + # create four tensors containing one element each. But we can't + # do this with uint4 because such a tensor's size is not divisible + # by bytes. What I am going to do instead is promote to uint8 + # when this happens + self, dim = fill_defaults(args, 2, [0]) + if dim != self.dim() - 1: + raise NotImplementedError(f"unbind dim={dim}") + else: + # We're unbinding the last dimension, need to promote + return torch.ops.aten._to_copy.default(self, dtype=torch.uint8).unbind( + dim + ) + elif func is torch.ops.aten.select.int: + self, dim, index = args + if dim != self.dim() - 1: + return UInt4Tensor(torch.ops.aten.select.int(self.elem, dim, index)) + else: + raise NotImplementedError(f"select dim={dim}") + elif func is torch.ops.aten.slice.Tensor: + self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1]) + if dim == self.dim() - 1: + # hard case + if step != 1: + raise NotImplementedError(f"slice step={step}") + assert start % 2 == 0, start + assert end >= self.shape[dim] or end % 2 == 0, end + return UInt4Tensor( + torch.ops.aten.slice.Tensor(self.elem, dim, start // 2, end // 2, 1) + ) + else: + # easy case + return UInt4Tensor( + torch.ops.aten.slice.Tensor(self.elem, dim, start, end, step) + ) + elif func is torch.ops.aten.t.default: + # assert False, "transpose is not properly implemented currently" + (self,) = args + unpacked = unpack_uint4(self.elem) + transposed = torch.ops.aten.t.default(unpacked) + transposed_and_packed = pack_uint4(transposed) + return UInt4Tensor(transposed_and_packed) + elif func is torch.ops.aten.transpose_copy.int: + self, dim0, dim1 = args + unpacked = unpack_uint4(self.elem).view(self.shape) + transposed = torch.ops.aten.transpose_copy.int(unpacked, dim0, dim1) + transposed_and_packed = pack_uint4(transposed) + return UInt4Tensor(transposed_and_packed) + elif func is torch.ops.aten.as_strided.default: + # size, stride, storage_offset are referring to tensor elements, not physical bytes + self, size, stride, storage_offset = args + size = down_size(size) + + new_stride = [] + for s in stride: + if s != 1: + # since two int4 equals to 1 uint8 + new_stride.append(s // 2) + else: + new_stride.append(s) + stride = new_stride + + storage_offset //= 2 + return UInt4Tensor( + torch.ops.aten.as_strided.default( + self.elem, size, stride, storage_offset + ) + ) + + raise NotImplementedError(f"{func}") + + __torch_function__ = torch._C._disabled_torch_function_impl + + +def _dynamically_quantize_per_channel_int4(x, quant_min, quant_max, target_dtype): + # assumes symmetric quantization + # assumes axis == 0 + # assumes dense memory format + # TODO(future): relax ^ as needed + + # default setup for affine quantization of activations + eps = torch.finfo(torch.float32).eps + + # get min and max + min_val, max_val = torch.aminmax(x, dim=1) + + # calculate scale and zero point based on min and max + # reference: https://fburl.com/code/srbiybme + min_val_neg = torch.min(min_val, torch.zeros_like(min_val)) + max_val_pos = torch.max(max_val, torch.zeros_like(max_val)) + device = min_val_neg.device + + # reference: https://fburl.com/code/4wll53rk + max_val_pos = torch.max(-min_val_neg, max_val_pos) + scale = max_val_pos / (float(quant_max - quant_min) / 2) + # ensure scale is the same dtype as the original tensor + scale = torch.clamp(scale, min=eps).to(x.dtype) + zero_point = torch.zeros(min_val_neg.size(), dtype=torch.int64, device=device) + + # quantize based on qmin/qmax/scale/zp + # reference: torch/ao/quantization/fx/_decomposed.py?lines=63 + x_div = x.transpose(0, 1) / scale + x_round = torch.round(x_div) + x_zp = x_round + zero_point + x_zp = x_zp.transpose(0, 1) + quant = torch.clamp(x_zp, quant_min, quant_max) + + if target_dtype == torch.uint4: + # TODO: simplify (maybe implement to) + quant = PerChannelSymmetricWeightUInt4Tensor.from_unpacked( + quant.to(torch.uint8), scale + ) + else: + quant = quant.to(target_dtype) + + return quant, scale, zero_point + + +class PerChannelSymmetricWeightUInt4Tensor(UInt4Tensor): + @staticmethod + def __new__(cls, elem, scales, **kwargs): + return super().__new__(cls, elem, **kwargs) + + def __init__(self, elem, scales, **kwargs): + super().__init__(elem, **kwargs) + + self.scales = scales + + def __tensor_flatten__(self): + return ["elem", "scales"], None + + @staticmethod + def __tensor_unflatten__(flattened, meta, outer_size, outer_stride): + assert meta is None + elem = flattened["elem"] + scales = flattened["scales"] + return PerChannelSymmetricWeightUInt4Tensor(elem, scales) + + @classmethod + + # inconsistently. + + def from_unpacked(cls, unpacked, scales): + return cls(pack_uint4(unpacked), scales) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if func is torch.ops.aten.addmm.default: + bias, x, weight = args + x_view = x.view(-1, x.shape[-1]) + y = torch.mm(x_view, weight.to(torch.uint8).to(x.dtype)) * weight.scales + y = y.reshape(*x.shape[:-1], -1) + if bias is not None: + y += bias + return y + elif func is torch.ops.aten.t.default: + # TODO: add proper support for transpose + (self,) = args + unpacked = unpack_uint4(self.elem) + transposed = torch.ops.aten.t.default(unpacked) + return PerChannelSymmetricWeightUInt4Tensor.from_unpacked( + transposed, self.scales + ) + elif func is torch.ops.aten.detach.default: + (self,) = args + return self + return super().__torch_dispatch__(func, types, args, kwargs) + + @classmethod + def from_float(cls, w_fp32): + w_int4, scales, _zp = _dynamically_quantize_per_channel_int4( + w_fp32, 0, 15, torch.uint4 + ) + w_int4 = w_int4.to(device=w_fp32.device) + return w_int4 diff --git a/lib/python3.12/site-packages/torchao/dtypes/uintx/uintx_layout.py b/lib/python3.12/site-packages/torchao/dtypes/uintx/uintx_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..96e5401de590bd3072302da163b0cf771cfbc645 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/uintx/uintx_layout.py @@ -0,0 +1,255 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import List, Tuple + +import torch +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.dtypes.affine_quantized_tensor import register_layout +from torchao.dtypes.uintx.plain_layout import PlainAQTTensorImpl +from torchao.dtypes.utils import ( + Layout, +) +from torchao.utils import TORCH_VERSION_AT_LEAST_2_3, TorchAOBaseTensor + +from .bitpacking import pack, unpack + +aten = torch.ops.aten + +# Note: Uintx does not work for torch 2.3 and below +_DTYPE_TO_BIT_WIDTH = {} +_BIT_WIDTH_TO_DTYPE = {} + +if TORCH_VERSION_AT_LEAST_2_3: + _DTYPE_TO_BIT_WIDTH = { + torch.uint1: 1, + torch.uint2: 2, + torch.uint3: 3, + torch.uint4: 4, + torch.uint5: 5, + torch.uint6: 6, + torch.uint7: 7, + } + + _BIT_WIDTH_TO_DTYPE = {v: k for k, v in _DTYPE_TO_BIT_WIDTH.items()} +else: + print("uintx feature requires torch 2.3+, please upgrade pytorch") + + +class UintxTensor(TorchAOBaseTensor): + """ + Splits int data into packed shards based on bit size + fields: + int4_shard (torch.Tensor): 4 bit packed shard + int2_shard (torch.Tensor): 2 bit packed shard + int1_shard (torch.Tensor): 1 bit packed shard + bit_width (int): number of bits for each element + pack_dim: (int) dimension to pack along + """ + + bits_to_shard = { + 1: ["int1_shard"], + 2: ["int2_shard"], + 3: ["int2_shard", "int1_shard"], + 4: ["int4_shard"], + 5: ["int4_shard", "int1_shard"], + 6: ["int4_shard", "int2_shard"], + 7: ["int4_shard", "int2_shard", "int1_shard"], + } + + def __new__( + cls, + shards: List[torch.Tensor], + packed_shape: List[int], + bit_width: int, + pack_dim: int = -1, + ): + kwargs = {"device": shards[0].device} + kwargs["device"] = shards[0].device + kwargs["layout"] = shards[0].layout + kwargs["requires_grad"] = False + kwargs["dtype"] = torch.uint8 + return torch.Tensor._make_wrapper_subclass(cls, packed_shape, **kwargs) + + def __init__( + self, + shards: List[torch.Tensor], + packed_shape: List[int], + bit_width: int, + pack_dim: int = -1, + ): + for i, attrib in enumerate(self.bits_to_shard[bit_width]): + setattr(self, attrib, shards[i]) + + self.packed_shape = packed_shape + self.bit_width = bit_width + self.pack_dim = pack_dim + + def get_shards(self): + return [getattr(self, i) for i in self.__class__.bits_to_shard[self.bit_width]] + + def __repr__(self): + return f"Int{self.bit_width}Tensor(shape = {self.packed_shape}, data = {unpack(self.get_shards(), self.bit_width, dim=self.pack_dim)})" + + def __tensor_flatten__(self): + return self.__class__.bits_to_shard[self.bit_width], [ + self.packed_shape, + self.bit_width, + self.pack_dim, + ] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + shards = list(tensor_data_dict.values()) + packed_shape, bit_width, pack_dim = tensor_attributes + return cls(shards, packed_shape, bit_width, pack_dim) + + def get_plain(self): + return unpack(self.get_shards(), self.bit_width, dim=self.pack_dim) + + # temporary until kernels on packed tensors are created + def apply_transformation(self, fn): + og = self.get_plain() + new = fn(og) + dtype = _BIT_WIDTH_TO_DTYPE[self.bit_width] + return self.from_uint8(new, dtype, self.pack_dim) + + # temporary until kernels on packed tensors are created + def apply_fn_to_shards(self, fn): + new_shards = [fn(shard) for shard in self.get_shards()] + return self.__class__( + new_shards, self.packed_shape, self.bit_width, self.pack_dim + ) + + @classmethod + def from_uint8(cls, int_data: torch.Tensor, dtype: torch.dtype, pack_dim: int = -1): + assert dtype in _DTYPE_TO_BIT_WIDTH.keys(), ( + "Expected dtype to be one of {_DTYPE_TO_BIT_WIDTH.keys()}" + ) + bit_width = _DTYPE_TO_BIT_WIDTH[dtype] + shards = pack(int_data, bit_width, dim=pack_dim) + shape = list(int_data.shape) + shape[pack_dim] = shape[pack_dim] * bit_width // 8 + return cls(shards, int_data.shape, bit_width, pack_dim) + + def _get_to_kwargs(self, *args, **kwargs): + device, dtype, _, memory_format = torch._C._nn._parse_to(*args, **kwargs) + device = self.device if device is None else device + dtype = self.dtype if dtype is None else dtype + memory_format = ( + memory_format if memory_format is not None else torch.preserve_format + ) + kwargs = { + "device": device, + "dtype": dtype, + "memory_format": memory_format, + } + return kwargs + + def to(self, *args, **kwargs): + if "copy" in kwargs: + return super().to(*args, **kwargs) + kwargs = self._get_to_kwargs(*args, **kwargs) + if "device" in kwargs: + return self.__class__( + list(shard.to(kwargs["device"]) for shard in self.get_shards()), + self.packed_shape, + self.bit_width, + self.pack_dim, + ) + return super().to(*args, **kwargs) + + +implements = UintxTensor.implements + + +@implements(aten.detach.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0].apply_fn_to_shards(torch.detach) + ) + + +@implements(aten.view.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0].apply_transformation(lambda x: x.view(*args[1:])) + ) + + +@implements(aten._to_copy.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing(func, args, kwargs, args[0]) + + +@implements(aten.sub.Tensor) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].apply_transformation(lambda x: (x - args[1]).to(torch.uint8)), + ) + + +@implements(aten.mul.Tensor) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].apply_transformation(lambda x: (x * args[1]).to(torch.uint8)), + ) + + +# quantization api integrations +to_uintx = UintxTensor.from_uint8 + + +@dataclass(frozen=True) +class UintxLayout(Layout): + """A layout class for Uintx tensors, which are tensors with elements packed into + smaller bit-widths than the standard 8-bit byte. This layout is used to define + how the data is stored and processed in UintxTensor objects. + + Attributes: + dtype (torch.dtype): The data type of the tensor elements, which determines + the bit-width used for packing. + pack_dim (int): The dimension along which the data is packed. Default is -1, + which indicates the last dimension. + """ + + dtype: torch.dtype + pack_dim: int = -1 + + def post_process( + self, + input: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + block_size: Tuple[int, ...], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return to_uintx(input, self.dtype, self.pack_dim), scale, zero_point + + +@register_layout(UintxLayout) +class UintxAQTTensorImpl(PlainAQTTensorImpl): + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.int_data.get_plain(), self.scale, self.zero_point + + @classmethod + def from_plain( + cls, + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + _layout: Layout, + ): + assert isinstance(_layout, UintxLayout) + return cls(int_data, scale, zero_point, _layout) diff --git a/lib/python3.12/site-packages/torchao/dtypes/utils.py b/lib/python3.12/site-packages/torchao/dtypes/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a07188a18d9b2c0b64c64b5469e5ce67d12a432c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/dtypes/utils.py @@ -0,0 +1,135 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch + +from torchao.utils import TorchAOBaseTensor + +""" +Base class for different layout, following the same design of PyTorch layout +https://pytorch.org/docs/stable/tensor_attributes.html#torch-layout, used to represent different +data layout of a Tensor, it's used in conjunction with TensorImpl to represent custom data layout. + +As a native PyTorch example, Sparse Coordinate format Tensor (https://pytorch.org/docs/stable/generated/torch.sparse_coo_tensor.html#torch-sparse-coo-tensor) has `torch.sparse_coo` layout, which is backed up by +`SparseImpl`: https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/SparseTensorImpl.h which stores two Tensors (indices_ and values_) + +We extended the layout in torchao with Layout class (instead of torch.layout objects), also we use tensor subclass to implement TensorImpl classes. + +Layout also allows users to pass around configurations for the TensorImpl, +e.g. inner_k_tiles for int4 tensor core tiled TensorImpl + +Note: Layout is an abstraction not only for custom data representation, it is also used for how the +Tensor interacts with different operators, e.g. the same data representation can have different +behaviors when running the same operator, e.g. transpose, quantized_linear. This is the same as layout +in PyTorch native Tensor +""" + + +@dataclass(frozen=True) +class Layout: + """The Layout class serves as a base class for defining different data layouts for tensors. + It provides methods for pre-processing and post-processing tensors, as well as static + pre-processing with additional parameters like scale, zero_point, and block_size. + + The Layout class is designed to be extended by other layout classes that define specific + data representations and behaviors for tensors. It is used in conjunction with TensorImpl + classes to represent custom data layouts and how tensors interact with different operators. + """ + + def pre_process(self, input: torch.Tensor) -> torch.Tensor: + return input + + def post_process( + self, + input: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + block_size: Tuple[int, ...], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return input, scale, zero_point + + def pre_process_static( + self, + input: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + block_size: Tuple[int, ...], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.pre_process(input), scale, zero_point + + def __repr__(self): + return f"{self.__class__.__name__}({self.extra_repr()})" + + def extra_repr(self) -> str: + return "" + + +@dataclass(frozen=True) +class PlainLayout(Layout): + """PlainLayout is the most basic layout class, inheriting from the Layout base class. + It does not add any additional metadata or processing steps to the tensor. + Typically, this layout is used as the default when no specific layout is required. + """ + + pass + + +def is_device(target_device_str: str, device: Union[str, torch.device]): + return torch.device(device).type == target_device_str + + +def get_out_shape(input_shape: Tuple[int], weight_shape: Tuple[int]) -> Tuple[int, int]: + """Returns the unflattened shape of the input tensor. + Args: + input_shape: The input tensor shape possibly more than 2 dimensions + weight_shape: The weight tensor shape. + Returns: + The unflattened shape of the input tensor. + """ + out_dim = weight_shape[0] + inpt_dims = input_shape[:-1] + return (*inpt_dims, out_dim) + + +############################### +# Base Tensor Impl Subclass # +############################### +class AQTTensorImpl(TorchAOBaseTensor): + """ + Base class for the tensor impl for `AffineQuantizedTensor` + + Note: This is not a user facing API, it's used by AffineQuantizedTensor to construct + the underlying implementation of a AQT based on layout + """ + + def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Get the plain (unpacked) Tensor for the tensor impl + + Returns data, scale and zero_point + Can be overwritten if other types of AQTTensorImpl has different numbers of plain tensors + """ + pass + + def get_layout(self) -> Layout: + pass + + @classmethod + def from_plain( + cls, + data: torch.Tensor, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + _layout: Layout, + ): + """Construct a TensorImpl from data, scale, zero_point and the _layout""" + pass + + def __repr__(self): + data, scale, zero_point = self.get_plain() + _layout = self.get_layout() + return f"{self.__class__.__name__}(data={str(data)}... , scale={str(scale)}... , zero_point={str(zero_point)}... , _layout={_layout})" diff --git a/lib/python3.12/site-packages/torchao/experimental/__init__.py b/lib/python3.12/site-packages/torchao/experimental/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/experimental/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/experimental/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..367285467399d3190d026dba2d373302a639119f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/experimental/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/experimental/__pycache__/op_lib.cpython-312.pyc b/lib/python3.12/site-packages/torchao/experimental/__pycache__/op_lib.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a27cd95af4b424a94c3e78fbcde543e7e6a91b10 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/experimental/__pycache__/op_lib.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/experimental/__pycache__/op_lib_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/experimental/__pycache__/op_lib_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf6a696ca40c2b53ef3173a910061a889e97991a Binary files /dev/null and b/lib/python3.12/site-packages/torchao/experimental/__pycache__/op_lib_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/experimental/__pycache__/packed_linear_int8_dynamic_activation_intx_weight_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/experimental/__pycache__/packed_linear_int8_dynamic_activation_intx_weight_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70dc841c1cf193dd83a081480ccc5d1e2b21dcda Binary files /dev/null and b/lib/python3.12/site-packages/torchao/experimental/__pycache__/packed_linear_int8_dynamic_activation_intx_weight_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/experimental/__pycache__/q_dq_layout.cpython-312.pyc b/lib/python3.12/site-packages/torchao/experimental/__pycache__/q_dq_layout.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..188c716c8f1e0f4a16d9e2ac432abc846425131a Binary files /dev/null and b/lib/python3.12/site-packages/torchao/experimental/__pycache__/q_dq_layout.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/experimental/__pycache__/quant_api.cpython-312.pyc b/lib/python3.12/site-packages/torchao/experimental/__pycache__/quant_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24c76bf88fcbfd496b250434b3817561cbdb9c74 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/experimental/__pycache__/quant_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/experimental/__pycache__/quant_passes.cpython-312.pyc b/lib/python3.12/site-packages/torchao/experimental/__pycache__/quant_passes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94434da3563706ed1d9576e9209f694bf3eb496b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/experimental/__pycache__/quant_passes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/experimental/__pycache__/temp_build.cpython-312.pyc b/lib/python3.12/site-packages/torchao/experimental/__pycache__/temp_build.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8abd69ec2f7905d6cd61b0ad306002989a2b55b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/experimental/__pycache__/temp_build.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/experimental/op_lib.py b/lib/python3.12/site-packages/torchao/experimental/op_lib.py new file mode 100644 index 0000000000000000000000000000000000000000..4fe478d1e8c7878e7434e26fc7049fd10e027bd6 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/experimental/op_lib.py @@ -0,0 +1,57 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import torch +from torch import Tensor +from torch.library import impl + +# Load C++ ops +lib_path = Path(__file__).parent.parent +libs = list(lib_path.glob("libtorchao_ops_aten.*")) +assert len(libs) == 1, ( + f"Expected to find one libtorchao_ops_aten.* library at {lib_path}, but found {len(libs)}" +) +torch.ops.load_library(str(libs[0])) + + +# Define meta ops. To support dynamic shapes, some meta ops need to +# be defined in python instead of C++. +torchao_lib = torch.library.Library("torchao", "IMPL") +for weight_nbit in range(1, 9): + + @impl(torchao_lib, f"_linear_8bit_act_{weight_nbit}bit_weight", "Meta") + def _( + activations: Tensor, + packed_weights: Tensor, + group_size: int, + n: int, + k: int, + ): + assert activations.dim() == 2 + m, k_ = activations.shape + assert k_ == k + return torch.empty(m, n, dtype=activations.dtype, device="meta") + + @impl(torchao_lib, f"_embedding_{weight_nbit}bit", "Meta") + def _( + packed_weight_qvals: Tensor, + num_embeddings: int, + embedding_dim: int, + weight_scales: Tensor, + weight_zeros: Tensor, + indices: Tensor, + ): + assert indices.dim() == 1 + num_out = indices.shape[0] + return torch.empty(num_out, embedding_dim, dtype=torch.float32, device="meta") + + @impl(torchao_lib, f"_shared_embedding_{weight_nbit}bit", "Meta") + def _(packed_weights: Tensor, group_size: int, n: int, k: int, indices: Tensor): + assert indices.dim() == 1 + num_out = indices.shape[0] + return torch.empty(num_out, k, dtype=torch.float32, device="meta") diff --git a/lib/python3.12/site-packages/torchao/experimental/op_lib_utils.py b/lib/python3.12/site-packages/torchao/experimental/op_lib_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..25cb8a1ed28dfda04d717e41169367735e26b479 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/experimental/op_lib_utils.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch + + +def _check_torchao_ops_loaded(): + # Check kernels are installed/loaded + try: + torch.ops.torchao._pack_8bit_act_4bit_weight + except AttributeError: + raise Exception( + "TorchAO experimental kernels are not loaded. To install the kernels, run `USE_CPP=1 pip install .` from ao on a machine with an ARM CPU." + + " You can also set target to 'aten' if you are using ARM CPU." + ) diff --git a/lib/python3.12/site-packages/torchao/experimental/packed_linear_int8_dynamic_activation_intx_weight_layout.py b/lib/python3.12/site-packages/torchao/experimental/packed_linear_int8_dynamic_activation_intx_weight_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..b6b9fcbcc523c2b6c7effbf1033a3d229e96c653 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/experimental/packed_linear_int8_dynamic_activation_intx_weight_layout.py @@ -0,0 +1,22 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# TODO: delete this file. +# File is kept in torchao/experimental to avoid breaking existing code +import logging + +logging.warning( + "torchao.experimental.packed_linear_int8_dynamic_activation_intx_weight_layout.py is deprecated and will be removed. Please use torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout.py instead." +) +from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import ( + PackedLinearInt8DynamicActivationIntxWeightLayout, + Target, +) + +__all__ = [ + "PackedLinearInt8DynamicActivationIntxWeightLayout", + "Target", +] diff --git a/lib/python3.12/site-packages/torchao/experimental/q_dq_layout.py b/lib/python3.12/site-packages/torchao/experimental/q_dq_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..5eeea7f4bdd44bcbe3e6cbfef31541e227eae475 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/experimental/q_dq_layout.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# TODO: delete this file. +# File is kept in torchao/experimental to avoid breaking existing code +import logging + +logging.warning( + "torchao.experimental.q_dq_layout.py is deprecated and will be removed. Please use torchao.dtypes.uintx.q_dq_layout.py instead." +) +from torchao.dtypes import QDQLayout + +__all__ = [ + "QDQLayout", +] diff --git a/lib/python3.12/site-packages/torchao/experimental/quant_api.py b/lib/python3.12/site-packages/torchao/experimental/quant_api.py new file mode 100644 index 0000000000000000000000000000000000000000..b7630cada3beb04e955c36616e9173d53f2fddd8 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/experimental/quant_api.py @@ -0,0 +1,638 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import sys +from typing import Callable, List, Mapping, Optional, Tuple, Union + +import torch +import torch.nn as nn +from torch.ao.quantization.fx._decomposed import ( + quantize_per_channel_group, +) + +logger = logging.getLogger(__name__) +logger.setLevel(logging.WARNING) + + +handler = logging.StreamHandler(sys.stdout) +formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") +handler.setFormatter(formatter) +logger.addHandler(handler) + +from dataclasses import dataclass + +from torchao.core.config import AOBaseConfig +from torchao.dtypes.affine_quantized_tensor import ( + AffineQuantizedTensor, +) +from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import ( + PackedLinearInt8DynamicActivationIntxWeightLayout, + Target, +) +from torchao.experimental.op_lib_utils import _check_torchao_ops_loaded +from torchao.quantization.granularity import Granularity, PerAxis, PerGroup, PerRow +from torchao.quantization.quant_api import ( + Int8DynamicActivationIntxWeightConfig as Int8DynamicActivationIntxWeightConfig_NonExperimental, +) +from torchao.quantization.quant_api import ( + IntxWeightOnlyConfig, + MappingType, + quantize_, +) +from torchao.quantization.quant_primitives import _DTYPE_TO_BIT_WIDTH + + +@dataclass +class Int8DynamicActivationIntxWeightConfig(AOBaseConfig): + weight_dtype: torch.dtype = torch.int4 + granularity: Union[PerRow, PerGroup] = PerRow() + has_weight_zeros: bool = False + weight_mapping_type: MappingType = MappingType.ASYMMETRIC + act_mapping_type: MappingType = MappingType.ASYMMETRIC + round_weight_scale_to_bf16: bool = True + layout = PackedLinearInt8DynamicActivationIntxWeightLayout(target=Target.AUTO) + + def __post_init__(self): + raise NotImplementedError( + "Int8DynamicActivationIntxWeightConfig has moved from torchao.experimental.quant_api to torchao.quantization.quant_api.\n" + "Please migrate to using the new version. The following args are renamed in the new version:\n" + "* granularity -> weight_granularity\n" + "* has_weight_zeros=True -> weight_mapping_type=torchao.quantization.quant_api.MappingType.ASYMMETRIC\n" + "* has_weight_zeros=False -> weight_zero_point_domain=torchao.quantization.quant_api.MappingType.SYMMETRIC\n" + "* round_weight_scale_to_bf16=True -> weight_scale_dtype=torch.bfloat16\n" + "* layout default has changed to QDQLayout(). IF YOU WANT CPU PERFORMANCE, USE layout=PackedLinearInt8DynamicActivationIntxWeightLayout()." + ) + + +# For BC +int8_dynamic_activation_intx_weight = Int8DynamicActivationIntxWeightConfig + + +class QuantizedEmbedding(nn.Module): + def __init__( + self, + bit_width, + ): + super().__init__() + self.bit_width = bit_width + + def quantize_and_pack_weights(self, weights, group_size, mapping_type): + num_embeddings, embedding_dim = weights.shape + + embedding = torch.nn.Embedding(num_embeddings, embedding_dim) + embedding.weight = weights + quantize_( + embedding, + IntxWeightOnlyConfig( + weight_dtype=getattr(torch, f"int{self.bit_width}"), + granularity=PerGroup(group_size) if group_size > 0 else PerAxis(0), + mapping_type=mapping_type, + ), + lambda m, fqn: isinstance(m, torch.nn.Embedding), + ) + weight_qvals, weight_scales, weight_zeros = ( + embedding.weight.tensor_impl.get_plain() + ) + assert weight_zeros is not None + weight_scales = weight_scales.reshape(num_embeddings, -1) + weight_zeros = weight_zeros.reshape(num_embeddings, -1).to(torch.int8) + self.register_buffer( + "packed_weight_qvals", + getattr(torch.ops.torchao, f"_pack_embedding_{self.bit_width}bit")( + weight_qvals.to(torch.int8) + ), + ) + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + self.register_buffer("weight_scales", weight_scales) + self.register_buffer("weight_zeros", weight_zeros) + + def forward(self, x): + shape = x.shape + return getattr(torch.ops.torchao, f"_embedding_{self.bit_width}bit")( + self.packed_weight_qvals, + self.num_embeddings, + self.embedding_dim, + self.weight_scales, + # embedding op requires weight_zeros be passed, even if they are all 0 + self.weight_zeros, + x.reshape(-1), + ).reshape(*shape, -1) + + +class QuantizedEmbeddingFallback(nn.Module): + def __init__( + self, + bit_width, + ): + super().__init__() + self.bit_width = bit_width + + def quantize_and_pack_weights(self, weights, group_size, mapping_type): + self.embedding = torch.nn.Embedding(*weights.shape) + self.embedding.weight = weights + quantize_( + self.embedding, + IntxWeightOnlyConfig( + weight_dtype=getattr(torch, f"int{self.bit_width}"), + granularity=PerGroup(group_size) if group_size > 0 else PerAxis(0), + mapping_type=mapping_type, + ), + lambda m, fqn: isinstance(m, torch.nn.Embedding), + ) + + def forward(self, x): + return self.embedding(x) + + +class QuantizedSharedEmbedding(nn.Module): + def __init__(self, bit_width, unembedding_packed_weights, group_size, n, k): + super().__init__() + self.bit_width = bit_width + self.register_buffer("unembedding_packed_weights", unembedding_packed_weights) + self.n = n + self.k = k + if group_size == -1: + self.group_size = k + else: + self.group_size = group_size + self.shared_embedding_op = getattr( + torch.ops.torchao, f"_shared_embedding_{bit_width}bit" + ) + + def forward(self, x): + shape = x.shape + return self.shared_embedding_op( + self.unembedding_packed_weights, + self.group_size, + self.n, + self.k, + x.reshape(-1), + ).reshape(*shape, -1) + + +def _replace_embedding_with_quantized_embedding( + module: nn.Module, + kwargs={}, + fqn: str = "", +): + group_size = kwargs.get("group_size", None) + bit_width = kwargs.get("bit_width", None) + use_fallback = kwargs.get("use_fallback", None) + mapping_type = kwargs.get("mapping_type", None) + embedding_fqn_to_quantized_unembedding = kwargs.get( + "embedding_fqn_to_quantized_unembedding", None + ) + + assert not isinstance(module, nn.Embedding) + for name, child in module.named_children(): + child_fqn = f"{fqn}.{name}" if fqn != "" else name + + if not isinstance(child, nn.Embedding): + _replace_embedding_with_quantized_embedding(child, kwargs, child_fqn) + else: + assert child.weight.device == torch.device("cpu"), "Only CPU is supported" + assert child.weight.dtype == torch.float32, "Only float32 is supported" + + if use_fallback: + qembedding = QuantizedEmbeddingFallback(bit_width) + setattr(module, name, qembedding) + getattr(module, name).quantize_and_pack_weights( + child.weight, + group_size, + mapping_type, + ) + else: + _check_torchao_ops_loaded() + if embedding_fqn_to_quantized_unembedding is None: + qembedding = QuantizedEmbedding(bit_width) + setattr(module, name, qembedding) + getattr(module, name).quantize_and_pack_weights( + child.weight, + group_size, + mapping_type, + ) + else: + if child_fqn not in embedding_fqn_to_quantized_unembedding: + continue + weight_tensor = embedding_fqn_to_quantized_unembedding[child_fqn] + n, k = weight_tensor.shape + group_size = weight_tensor.tensor_impl.get_layout().group_size + packed_weight = weight_tensor.tensor_impl.packed_weight + bit_width = weight_tensor.tensor_impl.get_layout().bit_width + + assert n == child.num_embeddings, ( + "num_embeddings must match n in shared_unembedding" + ) + assert k == child.embedding_dim, ( + "embedding_dim must match k in shared_unembedding" + ) + qembedding = QuantizedSharedEmbedding( + bit_width, + packed_weight, + group_size, + n, + k, + ) + setattr(module, name, qembedding) + + +class EmbeddingQuantizer: + def __init__( + self, + weight_dtype: torch.dtype = torch.int4, + granularity: Granularity = PerAxis(0), + mapping_type: MappingType = MappingType.ASYMMETRIC, + use_fallback: bool = False, + ): + assert weight_dtype in [getattr(torch, f"int{i}") for i in range(1, 9)] + bit_width = _DTYPE_TO_BIT_WIDTH[weight_dtype] + + if isinstance(granularity, PerGroup): + group_size = granularity.group_size + elif isinstance(granularity, PerAxis): + assert granularity.axis == 0 + group_size = -1 + else: + raise ValueError(f"Unsupported granularity: {granularity}") + + self.bit_width = bit_width + self.group_size = group_size + self.use_fallback = use_fallback + self.mapping_type = mapping_type + + def quantize(self, model: nn.Module) -> nn.Module: + _replace_embedding_with_quantized_embedding( + model, + kwargs={ + "group_size": self.group_size, + "bit_width": self.bit_width, + "use_fallback": self.use_fallback, + "mapping_type": self.mapping_type, + }, + ) + return model + + +def _get_fqns_with_filter( + module: nn.Module, + filter_fn: Callable[Tuple[str, nn.Module], bool], + fqn: str, + fqns: List[str], +): + for name, child in module.named_children(): + child_fqn = f"{fqn}.{name}" if fqn != "" else name + if filter_fn(child, child_fqn): + fqns.append(child_fqn) + else: + _get_fqns_with_filter(child, filter_fn, child_fqn, fqns) + + +def get_fqns_with_filter( + module: nn.Module, filter_fn: Callable[Tuple[str, nn.Module], bool] +) -> List[str]: + fqns = [] + _get_fqns_with_filter(module, filter_fn, "", fqns) + return fqns + + +class QuantizedLinear(nn.Module): + def __init__(self, packed_weight, n, k, group_size, bit_width, bias): + super().__init__() + self.register_buffer("packed_weight", packed_weight) + self.n = n + self.k = k + self.group_size = group_size + self.bit_width = bit_width + self.bias = bias + + def _forward_2d(self, x): + assert x.dim() == 2 + m, k = x.shape + assert k == self.k + return getattr( + torch.ops.torchao, f"_linear_8bit_act_{self.bit_width}bit_weight" + )(x, self.packed_weight, self.group_size, self.n, self.k) + + def forward(self, x): + if x.dim() == 2: + res = self._forward_2d(x) + else: + assert x.dim() >= 3 + lead_shape = x.shape[0:-2] + m, k = x.shape[-2], x.shape[-1] + assert k == self.k + res = self._forward_2d(x.reshape(-1, k)) + res = res.reshape(*lead_shape, m, self.n) + + if self.bias is not None: + res = res + self.bias + return res + + +def quantized_linear_from_aqt( + weight: Optional[torch.Tensor], bias: Optional[torch.Tensor] +): + n, k = weight.shape + group_size = weight.tensor_impl.get_layout().group_size + bit_width = weight.tensor_impl.get_layout().bit_width + packed_weight = weight.tensor_impl.packed_weight + if weight.tensor_impl.get_layout().has_bias: + assert bias is None + return QuantizedLinear(packed_weight, n, k, group_size, bit_width, bias) + + +def replace_linear_tensor_subclass_with_module(module: nn.Module): + assert not isinstance(module, nn.Linear) + for name, child in module.named_children(): + if not isinstance(child, nn.Linear): + replace_linear_tensor_subclass_with_module(child) + else: + if not isinstance(child.weight, AffineQuantizedTensor): + continue + if not isinstance( + child.weight.tensor_impl.get_layout(), + PackedLinearInt8DynamicActivationIntxWeightLayout, + ): + continue + if child.weight.tensor_impl.get_layout().target == Target.ATEN: + continue + setattr(module, name, quantized_linear_from_aqt(child.weight, child.bias)) + + +class SharedEmbeddingQuantizer: + def __init__( + self, + weight_dtype: torch.dtype = torch.int4, + granularity: Granularity = PerAxis(0), + mapping_type: MappingType = MappingType.ASYMMETRIC, + ): + self.weight_dtype = weight_dtype + self.granularity = granularity + self.mapping_type = mapping_type + + def quantize( + self, + model: nn.Module, + embedding_to_unembedding: Optional[Mapping[str, str]] = None, + ): + embedding_fqns = get_fqns_with_filter( + model, lambda m, fqn: isinstance(m, nn.Embedding) + ) + linear_fqns = get_fqns_with_filter( + model, lambda m, fqn: isinstance(m, nn.Linear) + ) + state_dict = model.state_dict() + + # If embedding_to_unembedding is not provided, automatically detect shared embeddings and unembeddings + if embedding_to_unembedding is None: + embedding_to_unembedding = {} + for embedding_fqn in embedding_fqns: + embedding_w = state_dict[embedding_fqn + ".weight"] + for linear_fqn in linear_fqns: + linear_w = state_dict[linear_fqn + ".weight"] + if embedding_w.shape == linear_w.shape and torch.allclose( + embedding_w, linear_w + ): + print( + f"Found shared embedding {embedding_fqn} and unembedding {linear_fqn}" + ) + if embedding_fqn not in embedding_to_unembedding: + embedding_to_unembedding[embedding_fqn] = linear_fqn + else: + raise ValueError( + f"Found multiple candidate unembeddings ({embedding_to_unembedding[embedding_fqn]}, {linear_fqn}) for embedding {embedding_fqn}. This is not supported yet. Please explicitly define the input embedding_to_unembedding." + ) + + # Construct reverse mapping + unembedding_to_embedding = {} + for v, k in embedding_to_unembedding.items(): + if k not in unembedding_to_embedding: + unembedding_to_embedding[k] = v + else: + raise ValueError( + f"Found multiple candidate embeddings ({unembedding_to_embedding[k]}, {v}) for unembedding {k}. This is not supported yet." + ) + + # Check that embeddings are shared, embeddings are embeddings, and unembeddings are linear ops + for embedding_fqn, unembedding_fqn in embedding_to_unembedding.items(): + assert embedding_fqn in embedding_fqns, ( + f"Embedding {embedding_fqn} is not found in model" + ) + assert unembedding_fqn in linear_fqns, ( + f"Unembedding {unembedding_fqn} is not found in model" + ) + assert torch.allclose( + state_dict[embedding_fqn + ".weight"], + state_dict[unembedding_fqn + ".weight"], + ), ( + f"Embedding {embedding_fqn} does not share weights with unembedding {unembedding_fqn}" + ) + + # Quantize unembeddings + quantize_( + model, + Int8DynamicActivationIntxWeightConfig_NonExperimental( + weight_dtype=self.weight_dtype, + weight_granularity=self.granularity, + weight_mapping_type=self.mapping_type, + # Only universal layout is supported for shared embedding + layout=PackedLinearInt8DynamicActivationIntxWeightLayout( + target="universal" + ), + ), + filter_fn=lambda m, fqn: isinstance(m, nn.Linear) + and fqn in list(embedding_to_unembedding.values()), + ) + + embedding_fqn_to_quantized_unembedding = {} + for fqn, t in model.state_dict().items(): + if ( + fqn.endswith(".weight") + and fqn[: -len(".weight")] in unembedding_to_embedding + ): + embedding_fqn = unembedding_to_embedding[fqn[: -len(".weight")]] + embedding_fqn_to_quantized_unembedding[embedding_fqn] = t + + _replace_embedding_with_quantized_embedding( + model, + kwargs={ + "embedding_fqn_to_quantized_unembedding": embedding_fqn_to_quantized_unembedding, + }, + ) + + # Remove subclasses. Otherwise there are two packed_weight objects in exported model, + # even though they have the same id in eager mode + replace_linear_tensor_subclass_with_module(model) + + +def _quantize( + vals: torch.Tensor, group_size: int, nbit: int, has_weight_zeros: bool, signed=True +): + assert nbit >= 1 and nbit <= 8 + if signed: + qmin = -(1 << (nbit - 1)) + qmax = (1 << (nbit - 1)) - 1 + else: + qmin = 0 + qmax = (1 << nbit) - 1 + + n, k = vals.shape + vals = vals.reshape(-1, group_size) + vmins, _ = torch.min(vals, axis=1) + vmaxs, _ = torch.max(vals, axis=1) + group_scales = (vmaxs - vmins) / (qmax - qmin) + + if not has_weight_zeros: + group_zeros = torch.zeros_like(group_scales) + else: + group_zeros = qmin - torch.round(vmins / group_scales) + + vals = vals.reshape(n, k) + group_scales = group_scales.reshape(n, -1) + group_zeros = group_zeros.reshape(n, -1) + + group_qvals = quantize_per_channel_group( + input=vals, + scales=group_scales, + zero_points=group_zeros, + quant_min=qmin, + quant_max=qmax, + dtype=torch.int8 if signed else torch.uint8, + group_size=group_size, + ) + + if not has_weight_zeros: + group_zeros = None + + return group_qvals, group_scales, group_zeros + + +class UIntxWeightOnlyQuantizedLinear(nn.Module): + def __init__( + self, + pack_weight_op, + linear_op, + ): + super().__init__() + self._pack_weights_op = pack_weight_op + self._linear_op = linear_op + + def quantize_and_pack_weights(self, weights, nbit, group_size): + self.nbit = nbit + self.group_size = group_size + + weight_qvals, weight_scales, weight_zeros = _quantize( + weights, self.group_size, self.nbit, has_weight_zeros=True, signed=False + ) + weight_scales = torch.transpose_copy(weight_scales, 1, 0) + weight_zeros = torch.transpose_copy(weight_zeros, 1, 0) + weight_zeros = -weight_zeros * weight_scales + self.weight_scales = nn.Parameter(weight_scales, requires_grad=False) + self.weight_zeros = nn.Parameter(weight_zeros, requires_grad=False) + packed_weights = self._pack_weights_op(weight_qvals.cpu()).to(device="mps") + self.packed_weights = nn.Parameter(packed_weights, requires_grad=False) + + def forward(self, x): + assert x.dim() >= 2 + if x.dim() == 2: + return self._linear_op( + x, + self.packed_weights, + self.group_size, + self.weight_scales, + self.weight_zeros, + ) + + lead_shape = x.shape[0:-1] + k = x.shape[-1] + n = self.weight_scales.shape[1] + return self._linear_op( + x.reshape(-1, k), + self.packed_weights, + self.group_size, + self.weight_scales, + self.weight_zeros, + ).reshape(*lead_shape, n) + + +# TODO(mcandales): Consolidate with _replace_linear_with_quantized_linear +def _replace_linear_with_quantized_linear_mps(module: nn.Module, kwargs={}): + group_size = kwargs["group_size"] + nbit = kwargs["nbit"] + + assert not isinstance(module, nn.Linear) + assert nbit >= 1 and nbit <= 7 + + for name, child in module.named_children(): + if not isinstance(child, nn.Linear): + _replace_linear_with_quantized_linear_mps(child, kwargs) + else: + assert child.bias is None + qlinear = UIntxWeightOnlyQuantizedLinear( + pack_weight_op=getattr(torch.ops.torchao, f"_pack_weight_{nbit}bit"), + linear_op=getattr( + torch.ops.torchao, f"_linear_fp_act_{nbit}bit_weight" + ), + ) + setattr(module, name, qlinear) + qlinear.quantize_and_pack_weights(child.weight, nbit, group_size) + + +class UIntxWeightOnlyLinearQuantizer: + def __init__( + self, + device, + precision, + *, + bitwidth: Optional[int] = None, + groupsize: Optional[int] = None, + ): + if device != "mps": + raise NotImplementedError( + "Only device=mps is currently supported in UIntxWeightOnlyLinearQuantizer" + ) + else: + self.device = device + + if precision not in [torch.float32, torch.float16, torch.bfloat16]: + raise ValueError( + "Only precisions float32, float16 & bfloat16 are supported in UIntxWeightOnlyLinearQuantizer" + ) + else: + self.precision = precision + + if bitwidth is None: + bitwidth = 4 + logger.warning(f"bitwidth not specified, defaulting to {bitwidth}.") + if bitwidth not in range(1, 8): + raise ValueError( + "Only bitwidts 1 to 7 are supported in UIntxWeightOnlyLinearQuantizer" + ) + else: + self.bitwidth = bitwidth + + if groupsize is None: + groupsize = 128 + logger.warning(f"groupsize not specified, defaulting to {groupsize}.") + if groupsize not in [32, 64, 128, 256]: + raise ValueError( + "Only groupsizes 32, 64, 128 & 256 are supported in UIntxWeightOnlyLinearQuantizer" + ) + else: + self.groupsize = groupsize + + def quantize(self, model: nn.Module) -> nn.Module: + model = model.to(self.device).to(self.precision) + _replace_linear_with_quantized_linear_mps( + model, + kwargs={ + "group_size": self.groupsize, + "nbit": self.bitwidth, + }, + ) + return model diff --git a/lib/python3.12/site-packages/torchao/experimental/quant_passes.py b/lib/python3.12/site-packages/torchao/experimental/quant_passes.py new file mode 100644 index 0000000000000000000000000000000000000000..13a0a755fb86d09ce235eb9f76029c3eea9c2b63 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/experimental/quant_passes.py @@ -0,0 +1,318 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + + +import itertools +from collections import defaultdict +from typing import Callable, Optional + +import torch + +# import this for pt2e_quant.dequantize_affine op definition +# should be removed after removing dep on `torch._export.passes.constant_folding` +import torch.ao.quantization.pt2e._affine_quantization # noqa: F401 + +# TODO: remove dependency on ConstantFolder +from torch._export.passes.constant_folding import ( + ConstantFolder, + replace_node_with_constant, +) +from torch.fx import subgraph_rewriter + + +def constant_fold( + gm: torch.fx.GraphModule, + constraint_fn: Optional[Callable[[torch.fx.Node], bool]] = None, + skip_constructors: bool = False, +): + with torch.utils._python_dispatch._disable_current_modes(): + # The ConstantFolder has a bug where it throws if dequantize_affine is not defined + # TODO: fix upstream + try: + getattr(torch.ops.torchao, "dequantize_affine") + except AttributeError: + setattr(torch.ops.torchao, "dequantize_affine", None) + + cf = ConstantFolder(gm, skip_constructors) + cf.run() + + for node, constant in cf.node_replacements.items(): + if constraint_fn is not None and not constraint_fn(node): + continue + replace_node_with_constant(gm, node, constant) + + erased_params = [] + # Get all attr users by looking up the graph instead from node.users, because in this case + # _tensor_constant0 and _tensor_constant0_1 are actually refereing to the same tensor. + + # opcode name target args kwargs + # ------------- ------------------- ---------------- --------------------------- -------- + # placeholder arg0_1 arg0 () {} + # get_attr _tensor_constant0 state () {} + # call_function add aten.add.Tensor (arg0_1, _tensor_constant0) {} + # get_attr _tensor_constant0_1 state () {} + # call_function add_ aten.add_.Tensor (_tensor_constant0_1, 1) {} + # output output output ([add],) {} + + get_attr_node_users = defaultdict(list) + for node in gm.graph.nodes: + if node.op == "get_attr": + get_attr_node_users[node.target].extend(node.users.keys()) + for node in gm.graph.find_nodes(op="get_attr"): + if node.op == "get_attr" and len(get_attr_node_users[node.target]) == 0: + if hasattr(gm, node.target): + delattr(gm, node.target) + erased_params.append(node) + for node in erased_params: + gm.graph.erase_node(node) + + gm.graph.eliminate_dead_code() + gm.graph.lint() + gm.recompile() + + +def _get_q_dq_linear_patterns_replacements_and_filters( + weight_bit_width, has_weight_zeros, target +): + glbs = globals() + glbs["weight_bit_width"] = weight_bit_width + glbs["target"] = target + glbs["w_quant_min"] = -(1 << (weight_bit_width - 1)) + glbs["w_quant_max"] = (1 << (weight_bit_width - 1)) - 1 + glbs["a_target_dtype"] = torch.int8 + glbs["a_quant_min"] = None + glbs["a_quant_max"] = None + glbs["a_mapping_type"] = "ASYMMETRIC" + glbs["a_scale_dtype"] = torch.float32 + glbs["a_eps"] = None + + lcls = {} + + pattern_str = f""" +def pattern( + a, a_block_size, a_zero_point_dtype, + w_int_data, w_block_size, w_scale, w_zero_point, w_target_dtype, + bias): + a_scale, a_zero_point = torch.ops.torchao.choose_qparams_affine.default( + a, + a_mapping_type, + a_block_size, + a_target_dtype, + a_quant_min, + a_quant_max, + a_eps, + a_scale_dtype, + a_zero_point_dtype, + ) + a_int_data = torch.ops.torchao.quantize_affine.default( + a, a_block_size, a_scale, a_zero_point, a_target_dtype, a_quant_min, a_quant_max, + ) + dq_a = torch.ops.torchao.dequantize_affine.default( + a_int_data, a_block_size, a_scale, a_zero_point, a_target_dtype, a_quant_min, a_quant_max + ) + dq_w = torch.ops.torchao.dequantize_affine.default( + w_int_data, + w_block_size, + w_scale, + w_zero_point, + w_target_dtype, + w_quant_min, + w_quant_max, + {"'INT'" if has_weight_zeros else "'NONE'"} + ) + return torch.ops.aten.linear.default(dq_a, dq_w, bias) +""" + exec(pattern_str, glbs, lcls) + pattern = lcls["pattern"] + + replacement_str = f""" +def replacement( + a, a_block_size, a_zero_point_dtype, + w_int_data, w_block_size, w_scale, w_zero_point, w_target_dtype, + bias,): + n = w_int_data.size(0) + k = a_block_size[-1] + group_size = w_block_size[-1] + out_shape = a.shape[:-1] + (n,) + packed_weight = getattr( + torch.ops.torchao, + f"_pack_8bit_act_{weight_bit_width}bit_weight", + )( + w_int_data.to(torch.int8), + w_scale.reshape(-1), + {"w_zero_point.reshape(-1).to(torch.int8)" if has_weight_zeros else "None"}, + group_size, + bias, + target, + ) + return getattr( + torch.ops.torchao, f"_linear_8bit_act_{weight_bit_width}bit_weight" + )(a.reshape(-1, k), packed_weight, group_size, n, k).reshape(out_shape) +""" + + exec(replacement_str, glbs, lcls) + replacement = lcls["replacement"] + + def match_filter(match, x, y): + def get_val(name): + node = [n for n in match.nodes_map if n.name == name][0] + return match.nodes_map[node] + + int_types = [torch.int8, torch.int16, torch.int32, torch.int64] + + a_zero_point_dtype = get_val("a_zero_point_dtype") + if a_zero_point_dtype not in int_types: + return False + + # We only want a_block_size with shape [1, ..., 1, k] + a_block_size = get_val("a_block_size") + for d in a_block_size[0:-1]: + if d != 1: + print("a_block_size not [1, ..., 1, k]") + return False + + # We only want w_block_size with shape [1, group_size] + w_block_size = get_val("w_block_size") + if len(w_block_size) != 2 or w_block_size[0] != 1: + return False + + return True + + return pattern, replacement, match_filter + + +def replace_q_dq_patterns_with_quantized_linear_ops_pass( + ep: torch.export.ExportedProgram, + target=None, +) -> torch.export.ExportedProgram: + """ + This replaces Q/DQ patterns with torchao quantized linear ops. + It is intended for converting Q/DQ nodes exported with QDQLayout to using + the lowbit quantized linear ops. + """ + # TODO: figure out how to do this with dynamic_shapes (not saved on EP for easy re-export) + # See https://fb.workplace.com/groups/1028545332188949/permalink/1185289956514485/ + assert len(ep.range_constraints) == 0, ( + "ExportedProgram with range constraints are not supported" + ) + + # ep.module() unlifts the weight inputs, which we need for constant folding + gm = ep.module() + for weight_bit_width, has_weight_zeros in itertools.product( + range(1, 9), [True, False] + ): + pattern, replacement, match_filter = ( + _get_q_dq_linear_patterns_replacements_and_filters( + weight_bit_width, has_weight_zeros, target + ) + ) + subgraph_rewriter.replace_pattern_with_filters( + gm, pattern, replacement, match_filters=[match_filter] + ) + + # Constant fold evaluates and removes the packing ops + constant_fold(gm) + + # Re-export + return torch.export.export(gm, *ep.example_inputs) + + +def _get_q_dq_embedding_patterns_replacements_and_filters( + weight_bit_width, +): + w_quant_min = -(1 << (weight_bit_width - 1)) + w_quant_max = (1 << (weight_bit_width - 1)) - 1 + w_target_dtype = torch.int8 + + def pattern( + indices, + w_int_data, + w_block_size, + w_scale, + w_zero_point, + ): + dq_w = torch.ops.torchao.dequantize_affine.default( + w_int_data, + w_block_size, + w_scale, + w_zero_point, + w_target_dtype, + w_quant_min, + w_quant_max, + ) + return torch.ops.aten.embedding.default(dq_w, indices) + + def replacement( + indices, + w_int_data, + w_block_size, + w_scale, + w_zero_point, + ): + num_embeddings, embedding_dim = w_int_data.size() + packed_weight_qvals = getattr( + torch.ops.torchao, f"_pack_embedding_{weight_bit_width}bit" + )(w_int_data) + out_shape = indices.shape + (embedding_dim,) + group_size = w_block_size[-1] + n_groups = embedding_dim // group_size + w_scale = w_scale.reshape(-1, n_groups) + w_zero_point = w_zero_point.reshape(-1, n_groups) + return getattr(torch.ops.torchao, f"_embedding_{weight_bit_width}bit")( + packed_weight_qvals, + num_embeddings, + embedding_dim, + w_scale, + w_zero_point, + indices.reshape(-1), + ).reshape(out_shape) + + def match_filter(match, x, y): + def get_val(name): + node = [n for n in match.nodes_map if n.name == name][0] + return match.nodes_map[node] + + # We only want w_block_size with shape [1, group_size] + w_block_size = get_val("w_block_size") + if len(w_block_size) != 2 or w_block_size[0] != 1: + return False + + return True + + return pattern, replacement, match_filter + + +def replace_q_dq_patterns_with_quantized_embedding_ops_pass( + ep: torch.export.ExportedProgram, +) -> torch.export.ExportedProgram: + """ + This replaces Q/DQ patterns with torchao quantized embedding ops. + It is intended for converting Q/DQ nodes exported with QDQLayout to using + the lowbit quantized embedding ops. + """ + # TODO: figure out how to do this with dynamic_shapes (not saved on EP for easy re-export) + # See https://fb.workplace.com/groups/1028545332188949/permalink/1185289956514485/ + assert len(ep.range_constraints) == 0, ( + "ExportedProgram with range constraints are not supported" + ) + + # ep.module() unlifts the weight inputs, which we need for constant folding + gm = ep.module() + for weight_bit_width in range(1, 9): + pattern, replacement, match_filter = ( + _get_q_dq_embedding_patterns_replacements_and_filters( + weight_bit_width, + ) + ) + subgraph_rewriter.replace_pattern_with_filters( + gm, pattern, replacement, match_filters=[match_filter] + ) + + # Constant fold evaluates and removes the packing ops + constant_fold(gm) + + # Re-export + return torch.export.export(gm, *ep.example_inputs) diff --git a/lib/python3.12/site-packages/torchao/experimental/temp_build.py b/lib/python3.12/site-packages/torchao/experimental/temp_build.py new file mode 100644 index 0000000000000000000000000000000000000000..3195e24581ef85b60bb9f10dfb3e51f081e80bb2 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/experimental/temp_build.py @@ -0,0 +1,47 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import glob +import subprocess +import tempfile + +import torch + + +def cmake_build_torchao_ops(cmake_lists_path, temp_build_dir): + from distutils.sysconfig import get_python_lib + + print("Building torchao ops for ATen target") + cmake_prefix_path = get_python_lib() + subprocess.run( + [ + "cmake", + "-DCMAKE_PREFIX_PATH=" + cmake_prefix_path, + "-DCMAKE_INSTALL_PREFIX=" + temp_build_dir.name, + "-S " + cmake_lists_path, + "-B " + temp_build_dir.name, + ] + ) + subprocess.run( + [ + "cmake", + "--build", + temp_build_dir.name, + "-j 16", + "--target install", + "--config Release", + ] + ) + + +def temp_build_and_load_torchao_ops(cmake_lists_path): + temp_build_dir = tempfile.TemporaryDirectory() + cmake_build_torchao_ops(cmake_lists_path, temp_build_dir) + libs = glob.glob(f"{temp_build_dir.name}/lib/libtorchao_ops_aten.*") + libs = list(filter(lambda l: (l.endswith("so") or l.endswith("dylib")), libs)) + assert len(libs) == 1 + torch.ops.load_library(libs[0]) + print(f"TorchAO ops are loaded from {libs[0]}") diff --git a/lib/python3.12/site-packages/torchao/optim/__init__.py b/lib/python3.12/site-packages/torchao/optim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9a22507b4ce042018f3df1d48acb28763356391c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/optim/__init__.py @@ -0,0 +1,13 @@ +from .adam import Adam4bit, Adam8bit, AdamFp8, AdamW4bit, AdamW8bit, AdamWFp8, _AdamW +from .cpu_offload import CPUOffloadOptimizer + +__all__ = [ + "Adam4bit", + "Adam8bit", + "AdamFp8", + "AdamW4bit", + "AdamW8bit", + "AdamWFp8", + "_AdamW", + "CPUOffloadOptimizer", +] diff --git a/lib/python3.12/site-packages/torchao/optim/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/optim/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..185ce1b890beb2477fc26a474bf1cd35fd10c05e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/optim/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/optim/__pycache__/adam.cpython-312.pyc b/lib/python3.12/site-packages/torchao/optim/__pycache__/adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51d6409eaadd88d130fd1e0423c6a8b687c2db6e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/optim/__pycache__/adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/optim/__pycache__/cpu_offload.cpython-312.pyc b/lib/python3.12/site-packages/torchao/optim/__pycache__/cpu_offload.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..451d33867097c36e165edba8a4e6d6b31ed34dae Binary files /dev/null and b/lib/python3.12/site-packages/torchao/optim/__pycache__/cpu_offload.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/optim/__pycache__/quant_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/optim/__pycache__/quant_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a4d9093133654d5524f8f79f70b8871751a6723 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/optim/__pycache__/quant_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_4bit.cpython-312.pyc b/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_4bit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e628652e3acb45e1c6ade0f797b6b2db31f43ee6 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_4bit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_8bit.cpython-312.pyc b/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_8bit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..013b31c68a2596b29d43cbd5f80e83947670490e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_8bit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_fp8.cpython-312.pyc b/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dedbb600ba562c2ee6555cd430f830c53c2b7a94 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/optim/__pycache__/subclass_fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/optim/adam.py b/lib/python3.12/site-packages/torchao/optim/adam.py new file mode 100644 index 0000000000000000000000000000000000000000..cba86897d9237183b91a1b8170329bd79f395a78 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/optim/adam.py @@ -0,0 +1,403 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Optional + +import torch +from torch import Tensor +from torch.distributed._tensor import DTensor +from torch.optim import Optimizer + +from .quant_utils import _fp32_to_bf16_sr +from .subclass_4bit import OptimState4bit +from .subclass_8bit import OptimState8bit +from .subclass_fp8 import OptimStateFp8 + + +class _AdamBase(Optimizer): + def __init__( + self, + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + *, + block_size, + bf16_stochastic_round, + is_adamw, + ) -> None: + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) + if not 0.0 <= betas[1] < 1.0: + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) + defaults = dict( + lr=torch.tensor(lr), + betas=betas, + eps=eps, + weight_decay=weight_decay, + amsgrad=amsgrad, + ) + super().__init__(params, defaults) + self.block_size = block_size + self.bf16_stochastic_round = bf16_stochastic_round + self.is_adamw = is_adamw + + def __setstate__(self, state): + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("amsgrad", False) + + # bring your own function to create zero-filled subclass + @staticmethod + def _subclass_zeros(p: Tensor, signed: bool, block_size: int): + raise NotImplementedError + + def _new_buffer(self, p: Tensor, signed: bool): + local_p = p.to_local() if isinstance(p, DTensor) else p + + # follow bitsandbytes, only quantize tensors >= 4096 values + if local_p.numel() >= 4096 and local_p.numel() % self.block_size == 0: + out = self._subclass_zeros(local_p, signed, self.block_size) + else: + out = torch.zeros_like(local_p) + + # wrap subclass in DTensor as needed + # NOTE: local tensor may have different shapes across ranks. + # this happens when the 1st dim is not divisible by WORLD_SIZE. + # thus, we must supply shape (and stride) to DTensor.from_local() + if isinstance(p, DTensor): + out = DTensor.from_local( + local_tensor=out, + device_mesh=p.device_mesh, + placements=p.placements, + run_check=False, + shape=p.shape, + stride=p.stride(), + ) + + return out + + @torch.no_grad() + def step(self, closure=None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # for a given model, the number of different argument combinations to single_param_adam() is fixed. + # thus, it is safe to disable cache limit without the risk of always re-compiling. + with torch._dynamo.utils.disable_cache_limit(): + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad + if grad.is_sparse: + raise RuntimeError("Sparse gradient is not supported") + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = torch.tensor(0.0) + state["exp_avg"] = self._new_buffer(p, True) + state["exp_avg_sq"] = self._new_buffer(p, False) + if group["amsgrad"]: + state["max_exp_avg_sq"] = self._new_buffer(p, False) + + state["step"] += 1 + + if not isinstance(group["lr"], Tensor): + raise RuntimeError( + "lr was changed to a non-Tensor object. If you want to update lr, please use " + "optim.param_groups[0]['lr'].fill_(new_lr)" + ) + + # without calling p.detach(), torch.compile() will have issues with FSDP2 in some cases + # https://github.com/pytorch/ao/issues/652#issuecomment-2285040894 + # thus, by calling p.detach(), DTensor won't have .grad anymore, which is ok since we + # are passing grad separately anyway. + torch.compile(single_param_adam, fullgraph=True, dynamic=False)( + p.detach(), + grad, + state["step"], + state["exp_avg"], + state["exp_avg_sq"], + state.get("max_exp_avg_sq", None), + group["lr"], + group["betas"][0], + group["betas"][1], + group["weight_decay"], + group["eps"], + self.is_adamw, + self.bf16_stochastic_round and p.dtype is torch.bfloat16, + ) + + return loss + + +# this will work with any optim state tensor subclass that implements aten.lerp.Scalar and aten.copy_.default +# and param tensor subclass that implements aten.add_.Tensor, and aten.addcdiv_.default +def single_param_adam( + p: Tensor, + grad: Tensor, + step: Tensor, + exp_avg: Tensor, + exp_avg_sq: Tensor, + max_exp_avg_sq: Optional[Tensor], + lr: Tensor, + beta1: float, + beta2: float, + weight_decay: float, + eps: float, + IS_ADAMW: bool, + BF16_STOCHASTIC_ROUND: bool, +): + # compute in FP32 for accurate calculations + p_f32 = p.float() + grad_f32 = grad.float() + + if IS_ADAMW: + p_f32 = p_f32 - lr * weight_decay * p_f32 + else: + grad_f32 = grad_f32 + weight_decay * p_f32 + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + + # keep high precision copy for param update + exp_avg_f32 = exp_avg.float().lerp(grad_f32, 1 - beta1) + exp_avg_sq_f32 = exp_avg_sq.float().lerp(grad_f32.square(), 1 - beta2) + + exp_avg.copy_(exp_avg_f32) + exp_avg_sq.copy_(exp_avg_sq_f32) + + if max_exp_avg_sq is not None: + max_exp_avg_sq_f32 = torch.maximum(max_exp_avg_sq.float(), exp_avg_sq_f32) + max_exp_avg_sq.copy_(max_exp_avg_sq_f32) + denom = (max_exp_avg_sq_f32.sqrt() / bias_correction2.sqrt()) + eps + else: + denom = (exp_avg_sq_f32.sqrt() / bias_correction2.sqrt()) + eps + + p_f32 = p_f32 - lr * (exp_avg_f32 / bias_correction1) / denom + + if BF16_STOCHASTIC_ROUND: + p.copy_(_fp32_to_bf16_sr(p_f32)) + else: + p.copy_(p_f32) + + +class Adam8bit(_AdamBase): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + amsgrad=False, + *, + block_size=256, + bf16_stochastic_round=False, + ) -> None: + super().__init__( + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + block_size=block_size, + bf16_stochastic_round=bf16_stochastic_round, + is_adamw=False, + ) + + @staticmethod + def _subclass_zeros(p: Tensor, signed: bool, block_size: int): + return OptimState8bit.zeros(p.shape, signed, block_size, p.device) + + +class Adam4bit(_AdamBase): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + amsgrad=False, + *, + block_size=128, + bf16_stochastic_round=False, + ) -> None: + super().__init__( + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + block_size=block_size, + bf16_stochastic_round=bf16_stochastic_round, + is_adamw=False, + ) + + @staticmethod + def _subclass_zeros(p: Tensor, signed: bool, block_size: int): + return OptimState4bit.zeros(p.shape, signed, block_size, p.device) + + +class AdamFp8(_AdamBase): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + amsgrad=False, + *, + block_size=256, + bf16_stochastic_round=False, + ) -> None: + super().__init__( + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + block_size=block_size, + bf16_stochastic_round=bf16_stochastic_round, + is_adamw=False, + ) + + @staticmethod + def _subclass_zeros(p: Tensor, signed: bool, block_size: int): + return OptimStateFp8.zeros(p.shape, block_size, p.device) + + +class AdamW8bit(_AdamBase): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=1e-2, + amsgrad=False, + *, + block_size=256, + bf16_stochastic_round=False, + ) -> None: + super().__init__( + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + block_size=block_size, + bf16_stochastic_round=bf16_stochastic_round, + is_adamw=True, + ) + + @staticmethod + def _subclass_zeros(p: Tensor, signed: bool, block_size: int): + return OptimState8bit.zeros(p.shape, signed, block_size, p.device) + + +class AdamW4bit(_AdamBase): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=1e-2, + amsgrad=False, + *, + block_size=128, + bf16_stochastic_round=False, + ) -> None: + super().__init__( + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + block_size=block_size, + bf16_stochastic_round=bf16_stochastic_round, + is_adamw=True, + ) + + @staticmethod + def _subclass_zeros(p: Tensor, signed: bool, block_size: int): + return OptimState4bit.zeros(p.shape, signed, block_size, p.device) + + +class AdamWFp8(_AdamBase): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=1e-2, + amsgrad=False, + *, + block_size=256, + bf16_stochastic_round=False, + ) -> None: + super().__init__( + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + block_size=block_size, + bf16_stochastic_round=bf16_stochastic_round, + is_adamw=True, + ) + + @staticmethod + def _subclass_zeros(p: Tensor, signed: bool, block_size: int): + return OptimStateFp8.zeros(p.shape, block_size, p.device) + + +class _AdamW(_AdamBase): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=1e-2, + amsgrad=False, + *, + bf16_stochastic_round=False, + ) -> None: + """AdamW optimizer that supports quantized training (parameter is quantized). This optimizer should + only be used with torchao's quantized training.""" + super().__init__( + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + block_size=float("inf"), + bf16_stochastic_round=bf16_stochastic_round, + is_adamw=True, + ) diff --git a/lib/python3.12/site-packages/torchao/optim/cpu_offload.py b/lib/python3.12/site-packages/torchao/optim/cpu_offload.py new file mode 100644 index 0000000000000000000000000000000000000000..cca55749db3db70890c2e57c8ca00849e409463c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/optim/cpu_offload.py @@ -0,0 +1,184 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Type + +import torch +from torch.optim.optimizer import Optimizer, ParamsT + +from torchao.utils import TORCH_VERSION_AT_LEAST_2_4, get_available_devices + + +# NOTE: We make this inherit Optimizer so it works with PyTorch's built-in LR +# schedulers. (those schedulers specifically check for instances of Optimizer). +# However, it won't behave exactly like Optimizer e.g. we don't call +# Optimizer.__init__(), there is no self.defaults. +class CPUOffloadOptimizer(Optimizer): + def __init__( + self, + params: ParamsT, + optimizer_class: Type[Optimizer] = torch.optim.AdamW, + *, + offload_gradients: bool = False, + minimal_size: int = 4096, + **kwargs, + ) -> None: + """Offload optimizer to CPU for single-GPU training. This will reduce GPU memory by the size of optimizer state. + Optimizer step will be done on CPU. + + Args + params: a list of parameters or parameter groups. + optimizer_class: constructor of the base optimizer. Defaults to :class:`torch.optim.AdamW`. + offload_gradients: free GPU gradients once they are moved to CPU. Not compatible with gradient accumulation. + minimal_size: tensors smaller than this are kept on the GPU, to avoid excessively many small transfers. + kwargs: other keyword arguments to be passed to the base optimizer e.g. `lr`, `weight_decay`. + """ + # default to fused CPU AdamW + if ( + optimizer_class is torch.optim.AdamW + and TORCH_VERSION_AT_LEAST_2_4 + and "fused" not in kwargs + ): + kwargs.update(fused=True) + + param_groups = list(params) + if len(param_groups) == 0: + raise ValueError("optimizer got an empty parameter list") + if not isinstance(param_groups[0], dict): + param_groups = [{"params": param_groups}] + + # any parameter smaller than minimal size will be handled by the on-device optimizer d_opt + self.minimal_size = minimal_size + self.d_opt = None + self.d_param_groups = [] + + self.param_d2h_map = dict() + self.optim_dict = dict() + self.device = get_available_devices()[-1] + assert self.device in [ + "cuda", + "xpu", + ], "CPU Offload currently only supports CUDA & XPU" + self.stream = getattr(torch, self.device).Stream() + + # the queue maintains the order which param we should do optim step on first. + self.queue = dict() + + def backward_hook(p_device): + if p_device.grad is not None: + p_host = self.param_d2h_map[p_device] + + # make sure backward for this param finishes + self.stream.wait_stream(getattr(torch, self.device).current_stream()) + with getattr(torch, self.device).stream(self.stream): + p_host.grad.copy_(p_device.grad, non_blocking=True) + + # we rely on CPython implementation of dictionary, which preserves insertion order. + # if a param is added again (e.g. due to gradient accumulation), it is moved to the + # end of the queue by removing and inserting it again. + if p_device in self.queue: + del self.queue[p_device] + self.queue[p_device] = self.stream.record_event() + + # deallocate DEVICE gradients once D2H transfer finishes. + if offload_gradients: + p_device.grad.record_stream(self.stream) + p_device.grad = None + + for param_group in param_groups: + params = param_group.pop("params") + retained_params = [] + + for p_device in params: + if not p_device.requires_grad: + continue + + if p_device.numel() < self.minimal_size: + retained_params.append(p_device) + continue + + # pre-allocate CPU params and grads + p_host = torch.empty_like(p_device, device="cpu", pin_memory=True) + p_host.grad = torch.empty_like(p_host, pin_memory=True) + + p_host.copy_(p_device.detach(), non_blocking=True) + self.param_d2h_map[p_device] = p_host + + p_device.register_post_accumulate_grad_hook(backward_hook) + self.optim_dict[p_device] = optimizer_class( + [{"params": p_host, **param_group}], **kwargs + ) + + if len(retained_params) > 0: + self.d_param_groups.append({"params": retained_params, **param_group}) + + if len(self.d_param_groups) > 0: + self.d_opt = optimizer_class(self.d_param_groups, **kwargs) + + @torch.no_grad() + def step(self, closure=None): + loss = None + if closure is not None: + loss = closure() + + # handle small parameters on the GPU, in parallel with the CPU calls below + if self.d_opt is not None: + self.d_opt.step() + + for p_device, grad_d2h_event in self.queue.items(): + grad_d2h_event.synchronize() + self.optim_dict[p_device].step() + + # submit more job to self.stream. it guarantees that we only start + # moving param H2D once all backwards finish, since self.stream + # will wait for current_stream when moving grad D2H. + p_host = self.param_d2h_map[p_device] + with getattr(torch, self.device).stream(self.stream): + p_device.copy_(p_host, non_blocking=True) + + # make sure param H2D finishes before the next forward pass + self.stream.synchronize() + self.queue.clear() + return loss + + def zero_grad(self, set_to_none=True): + assert set_to_none + + # only clear DEVICE grad. CPU grad will always be overwritten by DEVICE grad. + for p_device in self.param_d2h_map.keys(): + p_device.grad = None + + if self.d_opt is not None: + self.d_opt.zero_grad(set_to_none=set_to_none) + + @property + def param_groups(self): + # each param group will only has 1 parameter + # TODO: we might want to return the original param_groups instead. + return sum( + (optim.param_groups for optim in self.optim_dict.values()), + start=self.d_param_groups, + ) + + def state_dict(self): + state_dict = { + "offloaded": [optim.state_dict() for optim in self.optim_dict.values()] + } + if self.d_opt: + state_dict["on-device"] = self.d_opt.state_dict() + return state_dict + + def load_state_dict(self, state_dict): + for optim, optim_state_dict in zip( + self.optim_dict.values(), state_dict["offloaded"] + ): + optim.load_state_dict(optim_state_dict) + + if self.d_opt: + self.d_opt.load_state_dict(state_dict["on-device"]) + elif "on-device" in state_dict: + raise ValueError( + "loaded state dict has a 'on-device' parameter group not present in the optimizer" + ) diff --git a/lib/python3.12/site-packages/torchao/optim/quant_utils.py b/lib/python3.12/site-packages/torchao/optim/quant_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a4035fde1c33c5c7fb44bdb7426e19be62144bcd --- /dev/null +++ b/lib/python3.12/site-packages/torchao/optim/quant_utils.py @@ -0,0 +1,145 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +from torch import Tensor + + +# https://github.com/TimDettmers/bitsandbytes/blob/dada530149212d64d4b69534716202659ef37ec8/bitsandbytes/functional.py#L339-L391 +# NOTE: zero padding is removed so this function can work with 4-bit qmap +def create_dynamic_map(signed=True, max_exponent_bits=7, total_bits=8): + """ + Creates the dynamic quantiztion map. + + The dynamic data type is made up of a dynamic exponent and + fraction. As the exponent increase from 0 to -7 the number + of bits available for the fraction shrinks. + + This is a generalization of the dynamic type where a certain + number of the bits and be reserved for the linear quantization + region (the fraction). n determines the maximum number of + exponent bits. + + For more details see + (8-Bit Approximations for Parallelism in Deep Learning)[https://arxiv.org/abs/1511.04561] + """ + + data = [] + # these are additional items that come from the case + # where all the exponent bits are zero and no + # indicator bit is present + non_sign_bits = total_bits - (1 if signed else 1) + additional_items = 2 ** (non_sign_bits - max_exponent_bits) - 1 + for i in range(max_exponent_bits): + fraction_items = int( + 2 ** (i + non_sign_bits - max_exponent_bits) + 1 + if signed + else 2 ** (i + non_sign_bits - max_exponent_bits + 1) + 1, + ) + boundaries = torch.linspace(0.1, 1, fraction_items) + means = (boundaries[:-1] + boundaries[1:]) / 2.0 + data += ((10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() + if signed: + data += (-(10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() + + if additional_items > 0: + boundaries = torch.linspace(0.1, 1, additional_items + 1) + means = (boundaries[:-1] + boundaries[1:]) / 2.0 + data += ((10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() + if signed: + data += (-(10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() + + data.append(0) + data.append(1.0) + + assert len(data) == 2**total_bits + + data.sort() + return data + + +def scale_tensor(input: Tensor, block_size: int): + """Scale tensor so that max(abs(input)) = 1""" + shape = input.shape + + # section 2.1 from https://arxiv.org/abs/2110.02861 + input = input.view(-1, block_size) + scale = input.abs().amax(-1).clip(1e-12) + input = input / scale.view(-1, 1) + return input.view(shape), scale + + +def quantize_8bit_with_qmap(input: Tensor, qmap: Tensor): + # GPU-friendly binary search + # https://blog.demofox.org/2017/06/20/simd-gpu-friendly-branchless-binary-search/ + codes = torch.where(input >= qmap[128], 128, 0) + codes += torch.where(input >= qmap[codes + 64], 64, 0) + codes += torch.where(input >= qmap[codes + 32], 32, 0) + codes += torch.where(input >= qmap[codes + 16], 16, 0) + codes += torch.where(input >= qmap[codes + 8], 8, 0) + codes += torch.where(input >= qmap[codes + 4], 4, 0) + codes += torch.where(input >= qmap[codes + 2], 2, 0) + codes += torch.where(input >= qmap[codes + 1], 1, 0) + + # rounding + codes_up = (codes + 1).clip(max=255) + val_down = qmap[codes] + val_up = qmap[codes_up] + residual = input - val_down + codes = torch.where(residual >= (val_up - val_down) * 0.5, codes_up, codes) + + return codes.to(torch.uint8) + + +def quantize_4bit_with_qmap(input: Tensor, qmap: Tensor): + # GPU-friendly binary search + # https://blog.demofox.org/2017/06/20/simd-gpu-friendly-branchless-binary-search/ + codes = torch.where(input >= qmap[8], 8, 0) + codes += torch.where(input >= qmap[codes + 4], 4, 0) + codes += torch.where(input >= qmap[codes + 2], 2, 0) + codes += torch.where(input >= qmap[codes + 1], 1, 0) + + # rounding + codes_up = (codes + 1).clip(max=15) + val_down = qmap[codes] + val_up = qmap[codes_up] + residual = input - val_down + codes = torch.where(residual >= (val_up - val_down) * 0.5, codes_up, codes) + + return codes.to(torch.uint8) + + +def dequant_with_qmap(codes: Tensor, qmap: Tensor, scale: Tensor): + # torch.compile() cannot use uint8 as index + out = qmap[codes.int()].view(scale.shape[0], -1) * scale.view(-1, 1) + return out.view(codes.shape) + + +def _fp32_to_bf16_sr(x_f32: Tensor) -> Tensor: + # For an FP32 number [a31, ..., a16, a15, ..., a0] to be converted to BF16 + # - Round towards zero: [a31, ..., a16, 0, ..., 0] + # - Round away from zero: [a31, ..., a16+1, 0, ..., 0] + # (since the value can be negative, we use round towards/away from zero instead of round up/down) + # + # For stochastic rounding, we round away from zero with the probability of + # [a15, ..., a0] / 2^16, where the bit pattern [a15, ..., a0] is interpreted as uint16 + # + # we have to use int32 since most arithmetic ops are not implemented for uint32/int16/uint16 + rand_16bit = torch.randint( + 0, 1 << 16, x_f32.shape, device=x_f32.device, dtype=torch.int32 + ) + x_f32_bits = x_f32.view(torch.int32) + x_fraction = x_f32_bits & 0xFFFF # lower 16 bits + x_bf16_towards_zero = x_f32_bits & 0xFFFF0000 # upper 16 bits + + x_f32_bits = torch.where( + rand_16bit < x_fraction, # this is True with the probability of p_fraction + x_bf16_towards_zero + + 0x10000, # this might overflow, which will result in UB due to signed integer + x_bf16_towards_zero, + ) + # alternative, slightly faster + # x_f32_bits = (x_f32_bits + rand_16bit) & 0xFFFF0000 + return x_f32_bits.view(torch.float32).bfloat16() diff --git a/lib/python3.12/site-packages/torchao/optim/subclass_4bit.py b/lib/python3.12/site-packages/torchao/optim/subclass_4bit.py new file mode 100644 index 0000000000000000000000000000000000000000..55d9bf356fe9fa2650539861f27f54683b10d0f2 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/optim/subclass_4bit.py @@ -0,0 +1,261 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import math + +import torch +from torch import Tensor +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_4, + TORCH_VERSION_AT_LEAST_2_5, + TorchAOBaseTensor, +) + +from .quant_utils import ( + create_dynamic_map, + dequant_with_qmap, + quantize_4bit_with_qmap, + scale_tensor, +) + +aten = torch.ops.aten +c10d_functional = torch.ops.c10d_functional +_c10d_functional = torch.ops._c10d_functional + +# https://github.com/thu-ml/low-bit-optimizers/blob/e3e2854728e498c2a606e3fdb88daa27ae94f9a6/lpmm/configs/2nd_moment_group_128.yml +# NOTE: power-1 is linear +# TODO: since QMAP_UNSIGNED is linear, perhaps doing affine quantize is faster? +QMAP_SIGNED = create_dynamic_map(True, 3, 4) +QMAP_UNSIGNED = torch.linspace(0, 1, 17)[1:].tolist() # no zero + + +class OptimState4bit(TorchAOBaseTensor): + tensor_attrs = ["codes", "scale", "qmap"] + + @staticmethod + def __new__(cls, codes: Tensor, scale: Tensor, qmap: Tensor, signed: bool, shape): + return Tensor._make_wrapper_subclass(cls, shape, device=codes.device) + + def __init__(self, codes: Tensor, scale: Tensor, qmap: Tensor, signed: bool, shape): + """Create quantized 4-bit optimizer state as proposed in https://arxiv.org/abs/2309.01507 + + Args + codes: quantized and packed 4-bit data stored as uint8. + scale: scale data for block-wise quantization. + qmap: lookup table that maps between quantized value (code) and float value. + signed: whether the tensor is signed or unsigned. + shape: shape of original float tensor. + + NOTE: To get block-wise scale, the original float tensor is first reshape to (-1, block_size). + Thus, the last dimension of the original float tensor is not necessarily divisible by block size. + Given `codes` and `scale`, `block_size` is calculated as `codes.numel() * 2 // scale.numel()`. + The extra `* 2` is because `codes` is 4-bit data packed in 8-bit storage. + """ + assert codes.dtype is torch.uint8 + assert codes.ndim == 1 # flattened buffer + assert scale.ndim == 1 + self.codes = codes + self.scale = scale + self.qmap = qmap + self.signed = signed + self._shape = shape + self.block_size = codes.numel() * 2 // scale.numel() + + def __tensor_flatten__(self): + return self.tensor_attrs, [self.signed, self._shape] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size=None, outer_stride=None + ): + return cls( + *[tensor_data_dict[name] for name in cls.tensor_attrs], *tensor_attributes + ) + + def dequantize(self, output_dtype=None): + codes = torch.stack([self.codes >> 4, self.codes & 0b1111], dim=-1) # unpack + float_data = dequant_with_qmap(codes, self.qmap, self.scale) + if output_dtype is not None: + float_data = float_data.to(output_dtype) + return float_data.view(self._shape) + + @classmethod + def zeros(cls, shape, signed: bool = True, block_size: int = 128, device=None): + shape = (shape,) if isinstance(shape, int) else shape + n_elems = math.prod(shape) + + codes = torch.zeros(n_elems // 2, dtype=torch.uint8, device=device) + scale = torch.zeros(n_elems // block_size, device=device) + qmap = torch.tensor(QMAP_SIGNED if signed else QMAP_UNSIGNED, device=device) + return cls(codes, scale, qmap, signed, shape) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(signed={self.signed}, block_size={self.block_size}, " + f"shape={tuple(self.shape)}, device={self.device}, requires_grad={self.requires_grad})" + ) + + +# in pre-2.4, calling .to(device, dtype) will not dispatch aten._to_copy.default when +# dtype is the same but device is different. thus, we must override .to() method instead. +if not TORCH_VERSION_AT_LEAST_2_4: + + def _to(self, *args, **kwargs): + # ignore other args/kwargs + device = kwargs.pop("device", None) + return OptimState4bit( + self.codes.to(device), + self.scale.to(device), + self.qmap.to(device), + self.signed, + self.shape, + ) + + OptimState4bit.to = _to + del _to # make sure to not re-use + + +@OptimState4bit.implements(aten.copy_.default) +def _(func, types, args, kwargs): + dst = args[0] + src = args[1] + + if isinstance(dst, OptimState4bit) and isinstance(src, OptimState4bit): + assert ( + dst.signed == src.signed + and dst.block_size == src.block_size + and dst._shape == src._shape + ) + dst.codes.copy_(src.codes) + dst.scale.copy_(src.scale) + # qmap should be the same, don't need to copy + + elif isinstance(dst, OptimState4bit): + scaled_src, scale = scale_tensor(src.view(-1), dst.block_size) + codes = quantize_4bit_with_qmap(scaled_src, dst.qmap) + dst.codes.copy_((codes[::2] << 4) | codes[1::2]) # packing + dst.scale.copy_(scale) + + else: + dst.copy_(src.dequantize()) + + return dst + + +@OptimState4bit.implements(aten._to_copy.default) +def _(func, types, args, kwargs): + # ignore dtype + device = kwargs.get("device", None) + out = OptimState4bit( + args[0].codes.to(device=device), + args[0].scale.to(device=device), + args[0].qmap.to(device=device), + args[0].signed, + args[0].shape, + ) + return return_and_correct_aliasing(func, args, kwargs, out) + + +@OptimState4bit.implements(aten.lerp.Scalar) +def _(func, types, args, kwargs): + args = [x.dequantize() if isinstance(x, OptimState4bit) else x for x in args] + return func(*args, **kwargs) + + +# this is needed for DTensor.from_local() and for flattening tensor +@OptimState4bit.implements(aten.view.default) +def _(func, types, args, kwargs): + x, shape = args + + if tuple(x.shape) == tuple(shape): + return OptimState4bit(x.codes, x.scale, x.qmap, x.signed, x._shape) + + if len(shape) == 1 and shape[0] == -1: + return OptimState4bit(x.codes, x.scale, x.qmap, x.signed, (x.numel(),)) + + raise ValueError( + f"{x.__class__.__name__} only supports .view() with same shape or shape=[-1]" + ) + + +@OptimState4bit.implements( + [ + # required by DTensor.full_tensor() + c10d_functional.all_gather_into_tensor.default, + _c10d_functional.all_gather_into_tensor.default, + c10d_functional.wait_tensor.default, + _c10d_functional.wait_tensor.default, + # required by torch.distributed.checkpoint.save + aten.detach.default, + ] +) +def _(func, types, args, kwargs): + x = args[0] + if not isinstance(x, OptimState4bit): + raise ValueError(f"expecting a OptimState4bit but found {type(x)}") + + codes = func(x.codes, *args[1:], **kwargs) + scale = func(x.scale, *args[1:], **kwargs) + + # adjust the first dim + shape = (x._shape[0] * codes.numel() // x.codes.numel(),) + x._shape[1:] + + # assume tensors from all ranks have the same signedness + return OptimState4bit(codes, scale, x.qmap.clone(), x.signed, shape) + + +# required by torch.distributed.checkpoint.save +# note that we don't actually implement pin memory for this tensor subclass +# (pin_memory argument is ignored in aten._to_copy) +@OptimState4bit.implements(aten.is_pinned.default) +def _(func, types, args, kwargs): + return ( + args[0].codes.is_pinned() + and args[0].scale.is_pinned() + and args[0].qmap.is_pinned() + ) + + +# required by torch.distributed.checkpoint.load when world size changes i.e. re-sharding +@OptimState4bit.implements(aten.slice.Tensor) +def _(func, types, args, kwargs): + x, dim, start, end = args[:4] + step = args[4] if len(args) > 4 else 1 + + # input validation + if dim != 0: + raise ValueError("Only support aten.slice along the first dim") + if step != 1: + raise ValueError("Only support aten.slice with step=1") + + block_size = x.block_size + stride = math.prod(x.shape[1:]) + + # for 1 increment in x along the first dim, + # (flattened) scale will increment by stride / block_size + if (start * stride) % block_size != 0 or (end * stride) % block_size != 0: + raise ValueError( + f"Invalid start or end for shape={x.shape} and block_size={block_size}. " + f"Make sure start and end align with block boundary. " + f"Received start={start}, end={end}." + ) + + # note that for 4-bit, we store .codes as flattened buffer + # divide by 2 since we store 2x 4-bit in 1x uint8 + codes = x.codes[start * stride // 2 : end * stride // 2] + scale = x.scale[start * stride // block_size : end * stride // block_size] + + # adjust the first dim + shape = (x.shape[0] * codes.numel() // x.codes.numel(),) + x.shape[1:] + + return OptimState4bit(codes, scale, x.qmap.clone(), x.signed, shape) + + +if TORCH_VERSION_AT_LEAST_2_5: + from torch.serialization import add_safe_globals + + add_safe_globals([OptimState4bit]) diff --git a/lib/python3.12/site-packages/torchao/optim/subclass_8bit.py b/lib/python3.12/site-packages/torchao/optim/subclass_8bit.py new file mode 100644 index 0000000000000000000000000000000000000000..66d31432b81d69b062b2792a2d6d031a38a7fe5a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/optim/subclass_8bit.py @@ -0,0 +1,231 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import math + +import torch +from torch import Tensor +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_4, + TORCH_VERSION_AT_LEAST_2_5, + TorchAOBaseTensor, +) + +from .quant_utils import ( + create_dynamic_map, + dequant_with_qmap, + quantize_8bit_with_qmap, + scale_tensor, +) + +aten = torch.ops.aten +c10d_functional = torch.ops.c10d_functional +_c10d_functional = torch.ops._c10d_functional + +QMAP_SIGNED = create_dynamic_map(signed=True) +QMAP_UNSIGNED = create_dynamic_map(signed=False) + + +class OptimState8bit(TorchAOBaseTensor): + tensor_attrs = ["codes", "scale", "qmap"] + + @staticmethod + def __new__(cls, codes: Tensor, scale: Tensor, qmap: Tensor, signed: bool): + return Tensor._make_wrapper_subclass(cls, codes.shape, device=codes.device) + + def __init__(self, codes: Tensor, scale: Tensor, qmap: Tensor, signed: bool): + """Create quantized 8-bit optimizer state as proposed in https://arxiv.org/abs/2110.02861 + + Args + codes: quantized 8-bit data stored as uint8. Has the same shape as the original float tensor. + scale: scale data for block-wise quantization. + qmap: lookup table that maps between quantized value (code) and float value. + signed: whether the tensor is signed or unsigned. + + NOTE: To get block-wise scale, the original float tensor is first reshape to (-1, block_size). + Thus, the last dimension of the original float tensor is not necessarily divisible by block size. + Given `codes` and `scale`, `block_size` is calculated as `codes.numel() // scale.numel()`. + """ + assert codes.dtype is torch.uint8 + assert scale.ndim == 1 + self.codes = codes + self.scale = scale + self.qmap = qmap + self.signed = signed + self.block_size = codes.numel() // scale.numel() + + def __tensor_flatten__(self): + return self.tensor_attrs, [self.signed] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size=None, outer_stride=None + ): + return cls( + *[tensor_data_dict[name] for name in cls.tensor_attrs], *tensor_attributes + ) + + def dequantize(self, output_dtype=None): + float_data = dequant_with_qmap(self.codes, self.qmap, self.scale) + if output_dtype is not None: + float_data = float_data.to(output_dtype) + return float_data + + @classmethod + def zeros(cls, shape, signed: bool = True, block_size: int = 256, device=None): + codes = torch.zeros(shape, dtype=torch.uint8, device=device) + scale = torch.zeros(codes.numel() // block_size, device=device) + qmap = torch.tensor(QMAP_SIGNED if signed else QMAP_UNSIGNED, device=device) + return cls(codes, scale, qmap, signed) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(signed={self.signed}, block_size={self.block_size}, " + f"shape={tuple(self.shape)}, device={self.device}, requires_grad={self.requires_grad})" + ) + + +# in pre-2.4, calling .to(device, dtype) will not dispatch aten._to_copy.default when +# dtype is the same but device is different. thus, we must override .to() method instead. +if not TORCH_VERSION_AT_LEAST_2_4: + + def _to(self, *args, **kwargs): + # ignore other args/kwargs + device = kwargs.pop("device", None) + return OptimState8bit( + self.codes.to(device), + self.scale.to(device), + self.qmap.to(device), + self.signed, + ) + + OptimState8bit.to = _to + del _to # make sure to not re-use + + +@OptimState8bit.implements(aten.copy_.default) +def _(func, types, args, kwargs): + dst = args[0] + src = args[1] + + if isinstance(dst, OptimState8bit) and isinstance(src, OptimState8bit): + assert dst.signed == src.signed and dst.block_size == src.block_size + dst.codes.copy_(src.codes) + dst.scale.copy_(src.scale) + # qmap should be the same, don't need to copy + + elif isinstance(dst, OptimState8bit): + scaled_src, scale = scale_tensor(src, dst.block_size) + codes = quantize_8bit_with_qmap(scaled_src, dst.qmap) + dst.codes.copy_(codes) + dst.scale.copy_(scale) + + else: + dst.copy_(src.dequantize()) + + return dst + + +@OptimState8bit.implements(aten._to_copy.default) +def _(func, types, args, kwargs): + # ignore dtype + device = kwargs.get("device", None) + out = OptimState8bit( + args[0].codes.to(device=device), + args[0].scale.to(device=device), + args[0].qmap.to(device=device), + args[0].signed, + ) + return return_and_correct_aliasing(func, args, kwargs, out) + + +@OptimState8bit.implements(aten.lerp.Scalar) +def _(func, types, args, kwargs): + args = [x.dequantize() if isinstance(x, OptimState8bit) else x for x in args] + return func(*args, **kwargs) + + +# this is needed for DTensor.from_local() +@OptimState8bit.implements(aten.view.default) +def _(func, types, args, kwargs): + x, shape = args + return OptimState8bit(x.codes.view(shape), x.scale, x.qmap, x.signed) + + +@OptimState8bit.implements( + [ + # required by DTensor.full_tensor() + c10d_functional.all_gather_into_tensor.default, + _c10d_functional.all_gather_into_tensor.default, + c10d_functional.wait_tensor.default, + _c10d_functional.wait_tensor.default, + # required by torch.distributed.checkpoint.save + aten.detach.default, + ] +) +def _(func, types, args, kwargs): + x = args[0] + if not isinstance(x, OptimState8bit): + raise ValueError(f"expecting a OptimState8bit but found {type(x)}") + + # assume tensors from all ranks have the same signedness + return OptimState8bit( + func(x.codes, *args[1:], **kwargs), + func(x.scale, *args[1:], **kwargs), + x.qmap.clone(), + x.signed, + ) + + +# required by torch.distributed.checkpoint.save +# note that we don't actually implement pin memory for this tensor subclass +# (pin_memory argument is ignored in aten._to_copy) +@OptimState8bit.implements(aten.is_pinned.default) +def _(func, types, args, kwargs): + return ( + args[0].codes.is_pinned() + and args[0].scale.is_pinned() + and args[0].qmap.is_pinned() + ) + + +# required by torch.distributed.checkpoint.load when world size changes i.e. re-sharding +@OptimState8bit.implements(aten.slice.Tensor) +def _(func, types, args, kwargs): + x, dim, start, end = args[:4] + step = args[4] if len(args) > 4 else 1 + + # input validation + if dim != 0: + raise ValueError("Only support aten.slice along the first dim") + if step != 1: + raise ValueError("Only support aten.slice with step=1") + + block_size = x.block_size + stride = math.prod(x.shape[1:]) + + # for 1 increment in x along the first dim, + # (flattened) scale will increment by stride / block_size + if (start * stride) % block_size != 0 or (end * stride) % block_size != 0: + raise ValueError( + f"Invalid start or end for shape={x.shape} and block_size={block_size}. " + f"Make sure start and end align with block boundary. " + f"Received start={start}, end={end}." + ) + + return OptimState8bit( + x.codes[start:end], + x.scale[start * stride // block_size : end * stride // block_size], + x.qmap.clone(), + x.signed, + ) + + +if TORCH_VERSION_AT_LEAST_2_5: + from torch.serialization import add_safe_globals + + add_safe_globals([OptimState8bit]) diff --git a/lib/python3.12/site-packages/torchao/optim/subclass_fp8.py b/lib/python3.12/site-packages/torchao/optim/subclass_fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..1ae670dd6d23532b0991d81a0fd7565d861d55fd --- /dev/null +++ b/lib/python3.12/site-packages/torchao/optim/subclass_fp8.py @@ -0,0 +1,198 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import math + +import torch +from torch import Tensor +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, TorchAOBaseTensor + +aten = torch.ops.aten +c10d_functional = torch.ops.c10d_functional +_c10d_functional = torch.ops._c10d_functional + +DTYPE = torch.float8_e4m3fn + + +def quantize_fp8(input: Tensor, block_size: int): + shape = input.shape + input = input.view(-1, block_size) + scale = input.abs().amax(-1).clip(1e-12) / torch.finfo(DTYPE).max + input = input / scale.view(-1, 1) + codes = input.to(DTYPE).view(-1) + return codes.view(shape), scale + + +# NOTE: FP8 sign bit is redundant for unsigned optim state. +# we may investigate how to use it to increase range/precision for unsigned optim state. +# https://arxiv.org/abs/2409.12517 uses FP8 E5M2 for 2nd Adam buffer +class OptimStateFp8(TorchAOBaseTensor): + tensor_attrs = ["codes", "scale"] + + @staticmethod + def __new__(cls, codes: Tensor, scale: Tensor): + return Tensor._make_wrapper_subclass(cls, codes.shape, device=codes.device) + + def __init__(self, codes: Tensor, scale: Tensor): + """Create quantized FP8 optimizer state. + + Args + codes: quantized FP8 E4M3FN data. Has the same shape as the original float tensor. + scale: scale data for block-wise quantization. + + NOTE: To get block-wise scale, the original float tensor is first reshape to (-1, block_size). + Thus, the last dimension of the original float tensor is not necessarily divisible by block size. + Given `codes` and `scale`, `block_size` is calculated as `codes.numel() // scale.numel()`. + """ + assert codes.dtype is DTYPE + assert scale.ndim == 1 + self.codes = codes + self.scale = scale + self.block_size = codes.numel() // scale.numel() + + def __tensor_flatten__(self): + return self.tensor_attrs, [] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size=None, outer_stride=None + ): + return cls( + *[tensor_data_dict[name] for name in cls.tensor_attrs], *tensor_attributes + ) + + def dequantize(self, output_dtype=None): + float_data = self.codes.float() + float_data = float_data.view(-1, self.block_size) * self.scale.view(-1, 1) + + if output_dtype is not None: + float_data = float_data.to(output_dtype) + return float_data.view(self.codes.shape) + + @classmethod + def zeros(cls, shape, block_size: int = 256, device=None): + codes = torch.zeros(shape, dtype=DTYPE, device=device) + scale = torch.zeros(codes.numel() // block_size, device=device) + return cls(codes, scale) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(block_size={self.block_size}, " + f"shape={tuple(self.shape)}, device={self.device}, requires_grad={self.requires_grad})" + ) + + +@OptimStateFp8.implements(aten.copy_.default) +def _(func, types, args, kwargs): + dst = args[0] + src = args[1] + + if isinstance(dst, OptimStateFp8) and isinstance(src, OptimStateFp8): + assert dst.block_size == src.block_size + dst.codes.copy_(src.codes) + dst.scale.copy_(src.scale) + + elif isinstance(dst, OptimStateFp8): + codes, scale = quantize_fp8(src, dst.block_size) + dst.codes.copy_(codes) + dst.scale.copy_(scale) + + else: + dst.copy_(src.dequantize()) + + return dst + + +@OptimStateFp8.implements(aten._to_copy.default) +def _(func, types, args, kwargs): + # ignore dtype + device = kwargs.get("device", None) + out = OptimStateFp8( + args[0].codes.to(device=device), + args[0].scale.to(device=device), + ) + return return_and_correct_aliasing(func, args, kwargs, out) + + +@OptimStateFp8.implements(aten.lerp.Scalar) +def _(func, types, args, kwargs): + args = [x.dequantize() if isinstance(x, OptimStateFp8) else x for x in args] + return func(*args, **kwargs) + + +# this is needed for DTensor.from_local() +@OptimStateFp8.implements(aten.view.default) +def _(func, types, args, kwargs): + x, shape = args + return OptimStateFp8(x.codes.view(shape), x.scale) + + +@OptimStateFp8.implements( + [ + # required by DTensor.full_tensor() + c10d_functional.all_gather_into_tensor.default, + _c10d_functional.all_gather_into_tensor.default, + c10d_functional.wait_tensor.default, + _c10d_functional.wait_tensor.default, + # required by torch.distributed.checkpoint.save + aten.detach.default, + ] +) +def _(func, types, args, kwargs): + x = args[0] + if not isinstance(x, OptimStateFp8): + raise ValueError(f"expecting a OptimStateFp8 but found {type(x)}") + + # assume tensors from all ranks have the same signedness + return OptimStateFp8( + func(x.codes, *args[1:], **kwargs), + func(x.scale, *args[1:], **kwargs), + ) + + +# required by torch.distributed.checkpoint.save +# note that we don't actually implement pin memory for this tensor subclass +# (pin_memory argument is ignored in aten._to_copy) +@OptimStateFp8.implements(aten.is_pinned.default) +def _(func, types, args, kwargs): + return args[0].codes.is_pinned() and args[0].scale.is_pinned() + + +# required by torch.distributed.checkpoint.load when world size changes i.e. re-sharding +@OptimStateFp8.implements(aten.slice.Tensor) +def _(func, types, args, kwargs): + x, dim, start, end = args[:4] + step = args[4] if len(args) > 4 else 1 + + # input validation + if dim != 0: + raise ValueError("Only support aten.slice along the first dim") + if step != 1: + raise ValueError("Only support aten.slice with step=1") + + block_size = x.block_size + stride = math.prod(x.shape[1:]) + + # for 1 increment in x along the first dim, + # (flattened) scale will increment by stride / block_size + if (start * stride) % block_size != 0 or (end * stride) % block_size != 0: + raise ValueError( + f"Invalid start or end for shape={x.shape} and block_size={block_size}. " + f"Make sure start and end align with block boundary. " + f"Received start={start}, end={end}." + ) + + return OptimStateFp8( + x.codes[start:end], + x.scale[start * stride // block_size : end * stride // block_size], + ) + + +if TORCH_VERSION_AT_LEAST_2_5: + from torch.serialization import add_safe_globals + + add_safe_globals([OptimStateFp8]) diff --git a/lib/python3.12/site-packages/torchao/prototype/__init__.py b/lib/python3.12/site-packages/torchao/prototype/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/prototype/custom_fp_utils.py b/lib/python3.12/site-packages/torchao/prototype/custom_fp_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3d8de6f0dea7cae06948ba66d9d28ece74495d94 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/custom_fp_utils.py @@ -0,0 +1,235 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# This script was initially developed for sub-byte MX dtypes (FP4 E2M1, FP6 E3M2, and FP6 E2M3). +# It has been refactored to support any sub-byte FP dtypes. However, some behaviors of MX dtypes remain: +# 1. No encodings are reserved for special values (+/-inf, NaN). +# 2. When downcasting from FP32 to Floatx, +# - Rounding mode is round to nearest, ties to even. +# - Values outside the representable range of Floatx after rounding are clamped to the maximum Floatx +# magnitude (sign is preserved). + +import torch +from torch import Tensor + + +def _n_ones(n: int) -> int: + return (1 << n) - 1 + + +EBITS_F32, MBITS_F32 = 8, 23 +F32_EXP_BIAS = _n_ones(EBITS_F32 - 1) + + +def _f32_to_floatx_unpacked(x: Tensor, ebits: int, mbits: int) -> Tensor: + """Convert FP32 numbers to sub-byte floating point numbers with the given + number of exponent and mantissa bits. + + Input: torch.Tensor of dtype torch.float + Output: torch.Tensor of dtype torch.uint8, where the bit encoding is stored + in the least significant bits. e.g. + fp4: bits 0-3 empty and bits 4-7 in fp4_e2m1 encoding + fp6: bits 0-1 empty and bits 2-7 in fp6_e2m3 or fp6_e3m2 encoding + + Note: there are no special values (NaN, inf) support in this code. Values + outside the representable range of Floatx after rounding are clamped to the + maximum Floatx magnitude (sign is preserved). + + Code below is an adaptation of https://fburl.com/code/ciwofcg4 + + Background 1: last answer in https://stackoverflow.com/questions/8981913/how-to-perform-round-to-even-with-floating-point-numbers # noqa: E501 + Background 2: Computer Organization and Design, RISC-V edition, Chapter 3.5 + """ + assert x.dtype == torch.float + assert 1 + ebits + mbits <= 8 + + # calculate constants + exp_bias = _n_ones(ebits - 1) + max_int = _n_ones(ebits + mbits) + sign_mask = 1 << (ebits + mbits) + + # TODO document this better + magic_adder = _n_ones(MBITS_F32 - mbits - 1) + + # all E bits and M bits are 1s + max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2**mbits)) + + # E bits = 1, M bits = 0 + min_normal = 2 ** (1 - exp_bias) + + denorm_exp = ( + # exp bias conversion between formats + (F32_EXP_BIAS - exp_bias) + # mantissa length difference between formats + + (MBITS_F32 - mbits) + # add one to encoded exponent for denormalized numbers + + 1 + ) + denorm_mask_int = denorm_exp << MBITS_F32 + + # reinterpret int32 as float32 + denorm_mask_float = torch.tensor(denorm_mask_int, dtype=torch.int32).view( + torch.float32 + ) + + # save the sign + # Note that we have torch.uint32, but some ops like cpu bit shifts + # do not work on it. So, we stay in int32. + x = x.view(torch.int32) + sign = x & 0x80000000 + + # set everything to positive, will add sign back at the end + x = x ^ sign + + # TODO: can the branch floating point comparisons below be done without + # converting to float? probably but need to verify + x = x.view(torch.float) + + # rewrite saturate/denorm/norm branches without explicit data dependent + # control flow, to be more compiler friendly + saturate_mask = x >= max_normal + denormal_mask = torch.logical_and(torch.logical_not(saturate_mask), x < min_normal) + normal_mask = torch.logical_not(torch.logical_or(saturate_mask, denormal_mask)) + + # + # branch 1: saturate to max val - handled later in the code which combines + # the branches + # + + # + # branch 2: to conversion to denormal as well as rounding up to normal + # + denormal_x = x + denorm_mask_float + denormal_x = denormal_x.view(torch.int32) + denormal_x -= denorm_mask_int + denormal_x = denormal_x.to(torch.uint8) + + # + # branch 3: stay in normal range, adjust the exponent and round + # + normal_x = x.view(torch.int32) + # resulting mantissa is odd + mant_odd = (normal_x >> (MBITS_F32 - mbits)) & 1 + # update exponent, rounding bias part 1 + val_to_add = ((exp_bias - F32_EXP_BIAS) << MBITS_F32) + magic_adder + normal_x += val_to_add + # rounding bias part 2 + normal_x += mant_odd + # take the bits! + normal_x = normal_x >> (MBITS_F32 - mbits) + normal_x = normal_x.to(torch.uint8) + + # + # combine the branches + # + x = torch.full_like(x, max_int, dtype=torch.uint8) + x = torch.where(denormal_mask, denormal_x, x) + x = torch.where(normal_mask, normal_x, x) + + # add sign back + sign_lp = sign >> (MBITS_F32 + EBITS_F32 - mbits - ebits) + sign_lp = sign_lp.to(torch.uint8) + # Right shift of a negative signed integer can fill the least significant + # bits with either 1s or 0s, depending on the implementation. Since PyTorch + # doesn't have an uint32 dtype, we mask out these bits to get just the + # f4 sign bit + sign_lp = sign_lp & sign_mask + x = x | sign_lp + + return x.to(torch.uint8) + + +# TODO(future): check if LUT for everything is faster than bit shifting, +# especially for fp4 (only 2^4=16 unique values). +def _floatx_unpacked_to_f32(x: Tensor, ebits: int, mbits: int) -> Tensor: + """Convert sub-byte floating point numbers with the given number of exponent + and mantissa bits to FP32. + + Input: torch.Tensor of dtype uint8, where the bit encoding is stored + in the least significant bits. e.g. + fp4: bits 0-3 empty and bits 4-7 in fp4_e2m1 encoding + fp6: bits 0-1 empty and bits 2-7 in fp6_e2m3 or fp6_e3m2 encoding + Output: torch.Tensor of dtype fp32 with the dequantized value + """ + assert x.dtype == torch.uint8 + assert 1 + ebits + mbits <= 8 + + sign_mask = 1 << (ebits + mbits) + exp_bias = _n_ones(ebits - 1) + mantissa_mask = _n_ones(mbits) + + # save the sign + sign_lp = x & sign_mask + + # set everything to positive, will add sign back at the end + x_pos = x ^ sign_lp + + # + # 1. Calculate zero mask + # + zero_mask = x_pos == 0 + + # + # 2. Calculate the denormal path mask + # + denormal_mask = torch.logical_and((x_pos > 0), ((x_pos >> mbits) == 0)) + + # + # 3. Calculate the normal path + # + + # calculate the new exponent and shift it to bits 2:9 of the result + exp_biased_lp = x_pos >> mbits + exp_biased_f32 = exp_biased_lp - exp_bias + F32_EXP_BIAS + exp_biased_f32 = exp_biased_f32.to(torch.int32) << MBITS_F32 + + # shift the mantissa to bits 10:32 of the result + mantissa_lp_int32 = (x_pos & mantissa_mask).to(torch.int32) + mantissa_f32 = mantissa_lp_int32 << (MBITS_F32 - mbits) + result = exp_biased_f32 | mantissa_f32 + + # + # 4. Add the zero and denormal casts to the already casted normal path + # + result[zero_mask] = 0 + + denormal_exp_biased = 1 - exp_bias + F32_EXP_BIAS + + # fast path. + # without this, performance for FP4_E2M1 is slower by 2x + if mbits == 1: + result[denormal_mask] = (denormal_exp_biased - mbits) << MBITS_F32 + + else: + # iterate over all possible values of mantissa + # i=0, j=1 + # i=1, j=10,11 + # i=2, j=100,101,110,111 + # and so on + for i in range(mbits): + for mantissa_cmp in range(1 << i, 1 << (i + 1)): + # left shift mantissa until it overflows (create an implicit 1) + # subtract exponent by the same amount + left_shift = mbits - i + mantissa_f32 = (mantissa_cmp - (1 << i)) << ( + left_shift + MBITS_F32 - mbits + ) + exp_biased_f32 = (denormal_exp_biased - left_shift) << MBITS_F32 + + # we can update this in-place since the values won't overlap + # torch.compile() may complain unsupported operand type(s) for |: 'SymInt' and 'int' + # thus we use + instead of | here + mantissa_lp_int32[mantissa_lp_int32 == mantissa_cmp] = ( + exp_biased_f32 + mantissa_f32 + ) + + result = torch.where(denormal_mask, mantissa_lp_int32, result) + + # add sign back + sign_f32 = sign_lp.to(torch.int32) << (MBITS_F32 - mbits + EBITS_F32 - ebits) + result = result | sign_f32 + + return result.view(torch.float) diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__init__.py b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..12b940e9b93f6b7c9e513a4c1024b3e70e765f4d Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe1ae0a18ff173d4058119cf27a1619844b03bd4 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear_test.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5f9481f0906993a73d21687c06b546e22043af8 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e122e70f40ad7a2e273a326fbd10112a9fd32452 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_linear_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_scaling_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_scaling_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c830ab37ab1a6ed78dec72c8e5a64b42fc8a411 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/__pycache__/float8nocompile_scaling_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear.py b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..7e0eb85022c4224fcd45d1a1020a426029b33ecd --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear.py @@ -0,0 +1,318 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +""" +A simple module swap UX for a float8 version of `torch.nn.Linear` which +does not require `torch.compile` to be performant. +""" + +import torch + +from torchao.float8.config import Float8LinearConfig +from torchao.float8.float8_tensor import GemmInputRole, LinearMMConfig, ScaledMMConfig +from torchao.prototype.float8nocompile.float8nocompile_scaling_utils import ( + ToFP8ColumnMajor, + ToFP8ColumnMajorT, + ToFP8RowAndColumnMajor, + ToFP8RowMajor, + ToFP8RowMajorTAndNonT, +) +from torchao.prototype.float8nocompile.kernels.fp8_dynamic_tensorwise import ( + KernelAlgorithm, +) + + +class Float8LinearNoCompile(torch.nn.Linear): + """ + Float8LinearNoCompile is a version of Float8Linear that does not require + the use of torch.compile to be performant. + + Note: this is **prototype** and not suitable for production use. + """ + + def __init__(self, *args, **kwargs): + """ + Additional arguments on top of `torch.nn.Linear`'s arguments: + * `config`: Float8LinearConfig + """ + self.config = kwargs.pop("config") + self.kernel_algo = kwargs.pop("kernel_algo") + self.no_precompute_for_backward = kwargs.pop( + "no_precompute_for_backward", False + ) + super().__init__(*args, **kwargs) + + self.linear_mm_config = LinearMMConfig( + # output + ScaledMMConfig( + self.config.emulate, + self.config.gemm_config_output.use_fast_accum, + False, + self.config.pad_inner_dim, + ), + # grad_input + ScaledMMConfig( + self.config.emulate, + self.config.gemm_config_grad_input.use_fast_accum, + False, + self.config.pad_inner_dim, + ), + # grad_weight + ScaledMMConfig( + self.config.emulate, + self.config.gemm_config_grad_weight.use_fast_accum, + False, + self.config.pad_inner_dim, + ), + ) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + if self.no_precompute_for_backward: + output = matmul_with_args_in_hp_no_precompute_for_backward.apply( + input, + self.weight, + self.config, + self.linear_mm_config, + self.kernel_algo, + ) + else: + output = matmul_with_args_in_hp.apply( + input, + self.weight, + self.config, + self.linear_mm_config, + self.kernel_algo, + ) + return output + + @classmethod + def from_float( + cls, + mod, + config: Float8LinearConfig, # only default config is supported, non-defaults silently ignored + kernel_algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, + no_precompute_for_backward: bool = False, + ): + """ + Create an nn.Linear with fp8 compute from a regular nn.Linear + + Args: + mod (torch.nn.Linear): nn.Linear to convert + config (Optional[Float8LinearConfig]): configuration for conversion to float8 (note: only + default config is supported, non-defaults silently ignored) + """ + with torch.device("meta"): + new_mod = cls( + mod.in_features, + mod.out_features, + bias=False, + config=config, + kernel_algo=kernel_algo, + no_precompute_for_backward=no_precompute_for_backward, + ) + new_mod.weight = mod.weight + new_mod.bias = mod.bias + + # TODO(danielvegamyhre): support for FSDP once dependencies are implemented + return new_mod + + +class matmul_with_args_in_hp(torch.autograd.Function): + """FP8 matmul with args in high precision to be used in a region without AC. + FP8 tensors only needed for backward are computed as part of kernels in the forward pass, + to reduce number of kernel dispatches and increase throughput, at the cost of higher + peak memory usage.""" + + @staticmethod + def forward( + ctx, + input_hp: torch.Tensor, + weight_hp: torch.Tensor, + config: Float8LinearConfig, + linear_mm_config: LinearMMConfig, + kernel_algo: KernelAlgorithm, + ): + # reshape to be 2D for triton kernels + orig_input_shape = input_hp.shape + input_hp = input_hp.reshape(-1, input_hp.shape[-1]) + + # output = input @ weight_t + input_fp8_row_major, input_fp8_col_major = ToFP8RowAndColumnMajor.apply( + input_hp, + config.cast_config_input.target_dtype, + linear_mm_config, + GemmInputRole.INPUT, + kernel_algo, + ) + weight_t_fp8_col_major = ToFP8ColumnMajorT.apply( + weight_hp, + config.cast_config_weight.target_dtype, + linear_mm_config, + GemmInputRole.WEIGHT, + kernel_algo, + ) + output = torch.mm(input_fp8_row_major, weight_t_fp8_col_major) + + # save data for backward before returning + ctx.save_for_backward(input_fp8_col_major, weight_hp) + ctx.config = config + ctx.linear_mm_config = linear_mm_config + ctx.kernel_algo = kernel_algo + ctx.no_precompute_for_backward = False + + # reshape back to expected dims + output = output.reshape(*orig_input_shape[:-1], output.shape[-1]) + return output + + @staticmethod + def backward(ctx, grad_output): + # grad_output may not be contiguous in cases like: + # output.sum().backward() where grad is all 1s, so the (M,N) view of the scalar "1" + # results in a non-contiguous tensor with stride (0,0). + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + input_fp8_col_major, weight_hp = ctx.saved_tensors + + # reshsape to be 2D for triton kernels + orig_grad_output_shape = grad_output.shape + grad_output = grad_output.reshape(-1, grad_output.shape[-1]) + + # cast grad output to float8_e5m2 for backward + grad_output_fp8_row_major, grad_output_t_row_major = ( + ToFP8RowMajorTAndNonT.apply( + grad_output, + ctx.config.cast_config_grad_output.target_dtype, + ctx.linear_mm_config, + GemmInputRole.GRAD_OUTPUT, + ctx.kernel_algo, + ) + ) + + # grad_input = grad_output @ weight + weight_fp8_col_major = ToFP8ColumnMajor.apply( + weight_hp, + ctx.config.cast_config_weight.target_dtype, + ctx.linear_mm_config, + GemmInputRole.WEIGHT, + ctx.kernel_algo, + ) + grad_input = torch.mm(grad_output_fp8_row_major, weight_fp8_col_major) + + # reshape grad input to match original shape + grad_input = grad_input.reshape( + *orig_grad_output_shape[:-1], grad_input.shape[-1] + ) + + # grad_weight = grad_output_t @ input + # apparently this variant is slightly faster than `grad_weight_t = input_t @ grad_output` + # source: https://github.com/pytorch/ao/blob/fe5f11b2c58b452e01ba9ec7359629928b143619/torchao/float8/float8_linear.py#L84-L85 + grad_weight = torch.mm(grad_output_t_row_major, input_fp8_col_major) + + # grad input shape + return grad_input, grad_weight, None, None, None, None + + +class matmul_with_args_in_hp_no_precompute_for_backward(torch.autograd.Function): + """FP8 matmul with args in high precision to be used in a region with AC. + FP8 tensors only needed for backward are only computed in the backward pass + when needed, to reduce peak memory usage.""" + + @staticmethod + def forward( + ctx, + input_hp: torch.Tensor, + weight_hp: torch.Tensor, + config: Float8LinearConfig, + linear_mm_config: LinearMMConfig, + kernel_algo: KernelAlgorithm, + ): + # reshape to be 2D for triton kernels + orig_input_shape = input_hp.shape + input_hp = input_hp.reshape(-1, input_hp.shape[-1]) + + # output = input @ weight_t + input_fp8_row_major = ToFP8RowMajor.apply( + input_hp, + config.cast_config_input.target_dtype, + linear_mm_config, + GemmInputRole.INPUT, + kernel_algo, + ) + weight_t_fp8_col_major = ToFP8ColumnMajorT.apply( + weight_hp, + config.cast_config_weight.target_dtype, + linear_mm_config, + GemmInputRole.WEIGHT, + kernel_algo, + ) + output = torch.mm(input_fp8_row_major, weight_t_fp8_col_major) + + # with AC we only will save the original hp input tensor and weight for backward, + # and do the necessary fp8 conversions during the backward pass. + ctx.save_for_backward(input_hp, weight_hp) + ctx.config = config + ctx.linear_mm_config = linear_mm_config + ctx.kernel_algo = kernel_algo + ctx.no_precompute_for_backward = True + + # reshape back to expected dims + output = output.reshape(*orig_input_shape[:-1], output.shape[-1]) + return output + + @staticmethod + def backward(ctx, grad_output): + # grad_output may not be contiguous in cases like: + # output.sum().backward() where grad is all 1s, so the (M,N) view of the scalar "1" + # results in a non-contiguous tensor with stride (0,0). + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + input_hp, weight_hp = ctx.saved_tensors + + # reshsape to be 2D for triton kernels + orig_grad_output_shape = grad_output.shape + grad_output = grad_output.reshape(-1, grad_output.shape[-1]) + + # cast grad output to float8_e5m2 for backward + grad_output_fp8_row_major, grad_output_t_row_major = ( + ToFP8RowMajorTAndNonT.apply( + grad_output, + ctx.config.cast_config_grad_output.target_dtype, + ctx.linear_mm_config, + GemmInputRole.GRAD_OUTPUT, + ctx.kernel_algo, + ) + ) + + # grad_input = grad_output @ weight + weight_fp8_col_major = ToFP8ColumnMajor.apply( + weight_hp, + ctx.config.cast_config_weight.target_dtype, + ctx.linear_mm_config, + GemmInputRole.WEIGHT, + ctx.kernel_algo, + ) + grad_input = torch.mm(grad_output_fp8_row_major, weight_fp8_col_major) + + # reshape grad input to match original shape + grad_input = grad_input.reshape( + *orig_grad_output_shape[:-1], grad_input.shape[-1] + ) + + # grad_weight = grad_output_t @ input + # apparently this variant is slightly faster than `grad_weight_t = input_t @ grad_output` + # source: https://github.com/pytorch/ao/blob/fe5f11b2c58b452e01ba9ec7359629928b143619/torchao/float8/float8_linear.py#L84-L85 + input_fp8_col_major = ToFP8ColumnMajor.apply( + input_hp, + ctx.config.cast_config_input.target_dtype, + ctx.linear_mm_config, + GemmInputRole.INPUT, + ctx.kernel_algo, + ) + grad_weight = torch.mm(grad_output_t_row_major, input_fp8_col_major) + + # grad input shape + return grad_input, grad_weight, None, None, None, None diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear_test.py b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7df5ce768cf13bb4c9f08702108440faefa14b3b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear_test.py @@ -0,0 +1,102 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import pytest +import torch + +from torchao.float8.config import Float8LinearConfig +from torchao.float8.float8_linear import matmul_with_hp_or_float8_args +from torchao.float8.float8_tensor import LinearMMConfig, ScaledMMConfig +from torchao.prototype.float8nocompile.float8nocompile_linear import ( + matmul_with_args_in_hp, +) +from torchao.prototype.float8nocompile.float8nocompile_scaling_utils import ( + KernelAlgorithm, +) + + +# unit test comparing the two implementations +@pytest.mark.parametrize( + "input_shape", + [(32, 16), (1, 32, 16), (2, 32, 16)], +) +def test_matmul_with_args_in_hp(input_shape: tuple[int, int]): + assert torch.cuda.is_available() + device = "cuda" + + # high precision inputs + input_bf16 = torch.randn( + input_shape, dtype=torch.bfloat16, device=device, requires_grad=True + ) + prod_input_bf16 = input_bf16.clone().detach().to(device).requires_grad_(True) + prototype_input_bf16 = input_bf16.clone().detach().to(device).requires_grad_(True) + + # high precision weights + # nn.Linear stores weights in transposed form + weight_bf16 = torch.randn( + (32, input_bf16.shape[-1]), + dtype=torch.bfloat16, + device=device, + requires_grad=True, + ) + prod_weight_bf16 = weight_bf16.clone().detach().to(device).requires_grad_(True) + prototype_weight_bf16 = weight_bf16.clone().detach().to(device).requires_grad_(True) + + # default configs + config = Float8LinearConfig() + emulate = False + linear_mm_config = linear_mm_config = LinearMMConfig( + # output + ScaledMMConfig( + emulate, + config.gemm_config_output.use_fast_accum, + False, + config.pad_inner_dim, + ), + # grad_input + ScaledMMConfig( + emulate, + config.gemm_config_grad_input.use_fast_accum, + False, + config.pad_inner_dim, + ), + # grad_weight + ScaledMMConfig( + emulate, + config.gemm_config_grad_weight.use_fast_accum, + False, + config.pad_inner_dim, + ), + ) + + # prod forward. expects transposed weight. + out_prod = matmul_with_hp_or_float8_args.apply( + prod_input_bf16, prod_weight_bf16.t(), linear_mm_config, config + ) + + # prototype forward. expects non-transposed weight + out_prototype = matmul_with_args_in_hp.apply( + prototype_input_bf16, + prototype_weight_bf16, + config, + linear_mm_config, + KernelAlgorithm.ATOMIC_MAX, + ) + + # compare model outputs + assert torch.allclose(out_prod, out_prototype, atol=0, rtol=0) + + out_prod.sum().backward() + out_prototype.sum().backward() + + # compare input gradients + assert torch.allclose( + prod_input_bf16.grad, prototype_input_bf16.grad, atol=0, rtol=0 + ) + + # compare weight gradients + assert torch.allclose( + prod_weight_bf16.grad, prototype_weight_bf16.grad, atol=0, rtol=0 + ) diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear_utils.py b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7e121c559e0be0dd81b5e5cc48c3f189c55bb989 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_linear_utils.py @@ -0,0 +1,58 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import logging +from typing import Callable, Optional + +import torch.nn as nn + +from torchao.float8.config import Float8LinearConfig +from torchao.float8.float8_linear_utils import swap_linear_layers +from torchao.prototype.float8nocompile.float8nocompile_linear import ( + Float8LinearNoCompile, +) +from torchao.prototype.float8nocompile.kernels.fp8_dynamic_tensorwise import ( + KernelAlgorithm, +) + +log = logging.getLogger(__name__) +log.addHandler(logging.NullHandler()) + + +def convert_to_float8_nocompile_training( + module: nn.Module, + *, + config: Float8LinearConfig = None, + module_filter_fn: Optional[Callable[[nn.Module, str], bool]] = None, + kernel_algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, + no_precompute_for_backward: bool = False, +) -> nn.Module: + """ + Swaps `torch.nn.Linear` in `module` with `Float8LinearNoCompile`. + + Args: + module: Module to modify. + module_filter_fn: If specified, only the `torch.nn.Linear` subclasses that + that pass the filter function will be swapped. The inputs to the + filter function are the module instance and the FQN. + config (Float8LinearConfig): configuration for conversion to float8 + + Returns: + nn.Module: The modified module with swapped linear layers. + """ + if config is None: + config = Float8LinearConfig() + + from_float = lambda m: Float8LinearNoCompile.from_float( + m, + config=config, + kernel_algo=kernel_algo, + no_precompute_for_backward=no_precompute_for_backward, + ) + return swap_linear_layers( + module, + from_float, + module_filter_fn=module_filter_fn, + ) diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_scaling_utils.py b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_scaling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7b6a25e3f9e9a1cc5b0d7be3da72434aeba5746b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/float8nocompile_scaling_utils.py @@ -0,0 +1,202 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +""" +Utilities for scaling high precision tensors to float8. +""" + +import torch + +from torchao.float8.float8_tensor import GemmInputRole, LinearMMConfig +from torchao.prototype.float8nocompile.kernels.fp8_dynamic_tensorwise import ( + KernelAlgorithm, + hp_to_fp8_col_major, + hp_to_fp8_col_major_t, + hp_to_fp8_row_and_col_major, + hp_to_fp8_row_major, + hp_to_fp8_row_major_t, + hp_to_fp8_row_major_t_and_non_t, +) + + +class ToFP8RowAndColumnMajor(torch.autograd.Function): + """ + A differentiable conversion to fp8. + * forward: convert from high precision to float8 and produces both row-major and column-major outputs + * backward: pass the gradient without changes + """ + + @staticmethod + def forward( + ctx, + tensor: torch.Tensor, + float8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole, + kernel_algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, + ): + fp8_row_major, fp8_col_major = hp_to_fp8_row_and_col_major( + tensor, + float8_dtype, + linear_mm_config, + gemm_input_role, + algo=kernel_algo, + ) + return fp8_row_major, fp8_col_major + + @staticmethod + def backward(ctx, g): + return g, None, None, None, None + + +class ToFP8RowMajor(torch.autograd.Function): + """ + A differentiable conversion to fp8 in row-major layout. + * forward: convert from high precision to float8 with row-major memory layout + * backward: pass the gradient without changes + """ + + @staticmethod + def forward( + ctx, + tensor: torch.Tensor, + float8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole, + kernel_algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, + ): + fp8_row_major = hp_to_fp8_row_major( + tensor, + float8_dtype, + linear_mm_config, + gemm_input_role, + algo=kernel_algo, + ) + return fp8_row_major + + @staticmethod + def backward(ctx, g): + return g, None, None, None, None + + +class ToFP8RowMajorT(torch.autograd.Function): + """ + A differentiable conversion to fp8 with transposed dimensions in row-major layout. + * forward: convert from high precision to float8 with transposed dimensions with row-major memory layout + * backward: pass the gradient without changes + """ + + @staticmethod + def forward( + ctx, + tensor: torch.Tensor, + float8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole, + kernel_algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, + ): + fp8_row_major_t = hp_to_fp8_row_major_t( + tensor, + float8_dtype, + linear_mm_config, + gemm_input_role, + algo=kernel_algo, + ) + return fp8_row_major_t + + @staticmethod + def backward(ctx, g): + return g, None, None, None, None + + +class ToFP8ColumnMajor(torch.autograd.Function): + """ + A differentiable conversion to fp8 in column-major layout. + * forward: convert from high precision to float8 with column-major memory layout + * backward: pass the gradient without changes + """ + + @staticmethod + def forward( + ctx, + tensor: torch.Tensor, + float8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole, + kernel_algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, + ): + fp8_col_major = hp_to_fp8_col_major( + tensor, + float8_dtype, + linear_mm_config, + gemm_input_role, + algo=kernel_algo, + ) + return fp8_col_major + + @staticmethod + def backward(ctx, g): + return g, None, None, None, None + + +class ToFP8ColumnMajorT(torch.autograd.Function): + """ + A differentiable conversion to fp8 with transposed dimensions in column-major layout. + * forward: convert from high precision to float8 with transposed dimensions in column-major memory layout. + * backward: pass the gradient without changes + """ + + @staticmethod + def forward( + ctx, + tensor: torch.Tensor, + float8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole, + kernel_algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, + ): + fp8_col_major_t = hp_to_fp8_col_major_t( + tensor, + float8_dtype, + linear_mm_config, + gemm_input_role, + algo=kernel_algo, + ) + return fp8_col_major_t + + @staticmethod + def backward(ctx, g): + return g, None, None, None, None + + +class ToFP8RowMajorTAndNonT(torch.autograd.Function): + """ + A differentiable conversion to fp8. + * forward: convert from high precision to float8 and produces both row-major (transposed) and row-major (non-transposed) outputs + * backward: pass the gradient without changes + """ + + @staticmethod + def forward( + ctx, + tensor: torch.Tensor, + float8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole, + kernel_algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, + ): + fp8_row_major, fp8_row_major_t = hp_to_fp8_row_major_t_and_non_t( + tensor, + float8_dtype, + linear_mm_config, + gemm_input_role, + algo=kernel_algo, + ) + return fp8_row_major, fp8_row_major_t + + @staticmethod + def backward(ctx, g): + return g, None, None, None, None diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__init__.py b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b9839c9d13e620826c19b7782c9df5edca0f0c39 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/fp8_dynamic_tensorwise.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/fp8_dynamic_tensorwise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b5351ed5b31877eddcb1ab03a552e8a41cdabe1 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/fp8_dynamic_tensorwise.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/fp8_dynamic_tensorwise_test.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/fp8_dynamic_tensorwise_test.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d9708a7f4630f025e328cc31354c2bac7662b81 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/__pycache__/fp8_dynamic_tensorwise_test.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise.py b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise.py new file mode 100644 index 0000000000000000000000000000000000000000..3786b52eb53db5ade1f51da5f42a56e9506b6ff7 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise.py @@ -0,0 +1,1082 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +""" +Triton kernels for scaling high precision tensors to float8. +""" + +from enum import Enum + +import torch +import triton +import triton.language as tl + +from torchao.float8.float8_tensor import Float8Tensor, GemmInputRole, LinearMMConfig + +EPS = 1e-12 + +FP8_DTYPE_MAP = { + torch.int8: tl.int8, + torch.int16: tl.int16, + torch.int32: tl.int32, + torch.int64: tl.int64, + torch.float8_e4m3fn: tl.float8e4nv, + torch.float8_e5m2: tl.float8e5, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, + torch.float32: tl.float32, + torch.float64: tl.float64, +} + + +kernel_configs_1D = [ + triton.Config({"BLOCK_SIZE": 128}, num_warps=1), + triton.Config({"BLOCK_SIZE": 256}, num_warps=2), + triton.Config({"BLOCK_SIZE": 512}, num_warps=4), + triton.Config({"BLOCK_SIZE": 1024}, num_warps=4), + triton.Config({"BLOCK_SIZE": 2048}, num_warps=8), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=8), +] + +kernel_configs_2D = [ + triton.Config({"BLOCK_SIZE_ROWS": 32, "BLOCK_SIZE_COLS": 32}, num_warps=1), + triton.Config({"BLOCK_SIZE_ROWS": 64, "BLOCK_SIZE_COLS": 64}, num_warps=8), + triton.Config({"BLOCK_SIZE_ROWS": 128, "BLOCK_SIZE_COLS": 128}, num_warps=4), +] + + +class KernelAlgorithm(Enum): + """Enum for FP8 conversion strategy.""" + + # use atomic max to compute global amax between blocks + ATOMIC_MAX = "atomic_max" + + # reduce shared buffer containing local block amaxes to find global amax + REDUCTION = "reduction" + + +@triton.autotune(configs=kernel_configs_1D, key=["num_elements"]) +@triton.jit +def _to_fp8_row_major( + input_ptr, + out_ptr, + scale_ptr, + num_elements: int, + fp8_dtype_min: float, + fp8_dtype_max: float, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + EPS: tl.constexpr, +): + block_id = tl.program_id(axis=0) + + # load scaling factor + scale = tl.load(scale_ptr).to(tl.float32) + + # load block of input tensor + block_start = block_id * BLOCK_SIZE + block_offs = block_start + tl.arange(0, BLOCK_SIZE) + mask = block_offs < num_elements + vals = tl.load(input_ptr + block_offs, mask=mask).to(input_dtype) + + # perform conversion + vals = vals * scale + fp8_vals = tl.clamp(vals, min=fp8_dtype_min, max=fp8_dtype_max).to(output_dtype) + tl.store(out_ptr + block_offs, fp8_vals, mask=mask) + + +@triton.autotune( + configs=kernel_configs_2D, + key=["num_elements"], +) +@triton.jit +def _to_fp8_row_major_t( + input_ptr, + out_ptr, + scale_ptr, + num_elements: int, + fp8_dtype_min: float, + fp8_dtype_max: float, + input_num_rows: int, + input_num_cols: int, + output_num_rows: int, + output_num_cols: int, + input_stride_row: int, + input_stride_col: int, + output_stride_row: int, + output_stride_col: int, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + BLOCK_SIZE_ROWS: tl.constexpr, + BLOCK_SIZE_COLS: tl.constexpr, + EPS: tl.constexpr, +): + block_row_id = tl.program_id(axis=0) + block_col_id = tl.program_id(axis=1) + + # load scaling factor + scale = tl.load(scale_ptr).to(tl.float32) + + # load block of input tensor + block_row_start = block_row_id * BLOCK_SIZE_ROWS + block_col_start = block_col_id * BLOCK_SIZE_COLS + block_row_offs = block_row_start + tl.arange(0, BLOCK_SIZE_ROWS) + block_col_offs = block_col_start + tl.arange(0, BLOCK_SIZE_COLS) + input_offs = ( + block_row_offs[:, None] * input_stride_row + + block_col_offs[None, :] * input_stride_col + ) + input_mask = (block_row_offs[:, None] < input_num_rows) & ( + block_col_offs[None, :] < input_num_cols + ) + vals = tl.load(input_ptr + input_offs, mask=input_mask).to(input_dtype) + + # perform conversion + vals = vals * scale + fp8_vals = tl.clamp(vals, min=fp8_dtype_min, max=fp8_dtype_max).to(output_dtype) + + # write back in tranposed output tensor + out_offs = ( + block_col_offs[:, None] * output_stride_row + + block_row_offs[None, :] * output_stride_col + ) + out_mask = (block_row_offs[:, None] < output_num_rows) & ( + block_col_offs[None, :] < output_num_cols + ) + tl.store(out_ptr + out_offs, fp8_vals.trans(1, 0), mask=out_mask) + + +@triton.autotune( + configs=kernel_configs_2D, + key=["num_elements"], +) +@triton.jit +def _to_fp8_col_major( + input_ptr, + out_ptr, + scale_ptr, + num_elements: int, + fp8_dtype_min: float, + fp8_dtype_max: float, + num_rows: int, + num_cols: int, + out_stride_row: int, + out_stride_col: int, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + BLOCK_SIZE_ROWS: tl.constexpr, + BLOCK_SIZE_COLS: tl.constexpr, + EPS: tl.constexpr, +): + block_row_id = tl.program_id(axis=0) + block_col_id = tl.program_id(axis=1) + + # load scaling factor + scale = tl.load(scale_ptr).to(tl.float32) + + # load block of input tensor + block_row_start = block_row_id * BLOCK_SIZE_ROWS + block_col_start = block_col_id * BLOCK_SIZE_COLS + block_row_offs = block_row_start + tl.arange(0, BLOCK_SIZE_ROWS) + block_col_offs = block_col_start + tl.arange(0, BLOCK_SIZE_COLS) + block_offs = block_row_offs[:, None] * num_cols + block_col_offs[None, :] + mask = (block_row_offs[:, None] < num_rows) & (block_col_offs[None, :] < num_cols) + vals = tl.load(input_ptr + block_offs, mask=mask).to(input_dtype) + + # perform conversion + vals = vals * scale + fp8_vals = tl.clamp(vals, min=fp8_dtype_min, max=fp8_dtype_max).to(output_dtype) + out_offs = block_col_offs[None, :] * num_rows + block_row_offs[:, None] + tl.store(out_ptr + out_offs, fp8_vals, mask=mask) + + +@triton.autotune( + configs=kernel_configs_2D, + key=["num_elements"], +) +@triton.jit +def to_fp8_col_major_t( + input_ptr, + out_ptr, + scale_ptr, + num_elements: int, + fp8_dtype_min: float, + fp8_dtype_max: float, + input_num_rows: int, + input_num_cols: int, + output_num_rows: int, + output_num_cols: int, + input_stride_row: int, + input_stride_col: int, + output_stride_row: int, + output_stride_col: int, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + BLOCK_SIZE_ROWS: tl.constexpr, + BLOCK_SIZE_COLS: tl.constexpr, + EPS: tl.constexpr, +): + block_row_id = tl.program_id(axis=0) + block_col_id = tl.program_id(axis=1) + + # load scaling factor + scale = tl.load(scale_ptr).to(tl.float32) + + # load block of input tensor + block_row_start = block_row_id * BLOCK_SIZE_ROWS + block_col_start = block_col_id * BLOCK_SIZE_COLS + block_row_offs = block_row_start + tl.arange(0, BLOCK_SIZE_ROWS) + block_col_offs = block_col_start + tl.arange(0, BLOCK_SIZE_COLS) + input_offs = ( + block_row_offs[:, None] * input_stride_row + + block_col_offs[None, :] * input_stride_col + ) + input_mask = (block_row_offs[:, None] < input_num_rows) & ( + block_col_offs[None, :] < input_num_cols + ) + vals = tl.load(input_ptr + input_offs, mask=input_mask).to(input_dtype) + + # perform conversion + vals = vals * scale + fp8_vals = tl.clamp(vals, min=fp8_dtype_min, max=fp8_dtype_max).to(output_dtype) + fp8_vals = fp8_vals.trans(1, 0) + + # write back in tranposed output tensor + out_offs = ( + block_col_offs[:, None] * output_stride_row + + block_row_offs[None, :] * output_stride_col + ) + out_mask = (block_col_offs[:, None] < output_num_rows) & ( + block_row_offs[None, :] < output_num_cols + ) + tl.store(out_ptr + out_offs, fp8_vals, mask=out_mask) + + +@triton.autotune( + configs=kernel_configs_2D, + key=["num_elements"], +) +@triton.jit +def _to_fp8_row_and_col_major( + input_ptr, + row_major_out_ptr, + col_major_out_ptr, + scale_ptr, + num_elements: int, + fp8_dtype_min: float, + fp8_dtype_max: float, + num_rows: int, + num_cols: int, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + BLOCK_SIZE_ROWS: tl.constexpr, + BLOCK_SIZE_COLS: tl.constexpr, + EPS: tl.constexpr, +): + block_row_id = tl.program_id(axis=0) + block_col_id = tl.program_id(axis=1) + + # load scaling factor + scale = tl.load(scale_ptr).to(tl.float32) + + # load block of input tensor + block_row_start = block_row_id * BLOCK_SIZE_ROWS + block_col_start = block_col_id * BLOCK_SIZE_COLS + block_row_offs = block_row_start + tl.arange(0, BLOCK_SIZE_ROWS) + block_col_offs = block_col_start + tl.arange(0, BLOCK_SIZE_COLS) + block_offs = block_row_offs[:, None] * num_cols + block_col_offs[None, :] + mask = (block_row_offs[:, None] < num_rows) & (block_col_offs[None, :] < num_cols) + vals = tl.load(input_ptr + block_offs, mask=mask).to(input_dtype) + + # perform conversion + vals = vals * scale + fp8_vals = tl.clamp(vals, min=fp8_dtype_min, max=fp8_dtype_max).to(output_dtype) + + # write row major output + row_major_offs = block_row_offs[:, None] * num_cols + block_col_offs[None, :] + tl.store(row_major_out_ptr + row_major_offs, fp8_vals, mask=mask) + + # write column major output + col_major_offs = block_col_offs[None, :] * num_rows + block_row_offs[:, None] + tl.store(col_major_out_ptr + col_major_offs, fp8_vals, mask=mask) + + +@triton.autotune( + configs=kernel_configs_2D, + key=["num_elements"], +) +@triton.jit +def _to_fp8_row_major_t_and_non_t( + input_ptr, + row_major_out_ptr, + row_major_t_out_ptr, + scale_ptr, + num_elements: int, + fp8_dtype_min: float, + fp8_dtype_max: float, + input_num_rows: int, + input_num_cols: int, + input_stride_row: int, + input_stride_col: int, + row_major_out_stride_row: int, + row_major_out_stride_col: int, + row_major_t_out_stride_row: int, + row_major_t_out_stride_col: int, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + BLOCK_SIZE_ROWS: tl.constexpr, + BLOCK_SIZE_COLS: tl.constexpr, + EPS: tl.constexpr, +): + """ + Reads a row-major, high precision input tensor and writes 2 output tensors: + 1) fp8 row major tensor (transposed) + 2) fp8 row major tensor + """ + block_row_id = tl.program_id(axis=0) + block_col_id = tl.program_id(axis=1) + + # load scaling factor + scale = tl.load(scale_ptr).to(tl.float32) + + # load block of input tensor + block_row_start = block_row_id * BLOCK_SIZE_ROWS + block_col_start = block_col_id * BLOCK_SIZE_COLS + block_row_offs = block_row_start + tl.arange(0, BLOCK_SIZE_ROWS) + block_col_offs = block_col_start + tl.arange(0, BLOCK_SIZE_COLS) + input_offs = ( + block_row_offs[:, None] * input_stride_row + + block_col_offs[None, :] * input_stride_col + ) + mask = (block_row_offs[:, None] < input_num_rows) & ( + block_col_offs[None, :] < input_num_cols + ) + vals = tl.load(input_ptr + input_offs, mask=mask).to(input_dtype) + + # perform conversion + vals = vals * scale + fp8_vals = tl.clamp(vals, min=fp8_dtype_min, max=fp8_dtype_max).to(output_dtype) + + # write row-major output + row_major_offs = ( + block_row_offs[:, None] * row_major_out_stride_row + + block_col_offs[None, :] * row_major_out_stride_col + ) + tl.store(row_major_out_ptr + row_major_offs, fp8_vals, mask=mask) + + # write tranposed row-major output + row_major_t_num_rows = input_num_cols + row_major_t_num_cols = input_num_rows + row_major_t_offs = ( + block_col_offs[:, None] * row_major_t_out_stride_row + + block_row_offs[None, :] * row_major_t_out_stride_col + ) + mask = (block_col_offs[:, None] < row_major_t_num_rows) & ( + block_row_offs[None, :] < row_major_t_num_cols + ) + tl.store(row_major_t_out_ptr + row_major_t_offs, fp8_vals.trans(1, 0), mask=mask) + + +@triton.autotune(configs=kernel_configs_2D, key=["num_elements"]) +@triton.jit +def _to_fp8_col_major_t_and_non_t( + input_ptr, + col_major_out_ptr, + col_major_t_out_ptr, + scale_ptr, + num_elements: int, + fp8_dtype_min: float, + fp8_dtype_max: float, + input_num_rows: int, + input_num_cols: int, + input_stride_row: int, + input_stride_col: int, + col_major_out_stride_row: int, + col_major_out_stride_col: int, + col_major_t_out_stride_row: int, + col_major_t_out_stride_col: int, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + BLOCK_SIZE_ROWS: tl.constexpr, + BLOCK_SIZE_COLS: tl.constexpr, + EPS: tl.constexpr, +): + """ + Reads a row-major, high precision input tensor and writes 2 output tensors: + 1) fp8 col major tensor (transposed) + 2) fp8 col major tensor + """ + # col major tranposed + block_row_id = tl.program_id(axis=0) + block_col_id = tl.program_id(axis=1) + + # load scaling factor + scale = tl.load(scale_ptr).to(tl.float32) + + # load block of input tensor + block_row_start = block_row_id * BLOCK_SIZE_ROWS + block_col_start = block_col_id * BLOCK_SIZE_COLS + block_row_offs = block_row_start + tl.arange(0, BLOCK_SIZE_ROWS) + block_col_offs = block_col_start + tl.arange(0, BLOCK_SIZE_COLS) + input_offs = ( + block_row_offs[:, None] * input_stride_row + + block_col_offs[None, :] * input_stride_col + ) + mask = (block_row_offs[:, None] < input_num_rows) & ( + block_col_offs[None, :] < input_num_cols + ) + vals = tl.load(input_ptr + input_offs, mask=mask).to(input_dtype) + + # perform conversion + vals = vals * scale + fp8_vals = tl.clamp(vals, min=fp8_dtype_min, max=fp8_dtype_max).to(output_dtype) + + # 1. write col-major output + out_offs = block_row_offs[:, None] + block_col_offs[None, :] * input_num_rows + tl.store(col_major_out_ptr + out_offs, fp8_vals, mask=mask) + + # 2. write tranposed col-major output + col_major_t_num_rows = input_num_cols + col_major_t_num_cols = input_num_rows + out_offs = ( + block_col_offs[:, None] * col_major_t_out_stride_row + + block_row_offs[None, :] * col_major_t_out_stride_col + ) + out_mask = (block_col_offs[:, None] < col_major_t_num_rows) & ( + block_row_offs[None, :] < col_major_t_num_cols + ) + tl.store(col_major_t_out_ptr + out_offs, fp8_vals.trans(1, 0), mask=out_mask) + + +@triton.autotune(configs=kernel_configs_1D, key=["num_elements"]) +@triton.jit +def _amax_atomic( + input_ptr, + amax_ptr, + num_elements, + input_dtype: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + EPS: tl.constexpr, +): + # compute local amax for each block + block_id = tl.program_id(axis=0) + block_start = block_id * BLOCK_SIZE + block_offs = block_start + tl.arange(0, BLOCK_SIZE) + block_mask = block_offs < num_elements + vals = tl.load(input_ptr + block_offs, mask=block_mask).to(input_dtype) + block_amax = tl.max(tl.abs(vals)) + tl.atomic_max(amax_ptr, block_amax) + + +@triton.jit +def _scale_atomic( + amax_ptr, + scale_out_ptr, + fp8_dtype_max, + EPS: tl.constexpr, +): + # load previously computed global amax + global_amax = tl.load(amax_ptr).to(tl.float64) + + # compute scale, must be fp32 + scale = (fp8_dtype_max / tl.clamp(global_amax, min=EPS, max=float("inf"))).to( + tl.float32 + ) + + # store scale for use in Float8Tensor constructor + scale_off = tl.arange(0, 1) + tl.store(scale_out_ptr + scale_off, scale) + + +@triton.jit +def _amax_reduction( + input_ptr, + block_amaxes_ptr, + num_elements, + input_dtype: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + EPS: tl.constexpr, +): + # compute local amax for each block + block_id = tl.program_id(axis=0) + block_start = block_id * BLOCK_SIZE + block_offs = block_start + tl.arange(0, BLOCK_SIZE) + block_mask = block_offs < num_elements + vals = tl.load(input_ptr + block_offs, mask=block_mask).to(input_dtype) + block_amax = tl.max(tl.abs(vals)) + tl.store(block_amaxes_ptr + block_id, block_amax) + + +@triton.jit +def _scale_reduction( + block_amaxes_ptr, + scale_out_ptr, + num_elements, + fp8_dtype_max, + BLOCK_SIZE: tl.constexpr, + EPS: tl.constexpr, +): + # calculate global amax across all blocks + global_amax = tl.zeros([1], dtype=tl.float64) + num_blocks = tl.cdiv(num_elements, BLOCK_SIZE) + for i in range(num_blocks): + block_max = tl.load(block_amaxes_ptr + i) + global_amax = tl.maximum(global_amax, block_max) + + # compute scale, must be fp32 + scale = (fp8_dtype_max / tl.clamp(global_amax, min=EPS, max=float("inf"))).to( + tl.float32 + ) + scale_off = tl.arange(0, 1) + tl.store(scale_out_ptr + scale_off, scale) + + +def hp_to_fp8_row_major( + hp_tensor: torch.Tensor, + fp8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole = GemmInputRole.INPUT, + algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, +) -> Float8Tensor: + assert hp_tensor.is_contiguous(), "input tensor must be contiguous" + + num_elements = hp_tensor.numel() + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[fp8_dtype] + + fp8_dtype_min = torch.finfo(fp8_dtype).min + fp8_dtype_max = torch.finfo(fp8_dtype).max + + # compute scaling factor for tensor + scale = _hp_tensor_to_scale( + hp_tensor, + tl_input_dtype, + fp8_dtype_max, + algo, + ) + + # perform fp8 conversion + output_buffer = torch.empty_like( + hp_tensor, dtype=fp8_dtype, device=hp_tensor.device + ) + grid = lambda meta: (triton.cdiv(num_elements, meta["BLOCK_SIZE"]),) + _to_fp8_row_major[grid]( + hp_tensor, + output_buffer, + scale, + num_elements, + fp8_dtype_min, + fp8_dtype_max, + tl_input_dtype, + tl_output_dtype, + EPS=EPS, + ) + + # wrap output tensor in Float8Tensor + fp8_tensor_row_major = Float8Tensor( + output_buffer, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + return fp8_tensor_row_major + + +def hp_to_fp8_row_major_t( + hp_tensor: torch.Tensor, + fp8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole = GemmInputRole.INPUT, + algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, +) -> Float8Tensor: + assert hp_tensor.is_contiguous(), "input tensor must be contiguous" + + num_elements = hp_tensor.numel() + input_num_rows, input_num_cols = hp_tensor.shape + output_num_rows, output_num_cols = input_num_cols, input_num_rows + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[fp8_dtype] + + fp8_dtype_min = torch.finfo(fp8_dtype).min + fp8_dtype_max = torch.finfo(fp8_dtype).max + + # compute scaling factor for tensor + scale = _hp_tensor_to_scale( + hp_tensor, + tl_input_dtype, + fp8_dtype_max, + algo, + ) + + # perform conversion + output_buffer = torch.empty( + (output_num_rows, output_num_cols), dtype=fp8_dtype, device=hp_tensor.device + ) + grid = lambda meta: ( + triton.cdiv(input_num_rows, meta["BLOCK_SIZE_ROWS"]), + triton.cdiv(input_num_cols, meta["BLOCK_SIZE_COLS"]), + ) + _to_fp8_row_major_t[grid]( + hp_tensor, + output_buffer, + scale, + num_elements, + fp8_dtype_min, + fp8_dtype_max, + input_num_rows, + input_num_cols, + output_num_rows, + output_num_cols, + hp_tensor.stride(0), + hp_tensor.stride(1), + output_buffer.stride(0), + output_buffer.stride(1), + input_dtype=tl_input_dtype, + output_dtype=tl_output_dtype, + EPS=EPS, + ) + + # wrap output tensor in Float8Tensor + fp8_tensor_row_major_t = Float8Tensor( + output_buffer, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + return fp8_tensor_row_major_t + + +def hp_to_fp8_col_major( + hp_tensor: torch.Tensor, + fp8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole = GemmInputRole.INPUT, + algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, +) -> Float8Tensor: + assert hp_tensor.is_contiguous(), "input tensor must be contiguous" + + num_elements = hp_tensor.numel() + num_rows, num_cols = hp_tensor.shape + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[fp8_dtype] + + fp8_dtype_min = torch.finfo(fp8_dtype).min + fp8_dtype_max = torch.finfo(fp8_dtype).max + + # compute scaling factor for tensor + scale = _hp_tensor_to_scale( + hp_tensor, + tl_input_dtype, + fp8_dtype_max, + algo, + ) + + # perform fp8 conversion + output_buffer = torch.empty_like( + hp_tensor, dtype=fp8_dtype, device=hp_tensor.device + ) + grid = lambda meta: ( + triton.cdiv(num_rows, meta["BLOCK_SIZE_ROWS"]), + triton.cdiv(num_cols, meta["BLOCK_SIZE_COLS"]), + ) + _to_fp8_col_major[grid]( + hp_tensor, + output_buffer, + scale, + num_elements, + fp8_dtype_min, + fp8_dtype_max, + num_rows, + num_cols, + output_buffer.stride(0), + output_buffer.stride(1), + input_dtype=tl_input_dtype, + output_dtype=tl_output_dtype, + EPS=EPS, + ) + + # for col major we need to update the strides to reflect the new memory layout + col_major_strides = (1, num_rows) + output_buffer = output_buffer.as_strided(output_buffer.size(), col_major_strides) + + # wrap output tensor in Float8Tensor + fp8_tensor_col_major = Float8Tensor( + output_buffer, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + return fp8_tensor_col_major + + +def hp_to_fp8_col_major_t( + hp_tensor: torch.Tensor, + fp8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole = GemmInputRole.INPUT, + algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, +) -> Float8Tensor: + assert hp_tensor.is_contiguous(), "input tensor must be contiguous" + + num_elements = hp_tensor.numel() + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[fp8_dtype] + + fp8_dtype_min = torch.finfo(fp8_dtype).min + fp8_dtype_max = torch.finfo(fp8_dtype).max + + # compute scaling factor for tensor + scale = _hp_tensor_to_scale( + hp_tensor, + tl_input_dtype, + fp8_dtype_max, + algo, + ) + + # perform conversion + output_buffer = torch.empty_like( + hp_tensor.t(), dtype=fp8_dtype, device=hp_tensor.device + ) + grid = lambda meta: (triton.cdiv(num_elements, meta["BLOCK_SIZE"]),) + _to_fp8_row_major[grid]( + hp_tensor, + output_buffer, + scale, + num_elements, + fp8_dtype_min, + fp8_dtype_max, + tl_input_dtype, + tl_output_dtype, + EPS=EPS, + ) + + # wrap output tensor in Float8Tensor + fp8_tensor_col_major_t = Float8Tensor( + output_buffer, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + return fp8_tensor_col_major_t + + +def hp_to_fp8_row_and_col_major( + hp_tensor: torch.Tensor, + fp8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole = GemmInputRole.INPUT, + algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, +) -> Float8Tensor: + assert hp_tensor.is_contiguous(), "input tensor must be contiguous" + + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[fp8_dtype] + + fp8_dtype_min = torch.finfo(fp8_dtype).min + fp8_dtype_max = torch.finfo(fp8_dtype).max + + # compute scaling factor for tensor + scale = _hp_tensor_to_scale( + hp_tensor, + tl_input_dtype, + fp8_dtype_max, + algo, + ) + + # perform fp8 conversion + orig_shape = hp_tensor.shape + num_elements = hp_tensor.numel() + + # preallocate necessary output tensors + fp8_output_row_major = torch.empty( + orig_shape, dtype=fp8_dtype, device=hp_tensor.device + ) + fp8_output_col_major = torch.empty( + orig_shape, dtype=fp8_dtype, device=hp_tensor.device + ) + + # launch triton kernel to perform conversion + num_rows, num_cols = orig_shape + grid = lambda meta: ( + triton.cdiv(num_rows, meta["BLOCK_SIZE_ROWS"]), + triton.cdiv(num_cols, meta["BLOCK_SIZE_COLS"]), + ) + _to_fp8_row_and_col_major[grid]( + hp_tensor, + fp8_output_row_major, + fp8_output_col_major, + scale, + num_elements, + fp8_dtype_min, + fp8_dtype_max, + num_rows, + num_cols, + input_dtype=tl_input_dtype, + output_dtype=tl_output_dtype, + EPS=EPS, + ) + + # for col major we need to update the strides to reflect the new memory layout + col_major_strides = (1, num_rows) + fp8_output_col_major = fp8_output_col_major.as_strided( + fp8_output_col_major.size(), col_major_strides + ) + + # wrap outputs in Float8Tensors + fp8_tensor_row_major = Float8Tensor( + fp8_output_row_major, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + fp8_tensor_col_major = Float8Tensor( + fp8_output_col_major, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + return fp8_tensor_row_major, fp8_tensor_col_major + + +def hp_to_fp8_row_major_t_and_non_t( + hp_tensor: torch.Tensor, + fp8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole = GemmInputRole.INPUT, + algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, +) -> Float8Tensor: + assert hp_tensor.is_contiguous(), "input tensor must be contiguous" + + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[fp8_dtype] + + fp8_dtype_min = torch.finfo(fp8_dtype).min + fp8_dtype_max = torch.finfo(fp8_dtype).max + + # compute scaling factor for tensor + scale = _hp_tensor_to_scale( + hp_tensor, + tl_input_dtype, + fp8_dtype_max, + algo, + ) + + # perform fp8 conversion + input_num_rows, input_num_cols = hp_tensor.shape + transposed_num_rows, transposed_num_cols = input_num_cols, input_num_rows + num_elements = hp_tensor.numel() + + # preallocate necessary output tensors + fp8_output_row_major = torch.empty( + (input_num_rows, input_num_cols), dtype=fp8_dtype, device=hp_tensor.device + ) + fp8_output_row_major_t = torch.empty( + (transposed_num_rows, transposed_num_cols), + dtype=fp8_dtype, + device=hp_tensor.device, + ) + + # launch triton kernel to perform conversion + grid = lambda meta: ( + triton.cdiv(input_num_rows, meta["BLOCK_SIZE_ROWS"]), + triton.cdiv(input_num_cols, meta["BLOCK_SIZE_COLS"]), + ) + _to_fp8_row_major_t_and_non_t[grid]( + hp_tensor, + fp8_output_row_major, + fp8_output_row_major_t, + scale, + num_elements, + fp8_dtype_min, + fp8_dtype_max, + input_num_rows, + input_num_cols, + hp_tensor.stride(0), + hp_tensor.stride(1), + fp8_output_row_major.stride(0), + fp8_output_row_major.stride(1), + fp8_output_row_major_t.stride(0), + fp8_output_row_major_t.stride(1), + input_dtype=tl_input_dtype, + output_dtype=tl_output_dtype, + EPS=EPS, + ) + + # wrap outputs in Float8Tensors + fp8_tensor_row_major = Float8Tensor( + fp8_output_row_major, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + fp8_tensor_row_major_t = Float8Tensor( + fp8_output_row_major_t, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + return fp8_tensor_row_major, fp8_tensor_row_major_t + + +def hp_to_fp8_col_major_t_and_non_t( + hp_tensor: torch.Tensor, + fp8_dtype: torch.dtype, + linear_mm_config: LinearMMConfig, + gemm_input_role: GemmInputRole = GemmInputRole.INPUT, + algo: KernelAlgorithm = KernelAlgorithm.ATOMIC_MAX, +) -> Float8Tensor: + assert hp_tensor.is_contiguous(), "input tensor must be contiguous" + + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[fp8_dtype] + + fp8_dtype_min = torch.finfo(fp8_dtype).min + fp8_dtype_max = torch.finfo(fp8_dtype).max + + # compute scaling factor for tensor + scale = _hp_tensor_to_scale( + hp_tensor, + tl_input_dtype, + fp8_dtype_max, + algo, + ) + + # perform fp8 conversion + input_num_rows, input_num_cols = hp_tensor.shape + num_elements = hp_tensor.numel() + + # preallocate necessary output tensors + fp8_output_col_major = torch.empty( + (input_num_rows, input_num_cols), dtype=fp8_dtype, device=hp_tensor.device + ) + fp8_output_col_major_t = torch.empty_like( + hp_tensor.t(), + dtype=fp8_dtype, + device=hp_tensor.device, + ) + + # launch triton kernel to perform conversion + grid = lambda meta: ( + triton.cdiv(input_num_rows, meta["BLOCK_SIZE_ROWS"]), + triton.cdiv(input_num_cols, meta["BLOCK_SIZE_COLS"]), + ) + _to_fp8_col_major_t_and_non_t[grid]( + hp_tensor, + fp8_output_col_major, + fp8_output_col_major_t, + scale, + num_elements, + fp8_dtype_min, + fp8_dtype_max, + input_num_rows, + input_num_cols, + hp_tensor.stride(0), + hp_tensor.stride(1), + fp8_output_col_major.stride(0), + fp8_output_col_major.stride(1), + fp8_output_col_major_t.stride(0), + fp8_output_col_major_t.stride(1), + input_dtype=tl_input_dtype, + output_dtype=tl_output_dtype, + EPS=EPS, + ) + + # for col major we need to update the strides to reflect the new memory layout + col_major_strides = (1, input_num_rows) + fp8_output_col_major = fp8_output_col_major.as_strided( + fp8_output_col_major.size(), col_major_strides + ) + + # wrap outputs in Float8Tensors + fp8_tensor_col_major = Float8Tensor( + fp8_output_col_major, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + fp8_tensor_col_major_t = Float8Tensor( + fp8_output_col_major_t, + scale, + orig_dtype=hp_tensor.dtype, + linear_mm_config=linear_mm_config, + gemm_input_role=gemm_input_role, + ) + return fp8_tensor_col_major, fp8_tensor_col_major_t + + +def _hp_tensor_to_scale( + hp_tensor: torch.Tensor, + tl_input_dtype: tl.core.dtype, + fp8_dtype_max: float, + algo: KernelAlgorithm, +) -> torch.Tensor: + num_elements = hp_tensor.numel() + scale_out = torch.empty((), dtype=torch.float32, device=hp_tensor.device) + grid = lambda meta: (triton.cdiv(num_elements, meta["BLOCK_SIZE"]),) + + # compute the fp8 scale using the given algorithm + if algo == KernelAlgorithm.ATOMIC_MAX: + global_amax = torch.zeros((1,), dtype=torch.float32, device=hp_tensor.device) + # compute global amax to be used for scaling + _amax_atomic[grid]( + hp_tensor, + global_amax, + num_elements, + input_dtype=tl_input_dtype, + EPS=EPS, + ) + + # compute scale for fp8 conversion + _scale_atomic[1, 1, 1]( + global_amax, + scale_out, + fp8_dtype_max, + EPS=EPS, + ) + + elif algo == KernelAlgorithm.REDUCTION: + # max block size and num warps values determined via manual tuning + max_block_size = 4096 + num_warps = 8 + block_size = triton.next_power_of_2(min(max_block_size, num_elements)) + amax_buffer_size = triton.cdiv(num_elements, block_size) + block_amaxes = torch.zeros( + (amax_buffer_size,), dtype=torch.float32, device=hp_tensor.device + ) + # compute local amax for each block + _amax_reduction[grid]( + hp_tensor, + block_amaxes, + num_elements, + input_dtype=tl_input_dtype, + BLOCK_SIZE=block_size, + EPS=EPS, + num_warps=num_warps, + ) + + # calculate global amax across all blocks and use it to compute scale + _scale_reduction[(1, 1, 1)]( + block_amaxes, + scale_out, + num_elements, + fp8_dtype_max, + BLOCK_SIZE=block_size, + EPS=EPS, + ) + else: + raise ValueError(f"Unsupported kernel algorithm: {algo}") + + return scale_out diff --git a/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise_test.py b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2348877d5c78b008219211cc4b307c77d935ef46 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise_test.py @@ -0,0 +1,493 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import pytest +import torch + +from torchao.float8.float8_scaling_utils import hp_tensor_to_float8_dynamic +from torchao.float8.float8_tensor import LinearMMConfig +from torchao.float8.float8_utils import is_row_major +from torchao.prototype.float8nocompile.kernels.fp8_dynamic_tensorwise import ( + KernelAlgorithm, + hp_to_fp8_col_major, + hp_to_fp8_col_major_t, + hp_to_fp8_col_major_t_and_non_t, + hp_to_fp8_row_and_col_major, + hp_to_fp8_row_major, + hp_to_fp8_row_major_t, + hp_to_fp8_row_major_t_and_non_t, +) + + +@pytest.mark.parametrize( + "algo", + [KernelAlgorithm.REDUCTION, KernelAlgorithm.ATOMIC_MAX], +) +@pytest.mark.parametrize( + "input_shape", + [(2, 4), (32, 16), (512, 512)], +) +def test_fp8_hp_to_fp8_row_major(input_shape: tuple[int, int], algo: KernelAlgorithm): + assert torch.cuda.is_available() + device = "cuda" + input_bf16 = torch.randn(input_shape, dtype=torch.bfloat16, device=device) + x_bf16 = input_bf16.clone().detach().to(device) + y_bf16 = input_bf16.clone().detach().to(device) + + # production implementation + x_fp8_row_major = hp_tensor_to_float8_dynamic( + x_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + ) + + # float8nocompile triton implementation + y_fp8_row_major = hp_to_fp8_row_major( + y_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + algo=algo, + ) + + # check scales + assert torch.eq(x_fp8_row_major._scale, y_fp8_row_major._scale) + + # check data + assert torch.all(torch.eq(x_fp8_row_major._data, y_fp8_row_major._data)) + + # check shapes + assert x_fp8_row_major.shape == y_fp8_row_major.shape + + # check strides + assert x_fp8_row_major.stride() == y_fp8_row_major.stride() + + # check memory layout + assert is_row_major(x_fp8_row_major.stride()) + assert is_row_major(y_fp8_row_major.stride()) + + # check underlying memory layout + assert ( + x_fp8_row_major._data.storage().tolist() + == y_fp8_row_major._data.storage().tolist() + ) + + # assert that error is raised when input tensor is not contiguous + with pytest.raises(AssertionError, match="tensor must be contiguous"): + hp_to_fp8_row_major( + y_bf16.t(), # transpose so tensor memory layout is no longer contiguous + torch.float8_e4m3fn, + LinearMMConfig(), + ) + + +@pytest.mark.parametrize( + "algo", + [KernelAlgorithm.REDUCTION, KernelAlgorithm.ATOMIC_MAX], +) +@pytest.mark.parametrize( + "input_shape", + [(2, 4), (32, 16), (512, 512)], +) +def test_fp8_hp_to_fp8_row_major_t(input_shape: tuple[int, int], algo: KernelAlgorithm): + assert torch.cuda.is_available() + device = "cuda" + input_bf16 = torch.tensor( + [[1, 2, 3], [4, 5, 6]], dtype=torch.bfloat16, device=device + ) + x_bf16 = input_bf16.clone().detach().to(device) + y_bf16 = input_bf16.clone().detach().to(device) + + # production implementation + x_fp8_row_major = hp_tensor_to_float8_dynamic( + x_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + ) + x_fp8_row_major_t = x_fp8_row_major.t().contiguous() + + # float8nocompile triton implementation + y_fp8_row_major_t = hp_to_fp8_row_major_t( + y_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + algo=algo, + ) + + # check scales + assert torch.eq(x_fp8_row_major_t._scale, y_fp8_row_major_t._scale) + + # check data + assert torch.all(torch.eq(x_fp8_row_major_t._data, y_fp8_row_major_t._data)) + + # check shapes + assert x_fp8_row_major_t.shape == y_fp8_row_major_t.shape + + # check strides + assert x_fp8_row_major_t.stride() == y_fp8_row_major_t.stride() + + # check memory layout + assert is_row_major(x_fp8_row_major_t.stride()) + assert is_row_major(y_fp8_row_major_t.stride()) + + # check underlying memory layout + assert ( + x_fp8_row_major_t._data.storage().tolist() + == y_fp8_row_major_t._data.storage().tolist() + ) + + # assert that error is raised when input tensor is not contiguous + with pytest.raises(AssertionError, match="tensor must be contiguous"): + hp_to_fp8_row_major( + y_bf16.t(), # transpose so tensor memory layout is no longer contiguous + torch.float8_e4m3fn, + LinearMMConfig(), + ) + + +@pytest.mark.parametrize( + "algo", + [KernelAlgorithm.REDUCTION, KernelAlgorithm.ATOMIC_MAX], +) +@pytest.mark.parametrize( + "input_shape", + [(2, 4), (32, 16), (512, 512)], +) +def test_fp8_hp_to_fp8_col_major(input_shape: tuple[int, int], algo: KernelAlgorithm): + assert torch.cuda.is_available() + device = "cuda" + input_bf16 = torch.randn(input_shape, dtype=torch.bfloat16, device=device) + x_bf16 = input_bf16.clone().detach().to(device) + y_bf16 = input_bf16.clone().detach().to(device) + + # production implementation + x_fp8_row_major = hp_tensor_to_float8_dynamic( + x_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + ) + x_fp8_col_major = x_fp8_row_major.t().contiguous().t() + + # float8nocompile triton implementation + y_fp8_col_major = hp_to_fp8_col_major( + y_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + algo=algo, + ) + + # check scales + assert torch.eq(x_fp8_col_major._scale, y_fp8_col_major._scale) + + # check data + assert torch.all(torch.eq(x_fp8_col_major._data, y_fp8_col_major._data)) + + # check shapes + assert x_fp8_col_major.shape == y_fp8_col_major.shape + + # check strides + assert x_fp8_col_major.stride() == y_fp8_col_major.stride() + + # check memory layout + assert not is_row_major(x_fp8_col_major.stride()) + assert not is_row_major(y_fp8_col_major.stride()) + + # check underlying memory layout + assert ( + x_fp8_col_major._data.storage().tolist() + == y_fp8_col_major._data.storage().tolist() + ) + + # assert that error is raised when input tensor is not contiguous + with pytest.raises(AssertionError, match="tensor must be contiguous"): + hp_to_fp8_col_major( + y_bf16.t(), # transpose so tensor memory layout is no longer contiguous + torch.float8_e4m3fn, + LinearMMConfig(), + ) + + +@pytest.mark.parametrize( + "algo", + [KernelAlgorithm.REDUCTION, KernelAlgorithm.ATOMIC_MAX], +) +@pytest.mark.parametrize( + "input_shape", + [(2, 4), (32, 16), (512, 512)], +) +def test_fp8_hp_to_fp8_col_major_t(input_shape: tuple[int, int], algo: KernelAlgorithm): + assert torch.cuda.is_available() + device = "cuda" + input_bf16 = torch.randn(input_shape, dtype=torch.bfloat16, device=device) + x_bf16 = input_bf16.clone().detach().to(device) + y_bf16 = input_bf16.clone().detach().to(device) + + # production implementation + x_fp8_row_major = hp_tensor_to_float8_dynamic( + x_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + ) + x_fp8_col_major_t = x_fp8_row_major.contiguous().t() + + # float8nocompile triton implementation + y_fp8_col_major_t = hp_to_fp8_col_major_t( + y_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + algo=algo, + ) + + # check scales + assert torch.eq(x_fp8_col_major_t._scale, y_fp8_col_major_t._scale) + + # check data + assert torch.all(torch.eq(x_fp8_col_major_t._data, y_fp8_col_major_t._data)) + + # check shapes + assert x_fp8_col_major_t.shape == y_fp8_col_major_t.shape + + # check strides + assert x_fp8_col_major_t.stride() == y_fp8_col_major_t.stride() + + # check memory layout + assert not is_row_major(x_fp8_col_major_t.stride()) + assert not is_row_major(y_fp8_col_major_t.stride()) + + # check underlying memory layout + assert ( + x_fp8_col_major_t._data.storage().tolist() + == y_fp8_col_major_t._data.storage().tolist() + ) + + # assert that error is raised when input tensor is not contiguous + with pytest.raises(AssertionError, match="tensor must be contiguous"): + hp_to_fp8_col_major( + y_bf16.t(), # transpose so tensor memory layout is no longer contiguous + torch.float8_e4m3fn, + LinearMMConfig(), + ) + + +@pytest.mark.parametrize( + "algo", + [KernelAlgorithm.REDUCTION, KernelAlgorithm.ATOMIC_MAX], +) +@pytest.mark.parametrize( + "input_shape", + [(2, 4), (32, 16), (512, 512)], +) +def test_fp8_hp_to_fp8_row_and_col_major( + input_shape: tuple[int, int], algo: KernelAlgorithm +): + assert torch.cuda.is_available() + device = "cuda" + input_bf16 = torch.randn(input_shape, dtype=torch.bfloat16, device=device) + x_bf16 = input_bf16.clone().detach().to(device) + y_bf16 = input_bf16.clone().detach().to(device) + + # production implementation + x_fp8_row_major = hp_tensor_to_float8_dynamic( + x_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + ) + x_fp8_col_major = x_fp8_row_major.t().contiguous().t() + + # float8nocompile triton implementation + y_fp8_row_major, y_fp8_col_major = hp_to_fp8_row_and_col_major( + y_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + algo=algo, + ) + + # check scales + assert torch.eq(x_fp8_row_major._scale, y_fp8_row_major._scale) + assert torch.eq(x_fp8_col_major._scale, y_fp8_col_major._scale) + + # check data + assert torch.all(torch.eq(x_fp8_row_major._data, y_fp8_row_major._data)) + assert torch.all(torch.eq(x_fp8_col_major._data, y_fp8_col_major._data)) + + # check shapes + assert x_fp8_row_major.shape == y_fp8_row_major.shape + assert x_fp8_col_major.shape == y_fp8_col_major.shape + + # check strides + assert x_fp8_row_major.stride() == y_fp8_row_major.stride() + assert x_fp8_col_major.stride() == y_fp8_col_major.stride() + + # check memory layout + assert is_row_major(x_fp8_row_major.stride()) + assert is_row_major(y_fp8_row_major.stride()) + assert not is_row_major(x_fp8_col_major.stride()) + assert not is_row_major(y_fp8_col_major.stride()) + + # check underlying memory layout + assert ( + x_fp8_row_major._data.storage().tolist() + == y_fp8_row_major._data.storage().tolist() + ) + assert ( + x_fp8_col_major._data.storage().tolist() + == y_fp8_col_major._data.storage().tolist() + ) + + # assert that error is raised when input tensor is not contiguous + with pytest.raises(AssertionError, match="tensor must be contiguous"): + hp_to_fp8_row_and_col_major( + y_bf16.t(), # transpose so tensor memory layout is no longer contiguous + torch.float8_e4m3fn, + LinearMMConfig(), + ) + + +@pytest.mark.parametrize( + "algo", + [KernelAlgorithm.REDUCTION, KernelAlgorithm.ATOMIC_MAX], +) +@pytest.mark.parametrize( + "input_shape", + [(2, 4), (32, 16), (512, 512)], +) +def test_fp8_hp_to_fp8_row_major_t_and_non_t( + input_shape: tuple[int, int], algo: KernelAlgorithm +): + assert torch.cuda.is_available() + device = "cuda" + input_bf16 = torch.randn(input_shape, dtype=torch.bfloat16, device=device) + x_bf16 = input_bf16.clone().detach().to(device) + y_bf16 = input_bf16.clone().detach().to(device) + + # production implementation + x_fp8_row_major = hp_tensor_to_float8_dynamic( + x_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + ) + x_fp8_row_major_t = x_fp8_row_major.t().contiguous() + + # float8nocompile triton implementation + y_fp8_row_major, y_fp8_row_major_t = hp_to_fp8_row_major_t_and_non_t( + y_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + algo=algo, + ) + + # check scales + assert torch.eq(x_fp8_row_major._scale, y_fp8_row_major._scale) + assert torch.eq(x_fp8_row_major_t._scale, y_fp8_row_major_t._scale) + + # check data + assert torch.all(torch.eq(x_fp8_row_major._data, y_fp8_row_major._data)) + assert torch.all(torch.eq(x_fp8_row_major_t._data, y_fp8_row_major_t._data)) + + # check shapes + assert x_fp8_row_major.shape == y_fp8_row_major.shape + assert x_fp8_row_major_t.shape == y_fp8_row_major_t.shape + + # check strides + assert x_fp8_row_major.stride() == y_fp8_row_major.stride() + assert x_fp8_row_major_t.stride() == y_fp8_row_major_t.stride() + + # check memory layout + assert is_row_major(x_fp8_row_major.stride()) + assert is_row_major(y_fp8_row_major.stride()) + assert is_row_major(x_fp8_row_major_t.stride()) + assert is_row_major(y_fp8_row_major_t.stride()) + + # check underlying memory layout + assert ( + x_fp8_row_major._data.storage().tolist() + == y_fp8_row_major._data.storage().tolist() + ) + assert ( + x_fp8_row_major_t._data.storage().tolist() + == y_fp8_row_major_t._data.storage().tolist() + ) + + # assert that error is raised when input tensor is not contiguous + with pytest.raises(AssertionError, match="tensor must be contiguous"): + hp_to_fp8_row_major_t_and_non_t( + y_bf16.t(), # transpose so tensor memory layout is no longer contiguous + torch.float8_e4m3fn, + LinearMMConfig(), + ) + + +@pytest.mark.parametrize( + "algo", + [KernelAlgorithm.REDUCTION, KernelAlgorithm.ATOMIC_MAX], +) +@pytest.mark.parametrize( + "input_shape", + [(2, 4), (32, 16), (512, 512)], +) +def test_fp8_hp_to_fp8_col_major_t_and_non_t( + input_shape: tuple[int, int], algo: KernelAlgorithm +): + assert torch.cuda.is_available() + device = "cuda" + input_bf16 = torch.randn(input_shape, dtype=torch.bfloat16, device=device) + x_bf16 = input_bf16.clone().detach().to(device) + y_bf16 = input_bf16.clone().detach().to(device) + + # production implementation + x_fp8_row_major = hp_tensor_to_float8_dynamic( + x_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + ) + x_fp8_col_major = x_fp8_row_major.t().contiguous().t() + x_fp8_col_major_t = x_fp8_row_major.t() + + # float8nocompile triton implementation + y_fp8_col_major, y_fp8_col_major_t = hp_to_fp8_col_major_t_and_non_t( + y_bf16, + torch.float8_e4m3fn, + LinearMMConfig(), + algo=algo, + ) + + # check scales + assert torch.eq(x_fp8_col_major._scale, y_fp8_col_major._scale) + assert torch.eq(x_fp8_col_major_t._scale, y_fp8_col_major_t._scale) + + # check data + assert torch.all(torch.eq(x_fp8_col_major._data, y_fp8_col_major._data)) + assert torch.all(torch.eq(x_fp8_col_major_t._data, y_fp8_col_major_t._data)) + + # check shapes + assert x_fp8_col_major.shape == y_fp8_col_major.shape + assert x_fp8_col_major_t.shape == y_fp8_col_major_t.shape + + # check strides + assert x_fp8_col_major.stride() == y_fp8_col_major.stride() + assert x_fp8_col_major_t.stride() == y_fp8_col_major_t.stride() + + # check memory layout + assert not is_row_major(x_fp8_col_major.stride()) + assert not is_row_major(y_fp8_col_major.stride()) + assert not is_row_major(x_fp8_col_major_t.stride()) + assert not is_row_major(y_fp8_col_major_t.stride()) + + # check underlying memory layout + assert ( + x_fp8_col_major._data.storage().tolist() + == y_fp8_col_major._data.storage().tolist() + ) + assert ( + x_fp8_col_major_t._data.storage().tolist() + == y_fp8_col_major_t._data.storage().tolist() + ) + + # assert that error is raised when input tensor is not contiguous + with pytest.raises(AssertionError, match="tensor must be contiguous"): + hp_to_fp8_col_major_t_and_non_t( + y_bf16.t(), # transpose so tensor memory layout is no longer contiguous + torch.float8_e4m3fn, + LinearMMConfig(), + ) diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/__init__.py b/lib/python3.12/site-packages/torchao/prototype/galore/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a769e11a162347b1774a61e64f720a6e3a1a3c3c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/__init__.py @@ -0,0 +1 @@ +from .kernels import * # noqa: F403 diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/galore/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c36dd9b44c16a7193175981b287559d84e813a84 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/galore/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/galore/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..862123b7793e5b94491032ac5589563bd4afc0e1 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/galore/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__init__.py b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..71129d34a28e9228b5c80a4c677084481b9137ef --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__init__.py @@ -0,0 +1,12 @@ +from .adam_downproj_fused import fused_adam_mm_launcher +from .adam_step import triton_adam_launcher +from .matmul import triton_mm_launcher +from .quant import triton_dequant_blockwise, triton_quantize_blockwise + +__all__ = [ + "fused_adam_mm_launcher", + "triton_adam_launcher", + "triton_mm_launcher", + "triton_dequant_blockwise", + "triton_quantize_blockwise", +] diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/adam_downproj_fused.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/adam_downproj_fused.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f9241fa84c519e13d1a038ac8d3f4e92d26bafb Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/adam_downproj_fused.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/adam_step.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/adam_step.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef340de3f4e8cdb6422c8014f0c8348f1036b6af Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/adam_step.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/custom_autotune.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/custom_autotune.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f41b4f02403242762bb818366846cee965cede41 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/custom_autotune.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/matmul.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/matmul.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..090cb2c3b8043599c36e87e40642eca8b1dcf158 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/__pycache__/matmul.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/adam_downproj_fused.py b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/adam_downproj_fused.py new file mode 100644 index 0000000000000000000000000000000000000000..c45fc5d2384f6ac9eb358a710b91fa6ca9549d64 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/adam_downproj_fused.py @@ -0,0 +1,365 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import logging + +import torch +import triton +import triton.language as tl + +from torchao.prototype.common.triton.matmul_perf_model import ( + early_config_prune, + estimate_matmul_time, +) + +from .adam_step import BETA1, BETA2, EPS +from .custom_autotune import Config, autotune +from .matmul import ( + TRITON_ACC_TYPES, + get_higher_dtype, + get_mm_heuristics, + init_to_zero, + to_tl_type, +) +from .matmul import get_autotuner as default_mm_autotuner + +logger = logging.getLogger(__name__) + +AUTOTUNER_TOP_K = 50 + + +def set_tuner_top_k(k): + global AUTOTUNER_TOP_K + AUTOTUNER_TOP_K = k + + +@triton.jit +def _fused_adam_mm_kernel( + # matmul args + A, + B, + C, + M, + N, + K, # + stride_am, + stride_ak, # + stride_bk, + stride_bn, # + stride_cm, + stride_cn, # + # adam epilogue, + exp_avg_ptr, # these will be updated inplace + exp_avg2_ptr, + store, + # grad_ptr, # low rank grad output -- not needed, C is the output + # meta params + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, # + SPLIT_K: tl.constexpr, + EVEN_K: tl.constexpr, + GROUP_M: tl.constexpr, + # Adam-specific params + BETA1: tl.constexpr = BETA1, + BETA2: tl.constexpr = BETA2, + EPS: tl.constexpr = EPS, + # matmul kernel settings + acc_dtype: tl.constexpr = tl.float32, # + allow_tf32: tl.constexpr = False, # higher precision for this phase + fp8_fast_accum: tl.constexpr = False, # + AB_DTYPE: tl.constexpr = None, # +): + # matrix multiplication + pid = tl.program_id(0) + pid_z = tl.program_id(1) + grid_m = tl.cdiv(M, BLOCK_M) + grid_n = tl.cdiv(N, BLOCK_N) + # re-order program ID for better L2 performance + width = GROUP_M * grid_n + group_id = pid // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // (group_size) + # do matrix multiplication + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K) + # pointers + A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) + B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype) + for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)): + if EVEN_K: + a = tl.load(A) + b = tl.load(B) + else: + k_remaining = K - k * (BLOCK_K * SPLIT_K) + _0 = tl.zeros((1, 1), dtype=C.dtype.element_ty) + a = tl.load(A, mask=rk[None, :] < k_remaining, other=_0) + b = tl.load(B, mask=rk[:, None] < k_remaining, other=_0) + if AB_DTYPE is not None: + a = a.to(AB_DTYPE) + b = b.to(AB_DTYPE) + if fp8_fast_accum: + acc = tl.dot(a, b, acc, out_dtype=acc_dtype, allow_tf32=allow_tf32) + else: + acc += tl.dot(a, b, out_dtype=acc_dtype, allow_tf32=allow_tf32) + A += BLOCK_K * SPLIT_K * stride_ak + B += BLOCK_K * SPLIT_K * stride_bk + # acc = acc.to(C.dtype.element_ty) + + # rematerialize rm and rn to save registers + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + epilogue_offsets = rm[:, None] * stride_cm + rn[None, :] * stride_cn + mask = (rm < M)[:, None] & (rn < N)[None, :] + + # Load adam state + exp_avg = tl.load(exp_avg_ptr + epilogue_offsets, mask=mask) + exp_avg2 = tl.load(exp_avg2_ptr + epilogue_offsets, mask=mask) + + # Perform update + exp_avg = BETA1 * exp_avg.to(acc.dtype) + (1.0 - BETA1) * acc + exp_avg2 = BETA2 * exp_avg2.to(acc.dtype) + (1.0 - BETA2) * (acc * acc) + denom = tl.sqrt(exp_avg2) + EPS + norm_grad = exp_avg / denom + # Convert to output type + norm_grad = norm_grad.to(C.dtype.element_ty) + + # acc = acc.to(C.dtype.element_ty) + C = C + epilogue_offsets + + # handles write-back with reduction-splitting + if SPLIT_K == 1: + tl.store(C, norm_grad, mask=mask) + else: + tl.atomic_add(C, norm_grad, mask=mask) + + if store: + tl.store( + exp_avg_ptr + epilogue_offsets, + exp_avg, + mask=mask, + ) + tl.store( + exp_avg2_ptr + epilogue_offsets, + exp_avg2, + mask=mask, + ) + + +def _get_configs_splitk_all(): + """ + Configs specific to split-k matmuls + Not used currently + """ + configs = [] + for num_stages in [2, 3, 4, 5]: + for block_m in [16, 32, 64, 128]: + for block_k in [16, 32, 64, 128, 256]: + for block_n in [16, 32, 64, 128]: + num_warps = 2 if block_n <= 64 else 4 + configs.append( + Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": 1, + }, + num_stages=num_stages, + num_warps=num_warps, + ) + ) + # split_k + for split_k in [2, 4, 8]: + configs.append( + Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": split_k, + }, + num_stages=num_stages, + num_warps=num_warps, + pre_hook=init_to_zero("C"), + ) + ) + return configs + + +def _get_configs_splitk_small(): + """Configs for split-k, smaller version than above + Not used currently + """ + configs = [] + for num_stages in [2, 3, 4]: + for block_m in [64, 128]: + for block_k in [16, 32, 64]: + for block_n in [64, 128]: + num_warps = 2 if block_n <= 64 else 4 + configs.append( + Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": 1, + }, + num_stages=num_stages, + num_warps=num_warps, + ) + ) + # split_k + for split_k in [2, 4, 8]: + configs.append( + Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": split_k, + }, + num_stages=num_stages, + num_warps=num_warps, + pre_hook=init_to_zero("C"), + ) + ) + return configs + + +def _splitk_autotuner( + configs=_get_configs_splitk_small(), + key=["M", "N", "K"], + early_config_prune=early_config_prune, + perf_model=estimate_matmul_time, + top_k=AUTOTUNER_TOP_K, +): + """Autotuner for splitk matmuls + Not used currently + """ + autotuner = autotune( + configs=configs, + key=key, + prune_configs_by={ + "early_config_prune": early_config_prune, + "perf_model": perf_model, + "top_k": top_k, + }, + ) + + return autotuner + + +def _get_kernel( + tuner_fn=default_mm_autotuner, heuristics_fn=get_mm_heuristics, topk=AUTOTUNER_TOP_K +): + tuner = tuner_fn() + tuner.topk = topk + heuristics = heuristics_fn() + return tuner(heuristics(_fused_adam_mm_kernel)) + + +DEFAULT_KERNEL = _get_kernel() + + +def fused_adam_mm_launcher( + a, + b, + *, + exp_avg, + exp_avg2, + store=True, + BETA1=BETA1, + BETA2=BETA2, + EPS=EPS, + allow_tf32=False, + fp8_fast_accum=False, + acc_dtype=None, + output_dtype=None, + kernel=None, +): + device = a.device + # handle non-contiguous inputs if necessary + # a = grad + # b = galore_proj.ortho_matrix.t() + if a.stride(0) > 1 and a.stride(1) > 1: + a = a.contiguous() + if b.stride(0) > 1 and b.stride(1) > 1: + b = b.contiguous() + # checks constraints + assert a.shape[1] == b.shape[0], "incompatible dimensions" + M, K = a.shape + _, N = b.shape + + # common type between a and b + ab_dtype = get_higher_dtype(a.dtype, b.dtype) + + # allocates output + if output_dtype is None: + output_dtype = ab_dtype + + c = torch.empty((M, N), device=device, dtype=output_dtype) + + if acc_dtype is None: + acc_dtype = [ab_dtype][0] + else: + assert isinstance(acc_dtype, torch.dtype), "acc_dtype must be a torch.dtype" + assert acc_dtype in TRITON_ACC_TYPES[a.dtype], ( + "acc_dtype not compatible with the type of a" + ) + assert acc_dtype in TRITON_ACC_TYPES[b.dtype], ( + "acc_dtype not compatible with the type of b" + ) + + acc_dtype = to_tl_type(acc_dtype) + ab_dtype = to_tl_type(ab_dtype) + output_dtype = to_tl_type(output_dtype) + + # Tensor cores support input with mixed float8 types. + if a.dtype in [tl.float8e4nv, tl.float8e5] and b.dtype in [ + tl.float8e4nv, + tl.float8e5, + ]: + ab_dtype = None + + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]), + META["SPLIT_K"], + ) + + if kernel is None: + kernel = DEFAULT_KERNEL + kernel[grid]( + a, + b, + c, + M, + N, + K, # + a.stride(0), + a.stride(1), # + b.stride(0), + b.stride(1), # + c.stride(0), + c.stride(1), # + exp_avg, + exp_avg2, + store=store, + BETA1=BETA1, # , # + BETA2=BETA2, # , # + EPS=EPS, # + acc_dtype=acc_dtype, # + allow_tf32=allow_tf32, # + fp8_fast_accum=fp8_fast_accum, # + GROUP_M=8, + AB_DTYPE=ab_dtype, + ) + return exp_avg, exp_avg2, c # c -> normalized low rank grad diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/adam_step.py b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/adam_step.py new file mode 100644 index 0000000000000000000000000000000000000000..fcb554bd1dac5eac5b7bc013a2bbc5035fe5be2b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/adam_step.py @@ -0,0 +1,181 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +import triton +import triton.language as tl +from triton.language.math import sqrt +from triton.runtime.autotuner import heuristics + +from .custom_autotune import Config, autotune + +BETA1, BETA2 = 0.9, 0.999 +EPS = 1e-8 + +AUTOTUNER_TOP_K = 10 + + +def get_configs_for_adam(num_warps=[2, 4, 8], block_sizes=[512, 1024, 2048]): + configs = [] + for w in num_warps: + for bs in block_sizes: + configs.append(Config({"BLOCK_SIZE": bs}, num_warps=w)) + return configs + + +def early_adam_prune(configs, named_args): + numels = named_args["numels"] + pruned_configs = [cfg for cfg in configs if numels % cfg.kwargs["BLOCK_SIZE"] == 0] + # print("Pruned configs:\n") + for cfg in pruned_configs: + print(f"{cfg}\n") + return pruned_configs + + +def get_adam_tuner( + configs=get_configs_for_adam(), + early_config_prune=None, # early_adam_prune, + top_k=AUTOTUNER_TOP_K, +): + return autotune( + configs=configs, + prune_configs_by={ + "early_config_prune": early_config_prune, + "top_k": top_k, + }, + key=["numels"], + ) + + +def get_adam_heuristics(): + return { + "USE_MASK": lambda args: args["numels"] % args["BLOCK_SIZE"] != 0, + } + + +@autotune(configs=get_configs_for_adam(), key=["numels"]) +@heuristics(get_adam_heuristics()) +@triton.jit +def _adam_update( + avg_ptr, + avg2_ptr, + grad_ptr, + # avg_out_ptr, + # avg2_out_ptr, + # grad_out_ptr, + numels, + store, + BLOCK_SIZE: tl.constexpr, + USE_MASK: tl.constexpr, + BETA1: tl.constexpr = BETA1, + BETA2: tl.constexpr = BETA2, + EPS: tl.constexpr = EPS, +): + pid_m = tl.program_id(0) + offset = pid_m * BLOCK_SIZE + offset = offset + tl.arange(0, BLOCK_SIZE) + # load_idx = offset + tl.arange(0, BLOCK_SIZE) + load_idx = tl.max_contiguous(tl.multiple_of(offset, BLOCK_SIZE), BLOCK_SIZE) + + mask = None + if USE_MASK: + mask = load_idx < numels + avg = tl.load(avg_ptr + load_idx, mask=mask) + avg2 = tl.load(avg2_ptr + load_idx, mask=mask) + grad = tl.load(grad_ptr + load_idx, mask=mask) + + avg = BETA1 * avg + (1.0 - BETA1) * grad + avg2 = BETA2 * avg2 + (1.0 - BETA2) * (grad * grad) + + denom = sqrt(avg2) + EPS + # denom = tl.sqrt(avg2) + EPS + + norm_grad = avg / denom + + if store: + tl.store(avg_ptr + load_idx, avg, mask=mask) + tl.store(avg2_ptr + load_idx, avg2, mask=mask) + tl.store(grad_ptr + load_idx, norm_grad, mask=mask) + # tl.store(avg_out_ptr + load_idx, avg, mask=mask) + # tl.store(avg2_out_ptr + load_idx, avg2, mask=mask) + # tl.store(grad_out_ptr + load_idx, norm_grad, mask=mask) + + +adam_update = _adam_update + + +def triton_adam_launcher( + avg, + avg2, + grad, + store=True, + beta1=BETA1, + beta2=BETA2, + eps=EPS, +): + M, N = avg.shape + # avg_out = torch.empty_like(avg) + # avg2_out = torch.empty_like(avg2) + # grad_out = torch.empty_like(grad) + + grid = lambda META: (triton.cdiv(M * N, META["BLOCK_SIZE"]),) + adam_update[grid]( + avg, + avg2, + grad, + # avg_out, + # avg2_out, + # grad_out, + avg.numel(), + store=store, + BETA1=beta1, + BETA2=beta2, + EPS=eps, + # BLOCK_SIZE=1024, + # USE_MASK=USE_MASK, + ) + return avg, avg2, grad + + +def ref_adam_step(exp_avg, exp_avg2, grad, beta1=BETA1, beta2=BETA2, eps=EPS): + exp_avg = beta1 * exp_avg + (1 - beta1) * grad + exp_avg2 = beta2 * exp_avg2 + (1 - beta2) * torch.square(grad) + denom = exp_avg2.sqrt() + eps + norm_grad = exp_avg / denom + return exp_avg, exp_avg2, norm_grad + + +def make_data(M, N, rank, dtype): + # full_grad = torch.randn(M, N, device="cuda", dtype=dtype) + params = torch.randn(M, N, device="cuda", dtype=dtype) + + if M >= N: + exp_avg = torch.randn(M, rank, device="cuda", dtype=dtype) + else: + exp_avg = torch.randn(rank, N, device="cuda", dtype=dtype) + exp_avg2 = exp_avg**2 + down_grad = torch.randn_like(exp_avg) + + return exp_avg, exp_avg2, down_grad, params + + +if __name__ == "__main__": + M = N = 4096 + rank = 128 + dtype = torch.float32 + exp_avg, exp_avg2, grad, params = make_data(M, N, rank, dtype=dtype) + exp_avg_copy, exp_avg2_copy, grad_copy = ( + exp_avg.clone(), + exp_avg2.clone(), + grad.clone(), + ) + ref_out = ref_adam_step(exp_avg, exp_avg2, grad) + + # Autotune run -- changes exp_avg, exp_avg2, grad in-place + _ = triton_adam_launcher(exp_avg, exp_avg2, grad) + triton_out = triton_adam_launcher(exp_avg_copy, exp_avg2_copy, grad_copy) + + for ref, tt in zip(ref_out, triton_out): + print(torch.max(torch.abs(ref - tt))) diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/custom_autotune.py b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/custom_autotune.py new file mode 100644 index 0000000000000000000000000000000000000000..8ecfabf849b9319f79dfc8c1042cbbd1060fa61b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/custom_autotune.py @@ -0,0 +1,395 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import builtins +import logging +import os +import time +from typing import Dict + +import numpy as np +from triton.runtime.cache import default_cache_dir +from triton.runtime.errors import OutOfResources +from triton.runtime.jit import KernelInterface +from triton.testing import do_bench + +logger = logging.getLogger(__file__) + + +class Autotuner(KernelInterface): + def __init__( + self, + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + prune_configs_by: Dict = None, + warmup=25, + rep=100, + ): + """ + :param prune_configs_by: a dict of functions that are used to prune configs, fields: + 'perf_model': performance model used to predicate running time with different configs, returns running time + 'top_k': number of configs to bench + 'prune_num_stages_by'(optional): a function used to prune num_stages. It takes configs:List[Config] as its input, and returns pruned configs. + """ + if not configs: + self.configs = [Config({}, num_warps=4, num_stages=2, num_ctas=1)] + else: + self.configs = configs + self.key_idx = [arg_names.index(k) for k in key] + self.cache = {} + self.arg_names = arg_names + + # Reset to zero or restore values + self.reset_idx = [] + if reset_to_zero is not None: + self.reset_idx = [arg_names.index(k) for k in reset_to_zero] + self.restore_idx = [] + if restore_value is not None: + self.restore_idx = [arg_names.index(k) for k in restore_value] + + # Hook to reset or restore for required tensors + self.pre_hook = lambda args, reset_only=False: 0 + self.post_hook = lambda args: 0 + if len(self.reset_idx) > 0 or len(self.restore_idx) > 0: + + def _pre_hook(args, reset_only=False): + for i in self.reset_idx: + args[i].zero_() + if not reset_only: + self.restore_copies = [args[i].clone() for i in self.restore_idx] + + self.pre_hook = _pre_hook + if len(self.restore_idx) > 0: + + def _post_hook(args): + for i, j in enumerate(self.restore_idx): + args[j].copy_(self.restore_copies[i]) + self.restore_copies = [] + + self.post_hook = _post_hook + + self.perf_model = None + self.configs_top_k = 1.0 + self.early_config_prune = None + if prune_configs_by: + self.perf_model = prune_configs_by.get("perf_model", self.perf_model) + self.configs_top_k = prune_configs_by.get("top_k", self.configs_top_k) + self.early_config_prune = prune_configs_by.get( + "early_config_prune", self.early_config_prune + ) + + self.fn = fn + self.num_warmups = warmup + self.num_reps = rep + # self.autotune_log_path = os.path.join(default_cache_dir(), autotune_log_file) + self.kernel_name = self._find_kernel_name() + + def _find_kernel_name(self): + try: + kernel_name = self.fn.__name__ + except AttributeError: + try: # in case JITfn is wrapped in both autotune and heuristic + kernel_name = self.fn.fn.__name__ + except: # noqa + kernel_name = self.fn.__name__ + return kernel_name + + def _get_key_combination(self, args, as_str=True, sep=" "): + key_vals = [f"{self.arg_names[i]}={args[i]}" for i in self.key_idx] + return f"{sep}".join(key_vals) if as_str else key_vals + + def _bench(self, *args, config, **meta): + # check for conflicts, i.e. meta-parameters both provided + # as kwargs and by the autotuner + conflicts = meta.keys() & config.kwargs.keys() + if conflicts: + raise ValueError( + f"Conflicting meta-parameters: {', '.join(conflicts)}." + " Make sure that you don't re-define auto-tuned symbols." + ) + # augment meta-parameters with tunable ones + current = dict(meta, **config.kwargs) + full_nargs = {**self.nargs, **current} + + def kernel_call(): + if config.pre_hook: + config.pre_hook(full_nargs) + self.pre_hook(args) + self.fn.run( + *args, + num_warps=config.num_warps, + num_stages=config.num_stages, + num_ctas=config.num_ctas, + **current, + ) + self.post_hook(args) + + try: + return do_bench( + kernel_call, + warmup=self.num_warmups, + rep=self.num_reps, + quantiles=(0.5, 0.2, 0.8), + ) + except OutOfResources: + return [float("inf"), float("inf"), float("inf")] + + def run(self, *args, **kwargs): + self.nargs = dict(zip(self.arg_names, args)) + if len(self.configs) > 1: + all_args = {**self.nargs, **kwargs} + _args = [] + for name in self.arg_names: + if name in all_args: + _args.append(all_args[name]) + key = [_args[i] for i in self.key_idx] + for arg in _args: + if hasattr(arg, "dtype"): + key.append(str(arg.dtype)) + key = tuple(key) + if key not in self.cache: + logger.debug("Cache miss!\n") + logger.info( + f"\n==== Autotune ====\nRunning autotune for {self.kernel_name} for {len(self.configs)} total configs" + f" for key combination {self._get_key_combination(args)}..." + ) + # prune configs + pruned_configs = self.prune_configs(kwargs) + logger.info(f"\nNum configs after pruning {len(pruned_configs)}") + bench_start = time.time() + timings = {} + for config in pruned_configs: + timings[config] = self._bench(*args, config=config, **kwargs) + # timings = { + # config: self._bench(*args, config=config, **kwargs) + # for config in pruned_configs + # } + bench_end = time.time() + self.bench_time = bench_end - bench_start + self.cache[key] = builtins.min(timings, key=timings.get) + self.pre_hook(args, reset_only=True) + self.configs_timings = timings + + sorted_timings = dict( + sorted(timings.items(), key=lambda x: np.mean(x[1])) + ) + _key_suffix = self._get_key_combination(args, sep="-") + autotune_file = f"autotune_{self.kernel_name}_{_key_suffix}.log" + autotune_log_path = os.path.join(default_cache_dir(), autotune_file) + + logger.info(f"\nFinished autotune, writing log to {autotune_log_path}") + + with open(f"{autotune_log_path}", "w") as f: + f.write( + f" ==== Autotune Results ====\nKernel name: {self.kernel_name}\nArgs: {self.arg_names}\nKeys: {self._get_key_combination(args)}\n" + ) + f.write("\nPruned configs:\n") + for cfg in pruned_configs: + f.write(f"{cfg}\n") + f.write("Timings:\n") + for cfg, timing in sorted_timings.items(): + f.write(f"{cfg} {timing} \n") + f.write(f"Best config: {self.cache[key]}\n") + config = self.cache[key] + logger.debug("\nAutotune: Cache hit! Running best config...") + else: + config = self.configs[0] + self.best_config = config + logger.info(f"\nAutotune Best Config: {config}\n") + + full_nargs = {**self.nargs, **kwargs, **self.best_config.kwargs} + if config.pre_hook is not None: + config.pre_hook(full_nargs) + ret = self.fn.run( + *args, + num_warps=config.num_warps, + num_stages=config.num_stages, + num_ctas=config.num_ctas, + **kwargs, + **config.kwargs, + ) + self.nargs = None + return ret + + def prune_configs(self, kwargs): + pruned_configs = self.configs + if self.early_config_prune: + pruned_configs = self.early_config_prune(self.configs, self.nargs) + if self.perf_model: + top_k = self.configs_top_k + if isinstance(top_k, float) and top_k <= 1.0: + top_k = int(len(self.configs) * top_k) + if len(pruned_configs) > top_k: + est_timing = { + config: self.perf_model( + **self.nargs, + **kwargs, + **config.kwargs, + num_stages=config.num_stages, + num_warps=config.num_warps, + num_ctas=config.num_ctas, + ) + for config in pruned_configs + } + pruned_configs = sorted(est_timing.keys(), key=lambda x: est_timing[x])[ + :top_k + ] + return pruned_configs + + def warmup(self, *args, **kwargs): + self.nargs = dict(zip(self.arg_names, args)) + ret = [] + for config in self.prune_configs(kwargs): + ret.append( + self.fn.warmup( + *args, + num_warps=config.num_warps, + num_ctas=config.num_ctas, + num_stages=config.num_stages, + **kwargs, + **config.kwargs, + ) + ) + self.nargs = None + return ret + + +class Config: + """ + An object that represents a possible kernel configuration for the auto-tuner to try. + + :ivar meta: a dictionary of meta-parameters to pass to the kernel as keyword arguments. + :type meta: dict[Str, Any] + :ivar num_warps: the number of warps to use for the kernel when compiled for GPUs. For example, if + `num_warps=8`, then each kernel instance will be automatically parallelized to + cooperatively execute using `8 * 32 = 256` threads. + :type num_warps: int + :ivar num_stages: the number of stages that the compiler should use when software-pipelining loops. + Mostly useful for matrix multiplication workloads on SM80+ GPUs. + :type num_ctas: int + :ivar num_ctas: number of blocks in a block cluster. SM90+ only. + :ivar pre_hook: a function that will be called before the kernel is called. Parameters of this + function are args. + """ + + def __init__(self, kwargs, num_warps=4, num_stages=2, num_ctas=1, pre_hook=None): + self.kwargs = kwargs + self.num_warps = num_warps + self.num_ctas = num_ctas + self.num_stages = num_stages + self.pre_hook = pre_hook + + def __str__(self): + res = [] + for k, v in self.kwargs.items(): + res.append(f"{k}: {v}") + res.append(f"num_warps: {self.num_warps}") + res.append(f"num_ctas: {self.num_ctas}") + res.append(f"num_stages: {self.num_stages}") + return ", ".join(res) + + +def autotune( + configs, + key, + prune_configs_by=None, + reset_to_zero=None, + restore_value=None, + warmup=25, + rep=100, +): + """ + Decorator for auto-tuning a :code:`triton.jit`'d function. + + .. highlight:: python + .. code-block:: python + + @triton.autotune(configs=[ + triton.Config(meta={'BLOCK_SIZE': 128}, num_warps=4), + triton.Config(meta={'BLOCK_SIZE': 1024}, num_warps=8), + ], + key=['x_size'] # the two above configs will be evaluated anytime + # the value of x_size changes + ) + @triton.jit + def kernel(x_ptr, x_size, **META): + BLOCK_SIZE = META['BLOCK_SIZE'] + :note: When all the configurations are evaluated, the kernel will run multiple times. + This means that whatever value the kernel updates will be updated multiple times. + To avoid this undesired behavior, you can use the `reset_to_zero` argument, which + resets the value of the provided tensor to `zero` before running any configuration. + :param configs: a list of :code:`triton.Config` objects + :type configs: list[triton.Config] + :param key: a list of argument names whose change in value will trigger the evaluation of all provided configs. + :type key: list[str] + :param prune_configs_by: a dict of functions that are used to prune configs, fields: + 'perf_model': performance model used to predicate running time with different configs, returns running time + 'top_k': number of configs to bench + 'early_config_prune'(optional): a function used to do early prune (eg, num_stages). It takes configs:List[Config] as its input, and returns pruned configs. + :param reset_to_zero: a list of argument names whose value will be reset to zero before evaluating any configs. + :type reset_to_zero: list[str] + :param restore_value: a list of argument names whose value will be restored after evaluating any configs. + :type restore_value: list[str] + :param warmup: Warmup time (in ms) to pass to benchmarking, defaults to 25. + :type warmup: int + :param rep: Repetition time (in ms) to pass to benchmarking, defaults to 100. + :type rep: int + """ + + def decorator(fn): + return Autotuner( + fn, + fn.arg_names, + configs, + key, + reset_to_zero, + restore_value, + prune_configs_by, + warmup, + rep, + ) + + return decorator + + +class Heuristics(KernelInterface): + def __init__(self, fn, arg_names, values) -> None: + self.fn = fn + self.values = values + self.arg_names = arg_names + + def run(self, *args, **kwargs): + for v, heur in self.values.items(): + kwargs[v] = heur({**dict(zip(self.arg_names, args)), **kwargs}) + return self.fn.run(*args, **kwargs) + + +def heuristics(values): + """ + Decorator for specifying how the values of certain meta-parameters may be computed. + This is useful for cases where auto-tuning is prohibitevely expensive, or just not applicable. + + .. highlight:: python + .. code-block:: python + + @triton.heuristics(values={'BLOCK_SIZE': lambda args: 2 ** int(math.ceil(math.log2(args[1])))}) + @triton.jit + def kernel(x_ptr, x_size, **META): + BLOCK_SIZE = META['BLOCK_SIZE'] # smallest power-of-two >= x_size + :param values: a dictionary of meta-parameter names and functions that compute the value of the meta-parameter. + each such function takes a list of positional arguments as input. + :type values: dict[str, Callable[[list[Any]], Any]] + """ + + def decorator(fn): + return Heuristics(fn, fn.arg_names, values) + + return decorator diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/matmul.py b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/matmul.py new file mode 100644 index 0000000000000000000000000000000000000000..0a7c830f02edc9127cba8518ad527e3f99ea5350 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/matmul.py @@ -0,0 +1,417 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +import triton +import triton.language as tl + +from torchao.prototype.common.triton.matmul_perf_model import ( + early_config_prune, + estimate_matmul_time, +) + +from .custom_autotune import Config, autotune, heuristics + +# Allowed types for acc_type given the types of a and b. +TRITON_ACC_TYPES = { + torch.float16: (torch.float32, torch.float16), + torch.bfloat16: (torch.float32, torch.bfloat16), + torch.float32: (torch.float32,), + torch.int8: (torch.int32,), +} + +AUTOTUNER_TOP_K = 50 +_ordered_datatypes = [torch.int8, torch.float16, torch.bfloat16, torch.float32] + + +def upcast_if_fp8(a): + if "fp8" in str(a): + return torch.float16 + return a + + +def get_higher_dtype(a, b): + a = upcast_if_fp8(a) + b = upcast_if_fp8(b) + if a is b: + return a + + assert a in _ordered_datatypes + assert b in _ordered_datatypes + + for d in _ordered_datatypes: + if a is d: + return b + if b is d: + return a + + +def init_to_zero(name): + return lambda nargs: nargs[name].zero_() + + +def set_tuner_top_k(k): + global AUTOTUNER_TOP_K + AUTOTUNER_TOP_K = k + + +def to_tl_type(ty): + return getattr(tl, str(ty).split(".")[-1]) + + +def get_configs_io_bound(): + configs = [] + for num_stages in [2, 3, 4, 5, 6]: + for block_m in [16, 32]: + for block_k in [32, 64]: + for block_n in [32, 64, 128, 256]: + num_warps = 2 if block_n <= 64 else 4 + configs.append( + Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": 1, + }, + num_stages=num_stages, + num_warps=num_warps, + ) + ) + # split_k + for split_k in [2, 4, 8, 16]: + configs.append( + Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": split_k, + }, + num_stages=num_stages, + num_warps=num_warps, + pre_hook=init_to_zero("C"), + ) + ) + return configs + + +def get_configs_compute_bound(): + configs = [ + # basic configs for compute-bound matmuls + Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + Config( + {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + Config( + {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=5, + num_warps=2, + ), + # good for int8 + Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + Config( + {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + Config( + {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, + num_stages=5, + num_warps=2, + ), + ] + return configs + + +def get_autotuner( + configs=get_configs_compute_bound() + get_configs_io_bound(), + key=["M", "N", "K"], + early_config_prune=early_config_prune, + perf_model=estimate_matmul_time, + top_k=AUTOTUNER_TOP_K, +): + autotuner = autotune( + configs=configs, + key=key, + prune_configs_by={ + "early_config_prune": early_config_prune, + "perf_model": perf_model, + "top_k": top_k, + }, + ) + + return autotuner + + +def get_mm_heuristics(): + return heuristics( + { + "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, + } + ) + + +@triton.jit +def _matmul_kernel( + A, + B, + C, + M, + N, + K, # + stride_am, + stride_ak, # + stride_bk, + stride_bn, # + stride_cm, + stride_cn, # + # meta params + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, # + SPLIT_K: tl.constexpr, + EVEN_K: tl.constexpr, + GROUP_M: tl.constexpr, + # epilogue + epilogue_alpha=None, + epilogue_beta=None, + epilogue_source=None, # Corresponds to C in GEMM convention of D = AB + C + # matmul kernel settings + acc_dtype: tl.constexpr = tl.float32, # + allow_tf32: tl.constexpr = True, # + fp8_fast_accum: tl.constexpr = True, # + AB_DTYPE: tl.constexpr = None, # + EPILOGUE: tl.constexpr = False, +): + # matrix multiplication + pid = tl.program_id(0) + pid_z = tl.program_id(1) + grid_m = tl.cdiv(M, BLOCK_M) + grid_n = tl.cdiv(N, BLOCK_N) + # re-order program ID for better L2 performance + width = GROUP_M * grid_n + group_id = pid // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // (group_size) + # do matrix multiplication + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K) + # pointers + A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) + B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype) + for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)): + if EVEN_K: + a = tl.load(A) + b = tl.load(B) + else: + k_remaining = K - k * (BLOCK_K * SPLIT_K) + _0 = tl.zeros((1, 1), dtype=C.dtype.element_ty) + a = tl.load(A, mask=rk[None, :] < k_remaining, other=_0) + b = tl.load(B, mask=rk[:, None] < k_remaining, other=_0) + if AB_DTYPE is not None: + a = a.to(AB_DTYPE) + b = b.to(AB_DTYPE) + if fp8_fast_accum: + acc = tl.dot(a, b, acc, out_dtype=acc_dtype, allow_tf32=allow_tf32) + else: + acc += tl.dot(a, b, out_dtype=acc_dtype, allow_tf32=allow_tf32) + A += BLOCK_K * SPLIT_K * stride_ak + B += BLOCK_K * SPLIT_K * stride_bk + # acc = acc.to(C.dtype.element_ty) + + # rematerialize rm and rn to save registers + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + if EPILOGUE: + if epilogue_alpha is not None: + acc = epilogue_alpha.to(acc_dtype) * acc + if epilogue_source is not None: + epilogue_src = tl.load( + epilogue_source + rm[:, None] * stride_cm + rn[None, :] * stride_cn + ) + if epilogue_beta is not None: + epilogue_src = epilogue_src.to(acc_dtype) * epilogue_beta.to(acc_dtype) + acc = acc + epilogue_src + + acc = acc.to(C.dtype.element_ty) + C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn) + mask = (rm < M)[:, None] & (rn < N)[None, :] + # handles write-back with reduction-splitting + if SPLIT_K == 1: + tl.store(C, acc, mask=mask) + else: + tl.atomic_add(C, acc, mask=mask) + + +_autotuner = get_autotuner() +_heuristics = get_mm_heuristics() +matmul = _autotuner(_heuristics(_matmul_kernel)) + + +def triton_mm_launcher( + a, + b, + epilogue_alpha=None, + epilogue_beta=None, + epilogue_source=None, + allow_tf32=True, + fp8_fast_accum=True, + acc_dtype=None, + output_dtype=None, + kernel=matmul, +): + device = a.device + # handle non-contiguous inputs if necessary + # a = grad + # b = galore_proj.ortho_matrix.t() + if a.stride(0) > 1 and a.stride(1) > 1: + a = a.contiguous() + if b.stride(0) > 1 and b.stride(1) > 1: + b = b.contiguous() + # checks constraints + assert a.shape[1] == b.shape[0], "incompatible dimensions" + M, K = a.shape + _, N = b.shape + + # common type between a and b + ab_dtype = get_higher_dtype(a.dtype, b.dtype) + + # allocates output + if output_dtype is None: + output_dtype = ab_dtype + + c = torch.empty((M, N), device=device, dtype=output_dtype) + + if acc_dtype is None: + acc_dtype = [ab_dtype][0] + else: + assert isinstance(acc_dtype, torch.dtype), "acc_dtype must be a torch.dtype" + assert acc_dtype in TRITON_ACC_TYPES[a.dtype], ( + "acc_dtype not compatible with the type of a" + ) + assert acc_dtype in TRITON_ACC_TYPES[b.dtype], ( + "acc_dtype not compatible with the type of b" + ) + + acc_dtype = to_tl_type(acc_dtype) + ab_dtype = to_tl_type(ab_dtype) + output_dtype = to_tl_type(output_dtype) + + # Tensor cores support input with mixed float8 types. + if a.dtype in [tl.float8e4nv, tl.float8e5] and b.dtype in [ + tl.float8e4nv, + tl.float8e5, + ]: + ab_dtype = None + # launch kernel + # print( + # f"{__file__} triton matmul args: (AB dtype {ab_dtype}) (C dtype {c.dtype}) (allow_tf32 {allow_tf32}) (fp8_fast_accum {fp8_fast_accum})" + # ) + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]), + META["SPLIT_K"], + ) + + matmul[grid]( + a, + b, + c, + M, + N, + K, # + a.stride(0), + a.stride(1), # + b.stride(0), + b.stride(1), # + c.stride(0), + c.stride(1), # + epilogue_alpha=epilogue_alpha, # + epilogue_beta=epilogue_beta, # + epilogue_source=epilogue_source, # + acc_dtype=acc_dtype, # + allow_tf32=allow_tf32, # + fp8_fast_accum=fp8_fast_accum, # + GROUP_M=8, + AB_DTYPE=ab_dtype, + EPILOGUE=any([epilogue_alpha, epilogue_beta, epilogue_source]), + ) + return c diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/kernels/quant.py b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/quant.py new file mode 100644 index 0000000000000000000000000000000000000000..7c8a58e864c4a17ffd0351abfd507ceafb89a6a1 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/kernels/quant.py @@ -0,0 +1,189 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +import triton +import triton.language as tl + + +@triton.jit +def _dequant_kernel( + q_idx_ptr, + absmax_ptr, + qmap_ptr, + dq_ptr, + stride_qm, + stride_qn, + M, + N, + GROUP_SIZE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + # rm = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + # rn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + offsets = rm[:, None] * stride_qm + rn[None, :] * stride_qn + mask = (rm[:, None] < M) & (rn[None, :] < N) + tl.static_print(offsets) + group_offsets = offsets // GROUP_SIZE + tl.static_print("group_offsets", group_offsets) + q_idx = tl.load(q_idx_ptr + offsets, mask=mask) + tl.static_print(q_idx) + # NOTE: Must upcast q_idx to int32 (q_idx is tl.uint8, which does not work for pointer indexing) + q_vals = tl.load(qmap_ptr + q_idx.to(tl.int32)) + absmax = tl.load( + absmax_ptr + group_offsets, mask=group_offsets < (M * N // GROUP_SIZE) + ) + + dq = q_vals * absmax + tl.store(dq_ptr + offsets, dq, mask=mask) + + +def triton_dequant_blockwise( + q: torch.Tensor, qmap: torch.Tensor, absmax: torch.Tensor, group_size: int +): + M, N = q.shape + dq = torch.empty_like(q).to(absmax.dtype) + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_M"]), + triton.cdiv(N, META["BLOCK_N"]), + ) + _dequant_kernel[grid]( + q, + absmax, + qmap, + dq, + q.stride(0), + q.stride(1), + M, + N, + BLOCK_M=1, + BLOCK_N=group_size, + GROUP_SIZE=group_size, + ) + return dq + + +@triton.heuristics( + values={ + "USE_MASK": lambda args: args["numels"] % args["BLOCK_SIZE"] != 0, + "NUM_GROUPS": lambda args: triton.cdiv(args["numels"], args["BLOCK_SIZE"]), + } +) +@triton.jit +def _quantize_blockwise_kernel( + t_ptr, + cutoffs_ptr, + q_ptr, + absmax_ptr, + norm_ptr, + numels, + BLOCK_SIZE: tl.constexpr, + NUM_BUCKETS: tl.constexpr, + USE_MASK: tl.constexpr, + NUM_GROUPS: tl.constexpr, + RETURN_NORM: tl.constexpr = False, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = None + absmax_mask = None + if USE_MASK: + mask = offsets < numels + absmax_mask = pid < NUM_GROUPS + t = tl.load(t_ptr + offsets, mask=mask) + + absmax = tl.max(tl.abs(t), axis=0) + normalized = t / absmax + + # Load code buckets + cutoffs = tl.load(cutoffs_ptr + tl.arange(0, NUM_BUCKETS)) + q = tl.reshape(normalized, (BLOCK_SIZE, 1)) > cutoffs + + # NOTE: explicit cast is needed, addition on tl.int1 (bool) does not work as per torch / numpy + q = q.to(tl.uint8) + q = tl.sum(q, axis=1) + + tl.store(q_ptr + offsets, q, mask=mask) + # Each block processes one group_size number of elements, hence 1 absmax + tl.store(absmax_ptr + pid, absmax, mask=absmax_mask) + + if RETURN_NORM: + tl.store(norm_ptr + offsets, normalized, mask=mask) + + +# NOTE: Each block processes one group_size number of elements, hence BLOCK_SIZE = group_size +# where group_size corresponds to the groupwise quantization blocksize +def triton_quantize_blockwise( + t: torch.Tensor, code, group_size=2048, return_normalized=False +): + """ + Params: + t: torch.Tensor, tensor to quantize + code: torch.Tensor, quantization codebook for bitsandbytes, output of `bitsandbytes.functional.create_dynamic_map` + # absmax: torch.Tensor, absolute max values for each block, if None, will be calculated from the input tensor + group_size: int, groupwise quantization blocksize, default 2048, the hardcoded blocksize for bitsandbytes 8-bit optimizers + return_normalized: bool, if True, will return the normalized tensor, primarily for debugging + """ + numel = t.numel() + q = torch.empty(numel, dtype=torch.uint8, device=t.device) + normalized = torch.empty_like(t) if return_normalized else None + num_groups = numel // group_size + abs_max = torch.empty(num_groups, dtype=t.dtype, device="cuda") + # Cutoffs for quantization + # code corresponds to actual (normalized) quant codes + # Cutoffs are used to calculate which of these codes a value belongs to + # E.g., for consecutive codes C1 and C2, the corresponding cutoff is C1 + C2 / 2 + # Hence, if a value is greater is assigned C1 if it is less than all cutoffs up to this cutoff + cutoffs = (code[:-1] + code[1:]) / 2 + + # Need to make cutoffs multiple of 2 for triton reduce + MAX_CUTOFF = torch.tensor( + torch.finfo(cutoffs.dtype).max, dtype=cutoffs.dtype, device=cutoffs.device + ).reshape( + 1, + ) + cutoffs = torch.cat([cutoffs, MAX_CUTOFF], dim=-1) + assert cutoffs.numel() % 2 == 0 + + grid = lambda META: (triton.cdiv(t.numel(), META["BLOCK_SIZE"]),) + # assert t.numel() % group_size == 0 + _quantize_blockwise_kernel[grid]( + t.view(-1), + cutoffs, + q, + abs_max, + normalized.view(-1) if return_normalized else None, + numel, + NUM_BUCKETS=len(cutoffs), + BLOCK_SIZE=group_size, + RETURN_NORM=return_normalized, + ) + return ( + q.reshape(t.shape), + normalized.reshape(t.shape) if return_normalized else None, + abs_max, + ) + + +# Reference implementation +def _torch_quantize_blockwise(tensor: torch.Tensor, code, absmax=None, blocksize=2048): + # Flatten values first + + # If not absmax, need to first normalize -> reshape to (-1, blocksize) -> max over the last dim + + # Quantize by flattening A to [numels, 1] > code[:, None], sum, then reshape back to original shape + if absmax is None: + absmax = tensor.reshape(-1, blocksize).abs().max(dim=-1).values + + normalized = tensor.reshape(-1, blocksize) / absmax[:, None] + buckets = (code[:-1] + code[1:]) / 2 + q = normalized.reshape(normalized.numel(), 1) > buckets + q = q.sum(dim=1).reshape(tensor.shape) + return q.to(torch.uint8), normalized.reshape(tensor.shape), absmax diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/optim/__init__.py b/lib/python3.12/site-packages/torchao/prototype/galore/optim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/optim/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/galore/optim/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..511305ecfccbcf06a186f5def20ca898cb1e62f8 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/galore/optim/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/optim/__pycache__/galore_torch.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/galore/optim/__pycache__/galore_torch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9dcf1e11ca9a5eb52f649ac6899a42bb9e411b50 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/galore/optim/__pycache__/galore_torch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/optim/galore_torch.py b/lib/python3.12/site-packages/torchao/prototype/galore/optim/galore_torch.py new file mode 100644 index 0000000000000000000000000000000000000000..876c40292d869bfff5ee934da0777af788e19cbb --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/optim/galore_torch.py @@ -0,0 +1,401 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +"""Reference implementations +See https://github.com/jiaweizzhao/GaLore/tree/master/galore_torch +""" + +# copy dependencies from transformers/optimization.py +import math +import warnings +from typing import Callable, Iterable, Tuple + +import torch +from bitsandbytes.optim.optimizer import Optimizer2State +from torch import nn +from torch.optim import Optimizer + + +class GaLoreProjector: + def __init__( + self, rank, verbose=False, update_proj_gap=200, scale=1.0, proj_type="std" + ): + self.rank = rank + self.verbose = verbose + self.update_proj_gap = update_proj_gap + self.scale = scale + self.ortho_matrix = None + self.proj_type = proj_type + + def project(self, full_rank_grad, iter): + if self.proj_type == "std": + if full_rank_grad.shape[0] >= full_rank_grad.shape[1]: + if self.ortho_matrix is None or iter % self.update_proj_gap == 0: + self.ortho_matrix = self.get_orthogonal_matrix( + full_rank_grad, self.rank, type="right" + ) + low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t()) + else: + if self.ortho_matrix is None or iter % self.update_proj_gap == 0: + self.ortho_matrix = self.get_orthogonal_matrix( + full_rank_grad, self.rank, type="left" + ) + low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad) + elif self.proj_type == "reverse_std": + if full_rank_grad.shape[0] >= full_rank_grad.shape[1]: + if self.ortho_matrix is None or iter % self.update_proj_gap == 0: + self.ortho_matrix = self.get_orthogonal_matrix( + full_rank_grad, self.rank, type="left" + ) + low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad) + else: + if self.ortho_matrix is None or iter % self.update_proj_gap == 0: + self.ortho_matrix = self.get_orthogonal_matrix( + full_rank_grad, self.rank, type="right" + ) + low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t()) + elif self.proj_type == "right": + if self.ortho_matrix is None or iter % self.update_proj_gap == 0: + self.ortho_matrix = self.get_orthogonal_matrix( + full_rank_grad, self.rank, type="right" + ) + low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t()) + elif self.proj_type == "left": + if self.ortho_matrix is None or iter % self.update_proj_gap == 0: + self.ortho_matrix = self.get_orthogonal_matrix( + full_rank_grad, self.rank, type="left" + ) + low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad) + elif self.proj_type == "full": + if self.ortho_matrix is None or iter % self.update_proj_gap == 0: + self.ortho_matrix = self.get_orthogonal_matrix( + full_rank_grad, self.rank, type="full" + ) + low_rank_grad = ( + torch.matmul(self.ortho_matrix[0].t(), full_rank_grad) + @ self.ortho_matrix[1].t() + ) + + return low_rank_grad + + def project_back(self, low_rank_grad): + if self.proj_type == "std": + if low_rank_grad.shape[0] >= low_rank_grad.shape[1]: + full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix) + else: + full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad) + elif self.proj_type == "reverse_std": + if ( + low_rank_grad.shape[0] <= low_rank_grad.shape[1] + ): # note this is different from std + full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad) + else: + full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix) + elif self.proj_type == "right": + full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix) + elif self.proj_type == "left": + full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad) + elif self.proj_type == "full": + full_rank_grad = ( + torch.matmul(self.ortho_matrix[0], low_rank_grad) @ self.ortho_matrix[1] + ) + + return full_rank_grad * self.scale + + # svd decomposition + def get_orthogonal_matrix(self, weights, rank, type): + module_params = weights + + if module_params.data.dtype != torch.float: + float_data = False + original_type = module_params.data.dtype + original_device = module_params.data.device + matrix = module_params.data.float() + else: + float_data = True + matrix = module_params.data + + U, s, Vh = torch.linalg.svd(matrix, full_matrices=False) + + # make the smaller matrix always to be orthogonal matrix + if type == "right": + # A = U[:, :rank] @ torch.diag(s[:rank]) + B = Vh[:rank, :] + + if not float_data: + B = B.to(original_device).type(original_type) + return B + elif type == "left": + A = U[:, :rank] + # B = torch.diag(s[:rank]) @ Vh[:rank, :] + if not float_data: + A = A.to(original_device).type(original_type) + return A + elif type == "full": + A = U[:, :rank] + B = Vh[:rank, :] + if not float_data: + A = A.to(original_device).type(original_type) + B = B.to(original_device).type(original_type) + return [A, B] + else: + raise ValueError("type should be left, right or full") + + +class AdamW(Optimizer): + """ + Implements Adam algorithm with weight decay fix as introduced in [Decoupled Weight Decay + Regularization](https://arxiv.org/abs/1711.05101). + + Parameters: + params (`Iterable[nn.parameter.Parameter]`): + Iterable of parameters to optimize or dictionaries defining parameter groups. + lr (`float`, *optional*, defaults to 0.001): + The learning rate to use. + betas (`Tuple[float,float]`, *optional*, defaults to `(0.9, 0.999)`): + Adam's betas parameters (b1, b2). + eps (`float`, *optional*, defaults to 1e-06): + Adam's epsilon for numerical stability. + weight_decay (`float`, *optional*, defaults to 0.0): + Decoupled weight decay to apply. + correct_bias (`bool`, *optional*, defaults to `True`): + Whether or not to correct bias in Adam (for instance, in Bert TF repository they use `False`). + no_deprecation_warning (`bool`, *optional*, defaults to `False`): + A flag used to disable the deprecation warning (set to `True` to disable the warning). + """ + + def __init__( + self, + params: Iterable[nn.parameter.Parameter], + lr: float = 1e-3, + betas: Tuple[float, float] = (0.9, 0.999), + eps: float = 1e-6, + weight_decay: float = 0.0, + correct_bias: bool = True, + no_deprecation_warning: bool = False, + ): + if not no_deprecation_warning: + warnings.warn( + "This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch" + " implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this" + " warning", + FutureWarning, + ) + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr} - should be >= 0.0") + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)" + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)" + ) + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps} - should be >= 0.0") + defaults = { + "lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": weight_decay, + "correct_bias": correct_bias, + } + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure: Callable = None): + """ + Performs a single optimization step. + + Arguments: + closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if grad.is_sparse: + raise RuntimeError( + "Adam does not support sparse gradients, please consider SparseAdam instead" + ) + + state = self.state[p] + + if "step" not in state: + state["step"] = 0 + + # GaLore Projection + if "rank" in group: + if "projector" not in state: + state["projector"] = GaLoreProjector( + group["rank"], + update_proj_gap=group["update_proj_gap"], + scale=group["scale"], + proj_type=group["proj_type"], + ) + + grad = state["projector"].project(grad, state["step"]) + + # State initialization + if "exp_avg" not in state: + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(grad) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like(grad) + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + beta1, beta2 = group["betas"] + + state["step"] += 1 + + # Decay the first and second moment running average coefficient + # In-place operations to update the averages at the same time + exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1)) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) + denom = exp_avg_sq.sqrt().add_(group["eps"]) + + step_size = group["lr"] + if group["correct_bias"]: # No bias correction for Bert + bias_correction1 = 1.0 - beta1 ** state["step"] + bias_correction2 = 1.0 - beta2 ** state["step"] + step_size = ( + step_size * math.sqrt(bias_correction2) / bias_correction1 + ) + + # compute norm gradient + norm_grad = exp_avg / denom + + # GaLore Projection Back + if "rank" in group: + norm_grad = state["projector"].project_back(norm_grad) + + p.add_(norm_grad, alpha=-step_size) + + # Just adding the square of the weights to the loss function is *not* + # the correct way of using L2 regularization/weight decay with Adam, + # since that will interact with the m and v parameters in strange ways. + # + # Instead we want to decay the weights in a manner that doesn't interact + # with the m/v parameters. This is equivalent to adding the square + # of the weights to the loss with plain (non-momentum) SGD. + # Add weight decay at the end (fixed version) + if group["weight_decay"] > 0.0: + p.add_(p, alpha=(-group["lr"] * group["weight_decay"])) + + return loss + + +class AdamW8bit(Optimizer2State): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=1e-2, + amsgrad=False, + optim_bits=32, + args=None, + min_8bit_size=4096, + percentile_clipping=100, + block_wise=True, + is_paged=False, + ): + super().__init__( + "adam", + params, + lr, + betas, + eps, + weight_decay, + 8, + args, + min_8bit_size, + percentile_clipping, + block_wise, + is_paged=is_paged, + ) + + @torch.no_grad() + def step(self, closure=None): + """Performs a single optimization step. + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + if not self.initialized: + self.check_overrides() + self.to_gpu() # needed for fairseq pure fp16 training + self.initialized = True + + # if self.is_paged: self.page_mng.prefetch_all() + for gindex, group in enumerate(self.param_groups): + for pindex, p in enumerate(group["params"]): + if p.grad is None: + continue + state = self.state[p] + + if "step" not in state: + state["step"] = 0 + + # GaLore Projection + if "rank" in group: + if "projector" not in state: + state["projector"] = GaLoreProjector( + group["rank"], + update_proj_gap=group["update_proj_gap"], + scale=group["scale"], + proj_type=group["proj_type"], + ) + + if "weight_decay" in group and group["weight_decay"] > 0: + # ensure that the weight decay is not applied to the norm grad + group["weight_decay_saved"] = group["weight_decay"] + group["weight_decay"] = 0 + + grad = state["projector"].project(p.grad, state["step"]) + + # suboptimal implementation + p.saved_data = p.data.clone() + p.data = grad.clone().to(p.data.dtype).to(p.data.device) + p.data.zero_() + p.grad = grad + + if "state1" not in state: + self.init_state(group, p, gindex, pindex) + + self.prefetch_state(p) + self.update_step(group, p, gindex, pindex) + torch.cuda.synchronize() + + # GaLore Projection Back + if "rank" in group: + p.data = p.saved_data.add_(state["projector"].project_back(p.data)) + + # apply weight decay + if "weight_decay_saved" in group: + p.data.add_( + p.data, alpha=-group["lr"] * group["weight_decay_saved"] + ) + group["weight_decay"] = group["weight_decay_saved"] + del group["weight_decay_saved"] + + if self.is_paged: + # all paged operation are asynchronous, we need + # to sync to make sure all tensors are in the right state + torch.cuda.synchronize() + + return loss diff --git a/lib/python3.12/site-packages/torchao/prototype/galore/utils.py b/lib/python3.12/site-packages/torchao/prototype/galore/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6e9db05d30ea90c107c2b83b26fbdfb8093b8973 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/galore/utils.py @@ -0,0 +1,114 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch + + +def get_orthogonal_matrix(weights, rank, type): + module_params = weights + + if module_params.data.dtype != torch.float: + float_data = False + original_type = module_params.data.dtype + original_device = module_params.data.device + matrix = module_params.data.float() + else: + float_data = True + matrix = module_params.data + + U, s, Vh = torch.linalg.svd(matrix, full_matrices=False) + + # make the smaller matrix always to be orthogonal matrix + if type == "right": + # A = U[:, :rank] @ torch.diag(s[:rank]) + B = Vh[:rank, :] + + if not float_data: + B = B.to(original_device).type(original_type) + return B + elif type == "left": + A = U[:, :rank] + # B = torch.diag(s[:rank]) @ Vh[:rank, :] + if not float_data: + A = A.to(original_device).type(original_type) + return A + elif type == "full": + A = U[:, :rank] + B = Vh[:rank, :] + if not float_data: + A = A.to(original_device).type(original_type) + B = B.to(original_device).type(original_type) + return [A, B] + else: + raise ValueError("type should be left, right or full") + + +class TestGaLoreProjector: + def __init__( + self, + rank=128, + scale=1.0, + proj_type="std", + ): + self.rank = rank + self.scale = scale + + if proj_type != "std": + raise ("Only std projection is supported") + + self.proj_type = proj_type + + self.ortho_matrix = None + + def update_orthogonal_matrix(self, full_rank_grad): + if full_rank_grad.shape[0] >= full_rank_grad.shape[1]: + self.ortho_matrix = get_orthogonal_matrix( + full_rank_grad, self.rank, type="right" + ) + else: + self.ortho_matrix = get_orthogonal_matrix( + full_rank_grad, self.rank, type="left" + ) + + def project(self, full_rank_grad): + if full_rank_grad.shape[0] >= full_rank_grad.shape[1]: + low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t()) + else: + low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad) + + return low_rank_grad + + def project_back(self, low_rank_grad): + if low_rank_grad.shape[0] >= low_rank_grad.shape[1]: + full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix) + else: + full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad) + + return full_rank_grad * self.scale + + +def make_copy(*args): + return [t.detach().clone() for t in args] + + +# def adam_step( +# exp_avg, +# exp_avg2, +# grad, +# galore_proj, +# params, +# step_size=1e-4, +# beta1=BETA1, +# beta2=BETA2, +# eps=EPS, +# ): +# grad = galore_proj.project(grad) +# exp_avg = beta1 * exp_avg + (1 - beta1) * grad +# exp_avg2 = beta2 * exp_avg2 + (1 - beta2) * torch.square(grad) +# denom = exp_avg2.sqrt() + eps +# norm_grad = exp_avg / denom +# norm_grad = galore_proj.project_back(norm_grad) +# # params = params - step_size * norm_grad +# return exp_avg, exp_avg2, denom, norm_grad diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/__init__.py b/lib/python3.12/site-packages/torchao/prototype/hqq/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c8be63e6dcd483c9195ceca24f6fa9bcaed590f --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/hqq/__init__.py @@ -0,0 +1,6 @@ +from .mixed_mm import pack_2xint4, triton_mixed_mm + +__all__ = [ + "pack_2xint4", + "triton_mixed_mm", +] diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08b52b7620822e9683aece2e034ed6d14388e462 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/example.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f9b25905a3b740425f4fe4fbc9ffd9714b08a6c Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/hqq_tinygemm_linear.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/hqq_tinygemm_linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32f99d4813f6d180e74d4ba81abdea1ee37dc331 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/hqq_tinygemm_linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/kernels.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/kernels.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8a264a2dbaf7479c3fc7de9e6f0f9d7108be9ae Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/kernels.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/mixed_mm.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/mixed_mm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4472189cbac1f2349eded085d0770d35909a7f0 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/hqq/__pycache__/mixed_mm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/example.py b/lib/python3.12/site-packages/torchao/prototype/hqq/example.py new file mode 100644 index 0000000000000000000000000000000000000000..46fae4bfe98f984089928800a9610b62be2295fe --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/hqq/example.py @@ -0,0 +1,174 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch + +from torchao.dtypes import PlainLayout, TensorCoreTiledLayout +from torchao.dtypes.affine_quantized_tensor import ( + to_affine_quantized_intx, +) +from torchao.quantization import ( + MappingType, + ZeroPointDomain, +) + +# Parameters +device, compute_dtype = "cuda:0", torch.bfloat16 +group_size, axis = 64, 1 +in_features, out_features = 4096, 11800 + +torch.random.manual_seed(100) +linear_layer = torch.nn.Linear(in_features, out_features, bias=False, device=device) +x = torch.randn((1, linear_layer.in_features), dtype=torch.float, device=device) / 20.0 +y_ref = linear_layer(x) +W = linear_layer.weight.data.clone().to(device=device, dtype=compute_dtype) +del linear_layer.weight + +################################################################################################ +# AffineQuantizedTensor example +################################################################################################ +print("-------------------------------------------------------------------") +print("AffineQuantizedTensor example") +print("-------------------------------------------------------------------") +mapping_type = MappingType.ASYMMETRIC +block_size = (1, group_size) +target_dtype = torch.uint8 # until sub-byte dtypes are supported +preserve_zero = False +zero_point_domain = ZeroPointDomain.FLOAT +zero_point_dtype = compute_dtype +_layout = PlainLayout() + +for nbits in list(range(2, 9))[::-1]: + print( + "------------------------------------------------------------------------------" + ) + q_tensor_default = to_affine_quantized_intx( + input_float=W, + mapping_type=mapping_type, + block_size=block_size, + target_dtype=target_dtype, + quant_min=0, + quant_max=2**nbits - 1, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + _layout=_layout, + ) + + linear_layer.weight = q_tensor_default + print( + "nbits", + nbits, + "| Default dequantization error", + (W - q_tensor_default.dequantize()).abs().mean().item(), + ) + print( + "nbits", + nbits, + "| Default Dot product error", + (y_ref - linear_layer(x.to(compute_dtype))).abs().mean().item(), + ) + # nbits 4 | Default dequantization error 0.001953125 + # nbits 4 | Default Dot product error 0.005926903802901506 + + q_tensor_hqq = to_affine_quantized_intx( + input_float=W, + mapping_type=mapping_type, + block_size=block_size, + target_dtype=target_dtype, + quant_min=0, + quant_max=2**nbits - 1, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + _layout=_layout, + use_hqq=True, + ) + + linear_layer.weight = q_tensor_hqq + print( + "nbits", + nbits, + "| HQQ dequantization error", + (W - q_tensor_hqq.dequantize()).abs().mean().item(), + ) + print( + "nbits", + nbits, + "| HQQ Dot product error", + (y_ref - linear_layer(x.to(compute_dtype))).abs().mean().item(), + ) + # nbits 4 | HQQ dequantization error 0.0004863739013671875 + # nbits 4 | HQQ Dot product error 0.0014713306445628405 + +################################################################################################ +# quant_api example +################################################################################################ +print("-------------------------------------------------------------------") +print("Quant API example") +print("-------------------------------------------------------------------") + +from torchao.quantization.quant_api import int4_weight_only + +nbits = 4 +target_dtype = torch.int32 +inner_k_tiles = 8 +_layout = TensorCoreTiledLayout(inner_k_tiles=inner_k_tiles) + +int4_weight_only_patch_fct = int4_weight_only( + group_size=group_size, inner_k_tiles=inner_k_tiles +) +linear_layer_default = torch.nn.Linear( + in_features, out_features, bias=False, device=device +) +linear_layer_default.weight.data = W.clone() +linear_layer_default = int4_weight_only_patch_fct(linear_layer_default) +print( + "nbits", + nbits, + "| Default dequantization error", + (W - linear_layer_default(torch.eye(W.shape[1], dtype=W.dtype, device=W.device)).T) + .abs() + .mean() + .item(), +) +print( + "nbits", + nbits, + "| Default Dot product error", + (y_ref - linear_layer_default(x.to(compute_dtype))).abs().mean().item(), +) +# nbits 4 | Default dequantization error 0.000492095947265625 +# nbits 4 | Default Dot product error 0.0015244047390297055 + + +q_tensor_hqq = to_affine_quantized_intx( + input_float=W, + mapping_type=mapping_type, + block_size=block_size, + target_dtype=target_dtype, + quant_min=0, + quant_max=2**nbits - 1, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + _layout=_layout, + use_hqq=True, +) +linear_layer.weight = q_tensor_hqq +print( + "nbits", + nbits, + "| HQQ dequantization error", + (W - linear_layer(torch.eye(W.shape[1], dtype=W.dtype, device=W.device)).T) + .abs() + .mean() + .item(), +) +print( + "nbits", + nbits, + "| HQQ Dot product error", + (y_ref - linear_layer(x.to(compute_dtype))).abs().mean().item(), +) +# nbits 4 | HQQ dequantization error 0.0004863739013671875 +# nbits 4 | HQQ Dot product error 0.0014699687017127872 diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/hqq_tinygemm_linear.py b/lib/python3.12/site-packages/torchao/prototype/hqq/hqq_tinygemm_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..f15c9a8104aec62893a53f5cb8f44c217d5eef9a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/hqq/hqq_tinygemm_linear.py @@ -0,0 +1,275 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +# mobicham's tinygemm hqq eval script +import torch + +device = "cuda" + + +import copy + +import torch.nn.functional as F +from hqq.core.quantize import Quantizer +from hqq.core.utils import * # noqa: F401, F403 +from torch import Tensor, nn + +from torchao.dtypes.utils import is_device +from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, check_cpu_version + + +class HQQLinearTorchWeightOnlyInt4(torch.nn.Module): + def __init__( + self, + linear_layer: nn.Module | None, + quant_config: dict, + del_orig: bool = True, + compute_dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + initialize: bool = True, + inner_k_tiles=8, + padding=True, + ): + super().__init__() + + self.ready = False + self.in_gpu = False + self.bias = None + self.device = device + self.compute_dtype = compute_dtype + self.quant_config = copy.deepcopy(quant_config) + self.del_orig = del_orig + + weight_quant_params = self.quant_config["weight_quant_params"] + self.groupsize = weight_quant_params["group_size"] + self.nbits = weight_quant_params["nbits"] + self.inner_k_tiles = inner_k_tiles + self.padding = padding + + assert self.nbits in [1, 2, 4], "Unsupported nbits" + assert self.groupsize in [None, 32, 64, 128, 256], "Unsupported groupsize" + assert self.inner_k_tiles in [2, 4, 8], "Unsupported tile" + + self.linear_layer = linear_layer + self.compute_dtype = compute_dtype + + if initialize: + self.initialize() + + ###################### Initializers ###################### + def initialize_with_hqq_quants(self, W_q, meta, bias=None): + self.padding = ( + False # Force padding off, a bit tricky to post-pad with grouping + ) + + self.set_shape(meta["shape"]) + self.process_hqq_quants(W_q, meta) + self.bias = bias + self.ready = True + self.in_gpu = True + torch.cuda.empty_cache() + + return self + + def initialize(self): + if self.linear_layer is not None: + W = self.linear_layer.weight.data + self.set_shape(W.shape) + + if self.in_features_diff > 0: + W = F.pad(W, pad=(0, self.in_features_diff), value=0) + + W_q, meta = self.quantize(W, **self.quant_config) + self.process_hqq_quants(W_q, meta) + del W_q, meta + + self.bias = ( + None + if (self.linear_layer.bias is None) + else self.linear_layer.bias.to( + dtype=self.compute_dtype, device=self.device + ) + ) + + if self.del_orig: + del self.linear_layer + + self.ready = True + self.in_gpu = True + torch.cuda.empty_cache() + + return self + + ###################### Quantize/packing ###################### + + def quantize( + self, + W: Tensor, + weight_quant_params: dict, + scale_quant_params=dict | None, + zero_quant_params=dict | None, + offload_meta=False, + ): + W_q, meta = Quantizer.quantize( + W, + **weight_quant_params, + device=self.device, + compute_dtype=self.compute_dtype, + bitpack=False, + ) + + # ToDO: meta quantization + + return W_q, meta + + # TODO: move these to utils + @torch.no_grad() + def reshape_meta_axis1(self, meta_tensor, new_group_size, shape): + meta_tensor = meta_tensor.repeat([1, shape[1]]).reshape(shape) + meta_tensor = torch.mean( + meta_tensor.reshape([-1, new_group_size]), axis=1, keepdim=True + ) + return meta_tensor + + def find_multiple(self, n: int, k: int) -> int: + if n % k == 0: + return n + return n + k - (n % k) + + def set_shape(self, shape): + self.shape = shape + self.in_features = shape[1] + self.out_features = shape[0] + + self.origin_in_features = self.in_features + if self.padding: + self.in_features = self.find_multiple(self.in_features, 1024) + + self.in_features_diff = self.in_features - self.origin_in_features + + @torch.no_grad() + def process_hqq_quants(self, W_q, meta): + scales = meta["scale"] + zeros = meta["zero"] + shape = meta["shape"] + + if meta["packing"] is not None: + W_q = Quantizer.unpack[meta["packing"]](W_q) + + if self.groupsize is None: + self.groupsize = 128 + W_q = W_q.reshape([-1, self.groupsize]) + scales = self.reshape_meta_axis1(scales, self.groupsize, shape) + zeros = self.reshape_meta_axis1(zeros, self.groupsize, shape) + + W_q_torch, scales_torch, zeros_torch = self.hqq_quants_to_torch_quants( + W_q=W_q, scales=scales, zeros=zeros, shape=shape, nbits=self.nbits + ) + if check_cpu_version(W_q.device): + self.weight_int4pack = torch.ops.aten._convert_weight_to_int4pack_for_cpu( + W_q_torch, self.inner_k_tiles + ) + else: + self.weight_int4pack = torch.ops.aten._convert_weight_to_int4pack( + W_q_torch, self.inner_k_tiles + ) + self.scales_and_zeros = self.pack_scales_and_zeros(scales_torch, zeros_torch) + + del W_q_torch, scales_torch, zeros_torch + torch.cuda.empty_cache() + + @torch.no_grad() + def hqq_quants_to_torch_quants( + self, W_q: Tensor, scales: Tensor, zeros: Tensor, shape, nbits=4 + ): + W_q = W_q.to(dtype=self.compute_dtype, device=self.device) + scales = scales.to(dtype=self.compute_dtype, device=self.device) + zeros = zeros.to(dtype=self.compute_dtype, device=self.device) + + max_int = 2**nbits - 1 + min_int = 0 + dump = 2 ** (nbits - 1) + + # HQQ -> torch logic + new_zeros = (scales * dump) - zeros * scales + + min_val = new_zeros - scales * dump + + # group_quantize_tensor_from_qparams + W_r = (W_q - zeros) * scales + + W_q = ( + W_r.sub(min_val) + .div(scales) + .round() + .clamp_(min_int, max_int) + .to(torch.int32) + .reshape(shape) + .contiguous() + ) + if TORCH_VERSION_AT_LEAST_2_5: + if not is_device(W_q.device.type, "cpu"): + W_q = (W_q[::, ::2] << 4 | W_q[::, 1::2]).to(torch.uint8) + + # group_dequantize_tensor_from_qparams + # W_r = W_q*scales + min_val + + scales = scales.contiguous().reshape(shape[0], -1) + new_zeros = new_zeros.contiguous().reshape(shape[0], -1) + + return W_q, scales, new_zeros + + def pack_scales_and_zeros(self, scales, zeros): + assert scales.shape == zeros.shape + assert scales.dtype == torch.bfloat16 + assert zeros.dtype == torch.bfloat16 + return ( + torch.cat( + [ + scales.reshape(scales.size(0), scales.size(1), 1), + zeros.reshape(zeros.size(0), zeros.size(1), 1), + ], + 2, + ) + .transpose(0, 1) + .contiguous() + ) + + ###################### Forward/matmul ###################### + + @torch.jit.ignore() + def matmul(self, x): + origin_x_size = x.size() + x = x.reshape(-1, origin_x_size[-1]) + if check_cpu_version(x.device): + c = torch.ops.aten._weight_int4pack_mm_for_cpu( + x, self.weight_int4pack, self.groupsize, self.scales_and_zeros + ) + else: + c = torch.ops.aten._weight_int4pack_mm( + x, self.weight_int4pack, self.groupsize, self.scales_and_zeros + ) + new_shape = origin_x_size[:-1] + (self.out_features,) + c = c.reshape(new_shape) + return c + + # TODO without matmul + def dequantize(self): + return self.matmul( + torch.eye(self.in_features, dtype=self.compute_dtype, device=self.device) + )[: self.origin_in_features].t() + + # TODO: backward + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.to(self.compute_dtype) + + if self.in_features_diff > 0: + x = F.pad(x, pad=(0, self.in_features_diff)) + + out = self.matmul(x) + + if self.bias is not None: + out += self.bias + return out diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/kernels.py b/lib/python3.12/site-packages/torchao/prototype/hqq/kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..d32fab3f2c66a5659b1351ab8abd4ac028eba89b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/hqq/kernels.py @@ -0,0 +1,402 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import triton +import triton.language as tl +from triton import Config + +# TODO: add early config prune and estimate_matmul_time to reduce autotuning time +# from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time + + +def get_configs_io_bound(): + configs = [] + for num_stages in [2, 3, 4, 5, 6]: + for block_m in [16, 32]: + for block_k in [32, 64]: + for block_n in [32, 64, 128, 256]: + num_warps = 2 if block_n <= 64 else 4 + configs.append( + Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": 1, + }, + num_stages=num_stages, + num_warps=num_warps, + ) + ) + # split_k + for split_k in [2, 4, 8, 16]: + configs.append( + Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": split_k, + }, + num_stages=num_stages, + num_warps=num_warps, + pre_hook=init_to_zero("C"), + ) + ) + return configs + + +def get_configs_compute_bound(): + configs = [ + # basic configs for compute-bound matmuls + Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + Config( + {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + Config( + {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, + num_stages=5, + num_warps=2, + ), + # good for int8 + Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + Config( + {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + Config( + {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + Config( + {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, + num_stages=5, + num_warps=2, + ), + ] + return configs + + +def init_to_zero(name): + return lambda nargs: nargs[name].zero_() + + +MIXED_MM_HEURISTICS = { + "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, + "BLOCK_K": lambda args: min(args["BLOCK_K"], args["QGROUP_SIZE"]) + if not args["TRANSPOSED"] + else args["BLOCK_K"], + "BLOCK_N": lambda args: min(args["BLOCK_N"], args["QGROUP_SIZE"]) + if args["TRANSPOSED"] + else args["BLOCK_N"], + "SPLIT_K": lambda args: 1 + if args["IS_BFLOAT16"] + else args["SPLIT_K"], # atomic add not supported for bfloat16 +} + + +@triton.jit +def _mixed_mm_kernel( + # Operands + A, + B, + scales_ptr, + zeros_ptr, + C, + # Matrix dims. + M, + N, + K, + # a, b, c, scales / zeros strides + stride_am, + stride_ak, # + stride_bk, + stride_bn, # + stride_cm, + stride_cn, + stride_scale_k, + stride_scale_n, + # Meta-params + IS_BFLOAT16: tl.constexpr, + QGROUP_SIZE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + SPLIT_K: tl.constexpr, + EVEN_K: tl.constexpr, + TRANSPOSED: tl.constexpr = False, + GROUP_M: tl.constexpr = 8, + # tl.dot options + acc_dtype: tl.constexpr = tl.float32, + input_precision: tl.constexpr = "ieee", + fp8_fast_accum: tl.constexpr = False, + # Only used for debugging + DEBUG: tl.constexpr = False, +): + """Mixed matmul kernel + + A has shape (M, K) and is float16, bfloat16, or float32 + + B is i4 / s4 and has shape (K // 2, N) and is packed as uint8 / int8. See `packed_2xint4` for details. + + Scales and zeros are of shape (NUM_GROUPS, N) and are same dtype as A, where NUM_GROUPS = (K // QGROUP_SIZE) + QGROUP_SIZE should be a multiple of BLOCK_K such that a vector of scales / zeros is loaded and broadcasted to block shape + per mainloop iteration. + + In the transposed case, A is M x N and B is K x N, and we reduce along "N": + - TLDR: we are loading rows of A and B blocks at a time, dequantizing and transposing each block of B to achieve the overall + effect of a transposed matmul. This is necessary to perform a transposed matmul without unpacking and repacking the B matrix. + - Indexing remains the same for A (the reduction dim (BLK_K / K) corresponds to axis 1 of A -- "N" above) + - We load a BLK_M x BLK_K block of A + - Indexing for B is now flipped: N <-> K + - We load BLK_N x BLK_K block of B (remembering that the reduction dimension is axis 1 of B) + - We dequantize and transpose to BLK_K x BLK_N + - scale / zero indexing also change, since we are now iterating along the non-grouping dim within the mac loop and along + the grouping dim across blocks. + - Each mac loop calculates BLK_M x BLK_N -> M x "N"(= K) + - Within the mac loop for each block, we iterate along axis=1 for **both** A and B since axis = 1 is now the reduction dim for B. + + NOTE: Assumes that the quantization grouping was done along the K dimension originally (i.e., QGROUP_SIZE consecutive elements + of original weight matrix in the K dimension were grouped together when calculating min / max scaling factors). + """ + + if not TRANSPOSED: + tl.static_assert(QGROUP_SIZE % BLOCK_K == 0) + else: + tl.static_assert(QGROUP_SIZE % BLOCK_N == 0) + + # TODO: refactor swizzling to separate function + # Threadblock swizzling + pid = tl.program_id(0) + pid_z = tl.program_id(1) + + grid_m = tl.cdiv(M, BLOCK_M) + grid_n = tl.cdiv(N, BLOCK_N) + + width = GROUP_M * grid_n + group_id = pid // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // group_size + + rm = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M + if not DEBUG: + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + else: + ram = rm + # rn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N + # rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rak = pid_z * BLOCK_K + tl.arange(0, BLOCK_K) + + # BLOCK_K for b is effectively BLOCK_K // 2 + if not TRANSPOSED: + rn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N + if not DEBUG: + rbn = tl.max_contiguous( + tl.multiple_of(rn % N, BLOCK_N), BLOCK_N + ) # rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + else: + rbn = rn + rbk = pid_z * BLOCK_K // 2 + tl.arange(0, BLOCK_K // 2) + else: + rn = (pid_n * BLOCK_N // 2 + tl.arange(0, BLOCK_N // 2)) % N + if not DEBUG: + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N // 2), BLOCK_N // 2) + else: + rbn = rn + rbk = rak + + A = A + (ram[:, None] * stride_am + rak[None, :] * stride_ak) + + if not TRANSPOSED: + B = B + (rbk[:, None] * stride_bk + rbn[None, :] * stride_bn) + else: + # Note: in the transposed case, we are loading BLK_N x BLK_K, but we need to transpose to BLK_K x BLK_N + # the strides are adjusted accordingly, since we to stride by stride_bk to get rows of BLK_N + # and stride_bn to get columns of BLK_K + B = B + (rbn[:, None] * stride_bk + rbk[None, :] * stride_bn) + + # Grouping is along K, so in the forward pass, each block loads a row vector of BLK_K x BLK_N + # where grouping varies along N, hence the mainloop marches down the K dimension, where + # group idx is given by K // QGROUP_SIZE + + if not TRANSPOSED: + offsets_scale_n = ( + pid_n * stride_scale_n * BLOCK_N + tl.arange(0, BLOCK_N) * stride_scale_n + ) + else: + scale_offset_k = pid_n * BLOCK_N * stride_scale_k // QGROUP_SIZE + offsets_scale_n = tl.arange(0, BLOCK_K) * stride_scale_n + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype) + for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)): + if EVEN_K: + a = tl.load(A) + qb = tl.load(B) + else: + k_remaining_a = K - k * (BLOCK_K * SPLIT_K) + if not TRANSPOSED: + k_remaining_b = ( + K - k * (BLOCK_K * SPLIT_K) // 2 + ) # Note the division by 2 + else: + k_remaining_b = K - k * (BLOCK_K * SPLIT_K) # = k_remaining_a + + _0 = tl.zeros((1, 1), dtype=C.dtype.element_ty) + a = tl.load(A, mask=rak[None, :] < k_remaining_a, other=_0) + qb = tl.load(B, mask=rbk[:, None] < k_remaining_b, other=_0) + + if not TRANSPOSED: + scale_offset_k = k * BLOCK_K * SPLIT_K * stride_scale_k // QGROUP_SIZE + else: + offsets_scale_n = ( + k * stride_scale_n * BLOCK_K + tl.arange(0, BLOCK_K) * stride_scale_n + ) + + scales = tl.load(scales_ptr + offsets_scale_n + scale_offset_k) + zeros = tl.load(zeros_ptr + offsets_scale_n + scale_offset_k) + + # Unpack qweights -- h/t jlebar! + _4_i8 = tl.full((1,), 4, dtype=tl.int8) + qb_lo = (qb << _4_i8) >> _4_i8 + qb_hi = qb >> _4_i8 + + # Upcast to fp16 + # TODO: better bfloat16 conversion? compilation error if direct conversion from int8 to bfloat16 + if IS_BFLOAT16: + dq_b = ( + tl.join( + qb_lo.to(tl.float16).to(A.dtype.element_ty), + qb_hi.to(tl.float16).to(A.dtype.element_ty), + ).permute(0, 2, 1) + # .reshape(BLOCK_K, BLOCK_N) + ) + else: + dq_b = ( + tl.join( + qb_lo.to(A.dtype.element_ty), + qb_hi.to(A.dtype.element_ty), + ).permute(0, 2, 1) + # .reshape(BLOCK_K, BLOCK_N) + ) + if not TRANSPOSED: + dq_b = dq_b.reshape(BLOCK_K, BLOCK_N) + else: + dq_b = dq_b.reshape(BLOCK_N, BLOCK_K) + + # Scale upcasted weights + # Note that we broadcast the scales --> the assumption is that all scales fall within a single QGROUP + # This condition is statically check (see assertions above) + + zeros = zeros[None, :] + scales = scales[None, :] + + dq_b = (dq_b - zeros) * scales + + if TRANSPOSED: + dq_b = tl.trans(dq_b) + + if fp8_fast_accum: + acc = tl.dot( + a, dq_b, acc, out_dtype=acc_dtype, input_precision=input_precision + ) + else: + acc += tl.dot(a, dq_b, out_dtype=acc_dtype, input_precision=input_precision) + A += BLOCK_K * SPLIT_K * stride_ak + + # Advance by half the block size, since each block is unpacked and upcasted into two fp16 values + if not TRANSPOSED: + B += BLOCK_K * SPLIT_K * stride_bk // 2 + else: + # we iterating across a row of B (non-packing dim, hence no need for div 2) + B += BLOCK_K * SPLIT_K * stride_bn + acc = acc.to(C.dtype.element_ty) + + offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + C = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + + if SPLIT_K == 1: + tl.store(C, acc, mask=mask) + else: + tl.atomic_add(C, acc, mask=mask) + + +_mixed_mm = triton.heuristics(MIXED_MM_HEURISTICS)(_mixed_mm_kernel) +mixed_mm_kernel_max_autotune = triton.autotune( + configs=get_configs_compute_bound() + get_configs_io_bound(), key=["M", "N", "K"] +)(_mixed_mm) +mixed_mm_kernel_compute_bound = triton.autotune( + configs=get_configs_compute_bound(), key=["M", "N", "K"] +)(_mixed_mm) +_mixed_mm_debug = _mixed_mm diff --git a/lib/python3.12/site-packages/torchao/prototype/hqq/mixed_mm.py b/lib/python3.12/site-packages/torchao/prototype/hqq/mixed_mm.py new file mode 100644 index 0000000000000000000000000000000000000000..5ae2f28f0cc43ddc6b94613778b7768aaa1941d1 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/hqq/mixed_mm.py @@ -0,0 +1,179 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +import triton.language as tl +from triton import cdiv + +from .kernels import ( + mixed_mm_kernel_compute_bound, + mixed_mm_kernel_max_autotune, +) + + +# h/t jlebar for the bit packing / unpacking logic (source: Triton Slack thread) +# https://gist.github.com/jlebar/3435b2c00deea53258887ce37231e5e2 +def pack_2xint4(t): + """ + The packing format is such that consecutive rows are packed into a lower / upper bits + E.g., + Original, unpacked B (dtype i8): + [ + [0, 1, 2, 3] + [4, 5, 6, 7] + [8, 9, 10, 11] + [12, 13, 14, 15] + ] + Packed B: + [ + [0|4, 1|5, 2|6, 3|7] + [8|12, 9|13, 10|14, 11|15] + ] + (Note each entry in `Packed B` is shown lsb->msb) + """ + assert t.dtype == torch.int8 or t.dtype == torch.uint8 + t = t.reshape(t.shape[0] // 2, 2, t.shape[1]).permute(1, 0, 2) + return (t[0] & 0xF) | (t[1] << 4) + + +def triton_mixed_mm( + a, + b, + scales, + zeros, + group_size, + transposed=False, + acc_dtype=None, + input_precision="ieee", + fp8_fast_accum=False, + kernel_type="compute_bound", + # For debugging only + BLOCK_M=None, + BLOCK_N=None, + BLOCK_K=None, +): + """Run fused int4 / fp16 dequant GEMM + + Args: + a (torch.Tensor): M x K if not transposed, M x N if transposed + b (torch.Tensor): (K // 2) x N, packed such that 2 int4's are packed into 1 uint8 (see pack_2xint4) + scales (torch.Tensor): (num_groups x N), where num_groups = (N * K / group_size) + zeros (torch.Tensor): same shape as scales + group_size (torch.Tensor): size of group in groupwise quantization -- MUST be along axis 1 of an N x K matrix + transposed (bool, optional): Whether to run a transposed matmul where shapes are (M x N) x (K x N) => (M x K) + acc_dtype (_type_, optional): dtype of accumulator. Defaults to None, which corresponds to tl.float32. + input_precision (str, optional): Only relevant when dtype of a is torch.float32. Defaults to "ieee". + kernel_type (str, optional): Type of autoconfig to use. Either "max_autotune" or "compute_bound". + BLOCK_M (int, optional): Only for debugging. Defaults to None. + BLOCK_N (int, optional): Only for debugging. Defaults to None. + BLOCK_K (int, optional): Only for debugging. Defaults to None. + + Returns: + c (torch.Tensor): M x N + """ + device = a.device + # handle non-contiguous inputs if necessary + if a.stride(0) > 1 and a.stride(1) > 1: + a = a.contiguous() + if b.stride(0) > 1 and b.stride(1) > 1: + b = b.contiguous() + # checks constraints + if not transposed: + assert a.shape[1] == b.shape[0] * 2, "incompatible dimensions" + + assert b.dtype == torch.int8 or b.dtype == torch.uint8, "b must be int8 or uint8" + assert scales.ndim == 2 + if transposed: + assert a.shape[1] == b.shape[1], ( + "transpose requires (M x N) x (K x N), where reduction dim is N" + ) + + M, K = a.shape + N = b.shape[1] if not transposed else b.shape[0] * 2 + # assert scales.shape[1] == N if not transposed else scales.shape[0] == N + # assert scales.shape[0] == K // group_size if not transposed else scales.shape[1] == K // group_size + assert scales.dtype == a.dtype + assert scales.shape == zeros.shape + assert zeros.dtype == a.dtype + + # Assumes c is same type as a + c = torch.empty((M, N), device=device, dtype=a.dtype) + if acc_dtype is None: + acc_dtype = tl.float32 + + grid = lambda META: ( + cdiv(M, META["BLOCK_M"]) * cdiv(N, META["BLOCK_N"]), + META["SPLIT_K"], + ) + + if kernel_type == "max_autotune": + kernel = mixed_mm_kernel_max_autotune + elif kernel_type == "compute_bound": + kernel = mixed_mm_kernel_compute_bound + else: + from .kernels import _mixed_mm_debug + + kernel = _mixed_mm_debug + + if kernel_type == "max_autotune" or kernel_type == "compute_bound": + kernel[grid]( + a, + b, + scales, + zeros, + c, + M, + N, + K, # + a.stride(0), + a.stride(1), # + b.stride(0), + b.stride(1), # + c.stride(0), + c.stride(1), + scales.stride(0), + scales.stride(1), + TRANSPOSED=transposed, + IS_BFLOAT16=a.dtype == torch.bfloat16, + QGROUP_SIZE=group_size, + acc_dtype=acc_dtype, + input_precision=input_precision, + fp8_fast_accum=fp8_fast_accum, + ) + else: + assert all([BLOCK_M is not None, BLOCK_N is not None, BLOCK_K is not None]) + grid = (M // BLOCK_M * N // BLOCK_N, 1, 1) + kernel[grid]( + a, + b, + scales, + zeros, + c, + M, + N, + K, # + a.stride(0), + a.stride(1), # + b.stride(0), + b.stride(1), # + c.stride(0), + c.stride(1), + scales.stride(0), + scales.stride(1), + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + SPLIT_K=1, + EVEN_K=True, + TRANSPOSED=transposed, + IS_BFLOAT16=a.dtype == torch.bfloat16, + QGROUP_SIZE=group_size, + acc_dtype=acc_dtype, + input_precision=input_precision, + fp8_fast_accum=fp8_fast_accum, + DEBUG=True, + ) + + return c diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/__init__.py b/lib/python3.12/site-packages/torchao/prototype/parq/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d254b1395a83cbd1aff0c20c35fc05529a9823e6 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from .optim import ( # noqa: F401 + ProxBinaryRelax, + ProxHardQuant, + ProxMap, + ProxPARQ, + QuantOptimizer, +) +from .quant import ( # noqa: F401 + Int4UnifTorchaoQuantizer, + LSBQuantizer, + MaxUnifQuantizer, + Quantizer, + TernaryUnifQuantizer, + UnifQuantizer, + UnifTorchaoQuantizer, +) diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bbe02353526e1a93f19741fe8a8d65cc4f30d53 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c38930b77592987b445ac37b8b70d37322741fc6 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/__init__.py b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..237a058a12fedd2f23ee784789a0bae129074ba2 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from .binarelax import ProxBinaryRelax # noqa: F401 +from .parq import ProxPARQ # noqa: F401 +from .proxmap import ProxHardQuant, ProxMap # noqa: F401 +from .quantopt import QuantOptimizer # noqa: F401 diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c9e1d6cd02fefb314021f4b5ac7b3b008bd8e4b5 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/binarelax.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/binarelax.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb1a30693a8000227828900a77bae1f8441b7610 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/binarelax.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/parq.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/parq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e4ce8e0e0fbd5189f828fcd8d656cf4737cfe29 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/parq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/proxmap.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/proxmap.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b88c4b768d4314717b4cd21f960e36243bc98cbb Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/proxmap.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/quantopt.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/quantopt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1adf8b652fc3f3f1b06ab4f330ea998ea6518a04 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/optim/__pycache__/quantopt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/binarelax.py b/lib/python3.12/site-packages/torchao/prototype/parq/optim/binarelax.py new file mode 100644 index 0000000000000000000000000000000000000000..2cc6611f6b9871be440675676a83c81a9f162b2c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/optim/binarelax.py @@ -0,0 +1,50 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Optional + +import torch +from torch import Tensor + +from ..utils import channel_bucketize +from .proxmap import ProxMap + + +class ProxBinaryRelax(ProxMap): + """Prox-map of Binary Relax, Q may not be evenly spaced.""" + + def __init__(self, anneal_start: int, anneal_end: int) -> None: + self.anneal_start = anneal_start + self.anneal_end = anneal_end + + @torch.no_grad() + def apply_( + self, + p: Tensor, + q: Tensor, + Q: Tensor, + step_count: int, + dim: Optional[int] = None, + ) -> None: + if step_count < self.anneal_start: + return + + if q is None: + # hard quantization to the nearest point in Q + Q_mid = (Q[..., :-1] + Q[..., 1:]) / 2 + if dim is None: + q = Q[torch.bucketize(p, Q_mid)] + else: + q = Q.gather(1, channel_bucketize(p, Q_mid)) + + if step_count >= self.anneal_end: + p.copy_(q) + else: + # linear annealing of relaxation coefficient + theta = (step_count - self.anneal_start) / ( + self.anneal_end - self.anneal_start + ) + p.mul_(1 - theta).add_(q, alpha=theta) diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/parq.py b/lib/python3.12/site-packages/torchao/prototype/parq/optim/parq.py new file mode 100644 index 0000000000000000000000000000000000000000..ade403a87d821804bab3bd4b3253f726eb405801 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/optim/parq.py @@ -0,0 +1,113 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import math +from functools import partial +from typing import Optional + +import torch +from torch import Tensor + +from ..utils import channel_bucketize +from .proxmap import ProxMap + + +def amp_custom_fwd(cast_inputs: Optional[torch.types._dtype] = None): + try: + return partial( + torch.amp.custom_fwd, device_type="cuda", cast_inputs=cast_inputs + ) + except AttributeError: + return partial(torch.cuda.amp.custom_fwd, cast_inputs=cast_inputs) + + +def normalized_mirror_sigmoid( + t: float, t1: float, t2: float, s: float, c: float +) -> float: + """Sigmoid-like function decreasing from 1 to 0 over interval [t1, t2). + s is steepness of the sigmoid-like function, almost linear for s < 1. + 'mirror' means decreasing instead of increasing as true sigmoid, + 'normalized' means value 1 at starting point t1 and 0 at end point t2.""" + assert t >= t1 and t < t2, "Normalized sigmoid: ensure t1 <= t < t2" + ft = (t - t1) / (t2 - t1) # fraction of progress from t1 to t2 + st = 1 / (1 + math.exp(s * (ft - c))) # scaled and shifted mirror sigmoid + s1 = 1 / (1 + math.exp(-0.5 * s)) # st value when t = t1 -> ft = 0 + s2 = 1 / (1 + math.exp(0.5 * s)) # st value when t = t2 -> ft = 1 + return (st - s2) / (s1 - s2) # shift and scale to range (0, 1] + + +class ProxPARQ(ProxMap): + def __init__( + self, + anneal_start: int, + anneal_end: int, + steepness: float = 10, + anneal_center: float = 0.5, + ) -> None: + assert anneal_start < anneal_end, "PARQ annealing: start before end." + assert steepness > 0, "PARQ annealing steepness should be positive." + self.anneal_start = anneal_start + self.anneal_end = anneal_end + self.steepness = steepness + self.anneal_center = anneal_center + + @torch.no_grad() + @amp_custom_fwd(cast_inputs=torch.float32) + def apply_( + self, + p: Tensor, + q: Tensor, + Q: Tensor, + step_count: int, + dim: Optional[int] = None, + ) -> float: + """Prox-map of PARQ with gradual annealing to hard quantization.""" + + if step_count < self.anneal_start: + inv_slope = 1.0 + elif step_count >= self.anneal_end: + inv_slope = 0.0 + if q is None: + # hard quantization to the nearest point in Q + Q_mid = (Q[..., :-1] + Q[..., 1:]) / 2 + if dim is None: + q = Q[torch.bucketize(p, Q_mid)] + else: + q = Q.gather(1, channel_bucketize(p, Q_mid)) + p.copy_(q) + else: + inv_slope = normalized_mirror_sigmoid( + step_count, + self.anneal_start, + self.anneal_end, + self.steepness, + self.anneal_center, + ) + inv_slope = max(torch.finfo(p.dtype).tiny, inv_slope) + # it is important to clamp idx-1 and then clamping idx itself + # idx_1[k] == idx[k] iff p[k] > Q.max() or p[k] <= Q.min() + if dim is None: + idx = torch.bucketize(p, Q) # locate quant interval + idx_lower = (idx - 1).clamp_(min=0) # index of lower bound + idx_upper = idx.clamp(max=Q.numel() - 1) # index of upper bound + q_lower = Q[idx_lower] # lower boundary of interval + q_upper = Q[idx_upper] # upper boundary of interval + center = (q_lower + q_upper) / 2 # center of interval + # concise implementation of piecewise-affine prox map + q = (center + (p - center) / inv_slope).clamp_(min=q_lower, max=q_upper) + else: + idx = channel_bucketize(p, Q) + idx_lower = (idx - 1).clamp_(min=0) + idx_upper = idx.clamp(max=Q.size(1) - 1) + q_lower = Q.gather(1, idx_lower) + q_upper = Q.gather(1, idx_upper) + center = (q_lower + q_upper) / 2 + q = torch.minimum( + torch.maximum(center + (p - center) / inv_slope, q_lower), q_upper + ) + # in-place update of model parameters + p.copy_(q) + return inv_slope diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/proxmap.py b/lib/python3.12/site-packages/torchao/prototype/parq/optim/proxmap.py new file mode 100644 index 0000000000000000000000000000000000000000..da867cc5dba6cf824ffcfc02c5f6972fad96384a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/optim/proxmap.py @@ -0,0 +1,49 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from abc import ABC, abstractmethod +from typing import Optional + +import torch +from torch import Tensor + +from ..utils import channel_bucketize + + +# Create an abstract class to provide proximal-mapping interface +class ProxMap(ABC): + @abstractmethod + def apply_(self, p: Tensor, q: Tensor, Q: Tensor, step_count: int) -> None: + """Provide interface for proximal mapping (modify p in-place): + prox_map.apply_(p, q, Q, step_count) + Inputs: + p (Tensor): tensor to be quantized + q (Tensor): None or hard quantized tensor of same size as p + Q (Tensor): set of target quantization values + step_count: trigger iteration-dependent mapping if needed + """ + + +class ProxHardQuant(ProxMap): + """Prox-map of hard quantization, Q may not be evenly spaced.""" + + @torch.no_grad() + def apply_( + self, + p: Tensor, + q: Tensor, + Q: Tensor, + step_count: int, + dim: Optional[int] = None, + ) -> None: + if q is None: + # quantize to the nearest point in Q + Q_mid = (Q[..., :-1] + Q[..., 1:]) / 2 + if dim is None: + q = Q[torch.bucketize(p, Q_mid)] + else: + q = Q.gather(1, channel_bucketize(p, Q_mid)) + p.copy_(q) diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/optim/quantopt.py b/lib/python3.12/site-packages/torchao/prototype/parq/optim/quantopt.py new file mode 100644 index 0000000000000000000000000000000000000000..2cdd34536db6169becc9fd8833433974664f0dd1 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/optim/quantopt.py @@ -0,0 +1,310 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from collections import defaultdict +from collections.abc import Callable +from functools import partial +from typing import Any, Optional + +import torch +from torch import Tensor +from torch.optim import Optimizer + +from ..quant import Quantizer +from ..utils import HAS_DTENSOR, is_dtensor +from .proxmap import ProxMap + +if HAS_DTENSOR: + from torch.distributed.tensor import distribute_tensor + from torch.distributed.tensor.experimental import local_map + from torch.distributed.tensor.placement_types import Shard + + +class QuantOptimizer(Optimizer): + """QuantOptimizer assembles functionalities of the following objects: + a base optimizer (e.g., SGD or AdamW) + - update the latent variables for QAT + a quantizer (e.g., UnifQuantizer, LSBQuantizer, LearnedScale) + - update target quantization values for model parameters + a proximal mapping (e.g, HardQuant/STE, PARQ, BinaryRelax) + - update model parameters based on the above two updates + Other parameters: + - warmup_steps: int >= 0 + - quant_period: int > 0 + - quant_per_channel: True or False + - quant_shrink: True or False + """ + + def __init__( + self, + base_optimizer: Optimizer, + quantizer: Quantizer, + prox_map: ProxMap, + warmup_steps: int = 0, + quant_period: int = 10, + quant_per_channel: bool = False, + quant_shrink: bool = False, + anneal_wd_frac: float = 0.0, + ) -> None: + if not 0 <= anneal_wd_frac <= 1: + raise ValueError(f"Invalid {anneal_wd_frac=} outside range [0.0, 1.0]") + + # need to reconstruct these objects if loading checkpoint + self.base_optimizer = base_optimizer + self.quantizer = quantizer + self.prox_map = prox_map + + # need to store these attributes in state_dict for checkpoint + self.warmup_steps = warmup_steps + self.quant_period = quant_period + self.quant_per_channel = quant_per_channel + self.quant_shrink = quant_shrink + self.anneal_wd_frac = anneal_wd_frac + self.num_steps = 0 + + # Initialize "cumu_lr" and latent params in optimizer states + for group in self.regularized_param_groups(): + group["cumu_lr"] = 0.0 + if self.anneal_wd_frac > 0: + group["initial_wd"] = group["weight_decay"] + # NOTE: Filling state dict here cause Adam(W) error, which assumes + # empty state[p] at first step() where optimizer states are initialized + + def __getattribute__(self, name: str): + try: + attr = super(Optimizer, self).__getattribute__(name) + except AttributeError: + attr = self.base_optimizer.__getattribute__(name) + return attr + + def __repr__(self) -> str: + base_optimizer = "\n ".join(self.base_optimizer.__repr__().split("\n")) + quantizer = self.quantizer.__class__.__name__ + prox_map = self.prox_map.__class__.__name__ + extra_repr = "\n ".join(("(", base_optimizer, f"{quantizer=}", f"{prox_map=}")) + return f"{self.__class__.__name__} {extra_repr}\n)" + + @property + def state(self) -> defaultdict[Tensor, Any]: # pyre-ignore[3] + return self._state if hasattr(self, "_state") else self.base_optimizer.state + + @staticmethod + def quantize_( + p: Tensor, + quants: Tensor, + quantizer: Quantizer, + b: int, + dim: Optional[int] = None, + ) -> Optional[Tensor]: + """Optionally update the quantization targets `quants` in place. + Return the quantized `p` as a by-product if `quant_update=True`. + """ + q, Q = quantizer.quantize(p, b, dim=dim) # pyre-ignore[28] + quants.copy_(Q) + return q + + def regularized_param_groups(self): # pyre-ignore[3] + """Yield parameter groups that need to be quantized.""" + for group in self.param_groups: + if group.get("quant_bits", 16) < 16: + yield group + + @property + def _param_set(self) -> set[int]: + return { + p.data_ptr() + for group in self.regularized_param_groups() + for p in group["params"] + } + + def get_filter_fn( + self, module: torch.nn.Module + ) -> Callable[[torch.nn.Module], bool]: + param_set = self._param_set + + def _filter_fn(module: torch.nn.Module, *args) -> bool: + for p in module.parameters(recurse=False): + if p.data_ptr() in param_set: + return True + return False + + return _filter_fn + + @torch._disable_dynamo + def state_dict(self) -> dict[str, Any]: + state_dict = self.base_optimizer.state_dict() + state_dict["qat_state"] = {"num_steps": self.num_steps} + # quantizer and prox_map may also need to save states, can add here + return state_dict + + @torch._disable_dynamo + def load_state_dict( + self, state_dict: dict[str, Any], start_step: Optional[int] = None + ) -> None: + qat_state = state_dict.get("qat_state") + # resume from check points usually not corresponds to saved num_steps + # so allow explicit start_step computed from epochs * steps_per_epoc + if start_step is not None: + self.num_steps = start_step + elif qat_state is not None: + # hope discrepancy in num_steps does not cause major problem! + self.num_steps = qat_state["num_steps"] + self.base_optimizer.load_state_dict(state_dict) + + @torch.no_grad() + def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]: + """Performs a single optimization step. + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + if self.num_steps < self.warmup_steps: + # warmup stage: running the base optimizer only + loss = self.base_optimizer.step(closure=closure) # pyre-ignore[6] + self.num_steps += 1 + return loss + + if self.num_steps == self.warmup_steps: + # first step of qat, save latent params, instead of restore + self.save_latent_params() + else: + # qat: restore latent params for update by the base optimizer + self.restore_latent_params() + + # call base optimizer step() method to update latent parameters + loss = self.base_optimizer.step(closure=closure) # pyre-ignore[6] + + if hasattr(self, "_state"): + assert self.warmup_steps == 0 + # restore the temporary state to the base optimizer's state + for p in self._state.keys(): + self.base_optimizer.state[p]["latent"] = self._state[p]["latent"] + del self._state + + # check if it is time to update set of quantization values Q + if (self.num_steps - self.warmup_steps) % self.quant_period == 0: + quant_update = True + else: + quant_update = False + + for group in self.regularized_param_groups(): + # AProx in practice: ensure shrinkage coefficient >= 1 + group["cumu_lr"] += group["lr"] + gamma = max(1.0, group["cumu_lr"]) + b = group["quant_bits"] + block_size = group.get("quant_block_size") + inv_slope = 0.0 + for p in group["params"]: + if not p.requires_grad: + continue + state = self.state[p] + # save latent parameters, need detach()? or copy p) + state["latent"].copy_(p) + + # in-place scaling of parameters by 1/gamma if specified + if self.quant_shrink: + p.div_(gamma) + + # reshape p according to block size if specified + if block_size is not None: + assert p.size(-1) % block_size == 0, ( + f"{p.size(-1)=} is not divisible by {block_size=}" + ) + assert p.dim() <= 2, f"Invalid {p.dim()=} for {block_size=}" + if p.dim() == 1: + p = p.unsqueeze(0) + + # row-major ordering ensures this is correct + p = p.view(-1, block_size) + + # quantization by channel or by layer + # update quantization targets periodically + per_channel = self.quant_per_channel and p.dim() > 1 + if quant_update: + quant_size = self.quantizer.get_quant_size(b) + + if per_channel: + quant_size = (p.size(0), quant_size) + state["quants"] = torch.empty(quant_size, device=p.device) + if is_dtensor(p): + state["quants"] = distribute_tensor( + state["quants"], + device_mesh=p.device_mesh, + placements=p.placements, + ) + + dim = -1 if per_channel else None + if per_channel and p.dim() > 2: + p = p.flatten(start_dim=1) + + q = None + if quant_update: + qfunc = partial( + self.quantize_, quantizer=self.quantizer, b=b, dim=dim + ) + if is_dtensor(p): + qfunc = local_map( + qfunc, + out_placements=[*p.placements], + in_placements=([Shard(0)], [Shard(0)]), + ) + q = qfunc(p, state["quants"]) + + # apply (step-dependent) proximal mapping in place + pfunc = partial( + self.prox_map.apply_, step_count=self.num_steps, dim=dim + ) + if is_dtensor(p): + pfunc = local_map( + pfunc, + out_placements=None, + in_placements=( + [Shard(0)], + None if q is None else [Shard(0)], + [Shard(0)], + ), + ) + inv_slope = pfunc(p, q, state["quants"]) + + # quantized parameters share the same PARQ inverse slope + if inv_slope: + if self.anneal_wd_frac > 0: + group["weight_decay"] = ( + inv_slope * self.anneal_wd_frac * group["initial_wd"] + + (1 - self.anneal_wd_frac) * group["initial_wd"] + ) + group["inv_slope"] = inv_slope # save for tensorboard + + self.num_steps += 1 + return loss + + @torch._disable_dynamo + @torch.no_grad() + def restore_latent_params(self) -> None: + """Restore latent parameters as optimizer parameters""" + for group in self.regularized_param_groups(): + for p in group["params"]: + if p.requires_grad: + p.copy_(self.state[p]["latent"]) + + @torch._disable_dynamo + @torch.no_grad() + def save_latent_params(self) -> None: + """Save updated latent parameters before applying prox-map""" + if self.warmup_steps == 0: + assert len(self.state) == 0, "Expected empty state at first step()" + # Maintain the invariant that `len(self.state) == 0` before first + # self.base_optimizer.step() call by using a temporary state buffer + self._state = defaultdict(dict) + + for group in self.regularized_param_groups(): + for p in group["params"]: + if p.requires_grad: + state = self.state[p] + if "latent" not in state: + state["latent"] = p.detach().clone() + else: + state["latent"].copy_(p) diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/__init__.py b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c8b83657259c8e9684b21d2397edefe5240c846c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from .lsbq import LSBQuantizer # noqa: F401 +from .quantizer import Quantizer # noqa: F401 +from .uniform import ( # noqa: F401 + MaxUnifQuantizer, + TernaryUnifQuantizer, + UnifQuantizer, +) +from .uniform_torchao import ( # noqa: F401 + Int4UnifTorchaoQuantizer, + UnifTorchaoQuantizer, +) diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f2569de5b5a9f77af92e73b9158724cb7d588f9 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/lsbq.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/lsbq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adac665eff3c4117fe2ddcfe25769b9cb3681ce7 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/lsbq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93b8616f781c22352f06bf17b19f4846afa432ee Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/uniform.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/uniform.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59ece3de0ec703f064b8eb2947b99d91c73f3907 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/uniform.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/uniform_torchao.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/uniform_torchao.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bc7357f2f65ff3e180fff3e72096b96e5748357 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/parq/quant/__pycache__/uniform_torchao.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/lsbq.py b/lib/python3.12/site-packages/torchao/prototype/parq/quant/lsbq.py new file mode 100644 index 0000000000000000000000000000000000000000..2d9f4e4c1eb0d78de98427b1f5e6fa4725f40613 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/quant/lsbq.py @@ -0,0 +1,276 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import itertools +from collections.abc import Iterable +from typing import Optional + +import torch +from torch import Tensor + +from ..utils import channel_bucketize +from .quantizer import Quantizer + + +def binary_sign(input: Tensor) -> Tensor: + """Same as `torch.sign(input)` but map 0 to 1.""" + return torch.where(input == 0, 1.0, input.sign()) + + +def binary_quant_residue(u: Tensor, vs: Iterable[float]) -> Tensor: + """Return residue for foldable binary quantization""" + r = u.detach().clone() + for v in vs: + r.sub_(v * binary_sign(r)) + return r + + +def compute_v_per_channel(p: Tensor, dim: Optional[int] = None, ternary: bool = False): + """Vectorized computation of optimal `v` for ternary/2-bit algorithm.""" + v_cands = p.abs().sort(dim=dim).values + cumsum = v_cands.cumsum(dim=dim) + cumsum, total_sum = cumsum[:, 1:-1], cumsum[:, -1:] + + # compute cumulative mean from right to left + counts = torch.arange(1, p.size(dim=dim), device=p.device) + counts_r2l = counts[:-1].flip((-1,)) + cmean_r2l = (total_sum - cumsum).div_(counts_r2l.mul_(2)) + v_cands, v_cands2 = v_cands[:, 1:-1], v_cands[:, 2:] + + # mask to estimate conditional expectation + mask = (v_cands <= cmean_r2l).logical_and_(v_cands2 >= cmean_r2l) + if ternary: + # detect and fix any edge cases + optimal_v = p.mean(dim=dim, keepdim=True).div_(2) + row_invalid = optimal_v < p.min(dim=dim, keepdim=True).values + if row_invalid.any(): + extra_col = row_invalid.to(p.dtype).mul(optimal_v) + v_cands = torch.cat((v_cands, extra_col), -1) + mask = torch.cat((mask, row_invalid), -1) + else: + # compute cumulative mean from left to right + cmean_l2r = cumsum.div_(counts[1:].mul_(2)).add_(cmean_r2l) + mask.logical_or_((v_cands <= cmean_l2r).logical_and_(v_cands2 >= cmean_l2r)) + + # handle variable number of candidates per channel + split_sizes = mask.sum(dim=dim).tolist() + v_cands = v_cands[mask].split(split_sizes) + v_cands = torch.nested.nested_tensor(list(v_cands)) + v_cands = torch.nested.to_padded_tensor(v_cands, 0.0) + + # update residual for each candidate `v` + r = p.unsqueeze(dim - 1) + v = v_cands.unsqueeze(-1) + r = r.sub(v * binary_sign(r)) + if not ternary: + v = v.mean(dim=dim, keepdim=True) + r = r.sub(v * binary_sign(r)) + + # compute least squares error, then select the `v` minimizes it + costs = r.norm(dim=dim) + indices = costs.argmin(dim=dim, keepdim=True) + v_best = v_cands.gather(1, indices) + return v_best + + +class LSBQuantizer(Quantizer): + """Least-Square Binary Quantizer, using greedy algorithm by default. + Optimal solution available for three cases: b=1, b=2 and ternary. + """ + + def __init__( + self, + center: bool = False, + optimal: bool = False, + ternary_multiplier: float = 1.1, + ) -> None: + """ternary_mult is factor multipled to 1-bit range as heuristic.""" + super().__init__(center) + # optimal choice only meaningful for b=1, b=2 and ternary + self.optimal = optimal + self.ternary_multiplier = ternary_multiplier + + def get_quant_size(self, b: int) -> int: + return 2**b if b > 0 else 3 + + def quantize( + self, p: Tensor, b: int, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + """Instantiation of Quantizer.quantize(), with b=0 for ternary""" + if b < 0: + raise ValueError(f"Invalid {b=}; must be nonnegative") + if self.optimal and b > 2: + raise NotImplementedError(f"Unsupported {self.optimal=} for {b=}") + + if self.center: + q, mean = super().remove_mean(p.detach(), dim=dim) + else: + q = p.detach().clone() + mean = torch.zeros(1, dtype=p.dtype, device=p.device) + + if self.optimal and b != 1: # b == 1 optimal is the same as greedy + if b == 0: + q, Q = self.quantize_optimal_ternary(q, dim=dim) + elif b == 2: + q, Q = self.quantize_optimal_2bits(q, dim=dim) + elif b == 0: + q, Q = self.quantize_simple_ternary(q, self.ternary_multiplier, dim=dim) + else: + q, Q = self.quantize_greedy(q, b, dim=dim) + + # return quantized tensor and set of quantization values + if self.center: + q += mean + Q += mean + return q, Q + + @staticmethod + def quantize_greedy( + p: Tensor, b: int, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + r = p.detach().clone() + vs = [] + keepdim = dim is not None + for _ in range(b): + v = r.abs().mean(dim=dim, keepdim=keepdim) + r.sub_(binary_sign(r).mul_(v)) + vs.append(v) + q = p - r + + # generate 2^b by b basis tensor B + basis = list(itertools.product((-1, 1), repeat=b)) + B = torch.tensor(basis, dtype=p.dtype, device=p.device) + if dim is not None: + V = torch.concat(vs, dim=1) # [dim0, b] + Q = torch.sort(V @ B.T, dim=dim).values # [dim0, 2^b] + else: + V = torch.tensor(vs, dtype=p.dtype, device=p.device) + Q = torch.msort(B.matmul(V)) # [2^b] + return q, Q + + @staticmethod + def quantize_optimal_2bits( + p: Tensor, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + # generate 4 x 2 basis tensor B, sorted lexicographically along dim 0 + basis = list(itertools.product((-1, 1), repeat=2)) + B = torch.tensor(basis, dtype=p.dtype, device=p.device) + if dim is not None: + v1 = compute_v_per_channel(p, dim=dim, ternary=False) + s = binary_sign(p).mul_(v1) + r = p.sub(s) + v2 = r.abs().mean(dim=dim, keepdim=True) + q = s.add_(binary_sign(r).mul_(v2)) + + V = torch.cat((v1, v2), dim=-1) # [dim0, b] + Q = V @ B.T # [dim0, 2^b] + return q, Q + + # first form the cumulative sum of sorted absolute values of p + p_abs_sorted = p.abs().flatten().sort().values + cumsum = torch.cumsum(p_abs_sorted, dim=0) + n = cumsum.numel() + # find all solutions v1 to an inclusion problem (after sorting |p|) + # |p|_{i} <= v1 < |p|_{i+1} + # where v1 = (1/2) * (avg_{0:i}(|p|) + avg_{i+1:n-1}(|p|)) + V1V2 = [] + v1 = cumsum[-1] / n / 2 + v2 = torch.zeros_like(v1) + q = torch.zeros_like(p) + # check if v1 < |p|.min(), which means v1 == v2, effectively ternary + if v1 < p_abs_sorted[0]: + V1V2.append((v1, v1)) + for i in range(n - 1): + E_above = (cumsum[-1] - cumsum[i]) / (n - (i + 1)) + E_below = cumsum[i] / (i + 1) + # skip if E_above < E_below, implying v2 < 0, invalid solution + if E_above < E_below: + continue + v1 = (E_above + E_below) / 2 + if v1 >= p_abs_sorted[i] and v1 < p_abs_sorted[i + 1]: + v2 = (E_above - E_below) / 2 + assert v1 >= v2, "LSBQ 2-bit optimal: v1 should >= v2." + V1V2.append((v1, v2)) + assert len(V1V2) > 0, "LSBQ 2-bit optimal: No solution found." + # find the best solution with least-square quantization error + min_error = p.norm() + for v1v2 in V1V2: + r = binary_quant_residue(p, v1v2) + error = r.norm() + if error < min_error: + min_error = error + q = p - r + v1, v2 = v1v2 + + V = torch.tensor((v1, v2), dtype=p.dtype, device=p.device) + Q = B @ V + return q, Q + + @staticmethod + def quantize_optimal_ternary( + p: Tensor, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + """Formula look reasonable, but derivation in reference incorrect?""" + if dim is not None: + v = compute_v_per_channel(p, dim=dim, ternary=True) + p_sign = binary_sign(p) + r = p.sub(p_sign.mul(v)) + + # 0 if sign(p) != sign(r), else sign(p) * 2v + q = p_sign.add_(binary_sign(r)).mul_(v) + + # each channel can take values [-2v, 0, 2v] + v.mul_(2) + Q = torch.cat((-v, torch.zeros_like(v), v), dim=-1) # [dim0, 3] + return q, Q + + # first form the cumulative sum of sorted absolute values of p + p_abs_sorted = p.abs().flatten().sort().values + cumsum = torch.cumsum(p_abs_sorted, dim=0) + n = cumsum.numel() + # find all solutions v1 to an inclusion problem (after sorting |p|) + # |p|_{i} <= v < |p|_{i+1} where v = (1/2) * avg_{i+1:n-1}(|p|)) + v_feasible = [] + v = cumsum[-1] / n / 2 + # check if v < |p|.min(), which means v1 == v2, effectively ternary + if v < p_abs_sorted[0]: + v_feasible.append(v) + for i in range(n - 1): + v = ((cumsum[-1] - cumsum[i]) / (n - (i + 1))) / 2 + if v >= p_abs_sorted[i] and v < p_abs_sorted[i + 1]: + v_feasible.append(v) + assert len(v_feasible) > 0, "LSBQ ternary optimal: No solution found." + # find the best solution with least-square quantization error + min_error = p.norm() + q_best = torch.zeros_like(p) + v_best = torch.zeros_like(v) + for v in v_feasible: + Q = v * torch.tensor([-1.0, 0.0, 1.0], device=p.device) + boundaries = v * torch.tensor([-0.5, 0.5], device=p.device) + q = Q[torch.bucketize(p, boundaries)] + error = torch.linalg.norm(p - q) + if error < min_error: + min_error = error + q_best = q + v_best = v + Q = v_best * torch.tensor([-1, 0, 1], dtype=p.dtype, device=p.device) + return q_best, Q + + @staticmethod + def quantize_simple_ternary( + p: Tensor, multiplier: float, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + """Heuristic by setting v as given multiple of |p|.abs().mean()""" + v = multiplier * p.abs().mean(dim=dim, keepdim=dim is not None) + Q = v * torch.tensor([-1.0, 0.0, 1.0], device=p.device) + if dim is None: + boundaries = v * torch.tensor([-0.5, 0.5], device=p.device) + q = Q[torch.bucketize(p, boundaries)] + else: + # for each row, find the element-wise index of least upper bound + boundaries = v * torch.tensor([[-0.5, 0.5]], device=p.device) + q = Q.gather(1, channel_bucketize(p, boundaries)) + return q, Q diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/quantizer.py b/lib/python3.12/site-packages/torchao/prototype/parq/quant/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..b44050e77390a244b23fe380c99ac83ff1bb70c5 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/quant/quantizer.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from abc import ABC, abstractmethod +from typing import Optional + +from torch import Tensor + + +class Quantizer(ABC): + """Abstract base class that defines the quantization interface""" + + def __init__(self, center: bool = False) -> None: + self.center = center + + @abstractmethod + def get_quant_size(self, b: int) -> int: + """Given number of bits b, return total number of quantization values""" + + @abstractmethod + def quantize(self, p: Tensor, b: int) -> tuple[Tensor, Tensor]: + """Provide interface for quantization: + q, Q = Quantizer.quantize(p) + Inputs: + p (Tensor): tensor to be quantized + Outputs: + q (Tensor): quantized tensor of same size as p + Q (Tensor): set of 2^b quantization values + Instantiation should not modify p, leaving update to ProxMap. + """ + + @staticmethod + def remove_mean(p: Tensor, dim: Optional[int] = None) -> tuple[Tensor, Tensor]: + """Center parameters in a Tensor, called if self.center == True. + Note that this is different from direct asymmetric quantization, + and may lead to (hopefully only) slightly different performance. + """ + mean = p.mean(dim=dim, keepdim=dim is not None) + q = p - mean + return q, mean diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/uniform.py b/lib/python3.12/site-packages/torchao/prototype/parq/quant/uniform.py new file mode 100644 index 0000000000000000000000000000000000000000..f2648941159ea8b0a84e3b902749d2f8d65f32fb --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/quant/uniform.py @@ -0,0 +1,214 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import math +from typing import Optional + +import torch +from torch import Tensor + +from .quantizer import Quantizer + + +def get_q_max( + q: Tensor, b: int, dim: Optional[int] = None, scale_method: str = "mean" +) -> Tensor: + if scale_method == "mean": + # set range of quantization: min(b * |q|.mean(), |q|.max()) + q_abs = q.abs() + if dim is not None: + q_max = torch.minimum( + b * q_abs.mean(dim=dim, keepdim=True), # pyre-ignore[6,9] + torch.max(q_abs, dim=dim, keepdim=True).values, # pyre-ignore[6] + ) + else: + q_max = torch.minimum(b * q_abs.mean(), torch.max(q_abs)) # pyre-ignore[6] + elif scale_method == "max": + q_max = ( + q.abs().max(dim=dim, keepdim=True).values + if dim is not None + else q.abs().max() + ) + else: + raise NotImplementedError(f"Invalid {scale_method=}, choices=('mean','max')") + return q_max + + +class UnifQuantizer(Quantizer): + """Uniform and symmetric quantizer""" + + def __init__( + self, + center: bool = False, + scale_method: str = "mean", + int_shift: float = 0.5, + zero_point: float = 0.5, + ): + """Set quantization function parameters. + + Args: + center: whether to subtract p.mean() prior to quantization + scale_method: compute scale based 'mean', multiples of |p|.mean(), + or 'max', |p|.max() (default: 'mean') + int_shift: float value to shift the lower bound of integer range by: + -2^{b - 1} + int_shift (default: 0.5). Using 0.5 results in 2^b + values. E.g., [-1.5, -0.5, 0.5, 1.5] for b=2. + zero_point: float value to shift p by after scale and round. + """ + assert scale_method in ("max", "mean"), f"Invalid {scale_method=}" + super().__init__(center=center) + + self.scale_method = scale_method + self.int_shift = int_shift + self.zero_point = zero_point + + def get_quant_size(self, b: int) -> int: + """Levels in [-2^{b-1} + self.int_shift, 2^{b-1} - self.int_shift]. + + Note that range_absmax = 2^{b-1} - self.int_shift on both ends of the + boundary and the interval is closed.""" + return math.floor(2**b - 2 * self.int_shift) + 1 + + def quantize( + self, p: Tensor, b: int, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + """Instantiation of Quantizer.quantize() method""" + assert b != 0, "Please use TernaryUnifQuantizer instead" + + if self.center: + q, mean = super().remove_mean(p.detach(), dim=dim) + else: + q = p.detach().clone() + mean = torch.zeros(1, dtype=p.dtype, device=p.device) + q_max = get_q_max(q, b, dim=dim, scale_method=self.scale_method) + q_max.clamp_(min=torch.finfo(q.dtype).tiny) + + # clamp to quantization range + q.copy_(torch.minimum(torch.maximum(q, -q_max), q_max)) + + # scale from [-2^{b-1}+int_shift, 2^{b-1}-int_shift] to [-q_max, q_max] + range_absmax = 2 ** (b - 1) - self.int_shift + s = q_max / range_absmax + + # scale by 1/s -> shift -zero_point -> round -> shift +zero_point -> + # scale by s, where shift ensures rounding to integers + q.div_(s).sub_(self.zero_point).round_().add_(self.zero_point).mul_(s) + + # set of all target quantization values + Q = torch.arange( + -range_absmax, range_absmax + 1e-5, dtype=p.dtype, device=p.device + ) + if dim is not None: + Q = Q.unsqueeze(0).mul(s) # broadcasted multiply requires copy + else: + Q.mul_(s) + + # return quantized tensor and set of possible quantization values + if self.center: + q += mean + Q += mean + return q, Q + + +class MaxUnifQuantizer(UnifQuantizer): + def __init__( + self, + center: bool = False, + scale_method: str = "max", + int_shift: float = 1.0, + zero_point: float = 0.0, + ): + """Set quantization function with int_shift=1.0. + + The final quantization range includes 2^b - 1 quantized values. E.g., + [-1, 0, 1] for b=2. The quantization scale is determined by |p|.max() + by default and zero point is 0.0. + """ + super().__init__( + center=center, + scale_method=scale_method, + int_shift=int_shift, + zero_point=zero_point, + ) + + +class AsymUnifQuantizer(Quantizer): + def get_quant_size(self, b: int) -> int: + """Equivalent to int_max - int_min + 1, where int_min = -2^{b-1} and + int_max = 2^{b-1} - 1.""" + return 2**b + + def quantize( + self, p: Tensor, b: int, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + assert b != 0, "Please use TernaryUnifQuantizer instead" + + if self.center: + q, mean = super().remove_mean(p.detach(), dim=dim) + else: + q = p.detach().clone() + mean = torch.zeros(1, dtype=p.dtype, device=p.device) + + if dim is not None: + q_min = q.min(dim=dim, keepdim=True).values + q_max = q.max(dim=dim, keepdim=True).values + else: + q_min = q.min() + q_max = q.max() + + int_min = -(2 ** (b - 1)) + int_max = 2 ** (b - 1) - 1 + s = (q_max - q_min) / (int_max - int_min) + s.clamp_(min=torch.finfo(q.dtype).tiny) + + zero_point = q_min.div_(s).round_() + q.div_(s).round_().sub_(zero_point).add_(zero_point).mul_(s) + + Q = torch.arange(int_min, int_max + 1, dtype=p.dtype, device=p.device) + if dim is not None: + Q = Q.unsqueeze(0).mul(s) # broadcasted multiply requires copy + else: + Q.mul_(s) + + # return quantized tensor and set of possible quantization values + if self.center: + q += mean + Q += mean + return q, Q + + +class TernaryUnifQuantizer(Quantizer): + """Uniform quantizer for ternary bit case. Quantization range is [-1, 1].""" + + def get_quant_size(self, b: int) -> int: + return 3 + + def quantize( + self, p: Tensor, b: int, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + assert b == 0, f"Unexpected {b=} for ternary case" + + if self.center: + q, mean = super().remove_mean(p.detach(), dim=dim) + else: + q = p.detach().clone() + mean = torch.zeros(1, dtype=p.dtype, device=p.device) + + q_max = get_q_max(q, b, dim=dim, scale_method="max") + q_max.clamp_(min=torch.finfo(q.dtype).tiny) + s = q_max / 1.5 + q.div_(s).round_().clamp_(min=-1, max=1).mul_(s) + + Q = torch.tensor([-1, 0, 1], dtype=p.dtype, device=p.device) + if dim is not None: + Q = Q.unsqueeze(0).mul(s) + else: + Q.mul_(s) + + if self.center: + q += mean + Q += mean + return q, Q diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/quant/uniform_torchao.py b/lib/python3.12/site-packages/torchao/prototype/parq/quant/uniform_torchao.py new file mode 100644 index 0000000000000000000000000000000000000000..0579a23b025510331647631579efea9e70864f36 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/quant/uniform_torchao.py @@ -0,0 +1,131 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Optional, Union + +import torch +from torch import Tensor + +from torchao.quantization.quant_primitives import ( + _DTYPE_TO_BIT_WIDTH, + _DTYPE_TO_QVALUE_BOUNDS, + MappingType, + ZeroPointDomain, + choose_qparams_affine, + dequantize_affine, + quantize_affine, +) + +from .quantizer import Quantizer + +_BIT_WIDTH_TO_DTYPE = {v: k for k, v in _DTYPE_TO_BIT_WIDTH.items()} + + +class UnifTorchaoQuantizer(Quantizer): + """Uniform quantizer that uses torchao's quantization primitives""" + + def __init__( + self, + mapping_type: MappingType = MappingType.SYMMETRIC, + target_dtype: torch.dtype = torch.int8, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + eps: Optional[float] = None, + preserve_zero: bool = True, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + ) -> None: + super().__init__(center=False) + + self.mapping_type = mapping_type + self.target_dtype = target_dtype + self.quant_min = quant_min + self.quant_max = quant_max + self.eps = eps + self.preserve_zero = preserve_zero + self.zero_point_domain = zero_point_domain + + def _init_quant_min_max(self, b: int) -> None: + if self.quant_min is None or self.quant_max is None: + assert b in _BIT_WIDTH_TO_DTYPE, f"Unsupported bitwidth {b}" + self.quant_min, self.quant_max = _DTYPE_TO_QVALUE_BOUNDS[ + _BIT_WIDTH_TO_DTYPE[b] + ] + + def get_quant_size(self, b: int) -> int: + self._init_quant_min_max(b) + return self.quant_max - self.quant_min + 1 + + def quantize( + self, p: Tensor, b: int, dim: Optional[int] = None + ) -> tuple[Tensor, Tensor]: + self._init_quant_min_max(b) + if self.eps is None: + self.eps = torch.finfo(p.dtype).eps + + # assume that p has already been grouped in QuantOptimizer.step + block_size = (1, p.size(-1)) if dim is not None else p.size() + s, zero_point = choose_qparams_affine( + p, + self.mapping_type, + block_size, + self.target_dtype, + eps=self.eps, + preserve_zero=self.preserve_zero, + quant_min=self.quant_min, + quant_max=self.quant_max, + zero_point_domain=self.zero_point_domain, + ) + q_args = (block_size, s, zero_point, self.target_dtype) + q = quantize_affine( + p, + *q_args, + quant_min=self.quant_min, + quant_max=self.quant_max, + zero_point_domain=self.zero_point_domain, + ) + q = dequantize_affine( + q, + *q_args, + output_dtype=p.dtype, + quant_min=self.quant_min, + quant_max=self.quant_max, + zero_point_domain=self.zero_point_domain, + ) + + Q = torch.arange( + self.quant_min, self.quant_max + 1, dtype=self.target_dtype, device=p.device + ) + if dim is not None: + Q = Q.view(1, -1).expand(q.size(0), -1) + block_size = (1, Q.size(-1)) + else: + block_size = Q.shape + + Q = dequantize_affine( + Q, + block_size, + *q_args[1:], + output_dtype=p.dtype, + quant_min=self.quant_min, + quant_max=self.quant_max, + zero_point_domain=self.zero_point_domain, + ) + return q, Q + + +class Int4UnifTorchaoQuantizer(UnifTorchaoQuantizer): + """Based on torchao.quantization.quant_api._int4_weight_only_transform""" + + def __init__(self) -> None: + super().__init__( + mapping_type=MappingType.ASYMMETRIC, + target_dtype=torch.int32, + quant_min=0, + quant_max=15, + eps=1e-6, + preserve_zero=False, + zero_point_domain=ZeroPointDomain.FLOAT, + ) diff --git a/lib/python3.12/site-packages/torchao/prototype/parq/utils.py b/lib/python3.12/site-packages/torchao/prototype/parq/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ac5024fb5d3740526d658fcdeb942029bcb55aca --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/parq/utils.py @@ -0,0 +1,33 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import Tensor + +try: + from torch.distributed.tensor import DTensor + + HAS_DTENSOR = True +except ImportError: + HAS_DTENSOR = False + + +def is_dtensor(x): + return HAS_DTENSOR and isinstance(x, DTensor) + + +def channel_bucketize(input: Tensor, boundaries: Tensor, right: bool = False) -> Tensor: + """Generalizes torch.bucketize to run on 2-D boundaries.""" + inf_pad = torch.full_like(boundaries[:, :1], torch.inf) + boundaries = ( + torch.cat((-inf_pad, boundaries), dim=1) + if right + else torch.cat((boundaries, inf_pad), dim=1) + ) + boundaries = boundaries.unsqueeze(1) + input = input.unsqueeze(-1) + mask = input.ge(boundaries) if right else input.le(boundaries) + return mask.to(torch.uint8).argmax(dim=-1) diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__init__.py b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c6278884a05798e4615b1b456d9809a752f6c43 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__init__.py @@ -0,0 +1,3 @@ +from torchao.prototype.scaled_grouped_mm.scaled_grouped_mm import ( + _scaled_grouped_mm as _scaled_grouped_mm, +) diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b16db7c7aed4c55f4913deaff853218ee290fb4 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/scaled_grouped_mm.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/scaled_grouped_mm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f3bb85816a594a0b5fe385a260a0a2ce7fe5355 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/scaled_grouped_mm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d4e85ec645cfd63bfb9d74ac6bf0e26e6febd54 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__init__.py b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1c753035686bad095b6f07e37640af04414e444f --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__init__.py @@ -0,0 +1,6 @@ +from torchao.prototype.scaled_grouped_mm.kernels.jagged_float8_scales import ( + triton_fp8_col_major_jagged_colwise_scales as triton_fp8_col_major_jagged_colwise_scales, +) +from torchao.prototype.scaled_grouped_mm.kernels.jagged_float8_scales import ( + triton_fp8_row_major_jagged_rowwise_scales as triton_fp8_row_major_jagged_rowwise_scales, +) diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..000994db80afb6372fd4b9ead2b38f199cf77367 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__pycache__/jagged_float8_scales.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__pycache__/jagged_float8_scales.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2ef72e3a6b4eb6268a0105fd2c362d4662dd3bd Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/__pycache__/jagged_float8_scales.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/jagged_float8_scales.py b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/jagged_float8_scales.py new file mode 100644 index 0000000000000000000000000000000000000000..4cc6177a48ecb35ae85292afa39542dffd159fd3 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/kernels/jagged_float8_scales.py @@ -0,0 +1,366 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +""" +Triton kernels for scaling high precision tensors to float8 using "jagged" +rowwise scales (i.e., separate scales for each group/subtensor as determined by +the offsets). +""" + +from typing import Tuple + +import torch +import triton +import triton.language as tl + +from torchao.prototype.scaled_grouped_mm.utils import _is_column_major + +EPS = 1e-12 + +FP8_DTYPE_MAP = { + torch.int8: tl.int8, + torch.int16: tl.int16, + torch.int32: tl.int32, + torch.int64: tl.int64, + torch.float8_e4m3fn: tl.float8e4nv, + torch.float8_e5m2: tl.float8e5, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, + torch.float32: tl.float32, + torch.float64: tl.float64, +} + +block_sizes = [128, 256] +kernel_configs_2D = [ + triton.Config( + {"BLOCK_SIZE_ROWS": block_size_rows, "BLOCK_SIZE_COLS": block_size_cols} + ) + for block_size_rows in block_sizes + for block_size_cols in block_sizes +] + + +def triton_fp8_row_major_jagged_rowwise_scales( + hp_tensor: torch.Tensor, + offsets: torch.Tensor, + output_dtype: torch.dtype = torch.float8_e4m3fn, + round_scales_to_power_of_2: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Converts a high precision tensor to a float8 tensor in row-major memory layout, + using 'jagged' rowwise scales (i.e., separate scales for each group/subtensor as + determined by the offsets). + + Args: + - hp_tensor: 2D high precision tensor to be converted + - offsets: end index for each group/subtensor along dim 0 + - output_dtype: desired float8 dtype for the output tensor + - round_scales_to_power_of_2: boolean indicating if scales should be rounded + down to the nearest power of 2. + Returns: + - float8 tensor + - jagged rowwise scales (i.e., rowwise scales for each group) + """ + assert hp_tensor.ndim == 2, "input tensor must be 2D" + assert hp_tensor.is_contiguous(), "input tensor must be contiguous" + + num_elements = hp_tensor.numel() + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[output_dtype] + + fp8_dtype_min = torch.finfo(output_dtype).min + fp8_dtype_max = torch.finfo(output_dtype).max + + m, k = hp_tensor.shape + n_groups = offsets.numel() + + # allocate on-device buffers for output and scales + output_buffer = torch.empty_like( + hp_tensor, dtype=output_dtype, device=hp_tensor.device + ) + scales_buffer = torch.empty( + (m * n_groups), dtype=torch.float32, device=hp_tensor.device + ) + + # parallelize across rows and groups (offsets) + grid = lambda meta: ( + triton.cdiv(m, meta["BLOCK_SIZE_ROWS"]), + offsets.numel(), + ) + _triton_fp8_row_major_jagged_rowwise_scales[grid]( + hp_tensor, + offsets, + output_buffer, + scales_buffer, + m, + k, + hp_tensor.stride(0), + hp_tensor.stride(1), + output_buffer.stride(0), + output_buffer.stride(1), + num_elements, + fp8_dtype_min, + fp8_dtype_max, + tl_input_dtype, + tl_output_dtype, + round_scales_to_power_of_2, + EPS=EPS, + ) + return output_buffer, scales_buffer + + +@triton.autotune(configs=kernel_configs_2D, key=["num_elements"]) +@triton.jit +def _triton_fp8_row_major_jagged_rowwise_scales( + input_ptr, + offsets_ptr, + out_ptr, + scales_ptr, + M: int, + K: int, + stride_input_row: int, + stride_input_col: int, + stride_output_row: int, + stride_output_col: int, + num_elements: int, + fp8_dtype_min: tl.constexpr, + fp8_dtype_max: tl.constexpr, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + round_scales_to_power_of_2: tl.constexpr, + BLOCK_SIZE_ROWS: tl.constexpr, + BLOCK_SIZE_COLS: tl.constexpr, + EPS: tl.constexpr, +): + # parallel across rows and groups (offsets) + block_row_id = tl.program_id(axis=0) + offset_idx = tl.program_id(axis=1) + + # determine start and end column idx for this group + group_col_start_idx = tl.load( + offsets_ptr + offset_idx - 1, mask=offset_idx > 0, other=0 + ) + group_col_end_idx = tl.load(offsets_ptr + offset_idx) + block_row_offs = block_row_id * BLOCK_SIZE_ROWS + tl.arange(0, BLOCK_SIZE_ROWS) + + # compute rowwise amaxes for this group + amax_buffer = tl.zeros((BLOCK_SIZE_ROWS,), dtype=input_dtype) + for col_start_idx in range(group_col_start_idx, group_col_end_idx, BLOCK_SIZE_COLS): + block_col_offs = col_start_idx + tl.arange(0, BLOCK_SIZE_COLS) + block_offs = ( + block_row_offs[:, None] * stride_input_row + + block_col_offs[None, :] * stride_input_col + ) + block_mask = (block_row_offs[:, None] < M) & ( + block_col_offs[None, :] < group_col_end_idx + ) + data = tl.load(input_ptr + block_offs, mask=block_mask, other=0.0).to( + input_dtype + ) + # we need to cast back to input dtype since triton promotes bf16 to fp32: + # https://github.com/triton-lang/triton/blob/981e987eed9053b952f81153bc0779c99d8c642e/python/triton/language/standard.py#L173 + amax_buffer = tl.maximum(amax_buffer, tl.max(tl.abs(data), axis=1)).to( + input_dtype + ) + + # compute rowwise scales for this group. round scales to nearest power of 2. + amax_buffer = amax_buffer.to(tl.float64) + scales = (fp8_dtype_max / tl.clamp(amax_buffer, min=EPS, max=float("inf"))).to( + tl.float32 + ) + if round_scales_to_power_of_2: + scales = tl.exp2(tl.floor(tl.log2(scales))) + + # store rowwise scales for each group in contiguous memory: + # [group0_row0, group_0_row1, ..., group2_row0, group2_row1] + scales_offs = block_row_offs + (M * offset_idx) + scales_mask = tl.arange(0, BLOCK_SIZE_ROWS) < M + tl.store(scales_ptr + scales_offs, scales, mask=scales_mask) + + # perform float8 conversion for this group + for col_start_idx in range(group_col_start_idx, group_col_end_idx, BLOCK_SIZE_COLS): + block_col_offs = col_start_idx + tl.arange(0, BLOCK_SIZE_COLS) + block_offs = ( + block_row_offs[:, None] * stride_input_row + + block_col_offs[None, :] * stride_input_col + ) + block_mask = (block_row_offs[:, None] < M) & ( + block_col_offs[None, :] < group_col_end_idx + ) + data = tl.load(input_ptr + block_offs, mask=block_mask, other=0.0).to( + input_dtype + ) + scaled_data = data * scales[:, None] + fp8_data = tl.clamp(scaled_data, min=fp8_dtype_min, max=fp8_dtype_max).to( + output_dtype + ) + out_offs = ( + block_row_offs[:, None] * stride_output_row + + block_col_offs[None, :] * stride_output_col + ) + tl.store(out_ptr + out_offs, fp8_data, mask=block_mask) + + +def triton_fp8_col_major_jagged_colwise_scales( + hp_tensor: torch.Tensor, + offsets: torch.Tensor, + output_dtype: torch.dtype = torch.float8_e4m3fn, + round_scales_to_power_of_2: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Converts a high precision tensor to a float8 tensor in row-major memory layout, + using 'jagged' column-wise scales (i.e., separate scales for each group/subtensor as + determined by the offsets). + + Args: + - hp_tensor: 2D high precision tensor to be converted + - offsets: end index for each group/subtensor along dim 0 + - output_dtype: desired float8 dtype for the output tensor + - round_scales_to_power_of_2: boolean indicating if scales should be rounded + down to the nearest power of 2. + Returns: + - float8 tensor + - jagged column-wise scales (i.e., column-wise scales for each group) + """ + assert hp_tensor.ndim == 2, "input tensor must be 2D" + assert _is_column_major(hp_tensor), "input tensor must be column-major" + + num_elements = hp_tensor.numel() + tl_input_dtype = FP8_DTYPE_MAP[hp_tensor.dtype] + tl_output_dtype = FP8_DTYPE_MAP[output_dtype] + + fp8_dtype_min = torch.finfo(output_dtype).min + fp8_dtype_max = torch.finfo(output_dtype).max + + k, n = hp_tensor.shape + n_groups = offsets.numel() + + # allocate on-device buffers for output and scales + output_buffer = torch.empty_like( + hp_tensor, dtype=output_dtype, device=hp_tensor.device + ) + scales_buffer = torch.empty( + (n * n_groups), dtype=torch.float32, device=hp_tensor.device + ) + + # parallelize across columns and groups (offsets) + grid = lambda meta: ( + triton.cdiv(n, meta["BLOCK_SIZE_COLS"]), + offsets.numel(), + ) + _triton_fp8_col_major_jagged_colwise_scales[grid]( + hp_tensor, + offsets, + output_buffer, + scales_buffer, + k, + n, + hp_tensor.stride(0), + hp_tensor.stride(1), + output_buffer.stride(0), + output_buffer.stride(1), + num_elements, + fp8_dtype_min, + fp8_dtype_max, + tl_input_dtype, + tl_output_dtype, + round_scales_to_power_of_2, + EPS=EPS, + ) + return output_buffer, scales_buffer + + +@triton.autotune(configs=kernel_configs_2D, key=["num_elements"]) +@triton.jit +def _triton_fp8_col_major_jagged_colwise_scales( + input_ptr, + offsets_ptr, + out_ptr, + scales_ptr, + K: int, + N: int, + stride_input_row: int, + stride_input_col: int, + stride_output_row: int, + stride_output_col: int, + num_elements: int, + fp8_dtype_min: tl.constexpr, + fp8_dtype_max: tl.constexpr, + input_dtype: tl.constexpr, + output_dtype: tl.constexpr, + round_scales_to_power_of_2: tl.constexpr, + BLOCK_SIZE_ROWS: tl.constexpr, + BLOCK_SIZE_COLS: tl.constexpr, + EPS: tl.constexpr, +): + # parallel across columns and groups (offsets) + block_col_id = tl.program_id(axis=0) + offset_idx = tl.program_id(axis=1) + + # determine start and end row idx for this group + group_row_start_idx = tl.load( + offsets_ptr + offset_idx - 1, mask=offset_idx > 0, other=0 + ) + group_row_end_idx = tl.load(offsets_ptr + offset_idx) + block_col_offs = block_col_id * BLOCK_SIZE_COLS + tl.arange(0, BLOCK_SIZE_COLS) + + # compute colwise amaxes for this group + amax_buffer = tl.zeros((BLOCK_SIZE_COLS,), dtype=input_dtype) + for row_start_idx in range(group_row_start_idx, group_row_end_idx, BLOCK_SIZE_ROWS): + block_row_offs = row_start_idx + tl.arange(0, BLOCK_SIZE_ROWS) + block_offs = ( + block_row_offs[:, None] * stride_input_row + + block_col_offs[None, :] * stride_input_col + ) + block_mask = (block_row_offs[:, None] < group_row_end_idx) & ( + block_col_offs[None, :] < N + ) + data = tl.load(input_ptr + block_offs, mask=block_mask, other=0.0).to( + input_dtype + ) + # we need to cast back to input dtype since triton promotes bf16 to fp32: + # https://github.com/triton-lang/triton/blob/981e987eed9053b952f81153bc0779c99d8c642e/python/triton/language/standard.py#L173 + amax_buffer = tl.maximum(amax_buffer, tl.max(tl.abs(data), axis=0)).to( + input_dtype + ) + + # compute rowwise scales for this group. + amax_buffer = amax_buffer.to(tl.float64) + scales = (fp8_dtype_max / tl.clamp(amax_buffer, min=EPS, max=float("inf"))).to( + tl.float32 + ) + if round_scales_to_power_of_2: + scales = tl.exp2(tl.floor(tl.log2(scales))) + + # store colwise scales for each group in contiguous memory: + # [group0_col0, group_0_col1, ..., group2_col0, group2_col1] + # note: input tensor is in col-major memory layout. + scales_offs = block_col_offs + (N * offset_idx) + scales_mask = tl.arange(0, BLOCK_SIZE_COLS) < N + tl.store(scales_ptr + scales_offs, scales, mask=scales_mask) + + # perform float8 conversion for this group + for row_start_idx in range(group_row_start_idx, group_row_end_idx, BLOCK_SIZE_ROWS): + block_row_offs = row_start_idx + tl.arange(0, BLOCK_SIZE_ROWS) + block_offs = ( + block_row_offs[:, None] * stride_input_row + + block_col_offs[None, :] * stride_input_col + ) + block_mask = (block_row_offs[:, None] < group_row_end_idx) & ( + block_col_offs[None, :] < N + ) + data = tl.load(input_ptr + block_offs, mask=block_mask, other=0.0).to( + input_dtype + ) + scaled_data = data * scales[None, :] + fp8_data = tl.clamp(scaled_data, min=fp8_dtype_min, max=fp8_dtype_max).to( + output_dtype + ) + out_offs = ( + block_row_offs[:, None] * stride_output_row + + block_col_offs[None, :] * stride_output_col + ) + tl.store(out_ptr + out_offs, fp8_data, mask=block_mask) diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/scaled_grouped_mm.py b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/scaled_grouped_mm.py new file mode 100644 index 0000000000000000000000000000000000000000..169e2c5407c12bbbba0fd4ea3886a5874a2c9d9f --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/scaled_grouped_mm.py @@ -0,0 +1,223 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Optional + +import torch + +from torchao.float8.config import ScalingGranularity +from torchao.float8.float8_utils import tensor_to_scale, to_fp8_saturated +from torchao.prototype.scaled_grouped_mm.kernels import ( + triton_fp8_col_major_jagged_colwise_scales, + triton_fp8_row_major_jagged_rowwise_scales, +) +from torchao.prototype.scaled_grouped_mm.utils import _is_column_major + + +def _scaled_grouped_mm( + A: torch.Tensor, + B_t: torch.Tensor, + offs: torch.Tensor, + out_dtype: Optional[torch.dtype] = torch.bfloat16, +) -> torch.Tensor: + """ + This function performs dynamic float8 quantization with row-wise scaling + on the input tensors A and B, then performs a scaled grouped GEMM and returns the results. + + Args: + A (bf16/float32 torch.Tensor): The first high-precision input tensor, which must be a 2D tensor of shape (M * num_groups, K) + and in row-major memory layout. + B_t (bf16/float32 torch.Tensor): The second high-precision input tensor which must be 3D, which must be shape (B, K, N) + and in column-major memory layout. + offs (int32 torch.Tensor): The offsets to use to mark the starting index of each group along dim0 of the A tensor. + out_dtype (Optional[torch.dtype]): The dtype of the output tensor. Currently only torch.bfloat16 is supported. + """ + return _Float8GroupedMM.apply( + A, + B_t, + offs, + out_dtype, + ) + + +class _Float8GroupedMM(torch.autograd.Function): + """Differentiable implementation of grouped GEMM with dynamic float8 quantization.""" + + @staticmethod + def forward( + ctx, + A: torch.Tensor, + B_t: torch.Tensor, + offs: torch.Tensor, + out_dtype: Optional[torch.dtype] = torch.bfloat16, + ) -> torch.Tensor: + # torchao _scaled_grouped_mm only supports A=2D, B=3D. + assert A.ndim == 2, "A must be 2D" + assert B_t.ndim == 3, "B must be 3D" + + assert A.size(-1) % 16 == 0, ( + f"A must have a last dim divisible by 16, but got shape: {A.shape}" + ) + assert B_t.size(-2) % 16 == 0 and B_t.size(-1) % 16 == 0, ( + f"B must have last 2 dims divisible by 16, but got shape: {B_t.shape}" + ) + + # Assert input tensors are in high-precision dtypes. + assert A.dtype == torch.float32 or A.dtype == torch.bfloat16, ( + "A must be float32 or bfloat16" + ) + assert B_t.dtype == torch.float32 or B_t.dtype == torch.bfloat16, ( + "B must be float32 or bfloat16" + ) + assert offs.dtype == torch.int32, "offs must be int32" + + # Assert A and B dims are compatible for a scaled grouped GEMM. + assert A.size(-1) == B_t.size(-2), ( + f"shape {A.shape} and {B_t.shape} are not compatible for _scaled_grouped_mm" + ) + + # The left operand in the scaled grouped GEMM must be row-major due to hardware requirements. + assert not _is_column_major(A), "A must be row-major" + + # Due to hardware requirements, the right operand in a scaled grouped GEMM must be column-major. + assert _is_column_major(B_t), "B must be column-major" + + # Convert high precision input tensor to float8, row-major for left operand of grouped GEMM. + # A shape: (M, K) + # A_scales shape: (M,1) + A_scales = tensor_to_scale( + A, + torch.float8_e4m3fn, + scaling_granularity=ScalingGranularity.AXISWISE, + axiswise_dim=-1, + round_scales_to_power_of_2=True, + ) + A_scaled = A.to(torch.float32) * A_scales + A_fp8_row_major = to_fp8_saturated(A_scaled, torch.float8_e4m3fn) + + # Convert B to float8, column-major for right operand of grouped GEMM. + # B shape: (B, K, N) + # B scales must be computed rowwise keeping the outer/final dim, so: + # B_scales shape: (B, 1, N) + B_t_scales = tensor_to_scale( + B_t, + torch.float8_e4m3fn, + scaling_granularity=ScalingGranularity.AXISWISE, + axiswise_dim=-2, + round_scales_to_power_of_2=True, + ) + B_t_scaled = B_t.to(torch.float32) * B_t_scales + B_t_fp8_col_major = to_fp8_saturated(B_t_scaled, torch.float8_e4m3fn) + + # Precompute non-transposed B column-major for backward, to save memory by storing the + # low precision B tensor instead of the high precision B tensor. + # In the backward this is needed for grad_A: grad_output @ B. + B = B_t.contiguous().transpose(-2, -1) + + # - B shape: (B, K, N) + # - B scales must be computed rowwise keeping the outer/final dim, so: + # - B_scale shape: (B, 1, N) + B_scales = tensor_to_scale( + B, + torch.float8_e4m3fn, + scaling_granularity=ScalingGranularity.AXISWISE, + axiswise_dim=-2, + round_scales_to_power_of_2=True, + ) + B_scaled = B.to(torch.float32) * B_scales + B_fp8_col_major = to_fp8_saturated(B_scaled, torch.float8_e4m3fn) + + # Store what we need for backward. + ctx.save_for_backward(A, B_fp8_col_major, B_scales, offs) + ctx.out_dtype = out_dtype + + # Perform scaled grouped GEMM and return result. + # output shape: scaled grouped mm of (M,K) @ (B,K,N) = (M,N) + return torch._scaled_grouped_mm( + A_fp8_row_major, + B_t_fp8_col_major, + A_scales.squeeze().reciprocal(), + B_t_scales.squeeze().reciprocal(), + offs, + out_dtype=out_dtype, + use_fast_accum=True, + ) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + A, B_fp8_col_major, B_scales, offs = ctx.saved_tensors + out_dtype = ctx.out_dtype + + # Convert grad_output to float8, row-major for left operand of grouped GEMM + # needed for grad_A: grad_output @ B + # + # grad_output shape: (M, N) + # grad_output_scale shape: (M, 1) + grad_output_scales = tensor_to_scale( + grad_output, + torch.float8_e4m3fn, + scaling_granularity=ScalingGranularity.AXISWISE, + axiswise_dim=-1, + round_scales_to_power_of_2=True, + ) + grad_output_scaled = grad_output.to(torch.float32) * grad_output_scales + grad_output_fp8_row_major = to_fp8_saturated( + grad_output_scaled, torch.float8_e4m3fn + ) + + # Compute grad_A. + # + # grad_A = grad_output @ B + # grad_A = scaled grouped mm of (M,N) @ (B,N,K) = (M,K) + grad_A = torch._scaled_grouped_mm( + grad_output_fp8_row_major, + B_fp8_col_major, + grad_output_scales.squeeze().reciprocal(), + B_scales.squeeze().reciprocal(), + offs, + out_dtype=out_dtype, + use_fast_accum=True, + ) + + # Convert tranpose of grad_output to float8, row-major for left operand of grouped GEMM + # needed for grad_B: grad_output_t @ A + grad_output_t_row_major = grad_output.transpose(-2, -1).contiguous() + + # Convert A to float8, column-major for right operand of grouped GEMM: + # needed for grad_B: grad_output @ A + A_col_major = A.transpose(-2, -1).contiguous().transpose(-2, -1) + + # grad_B is a special case. both operands of the grouped gemm will be 2D with offsets determing the "groups." + # Compute scales for grad_output_t and A, which are both 2D tensors with offsets which define the "jagged" groups. + grad_output_t_fp8_row_major, grad_output_t_scales = ( + triton_fp8_row_major_jagged_rowwise_scales( + grad_output_t_row_major, + offs, + output_dtype=torch.float8_e4m3fn, + round_scales_to_power_of_2=True, + ) + ) + + A_fp8_col_major, A_scales = triton_fp8_col_major_jagged_colwise_scales( + A_col_major, + offs, + output_dtype=torch.float8_e4m3fn, + round_scales_to_power_of_2=True, + ) + + # Compute grad_B = grad_output_t @ A. + # grad_B = grad_output_t @ A + # grad_B = (N,M) @ (M,K) = (N,K) + grad_B = torch._scaled_grouped_mm( + grad_output_t_fp8_row_major, + A_fp8_col_major, + grad_output_t_scales.reciprocal(), + A_scales.reciprocal(), + offs, + out_dtype=out_dtype, + use_fast_accum=True, + ) + return grad_A, grad_B.transpose(-2, -1), None, None, None, None diff --git a/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/utils.py b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..038c379d6250b8504ef971dc524dd1edc66d126d --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/scaled_grouped_mm/utils.py @@ -0,0 +1,156 @@ +from typing import Tuple + +import torch + +from torchao.float8.config import ScalingGranularity +from torchao.float8.float8_utils import tensor_to_scale, to_fp8_saturated + + +def _to_2d_jagged_float8_tensor_colwise( + A_col_major: torch.Tensor, + offs: torch.Tensor, + target_dtype: torch.dtype = torch.float8_e4m3fn, + round_scales_to_power_of_2: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + This function converts the 2D input tensor A to a jagged float8 tensor, + with scales computed along *logical columns* for each group individually, + where groups are determined based on the offsets. + + For the right operand of a normal scaled GEMM, the rowwise scales are computed over logical columns. + (i.e., a tensor of (K,N) will have scales of shape (1,N). + + However, for a 2D right operand of a grouped GEMM, these logical columns go through multiple distinct + groups/subtensors, for which we want to compute scales individually. So we cannot take one set of scales + along the logical columns and apply it to the entire tensor. + + Instead, we need to compute scales for each subtensor individually. For a tensor of shape (K,N) this results + in scales of shape (1,N * num_groups). + + Args: + A (torch.Tensor): The input tensor to be converted to a jagged float8 tensor. + + Returns: + A tuple containing the jagged float8 tensor and the scales used for the conversion. + """ + assert A_col_major.ndim == 2, "A must be 2D" + + num_groups = offs.numel() + A_fp8_col_major = torch.empty_like(A_col_major, dtype=target_dtype) + A_scales = torch.empty( + A_fp8_col_major.size(1) * num_groups, + dtype=torch.float32, + device=A_fp8_col_major.device, + ) + + start_idx = 0 + next_scale_idx = 0 + for end_idx in offs.tolist(): + # Get the subtensor of A for this group, fetching the next group of rows, with all columns for each. + subtensor = A_col_major[start_idx:end_idx, :] # (local_group_size, K) + + # Compute local rowwise scales for this subtensor, which are along logical columns for the right operand. + subtensor_scales = tensor_to_scale( + subtensor, + target_dtype, + scaling_granularity=ScalingGranularity.AXISWISE, + axiswise_dim=0, + round_scales_to_power_of_2=round_scales_to_power_of_2, + ) + + # Apply scales to subtensor and convert to float8. + tensor_scaled = subtensor.to(torch.float32) * subtensor_scales + float8_subtensor = to_fp8_saturated(tensor_scaled, target_dtype) + + # Store this portion of the resulting float8 tensor and scales. + A_fp8_col_major[start_idx:end_idx, :] = float8_subtensor + A_scales[next_scale_idx : next_scale_idx + subtensor_scales.numel()] = ( + subtensor_scales.squeeze() + ) + + # Update start index for next group. + start_idx = end_idx + next_scale_idx += subtensor_scales.numel() + + return A_fp8_col_major, A_scales + + +def _to_2d_jagged_float8_tensor_rowwise( + x: torch.Tensor, + offs: torch.Tensor, + target_dtype: torch.dtype, + round_scales_to_power_of_2: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + This function converts the 2D input tensor to a jagged float8 tensor, + with scales computed along *logical rows* for each group individually, + where groups are determined based on the offsets. + + For a 2D *left* operand of a normal scaled GEMM, the rowwise scales are computed over logical rows. + (i.e., a tensor of (M,K) will have scales of shape (M,1). + + However, for a 2D left operand of a grouped GEMM, these logical rows go through multiple distinct + groups/subtensors, for which we want to compute scales individually. So we cannot take one set of scales + along the logical rows and apply it to the entire tensor. + + Instead, we need to compute scales for each subtensor individually. For a tensor of shape (M,K) this results + in scales of shape (M * num_groups, 1). + + Args: + A (torch.Tensor): The input tensor to be converted to a jagged float8 tensor. + + Returns: + A tuple containing the jagged float8 tensor and the scales used for the conversion. + """ + assert x.ndim == 2, "input tensor must be 2D" + + num_groups = offs.numel() + x_fp8 = torch.empty_like(x, dtype=target_dtype) + x_scales = torch.empty( + x_fp8.size(0) * num_groups, dtype=torch.float32, device=x_fp8.device + ) + + start_idx = 0 + next_scale_idx = 0 + for end_idx in offs.tolist(): + # Get the subtensor of A for this group, fetching all rows with the next group of rows. + subtensor = x[:, start_idx:end_idx] # (M, local_group_size) + + # Compute local rowwise scales for this subtensor, which are along logical rows for the left operand. + subtensor_scales = tensor_to_scale( + subtensor, + target_dtype, + scaling_granularity=ScalingGranularity.AXISWISE, + axiswise_dim=-1, + round_scales_to_power_of_2=round_scales_to_power_of_2, + ) + + # Apply scales to subtensor and convert to float8. + tensor_scaled = subtensor.to(torch.float32) * subtensor_scales + float8_subtensor = to_fp8_saturated(tensor_scaled, target_dtype) + + # Store this portion of the resulting float8 tensor and scales. + x_fp8[:, start_idx:end_idx] = float8_subtensor + x_scales[next_scale_idx : next_scale_idx + subtensor_scales.numel()] = ( + subtensor_scales.squeeze() + ) + + # Update start index for next group. + start_idx = end_idx + next_scale_idx += subtensor_scales.numel() + + return x_fp8, x_scales + + +def _is_column_major(x: torch.Tensor) -> bool: + """ + This function checks if the input tensor is column-major. + + Args: + x (torch.Tensor): The input tensor to be checked. + + Returns: + A boolean indicating whether the input tensor is column-major. + """ + assert x.ndim == 2 or x.ndim == 3, "input tensor must be 2D or 3D" + return x.stride(-2) == 1 and x.stride(-1) > 1 diff --git a/lib/python3.12/site-packages/torchao/prototype/smoothquant/__init__.py b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..948a99c080b941ebe687f16dd99acf58df42b26b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__init__.py @@ -0,0 +1,15 @@ +from .api import ( + SmoothQuantConfig, + insert_smooth_quant_observer_, + load_smooth_quant_recipe, + save_smooth_quant_recipe, +) +from .core import SmoothQuantObservedLinear + +__all__ = [ + "insert_smooth_quant_observer_", + "load_smooth_quant_recipe", + "save_smooth_quant_recipe", + "SmoothQuantConfig", + "SmoothQuantObservedLinear", +] diff --git a/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5754cf521c021592cbcece997a930bf430ff05d4 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ca3da9309e6845fef1efaf479051951542cd1b5 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/core.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/core.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3950fc7749f1833c82e78b43c1795146a840d663 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/core.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/example.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fb7cea7a365ff2a2bbbb911f73d84369bc2660f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/smoothquant/__pycache__/example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/smoothquant/api.py b/lib/python3.12/site-packages/torchao/prototype/smoothquant/api.py new file mode 100644 index 0000000000000000000000000000000000000000..9397b340b36a4eeddfdc96fc055837f1daea41f7 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/smoothquant/api.py @@ -0,0 +1,231 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import types +from dataclasses import dataclass +from typing import Dict, Optional + +import torch + +import torchao +from torchao.core.config import AOBaseConfig +from torchao.dtypes import to_affine_quantized_intx, to_affine_quantized_intx_static +from torchao.prototype.smoothquant.core import ( + SmoothQuantObservedLinear, + SmoothQuantObserver, +) +from torchao.quantization import quantize_ +from torchao.quantization.linear_activation_quantized_tensor import ( + to_linear_activation_quantized, +) +from torchao.quantization.linear_activation_scale import ( + to_weight_tensor_with_linear_activation_scale_metadata, +) +from torchao.quantization.quant_api import ( + _linear_extra_repr, + _replace_with_custom_fn_if_matches_filter, +) +from torchao.quantization.quant_primitives import MappingType +from torchao.quantization.transform_module import ( + register_quantize_module_handler, +) +from torchao.quantization.utils import _get_per_token_block_size +from torchao.quantization.weight_tensor_linear_activation_quantization import ( + to_weight_tensor_with_linear_activation_quantization_metadata, +) + + +def insert_smooth_quant_observer_( + model: torch.nn.Module, alpha: Optional[float] = 0.5, quant_mode: str = "dynamic" +): + """ + Inserts SmoothQuantObserver into Linear layers of a given model. + + Args: + model: The model to be modified (in place). Ensure model is on the desired device for calibration + alpha: The alpha value to determine smoothing factor. Factor = 1 if alpha is None, which means + falling back to conventional quantization. + quant_mode: dynamic or static quantization of activation + """ + _is_linear = lambda m, fqn: isinstance(m, torch.nn.Linear) + + quant_min, quant_max = -127, 127 + eps = torch.finfo(torch.float32).eps + + def replace_with_observer(layer): + # creates observer and replaces linear layers with observed linear layers + observer = SmoothQuantObserver( + layer.weight, + alpha, + quant_mode, + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + ) + return SmoothQuantObservedLinear.from_float(layer, observer) + + _replace_with_custom_fn_if_matches_filter(model, replace_with_observer, _is_linear) + + +def save_smooth_quant_recipe( + model: torch.nn.Module, save_path: str +) -> Dict[str, torch.Tensor]: + """ + Save smoothing_factors, act_scales, and wei_scales for each SmoothQuantObservedLinear layer in the model. + """ + result = {} + + def recurse(module: torch.nn.Module, name: str = ""): + for child_name, child in module.named_children(): + full_name = f"{name}.{child_name}" if name else child_name + + # Apply the analysis function to this layer + if isinstance(child, SmoothQuantObservedLinear): + smoothing_factor, act_scales, wei_scales = child.obs.calculate_qparams() + result[full_name + ".smoothing_factor"] = smoothing_factor + result[full_name + ".act_scales"] = act_scales + result[full_name + ".wei_scales"] = wei_scales + + # Recurse into child modules + recurse(child, full_name) + + recurse(model) + + torch.save(result, save_path) + + +def load_smooth_quant_recipe( + model: torch.nn.Module, recipe_path: str, device=None +) -> torch.nn.Module: + recipe = torch.load(recipe_path, weights_only=True) + + def recurse(module: torch.nn.Module, name: str = ""): + if isinstance(module, SmoothQuantObservedLinear): + smoothing_factor = recipe.get(name + ".smoothing_factor", None) + act_scales = recipe.get(name + ".act_scales", None) + wei_scales = recipe.get(name + ".wei_scales", None) + if device is not None: + module.to(device=device) + # act_scales is None for dynamic quantization + if any(x is None for x in (smoothing_factor, wei_scales)): + return module + is_observed_linear = lambda m, fqn: isinstance(m, SmoothQuantObservedLinear) + wrapper = torch.nn.Sequential(module) + quantize_( + wrapper, + SmoothQuantConfig(smoothing_factor, act_scales, wei_scales), + is_observed_linear, + ) + return wrapper[0] + + mod_new = module + + for child_name, child in module.named_children(): + full_name = f"{name}.{child_name}" if name else child_name + setattr(mod_new, child_name, recurse(child, full_name)) + return mod_new + + recurse(model) + + +class _ActQuantizer: + def __init__(self, target_dtype, quant_min=-127): + self.target_dtype = target_dtype + self.quant_min = quant_min + + def dynamic_quantize(self, input): + return to_affine_quantized_intx( + input, + MappingType.SYMMETRIC, + _get_per_token_block_size(input), + self.target_dtype, + self.quant_min, + ) + + def static_quantize(self, input, scale, zero_point): + return to_affine_quantized_intx_static( + input, + scale, + zero_point, + list(input.shape), + self.target_dtype, + self.quant_min, + ) + + +@dataclass +class SmoothQuantConfig(AOBaseConfig): + """ + Configuration for quantizing linear layers when passed into quantize_() + + Args: + smoothing_factor: The smoothing factor for the layer. Acquired from the layer's observer if None. + act_scales: The activation scales for the layer. Acquired from the layer's observer if None. + wei_scales: The weight scales for the layer. Acquired from the layer's observer if None. + set_inductor_config: if True, adjusts `torchinductor` settings to recommended values. + """ + + smoothing_factor: Optional[torch.Tensor] = None + act_scales: Optional[torch.Tensor] = None + wei_scales: Optional[torch.Tensor] = None + set_inductor_config: bool = True + + +@register_quantize_module_handler(SmoothQuantConfig) +def _smooth_quant_transform( + module: torch.nn.Module, + config: SmoothQuantConfig, +): + smoothing_factor = config.smoothing_factor + act_scales = config.act_scales + wei_scales = config.wei_scales + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + observed_linear = module + + linear = torch.nn.Linear( + observed_linear.in_features, + observed_linear.out_features, + observed_linear.bias is not None, + device=observed_linear.weight.device, + dtype=observed_linear.weight.dtype, + ) + linear.bias = observed_linear.bias + + target_dtype = torch.int8 + # act_scales is None for dynamic quantization thus not checked + if any(x is None for x in (smoothing_factor, wei_scales)): + factor, x_scale, w_scales = observed_linear.obs.calculate_qparams() + weight = observed_linear.obs.weight * factor + else: + factor, x_scale, w_scales = smoothing_factor, act_scales, wei_scales + weight = observed_linear.weight * factor + weight = weight.to(observed_linear.weight.dtype) + block_size = (1, weight.size(1)) + wei_zero_points = torch.zeros_like(w_scales, dtype=torch.int64) + qw = to_affine_quantized_intx_static( + weight, + w_scales, + wei_zero_points, + block_size, + target_dtype, + ) + + if x_scale is None: + # dynamic quant + qw = to_linear_activation_quantized( + qw, _ActQuantizer(target_dtype).dynamic_quantize + ) + else: + # static quant + x_zero_point = torch.zeros_like(x_scale, dtype=torch.int64) + qw = to_weight_tensor_with_linear_activation_quantization_metadata( + qw, _ActQuantizer(target_dtype).static_quantize, x_scale, x_zero_point + ) + + qw = to_weight_tensor_with_linear_activation_scale_metadata(qw, factor.to(qw.dtype)) + linear.weight = torch.nn.Parameter(qw, requires_grad=False) + linear.extra_repr = types.MethodType(_linear_extra_repr, module) + return linear diff --git a/lib/python3.12/site-packages/torchao/prototype/smoothquant/core.py b/lib/python3.12/site-packages/torchao/prototype/smoothquant/core.py new file mode 100644 index 0000000000000000000000000000000000000000..3e6c6ea5d5e7ebb7a38662d86a3b8c43fe2448a7 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/smoothquant/core.py @@ -0,0 +1,162 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Optional + +import torch +import torch.nn.functional as F + +from torchao.quantization.observer import AffineQuantizedMinMaxObserver, PerAxis +from torchao.quantization.quant_primitives import ( + MappingType, +) + + +class SmoothQuantObserver(torch.nn.Module): + def __init__( + self, + weight: torch.Tensor, + alpha: Optional[float] = 0.5, + quant_mode: str = "static", # or dynamic + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + ): + """ + A custom observer for SmoothQuant + + Args: + weight: The weight tensor to be observed. + alpha: The alpha value to determine smoothing factor, normally between 0 and 1. + Fall back to conventional quantization if alpha is None. + quant_mode: The mode of activation quantization, either static or dynamic + quant_min: The minimum quantized value + quant_max: The maximum quantized value + eps: The minimum scale to avoid dividing by zero. + """ + super().__init__() + assert weight.ndim == 2 + self.weight = weight + self.inputs = [] + self.device = self.weight.device + self.alpha = alpha + assert quant_mode in ["static", "dynamic"] + self.quant_mode = quant_mode + self.quant_min = quant_min + self.quant_max = quant_max + self.eps = eps + # act.shape = [mb, ic] (reshape if needed), wei.shape = [oc, ic] + # *_ic_obs are used to determine smoothing_factor + # wei_oc_obs is used to find qparams for quantization + self.act_ic_obs = AffineQuantizedMinMaxObserver( + MappingType.SYMMETRIC, + torch.int8, + PerAxis(-1), + eps=eps, + ) + self.wei_ic_obs = AffineQuantizedMinMaxObserver( + MappingType.SYMMETRIC, + torch.int8, + PerAxis(-1), + eps=eps, + ) + self.wei_oc_obs = AffineQuantizedMinMaxObserver( + MappingType.SYMMETRIC, + torch.int8, + PerAxis(0), + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + ) + self.wei_ic_obs(self.weight) + + @torch.no_grad() + def forward(self, input: torch.Tensor): + self.act_ic_obs(input.to("cpu")) + return input + + def calculate_qparams(self): + # 1 Get min/max per IC from observers + wei_min_per_ic = self.wei_ic_obs.min_val + wei_max_per_ic = self.wei_ic_obs.max_val + act_min_per_ic = self.act_ic_obs.min_val + act_max_per_ic = self.act_ic_obs.max_val + x_abs_max_per_ic = ( + torch.max(torch.abs(act_min_per_ic), torch.abs(act_max_per_ic)) + self.eps + ) + w_abs_max_per_ic = ( + torch.max(torch.abs(wei_min_per_ic), torch.abs(wei_max_per_ic)) + self.eps + ) + # 2 calculate the smoothing factor + if self.alpha is None: + # fall back to conventional quantization if alpha is None + smoothing_factor = torch.ones_like( + x_abs_max_per_ic, + dtype=x_abs_max_per_ic.dtype, + device=x_abs_max_per_ic.device, + ) + else: + smoothing_factor = torch.pow(x_abs_max_per_ic, self.alpha) / torch.pow( + w_abs_max_per_ic.to(x_abs_max_per_ic.device), 1 - self.alpha + ) + # 3 apply smoothing factor to activations and find scales for static quantization + act_scales = None + if self.quant_mode == "static": + act_min_per_ic_new = act_min_per_ic / smoothing_factor.reshape( + act_min_per_ic.shape + ) + act_max_per_ic_new = act_max_per_ic / smoothing_factor.reshape( + act_max_per_ic.shape + ) + min_val_per_tensor = torch.min(act_min_per_ic_new) + max_val_per_tensor = torch.max(act_max_per_ic_new) + min_val_neg = torch.min( + min_val_per_tensor, torch.zeros_like(min_val_per_tensor) + ) + max_val_pos = torch.max( + max_val_per_tensor, torch.zeros_like(max_val_per_tensor) + ) + max_val_pos = torch.max(-min_val_neg, max_val_pos) + act_scale = max_val_pos / (float(self.quant_max - self.quant_min) / 2) + act_scales = act_scale.to(self.device) + # 4 update weight and find scales + self.wei_oc_obs(self.weight * smoothing_factor.to(self.device)) + wei_scales, _ = self.wei_oc_obs.calculate_qparams() + # 5 return results + return smoothing_factor.to(self.device), act_scales, wei_scales.to(self.device) + + +class SmoothQuantObservedLinear(torch.nn.Linear): + def __init__( + self, + in_features: int, + out_features: int, + bias: bool, + obs: SmoothQuantObserver, + device=None, + dtype=None, + ): + super().__init__(in_features, out_features, bias, device, dtype) + assert isinstance(obs, SmoothQuantObserver) + self.obs = obs + + def forward(self, input: torch.Tensor): + input = self.obs(input) + output = F.linear(input, self.weight, self.bias) + return output + + @classmethod + def from_float(cls, float_linear: torch.nn.Linear, obs: SmoothQuantObserver): + observed_linear = cls( + float_linear.in_features, + float_linear.out_features, + float_linear.bias is not None, + obs, + device=float_linear.weight.device, + dtype=float_linear.weight.dtype, + ) + observed_linear.weight = float_linear.weight + observed_linear.bias = float_linear.bias + return observed_linear diff --git a/lib/python3.12/site-packages/torchao/prototype/smoothquant/example.py b/lib/python3.12/site-packages/torchao/prototype/smoothquant/example.py new file mode 100644 index 0000000000000000000000000000000000000000..de1e4ed93ec09778d48f57e32743f895259d8ff0 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/smoothquant/example.py @@ -0,0 +1,247 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import argparse +import os +import time +from typing import Optional + +import torch +from datasets import load_dataset +from tqdm import tqdm +from transformers import AutoModelForCausalLM, AutoTokenizer + +from torchao.prototype.smoothquant import ( + SmoothQuantConfig, + SmoothQuantObservedLinear, + insert_smooth_quant_observer_, +) +from torchao.quantization import quantize_ + + +def get_calib_dataset(tokenizer=None, n_samples=100, block_size=512): + dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation") + samples = [] + n_tokens = n_samples * block_size + n_run = n_tokens + for data in dataset: + line = data["text"] + line = line.strip() + line_encoded = tokenizer.encode(line) + if len(line_encoded) > 512: + continue + sample = torch.tensor([line_encoded]) + if sample.numel() == 0: + continue + samples.append(sample) + n_run -= len(line_encoded) + if n_run <= n_samples: + break + + cat_samples = torch.cat(samples, dim=1) + return [ + cat_samples[:, i * block_size : (i + 1) * block_size] for i in range(n_samples) + ] + + +def wiki2_eval( + model, tokenizer, sequence_length, stride=512, verbose=True, device="cuda" +): + model.eval() + tokenizer.pad_token = tokenizer.eos_token + tokenizer.padding_side = "right" + tokenizer.add_eos_token = False + + print("Loading dataset") + t0 = time.time() + dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test") + encodings = tokenizer("\n\n".join(dataset["text"]), return_tensors="pt") + print(f"Time to load dataset: {time.time() - t0:.02f} seconds") + + encodings["input_ids"] = encodings["input_ids"].to(device) + + print("Running evaluation") + lls, t = [], [] + for i in tqdm( + range(0, encodings["input_ids"].size(1), stride), disable=not verbose + ): + begin_loc = max(i + stride - sequence_length, 0) + end_loc = min(i + stride, encodings["input_ids"].size(1)) + trg_len = end_loc - i + input_ids = encodings["input_ids"][:, begin_loc:end_loc] + target_ids = input_ids.clone() + target_ids[:, :-trg_len] = -100 # ignore context + + t1 = time.time() + with torch.no_grad(): + log_likelihood = model(input_ids, labels=target_ids).loss * trg_len + if device == "cuda": + torch.cuda.synchronize() + t2 = time.time() + t.append((t2 - t1)) + lls.append(log_likelihood) + + del input_ids, target_ids + + ppl = float(torch.exp(torch.stack(lls).sum() / end_loc)) + pred_time = sum(t) / len(t) + if verbose: + print("perplexity", ppl) + print("time", str(pred_time) + " sec/it") + + return {"perplexity": ppl, "prediction_time": pred_time} + + +def benchmark(model, tokenizer, max_length, tasks=None, device="cuda"): + model.eval() + model.config.use_cache = False + if tasks is None: + tasks = ["PPL"] + results = {} + if "PPL" in tasks: + results["perplexity"] = wiki2_eval( + model, tokenizer, 512, verbose=True, device=device + ) + return results + + +def wikitext2_ppl( + model_id: str, + alpha: Optional[float], + quant_mode: str, + calibration_size: int, + device: str, + precision: torch.dtype, + sequence_length: int, + compile: bool, + model_load_path: str, + model_save_path: str, +): + print(f"Loading model on {device}...") + torch.manual_seed(34) + t0 = time.time() + tokenizer = AutoTokenizer.from_pretrained(model_id) + if model_load_path is not None and os.path.exists(model_load_path): + print(f"Loading quantized model from {model_load_path}") + t0 = time.time() + model = torch.load(model_load_path, weights_only=False).to(device) + print(f"Time to load quantized model: {time.time() - t0:.02f} seconds") + else: + model = ( + AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=precision) + .eval() + .to(device) + ) + print(f"Time to load model: {time.time() - t0:.02f} seconds") + print("running calibration") + t0 = time.time() + # insert observers to find average magnitude and calculate scales + insert_smooth_quant_observer_(model, alpha, quant_mode) + calibration_data = get_calib_dataset( + tokenizer=tokenizer, n_samples=calibration_size, block_size=sequence_length + ) + for batch in calibration_data: + model(batch.to(device)) + batch.to("cpu") + print(f"time for calibration: {time.time() - t0:.02f} seconds") + + is_observed_linear = lambda m, fqn: isinstance(m, SmoothQuantObservedLinear) + print(f"running SmoothQuant with {quant_mode} quantization") + t0 = time.time() + quantize_(model, SmoothQuantConfig(), is_observed_linear) + print(f"time for quantization: {time.time() - t0:.02f} seconds") + if model_save_path is not None: + print(f"Saving quantized model to {model_save_path}") + t0 = time.time() + torch.save(model, model_save_path) + print(f"Time to save quantized model: {time.time() - t0:.02f} seconds") + if compile: + model = torch.compile(model, dynamic=True) + + return benchmark(model, tokenizer, sequence_length, tasks=["PPL"], device=device) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Evaluate a model with the specified parameters." + ) + + # Optional arguments with default values + parser.add_argument( + "--model-id", "-m", type=str, help="Repository ID of the model." + ) + parser.add_argument( + "--alpha", + type=float, + default=0.5, + help="The alpha hyperparameter for SmoothQuant.", + ) + parser.add_argument( + "--quant-mode", type=str, help="Quantization mode, either static or dynamic." + ) + parser.add_argument( + "--calibration-samples", + type=int, + default=10, + help="Number of samples to use for calibration. Default is 10.", + ) + parser.add_argument( + "--device", + type=str, + default="cuda", + help="Device to run the evaluation on. Default is 'cuda'.", + ) + parser.add_argument( + "--precision", + type=str, + default="bfloat16", + help="Precision type. Default is 'bfloat16'.", + ) + parser.add_argument( + "--seq_len", + type=int, + default=512, + help="Length of examples to calibrate and evaluate model on. Default is 512", + ) + parser.add_argument( + "--compile", + action="store_true", + help="Flag to indicate if compilation is required.", + ) + parser.add_argument( + "--model-load-path", + type=str, + default=None, + help="Path to load quantized model. If this is provided, " + "the model will be loaded from this path instead of quantizing the model.", + ) + parser.add_argument( + "--model-save-path", + type=str, + default=None, + help="Path to store quantized model.", + ) + parser.add_argument( + "--disable-smooth-quant", + action="store_true", + help="Run conventional dynamic or static quantization for testing or debugging.", + ) + + args = parser.parse_args() + + # Convert precision argument to torch dtype + precision_dtype = getattr(torch, args.precision, torch.bfloat16) + ppl = wikitext2_ppl( + args.model_id, + None if args.disable_smooth_quant else args.alpha, + args.quant_mode, + args.calibration_samples, + args.device, + args.precision, + args.seq_len, + args.compile, + args.model_load_path, + args.model_save_path, + ) diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/__init__.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..821e5049e0750e531ab95342770d9b308ad3965a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/__init__.py @@ -0,0 +1,33 @@ +# Sparsifier +# Scheduler +from torchao.prototype.sparsity.scheduler.base_scheduler import BaseScheduler +from torchao.prototype.sparsity.scheduler.cubic_scheduler import CubicSL +from torchao.prototype.sparsity.scheduler.lambda_scheduler import LambdaSL +from torchao.prototype.sparsity.sparsifier.base_sparsifier import BaseSparsifier +from torchao.prototype.sparsity.sparsifier.nearly_diagonal_sparsifier import ( + NearlyDiagonalSparsifier, +) + +# Parametrizations +from torchao.prototype.sparsity.sparsifier.utils import ( + FakeSparsity, + fqn_to_module, + get_arg_info_from_tensor_fqn, + module_to_fqn, +) +from torchao.prototype.sparsity.sparsifier.weight_norm_sparsifier import ( + WeightNormSparsifier, +) + +__all__ = [ + "BaseScheduler", + "CubicSL", + "LambdaSL", + "BaseSparsifier", + "NearlyDiagonalSparsifier", + "FakeSparsity", + "fqn_to_module", + "get_arg_info_from_tensor_fqn", + "module_to_fqn", + "WeightNormSparsifier", +] diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a29f648366d3ddf61e78cc66561e683d6f09720 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/FPGM_pruner.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/FPGM_pruner.py new file mode 100644 index 0000000000000000000000000000000000000000..ae46a8df1140df4d61956db21560b3741b3cb9cf --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/FPGM_pruner.py @@ -0,0 +1,98 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Callable, Optional, Union + +import torch + +from .base_structured_sparsifier import BaseStructuredSparsifier + +__all__ = ["FPGMPruner"] + + +class FPGMPruner(BaseStructuredSparsifier): + r"""Filter Pruning via Geometric Median (FPGM) Structured Pruner + This sparsifier prune fliter (row) in a tensor according to distances among filters according to + `Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration `_. + + This sparsifier is controlled by three variables: + 1. `sparsity_level` defines the number of filters (rows) that are zeroed-out. + 2. `dist` defines the distance measurement type. Default: 3 (L2 distance). + Available options are: [1, 2, (custom callable distance function)]. + + Note:: + Inputs should be a 4D convolutional tensor of shape (N, C, H, W). + - N: output channels size + - C: input channels size + - H: height of kernel + - W: width of kernel + """ + + def __init__( + self, sparsity_level: float = 0.5, dist: Optional[Union[Callable, int]] = None + ): + defaults = { + "sparsity_level": sparsity_level, + } + + if dist is None: + dist = 2 + + if callable(dist): + self.dist_fn = dist + elif dist == 1: + self.dist_fn = lambda x: torch.cdist(x, x, p=1) + elif dist == 2: + self.dist_fn = lambda x: torch.cdist(x, x, p=2) + else: + raise NotImplementedError("Distance function is not yet implemented.") + super().__init__(defaults=defaults) + + def _compute_distance(self, t): + r"""Compute distance across all entries in tensor `t` along all dimension + except for the one identified by dim. + Args: + t (torch.Tensor): tensor representing the parameter to prune + Returns: + distance (torch.Tensor): distance computed across filtters + """ + dim = 0 # prune filter (row) + + size = t.size(dim) + slc = [slice(None)] * t.dim() + + # flatten the tensor along the dimension + t_flatten = [ + t[tuple(slc[:dim] + [slice(i, i + 1)] + slc[dim + 1 :])].reshape(-1) + for i in range(size) + ] + t_flatten = torch.stack(t_flatten) + + # distance measurement + dist_matrix = self.dist_fn(t_flatten) + + # more similar with other filter indicates large in the sum of row + distance = torch.sum(torch.abs(dist_matrix), 1) + + return distance + + def update_mask(self, module, tensor_name, sparsity_level, **kwargs): + tensor_weight = getattr(module, tensor_name) + mask = getattr(module.parametrizations, tensor_name)[0].mask + + if sparsity_level <= 0: + mask.data = torch.ones_like(mask).bool() + elif sparsity_level >= 1.0: + mask.data = torch.zeros_like(mask).bool() + else: + distance = self._compute_distance(tensor_weight) + + tensor_size = tensor_weight.shape[0] # prune filter (row) + nparams_toprune = round(sparsity_level * tensor_size) + nparams_toprune = min( + max(nparams_toprune, 0), tensor_size + ) # clamp to [0, tensor_size] + topk = torch.topk(distance, k=nparams_toprune, largest=False) + mask[topk.indices] = False diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__init__.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9d7f7753891a503419cb2be080f00e67bd62e012 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__init__.py @@ -0,0 +1,17 @@ +from .base_structured_sparsifier import BaseStructuredSparsifier +from .FPGM_pruner import FPGMPruner +from .lstm_saliency_pruner import LSTMSaliencyPruner +from .parametrization import ( + BiasHook, + FakeStructuredSparsity, +) +from .saliency_pruner import SaliencyPruner + +__all__ = [ + "BaseStructuredSparsifier", + "FPGMPruner", + "LSTMSaliencyPruner", + "BiasHook", + "FakeStructuredSparsity", + "SaliencyPruner", +] diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/FPGM_pruner.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/FPGM_pruner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b4c42a03edf57753e77da3ef0a69d89c8ecd9bb Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/FPGM_pruner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e385922105fd6ecedd69fa63470087ae11f4802 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/base_structured_sparsifier.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/base_structured_sparsifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..327ccac89a7eb673ecdffdf325731205b3c3f53e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/base_structured_sparsifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/lstm_saliency_pruner.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/lstm_saliency_pruner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1019b33a80a8f2649bcd129955669e1bca5ac163 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/lstm_saliency_pruner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/match_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/match_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21093ee0fe1aee25933d5519f9669a0d5b207cb9 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/match_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/parametrization.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/parametrization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c96379663999f4639bc59dce3e8ba15a78bf70e5 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/parametrization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/prune_functions.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/prune_functions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..503730e001ed1fcfff8764f3b93dd9af0df313ef Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/prune_functions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/saliency_pruner.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/saliency_pruner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4f1a17f06aa343f78afd99f7184e82da3ec7dfd Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/__pycache__/saliency_pruner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/base_structured_sparsifier.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/base_structured_sparsifier.py new file mode 100644 index 0000000000000000000000000000000000000000..c3c097690b00175f70a4beb0eed49e4036f90c78 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/base_structured_sparsifier.py @@ -0,0 +1,319 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from itertools import chain +from operator import getitem +from typing import Callable, Dict, Optional, Set, Tuple, Type, Union + +import torch +import torch.nn.functional as F +from torch import nn +from torch.fx import symbolic_trace +from torch.nn.utils import parametrize + +from torchao.prototype.sparsity import BaseSparsifier + +from .match_utils import MatchAllNode, apply_match +from .parametrization import BiasHook, FakeStructuredSparsity, module_contains_param +from .prune_functions import ( + prune_conv2d, + prune_conv2d_activation_conv2d, + prune_conv2d_activation_pool_conv2d, + prune_conv2d_conv2d, + prune_conv2d_pool_activation_conv2d, + prune_conv2d_pool_flatten_linear, + prune_linear, + prune_linear_activation_linear, + prune_linear_linear, + prune_lstm_output_layernorm_linear, + prune_lstm_output_linear, +) + + +def _get_supported_structured_pruning_modules(): + SUPPORTED_STRUCTURED_PRUNING_MODULES = { # added to config if None given + nn.Linear, + nn.Conv2d, + nn.LSTM, + } + return SUPPORTED_STRUCTURED_PRUNING_MODULES + + +def _get_supported_activation_functions(): + SUPPORTED_ACTIVATION_FUNCTIONS = { + F.relu, + F.rrelu, + F.hardtanh, + F.relu6, + F.sigmoid, + F.hardsigmoid, + F.tanh, + F.silu, + F.mish, + F.hardswish, + F.elu, + F.celu, + F.selu, + F.hardshrink, + F.leaky_relu, + F.logsigmoid, + F.softplus, + F.prelu, + F.softsign, + F.tanhshrink, + F.gelu, + F.dropout, + } + return SUPPORTED_ACTIVATION_FUNCTIONS + + +def _get_supported_activation_modules(): + SUPPORTED_ACTIVATION_MODULES = { + nn.ReLU, + nn.RReLU, + nn.Hardtanh, + nn.ReLU6, + nn.Sigmoid, + nn.Hardsigmoid, + nn.Tanh, + nn.SiLU, + nn.Mish, + nn.Hardswish, + nn.ELU, + nn.CELU, + nn.SELU, + nn.Hardshrink, + nn.LeakyReLU, + nn.LogSigmoid, + nn.Softplus, + nn.PReLU, + nn.Softsign, + nn.Tanhshrink, + nn.GELU, + nn.Dropout, + } + return SUPPORTED_ACTIVATION_MODULES + + +def _get_default_structured_pruning_patterns() -> Dict[ + Tuple[Union[Type[nn.Module], Callable, MatchAllNode, str], ...], + Callable[..., None], +]: + """ + Returns the patterns for conv2d / linear conversion for each element in the activation functions/modules defined above. + """ + patterns: Dict[ + Tuple[Union[Type[nn.Module], Callable, MatchAllNode, str], ...], + Callable[..., None], + ] = { + # linear -> linear + (nn.Linear, "output"): prune_linear, + (nn.Linear, nn.Linear): prune_linear_linear, + # conv2d -> conv2d + (nn.Conv2d, "output"): prune_conv2d, + (nn.Conv2d, nn.Conv2d): prune_conv2d_conv2d, + # TODO LSTM Structured pruning does not support returned state currently. + # Should find a way to explicitly match getitem(0) instead of getitem. + # This will also require changing the pruning function. + # lstm -> getitem(0) -> linear + (nn.LSTM, getitem, nn.Linear): prune_lstm_output_linear, + # lstm -> getitem(0) -> layernorm -> linear + (nn.LSTM, getitem, nn.LayerNorm, nn.Linear): prune_lstm_output_layernorm_linear, + } + + for activation in chain( + _get_supported_activation_functions(), _get_supported_activation_modules() + ): + patterns.update( + { + # linear -> activation -> linear + (nn.Linear, activation, nn.Linear): prune_linear_activation_linear, + # conv2d -> activation -> conv2d + (nn.Conv2d, activation, nn.Conv2d): prune_conv2d_activation_conv2d, + # conv2d -> activation -> pool -> conv2d + ( + nn.Conv2d, + activation, + nn.AvgPool2d, + nn.Conv2d, + ): prune_conv2d_activation_pool_conv2d, + ( + nn.Conv2d, + activation, + F.avg_pool2d, + nn.Conv2d, + ): prune_conv2d_activation_pool_conv2d, + ( + nn.Conv2d, + activation, + nn.MaxPool2d, + nn.Conv2d, + ): prune_conv2d_activation_pool_conv2d, + ( + nn.Conv2d, + activation, + F.max_pool2d, + nn.Conv2d, + ): prune_conv2d_activation_pool_conv2d, + # conv2d -> pool -> activation -> conv2d + ( + nn.Conv2d, + nn.AvgPool2d, + activation, + nn.Conv2d, + ): prune_conv2d_pool_activation_conv2d, + ( + nn.Conv2d, + F.avg_pool2d, + activation, + nn.Conv2d, + ): prune_conv2d_pool_activation_conv2d, + ( + nn.Conv2d, + nn.MaxPool2d, + activation, + nn.Conv2d, + ): prune_conv2d_pool_activation_conv2d, + ( + nn.Conv2d, + F.max_pool2d, + activation, + nn.Conv2d, + ): prune_conv2d_pool_activation_conv2d, + # conv2d -> adaptive pool -> flatten -> linear + ( + nn.Conv2d, + nn.AdaptiveAvgPool2d, + nn.Flatten, + nn.Linear, + ): prune_conv2d_pool_flatten_linear, + ( + nn.Conv2d, + nn.AdaptiveAvgPool2d, + torch.flatten, + nn.Linear, + ): prune_conv2d_pool_flatten_linear, + ( + nn.Conv2d, + nn.AdaptiveMaxPool2d, + nn.Flatten, + nn.Linear, + ): prune_conv2d_pool_flatten_linear, + ( + nn.Conv2d, + nn.AdaptiveMaxPool2d, + torch.flatten, + nn.Linear, + ): prune_conv2d_pool_flatten_linear, + } + ) + return patterns + + +class BaseStructuredSparsifier(BaseSparsifier): + r"""Base class for structured pruning. + + Abstract methods that need to be implemented: + - update_mask: Function to compute a new mask for all keys in the + `groups` attribute. + + Args: + - defaults [dict]: default configurations will be attached to the + configuration. Only the keys that don't exist in the `config` will + be updated. + """ + + def __init__(self, defaults, patterns=None): + super().__init__(defaults) + if patterns is None: + patterns = _get_default_structured_pruning_patterns() + self.patterns = patterns + + def make_config_from_model( + self, + model: nn.Module, + SUPPORTED_MODULES: Optional[Set[Type]] = None, + ) -> None: + if SUPPORTED_MODULES is None: + SUPPORTED_MODULES = _get_supported_structured_pruning_modules() + super().make_config_from_model(model, SUPPORTED_MODULES=SUPPORTED_MODULES) + + def _prepare(self, *args, **kwargs) -> None: + r"""This function will attach the FakeStructuredSparsity parameterizations + and BiasHooks at the appropriate points in the model. + """ + for config in self.groups: + module = config["module"] + tensor_name = config["tensor_name"] + parametrization = config.get("parametrization", FakeStructuredSparsity) + tensor = getattr(module, tensor_name) + + mask = config.get( + "mask", + torch.ones(tensor.shape[0], dtype=torch.bool, device=tensor.device), + ) + self.state[config["tensor_fqn"]]["mask"] = mask + parametrize.register_parametrization( + module, tensor_name, parametrization(mask) + ) + + # if linear / conv, we add in bias hooks + if isinstance(module, (nn.Linear, nn.Conv2d)): + prune_bias = config.get("prune_bias", True) + if module.bias is not None: + module.register_parameter( + "_bias", nn.Parameter(module.bias.detach()) + ) + module.bias = None + module.prune_bias = prune_bias + + module.register_forward_hook( + BiasHook(module.parametrizations.weight[0], prune_bias) + ) + + def prune(self) -> None: + r""" + This function will FX symbolically trace the model and then find instances of the patterns + defined in self.patterns (by default SUPPORTED_STRUCTURED_PRUNING_PATTERNS ). + + For each pattern, it will apply to corresponding conversion function, which will modify the output + and input size expected by the modules within the pattern + """ + + self.traced = symbolic_trace(self.model) + modules = dict(self.traced.named_modules()) + + # Right now we check for matches simply by iterating across all the patterns + # if this is slow we can store patterns in a trie-structure and modify this code for faster lookup + for node in self.traced.graph.nodes: + for pattern, convert_fn in self.patterns.items(): + matched = apply_match(modules, pattern, node, []) + if matched is None: + continue + + first_module = modules.get(node.target) + # check if first module exists and has appropriate parameterization, otherwise skip + if ( + first_module is not None + and parametrize.is_parametrized(first_module) + and module_contains_param(first_module, FakeStructuredSparsity) + ): + convert_block = [] + for node in matched: + if node.op == "call_module": + convert_block.append(modules.get(node.target)) + elif node.op == "call_function": + convert_block.append(node.target) + convert_fn(*convert_block) + + for module in self.traced.modules(): + if module_contains_param(module, FakeStructuredSparsity): + raise Exception( + f"Error: {module} still contains FakeStructuredSparsity parametrizations!" + ) + + self.traced.graph.lint() + self.traced.recompile() + return self.traced diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/lstm_saliency_pruner.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/lstm_saliency_pruner.py new file mode 100644 index 0000000000000000000000000000000000000000..c61a00b8e14956fc22d25d00df59b188dd8493b7 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/lstm_saliency_pruner.py @@ -0,0 +1,57 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import cast + +import torch + +from .base_structured_sparsifier import BaseStructuredSparsifier, FakeStructuredSparsity + + +class LSTMSaliencyPruner(BaseStructuredSparsifier): + """ + Prune packed LSTM weights based on saliency. + For each layer {k} inside a LSTM, we have two packed weight matrices + - weight_ih_l{k} + - weight_hh_l{k} + + These tensors pack the weights for the 4 linear layers together for efficiency. + + [W_ii | W_if | W_ig | W_io] + + Pruning this tensor directly will lead to weights being misassigned when unpacked. + To ensure that each packed linear layer is pruned the same amount: + 1. We split the packed weight into the 4 constituent linear parts + 2. Update the mask for each individual piece using saliency individually + + This applies to both weight_ih_l{k} and weight_hh_l{k}. + """ + + def update_mask(self, module, tensor_name, **kwargs): + weights = getattr(module, tensor_name) + + for p in getattr(module.parametrizations, tensor_name): + if isinstance(p, FakeStructuredSparsity): + mask = cast(torch.Tensor, p.mask) + + # select weights based on magnitude + if weights.dim() <= 1: + raise Exception( + "Structured pruning can only be applied to a 2+dim weight tensor!" + ) + # take norm over all but first dim + dims = tuple(range(1, weights.dim())) + saliency = weights.norm(dim=dims, p=1) + + # handle weights in 4 groups + split_size = len(mask) // 4 + masks = torch.split(mask, split_size) + saliencies = torch.split(saliency, split_size) + + for keep_mask, sal in zip(masks, saliencies): + # mask smallest k values to be removed + k = int(len(keep_mask) * kwargs["sparsity_level"]) + prune = sal.topk(k, largest=False, sorted=False).indices + keep_mask.data[prune] = False # modifies underlying p.mask directly diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/match_utils.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/match_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..353ee02a2fa90884817d520ef33f97989487c2d9 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/match_utils.py @@ -0,0 +1,70 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +""" +Contains utility functions to check if a pattern is in the graph and return the matching nodes +""" + +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from torch import nn +from torch.ao.quantization.utils import ( + MatchAllNode, +) +from torch.fx import Node +from torch.nn.utils import parametrize + + +def _match( + modules: Dict[str, nn.ModuleDict], node: Node, current: Union[nn.Module, Any] +) -> bool: + r""" + checks to see if a single node of a pattern matches + """ + if isinstance(current, type) and issubclass(current, MatchAllNode): + return True + if not isinstance(node, Node): + return False + if isinstance(current, type) and issubclass(current, torch.nn.Module): + return ( + node.op == "call_module" + and parametrize.type_before_parametrizations(modules[node.target]) + == current + ) + elif callable(current): + return node.op == "call_function" and node.target is current + elif isinstance(current, str): + return node.target == current + return False + + +def apply_match( + modules: Dict[str, nn.ModuleDict], + pattern: Union[Tuple[Any], Any], + node: Node, + matched_node_pattern: List[Node], +) -> Optional[List[Node]]: + r""" + This function will return the matched nodes if the pattern matches the node given + If there is no match, it will return None + """ + if isinstance(pattern, tuple): + if len(pattern) == 1: + if _match(modules, node, pattern[0]): + return matched_node_pattern + [node] + + first, *rest = pattern + if _match(modules, node, first): + if rest is None: + return matched_node_pattern + [node] + + for user in node.users: + return apply_match( + modules, tuple(rest), user, matched_node_pattern + [node] + ) + elif _match(modules, node, pattern): + return [node] + return None diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/parametrization.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/parametrization.py new file mode 100644 index 0000000000000000000000000000000000000000..47d485009c09c749d95cc0fa26f0ee6147ae475e --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/parametrization.py @@ -0,0 +1,63 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +from torch import nn +from torch.nn.utils.parametrize import is_parametrized + + +def module_contains_param(module, parametrization): + if is_parametrized(module): + # see if any of the module tensors have a parametriztion attached that matches the one passed in + return any( + any(isinstance(param, parametrization) for param in param_list) + for key, param_list in module.parametrizations.items() + ) + return False + + +# Structured Pruning Parameterizations +class FakeStructuredSparsity(nn.Module): + r""" + Parametrization for Structured Pruning. Like FakeSparsity, this should be attached to + the 'weight' or any other parameter that requires a mask. + + Instead of an element-wise bool mask, this parameterization uses a row-wise bool mask. + """ + + def __init__(self, mask): + super().__init__() + self.register_buffer("mask", mask) + + def forward(self, x): + assert isinstance(self.mask, torch.Tensor) + assert self.mask.shape[0] == x.shape[0] + shape = [1] * len(x.shape) + shape[0] = -1 + return self.mask.reshape(shape) * x + + def state_dict(self, *args, **kwargs): + # avoid double saving masks + return {} + + +class BiasHook: + def __init__(self, parametrization, prune_bias): + self.param = parametrization + self.prune_bias = prune_bias + + def __call__(self, module, input, output): + if getattr(module, "_bias", None) is not None: + bias = module._bias.data + if self.prune_bias: + bias[~self.param.mask] = 0 + + # reshape bias to broadcast over output dimensions + idx = [1] * len(output.shape) + idx[1] = -1 + bias = bias.reshape(idx) + + output += bias + return output diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/prune_functions.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/prune_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..7197c30cc8b45395e29c82e5ebd312fc87af7c3d --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/prune_functions.py @@ -0,0 +1,483 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +""" +Collection of conversion functions for linear / conv2d structured pruning +Also contains utilities for bias propagation +""" + +from typing import Callable, List, Optional, Tuple, cast + +import torch +from torch import Tensor, nn +from torch.nn.utils import parametrize +from torch.nn.utils.parametrize import ParametrizationList + +from .parametrization import BiasHook, FakeStructuredSparsity + + +# BIAS PROPAGATION +def _remove_bias_handles(module: nn.Module) -> None: + if hasattr(module, "_forward_hooks"): + bias_hooks: List[int] = [] + for key, hook in module._forward_hooks.items(): + if isinstance(hook, BiasHook): + bias_hooks.append(key) + + for key in bias_hooks: + del module._forward_hooks[key] + + +def _get_adjusted_next_layer_bias( + next_layer: nn.Module, pruned_biases: Tensor, mask: Tensor +) -> nn.Parameter: + r"""Returns new adjusted bias for the second supported module""" + if parametrize.is_parametrized(next_layer): + # need to access original weight + parametrization_dict = cast(nn.ModuleDict, next_layer.parametrizations) + weight_parameterizations = cast( + ParametrizationList, parametrization_dict.weight + ) + next_weight = weight_parameterizations.original + else: + next_weight = cast(Tensor, next_layer.weight) + + scaling_weight = next_weight[:, ~mask] + if isinstance(next_layer, nn.Conv2d): # checking for Conv2d + # Propagating first layer pruned biases and calculating the new second layer bias + # involves more steps since the Conv2d scaling weight has extra dimensions, + # so adding bias involves broadcasting, logically: + # for each channel k in range(oC): + # scaled_biases = sum(first_bias[pruned_idx] @ next_weight[k, pruned_idx, :, :].T) + # new_next_bias[k] = old_next_bias[k] + scaled_biases + scaling_product = torch.matmul( + pruned_biases.reshape(1, -1), torch.transpose(scaling_weight, 1, 2) + ) + sum_range = list(range(len(scaling_product.shape)))[ + 1: + ] # all but the first dimension + scaled_biases = torch.sum(scaling_product, sum_range) + elif isinstance(next_layer, nn.Linear): # Linear + scaled_biases = torch.matmul( + pruned_biases, torch.transpose(scaling_weight, 0, 1) + ) # recall b2_new = b1 @ w2.T + b2 + else: + raise NotImplementedError(f"Type {type(next_layer)} not supported yet.") + + if ( + parametrize.is_parametrized(next_layer) + and getattr(next_layer, "_bias", None) is not None + ): # next_layer is parametrized & has original bias ._bias + adjusted_bias = nn.Parameter(scaled_biases + next_layer._bias) + elif ( + not parametrize.is_parametrized(next_layer) and next_layer.bias is not None + ): # next_layer not parametrized & has .bias + adjusted_bias = nn.Parameter(scaled_biases + next_layer.bias) + else: # next_layer has no bias + adjusted_bias = nn.Parameter(scaled_biases) + return adjusted_bias + + +def _prune_module_bias(module: nn.Module, mask: Tensor) -> None: + r"""Applies mask to given modules bias""" + # prune bias along with weights, discard pruned indices of bias + original_bias = cast(Tensor, getattr(module, "_bias", module.bias)) + if original_bias is not None: + module.bias = nn.Parameter(original_bias[mask]) + + # remove _bias parameter + if hasattr(module, "_bias"): + delattr(module, "_bias") + + +def _propogate_module_bias(module: nn.Module, mask: Tensor) -> Optional[Tensor]: + r""" + In the case that we need to propagate biases, this function will return the biases we need + """ + # set current module bias + if module.bias is not None: + module.bias = nn.Parameter(cast(Tensor, module.bias)[mask]) + elif getattr(module, "_bias", None) is not None: + module.bias = nn.Parameter(cast(Tensor, module._bias)[mask]) + + # get pruned biases to propagate to subsequent layer + if getattr(module, "_bias", None) is not None: + pruned_biases = cast(Tensor, module._bias)[~mask] + else: + pruned_biases = None + + if hasattr(module, "_bias"): + delattr(module, "_bias") + + return pruned_biases + + +# LINEAR +def _prune_linear_helper(linear: nn.Linear) -> Tensor: + # expects linear to be a parameterized linear module + parametrization_dict = cast(nn.ModuleDict, linear.parametrizations) + weight_parameterizations = cast(ParametrizationList, parametrization_dict.weight) + for p in weight_parameterizations: + if isinstance(p, FakeStructuredSparsity): + mask = cast(Tensor, p.mask) + + with torch.no_grad(): + parametrize.remove_parametrizations(linear, "weight", leave_parametrized=True) + linear.weight = nn.Parameter(linear.weight[mask]) # type: ignore[possibly-undefined] + linear.out_features = linear.weight.shape[0] + _remove_bias_handles(linear) + + return mask + + +def prune_linear(linear: nn.Linear) -> None: + mask = _prune_linear_helper(linear) + if getattr(linear, "prune_bias", False): + _prune_module_bias(linear, mask) + + +def prune_linear_linear(linear1: nn.Linear, linear2: nn.Linear) -> None: + prune_linear_activation_linear(linear1, None, linear2) + + +def prune_linear_activation_linear( + linear1: nn.Linear, + activation: Optional[Callable[[Tensor], Tensor]], + linear2: nn.Linear, +): + mask = _prune_linear_helper(linear1) + if getattr(linear1, "prune_bias", False): + _prune_module_bias(linear1, mask) + else: + pruned_biases = _propogate_module_bias(linear1, mask) + if pruned_biases is not None: + if activation: + pruned_biases = activation(pruned_biases) + linear2.bias = _get_adjusted_next_layer_bias(linear2, pruned_biases, mask) + + with torch.no_grad(): + if parametrize.is_parametrized(linear2): + parametrization_dict = cast(nn.ModuleDict, linear2.parametrizations) + weight_parameterizations = cast( + ParametrizationList, parametrization_dict.weight + ) + + weight_parameterizations.original = nn.Parameter( + weight_parameterizations.original[:, mask] + ) + linear2.in_features = weight_parameterizations.original.shape[1] + else: + linear2.weight = nn.Parameter(linear2.weight[:, mask]) + linear2.in_features = linear2.weight.shape[1] + + +# CONV2D +def _prune_conv2d_helper(conv2d: nn.Conv2d) -> Tensor: + parametrization_dict = cast(nn.ModuleDict, conv2d.parametrizations) + weight_parameterizations = cast(ParametrizationList, parametrization_dict.weight) + for p in weight_parameterizations: + if isinstance(p, FakeStructuredSparsity): + mask = cast(Tensor, p.mask) + + with torch.no_grad(): + parametrize.remove_parametrizations(conv2d, "weight", leave_parametrized=True) + conv2d.weight = nn.Parameter(conv2d.weight[mask]) # type: ignore[possibly-undefined] + conv2d.out_channels = conv2d.weight.shape[0] + + _remove_bias_handles(conv2d) + return mask + + +def prune_conv2d_padded(conv2d_1: nn.Conv2d) -> None: + parametrization_dict = cast(nn.ModuleDict, conv2d_1.parametrizations) + weight_parameterizations = cast(ParametrizationList, parametrization_dict.weight) + for p in weight_parameterizations: + if isinstance(p, FakeStructuredSparsity): + mask = cast(Tensor, p.mask) + + with torch.no_grad(): + parametrize.remove_parametrizations(conv2d_1, "weight", leave_parametrized=True) + + if getattr(conv2d_1, "_bias", None) is not None: + if ( + conv2d_1.bias is not None + ): # conv2d_1 has original bias and bias propagated from previous layer + new_bias = torch.zeros(conv2d_1.bias.shape) + new_bias[mask] = conv2d_1.bias[mask] # type: ignore[possibly-undefined] + # adjusted bias that to keep in conv2d_1 + new_bias[~mask] = cast(Tensor, conv2d_1._bias)[~mask] + # pruned biases that are kept instead of propagated + conv2d_1.bias = nn.Parameter(new_bias) + else: # conv2d_1 has only original bias + conv2d_1.bias = nn.Parameter(cast(Tensor, conv2d_1._bias)) + else: + # no original bias, only propagated bias + if ( + conv2d_1.bias is not None + ): # conv2d_1 has bias propagated from previous layer + conv2d_1.bias.data[~mask] = 0 # type: ignore[possibly-undefined] + + if hasattr(conv2d_1, "_bias"): + delattr(conv2d_1, "_bias") + + +def prune_conv2d(conv2d: nn.Conv2d) -> None: + mask = _prune_conv2d_helper(conv2d) + if getattr(conv2d, "prune_bias", False): + _prune_module_bias(conv2d, mask) + + +def prune_conv2d_conv2d(conv2d_1: nn.Conv2d, conv2d_2: nn.Conv2d) -> None: + prune_conv2d_activation_conv2d(conv2d_1, None, conv2d_2) + + +def prune_conv2d_activation_conv2d( + conv2d_1: nn.Conv2d, + activation: Optional[Callable[[Tensor], Tensor]], + conv2d_2: nn.Conv2d, +): + r""" + Fusion Pattern for conv2d -> some activation module / function -> conv2d layers + """ + parametrization_dict = cast(nn.ModuleDict, conv2d_1.parametrizations) + weight_parameterizations = cast(ParametrizationList, parametrization_dict.weight) + for p in weight_parameterizations: + if isinstance(p, FakeStructuredSparsity): + mask = cast(Tensor, p.mask) + + prune_bias = getattr(conv2d_1, "prune_bias", False) + if ( + hasattr(conv2d_2, "padding") + and cast(Tuple[int], conv2d_2.padding) > (0, 0) + and (conv2d_1.bias is not None or getattr(conv2d_1, "_bias", None) is not None) + ): + prune_conv2d_padded(conv2d_1) + else: + mask = _prune_conv2d_helper(conv2d_1) + if prune_bias: + _prune_module_bias(conv2d_1, mask) + else: + pruned_biases = _propogate_module_bias(conv2d_1, mask) + if pruned_biases is not None: + if activation: + pruned_biases = activation(pruned_biases) + conv2d_2.bias = _get_adjusted_next_layer_bias( + conv2d_2, pruned_biases, mask + ) + + if ( + not ( + hasattr(conv2d_2, "padding") + and cast(Tuple[int], conv2d_2.padding) > (0, 0) + ) + or conv2d_1.bias is None + ): + with torch.no_grad(): + if parametrize.is_parametrized(conv2d_2): + parametrization_dict = cast( + nn.ModuleDict, conv2d_2.parametrizations + ) + weight_parameterizations = cast( + ParametrizationList, parametrization_dict.weight + ) + weight_parameterizations.original = nn.Parameter( + weight_parameterizations.original[:, mask] + ) + conv2d_2.in_channels = weight_parameterizations.original.shape[1] + else: + conv2d_2.weight = nn.Parameter(conv2d_2.weight[:, mask]) + conv2d_2.in_channels = conv2d_2.weight.shape[1] + + +def prune_conv2d_pool_activation_conv2d( + c1: nn.Conv2d, + pool: nn.Module, + activation: Optional[Callable[[Tensor], Tensor]], + c2: nn.Conv2d, +) -> None: + prune_conv2d_activation_conv2d(c1, activation, c2) + + +def prune_conv2d_activation_pool_conv2d( + c1: nn.Conv2d, + activation: Optional[Callable[[Tensor], Tensor]], + pool: nn.Module, + c2: nn.Conv2d, +) -> None: + prune_conv2d_activation_conv2d(c1, activation, c2) + + +def prune_conv2d_pool_flatten_linear( + conv2d: nn.Conv2d, + pool: nn.Module, + flatten: Optional[Callable[[Tensor], Tensor]], + linear: nn.Linear, +) -> None: + mask = _prune_conv2d_helper(conv2d) + + # We map the pruned indices of the Conv2d output to the flattened indices of the Linear following the Flatten layer. + # we determine the flattening scale (h * w), and readjust `first_pruned_indices` + # (each idx maps to range idx * h * w to (idx+1) * h * w), `first_valid_indices`, + # and `pruned_biases` (repeat each bias by h * w). + if parametrize.is_parametrized(linear): + parametrization_dict = cast(nn.ModuleDict, linear.parametrizations) + weight_parameterizations = cast( + ParametrizationList, parametrization_dict.weight + ) + linear_ic = weight_parameterizations.original.shape[1] + else: + linear_ic = linear.weight.shape[1] + + conv2d_oc = len(mask) + assert linear_ic % conv2d_oc == 0, ( + f"Flattening from dimensions {conv2d_oc} to {linear_ic} not supported" + ) + + flatten_scale = linear_ic // conv2d_oc + flattened_mask = torch.tensor( + [[val] * flatten_scale for val in mask], dtype=torch.bool, device=mask.device + ).flatten() + + if getattr(conv2d, "prune_bias", False): + _prune_module_bias(conv2d, mask) + else: + pruned_biases = cast(Tensor, _propogate_module_bias(conv2d, mask)) + flattened_pruned_biases = torch.tensor( + [[bias] * flatten_scale for bias in pruned_biases], device=mask.device + ).flatten() + linear.bias = _get_adjusted_next_layer_bias( + linear, flattened_pruned_biases, flattened_mask + ) + + with torch.no_grad(): + if parametrize.is_parametrized(linear): + parametrization_dict = cast(nn.ModuleDict, linear.parametrizations) + weight_parameterizations = cast( + ParametrizationList, parametrization_dict.weight + ) + weight_parameterizations.original = nn.Parameter( + weight_parameterizations.original[:, flattened_mask] + ) + linear.in_features = weight_parameterizations.original.shape[1] + else: + linear.weight = nn.Parameter(linear.weight[:, flattened_mask]) + linear.in_features = linear.weight.shape[1] + + +def prune_lstm_output_linear( + lstm: nn.LSTM, getitem: Callable, linear: nn.Linear +) -> None: + prune_lstm_output_layernorm_linear(lstm, getitem, None, linear) + + +def prune_lstm_output_layernorm_linear( + lstm: nn.LSTM, + getitem: Callable, + layernorm: Optional[nn.LayerNorm], + linear: nn.Linear, +) -> None: + for i in range(lstm.num_layers): + if parametrize.is_parametrized(lstm, f"weight_ih_l{i}"): + parametrization_dict = cast(nn.ModuleDict, lstm.parametrizations) + weight_parameterizations = cast( + ParametrizationList, parametrization_dict[f"weight_ih_l{i}"] + ) + mask = weight_parameterizations[0].mask + + with torch.no_grad(): + parametrize.remove_parametrizations( + lstm, f"weight_ih_l{i}", leave_parametrized=True + ) + setattr( + lstm, + f"weight_ih_l{i}", + nn.Parameter(getattr(lstm, f"weight_ih_l{i}")[mask]), + ) + setattr( + lstm, + f"bias_ih_l{i}", + nn.Parameter(getattr(lstm, f"bias_ih_l{i}")[mask]), + ) + + if parametrize.is_parametrized(lstm, f"weight_hh_l{i}"): + parametrization_dict = cast(nn.ModuleDict, lstm.parametrizations) + weight_parameterizations = cast( + ParametrizationList, parametrization_dict[f"weight_hh_l{i}"] + ) + mask = weight_parameterizations[0].mask + + with torch.no_grad(): + parametrize.remove_parametrizations( + lstm, f"weight_hh_l{i}", leave_parametrized=True + ) + # splitting out hidden-hidden masks + W_hi, W_hf, W_hg, W_ho = torch.split( + getattr(lstm, f"weight_hh_l{i}"), lstm.hidden_size + ) + M_hi, M_hf, M_hg, M_ho = torch.split(mask, lstm.hidden_size) + + # resize each individual weight separately + W_hi = W_hi[M_hi][:, M_hi] + W_hf = W_hf[M_hf][:, M_hf] + W_hg = W_hg[M_hg][:, M_hg] + W_ho = W_ho[M_ho][:, M_ho] + + # concat, use this as new weight + new_weight = torch.cat((W_hi, W_hf, W_hg, W_ho)) + setattr(lstm, f"weight_hh_l{i}", nn.Parameter(new_weight)) + setattr( + lstm, + f"bias_hh_l{i}", + nn.Parameter(getattr(lstm, f"bias_hh_l{i}")[mask]), + ) + + # If this is the final layer, then we need to prune linear layer columns + if i + 1 == lstm.num_layers: + lstm.hidden_size = int(M_hi.sum()) + with torch.no_grad(): + if parametrize.is_parametrized(linear): + parametrization_dict = cast( + nn.ModuleDict, linear.parametrizations + ) + weight_parameterizations = cast( + ParametrizationList, parametrization_dict.weight + ) + + weight_parameterizations.original = nn.Parameter( + weight_parameterizations.original[:, M_ho] + ) + linear.in_features = weight_parameterizations.original.shape[1] + else: + linear.weight = nn.Parameter(linear.weight[:, M_ho]) + linear.in_features = linear.weight.shape[1] + + # if layernorm module, prune weight and bias + if layernorm is not None: + layernorm.normalized_shape = (linear.in_features,) + layernorm.weight = nn.Parameter(layernorm.weight[M_ho]) + layernorm.bias = nn.Parameter(layernorm.bias[M_ho]) + + # otherwise need to prune the columns of the input of the next LSTM layer + else: + with torch.no_grad(): + if parametrize.is_parametrized(lstm, f"weight_ih_l{i + 1}"): + parametrization_dict = cast( + nn.ModuleDict, lstm.parametrizations + ) + weight_parameterizations = cast( + ParametrizationList, + getattr(parametrization_dict, f"weight_ih_l{i + 1}"), + ) + + weight_parameterizations.original = nn.Parameter( + weight_parameterizations.original[:, M_ho] + ) + else: + next_layer_weight = getattr(lstm, f"weight_ih_l{i + 1}") + setattr( + lstm, + f"weight_ih_l{i + 1}", + nn.Parameter(next_layer_weight[:, M_ho]), + ) diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/saliency_pruner.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/saliency_pruner.py new file mode 100644 index 0000000000000000000000000000000000000000..0c0af152fa9d956e6abdccdc704b1c15ad508da5 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/pruner/saliency_pruner.py @@ -0,0 +1,36 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from .base_structured_sparsifier import BaseStructuredSparsifier + + +class SaliencyPruner(BaseStructuredSparsifier): + """ + Prune rows based on the saliency (L1 norm) of each row. + + This pruner works on N-Dimensional weight tensors. + For each row, we will calculate the saliency, whic is the sum the L1 norm of all weights in that row. + We expect that the resulting saliency vector has the same shape as our mask. + We then pick elements to remove until we reach the target sparsity_level. + """ + + def update_mask(self, module, tensor_name, **kwargs): + # tensor_name will give you the FQN, all other entries in sparse config is present in kwargs + weights = getattr(module, tensor_name) + mask = getattr(module.parametrizations, tensor_name)[0].mask + + # use negative weights so we can use topk (we prune out the smallest) + if weights.dim() <= 1: + raise Exception( + "Structured pruning can only be applied to a 2+dim weight tensor!" + ) + saliency = -weights.norm(dim=tuple(range(1, weights.dim())), p=1) + assert saliency.shape == mask.shape + + num_to_pick = int(len(mask) * kwargs["sparsity_level"]) + prune = saliency.topk(num_to_pick).indices + + # Set the mask to be false for the rows we want to prune + mask.data[prune] = False diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__init__.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2fffde9938a51e07546f03469a31215176950a22 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/base_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/base_scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d038f7fdc9cc86cf41cdf124a6be071abddf38c9 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/base_scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/cubic_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/cubic_scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85de8715d4c745faac5841a5e1fe530f49a8578c Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/cubic_scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/lambda_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/lambda_scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..632f0703200c7e74e56a7e5b0ca83df03dd6bec1 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/__pycache__/lambda_scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/base_scheduler.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/base_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..0c500482ee6a82767e2d8f7822616e50ec535d15 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/base_scheduler.py @@ -0,0 +1,172 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import warnings +import weakref +from functools import wraps + +from torchao.prototype.sparsity.sparsifier.base_sparsifier import BaseSparsifier + +__all__ = ["BaseScheduler"] + + +class BaseScheduler: + def __init__(self, sparsifier, last_epoch=-1, verbose=False): + # Attach sparsifier + if not isinstance(sparsifier, BaseSparsifier): + raise TypeError( + f"{type(sparsifier).__name__} is not an instance of torch.ao.pruning.BaseSparsifier" + ) + self.sparsifier = sparsifier + + # Initialize epoch and base sparsity levels + + self.base_sl = [group["sparsity_level"] for group in sparsifier.groups] + self.last_epoch = last_epoch + + # Following https://github.com/pytorch/pytorch/issues/20124 + # We would like to ensure that `scheduler.step()` is called after + # `sparsifier.step()` + def with_counter(method): + if getattr(method, "_with_counter", False): + # `sparsifier.step()` has already been replaced, return. + return method + + # Keep a weak reference to the sparsifier instance to prevent + # cyclic references. + instance_ref = weakref.ref(method.__self__) + # Get the unbound method for the same purpose. + func = method.__func__ + cls = instance_ref().__class__ + del method + + @wraps(func) + def wrapper(*args, **kwargs): + instance = instance_ref() + instance._step_count += 1 # type: ignore[union-attr] + wrapped = func.__get__(instance, cls) + return wrapped(*args, **kwargs) + + # Note that the returned function here is no longer a bound method, + # so attributes like `__func__` and `__self__` no longer exist. + wrapper._with_counter = True # type: ignore[attr-defined] + return wrapper + + self.sparsifier.step = with_counter(self.sparsifier.step) # type: ignore[assignment] + self.sparsifier._step_count = 0 # type: ignore[attr-defined] + self._step_count: int = 0 + self.verbose = verbose + + # Housekeeping + self._get_sl_called_within_step: bool = False + + self.step() + + def state_dict(self): + """Returns the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in self.__dict__ which + is not the sparsifier. + """ + return { + key: value for key, value in self.__dict__.items() if key != "sparsifier" + } + + def load_state_dict(self, state_dict): + """Loads the schedulers state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + self.__dict__.update(state_dict) + + def get_last_sl(self): + """Return last computed sparsity level by current scheduler.""" + return self._last_sl + + def get_sl(self): + # Compute sparsity level using chainable form of the scheduler + # Note: This method is not intended to be called directly, and is only + # used by the ".step" method. Use .get_last_sl() instead. + if not self._get_sl_called_within_step: + warnings.warn( + "To get the last sparsity level computed by the scheduler, " + "please use `get_last_sl()`." + ) + raise NotImplementedError + + def print_sl(self, is_verbose, group, sl, epoch=None): + """Display the current sparsity level.""" + if is_verbose: + if epoch is None: + print(f"Adjusting sparsity level of group {group} to {sl:.4e}.") + else: + print( + f"Epoch {epoch:5d}: adjusting sparsity level of group {group} to {sl:.4e}." + ) + + def __repr__(self): + format_string = self.__class__.__name__ + " (" + format_string += "\n" + format_string += f"Sparsifier {self.sparsifier}\n" + format_string += f" base_sl: {self.base_sl}\n" + format_string += ")" + return format_string + + def step(self, epoch=None): + # Raise warning if trying to call scheduler step before the sparsifier. + # https://github.com/pytorch/pytorch/issues/20124 + if self._step_count == 1: + if not hasattr(self.sparsifier.step, "_with_counter"): + warnings.warn( + "Seems like `sparsifier.step()` has been overridden after sparsity scheduler " + "initialization. Please, make sure to call `sparsifier.step()` before " + "`scheduler.step()`.", + UserWarning, + ) + + # Just check if there were two first scheduler.step() calls before sparsifier.step() + elif self.sparsifier._step_count < 1: # type: ignore[attr-defined] + warnings.warn( + "Detected call of `scheduler.step()` before `sparsifier.step()`. " + "You have to make sure you run the sparsifier.step() BEFORE any " + "calls to the scheduler.step().", + UserWarning, + ) + self._step_count += 1 + + class _enable_get_sl_call: + def __init__(self, o): + self.o = o + + def __enter__(self): + self.o._get_sl_called_within_step = True + return self + + def __exit__(self, type, value, traceback): + self.o._get_sl_called_within_step = False + + with _enable_get_sl_call(self): + self.last_epoch += 1 + values = self.get_sl() + + for i, data in enumerate(zip(self.sparsifier.groups, values)): + param_group, sl = data + param_group["sparsity_level"] = sl + self.print_sl(self.verbose, i, sl, epoch) + + self._last_sl = [group["sparsity_level"] for group in self.sparsifier.groups] + self.sparsifier.enable_mask_update = True + + def _make_sure_a_list(self, var): + r"""Utility that extends it to the same length as the .groups, ensuring it is a list""" + n = len(self.sparsifier.groups) + if not isinstance(var, (list, tuple)): + return [var] * n + else: + if len(var) != n: + raise ValueError(f"Expected variable of length {n}, but got {len(var)}") + return list(var) # We want the result to be in a list, not tuple diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/cubic_scheduler.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/cubic_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..ea67db09a93978fe0ce70f479276c8b39548b446 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/cubic_scheduler.py @@ -0,0 +1,116 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import warnings + +from .base_scheduler import BaseScheduler + +__all__ = ["CubicSL"] + + +def _clamp(x, lo, hi): + return max(lo, min(hi, x)) + + +class CubicSL(BaseScheduler): + r"""Sets the sparsity level of each parameter group to the final sl + plus a given exponential function. + + .. math:: + + s_i = s_f + (s_0 - s_f) \cdot \left( 1 - \frac{t - t_0}{n\Delta t} \right)^3 + + where :math:`s_i` is the sparsity at epoch :math:`t`, :math;`s_f` is the final + sparsity level, :math:`f(i)` is the function to be applied to the current epoch + :math:`t`, initial epoch :math:`t_0`, and final epoch :math:`t_f`. + :math:`\Delta t` is used to control how often the update of the sparsity level + happens. By default, + + Args: + sparsifier (BaseSparsifier): Wrapped sparsifier. + init_sl (int, list): Initial level of sparsity + init_t (int, list): Initial step, when pruning starts + delta_t (int, list): Pruning frequency + total_t (int, list): Total number of pruning steps + initially_zero (bool, list): If True, sets the level of sparsity to 0 + before init_t (:math:`t_0`). Otherwise, the sparsity level before + init_t (:math:`t_0`) is set to init_sl(:math:`s_0`) + last_epoch (int): The index of last epoch. Default: -1. + verbose (bool): If ``True``, prints a message to stdout for + each update. Default: ``False``. + """ + + def __init__( + self, + sparsifier, + init_sl=0.0, + init_t=0, + delta_t=10, + total_t=100, + initially_zero=False, + last_epoch=-1, + verbose=False, + ): + self.sparsifier = sparsifier + + self.init_sl = self._make_sure_a_list(init_sl) + self.init_t = self._make_sure_a_list(init_t) + self.delta_t = self._make_sure_a_list(delta_t) + self.total_t = self._make_sure_a_list(total_t) + + self.initially_zero = self._make_sure_a_list(initially_zero) + + super().__init__(sparsifier, last_epoch, verbose) + + @staticmethod + def sparsity_compute_fn(s_0, s_f, t, t_0, dt, n, initially_zero=False): + r""" "Computes the current level of sparsity. + + Based on https://arxiv.org/pdf/1710.01878.pdf + + Args: + s_0: Initial level of sparsity, :math:`s_i` + s_f: Target level of sparsity, :math:`s_f` + t: Current step, :math:`t` + t_0: Initial step, :math:`t_0` + dt: Pruning frequency, :math:`\Delta T` + n: Pruning steps, :math:`n` + initially_zero: Sets the level of sparsity to 0 before t_0. + If False, sets to s_0 + + Returns: + The sparsity level :math:`s_t` at the current step :math:`t` + """ + if initially_zero and t < t_0: + return 0 + s_t = s_f + (s_0 - s_f) * (1.0 - (t - t_0) / (dt * n)) ** 3 + s_t = _clamp(s_t, s_0, s_f) + return s_t + + def get_sl(self): + if not self._get_sl_called_within_step: + warnings.warn( + "To get the last sparsity level computed by the scheduler, " + "please use `get_last_sl()`." + ) + return [ + self.sparsity_compute_fn( + s_0=initial_sparsity, + s_f=final_sparsity, + t=self.last_epoch, + t_0=initial_epoch, + dt=delta_epoch, + n=interval_epochs, + initially_zero=initially_zero, + ) + for initial_sparsity, final_sparsity, initial_epoch, delta_epoch, interval_epochs, initially_zero in zip( + self.init_sl, + self.base_sl, + self.init_t, + self.delta_t, + self.total_t, + self.initially_zero, + ) + ] diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/lambda_scheduler.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/lambda_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..d5bed639a459f11fee6165d41f34aa32cf3abae0 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/scheduler/lambda_scheduler.py @@ -0,0 +1,58 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import warnings + +from .base_scheduler import BaseScheduler + +__all__ = ["LambdaSL"] + + +class LambdaSL(BaseScheduler): + """Sets the sparsity level of each parameter group to the final sl + times a given function. When last_epoch=-1, sets initial sl as zero. + Args: + sparsifier (BaseSparsifier): Wrapped sparsifier. + sl_lambda (function or list): A function which computes a multiplicative + factor given an integer parameter epoch, or a list of such + functions, one for each group in sparsifier.param_groups. + last_epoch (int): The index of last epoch. Default: -1. + verbose (bool): If ``True``, prints a message to stdout for + each update. Default: ``False``. + Example: + >>> # Assuming sparsifier has two groups. + >>> lambda1 = lambda epoch: epoch // 30 + >>> lambda2 = lambda epoch: 0.95 ** epoch + >>> # xdoctest: +SKIP + >>> scheduler = LambdaSL(sparsifier, sl_lambda=[lambda1, lambda2]) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + """ + + def __init__(self, sparsifier, sl_lambda, last_epoch=-1, verbose=False): + self.sparsifier = sparsifier + + if not isinstance(sl_lambda, list) and not isinstance(sl_lambda, tuple): + self.sl_lambdas = [sl_lambda] * len(sparsifier.groups) + else: + if len(sl_lambda) != len(sparsifier.groups): + raise ValueError( + f"Expected {len(sparsifier.groups)} lr_lambdas, but got {len(sl_lambda)}" + ) + self.sl_lambdas = list(sl_lambda) + super().__init__(sparsifier, last_epoch, verbose) + + def get_sl(self): + if not self._get_sl_called_within_step: + warnings.warn( + "To get the last sparsity level computed by the scheduler, " + "please use `get_last_sl()`." + ) + return [ + base_sl * lmbda(self.last_epoch) + for lmbda, base_sl in zip(self.sl_lambdas, self.base_sl) + ] diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__init__.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f4466377874cab3f6c62846c791e4802535e460 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/base_sparsifier.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/base_sparsifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b079a1db25254ae9736eb19743d1f589585d9af Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/base_sparsifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/nearly_diagonal_sparsifier.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/nearly_diagonal_sparsifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e86d6db5470f3395d8962d4cb3f9062cdfd648f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/nearly_diagonal_sparsifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..639a0e29340ae3592cdc7aa261b17d4935722a82 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/weight_norm_sparsifier.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/weight_norm_sparsifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb8fa23fcb30b23b2fa80fe1268fcb86310270dd Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/__pycache__/weight_norm_sparsifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/base_sparsifier.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/base_sparsifier.py new file mode 100644 index 0000000000000000000000000000000000000000..035921e0f6263f0aa5f4360a5530b592e20a5353 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/base_sparsifier.py @@ -0,0 +1,358 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import abc +import copy +from collections import defaultdict +from typing import Any, Dict, List, Optional, Set, Tuple, Type + +import torch +from torch import nn +from torch.nn.utils import parametrize +from torch.nn.utils.parametrize import type_before_parametrizations + +from .utils import ( + FakeSparsity, + get_arg_info_from_tensor_fqn, + module_contains_param, + module_to_fqn, + swap_module, +) + +__all__ = ["BaseSparsifier"] + +SUPPORTED_MODULES = {nn.Linear} + +KEYS_NOT_IN_STATE_DICT = ["module", "module_fqn", "tensor_name"] + +__all__ = ["BaseSparsifier"] + + +# TODO update desc with new config args +class BaseSparsifier(abc.ABC): + r"""Base class for all sparsifiers. + + Abstract methods that need to be implemented: + + - update_mask: Function to compute a new mask for all keys in the + `groups`. + + Args: + - model [nn.Module]: model to configure. The model itself is not saved + but used for the state_dict saving / loading. + - config [list]: configuration elements should be a dict map that includes + `tensor_fqn` of tensors to sparsify + - defaults [dict]: default configurations will be attached to the + configuration. Only the keys that don't exist in the `config` will + be updated. + + Example:: + + >>> # xdoctest: +SKIP("Can't instantiate abstract class BaseSparsifier with abstract method update_mask") + >>> config = [{'tensor_fqn': 'layer1.weight', 'tensor_fqn': 'linear2.weight2', 'sparsity_level': 0.5}] + >>> defaults = {'sparsity_level': 0.7} + >>> # model.layer1.weight will have `sparsity_level` = 0.7 (getting default) + >>> sparsifier = BaseSparsifier(config, defaults) + """ + + def __init__(self, defaults: Optional[Dict[str, Any]] = None): + super().__init__() + self.defaults: Dict[str, Any] = defaults or {} + + self.state: Dict[str, Dict] = defaultdict(dict) + self.groups: List[Dict[str, Any]] = [] + self.enable_mask_update = True + + def __getstate__(self) -> Dict[str, Any]: + return { + "defaults": self.defaults, + "state": self.state, + "groups": self.groups, + } + + def __setstate__(self, state: Dict[str, Dict[str, Any]]) -> None: + self.__dict__.update(state) + + def __repr__(self): + format_string = self.__class__.__name__ + " (" + for i, sparse_args in enumerate(self.groups): + module = sparse_args["module"] + format_string += "\n" + format_string += f"\tGroup {i}\n" + format_string += f"\t module: {module}\n" + for key in sorted(sparse_args.keys()): + if key == "module": + continue + format_string += f"\t {key}: {sparse_args[key]}\n" + format_string += ")" + return format_string + + def state_dict(self) -> Dict[str, Any]: + r"""Returns the state of the optimizer as a :class:`dict`. + + It contains: + * state - current state of the sparsification. + * groups - a list containing all sparsity configuration groups + with the key 'tensor_fqn' specifying the path to the sparsified tensor within a model + + TODO: Need a clean way of loading the state of the "prepared" module + """ + + groups: List[Dict[str, Any]] = [ + dict( + filter( + lambda key_value: key_value[0] not in KEYS_NOT_IN_STATE_DICT, + mg.items(), + ) + ) + for mg in self.groups + ] + + return { + "state": self.state, + "groups": groups, + } + + def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True): + groups = copy.deepcopy(state_dict["groups"]) + states = state_dict["state"] + for tensor_fqn, s in states.items(): + arg_info = get_arg_info_from_tensor_fqn(self.model, tensor_fqn) + module = arg_info["module"] + tensor_name = arg_info["tensor_name"] + if strict and module is None: + raise RuntimeError(f"Error loading {tensor_fqn} into the model") + + found = False + for p in module.parametrizations[tensor_name]: + if isinstance(p, FakeSparsity): + found = True + break + if not found: + p = FakeSparsity(torch.ones(getattr(module, tensor_name).shape)) + parametrize.register_parametrization(module, tensor_name, p) + if s.get("mask", None) is not None: + mask = s.pop("mask") + p.mask = mask + + for mg in groups: + if mg["tensor_fqn"] == tensor_fqn: + mg.update(arg_info) + self.__setstate__({"state": states, "groups": groups}) + + def make_config_from_model( + self, + model: nn.Module, + SUPPORTED_MODULES: Set[Type] = SUPPORTED_MODULES, + ) -> None: + self.config = [] + stack = [model] + while stack: + module = stack.pop() + for name, child in module.named_children(): + if type(child) in SUPPORTED_MODULES: + module_fqn = module_to_fqn(model, child) + assert isinstance(module_fqn, str) # for mypy + self.config.append({"tensor_fqn": module_fqn + ".weight"}) + else: + stack.append(child) + + def prepare(self, model, config): + r"""Prepares a model, by adding the parametrizations. + + Note:: + + The model is modified inplace. If you need to preserve the original + model, use copy.deepcopy. + """ + self.model = model # TODO: Need to figure out how to load without this. + self.config = config + + # If no config -- try getting all the supported layers + if self.config is None: + self.make_config_from_model(model) + + # TODO: Remove the configuration by reference ('module') + for module_config in self.config: + assert isinstance(module_config, dict), ( + "config elements should be dicts not modules i.e.:" + "[{`tensor_fqn`: `foo.bar.weight`}, {`tensor_fqn`: ... }, ...]" + ) + + assert isinstance(self.defaults, Dict) # for mypy + local_args = copy.deepcopy(self.defaults) + local_args.update(module_config) + + tensor_fqn = local_args.get("tensor_fqn", None) + assert tensor_fqn is not None, ( + "tensor_fqn is a required argument in the sparsity config which" + "replaces previous `module` and [module]`fqn` arguments" + ) + + # populate all information from tensor_fqn + info_from_tensor_fqn = get_arg_info_from_tensor_fqn(model, tensor_fqn) + + # check that whatever was put into local_args agrees with what was obtained + # from tensor_fqn + for key in info_from_tensor_fqn.keys(): + if key in local_args: + assert ( + info_from_tensor_fqn[key] == local_args[key] + or ( + key == "tensor_fqn" + and "." + info_from_tensor_fqn[key] == local_args[key] + ) + # info_from_tensor_fqn will chop leading '.' from tensor_fqn so ignore that + ), ( + f"Given both `{key}` and `tensor_fqn` in the config, it is expected them to agree!" + ) + local_args.update(info_from_tensor_fqn) + self.groups.append(local_args) + self._prepare() + + def _prepare(self, *args, **kwargs): + r"""Adds mask parametrization to the layer weight""" + for config in self.groups: + module = config["module"] + tensor_name = config["tensor_name"] + parametrization = config.get("parametrization", FakeSparsity) + mask = config.get("mask", torch.ones_like(getattr(module, tensor_name))) + self.state[config["tensor_fqn"]]["mask"] = mask + parametrize.register_parametrization( + module, tensor_name, parametrization(mask) + ) + + def squash_mask( + self, + params_to_keep: Optional[Tuple[str, ...]] = None, + params_to_keep_per_layer: Optional[Dict[str, Tuple[str, ...]]] = None, + *args, + **kwargs, + ): + r"""Squashes the sparse masks into the appropriate tensors. + + If either the `params_to_keep` or `params_to_keep_per_layer` is set, + the module will have a `sparse_params` dict attached to it. + + Args: + params_to_keep: List of keys to save in the module or a dict + representing the modules and keys that will have + sparsity parameters saved + params_to_keep_per_layer: Dict to specify the params that should be + saved for specific layers. The keys in the dict + should be the module fqn, while the values should + be a list of strings with the names of the variables + to save in the `sparse_params` + + Examples: + >>> # xdoctest: +SKIP("locals are undefined") + >>> # Don't save any sparse params + >>> sparsifier.squash_mask() + >>> hasattr(model.submodule1, 'sparse_params') + False + + >>> # Keep sparse params per layer + >>> sparsifier.squash_mask( + ... params_to_keep_per_layer={ + ... 'submodule1.linear1': ('foo', 'bar'), + ... 'submodule2.linear42': ('baz',) + ... }) + >>> print(model.submodule1.linear1.sparse_params) + {'foo': 42, 'bar': 24} + >>> print(model.submodule2.linear42.sparse_params) + {'baz': 0.1} + + >>> # Keep sparse params for all layers + >>> sparsifier.squash_mask(params_to_keep=('foo', 'bar')) + >>> print(model.submodule1.linear1.sparse_params) + {'foo': 42, 'bar': 24} + >>> print(model.submodule2.linear42.sparse_params) + {'foo': 42, 'bar': 24} + + >>> # Keep some sparse params for all layers, and specific ones for + >>> # some other layers + >>> sparsifier.squash_mask( + ... params_to_keep=('foo', 'bar'), + ... params_to_keep_per_layer={ + ... 'submodule2.linear42': ('baz',) + ... }) + >>> print(model.submodule1.linear1.sparse_params) + {'foo': 42, 'bar': 24} + >>> print(model.submodule2.linear42.sparse_params) + {'foo': 42, 'bar': 24, 'baz': 0.1} + """ + for config in self.groups: + module = config["module"] + tensor_name = config["tensor_name"] + parametrize.remove_parametrizations( + module, tensor_name, leave_parametrized=True + ) + sparse_params = {} + if params_to_keep is not None: + global_params = {k: config[k] for k in params_to_keep} + sparse_params.update(global_params) + if params_to_keep_per_layer is not None: + params = params_to_keep_per_layer.get(config["module_fqn"], None) + if params is not None: + per_layer_params = {k: config[k] for k in params} + sparse_params.update(per_layer_params) + if sparse_params: + # TODO handle multiple tensor being quantized on a single module, where to store sparse_params? + module.sparse_params = sparse_params + + def convert( + self, + module: nn.Module, + mapping: Optional[Dict[Type[nn.Module], Type[nn.Module]]] = None, + inplace: bool = False, + parameterization: Type[nn.Module] = FakeSparsity, + ): + r"""Converts submodules in input module to a different module according to `mapping` + by calling `from_dense` method on the target module class + Args: + module: input module + mapping: a dictionary that maps from source module type to target + module type, can be overwritten to allow swapping user defined + Modules + inplace: carry out model transformations in-place, the original module + is mutated + """ + if mapping is None: + raise NotImplementedError("Need to auto generate mapping ") + if not inplace: + module = copy.deepcopy(module) + + reassign = {} + for name, mod in module.named_children(): + # leaf node + if ( + module_contains_param(mod, parameterization) + and type_before_parametrizations(mod) in mapping + ): + reassign[name] = swap_module(mod, mapping) + else: + # recurse + reassign[name] = self.convert( + mod, + mapping=mapping, + inplace=True, + parameterization=parameterization, + ) + + for key, value in reassign.items(): + module._modules[key] = value + + return module + + def step(self, use_path: bool = True) -> None: + if not self.enable_mask_update: + return + with torch.no_grad(): + for config in self.groups: + self.update_mask(**config) + + @abc.abstractmethod + def update_mask(self, module: nn.Module, tensor_name: str, **kwargs): + pass diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/nearly_diagonal_sparsifier.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/nearly_diagonal_sparsifier.py new file mode 100644 index 0000000000000000000000000000000000000000..a1f824c3732d82c9033681d1a2ecd4b77df30c3a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/nearly_diagonal_sparsifier.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch + +from . import base_sparsifier + + +class NearlyDiagonalSparsifier(base_sparsifier.BaseSparsifier): + r"""Nearly Diagonal Sparsifier + + This sparsifier creates a nearly diagonal mask to be applied to the weight matrix. + Nearly Diagonal Matrix is a matrix that contains non-zero elements near the diagonal and the rest are zero. + An example of a nearly diagonal matrix with degree (or nearliness) 3 and 5 are follows respectively. + 1 1 0 0 1 1 1 0 + 1 1 1 0 1 1 1 1 + 0 1 1 1 1 1 1 1 + 0 0 1 1 0 1 1 1 + Note that a nearly diagonal matrix with degree 1 is just a matrix with main diagonal populated + + This sparsifier is controlled by one variable: + 1. `nearliness` defines the number of non-zero diagonal lines that are closest to the main diagonal. + Currently - supports only odd number + + Note: + This can be accelerated (vectorized) once the Spdiagonal feature (PR: #78439) is landed or the banded matrix + feature is landed: https://stackoverflow.com/questions/52463972/generating-banded-matrices-using-numpy + + Args: + nearliness: The degree of nearliness (default = 1) + + """ + + def __init__(self, nearliness: int = 1): + defaults = {"nearliness": nearliness} + super().__init__(defaults=defaults) + + def update_mask(self, module, tensor_name, nearliness, **kwargs): + mask = getattr(module.parametrizations, tensor_name)[0].mask + mask.data = torch.zeros_like(mask) + if nearliness <= 0: + return + + tensor = getattr(module, tensor_name) + height, width = tensor.shape + + if nearliness % 2 == 0: + raise ValueError("nearliness can only be an odd number") + dist_to_diagonal = nearliness // 2 + # check + if dist_to_diagonal >= min(height, width): + raise ValueError( + "nearliness cannot be larger than the dimensions of tensor." + ) + + for row in range(0, height): + # Bounds of entries that needs to be set to 1 + low = max(0, row - dist_to_diagonal) + high = min(width, row + dist_to_diagonal + 1) + mask[row, low:high].fill_(1) diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/utils.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..855658a074c196b7bcf4daa6e7760999ea1b2426 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/utils.py @@ -0,0 +1,135 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from itertools import chain +from typing import Any, Dict, Optional, Type + +from torch import nn +from torch.nn.utils.parametrize import is_parametrized, type_before_parametrizations + +__all__ = [ + "module_contains_param", + "swap_module", + "module_to_fqn", + "fqn_to_module", + "get_arg_info_from_tensor_fqn", + "FakeSparsity", +] + + +def module_contains_param(module: nn.Module, parametrization: Type[nn.Module]) -> bool: + if is_parametrized(module): + # see if any of the module tensors have a parametriztion attached that matches the one passed in + return any( + any(isinstance(param, parametrization) for param in param_list) + for key, param_list in module.parametrizations.items() # type: ignore[union-attr,operator] + ) + return False + + +def swap_module( + mod: nn.Module, mapping: Dict[Type[nn.Module], Type[nn.Module]] +) -> nn.Module: + r"""Swaps the module using from_dense according to the mapping passed in. + Args: + mod: input module + mapping: a dictionary that maps from nn module to sparse nn module + Return: + The corresponding sparse module of `mod` according to mapping, created using from_dense + """ + if type_before_parametrizations(mod) in mapping: + sparse_mod = mapping[type_before_parametrizations(mod)] + + # TODO Fix this typing, as Type[Module] has no attribute "from_dense" + new_mod = sparse_mod.from_dense(mod) # type: ignore[attr-defined] + + # Preserve module's pre forward hooks. They'll be called on quantized input + for pre_hook_fn in mod._forward_pre_hooks.values(): + new_mod.register_forward_pre_hook(pre_hook_fn) + # Preserve module's post forward hooks except _observer_forward_hook + # After convert they'll work with quantized output + for hook_fn in mod._forward_hooks.values(): + new_mod.register_forward_hook(hook_fn) + + # respect device affinity when swapping modules + devices = {p.device for p in chain(mod.parameters(), mod.buffers())} + assert len(devices) <= 1, ( + f"swap_module only works with cpu or single-device CUDA modules, but got devices {devices}" + ) + device = next(iter(devices)) if len(devices) > 0 else None + if device: + new_mod.to(device) + + return new_mod + + else: + return mod + + +def module_to_fqn( + model: nn.Module, module: nn.Module, prefix: str = "" +) -> Optional[str]: + """ + Returns the fqn for a module or None if module not a descendent of model. + """ + if module is model: + return "" + for name, child in model.named_children(): + fqn = module_to_fqn(child, module, ".") + if isinstance(fqn, str): + return prefix + name + fqn + return None + + +def fqn_to_module(model: Optional[nn.Module], path: str) -> Optional[nn.Module]: + """ + Given an fqn, returns the corresponding module or tensor or None if the fqn given by `path` + doesn't correspond to anything. Similar to model.get_submodule(path) but works for tensors. + """ + if path != "": + for name in path.split("."): + model = getattr(model, name, None) + return model + + +def get_arg_info_from_tensor_fqn(model: nn.Module, tensor_fqn: str) -> Dict[str, Any]: + """ + Uses tensor_fqn to obtain a dict containing module_fqn, module and tensor_name + """ + # string manip to split tensor_fqn into module_fqn and tensor_name + # if tensor_fqn is 'weight' then module_fqn and tensor_name are '' and 'weight' + # if tensor_fqn is 'linear.weight' then module_fqn and tensor_name are 'linear' and 'weight' + tensor_name = tensor_fqn.split(".")[-1] + module_fqn = tensor_fqn[: -len(tensor_name) - ("." in tensor_fqn)] + + module = fqn_to_module(model, module_fqn) + + return { + "module_fqn": module_fqn, + "module": module, + "tensor_name": tensor_name, + "tensor_fqn": tensor_fqn, + } + + +# Parametrizations +class FakeSparsity(nn.Module): + r"""Parametrization for the weights. Should be attached to the 'weight' or + any other parameter that requires a mask applied to it. + + Note:: + + Once the mask is passed, the variable should not change the id. The + contents of the mask can change, but the mask reference itself should + not. + """ + + def __init__(self, mask): + super().__init__() + self.register_buffer("mask", mask) + + def forward(self, x): + assert self.mask.shape == x.shape + return self.mask * x diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/weight_norm_sparsifier.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/weight_norm_sparsifier.py new file mode 100644 index 0000000000000000000000000000000000000000..d86e6c63cf8e327de9343a0cf96d69c7ed8c88c2 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/sparsifier/weight_norm_sparsifier.py @@ -0,0 +1,251 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import operator +from functools import reduce +from typing import Callable, Optional, Tuple, Union + +import torch +import torch.nn.functional as F + +from .base_sparsifier import BaseSparsifier + +__all__ = ["WeightNormSparsifier"] + + +def _flat_idx_to_2d(idx, shape): + rows = idx // shape[1] + cols = idx % shape[1] + return rows, cols + + +class WeightNormSparsifier(BaseSparsifier): + r"""Weight-Norm Sparsifier + + This sparsifier computes the norm of every sparse block and "zeroes-out" the + ones with the lowest norm. The level of sparsity defines how many of the + blocks is removed. + + This sparsifier is controlled by three variables: + 1. `sparsity_level` defines the number of *sparse blocks* that are zeroed-out + 2. `sparse_block_shape` defines the shape of the sparse blocks. Note that + the sparse blocks originate at the zero-index of the tensor. + 3. `zeros_per_block` is the number of zeros that we are expecting in each + sparse block. By default we assume that all elements within a block are + zeroed-out. However, setting this variable sets the target number of + zeros per block. The zeros within each block are chosen as the *smallest + absolute values*. + + Args: + + sparsity_level: The target level of sparsity + sparse_block_shape: The shape of a sparse block (see note below) + zeros_per_block: Number of zeros in a sparse block + norm: Norm to use. Could be either `int` or a callable. + If `int`, only L1 and L2 are implemented. + + Note:: + The `sparse_block_shape` is tuple representing (block_ROWS, block_COLS), + irrespective of what the rows / cols mean in the data tensor. That means, + if you were to sparsify a weight tensor in the nn.Linear, which has a + weight shape `(Cout, Cin)`, the `block_ROWS` would refer to the output + channels, while the `block_COLS` would refer to the input channels. + + Note:: + All arguments to the WeightNormSparsifier constructor are "default" + arguments and could be overriden by the configuration provided in the + `prepare` step. + """ + + def __init__( + self, + sparsity_level: float = 0.5, + sparse_block_shape: Tuple[int, int] = (1, 4), + zeros_per_block: Optional[int] = None, + norm: Optional[Union[Callable, int]] = None, + ): + if zeros_per_block is None: + zeros_per_block = reduce(operator.mul, sparse_block_shape) + defaults = { + "sparsity_level": sparsity_level, + "sparse_block_shape": sparse_block_shape, + "zeros_per_block": zeros_per_block, + } + if norm is None: + norm = 2 + if callable(norm): + self.norm_fn = norm + elif norm == 1: + self.norm_fn = lambda T: T.abs() + elif norm == 2: + self.norm_fn = lambda T: T * T + else: + raise NotImplementedError(f"L-{norm} is not yet implemented.") + super().__init__(defaults=defaults) + + def _scatter_fold_block_mask( + self, + output_shape, + dim, + indices, + block_shape, + mask=None, + input_shape=None, + device=None, + ): + r"""Creates patches of size `block_shape` after scattering the indices.""" + if mask is None: + assert input_shape is not None + mask = torch.ones(input_shape, device=device) + mask.scatter_(dim=dim, index=indices, value=0) + mask.data = F.fold( + mask, output_size=output_shape, kernel_size=block_shape, stride=block_shape + ) + return mask + + def _make_tensor_mask( + self, data, input_shape, sparsity_level, sparse_block_shape, mask=None + ): + r"""Creates a tensor-level mask. + + Tensor-level mask is described as a mask, where the granularity of sparsification of the + smallest patch is the sparse_block_shape. That means, that for a given mask and a + sparse_block_shape, the smallest "patch" of zeros/ones could be the sparse_block_shape. + + In this context, `sparsity_level` describes the fraction of sparse patches. + """ + h, w = data.shape[-2:] + block_h, block_w = sparse_block_shape + dh = (block_h - h % block_h) % block_h + dw = (block_w - w % block_w) % block_w + + if mask is None: + mask = torch.ones(h + dh, w + dw, device=data.device) + + if sparsity_level >= 1.0: + mask.data = torch.zeros_like(mask) + return mask + elif sparsity_level <= 0.0: + mask.data = torch.ones_like(mask) + return mask + + values_per_block = reduce(operator.mul, sparse_block_shape) + if values_per_block > 1: + # Reduce the data + data = F.avg_pool2d( + data[None, None, :], + kernel_size=sparse_block_shape, + stride=sparse_block_shape, + ceil_mode=True, + ) + data = data.flatten() + num_blocks = len(data) + + data = data.repeat(1, values_per_block, 1) + + threshold_idx = int(round(sparsity_level * num_blocks)) + threshold_idx = max(0, min(num_blocks - 1, threshold_idx)) # Sanity check + _, sorted_idx = torch.topk(data, k=threshold_idx, dim=2, largest=False) + + # Temp reshape for mask + mask_reshape = mask.reshape(data.shape) # data might be reshaped + self._scatter_fold_block_mask( + dim=2, + output_shape=(h + dh, w + dw), + indices=sorted_idx, + block_shape=sparse_block_shape, + mask=mask_reshape, + ) + mask.data = mask_reshape.squeeze().reshape(mask.shape)[:h, :w].contiguous() + return mask + + def _make_block_mask(self, data, sparse_block_shape, zeros_per_block, mask=None): + r"""Creates a block-level mask. + + Block-level mask is described as a mask, where the granularity of sparsification of the + largest patch is the sparse_block_shape. That means that for a given mask and a + sparse_block_shape, the sparsity is computed only within a patch of a size sparse_block_shape. + + In this context the `zeros_per_block` describes the number of zeroed-out elements within a patch. + """ + h, w = data.shape[-2:] + block_h, block_w = sparse_block_shape + dh = (block_h - h % block_h) % block_h + dw = (block_w - w % block_w) % block_w + values_per_block = reduce(operator.mul, sparse_block_shape) + + if mask is None: + mask = torch.ones((h + dh, w + dw), device=data.device) + + if values_per_block == zeros_per_block: + # Everything should be sparsified + mask.data = torch.zeros_like(mask) + return mask + + # create a new padded tensor like data (to match the block_shape) + padded_data = torch.ones(h + dh, w + dw, dtype=data.dtype, device=data.device) + padded_data.fill_(torch.nan) + padded_data[:h, :w] = data + unfolded_data = F.unfold( + padded_data[None, None, :], + kernel_size=sparse_block_shape, + stride=sparse_block_shape, + ) + + # Temp reshape for mask + mask_reshape = mask.reshape(unfolded_data.shape) + _, sorted_idx = torch.topk( + unfolded_data, k=zeros_per_block, dim=1, largest=False + ) + + self._scatter_fold_block_mask( + dim=1, + indices=sorted_idx, + output_shape=padded_data.shape, + block_shape=sparse_block_shape, + mask=mask_reshape, + ) + + mask.data = mask_reshape.squeeze().reshape(mask.shape).contiguous() + return mask + + def update_mask( + self, + module, + tensor_name, + sparsity_level, + sparse_block_shape, + zeros_per_block, + **kwargs, + ): + values_per_block = reduce(operator.mul, sparse_block_shape) + if zeros_per_block > values_per_block: + raise ValueError( + "Number of zeros per block cannot be more than the total number of elements in that block." + ) + if zeros_per_block < 0: + raise ValueError("Number of zeros per block should be positive.") + + mask = getattr(module.parametrizations, tensor_name)[0].mask + if sparsity_level <= 0 or zeros_per_block == 0: + mask.data = torch.ones_like(mask) + elif sparsity_level >= 1.0 and (zeros_per_block == values_per_block): + mask.data = torch.zeros_like(mask) + else: + ww = self.norm_fn(getattr(module, tensor_name)) + tensor_mask = self._make_tensor_mask( + data=ww, + input_shape=ww.shape, + sparsity_level=sparsity_level, + sparse_block_shape=sparse_block_shape, + ) + if values_per_block != zeros_per_block: + block_mask = self._make_block_mask( + data=ww, + sparse_block_shape=sparse_block_shape, + zeros_per_block=zeros_per_block, + ) + tensor_mask = torch.logical_or(tensor_mask, block_mask) + mask.data = tensor_mask diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__init__.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7030f5ad445885874f68971c2c8e3c2bf01397af Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/benchmark.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/benchmark.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f35a4242da22fa3447370fe8d466d31f661213d Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/benchmark.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/evaluate.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/evaluate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae8f09eb8f7744fdc05265061d3a0609adbba722 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/evaluate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/train.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/train.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9c0ac11f02e6c24e9a50b8f8917584f22b6fb67 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/train.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f602bafacd5cfd62998c8449ba1afef4d9a6d23 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/benchmark.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..b87834afae4d29fa585de3bb3fa1b3513a7dc5b1 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/benchmark.py @@ -0,0 +1,140 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +import torch +import torchvision +from torch.sparse._triton_ops_meta import ( + optimize_bsr_dense_addmm, +) + +from torchao.prototype.sparsity.superblock.utils import ( + accelerate_with_sparsity, + get_args_parser, + simulate_sparsity, +) +from torchao.utils import benchmark_model, profiler_runner + +torch.sparse.SparseSemiStructuredTensor._FORCE_CUTLASS = False +torch.backends.mha.set_fastpath_enabled(False) + + +@torch.inference_mode +def main(args): + device = torch.device(args.device) + + # We disable the cudnn benchmarking because it can noticeably affect the accuracy + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + num_classes = 1000 + + dtype = getattr(torch, args.dtype) + + # BSR kernel tuning + if args.bsr and args.tune_kernel_params: + kwargs = dict( + dtype=torch.int8 if args.quantization else dtype, + sparsity=args.sparsity_linear, + verbose=True, + # per blocksparse_int_addmm: + alpha=1, + beta=0, + use_left_alpha=True, + use_right_alpha=True, + # force tuning because existing tuning parameters are + # computed for use_left/right_alpha=False, however, it + # turns out that re-tuning for use_left/right_alpha=False + # leads to the same set of tuning parametes: + # force=True + ) + if args.model == "vit_b_16": + optimize_bsr_dense_addmm(3072, 768, 50432, args.bsr, args.bsr, **kwargs) + optimize_bsr_dense_addmm(768, 3072, 50432, args.bsr, args.bsr, **kwargs) + elif args.model == "vit_h_14": + optimize_bsr_dense_addmm(5120, 1280, 65792, args.bsr, args.bsr, **kwargs) + optimize_bsr_dense_addmm(1280, 5120, 65792, args.bsr, args.bsr, **kwargs) + else: + raise NotImplementedError( + "Tuning kernel params for this model is not supported yet." + ) + # Warning: the following call will overwrite the source code + # of torch.sparse._triton_ops_meta (hence it is commented out + # by default) but when used, it'll enables reusing the tuned + # parameters in subsequent runs of this script: + # store_tuned_kernel_params() + model = torchvision.models.get_model( + args.model, weights=args.weights, num_classes=num_classes + ).eval() + + # Fake sparsity necessary for BSR, since we find based on SuperBlock + sparsifier_or_none = simulate_sparsity(model, args) + if sparsifier_or_none is not None: + sparsifier_or_none.squash_mask() + + if args.weights_path: + try: + checkpoint = torch.load(args.weights_path, map_location="cpu") + model.load_state_dict(checkpoint["model"]) + except FileNotFoundError: + raise FileNotFoundError(f"No checkpoint found at {args.weights_path}.") + + model.to(device).to(dtype) + + # With quantization, we must use cuSPARSELt to fuse one of the scalar matmuls. + # Otherwise, we observe the CUTLASS kernels to be faster, so we use those instead. + accelerate_with_sparsity(model, args) + + # compile + model = torch.compile(model, mode="max-autotune", fullgraph=True) + + # define image + image = torch.randn( + args.batch_size, + 3, + args.val_crop_size, + args.val_crop_size, + dtype=dtype, + device=device, + ) + + # warmup + benchmark_model(model, 10, args=(image,)) + if args.profile: + return profiler_runner("test.json.gz", benchmark_model, model, 10, (image,)) + else: + return benchmark_model(model, 100, args=(image,)) + + +if __name__ == "__main__": + args = get_args_parser(benchmark=True).parse_args() + result = main(args) + header = [ + "model", + "batch_size", + "dtype", + "sparsity", + "bsr", + "sparsity_level", + "quantization", + "tune_kernel_params", + "latency", + "img/s", + ] + result_string = ",".join( + str(_) + for _ in [ + args.model, + args.batch_size, + args.dtype, + args.sparsity, + args.bsr, + args.sparsity_linear, + args.quantization, + args.tune_kernel_params, + result, + 1000 / result, + ] + ) + with open("benchmark_results.txt", "a") as f: + if args.header: + f.write(",".join(header) + "\n") + f.write(result_string + "\n") + print(result_string) diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/evaluate.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..04701e518ea02daa6867c1201cbca9e816b2864c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/evaluate.py @@ -0,0 +1,104 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +import os + +import torch +import torchvision + +from torchao.prototype.sparsity.superblock.train import evaluate, load_data +from torchao.prototype.sparsity.superblock.utils import ( + accelerate_with_sparsity, + get_args_parser, + init_distributed_mode, + simulate_sparsity, +) + +torch.sparse.SparseSemiStructuredTensor._FORCE_CUTLASS = False +torch.backends.mha.set_fastpath_enabled(False) + + +def main(args): + init_distributed_mode(args) + print(args) + + device = torch.device(args.device) + # We disable the cudnn benchmarking because it can noticeably affect the accuracy + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + + # Load validation data + val_dir = os.path.join(args.data_path, "val") + dataset_test, test_sampler = load_data(None, val_dir, args) + data_loader_test = torch.utils.data.DataLoader( + dataset_test, + batch_size=args.batch_size, + sampler=test_sampler, + num_workers=args.workers, + pin_memory=True, + drop_last=True, + ) + num_classes = len(dataset_test.classes) + + # Create Model + print("Creating model") + model = torchvision.models.get_model( + args.model, weights=args.weights, num_classes=num_classes + ) + + sparsifier_or_none = simulate_sparsity(model, args) + + if args.weights_path: + try: + checkpoint = torch.load(args.weights_path, map_location="cpu") + model.load_state_dict(checkpoint["model"]) + print(f"Loaded checkpoint successfully from: {args.weights_path}") + except FileNotFoundError: + raise FileNotFoundError(f"No checkpoint found at {args.weights_path}") + + model.to(device).bfloat16() + + if sparsifier_or_none is not None: + sparsifier_or_none.squash_mask() + accelerate_with_sparsity(model, args) + model = torch.compile(model, mode="max-autotune", fullgraph=True) + + criterion = torch.nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + return evaluate( + model, criterion, data_loader_test, device=device, dtype=torch.bfloat16 + ) + + +if __name__ == "__main__": + args = get_args_parser(evaluate=True).parse_args() + accuracy, throughput, max_mem = main(args) + header = [ + "model", + "batch_size", + "dtype", + "sparsity", + "bsr", + "sparsity_level", + "quantization", + "top-1_acc", + "encoder img/s", + "max_mem (MB)", + ] + result_string = ",".join( + str(_) + for _ in [ + args.model, + args.batch_size, + "bfloat16", + args.sparsity, + args.bsr, + args.sparsity_linear, + args.quantization, + accuracy, + throughput, + max_mem, + ] + ) + with open("evaluation_results.txt", "a") as f: + if args.header: + f.write(",".join(header) + "\n") + f.write(result_string + "\n") + print(result_string) diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/train.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/train.py new file mode 100644 index 0000000000000000000000000000000000000000..330dba4b6dfd7d581ec6be4d6f76543b7a27ea2b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/train.py @@ -0,0 +1,542 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +import datetime +import glob +import os +import time +import warnings + +import torch +import torch.utils.data +import torchvision +import utils +from torch import nn +from torch.utils.data.dataloader import default_collate +from torchvision.transforms.functional import InterpolationMode +from utils import RASampler + +from torchao.prototype.sparsity.superblock.utils import simulate_sparsity + + +def train_one_epoch( + model, + criterion, + optimizer, + data_loader, + device, + epoch, + args, + model_ema=None, + scaler=None, +): + model.train() + metric_logger = utils.MetricLogger(delimiter=" ") + metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value}")) + metric_logger.add_meter("img/s", utils.SmoothedValue(window_size=10, fmt="{value}")) + + header = f"Epoch: [{epoch}]" + accumulation_counter = 0 # Counter for tracking accumulated gradients + + for i, (image, target) in enumerate( + metric_logger.log_every(data_loader, args.print_freq, header) + ): + start_time = time.time() + image, target = image.to(device), target.to(device) + + with torch.cuda.amp.autocast(enabled=scaler is not None): + output = model(image) + loss = criterion(output, target) / args.accumulation_steps # Scale loss + + if scaler is not None: + scaler.scale(loss).backward() + else: + loss.backward() + + accumulation_counter += 1 + + if accumulation_counter % args.accumulation_steps == 0: + if scaler is not None: + if args.clip_grad_norm is not None: + scaler.unscale_(optimizer) # Unscale gradients before clipping + nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad_norm) + scaler.step(optimizer) + scaler.update() + else: + if args.clip_grad_norm is not None: + nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad_norm) + optimizer.step() + + optimizer.zero_grad() # Zero out gradients after optimization step + + if model_ema and i % args.model_ema_steps == 0: + model_ema.update_parameters(model) + if epoch < args.lr_warmup_epochs: + model_ema.n_averaged.fill_(0) + + acc1, acc5 = utils.accuracy(output, target, topk=(1, 5)) + batch_size = image.shape[0] + metric_logger.update( + loss=loss.item() * args.accumulation_steps, + lr=optimizer.param_groups[0]["lr"], + ) # Scale back up for logging + metric_logger.meters["acc1"].update(acc1.item(), n=batch_size) + metric_logger.meters["acc5"].update(acc5.item(), n=batch_size) + metric_logger.meters["img/s"].update(batch_size / (time.time() - start_time)) + + +def evaluate( + model, + criterion, + data_loader, + device, + print_freq=100, + log_suffix="", + dtype=torch.float32, +): + model.eval() + metric_logger = utils.MetricLogger(delimiter=" ") + header = f"Test: {log_suffix}" + encoder_time = 0 + num_processed_samples = 0 + with torch.inference_mode(): + for image, target in metric_logger.log_every(data_loader, print_freq, header): + image = image.to(device, non_blocking=True).to(dtype) + target = target.to(device, non_blocking=True).to(dtype) + # intialize encoder measurements + torch.cuda.reset_max_memory_allocated() + torch.cuda.synchronize() + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + + # run encoder + output = model(image) + + # measure time in encoder + end_event.record() + torch.cuda.synchronize() + encoder_time += start_event.elapsed_time(end_event) + max_mem = torch.cuda.max_memory_allocated() / (1024**2) + + acc1, acc5 = utils.accuracy(output, target, topk=(1, 5)) + # FIXME need to take into account that the datasets + # could have been padded in distributed setup + batch_size = image.shape[0] + # metric_logger.update(loss=loss.item()) + metric_logger.meters["acc1"].update(acc1.item(), n=batch_size) + metric_logger.meters["acc5"].update(acc5.item(), n=batch_size) + metric_logger.meters["batch_time"].update(encoder_time, n=batch_size) + num_processed_samples += batch_size + # gather the stats from all processes + + num_processed_samples = utils.reduce_across_processes(num_processed_samples) + if ( + hasattr(data_loader.dataset, "__len__") + and len(data_loader.dataset) != num_processed_samples + ): + # See FIXME above + warnings.warn( + f"It looks like the dataset has {len(data_loader.dataset)} samples, but {num_processed_samples} " + "samples were used for the validation, which might bias the results. " + "Try adjusting the batch size and / or the world size. " + "Setting the world size to 1 is always a safe bet." + ) + + metric_logger.synchronize_between_processes() + + print( + f"{header} Acc@1 {metric_logger.acc1.global_avg:.3f} Acc@5 {metric_logger.acc5.global_avg:.3f}" + ) + total_time = encoder_time / 1000.0 + return ( + metric_logger.acc1.global_avg, + num_processed_samples.item() / total_time, + max_mem, + ) + + +def _get_cache_path(filepath): + import hashlib + + h = hashlib.sha1(filepath.encode()).hexdigest() + cache_path = os.path.join( + "~", ".torch", "vision", "datasets", "imagefolder", h[:10] + ".pt" + ) + cache_path = os.path.expanduser(cache_path) + return cache_path + + +def load_data(traindir, valdir, args): + # Data loading code + print("Loading data") + ( + val_resize_size, + val_crop_size, + ) = ( + args.val_resize_size, + args.val_crop_size, + ) + interpolation = InterpolationMode(args.interpolation) + if traindir is not None: + train_crop_size = args.train_crop_size + print("Loading training data") + st = time.time() + cache_path = _get_cache_path(traindir) + if args.cache_dataset and os.path.exists(cache_path): + # Attention, as the transforms are also cached! + print(f"Loading dataset_train from {cache_path}") + dataset, _ = torch.load(cache_path) + else: + auto_augment_policy = getattr(args, "auto_augment", None) + random_erase_prob = getattr(args, "random_erase", 0.0) + ra_magnitude = args.ra_magnitude + augmix_severity = args.augmix_severity + preprocessing = utils.ClassificationPresetTrain( + crop_size=train_crop_size, + interpolation=interpolation, + auto_augment_policy=auto_augment_policy, + random_erase_prob=random_erase_prob, + ra_magnitude=ra_magnitude, + augmix_severity=augmix_severity, + ) + dataset = torchvision.datasets.ImageFolder(traindir, preprocessing) + # ) if args.meta else torchvision.datasets.ImageNet( + # traindir, + # split="train", + # transform=preprocessing, + # ) + if args.cache_dataset: + print(f"Saving dataset_train to {cache_path}") + utils.mkdir(os.path.dirname(cache_path)) + utils.save_on_master((dataset, traindir), cache_path) + print("Took", time.time() - st) + print(f"Number of training images: {len(dataset)}") + if args.distributed: + if hasattr(args, "ra_sampler") and args.ra_sampler: + train_sampler = RASampler( + dataset, shuffle=True, repetitions=args.ra_reps + ) + else: + train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) + else: + train_sampler = torch.utils.data.RandomSampler(dataset) + + print("Loading validation data") + cache_path = _get_cache_path(valdir) + if args.cache_dataset and os.path.exists(cache_path): + # Attention, as the transforms are also cached! + print(f"Loading dataset_test from {cache_path}") + dataset_test, test_sampler = torch.load(cache_path) + else: + if args.weights: + weights = torchvision.models.get_weight(args.weights) + preprocessing = weights.transforms() + else: + preprocessing = utils.ClassificationPresetEval( + crop_size=val_crop_size, + resize_size=val_resize_size, + interpolation=interpolation, + ) + dataset_test = ( + torchvision.datasets.ImageFolder( + valdir, + preprocessing, + ) + if args.meta + else torchvision.datasets.ImageNet( + valdir, split="val", transform=preprocessing + ) + ) + if args.cache_dataset: + print(f"Saving dataset_test to {cache_path}") + utils.mkdir(os.path.dirname(cache_path)) + utils.save_on_master((dataset_test, valdir), cache_path) + + print(f"Number of validation images: {len(dataset_test)}") + test_sampler = ( + torch.utils.data.distributed.DistributedSampler(dataset_test, shuffle=False) + if args.distributed + else torch.utils.data.SequentialSampler(dataset_test) + ) + + # for evaluation + if traindir is None: + return dataset_test, test_sampler + + return dataset, dataset_test, train_sampler, test_sampler + + +def main(args): + if args.output_dir: + utils.mkdir(args.output_dir) + + utils.init_distributed_mode(args) + print(args) + + device = torch.device(args.device) + + if args.use_deterministic_algorithms: + torch.backends.cudnn.benchmark = False + torch.use_deterministic_algorithms(True) + else: + torch.backends.cudnn.benchmark = True + + train_dir = os.path.join(args.data_path, "train_blurred") + val_dir = os.path.join(args.data_path, "val") + dataset, dataset_test, train_sampler, test_sampler = load_data( + train_dir, val_dir, args + ) + + collate_fn = None + num_classes = len(dataset.classes) + mixup_transforms = [] + if args.mixup_alpha > 0.0: + mixup_transforms.append( + utils.RandomMixup(num_classes, p=1.0, alpha=args.mixup_alpha) + ) + if args.cutmix_alpha > 0.0: + mixup_transforms.append( + utils.RandomCutmix(num_classes, p=1.0, alpha=args.cutmix_alpha) + ) + if mixup_transforms: + mixupcutmix = torchvision.transforms.RandomChoice(mixup_transforms) + + def collate_fn(batch): + return mixupcutmix(*default_collate(batch)) + + data_loader = torch.utils.data.DataLoader( + dataset, + batch_size=args.batch_size, + sampler=train_sampler, + num_workers=args.workers, + pin_memory=True, + collate_fn=collate_fn, + ) + data_loader_test = torch.utils.data.DataLoader( + dataset_test, + batch_size=args.batch_size, + sampler=test_sampler, + num_workers=args.workers, + pin_memory=True, + ) + + print("Creating model") + model = torchvision.models.get_model( + args.model, weights=args.weights, num_classes=num_classes + ) + + if args.weights_path is not None: + sd = torch.load(args.weights_path, map_location="cpu") + model.load_state_dict(sd) + + model.to(device) + if args.distributed and args.sync_bn: + model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) + + sparsifier = simulate_sparsity(model, args) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + + custom_keys_weight_decay = [] + if args.bias_weight_decay is not None: + custom_keys_weight_decay.append(("bias", args.bias_weight_decay)) + if args.transformer_embedding_decay is not None: + for key in [ + "class_token", + "position_embedding", + "relative_position_bias_table", + ]: + custom_keys_weight_decay.append((key, args.transformer_embedding_decay)) + parameters = utils.set_weight_decay( + model, + args.weight_decay, + norm_weight_decay=args.norm_weight_decay, + custom_keys_weight_decay=( + custom_keys_weight_decay if len(custom_keys_weight_decay) > 0 else None + ), + ) + + opt_name = args.opt.lower() + if opt_name.startswith("sgd"): + optimizer = torch.optim.SGD( + parameters, + lr=args.lr, + momentum=args.momentum, + weight_decay=args.weight_decay, + nesterov="nesterov" in opt_name, + ) + elif opt_name == "rmsprop": + optimizer = torch.optim.RMSprop( + parameters, + lr=args.lr, + momentum=args.momentum, + weight_decay=args.weight_decay, + eps=0.0316, + alpha=0.9, + ) + elif opt_name == "adamw": + optimizer = torch.optim.AdamW( + parameters, lr=args.lr, weight_decay=args.weight_decay + ) + else: + raise RuntimeError( + f"Invalid optimizer {args.opt}. Only SGD, RMSprop and AdamW are supported." + ) + + scaler = torch.cuda.amp.GradScaler() if args.amp else None + + args.lr_scheduler = args.lr_scheduler.lower() + if args.lr_scheduler == "steplr": + main_lr_scheduler = torch.optim.lr_scheduler.StepLR( + optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma + ) + elif args.lr_scheduler == "cosineannealinglr": + main_lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=args.epochs - args.lr_warmup_epochs, eta_min=args.lr_min + ) + elif args.lr_scheduler == "exponentiallr": + main_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( + optimizer, gamma=args.lr_gamma + ) + else: + raise RuntimeError( + f"Invalid lr scheduler '{args.lr_scheduler}'. Only StepLR, CosineAnnealingLR and ExponentialLR " + "are supported." + ) + + if args.lr_warmup_epochs > 0: + if args.lr_warmup_method == "linear": + warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR( + optimizer, + start_factor=args.lr_warmup_decay, + total_iters=args.lr_warmup_epochs, + ) + elif args.lr_warmup_method == "constant": + warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR( + optimizer, + factor=args.lr_warmup_decay, + total_iters=args.lr_warmup_epochs, + ) + else: + raise RuntimeError( + f"Invalid warmup lr method '{args.lr_warmup_method}'. Only linear and constant are supported." + ) + lr_scheduler = torch.optim.lr_scheduler.SequentialLR( + optimizer, + schedulers=[warmup_lr_scheduler, main_lr_scheduler], + milestones=[args.lr_warmup_epochs], + ) + else: + lr_scheduler = main_lr_scheduler + + model_without_ddp = model + if args.distributed: + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[args.gpu], find_unused_parameters=True + ) + model_without_ddp = model.module + + model_ema = None + if args.model_ema: + # Decay adjustment that aims to keep the decay independent from other hyper-parameters originally proposed at: + # https://github.com/facebookresearch/pycls/blob/f8cd9627/pycls/core/net.py#L123 + # + # total_ema_updates = (Dataset_size / n_GPUs) * epochs / (batch_size_per_gpu * EMA_steps) + # We consider constant = Dataset_size for a given dataset/setup and ommit it. Thus: + # adjust = 1 / total_ema_updates ~= n_GPUs * batch_size_per_gpu * EMA_steps / epochs + adjust = args.world_size * args.batch_size * args.model_ema_steps / args.epochs + alpha = 1.0 - args.model_ema_decay + alpha = min(1.0, alpha * adjust) + model_ema = utils.ExponentialMovingAverage( + model_without_ddp, device=device, decay=1.0 - alpha + ) + + # TODO: need to test resume functionality + if args.resume: + checkpoint_pattern = os.path.join(args.output_dir, "model_*.pth") + checkpoint_files = glob.glob(checkpoint_pattern) + epochs = [int(f.split("_")[-1].split(".")[0]) for f in checkpoint_files] + if epochs: + latest_epoch = max(epochs) + latest_checkpoint = os.path.join( + args.output_dir, f"model_{latest_epoch}.pth" + ) + try: + checkpoint = torch.load(latest_checkpoint, map_location="cpu") + model_without_ddp.load_state_dict(checkpoint["model"]) + optimizer.load_state_dict(checkpoint["optimizer"]) + lr_scheduler.load_state_dict(checkpoint["lr_scheduler"]) + args.start_epoch = checkpoint["epoch"] + 1 + if model_ema: + model_ema.load_state_dict(checkpoint["model_ema"]) + if scaler: + scaler.load_state_dict(checkpoint["scaler"]) + print(f"Resumed training from epoch {args.start_epoch}.") + except FileNotFoundError: + print( + f"No checkpoint found at {latest_checkpoint}. Starting training from scratch." + ) + args.start_epoch = 0 + else: + print("No checkpoint found. Starting training from scratch.") + args.start_epoch = 0 + else: + args.start_epoch = 0 + print("Zero-shot evaluation") + if model_ema: + evaluate( + model_ema, criterion, data_loader_test, device=device, log_suffix="EMA" + ) + else: + evaluate(model, criterion, data_loader_test, device=device) + + print("Start training") + start_time = time.time() + for epoch in range(args.start_epoch, args.epochs): + if args.distributed: + train_sampler.set_epoch(epoch) + train_one_epoch( + model, + criterion, + optimizer, + data_loader, + device, + epoch, + args, + model_ema, + scaler, + ) + lr_scheduler.step() + evaluate(model, criterion, data_loader_test, device=device) + if model_ema: + evaluate( + model_ema, criterion, data_loader_test, device=device, log_suffix="EMA" + ) + if args.output_dir: + checkpoint = { + "model": model_without_ddp.state_dict(), + "optimizer": optimizer.state_dict(), + "lr_scheduler": lr_scheduler.state_dict(), + "epoch": epoch, + "args": args, + } + if sparsifier: + checkpoint["sparsifier"] = sparsifier.state_dict() + if model_ema: + checkpoint["model_ema"] = model_ema.state_dict() + if scaler: + checkpoint["scaler"] = scaler.state_dict() + utils.save_on_master( + checkpoint, os.path.join(args.output_dir, f"model_{epoch}.pth") + ) + utils.save_on_master( + checkpoint, os.path.join(args.output_dir, "checkpoint.pth") + ) + + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print(f"Training time {total_time_str}") + + +if __name__ == "__main__": + args = utils.get_args_parser(train=True).parse_args() + main(args) diff --git a/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/utils.py b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d00cdea5108bebe9ea2f2ad87aef16a7d59acf3b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/prototype/sparsity/superblock/utils.py @@ -0,0 +1,1297 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +import argparse +import copy +import datetime +import errno +import hashlib +import math +import os +import time +from collections import OrderedDict, defaultdict, deque +from typing import List, Optional, Tuple + +import torch +from torchvision.transforms import autoaugment, transforms +from torchvision.transforms import functional as F +from torchvision.transforms.functional import InterpolationMode + +from torchao.prototype.sparsity.sparsifier.weight_norm_sparsifier import ( + WeightNormSparsifier, +) +from torchao.prototype.sparsity.superblock.blocksparse import block_sparse_weight +from torchao.prototype.sparsity.superblock.supermask import ( + SupermaskLinear, + apply_supermask, +) +from torchao.quantization import int8_dynamic_activation_int8_weight, quantize_ +from torchao.sparsity import semi_sparse_weight, sparsify_ + + +def get_args_parser(train=False, evaluate=False, benchmark=False): + assert sum([train, evaluate, benchmark]) == 1, ( + "One and only one of training, evaluation, or benchmark can be true" + ) + + # Shared common args + parser = argparse.ArgumentParser( + description="SuperBlock Imagenet Training/Evaluation/Benchmarking Script", + add_help=True, + ) + parser.add_argument("--data-path", type=str, help="IMAGENET dataset path") + parser.add_argument( + "--model", + default="vit_b_16", + choices=["vit_b_16", "vit_h_14"], + type=str, + help="ViT base model", + ) + parser.add_argument( + "--device", default="cuda", type=str, help="device (Default: cuda)" + ) + parser.add_argument( + "-b", "--batch-size", default=32, type=int, help="per device batch size" + ) + parser.add_argument( + "--val-crop-size", + default=224, + type=int, + help="the central crop size used for validation (default: 224)", + ) + parser.add_argument( + "--sparsity", + choices=["bsr", "semi_structured"], + default=None, + help="weight sparsification to apply", + ) + parser.add_argument( + "--bsr", + type=int, + nargs="?", + const=256, + default=None, + help="Convert sparsified weights to BSR format with optional block size (default: 256)", + ) + parser.add_argument("--sparsity-linear", type=float, default=0.0) + parser.add_argument("--sparsity-conv1x1", type=float, default=0.0) + parser.add_argument("--sparsity-conv", type=float, default=0.0) + parser.add_argument( + "--skip-last-layer-sparsity", + action="store_true", + help="Skip applying sparsity to the last linear layer (for vit only)", + ) + parser.add_argument( + "--skip-first-transformer-sparsity", + action="store_true", + help="Skip applying sparsity to the first transformer layer (for vit only)", + ) + parser.add_argument( + "--quantization", action="store_true", help="Run with int8 dynamic quantization" + ) + parser.add_argument( + "--weights", default=None, type=str, help="the weights enum name to load" + ) + parser.add_argument( + "--weights-path", + type=str, + help="optional checkpoint to load weights after intialization", + ) + parser.add_argument( + "--header", action="store_true", help="Print header for first run" + ) + + # Eval a subset of training args + # lots of training args + if train or evaluate: + parser.add_argument( + "-j", + "--workers", + default=16, + type=int, + metavar="N", + help="number of data loading workers", + ) + parser.add_argument( + "--accumulation-steps", + default=1, + type=int, + help="Number of steps to accumulate gradients over", + ) + parser.add_argument( + "--epochs", + default=90, + type=int, + metavar="N", + help="number of total epochs to run", + ) + parser.add_argument("--opt", default="sgd", type=str, help="optimizer") + parser.add_argument( + "--lr", default=0.1, type=float, help="initial learning rate" + ) + parser.add_argument( + "--momentum", default=0.9, type=float, metavar="M", help="momentum" + ) + parser.add_argument( + "--wd", + "--weight-decay", + default=1e-4, + type=float, + metavar="W", + help="weight decay", + dest="weight_decay", + ) + parser.add_argument( + "--norm-weight-decay", + default=None, + type=float, + help="weight decay for Normalization layers (default: None, same value as --wd)", + ) + parser.add_argument( + "--bias-weight-decay", + default=None, + type=float, + help="weight decay for bias parameters of all layers (default: None, same value as --wd)", + ) + parser.add_argument( + "--transformer-embedding-decay", + default=None, + type=float, + help="weight decay for embedding parameters for vision transformer models (default: None, same value as --wd)", + ) + parser.add_argument( + "--label-smoothing", + default=0.0, + type=float, + help="label smoothing (default: 0.0)", + dest="label_smoothing", + ) + parser.add_argument( + "--mixup-alpha", default=0.0, type=float, help="mixup alpha (default: 0.0)" + ) + parser.add_argument( + "--cutmix-alpha", + default=0.0, + type=float, + help="cutmix alpha (default: 0.0)", + ) + parser.add_argument( + "--lr-scheduler", + default="steplr", + type=str, + help="the lr scheduler (default: steplr)", + ) + parser.add_argument( + "--lr-warmup-epochs", + default=0, + type=int, + help="the number of epochs to warmup (default: 0)", + ) + parser.add_argument( + "--lr-warmup-method", + default="constant", + type=str, + help="the warmup method (default: constant)", + ) + parser.add_argument( + "--lr-warmup-decay", default=0.01, type=float, help="the decay for lr" + ) + parser.add_argument( + "--lr-step-size", + default=30, + type=int, + help="decrease lr every step-size epochs", + ) + parser.add_argument( + "--lr-gamma", + default=0.1, + type=float, + help="decrease lr by a factor of lr-gamma", + ) + parser.add_argument( + "--lr-min", + default=0.0, + type=float, + help="minimum lr of lr schedule (default: 0.0)", + ) + parser.add_argument( + "--print-freq", default=10, type=int, help="print frequency" + ) + parser.add_argument( + "--output-dir", default=".", type=str, help="path to save outputs" + ) + parser.add_argument( + "--resume", + action="store_true", + help='Resumes training from latest available checkpoint ("model_.pth")', + ) + parser.add_argument( + "--start-epoch", default=0, type=int, metavar="N", help="start epoch" + ) + parser.add_argument( + "--cache-dataset", + dest="cache_dataset", + help="Cache the datasets for quicker initialization. It also serializes the transforms", + action="store_true", + ) + parser.add_argument( + "--sync-bn", dest="sync_bn", help="Use sync batch norm", action="store_true" + ) + parser.add_argument( + "--auto-augment", + default=None, + type=str, + help="auto augment policy (default: None)", + ) + parser.add_argument( + "--ra-magnitude", + default=9, + type=int, + help="magnitude of auto augment policy", + ) + parser.add_argument( + "--augmix-severity", default=3, type=int, help="severity of augmix policy" + ) + parser.add_argument( + "--random-erase", + default=0.0, + type=float, + help="random erasing probability (default: 0.0)", + ) + # Mixed precision training parameters + parser.add_argument( + "--amp", + action="store_true", + help="Use torch.cuda.amp for mixed precision training", + ) + # distributed training parameters + parser.add_argument( + "--world-size", default=1, type=int, help="number of distributed processes" + ) + parser.add_argument( + "--dist-url", + default="env://", + type=str, + help="url used to set up distributed training", + ) + parser.add_argument( + "--model-ema", + action="store_true", + help="enable tracking Exponential Moving Average of model parameters", + ) + parser.add_argument( + "--model-ema-steps", + type=int, + default=32, + help="the number of iterations that controls how often to update the EMA model (default: 32)", + ) + parser.add_argument( + "--model-ema-decay", + type=float, + default=0.99998, + help="decay factor for Exponential Moving Average of model parameters (default: 0.99998)", + ) + parser.add_argument( + "--use-deterministic-algorithms", + action="store_true", + help="Forces the use of deterministic algorithms only.", + ) + parser.add_argument( + "--interpolation", + default="bilinear", + type=str, + help="the interpolation method (default: bilinear)", + ) + parser.add_argument( + "--val-resize-size", + default=256, + type=int, + help="the resize size used for validation (default: 256)", + ) + parser.add_argument( + "--train-crop-size", + default=224, + type=int, + help="the random crop size used for training (default: 224)", + ) + parser.add_argument( + "--clip-grad-norm", + default=None, + type=float, + help="the maximum gradient norm (default None)", + ) + parser.add_argument( + "--ra-reps", + default=3, + type=int, + help="number of repetitions for Repeated Augmentation (default: 3)", + ) + parser.add_argument( + "--meta", action="store_true", help="Use Meta internal imagenet structure" + ) + + if benchmark: + parser.add_argument( + "--dtype", + choices=["float32", "bfloat16", "float16"], + help="Data type", + default="bfloat16", + ) + parser.add_argument( + "--tune-kernel-params", + action="store_true", + help="Tune kernel params for BSR", + ) + parser.add_argument( + "--profile", action="store_true", help="Dump Prefetto trace" + ) + + return parser + + +# filter functions +def mlp_0_only(mod, name): + return isinstance(mod, torch.nn.Linear) and "mlp.0" in name + + +def mlp_3_only(mod, name): + return isinstance(mod, torch.nn.Linear) and "mlp.3" in name + + +def mlp_only(mod, name): + return isinstance(mod, torch.nn.Linear) and "mlp" in name + + +def superblock_only(mod, name): + return isinstance(mod, SupermaskLinear) and "mlp" in name + + +def mlp_only_with_args( + mod, name, skip_last_layer_sparsity=False, skip_first_transformer_sparsity=False +): + if skip_last_layer_sparsity and "heads.head" in name: + return False + if skip_first_transformer_sparsity and "encoder.layers.encoder_layer_0" in name: + return False + if isinstance(mod, torch.nn.Linear) and "mlp" in name: + return True + return False + + +### Custom sparsification utils +def apply_sparsity(model): + for name, module in model.named_modules(): + if isinstance(module, SupermaskLinear) and "mlp" in name: + module.sparsify_offline() + + +def accelerate_with_sparsity(model, args): + if args.sparsity == "bsr": + apply_sparsity(model) + if args.quantization: + from torchao.dtypes import BlockSparseLayout + + quantize_( + model, + int8_dynamic_activation_int8_weight( + _layout=BlockSparseLayout(blocksize=args.bsr) + ), + superblock_only, + ) + else: + assert args.bsr is not None, "BSR requires a block size" + sparsify_(model, block_sparse_weight(blocksize=args.bsr), superblock_only) + elif args.sparsity == "semi_structured": + if args.quantization: + from torchao.dtypes import SemiSparseLayout + + quantize_( + model, + int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()), + mlp_0_only, + ) + sparsify_(model, semi_sparse_weight(), mlp_3_only) + else: + sparsify_(model, semi_sparse_weight(), mlp_only) + else: + if args.quantization: + quantize_(model, int8_dynamic_activation_int8_weight(), mlp_only) + + +def simulate_sparsity(model, args): + if args.sparsity == "bsr": + apply_supermask( + model, + linear_sparsity=args.sparsity_linear, + linear_sp_tilesize=args.bsr, + conv1x1_sparsity=args.sparsity_conv1x1, + conv1x1_sp_tilesize=args.bsr, + conv_sparsity=args.sparsity_conv, + conv_sp_tilesize=args.bsr, + skip_last_layer_sparsity=args.skip_last_layer_sparsity, + skip_first_transformer_sparsity=args.skip_first_transformer_sparsity, + device=args.device, + verbose=False, + ) + elif args.sparsity == "semi_structured": + sparse_config = [] + for name, mod in model.named_modules(): + if mlp_only_with_args( + mod, + name, + skip_first_transformer_sparsity=args.skip_first_transformer_sparsity, + skip_last_layer_sparsity=args.skip_last_layer_sparsity, + ): + sparse_config.append({"tensor_fqn": f"{name}.weight"}) + + sparsifier = WeightNormSparsifier( + sparsity_level=1.0, sparse_block_shape=(1, 4), zeros_per_block=2 + ) + sparsifier.prepare(model, sparse_config) + sparsifier.step() + return sparsifier + + +# ------------------------------------------------------------ +# The following code contains torchvision reference code, +# largely copied from: https://github.com/pytorch/vision/tree/main/references/classification +# Please open issues in the original repository if you have questions. + + +class SmoothedValue: + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + t = reduce_across_processes([self.count, self.total]) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self): + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + return self.total / self.count + + @property + def max(self): + return max(self.deque) + + @property + def value(self): + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value, + ) + + +class MetricLogger: + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, **kwargs): + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{attr}'" + ) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + loss_str.append(f"{name}: {str(meter)}") + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None): + i = 0 + if not header: + header = "" + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt="{avg:.4f}") + data_time = SmoothedValue(fmt="{avg:.4f}") + space_fmt = ":" + str(len(str(len(iterable)))) + "d" + if torch.cuda.is_available(): + log_msg = self.delimiter.join( + [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + "max mem: {memory:.0f}", + ] + ) + else: + log_msg = self.delimiter.join( + [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + ] + ) + MB = 1024.0 * 1024.0 + for obj in iterable: + data_time.update(time.time() - end) + yield obj + iter_time.update(time.time() - end) + if i % print_freq == 0: + eta_seconds = iter_time.global_avg * (len(iterable) - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print( + log_msg.format( + i, + len(iterable), + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB, + ) + ) + else: + print( + log_msg.format( + i, + len(iterable), + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + ) + ) + i += 1 + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print(f"{header} Total time: {total_time_str}") + + +class ExponentialMovingAverage(torch.optim.swa_utils.AveragedModel): + """Maintains moving averages of model parameters using an exponential decay. + ``ema_avg = decay * avg_model_param + (1 - decay) * model_param`` + `torch.optim.swa_utils.AveragedModel `_ + is used to compute the EMA. + """ + + def __init__(self, model, decay, device="cpu"): + def ema_avg(avg_model_param, model_param, num_averaged): + return decay * avg_model_param + (1 - decay) * model_param + + super().__init__(model, device, ema_avg, use_buffers=True) + + +def accuracy(output, target, topk=(1,)): + """Computes the accuracy over the k top predictions for the specified values of k""" + with torch.inference_mode(): + maxk = max(topk) + batch_size = target.size(0) + if target.ndim == 2: + target = target.max(dim=1)[1] + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target[None]) + + res = [] + for k in topk: + correct_k = correct[:k].flatten().sum(dtype=torch.float32) + res.append(correct_k * (100.0 / batch_size)) + return res + + +def mkdir(path): + try: + os.makedirs(path) + except OSError as e: + if e.errno != errno.EEXIST: + raise + + +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print + + +def is_dist_avail_and_initialized(): + if not torch.distributed.is_available(): + return False + if not torch.distributed.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return torch.distributed.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return torch.distributed.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def save_on_master(*args, **kwargs): + if is_main_process(): + torch.save(*args, **kwargs) + + +def init_distributed_mode(args): + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ["WORLD_SIZE"]) + args.gpu = int(os.environ["LOCAL_RANK"]) + elif "SLURM_PROCID" in os.environ: + args.rank = int(os.environ["SLURM_PROCID"]) + args.gpu = args.rank % torch.cuda.device_count() + elif hasattr(args, "rank"): + pass + else: + print("Not using distributed mode") + args.distributed = False + return + + args.distributed = True + + torch.cuda.set_device(args.gpu) + args.dist_backend = "nccl" + print(f"| distributed init (rank {args.rank})", flush=True) + torch.distributed.init_process_group( + backend=args.dist_backend, + init_method=args.dist_url, + world_size=args.world_size, + rank=args.rank, + ) + torch.distributed.barrier() + setup_for_distributed(args.rank == 0) + + +def average_checkpoints(inputs): + """Loads checkpoints from inputs and returns a model with averaged weights. Original implementation taken from: + https://github.com/pytorch/fairseq/blob/a48f235636557b8d3bc4922a6fa90f3a0fa57955/scripts/average_checkpoints.py#L16 + + Args: + inputs (List[str]): An iterable of string paths of checkpoints to load from. + Returns: + A dict of string keys mapping to various values. The 'model' key + from the returned dict should correspond to an OrderedDict mapping + string parameter names to torch Tensors. + """ + params_dict = OrderedDict() + params_keys = None + new_state = None + num_models = len(inputs) + for fpath in inputs: + with open(fpath, "rb") as f: + state = torch.load( + f, + map_location=( + lambda s, _: torch.serialization.default_restore_location(s, "cpu") + ), + ) + # Copies over the settings from the first checkpoint + if new_state is None: + new_state = state + model_params = state["model"] + model_params_keys = list(model_params.keys()) + if params_keys is None: + params_keys = model_params_keys + elif params_keys != model_params_keys: + raise KeyError( + f"For checkpoint {f}, expected list of params: {params_keys}, but found: {model_params_keys}" + ) + for k in params_keys: + p = model_params[k] + if isinstance(p, torch.HalfTensor): + p = p.float() + if k not in params_dict: + params_dict[k] = p.clone() + # NOTE: clone() is needed in case of p is a shared parameter + else: + params_dict[k] += p + averaged_params = OrderedDict() + for k, v in params_dict.items(): + averaged_params[k] = v + if averaged_params[k].is_floating_point(): + averaged_params[k].div_(num_models) + else: + averaged_params[k] //= num_models + new_state["model"] = averaged_params + return new_state + + +def store_model_weights(model, checkpoint_path, checkpoint_key="model", strict=True): + """ + This method can be used to prepare weights files for new models. It receives as + input a model architecture and a checkpoint from the training script and produces + a file with the weights ready for release. + + Examples: + from torchvision import models as M + + # Classification + model = M.mobilenet_v3_large(weights=None) + print(store_model_weights(model, './class.pth')) + + # Quantized Classification + model = M.quantization.mobilenet_v3_large(weights=None, quantize=False) + model.fuse_model(is_qat=True) + model.qconfig = torch.ao.quantization.get_default_qat_qconfig('qnnpack') + _ = torch.ao.quantization.prepare_qat(model, inplace=True) + print(store_model_weights(model, './qat.pth')) + + # Object Detection + model = M.detection.fasterrcnn_mobilenet_v3_large_fpn(weights=None, weights_backbone=None) + print(store_model_weights(model, './obj.pth')) + + # Segmentation + model = M.segmentation.deeplabv3_mobilenet_v3_large(weights=None, weights_backbone=None, aux_loss=True) + print(store_model_weights(model, './segm.pth', strict=False)) + + Args: + model (pytorch.nn.Module): The model on which the weights will be loaded for validation purposes. + checkpoint_path (str): The path of the checkpoint we will load. + checkpoint_key (str, optional): The key of the checkpoint where the model weights are stored. + Default: "model". + strict (bool): whether to strictly enforce that the keys + in :attr:`state_dict` match the keys returned by this module's + :meth:`~torch.nn.Module.state_dict` function. Default: ``True`` + + Returns: + output_path (str): The location where the weights are saved. + """ + # Store the new model next to the checkpoint_path + checkpoint_path = os.path.abspath(checkpoint_path) + output_dir = os.path.dirname(checkpoint_path) + + # Deep copy to avoid side-effects on the model object. + model = copy.deepcopy(model) + checkpoint = torch.load(checkpoint_path, map_location="cpu") + + # Load the weights to the model to validate that everything works + # and remove unnecessary weights (such as auxiliaries, etc) + if checkpoint_key == "model_ema": + del checkpoint[checkpoint_key]["n_averaged"] + torch.nn.modules.utils.consume_prefix_in_state_dict_if_present( + checkpoint[checkpoint_key], "module." + ) + model.load_state_dict(checkpoint[checkpoint_key], strict=strict) + + tmp_path = os.path.join(output_dir, str(model.__hash__())) + torch.save(model.state_dict(), tmp_path) + + sha256_hash = hashlib.sha256() + with open(tmp_path, "rb") as f: + # Read and update hash string value in blocks of 4K + for byte_block in iter(lambda: f.read(4096), b""): + sha256_hash.update(byte_block) + hh = sha256_hash.hexdigest() + + output_path = os.path.join(output_dir, "weights-" + str(hh[:8]) + ".pth") + os.replace(tmp_path, output_path) + + return output_path + + +def reduce_across_processes(val): + if not is_dist_avail_and_initialized(): + # nothing to sync, but we still convert to tensor for consistency with the distributed case. + return torch.tensor(val) + + t = torch.tensor(val, device="cuda") + torch.distributed.barrier() + torch.distributed.all_reduce(t) + return t + + +def set_weight_decay( + model: torch.nn.Module, + weight_decay: float, + norm_weight_decay: Optional[float] = None, + norm_classes: Optional[List[type]] = None, + custom_keys_weight_decay: Optional[List[Tuple[str, float]]] = None, +): + if not norm_classes: + norm_classes = [ + torch.nn.modules.batchnorm._BatchNorm, + torch.nn.LayerNorm, + torch.nn.GroupNorm, + torch.nn.modules.instancenorm._InstanceNorm, + torch.nn.LocalResponseNorm, + ] + norm_classes = tuple(norm_classes) + + params = { + "other": [], + "norm": [], + } + params_weight_decay = { + "other": weight_decay, + "norm": norm_weight_decay, + } + custom_keys = [] + if custom_keys_weight_decay is not None: + for key, weight_decay in custom_keys_weight_decay: + params[key] = [] + params_weight_decay[key] = weight_decay + custom_keys.append(key) + + def _add_params(module, prefix=""): + for name, p in module.named_parameters(recurse=False): + if not p.requires_grad: + continue + is_custom_key = False + for key in custom_keys: + target_name = ( + f"{prefix}.{name}" if prefix != "" and "." in key else name + ) + if key == target_name: + params[key].append(p) + is_custom_key = True + break + if not is_custom_key: + if norm_weight_decay is not None and isinstance(module, norm_classes): + params["norm"].append(p) + else: + params["other"].append(p) + + for child_name, child_module in module.named_children(): + child_prefix = f"{prefix}.{child_name}" if prefix != "" else child_name + _add_params(child_module, prefix=child_prefix) + + _add_params(model) + + param_groups = [] + for key in params: + if len(params[key]) > 0: + param_groups.append( + {"params": params[key], "weight_decay": params_weight_decay[key]} + ) + return param_groups + + +# Presets for ImageNet training/eval taken from: https://github.com/pytorch/vision/blob/main/references/classification/presets.py + + +class ClassificationPresetTrain: + def __init__( + self, + *, + crop_size, + mean=(0.485, 0.456, 0.406), + std=(0.229, 0.224, 0.225), + interpolation=InterpolationMode.BILINEAR, + hflip_prob=0.5, + auto_augment_policy=None, + ra_magnitude=9, + augmix_severity=3, + random_erase_prob=0.0, + ): + trans = [transforms.RandomResizedCrop(crop_size, interpolation=interpolation)] + if hflip_prob > 0: + trans.append(transforms.RandomHorizontalFlip(hflip_prob)) + if auto_augment_policy is not None: + if auto_augment_policy == "ra": + trans.append( + autoaugment.RandAugment( + interpolation=interpolation, magnitude=ra_magnitude + ) + ) + elif auto_augment_policy == "ta_wide": + trans.append( + autoaugment.TrivialAugmentWide(interpolation=interpolation) + ) + elif auto_augment_policy == "augmix": + trans.append( + autoaugment.AugMix( + interpolation=interpolation, severity=augmix_severity + ) + ) + else: + aa_policy = autoaugment.AutoAugmentPolicy(auto_augment_policy) + trans.append( + autoaugment.AutoAugment( + policy=aa_policy, interpolation=interpolation + ) + ) + trans.extend( + [ + transforms.PILToTensor(), + transforms.ConvertImageDtype(torch.float), + transforms.Normalize(mean=mean, std=std), + ] + ) + if random_erase_prob > 0: + trans.append(transforms.RandomErasing(p=random_erase_prob)) + + self.transforms = transforms.Compose(trans) + + def __call__(self, img): + return self.transforms(img) + + +class ClassificationPresetEval: + def __init__( + self, + *, + crop_size, + resize_size=256, + mean=(0.485, 0.456, 0.406), + std=(0.229, 0.224, 0.225), + interpolation=InterpolationMode.BILINEAR, + ): + self.transforms = transforms.Compose( + [ + transforms.Resize(resize_size, interpolation=interpolation), + transforms.CenterCrop(crop_size), + transforms.PILToTensor(), + transforms.ConvertImageDtype(torch.float), + transforms.Normalize(mean=mean, std=std), + ] + ) + + def __call__(self, img): + return self.transforms(img) + + +# transforms taken from: https://github.com/pytorch/vision/blob/main/references/classification/transforms.py + + +class RandomMixup(torch.nn.Module): + """Randomly apply Mixup to the provided batch and targets. + The class implements the data augmentations as described in the paper + `"mixup: Beyond Empirical Risk Minimization" `_. + + Args: + num_classes (int): number of classes used for one-hot encoding. + p (float): probability of the batch being transformed. Default value is 0.5. + alpha (float): hyperparameter of the Beta distribution used for mixup. + Default value is 1.0. + inplace (bool): boolean to make this transform inplace. Default set to False. + """ + + def __init__( + self, + num_classes: int, + p: float = 0.5, + alpha: float = 1.0, + inplace: bool = False, + ) -> None: + super().__init__() + + if num_classes < 1: + raise ValueError( + f"Please provide a valid positive value for the num_classes. Got num_classes={num_classes}" + ) + + if alpha <= 0: + raise ValueError("Alpha param can't be zero.") + + self.num_classes = num_classes + self.p = p + self.alpha = alpha + self.inplace = inplace + + def forward( + self, batch: torch.Tensor, target: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + batch (Tensor): Float tensor of size (B, C, H, W) + target (Tensor): Integer tensor of size (B, ) + + Returns: + Tensor: Randomly transformed batch. + """ + if batch.ndim != 4: + raise ValueError(f"Batch ndim should be 4. Got {batch.ndim}") + if target.ndim != 1: + raise ValueError(f"Target ndim should be 1. Got {target.ndim}") + if not batch.is_floating_point(): + raise TypeError(f"Batch dtype should be a float tensor. Got {batch.dtype}.") + if target.dtype != torch.int64: + raise TypeError(f"Target dtype should be torch.int64. Got {target.dtype}") + + if not self.inplace: + batch = batch.clone() + target = target.clone() + + if target.ndim == 1: + target = torch.nn.functional.one_hot( + target, num_classes=self.num_classes + ).to(dtype=batch.dtype) + + if torch.rand(1).item() >= self.p: + return batch, target + + # It's faster to roll the batch by one instead of shuffling it to create image pairs + batch_rolled = batch.roll(1, 0) + target_rolled = target.roll(1, 0) + + # Implemented as on mixup paper, page 3. + lambda_param = float( + torch._sample_dirichlet(torch.tensor([self.alpha, self.alpha]))[0] + ) + batch_rolled.mul_(1.0 - lambda_param) + batch.mul_(lambda_param).add_(batch_rolled) + + target_rolled.mul_(1.0 - lambda_param) + target.mul_(lambda_param).add_(target_rolled) + + return batch, target + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"num_classes={self.num_classes}" + f", p={self.p}" + f", alpha={self.alpha}" + f", inplace={self.inplace}" + f")" + ) + return s + + +class RandomCutmix(torch.nn.Module): + """Randomly apply Cutmix to the provided batch and targets. + The class implements the data augmentations as described in the paper + `"CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features" + `_. + + Args: + num_classes (int): number of classes used for one-hot encoding. + p (float): probability of the batch being transformed. Default value is 0.5. + alpha (float): hyperparameter of the Beta distribution used for cutmix. + Default value is 1.0. + inplace (bool): boolean to make this transform inplace. Default set to False. + """ + + def __init__( + self, + num_classes: int, + p: float = 0.5, + alpha: float = 1.0, + inplace: bool = False, + ) -> None: + super().__init__() + if num_classes < 1: + raise ValueError( + "Please provide a valid positive value for the num_classes." + ) + if alpha <= 0: + raise ValueError("Alpha param can't be zero.") + + self.num_classes = num_classes + self.p = p + self.alpha = alpha + self.inplace = inplace + + def forward( + self, batch: torch.Tensor, target: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + batch (Tensor): Float tensor of size (B, C, H, W) + target (Tensor): Integer tensor of size (B, ) + + Returns: + Tensor: Randomly transformed batch. + """ + if batch.ndim != 4: + raise ValueError(f"Batch ndim should be 4. Got {batch.ndim}") + if target.ndim != 1: + raise ValueError(f"Target ndim should be 1. Got {target.ndim}") + if not batch.is_floating_point(): + raise TypeError(f"Batch dtype should be a float tensor. Got {batch.dtype}.") + if target.dtype != torch.int64: + raise TypeError(f"Target dtype should be torch.int64. Got {target.dtype}") + + if not self.inplace: + batch = batch.clone() + target = target.clone() + + if target.ndim == 1: + target = torch.nn.functional.one_hot( + target, num_classes=self.num_classes + ).to(dtype=batch.dtype) + + if torch.rand(1).item() >= self.p: + return batch, target + + # It's faster to roll the batch by one instead of shuffling it to create image pairs + batch_rolled = batch.roll(1, 0) + target_rolled = target.roll(1, 0) + + # Implemented as on cutmix paper, page 12 (with minor corrections on typos). + lambda_param = float( + torch._sample_dirichlet(torch.tensor([self.alpha, self.alpha]))[0] + ) + _, H, W = F.get_dimensions(batch) + + r_x = torch.randint(W, (1,)) + r_y = torch.randint(H, (1,)) + + r = 0.5 * math.sqrt(1.0 - lambda_param) + r_w_half = int(r * W) + r_h_half = int(r * H) + + x1 = int(torch.clamp(r_x - r_w_half, min=0)) + y1 = int(torch.clamp(r_y - r_h_half, min=0)) + x2 = int(torch.clamp(r_x + r_w_half, max=W)) + y2 = int(torch.clamp(r_y + r_h_half, max=H)) + + batch[:, :, y1:y2, x1:x2] = batch_rolled[:, :, y1:y2, x1:x2] + lambda_param = float(1.0 - (x2 - x1) * (y2 - y1) / (W * H)) + + target_rolled.mul_(1.0 - lambda_param) + target.mul_(lambda_param).add_(target_rolled) + + return batch, target + + def __repr__(self) -> str: + s = ( + f"{self.__class__.__name__}(" + f"num_classes={self.num_classes}" + f", p={self.p}" + f", alpha={self.alpha}" + f", inplace={self.inplace}" + f")" + ) + return s + + +# RA Sampler implementaion taken from: https://github.com/pytorch/vision/blob/main/references/classification/sampler.py + + +class RASampler(torch.utils.data.Sampler): + """Sampler that restricts data loading to a subset of the dataset for distributed, + with repeated augmentation. + It ensures that different each augmented version of a sample will be visible to a + different process (GPU). + Heavily based on 'torch.utils.data.DistributedSampler'. + + This is borrowed from the DeiT Repo: + https://github.com/facebookresearch/deit/blob/main/samplers.py + """ + + def __init__( + self, dataset, num_replicas=None, rank=None, shuffle=True, seed=0, repetitions=3 + ): + if num_replicas is None: + if not torch.distributed.is_available(): + raise RuntimeError("Requires distributed package to be available!") + num_replicas = torch.distributed.get_world_size() + if rank is None: + if not torch.distributed.is_available(): + raise RuntimeError("Requires distributed package to be available!") + rank = torch.distributed.get_rank() + self.dataset = dataset + self.num_replicas = num_replicas + self.rank = rank + self.epoch = 0 + self.num_samples = int( + math.ceil(len(self.dataset) * float(repetitions) / self.num_replicas) + ) + self.total_size = self.num_samples * self.num_replicas + self.num_selected_samples = int( + math.floor(len(self.dataset) // 256 * 256 / self.num_replicas) + ) + self.shuffle = shuffle + self.seed = seed + self.repetitions = repetitions + + def __iter__(self): + if self.shuffle: + # Deterministically shuffle based on epoch + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + indices = torch.randperm(len(self.dataset), generator=g).tolist() + else: + indices = list(range(len(self.dataset))) + + # Add extra samples to make it evenly divisible + indices = [ele for ele in indices for i in range(self.repetitions)] + indices += indices[: (self.total_size - len(indices))] + assert len(indices) == self.total_size + + # Subsample + indices = indices[self.rank : self.total_size : self.num_replicas] + assert len(indices) == self.num_samples + + return iter(indices[: self.num_selected_samples]) + + def __len__(self): + return self.num_selected_samples + + def set_epoch(self, epoch): + self.epoch = epoch diff --git a/lib/python3.12/site-packages/torchao/quantization/GPTQ.py b/lib/python3.12/site-packages/torchao/quantization/GPTQ.py new file mode 100644 index 0000000000000000000000000000000000000000..94fccc7bf1af07d7e88c1a4aa99d97b731352c82 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/GPTQ.py @@ -0,0 +1,1311 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from typing import Any, Callable, Dict, List, Optional, Type + +import torch +import torch.fx as fx +import torch.nn as nn +import torch.nn.functional as F +from torch.utils._pytree import tree_flatten, tree_unflatten + +from torchao.dtypes.utils import is_device +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_3, + TORCH_VERSION_AT_LEAST_2_6, + find_multiple, +) + +from .quant_primitives import ( + MappingType, + dequantize_affine, +) +from .unified import Quantizer +from .utils import ( + _MultiInput, + get_group_qparams_symmetric, + get_groupwise_affine_qparams, + group_quantize_tensor_symmetric, + groupwise_affine_dequantize_tensor_from_qparams, + groupwise_affine_quantize_tensor, + groupwise_affine_quantize_tensor_from_qparams, + pack_tinygemm_scales_and_zeros, + per_token_dynamic_quant, +) + +aten = torch.ops.aten + +add_ons = [] + +if TORCH_VERSION_AT_LEAST_2_3: + add_ons += ["Int8DynActInt4WeightQuantizer", "Int8DynActInt4WeightGPTQQuantizer"] + + +__all__ = [ + "Int4WeightOnlyGPTQQuantizer", + "Int4WeightOnlyQuantizer", +] + add_ons + + +class GenericGPTQRunner(fx.Interpreter): + """ + This is a generic GPTQ runner that takes an existing model and applies GPTQ. + It uses torch._dynamo.export to obtain a graph of the model and then hooks + into function calls and when it detects a linear, it applies GPTQ to the weight + given the calibration of inputs passed in at initialization. It puts the results + into the state_dict so that the quantized model weights/qparams can be loaded + directly into the model. + + intended to be used in concert with a GPTQQuantizer class to define the quantization mode. + """ + + def __init__( + self, + model, + inputs: _MultiInput, + blocksize=128, + percdamp=0.01, + groupsize=128, + ): + self.id_to_name = { + id(value): name for name, value in dict(model.named_parameters()).items() + } + + # trace model for one input + one_input = [multi.values[0].cpu() for multi in inputs] # pyre-ignore[16] + # needed for GPTQ on the torchao llama model + import torchao + + torchao._models.llama.model.use_index_put_for_kv_cache = True + exported_model = torch._dynamo.export( + model.cpu(), aten_graph=True, pre_dispatch=True, tracing_mode="fake" + )(*one_input) + super().__init__(exported_model.graph_module) + + self.new_state_dict = model.state_dict() + + self.blocksize = blocksize + + self.percdamp = percdamp + + self.groupsize = groupsize + self.inputs = inputs + self.gptq_done = False + self.debug = False + + def configure_quantization_mode( + self, + get_qparams_func, + quantize_func, + dequantize_func, + combine_qparams_list_func, + make_names_and_values_dict_func, + skip_layer_func, + act_fake_quant_func=None, + ): + # these functions need to already be curried with all inputs other than weight, qparams + + self.get_qparams_func = ( + get_qparams_func # accepts [2d weight tensor], outputs qparams. + ) + + self.quantize_func = quantize_func # accepts [2d weight tensor], [qparams], outputs a 2d quantized tensor of desired dtype + + self.dequantize_func = dequantize_func + # accepts [quantized] tensor and [qparams], outputs a 2d dequantized tensor of type float, + # assumes this output .to(w_orig_dtype) is ~eventual desired dequant behavior + + # `combine_qparams_list_func`. + self.combine_qparams_list_func = combine_qparams_list_func + # accepts [`list` of qparams] from quantizing one group at a time, + # outputs a qparams object that could be passed into quant/dequantize_func + + self.skip_layer_func = skip_layer_func # accepts [weight tensor], outputs a bool on whether or not to apply gptq to this layer + + # `make_names_and_values_dict_func`. + self.make_names_and_values_dict_func = make_names_and_values_dict_func # accepts [2d quantized tensor], [qparams], returns a dict of names, values to put in state_dict + # note any final packing for storage should happen here + + # `act_fake_quant_func` + if act_fake_quant_func is None: + self.act_fake_quant_func = lambda x: x + else: + self.act_fake_quant_func = act_fake_quant_func # accepts [activation tensor], returns a fake-quantized activation tensor + return self + + def run(self): + assert self.get_qparams_func is not None, ( + "need to configure quantization mode before running" + ) + self.gptq_done = True + super().run(*self.inputs) + + def get_quantized_state_dict(self): + assert self.gptq_done, ( + "need to run GPTQRunner before you can get_quantized_state_dict" + ) + quantized_state_dict = self.new_state_dict + # Don't want to store/load the kv_cache so remove it from the state_dict + del_list = [] + for param_fqn in quantized_state_dict: + if "kv_cache" in param_fqn: + del_list.append(param_fqn) + for param_fqn in del_list: + quantized_state_dict.pop(param_fqn) + return quantized_state_dict + + def call_function(self, target, args, kwargs, already_quantized=False): # noqa: C901 + def tensors_to_cuda(args): + new_args = [] + for x in args: + new_args.append(x.cuda() if isinstance(x, torch.Tensor) else x) + return new_args + + # flatten args and kwargs together + flat_args, spec = tree_flatten((args, kwargs)) + # move all single tensors to cuda, will move _MultiInputs to cuda one at a time + flat_args = tensors_to_cuda(flat_args) + + has_multi_input = _MultiInput in [type(x) for x in flat_args] + if has_multi_input: + # Just some trickery to convert + # [_MultiInput[a, a, a], _MultiInput(b, b, b)] => [a, b], [a, b], [a, b] + multi_input_count = max( + [len(x.values) if isinstance(x, _MultiInput) else 1 for x in flat_args] + ) + transposed_args = list( + zip( + *[ + ( + x.values + if isinstance(x, _MultiInput) + else [x] * multi_input_count + ) + for x in flat_args + ] + ) + ) + else: + transposed_args = [flat_args] + outputs = [] + + # check whether we apply GPTQ to this module + quantize_linear = ( + (target == aten.linear.default) # if its a linear + and id(args[1]) in self.id_to_name # and if we know the layer name + # and we haven't already quantized this layer + and not already_quantized + # and if the skip_layer_func doesn't say we should skip + and not (self.skip_layer_func is not None and self.skip_layer_func(args[1])) + ) # then we will quantize this linear layer/weight + + if quantize_linear: # instantiate variables for GPTQ + H = 0 + total_batches = 0 + + for inp in transposed_args: + inp = tensors_to_cuda(inp) + cur_args, cur_kwargs = tree_unflatten(inp, spec) + + if quantize_linear: # calculate H instead of output (will run the linear eventually with updated weight) + x = cur_args[0].float() + x = self.act_fake_quant_func(x) + shape = x.shape + n = 1 if len(shape) == 2 else shape[0] + H *= total_batches / (total_batches + n) + total_batches += n + x = ((2 / total_batches) ** (1 / 2)) * x.reshape( + -1, shape[-1] + ).t().float() + H += x.matmul(x.t()) + else: + # weight has already been quantized but still need to apply + # activation quant for final calculation + if already_quantized: + cur_args = (self.act_fake_quant_func(cur_args[0]), *cur_args[1:]) + + # get output if its not a linear + out = super().call_function(target, cur_args, cur_kwargs) + if isinstance(out, torch.Tensor): + outputs.append(out.cpu()) + else: + outputs.append(out) + + if quantize_linear: + mod_fqn = ".".join(self.id_to_name[id(args[1])].split(".")[:-1]) + + W = args[1].to(H.device) + + Q, DQ, qparams = self.faster_quant(H, W.detach()) + print(mod_fqn) + + # `make_names_and_values_dict_func`. + names_and_values_dict = self.make_names_and_values_dict_func(Q, qparams) + + # delete old weight + if mod_fqn + ".weight" in self.new_state_dict: + self.new_state_dict.pop(mod_fqn + ".weight") + if len(args) > 2: + self.new_state_dict[mod_fqn + ".bias"] = args[2] + for name, value in names_and_values_dict.items(): + self.new_state_dict[mod_fqn + "." + name] = value + + # run linear with new weight to get corrected output + new_out = self.call_function( + target, (args[0], DQ, *args[2:]), kwargs, already_quantized=True + ) + + if self.debug: + old_out = self.call_function( + target, + (args[0][:2], args[1], *args[2:]), + kwargs, + already_quantized=True, + ) + + def SQNR(x, y): + # TODO: Use of deprecated function torch.norm + return 20 * torch.log10( + torch.linalg.norm(x) / torch.linalg.norm(x - y) + ) + + # `dequantize_func`. + DQ_after = self.dequantize_func(Q, qparams).to(W.dtype) + print( + "SQNR for QDQ (this should be inf)", SQNR(DQ, DQ_after) + ) # matches + print( + "SQNR for weight (can be low)", SQNR(W, DQ.cuda()) + ) # fine to not match + print( + "SQNR for output with GPTQ (hopefully 35+)", + torch.cat( + [ + SQNR(old.cpu(), new.cpu()).unsqueeze(0) + for (old, new) in zip(old_out.values, new_out.values[:2]) + ] + ).mean(), + ) + + # `get_qparams_func`. + qparams2 = self.get_qparams_func(W) + + Q2 = self.quantize_func(W, qparams2) + DQ2 = self.dequantize_func(Q2, qparams2).to(W.dtype) + old_q_out = self.call_function( + target, + (args[0][:2], DQ2, *args[2:]), + kwargs, + already_quantized=True, + ) + + print( + "SQNR for output without GPTQ (should be less than above)", + torch.cat( + [ + SQNR(old.cpu(), old_q.cpu()).unsqueeze(0) + for (old, old_q) in zip(old_out.values, old_q_out.values) + ] + ).mean(), + ) + return new_out + + return _MultiInput(outputs) if has_multi_input else outputs[0] + + def faster_quant(self, H, W): + percdamp = self.percdamp + blocksize = self.blocksize + groupsize = self.groupsize + orig_dtype = W.dtype + W = W.detach().float() + _, columns = W.shape[0], W.shape[1] + device = W.device + + if groupsize == -1: + cur_qparams = self.get_qparams_func(W) + dead = torch.diag(H) == 0 + H[dead, dead] = 1 + W[:, dead] = 0 + + Losses = torch.zeros_like(W) + DQ = torch.zeros_like(W) + + damp = percdamp * torch.mean(torch.diag(H)) + diag = torch.arange(columns, device=device) + H[diag, diag] += damp + H = torch.linalg.cholesky(H) + H = torch.cholesky_inverse(H) + H = torch.linalg.cholesky(H, upper=True) + Hinv = H + + all_qparams = [] + for i1 in range(0, columns, blocksize): + i2 = min(i1 + blocksize, columns) + count = i2 - i1 + W1 = W[:, i1:i2].clone() + DQ1 = torch.zeros_like(W1) + Err1 = torch.zeros_like(W1) + Losses1 = torch.zeros_like(W1) + Hinv1 = Hinv[i1:i2, i1:i2] + for i in range(count): + w = W1[:, i] + d = Hinv1[i, i] + + if groupsize != -1 and (i1 + i) % groupsize == 0: # start of new group + cur_qparams = self.get_qparams_func( + W[:, (i1 + i) : (i1 + i + groupsize)] + ) + all_qparams.append(cur_qparams) + + q = self.quantize_func(w.unsqueeze(1), cur_qparams).flatten() + + # `dequantize_func`. + + dq = self.dequantize_func(q.unsqueeze(1), cur_qparams).flatten() + + DQ1[:, i] = dq + Losses1[:, i] = (w - dq) ** 2 / d**2 + + err1 = (w - dq) / d + W1[:, i:] -= ( + err1.to(Hinv1.dtype).unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0)) + ) + Err1[:, i] = err1 + + DQ[:, i1:i2] = DQ1 + Losses[:, i1:i2] = Losses1 / 2 + + W[:, i2:] -= Err1.to(Hinv.dtype).matmul(Hinv[i1:i2, i2:]) + + torch.cuda.synchronize() + + if all_qparams == []: + all_qparams.append(cur_qparams) + + # convert a list of qparams objects into a single one. enerally by + # concatenating a bunch of n,1 scale/zeros tensors into a n,num_groups tensor + + # `combine_qparams_list_func`. + all_qparams = self.combine_qparams_list_func(all_qparams) + Q = self.quantize_func(DQ, all_qparams) + return Q, DQ.to(orig_dtype), all_qparams + + +class GPTQQuantizer(Quantizer): + """ + This class implements a GPTQ Quantizer that can be used to apply GPTQ to a model in concert with the GenericGPTQRunner class. + Unlike the base Quantizer class, the user does not need to implement the create_quantized_state_dict, instead they have to reimplement + __init__ such that it defines the functions for the quantization mode. User is expected to reimplement convert_for_runtime. + + The following functions (which must be defined in __init__) are used to define the quantization mode for both GPTQ and + create_quantized_state_dict. Here is a description of each function. + + get_qparams_func: + A function that calculates the quantization qparams for an input tensor. + Args: + weight: A 2d weight tensor with non-integer dtype. + Returns: + qparams: it can have any format but will need to be handled by the other defined functions below. + + quantize_func: + A function that applies quantization to an input tensor. It should be noted + that this function needs to be able to handle quantizing the entire weight tensor, a single group, + or a single column. + Args: + weight: A 2d weight tensor with non-integer dtype. + qparams: the output from get_qparams_func + Returns: + quantized_weight: A 2d quantized weight tensor (generally with an integer dtype) + + + dequantize_func: + A function that dequantizes an input quantized weight tensor. It should be noted + that this function needs to be able to handle dequantizing the entire weight tensor, a single group, + or a single column. + Args: + quantized_weight: A 2d quantized weight tensor (generally with an integer dtype) + qparams: the output from get_qparams_func + Returns: + weight: A 2d weight tensor with non-integer dtype. + + act_fake_quant_func (optional): + A function that (dynamically) quantizes activation to input + Args: + input: input Tensor in f32/bf16/f16 + Returns: + output: dynamically quantized and dequantized Tensor (with the same dtype as input) + + combine_qparams_list_func: + A function that combines several qparams into one qparam. + Args: + qparams_list: a list of qparams objects, each obtained by calling get_qparams_func + on a single group from a weight tensor + Returns: + qparams: an object of the same format as the qparams above. + + skip_layer_func: + A function that determines which linear layers should be skipped during GPTQ + Args: + weight: A 2d weight tensor with non-integer dtype. + Returns: + skip: boolean indicating whether layer should be skipped + + make_names_and_values_dict_func: + A function that prepares the qparams and quantized_weight and creates a dictionary indicating how they + should be inserted into the state_dict. Generally any packing of the weight and qparams should be done here. + Args: + quantized_weight: A 2d quantized weight tensor (generally with an integer dtype) + qparams: the output from get_qparams_func + Returns: + names_and_values_dict: a dictionary mapping the name of the parameters of the quantized module to the + corresponding quantized weights and qparams. + """ + + def __init__(self): + assert self.get_qparams_func is not None + + assert self.quantize_func is not None + + assert self.dequantize_func is not None + + assert self.combine_qparams_list_func is not None + + # `make_names_and_values_dict_func`. + assert self.make_names_and_values_dict_func is not None + + @torch.no_grad() + def _create_quantized_state_dict( + self, + model, + inputs, + blocksize, + percdamp, + groupsize, + # `typing.Dict[, ]` to avoid runtime subscripting errors. + ) -> Dict: + print("Tracing model for GPTQ") + GPTQ_runner = GenericGPTQRunner( + model, + inputs, + blocksize, + percdamp, + groupsize, + ).configure_quantization_mode( + self.get_qparams_func, # pyre-ignore[16] + self.quantize_func, # pyre-ignore[16] + self.dequantize_func, # pyre-ignore[16] + self.combine_qparams_list_func, # pyre-ignore[16] + self.make_names_and_values_dict_func, # pyre-ignore[16] + self.skip_layer_func, # pyre-ignore[16] + self.act_fake_quant_func + if hasattr(self, "act_fake_quant_func") + else None, # pyre-ignore[16] + ) + print("Applying GPTQ to weights") + GPTQ_runner.run() + return GPTQ_runner.get_quantized_state_dict() + + def _convert_for_runtime(self, model: torch.nn.Module) -> "nn.Module": + raise NotImplementedError("_convert_for_runtime not implemented") + + @torch.no_grad() + def quantize( + self, model: torch.nn.Module, inputs: List[_MultiInput], **kwargs: Any + ) -> torch.nn.Module: + pass + + +def _check_linear_int4_k(k, groupsize=1, inner_k_tiles=None): + k_divisible_by_groupsize = k % groupsize == 0 + if inner_k_tiles is not None: + k_divisible_by_16_times_inner_k_tiles = k % (inner_k_tiles * 16) == 0 + return k_divisible_by_groupsize and k_divisible_by_16_times_inner_k_tiles + return k_divisible_by_groupsize + + +def linear_forward_int4( + x: torch.Tensor, + weight_int4pack: torch.Tensor, + scales_and_zeros: torch.Tensor, + out_features: int, + groupsize: int, + precision: torch.dtype = torch.bfloat16, + scales_precision: torch.dtype = torch.bfloat16, +): + origin_x_size = x.size() + x = x.reshape(-1, origin_x_size[-1]) + if is_device(x.device.type, "cpu") and TORCH_VERSION_AT_LEAST_2_6: + c = torch.ops.aten._weight_int4pack_mm_for_cpu( + x.to(precision), + weight_int4pack, + groupsize, + scales_and_zeros.to(scales_precision), + ).to(dtype=x.dtype) + else: + c = torch.ops.aten._weight_int4pack_mm( + x.to(precision), + weight_int4pack, + groupsize, + scales_and_zeros.to(scales_precision), + ).to(dtype=x.dtype) + new_shape = origin_x_size[:-1] + (out_features,) + c = c.reshape(new_shape) + return c + + +class WeightOnlyInt4Linear(torch.nn.Module): + __constants__ = ["in_features", "out_features"] + in_features: int + out_features: int + weight: torch.Tensor + + def __init__( + self, + in_features: int, + out_features: int, + # TODO: remove dtype field, not used + bias=False, + device=None, + dtype=None, + groupsize: int = 128, + inner_k_tiles: int = 8, + precision: torch.dtype = torch.bfloat16, + scales_precision: torch.dtype = torch.bfloat16, + ) -> None: + super().__init__() + self.padding = not _check_linear_int4_k(in_features, groupsize, inner_k_tiles) + if self.padding: + self.origin_in_features = in_features + in_features = find_multiple(in_features, 1024) + + self.in_features = in_features + self.out_features = out_features + assert not bias, "require bias=False" + self.device = device + self.groupsize = groupsize + self.inner_k_tiles = inner_k_tiles + self.precision = precision + self.scales_precision = scales_precision + + if dtype is not None: + raise ValueError("Please specify 'precision' instead of 'dtype'") + + assert out_features % 8 == 0, "require out_features % 8 == 0" + assert in_features % (inner_k_tiles * 16) == 0, ( + "require in_features % (innerKTiles * 16) == 0" + ) + if is_device(device.type, "cpu"): + self.register_buffer( + "weight", + torch.zeros( + ( + out_features, + in_features // 2, + ), + dtype=torch.uint8, + device=device, + ), + ) + else: + self.register_buffer( + "weight", + torch.zeros( + ( + out_features // 8, + in_features // (inner_k_tiles * 16), + 32, + inner_k_tiles // 2, + ), + dtype=torch.int32, + device=device, + ), + ) + self.dtype = dtype + self.register_buffer( + "scales_and_zeros", + torch.zeros( + (in_features // groupsize, out_features, 2), + dtype=self.scales_precision, + device=device, + ), + ) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + if self.padding: + input = F.pad(input, pad=(0, self.in_features - self.origin_in_features)) + return linear_forward_int4( + input, + self.weight, + self.scales_and_zeros, + self.out_features, + self.groupsize, + self.precision, + self.scales_precision, + ) + + +def _replace_linear_int4( + module: torch.nn.Module, + groupsize: int, + inner_k_tiles: Optional[int], + padding_allowed: bool, + skip_layer_func: Optional[Callable] = None, + precision: torch.dtype = torch.bfloat16, + scales_precision: torch.dtype = torch.bfloat16, + linear_class: Type[torch.nn.Module] = WeightOnlyInt4Linear, + copy_weights: bool = False, +): + for name, child in module.named_children(): + # TODO: support linear bias + if ( + isinstance(child, nn.Linear) + and child.bias is None + and (skip_layer_func is None or not skip_layer_func(child.weight)) + ): + if ( + _check_linear_int4_k(child.in_features, groupsize, inner_k_tiles) + or padding_allowed + ): + new_linear = linear_class( + child.in_features, + child.out_features, + bias=False, + device=child.weight.device, + groupsize=groupsize, + inner_k_tiles=inner_k_tiles, + precision=precision, + scales_precision=scales_precision, + ) + # TODO: merge with 8da4w? + # In distributed training, the model may be instantiated + # on the meta device, in which case there is no need to + # copy the weights, and doing so will result in an error + if copy_weights and child.weight.device != torch.device("meta"): + new_linear.weight = child.weight + setattr(module, name, new_linear) + else: + _replace_linear_int4( + child, + groupsize, + inner_k_tiles, + padding_allowed, + skip_layer_func, + precision, + scales_precision, + linear_class, + copy_weights, + ) + + +def replace_linear_int4( + module, groupsize, inner_k_tiles, padding_allowed, skip_layer_func=None +): + _replace_linear_int4( + module, + groupsize, + inner_k_tiles, + padding_allowed, + skip_layer_func, + linear_class=WeightOnlyInt4Linear, + ) + + +class Int4WeightOnlyQuantizer(Quantizer): + def __init__( + self, + groupsize: int = 256, + padding_allowed: bool = True, + inner_k_tiles: Optional[int] = 8, + device: torch.device = torch.device("cuda"), + precision: torch.dtype = torch.bfloat16, + ) -> None: + super().__init__() + assert inner_k_tiles in [2, 4, 8] + assert groupsize in [32, 64, 128, 256] + + self.inner_k_tiles = inner_k_tiles + self.groupsize: int = groupsize + self.padding_allowed: bool = padding_allowed + self.device: torch.device = device + # precision and dtype are being used interchangeably here + self.precision: torch.dtype = precision + + @torch.no_grad() + def _create_quantized_state_dict( + self, model: torch.nn.Module + ) -> Dict[str, torch.Tensor]: + cur_state_dict = model.state_dict() + for fqn, mod in model.named_modules(): + if isinstance(mod, torch.nn.Linear) and mod.bias is None: + out_features = mod.out_features + in_features = mod.in_features + # assert out_features % 8 == 0, "require out_features % 8 == 0" + logging.info(f"linear: {fqn}, in={in_features}, out={out_features}") + + assert in_features % self.groupsize == 0, ( + f"require in_features:{in_features} % self.groupsize:{self.groupsize} == 0" + ) + + weight = mod.weight.data + if not _check_linear_int4_k( + in_features, self.groupsize, self.inner_k_tiles + ): + if self.padding_allowed: + import torch.nn.functional as F + + logging.warning( + f"warning: {fqn} is padded to satisfy in_features % 1024 == 0" + ) + padded_in_features = find_multiple(in_features, 1024) + weight = F.pad( + weight, pad=(0, padded_in_features - in_features) + ) + else: + logging.warning( + f"warning: {fqn} is skipped, int4 requires that in_features is 32, 64, or is divisible by 1024, " + + "and that groupsize and inner_k_tiles*16 evenly divide into it" + ) + continue + (w_int4x8, scales_and_zeros) = groupwise_affine_quantize_tensor( + weight, + 4, # n_bit + self.groupsize, + self.precision, # dtype for scales_and_zeros + ) + # TODO: just get the device from mod.weight.device? + if ( + is_device(w_int4x8.device.type, "cpu") + and TORCH_VERSION_AT_LEAST_2_6 + ): + weight_int4pack = ( + torch.ops.aten._convert_weight_to_int4pack_for_cpu( + w_int4x8.to(self.device), self.inner_k_tiles + ) + ) + else: + weight_int4pack = torch.ops.aten._convert_weight_to_int4pack( + w_int4x8.to(self.device), self.inner_k_tiles + ) + cur_state_dict[f"{fqn}.weight"] = weight_int4pack.to(self.device) + cur_state_dict[f"{fqn}.scales_and_zeros"] = scales_and_zeros.to( + self.device + ) + return cur_state_dict + + def _convert_for_runtime(self, model: torch.nn.Module) -> torch.nn.Module: + _replace_linear_int4( + model, + self.groupsize, + self.inner_k_tiles, + self.padding_allowed, + skip_layer_func=None, + precision=self.precision, + scales_precision=self.precision, + ) + return model + + def quantize( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + state_dict = self._create_quantized_state_dict(model) + model = self._convert_for_runtime(model) + # TODO: make it strict + model.load_state_dict(state_dict, strict=False) + return model + + +class Int4WeightOnlyGPTQQuantizer(GPTQQuantizer): + def __init__( + self, + blocksize=128, + percdamp=0.01, + groupsize=64, + inner_k_tiles=8, + padding_allowed=True, + device: torch.device = torch.device("cuda"), + ): + self.blocksize = blocksize + self.percdamp = percdamp + self.groupsize = groupsize + self.inner_k_tiles = inner_k_tiles + self.padding_allowed = padding_allowed + self.device = device + self.act_fake_quant_func = None + n_bit = 4 + self.get_qparams_func = lambda w: get_groupwise_affine_qparams( + w, n_bit, groupsize + ) + self.quantize_func = ( + lambda w, qparams: groupwise_affine_quantize_tensor_from_qparams( + w, qparams[0], qparams[1], n_bit, groupsize + ) + ) + self.dequantize_func = ( + lambda q, qparams: groupwise_affine_dequantize_tensor_from_qparams( + q, + qparams[0], + qparams[1], + n_bit, + groupsize, + ) + ) + self.combine_qparams_list_func = lambda qparams_list: [ + torch.cat(x, dim=1) for x in zip(*qparams_list) + ] + # skip unless padding_allowed=True or its correctly sized + self.skip_layer_func = lambda linear_weight: not ( + _check_linear_int4_k(linear_weight.shape[-1], groupsize) or padding_allowed + ) + + # we need to do the padding here, both for q and the qparams if necessary + + # TODO: this is the gpt-fast version, merge with the main version later + def make_names_and_values_dict_func(q, qparams): + k = q.shape[1] * 2 + if not _check_linear_int4_k(k, groupsize): + new_k = find_multiple(k, 1024) + else: + new_k = k + # how much we need to pad the weight + delta_k = int((new_k - k) / 2) + q = q.to(self.device) + if is_device(self.device.type, "cpu") and TORCH_VERSION_AT_LEAST_2_6: + final_q = torch.ops.aten._convert_weight_to_int4pack_for_cpu( + F.pad(q, pad=(0, delta_k)), inner_k_tiles + ) + else: + final_q = torch.ops.aten._convert_weight_to_int4pack( + F.pad(q, pad=(0, delta_k)), inner_k_tiles + ) + scales = qparams[0].to(torch.bfloat16).to(self.device) + zeros = qparams[1].to(torch.bfloat16).to(self.device) + scales_and_zeros = pack_tinygemm_scales_and_zeros(scales, zeros) + # how many new groups we need for padded weight + delta_groups = new_k // groupsize - scales_and_zeros.shape[0] + final_s_and_z = F.pad( + scales_and_zeros, pad=(0, 0, 0, 0, 0, delta_groups), value=1 + ) + return {"weight": final_q, "scales_and_zeros": final_s_and_z} + + self.make_names_and_values_dict_func = make_names_and_values_dict_func + super().__init__() + + def _convert_for_runtime(self, model): + replace_linear_int4( + model, + self.groupsize, + self.inner_k_tiles, + self.padding_allowed, + skip_layer_func=self.skip_layer_func, + ) + return model + + def quantize( + self, model: torch.nn.Module, inputs: List[_MultiInput], **kwargs: Any + ) -> torch.nn.Module: + state_dict = self._create_quantized_state_dict( + model, + inputs, + self.blocksize, + self.percdamp, + self.groupsize, + ) + model = self._convert_for_runtime(model) + model.load_state_dict(state_dict, strict=False) + return model + + +def linear_forward_8da4w( + x, + weight_int8, + bias, + scales, + zeros, + out_features, + groupsize, + output_precision, +): + # uses fp32 to match torchao.quantization.quant_api._int8_asymm_per_token_quant + # and activation_scale_dtype in QAT configs + # TODO: in future add ability to specify activation_scale_dtype to PTQ configs + # and enable similar change here + x = per_token_dynamic_quant( + x, scale_dtype=torch.float32, zero_point_dtype=torch.float32 + ) + + # TODO: verify and remove following reshape code + # origin_x_size = x.size() + # x = x.reshape(-1, origin_x_size[-1]) + + # TODO: better API + # weight_int8 = torch.ops.quantized_decomposed.unpack_int4_to_int8(weight_int4packed) + n_bit = 4 + quant_min = -(2 ** (n_bit - 1)) + quant_max = 2 ** (n_bit - 1) - 1 + block_size = (1, groupsize) + + w_dq = dequantize_affine( + weight_int8, + block_size, + scales, + zeros, + torch.int8, + quant_min, + quant_max, + output_dtype=output_precision, + ) + + # x = x.to(torch.float16) + # w_dq = w_dq.to(torch.float16) + c = torch.nn.functional.linear(x, w_dq, bias) + + # new_shape = origin_x_size[:-1] + (out_features,) + # c = c.reshape(new_shape) + + return c + + +class Int8DynActInt4WeightLinear(torch.nn.Module): + __constants__ = ["in_features", "out_features"] + + in_features: int + out_features: int + weight: torch.Tensor + bias: torch.Tensor + + """ + This module implements a dynamic quantized linear layer with int4 weight. + Weights are per channel groupwise quantized. Parameters of importance + groupsize: the number of elements in each quantized group + precision: precision of input and output. e.g. torch.float32 means input + activation is float32 and output is float32. + scales_precision: precision of per group scale. + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias=True, + device=None, + # TODO: remove this field, not used + dtype=None, + groupsize: int = 256, + precision: torch.dtype = torch.float32, + scales_precision: torch.dtype = torch.float32, + ) -> None: + super().__init__() + # always pad if needed since it becomes a noop at runtime if not needed + # self.origin_in_features = in_features + assert in_features % groupsize == 0, ( + f"require in_features:{in_features} % groupsize:{groupsize} == 0" + ) + # in_features = _calc_padded_size_linear_int4( + # in_features, groupsize + # ) + self.in_features = in_features + self.out_features = out_features + # TODO: align groupsize naming + self.groupsize = groupsize + # Precision of the activation which also indicates + # output precision of the dynamically quantized linear layer + # that his module represents. + self.precision = precision + + if dtype is not None: + raise ValueError("Please specify 'precision' instead of 'dtype'") + + # currently storing unpacked int8 weights + self.register_buffer( + "weight", + torch.zeros((out_features, in_features), dtype=torch.int8), + ) + self.register_buffer( + "scales", + torch.zeros( + (out_features, in_features // groupsize), + dtype=scales_precision, + ), + ) + self.register_buffer( + "zeros", + torch.zeros( + (out_features, in_features // groupsize), + dtype=scales_precision, + ), + ) + + if bias: + self.register_buffer("bias", torch.zeros(out_features, dtype=precision)) + else: + self.bias = None + + def forward(self, input: torch.Tensor) -> torch.Tensor: + input = input.to(self.precision) + # padding is removed for perf + # input = F.pad(input, pad=(0, self.in_features - self.origin_in_features)) + return linear_forward_8da4w( + input, + self.weight, + self.bias, + self.scales, + self.zeros, + self.out_features, + self.groupsize, + self.precision, + ) + + +def _replace_linear_8da4w( + module: torch.nn.Module, + groupsize: int, + padding_allowed: bool, + precision: torch.dtype, + scales_precision: torch.dtype, + linear_class: Type[torch.nn.Module], + copy_weights: bool = False, +): + # import the util function here to avoid circular dependency + from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter + + def filter_fn(child: torch.nn.Module, cur_fqn: str) -> bool: + return isinstance(child, nn.Linear) and ( + _check_linear_int4_k(child.in_features, groupsize) or padding_allowed + ) + + def replacement_fn(child: torch.nn.Module) -> torch.nn.Module: + new_linear = linear_class( + child.in_features, + child.out_features, + bias=child.bias is not None, + device=child.weight.device, + groupsize=groupsize, + precision=precision, + scales_precision=scales_precision, + ) + # In distributed training, the model may be instantiated + # on the meta device, in which case there is no need to + # copy the weights, and doing so will result in an error + if copy_weights and child.weight.device != torch.device("meta"): + new_linear.weight = child.weight + new_linear.bias = child.bias + return new_linear + + _replace_with_custom_fn_if_matches_filter(module, replacement_fn, filter_fn) + + +def replace_linear_8da4w( + module: torch.nn.Module, + groupsize: int, + padding_allowed: bool, + precision: torch.dtype, + scales_precision: torch.dtype, +): + _replace_linear_8da4w( + module, + groupsize, + padding_allowed, + precision, + scales_precision, + Int8DynActInt4WeightLinear, + ) + + +class Int8DynActInt4WeightQuantizer(Quantizer): + def __init__( + self, + groupsize: int = 256, + padding_allowed: bool = False, + precision: torch.dtype = torch.float32, + scales_precision: torch.dtype = torch.float32, + device: torch.device = torch.device("cpu"), + mapping_type: MappingType = MappingType.SYMMETRIC, + ) -> None: + super().__init__() + self.groupsize: int = groupsize + self.padding_allowed: bool = padding_allowed + self.precision: torch.dtype = precision + self.scales_precision: torch.dtype = scales_precision + self.device: torch.device = device + self.mapping_type: MappingType = mapping_type + + @torch.no_grad() + def _create_quantized_state_dict( + self, model: torch.nn.Module + ) -> Dict[str, torch.Tensor]: + cur_state_dict = model.state_dict() + for fqn, mod in model.named_modules(): + if isinstance(mod, torch.nn.Linear): + out_features = mod.out_features + in_features = mod.in_features + # assert out_features % 8 == 0, "require out_features % 8 == 0" + logging.info(f"linear: {fqn}, in={in_features}, out={out_features}") + + assert in_features % self.groupsize == 0, ( + f"require in_features:{in_features} % self.groupsize:{self.groupsize} == 0" + ) + + weight = mod.weight.data + if not _check_linear_int4_k(in_features, self.groupsize): + if self.padding_allowed: + import torch.nn.functional as F + + logging.warning( + f"warning: {fqn} is padded to satisfy in_features % 1024 == 0" + ) + padded_in_features = find_multiple(in_features, 1024) + weight = F.pad( + weight, pad=(0, padded_in_features - in_features) + ) + else: + logging.warning( + f"warning: {fqn} is skipped, int4 requires that in_features is 32, 64, or is divisible by 1024, " + + "and that groupsize and inner_k_tiles*16 evenly divide into it" + ) + continue + ( + weight_int8, + scales, + zeros, + ) = group_quantize_tensor_symmetric( + weight.to(self.precision), + 4, # n_bit + self.groupsize, + self.scales_precision, + mapping_type=self.mapping_type, + ) + cur_state_dict[f"{fqn}.weight"] = weight_int8.to(self.device) + cur_state_dict[f"{fqn}.scales"] = scales.to(self.device) + cur_state_dict[f"{fqn}.zeros"] = zeros.to(self.device) + + return cur_state_dict + + def _convert_for_runtime(self, model: torch.nn.Module) -> torch.nn.Module: + replace_linear_8da4w( + model, + self.groupsize, + self.padding_allowed, + self.precision, + # TODO: this should be self.scales_precision? + self.precision, + ) + return model + + def quantize( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + state_dict = self._create_quantized_state_dict(model) + model = self._convert_for_runtime(model) + # TODO: make it strict + model.load_state_dict(state_dict, strict=False) + return model + + +class Int8DynActInt4WeightGPTQQuantizer(GPTQQuantizer): + def __init__( + self, + blocksize=128, + percdamp=0.01, + groupsize=64, + inner_k_tiles=8, + padding_allowed=True, + precision=torch.float32, + ): + self.blocksize = blocksize + self.percdamp = percdamp + self.groupsize = groupsize + self.inner_k_tiles = inner_k_tiles + self.padding_allowed = padding_allowed + self.precision = precision + + self.act_fake_quant_func = per_token_dynamic_quant + n_bit = 4 + self.get_qparams_func = lambda w: get_group_qparams_symmetric( + w, n_bit, groupsize, self.precision + ) + quant_min = -(2 ** (n_bit - 1)) + quant_max = 2 ** (n_bit - 1) - 1 + + from torchao._executorch_ops import ( + _quantized_decomposed_quantize_per_channel_group_wrapper, + ) + + self.quantize_func = ( + lambda w, qparams: _quantized_decomposed_quantize_per_channel_group_wrapper( + w, qparams[0], qparams[1], quant_min, quant_max, torch.int8, groupsize + ) + ) + + from torchao._executorch_ops import ( + _quantized_decomposed_dequantize_per_channel_group_wrapper, + ) + + self.dequantize_func = ( + lambda q, + qparams: _quantized_decomposed_dequantize_per_channel_group_wrapper( + q, + qparams[0], + qparams[1], + quant_min, + quant_max, + torch.int8, + groupsize, + self.precision, + ) + ) + + self.combine_qparams_list_func = lambda qparams_list: [ + torch.cat(x, dim=1) for x in zip(*qparams_list) + ] + # skip unless padding_allowed=True or its correctly sized + + self.skip_layer_func = lambda linear_weight: not ( + _check_linear_int4_k(linear_weight.shape[-1], groupsize) or padding_allowed + ) + + # we need to do the padding here, both for q and the qparams if necessary + def make_names_and_values_dict_func(q, qparams): + k = q.shape[1] + new_k = find_multiple(k, 1 if groupsize is None else groupsize) + # how much we need to pad the weight + delta_k = new_k - q.shape[1] + final_q = F.pad(q, pad=(0, delta_k)) + scales = qparams[0].to(self.precision) + zeros = qparams[1].to(self.precision) + return {"weight": final_q, "scales": scales, "zeros": zeros} + + self.make_names_and_values_dict_func = make_names_and_values_dict_func + super().__init__() + + def _convert_for_runtime(self, model): + replace_linear_8da4w( + model, + self.groupsize, + self.padding_allowed, + self.precision, + # TODO: this should be self.scales_precision? + self.precision, + ) + return model + + def quantize( + self, model: torch.nn.Module, inputs: List[_MultiInput], **kwargs: Any + ) -> torch.nn.Module: + state_dict = self._create_quantized_state_dict( + model, + inputs, + self.blocksize, + self.percdamp, + self.groupsize, + ) + model = self._convert_for_runtime(model) + model.load_state_dict(state_dict, strict=False) + return model diff --git a/lib/python3.12/site-packages/torchao/quantization/GPTQ_MT.py b/lib/python3.12/site-packages/torchao/quantization/GPTQ_MT.py new file mode 100644 index 0000000000000000000000000000000000000000..664ebdfc17f933adeb0aab30c9997956dee121ef --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/GPTQ_MT.py @@ -0,0 +1,788 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union + +import torch +import torch.nn as nn +from torch.utils._pytree import tree_flatten, tree_unflatten + +from torchao.dtypes import TensorCoreTiledLayout, to_affine_quantized_intx_static +from torchao.quantization.quant_primitives import ( + ZeroPointDomain, +) +from torchao.quantization.unified import Quantizer +from torchao.quantization.utils import compute_error as SQNR +from torchao.quantization.utils import ( + get_groupwise_affine_qparams, + groupwise_affine_dequantize_tensor_from_qparams, + groupwise_affine_quantize_tensor_from_qparams, +) + + +def _check_linear_int4_k(k, group_size=1, inner_k_tiles=None): + k_divisible_by_group_size = k % group_size == 0 + if inner_k_tiles is not None: + k_divisible_by_16_times_inner_k_tiles = k % (inner_k_tiles * 16) == 0 + return k_divisible_by_group_size and k_divisible_by_16_times_inner_k_tiles + return k_divisible_by_group_size + + +NON_IN_PLACE_OPS = {} + + +class MultiTensor(torch.Tensor): + get_qparams_func = None + quantize_func = None + dequantize_func = None + combine_qparams_list_func = None + make_qtensor = None + skip_layer_func = None + act_fake_quant_func = None + percdamp = 0.01 + blocksize = 128 + group_size = -1 + in_place_threshold = ( + 5 # Number of times to see a function before assuming it's not in-place + ) + + @staticmethod + def __new__( + cls, input: Union[torch.Tensor, Sequence[torch.Tensor]], **kwargs: Any + ) -> "MultiTensor": + if isinstance(input, (list, tuple)): + input = input[0] + kwargs["dtype"] = kwargs.get("dtype", input.dtype) + shape = kwargs.pop("shape", input.shape) + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) + + def __init__( + self, input: Union[torch.Tensor, Sequence[torch.Tensor]], **kwargs: Any + ) -> None: + self.values: List[torch.Tensor] = [] + self.state_dict_manager = StateDictManager.get_instance() + self.count: int = 0 + self.add_tensors(input) + self.debug: bool = True + self.gptq_done = False + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(data={self.values})" + + def add_tensors( + self, input: Union[torch.Tensor, Sequence[torch.Tensor]] + ) -> "MultiTensor": + if isinstance(input, (tuple, list)): + for inp in input: + self.add_tensors(inp) + else: + assert isinstance(input, torch.Tensor), ( + f"MultiTensor can only use add_tensors for Tensors or lists of tensors but got {type(input)}" + ) + self.count += 1 + self.values.append(input) + return self + + def pad_to_length(self, length, pad_in_place=True): + if self.count < length: + if pad_in_place: + for _ in range(length - self.count): + # we need to handle in place ops where we do want the model's value to stay changed. + # e.g. if someone does z[1,:]=x, if z were a size 1 multiTensor and x were size 3, + # we want z to become a multi tensor of size 3. Thus we pad the MultiTensor to the correct + # size by adding new tensor instances (and not just instances of the pointers to the same original tensor.. + # otherwise changes to one would change all of them) + self.add_tensors(self.values[-1].clone()) + else: + # for non in place ops, no need to bloat memory, can just pad with same tensor instance + return self.__class__(self.values).add_tensors( + [self.values[-1]] * (length - self.count) + ) + return self + + def unpad(self, count=1, force=False): + count = min(count, self.count) + if force or all((self.values[0] == x).all().item() for x in self.values): + self.values = self.values[:count] + self.count = count + else: + return self + + @classmethod + def configure_quantization_mode( + cls, + get_qparams_func, + quantize_func, + dequantize_func, + combine_qparams_list_func, + make_qtensor, + skip_layer_func, + act_fake_quant_func=None, + percdamp=0.01, + blocksize=128, + group_size=-1, + ): + cls.get_qparams_func = get_qparams_func + cls.quantize_func = quantize_func + cls.dequantize_func = dequantize_func + cls.combine_qparams_list_func = combine_qparams_list_func + cls.make_qtensor = make_qtensor + cls.skip_layer_func = skip_layer_func + cls.act_fake_quant_func = ( + act_fake_quant_func if act_fake_quant_func is not None else lambda x: x + ) + cls.percdamp = percdamp + cls.blocksize = blocksize + cls.group_size = group_size + + @classmethod + def __torch_function__( + cls, + func: Callable, + types: Tuple[type, ...], + args: Tuple[Any, ...] = (), + kwargs: Optional[Dict[str, Any]] = None, + skip_gptq: bool = False, + ) -> Any: + def flat_to_grouped( + flat: List[Any], pad_in_place=True + ) -> List[Tuple[Any, ...]]: + # convert [A, MultiTensor(b1,b2,b3), MultiTensor(c1,c2,c3)] => [[A,b1,c1], [A,b2,c2] [A,b3,c3]] + multi_tensor_size = max( + [x.count if isinstance(x, MultiTensor) else 1 for x in flat] + ) + grouped = list( + zip( + *[ + x.pad_to_length( + multi_tensor_size, pad_in_place=pad_in_place + ).values + if isinstance(x, MultiTensor) + else [x] * multi_tensor_size + for x in flat + ] + ) + ) + return grouped + + def grouped_to_flat(grouped: List[Tuple[Any, ...]]) -> Tuple[List[Any], bool]: + # convert [[A,b1,c1], [A,b2,c2] [A,b3,c3]] => [(A,A,A), (b1,b2,b3), (c1,c2,c3)] + flat_tups = list(zip(*grouped)) + # convert [(A,A,A), (b1,b2,b3), (c1,c2,c3)] => [A, MultiTensor(b1,b2,b3), MultiTensor(c1,c2,c3)] + flattened = [ + cls(tup).cpu() if isinstance(tup[0], torch.Tensor) else tup[0] + for tup in flat_tups + ] + non_tensors_equal = all( + all(x == tup[0] for x in tup) + for tup in flat_tups + if not isinstance(tup[0], torch.Tensor) + ) + return flattened, non_tensors_equal + + def tensors_to_cuda(args): + # this is needed because we want to execute the actual ops in cuda so they don't take forever + new_args = [] + for x in args: + if isinstance(x, MultiTensor) and x.count == 1: + new_args.append(x.__class__(x.values[0].cuda())) + else: + new_args.append( + x.cuda() + if isinstance(x, torch.Tensor) + and not isinstance(x, MultiTensor) + else x + ) + return new_args + + def maybe_copy_new_values(orig_inp, new_inp): + detected_difference = False + for x, new_x in zip(orig_inp, new_inp): + if isinstance(x, torch.Tensor): + new_x = new_x.to(x.device) + if (x != new_x).any(): + x.copy_(new_x) + detected_difference = True + return detected_difference + + def unpad(args, orig_counts, force=False): + for arg, count in zip(args, orig_counts): + if isinstance(arg, MultiTensor) and arg.count > count: + arg.unpad(count, force) + + # The way MultiTensor handles various functions is as follows. Normally when you apply a function on a MultiTensor that has n Tensors inside, we want + # the function handling here to run that function once for each of the MultiTensor inputs. We also want it to happen in the same way as if you ran the function + # first input and the second independently, i.e. if you ran model(input1)=out1 and model(input2), vs model(MultiTensor(input1, input2)) the output for the + # second case should be MultiTensor(out1, out2) and all the activations along the way should be the same. Normally it is easy enough to handle MultiTensors + # running the func for each set of inputs and then combine the outputs back into MultiTensors at the end if applicable. Since we can have a lot of tensors in a MultiTensor but we + # normally want to execute the function on cuda, we can move them over to cuda before we evaluate the function. + # this scheme works pretty well but has a few issues + # 1) We end up moving the same tensors to cuda over and over again which can make things slow. if we have a function input like (MultiTensor(a), MultiTensor(b1, b2, ...b_n) + # when we pad and group MultiTensors into Tensor-only function inputs i.e. (a, b1), (a, b2), ..., (a, b_n), since each Tensor-only input uses a common tensor a but a unique tensor b, + # we would be moving the same tensor a to cuda n times for no reason. This is easy to fix normally, just move all singular MultiTensors to cuda before padding/grouping but... + # 2) In place ops are tricky to handle (funcs that modify the inputs). We want in_place operations like k_cache[:, indices] = k_val (where k_cache, k_val are MultiTensor inputs) + # to be supported but if we move tensors to cuda then any modifications to the inputs will be applied to the cuda tensors rather than the originals. Additionally the originals are often + # singular MultiTensors i.e. MultiTensor(a) that only has a single value, we don't want each in place op to overwrite the value of a, we need each change to the inputs to be recorded. So + # we can modify MultiTensor(a) -> MultiTensor(a1, a2, ...) at the start so that each change of a can be recorded and won't overwrite the value of a for other ops, but then problem 1 comes back + # since we have to move a1, a2, a3...etc to cuda over and over again. despite them all being the same value initially. And if we move them to cuda then the in place value changes will + # be applied to a1_cuda not a1 and a1_cuda isn't in the MultiTensor. So we have to manually copy those values back a1.copy_(a1_cuda). + # There's not really a great way to resolve the 2 issues, when there's an in place op you have to do the slow thing with cuda and checking for modified values....etc, when + # there's not an in place op you can throw all singular tensors onto cuda at the start and go much faster. + # This brings up the final issue, how do we know if we have an in place op? In general we don't so I added handling to MultiTensor to resolve that as well as can be hoped. + # we have a dict that contains all the funcs we see NON_IN_PLACE_OPS, we initially treat ops as in place and see if any of the inputs got modified if they do then it gets + # set to always be handled as an in place op. If nothing changes then once we've seen the op enough times that we're confident its not an in place op (cls.in_place_threshold) + # then we can do the fast thing. + + quantize_linear = not skip_gptq and cls.is_linear_layer(func) + # Determine if function is in-place + + # initialize function tracking + if func not in NON_IN_PLACE_OPS: + NON_IN_PLACE_OPS[func] = {"count": 0, "is_in_place": None} + NON_IN_PLACE_OPS[func]["count"] += 1 + + if NON_IN_PLACE_OPS[func]["is_in_place"] is not None: + is_in_place = NON_IN_PLACE_OPS[func]["is_in_place"] + elif ( + NON_IN_PLACE_OPS[func]["count"] >= cls.in_place_threshold or quantize_linear + ): + is_in_place = False # Assume not in-place after threshold + else: + is_in_place = True + + kwargs = {} if kwargs is None else kwargs + # combine args and kwargs into a single tuple + # flat_args holds all the actual inputs, spec stores the original structure + flat_args, spec = tree_flatten((args, kwargs)) + + orig_counts = [x.count if isinstance(x, MultiTensor) else 1 for x in flat_args] + + # if we're not doing an in place op, move singular tensors to cuda now + if not is_in_place: + flat_args = tensors_to_cuda(flat_args) + + # convert [A, MultiTensor(b), MultiTensor(c1,c2,c3)] => [[A,b,c1], [A,b,c2] [A,b,c3]] + # if its in place then instead we first convert MultiTensor(b) => MultiTensor(b1, b2, b3) + # then proceed as normal. + grouped_args = flat_to_grouped(flat_args, is_in_place) + + with torch._C.DisableTorchFunctionSubclass(): + if quantize_linear: + H = 0 + total_batches = 0 + + outputs = [] + for inp in grouped_args: + # we move all remaining cpu tensors to cuda + cuda_inp = tensors_to_cuda(inp) + + # return input to original structure + cur_args, cur_kwargs = tree_unflatten(cuda_inp, spec) + + if quantize_linear: + # Construct Hessian matrix for quantization + x = cur_args[0].float() + # x = self.act_fake_quant_func(x) + shape = x.shape + n = 1 if len(shape) == 2 else shape[0] + H *= total_batches / (total_batches + n) + total_batches += n + + x = ((2 / total_batches) ** (1 / 2)) * x.reshape( + -1, shape[-1] + ).t().float() + + H += x.matmul(x.t()) + else: + try: + out = func(*cur_args, **cur_kwargs) + except Exception: + breakpoint() + + outputs.append(out.cpu() if isinstance(out, torch.Tensor) else out) + + # if we're doing an in place op, here is where we copy modifications + # back to the original tensors, if we saw any differences, immediately + # categortize func as in place, otherwise we can treat as not in + # place (especially for the upcoming unpad step) + if is_in_place: + detected_difference = maybe_copy_new_values(inp, cuda_inp) + # if detected_difference and not isinstance(NON_IN_PLACE_OPS[func], bool): + # print("THIS OP IS IN PLACE", func) + # NON_IN_PLACE_OPS[func] = False + if detected_difference and not NON_IN_PLACE_OPS[func]: + NON_IN_PLACE_OPS[func]["is_in_place"] = True + print(f"Function {func} is in-place") + + elif NON_IN_PLACE_OPS[func]["count"] >= cls.in_place_threshold: + NON_IN_PLACE_OPS[func]["is_in_place"] = False + + if quantize_linear: + # turn weight MultiTensor into single cuda tensor + W = args[1] + if isinstance(W, MultiTensor): + W = W.values[0] + W = W.to(H.device) + + Q, DQ, all_qparams = cls.faster_quant(H, W.detach()) + + # make quantized tensor subclass + qtensor = cls.make_qtensor(Q, all_qparams) + + # Get the original parameter name + state_dict_manager = StateDictManager.get_instance() + original_param_name = state_dict_manager.get_name_for_param(args[1]) + state_dict_manager.update_param(original_param_name, qtensor) + print(original_param_name) + + # Run the function again with updated weights and skip_gptq=True + out = cls.__torch_function__( + func, types, (args[0], DQ.cpu(), *args[2:]), kwargs, skip_gptq=True + ) + if args[0].debug: + act = args[0].values[0].to("cuda") + bias = ( + args[2].values[0].to("cuda") if args[2] is not None else args[2] + ) + + new_out = out.values[0].cpu() + old_out = ( + cls.__torch_function__( + func, + types, + (act, args[1].values[0], bias), + kwargs, + skip_gptq=True, + ) + .values[0] + .cpu() + ) + + DQ_after = cls.dequantize_func(Q, all_qparams).to(W.dtype) + print( + "SQNR for QDQ (this should be inf)", SQNR(DQ, DQ_after) + ) # matches + print( + "SQNR for weight (can be low)", SQNR(W, DQ.cuda()) + ) # fine to not match + print( + "SQNR for output with GPTQ (hopefully 35+)", + SQNR(old_out, new_out), + ) + + DQ_from_qtensor = qtensor.dequantize() + qtensor_out = torch.nn.functional.linear(act, qtensor, bias).cpu() + print( + "SQNR for output from qtensor vs output from DQ (should be high)", + SQNR(qtensor_out, new_out), + ) + print( + "SQNR for DQ vs DQ from qtensor (should be inf)", + SQNR(DQ, DQ_from_qtensor), + ) + + qparams2 = cls.get_qparams_func(W) + Q2 = cls.quantize_func(W, qparams2) + DQ2 = cls.dequantize_func(Q2, qparams2).to(W.dtype) + old_q_out = ( + cls.__torch_function__( + func, types, (act, DQ2, bias), kwargs, skip_gptq=True + ) + .values[0] + .cpu() + ) + + print( + "SQNR for output without GPTQ (should be less than above)", + SQNR(old_out, old_q_out), + ) + unpad(flat_args, orig_counts=orig_counts, force=True) + return out + else: + # we padded each of the MultiTensors to match the largest multitensor so that if we had in place ops, we would be able + # to store the many changed value and have those updates be reflected in the model. However if there are no in place ops, then + # we just increased the size of all parameters/buffers by n times for no reason. To avoid issues, go back and unpad + # everything where possible. i.e. all the multi tensor values are the same. We already checked for mutations and + # if we detected them, we updated NON_IN_PLACE_OPS to be False, so we can just check that see if we need + # to be careful during unpadding. + unpad( + flat_args, + orig_counts=orig_counts, + force=(not isinstance(NON_IN_PLACE_OPS[func], bool)), + ) + + grouped_outputs = [tree_flatten(x)[0] for x in outputs] + out_spec = tree_flatten(outputs[0])[1] + # conslidate out into MultiTensors [[A,b1,c1], [A,b2,c2] [A,b3,c3]] => [A, MultiTensor(b1,b2,b3), MultiTensor(c1,c2,c3)] + flat_outputs, non_tensors_equal = grouped_to_flat(grouped_outputs) + assert non_tensors_equal, ( + f"ERR: found a function in model: {func} which " + + "caused an error in GPTQ MultiTensor, the function dispatch only works for functions" + + " with Tensor outputs or that have the same non-Tensor output value for all across all inputs" + ) + final_out = tree_unflatten(flat_outputs, out_spec) + return final_out + + @classmethod + def faster_quant(cls, H, W): + percdamp = cls.percdamp + blocksize = cls.blocksize + group_size = cls.group_size + orig_dtype = W.dtype + W = W.detach().float() + _, columns = W.shape[0], W.shape[1] + device = W.device + + if group_size == -1: + cur_qparams = cls.get_qparams_func(W) + + dead = torch.diag(H) == 0 + H[dead, dead] = 1 + W[:, dead] = 0 + + Losses = torch.zeros_like(W) + DQ = torch.zeros_like(W) + + damp = percdamp * torch.mean(torch.diag(H)) + diag = torch.arange(columns, device=device) + H[diag, diag] += damp + H = torch.linalg.cholesky(H) + H = torch.cholesky_inverse(H) + H = torch.linalg.cholesky(H, upper=True) + Hinv = H + + all_qparams = [] + for i1 in range(0, columns, blocksize): + i2 = min(i1 + blocksize, columns) + count = i2 - i1 + W1 = W[:, i1:i2].clone() + DQ1 = torch.zeros_like(W1) + Err1 = torch.zeros_like(W1) + Losses1 = torch.zeros_like(W1) + Hinv1 = Hinv[i1:i2, i1:i2] + for i in range(count): + w = W1[:, i] + d = Hinv1[i, i] + + if ( + group_size != -1 and (i1 + i) % group_size == 0 + ): # start of new group + cur_qparams = cls.get_qparams_func( + W[:, (i1 + i) : (i1 + i + group_size)] + ) + all_qparams.append(cur_qparams) + + q = cls.quantize_func(w.unsqueeze(1), cur_qparams).flatten() + dq = cls.dequantize_func(q.unsqueeze(1), cur_qparams).flatten() + + DQ1[:, i] = dq + Losses1[:, i] = (w - dq) ** 2 / d**2 + + err1 = (w - dq) / d + W1[:, i:] -= ( + err1.to(Hinv1.dtype).unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0)) + ) + Err1[:, i] = err1 + + DQ[:, i1:i2] = DQ1 + Losses[:, i1:i2] = Losses1 / 2 + + W[:, i2:] -= Err1.to(Hinv.dtype).matmul(Hinv[i1:i2, i2:]) + + torch.cuda.synchronize() + + if all_qparams == []: + all_qparams.append(cur_qparams) + + all_qparams = cls.combine_qparams_list_func(all_qparams) + Q = cls.quantize_func(DQ, all_qparams) + return Q, DQ.to(orig_dtype), all_qparams + + @classmethod + def __torch_dispatch__( + cls, + func: Callable, + types: Tuple[type, ...], + args: Tuple[Any, ...] = (), + kwargs: Dict[str, Any] = {}, + skip_gptq: bool = False, + ) -> Any: + pass + + def __tensor_flatten__(self) -> Tuple[List[str], Optional[Any]]: + return ["values"], None + + @classmethod + def __tensor_unflatten__( + cls, + tensor_data_dict: Dict[str, Any], + tensor_attributes: Optional[Any], + outer_size: torch.Size, + outer_stride: Tuple[int, ...], + ) -> "MultiTensor": + return cls(tensor_data_dict["values"]) + + @classmethod + def is_linear_layer(cls, func: Callable) -> bool: + return func == torch.nn.functional.linear + + +class StateDictManager: + _instance = None + + @staticmethod + def get_instance(): + if StateDictManager._instance is None: + StateDictManager._instance = StateDictManager() + return StateDictManager._instance + + def __init__(self): + self.state_dict = {} + self.id_to_name = {} + + def set_state_dict(self, model): + self.state_dict = model.state_dict() + self.id_to_name = {id(v): k for k, v in model.named_parameters()} + + def update_id_to_name(self, model): + self.id_to_name = {id(v): k for k, v in model.named_parameters()} + + def get_name_for_param(self, param): + return self.id_to_name.get(id(param), None) + + def update_param(self, name, new_value): + if name in self.state_dict: + if isinstance(new_value, MultiTensor): + self.state_dict[name] = new_value.values[ + 0 + ] # Convert MultiTensor to regular tensor + else: + self.state_dict[name] = new_value + else: + raise KeyError(f"Parameter {name} not found in state_dict") + + def get_state_dict(self): + return self.state_dict + + +class GPTQQuantizer(Quantizer): + def __init__(self): + super().__init__() + self.state_dict_manager = StateDictManager.get_instance() + self.get_qparams_func = None + self.quantize_func = None + self.dequantize_func = None + self.combine_qparams_list_func = None + self.make_qtensor = None + self.skip_layer_func = None + self.act_fake_quant_func = None + + def _check_functions(self): + assert self.get_qparams_func is not None, "get_qparams_func must be set" + assert self.quantize_func is not None, "quantize_func must be set" + assert self.dequantize_func is not None, "dequantize_func must be set" + assert self.combine_qparams_list_func is not None, ( + "combine_qparams_list_func must be set" + ) + assert self.make_qtensor is not None, "make_qtensor must be set" + assert self.skip_layer_func is not None, "skip_layer_func must be set" + + # this doesn't work + def _replace_parameters_with_multitensor(self, model): + for name, param in model.named_parameters(): + setattr(model, name.split(".")[-1], MultiTensor(param)) + + def covert_multi_tensors_to_tensors(self, state_dict): + for key, value in state_dict.items(): + if isinstance(value, MultiTensor): + state_dict[key] = value.values[0] + return state_dict + + @torch.no_grad() + def _create_quantized_state_dict( + self, + model, + inputs, + blocksize=128, + percdamp=0.01, + group_size=64, + # `typing.Dict[, ]` to avoid runtime subscripting errors. + ) -> Dict: + MultiTensor.configure_quantization_mode( + get_qparams_func=self.get_qparams_func, + quantize_func=self.quantize_func, + dequantize_func=self.dequantize_func, + combine_qparams_list_func=self.combine_qparams_list_func, + make_qtensor=self.make_qtensor, + skip_layer_func=self.skip_layer_func, + percdamp=percdamp, + blocksize=blocksize, + group_size=group_size, + ) + # Set the state dict for the original model + self.state_dict_manager.set_state_dict(model) + # Replace parameters with MultiTensor + # self._replace_parameters_with_multitensor(model) + # Replace buffers and parameters with MultiTensor + with torch.no_grad(): + _replace_with_custom_fn_if_matches_filter( + model=model, + replacement_fn=replace_buffers_and_params_with_multitensors, + filter_fn=lambda x, y: True, + ) + self.state_dict_manager.update_id_to_name(model) + # Run the model + with torch.no_grad(): + model(*inputs) + state_dict = self.state_dict_manager.get_state_dict() + return state_dict + + +class Int4WeightOnlyGPTQQuantizer(GPTQQuantizer): + def __init__( + self, + blocksize=128, + percdamp=0.01, + group_size=64, + inner_k_tiles=8, + padding_allowed=True, + device: torch.device = torch.device("cuda"), + ): + super().__init__() + self.blocksize = blocksize + self.percdamp = percdamp + self.group_size = group_size + self.inner_k_tiles = inner_k_tiles + self.padding_allowed = padding_allowed + self.device = device + self.act_fake_quant_func = None + n_bit = 4 + self.get_qparams_func = lambda w: get_groupwise_affine_qparams( + w, n_bit, group_size + ) + self.quantize_func = ( + lambda w, qparams: groupwise_affine_quantize_tensor_from_qparams( + w, qparams[0], qparams[1], n_bit, group_size + ) + ) + self.dequantize_func = ( + lambda q, qparams: groupwise_affine_dequantize_tensor_from_qparams( + q, + qparams[0], + qparams[1], + n_bit, + group_size, + ) + ) + self.combine_qparams_list_func = lambda qparams_list: [ + torch.cat(x, dim=1) for x in zip(*qparams_list) + ] + # skip unless padding_allowed=True or its correctly sized + self.skip_layer_func = lambda linear_weight: not ( + _check_linear_int4_k(linear_weight.shape[-1], group_size) or padding_allowed + ) + + def make_qtensor(q, qparams): + # this should be setup to just use the quantized tensor and qparams directly to make + # the aqt int4 tensor but i don't think we have that functionality atm so just dequant + # then requant + weight = self.dequantize_func(q, qparams) + scale = qparams[0] + zero_point = qparams[1] + + # copied from quant_api apply_int4_weight_only_quant (this should probably be made into a utility fn at some point) + # mapping_type = MappingType.ASYMMETRIC + block_size = (1, group_size) + target_dtype = torch.int32 + quant_min = 0 + quant_max = 15 + zero_point_domain = ZeroPointDomain.FLOAT + _layout = TensorCoreTiledLayout(inner_k_tiles=8) + # at least the big up to here should be a util + + quantized_tensor = to_affine_quantized_intx_static( + weight, + scale=scale, + zero_point=zero_point, + block_size=block_size, + target_dtype=target_dtype, + quant_min=quant_min, + quant_max=quant_max, + zero_point_domain=zero_point_domain, + _layout=_layout, + ) + return quantized_tensor + + self.make_qtensor = make_qtensor + + self._check_functions() + + def quantize( + self, model: torch.nn.Module, inputs: List[MultiTensor], **kwargs: Any + ) -> torch.nn.Module: + state_dict = self._create_quantized_state_dict( + model, + inputs, + self.blocksize, + self.percdamp, + self.group_size, + ) + + # this is hacky and potentially wrong, better to just make the flow return a state dict and let user + # do with it what they will + + model = _replace_with_custom_fn_if_matches_filter( + model=model, + replacement_fn=remove_multitensors_from_buffers_and_params, + filter_fn=lambda x, y: True, + ) + remove = [k for k in state_dict if "kv_cache" in k] + for k in remove: + del state_dict[k] + + model.load_state_dict(state_dict, assign=True, strict=False) + + return model + + +# this should probably be a multitensor method that can be applied and we just traverse +# and look for multitensors and unpack them +def remove_multitensors_from_buffers_and_params(model: nn.Module) -> nn.Module: + for name, buf in model.named_buffers(recurse=False): + if isinstance(buf, MultiTensor): + setattr(model, name, buf.values[0]) + for name, param in model.named_parameters(recurse=False): + if isinstance(param, MultiTensor): + setattr( + model, + name, + nn.Parameter(param.values[0], param.values[0].requires_grad), + ) + return model + + +def replace_buffers_and_params_with_multitensors(model: nn.Module) -> nn.Module: + for name, buf in model.named_buffers(recurse=False): + setattr(model, name, MultiTensor([buf])) + for name, param in model.named_parameters(recurse=False): + setattr(model, name, nn.Parameter(MultiTensor([param]), param.requires_grad)) + return model + + +def _replace_with_custom_fn_if_matches_filter( + model: nn.Module, + replacement_fn: Callable[[nn.Module], nn.Module], + filter_fn: Callable[[nn.Module, str], bool], + cur_fqn: str = "", +) -> None: + if filter_fn(model, cur_fqn[:-1]): + model = replacement_fn(model) + for name, child in model.named_children(): + new_child = _replace_with_custom_fn_if_matches_filter( + child, replacement_fn, filter_fn, f"{cur_fqn}{name}." + ) + if new_child is not child: + setattr(model, name, new_child) + return model diff --git a/lib/python3.12/site-packages/torchao/quantization/__init__.py b/lib/python3.12/site-packages/torchao/quantization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fa156691cadd8b1b61fe7f78a60b8079fe7f1adf --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/__init__.py @@ -0,0 +1,199 @@ +from torchao.kernel import ( + int_scaled_matmul, + safe_int_mm, +) + +from .autoquant import ( + ALL_AUTOQUANT_CLASS_LIST, + DEFAULT_AUTOQUANT_CLASS_LIST, + DEFAULT_FLOAT_AUTOQUANT_CLASS_LIST, + DEFAULT_INT4_AUTOQUANT_CLASS_LIST, + DEFAULT_SPARSE_AUTOQUANT_CLASS_LIST, + GEMLITE_INT4_AUTOQUANT_CLASS_LIST, + OTHER_AUTOQUANT_CLASS_LIST, + autoquant, +) +from .GPTQ import ( + Int4WeightOnlyGPTQQuantizer, + Int4WeightOnlyQuantizer, + Int8DynActInt4WeightGPTQQuantizer, + Int8DynActInt4WeightLinear, + Int8DynActInt4WeightQuantizer, +) +from .granularity import ( + PerAxis, + PerGroup, + PerRow, + PerTensor, + PerToken, +) +from .linear_activation_quantized_tensor import ( + LinearActivationQuantizedTensor, + to_linear_activation_quantized, +) +from .linear_activation_scale import ( + to_weight_tensor_with_linear_activation_scale_metadata, +) +from .observer import ( + AffineQuantizedMinMaxObserver, + AffineQuantizedObserverBase, +) +from .quant_api import ( + AOPerModuleConfig, + CutlassInt4PackedLayout, + Float8DynamicActivationFloat8SemiSparseWeightConfig, + Float8DynamicActivationFloat8WeightConfig, + Float8MMConfig, + Float8StaticActivationFloat8WeightConfig, + Float8WeightOnlyConfig, + FPXWeightOnlyConfig, + GemliteUIntXWeightOnlyConfig, + Int4DynamicActivationInt4WeightConfig, + Int4WeightOnlyConfig, + Int8DynamicActivationInt4WeightConfig, + Int8DynamicActivationInt8WeightConfig, + Int8WeightOnlyConfig, + IntxWeightOnlyConfig, + PlainLayout, + TensorCoreTiledLayout, + UIntXWeightOnlyConfig, + float8_dynamic_activation_float8_weight, + float8_static_activation_float8_weight, + float8_weight_only, + fpx_weight_only, + gemlite_uintx_weight_only, + int4_dynamic_activation_int4_weight, + int4_weight_only, + int8_dynamic_activation_int4_weight, + int8_dynamic_activation_int8_semi_sparse_weight, + int8_dynamic_activation_int8_weight, + int8_weight_only, + intx_quantization_aware_training, + quantize_, + swap_conv2d_1x1_to_linear, + uintx_weight_only, +) +from .quant_primitives import ( + MappingType, + TorchAODType, + ZeroPointDomain, + choose_qparams_affine, + choose_qparams_affine_floatx, + choose_qparams_affine_with_min_max, + choose_qparams_and_quantize_affine_hqq, + dequantize_affine, + dequantize_affine_floatx, + fake_quantize_affine, + fake_quantize_affine_cachemask, + quantize_affine, + quantize_affine_floatx, +) +from .smoothquant import ( + SmoothFakeDynamicallyQuantizedLinear, + SmoothFakeDynQuantMixin, + get_scale, + set_smooth_fq_attribute, + smooth_fq_linear_to_inference, + swap_linear_with_smooth_fq_linear, +) +from .subclass import * # noqa: F403 +from .transform_module import register_quantize_module_handler +from .unified import Quantizer, TwoStepQuantizer +from .utils import ( + compute_error, +) +from .weight_only import WeightOnlyInt8QuantLinear + +__all__ = [ + # top level API - auto + "autoquant", + "DEFAULT_AUTOQUANT_CLASS_LIST", + "DEFAULT_INT4_AUTOQUANT_CLASS_LIST", + "GEMLITE_INT4_AUTOQUANT_CLASS_LIST", + "DEFAULT_FLOAT_AUTOQUANT_CLASS_LIST", + "DEFAULT_SPARSE_AUTOQUANT_CLASS_LIST", + "OTHER_AUTOQUANT_CLASS_LIST", + "ALL_AUTOQUANT_CLASS_LIST", + # top level API - manual + "quantize_", + "int4_dynamic_activation_int4_weight", + "int8_dynamic_activation_int4_weight", + "int8_dynamic_activation_int8_weight", + "int8_dynamic_activation_int8_semi_sparse_weight", + "int4_weight_only", + "int8_weight_only", + "intx_quantization_aware_training", + "float8_weight_only", + "float8_dynamic_activation_float8_weight", + "float8_static_activation_float8_weight", + "uintx_weight_only", + "fpx_weight_only", + "gemlite_uintx_weight_only", + "swap_conv2d_1x1_to_linear", + "Int4DynamicActivationInt4WeightConfig", + "Int8DynamicActivationInt4WeightConfig", + "Int8DynamicActivationInt8WeightConfig", + "Int4WeightOnlyConfig", + "Int8WeightOnlyConfig", + "Float8WeightOnlyConfig", + "Float8DynamicActivationFloat8WeightConfig", + "Float8StaticActivationFloat8WeightConfig", + "Float8DynamicActivationFloat8SemiSparseWeightConfig", + "UIntXWeightOnlyConfig", + "IntxWeightOnlyConfig", + "FPXWeightOnlyConfig", + "GemliteUIntXWeightOnlyConfig", + "AOPerModuleConfig", + # smooth quant - subject to change + "get_scale", + "SmoothFakeDynQuantMixin", + "SmoothFakeDynamicallyQuantizedLinear", + "swap_linear_with_smooth_fq_linear", + "smooth_fq_linear_to_inference", + "set_smooth_fq_attribute", + "compute_error", + # building blocks + "to_linear_activation_quantized", + "to_weight_tensor_with_linear_activation_scale_metadata", + "AffineQuantizedMinMaxObserver", + "AffineQuantizedObserverBase", + # quant primitive ops + "choose_qparams_affine", + "choose_qparams_affine_with_min_max", + "choose_qparams_affine_floatx", + "quantize_affine", + "quantize_affine_floatx", + "dequantize_affine", + "dequantize_affine_floatx", + "choose_qparams_and_quantize_affine_hqq", + "fake_quantize_affine", + "fake_quantize_affine_cachemask", + # operators/kernels + "safe_int_mm", + "int_scaled_matmul", + # registration of module transforms for quantize_ + "register_quantize_module_handler", + # dataclasses and types + "MappingType", + "ZeroPointDomain", + "TorchAODType", + "PerTensor", + "PerAxis", + "PerGroup", + "PerRow", + "PerToken", + "LinearActivationQuantizedTensor", + "Int4WeightOnlyGPTQQuantizer", + "Int4WeightOnlyQuantizer", + "Int8DynActInt4WeightGPTQQuantizer", + "Int8DynActInt4WeightQuantizer", + "Int8DynActInt4WeightLinear", + "WeightOnlyInt8QuantLinear", + "TwoStepQuantizer", + "Quantizer", + # Layouts for quant_api + "PlainLayout", + "TensorCoreTiledLayout", + "CutlassInt4PackedLayout", + "Float8MMConfig", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/GPTQ.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/GPTQ.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0fe460df33c86bc721bc52214d6db156c79e505 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/GPTQ.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/GPTQ_MT.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/GPTQ_MT.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b33e5e033f0feda067e62b52617d3fcc8e8e0821 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/GPTQ_MT.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14cdc795876d319b872d9459ec03fa2d1676f592 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/autoquant.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/autoquant.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e0a9dffb65964b3f2341774f43ceaa3a703d739 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/autoquant.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/dynamic_quant.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/dynamic_quant.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58ee9a8aa7dfafda655fa1bf8580f5c7a9ea2fd4 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/dynamic_quant.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/granularity.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/granularity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70a320afbd332a04dd59055261365c3b73322c33 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/granularity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_quantized_tensor.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_quantized_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..789f4e4bcddeee3c5b717403713e34e0886391e0 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_quantized_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_scale.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_scale.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..17630fd7c2da991663b49de3303ed2b81ede5e4f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_scale.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_weight_observed_tensor.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_weight_observed_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94a9face09915a1f69d51236aa58372e80ae0f3e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/linear_activation_weight_observed_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/observer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/observer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7663c8eb2669018fbee45702c03748bed918e90e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/observer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/quant_api.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/quant_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..313bb63d2fca75731a7401c6153d870fcc22457f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/quant_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/quant_primitives.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/quant_primitives.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0415ead2d1fb102b1e37b5f78031210229a7018 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/quant_primitives.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/smoothquant.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/smoothquant.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..efc20f2c70e2ef550524c2aa74d3733142cec233 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/smoothquant.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/subclass.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/subclass.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5e2b509e12e1e067957b81c15f5a992a5f26275 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/subclass.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/transform_module.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/transform_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..acf6142790518949e8bde3dd61ab40b9b1319423 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/transform_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/unified.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/unified.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6fc9c772a816e3e98ea675f9a884db4804bdb7cb Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/unified.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14d0547fd5f22da4e004607b875c0e97bde18bbf Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/weight_only.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/weight_only.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23a41d911058b01d3e1d395254d8bf9afcc96080 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/weight_only.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/__pycache__/weight_tensor_linear_activation_quantization.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/__pycache__/weight_tensor_linear_activation_quantization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b1dbd883d1258c764b9fcb104b83ad769e80b07 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/__pycache__/weight_tensor_linear_activation_quantization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/autoquant.py b/lib/python3.12/site-packages/torchao/quantization/autoquant.py new file mode 100644 index 0000000000000000000000000000000000000000..f018d5dffee294a44a427d09ff77121f728e03bc --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/autoquant.py @@ -0,0 +1,1364 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +import torch.nn.functional as F +from torch.utils._python_dispatch import return_and_correct_aliasing + +import torchao +from torchao.dtypes import ( + AffineQuantizedTensor, + Float8Layout, + MarlinSparseLayout, + PlainLayout, + SemiSparseLayout, + TensorCoreTiledLayout, +) +from torchao.dtypes.utils import Layout +from torchao.float8.inference import Float8MMConfig +from torchao.kernel import safe_int_mm +from torchao.quantization.linear_activation_quantized_tensor import ( + LinearActivationQuantizedTensor, + to_linear_activation_quantized, +) +from torchao.quantization.quant_primitives import ( + MappingType, + ZeroPointDomain, +) +from torchao.quantization.utils import ( + compute_error, + quantize_activation_per_token_absmax, +) +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_3, + TORCH_VERSION_AT_LEAST_2_5, + TorchAOBaseTensor, + is_sm_at_least_89, + is_sm_at_least_90, +) + +from .granularity import ( + PerRow, + PerTensor, +) +from .subclass import ( # noqa + Int8DynamicallyQuantizedLinearWeight, + Int8WeightOnlyQuantizedLinearWeight, + QuantizedLinearWeightBase, +) + +__all__ = [ + "AutoQuantizableLinearWeight", + "autoquant", + "DEFAULT_AUTOQUANT_CLASS_LIST", + "DEFAULT_INT4_AUTOQUANT_CLASS_LIST", + "GEMLITE_INT4_AUTOQUANT_CLASS_LIST", + "DEFAULT_FLOAT_AUTOQUANT_CLASS_LIST", + "DEFAULT_SPARSE_AUTOQUANT_CLASS_LIST", + "OTHER_AUTOQUANT_CLASS_LIST", + "ALL_AUTOQUANT_CLASS_LIST", +] + + +aten = torch.ops.aten + +AUTOQUANT_CACHE = {} + + +def check_cache(cls, shapes_and_dtype): + return AUTOQUANT_CACHE.get((cls,) + shapes_and_dtype, None) + + +def update_cache(cls, shapes_and_dtype, res): + AUTOQUANT_CACHE[(cls,) + shapes_and_dtype] = res + + +# TODO: Document the methods +class AutoQuantizableLinearWeight(torch.Tensor): + """ + A subclass of torch.Tensor that, when run, finds the best type of quantization for itself and swaps + its data with the quantized version. + + Args: + weight (torch.Tensor): The initial weight tensor. + qtensor_class_list (list): A list of tensor classes to be considered for quantization. + *args: Additional positional arguments. + mode (list, optional): A list containing mode settings for quantization. The first element is the mode type + (e.g., "relu"), and the second element is the mode value (e.g., None). Defaults to ["relu", None]. + **kwargs: Additional keyword arguments. + """ + + @staticmethod + def __new__( + cls, + weight, + qtensor_class_list, + *args, + mode=["relu", None], + min_sqnr=None, + **kwargs, + ): + kwargs["device"] = weight.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else weight.layout + ) + kwargs["dtype"] = ( + kwargs.get("dtype") if kwargs.get("dtype", False) else weight.dtype + ) + kwargs["requires_grad"] = False + shape = kwargs.pop("shape", weight.shape) + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + weight, + qtensor_class_list, + *args, + mode=["relu", None], + min_sqnr=None, + **kwargs, + ): + self.weight = weight + self.qtensor_class_list = qtensor_class_list + self.logged_data = {} + self.mode = mode + self.min_sqnr = min_sqnr + + def __repr__(self): + return ( + f"{self.__class__.__name__}(data={self.weight}, shape={self.shape}, " + f"device={self.device}, dtype={self.dtype}, qtensor_class_list={self.qtensor_class_list})" + ) + + @staticmethod + def log_shape(act_mat, w_autoquant, bias): + act_mat = act_mat.reshape(-1, act_mat.shape[-1]) + logged_dtype = act_mat.dtype + logged_shapes = ( + act_mat.shape, + w_autoquant.shape, + None if bias is None else bias.shape, + ) + shapes_and_dtype = logged_shapes + (logged_dtype,) + w_autoquant.logged_data[shapes_and_dtype] = 1 + w_autoquant.logged_data.get( + shapes_and_dtype, 0 + ) + for q_cls in w_autoquant.qtensor_class_list: + if check_cache(q_cls, shapes_and_dtype) is None: + update_cache(q_cls, shapes_and_dtype, None) + + def tune_autoquant(self, q_cls, shapes_and_dtype, best_time): + act_shape, w_shape, bias_shape, act_dtype = shapes_and_dtype + if check_cache(q_cls, shapes_and_dtype) is None: + with torch.no_grad(): + act_mat = torch.randn(act_shape, dtype=act_dtype, device=self.device) + bias = ( + None + if bias_shape is None + else torch.randn(bias_shape, dtype=act_dtype, device=self.device) + ) + try: + ref_output = AQDefaultLinearWeight._quantized_linear_op( + act_mat, self.weight, bias + ) + q_output = q_cls._quantized_linear_op( + act_mat, q_cls.from_float(self.weight), bias + ) + if ( + self.min_sqnr is not None + and (sqnr := compute_error(q_output, ref_output)) + < self.min_sqnr + ): + print( + f"skipping q_cls: {q_cls} because the sqnr is too small, minimum expected sqnr: {self.min_sqnr}, got {sqnr}" + ) + res = torch.inf + else: + res = q_cls._autoquant_test( + act_mat, self.weight, bias, best_time, self.mode + ) + except Exception as e: + print( + f"warning: failed to autoquant {q_cls.__name__} for shape: {shapes_and_dtype} due to {e}" + ) + res = torch.inf + update_cache(q_cls, shapes_and_dtype, res) + + @torch.no_grad() + def to_quantized(self, error_on_unseen, **kwargs): + if error_on_unseen and self.logged_data == {}: + raise RuntimeError( + "must run module normally to get shape, dtype info for autoquant" + ) + elif (self.logged_data == {}) and not error_on_unseen: + # default back to non-quantized weight if not seen + self = AQDefaultLinearWeight.from_float(self.weight) + return self + + # only want to print shape (at start) and final result (at end) + # once per shape+quantization subclass combination. + ran_new_benchmarks = False + print_shape_once = True + + def count_shapes(self, do_print=True): + differe_shape_count = 0 + for shapes_and_dtype, times_seen in self.logged_data.items(): + differe_shape_count += 1 + if do_print: + act_shape, weight_shape, bias_shape, dtype = shapes_and_dtype + print(f"activation_shapes: {act_shape}, times_seen: {times_seen}") + if do_print: + print( + f"weight_shape: {weight_shape}, dtype: {dtype}, bias_shape: {bias_shape}" + ) + return differe_shape_count + + # check each class + best_time = torch.inf + best_cls = None + for q_cls in self.qtensor_class_list: + # for each logged shape+dtype, benchmark + cur_time = 0 + total_seen = 0 + shape_count = count_shapes(self, do_print=False) + for shapes_and_dtype, times_seen in self.logged_data.items(): + if check_cache(q_cls, shapes_and_dtype) is None: + # only print shapes once + if print_shape_once: + print_shape_once = False + count_shapes(self, do_print=True) + + time_for_best_shape = check_cache(best_cls, shapes_and_dtype) + time_for_best_shape = ( + torch.inf + if time_for_best_shape is None + else time_for_best_shape + ) + self.tune_autoquant(q_cls, shapes_and_dtype, time_for_best_shape) + ran_new_benchmarks = True + torch._dynamo.reset() + cur_time += check_cache(q_cls, shapes_and_dtype) * times_seen + total_seen += times_seen + cur_time = cur_time / total_seen + # print aggregated time if there were multiple shapes to aggregate and some new benchmarking was done + if shape_count is not None and shape_count > 1 and ran_new_benchmarks: + print( + f">time (all shapes): {cur_time:0.4f}ms for {q_cls}, prev_best: {best_time:0.4f}ms" + ) + if cur_time != torch.inf and best_time >= cur_time: + best_time = cur_time + best_cls = q_cls + # if no new benchmarking was done, don't print the final result, it will be the same as for another layer + if ran_new_benchmarks: + print(f"best_cls={best_cls}\n") + + if best_cls is None: + best_cls = AQDefaultLinearWeight + + # TODO handle random cls args/kwargs? or should they be curried? + self = best_cls.from_float(self.weight) + return self + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.weight), + self.qtensor_class_list, + dtype=self.dtype, + mode=self.mode, + min_sqnr=self.min_sqnr, + ) + + def __tensor_flatten__(self): + return ["weight"], [ + self.qtensor_class_list, + self.mode, + self.min_sqnr, + self.dtype, + self.shape, + ] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size=None, outer_stride=None + ): + weight = tensor_data_dict["weight"] + qtensor_class_list, mode, min_sqnr, dtype, shape = tensor_attributes + return cls( + weight, + qtensor_class_list, + mode=mode, + min_sqnr=min_sqnr, + shape=shape if outer_size is None else outer_size, + dtype=dtype, + strides=outer_stride, + ) + + @classmethod + def from_float(cls, weight, qtensor_class_list, **kwargs): + return cls(weight, qtensor_class_list, **kwargs) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + kwargs = {} if kwargs is None else kwargs + + if func is torch.nn.functional.linear: + mat1, w_autoquant, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + cls.log_shape(mat1, w_autoquant, bias) + return func(mat1, w_autoquant.weight, bias) + try: + with torch._C.DisableTorchFunctionSubclass(): + return func(*args, **kwargs) + except Exception: + print(f"ERR: subclass doesn't implement {func}") + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + +@torch.no_grad() +def do_autoquant_bench(op, *args, **kwargs): + """ + runs benchmark op(*args, **kwargs) avoiding torch.compile overhead + """ + rep = kwargs.pop("rep", 100) + warmup = kwargs.pop("warmup", 25) + with torch.no_grad(): + torch.cuda.synchronize() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(*args, **kwargs) + stream.synchronize() + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + op(*args, **kwargs) + if TORCH_VERSION_AT_LEAST_2_5: + from torch._inductor.runtime.benchmarking import benchmarker + + res = benchmarker.benchmark_gpu( + lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median" + ) + elif TORCH_VERSION_AT_LEAST_2_3: + from torch._inductor.runtime.runtime_utils import do_bench_gpu + + res = do_bench_gpu( + lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median" + ) + else: + from torch._inductor.utils import do_bench + + res = do_bench( + lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median" + ) + return res + + +def _is_interpolate_mode(mode): + if ( + isinstance(mode, list) + and mode[0] == "interpolate" + and len(mode) == 2 + and isinstance(mode[1], float) + ): + return True + return False + + +def _to_float16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float16) + + +def _to_bfloat16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16) + + +def _identity(x: torch.Tensor) -> torch.Tensor: + return x + + +class AQMixin: + """ + Tests and benchmarks the autoquantization process for the given activation matrix, weight, and bias. + + Args: + act_mat (torch.Tensor): The activation matrix. + weight (torch.Tensor): The weight tensor. + bias (torch.Tensor or None): The bias tensor. + best_time (float): The best time to beat for the quantization process. + mode (list, optional): A list containing mode settings for quantization. The first element is the mode type + (e.g., "relu"), and the second element is the mode value (e.g., None). Defaults to ["relu", None]. + + Returns: + float: The benchmarked time for the autoquantization process. + """ + + @classmethod + def _autoquant_test(cls, act_mat, weight, bias, best_time, mode=["relu", None]): + w_qtensor = cls.from_float(weight) + if _is_interpolate_mode(mode): + q_c_op = torch.compile( + cls._quantized_linear_op, mode="max-autotune-no-cudagraphs" + ) + else: + func = lambda a, b, c: F.relu(cls._quantized_linear_op(F.relu(a), b, c)) + q_c_op = torch.compile(func, mode="max-autotune-no-cudagraphs") + res = do_autoquant_bench(q_c_op, act_mat, w_qtensor, bias, warmup=25, rep=100) + if res < best_time * 1.1: + res2 = do_autoquant_bench( + q_c_op, act_mat, w_qtensor, bias, warmup=25, rep=900 + ) + res = res2 * 0.9 + res * 0.1 + print(f">>time: {res:0.3f}ms for {cls}, to_beat: {best_time:0.3f}ms ") + return res + + +class AQInt8DynamicallyQuantizedLinearWeight(AQMixin, LinearActivationQuantizedTensor): + """ + AutoQuantizable version of Int8DynamicallyQuantizedLinearWeight + """ + + aq_layout: Layout = PlainLayout() + + @classmethod + def from_float(cls, weight): + if weight.dim() != 2: + return weight + + # TODO test if this is valid + # in_features = weight.shape[1] + # int8 dynamic quantization only has benefit when in_feature > 16 + # if in_features <= 16: + # return weight + + # avoid circular dep + from torchao.dtypes import to_affine_quantized_intx + from torchao.quantization.quant_api import ( + _int8_symm_per_token_reduced_range_quant, + ) + + # input settings + input_quant_func = _int8_symm_per_token_reduced_range_quant + + # weight settings + mapping_type = MappingType.SYMMETRIC + + def get_weight_block_size(x): + return (1, x.shape[1]) + + target_dtype = torch.int8 + eps = torch.finfo(torch.float32).eps + zero_point_dtype = torch.int64 + _layout = cls.aq_layout + block_size = get_weight_block_size(weight) + + weight = to_affine_quantized_intx( + weight, + mapping_type, + block_size, + target_dtype, + eps=eps, + zero_point_dtype=zero_point_dtype, + _layout=_layout, + ) + weight = super(AQInt8DynamicallyQuantizedLinearWeight, cls).from_float( + weight, input_quant_func + ) + return weight + + @classmethod + def _autoquant_test(cls, act_mat, weight, bias, best_time, mode=["relu", None]): + """ + Tests and benchmarks the autoquantization process with special handling for interpolate mode. + + Args: + act_mat (torch.Tensor): The activation matrix. + weight (torch.Tensor): The weight tensor. + bias (torch.Tensor or None): The bias tensor. + best_time (float): The best time to beat for the quantization process. + mode (list, optional): A list containing mode settings for quantization. The first element is the mode type + (e.g., "relu"), and the second element is the mode value (e.g., None). Defaults to ["relu", None]. + + Returns: + float: The benchmarked time for the autoquantization process. + """ + if not _is_interpolate_mode(mode): + return super()._autoquant_test(act_mat, weight, bias, best_time, mode) + + # SAM best is between .8 and 1, SDXL also performs best in this range + INTERPOLATION_CONSTANT = mode[1] + w_qtensor = cls.from_float(weight) + x_vals_int8, x_scales = quantize_activation_per_token_absmax( + act_mat.reshape(-1, act_mat.shape[-1]) + ) + quantized_matmul = ( + lambda x_vals_int8, x_scales, w_vals_int8: safe_int_mm( + x_vals_int8, w_vals_int8 + ) + * x_scales + ) + q_c_matmul = torch.compile(quantized_matmul, mode="max-autotune-no-cudagraphs") + with torch.no_grad(): + w_vals_int8 = ( + w_qtensor.original_weight_tensor.tensor_impl.int_data.contiguous().t() + ) + res_matmul = do_autoquant_bench( + q_c_matmul, x_vals_int8, x_scales.reshape(-1, 1), w_vals_int8 + ) + print( + f">>time: {res_matmul:0.3f}ms for {cls} matmul, to_beat: {best_time:0.3f}ms" + ) + + # if the (much faster) matmul kernel is already beat, don't bother benchmarking full op + if res_matmul >= best_time: + return res_matmul + + # calculate what time full op needs to beat for dynamic quant to be best given INTERPOLATION_CONSTANT + to_beat = best_time + INTERPOLATION_CONSTANT / (1 - INTERPOLATION_CONSTANT) * ( + best_time - res_matmul + ) + res = super()._autoquant_test(act_mat, weight, bias, to_beat) + max_int_const_win = (best_time - res_matmul) / (res - res_matmul) + res_f = INTERPOLATION_CONSTANT * res + (1 - INTERPOLATION_CONSTANT) * res_matmul + print( + f">>time: {res_f:0.3f}ms for {cls} interpolated, breakeven constant: {max_int_const_win:0.2f}" + ) + return res_f + + +class AQInt8DynamicallyQuantizedSemiSparseLinearWeight( + AQInt8DynamicallyQuantizedLinearWeight +): + aq_layout: Layout = SemiSparseLayout() + + @classmethod + def _autoquant_test(cls, act_mat, weight, bias, best_time, mode=["relu", None]): + return super()._autoquant_test(act_mat, weight, bias, best_time, None) + + +class AQInt8WeightOnlyQuantizedLinearWeight(AffineQuantizedTensor, AQMixin): + """ + AutoQuantizable version of Int8WeightOnlyQuantizedLinearWeight + """ + + @classmethod + def from_float(cls, weight): + mapping_type = MappingType.SYMMETRIC + target_dtype = torch.int8 + eps = torch.finfo(torch.float32).eps + zero_point_dtype = torch.int64 + block_size = (1, weight.shape[1]) + return super(AQInt8WeightOnlyQuantizedLinearWeight, cls).from_hp_to_intx( + weight, + mapping_type, + block_size, + target_dtype, + eps=eps, + zero_point_dtype=zero_point_dtype, + ) + + +class AQInt8WeightOnlyQuantizedLinearWeight2( + AQInt8WeightOnlyQuantizedLinearWeight, AQMixin +): + """ + AutoQuantizable version of Int8WeightOnlyQuantizedLinearWeight that + uses a different kernel + """ + + @staticmethod + def _quantized_linear_op(act_mat, w_qtensor, bias): + """ + Performs the quantized linear operations + + Args: + act_mat (torch.Tensor): The activation matrix. + w_qtensor (torch.Tensor): The quantized weight tensor. + bias (torch.Tensor or None): The bias tensor. + + Returns: + torch.Tensor: The result of the quantized operation. + """ + orig_dtype = act_mat.dtype + orig_shape = act_mat.shape + act_mat = act_mat.reshape(-1, act_mat.shape[-1], 1) + y = (act_mat * w_qtensor.tensor_impl.int_data.t().unsqueeze(0)).sum(dim=-2) + y = y.reshape(*orig_shape[:-1], y.shape[-1]) * w_qtensor.tensor_impl.scale + if bias is not None: + y += bias + return y.to(orig_dtype) + + @classmethod + def _autoquant_test(cls, act_mat, *args): + # if act_mat has batchsize>2 don't use this kernel + if act_mat.reshape(-1, act_mat.shape[-1]).shape[0] > 32: + return torch.inf + return super()._autoquant_test(act_mat, *args) + + +class AQInt8WeightOnlyQuantizedLinearWeight3( + AQInt8WeightOnlyQuantizedLinearWeight, AQMixin +): + """ + AutoQuantizable version of Int8WeightOnlyQuantizedLinearWeight that + uses a different kernel + """ + + @staticmethod + def _quantized_linear_op(act_mat, w_qtensor, bias): + orig_shape = act_mat.shape + y = torch.mm( + act_mat.reshape(-1, orig_shape[-1]), + w_qtensor.tensor_impl.int_data.t() * w_qtensor.tensor_impl.scale, + ) + y = y.reshape(*orig_shape[:-1], y.shape[-1]) + if bias is not None: + y += bias + return y + + +class AQInt4G32WeightOnlyQuantizedLinearWeight( + LinearActivationQuantizedTensor, AQMixin +): + """ + AutoQuantizable version of int4_weight_only + """ + + group_size: int = 32 + # can't override the `layout` attribute + aq_layout: Layout = TensorCoreTiledLayout(inner_k_tiles=8) + + @classmethod + def from_float(cls, weight): + from torchao.dtypes import to_affine_quantized_intx + + group_size = cls.group_size + _layout = cls.aq_layout + + if weight.shape[-1] % group_size != 0: + return weight + + input_quant_func = None + + # NOTE: we only convert activation dtype and weight dtype here + # because the kernel implementation for both TensorCoreTiledLayout and MarlinSparseLayout + # can work with multiple bias dtypes (by converting bias to the dtype of activation) + if ( + isinstance(_layout, TensorCoreTiledLayout) + and weight.dtype != torch.bfloat16 + ): + weight = weight.to(torch.bfloat16) + input_quant_func = _to_bfloat16 + elif isinstance(_layout, MarlinSparseLayout) and weight.dtype != torch.float16: + weight = weight.to(torch.float16) + input_quant_func = _to_float16 + else: + input_quant_func = _identity + + use_hqq = True + mapping_type = MappingType.ASYMMETRIC + block_size = (1, group_size) + target_dtype = torch.int32 + quant_min = 0 + quant_max = 15 + eps = 1e-6 + preserve_zero = False + zero_point_dtype = torch.bfloat16 + zero_point_domain = ZeroPointDomain.FLOAT + + if isinstance(_layout, MarlinSparseLayout): + mapping_type = MappingType.SYMMETRIC + preserve_zero = True + zero_point_domain = ZeroPointDomain.INT + use_hqq = False + + weight = to_affine_quantized_intx( + weight, + mapping_type, + block_size, + target_dtype, + quant_min, + quant_max, + eps, + zero_point_dtype=zero_point_dtype, + preserve_zero=preserve_zero, + zero_point_domain=zero_point_domain, + _layout=_layout, + use_hqq=use_hqq, + ) + + return super(AQInt4G32WeightOnlyQuantizedLinearWeight, cls).from_float( + weight, input_quant_func + ) + + +class AQInt4G64WeightOnlyQuantizedLinearWeight( + AQInt4G32WeightOnlyQuantizedLinearWeight +): + group_size: int = 64 + + +class AQInt4G128WeightOnlyQuantizedLinearWeight( + AQInt4G32WeightOnlyQuantizedLinearWeight +): + group_size: int = 128 + + +class AQInt4G256WeightOnlyQuantizedLinearWeight( + AQInt4G32WeightOnlyQuantizedLinearWeight +): + group_size: int = 256 + + +class AQInt4G128WeightOnlyQuantizedMarlinSparseLinearWeight( + AQInt4G32WeightOnlyQuantizedLinearWeight +): + group_size: int = 128 + aq_layout: Layout = MarlinSparseLayout() + + +class AQGemliteInt4G32WeightOnlyQuantizedLinearWeight( + LinearActivationQuantizedTensor, AQMixin +): + group_size: int = 32 + + @classmethod + def from_float(cls, weight): + from torchao.dtypes import to_affine_quantized_intx + from torchao.dtypes.uintx.gemlite_layout import get_gemlite_aqt_kwargs + + if weight.dtype != torch.float16: + weight = weight.to(torch.float16) + + bit_width = 4 + packing_bitwidth = 32 + contiguous = None + use_hqq = True + aqt_kwargs = get_gemlite_aqt_kwargs( + weight, cls.group_size, bit_width, packing_bitwidth, contiguous, use_hqq + ) + weight = to_affine_quantized_intx(weight, **aqt_kwargs) + input_quant_func = _to_float16 + + return super(AQGemliteInt4G32WeightOnlyQuantizedLinearWeight, cls).from_float( + weight, input_quant_func + ) + + +class AQGemliteInt4G64WeightOnlyQuantizedLinearWeight( + AQGemliteInt4G32WeightOnlyQuantizedLinearWeight +): + group_size: int = 64 + + +class AQGemliteInt4G128WeightOnlyQuantizedLinearWeight( + AQGemliteInt4G32WeightOnlyQuantizedLinearWeight +): + group_size: int = 128 + + +class AQGemliteInt4G256WeightOnlyQuantizedLinearWeight( + AQGemliteInt4G32WeightOnlyQuantizedLinearWeight +): + group_size: int = 256 + + +class AQDefaultLinearWeight(torch.Tensor, AQMixin): + """ + A class to be used in concert with AutoQuantizableLinearWeight to provide a + default/non-quantized option. Only implements the bare minimum needed to work with the + AutoQuantizableLinearWeight class using the same interfaces that would normally be + used by QTensor subclasses but for a default linear op instead. Result of from_float + is not a tensor subclass, but rather the float tensor. + """ + + def __init__(self): + super().__init__() + + @staticmethod + def _quantized_linear_op(act_mat, w_qtensor, bias): + return torch.nn.functional.linear(act_mat, w_qtensor, bias) + + @classmethod + def from_float(cls, weight): + return weight + + +# TODO: remove skip_weight_conversion arg +class Float32Tensor(TorchAOBaseTensor): + """Tensor subclass tensor for fp32 dtype""" + + @staticmethod + def __new__(cls, weight, skip_weight_conversion=False): + kwargs = {} + kwargs["device"] = weight.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else weight.layout + ) + kwargs["dtype"] = weight.dtype + kwargs["requires_grad"] = False + shape = weight.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) + + def __init__(self, weight, skip_weight_conversion=False): + self.weight = weight if skip_weight_conversion else weight.to(torch.float32) + + @staticmethod + def _quantized_linear_op(act_mat, w_qtensor, bias): + _DTYPE = torch.float32 + orig_dtype = act_mat.dtype + return torch.nn.functional.linear( + act_mat.to(_DTYPE), + w_qtensor.weight, + bias.to(_DTYPE) if bias is not None else bias, + ).to(dtype=orig_dtype) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.weight), + ) + + @classmethod + def from_float(cls, weight): + return cls(weight) + + +@Float32Tensor.implements([torch.nn.functional.linear, aten.linear.default]) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias) + + +@Float32Tensor.implements(aten.detach.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + +@Float32Tensor.implements(aten.clone.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + +@Float32Tensor.implements(aten._to_copy.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].to(*args[1:], **kwargs)._apply_fn_to_data(torch.clone), + ) + + +class BFloat16Tensor(Float32Tensor): + def __init__(self, weight, skip_weight_conversion=False): + self.weight = weight if skip_weight_conversion else weight.to(torch.bfloat16) + + @staticmethod + def _quantized_linear_op(act_mat, w_qtensor, bias): + _DTYPE = torch.bfloat16 + orig_dtype = act_mat.dtype + return torch.nn.functional.linear( + act_mat.to(_DTYPE), + w_qtensor.weight, + bias.to(_DTYPE) if bias is not None else bias, + ).to(dtype=orig_dtype) + + @classmethod + def from_float(cls, weight, skip_weight_conversion=False): + return cls(weight, skip_weight_conversion) + + +class Float16Tensor(Float32Tensor): + def __init__(self, weight, skip_weight_conversion=False): + self.weight = weight if skip_weight_conversion else weight.to(torch.float16) + + @staticmethod + def _quantized_linear_op(act_mat, w_qtensor, bias): + _DTYPE = torch.float16 + orig_dtype = act_mat.dtype + return torch.nn.functional.linear( + act_mat.to(_DTYPE), + w_qtensor.weight, + bias.to(_DTYPE) if bias is not None else bias, + ).to(dtype=orig_dtype) + + @classmethod + def from_float(cls, weight, skip_weight_conversion=False): + return cls(weight, skip_weight_conversion) + + +class AQFloat32LinearWeight(Float32Tensor, AQMixin): + """ + AutoQuantizable version for float32 precision weight + + (also converts input activation and bias to float32, and restores the original precision after + linear) + """ + + @classmethod + def from_float(cls, weight): + return super(AQFloat32LinearWeight, cls).from_float(weight) + + +class AQBFloat16LinearWeight(BFloat16Tensor, AQMixin): + """ + AutoQuantizable version for bfloat16 precision weight + + (also converts input activation and bias to bfloat16, and restores the original precision after + linear) + """ + + @classmethod + def from_float(cls, weight): + return super(AQBFloat16LinearWeight, cls).from_float(weight) + + +class AQFloat16LinearWeight(Float16Tensor, AQMixin): + """ + AutoQuantizable version for float16 precision weight + + (also converts input activation and bias to float16, and restores the original precision after + linear) + """ + + @classmethod + def from_float(cls, weight): + return super(AQFloat16LinearWeight, cls).from_float(weight) + + +class AQFloat8WeightOnlyQuantizedLinearWeight(AffineQuantizedTensor, AQMixin): + """ + AutoQuantizable version of Float8WeightOnlyQuantizedLinearWeight for target_dtype=torch.float8_e4m3fn + """ + + target_dtype: torch.dtype = torch.float8_e4m3fn + + @staticmethod + def _quantized_linear_op(act_mat, w_qtensor, bias): + return torch.nn.functional.linear(act_mat, w_qtensor.dequantize(), bias) + + @classmethod + def from_float(cls, weight): + block_size = (1, weight.shape[1]) + return super(AQFloat8WeightOnlyQuantizedLinearWeight, cls).from_hp_to_floatx( + weight, block_size, target_dtype=cls.target_dtype, _layout=Float8Layout() + ) + + +class AQFloat8PerRowScalingDynamicallyQuantizedLinearWeight(AQMixin, BFloat16Tensor): + """ + AutoQuantizable version of Float8DynamicallyQuantizedLinearWeight using per row scaling + """ + + activation_granularity = PerRow() + + @classmethod + def from_float(cls, weight): + # avoid circular dep + from torchao.dtypes import to_affine_quantized_floatx + from torchao.quantization.quant_api import _input_activation_quant_func_fp8 + + # weight settings + def get_weight_block_size(x): + return (1, x.shape[1]) + + target_dtype = torch.float8_e4m3fn + + # input settings + def get_per_token_block_size(x): + block_size = list(x.shape) + for i in range(len(block_size) - 1): + block_size[i] = 1 + return block_size + + input_target_dtype = torch.float8_e4m3fn + _layout = Float8Layout(mm_config=Float8MMConfig(use_fast_accum=True)) + # TODO: make this serializable + input_quant_func = _input_activation_quant_func_fp8 + input_quant_kwargs = { + "activation_granularity": cls.activation_granularity, + "activation_dtype": input_target_dtype, + } + block_size = get_weight_block_size(weight) + + weight = to_affine_quantized_floatx( + input_float=weight, + block_size=block_size, + target_dtype=target_dtype, + _layout=_layout, + scale_dtype=torch.float32, + ) + weight = to_linear_activation_quantized( + weight, input_quant_func, quant_kwargs=input_quant_kwargs + ) + # at inference time, + # we first convert the input, weight and bias to bfloat16, and then quantize activation + # and then dispatch to the quantized ops + return super( + AQFloat8PerRowScalingDynamicallyQuantizedLinearWeight, cls + ).from_float(weight, skip_weight_conversion=True) + + +class AQFloat8PerTensorScalingDynamicallyQuantizedLinearWeight( + AQMixin, LinearActivationQuantizedTensor +): + """ + AutoQuantizable version of Float8DynamicallyQuantizedLinearWeight using per tensor scaling + """ + + activation_granularity = PerTensor() + + @classmethod + def from_float(cls, weight): + # avoid circular dep + from torchao.dtypes import to_affine_quantized_floatx + from torchao.quantization.quant_api import _input_activation_quant_func_fp8 + + # weight settings + def get_weight_block_size(x): + assert x.ndim == 2, "Only works for 2D tensors" + return x.shape + + target_dtype = torch.float8_e4m3fn + input_target_dtype = torch.float8_e4m3fn + _layout = Float8Layout(mm_config=Float8MMConfig(use_fast_accum=True)) + # TODO: test serializable + input_quant_func = _input_activation_quant_func_fp8 + input_quant_args = { + "activation_granularity": cls.activation_granularity, + "activation_dtype": input_target_dtype, + } + block_size = get_weight_block_size(weight) + weight = to_affine_quantized_floatx( + input_float=weight, + block_size=block_size, + target_dtype=target_dtype, + _layout=_layout, + scale_dtype=torch.float32, + ) + weight = super( + AQFloat8PerTensorScalingDynamicallyQuantizedLinearWeight, cls + ).from_float(weight, input_quant_func, input_quant_args) + return weight + + +# here we don't include int4 quantization in since int8 tends to be a better apples to apples comparison +DEFAULT_AUTOQUANT_CLASS_LIST = [ + AQDefaultLinearWeight, + AQInt8WeightOnlyQuantizedLinearWeight, + AQInt8WeightOnlyQuantizedLinearWeight2, + # AQInt8WeightOnlyQuantizedLinearWeight3, + # TODO this gets picked in places where it makes perf worse, why? + AQInt8DynamicallyQuantizedLinearWeight, +] + +DEFAULT_INT4_AUTOQUANT_CLASS_LIST = [ + AQDefaultLinearWeight, + AQInt8DynamicallyQuantizedLinearWeight, + AQInt4G64WeightOnlyQuantizedLinearWeight, +] + +GEMLITE_INT4_AUTOQUANT_CLASS_LIST = [ + AQDefaultLinearWeight, + AQInt8DynamicallyQuantizedLinearWeight, + AQGemliteInt4G64WeightOnlyQuantizedLinearWeight, +] + +DEFAULT_FLOAT_AUTOQUANT_CLASS_LIST = [ + AQFloat32LinearWeight, + AQBFloat16LinearWeight, + AQFloat16LinearWeight, +] + +OTHER_AUTOQUANT_CLASS_LIST = [ + AQDefaultLinearWeight, + AQFloat8WeightOnlyQuantizedLinearWeight, + AQFloat8PerRowScalingDynamicallyQuantizedLinearWeight, + AQFloat8PerTensorScalingDynamicallyQuantizedLinearWeight, +] + +DEFAULT_SPARSE_AUTOQUANT_CLASS_LIST = [ + AQDefaultLinearWeight, + # TODO: investigate why there are some problems when adding sparse kernels for sam2 + AQInt4G128WeightOnlyQuantizedMarlinSparseLinearWeight, + # some errors when calling cusparse kernels when running on sam2 + AQInt8DynamicallyQuantizedSemiSparseLinearWeight, +] + +ALL_AUTOQUANT_CLASS_LIST = ( + DEFAULT_AUTOQUANT_CLASS_LIST + + DEFAULT_INT4_AUTOQUANT_CLASS_LIST + + DEFAULT_FLOAT_AUTOQUANT_CLASS_LIST +) + +# add gemlite options +ALL_AUTOQUANT_CLASS_LIST += [ + AQGemliteInt4G64WeightOnlyQuantizedLinearWeight, +] + +if is_sm_at_least_89(): + ALL_AUTOQUANT_CLASS_LIST += [ + AQFloat8WeightOnlyQuantizedLinearWeight, + AQFloat8PerTensorScalingDynamicallyQuantizedLinearWeight, + ] + +if is_sm_at_least_90(): + ALL_AUTOQUANT_CLASS_LIST += [AQFloat8PerRowScalingDynamicallyQuantizedLinearWeight] + +if not is_sm_at_least_89(): + ALL_AUTOQUANT_CLASS_LIST += DEFAULT_SPARSE_AUTOQUANT_CLASS_LIST + +# deduplicate +ALL_AUTOQUANT_CLASS_LIST = list(set(ALL_AUTOQUANT_CLASS_LIST)) + + +def _change_linears_to_autoquantizable(model, **kwargs): + """ + Converts all linear weight tensors to the + AutoQuantizableLinearWeight tensor subclass. Expectation is that this is followed + by running the model and then calling _change_autoquantizable_to_quantized + """ + from torchao.quantization.quant_api import _is_linear + + filter_fn = kwargs.pop("filter_fn", _is_linear) + _ = kwargs.pop( + "error_on_unseen", True + ) # same kwargs used for this and to_quantized + kwargs["qtensor_class_list"] = kwargs.get( + "qtensor_class_list", DEFAULT_AUTOQUANT_CLASS_LIST + ) + kwargs["mode"] = kwargs.get("mode", ["relu", None]) + kwargs["min_sqnr"] = kwargs.get("min_sqnr", None) + from torchao.quantization.quant_api import ( + _get_subclass_inserter, + _replace_with_custom_fn_if_matches_filter, + ) + + _replace_with_custom_fn_if_matches_filter( + model, + _get_subclass_inserter(AutoQuantizableLinearWeight, **kwargs), + filter_fn if filter_fn is not None else _is_linear, + ) + + +def _change_autoquantizable_to_quantized( + model, supress_autoquant_errors=True, **kwargs +): + """ + Converts AutoQuantizableLinearWeight tensor subclasses + to various quantized/non-quantized tensor subclasses depending + on benchmark results. Expectation is that these modules are + torch.compiled afterwards. + """ + hold_automatic_dynamic_shapes = torch._dynamo.config.automatic_dynamic_shapes + torch._dynamo.config.automatic_dynamic_shapes = False + + if supress_autoquant_errors: + hold_supress_errors = torch._dynamo.config.suppress_errors + torch._dynamo.config.suppress_errors = True + import logging + + torch._logging.set_logs(inductor=logging.CRITICAL, dynamo=logging.CRITICAL) + filter_fn = kwargs.pop( + "filter_fn", + lambda mod, *args: hasattr(mod, "weight") + and isinstance(mod.weight, AutoQuantizableLinearWeight), + ) + error_on_unseen = kwargs.pop("error_on_unseen", True) + from torchao.quantization.quant_api import ( + _get_subclass_inserter, + _replace_with_custom_fn_if_matches_filter, + ) + + _replace_with_custom_fn_if_matches_filter( + model, + _get_subclass_inserter( + AutoQuantizableLinearWeight, + method="to_quantized", + error_on_unseen=error_on_unseen, + **kwargs, + ), + filter_fn, + ) + # undo dynamic shape change + torch._dynamo.config.automatic_dynamic_shapes = hold_automatic_dynamic_shapes + + # undo error supression + if supress_autoquant_errors: + torch._dynamo.config.suppress_errors = hold_supress_errors + torch._logging.set_logs() + torch._dynamo.reset() + + +# TODO: example_input seems weird to include in the API +# TODO: Document all the modes +# TODO: Mode being a list is weird, should be a string or some object +@torch.no_grad() +def autoquant( + model, + example_input=None, + qtensor_class_list=DEFAULT_AUTOQUANT_CLASS_LIST, + filter_fn=None, + mode=["interpolate", 0.85], + manual=False, + set_inductor_config=True, + supress_autoquant_errors=True, + min_sqnr=None, + **aq_kwargs, +): + """ + Autoquantization is a process which identifies the fastest way to quantize each layer of a model over some set of potential + qtensor subclasses. + + Autoquantization happens in three steps: + + 1-Prepare Model: the model is searched for Linear layers whose weights are exchanged for AutoQuantizableLinearWeight. + 2-Shape Calibration: the user runs the model on one or more inputs, the details of the activation shape/dtype seen by + the AutoQuantizableLinearWeight are recorded so we know what shapes/dtypes to use in order to optimize the quantized op in step 3 + 3-Finalize Autoquantization: for each AutoQuantizableLinearWeight, benchmarks are run for each shape/dtype on each member of the qtensor_class_list. + the fastest option is picked, resulting in a highly performant model + + This autoquant function performs step 1. Steps 2 and 3 can be completed by simply running the model. + If `example_input` is provided, this function also runs the model (which completes steps 2 and 3). + This autoquant api can handle models which have already had torch.compile applied to them, in which case, once the model is run and quantized, + the torch.compile process normally proceeds as well. + + To optimize over a combination of input shapes/dtypes, the user can set manual=True, run the model with all desired shapes/dtypes, then + call model.finalize_autoquant to finalize the quantization once the desired set of inputs have been logged. + + Args: + model (torch.nn.Module): The model to be autoquantized. + example_input (Any, optional): An example input for the model. If provided, the function performs a forward pass + on this input (which fully autoquantizes the model unless manual=True). Defaults to None. + qtensor_class_list (list, optional): A list of tensor classes to be used for quantization. Defaults to DEFAULT_AUTOQUANT_CLASS_LIST. + filter_fn (callable, optional): A filter function to apply to the model parameters. Defaults to None. + mode (list, optional): A list containing mode settings for quantization. The first element is the mode type (e.g., "interpolate"), + and the second element is the mode value (e.g., 0.85). Defaults to ["interpolate", .85]. + manual (bool, optional): Whether to stop shape calibration and do autoquant after a single run (default, False) or to wait for + the user to call model.finalize_autoquant (True) so inputs with several shapes/dtypes can be logged. + set_inductor_config (bool, optional): Whether to automatically use recommended inductor config settings (defaults to True) + supress_autoquant_errors (bool, optional): Whether to suppress errors during autoquantization. (defaults to True) + min_sqnr (float, optional): minimum acceptable signal to quantization noise ration (https://en.wikipedia.org/wiki/Signal-to-quantization-noise_ratio) for output of quantized layer v.s. non-quantized layer, this is used to filter + out quantization methods that causes too large numerical impact, user can start with a resaonable + number like 40 and adjust depending on the result + **aq_kwargs: Additional keyword arguments for the autoquantization process. + + Returns: + torch.nn.Module: The autoquantized and wrapped model. If `example_input` is provided, the function performs a forward pass + on the input and returns the result of the forward pass. + + Example usage: + torchao.autoquant(torch.compile(model)) + model(*example_input) + + # multiple input shapes + torchao.autoquant(model, manual=True) + model(*example_input1) + model(*example_input2) + model.finalize_autoquant() + """ + if set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + if qtensor_class_list is OTHER_AUTOQUANT_CLASS_LIST: + assert torch.cuda.is_available() and torch.cuda.get_device_capability() >= ( + 8, + 9, + ), "float8 requires CUDA arch >= 8.9" + + # perform initial swap from linear weights + # to AutoQuantizableLinearWeight + _change_linears_to_autoquantizable( + model, + filter_fn=filter_fn, + qtensor_class_list=qtensor_class_list, + mode=mode, + min_sqnr=min_sqnr, + **aq_kwargs, + ) + + # access actual model of torch.compile wrapper if needed + is_compiled = isinstance(model, torch._dynamo.eval_frame.OptimizedModule) + if is_compiled: + real_model = model._orig_mod + else: + real_model = model + + if manual: + # we don't want model.forward to trigger + # torch.compilation + if is_compiled: + real_model.old_forward = model.forward + model.forward = real_model.forward + + # we want to automatically do autoquant after a single model run + # and have it occur before torch.compilation if applicable + else: + # the hook we will use to intercept the model forward and perform + # autoquantization + def autoquant_prehook(module, args, kwargs): + real_model.forward(*args, **kwargs) + module.finalize_autoquant() + return args, kwargs + + # the autoquant_prehook intercepts the forward call, performs logging then + # does autoquantization. if model is a torch.compile wrapper, it then + # does the tracing/compile since the prehook is naturally followed by the normal. + # model run. + handle = model.register_forward_pre_hook(autoquant_prehook, with_kwargs=True) + + # note the torch.compile wrapper (eval_frame) moves the assignment of any assigned + # attributes to the inner model that didn't exist before, so we have to call delattr on the inner model + def finalize_autoquant(): + _change_autoquantizable_to_quantized( + real_model, + supress_autoquant_errors, + **aq_kwargs, + ) + if hasattr(real_model, "old_forward"): + model.forward = real_model.old_forward + delattr(real_model, "old_forward") + if hasattr(real_model, "finalize_autoquant"): + delattr(real_model, "finalize_autoquant") + if not manual: + handle.remove() + + real_model.finalize_autoquant = finalize_autoquant + + # if example input was provided, check it and run it + if isinstance(example_input, torch.Tensor): + example_input = [example_input] + if isinstance(example_input, (tuple, list)): + model(*example_input) + + return model + + +if TORCH_VERSION_AT_LEAST_2_5: + torch.serialization.add_safe_globals(ALL_AUTOQUANT_CLASS_LIST) + torch.serialization.add_safe_globals( + [ + _to_float16, + _to_bfloat16, + _identity, + ] + ) diff --git a/lib/python3.12/site-packages/torchao/quantization/dynamic_quant.py b/lib/python3.12/site-packages/torchao/quantization/dynamic_quant.py new file mode 100644 index 0000000000000000000000000000000000000000..61c6b0dc075a28807e837831aa204b3da6a4d466 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/dynamic_quant.py @@ -0,0 +1,85 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn + +from .utils import ( + dynamically_quantize_per_channel, + quant_int8_dynamic_per_token_linear, +) + +__all__ = ["DynamicallyPerAxisQuantizedLinear"] + + +class DynamicallyPerAxisQuantizedLinear(torch.nn.Linear): + """ + This class is a replacement for `torch.nn.Linear`. It implements a + quantized matmul using int8 dynamic symmetric per-token activation, + and int8 symmetric per-channel weight quantization + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + ) -> None: + super().__init__(in_features, out_features, bias) + + def forward(self, X: torch.Tensor, *args, **kwargs) -> torch.Tensor: + """ + Performs the forward pass of the quantized linear layer which consists + of int8 dynamic symmetric per-token activation and int8 symmetric per-channel weight + quantization + + Args: + X (torch.Tensor): The input floating point tensor to the quantized linear layer. + + Returns: + torch.Tensor: The output floating point tensor after the quantized matmul and rescale. + + """ + + Y = quant_int8_dynamic_per_token_linear( + X, self.W_int_repr_t, self.W_scales, self.bias, X.dtype + ) + return Y + + @classmethod + def from_float(cls, mod: torch.nn.Linear) -> "DynamicallyPerAxisQuantizedLinear": + """ + Converts a `mod` of class `torch.nn.Linear` to the + `DynamicallyPerAxisQuantizedLinear` class + + Args: + mod (torch.nn.Linear): The original `torch.nn.Linear` module to convert. + + Returns: + DynamicallyPerAxisQuantizedLinear: The converted quantized linear module. + + """ + + # create the new module with a toy size to ensure initialization is fast + fake_in_features, fake_out_features = 8, 8 + new_mod = cls( + fake_in_features, + fake_out_features, + bias=mod.bias is not None, + ) + new_mod.in_features = mod.in_features + new_mod.out_features = mod.out_features + W_int_repr, W_scales, _W_zps = dynamically_quantize_per_channel( + mod.weight, -128, 127, torch.int8 + ) + new_mod.register_buffer("W_int_repr_t", W_int_repr.contiguous().t()) + new_mod.W_scales = nn.Parameter(W_scales) + new_mod.bias = mod.bias + del new_mod.weight + + device_to_use = next(mod.parameters()).device + new_mod.to(device_to_use) + return new_mod diff --git a/lib/python3.12/site-packages/torchao/quantization/granularity.py b/lib/python3.12/site-packages/torchao/quantization/granularity.py new file mode 100644 index 0000000000000000000000000000000000000000..8cdae9caaa330bdbfb842988a7c8a2f77a3d90f2 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/granularity.py @@ -0,0 +1,101 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class Granularity: + """ + Base class for representing the granularity of quantization. + + This class serves as a parent for specific granularity types used in + quantization operations, such as per-tensor or per-axis quantization. + """ + + pass + + +@dataclass(frozen=True) +class PerTensor(Granularity): + """ + Represents per-tensor granularity in quantization. + + This granularity type calculates the quantization parameters + based off the entire tensor. + """ + + pass + + +@dataclass(frozen=True) +class PerAxis(Granularity): + """ + Represents per-axis granularity in quantization. + + This granularity type calculates different quantization parameters + along a specified axis of the tensor. + + For example if the input tensor is shape [8, 16] and axis=0, then + the quantization parameters are calculated for each row of the tensor. + Giving a total of 8 quantization parameters. + + Attributes: + axis (int): The axis along which reduction is performed. + """ + + axis: int + + +@dataclass(frozen=True) +class PerGroup(Granularity): + """ + Represents per-channel group granularity in quantization. + + This granularity type calculates different quantization parameters + for each group of elements. + + For example if the input tensor is shape [8, 16], and the group size is 4, then + the input tensor is reshaped to [64, 4] + quantization parameters are calculated for each group of 4 elements, + giving a total of 64 quantization parameters. + + Attributes: + group_size (int): The size of each quantization group + + """ + + group_size: int + + +class PerRow(Granularity): + """ + Represents row-wise granularity in quantization. + + This is a special case of per-axis quantization and is unique to Float8 matmuls + where the input is quantized with a block_size of (1, ..., input.shape[-1]). And the weight + is quantized with a block_size of (1, weight.shape[1]). + """ + + pass + + +class PerToken(Granularity): + """ + Represents per-token granularity in quantization. + + This granularity type calculates a different set of quantization parameters + for each token, which is represented as the last dimension of the tensor. + + For example, if the input tensor has shape [2, 3, 4], then there are 6 tokens + with 4 elements each, and we will calculate 6 sets of quantization parameters, + one for each token. + + If the input tensor has only two dimensions, e.g. [8, 16], then this is + equivalent to `PerAxis(axis=0)`, which yields 8 sets of quantization parameters. + """ + + pass diff --git a/lib/python3.12/site-packages/torchao/quantization/linear_activation_quantized_tensor.py b/lib/python3.12/site-packages/torchao/quantization/linear_activation_quantized_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..aa946c064ffe6998f3b54595442746b68640325a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/linear_activation_quantized_tensor.py @@ -0,0 +1,295 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Any, Callable, Dict, Optional + +import torch +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_5, + TorchAOBaseTensor, +) + +__all__ = [ + "LinearActivationQuantizedTensor", + "to_linear_activation_quantized", +] + +aten = torch.ops.aten + + +class LinearActivationQuantizedTensor(TorchAOBaseTensor): + """ + Applies activation quantization for linear operator, this is used to support + dynamic quantization, user can pass in a `input_quant_func` + that is used to quantize the activation + + Args: + `original_weight_tensor`: the weight tensor, if weight need to be quantized as well, we'd need + to apply quantization to weight first, e.g. for int8 dynamic activation int8 weight quantization + we will first apply int8 quantization to weight and then apply LinearActivationQuantizedTensor + on top of it + `input_quant_func` (Callable[[torch.Tensor], torch.Tensor]): a function that takes a high precision floating point tensor and returns + a quantized tensor, this is used to quantize input + `quant_kwargs` (Dict[str, Any]): Additional keyword arguments for the quantization function. + Restriction: Must not contain tensor values. + """ + + quant_kwargs: Dict[str, Any] + + def __new__( + cls, + original_weight_tensor: torch.Tensor, + input_quant_func: Callable, + quant_kwargs: Dict[str, Any], + ): + kwargs = {} + dtype = original_weight_tensor.dtype + kwargs["dtype"] = dtype + kwargs["requires_grad"] = False + kwargs["device"] = original_weight_tensor.device + shape = original_weight_tensor.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + original_weight_tensor: torch.Tensor, + input_quant_func: Callable[[torch.Tensor], torch.Tensor], + quant_kwargs: Dict[str, Any], + ): + self.original_weight_tensor = original_weight_tensor + self.input_quant_func = input_quant_func + self.quant_kwargs = quant_kwargs + + def __repr__(self): + return f"{self.__class__.__name__}({self.original_weight_tensor}, {self.input_quant_func}, quant_kwargs={self.quant_kwargs}))" + + def __tensor_flatten__(self): + return ["original_weight_tensor"], [self.input_quant_func, self.quant_kwargs] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + original_weight_tensor = tensor_data_dict["original_weight_tensor"] + input_quant_func, quant_kwargs = tensor_attributes + return cls(original_weight_tensor, input_quant_func, quant_kwargs) + + @staticmethod + def _quantized_linear_op( + input_tensor: torch.Tensor, weight_tensor: torch.Tensor, bias: torch.Tensor + ): + if input_tensor.numel() == 0: + return input_tensor + input_quant_func = weight_tensor.input_quant_func + original_weight_tensor = weight_tensor.original_weight_tensor + quant_kwargs = weight_tensor.quant_kwargs + quantized_tensor = input_quant_func(input_tensor, **quant_kwargs) + return torch.nn.functional.linear( + quantized_tensor, original_weight_tensor, bias + ) + + @classmethod + def from_float( + cls, + input_float: torch.Tensor, + input_quant_func: Callable, + quant_kwargs: Optional[Dict[str, Any]] = None, + ): + if quant_kwargs is None: + quant_kwargs = {} + return cls(input_float, input_quant_func, quant_kwargs) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.original_weight_tensor), + self.input_quant_func, + self.quant_kwargs, + ) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + return self.__class__( + self.original_weight_tensor.to(**kwargs), + self.input_quant_func, + self.quant_kwargs, + ) + + +def _same_metadata( + self: LinearActivationQuantizedTensor, src: LinearActivationQuantizedTensor +): + return ( + isinstance(self, LinearActivationQuantizedTensor) + and isinstance(src, LinearActivationQuantizedTensor) + and self.shape == src.shape + and self.input_quant_func == src.input_quant_func + and self.quant_kwargs == src.quant_kwargs + ) + + +implements = LinearActivationQuantizedTensor.implements + + +@implements([torch.nn.functional.linear, aten.linear.default]) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + if isinstance(weight_tensor, LinearActivationQuantizedTensor): + return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias) + + raise NotImplementedError( + "LinearActivationQuantizedTensor: No specialized dispatch found for linear op" + ) + + +@implements([aten.mm.default, aten.addmm.default]) +def _(func, types, args, kwargs): + if not args[0].is_floating_point(): + raise NotImplementedError( + "LinearActivationQuantizedTensor: expecting a floating point input" + ) + + if func == aten.addmm.default: + assert args[1].shape[-1] == args[2].shape[0], ( + f"need mat1 shape: {args[1].shape} final" + f"dim to match mat2 shape: {args[2].shape} first dim " + ) + input_tensor, weight_tensor, bias = ( + args[1], + args[2], + args[0], + ) + input_quant_func = weight_tensor.input_quant_func + original_weight_tensor = weight_tensor.original_weight_tensor + qtensor = input_quant_func(input_tensor, **weight_tensor.quant_kwargs) + return func(bias, qtensor, original_weight_tensor) + else: + # aten.mm.default + assert args[0].shape[-1] == args[1].shape[0], ( + f"need mat1 shape: {args[0].shape} final dim" + f"to match mat2 shape: {args[1].shape} first dim" + ) + input_tensor, weight_tensor = ( + args[0], + args[1], + ) + input_quant_func = weight_tensor.input_quant_func + original_weight_tensor = weight_tensor.original_weight_tensor + qtensor = input_quant_func(input_tensor, **weight_tensor.quant_kwargs) + return func(qtensor, original_weight_tensor) + + +@implements([aten.detach.default, aten.alias.default]) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(func) + ) + + +@implements(aten.clone.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + +@implements(aten._to_copy.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].to(*args[1:], **kwargs)._apply_fn_to_data(torch.clone), + ) + + +@implements(aten.copy_.default) +def _(func, types, args, kwargs): + self = args[0] + src = args[1] + if _same_metadata(self, src): + self_tensors = self.__tensor_flatten__()[0] + for tensor_name in self_tensors: + getattr(self, tensor_name).copy_(getattr(src, tensor_name)) + return + raise ValueError( + f"Not supported args for copy_ due to metadata mistach: {args[0], args[1]}" + ) + + +@implements(aten.t.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.t) + ) + + +@implements(aten.slice.Tensor) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + LinearActivationQuantizedTensor( + func(args[0].original_weight_tensor, *args[1:]), + args[0].input_quant_func, + args[0].quant_kwargs, + ), + ) + + +@implements(aten.select.int) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + LinearActivationQuantizedTensor( + func(args[0].original_weight_tensor, *args[1:]), + args[0].input_quant_func, + args[0].quant_kwargs, + ), + ) + + +@implements(aten.index.Tensor) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + LinearActivationQuantizedTensor( + func(args[0].original_weight_tensor, *args[1:]), + args[0].input_quant_func, + args[0].quant_kwargs, + ), + ) + + +# this is needed for DTensor.from_local() and for flattening tensor +@implements(aten.view.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + LinearActivationQuantizedTensor( + func(args[0].original_weight_tensor, *args[1:]), + args[0].input_quant_func, + args[0].quant_kwargs, + ), + ) + + +to_linear_activation_quantized = LinearActivationQuantizedTensor.from_float + +if TORCH_VERSION_AT_LEAST_2_5: + # Allow a model with LinearActivationQuantizedTensor weights to be loaded with `weights_only=True` + torch.serialization.add_safe_globals([LinearActivationQuantizedTensor]) diff --git a/lib/python3.12/site-packages/torchao/quantization/linear_activation_scale.py b/lib/python3.12/site-packages/torchao/quantization/linear_activation_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..6c433844a632caef90c50dd432dc3bd7dcef5c69 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/linear_activation_scale.py @@ -0,0 +1,168 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import torch +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_5, + TorchAOBaseTensor, +) + +__all__ = [ + "WeightTensorWithLinearActivationScaleMetadata", + "to_weight_tensor_with_linear_activation_scale_metadata", +] + +aten = torch.ops.aten + + +class WeightTensorWithLinearActivationScaleMetadata(TorchAOBaseTensor): + """ + Tensor subclass that wraps a weight tensor and provides metadata for linear activation scaling. + Right now we hardcode how we apply the scale: + scaled_linear_act = input_act / scale + out = F.linear(scaled_linear_act, weight, ...) + + We can generalize this to accept a function as well if needed. + + Args: + original_weight_tensor (torch.Tensor): The weight tensor to be wrapped. + scale (torch.Tensor): The scale tensor to be applied to activation. + """ + + original_weight_tensor: torch.Tensor + scale: torch.Tensor + + def __new__( + cls, + original_weight_tensor: torch.Tensor, + scale: torch.Tensor, + ): + kwargs = {} + dtype = original_weight_tensor.dtype + kwargs["dtype"] = dtype + kwargs["requires_grad"] = False + kwargs["device"] = original_weight_tensor.device + shape = original_weight_tensor.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + original_weight_tensor: torch.Tensor, + scale: torch.Tensor, + ): + self.original_weight_tensor = original_weight_tensor + self.scale = scale + + def __repr__(self): + return f"WeightTensorWithLinearActivationScaleMetadata({self.original_weight_tensor}, scale={self.scale}" + + def __tensor_flatten__(self): + tensor_data = ["original_weight_tensor", "scale"] + return tensor_data, [] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + return cls( + tensor_data_dict["original_weight_tensor"], + tensor_data_dict["scale"], + ) + + @staticmethod + def _quantized_linear_op( + input_tensor: torch.Tensor, weight_tensor: torch.Tensor, bias: torch.Tensor + ): + original_weight_tensor = weight_tensor.original_weight_tensor + scale = weight_tensor.scale + # Note: we can make this function configurable as well + scaled_input_act = input_tensor / scale + return torch.nn.functional.linear( + scaled_input_act, original_weight_tensor, bias + ) + + @classmethod + def from_float( + cls, + input_float: torch.Tensor, + scale: torch.Tensor, + ): + return cls(input_float, scale) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.original_weight_tensor), + fn(self.scale), + ) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs.pop("device") + return self.__class__( + self.original_weight_tensor.to(device), + self.scale.to(device), + ) + + +implements = WeightTensorWithLinearActivationScaleMetadata.implements + + +@implements(torch.nn.functional.linear) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + if isinstance(weight_tensor, WeightTensorWithLinearActivationScaleMetadata): + return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias) + + raise NotImplementedError( + "LinearActivationQuantizedTensor: No specialized dispatch found for linear op" + ) + + +@implements(aten.detach.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + +@implements(aten.clone.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + +@implements(aten._to_copy.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].to(*args[1:], **kwargs)._apply_fn_to_data(torch.clone), + ) + + +@implements(aten.t.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.t) + ) + + +to_weight_tensor_with_linear_activation_scale_metadata = ( + WeightTensorWithLinearActivationScaleMetadata.from_float +) + +if TORCH_VERSION_AT_LEAST_2_5: + # Allow a model with LinearActivationQuantizedTensor weights to be loaded with `weights_only=True` + torch.serialization.add_safe_globals( + [WeightTensorWithLinearActivationScaleMetadata] + ) diff --git a/lib/python3.12/site-packages/torchao/quantization/linear_activation_weight_observed_tensor.py b/lib/python3.12/site-packages/torchao/quantization/linear_activation_weight_observed_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..029b89e54bb6be8900c1ea5883969d079ddf4300 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/linear_activation_weight_observed_tensor.py @@ -0,0 +1,158 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Callable, Dict, Optional + +import torch +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.quantization.observer import AffineQuantizedObserverBase +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_5, + TorchAOBaseTensor, +) + +__all__ = [ + "LinearActivationWeightObservedTensor", +] + +aten = torch.ops.aten +Tensor = torch.Tensor + + +class LinearActivationWeightObservedTensor(TorchAOBaseTensor): + """ + This subclass of Tensor is used in conjuction with a static calibration flow. + The flow is broken up into 3 parts; + 1. Insert the LinearActivationWeightObservedTensor subclass into the model's nn.Linear layers + 2. Run the model with a calibration dataset, the observer will record the min/max of the input and weight + 3. quantize_ the model to static using the statistics recorded by the observer + + This subclass wraps the original weight tensor on the nn.Linear layer. When forward is called, the observer + will first calculat statistics on BOTH the input and weight, and then run the linear op. + """ + + original_weight_tensor: torch.Tensor + input_observer: Optional[AffineQuantizedObserverBase] + weight_observer: Optional[AffineQuantizedObserverBase] + + def __new__( + cls, + original_weight_tensor: torch.Tensor, + input_observer: Optional[AffineQuantizedObserverBase] = None, + weight_observer: Optional[AffineQuantizedObserverBase] = None, + ): + kwargs = {} + dtype = original_weight_tensor.dtype + kwargs["dtype"] = dtype + kwargs["requires_grad"] = False + kwargs["device"] = original_weight_tensor.device + shape = original_weight_tensor.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + original_weight_tensor: torch.Tensor, + input_observer: Optional[AffineQuantizedObserverBase] = None, + weight_observer: Optional[AffineQuantizedObserverBase] = None, + ): + self.original_weight_tensor = original_weight_tensor + self.input_observer = input_observer + self.weight_observer = weight_observer + + def __repr__(self): + return ( + f"LinearActivationWeightObservedTensor(\n" + f"original_weight={self.original_weight_tensor}\n" + f"input_observer={self.input_observer.__class__.__name__ if self.input_observer else None}\n" + f"weight_observer={self.weight_observer.__class__.__name__ if self.weight_observer else None}\n)" + ) + + def __tensor_flatten__(self): + return ["original_weight_tensor"], [self.input_observer, self.weight_observer] + + @classmethod + def __tensor_unflatten__( + cls, + tensor_data_dict: Dict[str, Tensor], + tensor_attributes, + outer_size, + outer_stride, + ): + original_weight_tensor = tensor_data_dict["original_weight_tensor"] + (input_observer, weight_observer) = tensor_attributes + return cls(original_weight_tensor, input_observer, weight_observer) + + @classmethod + def from_float( + cls, + original_weight_tensor: Tensor, + input_observer: Optional[AffineQuantizedObserverBase] = None, + weight_observer: Optional[AffineQuantizedObserverBase] = None, + ): + return cls(original_weight_tensor, input_observer, weight_observer) + + def _apply_fn_to_data(self, fn: Callable): + """Applies a fn to the tensor component of the LinearActivationWeightObservedTensor""" + return self.__class__( + fn(self.original_weight_tensor), + self.input_observer, + self.weight_observer, + ) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + return self._apply_fn_to_data(lambda x: x.to(**kwargs)) + + +implements = LinearActivationWeightObservedTensor.implements + + +@implements(torch.nn.functional.linear) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + if weight_tensor.input_observer is not None: + input_tensor = weight_tensor.input_observer(input_tensor) + if weight_tensor.weight_observer is not None: + weight_tensor = weight_tensor.weight_observer( + weight_tensor.original_weight_tensor + ) + else: + weight_tensor = weight_tensor.original_weight_tensor + + return torch.nn.functional.linear(input_tensor, weight_tensor, bias) + + +@implements(aten.detach.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + +@implements(aten.clone.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + +@implements(aten._to_copy.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].to(*args[1:], **kwargs)._apply_fn_to_data(torch.clone), + ) + + +if TORCH_VERSION_AT_LEAST_2_5: + # Allow a model with LinearActivationQuantizedTensor weights to be loaded with `weights_only=True` + torch.serialization.add_safe_globals([LinearActivationWeightObservedTensor]) diff --git a/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__init__.py b/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7251771cc377d3181158587caa4137ccdba835bd --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__init__.py @@ -0,0 +1,282 @@ +from typing import Tuple + +import torch + +from torchao.quantization.granularity import ( + PerAxis, + PerGroup, +) +from torchao.quantization.marlin_qqq.utils import ( + const, + get_pack_factor, + get_qqq_scale_perms, + get_qqq_scale_reverse_perms, + get_qqq_weight_perm, + get_qqq_weight_reverse_perm, + marlin_permute_weights, + reverse_marlin_permute_weights, +) + +__all__ = [ + "marlin_qqq_workspace", + "pack_to_marlin_qqq", + "unpack_from_marlin_qqq", +] + + +def marlin_qqq_workspace( + out_features: int, + min_thread_n: int = const.MIN_THREAD_N, + max_parallel: int = const.MAX_PARALLEL, +) -> torch.Tensor: + """Creates a workspace for marlin qqq. The workspace is used to coordinate the locks + during the execution of the kernel. + + Args: + out_features (int): The number of output features. + min_thread_n (int, optional): The minimum number of threads per block. Defaults to `MARLIN_QQQ_MIN_THREAD_N`. + max_parallel (int, optional): The maximum number of parallel threads. Defaults to `MARLIN_QQQ_MAX_PARALLEL`. + Returns: + torch.Tensor: The workspace tensor fully initialized with zeros. + """ + assert out_features % min_thread_n == 0, ( + f"out_features = {out_features}, min_thread_n = {min_thread_n}" + ) + max_workspace_size = (out_features // min_thread_n) * max_parallel + return torch.zeros(max_workspace_size, dtype=torch.int, device="cuda") + + +def pack_to_marlin_qqq( + q_w: torch.Tensor, + s_group: torch.Tensor, + s_channel: torch.Tensor, + num_bits: int, + group_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pack the quantized weights and scales to the marlin format. + + Args: + q_w (torch.Tensor): The quantized weight. + s_group (torch.Tensor): The per-group quantization scale. + s_channel (torch.Tensor): The per-channel quantization scale. + num_bits (int): The number of bits used for quantization. + group_size (int): The group size of quantization. + Returns: + Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: The packed quantized weight in marlin format, the packed per-group scale in marlin format, and the packed per-channel scale in marlin format. + """ + in_features, out_features = q_w.shape + + assert num_bits == 4, "Marlin QQQ only supports 4-bit for now." + + # Reformat to marlin + marlin_qqq_q_w = _to_marlin_weights( + q_w, in_features, out_features, num_bits, group_size + ) + marlin_qqq_s_group, marlin_qqq_s_channel = _to_marlin_scales( + s_group, s_channel, in_features, out_features, num_bits, group_size + ) + + return marlin_qqq_q_w, marlin_qqq_s_group, marlin_qqq_s_channel + + +def _to_marlin_weights( + q_w: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + group_size: int, +) -> torch.Tensor: + """Converts a quantized weight tensor to the marlin format. + + Args: + q_w (torch.Tensor): The quantized weight. + size_k (int): The number of input features. + size_n (int): The number of output features. + num_bits (int): The number of bits used for quantization. + group_size (int): The group size of quantization. + Returns: + torch.Tensor: The packed quantized weight in marlin format. + """ + if group_size == -1: + group_size = size_k + granularity = PerAxis(1) if group_size == size_k else PerGroup(group_size) + # Permute + perm = get_qqq_weight_perm(num_bits, granularity) + q_w = marlin_permute_weights(q_w, size_k, size_n, perm) + + # Pack + pack_factor = get_pack_factor(num_bits) + orig_device = q_w.device + + # q_w is torch.uint32 originally, but torch does not support lshift_cuda or lshift_cpu, we have to + # convert it to torch.int64 + q_w = q_w.to(torch.int64) + q_packed = torch.zeros( + (q_w.shape[0], q_w.shape[1] // pack_factor), + dtype=torch.int64, + device=q_w.device, + ) + + if group_size == size_k: + for i in range(pack_factor): + q_packed |= (q_w[:, i::pack_factor] & 0xF) << (num_bits * i) + else: + for i in range(pack_factor): + q_packed |= q_w[:, i::pack_factor] << (num_bits * i) + + q_packed = q_packed.to(torch.int32).to(orig_device) + return q_packed + + +def _to_marlin_scales( + s_group: torch.Tensor, + s_channel: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + group_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Converts the per-group scale and the per-channel scale to the format necessary for marlin. + + Args: + s_group (torch.Tensor): The per-group quantization scale. + s_channel (torch.Tensor): The per-channel quantization scale. + size_k (int): The number of input features. + size_n (int): The number of output features. + num_bits (int): The number of bits used for quantization. + group_size (int): The group size of quantization. + Returns: + Tuple[torch.Tensor, torch.Tensor]: The scale tensors in the marlin format. + """ + if group_size == -1: + group_size = size_k + scale_perm, scale_perm_single = get_qqq_scale_perms(num_bits) + if group_size < size_k: + s_group = s_group.reshape((-1, len(scale_perm)))[:, scale_perm] + s_channel = s_channel.reshape((-1, len(scale_perm_single)))[ + :, scale_perm_single + ] + s_group = s_group.reshape((-1, size_n)).contiguous() + else: + s_channel = s_channel.reshape((-1, len(scale_perm_single)))[ + :, scale_perm_single + ] + s_channel = s_channel.reshape((-1, size_n)).contiguous() + + return s_group, s_channel + + +def unpack_from_marlin_qqq( + q_w: torch.Tensor, + s_group: torch.Tensor, + s_channel: torch.Tensor, + original_shape: torch.Size, + num_bits: int, + group_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Unpacks the quantized weights and scales from the marlin format. + Args: + q_w (torch.Tensor): The packed quantized weights. + s_group (torch.Tensor): The per-group quantization scale. + s_channel (torch.Tensor): The per-channel quantization scale. + original_shape (torch.Size): The original shape of the weight tensor. + num_bits (int): The number of bits used for quantization. + group_size (int): The group size of quantization. + Returns: + Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: The unpacked quantized weights and scales. + """ + in_features, out_features = original_shape + + assert num_bits == 4, "Marlin QQQ only supports 4-bit for now." + + # Unpacks the scales + unpacked_s_group, unpacked_s_channel = _from_marlin_scales( + s_group, s_channel, in_features, out_features, num_bits, group_size + ) + + # Unpacks the weights + unpacked_q_w = _from_marlin_weights( + q_w, in_features, out_features, num_bits, group_size + ) + + return unpacked_q_w, unpacked_s_group, unpacked_s_channel + + +def _from_marlin_weights( + q_w: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + group_size: int, +) -> torch.Tensor: + """Converts a weight tensor in the marlin format to a regular format. + Args: + q_w (torch.Tensor): The packed quantized weights. + size_k (int): The number of input features. + size_n (int): The number of output features. + num_bits (int): The number of bits used for quantization. + group_size (int): The group size of quantization. + Returns: + torch.Tensor: The unpacked quantized weights. + """ + if group_size == -1: + group_size = size_k + granularity = PerAxis(1) if group_size == size_k else PerGroup(group_size) + # Permute + perm = get_qqq_weight_reverse_perm(num_bits, granularity) + + orig_device = q_w.device + + wf = ( + torch.tensor(list(range(0, 32, num_bits)), dtype=torch.int32) + .unsqueeze(0) + .to(orig_device) + ) + # unpack weight + weight = torch.bitwise_right_shift( + torch.unsqueeze(q_w, 2).expand(-1, -1, 32 // num_bits), + wf.unsqueeze(0), + ) + weight = torch.bitwise_and(weight, (2**num_bits) - 1) + weight = weight.reshape(weight.shape[0], weight.shape[1] * weight.shape[2]) + q_w_comp = reverse_marlin_permute_weights(weight, size_k, size_n, perm) + + return q_w_comp + + +def _from_marlin_scales( + s_group: torch.Tensor, + s_channel: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + group_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Converts the quantization scales from the marlin format to their original format. + Args: + s_group (torch.Tensor): The per-group quantization scale in marlin format. + s_channel (torch.Tensor): The per-channel quantization scale in marlin format. + size_k (int): The number of input features. + size_n (int): The number of output features. + num_bits (int): The number of bits used for quantization. + group_size (int): The group size of quantization. + Returns: + Tuple[torch.Tensor, torch.Tensor]: The per-group quantization scale in the original format and + the per-channel quantization scale in the original format. + """ + if group_size == -1: + group_size = size_k + scale_perm, scale_perm_single = get_qqq_scale_reverse_perms(num_bits) + if group_size < size_k: + s_group = s_group.reshape((-1, len(scale_perm)))[:, scale_perm] + s_channel = s_channel.reshape((-1, len(scale_perm_single)))[ + :, scale_perm_single + ] + s_group = s_group.reshape((-1, size_n)).contiguous() + else: + s_channel = s_channel.reshape((-1, len(scale_perm_single)))[ + :, scale_perm_single + ] + s_channel = s_channel.reshape((-1, size_n)).contiguous() + + return s_group, s_channel diff --git a/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b4f43bdcdf0350ddace4e216ab4d0e1045ef79b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..475c5efe6857a0c4a85033a4607ff8c0804e1a6b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/utils.py b/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c1fe817dba7685fcbde1667e010851e68ade02cb --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/marlin_qqq/utils.py @@ -0,0 +1,198 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from dataclasses import dataclass, field +from typing import List, Tuple + +import numpy +import torch + +from torchao.quantization.granularity import ( + Granularity, + PerAxis, +) + + +@dataclass(frozen=True) +class MarlinQQQConstants: + TILE: int = 16 + MIN_THREAD_N: int = 64 + MAX_PARALLEL: int = 16 + + SUPPORTED_NUM_BITS: List[int] = field(default_factory=lambda: [4]) + SUPPORTED_GROUP_SIZES: List[int] = field(default_factory=lambda: [-1, 128]) + + +const = MarlinQQQConstants() + + +def get_pack_factor(num_bits: int) -> int: + """Compute the packing factor for a given number of bits. + + Args: + num_bits (int): Number of bits to pack. + Returns: + int: The packing factor. + """ + assert 32 % num_bits == 0, f"Unsupported num_bits = {num_bits}" + return 32 // num_bits + + +def marlin_permute_weights( + q_w: torch.Tensor, + size_k: int, + size_n: int, + perm: torch.Tensor, + tile: int = const.TILE, +) -> torch.Tensor: + """Permute weights to 16x64 Marlin tiles. + + Args: + q_w (torch.Tensor): Quantized weights. + size_k (int): Number of input features. + size_n (int): Number of output features. + perm (torch.Tensor): The computed permutation tensor to be applied. + tile (int, optional): Tile size. Defaults to `TILE`. + Returns: + torch.Tensor: Weight tensor permuted to Marlin tiles. + """ + assert q_w.shape == (size_k, size_n) + assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}" + assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}" + + # Permute weights to 16x64 marlin tiles + q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile)) + q_w = q_w.permute((0, 2, 1, 3)) + q_w = q_w.reshape((size_k // tile, size_n * tile)) + + q_w = q_w.reshape((-1, perm.numel()))[:, perm].reshape(q_w.shape) + + return q_w + + +def reverse_marlin_permute_weights( + q_w_unpacked: torch.Tensor, + size_k: int, + size_n: int, + reverse_perm: torch.Tensor, + tile: int = const.TILE, +) -> torch.Tensor: + """Reverse permute weights from 16x64 Marlin tiles. + Args: + q_w_unpacked (torch.Tensor): Unpacked quantized weights. + size_k (int): Number of input features. + size_n (int): Number of output features. + reverse_perm (torch.Tensor): The computed reverse permutation tensor to be applied. + tile (int, optional): Tile size. Defaults to `TILE`. + Returns: + torch.Tensor: Weight tensor reverse permuted from Marlin tiles. + """ + + assert (q_w_unpacked.shape[0], size_n) == ( + size_k // tile, + q_w_unpacked.shape[1] // tile, + ) + assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}" + assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}" + + # Reverse permute weights to original shape + q_w_comp = q_w_unpacked.reshape((-1, reverse_perm.numel()))[ + :, reverse_perm + ].reshape(q_w_unpacked.shape) + q_w_comp = q_w_comp.reshape((size_k // tile, size_n // tile, tile, tile)) + q_w_comp = q_w_comp.permute((0, 2, 1, 3)) + q_w_comp = q_w_comp.reshape((size_k, size_n)) + + return q_w_comp + + +# NOTE(HandH1998): QQQ employs different perms for per-group and per-channel weight quantization. # noqa: E501 +def get_qqq_weight_perm(num_bits: int, granularity: Granularity) -> torch.Tensor: + """Precompute permutations for the marlin weight shuffling. + + Args: + num_bits (int): Number of bits to pack. + granularity (Granularity): The weight quantization granularity. + Returns: + torch.Tensor: The weight permutation tensor. + """ + perm_list: List[int] = [] + for i in range(32): + perm1: List[int] = [] + col = i // 4 + for block in [0, 1]: + for row in [ + 4 * (i % 4), + 4 * (i % 4) + 1, + 4 * (i % 4) + 2, + 4 * (i % 4) + 3, + ]: + perm1.append(16 * row + col + 8 * block) + for j in range(4): + perm_list.extend([p + 256 * j for p in perm1]) + + perm = numpy.array(perm_list) + + if num_bits == 4: + if isinstance(granularity, PerAxis): + interleave = numpy.array([4, 0, 5, 1, 6, 2, 7, 3]) + else: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + else: + raise Exception("num_bits must be 4, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + return perm + + +def get_qqq_scale_perms(num_bits: int) -> Tuple[List[int], List[int]]: + """Precompute permutations for the marlin scale shuffling. + Args: + num_bits (int): Number of bits to pack. + Returns: + Tuple[List[int], List[int]]: Scale permutation list and + scale permutation list for a single group. + """ + if num_bits != 4: + raise Exception("num_bits must be 4, got {}".format(num_bits)) + scale_perm: List[int] = [] + for i in range(8): + scale_perm.extend([i + 8 * j for j in range(8)]) + scale_perm_single: List[int] = [] + for i in range(4): + scale_perm_single.extend([2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]]) + return scale_perm, scale_perm_single + + +def get_qqq_weight_reverse_perm( + num_bits: int, granularity: Granularity +) -> torch.Tensor: + """Reverse permutation for Marlin weight shuffling from `get_qqq_weight_perm`. + Args: + num_bits (int): Number of bits to pack. + granularity (Granularity): The weight quantization granularity. + Returns: + torch.Tensor: The reversed weight permutation tensor. + """ + perm = get_qqq_weight_perm(num_bits, granularity) + perm = perm.argsort() + + return perm + + +def get_qqq_scale_reverse_perms(num_bits: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Reverse permutation for Marlin scale shuffling from `get_qqq_scale_perms`. + Args: + num_bits (int): Number of bits to pack. + Returns: + Tuple[List[int], List[int]]: The reversed scale permutation list and + the reversed scale permutation list for a single group. + """ + scale_perm, scale_perm_single = get_qqq_scale_perms(num_bits) + scale_perm = torch.tensor(scale_perm).argsort() + scale_perm_single = torch.tensor(scale_perm_single).argsort() + + return scale_perm, scale_perm_single diff --git a/lib/python3.12/site-packages/torchao/quantization/observer.py b/lib/python3.12/site-packages/torchao/quantization/observer.py new file mode 100644 index 0000000000000000000000000000000000000000..e103f0a59e9df94d591c32bab5e98678b913e802 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/observer.py @@ -0,0 +1,198 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import logging +from abc import ABCMeta, abstractmethod +from functools import partial +from typing import Any, Optional, Tuple + +import torch + +from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 + +from .granularity import ( + Granularity, + PerAxis, + PerRow, + PerTensor, +) +from .quant_primitives import ( + MappingType, + ZeroPointDomain, + _get_reduction_params, + choose_qparams_affine_with_min_max, +) + +logger = logging.getLogger(__name__) + + +# borrowed from torch.ao.quantization.observer +class _PartialWrapper: + def __init__(self, p): + self.p = p + + def __call__(self, *args, **keywords): + return self.p(*args, **keywords) + + def __repr__(self): + return self.p.__repr__() + + def with_args(self, *args, **kwargs): + return _with_args(self, *args, **kwargs) + + +def _with_args(cls_or_self, *args, **kwargs): + r"""Wrapper that allows creation of class factories. + + This can be useful when there is a need to create classes with the same + constructor arguments, but different instances. + + Example:: + + >>> # xdoctest: +SKIP("Undefined vars") + >>> Foo.with_args = classmethod(_with_args) + >>> foo_builder = Foo.with_args(a=3, b=4).with_args(answer=42) + >>> foo_instance1 = foo_builder() + >>> foo_instance2 = foo_builder() + >>> id(foo_instance1) == id(foo_instance2) + False + """ + r = _PartialWrapper(partial(cls_or_self, *args, **kwargs)) + return r + + +def get_block_size( + input_shape: Tuple[int, ...], granularity: Granularity +) -> Tuple[int, ...]: + """Get the block size based on the input shape and granularity type. + + Args: + input_shape: The input tensor shape possibly more than 2 dimensions + granularity: The granularity type of the quantization + """ + if isinstance(granularity, PerTensor): + return input_shape + elif isinstance(granularity, PerAxis): + block_size = list(input_shape) + block_size[granularity.axis] = 1 + return tuple(block_size) + elif isinstance(granularity, PerRow): + return (1,) * (len(input_shape) - 1) + (input_shape[-1],) + raise ValueError(f"Unsupported Granularity: {granularity}") + + +ABC: Any = ABCMeta("ABC", (object,), {}) # compatible with Python 2 *and* 3: + + +class AffineQuantizedObserverBase(ABC, torch.nn.Module): + """Observer module for affine quantization (https://github.com/pytorch/ao/tree/main/torchao/quantization#affine-quantization) + + Args: + `granularity` and `block_size`: The granularity of the quantization, + must specify at least one, if both are specified `block_size` takes precedence + Current supported granularity type are `PerTensor` and `PerAxis` + other args: please see `:class:torchao.dtypes.AffineQuantizedTensor` + """ + + with_args = classmethod(_with_args) + + def __init__( + self, + mapping_type: MappingType, + target_dtype: torch.dtype, + granularity: Granularity, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + ): + super().__init__() + assert granularity is not None, "granularity is None" + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + self.mapping_type = mapping_type + self.target_dtype = target_dtype + self.granularity = granularity + self.quant_min = quant_min + self.quant_max = quant_max + self.eps = eps + self.scale_dtype = scale_dtype + self.zero_point_dtype = zero_point_dtype + self.preserve_zero = preserve_zero + self.zero_point_domain = zero_point_domain + + @abstractmethod + def forward(self, input: torch.Tensor) -> torch.Tensor: + """forward function should take the input tensor + and updates internal stats and return the original input Tensor + """ + pass + + @abstractmethod + def calculate_qparams(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Calculate quantization parameter based on the stats attached to the observer module + and returns a tuple of scale and zero_point Tensor + """ + pass + + +class AffineQuantizedMinMaxObserver(AffineQuantizedObserverBase): + def forward(self, input: torch.Tensor): + if input.numel() == 0: + return input + + input_detached = input.detach() + assert self.granularity is not None, "granularity is None" + block_size = get_block_size(input_detached.shape, self.granularity) + + shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input_detached.size() + ) + input_detached = input_detached.view(shape_for_reduction) + min_val = torch.amin(input_detached, dim=reduction_dims, keepdim=False) + max_val = torch.amax(input_detached, dim=reduction_dims, keepdim=False) + if not hasattr(self, "min_val") or not hasattr(self, "max_val"): + self.min_val = min_val + self.max_val = max_val + else: + assert self.min_val.shape == min_val.shape, ( + f"Can't update existing min_val - shape mismatch, self.min_val:{self.min_val.shape} != min_val:{min_val.shape}" + ) + assert self.max_val.shape == max_val.shape, ( + f"Can't update existing max_val - shape mismatch, self.max_val {self.max_val.shape} != max_val:{max_val.shape}" + ) + min_val = torch.min(self.min_val, min_val) + max_val = torch.max(self.max_val, max_val) + self.min_val.copy_(min_val) + self.max_val.copy_(max_val) + # returning original input + return input + + def calculate_qparams(self) -> Tuple[torch.Tensor, torch.Tensor]: + assert hasattr(self, "min_val") and hasattr(self, "max_val"), ( + "Expecting the observer has min_val and max_val, please run the observer before calling calculate_qparams" + ) + return choose_qparams_affine_with_min_max( + self.min_val, + self.max_val, + self.mapping_type, + [], # BlockSize is not needed because the min/max are already reduced + self.target_dtype, + self.quant_min, + self.quant_max, + self.eps, + self.scale_dtype, + self.zero_point_dtype, + self.preserve_zero, + self.zero_point_domain, + ) + + +if TORCH_VERSION_AT_LEAST_2_5: + # Allow a model with LinearActivationQuantizedTensor weights to be loaded with `weights_only=True` + torch.serialization.add_safe_globals([PerRow, PerTensor]) diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/__init__.py b/lib/python3.12/site-packages/torchao/quantization/prototype/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/prototype/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7737ebde95db4f9d921072023ee3798f86c88e69 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/prototype/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__init__.py b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aa0084c6ea4c3ee4575999e6eecf177a172be4be --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__init__.py @@ -0,0 +1,25 @@ +from torchao.quantization.qat import ( + ComposableQATQuantizer, + Int4WeightOnlyEmbeddingQATQuantizer, + Int4WeightOnlyQATQuantizer, + Int8DynActInt4WeightQATQuantizer, +) +from torchao.quantization.qat.linear import ( + Int8DynActInt4WeightQATLinear, + disable_4w_fake_quant, + disable_8da4w_fake_quant, + enable_4w_fake_quant, + enable_8da4w_fake_quant, +) + +__all__ = [ + "disable_4w_fake_quant", + "disable_8da4w_fake_quant", + "enable_4w_fake_quant", + "enable_8da4w_fake_quant", + "ComposableQATQuantizer", + "Int4WeightOnlyQATQuantizer", + "Int4WeightOnlyEmbeddingQATQuantizer", + "Int8DynActInt4WeightQATQuantizer", + "Int8DynActInt4WeightQATLinear", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..812954601ccab76c4bac48e1ead1ac6a21b7d811 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/_module_swap_api.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/_module_swap_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8676c633f36cd9c5d167137888d9baaee5f3f1b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/_module_swap_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/affine_fake_quantized_tensor.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/affine_fake_quantized_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e195e9c6c1a9de61b88763fe6c7e0377521bab5 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/affine_fake_quantized_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7c8bb857d2580af047cdddc81197894803b66cb Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/embedding.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/embedding.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ab95f4ca617f8198db2ddf476552fb3a9fa5e92 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/embedding.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/fake_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/fake_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..53d926f3eec9319e5df8c84584c2277feb62b825 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/fake_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/linear.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49f3df0322ddeb4e3cdd8b36596a881ac9c88a98 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/__pycache__/linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/_module_swap_api.py b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/_module_swap_api.py new file mode 100644 index 0000000000000000000000000000000000000000..a0fbd877b2af1deeb60e0b72cefbb1cf52cb0641 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/_module_swap_api.py @@ -0,0 +1,30 @@ +# For backward compatibility only +# These will be removed in the future + +from torchao.quantization.qat.linear import ( + Int4WeightOnlyQATQuantizer as Int4WeightOnlyQATQuantizerModuleSwap, +) +from torchao.quantization.qat.linear import ( + Int8DynActInt4WeightQATQuantizer as Int8DynActInt4WeightQATQuantizerModuleSwap, +) +from torchao.quantization.qat.linear import ( + disable_4w_fake_quant as disable_4w_fake_quant_module_swap, +) +from torchao.quantization.qat.linear import ( + disable_8da4w_fake_quant as disable_8da4w_fake_quant_module_swap, +) +from torchao.quantization.qat.linear import ( + enable_4w_fake_quant as enable_4w_fake_quant_module_swap, +) +from torchao.quantization.qat.linear import ( + enable_8da4w_fake_quant as enable_8da4w_fake_quant_module_swap, +) + +__all__ = [ + "Int8DynActInt4WeightQATQuantizerModuleSwap", + "Int4WeightOnlyQATQuantizerModuleSwap", + "enable_8da4w_fake_quant_module_swap", + "disable_8da4w_fake_quant_module_swap", + "enable_4w_fake_quant_module_swap", + "disable_4w_fake_quant_module_swap", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/affine_fake_quantized_tensor.py b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/affine_fake_quantized_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..20d51912f0f424c02ad318c10a2101680f8ca1e3 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/affine_fake_quantized_tensor.py @@ -0,0 +1,9 @@ +from torchao.quantization.qat.affine_fake_quantized_tensor import ( + AffineFakeQuantizedTensor, + to_affine_fake_quantized, +) + +__all__ = [ + "AffineFakeQuantizedTensor", + "to_affine_fake_quantized", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/api.py b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/api.py new file mode 100644 index 0000000000000000000000000000000000000000..c2f1d6f8d7deaa51fdfb31133da3adaa257d330f --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/api.py @@ -0,0 +1,9 @@ +from torchao.quantization.qat.api import ( + ComposableQATQuantizer, + FakeQuantizeConfig, +) + +__all__ = [ + "ComposableQATQuantizer", + "FakeQuantizeConfig", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/embedding.py b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..98b4d05720eb16359ed30c4dace2cbfe52efa4c7 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/embedding.py @@ -0,0 +1,13 @@ +from torchao.quantization.qat.embedding import ( + FakeQuantizedEmbedding, + Int4WeightOnlyEmbedding, + Int4WeightOnlyEmbeddingQATQuantizer, + Int4WeightOnlyQATEmbedding, +) + +__all__ = [ + "FakeQuantizedEmbedding", + "Int4WeightOnlyEmbeddingQATQuantizer", + "Int4WeightOnlyEmbedding", + "Int4WeightOnlyQATEmbedding", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/fake_quantizer.py b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/fake_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..3bbe1fb704469e98dbc1bf964d68e2e621bc581d --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/fake_quantizer.py @@ -0,0 +1,7 @@ +from torchao.quantization.qat.fake_quantizer import ( + FakeQuantizer, +) + +__all__ = [ + "FakeQuantizer", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/prototype/qat/linear.py b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..366ab3565cfcb40fed671a81a64a03d0d1079f53 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/prototype/qat/linear.py @@ -0,0 +1,23 @@ +from torchao.quantization.qat.linear import ( + FakeQuantizedLinear, + Int4WeightOnlyQATLinear, + Int4WeightOnlyQATQuantizer, + Int8DynActInt4WeightQATLinear, + Int8DynActInt4WeightQATQuantizer, + disable_4w_fake_quant, + disable_8da4w_fake_quant, + enable_4w_fake_quant, + enable_8da4w_fake_quant, +) + +__all__ = [ + "disable_4w_fake_quant", + "disable_8da4w_fake_quant", + "enable_4w_fake_quant", + "enable_8da4w_fake_quant", + "FakeQuantizedLinear", + "Int4WeightOnlyQATLinear", + "Int4WeightOnlyQATQuantizer", + "Int8DynActInt4WeightQATLinear", + "Int8DynActInt4WeightQATQuantizer", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__init__.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3e4352dabdae407aed58bd589a7d46b9a47c03d2 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/__init__.py @@ -0,0 +1,188 @@ +# mypy: allow-untyped-defs + +from typing import Callable, Optional, Union + +import torch +from torch import Tensor + +from torchao.quantization.pt2e._numeric_debugger import ( # noqa: F401 + CUSTOM_KEY, + NUMERIC_DEBUG_HANDLE_KEY, + compare_results, + extract_results_from_loggers, + generate_numeric_debug_handle, + prepare_for_propagation_comparison, +) +from torchao.quantization.pt2e.export_utils import ( + WrapperModule, +) +from torchao.quantization.pt2e.export_utils import ( + _allow_exported_model_train_eval as allow_exported_model_train_eval, +) +from torchao.quantization.pt2e.export_utils import ( + _move_exported_model_to_eval as move_exported_model_to_eval, +) +from torchao.quantization.pt2e.export_utils import ( + _move_exported_model_to_train as move_exported_model_to_train, +) +from torchao.quantization.pt2e.graph_utils import ( + bfs_trace_with_node_process, + find_sequential_partitions, + get_equivalent_types, + update_equivalent_types_dict, +) + +from .fake_quantize import ( + FakeQuantize, + FakeQuantizeBase, + FixedQParamsFakeQuantize, + FusedMovingAvgObsFakeQuantize, + default_dynamic_fake_quant, + default_fake_quant, + enable_fake_quant, + enable_observer, +) +from .observer import ( + AffineQuantizedObserverBase, + FixedQParamsObserver, + Granularity, + HistogramObserver, + MappingType, + MinMaxObserver, + MovingAverageMinMaxObserver, + MovingAveragePerChannelMinMaxObserver, + NoopObserver, + ObserverBase, + PartialWrapper, + PerAxis, + PerBlock, + PerChannelMinMaxObserver, + PerGroup, + PerRow, + PerTensor, + PerToken, + PlaceholderObserver, + RecordingObserver, + ReuseInputObserver, + TorchAODType, + UniformQuantizationObserverBase, + ZeroPointDomain, + get_block_size, +) + +for _f in [ + compare_results, + extract_results_from_loggers, + generate_numeric_debug_handle, + prepare_for_propagation_comparison, +]: + _f.__module__ = "torchao.quantization.pt2e" + + +# ensure __module__ is set correctly for public APIs +ObserverOrFakeQuantize = Union[ObserverBase, FakeQuantizeBase] +ObserverOrFakeQuantize.__module__ = "torchao.quantization.pt2e" + +ObserverOrFakeQuantizeConstructor = Union[ + PartialWrapper, type[ObserverBase], type[FakeQuantizeBase] +] +ObserverOrFakeQuantizeConstructor.__module__ = "torchao.quantization.pt2e" + + +__all__ = [ + # old fake quantizers + "FakeQuantize", + "FakeQuantizeBase", + "FixedQParamsFakeQuantize", + "FixedQParamsObserver", + "FusedMovingAvgObsFakeQuantize", + # old observers + "HistogramObserver", + "MinMaxObserver", + "MovingAverageMinMaxObserver", + "MovingAveragePerChannelMinMaxObserver", + "NoopObserver", + "ObserverBase", + "ObserverOrFakeQuantize", + "PerChannelMinMaxObserver", + "PlaceholderObserver", + "RecordingObserver", + "ReuseInputObserver", + "UniformQuantizationObserverBase", + "ObserverOrFakeQuantizeConstructor", + "DerivedObserverOrFakeQuantize", + # utils + "enable_fake_quant", + "enable_observer", + # export_utils + "move_exported_model_to_eval", + "move_exported_model_to_train", + "allow_exported_model_train_eval", + "WrapperModule", + # graph_utils + "find_sequential_partitions", + "get_equivalent_types", + "update_equivalent_types_dict", + "bfs_trace_with_node_process", + # pt2e numeric debugger + "generate_numeric_debug_handle", + "CUSTOM_KEY", + "NUMERIC_DEBUG_HANDLE_KEY", + "prepare_for_propagation_comparison", + "extract_results_from_loggers", + "compare_results", + # should be merged with torchao/quantization/observer.py in the future + "AffineQuantizedObserverBase", + "Granularity", + "MappingType", + "PerAxis", + "PerBlock", + "PerGroup", + "PerRow", + "PerTensor", + "PerToken", + "TorchAODType", + "ZeroPointDomain", + "get_block_size", + "default_fake_quant", + "default_dynamic_fake_quant", +] + + +class DerivedObserverOrFakeQuantize(ObserverBase): + r"""This observer is used to describe an observer whose quantization parameters + are derived from other observers + """ + + def __init__( + self, + dtype: torch.dtype, + obs_or_fqs: list[ObserverOrFakeQuantize], + derive_qparams_fn: Callable[ + [list[ObserverOrFakeQuantize]], tuple[Tensor, Tensor] + ], + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + qscheme: Optional[torch.qscheme] = None, + ch_axis: Optional[int] = None, + ): + super().__init__(dtype) + self.obs_or_fqs = obs_or_fqs + self.derive_qparams_fn = derive_qparams_fn + self.quant_min = quant_min + self.quant_max = quant_max + self.qscheme = qscheme + self.ch_axis = ch_axis + + from .utils import is_per_channel + + if is_per_channel(self.qscheme): + assert self.ch_axis is not None, ( + "Must provide a valid ch_axis if qscheme is per channel" + ) + + def forward(self, x: Tensor) -> Tensor: + return x + + def calculate_qparams(self): # type:ignore[override] + return self.derive_qparams_fn(self.obs_or_fqs) diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7793974be4ddd7a397a702292d1c0cfe7e040f28 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/_affine_quantization.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/_affine_quantization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..850851b2aff1b085f5d4b897a5be2d36170f4de9 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/_affine_quantization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/_numeric_debugger.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/_numeric_debugger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5ef70f32318d199c81b8d6d60b55b6eeb44f735 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/_numeric_debugger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/constant_fold.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/constant_fold.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc338ab981f7e1b9ce2b57473b0595416940b06b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/constant_fold.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/convert.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/convert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f5201c92590e5b713bda3980cd80e43fae52dc0 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/convert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/export_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/export_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e430a2a9fe5857aa8ea011b51e7b6ba47827195 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/export_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/fake_quantize.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/fake_quantize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22bd49c3f327f040452649868ab138601d0bac64 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/fake_quantize.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/graph_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/graph_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b74bc75366b256aa51f96066f1e4d20b682e0773 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/graph_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/lowering.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/lowering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0c311ddfdfcb15897418e3ecee91ccc8e01c9b4 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/lowering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/observer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/observer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e57864533d956741ffbe4a1d0f1b2c37f14d8b0b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/observer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/prepare.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/prepare.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6639f29a2f608cb58c3e78435ae08ec0258d6eb8 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/prepare.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/qat_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/qat_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8aa3db9c5b4928eb2d8888d097061187123981e6 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/qat_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/quantize_pt2e.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/quantize_pt2e.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33c44f961cabf3cb5db1d46dae52c6c035116b63 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/quantize_pt2e.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/reference_representation_rewrite.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/reference_representation_rewrite.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c0190309a033a637b62f3cc25828231b9a453d43 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/reference_representation_rewrite.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f0a72f75faf299e08940f5be6b0321360545862 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/_affine_quantization.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/_affine_quantization.py new file mode 100644 index 0000000000000000000000000000000000000000..e02bee03ce9735164aad2c20f913c31ab45a1161 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/_affine_quantization.py @@ -0,0 +1,813 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# copied from https://github.com/pytorch/ao/blob/main/torchao/quantization/observer.py +# and https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_primitives.py +# PLESE DON'T MODIFY THIS FILE SO THAT WE DON'T GET OUT OF SYNC +import logging +from abc import ABCMeta +from typing import Any, List, Optional, Union + +import torch + +from torchao.quantization.pt2e.observer import ( + AffineQuantizedObserverBase, + Granularity, + MappingType, + TorchAODType, + ZeroPointDomain, + get_block_size, +) + +ABC: Any = ABCMeta("ABC", (object,), {}) # compatible with Python 2 *and* 3: + +logger = logging.getLogger(__name__) + +FP8_TYPES = { + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float8_e4m3fnuz, + torch.float8_e5m2fnuz, +} +_SUB_BYTE_UINT_BOUNDS = { + torch.uint1: (0, 2**1 - 1), + torch.uint2: (0, 2**2 - 1), + torch.uint3: (0, 2**3 - 1), + torch.uint4: (0, 2**4 - 1), + torch.uint5: (0, 2**5 - 1), + torch.uint6: (0, 2**6 - 1), + torch.uint7: (0, 2**7 - 1), +} + +""" +Map from dtype to the bound value of integers +TODO: maybe can replace this with call to torch.iinfo +""" +_DTYPE_TO_QVALUE_BOUNDS: dict[Union[torch.dtype, TorchAODType], tuple[int, int]] = { + torch.uint8: (0, 255), + torch.int8: (-128, 127), + torch.int16: (-(2**15), 2**15 - 1), + torch.int32: (-(2**31), 2**31 - 1), +} +_DTYPE_TO_QVALUE_BOUNDS.update(_SUB_BYTE_UINT_BOUNDS) + + +def _is_float8_type(dtype: torch.dtype) -> bool: + fp8_types = { + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + torch.float8_e5m2, + torch.float8_e5m2fnuz, + } + return dtype in fp8_types + + +# TODO: decide on if we want to allow custom quant_min/quant_max here +def _get_and_check_qmin_qmax(dtype, quant_min, quant_max): + """Get quant_min and quant_max args based on dtype and also + verify that they are within the range of possible quant_min/quant_max + for dtype + """ + if dtype in FP8_TYPES: + quant_min_lower_bound, quant_max_upper_bound = ( + torch.finfo(dtype).min, + torch.finfo(dtype).max, + ) + elif dtype not in _DTYPE_TO_QVALUE_BOUNDS: + raise ValueError(f"Unsupported dtype: {dtype}") + else: + quant_min_lower_bound, quant_max_upper_bound = _DTYPE_TO_QVALUE_BOUNDS[dtype] + if quant_min is None: + quant_min = quant_min_lower_bound + if quant_max is None: + quant_max = quant_max_upper_bound + + assert quant_min >= quant_min_lower_bound, ( + "quant_min out of bound for dtype, " + f"quant_min_lower_bound: {quant_min_lower_bound} quant_min: {quant_min}" + ) + + assert quant_max <= quant_max_upper_bound, ( + "quant_max out of bound for dtype, " + f"quant_max_upper_bound: {quant_max_upper_bound} quant_max: {quant_max}" + ) + return quant_min, quant_max + + +def _get_reduction_params(block_size, input_size): + """Given block_size and input size find the parameters for reduction: + + Output: + shape_for_reduction: the shape we use to `view` input to prepare it for reduction + reduction_dims: the dims we'll do reduction over + + Example:: + Input: + block_size: (3, 3, 2, 10) + input_size: (3, 3, 10, 10) + + Output: + shape_for_reduction: (3, 3, 5, 2, 10) + reduction_dim: [0, 1, 3, 4] + """ + assert len(block_size) == len(input_size) + shape_for_reduction = [] + reduction_dims = [] + cur_dim = 0 + for i in range(len(block_size)): + if block_size[i] != input_size[i] and block_size[i] > 1: + assert input_size[i] % block_size[i] == 0, ( + f"Expecting input size at {i} dimension: " + f"{input_size[i]} to be divisible by block_size at {i} dimension: {block_size[i]}" + ) + shape_for_reduction.append(input_size[i] // block_size[i]) + shape_for_reduction.append(block_size[i]) + # reduce over the block_size[i] dim + reduction_dims.append(cur_dim + 1) + cur_dim += 2 + else: + # block_size[i] == input_size[i] or block_size[i] == 1 + shape_for_reduction.append(input_size[i]) + # we only need to reduce over the dimension if block_size is greater than 1 + # otherwise it's already the same as reduced dimension + if block_size[i] != 1: + reduction_dims.append(cur_dim) + cur_dim += 1 + return shape_for_reduction, reduction_dims + + +def choose_qparams_affine_with_min_max( + min_val: torch.Tensor, + max_val: torch.Tensor, + mapping_type: MappingType, + block_size: tuple[int, ...], + target_dtype: torch.dtype, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.INT, +) -> tuple[torch.Tensor, torch.Tensor]: + """A variant of :func:`~torchao.quantization.quant_primitives.choose_qparams_affine` + operator that pass in min_val and max_val directly instead of deriving these from a single input. + This is used for observers in static quantization where min_val and max_val may be obtained through + tracking all the data in calibration data set. + + Args: + Mostly same as :func:`~torchao.quantization.quant_primitives.choose_qparams_affine`. with one + difference: instead of passing in `input` Tensor and use that to calculate min_val/max_val + and then scale/zero_point, we pass in min_val/max_val directly + """ + return _choose_qparams_affine( + None, + mapping_type.name, + block_size, + target_dtype, + quant_min, + quant_max, + eps, + scale_dtype, + zero_point_dtype, + preserve_zero, + zero_point_domain.name if zero_point_domain is not None else None, + min_val, + max_val, + ) + + +def _choose_qparams_affine( + input: Optional[torch.Tensor], + mapping_type: str, + block_size: List[int], + target_dtype: torch.dtype, + quant_min: Optional[Union[int, float, bool]] = None, + quant_max: Optional[Union[int, float, bool]] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: Optional[str] = "INT", + min_val: Optional[torch.Tensor] = None, + max_val: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """op definition that has compatible signatures with custom op library + + The op does the following: + 1. figure out the dimension for reduction based on block_size + 2. find min_val/max_val based on the dimension for reduction + 3. calculate quantization parameters based on min_val/max_val based on args like `preserve_zero` + and `zero_point_domain` + """ + quant_min, quant_max = _get_and_check_qmin_qmax(target_dtype, quant_min, quant_max) + assert mapping_type in [ + MappingType.SYMMETRIC.name, + MappingType.SYMMETRIC_NO_CLIPPING_ERR.name, + MappingType.ASYMMETRIC.name, + ], f"Unsupported mapping type: {mapping_type}" + if target_dtype in FP8_TYPES: + assert mapping_type == MappingType.SYMMETRIC.name, ( + f"Only symmetric quantization is supported for FP8 types, got {mapping_type}" + ) + + if input is not None: + if scale_dtype is None: + scale_dtype = input.dtype + if zero_point_dtype is None: + zero_point_dtype = input.dtype + if eps is None: + eps = torch.finfo(input.dtype).eps + + assert len(block_size) == input.dim(), ( + f"Got input dim:{input.dim()}, block_size: {block_size}" + ) + shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + input = input.view(shape_for_reduction) + + min_val = torch.amin(input, dim=reduction_dims, keepdim=False) + max_val = torch.amax(input, dim=reduction_dims, keepdim=False) + else: + assert min_val is not None and max_val is not None, ( + "Need to provide `min_val` and `max_val` when `input` is None, got: {min_val, max_val}" + ) + assert min_val.dtype == max_val.dtype, ( + "Expecting `min_val` and `max_val` to have the same dtype, got: {min_val.dtype, max_val.dtype}" + ) + + if scale_dtype is None: + scale_dtype = min_val.dtype + if zero_point_dtype is None: + zero_point_dtype = min_val.dtype + if eps is None: + eps = torch.finfo(min_val.dtype).eps + + if preserve_zero: + min_val_neg = torch.min(min_val, torch.zeros_like(min_val)) + max_val_pos = torch.max(max_val, torch.zeros_like(max_val)) + else: + min_val_neg = min_val + max_val_pos = max_val + + if ( + mapping_type == MappingType.SYMMETRIC.name + or mapping_type == MappingType.SYMMETRIC_NO_CLIPPING_ERR.name + ): + # scales + if mapping_type == MappingType.SYMMETRIC.name: + max_val_pos = torch.max(-min_val_neg, max_val_pos) + scale = max_val_pos / (float(quant_max - quant_min) / 2) + else: + assert mapping_type == MappingType.SYMMETRIC_NO_CLIPPING_ERR.name + # calculate smin and smax individually and choose the larger one. For example, if quant_min = -8 and + # quant_max = 7. + # - If smin is bigger: There would be coverage on negative values down to -8, and less rounding + # error than the existing SYMMETRIC case. + # - If smax is bigger: it covers the positive values up to 7. The round + # error may be bigger than the existing SYMMETRIC case. Either way, there's no out-of-range fp values after + # quantization. + smin = min_val_neg / float(quant_min) + smax = max_val_pos / float(quant_max) + mask = smin > smax + scale = torch.where(mask, smin, smax) + # zeros + if not preserve_zero: + raise ValueError( + "preserve_zero == False is not supported for symmetric quantization" + ) + if ( + zero_point_domain is not None + and zero_point_domain != ZeroPointDomain.INT.name + ): + raise ValueError( + "zero_point_domain != ZeroPointDomain.INT is not supported for symmetric quantization" + ) + scale = torch.clamp(scale, min=eps) + zero_point = torch.full_like(scale, int((quant_max + quant_min + 1) / 2)) + else: + assert mapping_type == MappingType.ASYMMETRIC.name + scale = (max_val_pos - min_val_neg) / float(quant_max - quant_min) + scale = torch.clamp(scale, min=eps) + if zero_point_domain == ZeroPointDomain.NONE.name: + zero_point = None + else: + if preserve_zero: + zero_point = quant_min - torch.round(min_val_neg / scale) + zero_point = torch.clamp(zero_point, quant_min, quant_max) + else: + assert zero_point_domain == ZeroPointDomain.FLOAT.name, ( + "if not preserve_zero, zero_point must be in FLOAT domain" + ) + mid_point = (quant_max + quant_min + 1) / 2 + zero_point = min_val_neg + scale * mid_point + + if zero_point is not None: + zero_point = zero_point.to(dtype=zero_point_dtype) + return scale.to(dtype=scale_dtype), zero_point + + +@torch.no_grad() +def quantize_affine( + input: torch.Tensor, + block_size: tuple[int, ...], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + output_dtype: torch.dtype, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.INT, +) -> torch.Tensor: + """ + Args: + input (torch.Tensor): original float32, float16 or bfloat16 Tensor + block_size: (Tuple[int, ...]): granularity of quantization, + this means the size of the tensor elements that's sharing the same qparam + e.g. when size is the same as the input tensor dimension, we are using per tensor quantization + scale (float): quantization parameter for affine quantization + zero_point (int): quantization parameter for affine quantization + output_dtype (torch.dtype): requested dtype (e.g. torch.uint8) for output Tensor + quant_min (Optional[int]): minimum quantized value for output Tensor, if not specified, it will be derived from dtype + quant_max (Optional[int]): maximum quantized value for output Tensor, if not specified, it will be derived from dtype + zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be either integer or float + if zero_point is in integer domain, zero point is added to the quantized integer value during + quantization + if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) + value during quantization + default is ZeroPointDomain.INT + + Note: + How can block_size represent different granularities? + let's say we have a Tensor of size: (3, 3, 10, 10), here is the table showing how block_size represents different + granularities: + + granularity type | block_size + per_tensor | (3, 3, 10, 10) + per_axis (axis=0) | (1, 3, 10, 10) + per_axis (axis=1) | (3, 1, 10, 10) + per_group (groupsize=2) | (3, 3, 10, 2) + per_group (groupsize=2) for axis = 3 | (3, 3, 2, 10) + + + Output: + quantized tensor with requested dtype + """ + return _quantize_affine( + input, + block_size, + scale, + zero_point, + output_dtype, + quant_min, + quant_max, + zero_point_domain.name if zero_point_domain is not None else None, + ) + + +def _quantize_affine( + input: torch.Tensor, + block_size: list[int], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + output_dtype: torch.dtype, + quant_min: Optional[Union[int, float, bool]] = None, + quant_max: Optional[Union[int, float, bool]] = None, + zero_point_domain: Optional[str] = ZeroPointDomain.INT.name, +) -> torch.Tensor: + """op definition that has compatible signatures with custom op library + + Note: + zero_point_domain is optional specifies how we quantize the floating point to quantized data: + INT: quantized_val = (float_val / scale) (integer) + zero_point (integer) + FLOAT: quantized_val = (float_val - (zero_point (float) - scale * mid_point)) / scale + None: quantized_val = (float_val / scale) | this is primarily used for floatx quantization + Where we do not want to round values to nearest integer and instead scale and cast. + """ + quant_min, quant_max = _get_and_check_qmin_qmax(output_dtype, quant_min, quant_max) + # workaround for uintx dtypes, since we don't have native Uintx dtype connected with + # torch.uintx dtypes yet + if output_dtype in _SUB_BYTE_UINT_BOUNDS: + output_dtype = torch.uint8 + return _quantize_affine_no_dtype_cast( + input, + block_size, + scale, + zero_point, + quant_min, + quant_max, + zero_point_domain, + ).to(output_dtype) + + +def _quantize_affine_no_dtype_cast( + input: torch.Tensor, + block_size: list[int], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_min: Union[int, float], + quant_max: Union[int, float], + zero_point_domain: Optional[str] = ZeroPointDomain.INT.name, +) -> torch.Tensor: + """ + The op does the following: + 1. figure out the dimension for reduction based on block_size, also reshape the input to align with + the shape after reduction + 2. quantize the input based on the quantization parameters scale and zero_point and args like zero_point_domain + 3. reshape the quantized result to origianl shape + """ + # TODO: validations + # TODO: validate scale/zero_point dimensions are compatible with block_size + assert input.dtype in [ + torch.float32, + torch.float16, + torch.bfloat16, + ], f"Unsupported input dtype: {input.dtype}" + assert len(block_size) == input.dim(), ( + f"Got input dim:{input.dim()}, block_size: {block_size}" + ) + shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + original_shape = input.shape + input = input.view(shape_for_reduction) + shape_after_reduction = shape_for_reduction + for i in reduction_dims: + shape_after_reduction[i] = 1 + scale = scale.view(shape_after_reduction) + if zero_point is not None: + zero_point = zero_point.view(shape_after_reduction) + + if zero_point_domain == ZeroPointDomain.INT.name: + quant = torch.clamp( + torch.round(input * (1.0 / scale)) + zero_point, quant_min, quant_max + ) + elif zero_point_domain == ZeroPointDomain.NONE.name: + assert zero_point is None, ( + "zero_point should be None when zero_point_domain is NONE" + ) + quant = torch.clamp(torch.round(input * (1.0 / scale)), quant_min, quant_max) + elif zero_point_domain is None: + # This case handles quantization for float8 we expect no zero point and no zero point domain + assert zero_point is None, ( + "zero_point should be None when zero_point_domain is None" + ) + quant = torch.clamp(input * scale.reciprocal(), quant_min, quant_max) + else: + assert zero_point_domain == ZeroPointDomain.FLOAT.name + mid_point = (quant_max + quant_min + 1) / 2 + min_val = zero_point - scale * mid_point + quant = torch.clamp( + torch.round((input - min_val) / scale), quant_min, quant_max + ) + quant = quant.view(original_shape) + + return quant + + +def dequantize_affine( + input: torch.Tensor, + block_size: tuple[int, ...], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + input_dtype: torch.dtype, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + *, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Args: + input (torch.Tensor): quantized tensor, should match the dtype `dtype` argument + block_size: (List[int]): granularity of quantization, + this means the size of the tensor elements that's sharing the same qparam + e.g. when size is the same as the input tensor dimension, we are using per tensor quantization + scale (Tensor): quantization parameter for affine quantization + zero_point (Tensor): quantization parameter for affine quantization + input_dtype (torch.dtype): requested dtype (e.g. torch.uint8) for output Tensor + quant_min (Optional[int]): minimum quantized value for input Tensor + quant_max (Optional[int]): maximum quantized value for input Tensor + output_dtype (torch.dtype): dtype for output Tensor, default is fp32 + zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be either integer or float + if zero_point is in integer domain, zero point is added to the quantized integer value during + quantization + if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) + value during quantization + default is ZeroPointDomain.INT + + Output: + dequantized Tensor, with requested dtype or fp32 + """ + return _dequantize_affine( + input, + block_size, + scale, + zero_point, + input_dtype, + quant_min, + quant_max, + zero_point_domain.name if zero_point_domain is not None else None, + output_dtype=output_dtype, + ) + + +def _dequantize_affine( + input: torch.Tensor, + block_size: list[int], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + input_dtype: torch.dtype, + quant_min: Optional[Union[int, float, bool]] = None, + quant_max: Optional[Union[int, float, bool]] = None, + zero_point_domain: Optional[str] = ZeroPointDomain.INT.name, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """op definition that has compatible signatures with custom op library""" + # TODO: validate scale/zero_point dimensions are compatible with block_size + if input_dtype not in _SUB_BYTE_UINT_BOUNDS: + assert input.dtype == input_dtype, ( + f"Expected: {input_dtype}, got: {input.dtype}" + ) + assert output_dtype in [ + torch.float32, + torch.float16, + torch.bfloat16, + ], f"Unsupported output dtype: {output_dtype}" + quant_min, quant_max = _get_and_check_qmin_qmax(input_dtype, quant_min, quant_max) + return _dequantize_affine_no_dtype_check( + input, + block_size, + scale, + zero_point, + quant_min, + quant_max, + zero_point_domain, + output_dtype, + ) + + +def _dequantize_affine_no_dtype_check( + input: torch.Tensor, + block_size: list[int], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_min: Union[int, float], + quant_max: Union[int, float], + zero_point_domain: Optional[str] = ZeroPointDomain.INT.name, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """This function converts AQT tensors to their high precision floating point representation + + The op does the following: + 1. figure out the dimension for reduction based on block_size, also reshape the input to align with + the shape after reduction + 2. dequantize the input based on the quantization parameters scale and zero_point and args like zero_point_domain + 3. reshape the quantized result to origianl shape and change dtype to the output_dtype + """ + assert len(block_size) == input.dim(), ( + f"Got input dim:{input.dim()}, block_size: {block_size}" + ) + shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + original_shape = input.shape + input = input.view(shape_for_reduction) + shape_after_reduction = shape_for_reduction + for i in reduction_dims: + shape_after_reduction[i] = 1 + scale = scale.view(shape_after_reduction) + + if zero_point is not None: + zero_point = zero_point.view(shape_after_reduction) + + if zero_point_domain == ZeroPointDomain.INT.name: + # Force a copy to avoid input modification due + # to upcoming in-place operations. + dequant = input.to(torch.int32, copy=True) + if zero_point is not None: + dequant = dequant - zero_point.to(torch.int32) + dequant = dequant.to(output_dtype) + dequant = dequant * scale + elif zero_point_domain == ZeroPointDomain.NONE.name: + assert zero_point is None, ( + "zero_point should be None when zero_point_domain is NONE" + ) + dequant = input.to(output_dtype) + dequant = dequant * scale + elif zero_point_domain is None: + # This case handles dequantization for float8 we expect no zero point and no zero point domain + assert zero_point is None, ( + "zero_point should be None when zero_point_domain is None" + ) + assert _is_float8_type(input.dtype), ( + f"dequantiztion with no zero point domain is only supported with FP8 types, got {input.dtype}" + ) + dequant = input.to(output_dtype) + dequant = dequant * scale + else: + assert zero_point_domain == ZeroPointDomain.FLOAT.name, ( + f"Unexpected zero point domain: {zero_point_domain}" + ) + # TODO: this seems to be a detail for tinygemm (converting from uint to int, probably need to refactor this) + mid_point = (quant_max + quant_min + 1) / 2 + # This should allocate new memory and avoid input modification + dequant = input - mid_point + dequant = dequant.to(output_dtype) + dequant *= scale + if zero_point is not None: + dequant += zero_point + + return dequant.view(original_shape).to(output_dtype) + + +class AffineQuantizedMinMaxObserver(AffineQuantizedObserverBase): + def forward(self, input: torch.Tensor): + if input.numel() == 0: + return input + + input_detached = input.detach() + self.original_dtype = input_detached.dtype + assert self.granularity is not None, "granularity is None" + self.block_size = get_block_size(input_detached.shape, self.granularity) + + shape_for_reduction, reduction_dims = _get_reduction_params( + self.block_size, input_detached.size() + ) + input_detached = input_detached.view(shape_for_reduction) + min_val = torch.amin(input_detached, dim=reduction_dims, keepdim=False) + max_val = torch.amax(input_detached, dim=reduction_dims, keepdim=False) + if not hasattr(self, "min_val") or not hasattr(self, "max_val"): + self.min_val = min_val + self.max_val = max_val + else: + assert self.min_val.shape == min_val.shape, ( + f"Can't update existing min_val - shape mismatch, self.min_val:{self.min_val.shape} != min_val:{min_val.shape}" + ) + assert self.max_val.shape == max_val.shape, ( + f"Can't update existing max_val - shape mismatch, self.max_val {self.max_val.shape} != max_val:{max_val.shape}" + ) + min_val = torch.min(self.min_val, min_val) + max_val = torch.max(self.max_val, max_val) + self.min_val.copy_(min_val) + self.max_val.copy_(max_val) + # returning original input + return input + + def calculate_qparams(self) -> tuple[torch.Tensor, torch.Tensor]: + assert hasattr(self, "min_val") and hasattr(self, "max_val"), ( + "Expecting the observer has min_val and max_val, please run the observer before calling calculate_qparams" + ) + return choose_qparams_affine_with_min_max( + self.min_val, + self.max_val, + self.mapping_type, + [], # BlockSize is not needed because the min/max are already reduced + self.target_dtype, + self.quant_min, + self.quant_max, + self.eps, + self.scale_dtype, + self.zero_point_dtype, + self.preserve_zero, + self.zero_point_domain, + ) + + +class AffineQuantizedMovingAverageMinMaxObserver(AffineQuantizedObserverBase): + def __init__( + self, + mapping_type: MappingType, + target_dtype: torch.dtype, + granularity: Granularity, + averaging_constant=0.01, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + is_dynamic=False, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.INT, + # there could be some extra args that's ignored + **kwargs, + ): + self.is_dynamic = is_dynamic + self.averaging_constant = averaging_constant + if is_dynamic and self.averaging_constant != 1: + raise NotImplementedError( + "MovingAverageMinMaxObserver doesn't support dynamic quantization for " + f"averaging constant of {self.averaging_constant}" + ) + + super().__init__( + mapping_type=mapping_type, + target_dtype=target_dtype, + granularity=granularity, + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + scale_dtype=scale_dtype, + zero_point_dtype=zero_point_dtype, + preserve_zero=preserve_zero, + zero_point_domain=zero_point_domain, + ) + + def forward(self, input: torch.Tensor): + if input.numel() == 0: + return input + + input_detached = input.detach() + self.original_dtype = input_detached.dtype + assert self.granularity is not None, "granularity is None" + self.block_size = get_block_size(input_detached.shape, self.granularity) + + shape_for_reduction, reduction_dims = _get_reduction_params( + self.block_size, input_detached.size() + ) + input_detached = input_detached.view(shape_for_reduction) + min_val = torch.amin(input_detached, dim=reduction_dims, keepdim=False) + max_val = torch.amax(input_detached, dim=reduction_dims, keepdim=False) + if not hasattr(self, "min_val") or not hasattr(self, "max_val"): + self.min_val = min_val + self.max_val = max_val + else: + assert self.min_val.shape == min_val.shape, ( + f"Can't update existing min_val - shape mismatch, self.min_val:{self.min_val.shape} != min_val:{min_val.shape}" + ) + assert self.max_val.shape == max_val.shape, ( + f"Can't update existing max_val - shape mismatch, self.max_val {self.max_val.shape} != max_val:{max_val.shape}" + ) + min_val = self.min_val + self.averaging_constant * (min_val - self.min_val) + max_val = self.max_val + self.averaging_constant * (max_val - self.max_val) + self.min_val.copy_(min_val) + self.max_val.copy_(max_val) + + # returning original input + return input + + def calculate_qparams(self) -> tuple[torch.Tensor, torch.Tensor]: + assert hasattr(self, "min_val") and hasattr(self, "max_val"), ( + "Expecting the observer has min_val and max_val, please run the observer before calling calculate_qparams" + ) + + return choose_qparams_affine_with_min_max( + self.min_val, + self.max_val, + self.mapping_type, + [], # BlockSize is not needed because the min/max are already reduced + self.target_dtype, + self.quant_min, + self.quant_max, + self.eps, + self.scale_dtype, + self.zero_point_dtype, + self.preserve_zero, + self.zero_point_domain, + ) + + +class AffineQuantizedPlaceholderObserver(AffineQuantizedObserverBase): + def __init__( + self, + mapping_type: MappingType, + target_dtype: torch.dtype, + granularity: Granularity, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + is_dynamic=False, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.INT, + # there could be some extra args that's ignored + **kwargs, + ): + self.is_dynamic = is_dynamic + + super().__init__( + mapping_type=mapping_type, + target_dtype=target_dtype, + granularity=granularity, + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + scale_dtype=scale_dtype, + zero_point_dtype=zero_point_dtype, + preserve_zero=preserve_zero, + zero_point_domain=zero_point_domain, + ) + + def forward(self, input): + self.block_size = get_block_size(input.shape, self.granularity) + self.original_dtype = input.dtype + return input + + def calculate_qparams(self): + raise Exception( # noqa: TRY002 + "calculate_qparams should not be called for PlaceholderObserver" + ) diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/_numeric_debugger.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/_numeric_debugger.py new file mode 100644 index 0000000000000000000000000000000000000000..0d66ca71eefcb7bab4bb49affbe0d3616d4e8e70 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/_numeric_debugger.py @@ -0,0 +1,348 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import copy +import logging +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Callable, Optional + +import torch +from torch.ao.ns.fx.utils import compute_sqnr +from torch.export import ExportedProgram +from torch.fx import GraphModule, Node +from torch.nn import functional as F + +from .graph_utils import bfs_trace_with_node_process + +NUMERIC_DEBUG_HANDLE_KEY = "numeric_debug_handle" +CUSTOM_KEY = "custom" + +log = logging.getLogger(__name__) + + +def generate_numeric_debug_handle(ep: ExportedProgram) -> None: + """ + Attach numeric_debug_handle_id for all nodes in the graph module of the given + ExportedProgram, like conv2d, squeeze, conv1d, etc, except for placeholder. + Notice that nodes like getattr are out of scope since they are not in the graph. + + The graph nodes of input exported program are modified inplace. + + Here's an example of using debug handle quantize flow:: + + ep = export_for_training(eager_model, example_inputs) + generate_numeric_debug_handle(ep) + + m = ep.module() + quantizer = XNNPACKQuantizer() + m = prepare_pt2e(m, quantizer) + m = convert_pt2e(m) + """ + + # Sanity check the input data type + if not isinstance(ep, ExportedProgram): + raise ValueError( + f"Expected ep to be ExportedProgram, got {type(ExportedProgram)}" + ) + + unique_id = 0 + + def _find_max_id(node: torch.fx.Node) -> None: + nonlocal unique_id + unique_id = max( + unique_id, node.meta.get(CUSTOM_KEY, {}).get(NUMERIC_DEBUG_HANDLE_KEY, 0) + ) + + def _assign_debug_handle(node: torch.fx.Node) -> None: + nonlocal unique_id + if CUSTOM_KEY not in node.meta: + node.meta[CUSTOM_KEY] = {} + + if NUMERIC_DEBUG_HANDLE_KEY not in node.meta[CUSTOM_KEY]: + node.meta[CUSTOM_KEY][NUMERIC_DEBUG_HANDLE_KEY] = unique_id + unique_id += 1 + + # Find the max ID that exists in the graph first, in case part of the graph + # has already been annotated. This way we guarantee there are no duplicate + # handle IDs. + bfs_trace_with_node_process(ep, _find_max_id) + + unique_id += 1 + + # Assign debug handles to all nodes in the graph that don't have one based on the + # max ID found in the previous step. + bfs_trace_with_node_process(ep, _assign_debug_handle) + + +def _detach(x: object) -> object: + detached: object = None + if isinstance(x, torch.Tensor): + detached = x.detach() + elif isinstance(x, (list, tuple)): + detached = type(x)([_detach(e) for e in x]) + elif isinstance(x, dict): + detached = {k: _detach(e) for k, e in x.items()} + else: + detached = x + return detached + + +def _tensor_shape_equals(x: object, y: object) -> bool: + if isinstance(x, torch.Tensor) and isinstance(y, torch.Tensor): + return x.shape == y.shape + elif isinstance(x, (list, tuple)) and isinstance(y, (list, tuple)): + return all(_tensor_shape_equals(e1, e2) for e1, e2 in zip(x, y)) + elif isinstance(x, dict) and isinstance(y, dict): + all_equal = True + for k in x: + all_equal = all_equal and k in y and (_tensor_shape_equals(x[k], y[k])) + return all_equal + else: + log.debug("Comparing non Tensors: %s and %s, they must be equal", x, y) + return type(x) == type(y) and x == y + + +def _loss_fn( + loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], x: object, y: object +) -> object: + """The returned loss will have the same structure as `x` and `y`, e.g. + if both are Tensor, we'll return a Tensor + if both are list, we'll return a list of Tensors + if both are dict, we'll return a dict with the same key, and value being the loss between the + two Tensors + """ + if isinstance(x, torch.Tensor) and isinstance(y, torch.Tensor): + return loss(x.to(torch.float32), y.to(torch.float32)) + elif isinstance(x, (list, tuple)) and isinstance(y, (list, tuple)): + return type(x)([_loss_fn(loss, e1, e2) for e1, e2 in zip(x, y)]) + elif isinstance(x, dict) and isinstance(y, dict): + return {k: _loss_fn(loss, e, y[k]) for k, e in x.items()} + else: + return None + + +class OutputLogger(torch.nn.Module): + """ + Base class for capturing output values for nodes in a GraphModule, it only captures + Tensor output currently, but we can extend it to work for other types of inputs later if needed + """ + + # Mark as impure so that calls to it will not be removed during DCE. + _is_impure = True + + def __init__( + self, + debug_handle: int, + node_name: Optional[str] = None, + nn_module_stack: Optional[object] = None, + ) -> None: + super().__init__() + self.node_name = node_name + self.nn_module_stack = nn_module_stack + self.debug_handle = debug_handle + self.stats: list[object] = [] + + def forward(self, x: object) -> object: + self.stats.append(_detach(x)) + return x + + def __extra_repr__(self) -> str: + return ( + f"debug_handle={self.debug_handle}, node_name={self.node_name}, " + "nn_module_stack={self.nn_module_stack}, num_stats={len(self.stats)})" + ) + + +def _insert_logger(model: GraphModule, node: Node, debug_handle: int) -> Node: + """For a given node, adds an OutputLogger that observes the output of that node, + and all its users use the OutputLogger output instead. + The OutputLogger will contain the debug_handle which can be used to compare + graphs after transforms""" + + # to avoid circular dep + from torchao.quantization.pt2e.utils import get_new_attr_name_with_prefix + + # add a logger after the node + with model.graph.inserting_after(node): + get_new_attr_name = get_new_attr_name_with_prefix(f"{node.name}_logger") + logger_name = get_new_attr_name(model) + setattr( + model, + logger_name, + OutputLogger(debug_handle, node.name, node.meta.get("nn_module_stack")), + ) + logger_node = model.graph.call_module(logger_name, (node,), {}) + + orig_users = list(node.users.keys()) + for user_node in orig_users: + if user_node is logger_node: + continue + user_node.replace_input_with(node, logger_node) + + return logger_node + + +def prepare_for_propagation_comparison(model: GraphModule) -> GraphModule: + """Add output loggers to node that has numeric_debug_handle + + Args: + model (GraphModule): original model + Returns: + a model with output loggers for all nodes that has numeric_debug_handle_id + """ + # don't change the original model + model = copy.deepcopy(model) + for n in model.graph.nodes: + if ( + CUSTOM_KEY not in n.meta + or NUMERIC_DEBUG_HANDLE_KEY not in n.meta[CUSTOM_KEY] + ): + continue + numeric_debug_handle = n.meta[CUSTOM_KEY][NUMERIC_DEBUG_HANDLE_KEY] + _insert_logger(model, n, numeric_debug_handle) + + model.recompile() + return model + + +@dataclass(frozen=True) +class QuantizationComparisonResult: + actual: torch.Tensor + ref: torch.Tensor + + @property + def mse_loss(self) -> object: + return self.loss(F.mse_loss) + + @property + def sqnr(self) -> object: + return self.loss(compute_sqnr) + + def loss( + self, loss_function: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] + ) -> object: + return _loss_fn(loss_function, self.actual, self.ref) + + def __repr__(self) -> str: + # Don't include the tensors themselves as they are quite large to print + # out. + return ( + f"QuantizationComparisonResult(mse_loss={self.mse_loss}, sqnr={self.sqnr})" + ) + + def __post_init__(self) -> None: + if not isinstance(self.actual, (torch.Tensor, list, tuple, dict)): + raise ValueError( + f"`self.actual` value must be a Tensor, list, tuple or dict, got: {self.actual}" + ) + + if not isinstance(self.ref, (torch.Tensor, list, tuple, dict)): + raise ValueError( + f"`self.ref` value must be a Tensor, list, tuple or dict, got: {self.ref}" + ) + + if not _tensor_shape_equals(self.ref, self.actual): + raise ValueError( + f"Cannot compare tensors with different shapes: ref={self.ref} vs actual={self.actual}" + ) + + +@dataclass(frozen=True) +class NodeAccuracySummary: + handle: int + actual_node_name: str + actual_module_stack: str + ref_node_name: str + ref_module_stack: str + results: Sequence[QuantizationComparisonResult] + + +def _module_stack_to_str(module_stack: object) -> str: + """Simplifies the stack from ("mod", "mod.foo", "mod.foo.0", "mod.foo.0.linear") + to "mod.foo.0.linear" + """ + if not isinstance(module_stack, dict): + return str(module_stack) + module_values_list = list(module_stack.values()) + if len(module_values_list) > 0: + owning_module = module_values_list[-1][0] + return str(owning_module) + else: + return str(module_stack) + + +def extract_results_from_loggers( + model: GraphModule, +) -> dict[int, tuple[Optional[str], object, list[object]]]: + """For a given model, extract the tensors stats and related information for each debug handle. + The reason we have a list of object, instead of Tensor is because the output of node may not be + a Tensor, it could be (nested) list, tuple or dict as well. + + Returns: + A dict is keyed by the debug_handle id and the values are a list of object recorded + in loggers + + """ + # Results maps debug handle to a tensor list for each model being compared. + handles: dict[int, tuple[Optional[str], object, list[object]]] = {} + for _name, module in model.named_children(): + if isinstance(module, OutputLogger) and len(module.stats) > 0: + handles[module.debug_handle] = ( + module.node_name, + module.nn_module_stack, + module.stats, + ) + + return handles + + +def compare_results( + ref_results: dict[int, tuple[Optional[str], object, list[torch.Tensor]]], + actual_results: dict[int, tuple[Optional[str], object, list[torch.Tensor]]], +) -> dict[int, NodeAccuracySummary]: + """Given two dict mapping from `debug_handle_id` (int) to list of tensors + return a map from `debug_handle_id` to `NodeAccuracySummary` that contains + comparison information like SQNR, MSE etc. + + Args: + ref_results (Dict[int, Tuple[str, object, List[torch.Tensor]]]): reference results for each debug_handle_id + actual_results (Dict[int, Tuple[str, object, List[torch.Tensor]]]): actual results for each debug_handle_id + + Returns: + Dict[int, NodeAccuracySummary] + """ + comparisons = {} + for debug_handle, (ref_name, ref_stack, ref_stats) in ref_results.items(): + if debug_handle not in actual_results: + log.debug( + "Cannot compare for handle %s because it wasn't found in the transformed model", + debug_handle, + ) + continue + actual_name, actual_stack, actual_stats = actual_results[debug_handle] + try: + results = [ + QuantizationComparisonResult(actual=a, ref=b) + for a, b in zip(actual_stats, ref_stats) + ] + except Exception as e: + # Add extra information for an exception from QuantizationComparisonResult + # if the shapes didn't match, to include the handle and the node names. + raise ValueError( + f"For numeric_debug_handle={debug_handle} from ref node {ref_name} and actual node {actual_name}" + ) from e + + comparisons[debug_handle] = NodeAccuracySummary( + handle=debug_handle, + actual_node_name=actual_name or "", + actual_module_stack=_module_stack_to_str(actual_stack), + ref_node_name=ref_name or "", + ref_module_stack=_module_stack_to_str(ref_stack), + results=results, + ) + + return comparisons diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/constant_fold.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/constant_fold.py new file mode 100644 index 0000000000000000000000000000000000000000..37a84c45bf49a288d7339518582d5a8038cfcebd --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/constant_fold.py @@ -0,0 +1,410 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import collections +from typing import Any, Callable, Optional + +import torch +import torch.utils._pytree as pytree +from torch._inductor.freezing_utils import maybe_set_is_frozen_param +from torch.utils._ordered_set import OrderedSet + +from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 + +aten = torch.ops.aten + +# We would like to split modules into two subgraphs for runtime weight updates to work correctly. +# The use case and more information could be found at: +# https://docs.google.com/document/d/1inZC-8KarJ6gKB7G9egmYLx1V_dKX_apxon0w4zPC0Q/edit?usp=sharing +META_TAG = "MODULE_TYPE" +MODULE_TAG = "_MAIN_MODULE" +CONST_MODULE_TAG = "_CONST_MODULE" + + +def replace_node_with_constant( + gm: torch.fx.GraphModule, + node: torch.fx.Node, + constant: Optional[torch.Tensor] = None, + name: Optional[str] = None, +) -> None: + g = gm.graph + + if name: + qualname = name + else: + if not hasattr(gm, "_frozen_param_count"): + gm._frozen_param_count = 0 # type: ignore[assignment] + i = gm._frozen_param_count + + while True: + qualname = f"_frozen_param{i}" + if not hasattr(gm, qualname): + break + i += 1 # type: ignore[assignment, operator] + + gm._frozen_param_count = i + 1 # type: ignore[assignment, operator] + + with g.inserting_before(node): + if constant is not None: + new_input_node = g.create_node("get_attr", qualname, (), {}) + else: + # this is the case for lifted constants + new_input_node = g.create_node("placeholder", qualname, (), {}) + node.replace_all_uses_with(new_input_node) + new_input_node.meta.update(node.meta) + g.erase_node(node) + new_input_node.name = node.name + + if constant is not None: + # needed to suppress `does not reference an nn.Module, nn.Parameter, or buffer` warning + gm.register_buffer(qualname, constant) + setattr(gm, qualname, constant) + # mark any constants created during freezing + maybe_set_is_frozen_param(constant) + + +def is_const_source( + node: torch.fx.Node, lifted_constant_names: Optional[list[str]] +) -> bool: + return node.op == "get_attr" or node.name in (lifted_constant_names or ()) + + +class ConstantFolder(torch.fx.Interpreter): + def __init__( + self, + gm: torch.fx.GraphModule, + skip_constructors: bool = False, + lifted_constant_names: Optional[list[str]] = None, + skip_folding_node_fn: Optional[Callable[[torch.fx.Node], bool]] = None, + ) -> None: + super().__init__(gm) + self.node_replacements: dict[torch.fx.Node, Any] = {} + self.replaced_uses: dict[torch.fx.Node, int] = collections.Counter() + self.unknown_value = object() + self.skip_constructors: bool = skip_constructors + + # overwrite this to deallocate env values if their only remaining use + # is the output + self.user_to_last_uses = self.node_to_last_non_output_use() + self.lifted_constant_names = lifted_constant_names + self.deferred_value = object() + self.skip_folding_node_fn = skip_folding_node_fn + + def _support_dynamic_shape(self) -> bool: + # ConstantFolder not support dynamic shape now + return False + + def _deduce_value(self, node: torch.fx.Node) -> Any: + if self.lifted_constant_names is None: + return super().run_node(node) + # if lifted_constant_names is passed in, no concrete value is available + # so we just check if all inputs have values + if self.skip_folding_node_fn is not None and self.skip_folding_node_fn(node): + return self.unknown_value + flattened_node_inps = pytree.arg_tree_leaves(*node.args, **node.kwargs) + for inp in flattened_node_inps: + if ( + isinstance(inp, torch.fx.Node) + and inp.name not in (self.lifted_constant_names or ()) + and self.env[inp] != self.deferred_value + ): + return self.unknown_value + return self.deferred_value + + def is_impure(self, node: torch.fx.node.Node) -> bool: + def is_woq_int8_pattern(node: torch.fx.node.Node) -> bool: + return ( + node.target == torch.ops.prims.convert_element_type.default # type: ignore[return-value] + and isinstance(node.args[0], torch.fx.Node) + and "val" in node.args[0].meta + and node.args[0].meta["val"].dtype == torch.int8 # type: ignore[union-attr] + and node.args[1] == torch.bfloat16 + ) + + if ( + is_woq_int8_pattern(node) + or ( + node.target == torch.ops.aten.permute.default + and len(node.users) == 1 + and is_woq_int8_pattern(next(iter(node.users))) + ) + ) and is_const_source( + node.args[0], # type: ignore[arg-type] + self.lifted_constant_names, + ): + # Case 1: int8_weight -> dq -> bf16_weight + # Case 2: int8_weight -> permute -> dq -> bf16_weight + return True + + DEQUANT_OPS = [ + torch.ops.quantized_decomposed.dequantize_per_channel.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor, + torch.ops.quantized_decomposed.convert_element_type.no_fuse, + ] + + if TORCH_VERSION_AT_LEAST_2_5: + DEQUANT_OPS += [ + torch.ops.torchao.dequantize_affine, + ] + + if node.target in DEQUANT_OPS: + # For the pattern fp32_weight -> q -> dq + # We only folding fp32_weight -> q + # int8_weight and leave dq in graph to be fused + return True + return False + + def node_to_last_non_output_use(self) -> dict[torch.fx.Node, list[torch.fx.Node]]: + last_non_output_use = collections.defaultdict(list) + seen_uses = OrderedSet[torch.fx.Node]() + output_node = next(iter(reversed(self.module.graph.nodes))) # type: ignore[arg-type, union-attr] + + for node in reversed(self.module.graph.nodes): # type: ignore[arg-type, union-attr] + if node.target == "output": + continue + + def add_use(inp: torch.fx.Node) -> None: + if inp in seen_uses: + return + + seen_uses.add(inp) + last_non_output_use[node].append(inp) + + # In-place is fine since we don't mutate + pytree.tree_map_only_(torch.fx.Node, add_use, (node.args, node.kwargs)) + + # if this node is only used in output, we want to gc it right away + if len(node.users) == 1 and output_node in node.users: + last_non_output_use[node].append(node) + + return last_non_output_use + + def run_node(self, node: torch.fx.Node) -> Any: + if node.target == "output": + # because we remove nodes from env on last non output use, + # re-define them now or we'll get error in interpreter + def set_env(arg: torch.fx.Node) -> None: + self.env[arg] = self.unknown_value + + # In-place is fine since we don't mutate + pytree.tree_map_only_(torch.fx.Node, set_env, node.args) + return super().run_node(node) + + args, kwargs = self.fetch_args_kwargs_from_env(node) + flattened_inputs = pytree.arg_tree_leaves(*args, **kwargs) + + # We need to do this weird thing because in cases where flattened_inputs + # contains a ScriptObject, equality checking results in a type error if + # the types are different. + if any( + type(self.unknown_value) == type(input_) and self.unknown_value == input_ + for input_ in flattened_inputs + ): + return self.unknown_value + + # TODO - fix errors with this + if ( + node.op == "call_function" + and node.target == aten._efficientzerotensor.default + ): + return self.unknown_value + + # TODO - constant folding triton kernel returns the inputs -- fix this + if ( + node.op == "call_function" + and node.name == "triton_kernel_wrapper_functional_proxy" + ): + return self.unknown_value + + # skip constructors, since inductor generates optimal code for them already + # and turning into tensor would result in an additional global memory read + # TODO - more complicated strategy + if ( + self.skip_constructors + and not is_const_source(node, self.lifted_constant_names) + and not any(isinstance(e, torch.Tensor) for e in flattened_inputs) + ): + return self.unknown_value + + # All mutations should either be removed or on inputs which we did not make constant + if ( + isinstance(node.target, torch._ops.OpOverload) + and torch.Tag.nondeterministic_seeded in node.target.tags + ): + return self.unknown_value + + if node.op == "call_function" and isinstance( + node.target, torch._ops.HigherOrderOperator + ): + return self.unknown_value + + out = self._deduce_value(node) + + if isinstance(out, torch._C.ScriptObject): + return out + + if out == self.unknown_value: + return self.unknown_value + + if not is_const_source(node, self.lifted_constant_names) and ( + isinstance(out, torch.Tensor) or out == self.deferred_value + ): + if out != self.deferred_value and out.device.type == "meta": + return out + + if not self.insertable_tensor_check(out): + return out + + if self.is_impure(node): + return self.unknown_value + + self.add_node_replacement(node, out) + + flattened_node_inps = pytree.arg_tree_leaves(*node.args, **node.kwargs) + + for n in flattened_node_inps: + if not isinstance(n, torch.fx.Node): + continue + + self.replaced_uses[n] += 1 + + for to_delete in self.user_to_last_uses.get(node, []): + if self.replaced_uses[to_delete] == len(to_delete.users): + self.node_replacements.pop(to_delete, None) + + return out + + def insertable_tensor_check(self, tensor: torch.Tensor) -> bool: + return True + + def add_node_replacement(self, node: torch.fx.Node, tensor: torch.Tensor) -> None: + self.node_replacements[node] = tensor + + def run(self) -> Any: # type: ignore[override] + env: dict[torch.fx.Node, Any] = {} + self.insert_placerholder_values(env) + return super().run(initial_env=env) + + def insert_placerholder_values(self, env: dict[torch.fx.Node, Any]) -> None: + for n in self.module.graph.find_nodes(op="placeholder"): # type: ignore[operator, union-attr] + env[n] = self.unknown_value # type: ignore[assignment] + if self.lifted_constant_names is None: + return + for n in self.module.graph.nodes: # type: ignore[union-attr] + if n.name in (self.lifted_constant_names or ()): + env[n] = self.deferred_value + + +def constant_fold( + gm: torch.fx.GraphModule, + constraint_fn: Optional[Callable[[torch.fx.Node], bool]] = None, +) -> None: + with torch.utils._python_dispatch._disable_current_modes(): + cf = ConstantFolder(gm, skip_constructors=True) + cf.run() + + for node, constant in cf.node_replacements.items(): + if constraint_fn is not None and not constraint_fn(node): + continue + replace_node_with_constant(gm, node, constant) + + erased_params = [] + for node in gm.graph.find_nodes(op="get_attr"): + if len(node.users) == 0: + if hasattr(gm, node.target): + delattr(gm, node.target) + erased_params.append(node) + + for node in erased_params: + gm.graph.erase_node(node) + + gm.graph.eliminate_dead_code() + gm.graph.lint() + gm.recompile() + + +def constant_graph_tag( + gm: torch.fx.GraphModule, + skip_constructors: bool = True, + lifted_constant_names: Optional[list[str]] = None, + skip_folding_node_fn: Optional[Callable[[torch.fx.Node], bool]] = None, +) -> None: + with torch.utils._python_dispatch._disable_current_modes(): + cf = ConstantFolder( + gm, + skip_constructors=skip_constructors, + lifted_constant_names=lifted_constant_names, + skip_folding_node_fn=skip_folding_node_fn, + ) + cf.run() + + for node in gm.graph.nodes: + if skip_folding_node_fn is not None and skip_folding_node_fn(node): + node.meta[META_TAG] = MODULE_TAG + continue + if ( + is_const_source(node, lifted_constant_names) + or node in cf.node_replacements + or node in cf.replaced_uses + ): + node.meta[META_TAG] = CONST_MODULE_TAG + else: + node.meta[META_TAG] = MODULE_TAG + + +def run_and_get_constant_graph( + gm: torch.fx.GraphModule, + skip_constructors: bool = True, + lifted_constant_names: Optional[list[str]] = None, + skip_folding_node_fn: Optional[Callable[[torch.fx.Node], bool]] = None, +) -> torch.fx.GraphModule: + """ + Construct a GraphModule which corresponds to the part which could be + constant folded in provided gm. + """ + + constant_graph_tag( + gm, skip_constructors, lifted_constant_names, skip_folding_node_fn + ) + + def untag(node: torch.fx.Node) -> bool: + used_to_fold = False + for u in node.users: + if u.meta[META_TAG] == CONST_MODULE_TAG: + used_to_fold = True + break + if not used_to_fold: + node.meta[META_TAG] = MODULE_TAG + return used_to_fold + + # We rewrite the tags, if it's a constant being directly consumed, without + # any folding opportunity, we keep it in main gm. + for node in gm.graph.nodes: + if node.op == "getattr" or (node.name in (lifted_constant_names or ())): + untag(node) + + new_graph = torch.fx.Graph() + + node_remapping: dict[torch.fx.Node, torch.fx.Node] = {} + output_nodes = [] + for node in gm.graph.nodes: + if node.meta[META_TAG] == MODULE_TAG: + continue + + new_node = new_graph.node_copy(node, lambda x: node_remapping[x]) + node_remapping[node] = new_node + + for user in node.users: + if user.meta[META_TAG] == MODULE_TAG: + output_nodes.append(new_node) + break + + new_graph.output(tuple(output_nodes)) + new_graph.lint() + new_gm = torch.fx.GraphModule(gm, new_graph) + + return new_gm diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/convert.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/convert.py new file mode 100644 index 0000000000000000000000000000000000000000..2f3e6e0461f6d09401d0cf1d7863e918aff573bf --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/convert.py @@ -0,0 +1,1274 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: ignore-errors + +# TODO: cleanup + +import copy +import operator +import warnings +from typing import Any, Callable, Optional, Union + +import torch +from torch.ao.quantization.backend_config import ( + BackendConfig, + get_native_backend_config, +) +from torch.ao.quantization.backend_config.utils import ( + get_fused_module_classes, + get_pattern_to_dtype_configs, + get_qat_module_classes, + get_root_module_to_quantized_reference_module, +) + +# importing the lib so that the quantized_decomposed ops are registered +from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib # noqa: F401 +from torch.ao.quantization.fx._equalize import ( + convert_eq_obs, + update_obs_for_equalization, +) +from torch.ao.quantization.fx.custom_config import ( + ConvertCustomConfig, + PrepareCustomConfig, +) +from torch.ao.quantization.fx.graph_module import ( + _is_observed_module, + _is_observed_standalone_module, +) +from torch.ao.quantization.fx.lower_to_fbgemm import lower_to_fbgemm +from torch.ao.quantization.fx.qconfig_mapping_utils import ( + _compare_prepare_convert_qconfig_mappings, + _generate_node_name_to_qconfig, + _is_qconfig_supported_by_dtype_configs, + _update_qconfig_for_fusion, + _update_qconfig_for_qat, +) +from torch.ao.quantization.fx.utils import ( + _get_module, + assert_and_get_unique_device, + collect_producer_nodes, + create_getattr_from_value, + graph_module_from_producer_nodes, + node_arg_is_weight, +) +from torch.ao.quantization.qconfig import QConfigAny, qconfig_equals +from torch.ao.quantization.qconfig_mapping import QConfigMapping +from torch.ao.quantization.quantize import _remove_qconfig +from torch.ao.quantization.stubs import DeQuantStub +from torch.ao.quantization.utils import ( + _parent_name, + get_qparam_dict, + is_per_channel, + to_underlying_dtype, + weight_is_quantized, +) +from torch.fx import GraphModule +from torch.fx.graph import Argument, Graph, Node +from torch.fx.graph_module import _USER_PRESERVED_ATTRIBUTES_KEY +from torch.nn.utils.parametrize import type_before_parametrizations + +from torchao.quantization.pt2e import CUSTOM_KEY, NUMERIC_DEBUG_HANDLE_KEY +from torchao.quantization.pt2e.observer import _is_activation_post_process + +__all__ = [ + "convert", + "convert_standalone_module", + "convert_weighted_module", +] + +SUPPORTED_QDTYPES = [ + torch.quint8, + torch.qint8, + torch.qint32, + torch.uint8, + torch.int8, + torch.uint16, + torch.int16, + torch.int32, + torch.float8_e5m2, + torch.float8_e4m3fn, +] + +_QSCHEME_TO_CHOOSE_QPARAMS_OP = { + torch.per_tensor_affine: torch.ops.quantized_decomposed.choose_qparams.tensor, + torch.per_tensor_symmetric: torch.ops.quantized_decomposed.choose_qparams_symmetric.tensor, +} + + +def attach_preserved_attrs_to_model( + model: Union[GraphModule, torch.nn.Module], + preserved_attrs: dict[str, Any], +) -> None: + """Store preserved attributes to the model.meta so that it can be preserved during deepcopy""" + model.meta[_USER_PRESERVED_ATTRIBUTES_KEY] = copy.copy(preserved_attrs) # type: ignore[operator, index, assignment] + # set the preserved attributes in the model so that user can call + # model.attr as they do before calling fx graph mode quantization + for attr_name, attr in model.meta[_USER_PRESERVED_ATTRIBUTES_KEY].items(): # type: ignore[index, union-attr] + setattr(model, attr_name, attr) + + +def _check_is_graph_module(model: torch.nn.Module) -> None: + if not isinstance(model, GraphModule): + raise ValueError( + "input model must be a GraphModule, " + + "Got type:" + + str(type(model)) + + " Please make " + + "sure to follow the tutorials." + ) + + +def _replace_observer_with_quantize_dequantize_node_decomposed( + model: torch.fx.GraphModule, + node: Node, + modules: dict[str, torch.nn.Module], + node_name_to_scope: dict[str, tuple[str, type]], + node_name_to_qconfig: dict[str, QConfigAny], +) -> None: + """Replace activation_post_process module call node with quantize and + dequantize node working with decomposed Tensor + + Before: + ... -> observer_0(x) -> ... + After: + ... -> torch.ops.quantized_decomposed.quantize_per_tensor(x, ...) -> + torch.ops.quantized_decomposed.dequantize_per_tensor() -> ... + + or quantize_per_channel and dequantize_per_channel + """ + graph = model.graph + assert modules is not None + assert isinstance(node.target, str) + module_path, prefix = _get_module_path_and_prefix( + node, node_name_to_scope, node_name_to_qconfig + ) + activation_post_process = modules[node.target] + if hasattr(activation_post_process, "convert"): + activation_post_process.convert(model, node) + return + # skip replacing observers to quant/dequant nodes if the qconfigs of all + # consumers and producers of this observer are None + skip_replacement = all( + _has_none_qconfig(n, node_name_to_qconfig) + for n in list(node.args) + list(node.users.keys()) + ) + if skip_replacement or not _is_conversion_supported(activation_post_process): + # didn't find corresponding quantize op and info for the activation_post_process + # so we just remove the observer + with graph.inserting_before(node): + node.replace_all_uses_with(node.args[0]) + graph.erase_node(node) + return + + # otherwise, we can convert the activation_post_process module call to quantize/dequantize node + + # 1. extract the information from activation_post_process module for generating + # the quantize and dequantize operator + dtype = activation_post_process.dtype # type: ignore[attr-defined] + + is_dynamic = False + if hasattr(activation_post_process, "is_dynamic"): + is_dynamic = activation_post_process.is_dynamic # type: ignore[assignment] + + def add_dequantize_op_kwargs(dequantize_op, input_node): + dequantize_op_kwargs = {} + if "val" in input_node.meta: + dq_out_dtype = input_node.meta["val"].dtype + if dq_out_dtype != torch.float32: + dequantize_op_kwargs = {"out_dtype": dq_out_dtype} + return dequantize_op_kwargs + + if dtype in SUPPORTED_QDTYPES and (not is_dynamic): + # TODO: probably should cleanup this condition check, it's hard + # to reason about this if and the following elif + + # uint8/int8/int32 static quantization branch + + # 1. extract information for inserting q/dq node from activation_post_process + node_type = "call_function" + quantize_op: Optional[Callable] = None + scale, zero_point = activation_post_process.calculate_qparams() # type: ignore[attr-defined, operator] + if is_per_channel(activation_post_process.qscheme): # type: ignore[attr-defined] + ch_axis = int(activation_post_process.ch_axis) # type: ignore[attr-defined, arg-type] + quantize_op = torch.ops.quantized_decomposed.quantize_per_channel.default + dequantize_op = ( + torch.ops.quantized_decomposed.dequantize_per_channel.default + ) + quant_min = activation_post_process.quant_min + quant_max = activation_post_process.quant_max + dtype_ = to_underlying_dtype(dtype) + qparams = { + "_scale_": scale, + "_zero_point_": zero_point, + "_axis_": ch_axis, + "_quant_min_": quant_min, + "_quant_max_": quant_max, + "_dtype_": dtype_, + } + else: + quantize_op = torch.ops.quantized_decomposed.quantize_per_tensor.default + dequantize_op = torch.ops.quantized_decomposed.dequantize_per_tensor.default + scale = float(scale) + zero_point = int(zero_point) + quant_min = activation_post_process.quant_min # type: ignore[attr-defined] + quant_max = activation_post_process.quant_max # type: ignore[attr-defined] + dtype_ = to_underlying_dtype(dtype) + qparams = { + "_scale_": scale, + "_zero_point_": zero_point, + "_quant_min_": quant_min, + "_quant_max_": quant_max, + "_dtype_": dtype_, + } + + # 2. replace activation_post_process node with quantize and dequantize + with graph.inserting_before(node): + input_node = node.args[0] + quantize_op_inputs = [input_node] + for key, value_or_node in qparams.items(): + # TODO: we can add the information of whether a value needs to + # be registered as an attribute in qparams dict itself + if key in ["_scale_", "_zero_point_"] and ( + not isinstance(value_or_node, (float, int)) + ): + # For scale and zero_point values we register them as buffers in the root module. + # However, note that when the values are not tensors, as in the case of + # per_tensor quantization, they will be treated as literals. + # However, registering them as a node seems to cause issue with dynamo + # tracing where it may consider tensor overload as opposed to default. + # With extra check of scale and zero_point being scalar, it makes + # sure that the default overload can be used. + # TODO: maybe need more complex attr name here + qparam_node = create_getattr_from_value( + model, graph, module_path + prefix + key, value_or_node + ) + quantize_op_inputs.append(qparam_node) + else: + # for qparams that are not scale/zero_point (like axis, dtype) we store them as literals in the graph. + quantize_op_inputs.append(value_or_node) + + quantized_node = graph.create_node( + node_type, quantize_op, tuple(quantize_op_inputs), {} + ) + # use the same qparams from quantize op + dq_inputs = [quantized_node] + quantize_op_inputs[1:] + dequantized_node = graph.call_function( + dequantize_op, + tuple(dq_inputs), + add_dequantize_op_kwargs(dequantize_op, input_node), + ) + + node.replace_all_uses_with(dequantized_node) + # propagate numeric debug handle from observer/fake_quant node to dequantize node + if ( + CUSTOM_KEY in node.meta + and NUMERIC_DEBUG_HANDLE_KEY in node.meta[CUSTOM_KEY] + ): + if CUSTOM_KEY not in dequantized_node.meta: + dequantized_node.meta[CUSTOM_KEY] = {} + dequantized_node.meta[CUSTOM_KEY][NUMERIC_DEBUG_HANDLE_KEY] = node.meta[ + CUSTOM_KEY + ][NUMERIC_DEBUG_HANDLE_KEY] + graph.erase_node(node) + elif is_dynamic: + # uint8/int8/fp16 dynamic quantization + + # 1. extract information for inserting q/dq node from activation_post_process + node_type = "call_function" + quantize_op = torch.ops.quantized_decomposed.quantize_per_tensor.tensor + # we only use choose_qparams for is_decomposed now, + # but we should probably align the non-decomposed path with this as well, + # and that can be done after we remove reduce_range flag + # 1. extract qparams from activation_post_process module + dtype_ = to_underlying_dtype(dtype) + assert dtype_ in [torch.uint8, torch.int8], ( + "only uint8 and int8 are supported in reference flow for " + "dynamic quantization right now" + ) + quant_min = activation_post_process.quant_min # type: ignore[attr-defined] + quant_max = activation_post_process.quant_max # type: ignore[attr-defined] + qscheme = getattr(activation_post_process, "qscheme", torch.per_tensor_affine) # type: ignore[attr-defined] + eps = getattr(activation_post_process, "eps", torch.finfo(torch.float32).eps) # type: ignore[attr-defined] + # note: scale and zero_point are missing for quantize_per_tensor op + # we'll need to get this from choose_qparams op, which we'll add after + # this step + qparams = { + "_quant_min_": quant_min, + "_quant_max_": quant_max, + "_eps_": eps, + "_dtype_": dtype_, + } + + choose_qparams_op = _QSCHEME_TO_CHOOSE_QPARAMS_OP[qscheme] + # 2. insert choose_qparams op and update the qparams list + with graph.inserting_before(node): + input_node = node.args[0] + choose_qparams_op_inputs = [node.args[0]] + for key, value in qparams.items(): + # we have quant_min, quant_max and dtype, all should be stored + # as literals + choose_qparams_op_inputs.append(value) + choose_qparams_node = graph.create_node( + "call_function", choose_qparams_op, tuple(choose_qparams_op_inputs), {} + ) + # choose_qparms returns (scale, zero_point) + scale_node = graph.create_node( + "call_function", operator.getitem, (choose_qparams_node, 0), {} + ) + zero_point_node = graph.create_node( + "call_function", operator.getitem, (choose_qparams_node, 1), {} + ) + quant_min = qparams["_quant_min_"] + quant_max = qparams["_quant_max_"] + dtype = qparams["_dtype_"] + qparams = { + "_scale_": scale_node, + "_zero_point_": zero_point_node, + "_quant_min_": quant_min, + "_quant_max_": quant_max, + "_dtype_": dtype, + } + + # 3. replace activation_post_process node to quantize and dequantize node + with graph.inserting_before(node): + input_node = node.args[0] + quantize_op_inputs = [input_node] + for key, value_or_node in qparams.items(): + # TODO: we can add the information of whether a value needs to + # be registered as an attribute in qparams dict itself + if key in ["_scale_", "_zero_point_"]: + # in this case we have a node in the graph since it's dynamically + # computed from the input, with choose_qparams op + qparam_node = value_or_node + quantize_op_inputs.append(qparam_node) + else: + # for qparams that are not scale/zero_point (like axis, dtype) we + # store them as literals in the graph. + quantize_op_inputs.append(value_or_node) + + quantized_node = graph.create_node( + node_type, quantize_op, tuple(quantize_op_inputs), {} + ) + # use the same qparams from quantize op + dq_inputs = [quantized_node] + quantize_op_inputs[1:] + # need to use the tensor variant of this op, since scale and zero_point + # from choose_qparam are Tensors, instead of float/int, this is to + # prevent these nodes being traced away by downstream systems + dequantize_op = torch.ops.quantized_decomposed.dequantize_per_tensor.tensor + dequantized_node = graph.call_function( + dequantize_op, + tuple(dq_inputs), + add_dequantize_op_kwargs(dequantize_op, input_node), + ) + + node.replace_all_uses_with(dequantized_node) + # propagate numeric debug handle from observer/fake_quant node to dequantize node + if NUMERIC_DEBUG_HANDLE_KEY in node.meta: + dequantized_node.meta[NUMERIC_DEBUG_HANDLE_KEY] = node.meta[ + NUMERIC_DEBUG_HANDLE_KEY + ] + graph.erase_node(node) + elif dtype == torch.float16: + # Insert to_fp16 -> to_fp32 node + dtype_convert_op = torch.ops.quantized_decomposed.convert_element_type.no_fuse + with graph.inserting_before(node): + input_node = node.args[0] + convert_fp16_node = graph.create_node( + "call_function", dtype_convert_op, (input_node, torch.float16), {} + ) + convert_fp32_node = graph.create_node( + "call_function", dtype_convert_op, (convert_fp16_node, torch.float), {} + ) + node.replace_all_uses_with(convert_fp32_node) + graph.erase_node(node) + + # should not reach since we have checks in the beginning to make sure the + # activation_post_process is supported + + +def _replace_observer_with_quantize_dequantize_node( + model: torch.fx.GraphModule, + node: Node, + modules: dict[str, torch.nn.Module], + node_name_to_scope: dict[str, tuple[str, type]], + node_name_to_qconfig: dict[str, QConfigAny], +) -> None: + """Replace activation_post_process module call node with quantize and + dequantize node + + Before: + ... -> observer_0(x) -> ... + After: + ... -> torch.quantize_per_tensor(x, ...) -> x.dequantize() -> ... + """ + assert modules is not None + assert isinstance(node.target, str) + graph = model.graph + module_path, prefix = _get_module_path_and_prefix( + node, node_name_to_scope, node_name_to_qconfig + ) + activation_post_process = modules[node.target] + # skip replacing observers to quant/dequant nodes if the qconfigs of all + # consumers and producers of this observer are None + skip_replacement = all( + _has_none_qconfig(n, node_name_to_qconfig) + for n in list(node.args) + list(node.users.keys()) + ) + if skip_replacement or not _is_conversion_supported(activation_post_process): + # didn't find corresponding quantize op and info for the activation_post_process + # so we just remove the observer + with graph.inserting_before(node): + node.replace_all_uses_with(node.args[0]) + graph.erase_node(node) + return + + # otherwise, we can convert the activation_post_process module call to quantize/dequantize node + dtype = activation_post_process.dtype # type: ignore[attr-defined] + + is_dynamic = False + if hasattr(activation_post_process, "is_dynamic"): + is_dynamic = activation_post_process.is_dynamic # type: ignore[attr-defined, assignment] + + if dtype in [ + torch.quint8, + torch.qint8, + torch.qint32, + torch.float8_e5m2, + torch.float8_e4m3fn, + ] and (not is_dynamic): + # TODO: probably should cleanup this condition check, it's hard + # to reason about this if and the following elif + + # uint8/int8/int32 static quantization branch + + # 1. extract the information from activation_post_process module for generating + # the quantize and dequantize operator + node_type = "call_function" + quantize_op: Optional[Callable] = None + scale, zero_point = activation_post_process.calculate_qparams() # type: ignore[attr-defined, operator] + if is_per_channel(activation_post_process.qscheme): # type: ignore[attr-defined] + ch_axis = int(activation_post_process.ch_axis) # type: ignore[attr-defined, arg-type] + qparams = { + "_scale_": scale, + "_zero_point_": zero_point, + "_axis_": ch_axis, + "_dtype_": dtype, + } + quantize_op = torch.quantize_per_channel + else: + scale = float(scale) + zero_point = int(zero_point) + qparams = {"_scale_": scale, "_zero_point_": zero_point, "_dtype_": dtype} + quantize_op = torch.quantize_per_tensor + + # 2. replace activation_post_process node with quantize and dequantize + with graph.inserting_before(node): + input_node = node.args[0] + quantize_op_inputs = [input_node] + for key, value_or_node in qparams.items(): + # TODO: we can add the information of whether a value needs to + # be registered as an attribute in qparams dict itself + if key in ["_scale_", "_zero_point_"]: + # For scale and zero_point values we register them as buffers in the root module. + # TODO: maybe need more complex attr name here + qparam_node = create_getattr_from_value( + model, graph, module_path + prefix + key, value_or_node + ) + quantize_op_inputs.append(qparam_node) + else: + # for qparams that are not scale/zero_point (like axis, dtype) we store them as literals in the graph. + quantize_op_inputs.append(value_or_node) + + quantized_node = graph.create_node( + node_type, quantize_op, tuple(quantize_op_inputs), {} + ) + dequantized_node = graph.call_method("dequantize", args=(quantized_node,)) + node.replace_all_uses_with(dequantized_node) + graph.erase_node(node) + elif is_dynamic: + # uint8/int8/fp16 dynamic quantization branch + + node_type = "call_function" + quantize_op = torch.quantize_per_tensor_dynamic + # TODO: get reduce range from observer + # reduce_range = activation_post_process.reduce_range + reduce_range = torch.backends.quantized.engine in ("fbgemm", "x86") + qparams = {"_dtype_": dtype, "_reduce_range_": reduce_range} + + with graph.inserting_before(node): + input_node = node.args[0] + quantize_op_inputs = [input_node] + for key, value in qparams.items(): + quantize_op_inputs.append(value) + + quantized_node = graph.create_node( + node_type, quantize_op, tuple(quantize_op_inputs), {} + ) + dequantized_node = graph.call_method("dequantize", args=(quantized_node,)) + node.replace_all_uses_with(dequantized_node) + graph.erase_node(node) + elif dtype == torch.float16: + node_type = "call_method" + quantize_op = "to" # type: ignore[assignment] + qparams = {"_dtype_": dtype} + with graph.inserting_before(node): + input_node = node.args[0] + quantize_op_inputs = [input_node] + for key, value in qparams.items(): + # TODO: we can add the information of whether a value needs to + # be registered as an attribute in qparams dict itself + quantize_op_inputs.append(value) + + quantized_node = graph.create_node( + node_type, quantize_op, tuple(quantize_op_inputs), {} + ) + dequantized_node = graph.call_method("dequantize", args=(quantized_node,)) + node.replace_all_uses_with(dequantized_node) + graph.erase_node(node) + + # should not reach since we have checks in the beginning to make sure the + # activation_post_process is supported + + +# this is a temporary hack for custom module, we may want to implement +# this properly after the custom module class design is finalized +# TODO: DeQuantStubs are currently inserted only after custom module LSTM, while observers are inserted +# after all other custom modules. In the future, we should simply insert QuantStubs before and DeQuantStubs +# after custom modules in general, and replace these with "quantize" and "dequantize" nodes respectively. +def _replace_observer_or_dequant_stub_with_dequantize_node( + node: Node, graph: Graph +) -> None: + call_custom_module_node = node.args[0] + assert isinstance(call_custom_module_node, Node), ( + f"Expecting the for call custom module node to be a Node, but got {call_custom_module_node}" + ) + node.replace_all_uses_with(call_custom_module_node) + graph.erase_node(node) + _insert_dequantize_node(call_custom_module_node, graph) + + +def _is_conversion_supported(activation_post_process: torch.nn.Module) -> bool: + dtype = activation_post_process.dtype # type: ignore[attr-defined] + + is_dynamic = False + if hasattr(activation_post_process, "is_dynamic"): + is_dynamic = activation_post_process.is_dynamic # type: ignore[attr-defined, assignment] + + return ( + (dtype in SUPPORTED_QDTYPES and (not is_dynamic)) + or is_dynamic # type: ignore[return-value] + or dtype == torch.float16 + ) + + +def _has_none_qconfig( + node: Argument, node_name_to_qconfig: dict[str, QConfigAny] +) -> bool: + """Check if a node has a qconfig of None, i.e. user requested to not quantize + the node + """ + return ( + isinstance(node, Node) + and node.name in node_name_to_qconfig + and node_name_to_qconfig[node.name] is None + ) + + +def _run_weight_observers(observed: GraphModule, backend_config: BackendConfig) -> None: + """Extract the subgraph that produces the weight for dynamic quant + or weight only quant node and run the subgraph to observe the weight. + Note that the observers of dynamic quant or weight only quant ops are + run during the convert step. + """ + for node in observed.graph.nodes: + if node.op != "call_function": + continue + for node_arg in node.args: + # node_arg is weight + if node_arg and node_arg_is_weight(node, node_arg): + weight_observer_nodes = collect_producer_nodes(node_arg) + if weight_observer_nodes is None: + continue + weight_observer_module = graph_module_from_producer_nodes( + observed, weight_observer_nodes + ) + # run the weight observer + weight_observer_module() + + +def _maybe_recursive_remove_dequantize(arg: Any, node: Node, graph: Graph) -> None: + """If the arg is a dequantize Node, or a list/tuple/dict of dequantize Node, + we'll recursively remove the dequantize Node + """ + if isinstance(arg, Node) and arg.op == "call_method" and arg.target == "dequantize": + quantize_node = arg.args[0] + # we only replace the specific use since dequantize could be used by other nodes + # as well + node.replace_input_with(arg, quantize_node) + elif isinstance(arg, (list, tuple)): + for arg_element in arg: + _maybe_recursive_remove_dequantize(arg_element, node, graph) + elif isinstance(arg, dict): + for arg_element in arg.values(): + _maybe_recursive_remove_dequantize(arg_element, node, graph) + else: + warnings.warn( + f"Unsupported node type in recursive remove dequantize: {type(arg)}" + ) + + +def _get_module_path_and_prefix( + obs_node: Node, + node_name_to_scope: dict[str, tuple[str, type]], + node_name_to_qconfig: dict[str, QConfigAny], +) -> tuple[str, str]: + """Given and observer node, get the `Scope` or the fully qualified name for + the submodule containing the observed node, also return a prefix of "_input" + when the observed node is an input of a F.linear op, and not the output of another + quantized op. + TODO: this logic is hacky, we should think about how to remove it or make it more + general + """ + observed_node = obs_node.args[0] + # an observer can be inserted for both input of the next operator or output of the previous + # operator (they can be the same) + # this flag identifies if the observer is inserted only because the observed node is + # the input of the next operator + assert isinstance(observed_node, Node), ( + f"Expecting observed node to be a Node, but got {observed_node}" + ) + is_input_observer_only = ( + node_name_to_qconfig[observed_node.name] is None + if observed_node.name in node_name_to_qconfig + else None + ) + if is_input_observer_only: + # if the quantize function is at the input of op, then we find the first user of the observer_node + # to get the path. If a linear call_function is in the user list, we return the first instance + # of linear node to get the FQN. + users = list(obs_node.users) + first_linear_use_or_first_use = users[0] if users else None + linear_node = None + for n in users: + if n.op == "call_function" and n.target == torch.nn.functional.linear: + linear_node = n + break + if linear_node: + first_linear_use_or_first_use = linear_node + prefix = "_input" + else: + # if the quantize function is at the output of the op, we use the observer input node to get the path + first_linear_use_or_first_use = observed_node + prefix = "" + + if ( + first_linear_use_or_first_use + and first_linear_use_or_first_use.name in node_name_to_scope + ): + module_path, _ = node_name_to_scope[first_linear_use_or_first_use.name] + else: + # TODO: it's not used, so actually we can skip quantization + # but this requires changing return type of quantize_node + # we can fix it later if needed + module_path = "" + return module_path, prefix + + +def _insert_dequantize_node(node: Node, graph: Graph) -> None: + """Inserts dequantize node for `node` in `graph`""" + with graph.inserting_after(node): + dequantize_node = graph.call_method("dequantize", (node,)) + for user_node in dict(node.users): + if user_node is not dequantize_node: + user_node.replace_input_with(node, dequantize_node) + + +def _maybe_get_observer_for_node( + node: Node, modules: dict[str, torch.nn.Module] +) -> Optional[torch.nn.Module]: + """ + If the node is observed, return the observer + instance. Otherwise, return None. + """ + for maybe_obs_node in node.users.keys(): + if maybe_obs_node.op == "call_module": + maybe_obs = modules[str(maybe_obs_node.target)] + if _is_activation_post_process(maybe_obs): + return maybe_obs + return None + + +def convert_standalone_module( + node: Node, + modules: dict[str, torch.nn.Module], + model: torch.fx.GraphModule, + is_reference: bool, + backend_config: Optional[BackendConfig], +) -> None: + """Converts a observed standalone module to a quantized standalone module by calling + the fx convert api, currently using the same `is_reference` flag as parent, but we may + changing this behavior in the future (e.g. separating quantization and lowering for + standalone module as well) + + Args: + - node: The call_module node of the observed standalone module + - modules: named_module of original model + - model: original model + - is_reference: a flag from parent provided by user to decide if we want to + produce a reference model or a fbgemm/qnnpack model + - backend_config: backend configuration of the target backend of quantization + """ + # TODO: remove is_reference flag + if is_reference: + convert_fn = torch.ao.quantization.quantize_fx.convert_to_reference_fx + else: + convert_fn = torch.ao.quantization.quantize_fx.convert_fx # type: ignore[attr-defined] + # We know that observed standalone module is a GraphModule since + # it's produced by us + observed_standalone_module: GraphModule = modules[str(node.target)] # type: ignore[assignment] + sm_input_quantized_idxs = observed_standalone_module.meta[ + "_observed_graph_module_attrs" + ].standalone_module_input_quantized_idxs + # remove the dequantize nodes for inputs + args = list(node.args) + for idx in range(len(args)): + if idx in sm_input_quantized_idxs: + arg = args[idx] + if arg.op == "call_method" and arg.target == "dequantize": # type: ignore[union-attr] + quantize_node = arg.args[0] # type: ignore[union-attr] + node.replace_input_with(arg, quantize_node) + if len(arg.users) == 0: # type: ignore[union-attr] + model.graph.erase_node(arg) + # add dequantize node for output + sm_output_quantized_idxs = observed_standalone_module.meta[ + "_observed_graph_module_attrs" + ].standalone_module_output_quantized_idxs + if len(sm_output_quantized_idxs) > 0: + assert sm_output_quantized_idxs[0] == 0, "Currently only quantized" + "output idxs = [0] is supported" + + # if it's non-empty, then it means the output is kept in quantized form + # we'll just add a dequantize node after this node + _insert_dequantize_node(node, model.graph) + + # TODO: allow convert_custom_config to override backend_config + # for standalone module + quantized_standalone_module = convert_fn( + observed_standalone_module, backend_config=backend_config + ) + parent_name, name = _parent_name(node.target) + # update the modules dict + setattr(modules[parent_name], name, quantized_standalone_module) + modules[str(node.target)] = quantized_standalone_module + + +def convert_weighted_module( + node: Node, + modules: dict[str, torch.nn.Module], + observed_node_names: set[str], + node_name_to_qconfig: dict[str, QConfigAny], + backend_config: BackendConfig, + is_decomposed: bool = False, + is_reference: bool = False, +) -> None: + """Convert a weighted module to reference quantized module in the model + If the QConfig of a QAT module is not set, the module will still be converted to + a float module. + + Args: + - node: The call_module node of the observed standalone module + - modules: named_module of original model + - observed_node_names: names for the set of observed fx node, we can skip + this conversion if the node is not observed + """ + original_module = modules[str(node.target)] + qconfig: QConfigAny = original_module.qconfig # type: ignore[assignment] + weight_post_process = None + qat_module_classes = get_qat_module_classes(backend_config) + + if isinstance(original_module, qat_module_classes): + # Converting qat module to a float module, we need to attach + # weight fake_quant to the module, weight fake_quant is assumed to be run during + # QAT so we don't need to run it again here + weight_post_process = original_module.weight_fake_quant + original_module = original_module.to_float() # type: ignore[operator] + # change qat module to float module + parent_name, name = _parent_name(node.target) + setattr(modules[parent_name], name, original_module) + + is_observed = node.name in observed_node_names + # If a qconfig is not defined for this node, then skip converting to a reference module + if ( + qconfig is None + or _has_none_qconfig(node, node_name_to_qconfig) + or not is_observed + ): + return + + # skip converting to reference quantized module if the qconfig is not supported + pattern_to_dtype_configs = get_pattern_to_dtype_configs(backend_config) + dtype_configs = pattern_to_dtype_configs.get(type(original_module), []) + if not _is_qconfig_supported_by_dtype_configs(qconfig, dtype_configs): + return + + # TODO: rename weight_is_statically_quantized to weight_is_int8_quantized + is_weight_quantized = weight_is_quantized(qconfig) + + # the condition for swapping the module to reference quantized module is: + # weights need to be quantized + if not is_weight_quantized: + return + + fused_module = None + float_module = original_module + # extract the individual float_module and fused module + if isinstance(original_module, torch.ao.nn.intrinsic._FusedModule): + fused_module = float_module + float_module = fused_module[0] # type: ignore[index] + + # TODO: move this to the reference quantized module + # weight_qparams or weight_qparams dict + wq_or_wq_dict = {"is_decomposed": is_decomposed} + if isinstance(float_module, torch.nn.RNNCellBase): + weight_post_process_ih = qconfig.weight() # type: ignore[union-attr, operator] + weight_post_process_hh = qconfig.weight() # type: ignore[union-attr, operator] + weight_post_process_ih(float_module.weight_ih) + weight_post_process_hh(float_module.weight_hh) + weight_qparams_ih = get_qparam_dict(weight_post_process_ih) + weight_qparams_hh = get_qparam_dict(weight_post_process_hh) + wq_or_wq_dict.update( + { + "weight_ih": weight_qparams_ih, + "weight_hh": weight_qparams_hh, + } + ) + elif isinstance(float_module, (torch.nn.LSTM, torch.nn.GRU)): + # format for wq_or_wq_dict (flattened attributes): + # {"weight_ih_l0_scale": ..., "weight_ih_l0_qscheme": ..., ...} + for wn in float_module._flat_weights_names: + if hasattr(float_module, wn) and wn.startswith("weight"): + weight = getattr(float_module, wn) + weight_post_process = qconfig.weight() # type: ignore[union-attr, operator] + if weight_post_process.dtype == torch.qint8: # type: ignore[union-attr] + weight_post_process(weight) # type: ignore[operator, misc] + wq_or_wq_dict[wn] = get_qparam_dict(weight_post_process) + else: + # weight_post_process is None means the original module is not a QAT module + # we need to get weight_post_process from qconfig in this case + is_ptq = weight_post_process is None + if is_ptq: + weight_post_process = qconfig.weight() # type: ignore[union-attr, operator] + device = assert_and_get_unique_device(float_module) + if device: + weight_post_process.to(device) + + # Call weight observer/fake_quant at least once to ensure the scales and zero points + # have the right shapes. Note: there are two cases where we don't have to do this: + # + # (1) QAT: The model's forward method already calls the weight observer/fake_quant, + # and this typically happens during training, so we don't need to do it here. + # + # (2) Non-reference (lowered) case: The quantized module's from_float method already + # calls the weight observer/fake_quant, so we don't have to do it here. + # + # Currently we ignore both cases and call the weight observer/fake_quant here + # regardless, which is technically incorrect. For (1), this is mainly to preserve BC + # in test code, which may not always train before convert. In the future, we should + # break BC for these two cases. See https://github.com/pytorch/pytorch/issues/73941. + # + # For PT2, however, we don't need to preserve BC here, so we can skip this hack + # for QAT. We identify this case as (is_decomposed + is_reference + is_qat). + # Note that we still need it for PTQ in the PT2 flow since the model's forward + # method doesn't call the weight observer. + is_qat = not is_ptq + if not (is_decomposed and is_reference and is_qat): + weight_post_process(float_module.weight) # type: ignore[operator] + + wq_or_wq_dict.update(get_qparam_dict(weight_post_process)) + + # We use the same reference module for all modes of quantization: static, dynamic, weight_only + # root_module_to_quantized_reference_module: module mapping from root (floating point) module class + # to quantized reference module class, e.g. nn.Conv2d to nn.quantized._reference.Conv2d + root_module_to_quantized_reference_module = ( + get_root_module_to_quantized_reference_module(backend_config) + ) + ref_qmodule_cls = root_module_to_quantized_reference_module.get( + type_before_parametrizations(float_module), None + ) + assert ref_qmodule_cls is not None, ( + f"No reference quantized module class configured for {type_before_parametrizations(float_module)}" + ) + ref_qmodule = ref_qmodule_cls.from_float(float_module, wq_or_wq_dict) # type: ignore[attr-defined] + if fused_module is not None: + fused_module[0] = ref_qmodule # type: ignore[operator] + else: + parent_name, name = _parent_name(node.target) + setattr(modules[parent_name], name, ref_qmodule) + + +def convert( + model: GraphModule, + is_reference: bool = False, + convert_custom_config: Union[ConvertCustomConfig, dict[str, Any], None] = None, + is_standalone_module: bool = False, + _remove_qconfig_flag: bool = True, + qconfig_mapping: Union[QConfigMapping, dict[str, Any], None] = None, + backend_config: Union[BackendConfig, dict[str, Any], None] = None, + is_decomposed: bool = False, + keep_original_weights: bool = False, +) -> GraphModule: + """ + We will convert an observed model (a module with observer calls) to a reference + quantized model, the rule is simple: + 1. for each observer module call in the graph, we'll convert it to calls to + quantize and dequantize functions based on the observer instance + 2. for weighted operations like linear/conv, we need to convert them to reference + quantized module, this requires us to know whether the dtype configured for the + weight is supported in the backend, this is done in prepare step and the result + is stored in observed_node_names, we can decide whether we need to swap the + module based on this set + + Args: + * `is_standalone_module`: when this flag is True, it means we are quantizing + a submodule that is not inlined in parent module, and will be quantized + separately as one unit. + + * `is_decomposed`: a boolean flag to indicate whether we want to use the + quantize operator for decomposed quantized tensor + (torch.ops.quantized_decomposed.quantize_per_tensor) or default/standalone + quantized tensor (torch.quantize_per_tensor) + + Returns: + a quantized standalone module, whether input/output is quantized is + specified by prepare_custom_config, with + input_quantized_idxs, output_quantized_idxs, please + see docs for :func:`~torch.ao.quantization.prepare_fx` for details + """ + if convert_custom_config is None: + convert_custom_config = ConvertCustomConfig() + + if isinstance(convert_custom_config, dict): + warnings.warn( + "Passing a convert_custom_config_dict to convert is deprecated and will not be supported " + "in a future version. Please pass in a ConvertCustomConfig instead.", + FutureWarning, + stacklevel=2, + ) + convert_custom_config = ConvertCustomConfig.from_dict(convert_custom_config) + + if isinstance(qconfig_mapping, dict): + warnings.warn( + "Passing a QConfig dictionary to convert is deprecated and will not be supported " + "in a future version. Please pass in a QConfigMapping instead.", + FutureWarning, + stacklevel=2, + ) + qconfig_mapping = ( + QConfigMapping.from_dict(qconfig_mapping) if qconfig_mapping else None + ) + qconfig_mapping = copy.deepcopy(qconfig_mapping) + assert qconfig_mapping is None or isinstance(qconfig_mapping, QConfigMapping) + + if isinstance(backend_config, dict): + warnings.warn( + "Passing a backend_config_dict to prepare is deprecated and will not be supported " + "in a future version. Please pass in a BackendConfig instead.", + FutureWarning, + stacklevel=2, + ) + backend_config = BackendConfig.from_dict(backend_config) + + if backend_config is None: + backend_config = get_native_backend_config() + + assert _is_observed_module(model), "incoming model must be produced by prepare_fx" + observed_graph_module_attrs = model.meta["_observed_graph_module_attrs"] + node_name_to_scope: dict[str, tuple[str, type]] = ( + observed_graph_module_attrs.node_name_to_scope + ) + prepare_custom_config: PrepareCustomConfig = ( + observed_graph_module_attrs.prepare_custom_config + ) + observed_node_names: set[str] = observed_graph_module_attrs.observed_node_names + node_name_to_qconfig: dict[str, QConfigAny] = ( + observed_graph_module_attrs.node_name_to_qconfig + ) # type: ignore[assignment] + + # mapping from fully qualified module name to module instance + # for example, + # { + # '': Model(...), + # 'linear': Linear(...), + # 'linear.weight_fake_quant': PerChannelMinMaxObserver(...), + # } + # We use remove_duplicate=False here because torch.cat uses + # the same activation_post_process module instance but different names + modules = dict(model.named_modules(remove_duplicate=False)) + + # TODO refactor this code once we update the prepare logic to have additional information on + # which graph nodes have been observed and share that with convert to decide which observers to ignore. + if qconfig_mapping: + prepare_qconfig_mapping: QConfigMapping = ( + observed_graph_module_attrs.qconfig_mapping + ) # type: ignore[assignment] + modules_copy = copy.deepcopy(modules) + + if observed_graph_module_attrs.is_qat: + _update_qconfig_for_qat(qconfig_mapping, backend_config) + _update_qconfig_for_fusion(model, qconfig_mapping) + + _compare_prepare_convert_qconfig_mappings( + prepare_qconfig_mapping, qconfig_mapping + ) # type: ignore[arg-type] + convert_node_name_to_qconfig = _generate_node_name_to_qconfig( + model, modules_copy, model.graph, qconfig_mapping, node_name_to_scope + ) + # check the convert_node_name_to_qconfig generated and ensure that + # all the values either match what was set in prepare node_name_to_qconfig + # or are set to None in the convert_node_name_to_qconfig. + for k, v in node_name_to_qconfig.items(): + assert k in convert_node_name_to_qconfig, ( + f"Expected key {k} in convert node_name_to_qconfig" + ) + if convert_node_name_to_qconfig[k] is not None: + assert qconfig_equals(v, convert_node_name_to_qconfig[k]), ( + f"Expected k {k} to have the same value in prepare and convert QConfigMappings, " + f"but {v} was updated to {convert_node_name_to_qconfig[k]}" + ) + node_name_to_qconfig = convert_node_name_to_qconfig + + if observed_graph_module_attrs.equalization_node_name_to_qconfig is not None: + # If we want to do equalization then do the following: + # Calculate the equalization scale, update the observers with the scaled + # inputs, and scale the weight + weight_eq_obs_dict = update_obs_for_equalization(model, modules) + convert_eq_obs(model, modules, weight_eq_obs_dict) + + # always run weight observers in the top level forward method + # for dynamic quant ops or weight only quant ops + _run_weight_observers(model, backend_config) + + # additional state to override inputs to be quantized, if specified + # by the user + placeholder_node_seen_cnt = 0 + input_quantized_idxs: list[int] = prepare_custom_config.input_quantized_indexes + output_quantized_idxs: list[int] = prepare_custom_config.output_quantized_indexes + + root_module_to_quantized_reference_module = ( + get_root_module_to_quantized_reference_module(backend_config) + ) + # convert tuples so that it can work with isinstance(module, tuple_of_classes) + root_module_classes = tuple(root_module_to_quantized_reference_module.keys()) + qat_module_classes = get_qat_module_classes(backend_config) + fused_module_classes = get_fused_module_classes(backend_config) + + for node in list(model.graph.nodes): + if node.op == "placeholder": + cur_placeholder_node_idx = placeholder_node_seen_cnt + placeholder_node_seen_cnt += 1 + if cur_placeholder_node_idx in input_quantized_idxs: + # Inputs are assumed to be quantized if the user specified the + # input_quantized_idxs override. + # we need to dequantize the inputs since all operators took + # floating point inputs in reference quantized models + _insert_dequantize_node(node, model.graph) + elif node.op == "output": + # If the argument is empty we don't need to do anything + if len(output_quantized_idxs) == 0: + continue + # Result are kept quantized if the user specified the + # output_quantized_idxs override. + # Remove the dequantize operator for the node in the end if any + return_node = node + output = node.args[0] + # outputs can be Node, list, tuple, dict, other cases are not supported yet + if isinstance(output, (list, tuple)): + for idx in output_quantized_idxs: + _maybe_recursive_remove_dequantize( + output[idx], return_node, model.graph + ) + elif isinstance(output, (Node, dict)): + # we treat dict as a single argument currently, but it can be extended + # to support {"key": dtype} after we change output_quantized_idxs to + # dict + if 0 in output_quantized_idxs: + _maybe_recursive_remove_dequantize(output, return_node, model.graph) + else: + warnings.warn( + f"Unsupported node type for output_quantized_idxs: {type(output)}" + ) + elif node.op == "call_module": + mod = _get_module(node, modules) + assert mod is not None + if _is_activation_post_process(mod): + if is_decomposed: + _replace_observer_with_quantize_dequantize_node_decomposed( + model, + node, + modules, + node_name_to_scope, + node_name_to_qconfig, + ) + else: + _replace_observer_with_quantize_dequantize_node( + model, + node, + modules, + node_name_to_scope, + node_name_to_qconfig, + ) + elif isinstance(mod, DeQuantStub): + _replace_observer_or_dequant_stub_with_dequantize_node( + node, model.graph + ) + elif _is_observed_standalone_module(mod): + convert_standalone_module( + node, modules, model, is_reference, backend_config + ) + # below this point `type_before_parametrizations` is used + # instead of `type` to handle situations with fx quant + sparsity + elif type_before_parametrizations(mod) in set(root_module_classes).union( + qat_module_classes + ).union(fused_module_classes): + # extra check for fused module classes to make sure they are fused module classes + # of target modules + if ( + type_before_parametrizations(mod) in fused_module_classes + and type_before_parametrizations(mod[0]) not in root_module_classes + ): # type: ignore[index] + continue + convert_weighted_module( + node, + modules, + observed_node_names, + node_name_to_qconfig, + backend_config, + is_decomposed, + is_reference, + ) + + # remove deadcode after converting observers to quant/dequant ops + model.graph.eliminate_dead_code() + model = GraphModule(model, model.graph) + + # TODO: maybe move this to quantize_fx.py + if not is_reference: + model = lower_to_fbgemm( + model, node_name_to_qconfig, node_name_to_scope, keep_original_weights + ) + + # TODO: this looks hacky, we want to check why we need this and see if we can + # remove this + # removes qconfig and activation_post_process modules + if _remove_qconfig_flag: + _remove_qconfig(model) + model.delete_all_unused_submodules() + model.meta.pop("_observed_graph_module_attrs", None) + return model + + +def _convert_fx( + graph_module: GraphModule, + is_reference: bool, + convert_custom_config: Union[ConvertCustomConfig, dict[str, Any], None] = None, + is_standalone_module: bool = False, + _remove_qconfig: bool = True, + qconfig_mapping: Union[QConfigMapping, dict[str, Any], None] = None, + backend_config: Union[BackendConfig, dict[str, Any], None] = None, + is_decomposed: bool = False, + keep_original_weights: bool = False, +) -> GraphModule: + """`is_standalone_module`: see docs in :func:`~torch.ao.quantization.prepare_standalone_module_fx`""" + if convert_custom_config is None: + convert_custom_config = ConvertCustomConfig() + + if isinstance(convert_custom_config, dict): + warnings.warn( + "Passing a convert_custom_config_dict to convert is deprecated and will not be supported " + "in a future version. Please pass in a ConvertCustomConfig instead.", + FutureWarning, + stacklevel=3, + ) + convert_custom_config = ConvertCustomConfig.from_dict(convert_custom_config) + + _check_is_graph_module(graph_module) + preserved_attr_names = convert_custom_config.preserved_attributes + preserved_attrs = { + attr: getattr(graph_module, attr) + for attr in preserved_attr_names + if hasattr(graph_module, attr) + } + + quantized = convert( + graph_module, + is_reference, + convert_custom_config, + is_standalone_module, + _remove_qconfig_flag=_remove_qconfig, + qconfig_mapping=qconfig_mapping, + backend_config=backend_config, + is_decomposed=is_decomposed, + keep_original_weights=keep_original_weights, + ) + + attach_preserved_attrs_to_model(quantized, preserved_attrs) + return quantized + + +def _convert_to_reference_decomposed_fx( + graph_module: GraphModule, + convert_custom_config: Union[ConvertCustomConfig, dict[str, Any], None] = None, + qconfig_mapping: Union[QConfigMapping, dict[str, Any], None] = None, + backend_config: Union[BackendConfig, dict[str, Any], None] = None, +) -> GraphModule: + r"""Convert a calibrated or trained model to a reference quantized model, with + decomposed representation for quantized Tensor + see https://github.com/pytorch/rfcs/blob/master/RFC-0019-Extending-PyTorch-Quantization-to-Custom-Backends.md for more details, + reference quantized model is a standard representation of a quantized model provided + by FX Graph Mode Quantization, it can be further lowered to run on the target + hardware, like accelerators + + Note: this is not public API + + Args: + * `graph_module` (GraphModule): A prepared and calibrated/trained model (GraphModule) + + * `convert_custom_config` (ConvertCustomConfig): custom configurations for convert function. + See :func:`~torch.ao.quantization.quantize_fx.convert_fx` for more details. + + * `_remove_qconfig` (bool): Option to remove the qconfig attributes in the model after convert. + + * `qconfig_mapping` (QConfigMapping): config for specifying how to convert a model for quantization. + See :func:`~torch.ao.quantization.quantize_fx.convert_fx` for more details. + + * `backend_config` (BackendConfig): A configuration for the backend which describes how + operators should be quantized in the backend. See + :func:`~torch.ao.quantization.quantize_fx.convert_fx` for more details. + + Return: + A reference quantized model (GraphModule) with operators working with decomposed quantized Tensor + + Example:: + + # prepared_model: the model after prepare_fx/prepare_qat_fx and calibration/training + # TODO: add backend_config after we split the backend_config for fbgemm and qnnpack + # e.g. backend_config = get_default_backend_config("fbgemm") + reference_quantized_model = _convert_to_reference_decomposed_fx(prepared_model) + + """ + torch._C._log_api_usage_once( + "quantization_api.quantize_fx._convert_to_reference_decomposed_fx" + ) + return _convert_fx( + graph_module, + is_reference=True, + convert_custom_config=convert_custom_config, + _remove_qconfig=False, + qconfig_mapping=qconfig_mapping, + backend_config=backend_config, + is_decomposed=True, + ) diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/export_utils.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/export_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..909fc92918ca460bd7d0073015265f4276fd0c8e --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/export_utils.py @@ -0,0 +1,246 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import types + +import torch +import torch.nn.functional as F + +from torchao.utils import _assert_and_get_unique_device + +__all__ = [ + "model_is_exported", +] + +_EXPORTED_TRAINING_ATTR = "_exported_training" + + +class WrapperModule(torch.nn.Module): + """Class to wrap a callable in an :class:`torch.nn.Module`. Use this if you + are trying to export a callable. + """ + + def __init__(self, fn): + super().__init__() + self.fn = fn + + def forward(self, *args, **kwargs): + """Simple forward that just calls the ``fn`` provided to :meth:`WrapperModule.__init__`.""" + return self.fn(*args, **kwargs) + + +def model_is_exported(m: torch.nn.Module) -> bool: + """ + Return True if the `torch.nn.Module` was exported, False otherwise + (e.g. if the model was FX symbolically traced or not traced at all). + """ + return isinstance(m, torch.fx.GraphModule) and any( + "val" in n.meta for n in m.graph.nodes + ) + + +def _replace_dropout(m: torch.fx.GraphModule, train_to_eval: bool): + """ + Switch dropout patterns in the model between train and eval modes. + + Dropout has different behavior in train vs eval mode. For exported models, + however, calling `model.train()` or `model.eval()` does not automatically switch + the dropout behavior between the two modes, so here we need to rewrite the aten + dropout patterns manually to achieve the same effect. + + See https://github.com/pytorch/pytorch/issues/103681. + """ + # Avoid circular dependencies + from .utils import _get_aten_graph_module_for_pattern + + # Needed to ensure subgraph matches are self-contained + m.graph.eliminate_dead_code() + m.recompile() + + for inplace in [False, True]: + + def dropout_train(x): + return F.dropout(x, p=0.5, training=True, inplace=inplace) + + def dropout_eval(x): + return F.dropout(x, p=0.5, training=False, inplace=inplace) + + example_inputs = (torch.randn(1),) + if train_to_eval: + match_pattern = _get_aten_graph_module_for_pattern( + WrapperModule(dropout_train), + example_inputs, + ) + replacement_pattern = _get_aten_graph_module_for_pattern( + WrapperModule(dropout_eval), + example_inputs, + ) + else: + match_pattern = _get_aten_graph_module_for_pattern( + WrapperModule(dropout_eval), + example_inputs, + ) + replacement_pattern = _get_aten_graph_module_for_pattern( + WrapperModule(dropout_train), + example_inputs, + ) + + from torch.fx.subgraph_rewriter import replace_pattern_with_filters + + replace_pattern_with_filters( + m, + match_pattern, + replacement_pattern, + match_filters=[], + ignore_literals=True, + ) + m.recompile() + + +def _replace_batchnorm(m: torch.fx.GraphModule, train_to_eval: bool): + """ + Switch batchnorm patterns in the model between train and eval modes. + + Batchnorm has different behavior in train vs eval mode. For exported models, + however, calling `model.train()` or `model.eval()` does not automatically switch + the batchnorm behavior between the two modes, so here we need to rewrite the aten + batchnorm patterns manually to achieve the same effect. + """ + # TODO(Leslie): This function still fails to support custom momentum and eps value. + # Enable this support in future updates. + + # Avoid circular dependencies + from .utils import _get_aten_graph_module_for_pattern + + # Needed to ensure subgraph matches are self-contained + m.graph.eliminate_dead_code() + m.recompile() + + def bn_train( + x: torch.Tensor, + bn_weight: torch.Tensor, + bn_bias: torch.Tensor, + bn_running_mean: torch.Tensor, + bn_running_var: torch.Tensor, + ): + return F.batch_norm( + x, bn_running_mean, bn_running_var, bn_weight, bn_bias, training=True + ) + + def bn_eval( + x: torch.Tensor, + bn_weight: torch.Tensor, + bn_bias: torch.Tensor, + bn_running_mean: torch.Tensor, + bn_running_var: torch.Tensor, + ): + return F.batch_norm( + x, bn_running_mean, bn_running_var, bn_weight, bn_bias, training=False + ) + + example_inputs = ( + torch.randn(1, 1, 3, 3), # x + torch.randn(1), # bn_weight + torch.randn(1), # bn_bias + torch.randn(1), # bn_running_mean + torch.randn(1), # bn_running_var + ) + + device = _assert_and_get_unique_device(m) + is_cuda = device is not None and device.type == "cuda" + bn_train_aten = _get_aten_graph_module_for_pattern( + WrapperModule(bn_train), + example_inputs, + is_cuda, + ) + bn_eval_aten = _get_aten_graph_module_for_pattern( + WrapperModule(bn_eval), + example_inputs, + is_cuda, + ) + + if train_to_eval: + match_pattern = bn_train_aten + replacement_pattern = bn_eval_aten + else: + match_pattern = bn_eval_aten + replacement_pattern = bn_train_aten + + from torch.fx.subgraph_rewriter import replace_pattern_with_filters + + replace_pattern_with_filters( + m, + match_pattern, + replacement_pattern, + match_filters=[], + ignore_literals=True, + ) + m.recompile() + + +# TODO: expose these under this namespace? +def _move_exported_model_to_eval(model: torch.fx.GraphModule): + """ + Move an exported GraphModule to eval mode. + + This is equivalent to model.eval() but only for certain special ops like dropout, batchnorm. + QAT users should call this before performing inference on the model. + + This call is idempotent; if the model is already in eval mode, nothing will happen. + """ + is_training = getattr(model, _EXPORTED_TRAINING_ATTR, True) + if not is_training: + return model + setattr(model, _EXPORTED_TRAINING_ATTR, False) + _replace_dropout(model, train_to_eval=True) + _replace_batchnorm(model, train_to_eval=True) + return model + + +def _move_exported_model_to_train(model: torch.fx.GraphModule): + """ + Move an exported GraphModule to train mode. + + This is equivalent to model.train() but only for certain special ops like dropout, batchnorm. + QAT users should call this before performing training on the model. + + This call is idempotent; if the model is already in train mode, nothing will happen. + """ + is_training = getattr(model, _EXPORTED_TRAINING_ATTR, False) + if is_training: + return model + setattr(model, _EXPORTED_TRAINING_ATTR, True) + _replace_dropout(model, train_to_eval=False) + _replace_batchnorm(model, train_to_eval=False) + return model + + +def _allow_exported_model_train_eval(model: torch.fx.GraphModule): + """ + Allow users to call `model.train()` and `model.eval()` on an exported model, + but with the effect of changing behavior between the two modes limited to special + ops only, which are currently dropout and batchnorm. + + Note: This does not achieve the same effect as what `model.train()` and `model.eval()` + does in eager models, but only provides an approximation. In particular, user code + branching on `training` flag will not function correctly in general because the branch + is already specialized at export time. Additionally, other ops beyond dropout and batchnorm + that have different train/eval behavior will also not be converted properly. + """ + + def _train(self, mode: bool = True): + if mode: + _move_exported_model_to_train(self) + else: + _move_exported_model_to_eval(self) + + def _eval(self): + _move_exported_model_to_eval(self) + + model.train = types.MethodType(_train, model) # type: ignore[method-assign] + model.eval = types.MethodType(_eval, model) # type: ignore[method-assign] + return model diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/fake_quantize.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/fake_quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..fd20523c6a8a757ea27a8ebcd156de8e2f0b979f --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/fake_quantize.py @@ -0,0 +1,656 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +"""Implements modules used to perform fake quantization.""" + +import re +from abc import ABC, abstractmethod +from typing import Any + +import torch +from torch.nn import Module + +from torchao.quantization.pt2e.observer import ( + FixedQParamsObserver, + HistogramObserver, + MovingAverageMinMaxObserver, + MovingAveragePerChannelMinMaxObserver, + _with_args, + default_fixed_qparams_range_0to1_observer, + default_fixed_qparams_range_neg1to1_observer, +) + +__all__ = [ + "FakeQuantizeBase", + "FakeQuantize", + "FixedQParamsFakeQuantize", + "FusedMovingAvgObsFakeQuantize", + "disable_fake_quant", + "disable_observer", + "enable_fake_quant", + "enable_observer", + "default_fake_quant", + "default_weight_fake_quant", + "default_dynamic_fake_quant", + "default_fixed_qparams_range_neg1to1_fake_quant", + "default_fixed_qparams_range_0to1_fake_quant", + "default_symmetric_fixed_qparams_fake_quant", + "default_affine_fixed_qparams_fake_quant", + "default_per_channel_weight_fake_quant", + "default_embedding_fake_quant", + "default_embedding_fake_quant_4bit", + "default_histogram_fake_quant", + "default_fused_act_fake_quant", + "default_fused_wt_fake_quant", + "default_fused_per_channel_wt_fake_quant", + "fused_wt_fake_quant_range_neg_127_to_127", + "fused_per_channel_wt_fake_quant_range_neg_127_to_127", +] + + +def _is_per_channel(qscheme: "torch.qscheme") -> bool: + return qscheme in [ + torch.per_channel_symmetric, + torch.per_channel_affine, + torch.per_channel_affine_float_qparams, + ] + + +def _is_per_tensor(qscheme: "torch.qscheme") -> bool: + return qscheme in [torch.per_tensor_symmetric, torch.per_tensor_affine] + + +def _is_symmetric_quant(qscheme: "torch.qscheme") -> bool: + return qscheme in [torch.per_tensor_symmetric, torch.per_channel_symmetric] + + +def _is_float_qparams(qscheme: "torch.qscheme") -> bool: + return qscheme in [ + torch.per_channel_affine_float_qparams, + ] + + +class FakeQuantizeBase(ABC, Module): + r"""Base fake quantize module. + + Base fake quantize module + Any fake quantize implementation should derive from this class. + + Concrete fake quantize module should follow the same API. In forward, they will update + the statistics of the observed Tensor and fake quantize the input. They should also provide a + `calculate_qparams` function that computes the quantization parameters given + the collected statistics. + + """ + + fake_quant_enabled: torch.Tensor + observer_enabled: torch.Tensor + + def __init__(self) -> None: + """Set fake_quant_enabled and observer_enabled.""" + super().__init__() + # fake_quant_enabled and observer_enabled are buffers to support their + # replication in DDP. Data type is uint8 because NCCL does not support + # bool tensors. + self.register_buffer("fake_quant_enabled", torch.tensor([1], dtype=torch.uint8)) + self.register_buffer("observer_enabled", torch.tensor([1], dtype=torch.uint8)) + + @abstractmethod + def forward(self, x): + pass + + @abstractmethod + def calculate_qparams(self, **kwargs): + pass + + @torch.jit.export + def enable_fake_quant(self, enabled: bool = True) -> None: + self.fake_quant_enabled[0] = 1 if enabled else 0 + + @torch.jit.export + def disable_fake_quant(self): + self.enable_fake_quant(False) + + @torch.jit.export + def enable_observer(self, enabled: bool = True) -> None: + self.observer_enabled[0] = 1 if enabled else 0 + + @torch.jit.export + def disable_observer(self): + self.enable_observer(False) + + @classmethod + def with_args(cls, **kwargs): + fake_quant_constructor = _with_args(cls, **kwargs) + # need to assign the correct module to fake_quantize + # constructors to satisfy public v private requirements + fake_quant_constructor.__module__ = "torchao.quantization.pt2e.fake_quantize" + return fake_quant_constructor + + +class FakeQuantize(FakeQuantizeBase): + r"""Simulate the quantize and dequantize operations in training time. + + The output of this module is given by:: + + x_out = ( + clamp(round(x/scale + zero_point), quant_min, quant_max) - zero_point + ) * scale + + * :attr:`is_dynamic` indicates whether the fake quantie is a placeholder for dynamic quantization + operators (choose_qparams -> q -> dq) or static quantization operators (q -> dq) + + * :attr:`scale` defines the scale factor used for quantization. + + * :attr:`zero_point` specifies the quantized value to which 0 in floating point maps to + + * :attr:`fake_quant_enabled` controls the application of fake quantization on tensors, note that + statistics can still be updated. + + * :attr:`observer_enabled` controls statistics collection on tensors + + * :attr:`dtype` specifies the quantized dtype that is being emulated with fake-quantization, + allowable values are torch.qint8 and torch.quint8. + + Args: + + observer (module): Module for observing statistics on input tensors and calculating scale + and zero-point. + observer_kwargs (optional): Arguments for the observer module + + Attributes: + activation_post_process (Module): User provided module that collects statistics on the input tensor and + provides a method to calculate scale and zero-point. + + """ + + scale: torch.Tensor + zero_point: torch.Tensor + + def __init__( + self, + observer=MovingAverageMinMaxObserver, + quant_min=None, + quant_max=None, + is_dynamic=False, + **observer_kwargs, + ): + super().__init__() + # Populate quant_min/quant_max to observer_kwargs if valid + if quant_min is not None and quant_max is not None: + assert quant_min <= quant_max, ( + "quant_min must be less than or equal to quant_max" + ) + dtype = observer_kwargs.get("dtype", torch.quint8) + if hasattr(observer, "p"): + # In case observer is _PartialWrapper, dtype can be stored in + # observer.p.keywords["dtype"] + dtype = getattr(getattr(observer, "p", {}), "keywords", {}).get( + "dtype", dtype + ) + assert torch.iinfo(dtype).min <= quant_min, "quant_min out of bound" + assert quant_max <= torch.iinfo(dtype).max, "quant_max out of bound" + observer_kwargs.update({"quant_min": quant_min, "quant_max": quant_max}) + observer_kwargs["is_dynamic"] = is_dynamic + self.activation_post_process = observer(**observer_kwargs) + # TODO: keeping self.quant_min/max for BC; remove after a couple releases + # Users should use self.activation_post_process.quant_min + self.quant_min = self.activation_post_process.quant_min + self.quant_max = self.activation_post_process.quant_max + self.is_dynamic = self.activation_post_process.is_dynamic + if _is_float_qparams(self.activation_post_process.qscheme): + zero_point_dtype = torch.float + else: + zero_point_dtype = torch.int + self.register_buffer("scale", torch.tensor([1.0], dtype=torch.float)) + self.register_buffer("zero_point", torch.tensor([0], dtype=zero_point_dtype)) + self.dtype = self.activation_post_process.dtype + self.qscheme = self.activation_post_process.qscheme + self.ch_axis = ( + self.activation_post_process.ch_axis + if hasattr(self.activation_post_process, "ch_axis") + else -1 + ) + assert _is_per_channel(self.qscheme) or _is_per_tensor(self.qscheme), ( + "Only per channel and per tensor quantization are supported in fake quantize" + + " got qscheme: " + + str(self.qscheme) + ) + self.is_per_channel = _is_per_channel(self.qscheme) + + @torch.jit.export + def calculate_qparams(self): + return self.activation_post_process.calculate_qparams() + + def forward(self, X): + if self.observer_enabled[0] == 1: + self.activation_post_process(X.detach()) + _scale, _zero_point = self.calculate_qparams() + _scale, _zero_point = ( + _scale.to(self.scale.device), + _zero_point.to(self.zero_point.device), + ) + if self.scale.shape != _scale.shape: + self.scale.resize_(_scale.shape) + self.zero_point.resize_(_zero_point.shape) + self.scale.copy_(_scale) + self.zero_point.copy_(_zero_point) + + if self.fake_quant_enabled[0] == 1: + if self.is_per_channel: + X = torch.fake_quantize_per_channel_affine( + X, + self.scale, + self.zero_point, + self.ch_axis, + self.activation_post_process.quant_min, + self.activation_post_process.quant_max, + ) + else: + X = torch.fake_quantize_per_tensor_affine( + X, + self.scale, + self.zero_point, + self.activation_post_process.quant_min, + self.activation_post_process.quant_max, + ) + return X + + @torch.jit.export + def extra_repr(self): + return ( + f"fake_quant_enabled={self.fake_quant_enabled}, observer_enabled={self.observer_enabled}, " + f"quant_min={self.activation_post_process.quant_min}, quant_max={self.activation_post_process.quant_max}, " + f"dtype={self.dtype}, qscheme={self.qscheme}, ch_axis={self.ch_axis}, " + f"scale={self.scale}, zero_point={self.zero_point}" + ) + + def _save_to_state_dict(self, destination, prefix, keep_vars): + # We cannot currently register scalar values as buffers, so need to manually + # specify serialization here. + super()._save_to_state_dict(destination, prefix, keep_vars) + destination[prefix + "scale"] = self.scale + destination[prefix + "zero_point"] = self.zero_point + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + # Removing this function throws an error that the size of the loaded tensor does not match the original size + # i.e., These buffers start out with numel 0 and become numel 1 once they have their first forward pass. + local_state = ["scale", "zero_point"] + for name in local_state: + key = prefix + name + if key in state_dict: + val = state_dict[key] + # Custom handling to allow loading scale and zero_point + # of size N into uninitialized buffers of size 0. The + # buffers are resized here, and the values are copied in + # the default state_dict loading code of the parent. + if name == "scale": + self.scale.resize_(val.shape) + else: + assert name == "zero_point" + self.zero_point.resize_(val.shape) + # For torchscript module we need to update the attributes here since we do not + # call the `_load_from_state_dict` function defined module.py + if torch.jit.is_scripting(): + if name == "scale": + self.scale.copy_(val) + else: + assert name == "zero_point" + self.zero_point.copy_(val) + elif strict: + missing_keys.append(key) + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + +class FixedQParamsFakeQuantize(FakeQuantize): + """Simulate quantize and dequantize in training time. + + Simulate quantize and dequantize with fixed quantization + parameters in training time. Only per tensor quantization + is supported. + """ + + # TODO: rename observer to observer_ctr + def __init__(self, observer): + super().__init__(observer=observer) + assert type(self.activation_post_process) == FixedQParamsObserver, ( + f"{self.__class__.__name__}'s observer must be a {FixedQParamsObserver.__name__}" + ) + self._observer_ctr = observer + self.scale = self.activation_post_process.scale + self.zero_point = self.activation_post_process.zero_point + assert _is_per_tensor(self.qscheme), ( + "Only per tensor quantization is supported" + + " FixedQParamsFakeQuantize module, got qscheme:" + + str(self.qscheme) + ) + + @torch.jit.export + def calculate_qparams(self): + return self.scale, self.zero_point + + @torch.jit.export + def extra_repr(self): + """Define a string representation of the object's attributes.""" + return ( + f"fake_quant_enabled={self.fake_quant_enabled}, observer_enabled={self.observer_enabled}, " + f"scale={self.scale}, zero_point={self.zero_point}, " + f"dtype={self.dtype}, quant_min={self.activation_post_process.quant_min}, " + f"quant_max={self.activation_post_process.quant_max}, qscheme={self.qscheme}" + ) + + +class FusedMovingAvgObsFakeQuantize(FakeQuantize): + r"""Define a fused module to observe the tensor. + + Fused module that is used to observe the input tensor (compute min/max), compute + scale/zero_point and fake_quantize the tensor. + This module uses calculation similar MovingAverageMinMaxObserver for the inputs, + to compute the min/max values in order to compute the scale/zero_point. + The qscheme input in the observer is used to differentiate between symmetric/affine + quantization scheme. + + The output of this module is given by + x_out = (clamp(round(x/scale + zero_point), quant_min, quant_max)-zero_point)*scale + + Similar to :class:`~torchao.quantization.pt2e.FakeQuantize`, and accepts the same attributes as the + base class. + + """ + + def __init__( + self, + observer: Any = MovingAverageMinMaxObserver, + quant_min: int = 0, + quant_max: int = 255, + **observer_kwargs: Any, + ) -> None: + super().__init__(observer, quant_min, quant_max, **observer_kwargs) + assert isinstance( + self.activation_post_process, + (MovingAverageMinMaxObserver, MovingAveragePerChannelMinMaxObserver), + ), ( + "Fused observer+fake_quant module only works with MovingAverageMinMaxObserver" + ) + self.register_buffer("fake_quant_enabled", torch.tensor([1], dtype=torch.long)) + self.register_buffer("observer_enabled", torch.tensor([1], dtype=torch.long)) + self.is_symmetric_quant = _is_symmetric_quant( + self.activation_post_process.qscheme + ) + + @torch.jit.export + def calculate_qparams(self) -> tuple[torch.Tensor, torch.Tensor]: + return self.activation_post_process.calculate_qparams() + + @torch.jit.export + def extra_repr(self) -> str: + return ( + f"fake_quant_enabled={self.fake_quant_enabled}, observer_enabled={self.observer_enabled}, " + f"scale={self.scale}, zero_point={self.zero_point}, dtype={self.dtype}, " + f"quant_min={self.activation_post_process.quant_min}, quant_max={self.activation_post_process.quant_max}, " + f"qscheme={self.qscheme}, reduce_range={self.activation_post_process.reduce_range}" + ) + + def forward(self, X: torch.Tensor) -> torch.Tensor: + return torch.fused_moving_avg_obs_fake_quant( + X, + self.observer_enabled, + self.fake_quant_enabled, + self.activation_post_process.min_val, + self.activation_post_process.max_val, + self.scale, + self.zero_point, + self.activation_post_process.averaging_constant, + self.activation_post_process.quant_min, + self.activation_post_process.quant_max, + self.ch_axis, + self.is_per_channel, + self.is_symmetric_quant, + ) + + +default_fake_quant = FakeQuantize.with_args( + observer=MovingAverageMinMaxObserver, + quant_min=0, + quant_max=255, + dtype=torch.quint8, + qscheme=torch.per_tensor_affine, + reduce_range=True, +) +""" +Default fake_quant for activations. +""" + +default_weight_fake_quant = FakeQuantize.with_args( + observer=MovingAverageMinMaxObserver, + quant_min=-128, + quant_max=127, + dtype=torch.qint8, + qscheme=torch.per_tensor_symmetric, + reduce_range=False, +) +""" +Default fake_quant for weights. +Observer is memoryless since averaging_constant is 1. +""" + +default_dynamic_fake_quant = FakeQuantize.with_args( + observer=MovingAverageMinMaxObserver, + quant_min=0, + quant_max=255, + is_dynamic=True, + dtype=torch.quint8, + averaging_constant=1, +) +""" +Default dynamic fake_quant for activations. +""" + +default_fixed_qparams_range_neg1to1_fake_quant = FixedQParamsFakeQuantize.with_args( + observer=default_fixed_qparams_range_neg1to1_observer +) +default_fixed_qparams_range_0to1_fake_quant = FixedQParamsFakeQuantize.with_args( + observer=default_fixed_qparams_range_0to1_observer +) +# TODO: the following 2 variables are kept for backwards compatibility; remove after a few releases +default_symmetric_fixed_qparams_fake_quant = ( + default_fixed_qparams_range_neg1to1_fake_quant +) +default_affine_fixed_qparams_fake_quant = default_fixed_qparams_range_0to1_fake_quant + +default_per_channel_weight_fake_quant = FakeQuantize.with_args( + observer=MovingAveragePerChannelMinMaxObserver, + quant_min=-128, + quant_max=127, + dtype=torch.qint8, + qscheme=torch.per_channel_symmetric, + reduce_range=False, + ch_axis=0, +) +""" +Default fake_quant for per-channel weights. +Observer is memoryless since averaging_constant is 1. +""" +default_embedding_fake_quant = FakeQuantize.with_args( + observer=MovingAveragePerChannelMinMaxObserver, + qscheme=torch.per_channel_affine_float_qparams, + dtype=torch.quint8, + quant_min=0, + quant_max=255, + ch_axis=0, + averaging_constant=1, +) +""" +Default fake_quant for embeddings. +Observer is memoryless since averaging_constant is 1. +""" + +default_embedding_fake_quant_4bit = FakeQuantize.with_args( + observer=MovingAveragePerChannelMinMaxObserver, + qscheme=torch.per_channel_affine_float_qparams, + ch_axis=0, + dtype=torch.quint4x2, + averaging_constant=1, +) + +default_histogram_fake_quant = FakeQuantize.with_args( + observer=HistogramObserver, + quant_min=0, + quant_max=255, + dtype=torch.quint8, + qscheme=torch.per_tensor_affine, + reduce_range=True, +) +""" +Fake_quant for activations using a histogram.. +""" + + +default_fused_act_fake_quant = FusedMovingAvgObsFakeQuantize.with_args( + observer=MovingAverageMinMaxObserver, + quant_min=0, + quant_max=255, + dtype=torch.quint8, +) + +""" +Fused version of `default_fake_quant`, with improved performance. +""" + + +default_fused_wt_fake_quant = FusedMovingAvgObsFakeQuantize.with_args( + observer=MovingAverageMinMaxObserver, + quant_min=-128, + quant_max=127, + dtype=torch.qint8, + qscheme=torch.per_tensor_symmetric, +) +""" +Fused version of `default_weight_fake_quant`, with improved performance. +""" + +default_fused_per_channel_wt_fake_quant = FusedMovingAvgObsFakeQuantize.with_args( + observer=MovingAveragePerChannelMinMaxObserver, + quant_min=-128, + quant_max=127, + dtype=torch.qint8, + qscheme=torch.per_channel_symmetric, +) +""" +Fused version of `default_per_channel_weight_fake_quant`, with improved performance. +""" + +fused_wt_fake_quant_range_neg_127_to_127 = FusedMovingAvgObsFakeQuantize.with_args( + observer=MovingAverageMinMaxObserver, + quant_min=-127, + quant_max=127, + dtype=torch.qint8, + qscheme=torch.per_tensor_symmetric, + eps=2**-12, +) +""" +Fused version of `default_weight_fake_quant`, with the 8-bit values restricted to [-127, +127], excluding -128. +""" + +fused_per_channel_wt_fake_quant_range_neg_127_to_127 = ( + FusedMovingAvgObsFakeQuantize.with_args( + observer=MovingAveragePerChannelMinMaxObserver, + quant_min=-127, + quant_max=127, + dtype=torch.qint8, + qscheme=torch.per_channel_symmetric, + eps=2**-12, + ) +) + +""" +Fused version of `default_per_channel_weight_fake_quant`, with the 8-bit values restricted to [-127, +127], excluding -128. +""" + + +def _is_fake_quant_script_module(mod): + """Return true if given mod is an instance of FakeQuantize script module.""" + if isinstance(mod, torch.jit.RecursiveScriptModule): + # qualified name looks like '__torch__.torchao.quantization.pt2e.fake_quantize.___torch_mangle_2.FakeQuantize' + suffix = mod._c.qualified_name.split(".", 1)[1] + name = re.sub(r"\.___torch_mangle_\d+", "", suffix) + return ( + name == "torchao.quantization.pt2e.fake_quantize.FakeQuantize" + or name + == "torchao.quantization.pt2e.fake_quantize.FusedMovingAvgObsFakeQuantize" + ) + return False + + +def disable_fake_quant(mod): + """Disable fake quantization for the module. + + Disable fake quantization for this module, if applicable. Example usage:: + + # model is any PyTorch model + model.apply(torchao.quantization.pt2e.disable_fake_quant) + + """ + if isinstance(mod, FakeQuantizeBase) or _is_fake_quant_script_module(mod): + mod.disable_fake_quant() + + +def enable_fake_quant(mod): + """Enable fake quantization for the module. + + Enable fake quantization for this module, if applicable. Example usage:: + + # model is any PyTorch model + model.apply(torchao.quantization.pt2e.enable_fake_quant) + + """ + if isinstance(mod, FakeQuantizeBase) or _is_fake_quant_script_module(mod): + mod.enable_fake_quant() + + +def disable_observer(mod): + """Disable observation for this module. + + Disable observation for this module, if applicable. Example usage:: + + # model is any PyTorch model + model.apply(torchao.quantization.pt2e.disable_observer) + + """ + if isinstance(mod, FakeQuantizeBase) or _is_fake_quant_script_module(mod): + mod.disable_observer() + + +def enable_observer(mod): + """Enable observation for this module. + + Enable observation for this module, if applicable. Example usage:: + + # model is any PyTorch model + model.apply(torchao.quantization.pt2e.enable_observer) + + """ + if isinstance(mod, FakeQuantizeBase) or _is_fake_quant_script_module(mod): + mod.enable_observer() diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/graph_utils.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/graph_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..18d4a1d0439ed789e4485261a13c0161dba249ce --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/graph_utils.py @@ -0,0 +1,186 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import itertools +import operator +from collections import OrderedDict +from collections.abc import Sequence +from typing import Any, Callable, Optional, Union + +import torch +from torch.export import ExportedProgram +from torch.fx import Node +from torch.fx.passes.utils.source_matcher_utils import ( + SourcePartition, + check_subgraphs_connected, + get_source_partitions, +) + +__all__ = [ + "find_sequential_partitions", + "get_equivalent_types", + "update_equivalent_types_dict", + "bfs_trace_with_node_process", +] + +_EQUIVALENT_TYPES: list[set] = [ + {torch.nn.Conv1d, torch.nn.functional.conv1d}, + {torch.nn.Conv2d, torch.nn.functional.conv2d}, + {torch.nn.AdaptiveAvgPool2d, torch.nn.functional.adaptive_avg_pool2d}, + {torch.nn.ReLU, torch.nn.functional.relu, torch.nn.functional.relu_}, + {torch.nn.BatchNorm2d, torch.nn.functional.batch_norm}, + {torch.nn.Hardtanh, torch.nn.functional.hardtanh, torch.nn.functional.hardtanh_}, + {torch.add, operator.add, operator.iadd, "add", "add_"}, + {torch.mul, operator.mul, operator.imul, "mul", "mul_"}, +] + + +def _create_equivalent_types_dict(): + _DICT = {} + for values in _EQUIVALENT_TYPES: + for v in values: + _DICT[v] = list(values) + return _DICT + + +_EQUIVALENT_TYPES_DICT = _create_equivalent_types_dict() + + +def get_equivalent_types() -> list[set]: + return _EQUIVALENT_TYPES + + +def update_equivalent_types_dict(customized_equivalent_types=None): + """Help function for user who wants to customize the _EQUIVALENT_TYPES and _EQUIVALENT_TYPES_DICT. + When customized_equivalent_types passes in, + re-generate _EQUIVALENT_TYPES and _EQUIVALENT_TYPES_DICT. + """ + if customized_equivalent_types is None: + raise ValueError("customized_equivalent_types should not be None") + global _EQUIVALENT_TYPES + global _EQUIVALENT_TYPES_DICT + _EQUIVALENT_TYPES = customized_equivalent_types + _EQUIVALENT_TYPES_DICT = _create_equivalent_types_dict() + + +def _partitions_sequential(partitions: Sequence[SourcePartition]): + prev_partition = None + for partition in partitions: + if prev_partition is not None and not check_subgraphs_connected( + prev_partition, partition + ): + return False + prev_partition = partition + return True + + +def _get_matching_types(partition_type): + matching_types = [partition_type] + if partition_type in _EQUIVALENT_TYPES_DICT: + matching_types.extend(_EQUIVALENT_TYPES_DICT[partition_type]) + return matching_types + + +def _valid_type_sequence(partition_types: list[Any]): + partition_types_set = set() # type: ignore[var-annotated] + for partition_type in partition_types: + matching_types = _get_matching_types(partition_type) + matching_types_set = set(matching_types) + if len(partition_types_set & matching_types_set) > 0: + return False + partition_types_set |= matching_types_set + return True + + +def find_sequential_partitions( + gm: torch.fx.GraphModule, + partition_types: list[Any], + include_functional_equivalent=True, + filter_fn: Optional[Callable[[Node], bool]] = None, +): + if not _valid_type_sequence(partition_types): + raise ValueError( + f"Invalid partition types: {partition_types}. Each type in the sequence must be unique" + ) + + typed_partitions: OrderedDict[Any, list[SourcePartition]] = OrderedDict() + for partition_type in partition_types: + types_to_match = _get_matching_types(partition_type) + partitions = get_source_partitions(gm.graph, types_to_match, filter_fn) + typed_partitions[partition_type] = list( + itertools.chain.from_iterable(partitions.values()) + ) + + typed_partitions_list = list(typed_partitions.values()) + fusion_candidates = itertools.product(*typed_partitions_list) + fused_partitions = [ + candidate + for candidate in fusion_candidates + if _partitions_sequential(candidate) + ] + return fused_partitions + + +def _get_submodule( + graph_module: torch.fx.GraphModule, node: torch.fx.Node, arg_index: int +) -> tuple[str, torch.nn.Module, torch.fx.Node]: + submod_node = node.args[arg_index] + assert isinstance(submod_node, torch.fx.Node) + assert submod_node.op == "get_attr" + assert isinstance(submod_node.target, str) + submodule = graph_module.get_submodule(submod_node.target) + # pyre-ignore + return submod_node.target, submodule, node + + +def _get_control_flow_submodules( + graph_module: torch.fx.GraphModule, +) -> list[tuple[str, torch.nn.Module, torch.fx.Node]]: + """ + Returns a list of submodules used for control flow operations + (torch.ops.higher_order.cond/map) that are in the given toplevel graph (does not look + into submodules). Specifically, the returned value is a list containing a + tuple of (name of the submodule that's stored in the graph module, the + submodule itself, and the fx node that uses this submodule). + """ + control_flow_submodules = [] + for node in graph_module.graph.nodes: + if node.op != "call_function": + continue + + if node.target is torch.ops.higher_order.cond: + control_flow_submodules.append(_get_submodule(graph_module, node, 1)) + control_flow_submodules.append(_get_submodule(graph_module, node, 2)) + if node.target is torch.ops.higher_order.map_impl: + control_flow_submodules.append(_get_submodule(graph_module, node, 0)) + + return control_flow_submodules + + +def bfs_trace_with_node_process( + model: Union[ExportedProgram, torch.fx.GraphModule], node_op: Callable +) -> None: + """Traverse the graph module and apply node_op to each node.""" + + assert isinstance(model, (ExportedProgram, torch.fx.GraphModule)), ( + f"Expected GraphModule or ExportedProgram, got {type(model)}" + ) + gm = model.graph_module if isinstance(model, ExportedProgram) else model + queue = [gm] + while queue: + current_graph_module = queue.pop(0) + for node in current_graph_module.graph.nodes: + if node.op in ["output", "placeholder"]: + continue + + node_op(node) + + control_flow_submodules = [ + submodule + for _, submodule, _ in _get_control_flow_submodules(current_graph_module) + ] + queue.extend(control_flow_submodules) diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/lowering.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/lowering.py new file mode 100644 index 0000000000000000000000000000000000000000..5491623b66b6de6b7046b982d7857fbcf123d009 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/lowering.py @@ -0,0 +1,65 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch._inductor.constant_folding import constant_fold +from torch._inductor.fx_passes.freezing_patterns import freezing_passes + +__all__ = [ + "lower_pt2e_quantized_to_x86", +] + + +def lower_pt2e_quantized_to_x86( + model: torch.fx.GraphModule, + example_inputs: tuple[torch.Tensor, ...], +) -> torch.fx.GraphModule: + """Lower a PT2E-qantized model to x86 backend. + + Args: + * `model` (torch.fx.GraphModule): a model quantized by PT2E quantization flow. + * `example_inputs` (tuple[torch.Tensor, ...]): example inputs for the model. + + Return: + A GraphModule lowered to x86 backend. + """ + + def _post_autograd_decomp_table(): # type: ignore[no-untyped-def] + decomp_table = torch.export.default_decompositions() + + # if we are post-autograd, we shouldn't + # decomp prim ops. + for k in list(decomp_table.keys()): + if not torch._export.utils._is_cia_op(k): + del decomp_table[k] + + return decomp_table + + def _node_replace(m): # type: ignore[no-untyped-def] + # Replace aten.t(x) with aten.permute(x, [1, 0]) + aten = torch.ops.aten + g = m.graph + for node in g.nodes: + if node.target == aten.t.default: + with g.inserting_before(node): + x = node.args[0] + dims = [1, 0] + perm_node = g.call_function(aten.permute.default, args=(x, dims)) + node.replace_all_uses_with(perm_node) + g.erase_node(node) + + g.lint() + m.recompile() + + lowered_model = ( + torch.export.export_for_training(model, example_inputs, strict=True) + .run_decompositions(_post_autograd_decomp_table()) + .module() + ) + _node_replace(lowered_model) + freezing_passes(lowered_model, example_inputs) + constant_fold(lowered_model) + return lowered_model diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/observer.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/observer.py new file mode 100644 index 0000000000000000000000000000000000000000..f6534308d88747930c0f39b88e6574453971bc2e --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/observer.py @@ -0,0 +1,2146 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +# temporarily skip RUF for this file for now, we can re-enable +# after move the affine quantization related things to torchao +# noqa: RUF +""" +This module implements observers which are used to collect statistics about +the values observed during calibration (PTQ) or training (QAT). +""" + +import operator +import re +import warnings +from abc import ABCMeta, abstractmethod +from collections import OrderedDict +from functools import partial +from typing import Any, Optional + +import torch +import torch.nn as nn +from torch.fx import Node + +import torchao +from torchao.quantization.pt2e.utils import ( + calculate_qmin_qmax, + check_min_max_valid, + is_per_channel, + is_per_tensor, + validate_qmin_qmax, +) + +__all__ = [ + "default_affine_fixed_qparams_observer", + "default_debug_observer", + "default_dynamic_quant_observer", + "default_fixed_qparams_range_0to1_observer", + "default_fixed_qparams_range_neg1to1_observer", + "default_float_qparams_observer", + "default_float_qparams_observer_4bit", + "default_histogram_observer", + "default_observer", + "default_per_channel_weight_observer", + "default_placeholder_observer", + "default_reuse_input_observer", + "default_symmetric_fixed_qparams_observer", + "default_weight_observer", + "get_observer_state_dict", + "load_observer_state_dict", + "per_channel_weight_observer_range_neg_127_to_127", + "weight_observer_range_neg_127_to_127", + "FixedQParamsObserver", + "HistogramObserver", + "MinMaxObserver", + "MovingAverageMinMaxObserver", + "MovingAveragePerChannelMinMaxObserver", + "NoopObserver", + "ObserverBase", + "PerChannelMinMaxObserver", + "PlaceholderObserver", + "RecordingObserver", + "ReuseInputObserver", + "UniformQuantizationObserverBase", + "AffineQuantizedObserverBase", + "Granularity", + "MappingType", + "PerAxis", + "PerBlock", + "PerGroup", + "PerRow", + "PerTensor", + "PerToken", + "TorchAODType", + "ZeroPointDomain", + "get_block_size", +] + + +class PartialWrapper: + def __init__(self, p): + self.p = p + self.callable_args = {} + + def __call__(self, *args, **keywords): + # call each arg in callable_args and add them partial, then run with keywords + # skip if arg_name in keywords so its possible to overwrite + for arg_name in self.callable_args: + if arg_name not in keywords: + keywords = {**keywords, arg_name: self.callable_args[arg_name]()} + return self.p(*args, **keywords) + + def __repr__(self): + return self.p.__repr__() + self.callable_args.__repr__() + + def with_args(self, **kwargs): + return _with_args(self, **kwargs) + + def with_callable_args(self, **kwargs): + result = PartialWrapper(p=self.p) + result.callable_args = {**self.callable_args, **kwargs} + return result + + +def _with_args(cls_or_self, **kwargs): + r"""Wrapper that allows creation of class factories. + + This can be useful when there is a need to create classes with the same + constructor arguments, but different instances. Can be used in conjunction with + _callable_args + + Example:: + + >>> # xdoctest: +SKIP("Undefined vars") + >>> Foo.with_args = classmethod(_with_args) + >>> foo_builder = Foo.with_args(a=3, b=4).with_args(answer=42) + >>> foo_instance1 = foo_builder() + >>> foo_instance2 = foo_builder() + >>> id(foo_instance1) == id(foo_instance2) + False + """ + r = PartialWrapper(partial(cls_or_self, **kwargs)) + return r + + +def _with_callable_args(cls_or_self, **kwargs): + r"""Wrapper that allows creation of class factories args that need to be + called at construction time. + + This can be useful when there is a need to create classes with the same + constructor arguments, but different instances and those arguments should only + be calculated at construction time. Can be used in conjunction with _with_args + + Example:: + + >>> # xdoctest: +SKIP("Undefined vars") + >>> Foo.with_callable_args = classmethod(_with_callable_args) + >>> Foo.with_args = classmethod(_with_args) + >>> foo_builder = Foo.with_callable_args(cur_time=get_time_func).with_args(name="dan") + >>> foo_instance1 = foo_builder() + >>> # wait 50 + >>> foo_instance2 = foo_builder() + >>> id(foo_instance1.creation_time) == id(foo_instance2.creation_time) + False + """ + r = PartialWrapper(partial(cls_or_self)) + return r.with_callable_args(**kwargs) + + +ABC: Any = ABCMeta("ABC", (object,), {}) # compatible with Python 2 *and* 3: + + +class ObserverBase(ABC, nn.Module): + r"""Base observer Module. + Any observer implementation should derive from this class. + + Concrete observers should follow the same API. In forward, they will update + the statistics of the observed Tensor. And they should provide a + `calculate_qparams` function that computes the quantization parameters given + the collected statistics. + + Args: + dtype: dtype argument to the `quantize` node needed to implement the + reference model spec. + is_dynamic: indicator for whether the observer is a placeholder for dynamic quantization + or static quantization + """ + + def __init__(self, dtype, is_dynamic: bool = False): + super().__init__() + self.dtype = dtype + self.is_dynamic = is_dynamic + + @abstractmethod + def forward(self, x): + pass + + @abstractmethod + def calculate_qparams(self, **kwargs): + pass + + with_args = classmethod(_with_args) + with_callable_args = classmethod(_with_callable_args) + + +class UniformQuantizationObserverBase(ObserverBase): + r"""Common base for all observers using uniform quantization to calculate + scale and zero_point. + + Args: + dtype: dtype argument to the `quantize` node needed to implement the + reference model spec. + qscheme: Quantization scheme to be used. + reduce_range: Reduces the range of the quantized data type by 1 bit. + This is sometimes required to avoid instruction overflow. + quant_min: Minimum quantization value. If unspecified, it will follow the 8-bit setup. + quant_max: Maximum quantization value. If unspecified, it will follow the 8-bit setup. + eps: Epsilon value for float32, Defaults to `torch.finfo(torch.float32).eps`. + + .. warning:: + + :attr:`dtype` can only take ``torch.qint8`` or ``torch.quint8``. + or `torch.int8` or `torch.uint8` + + .. warning:: + + :attr:`qscheme` can only take one of the following options: + + - ``torch.per_tensor_affine`` + - ``torch.per_tensor_symmetric`` + - ``torch.per_channel_affine`` + - ``torch.per_channel_symmetric`` + """ + + # Note: the version is shared by all observer types + # + # Version 1/None + # self + # + # Version 2 (base class only, does not include child class buffers) + # self + # |--- eps : Tensor + # + # Version 3 + # for HistogramObserver only, changed the shape of uninitialized + # min_val and max_val buffers from torch.Size([0]) to torch.Size([]) + # for PerChannelObservers, changed the name of the buffers from min_vals + # to min_val and from max_vals to max_val. + _version = 3 + + eps: torch.Tensor + + def __init__( + self, + dtype=torch.quint8, + qscheme=torch.per_tensor_affine, + reduce_range=False, + quant_min=None, + quant_max=None, + factory_kwargs=None, + eps=torch.finfo(torch.float32).eps, + is_dynamic=False, + **kwargs, + ) -> None: + factory_kwargs = torch.nn.factory_kwargs(factory_kwargs) + super().__init__(dtype=dtype, is_dynamic=is_dynamic, **kwargs) + self.qscheme = qscheme + if reduce_range: + warnings.warn( + "Please use quant_min and quant_max to specify the range for observers. \ + reduce_range will be deprecated in a future release of PyTorch." + ) + self.reduce_range = reduce_range + self.register_buffer("eps", torch.tensor([eps], **factory_kwargs)) + assert self.qscheme in ( + torch.per_tensor_affine, + torch.per_tensor_symmetric, + torch.per_channel_affine, + torch.per_channel_symmetric, + torch.per_channel_affine_float_qparams, + ), ( + "Default Observer only works for per_tensor_affine, \ + per_tensor_symmetric, per_channel_affine, \ + per_channel_symmetric and per_channel_float_qparams quantization scheme" + ) + + _ALLOWED_DTYPES = ( + torch.qint8, + torch.quint8, + torch.quint4x2, + torch.qint32, + torch.int8, + torch.uint8, + torch.int16, + torch.int32, + torch.float8_e5m2, + torch.float8_e4m3fn, + torch.uint16, + ) + + assert self.dtype in _ALLOWED_DTYPES, ( + f"Default Observer only works for {_ALLOWED_DTYPES} data type" + ) + self.has_customized_qrange = (quant_min is not None) and (quant_max is not None) + if self.has_customized_qrange: + validate_qmin_qmax(quant_min, quant_max) + self.quant_min, self.quant_max = calculate_qmin_qmax( + quant_min, + quant_max, + self.has_customized_qrange, + self.dtype, + self.reduce_range, + ) + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if version is None or version == 1: + # eps was moved to a buffer in version 2 + eps = torch.tensor([torch.finfo(torch.float32).eps]) + state_dict[prefix + "eps"] = eps + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + @torch.jit.export + def _validate_qmin_qmax(self, quant_min: int, quant_max: int) -> None: + r"""Validates that the user-specified quantization range is properly initialized + and within the given bound supported by the observer dtype. + + To accommodate lower-bit quantization with respect to the existing torch.qint8 and + torch.quint8 datatypes, the user can choose to use dynamic quantization range by passing + in a tuple of initial qmin and qmax values. One use case is these customized qmin and qmax + values are used to calculate static estimates of the scale and zero point for aggressive lower-bit + fake quantization. These estimates are compared against parameters learned through backpropagation. + The related literatures for scale and zero point via backpropagation are as follows: + + Learned Step Size Quantization: https://openreview.net/pdf?id=rkgO66VKDS + Trained Quantization Thresholds: https://arxiv.org/pdf/1903.08066.pdf + """ + # The variable names are prefixed with "initial" because their values (qmin and qmax) might be adjusted + # based on whether quantization range is reduced and the datatype (signed/unsigned) used by the observer. + assert quant_min <= 0 <= quant_max, ( + "Used-specified quantization range must include 0." + ) + assert quant_min < quant_max, ( + "qmin must be strictly less than qmax for user-specified quantization range." + ) + + @torch.jit.export + def _calculate_qparams( + self, min_val: torch.Tensor, max_val: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + r"""Calculates the quantization parameters, given min and max + value tensors. Works for both per tensor and per channel cases + + Args: + min_val: Minimum values per channel + max_val: Maximum values per channel + + Returns: + scales: Scales tensor of shape (#channels,) + zero_points: Zero points tensor of shape (#channels,) + """ + # Functionally equivalent to 'determine_qparams' in utils.py. Observers must be torchscriptable however and qscheme + # as far as I can tell is not allowed to passed as a parameter in torchscript functions. This makes refactoring observer + # to use this utility a massive pain and very gross. For now Im opting just to duplicate as this code + # seems unlikey to change (last update over 1 year ago) and when torchscript is fully deprecated we can refactor. + # TODO(jakeszwe, jerryzh168) + if not check_min_max_valid(min_val, max_val): + return torch.tensor([1.0], device=min_val.device.type), torch.tensor( + [0], device=min_val.device.type + ) + + quant_min, quant_max = self.quant_min, self.quant_max + min_val_neg = torch.min(min_val, torch.zeros_like(min_val)) + max_val_pos = torch.max(max_val, torch.zeros_like(max_val)) + + device = min_val_neg.device + scale = torch.ones(min_val_neg.size(), dtype=torch.float32, device=device) + zero_point = torch.zeros(min_val_neg.size(), dtype=torch.int64, device=device) + + if ( + self.qscheme == torch.per_tensor_symmetric + or self.qscheme == torch.per_channel_symmetric + ): + max_val_pos = torch.max(-min_val_neg, max_val_pos) + scale = max_val_pos / (float(quant_max - quant_min) / 2) + scale = torch.max(scale, self.eps) + if self.dtype in [torch.quint8, torch.uint8]: + if self.has_customized_qrange: + # When customized quantization range is used, down-rounded midpoint of the range is chosen. + zero_point = zero_point.new_full( + zero_point.size(), (quant_min + quant_max) // 2 + ) + else: + zero_point = zero_point.new_full(zero_point.size(), 128) + elif self.dtype in [torch.uint16]: + zero_point = zero_point.new_full(zero_point.size(), 2**15) + elif self.qscheme == torch.per_channel_affine_float_qparams: + scale = (max_val - min_val) / float(quant_max - quant_min) + scale = torch.where(scale > self.eps, scale, torch.ones_like(scale)) + # We use the quantize function + # xq = Round(Xf * inv_scale + zero_point), + # setting zero_point to (-1 * min *inv_scale) we get + # Xq = Round((Xf - min) * inv_scale) + zero_point = -1 * min_val / scale + else: + scale = (max_val_pos - min_val_neg) / float(quant_max - quant_min) + scale = torch.max(scale, self.eps) + zero_point = quant_min - torch.round(min_val_neg / scale).to(torch.int) + zero_point = torch.clamp(zero_point, quant_min, quant_max) + + # For scalar values, cast them to Tensors of size 1 to keep the shape + # consistent with default values in FakeQuantize. + if len(scale.shape) == 0: + # TODO: switch to scale.item() after adding JIT support + scale = torch.tensor([float(scale)], dtype=scale.dtype, device=device) + if len(zero_point.shape) == 0: + # TODO: switch to zero_point.item() after adding JIT support + zero_point = torch.tensor( + [int(zero_point)], dtype=zero_point.dtype, device=device + ) + if self.qscheme == torch.per_channel_affine_float_qparams: + zero_point = torch.tensor( + [float(zero_point)], dtype=zero_point.dtype, device=device + ) + + return scale, zero_point + + @torch.jit.export + def reset_min_max_vals(self): + raise NotImplementedError("Cannot reset min/max values in the given observer.") + + +# Originally, this class was called `_ObserverBase`. Keeping the old name around +# for backwards compatibility. +# TODO(after v1.13): delete this +_ObserverBase = UniformQuantizationObserverBase + + +class MinMaxObserver(UniformQuantizationObserverBase): + r"""Observer module for computing the quantization parameters based on the + running min and max values. + + This observer uses the tensor min/max statistics to compute the quantization + parameters. The module records the running minimum and maximum of incoming + tensors, and uses this statistic to compute the quantization parameters. + + Args: + dtype: dtype argument to the `quantize` node needed to implement the + reference model spec. + qscheme: Quantization scheme to be used + reduce_range: Reduces the range of the quantized data type by 1 bit + quant_min: Minimum quantization value. If unspecified, it will follow the 8-bit setup. + quant_max: Maximum quantization value. If unspecified, it will follow the 8-bit setup. + eps: Epsilon value for float32, Defaults to `torch.finfo(torch.float32).eps`. + + Given running min/max as :math:`x_\text{min}` and :math:`x_\text{max}`, + scale :math:`s` and zero point :math:`z` are computed as: + + The running minimum/maximum :math:`x_\text{min/max}` is computed as: + + .. math:: + + \begin{array}{ll} + x_\text{min} &= \begin{cases} + \min(X) & \text{if~}x_\text{min} = \text{None} \\ + \min\left(x_\text{min}, \min(X)\right) & \text{otherwise} + \end{cases}\\ + x_\text{max} &= \begin{cases} + \max(X) & \text{if~}x_\text{max} = \text{None} \\ + \max\left(x_\text{max}, \max(X)\right) & \text{otherwise} + \end{cases}\\ + \end{array} + + where :math:`X` is the observed tensor. + + The scale :math:`s` and zero point :math:`z` are then computed as: + + .. math:: + + \begin{aligned} + \text{if Symmetric:}&\\ + &s = 2 \max(|x_\text{min}|, x_\text{max}) / + \left( Q_\text{max} - Q_\text{min} \right) \\ + &z = \begin{cases} + 0 & \text{if dtype is qint8} \\ + 128 & \text{otherwise} + \end{cases}\\ + \text{Otherwise:}&\\ + &s = \left( x_\text{max} - x_\text{min} \right ) / + \left( Q_\text{max} - Q_\text{min} \right ) \\ + &z = Q_\text{min} - \text{round}(x_\text{min} / s) + \end{aligned} + + where :math:`Q_\text{min}` and :math:`Q_\text{max}` are the minimum and + maximum of the quantized data type. + + .. warning:: :attr:`dtype` can only take ``torch.qint8`` or ``torch.quint8``. + + .. note:: If the running minimum equals to the running maximum, the scale + and zero_point are set to 1.0 and 0. + """ + + min_val: torch.Tensor + max_val: torch.Tensor + + def __init__( + self, + dtype=torch.quint8, + qscheme=torch.per_tensor_affine, + reduce_range=False, + quant_min=None, + quant_max=None, + factory_kwargs=None, + eps=torch.finfo(torch.float32).eps, + is_dynamic=False, + **kwargs, + ) -> None: + if not is_per_tensor(qscheme): + raise NotImplementedError( + "MinMaxObserver's qscheme only support torch.per_tensor_symmetric \ + and torch.per_tensor_affine." + ) + # TODO: MinMaxObserver by itself doesn't support dynamic quantization, but + # if it's inherited by MovingAverageObserver, and averaging_constant is 1, it + # supports dynamic quantization, we may need to better error checking here + + # For x86 quantized kernels, we need to ensure that the vpmaddubsw + # instruction does not overflow. We allow for a reduce_range argument to + # observers that reduces the quantized range to (0,127) or (-64, 63). + # For more details see aten/src/ATen/native/quantized/cpu/qconv.cpp + # This is not an optimal choice for non x86 backends as it loses a bit + # of precision for activations. + super().__init__( + dtype=dtype, + qscheme=qscheme, + reduce_range=reduce_range, + quant_min=quant_min, + quant_max=quant_max, + factory_kwargs=factory_kwargs, + eps=eps, + is_dynamic=is_dynamic, + **kwargs, + ) + factory_kwargs = torch.nn.factory_kwargs(factory_kwargs) + self.register_buffer("min_val", torch.tensor(float("inf"), **factory_kwargs)) + self.register_buffer("max_val", torch.tensor(float("-inf"), **factory_kwargs)) + if ( + self.qscheme == torch.per_tensor_symmetric + and self.reduce_range + and self.dtype == torch.quint8 + ): + raise NotImplementedError( + "Cannot reduce range for symmetric \ + quantization for quint8" + ) + + def forward(self, x_orig): + r"""Records the running minimum and maximum of ``x``.""" + if x_orig.numel() == 0: + return x_orig + x = x_orig.detach() # avoid keeping autograd tape + x = x.to(self.min_val.dtype) + min_val_cur, max_val_cur = torch.aminmax(x) + min_val = torch.min(min_val_cur, self.min_val) + max_val = torch.max(max_val_cur, self.max_val) + self.min_val.copy_(min_val) + self.max_val.copy_(max_val) + return x_orig + + @torch.jit.export + def calculate_qparams(self): + r"""Calculates the quantization parameters.""" + return self._calculate_qparams(self.min_val, self.max_val) + + @torch.jit.export + def extra_repr(self): + return f"min_val={self.min_val}, max_val={self.max_val}" + + @torch.jit.export + def reset_min_max_vals(self): + """Resets the min/max values.""" + self.min_val.copy_(torch.tensor(float("inf"))) + self.max_val.copy_(torch.tensor(float("-inf"))) + + +class MovingAverageMinMaxObserver(MinMaxObserver): + r"""Observer module for computing the quantization parameters based on the + moving average of the min and max values. + + This observer computes the quantization parameters based on the moving + averages of minimums and maximums of the incoming tensors. The module + records the average minimum and maximum of incoming tensors, and uses this + statistic to compute the quantization parameters. + + Args: + averaging_constant: Averaging constant for min/max. + dtype: dtype argument to the `quantize` node needed to implement the + reference model spec. + qscheme: Quantization scheme to be used + reduce_range: Reduces the range of the quantized data type by 1 bit + quant_min: Minimum quantization value. If unspecified, it will follow the 8-bit setup. + quant_max: Maximum quantization value. If unspecified, it will follow the 8-bit setup. + eps: Epsilon value for float32, Defaults to `torch.finfo(torch.float32).eps`. + + The moving average min/max is computed as follows + + .. math:: + + \begin{array}{ll} + x_\text{min} = \begin{cases} + \min(X) & \text{if~}x_\text{min} = \text{None} \\ + (1 - c) x_\text{min} + c \min(X) & \text{otherwise} + \end{cases}\\ + x_\text{max} = \begin{cases} + \max(X) & \text{if~}x_\text{max} = \text{None} \\ + (1 - c) x_\text{max} + c \max(X) & \text{otherwise} + \end{cases}\\ + \end{array} + + where :math:`x_\text{min/max}` is the running average min/max, :math:`X` is + is the incoming tensor, and :math:`c` is the ``averaging_constant``. + + The scale and zero point are then computed as in + :class:`~torchao.quantization.pt2e.observer.MinMaxObserver`. + + .. note:: Only works with ``torch.per_tensor_affine`` quantization scheme. + + .. note:: If the running minimum equals to the running maximum, the scale + and zero_point are set to 1.0 and 0. + """ + + def __init__( + self, + averaging_constant=0.01, + dtype=torch.quint8, + qscheme=torch.per_tensor_affine, + reduce_range=False, + quant_min=None, + quant_max=None, + eps=torch.finfo(torch.float32).eps, + is_dynamic=False, + **kwargs, + ) -> None: + if not is_per_tensor(qscheme): + raise NotImplementedError( + f"MovingAverageMinMaxObserver's qscheme only support \ + torch.per_tensor_symmetric and torch.per_tensor_affine. \ + but got: {qscheme}" + ) + self.averaging_constant = averaging_constant + if is_dynamic and self.averaging_constant != 1: + raise NotImplementedError( + "MovingAverageMinMaxObserver doesn't support dynamic quantization for " + f"averaging constant of {self.averaging_constant}" + ) + super().__init__( + dtype=dtype, + qscheme=qscheme, + reduce_range=reduce_range, + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + is_dynamic=is_dynamic, + **kwargs, + ) + + def forward(self, x_orig): + if x_orig.numel() == 0: + return x_orig + x = x_orig.detach() # avoid keeping autograd tape + x = x.to(self.min_val.dtype) + min_val = self.min_val + max_val = self.max_val + if min_val == float("inf") and max_val == float("-inf"): + min_val, max_val = torch.aminmax(x) + else: + min_val_cur, max_val_cur = torch.aminmax(x) + min_val = min_val + self.averaging_constant * (min_val_cur - min_val) + max_val = max_val + self.averaging_constant * (max_val_cur - max_val) + self.min_val.copy_(min_val) + self.max_val.copy_(max_val) + return x_orig + + +class PerChannelMinMaxObserver(UniformQuantizationObserverBase): + r"""Observer module for computing the quantization parameters based on the + running per channel min and max values. + + This observer uses the tensor min/max statistics to compute the per channel + quantization parameters. The module records the running minimum and maximum + of incoming tensors, and uses this statistic to compute the quantization + parameters. + + Args: + ch_axis: Channel axis + dtype: dtype argument to the `quantize` node needed to implement the + reference model spec. + qscheme: Quantization scheme to be used + reduce_range: Reduces the range of the quantized data type by 1 bit + quant_min: Minimum quantization value. If unspecified, it will follow the 8-bit setup. + quant_max: Maximum quantization value. If unspecified, it will follow the 8-bit setup. + eps: Epsilon value for float32, Defaults to `torch.finfo(torch.float32).eps`. + + The quantization parameters are computed the same way as in + :class:`~torchao.quantization.pt2e.observer.MinMaxObserver`, with the difference + that the running min/max values are stored per channel. + Scales and zero points are thus computed per channel as well. + + .. note:: If the running minimum equals to the running maximum, the scales + and zero_points are set to 1.0 and 0. + """ + + min_val: torch.Tensor + max_val: torch.Tensor + + def __init__( + self, + ch_axis=0, + dtype=torch.quint8, + qscheme=torch.per_channel_affine, + reduce_range=False, + quant_min=None, + quant_max=None, + factory_kwargs=None, + eps=torch.finfo(torch.float32).eps, + is_dynamic=False, + **kwargs, + ) -> None: + if not is_per_channel(qscheme): + raise NotImplementedError( + "PerChannelMinMaxObserver's qscheme only support \ + torch.per_channel_symmetric, torch.per_channel_affine and torch.per_channel_affine_float_qparams." + ) + if is_dynamic: + raise NotImplementedError( + "PerChannelMinMaxObserver doesn't support dynamic quantization" + ) + super().__init__( + dtype=dtype, + qscheme=qscheme, + reduce_range=reduce_range, + quant_min=quant_min, + quant_max=quant_max, + factory_kwargs=factory_kwargs, + eps=eps, + is_dynamic=is_dynamic, + **kwargs, + ) + factory_kwargs = torch.nn.factory_kwargs(factory_kwargs) + self.ch_axis = ch_axis + self.register_buffer("min_val", torch.tensor([], **factory_kwargs)) + self.register_buffer("max_val", torch.tensor([], **factory_kwargs)) + if ( + self.qscheme == torch.per_channel_symmetric + and self.reduce_range + and self.dtype == torch.quint8 + ): + raise NotImplementedError( + "Cannot reduce range for symmetric quantization for quint8" + ) + + def forward(self, x_orig): + return self._forward(x_orig) + + def _forward(self, x_orig): + if x_orig.numel() == 0: + return x_orig + x = x_orig.detach() # avoid keeping autograd tape + min_val = self.min_val + max_val = self.max_val + x_dim = x.size() + + new_axis_list = [i for i in range(len(x_dim))] # noqa: C416 + new_axis_list[self.ch_axis] = 0 + new_axis_list[0] = self.ch_axis + y = x.permute(new_axis_list) + # Need to match dtype of min/max because the updates to buffers + # are done in place and types need to match for comparisons + y = y.to(self.min_val.dtype) + y = torch.flatten(y, start_dim=1) + if min_val.numel() == 0 or max_val.numel() == 0: + min_val, max_val = torch.aminmax(y, dim=1) + else: + min_val_cur, max_val_cur = torch.aminmax(y, dim=1) + min_val = torch.min(min_val_cur, min_val) + max_val = torch.max(max_val_cur, max_val) + self.min_val.resize_(min_val.shape) + self.max_val.resize_(max_val.shape) + self.min_val.copy_(min_val) + self.max_val.copy_(max_val) + return x_orig + + @torch.jit.export + def calculate_qparams(self): + return self._calculate_qparams(self.min_val, self.max_val) + + def extra_repr(self): + return f"min_val={self.min_val}, max_val={self.max_val}" + + def _load_from_state_dict( + self, + state_dict: dict[str, Any], + prefix: str, + local_metadata: dict[str, torch.Tensor], + strict: bool, + missing_keys: list[str], + unexpected_keys: list[str], + error_msgs: list[str], + ): + version = local_metadata.get("version", None) + if version is not None and version < 3: + local_state = ["min_vals", "max_vals"] + expected_min_name = "min_vals" + expected_max_name = "max_vals" + else: + local_state = ["min_val", "max_val"] + expected_min_name = "min_val" + expected_max_name = "max_val" + for name in local_state: + key = prefix + name + if key in state_dict: + val = state_dict[key] + # Custom handling to allow loading min_val or max_val + # of size N into uninitialized buffers of size 0. The + # buffers are resized here, and the values are copied in + # the default state_dict loading code of the parent. + if name == expected_min_name: + self.min_val.resize_(val.shape) + elif name == expected_max_name: + self.max_val.resize_(val.shape) + else: + warnings.warn( + f"Observer load_from_state_dict got unexpected name {name}" + ) + # For torchscript module we need to update the attributes here since we do not + # call the `_load_from_state_dict` function defined module.py + if torch.jit.is_scripting(): + if name == expected_min_name: + self.min_val.copy_(val) + elif name == expected_max_name: + self.max_val.copy_(val) + else: + warnings.warn( + f"Observer load_from_state_dict got unexpected name {name}" + ) + elif strict: + missing_keys.append(key) + + if not torch.jit.is_scripting(): + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + False, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def _load_from_state_dict_script( + self, + state_dict: dict[str, Any], + prefix: str, + local_metadata: dict[str, torch.Tensor], + strict: bool, + missing_keys: list[str], + unexpected_keys: list[str], + error_msgs: list[str], + ): + self._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + @torch.jit.export + def reset_min_max_vals(self): + """Resets the min/max values.""" + # This used to be torch.ones but that does not work because + # JIT compiler can optimize it via common subexpression elimination + # in which case both min_val and max_val point to the same tensor. + self.min_val = torch.rand( + 0, + ) + self.max_val = torch.rand( + 0, + ) + + +class MovingAveragePerChannelMinMaxObserver(PerChannelMinMaxObserver): + r"""Observer module for computing the quantization parameters based on the + running per channel min and max values. + + This observer uses the tensor min/max statistics to compute the per channel + quantization parameters. The module records the running minimum and maximum + of incoming tensors, and uses this statistic to compute the quantization + parameters. + + Args: + averaging_constant: Averaging constant for min/max. + ch_axis: Channel axis + dtype: Quantized data type + qscheme: Quantization scheme to be used + reduce_range: Reduces the range of the quantized data type by 1 bit + quant_min: Minimum quantization value. If unspecified, it will follow the 8-bit setup. + quant_max: Maximum quantization value. If unspecified, it will follow the 8-bit setup. + eps: Epsilon value for float32, Defaults to `torch.finfo(torch.float32).eps`. + + The quantization parameters are computed the same way as in + :class:`~torchao.quantization.pt2e.observer.MovingAverageMinMaxObserver`, with the + difference that the running min/max values are stored per channel. + Scales and zero points are thus computed per channel as well. + + .. note:: If the running minimum equals to the running maximum, the scales + and zero_points are set to 1.0 and 0. + """ + + def __init__( + self, + averaging_constant=0.01, + ch_axis=0, + dtype=torch.quint8, + qscheme=torch.per_channel_affine, + reduce_range=False, + quant_min=None, + quant_max=None, + eps=torch.finfo(torch.float32).eps, + is_dynamic=False, + **kwargs, + ) -> None: + if not is_per_channel(qscheme): + raise NotImplementedError( + "MovingAveragePerChannelMinMaxObserver's qscheme only support \ + torch.per_channel_symmetric, torch.per_channel_affine and torch.per_channel_affine_float_qparams." + ) + if is_dynamic: + raise NotImplementedError( + "MovingAveragePerChannelMinMaxObserver doesn't support dynamic quantization" + ) + super().__init__( + ch_axis=ch_axis, + dtype=dtype, + qscheme=qscheme, + reduce_range=reduce_range, + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + is_dynamic=is_dynamic, + **kwargs, + ) + self.averaging_constant = averaging_constant + + def forward(self, x_orig): + if x_orig.numel() == 0: + return x_orig + x = x_orig.detach() # avoid keeping autograd tape + x = x.to(self.min_val.dtype) + min_val = self.min_val + max_val = self.max_val + x_dim = x.size() + + new_axis_list = [i for i in range(len(x_dim))] # noqa: C416 + new_axis_list[self.ch_axis] = 0 + new_axis_list[0] = self.ch_axis + y = x.permute(new_axis_list) + y = torch.flatten(y, start_dim=1) + if min_val.numel() == 0 or max_val.numel() == 0: + min_val, max_val = torch.aminmax(y, dim=1) + else: + min_val_cur, max_val_cur = torch.aminmax(y, dim=1) + min_val = min_val + self.averaging_constant * (min_val_cur - min_val) + max_val = max_val + self.averaging_constant * (max_val_cur - max_val) + self.min_val.resize_(min_val.shape) + self.max_val.resize_(max_val.shape) + self.min_val.copy_(min_val) + self.max_val.copy_(max_val) + return x_orig + + +class HistogramObserver(UniformQuantizationObserverBase): + r""" + The module records the running histogram of tensor values along with + min/max values. ``calculate_qparams`` will calculate scale and zero_point. + + Args: + bins: Number of bins to use for the histogram + dtype: dtype argument to the `quantize` node needed to implement the + reference model spec + qscheme: Quantization scheme to be used + reduce_range: Reduces the range of the quantized data type by 1 bit + eps: Epsilon value for float32, Defaults to `torch.finfo(torch.float32).eps`. + + The scale and zero point are computed as follows: + + 1. Create the histogram of the incoming inputs. + The histogram is computed continuously, and the ranges per bin change + with every new tensor observed. + 2. Search the distribution in the histogram for optimal min/max values. + The search for the min/max values ensures the minimization of the + quantization error with respect to the floating point model. + 3. Compute the scale and zero point the same way as in the + :class:`~torchao.quantization.pt2e.MinMaxObserver` + """ + + histogram: torch.Tensor + min_val: torch.Tensor + max_val: torch.Tensor + + def __init__( + self, + bins: int = 2048, + dtype: torch.dtype = torch.quint8, + qscheme=torch.per_tensor_affine, + reduce_range=False, + quant_min=None, + quant_max=None, + factory_kwargs=None, + eps=torch.finfo(torch.float32).eps, + is_dynamic=False, + **kwargs, + ) -> None: + if not is_per_tensor(qscheme): + raise NotImplementedError( + "HistogramObserver's qscheme only support torch.per_tensor_symmetric \ + and torch.per_tensor_affine." + ) + if is_dynamic: + raise NotImplementedError( + "HistogramObserver doesn't support dynamic quantization" + ) + # bins: The number of bins used for histogram calculation. + super().__init__( + dtype=dtype, + qscheme=qscheme, + reduce_range=reduce_range, + quant_min=quant_min, + quant_max=quant_max, + factory_kwargs=factory_kwargs, + eps=eps, + is_dynamic=is_dynamic, + **kwargs, + ) + factory_kwargs = torch.nn.factory_kwargs(factory_kwargs) + self.bins = bins + self.register_buffer("histogram", torch.zeros(self.bins, **factory_kwargs)) + self.register_buffer("min_val", torch.tensor(float("inf"), **factory_kwargs)) + self.register_buffer("max_val", torch.tensor(float("-inf"), **factory_kwargs)) + self.dst_nbins = 2 ** torch.iinfo(self.dtype).bits + self.upsample_rate = ( + 16 # used to reduce quantization errors when upscaling histogram + ) + + def _get_norm( + self, delta_begin: torch.Tensor, delta_end: torch.Tensor, density: torch.Tensor + ) -> torch.Tensor: + r""" + Compute the norm of the values uniformaly distributed between + delta_begin and delta_end. + Currently only L2 norm is supported. + + norm = density * (integral_{begin, end} x^2) + = density * (end^3 - begin^3) / 3 + """ + norm = ( + delta_end * delta_end * delta_end - delta_begin * delta_begin * delta_begin + ) / 3 + return density * norm + + def _compute_quantization_error(self, next_start_bin: int, next_end_bin: int): + r""" + Compute the quantization error if we use start_bin to end_bin as the + min and max to do the quantization. + """ + bin_width = (self.max_val.item() - self.min_val.item()) / self.bins + + dst_bin_width = bin_width * (next_end_bin - next_start_bin + 1) / self.dst_nbins + if dst_bin_width == 0.0: + return 0.0 + + src_bin = torch.arange(self.bins, device=self.histogram.device) + # distances from the beginning of first dst_bin to the beginning and + # end of src_bin + src_bin_begin = (src_bin - next_start_bin) * bin_width + src_bin_end = src_bin_begin + bin_width + + # which dst_bins the beginning and end of src_bin belong to? + dst_bin_of_begin = torch.clamp( + torch.div(src_bin_begin, dst_bin_width, rounding_mode="floor"), + 0, + self.dst_nbins - 1, + ) + dst_bin_of_begin_center = (dst_bin_of_begin + 0.5) * dst_bin_width + + dst_bin_of_end = torch.clamp( + torch.div(src_bin_end, dst_bin_width, rounding_mode="floor"), + 0, + self.dst_nbins - 1, + ) + density = self.histogram / bin_width + + norm = torch.zeros(self.bins, device=self.histogram.device) + + delta_begin = src_bin_begin - dst_bin_of_begin_center + delta_end = dst_bin_width / 2 + norm += self._get_norm( + delta_begin, + torch.ones(self.bins, device=self.histogram.device) * delta_end, + density, + ) + + norm += (dst_bin_of_end - dst_bin_of_begin - 1) * self._get_norm( + torch.tensor(-dst_bin_width / 2), torch.tensor(dst_bin_width / 2), density + ) + + dst_bin_of_end_center = dst_bin_of_end * dst_bin_width + dst_bin_width / 2 + + delta_begin = -dst_bin_width / 2 + delta_end = src_bin_end - dst_bin_of_end_center + norm += self._get_norm(torch.tensor(delta_begin), delta_end, density) + + return norm.sum().item() + + def _non_linear_param_search(self) -> tuple[torch.Tensor, torch.Tensor]: + r"""Non-linear parameter search. + + An approximation for L2 error minimization for selecting min/max. + By selecting new min/max, we filter out outliers in input distribution. + This follows the implementation of NormMinimization::NonlinearQuantizationParamsSearch in + caffe2/quantization/server/norm_minimization.cc + """ + assert self.histogram.size()[0] == self.bins, "bins mismatch" + bin_width = (self.max_val - self.min_val) / self.bins + + # cumulative sum + total = torch.sum(self.histogram).item() + cSum = torch.cumsum(self.histogram, dim=0) + + stepsize = 1e-5 # granularity + alpha = 0.0 # lower bound + beta = 1.0 # upper bound + start_bin = 0 + end_bin = self.bins - 1 + norm_min = float("inf") + + while alpha < beta: + # Find the next step + next_alpha = alpha + stepsize + next_beta = beta - stepsize + + # find the left and right bins between the quantile bounds + l = start_bin + r = end_bin + while l < end_bin and cSum[l] < next_alpha * total: + l = l + 1 + while r > start_bin and cSum[r] > next_beta * total: + r = r - 1 + + # decide the next move + next_start_bin = start_bin + next_end_bin = end_bin + if (l - start_bin) > (end_bin - r): + # move the start bin + next_start_bin = l + alpha = next_alpha + else: + # move the end bin + next_end_bin = r + beta = next_beta + + if next_start_bin == start_bin and next_end_bin == end_bin: + continue + + # calculate the quantization error using next_start_bin and next_end_bin + norm = self._compute_quantization_error(next_start_bin, next_end_bin) + + if norm > norm_min: + break + norm_min = norm + start_bin = next_start_bin + end_bin = next_end_bin + + new_min = self.min_val + bin_width * start_bin + new_max = self.min_val + bin_width * (end_bin + 1) + return new_min, new_max + + def _upscale_histogram( + self, + histogram: torch.Tensor, + orig_min: torch.Tensor, + orig_max: torch.Tensor, + update_min: torch.Tensor, + update_max: torch.Tensor, + ): + # this turns the histogram into a more fine-coarsed histogram to reduce + # bin quantization errors + histogram = histogram.repeat_interleave(self.upsample_rate) / self.upsample_rate + bin_size = (orig_max - orig_min) / (self.bins * self.upsample_rate) + mid_points_histogram = ( + torch.linspace( + orig_min, + orig_max, + self.bins * self.upsample_rate + 1, + device=orig_min.device, + )[:-1].to(histogram.device) + + 0.5 * bin_size + ) + boundaries_new_histogram = torch.linspace( + update_min, update_max, self.bins + 1, device=update_min.device + ).to(histogram.device) + # this maps the mid-poits of the histogram to the new histogram's space + bucket_assignments = ( + torch.bucketize(mid_points_histogram, boundaries_new_histogram, right=True) + - 1 + ) + # this then maps the histogram mid-points in the new space, weighted by the original histogram's values + # this is just the old histogram in the new histogram's space + + # In case due to numerical issues the values land higher/lower than the maximum/minimum + bucket_assignments[bucket_assignments >= self.bins] = self.bins - 1 + bucket_assignments[bucket_assignments < 0] = 0 + + update_histogram = torch.bincount( + bucket_assignments, weights=histogram, minlength=self.bins + ) + return update_histogram + + def _combine_histograms( + self, + orig_hist: torch.Tensor, + orig_min: torch.Tensor, + orig_max: torch.Tensor, + update_hist: torch.Tensor, + update_min: torch.Tensor, + update_max: torch.Tensor, + ) -> torch.Tensor: + # If the new min and max are the same as the current min and max, + # we can just add the new histogram to the original histogram + if update_min == orig_min and update_max == orig_max: + return orig_hist + update_hist + + # If the orig hist only has one value (i.e., the min and max are the same) + # we can just add it into new histogram + if orig_min == orig_max: + bin_value = torch.sum(update_hist) + transformed_orig_hist = ( + torch.histc(orig_min, bins=self.bins, min=update_min, max=update_max) # type: ignore[arg-type] + * bin_value + ) + return transformed_orig_hist + update_hist + + # We assume the update_hist is already in the target range, we will map the orig_max to it + assert update_min <= orig_min + assert update_max >= orig_max + + # Now we need to turn the old_histogram, into the range of the new histogram + transformed_orig_hist = self._upscale_histogram( + orig_hist, + orig_min, + orig_max, + update_min, + update_max, + ) + + return update_hist + transformed_orig_hist + + def reset_histogram( + self, x: torch.Tensor, min_val: torch.Tensor, max_val: torch.Tensor + ) -> None: + self.min_val.resize_(min_val.shape) + self.min_val.copy_(min_val) + self.max_val.resize_(max_val.shape) + self.max_val.copy_(max_val) + assert min_val.numel() == 1 and max_val.numel() == 1, ( + "histogram min/max values must be scalar." + ) + new_histogram = torch.histc(x, self.bins, min=min_val, max=max_val) # type: ignore[arg-type] + self.histogram.detach_().resize_(new_histogram.shape) + self.histogram.copy_(new_histogram) + + def forward(self, x_orig: torch.Tensor) -> torch.Tensor: # pyre-ignore[14] + if x_orig.numel() == 0: + return x_orig + x = x_orig.detach() + x_min, x_max = torch.aminmax(x) + # want to ignore torch.inf since we don't actually + # want to make our quantization range infinite + # and in practice those values will be clamped + if x_min == -torch.inf or x_max == torch.inf: + warnings.warn("torch.inf detected in input tensor, ignoring input") + x = x[x.abs() != torch.inf] + if x.numel() == 0: + return x_orig + x_min, x_max = torch.aminmax(x) + + current_min = self.min_val + current_max = self.max_val + + is_uninitialized = self.min_val == float("inf") or self.max_val == float("-inf") + if is_uninitialized: + self.reset_histogram(x, x_min, x_max) + else: + update_min, update_max = x_min, x_max + new_min = torch.min(current_min, update_min) + new_max = torch.max(current_max, update_max) + + # TODO: For some reason, this is required for it to pass torchscript test + # new_min and new_max should already have requires_grad set to False + new_min, new_max = new_min.detach(), new_max.detach() + update_histogram = torch.histc( + x, + self.bins, + min=new_min, + max=new_max, # type: ignore[arg-type] + ).to(self.histogram.device) + if new_min == current_min and new_max == current_max: + combined_histogram = self.histogram + update_histogram + self.histogram.detach_().resize_(combined_histogram.shape) + self.histogram.copy_(combined_histogram) + else: + combined_histogram = self._combine_histograms( + self.histogram, + current_min, + current_max, + update_histogram, + new_min, + new_max, + ) + self.histogram.detach_().resize_(combined_histogram.shape) + self.histogram.copy_(combined_histogram) + self.min_val.detach_().resize_(new_min.shape) + self.min_val.copy_(new_min) + self.max_val.detach_().resize_(new_max.shape) + self.max_val.copy_(new_max) + + return x_orig + + @torch.jit.export + def calculate_qparams(self): + is_uninitialized = self.min_val == float("inf") and self.max_val == float( + "-inf" + ) + if is_uninitialized: + warnings.warn( + "must run observer before calling calculate_qparams.\ + Returning default scale and zero point " + ) + return torch.tensor([1.0], device=self.min_val.device.type), torch.tensor( + [0], device=self.min_val.device.type + ) + assert self.bins == len(self.histogram), ( + "The number of bins in histogram should be equal to the number of bins " + "supplied while making this observer" + ) + + new_min, new_max = self._non_linear_param_search() + + return self._calculate_qparams(new_min, new_max) + + def _save_to_state_dict(self, destination, prefix, keep_vars): + super()._save_to_state_dict(destination, prefix, keep_vars) + destination[prefix + "min_val"] = self.min_val + destination[prefix + "max_val"] = self.max_val + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + version = local_metadata.get("version", None) + + if version is None or version < 3: + # if min_val and max_val are not initialized, update their shape + # to account for the differences between v2 and v3 + min_val_name, max_val_name = prefix + "min_val", prefix + "max_val" + if min_val_name in state_dict: + if state_dict[min_val_name].shape == torch.Size([0]): + state_dict[min_val_name] = torch.tensor(float("inf")) + if max_val_name in state_dict: + if state_dict[max_val_name].shape == torch.Size([0]): + state_dict[max_val_name] = torch.tensor(float("-inf")) + + local_state = ["min_val", "max_val"] + for name in local_state: + key = prefix + name + if key in state_dict: + val = state_dict[key] + setattr(self, name, val) + elif strict: + missing_keys.append(key) + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def extra_repr(self): + return f"min_val={self.min_val}, max_val={self.max_val}" + + +class FixedQParamsObserver(ObserverBase): + r""" + Observer that simulates quantize and dequantize with fixed + quantization parameters in training time. Only per tensor + quantization is supported. + + Args: + `scale` (float): fixed scale for the observer + `zero_point` (int): fixed zero point for the observer + `dtype`, `qscheme`, `quant_min`, `quant_max` + """ + + scale: torch.Tensor + zero_point: torch.Tensor + + def __init__( + self, + scale, + zero_point, + dtype=torch.quint8, + qscheme=torch.per_tensor_affine, + quant_min=0, + quant_max=255, + is_dynamic=False, + **kwargs, + ): + if is_dynamic: + raise NotImplementedError( + "FixedQParamsObserver doesn't support dynamic quantization" + ) + super().__init__(dtype=dtype, is_dynamic=is_dynamic, **kwargs) + self.quant_min = quant_min + self.quant_max = quant_max + self.register_buffer("scale", torch.tensor([scale], dtype=torch.float)) + self.register_buffer("zero_point", torch.tensor([zero_point], dtype=torch.int)) + self.dtype = dtype + self.qscheme = qscheme + + def forward(self, X): + return X + + @torch.jit.export + def calculate_qparams(self): + return self.scale, self.zero_point + + +class PlaceholderObserver(ObserverBase): + r""" + Observer that doesn't do anything and just passes its configuration to the + quantized module's ``.from_float()``. + + Can be used for quantization to float16 which doesn't require determining + ranges. + + Args: + dtype: dtype argument to the `quantize` node needed to implement the + reference model spec. + quant_min: minimum value in quantized domain (TODO: align behavior with other observers) + quant_max: maximum value in quantized domain + custom_op_name: (temporary) specify this observer for an operator that doesn't require any observation + (Can be used in Graph Mode Passes for special case ops). + compute_dtype (deprecated): if set, marks the future quantize function to use + dynamic quantization instead of static quantization. + This field is deprecated, use `is_dynamic=True` instead. + is_dynamic: if True, the `quantize` function in the reference model + representation taking stats from this observer instance will + use dynamic quantization. + """ + + def __init__( + self, + dtype=torch.float32, + custom_op_name="", + compute_dtype=None, + quant_min=None, + quant_max=None, + qscheme=None, + eps=None, + is_dynamic=False, + ) -> None: + super().__init__(dtype=dtype, is_dynamic=is_dynamic) + if qscheme is None: + qscheme = torch.per_tensor_affine + if eps is None: + eps = torch.finfo(torch.float32).eps + + # dtype of input of the target operator, e.g. for dynamic quantization + # ops, the dtype will be float32 + self.dtype = dtype + self.qscheme = qscheme + self.quant_min = quant_min + self.quant_max = quant_max + self.eps = eps + self.custom_op = custom_op_name + # used for configuration of computation type for dynamic quantization + if compute_dtype: + is_dynamic = True + warnings.warn( + "Please use `is_dynamic` instead of `compute_dtype`. \ + `compute_dtype` will be deprecated in a future release \ + of PyTorch." + ) + + def forward(self, x): + return x + + @torch.jit.export + def extra_repr(self): + return f"dtype={self.dtype}, is_dynamic={self.is_dynamic}" + + @torch.jit.export + def calculate_qparams(self): + raise Exception( # noqa: TRY002 + "calculate_qparams should not be called for PlaceholderObserver" + ) + + +class RecordingObserver(ObserverBase): + r""" + The module is mainly for debug and records the tensor values during runtime. + + Args: + dtype: Quantized data type + qscheme: Quantization scheme to be used + reduce_range: Reduces the range of the quantized data type by 1 bit + """ + + __annotations__ = {"tensor_val": list[Optional[torch.Tensor]]} + + def __init__(self, dtype=torch.quint8): + super().__init__(dtype=dtype, is_dynamic=False) + self.tensor_val = [] + + def forward(self, x): + self.tensor_val.append(x.clone()) + return x + + @torch.jit.export + def calculate_qparams(self): + raise Exception( # noqa: TRY002 + "calculate_qparams should not be called for RecordingObserver" + ) + + @torch.jit.export + def get_tensor_value(self): + return self.tensor_val + + +class NoopObserver(ObserverBase): + r""" + Observer that doesn't do anything and just passes its configuration to the + quantized module's ``.from_float()``. + + Primarily used for quantization to float16 which doesn't require determining + ranges. + + Args: + dtype: Quantized data type + custom_op_name: (temporary) specify this observer for an operator that doesn't require any observation + (Can be used in Graph Mode Passes for special case ops). + """ + + def __init__(self, dtype=torch.float16, custom_op_name="") -> None: + super().__init__(dtype=dtype, is_dynamic=False) + self.dtype = dtype + self.custom_op = custom_op_name + + def forward(self, x): + return x + + @torch.jit.export + def calculate_qparams(self): + raise Exception( # noqa: TRY002 + "calculate_qparams should not be called for NoopObserver" + ) + + +class ReuseInputObserver(ObserverBase): + r"""This observer is used when we want to reuse the observer from the operator + that produces the input Tensor, typically used for operators like reshape, e.g. + ``` + x0 = ... + x1 = x0.reshape() + ``` + if we configure x0 to be observed by some observer, let's say MinMaxObserver, + and reshape is configured with ReuseInputObserver, we'll reuse the observer instance + for x0 for x1 (output of reshape). If x0 is not observed, we also won't observe x1. + + Note: this is only enabled in FX Graph Mode Quantization + """ + + def __init__(self) -> None: + super().__init__(torch.quint8, is_dynamic=False) + + def forward(self, x): + return x + + @torch.jit.export + def calculate_qparams(self): + raise Exception( # noqa: TRY002 + "calculate_qparams should not be called for ReuseInputObserver" + ) + + +""" +# Experimental Affine Quantization Feature START +We plan to merge the following with torchao repo after we move pt2e flow to torchao +copied from https://github.com/pytorch/ao/blob/main/torchao/quantization/observer.py +""" +from dataclasses import dataclass +from enum import Enum, auto + + +class MappingType(Enum): + """How floating point number is mapped to integer number + + symmetric mapping means floating point range is symmetrically mapped to integer range + let's say we have floating point range (-3.5, 10.2) and integer range (-8, 7) (int4) + we'll use (-10.2, 10.2) as the range for floating point and map that to (-8, 7) + e.g. scale = (10.2 - (-10.2)) / (7 - (-8)) + + SYMMETRIC_NO_CLIPPING_ERR is a variant of symmetric mapping, where the scale is the max of smin + and smax, where smin = min_val_neg / quant_min, and smax = max_val_pos / quant_max. By calculating + smin and smax individually, there can be less round error on negative values, and no out-of-range + of all floating point values. + + asymmetric mapping means we just directly map the floating point range to integer range, + for the above example, we will map (-3.5, 10.2) to (-8, 7) and calculate quantization parameter + based on this mapping + e.g. scale = (10.2 - (-3.5)) / (7 - (-8)) + """ + + SYMMETRIC = auto() + SYMMETRIC_NO_CLIPPING_ERR = auto() + ASYMMETRIC = auto() + + +class ZeroPointDomain(Enum): + """Enum that indicate whether zero_point is in integer domain or floating point domain + + integer domain: quantized_val = (float_val / scale) (integer) + zero_point (integer) + float domain: quantized_val = (float_val - (zero_point (float) - scale * mid_point)) / scale + none domain: quantized_val = (float_val / scale) + """ + + INT = auto() + FLOAT = auto() + NONE = auto() + + +class TorchAODType(Enum): + """ + Placeholder for dtypes that do not exist in PyTorch core yet. + """ + + # torch.int1 to torch.int7 will be added to PyTorch 2.6 + # These will remain here for BC with older PyTorch versions + INT1 = auto() + INT2 = auto() + INT3 = auto() + INT4 = auto() + INT5 = auto() + INT6 = auto() + INT7 = auto() + + +@dataclass(frozen=True) +class Granularity: + """ + Base class for representing the granularity of quantization. + + This class serves as a parent for specific granularity types used in + quantization operations, such as per-tensor or per-axis quantization. + """ + + +@dataclass(frozen=True) +class PerBlock(Granularity): + """ + Represents per-block granularity in quantization. See + :func:`~torchao.quantization.quant_primitives.quantize_affine` for docs for + `block_size` + + Attributes: + block_size (Tuple[int, ...]): The size of each quantization group + """ + + block_size: tuple[int, ...] + + +@dataclass(frozen=True) +class PerTensor(Granularity): + """ + Represents per-tensor granularity in quantization. + + This granularity type calculates the quantization parameters + based off the entire tensor. + + """ + + +@dataclass(frozen=True) +class PerAxis(Granularity): + """ + Represents per-axis granularity in quantization. + + This granularity type calculates different quantization parameters + along a specified axis of the tensor. + + For example if the input tensor is shape [8, 16] and axis=0, then + the quantization parameters are calculated for each row of the tensor. + Giving a total of 8 quantization parameters. + + Attributes: + axis (int): The axis along which reduction is performed. + """ + + axis: int + + +@dataclass(frozen=True) +class PerGroup(Granularity): + """ + Represents per-channel group granularity in quantization. + + This granularity type calculates different quantization parameters + for each group of elements. + + For example if the input tensor is shape [8, 16], and the group size is 4, then + the input tensor is reshaped to [64, 4] + quantization parameters are calculated for each group of 4 elements, + giving a total of 64 quantization parameters. + + Attributes: + group_size (int): The size of each quantization group + + """ + + group_size: int + + +class PerRow(Granularity): + """ + Represents row-wise granularity in quantization. + + This is a special case of per-axis quantization and is unique to Float8 matmuls + where the input is quantized with a block_size of (1, ..., input.shape[-1]). And the weight + is quantized with a block_size of (1, weight.shape[1]). + """ + + +class PerToken(Granularity): + """ + Represents per-token granularity in quantization. + + This granularity type calculates a different set of quantization parameters + for each token, which is represented as the last dimension of the tensor. + + For example, if the input tensor has shape [2, 3, 4], then there are 6 tokens + with 4 elements each, and we will calculate 6 sets of quantization parameters, + one for each token. + + If the input tensor has only two dimensions, e.g. [8, 16], then this is + equivalent to `PerAxis(axis=0)`, which yields 8 sets of quantization parameters. + """ + + +def get_block_size( + input_shape: tuple[int, ...], granularity: Granularity +) -> tuple[int, ...]: + """Get the block size based on the input shape and granularity type. + + Args: + input_shape: The input tensor shape possibly more than 2 dimensions + granularity: The granularity type of the quantization + """ + assert isinstance(granularity, Granularity), ( + "Please provide an instance of Granularity, not subclass of it" + ) + if isinstance(granularity, PerTensor): + return input_shape + elif isinstance(granularity, PerAxis): + block_size = list(input_shape) + block_size[granularity.axis] = 1 + return tuple(block_size) + elif isinstance(granularity, PerRow): + return (1,) * (len(input_shape) - 1) + (input_shape[-1],) + elif isinstance(granularity, PerGroup): + assert len(input_shape) == 2, ( + f"Expecting input shape dim to be 2 for per group quantization, gotinput shape: {input_shape}" + ) + return (1, granularity.group_size) + elif isinstance(granularity, PerToken): + block_size = [1] * len(input_shape) + block_size[-1] = input_shape[-1] + return tuple(block_size) + raise ValueError(f"Unsupported Granularity: {granularity}") + + +class AffineQuantizedObserverBase(ABC, torch.nn.Module): + """Observer module for affine quantization (https://github.com/pytorch/ao/tree/main/torchao/quantization#affine-quantization) + + Args: + `granularity` and `block_size`: The granularity of the quantization, + must specify at least one, if both are specified `block_size` takes precedence + Current supported granularity type are `PerTensor` and `PerAxis` + other args: please see `:class:torchao.dtypes.AffineQuantizedTensor` + """ + + with_args = classmethod(_with_args) + + def __init__( + self, + mapping_type: MappingType, + target_dtype: torch.dtype, + granularity: Granularity, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.INT, + # there could be some extra args that's ignored + **kwargs, + ): + super().__init__() + assert granularity is not None, "granularity is None" + + self.mapping_type = mapping_type + self.target_dtype = target_dtype + self.granularity = granularity + self.quant_min = quant_min + self.quant_max = quant_max + self.eps = eps + self.scale_dtype = scale_dtype + self.zero_point_dtype = zero_point_dtype + self.preserve_zero = preserve_zero + self.zero_point_domain = zero_point_domain + # populatd during forward + self.block_size = None + self.original_dtype = None + + @abstractmethod + def forward(self, input: torch.Tensor) -> torch.Tensor: + """forward function should take the input tensor + and updates internal stats and return the original input Tensor + """ + + @abstractmethod + def calculate_qparams(self) -> tuple[torch.Tensor, torch.Tensor]: + """Calculate quantization parameter based on the stats attached to the observer module + and returns a tuple of scale and zero_point Tensor + """ + + def convert(self, model: torch.fx.GraphModule, observer_node: Node): + """ + Converts the observer node in the graph into its quantized representation + + Args: + model: graph module to conver the observer node in + observer_node: the observer node to convert + """ + + from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 + + if not TORCH_VERSION_AT_LEAST_2_5: + raise NotImplementedError( + "convert for AffineQuantization is not implemented for pytorch version earlier than 2.5, please upgrade your pytorch to 2.5+." + ) + + from torchao.quantization.pt2e.utils import create_getattr_from_value + + with model.graph.inserting_before(observer_node): + assert self.block_size is not None, "Expecting block_size to be populated" + assert self.original_dtype is not None, ( + "Expecting original_dtype to be populated" + ) + if hasattr(self, "is_dynamic") and self.is_dynamic: + choose_qparams_affine = model.graph.call_function( + torch.ops.torchao.choose_qparams_affine, + ( + observer_node.args[0], + self.mapping_type.name, + self.block_size, + self.target_dtype, + self.quant_min, + self.quant_max, + self.eps, + self.scale_dtype, + self.zero_point_dtype, + self.preserve_zero, + self.zero_point_domain.name, + ), + ) + scale_node = model.graph.call_function( + operator.getitem, (choose_qparams_affine, 0) + ) + zero_point_node = model.graph.call_function( + operator.getitem, (choose_qparams_affine, 1) + ) + else: + scale, zero_point = self.calculate_qparams() + scale_node = create_getattr_from_value( + model, model.graph, "_scale", scale + ) + zero_point_node = create_getattr_from_value( + model, model.graph, "_zero_point", zero_point + ) + + q_node = model.graph.call_function( + torch.ops.torchao.quantize_affine, + ( + observer_node.args[0], + self.block_size, + scale_node, + zero_point_node, + self.target_dtype, + self.quant_min, + self.quant_max, + self.zero_point_domain.name, + ), + {}, + ) + dq_node = model.graph.call_function( + torch.ops.torchao.dequantize_affine, + ( + q_node, + self.block_size, + scale_node, + zero_point_node, + self.target_dtype, + self.quant_min, + self.quant_max, + self.zero_point_domain.name, + ), + {"output_dtype": self.original_dtype}, + ) + observer_node.replace_all_uses_with(dq_node) + model.graph.erase_node(observer_node) + + +def _is_observer_script_module(mod, obs_type_name): + """Returns true if given mod is an instance of Observer script module.""" + if isinstance(mod, torch.jit.RecursiveScriptModule): + # qualified name looks like '__torch__.torchao.quantization.pt2e.observer.___torch_mangle_2.MinMaxObserver' + suffix = mod._c.qualified_name.split(".", 1)[1] + name = re.sub(r"\.___torch_mangle_\d+", "", suffix) + return obs_type_name in name + return False + + +# Experimental Affine Quantization Feature END + + +def _is_activation_post_process(module): + return isinstance( + module, + ( + torchao.quantization.pt2e.ObserverBase, + torchao.quantization.pt2e.FakeQuantizeBase, + AffineQuantizedObserverBase, + ), + ) or _is_observer_script_module(module, "quantization.observer") + + +def _is_per_channel_script_obs_instance(module): + if isinstance(module, torch.jit.RecursiveScriptModule): + return _is_observer_script_module( + module, "quantization.observer.PerChannelMinMaxObserver" + ) or _is_observer_script_module( + module, "quantization.observer.MovingAveragePerChannelMinMaxObserver" + ) + return False + + +def get_observer_state_dict(mod): + r""" + Returns the state dict corresponding to the observer stats. + Traverse the model state_dict and extract out the stats. + """ + od = OrderedDict() + if isinstance(mod, torch.jit.RecursiveScriptModule): + for k, v in mod.state_dict().items(): + if "observer" in k: + od[k] = v + else: + # path for GraphModule and nn.Module (eager mode) + for k, v in mod.state_dict().items(): + if "activation_post_process" in k: + od[k] = v + od._metadata = mod.state_dict()._metadata # type: ignore[attr-defined] + return od + + +def load_observer_state_dict(mod, obs_dict): + r""" + Given input model and a state_dict containing model observer stats, + load the stats back into the model. The observer state_dict can be saved + using torchao.quantization.pt2e.get_observer_state_dict + """ + missing_keys: list[str] = [] + unexpected_keys: list[str] = [] + for name, module in mod.named_modules(): + prefix = name + "." + if _is_activation_post_process(module): + if _is_per_channel_script_obs_instance(module): + # For per-channel observers we need to call a custom load_from_state_dict to resize the tensor. + # However this is not called when the module is scripted and we end up calling the default one in module.py + module._load_from_state_dict_script( + obs_dict, prefix, {}, True, missing_keys, unexpected_keys, [] + ) + else: + module._load_from_state_dict( + obs_dict, prefix, {}, False, missing_keys, unexpected_keys, [] + ) + for k in missing_keys: + if "observer" in k or "activation_post_process" in k: + raise Exception( # noqa: TRY002 + f"Missing keys for observer {k} in state_dict" + ) + for k in unexpected_keys: + if "observer" in k or "activation_post_process" in k: + raise Exception( # noqa: TRY002 + f"Unexpected keys for observer {k} in state_dict" + ) + + +# Restrict activations to be in the range (0,127) +default_observer = MinMaxObserver.with_args(quant_min=0, quant_max=127) +""" +Default observer for static quantization, usually used for debugging. +""" + +default_placeholder_observer = PlaceholderObserver +""" +Default placeholder observer, usually used for quantization to torch.float16. +""" + +default_debug_observer = RecordingObserver +""" +Default debug-only observer. +""" + +default_weight_observer = MinMaxObserver.with_args( + dtype=torch.qint8, qscheme=torch.per_tensor_symmetric +) +""" +Default weight observer. +""" + +weight_observer_range_neg_127_to_127 = MinMaxObserver.with_args( + dtype=torch.qint8, + qscheme=torch.per_tensor_symmetric, + quant_min=-127, + quant_max=127, + eps=2**-12, +) +""" +Symmetric weight observer with the 8-bit values restricted to [-127, +127], excluding -128. +""" + +default_histogram_observer = HistogramObserver.with_args(quant_min=0, quant_max=127) +""" +Default histogram observer, usually used for PTQ. +""" + +default_per_channel_weight_observer = PerChannelMinMaxObserver.with_args( + dtype=torch.qint8, qscheme=torch.per_channel_symmetric +) +""" +Default per-channel weight observer, usually used on backends where per-channel +weight quantization is supported, such as `fbgemm`. +""" + +per_channel_weight_observer_range_neg_127_to_127 = PerChannelMinMaxObserver.with_args( + dtype=torch.qint8, + qscheme=torch.per_channel_symmetric, + quant_min=-127, + quant_max=127, + eps=2**-12, +) +""" +Per-channel, symmetric weight observer with the 8-bit values restricted to [-127, +127], excluding -128. +""" + +default_dynamic_quant_observer = PlaceholderObserver.with_args( + dtype=torch.quint8, + quant_min=0, + quant_max=255, + is_dynamic=True, +) +""" +Default observer for dynamic quantization. +""" + +default_float_qparams_observer = PerChannelMinMaxObserver.with_args( + dtype=torch.quint8, qscheme=torch.per_channel_affine_float_qparams, ch_axis=0 +) +""" +Default observer for a floating point zero-point. +""" + +default_float_qparams_observer_4bit = PerChannelMinMaxObserver.with_args( + dtype=torch.quint4x2, qscheme=torch.per_channel_affine_float_qparams, ch_axis=0 +) +""" +Default observer for a floating point zero-point and 4 bit activations. +""" + +# TODO(future PR): remove these defaults and enforce activation functions +# to explicitly specify their output range +default_fixed_qparams_range_neg1to1_observer = FixedQParamsObserver.with_args( + scale=2.0 / 256.0, zero_point=128, dtype=torch.quint8, quant_min=0, quant_max=255 +) +default_fixed_qparams_range_0to1_observer = FixedQParamsObserver.with_args( + scale=1.0 / 256.0, zero_point=0, dtype=torch.quint8, quant_min=0, quant_max=255 +) +# TODO: the following 2 variables are kept for backwards compatibility; remove after a few releases +default_symmetric_fixed_qparams_observer = default_fixed_qparams_range_neg1to1_observer +default_affine_fixed_qparams_observer = default_fixed_qparams_range_0to1_observer + +""" +Default observers for fixed qparams operations. +""" + +default_reuse_input_observer = ReuseInputObserver +""" +Default observer for operators like reshape that reuses the observer of input to +the operator +""" diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/prepare.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/prepare.py new file mode 100644 index 0000000000000000000000000000000000000000..22542c09e23efa0822b55e1088a7d9bd7289c57f --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/prepare.py @@ -0,0 +1,668 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import copy +from dataclasses import asdict +from typing import Any, Optional, Union + +import torch +from torch._subclasses import FakeTensor +from torch.ao.quantization import QConfigMapping +from torch.ao.quantization.fx.custom_config import PrepareCustomConfig +from torch.ao.quantization.fx.prepare import ( + _insert_obs_or_fq, + _save_state, +) +from torch.ao.quantization.qconfig import QConfigAny +from torch.fx import Graph, GraphModule, Node +from torch.fx.node import Argument + +from torchao.quantization.pt2e import ( + CUSTOM_KEY, + NUMERIC_DEBUG_HANDLE_KEY, + DerivedObserverOrFakeQuantize, + ObserverOrFakeQuantize, +) +from torchao.quantization.pt2e.fake_quantize import ( + FixedQParamsFakeQuantize, +) +from torchao.quantization.pt2e.observer import ( + FixedQParamsObserver, + PartialWrapper, + _is_activation_post_process, +) +from torchao.quantization.pt2e.quantizer import ( + DerivedQuantizationSpec, + EdgeOrNode, + FixedQParamsQuantizationSpec, + QuantizationSpec, + QuantizationSpecBase, + SharedQuantizationSpec, +) + +# TODO: make pt2e folder private? +__all__ = [ + "prepare", +] + + +def _is_activation_post_process_node( + node: Node, named_modules: dict[str, torch.nn.Module] +) -> bool: + return ( + isinstance(node, torch.fx.Node) + and node.op == "call_module" + and _is_activation_post_process(named_modules[str(node.target)]) + ) + + +def _get_observer_kwargs( + quant_spec: Union[QuantizationSpec, FixedQParamsQuantizationSpec], +): + kwargs_dict = asdict(quant_spec) + return copy.deepcopy(kwargs_dict) + + +def _create_obs_or_fq_from_qspec( + quantization_spec: Optional[QuantizationSpecBase], + obs_or_fq_map: dict[EdgeOrNode, ObserverOrFakeQuantize], + is_qat: bool, +): + """Create observer or fake quantize objects based on quantization spec + + Args: + quantization_spec: used to store parameters to create the observer or fake quantizer + obs_or_fq_map: this is a map from edge/output to the corresponding observer/fake_quant + instance, it may be reused for different edge/output depending on configuration + """ + if quantization_spec is None: + return None + if isinstance(quantization_spec, SharedQuantizationSpec): + edge_or_node = quantization_spec.edge_or_node + assert edge_or_node in obs_or_fq_map, ( + "please make sure only refer to edge or node that has " + f"observer/fake_quant inserted: '{edge_or_node}' not in\n{obs_or_fq_map.keys()}" + ) + return obs_or_fq_map[edge_or_node] + elif isinstance(quantization_spec, DerivedQuantizationSpec): + # can't use asdict, so not calling get_observer_kwargs here + kwargs = { + "dtype": quantization_spec.dtype, + "derive_qparams_fn": quantization_spec.derive_qparams_fn, + "quant_min": quantization_spec.quant_min, + "quant_max": quantization_spec.quant_max, + "qscheme": quantization_spec.qscheme, + "ch_axis": quantization_spec.ch_axis, + } + edge_or_nodes = quantization_spec.derived_from + obs_or_fqs = [obs_or_fq_map[k] for k in edge_or_nodes] + kwargs["obs_or_fqs"] = obs_or_fqs + return DerivedObserverOrFakeQuantize.with_args(**kwargs)() + elif isinstance(quantization_spec, FixedQParamsQuantizationSpec): + kwargs = _get_observer_kwargs(quantization_spec) + observer_ctr = FixedQParamsObserver.with_args(**kwargs) + if is_qat: + return FixedQParamsFakeQuantize.with_args(observer=observer_ctr)() + else: + return observer_ctr() + + assert isinstance(quantization_spec, QuantizationSpec), ( + f"Expected QuantizationSpec got: {quantization_spec}" + ) + observer_or_fake_quant_ctr = quantization_spec.observer_or_fake_quant_ctr + kwargs = _get_observer_kwargs(quantization_spec) + kwargs.pop("observer_or_fake_quant_ctr") + # we will remove is_dynamic from QuantizationSpec because + # it seems that dynamic range quantization + obs_or_fq_class = observer_or_fake_quant_ctr + if isinstance(observer_or_fake_quant_ctr, PartialWrapper): + obs_or_fq_class = observer_or_fake_quant_ctr.p.func # type: ignore[union-attr, assignment] + if "PerChannel" not in obs_or_fq_class.__name__: # type: ignore[operator, union-attr] + kwargs.pop("ch_axis") + return observer_or_fake_quant_ctr.with_args(**kwargs)() + + +def _find_root_edge_or_node( + edge_or_node: EdgeOrNode, shared_with_map: dict[EdgeOrNode, EdgeOrNode] +) -> EdgeOrNode: + """Find the root node for the sharing tree + Args: + edge_or_node: edge/node that we want to find the root + shared_with_map: each edge/node points to the parent, the root node will points to itself + + Returns: + root edge/node + """ + parent = shared_with_map[edge_or_node] + if parent == edge_or_node: + return edge_or_node + root = _find_root_edge_or_node(parent, shared_with_map) + # path compression + shared_with_map[edge_or_node] = root + return root + + +def _union( + parent: EdgeOrNode, + child: EdgeOrNode, + shared_with_map: dict[EdgeOrNode, EdgeOrNode], +) -> None: + """Merge the subtree for `child` with `parent`, the order is important here""" + root_parent = _find_root_edge_or_node(parent, shared_with_map) + root_child = _find_root_edge_or_node(child, shared_with_map) + # union the two trees by pointing the root of child to root of parent + shared_with_map[root_child] = root_parent + + +def _update_shared_with( + child: EdgeOrNode, + qspec: QuantizationSpecBase, + shared_with_map: dict[EdgeOrNode, EdgeOrNode], +): + """Update the `shared_with_map` based on the qspec, this applies the `SharedQuantizationSpec` + configuration and established the relationship between `edge_or_node` with the edge/node that it + is pointing to, we'll use this information in the end to get the group id + """ + if isinstance(qspec, SharedQuantizationSpec): + parent = qspec.edge_or_node + # we point from edge_or_node to the node that it is sharing_with, e.g. + # qspec for a = SharedQuantizationSpec(b) means `a` points to `b` + _union(parent, child, shared_with_map) + + +def _unwrap_shared_qspec( + qspec: QuantizationSpecBase, + edge_or_node_to_qspec: dict[EdgeOrNode, QuantizationSpecBase], + shared_with_map: dict[EdgeOrNode, EdgeOrNode], +) -> QuantizationSpecBase: + """Unwraps qspec to get the final root qspec (non SharedQuantizationSpec) + if qspec is SharedQuantizationSpec + (1). tries to find the root edge or node for the node that the qspec points to + (2). recursively find the root qspec based on the qspec for the root node + """ + if isinstance(qspec, SharedQuantizationSpec): + sharing_with = qspec.edge_or_node + root = _find_root_edge_or_node(sharing_with, shared_with_map) + qspec = edge_or_node_to_qspec[root] + return _unwrap_shared_qspec(qspec, edge_or_node_to_qspec, shared_with_map) + return qspec + + +def _has_same_attr( + qspec_a: QuantizationSpecBase, qspec_b: QuantizationSpecBase, attr_name: str +): + return ( + hasattr(qspec_a, attr_name) + and hasattr(qspec_b, attr_name) + and getattr(qspec_a, attr_name) == getattr(qspec_b, attr_name) + ) or (not hasattr(qspec_a, attr_name) and not hasattr(qspec_b, attr_name)) + + +def _get_edge_or_node_to_qspec( + model: torch.fx.GraphModule, +) -> dict[EdgeOrNode, QuantizationSpecBase]: + """Get a map from EdgeOrNode to quantization spec based on annotations on the nodes""" + edge_or_node_to_qspec: dict[EdgeOrNode, QuantizationSpecBase] = {} + for n in model.graph.nodes: + if hasattr(n, "meta") and "quantization_annotation" in n.meta: + qa = n.meta["quantization_annotation"] + for input_to_n, qspec in qa.input_qspec_map.items(): + input_edge = (input_to_n, n) + edge_or_node_to_qspec[input_edge] = qspec + if qa.output_qspec is not None: + output_node = n + qspec = qa.output_qspec + edge_or_node_to_qspec[output_node] = qspec + return edge_or_node_to_qspec + + +def _union_input_edge_with( + input_edge, + input_edge_root_qspec, + edge_or_node, + edge_or_node_to_qspec, + shared_with_map, +): + """Union input edge with another edge or node, used in implicit sharing to point the current input + edge to other user edges of the producer node, or the output of producer node since these are + referring to the same Tensor + """ + root_qspec = None + if edge_or_node in edge_or_node_to_qspec: + qspec = edge_or_node_to_qspec[edge_or_node] + root_qspec = _unwrap_shared_qspec(qspec, edge_or_node_to_qspec, shared_with_map) + # TODO: add assertions for types of root qspecs + if root_qspec is not None and all( + _has_same_attr(root_qspec, input_edge_root_qspec, attr) + for attr in [ + "dtype", + "is_dynamic", + "quant_min", + "quant_max", + "qscheme", + "ch_axis", + "scale", + "zero_point", + ] + ): + # the input arg to the node should reuse the existing output observer for arg + # since dtype is the same (we may want to extend this to be a more strict check + # in the future) + # so we point from `input_edge` to `arg` (output of the argument) + _union(edge_or_node, input_edge, shared_with_map) + + +def _get_edge_or_node_to_group_id( + edge_or_node_to_qspec: dict[EdgeOrNode, QuantizationSpecBase], +) -> dict[EdgeOrNode, int]: + """Map from edge/node to the group ID, generated from quantization annotations, + edge/node with the same group ID should use the same observer/fake_quant instance + + This is applying SharedQuantizationSpec configuration and map each edge/node to a group + There is another implicit sharing that's built in the quantization, when we have the following: + * op1 -> op2 + * output of op1: int8_qspec + * (op1 -> op2) input edge: int8_qspec + we'll assume sharing between the output of op1 and input of (op1 -> op2) since these are the same Tensor. + + Figuring out the correct group ID for all edge/node is a standard union find problem: + https://www.geeksforgeeks.org/introduction-to-disjoint-set-data-structure-or-union-find-algorithm/ + + Args: + edge_or_node_to_qspec: Dictionary from edge_or_node to the qspec, derived from annotations + Returns: + edge_or_node_to_group_id: Dictionary from edge_or_node to group_id (int), all edge or node that + belongs to the same group should have the same id + + Example: + op2 -> cat1 -> cat2 + op1 / / + op3 + edge_or_node_to_qspec: { + op1: int8_qspec, + op2: int8_qspec, + (op1, cat1): int8_qspc, + (op2, cat1): SharedQuantizationSpec((op1, cat1)), + cat1: SharedQuantizationSpec((op1, cat1)), + (op3, cat2): int8_qspec, + (cat1, cat2): SharedQuantizationSpec((op3, cat2)), + cat2: SharedQuantizationSpec((op3, cat2)), + } + + edge_or_node_to_group_id = _get_edge_or_node_to_group_id(edge_or_node_to_qspec) + edge_or_node_to_group_id: { + op1: 1, + op2: 1, + (op1, cat1): 1, + (op2, cat1): 1, + cat1: 1, + (op3, cat2): 1, + (cat1, cat2): 1, + cat2: 1, + } + # everything are in the same group because (cat1) and (cat1, cat2) are implicitly shared, which + # connects the two sharing group around cat1 and cat2 op due to transitive sharing + """ + # means the observer of key should be shared with observer with value, by default it will + # be shared with itself + shared_with_map: dict[EdgeOrNode, EdgeOrNode] = { + k: k for k in edge_or_node_to_qspec.keys() + } + for edge_or_node, qspec in edge_or_node_to_qspec.items(): + if isinstance(edge_or_node, torch.fx.Node): + output_node = edge_or_node + _update_shared_with(output_node, qspec, shared_with_map) + else: + input_edge = edge_or_node + input_edge_root_qspec = _unwrap_shared_qspec( + qspec, edge_or_node_to_qspec, shared_with_map + ) + + assert isinstance(input_edge, tuple) + arg, n = input_edge + if n.meta["quantization_annotation"].allow_implicit_sharing: + # NOTE: the order is important here, we first share with other users and then share with previous + # output because the reverse order could cause circular dependency + # e.g node1 -> node2 + # \ -> node3 + # when processing (node1, node2), if we first point (node1, node2) to node1 + # Step 1. shared_map = {(node1, node2): node1} + # Step 2. after that, we point the (node1, node2) to its other user (node1, node3) , + # which means shared_map = {(node1, node2): node1, node1: (node1, node3)} + # because we will point the root of (node1, node2) (in this case node1) to the root of (node1, node3) + # Step 3. and when we process (node1, node3), it can try to point to node1 as well, then we'll + # have a circular dependency + # the following order works around this issue, but this does not allow arbitrary configuration + # of sharing so it might break in a different case in the future, when it breaks + # quantizer writer can check the notes here to debug the issue + + # sharing with other users of the producer node + # (arg, user) + if not isinstance(arg, Node) or not isinstance(n, Node): + raise Exception( # noqa: TRY002 + f"Expected input_edge to have type Tuple[Node, Node], but got: {arg, n}" + ) + for user in arg.users: + if user is n: + continue + arg_to_user_edge = (arg, user) + _union_input_edge_with( + input_edge, + input_edge_root_qspec, + arg_to_user_edge, + edge_or_node_to_qspec, + shared_with_map, + ) + + # sharing with output of producer node + _union_input_edge_with( + input_edge, + input_edge_root_qspec, + arg, + edge_or_node_to_qspec, + shared_with_map, + ) + + _update_shared_with(input_edge, qspec, shared_with_map) + + # now that we get the sharing relations between all edges and nodes, we can assingn group ids + cur_group_id = 0 + edge_or_node_to_group_id: dict[EdgeOrNode, int] = {} + for edge_or_node in shared_with_map.keys(): + root = _find_root_edge_or_node(edge_or_node, shared_with_map) + if root not in edge_or_node_to_group_id: + edge_or_node_to_group_id[root] = cur_group_id + cur_group_id += 1 + edge_or_node_to_group_id[edge_or_node] = edge_or_node_to_group_id[root] + + return edge_or_node_to_group_id + + +def _get_obs_or_fq_map( + edge_or_node_to_group_id: dict[EdgeOrNode, int], + edge_or_node_to_qspec: dict[EdgeOrNode, QuantizationSpecBase], + is_qat: bool, +) -> dict[EdgeOrNode, ObserverOrFakeQuantize]: + """Generates the EdgeOrNode to observer/fake_quant instances + Makes sure that for EdgeOrNode that has the same group_id should have the same observer or fake quant + instances + """ + obs_or_fq_map: dict[EdgeOrNode, ObserverOrFakeQuantize] = {} + group_id_to_obs_or_fq: dict[int, ObserverOrFakeQuantize] = {} + for edge_or_node, qspec in edge_or_node_to_qspec.items(): + group_id = edge_or_node_to_group_id[edge_or_node] + if group_id not in group_id_to_obs_or_fq: + # TODO: maybe edge_or_node_to_qspec should be edge_or_node_to_root_qspec, this will simplify + # the implementation for _create_obs_or_fq_from_qspec + group_id_to_obs_or_fq[group_id] = _create_obs_or_fq_from_qspec( + qspec, obs_or_fq_map, is_qat + ) + obs_or_fq_map[edge_or_node] = group_id_to_obs_or_fq[group_id] + return obs_or_fq_map + + +def _maybe_insert_input_observer_for_arg_or_kwarg( + node: Union[Node, Any], + arg: Argument, + qconfig: QConfigAny, + model: torch.nn.Module, + named_modules: dict[str, torch.nn.Module], + obs_or_fq_map: dict[EdgeOrNode, ObserverOrFakeQuantize], + is_qat: bool, +) -> Argument: + """ + Given a `node` and an `arg`, inserts an input observer between + `node` and `arg` if necessary. + """ + # for ops such as torch.cat([x0, x1]), + # traverse through the list + if isinstance(arg, (list, tuple)): + new_arg_to_return = [] + for inner_arg in arg: + new_inner_arg = _maybe_insert_input_observer_for_arg_or_kwarg( + node, + inner_arg, + qconfig, + model, + named_modules, + obs_or_fq_map, + is_qat, + ) + new_arg_to_return.append(new_inner_arg) + return type(arg)(new_arg_to_return) + + if not isinstance(arg, Node): + return arg + assert isinstance(arg, Node) + # default (no observer) + new_arg = arg + + # find the original `arg` node to the current node, skipping inserted observer/fake_quant nodes + original_arg = arg + while _is_activation_post_process_node(original_arg, named_modules): + original_arg = original_arg.args[0] # type: ignore[assignment] + assert isinstance(original_arg, Node), ( + f"expect original argument to be a Node, but got: {type(original_arg)}" + ) + + input_edge = (original_arg, node) + if input_edge not in obs_or_fq_map: + return new_arg + # input_edge needs to be observed + input_edge_obs_or_fq = obs_or_fq_map[input_edge] + if input_edge_obs_or_fq is None: + return new_arg + + arg_as_output_obs_or_fq = obs_or_fq_map.get(original_arg, None) + # the arg is observed as the output and is using the same instance as the input_edge + # we'll reuse the inserted observer/fake_quant + if arg_as_output_obs_or_fq is not None and id(arg_as_output_obs_or_fq) == id( + input_edge_obs_or_fq + ): + return new_arg + + # otherwise, we'll insert a new observer/fake_quant node + + # skip inserting new observers if the same observer instance is inserted before for another user + # Example: + # conv1 -> obs1 -> existing_obs -> conv2 + # \ -> conv3 + # + # instead of inserting new observers we will have: + # conv1 -> obs1 -> existing_obs -> conv2 + # \ -> conv3 + for maybe_obs_node in arg.users.keys(): + if not _is_activation_post_process_node(maybe_obs_node, named_modules): + continue + maybe_obs_mod = named_modules[maybe_obs_node.target] # type: ignore[index] + if id(maybe_obs_mod) == id(input_edge_obs_or_fq): + return maybe_obs_node + + assert isinstance(model.graph, Graph) + new_arg = _insert_obs_or_fq( + arg, input_edge_obs_or_fq, model, named_modules, model.graph + ) + return new_arg + + +def _maybe_insert_input_observers_for_node( + node: Node, + qconfig: QConfigAny, + model: torch.nn.Module, + named_modules: dict[str, torch.nn.Module], + obs_or_fq_map: dict[EdgeOrNode, ObserverOrFakeQuantize], + is_qat: bool, +) -> None: + """ + If needed, inserts observers to the input args and kwargs of `node`. + Note: modifies `node` inplace. + + For example, if cur_node needs an observer after prev_node, we change from + + prev_node -> cur_node + + To + + prev_node -> obs -> cur_node + + """ + # Look through every input arg. If that arg's target dtype does not + # match the current node's target dtype, insert an observer. + new_args = [] + for arg in node.args: + new_arg = _maybe_insert_input_observer_for_arg_or_kwarg( + node, + arg, + qconfig, + model, + named_modules, + obs_or_fq_map, + is_qat, + ) + new_args.append(new_arg) + + # Clone has a memory_format kwarg, zeros_like has a pin_memory kwarg, and + # gelu has a has an approximate kwarg that persist in exported graph. + # This is just a work around for these. + assert ( + node.target == torch.ops.aten.clone.default + or node.target == torch.ops.aten.zeros_like.default + or node.target == torch.ops.aten.gelu.default + or len(node.kwargs) == 0 + ), " expecting kwargs for aten op IR to be empty" + + # assign the new args to the node, inplace + node.args = tuple(new_args) + + +def _maybe_insert_output_observer_for_node( + node: Node, + model: torch.nn.Module, + named_modules: dict[str, torch.nn.Module], + graph: Graph, + obs_or_fq_map: dict[EdgeOrNode, ObserverOrFakeQuantize], + is_qat: bool, +) -> Optional[Node]: + if node in obs_or_fq_map: + output_act_obs_or_fq = obs_or_fq_map[node] + new_output = _insert_obs_or_fq( + node, output_act_obs_or_fq, model, named_modules, graph + ) + # propagate numeric debug handle from original node to observer/fake_quant node + if ( + isinstance(node, Node) + and isinstance(new_output, Node) + and CUSTOM_KEY in node.meta + and NUMERIC_DEBUG_HANDLE_KEY in node.meta[CUSTOM_KEY] + ): + if CUSTOM_KEY not in new_output.meta: + new_output.meta[CUSTOM_KEY] = {} + new_output.meta[CUSTOM_KEY][NUMERIC_DEBUG_HANDLE_KEY] = node.meta[ + CUSTOM_KEY + ][NUMERIC_DEBUG_HANDLE_KEY] + return new_output + return None + + +def _maybe_insert_input_and_output_observers_for_node( + node: Node, + model: torch.fx.GraphModule, + obs_or_fq_map: dict[EdgeOrNode, ObserverOrFakeQuantize], + is_qat: bool, +): + this_node_quantization_annotation = ( + node.meta["quantization_annotation"] + if "quantization_annotation" in node.meta + else None + ) + if this_node_quantization_annotation is None: + return + + named_modules = dict(model.named_modules(remove_duplicate=False)) + _maybe_insert_input_observers_for_node( + node, + None, # qconfig + model, + named_modules, + obs_or_fq_map, + is_qat, + ) + + output_is_a_tensor = "val" in node.meta and isinstance(node.meta["val"], FakeTensor) + if not output_is_a_tensor: + return + + # this returns the new observer node if it was needed + maybe_output_obs_node = _maybe_insert_output_observer_for_node( + node, model, named_modules, model.graph, obs_or_fq_map, is_qat + ) + + if maybe_output_obs_node is None: + return + # Update users of original node to use the output observer + # instead. For example, change + # + # next_node + # / + # cur_node -> obs + # + # to + # + # next_node + # / + # cur_node -> obs + # + # We need to save orig users before updating uses because + # the list of users will change as we update uses + orig_users = list(node.users.keys()) + for user_node in orig_users: + if user_node is maybe_output_obs_node: + continue + user_node.replace_input_with(node, maybe_output_obs_node) + + +def prepare( + model: GraphModule, + node_name_to_scope: dict[str, tuple[str, type]], + is_qat: bool, + obs_or_fq_callback=None, +) -> GraphModule: + # Since we are mutating the graph as we go, we iterate over the original + # nodes before observer insertion, instead of model.graph.nodes. + nodes_before_observation = list(model.graph.nodes) + + # At the high level we construct a map from EdgeOrNode to a observer_or_fake_quant instance + # all edge/nodes that belongs to the same group will use the same instance + # and when we insert observers we'll just query this map to get the correct observer_or_fake_quant + # instance + edge_or_node_to_qspec = _get_edge_or_node_to_qspec(model) + edge_or_node_to_group_id = _get_edge_or_node_to_group_id(edge_or_node_to_qspec) + obs_or_fq_map = _get_obs_or_fq_map( + edge_or_node_to_group_id, edge_or_node_to_qspec, is_qat + ) + if obs_or_fq_callback: + obs_or_fq_callback(model, obs_or_fq_map) + + for node in nodes_before_observation: + # TODO: simplify logic for inserting observers + _maybe_insert_input_and_output_observers_for_node( + node, model, obs_or_fq_map, is_qat + ) + + model = GraphModule(model, model.graph) + + _save_state( + model, + {}, # node_name_to_qconfig + node_name_to_scope, + PrepareCustomConfig(), + {}, # equalization_node_name_to_qconfig + QConfigMapping(), + is_qat, + set(), # observed_node_names + ) + return model diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/qat_utils.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/qat_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..565c4fcacc6bc13da204c67fdd0b88e50b756226 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/qat_utils.py @@ -0,0 +1,996 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import copy +import dataclasses +import itertools +import operator +from typing import TYPE_CHECKING, Any, Callable, Optional + +import torch +import torch.nn.functional as F +from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib # noqa: F401 +from torch.fx import Graph, GraphModule, Node +from torch.fx.subgraph_rewriter import ReplacedPatterns, replace_pattern_with_filters + +from torchao.quantization.pt2e.export_utils import WrapperModule +from torchao.quantization.pt2e.quantizer import ( + DerivedQuantizationSpec, + EdgeOrNode, + QuantizationSpecBase, + SharedQuantizationSpec, +) +from torchao.quantization.pt2e.utils import ( + _get_aten_graph_module_for_pattern, + _is_bn_node, + _is_conv_or_conv_transpose_node, + _is_conv_transpose_fn, + fold_bn_weights_into_conv_node, +) + +if TYPE_CHECKING: + from torch.fx.passes.utils.matcher_with_name_node_map_utils import InternalMatch + +__all__ = [] # type: ignore[var-annotated] + + +def _get_quantized_conv_bn_example_inputs_kwargs( + is_per_channel: bool, + has_bias: bool, + bias_is_quantized: bool, + is_cuda: bool, +) -> dict[str, Any]: + """ + Optional example inputs for quantized and folded conv-bn patterns + used in convert, expressed as kwargs. + """ + kwargs = {} + # Per tensor quantization uses literals to represent scale and zero + # point, so there is no need to include them here as kwargs + if is_per_channel: + kwargs["weight_scale"] = torch.tensor([1], dtype=torch.float) + kwargs["weight_zero_point"] = torch.tensor([0], dtype=torch.int) + if has_bias and bias_is_quantized: + kwargs["bias_scale"] = torch.tensor([1], dtype=torch.float) + kwargs["bias_zero_point"] = torch.tensor([0], dtype=torch.int) + if has_bias: + kwargs["conv_bias"] = torch.randn(1) + if is_cuda: + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + kwargs[k] = v.cuda() + return kwargs + + +def _get_conv_bn_pattern(conv_fn: Callable) -> Callable: + def _conv_bn_pattern( + x: torch.Tensor, + conv_weight: torch.Tensor, + conv_bias: torch.Tensor, + bn_weight: torch.Tensor, + bn_bias: torch.Tensor, + bn_running_mean: torch.Tensor, + bn_running_var: torch.Tensor, + ) -> torch.Tensor: + x = conv_fn(x, conv_weight, conv_bias) + x = F.batch_norm( + x, bn_running_mean, bn_running_var, bn_weight, bn_bias, training=True + ) + return x + + return WrapperModule(_conv_bn_pattern) + + +# TODO: merge this with the `no_conv_bias` case +def _get_qat_conv_bn_pattern(conv_fn: Callable) -> Callable: + def _qat_conv_bn_pattern( + x: torch.Tensor, + conv_weight: torch.Tensor, + conv_bias: torch.Tensor, + bn_weight: torch.Tensor, + bn_bias: torch.Tensor, + bn_running_mean: torch.Tensor, + bn_running_var: torch.Tensor, + ) -> torch.Tensor: + """ + Approximated method to fuse conv and bn. It requires only one forward pass. + conv_orig = conv / scale_factor where scale_factor = bn.weight / running_std. + This is based on `nniqat.ConvBn2d._forward_approximate`. + """ + # TODO: allow setting eps + bn_eps = 1e-5 + running_std = torch.sqrt(bn_running_var + bn_eps) + scale_factor = bn_weight / running_std + weight_shape = [1] * len(conv_weight.shape) + weight_in_channel_axis = 1 if _is_conv_transpose_fn(conv_fn) else 0 + weight_shape[weight_in_channel_axis] = -1 + bias_shape = [1] * len(conv_weight.shape) + bias_shape[1] = -1 + scaled_weight = conv_weight * scale_factor.reshape(weight_shape) + zero_bias = torch.zeros_like(conv_bias, dtype=x.dtype) + x = conv_fn(x, scaled_weight, zero_bias) + x = x / scale_factor.reshape(bias_shape) + x = x + conv_bias.reshape(bias_shape) + x = F.batch_norm( + x, + bn_running_mean, + bn_running_var, + bn_weight, + bn_bias, + training=True, + eps=bn_eps, + ) + return x + + return WrapperModule(_qat_conv_bn_pattern) + + +def _get_qat_conv_bn_pattern_no_conv_bias(conv_fn: Callable) -> Callable: + def _qat_conv_bn_pattern_no_conv_bias( + x: torch.Tensor, + conv_weight: torch.Tensor, + # Not used, only for matching convenience + conv_bias: torch.Tensor, + bn_weight: torch.Tensor, + bn_bias: torch.Tensor, + bn_running_mean: torch.Tensor, + bn_running_var: torch.Tensor, + ) -> torch.Tensor: + """ + Same as `_get_qat_conv_bn_pattern`, but handles the case with no conv bias. + """ + # TODO: allow setting eps + bn_eps = 1e-5 + running_std = torch.sqrt(bn_running_var + bn_eps) + scale_factor = bn_weight / running_std + weight_shape = [1] * len(conv_weight.shape) + weight_in_channel_axis = 1 if _is_conv_transpose_fn(conv_fn) else 0 + weight_shape[weight_in_channel_axis] = -1 + bias_shape = [1] * len(conv_weight.shape) + bias_shape[1] = -1 + scaled_weight = conv_weight * scale_factor.reshape(weight_shape) + x = conv_fn(x, scaled_weight, None) + x = x / scale_factor.reshape(bias_shape) + x = F.batch_norm( + x, + bn_running_mean, + bn_running_var, + bn_weight, + bn_bias, + training=True, + eps=bn_eps, + ) + return x + + return WrapperModule(_qat_conv_bn_pattern_no_conv_bias) + + +def _append_qdq(x, is_per_channel, is_bias, kwargs): + """ + Helper function to append q-dq ops after `x`, using dummy values for the qparams + and qmin/qmax. We use dummy values here because we match with `ignore_literals=True` + and will manually replace these values after subgraph rewriting. + + Return the dq node. + """ + # Dummy args to be passed into q-dq ops + per_channel_axis = 0 + scale_key = "bias_scale" if is_bias else "weight_scale" + zp_key = "bias_zero_point" if is_bias else "weight_zero_point" + scale = kwargs[scale_key] if is_per_channel else 1.0 + zp = kwargs[zp_key] if is_per_channel else 0 + qmin = -127 + qmax = 127 + dtype = torch.int8 + + qd = torch.ops.quantized_decomposed + if is_per_channel: + x = qd.quantize_per_channel(x, scale, zp, per_channel_axis, qmin, qmax, dtype) + x = qd.dequantize_per_channel(x, scale, zp, per_channel_axis, qmin, qmax, dtype) + else: + x = qd.quantize_per_tensor(x, scale, zp, qmin, qmax, dtype) + x = qd.dequantize_per_tensor(x, scale, zp, qmin, qmax, dtype) + return x + + +def _get_quantized_qat_conv_bn_pattern( + is_per_channel: bool, + has_bias: bool, + bias_is_quantized: bool, + conv_fn: Callable, + bn_is_training: bool, +) -> Callable: + """ + Return the quantized version of QAT conv + BN pattern. + This is based on `nniqat.ConvBn2d._forward_approximate`, + used in QAT convert. We first match this pattern and replace + it with the normal [conv - bn] pattern, then fold the BN + weights into conv. + """ + # TODO: allow setting eps + bn_eps = 1e-5 + + def _quantized_qat_conv_bn_pattern( + x: torch.Tensor, + conv_weight: torch.Tensor, + bn_weight: torch.Tensor, + bn_bias: torch.Tensor, + bn_running_mean: torch.Tensor, + bn_running_var: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + running_std = torch.sqrt(bn_running_var + bn_eps) + scale_factor = bn_weight / running_std + weight_shape = [1] * len(conv_weight.shape) + weight_shape[0] = -1 + bias_shape = [1] * len(conv_weight.shape) + bias_shape[1] = -1 + scaled_weight = conv_weight * scale_factor.reshape(weight_shape) + scaled_weight = _append_qdq( + scaled_weight, + is_per_channel, + is_bias=False, + kwargs=kwargs, + ) + if has_bias: + zero_bias = torch.zeros_like(kwargs["conv_bias"], dtype=x.dtype) + if bias_is_quantized: + zero_bias = _append_qdq( + zero_bias, + is_per_channel, + is_bias=True, + kwargs=kwargs, + ) + x = conv_fn(x, scaled_weight, zero_bias) + else: + x = conv_fn(x, scaled_weight, None) + x = x / scale_factor.reshape(bias_shape) + if has_bias: + x = x + kwargs["conv_bias"].reshape(bias_shape) + x = F.batch_norm( + x, + bn_running_mean, + bn_running_var, + bn_weight, + bn_bias, + training=bn_is_training, + eps=bn_eps, + ) + return x + + return WrapperModule(_quantized_qat_conv_bn_pattern) + + +def _get_folded_quantized_qat_conv_bn_pattern( + is_per_channel: bool, + has_bias: bool, + bias_is_quantized: bool, + conv_fn: Callable, + bn_is_training: bool, +) -> Callable: + """ + Quantized QAT conv - bn pattern with bn weights being folded into conv. + """ + # TODO: allow setting eps + bn_eps = 1e-5 + + def _folded_quantized_qat_conv_bn_pattern( + x: torch.Tensor, + conv_weight: torch.Tensor, + bn_weight: torch.Tensor, + bn_bias: torch.Tensor, + bn_running_mean: torch.Tensor, + bn_running_var: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + conv_weight = _append_qdq( + conv_weight, + is_per_channel, + is_bias=False, + kwargs=kwargs, + ) + if has_bias: + bias = kwargs["conv_bias"] + if bias_is_quantized: + bias = _append_qdq( + bias, + is_per_channel, + is_bias=True, + kwargs=kwargs, + ) + else: + bias = None + x = conv_fn(x, conv_weight, bias) + x = F.batch_norm( + x, + bn_running_mean, + bn_running_var, + bn_weight, + bn_bias, + training=bn_is_training, + eps=bn_eps, + ) + return x + + return WrapperModule(_folded_quantized_qat_conv_bn_pattern) + + +def _has_conv_bias_filter( + match: "InternalMatch", + original_graph: Graph, + pattern_graph: Graph, +) -> bool: + """ + Match filter for the subgraph rewriter that returns True if the conv node in + the original graph has bias. + """ + for n in match.nodes_map.values(): + if _is_conv_or_conv_transpose_node(n): + return len(n.args) > 2 and n.args[2] is not None + raise ValueError("Could not find conv node in matched conv + bn pattern") + + +def _no_conv_bias_filter( + match: "InternalMatch", + original_graph: Graph, + pattern_graph: Graph, +) -> bool: + """ + Match filter for the subgraph rewriter that returns True if the conv node in + the original graph does NOT have bias. + """ + return not _has_conv_bias_filter(match, original_graph, pattern_graph) + + +def _is_quantize(n: Node) -> bool: + return n.target in [ + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor, + torch.ops.quantized_decomposed.quantize_per_channel.default, + ] + + +def _is_dequantize(n: Node) -> bool: + return n.target in [ + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor, + torch.ops.quantized_decomposed.dequantize_per_channel.default, + ] + + +def _get_conv_bn_pattern_nodes(r: ReplacedPatterns) -> dict[str, tuple[Node, Node]]: + """ + Helper function to extract the nodes in the conv-bn fusion pattern after + subgraph rewriting, in the form of a map: + + {name: (original_node, replacement_node)} + + The following names must exist in the map: + + "conv", "conv_weight", "conv_input", "bn", "getitem" + + The following names may exist in the map: + + "conv_weight_q", "conv_weight_dq", "conv_bias", + "conv_bias_q", "conv_bias_dq" + """ + + def _get_nodes(nodes: list[Node]) -> tuple[Node, Node, Optional[Node]]: + """ + Return a 3-tuple of (conv_node, bn_node, getitem_node). + This asserts that the match contains exactly one of each node. + """ + conv_node, bn_node, getitem_node = None, None, None + for n in nodes: + if n.op != "call_function": + continue + if _is_conv_or_conv_transpose_node(n): + assert conv_node is None + conv_node = n + if _is_bn_node(n): + assert bn_node is None + bn_node = n + if n.target == operator.getitem: + assert getitem_node is None + getitem_node = n + assert conv_node is not None + assert bn_node is not None + return (conv_node, bn_node, getitem_node) + + def _get_q_dq_nodes(n: Node) -> tuple[Node, Node, Node]: + """ + Return a 3-tuple of (orig_node, q_node, dq_node). + """ + assert _is_dequantize(n) + q_node = n.args[0] + assert isinstance(q_node, Node) + assert _is_quantize(q_node) + orig_node = q_node.args[0] + assert isinstance(orig_node, Node) + return (orig_node, q_node, n) + + original_nodes = list(_filter_nodes_map(r.nodes_map).values()) + o_conv, o_bn, o_getitem = _get_nodes(original_nodes) + r_conv, r_bn, r_getitem = _get_nodes(r.replacements) + + # Create the mapping from original node to replacement node + assert o_getitem is None + assert r_getitem is None + mapping = { + "conv": (o_conv, r_conv), + "bn": (o_bn, r_bn), + } + + # Extract conv input and weight + # Note: here we extract the original nodes indirectly through the pattern nodes + # because the args of the original nodes are no longer available after replacement + (p_conv, _, _) = _get_nodes(list(r.nodes_map.keys())) + (p_conv_input, p_conv_weight, *_) = p_conv.args + (r_conv_input, r_conv_weight, *_) = r_conv.args + assert isinstance(p_conv_input, Node) + assert isinstance(p_conv_weight, Node) + assert isinstance(r_conv_input, Node) + assert isinstance(r_conv_weight, Node) + o_conv_input = r.nodes_map[p_conv_input] + o_conv_weight = r.nodes_map[p_conv_weight] + + # If conv weight is quantized, extract the q - dq nodes + if _is_dequantize(p_conv_weight): + p_conv_weight, p_conv_weight_q, p_conv_weight_dq = _get_q_dq_nodes( + p_conv_weight + ) + r_conv_weight, r_conv_weight_q, r_conv_weight_dq = _get_q_dq_nodes( + r_conv_weight + ) + o_conv_weight = r.nodes_map[p_conv_weight] + o_conv_weight_q = r.nodes_map[p_conv_weight_q] + o_conv_weight_dq = r.nodes_map[p_conv_weight_dq] + mapping["conv_weight_q"] = (o_conv_weight_q, r_conv_weight_q) + mapping["conv_weight_dq"] = (o_conv_weight_dq, r_conv_weight_dq) + mapping["conv_input"] = (o_conv_input, r_conv_input) + mapping["conv_weight"] = (o_conv_weight, r_conv_weight) + + # Extract conv bias + if len(p_conv.args) > 2 and len(r_conv.args) > 2: + p_conv_bias = p_conv.args[2] + r_conv_bias = r_conv.args[2] + assert isinstance(p_conv_bias, Node) + assert isinstance(r_conv_bias, Node) + o_conv_bias = r.nodes_map[p_conv_bias] + + # If conv bias is quantized, extract the q - dq nodes + if _is_dequantize(p_conv_bias): + p_conv_bias, p_conv_bias_q, p_conv_bias_dq = _get_q_dq_nodes(p_conv_bias) + r_conv_bias, r_conv_bias_q, r_conv_bias_dq = _get_q_dq_nodes(r_conv_bias) + o_conv_bias = r.nodes_map[p_conv_bias] + o_conv_bias_q = r.nodes_map[p_conv_bias_q] + o_conv_bias_dq = r.nodes_map[p_conv_bias_dq] + mapping["conv_bias_q"] = (o_conv_bias_q, r_conv_bias_q) + mapping["conv_bias_dq"] = (o_conv_bias_dq, r_conv_bias_dq) + mapping["conv_bias"] = (o_conv_bias, r_conv_bias) + return mapping + + +def _filter_nodes_map(nodes_map: dict[Node, Node]) -> dict[Node, Node]: + """ + Return a filtered `nodes_map` returned from the subgraph rewriter. + The filtered `nodes_map` will contain only nodes that are actually + matched in the pattern, excluding None or placeholder nodes. + """ + new_nodes_map: dict[Node, Node] = {} + for pattern_node, graph_node in nodes_map.items(): + # bias can be None + if graph_node is None: + continue + # skip pattern placeholder nodes + if pattern_node.op == "placeholder": + continue + new_nodes_map[pattern_node] = graph_node + return new_nodes_map + + +# TODO: this is error prone, use the replace_literals_with_placeholders hack instead +def _copy_over_literal_conv_args(original_node: Node, new_node: Node): + """ + Copy over literal args in conv, such as stride and padding, from the matched node + in the original graph to its replacement in the new graph. + + This is needed due to the following limitation in the subgraph rewriter when used + with dynamo export: literal (non-tensor) args are not supported in the match and + replacement patterns. This is because dynamo export automatically inlines these + literal args, making them dead placeholder nodes. In the future, we should check + if dynamo export can optionally disable this inlining, or if subgraph rewriter + can do the copying for us. See https://github.com/pytorch/pytorch/issues/100419. + + Note: Unlike other tensor args like conv weights and biases, literal args are + preserved in the original nodes after replacement, so we can access them here. + """ + assert _is_conv_or_conv_transpose_node(original_node) + assert _is_conv_or_conv_transpose_node(new_node) + # x, weight, bias, [stride, padding, dilation, transposed, output_padding, groups] + new_args = list(new_node.args) + if len(new_args) < 3: + # bias is optional, when it is not present, it means it is None + new_args.append(None) + new_node.args = tuple(new_args[:3]) + original_node.args[3:] + + +def _update_conv_input_qspec_map_after_replacement( + original_node: Node, replacement_node: Node +): + """ + Update the `input_qspec_map` in the annotation after subgraph rewriting. + + The original annotation referred to the nodes in the original graph, + so the keys in the `input_qspec_map` will need to be updated to reflect + the corresponding nodes in the replacement graph. + """ + assert _is_conv_or_conv_transpose_node(original_node) + assert _is_conv_or_conv_transpose_node(replacement_node) + if "quantization_annotation" not in original_node.meta: + return + original_input_qspec_map = original_node.meta[ + "quantization_annotation" + ].input_qspec_map + input_qspec_map = {} + # get the list of configs, it should be ordered as input, weight, bias + # note: this is really hacky, we need a better solution, hopefully + # in subgraph_rewriter, issue tracking the problem: https://github.com/pytorch/pytorch/issues/101820 + all_configs = list(original_input_qspec_map.items()) + # input activation + input_qspec_map[replacement_node.args[0]] = all_configs[0][1] + # weight + input_qspec_map[replacement_node.args[1]] = all_configs[1][1] + # bias + if len(replacement_node.args) > 2 and len(all_configs) > 2: + input_qspec_map[replacement_node.args[2]] = all_configs[2][1] + replacement_node.meta["quantization_annotation"].input_qspec_map = input_qspec_map + + +def _update_special_qspecs_after_replacement( + node: Node, + original_to_replacement_node: dict[Node, Node], +): + """ + Update the `SharedQuantizationSpec`s and `DerivedQuantizationSpec`s + used in `node`'s quantization annotation after subgraph rewriting. + + The original annotation referred to the nodes in the original graph, + so the nodes used in these special quantization specs will need to + be updated to the corresponding nodes in the replacement graph. + """ + + def _get_new_edge_or_node(edge_or_node: EdgeOrNode): + if isinstance(edge_or_node, Node): + _node = edge_or_node + return original_to_replacement_node.get(_node, _node) + elif ( + isinstance(edge_or_node, tuple) + and len(edge_or_node) == 2 + and all(isinstance(x, Node) for x in edge_or_node) + ): + src, dest = edge_or_node + return ( + original_to_replacement_node.get(src, src), + original_to_replacement_node.get(dest, dest), + ) + else: + raise ValueError("unexpected type for edge_or_node: ", type(edge_or_node)) + + def _get_new_qspec(qspec: QuantizationSpecBase): + if isinstance(qspec, SharedQuantizationSpec): + new_edge_or_node = _get_new_edge_or_node(qspec.edge_or_node) + return SharedQuantizationSpec(new_edge_or_node) + elif isinstance(qspec, DerivedQuantizationSpec): + new_derived_from = [_get_new_edge_or_node(x) for x in qspec.derived_from] + return dataclasses.replace(qspec, derived_from=new_derived_from) + else: + return qspec + + if "quantization_annotation" not in node.meta: + return + annotation = node.meta["quantization_annotation"] + for input_node, qspec in annotation.input_qspec_map.items(): + annotation.input_qspec_map[input_node] = _get_new_qspec(qspec) + annotation.output_qspec = _get_new_qspec(annotation.output_qspec) + + +def _fuse_conv_bn_qat(m: GraphModule) -> GraphModule: + # Example inputs for conv-bn1d patterns + _conv1d_bn_example_inputs = ( + torch.randn(1, 1, 3), # x + torch.randn(1, 1, 1), # conv_weight + torch.randn(1), # conv_bias + torch.randn(1), # bn_weight + torch.randn(1), # bn_bias + torch.randn(1), # bn_running_mean + torch.randn(1), # bn_running_var + ) + + # Example inputs for conv-bn2d patterns + _conv2d_bn_example_inputs = ( + torch.randn(1, 1, 3, 3), # x + torch.randn(1, 1, 1, 1), # conv_weight + torch.randn(1), # conv_bias + torch.randn(1), # bn_weight + torch.randn(1), # bn_bias + torch.randn(1), # bn_running_mean + torch.randn(1), # bn_running_var + ) + + has_bn = any(_is_bn_node(n) for n in m.graph.nodes) + if not has_bn: + return m + is_cuda_options = [True, False] if torch.cuda.is_available() else [False] + for is_cuda in is_cuda_options: + m = _fuse_conv_bn_qat_helper( + m, F.conv1d, _conv1d_bn_example_inputs, is_cuda=is_cuda + ) + m = _fuse_conv_bn_qat_helper( + m, F.conv2d, _conv2d_bn_example_inputs, is_cuda=is_cuda + ) + m = _fuse_conv_bn_qat_helper( + m, F.conv_transpose1d, _conv1d_bn_example_inputs, is_cuda=is_cuda + ) + m = _fuse_conv_bn_qat_helper( + m, F.conv_transpose2d, _conv2d_bn_example_inputs, is_cuda=is_cuda + ) + return m + + +def _fuse_conv_bn_qat_helper( + m: GraphModule, + conv_fn: Callable, + example_inputs: tuple[Any, ...], + is_cuda: bool, +) -> GraphModule: + """ + Given a graph of decomposed aten ops, replace the (conv + bn) pattern with + the fused QAT subgraph equivalent. The input graph should already be annotated. + The annotations in the original nodes will be preserved in the corresponding + nodes in the new subgraph. + + Note: This also handles the (conv + bn + relu) pattern. + """ + m.graph.eliminate_dead_code() + m.recompile() + + conv_bn_pattern = _get_conv_bn_pattern(conv_fn) + match_pattern = _get_aten_graph_module_for_pattern( + conv_bn_pattern, + example_inputs, + is_cuda, + ) + + # Step (1): Replace patterns with conv bias + # + # Here we do replacement separately for cases with and without conv bias, since + # the replacement patterns for these two cases are substantially different. + # TODO: use the public replace_pattern API once it also returns replacement nodes + + qat_conv_bn_pattern = _get_qat_conv_bn_pattern(conv_fn) + replacement_pattern_with_conv_bias = _get_aten_graph_module_for_pattern( + qat_conv_bn_pattern, + example_inputs, + is_cuda, + ) + replacements_with_conv_bias = replace_pattern_with_filters( + m, + match_pattern, + replacement_pattern_with_conv_bias, + match_filters=[_has_conv_bias_filter], + ignore_literals=True, + ) + m.recompile() + + # Step (2): Replace patterns without conv bias + + qat_conv_bn_pattern_no_conv_bias = _get_qat_conv_bn_pattern_no_conv_bias(conv_fn) + replacement_pattern_no_conv_bias = _get_aten_graph_module_for_pattern( + qat_conv_bn_pattern_no_conv_bias, + example_inputs, + is_cuda, + ) + replacements_no_conv_bias = replace_pattern_with_filters( + m, + match_pattern, + replacement_pattern_no_conv_bias, + match_filters=[_no_conv_bias_filter], + ignore_literals=True, + ) + m.recompile() + + # Step (3): Post processing + # + # Due to limited functionality in the subgraph rewriter, here we manually + # update the replacement graph as follows: + # + # (a) Copy over metadata from original subgraph. This ensures the stack traces + # and annotations are preserved in the new subgraph + # + # (b) Copy over literal args for conv from the original subgraph + # TODO: do this for literal args for batchnorm as well + # + # (c) Update all references of the old nodes in the original subgraph to refer + # to the corresponding nodes in the new subgraph in the annotations + # + # In the future, we should try to push as much of this functionality into the + # subgraph rewriter as possible, so we don't have to manually copy anything over. + # For more detail, see https://github.com/pytorch/pytorch/issues/100419. + + all_original_to_replacement_nodes = {} + for r in replacements_with_conv_bias + replacements_no_conv_bias: + replacement_dict = _get_conv_bn_pattern_nodes(r) + # The original conv node's "nn_module_stack" + conv_nn_module = replacement_dict["conv"][0].meta.get("nn_module_stack", None) + for k, node_tuple in replacement_dict.items(): + original_node, replacement_node = node_tuple + # Step (3a): Copy over metadata for all nodes in [conv - bn - getitem] + replacement_node.meta = original_node.meta + # If original_node is a get_attr node, it doesn't have nn_module_stack. + # In this case, we copy nn_module_stack from the original conv node. + if ( + k in ["conv_input", "conv_weight"] + and conv_nn_module + and "nn_module_stack" not in replacement_node.meta + ): + replacement_node.meta["nn_module_stack"] = copy.deepcopy(conv_nn_module) + if _is_conv_or_conv_transpose_node(original_node): + # Step (3b): Copy over conv literal args + _copy_over_literal_conv_args(original_node, replacement_node) + # Step (3c): Update old references in the conv node's input_qspec_map + _update_conv_input_qspec_map_after_replacement( + original_node, replacement_node + ) + all_original_to_replacement_nodes[original_node] = replacement_node + + # Step (3c): Update old references in the special qspecs for all nodes in the graph + for n in m.graph.nodes: + _update_special_qspecs_after_replacement(n, all_original_to_replacement_nodes) + + return m + + +def _duplicate_dequantize_node(m: GraphModule): + """ + Helper function to duplicate all dequantize nodes in the graph if the + node has more than one user. For example: + + Before: + quantize -> dequantize -> a + \\--> b + \\--> c + + After: + quantize -> dequantize_1 -> a + \\--> dequantize_2 -> b + \\--> dequantize_3 -> c + + This is useful for subgraph rewriting. E.g. if we wish to match the + pattern [dequantize - a] above, subgraph matching would fail because + the dequantize node has users outside the matched portion of the graph. + Instead, we match [dequantize_1 - a], which is safe. + """ + dq_op = torch.ops.quantized_decomposed.dequantize_per_tensor + for n in m.graph.nodes: + if n.op != "call_function" or n.target != dq_op or len(n.users) == 1: + continue + for user in list(n.users): + with m.graph.inserting_before(n): + new_node = m.graph.create_node("call_function", dq_op, n.args, n.kwargs) + user.replace_input_with(n, new_node) + m.graph.erase_node(n) + m.recompile() + + +def _remove_extra_dequantize(m: GraphModule): + """ + Removes duplicate dequant nodes in the graph, for an operator that has + multiple dequant nodes as a user, replace them with a single dequant node + that can be shared across all the uses. This should be seen as the "reverse" + of `_duplicate_dequantize_node`. + """ + dq_op = torch.ops.quantized_decomposed.dequantize_per_tensor + for n in m.graph.nodes: + dq_users = [ + user + for user in n.users + if user.op == "call_function" and user.target == dq_op + ] + if len(dq_users) > 1: + with m.graph.inserting_after(dq_users[0]): + new_node = m.graph.create_node( + "call_function", dq_op, dq_users[0].args, {} + ) + for dq_user in dq_users: + dq_user.replace_all_uses_with(new_node) + m.graph.erase_node(dq_user) + m.recompile() + + +def _copy_over_q_dq_args(original_node: Node, replacement_node: Node): + """ + Given a pair of quantize or dequantize nodes, copy over all literal args + from the original node to the replacement node. + """ + # For quantize_per_tensor, scale and zp are literals and need to be copied + # For quantize_per_channel, scale and zp are get_attr nodes and should be skipped + assert original_node.target == replacement_node.target + if original_node.target in ( + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + ): + # Args: input, [scale, zp, qmin, qmax, dtype] + start_copy_arg_index = 1 + elif original_node.target in ( + torch.ops.quantized_decomposed.quantize_per_channel.default, + torch.ops.quantized_decomposed.dequantize_per_channel.default, + ): + # Args: input, scale, zp, [axis, qmin, qmax, dtype] + start_copy_arg_index = 3 + else: + raise ValueError( + f"Expected quantize/dequantize nodes, got '{original_node.target}'" + ) + replacement_node.args = ( + replacement_node.args[:start_copy_arg_index] + + original_node.args[start_copy_arg_index:] + ) + + +def _fold_conv_bn_qat(m: GraphModule) -> GraphModule: + # Example inputs for quantized and folded conv-bn1d patterns used in convert + _quantized_conv1d_bn_example_inputs = ( + torch.randn(1, 1, 3), # x + torch.randn(1, 1, 1), # conv_weight + torch.randn(1), # bn_weight + torch.randn(1), # bn_bias + torch.randn(1), # bn_running_mean + torch.randn(1), # bn_running_var + ) + + # Example inputs for quantized and folded conv-bn2d patterns used in convert + _quantized_conv2d_bn_example_inputs = ( + torch.randn(1, 1, 3, 3), # x + torch.randn(1, 1, 1, 1), # conv_weight + torch.randn(1), # bn_weight + torch.randn(1), # bn_bias + torch.randn(1), # bn_running_mean + torch.randn(1), # bn_running_var + ) + + has_bn = any(_is_bn_node(n) for n in m.graph.nodes) + if not has_bn: + return m + is_cuda_options = [True, False] if torch.cuda.is_available() else [False] + for is_cuda in is_cuda_options: + m = _fold_conv_bn_qat_helper( + m, F.conv1d, _quantized_conv1d_bn_example_inputs, is_cuda=is_cuda + ) + m = _fold_conv_bn_qat_helper( + m, F.conv2d, _quantized_conv2d_bn_example_inputs, is_cuda=is_cuda + ) + m = _fold_conv_bn_qat_helper( + m, F.conv_transpose1d, _quantized_conv1d_bn_example_inputs, is_cuda=is_cuda + ) + m = _fold_conv_bn_qat_helper( + m, F.conv_transpose2d, _quantized_conv2d_bn_example_inputs, is_cuda=is_cuda + ) + + # remove in place add from batchnorm tracking traning stats + for node in m.graph.nodes: + if ( + node.target == torch.ops.aten.add_.Tensor + and node.args[0].op == "get_attr" + and node.args[1] == 1 + and torch.nn.modules.batchnorm.BatchNorm2d + in [val[1] for val in node.meta["source_fn_stack"]] + ): + m.graph.erase_node(node) + + m.graph.eliminate_dead_code() + m.recompile() + + return m + + +def _fold_conv_bn_qat_helper( + m: GraphModule, + conv_fn: Callable, + example_inputs: tuple[Any, ...], + is_cuda: bool, +) -> GraphModule: + """ + Replace the quantized (conv + bn) pattern with conv with bn weights folded into the weights of conv. + """ + + m.graph.eliminate_dead_code() + m.recompile() + _duplicate_dequantize_node(m) + + # Step (1): Replace QAT pattern with simple [conv - bn] pattern + replacements = [] + replacement_options = itertools.product( + [True, False], # is_per_channel + [True, False], # has_bias + [True, False], # bias_is_quantized + [True, False], # bn_is_training + ) + for ( + is_per_channel, + has_bias, + bias_is_quantized, + bn_is_training, + ) in replacement_options: + # For the cases without bias, `bias_is_quantized` is irrelevant, so here we arbitrarily + # filter out one of the values for this flag to avoid having duplicate patterns + if not has_bias and bias_is_quantized: + continue + kwargs = _get_quantized_conv_bn_example_inputs_kwargs( + is_per_channel, has_bias, bias_is_quantized, is_cuda + ) + match_pattern = _get_quantized_qat_conv_bn_pattern( + is_per_channel, has_bias, bias_is_quantized, conv_fn, bn_is_training + ) + match_pattern = _get_aten_graph_module_for_pattern( + match_pattern, + example_inputs, + is_cuda, + **kwargs, + ) + replacement_pattern = _get_folded_quantized_qat_conv_bn_pattern( + is_per_channel, has_bias, bias_is_quantized, conv_fn, bn_is_training + ) + replacement_pattern = _get_aten_graph_module_for_pattern( + replacement_pattern, + example_inputs, + is_cuda, + **kwargs, + ) + replacements.extend( + replace_pattern_with_filters( + m, + match_pattern, + replacement_pattern, + ignore_literals=True, + ) + ) + m.recompile() + _remove_extra_dequantize(m) + + for r in replacements: + node_map = _get_conv_bn_pattern_nodes(r) + + # Step (2): Copy over metadata from original subgraph + for original_node, replacement_node in node_map.values(): + replacement_node.meta = original_node.meta + + # Step (3): Copy over args for weight (and optionally bias) q - dq nodes + _copy_over_q_dq_args(*node_map["conv_weight_q"]) + _copy_over_q_dq_args(*node_map["conv_weight_dq"]) + if "conv_bias_q" in node_map: + assert "conv_bias_dq" in node_map + _copy_over_q_dq_args(*node_map["conv_bias_q"]) + _copy_over_q_dq_args(*node_map["conv_bias_dq"]) + + # Step (4): Fold BN weights into conv + conv_bias = None + (_, conv_node) = node_map["conv"] + (_, bn_node) = node_map["bn"] + (_, conv_weight) = node_map["conv_weight"] + if "conv_bias" in node_map: + (_, conv_bias) = node_map["conv_bias"] + fold_bn_weights_into_conv_node(conv_node, conv_weight, conv_bias, bn_node, m) + + # Copy over literal args for conv + for original_node in _filter_nodes_map(r.nodes_map).values(): + if _is_conv_or_conv_transpose_node(original_node): + _copy_over_literal_conv_args(original_node, conv_node) + + m.graph.eliminate_dead_code() + m.recompile() + return m diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantize_pt2e.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantize_pt2e.py new file mode 100644 index 0000000000000000000000000000000000000000..d352d91f9bce563a3967a605902bb829f41f3986 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantize_pt2e.py @@ -0,0 +1,270 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, TORCH_VERSION_AT_LEAST_2_7 + +if TORCH_VERSION_AT_LEAST_2_7: + from .constant_fold import constant_fold + +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_manager import PassManager + +from torchao.quantization.pt2e.qat_utils import _fold_conv_bn_qat, _fuse_conv_bn_qat +from torchao.quantization.pt2e.quantizer import ( # noqa: F401 + DuplicateDQPass, + PortNodeMetaForQDQ, + Quantizer, +) +from torchao.quantization.pt2e.utils import ( + _disallow_eval_train, + _fuse_conv_bn_, + _get_node_name_to_scope, +) + +from .convert import _convert_to_reference_decomposed_fx +from .prepare import prepare +from .reference_representation_rewrite import reference_representation_rewrite + +__all__ = [ + "prepare_pt2e", + "prepare_qat_pt2e", + "convert_pt2e", +] + + +def prepare_pt2e( + model: GraphModule, + quantizer: Quantizer, +) -> GraphModule: + """Prepare a model for post training quantization + + Args: + * `model` (torch.fx.GraphModule): a model captured by `torch.export.export_for_training` API. + * `quantizer`: A backend specific quantizer that conveys how user want the + model to be quantized. Tutorial for how to write a quantizer can be found here: + https://pytorch.org/tutorials/prototype/pt2e_quantizer.html + + Return: + A GraphModule with observer (based on quantizer annotation), ready for calibration + + Example:: + + import torch + from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e + from torchao.quantization.pt2e.quantizer import ( + XNNPACKQuantizer, + get_symmetric_quantization_config, + ) + + class M(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(5, 10) + + def forward(self, x): + return self.linear(x) + + # initialize a floating point model + float_model = M().eval() + + # define calibration function + def calibrate(model, data_loader): + model.eval() + with torch.no_grad(): + for image, target in data_loader: + model(image) + + # Step 1. program capture + # NOTE: this API will be updated to torch.export API in the future, but the captured + # result shoud mostly stay the same + m = torch.export.export_for_training(m, *example_inputs).module() + # we get a model with aten ops + + # Step 2. quantization + # backend developer will write their own Quantizer and expose methods to allow + # users to express how they + # want the model to be quantized + quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config()) + m = prepare_pt2e(m, quantizer) + + # run calibration + # calibrate(m, sample_inference_data) + """ + torch._C._log_api_usage_once("quantization_api.quantize_pt2e.prepare_pt2e") + original_graph_meta = model.meta + node_name_to_scope = _get_node_name_to_scope(model) + # TODO: check qconfig_mapping to make sure conv and bn are both configured + # to be quantized before fusion + # TODO: (maybe) rewrite this with subgraph_rewriter + _fuse_conv_bn_(model) + model = quantizer.transform_for_annotation(model) + quantizer.annotate(model) + quantizer.validate(model) + model = prepare( + model, + node_name_to_scope, + is_qat=False, + obs_or_fq_callback=quantizer.prepare_obs_or_fq_callback, + ) + model.meta.update(original_graph_meta) + model = _disallow_eval_train(model) + return model + + +def prepare_qat_pt2e( + model: GraphModule, + quantizer: Quantizer, +) -> GraphModule: + """Prepare a model for quantization aware training + + Args: + * `model` (torch.fx.GraphModule): see :func:`~torchao.quantization.pt2e.quantize_pt2e.prepare_pt2e` + * `quantizer`: see :func:`~torchao.quantization.pt2e.quantize_pt2e.prepare_pt2e` + + Return: + A GraphModule with fake quant modules (based on quantizer annotation), ready for + quantization aware training + + Example:: + import torch + from torchao.quantization.pt2e.quantize_pt2e import prepare_qat_pt2e + from torchao.quantization.pt2e.quantizer import ( + XNNPACKQuantizer, + get_symmetric_quantization_config, + ) + + class M(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(5, 10) + + def forward(self, x): + return self.linear(x) + + # initialize a floating point model + float_model = M().eval() + + # define the training loop for quantization aware training + def train_loop(model, train_data): + model.train() + for image, target in data_loader: + ... + + # Step 1. program capture + # NOTE: this API will be updated to torch.export API in the future, but the captured + # result shoud mostly stay the same + m = torch.export.export_for_training(m, *example_inputs).module() + # we get a model with aten ops + + # Step 2. quantization + # backend developer will write their own Quantizer and expose methods to allow + # users to express how they + # want the model to be quantized + quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config()) + m = prepare_qat_pt2e(m, quantizer) + + # run quantization aware training + train_loop(prepared_model, train_loop) + + """ + torch._C._log_api_usage_once("quantization_api.quantize_pt2e.prepare_qat_pt2e") + original_graph_meta = model.meta + node_name_to_scope = _get_node_name_to_scope(model) + model = quantizer.transform_for_annotation(model) + quantizer.annotate(model) + quantizer.validate(model) + # Perform fusion after annotate to avoid quantizing ops in the new + # subgraph that don't need to be quantized + # TODO: only fuse if conv and bn are both configured to be quantized + _fuse_conv_bn_qat(model) + model = prepare( + model, + node_name_to_scope, + is_qat=True, + obs_or_fq_callback=quantizer.prepare_obs_or_fq_callback, + ) + model.meta.update(original_graph_meta) + model = _disallow_eval_train(model) + return model + + +_QUANT_OPS = [ + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor, + torch.ops.quantized_decomposed.quantize_per_channel.default, +] + +# ops are only registered after 2.5 +if TORCH_VERSION_AT_LEAST_2_5: + _QUANT_OPS += [ + torch.ops.torchao.quantize_affine, + ] + + +def _quant_node_constraint(n: Node) -> bool: + """If there is any pure ops between get_attr and quantize op they will be const propagated + e.g. get_attr(weight) -> transpose -> quantize -> dequantize* + (Note: dequantize op is not going to be constant propagated) + + This filter is added because we don't want to constant fold the things that are not + related to quantization + """ + return n.op == "call_function" and n.target in _QUANT_OPS + + +def convert_pt2e( + model: GraphModule, + use_reference_representation: bool = False, + fold_quantize: bool = True, +) -> GraphModule: + """Convert a calibrated/trained model to a quantized model + + Args: + * `model` (torch.fx.GraphModule): calibrated/trained model + * `use_reference_representation` (bool): boolean flag to indicate whether to produce referece representation or not + * `fold_quantize` (bool): boolean flag for whether fold the quantize op or not + + Returns: + quantized model, either in q/dq representation or reference representation + + Example:: + + # prepared_model: the model produced by `prepare_pt2e`/`prepare_qat_pt2e` and calibration/training + # `convert_pt2e` produces a quantized model that represents quantized computation with + # quantize dequantize ops and fp32 ops by default. + # Please refer to + # https://pytorch.org/tutorials/prototype/pt2e_quant_ptq_static.html#convert-the-calibrated-model-to-a-quantized-model + # for detailed explanation of output quantized model + quantized_model = convert_pt2e(prepared_model) + + """ + torch._C._log_api_usage_once("quantization_api.quantize_pt2e.convert_pt2e") + if not isinstance(use_reference_representation, bool): + raise ValueError( + "Unexpected argument type for `use_reference_representation`, " + f"please make sure you intend to pass argument {use_reference_representation} to convert_pt2e" + ) + original_graph_meta = model.meta + model = _convert_to_reference_decomposed_fx(model) + model = _fold_conv_bn_qat(model) + + pm = PassManager([DuplicateDQPass()]) + model = pm(model).graph_module + + pm = PassManager([PortNodeMetaForQDQ()]) + model = pm(model).graph_module + + if fold_quantize and TORCH_VERSION_AT_LEAST_2_7: + constant_fold(model, _quant_node_constraint) + + if use_reference_representation: + model = reference_representation_rewrite(model) + + model.meta.update(original_graph_meta) + model = _disallow_eval_train(model) + return model diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__init__.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c5d134b71d1f4ad9513567ad25208923b5e6255 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__init__.py @@ -0,0 +1,53 @@ +from .composable_quantizer import ComposableQuantizer +from .duplicate_dq_pass import DuplicateDQPass +from .port_metadata_pass import PortNodeMetaForQDQ +from .quantizer import ( + DerivedQuantizationSpec, + EdgeOrNode, + FixedQParamsQuantizationSpec, + QuantizationAnnotation, + QuantizationSpec, + QuantizationSpecBase, + Quantizer, + SharedQuantizationSpec, +) +from .utils import ( + OperatorConfig, + OperatorPatternType, + QuantizationConfig, + annotate_input_qspec_map, + annotate_output_qspec, + get_bias_qspec, + get_input_act_qspec, + get_module_name_filter, + get_output_act_qspec, + get_weight_qspec, + is_valid_annotation, +) + +__all__ = [ + # basic classes for quantizer and annotations + "Quantizer", + "ComposableQuantizer", + "EdgeOrNode", + "QuantizationSpec", + "QuantizationSpecBase", + "DerivedQuantizationSpec", + "FixedQParamsQuantizationSpec", + "SharedQuantizationSpec", + "QuantizationAnnotation", + # utils + "annotate_input_qspec_map", + "annotate_output_qspec", + "get_module_name_filter", + "is_valid_annotation", + "QuantizationConfig", + "OperatorPatternType", + "OperatorConfig", + "get_input_act_qspec", + "get_output_act_qspec", + "get_weight_qspec", + "get_bias_qspec", + "DuplicateDQPass", + "PortNodeMetaForQDQ", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6713a617a6ad8339fe75ac9a329d5bb57449f1ae Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/composable_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/composable_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47333b4c1aa973e1bd3cd121bd88ef66f06f8b04 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/composable_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/duplicate_dq_pass.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/duplicate_dq_pass.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5087465e28c5892d778712f13e57a68f2315b64 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/duplicate_dq_pass.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/embedding_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/embedding_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da263d65baea14f57c65d5f648f56fb6eef1c59e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/embedding_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/port_metadata_pass.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/port_metadata_pass.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e1b2e429631ecda78f935aa44a4dad637ab225e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/port_metadata_pass.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb53c21589f859f5551537517bf309264c987eee Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b7995f1d24309a74396efcd07b2e2af1ac786d21 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/x86_inductor_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/x86_inductor_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f24ac12d9e6d04f643a8cbca126f6d15969c9082 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/x86_inductor_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/xpu_inductor_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/xpu_inductor_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4712d5af00d8238d92a4bf306416902cc8fb24ed Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/__pycache__/xpu_inductor_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/composable_quantizer.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/composable_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..6602151e3f507ea2115e1f2302bafca01cf825af --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/composable_quantizer.py @@ -0,0 +1,84 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .quantizer import QuantizationAnnotation, Quantizer + +if TYPE_CHECKING: + import torch + from torch.fx import Node + +__all__ = [ + "ComposableQuantizer", +] + + +class ComposableQuantizer(Quantizer): + """ + ComposableQuantizer allows users to combine more than one quantizer into a single quantizer. + This allows users to quantize a model with multiple quantizers. E.g., embedding quantization + maybe supported by one quantizer while linear layers and other ops might be supported by another + quantizer. + + ComposableQuantizer is initialized with a list of `Quantizer` instances. + The order of the composition matters since that is the order in which the quantizers will be + applies. + Example: + ``` + embedding_quantizer = EmbeddingQuantizer() + linear_quantizer = MyLinearQuantizer() + xnnpack_quantizer = XNNPackQuantizer() # to handle ops not quantized by previous two quantizers + composed_quantizer = ComposableQuantizer([embedding_quantizer, linear_quantizer, xnnpack_quantizer]) + prepared_m = prepare_pt2e(model, composed_quantizer) + ``` + """ + + def __init__(self, quantizers: list[Quantizer]): + super().__init__() + self.quantizers = quantizers + self._graph_annotations: dict[Node, QuantizationAnnotation] = {} + + def _record_and_validate_annotations( + self, gm: torch.fx.GraphModule, quantizer: Quantizer + ) -> None: + for n in gm.graph.nodes: + if "quantization_annotation" in n.meta: + # check if the annotation has been changed by + # comparing QuantizationAnnotation object id + if n in self._graph_annotations and ( + id(self._graph_annotations[n]) + != id(n.meta["quantization_annotation"]) + ): + raise RuntimeError( + f"Quantizer {quantizer.__class__.__name__} has changed annotations on node {n}" + ) + else: + self._graph_annotations[n] = n.meta["quantization_annotation"] + else: + if n in self._graph_annotations: + raise RuntimeError( + f"Quantizer {quantizer.__class__.__name__} has removed annotations on node {n}" + ) + + def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule: + """just handling global spec for now""" + for quantizer in self.quantizers: + quantizer.annotate(model) + self._record_and_validate_annotations(model, quantizer) + return model + + def transform_for_annotation( + self, model: torch.fx.GraphModule + ) -> torch.fx.GraphModule: + for quantizer in self.quantizers: + model = quantizer.transform_for_annotation(model) + return model + + def validate(self, model: torch.fx.GraphModule) -> None: + pass diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/duplicate_dq_pass.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/duplicate_dq_pass.py new file mode 100644 index 0000000000000000000000000000000000000000..2bf4e732c1c53d2bb5614bc37687849505d52f5e --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/duplicate_dq_pass.py @@ -0,0 +1,91 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import logging +import operator + +import torch +from torch.fx.node import map_arg +from torch.fx.passes.infra.pass_base import PassBase, PassResult + +from torchao.quantization.pt2e.utils import ( + _filter_sym_size_users, +) + +from .utils import ( + is_valid_annotation, +) + +logger = logging.getLogger(__name__) +logger.setLevel(logging.WARNING) + +__all__ = ["DuplicateDQPass"] + +_QUANTIZE_OPS = [ + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor, + torch.ops.quantized_decomposed.quantize_per_channel.default, +] + +_DEQUANTIZE_OPS = [ + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor, + torch.ops.quantized_decomposed.dequantize_per_channel.default, +] + + +def _maybe_duplicate_dq( + gm: torch.fx.GraphModule, dq_node: torch.fx.Node, user: torch.fx.Node +): + annotation = user.meta.get("quantization_annotation", None) + if not is_valid_annotation(annotation): + return + with gm.graph.inserting_after(dq_node): + new_node = gm.graph.node_copy(dq_node) + + def maybe_replace_node(n: torch.fx.Node) -> torch.fx.Node: + if n == dq_node: + return new_node + else: + return n + + new_args = map_arg(user.args, maybe_replace_node) + new_kwargs = map_arg(user.kwargs, maybe_replace_node) + user.args = new_args # type: ignore[assignment] + user.kwargs = new_kwargs # type: ignore[assignment] + + +class DuplicateDQPass(PassBase): + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + for node in graph_module.graph.nodes: + if node.op == "call_function" and node.target in _DEQUANTIZE_OPS: + dq_users = _filter_sym_size_users(node) + if len(dq_users) <= 1: + continue + # Do not duplicate dq for dynamic quantization + # Pattern: choose_qparam - getitem - q - dq + q_node = node.args[0] + if q_node.op == "call_function" and q_node.target in _QUANTIZE_OPS: + getitem_node = q_node.args[1] + if ( + isinstance(getitem_node, torch.fx.node.Node) + and getitem_node.op == "call_function" + and getitem_node.target == operator.getitem + ): + choose_qparam_node = getitem_node.args[0] + if ( + isinstance(choose_qparam_node, torch.fx.node.Node) + and choose_qparam_node.op == "call_function" + and choose_qparam_node.target + == torch.ops.quantized_decomposed.choose_qparams.tensor + ): + continue + for user in dq_users: + _maybe_duplicate_dq(graph_module, node, user) + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/embedding_quantizer.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/embedding_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..40979f6fe820fad27a02a26c24b7a93d5a5c1d78 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/embedding_quantizer.py @@ -0,0 +1,101 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +from __future__ import annotations + +import copy + +import torch +import torch.nn.functional as F + +from torchao.quantization.pt2e.observer import PerChannelMinMaxObserver +from torchao.quantization.pt2e.quantizer import ( + OperatorConfig, + OperatorPatternType, + QuantizationAnnotation, + QuantizationConfig, + QuantizationSpec, + Quantizer, +) + +__all__ = [ + "get_embedding_operators_config", + "EmbeddingQuantizer", +] + + +def get_embedding_operators_config() -> OperatorConfig: + weight_quantization_spec = QuantizationSpec( + dtype=torch.uint8, + qscheme=torch.per_channel_affine_float_qparams, + ch_axis=0, + observer_or_fake_quant_ctr=PerChannelMinMaxObserver.with_args(eps=2**-12), + ) + quantization_config = QuantizationConfig(None, None, weight_quantization_spec, None) + ops: list[OperatorPatternType] = [[torch.nn.Embedding]] + ops.append([F.embedding]) + supported_config_and_operators = OperatorConfig( + config=quantization_config, operators=ops + ) + return copy.deepcopy(supported_config_and_operators) + + +class EmbeddingQuantizer(Quantizer): + def __init__(self) -> None: + super().__init__() + + @classmethod + def get_supported_quantization_configs(cls) -> list[QuantizationConfig]: + op_configs: set[QuantizationConfig] = { + spec for spec, _ in cls.get_supported_operators() + } + return list(op_configs) + + @classmethod + def get_supported_operator_for_quantization_config( + cls, quantization_config: QuantizationConfig + ) -> list[OperatorPatternType]: + for config, ops in cls.get_supported_operators(): + # note: this assumes each entry in cls.supported_spec_and_operators + # corresponds to one spec, e.g. we don't have + # [(spec1, op_list1), (spec1, op_list2), (spec2, op_list3)] + # where the first and second entry have the same spec but did not + # merge the op list + if config == quantization_config: + return ops + return [] + + def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule: + """just handling global spec for now""" + self._annotate_embedding_ops(model.graph) + return model + + def _annotate_embedding_ops(self, graph: torch.fx.Graph) -> None: + embedding_config: OperatorConfig = get_embedding_operators_config() + for node in graph.nodes: + # Keep node parsing based annotations instead of module partitioners + # just as an example of alternate ways of annotating + if ( + node.op == "call_function" + and node.target == torch.ops.aten.embedding.default + ): + if embedding_config.config.weight is None: + raise ValueError( + "Embedding config must have a valid weight quantization spec." + ) + node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map={ + node.args[0]: embedding_config.config.weight, + } + ) + + def validate(self, model: torch.fx.GraphModule) -> None: + pass + + @classmethod + def get_supported_operators(cls) -> list[OperatorConfig]: + return [get_embedding_operators_config()] diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/port_metadata_pass.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/port_metadata_pass.py new file mode 100644 index 0000000000000000000000000000000000000000..b0d910e603e2a8aff06d86fb5bec5d123a6ac7ec --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/port_metadata_pass.py @@ -0,0 +1,291 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import logging +from typing import Any, Optional + +import torch +from torch._export.error import InternalError +from torch.fx.passes.infra.pass_base import PassBase, PassResult + +from torchao.quantization.pt2e.utils import ( + _filter_sym_size_users, +) +from torchao.quantization.quant_primitives import quant_lib # noqa: F401 +from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 + +from .quantizer import ( + QuantizationSpecBase, +) +from .utils import ( + is_valid_annotation, +) + +logger = logging.getLogger(__name__) +logger.setLevel(logging.ERROR) + +__all__ = ["PortNodeMetaForQDQ"] + +_METADATA_TO_PORT = [ + "stack_trace", + "quantization_tag", +] + +_QUANTIZE_OPS = [ + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor, + torch.ops.quantized_decomposed.quantize_per_channel.default, +] + +_DEQUANTIZE_OPS = [ + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor, + torch.ops.quantized_decomposed.dequantize_per_channel.default, +] + +_CHOOSE_QPARAMS_OPS = [ + torch.ops.quantized_decomposed.choose_qparams.tensor, + torch.ops.quantized_decomposed.choose_qparams_symmetric.tensor, +] + + +# ops are only registered after 2.5 +if TORCH_VERSION_AT_LEAST_2_5: + _QUANTIZE_OPS += [torch.ops.torchao.quantize_affine] + _DEQUANTIZE_OPS += [torch.ops.torchao.dequantize_affine] + _CHOOSE_QPARAMS_OPS += [torch.ops.torchao.choose_qparams_affine] + + +def _add_metadata(to_node: torch.fx.Node, from_node: torch.fx.Node) -> None: + from_meta = from_node.meta + for meta_name in _METADATA_TO_PORT: + if meta_name in from_meta: + to_node.meta[meta_name] = from_meta[meta_name] + + +def _has_quant_annotation(node: torch.fx.Node) -> bool: + return "quantization_annotation" in node.meta + + +def _find_choose_qparams_node(node: torch.fx.Node) -> Optional[torch.fx.Node]: + # BFS to look for choose qparams + from collections import deque + + queue = deque(list(node.users.keys())) + while len(queue): + n = queue.popleft() + if n.op == "output": + continue + if n.op == "call_function" and n.target in _CHOOSE_QPARAMS_OPS: + return n + for k in n.users.keys(): + queue.append(k) + return None + + +def _is_connected(source: torch.fx.Node, dest: torch.fx.Node) -> bool: + """ + Assuming dest is one of the ops inserted by quant workflow, this function + finds if source and dest are connected. Assumption is that only quant workflow + inserted ops exist between source and dest + """ + quant_workflow_ops = _QUANTIZE_OPS + _DEQUANTIZE_OPS + quant_workflow_ops.append(torch.ops.quantized_decomposed.choose_qparams.tensor) + while dest.target in quant_workflow_ops: + if not isinstance(dest.args[0], torch.fx.Node): + raise ValueError( + f"expected arg[0] of quant workflow ops to be a node but found {dest.args[0]}" + ) + dest = dest.args[0] + return dest == source + + +def _find_q_dq_node_for_user( + produer: torch.fx.Node, user: torch.fx.Node +) -> tuple[Any, Any]: + """ + Find q, dq pair corresponding to [producer -> q -> dq -> user] + Utils works by finding dq arg of user and ensuring it is connected to + producer + """ + dq_node = None + for n in user.args: + if ( + isinstance(n, torch.fx.Node) + and n.op == "call_function" + and n.target in _DEQUANTIZE_OPS + ): + if _is_connected(produer, n): + dq_node = n + break + if dq_node is None: + for n in user.kwargs: + if ( + isinstance(n, torch.fx.Node) + and n.op == "call_function" + and n.target in _DEQUANTIZE_OPS + ): + if _is_connected(produer, n): + dq_node = n + break + if dq_node is None: + return (None, None) + + q_node = None + if ( + dq_node.args[0].op == "call_function" # type: ignore[union-attr] + and dq_node.args[0].target in _QUANTIZE_OPS # type: ignore[union-attr] + ): + q_node = dq_node.args[0] + return (q_node, dq_node) + + +def _port_metadata_for_input_quant_nodes( + input_node: torch.fx.Node, + node: torch.fx.Node, + qspec: Optional[QuantizationSpecBase], +): + if qspec is None: + return + + is_dynamic_quant = getattr(qspec, "is_dynamic", None) + if is_dynamic_quant is not None and is_dynamic_quant is True: + choose_qparams_node = _find_choose_qparams_node(input_node) + if choose_qparams_node is None: + raise ValueError(f"No chose qparams node found for {node}") + choose_qparam_users = _filter_sym_size_users(choose_qparams_node) + if len(choose_qparam_users) != 2: + raise InternalError(f"Expecting exactly two user for {choose_qparams_node}") + scale_node = choose_qparam_users.pop() + dynamic_q_node = next(iter(scale_node.users.keys())) + dynamic_q_node_users = _filter_sym_size_users(dynamic_q_node) + if len(dynamic_q_node_users) > 1: + raise InternalError(f"Expecting single user for {dynamic_q_node}") + dynamic_dq_node = dynamic_q_node_users.pop() + _add_metadata(choose_qparams_node, node) + _add_metadata(dynamic_q_node, node) + _add_metadata(dynamic_dq_node, node) + else: + q_node, dq_node = _find_q_dq_node_for_user(input_node, node) + if q_node is None or dq_node is None: + return + # add metadata for all the node between q_node and get_attr node + # if the q_node can be traced back to get_attr node + q_to_get_attr_nodes = [q_node] + q_node_input = q_node.args[0] + while ( + isinstance(q_node_input, torch.fx.Node) + and q_node_input.op == "call_function" + and q_node_input.target + in [ + torch.ops.aten.flatten.using_ints, + torch.ops.aten.permute.default, + torch.ops.aten.permute_copy.default, + torch.ops.aten.slice_copy.Tensor, + torch.ops.aten.squeeze.dim, + torch.ops.aten.squeeze_copy.dim, + torch.ops.aten.transpose.Dimname, + torch.ops.aten.transpose.int, + torch.ops.aten.transpose_, + torch.ops.aten.view_copy.default, + torch.ops.aten.view.default, + torch.ops.aten._mkldnn_transpose, + ] + ): + q_to_get_attr_nodes.append(q_node_input) + q_node_input = q_node_input.args[0] + if isinstance(q_node_input, torch.fx.Node) and q_node_input.op == "get_attr": + for n in q_to_get_attr_nodes: + _add_metadata(n, q_node_input) + _add_metadata(dq_node, node) + + +def _port_metadata_for_output_quant_nodes( + node: torch.fx.Node, qspec: Optional[QuantizationSpecBase] +): + if qspec is None: + return + + node_users = _filter_sym_size_users(node) + if len(node.users) == 0: + return + if len(node_users) != 1: + logger.warning(f"Expecting {node} to have single user") # noqa: G004 + q_node = node_users.pop() + if q_node.op != "call_function" or q_node.target not in _QUANTIZE_OPS: + logger.warning( + f"Expecting {node} user to be a quantized op but got {q_node}" # noqa: G004 + ) # noqa: G004 + return + + _add_metadata(q_node, node) + + +class PortNodeMetaForQDQ(PassBase): + """ + Port metadata for nodes added by quantization flow. + For static quant these are: + - quantizer_per_tensor.default, dequantize_per_tensor.default + - quantizer_per_channel.default, dequantize_per_channel.default + For dynamic quant these are: + - choose_qparams.tensor + - quantizer_per_tensor.tensor, dequantize_per_tensor.tensor + - quantizer_per_channel.default, dequantize_per_channel.default + + Rules of porting metadata: + - Metadata to be ported: + - nn_module_stack + - stack_trace + - quantization_tag + - Metadata to NOT be ported: + - Everything else + - Rules: + - Statically quantized patterns: + - Dequantize nodes on the inputs to be quantized inherit metadata of the consumer node. + - Quantize nodes on the outputs inherit metadata of the producer node. + - Example 1: + - Original: [Conv -> AvgPool -> Linear] + - Quantized [Q-> DQ -> Conv -> Q -> DQ -> AvgPool -> Q -> DQ -> Linear -> Q -> DQ] + - Inner brackets specify which nodes Q/DQ inherit metdata from + - [Q-> [DQ -> Conv -> Q] -> [DQ -> AvgPool -> Q] -> [DQ -> Linear -> Q] -> DQ] + - Note first Q and last DQ do not inherit metadata from any nodes + - Example 2: + - Original: [Conv -> AvgPool -> Linear] + - AvgPool is not quantized + - Quantized [Q-> DQ -> Conv -> Q -> DQ -> AvgPool -> Q -> DQ -> Linear -> Q -> DQ] + - Inner brackets specify which nodes Q/DQ inherit metdata from + - [Q-> [DQ -> Conv -> Q] -> DQ -> [AvgPool] -> Q -> [DQ -> Linear -> Q] -> DQ] + - Note DQ and Q nodes around AvgPool do not inherit metadata from AvgPool because + AvgPool was not supposed to be quantized. Metadata porting relies on quantization_annotation + on the nodes (in this case AvgPool node) to conclude if the node or patter was + supposed to be quantized. And subsequntly decide if the preceding Q, if any, should + inherit metadata from AvgPool. + - Dynamically quantized patterns: + - Input that are dynamically quantized have choose_qparams, quantize and dequantize nodes + - For example, below linear is dynamically quantized while rest statically: + - Original: [Conv -> AvgPool -> Linear] + - Quantized [Q-> DQ -> Conv -> Q -> DQ -> AvgPool -> Q -> DQ -> choose_params -> Q -> DQ -> Linear] + - Quantized [Q-> [DQ -> Conv -> Q] -> [DQ -> AvgPool -> Q] -> DQ -> [choose_params -> Q -> DQ -> Linear]] + - Note first Q does not inherit metadata from any nodes + NB: + - The best place for porting metadata is during observer conversion to q/dq. This is because it precisely + knows which quantization spec is converted to q/dq and thus from where the metadata should be ported. + However, since FX and PT2E quant workflow are on a common code-base, this hurts readability quite a bit. + Doing it via a separate pass, helps readability of the code. Once we are able to refactor PT2E quant + code, this pass should like to be integrated in the refactored variant of "convert" step. + """ + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + for node in graph_module.graph.nodes: + annotation = node.meta.get("quantization_annotation", None) + if is_valid_annotation(annotation): + input_qspec_map = node.meta["quantization_annotation"].input_qspec_map + output_qspec = node.meta["quantization_annotation"].output_qspec + for input_node, qspec in input_qspec_map.items(): + _port_metadata_for_input_quant_nodes(input_node, node, qspec) + _port_metadata_for_output_quant_nodes(node, output_qspec) + return PassResult(graph_module, True) diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/quantizer.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..479a2a678f179a298a3ee155afa2947676e93a12 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/quantizer.py @@ -0,0 +1,189 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Callable, Optional, Union + +import torch +from torch import Tensor +from torch.fx import Node + +from torchao.quantization.pt2e import ( + ObserverOrFakeQuantize, + ObserverOrFakeQuantizeConstructor, +) + +__all__ = [ + "Quantizer", + "QuantizationSpecBase", + "QuantizationSpec", + "FixedQParamsQuantizationSpec", + "EdgeOrNode", + "SharedQuantizationSpec", + "DerivedQuantizationSpec", + "QuantizationAnnotation", +] + + +class QuantizationSpecBase(ABC): # noqa: B024 + """Base class for different types of quantization specs that allows users to + specify how to quantize a Tensor (input/output of a Node) in the model + """ + + +@dataclass(eq=True, frozen=True) +class QuantizationSpec(QuantizationSpecBase): + """Quantization spec for common operators that allows user to specify how to + quantize a Tensor, this includes dtype, quant_min, quant_max etc. + """ + + dtype: torch.dtype + # observer or fake_quantize constructor such as + # MinMaxObserver, PerChannelHistogramObserver etc. + # or we can attach some custom args to them + # e.g. MinMaxObserver.with_args(eps=eps) + observer_or_fake_quant_ctr: ObserverOrFakeQuantizeConstructor + quant_min: Optional[int] = None + quant_max: Optional[int] = None + qscheme: Optional[torch.qscheme] = None + ch_axis: Optional[int] = None + is_dynamic: bool = False + + def __post_init__(self): + # TODO: add init for quant_min/quant_max + # quant_min must be less than quant_max + if ( + self.quant_min is not None + and self.quant_max is not None + and self.quant_min > self.quant_max + ): + raise ValueError( + f"quant_min {self.quant_min} must be <= quant_max {self.quant_max}." + ) + + # ch_axis must be less than the number of channels + # but no way to check here. Just check that it is not < 0. + if self.ch_axis is not None and self.ch_axis < 0: + raise ValueError("Ch_axis is < 0.") + + +@dataclass(eq=True, frozen=True) +class FixedQParamsQuantizationSpec(QuantizationSpecBase): + dtype: torch.dtype + scale: float + zero_point: int + quant_min: Optional[int] = None + quant_max: Optional[int] = None + qscheme: Optional[torch.qscheme] = None + is_dynamic: bool = False + + +""" +The way we refer to other points of quantization in the graph will be either +an input edge or an output value +input edge is the connection between input node and the node consuming the input, so it's a Tuple[Node, Node] +output value is an fx Node +""" +EdgeOrNode = Union[tuple[Node, Node], Node] +EdgeOrNode.__module__ = "torchao.quantization.pt2e.quantizer.quantizer" + + +@dataclass(eq=True, frozen=True) +class SharedQuantizationSpec(QuantizationSpecBase): + """ + Quantization spec for the Tensors whose quantization parameters are shared with other Tensors + """ + + # the edge or node to share observer or fake quant instances with + edge_or_node: EdgeOrNode + + +@dataclass(eq=True, frozen=True) +class DerivedQuantizationSpec(QuantizationSpecBase): + """Quantization spec for the Tensors whose quantization parameters are derived from other Tensors""" + + derived_from: list[EdgeOrNode] + derive_qparams_fn: Callable[[list[ObserverOrFakeQuantize]], tuple[Tensor, Tensor]] + dtype: torch.dtype + quant_min: Optional[int] = None + quant_max: Optional[int] = None + qscheme: Optional[torch.qscheme] = None + ch_axis: Optional[int] = None + is_dynamic: bool = False + + +@dataclass +class QuantizationAnnotation: + """How are input arguemnt or output should be quantized, + expressed as QuantizationSpec, this corresponds to how a Tensor in the + operator Graph is observed (PTQ) or fake quantized (QAT) + """ + + # a map from torch.fx.Node to a type of QuantizationSpecBase + input_qspec_map: dict[Node, Optional[QuantizationSpecBase]] = field( + default_factory=dict + ) + + # How the output of this node is quantized, expressed as QuantizationSpec + # TODO: change the value to QuantizationSpec in a separate PR + output_qspec: Optional[QuantizationSpecBase] = None + + # For a Node: node1 and edge: (node1, node2), since they are observing the same + # Tensor, we may want to implicitly share observers, this flag allows people to + # turn off this behavior for the output of the node + allow_implicit_sharing: bool = True + + # whether the node is annotated or not + _annotated: bool = False + + +class Quantizer(ABC): + def transform_for_annotation( + self, model: torch.fx.GraphModule + ) -> torch.fx.GraphModule: + """Allows for user defined transforms to run before annotating the graph. + This allows quantizer to allow quantizing part of the model that are otherwise not quantizable. + For example quantizer can + a) decompose a compound operator like scaled dot product attention, + into bmm and softmax if quantizer knows how to quantize bmm/softmax but not sdpa + or b) transform scalars to tensor to allow quantizing scalares. + + Note: this is an optional method + """ + return model + + # annotate nodes in the graph with observer or fake quant constructors + # to convey the desired way of quantization + @abstractmethod + def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule: + pass + + # validate the annotated graph is supported by the backend + @abstractmethod + def validate(self, model: torch.fx.GraphModule) -> None: + pass + + def prepare_obs_or_fq_callback( + self, + model: torch.fx.GraphModule, + edge_or_node_to_obs_or_fq: dict[EdgeOrNode, ObserverOrFakeQuantize], + ) -> None: + """A callback that will be called after the observers or fake quants are created + for each sharing group, but before they are inserted into the graph. The + callback can be used to make final quantization adjustments, such as enforcing + specific scale and zero point on model input or output. + + Args: + * `model`: the graph module being prepared. + * `edge_or_node_to_obs_or_fq`: a dictionary mapping each annotated edge and + node to the corresponding observer or fake quant object. Note that multiple + edges and/or nodes can map to the same observer / fake quant instance if + they were annotated with SharedQuantizationSpec. This dictionary can be + modified by the callback. + """ + return diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/utils.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f84ae44817a1af658909d3620e2c9a0668672699 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/utils.py @@ -0,0 +1,164 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs + +import typing +from dataclasses import dataclass +from typing import Callable, NamedTuple, Optional + +import torch +from torch.fx import Node + +from .quantizer import QuantizationAnnotation, QuantizationSpec + + +# In the absence of better name, just winging it with QuantizationConfig +@dataclass(eq=True, frozen=True) +class QuantizationConfig: + input_activation: Optional[QuantizationSpec] + output_activation: Optional[QuantizationSpec] + weight: Optional[QuantizationSpec] + bias: Optional[QuantizationSpec] + # TODO: remove, since we can use observer_or_fake_quant_ctr to express this + is_qat: bool = False + + +# Use Annotated because list[Callable].__module__ is read-only. +OperatorPatternType = typing.Annotated[list[Callable], None] +OperatorPatternType.__module__ = "torchao.quantization.pt2e.quantizer.utils" + + +class OperatorConfig(NamedTuple): + # fix List[str] with List[List[Union[nn.Module, FunctionType, BuiltinFunctionType]]] + # Basically we are mapping a quantization config to some list of patterns. + # a pattern is defined as a list of nn module, function or builtin function names + # e.g. [nn.Conv2d, torch.relu, torch.add] + # We have not resolved whether fusion can be considered internal details of the + # quantizer hence it does not need communication to user. + # Note this pattern is not really informative since it does not really + # tell us the graph structure resulting from the list of ops. + config: QuantizationConfig + operators: list[OperatorPatternType] + + +def get_input_act_qspec(quantization_config: Optional[QuantizationConfig]): + if quantization_config is None: + return None + if quantization_config.input_activation is None: + return None + quantization_spec: QuantizationSpec = quantization_config.input_activation + assert quantization_spec.qscheme in [ + torch.per_tensor_affine, + torch.per_tensor_symmetric, + ] + return quantization_spec + + +def get_output_act_qspec(quantization_config: Optional[QuantizationConfig]): + if quantization_config is None: + return None + if quantization_config.output_activation is None: + return None + quantization_spec: QuantizationSpec = quantization_config.output_activation + assert quantization_spec.qscheme in [ + torch.per_tensor_affine, + torch.per_tensor_symmetric, + ] + return quantization_spec + + +def get_weight_qspec(quantization_config: Optional[QuantizationConfig]): + if quantization_config is None: + return None + assert quantization_config is not None + if quantization_config.weight is None: + return None + quantization_spec: QuantizationSpec = quantization_config.weight + if quantization_spec.qscheme not in [ + torch.per_tensor_symmetric, + torch.per_channel_symmetric, + None, + ]: + raise ValueError( + f"Unsupported quantization_spec {quantization_spec} for weight" + ) + return quantization_spec + + +def get_bias_qspec(quantization_config: Optional[QuantizationConfig]): + if quantization_config is None: + return None + assert quantization_config is not None + if quantization_config.bias is None: + return None + quantization_spec: QuantizationSpec = quantization_config.bias + assert quantization_spec.dtype == torch.float, ( + "Only float dtype for bias is supported for bias right now" + ) + return quantization_spec + + +def annotate_input_qspec_map(node: Node, input_node: Node, qspec): + quantization_annotation = node.meta.get( + "quantization_annotation", QuantizationAnnotation() + ) + if quantization_annotation.input_qspec_map is None: + quantization_annotation.input_qspec_map = {} + quantization_annotation.input_qspec_map[input_node] = qspec + node.meta["quantization_annotation"] = quantization_annotation + + +def annotate_output_qspec(node: Node, qspec): + quantization_annotation = node.meta.get( + "quantization_annotation", QuantizationAnnotation() + ) + quantization_annotation.output_qspec = qspec + node.meta["quantization_annotation"] = quantization_annotation + + +def get_module_name_filter(module_name: str): + """Get the module_name_filter function for a given module name, the filter accepts + a node and checks if the node comes from a module that has certain module name + + For example: + node: linear_op = call_function[...](...) # comes from a module with name blocks.sub.linear1 + + + >> module_name_filter = _get_module_name_filter("blocks.sub") + >> print(module_name_filter(node)) + True # the node is from "blocks.sub" based on the fully qualified name "blocks.sub.linear1" + """ + + def module_name_filter(n: Node) -> bool: + # example: { + # 'L__self___sub': ("L['self'].sub", ), + # 'L__self___sub_linear': ("L['self'].sub.linear", ) + # } + # get_attr nodes doesn't have nn_module_stack? + nn_module_stack = n.meta.get("nn_module_stack", {}) + + def _normalize_path(n): + prefix = 0 + # TODO This is non standard behavior and should be removed when we migrate off capture_pre_autograd_graph. + if n.startswith("L['self']."): + prefix = len("L['self'].") + return n[prefix:] + + names = [_normalize_path(n) for n, _ in nn_module_stack.values()] + return module_name in names + + return module_name_filter + + +def is_valid_annotation(annotation: QuantizationAnnotation) -> bool: + if annotation is None: + return False + input_qspec_map = annotation.input_qspec_map + output_qspec = annotation.output_qspec + if len(input_qspec_map) == 0 and output_qspec is None: + return False + return True diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/x86_inductor_quantizer.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/x86_inductor_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..9b49cadf77ce1b6178093ae89fed1b1b2475cbea --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/x86_inductor_quantizer.py @@ -0,0 +1,1627 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import functools +import itertools +import operator +import warnings +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Optional, Union + +import torch +import torch.nn.functional as F +from torch.fx import Node +from torch.fx.passes.utils.source_matcher_utils import ( + SourcePartition, + get_source_partitions, +) +from typing_extensions import TypeAlias + +from torchao.quantization.pt2e import ( + FakeQuantize, + FusedMovingAvgObsFakeQuantize, + HistogramObserver, + MovingAverageMinMaxObserver, + MovingAveragePerChannelMinMaxObserver, + PerChannelMinMaxObserver, + PlaceholderObserver, + find_sequential_partitions, +) +from torchao.quantization.pt2e.quantizer import ( + QuantizationAnnotation, + QuantizationConfig, + QuantizationSpec, + Quantizer, + SharedQuantizationSpec, + get_bias_qspec, + get_input_act_qspec, + get_module_name_filter, + get_output_act_qspec, + get_weight_qspec, +) + +FilterFn: TypeAlias = Callable[[list[Node]], bool] + + +if TYPE_CHECKING: + from torchao.quantization.pt2e import _ObserverOrFakeQuantizeConstructor + +__all__ = [ + "X86InductorQuantizer", + "get_default_x86_inductor_quantization_config", + "get_x86_inductor_linear_dynamic_fp16_config", +] + + +@dataclass +class _X86InductorQuantizationAnnotation(QuantizationAnnotation): + # _is_output_of_quantized_pattern: + # * Node as output node of a fusion pattern. + # * The fusion pattern supports int8 data type. + # * The fusion pattern has inputs annotated to insert observer. + # * The quantization_config is not `None`. + _is_output_of_quantized_pattern: bool = False + + +# Operators that: +# 1. Operators are optimized to run with int8 when int8 input provided. +# 2. Operators do not support int8 input and produce fp32 output. +int8_in_int8_out_ops: set = { + torch.ops.aten.max_pool2d.default, + torch.ops.aten.cat.default, + torch.ops.aten.avg_pool2d.default, + torch.ops.aten.adaptive_avg_pool2d.default, + torch.ops.aten.flatten.using_ints, +} + +# Operators that support the int8 data type for quantization config propagation. +# A superset of int8_in_int8_out_ops incorporating additional operators. +propagation_quantizable_ops = int8_in_int8_out_ops + +# Operators support the int8 data type +# and recipe is configured by default in X86InductorQuantizer. +default_quantizable_ops = propagation_quantizable_ops | { + torch.ops.aten.conv1d.default, + torch.ops.aten.conv2d.default, + torch.ops.aten.linear.default, + torch.ops.aten.mul.Tensor, +} + +# A superset of default_quantizable_ops includes operators support the int8 data type +# but not enabled by default recipe of X86InductorQuantizer. +quantizable_ops = default_quantizable_ops | { + torch.ops.aten.matmul.default, +} + +QUANT_ANNOTATION_KEY = "quantization_annotation" + + +def _skip_annotate(nodes: list[Node], filter_fn: Optional[FilterFn] = None) -> bool: + """Determine whether to skip annotation for a list of nodes.""" + + # 1) Skip annotate if any node is already annotated + if _is_any_annotated(nodes): + return True + + # 2) Proceed annotate if a) a filter function is provided + # and b) the given nodes list passes the filter function check. + if filter_fn and filter_fn(nodes): + return False + + return True + + +def _create_module_name_filter(module_name: str) -> FilterFn: + """Create a filter function for a given module name. + + The filter function takes a list of nodes (as determined by the annotate function) + and return True if *all* nodes come from the specified module name, False otherwise. + + For example: + linear_1: "f32[3, 10]" = torch.ops.aten.linear.default(...) # comes from a module with name `sub.linear1` + relu: "f32[3, 10]" = torch.ops.aten.relu.default(linear_1); # comes from a module with name `sub.relu1` + + >> module_name_filter = _create_module_name_filter_inner("sub") + >> print(module_name_filter([relu, linear_1])) + # True # These two nodes are determined by `_annotate_linear_unary` function and from "sub". + """ + + filter_fn = get_module_name_filter(module_name) + + def check_all_nodes_from_module(nodes: list[Node]) -> bool: + all_nodes_from_module_name: bool = all(filter_fn(n) for n in nodes) + return all_nodes_from_module_name + + return check_all_nodes_from_module + + +def _create_operator_type_filter( + operator_type: Callable, +) -> FilterFn: + """Create a filter function for a given operator type. + + The filter function takes a list of nodes and returns True if it contains + exactly one node with the specified operator type, False otherwise. + + For example: + linear_1: "f32[3, 10]" = torch.ops.aten.linear.default(...) # comes from a module with name `sub.linear1` + relu: "f32[3, 10]" = torch.ops.aten.relu.default(linear_1); # comes from a module with name `sub.relu1` + + >> operator_type_filter = _create_operator_type_filter(torch.ops.aten.linear.default) + >> print(operator_type_filter([relu, linear_1])) + # True # These two nodes are determined by `_annotate_linear_unary` function and the second node is `linear`. + """ + + def operator_type_filter(nodes: list[Node]): + num_nodes_with_operator_type = sum( + node.target == operator_type for node in nodes + ) + if num_nodes_with_operator_type > 1: + raise NotImplementedError( + f"Several nodes within a single pattern are {operator_type}." + ) + return num_nodes_with_operator_type == 1 + + return operator_type_filter + + +def _global_config_filter(nodes: list[Node]) -> bool: + """Filter function for global configuration. + + This filter function takes a list of nodes and returns True if there is exactly one node + in the list that is a default quantizable operation, False otherwise. + """ + num_nodes_in_default_quantizable_ops = sum( + node.target in default_quantizable_ops for node in nodes + ) + if num_nodes_in_default_quantizable_ops > 1: + raise NotImplementedError( + "Several nodes within a single pattern are default quantizable operations." + ) + return num_nodes_in_default_quantizable_ops == 1 + + +def _map_module_function_to_aten_operator_type(): + module_function_to_aten_operator: dict[Callable, torch._ops.OpOverloadPacket] = {} + map_list = ( + ([torch.nn.Conv2d, F.conv1d], torch.ops.aten.conv1d.default), + ([torch.nn.Conv2d, F.conv2d], torch.ops.aten.conv2d.default), + ([torch.nn.Linear, F.linear], torch.ops.aten.linear.default), + ([torch.nn.MaxPool2d, F.max_pool2d], torch.ops.aten.max_pool2d.default), + ( + [ + torch.cat, + ], + torch.ops.aten.cat.default, + ), + ([torch.nn.AvgPool2d, F.avg_pool2d], torch.ops.aten.avg_pool2d.default), + ( + [torch.nn.AdaptiveAvgPool2d, F.adaptive_avg_pool2d], + torch.ops.aten.adaptive_avg_pool2d.default, + ), + ( + [ + torch.flatten, + ], + torch.ops.aten.flatten.using_ints, + ), + ( + [ + torch.matmul, + ], + torch.ops.aten.matmul.default, + ), + ( + [ + torch.mul, + ], + torch.ops.aten.mul.Tensor, + ), + ) + for map_item in map_list: + module_function_to_aten_operator.update(dict.fromkeys(map_item[0], map_item[1])) # type: ignore[arg-type, call-overload] + return module_function_to_aten_operator + + +def _mark_nodes_as_annotated(nodes: list[Node]): + for node in nodes: + if node is not None: + if QUANT_ANNOTATION_KEY not in node.meta: + node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation() + node.meta[QUANT_ANNOTATION_KEY]._annotated = True + + +def _is_node_annotated(_node): + """ + return True if the node is annotated, otherwise return False + """ + return ( + QUANT_ANNOTATION_KEY in _node.meta + and _node.meta[QUANT_ANNOTATION_KEY]._annotated + ) + + +def _is_any_annotated(nodes: list[Node]): + """ + Given a list of nodes (that represents an operator pattern), + check if any of the node is annotated, return True if any of the node + is annotated, otherwise return False. + """ + return any(_is_node_annotated(node) for node in nodes) + + +def _is_all_annotated(nodes: list[Node]): + """ + Given a list of nodes (that represents an operator pattern), + return True if all of the node is annotated, otherwise return False. + """ + return all(_is_node_annotated(node) for node in nodes) + + +def _is_quantized_op_pt2e(node: torch.fx.Node): + """ + Used for pt2e flow to check if the node is a quantized node: + Case1: the node has been annotated as output node of a fusion pattern. + Case2: the node has been annotated as single quantized node. + """ + if not _is_any_annotated([node]): + # The node has not been annotated, directly return False + return False + quantization_annotation = node.meta.get(QUANT_ANNOTATION_KEY, None) + assert isinstance(quantization_annotation, _X86InductorQuantizationAnnotation) + return quantization_annotation._is_output_of_quantized_pattern + + +@functools.lru_cache +def get_default_x86_inductor_quantization_config( + is_qat: bool = False, + is_dynamic: bool = False, + reduce_range: bool = False, +): + """ + reduce_range is False by default. Set it to True on earlier CPUs without VNNI to avoid accuracy issue. + """ + extra_args: dict[str, Any] = {"eps": 2**-12} + if is_qat: + if is_dynamic: + act_observer_or_fake_quant_ctr = FakeQuantize + dynamic_quant_observer = MovingAverageMinMaxObserver.with_args( + averaging_constant=1 + ) + extra_args["observer"] = dynamic_quant_observer + else: + act_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize # type: ignore[assignment] + else: + if is_dynamic: + act_observer_or_fake_quant_ctr = PlaceholderObserver # type: ignore[assignment] + else: + act_observer_or_fake_quant_ctr = HistogramObserver # type: ignore[assignment] + + # Copy from x86 default qconfig from torch/ao/quantization/qconfig.py + act_quantization_spec = QuantizationSpec( + dtype=torch.uint8, + quant_min=0, + quant_max=127 if reduce_range else 255, + qscheme=torch.per_tensor_affine, + is_dynamic=is_dynamic, + observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args( + **extra_args + ), + ) + + weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = ( + FusedMovingAvgObsFakeQuantize if is_qat else PerChannelMinMaxObserver + ) + + if is_qat: + # Only support per channel quant for now + extra_args["observer"] = MovingAveragePerChannelMinMaxObserver # type: ignore[dict-item] + weight_quantization_spec = QuantizationSpec( + dtype=torch.int8, + quant_min=-128, + quant_max=127, + qscheme=torch.per_channel_symmetric, + ch_axis=0, # 0 corresponding to weight shape = (oc, ic, kh, kw) of conv + is_dynamic=False, + observer_or_fake_quant_ctr=weight_observer_or_fake_quant_ctr.with_args( + **extra_args + ), + ) + bias_quantization_spec = None # will use placeholder observer by default + quantization_config = QuantizationConfig( + act_quantization_spec, + act_quantization_spec, + weight_quantization_spec, + bias_quantization_spec, + is_qat, + ) + return quantization_config + + +@functools.lru_cache +def get_x86_inductor_linear_dynamic_fp16_config(): + """ + For linear_dynamic_fp16. The name may be confusing. + The op's behavior is fp32_input * (fp16_weight -> to_fp32) -> fp32_output. + """ + weight_quantization_spec = QuantizationSpec( + dtype=torch.float16, + observer_or_fake_quant_ctr=PlaceholderObserver, + ) + quantization_config = QuantizationConfig( + None, # input_quantization_spec + None, # output_quantization_spec + weight_quantization_spec, + None, # bias_quantization_spec + ) + return quantization_config + + +def _annotate_nodes_not_quantize(nodes: Union[Node, list[Node]]) -> None: + """Annotate nodes to exclude them from quantization (their `quantization_config` is `None`).""" + if not isinstance(nodes, list): + nodes = [nodes] + for node in nodes: + node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + _annotated=True + ) + + +def _config_checker(method: Callable) -> Callable: + @functools.wraps(method) + def wrapper( + quantizer: "X86InductorQuantizer", + name: Any, + quantization_config: Optional["QuantizationConfig"], + ) -> "X86InductorQuantizer": + if quantizer._need_skip_config(quantization_config): + warnings.warn( + f"Skip the quantization config for {name}.", + ) + return quantizer + return method(quantizer, name, quantization_config) + + return wrapper + + +@dataclass +class _CurrentQuantizationMode: + r"""Configuration defining the current quantization mode for the quantizer. + + All possible current quantization modes are listed below: + ---------------------------------------------------------------------------------------------------------- + | dynamic_state + qat_state |--------------------------------------------------------------------------------------------- + | None | True | False + ---------------------------------------------------------------------------------------------------------- + None | quantizer does not receive a non-None `quantization_config` | \ | \ + False | quantizer will not do QAT | dynamic | static + True | quantizer will do QAT | QAT + dynamic | QAT + static + """ + + qat_state: Optional[bool] + dynamic_state: Optional[bool] + + +class X86InductorQuantizer(Quantizer): + module_function_to_aten_operator_type = _map_module_function_to_aten_operator_type() + + def __init__(self) -> None: + super().__init__() + self.global_config: Optional[QuantizationConfig] = None + self.operator_type_qconfig: dict[ + torch._ops.OpOverloadPacket, Optional[QuantizationConfig] + ] = {} + self.module_name_qconfig: dict[str, Optional[QuantizationConfig]] = {} + + def _get_current_quantization_mode(self) -> _CurrentQuantizationMode: + """Retrieves the current quantization mode based on all configurations.""" + qat_state = None + dynamic_state = None + + # As we use `_need_skip_config` to skip all invalid configurations, + # we can safely assume that the all existing non-None configurations + # have the same quantization mode. + for qconfig in ( + list(self.module_name_qconfig.values()) + + list(self.operator_type_qconfig.values()) + + [self.global_config] + ): + if qconfig is not None: + # Query the `is_qat` state + if qat_state is None: + qat_state = qconfig.is_qat + else: + assert qat_state == qconfig.is_qat, ( + f"All non-None quantization configs should have the same `is_qat`," + f"but got {qat_state} and {qconfig.is_qat}." + ) + # Query the `is_dynamic` state + input_activation_spec = qconfig.input_activation + if input_activation_spec is not None: + if dynamic_state is None: + dynamic_state = input_activation_spec.is_dynamic + else: + assert dynamic_state == input_activation_spec.is_dynamic, ( + f"All non-None `input_activation_spec` should have the same `is_dynamic`," + f"but got {dynamic_state} and {input_activation_spec.is_dynamic}." + ) + return _CurrentQuantizationMode( + qat_state=qat_state, dynamic_state=dynamic_state + ) + + def _need_skip_config( + self, quantization_config: Optional[QuantizationConfig] + ) -> bool: + """Check if the provided quantization config is valid for X86InductorQuantizer. + + Mixed static/dynamic configurations or mixed QAT/non-QAT configurations are not supported. + To avoid such a mix, we compare the incoming configuration with current configuration status. + Refer the `_CurrentQuantizationMode` definition for all possible modes. + """ + if quantization_config is None: + return False + + need_skip = False + current_mode = self._get_current_quantization_mode() + if ( + current_mode.qat_state is not None + and current_mode.qat_state != quantization_config.is_qat + ): + warnings.warn("Mixed QAT and Non-QAT quantization config is not supported.") + need_skip = True + if current_mode.dynamic_state is not None: + input_activation_spec = quantization_config.input_activation + if ( + input_activation_spec is not None + and current_mode.dynamic_state != input_activation_spec.is_dynamic + ): + warnings.warn( + "Mixed dynamic and static quantization config is not supported." + ) + need_skip = True + return need_skip + + def set_global(self, quantization_config: QuantizationConfig): + if self._need_skip_config(quantization_config): + warnings.warn("Skip the global quantization config.") + return self + self.global_config = quantization_config + return self + + def get_global_quantization_config(self): + if not isinstance(self.global_config, QuantizationConfig): + warnings.warn( + "The global_config for X86InductorQuantizer is currently invalid. \ + Please ensure that you use set_global to establish the global quantization configuration." + ) + return self.global_config + + @_config_checker + def set_function_type_qconfig( + self, + function_type: Callable, + quantization_config: Optional[QuantizationConfig], + ) -> "X86InductorQuantizer": + if function_type in X86InductorQuantizer.module_function_to_aten_operator_type: + self._set_aten_operator_qconfig( + X86InductorQuantizer.module_function_to_aten_operator_type[ + function_type + ], + quantization_config, + ) + else: + warnings.warn( + f"function: Unable to customize quantization config for {function_type} by X86InductorQuantizer." + ) + return self + + @_config_checker + def set_module_type_qconfig( + self, + module_type: torch.nn.Module, + quantization_config: Optional[QuantizationConfig], + ) -> "X86InductorQuantizer": + if module_type in X86InductorQuantizer.module_function_to_aten_operator_type: + self._set_aten_operator_qconfig( + X86InductorQuantizer.module_function_to_aten_operator_type[module_type], + quantization_config, + ) + else: + warnings.warn( + f"Module: Unable to customize quantization config for {module_type} by X86InductorQuantizer." + ) + return self + + @_config_checker + def set_module_name_qconfig( + self, module_name: str, quantization_config: Optional[QuantizationConfig] + ): + """Set quantization_config for a submodule with name: `module_name`, for example: + quantizer.set_module_name_qconfig("blocks.sub"), it will quantize all supported operator/operator + patterns in the submodule with this module name with the given `quantization_config` + + The supported operators include `quantizable_ops` and `propagation_quantizable_ops`. + """ + self.module_name_qconfig[module_name] = quantization_config + return self + + def _set_aten_operator_qconfig( + self, + operator_type: torch._ops.OpOverloadPacket, + quantization_config: Optional[QuantizationConfig], + ) -> "X86InductorQuantizer": + if operator_type in quantizable_ops: + self.operator_type_qconfig[operator_type] = quantization_config + else: + warnings.warn( + f"operator: Unable to quantize {operator} by X86InductorQuantizer." + ) + return self + + def _annotate_conv_node_helper( + self, + conv_node: torch.fx.Node, + annotate_output: bool, + quantization_config: Optional[QuantizationConfig], + ) -> None: + """Helper function to annotate the conv node""" + if quantization_config is None: + _annotate_nodes_not_quantize(conv_node) + return + input_qspec_map = {} + input_node = conv_node.args[0] + assert isinstance(input_node, Node) + input_qspec_map[input_node] = get_input_act_qspec(quantization_config) + weight_node = conv_node.args[1] + assert isinstance(weight_node, Node) + input_qspec_map[weight_node] = get_weight_qspec(quantization_config) + bias_node = None if len(conv_node.args) == 2 else conv_node.args[2] + if isinstance(bias_node, Node): + input_qspec_map[bias_node] = get_bias_qspec(quantization_config) + if annotate_output: + conv_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + else: + conv_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + ) + + def _annotate_linear_node_helper( + self, + linear_node: torch.fx.Node, + annotate_output: bool, + quantization_config: Optional[QuantizationConfig], + ) -> None: + """Helper function to annotate the linear node""" + if quantization_config is None: + _annotate_nodes_not_quantize(linear_node) + return + input_qspec_map = {} + assert linear_node.target in (torch.ops.aten.linear.default,) + has_bias = len(linear_node.args) == 3 + input_index = 0 + weight_index = 1 + bias_index = 2 + + input_node = linear_node.args[input_index] + assert isinstance(input_node, Node) + input_qspec_map[input_node] = get_input_act_qspec(quantization_config) + + weight_node = linear_node.args[weight_index] + assert isinstance(weight_node, Node) + input_qspec_map[weight_node] = get_weight_qspec(quantization_config) + + bias_node = linear_node.args[bias_index] if has_bias else None + if isinstance(bias_node, Node): + input_qspec_map[bias_node] = get_bias_qspec(quantization_config) + + if annotate_output: + linear_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + else: + linear_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, _annotated=True + ) + + def _get_output_nodes_of_partitions( + self, + partition_list: list[SourcePartition], + ) -> list[torch.fx.Node]: + """Helper function to get the output node list from partition list""" + output_node_list = [] + for partition in partition_list: + if len(partition.output_nodes) > 1: + raise ValueError("Input partition has more than one output node") + output_node = partition.output_nodes[0] + assert isinstance(output_node, Node) + output_node_list.append(output_node) + if len(output_node_list) != len(partition_list): + raise ValueError( + "length of output_node_list should equal to length of partition_list" + ) + return output_node_list + + def _get_input_idx_for_binary_node( + self, + conv_gemm_node: torch.fx.Node, + binary_node: torch.fx.Node, + ): + """Helper function to check conv_gemm and extra input node index + for binary node fused with conv_gemm. + """ + conv_gemm_node_idx = None + extra_input_node_idx = None + if (binary_node.args[0].op == "call_function") and ( # type: ignore[union-attr] + binary_node.args[0] == conv_gemm_node + ): + conv_gemm_node_idx = 0 + extra_input_node_idx = 1 + elif (binary_node.args[1].op == "call_function") and ( # type: ignore[union-attr] + binary_node.args[1] == conv_gemm_node + ): + conv_gemm_node_idx = 1 + extra_input_node_idx = 0 + extra_input_node = binary_node.args[extra_input_node_idx] # type: ignore[index] + assert isinstance(extra_input_node, Node) + return conv_gemm_node_idx, extra_input_node_idx + + def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule: + """Annotate the given model with quantization configurations. + + Annotation contracts: + 1. Annotate each node according to the user's qconfig in the following order: + `module_name_qconfig`, `operator_type_qconfig`, and `global_config`. + 2. Avoid re-annotating nodes already annotated in prior stages. For example, + if `linear1` has been annotated by `module_name_qconfig`, it won't be annotated again + during the processing of the 'operator_type_qconfig' or 'global_config'. + 3. For config is `None`, the node will be annotated with `_X86InductorQuantizationAnnotation(_annotated=True)`. + + For each pair of (module_name_or_operator_type_or_global, qconfig), a filter function is created. + This filter function checks if the node is marked by current stage and not annotated by the previous stage. + """ + for module_name, quantization_config in self.module_name_qconfig.items(): + self._annotate_with_config( + model, quantization_config, _create_module_name_filter(module_name) + ) + + for operator_type, quantization_config in self.operator_type_qconfig.items(): + self._annotate_with_config( + model, quantization_config, _create_operator_type_filter(operator_type) + ) + + if self.global_config: + self._annotate_with_config( + model, + self.global_config, + _global_config_filter, + ) + + # Once we've annotated the model with quantization configurations, we also need to annotate + # the output of quantizable operations. For example, if we annotated `maxpool2d` to quantize its inputs, + # we will quantize its output accordingly. This enables us to fuse the dq-operator-q into a quantized op. + # Refer to https://github.com/intel/intel-extension-for-pytorch/blob/ + # 90d19323d96afc53fcc22ba5a7bb3fb07fdd6c1c/intel_extension_for_pytorch/quantization/_recipe.py#L487 + + self._annotate_output_for_int8_in_int8_out_pattern_entry(model) + + return model + + def _annotate_with_config( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: FilterFn, + ) -> None: + """Annotate the model with the given quantization configuration. + + High-level description of quantization recipe for X86 Inductor Backend: + Step 1: Apply quantization recipe for fusion patterns of conv/linear to enable int8 data type actively. + Step 2: Propagate quantization annotation for patterns besides conv/linear. Go through the pattern in model + from start to the end. If a pattern supports computation with int8 data type and inputs connected to + quantized patterns, annotate its inputs as quantized pattern. + """ + + # Step1: Recipe of fusion patterns like conv/linear. + self._annotate_conv2d_fusion_pattern(model, quantization_config, filter_fn) + self._annotate_linear_fusion_pattern(model, quantization_config, filter_fn) + self._annotate_matmul(model, quantization_config, filter_fn) + self._annotate_mul_tensor(model, quantization_config, filter_fn) + + # Step2: Recipe to propagate annotation for patterns beside conv/linear. + # Go through all the nodes from start to end. + # Recipe refer to https://github.com/intel/intel-extension-for-pytorch/blob/ + # 90d19323d96afc53fcc22ba5a7bb3fb07fdd6c1c/intel_extension_for_pytorch/quantization/_recipe.py#L538 + + self._annotate_propagation_quantizable_pattern_entry( + model, quantization_config, filter_fn + ) + + def _annotate_qat_conv2d_fusion_pattern( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ): + # Annotate QAT Specific patterns + self._annotate_qat_conv2d_bn_binary_unary(model, quantization_config, filter_fn) + self._annotate_qat_conv2d_bn_binary(model, quantization_config, filter_fn) + self._annotate_qat_conv2d_bn_unary(model, quantization_config, filter_fn) + self._annotate_qat_conv2d_bn(model, quantization_config, filter_fn) + + def _annotate_qat_conv2d_bn_binary_unary( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + fused_partitions = find_sequential_partitions( + gm, [torch.nn.Conv2d, torch.nn.BatchNorm2d, operator.add, torch.nn.ReLU] + ) + for fused_partition in fused_partitions: + ( + conv_partition, + bn_partition, + binary_partition, + unary_partition, + ) = fused_partition + + ( + conv_node, + bn_output_node, + binary_node, + unary_node, + ) = self._get_output_nodes_of_partitions( + [conv_partition, bn_partition, binary_partition, unary_partition] + ) + if len(bn_output_node.users) != 1: + # Conv BN pattern should only has 1 user. + continue + ( + bn_output_node_idx, + extra_input_node_idx, + ) = self._get_input_idx_for_binary_node(bn_output_node, binary_node) + if (bn_output_node_idx is None) or (extra_input_node_idx is None): + continue + if bn_output_node != binary_node.args[bn_output_node_idx]: + raise ValueError(f"{bn_output_node} doesn't match input of binary node") + extra_input_node = binary_node.args[extra_input_node_idx] + + if ( + conv_node.op != "call_function" + or conv_node.target != torch.ops.aten.conv2d.default + ): + continue + + if _skip_annotate( + [unary_node, binary_node, bn_output_node, conv_node], filter_fn + ): + continue + + self._annotate_conv_node_helper(conv_node, False, quantization_config) + + if quantization_config is not None: + binary_node_input_qspec_map = {} + binary_node_input_qspec_map[extra_input_node] = get_input_act_qspec( + quantization_config + ) + binary_node.meta[QUANT_ANNOTATION_KEY] = ( + _X86InductorQuantizationAnnotation( + input_qspec_map=binary_node_input_qspec_map, + _annotated=True, + ) + ) + unary_node.meta[QUANT_ANNOTATION_KEY] = ( + _X86InductorQuantizationAnnotation( + # TODO Remove the annotate of output in QAT when qat util support pattern matcher. + output_qspec=get_output_act_qspec(quantization_config), # type: ignore[arg-type] + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + ) + else: + _annotate_nodes_not_quantize([binary_node, unary_node]) + nodes_to_mark_annotated = list(conv_partition.nodes) + nodes_to_mark_annotated.extend(list(bn_partition.nodes)) + nodes_to_mark_annotated.extend(list(binary_partition.nodes)) + nodes_to_mark_annotated.extend(list(unary_partition.nodes)) + _mark_nodes_as_annotated(nodes_to_mark_annotated) + + def _annotate_qat_conv2d_bn_binary( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + fused_partitions = find_sequential_partitions( + gm, [torch.nn.Conv2d, torch.nn.BatchNorm2d, operator.add] + ) + for fused_partition in fused_partitions: + conv_partition, bn_partition, binary_partition = fused_partition + ( + conv_node, + bn_output_node, + binary_node, + ) = self._get_output_nodes_of_partitions( + [conv_partition, bn_partition, binary_partition] + ) + if len(bn_output_node.users) != 1: + # Conv BN pattern should only has 1 user. + continue + ( + bn_output_node_idx, + extra_input_node_idx, + ) = self._get_input_idx_for_binary_node(bn_output_node, binary_node) + if (bn_output_node_idx is None) or (extra_input_node_idx is None): + continue + if bn_output_node != binary_node.args[bn_output_node_idx]: + raise ValueError(f"{bn_output_node} doesn't match input of binary node") + + extra_input_node = binary_node.args[extra_input_node_idx] + + if ( + conv_node.op != "call_function" + or conv_node.target != torch.ops.aten.conv2d.default + ): + continue + + if _skip_annotate([binary_node, bn_output_node, conv_node], filter_fn): + continue + + self._annotate_conv_node_helper(conv_node, False, quantization_config) + + if quantization_config is not None: + binary_node_input_qspec_map = {} + binary_node_input_qspec_map[extra_input_node] = get_input_act_qspec( + quantization_config + ) + binary_node.meta[QUANT_ANNOTATION_KEY] = ( + _X86InductorQuantizationAnnotation( + input_qspec_map=binary_node_input_qspec_map, + # TODO Remove the annotate of output in QAT when qat util support pattern matcher. + output_qspec=get_output_act_qspec(quantization_config), # type: ignore[arg-type] + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + ) + else: + _annotate_nodes_not_quantize(binary_node) + nodes_to_mark_annotated = list(conv_partition.nodes) + nodes_to_mark_annotated.extend(list(bn_partition.nodes)) + nodes_to_mark_annotated.extend(list(binary_partition.nodes)) + _mark_nodes_as_annotated(nodes_to_mark_annotated) + + def _annotate_qat_conv2d_bn_unary( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + fused_partitions = [] + unary_patterns = [ + [torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU], + [torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.Hardtanh], + [torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.Hardswish], + [torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU6], + [torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.SiLU], + ] + for unary_pattern in unary_patterns: + partitions = find_sequential_partitions(gm, unary_pattern) + if partitions: + # Extend the fused_partitions if partitions is not empty + fused_partitions.extend(partitions) + + for fused_partition in fused_partitions: + conv_partition, bn_partition, unary_partition = fused_partition + ( + conv_node, + bn_output_node, + unary_node, + ) = self._get_output_nodes_of_partitions( + [conv_partition, bn_partition, unary_partition] + ) + + if ( + conv_node.op != "call_function" + or conv_node.target != torch.ops.aten.conv2d.default + ): + continue + + if _skip_annotate([unary_node, bn_output_node, conv_node], filter_fn): + continue + + self._annotate_conv_node_helper(conv_node, False, quantization_config) + if quantization_config is not None: + unary_node.meta[QUANT_ANNOTATION_KEY] = ( + _X86InductorQuantizationAnnotation( + # TODO Remove the annotate of output in QAT when qat util support pattern matcher. + output_qspec=get_output_act_qspec(quantization_config), # type: ignore[arg-type] + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + ) + else: + _annotate_nodes_not_quantize(unary_node) + nodes_to_mark_annotated = list(conv_partition.nodes) + nodes_to_mark_annotated.extend(list(bn_partition.nodes)) + nodes_to_mark_annotated.extend(list(unary_partition.nodes)) + _mark_nodes_as_annotated(nodes_to_mark_annotated) + + def _annotate_qat_conv2d_bn( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + fused_partitions = find_sequential_partitions( + gm, [torch.nn.Conv2d, torch.nn.BatchNorm2d] + ) + for fused_partition in fused_partitions: + conv_partition, bn_partition = fused_partition + conv_node, bn_output_node = self._get_output_nodes_of_partitions( + [conv_partition, bn_partition] + ) + + if ( + conv_node.op != "call_function" + or conv_node.target != torch.ops.aten.conv2d.default + ): + continue + + if _skip_annotate([bn_output_node, conv_node], filter_fn): + continue + + self._annotate_conv_node_helper(conv_node, False, quantization_config) + if quantization_config is not None: + bn_output_node.meta[QUANT_ANNOTATION_KEY] = ( + _X86InductorQuantizationAnnotation( + # TODO Remove the annotate of output in QAT when qat util support pattern matcher. + output_qspec=get_output_act_qspec(quantization_config), # type: ignore[arg-type] + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + ) + else: + _annotate_nodes_not_quantize(bn_output_node) + nodes_to_mark_annotated = list(conv_partition.nodes) + nodes_to_mark_annotated.extend(list(bn_partition.nodes)) + _mark_nodes_as_annotated(nodes_to_mark_annotated) + + def _annotate_conv2d_fusion_pattern( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ): + if (quantization_config is None) or (quantization_config.is_qat): + # Annotate QAT specific pattern: mainly due to BN not folded in prepare_qat + self._annotate_qat_conv2d_fusion_pattern( + model, quantization_config, filter_fn + ) + self._annotate_conv2d_binary_unary(model, quantization_config, filter_fn) + self._annotate_conv2d_binary(model, quantization_config, filter_fn) + self._annotate_conv2d_unary(model, quantization_config, filter_fn) + self._annotate_conv2d(model, quantization_config, filter_fn) + + def _annotate_linear_fusion_pattern( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ): + self._annotate_linear_binary_unary(model, quantization_config, filter_fn) + self._annotate_linear_unary(model, quantization_config, filter_fn) + self._annotate_linear(model, quantization_config, filter_fn) + + def _annotate_matmul( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ): + for node in model.graph.nodes: + if node.target != torch.ops.aten.matmul.default: + continue + if _skip_annotate([node], filter_fn): + continue + + if quantization_config is None: + _annotate_nodes_not_quantize(node) + continue + + input_qspec_map = {} + matmul_node = node + for input_node in matmul_node.args: + input_qspec_map[input_node] = get_input_act_qspec(quantization_config) + matmul_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + + def _annotate_conv2d_binary_unary( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + # Conv2d + add + unary op + fused_partitions = find_sequential_partitions( + gm, [torch.nn.Conv2d, operator.add, torch.nn.ReLU] + ) + for fused_partition in fused_partitions: + conv_partition, binary_partition, unary_partition = fused_partition + conv_node, binary_node, unary_node = self._get_output_nodes_of_partitions( + [conv_partition, binary_partition, unary_partition] + ) + if len(conv_node.users) != 1: + # Conv Node should only has 1 user node + continue + conv_node_idx, extra_input_node_idx = self._get_input_idx_for_binary_node( + conv_node, binary_node + ) + if (conv_node_idx is None) or (extra_input_node_idx is None): + continue + if conv_node != binary_node.args[conv_node_idx]: + raise ValueError(f"{conv_node} doesn't match input of binary node") + extra_input_node = binary_node.args[extra_input_node_idx] + if ( + conv_node.op != "call_function" + or conv_node.target != torch.ops.aten.conv2d.default + ): + # No conv node found to be fused with add + continue + if _skip_annotate([unary_node, binary_node, conv_node], filter_fn): + continue + + if quantization_config is None: + _annotate_nodes_not_quantize([conv_node, binary_node, unary_node]) + continue + + self._annotate_conv_node_helper(conv_node, False, quantization_config) + binary_node_input_qspec_map = {} + binary_node_input_qspec_map[extra_input_node] = get_input_act_qspec( + quantization_config + ) + binary_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=binary_node_input_qspec_map, + _annotated=True, + ) + unary_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + + def _annotate_conv2d_binary( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + # Conv2d + add + fused_partitions = find_sequential_partitions( + gm, [torch.nn.Conv2d, operator.add] + ) + for fused_partition in fused_partitions: + conv_partition, binary_partition = fused_partition + conv_node, binary_node = self._get_output_nodes_of_partitions( + [conv_partition, binary_partition] + ) + if len(conv_node.users) != 1: + # Conv Node should only has 1 user node + continue + conv_node_idx, extra_input_node_idx = self._get_input_idx_for_binary_node( + conv_node, binary_node + ) + if (conv_node_idx is None) or (extra_input_node_idx is None): + continue + if conv_node != binary_node.args[conv_node_idx]: + raise ValueError(f"{conv_node} doesn't match input of binary node") + extra_input_node = binary_node.args[extra_input_node_idx] + assert isinstance(conv_node, Node) + if ( + conv_node.op != "call_function" + or conv_node.target != torch.ops.aten.conv2d.default + ): + # No conv node found to be fused with add + continue + if _skip_annotate([binary_node, conv_node], filter_fn): + continue + + if quantization_config is None: + _annotate_nodes_not_quantize([conv_node, binary_node]) + continue + + self._annotate_conv_node_helper(conv_node, False, quantization_config) + binary_node_input_qspec_map = {} + binary_node_input_qspec_map[extra_input_node] = get_input_act_qspec( + quantization_config + ) + binary_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=binary_node_input_qspec_map, + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + + def _annotate_conv2d_unary( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + fused_partitions = [] + unary_patterns = [ + [torch.nn.Conv2d, torch.nn.ReLU], + [torch.nn.Conv2d, torch.nn.Hardtanh], + [torch.nn.Conv2d, torch.nn.Hardswish], + [torch.nn.Conv2d, torch.nn.ReLU6], + [torch.nn.Conv2d, torch.nn.SiLU], + [torch.nn.Conv1d, torch.nn.ReLU], + ] + for unary_pattern in unary_patterns: + partitions = find_sequential_partitions(gm, unary_pattern) + if partitions: + # Extend the fused_partitions if partitions is not empty + fused_partitions.extend(partitions) + + for fused_partition in fused_partitions: + conv_partition, unary_partition = fused_partition + conv_node, unary_node = self._get_output_nodes_of_partitions( + [conv_partition, unary_partition] + ) + if conv_node.op != "call_function" or conv_node.target not in ( + torch.ops.aten.conv2d.default, + torch.ops.aten.conv1d.default, + ): + continue + if _skip_annotate([unary_node, conv_node], filter_fn): + continue + + if quantization_config is None: + _annotate_nodes_not_quantize([conv_node, unary_node]) + continue + + self._annotate_conv_node_helper(conv_node, False, quantization_config) + unary_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + + def _annotate_conv2d( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + conv_partitions = get_source_partitions( + gm.graph, [torch.nn.Conv2d, torch.nn.functional.conv2d] + ) + conv_partitions = list(itertools.chain.from_iterable(conv_partitions.values())) + for conv_partition in conv_partitions: + if len(conv_partition.output_nodes) > 1: + raise ValueError("conv partition has more than one output node") + conv_node = conv_partition.output_nodes[0] + if ( + conv_node.op != "call_function" + or conv_node.target != torch.ops.aten.conv2d.default + ): + raise ValueError(f"{conv_node} is not an aten conv2d operator") + # skip annotation if it is already annotated + if _skip_annotate([conv_node], filter_fn): + continue + self._annotate_conv_node_helper(conv_node, True, quantization_config) + + def _annotate_maxpool2d( + self, + node: Node, + quantization_config: Optional[QuantizationConfig], + ) -> None: + if node.target is not torch.ops.aten.max_pool2d.default: + return + if quantization_config is None: + _annotate_nodes_not_quantize(node) + return + + maxpool_node = node + if _is_any_annotated( + [ + maxpool_node, + ] + ): + return + + input_node = maxpool_node.args[0] + assert isinstance(input_node, Node) + input_qspec_map = {} + input_qspec_map[input_node] = get_input_act_qspec(quantization_config) + maxpool_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + + def _annotate_cat( + self, node: Node, quantization_config: QuantizationConfig + ) -> None: + if quantization_config is None: + _annotate_nodes_not_quantize(node) + return + cat_node = node + input_nodes = cat_node.args[0] + assert isinstance(input_nodes, Sequence) + first_input_node = input_nodes[0] + input_qspec_map = {} + assert isinstance(first_input_node, Node) + assert isinstance(cat_node, Node) + input_qspec_map[first_input_node] = get_input_act_qspec(quantization_config) + share_qparams_with_input_act0_qspec = SharedQuantizationSpec( + (first_input_node, cat_node) + ) + + for input_node in input_nodes[1:]: + if input_node not in input_qspec_map: + # There has the case of cat same nodes: torch.cat([input0, input0], 1) + assert isinstance(input_node, Node) + input_qspec_map[input_node] = share_qparams_with_input_act0_qspec + + cat_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + + def _annotate_propagation_quantizable_pattern_entry( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ): + for node in gm.graph.nodes: + self._annotate_propagation_quantizable_pattern( + node, quantization_config, filter_fn + ) + + def _annotate_propagation_quantizable_pattern( + self, node: Node, quantization_config, filter_fn + ) -> None: + # Propagate annotation to quantizable patterns. + if ( + (node.target in propagation_quantizable_ops) + and (not _is_any_annotated([node])) + and (node.op == "call_function") + ): + + def is_all_inputs_connected_to_quantized_op(input_nodes): + # Ensure all the inputs connect to fusion pattern or quantized node + for input_node in input_nodes: + if not _is_quantized_op_pt2e(input_node): + return False + return True + + if _skip_annotate([node], filter_fn): + return + + if quantization_config is None: + _annotate_nodes_not_quantize(node) + return + + if node.target is torch.ops.aten.max_pool2d.default: + # Recipe of maxpool2d: check input arg[0] of maxpool2d is quantized or not + input_nodes_to_check = [node.all_input_nodes[0]] + if not is_all_inputs_connected_to_quantized_op(input_nodes_to_check): + if quantization_config is not None: + warnings.warn( + f"The input of maxpool2d is not quantized, skip annotate maxpool2d with config {quantization_config}." + ) + return + + self._annotate_maxpool2d(node, quantization_config) + return + elif node.target is torch.ops.aten.cat.default: + input_nodes_to_check = node.all_input_nodes + if not is_all_inputs_connected_to_quantized_op(input_nodes_to_check): + return + self._annotate_cat(node, quantization_config) + elif ( + node.target is torch.ops.aten.flatten.using_ints + and len(node.users) > 0 + and not any( + user.target in quantizable_ops for user in node.users.keys() + ) + ): + # Recipe of flatten: check if any users of flatten node are quantizable ops or not + return + else: + input_node = node.all_input_nodes[0] + if not is_all_inputs_connected_to_quantized_op( + [ + input_node, + ] + ): + return + input_qspec_map = {} + input_qspec_map[input_node] = get_input_act_qspec(quantization_config) + node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + return + + def _annotate_output_share_observer_as_input( + self, input_node: Node, source_node: Node + ): + source_node_quantization_annotation = ( + source_node.meta[QUANT_ANNOTATION_KEY] + if QUANT_ANNOTATION_KEY in source_node.meta + else None + ) + if ( + source_node_quantization_annotation + and source_node_quantization_annotation._is_output_of_quantized_pattern + ): + edge_or_node = (input_node, source_node) + source_node_quantization_annotation.output_qspec = SharedQuantizationSpec( + edge_or_node + ) + return + + def _annotate_output_for_int8_in_int8_out_pattern_entry( + self, + model: torch.fx.GraphModule, + ): + for node in model.graph.nodes: + self._annotate_output_for_int8_in_int8_out_pattern(node) + + def _annotate_output_for_int8_in_int8_out_pattern( + self, + node: Node, + ) -> None: + r""" + Check and insert observer at output of node in int8_in_int8_out_ops if needed. + Recipe refers to https://github.com/intel/intel-extension-for-pytorch/blob/ + 90d19323d96afc53fcc22ba5a7bb3fb07fdd6c1c/intel_extension_for_pytorch/quantization/_utils.py#L495 + """ + edge_or_node: tuple[Node, Node] + if (node.target in int8_in_int8_out_ops) and (_is_any_annotated([node])): + if node.target == torch.ops.aten.max_pool2d.default: + maxpool_node = node + if not _is_all_annotated( + [ + maxpool_node, + ] + ): + return + + # Get the quantization_annotation from getitem_node + maxpool_node_quantization_annotation = ( + maxpool_node.meta[QUANT_ANNOTATION_KEY] + if QUANT_ANNOTATION_KEY in maxpool_node.meta + else None + ) + if ( + maxpool_node_quantization_annotation + and maxpool_node_quantization_annotation._is_output_of_quantized_pattern + ): + # Annotate the output_qspec of getitem_node + input_act = maxpool_node.args[0] + assert isinstance(input_act, Node) + assert isinstance(maxpool_node, Node) + edge_or_node = (input_act, maxpool_node) + maxpool_node_quantization_annotation.output_qspec = ( + SharedQuantizationSpec(edge_or_node) + ) + else: + input_node = node.all_input_nodes[0] + self._annotate_output_share_observer_as_input(input_node, node) + return + + def _annotate_linear( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + linear_partitions = get_source_partitions( + gm.graph, [torch.nn.Linear, torch.nn.functional.linear] + ) + linear_partitions = list( + itertools.chain.from_iterable(linear_partitions.values()) + ) + for partition in linear_partitions: + if len(partition.output_nodes) > 1: + raise ValueError( + "Linear partition cannot have more than one output node" + ) + linear_node = partition.output_nodes[0] + if linear_node.op != "call_function" or linear_node.target not in ( + torch.ops.aten.linear.default, + ): + raise ValueError(f"{linear_node} is not an aten linear operator") + # skip annotation if it is already annotated + if _skip_annotate([linear_node], filter_fn): + continue + self._annotate_linear_node_helper(linear_node, True, quantization_config) + + def _annotate_linear_unary( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + postop_list = [ + torch.nn.ReLU, + torch.nn.LeakyReLU, + torch.nn.Tanh, + torch.nn.GELU, + ] + fused_partitions: list[tuple] = [] + for postop in postop_list: + fused_partitions = fused_partitions + find_sequential_partitions( + gm, [torch.nn.Linear, postop] + ) + for fused_partition in fused_partitions: + linear_partition, unary_partition = fused_partition + linear_node, unary_node = self._get_output_nodes_of_partitions( + [linear_partition, unary_partition] + ) + if linear_node.op != "call_function" or linear_node.target not in ( + torch.ops.aten.linear.default, + ): + continue + if _skip_annotate([unary_node, linear_node], filter_fn): + continue + + if quantization_config is None: + _annotate_nodes_not_quantize([linear_node, unary_node]) + continue + + self._annotate_linear_node_helper(linear_node, False, quantization_config) + unary_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + + def _annotate_linear_binary_unary( + self, + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ) -> None: + # linear + binary_op + (optional) unary op + binary_op_list = [operator.add] + unary_op_list = [torch.nn.ReLU, None] + combinations = itertools.product(binary_op_list, unary_op_list) + for binary_op, unary_op in combinations: + has_unary = unary_op is not None + seq_partition = [torch.nn.Linear, binary_op] + if has_unary: + seq_partition.append(unary_op) + fused_partitions = find_sequential_partitions(gm, seq_partition) + for fused_partition in fused_partitions: + unary_partition, unary_node = None, None + if has_unary: + ( + linear_partition, + binary_partition, + unary_partition, + ) = fused_partition + ( + linear_node, + binary_node, + unary_node, + ) = self._get_output_nodes_of_partitions( + [linear_partition, binary_partition, unary_partition] + ) + else: + linear_partition, binary_partition = fused_partition + linear_node, binary_node = self._get_output_nodes_of_partitions( + [linear_partition, binary_partition] + ) + if len(linear_node.users) != 1: + # Linear Node should only has 1 user node + continue + ( + linear_node_idx, + extra_input_node_idx, + ) = self._get_input_idx_for_binary_node(linear_node, binary_node) + if (linear_node_idx is None) or (extra_input_node_idx is None): + continue + if linear_node != binary_node.args[linear_node_idx]: + raise ValueError( + f"{linear_node} doesn't match input of binary node" + ) + assert isinstance(linear_node, Node) + if ( + linear_node.op != "call_function" + or linear_node.target != torch.ops.aten.linear.default + ): + # No linear node found to be fused with add + continue + node_list = ( + [binary_node, linear_node] + if unary_node is None + else [unary_node, binary_node, linear_node] + ) + if _skip_annotate(node_list, filter_fn): + continue + + if quantization_config is None: + _annotate_nodes_not_quantize(node_list) + continue + + self._annotate_linear_node_helper( + linear_node, False, quantization_config + ) + # We don't insert q-dq before the binary input node due to accuracy issues + binary_node.meta[QUANT_ANNOTATION_KEY] = ( + _X86InductorQuantizationAnnotation( + input_qspec_map={}, + _annotated=True, + _is_output_of_quantized_pattern=(not has_unary), + ) + ) + if unary_node is not None: + unary_node.meta[QUANT_ANNOTATION_KEY] = ( + _X86InductorQuantizationAnnotation( + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + ) + + def _annotate_mul_tensor( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ): + def _is_tensor(n: Node): + return isinstance(n, Node) and isinstance( + n.meta["val"], torch._subclasses.fake_tensor.FakeTensor + ) + + def _same_shape(n1: Node, n2: Node): + return n1.meta["val"].shape == n2.meta["val"].shape + + for node in model.graph.nodes: + if node.target != torch.ops.aten.mul.Tensor: + continue + + if _skip_annotate([node], filter_fn): + continue + + if quantization_config is None: + _annotate_nodes_not_quantize(node) + continue + + assert len(node.args) == 2 + if not (_is_tensor(node.args[0]) and _is_tensor(node.args[1])): + continue + + if not _same_shape(node.args[0], node.args[1]): + continue + + input_qspec_map = {} + mul_node = node + for input_node in mul_node.args: + input_qspec_map[input_node] = get_input_act_qspec(quantization_config) + mul_node.meta[QUANT_ANNOTATION_KEY] = _X86InductorQuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + _is_output_of_quantized_pattern=True, + ) + + def validate(self, model: torch.fx.GraphModule) -> None: + pass diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/xpu_inductor_quantizer.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/xpu_inductor_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..d43d2bd8566b80a9dee3d134a2eb39b32cea8751 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/quantizer/xpu_inductor_quantizer.py @@ -0,0 +1,131 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import functools +from typing import TYPE_CHECKING, Any, Optional + +import torch +from torch.fx import Node + +from torchao.quantization.pt2e.observer import ( + HistogramObserver, + PerChannelMinMaxObserver, +) +from torchao.quantization.pt2e.quantizer import ( + QuantizationConfig, + QuantizationSpec, +) +from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import ( + FilterFn, + X86InductorQuantizer, + _is_any_annotated, + int8_in_int8_out_ops, +) + +if TYPE_CHECKING: + from torchao.quantization.pt2e import _ObserverOrFakeQuantizeConstructor + +__all__ = [ + "XPUInductorQuantizer", + "get_default_xpu_inductor_quantization_config", +] + + +@functools.lru_cache +def get_default_xpu_inductor_quantization_config(): + extra_args: dict[str, Any] = {"eps": 2**-12} + act_observer_or_fake_quant_ctr = HistogramObserver + act_quantization_spec = QuantizationSpec( + dtype=torch.int8, + quant_min=-128, + quant_max=127, + qscheme=torch.per_tensor_affine, + is_dynamic=False, + observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args( + **extra_args + ), + ) + + weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = ( + PerChannelMinMaxObserver + ) + + weight_quantization_spec = QuantizationSpec( + dtype=torch.int8, + quant_min=-128, + quant_max=127, + qscheme=torch.per_channel_symmetric, + ch_axis=0, # 0 corresponding to weight shape = (oc, ic, kh, kw) of conv + is_dynamic=False, + observer_or_fake_quant_ctr=weight_observer_or_fake_quant_ctr.with_args( + **extra_args + ), + ) + + bias_quantization_spec = None # will use placeholder observer by default + quantization_config = QuantizationConfig( + act_quantization_spec, + act_quantization_spec, + weight_quantization_spec, + bias_quantization_spec, + False, + ) + return quantization_config + + +class XPUInductorQuantizer(X86InductorQuantizer): + """ + XPUInductorQuantizer is a class designed to facilitate + quantization capability at Intel GPU backend. The class + highly reuses the existing implementation of + X86InductorQuantizer as both are intended to take advantage + of the optimized kernels in oneDNN library. + """ + + def __init__(self) -> None: + super().__init__() + + """ + Following annotate_xx overrides the impls in base class, as + no XPU implementation for these operators currently. We would + gradually enable the XPU implementation and remove following + overrides. We keep the annotate methods but make the function + body empty, aiming to let `_generate_qdq_quantized_model` + generate qdq around op and graph execute on fp32 dtype for + unspported operators. + """ + + def _annotate_qat_conv2d_fusion_pattern( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[FilterFn] = None, + ): + pass + + def _annotate_maxpool2d( + self, + node: Node, + quantization_config: Optional[QuantizationConfig], + ) -> None: + """ + Here we skip the annotate logic for maxpool at XPU backend + as the quantized::max_pool2d is only implemented for CPU. + """ + return + + def _annotate_output_for_int8_in_int8_out_pattern( + self, + node: Node, + ) -> None: + if (node.target in int8_in_int8_out_ops) and (_is_any_annotated([node])): + if node.target == torch.ops.aten.max_pool2d.default: + return + else: + input_node = node.all_input_nodes[0] + self._annotate_output_share_observer_as_input(input_node, node) + return diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/reference_representation_rewrite.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/reference_representation_rewrite.py new file mode 100644 index 0000000000000000000000000000000000000000..6526c6044fd4e3bcbee2ca048641e1bc5263d2fb --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/reference_representation_rewrite.py @@ -0,0 +1,835 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import contextlib +from dataclasses import dataclass +from functools import partial +from typing import Any, Callable, Optional + +import torch +from torch._higher_order_ops.out_dtype import out_dtype +from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib # noqa: F401 +from torch.fx import GraphModule +from torch.fx.subgraph_rewriter import replace_pattern + +from torchao.quantization.pt2e.export_utils import WrapperModule +from torchao.quantization.pt2e.utils import ( + _get_aten_graph_module_for_pattern, + _replace_literals_with_existing_placeholders, + _replace_literals_with_new_placeholders, + remove_tensor_overload_for_qdq_ops, +) + +try: + from torch._export.utils import _disable_aten_to_metadata_assertions +except: + _disable_aten_to_metadata_assertions = contextlib.nullcontext + + +__all__ = [ + "reference_representation_rewrite", +] + + +def _qdq_quantized_linear( + x_i8, + x_scale, + x_zero_point, + x_quant_min, + x_quant_max, + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + bias_fp32, + out_scale, + out_zero_point, + out_quant_min, + out_quant_max, +): + x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + x_i8, x_scale, x_zero_point, x_quant_min, x_quant_max, torch.int8 + ) + weight_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + torch.int8, + ) + out_fp32 = torch.ops.aten.linear.default(x_fp32, weight_fp32, bias_fp32) + out_i8 = torch.ops.quantized_decomposed.quantize_per_tensor( + out_fp32, out_scale, out_zero_point, out_quant_min, out_quant_max, torch.int8 + ) + return out_i8 + + +def _reference_quantized_linear( + x_i8, + x_scale, + x_zero_point, + x_quant_min, + x_quant_max, + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + bias_fp32, + out_scale, + out_zero_point, + out_quant_min, + out_quant_max, +): + # without using quant_min/max in clamp, the traced graph will not have quant_mi/max args. + # This results in failure to match the pattern. + # Therefore, we call a torch.ops.aten.clamp here + x_i8 = torch.ops.aten.clamp(x_i8, x_quant_min, x_quant_max) + weight_i8 = torch.ops.aten.clamp(weight_i8, weight_quant_min, weight_quant_max) + + x_i16 = x_i8.to(torch.int16) + weight_i16 = weight_i8.to(torch.int16) + # always set bias to None so that the same representation can work for the case + # no matter if bias_scale == x_scale * weight_scale or not + acc_i32 = out_dtype( + torch.ops.aten.linear.default, + torch.int32, + x_i16 - x_zero_point, + weight_i16 - weight_zero_point, + None, + ) + # TODO: change to mul.Scalar + # Note: we are quantizing bias with these scales without signal from user, but it might be OK + bias_scale = x_scale * weight_scale + bias_i32 = out_dtype(torch.ops.aten.div.Tensor, torch.int32, bias_fp32, bias_scale) + acc_i32 = acc_i32 + bias_i32 + # TODO: change to mul.Scalar when we make x_scale/weight_scale etc. Scalar values + acc_i32 = ( + out_dtype( + torch.ops.aten.mul.Tensor, + torch.int32, + acc_i32, + x_scale * weight_scale / out_scale, + ) + + out_zero_point + ) + out_i8 = torch.ops.aten.clamp(acc_i32, out_quant_min, out_quant_max).to(torch.int8) + return out_i8 + + +def _qdq_dynamic_quantized_linear( + x_fp32, + x_quant_min, + x_quant_max, + x_eps, + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + bias_fp32, +): + x_scale, x_zero_point = torch.ops.quantized_decomposed.choose_qparams( + x_fp32, x_quant_min, x_quant_max, x_eps, torch.int8 + ) + x_i8 = torch.ops.quantized_decomposed.quantize_per_tensor( + x_fp32, x_scale, x_zero_point, x_quant_min, x_quant_max, torch.int8 + ) + x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + x_i8, x_scale, x_zero_point, x_quant_min, x_quant_max, torch.int8 + ) + weight_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + torch.int8, + ) + out_fp32 = torch.ops.aten.linear.default(x_fp32, weight_fp32, bias_fp32) + return out_fp32 + + +def _reference_dynamic_quantized_linear( + x_fp32, + x_quant_min, + x_quant_max, + x_eps, + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + bias_fp32, +): + x_scale, x_zero_point = torch.ops.quantized_decomposed.choose_qparams( + x_fp32, x_quant_min, x_quant_max, x_eps, torch.int8 + ) + # decomposed representation for quantize_per_tensor + # TODO: use out_dtype(mul, ...) here when the op is ready + x_fp32 = x_fp32 / x_scale # fp32 + # round modes might be different here + # pytorch is rounding to even, which is also common for most of the backends + x_fp32 = torch.round(x_fp32) # fp32 + x_i32 = x_fp32.to(dtype=torch.int32) # int32 + x_i32 = x_i32 + x_zero_point # int32 + # clamp works for fp32, int32 and int8 dtypes + x_i32 = torch.clamp(x_i32, x_quant_min, x_quant_max) # int32 + x_i8 = x_i32.to(dtype=torch.int8) + + weight_i8 = torch.ops.aten.clamp(weight_i8, weight_quant_min, weight_quant_max) + + x_i16 = x_i8.to(torch.int16) + weight_i16 = weight_i8.to(torch.int16) + # always set bias to None so that the same representation can work for the case + # no matter if bias_scale == x_scale * weight_scale or not + acc_i32 = out_dtype( + torch.ops.aten.linear.default, + torch.int32, + x_i16 - x_zero_point, + weight_i16 - weight_zero_point, + None, + ) + bias_scale = x_scale * weight_scale + bias_i32 = out_dtype(torch.ops.aten.div.Tensor, torch.int32, bias_fp32, bias_scale) + acc_i32 = acc_i32 + bias_i32 + out_fp32 = acc_i32 * (x_scale * weight_scale) + return out_fp32 + + +def _qdq_quantized_conv2d( + x_i8, + x_scale, + x_zero_point, + x_quant_min, + x_quant_max, + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + bias_fp32, + out_scale, + out_zero_point, + out_quant_min, + out_quant_max, +): + stride = [1, 1] + padding = [0, 0] + dilation = [1, 1] + transposed = False + output_padding = [0, 0] + groups = 1 + x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + x_i8, x_scale, x_zero_point, x_quant_min, x_quant_max, torch.int8 + ) + weight_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + torch.int8, + ) + out_fp32 = torch.ops.aten.convolution.default( + x_fp32, + weight_fp32, + bias_fp32, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + out_i8 = torch.ops.quantized_decomposed.quantize_per_tensor( + out_fp32, out_scale, out_zero_point, out_quant_min, out_quant_max, torch.int8 + ) + return out_i8 + + +def _reference_quantized_conv2d( + x_i8, + x_scale, + x_zero_point, + x_quant_min, + x_quant_max, + weight_i8, + weight_scale, + weight_zero_point, + weight_quant_min, + weight_quant_max, + bias_fp32, + out_scale, + out_zero_point, + out_quant_min, + out_quant_max, +): + stride = [1, 1] + padding = [0, 0] + dilation = [1, 1] + transposed = False + output_padding = [0, 0] + groups = 1 + # without using quant_min/max in clamp, the traced graph will not have quant_mi/max args. + # This results in failure to match the pattern. + # Therefore, we call a torch.ops.aten.clamp here + x_i8 = torch.ops.aten.clamp(x_i8, x_quant_min, x_quant_max) + weight_i8 = torch.ops.aten.clamp(weight_i8, weight_quant_min, weight_quant_max) + + x_i16 = x_i8.to(torch.int16) + weight_i16 = weight_i8.to(torch.int16) + # always set bias to None so that the same representation can work for the case + # no matter if bias_scale == x_scale * weight_scale or not + acc_i32 = out_dtype( + torch.ops.aten.convolution.default, + torch.int32, + x_i16 - x_zero_point, + weight_i16 - weight_zero_point, + None, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + # Note: we are quantizing bias with these scales without signal from user, but it might be OK + bias_scale = x_scale * weight_scale + # bias quantization to int32 uses bias_scale = x_scale * weight_scale due to: + # Take linear calculation for example + # Out_(i, j)_fp32 = Sum_(over k)[X_(i, k)_fp32 * W_(i, k)_fp32] + bias_(i)_fp32 + # Represent X, W fp32 as their dequant transforms + # A_fp32 = (A_q - A_zero_point)/A_scale + # Out_(i, j)_fp32 = Sum_(over k)[(X_(i, k)_fp32 - X_zp) * X_scale * (W_(i, k)_fp32 - W_zp) * W_scale] + bias_(i)_fp32 + # Factor out X_scale and W_scale + # Out_(i, j)_fp32 = ((X_scale * W_scale) * Sum_(over k)[(X_(i, k)_fp32 - X_zp) * (W_(i, k)_fp32 - W_zp)]) + bias_(i)_fp32 + # In order to addition of bias_(i)_fp32 inside, we must do + # Out_(i, j)_fp32 = (X_scale * W_scale) * (Sum_(over k)[(X_(i, k)_fp32 - X_zp) * (W_(i, k)_fp32 - W_zp)] + (1 / (X_scale * W_scale)) * bias_(i)_fp32)W_scale # noqa: B950 + # Note we had to multiply bias_fp32 qith X_scale * W_scale = bias_scale + # Thus bias quantization to int32 must be with X_scale * W_scale + + bias_i32 = out_dtype(torch.ops.aten.div.Tensor, torch.int32, bias_fp32, bias_scale) + # Unsqueeze to match broadcast dims + # Unfortnuately I cannot do bias_i32.unsqueeze(0) due to literal matching nightmare + # in graph pattern replacement + bias_i32 = bias_i32.unsqueeze(-1) + bias_i32 = bias_i32.unsqueeze(-1) + acc_i32 = acc_i32 + bias_i32 + # TODO: change to mul.Scalar when we make x_scale/weight_scale etc. Scalar values + acc_i32 = ( + out_dtype( + torch.ops.aten.mul.Tensor, + torch.int32, + acc_i32, + x_scale * weight_scale / out_scale, + ) + + out_zero_point + ) + out_i8 = torch.ops.aten.clamp(acc_i32, out_quant_min, out_quant_max).to(torch.int8) + return out_i8 + + +def _qdq_quantized_add_relu( + x_i8, + x_scale, + x_zero_point, + y_i8, + y_scale, + y_zero_point, + out_scale, + out_zero_point, + quant_min, + quant_max, +): + x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + x_i8, x_scale, x_zero_point, quant_min, quant_max, torch.int8 + ) + y_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + y_i8, y_scale, y_zero_point, quant_min, quant_max, torch.int8 + ) + out_fp32 = x_fp32 + y_fp32 + out_fp32 = torch.ops.aten.relu(out_fp32) + out_i8 = torch.ops.quantized_decomposed.quantize_per_tensor( + out_fp32, out_scale, out_zero_point, quant_min, quant_max, torch.int8 + ) + return out_i8 + + +def _reference_quantized_add_relu( + x_i8, + x_scale, + x_zero_point, + y_i8, + y_scale, + y_zero_point, + out_scale, + out_zero_point, + quant_min, + quant_max, +): + """ + See comments for `_reference_quantized_add` for more information on + how to derive the formula for out_i8 based on x_i8 and y_i8 + """ + x_i32 = x_i8.to(torch.int32) + y_i32 = y_i8.to(torch.int32) + # TODO: change this to mul.Scalar? + x_i32 = out_dtype( + torch.ops.aten.mul.Tensor, + torch.int32, + (x_i32 - x_zero_point), + (x_scale / out_scale), + ) + y_i32 = out_dtype( + torch.ops.aten.mul.Tensor, + torch.int32, + (y_i32 - y_zero_point), + (y_scale / out_scale), + ) + out_i32 = x_i32 + y_i32 + out_zero_point + # out_i32 = torch.ops.aten.clamp(out_i32, out_zero_point) + out_i8 = torch.ops.aten.clamp(out_i32, out_zero_point, quant_max).to(torch.int8) + return out_i8 + + +def _qdq_quantized_add( + x_i8, + x_scale, + x_zero_point, + y_i8, + y_scale, + y_zero_point, + out_scale, + out_zero_point, + quant_min, + quant_max, +): + x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + x_i8, x_scale, x_zero_point, quant_min, quant_max, torch.int8 + ) + y_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + y_i8, y_scale, y_zero_point, quant_min, quant_max, torch.int8 + ) + out_fp32 = x_fp32 + y_fp32 + out_i8 = torch.ops.quantized_decomposed.quantize_per_tensor( + out_fp32, out_scale, out_zero_point, quant_min, quant_max, torch.int8 + ) + return out_i8 + + +def _reference_quantized_add( + x_i8, + x_scale, + x_zero_point, + y_i8, + y_scale, + y_zero_point, + out_scale, + out_zero_point, + quant_min, + quant_max, +): + """ + # How to Derive the formula for out_i8 based on x_i8 and y_i8 + # (since quantized add takes x_i8, y_i8 and their quantization parameters, and produce an out_i8) + + # out_i8 is quantized output, we can write down the formula for it first: + out_i8 = out_f32 / out_scale + out_zero_point (1) + + # then out_fp32 is computed from x_f32 + y_f32, and the x_fp32 and y_fp32 are the dequantized x_i8 and y_i8 + out_f32 = x_f32 + y_f32 (2) + x_fp32 = (x_i8 - x_zero_point) * x_scale (3) + y_fp32 = (y_i8 - y_zero_point) * y_scale (4) + + # applying the above fomula to the out_i8 equation we can get the following: + out_i8 = out_fp32 / out_scale + out_zero_point # (1) + = (x_f32 + y_f32) / out_scale + out_zero_point # applying (2) to substitute out_fp32 with x_fp32 + y_fp32 + = ((x_i8 - x_zero_point) * x_scale + (y_i8 - y_zero_point) * y_scale) / out_scale + out_zero_point # apply (3) and (4) + """ + x_i32 = x_i8.to(torch.int32) + y_i32 = y_i8.to(torch.int32) + # TODO: use out_dtype op + x_i32 = torch.round((x_scale / out_scale) * (x_i32 - x_zero_point)).to(torch.int32) + y_i32 = torch.round((y_scale / out_scale) * (y_i32 - y_zero_point)).to(torch.int32) + out_i32 = x_i32 + y_i32 + out_zero_point + quant_min = -128 + quant_max = 127 + out_i8 = torch.ops.aten.clamp(out_i32, quant_min, quant_max).to(torch.int8) + return out_i8 + + +def _qdq_quantized_max_pool2d( + x_i8, + x_scale, + x_zero_point, + x_quant_min, + x_quant_max, + out_scale, + out_zero_point, + out_quant_min, + out_quant_max, +): + kernel_size = 1 + stride = 1 + padding = 0 + dilation = 1 + ceil_mode = False + x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + x_i8, x_scale, x_zero_point, x_quant_min, x_quant_max, torch.int8 + ) + out_fp32, _ = torch.ops.aten.max_pool2d_with_indices.default( + x_fp32, kernel_size, stride, padding, dilation, ceil_mode + ) + out_i8 = torch.ops.quantized_decomposed.quantize_per_tensor( + out_fp32, out_scale, out_zero_point, out_quant_min, out_quant_max, torch.int8 + ) + return out_i8 + + +def _reference_quantized_max_pool2d( + x_i8, + x_scale, + x_zero_point, + x_quant_min, + x_quant_max, + out_scale, + out_zero_point, + out_quant_min, + out_quant_max, +): + kernel_size = 1 + stride = 1 + padding = 0 + dilation = 1 + ceil_mode = False + # to preserve x_quant_min, x_quant_max in the graph for pattern matching + x_i8 = torch.clamp(x_i8, x_quant_min, x_quant_max) + x_i32 = x_i8.to(torch.int32) + out_i32, _ = torch.ops.aten.max_pool2d_with_indices.default( + x_i32 - x_zero_point, kernel_size, stride, padding, dilation, ceil_mode + ) + out_fp32 = out_i32 * (x_scale / out_scale) + out_zero_point + out_fp32 = torch.clamp(out_fp32, out_quant_min, out_quant_max) + out_i8 = out_fp32.to(torch.int8) + return out_i8 + + +def _quantize_per_tensor_int8(x_fp32, scale, zero_point, quant_min, quant_max): + x = torch.ops.quantized_decomposed.quantize_per_tensor( + x_fp32, scale, zero_point, quant_min, quant_max, torch.int8 + ) + return x + + +def _reference_quantize_per_tensor_int8( + x_fp32, scale, zero_point, quant_min, quant_max +): + # TODO: use out_dtype(mul, ...) here when the op is ready + x = x_fp32 / scale # fp32 + # round modes might be different here + # pytorch is rounding to even, which is also common for most of the backends + x = torch.round(x) # fp32 + x = x.to(dtype=torch.int32) # int32 + x = x + zero_point # int32 + # clamp works for fp32, int32 and int8 dtypes + x = torch.clamp(x, quant_min, quant_max) # int32 + x = x.to(dtype=torch.int8) + return x + + +def _dequantize_per_tensor_int8(x_i8, scale, zero_point, quant_min, quant_max): + x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor( + x_i8, scale, zero_point, quant_min, quant_max, torch.int8 + ) + return x_fp32 + + +def _reference_dequantize_per_tensor_int8( + x_i8, scale, zero_point, quant_min, quant_max +): + # without using quant_min/max in clamp, the traced graph will not have quant_mi/max args. + # This results in failure to match the pattern. + # Therefore, we call a torch.ops.aten.clamp here + x_i8 = torch.ops.aten.clamp(x_i8, quant_min, quant_max) + # TODO: use out_dtype op + # note: x_i8.to(torch.int32) does not work here + # TODO: debug the implementation later when torchdynamo time out issue is resolved + return ((x_i8.to(torch.float32) - zero_point) * scale).to(dtype=torch.float32) + + +def _quantize_per_channel_int8( + x_fp32, scales, zero_points, ch_axis, quant_min, quant_max +): + out_i8 = torch.ops.quantized_decomposed.quantize_per_channel( + x_fp32, scales, zero_points, ch_axis, quant_min, quant_max, torch.int8 + ) + return out_i8 + + +def _reference_quantize_per_channel_int8( + x_fp32, scales, zero_points, ch_axis, quant_min, quant_max +): + x_fp32 = torch.transpose(x_fp32, ch_axis, -1) + out_i32 = torch.ops.aten.clamp( + torch.round(x_fp32 / scales).to(torch.int32) + zero_points, quant_min, quant_max + ) + out_i32 = torch.transpose(out_i32, ch_axis, -1) + return out_i32.to(torch.int8) + + +def _dequantize_per_channel_int8( + x_i8, scales, zero_points, ch_axis, quant_min, quant_max +): + # the following will be replaced as placeholders + out_fp32 = torch.ops.quantized_decomposed.dequantize_per_channel( + x_i8, scales, zero_points, ch_axis, quant_min, quant_max, torch.int8 + ) + return out_fp32 + + +def _reference_dequantize_per_channel_int8( + x_i8, scales, zero_points, ch_axis, quant_min, quant_max +): + # the following will be replaced as placeholders + # in order to preserve the quant_min/quant_max args for pattern matching (e.g. matching for int4 quantized ops) + # we call a torch.ops.aten.clamp here + x_i8 = torch.ops.aten.clamp(x_i8, quant_min, quant_max) + x_i8 = torch.transpose(x_i8, ch_axis, -1) + x_i32 = x_i8.to(torch.int32) + out_fp32 = (x_i32 - zero_points).to(torch.float) * scales + out_fp32 = torch.transpose(out_fp32, ch_axis, -1) + return out_fp32 + + +def _replace_ph_qdq_per_channel_replacement(gm: torch.fx.GraphModule): + return _replace_literals_with_existing_placeholders( + gm, exclude_literals=[-1], literal_to_ph_idx={1: 3, -128: 4, 127: 5} + ) + + +@dataclass +class _RewriteInfo: + """Data needed for rewrite, this includes example inputs, pattern and replacement functions + and post transformation functions for the exported pattern and replacement GraphModule + """ + + # example inputs used for exporting the pattern into GraphModule + example_inputs: tuple[Any, ...] + pattern: Callable + replacement: Callable + # post transformation on the exported pattern and replacement GraphModule + pattern_post_trans: Optional[Callable[[GraphModule], GraphModule]] = None + replacement_post_trans: Optional[Callable[[GraphModule], GraphModule]] = None + + +def reference_representation_rewrite(model: GraphModule) -> GraphModule: + _QUANTIZED_LINEAR_EXAMPLE_INPUTS = ( + torch.randint(-128, 127, (2, 5), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + torch.randint(-128, 127, (5, 5), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-127], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + torch.randn(1, dtype=torch.float), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + ) + + _DYNAMIC_QUANTIZED_LINEAR_EXAMPLE_INPUTS = ( + torch.randn((2, 5), dtype=torch.float), + -128, + 127, + torch.finfo(torch.float32).eps, + torch.randint(-128, 127, (5, 5), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-127], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + torch.randn(1, dtype=torch.float), + ) + + _QUANTIZED_CONV2d_EXAMPLE_INPUTS = ( + torch.randint(-128, 127, (1, 3, 3, 3), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + torch.randint(-128, 127, (1, 3, 3, 3), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-127], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + torch.randn(1, dtype=torch.float), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + ) + + _QUANTIZED_ADD_OR_ADD_RELU_EXAMPLE_INPUTS = ( + torch.randint(-128, 127, (1, 3, 3, 3), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.randint(-128, 127, (1, 3, 3, 3), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + ) + + _QUANTIZED_MAX_POOL2D_EXAMPLE_INPUTS = ( + torch.randint(-128, 127, (1, 3, 3, 3), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + ) + + _QUANTIZE_PER_TENSOR_INT8_EXAMPLE_INPUTS = ( + torch.randn(1, 3, 3, 3, dtype=torch.float), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + ) + + _DEQUANTIZE_PER_TENSOR_INT8_EXAMPLE_INPUTS = ( + torch.randint(-128, 127, (1, 3, 3, 3), dtype=torch.int8), + torch.randn(1, dtype=torch.float), + torch.zeros(1, dtype=torch.int), + torch.tensor([-128], dtype=torch.int), + torch.tensor([127], dtype=torch.int), + ) + + _QUANTIZE_PER_CHANNEL_INT8_EXAMPLE_INPUTS = ( + torch.randn(1, 3, 3, 3, dtype=torch.float), + torch.randn(3, dtype=torch.float), + torch.zeros(3, dtype=torch.int), + 1, + -128, + 127, + ) + + _DEQUANTIZE_PER_CHANNEL_INT8_EXAMPLE_INPUTS = ( + torch.randint(-128, 127, (1, 3, 3, 3), dtype=torch.int8), + torch.randn(3, dtype=torch.float), + torch.zeros(3, dtype=torch.int), + 1, + -128, + 127, + ) + + _REWRITE_INFO_LIST = [ + _RewriteInfo( + _DYNAMIC_QUANTIZED_LINEAR_EXAMPLE_INPUTS, + WrapperModule(_qdq_dynamic_quantized_linear), + WrapperModule(_reference_dynamic_quantized_linear), + partial( + _replace_literals_with_existing_placeholders, + literal_to_ph_idx={-128: 1, 127: 2, torch.finfo(torch.float32).eps: 3}, + ), + partial( + _replace_literals_with_existing_placeholders, + literal_to_ph_idx={-128: 1, 127: 2, torch.finfo(torch.float32).eps: 3}, + ), + ), + _RewriteInfo( + _QUANTIZED_LINEAR_EXAMPLE_INPUTS, + WrapperModule(_qdq_quantized_linear), + WrapperModule(_reference_quantized_linear), + _replace_literals_with_new_placeholders, + _replace_literals_with_new_placeholders, + ), + _RewriteInfo( + _QUANTIZED_CONV2d_EXAMPLE_INPUTS, + WrapperModule(_qdq_quantized_conv2d), + WrapperModule(_reference_quantized_conv2d), + partial(_replace_literals_with_new_placeholders, exclude_literals=[-1]), + partial(_replace_literals_with_new_placeholders, exclude_literals=[-1]), + ), + _RewriteInfo( + _QUANTIZED_ADD_OR_ADD_RELU_EXAMPLE_INPUTS, + WrapperModule(_qdq_quantized_add_relu), + WrapperModule(_reference_quantized_add_relu), + ), + _RewriteInfo( + _QUANTIZED_ADD_OR_ADD_RELU_EXAMPLE_INPUTS, + WrapperModule(_qdq_quantized_add), + WrapperModule(_reference_quantized_add), + ), + _RewriteInfo( + _QUANTIZED_MAX_POOL2D_EXAMPLE_INPUTS, + WrapperModule(_qdq_quantized_max_pool2d), + WrapperModule(_reference_quantized_max_pool2d), + _replace_literals_with_new_placeholders, + _replace_literals_with_new_placeholders, + ), + _RewriteInfo( + _QUANTIZE_PER_TENSOR_INT8_EXAMPLE_INPUTS, + WrapperModule(_quantize_per_tensor_int8), + WrapperModule(_reference_quantize_per_tensor_int8), + ), + _RewriteInfo( + _DEQUANTIZE_PER_TENSOR_INT8_EXAMPLE_INPUTS, + WrapperModule(_dequantize_per_tensor_int8), + WrapperModule(_reference_dequantize_per_tensor_int8), + ), + _RewriteInfo( + _QUANTIZE_PER_CHANNEL_INT8_EXAMPLE_INPUTS, + WrapperModule(_quantize_per_channel_int8), + WrapperModule(_reference_quantize_per_channel_int8), + _replace_ph_qdq_per_channel_replacement, + _replace_ph_qdq_per_channel_replacement, + ), + _RewriteInfo( + _DEQUANTIZE_PER_CHANNEL_INT8_EXAMPLE_INPUTS, + WrapperModule(_dequantize_per_channel_int8), + WrapperModule(_reference_dequantize_per_channel_int8), + _replace_ph_qdq_per_channel_replacement, + _replace_ph_qdq_per_channel_replacement, + ), + ] + + remove_tensor_overload_for_qdq_ops(model) + + with _disable_aten_to_metadata_assertions(): + for rewrite_info in _REWRITE_INFO_LIST: + example_inputs = rewrite_info.example_inputs + pattern = rewrite_info.pattern + replacement = rewrite_info.replacement + pattern_post_trans = rewrite_info.pattern_post_trans + replacement_post_trans = rewrite_info.replacement_post_trans + pattern = _get_aten_graph_module_for_pattern(pattern, example_inputs) # type: ignore[arg-type, assignment] + remove_tensor_overload_for_qdq_ops(pattern) # type: ignore[arg-type] + replacement = _get_aten_graph_module_for_pattern( + replacement, example_inputs + ) # type: ignore[arg-type, assignment] + remove_tensor_overload_for_qdq_ops(replacement) # type: ignore[arg-type] + if pattern_post_trans: + pattern = pattern_post_trans(pattern) + if replacement_post_trans: + replacement = replacement_post_trans(replacement) + pattern.recompile() # type: ignore[attr-defined] + replacement.recompile() # type: ignore[attr-defined] + replace_pattern(model, pattern, replacement) + + return model diff --git a/lib/python3.12/site-packages/torchao/quantization/pt2e/utils.py b/lib/python3.12/site-packages/torchao/quantization/pt2e/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1cc5d4f4e96088efbdc36e9cb817bc737d84312c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/pt2e/utils.py @@ -0,0 +1,1084 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +""" +Utils shared by different modes of quantization (eager/graph) +""" + +import functools + +# mypy: allow-untyped-defs +import operator +import types +import warnings +from collections import OrderedDict +from inspect import getfullargspec, signature +from typing import Any, Callable, Optional, Union + +import torch +import torch.nn.functional as F + +# Makes sure that quantized_decomposed ops are registered +from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib # noqa: F401 +from torch.export.unflatten import _assign_attr, _AttrKind +from torch.fx import Graph, GraphModule, Node +from torch.nn.utils.fusion import fuse_conv_bn_weights +from torch.nn.utils.parametrize import is_parametrized +from torch.utils._pytree import LeafSpec + +from torchao.utils import _assert_and_get_unique_device + +__all__ = [ + "is_per_tensor", + "is_per_channel", + "getattr_from_fqn", + "get_qparam_dict", + "check_min_max_valid", + "calculate_qmin_qmax", + "has_no_children_ignoring_parametrizations", + "get_fqn_to_example_inputs", + "to_underlying_dtype", + "determine_qparams", + "validate_qmin_qmax", + "get_new_attr_name_with_prefix", + "create_getattr_from_value", + "_get_aten_graph_module_for_pattern", + "_is_conv_node", + "_is_conv_transpose_node", + "_is_sym_size_node", + "_filter_sym_size_users", +] + + +# TODO: remove unused +def is_per_tensor(qscheme): + return qscheme == torch.per_tensor_affine or qscheme == torch.per_tensor_symmetric + + +def is_per_channel(qscheme): + return qscheme in [ + torch.per_channel_affine, + torch.per_channel_affine_float_qparams, + torch.per_channel_symmetric, + ] + + +def getattr_from_fqn(obj: Any, fqn: str) -> Any: + """ + Given an obj and a fqn such as "foo.bar.baz", returns gm.foo.bar.baz. + """ + return functools.reduce(getattr, fqn.split("."), obj) + + +def to_underlying_dtype(qdtype): + DTYPE_MAPPING = { + torch.quint8: torch.uint8, + torch.qint8: torch.int8, + torch.qint32: torch.int32, + torch.quint4x2: torch.uint8, + torch.quint2x4: torch.uint8, + torch.uint8: torch.uint8, + torch.int8: torch.int8, + torch.uint16: torch.uint16, + torch.int16: torch.int16, + torch.int32: torch.int32, + torch.float8_e5m2: torch.float8_e5m2, + torch.float8_e4m3fn: torch.float8_e4m3fn, + } + assert qdtype in DTYPE_MAPPING, "Unsupported dtype: " + str(qdtype) + return DTYPE_MAPPING[qdtype] + + +def get_qparam_dict(observer_or_fake_quant): + from torchao.quantization.pt2e.observer import PlaceholderObserver + + qscheme = getattr(observer_or_fake_quant, "qscheme", None) + dtype = observer_or_fake_quant.dtype + qparams = {"qscheme": qscheme, "dtype": dtype} + + if not qscheme or isinstance(observer_or_fake_quant, PlaceholderObserver): + return {"qscheme": None, "dtype": dtype} + + if is_per_tensor(qscheme): + qscheme = torch.per_tensor_affine + elif is_per_channel(qscheme): + # change symmetric to affine since we do not have symmetric + # quantized Tensor + if qscheme == torch.per_channel_symmetric: + qscheme = torch.per_channel_affine + qparams["axis"] = observer_or_fake_quant.ch_axis + else: + raise RuntimeError(f"Unrecognized qscheme: {qscheme}") + # update qscheme, since we don't have symmetric quant qscheme + # in quantized Tensor + qparams["qscheme"] = qscheme + + scale, zero_point = observer_or_fake_quant.calculate_qparams() + qparams["scale"] = scale + qparams["zero_point"] = zero_point + + if hasattr(observer_or_fake_quant, "quant_min"): + qparams["quant_min"] = observer_or_fake_quant.quant_min + if hasattr(observer_or_fake_quant, "quant_max"): + qparams["quant_max"] = observer_or_fake_quant.quant_max + + return qparams + + +def check_min_max_valid(min_val: torch.Tensor, max_val: torch.Tensor) -> bool: + """Checks if the given minimum and maximum values are valid, meaning that + they exist and the min value is less than the max value. + """ + if min_val.numel() == 0 or max_val.numel() == 0: + warnings.warn( + "must run observer before calling calculate_qparams. " + + "Returning default values." + ) + return False + + if min_val.dim() == 0 or max_val.dim() == 0: + if min_val == float("inf") and max_val == float("-inf"): + warnings.warn( + "must run observer before calling calculate_qparams. " + + "Returning default values." + ) + + return False + + assert min_val <= max_val, f"min {min_val} should be less than max {max_val}" + else: + assert torch.all(min_val <= max_val), ( + f"min {min_val} should be less than max {max_val}" + ) + + return True + + +def calculate_qmin_qmax( + quant_min: int, + quant_max: int, + has_customized_qrange: bool, + dtype: torch.dtype, + reduce_range: bool, +) -> tuple[int, int]: + r"""Calculates actual qmin and qmax based on the quantization range, + observer datatype and if range is reduced. + """ + # TODO(jerryzh): Figure out why custom quant_min/quant_max are still adjusted. + if has_customized_qrange: + # This initialization here is to be resolve TorchScript compilation issues and allow + # using of refinement to decouple initial_qmin and initial_qmax from quantization range. + # The actual values of initial_qmin and initial_qmax will be reset below. + if dtype in [torch.qint32, torch.int32]: + initial_quant_min, initial_quant_max = 0, 2**32 - 1 + else: + initial_quant_min, initial_quant_max = 0, 255 + # The following assignment of self.qmin and self.qmax to the local variables and the if check refine the + # attribute from Optional valid integers for use, based on TorchScript's requirements. + custom_quant_min, custom_quant_max = quant_min, quant_max + if custom_quant_min is not None and custom_quant_max is not None: + initial_quant_min, initial_quant_max = ( + custom_quant_min, + custom_quant_max, + ) + + qrange_len = initial_quant_max - initial_quant_min + 1 + if dtype in [torch.qint8, torch.int8]: + assert 0 < qrange_len <= 256, ( + "quantization range should be positive and not exceed the maximum bit range (=256)." + ) + elif dtype in [torch.qint32, torch.int32]: + assert 0 < qrange_len <= 2**32, ( + "quantization range should be positive and not exceed the maximum bit range (=4294967296)." + ) + if reduce_range: + quant_min, quant_max = quant_min // 2, quant_max // 2 + else: + # Fallback onto default 8-bit qmin and qmax calculation if dynamic range is not used. + if dtype in [torch.qint8, torch.int8]: + if reduce_range: + quant_min, quant_max = -64, 63 + else: + quant_min, quant_max = -128, 127 + elif dtype in [torch.quint8, torch.uint8]: + if reduce_range: + quant_min, quant_max = 0, 127 + else: + quant_min, quant_max = 0, 255 + elif dtype in [torch.qint32, torch.int32]: + quant_min, quant_max = -1 * (2**31), (2**31) - 1 + elif dtype in [torch.uint16]: + quant_min, quant_max = 0, 2**16 - 1 + elif dtype in [torch.int16]: + quant_min, quant_max = -(2**15), 2**15 - 1 + else: + quant_min, quant_max = 0, 15 + return quant_min, quant_max + + +def _parent_name(target): + """ + Turn 'foo.bar' into ['foo', 'bar'] + """ + r = target.rsplit(".", 1) + if len(r) == 1: + return "", r[0] + else: + return r[0], r[1] + + +def has_no_children_ignoring_parametrizations(module): + """ + Checks if module._modules is empty or + if module is a parametrization, checks that module._modules only has + the 'parametrizations' module + """ + if len(module._modules) == 0: + return True + elif is_parametrized(module): + return len(module._modules) == 1 and "parametrizations" in module._modules + else: + return False + + +def _get_path_of_module( + root: torch.nn.Module, submodule: torch.nn.Module +) -> Optional[str]: + """Get the path (fully qualified name) of a submodule + + Example:: + + >> class M(torch.nn.Module): + def __init__(self) -> None: + self.linear = torch.nn.Linear(5, 5) + def forward(self, x): + return self.linear(x) + + >> m = M() + >> l = m.linear + >> _get_path_of_module(m, l) + "linear" + """ + for n, p in root.named_modules(): + if submodule is p: + return n + return None + + +def _get_signature_locals(f: Callable, loc: dict[str, Any]) -> dict[str, Any]: + """Get local keyword arguments + + Example:: + + >> def f(self, a, b=9): + pass + >> loc = {"a": 6, "c": 7} + >> _get_signature_locals(f, loc) + {"a": 6} + """ + return {k: v for k, v in loc.items() if k in signature(f).parameters} + + +def _get_default_kwargs(f: Callable) -> "OrderedDict[str, Any]": + """Get all default keyword arguments from function signature + + Example:: + + >> def f(self, a, b=9): + pass + >> _get_default_kwargs(f) + {"b": 9} + """ + kwargs = {} + for name, param in signature(f).parameters.items(): + if param.default is not param.empty: + kwargs[name] = param.default + elif param.kind is param.VAR_POSITIONAL: + kwargs[name] = () + elif param.kind is param.VAR_KEYWORD: + kwargs[name] = {} + return OrderedDict(kwargs) + + +def _normalize_kwargs(func: Callable, loc: dict[str, Any]) -> "OrderedDict[str, Any]": + """Given a function and local function arguments, normalize the keyword + arguments by filling in default arguments from function signature + + Example:: + + >> def f(self, key1=3, key2=3): + pass + >> loc = {"key2": 6} + >> _normalize_kwargs(f, loc) + {"key1": 3, "key2": 6} + """ + default_kwargs = _get_default_kwargs(func) + local_kwargs = _get_signature_locals(func, loc) + normalized_kwargs = default_kwargs.copy() + for attr, val in local_kwargs.items(): + if attr in normalized_kwargs: + # override the default keyword arguments + normalized_kwargs[attr] = val + return normalized_kwargs + + +def validate_qmin_qmax(quant_min: int, quant_max: int) -> None: + r"""Validates that the user-specified quantization range is properly initialized + and within the given bound supported by the observer dtype. + + To accommodate lower-bit quantization with respect to the existing torch.qint8 and + torch.quint8 datatypes, the user can choose to use dynamic quantization range by passing + in a tuple of initial qmin and qmax values. One use case is these customized qmin and qmax + values are used to calculate static estimates of the scale and zero point for aggressive lower-bit + fake quantization. These estimates are compared against parameters learned through backpropagation. + The related literatures for scale and zero point via backpropagation are as follows: + + Learned Step Size Quantization: https://openreview.net/pdf?id=rkgO66VKDS + Trained Quantization Thresholds: https://arxiv.org/pdf/1903.08066.pdf + """ + # The variable names are prefixed with "initial" because their values (qmin and qmax) might be adjusted + # based on whether quantization range is reduced and the datatype (signed/unsigned) used by the observer. + assert quant_min <= 0 <= quant_max, ( + "Used-specified quantization range must include 0." + ) + assert quant_min < quant_max, ( + "qmin must be strictly less than qmax for user-specified quantization range." + ) + + +# Functionally equivalent to '_calculate_qparams' in observer.py. Observers must be torchscriptable however and qscheme +# as far as I can tell is not allowed to passed as a parameter in torchscript functions. This makes refactoring observer +# to use this utility a massive pain and very gross. For now Im opting just to duplicate as this code seems unlikey to change +# (last update over 1 year ago) and when torchscript is fully deprecated we can refactor. TODO(jakeszwe, jerryzh168) +def determine_qparams( + min_val: torch.Tensor, + max_val: torch.Tensor, + quant_min: int, + quant_max: int, + dtype: torch.dtype, + eps: torch.Tensor, + has_customized_qrange: bool, + qscheme: torch.qscheme = torch.per_tensor_affine, +) -> tuple[torch.Tensor, torch.Tensor]: + r"""Calculates the quantization parameters, given min and max + value tensors. Works for both per tensor and per channel cases + + Args: + min_val: Minimum values per channel + max_val: Maximum values per channel + + Returns: + scales: Scales tensor of shape (#channels,) + zero_points: Zero points tensor of shape (#channels,) + """ + if not check_min_max_valid(min_val, max_val): + return torch.tensor([1.0], device=min_val.device.type), torch.tensor( + [0], device=min_val.device.type + ) + + min_val_neg = torch.min(min_val, torch.zeros_like(min_val)) + max_val_pos = torch.max(max_val, torch.zeros_like(max_val)) + + device = min_val_neg.device + scale = torch.ones(min_val_neg.size(), dtype=torch.double, device=device) + zero_point = torch.zeros(min_val_neg.size(), dtype=torch.int64, device=device) + eps = eps.to(device) + + if qscheme == torch.per_tensor_symmetric or qscheme == torch.per_channel_symmetric: + max_val_pos = torch.max(-min_val_neg, max_val_pos) + scale = max_val_pos / (float(quant_max - quant_min) / 2) + scale = torch.max(scale, eps) + if dtype in [torch.uint8, torch.quint8]: + if has_customized_qrange: + # When customized quantization range is used, down-rounded midpoint of the range is chosen. + zero_point = zero_point.new_full( + zero_point.size(), (quant_min + quant_max) // 2 + ) + else: + zero_point = zero_point.new_full(zero_point.size(), 128) + elif qscheme == torch.per_channel_affine_float_qparams: + scale = (max_val - min_val) / float(quant_max - quant_min) + scale = torch.where(scale > eps, scale, torch.ones_like(scale)) + # We use the quantize function + # xq = Round(Xf * inv_scale + zero_point), + # setting zero_point to (-1 * min *inv_scale) we get + # Xq = Round((Xf - min) * inv_scale) + zero_point = -1 * min_val / scale + else: + scale = (max_val_pos - min_val_neg) / float(quant_max - quant_min) + scale = torch.max(scale, eps) + zero_point = quant_min - torch.round(min_val_neg / scale).to(torch.int) + zero_point = torch.clamp(zero_point, quant_min, quant_max) + + # For scalar values, cast them to Tensors of size 1 to keep the shape + # consistent with default values in FakeQuantize. + if len(scale.shape) == 0: + # TODO: switch to scale.item() after adding JIT support + scale = torch.tensor([float(scale)], dtype=scale.dtype, device=device) + if len(zero_point.shape) == 0: + # TODO: switch to zero_point.item() after adding JIT support + zero_point = torch.tensor( + [int(zero_point)], dtype=zero_point.dtype, device=device + ) + if qscheme == torch.per_channel_affine_float_qparams: + zero_point = torch.tensor( + [float(zero_point)], dtype=zero_point.dtype, device=device + ) + + return scale.to(torch.double), zero_point.to(torch.int64) + + +def _get_num_pos_args(f: Callable) -> int: + """Get number of positional args for a function + + Example:: + + >> def f(self, key1=3, key2=3): + pass + >> _get_num_pos_args(f) + 3 + """ + return len(getfullargspec(f).args) + + +def get_fqn_to_example_inputs( + model: torch.nn.Module, example_inputs: tuple[Any, ...] +) -> dict[str, tuple[Any, ...]]: + """Given a model and its example inputs, return a dictionary from + fully qualified name of submodules to example_inputs for that submodule, + e.g. {"linear1": (tensor1,), "linear2": (tensor2,), "sub": (tensor3,), + "sub.linear1": (tensor4,), ...} + + Used to make quantizing submodules easier now that FX Graph Mode Quantization requires + example inputs. + + Also works for keyword arguments with default values, we would flatten keyword + arguments as positional arguments and fill in the missing keyword args with default + values, e.g. if we have a forward function: + def forward(self, x, key1=3, key2=3): + ... + + and we call it with self.submodule(x, key2=6) + we'll get example_inputs: (x, 3, 6) + + user can also override `key1` with positional arguments as well: + for self.submodule(x, 5, key2=6) + we'll get: (x, 5, 6) + + variable positional arguments and variable positional keyword arguments in forward + function are not supported currently, so please make sure no submodules is using + them. + """ + root = model + fqn_to_example_inputs = {} + + def _patched_module_call(self, *args, **kwargs): + submodule_example_inputs = list(args).copy() + normalized_kwargs = _normalize_kwargs(self.forward, kwargs) + # minus 1 to skipping counting `self` + num_args = _get_num_pos_args(self.forward) - 1 + num_to_pop = num_args - len(submodule_example_inputs) + while num_to_pop and normalized_kwargs: + normalized_kwargs.popitem(last=False) + num_to_pop -= 1 + submodule_example_inputs.extend(normalized_kwargs.values()) + submodule_example_inputs_tuple = tuple(submodule_example_inputs) + fqn = _get_path_of_module(root, self) + if fqn is not None: + fqn_to_example_inputs[fqn] = submodule_example_inputs_tuple + return orig_module_call(self, *args, **kwargs) + + orig_module_call = torch.nn.Module.__call__ + torch.nn.Module.__call__ = _patched_module_call # type: ignore[method-assign] + try: + model(*example_inputs) + finally: + # restore the module call even if there is an exception + torch.nn.Module.__call__ = orig_module_call # type: ignore[method-assign] + return fqn_to_example_inputs + + +# Returns a function that can get a new attribute name for module with given +# prefix, for example, +# >> get_new_observer_name = get_new_attr_name_with_prefix('_observer') +# >> new_name = get_new_observer_name(module) +# new_name will be an unused attribute name on module, e.g. `_observer_1` +def get_new_attr_name_with_prefix(prefix: str) -> Callable: + prefix = prefix.replace(".", "_") + + def get_new_attr_name(module: torch.nn.Module): + def get_attr_name(i: int): + return prefix + str(i) + + i = 0 + attr_name = get_attr_name(i) + while hasattr(module, attr_name): + i += 1 + attr_name = get_attr_name(i) + return attr_name + + return get_new_attr_name + + +def create_getattr_from_value( + module: torch.nn.Module, graph: Graph, prefix: str, value: Any +) -> Node: + """ + Given a value of any type, creates a getattr node corresponding to the value and + registers the value as a buffer to the module. + """ + get_new_attr_name = get_new_attr_name_with_prefix(prefix) + attr_name = get_new_attr_name(module) + device = _assert_and_get_unique_device(module) + new_value = ( + value.detach().clone() + if isinstance(value, torch.Tensor) + else torch.tensor(value, device=device) + ) + module.register_buffer(attr_name, new_value) + # Create get_attr with value + attr_node = graph.create_node("get_attr", attr_name) + return attr_node + + +_QUANTIZE_OPS = [ + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor, + torch.ops.quantized_decomposed.quantize_per_channel.default, +] + + +_DEQUANTIZE_OPS = [ + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor, + torch.ops.quantized_decomposed.dequantize_per_channel.default, +] + + +def _is_connected(source: torch.fx.Node, dest: torch.fx.Node) -> bool: + """ + Assuming dest is one of the ops inserted by quant workflow, this function + finds if source and dest are connected. Assumption is that only quant workflow + inserted ops exist between source and dest + """ + quant_workflow_ops = _QUANTIZE_OPS + _DEQUANTIZE_OPS + quant_workflow_ops.append(torch.ops.quantized_decomposed.choose_qparams.tensor) + while dest.target in quant_workflow_ops: + if not isinstance(dest.args[0], torch.fx.Node): + raise ValueError( + f"expected arg[0] of quant workflow ops to be a node but found {dest.args[0]}" + ) + dest = dest.args[0] + return dest == source + + +def _get_tensor_constant_from_node(node, m): + if node is None: + return None + assert node.op == "get_attr" + target_atoms = node.target.split(".") + attr_itr = m + for i, atom in enumerate(target_atoms): + if not hasattr(attr_itr, atom): + raise RuntimeError( + f"Node referenced nonexistent target {'.'.join(target_atoms[:i])}" + ) + attr_itr = getattr(attr_itr, atom) + return attr_itr + + +def _get_all_arguments(orig_args, orig_kwargs, args_schema): + all_args = [] + for i, schema in enumerate(args_schema): + if schema.name in orig_kwargs: + all_args.append(orig_kwargs[schema.name]) + elif not schema.kwarg_only and i < len(orig_args): + all_args.append(orig_args[i]) + else: + all_args.append(schema.default_value) + return all_args + + +def _is_supported_batch_norm_for_training(node: Node): + """ + Return True if the given node refers to an aten batch norm op QAT supports. + """ + supported_ops = [ + torch.ops.aten.batch_norm.default, + torch.ops.aten._native_batch_norm_legit.default, + # Note: we won't need this op anymore after batch norm consolidation + # For now, we need to continue to support it because it gives better + # training numerics than `_native_batch_norm_legit` + torch.ops.aten.cudnn_batch_norm.default, + torch.ops.aten.miopen_batch_norm.default, + ] + return node.target in supported_ops + + +def _is_conv_node(n: Node): + """ + Return whether the node refers to an aten conv op. + """ + return n.op == "call_function" and n.target in [ + torch.ops.aten.conv1d.default, + torch.ops.aten.conv2d.default, + ] + + +def _is_conv_transpose_node(n: Node): + """ + Return whether the node refers to an aten conv_transpose op. + """ + return n.op == "call_function" and n.target in [ + torch.ops.aten.conv_transpose1d, + torch.ops.aten.conv_transpose1d.default, + torch.ops.aten.conv_transpose2d, + torch.ops.aten.conv_transpose2d.input, + ] + + +def _is_conv_or_conv_transpose_node(n: Node): + """ + Return whether the node refers to an aten conv or conv transpose op. + """ + return _is_conv_node(n) or _is_conv_transpose_node(n) + + +def _is_conv_transpose_fn(conv_fn: Callable): + return conv_fn in [F.conv_transpose1d, F.conv_transpose2d] + + +def _is_bn_node(n: Node): + return ( + _is_supported_batch_norm_for_training(n) + or n.target == torch.ops.aten._native_batch_norm_legit_no_training.default + ) + + +def fold_bn_weights_into_conv_node( + conv_node: Node, + conv_weight_node: Node, + conv_bias_node: Optional[Node], + bn_node: Node, + m: GraphModule, +) -> None: + # conv args: input, weight, bias, stride, padding, dilation, ... + conv_w = _get_tensor_constant_from_node(conv_weight_node, m) + conv_b = _get_tensor_constant_from_node(conv_bias_node, m) + transpose = _is_conv_transpose_node(conv_node) + + # eval bn args: input, weight, bias, running mean, running var, momentum, eps + # train bn args: input, weight, bias, running mean, running var, training, momentum, eps + bn_args_schema = bn_node.target._schema.arguments # type: ignore[union-attr] + bn_args = _get_all_arguments(bn_node.args, bn_node.kwargs, bn_args_schema) + bn_w = _get_tensor_constant_from_node(bn_args[1], m) + bn_b = _get_tensor_constant_from_node(bn_args[2], m) + bn_rm = _get_tensor_constant_from_node(bn_args[3], m) + bn_rv = _get_tensor_constant_from_node(bn_args[4], m) + if bn_node.target == torch.ops.aten._native_batch_norm_legit_no_training.default: + eps_arg_index = 6 + elif _is_supported_batch_norm_for_training(bn_node): + eps_arg_index = 7 + else: + raise ValueError("BN node target is unexpected ", bn_node.target) + bn_eps = bn_args[eps_arg_index] + + fused_weight, fused_bias = fuse_conv_bn_weights( + conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b, transpose=transpose + ) + + # update the weight and bias for conv + conv_args = list(conv_node.args) + # filling in the default bias argument + if len(conv_args) == 2: + conv_args.append(None) + + # calling data since the fused_weight and fused_bias are nn.Parameter + weight_attr_name = conv_weight_node.target + assert isinstance(weight_attr_name, str) + _assign_attr(fused_weight, m, weight_attr_name, _AttrKind.PARAMETER) + if conv_bias_node is not None: + bias_attr_name = conv_bias_node.target + _assign_attr(fused_bias, m, str(bias_attr_name), _AttrKind.PARAMETER) + else: + bias_attr_name = weight_attr_name + "_bias" + _assign_attr(fused_bias, m, bias_attr_name, _AttrKind.PARAMETER) + with m.graph.inserting_before(conv_node): + get_bias_node = m.graph.get_attr(bias_attr_name) + # NOTE: here we assume the bias of conv is not quantized! + conv_args[2] = get_bias_node + conv_node.args = tuple(conv_args) + + # native_batch_norm has 3 outputs, we expect getitem calls on the output + # and we want to replace the uses of getitem 0 with the output of conv + # + if bn_node.target == torch.ops.aten.batch_norm.default: + # With the new training ir, instead of batch_norm + getitem, + # we only have the batch_norm node. + # + # Before: + # conv -> bn -> users + # After: + # conv -> users + # bn has no users now + bn_node.replace_all_uses_with(conv_node) + else: + # Before: + # conv -> bn - (first output) -> users1 + # \ - (second output) -> users2 + # \ - (third output) -> users3 + # After: + # conv -> (first output) -> users1 + # bn - + # \ - (second output) -> users2 + # \ - (third output) -> users3 + # if users2 and users3 are empty then bn will be removed through dead code elimination + for user in bn_node.users: + if ( + user.op != "call_function" + or user.target != operator.getitem + or user.args[1] != 0 + ): + continue + user.replace_all_uses_with(conv_node) + + # If the BN node does not have users, erase it from the graph + # Note: we need to do this manually because the model can still be in train + # mode at this point, in which case DCE won't erase the BN node automatically + # since the node refers to a mutating op. Here we still need to call DCE first + # to get rid of the unused getitem nodes that consume the BN node. + m.graph.eliminate_dead_code() + if len(bn_node.users) == 0: + m.graph.erase_node(bn_node) + + +# fuse conv bn weights, inplace modification of the graph_module and graph +def _fuse_conv_bn_(m: GraphModule) -> None: + has_bn = any(_is_bn_node(n) for n in m.graph.nodes) + if not has_bn: + return + for n in m.graph.nodes: + if n.op != "call_function" or n.target not in ( + torch.ops.aten._native_batch_norm_legit_no_training.default, + torch.ops.aten.batch_norm.default, + ): + continue + bn_node = n + n = bn_node.args[0] + if not _is_conv_or_conv_transpose_node(n): + continue + conv_node = n + conv_weight_node = conv_node.args[1] + conv_bias_node = conv_node.args[2] if len(conv_node.args) > 2 else None + fold_bn_weights_into_conv_node( + conv_node, conv_weight_node, conv_bias_node, bn_node, m + ) + + m.graph.eliminate_dead_code() + m.recompile() + + +def _get_node_name_to_scope(model: GraphModule) -> dict[str, tuple[str, type]]: + # TODO: move this information to fx node itself + node_name_to_scope: dict[str, tuple[str, type]] = {} + for n in model.graph.nodes: + nn_module_stack = n.meta.get("nn_module_stack", None) + current_scope = ("", type(None)) + if nn_module_stack: + bt = list(nn_module_stack.values())[-1] + current_scope = (bt[0].split(".")[-1], bt[1]) + node_name_to_scope[n.name] = current_scope + return node_name_to_scope + + +def _get_aten_graph_module_for_pattern( + pattern: Callable, + example_inputs: tuple[Any, ...], + is_cuda: bool = False, + **kwargs, +) -> GraphModule: + """ + Convert the pattern to an FX graph with decomposed aten ops. + """ + if is_cuda: + example_inputs = tuple( + [x.cuda() if isinstance(x, torch.Tensor) else x for x in example_inputs] + ) + + aten_pattern = torch.export.export_for_training( + pattern, # type: ignore[arg-type] + example_inputs, + kwargs, + strict=True, + ).module() + + aten_pattern.graph.eliminate_dead_code() # type: ignore[operator, union-attr] + aten_pattern.recompile() # type: ignore[operator] + + # ep.module() adds copy_ nodes for the mutated inputs. + # For patterns, it doesn't matter + for node in aten_pattern.graph.nodes: # type: ignore[union-attr] + if ( + node.op == "call_function" + and node.target == torch.ops.aten.copy_.default + and len(node.users) == 0 + ): + aten_pattern.graph.erase_node(node) # type: ignore[operator, union-attr] + + aten_pattern.graph.eliminate_dead_code() # type: ignore[operator, union-attr] + aten_pattern.recompile() # type: ignore[operator] + + return aten_pattern # type: ignore[return-value] + + +def remove_tensor_overload_for_qdq_ops(match_pattern: GraphModule) -> None: + """Remove .tensor overload for quantize/dequantize ops so that we can + use the match_pattern that we get from torchdynamo export to match the output of convert_pt2e + """ + _MAP = { + torch.ops.quantized_decomposed.quantize_per_tensor.default: torch.ops.quantized_decomposed.quantize_per_tensor, + torch.ops.quantized_decomposed.dequantize_per_tensor.default: torch.ops.quantized_decomposed.dequantize_per_tensor, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor: torch.ops.quantized_decomposed.quantize_per_tensor, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor: torch.ops.quantized_decomposed.dequantize_per_tensor, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor2: torch.ops.quantized_decomposed.quantize_per_tensor, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor2: torch.ops.quantized_decomposed.dequantize_per_tensor, + torch.ops.quantized_decomposed.quantize_per_channel.default: torch.ops.quantized_decomposed.quantize_per_channel, + torch.ops.quantized_decomposed.dequantize_per_channel.default: torch.ops.quantized_decomposed.dequantize_per_channel, + torch.ops.aten.clamp.Tensor: torch.ops.aten.clamp, + } + for n in match_pattern.graph.nodes: + if n.op != "call_function": + continue + if n.target in _MAP: + n.target = _MAP[n.target] + + +def _is_literal(arg): + if isinstance(arg, (int, float)): + return True + if isinstance(arg, (tuple, list)): + return all(map(_is_literal, arg)) + return False + + +def _replace_literals_with_new_placeholders( + gm: torch.fx.GraphModule, + merge_dup: bool = False, + exclude_literals: Optional[list[Any]] = None, +): + """Replace the literals in the graph with placeholder nodes that's created on the fly while we + traverse the graph, so that the literal arguments in the graph can be matched and replaced + + To use this, the pattern and replacement graph should have the exact same number of literal args + and they should be used in the exact same order in the pattern and replacement graph. + + If the literal arguments are not used in the same order in pattern and replacement graph, please + use `_replace_literals_with_existing_placeholders` instead + + Args: + `gm`: input GraphModule that we'll transform + `merge_dup`: boolean flag to indicate that if the same literal appears multiple times in + the graph, whether they should correspond to the same placeholder or not + `exclude_literals`: a list of literals that will not be replaced with placeholders + + Example: + + # 1. Original Graph + def pattern(self, x): + return x + 3 + + def replacement(self, x): + return x - 3 + + example_inputs = (torch.randn(1, 3, 3, 3),) + pattern_gm = _get_aten_graph_module_for_pattern(pattern, example_inputs) + replacement_gm = _get_aten_graph_module_for_pattern(pattern, example_inptus) + + # 2. Before calling replace literals we'll see the following graph: + def pattern(self, x): + return x + 3 + + def replacement(self, x): + return x - 3 + + pattern_gm = _replace_literals_with_new_placeholders(pattern_gm) + replacement_gm = _replace_literals_with_new_placeholders(replacement_gm) + + # 3. After replacing literals with new placeholder nodes + + def pattern(self, x, new_ph): + return x + new_ph + + def pattern(self, x, new_ph): + return x - new_ph + + """ + last_ph = None + cnt = 0 + literal_to_ph: dict[Union[float, bool, int, torch.dtype], Node] = {} + if exclude_literals is None: + exclude_literals = [] + + in_spec = gm._in_spec + args_spec = in_spec.children_specs[0] + for node in gm.graph.nodes: + if node.op == "placeholder": + last_ph = node + cnt += 1 + continue + with gm.graph.inserting_after(last_ph): + new_args = [] + for arg in node.args: + if _is_literal(arg) and arg not in exclude_literals: + if merge_dup and arg in literal_to_ph: + new_args.append(literal_to_ph[arg]) + else: + ph_node = gm.graph.placeholder("arg" + str(cnt)) + new_args.append(ph_node) + args_spec.children_specs.append(LeafSpec()) + cnt += 1 + if merge_dup: + literal_to_ph[arg] = ph_node + else: + new_args.append(arg) + new_args = tuple(new_args) + + node.args = new_args + + # Update `num_nodes`, `num_leaves`, `num_children`. + args_spec.__post_init__() + in_spec.__post_init__() + return gm + + +def _replace_literals_with_existing_placeholders( + gm: torch.fx.GraphModule, + exclude_literals: Optional[list[Any]] = None, + literal_to_ph_idx: Optional[dict[Union[float, int, bool, torch.dtype], int]] = None, +): + """Replace the literals in the graph with **existing** placeholder nodes, so that the literal arguments + in the graph can be matched and replaced + + To use this, all literal args in the graph should be unique and each of them should correspond + to exactly one placeholder node + + # 1. Original Graph + def pattern(self, x_i8, scale, zero_point, quant_min, quant_max): + return torch.dequantize_per_tensor(x_i8, scale, zero_point, quant_min, quant_max) + + def replacement(x_i8, scale, zero_point, quant_min, quant_max): + x_i8 = torch.clamp(x_i8, quant_min, quant_max) + return ((x_i8.to(torch.float32) - zero_point) * scale).to(dtype=torch.float32) + + example_inputs = ( + torch.randn(1, 3, 3, 3), + 1.0, + 0, + -128, + 127, + ) + pattern_gm = _get_aten_graph_module_for_pattern(pattern, example_inputs) + replacement_gm = _get_aten_graph_module_for_pattern(pattern, example_inptus) + + # 2. Before calling replace literals we'll see the following graph: + def pattern(self, x_i8, scale, zero_point, quant_min, quant_max): + # scale/zero_point/quant_min/quant_max are burnt in since they are scalar values + return torch.dequantize_per_tensor(x_i8, 1.0, 0, -128, 127) + + def replacement(x_i8, scale, zero_point, quant_min, quant_max): + # scale/zero_point/quant_min/quant_max are burnt in since they are scalar values + x_i8 = torch.clamp(x_i8, -128, 127) + return ((x_i8.to(torch.float32) - 0) * 1.0).to(dtype=torch.float32) + + # Note that literal args appear in different order in pattern and replacement graph, so + # we can't use _replace_literals_with_new_placeholders + + literal_to_ph_idx = {1.0: 1, 0: 2, -128: 3, 127: 4} + pattern_gm = _replace_literals_with_existing_placeholders(pattern_gm, literal_to_ph_idx) + replacement_gm = _replace_literals_with_existing_placeholders(replacement_gm, literal_to_ph_idx) + + # 3. After replacing literals with existing placeholder nodes + + def pattern(self, x_i8, scale, zero_point, quant_min, quant_max): + # scale/zero_point/quant_min/quant_max are burnt in since they are scalar values + return torch.dequantize_per_tensor(x_i8, scale, zero_point, quant_min, quant_max) + + def replacement(x_i8, scale, zero_point, quant_min, quant_max): + # scale/zero_point/quant_min/quant_max are burnt in since they are scalar values + x_i8 = torch.clamp(x_i8, quant_min, quant_max) + return ((x_i8.to(torch.float32) - zero_point) * scale).to(dtype=torch.float32) + """ + if exclude_literals is None: + exclude_literals = [] + + if literal_to_ph_idx is None: + literal_to_ph_idx = {} + + phs = [node for node in gm.graph.nodes if node.op == "placeholder"] + + for node in gm.graph.nodes: + if node.op != "call_function": + continue + new_args = [] + for arg in node.args: + if ( + _is_literal(arg) + and arg not in exclude_literals + and arg in literal_to_ph_idx + ): + ph_idx = literal_to_ph_idx[arg] + ph_node = phs[ph_idx] + new_args.append(ph_node) + else: + new_args.append(arg) + new_args = tuple(new_args) + node.args = new_args + return gm + + +# TODO: Handle this in export itself and don't wrap the model in another GraphModule +# in prepare and convert +def _disallow_eval_train(model: GraphModule): + """ + Disallow calling `model.train()` or `model.eval()` on the given GraphModule. + This is useful for exported models, where these methods don't actually behave as expected. + """ + error_message = """ + Calling train() or eval() is not supported for exported models. + Please call `torchao.quantization.pt2e.move_exported_model_to_train(model)` (or eval) instead. + + If you cannot replace the calls to `model.train()` and `model.eval()`, you may override + the behavior for these methods by calling `torchao.quantization.pt2e.allow_exported_model_train_eval(model)`, + which does the above automatically for you. Note that this has limited effect on switching + behavior between train and eval modes, and should be used only for special ops such as dropout + and batchnorm. + """ + + def _train(self, mode: bool = True): + raise NotImplementedError(error_message) + + def _eval(self, mode: bool = True): + raise NotImplementedError(error_message) + + model.train = types.MethodType(_train, model) # type: ignore[method-assign] + model.eval = types.MethodType(_eval, model) # type: ignore[method-assign] + return model + + +def _is_sym_size_node(node: Node): + return ( + node.op == "call_function" + and node.target == torch.ops.aten.sym_size.default + or node.target == torch.ops.aten.sym_numel.default + or node.target == torch.ops.aten.sym_numel + or node.target == torch.ops.aten.sym_size + ) + + +def _filter_sym_size_users(node: torch.fx.Node) -> list[torch.fx.Node]: + node_users = list(filter((lambda x: (_is_sym_size_node(x) is False)), node.users)) + return node_users diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/__init__.py b/lib/python3.12/site-packages/torchao/quantization/qat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5dc3d8e00806befaf82c715d26474f95e1ad8352 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/qat/__init__.py @@ -0,0 +1,27 @@ +from .api import ( + ComposableQATQuantizer, + FakeQuantizeConfig, + FromIntXQuantizationAwareTrainingConfig, + IntXQuantizationAwareTrainingConfig, + from_intx_quantization_aware_training, + intx_quantization_aware_training, +) +from .embedding import ( + Int4WeightOnlyEmbeddingQATQuantizer, +) +from .linear import ( + Int4WeightOnlyQATQuantizer, + Int8DynActInt4WeightQATQuantizer, +) + +__all__ = [ + "ComposableQATQuantizer", + "FakeQuantizeConfig", + "Int4WeightOnlyQATQuantizer", + "Int4WeightOnlyEmbeddingQATQuantizer", + "Int8DynActInt4WeightQATQuantizer", + "intx_quantization_aware_training", + "from_intx_quantization_aware_training", + "FromIntXQuantizationAwareTrainingConfig", + "IntXQuantizationAwareTrainingConfig", +] diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..624236b634fc937a823629c673e53df3f18e32c1 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/affine_fake_quantized_tensor.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/affine_fake_quantized_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..085fbbaf01b06f4fbb1beec1f9ab31ae96eeb44d Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/affine_fake_quantized_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..99468c8414b4c0807e64e5d16949599a6a5c1a2c Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/embedding.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/embedding.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa4e199b93d6d963df4f21d4ec7182fa7f7e93cb Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/embedding.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/fake_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/fake_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..073efaf085cbd40e19d29196983da2846e9d50a3 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/fake_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/linear.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33022c13a3f0fef5df496b31bf272ecb6153110a Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d951ceaee3d12044c814ab2d32db88d3c180f17e Binary files /dev/null and b/lib/python3.12/site-packages/torchao/quantization/qat/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/affine_fake_quantized_tensor.py b/lib/python3.12/site-packages/torchao/quantization/qat/affine_fake_quantized_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..6829441f517e4407524df30043d4624187384e9d --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/qat/affine_fake_quantized_tensor.py @@ -0,0 +1,361 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Callable, Optional, Tuple + +import torch +import torch.utils._pytree as pytree +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.quantization.quant_primitives import ( + MappingType, + ZeroPointDomain, + _get_and_check_qmin_qmax, + choose_qparams_affine, +) +from torchao.utils import TorchAOBaseTensor + +from .utils import ( + _GenericFakeQuantize, + _UnwrapAffineFakeQuantizedTensor, +) + +aten = torch.ops.aten + + +class _ToAffineFakeQuantized(torch.autograd.Function): + """ + Differentiable constructor for `AffineFakeQuantizedTensor`, + needed for input activation fake quantization. + """ + + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + original_tensor: torch.Tensor, + mapping_type: MappingType, + block_size: Tuple[int, ...], + target_dtype: torch.dtype, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + ) -> "AffineFakeQuantizedTensor": + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + + def apply_fake_quant_fn(t: torch.Tensor): + assert isinstance(t, AffineFakeQuantizedTensor) + qmin, qmax = _get_and_check_qmin_qmax(target_dtype, quant_min, quant_max) + scale, zero_point = choose_qparams_affine( + t.original_tensor, + mapping_type, + block_size, + target_dtype, + qmin, + qmax, + eps, + scale_dtype, + zero_point_dtype, + preserve_zero, + zero_point_domain, + ) + fq = _GenericFakeQuantize.apply( + t, + block_size, + scale, + zero_point, + qmin, + qmax, + zero_point_domain, + ) + return fq + + return AffineFakeQuantizedTensor( + original_tensor, + apply_fake_quant_fn, + fake_quant_enabled=True, + ) + + @staticmethod + def backward(ctx, gy): + return gy, None, None, None, None, None, None, None, None, None, None + + +class AffineFakeQuantizedTensor(TorchAOBaseTensor): + """ + Affine fake quantized tensor subclass. Affine quantization means we quantize the floating point tensor + with an affine transformation: + quantized_tensor = float_tensor / scale + zero_point + + Fake quantization refers to performing the quantization math without actually casting the floating point + tensor into lower bit-width dtypes. It is commonly used for quantization-aware training (QAT). + + The shape and dtype of the tensor subclass represent how the tensor subclass looks externally, + regardless of the internal representation's type or orientation. + + fields: + original_tensor (torch.Tensor): tensor holding the original float values, needed for actual quantization later + apply_fake_quant_fn (Callable): function that transforms `original_tensor` to fake quantized values + """ + + @staticmethod + def __new__( + cls, + original_tensor: torch.Tensor, + apply_fake_quant_fn: Callable, + fake_quant_enabled: bool = True, + **kwargs, + ): + kwargs.setdefault("dtype", original_tensor.dtype) + kwargs.setdefault("device", original_tensor.device) + kwargs.setdefault("requires_grad", original_tensor.requires_grad) + return torch.Tensor._make_wrapper_subclass( + cls, + original_tensor.shape, + **kwargs, + ) + + def __init__( + self, + original_tensor: torch.Tensor, + apply_fake_quant_fn: Callable, + fake_quant_enabled: bool = True, + **kwargs, + ): + self.original_tensor = original_tensor + self.apply_fake_quant_fn = apply_fake_quant_fn + self.fake_quant_enabled = fake_quant_enabled + + def __tensor_flatten__(self): + return ["original_tensor"], [self.apply_fake_quant_fn, self.fake_quant_enabled] + + @classmethod + def __tensor_unflatten__( + cls, + tensor_data_dict, + tensor_attributes, + outer_size, + outer_stride, + ): + original_tensor = tensor_data_dict["original_tensor"] + (apply_fake_quant_fn, fake_quant_enabled) = tensor_attributes + return cls( + original_tensor, + apply_fake_quant_fn, + fake_quant_enabled, + ) + + @classmethod + def from_float( + cls, + original_input: torch.Tensor, + mapping_type: MappingType, + block_size: Tuple[int, ...], + target_dtype: torch.dtype, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + ): + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + return _ToAffineFakeQuantized.apply( + original_input, + mapping_type, + block_size, + target_dtype, + quant_min, + quant_max, + eps, + scale_dtype, + zero_point_dtype, + preserve_zero, + zero_point_domain, + ) + + def get_value(self) -> torch.Tensor: + if self.fake_quant_enabled: + return self.apply_fake_quant_fn(self) + else: + return _UnwrapAffineFakeQuantizedTensor.apply(self) + + def _get_to_kwargs(self, *args, **kwargs): + device, dtype, _, memory_format = torch._C._nn._parse_to(*args, **kwargs) + device = self.device if device is None else device + dtype = self.dtype if dtype is None else dtype + memory_format = ( + memory_format if memory_format is not None else torch.preserve_format + ) + kwargs = { + "device": device, + "dtype": dtype, + "memory_format": memory_format, + "requires_grad": self.requires_grad, + } + return kwargs + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs.pop("device") + # not supported yet + kwargs.pop("memory_format") + return self.__class__( + self.original_tensor.to(device), + self.apply_fake_quant_fn, + self.fake_quant_enabled, + **kwargs, + ) + + def _apply_fn_to_data(self, fn: Callable): + """ + Create a new `AffineFakeQuantizedTensor` with `fn` applied to the + original tensor, to be called within __torch_dispatch__. + """ + return self._create_new(fn(self.original_tensor)) + + def _create_new(self, new_value: torch.Tensor): + """ + Create a new `AffineFakeQuantizedTensor` with a new value, + to be called within __torch_dispatch__. + + Note: `requires_grad` must be False here because tensors created + in `__torch_dispatch__` cannot produce gradients, since autograd + will try to attach autograd metadata to these tensors when we exit + `__torch_dispatch__`, but if these tensors already have metadata + attached then autograd will throw an error. + """ + return self.__class__( + new_value, + self.apply_fake_quant_fn, + self.fake_quant_enabled, + requires_grad=False, + ) + + +implements = AffineFakeQuantizedTensor.implements + + +@implements(torch.nn.functional.linear) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + if isinstance(input_tensor, AffineFakeQuantizedTensor): + input_tensor = input_tensor.get_value() + if isinstance(weight_tensor, AffineFakeQuantizedTensor): + weight_tensor = weight_tensor.get_value() + return torch.nn.functional.linear(input_tensor, weight_tensor, bias) + + +@implements(aten.mm.default) +def _(func, types, args, kwargs): + input_tensor = args[0] + weight_tensor = args[1] + if isinstance(input_tensor, AffineFakeQuantizedTensor): + input_tensor = input_tensor.get_value() + if isinstance(weight_tensor, AffineFakeQuantizedTensor): + weight_tensor = weight_tensor.get_value() + return func(input_tensor, weight_tensor) + + +@implements(aten.addmm.default) +def _(func, types, args, kwargs): + bias = args[0] + input_tensor = args[1] + weight_tensor = args[2] + if isinstance(input_tensor, AffineFakeQuantizedTensor): + input_tensor = input_tensor.get_value() + if isinstance(weight_tensor, AffineFakeQuantizedTensor): + weight_tensor = weight_tensor.get_value() + return func(bias, input_tensor, weight_tensor) + + +@implements(aten.detach.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data(torch.detach), + ) + + +@implements(aten.clone.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data(torch.clone), + ) + + +@implements(aten.t.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0]._apply_fn_to_data(torch.t), + ) + + +@implements( + [ + aten.add.Tensor, + aten.add_.Tensor, + aten.mul_.Tensor, + aten.copy_.default, + ] +) +def _(func, types, args, kwargs): + assert len(args) == 2, f"dispatched the wrong op to the binary handler: {func}" + new_args = pytree.tree_map_only( + AffineFakeQuantizedTensor, lambda x: x.original_tensor, args + ) + first_afq_tensor = ( + args[0] if isinstance(args[0], AffineFakeQuantizedTensor) else args[1] + ) + new_value = func(*new_args, **kwargs) + out = first_afq_tensor._create_new(new_value) + return return_and_correct_aliasing(func, args, kwargs, out) + + +# Needed by FSDP: + + +@implements(aten.empty_like.default) +def _(func, types, args, kwargs): + out = torch.empty_like(args[0].original_tensor, **kwargs) + return return_and_correct_aliasing(func, args, kwargs, out) + + +@implements(aten.split.Tensor) +def _(func, types, args, kwargs): + new_values = torch.split(args[0].original_tensor, *args[1:], **kwargs) + + def make_new_tensor(value): + out = args[0]._create_new(value) + return return_and_correct_aliasing(func, args, kwargs, out) + + return list(map(make_new_tensor, new_values)) + + +@implements(aten.new_zeros.default) +def _(func, types, args, kwargs): + out = args[0].original_tensor.new_zeros(*args[1:], **kwargs) + return return_and_correct_aliasing(func, args, kwargs, out) + + +to_affine_fake_quantized = AffineFakeQuantizedTensor.from_float diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/api.py b/lib/python3.12/site-packages/torchao/quantization/qat/api.py new file mode 100644 index 0000000000000000000000000000000000000000..d7e8f204cc34ae349abbce39015afad82e29a7c8 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/qat/api.py @@ -0,0 +1,393 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass +from typing import Any, List, Optional, Union + +import torch + +from torchao.core.config import AOBaseConfig +from torchao.quantization.granularity import ( + Granularity, + PerAxis, + PerGroup, + PerToken, +) +from torchao.quantization.quant_primitives import ( + _SUB_BYTE_INT_BOUNDS, + _SUB_BYTE_UINT_BOUNDS, + MappingType, + TorchAODType, + ZeroPointDomain, +) +from torchao.quantization.transform_module import ( + register_quantize_module_handler, +) +from torchao.quantization.unified import TwoStepQuantizer + + +@dataclass +class FakeQuantizeConfig: + """ + Config for how to fake quantize weights or activations. + + args: + dtype: dtype to simulate during fake quantization, e.g. torch.int8. + For PyTorch versions older than 2.6, you may use `TorchAODType` to represent + torch.int1 to torch.int7 instead, e.g. TorchAODType.INT4. + granularity: granularity of scales and zero points, e.g. PerGroup(32). + We also support the following strings: + 1) 'per_token': equivalent to PerToken() + 2) 'per_channel': equivalent to PerAxis(0) + 3) 'per_group': equivalent to PerGroup(group_size), must be combined + with separate `group_size` kwarg, Alternatively, just set the + `group_size` kwarg and leave this field empty. + mapping_type: whether to use symmetric (default) or asymmetric quantization + Alternatively, set `is_symmetric` (bool) and leave this field empty. + scale_precision: scale dtype (default torch.fp32) + zero_point_precision: zero point dtype (default torch.int32) + zero_point_domain: whether zero point is in integer (default) or float domain + is_dynamic: whether to use dynamic (default) or static scale and zero points + range_learning: whether to learn scale and zero points during training (coming soon) + + kwargs (optional): + group_size: size of each group in per group fake quantization, + can be set instead of `granularity` + is_symmetric: whether to use symmetric or asymmetric quantization, + can be set instead of `mapping_type` + + Example usage:: + + # Per token asymmetric quantization + FakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False) + FakeQuantizeConfig(torch.int8, PerToken(), MappingType.ASYMMETRIC) + + # Per channel symmetric quantization + FakeQuantizeConfig(torch.int4, "per_channel") + FakeQuantizeConfig(torch.int4, "per_channel", is_symmetric=True) + FakeQuantizeConfig(torch.int4, PerAxis(0), MappingType.SYMMETRIC) + + # Per group symmetric quantization + FakeQuantizeConfig(torch.int4, group_size=32) + FakeQuantizeConfig(torch.int4, group_size=32, is_symmetric=True) + FakeQuantizeConfig(torch.int4, "per_group", group_size=32, is_symmetric=True) + FakeQuantizeConfig(torch.int4, PerGroup(32), MappingType.SYMMETRIC) + """ + + dtype: Union[torch.dtype, TorchAODType] + granularity: Granularity + mapping_type: MappingType + scale_precision: torch.dtype + zero_point_precision: torch.dtype + zero_point_domain: ZeroPointDomain + is_dynamic: bool = True + range_learning: bool = False + + def __init__( + self, + dtype: Union[torch.dtype, TorchAODType], + granularity: Union[Granularity, str, None] = None, + mapping_type: Optional[MappingType] = None, + scale_precision: torch.dtype = torch.float32, + zero_point_precision: torch.dtype = torch.int32, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + is_dynamic: bool = True, + range_learning: bool = False, + *, + group_size: Optional[int] = None, + is_symmetric: Optional[bool] = None, + ): + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + self.dtype = dtype + self.granularity = self._get_granularity(granularity, group_size) + self.mapping_type = self._get_mapping_type(mapping_type, is_symmetric) + self.scale_precision = scale_precision + self.zero_point_precision = zero_point_precision + self.zero_point_domain = zero_point_domain + self.is_dynamic = is_dynamic + self.range_learning = range_learning + + # Validate dtype + all_dtypes = [torch.int8, torch.uint8] + all_dtypes.extend(list(_SUB_BYTE_INT_BOUNDS.keys())) + all_dtypes.extend(list(_SUB_BYTE_UINT_BOUNDS.keys())) + if dtype not in all_dtypes: + raise ValueError( + "Unsupported dtype '%s', choose from %s" % (dtype, all_dtypes) + ) + + def _get_granularity( + self, + granularity: Union[Granularity, str, None], + group_size: Optional[int], + ) -> Granularity: + """ + Parse the `Granularity` represented in the args. + + Granularity can be specified in one of three ways: + 1) `Granularity` object: one of PerToken(), PerAxis(), and PerGroup(group_size) + 2) str: one of 'per_token', 'per_channel', and 'per_group' + 3) None: `group_size` must be set instead, represents per group granularity + """ + # If group_size is set, then granularity must be either "per_group" or None + if ( + group_size is not None + and granularity != "per_group" + and granularity is not None + ): + raise ValueError( + "`group_size` conflicts with granularity '%s'" % granularity + ) + + # Case 1: Granularity object + if isinstance(granularity, Granularity): + if not isinstance(granularity, (PerToken, PerAxis, PerGroup)): + raise ValueError("Granularity '%s' is not supported" % granularity) + if isinstance(granularity, PerAxis) and granularity.axis != 0: + raise ValueError("Only axis=0 is supported for PerAxis granularity") + return granularity + + # Case 2: str granularity + if granularity == "per_token": + return PerToken() + elif granularity == "per_channel": + return PerAxis(axis=0) + elif granularity == "per_group": + if group_size is None: + raise ValueError( + "Granularity was 'per_group' but no `group_size` was set" + ) + return PerGroup(group_size) + elif isinstance(granularity, str): + raise ValueError( + "Unexpected granularity: '%s', must be one of %s" + % (granularity, ["per_token", "per_channel", "per_group"]) + ) + + # Case 3: None granularity + group_size was specified + if granularity is not None: + raise ValueError( + "Granularity '%s' has unexpected type %s" + % (granularity, type(granularity)) + ) + if group_size is None: + raise ValueError( + "At least one of `granularity` or `group_size` must be set" + ) + return PerGroup(group_size) + + def _get_mapping_type( + self, + mapping_type: Optional[MappingType], + is_symmetric: Optional[bool], + ) -> MappingType: + """ + Parse the `MappingType` represented in the args. + + Mapping type can be specified in one of two ways: + 1): `MappingType` object: one of SYMMETRIC or ASYMMETRIC + 2): is_symmetric bool + """ + if mapping_type is not None and is_symmetric is not None: + raise ValueError("Cannot set both `mapping_type` and `is_symmetric`") + + # Case 0: Default to symmetric + if mapping_type is None and is_symmetric is None: + return MappingType.SYMMETRIC + + # Case 1: MappingType object + if mapping_type is not None: + if mapping_type not in [MappingType.SYMMETRIC, MappingType.ASYMMETRIC]: + raise ValueError("MappingType '%s' is not supported" % mapping_type) + return mapping_type + + # Case 2: is_symmetric flag + assert is_symmetric is not None + if is_symmetric: + return MappingType.SYMMETRIC + else: + return MappingType.ASYMMETRIC + + @property + def group_size(self) -> int: + """ + If this is per group granularity, return the group size. + Otherwise, throw an error. + """ + if isinstance(self.granularity, PerGroup): + return self.granularity.group_size + else: + raise ValueError( + "`group_size` is undefined for %s granularity" % self.granularity + ) + + @property + def is_symmetric(self) -> bool: + """ + Return True if mapping type is symmetric, else False (asymmetric). + """ + return self.mapping_type == MappingType.SYMMETRIC + + def __setattr__(self, name: str, value: Any): + """ + Support setting `group_size` and `is_symmetric`. + """ + if name == "group_size": + super().__setattr__("granularity", PerGroup(value)) + elif name == "is_symmetric": + mapping_type = MappingType.SYMMETRIC if value else MappingType.ASYMMETRIC + super().__setattr__("mapping_type", mapping_type) + else: + super().__setattr__(name, value) + + +@dataclass +class IntXQuantizationAwareTrainingConfig(AOBaseConfig): + activation_config: Optional[FakeQuantizeConfig] = None + weight_config: Optional[FakeQuantizeConfig] = None + + +# for BC +intx_quantization_aware_training = IntXQuantizationAwareTrainingConfig + + +@register_quantize_module_handler(IntXQuantizationAwareTrainingConfig) +def _intx_quantization_aware_training_transform( + module: torch.nn.Module, + config: IntXQuantizationAwareTrainingConfig, +) -> torch.nn.Module: + """ + THIS IS NOT A PUBLIC API - any usage of this outside of torchao + can break at any time. + + Apply fake quantization to a `torch.nn.Module`. + to be used with :func:`~torchao.quantization.quant_api.quantize_`. + + Example usage:: + + from torchao.quantization import quantize_ + from torchao.quantization.qat import FakeQuantizeConfig + activation_config = FakeQuantizeConfig( + torch.int8, "per_token", is_symmetric=False, + ) + weight_config = FakeQuantizeConfig( + torch.int4, group_size=32, is_symmetric=True, + ) + quantize_( + model, + IntXQuantizationAwareTrainingConfig(activation_config, weight_config), + ) + + Note: If the returned function is applied on a module that is not + `torch.nn.Linear` or `torch.nn.Embedding`, or it is applied on + `torch.nn.Embedding` with an activation config, then we will raise + ValueError as these are not supported. + """ + from .embedding import FakeQuantizedEmbedding + from .linear import FakeQuantizedLinear + + mod = module + activation_config = config.activation_config + weight_config = config.weight_config + + if isinstance(mod, torch.nn.Linear): + return FakeQuantizedLinear.from_linear( + mod, + activation_config, + weight_config, + ) + elif isinstance(mod, torch.nn.Embedding): + if activation_config is not None: + raise ValueError( + "Activation fake quantization is not supported for embedding" + ) + return FakeQuantizedEmbedding.from_embedding(mod, weight_config) + else: + raise ValueError("Module of type '%s' does not have QAT support" % type(mod)) + + +class FromIntXQuantizationAwareTrainingConfig(AOBaseConfig): + """ + Object that knows how to convert a model with fake quantized modules, + such as :func:`~torchao.quantization.qat.linear.FakeQuantizedLinear` + and :func:`~torchao.quantization.qat.linear.FakeQuantizedEmbedding`, + back to model with the original, corresponding modules without + fake quantization. This should be used with + :func:`~torchao.quantization.quant_api.quantize_`. + + Example usage:: + + from torchao.quantization import quantize_ + quantize_( + model_with_fake_quantized_linears, + FromIntXQuantizationAwareTrainingConfig(), + ) + """ + + pass + + +# for BC +from_intx_quantization_aware_training = FromIntXQuantizationAwareTrainingConfig + + +@register_quantize_module_handler(FromIntXQuantizationAwareTrainingConfig) +def _from_intx_quantization_aware_training_transform( + mod: torch.nn.Module, + config: FromIntXQuantizationAwareTrainingConfig, +) -> torch.nn.Module: + """ + If the given module is a fake quantized module, return the original + corresponding version of the module without fake quantization. + """ + from .embedding import FakeQuantizedEmbedding + from .linear import FakeQuantizedLinear + + if isinstance(mod, FakeQuantizedLinear): + return mod.to_linear() + elif isinstance(mod, FakeQuantizedEmbedding): + return mod.to_embedding() + else: + return mod + + +class ComposableQATQuantizer(TwoStepQuantizer): + """ + Composable quantizer that users can use to apply multiple QAT quantizers easily. + Quantizers will be applied in the order they are specified in the constructor. + + Note: the quantizers provided must apply to different modules in the model, + e.g. nn.Linear and nn.Embedding, otherwise the behavior will be undefined. + + Example usage:: + + my_quantizer = ComposableQATQuantizer([ + QATQuantizer1(), + QATQuantizer2(), + QATQuantizer3(), + ]) + model = my_quantizer.prepare(model) + train(model) + model = my_quantizer.convert(model) + """ + + def __init__(self, quantizers: List[TwoStepQuantizer]): + self.quantizers = quantizers + + def prepare( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + for quantizer in self.quantizers: + model = quantizer.prepare(model) + return model + + def convert( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + for quantizer in self.quantizers: + model = quantizer.convert(model) + return model diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/embedding.py b/lib/python3.12/site-packages/torchao/quantization/qat/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..2770956a2c890fe5adc46dc5494c0c0aff36e415 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/qat/embedding.py @@ -0,0 +1,405 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Any, Optional + +import torch +import torch.nn.functional as F + +from torchao.quantization.quant_primitives import TorchAODType +from torchao.quantization.unified import TwoStepQuantizer +from torchao.quantization.utils import get_group_qparams_symmetric + +from .api import FakeQuantizeConfig +from .fake_quantizer import FakeQuantizer +from .utils import ( + _get_qmin_qmax, +) + + +class FakeQuantizedEmbedding(torch.nn.Embedding): + """ + General embedding layer with fake quantized weights. + + Specific target dtypes, granularity, schemes etc. are specified + through separate configs for weights and activations. + + Example usage:: + + weight_config = FakeQuantizeConfig( + dtype=torch.int4, + group_size=8, + symmetric=True, + ) + fq_embedding = FakeQuantizedEmbedding(5, 10, weight_config) + fq_embedding(torch.LongTensor([3])) + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + padding_idx: Optional[int] = None, + max_norm: Optional[float] = None, + norm_type: float = 2.0, + scale_grad_by_freq: bool = False, + sparse: bool = False, + weight_config: Optional[FakeQuantizeConfig] = None, + *args, + **kwargs, + ) -> None: + super().__init__( + num_embeddings, + embedding_dim, + padding_idx, + max_norm, + norm_type, + scale_grad_by_freq, + sparse, + *args, + **kwargs, + ) + if weight_config is not None: + self.weight_fake_quantizer = FakeQuantizer(weight_config) + else: + self.weight_fake_quantizer = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.weight_fake_quantizer is not None: + w = self.weight_fake_quantizer(self.weight) + else: + w = self.weight + return F.embedding( + x, + w, + self.padding_idx, + self.max_norm, + self.norm_type, + self.scale_grad_by_freq, + self.sparse, + ) + + def to_embedding(self) -> torch.nn.Embedding: + new_embedding = torch.nn.Embedding( + self.num_embeddings, + self.embedding_dim, + self.padding_idx, + self.max_norm, + self.norm_type, + self.scale_grad_by_freq, + self.sparse, + device=self.weight.device, + ) + # In distributed training, the model may be instantiated + # on the meta device, in which case there is no need to + # copy the weights, and doing so will result in an error + if self.weight.device != torch.device("meta"): + new_embedding.weight = self.weight + return new_embedding + + @classmethod + def from_embedding( + cls, + mod: torch.nn.Embedding, + weight_config: Optional[FakeQuantizeConfig] = None, + ): + new_embedding = FakeQuantizedEmbedding( + mod.num_embeddings, + mod.embedding_dim, + mod.padding_idx, + mod.max_norm, + mod.norm_type, + mod.scale_grad_by_freq, + mod.sparse, + weight_config=weight_config, + device=mod.weight.device, + ) + # In distributed training, the model may be instantiated + # on the meta device, in which case there is no need to + # copy the weights, and doing so will result in an error + if mod.weight.device != torch.device("meta"): + new_embedding.weight = mod.weight + return new_embedding + + +# ====================================== +# | Embedding int4 weight-only QAT | +# ====================================== + + +class Int4WeightOnlyEmbeddingQATQuantizer(TwoStepQuantizer): + """ + Quantizer for performing QAT on a model, where embedding layers have + int4 fake quantized grouped per channel weights. + """ + + def __init__( + self, + group_size: int = 256, + scale_precision: torch.dtype = torch.float32, + zero_point_precision: torch.dtype = torch.int32, + ) -> None: + super().__init__() + self.bit_width = 4 + self.group_size: int = group_size + self.scale_precision: torch.dtype = scale_precision + self.zero_point_precision: torch.dtype = zero_point_precision + + def prepare( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + """ + Swap `nn.Embedding` modules with `Int4WeightOnlyQATEmbedding`. + """ + # avoid circular imports + from torchao.quantization.quant_api import ( + _replace_with_custom_fn_if_matches_filter, + ) + + def filter_fn(child: torch.nn.Module, cur_fqn: str) -> bool: + return isinstance(child, torch.nn.Embedding) + + def replacement_fn(child: torch.nn.Module) -> torch.nn.Module: + new_embedding = Int4WeightOnlyQATEmbedding( + # nn.Embedding args + num_embeddings=child.num_embeddings, + embedding_dim=child.embedding_dim, + padding_idx=child.padding_idx, + max_norm=child.max_norm, + norm_type=child.norm_type, + scale_grad_by_freq=child.scale_grad_by_freq, + sparse=child.sparse, + # quantization args + group_size=self.group_size, + scale_precision=self.scale_precision, + zero_point_precision=self.zero_point_precision, + device=child.weight.device, + dtype=child.weight.dtype, + ) + # In distributed training, the model may be instantiated + # on the meta device, in which case there is no need to + # copy the weights, and doing so will result in an error + if child.weight.device != torch.device("meta"): + new_embedding.weight = child.weight + return new_embedding + + _replace_with_custom_fn_if_matches_filter(model, replacement_fn, filter_fn) + return model + + def convert( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + """ + Swap all `Int4WeightOnlyQATEmbedding` modules with `Int4WeightOnlyEmbedding`. + """ + self._convert_helper(model) + return model + + def _convert_helper(self, module: torch.nn.Module): + """ + Helper function to recursively swap `Int4WeightOnlyQATEmbedding` + modules with `Int4WeightOnlyEmbedding` + """ + from torchao._executorch_ops import ( + _quantized_decomposed_quantize_per_channel_group_wrapper, + ) + + for name, child in module.named_children(): + if isinstance(child, Int4WeightOnlyQATEmbedding): + group_size = child.weight_fake_quantizer.config.group_size + scale_precision = child.weight_fake_quantizer.config.scale_precision + zero_point_precision = ( + child.weight_fake_quantizer.config.zero_point_precision + ) + quantized_embedding = Int4WeightOnlyEmbedding( + # nn.Embedding args + num_embeddings=child.num_embeddings, + embedding_dim=child.embedding_dim, + padding_idx=child.padding_idx, + max_norm=child.max_norm, + norm_type=child.norm_type, + scale_grad_by_freq=child.scale_grad_by_freq, + sparse=child.sparse, + # quantization args + group_size=group_size, + scale_precision=scale_precision, + zero_point_precision=zero_point_precision, + device=child.weight.device, + output_dtype=child.weight.dtype, + ) + setattr(module, name, quantized_embedding) + + # Load weights and qparams into quantized embedding + (qmin, qmax) = _get_qmin_qmax(self.bit_width) + (s, zp) = get_group_qparams_symmetric( + child.weight, + self.bit_width, + group_size, + precision=scale_precision, + ) + zp = zp.to(zero_point_precision) + q_weight = _quantized_decomposed_quantize_per_channel_group_wrapper( + child.weight, + s, + zp, + qmin, + qmax, + torch.int8, + group_size, + ) + quantized_embedding.weight = q_weight + quantized_embedding.scale = s.to(scale_precision) + quantized_embedding.zero_point = zp.to(zero_point_precision) + else: + self._convert_helper(child) + + +class Int4WeightOnlyQATEmbedding(FakeQuantizedEmbedding): + """ + This module implements a embedding layer with int4 fake quantized + grouped per channel weights. + + args: + group_size: the number of elements in each quantized group for weights + scale_precision: precision of per group scales + zero_point_precision: precision of per group zero points + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + padding_idx: Optional[int] = None, + max_norm: Optional[float] = None, + norm_type: float = 2.0, + scale_grad_by_freq: bool = False, + sparse: bool = False, + group_size: int = 32, + scale_precision: torch.dtype = torch.float32, + zero_point_precision: torch.dtype = torch.int32, + *args, + **kwargs, + ): + weight_config = FakeQuantizeConfig( + dtype=TorchAODType.INT4, + group_size=group_size, + is_symmetric=True, + is_dynamic=True, + scale_precision=scale_precision, + zero_point_precision=zero_point_precision, + ) + super().__init__( + num_embeddings, + embedding_dim, + padding_idx, + max_norm, + norm_type, + scale_grad_by_freq, + sparse, + weight_config, + *args, + **kwargs, + ) + + def enable_fake_quant(self, enabled: bool = True): + self.weight_fake_quantizer.enabled = enabled + + def disable_fake_quant(self): + self.enable_fake_quant(False) + + +class Int4WeightOnlyEmbedding(torch.nn.Module): + """ + This module implements a embedding layer with int4 quantized + grouped per channel weights. + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + padding_idx: Optional[int] = None, + max_norm: Optional[float] = None, + norm_type: float = 2.0, + scale_grad_by_freq: bool = False, + sparse: bool = False, + group_size: int = 32, + scale_precision: torch.dtype = torch.float32, + zero_point_precision: torch.dtype = torch.int32, + device: torch.device = None, + output_dtype: torch.dtype = torch.float32, + ): + super().__init__() + + # nn.Embedding args + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + self.padding_idx = padding_idx + self.max_norm = max_norm + self.norm_type = norm_type + self.scale_grad_by_freq = scale_grad_by_freq + self.sparse = sparse + + # quantization args + self.bit_width = 4 + self.group_size = group_size + self.scale_precision = scale_precision + self.zero_point_precision = zero_point_precision + self.output_dtype = output_dtype + + # currently storing unpacked int8 weights + self.register_buffer( + "weight", + torch.empty( + (num_embeddings, embedding_dim), dtype=torch.int8, device=device + ), + ) + self.register_buffer( + "scale", + torch.empty( + (num_embeddings, embedding_dim // group_size), + dtype=scale_precision, + device=device, + ), + ) + self.register_buffer( + "zero_point", + torch.empty( + (num_embeddings, embedding_dim // group_size), + dtype=zero_point_precision, + device=device, + ), + ) + + def forward(self, x): + from torchao.quantization.quant_primitives import ( + dequantize_affine, + ) + + qmin, qmax = _get_qmin_qmax(self.bit_width) + + # dequantize_affine casts to output_dtype before scaling + # dequantize_per_channel_group scales and then casts to output_dtype + # The two do not agree when dtype != torch.float32 + w_dq = dequantize_affine( + self.weight, + [1, self.group_size], + self.scale, + self.zero_point, + torch.int8, + qmin, + qmax, + output_dtype=self.output_dtype, + ) + return F.embedding( + x, + w_dq, + self.padding_idx, + self.max_norm, + self.norm_type, + self.scale_grad_by_freq, + self.sparse, + ) diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/fake_quantizer.py b/lib/python3.12/site-packages/torchao/quantization/qat/fake_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..1747f8a4a64c1f441e94c060bfd4c92c6cd7c725 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/qat/fake_quantizer.py @@ -0,0 +1,151 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Optional + +import torch + +from torchao.quantization.granularity import ( + PerAxis, + PerGroup, + PerToken, +) +from torchao.quantization.quant_primitives import ( + _DTYPE_TO_BIT_WIDTH, + _DTYPE_TO_QVALUE_BOUNDS, + MappingType, + choose_qparams_affine, +) +from torchao.quantization.utils import ( + _get_per_token_block_size, + get_group_qparams_symmetric, + get_groupwise_affine_qparams, +) + +from .api import ( + FakeQuantizeConfig, +) +from .utils import ( + _fake_quantize_per_channel_group, + _fake_quantize_per_token, +) + + +class FakeQuantizer(torch.nn.Module): + """ + Generic module for applying fake quantization to a tensor, as specified in the config. + """ + + def __init__(self, config: FakeQuantizeConfig): + super().__init__() + self.config = config + self.enabled = True + self.scale: Optional[torch.Tensor] = None + self.zero_point: Optional[torch.Tensor] = None + + # TODO: support range learinng + if self.config.range_learning: + raise NotImplementedError("Range learning is not supported yet") + + def forward(self, x: torch.Tensor): + """ + Apply fake quantization to the tensor based on the bit-width, + granularity, symmetry, and other properties specified in the config. + """ + if not self.enabled: + return x + + if isinstance(self.config.granularity, PerToken): + return self._per_token_forward(x) + elif isinstance(self.config.granularity, (PerAxis, PerGroup)): + return self._per_channel_or_group_forward(x) + else: + raise ValueError("Unknown granularity '%s'" % self.config.granularity) + + def _per_token_forward(self, x: torch.Tensor): + """ + Perform per token fake quantization on the tensor. + """ + if self.config.is_symmetric: + raise NotImplementedError("Symmetric per token is not supported yet") + + qmin, qmax = _DTYPE_TO_QVALUE_BOUNDS[self.config.dtype] + if self._should_compute_qparams(): + self.scale, self.zero_point = choose_qparams_affine( + x, + mapping_type=MappingType.ASYMMETRIC, + block_size=_get_per_token_block_size(x), + target_dtype=self.config.dtype, + quant_min=qmin, + quant_max=qmax, + scale_dtype=self.config.scale_precision, + zero_point_dtype=self.config.zero_point_precision, + ) + return _fake_quantize_per_token(x, self.scale, self.zero_point, qmin, qmax) + + def _per_channel_or_group_forward(self, x: torch.Tensor): + """ + Perform per channel or per group fake quantization on the tensor. + We express per channel using per group where the group size is the size + of the last dimension of the tensor. + """ + granularity = self.config.granularity + scale_precision = self.config.scale_precision + zero_point_precision = self.config.zero_point_precision + zero_point_domain = self.config.zero_point_domain + is_symmetric = self.config.is_symmetric + + # get group size + if isinstance(granularity, PerAxis): + assert granularity.axis == 0 + group_size = x.size()[-1] + elif isinstance(granularity, PerGroup): + group_size = granularity.group_size + else: + raise ValueError("Unexpected granularity '%s'" % granularity) + + # get scales and zero points + # TODO: refactor this to use `choose_qparams_affine` + if self._should_compute_qparams(): + bit_width = _DTYPE_TO_BIT_WIDTH[self.config.dtype] + if is_symmetric: + (self.scale, self.zero_point) = get_group_qparams_symmetric( + x, + bit_width, + group_size, + scale_precision, + ) + else: + (self.scale, self.zero_point) = get_groupwise_affine_qparams( + x, + bit_width, + group_size, + scale_precision, + ) + self.zero_point = self.zero_point.to(zero_point_precision) + + qmin, qmax = _DTYPE_TO_QVALUE_BOUNDS[self.config.dtype] + return _fake_quantize_per_channel_group( + x, + self.scale, + self.zero_point, + qmin, + qmax, + group_size, + zero_point_domain, + ) + + def _should_compute_qparams(self) -> bool: + """ + Return whether we need to compute new scales and zero points. + """ + return self.config.is_dynamic or self.scale is None or self.zero_point is None + + def __repr__(self) -> str: + """ + Return a human readable representation of this `FakeQuantizer` with config details. + """ + return "FakeQuantizer(%s)" % self.config diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/linear.py b/lib/python3.12/site-packages/torchao/quantization/qat/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..4b0b160c53a698139012f59fada0a69606ed31ad --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/qat/linear.py @@ -0,0 +1,525 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Any, Optional + +import torch +import torch.nn.functional as F + +from torchao.dtypes.utils import is_device +from torchao.quantization.GPTQ import ( + Int8DynActInt4WeightLinear, + WeightOnlyInt4Linear, + _check_linear_int4_k, + _replace_linear_8da4w, + _replace_linear_int4, + groupwise_affine_quantize_tensor, +) +from torchao.quantization.quant_primitives import ( + TorchAODType, + ZeroPointDomain, +) +from torchao.quantization.unified import TwoStepQuantizer +from torchao.quantization.utils import get_group_qparams_symmetric +from torchao.utils import TORCH_VERSION_AT_LEAST_2_6 + +from .api import FakeQuantizeConfig +from .fake_quantizer import FakeQuantizer +from .utils import ( + _get_qmin_qmax, +) + + +class FakeQuantizedLinear(torch.nn.Linear): + """ + General linear layer with fake quantized weights and activations. + + Specific target dtypes, granularity, schemes etc. are specified + through separate configs for weights and activations. + + Example usage:: + + activation_config = FakeQuantizeConfig( + dtype=torch.int8, + granularity="per_token", + is_symmetric=False, + ) + weight_config = FakeQuantizeConfig( + dtype=torch.int4, + group_size=8, + is_symmetric=True, + ) + fq_linear = FakeQuantizedLinear( + 16, 32, False, activation_config, weight_config, + ) + fq_linear(torch.randn(16)) + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + activation_config: Optional[FakeQuantizeConfig] = None, + weight_config: Optional[FakeQuantizeConfig] = None, + *args, + **kwargs, + ) -> None: + super().__init__( + in_features, + out_features, + bias, + *args, + **kwargs, + ) + # initialize activation fake quantizer + if activation_config is not None: + self.activation_fake_quantizer = FakeQuantizer(activation_config) + else: + self.activation_fake_quantizer = None + + # initialize weight fake quantizer + if weight_config is not None: + group_size = weight_config.group_size + if group_size is not None and in_features % group_size != 0: + raise ValueError( + "in_features (%s) %% group_size (%s) must be == 0" + % (in_features, group_size) + ) + self.weight_fake_quantizer = FakeQuantizer(weight_config) + else: + self.weight_fake_quantizer = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.activation_fake_quantizer is not None: + x = self.activation_fake_quantizer(x) + if self.weight_fake_quantizer is not None: + w = self.weight_fake_quantizer(self.weight) + else: + w = self.weight + return F.linear(x, w, self.bias) + + def to_linear(self) -> torch.nn.Linear: + new_linear = torch.nn.Linear( + self.in_features, + self.out_features, + self.bias is not None, + device=self.weight.device, + ) + # In distributed training, the model may be instantiated + # on the meta device, in which case there is no need to + # copy the weights, and doing so will result in an error + if self.weight.device != torch.device("meta"): + new_linear.weight = self.weight + new_linear.bias = self.bias + return new_linear + + @classmethod + def from_linear( + cls, + mod: torch.nn.Linear, + activation_config: Optional[FakeQuantizeConfig] = None, + weight_config: Optional[FakeQuantizeConfig] = None, + ): + new_linear = FakeQuantizedLinear( + mod.in_features, + mod.out_features, + mod.bias is not None, + activation_config=activation_config, + weight_config=weight_config, + device=mod.weight.device, + ) + # In distributed training, the model may be instantiated + # on the meta device, in which case there is no need to + # copy the weights, and doing so will result in an error + if mod.weight.device != torch.device("meta"): + new_linear.weight = mod.weight + new_linear.bias = mod.bias + return new_linear + + +class _LegacyQATQuantizer(TwoStepQuantizer): + """ + Base class for sharing common methods across legacy QAT quantizers. + """ + + def get_activation_fake_quantize_config(self) -> Optional[FakeQuantizeConfig]: + return None + + def get_weight_fake_quantize_config(self) -> Optional[FakeQuantizeConfig]: + return None + + +# ========================================================= +# | Linear int8 dynamic activations + int4 weight QAT | +# ========================================================= + + +class Int8DynActInt4WeightQATQuantizer(_LegacyQATQuantizer): + """ + Quantizer for performing QAT on a model, where linear layers have int8 + dynamic per token fake quantized activations and int4 fake quantized + grouped per channel weights. + """ + + def __init__( + self, + groupsize: int = 256, + padding_allowed: bool = False, + precision: torch.dtype = torch.float32, + scales_precision: torch.dtype = torch.float32, + ) -> None: + super().__init__() + self.groupsize: int = groupsize + self.padding_allowed: bool = padding_allowed + self.precision: torch.dtype = precision + self.scales_precision: torch.dtype = scales_precision + + def prepare( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + _replace_linear_8da4w( + model, + self.groupsize, + self.padding_allowed, + self.precision, + self.scales_precision, + Int8DynActInt4WeightQATLinear, + copy_weights=True, + ) + return model + + def convert( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + self._convert_qat_linear_8da4w(model) + return model + + def _convert_qat_linear_8da4w(self, module: torch.nn.Module): + """ + Replace all `Int8DynActInt4WeightQATLinear` with `Int8DynActInt4WeightLinear`. + """ + for name, child in module.named_children(): + if isinstance(child, Int8DynActInt4WeightQATLinear): + config = child.weight_fake_quantizer.config + quantized_linear = Int8DynActInt4WeightLinear( + child.in_features, + child.out_features, + child.bias is not None, + groupsize=config.group_size, + precision=child.weight.dtype, + scales_precision=config.scale_precision, + ) + setattr(module, name, quantized_linear) + + # Load weights and qparams into quantized linear + n_bit = 4 + (qmin, qmax) = _get_qmin_qmax(n_bit) + (s, zp) = get_group_qparams_symmetric( + child.weight, + n_bit, + config.group_size, + precision=config.scale_precision, + ) + zp = zp.to(config.zero_point_precision) + from torchao._executorch_ops import ( + _quantized_decomposed_quantize_per_channel_group_wrapper, + ) + + q_weight = _quantized_decomposed_quantize_per_channel_group_wrapper( + child.weight, + s, + zp, + qmin, + qmax, + torch.int8, + config.group_size, + ) + quantized_linear.weight = q_weight + quantized_linear.scales = s + quantized_linear.zeros = zp + if child.bias is not None: + quantized_linear.bias = child.bias + else: + self._convert_qat_linear_8da4w(child) + + def get_activation_fake_quantize_config(self) -> Optional[FakeQuantizeConfig]: + return _get_8da4w_activation_config(self.scales_precision) + + def get_weight_fake_quantize_config(self) -> Optional[FakeQuantizeConfig]: + return _get_8da4w_weight_config(self.groupsize, self.scales_precision) + + +class Int8DynActInt4WeightQATLinear(FakeQuantizedLinear): + """ + This module implements a linear layer with int8 dynamic per token fake + quantized activations with int4 fake quantized grouped per channel weights. + + args: + groupsize: the number of elements in each quantized group for weights + precision: precision of weights + scales_precision: precision of per group scales and zero points + + Note: we hardcode activation scales to use torch.fp32, but allow users to specify the weight scales (defaults to torch.fp32). + To get an exact numerical match with Int8DynamicActivationInt4WeightConfig, users must use the same dtype for both the weights + and the scales. Here scales_precision refers specifically to the weight scales only, not the activation scales. + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + device: torch.device = None, + groupsize: int = 256, + precision: torch.dtype = torch.float32, + scales_precision: torch.dtype = torch.float32, + ) -> None: + # Use torch.float32 to match torchao.quantization.quant_api._int8_asymm_per_token_quant, + # which is used in PTQ routines + activation_config = _get_8da4w_activation_config(torch.float32) + weight_config = _get_8da4w_weight_config(groupsize, scales_precision) + super().__init__( + in_features, + out_features, + bias, + activation_config, + weight_config, + device=device, + dtype=precision, + ) + + def enable_fake_quant(self, enabled: bool = True): + self.activation_fake_quantizer.enabled = enabled + self.weight_fake_quantizer.enabled = enabled + + def disable_fake_quant(self): + self.enable_fake_quant(False) + + +def enable_8da4w_fake_quant(mod: torch.nn.Module): + """ + Enable fake quantization for `Int8DynActInt4WeightQATLinear`. + """ + if isinstance(mod, Int8DynActInt4WeightQATLinear): + mod.enable_fake_quant() + + +def disable_8da4w_fake_quant(mod: torch.nn.Module): + """ + Disable fake quantization for `Int8DynActInt4WeightQATLinear`. + """ + if isinstance(mod, Int8DynActInt4WeightQATLinear): + mod.disable_fake_quant() + + +def _get_8da4w_activation_config(qparams_precision: torch.dtype) -> FakeQuantizeConfig: + """ + Return the activation `FakeQuantizeConfig` for `Int8DynActInt4WeightQATQuantizer`. + """ + return FakeQuantizeConfig( + dtype=torch.int8, + granularity="per_token", + is_symmetric=False, + is_dynamic=True, + scale_precision=qparams_precision, + zero_point_precision=qparams_precision, + ) + + +def _get_8da4w_weight_config( + group_size: int, + qparams_precision: torch.dtype, +) -> FakeQuantizeConfig: + """ + Return the weight `FakeQuantizeConfig` for `Int8DynActInt4WeightQATQuantizer`. + """ + return FakeQuantizeConfig( + dtype=TorchAODType.INT4, + group_size=group_size, + is_symmetric=True, + is_dynamic=True, + scale_precision=qparams_precision, + zero_point_precision=qparams_precision, + ) + + +# =================================== +# | Linear int4 weight-only QAT | +# =================================== + + +class Int4WeightOnlyQATQuantizer(_LegacyQATQuantizer): + """ + Quantizer for performing QAT on a model, where linear layers have + int4 fake quantized grouped per channel weights. + """ + + def __init__( + self, + groupsize: int = 256, + inner_k_tiles: Optional[int] = 8, + precision: torch.dtype = torch.bfloat16, + scales_precision: torch.dtype = torch.bfloat16, + ) -> None: + super().__init__() + assert inner_k_tiles in [2, 4, 8] + assert groupsize in [32, 64, 128, 256] + self.inner_k_tiles = inner_k_tiles + self.groupsize = groupsize + self.precision = precision + self.scales_precision = scales_precision + + def prepare( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + _replace_linear_int4( + model, + self.groupsize, + self.inner_k_tiles, + padding_allowed=True, + precision=self.precision, + scales_precision=self.scales_precision, + linear_class=Int4WeightOnlyQATLinear, + copy_weights=True, + ) + return model + + def convert( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + self._convert_qat_linear_4w(model) + return model + + def _convert_qat_linear_4w(self, module: torch.nn.Module): + """ + Replace all `Int4WeightOnlyQATLinear` with `WeightOnlyInt4Linear`. + """ + for name, child in module.named_children(): + if isinstance(child, Int4WeightOnlyQATLinear): + in_features = child.in_features + out_features = child.out_features + inner_k_tiles = child.inner_k_tiles + config = child.weight_fake_quantizer.config + quantized_linear = WeightOnlyInt4Linear( + in_features, + out_features, + bias=False, + groupsize=config.group_size, + inner_k_tiles=inner_k_tiles, + precision=child.weight.dtype, + scales_precision=config.scale_precision, + device=next(child.parameters()).device, + ) + setattr(module, name, quantized_linear) + + # Load weights and qparams into quantized linear + n_bit = 4 + (q_weight, scales_and_zeros) = groupwise_affine_quantize_tensor( + child.weight, + n_bit, + config.group_size, + ) + if ( + is_device(q_weight.device.type, "cpu") + and TORCH_VERSION_AT_LEAST_2_6 + ): + q_weight = torch.ops.aten._convert_weight_to_int4pack_for_cpu( + q_weight.to(child.weight.device), + child.inner_k_tiles, + ) + else: + q_weight = torch.ops.aten._convert_weight_to_int4pack( + q_weight.to(child.weight.device), + child.inner_k_tiles, + ) + quantized_linear.weight = q_weight + quantized_linear.scales_and_zeros = scales_and_zeros + else: + self._convert_qat_linear_4w(child) + + def get_weight_fake_quantize_config(self) -> Optional[FakeQuantizeConfig]: + return _get_4w_weight_config(self.groupsize, self.scales_precision) + + +class Int4WeightOnlyQATLinear(FakeQuantizedLinear): + """ + This module implements a linear layer with int4 fake quantized grouped + per channel weights, with forward numerics matching `WeightOnlyInt4Linear`, + which uses the efficient int4 tinygemm kernel. + + args: + groupsize: the number of elements in each quantized group for weights + precision: precision of weights + scales_precision: precision of per group scales and zero points + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + device: torch.device = None, + groupsize: int = 256, + inner_k_tiles: int = 8, + precision: torch.dtype = torch.bfloat16, + scales_precision: torch.dtype = torch.bfloat16, + ) -> None: + assert scales_precision == torch.bfloat16, "only bf16 is supported for scales" + if not _check_linear_int4_k(in_features, groupsize, inner_k_tiles): + raise ValueError("Padding for QAT 4w is not supported yet") + self.inner_k_tiles = inner_k_tiles + weight_config = _get_4w_weight_config(groupsize, scales_precision) + super().__init__( + in_features, + out_features, + bias, + activation_config=None, + weight_config=weight_config, + device=device, + dtype=precision, + ) + + def enable_fake_quant(self, enabled: bool = True): + self.activation_fake_quantizer.enabled = enabled + self.weight_fake_quantizer.enabled = enabled + + def disable_fake_quant(self): + self.enable_fake_quant(False) + + +def enable_4w_fake_quant(mod: torch.nn.Module): + """ + Enable fake quantization for `Int4WeightOnlyQATLinear`. + """ + if isinstance(mod, Int4WeightOnlyQATLinear): + mod.enable_fake_quant() + + +def disable_4w_fake_quant(mod: torch.nn.Module): + """ + Disable fake quantization for `Int4WeightOnlyQATLinear`. + """ + if isinstance(mod, Int4WeightOnlyQATLinear): + mod.disable_fake_quant() + + +def _get_4w_weight_config( + group_size: int, + qparams_precision: torch.dtype, +) -> FakeQuantizeConfig: + """ + Return the weight `FakeQuantizeConfig` for `Int4WeightOnlyQATQuantizer`. + """ + return FakeQuantizeConfig( + dtype=torch.uint4, + group_size=group_size, + is_symmetric=False, + is_dynamic=True, + scale_precision=qparams_precision, + zero_point_precision=qparams_precision, + zero_point_domain=ZeroPointDomain.FLOAT, + ) diff --git a/lib/python3.12/site-packages/torchao/quantization/qat/utils.py b/lib/python3.12/site-packages/torchao/quantization/qat/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..12e9097adaa9c7937bf4f0078099d3b639aa7854 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/qat/utils.py @@ -0,0 +1,147 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import List + +import torch + +from torchao.quantization.quant_primitives import ( + ZeroPointDomain, + fake_quantize_affine_cachemask, +) +from torchao.quantization.utils import ( + _get_per_token_block_size, +) + + +class _GenericFakeQuantize(torch.autograd.Function): + """ + Implementation of generic fake quantize with backward STE. + + With the appropriate input tensor shape, this can be used to express + grouped per channel fake quantize or per token fake quantize. + """ + + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + input: torch.Tensor, + block_size: List[int], + scales: torch.Tensor, + zero_points: torch.Tensor, + quant_min: int, + quant_max: int, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + ) -> torch.Tensor: + # avoid circular dependencies + from torchao.quantization.qat.affine_fake_quantized_tensor import ( + AffineFakeQuantizedTensor, + ) + + if isinstance(input, AffineFakeQuantizedTensor): + _input = input.original_tensor + else: + _input = input + + (fq, mask) = fake_quantize_affine_cachemask( + _input, + block_size, + scales, + zero_points, + torch.int32, + quant_min, + quant_max, + zero_point_domain, + ) + + ctx.save_for_backward(mask) + return fq + + @staticmethod + def backward(ctx, gy): + (mask,) = ctx.saved_tensors + return gy * mask, None, None, None, None, None, None + + +class _UnwrapAffineFakeQuantizedTensor(torch.autograd.Function): + """ + Helper autograd function to unwrap `AffineFakeQuantizedTensor` while ensuring + gradients are still passed to the tensor subclass. This is used in place of + `_GenericFakeQuantize` when fake quant is disabled. + """ + + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + input: torch.Tensor, + ) -> torch.Tensor: + # avoid circular dependencies + from torchao.quantization.qat.affine_fake_quantized_tensor import ( + AffineFakeQuantizedTensor, + ) + + assert isinstance(input, AffineFakeQuantizedTensor) + return input.original_tensor + + @staticmethod + def backward(ctx, gy): + return (gy,) + + +def _fake_quantize_per_channel_group( + input: torch.Tensor, + scales: torch.Tensor, + zero_points: torch.Tensor, + quant_min: int, + quant_max: int, + group_size: int, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, +) -> torch.Tensor: + assert group_size > 1 + assert input.shape[-1] % group_size == 0 + assert input.dim() == 2 + block_size = (1, group_size) + return _GenericFakeQuantize.apply( + input, + block_size, + scales, + zero_points, + quant_min, + quant_max, + zero_point_domain, + ) + + +def _fake_quantize_per_token( + input: torch.Tensor, + scales: torch.Tensor, + zero_points: torch.Tensor, + quant_min: int, + quant_max: int, +) -> torch.Tensor: + from torch.ao.quantization.fx._decomposed import _per_token_quant_qparam_dim_check + + _per_token_quant_qparam_dim_check(input, scales, zero_points) + block_size = _get_per_token_block_size(input) + fq = _GenericFakeQuantize.apply( + input, + block_size, + scales, + zero_points, + quant_min, + quant_max, + ) + return fq.reshape_as(input).to(input.dtype) + + +def _get_qmin_qmax(n_bit: int, symmetric: bool = True): + if symmetric: + qmin = -(2 ** (n_bit - 1)) + qmax = 2 ** (n_bit - 1) - 1 + else: + qmin = 0 + qmax = 2**n_bit - 1 + return (qmin, qmax) diff --git a/lib/python3.12/site-packages/torchao/quantization/quant_api.py b/lib/python3.12/site-packages/torchao/quantization/quant_api.py new file mode 100644 index 0000000000000000000000000000000000000000..c20c37a19410921feb7e01d406313d7383967f29 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/quant_api.py @@ -0,0 +1,2054 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +Quantization APIs + +Generally these APIs can be applied directly to any model +with Linear modules to obtain quantized linear ops. The intended +usage involves applying torch.compile to the model afterwards +both because primitives were designed based on the fusions that +come along with it and because that is how we access the intended quantized +and mixed GEMM kernels +""" + +import logging +import types +import warnings +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.utils.parametrize as parametrize + +import torchao +from torchao.core.config import AOBaseConfig +from torchao.dtypes import ( + AffineQuantizedTensor, + CutlassInt4PackedLayout, + CutlassSemiSparseLayout, + Float8Layout, + Int4CPULayout, + Int4XPULayout, + MarlinQQQLayout, + MarlinSparseLayout, + PackedLinearInt8DynamicActivationIntxWeightLayout, + PlainLayout, + QDQLayout, + SemiSparseLayout, + TensorCoreTiledLayout, + UintxLayout, + to_affine_quantized_floatx, + to_affine_quantized_floatx_static, + to_affine_quantized_intx, + to_marlinqqq_quantized_intx, +) +from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import ( + Target, + make_packed_linear_int8_dynamic_activation_intx_weight_tensor, +) +from torchao.dtypes.utils import Layout +from torchao.float8.float8_linear import Float8Linear +from torchao.float8.inference import Float8MMConfig +from torchao.quantization.linear_activation_weight_observed_tensor import ( + LinearActivationWeightObservedTensor, +) +from torchao.quantization.observer import AffineQuantizedObserverBase, get_block_size +from torchao.quantization.transform_module import ( + _QUANTIZE_CONFIG_HANDLER, + register_quantize_module_handler, +) +from torchao.quantization.weight_tensor_linear_activation_quantization import ( + to_weight_tensor_with_linear_activation_quantization_metadata, +) +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_4, + TORCH_VERSION_AT_LEAST_2_5, + TORCH_VERSION_AT_LEAST_2_6, + is_MI300, + is_sm_at_least_89, + is_sm_at_least_90, +) + +from .autoquant import AutoQuantizableLinearWeight, autoquant +from .GPTQ import ( + Int4WeightOnlyGPTQQuantizer, + Int4WeightOnlyQuantizer, + Int8DynActInt4WeightGPTQQuantizer, + Int8DynActInt4WeightQuantizer, +) +from .granularity import ( + Granularity, + PerAxis, + PerGroup, + PerRow, + PerTensor, +) +from .linear_activation_quantized_tensor import ( + LinearActivationQuantizedTensor, + to_linear_activation_quantized, +) +from .qat import ( + intx_quantization_aware_training, +) +from .quant_primitives import ( + _DTYPE_TO_QVALUE_BOUNDS, + MappingType, + ZeroPointDomain, +) +from .subclass import ( + Int4WeightOnlyQuantizedLinearWeight, + Int8DynamicallyQuantizedLinearWeight, + Int8WeightOnlyQuantizedLinearWeight, + QuantizedLinearWeightBase, +) +from .unified import Quantizer, TwoStepQuantizer +from .utils import _get_per_token_block_size + +logger = logging.getLogger(__name__) + +__all__ = [ + "swap_conv2d_1x1_to_linear", + "Quantizer", + "TwoStepQuantizer", + "Int4WeightOnlyGPTQQuantizer", + "Int4WeightOnlyQuantizer", + "autoquant", + "_get_subclass_inserter", + "quantize_", + "int8_dynamic_activation_int4_weight", + "int8_dynamic_activation_int8_weight", + "int8_dynamic_activation_int8_semi_sparse_weight", + "int4_weight_only", + "int8_weight_only", + "intx_quantization_aware_training", + "float8_weight_only", + "uintx_weight_only", + "fpx_weight_only", + "gemlite_uintx_weight_only", + "float8_dynamic_activation_float8_weight", + "float8_static_activation_float8_weight", + "Int8DynActInt4WeightQuantizer", + "Int8DynActInt4WeightGPTQQuantizer", + "Float8DynamicActivationFloat8SemiSparseWeightConfig", +] + +LAYOUT_TO_ZERO_POINT_DOMAIN = { + TensorCoreTiledLayout: [ZeroPointDomain.FLOAT], + MarlinSparseLayout: [ZeroPointDomain.INT], + Int4CPULayout: [ZeroPointDomain.FLOAT], + Int4XPULayout: [ZeroPointDomain.FLOAT, ZeroPointDomain.INT], +} + +LAYOUT_TO_PRESERVE_ZEROS = { + TensorCoreTiledLayout: False, + MarlinSparseLayout: True, + Int4CPULayout: False, + Int4XPULayout: False, +} + + +###### +# TO BE DEPRECATED START +###### +def _in_features_greater_than_16(mod, *args): + return hasattr(mod, "in_features") and mod.in_features > 16 + + +def change_linear_weights_to_int8_dqtensors(model, filter_fn=None, **kwargs): + """ + Converts all linear weight tensors to the `Int8DynamicallyQuantizedLinearWeight` + Tensor subclass, effectively applying the same form of quantization + as apply_dynamic_quant while not modifying the linear modules. + """ + if TORCH_VERSION_AT_LEAST_2_4: + raise ImportError( + "This API is deprecated for pytorch 2.4+, please checkout quantization/README.md for most up to date APIs" + ) + + if filter_fn is None: + filter_fn = lambda *args: _is_linear(*args) and _in_features_greater_than_16( + *args + ) + + _replace_with_custom_fn_if_matches_filter( + model, + _get_subclass_inserter( + Int8DynamicallyQuantizedLinearWeight, enable_parametrization=False, **kwargs + ), + filter_fn, + ) + + +def change_linear_weights_to_int8_woqtensors(model, filter_fn=None, **kwargs): + """ + Converts all linear weight tensors to the + `Int8WeightOnlyQuantizedLinearWeight` tensor subclass, + effectively applying the same form of quantization + as apply_weight_only_int8_quant while not modifying the linear modules. + """ + if TORCH_VERSION_AT_LEAST_2_4: + raise ImportError( + "This API is deprecated for pytorch 2.4+, please checkout quantization/README.md for most up to date APIs" + ) + + _replace_with_custom_fn_if_matches_filter( + model, + _get_subclass_inserter( + Int8WeightOnlyQuantizedLinearWeight, enable_parametrization=False, **kwargs + ), + _is_linear if filter_fn is None else filter_fn, + ) + + +def change_linear_weights_to_int4_woqtensors( + model, + groupsize=128, + inner_k_tiles=8, + filter_fn=None, + zero_point_domain=ZeroPointDomain.FLOAT, + preserve_zero=False, +): + """ + Converts all linear weight tensors to the + `Int4WeightOnlyQuantizedLinearWeight` tensor subclass, + effectively applying the same form of quantization + as apply_dynamic_quant while not modifying the linear modules. + Args: + `groupsize`: parameter for quantization, controls the granularity of quantization, smaller + size is more fine grained, choices are [256, 128, 64, 32] + `inner_k_tiles`: parameter for int4 mm kernel, choices are [8, 4, 2] + `filter_fn`: function that takes a nn.Module instance and fully qualified name of the module, \ + returns True if we want to run `config` on + `zero_point_domain`: data type of zeros points, choices are [ZeroPointDomain.FLOAT, \ + ZeroPointDomain.INT, ZeroPointDomain.NONE] + `preserve_zero`: whether to preserve zero, default is False + """ + if TORCH_VERSION_AT_LEAST_2_4: + raise ImportError( + "This API is deprecated for pytorch 2.4+, please checkout quantization/README.md for most up to date APIs" + ) + + if filter_fn is None: + filter_fn = _is_linear + + _replace_with_custom_fn_if_matches_filter( + model, + _get_subclass_inserter( + Int4WeightOnlyQuantizedLinearWeight, + enable_parametrization=False, + groupsize=groupsize, + inner_k_tiles=inner_k_tiles, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + ), + filter_fn, + ) + + +######## +# TO BE DEPRECATED END +######## + + +def _replace_with_custom_fn_if_matches_filter( + model, + replacement_fn, + filter_fn, + cur_fqn="", + device=None, + extra_args: Optional[Tuple[Any, ...]] = (), +) -> None: + """ + Recursively replaces each child module in `model` with the result of `replacement_fn(child)` + if `filter_fn(child)` returns `True`. + + Args: + model (torch.nn.Module): The model containing modules to be replaced. + replacement_fn (Callable[[torch.nn.Module], torch.nn.Module]): The function to replace matching modules. + filter_fn (Callable[[torch.nn.Module], bool]): The filter function to determine which modules to replace. + cur_fqn (str, optional): The current fully qualified name of the module being processed. Defaults to "". + device (device, optional): Device to move the model to before applying `filter_fn`. Defaults to None. + extra_args (Tuple[Any, ...], optional): optional extra args to pass to `replacement_fn`. + + Returns: + None + """ + if isinstance(model, Float8Linear): + with torch.device("meta"): + new_module = nn.Linear(model.in_features, model.out_features) + new_module.weight = model.weight + new_module.bias = model.bias + model = new_module + if filter_fn(model, cur_fqn[:-1]): + if device is not None: + model.to(device=device) # move to device before quantization + model = replacement_fn(model, *extra_args) + return model + else: + named_children_list = list(model.named_children()) + for name, child in named_children_list: + new_child = _replace_with_custom_fn_if_matches_filter( + child, + replacement_fn, + filter_fn, + f"{cur_fqn}{name}.", + device, + extra_args, + ) + if new_child is not child and new_child is not None: + setattr(model, name, new_child) + if device is not None: + model.to(device=device) # move parent module to device + return model + + +def _replace_with_custom_fn_if_matches_filter_with_name( + model, + replacement_fn, + filter_fn, + cur_fqn="", + device=None, + extra_args: Optional[Tuple[Any, ...]] = (), +) -> None: + """ + A variant of _replace_with_custom_fn_if_matches_filter where replacement_fn takes module name as well + ... + replacement_fn (Callable[[torch.nn.Module, str], torch.nn.Module]): The function to replace matching modules. + ... + + Returns: + None + """ + if isinstance(model, Float8Linear): + with torch.device("meta"): + new_module = nn.Linear(model.in_features, model.out_features) + new_module.weight = model.weight + new_module.bias = model.bias + model = new_module + if filter_fn(model, cur_fqn[:-1]): + if device is not None: + model.to(device=device) # move to device before quantization + model = replacement_fn(model, cur_fqn[:-1], *extra_args) + return model + else: + named_children_list = list(model.named_children()) + for name, child in named_children_list: + new_child = _replace_with_custom_fn_if_matches_filter_with_name( + child, + replacement_fn, + filter_fn, + f"{cur_fqn}{name}.", + device, + extra_args, + ) + if new_child is not child: + setattr(model, name, new_child) + if device is not None: + model.to(device=device) # move parent module to device + return model + + +def _is_linear(mod, *args): + # avoid circular dependencies + from torchao.quantization.qat.affine_fake_quantized_tensor import ( + AffineFakeQuantizedTensor, + ) + + # adding weight tensor subclass isinstance check to make sure the weight is only quantized once + # when it is shared by multiple linear modules + return ( + isinstance(mod, torch.nn.Linear) + and hasattr(mod, "weight") + and not isinstance(mod.weight, QuantizedLinearWeightBase) + and not isinstance(mod.weight, AutoQuantizableLinearWeight) + and not isinstance(mod.weight, AffineQuantizedTensor) + and not isinstance(mod.weight, LinearActivationQuantizedTensor) + and not isinstance(mod.weight, AffineFakeQuantizedTensor) + and not isinstance(mod, nn.modules.linear.NonDynamicallyQuantizableLinear) + ) + + +def _get_subclass_inserter(cls, enable_parametrization=False, **kwargs): + """ + Returns a function which inserts the given subclass into all linear modules + in the model. The inserted module will have its weight set to the result of + `cls(mod.weight, **kwargs)`. If parametrization is enabled then this will be done using + torch.nn.utils.parametrize instead of directly setting the attribute on the module. + + Args: + cls (torch.Tensor): The class to insert as a child module. + kwargs (Any): Any additional arguments for the constructor. + """ + constructor = kwargs.pop("constructor", "subclass_constructor") + from_float = kwargs.pop("method", "from_float") + + def insert_subclass(lin): + if enable_parametrization: + lin.weight = torch.nn.Parameter( + cls.from_float(lin.weight, **kwargs), requires_grad=False + ) + _, args = lin.weight.__tensor_flatten__() + parametrize.register_parametrization( + lin, "weight", getattr(cls, constructor)(*args) + ) + else: + lin.weight = torch.nn.Parameter( + # cls.from_float(...) + getattr(cls, from_float)(lin.weight, **kwargs), + requires_grad=False, + ) + return lin + + return insert_subclass + + +def swap_conv2d_1x1_to_linear(model, filter_fn=None): + """ + Changes all conv2d 1x1 modules to equivalent linear modules so that they can then be quantized. + """ + + class PermuteSandwich(torch.nn.Module): + def __init__(self, mod): + super().__init__() + self.mod = mod + + def forward(self, *args): + return self.mod(args[0].permute(0, 2, 3, 1)).permute(-0, 3, 1, 2) + + def replace_conv2d_1x1(conv): + assert conv.kernel_size == (1, 1) + lin = torch.nn.Linear( + conv.in_channels, conv.out_channels, bias=(conv.bias is None) + ) + lin.weight = torch.nn.Parameter(conv.weight.squeeze(-1, -2)) + lin.bias = conv.bias + return PermuteSandwich(lin) + + if filter_fn is None: + filter_fn = lambda mod, *args: isinstance( + mod, torch.nn.Conv2d + ) and mod.kernel_size == (1, 1) + + _replace_with_custom_fn_if_matches_filter( + model, replace_conv2d_1x1, filter_fn=filter_fn + ) + + +def insert_observers_( + model: nn.Module, + input_observer: Optional[AffineQuantizedObserverBase], + weight_observer: Optional[AffineQuantizedObserverBase], + *, + filter_fn: Optional[Callable[[torch.nn.Module, str], bool]] = None, +): + """ + Converts the weight of a linear module to a LinearActivationWeightObservedTensor. + + This function wraps the weight of the given linear module with a LinearActivationWeightObservedTensor, + which enables observation of both input and weight tensors during forward passes. + The wrapped weight is then re-wrapped as a nn.Parameter to maintain compatibility + with PyTorch's module system. + + Example:: + + ``` + import torch + import torch.nn as nn + from torchao.quantization.linear_observer_tensor import insert_observers_ + from torchao.quantization.observer import ( + AffineQuantizedMinMaxObserver, + PerTensor, + MappingType + ) + + # Create observers + input_observer = AffineQuantizedMinMaxObserver( + MappingType.SYMMETRIC, + torch.float8_e4m3fn, + granularity_type=PerTensor(), + eps=torch.finfo(torch.float32).eps, + scale_dtype=torch.float, + zero_point_dtype=torch.int, + zero_point_domain=ZeroPointDomain.NONE, + ) + + # Create a linear module + linear_module = nn.Linear(10, 20) + + # Convert the linear module's weight to an observed tensor + insert_observers_(linear_module, input_observer, weight_observer=None) + + # The linear_module can now be used as usual, with observers calculating statistics + output = linear_module(torch.randn(10, 10)) + + # Get the scale and zero point of the input observer + scale, zero_point = linear_module.weight.input_observer.calculate_qparams() + ``` + + Args: + model (nn.Module): The nn.Module to convert. + input_observer (Optional[AffineQuantizedObserverBase]): Observer for input tensor. + weight_observer (Optional[AffineQuantizedObserverBase]): Observer for weight tensor. + filter_fn (Optional[Callable[[torch.nn.Module, str], bool]]): Filter function to select which modules to convert. + If not provided, all linear modules will be converted. This function should take a module and its fully qualified name. + + Returns: + nn.Linear: The modified linear module with its weight wrapped in a LinearActivationWeightObservedTensor. + """ + + def convert_to_linear_observer(linear_module: nn.Linear): + # Wrap the weight with LinearActivationWeightObservedTensor and then with nn.Parameter + linear_module.weight = nn.Parameter( + LinearActivationWeightObservedTensor.from_float( + linear_module.weight, + input_observer=input_observer, + weight_observer=weight_observer, + ), + requires_grad=linear_module.weight.requires_grad, + ) + return linear_module + + _replace_with_custom_fn_if_matches_filter( + model, + convert_to_linear_observer, + _is_linear if filter_fn is None else filter_fn, + ) + + +def _quantization_type(weight: torch.Tensor): + if isinstance(weight, AffineQuantizedTensor): + return f"{weight.__class__.__name__}({weight._quantization_type()})" + + if isinstance(weight, LinearActivationQuantizedTensor): + return f"{weight.__class__.__name__}(activation={weight.input_quant_func}, weight={_quantization_type(weight.original_weight_tensor)})" + + if type(weight) is torch.Tensor: + return "not quantized" + + return "not recognized" + + +def _linear_extra_repr(self): + return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight={_quantization_type(self.weight)}" + + +def _get_linear_subclass_inserter( + constructor, *, allow_requires_grad=False, propagate_bias=False, **kwargs +): + """Helper function to apply the constructor that quantizes the weight Tensor (with additional kwargs) + to the weight of linear module + """ + + def insert_subclass(lin): + requires_grad = allow_requires_grad and lin.weight.requires_grad + if propagate_bias == True: + kwargs["bias"] = lin.bias + lin.weight = torch.nn.Parameter( + constructor(lin.weight, **kwargs), requires_grad=requires_grad + ) + lin.extra_repr = types.MethodType(_linear_extra_repr, lin) + return lin + + return insert_subclass + + +def quantize_( + model: torch.nn.Module, + config: AOBaseConfig, + filter_fn: Optional[Callable[[torch.nn.Module, str], bool]] = None, + device: Optional[torch.types.Device] = None, +): + """Convert the weight of linear modules in the model with `config`, model is modified inplace + + Args: + model (torch.nn.Module): input model + config (AOBaseConfig): a workflow configuration object. + filter_fn (Optional[Callable[[torch.nn.Module, str], bool]]): function that takes a nn.Module instance and fully qualified name of the module, returns True if we want to run `config` on + the weight of the module + device (device, optional): Device to move module to before applying `filter_fn`. This can be set to `"cuda"` to speed up quantization. The final model will be on the specified `device`. + Defaults to None (do not change device). + + Example:: + + import torch + import torch.nn as nn + from torchao import quantize_ + + # quantize with some predefined `config` method that corresponds to + # optimized execution paths or kernels (e.g. int4 tinygemm kernel) + # also customizable with arguments + # currently options are + # int8_dynamic_activation_int4_weight (for executorch) + # int8_dynamic_activation_int8_weight (optimized with int8 mm op and torch.compile) + # int4_weight_only (optimized with int4 tinygemm kernel and torch.compile) + # int8_weight_only (optimized with int8 mm op and torch.compile + from torchao.quantization.quant_api import int4_weight_only + + m = nn.Sequential(nn.Linear(32, 1024), nn.Linear(1024, 32)) + quantize_(m, int4_weight_only(group_size=32)) + + """ + filter_fn = _is_linear if filter_fn is None else filter_fn + + if isinstance(config, AOPerModuleConfig): + _replace_with_custom_fn_if_matches_filter_with_name( + model, + _ao_per_module_config_handler, + filter_fn, + device=device, + extra_args=(config,), + ) + return + + if isinstance(config, AOBaseConfig): + handler = _QUANTIZE_CONFIG_HANDLER[type(config)] + # for each linear in the model, apply the transform if filtering passes + _replace_with_custom_fn_if_matches_filter( + model, + handler, + filter_fn, + device=device, + extra_args=(config,), + ) + + else: + raise AssertionError( + """Passing a generic Callable to `quantize_` is no longer recommended and will be deprecated at a later release. Please see https://github.com/pytorch/ao/issues/1690 for instructions on how to pass in workflow configuration instead.""" + ) + + +def _int8_asymm_per_token_quant(x: torch.Tensor) -> torch.Tensor: + """This is defined here instead of local function to support serialization""" + mapping_type = MappingType.ASYMMETRIC + target_dtype = torch.int8 + scale_dtype = torch.float32 + zero_point_dtype = torch.int8 + if TORCH_VERSION_AT_LEAST_2_6: + return to_affine_quantized_intx( + x, + mapping_type, + _get_per_token_block_size(x), + target_dtype, + scale_dtype=scale_dtype, + zero_point_dtype=zero_point_dtype, + ) + else: + return to_affine_quantized_intx( + x, mapping_type, _get_per_token_block_size(x), target_dtype + ) + + +def _int8_symm_per_token_quant(x: torch.Tensor) -> torch.Tensor: + mapping_type = MappingType.SYMMETRIC + target_dtype = torch.int8 + eps = 1e-5 + quant_min = -127 + quant_max = 127 + + return to_affine_quantized_intx( + x, + mapping_type, + _get_per_token_block_size(x), + target_dtype, + eps=eps, + quant_min=quant_min, + quant_max=quant_max, + scale_dtype=torch.float32, + ) + + +@dataclass +class Int8DynamicActivationInt4WeightConfig(AOBaseConfig): + """Configuration for applying int8 dynamic per token asymmetric activation quantization and int4 per group weight symmetric quantization to linear + This is used to produce a model for executorch backend, but currently executorch did not + support lowering for the quantized model from this flow yet + + Args: + `group_size`: parameter for quantization, controls the granularity of quantization, smaller + size is more fine grained + `layout`: layout type for quantized weight tensor, only supports `MarlinQQQLayout()` and `CutlassInt4PackedLayout()` for now + `mapping_type`: quantization type for weight, controls the weight quantization is symmetric or asymmetric + `act_mapping_type`: quantization type for activation, controls the activation quantization is symmetric or asymmetric + `set_inductor_config`: if True, adjusts `torchinductor` settings to recommended values. + """ + + group_size: int = 32 + layout: Layout = PlainLayout() + mapping_type: MappingType = MappingType.SYMMETRIC + act_mapping_type: MappingType = MappingType.ASYMMETRIC + set_inductor_config: bool = True + + +# for BC +int8_dynamic_activation_int4_weight = Int8DynamicActivationInt4WeightConfig + + +@register_quantize_module_handler(Int8DynamicActivationInt4WeightConfig) +def _int8_dynamic_activation_int4_weight_transform( + module: torch.nn.Module, config: Int8DynamicActivationInt4WeightConfig +): + group_size = config.group_size + layout = config.layout + mapping_type = config.mapping_type + act_mapping_type = config.act_mapping_type + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + weight = module.weight + + if group_size is None or group_size == -1: + group_size = weight.shape[-1] + if weight.shape[-1] % group_size != 0: + return module + + # weight settings + block_size = (1, group_size) + target_dtype = torch.int8 + quant_min = -8 + quant_max = 7 + + # input settings + if act_mapping_type == MappingType.ASYMMETRIC: + input_quant_func = _int8_asymm_per_token_quant + elif act_mapping_type == MappingType.SYMMETRIC: + if isinstance(layout, MarlinQQQLayout): + input_quant_func = _int8_symm_per_token_quant + elif isinstance(layout, CutlassInt4PackedLayout): + input_quant_func = _int8_symm_cutlass_quant + else: + input_quant_func = _int8_symm_per_token_quant + else: + assert False, f"Unsupported activation mapping type: {act_mapping_type}" + + if isinstance(layout, MarlinQQQLayout): + weight = to_marlinqqq_quantized_intx( + weight, block_size, quant_min, quant_max, _layout=layout + ) + elif isinstance(layout, CutlassInt4PackedLayout): + weight = _int4_symm_cutlass_quant(weight) + else: + weight = to_affine_quantized_intx( + weight, + mapping_type, + block_size, + target_dtype, + quant_min, + quant_max, + _layout=layout, + ) + weight = to_linear_activation_quantized(weight, input_quant_func) + module.weight = torch.nn.Parameter(weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class Int8DynamicActivationIntxWeightConfig(AOBaseConfig): + """ + Configuration for dynamically quantizing activations to torch.int8 and weights to torch.intx, with 1 <= x <= 8. + More specifically, activations are dynamically quantized to 8-bits at a per-token granularity with scales/zeros. + Weights are quantized with scales/zeros in a groupwise or channelwise manner using the number of bits specified by weight_dtype. + + This layout is identical to Int8DynamicActivationInt4WeightConfig when weight_dtype is torch.int4 and other args + are the same. However, this layout is more general and supports other weight dtypes. + + args: + weight_dtype: The dtype to use for weight quantization. Must be torch.intx, where 1 <= x <= 8. + torch.intx with x < 8 requires TORCH_VERSION_AT_LEAST_2_6 + weight_granularity: The granularity to use for weight quantization. Must be PerGroup or PerAxis(axis=0). + weight_mapping_type: The type of mapping to use for the weight quantization. + Must be one of MappingType.ASYMMETRIC or MappingType.SYMMETRIC. MappingType.SYMMETRIC requires ZeroPointDomain.NONE + weight_scale_dtype: The dtype to use for the weight scale. + act_mapping_type: The type of mapping to use for the activation quantization. + Must be one of MappingType.ASYMMETRIC or MappingType.SYMMETRIC. + layout: The layout to use for the packed weight tensor: + - PackedLinearInt8DynamicActivationIntxWeightLayout: this layout is optimized for CPU performance. + - QDQLayout: this layout represents the quantization with Q/DQ quant primitives, and is intended for + export applications like ExecuTorch. + """ + + weight_dtype: torch.dtype = torch.int8 + weight_granularity: Granularity = PerGroup(32) + weight_mapping_type: MappingType = MappingType.SYMMETRIC + # TODO: add weight_scale_dtype to Int8DynamicActivationInt4WeightConfig + weight_scale_dtype: Optional[torch.dtype] = None + act_mapping_type: MappingType = MappingType.ASYMMETRIC + layout: Layout = QDQLayout() + + def __post_init__(self): + assert TORCH_VERSION_AT_LEAST_2_6, ( + "Int8DynamicActivationIntxWeightConfig requires torch 2.6+" + ) + assert self.weight_dtype in [getattr(torch, f"int{b}") for b in range(1, 9)], ( + f"weight_dtype must be torch.intx, where 1 <= x <= 8, but got {self.weight_dtype}" + ) + assert isinstance(self.weight_granularity, (PerAxis, PerGroup)), ( + f"weight_granularity must be PerAxis or PerGroup, but got {self.weight_granularity}" + ) + if isinstance(self.weight_granularity, PerAxis): + assert self.weight_granularity.axis == 0, ( + f"axis must be 0, but got {self.weight_granularity.axis}" + ) + assert self.weight_mapping_type in [ + MappingType.ASYMMETRIC, + MappingType.SYMMETRIC, + ], ( + f"weight_mapping_type must be MappingType.ASYMMETRIC or MappingType.SYMMETRIC, but got {self.weight_mapping_type}" + ) + assert self.act_mapping_type in [ + MappingType.ASYMMETRIC, + MappingType.SYMMETRIC, + ], ( + f"act_mapping_type must be MappingType.ASYMMETRIC or MappingType.SYMMETRIC, but got {self.act_mapping_type}" + ) + assert isinstance( + self.layout, (PackedLinearInt8DynamicActivationIntxWeightLayout, QDQLayout) + ), ( + f"layout must be PackedLinearInt8DynamicActivationIntxWeightLayout or QDQLayout, but got {self.layout}" + ) + + if isinstance(self.layout, PackedLinearInt8DynamicActivationIntxWeightLayout): + if self.layout.target in [Target.AUTO, Target.KLEIDIAI, Target.ATEN]: + if (self.weight_scale_dtype) is None or ( + self.weight_scale_dtype != torch.bfloat16 + ): + logging.warning( + f"When using layout PackedLinearInt8DynamicActivationIntxWeightLayout with target {self.layout.target}, " + f"the weight scale may be cast to bfloat16 by the kernel, but weight_scale_dtype is set to {self.weight_scale_dtype}. " + "Explicitly set weight_scale_dtype to torch.bfloat16 to suppress this warning. " + "If you need weight_scale_dtype = torch.float32, use target=Target.UNIVERSAL instead." + ) + + +@register_quantize_module_handler(Int8DynamicActivationIntxWeightConfig) +def _int8_dynamic_activation_intx_weight_transform( + module: torch.nn.Module, config: Int8DynamicActivationIntxWeightConfig +) -> torch.nn.Module: + weight = module.weight + bias = module.bias + weight_dtype = config.weight_dtype + weight_granularity = config.weight_granularity + weight_mapping_type = config.weight_mapping_type + weight_scale_dtype = config.weight_scale_dtype + act_mapping_type = config.act_mapping_type + layout = config.layout + + assert weight.dim() == 2, f"weight must be 2D, but got {weight.dim()}D" + if isinstance(weight_granularity, PerGroup): + group_size = weight_granularity.group_size + elif isinstance(weight_granularity, PerAxis): + assert weight_granularity.axis == 0, "axis must be 0" + group_size = weight.shape[-1] + else: + raise ValueError( + f"weight_granularity must be PerGroup or PerAxis, got {weight_granularity}" + ) + + quant_min, quant_max = _DTYPE_TO_QVALUE_BOUNDS[weight_dtype] + + # We quantize with QDQLayout, and then construct the packed weight tensor later + weight = to_affine_quantized_intx( + input_float=weight, + mapping_type=weight_mapping_type, + block_size=(1, group_size), + target_dtype=torch.int8, + quant_min=quant_min, + quant_max=quant_max, + scale_dtype=weight_scale_dtype, + zero_point_dtype=torch.int8, + preserve_zero=(weight_mapping_type == MappingType.SYMMETRIC), + zero_point_domain=ZeroPointDomain.INT, + _layout=QDQLayout(), + ) + if isinstance(layout, QDQLayout): + # TODO: _int8_asymm_per_token_quant uses scale_dtype=torch.float64, zero_point_dtype=torch.int64, + # which is not great for export with QDQLayout. It is also not consistent with _int8_symm_per_token_quant, + # which uses scale_dtype=torch.float32, zero_point_dtype=torch.int32. + # Maybe introduce new fp32/int32 versions of _int8_asymm_per_token_quant? + if act_mapping_type == MappingType.ASYMMETRIC: + activation_quant_func = _int8_asymm_per_token_quant + elif act_mapping_type == MappingType.SYMMETRIC: + activation_quant_func = _int8_symm_per_token_quant + else: + assert False, f"Unsupported activation mapping type: {act_mapping_type}" + weight = to_linear_activation_quantized(weight, activation_quant_func) + elif isinstance(layout, PackedLinearInt8DynamicActivationIntxWeightLayout): + # PackedLinearInt8DynamicActivationIntxWeightLayout has dynamic activation quantization + # fused with the kernel and it should not be applied separately + assert act_mapping_type == MappingType.ASYMMETRIC, ( + "PackedLinearInt8DynamicActivationIntxWeightLayout requires act_mapping_type=MappingType.ASYMMETRIC" + ) + data, scale, zero_point = weight.tensor_impl.get_plain() + groups_per_row = weight.shape[-1] // group_size + scale = scale.reshape(-1, groups_per_row) + + assert zero_point is not None + zero_point = zero_point.reshape(-1, groups_per_row) + has_weight_zeros = (zero_point != 0).any() + weight = make_packed_linear_int8_dynamic_activation_intx_weight_tensor( + data, + scale, + zero_point if has_weight_zeros else None, + bias, + weight_dtype, + layout.target, + validate_inputs=False, + ) + # bias is packed with weights if present + bias = None + + module.weight = torch.nn.Parameter(weight, requires_grad=False) + module.bias = bias + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class Int4DynamicActivationInt4WeightConfig(AOBaseConfig): + """Applies int4 dynamic per token symmetric activation quantization and int4 per row weight symmetric quantization to linear + + Args: + `layout`: layout type for quantized weight tensor, only supports `MarlinQQQLayout()` and `CutlassInt4PackedLayout()` for now + `mapping_type`: quantization type for weight, controls the weight quantization is symmetric or asymmetric + `act_mapping_type`: quantization type for activation, controls the activation quantization is symmetric or asymmetric + `set_inductor_config`: if True, adjusts `torchinductor` settings to recommended values. + """ + + layout: Layout = CutlassInt4PackedLayout() + mapping_type: MappingType = MappingType.SYMMETRIC + act_mapping_type: MappingType = MappingType.SYMMETRIC + set_inductor_config: bool = True + + +# for bc +int4_dynamic_activation_int4_weight = Int4DynamicActivationInt4WeightConfig + + +@register_quantize_module_handler(Int4DynamicActivationInt4WeightConfig) +def _int4_dynamic_activation_int4_weight_transform( + module: torch.nn.Module, config: Int4DynamicActivationInt4WeightConfig +) -> torch.nn.Module: + weight = module.weight + layout = config.layout + mapping_type = config.mapping_type + act_mapping_type = config.act_mapping_type + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + if not isinstance(layout, CutlassInt4PackedLayout): + raise NotImplementedError( + f"Only CutlassInt4PackedLayout layout is supported. Received {layout}." + ) + if mapping_type != MappingType.SYMMETRIC: + raise NotImplementedError("Only mapping_type=SYMMETRIC is supported.") + if act_mapping_type != MappingType.SYMMETRIC: + raise NotImplementedError("Only act_mapping_type=SYMMETRIC is supported.") + + weight = _int4_symm_cutlass_quant(weight) + weight = to_linear_activation_quantized( + weight, + _int4_symm_cutlass_quant, + ) + module.weight = torch.nn.Parameter(weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class GemliteUIntXWeightOnlyConfig(AOBaseConfig): + """ + applies weight only 4 or 8 bit integer quantization and utilizes the gemlite triton kernel and its associated weight packing format. + This only works for fp16 models. 8 bit quantization is symmetric, 4 bit quantization is asymmetric. + + Args: + `group_size`: parameter for quantization, controls the granularity of quantization, smaller + size is more fine grained + `bit_width`: bit width of the quantized weight. + `packing_bitwidth`: bit width of the packed weight, should be 8 or 32. Can have performance impacts depending on hardware. + `contiguous`: if set, the weight will be packed as specified. Leaving it as None lets gemlite determine the best choice. + `set_inductor_config`: if True, adjusts `torchinductor` settings to recommended values. + """ + + group_size: Optional[int] = 64 + bit_width: int = 4 + packing_bitwidth: int = 32 + contiguous: Optional[bool] = None + set_inductor_config: bool = True + + +# for BC +gemlite_uintx_weight_only = GemliteUIntXWeightOnlyConfig + + +@register_quantize_module_handler(GemliteUIntXWeightOnlyConfig) +def _gemlite_uintx_weight_only_transform( + module: torch.nn.Module, config: GemliteUIntXWeightOnlyConfig +): + group_size = config.group_size + bit_width = config.bit_width + packing_bitwidth = config.packing_bitwidth + contiguous = config.contiguous + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + weight = module.weight + + from torchao.dtypes.uintx.gemlite_layout import get_gemlite_aqt_kwargs + + use_hqq = True if bit_width == 4 else False + new_weight = to_affine_quantized_intx( + weight, + **get_gemlite_aqt_kwargs( + weight, group_size, bit_width, packing_bitwidth, contiguous, use_hqq + ), + ) + module.weight = torch.nn.Parameter(new_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class Int4WeightOnlyConfig(AOBaseConfig): + """ + Configuration for applying uint4 weight-only asymmetric per-group quantization to linear layers, using + "tensor_core_tiled" layout for speedup with tinygemm kernel + + Note: + This is targeting `tinygemm` int4mm kernel (`torch.ops.aten._weight_int4pack_mm` + and `torch.ops.aten._weight_int4pack_mm_for_cpu`), the main difference + of quantization algorithm compared to the more traditional type of integer quantization is the following: + 1). zero_point is in floating point domain instead of integer domain (`zero_point_domain`=`ZeroPointDomain.FLOAT`) + 2). floating point zero does not have to be exactly representable (`preserve_zero`=False in `choose_qparams_affine`) + please follow the relevant code in `choose_qparams_affine`, `quantize_affine` and `dequantize_affine` + to learn about how the quantization parameters are chosen and how the Tensor is quantized/dequantized for tinygemm + + Args: + `group_size`: parameter for quantization, controls the granularity of quantization, smaller + size is more fine grained, choices are [256, 128, 64, 32] + `layout`: layout type for quantized tensor, default is `TensorCoreTiledLayout(inner_k_tiles=8)` + `use_hqq`: whether to use hqq or default quantization mode, default is False + `zero_point_domain`: data type of zeros points, choices are [ZeroPointDomain.FLOAT, ZeroPointDomain.INT, ZeroPointDomain.NONE] + `set_inductor_config`: if True, adjusts `torchinductor` settings to recommended values. + `preserve_zero`: whether to preserve zero, default is None. Will be set to True if zero_point_domain is ZeroPointDomain.INT + """ + + group_size: int = 128 + layout: Optional[TensorCoreTiledLayout] = TensorCoreTiledLayout(inner_k_tiles=8) + use_hqq: bool = False + zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.NONE + set_inductor_config: bool = True + preserve_zero: Optional[bool] = None + + +# for BC +# TODO maybe change other callsites +int4_weight_only = Int4WeightOnlyConfig + + +def _int4_weight_only_quantize_tensor(weight, config): + # TODO(future PR): perhaps move this logic to a different file, to keep the API + # file clean of implementation details + + # for now, make these local variables to allow the rest of the function + # to be a direct copy-paste + group_size = config.group_size + layout = config.layout + use_hqq = config.use_hqq + zero_point_domain = config.zero_point_domain + + if weight.shape[-1] % group_size != 0: + logger.info( + f"Skipping quantizing weight with int4 weight only quantization because the shape of weight {weight.shape} is not compatible with group_size {group_size}" + ) + return weight + + mapping_type = MappingType.ASYMMETRIC + block_size = tuple([1 for _ in range(weight.dim() - 1)] + [group_size]) + target_dtype = torch.int32 + quant_min = 0 + quant_max = 15 + eps = 1e-6 + zero_point_dtype = ( + weight.dtype if isinstance(layout, Int4CPULayout) else torch.bfloat16 + ) + + # nonlocal zero_point_domain + assert type(layout) in LAYOUT_TO_ZERO_POINT_DOMAIN.keys(), ( + f"Only support layout: {LAYOUT_TO_ZERO_POINT_DOMAIN.keys()}" + ) + if zero_point_domain == ZeroPointDomain.NONE: + # the first value is the default one + zero_point_domain = LAYOUT_TO_ZERO_POINT_DOMAIN[type(layout)][0] + else: + assert zero_point_domain in LAYOUT_TO_ZERO_POINT_DOMAIN[type(layout)], ( + f"Layout only support {LAYOUT_TO_ZERO_POINT_DOMAIN[layout]}" + ) + + if zero_point_domain == ZeroPointDomain.INT and isinstance(layout, Int4XPULayout): + zero_point_dtype = torch.int32 + + preserve_zero = ( + config.preserve_zero + if config.preserve_zero is not None + else LAYOUT_TO_PRESERVE_ZEROS[type(layout)] + ) + # Sparse Marlin only supports symmetric quantization. + # NOTE: If we start having lots of layouts that require different configurations, + # we should consider moving this logic somewhere else. + if isinstance(layout, MarlinSparseLayout): + mapping_type = MappingType.SYMMETRIC + assert group_size == 128 or group_size == weight.shape[-1], ( + f"MarlinSparseLayout only supports 128 group size or per channel quantization, got {group_size}" + ) + + new_weight = to_affine_quantized_intx( + weight, + mapping_type, + block_size, + target_dtype, + quant_min, + quant_max, + eps, + zero_point_dtype=zero_point_dtype, + preserve_zero=preserve_zero, + zero_point_domain=zero_point_domain, + _layout=layout, + use_hqq=use_hqq, + ) + return new_weight + + +@register_quantize_module_handler(Int4WeightOnlyConfig) +def _int4_weight_only_transform( + module: torch.nn.Module, config: Int4WeightOnlyConfig +) -> torch.nn.Module: + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + assert hasattr(module, "weight"), ( + "applying int8 weight only quant requires module to have weight attribute" + + " but {module} does not have one" + ) + new_weight = _int4_weight_only_quantize_tensor(module.weight, config) + module.weight = torch.nn.Parameter(new_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class Int8WeightOnlyConfig(AOBaseConfig): + """ + Configuration for applying int8 weight-only symmetric per-channel quantization to linear layers. + """ + + group_size: Optional[int] = None + set_inductor_config: bool = True + + +# for BC +int8_weight_only = Int8WeightOnlyConfig + + +def _int8_weight_only_quantize_tensor(weight, config): + mapping_type = MappingType.SYMMETRIC + target_dtype = torch.int8 + eps = torch.finfo(torch.float32).eps + zero_point_dtype = torch.int64 + group_size = config.group_size + if group_size is None: + group_size = weight.shape[-1] + block_size = tuple([1 for x in range(weight.dim() - 1)] + [group_size]) + new_weight = to_affine_quantized_intx( + weight, + mapping_type, + block_size, + target_dtype, + eps=eps, + zero_point_dtype=zero_point_dtype, + ) + return new_weight + + +@register_quantize_module_handler(Int8WeightOnlyConfig) +def _int8_weight_only_transform(module: torch.nn.Module, config: Int8WeightOnlyConfig): + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + assert hasattr(module, "weight"), ( + "applying int8 weight only quant requires module to have weight attribute" + + " but {module} does not have one" + ) + new_weight = _int8_weight_only_quantize_tensor(module.weight, config) + module.weight = torch.nn.Parameter(new_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +def _int8_symm_per_token_reduced_range_quant(x: torch.Tensor) -> torch.Tensor: + mapping_type = MappingType.SYMMETRIC + target_dtype = torch.int8 + eps = 1e-5 + quant_min = -127 + quant_max = 127 + return to_affine_quantized_intx( + x, + mapping_type, + _get_per_token_block_size(x), + target_dtype, + eps=eps, + quant_min=quant_min, + quant_max=quant_max, + scale_dtype=torch.float32 if x.dtype == torch.float16 else None, + ) + + +def _int8_symm_per_token_reduced_range_quant_noop_decode( + x: torch.Tensor, +) -> torch.Tensor: + mapping_type = MappingType.SYMMETRIC + target_dtype = torch.int8 + eps = 1e-5 + quant_min = -127 + quant_max = 127 + if x.shape[1] == 1: + return x + else: + return to_affine_quantized_intx( + x, + mapping_type, + _get_per_token_block_size(x), + target_dtype, + eps=eps, + quant_min=quant_min, + quant_max=quant_max, + scale_dtype=torch.float32 if x.dtype == torch.float16 else None, + ) + + +def _int8_symm_cutlass_quant(x: torch.Tensor) -> torch.Tensor: + return to_affine_quantized_intx( + x, + mapping_type=MappingType.SYMMETRIC, + block_size=_get_per_token_block_size(x), + target_dtype=torch.int8, + scale_dtype=torch.float32, + eps=torch.finfo(torch.float32).eps, + zero_point_domain=ZeroPointDomain.NONE, + ) + + +def _int4_symm_cutlass_quant(x: torch.Tensor) -> torch.Tensor: + return to_affine_quantized_intx( + x, + mapping_type=MappingType.SYMMETRIC, + block_size=_get_per_token_block_size(x), + target_dtype=torch.int8, + quant_min=-8, + quant_max=7, + scale_dtype=torch.float32, + eps=torch.finfo(torch.float32).eps, + zero_point_domain=ZeroPointDomain.NONE, + _layout=CutlassInt4PackedLayout(), + ) + + +def _float8_cutlass_quant( + x: torch.Tensor, + target_dtype: torch.dtype, +) -> torch.Tensor: + return to_affine_quantized_floatx( + x, + block_size=_get_per_token_block_size(x), + scale_dtype=torch.float32, + target_dtype=target_dtype, + _layout=Float8Layout(mm_config=None), + ) + + +def _float8_cutlass_quant_sparse( + x: torch.Tensor, + target_dtype: torch.dtype, +) -> (torch.Tensor, torch.Tensor): + return to_affine_quantized_floatx( + x, + block_size=_get_per_token_block_size(x), + scale_dtype=torch.float32, + target_dtype=target_dtype, + _layout=CutlassSemiSparseLayout(), + ) + + +@dataclass +class Int8DynamicActivationInt8WeightConfig(AOBaseConfig): + """ + Configuration for applying int8 dynamic symmetric per-token activation and int8 per-channel weight + quantization to linear layers + """ + + layout: Optional[Layout] = PlainLayout() + act_mapping_type: Optional[MappingType] = MappingType.SYMMETRIC + weight_only_decode: bool = False + set_inductor_config: bool = True + + +# for BC +int8_dynamic_activation_int8_weight = Int8DynamicActivationInt8WeightConfig + + +def _int8_dynamic_activation_int8_weight_quantize_tensor(weight, config): + layout = config.layout + act_mapping_type = config.act_mapping_type + weight_only_decode = config.weight_only_decode + + in_features = weight.shape[-1] + # int8 dynamic quantization only has benefit when in_feature > 16 + if in_features <= 16: + logger.info( + f"Skipping applying int8_dynamic_activation_int8_weight to weight of shape {weight.shape}" + f" because `in_feature` is <= 16: {in_features}" + ) + return weight + + # weight settings + mapping_type = MappingType.SYMMETRIC + weight_zero_point_domain = ZeroPointDomain.NONE + + def get_weight_block_size(x): + return tuple([1 for _ in range(x.dim() - 1)] + [x.shape[-1]]) + + target_dtype = torch.int8 + eps = torch.finfo(torch.float32).eps + zero_point_dtype = torch.int64 + + if weight_only_decode: + input_quant_func = _int8_symm_per_token_reduced_range_quant_noop_decode + else: + # input settings + if act_mapping_type == MappingType.SYMMETRIC: + input_quant_func = _int8_symm_per_token_reduced_range_quant + else: + input_quant_func = _int8_asymm_per_token_quant + + block_size = get_weight_block_size(weight) + new_weight = to_affine_quantized_intx( + weight, + mapping_type, + block_size, + target_dtype, + eps=eps, + zero_point_dtype=zero_point_dtype, + _layout=layout, + zero_point_domain=weight_zero_point_domain, + ) + new_weight = to_linear_activation_quantized(new_weight, input_quant_func) + return new_weight + + +@register_quantize_module_handler(Int8DynamicActivationInt8WeightConfig) +def _int8_dynamic_activation_int8_weight_transform( + module: torch.nn.Module, config: Int8DynamicActivationInt8WeightConfig +) -> torch.nn.Module: + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + assert hasattr(module, "weight"), ( + "applying int8 dynamic activation int8 weight quant requires module to have weight attribute" + + "but {module} does not have one" + ) + new_weight = _int8_dynamic_activation_int8_weight_quantize_tensor( + module.weight, config + ) + module.weight = torch.nn.Parameter(new_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +def int8_dynamic_activation_int8_semi_sparse_weight(): + """ + Applies int8 dnynamic symmetric per-token activation and int8 per-channel weight + quantization + 2:4 sparsity to linear layers. + """ + warnings.warn("""int8_dyanmic_activation_int8_semi_sparse_weight() will be deprecated at a later release. Please use the layout kwarg in int8_dynamic_activation_int8_weight instead. + + from torchao.dtypes import SemiSparseLayout + int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()""") + + return int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()) + + +@dataclass +class Float8WeightOnlyConfig(AOBaseConfig): + """ + Configuration for applying float8 weight-only symmetric per-channel quantization to linear layers. + + Args: + weight_dtype (torch.dtype): The target data type for weight quantization. Default is torch.float8_e4m3fn. + set_inductor_config (bool): if True, adjusts `torchinductor` settings to recommended values. + + Note: + The actual matmul will be computed in original precision of the weight tensor. + """ + + weight_dtype: torch.dtype = torch.float8_e4m3fn + set_inductor_config: bool = True + + +# for BC +float8_weight_only = Float8WeightOnlyConfig + + +def _float8_weight_only_quant_tensor(weight, config): + from torchao.dtypes import to_affine_quantized_floatx + + block_size = tuple([1 for _ in range(weight.dim() - 1)] + [weight.shape[-1]]) + new_weight = to_affine_quantized_floatx( + input_float=weight, + block_size=block_size, + target_dtype=config.weight_dtype, + scale_dtype=None, + _layout=Float8Layout(mm_config=None), + ) + return new_weight + + +@register_quantize_module_handler(Float8WeightOnlyConfig) +def _float8_weight_only_transform( + module: torch.nn.Module, config: Float8WeightOnlyConfig +) -> torch.nn.Module: + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + assert hasattr(module, "weight"), ( + "applying int8 weight only quant requires module to have weight attribute" + + " but {module} does not have one" + ) + new_weight = _float8_weight_only_quant_tensor(module.weight, config) + + module.weight = torch.nn.Parameter(new_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +_fp8_granularities = Union[PerTensor, PerRow] + + +# Validate and process granularity input +def _normalize_granularity( + granularity: Optional[ + Union[_fp8_granularities, Tuple[_fp8_granularities, _fp8_granularities]] + ], +) -> Tuple[_fp8_granularities, _fp8_granularities]: + processed_granularity = None + if granularity is None: + processed_granularity = (PerTensor(), PerTensor()) + elif isinstance(granularity, (PerTensor, PerRow)): + processed_granularity = (granularity, granularity) + elif isinstance(granularity, tuple) and len(granularity) == 2: + if not ( + isinstance(granularity[0], (PerTensor, PerRow)) + and isinstance(granularity[1], (PerTensor, PerRow)) + ): + raise ValueError( + f"Invalid granularity types: {granularity}, only PerTensor or PerRow are supported." + ) + if not isinstance(granularity[0], type(granularity[1])): + raise ValueError( + f"Different granularities for activation and weight are not supported: {granularity}, only PerTensor or PerRow are supported." + ) + processed_granularity = granularity + else: + raise ValueError( + f"Invalid granularity specification: {granularity}, only PerTensor or PerRow are supported." + ) + # Validate granularity with supported Hardware + for _granularity in processed_granularity: + if isinstance(_granularity, PerTensor): + assert is_sm_at_least_89() or is_MI300(), ( + "PerTensor quantization only works for CUDA>=8.9 and MI300+" + ) + elif isinstance(_granularity, PerRow): + assert is_sm_at_least_90() or is_MI300(), ( + "PerRow quantization only works for CUDA>=9.0 and MI300+" + ) + else: + raise ValueError(f"Invalid granularity type: {_granularity}") + + return processed_granularity + + +def _input_activation_quant_func_fp8( + x: torch.Tensor, + activation_granularity: _fp8_granularities, + activation_dtype: torch.dtype, + scale: Optional[torch.Tensor] = None, + zero_point: Optional[torch.Tensor] = None, +): + """This function is used to quantize the input activation tensor for an aqt_float variant. If scale + is not provided it will be dynamically calculate the scales otherwise it will use the provided scale. + """ + assert zero_point is None, ( + "Zero point is not supported for dynamic FP8 quantization" + ) + if isinstance(activation_granularity, PerRow): + assert x.dtype == torch.bfloat16, ( + "PerRow quantization only works for bfloat16 precision input activation" + ) + + block_size = get_block_size(x.shape, activation_granularity) + if scale is None: + activation = to_affine_quantized_floatx( + input_float=x, + block_size=block_size, + target_dtype=activation_dtype, + scale_dtype=torch.float32, + _layout=Float8Layout(mm_config=None), # Config is stored on weight + ) + else: + assert isinstance(activation_granularity, PerTensor), ( + "Static quantization only supports PerTensor granularity" + ) + activation = to_affine_quantized_floatx_static( + input_float=x, + block_size=block_size, + scale=scale, + target_dtype=activation_dtype, + _layout=Float8Layout(mm_config=None), # Config is stored on weight + ) + return activation + + +def _fp8_mm_compat(weight: torch.Tensor) -> bool: + """ + Check if a weight tensor meets float8 quantization requirements. + + Args: + weight (torch.Tensor): The weight tensor to check + + Returns: + bool: True if the tensor can be quantized to float8, False otherwise + """ + assert weight.dim() in [ + 2, + 3, + ], f"float8 quantization only works for 2/3-D tensors, got {weight.dim()}D tensor" + + out_dim, in_dim = weight.shape[-2:] + is_compatible = (in_dim % 16 == 0) and (out_dim % 16 == 0) + + if not is_compatible: + logger.info( + f"Skipping float8 quantization: weight shape {weight.shape} is not compatible with _scaled_mm. " + f"Both input dimension ({in_dim}) and output dimension ({out_dim}) must be multiples of 16. " + ) + + return is_compatible + + +@dataclass +class Float8DynamicActivationFloat8WeightConfig(AOBaseConfig): + """ + Configuration for applying float8 dynamic symmetric quantization to both activations and weights of linear layers. + + Args: + activation_dtype (torch.dtype): The target data type for activation quantization. Default is torch.float8_e4m3fn. + weight_dtype (torch.dtype): The target data type for weight quantization. Default is torch.float8_e4m3fn. + granularity: + The granularity for quantization. Can be either a single granularity (applied to both + activations and weights) or a tuple of two granularities (one for activations, one for weights). + If None, defaults to PerTensor for both. Currently both quantizations need to be the same type. And + only PerTensor and PerRow are supported. + mm_config (Float8MMConfig): Configuration for the matrix multiplication. Default uses fast accumulation. + set_inductor_config (bool): if True, adjusts `torchinductor` settings to recommended values. + + """ + + activation_dtype: torch.dtype = torch.float8_e4m3fn + weight_dtype: torch.dtype = torch.float8_e4m3fn + granularity: Optional[ + Union[_fp8_granularities, Tuple[_fp8_granularities, _fp8_granularities]] + ] = None + mm_config: Optional[Float8MMConfig] = None + set_inductor_config: bool = True + + def __post_init__(self): + if self.mm_config is None: + self.mm_config = Float8MMConfig(use_fast_accum=True) + + +# for bc +float8_dynamic_activation_float8_weight = Float8DynamicActivationFloat8WeightConfig + + +def _float8_dynamic_activation_float8_weight_quantize_tensor(weight, config): + activation_dtype = config.activation_dtype + weight_dtype = config.weight_dtype + granularity = config.granularity + mm_config = config.mm_config + + activation_granularity, weight_granularity = _normalize_granularity(granularity) + + if not _fp8_mm_compat(weight): + # TODO(future PR): this should really throw an exception instead of silently + # not doing what the user asked + return weight + if isinstance(weight_granularity, PerRow): + assert weight.dtype == torch.bfloat16, ( + "PerRow quantization only works for bfloat16 precision input weight" + ) + + block_size = get_block_size(weight.shape[-2:], weight_granularity) + if weight.dim() == 3: + block_size = tuple([1] + list(block_size)) + quantized_weight = to_affine_quantized_floatx( + input_float=weight, + block_size=block_size, + target_dtype=weight_dtype, + scale_dtype=torch.float32, + _layout=Float8Layout(mm_config=mm_config), + ) + + input_quant_func = _input_activation_quant_func_fp8 + input_quant_kwargs = { + "activation_granularity": activation_granularity, + "activation_dtype": activation_dtype, + } + + quantized_weight = to_linear_activation_quantized( + quantized_weight, input_quant_func, quant_kwargs=input_quant_kwargs + ) + return quantized_weight + + +@register_quantize_module_handler(Float8DynamicActivationFloat8WeightConfig) +def _float8_dynamic_activation_float8_weight_transform( + module: torch.nn.Module, config: Float8DynamicActivationFloat8WeightConfig +): + assert is_sm_at_least_89() or is_MI300(), ( + "Float8 dynamic activation quantization is only supported on CUDA>=8.9 and MI300+" + ) + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + assert hasattr(module, "weight"), ( + "applying float8 dynamic activation quant requires module to have weight attribute" + + f"but {module} does not have one" + ) + quantized_weight = _float8_dynamic_activation_float8_weight_quantize_tensor( + module.weight, config + ) + module.weight = torch.nn.Parameter(quantized_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class Float8DynamicActivationFloat8SemiSparseWeightConfig(AOBaseConfig): + """ + Applies float8 dynamic quantization to activations and float8 quantization followed by compression to sparse semi-structured tensor to weights of linear layers. + + Args: + `layout`: layout type for quantized weight tensor, only supports `CutlassSemiSparseLayout` at the moment. + `activation_dtype`: data type for quantized activation tensor. + `weight_dtype`: data type for quantized weight tensor. + """ + + layout: Layout = CutlassSemiSparseLayout() + activation_dtype: torch.dtype = torch.float8_e5m2 + weight_dtype: torch.dtype = torch.float8_e4m3fn + + +@register_quantize_module_handler(Float8DynamicActivationFloat8SemiSparseWeightConfig) +def _float8_dynamic_activation_float8_semi_sparse_weight_transform( + module: torch.nn.Module, config: Float8DynamicActivationFloat8SemiSparseWeightConfig +): + assert is_sm_at_least_90(), "Float8 quantization is only supported on CUDA>=9.0" + + weight = module.weight + weight_dtype = config.weight_dtype + activation_dtype = config.activation_dtype + layout = config.layout + + if not isinstance(layout, CutlassSemiSparseLayout): + raise NotImplementedError( + f"Only CutlassSemiSparseLayout layout is supported. Received {layout}." + ) + + weight = _float8_cutlass_quant_sparse(weight, weight_dtype) + weight = to_linear_activation_quantized( + weight, + _float8_cutlass_quant, + quant_kwargs={"target_dtype": activation_dtype}, + ) + + module.weight = torch.nn.Parameter(weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class Float8StaticActivationFloat8WeightConfig(AOBaseConfig): + """ + Configuration for applying float8 static symmetric quantization to + + Args: + scale (torch.Tensor): The scale tensor for activation quantization. + activation_dtype (torch.dtype): The target data type for activation quantization. Default is torch.float8_e4m + weight_dtype (torch.dtype): The target data type for weight quantization. Default is torch.float8_e4m + mm_config (Float8MMConfig): Configuration for the matrix multiplication. Default uses fast accumulation. + set_inductor_config (bool): if True, adjusts `torchinductor` settings to recommended values. + """ + + scale: torch.Tensor + activation_dtype: torch.dtype = torch.float8_e4m3fn + weight_dtype: torch.dtype = torch.float8_e4m3fn + granularity: Optional[ + Union[_fp8_granularities, Tuple[_fp8_granularities, _fp8_granularities]] + ] = None + mm_config: Optional[Float8MMConfig] = None + set_inductor_config: bool = True + + def __post_init__(self): + if self.mm_config is None: + self.mm_config = Float8MMConfig(use_fast_accum=True) + + +# for bc +float8_static_activation_float8_weight = Float8StaticActivationFloat8WeightConfig + + +@register_quantize_module_handler(Float8StaticActivationFloat8WeightConfig) +def _float8_static_activation_float8_weight_transform( + module: torch.nn.Module, config: Float8StaticActivationFloat8WeightConfig +): + assert is_sm_at_least_89() or is_MI300(), ( + "Float8 static activation quantization is only supported on CUDA 8.9 and above" + ) + + scale = config.scale + activation_dtype = config.activation_dtype + weight_dtype = config.weight_dtype + granularity = config.granularity + mm_config = config.mm_config + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + weight = module.weight + activation_granularity, weight_granularity = _normalize_granularity(granularity) + assert isinstance(activation_granularity, PerTensor), ( + "Static quantization only supports PerTensor granularity" + ) + + if not _fp8_mm_compat(weight): + # TODO(future PR): this should really throw an exception instead of silently + # not doing what the user asked + return module + block_size = get_block_size(weight.shape, weight_granularity) + quantized_weight = to_affine_quantized_floatx( + input_float=weight, + block_size=block_size, + target_dtype=weight_dtype, + scale_dtype=torch.float32, + _layout=Float8Layout(mm_config=mm_config), + ) + + input_quant_func = _input_activation_quant_func_fp8 + input_quant_kwargs = { + "activation_granularity": activation_granularity, + "activation_dtype": activation_dtype, + } + + quantized_weight = to_weight_tensor_with_linear_activation_quantization_metadata( + quantized_weight, + input_quant_func, + scale=scale, + zero_point=None, + quant_kwargs=input_quant_kwargs, + ) + + module.weight = torch.nn.Parameter(quantized_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class UIntXWeightOnlyConfig(AOBaseConfig): + """ + Configuration for applying uintx weight-only asymmetric per-group quantization to linear layers, using uintx quantization where + x is the number of bits specified by `dtype` + + Args: + `dtype`: torch.uint1 to torch.uint7 sub byte dtypes + `group_size`: parameter for quantization, controls the granularity of quantization, smaller + size is more fine grained, defaults to 64 + `pack_dim`: the dimension we use for packing, defaults to -1 + `use_hqq`: whether to use hqq algorithm or the default algorithm to quantize the weight + `set_inductor_config`: if True, adjusts `torchinductor` settings to recommended values. + """ + + dtype: torch.dtype + group_size: int = 64 + pack_dim: int = -1 + use_hqq: bool = False + set_inductor_config: bool = True + + +# for BC +uintx_weight_only = UIntXWeightOnlyConfig + + +@register_quantize_module_handler(UIntXWeightOnlyConfig) +def _uintx_weight_only_transform( + module: torch.nn.Module, config: UIntXWeightOnlyConfig +): + dtype = config.dtype + group_size = config.group_size + pack_dim = config.pack_dim + use_hqq = config.use_hqq + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + weight = module.weight + + from torchao.quantization.quant_primitives import _DTYPE_TO_QVALUE_BOUNDS + + SUPPORTED_DTYPES = { + torch.uint1, + torch.uint2, + torch.uint3, + torch.uint4, + torch.uint5, + torch.uint6, + torch.uint7, + torch.uint8, + } + assert dtype in SUPPORTED_DTYPES, f"Unsupported dtype for hqq: {dtype}" + + mapping_type = MappingType.ASYMMETRIC + block_size = (1, group_size) + + if use_hqq: + if dtype == torch.uint4: + logger.warn( + "Recommended to use `int4_weight_only(group_size, use_hqq=True)` for the best performance" + ) + quant_min, quant_max = _DTYPE_TO_QVALUE_BOUNDS[dtype] + dtype = torch.uint8 + eps = None + zero_point_dtype = None + zero_point_domain = ZeroPointDomain.FLOAT + preserve_zero = False + _layout = PlainLayout() + else: + quant_min, quant_max = None, None + eps = torch.finfo(torch.float32).eps + zero_point_dtype = torch.int32 + zero_point_domain = ZeroPointDomain.INT + preserve_zero = True + _layout = UintxLayout(dtype=dtype, pack_dim=pack_dim) + + new_weight = to_affine_quantized_intx( + weight, + mapping_type, + block_size, + dtype, + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + zero_point_dtype=zero_point_dtype, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + _layout=_layout, + use_hqq=use_hqq, + ) + module.weight = torch.nn.Parameter(new_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class IntxWeightOnlyConfig(AOBaseConfig): + """ + Configuration for quantizing weights to torch.intx, with 1 <= x <= 8. + Weights are quantized with scales/zeros in a groupwise or channelwise + manner using the number of bits specified by weight_dtype. + args: + weight_dtype: The dtype to use for weight quantization. Must be torch.intx, where 1 <= x <= 8. + torch.intx with x < 8 requires TORCH_VERSION_AT_LEAST_2_6 + granularity: The granularity to use for weight quantization. Must be PerGroup or PerAxis(0). + mapping_type: The type of mapping to use for the weight quantization. + Must be one of MappingType.ASYMMETRIC or MappingType.SYMMETRIC. + scale_dtype: The dtype to use for the weight scale. + layout: The layout to use for the packed weight tensor: + - QDQLayout: this layout is designed for export to ExecuTorch.this layout represents the quantization with Q/DQ quant primitives, + and is intended for export applications like ExecuTorch. + """ + + weight_dtype: torch.dtype = torch.int8 + granularity: Granularity = PerAxis(0) + mapping_type: MappingType = MappingType.SYMMETRIC + scale_dtype: Optional[torch.dtype] = None + layout: Layout = QDQLayout() + + def __post_init__(self): + assert TORCH_VERSION_AT_LEAST_2_6, "IntxWeightOnlyConfig requires torch 2.6+" + assert self.weight_dtype in [getattr(torch, f"int{b}") for b in range(1, 9)], ( + f"weight_dtype must be torch.intx, where 1 <= x <= 8, but got {self.weight_dtype}" + ) + assert isinstance(self.granularity, (PerAxis, PerGroup)), ( + f"granularity must be PerAxis or PerGroup, but got {self.granularity}" + ) + if isinstance(self.granularity, PerAxis): + assert self.granularity.axis == 0, ( + f"axis must be 0 with PerAxis, but got {self.granularity.axis}" + ) + assert self.mapping_type in [MappingType.ASYMMETRIC, MappingType.SYMMETRIC], ( + f"mapping_type must be MappingType.ASYMMETRIC or MappingType.SYMMETRIC, but got {self.mapping_type}" + ) + + +@register_quantize_module_handler(IntxWeightOnlyConfig) +def _intx_weight_only_transform( + module: torch.nn.Module, config: IntxWeightOnlyConfig +) -> torch.nn.Module: + weight = module.weight + weight_dtype = config.weight_dtype + granularity = config.granularity + mapping_type = config.mapping_type + scale_dtype = config.scale_dtype + layout = config.layout + + assert weight.dim() == 2, ( + f"IntxWeightOnlyConfig only works for 2-d Tensor, got: {weight.dim()}" + ) + if isinstance(granularity, PerGroup): + group_size = granularity.group_size + elif isinstance(granularity, PerAxis): + assert granularity.axis == 0, ( + f"axis must be 0 with PerAxis, but got {granularity.axis}" + ) + group_size = weight.shape[-1] + else: + raise ValueError(f"granularity must be PerGroup or PerAxis, got {granularity}") + + quant_min, quant_max = _DTYPE_TO_QVALUE_BOUNDS[weight_dtype] + weight = to_affine_quantized_intx( + input_float=weight, + mapping_type=mapping_type, + block_size=(1, group_size), + target_dtype=torch.int8, + quant_min=quant_min, + quant_max=quant_max, + scale_dtype=scale_dtype, + zero_point_dtype=torch.int8, + preserve_zero=(mapping_type == MappingType.SYMMETRIC), + zero_point_domain=ZeroPointDomain.INT, + _layout=layout, + ) + module.weight = torch.nn.Parameter(weight, requires_grad=False) + return module + + +@dataclass +class FPXWeightOnlyConfig(AOBaseConfig): + """Sub-byte floating point dtypes defined by `ebits`: exponent bits and `mbits`: mantissa bits + e.g. fp6_e3_m2, fp6_e2_m3, ... + The packing format and kernels are from the fp6-llm paper: https://arxiv.org/abs/2401.14112 + github repo: https://github.com/usyd-fsalab/fp6_llm, now renamed to quant-llm + For more details for packing please see: :class:`~torchao.dtypes.fpx.FpxTensorCoreAQTTensorImpl` + + This is experimental, will be merged with `to_affine_quantized_floatx` + in the future + """ + + ebits: int + mbits: int + set_inductor_config: bool = True + + +# for BC +fpx_weight_only = FPXWeightOnlyConfig + + +@register_quantize_module_handler(FPXWeightOnlyConfig) +def _fpx_weight_only_transform( + module: torch.nn.Module, config: FPXWeightOnlyConfig +) -> torch.nn.Module: + ebits = config.ebits + mbits = config.mbits + weight = module.weight + if config.set_inductor_config: + torchao.quantization.utils.recommended_inductor_config_setter() + + from torchao.dtypes import to_affine_quantized_fpx + from torchao.dtypes.floatx import FloatxTensorCoreLayout + + assert weight.dim() == 2, f"floatx only works for 2-d Tensor, got: {weight.dim()}" + out_dim, in_dim = weight.shape + if (in_dim % 64 != 0) or (out_dim % 256 != 0): + logger.info( + f"Skipping floatx quantization float{ebits + mbits + 1}_{ebits}_{mbits} because " + f"the shape is not compatible with the kernel: in_dim={in_dim}, out_dim={out_dim} " + "expected in_dim % 64 == 0 and out_dim % 256 == 0" + ) + return module + + _layout = FloatxTensorCoreLayout(ebits, mbits) + new_weight = to_affine_quantized_fpx(weight, _layout) + module.weight = torch.nn.Parameter(new_weight, requires_grad=False) + module.extra_repr = types.MethodType(_linear_extra_repr, module) + return module + + +@dataclass +class AOPerModuleConfig(AOBaseConfig): + """Per module configurations for torchao quantize_ API + + Args: + `module_fqn_to_config`: Dict[str, Optional[AOBaseConfig]]: a dictionary from + the fully qualified name of module to the AOBaseConfig that we want to apply to the module. + Also has a special key: "_default", if "_default" is present in the dictionary, + the config for "_default" will be applied to all the remaining modules that does not have + per module configuration specified. + """ + + module_fqn_to_config: Dict[str, Optional[AOBaseConfig]] = field( + default_factory=dict + ) + + +def _ao_per_module_config_handler( + module: torch.nn.Module, module_fqn: str, config: AOPerModuleConfig +): + c = None + if module_fqn in config.module_fqn_to_config: + # Maybe: we can add module type specific config in the future, in needed + c = config.module_fqn_to_config[module_fqn] + else: + # fallback to use default if no module specific config is provided + c = config.module_fqn_to_config.get("_default", None) + + if c is not None: + handler = _QUANTIZE_CONFIG_HANDLER[type(c)] + return handler(module, c) + + return module + + +if TORCH_VERSION_AT_LEAST_2_5: + torch.serialization.add_safe_globals( + [ + _int8_asymm_per_token_quant, + _int8_symm_per_token_reduced_range_quant, + _input_activation_quant_func_fp8, + _int4_symm_cutlass_quant, + _int8_symm_cutlass_quant, + _float8_cutlass_quant, + _float8_cutlass_quant_sparse, + Target, + ] + ) diff --git a/lib/python3.12/site-packages/torchao/quantization/quant_primitives.py b/lib/python3.12/site-packages/torchao/quantization/quant_primitives.py new file mode 100644 index 0000000000000000000000000000000000000000..d13ac330a0b3c0bfc8c00caa0d52acdec7be674c --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/quant_primitives.py @@ -0,0 +1,1592 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import math +from enum import Enum, auto +from typing import Callable, Dict, List, Optional, Tuple, Union + +import torch + +from torchao.prototype.custom_fp_utils import ( + _f32_to_floatx_unpacked, + _floatx_unpacked_to_f32, + _n_ones, +) +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_3, + TORCH_VERSION_AT_LEAST_2_5, + TORCH_VERSION_AT_LEAST_2_6, + _is_float8_type, + _register_custom_op, +) + +__all__ = [ + "choose_qparams_affine", + "choose_qparams_affine_with_min_max", + "choose_qparams_affine_floatx", + "quantize_affine", + "dequantize_affine", + "quantize_affine_floatx", + "dequantize_affine_floatx", + "fake_quantize_affine", + "fake_quantize_affine_cachemask", + "choose_qparams_and_quantize_affine_hqq", + "choose_qparams_and_quantize_affine_qqq", + "dequantize_affine_qqq", + "MappingType", + "ZeroPointDomain", + "TorchAODType", + "choose_qparams_affine_float8", + "quantize_affine_float8", + "dequantize_affine_float8", + "choose_qparams_gguf", + "quantize_gguf", + "dequantize_gguf", +] + + +class MappingType(Enum): + """How floating point number is mapped to integer number + + symmetric mapping means floating point range is symmetrically mapped to integer range + let's say we have floating point range (-3.5, 10.2) and integer range (-8, 7) (int4) + we'll use (-10.2, 10.2) as the range for floating point and map that to (-8, 7) + e.g. scale = (10.2 - (-10.2)) / (7 - (-8)) + + SYMMETRIC_NO_CLIPPING_ERR is a variant of symmetric mapping, where the scale is the max of smin + and smax, where smin = min_val_neg / quant_min, and smax = max_val_pos / quant_max. By calculating + smin and smax individually, there can be less round error on negative values, and no out-of-range + of all floating point values. + + asymmetric mapping means we just directly map the floating point range to integer range, + for the above example, we will map (-3.5, 10.2) to (-8, 7) and calculate quantization parameter + based on this mapping + e.g. scale = (10.2 - (-3.5)) / (7 - (-8)) + """ + + SYMMETRIC = auto() + SYMMETRIC_NO_CLIPPING_ERR = auto() + ASYMMETRIC = auto() + + +class ZeroPointDomain(Enum): + """Enum that indicate whether zero_point is in integer domain or floating point domain + + integer domain: quantized_val = (float_val / scale) (integer) + zero_point (integer) + float domain: quantized_val = (float_val - (zero_point (float) - scale * mid_point)) / scale + none domain: quantized_val = (float_val / scale) + """ + + INT = auto() + FLOAT = auto() + NONE = auto() + + +class TorchAODType(Enum): + """ + Placeholder for dtypes that do not exist in PyTorch core yet. + """ + + # torch.int1 to torch.int7 will be added to PyTorch 2.6 + # These will remain here for BC with older PyTorch versions + INT1 = auto() + INT2 = auto() + INT3 = auto() + INT4 = auto() + INT5 = auto() + INT6 = auto() + INT7 = auto() + + +if TORCH_VERSION_AT_LEAST_2_5: + torch.serialization.add_safe_globals([MappingType, ZeroPointDomain]) + +FP8_TYPES = { + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float8_e4m3fnuz, + torch.float8_e5m2fnuz, +} + +""" +Map from dtype to the bound value of integers +TODO: maybe can replace this with call to torch.iinfo +""" +_DTYPE_TO_QVALUE_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = { + torch.uint8: (0, 255), + torch.int8: (-128, 127), + torch.int16: (-(2**15), 2**15 - 1), + torch.int32: (-(2**31), 2**31 - 1), +} +_DTYPE_TO_BIT_WIDTH: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = { + TorchAODType.INT1: 1, + TorchAODType.INT2: 2, + TorchAODType.INT3: 3, + TorchAODType.INT4: 4, + TorchAODType.INT5: 5, + TorchAODType.INT6: 6, + TorchAODType.INT7: 7, + torch.uint8: 8, + torch.int8: 8, + torch.int16: 16, + torch.int32: 32, +} + +_SUB_BYTE_UINT_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = {} +_SUB_BYTE_INT_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = { + TorchAODType.INT1: (-(2**0), 2**0 - 1), + TorchAODType.INT2: (-(2**1), 2**1 - 1), + TorchAODType.INT3: (-(2**2), 2**2 - 1), + TorchAODType.INT4: (-(2**3), 2**3 - 1), + TorchAODType.INT5: (-(2**4), 2**4 - 1), + TorchAODType.INT6: (-(2**5), 2**5 - 1), + TorchAODType.INT7: (-(2**6), 2**6 - 1), +} + +# torch.uintX available only in PyTorch 2.3+ +if TORCH_VERSION_AT_LEAST_2_3: + _SUB_BYTE_UINT_BOUNDS = { + torch.uint1: (0, 2**1 - 1), + torch.uint2: (0, 2**2 - 1), + torch.uint3: (0, 2**3 - 1), + torch.uint4: (0, 2**4 - 1), + torch.uint5: (0, 2**5 - 1), + torch.uint6: (0, 2**6 - 1), + torch.uint7: (0, 2**7 - 1), + } + _DTYPE_TO_BIT_WIDTH.update( + { + torch.uint1: 1, + torch.uint2: 2, + torch.uint3: 3, + torch.uint4: 4, + torch.uint5: 5, + torch.uint6: 6, + torch.uint7: 7, + } + ) + +# torch.intX available only in PyTorch 2.6+ +if TORCH_VERSION_AT_LEAST_2_6: + _SUB_BYTE_INT_BOUNDS.update( + { + torch.int1: (-(2**0), 2**0 - 1), + torch.int2: (-(2**1), 2**1 - 1), + torch.int3: (-(2**2), 2**2 - 1), + torch.int4: (-(2**3), 2**3 - 1), + torch.int5: (-(2**4), 2**4 - 1), + torch.int6: (-(2**5), 2**5 - 1), + torch.int7: (-(2**6), 2**6 - 1), + } + ) + _DTYPE_TO_BIT_WIDTH.update( + { + torch.int1: 1, + torch.int2: 2, + torch.int3: 3, + torch.int4: 4, + torch.int5: 5, + torch.int6: 6, + torch.int7: 7, + } + ) + +_DTYPE_TO_QVALUE_BOUNDS.update(_SUB_BYTE_UINT_BOUNDS) +_DTYPE_TO_QVALUE_BOUNDS.update(_SUB_BYTE_INT_BOUNDS) +assert _DTYPE_TO_BIT_WIDTH.keys() == _DTYPE_TO_QVALUE_BOUNDS.keys() + +_GGUF_QK_K = 256 + +_ONES_TABLE = [_n_ones(i) for i in range(8)] + +quant_lib = torch.library.Library("torchao", "FRAGMENT") + +register_custom_op = _register_custom_op(quant_lib) + + +# TODO: decide on if we want to allow custom quant_min/quant_max here +def _get_and_check_qmin_qmax(dtype, quant_min, quant_max): + """Get quant_min and quant_max args based on dtype and also + verify that they are within the range of possible quant_min/quant_max + for dtype + """ + if dtype in FP8_TYPES: + quant_min_lower_bound, quant_max_upper_bound = ( + torch.finfo(dtype).min, + torch.finfo(dtype).max, + ) + elif dtype not in _DTYPE_TO_QVALUE_BOUNDS: + raise ValueError(f"Unsupported dtype: {dtype}") + else: + quant_min_lower_bound, quant_max_upper_bound = _DTYPE_TO_QVALUE_BOUNDS[dtype] + if quant_min is None: + quant_min = quant_min_lower_bound + if quant_max is None: + quant_max = quant_max_upper_bound + + assert quant_min >= quant_min_lower_bound, ( + "quant_min out of bound for dtype, " + f"quant_min_lower_bound: {quant_min_lower_bound} quant_min: {quant_min}" + ) + + assert quant_max <= quant_max_upper_bound, ( + "quant_max out of bound for dtype, " + f"quant_max_upper_bound: {quant_max_upper_bound} quant_max: {quant_max}" + ) + return quant_min, quant_max + + +def _get_reduction_params(block_size, input_size): + """Given block_size and input size find the parameters for reduction: + + Output: + shape_for_reduction: the shape we use to `view` input to prepare it for reduction + reduction_dims: the dims we'll do reduction over + + Example:: + Input: + block_size: (3, 3, 2, 10) + input_size: (3, 3, 10, 10) + + Output: + shape_for_reduction: (3, 3, 5, 2, 10) + reduction_dim: [0, 1, 3, 4] + """ + assert len(block_size) == len(input_size) + shape_for_reduction = [] + reduction_dims = [] + cur_dim = 0 + for i in range(len(block_size)): + if block_size[i] != input_size[i] and block_size[i] > 1: + assert input_size[i] % block_size[i] == 0, ( + f"Expecting input size at {i} dimension: {input_size[i]} to be divisible by block_size at {i} dimension: {block_size[i]}" + ) + shape_for_reduction.append(input_size[i] // block_size[i]) + shape_for_reduction.append(block_size[i]) + # reduce over the block_size[i] dim + reduction_dims.append(cur_dim + 1) + cur_dim += 2 + else: + # block_size[i] == input_size[i] or block_size[i] == 1 + shape_for_reduction.append(input_size[i]) + # we only need to reduce over the dimension if block_size is greater than 1 + # otherwise it's already the same as reduced dimension + if block_size[i] != 1: + reduction_dims.append(cur_dim) + cur_dim += 1 + return shape_for_reduction, reduction_dims + + +@torch.no_grad() +def quantize_affine( + input: torch.Tensor, + block_size: Tuple[int, ...], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + output_dtype: torch.dtype, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, +) -> torch.Tensor: + """ + Args: + input (torch.Tensor): original float32, float16 or bfloat16 Tensor + block_size: (Tuple[int, ...]): granularity of quantization, this means the size of the tensor elements that's sharing the same qparam + e.g. when size is the same as the input tensor dimension, we are using per tensor quantization + scale (float): quantization parameter for affine quantization + zero_point (int): quantization parameter for affine quantization + output_dtype (torch.dtype): requested dtype (e.g. torch.uint8) for output Tensor + quant_min (Optional[int]): minimum quantized value for output Tensor, if not specified, it will be derived from dtype + quant_max (Optional[int]): maximum quantized value for output Tensor, if not specified, it will be derived from dtype + zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be either integer or float + if zero_point is in integer domain, zero point is added to the quantized integer value during + quantization + if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) + value during quantization + default is ZeroPointDomain.INT + + Note: + How can block_size represent different granularities? + let's say we have a Tensor of size: (3, 3, 10, 10), here is the table showing how block_size represents different + granularities: + + granularity type | block_size + per_tensor | (3, 3, 10, 10) + per_axis (axis=0) | (1, 3, 10, 10) + per_axis (axis=1) | (3, 1, 10, 10) + per_group (groupsize=2) | (3, 3, 10, 2) + per_group (groupsize=2) for axis = 3 | (3, 3, 2, 10) + + + Output: + quantized tensor with requested dtype + """ + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + elif zero_point_domain is ZeroPointDomain.NONE and zero_point is not None: + raise ValueError("zero_point should be None when zero_point_domain is NONE") + return _quantize_affine( + input, + block_size, + scale, + zero_point, + output_dtype, + quant_min, + quant_max, + zero_point_domain.name, + ) + + +@register_custom_op +def _quantize_affine( + input: torch.Tensor, + block_size: List[int], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + output_dtype: torch.dtype, + quant_min: Optional[Union[int, float, bool]] = None, + quant_max: Optional[Union[int, float, bool]] = None, + zero_point_domain: str = ZeroPointDomain.INT.name, +) -> torch.Tensor: + """op definition that has compatible signatures with custom op library + + Note: + zero_point_domain is optional specifies how we quantize the floating point to quantized data: + INT: quantized_val = (float_val / scale) (integer) + zero_point (integer) + FLOAT: quantized_val = (float_val - (zero_point (float) - scale * mid_point)) / scale + None: quantized_val = (float_val / scale) | this is primarily used for floatx quantization + Where we do not want to round values to nearest integer and instead scale and cast. + """ + quant_min, quant_max = _get_and_check_qmin_qmax(output_dtype, quant_min, quant_max) + # workaround for uintx dtypes, since we don't have native Uintx dtype connected with + # torch.uintx dtypes yet + if output_dtype in _SUB_BYTE_UINT_BOUNDS: + output_dtype = torch.uint8 + return _quantize_affine_no_dtype_cast( + input, + block_size, + scale, + zero_point, + quant_min, + quant_max, + output_dtype, + zero_point_domain, + ).to(output_dtype) + + +def _quantize_affine_no_dtype_cast( + input: torch.Tensor, + block_size: List[int], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_min: Union[int, float], + quant_max: Union[int, float], + quant_dtype: torch.dtype, + zero_point_domain: str = ZeroPointDomain.INT.name, +) -> torch.Tensor: + """ + The op does the following: + 1. figure out the dimension for reduction based on block_size, also reshape the input to align with + the shape after reduction + 2. quantize the input based on the quantization parameters scale and zero_point and args like zero_point_domain + 3. reshape the quantized result to origianl shape + """ + # TODO: validations + # TODO: validate scale/zero_point dimensions are compatible with block_size + assert input.dtype in [ + torch.float32, + torch.float16, + torch.bfloat16, + ], f"Unsupported input dtype: {input.dtype}" + assert len(block_size) == input.dim(), ( + f"Got input dim:{input.dim()}, block_size: {block_size}" + ) + shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + original_shape = input.shape + input = input.view(shape_for_reduction) + shape_after_reduction = shape_for_reduction + for i in reduction_dims: + shape_after_reduction[i] = 1 + scale = scale.view(shape_after_reduction) + + if zero_point is not None and zero_point.numel() > 0: + zero_point = zero_point.view(shape_after_reduction) + else: + # in some cases zero_point being a non-value shows as a tensor + # with numel=0 which we handle by unifying the two + zero_point = None + + if zero_point_domain == ZeroPointDomain.INT.name: + quant = torch.clamp( + torch.round(input * (1.0 / scale)) + zero_point, quant_min, quant_max + ) + elif zero_point_domain == ZeroPointDomain.NONE.name: + assert zero_point is None, ( + "zero_point should be None when zero_point_domain is NONE" + ) + if _is_float8_type(quant_dtype): + quant = torch.clamp(input * scale.reciprocal(), quant_min, quant_max) + else: + quant = torch.clamp( + torch.round(input * (1.0 / scale)), quant_min, quant_max + ) + else: + assert zero_point_domain == ZeroPointDomain.FLOAT.name + mid_point = (quant_max + quant_min + 1) / 2 + min_val = zero_point - scale * mid_point + quant = torch.clamp( + torch.round((input - min_val) / scale), quant_min, quant_max + ) + quant = quant.view(original_shape) + + return quant + + +def dequantize_affine( + input: torch.Tensor, + block_size: Tuple[int, ...], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + input_dtype: torch.dtype, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, + *, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Args: + input (torch.Tensor): quantized tensor, should match the dtype `dtype` argument + block_size: (List[int]): granularity of quantization, this means the size of the tensor elements that's sharing the same qparam + e.g. when size is the same as the input tensor dimension, we are using per tensor quantization + scale (Tensor): quantization parameter for affine quantization + zero_point (Tensor): quantization parameter for affine quantization + input_dtype (torch.dtype): requested dtype (e.g. torch.uint8) for output Tensor + quant_min (Optional[int]): minimum quantized value for input Tensor + quant_max (Optional[int]): maximum quantized value for input Tensor + output_dtype (torch.dtype): dtype for output Tensor, default is fp32 + zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be either integer or float + if zero_point is in integer domain, zero point is added to the quantized integer value during + quantization + if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) + value during quantization + default is ZeroPointDomain.INT + + Output: + dequantized Tensor, with requested dtype or fp32 + """ + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + elif zero_point_domain is ZeroPointDomain.NONE and zero_point is not None: + raise ValueError("zero_point should be None when zero_point_domain is NONE") + return _dequantize_affine( + input, + block_size, + scale, + zero_point, + input_dtype, + quant_min, + quant_max, + zero_point_domain.name, + output_dtype=output_dtype, + ) + + +@register_custom_op +def _dequantize_affine( + input: torch.Tensor, + block_size: List[int], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + input_dtype: torch.dtype, + quant_min: Optional[Union[int, float, bool]] = None, + quant_max: Optional[Union[int, float, bool]] = None, + zero_point_domain: Optional[str] = ZeroPointDomain.INT.name, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """op definition that has compatible signatures with custom op library""" + # TODO: validate scale/zero_point dimensions are compatible with block_size + if input_dtype not in _SUB_BYTE_UINT_BOUNDS: + assert input.dtype == input_dtype, ( + f"Expected: {input_dtype}, got: {input.dtype}" + ) + assert output_dtype in [ + torch.float32, + torch.float16, + torch.bfloat16, + ], f"Unsupported output dtype: {output_dtype}" + quant_min, quant_max = _get_and_check_qmin_qmax(input_dtype, quant_min, quant_max) + return _dequantize_affine_no_dtype_check( + input, + block_size, + scale, + zero_point, + quant_min, + quant_max, + zero_point_domain, + output_dtype, + ) + + +def _dequantize_affine_no_dtype_check( + input: torch.Tensor, + block_size: List[int], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_min: Union[int, float], + quant_max: Union[int, float], + zero_point_domain: Optional[str] = ZeroPointDomain.INT.name, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """This function converts AQT tensors to their high precision floating point representation + + The op does the following: + 1. figure out the dimension for reduction based on block_size, also reshape the input to align with + the shape after reduction + 2. dequantize the input based on the quantization parameters scale and zero_point and args like zero_point_domain + 3. reshape the quantized result to origianl shape and change dtype to the output_dtype + """ + assert len(block_size) == input.dim(), ( + f"Got input dim:{input.dim()}, block_size: {block_size}" + ) + shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + original_shape = input.shape + input = input.view(shape_for_reduction) + shape_after_reduction = shape_for_reduction + for i in reduction_dims: + shape_after_reduction[i] = 1 + scale = scale.view(shape_after_reduction) + + if zero_point is not None: + zero_point = zero_point.view(shape_after_reduction) + + if zero_point_domain == ZeroPointDomain.INT.name: + # Force a copy to avoid input modification due + # to upcoming in-place operations. + dequant = input.to(torch.int32, copy=True) + if zero_point is not None: + dequant = dequant - zero_point.to(torch.int32) + dequant = dequant.to(output_dtype) + dequant = dequant * scale + elif zero_point_domain == ZeroPointDomain.NONE.name: + assert zero_point is None, ( + "zero_point should be None when zero_point_domain is NONE" + ) + dequant = input.to(output_dtype) + dequant = dequant * scale + else: + assert zero_point_domain == ZeroPointDomain.FLOAT.name, ( + f"Unexpected zero point domain: {zero_point_domain}" + ) + # TODO: this seems to be a detail for tinygemm (converting from uint to int, probably need to refactor this) + mid_point = (quant_max + quant_min + 1) / 2 + # This should allocate new memory and avoid input modification + dequant = input - mid_point + dequant = dequant.to(output_dtype) + dequant *= scale + if zero_point is not None: + dequant += zero_point + + return dequant.view(original_shape).to(output_dtype) + + +def fake_quantize_affine( + input: torch.Tensor, + block_size: Tuple[int, ...], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_dtype: torch.dtype, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, +) -> torch.Tensor: + """ + General fake quantize op for quantization-aware training (QAT). + This is equivalent to calling `quantize_affine` + `dequantize_affine` + but without the dtype casts. + + Args: + input (torch.Tensor): original float32, float16 or bfloat16 Tensor + block_size: (Tuple[int, ...]): granularity of quantization, this means the size of the tensor elements that's sharing the same qparam + e.g. when size is the same as the input tensor dimension, we are using per tensor quantization + scale (float): quantization parameter for affine quantization + zero_point (int): quantization parameter for affine quantization + quant_dtype (torch.dtype): desired quantized dtype for determining and validating quant_min and quant_max values. + quant_min (Optional[int]): minimum quantized value for output Tensor, if not specified, it will be derived from dtype + quant_max (Optional[int]): maximum quantized value for output Tensor, if not specified, it will be derived from dtype + zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be either integer or float + if zero_point is in integer domain, zero point is added to the quantized integer value during + quantization + if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) + value during quantization + default is ZeroPointDomain.INT + """ + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + elif zero_point_domain is ZeroPointDomain.NONE and zero_point is not None: + raise ValueError("zero_point should be None when zero_point_domain is NONE") + (_, fq) = _do_fake_quantize_affine( + input, + block_size, + scale, + zero_point, + quant_dtype, + quant_min, + quant_max, + zero_point_domain, + ) + return fq + + +def fake_quantize_affine_cachemask( + input: torch.Tensor, + block_size: Tuple[int, ...], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_dtype: torch.dtype, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + General fake quantize op for quantization-aware training (QAT). + This is equivalent to calling `quantize_affine` + `dequantize_affine` + but without the dtype casts. + + Note: Compared to :func:`~torchao.quantization.quant_primitives.fake_quantize_affine`, + this consumes more memory and returns an additional outlier mask for + intermediate quantized values. + + Args: + Same as :func:`~torchao.quantization.quant_primitives.fake_quantize_affine`. + + Returns: + A 2-tuple of ( + final fake quantized values, + outlier mask for intermediate quantized values + ) + + """ + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + elif zero_point_domain is None and zero_point is not None: + raise ValueError("zero_point should be None when zero_point_domain is NONE") + (q, dq) = _do_fake_quantize_affine( + input, + block_size, + scale, + zero_point, + quant_dtype, + quant_min, + quant_max, + zero_point_domain, + ) + mask = torch.logical_and((q >= quant_min), (q <= quant_max)) + return (dq, mask) + + +def _do_fake_quantize_affine( + input: torch.Tensor, + block_size: Tuple[int, ...], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_dtype: torch.dtype, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Helper function for `fake_quantize_affine` that returns both the + intermediate quantized values and the final dequantized values. + """ + input_dtype = input.dtype + quant_min, quant_max = _get_and_check_qmin_qmax(quant_dtype, quant_min, quant_max) + q = _quantize_affine_no_dtype_cast( + input, + block_size, + scale, + zero_point, + quant_min, + quant_max, + quant_dtype, + zero_point_domain.name, + ) + dq = _dequantize_affine_no_dtype_check( + q, + block_size, + scale, + zero_point, + quant_min, + quant_max, + zero_point_domain.name, + output_dtype=input_dtype, + ) + return (q, dq) + + +@torch.no_grad() +def choose_qparams_affine( + input: torch.Tensor, + mapping_type: MappingType, + block_size: Tuple[int, ...], + target_dtype: torch.dtype, + quant_min: Optional[Union[int, float]] = None, + quant_max: Optional[Union[int, float]] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + input (torch.Tensor): fp32, bf16, fp16 input Tensor + mapping_type (MappingType): determines how the qparams are calculated, symmetric or asymmetric + block_size: (Tuple[int, ...]): granularity of quantization, this means the size of the tensor elements that's sharing the same qparam + e.g. when size is the same as the input tensor dimension, we are using per tensor quantization + target_dtype (torch.dtype): dtype for target quantized Tensor + quant_min (Optional[int]): minimum quantized value for target quantized Tensor + quant_max (Optioanl[int]): maximum quantized value for target quantized Tensor + eps (Optional[float]): minimum scale, if not provided, default to eps of input.dtype + scale_dtype (torch.dtype): dtype for scale Tensor + zero_point_dtype (torch.dtype): dtype for zero_point Tensor + preserve_zero (bool): a flag to indicate whether we need zero to be exactly + representable or not, this is typically required for ops that needs zero padding, like convolution + it's less important for ops that doesn't have zero padding in the op itself, like linear. + + For example, given a floating point Tensor [1.2, 0.1, 3.0, 4.0, 0.4, 0], if `preserve_zero` is True, + we'll make sure there is a integer value corresponding to the floating point 0, e.g. [-3, -8, 3, 7, -7, -8], 0 will be mapped to `-8` without loss. But if `preserve_zero` is not True, there won't be such + gurantee. + + If we don't need zero to be exactly representable, we won't do rounding and clamping for zero_point + + zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be either integer or float + if zero_point is in integer domain, zero point is added to the quantized integer value during + quantization + if zero_point is in floating point domain, zero point is subtracted from the floating point (unquantized) + value during quantization + default is ZeroPointDomain.INT + + Output: + Tuple of scales and zero_points Tensor with requested dtype + """ + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + + return _choose_qparams_affine( + input, + mapping_type.name, + block_size, + target_dtype, + quant_min, + quant_max, + eps, + scale_dtype, + zero_point_dtype, + preserve_zero, + zero_point_domain.name, + ) + + +def choose_qparams_affine_with_min_max( + min_val: torch.Tensor, + max_val: torch.Tensor, + mapping_type: MappingType, + block_size: Tuple[int, ...], + target_dtype: torch.dtype, + quant_min: Optional[int] = None, + quant_max: Optional[int] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT, +) -> Tuple[torch.Tensor, torch.Tensor]: + """A variant of :func:`~torchao.quantization.quant_primitives.choose_qparams_affine` + operator that pass in min_val and max_val directly instead of deriving these from a single input. + This is used for observers in static quantization where min_val and max_val may be obtained through + tracking all the data in calibration data set. + + Args: + Mostly same as :func:`~torchao.quantization.quant_primitives.choose_qparams_affine`. with one + difference: instead of passing in `input` Tensor and use that to calculate min_val/max_val + and then scale/zero_point, we pass in min_val/max_val directly + """ + if zero_point_domain is None: + raise ValueError("Please use ZeroPointDomain.NONE instead of None") + return _choose_qparams_affine( + None, + mapping_type.name, + block_size, + target_dtype, + quant_min, + quant_max, + eps, + scale_dtype, + zero_point_dtype, + preserve_zero, + zero_point_domain.name, + min_val, + max_val, + ) + + +@register_custom_op +def _choose_qparams_affine( + input: Optional[torch.Tensor], + mapping_type: str, + block_size: List[int], + target_dtype: torch.dtype, + quant_min: Optional[Union[int, float, bool]] = None, + quant_max: Optional[Union[int, float, bool]] = None, + eps: Optional[float] = None, + scale_dtype: Optional[torch.dtype] = None, + zero_point_dtype: Optional[torch.dtype] = None, + preserve_zero: bool = True, + zero_point_domain: Optional[str] = "INT", + min_val: Optional[torch.Tensor] = None, + max_val: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """op definition that has compatible signatures with custom op library + + The op does the following: + 1. figure out the dimension for reduction based on block_size + 2. find min_val/max_val based on the dimension for reduction + 3. calculate quantization parameters based on min_val/max_val based on args like `preserve_zero` + and `zero_point_domain` + """ + quant_min, quant_max = _get_and_check_qmin_qmax(target_dtype, quant_min, quant_max) + assert mapping_type in [ + MappingType.SYMMETRIC.name, + MappingType.SYMMETRIC_NO_CLIPPING_ERR.name, + MappingType.ASYMMETRIC.name, + ], f"Unsupported mapping type: {mapping_type}" + if target_dtype in FP8_TYPES: + assert mapping_type == MappingType.SYMMETRIC.name, ( + f"Only symmetric quantization is supported for FP8 types, got {mapping_type}" + ) + + if input is not None: + if scale_dtype is None: + scale_dtype = input.dtype + if eps is None: + eps = torch.finfo(input.dtype).eps + + assert len(block_size) == input.dim(), ( + f"Got input dim:{input.dim()}, block_size: {block_size}" + ) + shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + input = input.view(shape_for_reduction) + + min_val = torch.amin(input, dim=reduction_dims, keepdim=False) + max_val = torch.amax(input, dim=reduction_dims, keepdim=False) + else: + assert min_val is not None and max_val is not None, ( + "Need to provide `min_val` and `max_val` when `input` is None, got: {min_val, max_val}" + ) + assert min_val.dtype == max_val.dtype, ( + "Expecting `min_val` and `max_val` to have the same dtype, got: {min_val.dtype, max_val.dtype}" + ) + + if scale_dtype is None: + scale_dtype = min_val.dtype + if eps is None: + eps = torch.finfo(min_val.dtype).eps + + if preserve_zero: + min_val_neg = torch.min(min_val, torch.zeros_like(min_val)) + max_val_pos = torch.max(max_val, torch.zeros_like(max_val)) + else: + min_val_neg = min_val + max_val_pos = max_val + + if ( + mapping_type == MappingType.SYMMETRIC.name + or mapping_type == MappingType.SYMMETRIC_NO_CLIPPING_ERR.name + ): + # scales + if mapping_type == MappingType.SYMMETRIC.name: + max_val_pos = torch.max(-min_val_neg, max_val_pos) + scale = max_val_pos / (float(quant_max - quant_min) / 2) + else: + assert mapping_type == MappingType.SYMMETRIC_NO_CLIPPING_ERR.name + # calculate smin and smax individually and choose the larger one. For example, if quant_min = -8 and + # quant_max = 7. + # - If smin is bigger: There would be coverage on negative values down to -8, and less rounding + # error than the existing SYMMETRIC case. + # - If smax is bigger: it covers the positive values up to 7. The round + # error may be bigger than the existing SYMMETRIC case. Either way, there's no out-of-range fp values after + # quantization. + smin = min_val_neg / float(quant_min) + smax = max_val_pos / float(quant_max) + mask = smin > smax + scale = torch.where(mask, smin, smax) + # zeros + if not preserve_zero: + raise ValueError( + "preserve_zero == False is not supported for symmetric quantization" + ) + if zero_point_domain == ZeroPointDomain.FLOAT.name: + # TODO INT should not be a valid ZeroPointDomain for symmetric quantization since + # symmetric quant doesn't have a zero_point + raise ValueError( + "zero_point_domain should be ZeroPointDomain.INT or ZeroPointDomain.NONE for symmetric quantization" + ) + if zero_point_domain == ZeroPointDomain.NONE.name: + zero_point = None + else: + zero_point = torch.full_like(scale, int((quant_max + quant_min + 1) / 2)) + scale = torch.clamp(scale, min=eps) + else: + assert mapping_type == MappingType.ASYMMETRIC.name + scale = (max_val_pos - min_val_neg) / float(quant_max - quant_min) + scale = torch.clamp(scale, min=eps) + if zero_point_domain == ZeroPointDomain.NONE.name: + zero_point = None + elif zero_point_domain == ZeroPointDomain.INT.name: + zero_point = quant_min - torch.round(min_val_neg / scale) + zero_point = torch.clamp(zero_point, quant_min, quant_max) + if zero_point_dtype is None: + zero_point_dtype = torch.int32 + else: + assert zero_point_domain == ZeroPointDomain.FLOAT.name, ( + "zero_point must be in FLOAT/INT/None domain for asymmetric quantization" + ) + mid_point = (quant_max + quant_min + 1) / 2 + # this is not preserving zero_point, this is converting to TensorCoreTiledFormat + # TODO move the conversion of zero_point out of quant_primitives + # and into TensorCoreTiledLayout.from_plain + zero_point = min_val_neg + scale * mid_point + + if zero_point is not None: + zero_point = zero_point.to(dtype=zero_point_dtype) + return scale.to(dtype=scale_dtype), zero_point + + +def choose_qparams_and_quantize_affine_qqq( + w: torch.Tensor, + num_bits: int, + group_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert num_bits == 4, f"Unsupported num_bits = {num_bits}" + size_n, size_k = w.shape + assert group_size in [-1, 128, size_k], f"Unsupported groupsize = {group_size}" + orig_device = w.device + if group_size == -1: + group_size = size_k + + if group_size < size_k: + # Reshape to [-1, group_size] + w = w.reshape((-1, group_size)) + + max_q_val = 2**num_bits - 1 + half_q_val = (max_q_val + 1) // 2 + + # Compute scale for each group + s_group = torch.amax(torch.abs(w), -1, keepdim=True) + s_group *= 2 / max_q_val # 2 => symmetric + + # Quantize + q_w = torch.round(w / s_group).int() + q_w += half_q_val + q_w = torch.clamp(q_w, 0, max_q_val) + # Compute ref (dequantized) + w_ref = (q_w - half_q_val).half() * s_group + + # Restore original shapes + def reshape_w(w): + w = w.reshape((size_n, size_k)).contiguous() + return w + + q_w = reshape_w(q_w) + w_ref = reshape_w(w_ref) + + # Compute int8 quantization scale for each channel + s_channel = torch.amax(torch.abs(w_ref), -1, keepdim=True) + s_channel /= 127.0 + t_int8 = (w_ref / s_channel).round().clamp(-128, 127).to(torch.int8) + w_ref = t_int8.half() * s_channel + s_channel = s_channel.reshape(-1, 1).to(dtype=torch.float) + + # Fuse scales + s_group = (s_group.reshape(size_n, -1).contiguous() / s_channel).to( + dtype=torch.half + ) + else: + max_q_val = 2 ** (num_bits - 1) - 1 + + # Compute scale for each channel + s_channel = torch.amax(torch.abs(w), -1, keepdim=True) + s_channel /= max_q_val + + # Quantize + q_w = torch.round(w / s_channel).int() + q_w = torch.clamp(q_w, -max_q_val, max_q_val) + # Compute ref (dequantized) + w_ref = q_w.half() * s_channel + + s_group = torch.tensor([], dtype=torch.half, device=orig_device) + # div 2 ** (8 - self.bits)) to offset right shift in unpacking + s_channel /= 2 ** (8 - num_bits) + s_channel = s_channel.reshape(size_n, -1).contiguous().to(torch.float) + + return q_w, s_group, s_channel, w_ref + + +def choose_qparams_gguf( + input: Optional[torch.Tensor], + block_size: List[int], + target_dtype: torch.dtype, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + There are two sets of qparams: quantized_block_scale, quantized_block_min and super_block_scale_scale and super_block_min_scale + the relationship is the following: + block_scale = quantized_block_scale * super_block_sclae + block_min = quantized_block_min * super_block_min + quantized_val = (float_val - block_min) / block_scale + quant_min + first we calculate block_scale and block_min + then we calculate super_block_scale_scale and super_block_min_scale + after that we can calculate quantized_block_scale and quantized_min_scale + the returned values are: super_block_scale_scale, super_block_min_scale, quantized_block_scale + and quantized_min_scale + """ + dtype = input.dtype + + # 1. get block_scale block_min + shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + input = input.view(shape_for_reduction) + min_val = torch.amin(input, dim=reduction_dims, keepdim=False) + max_val = torch.amax(input, dim=reduction_dims, keepdim=False) + quant_max = 15 + quant_min = 0 + # asymmetric quant to fully utilize the range + block_scale = max_val / (float(quant_max - quant_min) / 2) + block_scale = (max_val - min_val) / float(quant_max - quant_min) + block_min = min_val + + # 2. get super_block_scale_scale and super_block_min_scale + assert _GGUF_QK_K % block_size[-1] == 0 + super_block_size = (1, _GGUF_QK_K // block_size[-1]) + shape_for_reduction, reduction_dims = _get_reduction_params( + super_block_size, block_scale.size() + ) + block_scale = block_scale.view(shape_for_reduction) + block_min = block_min.view(shape_for_reduction) + + shape_after_reduction = shape_for_reduction.copy() + for i in reduction_dims: + shape_after_reduction[i] = 1 + + block_scale_absmax = torch.amax( + torch.abs(block_scale), dim=reduction_dims, keepdim=False + ) + block_min_absmax = torch.amax( + torch.abs(block_min), dim=reduction_dims, keepdim=False + ) + + # 2. get super_block_scale_scale and super_block_min_scale + # TODO: make this configurable + # we also quantize the quantization parameters (scale and min) for each block to 6 bit + # for Q4_K + qparam_quant_max = 2**6 - 1 + qparam_quant_min = 0 + super_block_scale_scale = block_scale_absmax / float( + qparam_quant_max - qparam_quant_min + ) + super_block_min_scale = block_min_absmax / float( + qparam_quant_max - qparam_quant_min + ) + super_block_scale_scale_view = super_block_scale_scale.view(shape_after_reduction) + super_block_min_scale_view = super_block_min_scale.view(shape_after_reduction) + + # 3. quantize block scale and min are stored in 6 bits using super_block_scale_scale and super_block_min_scale + quantized_block_scale = torch.clamp( + block_scale / super_block_scale_scale_view, qparam_quant_min, qparam_quant_max + ) + quantized_block_min = torch.clamp( + block_min / super_block_min_scale_view, qparam_quant_min, qparam_quant_max + ) + return ( + super_block_scale_scale.to(dtype), + super_block_min_scale.to(dtype), + quantized_block_scale.to(dtype), + quantized_block_min.to(dtype), + ) + + +def quantize_gguf( + input: torch.Tensor, + block_size: List[int], + target_dtype: torch.dtype, + super_block_scale_scale: torch.Tensor, + super_block_min_scale: torch.Tensor, + quantized_block_scale: torch.Tensor, + quantized_block_min: torch.Tensor, +) -> torch.Tensor: + assert target_dtype == torch.uint4 + + # step 1: first order quantization + # just going through shape calculation for block_scale and block_min to get the correct shape + input_shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + block_qparam_shape_after_reduction = input_shape_for_reduction.copy() + for i in reduction_dims: + block_qparam_shape_after_reduction[i] = 1 + original_shape = input.shape + input = input.view(input_shape_for_reduction) + quantized_block_scale = quantized_block_scale.view( + block_qparam_shape_after_reduction + ) + quantized_block_min = quantized_block_min.view(block_qparam_shape_after_reduction) + + # step 2: second order quantization, recover unquantized block_scale and block_min + super_block_size = (1, _GGUF_QK_K // block_size[-1], 1) + super_block_input_shape_for_reduction, reduction_dims = _get_reduction_params( + super_block_size, quantized_block_scale.size() + ) + super_block_qparam_shape_after_reduction = ( + super_block_input_shape_for_reduction.copy() + ) + for i in reduction_dims: + super_block_qparam_shape_after_reduction[i] = 1 + + quantized_block_scale = quantized_block_scale.view( + super_block_input_shape_for_reduction + ) + quantized_block_min = quantized_block_min.view( + super_block_input_shape_for_reduction + ) + super_block_scale_scale = super_block_scale_scale.view( + super_block_qparam_shape_after_reduction + ) + super_block_min_scale = super_block_min_scale.view( + super_block_qparam_shape_after_reduction + ) + + block_scale = super_block_scale_scale * quantized_block_scale + block_min = super_block_min_scale * quantized_block_min + + # step 3: quantization with the unquantized block_scale and block_min + block_scale = block_scale.view(block_qparam_shape_after_reduction) + block_min = block_min.view(block_qparam_shape_after_reduction) + int_data = (input - block_min) / block_scale + int_data = int_data.view(original_shape) + + return int_data + + +def dequantize_gguf( + input: torch.Tensor, + block_size: List[int], + target_dtype: torch.dtype, + super_block_scale_scale: torch.Tensor, + super_block_min_scale: torch.Tensor, + quantized_block_scale: torch.Tensor, + quantized_block_min: torch.Tensor, + output_dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + # step 1. reshape input and quantized block scale and min to the shape + # after first quantization + input_shape_for_reduction, reduction_dims = _get_reduction_params( + block_size, input.size() + ) + block_qparam_shape_after_reduction = input_shape_for_reduction.copy() + for i in reduction_dims: + block_qparam_shape_after_reduction[i] = 1 + + original_shape = input.shape + input = input.view(input_shape_for_reduction) + quantized_block_scale = quantized_block_scale.view( + block_qparam_shape_after_reduction + ) + quantized_block_min = quantized_block_min.view(block_qparam_shape_after_reduction) + + # step 2. calculate and reshape block_qparams for second quantization step + super_block_size = (1, _GGUF_QK_K // block_size[-1], 1) + super_block_input_shape_for_reduction, reduction_dims = _get_reduction_params( + super_block_size, quantized_block_scale.size() + ) + super_block_qparam_shape_after_reduction = ( + super_block_input_shape_for_reduction.copy() + ) + for i in reduction_dims: + super_block_qparam_shape_after_reduction[i] = 1 + quantized_block_scale = quantized_block_scale.view( + super_block_input_shape_for_reduction + ) + quantized_block_min = quantized_block_min.view( + super_block_input_shape_for_reduction + ) + super_block_scale_scale = super_block_scale_scale.view( + super_block_qparam_shape_after_reduction + ) + super_block_min_scale = super_block_min_scale.view( + super_block_qparam_shape_after_reduction + ) + + block_scale = super_block_scale_scale * quantized_block_scale + block_min = super_block_min_scale * quantized_block_min + + # step 3. dequantize with block_scale and block_min + block_scale = block_scale.view(block_qparam_shape_after_reduction) + block_min = block_min.view(block_qparam_shape_after_reduction) + dequant = input * block_scale + block_min + dequant = dequant.view(original_shape) + if output_dtype is not None: + dequant = dequant.to(output_dtype) + + return dequant + + +def dequantize_affine_qqq( + w: torch.Tensor, + s_group: torch.Tensor, + s_channel: torch.Tensor, + num_bits: int, + group_size: int, + output_dtype: Optional[torch.dtype] = None, +): + assert num_bits == 4, f"Unsupported num_bits = {num_bits}" + size_n, size_k = w.shape + assert group_size in [-1, 128, size_k], f"Unsupported groupsize = {group_size}" + if group_size == -1: + group_size = size_k + + if group_size < size_k: + # Reshape to [-1, group_size] + w = w.reshape((-1, group_size)) + + max_q_val = 2**num_bits - 1 + half_q_val = (max_q_val + 1) // 2 + + s_group = s_group * s_channel.half() + w_dq = (w - half_q_val).half() * s_group.reshape(-1, 1) + + # Restore original shapes + def reshape_w(w): + w = w.reshape((size_n, size_k)).contiguous() + return w + + w_dq = reshape_w(w_dq) + + else: + s_channel = s_channel * (2 ** (8 - num_bits)) + w_dq = w.half() * s_channel + + if output_dtype is None: + w_dq = w_dq.to(torch.float16) + else: + w_dq = w_dq.to(output_dtype) + + return w_dq + + +# HQQ +############################################################################ +# Shrinking operator (proximal operator for the lp norm) +def _shrink_lp_op(x: torch.Tensor, beta: float, lp_norm: float) -> torch.Tensor: + if lp_norm == 1: + return torch.sign(x) * torch.nn.functional.relu(torch.abs(x) - 1.0 / beta) + else: + return torch.sign(x) * torch.nn.functional.relu( + torch.abs(x) - (1.0 / beta) * torch.pow(torch.abs(x), lp_norm - 1) + ) + + +# Proximal solver || W - dequantize(quantize(W))||_p^p +@torch.inference_mode() +def optimize_weights_proximal_legacy( + tensor: torch.Tensor, + scale: torch.Tensor, + zero: torch.Tensor, + min_max: list, + axis: int = 0, + dtype: Union[torch.dtype, None] = None, + device: Union[str, None] = None, + verbose: bool = False, + opt_params: dict = { + "lp_norm": 0.7, + "beta": 1e1, + "kappa": 1.01, + "iters": 20, + "early_stop": True, + }, +) -> tuple: + lp_norm, beta, kappa, iters, early_stop = ( + opt_params["lp_norm"], + opt_params["beta"], + opt_params["kappa"], + opt_params["iters"], + opt_params["early_stop"], + ) + + device = tensor.device if (device is None) else torch.device(device) + + if dtype is None: + dtype = torch.float16 if (device.type == "cuda") else torch.float32 + + W_f = tensor.to(dtype=dtype, device=device) + scale = scale.to(dtype=dtype, device=device) + zero = zero.to(dtype=dtype, device=device) + + best_error = 1e4 + for i in range(iters): + W_q = torch.round(W_f * scale + zero).clamp(min_max[0], min_max[1]) + W_r = (W_q - zero) / scale + W_e = _shrink_lp_op(W_f - W_r, beta, lp_norm) + zero = torch.mean(W_q - (W_f - W_e) * scale, axis=axis, keepdim=True) + beta *= kappa + + current_error = float(torch.abs(W_f - W_r).mean()) + if verbose: + print("Iter " + str(i + 1), " | Error: " + str(current_error)) + if early_stop: + if current_error < best_error: + best_error = current_error + else: + break + + scale = scale.to(tensor.device) + zero = zero.to(tensor.device) + del W_f, W_q, W_r, W_e + torch.cuda.empty_cache() + + W_q = torch.round(tensor * scale + zero).clamp(min_max[0], min_max[1]) + return W_q, scale, zero + + +# Mainly used to check if the group-size is divisible by numel() +def _is_divisible(val1: int, val2: int) -> bool: + return int(val2 * math.ceil(val1 / val2)) == val1 + + +# Converts hqq format W_dequant = (W_q - zero)*scale into affinequantized format: (W_q - mid_point)*scale_ao + zero_ao +def _convert_to_affinequantized_format( + W_q: torch.Tensor, + scale: torch.Tensor, + zero: torch.Tensor, + nbits: int, + shape: Union[List, Tuple, torch.Size], +) -> Tuple: + quant_min = 0 + quant_max = 2**nbits - 1 + mid_point = (quant_max + quant_min + 1) / 2 + zero_ao = ((mid_point - zero.float()) * scale.float()).to(zero.dtype) + scale_ao = scale + W_q_ao = W_q.view(shape) + return W_q_ao, scale_ao, zero_ao + + +# Main hqq quantizer function +def choose_qparams_and_quantize_affine_hqq( + tensor: torch.Tensor, + nbits: float = 4, + group_size: int = 64, + optimize: bool = True, + axis: int = 1, + compute_dtype: torch.dtype = torch.float16, + device: str = "cuda", + verbose: bool = False, # to check the optimizer error + raw_output: bool = False, # If True, it will return the quant params in hqq lib format + optimize_weights: Callable = optimize_weights_proximal_legacy, # weights proximal optimizer function +) -> tuple: + assert axis in [0, 1], "axis should be either 0 or 1" + if group_size is not None: + assert _is_divisible(tensor.numel(), group_size), ( + "group_size should be divisble by the total tensor dimensions. shape: " + + str(tensor.shape) + + ", group_size: " + + str(group_size) + ) + + # It's better to work with float32 here + W = tensor.to(device=device, dtype=torch.float32) + shape = W.shape + + # Reshape for grouping + if group_size is not None: + W = W.reshape([-1, group_size]) if (axis == 1) else W.reshape([group_size, -1]) + + # Get min/max values + _min = W.min(axis=axis, keepdim=True)[0] + _max = W.max(axis=axis, keepdim=True)[0] + + max_v = round(2**nbits - 1) + min_v = 0 + min_max = [min_v, max_v] + + # Clamp to avoid fp16 issues + scale = (max_v / (_max - _min)).clamp(max=2e4) + zero = -_min * scale + + # Round zero as in: https://github.com/casper-hansen/AutoAWQ/blob/main/awq/quantize/quantizer.py#L42C9-L42C14 + if nbits in [4]: + zero = torch.round(zero) + + # Fine-tune weights + if optimize: + W_q, scale, zero = optimize_weights( + tensor=W, + scale=scale, + zero=zero, + min_max=min_max, + axis=axis, + device=device, + verbose=verbose, + ) + else: + zero = zero.to(compute_dtype) + scale = scale.to(compute_dtype) + W_q = torch.round(W * scale + zero).clamp(min_max[0], min_max[1]) + + # Store meta-data (we invert the scale for dequantization) + scale = 1.0 / scale + + # Convert to TensorCoreTiled format + # TODO move the conversion of zero_point out of quant_primitives + # and into TensorCoreTiledLayout.from_plain and rename this + # helper function correctly. + if raw_output is False: + W_q, scale, zero = _convert_to_affinequantized_format( + W_q, scale, zero, nbits, shape + ) + else: + # this path was not used before, the way hqq sets up scale/zero is transposed + # compared to the rest of our utils so we need to reshape them acccordingly. + W_q = W_q.reshape(shape) + if axis == 1: + scale = scale.reshape(shape[0], -1) + zero = zero.reshape(shape[0], -1) + else: + scale = scale.reshape(-1, shape[-1]) + zero = zero.reshape(-1, shape[-1]) + # Make sure all the weights are in the right compute_dtype/device + W_q = W_q.to(dtype=torch.uint8, device=device) + scale = scale.to(dtype=compute_dtype, device=device) + zero = zero.to(dtype=compute_dtype, device=device) + + # cleanup + del W, _min, _max + torch.cuda.empty_cache() + + return W_q, scale, zero, shape + + +def choose_qparams_affine_floatx( + tensor: torch.Tensor, ebits: int, mbits: int +) -> torch.Tensor: + # _n_ones() is not compatible with torch.compile() due to << operator + # https://github.com/pytorch/pytorch/issues/119152 + # exp_bias = _n_ones(ebits - 1) + # max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2 ** mbits)) + + # workaround: global lookup table + exp_bias = _ONES_TABLE[ebits - 1] + max_normal = 2 ** (_ONES_TABLE[ebits] - exp_bias) * ( + _ONES_TABLE[mbits + 1] / (2**mbits) + ) + + dtype = tensor.dtype + tensor = tensor.float() + scale = tensor.abs().amax(1).clamp(min=1e-12) / max_normal + return scale.to(dtype) + + +def quantize_affine_floatx( + tensor: torch.Tensor, scale: torch.Tensor, ebits: int, mbits: int +) -> torch.Tensor: + """Quantizes the float32 high precision floating point tensor to low precision floating point number and + converts the result to unpacked floating point format with the format of 00SEEEMM (for fp6_e3m2) where S means sign bit, e means exponent bit and m means mantissa bit + """ + tensor = tensor.float() + tensor_floatx = _f32_to_floatx_unpacked(tensor / scale.view(-1, 1), ebits, mbits) + return tensor_floatx + + +def dequantize_affine_floatx( + tensor: torch.Tensor, + scale: torch.Tensor, + ebits: int, + mbits: int, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + tensor = _floatx_unpacked_to_f32(tensor, ebits, mbits) + tensor = tensor * scale.float().view(-1, 1) + tensor = tensor.to(dtype=output_dtype) + return tensor + + +def choose_qparams_affine_float8( + tensor: torch.Tensor, + float8_dtype: torch.dtype = torch.float8_e4m3fn, +) -> torch.Tensor: + """ + Calculates float8 scaling factor for the given high precision tensor, using tensorwise granularity. + + Args: + tensor (torch.Tensor): Input tensor to be quantized. + float8_dtype (torch.dtype): Data type of the quantized tensor (e.g., torch.float8_e4m3fn, torch.float8_e5m2). + """ + # only tensorwise scaling is supported for now: + quant_min, quant_max = torch.finfo(float8_dtype).min, torch.finfo(float8_dtype).max + min_val_neg = torch.min(tensor) + max_val_pos = torch.max(tensor) + max_val_pos = torch.max(-min_val_neg, max_val_pos) + scale = max_val_pos / (float(quant_max - quant_min) / 2) + return scale.to(dtype=torch.float32) + + +def quantize_affine_float8( + tensor: torch.Tensor, + scale: torch.Tensor, + float8_dtype: torch.dtype = torch.float8_e4m3fn, +) -> torch.Tensor: + """ + Quantizes the high precision floating point tensor to a float8 tensor, using the given scaling factor. + + Args: + tensor (torch.Tensor): Input tensor to be quantized. + scale (torch.Tensor): Scaling factor for the quantization. + float8_dtype (torch.dtype): Data type of the quantized tensor (e.g., torch.float8_e4m3fn, torch.float8_e5m2). + """ + # Note: when the line below is compiled with `torch.compile`, `tensor` is automatically + # upcasted to `float32` to multiply with the scale, since scale is a fp32 tensor in float8 quantization. + # In order to match numerics between eager and compile, we upcast manually here. + tensor_scaled = tensor.to(torch.float32) / scale + max_value = torch.finfo(float8_dtype).max + tensor_clamped = tensor_scaled.clamp(min=-max_value, max=max_value) + fp8_tensor = tensor_clamped.to(float8_dtype) + return fp8_tensor + + +def dequantize_affine_float8( + tensor: torch.Tensor, + scale: torch.Tensor, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Dequantizes the float8 tensor to high precision tensor. + + Args: + tensor (torch.Tensor): Input float8 tensor to be dequantized. + scale (torch.Tensor): Scaling factor for the dequantization. + output_dtype (torch.dtype): Data type of the output tensor (e.g., torch.float32). + """ + # Note: when the line below is compiled with `torch.compile`, `tensor` is automatically + # upcasted to `float32` to divide by the scale, since scale is a fp32 for float8 quantization. + # In order to match numerics between eager and compile, we upcast manually here. + fp8_tensor = tensor.to(torch.float32) + hp_tensor = fp8_tensor * scale + return hp_tensor.to(output_dtype) diff --git a/lib/python3.12/site-packages/torchao/quantization/smoothquant.py b/lib/python3.12/site-packages/torchao/quantization/smoothquant.py new file mode 100644 index 0000000000000000000000000000000000000000..972f0cc6ec1b142e002f76370f46f45f6b4aabf5 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/smoothquant.py @@ -0,0 +1,266 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +Testing out accuracy-only implementation of SmoothQuant +(https://arxiv.org/pdf/2211.10438.pdf) +Note: this is an application of input-weight equalization, with the addition that the +multiplication by scale is fused into the preceding layer, specifically for relevant +parts of transformer blocks. +""" + +import torch +import torch.nn.functional as F + +from .utils import ( + dynamically_quantize_per_channel, + quant_int8_dynamic_per_token_linear, +) + +__all__ = [ + "get_scale", + "SmoothFakeDynQuantMixin", + "SmoothFakeDynamicallyQuantizedLinear", + "swap_linear_with_smooth_fq_linear", + "smooth_fq_linear_to_inference", + "set_smooth_fq_attribute", +] + + +def get_scale(X_absmax, W_absmax, alpha=0.5): + """ + Calculate the scale based on abs(max(X)), abs(max(W)), and alpha. + + Args: + X_absmax (torch.Tensor): Absolute maximum values of the input tensor X. + W_absmax (torch.Tensor): Absolute maximum values of the weight tensor W. + alpha (float, optional): Scaling factor. Defaults to 0.5. + + Returns: + torch.Tensor: The calculated scale of dimension `k` if X is of dimension `b*n*k` and W is of dimension `k*m`. + """ + X_pow = torch.pow(X_absmax, alpha) + W_pow = torch.pow(W_absmax, 1.0 - alpha) + div = X_pow / W_pow + return div.reshape(-1) + + +class SmoothFakeDynQuantMixin(torch.nn.Module): + def init_smoothquant_variables(self, alpha): + self.calibrating = True + self.x_running_abs_max = None + self.register_buffer("smooth_scale", None) + self.alpha = alpha + # debug only + self.debug_skip_scaling = False + # self.debug_skip_scaling = True + + # Currently torch._int_mm cuBLAS underlying kernel does not work with + # non-contiguous weight. However, torch.compil'ing through + # torch._int_mm leads to triton code which is ~2x faster if the weight + # is transposed. So, for now we have a debug flag to toggle whether + # we store the quantized weight transposed, so that we can get correct + # numerics both in eager mode and after torch.compile. + # The default is True for cuBLAS / eager mode, set to False for + # torch.compile. + # self.store_w_int_repr_t = True + self.store_w_int_repr_t = False + + def update_x_running_abs_max(self, X): + # update the running max of incoming activations + all_dims_except_last = tuple(range(len(X.shape) - 1)) + cur_abs_max = torch.amax(torch.abs(X), dim=all_dims_except_last) + if self.x_running_abs_max is None: + self.x_running_abs_max = cur_abs_max + else: + self.x_running_abs_max = torch.max(cur_abs_max, self.x_running_abs_max) + + def get_scaled_quantized_w(self): + # inference + assert self.smooth_scale is not None, ( + "self.smooth_scale is None, did you turn on inference?" + ) + W = self.weight + + # scale weight + # in the future, this can be done ahead of time instead of + # during inference + if not self.debug_skip_scaling: + # TODO(future): do below in `to_inference` instead of here + W = torch.matmul( + torch.diag(self.smooth_scale), W.transpose(0, 1) + ).transpose(0, 1) + + # fake quantize input and weight, and then do matmul in fp32/fp16 + # in the future, this should be replaced with quantized kernels which + # work on NVIDIA GPUs (such as protoquant's implementation) + W_int_repr, W_scales, W_zps = dynamically_quantize_per_channel( + W, -128, 127, torch.int8 + ) + W_int_repr = W_int_repr.contiguous() + return W_int_repr, W_scales, W_zps + + def to_inference(self): + raise NotImplementedError() + + def fold_weight(self): + # note: _W_zps are zeroes and they are ignored + # TODO(future PR): set up serialization for this + W_int_repr, self.W_scales, _W_zps = self.get_scaled_quantized_w() + # need to store transposed weights to make eager mode matmul + # op work in cuBlas, or non-transposed to make it fast in torch.compile + if self.store_w_int_repr_t: + self.register_buffer("W_int_repr", W_int_repr.transpose(0, 1).contiguous()) + else: + self.register_buffer("W_int_repr", W_int_repr.contiguous()) + del self.weight + + def set_debug_x_absmax(self): + """ + Sets `self.x_running_abs_max` to a value which will lead to smooth scale + of all ones if `alpha=0.5`, to enable performance benchmarking without + calibration. + """ + raise NotImplementedError() + + +class SmoothFakeDynamicallyQuantizedLinear(SmoothFakeDynQuantMixin, torch.nn.Linear): + """ + This is a replacement for `torch.nn.Linear` which implements dynamic per-token + activation quantization and dynamic per-channel weight quantization based on + Smoothquant scaling. + """ + + def __init__(self, *args, **kwargs): + alpha = kwargs.pop("alpha") + super().__init__(*args, **kwargs) + self.init_smoothquant_variables(alpha) + + def forward(self, X, *args, **kwargs): + if self.calibrating: + self.update_x_running_abs_max(X) + Y = F.linear(X, self.weight, self.bias) + else: + if not self.debug_skip_scaling: + # Ideally this would be fused into preceding layers + # but in practice torch.compile fuses it with other + # ops so the slowdown is minimal + X = X / self.smooth_scale + W_int_repr_t = ( + self.W_int_repr if self.store_w_int_repr_t else self.W_int_repr.t() + ) + Y = quant_int8_dynamic_per_token_linear( + X, W_int_repr_t, self.W_scales, self.bias, X.dtype + ) + return Y + + @classmethod + def from_float(cls, mod, alpha=0.5): + """ + Converts a `mod` of class `torch.nn.Linear` to the smooth fake quantized + version of it. Note: requires calibration. + """ + # create the new module with a toy size to ensure initialization is fast + fake_in_features, fake_out_features = 8, 8 + new_mod = cls( + fake_in_features, fake_out_features, bias=mod.bias is not None, alpha=alpha + ) + new_mod.in_features = mod.in_features + new_mod.out_features = mod.out_features + new_mod.weight = mod.weight + new_mod.bias = mod.bias + # TODO: test when creation is on cuda + device_to_use = next(mod.parameters()).device + new_mod.to(device_to_use) + return new_mod + + def to_inference(self): + """ + Calculates the smoothquant scale based on calibration + in preparation for inference + """ + assert self.x_running_abs_max is not None, "no calibration data found" + self.calibrating = False + self.smooth_scale = get_scale( + self.x_running_abs_max, + torch.max(torch.abs(self.weight.transpose(0, 1)), dim=1).values, + alpha=self.alpha, + ) + self.fold_weight() + + def set_debug_x_absmax(self): + w_absmax = torch.max(torch.abs(self.weight.transpose(0, 1)), dim=1).values + self.x_running_abs_max = w_absmax + + +# +# utils to use the smooth linear on real models +# + +source_cls_to_target_cls = { + torch.nn.Linear: SmoothFakeDynamicallyQuantizedLinear, + torch.nn.modules.linear.NonDynamicallyQuantizableLinear: SmoothFakeDynamicallyQuantizedLinear, +} + + +def swap_linear_with_smooth_fq_linear( + model, skip_fqn_list=None, cur_fqn="", alpha=0.5 +) -> None: + """ + Replaces linear layers in the model with their SmoothFakeDynamicallyQuantizedLinear equivalents. + + Args: + model (torch.nn.Module): The model containing linear layers to be replaced. + skip_fqn_list (list of str, optional): List of fully qualified names to skip during replacement. Defaults to None. + cur_fqn (str, optional): The current fully qualified name of the module being processed. Defaults to "". + alpha (float, optional): The scaling factor for SmoothQuant. Defaults to 0.5. + + Returns: + None + """ + + name_to_child = dict(model.named_children()) + for name, child in name_to_child.items(): + if cur_fqn == "": + new_fqn = name + else: + new_fqn = f"{cur_fqn}.{name}" + if ((skip_fqn_list is None) or (new_fqn not in skip_fqn_list)) and ( + type(child) in source_cls_to_target_cls.keys() + ): + target_cls = source_cls_to_target_cls[type(child)] + new_child = target_cls.from_float(child, alpha=alpha) + setattr(model, name, new_child) + else: + swap_linear_with_smooth_fq_linear(child, skip_fqn_list, new_fqn, alpha) + + +def smooth_fq_linear_to_inference(model, debug_skip_calibration=False) -> None: + """ + Prepares the model for inference by calculating the smoothquant scale for each SmoothFakeDynamicallyQuantizedLinear layer. + + Args: + model (torch.nn.Module): The model containing SmoothFakeDynamicallyQuantizedLinear layers. + debug_skip_calibration (bool, optional): If True, sets the running maximum of activations to a debug value for performance benchmarking. + Defaults to False. + + Returns: + None + """ + for _, mod in model.named_modules(): + if isinstance(mod, tuple(source_cls_to_target_cls.values())): + if debug_skip_calibration: + mod.set_debug_x_absmax() + mod.to_inference() + + +# useful for quickly toggling smoothquant debug settings on all smoothquant +# modules in a model +def set_smooth_fq_attribute(model, attribute_name, new_attribute_val): + for _, mod in model.named_modules(): + if isinstance(mod, tuple(source_cls_to_target_cls.values())): + if hasattr(mod, attribute_name): + setattr(mod, attribute_name, new_attribute_val) diff --git a/lib/python3.12/site-packages/torchao/quantization/subclass.py b/lib/python3.12/site-packages/torchao/quantization/subclass.py new file mode 100644 index 0000000000000000000000000000000000000000..abaad317eba4de7ff95a7ded6fc0b51218621426 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/subclass.py @@ -0,0 +1,709 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + + +import torch +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.quantization.utils import ( + dequantize_per_channel, + dynamically_quantize_per_channel, + groupwise_affine_quantize_tensor, + quant_int8_dynamic_per_token_linear, + unpack_tinygemm_scales_and_zeros, +) +from torchao.utils import ( + check_cpu_version, + check_xpu_version, + find_multiple, +) + +from .quant_primitives import ( + ZeroPointDomain, +) + +__all__ = [ + "Int8DynamicallyQuantizedLinearWeight", + "Int8WeightOnlyQuantizedLinearWeight", + "Int4WeightOnlyQuantizedLinearWeight", +] + + +aten = torch.ops.aten + + +class QuantizedLinearWeightBase(torch.Tensor): + """ + Base quantized tensor subclass for quantized linear weights. When the from_float method is used, + to create an instance of any QuantizedLinearWeightBase, we assume the input + weight is oriented the way it is in a normal linear op, i.e. out-channels x in-channels. + + The shape and dtype of the tensor subclass represent how the tensor subclass looks externally, + regardless of the internal representation's type or orientation. + """ + + @staticmethod + def __new__(cls, int_data, transposed, shape, *args, **kwargs): + kwargs["device"] = int_data.device + kwargs["layout"] = ( + kwargs.get("layout") if kwargs.get("layout", False) else int_data.layout + ) + assert "dtype" in kwargs + assert not kwargs.get("requires_grad", False) + kwargs["requires_grad"] = False + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__(self, int_data, transposed, *args, **kwargs): + self.int_data = int_data + + self.transposed = transposed + + @staticmethod + def _quantized_op(act_mat, w_qtensor, bias): + pass + + def __repr__(self): + return ( + f"{self.__class__.__name__}(data={self.dequantize()}, shape={self.shape}, " + f"device={self.device}, dtype={self.dtype}, requires_grad={self.requires_grad})" + ) + + def dequantize(self): + pass + + def int_repr(self): + pass + + def q_params(self): + pass + + def half(self): + return self.to(torch.float16) + + def _get_to_kwargs(self, *args, **kwargs): + device, dtype, _, memory_format = torch._C._nn._parse_to(*args, **kwargs) + device = self.device if device is None else device + dtype = self.dtype if dtype is None else dtype + memory_format = ( + memory_format if memory_format is not None else torch.preserve_format + ) + kwargs = { + "device": device, + "dtype": dtype, + "memory_format": memory_format, + } + return kwargs + + def _apply_fn_to_data(self, fn): + pass + + def _change_shape(self): + pass + + def __tensor_flatten__(self): + pass + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + pass + + @classmethod + def from_float(cls, input_float): + pass + + # __torch_function__ = torch._C._disabled_torch_function_impl + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + kwargs = {} if kwargs is None else kwargs + + if func is torch.nn.functional.linear: + mat1, w_qtensor, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + assert not w_qtensor.transposed + return cls._quantized_op(mat1, w_qtensor, bias) + + try: + with torch._C.DisableTorchFunctionSubclass(): + return func(*args, **kwargs) + except Exception: + print(f"ERR: subclass doesn't implement {func}") + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + # two scenarios where we currently fall back to vanilla mm: + # 1 - when tensor is on CPU: we are missing qmm for CPU, but we should have a CPU implementation + # for consistency and to allow people to test + # 2 - we're given non-floats - quantizing long to int8 is crazy + if ( + func in [aten.mm.default, aten.addmm.default] + and args[0].is_floating_point() + and args[0].is_cuda + ): + if func == aten.addmm.default: + assert args[1].shape[-1] == args[2].shape[0], ( + f"need mat1 shape: {args[1].shape} final" + f"dim to match mat2 shape: {args[2].shape} first dim " + ) + mat1, w_qtensor, bias = ( + args[1], + args[2], + args[0], + ) + else: + assert args[0].shape[-1] == args[1].shape[0], ( + f"need mat1 shape: {args[0].shape} final dim" + f"to match mat2 shape: {args[1].shape} first dim" + ) + mat1, w_qtensor, bias = ( + args[0], + args[1], + None if len(args) == 2 else args[2], + ) + # call the quantized op for the specific type + # of quantized tensor subclass + return cls._quantized_op(mat1, w_qtensor, bias) + + if func is aten.detach.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + if func is aten.clone.default: + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + if func is aten.t.default: + args[0].transposed = not args[0].transposed + new = args[0]._change_shape(args[0].shape[::-1]) + return return_and_correct_aliasing(func, args, kwargs, new) + + if func is aten._to_copy.default: + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].to(*args[1:], **kwargs)._apply_fn_to_data(torch.clone), + ) + + +class ConstructTensorSubclass(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.args = args + self.kwargs = kwargs + + def forward(self, x): + pass + + def right_inverse(self, tensor_subclass_instance): + fields, _ = tensor_subclass_instance.__tensor_flatten__() + return [getattr(tensor_subclass_instance, field) for field in fields] + + +@torch._dynamo.allow_in_graph +def from_qtensor_components_int8dyn(*args, **kwargs): + return Int8DynamicallyQuantizedLinearWeight(*args, **kwargs) + + +class ConstructTensorSubclassInt8Dyn(ConstructTensorSubclass): + def forward(self, int_data, q_scales): + return from_qtensor_components_int8dyn( + int_data, q_scales, *self.args, **self.kwargs + ) + + +class Int8DynamicallyQuantizedLinearWeight(QuantizedLinearWeightBase): + """ + A Tensor subclass that when applied to a weight used in a linear op/module, changes the + linear op to a dynamically quantized linear op with symmetric per-token and per-channel + quantization on the activation and weight respectively. + """ + + subclass_constructor = ConstructTensorSubclassInt8Dyn + + @staticmethod + def __new__(cls, int_data, q_scales, transposed, shape, dtype=None, **kwargs): + if dtype is None: + dtype = q_scales.dtype + kwargs["dtype"] = dtype + return super().__new__(cls, int_data, transposed, shape, **kwargs) # type: ignore[attr-defined] + + def __init__(self, int_data, q_scales, transposed, shape, dtype=None, **kwargs): + self.q_scales = q_scales + super().__init__(int_data, transposed) + + @staticmethod + def _quantized_op(act_mat, w_qtensor, bias): + return quant_int8_dynamic_per_token_linear( + act_mat, w_qtensor.int_data, w_qtensor.q_scales, bias, act_mat.dtype + ) + + def dequantize(self, dtype=None): + """ + Obtain the dequantized version of the quantized tensor subclass + """ + zero_points = torch.zeros( + self.q_scales.shape, device=self.q_scales.device, dtype=self.q_scales.dtype + ) + # zero_points = 0 + # TODO: fix dtype here? `to(self.dtype)` is not overwritten by `dtype` arg? + dq_t = dequantize_per_channel( + self.int_data.t(), + self.q_scales, + zero_points, + self.dtype if dtype is None else dtype, + ).to(self.dtype) + # data was transposed to dequantize so make sure shape is correct + return dq_t if not self.transposed else dq_t.t() + + def int_repr(self): + """ + Get the internal integer representation of the quantized tensor + """ + return self.int_data if self.transposed else self.int_data.t() + + def q_params(self): + """ + Get the quantization scales for the quantized tensor + """ + return {"q_scales": self.q_scales} + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + return self.__class__( + self.int_data.to(kwargs["device"]), + self.q_scales.to(kwargs["device"]), + self.transposed, + self.shape, + **kwargs, + ) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.int_data), + fn(self.q_scales), + self.transposed, + self.shape, + dtype=self.dtype, + ) + + # `QuantizedLinearWeightBase` inconsistently. + + def _change_shape(self, shape): + return self.__class__( + self.int_data, self.q_scales, self.transposed, shape, dtype=self.dtype + ) + + def __tensor_flatten__(self): + # note: the order of args must match the order of args in __init__ + return ["int_data", "q_scales"], [self.transposed, self.shape, self.dtype] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size=None, outer_stride=None + ): + int_data, q_scales = tensor_data_dict["int_data"], tensor_data_dict["q_scales"] + transposed, shape, dtype = tensor_attributes + return cls( + int_data, + q_scales, + transposed, + shape if outer_size is None else outer_size, + dtype=dtype, + strides=outer_stride, + ) + + @classmethod + def from_float(cls, input_float, qmin=-128, qmax=127, dtype=None): + """ + Method used to convert a linear weight tensor to an instance of the + Int8DynamicallyQuantizedLinearWeight subclass. + + Example usage:: + + model.lin_mod.weight = ( + Int8DynamicallyQuantizedLinearWeight.from_float(model.lin_mod.weight) + ) + """ + if dtype is None: + dtype = input_float.dtype + + # because we call transpose in dequantization + w_int_repr, w_scales, _ = dynamically_quantize_per_channel( + input_float, qmin, qmax, torch.int8 + ) + # the desired representation shape for fast quantized matmul is + # transposed compared to how it's stored as a linear weight, + # i.e. we want in_channels as dim=0 and out_channels (and quantized axis) as dim=1 + # however the external representation of our tensor will maintain the correct + # shape attribute which needs to be tracked directly. + int_data = w_int_repr.contiguous().t() + if not issubclass(cls, Int8DynamicallyQuantizedLinearWeight): + int_data = int_data.contiguous() + return cls( + int_data, + w_scales, + False, + input_float.shape, + dtype=dtype, + ) + + +@torch._dynamo.allow_in_graph +def from_qtensor_components_int8wo(*args, **kwargs): + return Int8WeightOnlyQuantizedLinearWeight(*args, **kwargs) + + +class ConstructTensorSubclassInt8wo(ConstructTensorSubclass): + def forward(self, int_data, q_scales): + return from_qtensor_components_int8wo( + int_data, q_scales, *self.args, **self.kwargs + ) + + +class Int8WeightOnlyQuantizedLinearWeight(Int8DynamicallyQuantizedLinearWeight): + """ + A Tensor subclass that when applied to a weight used in a linear op/module, + changes the linear op to a weight-only quantized linear op with symmetric + per-channel quantization on the weight. + """ + + subclass_constructor = ConstructTensorSubclassInt8wo + + @staticmethod + def _quantized_op(act_mat, w_qtensor, bias): + orig_dtype = act_mat.dtype + y = ( + torch.mm( + act_mat.reshape(-1, act_mat.shape[-1]), + w_qtensor.int_data.to(act_mat.dtype), + ) + * w_qtensor.q_scales + ) + y = y.reshape(*act_mat.shape[:-1], y.shape[-1]) + if bias is not None: + y += bias + return y.to(orig_dtype) + + +@torch._dynamo.allow_in_graph +def from_qtensor_components_int4wo(*args, **kwargs): + return Int4WeightOnlyQuantizedLinearWeight(*args, **kwargs) + + +class ConstructTensorSubclassInt4wo(ConstructTensorSubclass): + def forward(self, int_data, scales_and_zeros): + return from_qtensor_components_int4wo( + int_data, scales_and_zeros, *self.args, **self.kwargs + ) + + +class Int4WeightOnlyQuantizedLinearWeight(QuantizedLinearWeightBase): + """ + A Tensor subclass that when applied to a weight used in a linear op/module, + changes that linear op to a weight-only int4 quantized linear op with groupwise + affine quantization on the weight. + """ + + subclass_constructor = ConstructTensorSubclassInt4wo + + @staticmethod + def __new__( + cls, + int_data, + scales_and_zeros, + transposed, + shape, + groupsize=128, + inner_k_tiles=8, + zero_point_domain=ZeroPointDomain.FLOAT, + preserve_zero=False, + dtype=None, + **kwargs, + ): + if dtype is None: + dtype = scales_and_zeros.dtype + kwargs["dtype"] = dtype + return super().__new__(cls, int_data, transposed, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + int_data, + scales_and_zeros, + transposed, + shape, + groupsize, + inner_k_tiles, + zero_point_domain, + preserve_zero, + dtype, + **kwargs, + ): + # the transposed flag tracks whether the tensor subclass has been transposed relative + # to how a weight is normally stored in a linear i.e. [out_features, in_features]. + # tracking both transposed and shape is slightly redundant but corner cases like + # square matrices can cause issues otherwise + + self.scales_and_zeros = scales_and_zeros + self.groupsize = groupsize + self.inner_k_tiles = inner_k_tiles + self.zero_point_domain = zero_point_domain + self.preserve_zero = preserve_zero + super().__init__(int_data, transposed) + + @staticmethod + def _quantized_op(act_mat, w_qtensor, bias): + orig_act_size = act_mat.size() + orig_dtype = act_mat.dtype + + # reshape and pad activation + act_mat = act_mat.reshape(-1, act_mat.shape[-1]).to(torch.bfloat16) + pad_size = find_multiple(act_mat.shape[-1], 1024) + act_mat = torch.nn.functional.pad(act_mat, (0, pad_size - act_mat.shape[-1])) + + # matmul + if check_cpu_version(act_mat.device): + y = aten._weight_int4pack_mm_for_cpu( + act_mat.contiguous(), + w_qtensor.int_data, + w_qtensor.groupsize, + w_qtensor.scales_and_zeros, + ) + elif check_xpu_version(act_mat.device): + if not w_qtensor.zero_point_domain == ZeroPointDomain.INT: + y = aten._weight_int4pack_mm( + act_mat.contiguous(), + w_qtensor.int_data, + w_qtensor.groupsize, + w_qtensor.scales_and_zeros, + ) + else: + y = aten._weight_int4pack_mm_with_scales_and_zeros( + act_mat.contiguous(), + w_qtensor.int_data, + w_qtensor.groupsize, + w_qtensor.scales_and_zeros[0], + w_qtensor.scales_and_zeros[1], + ) + else: + y = aten._weight_int4pack_mm( + act_mat.contiguous(), + w_qtensor.int_data, + w_qtensor.groupsize, + w_qtensor.scales_and_zeros, + ) + + # remove out_feature padding + orig_out_features = ( + w_qtensor.shape[-1] if w_qtensor.transposed else w_qtensor.shape[-2] + ) + y = y[:, :orig_out_features] + + y = y.reshape(*orig_act_size[:-1], orig_out_features) + if bias is not None: + y += bias + return y.to(orig_dtype) + + def dequantize(self): + eye_shape = self.shape[1] if not self.transposed else self.shape[0] + w_dq = self._quantized_op( + torch.eye(eye_shape, device=self.device, dtype=self.dtype), self, None + ) + # we dequantized using linear with the identity matrix, output has shape [in_channels, out_channels] + # so we need to transpose back to get the original shape unless self.transposed is set. + w_dq = w_dq if self.transposed else w_dq.t() + return w_dq.to(self.dtype) + + def int_repr(self): + return self.int_data + + def q_params(self): + scales, zero_points = unpack_tinygemm_scales_and_zeros( + self.scales_and_zeros, + ) + return {"q_scales": scales, "q_zero_points": zero_points} + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + return self.__class__( + self.int_data.to(kwargs["device"]), + self.scales_and_zeros.to(kwargs["device"]), + self.transposed, + self.shape, + self.groupsize, + self.inner_k_tiles, + self.zero_point_domain, + self.preserve_zero, + **kwargs, + ) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.int_data), + fn(self.scales_and_zeros), + self.transposed, + self.shape, + self.groupsize, + self.inner_k_tiles, + self.zero_point_domain, + self.preserve_zero, + dtype=self.dtype, + ) + + # `QuantizedLinearWeightBase` inconsistently. + + def _change_shape(self, shape): + return self.__class__( + self.int_data, + self.scales_and_zeros, + self.transposed, + shape, + self.groupsize, + self.inner_k_tiles, + self.zero_point_domain, + self.preserve_zero, + dtype=self.dtype, + ) + + def __tensor_flatten__(self): + return ["int_data", "scales_and_zeros"], ( + self.transposed, + self.shape, + self.groupsize, + self.inner_k_tiles, + self.zero_point_domain, + self.preserve_zero, + self.dtype, + ) + + @classmethod + + # `QuantizedLinearWeightBase` inconsistently. + + def __tensor_unflatten__( + cls, tensor_data_dict, attributes, outer_size=None, outer_stride=None + ): + int_data, scales_and_zeros = ( + tensor_data_dict["int_data"], + tensor_data_dict["scales_and_zeros"], + ) + ( + transposed, + shape, + groupsize, + inner_k_tiles, + zero_point_domain, + preserve_zero, + dtype, + ) = attributes + return cls( + int_data, + scales_and_zeros, + transposed, + shape if outer_size is None else outer_size, + groupsize, + inner_k_tiles, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + dtype=dtype, + strides=outer_stride, + ) + + @classmethod + def from_float( + cls, + input_float, + groupsize=128, + inner_k_tiles=8, + zero_point_domain=ZeroPointDomain.FLOAT, + preserve_zero=False, + dtype=None, + ): + """ + Method used to convert a linear weight tensor to an instance of the + Int4WeightOnlyQuantizedLinearWeight subclass. + + Example usage:: + + model.lin_mod.weight = ( + Int4WeightOnlyQuantizedLinearWeight.from_float(model.lin_mod.weight) + ) + """ + if dtype is None: + dtype = input_float.dtype + + int_data, scales_and_zeros, transposed, groupsize, inner_k_tils = ( + cls.to_qtensor_components( + input_float, + groupsize, + inner_k_tiles, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + ) + ) + return cls( + int_data, + scales_and_zeros, + transposed, + input_float.shape, + groupsize, + inner_k_tiles, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + dtype=dtype, + ) + + @classmethod + def to_qtensor_components( + cls, + input_float, + groupsize=128, + inner_k_tiles=8, + zero_point_domain=ZeroPointDomain.FLOAT, + preserve_zero=False, + ): + assert groupsize in [256, 128, 64, 32] + assert inner_k_tiles in [8, 4, 2] + orig_out_features, orig_in_features = input_float.shape + + # padding + in_features = find_multiple(orig_in_features, 1024) + out_features = find_multiple(orig_out_features, 8) + input_float = torch.nn.functional.pad( + input_float, + (0, in_features - orig_in_features, 0, out_features - orig_out_features), + ) + + # quantization and packing + input_int4x8, scales_and_zeros = groupwise_affine_quantize_tensor( + input_float, + 4, + groupsize, + dtype=input_float.dtype, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + ) + if check_cpu_version(input_float.device): + int_data = aten._convert_weight_to_int4pack_for_cpu( + input_int4x8, inner_k_tiles + ) + if check_xpu_version(input_float.device): + from torchao.quantization.utils import convert_weight_to_int4pack_xpu + + int_data = convert_weight_to_int4pack_xpu( + input_int4x8, + zero_point_domain_is_int=zero_point_domain == ZeroPointDomain.INT, + ) + else: + int_data = aten._convert_weight_to_int4pack(input_int4x8, inner_k_tiles) + return int_data, scales_and_zeros, False, groupsize, inner_k_tiles diff --git a/lib/python3.12/site-packages/torchao/quantization/transform_module.py b/lib/python3.12/site-packages/torchao/quantization/transform_module.py new file mode 100644 index 0000000000000000000000000000000000000000..339d46be3592f66dae77c6c0fe65ab5504cfb992 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/transform_module.py @@ -0,0 +1,52 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import functools +from typing import Callable, Dict + +import torch + +from torchao.core.config import AOBaseConfig + +_QUANTIZE_CONFIG_HANDLER: Dict[ + AOBaseConfig, + Callable[[torch.nn.Module, AOBaseConfig], torch.nn.Module], +] = {} + + +def register_quantize_module_handler(config_type): + """ + A decorator to register a transform function to map from a workflow + configuration (child of `AOBaseConfig`) to a function that transforms + a `torch.nn.Module` according to the specified configuration. + + For example:: + + # user facing code + class WorkflowFooConfig(AOBaseConfig): ... + # configuration for workflow `Foo` is defined here + bar = 'baz' + + # non user facing code + @register_quantize_module_handler(WorkflowFooConfig) + def _transform( + mod: torch.nn.Module, + config: WorkflowFooConfig, + ) -> torch.nn.Module: + # the transform is implemented here, usually a tensor sublass + # weight swap or a module swap + ... + + # then, the user calls `quantize_` with a config, and `_transform` is called + # under the hood by `quantize_. + + """ + + @functools.wraps(config_type) + def decorator(func): + _QUANTIZE_CONFIG_HANDLER[config_type] = func + return func # needed to make the functions usable externally + + return decorator diff --git a/lib/python3.12/site-packages/torchao/quantization/unified.py b/lib/python3.12/site-packages/torchao/quantization/unified.py new file mode 100644 index 0000000000000000000000000000000000000000..948d5f2fe843d7c461daa8ca14a9c8cb559480ee --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/unified.py @@ -0,0 +1,41 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from abc import ABC, abstractmethod +from typing import Any + +import torch + +""" +The vast majority of quantization algorithms follow one of two patterns +1. Single quantize call to create a quantized model with quantized state_dict +2. Flow that needs calibration or training + +This file defines the API for both patterns +""" + + +# API 1, single quantize call to create a quantized model with quantized state_dict +class Quantizer(ABC): + @abstractmethod + def quantize( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + pass + + +# API 2, flow that needs calibration or training +class TwoStepQuantizer: + @abstractmethod + def prepare( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + pass + + @abstractmethod + def convert( + self, model: torch.nn.Module, *args: Any, **kwargs: Any + ) -> torch.nn.Module: + pass diff --git a/lib/python3.12/site-packages/torchao/quantization/utils.py b/lib/python3.12/site-packages/torchao/quantization/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a9cad8060e906dcc4896b5a9d186c9834509538e --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/utils.py @@ -0,0 +1,649 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +import importlib.util +from typing import Dict, List, Optional + +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from torchao.kernel import ( + int_scaled_matmul, +) +from torchao.quantization.quant_primitives import ( + MappingType, + ZeroPointDomain, + choose_qparams_affine, + dequantize_affine, + quantize_affine, +) +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_5, + check_cpu_version, + check_xpu_version, +) + +__all__ = [ + "compute_error", + "_apply_logging_hook", + "quantize_activation_per_token_absmax", + "quant_int8_dynamic_per_token_linear", + "quant_int8_per_token_matmul", + "dynamically_quantize_per_channel", + "dequantize_per_tensor", + "dequantize_per_channel", + "get_groupwise_affine_qparams", + "pack_tinygemm_scales_and_zeros", + "unpack_tinygemm_scales_and_zeros", + "groupwise_affine_quantize_tensor_from_qparams", + "groupwise_affine_dequantize_tensor_from_qparams", + "groupwise_affine_quantize_tensor", + "groupwise_affine_dequantize_tensor", + "per_token_dynamic_quant", + "get_group_qparams_symmetric", + "recommended_inductor_config_setter", +] + +_lm_eval_available = importlib.util.find_spec("lm_eval") is not None + + +# basic SQNR +def compute_error(x, y): + Ps = torch.linalg.norm(x) + Pn = torch.linalg.norm(x - y) + return 20 * torch.log10(Ps / Pn) + + +# logger for fqn + op + shape +# note: not safe for any kind of multithreading +_cur_fqn: Optional[str] = None + + +def _get_logging_hook(fqn): + def forward_hook(module, input): + global _cur_fqn + _cur_fqn = fqn + + return forward_hook + + +def _apply_logging_hook(model): + for name, mod in model.named_modules(): + mod.register_forward_pre_hook(_get_logging_hook(name)) + + +# collections.defaultdict printing is weird with lambdas, so hand writing for now +_fqn_to_op_to_shape_to_count: Dict[ + Optional[str], Dict[Optional[str], Dict[Optional[str], int]] +] = {} + + +class LoggingTensorMode(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + rs = func(*args, **kwargs) + global _cur_fqn + op_name: str = f"{func.__module__}.{func.__name__}" + shape_str = "" + for arg in args: + if isinstance(arg, torch.Tensor): + shape_str += str(list(arg.shape)) + ", " + if shape_str != "": + shape_str = shape_str[:-2] + + if _cur_fqn not in _fqn_to_op_to_shape_to_count: + _fqn_to_op_to_shape_to_count[_cur_fqn] = {} + if op_name not in _fqn_to_op_to_shape_to_count[_cur_fqn]: + _fqn_to_op_to_shape_to_count[_cur_fqn][op_name] = {} + if shape_str not in _fqn_to_op_to_shape_to_count[_cur_fqn][op_name]: + _fqn_to_op_to_shape_to_count[_cur_fqn][op_name][shape_str] = 0 + _fqn_to_op_to_shape_to_count[_cur_fqn][op_name][shape_str] += 1 + + return rs + + +class _MultiInput: + def __init__(self, inputs): + self.values = list(inputs) + + def add_input(self, input): + self.values.append(input) + return self + + def __getitem__(self, slice): + return _MultiInput(self.values[slice]) + + def cuda(self): + self.values = [ + val.cuda() if isinstance(val, torch.Tensor) else val for val in self.values + ] + + +def guard_dtype_size(tensor_arg, arg_name, dtype=None, size=None): + if dtype is not None and tensor_arg.dtype != dtype: + raise ValueError( + f"Expected Tensor argument {arg_name} to have dtype {dtype}, but got {tensor_arg.dtype} instead." + ) + if size is not None and tensor_arg.size() != size: + raise ValueError( + f"Expected Tensor argument {arg_name} to have size {size}, but got {tensor_arg.size()} instead." + ) + + +def _get_per_token_block_size(x: torch.Tensor) -> List[int]: + block_size = [] + for _ in range(len(x.shape) - 1): + block_size.append(1) + block_size.append(x.shape[-1]) + return block_size + + +# taken from +# https://github.com/mit-han-lab/smoothquant/blob/2f87951dacfb9238d8d657f52ae83a82a3c9ba0c/smoothquant/fake_quant.py#L26 +# and slightly modified +def quantize_activation_per_token_absmax(t): + # if the shape of t is [B, N, K], the shape of scales will be [B, N, 1] + mapping_type = MappingType.SYMMETRIC + block_size = list(t.shape) + for i in range(len(block_size) - 1): + block_size[i] = 1 + dtype = torch.int8 + eps = 1e-5 + # Note: the original smoothquant does not clamp to qmin/qmax here, + # but some of the tests with bfloat16 ended up with a flipped sign + # if we don't clamp. TODO(future) look into this further. + quant_min = -127 + quant_max = 127 + scale_dtype = torch.float32 if t.dtype == torch.float16 else None + + scale, zero_point = choose_qparams_affine( + t, + mapping_type, + block_size, + dtype, + quant_min, + quant_max, + eps, + scale_dtype=scale_dtype, + ) + + quantized = quantize_affine( + t, block_size, scale, zero_point, dtype, quant_min, quant_max + ) + + return quantized, scale + + +def quant_int8_dynamic_per_token_linear( + x, + w_vals_int8_t, + w_scales, + bias, + out_dtype, +): + """ + like F.linear, but with int8 dynamic quantization of activation, + and a quantized weight + """ + x_vals_int8, x_scales = quantize_activation_per_token_absmax(x) + mm_out = quant_int8_per_token_matmul( + x_vals_int8, x_scales, w_vals_int8_t, w_scales, out_dtype + ) + if bias is not None: + mm_out = mm_out + bias + return mm_out + + +def quant_int8_per_token_matmul( + x_vals_int8, + x_scales, + w_vals_int8_t, + w_scales, + output_dtype=torch.float32, +): + """ + Quantized matmul of int8 operands that accumulates to int32 and returns + output_dtype. For now, this is written for approximate numerical + Assumes that activation and weight quantization are symmetric, + i.e. act_zp and w_zp is 0. + Assumes that weight quantization is per-channel. + + see + https://github.com/google/gemmlowp/blob/master/doc/quantization.md + for an overview of quantized matmul compute + + in scalar form, assuming output_dtype is fp32 and zw == 0: + + Y_i_j_fp32 = sx * sw dot(X_i, W_j) + """ + + assert x_vals_int8.dtype == torch.int8, ( + f"x dtype {x_vals_int8.dtype} not yet supported" + ) + assert w_vals_int8_t.dtype == torch.int8, ( + f"w dtype {w_vals_int8_t.dtype} not yet supported" + ) + + assert x_scales.dtype in [ + torch.float, + torch.bfloat16, + ], ( + f"x_scales needs to be a torch.float32 or torch.bfloat16 but got {x_scales.dtype}" + ) + + # + # 1. do the matrix form of dot(X_i, W_j) + # + # + # 2. rescale the output + # + # in cases with large matrices, y_dot_int32 can grow sufficiently + # large that y_dot_int32 * a float16 scale is greater than the maximum + # value of a float 16, (which results in a value of inf even if multiplying + # by the other scale would bring it within the expected range) + + tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1]) + y_dot_scaled = int_scaled_matmul(tmp, w_vals_int8_t, x_scales.reshape(-1, 1)) + + y = (y_dot_scaled * w_scales).reshape( + *x_vals_int8.shape[:-1], y_dot_scaled.shape[-1] + ) + + # can downcast only at the very end + y = y.to(output_dtype) + return y + + +def dynamically_quantize_per_channel(x, quant_min, quant_max, target_dtype): + """ + assumes symmetric quantization + assumes axis == 0 + assumes dense memory format + TODO(future): relax ^ as needed + """ + + assert x.dim() == 2, "only support 2d Tensors" + + eps = torch.finfo(torch.float32).eps + block_size = (1, x.shape[1]) + zero_point_dtype = torch.int64 + + mapping_type = MappingType.SYMMETRIC + scale, zero_point = choose_qparams_affine( + x, + mapping_type, + block_size, + target_dtype=target_dtype, + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + zero_point_dtype=zero_point_dtype, + ) + quant = quantize_affine( + x, block_size, scale, zero_point, target_dtype, quant_min, quant_max + ) + return quant, scale, zero_point + + +# reference: https://fburl.com/code/vfsygwd0 +def dequantize_per_tensor(int_repr, scale, zero_point, out_dtype=torch.float32): + block_size = int_repr.shape + input_dtype = int_repr.dtype + assert scale.numel() == 1, f"scale size: {scale.numel()}" + dequantized = dequantize_affine( + int_repr, block_size, scale, zero_point, input_dtype, output_dtype=out_dtype + ) + return dequantized + + +# reference: https://fburl.com/code/org0fmi3 +def dequantize_per_channel(int_repr, scales, zero_points, out_dtype=torch.float32): + assert int_repr.dim() == 2, "only support 2d Tensors" + # channel axis == 0 + # block_size before transpose should be (1, int_repr.shape[1]) for axis == 0 per channel quant + + # TODO: transpose is for perf reasons for torch.compile, we should separate this to lowering step + int_repr = int_repr.t() + # transpose for block_size as well + block_size = (int_repr.shape[0], 1) + input_dtype = int_repr.dtype + dequantized = dequantize_affine( + int_repr, block_size, scales, zero_points, input_dtype, output_dtype=out_dtype + ) + dequantized = dequantized.t() + return dequantized + + +def get_groupwise_affine_qparams( + w, + n_bit=4, + groupsize=128, + dtype=torch.bfloat16, + zero_point_domain=ZeroPointDomain.FLOAT, + preserve_zero=False, +): + if groupsize > w.shape[-1]: + groupsize = w.shape[-1] + assert groupsize > 1 + assert w.shape[-1] % groupsize == 0 + assert w.dim() == 2 + assert n_bit <= 8, f"only n_bit smaller than 8 is supported, got: {n_bit}" + + mapping_type = MappingType.ASYMMETRIC + target_dtype = torch.int32 + block_size = (1, groupsize) + quant_min = 0 + quant_max = 2**n_bit - 1 + eps = 1e-6 + scale_dtype = dtype + zero_point_dtype = ( + dtype if zero_point_domain != ZeroPointDomain.INT else torch.int32 + ) + + scale, zero_point = choose_qparams_affine( + w, + mapping_type, + block_size, + target_dtype, + quant_min, + quant_max, + eps, + scale_dtype=scale_dtype, + zero_point_dtype=zero_point_dtype, + preserve_zero=preserve_zero, + zero_point_domain=zero_point_domain, + ) + + return scale.to(dtype=dtype).reshape(w.shape[0], -1), zero_point.to( + dtype=zero_point_dtype + ).reshape(w.shape[0], -1) + + +def pack_tinygemm_scales_and_zeros(scales, zeros, dtype=torch.bfloat16): + guard_dtype_size(scales, "scales", dtype=dtype, size=zeros.size()) + guard_dtype_size(zeros, "zeros", dtype=dtype) + dim = scales.dim() + return ( + torch.cat( + [ + scales.unsqueeze(-1), + zeros.unsqueeze(-1), + ], + dim, + ) + .transpose(-3, -2) + .contiguous() + ) + + +def unpack_tinygemm_scales_and_zeros(scales_and_zeros): + assert scales_and_zeros.shape[-1] == 2 + return torch.split(scales_and_zeros.transpose(-3, -2), 1, -1) + + +def convert_weight_to_int4pack_xpu(weight, zero_point_domain_is_int=False): + assert weight.device.type == "xpu" + + if zero_point_domain_is_int: + # int_data = weight.to(dtype=torch.uint8) + int_data = (weight[::, 1::2] << 4 | weight[::, ::2]).to(torch.uint8) + packed_weight = torch.ops.aten._convert_weight_to_int4pack( + int_data, + 8, # TODO:remove + ) + else: + out = weight.to(dtype=torch.uint8) + out = (out[::, 1::2] << 4 | out[::, ::2]).to(torch.uint8) + packed_weight = out.view(torch.int32) + + # Second, N * K/2 uint8 -> N * K/8 int32 + return packed_weight + + +def groupwise_affine_quantize_tensor_from_qparams( + w, scales, zeros, n_bit=4, groupsize=128, zero_point_domain=ZeroPointDomain.FLOAT +): + assert groupsize > 1 + # needed for GPTQ single column quantize + if groupsize > w.shape[-1] and scales.shape[-1] == 1: + groupsize = w.shape[-1] + + assert w.shape[-1] % groupsize == 0 + assert w.dim() == 2 + + block_size = (1, groupsize) + output_dtype = torch.int32 + quant_min = 0 + quant_max = 2**n_bit - 1 + + int_data = quantize_affine( + w, + block_size, + scales, + zeros, + output_dtype, + quant_min, + quant_max, + zero_point_domain=zero_point_domain, + ) + if TORCH_VERSION_AT_LEAST_2_5 and w.shape[-1] > 1: + if (not (check_cpu_version(int_data.device))) and ( + not (check_xpu_version(int_data.device)) + ): + int_data = (int_data[::, ::2] << 4 | int_data[::, 1::2]).to(torch.uint8) + return int_data + + +def groupwise_affine_dequantize_tensor_from_qparams( + w_int4x8, + scales, + zeros, + n_bit=4, + groupsize=128, + zero_point_domain=ZeroPointDomain.FLOAT, +): + assert groupsize > 1 + assert w_int4x8.dim() == 2 + # need to handle single column case so check for dtype/size from groupwise_affine_quantize_tensor_from_qparams path + if ( + TORCH_VERSION_AT_LEAST_2_5 + and (w_int4x8.dtype == torch.uint8 or w_int4x8.shape[-1] > 1) + and not (check_cpu_version(w_int4x8.device)) + and not (check_xpu_version(w_int4x8.device)) + ): + data = w_int4x8.to(torch.int32) + high_bits = data >> 4 + low_bits = data & 0x0F + w_int32 = torch.zeros( + (w_int4x8.shape[0], w_int4x8.shape[1] * 2), + dtype=torch.int32, + device=w_int4x8.device, + ) + w_int32[::, ::2] = high_bits + w_int32[::, 1::2] = low_bits + else: + w_int32 = w_int4x8 + + # needed for GPTQ single column dequantize + if groupsize > w_int32.shape[-1] and scales.shape[-1] == 1: + groupsize = w_int32.shape[-1] + assert w_int32.shape[-1] % groupsize == 0 + block_size = (1, groupsize) + input_dtype = torch.int32 + quant_min = 0 + quant_max = 2**n_bit - 1 + return dequantize_affine( + w_int32, + block_size, + scales, + zeros, + input_dtype, + quant_min, + quant_max, + zero_point_domain=zero_point_domain, + output_dtype=scales.dtype, + ) + + +def groupwise_affine_quantize_tensor( + w, + n_bit=4, + groupsize=128, + dtype=torch.bfloat16, + zero_point_domain=ZeroPointDomain.FLOAT, + preserve_zero=False, +): + scales, zeros = get_groupwise_affine_qparams( + w, + n_bit, + groupsize, + dtype, + zero_point_domain=zero_point_domain, + preserve_zero=preserve_zero, + ) + w_int4x8 = groupwise_affine_quantize_tensor_from_qparams( + w, scales, zeros, n_bit, groupsize, zero_point_domain=zero_point_domain + ) + scales_and_zeros = pack_tinygemm_scales_and_zeros(scales, zeros, dtype) + return w_int4x8, scales_and_zeros + + +def groupwise_affine_dequantize_tensor( + w_int4x8, + scales_and_zeros, + n_bit=4, + groupsize=128, +): + scales, zeros = unpack_tinygemm_scales_and_zeros(scales_and_zeros) + return groupwise_affine_dequantize_tensor_from_qparams( + w_int4x8, scales, zeros, n_bit, groupsize + ) + + +# TODO: separate scale and zero point precision +def get_group_qparams_symmetric( + w, + n_bit=4, + groupsize=128, + precision=torch.float32, + mapping_type=MappingType.SYMMETRIC, +): + # needed for GPTQ with padding + if groupsize > w.shape[-1]: + groupsize = w.shape[-1] + assert groupsize > 1 + assert w.shape[-1] % groupsize == 0 + assert w.dim() == 2 + assert n_bit <= 8, f"unsupported n_bit: {n_bit}" + + block_size = (1, groupsize) + eps = torch.finfo(w.dtype).eps + ranges = {} + ranges[1] = (-1, 0) + # generating ranges for bit 2 to 8 + for i in range(2, 9): + ranges[i] = (-(2 ** (i - 1)), 2 ** (i - 1) - 1) + quant_min, quant_max = ranges[n_bit] + scale, zero_point = choose_qparams_affine( + w, + mapping_type, + block_size, + target_dtype=torch.int8, + quant_min=quant_min, + quant_max=quant_max, + eps=eps, + scale_dtype=precision, + zero_point_dtype=precision, + ) + return scale.reshape(w.shape[0], -1), zero_point.reshape(w.shape[0], -1) + + +def group_quantize_tensor_symmetric( + w, + n_bit=4, + group_size=128, + precision=torch.float32, + mapping_type=MappingType.SYMMETRIC, +): + scales, zeros = get_group_qparams_symmetric( + w, n_bit, group_size, precision, mapping_type + ) + n_bit = 4 + max_int = 2 ** (n_bit - 1) - 1 + min_int = -(2 ** (n_bit - 1)) + # TODO: currently we don't know how to express torch.int4, we'll + # add torch.int4 to core later + from torchao._executorch_ops import ( + _quantized_decomposed_quantize_per_channel_group_wrapper, + ) + + w_int8 = _quantized_decomposed_quantize_per_channel_group_wrapper( + w, scales, zeros, min_int, max_int, torch.int8, group_size + ) + + return w_int8, scales, zeros + + +def per_token_dynamic_quant( + input: torch.Tensor, + scale_dtype: torch.dtype = torch.float32, + zero_point_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + mapping_type = MappingType.ASYMMETRIC + block_size = _get_per_token_block_size(input) + quant_min = -128 + quant_max = 127 + quant_dtype = torch.int8 + output_dtype = input.dtype + + scales, zero_points = choose_qparams_affine( + input, + mapping_type, + block_size, + quant_dtype, + quant_min, + quant_max, + scale_dtype=scale_dtype, + zero_point_dtype=zero_point_dtype, + ) + q = quantize_affine( + input, + block_size, + scales, + zero_points, + quant_dtype, + quant_min, + quant_max, + ) + dq = dequantize_affine( + q, + block_size, + scales, + zero_points, + quant_dtype, + quant_min, + quant_max, + output_dtype=output_dtype, + ) + return dq + + +def recommended_inductor_config_setter(): + """ + Set inductor config to use the following optimizations which have been showed to improve performance for quantized models: + coordinate_descent_tuning = True + coordinate_descent_check_all_directions = True + force_fuse_int_mm_with_mul = True + fx_graph_cache = True + triton.unique_kernel_names = True + torch.set_float32_matmul_precision("high") + """ + torch._inductor.config.coordinate_descent_tuning = True + torch._inductor.config.coordinate_descent_check_all_directions = True + torch._inductor.config.force_fuse_int_mm_with_mul = True + torch._inductor.config.fx_graph_cache = True + torch._inductor.config.triton.unique_kernel_names = True + torch.set_float32_matmul_precision("high") diff --git a/lib/python3.12/site-packages/torchao/quantization/weight_only.py b/lib/python3.12/site-packages/torchao/quantization/weight_only.py new file mode 100644 index 0000000000000000000000000000000000000000..fb30c14936f04d8f1ae7577d245d803318ac5bed --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/weight_only.py @@ -0,0 +1,100 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from .utils import dynamically_quantize_per_channel + +__all__ = ["WeightOnlyInt8QuantLinear"] + + +class WeightOnlyInt8QuantLinear(torch.nn.Linear): + """ + This class is a replacement for `torch.nn.Linear`. It implements a + mixed dtype matrix multiplication using int8 symmetric per-channel weight quantization. + + The primary goal of this class is to leverage int8 quantization for weights to reduce the + memory footprint and computational requirements while performing linear transformations. + This can be particularly beneficial for deploying models in low latency environments + + Attributes: + w_int8 (torch.Tensor): The quantized weights in int8 format. + scales (torch.Tensor): The scaling factors for each channel to convert the quantized + weights back to floating point format during the forward pass. + """ + + def __init__(self, *args, **kwargs): + """ + Initializes the WeightOnlyInt8QuantLinear module. + + Args: + *args: Variable length argument list for `torch.nn.Linear`. + **kwargs: Arbitrary keyword arguments. + Must include 'w_int8' (int8 quantized weights) and 'scales' (scaling factors). + """ + w_int8 = kwargs.pop("w_int8") + scales = kwargs.pop("scales") + super().__init__(*args, **kwargs) + + self.register_buffer("w_int8", w_int8) + self.register_buffer("scales", scales) + + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + """ + Performs the forward pass of the quantized linear layer, which consists of + mixed dtype matrix multiplication using int8 symmetric per-channel weight quantization. + + Args: + x (torch.Tensor): The input floating point tensor to the quantized linear layer. + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + torch.Tensor: The output floating point tensor after the quantized matrix multiplication + and rescale. + """ + x_view = x.view(-1, x.shape[-1]) + y = torch.mm(x_view, self.w_int8.to(x.dtype)) * self.scales + y = y.reshape(*x.shape[:-1], -1) + if self.bias is not None: + y += self.bias + return y + + @classmethod + def from_float(cls, mod: torch.nn.Linear): + """ + Converts a `torch.nn.Linear` module to a `WeightOnlyInt8QuantLinear` module. + + This method performs the conversion by dynamically quantizing the weights of the original + floating point linear layer to int8 format and creating a new `WeightOnlyInt8QuantLinear` + instance with these quantized weights and the corresponding scaling factors. + + Args: + mod (torch.nn.Linear): The original `torch.nn.Linear` module to convert. + + Returns: + WeightOnlyInt8QuantLinear: The converted quantized linear module with int8 weights. + """ + w_fp32 = mod.weight + w_int8, scales, _zp = dynamically_quantize_per_channel( + w_fp32, -128, 127, torch.int8 + ) + # Create the new module with a toy size to ensure initialization is fast + fake_in_features, fake_out_features = 8, 8 + new_mod = cls( + fake_in_features, + fake_out_features, + bias=mod.bias is not None, + w_int8=w_int8.t().contiguous(), + scales=scales, + ) + new_mod.in_features = mod.in_features + new_mod.out_features = mod.out_features + del new_mod.weight + new_mod.bias = mod.bias + device_to_use = next(mod.parameters()).device + new_mod.to(device_to_use) + return new_mod diff --git a/lib/python3.12/site-packages/torchao/quantization/weight_tensor_linear_activation_quantization.py b/lib/python3.12/site-packages/torchao/quantization/weight_tensor_linear_activation_quantization.py new file mode 100644 index 0000000000000000000000000000000000000000..6612213bc1bb0b8142faec789c976d2fedbfda9b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/quantization/weight_tensor_linear_activation_quantization.py @@ -0,0 +1,208 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import Any, Callable, Dict, Optional + +import torch +from torch.utils._python_dispatch import return_and_correct_aliasing + +from torchao.utils import ( + TORCH_VERSION_AT_LEAST_2_5, + TorchAOBaseTensor, +) + +__all__ = [ + "WeightTensorWithLinearActivationQuantizationMetadata", + "to_weight_tensor_with_linear_activation_quantization_metadata", +] + +aten = torch.ops.aten + + +class WeightTensorWithLinearActivationQuantizationMetadata(TorchAOBaseTensor): + """ + Tensor subclass that wraps a weight tensor and provides metadata for linear activation static quantization. + + Args: + original_weight_tensor (torch.Tensor): The weight tensor to be wrapped. + input_quant_func_static (Callable): The quantization function for inputs. + Must have the signature: (Tensor, scale: Tensor, zero_point: Optional[Tensor], **quant_kwargs) -> Tensor + scale (torch.Tensor): The scale tensor for activation quantization. + zero_point (Optional[torch.Tensor]): The zero point tensor for activation quantization. Default is None. + quant_kwargs (Dict[str, Any]): Additional keyword arguments for the quantization function. + Restriction: Must not contain tensor values. + """ + + original_weight_tensor: torch.Tensor + input_quant_func_static: Callable + scale: torch.Tensor + zero_point: Optional[torch.Tensor] + quant_kwargs: Dict[str, Any] + + def __new__( + cls, + original_weight_tensor: torch.Tensor, + input_quant_func_static: Callable, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_kwargs: Dict[str, Any], + ): + kwargs = {} + dtype = original_weight_tensor.dtype + kwargs["dtype"] = dtype + kwargs["requires_grad"] = False + kwargs["device"] = original_weight_tensor.device + shape = original_weight_tensor.shape + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] + + def __init__( + self, + original_weight_tensor: torch.Tensor, + input_quant_func_static: Callable[ + [torch.Tensor, torch.Tensor, Optional[torch.Tensor], Dict[str, Any]], + torch.Tensor, + ], + scale: torch.Tensor, + zero_point: Optional[torch.Tensor], + quant_kwargs: Dict[str, Any], + ): + self.original_weight_tensor = original_weight_tensor + self.input_quant_func_static = input_quant_func_static + self.scale = scale + self.zero_point = zero_point + self.quant_kwargs = quant_kwargs + + def __repr__(self): + return f"{self.__class__.__name__}({self.original_weight_tensor}, {self.input_quant_func_static}, scale={self.scale}, zero_point={self.zero_point}, quant_kwargs={self.quant_kwargs})" + + def __tensor_flatten__(self): + tensor_data = ["original_weight_tensor", "scale"] + if self.zero_point is not None: + tensor_data.append("zero_point") + return tensor_data, [self.input_quant_func_static, self.quant_kwargs] + + @classmethod + def __tensor_unflatten__( + cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride + ): + original_weight_tensor = tensor_data_dict["original_weight_tensor"] + input_quant_func_static, quant_kwargs = tensor_attributes + zero_point = tensor_data_dict.get("zero_point", None) + return cls( + original_weight_tensor, + input_quant_func_static, + tensor_data_dict["scale"], + zero_point, + quant_kwargs, + ) + + @staticmethod + def _quantized_linear_op( + input_tensor: torch.Tensor, weight_tensor: torch.Tensor, bias: torch.Tensor + ): + input_quant_func_static = weight_tensor.input_quant_func_static + original_weight_tensor = weight_tensor.original_weight_tensor + scale = weight_tensor.scale + zero_point = weight_tensor.zero_point + quant_kwargs = weight_tensor.quant_kwargs + quantized_input_act = input_quant_func_static( + input_tensor, scale=scale, zero_point=zero_point, **quant_kwargs + ) + return torch.nn.functional.linear( + quantized_input_act, original_weight_tensor, bias + ) + + @classmethod + def from_float( + cls, + input_float: torch.Tensor, + input_quant_func: Callable, + scale: torch.Tensor, + zero_point: Optional[torch.Tensor] = None, + quant_kwargs: Optional[Dict[str, Any]] = None, + ): + if quant_kwargs is None: + quant_kwargs = {} + return cls(input_float, input_quant_func, scale, zero_point, quant_kwargs) + + def _apply_fn_to_data(self, fn): + return self.__class__( + fn(self.original_weight_tensor), + self.input_quant_func_static, + fn(self.scale), + fn(self.zero_point) if self.zero_point is not None else None, + self.quant_kwargs, + ) + + def to(self, *args, **kwargs): + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs.pop("device") + return self.__class__( + self.original_weight_tensor.to(device), + self.input_quant_func_static, + self.scale.to(device), + self.zero_point.to(device) if self.zero_point is not None else None, + self.quant_kwargs, + ) + + +implements = WeightTensorWithLinearActivationQuantizationMetadata.implements + + +@implements(torch.nn.functional.linear) +def _(func, types, args, kwargs): + input_tensor, weight_tensor, bias = ( + args[0], + args[1], + args[2] if len(args) > 2 else None, + ) + if isinstance(weight_tensor, WeightTensorWithLinearActivationQuantizationMetadata): + return weight_tensor._quantized_linear_op(input_tensor, weight_tensor, bias) + + raise NotImplementedError( + "LinearActivationQuantizedTensor: No specialized dispatch found for linear op" + ) + + +@implements(aten.detach.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.detach) + ) + + +@implements(aten.clone.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.clone) + ) + + +@implements(aten._to_copy.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, + args, + kwargs, + args[0].to(*args[1:], **kwargs)._apply_fn_to_data(torch.clone), + ) + + +@implements(aten.t.default) +def _(func, types, args, kwargs): + return return_and_correct_aliasing( + func, args, kwargs, args[0]._apply_fn_to_data(torch.t) + ) + + +to_weight_tensor_with_linear_activation_quantization_metadata = ( + WeightTensorWithLinearActivationQuantizationMetadata.from_float +) + +if TORCH_VERSION_AT_LEAST_2_5: + # Allow a model with LinearActivationQuantizedTensor weights to be loaded with `weights_only=True` + torch.serialization.add_safe_globals( + [WeightTensorWithLinearActivationQuantizationMetadata] + ) diff --git a/lib/python3.12/site-packages/torchao/sparsity/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ccb42d9946b70849b05856f9a3d51b581bdffffe Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/__pycache__/blocksparse.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/blocksparse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d2152857826348da910ebcc6a5252fb555d61be Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/blocksparse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/__pycache__/sparse_api.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/sparse_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e299b3fbb9ba4407ecbf66a3207c709c566533fc Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/sparse_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/__pycache__/supermask.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/supermask.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cfb70e8cbe00650a29afc498621664b46d9706f5 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/supermask.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a10ad9107f2276f231d44e25ed325dd8d17777f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/__pycache__/wanda.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/wanda.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5609abfef815981acc52411bedd417a0ea56d1de Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/__pycache__/wanda.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d01781e4d51b9a5306a4dd19c6d2be3356917b4b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/pruner/__pycache__/FPGM_pruner.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/pruner/__pycache__/FPGM_pruner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8490a2ac68cfeafe74be0b527d265f4e63aded53 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/pruner/__pycache__/FPGM_pruner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__init__.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e664bdd1a70a5e3127624736dea4bbc2ddfeef81 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/base_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/base_scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5e4f27c071b4c220aaef637586f71be9e9e973d Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/base_scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/cubic_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/cubic_scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14bce0980156ef586f3fa4af958f7442e39a45dc Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/cubic_scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/lambda_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/lambda_scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7998c59bc86b37b1966731ff94c37b3ca01281a5 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/__pycache__/lambda_scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/base_scheduler.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/base_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..30a29c47753c97de77bb4586ab7a85d4a75102db --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/base_scheduler.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.scheduler.base_scheduler import BaseScheduler + +__all__ = ["BaseScheduler"] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/cubic_scheduler.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/cubic_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..6ea7e95c36ea3db6864676a45d7dd2e1eabc2070 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/cubic_scheduler.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.scheduler.cubic_scheduler import CubicSL + +__all__ = ["CubicSL"] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/lambda_scheduler.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/lambda_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..97d0abf174e7323f0f941ba7b2091924c39897c9 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/scheduler/lambda_scheduler.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.scheduler.lambda_scheduler import LambdaSL + +__all__ = ["LambdaSL"] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__init__.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b9388fe5d1b6882eabae6308fa50d41e8ec7907 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/base_sparsifier.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/base_sparsifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad0968c2caa08b8a64f2c597686f98cf785f5d5b Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/base_sparsifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/nearly_diagonal_sparsifier.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/nearly_diagonal_sparsifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b9a643ef64339bccc1e4dd3b1e1793d36e28f63f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/nearly_diagonal_sparsifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd4bb24b436ae2ca9620a031f119ccec77325db1 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/weight_norm_sparsifier.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/weight_norm_sparsifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2a20a71b65af8bfd0fd143f6079030f73159955 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/__pycache__/weight_norm_sparsifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/base_sparsifier.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/base_sparsifier.py new file mode 100644 index 0000000000000000000000000000000000000000..3c9f0947d8a9fa688bc933d274f0adf7c0ae224a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/base_sparsifier.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.sparsifier.base_sparsifier import BaseSparsifier + +__all__ = ["BaseSparsifier"] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/nearly_diagonal_sparsifier.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/nearly_diagonal_sparsifier.py new file mode 100644 index 0000000000000000000000000000000000000000..e6caa9087d27af48e62b04f8fde32e0c20fd6be8 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/nearly_diagonal_sparsifier.py @@ -0,0 +1,10 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.sparsifier.nearly_diagonal_sparsifier import ( + NearlyDiagonalSparsifier, +) + +__all__ = ["NearlyDiagonalSparsifier"] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/utils.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8cb822513e960858afdd1643409f8efcd2917b23 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/utils.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.sparsifier.utils import FakeSparsity + +__all__ = ["FakeSparsity"] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/weight_norm_sparsifier.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/weight_norm_sparsifier.py new file mode 100644 index 0000000000000000000000000000000000000000..93d27993160ff7bc14508a3e3000eab0802dc84e --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/sparsifier/weight_norm_sparsifier.py @@ -0,0 +1,10 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.sparsifier.weight_norm_sparsifier import ( + WeightNormSparsifier, +) + +__all__ = ["WeightNormSparsifier"] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__init__.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..daae398a98c4485855fcc1797791f0bfb69e289f Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/blocksparse.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/blocksparse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1250bef7c253819d043bf0035a653e3ee8e266dd Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/blocksparse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/supermask.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/supermask.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c2056975d56e2faa092d18a36355b6a75abb463 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/supermask.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..813a498c50650ea9b43abdba951042c6fc21b2bd Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/blocksparse.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/blocksparse.py new file mode 100644 index 0000000000000000000000000000000000000000..845be888fc8ad4bff604e72085f091a6ec36269a --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/blocksparse.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.superblock.blocksparse import BlockSparseTensor + +__all__ = ["BlockSparseTensor"] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/supermask.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/supermask.py new file mode 100644 index 0000000000000000000000000000000000000000..1066271013132d012ca950caaa14140a1b0b8a7e --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/supermask.py @@ -0,0 +1,12 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.sparsity.supermask import ( + SupermaskLinear, +) + +__all__ = [ + "SupermaskLinear", +] diff --git a/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/utils.py b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e0ef628ebc2bc3932d18ea5e61762d7924c4e991 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/prototype/superblock/utils.py @@ -0,0 +1,26 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.prototype.sparsity.superblock.utils import ( + ClassificationPresetEval, + ClassificationPresetTrain, + ExponentialMovingAverage, + MetricLogger, + RandomCutmix, + RandomMixup, + RASampler, + SmoothedValue, +) + +__all__ = [ + "ClassificationPresetEval", + "ClassificationPresetTrain", + "ExponentialMovingAverage", + "MetricLogger", + "RandomCutmix", + "RandomMixup", + "RASampler", + "SmoothedValue", +] diff --git a/lib/python3.12/site-packages/torchao/sparsity/training/__init__.py b/lib/python3.12/site-packages/torchao/sparsity/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3c4212101b752d4b914d3e604ace57bbcbf0a768 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/training/__init__.py @@ -0,0 +1,105 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# +# This source code is licensed under the BSD license found in the +# LICENSE file in the root directory of this source tree. +import torch + +from torchao.sparsity.training.autograd import semi_structured_sparsify +from torchao.sparsity.training.pointwise_ops import CUTLASS_POINTWISE_OP_DISPATCH_TABLE +from torchao.utils import TORCH_VERSION_AT_LEAST_2_3 + +# load pointwise op support, which exists only for CUTLASS +if TORCH_VERSION_AT_LEAST_2_3: + from torch.sparse import SparseSemiStructuredTensorCUTLASS + + SparseSemiStructuredTensorCUTLASS._load_dispatch_table( + CUTLASS_POINTWISE_OP_DISPATCH_TABLE + ) + +__all__ = [ + "SemiSparseLinear", + "SemiSparseActivationLinear", + "swap_linear_with_semi_sparse_linear", + "swap_semi_sparse_linear_with_linear", +] + + +class SemiSparseLinear(torch.nn.Linear): + """ + Replacement nn.Linear that supports runtime weight sparsity + """ + + def forward(self, x): + sparse_weight = semi_structured_sparsify(self.weight, backend="cusparselt") + return torch.nn.functional.linear(x, sparse_weight, self.bias) + + @classmethod + def from_dense(cls, linear): + mod = cls(linear.in_features, linear.out_features) + mod.weight = linear.weight + mod.bias = linear.bias + return mod + + @classmethod + def to_dense(cls, semi_sparse_linear): + mod = torch.nn.Linear( + semi_sparse_linear.in_features, semi_sparse_linear.out_features + ) + mod.weight = semi_sparse_linear.weight + mod.bias = semi_sparse_linear.bias + return mod + + +class SemiSparseActivationLinear(torch.nn.Linear): + """ + Replacement nn.Linear that supports runtime activation sparsity + """ + + def forward(self, x): + sparse_x = semi_structured_sparsify(x, backend="cusparselt") + return torch.nn.functional.linear(sparse_x, self.weight, self.bias) + + @classmethod + def from_dense(cls, linear): + mod = cls(linear.in_features, linear.out_features) + mod.weight = linear.weight + mod.bias = linear.bias + return mod + + @classmethod + def to_dense(cls, semi_sparse_linear): + mod = torch.nn.Linear( + semi_sparse_linear.in_features, semi_sparse_linear.out_features + ) + mod.weight = semi_sparse_linear.weight + mod.bias = semi_sparse_linear.bias + return mod + + +def swap_linear_with_semi_sparse_linear(model, config, current=""): + """ + Public API for replacing nn.Linear with SemiSparseLinear + """ + name_to_child = dict(model.named_children()) + for name, child in name_to_child.items(): + fqn = f"{current}.{name}" if current else name + if isinstance(child, torch.nn.Linear): + if fqn in config: + setattr(model, name, config[fqn].from_dense(child)) + del child + else: + swap_linear_with_semi_sparse_linear(child, config, current=fqn) + + +def swap_semi_sparse_linear_with_linear(model, current=""): + """ + Public API for replacing instances of SemiSparseLinear/SemiSparseActivaitonLinear with nn.Linear + """ + name_to_child = dict(model.named_children()) + for name, child in name_to_child.items(): + fqn = f"{current}.{name}" if current else name + if isinstance(child, (SemiSparseLinear, SemiSparseActivationLinear)): + setattr(model, name, child.to_dense(child)) + del child + else: + swap_semi_sparse_linear_with_linear(child, current=fqn) diff --git a/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..470540a88b6ee3cee64f3e9b7436b9ed92f5b3a4 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/autograd.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/autograd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60f1f07e3d21616b36155821aabacedd6ee23ef0 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/autograd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/pointwise_ops.cpython-312.pyc b/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/pointwise_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a70b43f3d0101307fff18cdc69983624e338feff Binary files /dev/null and b/lib/python3.12/site-packages/torchao/sparsity/training/__pycache__/pointwise_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/sparsity/training/autograd.py b/lib/python3.12/site-packages/torchao/sparsity/training/autograd.py new file mode 100644 index 0000000000000000000000000000000000000000..fafbd7c3c38c4701be1d1f4bd2ebb95607fb2ff0 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/training/autograd.py @@ -0,0 +1,161 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from enum import Enum + +import torch +from torch.sparse import SparseSemiStructuredTensor + +from torchao.utils import TORCH_VERSION_AT_LEAST_2_3 + +if TORCH_VERSION_AT_LEAST_2_3: + from torch.sparse import ( + SparseSemiStructuredTensorCUSPARSELT, + SparseSemiStructuredTensorCUTLASS, + ) + + torch._dynamo.allow_in_graph(SparseSemiStructuredTensorCUSPARSELT) + torch._dynamo.allow_in_graph(SparseSemiStructuredTensorCUTLASS) + + +GRADIENT_TYPE = Enum("GRADIENT_TYPE", ["DENSE", "SPARSE", "STE"]) + + +class _SparsifyFunc(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, algo: str, backend: GRADIENT_TYPE): # type: ignore[override] + use_cutlass = backend == "cutlass" + if not isinstance(x, SparseSemiStructuredTensor): + (packed, meta, packed_t, meta_t, bitmask) = ( + torch._sparse_semi_structured_tile( + x, algorithm=algo, use_cutlass=use_cutlass + ) + ) + cls = ( + SparseSemiStructuredTensorCUTLASS + if use_cutlass + else SparseSemiStructuredTensorCUSPARSELT + ) + out = cls( + x.shape, + packed=packed, + meta=meta, + packed_t=packed_t, + meta_t=meta_t, + compressed_swizzled_bitmask=bitmask, + requires_grad=False, + fuse_transpose_cusparselt=True, + ) + else: + out = x.detach() + + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + # We just return grad_out, since we just use STE - straight through estimation + return grad_out, None, None + + +class _SparsifyLikeFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x: torch.Tensor, + pattern: SparseSemiStructuredTensor, + gradient=GRADIENT_TYPE.SPARSE, + ): # type: ignore[override] + assert isinstance(pattern, SparseSemiStructuredTensor) + + if not isinstance(pattern, SparseSemiStructuredTensorCUTLASS): + raise NotImplementedError( + "`sparsify_like(x, pattern)` is only implemented for CUTLASS backend" + ) + if not pattern.compressed_swizzled_bitmask.is_contiguous(): + raise NotImplementedError( + "`sparsify_like(x, pattern)` is not implemented when `bitmask` is transposed" + ) + + packed, packed_t = torch._sparse_semi_structured_apply( + x, pattern.compressed_swizzled_bitmask + ) + + # save for backwards + ctx.meta = pattern.meta + ctx.meta_t = pattern.meta_t + ctx.bitmask = pattern.compressed_swizzled_bitmask + ctx.gradient = gradient + + return pattern.__class__( + x.shape, + packed, + pattern.meta, + packed_t, + pattern.meta_t, + pattern.compressed_swizzled_bitmask, + requires_grad=x.requires_grad, + ) + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + if ctx.gradient == GRADIENT_TYPE.STE or isinstance( + grad_out, SparseSemiStructuredTensor + ): + return grad_out, None, None, None + assert not isinstance(grad_out, SparseSemiStructuredTensor) + assert grad_out.dtype == ctx.dtype + + if ctx.gradient == GRADIENT_TYPE.DENSE: + assert ctx.threads_masks.is_contiguous() + return ( + torch._sparse_semi_structured_apply_dense(grad_out, ctx.bitmask), + None, + None, + None, + ) + assert ctx.gradient == GRADIENT_TYPE.SPARSE + + packed, _, packed_t, _ = torch._sparse_semi_structured_tile( + grad_out, ctx.bitmask, backend="cutlass" + ) + return ( + SparseSemiStructuredTensorCUTLASS( + grad_out.shape, + packed, + ctx.meta, + packed_t, + ctx.meta_t, + ctx.bitmask, + requires_grad=grad_out.requires_grad, + ), + None, + None, + None, + ) + return grad_out, None + + +@torch._dynamo.allow_in_graph +def semi_structured_sparsify( + x: torch.Tensor, + algo: str = "", + backend: str = "cutlass", +) -> SparseSemiStructuredTensor: + """ + Sparsifies a dense tensor into a semi-structured tensor, according to the algo and backend passed. + """ + return _SparsifyFunc.apply(x, algo, backend) + + +@torch._dynamo.allow_in_graph +def semi_structured_sparsify_like( + x: torch.Tensor, + pattern: SparseSemiStructuredTensor, + gradient: GRADIENT_TYPE = GRADIENT_TYPE.SPARSE, +) -> SparseSemiStructuredTensor: + """ + Sparsifies a dense tensor into a semi-structured tensor, using the mask of the provided pattern. + """ + return _SparsifyLikeFunc.apply(x, pattern, gradient) diff --git a/lib/python3.12/site-packages/torchao/sparsity/training/pointwise_ops.py b/lib/python3.12/site-packages/torchao/sparsity/training/pointwise_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ec483fb0946db0e5121eecc8c8468b5acc53e727 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/sparsity/training/pointwise_ops.py @@ -0,0 +1,108 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from functools import partial + +import torch +from torch.sparse import SparseSemiStructuredTensor + +from torchao.sparsity.training.autograd import semi_structured_sparsify_like + + +def _semi_sparse_pointwise_op( + func, types, args=(), kwargs=None, sparsify_like_args_list=() +): + """ + adds pointwise op support for semi-structured tensors. + + Assumes that at least one of the arguments in arg is a SparseSemiStructuredTensor. + The last instance of a SparseSemiStructuredTensor is used as the reference mask to sparsify the others tensors passed in args. + sparsify_like_args_list is used to specify which arguments to sparsify like the reference tensor. + """ + reference_sparse_tensor = None + for tensor in args: + if isinstance(tensor, SparseSemiStructuredTensor): + reference_sparse_tensor = tensor + assert reference_sparse_tensor is not None + + def handle_arg(i, tensor): + if isinstance(tensor, torch.Tensor): + # For pointwise ops, dense tensors will be sparsified to match the sparsity pattern of the reference tensor + # if they are specified in `sparsify_like_args_list`. + if not isinstance(tensor, SparseSemiStructuredTensor): + if i in sparsify_like_args_list: + tensor = semi_structured_sparsify_like( + tensor, reference_sparse_tensor + ) + else: + raise ValueError( + f"Operation {func.__module__}.{func.__name__} on {type(reference_sparse_tensor)} requires all operands to " + f"be {type(reference_sparse_tensor)}, but operand {i} is a {type(tensor)}" + ) + + # If the tensor is a SparseSemiStructuredTensor, we make sure that the sparsity pattern is the same as the reference tensor. + # Pointwise ops on tensors containing two different sparsity patterns is not defined, as in the case of addition, where + # adding two semi-structured sparse tensors yields a result that is not semi-structured sparse. + else: + if ( + tensor.compressed_swizzled_bitmask is None + or reference_sparse_tensor.compressed_swizzled_bitmask is None + or tensor.compressed_swizzled_bitmask.data_ptr() + != reference_sparse_tensor.compressed_swizzled_bitmask.data_ptr() + or tensor.compressed_swizzled_bitmask.stride() + != reference_sparse_tensor.compressed_swizzled_bitmask.stride() + ): + raise ValueError( + f"Operation {func.__module__}.{func.__name__} on {type(reference_sparse_tensor)} requires all operands to be " + f"{type(reference_sparse_tensor)} with the same sparsity pattern" + ) + return tensor + + args_updated = [handle_arg(i, tensor) for i, tensor in enumerate(args)] + + return reference_sparse_tensor.__class__( + reference_sparse_tensor.shape, + func( + *[ + x.packed if isinstance(x, SparseSemiStructuredTensor) else x + for x in args_updated + ] + ), + reference_sparse_tensor.meta, + func( + *[ + x.packed_t if isinstance(x, SparseSemiStructuredTensor) else x + for x in args_updated + ] + ), + reference_sparse_tensor.meta_t, + reference_sparse_tensor.compressed_swizzled_bitmask, + ) + + +# Add pointwise ops to the dispatch table +CUTLASS_POINTWISE_OP_DISPATCH_TABLE = { + torch.ops.aten.relu: _semi_sparse_pointwise_op, + torch.ops.aten.gelu: _semi_sparse_pointwise_op, + torch.ops.aten.silu: _semi_sparse_pointwise_op, + torch.ops.aten.mul: partial( + # `mul` BW in swiglu + _semi_sparse_pointwise_op, + sparsify_like_args_list=(0, 1), + ), + torch.ops.aten.add: _semi_sparse_pointwise_op, + # Note: for these ops, we allow the gradient to come in as a `torch.Tensor` + # and we will run the sparsification right before calling the BW aten func + torch.ops.aten.gelu_backward: partial( + _semi_sparse_pointwise_op, sparsify_like_args_list=(0,) + ), + torch.ops.aten.silu_backward: partial( + _semi_sparse_pointwise_op, sparsify_like_args_list=(0, 1) + ), + torch.ops.aten.threshold_backward: partial( # relu BW + _semi_sparse_pointwise_op, + sparsify_like_args_list=(0,), + ), +} diff --git a/lib/python3.12/site-packages/torchao/testing/__init__.py b/lib/python3.12/site-packages/torchao/testing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/testing/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36243ccf65fd17f6bd15abaa0b78272064497b9d Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/__pycache__/model_architectures.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/__pycache__/model_architectures.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b5bfc40687486fba25223a9668db7ab19793a95 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/__pycache__/model_architectures.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42a389c8960260c9fe474cb34bbe4baa121abceb Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/float8/__init__.py b/lib/python3.12/site-packages/torchao/testing/float8/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..780055040ef6c251af41f540066662ebeda2428c Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/dtensor_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/dtensor_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77f3f5dde976906e8e3e378b5a3ad462943193be Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/dtensor_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/fsdp2_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/fsdp2_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..174e2ea1cfd619e6b4abbd593980120c33e419f6 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/fsdp2_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/roofline_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/roofline_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a38a4327edd9e0d9968a64ea71b4e379952e3ee Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/roofline_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/test_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/test_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8cc5adeaf0f9a0edb82291e9f92391b370d12f20 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/float8/__pycache__/test_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/float8/dtensor_utils.py b/lib/python3.12/site-packages/torchao/testing/float8/dtensor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..84e409526315b1f710254571d0e8b8ca6d051951 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/float8/dtensor_utils.py @@ -0,0 +1,30 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import torch.nn as nn +import torch.nn.functional as F + + +class FeedForward(nn.Module): + """MLP based model""" + + def __init__(self): + super(FeedForward, self).__init__() + self.w1 = nn.Linear(16, 32, bias=False) + self.w2 = nn.Linear(16, 32, bias=False) + self.out_proj = nn.Linear(32, 16, bias=False) + + def forward(self, x): + return self.out_proj(F.silu(self.w1(x)) * self.w2(x)) + + +class ToyModel(nn.Module): + def __init__(self): + super(ToyModel, self).__init__() + self.ffn = FeedForward() + + def forward(self, x): + return self.ffn(x) diff --git a/lib/python3.12/site-packages/torchao/testing/float8/fsdp2_utils.py b/lib/python3.12/site-packages/torchao/testing/float8/fsdp2_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9eef29604077922a19da57100dab5046f8e7975b --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/float8/fsdp2_utils.py @@ -0,0 +1,101 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from typing import List + +import torch +import torch.distributed as dist +import torch.nn as nn + +from torchao.float8.config import ( + Float8LinearConfig, + ScalingType, +) +from torchao.float8.fsdp_utils import precompute_float8_dynamic_scale_for_fsdp + + +def check_parity_no_mp( + test_cls, + ref_model: nn.Module, + ref_optim: torch.optim.Optimizer, + fsdp_model: nn.Module, + fsdp_optim: torch.optim.Optimizer, + local_inp: torch.Tensor, + config: Float8LinearConfig, + precompute: bool = False, + compile_transformer_block: bool = False, +): + # check that requires_grad matches ref module + for ref_param, fsdp_param in zip(ref_model.parameters(), fsdp_model.parameters()): + test_cls.assertEqual( + ref_param.requires_grad, + fsdp_param.requires_grad, + msg=f"ref_param.requires_grad: {ref_param.requires_grad}, fsdp_param.requires_grad: {fsdp_param.requires_grad}", + ) + + # TODO(before land): reorder args and make config not optional + for iter_idx in range(10): + losses: List[torch.Tensor] = [] + for model, optim in ((ref_model, ref_optim), (fsdp_model, fsdp_optim)): + optim.zero_grad(set_to_none=(iter_idx % 2 == 0)) + losses.append(model(local_inp).sum()) + losses[-1].backward() + if model is ref_model: + for param in model.parameters(): + if param.requires_grad: + dist.all_reduce(param.grad) + param.grad.div_(dist.get_world_size()) + + optim.step() + if ( + model is fsdp_model + and precompute + and config.cast_config_weight.scaling_type is ScalingType.DYNAMIC + ): + precompute_float8_dynamic_scale_for_fsdp(model) + + test_cls.assertEqual( + losses[0], + losses[1], + msg=f"iter: {iter_idx}, loss-ref: {losses[0]}, loss-fp8: {losses[1]}", + ) + + +def check_parity_bf16_mp( + test_cls, + ref_model: nn.Module, + ref_model_bf16: nn.Module, + ref_optim: torch.optim.Optimizer, + fsdp_model: nn.Module, + fsdp_optim: torch.optim.Optimizer, + local_inp: torch.Tensor, +): + for iter_idx in range(10): + losses: List[torch.Tensor] = [] + for model, optim in ( + (ref_model_bf16, ref_optim), + (fsdp_model, fsdp_optim), + ): + optim.zero_grad(set_to_none=(iter_idx % 2 == 0)) + losses.append(model(local_inp).sum()) + losses[-1].backward() + if model is ref_model_bf16: + for param_bf16, param_fp32 in zip( + ref_model_bf16.parameters(), ref_model.parameters() + ): + dist.all_reduce(param_bf16.grad) + param_bf16.grad.div_(dist.get_world_size()) + param_fp32.grad = param_bf16.grad.float() + param_bf16.grad = None + optim.step() + for param_fp32, param_bf16 in zip( + ref_model.parameters(), ref_model_bf16.parameters() + ): + param_bf16.detach().copy_(param_fp32) + test_cls.assertEqual( + losses[0], + losses[1], + msg=f"iter: {iter_idx}, loss-ref: {losses[0]}, loss-fp8: {losses[1]}", + ) diff --git a/lib/python3.12/site-packages/torchao/testing/float8/roofline_utils.py b/lib/python3.12/site-packages/torchao/testing/float8/roofline_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7bfb9887df58aedba1c822e216d453a1cb380a51 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/float8/roofline_utils.py @@ -0,0 +1,335 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +from typing import List, Optional, Union + +import sympy +import torch + +BYTES_PER_EL_FLOAT8 = 1 +BYTES_PER_EL_BF16 = 2 + +gpu_name_to_specs = { + "NVIDIA H100": { + # https://www.nvidia.com/en-us/data-center/h100/, divide by 2 because no sparsity + "bf16_peak_tops": 989e12, + "fp8_peak_tops": 1979e12, + # 2.4 TB per second, custom to Meta's H100 variant + "peak_mem_bw_bytes_sec": 2.4e12, + # based on experimental observation with sample large inputs + "pct_achievable_gemm_tops": 0.78, + # based on previous experience looking at pointwise triton kernels with large inputs, + # which would hit about 2.2k GBPS on Meta's H100 variant + "pct_achievable_mem_bw": 0.92, + }, + "NVIDIA B200": { + # https://resources.nvidia.com/en-us-blackwell-architecture, page 19, + # divide by 2 because no sparsity + "bf16_peak_tops": 2.25e15, + "fp8_peak_tops": 4.5e15, + "fp4_peak_tops": 9.0e15, + # https://resources.nvidia.com/en-us-blackwell-architecture, page 20 + # 8.0 TB per second + "peak_mem_bw_bytes_sec": 8.0e12, + # for now, copy over from H100 + # TODO(future): measure once we have the hardware + "pct_achievable_gemm_tops": 0.78, + # for now, copy over from H100 + # TODO(future): measure once we have the hardware + "pct_achievable_mem_bw": 0.92, + }, + "AMD Instinct MI300X": { + # https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/data-sheets/amd-instinct-mi300x-data-sheet.pdf, page 1, + "bf16_peak_tops": 1307e12, + "fp8_peak_tops": 2614e12, + # 5.3 TB per second + "peak_mem_bw_bytes_sec": 5.3e12, + # for now, copy over from H100 + # TODO(future): run measurement on hardware + "pct_achievable_gemm_tops": 0.78, + # for now, copy over from H100 + # TODO(future): run measurement on hardware + "pct_achievable_mem_bw": 0.92, + }, + # TODO(future): more GPU names +} + + +def get_specs(): + gpu_name = torch.cuda.get_device_name(0) + return gpu_name_to_specs[gpu_name] + + +# Source: run a triton kernel with a single element read/write on an H100 and +# measure GPU time from the trace +# TODO(future): audit this across different hardware and triton/non-triton +KERNEL_LAUNCH_OVERHEAD_SEC = 0.002 * 0.001 + + +def get_tensor_memory_traffic_ovhd_s( + specs, + dim0, + dim1, + tensor_role: str, + float8_recipe_name: Optional[str], + mx_recipe_name: Optional[str], + fuse_with_prev=False, +) -> List[Union[sympy.Symbol, float]]: + """ + Calculates the roofline estimate of casting one of the gemm inputs + (input, weight or grad_output) to float8 in fwd+bwd. + + Inputs: dim0 and dim1 (shape), tensor_role (input|weight|grad_output), recipe names + Outputs: list of read/write traffic overhead in seconds, one for each kernel + """ + # assumes input bf16, output f8 + numel = dim0 * dim1 + + res_bytes = None + if float8_recipe_name == "tensorwise": + if tensor_role == "weight": + # x_bf16 = ... + # kernel 1: x_bf16 -> max_abs_stage_1 -> tmp + # kernel 2 (mem traffic not modeled): tmp -> max_abs_stage_2 -> max_abs + # kernel 3 (fwd): x_bf16, max_abs -> to_float8 -> x_fp8_dim0 + # kernel 4 (bwd): x_bf16, max_abs -> to_float8 -> x_fp8_dim1 + if fuse_with_prev: + kernel_1_rw = 0 + else: + # kernel 1: read numel, write 0 (assume size(tmp) ~ 0) + kernel_1_rw = BYTES_PER_EL_BF16 * numel + # kernel 3: read in bf16, write twice in float8 (row-major and col-major) + kernel_3_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + kernel_4_rw = kernel_3_rw + res_bytes = [kernel_1_rw, 0, kernel_3_rw, kernel_4_rw] + else: + # x_bf16 = ... + # kernel 1: x_bf16 -> max_abs_stage_1 -> tmp + # kernel 2 (mem traffic not modeled): tmp -> max_abs_stage_2 -> max_abs + # kernel 3: x_bf16, max_abs -> to_float8 -> x_fp8_dim0, x_fp8_dim1 + if fuse_with_prev: + kernel_1_rw = 0 + else: + # kernel 1: read numel, write 0 (assume size(tmp) ~ 0) + kernel_1_rw = BYTES_PER_EL_BF16 * numel + # kernel 3: read in bf16, write twice in float8 (row-major and col-major) + kernel_3_rw = BYTES_PER_EL_BF16 * numel + 2 * BYTES_PER_EL_FLOAT8 * numel + res_bytes = [kernel_1_rw, 0, kernel_3_rw] + + elif float8_recipe_name == "rowwise": + if tensor_role == "weight": + # x_bf16 = ... + # kernel 1 (fwd): x_bf16_dim0 -> x_float8_dim0 + # kernel 2 (bwd): x_bf16_dim0 -> x_bf16_dim1 + # kernel 3 (bwd): x_bf16_dim1 -> x_float8_dim1 + # assume that we can't fuse 2 and 3 because that would require loading + # the entire tensor to shared memory + if fuse_with_prev: + # assume we can fuse one of the reads with previous op + kernel_1_rw = 0 + BYTES_PER_EL_FLOAT8 * numel + else: + kernel_1_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + kernel_2_rw = BYTES_PER_EL_BF16 * numel * 2 + kernel_3_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + res_bytes = [kernel_1_rw, kernel_2_rw, kernel_3_rw] + else: + # x_bf16 = ... + # kernel 1: x_bf16_dim0 -> x_float8_dim0, x_bf16_dim1 + # kernel 2: x_bf16_dim1 -> x_float8_dim1 + # assume that we can't fuse 1 and 2 because that would require loading + # the entire tensor to shared memory + if fuse_with_prev: + # assume we can fuse one of the reads with previous op + kernel_1_rw = ( + 0 + BYTES_PER_EL_FLOAT8 * numel + BYTES_PER_EL_BF16 * numel + ) + else: + kernel_1_rw = ( + BYTES_PER_EL_BF16 * numel + + BYTES_PER_EL_FLOAT8 * numel + + BYTES_PER_EL_BF16 * numel + ) + kernel_2_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + res_bytes = [kernel_1_rw, kernel_2_rw] + + elif float8_recipe_name == "rowwise_with_gw_hp": + if tensor_role in ("input", "grad_output"): + # x_bf16 = ... + # kernel 1 (fwd): x_bf16_dim0 -> x_float8_dim0 + # bwd: no-op + if fuse_with_prev: + kernel_1_rw = 0 + BYTES_PER_EL_FLOAT8 * numel + else: + kernel_1_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + res_bytes = [kernel_1_rw] + elif tensor_role == "weight": + # x_bf16 = ... + # kernel 1 (fwd): w_bf16 -> w_float8_dim0, w_scale_dim0 + # kernel 2 (bwd): w_scale_dim0 -> w_scale_tensorwise + # kernel 3 (bwd): w_bf16, w_scale_tensorwise -> w_float8_dim1 + kernel_1_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + kernel_2_rw = 0 + kernel_3_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + res_bytes = [kernel_1_rw, kernel_2_rw, kernel_3_rw] + else: + assert False, "unsupported" + + else: + assert mx_recipe_name in ( + "mxfp8_emulated", + "mxfp8_cublas", + ), "unsupported" + # For now, assume that we can't profitably fuse kernel 1 and kernel 2 + # x_bf16 = ... + # kernel 1: x_bf16 -> x_mxfp8_dim0 + # kernel 2: x_bf16 -> x_mxfp8_dim1 + if fuse_with_prev: + kernel_1_rw = 0 + BYTES_PER_EL_FLOAT8 * numel + else: + kernel_1_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + kernel_2_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + res_bytes = [kernel_1_rw, kernel_2_rw] + + # convert from bytes to seconds + res_s = [ + x / specs["peak_mem_bw_bytes_sec"] / specs["pct_achievable_mem_bw"] + for x in res_bytes + ] + + # take max of kernel_overhead, r/w time + res_s = [sympy.Max(x, KERNEL_LAUNCH_OVERHEAD_SEC) for x in res_s] + + return res_s + + +def get_individual_gemm_time_sympy( + M: sympy.Symbol, K: sympy.Symbol, N: sympy.Symbol, dtype, mx_recipe_name +) -> sympy.Symbol: + # compute bound + specs = get_specs() + gemm_ops = 2 * M * K * N + if dtype is torch.bfloat16: + peak_tops = specs["bf16_peak_tops"] + elif dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + peak_tops = specs["fp8_peak_tops"] + else: + assert False, "unsupported" + compute_gemm_time_s = gemm_ops / peak_tops / specs["pct_achievable_gemm_tops"] + + # memory bound + num_reads = M * K + K * N + num_writes = M * N + + if mx_recipe_name is not None: + assert mx_recipe_name in ( + "mxfp8_emulated", + "mxfp8_cublas", + ), "unsupported" + assert dtype in (torch.float8_e4m3fn, torch.float8_e5m2), "unsupported" + # adjust reads for MX scaling + block_size = 32 + num_scale_reads = num_reads // block_size + # note: e8m0 bytes per element is the same as for e4m3|e5m2 + num_reads = num_reads + num_scale_reads + + if dtype is torch.bfloat16: + bytes_rw = num_reads * BYTES_PER_EL_BF16 + num_writes * BYTES_PER_EL_BF16 + elif dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + # read in float8, output in bfloat16 + bytes_rw = num_reads * BYTES_PER_EL_FLOAT8 + num_writes * BYTES_PER_EL_BF16 + else: + assert False, "unsupported" + mem_gemm_time_s = ( + bytes_rw / specs["peak_mem_bw_bytes_sec"] / specs["pct_achievable_mem_bw"] + ) + + return sympy.Max(compute_gemm_time_s, mem_gemm_time_s, KERNEL_LAUNCH_OVERHEAD_SEC) + + +def get_gemm_time_sympy( + M: sympy.Symbol, + K: sympy.Symbol, + N: sympy.Symbol, + dtype, + float8_recipe_name: Optional[str], + mx_recipe_name: Optional[str], +): + # next: add rowwise_with_gw_hp here + # note: this function is currently not super accurate for small shapes: + # when M,K,N <= 1k,1k,1k it undercounts by around 2x + + gemm_dtype_input, gemm_dtype_grad_input, gemm_dtype_grad_weight = ( + dtype, + dtype, + dtype, + ) + if float8_recipe_name == "rowwise_with_gw_hp": + gemm_dtype_grad_weight = torch.bfloat16 + + gemm_output_time_s = get_individual_gemm_time_sympy( + M, K, N, gemm_dtype_input, mx_recipe_name + ) + gemm_grad_input_time_s = get_individual_gemm_time_sympy( + M, N, K, gemm_dtype_grad_input, mx_recipe_name + ) + gemm_grad_weight_time_s = get_individual_gemm_time_sympy( + K, M, N, gemm_dtype_grad_weight, mx_recipe_name + ) + total = gemm_output_time_s + gemm_grad_input_time_s + gemm_grad_weight_time_s + return total + + +def get_float8_mem_sympy( + M, + K, + N, + float8_recipe_name: Optional[str], + mx_recipe_name: Optional[str], + enable_fusion_modeling: bool, +): + specs = get_specs() + + # there are three gemms in the fwd/bwd of a linear: + # + # input @ weight_t = output + # MxK @ KxN => MxN + # + # grad_output @ weight = grad_input + # MxN @ NxK => MxK + # + # input_t @ grad_output = grad_weight + # KxM @ MxN => KxN + + fwd_fp8_input_mem = get_tensor_memory_traffic_ovhd_s( + specs, + M, + K, + tensor_role="input", + float8_recipe_name=float8_recipe_name, + mx_recipe_name=mx_recipe_name, + fuse_with_prev=enable_fusion_modeling, + ) + fwd_fp8_weight_mem = get_tensor_memory_traffic_ovhd_s( + specs, + K, + N, + tensor_role="weight", + float8_recipe_name=float8_recipe_name, + mx_recipe_name=mx_recipe_name, + fuse_with_prev=False, + ) + gi_fp8_grad_output_mem = get_tensor_memory_traffic_ovhd_s( + specs, + M, + N, + tensor_role="grad_output", + float8_recipe_name=float8_recipe_name, + mx_recipe_name=mx_recipe_name, + fuse_with_prev=enable_fusion_modeling, + ) + + res = sum([*fwd_fp8_input_mem, *fwd_fp8_weight_mem, *gi_fp8_grad_output_mem]) + return res diff --git a/lib/python3.12/site-packages/torchao/testing/float8/test_utils.py b/lib/python3.12/site-packages/torchao/testing/float8/test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4a30931c9e08165cfc1af0e2641c2d170d20f874 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/float8/test_utils.py @@ -0,0 +1,34 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +from torchao.float8.config import ( + CastConfig, + Float8LinearConfig, +) + + +def get_test_float8_linear_config( + scaling_type_input, + scaling_type_weight, + scaling_type_grad_output, + emulate: bool, +): + cast_config_input = CastConfig( + scaling_type=scaling_type_input, + ) + cast_config_weight = CastConfig( + scaling_type=scaling_type_weight, + ) + cast_config_grad_output = CastConfig( + scaling_type=scaling_type_grad_output, + ) + + config = Float8LinearConfig( + cast_config_input=cast_config_input, + cast_config_weight=cast_config_weight, + cast_config_grad_output=cast_config_grad_output, + emulate=emulate, + ) + return config diff --git a/lib/python3.12/site-packages/torchao/testing/model_architectures.py b/lib/python3.12/site-packages/torchao/testing/model_architectures.py new file mode 100644 index 0000000000000000000000000000000000000000..f59a1271b18c49b31818e43b8a2eba83af611244 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/model_architectures.py @@ -0,0 +1,179 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import re + +import torch +import torch.nn as nn + + +# TODO: Refactor torchao and tests to use these models +class ToyLinearModel(torch.nn.Module): + def __init__(self, k=64, n=32, dtype=torch.bfloat16): + super().__init__() + self.linear1 = torch.nn.Linear(k, n, bias=False).to(dtype) + + def forward(self, x): + x = self.linear1(x) + return x + + +class LNLinearActivationModel(nn.Module): + def __init__(self, fc_dim1, fc_dim2, dtype=torch.bfloat16, activation="sigmoid"): + super().__init__() + + activation = activation.lower() + activation_map = { + "relu": nn.ReLU(), + "sigmoid": nn.Sigmoid(), + "leakyrelu": nn.LeakyReLU(), + "relu6": nn.ReLU6(), + "gelu": nn.GELU(), + "silu": nn.SiLU(), + "hardswish": nn.Hardswish(), + } + + if activation not in activation_map: + raise ValueError(f"Unsupported activation: {activation}") + + self.ln = nn.LayerNorm(fc_dim1, elementwise_affine=False) + self.fc = nn.Linear(fc_dim1, fc_dim2, bias=False).to(dtype=dtype) + self.activation = activation_map[activation] + + def forward(self, x): + x = self.ln(x) + x = self.fc(x) + return self.activation(x) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def _norm(self, x): + return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = self._norm(x.float()).type_as(x) + return output * self.weight + + +class TransformerBlock(torch.nn.Module): + def __init__(self, hidden_dim, num_heads=8, mlp_ratio=4, dtype=torch.bfloat16): + super().__init__() + self.hidden_dim = hidden_dim + self.num_heads = num_heads + self.head_dim = hidden_dim // num_heads + + # Self-attention + self.qkv = torch.nn.Linear(hidden_dim, 3 * hidden_dim, bias=False).to(dtype) + self.proj = torch.nn.Linear(hidden_dim, hidden_dim, bias=False).to(dtype) + + # MLP + self.mlp_ratio = mlp_ratio + self.mlp_hidden_dim = int(hidden_dim * mlp_ratio) + self.mlp_fc1 = torch.nn.Linear(hidden_dim, self.mlp_hidden_dim, bias=False).to( + dtype + ) + self.mlp_fc2 = torch.nn.Linear(self.mlp_hidden_dim, hidden_dim, bias=False).to( + dtype + ) + + # Layer norms + self.norm1 = RMSNorm(hidden_dim).to(dtype) + self.norm2 = RMSNorm(hidden_dim).to(dtype) + + # Activation + self.activation = torch.nn.GELU() + + def forward(self, x): + batch_size, seq_len, _ = x.shape + + # Self-attention + residual = x + x = self.norm1(x) + + # Reshape qkv projection for better memory layout + qkv = self.qkv(x) # [batch_size, seq_len, 3 * hidden_dim] + qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim) + qkv = qkv.permute( + 2, 0, 3, 1, 4 + ) # [3, batch_size, num_heads, seq_len, head_dim] + q, k, v = qkv # Each has shape [batch_size, num_heads, seq_len, head_dim] + + # Scaled dot-product attention with proper reshaping + # Reshape for better memory layout and avoid broadcasting issues + q = q.reshape(batch_size * self.num_heads, seq_len, self.head_dim) + k = k.reshape(batch_size * self.num_heads, seq_len, self.head_dim) + v = v.reshape(batch_size * self.num_heads, seq_len, self.head_dim) + + # Compute attention scores + attn = (q @ k.transpose(-2, -1)) * (1.0 / (self.head_dim**0.5)) + attn = torch.softmax(attn, dim=-1) + + # Apply attention to values + x = attn @ v # [batch_size * num_heads, seq_len, head_dim] + + # Reshape back to original dimensions + x = x.reshape(batch_size, self.num_heads, seq_len, self.head_dim) + x = x.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_dim) + + # Project back to hidden dimension + x = self.proj(x) + x = residual + x + + # MLP + residual = x + x = self.norm2(x) + x = self.mlp_fc1(x) + x = self.activation(x) + x = self.mlp_fc2(x) + x = residual + x + + return x + + +def create_model_and_input_data( + model_type: str, + m: int, + k: int, + n: int, + high_precision_dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + activation: str = "relu", +): + """Create a model and input data for benchmarking. + + Args: + model_type (str): type of the model to be created + batch_size (int): batch size of the input data + device (str): device to run the model on + high_precision_dtype (torch.dtype): data type of the model + m, k, n (int): dimensions of the model and input data + """ + if model_type == "linear": + model = ToyLinearModel(k, n, high_precision_dtype).to(device) + input_data = torch.randn(m, k, device=device, dtype=high_precision_dtype) + elif "ln_linear" in model_type: + # Extract activation type from model_type string + match = re.search(r"ln_linear_?(\w+)?", model_type) + activation = match.group(1) if match and match.group(1) else "relu" + model = LNLinearActivationModel( + k, n, high_precision_dtype, activation=activation + ).to(device) + input_data = torch.randn(m, k, device=device, dtype=high_precision_dtype) + elif model_type == "transformer_block": + # For transformer block, k is the hidden dimension + model = TransformerBlock( + k, num_heads=8, mlp_ratio=4, dtype=high_precision_dtype + ).to(device) + # Input shape for transformer is [batch_size, seq_len, hidden_dim] + input_data = torch.randn(m, 16, k, device=device, dtype=high_precision_dtype) + else: + raise ValueError(f"Unknown model type: {model_type}") + return model, input_data diff --git a/lib/python3.12/site-packages/torchao/testing/pt2e/__init__.py b/lib/python3.12/site-packages/torchao/testing/pt2e/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d2c4068c8d6500f92d32597bd994ae5990986df1 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/_xnnpack_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/_xnnpack_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55a6f123bc622531026ebf3562b67d3f18cb54a8 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/_xnnpack_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/_xnnpack_quantizer_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/_xnnpack_quantizer_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86601fa58c21d7e144f805b345a204b41e7a90ae Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/_xnnpack_quantizer_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..acfd8259a77ffe568e3de75df0aee3276b1225c6 Binary files /dev/null and b/lib/python3.12/site-packages/torchao/testing/pt2e/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchao/testing/pt2e/_xnnpack_quantizer.py b/lib/python3.12/site-packages/torchao/testing/pt2e/_xnnpack_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..30d2e6178db23e3d266cb0ea6f8af269654a1fd2 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/pt2e/_xnnpack_quantizer.py @@ -0,0 +1,455 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +from __future__ import annotations + +import copy +import functools +import warnings +from typing import TYPE_CHECKING, Any, Callable, Optional + +import torch +import torch._dynamo as torchdynamo +import torch.nn.functional as F +from torch.fx._compatibility import compatibility + +from torchao.quantization.pt2e import ( + FakeQuantize, + FusedMovingAvgObsFakeQuantize, + HistogramObserver, + MinMaxObserver, + MovingAverageMinMaxObserver, + MovingAveragePerChannelMinMaxObserver, + PerChannelMinMaxObserver, + PlaceholderObserver, +) +from torchao.quantization.pt2e.quantizer import ( + QuantizationSpec, + Quantizer, + get_module_name_filter, +) +from torchao.testing.pt2e._xnnpack_quantizer_utils import ( + OP_TO_ANNOTATOR, + OperatorConfig, + OperatorPatternType, + QuantizationConfig, + _convert_scalars_to_attrs, + propagate_annotation, +) + +if TYPE_CHECKING: + from torch.fx import Node + + from torchao.quantization.pt2e import _ObserverOrFakeQuantizeConstructor + + +__all__ = [ + "XNNPACKQuantizer", + "get_symmetric_quantization_config", +] + + +def _get_dynamo_graph(function: Callable, inputs) -> torch.fx.Graph: + gm, _ = torchdynamo.export(function, aten_graph=True)(*inputs) + gm.graph.eliminate_dead_code() + return gm.graph + + +def _get_linear_patterns(input_size: list[int]): + in_channels = input_size[-1] + out_channels = 8 # hard coding but this should not matter + weight = torch.ones((out_channels, in_channels)) + bias = torch.ones((out_channels,)) + act = torch.ones(input_size) + + def linear_op(act, weight, bias=None): + return F.linear(act, weight, bias) + + pattern_w_bias = _get_dynamo_graph(linear_op, (act, weight, bias)) + pattern_wo_bias = _get_dynamo_graph(linear_op, (act, weight)) + return [pattern_w_bias, pattern_wo_bias] + + +def _supported_symmetric_quantized_operators() -> dict[str, list[OperatorPatternType]]: + supported_operators: dict[str, list[OperatorPatternType]] = { + # Both conv and linear should be able to handle relu + hardtanh fusion since + # those are clamp ops + "conv2d": [ + [torch.nn.Conv2d, torch.nn.ReLU], + [torch.nn.Conv2d, F.relu], + [F.conv2d, torch.nn.ReLU], + [F.conv2d, F.relu], + ], + "linear": [[torch.nn.Linear], [F.linear]], + "add": [[torch.add]], + "adaptive_avg_pool2d": [ + [torch.nn.AdaptiveAvgPool2d], + [F.adaptive_avg_pool2d], + ], + } + return copy.deepcopy(supported_operators) + + +def _get_supported_symmetric_config_and_operators() -> list[OperatorConfig]: + supported_config_and_operators: list[OperatorConfig] = [] + for quantization_config in [ + get_symmetric_quantization_config(), + get_symmetric_quantization_config(is_qat=True), + get_symmetric_quantization_config(is_per_channel=True), + get_symmetric_quantization_config(is_per_channel=True, is_qat=True), + ]: + ops = _supported_symmetric_quantized_operators() + supported_config_and_operators.extend( + OperatorConfig(quantization_config, pattern_list) + for pattern_list in ops.values() + ) + return copy.deepcopy(supported_config_and_operators) + + +@functools.lru_cache +def get_symmetric_quantization_config( + is_per_channel: bool = False, + is_qat: bool = False, + is_dynamic: bool = False, + act_qmin: int = -128, + act_qmax: int = 127, + weight_qmin: int = -127, + weight_qmax: int = 127, +): + extra_args: dict[str, Any] = {"eps": 2**-12} + if is_qat: + if is_dynamic: + act_observer_or_fake_quant_ctr = FakeQuantize + dynamic_quant_observer = MovingAverageMinMaxObserver.with_args( + averaging_constant=1 + ) + extra_args["observer"] = dynamic_quant_observer + else: + act_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize # type: ignore[assignment] + else: + if is_dynamic: + act_observer_or_fake_quant_ctr = PlaceholderObserver # type: ignore[assignment] + else: + act_observer_or_fake_quant_ctr = HistogramObserver # type: ignore[assignment] + + act_quantization_spec = QuantizationSpec( + dtype=torch.int8, + quant_min=act_qmin, + quant_max=act_qmax, + qscheme=torch.per_tensor_affine, + is_dynamic=is_dynamic, + observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args( + **extra_args, + ), + ) + weight_qscheme = ( + torch.per_channel_symmetric if is_per_channel else torch.per_tensor_symmetric + ) + weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = ( + MinMaxObserver + ) + if is_qat: + # TODO: qat + per channel? + weight_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize + elif is_per_channel: + weight_observer_or_fake_quant_ctr = PerChannelMinMaxObserver + + extra_args: dict[str, Any] = {"eps": 2**-12} + if is_qat: + if weight_qscheme == torch.per_tensor_symmetric: + extra_args["observer"] = MovingAverageMinMaxObserver + else: + extra_args["observer"] = MovingAveragePerChannelMinMaxObserver # type: ignore[dict-item] + weight_quantization_spec = QuantizationSpec( + dtype=torch.int8, + quant_min=weight_qmin, + quant_max=weight_qmax, + qscheme=weight_qscheme, + ch_axis=0, + is_dynamic=False, + observer_or_fake_quant_ctr=weight_observer_or_fake_quant_ctr.with_args( + **extra_args + ), + ) + + bias_quantization_spec = None + if is_dynamic: + quantization_config = QuantizationConfig( + act_quantization_spec, + None, + weight_quantization_spec, + bias_quantization_spec, + is_qat, + ) + else: + quantization_config = QuantizationConfig( + act_quantization_spec, + act_quantization_spec, + weight_quantization_spec, + bias_quantization_spec, + is_qat, + ) + return quantization_config + + +def _get_supported_config_and_operators() -> list[OperatorConfig]: + return _get_supported_symmetric_config_and_operators() + + +def _get_module_type_filter(tp: Callable): + """Get the module_type_filter function for a given module type, the filter accepts + a node and checks if the node comes from a module that has certain module type + + For example: + node: linear_op = call_function[...](...) # comes from a module with type Block -> Sub -> Linear + + + >> module_type_filter = _get_module_type_filter(Sub) # submodule with type `Sub`, under the `Block` submodule + >> print(module_type_filter(node)) + True # the node is from the submodule `Sub` (same for `Block` and `Linear` as well) + """ + + tp_str = tp.__module__ + "." + tp.__qualname__ + + def module_type_filter(n: Node) -> bool: + # example: { + # 'L__self___sub': ("L['self'].sub", ), + # 'L__self___sub_linear': ("L['self'].sub.linear", ) + # } + nn_module_stack = n.meta.get("nn_module_stack", {}) + types = [] + for _, t in nn_module_stack.values(): + # export() returns str, but older APIs (e.g. capture_pre_autograd_graph) + # return type. Handle both cases. + if isinstance(t, type): + t = t.__module__ + "." + t.__qualname__ + types.append(t) + return tp_str in types + + return module_type_filter + + +def _get_not_module_type_or_name_filter( + tp_list: list[Callable], module_name_list: list[str] +) -> Callable[[Node], bool]: + module_type_filters = [_get_module_type_filter(tp) for tp in tp_list] + module_name_list_filters = [get_module_name_filter(m) for m in module_name_list] + + def not_module_type_or_name_filter(n: Node) -> bool: + return not any(f(n) for f in module_type_filters + module_name_list_filters) + + return not_module_type_or_name_filter + + +@compatibility(is_backward_compatible=False) +class XNNPACKQuantizer(Quantizer): + """ + !!! DEPRECATED !!! + XNNPACKQuantizer is a marked as deprected. It will be removed in the future. + It has been moved to executorch.backends.xnnpack.quantizer.xnnpack_quantizer.XNNPACKQuantizer. + Please use the new quantizer instead. + """ + + supported_config_and_operators = _get_supported_config_and_operators() + STATIC_QAT_ONLY_OPS = [ + "conv_bn_relu", + "conv_bn", + "conv_transpose_bn_relu", + "conv_transpose_bn", + ] + + # static quantization ops (both PTQ and QAT) + # Preserve the order that fusions come before singular ops + STATIC_OPS = [ + "linear_relu", + "linear", + "conv_relu", + "conv", + "conv_transpose_relu", + "adaptive_avg_pool2d", + # TODO: move this to BoltNNQuantizer? + "gru_io_only", + "add_relu", + "add", + "mul_relu", + "mul", + "cat", + ] + + DYNAMIC_OPS = [ + "linear", + ] + + def __init__(self) -> None: + super().__init__() + warnings.warn(f"{self.__class__.__name__} is deprecated!") + self.global_config: Optional[QuantizationConfig] = None + self.operator_type_config: dict[ + torch._ops.OpOverloadPacket, Optional[QuantizationConfig] + ] = {} + self.module_type_config: dict[Callable, Optional[QuantizationConfig]] = {} + self.module_name_config: dict[str, Optional[QuantizationConfig]] = {} + + @classmethod + def get_supported_quantization_configs(cls) -> list[QuantizationConfig]: + op_configs: set[QuantizationConfig] = { + spec for spec, _ in cls.supported_config_and_operators + } + return list(op_configs) + + @classmethod + def get_supported_operator_for_quantization_config( + cls, quantization_config: Optional[QuantizationConfig] + ) -> list[OperatorPatternType]: + if quantization_config is None: + all_ops = [] + for _, ops in cls.supported_config_and_operators: + all_ops.extend(ops) + return all_ops + + for config, ops in cls.supported_config_and_operators: + # note: this assumes each entry in cls.supported_spec_and_operators + # corresponds to one spec, e.g. we don't have + # [(spec1, op_list1), (spec1, op_list2), (spec2, op_list3)] + # where the first and second entry have the same spec but did not + # merge the op list + if config == quantization_config: + return ops + return [] + + def set_global(self, quantization_config: QuantizationConfig) -> XNNPACKQuantizer: + self.global_config = quantization_config + return self + + def set_operator_type( + self, + operator_type: torch._ops.OpOverloadPacket, + quantization_config: QuantizationConfig, + ) -> XNNPACKQuantizer: + self.operator_type_config[operator_type] = quantization_config + return self + + def set_module_type( + self, module_type: Callable, quantization_config: QuantizationConfig + ): + """Set quantization_config for a submodule with type: `module_type`, for example: + quantizer.set_module_name(Sub) or quantizer.set_module_name(nn.Linear), it will quantize all supported operator/operator + patterns in the submodule with this module type with the given `quantization_config` + """ + self.module_type_config[module_type] = quantization_config + return self + + def set_module_name( + self, module_name: str, quantization_config: Optional[QuantizationConfig] + ): + """Set quantization_config for a submodule with name: `module_name`, for example: + quantizer.set_module_name("blocks.sub"), it will quantize all supported operator/operator + patterns in the submodule with this module name with the given `quantization_config` + """ + assert quantization_config is not None, ( + " quantization_config == None is not supported yet" + ) + self.module_name_config[module_name] = quantization_config + return self + + def transform_for_annotation( + self, model: torch.fx.GraphModule + ) -> torch.fx.GraphModule: + """Transforms scalar values to tensor attributes""" + return _convert_scalars_to_attrs(model) + + def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule: + """just handling global spec for now""" + # hacked for handling dynamic linear quant. will fix later. + if self.global_config and self.global_config.input_activation.is_dynamic: # type: ignore[union-attr] + model = self._annotate_for_dynamic_quantization_config(model) + else: + model = self._annotate_for_static_quantization_config(model) + propagate_annotation(model) + return model + + def _annotate_all_static_patterns( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, + ) -> torch.fx.GraphModule: + # TODO: implement the support for None to be canceling out previous annotations + if quantization_config is None: + return model + + if quantization_config.is_qat: + for op in self.STATIC_QAT_ONLY_OPS: + OP_TO_ANNOTATOR[op](model, quantization_config, filter_fn) + for op in self.STATIC_OPS: + OP_TO_ANNOTATOR[op](model, quantization_config, filter_fn) + return model + + def _annotate_all_dynamic_patterns( + self, + model: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, + ) -> torch.fx.GraphModule: + # TODO: implement the support for None to be canceling out previous annotations + if quantization_config is None: + return model + + for op in self.DYNAMIC_OPS: + OP_TO_ANNOTATOR[op](model, quantization_config, filter_fn) + return model + + def _annotate_for_static_quantization_config( + self, model: torch.fx.GraphModule + ) -> torch.fx.GraphModule: + module_name_list = list(self.module_name_config.keys()) + for module_name, config in self.module_name_config.items(): + self._annotate_all_static_patterns( + model, config, get_module_name_filter(module_name) + ) + + tp_list = list(self.module_type_config.keys()) + for module_type, config in self.module_type_config.items(): + self._annotate_all_static_patterns( + model, config, _get_module_type_filter(module_type) + ) + + self._annotate_all_static_patterns( + model, + self.global_config, + _get_not_module_type_or_name_filter(tp_list, module_name_list), + ) + return model + + def _annotate_for_dynamic_quantization_config( + self, model: torch.fx.GraphModule + ) -> torch.fx.GraphModule: + module_name_list = list(self.module_name_config.keys()) + for module_name, config in self.module_name_config.items(): + self._annotate_all_dynamic_patterns( + model, config, get_module_name_filter(module_name) + ) + + tp_list = list(self.module_type_config.keys()) + for module_type, config in self.module_type_config.items(): + self._annotate_all_dynamic_patterns( + model, config, _get_module_type_filter(module_type) + ) + + self._annotate_all_dynamic_patterns( + model, + self.global_config, + _get_not_module_type_or_name_filter(tp_list, module_name_list), + ) + return model + + def validate(self, model: torch.fx.GraphModule) -> None: + pass + + @classmethod + def get_supported_operators(cls) -> list[OperatorConfig]: + return cls.supported_config_and_operators diff --git a/lib/python3.12/site-packages/torchao/testing/pt2e/_xnnpack_quantizer_utils.py b/lib/python3.12/site-packages/torchao/testing/pt2e/_xnnpack_quantizer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c49caca28580de7dc31ee66198ebf6514f6e4afe --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/pt2e/_xnnpack_quantizer_utils.py @@ -0,0 +1,1133 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +# mypy: allow-untyped-defs +import itertools +import typing +from dataclasses import dataclass +from typing import Callable, NamedTuple, Optional + +import torch +import torch.nn.functional as F +from torch._subclasses import FakeTensor +from torch.fx import Node +from torch.fx.passes.utils.matcher_with_name_node_map_utils import ( + SubgraphMatcherWithNameNodeMap, +) +from torch.fx.passes.utils.source_matcher_utils import get_source_partitions + +from torchao.quantization.pt2e.export_utils import WrapperModule +from torchao.quantization.pt2e.quantizer import ( + QuantizationAnnotation, + QuantizationSpec, + SharedQuantizationSpec, +) +from torchao.quantization.pt2e.quantizer.utils import ( + annotate_input_qspec_map, + annotate_output_qspec, +) +from torchao.quantization.pt2e.utils import ( + _get_aten_graph_module_for_pattern, + _is_conv_node, + _is_conv_transpose_node, + get_new_attr_name_with_prefix, +) + +__all__ = [ + "OperatorConfig", + "OperatorPatternType", + "QuantizationConfig", + "get_input_act_qspec", + "get_output_act_qspec", + "get_weight_qspec", + "get_bias_qspec", + "OP_TO_ANNOTATOR", + "propagate_annotation", +] + + +# In the absence of better name, just winging it with QuantizationConfig +@dataclass(eq=True, frozen=True) +class QuantizationConfig: + input_activation: Optional[QuantizationSpec] + output_activation: Optional[QuantizationSpec] + weight: Optional[QuantizationSpec] + bias: Optional[QuantizationSpec] + # TODO: remove, since we can use observer_or_fake_quant_ctr to express this + is_qat: bool = False + + +# Use Annotated because list[Callable].__module__ is read-only. +OperatorPatternType = typing.Annotated[list[Callable], None] +OperatorPatternType.__module__ = ( + "torchao.quantization.pt2e.quantizer.xnnpack_quantizer_utils" +) + +AnnotatorType = Callable[ + [ + torch.fx.GraphModule, + Optional[QuantizationConfig], + Optional[Callable[[Node], bool]], + ], + Optional[list[list[Node]]], +] +OP_TO_ANNOTATOR: dict[str, AnnotatorType] = {} + + +def register_annotator(op: str) -> Callable[[AnnotatorType], None]: + def decorator(annotator: AnnotatorType) -> None: + OP_TO_ANNOTATOR[op] = annotator + + return decorator + + +class OperatorConfig(NamedTuple): + # fix List[str] with List[List[Union[nn.Module, FunctionType, BuiltinFunctionType]]] + # Basically we are mapping a quantization config to some list of patterns. + # a pattern is defined as a list of nn module, function or builtin function names + # e.g. [nn.Conv2d, torch.relu, torch.add] + # We have not resolved whether fusion can be considered internal details of the + # quantizer hence it does not need communication to user. + # Note this pattern is not really informative since it does not really + # tell us the graph structure resulting from the list of ops. + config: QuantizationConfig + operators: list[OperatorPatternType] + + +def _is_annotated(nodes: list[Node]): + """ + Given a list of nodes (that represents an operator pattern), + check if any of the node is annotated, return True if any of the node + is annotated, otherwise return False + """ + annotated = False + for node in nodes: + annotated = annotated or ( + "quantization_annotation" in node.meta + and node.meta["quantization_annotation"]._annotated + ) + return annotated + + +def _mark_nodes_as_annotated(nodes: list[Node]): + for node in nodes: + if node is not None: + if "quantization_annotation" not in node.meta: + node.meta["quantization_annotation"] = QuantizationAnnotation() + node.meta["quantization_annotation"]._annotated = True + + +def get_input_act_qspec(quantization_config: Optional[QuantizationConfig]): + if quantization_config is None: + return None + if quantization_config.input_activation is None: + return None + quantization_spec: QuantizationSpec = quantization_config.input_activation + assert quantization_spec.qscheme in [ + torch.per_tensor_affine, + torch.per_tensor_symmetric, + ] + return quantization_spec + + +def get_output_act_qspec(quantization_config: Optional[QuantizationConfig]): + if quantization_config is None: + return None + if quantization_config.output_activation is None: + return None + quantization_spec: QuantizationSpec = quantization_config.output_activation + assert quantization_spec.qscheme in [ + torch.per_tensor_affine, + torch.per_tensor_symmetric, + ] + return quantization_spec + + +def get_weight_qspec(quantization_config: Optional[QuantizationConfig]): + if quantization_config is None: + return None + assert quantization_config is not None + if quantization_config.weight is None: + return None + quantization_spec: QuantizationSpec = quantization_config.weight + if quantization_spec.qscheme not in [ + torch.per_tensor_symmetric, + torch.per_channel_symmetric, + None, + ]: + raise ValueError( + f"Unsupported quantization_spec {quantization_spec} for weight" + ) + return quantization_spec + + +def get_bias_qspec(quantization_config: Optional[QuantizationConfig]): + if quantization_config is None: + return None + assert quantization_config is not None + if quantization_config.bias is None: + return None + quantization_spec: QuantizationSpec = quantization_config.bias + assert quantization_spec.dtype == torch.float, ( + "Only float dtype for bias is supported for bias right now" + ) + return quantization_spec + + +@register_annotator("linear") +def _annotate_linear( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + annotated_partitions = [] + input_act_qspec = get_input_act_qspec(quantization_config) + output_act_qspec = get_output_act_qspec(quantization_config) + weight_qspec = get_weight_qspec(quantization_config) + bias_qspec = get_bias_qspec(quantization_config) + for node in gm.graph.nodes: + if node.op != "call_function" or node.target != torch.ops.aten.linear.default: + continue + if filter_fn and not filter_fn(node): + continue + act_node = node.args[0] + weight_node = node.args[1] + bias_node = None + if len(node.args) > 2: + bias_node = node.args[2] + + if _is_annotated([node]) is False: # type: ignore[list-item] + annotate_input_qspec_map( + node, + act_node, + input_act_qspec, + ) + annotate_input_qspec_map( + node, + weight_node, + weight_qspec, + ) + nodes_to_mark_annotated = [node, weight_node] + if bias_node: + annotate_input_qspec_map( + node, + bias_node, + bias_qspec, + ) + nodes_to_mark_annotated.append(bias_node) + annotate_output_qspec(node, output_act_qspec) + _mark_nodes_as_annotated(nodes_to_mark_annotated) + annotated_partitions.append(nodes_to_mark_annotated) + + return annotated_partitions + + +@register_annotator("linear_relu") +def _annotate_linear_relu( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + annotated_partitions = [] + input_act_qspec = get_input_act_qspec(quantization_config) + output_act_qspec = get_output_act_qspec(quantization_config) + weight_qspec = get_weight_qspec(quantization_config) + bias_qspec = get_bias_qspec(quantization_config) + for node in gm.graph.nodes: + if node.op != "call_function" or node.target not in [ + torch.ops.aten.relu.default, + torch.ops.aten.relu_.default, + ]: + continue + relu_node = node + maybe_linear_node = node.args[0] + if ( + not isinstance(maybe_linear_node, Node) + or maybe_linear_node.op != "call_function" + or maybe_linear_node.target != torch.ops.aten.linear.default + ): + continue + + linear_node = maybe_linear_node + if len(linear_node.users) > 1: + # if linear node has multiple users, then it can't be fused with relu + continue + + input_qspec_map = {} + input_act = linear_node.args[0] + assert isinstance(input_act, Node) + input_qspec_map[input_act] = input_act_qspec + + weight = linear_node.args[1] + assert isinstance(weight, Node) + input_qspec_map[weight] = weight_qspec + + # adding weight node to the partition as well + partition = [relu_node, linear_node, weight] + bias = linear_node.args[2] if len(linear_node.args) > 2 else None + if isinstance(bias, Node): + input_qspec_map[bias] = bias_qspec + partition.append(bias) + + if _is_annotated(partition): + continue + + if filter_fn and any(not filter_fn(n) for n in partition): + continue + + linear_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + ) + relu_node.meta["quantization_annotation"] = QuantizationAnnotation( + output_qspec=output_act_qspec, + _annotated=True, + ) + _mark_nodes_as_annotated(partition) + annotated_partitions.append(partition) + return annotated_partitions + + +@register_annotator("conv") +def _annotate_conv( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + annotated_partitions = [] + for n in gm.graph.nodes: + if n.op != "call_function" or n.target not in [ + torch.ops.aten.conv1d.default, + torch.ops.aten.conv2d.default, + ]: + continue + conv_node = n + + input_qspec_map = {} + input_act = conv_node.args[0] + assert isinstance(input_act, Node) + input_qspec_map[input_act] = get_input_act_qspec(quantization_config) + + weight = conv_node.args[1] + assert isinstance(weight, Node) + input_qspec_map[weight] = get_weight_qspec(quantization_config) + + # adding weight node to the partition as well + partition = [conv_node, conv_node.args[1]] + + bias = conv_node.args[2] if len(conv_node.args) > 2 else None + if isinstance(bias, Node): + input_qspec_map[bias] = get_bias_qspec(quantization_config) + partition.append(bias) + + if _is_annotated(partition): + continue + + if filter_fn and any(not filter_fn(n) for n in partition): + continue + + conv_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=get_output_act_qspec(quantization_config), + _annotated=True, + ) + _mark_nodes_as_annotated(partition) + annotated_partitions.append(partition) + return annotated_partitions + + +def _do_annotate_conv_relu( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, + is_conv_transpose: bool = False, +): + annotated_partitions = [] + for n in gm.graph.nodes: + if n.op != "call_function" or n.target not in [ + torch.ops.aten.relu.default, + torch.ops.aten.relu_.default, + ]: + continue + relu_node = n + maybe_conv_node = n.args[0] + + is_conv_node = _is_conv_transpose_node if is_conv_transpose else _is_conv_node + if not isinstance(maybe_conv_node, Node) or not is_conv_node(maybe_conv_node): + continue + conv_node = maybe_conv_node + + if len(conv_node.users) > 1: + # relu shouldn't be fuseable to conv if there are other users + # of convolution + continue + + input_qspec_map = {} + input_act = conv_node.args[0] + assert isinstance(input_act, Node) + input_qspec_map[input_act] = get_input_act_qspec(quantization_config) + + weight = conv_node.args[1] + assert isinstance(weight, Node) + input_qspec_map[weight] = get_weight_qspec(quantization_config) + + # adding weight node to the partition as well + partition = [relu_node, conv_node, conv_node.args[1]] + bias = conv_node.args[2] if len(conv_node.args) > 2 else None + if isinstance(bias, Node): + input_qspec_map[bias] = get_bias_qspec(quantization_config) + partition.append(bias) + + if _is_annotated(partition): + continue + + if filter_fn and any(not filter_fn(n) for n in partition): + continue + + conv_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, _annotated=True + ) + relu_node.meta["quantization_annotation"] = QuantizationAnnotation( + output_qspec=get_output_act_qspec(quantization_config), # type: ignore[arg-type] + _annotated=True, + ) + _mark_nodes_as_annotated(partition) + annotated_partitions.append(partition) + return annotated_partitions + + +@register_annotator("conv_relu") +def _annotate_conv_relu( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + return _do_annotate_conv_relu( + gm, quantization_config, filter_fn, is_conv_transpose=False + ) + + +@register_annotator("conv_transpose_relu") +def _annotate_conv_transpose_relu( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + return _do_annotate_conv_relu( + gm, quantization_config, filter_fn, is_conv_transpose=True + ) + + +@register_annotator("conv_bn") +def _annotate_conv_bn( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + """ + Find conv + batchnorm parititions + Note: This is only used for QAT. In PTQ, batchnorm should already be fused into the conv. + """ + return _do_annotate_conv_bn(gm, quantization_config, filter_fn, has_relu=False) + + +@register_annotator("conv_bn_relu") +def _annotate_conv_bn_relu( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + """ + Find conv + batchnorm + relu parititions + Note: This is only used for QAT. In PTQ, batchnorm should already be fused into the conv. + """ + return _do_annotate_conv_bn(gm, quantization_config, filter_fn, has_relu=True) + + +@register_annotator("conv_transpose_bn") +def _annotate_conv_transpose_bn( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + """ + Find conv_transpose + batchnorm parititions + Note: This is only used for QAT. In PTQ, batchnorm should already be fused into the conv. + """ + return _do_annotate_conv_bn( + gm, quantization_config, filter_fn, has_relu=False, is_conv_transpose=True + ) + + +@register_annotator("conv_transpose_bn_relu") +def _annotate_conv_transpose_bn_relu( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + """ + Find conv_transpose + batchnorm + relu parititions + Note: This is only used for QAT. In PTQ, batchnorm should already be fused into the conv. + """ + return _do_annotate_conv_bn( + gm, quantization_config, filter_fn, has_relu=True, is_conv_transpose=True + ) + + +def _do_annotate_conv_bn( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]], + has_relu: bool, + is_conv_transpose: bool = False, +) -> list[list[Node]]: + """ + Given a function that takes in a `conv_fn` and returns a conv-bn[-relu] pattern, + return a list of annotated partitions. + + The output of the pattern must include a dictionary from string name to node + for the following names: "input", "conv", "weight", "bias", and "output". + """ + + # Example inputs for conv-bn1d patterns + _conv1d_bn_example_inputs = ( + torch.randn(1, 1, 3), # x + torch.randn(1, 1, 1), # conv_weight + torch.randn(1), # conv_bias + torch.randn(1), # bn_weight + torch.randn(1), # bn_bias + torch.randn(1), # bn_running_mean + torch.randn(1), # bn_running_var + ) + + # Example inputs for conv-bn2d patterns + _conv2d_bn_example_inputs = ( + torch.randn(1, 1, 3, 3), # x + torch.randn(1, 1, 1, 1), # conv_weight + torch.randn(1), # conv_bias + torch.randn(1), # bn_weight + torch.randn(1), # bn_bias + torch.randn(1), # bn_running_mean + torch.randn(1), # bn_running_var + ) + + def get_pattern(conv_fn: Callable, relu_is_inplace: bool): + def _conv_bn(x, conv_weight, conv_bias, bn_weight, bn_bias, bn_rm, bn_rv): + conv = conv_fn(x, conv_weight, conv_bias) + bn = F.batch_norm(conv, bn_rm, bn_rv, bn_weight, bn_bias, training=True) + if has_relu: + output = F.relu_(bn) if relu_is_inplace else F.relu(bn) + else: + output = bn + return output, { + "input": x, + "conv": conv, + "weight": conv_weight, + "bias": conv_bias, + "output": output, + } + + return WrapperModule(_conv_bn) + + # Needed for matching, otherwise the matches gets filtered out due to unused + # nodes returned by batch norm + gm.graph.eliminate_dead_code() + gm.recompile() + + matches = [] + if is_conv_transpose: + combinations = [ + (F.conv_transpose1d, _conv1d_bn_example_inputs), + (F.conv_transpose2d, _conv2d_bn_example_inputs), + ] + else: + combinations = [ + (F.conv1d, _conv1d_bn_example_inputs), # type: ignore[list-item] + (F.conv2d, _conv2d_bn_example_inputs), # type: ignore[list-item] + ] + + # Add `is_cuda` and `relu_is_inplace` dimensions + combinations = itertools.product( # type: ignore[assignment] + combinations, + [True, False] if torch.cuda.is_available() else [False], # is_cuda + [True, False] if has_relu else [False], # relu_is_inplace + ) + + # Match against all conv dimensions and cuda variants + for (conv_fn, example_inputs), is_cuda, relu_is_inplace in combinations: # type: ignore[misc] + pattern = get_pattern(conv_fn, relu_is_inplace) # type: ignore[has-type] + pattern = _get_aten_graph_module_for_pattern(pattern, example_inputs, is_cuda) # type: ignore[has-type] + pattern.graph.eliminate_dead_code() + pattern.recompile() + matcher = SubgraphMatcherWithNameNodeMap(pattern, ignore_literals=True) + matches.extend(matcher.match(gm.graph)) + + # Annotate nodes returned in the matches + annotated_partitions = [] + for match in matches: + name_node_map = match.name_node_map + input_node = name_node_map["input"] + conv_node = name_node_map["conv"] + weight_node = name_node_map["weight"] + bias_node = name_node_map["bias"] + output_node = name_node_map["output"] + + # TODO: annotate the uses of input, weight, and bias separately instead + # of assuming they come from a single conv node. This is not possible today + # because input may have multiple users, and we can't rely on the conv node + # always being the first user. This was the case in models with skip + # connections like resnet18 + + # Validate conv args + if conv_node.args[0] is not input_node: + raise ValueError("Conv arg did not contain input node ", input_node) + if conv_node.args[1] is not weight_node: + raise ValueError("Conv arg did not contain weight node ", weight_node) + if len(conv_node.args) > 2 and conv_node.args[2] is not bias_node: + raise ValueError("Conv arg did not contain bias node ", bias_node) + + # Skip if the partition is already annotated or is filtered out by the user + partition = [conv_node, weight_node] + if bias_node is not None: + partition.append(bias_node) + if _is_annotated(partition): + continue + if filter_fn and any(not filter_fn(n) for n in partition): + continue + + # Annotate conv inputs and pattern output + input_qspec_map = {} + input_qspec_map[input_node] = get_input_act_qspec(quantization_config) + input_qspec_map[weight_node] = get_weight_qspec(quantization_config) + if bias_node is not None: + input_qspec_map[bias_node] = get_bias_qspec(quantization_config) + conv_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + ) + output_node.meta["quantization_annotation"] = QuantizationAnnotation( + output_qspec=get_output_act_qspec(quantization_config), # type: ignore[arg-type] + _annotated=True, + ) + _mark_nodes_as_annotated(partition) + annotated_partitions.append(partition) + return annotated_partitions + + +@register_annotator("gru_io_only") +def _annotate_gru_io_only( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + gru_partitions = get_source_partitions(gm.graph, [torch.nn.GRU], filter_fn) + gru_partitions = list(itertools.chain.from_iterable(gru_partitions.values())) + annotated_partitions = [] + for gru_partition in gru_partitions: + annotated_partitions.append(gru_partition.nodes) + output_nodes = gru_partition.output_nodes + input_nodes = gru_partition.input_nodes + # skip annotation if it is already annotated + if _is_annotated(input_nodes + output_nodes): + continue + # inside each GRU partition, we should be able to annotate each linear + # subgraph + input_act = input_nodes[0] + input_act_user = next(iter(input_act.users.keys())) + assert isinstance(input_act, Node) + assert isinstance(input_act_user, Node) + input_act_user.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map={ + input_act: get_input_act_qspec(quantization_config), + }, + _annotated=True, + ) + + hidden_state = input_nodes[1] + hidden_state_user = next(iter(hidden_state.users.keys())) + assert isinstance(hidden_state, Node) + assert isinstance(hidden_state_user, Node) + hidden_state_user.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map={ + hidden_state: get_input_act_qspec(quantization_config), + }, + _annotated=True, + ) + + assert len(output_nodes) == 2, "expecting GRU to have two outputs" + for output in output_nodes: + output.meta["quantization_annotation"] = QuantizationAnnotation( + output_qspec=get_output_act_qspec(quantization_config), + _annotated=True, + ) + nodes_to_mark_annotated = list(gru_partition.nodes) + _mark_nodes_as_annotated(nodes_to_mark_annotated) + return annotated_partitions + + +@register_annotator("adaptive_avg_pool2d") +def _annotate_adaptive_avg_pool2d( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + """Always annotate adaptive_avg_pool2d op""" + module_partitions = get_source_partitions( + gm.graph, [torch.nn.AdaptiveAvgPool2d, F.adaptive_avg_pool2d], filter_fn + ) + partitions = list(itertools.chain.from_iterable(module_partitions.values())) + annotated_partitions = [] + for partition in partitions: + pool_node = partition.output_nodes[0] + if ( + pool_node.op != "call_function" + or pool_node.target != torch.ops.aten.adaptive_avg_pool2d.default + ): + raise ValueError(f"{pool_node} is not an aten adaptive_avg_pool2d operator") + + if _is_annotated([pool_node]): + continue + + annotated_partitions.append(partition.nodes) + input_act = pool_node.args[0] + assert isinstance(input_act, Node) + + # only annotate input output sharing operator + # when the output of the input node is annotated + if ( + "quantization_annotation" not in input_act.meta + or not input_act.meta["quantization_annotation"]._annotated + or input_act.meta["quantization_annotation"].output_qspec is None + ): + input_act_qspec = get_input_act_qspec(quantization_config) + else: + input_act_qspec = SharedQuantizationSpec(input_act) + + # output sharing with input + output_act_qspec = SharedQuantizationSpec((input_act, pool_node)) + pool_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map={ + input_act: input_act_qspec, + }, + output_qspec=output_act_qspec, + _annotated=True, + ) + return annotated_partitions + + +def _is_input_large_scalar(node: Node, gm: torch.fx.GraphModule): + """Check if input is a large scalar value. So that we can skip quantization for the node + since histc op (in HistogramObserver) only works for values up to certain upper bound + """ + if node.op == "get_attr": + qualified_name = str(node.target) + module_path, _, name = qualified_name.rpartition(".") + submod = gm.get_submodule(module_path) + tensor = getattr(submod, name) + # torch.histc works until this upper bound + HISTC_UPPER_BOUND = 3.4028235e15 + return tensor.numel() == 1 and abs(tensor.item()) > HISTC_UPPER_BOUND + return False + + +def _is_input_non_float_tensor(node: Node): + """Check if the input is not a float tensor, so that we can skip quantization for the node + since observers only works with float Tensors + """ + if "val" not in node.meta or not isinstance(node.meta["val"], FakeTensor): + return True + return node.meta["val"].dtype != torch.float32 + + +@register_annotator("add_relu") +def _annotate_add_relu( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + annotated_partitions = [] + for node in gm.graph.nodes: + if node.op != "call_function" or node.target not in [ + torch.ops.aten.relu.default, + torch.ops.aten.relu_.default, + ]: + continue + relu_node = node + maybe_add = node.args[0] + if ( + not isinstance(maybe_add, Node) + or maybe_add.op != "call_function" + or maybe_add.target + not in [ + torch.ops.aten.add.Tensor, + torch.ops.aten.add_.Tensor, + ] + ): + continue + + add_node = maybe_add + + if len(add_node.users) > 1: + # add can't be fused with ReLU if the result of add is being used + # else where in the graph + continue + + partition = [relu_node, add_node] + + if _is_annotated(partition): + continue + + if filter_fn and any(not filter_fn(n) for n in partition): + continue + + input_act_qspec = get_input_act_qspec(quantization_config) + output_act_qspec = get_output_act_qspec(quantization_config) + + input_qspec_map = {} + input_act0 = add_node.args[0] + if isinstance(input_act0, Node): + if _is_input_large_scalar(input_act0, gm): + continue + if _is_input_non_float_tensor(input_act0): + continue + partition.append(input_act0) + input_qspec_map[input_act0] = input_act_qspec + + input_act1 = add_node.args[1] + if isinstance(input_act1, Node): + if _is_input_large_scalar(input_act1, gm): + continue + if _is_input_non_float_tensor(input_act1): + continue + partition.append(input_act1) + input_qspec_map[input_act1] = input_act_qspec + + add_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + ) + relu_node.meta["quantization_annotation"] = QuantizationAnnotation( + output_qspec=output_act_qspec, + _annotated=True, + ) + annotated_partitions.append(partition) + return annotated_partitions + + +@register_annotator("add") +def _annotate_add( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + annotated_partitions = [] + for node in gm.graph.nodes: + if node.op != "call_function" or node.target not in [ + torch.ops.aten.add.Tensor, + torch.ops.aten.add_.Tensor, + ]: + continue + add_node = node + partition = [add_node] + + if _is_annotated(partition): + continue + + if filter_fn and any(not filter_fn(n) for n in partition): + continue + + input_act_qspec = get_input_act_qspec(quantization_config) + output_act_qspec = get_output_act_qspec(quantization_config) + + input_qspec_map = {} + input_act0 = add_node.args[0] + if isinstance(input_act0, Node): + if _is_input_large_scalar(input_act0, gm): + continue + if _is_input_non_float_tensor(input_act0): + continue + input_qspec_map[input_act0] = input_act_qspec + partition.append(input_act0) + + input_act1 = add_node.args[1] + if isinstance(input_act1, Node): + if _is_input_large_scalar(input_act1, gm): + continue + if _is_input_non_float_tensor(input_act1): + continue + input_qspec_map[input_act1] = input_act_qspec + partition.append(input_act1) + + add_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=output_act_qspec, + _annotated=True, + ) + annotated_partitions.append(partition) + return annotated_partitions + + +@register_annotator("mul_relu") +def _annotate_mul_relu( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + annotated_partitions = [] + for node in gm.graph.nodes: + if node.op != "call_function" or node.target not in [ + torch.ops.aten.relu.default, + torch.ops.aten.relu_.default, + ]: + continue + relu_node = node + maybe_mul = node.args[0] + if ( + not isinstance(maybe_mul, Node) + or maybe_mul.op != "call_function" + or maybe_mul.target + not in [ + torch.ops.aten.mul.Tensor, + torch.ops.aten.mul_.Tensor, + ] + ): + continue + + mul_node = maybe_mul + if len(mul_node.users) > 1: + # mul can't be fused with ReLU if the result of mul is being used + # else where in the graph + continue + + partition = [relu_node, mul_node] + + if _is_annotated(partition): + continue + + if filter_fn and any(not filter_fn(n) for n in partition): + continue + + input_act_qspec = get_input_act_qspec(quantization_config) + output_act_qspec = get_output_act_qspec(quantization_config) + + input_qspec_map = {} + input_act0 = mul_node.args[0] + if isinstance(input_act0, Node): + if _is_input_large_scalar(input_act0, gm): + continue + if _is_input_non_float_tensor(input_act0): + continue + partition.append(input_act0) + input_qspec_map[input_act0] = input_act_qspec + + input_act1 = mul_node.args[1] + if isinstance(input_act1, Node): + if _is_input_large_scalar(input_act1, gm): + continue + if _is_input_non_float_tensor(input_act1): + continue + partition.append(input_act1) + input_qspec_map[input_act1] = input_act_qspec + + mul_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + ) + relu_node.meta["quantization_annotation"] = QuantizationAnnotation( + output_qspec=output_act_qspec, + _annotated=True, + ) + annotated_partitions.append(partition) + return annotated_partitions + + +@register_annotator("mul") +def _annotate_mul( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + annotated_partitions = [] + for node in gm.graph.nodes: + if node.op != "call_function" or node.target not in [ + torch.ops.aten.mul.Tensor, + torch.ops.aten.mul_.Tensor, + ]: + continue + + mul_node = node + partition = [mul_node] + if _is_annotated(partition): + continue + + if filter_fn and any(not filter_fn(n) for n in partition): + continue + + input_act_qspec = get_input_act_qspec(quantization_config) + output_act_qspec = get_output_act_qspec(quantization_config) + + input_qspec_map = {} + input_act0 = mul_node.args[0] + if isinstance(input_act0, Node): + if _is_input_large_scalar(input_act0, gm): + continue + if _is_input_non_float_tensor(input_act0): + continue + input_qspec_map[input_act0] = input_act_qspec + partition.append(input_act0) + + input_act1 = mul_node.args[1] + if isinstance(input_act1, Node): + if _is_input_large_scalar(input_act1, gm): + continue + if _is_input_non_float_tensor(input_act1): + continue + input_qspec_map[input_act1] = input_act_qspec + partition.append(input_act0) + + mul_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=output_act_qspec, + _annotated=True, + ) + annotated_partitions.append(partition) + return annotated_partitions + + +# TODO: remove Optional in return type, fix annotated_partitions logic +@register_annotator("cat") +def _annotate_cat( + gm: torch.fx.GraphModule, + quantization_config: Optional[QuantizationConfig], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Optional[list[list[Node]]]: + cat_partitions = get_source_partitions(gm.graph, [torch.cat], filter_fn) + cat_partitions = list(itertools.chain.from_iterable(cat_partitions.values())) + annotated_partitions = [] + for cat_partition in cat_partitions: + cat_node = cat_partition.output_nodes[0] + if _is_annotated([cat_node]): + continue + + if cat_node.target != torch.ops.aten.cat.default: + # TODO: change this to AnnotationException + raise Exception( # noqa: TRY002 + f"Expected cat node: torch.ops.aten.cat.default, but found {cat_node.target}" + " please check if you are calling the correct capture API" + ) + + annotated_partitions.append(cat_partition.nodes) + + input_act_qspec = get_input_act_qspec(quantization_config) + inputs = cat_node.args[0] + + input_qspec_map = {} + input_act0 = inputs[0] # type: ignore[index] + if isinstance(input_act0, Node): + input_qspec_map[input_act0] = input_act_qspec + + shared_with_input0_qspec = SharedQuantizationSpec((input_act0, cat_node)) # type: ignore[arg-type] + for input_act in inputs[1:]: # type: ignore[index, union-attr] + if input_act not in input_qspec_map: + input_qspec_map[input_act] = shared_with_input0_qspec # type: ignore[index] + + output_act_qspec = shared_with_input0_qspec + + cat_node.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=output_act_qspec, + _annotated=True, + ) + return annotated_partitions + + +def _is_share_obs_or_fq_op(op: Callable) -> bool: + return op in [ + torch.ops.aten.relu.default, + torch.ops.aten.hardtanh.default, + torch.ops.aten.hardtanh_.default, + torch.ops.aten.max_pool2d.default, + torch.ops.aten.mean.default, + torch.ops.aten.mean.dim, + torch.ops.aten.permute.default, + torch.ops.aten.permute_copy.default, + torch.ops.aten.squeeze.dim, + torch.ops.aten.squeeze_copy.dim, + # TODO: remove? + torch.ops.aten.adaptive_avg_pool2d.default, + torch.ops.aten.view_copy.default, + torch.ops.aten.view.default, + torch.ops.aten.slice_copy.Tensor, + torch.ops.aten.flatten.using_ints, + ] + + +def propagate_annotation(model: torch.fx.GraphModule) -> None: + for n in model.graph.nodes: + if n.op != "call_function" or not _is_share_obs_or_fq_op(n.target): + continue + + prev_node = n.args[0] + if not isinstance(prev_node, Node): + continue + + quantization_annotation = prev_node.meta.get("quantization_annotation", None) + if not quantization_annotation: + continue + + output_qspec = quantization_annotation.output_qspec + if not output_qspec: + continue + + # make sure current node is not annotated + if ( + "quantization_annotation" in n.meta + and n.meta["quantization_annotation"]._annotated + ): + continue + + shared_qspec = SharedQuantizationSpec(prev_node) + # propagate the previous output_qspec to the current node + n.meta["quantization_annotation"] = QuantizationAnnotation( + input_qspec_map={ + prev_node: shared_qspec, + }, + output_qspec=shared_qspec, + _annotated=True, + ) + + +# TODO: make the list of ops customizable +def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModule: + for n in model.graph.nodes: + if n.op != "call_function" or n.target not in [ + torch.ops.aten.add.Tensor, + torch.ops.aten.mul.Tensor, + ]: + continue + args = list(n.args) + new_args = [] + for i in range(len(args)): + if isinstance(args[i], torch.fx.Node): + new_args.append(args[i]) + continue + prefix = "_tensor_constant_" + get_new_attr_name = get_new_attr_name_with_prefix(prefix) + tensor_constant_name = get_new_attr_name(model) + float_tensor = torch.tensor(float(args[i])) + model.register_buffer(tensor_constant_name, float_tensor) + fake_mode = n.meta["val"].fake_mode + with model.graph.inserting_before(n): + get_attr_node = model.graph.create_node( + "get_attr", tensor_constant_name, (), {} + ) + get_attr_node.meta["val"] = fake_mode.from_tensor( + float_tensor, static_shapes=True + ) + new_args.append(get_attr_node) + n.args = tuple(new_args) + model.recompile() + return model diff --git a/lib/python3.12/site-packages/torchao/testing/pt2e/utils.py b/lib/python3.12/site-packages/torchao/testing/pt2e/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..41be460a402d775b54f23b4172bb25840a8737d9 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/pt2e/utils.py @@ -0,0 +1,134 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import copy +import unittest + +import torch +from torch.ao.quantization.backend_config import ( + get_executorch_backend_config, +) +from torch.ao.quantization.quantize_fx import ( + _convert_to_reference_decomposed_fx, + prepare_fx, +) +from torch.testing._internal.common_quantization import ( + NodeSpec, + QuantizationTestCase, +) + +from torchao.quantization.pt2e.quantize_pt2e import ( + convert_pt2e, + prepare_pt2e, + prepare_qat_pt2e, +) +from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 + +if TORCH_VERSION_AT_LEAST_2_5: + from torch.export import export_for_training + + +@unittest.skipIf( + not TORCH_VERSION_AT_LEAST_2_5, + "only works for torch 2.5+ since export_for_training is only supported after 2.5", +) +class PT2EQuantizationTestCase(QuantizationTestCase): + """ + Base QuantizationTestCase for PT2 with some helper methods. + """ + + _MAP_TO_FX_TRACED_OPS = { + torch.ops.quantized_decomposed.quantize_per_tensor: torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor: torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_channel: torch.ops.quantized_decomposed.quantize_per_channel.default, + torch.ops.quantized_decomposed.dequantize_per_channel: torch.ops.quantized_decomposed.dequantize_per_channel.default, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor: torch.ops.quantized_decomposed.quantize_per_tensor.tensor, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor: torch.ops.quantized_decomposed.dequantize_per_tensor.tensor, + } + + def _test_quantizer( + self, + model, + example_inputs, + quantizer, + expected_node_occurrence, + expected_node_list=None, + check_against_fx_quant=False, + # TODO: remove the test if fx quant is removed from pytorch + fx_qconfig_mapping=None, + export_with_dynamic_shape=False, + is_qat=False, + is_debug_mode=False, + training_ir_node_occurrence=None, + ): + # resetting dynamo cache + torch._dynamo.reset() + m_eager = model.eval() + + # program capture + m = copy.deepcopy(m_eager) + dynamic_shapes = tuple( + {0: torch.export.Dim("dim")} if i == 0 else None + for i in range(len(example_inputs)) + ) + m = export_for_training( + m, + example_inputs, + dynamic_shapes=dynamic_shapes if export_with_dynamic_shape else None, + ).module() + + if is_qat: + m = prepare_qat_pt2e(m, quantizer) + else: + m = prepare_pt2e(m, quantizer) + if is_debug_mode: + print("prepared model:", m) + # Calibrate + m(*example_inputs) + m = convert_pt2e(m) + if is_debug_mode: + print("quantized model", m) + + pt2_quant_output = m(*example_inputs) + ns = NodeSpec + node_occurrence = { + ns.call_function(k): v for k, v in expected_node_occurrence.items() + } + if expected_node_list is None: + expected_node_list = [] + node_list = [ns.call_function(n) for n in expected_node_list] + self.checkGraphModuleNodes( + m, expected_node_occurrence=node_occurrence, expected_node_list=node_list + ) + if check_against_fx_quant: + qconfig_mapping = fx_qconfig_mapping + backend_config = get_executorch_backend_config() + m_copy = copy.deepcopy(m_eager) + m_fx = prepare_fx( + m_copy, qconfig_mapping, example_inputs, backend_config=backend_config + ) + m_fx(*example_inputs) + m_fx = _convert_to_reference_decomposed_fx( + m_fx, backend_config=backend_config + ) + m_fx = export_for_training( + m_fx, + example_inputs, + dynamic_shapes=dynamic_shapes if export_with_dynamic_shape else None, + ).module() + node_occurrence = {} + for k, v in PT2EQuantizationTestCase._MAP_TO_FX_TRACED_OPS.items(): + if k in expected_node_occurrence: + node_occurrence[ns.call_function(v)] = expected_node_occurrence[k] + if training_ir_node_occurrence is not None: + node_occurrence = { + ns.call_function(k): v + for k, v in training_ir_node_occurrence.items() + } + self.checkGraphModuleNodes(m_fx, expected_node_occurrence=node_occurrence) + fx_quant_output = m_fx(*example_inputs) + self.assertEqual(fx_quant_output, pt2_quant_output) + return m diff --git a/lib/python3.12/site-packages/torchao/testing/utils.py b/lib/python3.12/site-packages/torchao/testing/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..26a738c53bdc5c8a621013d0e57d9641c40a8559 --- /dev/null +++ b/lib/python3.12/site-packages/torchao/testing/utils.py @@ -0,0 +1,430 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +import copy +import functools +import unittest + +import torch +from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard +from torch.testing._internal import common_utils +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) + +import torchao +from torchao.dtypes import AffineQuantizedTensor, to_affine_quantized_intx +from torchao.quantization import int8_weight_only, quantize_ +from torchao.quantization.quant_primitives import MappingType +from torchao.utils import TORCH_VERSION_AT_LEAST_2_6, get_compute_capability + +""" +How to use: + +import unittest +from torchao.testing.utils import TorchAOBasicTestCase, copy_tests +from torch.testing._internal import common_utils + +# TODO: currently there is no way to set COMMON_DEVICES/COMMON_DTYPES +# we can figure out this a bit later + +# change arguments +class MyTestCase(TorchAOBasicTestCase): + TENSOR_SUBCLASS = MyDTypeTensor + FACTOR_FN = to_my_dtype + kwargs = {"target_dtype": torch.uint8} + LINEAR_MIN_SQNR = 30 + +# copy the instantiated tests +copy_tests(TorchAOBasicTestCase, MyTestCase, "my_test_case") + +if __name__ == "__main__": + unittest.main() +""" + + +def skip_if_compute_capability_less_than(min_capability): + import unittest + + def decorator(test_func): + def wrapper(*args, **kwargs): + if get_compute_capability() < min_capability: + raise unittest.SkipTest( + f"Compute capability is less than {min_capability}" + ) + return test_func(*args, **kwargs) + + return wrapper + + return decorator + + +def skip_if_rocm(message=None): + """Decorator to skip tests on ROCm platform with custom message. + + Args: + message (str, optional): Additional information about why the test is skipped. + """ + import pytest + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + if torch.version.hip is not None: + skip_message = "Skipping the test in ROCm" + if message: + skip_message += f": {message}" + pytest.skip(skip_message) + return func(*args, **kwargs) + + return wrapper + + # Handle both @skip_if_rocm and @skip_if_rocm() syntax + if callable(message): + func = message + message = None + return decorator(func) + return decorator + + +def skip_if_no_cuda(): + import unittest + + def decorator(test_func): + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + if not torch.cuda.is_available(): + raise unittest.SkipTest("No cuda available") + return test_func(*args, **kwargs) + + return wrapper + + return decorator + + +def skip_if_no_gemlite(): + import unittest + + def decorator(test_func): + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + try: + import gemlite # noqa: F401 + except: + raise unittest.SkipTest("No gemlite available") + return test_func(*args, **kwargs) + + return wrapper + + return decorator + + +# copied from https://github.com/pytorch/pytorch/blob/941d094dd1b507dacf06ddc6ed3485a9537e09b7/test/inductor/test_torchinductor.py#L11389 +def copy_tests(my_cls, other_cls, suffix, test_failures=None, xfail_prop=None): # noqa: B902 + for name, value in my_cls.__dict__.items(): + if name.startswith("test_"): + # You cannot copy functions in Python, so we use closures here to + # create objects with different ids. Otherwise, unittest.skip + # would modify all methods sharing the same object id. Also, by + # using a default argument, we create a copy instead of a + # reference. Otherwise, we would lose access to the value. + + @functools.wraps(value) + def new_test(self, value=value): + return value(self) + + # Copy __dict__ which may contain test metadata + new_test.__dict__ = copy.deepcopy(value.__dict__) + + if xfail_prop is not None and hasattr(value, xfail_prop): + new_test = unittest.expectedFailure(new_test) + + tf = test_failures and test_failures.get(name) + if tf is not None and suffix in tf.suffixes: + skip_func = ( + unittest.skip("Skipped!") + if tf.is_skip + else unittest.expectedFailure + ) + new_test = skip_func(new_test) + + setattr(other_cls, f"{name}_{suffix}", new_test) + + +class TorchAOBasicTestCase(common_utils.TestCase): + COMMON_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + COMMON_DTYPES = [torch.float32, torch.float16, torch.bfloat16] + + TENSOR_SUBCLASS = AffineQuantizedTensor + FACTORY_FN = to_affine_quantized_intx + kwargs = { + "mapping_type": MappingType.ASYMMETRIC, + "block_size": (1, 32), + "target_dtype": torch.uint8, + } + # minimum sqnr for linear operation when the weight is quantized to low precision + # with the above setting + LINEAR_MIN_SQNR = 40 + + def test_flatten_unflatten(self): + hp_tensor = torch.randn(4, 128) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + tensor_data_name_dict, tensor_attributes = lp_tensor.__tensor_flatten__() + tensor_data_dict = { + name: getattr(lp_tensor, name) for name in tensor_data_name_dict + } + outer_size = lp_tensor.size() + outer_stride = lp_tensor.stride() + reconstructed = self.TENSOR_SUBCLASS.__tensor_unflatten__( + tensor_data_dict, tensor_attributes, outer_size, outer_stride + ) + self.assertEqual(lp_tensor.dequantize(), reconstructed.dequantize()) + + @common_utils.parametrize("device", COMMON_DEVICES) + @common_utils.parametrize("dtype", COMMON_DTYPES) + def test_hp_tensor_device_dtype(self, device, dtype): + hp_tensor = torch.randn(4, 128, device=device, dtype=dtype) + self.FACTORY_FN(hp_tensor, **self.kwargs) + + @common_utils.parametrize("device1", COMMON_DEVICES) + @common_utils.parametrize("device2", COMMON_DEVICES) + def test_device1_to_device2(self, device1, device2): + """Note: this should be parametrized with device1 and device2 + e.g. device1 = ["cpu", "cuda"], device2 = ["cpu", "cuda"] + """ + hp_tensor = torch.randn(4, 128, device=device1, dtype=torch.bfloat16) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + lp_tensor.to(device=device2) + + hp_tensor = torch.randn(4, 128, device=device1, dtype=torch.bfloat16) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + lp_tensor.to(device2) + + hp_tensor = torch.randn(4, 128, device=device1, dtype=torch.bfloat16) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + lp_tensor.cuda() + + hp_tensor = torch.randn(4, 128, device=device1, dtype=torch.bfloat16) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + lp_tensor.cpu() + + @common_utils.parametrize("device", COMMON_DEVICES) + @common_utils.parametrize("dtype", COMMON_DTYPES) + def test_transpose(self, device, dtype): + hp_tensor = torch.randn(4, 128, device=device, dtype=dtype) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + lp_tensor = lp_tensor.t() + self.assertEqual(lp_tensor.shape, (128, 4)) + + @common_utils.parametrize("device", COMMON_DEVICES) + @common_utils.parametrize("dtype", COMMON_DTYPES) + def test_linear(self, device, dtype): + hp_tensor = torch.randn(4, 128, device=device, dtype=dtype) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + + hp_act_tensor = torch.randn(32, 128, device=device, dtype=dtype) + hp_res = torch.nn.functional.linear(hp_act_tensor, hp_tensor) + lp_res = torch.nn.functional.linear(hp_act_tensor, lp_tensor) + self.assertGreater( + torchao.quantization.utils.compute_error(hp_res, lp_res), + self.LINEAR_MIN_SQNR, + ) + + +class TorchAOCompileTestCase(common_utils.TestCase): + COMMON_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + COMMON_DTYPES = [torch.float32, torch.float16, torch.bfloat16] + + TENSOR_SUBCLASS = AffineQuantizedTensor + FACTORY_FN = to_affine_quantized_intx + kwargs = { + "mapping_type": MappingType.ASYMMETRIC, + "block_size": (1, 32), + "target_dtype": torch.uint8, + } + # minimum sqnr for linear operation when the weight is quantized to low precision + # with the above setting + LINEAR_MIN_SQNR = 40 + COMPILE_MIN_SQNR = 50 + + @common_utils.parametrize("device", COMMON_DEVICES) + @common_utils.parametrize("dtype", COMMON_DTYPES) + def test_input_output_tensor_subclass(self, device, dtype): + hp_tensor = torch.randn(4, 128, device=device, dtype=dtype) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + + def f(tensor): + return tensor + + ref = f(lp_tensor) + f = torch.compile(f) + compiled = f(lp_tensor) + self.assertTrue(isinstance(f(lp_tensor), self.TENSOR_SUBCLASS)) + self.assertEqual(ref.dequantize(), compiled.dequantize()) + + @common_utils.parametrize("device", COMMON_DEVICES) + @common_utils.parametrize("dtype", COMMON_DTYPES) + def test_input_tensor_subclass(self, device, dtype): + hp_tensor = torch.randn(4, 128, device=device, dtype=dtype) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + + def f(tensor): + return tensor.dequantize() + + ref = f(lp_tensor) + f = torch.compile(f) + compiled = f(lp_tensor) + self.assertFalse(isinstance(f(lp_tensor), self.TENSOR_SUBCLASS)) + self.assertEqual(ref, compiled) + + @common_utils.parametrize("device", COMMON_DEVICES) + @common_utils.parametrize("dtype", COMMON_DTYPES) + def test_output_tensor_subclass(self, device, dtype): + hp_tensor = torch.randn(4, 128, device=device, dtype=dtype) + + def f(hp_tensor): + return self.FACTORY_FN(hp_tensor, **self.kwargs) + + ref = f(hp_tensor) + f = torch.compile(f) + compiled = f(hp_tensor) + self.assertTrue(isinstance(f(hp_tensor), self.TENSOR_SUBCLASS)) + # bfloat16 seems to result in much larger numerical differences + if dtype != torch.bfloat16: + self.assertGreater( + torchao.quantization.utils.compute_error( + ref.dequantize(), compiled.dequantize() + ), + self.COMPILE_MIN_SQNR, + ) + + @common_utils.parametrize("device", COMMON_DEVICES) + @common_utils.parametrize("dtype", COMMON_DTYPES) + def test_linear_compile(self, device, dtype): + hp_tensor = torch.randn(4, 128, device=device, dtype=dtype) + lp_tensor = self.FACTORY_FN(hp_tensor, **self.kwargs) + + hp_act_tensor = torch.randn(32, 128, device=device, dtype=dtype) + hp_res = torch.nn.functional.linear(hp_act_tensor, hp_tensor) + linear = torch.nn.Linear(128, 4, bias=False, device=device, dtype=dtype) + linear.weight = torch.nn.Parameter(lp_tensor) + lp_res = torch.compile(linear)(hp_act_tensor) + self.assertGreater( + torchao.quantization.utils.compute_error(hp_res, lp_res), + self.LINEAR_MIN_SQNR, + ) + + +class TorchAOTensorParallelTestCase(DTensorTestBase): + """Basic test case for tensor subclasses""" + + COMMON_DTYPES = [torch.float32, torch.float16, torch.bfloat16] + + TENSOR_SUBCLASS = AffineQuantizedTensor + QUANT_METHOD_FN = staticmethod(int8_weight_only) + QUANT_METHOD_KWARGS = {} + + @staticmethod + def colwise_shard(m: torch.nn.Module, mesh: DeviceMesh) -> torch.nn.Module: + """ + Shard linear layer of the model in column-wise fashion + """ + # Column-wise is wrt to A^T, so for A it is row-wise. + # Number of rows per rank + orig_weight = m.linear.weight + n_local_rows = orig_weight.size(0) // mesh.size() + rank = mesh.get_local_rank() + local_shard = orig_weight[rank * n_local_rows : (rank + 1) * n_local_rows, :] + # Construct DTensor from local shard + dtensor = DTensor.from_local(local_shard, mesh, [Shard(0)]) + # Replace parameter in module + m.linear.weight = torch.nn.Parameter(dtensor, requires_grad=False) + return m + + @staticmethod + def rowwise_shard(m: torch.nn.Module, mesh: DeviceMesh) -> torch.nn.Module: + """ + Shard linear layer of the model in row-wise fashion + """ + # Row-wise is wrt to A^T, so for A it is column-wise. + # Number of rows per rank + orig_weight = m.linear.weight + n_local_cols = orig_weight.size(1) // mesh.size() + rank = mesh.get_local_rank() + local_shard = orig_weight[:, rank * n_local_cols : (rank + 1) * n_local_cols] + # Construct DTensor from local shard + dtensor = DTensor.from_local(local_shard, mesh, [Shard(1)]) + # Replace parameter in module + m.linear.weight = torch.nn.Parameter(dtensor, requires_grad=False) + return m + + def quantize(self, m: torch.nn.Module) -> torch.nn.Module: + """ + Quantize the model + """ + quantize_(m, self.QUANT_METHOD_FN(**self.QUANT_METHOD_KWARGS)) + return m + + @common_utils.parametrize("dtype", COMMON_DTYPES) + @with_comms + @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available") + def test_tp(self, dtype): + device = "cuda" + # To make sure different ranks create the same module + torch.manual_seed(5) + + class M(torch.nn.Module): + def __init__(self, in_features, out_features, **kwargs) -> None: + super().__init__(**kwargs) + self.linear = torch.nn.Linear( + in_features, out_features, bias=False, device="cuda" + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + # Get rank and device + device = torch.device(f"cuda:{self.rank % torch.cuda.device_count()}") + + # Original model + proj_up = M(1024, 2048).to(device).to(dtype) + proj_dn = M(2048, 1024).to(device).to(dtype) + example_input = 100 * torch.randn(128, 1024, device=device, dtype=dtype) + proj_dn(proj_up(example_input)) + + # Quantize the model + up_quant = self.quantize(proj_up) + dn_quant = self.quantize(proj_dn) + dn_quant(up_quant(example_input)) + + mesh = self.build_device_mesh() + mesh.device_type = "cuda" + + # Shard the models + up_dist = self.colwise_shard(up_quant, mesh) + dn_dist = self.rowwise_shard(dn_quant, mesh) + + # We need to turn inputs into DTensor form as well -- just a format change + input_dtensor = DTensor.from_local(example_input, mesh, [Replicate()]) + + dn_dist(up_dist(input_dtensor)) + + if not TORCH_VERSION_AT_LEAST_2_6: + # Need torch 2.6 to support compiled tensor parallelism + return + + up_compiled = torch.compile(up_dist) + y_up = up_compiled(input_dtensor) + dn_compiled = torch.compile(dn_dist) + dn_compiled(y_up) + + +common_utils.instantiate_parametrized_tests(TorchAOBasicTestCase) +common_utils.instantiate_parametrized_tests(TorchAOCompileTestCase) +common_utils.instantiate_parametrized_tests(TorchAOTensorParallelTestCase) + +if __name__ == "__main__": + unittest.main()