foxhunt/QAT_CODE_PATTERNS_GUIDE.md

QAT Code Patterns & Architecture Deep Dive

Date: 2025-10-23
Purpose: Detailed pattern analysis for QAT implementation across 3 modules Audience: ML Engineers, Systems Architects

---

Code Organization Diagram

┌────────────────────────────────────────────────────────────────────┐
│                    QAT Training Pipeline                            │
└────────────────────────────────────────────────────────────────────┘

Phase 1: Model Creation
┌─────────────────────────────────────────────────────────────────────┐
│ FP32 Training (Standard)                                             │
│  • Use any standard training loop                                    │
│  • Supported models: TFT, MAMBA-2, DQN, PPO                         │
│  • Target: Full training for 50+ epochs                             │
│  • Duration: 4-5 minutes (TFT on RTX 3050 Ti)                      │
└─────────────────────────────────────────────────────────────────────┘
                              ↓
                    FP32 Trained Weights
                              ↓
Phase 2: QAT Wrapper Creation
┌─────────────────────────────────────────────────────────────────────┐
│ QATTemporalFusionTransformer::new_from_fp32(fp32_model)             │
│  • Zero-copy wrapper (no weight duplication)                        │
│  • Creates HashMap<String, FakeQuantize> observers                  │
│  • Initializes 11 observer layers                                   │
│  • Duration: <100ms                                                 │
│  • Memory overhead: ~10KB (observer metadata)                       │
└─────────────────────────────────────────────────────────────────────┘
                              ↓
                      QAT Model Ready
                              ↓
Phase 3: Calibration (100-500 batches)
┌─────────────────────────────────────────────────────────────────────┐
│ qat_model.calibrate(calibration_data)                               │
│  • Run forward passes without gradient updates                      │
│  • Observers collect running min/max via EMA                        │
│  • Freeze scale/zero_point parameters                               │
│  • Duration: 30-200 seconds (depends on batch count)               │
└─────────────────────────────────────────────────────────────────────┘
                              ↓
                   Calibrated Observers
                              ↓
Phase 4: QAT Fine-tuning (5-10 epochs)
┌─────────────────────────────────────────────────────────────────────┐
│ Training Loop with Fake Quantization                                │
│  • Standard SGD/Adam optimizer                                      │
│  • Forward: Apply FakeQuantize to linear outputs                    │
│  • Backward: Standard FP32 gradients (STE)                          │
│  • Result: Model learns quantization-robust weights                 │
│  • Duration: 20-40 seconds (5 epochs, small dataset)               │
│  • Memory: Same as FP32 (~1GB)                                      │
└─────────────────────────────────────────────────────────────────────┘
                              ↓
                    QAT Fine-tuned Model
                              ↓
Phase 5: INT8 Conversion
┌─────────────────────────────────────────────────────────────────────┐
│ QuantizedTemporalFusionTransformer::new_from_fp32(qat_model)        │
│  • Extract FP32 weights from trained model                          │
│  • Quantize to INT8 using calibrated scales                         │
│  • Create INT8 model for inference                                  │
│  • Duration: 10-15 seconds                                          │
│  • Memory: ~200MB (75% reduction)                                   │
└─────────────────────────────────────────────────────────────────────┘
                              ↓
                    INT8 Model (Production)
                              ↓
Phase 6: Validation & Deployment
┌─────────────────────────────────────────────────────────────────────┐
│ Quality Assurance                                                    │
│  • Compare INT8 vs FP32 accuracy (target: <1% diff)                │
│  • Benchmark inference latency (target: ~3.2ms)                    │
│  • Verify memory footprint (target: ~125MB)                        │
│  • Deploy to production                                             │
└─────────────────────────────────────────────────────────────────────┘

---

Tensor Operation Flow

FakeQuantize Forward Pass

