Large Language Models (LLMs)

Self-Attention & Multi-Head Attention

Python
import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = 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)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        B, T, C = x.shape
        # Project and reshape to [B, n_heads, T, d_k]
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1,2)
        K = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1,2)
        V = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1,2)

        # Scaled dot-product attention
        scores = (Q @ K.transpose(-2,-1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = self.dropout(torch.softmax(scores, dim=-1))
        out = (attn @ V).transpose(1,2).contiguous().view(B, T, C)
        return self.W_o(out)

# Causal mask for autoregressive (GPT-style)
def causal_mask(T):
    return torch.tril(torch.ones(T, T)).unsqueeze(0).unsqueeze(0)

Positional Encodings: Sinusoidal, RoPE & ALiBi

Python
import torch, math

# 1. Sinusoidal (Original Transformer)
def sinusoidal_pe(max_len, d_model):
    pe = torch.zeros(max_len, d_model)
    pos = torch.arange(max_len).unsqueeze(1).float()
    div = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000) / d_model))
    pe[:, 0::2] = torch.sin(pos * div)
    pe[:, 1::2] = torch.cos(pos * div)
    return pe

# 2. RoPE — Rotary Position Embeddings (LLaMA, GPT-NeoX)
def apply_rope(q, k, positions):
    """Rotate query/key vectors by position-dependent angles"""
    d = q.shape[-1]
    freqs = 1.0 / (10000 ** (torch.arange(0, d, 2).float() / d))
    angles = positions.unsqueeze(-1) * freqs
    cos_a, sin_a = torch.cos(angles), torch.sin(angles)

    # Rotate pairs of dimensions
    q_rot = torch.stack([q[..., 0::2]*cos_a - q[..., 1::2]*sin_a,
                         q[..., 0::2]*sin_a + q[..., 1::2]*cos_a], dim=-1).flatten(-2)
    k_rot = torch.stack([k[..., 0::2]*cos_a - k[..., 1::2]*sin_a,
                         k[..., 0::2]*sin_a + k[..., 1::2]*cos_a], dim=-1).flatten(-2)
    return q_rot, k_rot

# 3. ALiBi — Attention with Linear Biases (no embeddings!)
# Simply adds a linear bias to attention scores based on distance:
# score(i,j) = q_i · k_j - m · |i - j|
# where m is a head-specific slope. No learned parameters!

Layer Normalization, FFN & KV Cache

Python
class TransformerBlock(nn.Module):
    """Pre-LN Transformer Block (standard in GPT, LLaMA)"""
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)  # Pre-LN: normalize BEFORE attention
        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),                     # SwiGLU in LLaMA
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )

    def forward(self, x, mask=None):
        x = x + self.attn(self.ln1(x), mask)   # Residual + Pre-LN Attention
        x = x + self.ffn(self.ln2(x))           # Residual + Pre-LN FFN
        return x

KV Cache — Critical for Fast Inference

Python
class CachedAttention(nn.Module):
    """Attention with KV cache for autoregressive generation"""
    def forward(self, x, kv_cache=None):
        Q = self.W_q(x)  # Only compute Q for new tokens
        K_new = self.W_k(x)
        V_new = self.W_v(x)

        if kv_cache is not None:
            K = torch.cat([kv_cache[0], K_new], dim=1)  # Append to cache
            V = torch.cat([kv_cache[1], V_new], dim=1)
        else:
            K, V = K_new, V_new

        # Attention with full K,V but only new Q
        out = scaled_dot_product(Q, K, V)
        return out, (K, V)  # Return updated cache
    # Without cache: generating 1000 tokens recomputes ALL K,V for each token
    # With cache: each step is O(1) compute instead of O(n)

Encoder, Decoder & Encoder-Decoder Variants

Tokenization: BPE & SentencePiece

Python
# tiktoken — OpenAI's tokenizer (used in GPT-4)
import tiktoken

enc = tiktoken.get_encoding("cl100k_base")  # GPT-4 encoding
text = "Large Language Models are transforming AI"
tokens = enc.encode(text)
print(f"Text: {text}")
print(f"Tokens: {tokens}")
print(f"Token count: {len(tokens)}")
print(f"Decoded: {[enc.decode([t]) for t in tokens]}")

# BPE from scratch (simplified)
def get_pairs(tokens):
    return {(tokens[i], tokens[i+1]) for i in range(len(tokens)-1)}

