System Design 03 — Quantization Accuracy Recovery System

============================================================

Design for an automated system that detects, diagnoses, and recovers

from accuracy loss caused by quantization.

Problem Statement

Quantizing LLMs to INT8/INT4 often causes 1-8% accuracy regression. Build a system that:

1. Detects which layers are causing the regression
2. Automatically selects a recovery strategy (mixed precision, ADAROUND, QAT)
3. Applies the recovery with minimal human intervention
4. Validates the result and packages the final artifact

System Overview

Input: FP32 model + val_dataset + accuracy_target + latency_budget
                              │
                              ▼
                 [Stage 1: Baseline Assessment]
                 - Compute FP32 accuracy on val_dataset
                 - Compute naive INT8/INT4 accuracy
                 - If gap < tolerance → DONE (no recovery needed)
                 - Else → enter recovery pipeline
                              │
                              ▼
                 [Stage 2: Diagnosis]
                 - Per-layer sensitivity analysis
                 - Activation range / outlier detection
                 - Weight distribution analysis (bimodal, heavy-tailed)
                              │
                              ▼
                 [Stage 3: Strategy Selection]  ← Decision Engine
                              │
              ┌───────────────┼───────────────┐
              ▼               ▼               ▼
         [Mixed         [ADAROUND /        [QAT]
          Precision]     SmoothQuant]
              │               │               │
              └───────────────┴───────────────┘
                              │
                              ▼
                 [Stage 4: Validation]
                 - Re-evaluate accuracy after recovery
                 - If still failing: escalate to QAT or alert engineer
                              │
                              ▼
                 [Stage 5: Artifact Registry]

Stage 2: Diagnosis Engine

2a. Per-Layer Sensitivity Analysis

def sensitivity_analysis(model, val_data, baseline_acc, target_bits=8, n_examples=200):
    """
    Quantize each layer independently, measure accuracy drop.
    Returns sorted list of (layer_name, accuracy_drop).
    """
    sensitivity = {}
    model.eval()
    
    for name, module in model.named_modules():
        if not isinstance(module, QuantizableLayer):
            continue
        
        # Temporarily quantize just this layer
        with TemporaryQuantize(model, name, bits=target_bits):
            acc = fast_evaluate(model, val_data, n_examples=n_examples)
        
        sensitivity[name] = baseline_acc - acc
    
    return sorted(sensitivity.items(), key=lambda x: x[1], reverse=True)

Interpretation:

• Layer with sensitivity > 2%: "critical" → must be FP16 or FP32
• Layer with sensitivity 0.5-2%: "moderate" → INT8 preferred
• Layer with sensitivity < 0.5%: "low" → safe for INT4

2b. Outlier Detection

def detect_activation_outliers(model, calibration_data, threshold_sigma=5.0):
    """
    Identify layers with activation outliers that inflate quantization scale.
    Returns: {layer_name: {"outlier_fraction": float, "max_sigma": float}}
    """
    outlier_report = {}
    
    def make_hook(name):
        def hook(module, input, output):
            x = output.detach().float()
            mean = x.mean()
            std = x.std()
            sigma_max = (x.abs() - mean).max() / std
            outlier_frac = (x.abs() > threshold_sigma * std).float().mean().item()
            outlier_report[name] = {
                "outlier_fraction": outlier_frac,
                "max_sigma": sigma_max.item()
            }
        return hook
    
    hooks = [m.register_forward_hook(make_hook(n))
             for n, m in model.named_modules() if isinstance(m, nn.Linear)]
    
    with torch.no_grad():
        for batch in calibration_data:
            model(batch)
            break  # one batch is enough for outlier profiling
    
    for h in hooks:
        h.remove()
    
    return {k: v for k, v in outlier_report.items() if v["outlier_fraction"] > 0.001}

Stage 3: Strategy Decision Engine

