AxonLM

Production-grade GPT-2 (124M–1.5B) pre-training — Karpathy's "Let's reproduce GPT-2" lecture architecture, refactored to clean OOP/Design Pattern standards for Computational Neuroscience LLM research.

---

Overview

This project is a fully annotated, modular, production-ready implementation of GPT-2 pre-training from scratch in PyTorch.
It is based on Andrej Karpathy's landmark 4-hour lecture "Let's reproduce GPT-2 (124M)" and the accompanying build-nanogpt repository, refactored to:

• Clean OOP architecture with a single-responsibility class per concern.
• Design Patterns: Factory Method, Strategy, Observer, Facade, Repository, Template Method, Builder.
• Computational Neuroscience docstrings: every algorithmic decision is grounded in a neuroscience analogy.
• Full production features: DDP multi-GPU, torch.compile, Flash Attention, AMP, gradient accumulation, checkpoint management, HellaSwag evaluation.

---

Architecture

src/
├── config/
│   ├── model_config.py     — GPTConfig, ModelScale (Factory Method)
│   ├── training_config.py  — TrainingConfig + cosine LR schedule
│   ├── data_config.py      — DataConfig
│   └── system_config.py    — SystemConfig (auto DDP detection)
├── model/
│   └── gpt.py              — CausalSelfAttention, MLP, TransformerBlock, GPT
├── data/
│   ├── dataloader.py       — ShardedDataLoader, DataLoaderFactory
│   └── downloader.py       — FineWebDownloader (HuggingFace → binary shards)
├── training/
│   ├── trainer.py          — Trainer (Facade over full training loop)
│   └── checkpoint.py       — CheckpointManager (Repository Pattern)
├── evaluation/
│   └── hellaswag.py        — HellaSwagEvaluator
└── utils/
    └── __init__.py         — setup_logging, VRAMProfiler, set_seed

---

Quick Start

1. Install dependencies

pip install -r requirements.txt

2. Download & tokenize FineWeb-Edu (one-time)

python scripts/download_fineweb.py --target_dir data/fineweb_edu

This creates ~10 GB of pre-tokenized binary shard files (100M tokens each).

3. Train — Single GPU

python train.py --scale small --max_steps 19073 --eval_hellaswag

4. Train — Multi-GPU (8× A100)

torchrun --standalone --nproc_per_node=8 train.py \
    --scale small \
    --micro_batch 64 \
    --total_batch 524288 \
    --dtype bfloat16

5. Generate text

# From HuggingFace pre-trained weights:
python generate.py --pretrained small --prompt "The hippocampus encodes"

# From your trained checkpoint:
python generate.py --ckpt checkpoints/best.pt --prompt "Neurons communicate"

# Interactive REPL:
python generate.py --pretrained small --interactive

6. Run tests

pytest tests/ -v
pytest tests/ -v --cov=src --cov-report=term-missing

---

Key Features

Feature	Implementation
Flash Attention	F.scaled_dot_product_attention (PyTorch 2.0+)
Gradient accumulation	Configurable via total_batch_size / micro_batch_size
torch.compile	Enabled by default; --no_compile to disable
bfloat16 AMP	torch.autocast + no loss scaling needed
DDP multi-GPU	torchrun + DistributedDataParallel
Cosine LR + warmup	TrainingConfig.get_lr(step)
Weight tying	lm_head.weight ≡ wte.weight
Weight init scaling	1/√(2L) for residual projections
HuggingFace weights	GPT.from_pretrained(ModelScale.SMALL)
HellaSwag eval	Per-step accuracy tracking (no fine-tuning)
Checkpoint resume	Auto-detects latest checkpoint

---

Model Scales

Scale	Layers	Heads	Embd	Params
small	12	12	768	~124M
medium	24	16	1024	~355M
large	36	20	1280	~774M
xl	48	25	1600	~1558M

---

Training Recipe (GPT-3 paper, 124M scale)

Hyperparameter	Value
Effective batch size	524,288 tokens (2^19)
Peak learning rate	6e-4
Min learning rate	6e-5 (max_lr / 10)
Warmup steps	715 (~375M tokens)
Total steps	19,073 (~10B tokens)
Weight decay	0.1 (2D tensors only)
β₁, β₂	0.9, 0.95
Gradient clip	1.0

---

Computational Neuroscience Docstring Highlights

Every architectural component includes a biologically-grounded analogy:

• Residual stream ↔ apical dendritic summation in pyramidal neurons
• Multi-head attention ↔ parallel processing streams (dorsal/ventral)
• LayerNorm ↔ divisive normalization in sensory cortex
• GELU ↔ mean-field firing rate of sigmoidal neurons under Gaussian input
• Flash Attention ↔ working memory limitations (O(N) vs O(N²))
• Cosine LR schedule ↔ simulated annealing / exploration-exploitation
• Temperature sampling ↔ inverse temperature in Boltzmann neural decisions
• HellaSwag ↔ predictive coding / temporal sequence completion

---

References

• Radford et al. (2019). Language Models are Unsupervised Multitask Learners. OpenAI Blog. [GPT-2]
• Brown et al. (2020). Language Models are Few-Shot Learners. NeurIPS 2020. [GPT-3]
• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention. NeurIPS 2022.
• Zellers et al. (2019). HellaSwag: Can a Machine Really Finish Your Sentence? ACL 2019.
• Hoffmann et al. (2022). Training Compute-Optimal Large Language Models. [Chinchilla]
• Friston (2010). The free-energy principle: a unified brain theory? Nature Reviews Neuroscience.
• Karpathy (2024). Let's reproduce GPT-2 (124M). YouTube lecture