A pytorch Quantization Toolkit
import timeit

import torch


def mm(A, B):
    return torch._int_mm(A, B)


A = torch.randint(1, 10, [2400, 2400]).type(torch.int8).cuda()
B = torch.randint(1, 10, [2400, 2400]).type(torch.int8).cuda()
print(A)

# Warmup (slow)
mm(A, B)
# Average on several calls
it = 100
print(timeit.Timer(lambda: mm(A, B)).timeit(it) / it)

# Warmup (slow)
A = A.type(torch.float16)
B = B.type(torch.float16)
torch.matmul(A, B)
# Average on several calls
it = 100
print(timeit.Timer(lambda: torch.matmul(A, B)).timeit(it) / it)

bench/test_int_mm.py
huggingface/quanto
import timeit

import torch


def mm(a, b):
    return torch._int_mm(a, b)


A = torch.randint(1, 10, [2400, 2400]).type(torch.int8).cuda()
B = torch.randint(1, 10, [2400, 2400]).type(torch.int8).cuda()
it = 100

# Warmup (slow)
mm(A, B)
# Get a reference
print(timeit.Timer(lambda: mm(A, B)).timeit(it) / it)

cmm = torch.compile(mm, backend="inductor")
# First invocation will trigger the actual compilation
cmm(A, B)
# Now compare execution time
print(timeit.Timer(lambda: cmm(A, B)).timeit(it) / it)

bench/test_int_mm_inductor.py
huggingface/quanto
import timeit

import torch
from torch_int._CUDA import bmm_s8t_s8n_f32t


def mm(A, B):
    return torch._int_mm(A.squeeze(), B.squeeze().transpose(1, 0))


A = torch.randint(1, 10, [1, 512, 12288]).type(torch.int8).cuda()
B = torch.randint(1, 10, [1, 512, 12288]).type(torch.int8).cuda()
print(A)

# Using torch int matmul
# Warmup (slow)
mm(A, B)
# Average on several calls
it = 1000
print("torch _int_mm")
print(timeit.Timer(lambda: mm(A, B)).timeit(it) / it)

# Using torch_int custom kernels
# Warmup (slow)
bmm_s8t_s8n_f32t(A, B, 0.1)
# Average on several calls
it = 1000
print("torch_int kernels")
print(timeit.Timer(lambda: bmm_s8t_s8n_f32t(A, B, 0.1)).timeit(it) / it)

# Using torch f16 matmul
# Warmup (slow)
A = A.type(torch.float16)
B = B.type(torch.float16)
torch.matmul(A, B.transpose(2, 1))
# Average on several calls
it = 1000
print("torch fp16 matmul")
print(timeit.Timer(lambda: torch.matmul(A, B.transpose(2, 1))).timeit(it) / it)

bench/test_int_mm_torch_int.py
huggingface/quanto
import argparse
import os
from tempfile import TemporaryDirectory

import torch
import torch.nn.functional as F
from torchvision import datasets, transforms
from transformers import AutoModel

from quanto.quantization import QLinear, QTensor, freeze, quantize
from quanto.quantization.calibrate import calibration


def test(model, device, test_loader):
    model.to(device)
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            if isinstance(output, QTensor):
                output = output.dequantize()
            test_loss += F.nll_loss(output, target, reduction="sum").item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print(
        "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
            test_loss, correct, len(test_loader.dataset), 100.0 * correct / len(test_loader.dataset)
        )
    )


def train(log_interval, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        if isinstance(output, QTensor):
            output = output.dequantize()
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % log_interval == 0:
            print(
                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                    epoch,
                    batch_idx * len(data),
                    len(train_loader.dataset),
                    100.0 * batch_idx / len(train_loader),
                    loss.item(),
                )
            )


def print_quantization_stats(model):
    for name, m in model.named_modules():
        if isinstance(m, QLinear):
            print(f"{name} quantization stats:")
            qweight = QTensor.quantize(m.weight)
            weight_mae = torch.nn.L1Loss()(qweight.dequantize(), m.weight)
            weight_stats = f"  weight mae = {weight_mae}"
            if m.bias is not None:
                bias_scale = m.in_scale * qweight._scale
                qbias = QTensor.quantize(m.bias, torch.int32, bias_scale)
                bias_mae = torch.nn.L1Loss()(qbias.dequantize(), m.bias)
                weight_stats += f", bias mae = {bias_mae}"
            print(weight_stats)
            print(f"  scale: in = {m.in_scale}, out = {m.out_scale}")


