LLM Architecture & Internals

This comprehensive technical reference covers the architecture, training, and inference of Large Language Models. We'll explore tokenization, embeddings, transformer blocks, attention mechanisms, and the full pipeline from raw text to generated outputs.

Overview: The LLM Pipeline

Key Dimensions of a Modern LLM

Complete Forward Pass Overview

# Complete end-to-end forward pass for LLaMA-3 405B

❶ RAW TEXT INPUT
   Input: "The capital of France is"

❷ TOKENIZATION (SentencePiece BPE)
   Algorithm: Splits text into subword units, maps to integer IDs
   Output: [128000, 976, 6864, 315, 9822, 374]

❸ TOKEN EMBEDDING
   W_E ∈ ℝ^(128256 × 16384)
   Shape: [6] → [6, 16384]

❹ POSITIONAL ENCODING (RoPE)
   Applied as rotations to Q, K vectors

❺ TRANSFORMER BLOCKS (× 126 layers)
   Per layer: RMSNorm → Attention → Residual → FFN → Residual

❻ OUTPUT PROCESSING
   Final RMSNorm → Linear projection (LM Head)
   Temperature-scaled softmax → probabilities

❼ AUTOREGRESSIVE LOOP
   Output token → append to input → repeat

2. Tokenization

Tokenization is the very first step: converting raw text into a sequence of integer IDs. The quality of the tokenizer directly affects model efficiency, multilinguality, and downstream task performance.

Byte-Pair Encoding (BPE) Algorithm

def train_bpe(text_corpus, target_vocab_size):
    # Step 1: Initialize with all individual bytes
    vocab = {bytes([i]): i for i in range(256)}

    while len(vocab) < target_vocab_size:
        # Step 2: Count all adjacent token pairs
        pair_freqs = count_pairs(corpus_encoded_with_vocab(vocab))
        # Step 3: Find most frequent pair
        most_frequent = max(pair_freqs, key=pair_freqs.get)
        # Step 4: Merge all occurrences
        new_token_id = len(vocab)
        vocab[most_frequent] = new_token_id
        corpus = merge_pairs(corpus, most_frequent, new_token_id)

    return vocab

# Example progression:
"lower" → "l o w e r" → "lo w e r" → "low er" → "lower"

Tokenizer Comparison

Tokenization Implementation

import sentencepiece as spm

# Load pre-trained tokenizer
sp = spm.SentencePieceProcessor()
sp.Load("tokenizer.model")

# Encode text to token IDs
text = "The quick brown fox jumps"
token_ids = sp.Encode(text, out_type=int)
# → [256, 4910, 3892, 6841, 5910]

# Decode token IDs back to text
decoded = sp.Decode(token_ids)
# → "The quick brown fox jumps"

# Get tokens as strings
tokens = sp.Encode(text, out_type=str)
# → ['▁The', '▁quick', '▁brown', '▁fox', '▁jumps']

Encoding Example - "unhappiness"

"unhappiness"
  ↓ (character split)
  → ["u", "n", "h", "a", "p", "p", "i", "n", "e", "ss"]
  ↓ (apply merge rules from training)
  → ["un", "h", "app", "iness"]
  ↓
  → ["unhapp", "iness"]
  ↓
Final token IDs: [48291, 7274]

3. Embeddings & Positional Encoding

After tokenization, each token ID is mapped to a dense vector using an embedding lookup table. Modern LLMs use Rotary Position Embeddings (RoPE) which encode position information via rotations in the attention mechanism.

Token Embedding Implementation

import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        # Shape: (128256, 16384) for LLaMA-3

    def forward(self, token_ids):
        # token_ids: (batch, seq_len)
        # output: (batch, seq_len, d_model)
        return self.embed(token_ids)

# Usage:
embedding_layer = TokenEmbedding(128256, 16384)
token_ids = torch.tensor([[256, 4910, 3892]])
embeddings = embedding_layer(token_ids)  # (1, 3, 16384)

Rotary Position Embedding (RoPE)

RoPE encodes absolute position information by rotating query and key vectors. For a 2D subspace with indices (2i, 2i+1), the rotation matrix is applied with frequency θ_i = 10000^(-2i/d).

def apply_rope(x, freqs_cis):
    # x: (batch, seq_len, n_heads, head_dim)
    # freqs_cis: precomputed complex frequencies
    # View as complex pairs
    x_complex = torch.view_as_complex(
        x.float().reshape(*x.shape[:-1], -1, 2)
    )
    # Multiply by complex exponentials (rotation)
    x_rotated = x_complex * freqs_cis
    return torch.view_as_real(x_rotated).reshape_as(x)

def precompute_freqs(dim, max_seq_len, theta=10000.0):
    # Compute inverse frequencies: θ_i = 10000^(-2i/d)
    inv_freqs = 1.0 / (theta ** (torch.arange(
        0, dim, 2
    ).float() / dim))
    # Create position indices
    t = torch.arange(max_seq_len)
    # Outer product: m × θ_i
    freqs = torch.outer(t, inv_freqs)
    # Convert to complex exponentials: e^(i·m·θ_i)
    return torch.polar(torch.ones_like(freqs), freqs)

# Usage:
freqs_cis = precompute_freqs(128, 4096)
q_rotated = apply_rope(q, freqs_cis)
k_rotated = apply_rope(k, freqs_cis)

RoPE Frequency Scaling for Extended Context

4. The Transformer Block

The Transformer block is the core computational unit. A typical LLM stacks 80-128 identical blocks, each with self-attention and feed-forward layers with layer normalization and residual connections.

Pre-Norm Transformer Block Architecture

Self-Attention Mechanism

Self-attention allows every token to attend to every other token. The computation follows: Attention(Q, K, V) = softmax(QK^T / √d_k) V

def scaled_dot_product_attention(Q, K, V, mask=None):
    # Q, K, V: (batch, seq_len, d_k)
    # 1. Compute attention scores
    scores = (Q @ K.transpose(-2, -1)) / sqrt(128)
    # shape: (batch, seq_len, seq_len)

    # 2. Apply causal mask
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # 3. Softmax to get weights
    weights = torch.softmax(scores, dim=-1)

    # 4. Weighted sum of values
    output = weights @ V
    return output

# Causal mask for seq_len=4:
mask = torch.tril(torch.ones(4, 4))
# [[1, 0, 0, 0],
#  [1, 1, 0, 0],
#  [1, 1, 1, 0],
#  [1, 1, 1, 1]]

Multi-Head Attention Implementation

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        batch_size, seq_len, _ = x.shape

        # Project to Q, K, V
        Q = self.W_q(x)  # (batch, seq, d_model)
        K = self.W_k(x)
        V = self.W_v(x)

        # Reshape to multiple heads
        Q = Q.reshape(batch_size, seq_len, self.n_heads, self.d_k)
        K = K.reshape(batch_size, seq_len, self.n_heads, self.d_k)
        V = V.reshape(batch_size, seq_len, self.n_heads, self.d_k)

        # Transpose for attention: (batch, n_heads, seq, d_k)
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        # Apply attention
        attn_out = scaled_dot_product_attention(Q, K, V)
        # attn_out: (batch, n_heads, seq, d_k)

        # Concatenate heads
        attn_out = attn_out.transpose(1, 2)
        attn_out = attn_out.reshape(batch_size, seq_len, -1)

        # Final projection
        output = self.W_o(attn_out)
        return output

MHA vs GQA vs MQA Comparison

5. Attention Mechanisms & Variants

Beyond standard self-attention, several variants optimize for speed, memory, or context length. This section covers modern approaches used in production LLMs.

Flash Attention Algorithm

Flash Attention reduces memory I/O by computing attention in tiles, avoiding expensive HBM reads/writes of intermediate matrices.

# Simplified Flash Attention block computation
def flash_attention_block(Q, K, V, block_size=128):
    # Tile-by-tile computation
    seq_len = Q.shape[0]
    output = torch.zeros_like(Q)

    for i in range(0, seq_len, block_size):
        # Load Q tile
        Q_tile = Q[i:i+block_size]
        # Compute attention to all K,V (causal)
        scores = Q_tile @ K[:i+block_size].T / sqrt(128)
        attn_weights = torch.softmax(scores, dim=-1)
        # Accumulate output
        output[i:i+block_size] = attn_weights @ V[:i+block_size]

    return output

# Flash Attention benefits:
# - 2-4× faster than standard attention on modern hardware
# - Same numerical result (no approximation)
# - ~3× less memory bandwidth required

Sparse Attention Patterns

Attention Key-Value Cache (KV Cache)

class CachedAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.cache_k = None
        self.cache_v = None

    def forward(self, q, k, v, is_prefill=False):
        # Prefill: process entire prompt at once
        if is_prefill:
            # Save all K, V from prompt
            self.cache_k = k
            self.cache_v = v
            attention = scaled_dot_product_attention(q, k, v)

        # Decode: one token at a time
        else:
            # K, V cache already filled; only new q
            self.cache_k = torch.cat([self.cache_k, k], dim=1)
            self.cache_v = torch.cat([self.cache_v, v], dim=1)
            attention = scaled_dot_product_attention(
                q, self.cache_k, self.cache_v
            )

        return attention

# Memory savings with KV cache:
# Without cache: 2·L·batch·T·h·d_k bytes for all K,V tensors
# With cache: Only cache previous steps (reuse computation)
# For inference: ~50-70% reduction in computation

6. Activation Functions

Activation functions introduce non-linearity into the feed-forward network. Modern LLMs typically use SwiGLU or GELU rather than ReLU.

Common Activation Functions

SwiGLU Implementation

