Quantization-Aware Training: INT8/INT4 Models That Maintain Quality

# Symmetric quantization
scale = max(abs(weight)) / 127
weight_int8 = round(weight / scale).clamp(-128, 127)
weight_dequant = weight_int8 * scale
 
# Error
quantization_error = abs(weight - weight_dequant)

import torch
import matplotlib.pyplot as plt
 
# Real weights (e.g., from transformer layer)
weights = torch.randn(1000) * 0.1  # Typical scale
 
def quantize_symmetric(x, bits=8):
    """Symmetric quantization."""
    qmax = 2 ** (bits - 1) - 1
    scale = x.abs().max() / qmax
    x_quant = torch.round(x / scale).clamp(-qmax - 1, qmax)
    x_dequant = x_quant * scale
    return x_dequant, (x - x_dequant)
 
# Compare bit widths
for bits in [8, 4, 2]:
    dequant, error = quantize_symmetric(weights, bits=bits)
    print(f"INT{bits} - Mean error: {error.abs().mean():.6f}, "
          f"Max error: {error.abs().max():.6f}")
 
# INT8 - Mean error: 0.000831, Max error: 0.004123
# INT4 - Mean error: 0.006624, Max error: 0.032984
# INT2 - Mean error: 0.026497, Max error: 0.131937

# Simulate error propagation
def simulate_error_propagation(num_layers=32, bits=8):
    x = torch.randn(1, 512)  # Input
    
    for layer in range(num_layers):
        # Simulate layer computation with quantized weights
        W = torch.randn(512, 512) * 0.1
        W_quant, _ = quantize_symmetric(W, bits=bits)
        
        x = x @ W_quant
        x = torch.relu(x)
    
    return x
 
# Compare
output_fp16 = simulate_error_propagation(32, bits=16)
output_int8 = simulate_error_propagation(32, bits=8)
output_int4 = simulate_error_propagation(32, bits=4)
 
print(f"FP16 output norm: {output_fp16.norm():.4f}")
print(f"INT8 output norm: {output_int8.norm():.4f} "
      f"(diff: {(output_int8.norm() - output_fp16.norm()).abs():.4f})")
print(f"INT4 output norm: {output_int4.norm():.4f} "
      f"(diff: {(output_int4.norm() - output_fp16.norm()).abs():.4f})")

class FakeQuantize(torch.autograd.Function):
    """
    Fake quantization with straight-through estimator.
    
    Forward: Actually quantize
    Backward: Gradient flows as if no quantization (straight-through)
    """
    
    @staticmethod
    def forward(ctx, x, scale, bits=8):
        qmax = 2 ** (bits - 1) - 1
        x_quant = torch.round(x / scale).clamp(-qmax - 1, qmax)
        x_dequant = x_quant * scale
        return x_dequant
    
    @staticmethod
    def backward(ctx, grad_output):
        # Straight-through: gradient passes unchanged
        return grad_output, None, None
 
def fake_quantize(x, bits=8):
    """Apply fake quantization."""
    scale = x.abs().max() / (2 ** (bits - 1) - 1)
    return FakeQuantize.apply(x, scale, bits)
 
# Test gradient flow
x = torch.randn(10, requires_grad=True)
y = fake_quantize(x, bits=8)
loss = y.sum()
loss.backward()
 
print(f"Input: {x[:3]}")
print(f"Quantized: {y[:3]}")
print(f"Gradient: {x.grad[:3]}")  # Gradient flows!

class QuantizedLinear(torch.nn.Module):
    """Linear layer with QAT."""
    
    def __init__(self, in_features, out_features, bits=8):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.randn(out_features, in_features))
        self.bias = torch.nn.Parameter(torch.zeros(out_features))
        self.bits = bits
        
        # Learnable quantization parameters
        self.register_buffer('weight_scale', torch.tensor(1.0))
        
    def forward(self, x):
        if self.training:
            # Fake quantize weights during training
            w_quant = fake_quantize(self.weight, bits=self.bits)
        else:
            # Real quantization during inference
            w_quant = self.quantize_weights()
        
        return torch.nn.functional.linear(x, w_quant, self.bias)
    
    def quantize_weights(self):
        """Actual quantization for deployment."""
        qmax = 2 ** (self.bits - 1) - 1
        self.weight_scale = self.weight.abs().max() / qmax
        w_int = torch.round(self.weight / self.weight_scale).clamp(-qmax - 1, qmax)
        return w_int * self.weight_scale
    
    def export_quantized(self):
        """Export INT8 weights for deployment."""
        qmax = 2 ** (self.bits - 1) - 1
        scale = self.weight.abs().max() / qmax
        w_int = torch.round(self.weight / scale).clamp(-qmax - 1, qmax).to(torch.int8)
        return w_int, scale
 
# Training with QAT
model = QuantizedLinear(512, 512, bits=8)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
 