def main():
    # Training settings
    parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
    parser.add_argument(
        "--batch-size", type=int, default=250, metavar="N", help="input batch size for testing (default: 250)"
    )
    parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)")
    parser.add_argument("--model", type=str, default="dacorvo/mnist-mlp", help="The name of the trained Model.")
    parser.add_argument("--stats", action="store_true", default=False, help="Display quantization statistics")
    args = parser.parse_args()

    torch.manual_seed(args.seed)

    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    dataset_kwargs = {"batch_size": args.batch_size}
    if torch.cuda.is_available():
        cuda_kwargs = {"num_workers": 1, "pin_memory": True, "shuffle": True}
        dataset_kwargs.update(cuda_kwargs)

    transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,)),
            transforms.Lambda(lambda x: torch.flatten(x)),
        ]
    )
    dataset1 = datasets.MNIST("./data", train=True, download=True, transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1, **dataset_kwargs)
    dataset2 = datasets.MNIST("./data", train=False, download=True, transform=transform)
    test_loader = torch.utils.data.DataLoader(dataset2, **dataset_kwargs)
    model = AutoModel.from_pretrained(args.model, trust_remote_code=True)
    model.eval()
    # Test inference for reference
    print("Float model")
    test(model, device, test_loader)
    # Quantize model
    quantize(model)
    # Test inference (should be lower than float)
    print("Quantized model")
    test(model, device, test_loader)
    # Test inference with calibration (should be equivalent to float)
    print("Quantized calibrated model")
    with calibration():
        test(model, device, test_loader)
    print("Tuning quantized model for one epoch")
    optimizer = torch.optim.Adadelta(model.parameters(), lr=0.5)
    train(50, model, device, train_loader, optimizer, 1)
    print("Quantized tuned model")
    test(model, device, test_loader)
    print("Quantized frozen model")
    freeze(model)
    test(model, device, test_loader)
    if args.stats:
        print_quantization_stats(model)
    # Now save the model and reload it to verify quantized weights are restored
    with TemporaryDirectory() as tmpdir:
        mlp_file = os.path.join(tmpdir, "mlp.pt")
        torch.save(model.state_dict(), mlp_file)
        # Reinstantiate a model with float weights
        model_reloaded = AutoModel.from_pretrained(args.model, trust_remote_code=True)
        quantize(model_reloaded)
        # When reloading we must assign instead of copying to force quantized tensors assignment
        model_reloaded.load_state_dict(torch.load(mlp_file), assign=True)
    print("Quantized model with serialized integer weights")
    test(model_reloaded, device, test_loader)


if __name__ == "__main__":
    main()

examples/quantize_mnist_model.py
huggingface/quanto
from .calibrate import *
from .nn import *
from .qtensor import *
from .quantize import *

quanto/quantization/__init__.py
huggingface/quanto
from contextlib import contextmanager

import torch
from torch.nn.modules.module import (
    register_module_forward_hook,
    register_module_forward_pre_hook,
)

from .nn import QLinear
from .qtensor import QTensor


__all__ = ["calibration"]


momentum = 0.9


def calibrate_input(module: torch.nn.Module, input):
    if isinstance(module, QLinear):
        # We want to requantize with the most accurate scale
        input = input[0]
        if isinstance(input, QTensor):
            input = input.dequantize()
        input = QTensor.quantize(input, torch.int8)
        if torch.all(module.in_scale == 1):
            module.in_scale = input._scale
        else:
            module.in_scale = momentum * module.in_scale + input._scale * (1.0 - momentum)
        return input


def calibrate_output(module: torch.nn.Module, input, output):
    if isinstance(module, QLinear):
        # Reevaluate output using float path
        input = input[0].dequantize()
        output = super(module.__class__, module).forward(input)
        # Requantize output with the most accurate scale
        output = QTensor.quantize(output, torch.int8)
        if torch.all(module.out_scale == 1):
            module.out_scale = output._scale
        else:
            module.out_scale = momentum * module.out_scale + output._scale * (1.0 - momentum)
        return output


@contextmanager
def calibration():
    """A context to calibrate quantized modules."""
    try:
        pre_handle = register_module_forward_pre_hook(calibrate_input)
        post_handle = register_module_forward_hook(calibrate_output)
        yield
    finally:
        pre_handle.remove()
        post_handle.remove()

quanto/quantization/calibrate.py
huggingface/quanto
import torch

from ..qtensor import QTensor


class QLinear(torch.nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        dtype=None,
        device=None,
    ) -> None:
        super().__init__(in_features, out_features, bias, dtype, device)
        self.register_buffer("in_scale", torch.ones((), dtype=torch.float32))
        self.register_buffer("out_scale", torch.ones((), dtype=torch.float32))

    @classmethod
    def from_module(cls, module):
        qmodule = cls(module.in_features, module.out_features, module.bias is not None)
        with torch.no_grad():
            qmodule.weight.copy_(module.weight)
            if module.bias is not None:
                qmodule.bias.copy_(module.bias)
        return qmodule.to(module.weight.device)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        # If needed, quantize inputs, weights and bias
        if isinstance(input, QTensor):
            if input._data.dtype == torch.int32:
                # Reduce input bitwidth
                input = input.rescale(torch.int8, self.in_scale)
        else:
            input = QTensor.quantize(input, torch.int8, self.in_scale)
        weight = self.weight
        if not isinstance(weight, QTensor):
            weight = QTensor.quantize(weight)
        bias = self.bias
        if bias is not None:
            bias_scale = self.in_scale * weight._scale
            if isinstance(bias, QTensor):
                if bias._scale != bias_scale:
                    # This should only happen if we calibrate again a frozen module
                    bias = QTensor.rescale(torch.int32, bias_scale)
            else:
                bias = QTensor.quantize(bias, torch.int32, bias_scale)
        # Operate on quantized tensors
        out_int32 = torch.nn.functional.linear(input, weight, bias)
        # Downscale
        return out_int32.rescale(torch.int8, self.out_scale)

    def freeze(self):
        # Replace float weights by quantized weights
        self.weight = torch.nn.Parameter(QTensor.quantize(self.weight).to(self.weight.device))
        if self.bias is not None:
            bias_scale = self.in_scale * self.weight._scale
            self.bias = torch.nn.Parameter(QTensor.quantize(self.bias, torch.int32, bias_scale))