def bpe_train(text, num_merges):
    tokens = list(text.encode("utf-8"))
    merges = {}
    for i in range(num_merges):
        pairs = {}
        for j in range(len(tokens)-1):
            pair = (tokens[j], tokens[j+1])
            pairs[pair] = pairs.get(pair, 0) + 1
        if not pairs: break
        best = max(pairs, key=pairs.get)
        new_token = 256 + i
        merges[best] = new_token
        # Replace all occurrences of best pair with new token
        new_tokens = []
        j = 0
        while j < len(tokens):
            if j < len(tokens)-1 and (tokens[j], tokens[j+1]) == best:
                new_tokens.append(new_token); j += 2
            else:
                new_tokens.append(tokens[j]); j += 1
        tokens = new_tokens
    return merges, tokens

Language Modeling Objectives

Python
# Causal Language Modeling (GPT-style)
# Given: "The cat sat on the"
# Predict: "cat" "sat" "on" "the" "mat"

import torch.nn.functional as F

def causal_lm_loss(logits, targets):
    # logits: [B, T, vocab_size], targets: [B, T]
    # Shift: predict token t+1 from position t
    shift_logits = logits[:, :-1].contiguous()
    shift_labels = targets[:, 1:].contiguous()
    return F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1))

# Every position provides a training signal!
# A 4096-token document gives 4095 prediction tasks
# This efficiency is why causal LM scales so well

Data Collection, Cleaning & Deduplication

Python
# Data pipeline stages
def data_pipeline(raw_docs):
    # 1. Language filtering
    docs = [d for d in raw_docs if detect_language(d) == "en"]

    # 2. Quality filtering (perplexity-based)
    docs = [d for d in docs if quality_score(d) > 0.5]

    # 3. Deduplication (MinHash + LSH)
    docs = minhash_dedup(docs, threshold=0.8)

    # 4. Toxicity/PII removal
    docs = [remove_pii(d) for d in docs if not is_toxic(d)]

    # 5. Tokenize and pack into sequences
    tokens = tokenizer.encode_batch(docs)
    return pack_sequences(tokens, max_len=4096)

Compute Scaling & Scaling Laws

Supervised Fine-Tuning (SFT)

Python
# SFT training data format
sft_data = [
    {
        "instruction": "Explain quantum computing in simple terms",
        "response": "Quantum computing uses quantum bits (qubits) that can be both 0 and 1 simultaneously (superposition). This allows quantum computers to explore many solutions at once, making them exponentially faster for certain problems like cryptography and drug discovery..."
    },
    {
        "instruction": "Write a Python function to find prime numbers",
        "response": "def is_prime(n):\n    if n < 2: return False\n    for i in range(2, int(n**0.5)+1):\n        if n % i == 0: return False\n    return True"
    }
]

# Training: only compute loss on the response tokens!
def sft_loss(logits, labels, instruction_mask):
    # Mask out instruction tokens — don't train on the question
    loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                           labels.view(-1), reduction='none')
    loss = loss * (~instruction_mask).float().view(-1)
    return loss.sum() / (~instruction_mask).sum()

Python
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from datasets import load_dataset
from peft import LoraConfig, get_peft_model

# Load model and tokenizer
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)

# Add LoRA for efficient fine-tuning (Chapter 11)
lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj","v_proj"])
model = get_peft_model(model, lora_config)
print(f"Trainable params: {model.print_trainable_parameters()}")

# Load instruction dataset
dataset = load_dataset("tatsu-lab/alpaca")

# Training
training_args = TrainingArguments(
    output_dir="./sft_output", num_train_epochs=3,
    per_device_train_batch_size=4, gradient_accumulation_steps=8,
    learning_rate=2e-5, fp16=True, warmup_ratio=0.1)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset["train"])
trainer.train()

RLHF & DPO — Aligning with Human Preferences

RLHF Pipeline

Python
# DPO — Direct Preference Optimization (simpler than RLHF)
# Given pairs: (prompt, chosen_response, rejected_response)

import torch.nn.functional as F

def dpo_loss(policy_chosen_logps, policy_rejected_logps,
             ref_chosen_logps, ref_rejected_logps, beta=0.1):
    """
    policy_chosen_logps: log P_θ(chosen | prompt)
    ref_chosen_logps: log P_ref(chosen | prompt) — frozen reference model
    beta: temperature parameter
    """
    chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps)
    rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps)

    # The model should assign higher reward to chosen vs rejected
    loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
    return loss

# DPO advantage: no separate reward model, no PPO complexity!
# Trains directly on preference pairs: (prompt, winner, loser)

LoRA, QLoRA & PEFT Methods

Python
import torch.nn as nn