for epoch in range(10):
    # Forward pass uses fake quantization
    x = torch.randn(32, 512)
    y = model(x)
    loss = (y ** 2).mean()
    
    # Backward pass - gradients flow through fake quant
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
 
# Export quantized weights
w_int8, scale = model.export_quantized()
print(f"Quantized weights dtype: {w_int8.dtype}")
print(f"Scale: {scale:.6f}")

from transformers import LlamaConfig, LlamaForCausalLM
import torch.nn as nn
 
class QATLlamaModel:
    """
    Llama model with quantization-aware training.
    
    Converts all Linear layers to QAT versions.
    """
    
    def __init__(self, model_name_or_config, bits=8):
        if isinstance(model_name_or_config, str):
            self.model = LlamaForCausalLM.from_pretrained(model_name_or_config)
        else:
            self.model = LlamaForCausalLM(model_name_or_config)
        
        self.bits = bits
        self._convert_to_qat()
    
    def _convert_to_qat(self):
        """Replace Linear layers with QAT versions."""
        def convert_layer(module):
            for name, child in module.named_children():
                if isinstance(child, nn.Linear):
                    # Replace with QAT layer
                    qat_layer = QuantizedLinear(
                        child.in_features,
                        child.out_features,
                        bits=self.bits
                    )
                    # Copy weights
                    qat_layer.weight.data = child.weight.data.clone()
                    if child.bias is not None:
                        qat_layer.bias.data = child.bias.data.clone()
                    
                    setattr(module, name, qat_layer)
                else:
                    convert_layer(child)
        
        convert_layer(self.model)
        print(f"Converted model to QAT (INT{self.bits})")
    
    def train_qat(self, train_dataloader, num_steps=10000, lr=1e-4):
        """Train with QAT."""
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
        
        self.model.train()
        for step, batch in enumerate(train_dataloader):
            if step >= num_steps:
                break
            
            # Forward (uses fake quantization)
            outputs = self.model(**batch, labels=batch["input_ids"])
            loss = outputs.loss
            
            # Backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            if step % 100 == 0:
                print(f"Step {step}, Loss: {loss.item():.4f}")
    
    def export_quantized(self, output_path):
        """Export model with true INT8/INT4 weights."""
        self.model.eval()
        
        # Convert fake quant to real quant
        for module in self.model.modules():
            if isinstance(module, QuantizedLinear):
                w_int, scale = module.export_quantized()
                # Store for deployment
                module.weight_int = w_int
                module.weight_scale = scale
        
        torch.save(self.model.state_dict(), output_path)
        print(f"Exported quantized model to {output_path}")
 
# Usage
qat_model = QATLlamaModel("TinyLlama/TinyLlama-1.1B-Chat-v1.0", bits=8)
# qat_model.train_qat(train_dataloader)
# qat_model.export_quantized("model_int8.pt")

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
 
def train_gptq_model(
    model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    bits=4,
    group_size=128,
    calibration_samples=128
):
    """
    Train model with GPTQ quantization.
    
    Args:
        model_name: HuggingFace model ID
        bits: Target bit-width (4 or 8)
        group_size: Quantization group size (smaller = better quality, larger = faster)
        calibration_samples: Number of calibration samples
    """
    # Configure GPTQ
    quantize_config = BaseQuantizeConfig(
        bits=bits,
        group_size=group_size,
        desc_act=False,        # Disable activation reordering (faster)
        sym=True,              # Symmetric quantization
        damp_percent=0.01,     # Dampening for numerical stability
    )
    
    # Load model
    model = AutoGPTQForCausalLM.from_pretrained(
        model_name,
        quantize_config=quantize_config,
        device_map="auto"
    )
    
    # Prepare calibration data
    from datasets import load_dataset
    calibration_data = load_dataset(
        "allenai/c4",
        "en",
        split=f"train[:{calibration_samples}]",
        streaming=True
    )
    
    # Quantize
    print(f"Quantizing to INT{bits} with GPTQ...")
    model.quantize(
        calibration_data,
        use_triton=True,  # Use Triton kernels for speed
        batch_size=1
    )
    
    # Save
    model.save_quantized("./model-gptq-int4")
    print("GPTQ quantization complete!")
    
    return model
 
# Run GPTQ
gptq_model = train_gptq_model(bits=4)