quanto/quantization/nn/linear.py
huggingface/quanto
from functools import partial

import torch
from torch.autograd import Function


class Quantizer(Function):
    """A standard affine quantizer.

    If the quantization scale is not specified, then it uses the optimal scale
    for the base tensor value range.
    """

    @staticmethod
    def forward(ctx, base, int_dtype=torch.int8, scale=None):
        iinfo = torch.iinfo(int_dtype)
        if scale is None:
            scale = torch.max(torch.abs(base)) / torch.iinfo(int_dtype).max
        data = torch.clamp(torch.round(base / scale), min=iinfo.min, max=iinfo.max).to(int_dtype)
        # The instantiation of the quantized tensor must happen within the context of the Function
        # for the autograd magic to work.
        return QTensor(data, scale)

    @staticmethod
    def backward(ctx, gO):
        # For autograd, quantization is a no-op
        return gO, None, None


class Dequantizer(Function):
    @staticmethod
    def forward(ctx, t):
        return t._scale * t._data

    @staticmethod
    def backward(ctx, gO):
        # For autograd, dequantization is a no-op
        return gO


class ReQuantizer(Function):
    @staticmethod
    def forward(ctx, base, int_dtype=torch.int8, scale=None):
        dst_iinfo = torch.iinfo(int_dtype)
        if scale is None:
            if int_dtype == base._data.dtype:
                return base
            # Assuming the base scale is correct, simply project to the target integer range
            src_iinfo = torch.iinfo(base._data.dtype)
            int_rescale = dst_iinfo.max / src_iinfo.max
            scale = base._scale * src_iinfo.max / dst_iinfo.max
        else:
            # It is up to the caller to make sure the scale is consistent with the target int dtype
            int_rescale = base._scale / scale
        data = torch.clamp(torch.round(base._data * int_rescale), min=dst_iinfo.min, max=dst_iinfo.max).to(int_dtype)
        # The instantiation of the quantized tensor must happen within the context of the Function
        # for the autograd magic to work.
        return QTensor(data, scale)

    @staticmethod
    def backward(ctx, gO):
        # For autograd, requantization is a no-op
        return gO, None, None


_QTENSOR_DISPATCH_TABLE = {}


def register_dispatch(aten_ops):
    """
    Used for registering a new __torch_dispatch__ function to QTensor
    Called via _QTENSOR_DISPATCH_TABLE[func](func, *args, **kwargs)

    The code to register a new function looks like:

    @register_dispatch(list_of_ops)
    def foo(func, *args, **kwargs):
        <implementation>
    """

    def wrapper(func):
        for aten_op in aten_ops:
            _QTENSOR_DISPATCH_TABLE[aten_op] = partial(func, aten_op)

    return wrapper


def dequantize(*args):
    return [arg.dequantize() if isinstance(arg, QTensor) else arg for arg in args]


@register_dispatch([torch.ops.aten._to_copy])
def _to_copy(func, t, dtype=None, **kwargs):
    # Ignore dtype and use the inner data tensors dtypes instead
    out_data = func(t._data, dtype=t._data.dtype, **kwargs)
    out_scale = func(t._scale, dtype=t._scale.dtype, **kwargs)
    return QTensor(out_data, out_scale)


@register_dispatch([torch.ops.aten.detach])
def detach(func, t):
    # Detach both data and scale
    out_data = func(t._data)
    out_scale = func(t._scale)
    return QTensor(out_data, out_scale)


@register_dispatch([torch.ops.aten.add])
def add(func, input, other, alpha=1, out=None):
    if alpha != 1 or out is not None:
        raise ValueError("alpha and out parameters are not supported for quantized {func}.")
    if not torch.equal(input._scale, other._scale):
        raise ValueError("Quantized tensors with different scales cannot be added.")
    # We need to perform the operation in int16 because it might overflow
    out_data = func(input._data.to(torch.int16), other._data.to(torch.int16))
    out_scale = input._scale
    return QTensor(out_data, out_scale)


@register_dispatch([torch.ops.aten.addmm])
def addmm(func, input, mat1, mat2, beta=1, alpha=1):
    # Do the operation with int8 cast to float32
    out_data = func(
        input._data.to(torch.float32),
        mat1._data.to(torch.float32),
        mat2._data.to(torch.float32),
        beta=beta,
        alpha=alpha,
    )
    out_scale = mat1._scale * mat2._scale
    return QTensor(out_data.to(torch.int32), out_scale)


@register_dispatch([torch.ops.aten.copy_])
def copy_(func, dest, src):
    dest._data = func(dest._data, src._data)
    dest._scale = func(dest._scale, src._scale)
    return dest


@register_dispatch([torch.ops.aten.dot])
def dot(func, input, other):
    # Cast int8 data to float32 and do the operation
    out_data = func(input._data.to(torch.float32), other._data.to(torch.float32))
    out_scale = input._scale * other._scale
    return QTensor(out_data.to(torch.int32), out_scale)


@register_dispatch([torch.ops.aten.is_same_size])
def is_same_size(func, input, other):
    a = input._data if isinstance(input, QTensor) else input
    b = other._data if isinstance(other, QTensor) else other
    return func(a, b)


