# akshargrover/llm_from_scratch architecture

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

- Interactive diagram: https://gitdiagram.com/akshargrover/llm_from_scratch
- Repository: https://github.com/akshargrover/llm_from_scratch

## Overview

#### Purpose and Execution Model
This is a staged, educational PyTorch repository that builds a Transformer LLM from attention primitives through pretraining, alignment, and RL optimization. Each `part_N` is a self-contained progression stage, generally driven by an `orchestrator.py` entry point (with the Part 6 filename misspelled as `part_6/ochestrator.py`). Training, sampling, and evaluation modules are explicit scripts rather than a unified application.

#### 1. Transformer Foundations
Part 1 constructs causal decoder-Transformer mechanics from low-level attention components: masks, single- and multi-head attention, positional encoding, feed-forward networks, and residual/normalization blocks. The control flow is tensor inputs → masked attention → MLP → stacked block output.
Paths: `part_1/single_head.py`, `part_1/multi_head.py`, `part_1/block.py`, `part_1/orchestrator.py`.

#### 2. Baseline Autoregressive LM
Part 2 turns the block-level architecture into a trainable GPT-style language model. Tokenizers and datasets produce shifted next-token examples; training computes causal cross-entropy, while sampling autoregressively feeds generated tokens back into the model. Validation loss is separated from training.
Paths: `part_2/model_gpt.py`, `part_2/dataset.py`, `part_2/train.py`, `part_2/sample.py`.

#### 3. Modern Decoder and Inference Efficiency
Part 3 evolves the baseline model toward modern LLM conventions: RMSNorm, RoPE, SwiGLU, updated attention, and KV caching. The important inference boundary is between full-sequence prefill and cache-backed token-by-token decode; the cache supports efficient continuation and streaming-oriented attention behavior.
Paths: `part_3/model_modern.py`, `part_3/attn_modern.py`, `part_3/kv_cache.py`, `part_3/demo_generate.py`.

#### 4. Scalable Pretraining Operations
Part 4 adds the operational training layer around the model: BPE tokenization and dataset preparation, AMP plus gradient accumulation, learning-rate scheduling, checkpoint save/resume, and metric logging. Checkpoints form the continuity boundary for interrupted jobs and subsequent sampling or training resumption.
Paths: `part_4/train.py`, `part_4/amp_accum.py`, `part_4/checkpointing.py`, `part_4/tokenizer_bpe.py`.

#### 5. Sparse Mixture-of-Experts
Part 5 introduces a hybrid Transformer block whose MoE layer routes token activations through selected expert networks. Gating determines routing; experts execute conditional MLP computation; the MoE module recombines outputs, providing a sparse alternative to dense feed-forward layers.
Paths: `part_5/gating.py`, `part_5/experts.py`, `part_5/moe.py`, `part_5/block_hybrid.py`.

#### 6. Post-Training: SFT and Reward Models
Part 6 converts instruction examples into prompt/response sequences with response-only causal-loss masking, collates batches, and trains an SFT model. Part 7 consumes preference pairs, scores chosen and rejected responses with a reward model, and optimizes pairwise ranking losses. This establishes the policy-data boundary and the separately trained reward-signal boundary.
Paths: `part_6/formatters.py`, `part_6/collator_sft.py`, `part_6/train_sft.py`; `part_7/model_reward.py`, `part_7/loss_reward.py`, `part_7/train_rm.py`.

#### 7. Alignment Optimization
Part 8 implements PPO RLHF: a policy generates rollouts, a reward model scores them, and PPO updates a policy/value model while constraining divergence from a reference policy. Part 9 replaces the value baseline with group-relative advantages across multiple completions and uses a policy-only GRPO loss with explicit KL regularization.
Paths: `part_8/rollout.py`, `part_8/policy.py`, `part_8/ppo_loss.py`, `part_8/train_ppo.py`; `part_9/rollout.py`, `part_9/grpo_loss.py`, `part_9/train_grpo.py`.

#### Runtime and External Dependencies
Python runs PyTorch tensor computation, autograd, and training; NumPy supports numerical preprocessing. Tokenization may use Tiktoken and Hugging Face tooling. Training observability uses TQDM and optionally Weights & Biases, with logging concentrated in `part_4/logger.py`. Hardware acceleration, mixed precision, and checkpoint persistence are material to the scaling stage.

