# ankhu1610/llm-from-scratch architecture

> Architecture diagram of the GitHub repository ankhu1610/llm-from-scratch, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-19). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/ankhu1610/llm-from-scratch
- Repository: https://github.com/ankhu1610/llm-from-scratch

## Overview

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: [`02_tiny_gpt/model_gpt.py`](https://github.com/ankhu1610/llm-from-scratch/blob/main/02_tiny_gpt/model_gpt.py) (Tiny GPT model), [`03_modern_transformers/model_modern.py`](https://github.com/ankhu1610/llm-from-scratch/blob/main/03_modern_transformers/model_modern.py) (Modern model), [`05_moe/block_hybrid.py`](https://github.com/ankhu1610/llm-from-scratch/blob/main/05_moe/block_hybrid.py) (MoE block), [`04_scaling/train.py`](https://github.com/ankhu1610/llm-from-scratch/blob/main/04_scaling/train.py) (Scaling train)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_g01["01 Fundamentals"]
  node_n01["Transformer block<br/>fundamentals core<br/>[block.py]"]
  node_n02["Attention math<br/>fundamentals core<br/>[single_head.py]"]
  node_n03["Positional encoding<br/>fundamentals component<br/>[pos_encoding.py]"]
end

subgraph group_g02["02 Tiny GPT"]
  node_n04["Tiny GPT model<br/>trainable LM<br/>[model_gpt.py]"]
  node_n05["Text pipeline<br/>tokenize+dataset<br/>[dataset.py]"]
  node_n06["Train loop<br/>training entrypoint<br/>[train.py]"]
end

subgraph group_g03["03 Modern Transformers"]
  node_n07["Modern model<br/>modern LM<br/>[model_modern.py]"]
  node_n08["Modern attention<br/>runtime component<br/>[attn_modern.py]"]
  node_n09["Inference runtime<br/>generation entrypoint<br/>[demo_generate.py]"]
end

subgraph group_g04["04 Scaling"]
  node_n10["Scaling train<br/>training system<br/>[train.py]"]
  node_n11["BPE data<br/>tokenization+dataset<br/>[dataset_bpe.py]"]
  node_n12["Scale utilities<br/>infra components<br/>[amp_accum.py]"]
end

subgraph group_g05["05 MoE"]
  node_n13["MoE block<br/>expert routing<br/>[block_hybrid.py]"]
  node_n14["Gate + experts<br/>moe core<br/>[gating.py]"]
  node_n15["MoE demo<br/>validation demo<br/>[demo_moe.py]"]
end

subgraph group_g06["06 SFT"]
  node_n16["SFT model flow<br/>alignment pipeline<br/>[train_sft.py]"]
  node_n17["Prompt formatting<br/>data shaping<br/>[formatters.py]"]
  node_n18["SFT batching<br/>batching+curriculum<br/>[collator_sft.py]"]
  node_n19["SFT outputs<br/>eval+sample<br/>[evaluate.py]"]
end

node_n03 -->|"feeds masks"| node_n02
node_n02 -->|"composes"| node_n01
node_n05 -->|"feeds tokens"| node_n04
node_n04 -->|"trained by"| node_n06
node_n08 -->|"powers"| node_n07
node_n07 -->|"serves generation"| node_n09
node_n11 -->|"feeds corpus"| node_n10
node_n12 -->|"supports"| node_n10
node_n14 -->|"routes into"| node_n13
node_n13 -->|"demonstrated by"| node_n15
node_n17 -->|"formats input"| node_n18
node_n18 -->|"drives"| node_n16
node_n16 -->|"produces"| node_n19
node_n01 -->|"foundation for"| node_n04
node_n04 -->|"evolves into"| node_n07
node_n07 -->|"scaled by"| node_n10
node_n10 -->|"enables"| node_n13
node_n13 -->|"followed by"| node_n16

click node_n01 "https://github.com/ankhu1610/llm-from-scratch/blob/main/01_transformer_fundamentals/block.py"
click node_n02 "https://github.com/ankhu1610/llm-from-scratch/blob/main/01_transformer_fundamentals/single_head.py"
click node_n03 "https://github.com/ankhu1610/llm-from-scratch/blob/main/01_transformer_fundamentals/pos_encoding.py"
click node_n04 "https://github.com/ankhu1610/llm-from-scratch/blob/main/02_tiny_gpt/model_gpt.py"
click node_n05 "https://github.com/ankhu1610/llm-from-scratch/blob/main/02_tiny_gpt/dataset.py"
click node_n06 "https://github.com/ankhu1610/llm-from-scratch/blob/main/02_tiny_gpt/train.py"
click node_n07 "https://github.com/ankhu1610/llm-from-scratch/blob/main/03_modern_transformers/model_modern.py"
click node_n08 "https://github.com/ankhu1610/llm-from-scratch/blob/main/03_modern_transformers/attn_modern.py"
click node_n09 "https://github.com/ankhu1610/llm-from-scratch/blob/main/03_modern_transformers/demo_generate.py"
click node_n10 "https://github.com/ankhu1610/llm-from-scratch/blob/main/04_scaling/train.py"
click node_n11 "https://github.com/ankhu1610/llm-from-scratch/blob/main/04_scaling/dataset_bpe.py"
click node_n12 "https://github.com/ankhu1610/llm-from-scratch/blob/main/04_scaling/amp_accum.py"
click node_n13 "https://github.com/ankhu1610/llm-from-scratch/blob/main/05_moe/block_hybrid.py"
click node_n14 "https://github.com/ankhu1610/llm-from-scratch/blob/main/05_moe/gating.py"
click node_n15 "https://github.com/ankhu1610/llm-from-scratch/blob/main/05_moe/demo_moe.py"
click node_n16 "https://github.com/ankhu1610/llm-from-scratch/blob/main/06_sft/train_sft.py"
click node_n17 "https://github.com/ankhu1610/llm-from-scratch/blob/main/06_sft/formatters.py"
click node_n18 "https://github.com/ankhu1610/llm-from-scratch/blob/main/06_sft/collator_sft.py"
click node_n19 "https://github.com/ankhu1610/llm-from-scratch/blob/main/06_sft/evaluate.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_n01,node_n02,node_n03 toneBlue
class node_n04,node_n05,node_n06 toneAmber
class node_n07,node_n08,node_n09 toneMint
class node_n10,node_n11,node_n12 toneRose
class node_n13,node_n14,node_n15 toneIndigo
class node_n16,node_n17,node_n18,node_n19 toneTeal
```

## Components

### 01 Fundamentals

Transformer primitives and attention math

- [Transformer block](https://github.com/ankhu1610/llm-from-scratch/blob/main/01_transformer_fundamentals/block.py): fundamentals core · Stacks positional encoding, masking, attention, and FFN into a full block · `01_transformer_fundamentals/block.py`
- [Attention math](https://github.com/ankhu1610/llm-from-scratch/blob/main/01_transformer_fundamentals/single_head.py): fundamentals core · Single-head and multi-head attention with causal masking · `01_transformer_fundamentals/single_head.py`
- [Positional encoding](https://github.com/ankhu1610/llm-from-scratch/blob/main/01_transformer_fundamentals/pos_encoding.py): fundamentals component · Position signals and attention mask utilities · `01_transformer_fundamentals/pos_encoding.py`

### 02 Tiny GPT

Minimal trainable language model

- [Tiny GPT model](https://github.com/ankhu1610/llm-from-scratch/blob/main/02_tiny_gpt/model_gpt.py): trainable LM · Minimal next-token model built from the foundational primitives · `02_tiny_gpt/model_gpt.py`
- [Text pipeline](https://github.com/ankhu1610/llm-from-scratch/blob/main/02_tiny_gpt/dataset.py): tokenize+dataset · Raw text, tokenizer, and shifted token windows for causal training · `02_tiny_gpt/dataset.py`
- [Train loop](https://github.com/ankhu1610/llm-from-scratch/blob/main/02_tiny_gpt/train.py): training entrypoint · Optimization, evaluation, and sampling for the tiny language model · `02_tiny_gpt/train.py`

### 03 Modern Transformers

Inference-oriented transformer upgrades

- [Modern model](https://github.com/ankhu1610/llm-from-scratch/blob/main/03_modern_transformers/model_modern.py): modern LM · Transformer stack upgraded for contemporary inference patterns · `03_modern_transformers/model_modern.py`
- [Modern attention](https://github.com/ankhu1610/llm-from-scratch/blob/main/03_modern_transformers/attn_modern.py): runtime component · Attention path with RoPE and streaming KV cache support · `03_modern_transformers/attn_modern.py`
- [Inference runtime](https://github.com/ankhu1610/llm-from-scratch/blob/main/03_modern_transformers/demo_generate.py): generation entrypoint · Generate tokens using the modernized model stack · `03_modern_transformers/demo_generate.py`

### 04 Scaling

Training-system and reproducibility features

- [Scaling train](https://github.com/ankhu1610/llm-from-scratch/blob/main/04_scaling/train.py): training system · AMP, gradient accumulation, logging, checkpointing, and scheduling · `04_scaling/train.py`
- [BPE data](https://github.com/ankhu1610/llm-from-scratch/blob/main/04_scaling/dataset_bpe.py): tokenization+dataset · Byte-pair tokenization and corpus preparation for scaled training · `04_scaling/dataset_bpe.py`
- [Scale utilities](https://github.com/ankhu1610/llm-from-scratch/blob/main/04_scaling/amp_accum.py): infra components · Mixed precision, checkpoints, logger, and learning-rate scheduler · `04_scaling/amp_accum.py`

### 05 MoE

Mixture-of-experts model branch

- [MoE block](https://github.com/ankhu1610/llm-from-scratch/blob/main/05_moe/block_hybrid.py): expert routing · Hybrid transformer block that routes tokens through experts · `05_moe/block_hybrid.py`
- [Gate + experts](https://github.com/ankhu1610/llm-from-scratch/blob/main/05_moe/gating.py): moe core · Routing policy and expert networks used by the MoE subsystem · `05_moe/gating.py`
- [MoE demo](https://github.com/ankhu1610/llm-from-scratch/blob/main/05_moe/demo_moe.py): validation demo · Shows routed expert behavior and hybrid block execution · `05_moe/demo_moe.py`

### 06 SFT

Supervised fine-tuning pipeline

- [SFT model flow](https://github.com/ankhu1610/llm-from-scratch/blob/main/06_sft/train_sft.py): alignment pipeline · Instruction-tuning training, evaluation, and generation entrypoints · `06_sft/train_sft.py`
- [Prompt formatting](https://github.com/ankhu1610/llm-from-scratch/blob/main/06_sft/formatters.py): data shaping · Converts prompt-response examples into training sequences and masks · `06_sft/formatters.py`
- [SFT batching](https://github.com/ankhu1610/llm-from-scratch/blob/main/06_sft/collator_sft.py): batching+curriculum · Collation, masking, and data ordering for supervised fine-tuning · `06_sft/collator_sft.py`
- [SFT outputs](https://github.com/ankhu1610/llm-from-scratch/blob/main/06_sft/evaluate.py): eval+sample · Evaluation and sampled generations for aligned models · `06_sft/evaluate.py`

## Connections

- 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