@register_dispatch([torch.ops.aten.bmm, torch.ops.aten.mm])
def mm(func, input, other):
    if not isinstance(input, QTensor) or not isinstance(other, QTensor):
        return func(*dequantize(input, other))
    # Cast int8 data to float32 and do the operation
    out_data = func(input._data.to(torch.float32), other._data.to(torch.float32))
    out_scale = input._scale * other._scale
    return QTensor(out_data.to(torch.int32), out_scale)


@register_dispatch([torch.ops.aten.mul])
def mul(func, input, other):
    if not isinstance(input, QTensor) or not isinstance(other, QTensor):
        return func(*dequantize(input, other))
    # Cast int8 data to int32 and do the operation
    out_data = func(input._data.to(torch.int32), other._data.to(torch.int32))
    out_scale = input._scale * other._scale
    return QTensor(out_data, out_scale)


@register_dispatch([torch.ops.aten.relu])
def relu(func, input):
    out_data = func(input._data)
    return QTensor(out_data, input._scale)


@register_dispatch([torch.ops.aten._softmax])
def _softmax(func, input, dim, half_to_float):
    # Softmax must be performed in float
    out_data = func(input.dequantize(), dim, half_to_float)
    # Since softmax is normalized, we know the optimal scale
    out_scale = torch.tensor(1 / torch.iinfo(input._data.dtype).max, dtype=input._scale.dtype)
    return QTensor.quantize(out_data, input._data.dtype, out_scale)


@register_dispatch([torch.ops.aten.transpose, torch.ops.aten.t])
def transpose(func, input):
    # Transpose is not supported if the tensor is per-axis
    assert len(input._scale.shape) == 0
    out_data = func(input._data)
    return QTensor(out_data, input._scale)


@register_dispatch([torch.ops.aten.view, torch.ops.aten._unsafe_view])
def view(func, input, *shape):
    out_data = func(input._data, *shape)
    return QTensor(out_data, input._scale)


@register_dispatch([torch.ops.aten._softmax_backward_data])
def _softmax_backward_data(func, grad, output, dim, input_dtype):
    return func(grad, output.dequantize(), dim, input_dtype)


@register_dispatch([torch.ops.aten.threshold_backward])
def threshold_backward(func, grad, output, threshold):
    return func(grad, output.dequantize(), threshold)


class QTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, data, scale, requires_grad=False):
        # This constructor can ONLY create leaf Tensors wrt autograd.
        # Use QTensor.from_tensor(t) to get a non-leaf Tensor wrt autograd.
        return torch.Tensor._make_wrapper_subclass(cls, data.size(), dtype=scale.dtype, requires_grad=requires_grad)

    def __init__(self, data, scale, requires_grad=False):
        self._data = data
        self._scale = scale

    __torch_function__ = torch._C._disabled_torch_function_impl

    def __repr__(self):  # Zero out missing values for printing
        autograd_info = (
            f", grad_fn={self.grad_fn}" if self.grad_fn else ", requires_grad=True" if self.requires_grad else ""
        )
        return f"QTensor({self._data}, scale={self._scale}, public_dtype={self.dtype}{autograd_info})"

    @classmethod
    def quantize(cls, base, int_dtype=torch.int8, scale=None):
        """Differentiable quantization function"""
        return Quantizer.apply(base, int_dtype, scale)

    def dequantize(self):
        """Differentiable dequantization function"""
        return Dequantizer.apply(self)

    def rescale(self, int_dtype=torch.int8, scale=None):
        """Differentiable requantization function"""
        return ReQuantizer.apply(self, int_dtype, scale)

    @classmethod
    def __torch_dispatch__(cls, func, types, args, kwargs=None):
        # Do not use directly func, but rather its overload
        func = func.overloadpacket
        if func in _QTENSOR_DISPATCH_TABLE:
            dispatch = _QTENSOR_DISPATCH_TABLE[func]
            return dispatch(*args, **kwargs)
        # Identify the types of the args
        types = [type(arg).__name__ for arg in args]
        raise ValueError(f"{func} {types} is no supported for QTensor.")

    @property
    def device(self):
        return self._data.device

quanto/quantization/qtensor.py
huggingface/quanto
import torch

from .nn import QLinear


__all__ = ["quantize", "freeze"]


def set_module_by_name(parent_module, name, child_module):
    module_names = name.split(".")
    if len(module_names) == 1:
        setattr(parent_module, name, child_module)
    else:
        next_module = parent_module
        for idx in range(len(module_names) - 1):
            next_module_name = module_names[idx]
            if next_module_name.isnumeric():
                next_module = next_module[int(next_module_name)]
            else:
                next_module = getattr(next_module, next_module_name)
        setattr(next_module, module_names[-1], child_module)


def quantize(model):
    # Quantization happens in-place
    for name, m in model.named_modules():
        if isinstance(m, torch.nn.Linear):
            qlinear = QLinear.from_module(m)
            set_module_by_name(model, name, qlinear)


def freeze(model):
    for name, m in model.named_modules():
        if isinstance(m, QLinear):
            m.freeze()

quanto/quantization/quantize.py
huggingface/quanto
from setuptools import find_packages, setup

setup(
    name="quanto",
    long_description=open("README.md", "r", encoding="utf-8").read(),
    long_description_content_type="text/markdown",
    url="https://github.com/huggingface/quanto",
    author="David Corvoysier",
    packages=find_packages(),
    include_package_data=True,
    zip_safe=False,
)