Main entry points: [`part_4/train.py`](https://github.com/akshargrover/llm_from_scratch/blob/main/part_4/train.py) (Scalable pretraining), [`part_2/model_gpt.py`](https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/model_gpt.py) (Baseline GPT), [`part_6/train_sft.py`](https://github.com/akshargrover/llm_from_scratch/blob/main/part_6/train_sft.py) (SFT training), [`part_8/train_ppo.py`](https://github.com/akshargrover/llm_from_scratch/blob/main/part_8/train_ppo.py) (PPO RLHF)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_foundations["Transformer Foundations"]
  node_p1_attention["Attention primitives<br/>PyTorch module<br/>[multi_head.py]"]
  node_p1_block["Decoder block<br/>PyTorch module<br/>[block.py]"]
  node_p1_runner["Part 1 orchestrator<br/>entry point<br/>[orchestrator.py]"]
end

subgraph group_modeling["Language Model Evolution"]
  node_p2_data["Next-token dataset<br/>dataset pipeline<br/>[dataset.py]"]
  node_p2_model["Baseline GPT<br/>language model<br/>[model_gpt.py]"]
  node_p2_train["LM training<br/>training script<br/>[train.py]"]
  node_p2_sample["Autoregressive sampling<br/>inference script<br/>[sample.py]"]
  node_p3_model["Modern decoder<br/>language model<br/>[model_modern.py]"]
  node_p3_cache[("KV cache<br/>inference state<br/>[kv_cache.py]")]
  node_p3_generate["Cached generation demo<br/>inference script<br/>[demo_generate.py]"]
end

subgraph group_pretraining["Pretraining Operations"]
  node_p4_tokenizer["BPE preparation<br/>tokenization pipeline<br/>[tokenizer_bpe.py]"]
  node_p4_train["Scalable pretraining<br/>training script<br/>[train.py]"]
  node_p4_checkpoint[("Checkpoints<br/>training persistence<br/>[checkpointing.py]")]
  node_p5_moe["Sparse MoE layer<br/>model component<br/>[moe.py]"]
  node_p5_hybrid["Hybrid Transformer block<br/>model component<br/>[block_hybrid.py]"]
end

subgraph group_alignment["Post-Training &amp; Alignment"]
  node_p6_sft_data["SFT response batches<br/>alignment data pipeline<br/>[collator_sft.py]"]
  node_p6_train["SFT training<br/>training script<br/>[train_sft.py]"]
  node_p7_reward["Reward model training<br/>training script<br/>[train_rm.py]"]
  node_p8_ppo["PPO RLHF<br/>RL training script<br/>[train_ppo.py]"]
  node_p9_grpo["GRPO alignment<br/>RL training script<br/>[train_grpo.py]"]
end

node_torch{{"PyTorch runtime<br/>external dependency"}}
node_observability{{"Training observability<br/>external integration"}}

node_p1_attention -->|"feeds masked attention"| node_p1_block
node_p1_block -->|"demonstrated by"| node_p1_runner
node_p1_block -->|"evolves into"| node_p2_model
node_p2_data -->|"shifted token batches"| node_p2_train
node_p2_train -->|"optimizes"| node_p2_model
node_p2_model -->|"generates tokens"| node_p2_sample
node_p2_model -->|"modernized by"| node_p3_model
node_p3_model -->|"prefill and decode state"| node_p3_cache
node_p3_cache -->|"enables efficient continuation"| node_p3_generate
node_p4_tokenizer -->|"prepared token data"| node_p4_train
node_p3_model -->|"pretraining target"| node_p4_train
node_p4_train -->|"save and resume"| node_p4_checkpoint
node_p5_moe -->|"conditional MLP layer"| node_p5_hybrid
node_p4_train -.->|"scaling architecture option"| node_p5_hybrid
node_p4_checkpoint -->|"base model initialization"| node_p6_train
node_p6_sft_data -->|"masked instruction batches"| node_p6_train
node_p6_train -->|"policy responses for preferences"| node_p7_reward
node_p6_train -->|"initial policy"| node_p8_ppo
node_p7_reward -->|"scores rollouts"| node_p8_ppo
node_p8_ppo -.->|"alternative RL alignment"| node_p9_grpo
node_p7_reward -->|"reward signal"| node_p9_grpo
node_torch -->|"training runtime"| node_p4_train
node_torch -->|"RL runtime"| node_p8_ppo
node_p4_train -->|"metrics and logs"| node_observability

click node_p1_attention "https://github.com/akshargrover/llm_from_scratch/blob/main/part_1/multi_head.py"
click node_p1_block "https://github.com/akshargrover/llm_from_scratch/blob/main/part_1/block.py"
click node_p1_runner "https://github.com/akshargrover/llm_from_scratch/blob/main/part_1/orchestrator.py"
click node_p2_data "https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/dataset.py"
click node_p2_model "https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/model_gpt.py"
click node_p2_train "https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/train.py"
click node_p2_sample "https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/sample.py"
click node_p3_model "https://github.com/akshargrover/llm_from_scratch/blob/main/part_3/model_modern.py"
click node_p3_cache "https://github.com/akshargrover/llm_from_scratch/blob/main/part_3/kv_cache.py"
click node_p3_generate "https://github.com/akshargrover/llm_from_scratch/blob/main/part_3/demo_generate.py"
click node_p4_tokenizer "https://github.com/akshargrover/llm_from_scratch/blob/main/part_4/tokenizer_bpe.py"
click node_p4_train "https://github.com/akshargrover/llm_from_scratch/blob/main/part_4/train.py"
click node_p4_checkpoint "https://github.com/akshargrover/llm_from_scratch/blob/main/part_4/checkpointing.py"
click node_p5_moe "https://github.com/akshargrover/llm_from_scratch/blob/main/part_5/moe.py"
click node_p5_hybrid "https://github.com/akshargrover/llm_from_scratch/blob/main/part_5/block_hybrid.py"
click node_p6_sft_data "https://github.com/akshargrover/llm_from_scratch/blob/main/part_6/collator_sft.py"
click node_p6_train "https://github.com/akshargrover/llm_from_scratch/blob/main/part_6/train_sft.py"
click node_p7_reward "https://github.com/akshargrover/llm_from_scratch/blob/main/part_7/train_rm.py"
click node_p8_ppo "https://github.com/akshargrover/llm_from_scratch/blob/main/part_8/train_ppo.py"
click node_p9_grpo "https://github.com/akshargrover/llm_from_scratch/blob/main/part_9/train_grpo.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_p1_attention,node_p1_block,node_p1_runner toneBlue
class node_p2_data,node_p2_model,node_p2_train,node_p2_sample,node_p3_model,node_p3_cache,node_p3_generate toneAmber
class node_p4_tokenizer,node_p4_train,node_p4_checkpoint,node_p5_moe,node_p5_hybrid toneMint
class node_p6_sft_data,node_p6_train,node_p7_reward,node_p8_ppo,node_p9_grpo toneRose
class node_torch,node_observability toneNeutral
```

## Components

### Transformer Foundations

- [Attention primitives](https://github.com/akshargrover/llm_from_scratch/blob/main/part_1/multi_head.py): PyTorch module · Masked single- and multi-head attention. · `part_1/multi_head.py`
- [Decoder block](https://github.com/akshargrover/llm_from_scratch/blob/main/part_1/block.py): PyTorch module · Residual attention and MLP block. · `part_1/block.py`
- [Part 1 orchestrator](https://github.com/akshargrover/llm_from_scratch/blob/main/part_1/orchestrator.py): entry point · `part_1/orchestrator.py`

### Language Model Evolution

- [Next-token dataset](https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/dataset.py): dataset pipeline · Creates shifted autoregressive examples. · `part_2/dataset.py`
- [Baseline GPT](https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/model_gpt.py): language model · `part_2/model_gpt.py`
- [LM training](https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/train.py): training script · Causal cross-entropy with validation loss. · `part_2/train.py`
- [Autoregressive sampling](https://github.com/akshargrover/llm_from_scratch/blob/main/part_2/sample.py): inference script · `part_2/sample.py`
- [Modern decoder](https://github.com/akshargrover/llm_from_scratch/blob/main/part_3/model_modern.py): language model · Uses RMSNorm, RoPE, SwiGLU, and modern attention. · `part_3/model_modern.py`
- [KV cache](https://github.com/akshargrover/llm_from_scratch/blob/main/part_3/kv_cache.py): inference state · Separates full prefill from incremental decode. · `part_3/kv_cache.py`
- [Cached generation demo](https://github.com/akshargrover/llm_from_scratch/blob/main/part_3/demo_generate.py): inference script · `part_3/demo_generate.py`

### Pretraining Operations

- [BPE preparation](https://github.com/akshargrover/llm_from_scratch/blob/main/part_4/tokenizer_bpe.py): tokenization pipeline · `part_4/tokenizer_bpe.py`
- [Scalable pretraining](https://github.com/akshargrover/llm_from_scratch/blob/main/part_4/train.py): training script · Coordinates AMP, accumulation, scheduling, and logging. · `part_4/train.py`
- [Checkpoints](https://github.com/akshargrover/llm_from_scratch/blob/main/part_4/checkpointing.py): training persistence · Supports resume and downstream reuse. · `part_4/checkpointing.py`
- [Sparse MoE layer](https://github.com/akshargrover/llm_from_scratch/blob/main/part_5/moe.py): model component · Routes tokens to selected expert MLPs and recombines outputs. · `part_5/moe.py`
- [Hybrid Transformer block](https://github.com/akshargrover/llm_from_scratch/blob/main/part_5/block_hybrid.py): model component · `part_5/block_hybrid.py`

### Post-Training & Alignment

- [SFT response batches](https://github.com/akshargrover/llm_from_scratch/blob/main/part_6/collator_sft.py): alignment data pipeline · Formats instructions and masks loss to responses. · `part_6/collator_sft.py`
- [SFT training](https://github.com/akshargrover/llm_from_scratch/blob/main/part_6/train_sft.py): training script · `part_6/train_sft.py`
- [Reward model training](https://github.com/akshargrover/llm_from_scratch/blob/main/part_7/train_rm.py): training script · Ranks chosen responses above rejected responses. · `part_7/train_rm.py`
- [PPO RLHF](https://github.com/akshargrover/llm_from_scratch/blob/main/part_8/train_ppo.py): RL training script · Updates policy/value model from reward-scored rollouts. · `part_8/train_ppo.py`
- [GRPO alignment](https://github.com/akshargrover/llm_from_scratch/blob/main/part_9/train_grpo.py): RL training script · Uses group-relative advantages and policy-only KL-regularized loss. · `part_9/train_grpo.py`

### Other components

- **PyTorch runtime**: external dependency · Tensor computation, autograd, and hardware acceleration.
- **Training observability**: external integration · TQDM and optional Weights & Biases logging.

## Connections

- Attention primitives → Decoder block: feeds masked attention
- Decoder block → Part 1 orchestrator: demonstrated by
- Decoder block → Baseline GPT: evolves into
- Next-token dataset → LM training: shifted token batches
- LM training → Baseline GPT: optimizes
- Baseline GPT → Autoregressive sampling: generates tokens
- Baseline GPT → Modern decoder: modernized by
- Modern decoder → KV cache: prefill and decode state
- KV cache → Cached generation demo: enables efficient continuation
- BPE preparation → Scalable pretraining: prepared token data
- Modern decoder → Scalable pretraining: pretraining target
- Scalable pretraining → Checkpoints: save and resume
- Sparse MoE layer → Hybrid Transformer block: conditional MLP layer
- Scalable pretraining → Hybrid Transformer block: scaling architecture option
- Checkpoints → SFT training: base model initialization
- SFT response batches → SFT training: masked instruction batches
- SFT training → Reward model training: policy responses for preferences
- SFT training → PPO RLHF: initial policy
- Reward model training → PPO RLHF: scores rollouts
- PPO RLHF → GRPO alignment: alternative RL alignment
- Reward model training → GRPO alignment: reward signal
- PyTorch runtime → Scalable pretraining: training runtime
- PyTorch runtime → PPO RLHF: RL runtime
- Scalable pretraining → Training observability: metrics and logs