nanogpt-training by benchflow-ai

---
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

| Topic | Reference |
|-------|-----------|
| Model Architecture | [GPT Architecture](references/gpt-architecture.md) |
| Data Loading | [Tokenized Data](references/tokenized-data.md) |
| Optimizers | [Optimizers](references/optimizers.md) |
| Training Loop | [Training Loop](references/training-loop.md) |
| Hyperparameters | [Hyperparameters](references/hyperparameters.md) |

## Installation

```bash
pip install torch einops numpy huggingface_hub
```

## Minimal Example

```python
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

```python
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