class SwiGLU_FFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        # Dimension expansion: d_model → d_ff
        self.gate = nn.Linear(d_model, d_ff)
        self.up = nn.Linear(d_model, d_ff)
        # Dimension reduction: d_ff → d_model
        self.down = nn.Linear(d_ff, d_model)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        # Gate: d_model → d_ff, apply SiLU activation
        gate = torch.nn.functional.silu(self.gate(x))
        # Up projection: d_model → d_ff
        up = self.up(x)
        # Element-wise multiplication (gating)
        gated = gate * up
        # Down projection: d_ff → d_model
        output = self.down(gated)
        return output

# Compute parameter count:
# gate: d_model × d_ff = 16384 × 53248 = 872M params
# up: d_model × d_ff = 16384 × 53248 = 872M params
# down: d_ff × d_model = 53248 × 16384 = 872M params
# Total: 2.6B per layer × 126 layers = 327B params (81% of 405B model)

GELU Smooth Approximation

def gelu_approx(x):
    # Approximation of GELU (CDF of Gaussian)
    return 0.5 * x * (
        1 + torch.tanh(
            sqrt(2.0 / math.pi) * (x + 0.044715 * x**3)
        )
    )

def gelu_exact(x):
    # Use PyTorch's implementation
    return torch.nn.functional.gelu(x, approximate='none')

# Usage in FFN:
def gelu_ffn(x, W1, W2):
    return W2 @ gelu_approx(W1 @ x)

7. Normalization Techniques

Normalization stabilizes training by controlling activation distributions. Modern LLMs use RMSNorm instead of LayerNorm due to lower computation and similar effectiveness.

RMSNorm Implementation

class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        # Learnable scaling parameter
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        # Compute RMS (root mean square)
        rms = torch.sqrt(
            torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps
        )
        # Normalize and scale
        normalized = x / rms
        output = normalized * self.weight
        return output

# RMSNorm is simpler and faster than LayerNorm:
# - No mean subtraction (just RMS)
# - One less statistical moment to compute
# - ~8% faster on modern GPUs
# - Used in LLaMA, Mistral, Qwen, etc.

LayerNorm for Comparison

class LayerNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.bias = nn.Parameter(torch.zeros(d_model))
        self.eps = eps

    def forward(self, x):
        # Compute mean and variance
        mean = torch.mean(x, dim=-1, keepdim=True)
        var = torch.var(x, dim=-1, keepdim=True, unbiased=False)
        # Normalize
        normalized = (x - mean) / torch.sqrt(var + self.eps)
        # Scale and shift
        output = normalized * self.weight + self.bias
        return output

# LayerNorm properties:
# - Computes both mean and variance
# - Includes bias term (RMSNorm doesn't)
# - Slightly more expressive but slower

Normalization Comparison

Loss Functions & Training Objectives

The loss function defines what the model learns. Nearly all modern LLMs use autoregressive next-token prediction with cross-entropy loss, but understanding the nuances — label smoothing, auxiliary losses, and alternative objectives — is critical for training and fine-tuning.

Cross-Entropy Loss (Next-Token Prediction)

def autoregressive_loss(model, input_ids):
    # input_ids: (batch, seq_len) — token IDs
    # Model predicts next token at each position

    # Forward pass: get logits for all positions
    logits = model(input_ids[:, :-1])  # (batch, seq_len-1, vocab_size)

    # Targets: shifted by 1 position
    targets = input_ids[:, 1:]  # (batch, seq_len-1)

    # Cross-entropy loss: -log P(target_token | context)
    # CE = -Σ y_i · log(softmax(logit_i))
    loss = torch.nn.functional.cross_entropy(
        logits.view(-1, logits.size(-1)),  # flatten to (N, vocab)
        targets.view(-1),                   # flatten to (N,)
        ignore_index=PAD_TOKEN_ID            # skip padding
    )

    # Perplexity = exp(loss) — interpretable metric
    perplexity = torch.exp(loss)

    return loss, perplexity

# Key insight: loss is computed per-token, averaged over all positions
# A loss of 2.3 ≈ perplexity of 10 (model considers ~10 tokens equally likely)
# A loss of 1.0 ≈ perplexity of 2.7 (model is quite confident)

Label Smoothing

def cross_entropy_with_label_smoothing(logits, targets, smoothing=0.1):
    # Instead of hard targets [0, 0, 1, 0, ...], use soft targets
    # [ε/V, ε/V, 1-ε, ε/V, ...] where ε = smoothing, V = vocab_size

    vocab_size = logits.size(-1)
    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)

    # Smooth targets
    smooth_targets = torch.full_like(log_probs, smoothing / vocab_size)
    smooth_targets.scatter_(1, targets.unsqueeze(1), 1.0 - smoothing)

    # KL divergence as loss
    loss = -(smooth_targets * log_probs).sum(dim=-1).mean()
    return loss

# Benefits: prevents overconfident predictions, improves generalization
# Typical values: ε = 0.1 for pre-training, ε = 0.05 for fine-tuning
# Trade-off: slightly worse perplexity but better downstream performance

Training Objective Comparison

Auxiliary Losses in Modern LLMs

# Load balance = α · N · Σ(f_i · P_i)
# f_i = fraction of tokens routed to expert i
# P_i = average router probability for expert i
# α = 0.01 (small weight, added to main loss)

# z_loss = β · log²(Σ exp(z_i))
# β = 1e-4 (very small regularizer)
# Prevents logit magnitude explosion
# Critical for BF16 training stability

Optimizers & Learning Rate Schedules

The optimizer and learning rate schedule are the two most impactful hyperparameters after model architecture. Getting them wrong can mean training failure, instability, or 2-3× slower convergence.

AdamW: The Standard LLM Optimizer

class AdamW:
    """AdamW = Adam + decoupled weight decay (Loshchilov & Hutter, 2019)"""
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.95),
                 eps=1e-8, weight_decay=0.1):
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.wd = weight_decay
        # State: first moment (m), second moment (v), timestep (t)

    def step(self, param, grad):
        self.t += 1

        # 1. Update biased first moment estimate (momentum)
        self.m = self.beta1 * self.m + (1 - self.beta1) * grad

        # 2. Update biased second moment estimate (adaptive LR)
        self.v = self.beta2 * self.v + (1 - self.beta2) * grad**2

        # 3. Bias correction
        m_hat = self.m / (1 - self.beta1**self.t)
        v_hat = self.v / (1 - self.beta2**self.t)

        # 4. Parameter update (decoupled weight decay)
        param -= self.lr * (m_hat / (v_hat.sqrt() + self.eps))
        param -= self.lr * self.wd * param  # ← decoupled from gradient

        return param

# Key hyperparameters for LLM pre-training:
# lr: 1e-4 to 6e-4 (peak, depends on model size)
# betas: (0.9, 0.95) — β2=0.95 is more stable than 0.999 for LLMs
# weight_decay: 0.1 (standard for all transformer pre-training)
# eps: 1e-8 (FP32) or 1e-5 (BF16 for numerical stability)

Optimizer Comparison

Learning Rate Schedules

def cosine_schedule_with_warmup(
    step, warmup_steps=2000, total_steps=100000,
    peak_lr=3e-4, min_lr=3e-5
):
    """Standard LLM schedule: linear warmup → cosine decay"""

    if step < warmup_steps:
        # Phase 1: Linear warmup (0 → peak_lr)
        return peak_lr * (step / warmup_steps)
    else:
        # Phase 2: Cosine decay (peak_lr → min_lr)
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
        return min_lr + (peak_lr - min_lr) * cosine_decay

# Warmup is critical: prevents gradient explosion in early training
# Typical warmup: 0.5-2% of total steps (2000 steps for 100K total)
# min_lr: usually 10× smaller than peak_lr
# Some models use WSD schedule: Warmup → Stable → Decay

Schedule Comparison

Feed-Forward Networks & Gated Architectures

The FFN (feed-forward network) in each transformer block accounts for ~67% of model parameters. Modern LLMs use gated variants (SwiGLU, GeGLU) instead of the original two-layer MLP, achieving better performance at the same parameter count.

FFN Architecture Evolution

# Vaswani et al. FFN
# FFN(x) = W2 · ReLU(W1 · x + b1) + b2
# Expansion: d_model → 4·d_model → d_model

class VanillaFFN(nn.Module):
    def __init__(self, d=4096):
        self.w1 = nn.Linear(d, 4*d)
        self.w2 = nn.Linear(4*d, d)

    def forward(self, x):
        return self.w2(F.relu(self.w1(x)))

# SwiGLU(x) = (W1·x ⊙ SiLU(W_gate·x)) · W2
# ⊙ = element-wise multiply (gating)

class SwiGLU_FFN(nn.Module):
    def __init__(self, d=4096,
                 d_ff=14336):
        self.w1 = nn.Linear(d, d_ff)
        self.w_gate = nn.Linear(d, d_ff)
        self.w2 = nn.Linear(d_ff, d)

    def forward(self, x):
        return self.w2(
            self.w1(x) *
            F.silu(self.w_gate(x))
        )

# GeGLU(x) = (W1·x ⊙ GELU(W_gate·x)) · W2

class GeGLU_FFN(nn.Module):
    def __init__(self, d=4096,
                 d_ff=16384):
        self.w1 = nn.Linear(d, d_ff)
        self.w_gate = nn.Linear(d, d_ff)
        self.w2 = nn.Linear(d_ff, d)

    def forward(self, x):
        return self.w2(
            self.w1(x) *
            F.gelu(self.w_gate(x))
        )

Dimension Ratios in Real Models

Mixture of Experts (MoE) Architecture

MoE models replace the dense FFN with N parallel expert FFNs, routing each token to only the top-K experts. This enables massive parameter counts (e.g., 1.8T for Switch Transformer) while keeping FLOPs manageable — only 2 of 16 experts activate per token.

