NanoGPT Training Guide

NanoGPT Training Guide

Training Loop (~300 Lines)

NanoGPT's train.py is a self-contained training script with minimal dependencies.

Complete Training Script Structure

# train.py (simplified)
import os
import time
import math
import pickle
import torch
from model import GPTConfig, GPT

# Training config
batch_size = 12          # Micro batch size
block_size = 1024        # Context length
gradient_accumulation_steps = 5 * 8  # ~60K tokens per batch

# Model config
n_layer = 12
n_head = 12
n_embd = 768
dropout = 0.0

# Optimizer config
learning_rate = 6e-4
max_iters = 600000
weight_decay = 1e-1
beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0

# Learning rate schedule
warmup_iters = 2000
lr_decay_iters = 600000
min_lr = 6e-5

# System
device = 'cuda'
dtype = 'bfloat16' if torch.cuda.is_bf16_supported() else 'float16'
compile = True  # PyTorch 2.0

# Data loader
def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+1+block_size] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

# Learning rate schedule
def get_lr(it):
    # Warmup
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    # Decay to min_lr
    if it > lr_decay_iters:
        return min_lr
    # Cosine decay
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

# Init model
model = GPT(GPTConfig())
model.to(device)

# Compile model (PyTorch 2.0)
if compile:
    print("Compiling model...")
    model = torch.compile(model)

# Optimizer
optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device)

# Training loop
for iter_num in range(max_iters):
    # Set learning rate
    lr = get_lr(iter_num)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # Gradient accumulation
    for micro_step in range(gradient_accumulation_steps):
        X, Y = get_batch('train')
        with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
            logits, loss = model(X, Y)
            loss = loss / gradient_accumulation_steps
        loss.backward()

    # Clip gradients
    if grad_clip != 0.0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)

    # Update weights
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

    # Logging
    if iter_num % 100 == 0:
        print(f"iter {iter_num}: loss {loss.item():.4f}, lr {lr:.2e}")

Data Preparation

Shakespeare Character-Level

# Step 1: Download Shakespeare
cd data/shakespeare_char
python prepare.py

# Creates:
# - train.bin (90% of data, ~1MB)
# - val.bin (10% of data, ~110KB)
# - meta.pkl (vocab info)

prepare.py:

import os
import pickle
import requests
import numpy as np

# Download
input_file = 'input.txt'
if not os.path.exists(input_file):
    url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
    with open(input_file, 'w') as f:
        f.write(requests.get(url).text)

# Read and process
with open(input_file, 'r') as f:
    data = f.read()

print(f"Length: {len(data):,} characters")

# Create vocabulary
chars = sorted(list(set(data)))
vocab_size = len(chars)
print(f"Vocab size: {vocab_size}")

# Create mappings
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}

# Encode dataset
data_ids = [stoi[c] for c in data]

# Train/val split
n = len(data_ids)
train_ids = data_ids[:int(n*0.9)]
val_ids = data_ids[int(n*0.9):]

# Save as numpy arrays
train_ids = np.array(train_ids, dtype=np.uint16)
val_ids = np.array(val_ids, dtype=np.uint16)
train_ids.tofile('train.bin')
val_ids.tofile('val.bin')

# Save metadata
meta = {'vocab_size': vocab_size, 'itos': itos, 'stoi': stoi}
with open('meta.pkl', 'wb') as f:
    pickle.dump(meta, f)

OpenWebText (GPT-2 Reproduction)

# Step 1: Download OpenWebText (~12GB compressed)
cd data/openwebtext
python prepare.py

# Warning: Takes 1-2 hours, creates ~54GB of tokenized data

prepare.py:

import os
import numpy as np
import tiktoken
from datasets import load_dataset

# Download dataset
dataset = load_dataset("openwebtext", num_proc=8)

# Use GPT-2 tokenizer
enc = tiktoken.get_encoding("gpt2")

def tokenize(example):
    ids = enc.encode_ordinary(example['text'])
    ids.append(enc.eot_token)  # Add <|endoftext|>
    return {'ids': ids, 'len': len(ids)}

# Tokenize (parallel)
tokenized = dataset.map(
    tokenize,
    remove_columns=['text'],
    desc="Tokenizing",
    num_proc=8
)

