How akshargrover/llm_from_scratch works
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: train.py Scalable pretraining · model_gpt.py Baseline GPT · train_sft.py SFT training · train_ppo.py PPO RLHF
Components 22
Transformer Foundations
- Attention primitives · Masked single- and multi-head attention.
part_1/multi_head.py - Decoder block · Residual attention and MLP block.
part_1/block.py - Part 1 orchestrator
part_1/orchestrator.py
Language Model Evolution
- Next-token dataset · Creates shifted autoregressive examples.
part_2/dataset.py - Baseline GPT
part_2/model_gpt.py - LM training · Causal cross-entropy with validation loss.
part_2/train.py - Autoregressive sampling
part_2/sample.py - Modern decoder · Uses RMSNorm, RoPE, SwiGLU, and modern attention.
part_3/model_modern.py - KV cache · Separates full prefill from incremental decode.
part_3/kv_cache.py - Cached generation demo
part_3/demo_generate.py
Pretraining Operations
- BPE preparation
part_4/tokenizer_bpe.py - Scalable pretraining · Coordinates AMP, accumulation, scheduling, and logging.
part_4/train.py - Checkpoints · Supports resume and downstream reuse.
part_4/checkpointing.py - Sparse MoE layer · Routes tokens to selected expert MLPs and recombines outputs.
part_5/moe.py - Hybrid Transformer block
part_5/block_hybrid.py
Post-Training & Alignment
- SFT response batches · Formats instructions and masks loss to responses.
part_6/collator_sft.py - SFT training
part_6/train_sft.py - Reward model training · Ranks chosen responses above rejected responses.
part_7/train_rm.py - PPO RLHF · Updates policy/value model from reward-scored rollouts.
part_8/train_ppo.py - GRPO alignment · Uses group-relative advantages and policy-only KL-regularized loss.
part_9/train_grpo.py
Other components
- PyTorch runtime · Tensor computation, autograd, and hardware acceleration.
- Training observability · TQDM and optional Weights & Biases logging.
Connections 24
- 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