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

┌─────────────────────────────────────────────────────────────┐
│                    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

Technology stack

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

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

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

Data flow — instruction fine-tuning

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

Python dependencies (from requirements.txt)

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):

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

Using manual process launch:

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:

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:

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:

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:

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:

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:

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

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/)

Key class hierarchy

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

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

GPT-2 model configurations

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:

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:

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/):

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)

# 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:

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)

Memory estimator scripts:

Chapter 5 scripts:

Chapter 6 scripts:

Chapter 7 scripts:

Known issues from the community

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