- name
- nanogpt-training
- description
- Train GPT-2 scale models (~124M parameters) efficiently on a single GPU. Covers GPT-124M architecture, tokenized dataset loading (e.g., HuggingFace Hub shards), modern optimizers (Muon, AdamW), mixed precision training, and training loop implementation.
NanoGPT Training
Overview
Training GPT-2 scale models (~124M parameters) efficiently on a single GPU. It provides:
- GPT-124M Architecture: Standard transformer with RoPE and modern optimizations
- Tokenized Datasets: Loading pre-tokenized shards from HuggingFace Hub or local files
- Modern Optimizers: Muon optimizer with Newton-Schulz orthogonalization
- Mixed Precision: bfloat16 training on A100 for 2x speedup
Training options:
- Baseline GPT: Standard residual connections
- Experimental residual variants: Optional alternative residual schemes for stability/efficiency
Quick Reference
Installation
pip install torch einops numpy huggingface_hub
Minimal Example
import modal
app = modal.App("gpt-training")
image = modal.Image.debian_slim(python_version="3.11").pip_install(
"torch", "einops", "numpy", "huggingface_hub"
)
@app.function(gpu="A100", image=image, timeout=3600)
def train():
import torch
from dataclasses import dataclass
@dataclass
class GPTConfig:
block_size: int = 1024
vocab_size: int = 50257
n_layer: int = 12
n_head: int = 12
n_embd: int = 768
dropout: float = 0.0
bias: bool = False
# Download data, build model, train
# ... (see references for full implementation)
return {"final_loss": final_loss}
@app.local_entrypoint()
def main():
results = train.remote()
print(results)
Common Imports
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler
from dataclasses import dataclass
from einops import rearrange, repeat, reduce
import numpy as np
import math
When to Use What
| Scenario | Approach |
|---|
| Standard GPT training | Use baseline model with standard residuals |
| Stability experiments | Try alternative residual variants or extra streams |
| Small experiments | Use T4/A10G GPU |
| Full training | Use A100 with bfloat16 |
| Custom data | Modify the dataset loader class |
| Different model size | Adjust GPTConfig parameters |
Metrics to Monitor
| Metric | Typical Signal | Notes |
|---|
| Validation loss | Steady decrease | Absolute value depends on dataset/tokenizer |
| Grad norm | Moderate, stable range | Large spikes indicate instability |
| Training stability | Smooth curves | Frequent spikes suggest LR/batch issues |
| Throughput | Consistent tokens/sec | Use for comparing configs |
External Resources
- nanoGPT: https://github.com/karpathy/nanoGPT
- build-nanogpt: https://github.com/karpathy/build-nanogpt
- modded-nanogpt: https://github.com/KellerJordan/modded-nanogpt
- FineWeb-Edu token shards: https://huggingface.co/datasets/karpathy/fineweb-edu-100B-gpt2-token-shards