MoE Architecture Diagram

MoE Implementation

class MoELayer(nn.Module):
    def __init__(self, d_model=4096, d_ff=14336,
                 num_experts=16, top_k=2):
        super().__init__()
        self.experts = nn.ModuleList([
            SwiGLU_FFN(d_model, d_ff) for _ in range(num_experts)
        ])
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.top_k = top_k
        self.num_experts = num_experts

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        batch, seq_len, d = x.shape

        # Router: compute expert scores per token
        router_logits = self.router(x)  # (B, S, num_experts)

        # Select top-K experts per token
        weights, indices = torch.topk(
            torch.softmax(router_logits, dim=-1),
            k=self.top_k, dim=-1
        )
        # Normalize weights to sum to 1
        weights = weights / weights.sum(dim=-1, keepdim=True)

        # Compute weighted expert outputs
        output = torch.zeros_like(x)
        for k in range(self.top_k):
            expert_idx = indices[:, :, k]  # which expert for each token
            weight = weights[:, :, k].unsqueeze(-1)

            for e in range(self.num_experts):
                mask = (expert_idx == e)
                if mask.any():
                    expert_input = x[mask]
                    expert_output = self.experts[e](expert_input)
                    output[mask] += weight[mask] * expert_output

        return output

MoE Model Comparison

MoE Training Challenges

Pre-training Curriculum & Dynamics

Pre-training is the most expensive phase of LLM development ($2-50M+ for frontier models). Understanding training dynamics — loss curves, data scheduling, stability issues — is essential for efficient training runs.

Typical Pre-Training Setup

Data Curriculum (Multi-Stage Training)

# Modern pre-training uses phased data mixtures

class DataCurriculum:
    stages = [
        {
            "name": "Phase 1: General Pre-training",
            "tokens": "0 → 12T",
            "mix": {
                "web": 0.50,    # CommonCrawl, C4, refined
                "code": 0.17,   # GitHub, StackOverflow
                "books": 0.15,  # Gutenberg, books3
                "academic": 0.08, # arXiv, semantic scholar
                "math": 0.05,   # proofs, competition data
                "multilingual": 0.05
            }
        },
        {
            "name": "Phase 2: Quality Upsampling",
            "tokens": "12T → 14.5T",
            "mix": {
                "high_quality_web": 0.25,
                "code": 0.25,      # upsampled 2×
                "math": 0.15,      # upsampled 3×
                "academic": 0.15,  # upsampled 2×
                "books": 0.10,
                "curated_reasoning": 0.10
            }
        },
        {
            "name": "Phase 3: Long Context Extension",
            "tokens": "14.5T → 15T",
            "context_length": "8K → 128K",
            "note": "Gradual context extension with long documents"
        }
    ]

Common Training Instabilities

Loss Curve Interpretation

Tokenizer Chat Templates & Special Tokens

Every instruction-tuned LLM uses a specific chat template to format multi-turn conversations. Using the wrong template degrades performance significantly — the model was trained on a specific format and expects it exactly.

Common Special Tokens

Chat Template Formats

# LLaMA 3 Chat Template
<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>

What is the capital of France?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

The capital of France is Paris.<|eot_id|>

# Mistral / LLaMA 2 Template
<s>[INST] <<SYS>>
You are a helpful assistant.
<</SYS>>

What is the capital of France? [/INST] The capital of France is Paris. </s>

# ChatML Template (OpenAI-style, used by many open models)
<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
What is the capital of France?<|im_end|>
<|im_start|>assistant
The capital of France is Paris.<|im_end|>

Using HuggingFace Chat Templates

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What is the capital of France?"}
]

# apply_chat_template handles the format automatically
prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,        # return string (not token IDs)
    add_generation_prompt=True  # add assistant header for generation
)

# Tokenize for model input
inputs = tokenizer(prompt, return_tensors="pt")
# inputs.input_ids contains the correctly formatted token sequence

Tokenizer Algorithm Comparison

8. Training Methods & Distributed Training

Training large language models requires careful parallelism strategies to distribute computation and memory across multiple GPUs/TPUs.

Distributed Data Parallel (DDP)

import torch.distributed as dist
import torch.multiprocessing as mp

def setup_ddp(rank, world_size):
    # Initialize process group
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        rank=rank,
        world_size=world_size
    )
    torch.cuda.set_device(rank)

def train_ddp(rank, world_size, model, train_loader):
    setup_ddp(rank, world_size)

    # Wrap model with DDP
    model = model.cuda(rank)
    ddp_model = torch.nn.parallel.DistributedDataParallel(
        model,
        device_ids=[rank],
        find_unused_parameters=False
    )

    optimizer = torch.optim.AdamW(ddp_model.parameters())

    for epoch in range(10):
        for batch in train_loader:
            logits = ddp_model(batch)
            loss = compute_loss(logits, batch.labels)

            optimizer.zero_grad()
            loss.backward()

            # AllReduce: synchronize gradients across GPUs
            dist.all_reduce(torch.tensor(loss.item()))

            optimizer.step()

# DDP benefits:
# - Linear scaling: N GPUs → N× speedup
# - All GPUs have full model copy (no model sharding)
# - Synchronization via AllReduce on each backward pass

Parallelism Strategy Comparison

DeepSpeed ZeRO Optimization

# DeepSpeed ZeRO-2 Configuration (Gradient Partitioning)
{
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "contiguous_gradients": true
    },
    "fp16": {
        "enabled": true,
        "fp16_opt_level": "O2"
    },
    "gradient_clipping": 1.0,
    "train_micro_batch_size_per_gpu": 4,
    "gradient_accumulation_steps": 8
}

Gradient Accumulation & Checkpointing

def train_with_gradient_accumulation(
    model, train_loader, n_accumulation_steps=4
):
    # Accumulate gradients across multiple batches
    optimizer = torch.optim.AdamW(model.parameters())

    for batch_idx, batch in enumerate(train_loader):
        logits = model(batch.input_ids)
        loss = compute_loss(logits, batch.labels)

        # Scale loss by accumulation steps
        scaled_loss = loss / n_accumulation_steps
        scaled_loss.backward()

        # Update weights every n_accumulation_steps batches
        if (batch_idx + 1) % n_accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()

def train_with_gradient_checkpointing(model, batch):
    # Recompute activations instead of storing them
    def create_custom_forward(module):
        def custom_forward(*inputs):
            return module(*inputs)
        return custom_forward

    for layer in model.layers:
        # Checkpoint this layer
        x = torch.utils.checkpoint.checkpoint(
            create_custom_forward(layer), x
        )
    return x

# Gradient Accumulation benefits:
# - Larger effective batch size without GPU memory overhead
# - Gradient Checkpointing trades compute for memory (~30% slower, ~50% less VRAM)

9. Inference & Sampling Strategies

During inference, we generate tokens autoregressively using various sampling techniques to balance quality and diversity.

Temperature Scaling & Sampling

def sample_next_token(logits, temperature=0.8, top_k=50, top_p=0.9):
    # logits: (batch, vocab_size) from final layer

    # 1. Temperature scaling
    # - temp > 1: softer distribution (more diverse)
    # - temp = 1: unchanged
    # - temp < 1: sharper distribution (more deterministic)
    scaled_logits = logits / temperature

    # 2. Top-K filtering: keep only top k tokens
    topk_values, topk_indices = torch.topk(
        scaled_logits, k=top_k, dim=-1
    )
    # Mask out all other tokens
    mask = torch.full_like(scaled_logits, float('-inf'))
    mask.scatter_(dim=-1, index=topk_indices, src=topk_values)
    scaled_logits = mask

    # 3. Top-P (nucleus) sampling: cumulative probability threshold
    sorted_logits, sorted_indices = torch.sort(
        scaled_logits, descending=True, dim=-1
    )
    sorted_probs = torch.softmax(sorted_logits, dim=-1)
    cumsum_probs = torch.cumsum(sorted_probs, dim=-1)

    # Mask tokens where cumsum > top_p
    sorted_indices_to_remove = cumsum_probs > top_p
    # Keep at least 1 token
    sorted_indices_to_remove[..., 0] = False
    sorted_logits[sorted_indices_to_remove] = float('-inf')

    # 4. Sample from filtered distribution
    probs = torch.softmax(sorted_logits, dim=-1)
    next_token = torch.multinomial(probs, num_samples=1)

    return next_token

# Common strategies:
# - Greedy: argmax (deterministic, can get stuck)
# - Temperature: softmax with scaling (simple, controlled randomness)
# - Top-K: filter low-probability tokens (prevents nonsense)
# - Top-P (Nucleus): dynamic filtering (usually best results)

Sampling Comparison

Batched Inference with KV Caching

def generate_batched(
    model, prompt_ids, max_length=512, batch_size=16
):
    # prompt_ids: (batch, prompt_len)
    batch_size, prompt_len = prompt_ids.shape

    # Initialize KV cache
    kv_cache = {"k": [], "v": []} for _ in model.layers

    # Prefill: process prompt in parallel
    logits = model(prompt_ids, kv_cache, is_prefill=True)

    generated_ids = prompt_ids.clone()

    for _ range(max_length - prompt_len):
        # Get last token logits from last layer
        last_logits = logits[:, -1, :]

        # Sample next token
        next_tokens = sample_next_token(
            last_logits, temperature=0.7, top_p=0.9
        )  # (batch,)

        # Process only new tokens (decode)
        logits = model(
            next_tokens.unsqueeze(-1),  # (batch, 1)
            kv_cache,
            is_prefill=False
        )

        generated_ids = torch.cat([generated_ids, next_tokens.unsqueeze(-1)], dim=-1)

    return generated_ids

