How ankhu1610/llm-from-scratch works
This repository is organized as a teaching pipeline for building LLMs in stages. Each numbered directory is a mostly self-contained subsystem, and each one has its own entrypoint-style orchestrator.py plus small demos, training scripts, and tests. The repo is not a single monolith; it is a sequence of focused modules that progressively replace simpler components with more production-like ones.
1. Global structure and flow
The top-level README.md describes the intended learning path: transformer fundamentals → tiny language model → modern transformer improvements → scaling techniques → MoE → SFT, with later RLHF sections outlined but not yet present in the tree. Architecturally, the codebase is a progression of increasingly complete LM stacks rather than one shared framework.
2. 01_transformer_fundamentals: core transformer primitives
This chapter contains the lowest-level building blocks: positional encoding, causal attention masking, single-head attention, multi-head attention, feed-forward networks, and full transformer blocks. The presence of attn_numpy_demo.py, demo_mha_shapes.py, and demo_visualize_multi_head.py shows it is both math-first and visualization-oriented. tests/ validate attention math and causal masking, so the boundary here is “correctness of the transformer equations,” not end-to-end training.
3. Foundational data flow
In the fundamentals chapter, data flows from small tensors through attention math, then through stacking primitives into a complete block. vis_utils.py suggests an explicit visualization layer for inspecting shapes and head behavior. This chapter is about internal transformer mechanics, not tokenization or optimization.
4. 02_tiny_gpt: a minimal trainable LM
This chapter turns the primitives into a working next-token predictor. The main flow is: raw text files (tiny.txt, tiny_hi.txt) → tokenizer.py → dataset.py for shifted token windows → model_gpt.py → train.py → eval_loss.py and sample.py. This is the first end-to-end training boundary in the repo.
5. Tiny GPT training and evaluation boundary
The tiny GPT code separates training, evaluation, and generation. utils.py likely holds shared helpers, while tests focus on tokenizer behavior and dataset shifting, which are the key correctness boundaries for causal LM training. The architectural concern here is “does the data pipeline align with next-token loss?”
6. 03_modern_transformers: inference-oriented modernizations
This chapter refactors the GPT-style model around modern transformer components: rmsnorm.py, rope_custom.py, swiglu.py, and kv_cache.py. attn_modern.py, block_modern.py, and model_modern.py show that the architecture is still transformer-based, but optimized for current LLM design patterns. demo_generate.py indicates the main external behavior is faster/cleaner generation.
7. Modern transformer runtime path
The data flow here is generation-centric: tokenized input enters the modern model, attention uses RoPE and KV cache, and output tokens are produced incrementally. The tests check shape behavior for KV cache, RMSNorm, and RoPE application, which are the critical interfaces for correctness in streaming inference.
8. 04_scaling: training-system concerns
This chapter adds systems-level training features: mixed precision and gradient accumulation (amp_accum.py), checkpointing (checkpointing.py), learning-rate scheduling (lr_scheduler.py), logging (logger.py), and BPE tokenization (tokenizer_bpe.py, dataset_bpe.py). Compared with earlier chapters, the boundary shifts from model architecture to training infrastructure and reproducibility.
9. Scaling data and control flow
The scaling pipeline is: corpus/token stream → BPE dataset → training loop with accumulation/AMP → scheduler and logger → checkpoints → sampling. Tests verify scheduler behavior, BPE tokenization, and resuming shapes, which suggests the main architectural risk in this chapter is keeping training state consistent across restarts.
10. 05_moe: expert routing as a model subsystem
The MoE chapter introduces a new internal model branch: gating.py, experts.py, moe.py, and block_hybrid.py. demo_moe.py plus tests for gate shapes, MoE forward behavior, and hybrid blocks indicate the subsystem boundary is “dense transformer block plus routed expert computation.” This is a model-architecture extension rather than a training-system change.
11. 06_sft: supervised alignment pipeline
The SFT chapter implements instruction tuning. formatters.py and dataset_sft.py convert prompt/response data into model-ready sequences, collator_sft.py handles batching and masking, curriculum.py controls data ordering, and train_sft.py, evaluate.py, and sample_sft.py cover the training/evaluation loop. Tests focus on formatting and masking, which are the key correctness boundaries for causal loss on instruction data.
12. Cross-cutting patterns
Nearly every chapter follows the same internal pattern: an orchestrator.py entrypoint, core model/data modules, one or more demos or sampling scripts, and a small tests/ suite. This makes the repository diagram best understood as a layered training stack: mathematical primitives → minimal LM → modernized LM → scalable training → expert routing → supervised fine-tuning.
13. What is intentionally absent
The README outlines reward modeling and RLHF/GRPO, but those subsystems are not present in the file tree. For an architecture diagram, they should be treated as roadmap items, not implemented components. The current repository ends at SFT.
Main entry points: model_gpt.py Tiny GPT model · model_modern.py Modern model · block_hybrid.py MoE block · train.py Scaling train
Components 19
01 Fundamentals
Transformer primitives and attention math
- Transformer block · Stacks positional encoding, masking, attention, and FFN into a full block
01_transformer_fundamentals/block.py - Attention math · Single-head and multi-head attention with causal masking
01_transformer_fundamentals/single_head.py - Positional encoding · Position signals and attention mask utilities
01_transformer_fundamentals/pos_encoding.py
02 Tiny GPT
Minimal trainable language model
- Tiny GPT model · Minimal next-token model built from the foundational primitives
02_tiny_gpt/model_gpt.py - Text pipeline · Raw text, tokenizer, and shifted token windows for causal training
02_tiny_gpt/dataset.py - Train loop · Optimization, evaluation, and sampling for the tiny language model
02_tiny_gpt/train.py
03 Modern Transformers
Inference-oriented transformer upgrades
- Modern model · Transformer stack upgraded for contemporary inference patterns
03_modern_transformers/model_modern.py - Modern attention · Attention path with RoPE and streaming KV cache support
03_modern_transformers/attn_modern.py - Inference runtime · Generate tokens using the modernized model stack
03_modern_transformers/demo_generate.py
04 Scaling
Training-system and reproducibility features
- Scaling train · AMP, gradient accumulation, logging, checkpointing, and scheduling
04_scaling/train.py - BPE data · Byte-pair tokenization and corpus preparation for scaled training
04_scaling/dataset_bpe.py - Scale utilities · Mixed precision, checkpoints, logger, and learning-rate scheduler
04_scaling/amp_accum.py
05 MoE
Mixture-of-experts model branch
- MoE block · Hybrid transformer block that routes tokens through experts
05_moe/block_hybrid.py - Gate + experts · Routing policy and expert networks used by the MoE subsystem
05_moe/gating.py - MoE demo · Shows routed expert behavior and hybrid block execution
05_moe/demo_moe.py
06 SFT
Supervised fine-tuning pipeline
- SFT model flow · Instruction-tuning training, evaluation, and generation entrypoints
06_sft/train_sft.py - Prompt formatting · Converts prompt-response examples into training sequences and masks
06_sft/formatters.py - SFT batching · Collation, masking, and data ordering for supervised fine-tuning
06_sft/collator_sft.py - SFT outputs · Evaluation and sampled generations for aligned models
06_sft/evaluate.py
Connections 18
- Positional encoding → Attention math: feeds masks · Position and mask helpers support attention computation
- Attention math → Transformer block: composes · Attention math is assembled into a transformer block
- Text pipeline → Tiny GPT model: feeds tokens · Tokenized shifted windows drive the tiny GPT model
- Tiny GPT model → Train loop: trained by · Model is optimized, evaluated, and sampled through the train loop
- Modern attention → Modern model: powers · Modern attention is embedded in the upgraded transformer stack
- Modern model → Inference runtime: serves generation · The modern model is exposed through the generation demo
- BPE data → Scaling train: feeds corpus · BPE datasets supply the scaled training loop
- Scale utilities → Scaling train: supports · AMP, checkpointing, logging, and scheduling shape training behavior
- Gate + experts → MoE block: routes into · Gating and experts form the hybrid MoE block
- MoE block → MoE demo: demonstrated by · The demo exercises MoE routing and hybrid execution
- Prompt formatting → SFT batching: formats input · Prompt/response examples are converted into masked training batches
- SFT batching → SFT model flow: drives · Batched instruction data is consumed by the SFT loop
- SFT model flow → SFT outputs: produces · Aligned models are evaluated and sampled after training
- Transformer block → Tiny GPT model: foundation for · Transformer primitives are reused to build the tiny GPT model
- Tiny GPT model → Modern model: evolves into · The tiny GPT stack is refactored into modern transformer components
- Modern model → Scaling train: scaled by · Modern model training is extended with infrastructure features
- Scaling train → MoE block: enables · Scaled training infrastructure supports larger routed-model experiments
- MoE block → SFT model flow: followed by · The architecture progresses from MoE modeling to supervised alignment