Input Tensor (FP32)
  [batch=32, seq_len=60, hidden=256]
              ↓
         to_dtype(F32)         ← Already FP32, pass-through
              ↓
    Get device from input      ← KEY FIX: Use input.device()
              ↓
Create scale/zero_point tensors on SAME device
  scale_tensor: [0.01]  @ device
  zero_point:   [127]   @ device
              ↓
      Quantization Phase
    ┌──────────────────────────┐
    │ scaled = input / scale    │  broadcast_div
    │ shifted = scaled + zp     │  broadcast_add
    │ rounded = round(shifted)  │
    │ clamped = clamp(0, 255)   │
    └──────────────────────────┘
              ↓
   Dequantization Phase
    ┌──────────────────────────┐
    │ deshifted = clamped - zp  │  broadcast_sub
    │ dequantized = deshifted * │  broadcast_mul
    │            scale          │
    └──────────────────────────┘
              ↓
        to_dtype(original)
              ↓
  Output Tensor (FP32)
  [batch=32, seq_len=60, hidden=256]
  
  ✓ Same device as input
  ✓ Same shape as input
  ✓ Quantization noise simulated

---

Pattern Analysis: Device Consistency

Pattern: Correct Device Handling

Location: ml/src/memory_optimization/qat.rs (lines 329-331)

pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
    if !self.training {
        return Ok(input.clone());
    }

    // Convert to F32 for quantization
    let f32_input = input.to_dtype(DType::F32)?;

    // ✅ CORRECT: Get device from input tensor
    let input_device = f32_input.device();
    let scale_tensor = Tensor::new(&[self.scale], input_device)?;
    let zero_point_tensor = Tensor::new(&[self.zero_point as f32], input_device)?;

    // Broadcast operations now work correctly on GPU
    let scaled = f32_input.broadcast_div(&scale_tensor)?;
    let shifted = scaled.broadcast_add(&zero_point_tensor)?;
    
    // ... rest of quantization ...
}

Why This Works:

• input.device() returns the device the input tensor is on (CPU or CUDA)
• Tensor::new() creates tensors on the specified device
• broadcast_* operations automatically handle shape broadcasting
• GPU CUDA kernels only work when both operands are on GPU

Comparison: Incorrect Pattern

// ❌ WRONG: Hardcoded device in self
pub struct FakeQuantize {
    device: Device,  // Could be CPU, but input is GPU!
}

pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
    // ❌ Creates tensor on self.device (CPU)
    let scale_tensor = Tensor::new(&[self.scale], &self.device)?;
    
    // ❌ ERROR: GPU tensor / CPU tensor mismatch
    let scaled = input.broadcast_div(&scale_tensor)?;
    // Device mismatch error at runtime!
}

Error Message (What happens):

thread 'test_forward' panicked at 'attempt to divide tensors on different devices'
Location: candle_core/src/ops.rs:234

---

Pattern Analysis: Calibration Statistics

Pattern: EMA (Exponential Moving Average) Statistics

Location: ml/src/memory_optimization/qat.rs (lines 142-147)

Algorithm:

EMA Update Formula:
  new_value = momentum * old_value + (1 - momentum) * batch_value

With momentum = 0.99:
  new_value = 0.99 * old_value + 0.01 * batch_value

Effect:
  - New batch contributes only 1% to update
  - Running average has 99% "memory" of previous values
  - Converges slowly but stably to true min/max

Code:

match (self.running_min, self.running_max) {
    (Some(running_min), Some(running_max)) => {
        // EMA update: running_val = momentum * running_val + (1 - momentum) * new_val
        self.running_min = Some(
            self.ema_momentum * running_min + (1.0 - self.ema_momentum) * min_val,
        );
        self.running_max = Some(
            self.ema_momentum * running_max + (1.0 - self.ema_momentum) * max_val,
        );
    }
    _ => {
        // First sample: initialize running statistics
        self.running_min = Some(min_val);
        self.running_max = Some(max_val);
    }
}

Behavior Over Time:

