Optimum Quanto

Tensors

Quanto uses a Tensor subclass that handles projection and mapping:

• Projection: Maps source Tensors to the optimal range for the destination type to minimize saturated or zeroed values. Projection is symmetric per-tensor/per-channel for int8/float8, and group-wise affine for lower bitwidths.
• Mapping: Uses native PyTorch Tensor.to() for floating-point types and torch.round() for integer types.

Modules

Quanto replaces standard torch modules with quantized versions:

• Weight Quantization: Weights are typically quantized per-channel along the first dimension (output features). They are dynamically converted until the model is freeze()ed.
• Bias Handling: Biases are not quantized to preserve accuracy and avoid the complexity of extremely small scales.
• Activation Quantization: Activations are dynamically quantized per-tensor using static scales (defaulting to [-1, 1]). Calibration is recommended to find optimal scales.

Supported Modules:

• QLinear (from torch.nn.Linear): Weights are quantized; biases are not. Inputs/outputs can be quantized.
• QConv2D (from torch.nn.Conv2d): Weights are quantized; biases are not. Inputs/outputs can be quantized.
• LayerNorm (from torch.nn.LayerNorm): Weights and biases are not quantized. Outputs can be quantized.

When quantizing activations with optimum-quanto, keep the following constraints and best practices in mind:

• Per-tensor quantization: Activations are always quantized per-tensor because most linear algebra operations are incompatible with per-axis inputs.
• Weight quantization: Unlike activations, weights involved in matrix multiplications are always quantized along their first axis.
• Dequantization of outputs: The outputs of a quantized matrix multiplication are always dequantized. This is because accumulated values use a higher bitwidth (typically int32 or float32) than the activation bitwidth, and they may be combined with a float bias.
• Handling outliers: Quantizing activations per-tensor to int8 can cause significant errors if the tensors contain large outlier values, often resulting in quantized tensors where most values are zero except for the outliers.

Recommended solutions for activation outliers:

1. Use float8: Representing activations using float8 is the preferred option to mitigate outlier issues.
2. Static Smoothing: You can 'smooth' activations statically (similar to the SmoothQuant method). A script for smoothing certain model architectures is available in the external/smoothquant directory of this repository.

Quanto uses torch.library to provide device-specific implementations for its operations. Implementations are available for CPU (via generic C++), CUDA, MPS (Metal), and XPU (SYCL).

Overriding existing operations

To provide a specialized implementation for an operation that already has a default (e.g., unpack), use the @torch.library.impl decorator specifying the operation name and the target device(s).

Declaring new operations

To add a new device-specific operation to the quanto:: namespace, you must first declare its schema using torch.library.define before providing the implementation with @torch.library.impl.

# 1. Overriding an existing operation
@torch.library.impl("quanto::unpack", ["CPU", "CUDA"])
def unpack(packed: torch.Tensor, bits: int) -> torch.Tensor:
    return ext.unpack(t, bits)

# 2. Declaring and implementing a new operation
torch.library.define(
    "quanto::gemm_f16i4",
    "(Tensor input, Tensor other, Tensor other_scale, Tensor other_shift, int group_size) -> Tensor",
)

@torch.library.impl("quanto::gemm_f16i4", ["CUDA"])
def gemm_f16i4(
    input: torch.Tensor,
    other: torch.Tensor,
    scales: torch.Tensor,
    shift: torch.Tensor,
    group_size: int,
) -> torch.Tensor:
    ...

For non-Hugging Face models, follow this manual workflow. By default, weights are dynamically quantized until freeze() is called.

1. Quantize: Convert float model to a dynamic quantized model using quantize(model, weights=..., activations=...).
2. Calibrate (Optional): If quantizing activations, use the Calibration context manager to record activation ranges using representative samples.
3. Tune (Optional): Perform Quantization-Aware-Training (QAT) to recover accuracy.
4. Freeze: Replace float weights with integer weights using freeze(model).
5. Serialize: Save the state_dict (using safetensors recommended) and the quantization_map (using quantization_map(model)).
6. Reload: Use requantize(new_model, state_dict, quantization_map) to load weights into an empty model instance.

# 1. Quantize
from optimum.quanto import quantize, qint8
quantize(model, weights=qint8, activations=qint8)

# 2. Calibrate
from optimum.quanto import Calibration
with Calibration(momentum=0.9):
    model(samples)

# 3. Tune (QAT)
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    output = model(data).dequantize()
    loss = torch.nn.functional.nll_loss(output, target)
    loss.backward()
    optimizer.step()

# 4. Freeze
from optimum.quanto import freeze
freeze(model)

# 5. Serialize
from safetensors.torch import save_file
import json
from optimum.quanto import quantization_map

save_file(model.state_dict(), 'model.safetensors')
with open('quantization_map.json', 'w') as f:
    json.dump(quantization_map(model), f)

# 6. Reload
import torch
from safetensors.torch import load_file
from optimum.quanto import requantize

state_dict = load_file('model.safetensors')
with open('quantization_map.json', 'r') as f:
    q_map = json.load(f)

with torch.device('meta'):
    new_model = ... # Instantiate your model architecture
requantize(new_model, state_dict, q_map, device=torch.device('cuda'))