# LLMs-from-scratch Administrator & Developer Guide

## 1. Architecture Overview

**LLMs-from-scratch** is the companion code repository for *Build a Large Language Model From Scratch* by Sebastian Raschka (Manning Publications, ISBN 978-1633437166). The repository implements a complete GPT-style large language model pipeline — tokenization, attention, transformer architecture, pretraining, classification fine-tuning, and instruction-following fine-tuning — entirely from scratch in PyTorch. It also extends the GPT-2 base to modern architectures including LLaMA 3.2, Qwen3, Gemma3/4, OLMo3, Tiny-Aya, and Qwen3.5.

The codebase serves two roles: (1) a pedagogical resource where each chapter folder is self-contained with inline implementations and Jupyter notebooks, and (2) an installable Python package (`llms_from_scratch`) that consolidates all implementations into reusable, tested modules.

### Component architecture

```text
┌─────────────────────────────────────────────────────────────┐
│                    Repository Structure                      │
├────────────────────────┬────────────────────────────────────┤
│   Chapter Code (ch01-ch07, appendix-A..E)                   │
│   - Self-contained per chapter                              │
│   - previous_chapters.py carries forward dependencies       │
│   - Jupyter notebooks (.ipynb) + standalone scripts (.py)   │
├────────────────────────┼────────────────────────────────────┤
│   pkg/llms_from_scratch (installable package)               │
│   - ch02..ch07 modules consolidate chapter implementations  │
│   - llama3.py, qwen3.py: modern architecture implementations│
│   - kv_cache/, kv_cache_batched/: cached inference          │
│   - generate.py, utils.py: shared utilities                 │
│   - tests/: comprehensive test suite                        │
├────────────────────────┼────────────────────────────────────┤
│   setup/ (installation and environment)                     │
│   - pip, pixi, uv, Docker, AWS SageMaker                   │
│   - python_environment_check.py                             │
└────────────────────────┴────────────────────────────────────┘
```

### Component inventory

| Component | Path | Role |
|---|---|---|
| Chapter 1 | `ch01/` | Introduction (reading list only) |
| Chapter 2 | `ch02/` | Tokenization, BPE, data loading |
| Chapter 3 | `ch03/` | Self-attention, causal attention, multi-head attention |
| Chapter 4 | `ch04/` | GPTModel assembly, KV cache, GQA, MLA, SWA, MoE, DeltaNet |
| Chapter 5 | `ch05/` | Pretraining, weight loading, LLaMA/Qwen3/Gemma/OLMo extensions |
| Chapter 6 | `ch06/` | Classification fine-tuning (spam, IMDB), LoRA experiments |
| Chapter 7 | `ch07/` | Instruction SFT, DPO, evaluation, dataset generation |
| Appendix A | `appendix-A/` | Multi-GPU DDP training |
| Appendix D | `appendix-D/` | Learning rate scheduling (warmup + cosine annealing) |
| Appendix E | `appendix-E/` | LoRA parameter-efficient fine-tuning |
| Package | `pkg/llms_from_scratch/` | Installable consolidation of all implementations |
| Setup | `setup/` | Environment setup, Docker, AWS CloudFormation |
| Root config | `conftest.py`, `pixi.toml`, `pyproject.toml`, `requirements.txt` | Project configuration |

### Technology stack

| Layer | Technology |
|---|---|
| Language | Python 3.9+ |
| Deep learning | PyTorch >= 2.2.2 |
| Tokenization | tiktoken (OpenAI BPE) |
| Weight loading | TensorFlow (GPT-2 checkpoint parsing), safetensors (modern models) |
| Visualization | matplotlib |
| Data processing | pandas, numpy |
| Notebooks | JupyterLab >= 4.0 |
| Chat interfaces | Chainlit |
| Testing | pytest, nbval, pytest-ruff |
| Package management | pip, pixi, uv |

> [!NOTE]
> TensorFlow is a required dependency solely for parsing OpenAI's original GPT-2 TensorFlow checkpoints. It is not used for training or inference. Modern model weights (LLaMA, Qwen3, Gemma) use safetensors format instead.

### Data flow — pretraining pipeline

```text
Raw text → tiktoken BPE tokenizer → GPTDatasetV1 (sliding window)
    → DataLoader → GPTModel forward pass → cross-entropy loss
    → backprop → optimizer step → checkpoint save
```

### Data flow — classification fine-tuning

```text
CSV data → SpamDataset (tokenize + pad) → DataLoader
    → GPTModel (pretrained) + replaced output head
    → cross-entropy loss → train_classifier_simple() → saved .pth
```