# Concatenate all tokens
train_ids = np.concatenate([np.array(x['ids'], dtype=np.uint16) for x in tokenized['train']])
print(f"Train tokens: {len(train_ids):,}")  # ~9B tokens

# Save
train_ids.tofile('train.bin')

# Validation set (sample)
val_ids = np.concatenate([np.array(x['ids'], dtype=np.uint16) for x in tokenized['train'][:5000]])
val_ids.tofile('val.bin')

# Save metadata
meta = {'vocab_size': enc.n_vocab, 'eot_token': enc.eot_token}
with open('meta.pkl', 'wb') as f:
    pickle.dump(meta, f)

Learning Rate Schedules

Cosine Decay with Warmup (GPT-2 style)

def get_lr(it):
    # 1) Linear warmup
    if it < warmup_iters:
        return learning_rate * it / warmup_iters

    # 2) Constant at min_lr after decay
    if it > lr_decay_iters:
        return min_lr

    # 3) Cosine decay in between
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

# Example values
learning_rate = 6e-4  # Peak LR
min_lr = 6e-5         # Final LR (10% of peak)
warmup_iters = 2000   # Warmup steps
lr_decay_iters = 600000  # Total training steps

Visualization:

LR
^
|     Peak (6e-4)
|    /‾‾‾‾‾‾‾‾‾‾\
|   /            \
|  /              \_____ Min (6e-5)
| /
|/________________> Iteration
  Warmup  Cosine    Const
  (2K)    (598K)

Constant LR with Warmup (Simple)

def get_lr(it):
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    return learning_rate

# Good for small experiments

Gradient Accumulation

Effective batch size = batch_size × gradient_accumulation_steps × num_gpus

# Config
batch_size = 12  # Per-GPU micro batch
gradient_accumulation_steps = 40  # Accumulate gradients
# Effective batch: 12 × 40 = 480 sequences = ~0.5M tokens

# Training loop
optimizer.zero_grad()
for micro_step in range(gradient_accumulation_steps):
    X, Y = get_batch('train')
    logits, loss = model(X, Y)
    loss = loss / gradient_accumulation_steps  # Scale loss
    loss.backward()  # Accumulate gradients

# Update once
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
optimizer.step()

Why?

• Simulates large batch size without OOM
• GPT-2 (124M) uses effective batch ~0.5M tokens
• More stable training

Mixed Precision Training

BF16 (Best for A100/H100)

# Enable bfloat16
dtype = torch.bfloat16

# Training loop
for iter in range(max_iters):
    X, Y = get_batch('train')

    # Forward in BF16
    with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
        logits, loss = model(X, Y)

    # Backward in FP32 (automatic)
    loss.backward()
    optimizer.step()

Advantages:

• No gradient scaler needed
• Same dynamic range as FP32
• 2× faster, 50% memory reduction

FP16 (V100, older GPUs)

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

for iter in range(max_iters):
    X, Y = get_batch('train')

    # Forward in FP16
    with autocast():
        logits, loss = model(X, Y)

    # Scale loss, backward
    scaler.scale(loss).backward()

    # Unscale, clip gradients
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)

    # Update weights
    scaler.step(optimizer)
    scaler.update()

Distributed Data Parallel (DDP)

Single Node, Multiple GPUs

# train.py (DDP version)
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# Initialize
dist.init_process_group(backend='nccl')
ddp_rank = int(os.environ['RANK'])
ddp_local_rank = int(os.environ['LOCAL_RANK'])
ddp_world_size = int(os.environ['WORLD_SIZE'])
device = f'cuda:{ddp_local_rank}'
torch.cuda.set_device(device)

# Model
model = GPT(GPTConfig())
model.to(device)
model = DDP(model, device_ids=[ddp_local_rank])

# Training loop (same as before, DDP handles gradient sync)
for iter in range(max_iters):
    X, Y = get_batch('train')  # Each rank gets different data
    logits, loss = model(X, Y)
    loss.backward()  # DDP syncs gradients across GPUs
    optimizer.step()

Launch:

# 8 GPUs on single node
torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py

Multi-Node Training