class MixedPrecisionModel:
    """
    Model with mixed precision: some layers INT8, some FP16.
    
    Typically:
    - Attention: FP16 (critical for quality)
    - FFN: INT8 (less sensitive)
    - Embeddings: FP16 (vocabulary size makes INT8 ineffective)
    """
    
    def __init__(self, model):
        self.model = model
        self.sensitivity_scores = {}
    
    def analyze_sensitivity(self, calibration_dataloader):
        """
        Measure each layer's sensitivity to quantization.
        
        Returns layers ranked by sensitivity.
        """
        self.model.eval()
        
        # Get baseline outputs
        with torch.no_grad():
            baseline_outputs = []
            for batch in calibration_dataloader:
                outputs = self.model(**batch)
                baseline_outputs.append(outputs.logits)
        
        # Test each layer
        for name, module in self.model.named_modules():
            if not isinstance(module, nn.Linear):
                continue
            
            # Temporarily quantize this layer
            original_weight = module.weight.data.clone()
            module.weight.data = fake_quantize(original_weight, bits=8)
            
            # Measure output change
            with torch.no_grad():
                total_diff = 0
                for i, batch in enumerate(calibration_dataloader):
                    outputs = self.model(**batch)
                    diff = (outputs.logits - baseline_outputs[i]).abs().mean()
                    total_diff += diff.item()
            
            # Restore weights
            module.weight.data = original_weight
            
            # Store sensitivity
            self.sensitivity_scores[name] = total_diff / len(calibration_dataloader)
        
        # Rank by sensitivity
        ranked = sorted(self.sensitivity_scores.items(), key=lambda x: x[1], reverse=True)
        return ranked
    
    def apply_mixed_precision(self, keep_fp16_fraction=0.2):
        """
        Apply mixed precision based on sensitivity.
        
        Args:
            keep_fp16_fraction: Fraction of most sensitive layers to keep in FP16
        """
        # Get most sensitive layers
        ranked = sorted(self.sensitivity_scores.items(), key=lambda x: x[1], reverse=True)
        num_keep_fp16 = int(len(ranked) * keep_fp16_fraction)
        fp16_layers = set([name for name, _ in ranked[:num_keep_fp16]])
        
        print(f"Keeping {num_keep_fp16} layers in FP16:")
        for name in list(fp16_layers)[:5]:
            print(f"  {name} (sensitivity: {self.sensitivity_scores[name]:.6f})")
        
        # Convert others to INT8
        for name, module in self.model.named_modules():
            if name in fp16_layers:
                continue  # Keep FP16
            
            if isinstance(module, nn.Linear):
                # Convert to INT8
                qat_layer = QuantizedLinear(
                    module.in_features,
                    module.out_features,
                    bits=8
                )
                qat_layer.weight.data = module.weight.data.clone()
                if module.bias is not None:
                    qat_layer.bias.data = module.bias.data.clone()
                
                # Replace in model
                parent_name = '.'.join(name.split('.')[:-1])
                child_name = name.split('.')[-1]
                parent = self.model.get_submodule(parent_name)
                setattr(parent, child_name, qat_layer)
 
# Usage
mixed_precision = MixedPrecisionModel(model)
sensitivity_ranking = mixed_precision.analyze_sensitivity(calibration_loader)
mixed_precision.apply_mixed_precision(keep_fp16_fraction=0.2)

def export_onnx_int8(model, tokenizer, output_path="model_int8.onnx"):
    """Export quantized model to ONNX."""
    import torch.onnx
    
    # Prepare dummy input
    dummy_input = tokenizer("Hello world", return_tensors="pt")
    
    # Export
    torch.onnx.export(
        model,
        (dummy_input["input_ids"],),
        output_path,
        input_names=["input_ids"],
        output_names=["logits"],
        dynamic_axes={
            "input_ids": {0: "batch", 1: "sequence"},
            "logits": {0: "batch", 1: "sequence"}
        },
        opset_version=14
    )
    
    print(f"Exported to {output_path}")
 
# Quantize ONNX model further
from onnxruntime.quantization import quantize_dynamic
 
quantize_dynamic(
    "model_int8.onnx",
    "model_int8_quantized.onnx",
    weight_type=QuantType.QUInt8
)

def export_tensorrt(model, batch_size=1, seq_len=512):
    """Export to TensorRT for optimized inference."""
    import tensorrt as trt
    import torch2trt
    
    # Convert to TensorRT
    x = torch.ones((batch_size, seq_len), dtype=torch.long).cuda()
    
    model_trt = torch2trt.torch2trt(
        model,
        [x],
        fp16_mode=False,
        int8_mode=True,
        max_batch_size=batch_size
    )
    
    torch.save(model_trt.state_dict(), "model_trt_int8.pth")

config = {
    "bits": 8,
    "num_steps": 50000,
    "learning_rate": 5e-5,  # Lower than standard training
    "warmup": 2000,
    "batch_size": 4,
    "gradient_accumulation": 8,
}
# Expected: less than 1% quality loss

config = {
    "bits": 4,
    "method": "GPTQ",
    "group_size": 128,
    "calibration_samples": 256,
    "fine_tune_steps": 10000,  # Fine-tune after GPTQ
}
# Expected: 2-5% quality loss

config = {
    "attention_bits": 16,
    "ffn_bits": 8,
    "embedding_bits": 16,
    "output_bits": 8,
}
# Expected: Best quality/size trade-off