### Data flow — instruction fine-tuning

```text
JSON instructions → InstructionDataset (format + tokenize)
    → custom_collate_fn (pad + mask) → DataLoader
    → GPTModel (pretrained) → cross-entropy with ignore_index
    → train_model_simple() → saved .pth
```

---

## 2. System Requirements

| Resource | Minimum | Recommended | Notes |
|---|---|---|---|
| Python | 3.9 | 3.10-3.12 | pixi pins 3.10 |
| PyTorch | 2.2.2 | 2.5+ | Intel macOS pins `<2.6` |
| TensorFlow | 2.16.2 (Intel macOS) / 2.18.0 | 2.18+ | GPT-2 checkpoint loading only |
| RAM | 8 GB | 16 GB+ | For loading GPT-2 1558M |
| GPU VRAM | N/A (CPU works) | 8 GB+ | For training; 124M fits in ~2 GB |
| Disk | 2 GB | 20 GB+ | Full GPT-2 weights for all sizes |
| OS | Linux x64, macOS ARM64, Windows x64 | Linux with NVIDIA GPU | pixi supports all three |

### Python dependencies (from `requirements.txt`)

| Package | Version Constraint | Used By |
|---|---|---|
| `torch` | `>=2.2.2` (macOS x86: `<2.6`) | All chapters |
| `jupyterlab` | `>=4.0` | All notebooks |
| `tiktoken` | `>=0.5.1` | ch02, ch04, ch05 |
| `matplotlib` | `>=3.7.1` | ch04, ch06, ch07 |
| `tensorflow` | `>=2.18.0` (macOS x86: `>=2.16.2`) | ch05, ch06, ch07 |
| `tqdm` | `>=4.66.1` | ch05, ch07 |
| `numpy` | `>=1.26` | Transitive dependency |
| `pandas` | `>=2.2.1` | ch06 |
| `psutil` | `>=5.9.5` | ch07 |

### pixi environment (from `pixi.toml`)

The pixi environment specifies the `conda-forge` channel and pins:
- `python = "3.10.*"`
- `pytorch-cpu = ">=2.6.0,<3"` (CPU-only PyTorch from conda-forge)
- `tensorflow` via pypi (conda-forge feedstock broken for TF 2.18)
- `llms-from-scratch` installed as editable from project root

The `tests` feature environment adds: `blobfile`, `huggingface_hub`, `ipywidgets`, `safetensors`, `sentencepiece`, `transformers`, `pytest`, `nbval`, `pytest-ruff`, and `tensorflow-io-gcs-filesystem==0.31.0` (Windows only).

---

## 3. Deployment & Scaling

### Docker deployment

Docker environment setup is documented in `setup/03_optional-docker-environment/`. Consult the folder's `README.md` for the Dockerfile and compose configuration.

### AWS SageMaker deployment

A CloudFormation template at `setup/04_optional-aws-sagemaker-notebook/cloudformation-template.yml` provisions a SageMaker notebook instance with:
- Pre-installed Python dependencies
- GPU-enabled instance type support
- Access to the repository code

> [!IMPORTANT]
> The CloudFormation template creates billable AWS resources. Review the instance type and storage settings before deploying. Shut down instances when not in use to avoid charges.

### Multi-GPU training

Appendix A provides two approaches for Distributed Data Parallel (DDP) training:

**Using `torchrun` (recommended):**

```bash
torchrun --nproc_per_node=2 appendix-A/01_main-chapter-code/DDP-script-torchrun.py
```

**Using manual process launch:**

```bash
CUDA_VISIBLE_DEVICES=0,1 python appendix-A/01_main-chapter-code/DDP-script.py
```

The DDP scripts implement:
- `torch.distributed.init_process_group()` with NCCL backend
- `DistributedSampler` for data partitioning
- `DistributedDataParallel` model wrapper
- Rank-based logging (only rank 0 prints)

The optimized multi-GPU training script at `ch05/10_llm-training-speed/02_opt_multi_gpu_ddp.py` combines DDP with `torch.compile` and `PyTorchMultiHeadAttention` (FlashAttention) for maximum training throughput.

> [!WARNING]
> The DDP scripts default to 2 GPUs. Using a different GPU count requires uncommenting configuration code on lines 103-107 in the DDP scripts. Running with fewer GPUs than expected causes `ZeroDivisionError`.

---

## 4. Operations

### Environment verification

Run the environment checker to validate all dependencies:

```bash
python setup/02_installing-python-libraries/python_environment_check.py
```

The checker:
1. Validates Python version >= 3.9
2. Parses `requirements.txt` respecting environment markers
3. Compares installed versions against version specifiers
4. Reports `[OK]` or `[FAIL]` per package

### Model weight management

GPT-2 weights are downloaded on-demand by `download_and_load_gpt2()` in `gpt_download.py`:

| Size | Parameters | Download Size (approx.) | Download Location |
|---|---|---|---|
| `124M` | 124 million | ~500 MB | `gpt2/124M/` |
| `355M` | 355 million | ~1.4 GB | `gpt2/355M/` |
| `774M` | 774 million | ~3 GB | `gpt2/774M/` |
| `1558M` | 1.5 billion | ~6 GB | `gpt2/1558M/` |

Files downloaded per model:
- `checkpoint`, `encoder.json`, `hparams.json`
- `model.ckpt.data-00000-of-00001`, `model.ckpt.index`, `model.ckpt.meta`
- `vocab.bpe`

Primary URL: `https://openaipublic.blob.core.windows.net/gpt-2/models/{size}/`
Backup URL: `https://f001.backblazeb2.com/file/LLMs-from-scratch/gpt2/{size}/`

> [!TIP]
> The download function checks file size before re-downloading. If a file already exists with the correct size, it skips the download. This makes re-runs efficient.

Modern model weights (LLaMA, Qwen3, Gemma, OLMo, Tiny-Aya) are loaded via HuggingFace Hub's `safetensors` format using the `huggingface_hub` and `safetensors` packages (part of `requirements-extra.txt` or the pixi `tests` environment).

### Checkpoint management

Training scripts save checkpoints as `.pth` files:

| Script | Output | Location |
|---|---|---|
| `ch06/01_main-chapter-code/gpt_class_finetune.py` | `review_classifier.pth` | `ch06/01_main-chapter-code/` |
| `ch07/01_main-chapter-code/gpt_instruction_finetuning.py` | `gpt2-medium355M-sft.pth` | `ch07/01_main-chapter-code/` |
| `ch05/03_bonus_pretraining_on_gutenberg/pretraining_simple.py` | periodic checkpoints | `--output_dir` (default: `model_checkpoints/`) |

Checkpoints are standard PyTorch `state_dict` saves loaded via `torch.load(..., weights_only=True)`.

---

## 5. Performance

### Training speed optimization progression

The `ch05/10_llm-training-speed/` directory demonstrates a step-by-step optimization progression:

| Script | Technique | Speedup Factor |
|---|---|---|
| `00_orig.py` | Baseline: manual `MultiHeadAttention`, no compilation | 1x |
| `01_opt_single_gpu.py` | `PyTorchMultiHeadAttention` (FlashAttention), `torch.compile`, fused QKV | 2-4x |
| `02_opt_multi_gpu_ddp.py` | All above + `DistributedDataParallel` | Linear with GPU count |

**Key optimizations in `01_opt_single_gpu.py`:**

- Replaces the manual `MultiHeadAttention` with `PyTorchMultiHeadAttention` that uses `F.scaled_dot_product_attention` (which dispatches to FlashAttention when available)
- `FeedForwardFast` uses `nn.SiLU()` activation (hardware-accelerated)
- `torch.compile(model)` for kernel fusion
- `torch.set_float32_matmul_precision("high")` for tensor core utilization

**Memory estimation tools:**

Each architecture variant includes memory estimation scripts for comparing KV cache memory usage:

| Script | Estimates |
|---|---|
| `ch04/04_gqa/memory_estimator_gqa.py` | GQA KV cache vs. standard MHA |
| `ch04/05_mla/memory_estimator_mla.py` | MLA latent compression savings |
| `ch04/06_swa/memory_estimator_swa.py` | SWA window-bounded cache |
| `ch04/07_moe/memory_estimator_moe.py` | MoE parameter count vs. dense FFN |
| `ch04/08_deltanet/plot_memory_estimates_gated_deltanet.py` | DeltaNet recurrent state memory |

> [!NOTE]
> The memory estimators compute analytical estimates (not empirical measurements). They are useful for comparing architectural tradeoffs at different scales without needing the actual GPU hardware.

---

## 6. Security

### External network requests

The codebase makes network requests in these scenarios:

| Purpose | URL(s) | When |
|---|---|---|
| GPT-2 weight download | `openaipublic.blob.core.windows.net`, `f001.backblazeb2.com` | Chapter 5 `gpt_download.py` |
| Spam dataset download | External URL in `gpt_class_finetune.py` | Chapter 6 training |
| HuggingFace model download | `huggingface.co` | Bonus sections (LLaMA, Qwen3, Gemma) |
| Ollama API | `localhost:11434` | Chapter 7 evaluation |
| Link checking (CI) | Various (URLs in notebooks) | `conftest.py` with `--check-links` |

> [!CAUTION]
> The `gpt_download.py` scripts download model weights directly from external URLs without cryptographic verification (no hash checking). In sensitive environments, download weights manually, verify their integrity, and place them in the expected directory structure.

### Dependency security

- TensorFlow is included only for checkpoint loading compatibility. It does not run any training or inference computation.
- `conftest.py` implements retry logic with configurable timeout for link checking in CI, governed by environment variables: `CHECK_LINKS_TIMEOUT` (default 10.0), `CHECK_LINKS_RETRIES` (default 2), `CHECK_LINKS_RETRY_BACKOFF` (default 1.0).

### Data handling

- Training data (spam SMS, IMDB, instruction datasets) is downloaded on-demand by training scripts
- No user data is collected, stored, or transmitted
- Model checkpoints contain only model weights (no user data)
- The `--test_mode` flag in training scripts reduces dataset size for safe experimentation

---

## 7. Developer Internals

### Repository structure

```text
LLMs-from-scratch/
├── ch01/                          # Chapter 1: Introduction
│   └── reading-recommendations.md
├── ch02/                          # Chapter 2: Text data
│   ├── 01_main-chapter-code/      # Tokenization, dataloaders
│   ├── 02_bonus_bytepair-encoder/ # GPT-2 BPE reference
│   ├── 03_bonus_embedding-vs-matmul/
│   ├── 04_bonus_dataloader-intuition/
│   └── 05_bpe-from-scratch/       # BPE from scratch
├── ch03/                          # Chapter 3: Attention
│   ├── 01_main-chapter-code/      # Self-attention → MHA
│   ├── 02_bonus_efficient-multihead-attention/
│   └── 03_understanding-buffers/
├── ch04/                          # Chapter 4: GPT model
│   ├── 01_main-chapter-code/      # Full GPTModel
│   ├── 02_performance-analysis/   # FLOPs
│   ├── 03_kv-cache/               # KV caching
│   ├── 04_gqa/                    # Grouped-query attention
│   ├── 05_mla/                    # Multi-head latent attention
│   ├── 06_swa/                    # Sliding window attention
│   ├── 07_moe/                    # Mixture of experts
│   └── 08_deltanet/               # DeltaNet
├── ch05/                          # Chapter 5: Pretraining
│   ├── 01_main-chapter-code/      # GPT-2 weight loading
│   ├── 02_alternative_weight_loading/
│   ├── 03_bonus_pretraining_on_gutenberg/
│   ├── 04_learning_rate_schedulers/
│   ├── 05_bonus_hparam_tuning/
│   ├── 06_user_interface/         # Chainlit GPT-2 chat
│   ├── 07_gpt_to_llama/           # GPT → LLaMA conversion
│   ├── 08_memory_efficient_weight_loading/
│   ├── 09_extending-tokenizers/   # tiktoken extension
│   ├── 10_llm-training-speed/     # Training optimizations
│   ├── 11_qwen3/                  # Qwen3 from scratch
│   ├── 12_gemma3/                 # Gemma 3 from scratch
│   ├── 13_olmo3/                  # OLMo 3 from scratch
│   ├── 14_ch05_with_other_llms/   # ch05 with LLaMA/Qwen3
│   ├── 15_tiny-aya/               # Tiny-Aya from scratch
│   ├── 16_qwen3.5/                # Qwen3.5 (GatedDeltaNet)
│   └── 17_gemma4/                 # Gemma 4 from scratch
├── ch06/                          # Chapter 6: Classification FT
│   ├── 01_main-chapter-code/      # Spam classification
│   ├── 02_bonus_additional-experiments/ # LoRA, ablations
│   ├── 03_bonus_imdb-classification/   # IMDB + BERT comparison
│   └── 04_user_interface/         # Chainlit classifier
├── ch07/                          # Chapter 7: Instruction FT
│   ├── 01_main-chapter-code/      # SFT pipeline
│   ├── 02_dataset-utilities/      # Dedup, cleaning
│   ├── 03_model-evaluation/       # Eval with Ollama/OpenAI
│   ├── 04_preference-tuning-with-dpo/ # DPO from scratch
│   ├── 05_dataset-generation/     # Synthetic data generation
│   └── 06_user_interface/         # Chainlit instruction chat
├── appendix-A/                    # DDP multi-GPU training
├── appendix-B/                    # References
├── appendix-C/                    # Exercise pointers
├── appendix-D/                    # LR scheduling
├── appendix-E/                    # LoRA fine-tuning
├── pkg/llms_from_scratch/         # Installable package
├── reasoning-from-scratch/        # Experimental reasoning
├── setup/                         # Installation guides
├── conftest.py                    # pytest link-checking config
├── pixi.toml                      # pixi environment
├── pyproject.toml                 # Package metadata
└── requirements.txt               # Core dependencies
```