setup.py
huggingface/quanto
import pytest
import torch


devices = ["cpu"]
if torch.cuda.is_available():
    devices += ["cuda"]
elif torch.backends.mps.is_available():
    devices += ["mps"]


@pytest.fixture(scope="module", params=devices)
def device(request):
    return torch.device(request.param)

test/conftest.py
huggingface/quanto
import torch

from quanto.quantization import QTensor


def random_tensor(shape, dtype=torch.float32):
    # Return a random tensor between -1. and 1.
    return torch.rand(shape, dtype=dtype) * 2 - 1


def random_qtensor(shape, dtype=torch.float32):
    return QTensor.quantize(random_tensor(shape, dtype))


def q_assert_close(x: torch.Tensor, xq: QTensor):
    # Absolute error is the quantization scale
    atol = max(xq._scale, 1e-6)
    abs_error = torch.max(torch.abs(x - xq.dequantize()))
    assert abs_error <= atol, f"error {abs_error:.4f} exceeds atol {atol:.4f}"

test/helpers.py
huggingface/quanto
import os
from tempfile import TemporaryDirectory

import pytest
import torch
from helpers import random_qtensor

from quanto.quantization import QLinear, QTensor, freeze, quantize


class MLP(torch.nn.Module):
    def __init__(self, input_size, output_size, hidden_size):
        super().__init__()
        self.input_layer = torch.nn.Linear(input_size, hidden_size)
        self.mid_layer = torch.nn.Linear(hidden_size, hidden_size)
        self.output_layer = torch.nn.Linear(hidden_size, output_size)

    def forward(self, inputs):
        x = torch.nn.functional.relu(self.input_layer(inputs))
        x = torch.nn.functional.relu(self.mid_layer(x))
        return torch.nn.functional.softmax(self.output_layer(x))


def check_mlp(model, frozen):
    assert isinstance(model.input_layer, QLinear)
    assert isinstance(model.mid_layer, QLinear)
    assert isinstance(model.output_layer, QLinear)
    if frozen:
        assert isinstance(model.input_layer.weight, QTensor)
        assert isinstance(model.mid_layer.weight, QTensor)
        assert isinstance(model.output_layer.weight, QTensor)


def check_outputs(model, batch_size, input_features, device):
    qinputs = random_qtensor((batch_size, input_features), dtype=torch.float32).to(device)
    qout = model(qinputs)
    assert isinstance(qout, QTensor)


@pytest.mark.parametrize("frozen", [True, False], ids=["frozen", "non-frozen"])
def test_quantize_mlp(frozen, device):
    model = MLP(32, 10, 128).to(device)
    quantize(model)
    if frozen:
        freeze(model)
    check_mlp(model, frozen)
    check_outputs(model, 1, 32, device)


def test_serialize_quantized_mlp(device):
    input_features = 32
    hidden_features = 10
    output_features = 128
    model = MLP(input_features, hidden_features, output_features).to(device)
    quantize(model)
    freeze(model)
    with TemporaryDirectory() as tmpdir:
        mlp_file = os.path.join(tmpdir, "mlp.pt")
        torch.save(model.state_dict(), mlp_file)
        model_reloaded = MLP(input_features, hidden_features, output_features).to(device)
        quantize(model_reloaded)
        # When reloading we must assign instead of copying to force quantized tensors assignment
        model_reloaded.load_state_dict(torch.load(mlp_file), assign=True)
    check_mlp(model_reloaded, frozen=True)
    check_outputs(model_reloaded, 1, input_features, device)

test/model/test_quantize_mlp.py
huggingface/quanto
import os
from tempfile import TemporaryDirectory

import pytest
import torch
from helpers import q_assert_close, random_qtensor

from quanto.quantization import calibration, freeze
from quanto.quantization.nn import QLinear


@pytest.mark.parametrize("batch_size", [1, 10])
@pytest.mark.parametrize("tokens, embeddings", [(32, 32), (10, 32)])
@pytest.mark.parametrize("use_bias", [True, False], ids=["bias", "no-bias"])
def test_quantize_linear(batch_size, tokens, embeddings, use_bias, device):
    linear = torch.nn.Linear(embeddings, embeddings, bias=use_bias).to(device)
    qlinear = QLinear.from_module(linear)
    qinputs = random_qtensor((batch_size,) + (tokens, embeddings), dtype=torch.float32).to(device)
    # Calibrate and obtain quantized outputs
    with torch.no_grad(), calibration():
        qout = qlinear(qinputs)
    # Freeze to set quantized weights
    freeze(qlinear)
    # Align linear weights with quantized linear weights for comparison
    linear.weight = torch.nn.Parameter(qlinear.weight.dequantize())
    if use_bias:
        linear.bias = torch.nn.Parameter(qlinear.bias.dequantize())
    out = linear(qinputs.dequantize())
    q_assert_close(out, qout)
    # Now run an inference without calibrating
    with torch.no_grad():
        int_qout = qlinear(qinputs)
    assert qout._scale == int_qout._scale
    # There may be a slight difference, but of at most one quantization interval
    assert torch.max(torch.abs(qout._data - int_qout._data)) <= 1