Batch 1: min=-2.0, max=3.0
  running_min = -2.0
  running_max = 3.0

Batch 2: min=-1.5, max=2.5
  running_min = 0.99 * (-2.0) + 0.01 * (-1.5) = -1.995
  running_max = 0.99 * (3.0) + 0.01 * (2.5) = 2.995

Batch 3: min=-1.8, max=2.8
  running_min = 0.99 * (-1.995) + 0.01 * (-1.8) = -1.99305
  running_max = 0.99 * (2.995) + 0.01 * (2.8) = 2.99405

After 100 batches: Converged to stable estimate

Momentum Impact:

Momentum	Convergence Speed	Stability	Use Case
0.99	Slow (100-200 batches)	Very stable	Production
0.95	Medium (50-100 batches)	Stable	Good balance
0.90	Fast (30-50 batches)	Less stable	Quick testing
0.50	Very fast (10-20 batches)	Unstable	Not recommended

---

Pattern Analysis: Quantization Parameters

Pattern: Symmetric vs Asymmetric Quantization

Location: ml/src/tft/qat_tft.rs (lines 160-164)

Symmetric Quantization (TFT Default):

fn compute_quantization_params(&self, min_val: f32, max_val: f32) -> (f32, i8) {
    // Find largest absolute value
    let abs_max = min_val.abs().max(max_val.abs());
    
    // Map [-abs_max, +abs_max] to [-127, +127]
    let scale = abs_max / 127.0;           // Range per unit
    let zero_point = 127i8;                // Centered at 127
    
    (scale, zero_point)
}

Visualization:

FP32 Value Range: [-3.0, +3.0]
  abs_max = max(3.0, 3.0) = 3.0
  scale = 3.0 / 127 ≈ 0.02362

Quantization Mapping:
  FP32:   -3.0 ──────────── 0.0 ──────────── +3.0
  INT8:   0   ────────────── 127 ───────────  255
          
  Mapping:
    x_int8 = round(x_fp32 / scale) + 127
    x_fp32 = (x_int8 - 127) * scale

Why Symmetric for TFT:

1. TFT features are typically normalized (mean≈0, std≈1)
2. Distribution is roughly symmetric around zero
3. Simpler implementation (no learned zero_point)
4. Better for attention mechanisms (softmax produces symmetric outputs)

Asymmetric (Not Used):

// Maps [min, max] → [0, 255] with learned zero_point
let scale = (max_val - min_val) / 255.0;
let zero_point = (-min_val / scale).round() as i8;

When to Use Asymmetric:

• ReLU outputs (range [0, inf])
• Distributions heavily skewed one direction
• Clipped activations

---

Pattern Analysis: Broadcast Operations

Why Broadcasting is Needed

Problem: Shape Mismatch in Quantization

Input:         [batch=32, seq_len=60, hidden=256]  (3D)
Scale:         [0.01234]                           (scalar)
zero_point:    [127]                               (scalar)

Standard operation would fail:
  scaled = input / scale  ← 3D / scalar ❌

Solution: Broadcast scalars to match input shape
  scale_tensor: [0.01234] → expand to [32, 60, 256]
  scaled = input.broadcast_div(&scale_tensor)  ✅

Code Pattern (lines 333-346):

// Step 1: Create tensors (scalar shape)
let scale_tensor = Tensor::new(&[self.scale], device)?;           // [1]
let zero_point_tensor = Tensor::new(&[zero_point as f32], device)?; // [1]

// Step 2: Broadcast operations (expand to match input)
let scaled = f32_input.broadcast_div(&scale_tensor)?;      // [32,60,256] / [1]
let shifted = scaled.broadcast_add(&zero_point_tensor)?;   // [32,60,256] + [1]
let rounded = shifted.round()?;
let clamped = rounded.clamp(0.0, 255.0)?;

// Step 3: Continue broadcasting
let deshifted = clamped.broadcast_sub(&zero_point_tensor)?;  // [32,60,256] - [1]
let dequantized = deshifted.broadcast_mul(&scale_tensor)?;   // [32,60,256] * [1]

Broadcasting Rules (Candle):

• Scalars broadcast to any shape
• Dimension 1 broadcasts to match any size
• Example: [1,256] broadcasts with [32,60,256] → all operations valid

---

Pattern Analysis: Forward Pass Modes

Training vs Evaluation Mode

Location: ml/src/tft/qat_tft.rs (lines 179-208)

Two Distinct Modes:

pub fn forward(&mut self, x: &Tensor) -> Result<Tensor, MLError> {
    // MODE 1: CALIBRATION (collecting statistics)
    if self.calibration_mode {
        let x_vec = x.flatten_all()?.to_vec1::<f32>()?;
        let min_val = x_vec.iter().cloned().fold(f32::INFINITY, f32::min);
        let max_val = x_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        
        // Update running statistics (EMA)
        self.update_statistics(min_val, max_val);
        
        // Apply fake quantization with CURRENT statistics
        let (scale, zero_point) = self.compute_quantization_params(min_val, max_val);
        self.apply_fake_quantization(x, scale, zero_point)
    } 
    // MODE 2: TRAINING/EVALUATION (frozen parameters)
    else {
        match (self.scale, self.zero_point) {
            (Some(scale), Some(zero_point)) => {
                // Use frozen scale/zero_point from calibration
                self.apply_fake_quantization(x, scale, zero_point)
            }
            _ => {
                // No calibration: pass-through
                Ok(x.clone())
            }
        }
    }
}

State Transitions:

[CALIBRATION MODE]
     ↓ (after N batches)
disable_calibration()
     ↓
[EVALUATION MODE]
     ↓ (scale/zero_point frozen)
forward() uses frozen parameters
     ↓
[TRAINING MODE] (with gradient updates)
     ↓
backward() updates weights

---

Pattern Analysis: VarMap Quantization

Parallel Weight Quantization

Location: ml/src/tft/quantized_tft.rs (lines 131-162)

Purpose: Speed up INT8 weight conversion (10-15s vs 30-60s sequential)

Code:

pub fn new_from_fp32(fp32_model: &TemporalFusionTransformer) -> Result<Self, MLError> {
    // Extract FP32 weights
    let fp32_varmap = fp32_model.varmap();
    
    // Parallel quantization
    // - Each thread handles N weight tensors
    // - No dependencies between tensors
    // - 3-4x speedup on 8-core CPU
    let quantized_weights = quantize_varmap_parallel(fp32_varmap, &device)?;
    
    Ok(Self {
        quantized_weights,
        // ... rest of init ...
    })
}

Performance:

Sequential:  weight1 → weight2 → weight3 → ... → weight_N  (~30-60s)
Parallel:    weight1 ─┐
             weight2  ├─→ (in parallel) (~10-15s)
             weight3  ─┘

Speedup: 30s / 10s = 3x faster

---

Observer Graph Architecture

11-Layer Observer Topology

Location: ml/src/tft/qat_tft.rs (lines 328-348)

Variable Selection Networks (3x)
├── static_vsn.attention_weights
├── historical_vsn.attention_weights
└── future_vsn.attention_weights

LSTM Layers (2x)
├── lstm_encoder
└── lstm_decoder

Temporal Attention (4x)
├── temporal_attention.q_proj        (Query projection)
├── temporal_attention.k_proj        (Key projection)
├── temporal_attention.v_proj        (Value projection)
└── temporal_attention.o_proj        (Output projection)

Quantile Output (1x)
└── quantile_outputs.output_layer

Coverage:

• 10 observer layers + 1 output layer = 11 total
• Covers all major linear operations in TFT
• Each observer independently collects statistics

HashMap Structure:

fake_quant_observers: HashMap<String, FakeQuantize>
  Key: "static_vsn.attention_weights"
  Value: FakeQuantize { scale, zero_point, running_min, running_max, ... }

---

Integration Test Patterns

Pattern: End-to-End Workflow Test

Location: ml/tests/qat_tft_integration_test.rs (lines 396-450+)

#[test]
fn test_qat_end_to_end_workflow() {
    // Step 1: Create FP32 model
    let config = TFTConfig::default();
    let fp32_model = TemporalFusionTransformer::new_with_device(config, device)?;
    
    // Step 2: Wrap with QAT
    let mut qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model)?;
    
    // Step 3: Calibrate (100 samples)
    let calibration_data = generate_random_data(100);
    qat_model.calibrate(&calibration_data)?;
    
    // Step 4: Run forward with fake quantization
    let static_feat = Tensor::randn(...)?;
    let hist_feat = Tensor::randn(...)?;
    let fut_feat = Tensor::randn(...)?;
    let output = qat_model.forward(&static_feat, &hist_feat, &fut_feat)?;
    
    // Step 5: Verify output shape and device
    assert_eq!(output.dims()[0], batch_size);
    assert_eq!(output.dims()[1], horizon);
    assert_eq!(output.dims()[2], num_quantiles);
}

Test Coverage:

• Model creation ✓
• QAT wrapper initialization ✓
• Calibration with multiple batches ✓
• Forward pass with fake quantization ✓
• Output validation ✓

---

Device Mismatch Prevention

Test Pattern: Device Consistency

Location: ml/tests/qat_device_consistency_test.rs (lines 8-42)

#[test]
fn test_fake_quantize_device_consistency() {
    // Test on available device (GPU or CPU)
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    // Create observer and calibrate
    let observer = QuantizationObserver::new(config, device.clone());
    let batch = Tensor::randn(0.0, 1.0, (32, 64), &device)?;
    observer.observe(&batch)?;
    
    // Create FakeQuantize from observer
    let fake_quant = FakeQuantize::from_observer(&observer)?;
    
    // Create input on SAME device as observer
    let input = Tensor::randn(0.0, 1.0, (32, 64), &device)?;
    
    // Forward should NOT crash
    let output = fake_quant.forward(&input)?;
    
    // Verify output is on same device
    assert_eq!(
        format!("{:?}", input.device()),
        format!("{:?}", output.device())
    );
}

Why This Test Matters:

• Validates device awareness across GPU/CPU
• Detects hardcoded device assumptions
• Ensures broadcast operations work cross-device

---

Summary: Critical Code Patterns

Pattern	Location	Purpose	Key Insight
Device-Aware Tensors	qat.rs:329-331	Prevent CUDA/CPU mismatch	Use input.device() not self.device
EMA Statistics	qat.rs:142-147	Stable calibration	momentum=0.99 converges in 100 batches
Quantization Params	qat_tft.rs:160-164	Compute scale/zero_point	Symmetric maps [-abs_max, abs_max]
Broadcasting	qat.rs:333-346	Shape compatibility	Scalars broadcast to any shape
Forward Modes	qat_tft.rs:179-208	Calibration vs training	Two distinct paths based on mode
Parallel Quantization	quantized_tft.rs:131-162	Speed up INT8 conversion	3-4x faster than sequential
Observer Graph	qat_tft.rs:328-348	Track all linear layers	11 independent observers
Device Consistency	Test pattern	Validate cross-device ops	Test both GPU and CPU paths

---

Recommended Reading Order

1. QAT_COMPREHENSIVE_ANALYSIS.md - Architecture overview
2. This document - Code patterns and detailed analysis
3. ml/docs/QAT_GUIDE.md - Production usage guide
4. Source code:
   • Start: ml/src/memory_optimization/qat.rs (core algorithms)
   • Then: ml/src/tft/qat_tft.rs (TFT integration)
   • Finally: ml/src/tft/quantized_tft.rs (INT8 runtime)