### The `previous_chapters.py` contract

Each chapter from ch04 onward includes a `previous_chapters.py` file that consolidates all classes and functions from prior chapters. This design decision:

- Makes each chapter folder self-contained (no cross-chapter imports)
- Allows students to modify one chapter without breaking later chapters
- Duplicates code intentionally for pedagogical independence
- The `pkg/llms_from_scratch/` package serves as the single-source-of-truth consolidation

The `previous_chapters.py` files in chapters 4-7 each contain the full `GPTModel`, `MultiHeadAttention`, `LayerNorm`, `GELU`, `FeedForward`, `TransformerBlock`, and `GPTDatasetV1` implementations.

### Package module reference (`pkg/llms_from_scratch/`)

| Module | Exports | Description |
|---|---|---|
| `ch02.py` | `GPTDatasetV1`, `create_dataloader_v1` | Tokenization and data loading |
| `ch03.py` | `SelfAttention_v1/v2`, `CausalAttention`, `MultiHeadAttentionWrapper`, `MultiHeadAttention`, `PyTorchMultiHeadAttention` | Full attention mechanism progression |
| `ch04.py` | `LayerNorm`, `GELU`, `FeedForward`, `TransformerBlock`, `GPTModel`, `FeedForwardFast`, `TransformerBlockFast`, `GPTModelFast` | GPT architecture (standard + optimized) |
| `ch05.py` | `download_and_load_gpt2`, `load_weights_into_gpt`, `generate`, `text_to_token_ids`, `token_ids_to_text`, `calc_loss_batch`, `evaluate_model`, `generate_and_print_sample`, `train_model_simple` | Pretraining utilities |
| `ch06.py` | `SpamDataset`, `classify_review`, `calc_accuracy_loader`, `train_classifier_simple` | Classification fine-tuning |
| `ch07.py` | `InstructionDataset`, `format_input`, `custom_collate_fn` | Instruction fine-tuning |
| `llama3.py` | `Llama3Model`, `TransformerBlock`, `FeedForward`, `GroupedQueryAttention`, `Llama3Tokenizer`, `ChatFormat`, `GroupedQueryAttentionFast`, `TransformerBlockFast`, `Llama3ModelFast` | LLaMA 3 from-scratch implementation |
| `qwen3.py` | `Qwen3Model`, `TransformerBlock`, `FeedForward`, `MoEFeedForward`, `GroupedQueryAttention`, `RMSNorm`, `Qwen3Tokenizer` | Qwen3 from-scratch implementation (dense + MoE) |
| `generate.py` | Text generation utilities | Shared generation logic |
| `utils.py` | `download_file`, notebook execution utilities | Shared utility functions |
| `appendix_a.py` | `NeuralNetwork`, `ToyDataset` | DDP training example classes |
| `appendix_d.py` | `train_model` (with LR scheduling), `find_highest_gradient` | Cosine annealing trainer |
| `appendix_e.py` | `LoRALayer`, `LinearWithLoRA`, `replace_linear_with_lora` | LoRA implementation |

### Key class hierarchy

The core architecture follows this class hierarchy (all inherit from `nn.Module`):

```text
GPTModel
├── tok_emb: nn.Embedding(vocab_size, emb_dim)
├── pos_emb: nn.Embedding(context_length, emb_dim)
├── drop_emb: nn.Dropout(drop_rate)
├── trf_blocks: nn.Sequential
│   └── TransformerBlock (×n_layers)
│       ├── att: MultiHeadAttention
│       │   ├── W_query, W_key, W_value: nn.Linear(emb_dim, emb_dim)
│       │   ├── out_proj: nn.Linear(emb_dim, emb_dim)
│       │   └── mask: buffer (causal mask)
│       ├── ff: FeedForward
│       │   └── layers: nn.Sequential(Linear, GELU, Linear)
│       ├── norm1, norm2: LayerNorm(emb_dim)
│       └── drop_shortcut: nn.Dropout(drop_rate)
├── final_norm: LayerNorm(emb_dim)
└── out_head: nn.Linear(emb_dim, vocab_size)
```