# Batched inference speedups:
# - KV cache only stores previous tokens (not recomputed)
# - Prefill: vectorized computation on entire prompt
# - Decode: only 1 token at a time (memory-bound, not compute-bound)
# - Typical speedup: 10-30× faster than naive generation

10. Quantization

Quantization reduces model size and inference latency by converting weights from FP32 to lower-precision formats (INT8, INT4, etc.).

Quantization Techniques

INT8 Quantization

def quantize_to_int8(tensor, eps=1e-6):
    # Find per-channel (per-output) scaling factor
    # Assume tensor shape: (out_features, in_features)

    # Per-output scaling (most common)
    max_abs = torch.max(torch.abs(tensor), dim=-1, keepdim=True)[0]

    # Scale to [-128, 127]
    scales = max_abs / 127.0
    quantized = torch.round(tensor / scales).to(torch.int8)

    return quantized, scales

def dequantize_from_int8(quantized, scales):
    # Restore to original precision
    return quantized.float() * scales

# Example:
weights_fp32 = torch.randn(4096, 11008)  # LLaMA FFN layer
weights_int8, scales = quantize_to_int8(weights_fp32)
# Memory: 4096 × 11008 × 4 bytes = 180MB → 45MB (4× reduction)

INT4 GPTQ Quantization

def gptq_quantize(tensor, bits=4, group_size=128):
    # Quantize to INT4 with group-wise scaling
    # tensor: (out_features, in_features)
    out_features, in_features = tensor.shape

    quantized = torch.zeros_like(tensor, dtype=torch.uint8)
    scales = torch.zeros(out_features, (in_features + group_size - 1) // group_size)

    for i in range(out_features):
        for j in range(0, in_features, group_size):
            group = tensor[i, j:j+group_size]

            # Find optimal scale for this group
            max_val = torch.max(torch.abs(group))
            scale = max_val / ((1 << (bits-1)) - 1)

            # Quantize and store
            quantized[i, j:j+group_size] = torch.round(group / scale).int()
            scales[i, j // group_size] = scale

    return quantized, scales

# GPTQ properties:
# - Uses Hessian information for better quantization
# - Per-group scaling reduces error vs per-channel
# - 8× size reduction (FP32 to INT4)
# - Typically <1% accuracy loss

QLoRA for Efficient Fine-Tuning

class QLoRA_Linear(nn.Module):
    def __init__(self, in_features, out_features, r=8):
        super().__init__()
        # Original weights quantized to NF4 (frozen)
        self.W_q = None  # Loaded from quantized model

        # Low-rank adapters (trainable)
        self.lora_a = nn.Linear(in_features, r, bias=False)
        self.lora_b = nn.Linear(r, out_features, bias=False)

        # Initialize: lora_a ~ N(0, 1), lora_b = 0
        nn.init.normal_(self.lora_a.weight)
        nn.init.zeros_(self.lora_b.weight)

        self.scaling = 1.0 / r

    def forward(self, x):
        # x: (batch, seq_len, in_features)
        # Original (frozen) computation
        output = dequantize_and_apply(x, self.W_q)

        # Add low-rank update
        lora_out = self.lora_b(self.lora_a(x))

        output = output + self.scaling * lora_out
        return output

# QLoRA memory efficiency:
# - Base model: 13B params in NF4 = ~7GB VRAM
# - LoRA adapters: 13B × r/d = ~600MB (r=8)
# - Total: ~7.6GB vs 52GB (FP32 13B model)
# - Trainable on single 24GB GPU!

11. Fine-tuning Methods

Fine-tuning adapts pretrained models to specific tasks. Methods range from full fine-tuning to efficient parameter-efficient approaches like LoRA.

Fine-tuning Methods Comparison

LoRA Implementation

class LoRA_Linear(nn.Module):
    def __init__(self, original_linear, rank=8, alpha=16):
        super().__init__()
        self.original = original_linear
        # Freeze original weights
        original_linear.requires_grad = False

        # Low-rank decomposition: ΔW = B @ A (r << d)
        in_features = original_linear.in_features
        out_features = original_linear.out_features

        self.lora_a = nn.Linear(in_features, rank, bias=False)
        self.lora_b = nn.Linear(rank, out_features, bias=False)

        # Initialize: A ~ N(0, 1), B = 0
        nn.init.normal_(self.lora_a.weight, std=1.0 / rank)
        nn.init.zeros_(self.lora_b.weight)

        # Scaling factor
        self.scaling = alpha / rank

    def forward(self, x):
        # y = W·x + (B @ A)·x = W·x + α/r·B·(A·x)
        result = self.original(x)
        lora_update = self.lora_b(self.lora_a(x))
        return result + self.scaling * lora_update

# LoRA parameter count:
# - Original layer: 4096 × 11008 = 45M params
# - LoRA adapters: (4096 × 8) + (8 × 11008) = 32K + 88K = 120K
# - Reduction: 45M → 120K = 0.27% trainable params

Training with LoRA

def train_with_lora(base_model, train_dataset, lr=2e-4):
    # Step 1: Replace linear layers with LoRA versions
    for name, module in base_model.named_modules():
        if isinstance(module, nn.Linear) and "q_proj" in name:
            new_module = LoRA_Linear(module, rank=8, alpha=16)
            # Replace in parent

    # Step 2: Only train LoRA parameters
    lora_params = []
    for name, param in base_model.named_parameters():
        if "lora" in name:
            lora_params.append(param)
            param.requires_grad = True
        else:
            param.requires_grad = False

    optimizer = torch.optim.AdamW(lora_params, lr=lr)

    for batch in train_dataset:
        logits = base_model(batch.input_ids)
        loss = compute_loss(logits, batch.labels)

        loss.backward()
        torch.nn.utils.clip_grad_norm_(lora_params, 1.0)
        optimizer.step()
        optimizer.zero_grad()

# LoRA advantages:
# - 99.7% parameter reduction vs full fine-tuning
# - Can store multiple LoRA adapters (one per task)
# - Merge LoRA weights into original after training
# - Only marginal accuracy loss (often none)

Fine-tuning Methods Comparison

from peft import get_peft_model, LoraConfig, TaskType
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer

# 1. Load base model
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.3-8B",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.3-8B")
tokenizer.pad_token = tokenizer.eos_token

# 2. Configure LoRA
lora_config = LoraConfig(
    r=16,                        # Rank — higher = more capacity
    lora_alpha=32,               # Scaling factor (usually 2×r)
    lora_dropout=0.05,           # Regularization
    target_modules=[              # Which layers to adapt
        "q_proj", "k_proj",
        "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    task_type=TaskType.CAUSAL_LM,
    bias="none",                  # Don't train biases
)

# 3. Apply LoRA
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Output: "trainable params: 83,886,080 || all params: 8,114,212,864 || 1.03%"

# 4. Train with SFTTrainer
training_args = TrainingArguments(
    output_dir="./lora-llama",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,  # effective batch = 16
    learning_rate=2e-4,
    bf16=True,
    warmup_ratio=0.03,
    logging_steps=10,
    save_strategy="steps",
    save_steps=200,
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
    max_seq_length=2048,
)
trainer.train()

# After training: merge LoRA weights into base model for zero-overhead inference
from peft import PeftModel

# Load base + LoRA
base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.3-8B")
lora_model = PeftModel.from_pretrained(base_model, "./lora-llama")

# Merge: W' = W + (α/r) · B · A  →  single matrix, no LoRA overhead
merged_model = lora_model.merge_and_unload()

# Save merged model — identical interface to original, zero inference cost
merged_model.save_pretrained("./llama-merged")
tokenizer.save_pretrained("./llama-merged")

# Now deploy with vLLM, TensorRT-LLM, or any framework — no PEFT dependency needed

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training

# 1. 4-bit quantization config (NF4 + double quantization)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # NormalFloat4 — optimal for normal distributions
    bnb_4bit_use_double_quant=True,   # Quantize the quantization constants too
    bnb_4bit_compute_dtype=torch.bfloat16,  # Compute in BF16
)

# 2. Load model in 4-bit
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-70B",
    quantization_config=bnb_config,
    device_map="auto",
)

# 3. Prepare for k-bit training (handle gradient checkpointing + casting)
model = prepare_model_for_kbit_training(model)

# 4. Apply LoRA on top of quantized model
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# "trainable: 167M | all: 70B | 0.24%"

# VRAM usage breakdown:
# - 70B model in NF4: ~35 GB
# - LoRA adapters (BF16): ~1.5 GB
# - Optimizer states: ~3 GB (AdamW on LoRA params only)
# - Activations/gradients: ~5 GB (with gradient checkpointing)
# - Total: ~44 GB → fits on 1× A100 48GB or 2× A40 48GB

# DoRA: W' = m · (W + ΔW) / ||W + ΔW||
# where m = learnable magnitude, ΔW = B·A (standard LoRA)
# In HuggingFace PEFT:
lora_config = LoraConfig(
    r=16, use_dora=True,  # Enable DoRA
    target_modules=["q_proj", "v_proj"],
)

# rsLoRA: scaling = α/√r instead of α/r
# Prevents gradient explosion at high ranks
lora_config = LoraConfig(
    r=64, lora_alpha=64,
    use_rslora=True,  # Rank-stabilized scaling
)

# LoRA+: separate LR for A and B
optimizer = torch.optim.AdamW([
    {"params": lora_a_params, "lr": 1e-4},
    {"params": lora_b_params, "lr": 5e-4},
])

# QA-LoRA: r aligned to group_size
# group_size=32, r must be multiple of 32
# After merge: direct GPTQ/AWQ quantization
# Result: quantized + adapted in one shot

Adapters: Lightweight Modular Fine-Tuning