@pytest.mark.parametrize("use_bias", [True, False], ids=["bias", "no-bias"])
def test_qlinear_serialization(use_bias):
    tokens = 10
    embeddings = 32
    linear = torch.nn.Linear(embeddings, embeddings, bias=use_bias)
    qlinear = QLinear.from_module(linear)
    qinputs = random_qtensor((1,) + (tokens, embeddings), dtype=torch.float32)
    # Calibrate and obtain quantized outputs
    with torch.no_grad(), calibration():
        qlinear(qinputs)
    # Freeze linear to store quantized weights and biases
    qlinear.freeze()
    with TemporaryDirectory() as tmpdir:
        qlinear_file = os.path.join(tmpdir, "qlinear.pt")
        torch.save(qlinear.state_dict(), qlinear_file)
        qlinear_reloaded = QLinear(embeddings, embeddings, bias=use_bias)
        # When reloading we must assign instead of copying to force quantized tensors assignment
        qlinear_reloaded.load_state_dict(torch.load(qlinear_file), assign=True)
    for attr in ["weight", "bias"]:
        t = getattr(qlinear, attr)
        if t is not None:
            t_reloaded = getattr(qlinear_reloaded, attr)
            assert torch.equal(t._data, t_reloaded._data)
            assert torch.equal(t._scale, t_reloaded._scale)
    for attr in ["in_scale", "out_scale"]:
        v = getattr(qlinear, attr)
        v_reloaded = getattr(qlinear_reloaded, attr)
        assert torch.equal(v, v_reloaded)


def test_quantize_custom_module():
    tokens = 10
    embeddings = 32

    class TwoLinearModel(torch.nn.Module):
        def __init__(self, embeddings):
            super().__init__()
            self.linear1 = torch.nn.Linear(embeddings, embeddings)
            self.linear2 = torch.nn.Linear(embeddings, embeddings)

        def forward(self, input):
            return self.linear2(self.linear1(input))

    model = TwoLinearModel(embeddings)
    model.linear1 = QLinear.from_module(model.linear1)
    model.linear2 = QLinear.from_module(model.linear2)
    qinputs = random_qtensor((1,) + (tokens, embeddings), dtype=torch.float32)
    with torch.no_grad(), calibration():
        qout = model(qinputs)
    assert model.linear1.in_scale != 1
    assert model.linear1.out_scale != 1
    assert model.linear2.in_scale != 1
    assert model.linear2.out_scale != 1
    with torch.no_grad():
        int_qout = model(qinputs)
    assert torch.equal(qout._scale, int_qout._scale)
    # There may be a slight difference, but of at most one quantization interval
    assert torch.max(torch.abs(qout._data - int_qout._data)) <= 1


@pytest.mark.parametrize("tokens, embeddings", [(32, 32), (10, 32)])
def test_qlinear_gradient(tokens, embeddings, device):
    # We use a batch size of 1 to simplify gradient manual calculations
    batch_size = 1
    linear = torch.nn.Linear(embeddings, embeddings).to(device)
    qlinear = QLinear.from_module(linear)
    assert qlinear.weight.requires_grad is True
    assert qlinear.bias.requires_grad is True
    qinputs = random_qtensor((batch_size,) + (tokens, embeddings), dtype=torch.float32).to(device)
    qout = qlinear(qinputs)
    gradient = torch.randn(qout.size()).to(device)
    qout.backward(gradient)
    # Compute gradients manually and compare
    bias_gradient = torch.sum(gradient, axis=[0, 1])
    assert torch.allclose(qlinear.bias.grad, bias_gradient)
    weight_gradient = torch.matmul(gradient.squeeze().t(), qinputs.dequantize().squeeze())
    assert torch.allclose(qlinear.weight.grad, weight_gradient)

test/nn/test_qlinear.py
huggingface/quanto
import pytest
import torch
from helpers import q_assert_close, random_qtensor, random_tensor

from quanto.quantization import QTensor


def test_to_device(device):
    qa = random_qtensor((32, 32), dtype=torch.float)
    qa = qa.to(device)
    assert qa.device.type == device.type


@pytest.mark.parametrize("input_shape", [(10,), (1, 10), (10, 32, 32)])
def test_add(input_shape, device):
    qa = random_qtensor(input_shape, dtype=torch.float32).to(device)
    # Quantized sum will have int16 data
    qsum = qa + qa
    a = qa.dequantize()
    sum = a + a
    q_assert_close(sum, qsum)


@pytest.mark.parametrize("input_shape", [(10,), (1, 10), (10, 32, 32)])
def test_mul(input_shape, device):
    qa = random_qtensor(input_shape, dtype=torch.float32).to(device)
    qb = random_qtensor(input_shape, dtype=torch.float32).to(device)
    # Quantized product will have int32 data
    qprod = qa * qb
    prod = qa.dequantize() * qb.dequantize()
    q_assert_close(prod, qprod)


@pytest.mark.parametrize("input_shape", [(10, 10), (32, 32)])
def test_matmul(input_shape, device):
    qa = random_qtensor(input_shape, dtype=torch.float32).to(device)
    qb = random_qtensor(input_shape, dtype=torch.float32).to(device)
    qmatmul = torch.matmul(qa, qb)
    # The outputs should be almost identical if we use the dequantized inputs
    matmul = torch.matmul(qa.dequantize(), qb.dequantize())
    q_assert_close(matmul, qmatmul)