class LoRALayer(nn.Module):
    """Low-Rank Adaptation: W_new = W_frozen + A·B (rank r)"""
    def __init__(self, original_layer, r=16, alpha=32):
        super().__init__()
        self.original = original_layer
        self.original.weight.requires_grad = False  # Freeze original
        d_in, d_out = original_layer.weight.shape
        self.A = nn.Parameter(torch.randn(d_in, r) * 0.01)   # Down-project
        self.B = nn.Parameter(torch.zeros(r, d_out))          # Up-project
        self.scaling = alpha / r

    def forward(self, x):
        # Original output + low-rank adaptation
        return self.original(x) + (x @ self.A @ self.B) * self.scaling

# For LLaMA-7B:
# Full fine-tuning: 7B trainable params, needs 8× A100 80GB
# LoRA (r=16): ~4M trainable params (0.06%), needs 1× A100
# QLoRA (4-bit): same quality, needs 1× RTX 3090 24GB!

Constitutional AI & Instruction Following

Python
# Constitutional AI pipeline (Anthropic)
principles = [
    "Choose the response that is most helpful while being harmless.",
    "Choose the response that is most honest and factual.",
    "Avoid responses that are discriminatory or biased.",
]

def constitutional_critique(model, prompt, response):
    """Use the model itself to critique and revise responses"""
    critique_prompt = f"""
Critique this response based on the principle: {principles[0]}

Prompt: {prompt}
Response: {response}

Identify problems and suggest a revised response:"""
    critique = model.generate(critique_prompt)
    return critique

# Process:
# 1. Generate response to harmful prompt
# 2. Self-critique against principles
# 3. Revise response based on critique
# 4. Use (original, revised) as DPO preference pair
# = RLAIF — no human annotators needed!

Quantization, Pruning & Distillation

Python
# Quantization with bitsandbytes
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# 4-bit quantization (GPTQ/AWQ style)
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",            # NormalFloat4
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True       # Quantize the quantization constants!
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-70b-hf",
    quantization_config=quant_config,
    device_map="auto"
)
# 70B model: FP16 = 140 GB → INT4 = ~35 GB → fits on 2× RTX 3090!

Serving: Flash Attention, Parallelism & vLLM

Flash Attention

Python
# Flash Attention — fused kernel, O(n) memory instead of O(n²)
# Standard attention materializes the full n×n attention matrix
# Flash Attention computes attention in tiles, never storing the full matrix

# In PyTorch 2.0+:
from torch.nn.functional import scaled_dot_product_attention

# Automatically uses Flash Attention if available
output = scaled_dot_product_attention(Q, K, V, is_causal=True)

# For LLaMA with 128K context:
# Standard attention memory: 128K × 128K × 2 bytes = 32 GB per layer
# Flash Attention memory: ~O(n) = ~megabytes per layer
# Speedup: 2-4x on A100, enables 128K+ context windows

Parallelism Strategies

Serving with vLLM

Python
# vLLM — high-throughput LLM serving with PagedAttention
from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-2-7b-chat-hf",
           tensor_parallel_size=2,     # 2 GPUs
           gpu_memory_utilization=0.9)

params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=512)

# vLLM uses PagedAttention — KV cache is managed like virtual memory
# This eliminates memory waste from pre-allocated KV caches
# Throughput improvement: 2-4x over HuggingFace, 24x over naive approach

prompts = [
    "Explain machine learning in simple terms",
    "Write a Python quicksort function",
    "What causes climate change?",
]
outputs = llm.generate(prompts, params)  # Batched! Continuous batching
for out in outputs:
    print(out.outputs[0].text)

Prompt Engineering & Context Windows

Python
# System prompt engineering
system_prompt = """You are a senior financial analyst. Follow these rules:
1. Always cite your sources with specific data points
2. When uncertain, explicitly state your confidence level
3. Present both bullish and bearish perspectives
4. Use tables for comparisons"""

# RAG: Retrieve relevant context to fit in window
def rag_prompt(query, retrieved_docs, max_context=4000):
    context = "\n---\n".join(retrieved_docs[:max_context])
    return f"""Based on the following documents, answer the question.

Context:
{context}

Question: {query}

Answer based ONLY on the provided context. If the answer is not in the context, say so."""

Bash
# Start vLLM server (OpenAI-compatible API)
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-7b-chat-hf \
    --tensor-parallel-size 2 \
    --max-model-len 4096 \
    --port 8000

# Client code (works with any OpenAI SDK!)

Python
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")

response = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-chat-hf",
    messages=[{"role": "user", "content": "Explain transformers in 100 words"}],
    max_tokens=200, temperature=0.7
)
print(response.choices[0].message.content)