Adapters are small trainable bottleneck modules inserted between frozen Transformer layers. Unlike LoRA (which modifies existing weights), adapters add new parameters in a down-project → nonlinearity → up-project pattern.

Adapter Implementation

class Adapter(nn.Module):
    """Bottleneck adapter inserted after attention or FFN sub-layer."""
    def __init__(self, d_model: int, bottleneck: int = 64):
        super().__init__()
        self.down = nn.Linear(d_model, bottleneck, bias=False)
        self.act  = nn.GELU()
        self.up   = nn.Linear(bottleneck, d_model, bias=False)
        nn.init.zeros_(self.up.weight)  # init to identity residual

    def forward(self, x):
        return x + self.up(self.act(self.down(x)))  # residual + adapter

# Adapter param count (LLaMA-3 8B, bottleneck=64):
# - Per adapter: 4096 × 64 + 64 × 4096 = 524K params
# - 2 adapters per layer (after attn + after FFN) × 32 layers = 33.5M
# - Total: 33.5M / 8B = 0.4% of model params

class AdapterTransformerLayer(nn.Module):
    """Transformer layer with adapters after attention and FFN."""
    def __init__(self, original_layer, bottleneck=64):
        super().__init__()
        self.layer = original_layer  # frozen
        d = original_layer.self_attn.q_proj.in_features
        self.attn_adapter = Adapter(d, bottleneck)
        self.ffn_adapter  = Adapter(d, bottleneck)

    def forward(self, x, **kwargs):
        # Attention sub-layer + adapter
        h = self.layer.self_attn(x, **kwargs)
        h = self.attn_adapter(h)       # ← adapter after attention
        x = x + h                        # residual

        # FFN sub-layer + adapter
        h = self.layer.mlp(x)
        h = self.ffn_adapter(h)        # ← adapter after FFN
        return x + h                    # residual

Adapter vs LoRA vs Full Fine-Tuning

12. Safety & Alignment

Ensuring LLMs are safe, helpful, and honest requires multi-layered defenses spanning data curation, alignment training, runtime guardrails, and continuous adversarial testing.

The Alignment Pipeline

Alignment Techniques in Detail

Runtime Safety Stack

Known Safety Challenges

13. Evaluation & Benchmarks

LLM evaluation requires multiple complementary approaches — no single benchmark captures all capabilities. Production systems need both offline benchmarks and online quality monitoring.

Evaluation Framework

Major Benchmarks

Performance Trends (2020 → 2025)

Evaluation Pitfalls

Context Window & Long-Range Techniques

Modern LLMs need to process long documents, multi-turn conversations, and large codebases. These techniques extend effective context beyond the training window.

Context Extension Methods

Sliding Window + Sink Tokens

def sliding_window_attention(Q, K, V, window_size=4096, n_sink=4):
    """Attention with sliding window + sink tokens for streaming."""
    seq_len = Q.shape[1]
    attn_mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)

    for i in range(seq_len):
        # Always attend to first n_sink tokens (sink tokens)
        attn_mask[i, :n_sink] = True
        # Attend to local window
        start = max(n_sink, i - window_size)
        attn_mask[i, start:i+1] = True

    scores = (Q @ K.transpose(-2, -1)) / math.sqrt(Q.shape[-1])
    scores.masked_fill_(~attn_mask, float('-inf'))
    return F.softmax(scores, dim=-1) @ V

# Mistral uses W=4096 sliding window
# With 32 layers, effective context = 4096 × 32 = 131K
# But only attends locally — long-range info degrades

14. Frontier Research Directions

The field of LLMs is rapidly evolving with new approaches addressing context length, reasoning, multimodality, and interpretability.

Emerging Architectures

Interpretability Techniques

• Mechanistic Interpretability: Identify circuits and features that compute specific functions
• Attention Analysis: Visualize what tokens attend to; identify important attention heads
• Probing: Train classifiers on intermediate representations to decode semantic information
• Activation Decomposition: Use Sparse Autoencoders to find interpretable features in activations
• Causal Tracing: Ablate representations to find critical computation paths

Future Scaling Trends

Advanced Reasoning Techniques

# Chain-of-Thought prompting for better reasoning
def generate_with_reasoning(model, question, max_reasoning_tokens=2000):
    # Phase 1: Generate reasoning trace
    prompt = f"""Q: {question}
Let me work through this step by step:
"""
    reasoning = model.generate(
        prompt,
        max_tokens=max_reasoning_tokens,
        temperature=0.7
    )

    # Phase 2: Generate final answer based on reasoning
    full_prompt = prompt + reasoning + "\nTherefore, the answer is:"
    answer = model.generate(
        full_prompt,
        max_tokens=100,
        temperature=0.3
    )

    return reasoning, answer

# Training with reasoning supervision (o1 style):
# 1. Collect examples with explicit reasoning traces
# 2. SFT on reasoning_trace + answer pairs
# 3. RL: reward correct final answer (only reward if reasoning ∈ valid_steps)
# 4. Trade inference compute for accuracy (test-time scaling)

Model Merging & Mixture of Experts

def linear_interpolation_merge(model_a, model_b, alpha=0.5):
    # Merge two fine-tuned models via convex combination
    # θ_merged = α·θ_a + (1-α)·θ_b

    merged_model = copy.deepcopy(model_a)

    for (name_a, param_a), (name_b, param_b) in zip(
        model_a.named_parameters(), model_b.named_parameters()
    ):
        assert name_a == name_b
        merged_model.get_parameter(name_a).data = (
            alpha * param_a.data + (1 - alpha) * param_b.data
        )

    return merged_model

class MixtureOfExperts(nn.Module):
    def __init__(self, expert_models, num_experts_per_token=2):
        super().__init__()
        self.experts = nn.ModuleList(expert_models)
        # Router network learns to select experts
        self.router = nn.Linear(4096, len(expert_models))
        self.k = num_experts_per_token

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        # Get router logits
        router_logits = self.router(x)
        # Select top-k experts per token
        expert_weights, expert_indices = torch.topk(
            torch.softmax(router_logits, dim=-1), k=self.k
        )

        output = torch.zeros_like(x)
        for i in range(0, self.k):
            expert_id = expert_indices[..., i]
            weight = expert_weights[..., i].unsqueeze(-1)

            # Apply selected experts (per token)
            expert_output = apply_per_token_expert(
                x, expert_id, self.experts
            )
            output = output + weight * expert_output

        return output

Speculative Decoding

def speculative_decoding(
    large_model, small_model, prompt, max_tokens=256, gamma=4
):
    # gamma = number of tokens to speculate
    generated = prompt.clone()

    for _ range(max_tokens):
        # Step 1: Small model proposes gamma tokens
        draft_tokens = []
        logits_small = small_model(generated)

        for _ range(gamma):
            next_token = torch.argmax(logits_small[:, -1, :], dim=-1)
            draft_tokens.append(next_token)
            # Append to sequence for next prediction
            logits_small = small_model(
                torch.cat([generated, torch.stack(draft_tokens).T], dim=1)
            )

        # Step 2: Large model validates/corrects in parallel
        proposed_seq = torch.cat([generated, torch.stack(draft_tokens).T], dim=1)
        logits_large = large_model(proposed_seq)

        # Compare probabilities and accept/reject
        n_accepted = 0
        for i, draft_token in enumerate(draft_tokens):
            # Acceptance threshold based on large model confidence
            large_prob = torch.softmax(logits_large[:, generated.shape[1] + i], dim=-1)
            small_prob = torch.softmax(logits_small[:, generated.shape[1] + i - 1], dim=-1)

            # Rejection sampling
            accept_prob = min(1.0, large_prob[draft_token] / small_prob[draft_token])
            if torch.rand(1) < accept_prob:
                generated = torch.cat([generated, draft_token.unsqueeze(-1)], dim=1)
                n_accepted += 1
            else:
                break

        # If no tokens accepted, sample from large model
        if n_accepted == 0:
            logits = large_model(generated)
            next_token = torch.multinomial(
                torch.softmax(logits[:, -1, :], dim=-1), num_samples=1
            )
            generated = torch.cat([generated, next_token], dim=1)

    return generated

# Speculative decoding speedup:
# - Best case: 2-3× faster (gamma tokens generated with small model latency)
# - Worst case: ~same as large model (if small model is poor)
# - Accuracy: identical to large model (rejection sampling)

Appendix A: Data Scaling Laws

Empirical scaling laws reveal how model size, data, and compute relate to final loss. Chinchilla Optimal Compute (2022) suggests equal investment in model size and data tokens.

Chinchilla Scaling Laws

# Chinchilla optimal scaling: N ≈ D (equal parameters and tokens)
# This contradicts GPT-3 which was undertrained for its size

def estimate_optimal_compute(compute_budget_flops):
    # Compute budget for training = 6·N·D
    # (6 FLOPs per parameter per token)

    # Chinchilla optimal: N = D
    n_params = (compute_budget_flops / 6) ** 0.5
    n_tokens = (compute_budget_flops / 6) ** 0.5

    return n_params, n_tokens

# Example: 10^24 FLOPs budget
n, d = estimate_optimal_compute(1e24)
# → N ≈ 408B params, D ≈ 408B tokens (LLaMA 3 405B region!)

Loss Scaling Models

Typical Dataset Composition

Appendix B: Advanced Decoding Methods

Beam Search

