Quantization Deep Dive: From FP32 to INT4 - The Complete Guide

Quantized = round(Original / Scale + ZeroPoint)
Dequantized = (Quantized - ZeroPoint) × Scale

def minmax_calibration(tensor):
    min_val = tensor.min()
    max_val = tensor.max()
    scale = (max_val - min_val) / 255  # For INT8
    zero_point = round(-min_val / scale)
    return scale, zero_point

def percentile_calibration(tensor, percentile=99.9):
    min_val = torch.quantile(tensor, (100 - percentile) / 100)
    max_val = torch.quantile(tensor, percentile / 100)
    scale = (max_val - min_val) / 255
    zero_point = round(-min_val / scale)
    return scale, zero_point

def entropy_calibration(tensor, num_bins=2048):
    # Build histogram
    hist, bins = torch.histogram(tensor, bins=num_bins)
    
    # Find threshold that minimizes KL divergence
    best_threshold = find_optimal_threshold(hist, bins)
    
    scale = best_threshold / 127  # Symmetric quantization
    return scale, 0

class FakeQuantize(nn.Module):
    def __init__(self, num_bits=8):
        super().__init__()
        self.num_bits = num_bits
        self.scale = nn.Parameter(torch.tensor(1.0))
        self.zero_point = nn.Parameter(torch.tensor(0.0))
        
    def forward(self, x):
        if self.training:
            # Fake quantize: quantize then dequantize
            x_q = torch.round(x / self.scale + self.zero_point)
            x_q = torch.clamp(x_q, 0, 2**self.num_bits - 1)
            x_dq = (x_q - self.zero_point) * self.scale
            
            # Straight-through estimator for gradients
            return x + (x_dq - x).detach()
        else:
            # Real quantization during inference
            return self.quantize(x)

class LearnedQuantization(nn.Module):
    def __init__(self, num_features):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(num_features))
        self.zero_point = nn.Parameter(torch.zeros(num_features))
        
    def forward(self, x):
        # Per-channel quantization
        scale = F.softplus(self.scale)  # Ensure positive
        zero_point = self.zero_point
        
        x_q = torch.round(x / scale + zero_point)
        x_q = torch.clamp(x_q, -128, 127)
        return (x_q - zero_point) * scale

minimize ||WX - W_quantized × X||²

def gptq_quantize_layer(W, X, num_bits=4):
    """
    W: Weight matrix [out_features, in_features]
    X: Calibration data [batch_size, in_features]
    """
    # Compute Hessian H = X^T X
    H = X.T @ X
    
    # Add damping for numerical stability
    H_inv = torch.inverse(H + 0.01 * torch.eye(H.shape[0]))
    
    # Initialize quantized weights
    W_q = torch.zeros_like(W)
    
    # Process weights in blocks
    block_size = 128
    for i in range(0, W.shape[1], block_size):
        block = W[:, i:i+block_size]
        
        # Compute optimal quantization
        scale = compute_optimal_scale(block)
        W_q[:, i:i+block_size] = quantize(block, scale, num_bits)
        
        # Update remaining weights to compensate
        if i + block_size < W.shape[1]:
            error = (block - W_q[:, i:i+block_size]) @ H_inv[i:i+block_size, i+block_size:]
            W[:, i+block_size:] -= error
    
    return W_q

def awq_quantize(model, calibration_data):
    # Step 1: Identify salient weights
    salience_scores = compute_salience(model, calibration_data)
    
    # Step 2: Compute per-channel scales
    for layer in model.layers:
        # Find channels with high salience
        important_channels = salience_scores[layer] > threshold
        
        # Apply protective scaling
        scale = torch.ones(layer.out_features)
        scale[important_channels] *= protection_factor
        
        # Quantize with adjusted scales
        layer.weight = quantize_with_scale(layer.weight, scale)

Y = (W × diag(s)) × (diag(s)^(-1) × X) = W' × X'

def compute_smoothing_factor(W, X, alpha=0.5):
    """
    Balance quantization difficulty between weights and activations
    """
    # Compute per-channel statistics
    w_max = W.abs().max(dim=0).values
    x_max = X.abs().max(dim=0).values
    
    # Smoothing factor
    s = (x_max / w_max) ** alpha
    
    # Apply smoothing
    W_smooth = W * s.unsqueeze(0)
    X_smooth = X / s.unsqueeze(0)
    
    return W_smooth, X_smooth, s

