# 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 from typing import Any, List, Tuple import torch def int8_symmetric_quantize( fp32_tensor: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Symmetrically quantize the torch.float32 tensor into torch.int8. Return a 2-tuple of (quantized value, scale). """ quant_min = -128 quant_max = 127 min_val = torch.amin(fp32_tensor, dim=[1], keepdim=False) max_val = torch.amax(fp32_tensor, dim=[1], keepdim=False) min_val_neg = torch.min(min_val, torch.zeros_like(min_val)) max_val_pos = torch.max(max_val, torch.zeros_like(max_val)) max_val_pos = torch.max(-min_val_neg, max_val_pos) scale = max_val_pos / (float(quant_max - quant_min) / 2) scale = scale.view(fp32_tensor.shape[0], -1) out = torch.round(fp32_tensor * (1.0 / scale)) out = torch.clamp(out, quant_min, quant_max).to(torch.int8) return out, scale # Our subclass represents a tensor that has been quantized to int8 # It will hold two inner tensors: # - int_data: int8[M, N] # - scale: fp32[M, 1] class Int8SymmetricTensor(torch.Tensor): @staticmethod @torch._dynamo.disable def __new__(cls, int_data: torch.Tensor, scale: torch.Tensor): return torch.Tensor._make_wrapper_subclass( cls, int_data.shape, strides=int_data.stride(), storage_offset=int_data.storage_offset(), dtype=scale.dtype, device=int_data.device, ) @torch._dynamo.disable def __init__(self, int_data: torch.Tensor, scale: torch.Tensor): # inner data expected to be quantized already assert int_data.dtype is torch.int8 # we could do more work to support ndim > 2! assert int_data.ndim == 2 assert scale.ndim == 2 self.int_data = int_data self.scale = scale # __tensor_flatten__ returns a tuple of: # - names of all inner tensor attributes (two in our case) # - any other additional, non-tensor metadata. def __tensor_flatten__(self) -> Tuple[List[str], Any]: return ["int_data", "scale"], None # __tensor_unflatten__ should effectively undo __tensor_flatten__. # inputs: # - a dict mapping names of inner tensor attributes back to the tensors # - the constant metadata from __tensor_flatten__ # output: # - a new instance of your subclass @classmethod def __tensor_unflatten__( cls, tensor_data_dict, extra_metadata, outer_size=None, outer_stride=None ): assert extra_metadata is None int_data = tensor_data_dict["int_data"] scale = tensor_data_dict["scale"] return Int8SymmetricTensor(int_data, scale) def __repr__(self): return f"Int8SymmetricTensor(int_data={repr(self.int_data)}, scale={repr(self.scale)})" # Actually performs the symmetric quantization. # In our simple inference example we will quantize weights "ahead-of-time", # although later in a training example we can quantize/dequantize # during model execution, inside of our __torch_dispatch__ # input: # - float32 torch.Tensor # output: # - Int8SymmetricTensor @staticmethod def from_float(float_tensor): int8_tensor, scale = int8_symmetric_quantize(float_tensor) return Int8SymmetricTensor(int8_tensor, scale) # __torch_dispatch__ gets called for ATen operator # that our subclass is passed as an input to. # We need to define our own implementation for every operator here. @classmethod def __torch_dispatch__(cls, func, types, args, kwargs): if kwargs is None: kwargs = {} if func not in op_implementations_dict: raise AssertionError( f"Int8SymmetricTensor does not yet support op: {str(func)}" ) return op_implementations_dict[func](func, *args, **kwargs) # Convenience function for registering our own implementation # to every ATen operator in PyTorch op_implementations_dict = {} def register_op(ops: List[torch._ops.OpOverload]): def impl_decorator(op_impl): global op_implementations_dict for op in ops: op_implementations_dict[op] = op_impl return op_impl return impl_decorator from torch.utils._python_dispatch import return_and_correct_aliasing # matmul impl @register_op([torch.ops.aten.mm.default]) def int8_mm(func, x, weight): assert isinstance(weight, Int8SymmetricTensor), ( "Int8SymmetricTensor: matmul currently only supports the weight in low precision, not the input!" ) return torch.mm(x, weight.int_data.to(x.dtype)) * weight.scale # implementation of most view operations @register_op( [ torch.ops.aten.detach.default, torch.ops.aten.t.default, torch.ops.aten.view.default, torch.ops.aten._unsafe_view.default, ] ) def int8_view_ops(func, *args, **kwargs): assert isinstance(args[0], Int8SymmetricTensor) out_data = func(args[0].int_data, *args[1:], **kwargs) out_scale = func(args[0].scale, *args[1:], **kwargs) out = Int8SymmetricTensor(out_data, out_scale) # "return_and_correct_aliasing" here is needed for torch.compile support. # It effectively tells the compiler that the output of this view op aliases its input. # At some point, we're hoping to infer this automatically and kill this extra API! return return_and_correct_aliasing(func, args, kwargs, out) class ToyModel(torch.nn.Module): def __init__(self, m: int, n: int, k: int): super().__init__() self.linear1 = torch.nn.Linear(m, n, bias=False) self.linear2 = torch.nn.Linear(n, k, bias=False) def forward(self, x): x = self.linear1(x) x = self.linear2(x) return x if __name__ == "__main__": # Set up toy model float_model = ToyModel(64, 128, 32).cuda() quantized_model_subclass = copy.deepcopy(float_model) # Swap torch.nn.Linear weights with Int8SymmetricTensor subclasses for name, child in quantized_model_subclass.named_children(): if type(child) == torch.nn.Linear: subclass_param = Int8SymmetricTensor.from_float(child.weight) child.weight = torch.nn.Parameter(subclass_param, requires_grad=True) with torch.no_grad(): x = torch.randn(64, 64, 64, device="cuda") out = quantized_model_subclass(x) # We can also use torch.compile to fuse some of our quantized logic # run with TORCH_LOGS="output_code" to see the generated inductor code out_compiled = torch.compile(quantized_model_subclass)(x) print(torch.allclose(out, out_compiled))