def beam_search(model, prompt, beam_width=5, max_tokens=100):
    # Track top-k complete hypotheses
    hypotheses = [
        {"tokens": prompt, "score": 0.0, "finished": False}
        for _ range(beam_width)
    ]

    for step in range(max_tokens):
        all_candidates = []

        for hyp in hypotheses:
            if hyp["finished"]:
                all_candidates.append(hyp)
                continue

            # Get logits for next token
            logits = model(hyp["tokens"])
            log_probs = torch.log_softmax(logits[:, -1, :], dim=-1)

            # Expand: try all vocab tokens
            for token_id in range(vocab_size):
                new_hyp = {
                    "tokens": torch.cat([hyp["tokens"], torch.tensor([token_id])]),
                    "score": hyp["score"] + log_probs[token_id].item(),
                    "finished": token_id == EOS_TOKEN
                }
                all_candidates.append(new_hyp)

        # Keep top-k by score (length-normalized)
        all_candidates.sort(
            key=lambda x: x["score"] / len(x["tokens"]),
            reverse=True
        )
        hypotheses = all_candidates[:beam_width]

        # Stop if all hypotheses are finished
        if all(h["finished"] for h in hypotheses):
            break

    return hypotheses[0]["tokens"]

# Beam search properties:
# - Explores multiple paths (not greedy)
# - Finds higher-probability sequences than greedy
# - Slower: O(beam_width · vocab_size) per step
# - Trade-off: beam_width=1 (greedy), beam_width=10 (thorough but slow)

Contrastive Search

def contrastive_search_sampling(
    model, prompt, alpha=0.6, top_k=10
):
    # Balance model confidence and diversity
    # Score = α·P(token|context) + (1-α)·contrast(token)

    generated = prompt.clone()

    for _ range(100):
        # Get model logits
        logits = model(generated)
        log_probs = torch.log_softmax(logits[:, -1, :], dim=-1)

        # Get top-k candidates
        topk_log_probs, topk_indices = torch.topk(log_probs, k=top_k)

        # Get embeddings for context and candidates
        context_emb = model.get_last_hidden_state(generated)[:, -1, :]
        candidate_embs = model.embedding(topk_indices)

        # Compute contrast: similarity to context
        similarity = torch.nn.functional.cosine_similarity(
            context_emb.unsqueeze(1), candidate_embs, dim=-1
        )
        contrast = 1.0 - similarity

        # Combined score
        scores = alpha * topk_log_probs + (1 - alpha) * contrast

        # Select token with highest score
        best_idx = torch.argmax(scores)
        next_token = topk_indices[best_idx]

        generated = torch.cat([generated, next_token.unsqueeze(-1)], dim=-1)

        if next_token == EOS_TOKEN:
            break

    return generated

# Contrastive search benefits:
# - Avoids repetition and dull text
# - More diverse outputs than top-p sampling
# - Slight quality tradeoff vs model confidence alone

Appendix C: Implementation Details & Gotchas

Common Implementation Pitfalls

Efficient Attention Implementation

# Memory-efficient attention (avoid materializing Q×K^T)
def efficient_attention(Q, K, V, block_size=128):
    # Process in blocks to stay within GPU memory
    seq_len = Q.shape[0]
    output = torch.zeros_like(Q)

    for i in range(0, seq_len, block_size):
        # Load Q block
        Q_block = Q[i:i+block_size]

        # Compute attention to all K, V (causal)
        for j range(0, i+block_size, block_size):
            K_block = K[j:j+block_size]
            V_block = V[j:j+block_size]

            # Compute Q @ K^T for this block pair
            scores = Q_block @ K_block.T / sqrt(128)

            # Apply causal mask
            if j + block_size > i + block_size:
                mask = torch.ones(scores.shape, dtype=torch.bool)
                mask.triu_(j - i + 1)
                scores[mask] = float('-inf')

            # Softmax and aggregate
            weights = torch.softmax(scores, dim=-1)
            output[i:i+block_size] += weights @ V_block

    return output

Proper Weight Initialization

def init_transformer_weights(module):
    # Kaiming initialization for better convergence
    if isinstance(module, nn.Linear):
        # He/Kaiming init: std = √(2/n_in)
        nn.init.kaiming_normal_(module.weight, mode='fan_in', nonlinearity='relu')
        if module.bias is not None:
            nn.init.zeros_(module.bias)

    elif isinstance(module, nn.Embedding):
        # Normal dist for embeddings
        nn.init.normal_(module.weight, std=0.02)

    elif isinstance(module, (nn.LayerNorm, RMSNorm)):
        # Initialize scale to 1.0, no bias needed
        nn.init.ones_(module.weight)
        if hasattr(module, 'bias') and module.bias is not None:
            nn.init.zeros_(module.bias)

    else:
        if hasattr(module, 'weight'):
            nn.init.normal_(module.weight, std=0.02)

15. Multimodality: Vision, Audio & Beyond

Modern LLMs extend beyond text to process images, audio, video, and other modalities. This section covers the architectures and fusion strategies that enable foundation models to understand and reason across diverse data types.

Multimodal Encoder Architectures

# Image to patches (224x224 → 196 patches)
image: [B, 3, 224, 224]
patches: [B, 196, 768]  # 16×16 patches
with learnable [CLS] token: [B, 197, 768]
→ Transformer encoder
→ embeddings: [B, 768]

# DALL-E: image → tokens
image: [B, 3, 256, 256]
→ VAE encoder
→ latent: [B, 32, 32, 4]
→ discrete tokens (1-8192)
→ Transformer decoder

Cross-Modal Architectures

Audio & Speech Modalities

Multimodal Training Strategies

Practical Implementation: Vision-Language Adapter

class VisionLanguageAdapter:
    def __init__(self, vision_dim=1024, llm_dim=4096):
        self.vision_encoder = ViT(output_dim=vision_dim)
        self.adapter = nn.Sequential(
            nn.Linear(vision_dim, llm_dim // 2),
            nn.GELU(),
            nn.Linear(llm_dim // 2, llm_dim)
        )
        self.llm = LlamaForCausalLM.from_pretrained("llama-7b")
        # Freeze vision encoder, train adapter + LoRA on LLM
        for param in self.vision_encoder.parameters():
            param.requires_grad = False

    def forward(self, images, text_ids, text_mask):
        # 1. Encode images
        vision_features = self.vision_encoder(images)  # [B, N_patches, 1024]
        # 2. Project to LLM space
        vision_embeddings = self.adapter(vision_features)  # [B, N_patches, 4096]
        # 3. Concatenate with text embeddings
        text_embeddings = self.llm.embed_tokens(text_ids)
        # 4. Interleave or prepend vision tokens
        combined = torch.cat([vision_embeddings, text_embeddings], dim=1)
        # 5. Forward through LLM
        output = self.llm(inputs_embeds=combined, attention_mask=combined_mask)
        return output.logits

16. RAG Integration: Retrieval-Augmented Generation

RAG augments language models with external knowledge by retrieving relevant documents and conditioning generation on them. This section covers retrieval pipelines, vector databases, and orchestration frameworks.

The RAG Pipeline

RAG Frameworks & Tools

RAG Patterns

# Basic pattern
query_emb = embed(query)
docs = db.search(query_emb, k=5)
context = "\n".join([d.text for d in docs])
prompt = f"Context: {context}\nQ: {query}"
answer = llm.generate(prompt)

Practical Implementation: Basic RAG

import langchain
from langchain.vectorstores import Pinecone
from langchain.embeddings import OpenAIEmbeddings
from langchain.chains import RetrievalQA

# 1. Initialize components
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
vectorstore = Pinecone.from_documents(docs, embeddings, index_name="my-index")
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

# 2. Create QA chain
qa = RetrievalQA.from_chain_type(
    llm=ChatOpenAI(model="gpt-4"),
    chain_type="stuff",  # stuff, map_reduce, refine
    retriever=retriever
)

# 3. Query
response = qa.run("What are the key findings?")

17. Data Pipeline: From Raw Text to Training

Preparing high-quality data is critical to LLM performance. This section covers text cleaning, filtering, deduplication, and synthetic data generation—often 30-50% of project effort.

Data Processing Stages

def clean_text(text):
    # Remove null bytes, control characters
    text = text.replace("\x00", "")
    # Fix encoding artifacts
    text = ftfy.fix_text(text)
    # Normalize unicode (NFD → NFC)
    text = unicodedata.normalize("NFC", text)
    # Remove excessive whitespace
    text = re.sub(r"\s+", " ", text).strip()
    return text

def quality_score(document):
    score = 0
    # Language detection: keep en only
    if detect(document) != "en": return 0

    # Length heuristic
    if len(document.split()) < 50: return 0.1
    score += 0.3

    # Lexical diversity (unique unigrams / total)
    words = document.split()
    diversity = len(set(words)) / len(words)
    if diversity < 0.3: return 0  # Repetitive
    score += diversity * 0.3

    # Perplexity filter (reference LM)
    ppl = compute_perplexity(document)
    if ppl > 10000: return 0  # Gibberish
    score += 0.4

    return score

# MinHash for near-duplicate detection
from datasketch import MinHash

def dedup(documents):
    minhashes = {}
    for doc_id, text in documents:
        m = MinHash(num_perm=128)
        for token in text.split():
            m.update(token.encode())

        # Find similar existing documents
        duplicates = [d for d, mh in minhashes.items()
                      if m.jaccard(mh) > 0.8]
        if not duplicates:
            minhashes[doc_id] = m
            yield doc_id, text

Synthetic Data Generation

Reference: NVIDIA NeMo-Curator

# Large-scale data curation pipeline
from nemo_curator import Curator

curator = Curator(
    input_path="/data/raw",
    output_path="/data/processed",
    language="en",
    batch_size=10000,
    workers=128
)

# Multi-stage pipeline
curator.run([
    ("text_cleaning", {}),
    ("language_detection", {"lang": "en"}),
    ("quality_scoring", {"threshold": 0.5}),
    ("exact_dedup", {}),
    ("near_dedup", {"jaccard_threshold": 0.85}),
    ("pii_removal", {})
])

18. Deployment & Serving: From Lab to Production

Deploying LLMs requires careful consideration of latency, throughput, cost, and reliability. This section covers infrastructure, optimization frameworks, and CI/CD practices.

Deployment Strategies

Serving Engines & Optimization Frameworks

Practical Example: vLLM Deployment

from vllm import LLM, SamplingParams

# 1. Load model with paged attention
llm = LLM(
    model="meta-llama/Llama-2-70b-chat-hf",
    tensor_parallel_size=2,  # 2 GPUs
    gpu_memory_utilization=0.9,  # 90% VRAM
    dtype="float16",
    max_model_len=4096
)

# 2. Batch inference
prompts = ["What is AI?", "Explain transformers."]
sampling_params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=256)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)

