From cc5366624af65b7a80c9ea3e88b80bf88f1b5011 Mon Sep 17 00:00:00 2001 From: Madhur Panwar <39766613+mdrpanwar@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:43:00 +0200 Subject: [PATCH 1/2] Add KL divergence and SFT SmolLM problems New v3/alignment-training problems (kl-divergence, sft-smollm) with notebooks, solutions, grader specs, and regenerated indexes/docs. Co-Authored-By: Claude Opus 5.5 (1M context) --- PLANNED.toml | 10 - PROBLEMS.md | 18 +- README.md | 13 +- SOURCES.md | 6 +- STUDYPLAN.md | 18 +- mcp-server/src/data/hints.generated.ts | 20 + mcp-server/src/data/questions.ts | 42 +- mcp-server/src/tools/get-prerequisites.ts | 6 +- problems.json | 181 +++- python/test_grader.py | 166 ++++ python/torchleet/problems/kl_divergence.py | 242 +++++ python/torchleet/problems/sft_smollm.py | 181 ++++ space/problems.json | 181 +++- .../kl-divergence/kl-divergence.ipynb | 801 +++++++++++++++++ .../kl-divergence/kl-divergence_SOLN.ipynb | 830 ++++++++++++++++++ .../kl-divergence/problem.toml | 30 + v3/alignment-training/sft-smollm/problem.toml | 30 + .../sft-smollm/sft-smollm.ipynb | 623 +++++++++++++ .../sft-smollm/sft-smollm_SOLN.ipynb | 651 ++++++++++++++ .../data-parallel-loading/problem.toml | 17 +- website/src/data/problems.generated.ts | 8 +- 21 files changed, 4009 insertions(+), 65 deletions(-) create mode 100644 python/torchleet/problems/kl_divergence.py create mode 100644 python/torchleet/problems/sft_smollm.py create mode 100644 v3/alignment-training/kl-divergence/kl-divergence.ipynb create mode 100644 v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb create mode 100644 v3/alignment-training/kl-divergence/problem.toml create mode 100644 v3/alignment-training/sft-smollm/problem.toml create mode 100644 v3/alignment-training/sft-smollm/sft-smollm.ipynb create mode 100644 v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb diff --git a/PLANNED.toml b/PLANNED.toml index c5e4cbe..c8220d5 100644 --- a/PLANNED.toml +++ b/PLANNED.toml @@ -46,11 +46,6 @@ id = "v1-35" title = "Implement a Variational Autoencoder" difficulty = "hard" -[[planned]] -id = "v2-1" -title = "Implement KL Divergence Loss" -difficulty = "easy" - [[planned]] id = "v2-4" title = "Create a RAG Search of Embeddings" @@ -60,8 +55,3 @@ difficulty = "medium" id = "v2-13" title = "Implement Quantization (GPTQ)" difficulty = "hard" - -[[planned]] -id = "v2-20" -title = "Apply SFT on SmolLM" -difficulty = "hard" diff --git a/PROBLEMS.md b/PROBLEMS.md index e8e848d..718f93f 100644 --- a/PROBLEMS.md +++ b/PROBLEMS.md @@ -1,15 +1,19 @@ -# All 68 problems +# All 74 problems Generated from the per-problem `problem.toml` manifests β€” do not edit by hand. Run `python3 scripts/generate.py readme` to refresh. -**63 of 68 are auto-graded** (`pip install torchleet`). A βœ… in the +**65 of 74 are auto-graded** (`pip install torchleet`). A βœ… in the Graded column means `check()` will verify your solution. -### Basics (23) +### Basics (27) | Problem | Difficulty | Companies | Graded | Run | Links | |---|---|---|---|---|---| +| πŸ’‘ Implement a Shuffle Buffer for Streaming Data | medium | | | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/buffered-loading/buffered-loading.ipynb) | [Q](v3/data-loading/buffered-loading/buffered-loading.ipynb) / [S](v3/data-loading/buffered-loading/buffered-loading_SOLN.ipynb) | +| ⭐ Simulate DistributedSampler for Data Parallel Training | medium | Meta, Google, NVIDIA, Anthropic | | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/data-parallel-loading/data-parallel-loading.ipynb) | [Q](v3/data-loading/data-parallel-loading/data-parallel-loading.ipynb) / [S](v3/data-loading/data-parallel-loading/data-parallel-loading_SOLN.ipynb) | +| ⭐ Implement a Preloaded Map-Style Dataset | easy | | | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/preloaded-dataset/preloaded-dataset.ipynb) | [Q](v3/data-loading/preloaded-dataset/preloaded-dataset.ipynb) / [S](v3/data-loading/preloaded-dataset/preloaded-dataset_SOLN.ipynb) | +| ⭐ Implement a Streaming IterableDataset with Yield | medium | | | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/streaming-dataset/streaming-dataset.ipynb) | [Q](v3/data-loading/streaming-dataset/streaming-dataset.ipynb) / [S](v3/data-loading/streaming-dataset/streaming-dataset_SOLN.ipynb) | | πŸ’‘ Implement Linear Regression | basic | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/torch/basic/lin-regression/lin-regression.ipynb) | [Q](torch/basic/lin-regression/lin-regression.ipynb) / [S](torch/basic/lin-regression/lin-regression_SOLN.ipynb) | | πŸ’‘ Write a Custom Dataset and DataLoader | basic | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/torch/basic/custom-dataset/custom-dataset.ipynb) | [Q](torch/basic/custom-dataset/custom-dataset.ipynb) / [S](torch/basic/custom-dataset/custom-dataset_SOLN.ipynb) | | πŸ’‘ Write a Custom Activation Function | basic | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/torch/basic/custom-activation/custom-activation.ipynb) | [Q](torch/basic/custom-activation/custom-activation.ipynb) / [S](torch/basic/custom-activation/custom-activation_SOLN.ipynb) | @@ -34,10 +38,11 @@ Graded column means `check()` will verify your solution. | ⭐ Implement KNN in PyTorch | easy | Uber, LinkedIn, Meta | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/classical-ml/knn/knn.ipynb) | [Q](v3/classical-ml/knn/knn.ipynb) / [S](v3/classical-ml/knn/knn_SOLN.ipynb) | | ⭐ Implement Logistic Regression with Gradient Descent | easy | Google, Meta, Amazon | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/classical-ml/logistic-regression/logistic-regression.ipynb) | [Q](v3/classical-ml/logistic-regression/logistic-regression.ipynb) / [S](v3/classical-ml/logistic-regression/logistic-regression_SOLN.ipynb) | -### LLM Learning Path (21) +### LLM Learning Path (23) | Problem | Difficulty | Companies | Graded | Run | Links | |---|---|---|---|---|---| +| Implement KL Divergence Loss | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/kl-divergence/kl-divergence.ipynb) | [Q](v3/alignment-training/kl-divergence/kl-divergence.ipynb) / [S](v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb) | | πŸ’‘ Implement RMS Norm | easy | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/RMS-Norm/rms-norm.ipynb) | [Q](llm/RMS-Norm/rms-norm.ipynb) / [S](llm/RMS-Norm/rms-norm_SOLN.ipynb) | | πŸ’‘ Implement Byte Pair Encoding from Scratch | easy | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Byte-Pair-Encoder/BPE-q3.ipynb) | [Q](llm/Byte-Pair-Encoder/BPE-q3.ipynb) / [S](llm/Byte-Pair-Encoder/BPE-q3-SOLN.ipynb) | | πŸ’‘ Implement Attention from Scratch | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Implement-Attention-from-Scratch/attention-q4-Question.ipynb) | [Q](llm/Implement-Attention-from-Scratch/attention-q4-Question.ipynb) / [S](llm/Implement-Attention-from-Scratch/attention-q4.ipynb) | @@ -46,6 +51,7 @@ Graded column means `check()` will verify your solution. | πŸ’‘ Implement Sinusoidal Embeddings | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Sinusoidal-Positional-Embedding/sinusoidal-q7-Question.ipynb) | [Q](llm/Sinusoidal-Positional-Embedding/sinusoidal-q7-Question.ipynb) / [S](llm/Sinusoidal-Positional-Embedding/sinusoidal-q7.ipynb) | | πŸ’‘ Implement ROPE Embeddings | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Rotary-Positional-Embedding/rope-q8-Question.ipynb) | [Q](llm/Rotary-Positional-Embedding/rope-q8-Question.ipynb) / [S](llm/Rotary-Positional-Embedding/rope-q8.ipynb) | | πŸ’‘ Implement SmolLM from Scratch | hard | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/SmolLM/smollm-q12-Question.ipynb) | [Q](llm/SmolLM/smollm-q12-Question.ipynb) / [S](llm/SmolLM/smollm-q12.ipynb) | +| Apply SFT on SmolLM | hard | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/sft-smollm/sft-smollm.ipynb) | [Q](v3/alignment-training/sft-smollm/sft-smollm.ipynb) / [S](v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb) | | πŸ”₯ Implement Top-p (Nucleus) Sampling | medium | Anthropic, OpenAI, DeepMind, Perplexity, Cohere | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/llm-decoding/top-p-sampling/top-p-sampling.ipynb) | [Q](v3/llm-decoding/top-p-sampling/top-p-sampling.ipynb) / [S](v3/llm-decoding/top-p-sampling/top-p-sampling_SOLN.ipynb) | | πŸ”₯ Implement Top-k Sampling | medium | Anthropic, OpenAI, DeepMind, Cohere | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/llm-decoding/top-k-sampling/top-k-sampling.ipynb) | [Q](v3/llm-decoding/top-k-sampling/top-k-sampling.ipynb) / [S](v3/llm-decoding/top-k-sampling/top-k-sampling_SOLN.ipynb) | | πŸ”₯ Implement Temperature Sampling | easy | OpenAI, Anthropic, Cohere, Perplexity | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/llm-decoding/temperature-sampling/temperature-sampling.ipynb) | [Q](v3/llm-decoding/temperature-sampling/temperature-sampling.ipynb) / [S](v3/llm-decoding/temperature-sampling/temperature-sampling_SOLN.ipynb) | @@ -60,7 +66,7 @@ Graded column means `check()` will verify your solution. | ⭐ Implement GRPO (DeepSeek-R1 Algorithm) | expert | DeepMind, Anthropic, OpenAI | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/grpo/grpo.ipynb) | [Q](v3/alignment-training/grpo/grpo.ipynb) / [S](v3/alignment-training/grpo/grpo_SOLN.ipynb) | | πŸ”₯ Build a Complete LLM Inference Engine | expert | Perplexity, Together AI, Anyscale, Fireworks AI | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/llm-inference/inference-engine/inference-engine.ipynb) | [Q](v3/llm-inference/inference-engine/inference-engine.ipynb) / [S](v3/llm-inference/inference-engine/inference-engine_SOLN.ipynb) | -### Advanced (42) +### Advanced (44) | Problem | Difficulty | Companies | Graded | Run | Links | |---|---|---|---|---|---| @@ -76,12 +82,14 @@ Graded column means `check()` will verify your solution. | ⭐ Implement Muon from Scratch | hard | Meta, Google | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/torch/hard/optimizers/muon/muon.ipynb) | [Q](torch/hard/optimizers/muon/muon.ipynb) / [S](torch/hard/optimizers/muon/muon_SOLN.ipynb) | | πŸ’‘ Create Embeddings Out Of An Llm | medium | | | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Create-Embeddings-out-of-an-LLM/embeddings-q2.ipynb) | [Q](llm/Create-Embeddings-out-of-an-LLM/embeddings-q2.ipynb) | | πŸ’‘ Flash Attention | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Flash-Attention/flash-attention.ipynb) | [Q](llm/Flash-Attention/flash-attention.ipynb) | +| Implement KL Divergence Loss | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/kl-divergence/kl-divergence.ipynb) | [Q](v3/alignment-training/kl-divergence/kl-divergence.ipynb) / [S](v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb) | | πŸ’‘ Implement Attention from Scratch | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Implement-Attention-from-Scratch/attention-q4-Question.ipynb) | [Q](llm/Implement-Attention-from-Scratch/attention-q4-Question.ipynb) / [S](llm/Implement-Attention-from-Scratch/attention-q4.ipynb) | | πŸ’‘ Implement Multi-Head Attention from Scratch | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Multi-Head-Attention/multi-head-attention-q5-Question.ipynb) | [Q](llm/Multi-Head-Attention/multi-head-attention-q5-Question.ipynb) / [S](llm/Multi-Head-Attention/multi-head-attention-q5.ipynb) | | πŸ’‘ Implement Grouped Query Attention from Scratch | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Grouped-Query-Attention/grouped-query-attention-Question.ipynb) | [Q](llm/Grouped-Query-Attention/grouped-query-attention-Question.ipynb) / [S](llm/Grouped-Query-Attention/grouped-query-attention.ipynb) | | πŸ’‘ Implement Sinusoidal Embeddings | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Sinusoidal-Positional-Embedding/sinusoidal-q7-Question.ipynb) | [Q](llm/Sinusoidal-Positional-Embedding/sinusoidal-q7-Question.ipynb) / [S](llm/Sinusoidal-Positional-Embedding/sinusoidal-q7.ipynb) | | πŸ’‘ Implement ROPE Embeddings | medium | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Rotary-Positional-Embedding/rope-q8-Question.ipynb) | [Q](llm/Rotary-Positional-Embedding/rope-q8-Question.ipynb) / [S](llm/Rotary-Positional-Embedding/rope-q8.ipynb) | | πŸ’‘ Implement SmolLM from Scratch | hard | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/SmolLM/smollm-q12-Question.ipynb) | [Q](llm/SmolLM/smollm-q12-Question.ipynb) / [S](llm/SmolLM/smollm-q12.ipynb) | +| Apply SFT on SmolLM | hard | | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/sft-smollm/sft-smollm.ipynb) | [Q](v3/alignment-training/sft-smollm/sft-smollm.ipynb) / [S](v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb) | | πŸ”₯ Implement Contrastive Loss (InfoNCE) + CLIP Training Loop | medium | OpenAI, Anthropic, DeepMind, Midjourney, Apple | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/modern-architectures/contrastive-loss-clip/contrastive-loss-clip.ipynb) | [Q](v3/modern-architectures/contrastive-loss-clip/contrastive-loss-clip.ipynb) / [S](v3/modern-architectures/contrastive-loss-clip/contrastive-loss-clip_SOLN.ipynb) | | πŸ”₯ Implement 2D Positional Embeddings | medium | Anthropic, DeepMind, Midjourney, Runway | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/modern-architectures/2d-positional-embeddings/2d-positional-embeddings.ipynb) | [Q](v3/modern-architectures/2d-positional-embeddings/2d-positional-embeddings.ipynb) / [S](v3/modern-architectures/2d-positional-embeddings/2d-positional-embeddings_SOLN.ipynb) | | πŸ”₯ Implement Beam Search for LLM Decoding | medium | Google, DeepMind, Meta, Apple | βœ… | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/llm-decoding/beam-search/beam-search.ipynb) | [Q](v3/llm-decoding/beam-search/beam-search.ipynb) / [S](v3/llm-decoding/beam-search/beam-search_SOLN.ipynb) | diff --git a/README.md b/README.md index 98535f0..2616487 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ TorchLeet -**68 PyTorch problems from real ML/AI interviews at Google, Meta, Anthropic, and more.** +**74 PyTorch problems from real ML/AI interviews at Google, Meta, Anthropic, and more.** [![GitHub stars](https://img.shields.io/github/stars/Exorust/TorchLeet?style=social)](https://github.com/Exorust/TorchLeet) [![Website](https://img.shields.io/badge/website-torch--leet.vercel.app-8b5cf6)](https://torch-leet.vercel.app) @@ -20,7 +20,7 @@ I struggled to grind for ML/AI interviews so I went back to the basics and creat ## AI Tutor (NEW) -Turn any AI assistant into your PyTorch interview coach. The TorchLeet MCP server gives your AI access to all 68 problems, progressive hints, company prep plans, and learning paths, while enforcing a no-spoilers teaching style. +Turn any AI assistant into your PyTorch interview coach. The TorchLeet MCP server gives your AI access to all 74 problems, progressive hints, company prep plans, and learning paths, while enforcing a no-spoilers teaching style. ```bash # Clone the repo first @@ -66,7 +66,7 @@ Add this to your MCP config: --- -68 problems across three tracks: +74 problems across three tracks: | Track | Focus | Questions | |-------|-------|-----------| @@ -121,9 +121,9 @@ Build an LLM from scratch, one question at a time. Recommended order: ### 4. Alignment & Fine-Tuning | Problem | Companies | Links | |---------|-----------|-------| -| Implement KL Divergence Loss | | | +| Implement KL Divergence Loss | | [Q](v3/alignment-training/kl-divergence/kl-divergence.ipynb) / [S](v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb) | | Implement LoRA | Meta, Google, Anthropic, OpenAI | [Q](v3/alignment-training/lora/lora.ipynb) / [S](v3/alignment-training/lora/lora_SOLN.ipynb) | -| Apply SFT on SmolLM | | | +| Apply SFT on SmolLM | | [Q](v3/alignment-training/sft-smollm/sft-smollm.ipynb) / [S](v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb) | | Implement DPO Loss | Anthropic, OpenAI, DeepMind, Meta | [Q](v3/alignment-training/dpo/dpo.ipynb) / [S](v3/alignment-training/dpo/dpo_SOLN.ipynb) | | Implement PPO for RLHF | Anthropic, OpenAI, DeepMind, Meta | [Q](v3/alignment-training/ppo-rlhf/ppo-rlhf.ipynb) / [S](v3/alignment-training/ppo-rlhf/ppo-rlhf_SOLN.ipynb) | | Implement GRPO (DeepSeek-R1) | DeepMind, Anthropic, OpenAI | [Q](v3/alignment-training/grpo/grpo.ipynb) / [S](v3/alignment-training/grpo/grpo_SOLN.ipynb) | @@ -207,12 +207,13 @@ Company-tagged questions from real ML/AI interviews. Sorted by topic. ### Alignment & Training | Problem | Difficulty | Companies | Links | |---------|------------|-----------|-------| +| Implement KL Divergence Loss | Medium | | [Q](v3/alignment-training/kl-divergence/kl-divergence.ipynb) / [S](v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb) | | Implement LoRA | Medium | Meta, Google, Anthropic, OpenAI, Databricks | [Q](v3/alignment-training/lora/lora.ipynb) / [S](v3/alignment-training/lora/lora_SOLN.ipynb) | | Implement DPO Loss | Hard | Anthropic, OpenAI, DeepMind, Meta | [Q](v3/alignment-training/dpo/dpo.ipynb) / [S](v3/alignment-training/dpo/dpo_SOLN.ipynb) | | Implement PPO for RLHF | Hard | Anthropic, OpenAI, DeepMind, Meta | [Q](v3/alignment-training/ppo-rlhf/ppo-rlhf.ipynb) / [S](v3/alignment-training/ppo-rlhf/ppo-rlhf_SOLN.ipynb) | | Gradient Checkpointing | Hard | Meta, Google, NVIDIA, Tesla | [Q](v3/alignment-training/gradient-checkpointing/gradient-checkpointing.ipynb) / [S](v3/alignment-training/gradient-checkpointing/gradient-checkpointing_SOLN.ipynb) | | Implement GRPO (DeepSeek-R1) | Expert | DeepMind, Anthropic, OpenAI | [Q](v3/alignment-training/grpo/grpo.ipynb) / [S](v3/alignment-training/grpo/grpo_SOLN.ipynb) | -| Apply SFT on SmolLM | Hard | | | +| Apply SFT on SmolLM | Hard | | [Q](v3/alignment-training/sft-smollm/sft-smollm.ipynb) / [S](v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb) | ### LLM Inference & Systems | Problem | Difficulty | Companies | Links | diff --git a/SOURCES.md b/SOURCES.md index 0d5e06f..851a6ac 100644 --- a/SOURCES.md +++ b/SOURCES.md @@ -13,7 +13,7 @@ honestly than invent a precision we do not have. If you were asked one of these, [tell us](https://github.com/Exorust/TorchLeet/issues/new?template=interview-report.yml) and it becomes a **reported** tag with a real date. -**0 reported, 136 inferred across 27 companies.** +**0 reported, 140 inferred across 27 companies.** ### Adobe - Implement DDIM Sampling + Classifier-Free Guidance β€” _inferred_ @@ -40,6 +40,7 @@ If you were asked one of these, [tell us](https://github.com/Exorust/TorchLeet/i - Implement Temperature Sampling β€” _inferred_ - Implement Top-k Sampling β€” _inferred_ - Implement Top-p (Nucleus) Sampling β€” _inferred_ +- Simulate DistributedSampler for Data Parallel Training β€” _inferred_ ### Anyscale - Build a Complete LLM Inference Engine β€” _inferred_ @@ -102,6 +103,7 @@ If you were asked one of these, [tell us](https://github.com/Exorust/TorchLeet/i - Implement Speculative Decoding β€” _inferred_ - Implement Vision Transformer + MAE Pretraining β€” _inferred_ - Implement a Swin Transformer Block (Shifted Window Attention) β€” _inferred_ +- Simulate DistributedSampler for Data Parallel Training β€” _inferred_ - Write a Fused Softmax Kernel in Triton β€” _inferred_ ### LinkedIn @@ -128,6 +130,7 @@ If you were asked one of these, [tell us](https://github.com/Exorust/TorchLeet/i - Implement Softmax from Scratch β€” _inferred_ - Implement Vision Transformer + MAE Pretraining β€” _inferred_ - Implement a Swin Transformer Block (Shifted Window Attention) β€” _inferred_ +- Simulate DistributedSampler for Data Parallel Training β€” _inferred_ - Write a Fused Softmax Kernel in Triton β€” _inferred_ ### Microsoft @@ -147,6 +150,7 @@ If you were asked one of these, [tell us](https://github.com/Exorust/TorchLeet/i - Implement FSDP from Scratch β€” _inferred_ - Implement FlashAttention-2 in Triton β€” _inferred_ - Implement Gradient Checkpointing β€” _inferred_ +- Simulate DistributedSampler for Data Parallel Training β€” _inferred_ - Write a Fused Softmax Kernel in Triton β€” _inferred_ ### OpenAI diff --git a/STUDYPLAN.md b/STUDYPLAN.md index 9fcb8f4..7a2cff3 100644 --- a/STUDYPLAN.md +++ b/STUDYPLAN.md @@ -6,25 +6,25 @@ measurements. Treat them as relative effort, not a promise. ## Basics -**23 problems Β· ~11 hours Β· about 2 weeks at 5 h/week** +**27 problems Β· ~13 hours Β· about 3 weeks at 5 h/week** - 7 basic (~2h) -- 13 easy (~6h) -- 3 medium (~2h) +- 14 easy (~7h) +- 6 medium (~4h) ## LLM Learning Path -**21 problems Β· ~22 hours Β· about 4 weeks at 5 h/week** +**23 problems Β· ~24 hours Β· about 5 weeks at 5 h/week** - 3 easy (~1h) -- 10 medium (~7h) -- 6 hard (~9h) +- 11 medium (~8h) +- 7 hard (~10h) - 2 expert (~4h) ## Advanced -**42 problems Β· ~54 hours Β· about 11 weeks at 5 h/week** +**44 problems Β· ~56 hours Β· about 11 weeks at 5 h/week** -- 16 medium (~12h) -- 20 hard (~30h) +- 17 medium (~12h) +- 21 hard (~31h) - 6 expert (~12h) diff --git a/mcp-server/src/data/hints.generated.ts b/mcp-server/src/data/hints.generated.ts index 3ca5a19..1637bf7 100644 --- a/mcp-server/src/data/hints.generated.ts +++ b/mcp-server/src/data/hints.generated.ts @@ -233,6 +233,16 @@ export const authoredHints: Record = { "L is the logsumexp of the scaled scores: running_max + log(l) at the end.", "tl.advance the K and V block pointers each iteration; K's block pointer is already transposed, so it advances along its second axis." ], + "kl-divergence": [ + "Everything is sum(p * (log p - log q)) over the last dim. Use torch.xlogy when you have probabilities and F.log_softmax when you have logits, so 0 * log 0 and softmax underflow never turn into nan.", + "F.kl_div(input, target) is KL(target || exp(input)): input is the model's log-probs, the target comes first in the math, and 'batchmean' (sum / batch size) is the true KL while 'mean' divides by every element.", + "For samples x ~ p with log r = log q(x) - log p(x): k1 = -log r, k2 = (log r)^2 / 2, k3 = (r - 1) - log r. For kl_estimator_grad, weight every category by pi.detach() and let the gradient flow only through log pi." + ], + "v2-1": [ + "Everything is sum(p * (log p - log q)) over the last dim. Use torch.xlogy when you have probabilities and F.log_softmax when you have logits, so 0 * log 0 and softmax underflow never turn into nan.", + "F.kl_div(input, target) is KL(target || exp(input)): input is the model's log-probs, the target comes first in the math, and 'batchmean' (sum / batch size) is the true KL while 'mean' divides by every element.", + "For samples x ~ p with log r = log q(x) - log p(x): k1 = -log r, k2 = (log r)^2 / 2, k3 = (r - 1) - log r. For kl_estimator_grad, weight every category by pi.detach() and let the gradient flow only through log pi." + ], "byte-pair-encoder": [ "get_vocab splits each word into characters plus an end-of-word marker, and counts it.", "get_stats counts adjacent symbol pairs weighted by each word's frequency.", @@ -303,6 +313,16 @@ export const authoredHints: Record = { "repeat_kv expands each KV head n_rep times *in place* so head i of Q pairs with KV head i // n_rep β€” expand + reshape, not a plain repeat/tile.", "RopeAttention: project, split into num_heads (Q) and kv_heads (K/V), rotate Q and K, repeat_kv the K/V, then scaled dot-product attention plus the additive mask." ], + "sft-smollm": [ + "labels are a copy of input_ids (prompt + response + end-of-turn) with every prompt position set to -100; keep them aligned with input_ids and shift inside the loss.", + "Pad input_ids with pad_id but labels with -100, so a pad id that equals a real token id never erases a real label.", + "sft_loss: cross_entropy(logits[:, :-1], labels[:, 1:], ignore_index=-100, reduction='sum') divided by the trained-token count of the whole optimizer step, counted across all micro-batches before any backward." + ], + "v2-20": [ + "labels are a copy of input_ids (prompt + response + end-of-turn) with every prompt position set to -100; keep them aligned with input_ids and shift inside the loss.", + "Pad input_ids with pad_id but labels with -100, so a pad id that equals a real token id never erases a real label.", + "sft_loss: cross_entropy(logits[:, :-1], labels[:, 1:], ignore_index=-100, reduction='sum') divided by the trained-token count of the whole optimizer step, counted across all micro-batches before any backward." + ], "softmax": [ "exp() of a large number overflows. What can you subtract without changing the result?", "softmax(x) == softmax(x - c) for any constant c. Pick c = max(x) along `dim`.", diff --git a/mcp-server/src/data/questions.ts b/mcp-server/src/data/questions.ts index 66747d0..fa3ce18 100644 --- a/mcp-server/src/data/questions.ts +++ b/mcp-server/src/data/questions.ts @@ -447,11 +447,11 @@ const v2Questions: Omit = { "v1-38": "Implement the Muon optimizer from scratch on top of torch.optim.Optimizer: momentum orthogonalized by Newton–Schulz for matrix parameters, with an AdamW fallback for 1D parameters.", // V2 - "v2-1": "Compute KL divergence between two probability distributions from scratch, essential for VAEs and knowledge distillation.", + "v2-1": "Implement KL divergence every way it shows up in ML: from probabilities and logits, F.kl_div semantics, Jensen-Shannon, forward vs reverse, Gaussian closed form, masked token-level KL, and the k1/k2/k3 estimators used in PPO and GRPO.", "v2-2": "Build Root Mean Square Layer Normalization used in LLaMA and modern transformers β€” simpler and faster than LayerNorm.", "v2-3": "Build the BPE tokenizer algorithm that iteratively merges frequent character pairs to build a subword vocabulary.", "v2-4": "Build retrieval-augmented generation by encoding a corpus, performing similarity search, and conditioning generation on retrieved context.", @@ -1083,7 +1083,7 @@ const DESCRIPTIONS: Record = { "v2-17": "Scale logits by a temperature parameter before softmax to control the randomness of generated text.", "v2-18": "Add low-rank adapter matrices to a pretrained layer, enabling efficient fine-tuning with a fraction of the parameters.", "v2-19": "Build a gated MoE layer that routes tokens to specialized expert networks, scaling model capacity without proportional compute.", - "v2-20": "Fine-tune a language model on instruction-following data using supervised fine-tuning with cross-entropy loss.", + "v2-20": "Fine-tune the pretrained SmolLM-135M base model into a chat assistant: build prompt-masked examples, pad and collate, write the shifted token-level loss, and normalize gradient accumulation correctly.", "v2-21": "Align a language model with human preferences using a reward model and reinforcement learning with PPO.", "v2-22": "Train a language model directly on preference pairs without a separate reward model using Direct Preference Optimization.", "v2-23": "Implement dynamic batching that adds/removes sequences mid-batch to maximize GPU utilization during inference.", @@ -1148,19 +1148,21 @@ const LLM_PATH_CONFIG: Record = { // Full Model "v2-12": { order: 9, stage: "full-model" }, // SmolLM from Scratch // Alignment & PEFT - "v3-11": { order: 10, stage: "alignment" }, // LoRA - "v3-14": { order: 11, stage: "alignment" }, // DPO - "v3-15": { order: 12, stage: "alignment" }, // PPO - "v3-27": { order: 13, stage: "alignment" }, // GRPO + "v2-1": { order: 10, stage: "alignment" }, // KL Divergence + "v3-11": { order: 11, stage: "alignment" }, // LoRA + "v2-20": { order: 12, stage: "alignment" }, // SFT on SmolLM + "v3-14": { order: 13, stage: "alignment" }, // DPO + "v3-15": { order: 14, stage: "alignment" }, // PPO + "v3-27": { order: 15, stage: "alignment" }, // GRPO // Decoding & Inference - "v3-10": { order: 14, stage: "decoding-inference" }, // Temperature - "v3-8": { order: 15, stage: "decoding-inference" }, // Top-k - "v3-7": { order: 16, stage: "decoding-inference" }, // Top-p - "v3-18": { order: 17, stage: "decoding-inference" }, // Speculative Decoding - "v3-19": { order: 18, stage: "decoding-inference" }, // Continuous Batching - "v3-28": { order: 19, stage: "decoding-inference" }, // Full Inference Engine + "v3-10": { order: 16, stage: "decoding-inference" }, // Temperature + "v3-8": { order: 17, stage: "decoding-inference" }, // Top-k + "v3-7": { order: 18, stage: "decoding-inference" }, // Top-p + "v3-18": { order: 19, stage: "decoding-inference" }, // Speculative Decoding + "v3-19": { order: 20, stage: "decoding-inference" }, // Continuous Batching + "v3-28": { order: 21, stage: "decoding-inference" }, // Full Inference Engine // Systems (capstone for path) - "v3-17": { order: 20, stage: "systems" }, // MoE + "v3-17": { order: 22, stage: "systems" }, // MoE }; function assignTracksAndPath(q: any): Partial { diff --git a/mcp-server/src/tools/get-prerequisites.ts b/mcp-server/src/tools/get-prerequisites.ts index dc4ed08..362a0b9 100644 --- a/mcp-server/src/tools/get-prerequisites.ts +++ b/mcp-server/src/tools/get-prerequisites.ts @@ -21,9 +21,11 @@ const PREREQ_GRAPH: Record = { "v2-12": ["v2-3", "v2-11", "v2-8", "v3-12"], // Alignment chain + "v2-1": [], "v3-11": [], - "v3-14": ["v3-11"], - "v3-15": ["v3-14"], + "v2-20": ["v2-12"], + "v3-14": ["v3-11", "v2-20"], + "v3-15": ["v3-14", "v2-1"], "v3-27": ["v3-15"], // Decoding chain diff --git a/problems.json b/problems.json index a332029..43f5b7f 100644 --- a/problems.json +++ b/problems.json @@ -1,6 +1,100 @@ { - "count": 68, + "count": 74, "problems": [ + { + "id": "buffered-loading", + "legacy_ids": [], + "title": "Implement a Shuffle Buffer for Streaming Data", + "description": "Build a shuffle buffer that wraps an IterableDataset, providing approximate shuffling with O(buffer_size) memory.", + "difficulty": "medium", + "category": "data-loading", + "tracks": [ + "basics" + ], + "companies": [], + "question": "v3/data-loading/buffered-loading/buffered-loading.ipynb", + "solution": "v3/data-loading/buffered-loading/buffered-loading_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/buffered-loading/buffered-loading.ipynb", + "frequency": "emerging", + "company_confidence": {}, + "llm_path_order": null, + "graded": false, + "entries": [], + "hints": null + }, + { + "id": "data-parallel-loading", + "legacy_ids": [], + "title": "Simulate DistributedSampler for Data Parallel Training", + "description": "Implement a distributed sampler from scratch that shards a dataset across ranks with deterministic shuffling and padding.", + "difficulty": "medium", + "category": "data-loading", + "tracks": [ + "basics" + ], + "companies": [ + "Meta", + "Google", + "NVIDIA", + "Anthropic" + ], + "question": "v3/data-loading/data-parallel-loading/data-parallel-loading.ipynb", + "solution": "v3/data-loading/data-parallel-loading/data-parallel-loading_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/data-parallel-loading/data-parallel-loading.ipynb", + "frequency": "common", + "company_confidence": { + "Meta": "inferred", + "Google": "inferred", + "NVIDIA": "inferred", + "Anthropic": "inferred" + }, + "llm_path_order": null, + "graded": false, + "entries": [], + "hints": null + }, + { + "id": "preloaded-dataset", + "legacy_ids": [], + "title": "Implement a Preloaded Map-Style Dataset", + "description": "Build a map-style Dataset that loads all data into memory and integrates with DataLoader for batching and shuffling.", + "difficulty": "easy", + "category": "data-loading", + "tracks": [ + "basics" + ], + "companies": [], + "question": "v3/data-loading/preloaded-dataset/preloaded-dataset.ipynb", + "solution": "v3/data-loading/preloaded-dataset/preloaded-dataset_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/preloaded-dataset/preloaded-dataset.ipynb", + "frequency": "common", + "company_confidence": {}, + "llm_path_order": null, + "graded": false, + "entries": [], + "hints": null + }, + { + "id": "streaming-dataset", + "legacy_ids": [], + "title": "Implement a Streaming IterableDataset with Yield", + "description": "Build an IterableDataset that streams data line by line from a file using a generator, with multi-worker sharding support.", + "difficulty": "medium", + "category": "data-loading", + "tracks": [ + "basics" + ], + "companies": [], + "question": "v3/data-loading/streaming-dataset/streaming-dataset.ipynb", + "solution": "v3/data-loading/streaming-dataset/streaming-dataset_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/streaming-dataset/streaming-dataset.ipynb", + "frequency": "common", + "company_confidence": {}, + "llm_path_order": null, + "graded": false, + "entries": [], + "hints": null + }, { "id": "lin-regression", "legacy_ids": [ @@ -829,6 +923,44 @@ "L is the logsumexp of the scaled scores: running_max + log(l) at the end." ] }, + { + "id": "kl-divergence", + "legacy_ids": [ + "v2-1" + ], + "title": "Implement KL Divergence Loss", + "description": "Implement KL divergence every way it shows up in ML: from probabilities and logits, F.kl_div semantics, Jensen-Shannon, forward vs reverse, Gaussian closed form, masked token-level KL, and the k1/k2/k3 estimators used in PPO and GRPO.", + "difficulty": "medium", + "category": "alignment-training", + "tracks": [ + "advanced", + "llm-path" + ], + "companies": [], + "question": "v3/alignment-training/kl-divergence/kl-divergence.ipynb", + "solution": "v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/kl-divergence/kl-divergence.ipynb", + "frequency": "", + "company_confidence": {}, + "llm_path_order": 11, + "graded": true, + "entries": [ + "kl_categorical", + "kl_from_logits", + "my_kl_div", + "js_divergence", + "fit_gaussian", + "kl_gaussian", + "masked_token_kl", + "kl_estimators", + "kl_estimator_grad" + ], + "hints": [ + "Everything is sum(p * (log p - log q)) over the last dim. Use torch.xlogy when you have probabilities and F.log_softmax when you have logits, so 0 * log 0 and softmax underflow never turn into nan.", + "F.kl_div(input, target) is KL(target || exp(input)): input is the model's log-probs, the target comes first in the math, and 'batchmean' (sum / batch size) is the true KL while 'mean' divides by every element.", + "For samples x ~ p with log r = log q(x) - log p(x): k1 = -log r, k2 = (log r)^2 / 2, k3 = (r - 1) - log r. For kl_estimator_grad, weight every category by pi.detach() and let the gradient flow only through log pi." + ] + }, { "id": "rms-norm", "legacy_ids": [ @@ -844,14 +976,20 @@ ], "companies": [], "question": "llm/RMS-Norm/rms-norm.ipynb", - "solution": null, + "solution": "llm/RMS-Norm/rms-norm_SOLN.ipynb", "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/RMS-Norm/rms-norm.ipynb", "frequency": "emerging", "company_confidence": {}, "llm_path_order": null, - "graded": false, - "entries": [], - "hints": null + "graded": true, + "entries": [ + "RMSNorm" + ], + "hints": [ + "Square the input, take the mean along the last dimension, add epsilon, then take the square root.", + "Use keepdim=True in the mean so broadcasting works when you divide.", + "The scale parameter gamma should have shape (dim,) and be created with nn.Parameter(torch.ones(dim))." + ] }, { "id": "byte-pair-encoder", @@ -1071,6 +1209,39 @@ "RopeAttention: project, split into num_heads (Q) and kv_heads (K/V), rotate Q and K, repeat_kv the K/V, then scaled dot-product attention plus the additive mask." ] }, + { + "id": "sft-smollm", + "legacy_ids": [ + "v2-20" + ], + "title": "Apply SFT on SmolLM", + "description": "Fine-tune the pretrained SmolLM-135M base model into a chat assistant: build prompt-masked examples, pad and collate, write the shifted token-level loss, and normalize gradient accumulation correctly.", + "difficulty": "hard", + "category": "alignment-training", + "tracks": [ + "advanced", + "llm-path" + ], + "companies": [], + "question": "v3/alignment-training/sft-smollm/sft-smollm.ipynb", + "solution": "v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/sft-smollm/sft-smollm.ipynb", + "frequency": "", + "company_confidence": {}, + "llm_path_order": 13, + "graded": true, + "entries": [ + "build_sft_example", + "sft_collate", + "sft_loss", + "sft_train_step" + ], + "hints": [ + "labels are a copy of input_ids (prompt + response + end-of-turn) with every prompt position set to -100; keep them aligned with input_ids and shift inside the loss.", + "Pad input_ids with pad_id but labels with -100, so a pad id that equals a real token id never erases a real label.", + "sft_loss: cross_entropy(logits[:, :-1], labels[:, 1:], ignore_index=-100, reduction='sum') divided by the trained-token count of the whole optimizer step, counted across all micro-batches before any backward." + ] + }, { "id": "softmax", "legacy_ids": [ diff --git a/python/test_grader.py b/python/test_grader.py index 04f1e5c..16fd2c7 100644 --- a/python/test_grader.py +++ b/python/test_grader.py @@ -226,6 +226,163 @@ def ckpt_inputs_only(fn, *args): return CkptFnInputsOnly.apply(fn, *args) +# ------------------------------------------------------------- kl-divergence +def kl_categorical(p, q): + return (torch.xlogy(p, p) - torch.xlogy(p, q)).sum(-1) + + +def kl_from_logits(pl, ql): + lp, lq = F.log_softmax(pl, -1), F.log_softmax(ql, -1) + return (lp.exp() * (lp - lq)).sum(-1) + + +def kl_from_logits_lse(pl, ql): + """Correct, but normalizes by hand with logsumexp.""" + lp = pl - torch.logsumexp(pl, -1, keepdim=True) + lq = ql - torch.logsumexp(ql, -1, keepdim=True) + return torch.sum(torch.softmax(pl, -1) * (lp - lq), dim=-1) + + +def kl_from_logits_naive(pl, ql): + """Broken: goes through probabilities, so confident logits underflow to log(0).""" + p, q = pl.softmax(-1), ql.softmax(-1) + return (p * torch.log(p / q)).sum(-1) + + +def my_kl_div(inp, tgt, reduction="mean", log_target=False): + pw = tgt.exp() * (tgt - inp) if log_target else torch.xlogy(tgt, tgt) - tgt * inp + return {"none": lambda: pw, "sum": pw.sum, "mean": pw.mean, + "batchmean": lambda: pw.sum() / inp.size(0)}[reduction]() + + +def my_kl_div_mean_is_batchmean(inp, tgt, reduction="mean", log_target=False): + """Broken: 'fixes' mean to be the true KL, which is not what F.kl_div does.""" + return my_kl_div(inp, tgt, "batchmean" if reduction == "mean" else reduction, log_target) + + +def js_divergence(pl, ql): + lp, lq = F.log_softmax(pl, -1), F.log_softmax(ql, -1) + lm = torch.logsumexp(torch.stack([lp, lq]), 0) - math.log(2) + return 0.5 * (lp.exp() * (lp - lm)).sum(-1) + 0.5 * (lq.exp() * (lq - lm)).sum(-1) + + +def fit_gaussian(target_logp, grid, direction, steps=500, lr=0.05, mu_init=1.0, sigma_init=0.5): + mu = torch.tensor(float(mu_init), requires_grad=True) + ls = torch.tensor(math.log(sigma_init), requires_grad=True) + opt = torch.optim.Adam([mu, ls], lr=lr) + for _ in range(steps): + ql = -0.5 * ((grid - mu) / ls.exp()) ** 2 + loss = kl_from_logits(target_logp, ql) if direction == "forward" else kl_from_logits(ql, target_logp) + opt.zero_grad() + loss.backward() + opt.step() + return mu.item(), ls.exp().item() + + +def kl_gaussian(mp, lp, mq, lq): + return 0.5 * (lq - lp + (lp.exp() + (mp - mq) ** 2) / lq.exp() - 1).sum(-1) + + +def masked_token_kl(pl, ql, mask): + m = mask.float() + return (kl_from_logits(pl, ql) * m).sum() / m.sum().clamp(min=1) + + +def kl_estimators(logp, logq): + lr = logq - logp + return -lr, 0.5 * lr ** 2, lr.exp() - 1 - lr + + +def kl_estimator_grad(logits, ref, est): + th = logits.detach().requires_grad_(True) + lp, lq = F.log_softmax(th, -1), F.log_softmax(ref, -1).detach() + k = kl_estimators(lp, lq)[int(est[1]) - 1] + return torch.autograd.grad((lp.exp().detach() * k).sum(), th)[0] + + +def kl_estimator_grad_no_detach(logits, ref, est): + """Broken: differentiates through pi(x) as well, i.e. the exact KL, not the estimator-as-loss.""" + th = logits.detach().requires_grad_(True) + lp, lq = F.log_softmax(th, -1), F.log_softmax(ref, -1) + k = kl_estimators(lp, lq)[int(est[1]) - 1] + return torch.autograd.grad((lp.exp() * k).sum(), th)[0] + + +KL = dict(kl_categorical=kl_categorical, kl_from_logits=kl_from_logits, my_kl_div=my_kl_div, + js_divergence=js_divergence, fit_gaussian=fit_gaussian, kl_gaussian=kl_gaussian, + masked_token_kl=masked_token_kl, kl_estimators=kl_estimators, + kl_estimator_grad=kl_estimator_grad) + + +# ---------------------------------------------------------------- sft-smollm +def build_sft_example(prompt, response, eos, max_len=None): + prompt, response = list(map(int, prompt)), list(map(int, response)) + ids, lab = prompt + response + [eos], [-100] * len(prompt) + response + [eos] + return {"input_ids": torch.tensor(ids[:max_len]), "labels": torch.tensor(lab[:max_len])} + + +def sft_collate(exs, pad_id): + """Correct, but built with pad_sequence instead of a preallocated tensor.""" + ids = nn.utils.rnn.pad_sequence([e["input_ids"] for e in exs], batch_first=True, padding_value=pad_id) + lab = nn.utils.rnn.pad_sequence([e["labels"] for e in exs], batch_first=True, padding_value=-100) + am = nn.utils.rnn.pad_sequence([torch.ones_like(e["input_ids"]) for e in exs], batch_first=True) + return {"input_ids": ids, "attention_mask": am, "labels": lab} + + +def sft_collate_pads_labels_with_pad(exs, pad_id): + """Broken: labels = input_ids with pad positions masked, so a real eos == pad is lost.""" + b = sft_collate(exs, pad_id) + b["labels"] = b["labels"].masked_fill(b["input_ids"] == pad_id, -100) + return b + + +def sft_loss(logits, labels, num_items_in_batch=None): + lg, tg = logits[:, :-1], labels[:, 1:] + s = F.cross_entropy(lg.reshape(-1, lg.size(-1)).float(), tg.reshape(-1), ignore_index=-100, reduction="sum") + n = num_items_in_batch if num_items_in_batch is not None else int((tg != -100).sum()) + return s / max(int(n), 1) + + +def sft_loss_unshifted(logits, labels, num_items_in_batch=None): + """Broken: scores position t against its own token instead of the next one.""" + s = F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1), ignore_index=-100, reduction="sum") + n = num_items_in_batch if num_items_in_batch is not None else int((labels != -100).sum()) + return s / max(int(n), 1) + + +def sft_train_step(model, mbs, opt, max_grad_norm=None): + model.train() + n = sum(int((mb["labels"][:, 1:] != -100).sum()) for mb in mbs) + opt.zero_grad() + total = 0.0 + for mb in mbs: + loss = sft_loss(model(input_ids=mb["input_ids"], attention_mask=mb["attention_mask"]).logits, + mb["labels"], n) + loss.backward() + total += loss.item() + if max_grad_norm: + nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) + opt.step() + return total + + +def sft_train_step_mean_of_means(model, mbs, opt, max_grad_norm=None): + """Broken: the classic gradient-accumulation bug β€” each micro-batch is its own mean.""" + model.train() + opt.zero_grad() + total = 0.0 + for mb in mbs: + loss = sft_loss(model(input_ids=mb["input_ids"]).logits, mb["labels"]) / len(mbs) + loss.backward() + total += loss.item() + opt.step() + return total + + +SFT = dict(build_sft_example=build_sft_example, sft_collate=sft_collate, sft_loss=sft_loss, + sft_train_step=sft_train_step) + + CASES = [ # (label, expected, problem_id, args, kwargs) ("kv-cache canonical", True, "kv-cache", (), dict(KVCache=KVCache, CachedAttention=Attn)), @@ -242,6 +399,15 @@ def ckpt_inputs_only(fn, *args): ("ckpt correct", True, "gradient-checkpointing", (CkptFn, ckpt), {}), ("ckpt no dummy input", False, "gradient-checkpointing", (CkptFn, ckpt_no_dummy), {}), ("ckpt inputs-only grads", False, "gradient-checkpointing", (CkptFnInputsOnly, ckpt_inputs_only), {}), + ("kl canonical", True, "kl-divergence", (), KL), + ("kl divergent logsumexp", True, "kl-divergence", (), {**KL, "kl_from_logits": kl_from_logits_lse}), + ("kl naive softmax", False, "kl-divergence", (), {**KL, "kl_from_logits": kl_from_logits_naive}), + ("kl 'mean' as batchmean", False, "kl-divergence", (), {**KL, "my_kl_div": my_kl_div_mean_is_batchmean}), + ("kl grad without detach", False, "kl-divergence", (), {**KL, "kl_estimator_grad": kl_estimator_grad_no_detach}), + ("sft divergent collate", True, "sft-smollm", (), SFT), + ("sft labels padded via pad id", False, "sft-smollm", (), {**SFT, "sft_collate": sft_collate_pads_labels_with_pad}), + ("sft unshifted loss", False, "sft-smollm", (), {**SFT, "sft_loss": sft_loss_unshifted}), + ("sft mean of micro-batch means", False, "sft-smollm", (), {**SFT, "sft_train_step": sft_train_step_mean_of_means}), ] diff --git a/python/torchleet/problems/kl_divergence.py b/python/torchleet/problems/kl_divergence.py new file mode 100644 index 0000000..6518882 --- /dev/null +++ b/python/torchleet/problems/kl_divergence.py @@ -0,0 +1,242 @@ +"""kl-divergence β€” every way KL shows up: probabilities, logits, F.kl_div semantics, +JS, forward vs reverse, Gaussians, masked token-level KL, and k1/k2/k3 estimators.""" +import math +import warnings + +import torch +import torch.nn.functional as F +from torch.distributions import Categorical, Normal, kl_divergence + +ENTRIES = ["kl_categorical", "kl_from_logits", "my_kl_div", "js_divergence", "fit_gaussian", + "kl_gaussian", "masked_token_kl", "kl_estimators", "kl_estimator_grad"] +DEVICE = "cpu" +EXTRAS = [] + +HINTS = [ + "Everything is sum(p * (log p - log q)) over the last dim. Use torch.xlogy when you have " + "probabilities and F.log_softmax when you have logits, so 0 * log 0 and softmax underflow " + "never turn into nan.", + "F.kl_div(input, target) is KL(target || exp(input)): input is the model's log-probs, the " + "target comes first in the math, and 'batchmean' (sum / batch size) is the true KL while " + "'mean' divides by every element.", + "For samples x ~ p with log r = log q(x) - log p(x): k1 = -log r, k2 = (log r)^2 / 2, " + "k3 = (r - 1) - log r. For kl_estimator_grad, weight every category by pi.detach() and let " + "the gradient flow only through log pi.", +] + + +def _oracle_kl(p_logits, q_logits): + return kl_divergence(Categorical(logits=p_logits), Categorical(logits=q_logits)) + + +def _t(x): + return torch.as_tensor(x) + + +# ------------------------------------------------------------- probabilities +def check_categorical_matches_torch(ns): + torch.manual_seed(0) + p = torch.softmax(torch.randn(4, 3, 6), -1) + q = torch.softmax(torch.randn(4, 3, 6), -1) + got = _t(ns.kl_categorical(p, q)) + assert tuple(got.shape) == (4, 3), \ + f"kl_categorical should reduce only the last dim: expected (4, 3), got {tuple(got.shape)}" + ref = kl_divergence(Categorical(probs=p), Categorical(probs=q)) + assert torch.allclose(got, ref, atol=1e-5), \ + "kl_categorical(p, q) disagrees with KL(P || Q); check the direction and the formula" + + +def check_categorical_handles_zeros(ns): + p = torch.tensor([0.5, 0.5, 0.0]) + q = torch.tensor([0.25, 0.25, 0.5]) + got = _t(ns.kl_categorical(p, q)) + assert not torch.isnan(got), "p = 0 entries produced nan; 0 * log 0 must count as 0 (try torch.xlogy)" + assert abs(float(got) - math.log(2)) < 1e-5, f"expected log 2 = {math.log(2):.4f}, got {float(got):.4f}" + back = _t(ns.kl_categorical(q, p)) + assert torch.isinf(back) and back > 0, \ + f"q = 0 where p > 0 must give +inf, got {float(back)}" + + +# -------------------------------------------------------------------- logits +def check_logits_match_torch(ns): + torch.manual_seed(1) + pl, ql = torch.randn(3, 5, 10), torch.randn(3, 5, 10) + got = _t(ns.kl_from_logits(pl, ql)) + assert tuple(got.shape) == (3, 5), f"expected shape (3, 5), got {tuple(got.shape)}" + assert torch.allclose(got, _oracle_kl(pl, ql), atol=1e-5), \ + "kl_from_logits(p_logits, q_logits) disagrees with KL(softmax(p) || softmax(q))" + + +def check_logits_stable_for_large_logits(ns): + torch.manual_seed(2) + pl, ql = 100 * torch.randn(4, 12), 100 * torch.randn(4, 12) + got = _t(ns.kl_from_logits(pl, ql)) + assert torch.isfinite(got).all(), \ + "kl_from_logits returned inf/nan for large logits; compute log-probs with F.log_softmax" + lp, lq = F.log_softmax(pl.double(), -1), F.log_softmax(ql.double(), -1) + ref = (lp.exp() * (lp - lq)).sum(-1) + assert torch.allclose(got.double(), ref, rtol=1e-3, atol=1e-3), \ + "large-logit KL is finite but wrong; is a probability underflowing somewhere?" + + +# ------------------------------------------------------------------- F.kl_div +def check_kl_div_matches_torch_semantics(ns): + torch.manual_seed(3) + for shape in ((8, 5), (2, 3, 7)): + log_q = F.log_softmax(torch.randn(*shape), -1) + p = torch.softmax(torch.randn(*shape), -1) + p_zero = p.clone() + p_zero[..., 0] = 0.0 + p_zero = p_zero / p_zero.sum(-1, keepdim=True) + for log_target, tgt in ((False, p_zero), (True, p.log())): + for red in ("none", "sum", "mean", "batchmean"): + got = _t(ns.my_kl_div(log_q, tgt, reduction=red, log_target=log_target)) + with warnings.catch_warnings(): # torch warns on 'mean'; that is the point here + warnings.simplefilter("ignore") + exp = F.kl_div(log_q, tgt, reduction=red, log_target=log_target) + assert tuple(got.shape) == tuple(exp.shape) and torch.allclose(got, exp, atol=1e-6), ( + f"my_kl_div(reduction='{red}', log_target={log_target}) on shape {shape} " + "does not match F.kl_div; remember 'mean' averages over every element and " + "'batchmean' divides by input.size(0)") + + +# ------------------------------------------------------------------------ JS +def check_js_divergence(ns): + torch.manual_seed(4) + pl, ql = torch.randn(6, 7), torch.randn(6, 7) + got = _t(ns.js_divergence(pl, ql)) + P, Q = pl.softmax(-1), ql.softmax(-1) + M = Categorical(probs=0.5 * (P + Q)) + ref = 0.5 * kl_divergence(Categorical(probs=P), M) + 0.5 * kl_divergence(Categorical(probs=Q), M) + assert tuple(got.shape) == (6,), f"expected shape (6,), got {tuple(got.shape)}" + assert torch.allclose(got, ref, atol=1e-5), "js_divergence disagrees with 0.5*KL(P||M) + 0.5*KL(Q||M)" + assert torch.allclose(got, _t(ns.js_divergence(ql, pl)), atol=1e-6), "JS must be symmetric" + far_p = torch.tensor([50.0, -50.0, -50.0, -50.0]) + far_q = torch.tensor([-50.0, -50.0, -50.0, 50.0]) + far = float(_t(ns.js_divergence(far_p, far_q))) + assert abs(far - math.log(2)) < 1e-4, f"disjoint supports must give log 2, got {far:.4f}" + + +# ---------------------------------------------------------- forward vs reverse +def check_forward_covers_reverse_seeks(ns): + grid = torch.linspace(-6, 6, 241) + mix = torch.logsumexp(torch.stack([Normal(-2.5, 0.6).log_prob(grid), + Normal(2.5, 0.6).log_prob(grid)]), 0) + target = F.log_softmax(mix, -1) + kw = dict(steps=500, lr=0.05, mu_init=1.0, sigma_init=0.5) + mu_f, s_f = (float(v) for v in ns.fit_gaussian(target, grid, "forward", **kw)) + mu_r, s_r = (float(v) for v in ns.fit_gaussian(target, grid, "reverse", **kw)) + assert abs(mu_f) < 0.5 and s_f > 2.0, ( + f"forward KL(target || q) should cover both modes (mu ~ 0, wide sigma); got mu={mu_f:.2f}, " + f"sigma={s_f:.2f}") + assert abs(mu_r - 2.5) < 0.3 and s_r < 1.0, ( + f"reverse KL(q || target) from mu_init=1.0 should lock onto the +2.5 mode with a narrow sigma; " + f"got mu={mu_r:.2f}, sigma={s_r:.2f}. Is the direction swapped?") + + +# ------------------------------------------------------------------ Gaussians +def check_gaussian_closed_form(ns): + torch.manual_seed(5) + mp, lp, mq, lq = (torch.randn(4, 3) for _ in range(4)) + got = _t(ns.kl_gaussian(mp, lp, mq, lq)) + ref = kl_divergence(Normal(mp, (0.5 * lp).exp()), Normal(mq, (0.5 * lq).exp())).sum(-1) + assert tuple(got.shape) == (4,), f"sum over the last dim: expected (4,), got {tuple(got.shape)}" + assert torch.allclose(got, ref, atol=1e-5), \ + "kl_gaussian disagrees with the Normal closed form; arguments are log-variances" + z = torch.zeros_like(mp) + vae = -0.5 * (1 + lp - mp ** 2 - lp.exp()).sum(-1) + assert torch.allclose(_t(ns.kl_gaussian(mp, lp, z, z)), vae, atol=1e-5), \ + "against N(0, I) it should reduce to the VAE term -0.5 * sum(1 + logvar - mu^2 - exp(logvar))" + + +# ------------------------------------------------------------- token-level KL +def _token_setup(): + torch.manual_seed(6) + B, T, V = 3, 6, 11 + pl, ql = torch.randn(B, T, V), torch.randn(B, T, V) + pl[2] *= 4 + mask = torch.zeros(B, T) + for i, n in enumerate((6, 4, 2)): + mask[i, :n] = 1 + return pl, ql, mask + + +def check_masked_token_kl_is_token_mean(ns): + pl, ql, mask = _token_setup() + got = _t(ns.masked_token_kl(pl, ql, mask)) + assert got.ndim == 0, f"expected a scalar, got shape {tuple(got.shape)}" + ref = _oracle_kl(pl, ql)[mask.bool()].mean() + assert torch.allclose(got, ref, atol=1e-5), \ + f"expected the mean over the {int(mask.sum())} unmasked tokens ({float(ref):.4f}), got {float(got):.4f}" + as_bool = _t(ns.masked_token_kl(pl, ql, mask.bool())) + assert torch.allclose(as_bool, ref, atol=1e-5), "a boolean mask must work too" + + +def check_masked_positions_are_ignored(ns): + pl, ql, mask = _token_setup() + base = _t(ns.masked_token_kl(pl, ql, mask)) + pad = mask == 0 + pl2, ql2 = pl.clone(), ql.clone() + pl2[pad] = 1e3 * torch.randn(int(pad.sum()), pl.shape[-1]) + ql2[pad] = -1e3 * torch.randn(int(pad.sum()), pl.shape[-1]) + assert torch.allclose(_t(ns.masked_token_kl(pl2, ql2, mask)), base, atol=1e-5), \ + "changing the logits at masked positions changed the result" + qg = ql.clone().requires_grad_(True) + _t(ns.masked_token_kl(pl, qg, mask)).backward() + assert qg.grad is not None and qg.grad[pad].abs().max() == 0, "masked positions must get zero gradient" + assert qg.grad[~pad].abs().max() > 0, "no gradient reached the unmasked positions" + + +# ------------------------------------------------------------------ estimators +def check_estimator_formulas(ns): + torch.manual_seed(7) + logp, logq = -torch.rand(100) * 3, -torch.rand(100) * 3 + k1, k2, k3 = (_t(k) for k in ns.kl_estimators(logp, logq)) + lr = logq - logp + for name, got, exp in (("k1", k1, -lr), ("k2", k2, 0.5 * lr ** 2), ("k3", k3, lr.exp() - 1 - lr)): + assert tuple(got.shape) == (100,), f"{name}: expected one value per sample, got {tuple(got.shape)}" + assert torch.allclose(got, exp, atol=1e-5), \ + f"{name} is wrong; with log r = logq - logp, check the sign and the definition" + + +def check_estimators_unbiased_on_samples(ns): + torch.manual_seed(8) + pl = torch.randn(50) + ql = pl + torch.randn(50) + true = float(_oracle_kl(pl, ql)) + x = Categorical(logits=pl).sample((200_000,)) + k1, _, k3 = (_t(k) for k in ns.kl_estimators(F.log_softmax(pl, -1)[x], F.log_softmax(ql, -1)[x])) + for name, k in (("k1", k1), ("k3", k3)): + se = float(k.std()) / math.sqrt(k.numel()) + assert abs(float(k.mean()) - true) < 5 * se, \ + f"{name} should be unbiased for KL(p || q) with x ~ p: mean {float(k.mean()):.4f} vs true {true:.4f}" + assert (k3 >= -1e-6).all(), "k3 = (r - 1) - log r is non-negative for every sample" + + +def check_estimator_gradients(ns): + torch.manual_seed(9) + logits, ref = torch.randn(4, 10), torch.randn(4, 10) + before = logits.clone() + th = logits.clone().requires_grad_(True) + g_rev = torch.autograd.grad(_oracle_kl(th, ref).sum(), th)[0] + th = logits.clone().requires_grad_(True) + g_fwd = torch.autograd.grad(_oracle_kl(ref, th).sum(), th)[0] + got = {e: _t(ns.kl_estimator_grad(logits, ref, e)) for e in ("k1", "k2", "k3")} + assert torch.equal(logits, before), "kl_estimator_grad must not modify its input logits" + for e, g in got.items(): + assert tuple(g.shape) == (4, 10), f"{e}: gradient should have the logits' shape, got {tuple(g.shape)}" + assert torch.allclose(got["k1"], torch.zeros_like(logits), atol=1e-6), \ + "k1 as a loss has zero expected gradient; are you weighting by pi.detach() and differentiating only k?" + assert torch.allclose(got["k2"], g_rev, atol=1e-5), \ + "k2 as a loss should give the gradient of the reverse KL(pi || ref)" + assert torch.allclose(got["k3"], g_fwd, atol=1e-5), \ + "k3 as a loss should give the gradient of the forward KL(ref || pi)" + + +CHECKS = [check_categorical_matches_torch, check_categorical_handles_zeros, + check_logits_match_torch, check_logits_stable_for_large_logits, + check_kl_div_matches_torch_semantics, check_js_divergence, + check_forward_covers_reverse_seeks, check_gaussian_closed_form, + check_masked_token_kl_is_token_mean, check_masked_positions_are_ignored, + check_estimator_formulas, check_estimators_unbiased_on_samples, + check_estimator_gradients] diff --git a/python/torchleet/problems/sft_smollm.py b/python/torchleet/problems/sft_smollm.py new file mode 100644 index 0000000..8fc3d4e --- /dev/null +++ b/python/torchleet/problems/sft_smollm.py @@ -0,0 +1,181 @@ +"""sft-smollm β€” supervised fine-tuning data path: prompt-masked examples, padding, +the shifted token-level loss, and gradient-accumulation normalization. + +Checks use a tiny position-wise LM, so nothing is downloaded and no `transformers` +is needed; the SmolLM fine-tune itself is exercised by the notebook. +""" +import copy +from types import SimpleNamespace + +import torch +import torch.nn as nn +import torch.nn.functional as F + +ENTRIES = ["build_sft_example", "sft_collate", "sft_loss", "sft_train_step"] +DEVICE = "cpu" +EXTRAS = [] + +HINTS = [ + "labels are a copy of input_ids (prompt + response + end-of-turn) with every prompt position " + "set to -100; keep them aligned with input_ids and shift inside the loss.", + "Pad input_ids with pad_id but labels with -100, so a pad id that equals a real token id " + "never erases a real label.", + "sft_loss: cross_entropy(logits[:, :-1], labels[:, 1:], ignore_index=-100, reduction='sum') " + "divided by the trained-token count of the whole optimizer step, counted across all " + "micro-batches before any backward.", +] + +IGNORE = -100 + + +def _lists(ex): + return [int(v) for v in ex["input_ids"]], [int(v) for v in ex["labels"]] + + +# ------------------------------------------------------------------- examples +def check_example_masks_prompt_and_keeps_eos(ns): + ids, labels = _lists(ns.build_sft_example([5, 6, 7], [8, 9], 2)) + assert ids == [5, 6, 7, 8, 9, 2], f"input_ids should be prompt + response + eos, got {ids}" + assert labels == [IGNORE, IGNORE, IGNORE, 8, 9, 2], ( + f"labels should be -100 on the prompt and the response + eos elsewhere, got {labels}") + + +def check_example_truncates_and_accepts_tensors(ns): + ids, labels = _lists(ns.build_sft_example(torch.tensor([5, 6, 7]), torch.tensor([8, 9]), 2, max_len=4)) + assert ids == [5, 6, 7, 8] and labels == [IGNORE, IGNORE, IGNORE, 8], ( + f"max_len=4 should keep the first 4 tokens of both, got input_ids={ids}, labels={labels}") + ex = ns.build_sft_example([1], [3], 2) + assert torch.as_tensor(ex["input_ids"]).dtype == torch.long, "return torch.long tensors" + + +# ---------------------------------------------------------------------- collate +def check_collate_pads_right_with_ignored_labels(ns): + eos = pad = 2 + exs = [ns.build_sft_example([5, 6], [7], eos), ns.build_sft_example([5, 6, 7], [8, 9, 10], eos), + ns.build_sft_example([5], [], eos)] + b = ns.sft_collate(exs, pad_id=pad) + ids, am, lab = (torch.as_tensor(b[k]) for k in ("input_ids", "attention_mask", "labels")) + assert tuple(ids.shape) == (3, 7) and tuple(lab.shape) == (3, 7) and tuple(am.shape) == (3, 7), \ + f"expected (3, 7) tensors, got input_ids {tuple(ids.shape)}" + assert am.sum(1).tolist() == [4, 7, 2], f"attention_mask should mark the real tokens, got {am.tolist()}" + assert ids[0].tolist() == [5, 6, 7, 2, 2, 2, 2], f"right-pad with pad_id, got {ids[0].tolist()}" + assert lab[0].tolist() == [IGNORE, IGNORE, 7, 2, IGNORE, IGNORE, IGNORE], ( + "pad the labels with -100 and keep the real end-of-turn label even when it equals pad_id; " + f"got {lab[0].tolist()}") + + +# ------------------------------------------------------------------------ loss +def _loss_batch(): + torch.manual_seed(0) + B, T, V = 3, 6, 13 + logits = torch.randn(B, T, V) + labels = torch.randint(0, V, (B, T)) + labels[0, :2] = IGNORE + labels[1, :4] = IGNORE + labels[2, 3:] = IGNORE + return logits, labels + + +def check_loss_is_shifted_token_mean(ns): + logits, labels = _loss_batch() + got = torch.as_tensor(ns.sft_loss(logits, labels)) + assert got.ndim == 0, f"the loss must be a scalar, got shape {tuple(got.shape)}" + exp = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)), labels[:, 1:].reshape(-1), + ignore_index=IGNORE) + unshifted = F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1), ignore_index=IGNORE) + assert torch.allclose(got.float(), exp, atol=1e-5), ( + f"expected the token mean of shifted cross-entropy ({float(exp):.4f}), got {float(got):.4f}" + + (" β€” that is the UNshifted loss: logits[:, :-1] predict labels[:, 1:]" + if torch.allclose(got.float(), unshifted, atol=1e-5) else "")) + + +def check_loss_ignores_unused_logits(ns): + logits, labels = _loss_batch() + base = torch.as_tensor(ns.sft_loss(logits, labels)) + unused = torch.zeros(labels.shape, dtype=torch.bool) + unused[:, :-1] = labels[:, 1:] == IGNORE + unused[:, -1] = True + noisy = logits.clone() + noisy[unused] = 100 * torch.randn(int(unused.sum()), logits.size(-1)) + assert torch.allclose(torch.as_tensor(ns.sft_loss(noisy, labels)), base, atol=1e-5), \ + "logits whose next label is -100 (and the last logit) must not affect the loss" + all_ignored = torch.full(labels.shape, IGNORE) + assert torch.isfinite(torch.as_tensor(ns.sft_loss(logits, all_ignored))), \ + "a batch with nothing to train must not give nan" + + +def check_loss_num_items_makes_microbatches_add_up(ns): + logits, labels = _loss_batch() + full = torch.as_tensor(ns.sft_loss(logits, labels)) + n = int((labels[:, 1:] != IGNORE).sum()) + parts = (torch.as_tensor(ns.sft_loss(logits[:2], labels[:2], n)) + + torch.as_tensor(ns.sft_loss(logits[2:], labels[2:], n))) + assert torch.allclose(parts, full, atol=1e-5), ( + "with num_items_in_batch set to the step's token count, micro-batch losses must sum to the " + f"full-batch loss ({float(parts):.4f} vs {float(full):.4f})") + + +# ------------------------------------------------------------------ train step +class _TinyLM(nn.Module): + """Position-wise LM with the Hugging Face output contract.""" + + def __init__(self, vocab=20, dim=16): + super().__init__() + self.emb = nn.Embedding(vocab, dim) + self.head = nn.Linear(dim, vocab) + + def forward(self, input_ids, attention_mask=None): + return SimpleNamespace(logits=self.head(torch.tanh(self.emb(input_ids)))) + + +def _uneven_examples(ns): + eos = 19 + return [ns.build_sft_example([1, 2, 3], [4, 5, 6, 7, 8, 9, 10, 11], eos), + ns.build_sft_example([1, 2], [12], eos), + ns.build_sft_example([3, 4, 5, 6], [7, 8], eos), + ns.build_sft_example([2], [13, 14, 15], eos)] + + +def check_microbatching_matches_one_big_batch(ns): + torch.manual_seed(0) + exs = _uneven_examples(ns) + full = ns.sft_collate(exs, 0) + split = [ns.sft_collate(exs[:2], 0), ns.sft_collate(exs[2:], 0)] + + a = _TinyLM() + b, ref = copy.deepcopy(a), copy.deepcopy(a) + labels = torch.as_tensor(full["labels"]) + with torch.no_grad(): + lg = ref(input_ids=torch.as_tensor(full["input_ids"])).logits + expected = float(F.cross_entropy(lg[:, :-1].reshape(-1, lg.size(-1)), labels[:, 1:].reshape(-1), + ignore_index=IGNORE)) + + loss_split = float(ns.sft_train_step(a, split, torch.optim.SGD(a.parameters(), lr=0.5))) + loss_whole = float(ns.sft_train_step(b, [full], torch.optim.SGD(b.parameters(), lr=0.5))) + assert abs(loss_whole - expected) < 1e-4, \ + f"the step should return the token-mean loss ({expected:.4f}), got {loss_whole:.4f}" + assert abs(loss_split - loss_whole) < 1e-4, ( + f"micro-batched step loss {loss_split:.4f} != one-batch {loss_whole:.4f}; count trained tokens " + "over all micro-batches before the backward passes") + for (name, pa), pb in zip(a.named_parameters(), b.parameters()): + assert torch.allclose(pa, pb, atol=1e-6), ( + f"parameter '{name}' differs between micro-batched and one-batch steps; averaging per-" + "micro-batch means weights tokens unequally") + + +def check_train_step_learns(ns): + torch.manual_seed(1) + exs = _uneven_examples(ns) + model = _TinyLM() + opt = torch.optim.Adam(model.parameters(), lr=1e-2) + batches = [ns.sft_collate(exs[:2], 0), ns.sft_collate(exs[2:], 0)] + first = float(ns.sft_train_step(model, batches, opt)) + for _ in range(50): + last = float(ns.sft_train_step(model, batches, opt, max_grad_norm=1.0)) + assert last < 0.5 * first, f"loss should fall on a fixed batch ({first:.3f} -> {last:.3f})" + + +CHECKS = [check_example_masks_prompt_and_keeps_eos, check_example_truncates_and_accepts_tensors, + check_collate_pads_right_with_ignored_labels, check_loss_is_shifted_token_mean, + check_loss_ignores_unused_logits, check_loss_num_items_makes_microbatches_add_up, + check_microbatching_matches_one_big_batch, check_train_step_learns] diff --git a/space/problems.json b/space/problems.json index a332029..43f5b7f 100644 --- a/space/problems.json +++ b/space/problems.json @@ -1,6 +1,100 @@ { - "count": 68, + "count": 74, "problems": [ + { + "id": "buffered-loading", + "legacy_ids": [], + "title": "Implement a Shuffle Buffer for Streaming Data", + "description": "Build a shuffle buffer that wraps an IterableDataset, providing approximate shuffling with O(buffer_size) memory.", + "difficulty": "medium", + "category": "data-loading", + "tracks": [ + "basics" + ], + "companies": [], + "question": "v3/data-loading/buffered-loading/buffered-loading.ipynb", + "solution": "v3/data-loading/buffered-loading/buffered-loading_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/buffered-loading/buffered-loading.ipynb", + "frequency": "emerging", + "company_confidence": {}, + "llm_path_order": null, + "graded": false, + "entries": [], + "hints": null + }, + { + "id": "data-parallel-loading", + "legacy_ids": [], + "title": "Simulate DistributedSampler for Data Parallel Training", + "description": "Implement a distributed sampler from scratch that shards a dataset across ranks with deterministic shuffling and padding.", + "difficulty": "medium", + "category": "data-loading", + "tracks": [ + "basics" + ], + "companies": [ + "Meta", + "Google", + "NVIDIA", + "Anthropic" + ], + "question": "v3/data-loading/data-parallel-loading/data-parallel-loading.ipynb", + "solution": "v3/data-loading/data-parallel-loading/data-parallel-loading_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/data-parallel-loading/data-parallel-loading.ipynb", + "frequency": "common", + "company_confidence": { + "Meta": "inferred", + "Google": "inferred", + "NVIDIA": "inferred", + "Anthropic": "inferred" + }, + "llm_path_order": null, + "graded": false, + "entries": [], + "hints": null + }, + { + "id": "preloaded-dataset", + "legacy_ids": [], + "title": "Implement a Preloaded Map-Style Dataset", + "description": "Build a map-style Dataset that loads all data into memory and integrates with DataLoader for batching and shuffling.", + "difficulty": "easy", + "category": "data-loading", + "tracks": [ + "basics" + ], + "companies": [], + "question": "v3/data-loading/preloaded-dataset/preloaded-dataset.ipynb", + "solution": "v3/data-loading/preloaded-dataset/preloaded-dataset_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/preloaded-dataset/preloaded-dataset.ipynb", + "frequency": "common", + "company_confidence": {}, + "llm_path_order": null, + "graded": false, + "entries": [], + "hints": null + }, + { + "id": "streaming-dataset", + "legacy_ids": [], + "title": "Implement a Streaming IterableDataset with Yield", + "description": "Build an IterableDataset that streams data line by line from a file using a generator, with multi-worker sharding support.", + "difficulty": "medium", + "category": "data-loading", + "tracks": [ + "basics" + ], + "companies": [], + "question": "v3/data-loading/streaming-dataset/streaming-dataset.ipynb", + "solution": "v3/data-loading/streaming-dataset/streaming-dataset_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/streaming-dataset/streaming-dataset.ipynb", + "frequency": "common", + "company_confidence": {}, + "llm_path_order": null, + "graded": false, + "entries": [], + "hints": null + }, { "id": "lin-regression", "legacy_ids": [ @@ -829,6 +923,44 @@ "L is the logsumexp of the scaled scores: running_max + log(l) at the end." ] }, + { + "id": "kl-divergence", + "legacy_ids": [ + "v2-1" + ], + "title": "Implement KL Divergence Loss", + "description": "Implement KL divergence every way it shows up in ML: from probabilities and logits, F.kl_div semantics, Jensen-Shannon, forward vs reverse, Gaussian closed form, masked token-level KL, and the k1/k2/k3 estimators used in PPO and GRPO.", + "difficulty": "medium", + "category": "alignment-training", + "tracks": [ + "advanced", + "llm-path" + ], + "companies": [], + "question": "v3/alignment-training/kl-divergence/kl-divergence.ipynb", + "solution": "v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/kl-divergence/kl-divergence.ipynb", + "frequency": "", + "company_confidence": {}, + "llm_path_order": 11, + "graded": true, + "entries": [ + "kl_categorical", + "kl_from_logits", + "my_kl_div", + "js_divergence", + "fit_gaussian", + "kl_gaussian", + "masked_token_kl", + "kl_estimators", + "kl_estimator_grad" + ], + "hints": [ + "Everything is sum(p * (log p - log q)) over the last dim. Use torch.xlogy when you have probabilities and F.log_softmax when you have logits, so 0 * log 0 and softmax underflow never turn into nan.", + "F.kl_div(input, target) is KL(target || exp(input)): input is the model's log-probs, the target comes first in the math, and 'batchmean' (sum / batch size) is the true KL while 'mean' divides by every element.", + "For samples x ~ p with log r = log q(x) - log p(x): k1 = -log r, k2 = (log r)^2 / 2, k3 = (r - 1) - log r. For kl_estimator_grad, weight every category by pi.detach() and let the gradient flow only through log pi." + ] + }, { "id": "rms-norm", "legacy_ids": [ @@ -844,14 +976,20 @@ ], "companies": [], "question": "llm/RMS-Norm/rms-norm.ipynb", - "solution": null, + "solution": "llm/RMS-Norm/rms-norm_SOLN.ipynb", "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/RMS-Norm/rms-norm.ipynb", "frequency": "emerging", "company_confidence": {}, "llm_path_order": null, - "graded": false, - "entries": [], - "hints": null + "graded": true, + "entries": [ + "RMSNorm" + ], + "hints": [ + "Square the input, take the mean along the last dimension, add epsilon, then take the square root.", + "Use keepdim=True in the mean so broadcasting works when you divide.", + "The scale parameter gamma should have shape (dim,) and be created with nn.Parameter(torch.ones(dim))." + ] }, { "id": "byte-pair-encoder", @@ -1071,6 +1209,39 @@ "RopeAttention: project, split into num_heads (Q) and kv_heads (K/V), rotate Q and K, repeat_kv the K/V, then scaled dot-product attention plus the additive mask." ] }, + { + "id": "sft-smollm", + "legacy_ids": [ + "v2-20" + ], + "title": "Apply SFT on SmolLM", + "description": "Fine-tune the pretrained SmolLM-135M base model into a chat assistant: build prompt-masked examples, pad and collate, write the shifted token-level loss, and normalize gradient accumulation correctly.", + "difficulty": "hard", + "category": "alignment-training", + "tracks": [ + "advanced", + "llm-path" + ], + "companies": [], + "question": "v3/alignment-training/sft-smollm/sft-smollm.ipynb", + "solution": "v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb", + "colab": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/sft-smollm/sft-smollm.ipynb", + "frequency": "", + "company_confidence": {}, + "llm_path_order": 13, + "graded": true, + "entries": [ + "build_sft_example", + "sft_collate", + "sft_loss", + "sft_train_step" + ], + "hints": [ + "labels are a copy of input_ids (prompt + response + end-of-turn) with every prompt position set to -100; keep them aligned with input_ids and shift inside the loss.", + "Pad input_ids with pad_id but labels with -100, so a pad id that equals a real token id never erases a real label.", + "sft_loss: cross_entropy(logits[:, :-1], labels[:, 1:], ignore_index=-100, reduction='sum') divided by the trained-token count of the whole optimizer step, counted across all micro-batches before any backward." + ] + }, { "id": "softmax", "legacy_ids": [ diff --git a/v3/alignment-training/kl-divergence/kl-divergence.ipynb b/v3/alignment-training/kl-divergence/kl-divergence.ipynb new file mode 100644 index 0000000..41ad4bb --- /dev/null +++ b/v3/alignment-training/kl-divergence/kl-divergence.ipynb @@ -0,0 +1,801 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/kl-divergence/kl-divergence.ipynb)\n", + "\n", + "Check your work: `!pip install torchleet` then `from torchleet import check; check(\"kl-divergence\", ...)`\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Implement KL Divergence Loss\n", + "\n", + "**Difficulty**: 🟑 Medium\n", + "\n", + "---\n", + "\n", + "### Problem Statement\n", + "\n", + "KL divergence is everywhere in modern ML: the VAE regularizer, the knowledge-distillation loss, the KL penalty in RLHF/PPO, the `Ξ²Β·KL` term in GRPO, the trust region in TRPO. Interviewers love it because a \"one-line\" function hides a lot of decisions:\n", + "\n", + "- **Which direction?** KL(Pβ€–Q) β‰  KL(Qβ€–P), and the choice changes what a model learns.\n", + "- **What do you have?** Probabilities, logits, log-probs, Gaussian parameters, or only *samples*.\n", + "- **Is it numerically stable?** `0 Β· log 0`, softmax underflow, `-inf` logits.\n", + "- **How do you reduce?** Over classes, the batch, tokens, and padding.\n", + "- **For RL: is the gradient right?** The gradient of an estimator *used as a loss* is not always the gradient of KL.\n", + "\n", + "$$\\mathrm{KL}(P\\,\\|\\,Q) = \\sum_x P(x)\\log\\frac{P(x)}{Q(x)} = \\mathbb{E}_{x\\sim P}\\big[\\log P(x) - \\log Q(x)\\big] \\;\\ge\\; 0$$\n", + "\n", + "It is zero iff P = Q, it is **not symmetric**, and it is not a metric.\n", + "\n", + "This notebook covers every variant you're likely to be asked about, in nine short parts. Each part is one function followed by a test cell.\n", + "\n", + "| Part | Function | What it practises |\n", + "|---|---|---|\n", + "| 1 | `kl_categorical(p, q)` | The definition; `0Β·log 0`; support mismatch β†’ `inf` |\n", + "| 2 | `kl_from_logits(p_logits, q_logits)` | Working in log space; `log_softmax`; stability |\n", + "| 3 | `my_kl_div(input, target, reduction, log_target)` | Re-implementing `F.kl_div`: argument order and `batchmean` |\n", + "| 4 | `js_divergence(p_logits, q_logits)` | A symmetric, bounded divergence; log-space mixtures |\n", + "| 5 | `fit_gaussian(target_logp, grid, direction)` | Forward vs reverse KL: mode-covering vs mode-seeking |\n", + "| 6 | `kl_gaussian(mu_p, logvar_p, mu_q, logvar_q)` | Closed form for diagonal Gaussians; the VAE term |\n", + "| 7 | `masked_token_kl(p_logits, q_logits, mask)` | Token-level KL for LLMs with padding (distillation, RLHF) |\n", + "| 8 | `kl_estimators(logp, logq)` | Monte Carlo k1 / k2 / k3 from samples (PPO, GRPO) |\n", + "| 9 | `kl_estimator_grad(logits, ref_logits, estimator)` | Which estimator, used as a loss, has the right gradient |\n", + "\n", + "---\n", + "\n", + "### Constraints\n", + "\n", + "- βœ… Distributions live on the **last** dimension; keep every leading (batch) dimension unless told otherwise.\n", + "- βœ… Work in log space wherever you can.\n", + "- ❌ Do **not** call `F.kl_div`, `nn.KLDivLoss`, `torch.distributions.kl_divergence` or any other KL helper inside your implementations. The tests use them as oracles.\n", + "- ❌ No Python loops over batch elements, tokens or classes.\n", + "\n", + "---\n", + "\n", + "
\n", + " πŸ’‘ Hint\n", + "\n", + " - Everything reduces to `Ξ£ p Β· (log p βˆ’ log q)`. The parts differ in where `p`, `log p` and `log q` come from.\n", + " - `torch.xlogy(x, y)` computes `x Β· log(y)` and returns 0 wherever `x == 0`.\n", + " - `F.log_softmax` and `torch.logsumexp` do the max-subtraction trick for you.\n", + " - The k1/k2/k3 estimators come from John Schulman's note *\"Approximating KL Divergence\"* (2020).\n", + "\n", + "
\n", + "\n", + "---" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import math\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from torch.distributions import Categorical, Normal, kl_divergence # oracles for the tests only\n", + "\n", + "torch.manual_seed(0)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 1: KL between two categorical distributions (probabilities in)\n", + "\n", + "Implement `kl_categorical(p, q)`. Both inputs are probability tensors of shape `(..., K)` that sum to 1 over the last dim. Return KL(Pβ€–Q) with shape `(...)`.\n", + "\n", + "The edge cases are the point:\n", + "\n", + "| case | contribution | why |\n", + "|---|---|---|\n", + "| `p_k = 0` | `0` | `lim xβ†’0 xΒ·log x = 0`. Naive `p * torch.log(p / q)` gives `0 Β· (-inf) = nan` |\n", + "| `p_k > 0`, `q_k = 0` | `+inf` | Q says impossible something that P says happens |\n", + "| `p_k = 0`, `q_k = 0` | `0` | Neither distribution can produce it |\n", + "\n", + "Your function should return `inf`, never `nan`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_categorical(p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " KL(P || Q) for categorical distributions given as probabilities.\n", + "\n", + " Args:\n", + " p, q: probabilities of shape (..., K), each summing to 1 over the last dim.\n", + "\n", + " Returns:\n", + " Tensor of shape (...): sum_k p_k * (log p_k - log q_k), with 0 * log 0 = 0.\n", + " \"\"\"\n", + " # TODO: compute p*log(p) - p*log(q) elementwise so that p == 0 contributes exactly 0\n", + " # (hint: torch.xlogy(x, y) = x * log(y), and is 0 wherever x == 0)\n", + " # TODO: sum over the last dimension\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 1 tests ===\n", + "p = torch.softmax(torch.randn(4, 6), -1)\n", + "q = torch.softmax(torch.randn(4, 6), -1)\n", + "got = kl_categorical(p, q)\n", + "ref = kl_divergence(Categorical(probs=p), Categorical(probs=q))\n", + "assert got.shape == (4,), f\"expected shape (4,), got {tuple(got.shape)}\"\n", + "assert torch.allclose(got, ref, atol=1e-6), f\"mismatch vs torch.distributions: {got} vs {ref}\"\n", + "assert torch.allclose(kl_categorical(p, p), torch.zeros(4), atol=1e-6), \"KL(P||P) must be 0\"\n", + "assert (got >= 0).all(), \"KL is never negative (Gibbs' inequality)\"\n", + "assert not torch.allclose(got, kl_categorical(q, p)), \"KL is not symmetric; check the argument order\"\n", + "\n", + "# Zeros in the support\n", + "p0 = torch.tensor([0.5, 0.5, 0.0])\n", + "q0 = torch.tensor([0.25, 0.25, 0.5])\n", + "assert torch.isclose(kl_categorical(p0, q0), torch.tensor(math.log(2.0))), \\\n", + " f\"p = 0 entries must contribute 0; expected log 2, got {kl_categorical(p0, q0)}\"\n", + "assert torch.isinf(kl_categorical(q0, p0)), \"q = 0 where p > 0 must give +inf (not nan)\"\n", + "print(\"Part 1 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 2: KL from logits (the version you'll actually use)\n", + "\n", + "Models output **logits**. Implement `kl_from_logits(p_logits, q_logits)` = KL(softmax(p_logits) β€– softmax(q_logits)), with output shape `(...)`.\n", + "\n", + "Why not just `kl_categorical(p_logits.softmax(-1), q_logits.softmax(-1))`? With confident (large) logits, softmax **underflows to exactly 0**, `log 0 = -inf`, and you get `inf` or `nan` even though the true KL is finite.\n", + "\n", + "Instead, get `log P` and `log Q` directly with `F.log_softmax`, which uses `logsumexp` (the max-subtraction trick) and never takes the log of an underflowed probability:\n", + "\n", + "$$\\mathrm{KL} = \\sum_k e^{\\log P_k}\\,\\big(\\log P_k - \\log Q_k\\big)$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_from_logits(p_logits: torch.Tensor, q_logits: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " KL(P || Q) where P = softmax(p_logits), Q = softmax(q_logits) over the last dim.\n", + "\n", + " Args:\n", + " p_logits, q_logits: unnormalized logits of shape (..., K).\n", + "\n", + " Returns:\n", + " Tensor of shape (...).\n", + " \"\"\"\n", + " # TODO: log P and log Q with F.log_softmax over the last dim\n", + " # TODO: sum P * (log P - log Q) over the last dim (P = exp(log P))\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 2 tests ===\n", + "pl, ql = torch.randn(3, 5, 10), torch.randn(3, 5, 10)\n", + "got = kl_from_logits(pl, ql)\n", + "ref = kl_divergence(Categorical(logits=pl), Categorical(logits=ql))\n", + "assert got.shape == (3, 5), f\"expected shape (3, 5), got {tuple(got.shape)}\"\n", + "assert torch.allclose(got, ref, atol=1e-5), \"mismatch vs torch.distributions\"\n", + "assert torch.allclose(kl_from_logits(pl + 7.0, ql - 3.0), got, atol=1e-5), \\\n", + " \"softmax is shift-invariant, so adding a constant to the logits must not change the KL\"\n", + "\n", + "# Confident logits: the naive probability route breaks, the log-space route doesn't\n", + "big_p, big_q = 100 * pl, 100 * ql\n", + "naive = kl_categorical(big_p.softmax(-1), big_q.softmax(-1))\n", + "stable = kl_from_logits(big_p, big_q)\n", + "print(f\"naive route: {int((~torch.isfinite(naive)).sum())} of {naive.numel()} entries are inf/nan\")\n", + "assert torch.isfinite(stable).all(), \"kl_from_logits must stay finite for large logits; use log_softmax\"\n", + "assert (stable >= -1e-4).all(), \"KL must be non-negative\"\n", + "print(\"Part 2 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 3: Re-implement `F.kl_div` (the API everyone gets wrong)\n", + "\n", + "```python\n", + "F.kl_div(input, target, reduction=\"mean\", log_target=False)\n", + "```\n", + "\n", + "- `input` is the **model's log-probabilities**, `log Q`. It is *not* probabilities.\n", + "- `target` is the **reference distribution P**: probabilities, or `log P` if `log_target=True`.\n", + "- The pointwise loss is `P Β· (log P βˆ’ input)`, with 0 where `P = 0`.\n", + "- So `F.kl_div(log_q, p)` computes **KL(P β€– Q)**: the second argument is the first argument of the KL.\n", + "\n", + "The reductions:\n", + "\n", + "| `reduction` | returns |\n", + "|---|---|\n", + "| `'none'` | the pointwise tensor, same shape as `input` |\n", + "| `'sum'` | sum of everything |\n", + "| `'mean'` | mean over **all elements**, i.e. divided by `BΒ·K`, **not** the KL |\n", + "| `'batchmean'` | `sum / input.size(0)`, the true mean KL per example |\n", + "\n", + "Implement `my_kl_div` so that it matches `F.kl_div` for all four reductions and both values of `log_target`. Then read the test output to see how far off `'mean'` is.\n", + "\n", + "> PyTorch warns that `'mean'` will be changed to behave like `'batchmean'` \"in the next major release\". Through torch 2.x it still divides by `BΒ·K`, so always pass `reduction='batchmean'` explicitly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def my_kl_div(input: torch.Tensor, target: torch.Tensor,\n", + " reduction: str = \"mean\", log_target: bool = False) -> torch.Tensor:\n", + " \"\"\"\n", + " Same semantics as torch.nn.functional.kl_div.\n", + "\n", + " Args:\n", + " input: log-probabilities of the model distribution Q, any shape.\n", + " target: P as probabilities (log_target=False) or log-probabilities (log_target=True).\n", + " reduction: 'none' | 'sum' | 'mean' | 'batchmean'.\n", + " log_target: whether target is given in log space.\n", + " \"\"\"\n", + " # TODO: pointwise term P * (log P - input)\n", + " # - log_target=True: target is log P\n", + " # - log_target=False: target is P; make P == 0 contribute 0 (Part 1 trick)\n", + " # TODO: apply the reduction ('none', 'sum', 'mean' over all elements, 'batchmean' = sum / batch size)\n", + " # TODO: raise ValueError for an unknown reduction\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 3 tests ===\n", + "B, K = 8, 5\n", + "log_q = F.log_softmax(torch.randn(B, K), -1)\n", + "p = torch.softmax(torch.randn(B, K), -1)\n", + "p_pos = p.clone() # strictly positive: log P is finite\n", + "p[0, 0] = 0.0\n", + "p[0] = p[0] / p[0].sum() # a probability target with an exact zero\n", + "for log_target in (False, True):\n", + " tgt = p_pos.log() if log_target else p\n", + " for red in (\"none\", \"sum\", \"mean\", \"batchmean\"):\n", + " got = my_kl_div(log_q, tgt, reduction=red, log_target=log_target)\n", + " exp = F.kl_div(log_q, tgt, reduction=red, log_target=log_target)\n", + " assert torch.allclose(got, exp, atol=1e-6), \\\n", + " f\"reduction={red}, log_target={log_target}: {got} vs F.kl_div {exp}\"\n", + "\n", + "true_kl = kl_categorical(p, log_q.exp()).mean()\n", + "print(f\"true mean KL(P||Q) = {true_kl:.5f}\")\n", + "print(f\"F.kl_div 'batchmean' = {F.kl_div(log_q, p, reduction='batchmean'):.5f}\")\n", + "print(f\"F.kl_div 'mean' default = {F.kl_div(log_q, p):.5f} <- off by a factor of K = {K}\")\n", + "assert torch.allclose(my_kl_div(log_q, p, reduction=\"batchmean\"), true_kl, atol=1e-6)\n", + "print(\"Part 3 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 4: Jensen–Shannon divergence (symmetric and bounded)\n", + "\n", + "$$\\mathrm{JS}(P, Q) = \\tfrac12\\,\\mathrm{KL}(P\\,\\|\\,M) + \\tfrac12\\,\\mathrm{KL}(Q\\,\\|\\,M), \\qquad M = \\tfrac12(P + Q)$$\n", + "\n", + "Unlike KL, JS is **symmetric**, **bounded** in `[0, log 2]`, and **finite even when the supports don't overlap**, because M covers both. It is the divergence behind the original GAN objective, and its generalized (Ξ²-weighted) form is the loss in GKD distillation.\n", + "\n", + "Implement `js_divergence(p_logits, q_logits)`, with output shape `(...)`. Compute `log M` in log space: `log M = logsumexp([log P, log Q]) βˆ’ log 2`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def js_divergence(p_logits: torch.Tensor, q_logits: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " Jensen-Shannon divergence between softmax(p_logits) and softmax(q_logits).\n", + "\n", + " Args:\n", + " p_logits, q_logits: logits of shape (..., K).\n", + "\n", + " Returns:\n", + " Tensor of shape (...), in [0, log 2].\n", + " \"\"\"\n", + " # TODO: log P, log Q via log_softmax\n", + " # TODO: log M = logsumexp of the stacked log-probs, minus log 2\n", + " # TODO: 0.5 * KL(P || M) + 0.5 * KL(Q || M), computed from log-probs\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 4 tests ===\n", + "pl, ql = torch.randn(6, 7), torch.randn(6, 7)\n", + "got = js_divergence(pl, ql)\n", + "P, Q = pl.softmax(-1), ql.softmax(-1)\n", + "M = Categorical(probs=0.5 * (P + Q))\n", + "ref = 0.5 * kl_divergence(Categorical(probs=P), M) + 0.5 * kl_divergence(Categorical(probs=Q), M)\n", + "assert got.shape == (6,), f\"expected shape (6,), got {tuple(got.shape)}\"\n", + "assert torch.allclose(got, ref, atol=1e-5), \"mismatch vs 0.5*KL(P||M) + 0.5*KL(Q||M)\"\n", + "assert torch.allclose(got, js_divergence(ql, pl), atol=1e-6), \"JS must be symmetric\"\n", + "assert (got >= -1e-6).all() and (got <= math.log(2) + 1e-6).all(), \"JS must lie in [0, log 2]\"\n", + "\n", + "# Disjoint supports: KL explodes, JS saturates at log 2\n", + "far_p = torch.tensor([50.0, -50.0, -50.0, -50.0])\n", + "far_q = torch.tensor([-50.0, -50.0, -50.0, 50.0])\n", + "print(f\"KL = {kl_from_logits(far_p, far_q):.1f} JS = {js_divergence(far_p, far_q):.4f} log 2 = {math.log(2):.4f}\")\n", + "assert abs(js_divergence(far_p, far_q).item() - math.log(2)) < 1e-4, \"disjoint supports must give JS = log 2\"\n", + "print(\"Part 4 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 5: Forward vs reverse KL (mode-covering vs mode-seeking)\n", + "\n", + "Fit a **single** Gaussian `q_ΞΈ` (parameters `ΞΌ`, `Οƒ`) to a **bimodal** target `p` in two ways:\n", + "\n", + "| objective | penalizes | behaviour | where you meet it |\n", + "|---|---|---|---|\n", + "| **forward** KL(p β€– q_ΞΈ) = E_p[log p βˆ’ log q] | `q` small where `p` is large | **mode-covering**: stretches over both modes | MLE, cross-entropy, SFT, classic distillation |\n", + "| **reverse** KL(q_ΞΈ β€– p) = E_q[log q βˆ’ log p] | `q` large where `p` is small | **mode-seeking**: hides inside one mode | variational inference, KL-regularized RL, MiniLLM |\n", + "\n", + "To keep everything categorical, discretize the real line on a `grid`. `q_ΞΈ`'s logits on the grid are `-Β½((grid βˆ’ ΞΌ)/Οƒ)Β²`; `log_softmax` normalizes them for you. Then reuse `kl_from_logits`, remembering that it takes **logits**, and that `target_logp` is already a valid set of logits.\n", + "\n", + "Implement `fit_gaussian`: optimize `ΞΌ` and `log Οƒ` (so Οƒ stays positive) with Adam, starting from `mu_init` and `sigma_init`, and return `(ΞΌ, Οƒ)` as Python floats.\n", + "\n", + "Predict the two answers before you run the test." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def fit_gaussian(target_logp: torch.Tensor, grid: torch.Tensor, direction: str,\n", + " steps: int = 500, lr: float = 0.05,\n", + " mu_init: float = 1.0, sigma_init: float = 0.5) -> tuple:\n", + " \"\"\"\n", + " Fit a discretized Gaussian q to a target distribution by minimizing KL.\n", + "\n", + " Args:\n", + " target_logp: log-probabilities of the target on the grid, shape (G,).\n", + " grid: grid points, shape (G,).\n", + " direction: 'forward' -> minimize KL(target || q); 'reverse' -> minimize KL(q || target).\n", + " steps, lr: Adam iterations and learning rate.\n", + " mu_init, sigma_init: starting point.\n", + "\n", + " Returns:\n", + " (mu, sigma) as Python floats.\n", + " \"\"\"\n", + " # TODO: make mu and log_sigma leaf tensors with requires_grad=True\n", + " # TODO: Adam over [mu, log_sigma]\n", + " # TODO: each step: q_logits = -0.5 * ((grid - mu) / sigma) ** 2\n", + " # forward loss = kl_from_logits(target_logp, q_logits)\n", + " # reverse loss = kl_from_logits(q_logits, target_logp)\n", + " # TODO: return (mu, sigma) as floats\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 5 tests ===\n", + "grid = torch.linspace(-6, 6, 241)\n", + "# Two well-separated modes at -2.5 and +2.5, each with std 0.6. Built in log space so the\n", + "# valley between the modes stays finite (in probability space it would underflow to 0).\n", + "mixture = torch.logsumexp(torch.stack([Normal(-2.5, 0.6).log_prob(grid),\n", + " Normal(2.5, 0.6).log_prob(grid)]), dim=0)\n", + "target_logp = F.log_softmax(mixture, dim=-1)\n", + "\n", + "mu_f, sigma_f = fit_gaussian(target_logp, grid, \"forward\")\n", + "mu_r, sigma_r = fit_gaussian(target_logp, grid, \"reverse\")\n", + "print(f\"forward KL(p||q): mu = {mu_f:+.2f}, sigma = {sigma_f:.2f} (covers both modes)\")\n", + "print(f\"reverse KL(q||p): mu = {mu_r:+.2f}, sigma = {sigma_r:.2f} (locks onto one mode)\")\n", + "\n", + "assert abs(mu_f) < 0.5 and sigma_f > 2.0, \"forward KL should centre between the modes with a wide sigma\"\n", + "assert abs(mu_r - 2.5) < 0.3 and sigma_r < 1.0, \"reverse KL should sit on the +2.5 mode with a narrow sigma\"\n", + "print(\"Part 5 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> **Try it:** rerun the reverse fit with `mu_init=0.5, sigma_init=1.0`. It ends up wide and centred, stuck between the modes. Reverse KL is non-convex in ΞΈ, so *where you start decides which mode, or which compromise, you get*. That's one reason RL fine-tuning with a KL penalty is sensitive to initialization (the SFT model)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 6: Closed-form KL between diagonal Gaussians (the VAE term)\n", + "\n", + "For continuous distributions the sum becomes an integral. For two 1-D Gaussians it has a closed form:\n", + "\n", + "$$\\mathrm{KL}\\big(\\mathcal N(\\mu_p,\\sigma_p^2)\\,\\|\\,\\mathcal N(\\mu_q,\\sigma_q^2)\\big) = \\log\\frac{\\sigma_q}{\\sigma_p} + \\frac{\\sigma_p^2 + (\\mu_p-\\mu_q)^2}{2\\sigma_q^2} - \\frac12$$\n", + "\n", + "For **diagonal** Gaussians in `D` dimensions, the KL is the sum over dimensions.\n", + "\n", + "Implement `kl_gaussian` taking **log-variances** (that's what an encoder outputs, because it keeps σ² > 0 for free). Reduce over the last dim.\n", + "\n", + "Then check the special case every VAE uses, with `q = N(0, I)`:\n", + "\n", + "$$\\mathrm{KL}\\big(\\mathcal N(\\mu,\\sigma^2)\\,\\|\\,\\mathcal N(0, I)\\big) = -\\tfrac12\\sum_d\\big(1 + \\log\\sigma_d^2 - \\mu_d^2 - \\sigma_d^2\\big)$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_gaussian(mu_p: torch.Tensor, logvar_p: torch.Tensor,\n", + " mu_q: torch.Tensor, logvar_q: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " KL( N(mu_p, diag(exp(logvar_p))) || N(mu_q, diag(exp(logvar_q))) ).\n", + "\n", + " Args:\n", + " mu_p, logvar_p, mu_q, logvar_q: tensors of shape (..., D) (broadcastable).\n", + "\n", + " Returns:\n", + " Tensor of shape (...): KL summed over the D dimensions.\n", + " \"\"\"\n", + " # TODO: write the 1-D formula in terms of log-variances:\n", + " # 0.5 * (logvar_q - logvar_p + (var_p + (mu_p - mu_q)^2) / var_q - 1)\n", + " # TODO: sum over the last dimension\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 6 tests ===\n", + "mu_p, lv_p = torch.randn(4, 3), torch.randn(4, 3)\n", + "mu_q, lv_q = torch.randn(4, 3), torch.randn(4, 3)\n", + "got = kl_gaussian(mu_p, lv_p, mu_q, lv_q)\n", + "ref = kl_divergence(Normal(mu_p, (0.5 * lv_p).exp()), Normal(mu_q, (0.5 * lv_q).exp())).sum(-1)\n", + "assert got.shape == (4,), f\"expected shape (4,), got {tuple(got.shape)}\"\n", + "assert torch.allclose(got, ref, atol=1e-5), \"mismatch vs torch.distributions.Normal\"\n", + "assert torch.allclose(kl_gaussian(mu_p, lv_p, mu_p, lv_p), torch.zeros(4), atol=1e-6), \"KL(p||p) must be 0\"\n", + "\n", + "# The VAE special case: q = N(0, I)\n", + "zeros = torch.zeros_like(mu_p)\n", + "vae = -0.5 * (1 + lv_p - mu_p ** 2 - lv_p.exp()).sum(-1)\n", + "assert torch.allclose(kl_gaussian(mu_p, lv_p, zeros, zeros), vae, atol=1e-5), \"VAE closed form mismatch\"\n", + "print(\"Part 6 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 7: Token-level KL for language models (with padding)\n", + "\n", + "In LLM distillation and RLHF you compare two models **at every position** of a batch of sequences:\n", + "\n", + "- logits `(B, T, V)` from each model,\n", + "- a `mask` of shape `(B, T)`: 1 for real tokens, 0 for padding (or for prompt tokens you don't want to score).\n", + "\n", + "Implement `masked_token_kl(p_logits, q_logits, mask)`. It should:\n", + "\n", + "1. compute the exact, full-vocabulary KL(P_t β€– Q_t) at every position,\n", + "2. **ignore** masked positions completely, even if their logits are garbage, and\n", + "3. return a scalar: the **mean over unmasked tokens** (not over sequences).\n", + "\n", + "Directions in practice:\n", + "\n", + "- **forward distillation**: `masked_token_kl(teacher_logits.detach(), student_logits, mask)`\n", + "- **reverse distillation / RL penalty**: `masked_token_kl(student_logits, teacher_logits.detach(), mask)`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def masked_token_kl(p_logits: torch.Tensor, q_logits: torch.Tensor,\n", + " mask: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " Mean per-token KL(P_t || Q_t) over unmasked positions.\n", + "\n", + " Args:\n", + " p_logits, q_logits: logits of shape (B, T, V).\n", + " mask: (B, T), 1 (or True) for tokens to include, 0 for padding.\n", + "\n", + " Returns:\n", + " Scalar tensor.\n", + " \"\"\"\n", + " # TODO: per-token KL of shape (B, T) (reuse kl_from_logits)\n", + " # TODO: zero out masked positions and divide by the number of unmasked tokens\n", + " # (guard against an all-zero mask)\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 7 tests ===\n", + "B, T, V = 3, 6, 11\n", + "pl, ql = torch.randn(B, T, V), torch.randn(B, T, V)\n", + "pl[2] *= 4 # the short sequence disagrees much more\n", + "lengths = [6, 4, 2]\n", + "mask = torch.zeros(B, T)\n", + "for i, n in enumerate(lengths):\n", + " mask[i, :n] = 1\n", + "\n", + "got = masked_token_kl(pl, ql, mask)\n", + "oracle = torch.stack([kl_divergence(Categorical(logits=pl[i, t]), Categorical(logits=ql[i, t]))\n", + " for i in range(B) for t in range(lengths[i])]).mean()\n", + "assert got.ndim == 0, f\"expected a scalar, got shape {tuple(got.shape)}\"\n", + "assert torch.allclose(got, oracle, atol=1e-5), f\"expected the mean over the {sum(lengths)} real tokens\"\n", + "\n", + "# Garbage at padded positions must not leak in\n", + "pl_garbage, ql_garbage = pl.clone(), ql.clone()\n", + "pl_garbage[mask == 0] = 1e3 * torch.randn(int((mask == 0).sum()), V)\n", + "ql_garbage[mask == 0] = -1e3 * torch.randn(int((mask == 0).sum()), V)\n", + "assert torch.allclose(masked_token_kl(pl_garbage, ql_garbage, mask), got, atol=1e-5), \\\n", + " \"changing logits at padded positions changed the result\"\n", + "\n", + "# Gradients reach real tokens only\n", + "qg = ql.clone().requires_grad_(True)\n", + "masked_token_kl(pl, qg, mask).backward()\n", + "assert qg.grad[mask == 0].abs().max() == 0, \"padded positions must get zero gradient\"\n", + "assert qg.grad[mask == 1].abs().max() > 0, \"real positions must get a gradient\"\n", + "\n", + "# Token mean, not sequence mean\n", + "seq_mean = torch.stack([kl_from_logits(pl[i, :n], ql[i, :n]).mean() for i, n in enumerate(lengths)]).mean()\n", + "print(f\"token mean = {got:.4f} sequence mean = {seq_mean:.4f}\")\n", + "print(\"(the 2-token sequence is 2/12 of the token mean but 1/3 of the sequence mean)\")\n", + "print(\"Part 7 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 8: Estimating KL from samples (k1, k2, k3)\n", + "\n", + "In RL fine-tuning of LLMs, the exact KL over whole sequences is a sum over `V^T` possibilities, which is impossible. Often all you have is the **log-prob of the token that was actually sampled**, under the policy and under the reference model. So you estimate KL(Ο€ β€– Ο€_ref) by Monte Carlo, with `x ~ Ο€`.\n", + "\n", + "Let `logp = log Ο€(x)` and `logq = log Ο€_ref(x)`, and define the log-ratio `log r = logq βˆ’ logp`. Schulman's three estimators are:\n", + "\n", + "| estimator | formula | unbiased? | sign | variance |\n", + "|---|---|---|---|---|\n", + "| **k1** | `βˆ’log r` | βœ… | can be **negative** | high |\n", + "| **k2** | `Β½ (log r)Β²` | ❌ (slightly biased) | β‰₯ 0 | low |\n", + "| **k3** | `(r βˆ’ 1) βˆ’ log r` | βœ… | β‰₯ 0 | low |\n", + "\n", + "- **Why k3 is unbiased:** `E_{x~Ο€}[r] = Ξ£ Ο€ Β· (Ο€_ref/Ο€) = 1`, so `(r βˆ’ 1)` is a zero-mean **control variate** added to k1.\n", + "- **Why k3 is β‰₯ 0:** `log r ≀ r βˆ’ 1` for every r.\n", + "\n", + "PPO-RLHF uses k1 per token. GRPO uses k3.\n", + "\n", + "Implement `kl_estimators(logp, logq)`, returning the per-sample tuple `(k1, k2, k3)`, each with the same shape as the inputs. Then check the bias and variance experimentally." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_estimators(logp: torch.Tensor, logq: torch.Tensor) -> tuple:\n", + " \"\"\"\n", + " Per-sample Monte Carlo estimators of KL(p || q), for samples x ~ p.\n", + "\n", + " Args:\n", + " logp: log p(x) of the sampled x (the sampling distribution, e.g. the policy).\n", + " logq: log q(x) of the same x (e.g. the reference model). Same shape as logp.\n", + "\n", + " Returns:\n", + " (k1, k2, k3), each the same shape as logp.\n", + " \"\"\"\n", + " # TODO: log_ratio = logq - logp\n", + " # TODO: k1 = -log_ratio\n", + " # TODO: k2 = 0.5 * log_ratio ** 2\n", + " # TODO: k3 = (ratio - 1) - log_ratio, with ratio = exp(log_ratio)\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 8 tests ===\n", + "torch.manual_seed(0)\n", + "N = 200_000\n", + "print(f\"{'setting':>8} {'true KL':>9} | {'k1 bias/KL':>10} {'k1 std/KL':>9} | \"\n", + " f\"{'k2 bias/KL':>10} {'k2 std/KL':>9} | {'k3 bias/KL':>10} {'k3 std/KL':>9}\")\n", + "for name, scale in ((\"close\", 0.1), (\"far\", 1.0)):\n", + " policy_logits = torch.randn(50)\n", + " ref_logits = policy_logits + scale * torch.randn(50)\n", + " true = kl_from_logits(policy_logits, ref_logits).item()\n", + " x = Categorical(logits=policy_logits).sample((N,)) # samples from the policy\n", + " logp = F.log_softmax(policy_logits, -1)[x]\n", + " logq = F.log_softmax(ref_logits, -1)[x]\n", + " k1, k2, k3 = kl_estimators(logp, logq)\n", + " row = [f\"{(k.mean().item() - true) / true:>10.3f} {k.std().item() / true:>9.2f}\" for k in (k1, k2, k3)]\n", + " print(f\"{name:>8} {true:>9.5f} | \" + \" | \".join(row))\n", + "\n", + " se = lambda k: k.std().item() / math.sqrt(N)\n", + " assert abs(k1.mean().item() - true) < 5 * se(k1), \"k1 should be unbiased\"\n", + " assert abs(k3.mean().item() - true) < 5 * se(k3), \"k3 should be unbiased\"\n", + " assert (k2 >= 0).all() and (k3 >= -1e-6).all(), \"k2 and k3 are non-negative per sample\"\n", + " assert (k1 < 0).any(), \"k1 should produce some negative per-sample estimates\"\n", + " assert k3.std() < k1.std(), \"k3 should have lower variance than k1\"\n", + " if name == \"far\":\n", + " assert k2.mean().item() - true > 0.1 * true, \"k2's bias should show up when p and q are far apart\"\n", + "print(\"Part 8 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 9: KL as a loss. Does the estimator have the right gradient?\n", + "\n", + "Part 8 was about **values**. In training you **differentiate**. A common pattern (GRPO is one example) is to add `Ξ² Β· mean(k(x))` straight to the loss, with the samples `x ~ Ο€_ΞΈ` treated as fixed data. Autograd then differentiates `k` *only through `log Ο€_ΞΈ(x)`*, never through the sampling. In expectation the gradient you get is\n", + "\n", + "$$g_k = \\mathbb{E}_{x\\sim\\pi_\\theta}\\big[\\nabla_\\theta\\, k(x;\\theta)\\big].$$\n", + "\n", + "Is `g_k` equal to `βˆ‡_ΞΈ KL(Ο€_ΞΈ β€– Ο€_ref)`?\n", + "\n", + "For a single categorical \"policy\" with logits ΞΈ, you can compute that expectation **exactly**, without sampling noise, by weighting every category by its **detached** probability:\n", + "\n", + "$$\\text{surrogate} = \\sum_x \\mathrm{sg}\\big(\\pi_\\theta(x)\\big)\\, k(x;\\theta), \\qquad g_k = \\nabla_\\theta\\,\\text{surrogate}$$\n", + "\n", + "where `sg` is stop-gradient (`.detach()`).\n", + "\n", + "Implement `kl_estimator_grad(logits, ref_logits, estimator)`, which returns `g_k` w.r.t. `logits` (same shape as `logits`) for `estimator` ∈ `{'k1', 'k2', 'k3'}`. The test compares each one against autograd of the exact **reverse** KL(Ο€β€–Ο€_ref) and the exact **forward** KL(Ο€_refβ€–Ο€).\n", + "\n", + "**Write down your prediction for each of the three before running the test.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_estimator_grad(logits: torch.Tensor, ref_logits: torch.Tensor, estimator: str) -> torch.Tensor:\n", + " \"\"\"\n", + " Expected gradient (w.r.t. logits) of an on-policy KL estimator used as a loss.\n", + "\n", + " Args:\n", + " logits: policy logits, shape (..., V).\n", + " ref_logits: reference logits, shape (..., V); no gradient flows into them.\n", + " estimator: 'k1', 'k2' or 'k3' (see kl_estimators).\n", + "\n", + " Returns:\n", + " Tensor with the same shape as logits: grad of sum_x sg(pi(x)) * k(x) w.r.t. logits.\n", + " \"\"\"\n", + " # TODO: fresh leaf copy of logits with requires_grad=True\n", + " # TODO: logp = log_softmax(logits); logq = log_softmax(ref_logits), detached\n", + " # TODO: per-category k = kl_estimators(logp, logq)[...] for the chosen estimator\n", + " # TODO: surrogate = sum over everything of pi.detach() * k\n", + " # TODO: return torch.autograd.grad(surrogate, logits)[0]\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 9 tests ===\n", + "torch.manual_seed(0)\n", + "policy_logits, ref_logits = torch.randn(4, 10), torch.randn(4, 10)\n", + "\n", + "theta = policy_logits.clone().requires_grad_(True)\n", + "grad_reverse = torch.autograd.grad(kl_from_logits(theta, ref_logits).sum(), theta)[0] # KL(pi || ref)\n", + "grad_forward = torch.autograd.grad(kl_from_logits(ref_logits, theta).sum(), theta)[0] # KL(ref || pi)\n", + "\n", + "grads = {e: kl_estimator_grad(policy_logits, ref_logits, e) for e in (\"k1\", \"k2\", \"k3\")}\n", + "print(f\"{'as loss':>8} | {'|g|':>7} | {'max |g - grad reverse KL|':>26} | {'max |g - grad forward KL|':>26}\")\n", + "for e, g in grads.items():\n", + " assert g.shape == policy_logits.shape, f\"{e}: expected shape {tuple(policy_logits.shape)}\"\n", + " print(f\"{e:>8} | {g.norm():7.4f} | {(g - grad_reverse).abs().max():26.2e} | {(g - grad_forward).abs().max():26.2e}\")\n", + "\n", + "assert torch.allclose(grads[\"k1\"], torch.zeros_like(policy_logits), atol=1e-6), \\\n", + " \"k1 as a loss: E[grad log pi] = 0, so its expected gradient should vanish\"\n", + "assert torch.allclose(grads[\"k2\"], grad_reverse, atol=1e-5), \\\n", + " \"k2 as a loss should reproduce the gradient of the reverse KL(pi || ref)\"\n", + "assert torch.allclose(grads[\"k3\"], grad_forward, atol=1e-5), \\\n", + " \"k3 as a loss: its expected gradient is that of the forward KL(ref || pi)\"\n", + "print(\"Part 9 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**What just happened.** With samples treated as data, the gradient of each estimator is `c(x)Β·βˆ‡log Ο€(x)`:\n", + "\n", + "| used as a loss | per-sample gradient | expected gradient |\n", + "|---|---|---|\n", + "| k1 | `βˆ‡log Ο€` | **0**, because `E_Ο€[βˆ‡log Ο€] = Ξ£ βˆ‡Ο€ = βˆ‡1 = 0` |\n", + "| k2 | `(log Ο€ βˆ’ log Ο€_ref)Β·βˆ‡log Ο€` | βˆ‡ **reverse** KL(Ο€β€–Ο€_ref), exactly |\n", + "| k3 | `(1 βˆ’ Ο€_ref/Ο€)Β·βˆ‡log Ο€` | `βˆ’E_{Ο€_ref}[βˆ‡log Ο€]` = βˆ‡ **forward** KL(Ο€_refβ€–Ο€) |\n", + "\n", + "So **an unbiased estimator of the value (k1, k3) is not automatically a correct loss for the gradient**. The PPO-RLHF recipe, which puts k1 **in the reward** (detached, multiplied by the score function), has the same gradient as k2-as-loss. Both are correct for reverse KL. GRPO's k3-as-loss matches it only to first order near Ο€ = Ο€_ref." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Interview follow-ups\n", + "\n", + "Be ready to answer these out loud once the tests pass:\n", + "\n", + "1. `F.kl_div(a, b)` computes KL of *what* against *what*? What does the default `reduction='mean'` return, and why is it wrong?\n", + "2. Training with cross-entropy (or SFT) minimizes which KL direction, and against which distribution?\n", + "3. You're distilling a 70B teacher into a 1B student. Forward or reverse KL? What do MiniLLM and GKD do?\n", + "4. In Hinton-style distillation with temperature `T`, why is the soft loss multiplied by `TΒ²`?\n", + "5. PPO-RLHF subtracts `Ξ²Β·(log Ο€ βˆ’ log Ο€_ref)` from the **reward**. GRPO adds `Ξ²Β·k3` to the **loss**. Are these optimizing the same thing? (Use Part 9.)\n", + "6. Your top-k sampler sets filtered logits to `-inf`, and now your KL is `nan`. Why, and how do you fix it?\n", + "7. Write the KL between two **full-covariance** Gaussians. How do you compute it stably?\n", + "8. In `masked_token_kl`, why average over tokens rather than over sequences? When would you choose otherwise?\n", + "9. Your RL samples come from an *older* policy (off-policy, e.g. several PPO epochs per batch). Is the k3 estimate still unbiased?" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb b/v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb new file mode 100644 index 0000000..d962b5d --- /dev/null +++ b/v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb @@ -0,0 +1,830 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Implement KL Divergence Loss β€” Solution\n", + "\n", + "**Difficulty**: 🟑 Medium\n", + "\n", + "---\n", + "\n", + "### Problem Statement\n", + "\n", + "KL divergence is everywhere in modern ML: the VAE regularizer, the knowledge-distillation loss, the KL penalty in RLHF/PPO, the `Ξ²Β·KL` term in GRPO, the trust region in TRPO. Interviewers love it because a \"one-line\" function hides a lot of decisions:\n", + "\n", + "- **Which direction?** KL(Pβ€–Q) β‰  KL(Qβ€–P), and the choice changes what a model learns.\n", + "- **What do you have?** Probabilities, logits, log-probs, Gaussian parameters, or only *samples*.\n", + "- **Is it numerically stable?** `0 Β· log 0`, softmax underflow, `-inf` logits.\n", + "- **How do you reduce?** Over classes, the batch, tokens, and padding.\n", + "- **For RL: is the gradient right?** The gradient of an estimator *used as a loss* is not always the gradient of KL.\n", + "\n", + "$$\\mathrm{KL}(P\\,\\|\\,Q) = \\sum_x P(x)\\log\\frac{P(x)}{Q(x)} = \\mathbb{E}_{x\\sim P}\\big[\\log P(x) - \\log Q(x)\\big] \\;\\ge\\; 0$$\n", + "\n", + "It is zero iff P = Q, it is **not symmetric**, and it is not a metric.\n", + "\n", + "This notebook covers every variant you're likely to be asked about, in nine short parts. Each part is one function followed by a test cell.\n", + "\n", + "| Part | Function | What it practises |\n", + "|---|---|---|\n", + "| 1 | `kl_categorical(p, q)` | The definition; `0Β·log 0`; support mismatch β†’ `inf` |\n", + "| 2 | `kl_from_logits(p_logits, q_logits)` | Working in log space; `log_softmax`; stability |\n", + "| 3 | `my_kl_div(input, target, reduction, log_target)` | Re-implementing `F.kl_div`: argument order and `batchmean` |\n", + "| 4 | `js_divergence(p_logits, q_logits)` | A symmetric, bounded divergence; log-space mixtures |\n", + "| 5 | `fit_gaussian(target_logp, grid, direction)` | Forward vs reverse KL: mode-covering vs mode-seeking |\n", + "| 6 | `kl_gaussian(mu_p, logvar_p, mu_q, logvar_q)` | Closed form for diagonal Gaussians; the VAE term |\n", + "| 7 | `masked_token_kl(p_logits, q_logits, mask)` | Token-level KL for LLMs with padding (distillation, RLHF) |\n", + "| 8 | `kl_estimators(logp, logq)` | Monte Carlo k1 / k2 / k3 from samples (PPO, GRPO) |\n", + "| 9 | `kl_estimator_grad(logits, ref_logits, estimator)` | Which estimator, used as a loss, has the right gradient |\n", + "\n", + "---\n", + "\n", + "### Constraints\n", + "\n", + "- βœ… Distributions live on the **last** dimension; keep every leading (batch) dimension unless told otherwise.\n", + "- βœ… Work in log space wherever you can.\n", + "- ❌ Do **not** call `F.kl_div`, `nn.KLDivLoss`, `torch.distributions.kl_divergence` or any other KL helper inside your implementations. The tests use them as oracles.\n", + "- ❌ No Python loops over batch elements, tokens or classes.\n", + "\n", + "---\n", + "\n", + "
\n", + " πŸ’‘ Hint\n", + "\n", + " - Everything reduces to `Ξ£ p Β· (log p βˆ’ log q)`. The parts differ in where `p`, `log p` and `log q` come from.\n", + " - `torch.xlogy(x, y)` computes `x Β· log(y)` and returns 0 wherever `x == 0`.\n", + " - `F.log_softmax` and `torch.logsumexp` do the max-subtraction trick for you.\n", + " - The k1/k2/k3 estimators come from John Schulman's note *\"Approximating KL Divergence\"* (2020).\n", + "\n", + "
\n", + "\n", + "---" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import math\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from torch.distributions import Categorical, Normal, kl_divergence # oracles for the tests only\n", + "\n", + "torch.manual_seed(0)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 1: KL between two categorical distributions (probabilities in)\n", + "\n", + "Implement `kl_categorical(p, q)`. Both inputs are probability tensors of shape `(..., K)` that sum to 1 over the last dim. Return KL(Pβ€–Q) with shape `(...)`.\n", + "\n", + "The edge cases are the point:\n", + "\n", + "| case | contribution | why |\n", + "|---|---|---|\n", + "| `p_k = 0` | `0` | `lim xβ†’0 xΒ·log x = 0`. Naive `p * torch.log(p / q)` gives `0 Β· (-inf) = nan` |\n", + "| `p_k > 0`, `q_k = 0` | `+inf` | Q says impossible something that P says happens |\n", + "| `p_k = 0`, `q_k = 0` | `0` | Neither distribution can produce it |\n", + "\n", + "Your function should return `inf`, never `nan`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_categorical(p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " KL(P || Q) for categorical distributions given as probabilities.\n", + "\n", + " Args:\n", + " p, q: probabilities of shape (..., K), each summing to 1 over the last dim.\n", + "\n", + " Returns:\n", + " Tensor of shape (...): sum_k p_k * (log p_k - log q_k), with 0 * log 0 = 0.\n", + " \"\"\"\n", + " # xlogy(0, anything) == 0, so p == 0 entries vanish instead of producing 0 * -inf = nan.\n", + " # p > 0 with q == 0 gives xlogy(p, 0) = -inf, i.e. KL = +inf, which is the right answer.\n", + " return (torch.xlogy(p, p) - torch.xlogy(p, q)).sum(-1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 1 tests ===\n", + "p = torch.softmax(torch.randn(4, 6), -1)\n", + "q = torch.softmax(torch.randn(4, 6), -1)\n", + "got = kl_categorical(p, q)\n", + "ref = kl_divergence(Categorical(probs=p), Categorical(probs=q))\n", + "assert got.shape == (4,), f\"expected shape (4,), got {tuple(got.shape)}\"\n", + "assert torch.allclose(got, ref, atol=1e-6), f\"mismatch vs torch.distributions: {got} vs {ref}\"\n", + "assert torch.allclose(kl_categorical(p, p), torch.zeros(4), atol=1e-6), \"KL(P||P) must be 0\"\n", + "assert (got >= 0).all(), \"KL is never negative (Gibbs' inequality)\"\n", + "assert not torch.allclose(got, kl_categorical(q, p)), \"KL is not symmetric; check the argument order\"\n", + "\n", + "# Zeros in the support\n", + "p0 = torch.tensor([0.5, 0.5, 0.0])\n", + "q0 = torch.tensor([0.25, 0.25, 0.5])\n", + "assert torch.isclose(kl_categorical(p0, q0), torch.tensor(math.log(2.0))), \\\n", + " f\"p = 0 entries must contribute 0; expected log 2, got {kl_categorical(p0, q0)}\"\n", + "assert torch.isinf(kl_categorical(q0, p0)), \"q = 0 where p > 0 must give +inf (not nan)\"\n", + "print(\"Part 1 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 2: KL from logits (the version you'll actually use)\n", + "\n", + "Models output **logits**. Implement `kl_from_logits(p_logits, q_logits)` = KL(softmax(p_logits) β€– softmax(q_logits)), with output shape `(...)`.\n", + "\n", + "Why not just `kl_categorical(p_logits.softmax(-1), q_logits.softmax(-1))`? With confident (large) logits, softmax **underflows to exactly 0**, `log 0 = -inf`, and you get `inf` or `nan` even though the true KL is finite.\n", + "\n", + "Instead, get `log P` and `log Q` directly with `F.log_softmax`, which uses `logsumexp` (the max-subtraction trick) and never takes the log of an underflowed probability:\n", + "\n", + "$$\\mathrm{KL} = \\sum_k e^{\\log P_k}\\,\\big(\\log P_k - \\log Q_k\\big)$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_from_logits(p_logits: torch.Tensor, q_logits: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " KL(P || Q) where P = softmax(p_logits), Q = softmax(q_logits) over the last dim.\n", + "\n", + " Args:\n", + " p_logits, q_logits: unnormalized logits of shape (..., K).\n", + "\n", + " Returns:\n", + " Tensor of shape (...).\n", + " \"\"\"\n", + " logp = F.log_softmax(p_logits, dim=-1)\n", + " logq = F.log_softmax(q_logits, dim=-1)\n", + " # log-probs are finite for finite logits, so no log(0) ever happens here\n", + " return (logp.exp() * (logp - logq)).sum(-1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 2 tests ===\n", + "pl, ql = torch.randn(3, 5, 10), torch.randn(3, 5, 10)\n", + "got = kl_from_logits(pl, ql)\n", + "ref = kl_divergence(Categorical(logits=pl), Categorical(logits=ql))\n", + "assert got.shape == (3, 5), f\"expected shape (3, 5), got {tuple(got.shape)}\"\n", + "assert torch.allclose(got, ref, atol=1e-5), \"mismatch vs torch.distributions\"\n", + "assert torch.allclose(kl_from_logits(pl + 7.0, ql - 3.0), got, atol=1e-5), \\\n", + " \"softmax is shift-invariant, so adding a constant to the logits must not change the KL\"\n", + "\n", + "# Confident logits: the naive probability route breaks, the log-space route doesn't\n", + "big_p, big_q = 100 * pl, 100 * ql\n", + "naive = kl_categorical(big_p.softmax(-1), big_q.softmax(-1))\n", + "stable = kl_from_logits(big_p, big_q)\n", + "print(f\"naive route: {int((~torch.isfinite(naive)).sum())} of {naive.numel()} entries are inf/nan\")\n", + "assert torch.isfinite(stable).all(), \"kl_from_logits must stay finite for large logits; use log_softmax\"\n", + "assert (stable >= -1e-4).all(), \"KL must be non-negative\"\n", + "print(\"Part 2 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 3: Re-implement `F.kl_div` (the API everyone gets wrong)\n", + "\n", + "```python\n", + "F.kl_div(input, target, reduction=\"mean\", log_target=False)\n", + "```\n", + "\n", + "- `input` is the **model's log-probabilities**, `log Q`. It is *not* probabilities.\n", + "- `target` is the **reference distribution P**: probabilities, or `log P` if `log_target=True`.\n", + "- The pointwise loss is `P Β· (log P βˆ’ input)`, with 0 where `P = 0`.\n", + "- So `F.kl_div(log_q, p)` computes **KL(P β€– Q)**: the second argument is the first argument of the KL.\n", + "\n", + "The reductions:\n", + "\n", + "| `reduction` | returns |\n", + "|---|---|\n", + "| `'none'` | the pointwise tensor, same shape as `input` |\n", + "| `'sum'` | sum of everything |\n", + "| `'mean'` | mean over **all elements**, i.e. divided by `BΒ·K`, **not** the KL |\n", + "| `'batchmean'` | `sum / input.size(0)`, the true mean KL per example |\n", + "\n", + "Implement `my_kl_div` so that it matches `F.kl_div` for all four reductions and both values of `log_target`. Then read the test output to see how far off `'mean'` is.\n", + "\n", + "> PyTorch warns that `'mean'` will be changed to behave like `'batchmean'` \"in the next major release\". Through torch 2.x it still divides by `BΒ·K`, so always pass `reduction='batchmean'` explicitly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def my_kl_div(input: torch.Tensor, target: torch.Tensor,\n", + " reduction: str = \"mean\", log_target: bool = False) -> torch.Tensor:\n", + " \"\"\"\n", + " Same semantics as torch.nn.functional.kl_div.\n", + "\n", + " Args:\n", + " input: log-probabilities of the model distribution Q, any shape.\n", + " target: P as probabilities (log_target=False) or log-probabilities (log_target=True).\n", + " reduction: 'none' | 'sum' | 'mean' | 'batchmean'.\n", + " log_target: whether target is given in log space.\n", + " \"\"\"\n", + " if log_target:\n", + " pointwise = target.exp() * (target - input)\n", + " else:\n", + " pointwise = torch.xlogy(target, target) - target * input\n", + "\n", + " if reduction == \"none\":\n", + " return pointwise\n", + " if reduction == \"sum\":\n", + " return pointwise.sum()\n", + " if reduction == \"mean\": # mean over every element; NOT the KL\n", + " return pointwise.mean()\n", + " if reduction == \"batchmean\": # what the math actually means\n", + " return pointwise.sum() / input.size(0)\n", + " raise ValueError(f\"unknown reduction '{reduction}'\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 3 tests ===\n", + "B, K = 8, 5\n", + "log_q = F.log_softmax(torch.randn(B, K), -1)\n", + "p = torch.softmax(torch.randn(B, K), -1)\n", + "p_pos = p.clone() # strictly positive: log P is finite\n", + "p[0, 0] = 0.0\n", + "p[0] = p[0] / p[0].sum() # a probability target with an exact zero\n", + "for log_target in (False, True):\n", + " tgt = p_pos.log() if log_target else p\n", + " for red in (\"none\", \"sum\", \"mean\", \"batchmean\"):\n", + " got = my_kl_div(log_q, tgt, reduction=red, log_target=log_target)\n", + " exp = F.kl_div(log_q, tgt, reduction=red, log_target=log_target)\n", + " assert torch.allclose(got, exp, atol=1e-6), \\\n", + " f\"reduction={red}, log_target={log_target}: {got} vs F.kl_div {exp}\"\n", + "\n", + "true_kl = kl_categorical(p, log_q.exp()).mean()\n", + "print(f\"true mean KL(P||Q) = {true_kl:.5f}\")\n", + "print(f\"F.kl_div 'batchmean' = {F.kl_div(log_q, p, reduction='batchmean'):.5f}\")\n", + "print(f\"F.kl_div 'mean' default = {F.kl_div(log_q, p):.5f} <- off by a factor of K = {K}\")\n", + "assert torch.allclose(my_kl_div(log_q, p, reduction=\"batchmean\"), true_kl, atol=1e-6)\n", + "print(\"Part 3 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 4: Jensen–Shannon divergence (symmetric and bounded)\n", + "\n", + "$$\\mathrm{JS}(P, Q) = \\tfrac12\\,\\mathrm{KL}(P\\,\\|\\,M) + \\tfrac12\\,\\mathrm{KL}(Q\\,\\|\\,M), \\qquad M = \\tfrac12(P + Q)$$\n", + "\n", + "Unlike KL, JS is **symmetric**, **bounded** in `[0, log 2]`, and **finite even when the supports don't overlap**, because M covers both. It is the divergence behind the original GAN objective, and its generalized (Ξ²-weighted) form is the loss in GKD distillation.\n", + "\n", + "Implement `js_divergence(p_logits, q_logits)`, with output shape `(...)`. Compute `log M` in log space: `log M = logsumexp([log P, log Q]) βˆ’ log 2`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def js_divergence(p_logits: torch.Tensor, q_logits: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " Jensen-Shannon divergence between softmax(p_logits) and softmax(q_logits).\n", + "\n", + " Args:\n", + " p_logits, q_logits: logits of shape (..., K).\n", + "\n", + " Returns:\n", + " Tensor of shape (...), in [0, log 2].\n", + " \"\"\"\n", + " logp = F.log_softmax(p_logits, dim=-1)\n", + " logq = F.log_softmax(q_logits, dim=-1)\n", + " # log((P + Q) / 2) without leaving log space\n", + " logm = torch.logsumexp(torch.stack([logp, logq]), dim=0) - math.log(2.0)\n", + " kl_pm = (logp.exp() * (logp - logm)).sum(-1)\n", + " kl_qm = (logq.exp() * (logq - logm)).sum(-1)\n", + " return 0.5 * kl_pm + 0.5 * kl_qm" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 4 tests ===\n", + "pl, ql = torch.randn(6, 7), torch.randn(6, 7)\n", + "got = js_divergence(pl, ql)\n", + "P, Q = pl.softmax(-1), ql.softmax(-1)\n", + "M = Categorical(probs=0.5 * (P + Q))\n", + "ref = 0.5 * kl_divergence(Categorical(probs=P), M) + 0.5 * kl_divergence(Categorical(probs=Q), M)\n", + "assert got.shape == (6,), f\"expected shape (6,), got {tuple(got.shape)}\"\n", + "assert torch.allclose(got, ref, atol=1e-5), \"mismatch vs 0.5*KL(P||M) + 0.5*KL(Q||M)\"\n", + "assert torch.allclose(got, js_divergence(ql, pl), atol=1e-6), \"JS must be symmetric\"\n", + "assert (got >= -1e-6).all() and (got <= math.log(2) + 1e-6).all(), \"JS must lie in [0, log 2]\"\n", + "\n", + "# Disjoint supports: KL explodes, JS saturates at log 2\n", + "far_p = torch.tensor([50.0, -50.0, -50.0, -50.0])\n", + "far_q = torch.tensor([-50.0, -50.0, -50.0, 50.0])\n", + "print(f\"KL = {kl_from_logits(far_p, far_q):.1f} JS = {js_divergence(far_p, far_q):.4f} log 2 = {math.log(2):.4f}\")\n", + "assert abs(js_divergence(far_p, far_q).item() - math.log(2)) < 1e-4, \"disjoint supports must give JS = log 2\"\n", + "print(\"Part 4 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 5: Forward vs reverse KL (mode-covering vs mode-seeking)\n", + "\n", + "Fit a **single** Gaussian `q_ΞΈ` (parameters `ΞΌ`, `Οƒ`) to a **bimodal** target `p` in two ways:\n", + "\n", + "| objective | penalizes | behaviour | where you meet it |\n", + "|---|---|---|---|\n", + "| **forward** KL(p β€– q_ΞΈ) = E_p[log p βˆ’ log q] | `q` small where `p` is large | **mode-covering**: stretches over both modes | MLE, cross-entropy, SFT, classic distillation |\n", + "| **reverse** KL(q_ΞΈ β€– p) = E_q[log q βˆ’ log p] | `q` large where `p` is small | **mode-seeking**: hides inside one mode | variational inference, KL-regularized RL, MiniLLM |\n", + "\n", + "To keep everything categorical, discretize the real line on a `grid`. `q_ΞΈ`'s logits on the grid are `-Β½((grid βˆ’ ΞΌ)/Οƒ)Β²`; `log_softmax` normalizes them for you. Then reuse `kl_from_logits`, remembering that it takes **logits**, and that `target_logp` is already a valid set of logits.\n", + "\n", + "Implement `fit_gaussian`: optimize `ΞΌ` and `log Οƒ` (so Οƒ stays positive) with Adam, starting from `mu_init` and `sigma_init`, and return `(ΞΌ, Οƒ)` as Python floats.\n", + "\n", + "Predict the two answers before you run the test." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def fit_gaussian(target_logp: torch.Tensor, grid: torch.Tensor, direction: str,\n", + " steps: int = 500, lr: float = 0.05,\n", + " mu_init: float = 1.0, sigma_init: float = 0.5) -> tuple:\n", + " \"\"\"\n", + " Fit a discretized Gaussian q to a target distribution by minimizing KL.\n", + "\n", + " Args:\n", + " target_logp: log-probabilities of the target on the grid, shape (G,).\n", + " grid: grid points, shape (G,).\n", + " direction: 'forward' -> minimize KL(target || q); 'reverse' -> minimize KL(q || target).\n", + " steps, lr: Adam iterations and learning rate.\n", + " mu_init, sigma_init: starting point.\n", + "\n", + " Returns:\n", + " (mu, sigma) as Python floats.\n", + " \"\"\"\n", + " if direction not in (\"forward\", \"reverse\"):\n", + " raise ValueError(f\"direction must be 'forward' or 'reverse', got '{direction}'\")\n", + " mu = torch.tensor(float(mu_init), requires_grad=True)\n", + " log_sigma = torch.tensor(math.log(sigma_init), requires_grad=True)\n", + " opt = torch.optim.Adam([mu, log_sigma], lr=lr)\n", + " for _ in range(steps):\n", + " q_logits = -0.5 * ((grid - mu) / log_sigma.exp()) ** 2\n", + " if direction == \"forward\":\n", + " loss = kl_from_logits(target_logp, q_logits)\n", + " else:\n", + " loss = kl_from_logits(q_logits, target_logp)\n", + " opt.zero_grad()\n", + " loss.backward()\n", + " opt.step()\n", + " return mu.item(), log_sigma.exp().item()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 5 tests ===\n", + "grid = torch.linspace(-6, 6, 241)\n", + "# Two well-separated modes at -2.5 and +2.5, each with std 0.6. Built in log space so the\n", + "# valley between the modes stays finite (in probability space it would underflow to 0).\n", + "mixture = torch.logsumexp(torch.stack([Normal(-2.5, 0.6).log_prob(grid),\n", + " Normal(2.5, 0.6).log_prob(grid)]), dim=0)\n", + "target_logp = F.log_softmax(mixture, dim=-1)\n", + "\n", + "mu_f, sigma_f = fit_gaussian(target_logp, grid, \"forward\")\n", + "mu_r, sigma_r = fit_gaussian(target_logp, grid, \"reverse\")\n", + "print(f\"forward KL(p||q): mu = {mu_f:+.2f}, sigma = {sigma_f:.2f} (covers both modes)\")\n", + "print(f\"reverse KL(q||p): mu = {mu_r:+.2f}, sigma = {sigma_r:.2f} (locks onto one mode)\")\n", + "\n", + "assert abs(mu_f) < 0.5 and sigma_f > 2.0, \"forward KL should centre between the modes with a wide sigma\"\n", + "assert abs(mu_r - 2.5) < 0.3 and sigma_r < 1.0, \"reverse KL should sit on the +2.5 mode with a narrow sigma\"\n", + "print(\"Part 5 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> **Try it:** rerun the reverse fit with `mu_init=0.5, sigma_init=1.0`. It ends up wide and centred, stuck between the modes. Reverse KL is non-convex in ΞΈ, so *where you start decides which mode, or which compromise, you get*. That's one reason RL fine-tuning with a KL penalty is sensitive to initialization (the SFT model)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 6: Closed-form KL between diagonal Gaussians (the VAE term)\n", + "\n", + "For continuous distributions the sum becomes an integral. For two 1-D Gaussians it has a closed form:\n", + "\n", + "$$\\mathrm{KL}\\big(\\mathcal N(\\mu_p,\\sigma_p^2)\\,\\|\\,\\mathcal N(\\mu_q,\\sigma_q^2)\\big) = \\log\\frac{\\sigma_q}{\\sigma_p} + \\frac{\\sigma_p^2 + (\\mu_p-\\mu_q)^2}{2\\sigma_q^2} - \\frac12$$\n", + "\n", + "For **diagonal** Gaussians in `D` dimensions, the KL is the sum over dimensions.\n", + "\n", + "Implement `kl_gaussian` taking **log-variances** (that's what an encoder outputs, because it keeps σ² > 0 for free). Reduce over the last dim.\n", + "\n", + "Then check the special case every VAE uses, with `q = N(0, I)`:\n", + "\n", + "$$\\mathrm{KL}\\big(\\mathcal N(\\mu,\\sigma^2)\\,\\|\\,\\mathcal N(0, I)\\big) = -\\tfrac12\\sum_d\\big(1 + \\log\\sigma_d^2 - \\mu_d^2 - \\sigma_d^2\\big)$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_gaussian(mu_p: torch.Tensor, logvar_p: torch.Tensor,\n", + " mu_q: torch.Tensor, logvar_q: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " KL( N(mu_p, diag(exp(logvar_p))) || N(mu_q, diag(exp(logvar_q))) ).\n", + "\n", + " Args:\n", + " mu_p, logvar_p, mu_q, logvar_q: tensors of shape (..., D) (broadcastable).\n", + "\n", + " Returns:\n", + " Tensor of shape (...): KL summed over the D dimensions.\n", + " \"\"\"\n", + " # log(sigma_q / sigma_p) = 0.5 * (logvar_q - logvar_p); dividing by var_q is a\n", + " # multiplication by exp(-logvar_q), so no explicit sigma is ever formed.\n", + " per_dim = 0.5 * (logvar_q - logvar_p\n", + " + (logvar_p.exp() + (mu_p - mu_q) ** 2) * torch.exp(-logvar_q)\n", + " - 1.0)\n", + " return per_dim.sum(-1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 6 tests ===\n", + "mu_p, lv_p = torch.randn(4, 3), torch.randn(4, 3)\n", + "mu_q, lv_q = torch.randn(4, 3), torch.randn(4, 3)\n", + "got = kl_gaussian(mu_p, lv_p, mu_q, lv_q)\n", + "ref = kl_divergence(Normal(mu_p, (0.5 * lv_p).exp()), Normal(mu_q, (0.5 * lv_q).exp())).sum(-1)\n", + "assert got.shape == (4,), f\"expected shape (4,), got {tuple(got.shape)}\"\n", + "assert torch.allclose(got, ref, atol=1e-5), \"mismatch vs torch.distributions.Normal\"\n", + "assert torch.allclose(kl_gaussian(mu_p, lv_p, mu_p, lv_p), torch.zeros(4), atol=1e-6), \"KL(p||p) must be 0\"\n", + "\n", + "# The VAE special case: q = N(0, I)\n", + "zeros = torch.zeros_like(mu_p)\n", + "vae = -0.5 * (1 + lv_p - mu_p ** 2 - lv_p.exp()).sum(-1)\n", + "assert torch.allclose(kl_gaussian(mu_p, lv_p, zeros, zeros), vae, atol=1e-5), \"VAE closed form mismatch\"\n", + "print(\"Part 6 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 7: Token-level KL for language models (with padding)\n", + "\n", + "In LLM distillation and RLHF you compare two models **at every position** of a batch of sequences:\n", + "\n", + "- logits `(B, T, V)` from each model,\n", + "- a `mask` of shape `(B, T)`: 1 for real tokens, 0 for padding (or for prompt tokens you don't want to score).\n", + "\n", + "Implement `masked_token_kl(p_logits, q_logits, mask)`. It should:\n", + "\n", + "1. compute the exact, full-vocabulary KL(P_t β€– Q_t) at every position,\n", + "2. **ignore** masked positions completely, even if their logits are garbage, and\n", + "3. return a scalar: the **mean over unmasked tokens** (not over sequences).\n", + "\n", + "Directions in practice:\n", + "\n", + "- **forward distillation**: `masked_token_kl(teacher_logits.detach(), student_logits, mask)`\n", + "- **reverse distillation / RL penalty**: `masked_token_kl(student_logits, teacher_logits.detach(), mask)`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def masked_token_kl(p_logits: torch.Tensor, q_logits: torch.Tensor,\n", + " mask: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"\n", + " Mean per-token KL(P_t || Q_t) over unmasked positions.\n", + "\n", + " Args:\n", + " p_logits, q_logits: logits of shape (B, T, V).\n", + " mask: (B, T), 1 (or True) for tokens to include, 0 for padding.\n", + "\n", + " Returns:\n", + " Scalar tensor.\n", + " \"\"\"\n", + " per_token = kl_from_logits(p_logits, q_logits) # (B, T)\n", + " mask = mask.to(per_token.dtype)\n", + " # Multiplying by 0 is safe because per_token is finite for finite logits; with -inf\n", + " # logits you'd need torch.where to keep a nan at a pad position out of the sum.\n", + " return (per_token * mask).sum() / mask.sum().clamp(min=1.0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 7 tests ===\n", + "B, T, V = 3, 6, 11\n", + "pl, ql = torch.randn(B, T, V), torch.randn(B, T, V)\n", + "pl[2] *= 4 # the short sequence disagrees much more\n", + "lengths = [6, 4, 2]\n", + "mask = torch.zeros(B, T)\n", + "for i, n in enumerate(lengths):\n", + " mask[i, :n] = 1\n", + "\n", + "got = masked_token_kl(pl, ql, mask)\n", + "oracle = torch.stack([kl_divergence(Categorical(logits=pl[i, t]), Categorical(logits=ql[i, t]))\n", + " for i in range(B) for t in range(lengths[i])]).mean()\n", + "assert got.ndim == 0, f\"expected a scalar, got shape {tuple(got.shape)}\"\n", + "assert torch.allclose(got, oracle, atol=1e-5), f\"expected the mean over the {sum(lengths)} real tokens\"\n", + "\n", + "# Garbage at padded positions must not leak in\n", + "pl_garbage, ql_garbage = pl.clone(), ql.clone()\n", + "pl_garbage[mask == 0] = 1e3 * torch.randn(int((mask == 0).sum()), V)\n", + "ql_garbage[mask == 0] = -1e3 * torch.randn(int((mask == 0).sum()), V)\n", + "assert torch.allclose(masked_token_kl(pl_garbage, ql_garbage, mask), got, atol=1e-5), \\\n", + " \"changing logits at padded positions changed the result\"\n", + "\n", + "# Gradients reach real tokens only\n", + "qg = ql.clone().requires_grad_(True)\n", + "masked_token_kl(pl, qg, mask).backward()\n", + "assert qg.grad[mask == 0].abs().max() == 0, \"padded positions must get zero gradient\"\n", + "assert qg.grad[mask == 1].abs().max() > 0, \"real positions must get a gradient\"\n", + "\n", + "# Token mean, not sequence mean\n", + "seq_mean = torch.stack([kl_from_logits(pl[i, :n], ql[i, :n]).mean() for i, n in enumerate(lengths)]).mean()\n", + "print(f\"token mean = {got:.4f} sequence mean = {seq_mean:.4f}\")\n", + "print(\"(the 2-token sequence is 2/12 of the token mean but 1/3 of the sequence mean)\")\n", + "print(\"Part 7 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 8: Estimating KL from samples (k1, k2, k3)\n", + "\n", + "In RL fine-tuning of LLMs, the exact KL over whole sequences is a sum over `V^T` possibilities, which is impossible. Often all you have is the **log-prob of the token that was actually sampled**, under the policy and under the reference model. So you estimate KL(Ο€ β€– Ο€_ref) by Monte Carlo, with `x ~ Ο€`.\n", + "\n", + "Let `logp = log Ο€(x)` and `logq = log Ο€_ref(x)`, and define the log-ratio `log r = logq βˆ’ logp`. Schulman's three estimators are:\n", + "\n", + "| estimator | formula | unbiased? | sign | variance |\n", + "|---|---|---|---|---|\n", + "| **k1** | `βˆ’log r` | βœ… | can be **negative** | high |\n", + "| **k2** | `Β½ (log r)Β²` | ❌ (slightly biased) | β‰₯ 0 | low |\n", + "| **k3** | `(r βˆ’ 1) βˆ’ log r` | βœ… | β‰₯ 0 | low |\n", + "\n", + "- **Why k3 is unbiased:** `E_{x~Ο€}[r] = Ξ£ Ο€ Β· (Ο€_ref/Ο€) = 1`, so `(r βˆ’ 1)` is a zero-mean **control variate** added to k1.\n", + "- **Why k3 is β‰₯ 0:** `log r ≀ r βˆ’ 1` for every r.\n", + "\n", + "PPO-RLHF uses k1 per token. GRPO uses k3.\n", + "\n", + "Implement `kl_estimators(logp, logq)`, returning the per-sample tuple `(k1, k2, k3)`, each with the same shape as the inputs. Then check the bias and variance experimentally." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_estimators(logp: torch.Tensor, logq: torch.Tensor) -> tuple:\n", + " \"\"\"\n", + " Per-sample Monte Carlo estimators of KL(p || q), for samples x ~ p.\n", + "\n", + " Args:\n", + " logp: log p(x) of the sampled x (the sampling distribution, e.g. the policy).\n", + " logq: log q(x) of the same x (e.g. the reference model). Same shape as logp.\n", + "\n", + " Returns:\n", + " (k1, k2, k3), each the same shape as logp.\n", + " \"\"\"\n", + " log_ratio = logq - logp # log r, r = q(x) / p(x)\n", + " k1 = -log_ratio\n", + " k2 = 0.5 * log_ratio ** 2\n", + " # torch.expm1(x) = exp(x) - 1, accurate when r is close to 1 (the usual case)\n", + " k3 = torch.expm1(log_ratio) - log_ratio\n", + " return k1, k2, k3" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 8 tests ===\n", + "torch.manual_seed(0)\n", + "N = 200_000\n", + "print(f\"{'setting':>8} {'true KL':>9} | {'k1 bias/KL':>10} {'k1 std/KL':>9} | \"\n", + " f\"{'k2 bias/KL':>10} {'k2 std/KL':>9} | {'k3 bias/KL':>10} {'k3 std/KL':>9}\")\n", + "for name, scale in ((\"close\", 0.1), (\"far\", 1.0)):\n", + " policy_logits = torch.randn(50)\n", + " ref_logits = policy_logits + scale * torch.randn(50)\n", + " true = kl_from_logits(policy_logits, ref_logits).item()\n", + " x = Categorical(logits=policy_logits).sample((N,)) # samples from the policy\n", + " logp = F.log_softmax(policy_logits, -1)[x]\n", + " logq = F.log_softmax(ref_logits, -1)[x]\n", + " k1, k2, k3 = kl_estimators(logp, logq)\n", + " row = [f\"{(k.mean().item() - true) / true:>10.3f} {k.std().item() / true:>9.2f}\" for k in (k1, k2, k3)]\n", + " print(f\"{name:>8} {true:>9.5f} | \" + \" | \".join(row))\n", + "\n", + " se = lambda k: k.std().item() / math.sqrt(N)\n", + " assert abs(k1.mean().item() - true) < 5 * se(k1), \"k1 should be unbiased\"\n", + " assert abs(k3.mean().item() - true) < 5 * se(k3), \"k3 should be unbiased\"\n", + " assert (k2 >= 0).all() and (k3 >= -1e-6).all(), \"k2 and k3 are non-negative per sample\"\n", + " assert (k1 < 0).any(), \"k1 should produce some negative per-sample estimates\"\n", + " assert k3.std() < k1.std(), \"k3 should have lower variance than k1\"\n", + " if name == \"far\":\n", + " assert k2.mean().item() - true > 0.1 * true, \"k2's bias should show up when p and q are far apart\"\n", + "print(\"Part 8 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 9: KL as a loss. Does the estimator have the right gradient?\n", + "\n", + "Part 8 was about **values**. In training you **differentiate**. A common pattern (GRPO is one example) is to add `Ξ² Β· mean(k(x))` straight to the loss, with the samples `x ~ Ο€_ΞΈ` treated as fixed data. Autograd then differentiates `k` *only through `log Ο€_ΞΈ(x)`*, never through the sampling. In expectation the gradient you get is\n", + "\n", + "$$g_k = \\mathbb{E}_{x\\sim\\pi_\\theta}\\big[\\nabla_\\theta\\, k(x;\\theta)\\big].$$\n", + "\n", + "Is `g_k` equal to `βˆ‡_ΞΈ KL(Ο€_ΞΈ β€– Ο€_ref)`?\n", + "\n", + "For a single categorical \"policy\" with logits ΞΈ, you can compute that expectation **exactly**, without sampling noise, by weighting every category by its **detached** probability:\n", + "\n", + "$$\\text{surrogate} = \\sum_x \\mathrm{sg}\\big(\\pi_\\theta(x)\\big)\\, k(x;\\theta), \\qquad g_k = \\nabla_\\theta\\,\\text{surrogate}$$\n", + "\n", + "where `sg` is stop-gradient (`.detach()`).\n", + "\n", + "Implement `kl_estimator_grad(logits, ref_logits, estimator)`, which returns `g_k` w.r.t. `logits` (same shape as `logits`) for `estimator` ∈ `{'k1', 'k2', 'k3'}`. The test compares each one against autograd of the exact **reverse** KL(Ο€β€–Ο€_ref) and the exact **forward** KL(Ο€_refβ€–Ο€).\n", + "\n", + "**Write down your prediction for each of the three before running the test.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def kl_estimator_grad(logits: torch.Tensor, ref_logits: torch.Tensor, estimator: str) -> torch.Tensor:\n", + " \"\"\"\n", + " Expected gradient (w.r.t. logits) of an on-policy KL estimator used as a loss.\n", + "\n", + " Args:\n", + " logits: policy logits, shape (..., V).\n", + " ref_logits: reference logits, shape (..., V); no gradient flows into them.\n", + " estimator: 'k1', 'k2' or 'k3' (see kl_estimators).\n", + "\n", + " Returns:\n", + " Tensor with the same shape as logits: grad of sum_x sg(pi(x)) * k(x) w.r.t. logits.\n", + " \"\"\"\n", + " index = {\"k1\": 0, \"k2\": 1, \"k3\": 2}[estimator]\n", + " logits = logits.detach().requires_grad_(True)\n", + " logp = F.log_softmax(logits, dim=-1)\n", + " logq = F.log_softmax(ref_logits, dim=-1).detach()\n", + " k = kl_estimators(logp, logq)[index]\n", + " # Exact expectation over x ~ pi with pi detached: what autograd sees when the samples\n", + " # are fixed data and only log pi(x) carries gradient.\n", + " surrogate = (logp.exp().detach() * k).sum()\n", + " (grad,) = torch.autograd.grad(surrogate, logits)\n", + " return grad" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 9 tests ===\n", + "torch.manual_seed(0)\n", + "policy_logits, ref_logits = torch.randn(4, 10), torch.randn(4, 10)\n", + "\n", + "theta = policy_logits.clone().requires_grad_(True)\n", + "grad_reverse = torch.autograd.grad(kl_from_logits(theta, ref_logits).sum(), theta)[0] # KL(pi || ref)\n", + "grad_forward = torch.autograd.grad(kl_from_logits(ref_logits, theta).sum(), theta)[0] # KL(ref || pi)\n", + "\n", + "grads = {e: kl_estimator_grad(policy_logits, ref_logits, e) for e in (\"k1\", \"k2\", \"k3\")}\n", + "print(f\"{'as loss':>8} | {'|g|':>7} | {'max |g - grad reverse KL|':>26} | {'max |g - grad forward KL|':>26}\")\n", + "for e, g in grads.items():\n", + " assert g.shape == policy_logits.shape, f\"{e}: expected shape {tuple(policy_logits.shape)}\"\n", + " print(f\"{e:>8} | {g.norm():7.4f} | {(g - grad_reverse).abs().max():26.2e} | {(g - grad_forward).abs().max():26.2e}\")\n", + "\n", + "assert torch.allclose(grads[\"k1\"], torch.zeros_like(policy_logits), atol=1e-6), \\\n", + " \"k1 as a loss: E[grad log pi] = 0, so its expected gradient should vanish\"\n", + "assert torch.allclose(grads[\"k2\"], grad_reverse, atol=1e-5), \\\n", + " \"k2 as a loss should reproduce the gradient of the reverse KL(pi || ref)\"\n", + "assert torch.allclose(grads[\"k3\"], grad_forward, atol=1e-5), \\\n", + " \"k3 as a loss: its expected gradient is that of the forward KL(ref || pi)\"\n", + "print(\"Part 9 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**What just happened.** With samples treated as data, the gradient of each estimator is `c(x)Β·βˆ‡log Ο€(x)`:\n", + "\n", + "| used as a loss | per-sample gradient | expected gradient |\n", + "|---|---|---|\n", + "| k1 | `βˆ‡log Ο€` | **0**, because `E_Ο€[βˆ‡log Ο€] = Ξ£ βˆ‡Ο€ = βˆ‡1 = 0` |\n", + "| k2 | `(log Ο€ βˆ’ log Ο€_ref)Β·βˆ‡log Ο€` | βˆ‡ **reverse** KL(Ο€β€–Ο€_ref), exactly |\n", + "| k3 | `(1 βˆ’ Ο€_ref/Ο€)Β·βˆ‡log Ο€` | `βˆ’E_{Ο€_ref}[βˆ‡log Ο€]` = βˆ‡ **forward** KL(Ο€_refβ€–Ο€) |\n", + "\n", + "So **an unbiased estimator of the value (k1, k3) is not automatically a correct loss for the gradient**. The PPO-RLHF recipe, which puts k1 **in the reward** (detached, multiplied by the score function), has the same gradient as k2-as-loss. Both are correct for reverse KL. GRPO's k3-as-loss matches it only to first order near Ο€ = Ο€_ref." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Interview follow-ups (with answers)\n", + "\n", + "1. **`F.kl_div(input, target)` computes KL(target β€– exp(input)).** `input` must be the *model's* log-probabilities (log Q), and `target` is the *reference* P, as probabilities by default or as log-probs with `log_target=True`. The order is the opposite of the math notation, which is the classic bug. `reduction='mean'` averages over **every element**, so it returns `batchmean / K`, which is not a KL value at all. Use `reduction='batchmean'`. PyTorch's docs say this explicitly, and `'mean'` stays the default only for consistency with the other losses.\n", + "\n", + "2. **Forward KL from the data to the model.** Cross-entropy is `H(P_data, Q_ΞΈ) = H(P_data) + KL(P_data β€– Q_ΞΈ)`, and `H(P_data)` doesn't depend on ΞΈ. So maximum likelihood, SFT and classification all minimize KL(data β€– model). That makes them **mode-covering**: the model is punished hard for assigning low probability to anything the data contains.\n", + "\n", + "3. **Usually reverse KL (or something between the two).** Forward KL(teacher β€– student) forces a small student to spread mass over everything the teacher might say, including the teacher's low-probability tail. At generation time that produces samples the teacher would never produce. Reverse KL(student β€– teacher) is mode-seeking: the student commits to the teacher's major modes. **MiniLLM** (Gu et al., 2023) minimizes reverse KL with policy-gradient-style optimization on student samples. **GKD** (Agarwal et al., 2024) trains on the student's own samples, either fully or as a mixed fraction, with a generalized Jensen–Shannon divergence whose Ξ² interpolates between the forward and reverse directions. Part 5 is the 1-D version of this picture.\n", + "\n", + "4. **Gradient scale.** With soft targets at temperature `T`, the gradient of the soft loss w.r.t. the logits scales like `1/TΒ²`. Multiplying by `TΒ²` keeps its magnitude comparable to the hard-label term as `T` changes. See the `knowledge-distillation` problem.\n", + "\n", + "5. **Not exactly.** Part 9 shows that, on-policy:\n", + " - **k1 in the reward** (detached, multiplied by the score function) gives the exact gradient of **reverse** KL(Ο€β€–Ο€_ref). PPO-RLHF does this, per token.\n", + " - **k2 as a loss** has the *same* gradient, `(log Ο€ βˆ’ log Ο€_ref)Β·βˆ‡log Ο€`.\n", + " - **k1 as a loss** has expected gradient **zero**, because `E_Ο€[βˆ‡log Ο€] = 0`. It does nothing.\n", + " - **k3 as a loss** (GRPO) has expected gradient `E_Ο€[(1 βˆ’ Ο€_ref/Ο€)βˆ‡log Ο€] = βˆ’E_{Ο€_ref}[βˆ‡log Ο€]`, which is the gradient of the **forward** KL(Ο€_refβ€–Ο€). Forward and reverse KL agree to second order near Ο€ = Ο€_ref, so k3-as-loss is a first-order approximation of the reverse-KL gradient that drifts as the policy moves away.\n", + "\n", + " Liu et al. (2025), *\"Rethinking KL Regularization in RLHF\"*, make the same equivalences precise.\n", + "\n", + "6. With `logits = -inf`, `log_softmax` gives `-inf`. If **both** distributions are `-inf` at the same index, `logp βˆ’ logq = -inf βˆ’ (-inf) = nan`, and `0 Β· nan = nan`. If only Q is `-inf` where P > 0, the KL really is `inf`. Fix: compute the pointwise term with `torch.where(p > 0, p * (logp - logq), 0)`, or mask the filtered vocabulary out before reducing, or use a large finite negative value (e.g. `torch.finfo(dtype).min`) instead of `-inf`.\n", + "\n", + "7. **Full-covariance Gaussians:** `KL = Β½[ tr(Ξ£_q⁻¹Σ_p) + (ΞΌ_qβˆ’ΞΌ_p)α΅€Ξ£_q⁻¹(ΞΌ_qβˆ’ΞΌ_p) βˆ’ k + log det Ξ£_q βˆ’ log det Ξ£_p ]`. Never form an inverse: take Cholesky factors `L_q`, use `torch.linalg.solve_triangular` for both the trace term and the Mahalanobis term, and get `log det Ξ£ = 2 Ξ£ log diag(L)`.\n", + "\n", + "8. A **token mean** weights every token equally. A **sequence mean** (average within each sequence, then across sequences) gives each token in a short sequence more weight than a token in a long one, which introduces a length bias. Most LLM losses use the token mean. Recent RL recipes such as DAPO and Dr. GRPO changed their loss aggregation for exactly this reason. Choose a sequence mean only when every *example* should count equally regardless of length.\n", + "\n", + "9. **No.** k3 relies on `E_{xβˆΌΟ€}[Ο€_ref(x)/Ο€(x)] = 1`, which only holds when `x` really is sampled from Ο€. On samples from Ο€_old, you need an importance weight `Ο€/Ο€_old` (usually clipped), or you are estimating a different divergence. The masked mean of k3 is still non-negative, but it no longer estimates either KL direction." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/v3/alignment-training/kl-divergence/problem.toml b/v3/alignment-training/kl-divergence/problem.toml new file mode 100644 index 0000000..c92153b --- /dev/null +++ b/v3/alignment-training/kl-divergence/problem.toml @@ -0,0 +1,30 @@ +id = "kl-divergence" +legacy_ids = ["v2-1"] +title = "Implement KL Divergence Loss" +description = "Implement KL divergence every way it shows up in ML: from probabilities and logits, F.kl_div semantics, Jensen-Shannon, forward vs reverse, Gaussian closed form, masked token-level KL, and the k1/k2/k3 estimators used in PPO and GRPO." +difficulty = "medium" +number = 1 +set = "v2" +category = "alignment-training" +tracks = ["advanced", "llm-path"] + +# GENERATED from the grader spec β€” edit python/torchleet/problems/*.py, then run: generate.py hints +hints = [ + "Everything is sum(p * (log p - log q)) over the last dim. Use torch.xlogy when you have probabilities and F.log_softmax when you have logits, so 0 * log 0 and softmax underflow never turn into nan.", + "F.kl_div(input, target) is KL(target || exp(input)): input is the model's log-probs, the target comes first in the math, and 'batchmean' (sum / batch size) is the true KL while 'mean' divides by every element.", + "For samples x ~ p with log r = log q(x) - log p(x): k1 = -log r, k2 = (log r)^2 / 2, k3 = (r - 1) - log r. For kl_estimator_grad, weight every category by pi.detach() and let the gradient flow only through log pi.", +] + +[notebooks] +question = "kl-divergence.ipynb" +solution = "kl-divergence_SOLN.ipynb" + +[grading] +tier = "A" # A=auto-graded B=partial C=not auto-graded (SEEDED, confirm) +entries = ["kl_categorical", "kl_from_logits", "my_kl_div", "js_divergence", "fit_gaussian", "kl_gaussian", "masked_token_kl", "kl_estimators", "kl_estimator_grad"] +extras = [] +device = "cpu" + +[llm_path] +order = 11 +stage = "alignment" diff --git a/v3/alignment-training/sft-smollm/problem.toml b/v3/alignment-training/sft-smollm/problem.toml new file mode 100644 index 0000000..6606725 --- /dev/null +++ b/v3/alignment-training/sft-smollm/problem.toml @@ -0,0 +1,30 @@ +id = "sft-smollm" +legacy_ids = ["v2-20"] +title = "Apply SFT on SmolLM" +description = "Fine-tune the pretrained SmolLM-135M base model into a chat assistant: build prompt-masked examples, pad and collate, write the shifted token-level loss, and normalize gradient accumulation correctly." +difficulty = "hard" +number = 20 +set = "v2" +category = "alignment-training" +tracks = ["advanced", "llm-path"] + +# GENERATED from the grader spec β€” edit python/torchleet/problems/*.py, then run: generate.py hints +hints = [ + "labels are a copy of input_ids (prompt + response + end-of-turn) with every prompt position set to -100; keep them aligned with input_ids and shift inside the loss.", + "Pad input_ids with pad_id but labels with -100, so a pad id that equals a real token id never erases a real label.", + "sft_loss: cross_entropy(logits[:, :-1], labels[:, 1:], ignore_index=-100, reduction='sum') divided by the trained-token count of the whole optimizer step, counted across all micro-batches before any backward.", +] + +[notebooks] +question = "sft-smollm.ipynb" +solution = "sft-smollm_SOLN.ipynb" + +[grading] +tier = "A" # A=auto-graded B=partial C=not auto-graded (SEEDED, confirm) +entries = ["build_sft_example", "sft_collate", "sft_loss", "sft_train_step"] +extras = [] +device = "cpu" + +[llm_path] +order = 13 +stage = "alignment" diff --git a/v3/alignment-training/sft-smollm/sft-smollm.ipynb b/v3/alignment-training/sft-smollm/sft-smollm.ipynb new file mode 100644 index 0000000..221510c --- /dev/null +++ b/v3/alignment-training/sft-smollm/sft-smollm.ipynb @@ -0,0 +1,623 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/sft-smollm/sft-smollm.ipynb)\n", + "\n", + "Check your work: `!pip install torchleet` then `from torchleet import check; check(\"sft-smollm\", ...)`\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Apply SFT on SmolLM\n", + "\n", + "**Difficulty**: πŸ”΄ Hard\n", + "\n", + "---\n", + "\n", + "### Problem Statement\n", + "\n", + "A pretrained **base** model is a document completer. Ask it a question and it continues the document, usually with more questions, and it never stops. **Supervised fine-tuning (SFT)** turns it into an assistant that answers in a chat format and ends its turn.\n", + "\n", + "Mechanically, SFT is next-token cross-entropy. Interviewers still ask about it, because almost every bug is in the **data path**, not the model:\n", + "\n", + "- **Loss on the prompt.** Only the response should be trained, so the prompt positions get label `-100`.\n", + "- **Off-by-one.** The logits at position `t` predict the token at `t + 1`.\n", + "- **Padding.** Pads must be excluded from the loss, and real tokens must never be mistaken for pads.\n", + "- **End of turn.** If the end-of-turn token isn't in the labels, the model never learns to stop.\n", + "- **Normalization.** Averaging per-micro-batch means under gradient accumulation weights tokens unequally. Hugging Face fixed exactly this bug in `transformers` in late 2024.\n", + "\n", + "You'll fine-tune **`HuggingFaceTB/SmolLM-135M`**, the base model whose architecture you built in the *SmolLM from Scratch* problem, on a tiny instruction dataset. Then you'll show that it follows the chat format on questions it never saw.\n", + "\n", + "---\n", + "\n", + "### Requirements\n", + "\n", + "1. **`build_sft_example`**: prompt ids + response ids + end-of-turn id β†’ `input_ids` and prompt-masked `labels`.\n", + "2. **`sft_collate`**: right-pad a list of examples into `input_ids`, `attention_mask` and `labels`.\n", + "3. **`sft_loss`**: the shifted, masked, token-level cross-entropy, with an optional global token count for gradient accumulation.\n", + "4. **`sft_train_step`**: one optimizer step over several micro-batches, normalized correctly.\n", + "5. **Fine-tune SmolLM-135M**: write `encode_chat_example`, train, and evaluate on held-out prompts.\n", + "\n", + "---\n", + "\n", + "### Constraints\n", + "\n", + "- βœ… `labels` are aligned with `input_ids` (not pre-shifted), the Hugging Face convention. The shift happens inside the loss.\n", + "- βœ… Use `IGNORE_INDEX = -100` for positions that must not be trained.\n", + "- βœ… The loss is the mean over **all response tokens in the optimizer step**, however they're split into micro-batches.\n", + "- ❌ No `trl`, `SFTTrainer`, `Trainer` or HF data collators. Don't pass `labels=` to the HF model either, since it would compute the loss for you.\n", + "\n", + "**Runtime:** about a minute on a laptop CPU after a ~270 MB model download; faster on a GPU or Apple MPS.\n", + "\n", + "---\n", + "\n", + "
\n", + " πŸ’‘ Hint\n", + "\n", + " - Build `labels` as a copy of `input_ids` and overwrite the prompt part with `-100`.\n", + " - Padding `labels` with `-100` *in the collator* is what keeps a pad id that equals a real token id from erasing real labels.\n", + " - `F.cross_entropy(..., ignore_index=-100, reduction=\"sum\")` divided by the number of non-ignored targets is a token mean. With gradient accumulation, divide by the count over **all** micro-batches instead.\n", + "\n", + "
\n", + "\n", + "---" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import copy\n", + "import random\n", + "from types import SimpleNamespace\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "\n", + "IGNORE_INDEX = -100 # F.cross_entropy's default ignore_index; Hugging Face uses the same value\n", + "\n", + "torch.manual_seed(0)\n", + "random.seed(0)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 1: Build one SFT example\n", + "\n", + "An SFT example is the chat-formatted **prompt**, then the **response**, then the **end-of-turn** token. `labels` are a copy of `input_ids` with every prompt position set to `-100`:\n", + "\n", + "```\n", + "input_ids : <|im_start|> user \\n What is ... <|im_end|> \\n <|im_start|> assistant \\n | The capital ... . | <|im_end|>\n", + "labels : -100 -100 ... ... -100 | The capital ... . | <|im_end|>\n", + "```\n", + "\n", + "Labels stay **aligned** with `input_ids`, and the shift happens in the loss. If `max_len` is given, truncate both from the right.\n", + "\n", + "Accept Python lists or 1-D tensors, and return 1-D `torch.long` tensors." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def build_sft_example(prompt_ids, response_ids, eos_id: int, max_len: int = None) -> dict:\n", + " \"\"\"\n", + " Args:\n", + " prompt_ids: token ids of the chat-formatted prompt (list or 1-D tensor).\n", + " response_ids: token ids of the assistant's response (list or 1-D tensor).\n", + " eos_id: end-of-turn token id, appended after the response and trained on.\n", + " max_len: optional; truncate input_ids and labels to this many tokens.\n", + "\n", + " Returns:\n", + " {\"input_ids\": LongTensor (L,), \"labels\": LongTensor (L,)} with prompt labels = IGNORE_INDEX.\n", + " \"\"\"\n", + " # TODO: input_ids = prompt + response + [eos]\n", + " # TODO: labels = [IGNORE_INDEX] * len(prompt) + response + [eos]\n", + " # TODO: truncate both to max_len if given\n", + " # TODO: return both as torch.long tensors\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 1 tests ===\n", + "ex = build_sft_example([5, 6, 7], [8, 9], eos_id=2)\n", + "assert ex[\"input_ids\"].tolist() == [5, 6, 7, 8, 9, 2], f\"input_ids wrong: {ex['input_ids'].tolist()}\"\n", + "assert ex[\"labels\"].tolist() == [-100, -100, -100, 8, 9, 2], f\"labels wrong: {ex['labels'].tolist()}\"\n", + "assert ex[\"input_ids\"].dtype == torch.long and ex[\"labels\"].dtype == torch.long, \"use torch.long\"\n", + "\n", + "cut = build_sft_example([5, 6, 7], [8, 9], eos_id=2, max_len=4)\n", + "assert cut[\"input_ids\"].tolist() == [5, 6, 7, 8] and cut[\"labels\"].tolist() == [-100, -100, -100, 8], \\\n", + " \"truncate both tensors from the right\"\n", + "\n", + "from_tensors = build_sft_example(torch.tensor([5, 6, 7]), torch.tensor([8, 9]), eos_id=2)\n", + "assert from_tensors[\"labels\"].tolist() == ex[\"labels\"].tolist(), \"tensor inputs must work too\"\n", + "print(\"Part 1 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 2: Collate a batch (right padding)\n", + "\n", + "Examples have different lengths. Right-pad them to the longest one in the batch:\n", + "\n", + "- `input_ids`: pad with `pad_id`\n", + "- `attention_mask`: 1 for real tokens, 0 for pads\n", + "- `labels`: pad with `IGNORE_INDEX`, **not** `pad_id`\n", + "\n", + "The last point matters when `pad_id` is also a real token id. SmolLM's tokenizer has no pad token, so a common choice is to reuse `<|endoftext|>`. Many tokenizers reuse EOS. Because labels come from the examples and only the padding is `-100`, a genuine EOS in the labels survives." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def sft_collate(examples: list, pad_id: int) -> dict:\n", + " \"\"\"\n", + " Args:\n", + " examples: list of {\"input_ids\": (L_i,), \"labels\": (L_i,)} from build_sft_example.\n", + " pad_id: token id used to pad input_ids.\n", + "\n", + " Returns:\n", + " {\"input_ids\": (B, L), \"attention_mask\": (B, L), \"labels\": (B, L)}, all torch.long,\n", + " with L = max L_i and padding on the right.\n", + " \"\"\"\n", + " # TODO: L = longest example\n", + " # TODO: input_ids filled with pad_id, labels filled with IGNORE_INDEX, attention_mask with 0\n", + " # TODO: copy each example into the left part of its row and set its mask to 1\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 2 tests ===\n", + "EOS = PAD = 2 # deliberately the same id\n", + "exs = [build_sft_example([5, 6], [7], EOS), # length 4\n", + " build_sft_example([5, 6, 7], [8, 9, 10], EOS), # length 7\n", + " build_sft_example([5], [], EOS)] # length 2\n", + "batch = sft_collate(exs, pad_id=PAD)\n", + "assert batch[\"input_ids\"].shape == (3, 7), f\"expected (3, 7), got {tuple(batch['input_ids'].shape)}\"\n", + "assert batch[\"attention_mask\"].sum(1).tolist() == [4, 7, 2], \"attention_mask must count the real tokens\"\n", + "assert batch[\"input_ids\"][0].tolist() == [5, 6, 7, 2, 2, 2, 2], \"right-pad input_ids with pad_id\"\n", + "assert batch[\"labels\"][0].tolist() == [-100, -100, 7, 2, -100, -100, -100], \\\n", + " \"pad labels with -100, and keep the real end-of-turn label even though it equals pad_id\"\n", + "assert batch[\"labels\"][2].tolist()[:2] == [-100, 2], \"an empty response still trains the end-of-turn token\"\n", + "assert all(v.dtype == torch.long for v in batch.values()), \"all tensors should be torch.long\"\n", + "print(\"Part 2 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 3: The SFT loss\n", + "\n", + "The logits at position `t` predict token `t + 1`, so drop the last logit and the first label:\n", + "\n", + "```\n", + "logits[:, :-1] vs labels[:, 1:]\n", + "```\n", + "\n", + "Cross-entropy over the non-ignored targets gives the **token mean**: every trained token counts equally, whichever sequence it's in.\n", + "\n", + "**Gradient accumulation.** Suppose one optimizer step is split into micro-batches. If each one returns *its own* mean and you average those means, a micro-batch with 10 response tokens weighs as much as one with 1,000. Instead, compute the number of trained tokens across the **whole step** (`num_items_in_batch`, Hugging Face's name for it) and return `sum_of_token_losses / num_items_in_batch`. The micro-batch losses then **add up** to the full-batch token mean.\n", + "\n", + "Also: upcast the logits to float32 before the softmax, because models often run in bf16. Return `0` (not `nan`) if nothing is trained." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def sft_loss(logits: torch.Tensor, labels: torch.Tensor, num_items_in_batch: int = None) -> torch.Tensor:\n", + " \"\"\"\n", + " Shifted next-token cross-entropy over non-ignored labels.\n", + "\n", + " Args:\n", + " logits: (B, T, V) model outputs.\n", + " labels: (B, T), aligned with the inputs; IGNORE_INDEX where nothing is trained.\n", + " num_items_in_batch: if given, divide the summed loss by this (the trained-token count of the\n", + " whole optimizer step) instead of by this batch's own count.\n", + "\n", + " Returns:\n", + " Scalar loss.\n", + " \"\"\"\n", + " # TODO: shift: logits[:, :-1] predict labels[:, 1:]\n", + " # TODO: summed cross-entropy over all positions, ignoring IGNORE_INDEX (upcast logits to float32)\n", + " # TODO: divide by num_items_in_batch, or by this batch's number of trained targets (at least 1)\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 3 tests ===\n", + "torch.manual_seed(0)\n", + "B, T, V = 3, 6, 13\n", + "logits = torch.randn(B, T, V)\n", + "labels = torch.randint(0, V, (B, T))\n", + "labels[0, :2] = IGNORE_INDEX # prompt\n", + "labels[1, :4] = IGNORE_INDEX # longer prompt\n", + "labels[2, 3:] = IGNORE_INDEX # padding\n", + "\n", + "# Loop oracle: position t predicts labels[t + 1]\n", + "total, count = 0.0, 0\n", + "for b in range(B):\n", + " for t in range(T - 1):\n", + " y = labels[b, t + 1].item()\n", + " if y != IGNORE_INDEX:\n", + " total += -F.log_softmax(logits[b, t], -1)[y]\n", + " count += 1\n", + "oracle = total / count\n", + "\n", + "got = sft_loss(logits, labels)\n", + "assert got.ndim == 0, \"the loss must be a scalar\"\n", + "assert torch.allclose(got, oracle, atol=1e-5), f\"expected {oracle:.5f}, got {got:.5f}; check the shift\"\n", + "\n", + "# Logits whose next label is ignored (and the very last logit) must not matter\n", + "noisy = logits.clone()\n", + "unused = torch.zeros(B, T, dtype=torch.bool)\n", + "unused[:, :-1] = labels[:, 1:] == IGNORE_INDEX\n", + "unused[:, -1] = True\n", + "noisy[unused] = 100 * torch.randn(int(unused.sum()), V)\n", + "assert torch.allclose(sft_loss(noisy, labels), got, atol=1e-5), \"ignored positions leaked into the loss\"\n", + "\n", + "# Gradient accumulation: micro-batch losses with a global count must add up to the full loss\n", + "n_total = int((labels[:, 1:] != IGNORE_INDEX).sum())\n", + "# (rows 0-1 hold 6 trained tokens, row 2 holds 2, so the two micro-batches are unequal)\n", + "accumulated = sft_loss(logits[:2], labels[:2], n_total) + sft_loss(logits[2:], labels[2:], n_total)\n", + "assert torch.allclose(accumulated, got, atol=1e-5), \"with num_items_in_batch, micro-batch losses must sum to the full loss\"\n", + "naive = (sft_loss(logits[:2], labels[:2]) + sft_loss(logits[2:], labels[2:])) / 2\n", + "print(f\"full-batch token mean = {got:.4f} accumulated (correct) = {accumulated:.4f} mean of micro-batch means = {naive:.4f}\")\n", + "\n", + "all_ignored = torch.full((B, T), IGNORE_INDEX)\n", + "assert torch.isfinite(sft_loss(logits, all_ignored)), \"a batch with nothing to train must not give nan\"\n", + "print(\"Part 3 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 4: One optimizer step with gradient accumulation\n", + "\n", + "`micro_batches` is a list of collated batches that together form **one** optimizer step. Implement `sft_train_step`:\n", + "\n", + "1. Zero the gradients.\n", + "2. Count the trained tokens across **all** micro-batches first. Remember the shift: it's the labels from position 1 onward that count.\n", + "3. For each micro-batch, run the forward pass, compute `sft_loss(..., num_items_in_batch=total)`, and call `backward()`. Gradients add up across calls.\n", + "4. Optionally clip the gradient norm, then `optimizer.step()`.\n", + "5. Return the step's loss (the sum of the micro-batch losses) as a Python float.\n", + "\n", + "The model follows the Hugging Face contract: `model(input_ids=..., attention_mask=...)` returns an object with a `.logits` attribute. Move each batch to the model's device.\n", + "\n", + "The test checks the property that matters: **splitting a batch into micro-batches must give exactly the same update as not splitting it.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def sft_train_step(model: nn.Module, micro_batches: list, optimizer: torch.optim.Optimizer,\n", + " max_grad_norm: float = None) -> float:\n", + " \"\"\"\n", + " One optimizer step over several micro-batches.\n", + "\n", + " Args:\n", + " model: returns an object with .logits of shape (B, T, V) from model(input_ids=..., attention_mask=...).\n", + " micro_batches: list of dicts from sft_collate.\n", + " optimizer: optimizer over the model's parameters.\n", + " max_grad_norm: clip the total gradient norm to this, if given.\n", + "\n", + " Returns:\n", + " The token-mean loss of the whole step, as a float.\n", + " \"\"\"\n", + " # TODO: model.train(); find the model's device\n", + " # TODO: zero the gradients\n", + " # TODO: total trained tokens over ALL micro-batches (labels[:, 1:] != IGNORE_INDEX)\n", + " # TODO: for each micro-batch: forward, sft_loss(..., num_items_in_batch=total), backward, accumulate the float loss\n", + " # TODO: clip (if max_grad_norm), optimizer.step(), return the loss\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 4 tests ===\n", + "class TinyLM(nn.Module):\n", + " \"\"\"Embedding -> tanh -> Linear. Position-wise, so padding can't leak between tokens.\"\"\"\n", + " def __init__(self, vocab=20, dim=16):\n", + " super().__init__()\n", + " self.emb = nn.Embedding(vocab, dim)\n", + " self.head = nn.Linear(dim, vocab)\n", + "\n", + " def forward(self, input_ids, attention_mask=None):\n", + " return SimpleNamespace(logits=self.head(torch.tanh(self.emb(input_ids))))\n", + "\n", + "\n", + "torch.manual_seed(0)\n", + "EOS, PAD = 19, 0\n", + "# Very different response lengths, so a wrong normalization shows up\n", + "exs = [build_sft_example([1, 2, 3], [4, 5, 6, 7, 8, 9, 10, 11], EOS),\n", + " build_sft_example([1, 2], [12], EOS),\n", + " build_sft_example([3, 4, 5, 6], [7, 8], EOS),\n", + " build_sft_example([2], [13, 14, 15], EOS)]\n", + "\n", + "model_split = TinyLM()\n", + "model_whole = copy.deepcopy(model_split)\n", + "reference = copy.deepcopy(model_split)\n", + "full = sft_collate(exs, PAD)\n", + "with torch.no_grad():\n", + " expected_loss = sft_loss(reference(input_ids=full[\"input_ids\"]).logits, full[\"labels\"]).item()\n", + "\n", + "loss_split = sft_train_step(model_split, [sft_collate(exs[:2], PAD), sft_collate(exs[2:], PAD)],\n", + " torch.optim.SGD(model_split.parameters(), lr=0.5))\n", + "loss_whole = sft_train_step(model_whole, [full], torch.optim.SGD(model_whole.parameters(), lr=0.5))\n", + "\n", + "assert abs(loss_whole - expected_loss) < 1e-5, f\"step loss {loss_whole:.5f} != full-batch loss {expected_loss:.5f}\"\n", + "assert abs(loss_split - loss_whole) < 1e-5, \\\n", + " f\"split into micro-batches the step loss is {loss_split:.5f}, unsplit {loss_whole:.5f}\"\n", + "for (name, a), b in zip(model_split.named_parameters(), model_whole.parameters()):\n", + " assert torch.allclose(a, b, atol=1e-6), \\\n", + " f\"parameter '{name}' differs: micro-batching must give the same update as one big batch\"\n", + "\n", + "# And it trains\n", + "model = TinyLM()\n", + "opt = torch.optim.Adam(model.parameters(), lr=1e-2)\n", + "first = sft_train_step(model, [full], opt)\n", + "for _ in range(50):\n", + " last = sft_train_step(model, [sft_collate(exs[:2], PAD), sft_collate(exs[2:], PAD)], opt, max_grad_norm=1.0)\n", + "assert last < 0.5 * first, f\"loss should fall on a fixed batch ({first:.3f} -> {last:.3f})\"\n", + "print(f\"loss {first:.3f} -> {last:.3f}\")\n", + "print(\"Part 4 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 5: Fine-tune SmolLM-135M into a (tiny) assistant\n", + "\n", + "**Chat format.** Use the ChatML format that SmolLM's instruct models use:\n", + "\n", + "```\n", + "<|im_start|>user\n", + "What is the capital of France?<|im_end|>\n", + "<|im_start|>assistant\n", + "The capital of France is Paris.<|im_end|>\n", + "```\n", + "\n", + "The base model's vocabulary already contains `<|im_start|>` and `<|im_end|>`, but the model was never trained to use them. Here `<|im_end|>` is the end-of-turn token, and `<|endoftext|>` is reused as the pad.\n", + "\n", + "**Data.** 32 *\"What is the capital of X?\"* pairs for training, and 8 **held-out** countries for evaluation.\n", + "\n", + "**Your task.** Write `encode_chat_example`: tokenize the formatted prompt and the response **separately** (so the label boundary is exact), with `add_special_tokens=False`, and pass them to `build_sft_example`. Then run the cells below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "\n", + "MODEL_ID = \"HuggingFaceTB/SmolLM-135M\" # the *base* model, not the instruct one\n", + "device = \"cuda\" if torch.cuda.is_available() else (\"mps\" if torch.backends.mps.is_available() else \"cpu\")\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\n", + "model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32).to(device)\n", + "IM_END = tokenizer.convert_tokens_to_ids(\"<|im_end|>\")\n", + "PAD_ID = tokenizer.convert_tokens_to_ids(\"<|endoftext|>\")\n", + "print(f\"device={device} params={sum(p.numel() for p in model.parameters()) / 1e6:.0f}M \"\n", + " f\"<|im_end|>={IM_END} pad=<|endoftext|>={PAD_ID}\")\n", + "\n", + "\n", + "def chat_prompt(question: str) -> str:\n", + " return f\"<|im_start|>user\\n{question}<|im_end|>\\n<|im_start|>assistant\\n\"\n", + "\n", + "\n", + "CAPITALS = [\n", + " (\"France\", \"Paris\"), (\"Germany\", \"Berlin\"), (\"Italy\", \"Rome\"), (\"Spain\", \"Madrid\"),\n", + " (\"Japan\", \"Tokyo\"), (\"China\", \"Beijing\"), (\"India\", \"New Delhi\"), (\"Russia\", \"Moscow\"),\n", + " (\"Canada\", \"Ottawa\"), (\"Australia\", \"Canberra\"), (\"Egypt\", \"Cairo\"), (\"Kenya\", \"Nairobi\"),\n", + " (\"Brazil\", \"Brasilia\"), (\"Argentina\", \"Buenos Aires\"), (\"Mexico\", \"Mexico City\"), (\"Peru\", \"Lima\"),\n", + " (\"Chile\", \"Santiago\"), (\"Greece\", \"Athens\"), (\"Portugal\", \"Lisbon\"), (\"Poland\", \"Warsaw\"),\n", + " (\"Norway\", \"Oslo\"), (\"Sweden\", \"Stockholm\"), (\"Finland\", \"Helsinki\"), (\"Denmark\", \"Copenhagen\"),\n", + " (\"Ireland\", \"Dublin\"), (\"Austria\", \"Vienna\"), (\"Hungary\", \"Budapest\"), (\"Turkey\", \"Ankara\"),\n", + " (\"Iran\", \"Tehran\"), (\"Thailand\", \"Bangkok\"), (\"Vietnam\", \"Hanoi\"), (\"South Korea\", \"Seoul\"),\n", + " # held out: never trained on\n", + " (\"Netherlands\", \"Amsterdam\"), (\"Belgium\", \"Brussels\"), (\"Switzerland\", \"Bern\"), (\"Cuba\", \"Havana\"),\n", + " (\"Nigeria\", \"Abuja\"), (\"Morocco\", \"Rabat\"), (\"Indonesia\", \"Jakarta\"), (\"Ukraine\", \"Kyiv\"),\n", + "]\n", + "train_pairs, heldout_pairs = CAPITALS[:32], CAPITALS[32:]\n", + "\n", + "\n", + "def qa(country: str, capital: str) -> tuple:\n", + " return f\"What is the capital of {country}?\", f\"The capital of {country} is {capital}.\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def encode_chat_example(question: str, answer: str, max_len: int = 64) -> dict:\n", + " \"\"\"Chat-format one (question, answer) pair and build its SFT example.\"\"\"\n", + " # TODO: prompt_ids = tokenizer(chat_prompt(question), add_special_tokens=False).input_ids\n", + " # TODO: response_ids = tokenizer(answer, add_special_tokens=False).input_ids\n", + " # TODO: return build_sft_example(prompt_ids, response_ids, IM_END, max_len)\n", + " ..." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Look at what is actually trained ===\n", + "ex = encode_chat_example(*qa(\"France\", \"Paris\"))\n", + "trained = ex[\"labels\"][ex[\"labels\"] != IGNORE_INDEX]\n", + "print(\"full input :\", repr(tokenizer.decode(ex[\"input_ids\"])))\n", + "print(\"trained on :\", repr(tokenizer.decode(trained)))\n", + "assert tokenizer.decode(trained) == \"The capital of France is Paris.<|im_end|>\", \\\n", + " \"only the response and the end-of-turn token should be trained\"\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def greedy_answer(model, question: str, max_new_tokens: int = 30) -> tuple:\n", + " \"\"\"Greedy decoding until <|im_end|>. Returns (text, stopped).\"\"\"\n", + " model.eval()\n", + " ids = torch.tensor([tokenizer(chat_prompt(question), add_special_tokens=False).input_ids], device=device)\n", + " out = []\n", + " for _ in range(max_new_tokens):\n", + " next_id = model(input_ids=ids).logits[0, -1].argmax().item()\n", + " if next_id == IM_END:\n", + " return tokenizer.decode(out), True\n", + " out.append(next_id)\n", + " ids = torch.cat([ids, torch.tensor([[next_id]], device=device)], dim=1)\n", + " return tokenizer.decode(out), False\n", + "\n", + "\n", + "def evaluate(model, pairs) -> list:\n", + " return [(country, capital, *greedy_answer(model, qa(country, capital)[0])) for country, capital in pairs]\n", + "\n", + "\n", + "before = evaluate(model, heldout_pairs)\n", + "print(\"\\n=== Base model, before SFT ===\")\n", + "for country, _, text, stopped in before[:3]:\n", + " print(f\"{country:>12}: {text[:70]!r} stopped={stopped}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Fine-tune ===\n", + "torch.manual_seed(0)\n", + "random.seed(0)\n", + "train_examples = [encode_chat_example(*qa(c, k)) for c, k in train_pairs]\n", + "\n", + "EPOCHS, BATCH, MICRO = 3, 8, 4 # 2 micro-batches of 4 per optimizer step\n", + "optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.0)\n", + "losses = []\n", + "for epoch in range(EPOCHS):\n", + " random.shuffle(train_examples)\n", + " for i in range(0, len(train_examples), BATCH):\n", + " batch = train_examples[i:i + BATCH]\n", + " micro_batches = [sft_collate(batch[j:j + MICRO], PAD_ID) for j in range(0, len(batch), MICRO)]\n", + " losses.append(sft_train_step(model, micro_batches, optimizer, max_grad_norm=1.0))\n", + " print(f\"epoch {epoch}: step losses {[round(l, 3) for l in losses[-4:]]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Validation ===\n", + "after = evaluate(model, heldout_pairs)\n", + "print(\"=== After SFT, on held-out countries ===\")\n", + "for country, capital, text, stopped in after:\n", + " print(f\"{country:>12}: {text!r} stopped={stopped} {'βœ“' if capital in text else 'βœ— ' + capital}\")\n", + "\n", + "n = len(heldout_pairs)\n", + "stop_before = sum(s for *_, s in before) / n\n", + "stop_after = sum(s for *_, s in after) / n\n", + "formatted = sum(text.startswith(\"The capital of \") and text.endswith(\".\") for *_, text, _ in after) / n\n", + "correct = sum(k in text for _, k, text, _ in after) / n\n", + "print(f\"\\nends its turn : {stop_before:.0%} -> {stop_after:.0%}\")\n", + "print(f\"follows format: {formatted:.0%}\")\n", + "print(f\"right capital : {correct:.0%} (knowledge comes from pretraining, not from SFT)\")\n", + "\n", + "# Test 1: training loss should drop substantially\n", + "print(\"\\n=== Loss Test ===\")\n", + "print(f\"first step {losses[0]:.3f} -> last step {losses[-1]:.3f}\")\n", + "assert losses[-1] < 0.5 * losses[0], \"SFT loss should at least halve\"\n", + "print(\"PASSED\")\n", + "\n", + "# Test 2: the model learned to end its turn\n", + "print(\"=== End-of-Turn Test ===\")\n", + "assert stop_after >= 0.75 and stop_after > stop_before, \"after SFT the model should emit <|im_end|> on held-out prompts\"\n", + "print(\"PASSED\")\n", + "\n", + "# Test 3: the format generalizes to unseen prompts\n", + "print(\"=== Format Test ===\")\n", + "assert formatted >= 0.75, \"held-out answers should follow the trained format\"\n", + "print(\"PASSED\")\n", + "\n", + "print(\"\\nAll tests passed!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Interview follow-ups\n", + "\n", + "1. What happens if you *don't* mask the prompt? Is that ever a good idea? How would you mask a multi-turn conversation?\n", + "2. You trained without the end-of-turn token in the labels. What does the model do at inference time?\n", + "3. Your tokenizer has no pad token, so you set `pad_token = eos_token` and build labels as `input_ids` with pad positions set to `-100`. What silently breaks?\n", + "4. Why right-pad for training but left-pad for batched generation?\n", + "5. **Packing:** you concatenate many short examples into one 4k-token row to avoid padding. What must you change in attention and position ids, and why?\n", + "6. Gradient accumulation with 4 micro-batches on each of 8 GPUs. How do you get the token-level normalization exactly right across devices?\n", + "7. On the held-out set the model may answer \"The capital of Switzerland is Geneva.\" What did SFT teach, and what didn't it teach? What happens if you fine-tune on facts the base model doesn't know?\n", + "8. Which KL divergence does SFT minimize? (See the KL divergence problem.)\n", + "9. Full fine-tuning vs LoRA for this: memory, quality, and when you'd pick each." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb b/v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb new file mode 100644 index 0000000..ecf4520 --- /dev/null +++ b/v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb @@ -0,0 +1,651 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Apply SFT on SmolLM β€” Solution\n", + "\n", + "**Difficulty**: πŸ”΄ Hard\n", + "\n", + "---\n", + "\n", + "### Problem Statement\n", + "\n", + "A pretrained **base** model is a document completer. Ask it a question and it continues the document, usually with more questions, and it never stops. **Supervised fine-tuning (SFT)** turns it into an assistant that answers in a chat format and ends its turn.\n", + "\n", + "Mechanically, SFT is next-token cross-entropy. Interviewers still ask about it, because almost every bug is in the **data path**, not the model:\n", + "\n", + "- **Loss on the prompt.** Only the response should be trained, so the prompt positions get label `-100`.\n", + "- **Off-by-one.** The logits at position `t` predict the token at `t + 1`.\n", + "- **Padding.** Pads must be excluded from the loss, and real tokens must never be mistaken for pads.\n", + "- **End of turn.** If the end-of-turn token isn't in the labels, the model never learns to stop.\n", + "- **Normalization.** Averaging per-micro-batch means under gradient accumulation weights tokens unequally. Hugging Face fixed exactly this bug in `transformers` in late 2024.\n", + "\n", + "You'll fine-tune **`HuggingFaceTB/SmolLM-135M`**, the base model whose architecture you built in the *SmolLM from Scratch* problem, on a tiny instruction dataset. Then you'll show that it follows the chat format on questions it never saw.\n", + "\n", + "---\n", + "\n", + "### Requirements\n", + "\n", + "1. **`build_sft_example`**: prompt ids + response ids + end-of-turn id β†’ `input_ids` and prompt-masked `labels`.\n", + "2. **`sft_collate`**: right-pad a list of examples into `input_ids`, `attention_mask` and `labels`.\n", + "3. **`sft_loss`**: the shifted, masked, token-level cross-entropy, with an optional global token count for gradient accumulation.\n", + "4. **`sft_train_step`**: one optimizer step over several micro-batches, normalized correctly.\n", + "5. **Fine-tune SmolLM-135M**: write `encode_chat_example`, train, and evaluate on held-out prompts.\n", + "\n", + "---\n", + "\n", + "### Constraints\n", + "\n", + "- βœ… `labels` are aligned with `input_ids` (not pre-shifted), the Hugging Face convention. The shift happens inside the loss.\n", + "- βœ… Use `IGNORE_INDEX = -100` for positions that must not be trained.\n", + "- βœ… The loss is the mean over **all response tokens in the optimizer step**, however they're split into micro-batches.\n", + "- ❌ No `trl`, `SFTTrainer`, `Trainer` or HF data collators. Don't pass `labels=` to the HF model either, since it would compute the loss for you.\n", + "\n", + "**Runtime:** about a minute on a laptop CPU after a ~270 MB model download; faster on a GPU or Apple MPS.\n", + "\n", + "---\n", + "\n", + "
\n", + " πŸ’‘ Hint\n", + "\n", + " - Build `labels` as a copy of `input_ids` and overwrite the prompt part with `-100`.\n", + " - Padding `labels` with `-100` *in the collator* is what keeps a pad id that equals a real token id from erasing real labels.\n", + " - `F.cross_entropy(..., ignore_index=-100, reduction=\"sum\")` divided by the number of non-ignored targets is a token mean. With gradient accumulation, divide by the count over **all** micro-batches instead.\n", + "\n", + "
\n", + "\n", + "---" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import copy\n", + "import random\n", + "from types import SimpleNamespace\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "\n", + "IGNORE_INDEX = -100 # F.cross_entropy's default ignore_index; Hugging Face uses the same value\n", + "\n", + "torch.manual_seed(0)\n", + "random.seed(0)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 1: Build one SFT example\n", + "\n", + "An SFT example is the chat-formatted **prompt**, then the **response**, then the **end-of-turn** token. `labels` are a copy of `input_ids` with every prompt position set to `-100`:\n", + "\n", + "```\n", + "input_ids : <|im_start|> user \\n What is ... <|im_end|> \\n <|im_start|> assistant \\n | The capital ... . | <|im_end|>\n", + "labels : -100 -100 ... ... -100 | The capital ... . | <|im_end|>\n", + "```\n", + "\n", + "Labels stay **aligned** with `input_ids`, and the shift happens in the loss. If `max_len` is given, truncate both from the right.\n", + "\n", + "Accept Python lists or 1-D tensors, and return 1-D `torch.long` tensors." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def build_sft_example(prompt_ids, response_ids, eos_id: int, max_len: int = None) -> dict:\n", + " \"\"\"\n", + " Args:\n", + " prompt_ids: token ids of the chat-formatted prompt (list or 1-D tensor).\n", + " response_ids: token ids of the assistant's response (list or 1-D tensor).\n", + " eos_id: end-of-turn token id, appended after the response and trained on.\n", + " max_len: optional; truncate input_ids and labels to this many tokens.\n", + "\n", + " Returns:\n", + " {\"input_ids\": LongTensor (L,), \"labels\": LongTensor (L,)} with prompt labels = IGNORE_INDEX.\n", + " \"\"\"\n", + " prompt = torch.as_tensor(prompt_ids, dtype=torch.long).flatten()\n", + " response = torch.as_tensor(response_ids, dtype=torch.long).flatten()\n", + " eos = torch.tensor([eos_id], dtype=torch.long)\n", + "\n", + " input_ids = torch.cat([prompt, response, eos])\n", + " # The model must learn to emit eos, so it is a target like any response token.\n", + " labels = torch.cat([torch.full_like(prompt, IGNORE_INDEX), response, eos])\n", + " if max_len is not None:\n", + " input_ids, labels = input_ids[:max_len], labels[:max_len]\n", + " return {\"input_ids\": input_ids, \"labels\": labels}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 1 tests ===\n", + "ex = build_sft_example([5, 6, 7], [8, 9], eos_id=2)\n", + "assert ex[\"input_ids\"].tolist() == [5, 6, 7, 8, 9, 2], f\"input_ids wrong: {ex['input_ids'].tolist()}\"\n", + "assert ex[\"labels\"].tolist() == [-100, -100, -100, 8, 9, 2], f\"labels wrong: {ex['labels'].tolist()}\"\n", + "assert ex[\"input_ids\"].dtype == torch.long and ex[\"labels\"].dtype == torch.long, \"use torch.long\"\n", + "\n", + "cut = build_sft_example([5, 6, 7], [8, 9], eos_id=2, max_len=4)\n", + "assert cut[\"input_ids\"].tolist() == [5, 6, 7, 8] and cut[\"labels\"].tolist() == [-100, -100, -100, 8], \\\n", + " \"truncate both tensors from the right\"\n", + "\n", + "from_tensors = build_sft_example(torch.tensor([5, 6, 7]), torch.tensor([8, 9]), eos_id=2)\n", + "assert from_tensors[\"labels\"].tolist() == ex[\"labels\"].tolist(), \"tensor inputs must work too\"\n", + "print(\"Part 1 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 2: Collate a batch (right padding)\n", + "\n", + "Examples have different lengths. Right-pad them to the longest one in the batch:\n", + "\n", + "- `input_ids`: pad with `pad_id`\n", + "- `attention_mask`: 1 for real tokens, 0 for pads\n", + "- `labels`: pad with `IGNORE_INDEX`, **not** `pad_id`\n", + "\n", + "The last point matters when `pad_id` is also a real token id. SmolLM's tokenizer has no pad token, so a common choice is to reuse `<|endoftext|>`. Many tokenizers reuse EOS. Because labels come from the examples and only the padding is `-100`, a genuine EOS in the labels survives." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def sft_collate(examples: list, pad_id: int) -> dict:\n", + " \"\"\"\n", + " Args:\n", + " examples: list of {\"input_ids\": (L_i,), \"labels\": (L_i,)} from build_sft_example.\n", + " pad_id: token id used to pad input_ids.\n", + "\n", + " Returns:\n", + " {\"input_ids\": (B, L), \"attention_mask\": (B, L), \"labels\": (B, L)}, all torch.long,\n", + " with L = max L_i and padding on the right.\n", + " \"\"\"\n", + " B = len(examples)\n", + " L = max(len(e[\"input_ids\"]) for e in examples)\n", + " input_ids = torch.full((B, L), pad_id, dtype=torch.long)\n", + " labels = torch.full((B, L), IGNORE_INDEX, dtype=torch.long)\n", + " attention_mask = torch.zeros((B, L), dtype=torch.long)\n", + " for i, e in enumerate(examples):\n", + " n = len(e[\"input_ids\"])\n", + " input_ids[i, :n] = e[\"input_ids\"]\n", + " labels[i, :n] = e[\"labels\"]\n", + " attention_mask[i, :n] = 1\n", + " return {\"input_ids\": input_ids, \"attention_mask\": attention_mask, \"labels\": labels}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 2 tests ===\n", + "EOS = PAD = 2 # deliberately the same id\n", + "exs = [build_sft_example([5, 6], [7], EOS), # length 4\n", + " build_sft_example([5, 6, 7], [8, 9, 10], EOS), # length 7\n", + " build_sft_example([5], [], EOS)] # length 2\n", + "batch = sft_collate(exs, pad_id=PAD)\n", + "assert batch[\"input_ids\"].shape == (3, 7), f\"expected (3, 7), got {tuple(batch['input_ids'].shape)}\"\n", + "assert batch[\"attention_mask\"].sum(1).tolist() == [4, 7, 2], \"attention_mask must count the real tokens\"\n", + "assert batch[\"input_ids\"][0].tolist() == [5, 6, 7, 2, 2, 2, 2], \"right-pad input_ids with pad_id\"\n", + "assert batch[\"labels\"][0].tolist() == [-100, -100, 7, 2, -100, -100, -100], \\\n", + " \"pad labels with -100, and keep the real end-of-turn label even though it equals pad_id\"\n", + "assert batch[\"labels\"][2].tolist()[:2] == [-100, 2], \"an empty response still trains the end-of-turn token\"\n", + "assert all(v.dtype == torch.long for v in batch.values()), \"all tensors should be torch.long\"\n", + "print(\"Part 2 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 3: The SFT loss\n", + "\n", + "The logits at position `t` predict token `t + 1`, so drop the last logit and the first label:\n", + "\n", + "```\n", + "logits[:, :-1] vs labels[:, 1:]\n", + "```\n", + "\n", + "Cross-entropy over the non-ignored targets gives the **token mean**: every trained token counts equally, whichever sequence it's in.\n", + "\n", + "**Gradient accumulation.** Suppose one optimizer step is split into micro-batches. If each one returns *its own* mean and you average those means, a micro-batch with 10 response tokens weighs as much as one with 1,000. Instead, compute the number of trained tokens across the **whole step** (`num_items_in_batch`, Hugging Face's name for it) and return `sum_of_token_losses / num_items_in_batch`. The micro-batch losses then **add up** to the full-batch token mean.\n", + "\n", + "Also: upcast the logits to float32 before the softmax, because models often run in bf16. Return `0` (not `nan`) if nothing is trained." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def sft_loss(logits: torch.Tensor, labels: torch.Tensor, num_items_in_batch: int = None) -> torch.Tensor:\n", + " \"\"\"\n", + " Shifted next-token cross-entropy over non-ignored labels.\n", + "\n", + " Args:\n", + " logits: (B, T, V) model outputs.\n", + " labels: (B, T), aligned with the inputs; IGNORE_INDEX where nothing is trained.\n", + " num_items_in_batch: if given, divide the summed loss by this (the trained-token count of the\n", + " whole optimizer step) instead of by this batch's own count.\n", + "\n", + " Returns:\n", + " Scalar loss.\n", + " \"\"\"\n", + " shift_logits = logits[:, :-1, :]\n", + " shift_labels = labels[:, 1:]\n", + " total = F.cross_entropy(shift_logits.reshape(-1, shift_logits.size(-1)).float(),\n", + " shift_labels.reshape(-1),\n", + " ignore_index=IGNORE_INDEX, reduction=\"sum\")\n", + " if num_items_in_batch is None:\n", + " num_items_in_batch = (shift_labels != IGNORE_INDEX).sum()\n", + " return total / max(int(num_items_in_batch), 1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 3 tests ===\n", + "torch.manual_seed(0)\n", + "B, T, V = 3, 6, 13\n", + "logits = torch.randn(B, T, V)\n", + "labels = torch.randint(0, V, (B, T))\n", + "labels[0, :2] = IGNORE_INDEX # prompt\n", + "labels[1, :4] = IGNORE_INDEX # longer prompt\n", + "labels[2, 3:] = IGNORE_INDEX # padding\n", + "\n", + "# Loop oracle: position t predicts labels[t + 1]\n", + "total, count = 0.0, 0\n", + "for b in range(B):\n", + " for t in range(T - 1):\n", + " y = labels[b, t + 1].item()\n", + " if y != IGNORE_INDEX:\n", + " total += -F.log_softmax(logits[b, t], -1)[y]\n", + " count += 1\n", + "oracle = total / count\n", + "\n", + "got = sft_loss(logits, labels)\n", + "assert got.ndim == 0, \"the loss must be a scalar\"\n", + "assert torch.allclose(got, oracle, atol=1e-5), f\"expected {oracle:.5f}, got {got:.5f}; check the shift\"\n", + "\n", + "# Logits whose next label is ignored (and the very last logit) must not matter\n", + "noisy = logits.clone()\n", + "unused = torch.zeros(B, T, dtype=torch.bool)\n", + "unused[:, :-1] = labels[:, 1:] == IGNORE_INDEX\n", + "unused[:, -1] = True\n", + "noisy[unused] = 100 * torch.randn(int(unused.sum()), V)\n", + "assert torch.allclose(sft_loss(noisy, labels), got, atol=1e-5), \"ignored positions leaked into the loss\"\n", + "\n", + "# Gradient accumulation: micro-batch losses with a global count must add up to the full loss\n", + "n_total = int((labels[:, 1:] != IGNORE_INDEX).sum())\n", + "# (rows 0-1 hold 6 trained tokens, row 2 holds 2, so the two micro-batches are unequal)\n", + "accumulated = sft_loss(logits[:2], labels[:2], n_total) + sft_loss(logits[2:], labels[2:], n_total)\n", + "assert torch.allclose(accumulated, got, atol=1e-5), \"with num_items_in_batch, micro-batch losses must sum to the full loss\"\n", + "naive = (sft_loss(logits[:2], labels[:2]) + sft_loss(logits[2:], labels[2:])) / 2\n", + "print(f\"full-batch token mean = {got:.4f} accumulated (correct) = {accumulated:.4f} mean of micro-batch means = {naive:.4f}\")\n", + "\n", + "all_ignored = torch.full((B, T), IGNORE_INDEX)\n", + "assert torch.isfinite(sft_loss(logits, all_ignored)), \"a batch with nothing to train must not give nan\"\n", + "print(\"Part 3 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 4: One optimizer step with gradient accumulation\n", + "\n", + "`micro_batches` is a list of collated batches that together form **one** optimizer step. Implement `sft_train_step`:\n", + "\n", + "1. Zero the gradients.\n", + "2. Count the trained tokens across **all** micro-batches first. Remember the shift: it's the labels from position 1 onward that count.\n", + "3. For each micro-batch, run the forward pass, compute `sft_loss(..., num_items_in_batch=total)`, and call `backward()`. Gradients add up across calls.\n", + "4. Optionally clip the gradient norm, then `optimizer.step()`.\n", + "5. Return the step's loss (the sum of the micro-batch losses) as a Python float.\n", + "\n", + "The model follows the Hugging Face contract: `model(input_ids=..., attention_mask=...)` returns an object with a `.logits` attribute. Move each batch to the model's device.\n", + "\n", + "The test checks the property that matters: **splitting a batch into micro-batches must give exactly the same update as not splitting it.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def sft_train_step(model: nn.Module, micro_batches: list, optimizer: torch.optim.Optimizer,\n", + " max_grad_norm: float = None) -> float:\n", + " \"\"\"\n", + " One optimizer step over several micro-batches.\n", + "\n", + " Args:\n", + " model: returns an object with .logits of shape (B, T, V) from model(input_ids=..., attention_mask=...).\n", + " micro_batches: list of dicts from sft_collate.\n", + " optimizer: optimizer over the model's parameters.\n", + " max_grad_norm: clip the total gradient norm to this, if given.\n", + "\n", + " Returns:\n", + " The token-mean loss of the whole step, as a float.\n", + " \"\"\"\n", + " model.train()\n", + " device = next(model.parameters()).device\n", + " optimizer.zero_grad(set_to_none=True)\n", + "\n", + " # Count before any backward: every micro-batch must divide by the same global number.\n", + " num_items = sum(int((mb[\"labels\"][:, 1:] != IGNORE_INDEX).sum()) for mb in micro_batches)\n", + "\n", + " step_loss = 0.0\n", + " for mb in micro_batches:\n", + " out = model(input_ids=mb[\"input_ids\"].to(device),\n", + " attention_mask=mb[\"attention_mask\"].to(device))\n", + " loss = sft_loss(out.logits, mb[\"labels\"].to(device), num_items_in_batch=num_items)\n", + " loss.backward() # gradients add up across micro-batches\n", + " step_loss += loss.item()\n", + "\n", + " if max_grad_norm is not None:\n", + " torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n", + " optimizer.step()\n", + " return step_loss" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Part 4 tests ===\n", + "class TinyLM(nn.Module):\n", + " \"\"\"Embedding -> tanh -> Linear. Position-wise, so padding can't leak between tokens.\"\"\"\n", + " def __init__(self, vocab=20, dim=16):\n", + " super().__init__()\n", + " self.emb = nn.Embedding(vocab, dim)\n", + " self.head = nn.Linear(dim, vocab)\n", + "\n", + " def forward(self, input_ids, attention_mask=None):\n", + " return SimpleNamespace(logits=self.head(torch.tanh(self.emb(input_ids))))\n", + "\n", + "\n", + "torch.manual_seed(0)\n", + "EOS, PAD = 19, 0\n", + "# Very different response lengths, so a wrong normalization shows up\n", + "exs = [build_sft_example([1, 2, 3], [4, 5, 6, 7, 8, 9, 10, 11], EOS),\n", + " build_sft_example([1, 2], [12], EOS),\n", + " build_sft_example([3, 4, 5, 6], [7, 8], EOS),\n", + " build_sft_example([2], [13, 14, 15], EOS)]\n", + "\n", + "model_split = TinyLM()\n", + "model_whole = copy.deepcopy(model_split)\n", + "reference = copy.deepcopy(model_split)\n", + "full = sft_collate(exs, PAD)\n", + "with torch.no_grad():\n", + " expected_loss = sft_loss(reference(input_ids=full[\"input_ids\"]).logits, full[\"labels\"]).item()\n", + "\n", + "loss_split = sft_train_step(model_split, [sft_collate(exs[:2], PAD), sft_collate(exs[2:], PAD)],\n", + " torch.optim.SGD(model_split.parameters(), lr=0.5))\n", + "loss_whole = sft_train_step(model_whole, [full], torch.optim.SGD(model_whole.parameters(), lr=0.5))\n", + "\n", + "assert abs(loss_whole - expected_loss) < 1e-5, f\"step loss {loss_whole:.5f} != full-batch loss {expected_loss:.5f}\"\n", + "assert abs(loss_split - loss_whole) < 1e-5, \\\n", + " f\"split into micro-batches the step loss is {loss_split:.5f}, unsplit {loss_whole:.5f}\"\n", + "for (name, a), b in zip(model_split.named_parameters(), model_whole.parameters()):\n", + " assert torch.allclose(a, b, atol=1e-6), \\\n", + " f\"parameter '{name}' differs: micro-batching must give the same update as one big batch\"\n", + "\n", + "# And it trains\n", + "model = TinyLM()\n", + "opt = torch.optim.Adam(model.parameters(), lr=1e-2)\n", + "first = sft_train_step(model, [full], opt)\n", + "for _ in range(50):\n", + " last = sft_train_step(model, [sft_collate(exs[:2], PAD), sft_collate(exs[2:], PAD)], opt, max_grad_norm=1.0)\n", + "assert last < 0.5 * first, f\"loss should fall on a fixed batch ({first:.3f} -> {last:.3f})\"\n", + "print(f\"loss {first:.3f} -> {last:.3f}\")\n", + "print(\"Part 4 passed βœ“\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 5: Fine-tune SmolLM-135M into a (tiny) assistant\n", + "\n", + "**Chat format.** Use the ChatML format that SmolLM's instruct models use:\n", + "\n", + "```\n", + "<|im_start|>user\n", + "What is the capital of France?<|im_end|>\n", + "<|im_start|>assistant\n", + "The capital of France is Paris.<|im_end|>\n", + "```\n", + "\n", + "The base model's vocabulary already contains `<|im_start|>` and `<|im_end|>`, but the model was never trained to use them. Here `<|im_end|>` is the end-of-turn token, and `<|endoftext|>` is reused as the pad.\n", + "\n", + "**Data.** 32 *\"What is the capital of X?\"* pairs for training, and 8 **held-out** countries for evaluation.\n", + "\n", + "**Your task.** Write `encode_chat_example`: tokenize the formatted prompt and the response **separately** (so the label boundary is exact), with `add_special_tokens=False`, and pass them to `build_sft_example`. Then run the cells below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "\n", + "MODEL_ID = \"HuggingFaceTB/SmolLM-135M\" # the *base* model, not the instruct one\n", + "device = \"cuda\" if torch.cuda.is_available() else (\"mps\" if torch.backends.mps.is_available() else \"cpu\")\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\n", + "model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32).to(device)\n", + "IM_END = tokenizer.convert_tokens_to_ids(\"<|im_end|>\")\n", + "PAD_ID = tokenizer.convert_tokens_to_ids(\"<|endoftext|>\")\n", + "print(f\"device={device} params={sum(p.numel() for p in model.parameters()) / 1e6:.0f}M \"\n", + " f\"<|im_end|>={IM_END} pad=<|endoftext|>={PAD_ID}\")\n", + "\n", + "\n", + "def chat_prompt(question: str) -> str:\n", + " return f\"<|im_start|>user\\n{question}<|im_end|>\\n<|im_start|>assistant\\n\"\n", + "\n", + "\n", + "CAPITALS = [\n", + " (\"France\", \"Paris\"), (\"Germany\", \"Berlin\"), (\"Italy\", \"Rome\"), (\"Spain\", \"Madrid\"),\n", + " (\"Japan\", \"Tokyo\"), (\"China\", \"Beijing\"), (\"India\", \"New Delhi\"), (\"Russia\", \"Moscow\"),\n", + " (\"Canada\", \"Ottawa\"), (\"Australia\", \"Canberra\"), (\"Egypt\", \"Cairo\"), (\"Kenya\", \"Nairobi\"),\n", + " (\"Brazil\", \"Brasilia\"), (\"Argentina\", \"Buenos Aires\"), (\"Mexico\", \"Mexico City\"), (\"Peru\", \"Lima\"),\n", + " (\"Chile\", \"Santiago\"), (\"Greece\", \"Athens\"), (\"Portugal\", \"Lisbon\"), (\"Poland\", \"Warsaw\"),\n", + " (\"Norway\", \"Oslo\"), (\"Sweden\", \"Stockholm\"), (\"Finland\", \"Helsinki\"), (\"Denmark\", \"Copenhagen\"),\n", + " (\"Ireland\", \"Dublin\"), (\"Austria\", \"Vienna\"), (\"Hungary\", \"Budapest\"), (\"Turkey\", \"Ankara\"),\n", + " (\"Iran\", \"Tehran\"), (\"Thailand\", \"Bangkok\"), (\"Vietnam\", \"Hanoi\"), (\"South Korea\", \"Seoul\"),\n", + " # held out: never trained on\n", + " (\"Netherlands\", \"Amsterdam\"), (\"Belgium\", \"Brussels\"), (\"Switzerland\", \"Bern\"), (\"Cuba\", \"Havana\"),\n", + " (\"Nigeria\", \"Abuja\"), (\"Morocco\", \"Rabat\"), (\"Indonesia\", \"Jakarta\"), (\"Ukraine\", \"Kyiv\"),\n", + "]\n", + "train_pairs, heldout_pairs = CAPITALS[:32], CAPITALS[32:]\n", + "\n", + "\n", + "def qa(country: str, capital: str) -> tuple:\n", + " return f\"What is the capital of {country}?\", f\"The capital of {country} is {capital}.\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def encode_chat_example(question: str, answer: str, max_len: int = 64) -> dict:\n", + " \"\"\"Chat-format one (question, answer) pair and build its SFT example.\"\"\"\n", + " # Tokenizing the two parts separately makes the prompt/response boundary exact; tokenizing\n", + " # the concatenation could merge tokens across it and shift the mask by one.\n", + " prompt_ids = tokenizer(chat_prompt(question), add_special_tokens=False).input_ids\n", + " response_ids = tokenizer(answer, add_special_tokens=False).input_ids\n", + " return build_sft_example(prompt_ids, response_ids, IM_END, max_len)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Look at what is actually trained ===\n", + "ex = encode_chat_example(*qa(\"France\", \"Paris\"))\n", + "trained = ex[\"labels\"][ex[\"labels\"] != IGNORE_INDEX]\n", + "print(\"full input :\", repr(tokenizer.decode(ex[\"input_ids\"])))\n", + "print(\"trained on :\", repr(tokenizer.decode(trained)))\n", + "assert tokenizer.decode(trained) == \"The capital of France is Paris.<|im_end|>\", \\\n", + " \"only the response and the end-of-turn token should be trained\"\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def greedy_answer(model, question: str, max_new_tokens: int = 30) -> tuple:\n", + " \"\"\"Greedy decoding until <|im_end|>. Returns (text, stopped).\"\"\"\n", + " model.eval()\n", + " ids = torch.tensor([tokenizer(chat_prompt(question), add_special_tokens=False).input_ids], device=device)\n", + " out = []\n", + " for _ in range(max_new_tokens):\n", + " next_id = model(input_ids=ids).logits[0, -1].argmax().item()\n", + " if next_id == IM_END:\n", + " return tokenizer.decode(out), True\n", + " out.append(next_id)\n", + " ids = torch.cat([ids, torch.tensor([[next_id]], device=device)], dim=1)\n", + " return tokenizer.decode(out), False\n", + "\n", + "\n", + "def evaluate(model, pairs) -> list:\n", + " return [(country, capital, *greedy_answer(model, qa(country, capital)[0])) for country, capital in pairs]\n", + "\n", + "\n", + "before = evaluate(model, heldout_pairs)\n", + "print(\"\\n=== Base model, before SFT ===\")\n", + "for country, _, text, stopped in before[:3]:\n", + " print(f\"{country:>12}: {text[:70]!r} stopped={stopped}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Fine-tune ===\n", + "torch.manual_seed(0)\n", + "random.seed(0)\n", + "train_examples = [encode_chat_example(*qa(c, k)) for c, k in train_pairs]\n", + "\n", + "EPOCHS, BATCH, MICRO = 3, 8, 4 # 2 micro-batches of 4 per optimizer step\n", + "optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.0)\n", + "losses = []\n", + "for epoch in range(EPOCHS):\n", + " random.shuffle(train_examples)\n", + " for i in range(0, len(train_examples), BATCH):\n", + " batch = train_examples[i:i + BATCH]\n", + " micro_batches = [sft_collate(batch[j:j + MICRO], PAD_ID) for j in range(0, len(batch), MICRO)]\n", + " losses.append(sft_train_step(model, micro_batches, optimizer, max_grad_norm=1.0))\n", + " print(f\"epoch {epoch}: step losses {[round(l, 3) for l in losses[-4:]]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# === Validation ===\n", + "after = evaluate(model, heldout_pairs)\n", + "print(\"=== After SFT, on held-out countries ===\")\n", + "for country, capital, text, stopped in after:\n", + " print(f\"{country:>12}: {text!r} stopped={stopped} {'βœ“' if capital in text else 'βœ— ' + capital}\")\n", + "\n", + "n = len(heldout_pairs)\n", + "stop_before = sum(s for *_, s in before) / n\n", + "stop_after = sum(s for *_, s in after) / n\n", + "formatted = sum(text.startswith(\"The capital of \") and text.endswith(\".\") for *_, text, _ in after) / n\n", + "correct = sum(k in text for _, k, text, _ in after) / n\n", + "print(f\"\\nends its turn : {stop_before:.0%} -> {stop_after:.0%}\")\n", + "print(f\"follows format: {formatted:.0%}\")\n", + "print(f\"right capital : {correct:.0%} (knowledge comes from pretraining, not from SFT)\")\n", + "\n", + "# Test 1: training loss should drop substantially\n", + "print(\"\\n=== Loss Test ===\")\n", + "print(f\"first step {losses[0]:.3f} -> last step {losses[-1]:.3f}\")\n", + "assert losses[-1] < 0.5 * losses[0], \"SFT loss should at least halve\"\n", + "print(\"PASSED\")\n", + "\n", + "# Test 2: the model learned to end its turn\n", + "print(\"=== End-of-Turn Test ===\")\n", + "assert stop_after >= 0.75 and stop_after > stop_before, \"after SFT the model should emit <|im_end|> on held-out prompts\"\n", + "print(\"PASSED\")\n", + "\n", + "# Test 3: the format generalizes to unseen prompts\n", + "print(\"=== Format Test ===\")\n", + "assert formatted >= 0.75, \"held-out answers should follow the trained format\"\n", + "print(\"PASSED\")\n", + "\n", + "print(\"\\nAll tests passed!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Interview follow-ups (with answers)\n", + "\n", + "1. **No prompt masking:** the model also learns to *generate user prompts*. Capacity goes to modelling text it will never need to write, and with short responses most of the gradient comes from the prompt. Some studies report small gains from a *down-weighted* prompt loss when the responses are very short, but masking is the default. For **multi-turn** chats, mask every user (and system) turn and train every assistant turn, each including its end-of-turn token.\n", + "\n", + "2. **It never stops.** It keeps generating until `max_new_tokens`, often inventing the next user turn. The end-of-turn token has to be a *trained* target, which is why `build_sft_example` puts it in `labels`.\n", + "\n", + "3. **Every real EOS gets masked too**, because it has the same id as the pad. The model never learns to end its turn (see question 2). Fixes: build labels per example *before* padding and pad labels with `-100` (what `sft_collate` does), mask using the `attention_mask` rather than the token id, or use a dedicated pad token.\n", + "\n", + "4. **Training:** with right padding, real tokens sit at positions `0..n-1`, so position ids and the causal mask are trivially correct, and the pads (masked from the loss) come *after* everything that matters. **Generation:** each step appends a token at the end of every row. With right padding, a short row would generate after its pads. Left padding lines up every row's last real token in the final column, which is where the next token is read from.\n", + "\n", + "5. **Packing** puts several examples in one row. Without changes, tokens from example 2 attend to example 1 (cross-contamination), and position ids keep counting from the previous example. Fixes: a **block-diagonal** causal mask, or variable-length attention kernels (FlashAttention `varlen` with `cu_seqlens`); position ids that **reset** at each example boundary; and labels that are `-100` at the first token of each new example, so you never predict across a boundary.\n", + "\n", + "6. Count the response tokens over **all micro-batches on all ranks** (an `all_reduce` of the count) before the backward passes. Then divide each micro-batch's summed loss by that global count. DDP averages gradients over ranks, so multiply by the world size, or skip DDP's averaging. Only then does every token in the global step get equal weight.\n", + "\n", + "7. **SFT taught the format and when to stop, not the facts.** \"Geneva\" comes from what the base model learned in pretraining. Fine-tuning on facts the base model *doesn't* know tends to teach it to produce confident answers without grounding. Gekhman et al. (2024) found that such examples are learned slowly and increase hallucination.\n", + "\n", + "8. **Forward KL from the data distribution to the model**, restricted to response tokens: `CE = H(data) + KL(data β€– model)`. It is mode-covering, which is part of why SFT models can be \"averaged\" and verbose, and why preference optimization (DPO, RLHF) comes next.\n", + "\n", + "9. **Full fine-tuning** updates all 135M parameters. With AdamW in fp32 that's about 16 bytes per parameter (weights, gradients and two moments), which is trivial here but not at 70B. **LoRA** trains about 0.1–1% of the parameters, fits much bigger models on one GPU, makes it cheap to keep many task adapters, and forgets less. It can underfit when the task needs large changes. Use LoRA when memory or the number of tasks is the constraint, and full fine-tuning when you need maximum quality and can afford it." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/v3/data-loading/data-parallel-loading/problem.toml b/v3/data-loading/data-parallel-loading/problem.toml index 21eed0a..faccb24 100644 --- a/v3/data-loading/data-parallel-loading/problem.toml +++ b/v3/data-loading/data-parallel-loading/problem.toml @@ -5,7 +5,6 @@ difficulty = "medium" frequency = "common" category = "data-loading" tracks = ["basics"] -companies = ["Meta", "Google", "NVIDIA", "Anthropic"] [notebooks] question = "data-parallel-loading.ipynb" @@ -16,3 +15,19 @@ tier = "C" entries = [] extras = [] device = "cpu" + +[[companies]] +name = "Meta" +confidence = "inferred" # inferred | reported + +[[companies]] +name = "Google" +confidence = "inferred" # inferred | reported + +[[companies]] +name = "NVIDIA" +confidence = "inferred" # inferred | reported + +[[companies]] +name = "Anthropic" +confidence = "inferred" # inferred | reported diff --git a/website/src/data/problems.generated.ts b/website/src/data/problems.generated.ts index 743945f..da69b36 100644 --- a/website/src/data/problems.generated.ts +++ b/website/src/data/problems.generated.ts @@ -19,6 +19,10 @@ export interface GeneratedProblem { } export const generatedProblems: GeneratedProblem[] = [ + {"slug": "buffered-loading", "legacyIds": [], "title": "Implement a Shuffle Buffer for Streaming Data", "description": "Build a shuffle buffer that wraps an IterableDataset, providing approximate shuffling with O(buffer_size) memory.", "difficulty": "medium", "category": "data-loading", "tracks": ["basics"], "companies": [], "companyConfidence": {}, "questionPath": "v3/data-loading/buffered-loading/buffered-loading.ipynb", "solutionPath": "v3/data-loading/buffered-loading/buffered-loading_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/buffered-loading/buffered-loading.ipynb", "graded": false, "entries": []}, + {"slug": "data-parallel-loading", "legacyIds": [], "title": "Simulate DistributedSampler for Data Parallel Training", "description": "Implement a distributed sampler from scratch that shards a dataset across ranks with deterministic shuffling and padding.", "difficulty": "medium", "category": "data-loading", "tracks": ["basics"], "companies": ["Meta", "Google", "NVIDIA", "Anthropic"], "companyConfidence": {"Meta": "inferred", "Google": "inferred", "NVIDIA": "inferred", "Anthropic": "inferred"}, "questionPath": "v3/data-loading/data-parallel-loading/data-parallel-loading.ipynb", "solutionPath": "v3/data-loading/data-parallel-loading/data-parallel-loading_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/data-parallel-loading/data-parallel-loading.ipynb", "graded": false, "entries": []}, + {"slug": "preloaded-dataset", "legacyIds": [], "title": "Implement a Preloaded Map-Style Dataset", "description": "Build a map-style Dataset that loads all data into memory and integrates with DataLoader for batching and shuffling.", "difficulty": "easy", "category": "data-loading", "tracks": ["basics"], "companies": [], "companyConfidence": {}, "questionPath": "v3/data-loading/preloaded-dataset/preloaded-dataset.ipynb", "solutionPath": "v3/data-loading/preloaded-dataset/preloaded-dataset_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/preloaded-dataset/preloaded-dataset.ipynb", "graded": false, "entries": []}, + {"slug": "streaming-dataset", "legacyIds": [], "title": "Implement a Streaming IterableDataset with Yield", "description": "Build an IterableDataset that streams data line by line from a file using a generator, with multi-worker sharding support.", "difficulty": "medium", "category": "data-loading", "tracks": ["basics"], "companies": [], "companyConfidence": {}, "questionPath": "v3/data-loading/streaming-dataset/streaming-dataset.ipynb", "solutionPath": "v3/data-loading/streaming-dataset/streaming-dataset_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/data-loading/streaming-dataset/streaming-dataset.ipynb", "graded": false, "entries": []}, {"slug": "lin-regression", "legacyIds": ["v1-1"], "title": "Implement Linear Regression", "description": "Build a linear regression model from scratch using PyTorch tensors and autograd, fitting a line to synthetic data with gradient descent.", "difficulty": "basic", "category": "", "tracks": ["basics"], "companies": [], "companyConfidence": {}, "questionPath": "torch/basic/lin-regression/lin-regression.ipynb", "solutionPath": "torch/basic/lin-regression/lin-regression_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/torch/basic/lin-regression/lin-regression.ipynb", "graded": true, "entries": ["LinearRegressionModel"]}, {"slug": "custom-dataset", "legacyIds": ["v1-2"], "title": "Write a Custom Dataset and DataLoader", "description": "Create a custom Dataset class and use DataLoader to batch, shuffle, and iterate over your data for training.", "difficulty": "basic", "category": "", "tracks": ["basics"], "companies": [], "companyConfidence": {}, "questionPath": "torch/basic/custom-dataset/custom-dataset.ipynb", "solutionPath": "torch/basic/custom-dataset/custom-dataset_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/torch/basic/custom-dataset/custom-dataset.ipynb", "graded": true, "entries": ["LinearRegressionDataset"]}, {"slug": "custom-activation", "legacyIds": ["v1-3"], "title": "Write a Custom Activation Function", "description": "Implement a custom activation function as a PyTorch module and use it inside a neural network.", "difficulty": "basic", "category": "", "tracks": ["basics"], "companies": [], "companyConfidence": {}, "questionPath": "torch/basic/custom-activation/custom-activation.ipynb", "solutionPath": "torch/basic/custom-activation/custom-activation_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/torch/basic/custom-activation/custom-activation.ipynb", "graded": true, "entries": ["CustomActivationModel"]}, @@ -48,7 +52,8 @@ export const generatedProblems: GeneratedProblem[] = [ {"slug": "muon", "legacyIds": ["v1-38"], "title": "Implement Muon from Scratch", "description": "Implement the Muon optimizer from scratch on top of torch.optim.Optimizer: momentum orthogonalized by Newton\u2013Schulz for matrix parameters, with an AdamW fallback for 1D parameters.", "difficulty": "hard", "category": "", "tracks": ["advanced"], "companies": ["Meta", "Google"], "companyConfidence": {"Meta": "inferred", "Google": "inferred"}, "questionPath": "torch/hard/optimizers/muon/muon.ipynb", "solutionPath": "torch/hard/optimizers/muon/muon_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/torch/hard/optimizers/muon/muon.ipynb", "graded": true, "entries": ["newton_schulz", "MyMuon"]}, {"slug": "create-embeddings-out-of-an-llm", "legacyIds": [], "title": "Create Embeddings Out Of An Llm", "description": "", "difficulty": "medium", "category": "", "tracks": [], "companies": [], "companyConfidence": {}, "questionPath": "llm/Create-Embeddings-out-of-an-LLM/embeddings-q2.ipynb", "solutionPath": null, "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Create-Embeddings-out-of-an-LLM/embeddings-q2.ipynb", "graded": false, "entries": []}, {"slug": "flash-attention", "legacyIds": [], "title": "Flash Attention", "description": "", "difficulty": "medium", "category": "", "tracks": [], "companies": [], "companyConfidence": {}, "questionPath": "llm/Flash-Attention/flash-attention.ipynb", "solutionPath": null, "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Flash-Attention/flash-attention.ipynb", "graded": true, "entries": ["flash_attention_pytorch"]}, - {"slug": "rms-norm", "legacyIds": ["v2-2"], "title": "Implement RMS Norm", "description": "Build Root Mean Square Layer Normalization used in LLaMA and modern transformers \u2014 simpler and faster than LayerNorm.", "difficulty": "easy", "category": "modern-architectures", "tracks": ["basics", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "llm/RMS-Norm/rms-norm.ipynb", "solutionPath": null, "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/RMS-Norm/rms-norm.ipynb", "graded": false, "entries": []}, + {"slug": "kl-divergence", "legacyIds": ["v2-1"], "title": "Implement KL Divergence Loss", "description": "Implement KL divergence every way it shows up in ML: from probabilities and logits, F.kl_div semantics, Jensen-Shannon, forward vs reverse, Gaussian closed form, masked token-level KL, and the k1/k2/k3 estimators used in PPO and GRPO.", "difficulty": "medium", "category": "alignment-training", "tracks": ["advanced", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "v3/alignment-training/kl-divergence/kl-divergence.ipynb", "solutionPath": "v3/alignment-training/kl-divergence/kl-divergence_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/kl-divergence/kl-divergence.ipynb", "graded": true, "entries": ["kl_categorical", "kl_from_logits", "my_kl_div", "js_divergence", "fit_gaussian", "kl_gaussian", "masked_token_kl", "kl_estimators", "kl_estimator_grad"]}, + {"slug": "rms-norm", "legacyIds": ["v2-2"], "title": "Implement RMS Norm", "description": "Build Root Mean Square Layer Normalization used in LLaMA and modern transformers \u2014 simpler and faster than LayerNorm.", "difficulty": "easy", "category": "modern-architectures", "tracks": ["basics", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "llm/RMS-Norm/rms-norm.ipynb", "solutionPath": "llm/RMS-Norm/rms-norm_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/RMS-Norm/rms-norm.ipynb", "graded": true, "entries": ["RMSNorm"]}, {"slug": "byte-pair-encoder", "legacyIds": ["v2-3"], "title": "Implement Byte Pair Encoding from Scratch", "description": "Build the BPE tokenizer algorithm that iteratively merges frequent character pairs to build a subword vocabulary.", "difficulty": "easy", "category": "", "tracks": ["basics", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "llm/Byte-Pair-Encoder/BPE-q3.ipynb", "solutionPath": "llm/Byte-Pair-Encoder/BPE-q3-SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Byte-Pair-Encoder/BPE-q3.ipynb", "graded": true, "entries": ["get_vocab", "get_stats", "merge_vocab", "byte_pair_encoding"]}, {"slug": "implement-attention-from-scratch", "legacyIds": ["v2-6"], "title": "Implement Attention from Scratch", "description": "Build scaled dot-product attention from raw matrix operations \u2014 queries, keys, values, scaling, and softmax.", "difficulty": "medium", "category": "modern-architectures", "tracks": ["advanced", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "llm/Implement-Attention-from-Scratch/attention-q4-Question.ipynb", "solutionPath": "llm/Implement-Attention-from-Scratch/attention-q4.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Implement-Attention-from-Scratch/attention-q4-Question.ipynb", "graded": true, "entries": ["scaled_dot_product_attention"]}, {"slug": "multi-head-attention", "legacyIds": ["v2-7"], "title": "Implement Multi-Head Attention from Scratch", "description": "Split attention into multiple heads with independent projections, compute attention per head, and concatenate results.", "difficulty": "medium", "category": "modern-architectures", "tracks": ["advanced", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "llm/Multi-Head-Attention/multi-head-attention-q5-Question.ipynb", "solutionPath": "llm/Multi-Head-Attention/multi-head-attention-q5.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Multi-Head-Attention/multi-head-attention-q5-Question.ipynb", "graded": true, "entries": ["multi_head_attention"]}, @@ -56,6 +61,7 @@ export const generatedProblems: GeneratedProblem[] = [ {"slug": "sinusoidal-positional-embedding", "legacyIds": ["v2-10"], "title": "Implement Sinusoidal Embeddings", "description": "Build the fixed sinusoidal positional encoding from 'Attention Is All You Need' using sin/cos at different frequencies.", "difficulty": "medium", "category": "modern-architectures", "tracks": ["advanced", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "llm/Sinusoidal-Positional-Embedding/sinusoidal-q7-Question.ipynb", "solutionPath": "llm/Sinusoidal-Positional-Embedding/sinusoidal-q7.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Sinusoidal-Positional-Embedding/sinusoidal-q7-Question.ipynb", "graded": true, "entries": ["SinusoidalPositionalEmbedding"]}, {"slug": "rotary-positional-embedding", "legacyIds": ["v2-11"], "title": "Implement ROPE Embeddings", "description": "Build Rotary Position Embeddings that encode relative positions by rotating query and key vectors in complex space.", "difficulty": "medium", "category": "modern-architectures", "tracks": ["advanced", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "llm/Rotary-Positional-Embedding/rope-q8-Question.ipynb", "solutionPath": "llm/Rotary-Positional-Embedding/rope-q8.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/Rotary-Positional-Embedding/rope-q8-Question.ipynb", "graded": true, "entries": ["rotate_half", "apply_rotary_pos_emb"]}, {"slug": "smollm", "legacyIds": ["v2-12"], "title": "Implement SmolLM from Scratch", "description": "Build a complete small language model end-to-end: tokenizer integration, transformer blocks, and autoregressive text generation.", "difficulty": "hard", "category": "modern-architectures", "tracks": ["advanced", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "llm/SmolLM/smollm-q12-Question.ipynb", "solutionPath": "llm/SmolLM/smollm-q12.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/llm/SmolLM/smollm-q12-Question.ipynb", "graded": true, "entries": ["rotate_half", "apply_rotary_pos_emb", "repeat_kv", "RotaryEmbedder", "RopeAttention"]}, + {"slug": "sft-smollm", "legacyIds": ["v2-20"], "title": "Apply SFT on SmolLM", "description": "Fine-tune the pretrained SmolLM-135M base model into a chat assistant: build prompt-masked examples, pad and collate, write the shifted token-level loss, and normalize gradient accumulation correctly.", "difficulty": "hard", "category": "alignment-training", "tracks": ["advanced", "llm-path"], "companies": [], "companyConfidence": {}, "questionPath": "v3/alignment-training/sft-smollm/sft-smollm.ipynb", "solutionPath": "v3/alignment-training/sft-smollm/sft-smollm_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/alignment-training/sft-smollm/sft-smollm.ipynb", "graded": true, "entries": ["build_sft_example", "sft_collate", "sft_loss", "sft_train_step"]}, {"slug": "softmax", "legacyIds": ["v3-1"], "title": "Implement Softmax from Scratch", "description": "Build numerically stable softmax using the log-sum-exp trick, handling overflow and underflow in raw tensor math.", "difficulty": "easy", "category": "classical-ml", "tracks": ["basics"], "companies": ["Apple", "Meta", "Google", "Amazon"], "companyConfidence": {"Apple": "inferred", "Meta": "inferred", "Google": "inferred", "Amazon": "inferred"}, "questionPath": "v3/classical-ml/softmax/softmax.ipynb", "solutionPath": "v3/classical-ml/softmax/softmax_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/classical-ml/softmax/softmax.ipynb", "graded": true, "entries": ["softmax"]}, {"slug": "kmeans", "legacyIds": ["v3-2"], "title": "Implement K-Means Clustering in PyTorch", "description": "Code Lloyd's algorithm with PyTorch tensors: random init, distance computation, centroid updates, and convergence check.", "difficulty": "easy", "category": "classical-ml", "tracks": ["basics"], "companies": ["Uber", "LinkedIn", "Google", "Amazon"], "companyConfidence": {"Uber": "inferred", "LinkedIn": "inferred", "Google": "inferred", "Amazon": "inferred"}, "questionPath": "v3/classical-ml/kmeans/kmeans.ipynb", "solutionPath": "v3/classical-ml/kmeans/kmeans_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/classical-ml/kmeans/kmeans.ipynb", "graded": true, "entries": ["kmeans"]}, {"slug": "knn", "legacyIds": ["v3-3"], "title": "Implement KNN in PyTorch", "description": "Build k-nearest neighbors classification using vectorized distance computation and top-k selection on tensors.", "difficulty": "easy", "category": "classical-ml", "tracks": ["basics"], "companies": ["Uber", "LinkedIn", "Meta"], "companyConfidence": {"Uber": "inferred", "LinkedIn": "inferred", "Meta": "inferred"}, "questionPath": "v3/classical-ml/knn/knn.ipynb", "solutionPath": "v3/classical-ml/knn/knn_SOLN.ipynb", "colabUrl": "https://colab.research.google.com/github/Exorust/TorchLeet/blob/main/v3/classical-ml/knn/knn.ipynb", "graded": true, "entries": ["knn_predict"]}, From 8d62e3dfffa66732d0d4e83686d82ad5e9f70670 Mon Sep 17 00:00:00 2001 From: Madhur Panwar <39766613+mdrpanwar@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:43:00 +0200 Subject: [PATCH 2/2] My solutions: BPE and LoRA Co-Authored-By: Claude Opus 5.5 (1M context) --- llm/Byte-Pair-Encoder/BPE-q3.ipynb | 159 ++++++++++++++++++++++++-- v3/alignment-training/lora/lora.ipynb | 100 ++++++++++++++-- 2 files changed, 236 insertions(+), 23 deletions(-) diff --git a/llm/Byte-Pair-Encoder/BPE-q3.ipynb b/llm/Byte-Pair-Encoder/BPE-q3.ipynb index d897a65..b505231 100644 --- a/llm/Byte-Pair-Encoder/BPE-q3.ipynb +++ b/llm/Byte-Pair-Encoder/BPE-q3.ipynb @@ -13,34 +13,104 @@ { "cell_type": "markdown", "metadata": {}, - "source": "## Problem: Write a Byte Pair Encoder in Python\n\n### Problem Statement\nImplement the **Byte Pair Encoding (BPE)** algorithm in Python. BPE is a subword tokenization method used in modern NLP models (GPT, LLaMA, etc.) that iteratively merges the most frequent pair of adjacent symbols in a vocabulary to build a compact subword vocabulary from a corpus of words.\n\n### Requirements\n1. **`get_vocab(corpus)`** β€” Build an initial vocabulary from a list of words:\n - Split each word into individual characters plus a special end-of-word token (e.g., ``).\n - Count how many times each word appears in the corpus.\n\n2. **`get_stats(vocab)`** β€” Count adjacent symbol pairs:\n - For each word in the vocabulary, count every pair of adjacent symbols.\n - Weight each pair's count by the word's frequency.\n\n3. **`merge_vocab(pair, vocab)`** β€” Apply a merge operation:\n - Find every occurrence of the given pair of adjacent symbols in the vocabulary.\n - Replace the pair with a single merged symbol (concatenation of the two).\n - The merged vocabulary must spell the same words as before β€” only the chunking changes.\n\n4. **`byte_pair_encoding(corpus, num_merges=10)`** β€” Run the full BPE loop:\n - Build the initial vocabulary, then repeatedly find and merge the most frequent pair.\n - Return the final vocabulary and the list of merge operations performed.\n\n### Constraints\n- Each merge must pick the most frequent adjacent pair (greedy).\n- Merging must never change the text a word spells β€” only how it is segmented.\n- The algorithm must be deterministic: same corpus and `num_merges` β†’ same result." + "source": [ + "## Problem: Write a Byte Pair Encoder in Python\n", + "\n", + "### Problem Statement\n", + "Implement the **Byte Pair Encoding (BPE)** algorithm in Python. BPE is a subword tokenization method used in modern NLP models (GPT, LLaMA, etc.) that iteratively merges the most frequent pair of adjacent symbols in a vocabulary to build a compact subword vocabulary from a corpus of words.\n", + "\n", + "### Requirements\n", + "1. **`get_vocab(corpus)`** β€” Build an initial vocabulary from a list of words:\n", + " - Split each word into individual characters plus a special end-of-word token (e.g., ``).\n", + " - Count how many times each word appears in the corpus.\n", + "\n", + "2. **`get_stats(vocab)`** β€” Count adjacent symbol pairs:\n", + " - For each word in the vocabulary, count every pair of adjacent symbols.\n", + " - Weight each pair's count by the word's frequency.\n", + "\n", + "3. **`merge_vocab(pair, vocab)`** β€” Apply a merge operation:\n", + " - Find every occurrence of the given pair of adjacent symbols in the vocabulary.\n", + " - Replace the pair with a single merged symbol (concatenation of the two).\n", + " - The merged vocabulary must spell the same words as before β€” only the chunking changes.\n", + "\n", + "4. **`byte_pair_encoding(corpus, num_merges=10)`** β€” Run the full BPE loop:\n", + " - Build the initial vocabulary, then repeatedly find and merge the most frequent pair.\n", + " - Return the final vocabulary and the list of merge operations performed.\n", + "\n", + "### Constraints\n", + "- Each merge must pick the most frequent adjacent pair (greedy).\n", + "- Merging must never change the text a word spells β€” only how it is segmented.\n", + "- The algorithm must be deterministic: same corpus and `num_merges` β†’ same result." + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Final Vocabulary:\n", + "low : 1\n", + "lowest : 1\n", + "ne w er : 1\n", + "w i d er : 1\n" + ] + } + ], "source": [ "from collections import defaultdict, Counter\n", "\n", "def get_vocab(corpus):\n", " \"\"\"Creates a vocabulary with words split into characters and a special end-of-word token.\"\"\"\n", - " ...\n", + " vocab ={}\n", + " for word in corpus:\n", + " # Split the word into characters and add a special end-of-word token\n", + " chars = list(word) + ['']\n", + " vocab[tuple(chars)] = vocab.get(tuple(chars), 0) + 1\n", " return vocab\n", "\n", "def get_stats(vocab):\n", " \"\"\"Counts frequency of adjacent symbol pairs.\"\"\"\n", - " ...\n", + " pairs = Counter() # Initialize the counter for symbol pairs\n", + " for word, freq in vocab.items():\n", + " symbols = list(word)\n", + " for i in range(len(symbols) - 1):\n", + " pairs[(symbols[i], symbols[i + 1])] += freq\n", " return pairs\n", "\n", "def merge_vocab(pair, vocab):\n", " \"\"\"Merges the most frequent pair into a single symbol.\"\"\"\n", - " ...\n", + " new_vocab = {}\n", + " bigram = ''.join(pair)\n", + " for word, freq in vocab.items():\n", + " # Replace the pair with the merged symbol\n", + " new_word = []\n", + " i = 0\n", + " while i < len(word):\n", + " if i < len(word) - 1 and (word[i], word[i + 1]) == pair:\n", + " new_word.append(bigram)\n", + " i += 2 # Skip the next symbol since it's part of the merged pair\n", + " else:\n", + " new_word.append(word[i])\n", + " i += 1\n", + " new_vocab[tuple(new_word)] = freq\n", " return new_vocab\n", "\n", "def byte_pair_encoding(corpus, num_merges=10):\n", " \"\"\"Performs BPE on a corpus.\"\"\"\n", - " ...\n", + " vocab = get_vocab(corpus)\n", + " merges = []\n", + " for _ in range(num_merges):\n", + " pairs = get_stats(vocab)\n", + " if not pairs: # pairs was an empty dictionary\n", + " break\n", + " most_frequent_pair = max(pairs, key=pairs.get) # most_frequent_pair = min(pairs, key=lambda p: (-pairs[p], p)) -- this is independent of the order of the pairs. this outputs the pair with the highest frequency, and in case of a tie, it will choose the pair that comes first in lexicographical order.\n", + " vocab = merge_vocab(most_frequent_pair, vocab)\n", + " merges.append(most_frequent_pair)\n", " return vocab, merges\n", "\n", "# Example usage\n", @@ -54,17 +124,84 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "metadata": {}, - "outputs": [], - "source": "def test_get_vocab():\n corpus = [\"test\"]\n vocab = get_vocab(corpus)\n assert vocab == {('t', 'e', 's', 't', ''): 1}\n print(\"βœ“ test_get_vocab passed\")\n\ndef test_get_stats():\n vocab = {('t', 'e', 's', 't', ''): 1}\n stats = get_stats(vocab)\n expected = {\n ('t', 'e'): 1,\n ('e', 's'): 1,\n ('s', 't'): 1,\n ('t', ''): 1\n }\n assert stats == expected\n print(\"βœ“ test_get_stats passed\")\n\ndef test_merge_vocab():\n vocab = {('t', 'e', 's', 't', ''): 1}\n merged = merge_vocab(('e', 's'), vocab)\n expected = {('t', 'es', 't', ''): 1}\n assert merged == expected\n print(\"βœ“ test_merge_vocab passed\")\n\ndef test_bpe_sequence():\n corpus = [\"low\", \"lower\", \"newest\", \"widest\"]\n final_vocab, merges = byte_pair_encoding(corpus, num_merges=5)\n assert isinstance(final_vocab, dict)\n assert all(isinstance(pair, tuple) for pair in merges)\n assert len(merges) == 5\n print(\"βœ“ test_bpe_sequence passed\")\n\n# Run all tests\ntest_get_vocab()\ntest_get_stats()\ntest_merge_vocab()\ntest_bpe_sequence()\nprint(\"\\nAll tests passed!\")" + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "βœ“ test_get_vocab passed\n", + "βœ“ test_get_stats passed\n", + "βœ“ test_merge_vocab passed\n", + "βœ“ test_bpe_sequence passed\n", + "\n", + "All tests passed!\n" + ] + } + ], + "source": [ + "def test_get_vocab():\n", + " corpus = [\"test\"]\n", + " vocab = get_vocab(corpus)\n", + " assert vocab == {('t', 'e', 's', 't', ''): 1}\n", + " print(\"βœ“ test_get_vocab passed\")\n", + "\n", + "def test_get_stats():\n", + " vocab = {('t', 'e', 's', 't', ''): 1}\n", + " stats = get_stats(vocab)\n", + " expected = {\n", + " ('t', 'e'): 1,\n", + " ('e', 's'): 1,\n", + " ('s', 't'): 1,\n", + " ('t', ''): 1\n", + " }\n", + " assert stats == expected\n", + " print(\"βœ“ test_get_stats passed\")\n", + "\n", + "def test_merge_vocab():\n", + " vocab = {('t', 'e', 's', 't', ''): 1}\n", + " merged = merge_vocab(('e', 's'), vocab)\n", + " expected = {('t', 'es', 't', ''): 1}\n", + " assert merged == expected\n", + " print(\"βœ“ test_merge_vocab passed\")\n", + "\n", + "def test_bpe_sequence():\n", + " corpus = [\"low\", \"lower\", \"newest\", \"widest\"]\n", + " final_vocab, merges = byte_pair_encoding(corpus, num_merges=5)\n", + " assert isinstance(final_vocab, dict)\n", + " assert all(isinstance(pair, tuple) for pair in merges)\n", + " assert len(merges) == 5\n", + " print(\"βœ“ test_bpe_sequence passed\")\n", + "\n", + "# Run all tests\n", + "test_get_vocab()\n", + "test_get_stats()\n", + "test_merge_vocab()\n", + "test_bpe_sequence()\n", + "print(\"\\nAll tests passed!\")" + ] } ], "metadata": { + "kernelspec": { + "display_name": "torchleet", + "language": "python", + "name": "python3" + }, "language_info": { - "name": "python" + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.17" } }, "nbformat": 4, "nbformat_minor": 2 -} \ No newline at end of file +} diff --git a/v3/alignment-training/lora/lora.ipynb b/v3/alignment-training/lora/lora.ipynb index f2ab784..2d63abb 100644 --- a/v3/alignment-training/lora/lora.ipynb +++ b/v3/alignment-training/lora/lora.ipynb @@ -68,7 +68,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "metadata": {}, "outputs": [], "source": [ @@ -80,9 +80,25 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Model architecture:\n", + "SimpleModel(\n", + " (fc1): Linear(in_features=64, out_features=128, bias=True)\n", + " (fc2): Linear(in_features=128, out_features=128, bias=True)\n", + " (fc3): Linear(in_features=128, out_features=10, bias=True)\n", + " (relu): ReLU()\n", + ")\n", + "\n", + "Total parameters: 26,122\n" + ] + } + ], "source": [ "# A simple pretrained model to apply LoRA to\n", "class SimpleModel(nn.Module):\n", @@ -131,11 +147,20 @@ " d_out = linear.out_features\n", "\n", " # TODO: Freeze original linear layer weights (set requires_grad to False)\n", - " ...\n", + " self.linear.requires_grad_(False)\n", "\n", " # TODO: Create low-rank matrices A (d_in, rank) and B (rank, d_out) as nn.Parameter\n", " # Initialize A with kaiming_uniform_ and B with zeros\n", - " ...\n", + " A_tensor = torch.empty(d_in, self.rank)\n", + " B_tensor = torch.empty(self.rank, d_out)\n", + "\n", + " # init\n", + " nn.init.kaiming_uniform_(A_tensor)\n", + " nn.init.zeros_(B_tensor)\n", + "\n", + " # make param\n", + " self.A = nn.Parameter(A_tensor)\n", + " self.B = nn.Parameter(B_tensor)\n", "\n", " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", " \"\"\"\n", @@ -145,7 +170,7 @@ " # TODO: Compute original linear output\n", " # TODO: Compute LoRA contribution: (x @ A) @ B * scaling\n", " # TODO: Return sum of both\n", - " ...\n", + " return self.linear(x) + self.scaling * (x @ self.A) @ self.B\n", "\n", "\n", "def apply_lora(model: nn.Module, rank: int = 4, alpha: float = 1.0,\n", @@ -165,7 +190,16 @@ " # TODO: Iterate through named modules\n", " # For each module that matches target_modules and is nn.Linear,\n", " # replace it with a LoRALinear wrapper using setattr\n", - " ...\n", + " for mod in target_modules:\n", + " if hasattr(model, mod):\n", + " old_modl = getattr(model, mod) # Get the original model.mod\n", + "\n", + " if isinstance(old_modl, nn.Linear):\n", + " new_modl = LoRALinear(old_modl, rank = rank, alpha = alpha) # Create the new module using LoRALinear\n", + " setattr(model, mod, new_modl) # Replace model.mod with the new module\n", + " else:\n", + " print(\"Module 'm' not found directly on M.\")\n", + " return model\n", "\n", "\n", "def merge_lora(model: nn.Module) -> nn.Module:\n", @@ -181,7 +215,14 @@ " # For each, compute merged weight: W + scaling * (A @ B).T\n", " # Create a new nn.Linear with the merged weight and original bias\n", " # Replace the LoRALinear module with the new nn.Linear\n", - " ...\n", + " for name, module in model.named_modules():\n", + " if isinstance(module, LoRALinear):\n", + " new_module = nn.Linear(module.linear.in_features, module.linear.out_features)\n", + " with torch.no_grad():\n", + " new_module.weight = nn.Parameter(module.linear.weight + (module.scaling) * (module.A @ module.B).T)\n", + " new_module.bias = nn.Parameter(module.linear.bias)\n", + " setattr(model, name, new_module)\n", + " return model\n", "\n", "\n", "def count_parameters(model: nn.Module):\n", @@ -193,9 +234,36 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 16, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "=== Parameter Count Test ===\n", + "Original model: 26,122 total, 26,122 trainable\n", + "LoRA model: 27,914 total, 3,082 trainable\n", + "Trainable ratio: 11.0%\n", + "PASSED\n", + "\n", + "=== Training Test ===\n", + "Initial loss: 2.3145\n", + "Final loss: 2.0857\n", + "PASSED\n", + "\n", + "=== Merge Test ===\n", + "Max difference: 1.79e-07\n", + "PASSED\n", + "\n", + "=== Clean Merge Test ===\n", + "No LoRALinear modules remain after merge.\n", + "PASSED\n", + "\n", + "All tests passed!\n" + ] + } + ], "source": [ "# === Validation ===\n", "torch.manual_seed(42)\n", @@ -265,13 +333,21 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "torchleet", "language": "python", "name": "python3" }, "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", "name": "python", - "version": "3.10.0" + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.17" } }, "nbformat": 4,