# Node 0 (master)
torchrun --nproc_per_node=8 \
  --nnodes=4 --node_rank=0 \
  --master_addr=192.168.1.100 --master_port=29500 \
  train.py config/train_gpt2.py

# Node 1-3 (workers)
torchrun --nproc_per_node=8 \
  --nnodes=4 --node_rank=$RANK \
  --master_addr=192.168.1.100 --master_port=29500 \
  train.py config/train_gpt2.py

Checkpointing

Save Checkpoint

# Save every N iterations
if iter_num % 5000 == 0:
    checkpoint = {
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'model_args': model_args,
        'iter_num': iter_num,
        'best_val_loss': best_val_loss,
        'config': config,
    }
    torch.save(checkpoint, os.path.join(out_dir, f'ckpt_{iter_num}.pt'))

Resume from Checkpoint

# Load checkpoint
init_from = 'resume'  # or 'gpt2', 'gpt2-medium', etc.

if init_from == 'resume':
    ckpt_path = os.path.join(out_dir, 'ckpt_latest.pt')
    checkpoint = torch.load(ckpt_path, map_location=device)

    # Restore model
    model_args = checkpoint['model_args']
    model = GPT(GPTConfig(**model_args))
    model.load_state_dict(checkpoint['model'])

    # Restore optimizer
    optimizer.load_state_dict(checkpoint['optimizer'])

    # Restore iteration counter
    iter_num = checkpoint['iter_num']
    best_val_loss = checkpoint['best_val_loss']

Fine-Tuning Pretrained Models

Load OpenAI GPT-2 Weights

# model.py - from_pretrained method
@classmethod
def from_pretrained(cls, model_type):
    """Load pretrained GPT-2 model weights from HuggingFace."""
    from transformers import GPT2LMHeadModel

    # Download from HuggingFace
    model_hf = GPT2LMHeadModel.from_pretrained(model_type)
    sd_hf = model_hf.state_dict()

    # Filter out keys we don't need
    sd_hf_keys = [k for k in sd_hf.keys() if not k.endswith('.attn.masked_bias')]
    sd_hf_keys = [k for k in sd_hf_keys if not k.endswith('.attn.bias')]

    # Create our model
    config = GPTConfig.from_model_type(model_type)
    model = GPT(config)
    sd = model.state_dict()

    # Copy weights (transpose Conv1D → Linear)
    for k in sd_hf_keys:
        if any([k.endswith(w) for w in ['.c_attn.weight', '.c_proj.weight', '.c_fc.weight']]):
            sd[k] = sd_hf[k].t()  # Transpose
        else:
            sd[k] = sd_hf[k]  # Direct copy

    model.load_state_dict(sd)
    return model

# Usage
model = GPT.from_pretrained('gpt2')  # Load GPT-2 (124M)

Fine-Tune on Custom Data

# config/finetune_shakespeare.py
init_from = 'gpt2'  # Start from GPT-2
dataset = 'shakespeare_char'

# Fine-tuning hyperparameters
learning_rate = 3e-5  # Lower LR for fine-tuning
max_iters = 2000      # Short fine-tuning
warmup_iters = 100

# Regularization
weight_decay = 1e-1
dropout = 0.2  # Add dropout

# Run
# python train.py config/finetune_shakespeare.py

Evaluation

Perplexity

@torch.no_grad()
def estimate_loss():
    model.eval()
    losses = torch.zeros(eval_iters)

    for k in range(eval_iters):
        X, Y = get_batch('val')
        logits, loss = model(X, Y)
        losses[k] = loss.item()

    model.train()
    return losses.mean()

# Usage
val_loss = estimate_loss()
perplexity = math.exp(val_loss)
print(f"Val perplexity: {perplexity:.2f}")

Sample Generation

# sample.py
model.eval()

start = "ROMEO:"  # Prompt
start_ids = encode(start)
x = torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...]

# Generate
with torch.no_grad():
    y = model.generate(x, max_new_tokens=500, temperature=0.8, top_k=200)

print(decode(y[0].tolist()))

Training Times

Resources

• Training script: https://github.com/karpathy/nanoGPT/blob/master/train.py
• Configs: https://github.com/karpathy/nanoGPT/tree/master/config
• Video walkthrough: "Let's build GPT" (training section)
• GPT-2 paper: https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf