Mini PyTorch — E2E Architecture Reference
Projects: Autograd Engine · Matrix Multiplication Kernel · Softmax Kernel · Transformer from Scratch · Flash Attention Kernel · Vision Transformer · Diffusion Model (UNet + Scheduler)
What You're Building
A ground-up implementation of the compute stack that trains neural networks: automatic differentiation, hand-written GPU kernels, and the core architectures that define modern deep learning. After completing this system you will understand exactly what loss.backward() does, why Flash Attention exists, and how images can be generated from pure noise.
Project 1 — Autograd Engine (like Micrograd)
Conceptual Architecture
The autograd engine is a lazy computation graph where every operation records how to undo itself. Forward pass builds the graph; backward pass walks it in reverse to propagate gradients.
FORWARD PASS
────────────
User code: Graph built behind the scenes:
x = Value(2.0) x (data=2.0, grad=0)
w = Value(-3.0) w (data=-3.0, grad=0)
b = Value(5.0) b (data=5.0, grad=0)
n = x*w + b │
loss = n.tanh() n=x*w+b → loss=tanh(n)
BACKWARD PASS
─────────────
loss.backward()
Step 1: Topological sort (DFS, reverse postorder)
[x, w, b, x*w, n, loss] ← evaluation order
reversed = [loss, n, x*w, b, w, x]
Step 2: Set loss.grad = 1.0
Step 3: For each node in reverse topo order:
call node._backward() ← closure set during forward
Example _backward for multiplication (n = x * w):
x.grad += w.data * n.grad
w.grad += x.data * n.gradData Structure
class Value:
def __init__(self, data):
self.data = float(data) # actual value
self.grad = 0.0 # ∂loss/∂self, accumulates
self._backward = lambda: None # filled in by ops
self._prev = set() # parent nodes (for graph)
def __mul__(self, other):
out = Value(self.data * other.data, (self, other))
def _backward():
self.grad += other.data * out.grad # chain rule
other.grad += self.data * out.grad
out._backward = _backward
return out
def backward(self):
topo = []
visited = set()
def build(v):
if v not in visited:
visited.add(v)
for child in v._prev: build(child)
topo.append(v)
build(self)
self.grad = 1.0
for node in reversed(topo):
node._backward()Operations to Implement
| Op | Forward | Backward |
|---|---|---|
| add | a + b |
da += dout, db += dout |
| mul | a * b |
da += b*dout, db += a*dout |
| pow | a^k |
da += k*a^(k-1)*dout |
| relu | max(0, a) |
da += (a>0)*dout |
| tanh | tanh(a) |
da += (1-tanh²(a))*dout |
| exp | e^a |
da += e^a * dout |
| log | ln(a) |
da += (1/a)*dout |
Training Loop
for epoch in range(epochs):
# Forward
ypred = [model(x) for x in X_train]
loss = mean_squared_error(ypred, Y_train)
# Zero gradients (must reset each step — grads accumulate)
for p in model.parameters():
p.grad = 0.0
# Backward
loss.backward()
# Update (SGD)
for p in model.parameters():
p.data -= lr * p.gradExtension to Tensors
Replace float with np.ndarray. The _backward closures become vectorized operations. Most ops remain identical; broadcasting rules must be handled by summing gradients along broadcasted dims.
Common Gotchas
- Gradient accumulation:
grad +=notgrad =— same node can receive grad from multiple paths - In-place ops: Break the graph; avoid or clone before modifying
- Non-leaf grads: Usually don't need to be stored (set
requires_grad=Falsefor efficiency) - Numerical check: Finite difference verification —
(f(x+h) - f(x-h)) / 2h ≈ grad
Project 2 — Matrix Multiplication Kernel (CUDA)
Problem Statement
Compute C = A @ B where A is (M,K), B is (K,N), C is (M,N). This is the dominant operation in transformer inference and training.
Level 1 — Naive Kernel
Grid: dim3(ceil(N/32), ceil(M/32)) ← 2D grid of thread blocks
Block: dim3(32, 32) ← 1024 threads per block
Each thread (row, col) computes one output element:
float acc = 0;
for (int k = 0; k < K; k++)
acc += A[row*K + k] * B[k*N + col];
C[row*N + col] = acc;
Problem: Each output element reads K values from GLOBAL memory (HBM)
HBM bandwidth ≈ 2 TB/s, latency ≈ 400 cycles
Compute intensity too low → memory boundLevel 2 — Shared Memory Tiling
TILE = 32
Each block is responsible for a (TILE × TILE) output tile
__shared__ float As[TILE][TILE];
__shared__ float Bs[TILE][TILE];
float acc = 0.0;
for (int t = 0; t < K/TILE; t++) {
// Cooperatively load tile of A and B into SRAM
As[ty][tx] = A[row * K + (t*TILE + tx)];
Bs[ty][tx] = B[(t*TILE + ty) * N + col];
__syncthreads();
// Compute on fast shared memory
for (int k = 0; k < TILE; k++)
acc += As[ty][k] * Bs[k][tx];
__syncthreads();
}
C[row*N + col] = acc;
HBM reads reduced: O(M*N*K) → O(M*N*K / TILE)
SRAM bandwidth: ~19 TB/s (10x faster than HBM)Level 3 — Register Blocking + Vectorized Loads
Each thread owns a (Rv × Rc) = (8 × 8) output tile in REGISTERS
Threads per block: TILE/Rv × TILE/Rc
Load with float4 (128-bit):
float4 a4 = *reinterpret_cast<float4*>(&A[...]);
→ 4 floats in 1 memory transaction (4x bandwidth utilization)
Double buffering:
While computing tile[t], prefetch tile[t+1] into second shared mem buffer
Hides memory latency with compute
Achieved performance: ~60-70% of peak FLOPSLevel 4 — Tensor Cores (Warp Matrix Multiply)
Uses CUDA's wmma (Warp Matrix Multiply Accumulate) API
Operates on 16×16×16 matrix fragments in FP16
wmma::fragment<wmma::matrix_a, 16,16,16, half> a_frag;
wmma::fragment<wmma::matrix_b, 16,16,16, half> b_frag;
wmma::fragment<wmma::accumulator, 16,16,16, float> c_frag;
wmma::fill_fragment(c_frag, 0.0f);
wmma::load_matrix_sync(a_frag, A_ptr, K);
wmma::load_matrix_sync(b_frag, B_ptr, N);
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); ← single instruction
wmma::store_matrix_sync(C_ptr, c_frag, N, wmma::mem_row_major);
A100 Tensor Core peak: 312 TFLOPS (FP16), 77 TFLOPS (FP32)
cuBLAS achieves ~80% of peak via highly tuned tile sizes + schedulingFull Pipeline
Host Setup:
cudaMalloc(&d_A, M*K*sizeof(float))
cudaMalloc(&d_B, K*N*sizeof(float))
cudaMalloc(&d_C, M*N*sizeof(float))
cudaMemcpy(d_A, h_A, M*K*sizeof(float), cudaMemcpyHostToDevice)
cudaMemcpy(d_B, h_B, K*N*sizeof(float), cudaMemcpyHostToDevice)
Kernel Launch:
dim3 gridDim(ceil(N/TILE), ceil(M/TILE))
dim3 blockDim(TILE, TILE)
matmul_kernel<<<gridDim, blockDim>>>(d_A, d_B, d_C, M, N, K)
cudaDeviceSynchronize()
Copy Back:
cudaMemcpy(h_C, d_C, M*N*sizeof(float), cudaMemcpyDeviceToHost)
Profiling:
nsys profile ./matmul ← Nsight Systems (timeline)
ncu --metrics ... ./matmul ← Nsight Compute (kernel deep dive)
Check: memory throughput, achieved occupancy, warp efficiencyProject 3 — Softmax Kernel Optimization
The Three-Stage Evolution
STAGE 1: NAIVE (Numerically Unstable)
out[i] = exp(x[i]) / sum_j(exp(x[j]))
Problem: exp(x[i]) → ∞ for x[i] > 88 (FP32 overflow)
STAGE 2: NUMERICALLY STABLE (3-pass)
m = max(x) [pass 1: find max]
s = sum(exp(x[i] - m)) [pass 2: shifted sum]
out[i] = exp(x[i] - m) / s [pass 3: normalize]
Safe because exp(x[i] - m) ≤ exp(0) = 1
STAGE 3: ONLINE SOFTMAX (2-pass → fused)
Maintain running (max, normalizer) in single pass:
m = -∞, d = 0.0
for each x[i]:
m_prev = m
m = max(m, x[i])
d = d * exp(m_prev - m) + exp(x[i] - m)
┗━ rescale old sum when max updates ━━━━━━━┛
out[i] = exp(x[i] - m) / d
Why this works: invariant that d = Σexp(x[j]-m) for j seen so farCUDA Kernel Design
Each CUDA thread block handles one softmax row (one sequence position)
WARP-LEVEL REDUCTION:
// 32 threads in a warp reduce together
// Find max across warp
for (int mask = 16; mask > 0; mask >>= 1)
val = max(val, __shfl_xor_sync(0xffffffff, val, mask));
// Sum across warp (after subtracting max)
for (int mask = 16; mask > 0; mask >>= 1)
sum += __shfl_xor_sync(0xffffffff, sum, mask);
// Normalize
out[tid] = exp(x[tid] - warp_max) / warp_sum;
MEMORY LAYOUT:
Input: [batch, heads, seq_len, seq_len] (attention scores)
Each thread block: one (b, h, i) row → processes seq_len elements
Vectorized load: float4 → process 4 elements per thread per load
FUSION WITH ATTENTION:
Instead of: MatMul → store → Softmax → store → MatMul
Do: MatMul + Softmax + MatMul in one kernel pass (Flash Attention)
Eliminates 2 HBM round-trips per layerPerformance Targets
| Implementation | Throughput | HBM Reads |
|---|---|---|
| Naive 3-pass | baseline | 3× seq_len |
| Online 2-pass | 1.5× | 2× seq_len |
| Fused warp reduction | 3× | 1× seq_len |
| Flash Attention fused | 10×+ | O(1) per tile |
Project 4 — Transformer from Scratch
Full Architecture
INPUT
token_ids: (B, T) ← B=batch, T=sequence length
EMBEDDING LAYER
tok_emb = Embedding(vocab_size, d_model) → (B, T, d_model)
pos_emb = Embedding(max_seq_len, d_model) → (T, d_model) [learned]
OR
pos_enc = sinusoidal:
PE[pos, 2i] = sin(pos / 10000^(2i/d_model))
PE[pos, 2i+1] = cos(pos / 10000^(2i/d_model))
x = tok_emb + pos_emb → (B, T, d_model)
x = Dropout(x)
TRANSFORMER BLOCK × N_layers
┌───────────────────────────────────────────────────┐
│ [Pre-LayerNorm] │
│ x_norm = LayerNorm(x) │
│ │
│ [Multi-Head Self-Attention] │
│ Q = x_norm @ W_Q (d_model → d_model) │
│ K = x_norm @ W_K │
│ V = x_norm @ W_V │
│ │
│ Reshape: (B, T, d_model) → (B, n_heads, T, d_k) │
│ d_k = d_model / n_heads │
│ │
│ scores = Q @ K^T / sqrt(d_k) → (B, H, T, T) │
│ mask: upper triangle = -∞ (causal/autoregressive)│
│ attn = softmax(scores) → (B, H, T, T) │
│ attn = Dropout(attn) │
│ out = attn @ V → (B, H, T, d_k) │
│ out = reshape + W_O → (B, T, d_model) │
│ │
│ x = x + Dropout(out) ← residual connection │
│ │
│ [FFN] │
│ x_norm2 = LayerNorm(x) │
│ ffn = Linear(d_model, 4*d_model) │
│ ffn = GELU(ffn) │
│ ffn = Linear(4*d_model, d_model) │
│ ffn = Dropout(ffn) │
│ x = x + ffn ← residual connection │
└───────────────────────────────────────────────────┘
FINAL HEAD
x = LayerNorm(x)
logits = Linear(d_model, vocab_size) → (B, T, vocab_size)
LOSS (next-token prediction)
targets = token_ids shifted left by 1
loss = CrossEntropy(logits.view(-1, vocab_size), targets.view(-1))Hyperparameters (GPT-2 Small Scale)
| Param | Value |
|---|---|
| d_model | 768 |
| n_heads | 12 |
| n_layers | 12 |
| d_ff | 3072 |
| vocab_size | 50257 |
| max_seq_len | 1024 |
| dropout | 0.1 |
Training Loop
model = Transformer(config).to(device)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=3e-4, betas=(0.9, 0.95),
weight_decay=0.1
)
# Cosine LR schedule with linear warmup
scheduler = CosineAnnealingLR(optimizer, T_max=max_steps)
scaler = torch.cuda.amp.GradScaler() # mixed precision
for step, (x, y) in enumerate(dataloader):
with torch.cuda.amp.autocast():
logits, loss = model(x, y)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
scheduler.step()Inference / Generation
@torch.no_grad()
def generate(model, prompt_ids, max_new_tokens, temperature=1.0, top_k=50):
for _ in range(max_new_tokens):
# Crop context to max_seq_len
x = prompt_ids[:, -max_seq_len:]
logits = model(x)[:, -1, :] # last position only
logits = logits / temperature
# top-k sampling
v, _ = torch.topk(logits, top_k)
logits[logits < v[:, -1:]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
next_tok = torch.multinomial(probs, num_samples=1)
prompt_ids = torch.cat([prompt_ids, next_tok], dim=1)
return prompt_idsProject 5 — Flash Attention Kernel (CUDA)
The Core Problem
Standard attention: Attention(Q,K,V) = softmax(QK^T/√d) · V
- QK^T has shape (N, N) — quadratic in sequence length
- Must be stored in HBM for softmax then read back
- Memory: O(N²), compute: O(N²d)
- For N=8192, d=128: 8192² × 4 bytes = 256 MB per head per layer
Flash Attention Solution
Insight: Never materialize the full N×N attention matrix. Process it in tiles that fit in SRAM (on-chip, ~19 TB/s bandwidth).
TILING STRATEGY
Split Q into blocks of size Br rows
Split K,V into blocks of size Bc rows
Process each (Qi, Kj, Vj) combination
For each query block Qi:
Initialize: Oi = 0, ℓi = 0, mi = -∞
For each key/value block (Kj, Vj):
Load Qi, Kj, Vj from HBM → SRAM
Sij = Qi @ Kj^T / sqrt(d) ← local attention scores
m̃ij = rowmax(Sij) ← local max
P̃ij = exp(Sij - m̃ij) ← local softmax numerator
ℓ̃ij = rowsum(P̃ij) ← local sum
# Online softmax update (safe across tiles)
mi_new = max(mi, m̃ij)
ℓi_new = exp(mi - mi_new) * ℓi + exp(m̃ij - mi_new) * ℓ̃ij
# Rescale accumulated output and add new contribution
Oi = (exp(mi - mi_new) * ℓi * Oi + exp(m̃ij - mi_new) * P̃ij @ Vj) / ℓi_new
mi, ℓi = mi_new, ℓi_new
Write final Oi → HBM
Store (ℓi, mi) → HBM for backward passSRAM Tile Size Calculation
SRAM budget: 48 KB (typical for shared mem per block)
d = head_dim = 128
Tile holds: Qi (Br×d), Kj (Bc×d), Vj (Bc×d) in FP16
Memory: 2 * (Br + 2*Bc) * d bytes
Set Br = Bc = B:
2 * 3B * 128 = 768B bytes ≤ 48,000 bytes
B ≤ 62 → use B = 64 (power of 2)Backward Pass (Recomputation)
During forward: store only (O, ℓ, m) — NOT the N×N attention matrix
During backward:
Recompute Sij, Pij from stored Q, K, V, (ℓ, m)
Compute dV, dK, dQ using recomputed values
Memory for backward: O(N) instead of O(N²)
Compute overhead: ~33% more FLOPs (recomputation)
Net result: dramatically faster due to I/O reductionPerformance Comparison
Sequence Length N = 2048, d = 64, Batch = 8, Heads = 12:
Standard Attention:
HBM reads: QKV load + attn matrix write/read = 9 * N*d * 2 bytes = 4.7 GB
HBM writes: attn matrix = N²*B*H * 2 bytes = 3.2 GB
Total: ~8 GB of HBM traffic per forward pass
Flash Attention:
HBM reads: QKV load = 3 * N*d*B*H * 2 bytes = 4.7 GB
HBM writes: O = N*d*B*H * 2 bytes = 1.6 GB
Total: ~6.3 GB — but with NO N² term
At N=8192: 10-20x less HBM traffic
Wall-clock speedup: 2-4x on A100 (memory bandwidth bound)Project 6 — Vision Transformer (ViT)
Architecture
INPUT IMAGE: (B, 3, H, W) = (B, 3, 224, 224)
PATCH EMBEDDING
Divide image into P×P non-overlapping patches
P=16 → (224/16)² = 196 patches per image
Each patch: (3, 16, 16) → flatten → 768-dim vector via linear layer
Implementation: Conv2d(3, d_model, kernel_size=P, stride=P)
→ same as linear projection of flattened patches
Add [CLS] token (learnable) prepended → seq_len = 197
Add positional embeddings (learnable, shape 197×768)
x: (B, 197, 768)
TRANSFORMER ENCODER × 12 layers
Same as language transformer EXCEPT:
- No causal mask (bidirectional attention)
- [CLS] token attends to all patches; patches attend to each other
MSA (Multi-head Self-Attention):
Q, K, V projections (768 → 768)
12 heads, d_k = 64
scores = QK^T / 8.0 (= √64)
attn = softmax(scores) ← no mask
out = attn @ V → project back
MLP block:
GELU activation (not ReLU)
Expansion: 768 → 3072 → 768
CLASSIFICATION HEAD
Take [CLS] token output: (B, 768)
LayerNorm → Linear(768, num_classes)
→ Softmax → class probabilities
FULL FORWARD:
(B,3,224,224) → PatchEmbed → (B,197,768) → TransformerBlocks × 12
→ CLS extract → (B,768) → MLP Head → (B,1000)Pretraining Strategies
1. SUPERVISED (original ViT paper)
Train directly on ImageNet-21k (14M images) then fine-tune on IN-1k
Requires large data — ViT underperforms CNNs on small datasets
2. MAE (Masked Autoencoders, He et al 2021)
Mask 75% of patches at random
Encoder processes only visible patches (fast — 4× fewer tokens)
Decoder (lighter) reconstructs masked patches in pixel space
Loss: MSE on masked patches only
Masking strategies: random, block, grid
Key insight: image reconstruction is harder than it seems → good features
3. DINO (Self-supervised)
Student network sees augmented view A
Teacher network (EMA of student) sees augmented view B
Minimize cross-entropy between student/teacher CLS token distributions
Teacher never trained directly — only updated via:
θ_teacher = 0.996 * θ_teacher + 0.004 * θ_studentFine-tuning Protocol
Pretrained ViT-B/16 (ImageNet-21k) → Fine-tune on target dataset
Settings:
lr = 3e-3 (higher than training from scratch)
weight_decay = 0.01
batch_size = 512
Augmentation: RandAugment(num_ops=2, magnitude=9)
Mixup(alpha=0.2), CutMix(alpha=1.0)
Random erase probability = 0.25
Interpolate positional embeddings if changing resolution
(bicubic interpolation of pretrained pos_emb grid)Project 7 — Diffusion Model (UNet + Scheduler)
High-Level Intuition
Learn to reverse a gradual noising process. Training: corrupt data with noise, predict the noise. Inference: start from pure noise, iteratively denoise.
Forward Process (Fixed — No Learned Parameters)
NOISE SCHEDULE
β₁, β₂, ..., β_T = linear schedule from 1e-4 to 0.02 (T=1000)
αt = 1 - βt
ᾱt = ∏(α₁ · α₂ · ... · αt) ← cumulative product
CLOSED-FORM NOISING (key property — skip to any timestep directly)
q(xt | x₀) = N(√ᾱt · x₀, (1 - ᾱt) · I)
Sample: xt = √ᾱt · x₀ + √(1-ᾱt) · ε, ε ~ N(0,I)
At t=0: xt ≈ x₀ (clean)
At t=T: xt ≈ ε ~ N(0,I) (pure noise)UNet Architecture
INPUT: (B, C, H, W) = noisy image xt
+ timestep t (encoded as sinusoidal embedding → MLP → vector)
OUTPUT: (B, C, H, W) = predicted noise ε̂
ENCODER (downsampling path)
Resolution: 256 → 128 → 64 → 32 → 16 → 8
At each resolution:
ResBlock × 2
AttentionBlock (at resolutions ≤ 32)
Downsample: stride-2 Conv (or AvgPool + Conv)
ResBlock:
GroupNorm → SiLU → Conv3×3
+ time_emb projection (Linear → SiLU → Linear → add)
GroupNorm → SiLU → Dropout → Conv3×3
+ residual (1×1 conv if channels change)
BOTTLENECK
ResBlock → AttentionBlock → ResBlock
(operates at lowest spatial resolution)
DECODER (upsampling path)
At each resolution:
Upsample: bilinear or ConvTranspose2d
Concatenate skip connection from encoder (same resolution)
ResBlock × 2 (doubled channels due to skip)
AttentionBlock (at resolutions ≤ 32)
OUTPUT
GroupNorm → SiLU → Conv3×3(channels_out)
Output channels = input channels (predict noise ε̂)
CONDITIONING
Time: sinusoidal embed(t) → Linear(256, d_model) → inject via scale+shift in GroupNorm
Class: learned class embedding → add to time embedding
Text: cross-attention in attention blocks (Q from image features, KV from text encoder)Training
for x0 in dataloader: # clean images (B, C, H, W)
t = torch.randint(0, T, (B,)) # random timestep per sample
ε = torch.randn_like(x0) # random noise
# Noise the image (closed form — no loop needed)
sqrt_alpha_bar = extract(sqrt_alphas_bar, t, x0.shape)
sqrt_one_minus = extract(sqrt_one_minus_alphas_bar, t, x0.shape)
xt = sqrt_alpha_bar * x0 + sqrt_one_minus * ε
# Predict noise
ε_pred = unet(xt, t)
# Simple MSE loss (from DDPM paper)
loss = F.mse_loss(ε_pred, ε)
loss.backward()
optimizer.step()
optimizer.zero_grad()
# EMA of weights for sampling
ema.update(unet)Inference — DDPM Sampler (1000 steps)
@torch.no_grad()
def sample_ddpm(model, shape):
x = torch.randn(shape) # start from pure noise
for t in reversed(range(T)): # t: 999 → 0
t_batch = torch.full((B,), t)
ε_pred = model(x, t_batch) # predict noise
# Compute x₀ prediction
x0_pred = (x - sqrt_one_minus_bar[t] * ε_pred) / sqrt_bar[t]
x0_pred = x0_pred.clamp(-1, 1) # clip to valid image range
# Posterior mean
mu = posterior_mean_coef1[t] * x0_pred + posterior_mean_coef2[t] * x
if t > 0:
σ = posterior_variance[t] ** 0.5
x = mu + σ * torch.randn_like(x) # add noise (except final step)
else:
x = mu
return x # (B, C, H, W) generated imagesInference — DDIM Sampler (50 steps, deterministic)
@torch.no_grad()
def sample_ddim(model, shape, steps=50, eta=0.0):
# Select subset of timesteps (strided)
timesteps = torch.linspace(T-1, 0, steps).long()
x = torch.randn(shape)
for i, t in enumerate(timesteps):
ε_pred = model(x, t)
# DDIM step (deterministic when eta=0)
x0_pred = (x - sqrt_one_minus_bar[t] * ε_pred) / sqrt_bar[t]
if i < len(timesteps) - 1:
t_prev = timesteps[i+1]
x = sqrt_bar[t_prev] * x0_pred + sqrt_one_minus_bar[t_prev] * ε_pred
# eta > 0 adds stochasticity back (DDIM → DDPM as eta → 1)
else:
x = x0_pred
return xLatent Diffusion (Stable Diffusion Extension)
VAE ENCODER → latent z (8× smaller) → DIFFUSION IN LATENT SPACE → VAE DECODER → image
Benefit: 64×64 latents instead of 512×512 pixels
UNet processes 64×64 — 64× fewer elements
8× compression ratio in each spatial dimension
Text conditioning:
CLIP/T5 text encoder → text embeddings (77, 768)
Cross-attention in UNet decoder: Q=image_features, KV=text_embeddings
Classifier-Free Guidance:
ε_guided = ε_uncond + w * (ε_cond - ε_uncond)
w=7.5 typical — trades diversity for prompt adherenceHow These 7 Projects Form One System
┌─────────────────────────────────────┐
│ MINI PYTORCH │
│ │
Autograd ──▶ │ Gradient computation engine │
MatMul kernel ──▶ │ GPU compute primitive │
Softmax kernel ─▶ │ Numerical stability + fusion │
Transformer ──▶ │ Core architecture │
Flash Attention ▶ │ Memory-efficient attention │
ViT ──────────▶ │ Extends to vision │
Diffusion ────▶ │ Extends to generation │
└─────────────────────────────────────┘
Dependency graph:
Autograd
└─▶ Transformer
├─▶ Flash Attention (makes Transformer fast)
├─▶ ViT (swaps token embeddings for patch embeddings)
└─▶ Diffusion (UNet uses ResBlocks + Transformer attention blocks)
MatMul kernel
└─▶ Softmax kernel (both are inner loops of attention)
└─▶ Flash Attention (fuses them)Mini vLLM — E2E Architecture Reference
Projects: Inference Server (C++/Rust) · KV Cache Paging System · Speculative Decoding · Quantization Library (Int8/FP4) · Continuous Batching
What You're Building
The serving stack that takes a trained model and makes it run efficiently in production. This is what separates a model that runs on one GPU for one user from a system that serves thousands of concurrent requests, maximizing GPU utilization while minimizing latency. After completing this system you will understand every design decision behind vLLM, TensorRT-LLM, and TGI.
Project 1 — Inference Server (C++/Rust)
Full System Architecture
┌──────────────────────────────────────┐
│ INFERENCE SERVER │
│ │
Client ──HTTP/gRPC──▶ │ ┌─────────────┐ │
│ │ API Layer │ /v1/completions │
│ │ (REST/gRPC)│ /v1/chat │
│ └──────┬──────┘ /v1/embeddings │
│ │ │
│ ┌──────▼──────────────────────────┐ │
│ │ Request Queue │ │
│ │ Priority queue or FIFO │ │
│ │ Fields per request: │ │
│ │ - request_id (uuid) │ │
│ │ - prompt_token_ids │ │
│ │ - sampling_params │ │
│ │ - max_new_tokens │ │
│ │ - arrival_timestamp │ │
│ └──────┬──────────────────────────┘ │
│ │ │
│ ┌──────▼──────────────────────────┐ │
│ │ Scheduler │ │
│ │ Continuous batching engine │ │
│ │ - Groups prefill + decode │ │
│ │ - Preemption on OOM │ │
│ │ - Budget: max_batch_tokens │ │
│ └──────┬──────────────────────────┘ │
│ │ │
│ ┌──────▼────────────────────────┐ │
│ │ KV Cache Manager │ │
│ │ Page allocation/deallocation │ │
│ │ Prefix cache lookup │ │
│ └──────┬────────────────────────┘ │
│ │ │
│ ┌──────▼────────────────────────┐ │
│ │ Model Runner │ │
│ │ CUDA kernel dispatch │ │
│ │ Prefill + Decode batches │ │
│ └──────┬────────────────────────┘ │
│ │ │
│ ┌──────▼────────────────────────┐ │
│ │ Sampler │ │
│ │ temp / top-p / top-k / rep │ │
│ └──────┬────────────────────────┘ │
│ │ │
│ ┌──────▼────────────────────────┐ │
│ │ Detokenizer │ │
│ │ BPE decode → UTF-8 │ │
│ │ Streaming via SSE │ │
│ └───────────────────────────────┘ │
└──────────────────────────────────────┘Threading Model (C++)
// Main thread: HTTP server
auto server = httplib::Server();
server.Post("/v1/completions", [&](const Request& req, Response& res) {
auto r = parse_request(req.body);
r.id = generate_uuid();
request_queue.push(r);
// Wait on promise/future for response
auto result = r.promise.get_future().get();
res.set_content(result.to_json(), "application/json");
});
// Scheduler thread: continuous batching loop
void scheduler_loop() {
while (running) {
auto batch = scheduler.get_next_batch(request_queue, kv_cache_manager);
if (batch.empty()) { std::this_thread::sleep_for(1ms); continue; }
auto outputs = model_runner.forward(batch);
sampler.sample(outputs, batch);
for (auto& req : batch.finished_requests)
req.promise.set_value(req.output);
scheduler.update(batch); // update state, free completed requests
}
}Request Lifecycle
1. ARRIVE
Client sends POST /v1/completions
Server validates, assigns request_id, enqueues
2. WAIT IN QUEUE
Scheduler checks memory availability before admitting
If insufficient KV pages: request waits or lower-priority req preempted
3. PREFILL (first forward pass)
All prompt tokens processed in PARALLEL
Fills KV cache for all prompt positions
Most compute-intensive step: O(prompt_len²) attention
4. DECODE (subsequent forward passes)
ONE token generated per forward pass
Uses cached K,V from previous positions
Appends new K,V to cache
Memory-bandwidth bound (reads weights + KV cache each step)
5. SAMPLING
Logits (vocab_size) → apply temperature/top-p → sample token
Check stop conditions: max_tokens, stop strings, EOS token
6. STREAM / RETURN
SSE: send token immediately as generated
Non-streaming: buffer all tokens, return when done
7. CLEANUP
Free KV cache pages
Decrement reference counts (prefix cache sharing)
Update metrics (tokens/sec, latency p50/p99/p999)Metrics & Observability
Prometheus metrics to expose:
- inference_requests_total (counter, by model/status)
- inference_tokens_per_second (gauge)
- kv_cache_utilization (gauge, 0-1)
- request_queue_depth (gauge)
- time_to_first_token_seconds (histogram)
- inter_token_latency_seconds (histogram)
- batch_size (histogram)
Logging: structured JSON logs per request
{ request_id, prompt_tokens, completion_tokens, latency_ms, model }Project 2 — KV Cache Paging System (like vLLM)
The Problem
NAIVE APPROACH
Pre-allocate KV buffer: (max_seq_len, d_model) per request
For max_seq_len=2048, d=4096, 32 layers, FP16:
2048 * 4096 * 2 * 32 * 2 bytes = 1 GB per request
8x A100 (640 GB total) → ~640 concurrent requests MAX
Reality: average output < 200 tokens → 90% of allocation wasted
Memory fragmentation prevents packing more requests
PAGED ATTENTION SOLUTION
Manage KV memory like OS manages RAM:
Physical memory = fixed-size pages (blocks)
Virtual memory = logical view per request
Page table = mapping logical → physical blocksMemory Layout
PHYSICAL BLOCK POOL (GPU HBM)
Total HBM: 80 GB (A100)
Model weights: ~14 GB (7B param FP16)
Activation mem: ~4 GB (batch processing)
KV Cache Pool: ~60 GB available
Block size: 16 tokens
Per-block size: 2 (K+V) × d_kv × n_layers × dtype_size
= 2 × 128 × 32 × 2 = 16 KB per block
Total blocks: 60 GB / 16 KB = ~3.8M blocks
BLOCK STRUCTURE (in GPU memory)
block[b].key = float16[n_layers, n_heads, BLOCK_SIZE, d_head]
block[b].value = float16[n_layers, n_heads, BLOCK_SIZE, d_head]
Alternatively flat: block[b] = float16[2, n_layers, n_heads, BLOCK_SIZE, d_head]Block Manager (CPU-side)
class BlockManager:
def __init__(self, n_blocks, block_size):
self.n_blocks = n_blocks
self.block_size = block_size
self.free_blocks = deque(range(n_blocks)) # block_ids
self.ref_count = [0] * n_blocks
self.prefix_cache = {} # hash(tokens) → block_id
def allocate(self, n_tokens: int) -> List[int]:
"""Allocate physical blocks for n_tokens."""
n_blocks = ceil(n_tokens / self.block_size)
if len(self.free_blocks) < n_blocks:
raise OOMError("Insufficient KV cache pages")
allocated = []
for _ in range(n_blocks):
block_id = self.free_blocks.popleft()
self.ref_count[block_id] = 1
allocated.append(block_id)
return allocated
def free(self, block_table: List[int]):
"""Release blocks back to pool."""
for block_id in block_table:
self.ref_count[block_id] -= 1
if self.ref_count[block_id] == 0:
self.free_blocks.append(block_id)
def fork(self, block_table: List[int]) -> List[int]:
"""Copy-on-write fork (for beam search / prefix sharing)."""
new_table = block_table.copy()
for block_id in new_table:
self.ref_count[block_id] += 1
return new_table
def get_or_cache_prefix(self, tokens: List[int]) -> Optional[List[int]]:
"""Return cached blocks for a prompt prefix if available."""
for length in range(len(tokens), 0, -self.block_size):
prefix = tokens[:length]
key = hash(tuple(prefix))
if key in self.prefix_cache:
cached_blocks = self.prefix_cache[key]
# Increment ref counts (shared ownership)
for b in cached_blocks:
self.ref_count[b] += 1
return cached_blocks, length # blocks + how many tokens covered
return None, 0Paged Attention CUDA Kernel
Standard attention reads K,V from contiguous tensors:
K: [batch, heads, seq_len, d_head] ← contiguous
Paged attention reads K,V via indirection through block table:
block_table: [batch, max_blocks_per_seq] ← CPU-provided mapping
Kernel pseudocode:
For each query position (batch b, head h, seq position q):
acc = 0.0
for logical_block = 0 to num_blocks:
phys_block = block_table[b, logical_block]
for slot = 0 to BLOCK_SIZE:
k = K_cache[phys_block, h, slot, :]
v = V_cache[phys_block, h, slot, :]
score = dot(q_vec, k) / sqrt(d_head)
acc += softmax_weight * v
output[b, h, q, :] = acc
Memory access pattern: non-contiguous but coalesced within a block
Performance overhead vs contiguous: ~10-15% (acceptable for 3x+ capacity gain)Preemption Strategy
WHEN: new request admitted but KV cache runs OOM mid-generation
OPTIONS:
1. SWAP to CPU RAM
- Copy physical blocks from GPU HBM → CPU DRAM
- Resume later by swapping back
- Overhead: PCIe bandwidth (~32 GB/s) — slow
2. RECOMPUTE (drop and restart)
- Free all blocks for the preempted request
- When rescheduled: re-run prefill from scratch
- Simpler to implement, works well for short prompts
3. PARTIAL SWAP
- Keep recent blocks on GPU (needed for decode)
- Swap older/prefix blocks to CPU
- Best for long-context requests
SCHEDULING POLICY:
Preempt request with lowest priority (FCFS: oldest completed most tokens)
Never preempt the same request twice in a row (livelock prevention)Project 3 — Speculative Decoding System
Core Algorithm
PROBLEM: LLM decode is 1 token/forward-pass — memory BW bound
Target model (7B) runs at ~50 tokens/sec on A100
INSIGHT: GPU can run one 7B forward pass OR seven 7B forward passes
in nearly the same wall-clock time (latency bound, not throughput)
Instead: draft K tokens with tiny model, verify ALL in one pass
Draft Model Target Model
(68M params) (7B params)
──────────── ────────────
Step 1: [fast, K tokens] [idle]
Step 2: [idle] [verify K+1 positions in one pass]
Step 3: accept/reject each draft token
If all K accepted → K+1 tokens generated in cost of 1 target forward pass
If first token rejected → 1 token, discard rest
Average speedup: 2-3x depending on task/acceptance rateFull Algorithm (Speculative Sampling)
def speculative_decode(
target_model, draft_model, context, max_tokens, K=5
):
output = context.copy()
while len(output) - len(context) < max_tokens:
# ─── DRAFT PHASE ───────────────────────────────────────────
draft_tokens = []
draft_probs = []
draft_input = output.copy()
for _ in range(K):
logits = draft_model.forward(draft_input)
probs = softmax(logits[-1]) # last position
token = sample(probs) # draft sample
draft_tokens.append(token)
draft_probs.append(probs[token])
draft_input.append(token)
# ─── VERIFY PHASE (ONE TARGET FORWARD PASS) ────────────────
# Process output + all K draft tokens simultaneously
target_logits = target_model.forward(output + draft_tokens)
# target_logits has shape (K+1, vocab_size)
# position i gives distribution over token at position len(output)+i
# ─── ACCEPTANCE SAMPLING ───────────────────────────────────
n_accepted = 0
for i in range(K):
target_p = softmax(target_logits[i])[draft_tokens[i]]
draft_p = draft_probs[i]
r = random.uniform(0, 1)
if r < min(1.0, target_p / draft_p): # accept
output.append(draft_tokens[i])
n_accepted += 1
else: # reject
# Sample from adjusted distribution
# p_adjusted[x] = max(0, target_p[x] - draft_p[x]) / Z
adj_probs = relu(softmax(target_logits[i]) - draft_probs_full[i])
adj_probs /= adj_probs.sum()
output.append(sample(adj_probs))
break
else:
# All K accepted! Sample bonus token from final target position
output.append(sample(softmax(target_logits[K])))
n_accepted += 1 # always get at least 1 token
return output
# Key properties:
# - Provably identical distribution to running target model alone
# - Expected tokens per iteration = K * acceptance_rate + 1
# - No quality degradation, only speed changeDraft Model Selection
REQUIREMENTS:
- Same tokenizer as target (identical vocabulary essential)
- Much smaller (5-10x fewer params) for speed benefit
- Similar distribution to target (high acceptance rate)
STRATEGIES:
1. PURPOSE-TRAINED DRAFT
Distill from target model → dedicated small model
Best acceptance rate (trained to match target exactly)
Example: Medusa (multi-head on target), EAGLE (feature-based draft)
2. SAME-FAMILY SMALL MODEL
Llama-3.1-8B draft → Llama-3.1-70B target
Good acceptance for in-distribution text
Acceptance rate: 70-85% for typical generation tasks
3. N-GRAM / RETRIEVAL DRAFT
For repetitive text (code completion, structured output)
Match recent context to generate candidates
Zero parameters, instant draft — but narrow applicability
4. MEDUSA (multi-head speculative)
Attach K extra LM heads to target model
Each head predicts token at position +1, +2, ..., +K
Draft and verify in same forward pass (no separate model)
Verification uses tree attention to evaluate multiple token pathsSystem Integration
MEMORY:
Draft model fits in same GPU (68M << 7B)
KV cache: maintain SEPARATE caches for draft and target
On rejection: rollback target KV cache to last accepted position
Draft KV rollback: trivial (just reset position pointer)
Target KV rollback: must track position, truncate last K entries
BATCHING:
Speculative decoding is naturally sequential per request
Batch speculative: run draft for all requests, then batch verify with target
Complexity: accepted length varies per request → variable-length batches
DYNAMIC K:
Adjust K based on observed acceptance rate:
acceptance_rate > 0.8 → increase K
acceptance_rate < 0.5 → decrease K
Track per-request or per-model exponential moving averageProject 4 — Quantization Library (Int8/FP4)
Why Quantization
Llama-3.1-70B weights:
FP32: 70B × 4 bytes = 280 GB (requires 4× A100 just for weights)
BF16: 70B × 2 bytes = 140 GB
Int8: 70B × 1 byte = 70 GB ← fits in 1× H100 (80 GB)
FP4: 70B × 0.5 byte = 35 GB ← fits in 1× A100 (40 GB)
Memory bandwidth also reduced proportionally → faster decodePost-Training Quantization (PTQ) Pipeline
INPUT: Pre-trained FP32/BF16 model weights
OUTPUT: Quantized model + scale factors + zero points
STEP 1: Calibration
Run forward passes on ~512 calibration samples
Collect activation statistics: min, max, percentiles per layer
Important: outliers in activations (specific channels >> others)
STEP 2: Weight Quantization
For each weight tensor W of shape (out, in):
Per-tensor (cheapest):
scale = max(|W|) / 127
W_q = round(W / scale).clamp(-128, 127).to(int8)
W_dq = W_q.float() * scale # dequantized approximation
Per-channel (better accuracy):
scale_per_row = max(|W|, dim=1) / 127 # shape: (out,)
W_q = round(W / scale_per_row[:, None]).clamp(-128, 127).to(int8)
Per-group (best accuracy, ~4-8 tokens grouped):
group_size = 128
W reshaped: (out, in/128, 128)
scale per group
STEP 3: Activation Quantization (for Int8 GEMM)
Static: pre-computed scale from calibration data
Dynamic: compute scale at runtime from actual activation values
scale = max(|act|) / 127 (one pass over activation)
act_q = round(act / scale).clamp(-128, 127).to(int8)
STEP 4: Int8 GEMM
y = W_q @ x_q ← integer matrix multiply (fast)
y_float = y.float() * scale_W * scale_x ← dequantize outputGPTQ (Gradient-based PTQ, Weight-Only)
Insight: Minimizing per-layer output error is better than minimizing
weight quantization error directly
Algorithm (per linear layer):
for each column j in weight matrix W:
w_j = W[:, j] # column to quantize
quant(w_j) # quantize to Int4/Int8
error_j = w_j - dequant(quant(w_j)) # quantization error
# Compensate remaining columns for this error
# using second-order information (Hessian of activation covariance)
W[:, j+1:] -= error_j @ H_inv[j, j+1:]
# This is the "lazy batch" variant — update in groups of 128 cols
Result: quantized weights that minimize output error, not weight error
Layer-wise: run GPTQ on each transformer layer sequentially
Quality: FP16-comparable at Int4, with 4x memory reduction
Cost: ~1-4 GPU-hours for 7B model (one-time)FP4 (NF4 — NormalFloat4)
MOTIVATION: weights of quantized networks follow roughly normal distribution
→ allocate bits based on quantile positions, not linear range
NF4 QUANTILE LEVELS (16 levels for 4 bits):
Compute 16 quantiles of N(0,1):
[-1.0, -0.6962, -0.5251, -0.3949, -0.2844, -0.1848, -0.0911, 0.0,
0.0911, 0.1848, 0.2844, 0.3949, 0.5251, 0.6962, 1.0, 1.0]
These are equally-spaced PROBABILITY levels, not value levels
QUANTIZATION:
Normalize weights: w_norm = w / max(|w|)
For each w_norm: find nearest NF4 level (index 0-15)
Store 4-bit index
DEQUANTIZATION:
Look up NF4 table by index
Multiply by stored scale (max absolute value)
DOUBLE QUANTIZATION (QLoRA trick):
Scale factors themselves are quantized (FP8)
Scale of scales: one FP32 value per group
Memory saving: 0.5 bits per parameter for scale storage
BITSANDBYTES IMPLEMENTATION:
Uses block-wise quantization: 64-weight blocks
Each block has own scale factor
GPU kernel: blockwise_fp4_dequant_gemmQuantization-Aware Training (QAT)
IDEA: Simulate quantization during training so model learns to be robust
FAKE QUANTIZE op (differentiable):
Forward: x_q = round(x / scale) * scale ← quantize + dequantize
Backward: pass-through (Straight-Through Estimator)
∂L/∂x = ∂L/∂x_q (ignore quantization in backward)
TRAINING LOOP:
1. Initialize from FP32 pre-trained weights
2. Insert FakeQuantize nodes after linear layers
3. Fine-tune with small learning rate (1e-5 typical)
4. After training: extract quantized weights + scales
RESULT: better accuracy than PTQ, but requires training data + computeAccuracy Recovery Techniques
PROBLEM: Quantization kills accuracy for certain layers
LLM.int8() solution (Dettmers et al):
OBSERVATION: 0.1% of weight channels contain "outliers" (>>100x average)
SOLUTION: decompose matmul:
- 99.9% of columns → Int8 quantization (fast)
- 0.1% outlier columns → FP16 (keep full precision)
- Merge results
SmoothQuant:
INSIGHT: activations are harder to quantize than weights (more outliers)
SOLUTION: mathematically transfer quantization difficulty from activations to weights
X_new = X / s (easier to quantize activations)
W_new = W * s (slightly harder to quantize weights)
XW = X_new * W_new (same result)
where s is per-channel smoothing factor derived from activation statsProject 5 — Continuous Batching
The Problem with Static Batching
STATIC BATCHING (naive)
Collect B requests, run together, return when ALL finish
Request 1: ████████ (8 tokens)
Request 2: ██████████████████████ (22 tokens)
Request 3: ████ (4 tokens)
Request 4: ██████████████ (14 tokens)
All start together. All wait for Request 2 to finish.
GPU utilization during tail: 25% (only req 2 running)
Batch throughput bottlenecked by longest request in batch.Continuous Batching Solution
CONTINUOUS BATCHING (iteration-level scheduling)
After each decode step: check if any request finished
If yes: evict finished request, admit new request from queue
Timeline:
Step 1: [R1 R2 R3 R4] ← initial batch
Step 2: [R1 R2 R3 R4]
Step 4: [R1 R2 R3 R4] ← R3 finishes, R5 admitted
Step 5: [R1 R2 R5 R4]
Step 8: [R1 R2 R5 R4] ← R1 finishes, R6 admitted
Step 9: [R6 R2 R5 R4]
...
GPU always busy, no idle slots waiting for slow requests
Throughput improvement: 2-4x in practiceScheduler Implementation
class ContinuousBatchingScheduler:
def __init__(self, max_batch_tokens, kv_manager):
self.max_batch_tokens = max_batch_tokens # budget
self.kv_manager = kv_manager
self.running: List[Request] = [] # currently decoding
self.waiting: Queue[Request] = Queue() # admitted, waiting for KV
self.swapped: List[Request] = [] # preempted to CPU
def schedule(self) -> SchedulerOutput:
# 1. Check finished requests
for req in self.running:
if req.is_done():
self.kv_manager.free(req.block_table)
self.running = [r for r in self.running if not r.is_done()]
# 2. Try to resume swapped requests
while self.swapped and self._has_capacity():
req = self.swapped.pop(0)
self._swap_in(req)
self.running.append(req)
# 3. Admit new requests from waiting queue
while not self.waiting.empty() and self._has_capacity():
req = self.waiting.get()
try:
blocks = self.kv_manager.allocate(len(req.prompt_tokens))
req.block_table = blocks
self.running.append(req)
except OOMError:
# Preempt lowest priority running request
victim = min(self.running, key=lambda r: r.priority)
self._swap_out(victim)
self.running.remove(victim)
self.swapped.append(victim)
# 4. Construct batch
prefill_reqs = [r for r in self.running if r.phase == 'prefill']
decode_reqs = [r for r in self.running if r.phase == 'decode']
return SchedulerOutput(
prefill=prefill_reqs,
decode=decode_reqs,
)
def _has_capacity(self) -> bool:
used_tokens = sum(len(r.token_ids) for r in self.running)
free_blocks = self.kv_manager.n_free_blocks()
return used_tokens < self.max_batch_tokens and free_blocks > MIN_FREE_BLOCKSBatching Prefill and Decode Together
CHALLENGE: Prefill processes N tokens in parallel; Decode processes 1 token.
They have fundamentally different compute profiles.
CHUNKED PREFILL:
Long prompt would dominate a batch (thousands of tokens)
Solution: chunk prefill into pieces of max_chunk_size tokens
Interleave prefill chunks with decode steps
Schedule iteration:
Decode batch: all running requests (1 token each)
Prefill batch: one chunk from the next waiting request
Combined as single forward pass:
attention_mask handles that decode positions attend over full KV cache
while prefill chunk only attends over its own context
SEPARATION (simpler):
Two forward passes per scheduler step:
1. Prefill forward: all new requests in parallel
2. Decode forward: all decoding requests
Slightly less GPU utilization but simpler codeToken Budget Management
BUDGET CONSTRAINT: max_tokens_per_batch = max_batch_tokens
Decode batch:
Each request contributes exactly 1 input token
B decode requests → B tokens of compute (cheap per request)
KV reads: B requests × seq_len × d_kv (dominant cost)
Prefill batch:
New request contributes prompt_len tokens
10 requests of 200 tokens = 2000 tokens → quadratic attention
DYNAMIC BUDGET:
Reserve 80% of budget for decode (maintain latency)
Use remaining 20% for prefill (admit new requests)
decode_slots = floor(0.8 * max_batch_tokens)
prefill_budget = max_batch_tokens - actual_decode_tokens
admit_up_to = requests whose prompt_len fits in prefill_budgetHow These 5 Projects Form One System
┌──────────────────────────────────────────┐
│ MINI vLLM │
│ │
│ ┌─────────────┐ ┌────────────────┐ │
Request ────────▶ │ │ Inference │──▶│ Continuous │ │
│ │ Server │ │ Batching │ │
│ └─────────────┘ └───────┬────────┘ │
│ │ │
│ ┌────────▼───────┐ │
│ │ KV Cache │ │
│ │ Paging │ │
│ └────────┬───────┘ │
│ │ │
│ ┌──────────────────▼────────┐ │
│ │ Quantized Model (Int8) │ │
│ │ + Speculative Decoding │ │
│ └───────────────────────────┘ │
└──────────────────────────────────────────┘
Memory hierarchy of a production serving system:
GPU SRAM (~50 MB): Flash Attention tiles, active batch activations
GPU HBM (~80 GB): Model weights (quantized) + KV cache (paged)
CPU DRAM (~512 GB): Swapped KV pages, request queue, tokenizer
SSD (~10 TB): Model checkpoint, serving logs
Performance stack:
Quantization → 2-4x less memory → more requests in GPU simultaneously
KV Paging → 3x better GPU memory utilization vs static allocation
Continuous Batch → 2-4x throughput vs static batching
Speculative Dec → 2-3x latency reduction for individual requests
Combined: production system serving 10-50x more requests/sec vs naive implMini LangChain — E2E Architecture Reference
Projects: Chain-of-Thought Reasoner · ReAct Agent Loop · RAG Pipeline · Vector Database (HNSW) · Function Calling Router · Structured Output Parser · Semantic Router · Graph RAG · Knowledge Graph Builder · Code Interpreter Sandbox
What You're Building
The application layer that makes LLMs useful: orchestrating reasoning, retrieval, tool use, and structured outputs. After completing this system you understand how every major agent framework works internally, why retrieval quality determines application quality, and how to route and constrain LLM outputs programmatically.
Project 1 — Chain-of-Thought Reasoner
Architecture
STANDARD PROMPTING
Prompt: "What is 17 × 23?"
Output: "391"
Problem: Model pattern-matches, doesn't show work
Errors compound invisibly on multi-step problems
CHAIN-OF-THOUGHT PROMPTING
Prompt: "What is 17 × 23? Think step by step."
Output: "17 × 23 = 17 × 20 + 17 × 3 = 340 + 51 = 391"
Performance improvement:
GSM8K (math): 18% → 57% (GPT-3 scale)
MATH dataset: larger gains on harder problems
Works best at >100B params (emergent capability)Implementation Architecture
CoT SYSTEM ARCHITECTURE
┌───────────────────────────────────────────────────────────┐
│ REASONER │
│ │
│ Input: question │
│ ┌──────────────────────────────────────────────┐ │
│ │ PROMPT BUILDER │ │
│ │ │ │
│ │ System: "You are a careful reasoner. │ │
│ │ Think step by step before │ │
│ │ answering." │ │
│ │ │ │
│ │ Few-shot examples: │ │
│ │ Q: [example question] │ │
│ │ A: Let me think step by step. │ │
│ │ [reasoning steps...] │ │
│ │ Therefore, the answer is [X]. │ │
│ │ │ │
│ │ Q: {question} │ │
│ │ A: Let me think step by step. │ │
│ └──────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────┐ │
│ │ LLM API │ │
│ │ temperature = 0.0 │ ← deterministic │
│ │ stop = ["Q:"] │ ← stop before new Q │
│ └──────────┬──────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────┐ │
│ │ OUTPUT PARSER │ │
│ │ Extract reasoning │ │
│ │ Extract final ans │ │
│ └─────────────────────┘ │
└───────────────────────────────────────────────────────────┘CoT Variants
ZERO-SHOT CoT
Simply append: "Let's think step by step."
No examples needed. Surprisingly effective.
Best when: question domain is clear, model is large
FEW-SHOT CoT
Provide 4-8 examples with explicit reasoning chains
Examples must match domain (math examples for math tasks)
Best when: specialized domain, consistent format needed
SELF-CONSISTENCY CoT
Generate N reasoning paths (temperature > 0)
Final answer: majority vote across N outputs
N=40 paths typical. Trades compute for accuracy.
+3-5% accuracy over single chain
PROGRAM-OF-THOUGHT (PoT)
Instead of natural language reasoning → generate code
Execute code to get exact answer (no arithmetic errors)
for i, step in reasoning_steps:
# Each step is Python code
result = exec(step)
Especially powerful for math/computation
TREE-OF-THOUGHTS (ToT)
Maintain beam of k partial reasoning paths
Score each state with LLM: "Is this reasoning correct?"
Expand promising paths, prune bad ones
BFS or DFS through reasoning tree
Expensive but solves multi-step strategic problemsImplementation
class ChainOfThoughtReasoner:
def __init__(self, llm_client, examples: List[Example]):
self.llm = llm_client
self.examples = examples # (question, reasoning, answer) tuples
def build_prompt(self, question: str) -> str:
prompt = "You are a careful reasoner. Solve problems step by step.\n\n"
for ex in self.examples:
prompt += f"Q: {ex.question}\n"
prompt += f"A: Let me think step by step.\n{ex.reasoning}\n"
prompt += f"Therefore, the answer is: {ex.answer}\n\n"
prompt += f"Q: {question}\n"
prompt += "A: Let me think step by step.\n"
return prompt
def reason(self, question: str, n_samples: int = 1) -> str:
prompt = self.build_prompt(question)
if n_samples == 1:
response = self.llm.complete(prompt, temperature=0.0, stop=["Q:"])
return self.parse_answer(response)
else:
# Self-consistency: majority vote
answers = []
for _ in range(n_samples):
resp = self.llm.complete(prompt, temperature=0.7, stop=["Q:"])
answers.append(self.parse_answer(resp))
return Counter(answers).most_common(1)[0][0]
def parse_answer(self, response: str) -> str:
# Extract final answer after "Therefore" or "The answer is"
patterns = [
r"[Tt]herefore.*?(?:answer is|=)\s*(.+?)(?:\.|$)",
r"[Tt]he answer is:?\s*(.+?)(?:\.|$)",
r"=\s*(\d+\.?\d*)\s*$",
]
for p in patterns:
match = re.search(p, response)
if match: return match.group(1).strip()
# Fallback: last line
return response.strip().split('\n')[-1]Project 2 — ReAct Agent Loop
Architecture
ReAct = Reason + Act
Core loop: Thought → Action → Observation → Thought → ...
AGENT LOOP ARCHITECTURE
─────────────────────────
┌─────────────────────────────────────────────────────────┐
│ REACT AGENT │
│ │
│ User Input ──▶ "What's the weather in Tokyo today │
│ and should I bring an umbrella?" │
│ │
│ ┌───────────────────────────────────────────────────┐ │
│ │ CONTEXT WINDOW │ │
│ │ System prompt: "You are a helpful agent. │ │
│ │ Use tools to answer questions. │ │
│ │ Tools available: [weather_api, web_search] │ │
│ │ │ │
│ │ To use a tool, write: │ │
│ │ Thought: I need to check the weather │ │
│ │ Action: weather_api │ │
│ │ Action Input: {"city": "Tokyo"} │ │
│ │ The result will be returned as Observation. │ │
│ │ When done, write: Final Answer: [answer]" │ │
│ └─────────────────────┬─────────────────────────────┘ │
│ │ │
│ ┌──────────▼──────────┐ │
│ │ LLM │ │
│ └──────────┬──────────┘ │
│ │ generates: │
│ │ "Thought: I need to look up │
│ │ current weather in Tokyo │
│ │ Action: weather_api │
│ │ Action Input: {city: Tokyo}" │
│ │ │
│ ┌──────────▼──────────┐ │
│ │ OUTPUT PARSER │ parse action/input │
│ └──────────┬──────────┘ │
│ │ │
│ ┌──────────▼──────────┐ │
│ │ TOOL EXECUTOR │ │
│ │ weather_api(Tokyo) │ │
│ └──────────┬──────────┘ │
│ │ returns: │
│ │ "Rainy, 18°C, 90% precipitation│
│ │ │
│ Append to context: │
│ "Observation: Rainy, 18°C, 90% precipitation" │
│ │ │
│ [loop back to LLM] │
│ │ │
│ LLM generates: │
│ "Final Answer: Yes, bring an umbrella..." │
└─────────────────────────────────────────────────────────┘Full Implementation
class ReActAgent:
def __init__(self, llm, tools: List[Tool], max_steps=10):
self.llm = llm
self.tools = {t.name: t for t in tools}
self.max_steps = max_steps
def build_system_prompt(self) -> str:
tool_descriptions = "\n".join(
f"- {t.name}: {t.description}\n Input schema: {t.input_schema}"
for t in self.tools.values()
)
return f"""You are a helpful agent with access to tools.
Available tools:
{tool_descriptions}
To use a tool, respond with:
Thought: [your reasoning]
Action: [tool_name]
Action Input: [JSON input]
When you have the final answer, respond with:
Thought: [final reasoning]
Final Answer: [your answer]
Always think before acting. Always use exact tool names."""
def run(self, user_input: str) -> AgentResult:
messages = [
{"role": "system", "content": self.build_system_prompt()},
{"role": "user", "content": user_input}
]
trajectory = []
for step in range(self.max_steps):
# Generate next thought/action
response = self.llm.chat(messages)
messages.append({"role": "assistant", "content": response})
# Check if done
if "Final Answer:" in response:
answer = response.split("Final Answer:")[-1].strip()
return AgentResult(answer=answer, trajectory=trajectory)
# Parse action
try:
thought, action, action_input = self.parse_action(response)
except ParseError as e:
# Give LLM a chance to correct itself
messages.append({"role": "user", "content": f"Parse error: {e}. Please try again."})
continue
# Execute tool
if action not in self.tools:
observation = f"Error: tool '{action}' not found. Available: {list(self.tools.keys())}"
else:
try:
observation = self.tools[action].execute(action_input)
except Exception as e:
observation = f"Tool error: {str(e)}"
trajectory.append(Step(thought=thought, action=action,
action_input=action_input, observation=str(observation)))
# Append observation to messages
messages.append({"role": "user", "content": f"Observation: {observation}"})
# Exceeded max steps
return AgentResult(answer="Could not complete within max steps", trajectory=trajectory)
def parse_action(self, text: str) -> tuple:
thought_match = re.search(r"Thought:(.*?)(?:Action:|$)", text, re.DOTALL)
action_match = re.search(r"Action:\s*(\w+)", text)
input_match = re.search(r"Action Input:\s*(\{.*?\})", text, re.DOTALL)
if not action_match: raise ParseError("No Action found")
thought = thought_match.group(1).strip() if thought_match else ""
action = action_match.group(1).strip()
action_input = json.loads(input_match.group(1)) if input_match else {}
return thought, action, action_inputTool Interface
@dataclass
class Tool:
name: str
description: str
input_schema: dict # JSON schema for input validation
func: Callable
def execute(self, inputs: dict) -> Any:
# Validate inputs against schema
validate(inputs, self.input_schema)
return self.func(**inputs)
# Example tools
web_search = Tool(
name="web_search",
description="Search the internet for current information",
input_schema={"type": "object", "properties": {"query": {"type": "string"}}},
func=lambda query: search_api(query)
)
python_repl = Tool(
name="python_repl",
description="Execute Python code and return the output",
input_schema={"type": "object", "properties": {"code": {"type": "string"}}},
func=lambda code: sandbox_exec(code) # sandboxed!
)
calculator = Tool(
name="calculator",
description="Evaluate mathematical expressions",
input_schema={"type": "object", "properties": {"expression": {"type": "string"}}},
func=lambda expression: eval_math_safely(expression)
)Project 3 — RAG Pipeline
Full Architecture
┌─────────────────────────────────────────────┐
│ RAG SYSTEM │
│ │
INDEXING TIME: │ │
Documents ──▶ │ ┌──────────────────────────────────────┐ │
│ │ INGESTION PIPELINE │ │
│ │ │ │
│ │ Load → Chunk → Embed → Store │ │
│ │ │ │
│ │ Loaders: PDF, HTML, DOCX, MD, CSV │ │
│ │ Chunking: fixed-size or semantic │ │
│ │ Embedder: text-embedding-3-large │ │
│ │ Store: VectorDB + metadata store │ │
│ └──────────────────────────────────────┘ │
│ │
QUERY TIME: │ │
Question ────▶ │ ┌──────────────────────────────────────┐ │
│ │ RETRIEVAL ENGINE │ │
│ │ │ │
│ │ Query embed → ANN search → Rerank │ │
│ └─────────────────┬────────────────────┘ │
│ │ top-k chunks │
│ ┌─────────────────▼────────────────────┐ │
│ │ GENERATION ENGINE │ │
│ │ │ │
│ │ Build prompt with context + question │ │
│ │ LLM generates grounded answer │ │
│ │ Citation extraction │ │
│ └──────────────────────────────────────┘ │
└─────────────────────────────────────────────┘Chunking Strategies
FIXED-SIZE CHUNKING
Split by character count (e.g., 512 chars)
Overlap: 50-100 chars to avoid context loss at boundaries
chunk_size = 512
chunk_overlap = 50
chunks = []
for i in range(0, len(text), chunk_size - chunk_overlap):
chunks.append(text[i : i + chunk_size])
Pros: simple, predictable
Cons: splits mid-sentence, mid-concept
RECURSIVE CHARACTER SPLITTER (better)
Try to split on: \n\n, \n, ". ", " ", ""
Use largest separator that keeps chunks ≤ chunk_size
Preserves paragraph/sentence boundaries when possible
SEMANTIC CHUNKING
Embed each sentence
Split when cosine similarity between adjacent sentences drops
threshold = mean(similarities) - 1.5 * std(similarities)
Produces semantically coherent chunks
Pros: best retrieval quality; Cons: slower, requires embeddings
DOCUMENT-AWARE
Respect document structure: headings, paragraphs, tables
Keep tables together (don't split table rows)
Keep code blocks together
Requires document structure understanding (markdown, HTML parser)Retrieval Strategies
DENSE RETRIEVAL (standard)
embed(query) → cosine_sim with all chunk embeddings → top-k
SPARSE RETRIEVAL (BM25)
Term frequency + inverse document frequency scoring
Good for: exact keyword matches, rare terms
Implementation: elasticsearch, bm25s library
HYBRID RETRIEVAL (best of both)
scores_dense = cosine_sim(q_embed, chunk_embeds) # (n_chunks,)
scores_sparse = bm25_score(query_tokens, chunk_tokens) # (n_chunks,)
# Reciprocal Rank Fusion
rank_dense = argsort(argsort(-scores_dense)) + 1
rank_sparse = argsort(argsort(-scores_sparse)) + 1
rrf_score = 1/(60 + rank_dense) + 1/(60 + rank_sparse)
top_k = argsort(-rrf_score)[:k]
MULTI-QUERY RETRIEVAL
Generate 3-5 query variants using LLM:
"Write 5 different ways to ask: {original_query}"
Retrieve for each variant, deduplicate, merge results
Covers more semantic territory
HYDE (Hypothetical Document Embeddings)
Generate hypothetical answer: "Write a short passage that answers: {query}"
Embed the hypothetical answer (not the query)
Retrieve chunks similar to the hypothetical answer
Key insight: answer-to-document similarity > question-to-document similarity
CONTEXTUAL COMPRESSION
Retrieve top-20 chunks
For each chunk: ask LLM "Is this relevant to {query}? Extract relevant part."
Keep only relevant portions (reduces noise in context window)Reranking
CROSS-ENCODER RERANKER
Fast bi-encoder retrieval → slower cross-encoder refinement
Bi-encoder: embed query and chunks SEPARATELY → dot product
Cross-encoder: concatenate [query, chunk] → single classifier score
More accurate (full attention between query and chunk)
Too slow for all documents; fine for top-50 reranking
from sentence_transformers import CrossEncoder
reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
scores = reranker.predict([(query, chunk) for chunk in top_50_chunks])
top_k = sorted(zip(top_50_chunks, scores), key=lambda x: -x[1])[:k]
COLBERT (late interaction)
Embed query tokens and document tokens separately
Score = MaxSim: sum of max cosine similarities between query tokens and doc tokens
Faster than cross-encoder, better than bi-encoder
Stores token embeddings per document (more storage)Generation with Citations
def generate_answer(query: str, retrieved_chunks: List[Chunk]) -> Answer:
context_parts = []
for i, chunk in enumerate(retrieved_chunks):
context_parts.append(f"[{i+1}] Source: {chunk.source}\n{chunk.text}")
context = "\n\n".join(context_parts)
prompt = f"""Answer the question based on the provided sources.
After each claim, cite the source number like [1], [2], etc.
If the sources don't contain enough information, say so.
Sources:
{context}
Question: {query}
Answer:"""
response = llm.complete(prompt, temperature=0.0)
# Extract citations from response
citation_numbers = re.findall(r'\[(\d+)\]', response)
cited_chunks = [retrieved_chunks[int(n)-1] for n in set(citation_numbers)
if int(n) <= len(retrieved_chunks)]
return Answer(text=response, citations=cited_chunks, sources=[c.source for c in cited_chunks])Project 4 — Vector Database (HNSW Index)
HNSW Overview
HNSW (Hierarchical Navigable Small World) — the algorithm powering Pinecone, Weaviate, Chroma. Approximate nearest neighbor with logarithmic search complexity.
Data Structure
LAYERED GRAPH
Layer 0: ALL nodes (dense, short-range connections)
Layer 1: ~1/m fraction of nodes (sparser)
Layer 2: ~1/m² fraction of nodes
...
Layer L_max: ~1 node (entry point)
Each node:
- vector: float32[d] ← the embedding
- neighbors[layer]: List[node_id] ← up to M neighbors per layer
- layer_max: int ← max layer this node appears in
Node's max layer assigned randomly:
layer = floor(-ln(random()) * m_L) ← exponential distribution
m_L = 1/ln(M) (controls layer density)Insert Algorithm
def insert(hnsw, q_vec: np.ndarray, M=16, ef_construction=200):
"""Insert a new vector into the HNSW index."""
q_layer = random_layer() # sample from exponential distribution
q = Node(vector=q_vec, layer_max=q_layer)
# Start from entry point at top layer
ep = hnsw.entry_point
# Phase 1: Greedy descent to q_layer+1 (1 neighbor per level)
for lc in range(hnsw.L_max, q_layer + 1, -1):
ep = greedy_search_1(hnsw, q_vec, ep, layer=lc)
# Just move greedily to closest node, don't add connections
# Phase 2: Build connections from L_max down to layer 0
for lc in range(min(q_layer, hnsw.L_max), -1, -1):
# Beam search: find ef_construction candidates at this layer
candidates = beam_search(hnsw, q_vec, ep, ef=ef_construction, layer=lc)
# Select M best neighbors using heuristic
neighbors = select_neighbors_heuristic(q_vec, candidates, M)
# Add bidirectional connections
q.neighbors[lc] = neighbors
for nb in neighbors:
nb.neighbors[lc].append(q)
# Prune if degree exceeds M_max (keep M_max closest)
if len(nb.neighbors[lc]) > M_max:
nb.neighbors[lc] = select_neighbors_heuristic(
nb.vector, nb.neighbors[lc], M_max
)
ep = candidates[0] # closest found → entry for next layer
if q_layer > hnsw.L_max:
hnsw.entry_point = q
hnsw.L_max = q_layerSearch Algorithm
def search(hnsw, q_vec: np.ndarray, k=10, ef=100) -> List[Node]:
"""K nearest neighbor search."""
ep = hnsw.entry_point
# Phase 1: Greedy descent from top layer to layer 1
for lc in range(hnsw.L_max, 0, -1):
ep = greedy_search_1(hnsw, q_vec, ep, layer=lc)
# Phase 2: Beam search at layer 0 (full graph)
candidates = beam_search(hnsw, q_vec, ep, ef=max(ef, k), layer=0)
return candidates[:k]
def beam_search(hnsw, q_vec, entry_point, ef, layer):
"""Maintains a min-heap of ef candidates and max-heap of visited."""
visited = {entry_point.id}
candidates = MinHeap(key=lambda n: dist(q_vec, n.vector))
W = MaxHeap(key=lambda n: dist(q_vec, n.vector)) # result set
candidates.push(entry_point)
W.push(entry_point)
while candidates:
c = candidates.pop() # closest unprocessed
f = W.peek() # furthest in result set
if dist(q_vec, c.vector) > dist(q_vec, f.vector):
break # all remaining candidates are further than worst in W
for nb in c.neighbors[layer]:
if nb.id not in visited:
visited.add(nb.id)
f = W.peek()
if dist(q_vec, nb.vector) < dist(q_vec, f.vector) or len(W) < ef:
candidates.push(nb)
W.push(nb)
if len(W) > ef:
W.pop() # remove furthest
return sorted(W, key=lambda n: dist(q_vec, n.vector))Performance Characteristics
BUILD time: O(N log N) — each insertion is O(log N) with beam search
Query time: O(log N) amortized — due to hierarchical structure
Memory: O(N × M × L) — N nodes, M neighbors, L layers
Typical parameters:
M = 16 → each node has ~16 bidirectional connections
ef_construction = 200 → quality/speed tradeoff during build
ef (search) = 100 → quality/speed tradeoff during search
Benchmarks (SIFT-1M dataset, 128-dim):
Recall@10: 0.97+ with ef=100
QPS: ~1000-5000 on single core (vs ~5 for brute force at recall 0.97)
Disk layout (for persistence):
vectors.bin: flat float32 array (N, d)
graph.bin: adjacency lists per node per layer
metadata.json: params + entry_point + L_maxProject 5 — Function Calling Router
Architecture
INPUT: User message + available function schemas
OUTPUT: Which function to call with what arguments
APPROACHES:
1. Native function calling (OpenAI/Anthropic API format)
2. Prompt-based routing (for any LLM)
3. Semantic similarity routing (fast, no LLM needed)
4. Hybrid: semantic shortlist → LLM selectionImplementation
class FunctionCallingRouter:
def __init__(self, functions: List[FunctionDef], llm, embedder=None):
self.functions = {f.name: f for f in functions}
self.llm = llm
self.embedder = embedder
# Pre-compute function description embeddings for fast routing
if embedder:
self.func_embeddings = {
f.name: embedder.embed(f.description)
for f in functions
}
def route(self, user_message: str, top_k_candidates=3) -> FunctionCall:
# Fast pre-filter using semantic similarity
if self.embedder:
query_embed = self.embedder.embed(user_message)
candidates = self._semantic_shortlist(query_embed, top_k_candidates)
else:
candidates = list(self.functions.values())
# LLM-based selection from candidates
return self._llm_select(user_message, candidates)
def _semantic_shortlist(self, q_embed, k) -> List[FunctionDef]:
scores = {
name: cosine_similarity(q_embed, emb)
for name, emb in self.func_embeddings.items()
}
top_names = sorted(scores, key=scores.get, reverse=True)[:k]
return [self.functions[n] for n in top_names]
def _llm_select(self, message: str, candidates: List[FunctionDef]) -> FunctionCall:
# Build function schemas for the prompt
schemas = [f.to_json_schema() for f in candidates]
# Use LLM's native function calling
response = self.llm.chat(
messages=[{"role": "user", "content": message}],
tools=schemas,
tool_choice="auto"
)
if response.tool_calls:
tool_call = response.tool_calls[0]
return FunctionCall(
name=tool_call.function.name,
arguments=json.loads(tool_call.function.arguments)
)
return FunctionCall(name=None, arguments={}) # no function neededFunction Schema Format
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get current weather for a city",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The city name"
},
"units": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
"description": "Temperature units"
}
},
"required": ["city"]
}
}
}Project 6 — Structured Output Parser (Context Free Grammars)
The Problem
LLM output: "The answer is 42. Wait, let me recalculate... I think it's 43."
Expected JSON: {"answer": 42}
LLMs often:
- Add preamble ("Sure! Here's the JSON you requested:")
- Add explanation after JSON
- Use wrong quote types or trailing commas
- Nest incorrectly
SOLUTION: Grammar-constrained decoding
At each sampling step, MASK OUT logits for tokens that would
produce invalid output according to the grammar.
The model CAN ONLY generate valid structured output.Grammar-Constrained Decoding
JSON GRAMMAR (simplified BNF):
value ::= object | array | string | number | "true" | "false" | "null"
object ::= "{" (pair ("," pair)*)? "}"
pair ::= string ":" value
array ::= "[" (value ("," value)*)? "]"
string ::= '"' chars '"'
number ::= '-'? digits ('.' digits)?
PARSER STATE MACHINE:
Track current position in grammar derivation
At each step: compute VALID NEXT TOKENS
Mask logits: invalid_tokens → -∞
IMPLEMENTATION (llama.cpp grammar approach):
Grammar compiled to stack-based parser
Parser state = grammar stack + parse position
At each token generation step:
valid_tokens = grammar_parser.valid_next_tokens() → bitset
logits[~valid_tokens] = -inf
next_token = sample(softmax(logits))
grammar_parser.advance(next_token)Pydantic-Based Schema Extraction
from pydantic import BaseModel
from typing import List, Optional
class ExtractedPerson(BaseModel):
name: str
age: Optional[int]
occupation: str
location: str
class ExtractionResult(BaseModel):
people: List[ExtractedPerson]
confidence: float
class StructuredOutputParser:
def __init__(self, llm, schema: Type[BaseModel]):
self.llm = llm
self.schema = schema
self.json_schema = schema.model_json_schema()
def parse(self, text: str) -> BaseModel:
prompt = f"""Extract information from the text and return ONLY a JSON object
matching this schema:
{json.dumps(self.json_schema, indent=2)}
Text: {text}
JSON (no other text):"""
# Try with grammar constraint first, fall back to retry
for attempt in range(3):
response = self.llm.complete(
prompt,
temperature=0.0,
response_format={"type": "json_object"} # if API supports it
)
try:
# Clean response (remove markdown fences if present)
clean = re.sub(r'```json\n?|```\n?', '', response).strip()
data = json.loads(clean)
return self.schema(**data)
except (json.JSONDecodeError, ValidationError) as e:
if attempt == 2: raise
# Retry with error feedback
prompt += f"\n\nPrevious attempt failed with error: {e}\nTry again:"Project 7 — Semantic Router
Architecture
SEMANTIC ROUTER: Route queries to different handlers based on meaning
Without keyword matching or rule lists
EXAMPLES:
"What is machine learning?" → knowledge_handler
"Write me a poem" → creative_handler
"Book me a flight" → action_handler
"I'm feeling sad" → empathy_handler
COMPONENTS:
1. Route definitions (name + example utterances)
2. Encoder (converts utterances to vectors)
3. Index (fast similarity lookup)
4. Threshold logic (confidence scoring)Implementation
class SemanticRouter:
def __init__(self, routes: List[Route], encoder: Encoder, threshold=0.75):
self.routes = routes
self.encoder = encoder
self.threshold = threshold
self.index = self._build_index()
def _build_index(self):
"""Embed all example utterances per route."""
index = {}
for route in self.routes:
embeddings = self.encoder.batch_encode(route.utterances)
index[route.name] = embeddings
return index
def route(self, query: str) -> RouteDecision:
q_embed = self.encoder.encode(query)
# Score each route: max similarity to any utterance
route_scores = {}
for route_name, embeddings in self.index.items():
sims = cosine_similarity(q_embed, embeddings) # (n_utterances,)
route_scores[route_name] = float(sims.max())
best_route = max(route_scores, key=route_scores.get)
best_score = route_scores[best_route]
if best_score >= self.threshold:
return RouteDecision(route=best_route, score=best_score)
else:
return RouteDecision(route="fallback", score=best_score)
# Route definitions
routes = [
Route(
name="weather",
utterances=[
"What's the weather like?",
"Is it going to rain tomorrow?",
"Temperature in London",
"Should I bring a coat?",
"Weather forecast for next week",
]
),
Route(
name="code_help",
utterances=[
"Help me debug this code",
"Why is my function returning None?",
"How do I implement a binary search?",
"Explain this Python error",
]
),
]Project 8 — Graph RAG System
Why Graph RAG?
STANDARD RAG LIMITATION:
Query: "How did the acquisition of Company A affect Company B's market share
given that Company A was previously competing with Company C?"
Chunk retrieval: finds chunks about acquisition OR market share OR competition
But: multi-hop reasoning requires connecting MULTIPLE entities and relationships
Standard RAG: retrieves individually relevant chunks, misses connections
GRAPH RAG:
Extracts entities and relationships → builds knowledge graph
Query: traverse graph to find connected information
Answer: grounded in graph traversal + relevant chunksFull Pipeline
INDEXING PIPELINE
───────────────────
Documents
│
▼
┌─────────────────────────────────────┐
│ ENTITY + RELATION EXTRACTION │
│ │
│ Prompt LLM: │
│ "Extract entities and relations │
│ from this text as JSON: │
│ {entities: [{name, type, desc}], │
│ relations: [{from, rel, to}]}" │
└──────────────────┬──────────────────┘
│
▼
┌─────────────────────────────────────┐
│ ENTITY RESOLUTION │
│ │
│ "OpenAI" == "Open AI" == "OpenAI, │
│ Inc." → merge to canonical form │
│ │
│ Approach: │
│ 1. Exact match after normalization│
│ 2. Embedding similarity (>0.95) │
│ 3. LLM-based disambiguation │
└──────────────────┬──────────────────┘
│
▼
┌─────────────────────────────────────┐
│ COMMUNITY DETECTION │
│ │
│ Run Leiden/Louvain algorithm │
│ Groups tightly connected entities │
│ Generate community summaries │
│ (LLM summarizes each community) │
└──────────────────┬──────────────────┘
│
▼
┌─────────────────────────────────────┐
│ GRAPH STORAGE │
│ Neo4j / NetworkX / LightRAG │
│ Nodes: entities + properties │
│ Edges: typed relationships │
└─────────────────────────────────────┘
QUERY PIPELINE
───────────────
User Query
│
├──▶ Entity extraction from query
│ ["Company A", "Company B", "market share"]
│
├──▶ Local search: retrieve subgraph around query entities
│ 1-hop: direct neighbors
│ 2-hop: neighbors of neighbors
│ Prune by edge type relevance
│
├──▶ Global search: retrieve relevant community summaries
│ Embed query → find similar community summaries
│
├──▶ Vector search: find relevant chunks (standard RAG)
│
▼
CONTEXT ASSEMBLY
Subgraph (entities + relations as text)
Community summaries
Relevant chunks
│
▼
LLM GENERATION with grounded context
Citations link back to source chunks + graph pathsProject 9 — Code Interpreter Sandbox
Security Architecture
THREAT MODEL
Untrusted code from LLM or user executes on your server
Risks: file system access, network calls, fork bombs,
CPU exhaustion, memory exhaustion, container escape
ISOLATION LAYERS
Layer 1: Linux namespaces (separate pid, net, mnt, user)
Layer 2: seccomp syscall filtering (whitelist safe syscalls)
Layer 3: Resource limits (cgroups: CPU, RAM, disk)
Layer 4: Read-only filesystem (except /tmp)
Layer 5: No network access (isolated network namespace)
Layer 6: Timeout enforcement (SIGKILL after N seconds)Implementation (Python subprocess + seccomp)
import subprocess, resource, tempfile, os
class CodeSandbox:
def __init__(self, timeout_seconds=10, memory_limit_mb=256):
self.timeout = timeout_seconds
self.memory_limit = memory_limit_mb * 1024 * 1024
def execute(self, code: str, language="python") -> ExecutionResult:
with tempfile.TemporaryDirectory() as tmpdir:
# Write code to temp file
code_file = os.path.join(tmpdir, f"solution.{language[:2]}")
with open(code_file, 'w') as f:
f.write(code)
# Build sandbox command using bubblewrap (bwrap)
cmd = [
"bwrap",
"--ro-bind", "/usr", "/usr",
"--ro-bind", "/lib", "/lib",
"--ro-bind", "/lib64", "/lib64",
"--bind", tmpdir, "/sandbox",
"--tmpfs", "/tmp",
"--proc", "/proc",
"--dev", "/dev",
"--unshare-all", # all namespaces
"--die-with-parent",
"--",
"python3", "-u", "/sandbox/solution.py"
]
try:
result = subprocess.run(
cmd,
capture_output=True,
text=True,
timeout=self.timeout,
preexec_fn=self._set_limits
)
return ExecutionResult(
stdout=result.stdout,
stderr=result.stderr,
returncode=result.returncode,
timed_out=False
)
except subprocess.TimeoutExpired:
return ExecutionResult(stdout="", stderr="",
returncode=-1, timed_out=True)
def _set_limits(self):
"""Called in child process before exec — sets resource limits."""
resource.setrlimit(resource.RLIMIT_AS, (self.memory_limit, self.memory_limit))
resource.setrlimit(resource.RLIMIT_CPU, (self.timeout, self.timeout))
resource.setrlimit(resource.RLIMIT_NPROC, (50, 50))
resource.setrlimit(resource.RLIMIT_FSIZE, (10_000_000, 10_000_000)) # 10MBHow These 10 Projects Form One System
┌──────────────────────────────────────────────┐
│ MINI LANGCHAIN │
│ │
User Input ────▶ │ Semantic Router ──▶ route classification │
│ │ │
│ ┌────┴──────────────────────────────┐ │
│ │ │ │
│ RAG Path Agent Path │
│ │ │ │
│ Embed Query CoT Reasoning │
│ HNSW Search ReAct Loop │
│ Rerank Tool Calls: │
│ Graph RAG - Code Sandbox│
│ Context Build - Function Router│
│ │ - Web Search │
│ └────────────┬──────────────────────┘ │
│ │ │
│ Structured Output Parser │
│ (Grammar-constrained JSON) │
│ │ │
│ Final Response │
└──────────────────────────────────────────────┘
Data flow for complex query:
1. Semantic Router classifies intent
2. If knowledge query → RAG pipeline + Graph RAG
3. If action query → ReAct agent loop
4. CoT for multi-step reasoning within agent
5. Function calling for external tool invocation
6. Code sandbox for computation
7. Structured output parser for typed responseMini Hugging Face — E2E Architecture Reference
Projects: Tokenizer (BPE) · LoRA Trainer · PEFT Library · Eval Harness · Model Merger (SLERP) · Embedding Model · Logit Processor · Guardrails System · Prompt Caching · Whisper-style ASR · Text-to-Speech Pipeline · Audio Spectrogram Transformer
What You're Building
The model management and adaptation stack: how raw text becomes tokens, how models are fine-tuned efficiently without touching most parameters, how models are evaluated rigorously, and how they're configured for deployment. This is the layer Hugging Face made famous — transformers, tokenizers, PEFT, evaluate — built from scratch.
Project 1 — Tokenizer (BPE Implementation)
What BPE Does
RAW TEXT: "low lower newest widest"
VOCABULARY: starts with individual bytes/characters
GOAL: iteratively merge most frequent adjacent pairs
RESULT: vocabulary of subword tokens that compresses text efficiently
WHY SUBWORDS:
- "unbelievable" = "un" + "believ" + "able" (not 1 unknown word)
- Handles morphology, typos, new words
- Fixed vocabulary size with near-complete text coverageBPE Training Algorithm
def train_bpe(text: str, vocab_size: int) -> Tokenizer:
"""Train BPE tokenizer from scratch."""
# STEP 1: Initial vocabulary = all bytes (256 base tokens)
vocab = {i: bytes([i]) for i in range(256)}
merges = {} # (pair) → merged_token_id
# Pretokenize: split on whitespace, add word-boundary marker
# "low" → ["l", "o", "w", </w>] (</w> marks word end)
words = Counter(text.split())
word_tokens = {
word: list(word.encode()) + [WORD_END_TOKEN]
for word in words
}
word_freqs = dict(words)
# STEP 2: Iteratively merge until vocab_size reached
while len(vocab) < vocab_size:
# Count all adjacent pairs across corpus
pair_freqs = defaultdict(int)
for word, tokens in word_tokens.items():
freq = word_freqs[word]
for i in range(len(tokens) - 1):
pair_freqs[(tokens[i], tokens[i+1])] += freq
if not pair_freqs:
break
# Find most frequent pair
best_pair = max(pair_freqs, key=pair_freqs.get)
# Create new token for merge
new_token_id = len(vocab)
new_token_bytes = vocab[best_pair[0]] + vocab[best_pair[1]]
vocab[new_token_id] = new_token_bytes
merges[best_pair] = new_token_id
# Apply merge to all words
for word in word_tokens:
word_tokens[word] = apply_merge(word_tokens[word], best_pair, new_token_id)
return Tokenizer(vocab=vocab, merges=merges)
def apply_merge(tokens: List[int], pair: tuple, new_id: int) -> List[int]:
"""Replace all occurrences of pair in tokens with new_id."""
result = []
i = 0
while i < len(tokens):
if i < len(tokens)-1 and (tokens[i], tokens[i+1]) == pair:
result.append(new_id)
i += 2
else:
result.append(tokens[i])
i += 1
return resultBPE Encoding (Inference)
def encode(self, text: str) -> List[int]:
"""Encode text to token IDs using learned merges."""
# Apply pretokenization (same as training)
words = self.pretokenize(text)
all_tokens = []
for word in words:
# Start with byte-level tokens
tokens = list(word.encode())
# Apply merges in order of training (greedy, left-to-right)
while True:
# Find applicable merge with lowest training index
best = None
best_rank = float('inf')
for i in range(len(tokens) - 1):
pair = (tokens[i], tokens[i+1])
if pair in self.merges:
rank = self.merge_rank[pair] # training order
if rank < best_rank:
best_rank = rank
best = (i, pair)
if best is None:
break
i, pair = best
new_id = self.merges[pair]
tokens = tokens[:i] + [new_id] + tokens[i+2:]
all_tokens.extend(tokens)
return all_tokens
def decode(self, token_ids: List[int]) -> str:
"""Decode token IDs back to text."""
bytes_output = b"".join(self.vocab[id] for id in token_ids)
return bytes_output.decode('utf-8', errors='replace')Special Tokens & Chat Templates
SPECIAL TOKENS (must be added after BPE training):
<|endoftext|> ← EOS/separator
<|pad|> ← padding
<|unk|> ← unknown (rarely used in BPE)
[INST], [/INST] ← instruction markers (Llama)
<|im_start|>, <|im_end|> ← ChatML format
CHAT TEMPLATE (converts messages to flat string):
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hi"},
{"role": "assistant", "content": "Hello!"}
]
ChatML format:
<|im_start|>system
You are helpful.<|im_end|>
<|im_start|>user
Hi<|im_end|>
<|im_start|>assistant
Hello!<|im_end|>
<|im_start|>assistant
[model continues here]Project 2 — LoRA Trainer (Low-Rank Adaptation)
The Core Idea
FULL FINE-TUNING:
Update all W (d×d) weights — billions of parameters
Memory: 2× model size (model + gradients) + optimizer states (Adam: 3× params)
For 7B model: 7B × 4 bytes × (1 + 1 + 3) = ~140 GB minimum
LORA INSIGHT:
During fine-tuning, weight updates ΔW have low intrinsic rank
Instead of learning ΔW (d×d), learn: ΔW = A × B
where A is (d×r), B is (r×d), r << d
Full forward pass uses: W + AB
Only A and B are trained (r << d → 1000x fewer params)
NUMBERS:
d=4096, r=16
LoRA params per layer: 2 × 4096 × 16 = 131,072
Full params per layer: 4096 × 4096 = 16,777,216
Reduction: 128x per layerImplementation
import torch
import torch.nn as nn
class LoRALinear(nn.Module):
def __init__(self, base_layer: nn.Linear, r: int = 16, alpha: float = 32,
dropout: float = 0.1):
super().__init__()
self.base_layer = base_layer
self.r = r
self.alpha = alpha
self.scaling = alpha / r # scale factor (merged into B at init)
d_in = base_layer.in_features
d_out = base_layer.out_features
# LoRA matrices
self.lora_A = nn.Parameter(torch.randn(r, d_in) * 0.02) # init: small random
self.lora_B = nn.Parameter(torch.zeros(d_out, r)) # init: ZERO → ΔW=0 at start
self.dropout = nn.Dropout(dropout)
# Freeze base layer
for param in self.base_layer.parameters():
param.requires_grad = False
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Original: y = xW^T + b
base_output = self.base_layer(x)
# LoRA: y += x(AB)^T * scaling
lora_output = self.dropout(x) @ self.lora_A.T @ self.lora_B.T
return base_output + lora_output * self.scaling
def merge_weights(self):
"""Merge LoRA into base weights for inference (no overhead)."""
delta_W = (self.lora_B @ self.lora_A) * self.scaling
self.base_layer.weight.data += delta_W
self.merged = TrueLoRA Application to Transformer
def apply_lora(model: nn.Module, lora_config: LoRAConfig) -> nn.Module:
"""Replace target modules with LoRA versions."""
for name, module in model.named_modules():
# Apply LoRA to attention projections: Q, K, V, O
if isinstance(module, nn.Linear) and any(
target in name for target in lora_config.target_modules
):
parent = get_parent_module(model, name)
attr = name.split('.')[-1]
lora_layer = LoRALinear(
base_layer=module,
r=lora_config.r,
alpha=lora_config.alpha,
dropout=lora_config.dropout
)
setattr(parent, attr, lora_layer)
return model
# Which modules to apply LoRA to:
# Typical: ["q_proj", "v_proj"] (Llama) or ["query", "value"] (BERT)
# Aggressive: all attention + FFN layers
# Conservative: query and value only
# Count trainable parameters
def count_params(model):
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"Trainable: {trainable:,} ({100*trainable/total:.2f}%)")Training Loop
def train_lora(base_model, dataset, lora_config, training_args):
# Apply LoRA
model = apply_lora(base_model, lora_config)
# Only LoRA params in optimizer
optimizer = torch.optim.AdamW(
filter(lambda p: p.requires_grad, model.parameters()),
lr=training_args.learning_rate, # typically 2e-4 (higher than full FT)
weight_decay=0.01
)
# Gradient checkpointing + mixed precision (fits on single GPU)
model.gradient_checkpointing_enable()
scaler = torch.cuda.amp.GradScaler()
for step, batch in enumerate(dataloader):
with torch.cuda.amp.autocast():
outputs = model(**batch)
loss = outputs.loss
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
# Save only LoRA weights (tiny: ~100MB for 7B model)
save_lora_weights(model, "lora_checkpoint/")
# Merge for deployment
merged_model = merge_lora_weights(model)
merged_model.save_pretrained("merged_model/")Project 3 — PEFT Library (Parameter Efficient Fine-Tuning)
Taxonomy of PEFT Methods
PEFT METHODS
│
┌───────────────┼─────────────────┐
│ │ │
ADAPTERS SOFT PROMPTS SELECTIVE
│ │ FINE-TUNING
- LoRA - Prompt Tuning - BitFit
- Adapter layers - Prefix Tuning - Diff Pruning
- IA³ - P-Tuning v2 - Sparse FT
- AdaLoRA
- QLoRA
WHEN TO USE WHICH:
LoRA: Best all-around. Fast, effective, widely supported.
QLoRA: LoRA on quantized model. Fits 70B on single GPU.
Prefix Tuning: When you can't modify model architecture.
Prompt Tuning: Smallest overhead. Works for large models (>10B).
IA³: Extremely few params (3 vectors per layer). For massive models.Adapter Layers
class AdapterLayer(nn.Module):
"""Houlsby adapter: down-project, nonlinearity, up-project."""
def __init__(self, d_model: int, adapter_size: int = 64):
super().__init__()
self.down = nn.Linear(d_model, adapter_size)
self.act = nn.GELU()
self.up = nn.Linear(adapter_size, d_model)
# Initialize near identity (small init → minimal disruption)
nn.init.normal_(self.down.weight, std=1e-3)
nn.init.normal_(self.up.weight, std=1e-3)
nn.init.zeros_(self.down.bias)
nn.init.zeros_(self.up.bias)
def forward(self, x):
return x + self.up(self.act(self.down(x))) # residual
# Insert adapter after attention and after FFN in each transformer blockPrefix Tuning
class PrefixTuning(nn.Module):
"""Prepend learned soft tokens to K and V in every attention layer."""
def __init__(self, n_layers, n_heads, d_head, prefix_length=20):
super().__init__()
self.prefix_length = prefix_length
# Learnable prefix embeddings (reparameterized through MLP for stability)
self.prefix_embeddings = nn.Parameter(
torch.randn(prefix_length, n_layers * 2 * n_heads * d_head)
)
self.prefix_mlp = nn.Sequential(
nn.Linear(n_layers * 2 * n_heads * d_head, 512),
nn.Tanh(),
nn.Linear(512, n_layers * 2 * n_heads * d_head)
)
def get_prefix_kv(self, batch_size):
# Transform prefix embeddings through MLP
prefix = self.prefix_mlp(self.prefix_embeddings)
# Reshape to (n_layers, 2, batch, n_heads, prefix_len, d_head)
prefix = prefix.view(self.prefix_length, n_layers, 2, n_heads, d_head)
prefix = prefix.permute(1, 2, ...).expand(batch_size, ...)
# Split into key and value prefixes per layer
return prefix[:, 0], prefix[:, 1] # (n_layers, batch, n_heads, prefix_len, d_head)QLoRA
QLoRA = LoRA applied to a 4-bit quantized model
COMPONENTS:
1. NF4 quantization of base model (4-bit, see Quantization section)
2. Double quantization of scale factors (saves ~0.5 bits/param)
3. Paged optimizers (gradient checkpoints page to CPU RAM)
4. LoRA adapters trained in BF16 on top of frozen 4-bit weights
FORWARD PASS:
input (BF16) × W_nf4 (dequantized to BF16 just-in-time) + LoRA (BF16)
The NF4 base weights are NEVER updated
Only LoRA A and B are trained (BF16)
MEMORY CALCULATION (7B model, r=64):
Base model (NF4): 7B × 0.5 bytes = 3.5 GB
Scale factors (FP8): 7B/64 × 1 byte = 109 MB
LoRA params (BF16): ~160M × 2 bytes = 320 MB
Optimizer states: 320M × 8 bytes = 2.5 GB (Adam: m + v)
Activations (GC): ~1 GB
Total: ~7.5 GB → fits on single RTX 3090 (24 GB)Project 4 — LLM Eval Harness
Architecture
EVAL HARNESS PIPELINE
─────────────────────
┌──────────────────────────────────────────────────────────┐
│ EVAL HARNESS │
│ │
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────────┐ │
│ │ Dataset │ │ Model │ │ Metrics │ │
│ │ Loader │ │ Adapter │ │ Registry │ │
│ │ │ │ │ │ │ │
│ │ - MMLU │ │ - HF Model │ │ - Accuracy │ │
│ │ - GSM8K │ │ - OpenAI │ │ - Exact Match │ │
│ │ - HumanEval │ │ - Anthropic │ │ - BLEU/ROUGE │ │
│ │ - TruthfulQA│ │ - vLLM │ │ - Pass@k │ │
│ │ - HellaSwag │ │ - Local │ │ - Perplexity │ │
│ └──────┬──────┘ └──────┬──────┘ └────────┬────────┘ │
│ │ │ │ │
│ └────────────────┼───────────────────┘ │
│ │ │
│ ┌──────────▼──────────┐ │
│ │ TASK RUNNER │ │
│ │ │ │
│ │ For each sample: │ │
│ │ 1. Build prompt │ │
│ │ 2. Get model output│ │
│ │ 3. Score output │ │
│ │ 4. Aggregate │ │
│ └──────────┬──────────┘ │
│ │ │
│ ┌──────────▼──────────┐ │
│ │ RESULTS REPORTER │ │
│ │ JSON / leaderboard │ │
│ └─────────────────────┘ │
└──────────────────────────────────────────────────────────┘Task Types
1. MULTIPLE CHOICE (MMLU, HellaSwag, ARC)
Evaluation approach: log-likelihood scoring
For each answer choice A, B, C, D:
score = log P(choice | question)
= sum of log P(token_i | question + preceding_choice_tokens)
Predict: argmax(scores)
Why log-likelihood instead of generation:
Deterministic, no sampling variance
Faster (no autoregressive generation)
Normalized by token count (avoid length bias)
2. OPEN-ENDED GENERATION (GSM8K, TruthfulQA, NaturalQuestions)
Generate answer, compare to reference:
- Exact match: normalize (lowercase, strip punctuation) → compare
- F1 token overlap: for QA tasks
- BLEU/ROUGE: for generation quality
GSM8K: extract final numerical answer from chain-of-thought
Pattern: r'\d+\.?\d*$' on last line
3. CODE GENERATION (HumanEval, MBPP)
pass@k metric:
Generate k completions per problem
Problem is solved if ANY completion passes all unit tests
pass@k = 1 - C(n-c, k) / C(n, k)
where: n = total generations, c = correct generations
Estimate with n=20, compute pass@1, pass@5, pass@10
Execution:
Sandbox code + unit tests → run → check all tests pass
4. PERPLEXITY (language modeling quality)
PPL = exp(-1/N × Σ log P(token_i | context))
Lower = better
Rolling window for long documents:
stride = 512, context = 1024
Score only the new tokens (last stride tokens)
Average across documentImplementation
class EvalHarness:
def __init__(self, model: ModelAdapter, tasks: List[Task]):
self.model = model
self.tasks = tasks
def evaluate(self) -> EvalResults:
results = {}
for task in self.tasks:
dataset = task.load_dataset()
scores = []
for batch in DataLoader(dataset, batch_size=task.batch_size):
if task.eval_type == "multiple_choice":
batch_scores = self._eval_mc(task, batch)
elif task.eval_type == "generation":
batch_scores = self._eval_generation(task, batch)
elif task.eval_type == "code":
batch_scores = self._eval_code(task, batch)
scores.extend(batch_scores)
results[task.name] = {
"score": np.mean(scores),
"stderr": np.std(scores) / np.sqrt(len(scores)),
"n_samples": len(scores)
}
return EvalResults(results)
def _eval_mc(self, task, batch) -> List[float]:
"""Multiple choice via log-likelihood scoring."""
correct = []
for item in batch:
ctx = task.build_context(item)
choices = item["choices"]
gold = item["gold"]
# Get log-likelihood for each choice
lls = []
for choice in choices:
ll = self.model.loglikelihood(ctx, choice)
lls.append(ll)
# Normalize by continuation length (per-token)
normalized = [ll / len(enc) for ll, enc in zip(lls, encoded_choices)]
pred = np.argmax(normalized)
correct.append(int(pred == gold))
return correctProject 5 — Model Merger (SLERP / Model Soups)
Why Merge Models?
Scenario 1: Task arithmetic
Base model + code model + math model → merge → strong multi-task model
No additional training required
Scenario 2: Ensemble on budget
Average N fine-tuned models → better than any single model
Only 1 model at inference (averaged weights)
Scenario 3: Capability interpolation
Model A: strong at reasoning but verbose
Model B: concise but weaker reasoning
Merge at λ=0.5 → balanced modelModel Soup (Simple Averaging)
def model_soup(model_paths: List[str], weights: List[float] = None) -> Model:
"""Uniform or weighted average of model checkpoints."""
if weights is None:
weights = [1.0 / len(model_paths)] * len(model_paths)
models = [load_model(p) for p in model_paths]
base = models[0]
with torch.no_grad():
for name, param in base.named_parameters():
param.data = sum(
w * m.state_dict()[name].float()
for m, w in zip(models, weights)
)
return base
# Requirements for averaging to work:
# - Models must share SAME architecture (identical layer shapes)
# - Models must be in SAME loss basin (fine-tuned from same base)
# - Averaging interpolates between their loss basinsSLERP (Spherical Linear Interpolation)
WHY SLERP OVER LERP:
Weights lie on high-dimensional sphere (approximately)
Linear interpolation (LERP): shorter path through interior
Spherical interpolation (SLERP): follows arc on surface
SLERP maintains "norm" of weights better → better preserved features
FORMULA:
SLERP(w₁, w₂, t) = sin((1-t)θ)/sin(θ) × w₁ + sin(tθ)/sin(θ) × w₂
where θ = arccos(w₁·w₂ / (|w₁||w₂|))
t ∈ [0,1]: t=0 → w₁, t=1 → w₂
IMPLEMENTATION:
```python
def slerp(w1: torch.Tensor, w2: torch.Tensor, t: float) -> torch.Tensor:
"""Spherical linear interpolation between two weight tensors."""
w1_flat = w1.float().view(-1)
w2_flat = w2.float().view(-1)
# Normalize
w1_norm = F.normalize(w1_flat, dim=0)
w2_norm = F.normalize(w2_flat, dim=0)
# Angle between
dot = torch.clamp(torch.dot(w1_norm, w2_norm), -1.0, 1.0)
theta = torch.acos(dot)
if theta.abs() < 1e-6:
# Nearly parallel — fall back to linear interpolation
return (1-t) * w1 + t * w2
# SLERP formula
sin_theta = torch.sin(theta)
w1_coeff = torch.sin((1-t) * theta) / sin_theta
w2_coeff = torch.sin(t * theta) / sin_theta
result = w1_coeff * w1_flat + w2_coeff * w2_flat
# Rescale to original norm (interpolated between norms)
target_norm = (1-t) * w1_flat.norm() + t * w2_flat.norm()
result = F.normalize(result) * target_norm
return result.view(w1.shape).to(w1.dtype)DARE + TIES Merging (Advanced)
TIES (Trim, Elect, Disjoint Merge):
Problem: averaging causes interference from conflicting task vectors
Step 1: TRIM
Compute task vectors: τᵢ = θᵢ - θ_base
Keep only top-k% of values by magnitude, zero rest
Rationale: small deltas → noise, large deltas → meaningful updates
Step 2: ELECT
For each parameter: count sign agreement across models
Elect majority sign (resolve conflicts)
Step 3: MERGE
Average only parameters where sign agrees with elected sign
θ_merged = θ_base + mean(τᵢ where sign(τᵢ) == elected_sign)
DARE (Drop and REscale):
Randomly drop task vector parameters (p=0.9 typical)
Rescale remaining by 1/(1-p) to preserve magnitude
Reduces interference between task vectors
Works well when p is high (90% dropout)Project 6 — Embedding Model
Architecture
TEXT → EMBEDDING VECTOR (fixed-size, semantic representation)
BACKBONE: BERT-style encoder (bidirectional transformer)
Input: [CLS] token1 token2 ... tokenN [SEP]
Process: full bidirectional attention (all tokens see all)
Output: hidden states for all positions
POOLING STRATEGIES:
CLS pooling: use hidden state of [CLS] token
Mean pooling: average all token hidden states (usually better)
Max pooling: element-wise max across tokens
Mean pooling (weighted by attention mask):
token_embeddings = model_output.last_hidden_state # (B, T, D)
attention_mask = inputs['attention_mask'] # (B, T)
mask_expanded = attention_mask.unsqueeze(-1).expand_as(token_embeddings)
sum_embeddings = (token_embeddings * mask_expanded).sum(dim=1)
sum_mask = mask_expanded.sum(dim=1).clamp(min=1e-9)
embeddings = sum_embeddings / sum_mask # (B, D)
Normalize: F.normalize(embeddings, p=2, dim=-1) ← unit sphereTraining: Contrastive Learning (SimCSE / E5)
CONTRASTIVE LOSS (InfoNCE / NT-Xent):
Given: anchor sentence a, positive p, negatives n₁...nₖ
Goal: embed(a) close to embed(p), far from embed(nᵢ)
loss = -log [exp(sim(a,p)/τ) / (exp(sim(a,p)/τ) + Σ exp(sim(a,nᵢ)/τ))]
τ = temperature (0.05 typical)
sim = cosine similarity
DATA SOURCES:
Positive pairs:
- (question, answer) from QA datasets
- (premise, entailed_hypothesis) from NLI
- (query, relevant_doc) from MS-MARCO
- (sentence, paraphrase)
- (sentence, dropout-augmented same sentence) ← SimCSE unsupervised
Negative pairs:
- Random from batch (in-batch negatives, efficient)
- Hard negatives: similar-looking but semantically different
- Mined negatives: BM25 top results that are NOT relevant
IN-BATCH NEGATIVES:
Batch of (anchor, positive) pairs
Use all other positives in batch as negatives for each anchor
Batch size 256 → 255 negatives per sample (large effective batch)
scores = anchor_embs @ positive_embs.T # (B, B)
labels = torch.arange(B) # diagonal is positive
loss = CrossEntropyLoss(scores / temp, labels)Project 7 — Logit Processor
Logit Processor Chain
Raw logits from model
│
▼
┌───────────────────────────────────────────┐
│ LOGIT PROCESSOR CHAIN │
│ │
│ 1. TemperatureProcessor │
│ logits = logits / temperature │
│ │
│ 2. RepetitionPenaltyProcessor │
│ for tok in generated_tokens: │
│ if logits[tok] < 0: │
│ logits[tok] *= penalty │
│ else: │
│ logits[tok] /= penalty │
│ │
│ 3. TopKProcessor │
│ Remove all tokens outside top-k │
│ Set their logits to -inf │
│ │
│ 4. TopPProcessor (nucleus sampling) │
│ Sort by probability descending │
│ Find smallest set summing to > p │
│ Remove tokens outside that set │
│ │
│ 5. MinLengthProcessor │
│ If len(generated) < min_length: │
│ logits[EOS] = -inf │
│ │
│ 6. BadWordsProcessor │
│ logits[banned_token_ids] = -inf │
└───────────────────────────────────────────┘
│
▼
softmax → sampleImplementation
class LogitProcessorList:
def __init__(self, processors: List[LogitProcessor]):
self.processors = processors
def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
for processor in self.processors:
scores = processor(input_ids, scores)
return scores
class TopPLogitWarper:
def __init__(self, top_p: float, min_tokens_to_keep: int = 1):
self.top_p = top_p
self.min_tokens = min_tokens_to_keep
def __call__(self, input_ids, scores):
sorted_logits, sorted_indices = torch.sort(scores, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
# Remove tokens with cumulative prob above threshold (keep first token above)
sorted_remove = cumulative_probs - F.softmax(sorted_logits, dim=-1) > self.top_p
sorted_remove[..., :self.min_tokens] = False # keep at least min_tokens
# Scatter back to original indexing
indices_to_remove = sorted_remove.scatter(-1, sorted_indices, sorted_remove)
scores = scores.masked_fill(indices_to_remove, float('-inf'))
return scoresProject 8 — Guardrails System (Input/Output Filtering)
Architecture
GUARDRAILS PIPELINE
User Input
│
▼
┌─────────────────────────────────────────────┐
│ INPUT GUARDRAILS │
│ │
│ 1. PII Detector │
│ Regex + NER: SSN, credit cards, │
│ emails, phone numbers │
│ Action: redact or reject │
│ │
│ 2. Prompt Injection Detector │
│ Patterns: "ignore previous instructions"│
│ "you are now DAN", role-play bypasses │
│ ML classifier on input │
│ │
│ 3. Topic Filter │
│ Classifier: allowed vs blocked topics │
│ (violence, NSFW, medical, legal) │
│ │
│ 4. Language Detector │
│ Reject if not in supported_languages │
└─────────────────┬───────────────────────────┘
│ cleaned input
▼
LLM Generation
│
▼
┌─────────────────────────────────────────────┐
│ OUTPUT GUARDRAILS │
│ │
│ 1. Toxicity Classifier │
│ Score: hate speech, harassment, etc. │
│ Threshold-based: score > 0.8 → reject │
│ │
│ 2. Hallucination Detector │
│ Ground output against retrieved context │
│ NLI: does output contradict sources? │
│ │
│ 3. PII in Output │
│ Same as input PII detection │
│ Redact before returning to user │
│ │
│ 4. Brand Safety │
│ Custom word lists, competitor mentions │
└─────────────────┬───────────────────────────┘
│ safe output
▼
Final ResponseHow These 12 Projects Form One System
┌──────────────────────────────────────────────────┐
│ MINI HUGGING FACE │
│ │
Raw Text ──────▶ │ Tokenizer (BPE) ──▶ token_ids │
│ │ │
Model Weights ──▶ │ PEFT / LoRA / QLoRA ──▶ adapted model │
│ │ │
Adapted Model ──▶ │ Eval Harness ──▶ benchmark scores │
│ │ │
Multiple Models ▶ │ Model Merger (SLERP) ──▶ merged model │
│ │ │
│ Embedding Model ──▶ semantic vectors │
│ │ │
At Inference: │ Logit Processor ──▶ sampling control │
│ Guardrails ──▶ safety filtering │
│ Prompt Cache ──▶ KV reuse │
│ │
Audio Input ────▶ │ AST → Whisper → ASR transcript │
Text ──────────▶ │ TTS Pipeline ──▶ audio output │
└──────────────────────────────────────────────────┘
Production deployment checklist derived from this system:
✓ Tokenizer serialized and versioned alongside model
✓ PEFT adapters can be hot-swapped per request (LoRA multiplexing)
✓ Eval harness runs on every checkpoint (regression detection)
✓ Guardrails run async parallel to generation when possible
✓ Embedding model serves as backbone for retrieval and routing
✓ Multimodal: ASR converts voice input before text processingMini Research Lab — E2E Architecture Reference
Projects: RLHF Pipeline (PPO) · DPO Loss Function · Model Distillation · Synthetic Data Generator · MoE Routing Layer · Distributed Training (FSDP/Tensor Parallel) · State Space Model (Mamba) · Interpretability (SAE) · Neural Architecture Search · Data Curation Pipeline · Recommendation System (Two-Tower) · Adversarial Attack Generator · Multi-modal Projector (CLIP) · AI Gateway
What You're Building
The research infrastructure stack: how models are aligned to human preferences, how knowledge is transferred between models, how data is curated and generated, and how models are scaled and analyzed. This is the work that happens between pretraining and deployment — the layer that separates a capable model from a useful, safe, production model.
Project 1 — RLHF Pipeline (PPO Implementation)
Full RLHF Architecture
STAGE 1: SUPERVISED FINE-TUNING (SFT)
Base model → train on high-quality demonstrations
Result: SFT model (knows desired output format and style)
STAGE 2: REWARD MODEL TRAINING
SFT model → generate multiple responses per prompt
Humans rank responses: A > B, B > C, etc.
Train reward model: predict which response humans prefer
STAGE 3: PPO FINE-TUNING (RL)
Use reward model as environment
Update SFT model via PPO to maximize reward
KL penalty prevents drifting too far from SFT model
Result: RLHF-tuned model
ALL FOUR MODELS IN MEMORY DURING PPO:
1. Actor (policy model being trained)
2. Reference model (frozen SFT model, for KL)
3. Critic (value network, estimates return)
4. Reward model (frozen, provides reward signal)Reward Model Training
class RewardModel(nn.Module):
def __init__(self, base_model):
super().__init__()
self.transformer = base_model
self.reward_head = nn.Linear(base_model.config.d_model, 1)
def forward(self, input_ids, attention_mask):
outputs = self.transformer(input_ids, attention_mask=attention_mask)
last_hidden = outputs.last_hidden_state[:, -1, :] # last token
reward = self.reward_head(last_hidden).squeeze(-1)
return reward
def reward_model_loss(reward_model, chosen_ids, rejected_ids, chosen_mask, rejected_mask):
"""Bradley-Terry pairwise ranking loss."""
r_chosen = reward_model(chosen_ids, chosen_mask) # scalar per sample
r_rejected = reward_model(rejected_ids, rejected_mask)
# Loss: chosen should have higher reward than rejected
# -log sigmoid(r_chosen - r_rejected)
loss = -F.logsigmoid(r_chosen - r_rejected).mean()
# Accuracy: fraction of pairs where chosen > rejected
accuracy = (r_chosen > r_rejected).float().mean()
return loss, accuracy
# Training data format: (prompt, chosen_response, rejected_response)
# Source: human preference labels (Anthropic HH, OpenAI Instruct comparisons)PPO Training Loop
class PPOTrainer:
def __init__(self, actor, ref_model, critic, reward_model, tokenizer):
self.actor = actor # policy being trained
self.ref = ref_model # frozen SFT model
self.critic = critic # value function
self.rm = reward_model # frozen reward model
self.actor_optimizer = AdamW(actor.parameters(), lr=1e-5)
self.critic_optimizer = AdamW(critic.parameters(), lr=1e-5)
def collect_rollouts(self, prompts: List[str]) -> RolloutBuffer:
"""Generate responses and collect data for training."""
rollouts = []
for prompt in prompts:
# Actor generates response
prompt_ids = tokenize(prompt)
with torch.no_grad():
response_ids = self.actor.generate(prompt_ids, max_new_tokens=256)
new_tokens = response_ids[len(prompt_ids):]
# Get reward from reward model
full_ids = torch.cat([prompt_ids, response_ids])
reward = self.rm(full_ids).item()
# Get log probs from actor and reference
actor_logprobs = get_log_probs(self.actor, prompt_ids, new_tokens)
ref_logprobs = get_log_probs(self.ref, prompt_ids, new_tokens)
# KL penalty per token
kl = actor_logprobs - ref_logprobs # (n_new_tokens,)
kl_reward = reward - self.kl_coef * kl.sum() # scalar
# Value estimate (from critic)
value = self.critic(full_ids).item()
rollouts.append(Rollout(
prompt_ids=prompt_ids,
response_ids=response_ids,
reward=kl_reward,
actor_logprobs=actor_logprobs,
ref_logprobs=ref_logprobs,
value=value
))
return rollouts
def compute_advantages(self, rollouts):
"""GAE (Generalized Advantage Estimation)."""
for rollout in rollouts:
# Simple version: advantage = reward - value baseline
rollout.advantage = rollout.reward - rollout.value
rollout.returns = rollout.reward # (single-step; normally bootstrapped)
def train_step(self, rollouts):
"""PPO clip update."""
self.compute_advantages(rollouts)
for epoch in range(self.ppo_epochs): # typically 4 epochs
for batch in batch_iter(rollouts, self.batch_size):
# Recompute current log probs
current_logprobs = get_log_probs(self.actor, batch.prompt_ids, batch.response_ids)
# Probability ratio (importance sampling)
ratio = torch.exp(current_logprobs - batch.actor_logprobs)
# PPO clipped objective
adv = batch.advantage
obj_unclipped = ratio * adv
obj_clipped = torch.clamp(ratio, 1-self.clip_eps, 1+self.clip_eps) * adv
actor_loss = -torch.min(obj_unclipped, obj_clipped).mean()
# Value function loss
current_values = self.critic(batch.full_ids)
critic_loss = F.mse_loss(current_values, batch.returns)
# Entropy bonus (encourages exploration)
entropy = compute_entropy(current_logprobs)
total_loss = actor_loss + 0.5 * critic_loss - 0.01 * entropy
self.actor_optimizer.zero_grad()
self.critic_optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 1.0)
self.actor_optimizer.step()
self.critic_optimizer.step()Project 2 — DPO (Direct Preference Optimization)
DPO vs PPO
PPO ISSUES:
✗ 4 models in memory simultaneously
✗ Reward model training step
✗ Complex PPO hyperparameter tuning
✗ Training instability
✗ Reward hacking (model exploits reward model weaknesses)
DPO INSIGHT:
The optimal policy under KL-constrained RL has a closed form:
π*(y|x) ∝ π_ref(y|x) × exp(r(x,y) / β)
Rearranging: r(x,y) = β × log[π*(y|x)/π_ref(y|x)] + Z(x)
Substitute into Bradley-Terry preference model:
P(y_w ≻ y_l | x) = σ(r(x,y_w) - r(x,y_l))
= σ(β × log[π_θ(y_w|x)/π_ref(y_w|x)]
- β × log[π_θ(y_l|x)/π_ref(y_l|x)])
→ The reward MODEL is IMPLICIT in the policy ratio!
Train policy directly on preference data, no explicit reward model needed
RESULT: Same 3 stages → SFT → DPO (replaces RM training + PPO)
Only 2 models in memory: policy + frozen referenceDPO Loss Implementation
def dpo_loss(
policy_model,
reference_model,
chosen_ids: torch.Tensor, # (B, T)
rejected_ids: torch.Tensor, # (B, T)
chosen_mask: torch.Tensor, # (B, T)
rejected_mask:torch.Tensor, # (B, T)
prompt_length: int, # to separate prompt from response
beta: float = 0.1,
) -> torch.Tensor:
"""DPO loss function."""
def get_logprobs(model, input_ids, attention_mask):
"""Compute per-token log probabilities for the completion."""
with torch.no_grad() if model is reference_model else contextlib.nullcontext():
logits = model(input_ids, attention_mask=attention_mask).logits
# Shift for next-token prediction
shift_logits = logits[:, :-1, :] # (B, T-1, V)
shift_labels = input_ids[:, 1:] # (B, T-1)
# Get log prob of actual tokens
logprobs = F.log_softmax(shift_logits, dim=-1)
token_logprobs = logprobs.gather(-1, shift_labels.unsqueeze(-1)).squeeze(-1) # (B, T-1)
# Mask out prompt tokens (only score completion)
completion_mask = attention_mask[:, 1:].clone()
completion_mask[:, :prompt_length-1] = 0
# Sum log probs over completion tokens
return (token_logprobs * completion_mask).sum(-1) # (B,)
# Log probs under policy (trained) and reference (frozen)
policy_chosen_logps = get_logprobs(policy_model, chosen_ids, chosen_mask)
policy_rejected_logps = get_logprobs(policy_model, rejected_ids, rejected_mask)
ref_chosen_logps = get_logprobs(reference_model, chosen_ids, chosen_mask)
ref_rejected_logps = get_logprobs(reference_model, rejected_ids, rejected_mask)
# Log ratios (policy vs reference)
chosen_logratios = policy_chosen_logps - ref_chosen_logps # (B,)
rejected_logratios = policy_rejected_logps - ref_rejected_logps # (B,)
# DPO loss: push chosen logratios higher, rejected lower
loss = -F.logsigmoid(beta * (chosen_logratios - rejected_logratios)).mean()
# Diagnostics
chosen_rewards = beta * chosen_logratios.detach()
rejected_rewards = beta * rejected_logratios.detach()
reward_accuracy = (chosen_rewards > rejected_rewards).float().mean()
reward_margin = (chosen_rewards - rejected_rewards).mean()
return loss, {
"reward_accuracy": reward_accuracy,
"reward_margin": reward_margin,
"chosen_rewards": chosen_rewards.mean(),
"rejected_rewards": rejected_rewards.mean()
}DPO Variants
STANDARD DPO (Rafailov et al 2023):
Loss = -log σ(β × (log π_θ(y_w)/π_ref(y_w) - log π_θ(y_l)/π_ref(y_l)))
IPO (Identity Preference Optimization):
Loss = (log π_θ(y_w)/π_ref(y_w) - log π_θ(y_l)/π_ref(y_l) - 1/2β)²
More stable, doesn't have σ saturation issue
KTO (Kahneman-Tversky Optimization):
Uses unpaired preference data (just "good" or "bad" labels)
No need for paired (chosen, rejected) comparisons
More data-efficient in practice
ORPO (Odds Ratio Preference Optimization):
Doesn't need reference model at all
Combines SFT and preference loss in single stage
loss = -log P(y|x) - λ × log σ(log odds(y_w) - log odds(y_l))Project 3 — Model Distillation Pipeline
Knowledge Distillation Architecture
TEACHER: Large, accurate model (e.g., GPT-4, Llama-3-70B)
STUDENT: Small, fast model (e.g., Llama-3-1B, custom architecture)
GOAL: Student learns to mimic teacher's behavior
Student accuracy >> same-size model trained from scratch
THREE TYPES OF DISTILLATION:
1. Black-box: Only teacher outputs (API access)
2. Logit-level: Teacher's full probability distribution
3. Feature-level: Match internal representationsLogit Distillation (Hinton et al)
def distillation_loss(
student_logits: torch.Tensor, # (B, T, V)
teacher_logits: torch.Tensor, # (B, T, V) — from frozen teacher
labels: torch.Tensor, # (B, T) — ground truth tokens
temperature: float = 4.0,
alpha: float = 0.7, # weight: 1=pure distill, 0=pure CE
) -> torch.Tensor:
"""Combined hard-label CE and soft-label KL divergence."""
# Hard-label loss (student vs ground truth)
hard_loss = F.cross_entropy(
student_logits.view(-1, student_logits.size(-1)),
labels.view(-1)
)
# Soft-label distillation loss
# Temperature scaling: higher T → softer probabilities → more info in distribution
student_soft = F.log_softmax(student_logits / temperature, dim=-1)
teacher_soft = F.softmax(teacher_logits / temperature, dim=-1)
# KL divergence: how different is student distribution from teacher
# KL(teacher || student) = Σ teacher × log(teacher/student)
soft_loss = F.kl_div(student_soft, teacher_soft, reduction='batchmean')
soft_loss *= temperature ** 2 # scale factor from Hinton et al
return alpha * soft_loss + (1 - alpha) * hard_loss
# Training loop:
# 1. Load frozen teacher, trainable student
# 2. Forward both on same inputs
# 3. Compute distillation loss
# 4. Update only student parametersSequence-Level Distillation (No Teacher Access Needed)
OFFLINE DISTILLATION:
1. Use teacher to generate responses for all training prompts
2. Store (prompt, teacher_response) pairs as training data
3. Train student with standard SFT on this data
Advantage: No need to run teacher during student training
Use case: Teacher is commercial API (expensive), student is open model
DATA GENERATION PIPELINE:
prompts = load_dataset("openhermes", split="train")
for batch in prompts:
responses = teacher_api.batch_generate(
batch,
temperature=1.0, # sample diversity
n=1, # one response per prompt
max_tokens=2048
)
save_to_dataset(batch, responses)
# Result: (prompt, response) pairs matching teacher's style
train_student_sft(student, synthetic_dataset)Task-Specific Distillation
SPECULATIVE DECODING DISTILLATION:
Train draft model to match target model's next-token distribution
Loss: minimize KL(target_distribution || draft_distribution) per position
Result: high acceptance rate in speculative decoding
LAYER-LEVEL DISTILLATION:
Map student layer i → teacher layer j
MSE loss on attention patterns or hidden states
Enables deeper compression
ATTENTION TRANSFER:
student_attn = student_layer.attention_probs # (B, H, T, T)
teacher_attn = teacher_layer.attention_probs # (B, H', T, T)
# Match averaged attention patterns
student_avg = student_attn.mean(1) # avg over heads
teacher_avg = teacher_attn.mean(1)
attn_loss = F.mse_loss(student_avg, teacher_avg.detach())Project 4 — Synthetic Data Generator
Architecture
┌──────────────────────────────────────────────────────────────┐
│ SYNTHETIC DATA GENERATOR │
│ │
│ PIPELINE: │
│ │
│ 1. SEED DATA │
│ Real samples from domain │
│ Few high-quality examples │
│ Task description │
│ │ │
│ ▼ │
│ 2. GENERATION │
│ Template-based generation │
│ LLM-based generation (Self-Instruct, Magpie) │
│ Constrained generation (entity slots) │
│ │ │
│ ▼ │
│ 3. QUALITY FILTERING │
│ Perplexity filter (too easy/hard) │
│ Reward model scoring │
│ LLM-as-judge evaluation │
│ Deduplication (MinHash) │
│ │ │
│ ▼ │
│ 4. POST-PROCESSING │
│ Format standardization │
│ PII removal │
│ Train/val/test split │
└──────────────────────────────────────────────────────────────┘Self-Instruct Algorithm
class SelfInstructGenerator:
def __init__(self, generator_llm, judge_llm, seed_tasks: List[dict]):
self.generator = generator_llm
self.judge = judge_llm
self.task_pool = seed_tasks # initial high-quality tasks
def generate_new_task(self) -> dict:
# Sample 8 existing tasks as few-shot examples
examples = random.sample(self.task_pool, min(8, len(self.task_pool)))
prompt = self._build_generation_prompt(examples)
# Generate new instruction
new_instruction = self.generator.generate(prompt, temperature=0.9)
# Generate input/output for the instruction
input_prompt = f"Given the instruction: {new_instruction}\nGenerate a good input:"
synthetic_input = self.generator.generate(input_prompt)
output_prompt = f"Instruction: {new_instruction}\nInput: {synthetic_input}\nOutput:"
synthetic_output = self.generator.generate(output_prompt, temperature=0.0)
return {
"instruction": new_instruction,
"input": synthetic_input,
"output": synthetic_output
}
def filter_task(self, task: dict) -> bool:
"""Quality gate using LLM-as-judge."""
# ROUGE deduplication — reject if too similar to existing
for existing in self.task_pool:
rouge = compute_rouge_l(task["instruction"], existing["instruction"])
if rouge > 0.7:
return False
# LLM quality judge
judge_prompt = f"""Rate this instruction-following sample on a 1-5 scale:
Instruction: {task['instruction']}
Input: {task['input']}
Output: {task['output']}
Rate: (1=poor, 5=excellent). Only output the number."""
score = float(self.judge.generate(judge_prompt))
return score >= 4.0
def run(self, target_size: int = 50000):
while len(self.task_pool) < target_size:
task = self.generate_new_task()
if self.filter_task(task):
self.task_pool.append(task)
return self.task_poolPrivacy-Preserving Synthetic Data (Healthcare)
CLINICAL NOTE SYNTHESIS:
Problem: Real patient data can't be used for training
Solution: Generate realistic synthetic notes with no real PII
PIPELINE:
1. Extract schema from de-identified samples (structure only)
{patient_age, diagnosis_codes, symptoms, medications, ...}
2. Sample realistic field values:
age: Gaussian(65, 15) clipped to [0, 100]
diagnosis: sample ICD-10 codes by prevalence
medications: conditional on diagnosis (drug-disease correlation)
3. Fill template with sampled values:
"Patient is a {age}-year-old {sex} presenting with {symptoms}..."
4. LLM fluency pass:
"Rewrite this clinical note to be natural and professional: {template}"
5. NER verification: confirm no real names/SSNs leaked through
6. Medical validity check: LLM-as-judge for clinical plausibility
DIFFERENTIAL PRIVACY:
If training on real data, use DP-SGD:
- Clip per-sample gradients: ||g|| ≤ C
- Add Gaussian noise: g += N(0, σ²C²I)
- Privacy accounting: (ε, δ)-DP guarantee
Tradeoff: privacy ε vs model utilityProject 5 — MoE Routing Layer (Mixture of Experts)
Architecture
STANDARD FFN:
output = FFN(x) = W₂ × GELU(W₁ × x)
All tokens through same network
MOE FFN (Sparse):
K "expert" FFNs, each token routed to top-2
Total params: K × FFN_params
Active params per token: 2 × FFN_params (same as 1 FFN)
output = Σᵢ gate_i(x) × Expert_i(x) where i ∈ top_2_experts(x)
BENEFITS:
Capacity scales with K (more experts = more knowledge)
Compute stays constant (only 2 experts per token)
Mistral MoE: 8 experts, top-2 → 46.7B params, 12.9B activeRouter Implementation
class MoELayer(nn.Module):
def __init__(self, d_model: int, d_ff: int, n_experts: int = 8, top_k: int = 2):
super().__init__()
self.n_experts = n_experts
self.top_k = top_k
# Routing network
self.router = nn.Linear(d_model, n_experts, bias=False)
# Expert FFNs (each is independent FFN)
self.experts = nn.ModuleList([
FeedForward(d_model, d_ff) for _ in range(n_experts)
])
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
x: (batch, seq_len, d_model)
"""
B, T, D = x.shape
x_flat = x.view(-1, D) # (B*T, D)
# Compute routing weights
router_logits = self.router(x_flat) # (B*T, n_experts)
router_probs = F.softmax(router_logits, dim=-1) # (B*T, n_experts)
# Select top-k experts per token
top_k_probs, top_k_indices = router_probs.topk(self.top_k, dim=-1)
# top_k_probs: (B*T, top_k)
# top_k_indices: (B*T, top_k)
# Renormalize top-k weights
top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
# Compute expert outputs (sparse dispatch)
output = torch.zeros_like(x_flat)
for expert_idx in range(self.n_experts):
# Find tokens routed to this expert
token_mask = (top_k_indices == expert_idx).any(dim=-1) # (B*T,)
if not token_mask.any():
continue
# Get weights for these tokens
expert_positions = (top_k_indices[token_mask] == expert_idx).float() # (n_tokens, top_k)
weights = (top_k_probs[token_mask] * expert_positions).sum(-1) # (n_tokens,)
# Run expert on selected tokens
expert_output = self.experts[expert_idx](x_flat[token_mask]) # (n_tokens, D)
output[token_mask] += weights.unsqueeze(-1) * expert_output
return output.view(B, T, D)Load Balancing Loss
def auxiliary_load_balancing_loss(router_probs: torch.Tensor, n_experts: int) -> torch.Tensor:
"""
Prevent all tokens routing to same expert.
Loss = n_experts × Σᵢ (fraction_of_tokens_i × mean_routing_prob_i)
"""
B_T = router_probs.size(0)
# Fraction of tokens per expert (hard assignment)
top_1 = router_probs.argmax(dim=-1) # (B*T,)
tokens_per_expert = torch.bincount(top_1, minlength=n_experts).float() / B_T # (n_experts,)
# Mean routing probability per expert (soft)
mean_prob_per_expert = router_probs.mean(dim=0) # (n_experts,)
# Auxiliary loss: encourages uniform distribution
aux_loss = n_experts * torch.dot(tokens_per_expert, mean_prob_per_expert)
return aux_loss # add to main training loss with small coefficient (1e-2)Project 6 — Distributed Training (FSDP / Tensor Parallelism)
Distributed Strategy Overview
MODEL PARALLELISM STRATEGIES:
1. DATA PARALLELISM (DDP)
- Each GPU: full model copy
- Each GPU: different batch
- Sync: all-reduce gradients after backward
- Limit: model must fit in 1 GPU
- Efficiency: linear scaling up to communication overhead
2. FULLY SHARDED DATA PARALLELISM (FSDP)
- Shard model parameters, gradients, optimizer states across GPUs
- Each GPU: 1/N of each tensor
- Communication: all-gather before forward, reduce-scatter after backward
- Enables: models too large for 1 GPU
- Memory: nearly linear reduction with N GPUs
3. TENSOR PARALLELISM (TP)
- Split individual matrices across GPUs
- Matrix multiply distributed: partial results merged via all-reduce
- Used within a single node (high bandwidth NVLink)
4. PIPELINE PARALLELISM (PP)
- Split model layers across GPUs (GPU1: layers 1-10, GPU2: 11-20...)
- Micro-batching to fill pipeline and reduce bubble time
- Used across nodes (slower interconnect OK)
HYBRID (3D Parallelism, Megatron-LM):
DP × TP × PP for 1000+ GPU trainingFSDP Implementation Sketch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
def train_with_fsdp():
# Initialize process group
dist.init_process_group(backend='nccl')
rank = dist.get_rank()
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
# Load model on CPU first (avoid GPU OOM)
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-3-70b")
# Wrap with FSDP
# transformer_auto_wrap_policy: shard at TransformerBlock boundaries
model = FSDP(
model,
auto_wrap_policy=transformer_auto_wrap_policy,
mixed_precision=MixedPrecision(
param_dtype=torch.bfloat16,
reduce_dtype=torch.float32,
buffer_dtype=torch.bfloat16
),
sharding_strategy=ShardingStrategy.FULL_SHARD, # shard params + grads + optimizer
device_id=local_rank,
cpu_offload=CPUOffload(offload_params=False), # keep params on GPU
)
optimizer = AdamW(model.parameters(), lr=1e-5)
for batch in dataloader:
batch = batch.to(local_rank)
output = model(**batch)
loss = output.loss
loss.backward() # FSDP handles gradient collection
optimizer.step() # each rank updates its shard
optimizer.zero_grad()
# Communication pattern per step:
# Forward: all-gather full params → compute → free non-local shards
# Backward: all-gather full params → compute grad → reduce-scatter gradients
# Optimizer: each rank updates its own parameter shardTensor Parallelism (Megatron-LM Style)
COLUMN PARALLEL LINEAR (split output dimension):
W: (d_in, d_out) → GPU_i gets W[:, d_out/N × i : d_out/N × (i+1)]
y_partial = x @ W_i ← each GPU computes partial output
(no communication needed in forward if next layer is row parallel)
ROW PARALLEL LINEAR (split input dimension):
W: (d_in, d_out) → GPU_i gets W[d_in/N × i : d_in/N × (i+1), :]
x_partial already on GPU_i from column parallel above
y_partial = x_partial @ W_i ← each GPU computes partial output
y = all_reduce(y_partial) ← sum partial results → full output
ATTENTION HEADS:
Split n_heads across GPUs: GPU_i handles heads [i*h : (i+1)*h]
Q, K, V projections: column parallel
Output projection: row parallel
One all-reduce per attention layer (not per head)Project 7 — State Space Model (Mamba)
SSM Fundamentals
CONTINUOUS-TIME STATE SPACE:
h'(t) = Ah(t) + Bx(t) ← state update (A: n×n, B: n×1)
y(t) = Ch(t) + Dx(t) ← output (C: 1×n, D: scalar)
DISCRETIZATION (for sequences):
Given step size Δ (learned):
Ā = exp(ΔA) ← Zero-Order Hold (ZOH)
B̄ = (ΔA)⁻¹(exp(ΔA) - I)ΔB
Discrete recurrence:
h_t = Āh_{t-1} + B̄x_t ← O(1) per step at inference!
y_t = Ch_t
CONVOLUTIONAL VIEW (for training):
Unroll recurrence:
y_t = Σₖ C Āᵏ B̄ x_{t-k} ← convolution with kernel K
K = [CB̄, CAB̄, CA²B̄, ...] ← causal convolution kernel
Training: compute K once → parallel FFT convolution over all t
Inference: use recurrent form → constant memory
Best of both: parallel training like Transformers, O(1) inference memoryMamba (Selective State Space)
KEY INNOVATION: Make A, B, C, Δ INPUT-DEPENDENT
Standard SSM: time-invariant (same A,B,C for all positions)
Mamba: Δ, B, C = Linear(x) → each input selects its own dynamics
Why this matters:
Allows model to selectively "remember" or "forget" based on content
Equivalent to discrete-time attention with linear complexity
MAMBA BLOCK:
x: (B, T, D)
# Input projection and selective SSM
x_ssm = Linear(D → d_inner)(x) # expand dimension
z = Linear(D → d_inner)(x) # gating branch
x_ssm = depthwise_conv1d(x_ssm) # local context mixing
x_ssm = silu(x_ssm)
# Compute selection parameters (input-dependent)
Δ = softplus(Linear(d_inner → d_inner)(x_ssm) + Δ_bias)
B = Linear(d_inner → d_state)(x_ssm) # (B, T, d_state)
C = Linear(d_inner → d_state)(x_ssm) # (B, T, d_state)
# Discretize and apply SSM
# A is learnable but input-independent (log parameterization)
Ā = exp(Δ × A) # (B, T, d_inner, d_state)
B̄ = Δ × B # simplification
# Parallel scan (associative scan) over time dimension
y = selective_scan(x_ssm, Ā, B̄, C, D) # CUDA kernel
# Output gate
y = y * silu(z)
output = Linear(d_inner → D)(y)
MEMORY: O(d_state) hidden state at inference (vs O(T×d) for Transformer KV cache)
COMPUTE: O(T) per layer (vs O(T²) for attention)
Quality: ~Transformer level on language modeling
Weakness: struggles with tasks requiring exact long-range token recallProject 8 — Interpretability Tool (SAE — Sparse Autoencoders)
The Superposition Problem
OBSERVATION: Neural networks represent MORE features than they have neurons
Features compete for neurons → polysemanticity
Single neuron activates for multiple unrelated concepts
SOLUTION: Sparse Autoencoder learns a WIDER, SPARSER basis
d_hidden << d_SAE (SAE hidden dim >> residual stream dim)
At any time, only k << d_SAE features active (sparse)
Each feature is more interpretable (monosemantic)
ANALOGY: Holographic compression. Information stored in superposition.
SAE "decompresses" into interpretable components.SAE Architecture
class SparseAutoencoder(nn.Module):
def __init__(self, d_in: int, d_sae: int, k: int = 32):
"""
d_in: dimension of model activations (e.g., 4096 for residual stream)
d_sae: SAE hidden dimension (e.g., 16384 = 4× expansion)
k: sparsity: keep only top-k features active
"""
super().__init__()
self.k = k
self.W_enc = nn.Linear(d_in, d_sae, bias=True) # encoder
self.W_dec = nn.Linear(d_sae, d_in, bias=True) # decoder
# Initialize decoder columns to unit norm
nn.init.orthogonal_(self.W_dec.weight)
def encode(self, x: torch.Tensor) -> torch.Tensor:
"""Project activations into SAE space, apply sparsity."""
pre_acts = F.relu(self.W_enc(x)) # (B, d_sae)
# TopK sparsity: keep only k largest activations
topk_vals, topk_idx = torch.topk(pre_acts, self.k, dim=-1)
acts = torch.zeros_like(pre_acts)
acts.scatter_(-1, topk_idx, topk_vals)
return acts
def decode(self, acts: torch.Tensor) -> torch.Tensor:
return self.W_dec(acts)
def forward(self, x: torch.Tensor):
# Normalize decoder columns (unit norm constraint)
self.W_dec.weight.data = F.normalize(self.W_dec.weight.data, dim=0)
acts = self.encode(x)
x_hat = self.decode(acts)
# Loss: reconstruction + L1 sparsity penalty
recon_loss = F.mse_loss(x_hat, x)
l1_loss = acts.abs().mean() # encourage sparsity
loss = recon_loss + self.l1_coef * l1_loss
return x_hat, acts, loss
# Training:
# 1. Collect activations from model at a specific layer (e.g., residual stream)
# 2. Train SAE to reconstruct these activations with sparse features
# 3. Examine what each feature responds to → what concept does it represent?Feature Interpretation Pipeline
def interpret_features(sae, activation_dataset, tokenizer, top_examples=10):
"""For each SAE feature, find inputs that maximally activate it."""
# Forward all data through SAE
all_acts = []
all_tokens = []
for batch_tokens, batch_acts in activation_dataset:
_, acts, _ = sae(batch_acts) # (B, T, d_sae)
all_acts.append(acts.detach())
all_tokens.append(batch_tokens)
all_acts = torch.cat(all_acts, dim=0) # (N, T, d_sae)
all_tokens = torch.cat(all_tokens, dim=0) # (N, T)
feature_interpretations = {}
for feature_idx in range(sae.d_sae):
# Find positions where this feature fires most strongly
feature_acts = all_acts[:, :, feature_idx] # (N, T)
# Top-k (doc, position) pairs
flat_acts = feature_acts.view(-1)
top_positions = flat_acts.topk(top_examples).indices
doc_ids = top_positions // all_acts.size(1)
pos_ids = top_positions % all_acts.size(1)
# Show context around activation
contexts = []
for doc, pos in zip(doc_ids, pos_ids):
context_start = max(0, pos-5)
context_end = min(all_tokens.size(1), pos+5)
token_context = all_tokens[doc, context_start:context_end]
text = tokenizer.decode(token_context)
activation_value = feature_acts[doc, pos].item()
contexts.append((text, activation_value, pos.item()))
feature_interpretations[feature_idx] = {
"top_contexts": contexts,
"mean_activation": feature_acts[feature_acts > 0].mean().item(),
"sparsity": (feature_acts > 0).float().mean().item()
}
return feature_interpretationsProject 9 — Data Curation Pipeline (MinHash / Deduplication)
Why Deduplication Matters
Problem: Web crawl data (CommonCrawl) contains massive duplication
- Same article scraped from 1000 mirror sites
- Boilerplate text (cookie notices, navigation menus)
- Near-duplicate versions of same document
Effects of training on duplicate data:
- Model memorizes exact text instead of generalizing
- Inflated benchmark scores (test set contamination)
- Wasted compute training on redundant information
Scale: CommonCrawl → after dedup, 30-50% of data often removedMinHash LSH Pipeline
class MinHashDeduplicator:
def __init__(self, num_hashes=128, ngram_size=5, threshold=0.85):
self.num_hashes = num_hashes
self.ngram_size = ngram_size
self.threshold = threshold
# Random hash functions: h(x) = (ax + b) mod p
self.a = np.random.randint(1, 2**32, size=num_hashes)
self.b = np.random.randint(0, 2**32, size=num_hashes)
self.p = 2**31 - 1 # large prime
def compute_minhash(self, text: str) -> np.ndarray:
"""Compute MinHash signature for a document."""
# Generate n-grams (character-level)
tokens = text.lower().split()
ngrams = set()
for i in range(len(tokens) - self.ngram_size + 1):
ngrams.add(tuple(tokens[i:i+self.ngram_size]))
# For each hash function, find minimum hash value across all n-grams
signature = np.full(self.num_hashes, np.inf)
for ngram in ngrams:
ngram_hash = hash(ngram) & 0xFFFFFFFF # 32-bit hash
hashed = (self.a * ngram_hash + self.b) % self.p
signature = np.minimum(signature, hashed)
return signature.astype(np.uint32)
def jaccard_from_signatures(self, sig1: np.ndarray, sig2: np.ndarray) -> float:
"""Estimate Jaccard similarity from MinHash signatures."""
return float((sig1 == sig2).mean())
def build_lsh_buckets(self, signatures: np.ndarray, n_bands: int = 16):
"""LSH: group documents likely to be near-duplicates into same bucket."""
n_docs, n_hashes = signatures.shape
rows_per_band = n_hashes // n_bands
buckets = defaultdict(list)
for doc_id in range(n_docs):
for band in range(n_bands):
start = band * rows_per_band
end = start + rows_per_band
band_sig = tuple(signatures[doc_id, start:end])
bucket_key = hash((band, band_sig))
buckets[bucket_key].append(doc_id)
return buckets
def find_duplicates(self, documents: List[str]) -> List[Tuple[int, int]]:
"""Find all near-duplicate pairs."""
signatures = np.array([self.compute_minhash(d) for d in documents])
buckets = self.build_lsh_buckets(signatures)
duplicate_pairs = set()
for bucket_docs in buckets.values():
if len(bucket_docs) < 2:
continue
# Verify pairs in same bucket
for i, j in combinations(bucket_docs, 2):
sim = self.jaccard_from_signatures(signatures[i], signatures[j])
if sim >= self.threshold:
duplicate_pairs.add((min(i,j), max(i,j)))
return list(duplicate_pairs)Project 10 — Multi-modal Projector (CLIP)
CLIP Architecture
CLIP = Contrastive Language-Image Pretraining (Radford et al, OpenAI 2021)
IMAGE TEXT
│ │
┌──────▼──────┐ ┌─────▼─────┐
│ Vision │ │ Text │
│ Encoder │ │ Encoder │
│ (ViT-L/14) │ │ (GPT-like)│
└──────┬──────┘ └─────┬─────┘
│ image_features │ text_features
│ (N, d_vision) │ (N, d_text)
│ │
┌──────▼──────┐ ┌──────▼──────┐
│ Projection │ │ Projection │
│ W_i │ │ W_t │
│ d_vision→d │ │ d_text→d │
└──────┬──────┘ └──────┬──────┘
│ image_emb (N, d) │ text_emb (N, d)
└──────────────┬───────────────┘
│
L2 normalize both
│
N×N similarity matrix:
logit_scale × image_emb @ text_emb.T
│
Symmetric Cross-Entropy Loss:
diagonal = matching (i,i) pairs → label iCLIP Training Loss
def clip_loss(image_features, text_features, logit_scale=np.log(1/0.07)):
"""Symmetric InfoNCE loss for CLIP."""
B = image_features.size(0)
# Normalize
image_features = F.normalize(image_features, dim=-1)
text_features = F.normalize(text_features, dim=-1)
# Similarity matrix (all image-text pairs in batch)
logit_scale = logit_scale.exp()
logits_per_image = logit_scale * image_features @ text_features.T # (B, B)
logits_per_text = logits_per_image.T
# Labels: diagonal is the matching pair
labels = torch.arange(B, device=image_features.device)
# Cross-entropy in both directions (symmetric)
loss_i = F.cross_entropy(logits_per_image, labels)
loss_t = F.cross_entropy(logits_per_text, labels)
return (loss_i + loss_t) / 2
# Scaling: N=32768 images per batch (4096 GPUs × 8 images)
# Large batch = many negatives = hard training signal
# logit_scale learnable, initialized to log(1/0.07) = 2.659Multi-modal Projector (LLaVA-style)
GOAL: Connect vision encoder to language model
Vision tokens → language embedding space
ARCHITECTURE OPTIONS:
1. Linear projection (LLaVA v1):
vision_tokens (N_v, d_vision) @ W_proj → (N_v, d_llm)
Simple, fast, surprisingly effective
2. MLP projection (LLaVA v1.5):
vision_tokens → Linear → GELU → Linear → (N_v, d_llm)
Better alignment between modalities
3. Perceiver Resampler (Flamingo):
Q-Former: fixed K learned queries attend to vision tokens
Output: K tokens regardless of input image resolution
More compute, reduces vision token count
4. C-Abstractor:
Pooling + MLP to reduce vision tokens aggressively
TRAINING STAGES:
Stage 1: ALIGNMENT PRETRAINING
Freeze: LLM + vision encoder
Train: projection layer ONLY
Data: image-caption pairs (LAION, CC3M)
Goal: align vision features to LLM embedding space
Stage 2: INSTRUCTION FINE-TUNING
Freeze: vision encoder
Train: projection + LLM (full or LoRA)
Data: visual instruction following data (LLaVA-Instruct)
Goal: follow multimodal instructions
INFERENCE:
image → ViT patch embeddings → projection → visual_tokens
prompt → tokenizer → text_tokens
[image_tokens, text_tokens] → LLM → response
Image tokens treated like any other tokens (interleaved in sequence)How These Projects Form One System
┌──────────────────────────────────────────────────────┐
│ MINI RESEARCH LAB INFRASTRUCTURE │
│ │
Pretraining ────▶ │ Distributed Training (FSDP + TP) │
│ │ │
Raw Web Data ───▶ │ Data Curation Pipeline (MinHash dedup) │
│ │ │
│ ┌────────▼──────────────────────────────────┐ │
│ │ POST-TRAINING LOOP │ │
│ │ │ │
│ │ SFT on curated data │ │
│ │ ↓ │ │
│ │ Synthetic Data Generation (Self-Instruct)│ │
│ │ ↓ │ │
│ │ RLHF (PPO) or DPO │ │
│ │ ↓ │ │
│ │ Distillation (smaller student models) │ │
│ └────────────────────────────────────────────┘ │
│ │
│ ARCHITECTURE VARIANTS: │
│ MoE routing → sparse model │
│ Mamba → linear-time sequence model │
│ CLIP projector → multimodal model │
│ │
│ ANALYSIS + SAFETY: │
│ SAE interpretability → understand what model learns│
│ Adversarial attacks → find failure modes │
│ NAS → search for better architectures │
│ │
│ SERVING: │
│ AI Gateway → route across model versions │
└──────────────────────────────────────────────────────┘
Full research cycle timeline:
Week 1-2: Data curation (dedup + quality filter)
Week 3-4: Pretraining / fine-tuning on curated data
Week 5: RLHF or DPO alignment
Week 6: Eval (eval harness, adversarial testing)
Week 7: Distillation to production models
Week 8: Interpretability analysis (SAE)
Ongoing: AI Gateway serves traffic, collects preference data → loop