CI/CD & Model Registry

Monitoring & Observability

19. Failure Modes: Understanding & Mitigating LLM Risks

LLMs are powerful but fallible. Understanding common failure modes and mitigations is critical for responsible deployment.

Hallucinations

Bias & Toxicity

Cost Overruns

Model Collapse (Repeated RLHF)

Data Leakage & Privacy

Latency Spikes & SLA Violations

Summary: Failure Mode Matrix

20. Implementation Roadmap: From Idea to Production

Building an LLM system is complex. This roadmap breaks it into five phases with realistic timelines and checkpoints.

Phase 1: Requirements & Evaluation (Weeks 1-2)

Phase 2: Prototyping (Weeks 3-6)

# Typical prototype stack
- Model: OpenAI GPT-4 or open-source (Llama 2)
- Framework: LangChain or Anthropic SDK
- Vector DB: Pinecone or Weaviate (for RAG)
- Eval: DeepEval, RAGAS, custom metrics
- Logging: Weights & Biases or Langfuse

Phase 3: Optimization & Tuning (Weeks 7-12)

Phase 4: Safety & Alignment (Weeks 10-14, parallel)

Phase 5: Production Deployment (Weeks 13-16)

Timeline & Dependencies

Resource & Skill Requirements

Appendix D: Glossary of Terms

Key Terminology

Appendix E: Memory & Compute Requirements

Training Memory Breakdown

# Memory consumption breakdown for training a 405B model

def estimate_training_memory(
    n_params=405e9, batch_size=4, seq_len=4096
):
    # 1. Model weights (FP32): 4 bytes per parameter
    weight_memory = n_params * 4 / 1e9  # GB

    # 2. Optimizer state (AdamW: m, v, exp_avg)
    optimizer_memory = n_params * 8 / 1e9  # 2× FP32 per param

    # 3. Activations (gradient computation)
    # Roughly: batch·seq_len·d_model × num_layers × 2
    activation_memory = (batch_size * seq_len * 16384 * 126 * 4) / 1e9

    # 4. Gradients (same shape as weights)
    gradient_memory = n_params * 4 / 1e9

    total = weight_memory + optimizer_memory + activation_memory + gradient_memory

    return {
        "weights": weight_memory,
        "optimizer": optimizer_memory,
        "activations": activation_memory,
        "gradients": gradient_memory,
        "total_gb": total
    }

# Estimate:
# Weights: 1.6TB
# Optimizer: 3.2TB
# Activations: 0.5TB (with gradient checkpointing)
# Gradients: 1.6TB
# Total: ~6.9TB per GPU (before distributed sharding)
# Solution: ZeRO-2 or ZeRO-3 reduces to ~2TB per GPU across cluster

Inference Memory & Latency

def estimate_inference_memory(
    n_params=405e9, batch_size=1, seq_len=4096,
    precision="FP32", use_kv_cache=True
):
    # Model weights
    bytes_per_param = {"FP32": 4, "FP16": 2, "INT8": 1, "INT4": 0.5}
    weight_bytes = n_params * bytes_per_param[precision]

    # KV cache (seq_len × batch × h × d_k)
    # Assume h=128, d_k=128 for 405B
    kv_cache_bytes = 0
    if use_kv_cache:
        kv_cache_bytes = seq_len * batch_size * 128 * 128 * 2 * 126 * 2  # 2 for K and V

    # Total
    total_gb = (weight_bytes + kv_cache_bytes) / 1e9

    return total_gb

# Example: 405B with INT4 quantization
# Weights: 405B × 0.5 = 202.5GB (fits on 8× 40GB A100s)
# KV cache: ~40GB (for seq_len=4096)

FLOPs Calculation

def calculate_training_flops(
    n_params, n_tokens, n_layers=126
):
    # Forward pass: 2·N·T FLOPs (matrix multiply + activation)
    forward_flops = 2 * n_params * n_tokens

    # Backward pass: 2× forward (gradient computation)
    backward_flops = 4 * n_params * n_tokens

    # Total = 6·N·T
    total_flops = 6 * n_params * n_tokens

    return {
        "forward": forward_flops,
        "backward": backward_flops,
        "total": total_flops
    }

# Example: LLaMA 3 405B on 15T tokens
# Total: 6 × 405B × 15T = 3.6×10^24 FLOPs
# On 16K H100s @ 1.5 PFLOP/s each = 24 EFLOP/s
# Time: 3.6×10^24 / (24×10^18) ≈ 150K seconds ≈ 41.7 hours

Latency Analysis

Appendix F: Model Compression Techniques

Knowledge Distillation

def knowledge_distillation_loss(
    student_logits, teacher_logits, labels, temperature=4.0, alpha=0.7
):
    # Distillation Loss = α·CE(student, labels) + (1-α)·KL(soft_targets)

    # Hard target loss (standard cross-entropy)
    ce_loss = torch.nn.functional.cross_entropy(student_logits, labels)

    # Soft target loss (KL divergence with temperature)
    # Temperature > 1 softens the distributions
    student_probs = torch.log_softmax(student_logits / temperature, dim=-1)
    teacher_probs = torch.softmax(teacher_logits / temperature, dim=-1)

    kl_loss = torch.nn.functional.kl_div(
        student_probs, teacher_probs, reduction='batchmean'
    )

    # Combined loss
    total_loss = alpha * ce_loss + (1 - alpha) * (temperature ** 2) * kl_loss

    return total_loss

# Benefits:
# - Smaller student model can match 95%+ of teacher performance
# - Temperature parameter controls knowledge transfer
# - Effective for any teacher-student size ratio

Layer Pruning

def prune_layers(model, prune_ratio=0.2):
    # Remove least important layers (compute layer importance)

    n_layers = len(model.layers)
    n_to_prune = int(n_layers * prune_ratio)

    # Compute layer importance via gradient norm or activation magnitude
    layer_importance = []
    for i, layer in enumerate(model.layers):
        # Importance ≈ ||weight gradient|| × ||activation||
        importance = torch.norm(layer.weight.grad) * torch.norm(layer.activation)
        layer_importance.append((i, importance))

    # Sort by importance and remove least important
    layer_importance.sort(key=lambda x: x[1])
    layers_to_remove = [x[0] for x in layer_importance[:n_to_prune]]

    # Create pruned model
    pruned_model = copy.deepcopy(model)
    pruned_model.layers = nn.ModuleList([
        layer for i, layer in enumerate(model.layers)
        if i not in layers_to_remove
    ])

    return pruned_model

# Results:
# - Removing 20% of layers: ~1-2% accuracy loss
# - 15-20% speedup from fewer layer computations

Weight Pruning & Sparsity

Appendix G: Fine-tuning Recipes & Best Practices

SFT (Supervised Fine-Tuning) Recipe

# High-quality SFT training configuration

"""
Base model: Mistral 7B
Data: 100K high-quality instruction-response pairs
Optimization:
  - Learning rate: 5e-5
  - Warmup steps: 500
  - Max gradient norm: 1.0
  - Optimizer: AdamW (β1=0.9, β2=0.999)
  - Weight decay: 0.01
  - Batch size: 64 (per-device)
  - Gradient accumulation: 1
  - Num epochs: 3

Hardware:
  - 8x A100 80GB
  - Total training time: ~6 hours

Results:
  - MMLU: 63.2% → 72.8% (+9.6pp)
  - MT-Bench: 6.2 → 7.8 (+1.6pp)
"""

def train_sft(model, train_dataset, val_dataset):
    # Prepare data loaders
    train_loader = DataLoader(
        train_dataset, batch_size=64, shuffle=True
    )
    val_loader = DataLoader(
        val_dataset, batch_size=64, shuffle=False
    )

    # Setup optimizer
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=5e-5,
        weight_decay=0.01
    )

    # Learning rate scheduler
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=len(train_loader) * 3
    )

    for epoch in range(3):
        for batch in train_loader:
            logits = model(batch["input_ids"])
            loss = compute_loss(logits, batch["labels"])

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

        # Validation
        val_loss = evaluate(model, val_loader)
        print(f"Epoch {epoch} val_loss: {val_loss:.3f}")

RLHF Pipeline (Simplified)