class StrategySelector:
    def __init__(self, accuracy_gap, latency_budget_ms, available_gpu_hours):
        self.accuracy_gap = accuracy_gap      # % accuracy drop from PTQ
        self.latency_budget = latency_budget_ms
        self.gpu_hours = available_gpu_hours
    
    def select(self, sensitivity_report, outlier_report):
        strategies = []
        
        # Strategy 1: Mixed Precision (always try first — cheapest)
        n_critical = sum(1 for _, drop in sensitivity_report if drop > 1.0)
        mixed_prec_cost = n_critical * 0.1  # FP16 layers add ~10% latency
        if mixed_prec_cost < self.latency_budget * 0.1:
            strategies.append(("mixed_precision", cost=low, recovery_potential=0.5-1.5%))
        
        # Strategy 2: SmoothQuant (if outliers detected)
        n_outlier_layers = len(outlier_report)
        if n_outlier_layers > 3 and self.accuracy_gap > 1.5:
            strategies.append(("smoothquant", cost=medium, recovery_potential=1.0-2.0%))
        
        # Strategy 3: ADAROUND (if weight rounding is suboptimal)
        if self.accuracy_gap > 2.0 and self.gpu_hours >= 2:
            strategies.append(("adaround", cost=2 GPU hours, recovery_potential=0.5-1.0%))
        
        # Strategy 4: QAT (nuclear option, highest recovery)
        if self.accuracy_gap > 3.0 and self.gpu_hours >= 20:
            strategies.append(("qat", cost=20-50 GPU hours, recovery_potential=2.0-5.0%))
        
        # Pick lowest-cost strategy with sufficient recovery potential
        for strategy, cost, potential in sorted(strategies, key=lambda x: x[1]):
            if potential >= self.accuracy_gap * 0.8:
                return strategy
        
        return "alert_engineer"  # no automatic solution

Stage 3a: Mixed Precision Application

def apply_mixed_precision(model, sensitivity_report, target_bits=4):
    """
    Assign INT4/INT8/FP16 per layer based on sensitivity threshold.
    """
    precision_map = {}
    for name, drop in sensitivity_report:
        if drop > 1.0:
            precision_map[name] = "fp16"
        elif drop > 0.3:
            precision_map[name] = "int8"
        else:
            precision_map[name] = f"int{target_bits}"
    
    # Apply
    for name, precision in precision_map.items():
        layer = get_nested_attr(model, name)
        if precision == "fp16":
            layer.to(torch.float16)
        else:
            bits = int(precision.replace("int", ""))
            quantize_layer_inplace(layer, bits=bits)
    
    return model, precision_map

Typical result: top 8% of layers at FP16 (most sensitive), 22% at INT8, 70% at INT4.

Stage 3b: SmoothQuant

def apply_smoothquant(model, calibration_data, alpha=0.5):
    """
    Apply SmoothQuant activation-weight migration.
    For each Linear layer: W' = W * diag(s), x' = x / diag(s)
    Where s_j = max(|x_j|)^alpha / max(|W_j|)^(1-alpha)
    The s^{-1} factor is folded into the preceding LayerNorm.
    """
    # Collect activation scales
    act_scales = {}
    
    def make_hook(name):
        def hook(module, input, output):
            act_scales[name] = input[0].detach().abs().max(dim=0).values
        return hook
    
    hooks = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            hooks.append(module.register_forward_hook(make_hook(name)))
    
    with torch.no_grad():
        for batch in calibration_data:
            model(batch)
    
    for h in hooks:
        h.remove()
    
    # Apply smoothing
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and name in act_scales:
            a_max = act_scales[name]  # [in_features]
            w_max = module.weight.data.abs().max(dim=0).values  # [in_features]
            
            s = a_max.pow(alpha) / w_max.pow(1 - alpha)
            s = s.clamp(min=1e-4)
            
            # Scale weights up by s
            module.weight.data *= s.unsqueeze(0)
            
            # Find preceding LayerNorm and scale it down by s
            # (This requires knowing the model topology)
            # ... (model-specific implementation)
    
    return model