@pytest.mark.parametrize("input_shape", [(1, 32, 32)])
def test_bmm(input_shape, device):
    qa = random_qtensor(input_shape, dtype=torch.float32).to(device)
    qb = random_qtensor(input_shape, dtype=torch.float32).to(device)
    qbmm = torch.bmm(qa, qb)
    # The outputs should be almost identical if we use the dequantized inputs
    bmm = torch.bmm(qa.dequantize(), qb.dequantize())
    q_assert_close(bmm, qbmm)


@pytest.mark.parametrize("input_size", [1, 10, 32])
def test_dot(input_size, device):
    qa = random_qtensor((input_size,), dtype=torch.float32).to(device)
    qb = random_qtensor((input_size,), dtype=torch.float32).to(device)
    qdot = torch.dot(qa, qb)
    # The outputs should be almost identical if we use the dequantized inputs
    dot = torch.dot(qa.dequantize(), qb.dequantize())
    q_assert_close(dot, qdot)


@pytest.mark.parametrize("batch_size", [1, 10])
@pytest.mark.parametrize("tokens, embeddings", [(5, 5), (32, 32), (10, 32)])
@pytest.mark.parametrize("use_bias", [True, False], ids=["bias", "no-bias"])
def test_linear(batch_size, tokens, embeddings, use_bias, device):
    qinputs = random_qtensor((batch_size,) + (tokens, embeddings), dtype=torch.float32).to(device)
    qweight = random_qtensor((embeddings, embeddings), dtype=torch.float32).to(device)
    if use_bias:
        bias = random_tensor((embeddings,), dtype=torch.float32).to(device)
        bias_scale = qinputs._scale * qweight._scale
        # Bias must be quantized with a higher bitwidth as they are added to the product of two int8
        qbias = QTensor.quantize(bias, torch.int32, bias_scale)
        q_assert_close(bias, qbias)
    else:
        bias = None
        qbias = None
    out = torch.nn.functional.linear(qinputs.dequantize(), qweight.dequantize(), bias)
    qout = torch.nn.functional.linear(qinputs, qweight, qbias)
    q_assert_close(out, qout)


@pytest.mark.parametrize("batch_size", [1, 10])
@pytest.mark.parametrize("tokens, embeddings", [(5, 5), (32, 32), (10, 32)])
def test_relu(batch_size, tokens, embeddings, device):
    qinputs = random_qtensor((batch_size,) + (tokens, embeddings), dtype=torch.float32).to(device)
    qout = torch.nn.functional.relu(qinputs)
    assert isinstance(qout, QTensor)
    assert torch.equal(qout._data, torch.maximum(qinputs._data, torch.zeros((1,)).to(device)))


@pytest.mark.parametrize("batch_size", [1, 10])
@pytest.mark.parametrize("tokens, embeddings", [(5, 5), (32, 32), (10, 32)])
def test_softmax(batch_size, tokens, embeddings, device):
    qinputs = random_qtensor((batch_size,) + (tokens, embeddings), dtype=torch.float32).to(device)
    qout = torch.nn.functional.softmax(qinputs, dim=-1)
    assert isinstance(qout, QTensor)
    assert torch.min(qout.dequantize()) >= 0
    assert torch.max(qout.dequantize()) <= 1


@pytest.mark.parametrize("input_shape", [(10,), (10, 32)])
def test_view(input_shape, device):
    qinputs = random_qtensor(input_shape, dtype=torch.float32).to(device)
    qview = qinputs.view((1,) + input_shape)
    assert isinstance(qview, QTensor)

test/qtensor/test_quantized_dispatch.py
huggingface/quanto
import os
from tempfile import TemporaryDirectory

import pytest
import torch
from helpers import q_assert_close, random_qtensor, random_tensor

from quanto.quantization import QTensor


@pytest.mark.parametrize("input_shape", [(10,), (1, 10), (10, 32, 32)])
@pytest.mark.parametrize("int_dtype", [torch.int8, torch.int16], ids=["int8", "int16"])
def test_quantize_dequantize(input_shape, int_dtype, device):
    a = random_tensor(input_shape, dtype=torch.float32).to(device)
    qa = QTensor.quantize(a, int_dtype)
    q_assert_close(a, qa)


@pytest.mark.parametrize("input_shape", [(10,), (1, 10), (10, 32, 32)])
@pytest.mark.parametrize("int_dtype", [torch.int8, torch.int16, torch.int32], ids=["int8", "int16", "int32"])
def test_instantiate(input_shape, int_dtype, device):
    max_value = min(1024, torch.iinfo(int_dtype).max)
    data = torch.randint(-max_value, max_value, input_shape, dtype=int_dtype)
    qa = QTensor(data, scale=torch.tensor(1.0 / max_value)).to(device)
    assert torch.max(torch.abs(qa.dequantize())) <= 1


@pytest.mark.parametrize("input_shape", [(10,), (1, 10), (10, 32, 32)])
def test_rescale_int16_int8(input_shape, device):
    a = random_tensor((10,), dtype=torch.float32).to(device)
    qa = QTensor.quantize(a, int_dtype=torch.int16)
    # Rescale to int8
    qa_rescaled = qa.rescale(torch.int8)
    q_assert_close(a, qa_rescaled)