def sensitivity_analysis(model, calibration_data):
    sensitivities = {}
    
    for name, layer in model.named_modules():
        if isinstance(layer, nn.Linear):
            # Quantize this layer
            original_weight = layer.weight.clone()
            layer.weight.data = quantize(layer.weight, bits=4)
            
            # Measure accuracy drop
            accuracy_drop = evaluate(model) - baseline_accuracy
            sensitivities[name] = accuracy_drop
            
            # Restore original weight
            layer.weight.data = original_weight
    
    return sensitivities

def assign_bit_widths(sensitivities, bit_budget):
    # Assign more bits to sensitive layers
    sorted_layers = sorted(sensitivities.items(), key=lambda x: x[1], reverse=True)
    
    bit_assignment = {}
    for layer, sensitivity in sorted_layers:
        if sensitivity > threshold:
            bit_assignment[layer] = 8  # Keep sensitive layers at 8-bit
        else:
            bit_assignment[layer] = 4  # Aggressive quantization for others
    
    return bit_assignment

class GroupwiseQuantization:
    def __init__(self, group_size=128, bits=4):
        self.group_size = group_size
        self.bits = bits
        
    def quantize(self, tensor):
        # Reshape into groups
        orig_shape = tensor.shape
        tensor = tensor.reshape(-1, self.group_size)
        
        # Quantize each group independently
        scales = []
        quantized_groups = []
        
        for group in tensor:
            scale = group.abs().max() / (2**(self.bits-1) - 1)
            scales.append(scale)
            
            q_group = torch.round(group / scale)
            q_group = torch.clamp(q_group, -8, 7)  # INT4 range
            quantized_groups.append(q_group)
        
        return quantized_groups, scales

def pack_int4(tensor):
    """Pack two INT4 values into one INT8"""
    assert tensor.shape[-1] % 2 == 0
    
    # Reshape to separate pairs
    tensor = tensor.reshape(-1, 2)
    
    # Pack pairs into bytes
    packed = (tensor[:, 0] & 0xF) | ((tensor[:, 1] & 0xF) << 4)
    
    return packed.to(torch.int8)

def unpack_int4(packed):
    """Unpack INT8 into two INT4 values"""
    # Extract lower 4 bits
    low = (packed & 0xF).to(torch.int8)
    low = torch.where(low > 7, low - 16, low)  # Sign extend
    
    # Extract upper 4 bits
    high = ((packed >> 4) & 0xF).to(torch.int8)
    high = torch.where(high > 7, high - 16, high)  # Sign extend
    
    return torch.stack([low, high], dim=-1).reshape(-1)

def select_calibration_data(dataset, num_samples=1000):
    """Select representative calibration samples"""
    # Strategy 1: Random sampling
    random_samples = random.sample(dataset, num_samples)
    
    # Strategy 2: Diverse sampling (maximize coverage)
    diverse_samples = []
    embeddings = compute_embeddings(dataset)
    
    # K-means clustering for diversity
    clusters = KMeans(n_clusters=num_samples).fit(embeddings)
    for center in clusters.cluster_centers_:
        closest_idx = find_nearest(embeddings, center)
        diverse_samples.append(dataset[closest_idx])
    
    return diverse_samples

def optimize_bit_assignment(model, target_size_mb):
    """Find optimal per-layer bit widths given size constraint"""
    layer_sizes = get_layer_sizes(model)
    layer_sensitivities = compute_sensitivities(model)
    
    # Dynamic programming solution
    dp = {}  # (layer_idx, remaining_budget) -> (accuracy, assignment)
    
    def solve(idx, budget):
        if idx == len(layer_sizes):
            return 0, []
        
        if (idx, budget) in dp:
            return dp[(idx, budget)]
        
        best_accuracy = -float('inf')
        best_assignment = []
        
        # Try different bit widths
        for bits in [4, 6, 8]:
            size = layer_sizes[idx] * bits / 8
            if size <= budget:
                accuracy_loss = layer_sensitivities[idx][bits]
                future_acc, future_assign = solve(idx + 1, budget - size)
                
                total_acc = -accuracy_loss + future_acc
                if total_acc > best_accuracy:
                    best_accuracy = total_acc
                    best_assignment = [bits] + future_assign
        
        dp[(idx, budget)] = (best_accuracy, best_assignment)
        return best_accuracy, best_assignment
    
    return solve(0, target_size_mb)