def rlhf_pipeline(base_model, sft_model, dataset):
    # Step 1: SFT → already have sft_model

    # Step 2: Train reward model
    reward_model = train_reward_model(
        base_model, dataset.comparisons
    )

    # Step 3: RL training with PPO
    rl_model = copy.deepcopy(sft_model)
    optimizer = torch.optim.AdamW(rl_model.parameters(), lr=1e-5)

    for epoch in range(5):
        for prompt in dataset.prompts:
            # Generate completions
            completions = rl_model.generate(
                prompt, num_return_sequences=4, temperature=1.0
            )

            # Get rewards
            rewards = reward_model(completions)

            # Compute advantages (REINFORCE or A2C)
            advantages = rewards - torch.mean(rewards)

            # Policy gradient update
            logprobs = rl_model.compute_logprob(completions)
            loss = -torch.mean(logprobs * advantages)

            loss.backward()
            torch.nn.utils.clip_grad_norm_(rl_model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()

    return rl_model

# RLHF complexity:
# - Requires reward model training (expensive)
# - RL training is unstable (needs careful tuning)
# - But produces best human alignment results

Common Hyperparameter Ranges

Appendix H: Prompt Engineering & In-Context Learning

Few-Shot Prompting Strategies

# Zero-shot prompt
question = """Q: Solve this math problem step by step.
Problem: If a train travels 100 km in 2 hours, what is its speed?
A:"""

# Few-shot prompt (2 examples)
question_fewshot = """
Example 1:
Q: A car travels 60 km in 1 hour. What is its speed?
A: The car travels 60 km in 1 hour. Speed = distance/time = 60/1 = 60 km/h.

Example 2:
Q: A boat travels 150 km in 3 hours. What is its speed?
A: The boat travels 150 km in 3 hours. Speed = distance/time = 150/3 = 50 km/h.

Now solve this:
Q: A train travels 100 km in 2 hours. What is its speed?
A:"""

# Chain-of-Thought (explicit reasoning)
cot_prompt = """Let's solve this step by step.
Q: If a train travels 100 km in 2 hours, what is its speed?
A: Let me break this down:
   - We know distance = 100 km
   - We know time = 2 hours
   - Speed = distance / time
   - Speed = 100 / 2 = 50 km/h
   Therefore, the train's speed is 50 km/h."""

Prompt Optimization Techniques

In-Context Learning Example

# Demonstrate that models learn from context without weight updates

prompt = """
You will now learn a pattern from examples. Apply it to the new input.

Example 1: "hello" → "olleh"
Example 2: "world" → "dlrow"
Example 3: "test" → "tset"

Now apply the pattern to: "machine"
Output:"""

# The model has learned the pattern from examples:
# Pattern: reverse the string
# Output: "enihcam"

# This is in-context learning:
# - No weight updates
# - Learning from prompt examples alone
# - Emergent ability in large models (>1B params)
# - Accuracy improves with more examples (scaling in prompts)

Appendix I: Benchmarking & Evaluation Methodology

Proper Evaluation Setup

def evaluate_model_proper(model, benchmark_dataset):
    # 1. Ensure test set is separate (no data leakage)
    assert benchmark_dataset.split == "test"

    # 2. Set to evaluation mode (disable dropout, etc.)
    model.eval()

    # 3. Disable gradients (faster inference)
    with torch.no_grad():
        predictions = []
        for batch in benchmark_dataset:
            logits = model(batch["input_ids"])
            pred = torch.argmax(logits, dim=-1)
            predictions.append(pred)

    # 4. Compute metrics
    accuracy = (predictions == benchmark_dataset.labels).mean()
    f1 = compute_f1(predictions, benchmark_dataset.labels)

    # 5. Report with confidence intervals
    return {
        "accuracy": accuracy,
        "f1": f1,
        "n_samples": len(benchmark_dataset)
    }

# Common mistakes to avoid:
# 1. Using training data in evaluation (leakage)
# 2. Tuning on test set (overfitting)
# 3. Not reporting uncertainty/variance
# 4. Cherry-picking metrics (report multiple)
# 5. Not controlling for randomness (set seed, multiple runs)

Metric Definitions

Computing Perplexity

def compute_perplexity(model, dataset, batch_size=32):
    # Perplexity = exp(average cross-entropy loss)
    model.eval()
    total_loss = 0.0
    total_tokens = 0

    with torch.no_grad():
        for batch in DataLoader(dataset, batch_size=batch_size):
            logits = model(batch["input_ids"])
            # Shift: predict token i+1 from token i
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = batch["input_ids"][:, 1:].contiguous()

            # Compute loss
            loss = torch.nn.functional.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1)
            )

            total_loss += loss.item() * shift_labels.numel()
            total_tokens += shift_labels.numel()

    avg_loss = total_loss / total_tokens
    perplexity = math.exp(avg_loss)
    return perplexity

# Interpretation:
# - Perplexity = 10: model is ~10x better than random guessing
# - GPT-2 (small): ~20-30 on test sets
# - GPT-3: ~8-15 on test sets
# - State-of-the-art: ~4-8

Appendix J: Model Deployment & Serving

Efficient Model Serving

# Model serving with batching and KV cache management
import asyncio

class ModelServer:
    def __init__(self, model, batch_size=32, max_wait_ms=10):
        self.model = model
        self.batch_size = batch_size
        self.max_wait_ms = max_wait_ms
        self.request_queue = asyncio.Queue()

    async def add_request(self, prompt, max_tokens=256):
        # Add request to queue
        future = asyncio.Future()
        await self.request_queue.put({
            "prompt": prompt,
            "max_tokens": max_tokens,
            "future": future
        })
        return await future

    async def batch_processing_loop(self):
        # Process batches of requests
        while True:
            # Wait for batch or timeout
            batch = []
            start_time = time.time()

            try:
                while len(batch) < self.batch_size:
                    timeout = self.max_wait_ms / 1000.0
                    request = await asyncio.wait_for(
                        self.request_queue.get(),
                        timeout=timeout
                    )
                    batch.append(request)
            except asyncio.TimeoutError:
                pass

            if not batch:
                continue

            # Process batch in parallel
            prompts = [r["prompt"] for r in batch]
            outputs = await asyncio.gather(*[
                self.model.generate(p, 256)
                for p in prompts
            ])

            # Return results to requesters
            for request, output in zip(batch, outputs):
                request["future"].set_result(output)

# Batching benefits:
# - GPUs are 10-50× faster with batching
# - Trade-off: latency vs throughput
# - Dynamic batching: request-level optimization

Deployment Hardware Considerations

Appendix K: Advanced Architecture Variants

Vision Transformer (ViT) Integration

class MultimodalLLM(nn.Module):
    def __init__(self, vision_encoder, language_model):
        super().__init__()
        self.vision_encoder = vision_encoder  # ViT
        self.language_model = language_model  # LLM
        # Projection from vision to language space
        self.vision_to_lang = nn.Linear(768, 4096)

    def forward(self, images, text_ids):
        # Process images with vision encoder
        # images: (batch, 3, 224, 224)
        image_embeddings = self.vision_encoder(images)
        # → (batch, num_patches, 768)

        # Project to language model space
        image_features = self.vision_to_lang(image_embeddings)
        # → (batch, num_patches, 4096)

        # Process text with language model
        text_embeddings = self.language_model.embedding(text_ids)

        # Interleave vision and language tokens
        # Option 1: Prefix - prepend image tokens before text
        combined = torch.cat([image_features, text_embeddings], dim=1)

        # Process with transformer
        output = self.language_model.transformer(combined)
        logits = self.language_model.head(output)

        return logits

# Multimodal integration strategies:
# 1. Prefix: image tokens before text (simplest)
# 2. Interleaved: alternate image-text tokens (flexible)
# 3. Cross-attention: separate vision processing with attention to text
# 4. Adapter: small projection network between encoders

Mamba SSM Architecture

class MambaBlock(nn.Module):
    # State-Space Model: Linear O(n) attention alternative
    def __init__(self, d_model, d_state=16):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state

        # SSM parameters
        self.A = nn.Parameter(torch.randn(d_model, d_state))
        self.B = nn.Linear(d_model, d_state)
        self.C = nn.Linear(d_model, d_state)

        # Input/output projections
        self.D_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        # x: (batch, seq_len, d_model)

        # Compute B and C at each step (data-dependent)
        B = self.B(x)  # (batch, seq, d_state)
        C = self.C(x)  # (batch, seq, d_state)

        # State: h_t = A @ h_{t-1} + B_t @ x_t
        # Output: y_t = C_t @ h_t

        h = torch.zeros(x.shape[0], self.d_state, device=x.device)
        outputs = []

        for t in range(x.shape[1]):
            # State update
            h = self.A.unsqueeze(0) @ h.unsqueeze(-1) + \
                (B[:, t] * x[:, t]).unsqueeze(-1)
            h = h.squeeze(-1)

            # Output
            y = (C[:, t] * h).sum(dim=-1)
            outputs.append(y)

        output = torch.stack(outputs, dim=1)
        return self.out_proj(output)

# Mamba advantages:
# - O(n) complexity vs O(n²) for attention
# - Linear scaling with sequence length
# - Better for very long sequences (>1M tokens)
# - Trade-off: slightly lower quality than transformers

Appendix L: Active Research Directions

Research Areas & Recent Work

Training Objective Innovations

# Next-token prediction + auxiliary losses

def compute_hybrid_loss(logits, labels, hidden_states):
    # Main objective: cross-entropy on next token
    ce_loss = torch.nn.functional.cross_entropy(
        logits.view(-1, logits.size(-1)),
        labels.view(-1)
    )

    # Auxiliary: contrastive learning on representations
    # Encourage similar tokens to have similar representations
    sim_matrix = torch.nn.functional.cosine_similarity(
        hidden_states.unsqueeze(1),
        hidden_states.unsqueeze(0),
        dim=-1
    )
    contrastive_loss = compute_nt_xent_loss(sim_matrix, labels)

    # Auxiliary: auxiliary prediction heads
    # Predict sentence boundaries, document structure, etc.
    boundary_pred = self.boundary_head(hidden_states)
    boundary_loss = torch.nn.functional.binary_cross_entropy(
        boundary_pred, boundary_labels
    )

    # Combined loss
    total_loss = ce_loss + 0.1 * contrastive_loss + 0.05 * boundary_loss

    return total_loss

# Auxiliary losses benefits:
# - Improve convergence speed (~10% faster)
# - Better generalization on downstream tasks
# - Help with long-range dependencies