### Architecture variant classes

| Variant | Module | Key Class | Difference from Base |
|---|---|---|---|
| KV Cache | `ch04/03_kv-cache/` | `MultiHeadAttention` with `reset_cache()` | Caches key/value tensors across forward passes |
| GQA | `ch04/04_gqa/` | `GroupedQueryAttention` | `n_kv_groups` heads share key/value projections |
| MLA | `ch04/05_mla/` | `MultiHeadLatentAttention` | Compresses KV into low-rank latent space |
| SWA | `ch04/06_swa/` | `MultiHeadAttentionWithSWA` | Limits attention to sliding window + optional full layers |
| MoE | `ch04/07_moe/` | `MoEFeedForward` | Top-k expert routing replaces dense FFN |
| DeltaNet | `ch05/16_qwen3.5/` | `Qwen3_5GatedDeltaNet` | Gated linear attention (recurrent, no softmax) |

### GPT-2 model configurations

| Size | `vocab_size` | `context_length` | `emb_dim` | `n_heads` | `n_layers` | `drop_rate` | `qkv_bias` |
|---|---|---|---|---|---|---|---|
| 124M | 50257 | 1024 | 768 | 12 | 12 | 0.1 | True |
| 355M | 50257 | 1024 | 1024 | 16 | 24 | 0.0 | True |
| 774M | 50257 | 1024 | 1280 | 20 | 36 | 0.0 | True |
| 1558M | 50257 | 1024 | 1600 | 25 | 48 | 0.0 | True |

> [!NOTE]
> The 124M config uses `drop_rate=0.1` for training examples. When loading pretrained weights for inference, set `drop_rate=0.0` (as done in chapters 5-7).

### Qwen3 architecture implementation

The `pkg/llms_from_scratch/qwen3.py` implements the full Qwen3 architecture with:

- `Qwen3Model` — main model class with RoPE positional encoding
- `GroupedQueryAttention` — GQA with rotary positional embeddings
- `RMSNorm` — Root Mean Square Layer Normalization (replacing LayerNorm)
- `FeedForward` — SwiGLU-based feed-forward (three linear layers)
- `MoEFeedForward` — Mixture-of-experts variant with top-k routing
- `Qwen3Tokenizer` — tokenizer using `sentencepiece`

### LLaMA 3 architecture implementation

The `pkg/llms_from_scratch/llama3.py` implements:

- `Llama3Model` — with RoPE, RMSNorm, SwiGLU FFN
- `GroupedQueryAttention` — same GQA mechanism as Qwen3
- `Llama3Tokenizer` — using `sentencepiece` with `blobfile`
- `ChatFormat` — LLaMA chat template formatting
- Fast variants (`Llama3ModelFast`, `GroupedQueryAttentionFast`, `TransformerBlockFast`) using `F.scaled_dot_product_attention`

### Code patterns and conventions

**1. Configuration dictionaries:** All model configurations use Python dictionaries (not dataclasses or config files). Keys follow a consistent naming convention: `vocab_size`, `context_length`, `emb_dim`, `n_heads`, `n_layers`, `drop_rate`, `qkv_bias`.

**2. Weight loading pattern:** The `load_weights_into_gpt()` function maps TensorFlow checkpoint parameter names to PyTorch state dict keys via manual dictionary traversal and `assign()` calls with shape validation.

**3. Training loop pattern:** All training functions follow the same structure:
```python
for epoch in range(n_epochs):
    model.train()
    for input_batch, target_batch in train_loader:
        optimizer.zero_grad()
        loss = calc_loss_batch(input_batch, target_batch, model, device)
        loss.backward()
        optimizer.step()
```

**4. Classification adaptation:** The GPT-2 model is adapted for classification by replacing `out_head` with a new linear layer: `model.out_head = torch.nn.Linear(emb_dim, num_classes)`. Only the logits from the last token position are used for classification: `logits = model(input_batch)[:, -1, :]`.

**5. Instruction formatting:** Instructions use a consistent template:
```text
Below is an instruction that describes a task. Write a response
that appropriately completes the request.

### Instruction:
{instruction}

### Response:
{response}
```

**6. Device handling:** Scripts consistently use `torch.device("cuda" if torch.cuda.is_available() else "cpu")` with an `auto` option in `gpt_generate.py` that also checks for MPS.