@pytest.mark.parametrize("input_shape", [(10,), (1, 10), (10, 32, 32)])
def test_rescale_int32_int8(input_shape, device):
    # We need to generate a quantized tensor manually to avoid large int32 data
    int_max_value = 1000
    data = torch.randint(-int_max_value, int_max_value, input_shape, dtype=torch.int32)
    scale = torch.tensor(1.0 / int_max_value)
    qa = QTensor(data, scale).to(device)
    # Get the actual maximum
    a = qa.dequantize()
    float_max_value = torch.max(torch.abs(a))
    assert float_max_value <= 1
    # Rescale to int8
    qa_rescaled = qa.rescale(torch.int8, float_max_value / torch.iinfo(torch.int8).max)
    # Since we chose the optimal scale, we must have used the whole integer range
    assert torch.max(torch.abs(qa_rescaled._data)) == torch.iinfo(torch.int8).max
    q_assert_close(a, qa_rescaled)


def test_quantized_tensor_serialization():
    qinputs = random_qtensor((1, 10, 32), dtype=torch.float32)
    with TemporaryDirectory() as tmpdir:
        qinputs_file = os.path.join(tmpdir, "qinputs.pt")
        torch.save(qinputs, qinputs_file)
        qinputs_reloaded = torch.load(qinputs_file)
    assert torch.equal(qinputs._data, qinputs_reloaded._data)
    assert torch.equal(qinputs._scale, qinputs_reloaded._scale)


def test_quantized_tensor_requires_grad(device):
    weight = random_tensor((10,), dtype=torch.float32).to(device)
    weight.requires_grad = True
    qweight = QTensor.quantize(weight)
    assert qweight.requires_grad is True


def test_quantized_tensor_backward(device):
    weight = random_tensor((10,), dtype=torch.float32).to(device)
    weight.requires_grad = True
    qweight = QTensor.quantize(weight)
    gradient = torch.randn((10,)).to(device)
    # Backpropagate gradient to the inner float weights
    qweight.dequantize().backward(gradient)
    assert torch.equal(weight.grad, gradient)


def test_quantized_tensor_chained_backward(device):
    a = random_tensor((10,), dtype=torch.float32).to(device)
    a.requires_grad = True
    qa = QTensor.quantize(a)
    b = random_tensor((10,), dtype=torch.float32).to(device)
    b.requires_grad = True
    qb = QTensor.quantize(b)
    # Evaluate the product
    prod = qa * qb
    # Backpropagate
    gradient = torch.randn((10,)).to(device)
    prod.backward(gradient)
    assert torch.allclose(a.grad, qb.dequantize() * gradient)
    assert torch.allclose(b.grad, qa.dequantize() * gradient)


def test_rescale_backward(device):
    a = random_tensor((10,), dtype=torch.float32).to(device)
    a.requires_grad = True
    qa = QTensor.quantize(a, int_dtype=torch.int16)
    # Rescale to int8
    qa_rescaled = qa.rescale(torch.int8)
    assert qa_rescaled.requires_grad is True
    # Backpropagate
    gradient = torch.randn((10,)).to(device)
    qa_rescaled.backward(gradient)
    assert torch.allclose(a.grad, gradient)

test/qtensor/test_quantized_tensor.py
huggingface/quanto
import os
from tempfile import TemporaryDirectory

import pytest
import torch
from helpers import random_qtensor


def test_quantized_tensor_serialization():
    qinputs = random_qtensor((1, 10, 32), dtype=torch.float32)
    with TemporaryDirectory() as tmpdir:
        qinputs_file = os.path.join(tmpdir, "qinputs.pt")
        torch.save(qinputs, qinputs_file)
        qinputs_reloaded = torch.load(qinputs_file)
    assert torch.equal(qinputs._data, qinputs_reloaded._data)
    assert torch.equal(qinputs._scale, qinputs_reloaded._scale)


@pytest.mark.parametrize("quantize_before_load", [True, False])
def test_quantized_module_serialization(device, quantize_before_load):
    embeddings = 10

    def new_qlinear(embeddings, device):
        qweight = random_qtensor((embeddings, embeddings), dtype=torch.float32).to(device)
        qbias = random_qtensor((embeddings,), dtype=torch.float32).to(device)
        qlinear = torch.nn.Linear(embeddings, embeddings)
        with torch.no_grad():
            qlinear.weight = torch.nn.Parameter(qweight)
            qlinear.bias = torch.nn.Parameter(qbias)
        return qlinear

    qlinear = new_qlinear(embeddings, device)
    with TemporaryDirectory() as tmpdir:
        qlinear_file = os.path.join(tmpdir, "qlinear.pt")
        torch.save(qlinear.state_dict(), qlinear_file)
        torch.load(qlinear_file)
        if quantize_before_load:
            qlinear_reloaded = new_qlinear(embeddings, device)
            # Since the weights are already quantized tensors, we can copy instead of assigning
            assign = False
        else:
            qlinear_reloaded = torch.nn.Linear(embeddings, embeddings)
            # We need to force assignment instead of copy to replace weights by quantized weights
            assign = True
        qlinear_reloaded.load_state_dict(torch.load(qlinear_file), assign=assign)
    for attr in ["weight", "bias"]:
        t = getattr(qlinear, attr)
        t_reloaded = getattr(qlinear_reloaded, attr)
        assert torch.equal(t._data, t_reloaded._data)
        assert torch.equal(t._scale, t_reloaded._scale)

test/test_serialization.py
huggingface/quanto