class QuantizationPipeline:
    def __init__(self, method='gptq', bits=4):
        self.method = method
        self.bits = bits
        
    def quantize_model(self, model, calibration_loader):
        # Step 1: Prepare model
        model.eval()
        
        # Step 2: Collect statistics
        if self.method in ['gptq', 'awq']:
            statistics = self.collect_activation_statistics(
                model, calibration_loader
            )
        
        # Step 3: Apply quantization
        quantized_model = self.apply_quantization(model, statistics)
        
        # Step 4: Verify accuracy
        accuracy = self.validate(quantized_model, calibration_loader)
        
        return quantized_model, accuracy
    
    def collect_activation_statistics(self, model, loader):
        statistics = {}
        hooks = []
        
        def hook_fn(module, input, output, name):
            if name not in statistics:
                statistics[name] = []
            statistics[name].append(output.detach())
        
        # Register hooks
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                hook = module.register_forward_hook(
                    lambda m, i, o, n=name: hook_fn(m, i, o, n)
                )
                hooks.append(hook)
        
        # Run calibration
        with torch.no_grad():
            for batch in loader:
                model(batch)
        
        # Remove hooks
        for hook in hooks:
            hook.remove()
        
        return statistics

def detect_outliers(activations, threshold=6.0):
    """Detect activation outliers using z-score"""
    mean = activations.mean()
    std = activations.std()
    z_scores = torch.abs((activations - mean) / std)
    outliers = z_scores > threshold
    
    if outliers.any():
        print(f"Found {outliers.sum()} outliers")
        # Apply clipping or smoothing
        activations = torch.clamp(activations, 
                                  mean - threshold * std,
                                  mean + threshold * std)
    return activations

def correct_quantization_bias(weights, quantized_weights):
    """Correct systematic bias in quantization"""
    bias = (weights - quantized_weights).mean()
    
    # Adjust zero point to minimize bias
    corrected = quantized_weights + bias
    
    return corrected

def prevent_layer_collapse(model, min_scale=1e-5):
    """Prevent layers from quantizing to all zeros"""
    for module in model.modules():
        if hasattr(module, 'scale'):
            module.scale.data = torch.clamp(module.scale.data, min=min_scale)

import torch.quantization as quant

# Dynamic quantization (easiest)
model_int8 = quant.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)

# Static quantization (better performance)
model.qconfig = quant.get_default_qconfig('fbgemm')
quant.prepare(model, inplace=True)
# ... run calibration ...
quant.convert(model, inplace=True)

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = calibration_generator
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
tflite_model = converter.convert()

from onnxruntime.quantization import quantize_dynamic

quantize_dynamic(
    model_input='model.onnx',
    model_output='model_int8.onnx',
    weight_type=QuantType.QInt8
)

class QuantizedModelServer:
    def __init__(self, model_path, quantization_config):
        self.model = self.load_quantized_model(model_path)
        self.config = quantization_config
        
    def preprocess(self, input_data):
        # Scale inputs if needed
        if self.config.quantize_inputs:
            input_data = self.quantize_tensor(
                input_data, 
                self.config.input_scale,
                self.config.input_zero_point
            )
        return input_data
    
    def inference(self, input_data):
        with torch.no_grad():
            # Run quantized inference
            output = self.model(input_data)
            
            # Dequantize output if needed
            if self.config.quantized_output:
                output = self.dequantize_tensor(
                    output,
                    self.config.output_scale,
                    self.config.output_zero_point
                )
        return output
    
    def benchmark(self, num_runs=100):
        dummy_input = torch.randn(1, 512)
        
        # Warmup
        for _ in range(10):
            self.inference(dummy_input)
        
        # Benchmark
        start = time.time()
        for _ in range(num_runs):
            self.inference(dummy_input)
        
        avg_latency = (time.time() - start) / num_runs * 1000
        return avg_latency