> [!TIP]
> The `GPTModelFast` and `TransformerBlockFast` variants in `pkg/llms_from_scratch/ch04.py` use `PyTorchMultiHeadAttention` and `nn.SiLU()` for production-grade inference speed. Use these when performance matters more than pedagogical clarity.

### Testing strategy

The repository uses `pytest` with several test categories:

**Package unit tests** (`pkg/llms_from_scratch/tests/`):

| Test File | Coverage |
|---|---|
| `test_ch02.py` | Dataloader creation |
| `test_ch03.py` | Attention mechanisms |
| `test_ch04.py` | GPTModel forward pass |
| `test_ch05.py` | Weight loading, generation |
| `test_ch06.py` | Classification training |
| `test_ch07.py` | Instruction fine-tuning |
| `test_generate.py` | Text generation |
| `test_llama3.py` | LLaMA 3 implementation |
| `test_qwen3.py` | Qwen3 implementation |
| `test_appendix_a.py` | DDP utilities |
| `test_appendix_d.py` | LR scheduling |
| `test_appendix_e.py` | LoRA |

**Chapter-level tests** (in `tests.py` files within chapter folders): Validate that chapter code runs correctly end-to-end, typically with `--test_mode` for fast execution.

**Notebook validation**: The `nbval` pytest plugin validates notebook outputs. Combined with `pytest-ruff` for linting.

**Architecture variant tests** (in `tests/` subfolders within bonus sections): Specific tests for LLaMA, Qwen3, Gemma, OLMo, and Tiny-Aya implementations.

**Link checking**: `conftest.py` provides a `--check-links` pytest option that wraps `requests.Session.request` with configurable timeout and retry logic:
- `CHECK_LINKS_TIMEOUT`: Request timeout in seconds (default: 10.0)
- `CHECK_LINKS_RETRIES`: Max retries for GET/HEAD (default: 2)
- `CHECK_LINKS_RETRY_BACKOFF`: Backoff multiplier (default: 1.0)

```bash
# Run unit tests
pytest pkg/llms_from_scratch/tests/

# Run with link checking
pytest --check-links

# Run specific chapter tests
pytest ch05/01_main-chapter-code/tests.py
```

> [!NOTE]
> Many architecture variant tests require downloading model weights from HuggingFace Hub (several GB per model). Run these selectively with adequate disk space and network access.

### Error handling patterns

The codebase uses a consistent error handling approach:

| Error Type | Pattern | Example |
|---|---|---|
| Invalid model size | `ValueError` with allowed sizes list | `f'Model size not in {allowed_sizes}'` |
| Shape mismatch | `ValueError` with both shapes | `f'Shape mismatch. Left: {left.shape}, Right: {right.shape}'` |
| Missing file | `FileNotFoundError` with path | Model weight files not found |
| Invalid CLI arg | `ValueError` with valid options | `'Invalid --trainable_layers argument.'` |
| Missing service | `RuntimeError` with instructions | `'Ollama not running. Launch ollama before proceeding.'` |
| Math constraint | `ValueError` with constraint | `'n_kv_groups must divide n_heads exactly.'` |

### CLI reference — all scripts with arguments

**Chapter 4 architecture scripts:**

All accept: `--emb_dim` (int), `--n_heads` (int), `--n_layers` (int), `--max_new_tokens` (int, default 200)

| Script | Additional Args |
|---|---|
| `ch04/04_gqa/gpt_with_kv_gqa.py` | `--n_kv_groups` (int, default 2) |
| `ch04/05_mla/gpt_with_kv_mla.py` | `--latent_dim` (int, optional) |
| `ch04/06_swa/gpt_with_kv_swa.py` | `--sliding_window_size` (int, default 1024), `--sliding_window_stride` (int, default 2) |
| `ch04/07_moe/gpt_with_kv_ffn.py` | `--hidden_dim` (int), `--no_kv_cache` (flag) |
| `ch04/07_moe/gpt_with_kv_moe.py` | `--hidden_dim`, `--no_kv_cache`, `--num_experts` (int, default 0), `--num_experts_per_tok` (int, default 2) |

**Memory estimator scripts:**