Stage 3c: ADAROUND

Key algorithm (Nagel et al., 2020):

$$\min_{v \in {0,1}^{d}} |WX - \hat{W}X|_F^2 + \lambda \sum_i h(\hat{v}_i)$$

where $\hat{W} = (\lfloor W/\delta \rfloor + v) \cdot \delta$ and $h$ is a regularizer that pushes $v$ to be binary.

def adaround_layer(W, X, bits=4, n_iters=10000, lr=1e-3, reg_strength=0.01):
    """
    Optimize rounding decisions v for weight matrix W given input activations X.
    W: [out, in], X: [n_calib, in]
    """
    scale, _ = compute_qparams_minmax(W, bits, scheme="symmetric", per_channel=True)
    
    # Initialize v based on fractional part
    W_floor = torch.floor(W / scale.unsqueeze(1))
    v_raw = W / scale.unsqueeze(1) - W_floor  # values in [0, 1)
    
    # Learnable parameter: sigmoid(v_raw) → v ∈ (0, 1)
    v = nn.Parameter(torch.log(v_raw / (1 - v_raw + 1e-8)))  # inverse sigmoid init
    optimizer = torch.optim.Adam([v], lr=lr)
    
    for step in range(n_iters):
        # Soft rounding: sigmoid(v) approaches {0, 1} over training
        v_soft = torch.sigmoid(v)
        W_q = (W_floor + v_soft) * scale.unsqueeze(1)
        
        # Reconstruction loss
        loss_rec = ((W @ X.T) - (W_q @ X.T)).pow(2).mean()
        
        # Regularization: push v to binary (0 or 1)
        # h(v) = 1 - (2v - 1)^{2β}, β increases over training
        beta = 2.0 + (step / n_iters) * 18.0  # anneal beta from 2 to 20
        loss_reg = reg_strength * (1 - (2 * v_soft - 1).pow(2 * beta).abs()).sum()
        
        loss = loss_rec + loss_reg
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    # Binarize final v
    v_binary = (torch.sigmoid(v.detach()) > 0.5).float()
    W_adaround = (W_floor + v_binary) * scale.unsqueeze(1)
    return W_adaround

Expected recovery: ~0.3-0.8% accuracy improvement over round-to-nearest on INT4.

Stage 3d: QAT (Mini Fine-tuning)

def run_qat(model, train_data, val_data, n_steps=5000, lr=1e-4):
    """
    Quantization-Aware Training: insert fake-quant nodes and fine-tune.
    Uses STE (straight-through estimator) for gradient through quantization.
    """
    # Insert fake-quant nodes
    model = torch.quantization.prepare_qat(model.train())
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=n_steps)
    
    best_acc = 0
    for step, batch in enumerate(cycle(train_data)):
        if step >= n_steps:
            break
        
        loss = model(**batch).loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        scheduler.step()
        
        if step % 500 == 0:
            acc = evaluate(model, val_data)
            if acc > best_acc:
                best_acc = acc
                torch.save(model.state_dict(), "qat_best.pt")
    
    # Load best and convert to INT8
    model.load_state_dict(torch.load("qat_best.pt"))
    model = torch.quantization.convert(model.eval())
    return model

QAT guidelines:

• Start from the PTQ model (not FP32) — preserves knowledge
• LR = 1% to 10% of original training LR
• Steps = 1% to 5% of original training steps
• Use cosine LR schedule with warmup
• Monitor both QAT loss and validation accuracy

Decision Matrix

Monitoring & SLA

• Recovery SLA: Complete recovery pipeline in < 6 GPU-hours for INT8, < 30 GPU-hours for INT4 QAT
• Escalation: If no strategy achieves target accuracy, open engineering ticket with full sensitivity report
• Weekly audit: Re-run accuracy validation on all deployed quantized models; alert if accuracy drifts > 0.5% from acceptance criteria (can happen due to evaluation dataset distribution shift)