INT8 Quantization › Implementation
Language AI HandbookAll chapters INT8 Quantization › ImplementationOne section · Sections in this chapter10 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement INT8 quantization from scratch to solidify these concepts. The implementation will proceed from basic absmax quantization through per-channel quantization, then show the outlier problem and show how smooth quantization resolves it.
import numpy as np
import torch
# Set random seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)Absmax Quantization
We'll start with the basic absmax quantization and dequantization operations. The function returns both the integer tensor and the scale factor, since we need both to reconstruct the original values:
def absmax_quantize(x: torch.Tensor) -> tuple[torch.Tensor, float]:
"""
Quantize a tensor to INT8 using absmax symmetric quantization.
Returns:
quantized: INT8 tensor
scale: Scale factor for dequantization
"""
# Find the maximum absolute value
absmax = x.abs().max().item()
# Compute scale factor (avoid division by zero)
scale = absmax / 127.0 if absmax != 0 else 1.0
# Quantize: divide by scale and round to nearest integer
quantized = torch.round(x / scale).clamp(-128, 127).to(torch.int8)
return quantized, scale
def dequantize(quantized: torch.Tensor, scale: float) -> torch.Tensor:
"""Recover floating-point approximation from quantized tensor."""
return quantized.float() * scaleLet's test this on a simple tensor and observe the quantization error:
# Create a sample weight tensor
weights = torch.randn(4, 4) * 0.5
# Quantize
quantized, scale = absmax_quantize(weights)
# Dequantize
reconstructed = dequantize(quantized, scale)
# Compute error
error = (weights - reconstructed).abs()Original weights: [[ 0.9635 0.7436 0.4504 -1.0528] [ 0.3392 -0.6173 -0.0215 -0.8023] [-0.3761 0.8244 -0.1962 -0.7018] [-0.3639 -0.2797 -0.3844 0.3812]] Scale factor: 0.008289 Quantized (INT8): [[ 116 90 54 -127] [ 41 -74 -3 -97] [ -45 99 -24 -85] [ -44 -34 -46 46]] Reconstructed: [[ 0.9616 0.7461 0.4476 -1.0528] [ 0.3399 -0.6134 -0.0249 -0.8041] [-0.373 0.8207 -0.1989 -0.7046] [-0.3647 -0.2818 -0.3813 0.3813]] Absolute error: [[1.881e-03 2.409e-03 2.728e-03 0.000e+00] [6.580e-04 3.853e-03 3.335e-03 1.744e-03] [3.043e-03 3.705e-03 2.708e-03 2.800e-03] [7.950e-04 2.127e-03 3.105e-03 9.200e-05]] Max error: 0.003853 Mean error: 0.002186
The maximum error is bounded by half the quantization step size (), as values are rounded to the nearest integer. For this tensor with a scale around 0.01, our maximum possible error is about 0.005. The element with the largest magnitude (which determined the absmax) reconstructs with zero error, as expected from the theory.
Analyzing Quantization Error
Let's visualize how quantization error varies across different value ranges. This reveals the characteristic "sawtooth" pattern that emerges from uniform rounding:
# Generate values across a range
original_values = torch.linspace(-1.5, 1.5, 1000)
# Quantize
scale = 1.5 / 127
quantized_vals = torch.round(original_values / scale).clamp(-128, 127)
reconstructed_vals = quantized_vals * scale
# Compute errors
errors = original_values - reconstructed_vals
The error follows a uniform distribution between and , exactly what we expect from rounding to the nearest integer on a uniform grid. This uniformity has an important implication: the expected value of the quantization error is zero, meaning quantization introduces noise but not systematic bias. This property helps explain why quantized models often perform nearly as well as full-precision models despite the apparent information loss.
Per-Channel Quantization
Now let's implement per-channel quantization for weight matrices. This extends the basic absmax approach by computing a separate scale for each output channel:
def per_channel_quantize(
weight: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Quantize weight matrix with per-output-channel scales.
Args:
weight: Shape (out_features, in_features)
Returns:
quantized: INT8 weight matrix
scales: Scale factor for each output channel
"""
# Compute absmax for each row (output channel)
absmax_per_channel = weight.abs().max(dim=1).values
# Compute scales (avoid division by zero)
scales = absmax_per_channel / 127.0
scales = torch.where(scales == 0, torch.ones_like(scales), scales)
# Quantize each row by its scale
quantized = (
torch.round(weight / scales.unsqueeze(1))
.clamp(-128, 127)
.to(torch.int8)
)
return quantized, scales
def per_channel_dequantize(
quantized: torch.Tensor, scales: torch.Tensor
) -> torch.Tensor:
"""Dequantize using per-channel scales."""
return quantized.float() * scales.unsqueeze(1)Let's compare per-tensor vs per-channel quantization on a weight matrix with varying row magnitudes. We artificially scale different rows to simulate the kind of distribution variation you find in real neural networks:
# Create weight matrix where different rows have different magnitudes
# This simulates real neural network weights
weights_varied = torch.randn(8, 16)
# Artificially scale different rows
row_scales = torch.tensor([0.1, 0.2, 0.5, 1.0, 2.0, 0.3, 0.8, 0.15])
weights_varied = weights_varied * row_scales.unsqueeze(1)
# Per-tensor quantization
qt_tensor, scale_tensor = absmax_quantize(weights_varied)
recon_tensor = dequantize(qt_tensor, scale_tensor)
error_tensor = (weights_varied - recon_tensor).abs()
# Per-channel quantization
qt_channel, scales_channel = per_channel_quantize(weights_varied)
recon_channel = per_channel_dequantize(qt_channel, scales_channel)
error_channel = (weights_varied - recon_channel).abs()Row magnitude scales: [0.1 0.2 0.5 1. 2. 0.3 0.8 0.15] Per-tensor scale: 0.0271 Per-channel scales: [0.0013 0.0033 0.0073 0.0175 0.0271 0.0059 0.0092 0.0023] Per-tensor max error: 0.013417 Per-channel max error: 0.013251 Per-tensor mean error: 0.006463 Per-channel mean error: 0.002101 Mean error per row: Row | Magnitude | Per-tensor | Per-channel --------------------------------------------- 0 | 0.10 | 0.004851 | 0.000280 1 | 0.20 | 0.009060 | 0.000748 2 | 0.50 | 0.005872 | 0.002023 3 | 1.00 | 0.007094 | 0.003080 4 | 2.00 | 0.006667 | 0.006667 5 | 0.30 | 0.005832 | 0.001644 6 | 0.80 | 0.006436 | 0.001909 7 | 0.15 | 0.005891 | 0.000461
Per-channel quantization reaches substantially lower error, especially for rows with small magnitudes. The per-tensor approach is forced to use a scale dictated by the largest row, wasting precision for smaller rows. This matches our theoretical understanding: the row with magnitude 0.1 gets 20x worse relative precision under per-tensor quantization than the row with magnitude 2.0, while per-channel quantization gives both rows equally good relative precision.
Demonstrating the Outlier Problem
Let's simulate the outlier pattern seen in large language models. We inject a single high-magnitude channel into an otherwise normal activation tensor and observe the damage to quantization quality:
# Simulate activation tensor with outlier channel
batch_size, hidden_dim = 32, 128
activations = torch.randn(batch_size, hidden_dim)
# Inject outliers in channel 50 (simulating LLM behavior)
outlier_channel = 50
activations[:, outlier_channel] *= 20 # 20x larger than typical
# Per-tensor quantization
qt_act, scale_act = absmax_quantize(activations)
recon_act = dequantize(qt_act, scale_act)
error_act = (activations - recon_act).abs()
# Separate error for outlier vs normal channels
normal_mask = torch.ones(hidden_dim, dtype=torch.bool)
normal_mask[outlier_channel] = False
error_normal = error_act[:, normal_mask].mean().item()
error_outlier = error_act[:, outlier_channel].mean().item()Activation statistics: Normal channels max: 3.83 Outlier channel max: 38.07 Quantization scale: 0.2998 Mean absolute error: Normal channels: 0.075173 Outlier channel: 0.063882 Relative error (error/value): Normal channels: 9.4% Outlier channel: 0.4%
The outlier channel forces a large scale factor, causing severe relative error for normal channels. This is exactly the problem that smooth quantization solves. Despite having 20x larger values, the outlier channel reconstructs better in relative terms than the normal channels because the scale was set based on its magnitude. Meanwhile, the normal channels suffer because their values occupy only a tiny portion of the INT8 range.
Implementing Smooth Quantization
Now let's implement the smooth quantization transformation. The core logic fits in two functions: one that computes the smoothing scales from statistics, and one that applies the transformation to the tensors:
def compute_smooth_scales(
activation_absmax: torch.Tensor, # Per-channel activation max
weight_absmax: torch.Tensor, # Per-channel weight max
alpha: float = 0.5,
) -> torch.Tensor:
"""
Compute smoothing scale factors.
Args:
activation_absmax: Max absolute activation for each input channel
weight_absmax: Max absolute weight for each input channel (row of W)
alpha: Migration strength (0=no smoothing, 1=full smoothing)
Returns:
scales: Smoothing factors for each channel
"""
# s_j = (act_max^alpha) / (weight_max^(1-alpha))
scales = (activation_absmax**alpha) / (weight_absmax ** (1 - alpha))
# Clamp to avoid extreme values
scales = scales.clamp(min=1e-5)
return scales
def apply_smooth_transform(
activations: torch.Tensor, # (batch, in_features)
weights: torch.Tensor, # (in_features, out_features)
scales: torch.Tensor, # (in_features,)
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Apply smoothing transformation to activations and weights.
Returns:
smoothed_activations: X @ diag(1/s)
smoothed_weights: diag(s) @ W
"""
# Divide activations by scales (broadcast across batch)
smoothed_activations = activations / scales.unsqueeze(0)
# Multiply weight rows by scales
smoothed_weights = weights * scales.unsqueeze(1)
return smoothed_activations, smoothed_weightsLet's apply smooth quantization to our outlier problem and measure the improvement:
# Create weight matrix for the same linear layer
weights_for_layer = torch.randn(hidden_dim, 64) * 0.1 # (in, out)
# Compute channel-wise statistics
act_absmax = activations.abs().max(dim=0).values # (hidden_dim,)
weight_absmax = weights_for_layer.abs().max(dim=1).values # (hidden_dim,)
# Compute smooth scales with alpha=0.5
smooth_scales = compute_smooth_scales(act_absmax, weight_absmax, alpha=0.5)
# Apply transformation
smoothed_act, smoothed_weights = apply_smooth_transform(
activations, weights_for_layer, smooth_scales
)
# Now quantize the smoothed tensors
qt_smooth_act, scale_smooth_act = absmax_quantize(smoothed_act)
recon_smooth_act = dequantize(qt_smooth_act, scale_smooth_act)
error_smooth_act = (smoothed_act - recon_smooth_act).abs()
# Transform error back to original space for comparison
# (multiply back by scales to compare apples-to-apples)
error_smooth_original = error_smooth_act * smooth_scales.unsqueeze(0)Before smoothing:
Activation max (normal): 3.83
Activation max (outlier): 38.07
Ratio: 9.9x
After smoothing:
Activation max (normal): 1.08
Activation max (outlier): 3.30
Ratio: 3.1x
Smooth scale for outlier channel: 11.52
Quantization error (in original space):
Without smoothing:
Normal channels: 0.075173
Outlier channel: 0.063882
With smoothing:
Normal channels: 0.019993
Outlier channel: 0.063882Smooth quantization sharply reduces the activation range ratio, making per-tensor quantization much more effective. The error for normal channels drops materially because they are no longer penalized by the outlier's presence. The smoothing scales absorb the outlier information into the weights, which can accommodate the increased magnitude without sacrificing precision.
End-to-End Quantized Linear Layer
Let's put everything together into a complete quantized linear layer. This class combines per-channel weight quantization with optional smooth quantization to handle outlier activations:
class QuantizedLinear:
"""
INT8 quantized linear layer with optional smooth quantization.
"""
def __init__(
self,
weight: torch.Tensor, # Original FP32 weights (out, in)
bias: torch.Tensor = None,
smooth_scales: torch.Tensor = None, # If provided, applies smoothing
):
self.bias = bias
self.smooth_scales = smooth_scales
# If smoothing, transform weights
if smooth_scales is not None:
weight = weight * smooth_scales.unsqueeze(0) # (out, in) * (in,)
# Quantize weights (per-channel)
self.weight_quantized, self.weight_scales = per_channel_quantize(weight)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass with INT8 computation.
In real implementations, this would use INT8 GEMM kernels.
Here we simulate the numerical behavior.
"""
# Apply activation smoothing if configured
if self.smooth_scales is not None:
x = x / self.smooth_scales.unsqueeze(0)
# Quantize activations (per-tensor)
x_quantized, x_scale = absmax_quantize(x)
# Simulate INT8 matrix multiplication
# In hardware: INT8 x INT8 -> INT32, accumulated
# Here we do it in float but with quantized values
output_int = x_quantized.float() @ self.weight_quantized.float().T
# Dequantize: multiply by both scales
output = output_int * x_scale * self.weight_scales.unsqueeze(0)
if self.bias is not None:
output = output + self.bias
return outputNow let's compare the quantized layer against full precision on inputs with outlier activations:
# Create a test linear layer
in_features, out_features = 128, 64
weight = torch.randn(out_features, in_features) * 0.1
bias = torch.randn(out_features) * 0.01
# Test input (with outlier)
test_input = torch.randn(16, in_features)
test_input[:, 50] *= 20 # Inject outlier
# Full precision output
fp_output = test_input @ weight.T + bias
# Quantized without smoothing
quant_layer = QuantizedLinear(weight, bias, smooth_scales=None)
quant_output = quant_layer.forward(test_input)
# Compute smooth scales from test input
act_max = test_input.abs().max(dim=0).values
weight_max = weight.abs().max(dim=0).values
smooth_s = compute_smooth_scales(act_max, weight_max, alpha=0.5)
# Quantized with smoothing
smooth_quant_layer = QuantizedLinear(weight, bias, smooth_scales=smooth_s)
smooth_quant_output = smooth_quant_layer.forward(test_input)
# Calculate errors for comparison
error_quant = (fp_output - quant_output).abs()
error_smooth = (fp_output - smooth_quant_output).abs()Linear layer output comparison:
Output shape: torch.Size([16, 64])
Without smoothing:
Max error: 0.224213
Mean error: 0.049900
Relative error: 3.33%
With smooth quantization:
Max error: 0.070532
Mean error: 0.017435
Relative error: 1.16%Smooth quantization reduces both maximum and mean error, making INT8 quantization viable even with outlier activations. The combination of techniques we have implemented here, per-channel weight quantization plus smooth activation quantization, is the current best practice for deploying large language models with INT8 precision.