| Script | Key Args |
|---|---|
| `ch04/04_gqa/memory_estimator_gqa.py` | `--context_length`, `--emb_dim`, `--n_heads`, `--n_layers`, `--n_kv_groups`, `--batch_size`, `--dtype` |
| `ch04/05_mla/memory_estimator_mla.py` | Same + `--latent_dim` |
| `ch04/06_swa/memory_estimator_swa.py` | Same + `--sliding_window_size`, `--swa_ratio` (e.g., `5:1`) |
| `ch04/07_moe/memory_estimator_moe.py` | `--emb_dim`, `--hidden_dim`, `--ffn_type` (`gelu`/`swiglu`), `--num_experts`, `--top_k`, `--dtype`, `--match_dense` |
| `ch04/08_deltanet/plot_memory_estimates_gated_deltanet.py` | `--batch`, `--emb_dim`, `--n_heads`, `--n_layers`, `--dtype`, `--min_ctx`, `--max_ctx` |

**Chapter 5 scripts:**

| Script | Key Args |
|---|---|
| `ch05/01_main-chapter-code/gpt_generate.py` | `--prompt` (str), `--device` (str) |
| `ch05/03_bonus_pretraining_on_gutenberg/prepare_dataset.py` | `--data_dir`, `--max_size_mb`, `--output_dir` |
| `ch05/03_bonus_pretraining_on_gutenberg/pretraining_simple.py` | `--data_dir`, `--output_dir`, `--n_epochs`, `--print_sample_iter`, `--eval_freq`, `--save_ckpt_freq`, `--lr`, `--batch_size`, `--debug` |

**Chapter 6 scripts:**

| Script | Key Args |
|---|---|
| `ch06/01_main-chapter-code/gpt_class_finetune.py` | `--test_mode` |
| `ch06/02_bonus_additional-experiments/additional_experiments.py` | `--model_size`, `--weights`, `--trainable_layers`, `--trainable_token_pos`, `--average_embeddings`, `--context_length`, `--lora_rank`, `--lora_alpha`, `--no_padding`, `--num_epochs`, `--batch_size`, `--accumulation_steps`, `--disable_causal_mask`, `--ignore_index` |
| `ch06/03_bonus_imdb-classification/train_gpt.py` | `--model_size`, `--weights`, `--trainable_layers`, `--trainable_token_pos`, `--average_embeddings`, `--context_length`, `--num_epochs`, `--learning_rate`, `--compile` |
| `ch06/03_bonus_imdb-classification/train_bert_hf.py` | `--model`, `--trainable_layers`, `--use_attention_mask`, `--num_epochs`, `--learning_rate` |

**Chapter 7 scripts:**

| Script | Key Args |
|---|---|
| `ch07/01_main-chapter-code/gpt_instruction_finetuning.py` | `--test_mode` |
| `ch07/01_main-chapter-code/ollama_evaluate.py` | `--file_path` (required) |
| `ch07/01_main-chapter-code/exercise_experiments.py` | `--exercise_solution` (str, e.g., `baseline`) |
| `ch07/02_dataset-utilities/find-near-duplicates.py` | `--json_file`, `--threshold`, `--remove_duplicates`, `--json_output_file` |

### Known issues from the community

| Issue | Status | Description |
|---|---|---|
| #1008 | Closed | SWA window size calculation problem |
| #1002 | Closed | `find_freq_pair` crash on short sequences (<2 tokens) |
| #996 | Closed | Request to revert `--exclude-newer` pip constraint |
| #978 | Open | Feature request: DeepSeek Sparse Attention (DSA) |
| #977 | Open | Feature request: Apple M-series (MPS) device support in scripts |

### Contributing

The repository is licensed under Apache License 2.0. Contributions follow standard GitHub fork-and-PR workflow:

1. Fork the repository
2. Create a feature branch
3. Implement changes following the existing code patterns
4. Add tests if adding new functionality
5. Run the test suite: `pytest pkg/llms_from_scratch/tests/`
6. Submit a pull request

**Code style:** The pixi test environment includes `pytest-ruff` for linting. Code follows standard Python conventions with type hints used sparingly (matching the book's pedagogical style).

### Potential improvements

Based on community requests and code analysis:

- **MPS device support** (#977): Several scripts hardcode `cuda` as the only GPU option. Adding MPS detection would benefit Apple Silicon users
- **DeepSeek Sparse Attention** (#978): Community request for DSA implementation similar to existing SWA/GQA variants
- **Unified CLI entry point**: Individual chapter scripts could benefit from a unified CLI interface (currently each script is standalone)
- **Typed configurations**: Model configs currently use plain dicts; dataclasses would add validation

---

_Generated May 2026 from rasbt/LLMs-from-scratch at 28c65cd. AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it. Unofficial; not affiliated with or reviewed by the project._