Jesse A. Green / notebook

nanoGPT Administrator & Developer Guide

Generated April 2026 from karpathy/nanoGPT at 3adf61e

AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it.

This guide is for ML researchers, students, hobbyists, and operators who run nanoGPT in anger — multi-GPU pretraining, custom dataset pipelines, performance tuning, and deeper hacking on model.py and train.py. It documents every setting in every config file, every public method on the GPT class, the math of the transformer block, and the operational realities of training a GPT at this scale.

If you're looking for "how do I run this thing for the first time," start with the User Guide. This document picks up where the User Guide ends and goes deep on internals.

1. Architecture Overview

nanoGPT is a deliberately small, single-file research codebase. There are no abstractions, no plugin systems, no DI framework. Every behavior is one or two if statements in train.py away from the call site. That is the design intent.

Component inventory

Component Path Role
Training entry point train.py Single while-loop training, DDP detection, eval, checkpoint, wandb.
Sampling entry point sample.py Loads a checkpoint or HF GPT-2 weights, runs model.generate.
Benchmark entry point bench.py Stripped-down train.py for MFU and torch.profiler traces.
Configurator configurator.py Reads sys.argv, exec()s config files, applies --key=value overrides into globals.
Model definition model.py LayerNorm, CausalSelfAttention, MLP, Block, GPTConfig, GPT.
Training configs config/train_gpt2.py, config/train_shakespeare_char.py Reference pretraining recipes.
Finetune config config/finetune_shakespeare.py Constant-LR finetune from gpt2-xl.
Eval configs config/eval_gpt2{,_medium,_large,_xl}.py Eval-only runs against the four pretrained GPT-2 sizes.
Shakespeare-char data prep data/shakespeare_char/prepare.py Char-level encoding (vocab=65), 1.1M chars.
Shakespeare BPE data prep data/shakespeare/prepare.py tiktoken GPT-2 BPE, 301,966 train tokens.
OpenWebText data prep data/openwebtext/prepare.py HuggingFace datasets + tiktoken, ~9B train tokens.
Notebooks scaling_laws.ipynb, transformer_sizing.ipynb Karpathy's exploratory derivations of model-flop and Chinchilla relationships.

Technology stack

Layer Choice
Language Python 3.8+
Tensor framework PyTorch (≥2.0 recommended for Flash Attention)
Distributed training torch.distributed + torch.nn.parallel.DistributedDataParallel, launched via torchrun
Mixed precision torch.amp.autocast (bfloat16 or float16) + torch.cuda.amp.GradScaler for fp16
Optimizer torch.optim.AdamW (with the fused kernel when on CUDA)
Data loading numpy.memmap over train.bin / val.bin files
Tokenizers tiktoken (GPT-2 BPE) for Shakespeare and OpenWebText; in-line char-level for Shakespeare-char
Pretrained weight import HuggingFace transformers.GPT2LMHeadModel
Compilation torch.compile(model) (PyTorch 2.0+)
Logging Plain print to stdout; optional wandb

Training loop data flow

                                      data/{dataset}/{train,val}.bin
                                      (uint16 token IDs, np.memmap)
                                                  |
                                                  v
                                       +----------------------+
                                       |  get_batch(split)    |
                                       |  random offsets      |
                                       |  pin_memory + non-   |
                                       |  blocking copy to    |
                                       |  device              |
                                       +----------+-----------+
                                                  |
                                                  v
                                       +----------------------+
                                       |  GPT.forward(idx,    |
                                       |  targets)            |
                                       |   - token + pos emb  |
                                       |   - N x Block        |
                                       |   - LayerNorm        |
                                       |   - lm_head          |
                                       |   - F.cross_entropy  |
                                       +----------+-----------+
                                                  |
                                                  v
                                       +----------------------+
                                       |  loss / grad_accum   |
                                       |  scaler.scale(loss)  |
                                       |    .backward()       |
                                       +----------+-----------+
                                                  |
                          repeat micro-step       v
                  +-------------------+   gradient_accumulation_steps
                  |                   |   times
                  +-------+-----------+
                          v
                  +----------------------+
                  |  scaler.unscale_     |
                  |  clip_grad_norm_     |
                  |  scaler.step(optim)  |
                  |  scaler.update()     |
                  |  optim.zero_grad()   |
                  +----------+-----------+
                             |
                             v
            +---------------------------------+
            |  every eval_interval iters:     |
            |    estimate_loss()              |
            |    if val improved (or always): |
            |      torch.save(checkpoint,     |
            |        out_dir/ckpt.pt)         |
            +----------+----------------------+
                       |
                       v
            wandb.log(if wandb_log)
                       |
                       v
            iter_num++ until > max_iters

Note

The data loader recreates np.memmap on every get_batch call. The comment in train.py cites a known memory leak with reusing the same memmap across thousands of batches; rebuilding it per call costs nothing measurable and avoids the leak. Don't "optimize" this away.

GPT block architecture

                       idx (b, t)        targets (b, t) [optional]
                            |
                            v
                  +-----------------+
                  |  wte (vocab x   |  position ids 0..t-1
                  |     n_embd)     |        |
                  |  + wpe (block_  |        |
                  |     size x      |        |
                  |     n_embd)     |        |
                  |  + Dropout      |
                  +--------+--------+
                           |
                           v
                  +-----------------+
                  |  Block 1        |
                  |   x = x +       |
                  |     Attn(LN(x)) |
                  |   x = x +       |
                  |     MLP(LN(x))  |
                  +--------+--------+
                           |
                          ...
                           |
                  +--------+--------+
                  |  Block n_layer  |
                  +--------+--------+
                           |
                           v
                  +-----------------+
                  |  ln_f (LayerNorm)
                  +--------+--------+
                           |
                           v
                  +-----------------+
                  |  lm_head        |
                  |  (n_embd x      |
                  |  vocab_size,    |
                  |  no bias,       |
                  |  weight-tied to |
                  |  wte)           |
                  +--------+--------+
                           |
                  if targets is not None:
                      logits over all positions
                      F.cross_entropy(logits, targets, ignore_index=-1)
                  else (inference):
                      logits = lm_head(x[:, [-1], :])  # last position only
                  return logits, loss

Note

The wte weights and lm_head weights are tied (self.transformer.wte.weight = self.lm_head.weight). One of the most non-obvious tricks in this file. It nearly halves the parameter count of the embedding+head pair and consistently improves loss on small models.

2. System Requirements

Hardware tiers

Tier Hardware Suitable for
Hobbyist CPU 4+ core CPU, 8 GB RAM Shakespeare-char with size overrides
Hobbyist GPU Single consumer NVIDIA GPU (8-24 GB), bfloat16 capable Shakespeare-char, GPT-2 finetunes, small custom datasets
Workstation RTX 4090 / RTX 6000 Ada Custom datasets up to a few hundred million params
Single node training 8 × A100 40GB Reference: GPT-2 124M in ~5 days
Multi-node 2+ nodes × 8 × A100/H100 Larger models, faster wall-clock

Software requirements

Component Version Notes
Python 3.8+ (3.10+ recommended) nanoGPT uses dataclasses + f-strings + inspect.signature introspection.
PyTorch 2.0+ Required for scaled_dot_product_attention (Flash Attention path) and torch.compile. Older versions work via the manual attention fallback.
CUDA Matched to your PyTorch wheel If you want GPU. CPU and MPS work as well.
NCCL Bundled with PyTorch CUDA wheels DDP transport. Set NCCL_IB_DISABLE=1 if no Infiniband.
numpy any recent np.memmap and np.uint16 array IO.
transformers any recent Only used by GPT.from_pretrained.
datasets <2.20 recommended data/openwebtext/prepare.py calls load_dataset("openwebtext"). Newer versions broke this (issue #633).
tiktoken any GPT-2 BPE encoder.
wandb any (optional) Only required when wandb_log=True.
tqdm any Progress bar in data/openwebtext/prepare.py.
requests any Used by Shakespeare prepare scripts to download tinyshakespeare/input.txt.

Disk layout requirements

nanoGPT/
  data/
    {dataset_name}/
      train.bin       # uint16 token IDs, mmap'd by train.py
      val.bin         # uint16 token IDs, mmap'd by train.py
      meta.pkl        # optional, dict with vocab_size, stoi, itos
  out/                # default out_dir
    ckpt.pt           # checkpoint (overwritten on each save)

Important

train.py always writes its checkpoint to {out_dir}/ckpt.pt and always overwrites the file. Only one checkpoint per out_dir. If you want milestone snapshots, copy ckpt.pt somewhere else between eval intervals from a wrapper script.

3. Configuration Reference

train.py — complete config table

Every setting below is a module-level global declared in train.py before the call to configurator.py. The configurator only accepts keys that exist in this set.

I/O

Key Type Default Description
out_dir str 'out' Checkpoint and log directory. Created on master if missing.
eval_interval int 2000 Iters between train/val eval + checkpoint write.
log_interval int 1 Iters between per-step print of loss / time / mfu.
eval_iters int 200 Number of batches averaged in estimate_loss.
eval_only bool False If True, run a single eval and exit (used by config/eval_gpt2*.py).
always_save_checkpoint bool True If False, only save when losses['val'] < best_val_loss.
init_from str 'scratch' One of 'scratch', 'resume', 'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'.

Wandb logging

Key Type Default Description
wandb_log bool False Master switch for wandb.
wandb_project str 'owt' Wandb project.
wandb_run_name str 'gpt2' Wandb run name.

Data

Key Type Default Description
dataset str 'openwebtext' Folder under data/. The script reads data/{dataset}/train.bin and val.bin.
gradient_accumulation_steps int 40 (5 × 8) Number of micro-batches accumulated per optimizer step. Divided by WORLD_SIZE under DDP.
batch_size int 12 Micro-batch size. Effective batch is batch_size × gradient_accumulation_steps × WORLD_SIZE.
block_size int 1024 Context window in tokens. Hard upper bound at the position embedding (wpe) size.

Model

Key Type Default Description
n_layer int 12 Number of transformer Blocks.
n_head int 12 Number of attention heads. Must divide n_embd.
n_embd int 768 Embedding / hidden dim. Must be divisible by n_head.
dropout float 0.0 Dropout in attention residual, MLP residual, and embedding.
bias bool False If True, nn.Linear and LayerNorm use bias. GPT-2 uses bias; setting False is "a bit better and faster" per Karpathy's comment.

Optimizer

Key Type Default Description
learning_rate float 6e-4 Peak LR after warmup.
max_iters int 600000 Total iters; the loop terminates when iter_num > max_iters.
weight_decay float 1e-1 AdamW weight decay applied only to 2D tensors.
beta1 float 0.9 AdamW first moment beta.
beta2 float 0.95 AdamW second moment beta.
grad_clip float 1.0 Global gradient norm clip. 0.0 disables.

LR schedule

Key Type Default Description
decay_lr bool True Toggles the cosine schedule. If False, LR is held at learning_rate.
warmup_iters int 2000 Linear warmup iters from 0 → learning_rate.
lr_decay_iters int 600000 Iter at which cosine reaches min_lr. After this, get_lr returns min_lr.
min_lr float 6e-5 Floor LR (Chinchilla recommends learning_rate / 10).

DDP

Key Type Default Description
backend str 'nccl' torch.distributed backend. Use 'gloo' on CPU or for cross-platform.

System

Key Type Default Description
device str 'cuda' 'cuda', 'cpu', 'cuda:0', 'mps', etc. Under DDP this is overwritten with cuda:{LOCAL_RANK}.
dtype str 'bfloat16' if available else 'float16' Autocast dtype. Triggers GradScaler only when 'float16'.
compile bool True If True, calls torch.compile(model) (PyTorch 2.0+).

sample.py — complete config table

Key Type Default Description
init_from str 'resume' 'resume' loads out_dir/ckpt.pt. 'gpt2*' calls GPT.from_pretrained.
out_dir str 'out' Used only if init_from='resume'.
start str '\n' Prompt. "FILE:path" reads the prompt from disk.
num_samples int 10 Number of completions to generate.
max_new_tokens int 500 Tokens generated per sample.
temperature float 0.8 Sampling temperature.
top_k int or None 200 Truncate to top-k logits before softmax. None disables.
seed int 1337 RNG seed.
device str 'cuda' 'cuda', 'cpu', 'mps'.
dtype str autodetect Same semantics as train.py.
compile bool False torch.compile for inference; defaults to off because compile overhead rarely pays off on a single sampling pass.

bench.py — complete config table

Key Type Default Description
batch_size int 12 Same as train.py.
block_size int 1024 Same as train.py.
bias bool False Same as train.py.
real_data bool True If False, replaces the data loader with random int tensors.
seed int 1337 RNG seed.
device str 'cuda' Same as train.py.
dtype str autodetect Same as train.py.
compile bool True Same as train.py.
profile bool False Enables torch.profiler writing TensorBoard traces under ./bench_log.

Reference configurations shipped in config/

config/train_gpt2.py — GPT-2 124M reference run

wandb_log = True
wandb_project = 'owt'
wandb_run_name = 'gpt2-124M'
batch_size = 12
block_size = 1024
gradient_accumulation_steps = 5 * 8     # 12 * 1024 * 40 = 491,520 tokens / step on 8 GPUs
max_iters = 600000
lr_decay_iters = 600000                  # 300B tokens total
eval_interval = 1000
eval_iters = 200
log_interval = 10
weight_decay = 1e-1

Important

config/train_gpt2.py is tuned for 8 × A100 40GB. The accompanying comment says "wait ~5 days". On fewer GPUs, raise gradient_accumulation_steps to keep tokens-per-step constant, or scale down the model and learning rate.

config/train_shakespeare_char.py — baby model

out_dir = 'out-shakespeare-char'
eval_interval = 250
eval_iters = 200
log_interval = 10
always_save_checkpoint = False
wandb_log = False
wandb_project = 'shakespeare-char'
wandb_run_name = 'mini-gpt'

dataset = 'shakespeare_char'
gradient_accumulation_steps = 1
batch_size = 64
block_size = 256

n_layer = 6
n_head = 6
n_embd = 384
dropout = 0.2

learning_rate = 1e-3
max_iters = 5000
lr_decay_iters = 5000
min_lr = 1e-4
beta2 = 0.99
warmup_iters = 100

config/finetune_shakespeare.py — constant-LR finetune

out_dir = 'out-shakespeare'
eval_interval = 5
eval_iters = 40
wandb_log = False
wandb_project = 'shakespeare'
wandb_run_name = 'ft-' + str(time.time())

dataset = 'shakespeare'
init_from = 'gpt2-xl'
always_save_checkpoint = False

batch_size = 1
gradient_accumulation_steps = 32
max_iters = 20

learning_rate = 3e-5
decay_lr = False

config/eval_gpt2{,_medium,_large,_xl}.py — eval-only

Config init_from n_layer n_head n_embd Params
eval_gpt2.py 'gpt2' 12 12 768 124M
eval_gpt2_medium.py 'gpt2-medium' 24 16 1024 350M
eval_gpt2_large.py 'gpt2-large' 36 20 1280 774M
eval_gpt2_xl.py 'gpt2-xl' 48 25 1600 1558M

All four eval configs set batch_size=8, eval_iters=500, eval_only=True, wandb_log=False. They run a single eval pass and exit.

Environment variables

Variable Set by Read by Purpose
RANK torchrun train.py:82 Process global rank in DDP. -1 (default) means single-process.
LOCAL_RANK torchrun train.py Local rank within a node. Used to pick cuda:{LOCAL_RANK}.
WORLD_SIZE torchrun train.py Total number of DDP processes. Divides gradient_accumulation_steps.
NCCL_IB_DISABLE user NCCL Set to 1 on clusters without Infiniband.
NCCL_DEBUG user NCCL Set to INFO for verbose NCCL diagnostics.
HF_HOME user HuggingFace Override the cache root for HuggingFace datasets and transformers.

4. Deployment & Scaling

Single-process, single-GPU

The default. Run python train.py [config_file] [--overrides]. No environment setup needed beyond pip install.

Single-node, multi-GPU (DDP)

torchrun --standalone --nproc_per_node=N train.py [config] [--overrides]

--standalone chooses a free local port and skips --master_addr/ --master_port. --nproc_per_node=N spawns N processes, one per GPU. Inside train.py:

  • RANK, LOCAL_RANK, WORLD_SIZE are read from env.
  • device is set to cuda:{LOCAL_RANK}.
  • master_process = (ddp_rank == 0). Only the master process logs, evals, and writes checkpoints.
  • gradient_accumulation_steps //= WORLD_SIZE so the effective batch size is held constant.
  • seed_offset = ddp_rank so each rank sees a different RNG sequence.

Multi-node DDP

Pick a master node IP and port; both must be reachable across all worker nodes. Run torchrun on each node with matching --master_addr and --master_port, distinct --node_rank, and the same --nproc_per_node and --nnodes.

# master (rank 0)
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
  --master_addr=10.0.0.1 --master_port=29500 train.py config/train_gpt2.py

# worker (rank 1)
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 \
  --master_addr=10.0.0.1 --master_port=29500 train.py config/train_gpt2.py

Tip

Use screen or tmux on each node — long pretraining runs survive easier when the launching shell is detached. The config/train_gpt2.py header literally suggests "launch as the following (e.g. in a screen session) and wait ~5 days".

Effective batch size math

effective_batch_size_in_tokens =
    batch_size *
    block_size *
    gradient_accumulation_steps *
    WORLD_SIZE

The default config/train_gpt2.py produces: 12 * 1024 * 40 = 491,520 tokens per optimizer step on 8 GPUs (post-DDP it's 12 * 1024 * 5 per GPU = 61,440 tokens per micro-step).

Scaling down to fewer GPUs

Drop GPUs by raising gradient_accumulation_steps proportionally:

GPUs gradient_accumulation_steps Per-GPU micro-step tokens
8 40 (default in train.py) 5 micro-steps × 12 batch × 1024 = 61,440
4 80 unchanged effective size
2 160 unchanged effective size
1 320 unchanged effective size

You can also lower batch_size and block_size if memory is tight. The exact tradeoff is best found by running bench.py on the target hardware.

CPU / MPS deployments

For tiny demos, --device=cpu works. The Shakespeare-char config's trailing comment lists exactly the overrides needed:

# on macbook also add
# device = 'cpu'  # run on cpu only
# compile = False # do not torch compile the model

5. Operations

Day-to-day operations checklist

  • [x] Confirm data/{dataset}/{train.bin,val.bin} exist and are uint16.
  • [x] Pick or create a config under config/.
  • [x] Decide on out_dir (one checkpoint per directory, overwritten).
  • [x] Verify device, dtype, and compile choices match your hardware.
  • [x] If using DDP, confirm torchrun reachability and Infiniband settings.
  • [x] If using wandb, run wandb login once on the master node.
  • [x] Launch under screen / tmux for any run >30 minutes.

Monitoring an in-flight run

train.py prints two kinds of lines:

iter 1234: loss 2.7891, time 23.65ms, mfu 41.20%
step 2000: train loss 2.8123, val loss 2.9456

The first form fires every log_interval iters. The second fires every eval_interval iters and triggers checkpointing.

Note

mfu (Model FLOPs Utilization) prints as a running EMA after the first 5 iterations. It is computed in model.estimate_mfu against an A100 bfloat16 peak of 312 TFLOPS. On non-A100 hardware the reported number is still useful for relative comparisons but the absolute scale is offset.

Checkpoint contents

{
    'model':         raw_model.state_dict(),
    'optimizer':     optimizer.state_dict(),
    'model_args':    {n_layer, n_head, n_embd, block_size, bias, vocab_size, dropout},
    'iter_num':      int,
    'best_val_loss': float,
    'config':        {...all of train.py's config_keys...},
}

Saved via torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt')). The config field is a snapshot of every setting; reading it after the fact gives you the exact recipe used.

Tip

Add a wrapper like cp out/ckpt.pt out/ckpt-iter{N}.pt after each eval if you need rolling snapshots. The base train.py overwrites a single ckpt.pt to keep disk usage flat.

Resuming after a crash

Same out_dir, add --init_from=resume. The script restores iter_num, best_val_loss, optimizer state, and forces architecture args to match the checkpoint. Other knobs (dropout, batch_size, learning_rate, etc.) remain whatever your config / CLI specifies.

Common runbook recipes

Task Recipe
"I want to log to wandb mid-run" Stop, restart with --wandb_log=True --init_from=resume.
"I want to lower LR mid-run" Stop, restart with --learning_rate=3e-4 --init_from=resume. Cosine schedule will adapt.
"I want a milestone snapshot" Out-of-band cp out/ckpt.pt out/ckpt-snapshot.pt between eval intervals.
"I want to abort cleanly under DDP" Send SIGINT to all torchrun processes; they'll call destroy_process_group on next iteration check.
"I want to evaluate without training" --eval_only=True --init_from=resume. Loads the checkpoint, runs one estimate_loss, exits.

Tip

eval_only=True is exactly what config/eval_gpt2*.py does. Repurpose it for any "what's my val loss right now" question.

6. Performance

MFU (Model FLOPs Utilization)

model.estimate_mfu(fwdbwd_per_iter, dt) implements PaLM Appendix B's analytic FLOPs model:

N = self.get_num_params()
L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
flops_per_token   = 6*N + 12*L*H*Q*T
flops_per_fwdbwd  = flops_per_token * T
flops_per_iter    = flops_per_fwdbwd * fwdbwd_per_iter
flops_achieved    = flops_per_iter / dt
flops_promised    = 312e12              # A100 bfloat16 peak
mfu               = flops_achieved / flops_promised

The 6*N term is the standard "forward+backward of a parameter is roughly 6 FLOPs per token" approximation. The 12*L*H*Q*T term covers the attention QK^T and softmax-V multiplies that scale with sequence length.

Hyperparameter relationships

Variable Rule of thumb
batch_size × gradient_accumulation_steps × WORLD_SIZE × block_size Aim for ~0.5M tokens/step at GPT-2 124M scale (matches config/train_gpt2.py).
learning_rate Scale roughly linearly with batch size up to ~2-4M tokens/step.
min_lr Chinchilla heuristic: learning_rate / 10.
lr_decay_iters Equal to max_iters for cosine to land exactly at min_lr at the end.
warmup_iters ~0.3-1% of max_iters. The default 2000/600000 = 0.33%.
weight_decay 1e-1 for pretraining (default), 1e-2 for bench.py smoke testing.
dropout 0.0 for pretraining, 0.1+ for finetuning on small corpora.
n_head Must divide n_embd. Each head dim is n_embd / n_head.

Block-size considerations

Doubling block_size quadruples attention compute and memory. The Flash Attention path mitigates the memory cost (O(N) instead of O(N²)) but the compute cost is still quadratic in block_size. If you can't afford 1024 context, config/train_shakespeare_char.py shows a 256-block-size recipe that still produces a passable model.

compile=True tradeoffs

When compile=True is faster When compile=True is slower
Long pretraining runs (>1000 iters): 10-30% throughput improvement Short experiments where the ~1-minute compile dominates wall-clock
Stable model size and shapes Models with shape-dependent control flow (none in nanoGPT)
PyTorch 2.0+ on Ampere/Hopper/Ada PyTorch <2.0 (compile is a no-op there) or unstable MPS

Note

compile=True warning: "Compiling model... (takes a ~minute)". The first few iterations after compile are slower because the compiler is still warming up traces. The MFU readout in the User Guide example (mfu -100.00% then mfu 5.46% after warmup) is normal.

bfloat16 vs float16

Dtype Pros Cons
bfloat16 Same exponent range as fp32, no GradScaler, simpler training. Requires Ampere or newer (A100/A6000/H100/RTX 4090).
float16 Works on every CUDA-capable GPU. Narrow exponent range; relies on GradScaler to dynamically rescale; can NaN.

The autodetect at the top of train.py/sample.py/bench.py always prefers bfloat16 when available. Override with --dtype=... at your own risk on hardware that doesn't natively support it.

vocab_size=50304 padding

50,257 is the actual GPT-2 BPE vocab. The default train.py rounds up to 50304 (50304 % 64 == 0) because divisibility-by-64 unlocks faster matmul kernels on NVIDIA tensor cores. The 47 unused embedding rows cost a tiny amount of memory but pay for themselves in throughput. This is one of the small but high-leverage decisions throughout the codebase.

Fused AdamW

fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
use_fused = fused_available and device_type == 'cuda'
extra_args = dict(fused=True) if use_fused else dict()
optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
print(f"using fused AdamW: {use_fused}")

configure_optimizers introspects torch.optim.AdamW's signature to detect fused-kernel support and uses it on CUDA. Free perf, no config knob exposed.

gradient_accumulation under DDP

When DDP is detected, train.py divides gradient_accumulation_steps by WORLD_SIZE so the effective batch is invariant. It also disables gradient sync on every micro-step except the last:

model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)

This is a manual replacement for model.no_sync() that matches PyTorch's behavior without the context-manager bloat. It cuts DDP all-reduce calls by gradient_accumulation_steps× — a non-trivial speedup on multi-GPU runs.

7. Security

nanoGPT is a research repo. There is no network endpoint, no auth layer, no user-supplied content arrives over the wire at runtime. Security considerations are operational rather than application-layer:

Dataset preparation scripts execute external content

data/openwebtext/prepare.py calls datasets.load_dataset("openwebtext"), which (depending on datasets version) can exec() a Python loader script fetched from HuggingFace. data/shakespeare/prepare.py and data/shakespeare_char/prepare.py call requests.get against https://raw.githubusercontent.com/karpathy/char-rnn/... and write the response to disk.

Warning

Treat dataset preparation scripts as code from third parties. Inspect any prepare.py you didn't write before running it on a machine with access to credentials, source code, or other sensitive data. Consider a dedicated VM or container for first-time runs.

Model checkpoint loading uses torch.load (pickle)

train.py and sample.py both call torch.load(ckpt_path, map_location=...), which uses Python pickle under the hood. Loading a checkpoint from an untrusted source can execute arbitrary code.

Caution

Never load a .pt file from an unverified origin. If you must, use weights_only=True (PyTorch 1.13+) on torch.load. The base nanoGPT code does not pass that flag — patch it locally if you handle external checkpoints.

GPT.from_pretrained downloads from HuggingFace

from_pretrained triggers a HuggingFace download. The HuggingFace cache is trusted by default. If you operate in a regulated environment, mirror the GPT-2 weights internally and set HF_HOME to your mirror.

Wandb credentials

wandb login writes a token to ~/.netrc. On shared boxes use individual user accounts; on CI runners, prefer ephemeral env vars (WANDB_API_KEY=...) over persistent files.

Secrets in config files

Config files are plain Python and exec()'d by the configurator. Do not embed credentials in them and do not commit them with secrets in the config directory. Treat any file under config/ as committed source code.

Multi-tenant clusters

DDP uses torch.distributed over a TCP socket on --master_port. By default this is unauthenticated; rely on cluster-level network isolation (VLAN, security groups, Slurm ResvNet, etc.) rather than the framework.

Issue #717 — sandbox escape in nanochat

nanoGPT's successor project nanochat had an open issue (#717) titled "Critical Sandbox Escape Vulnerability in nanochat" at time of writing. This does not directly affect nanoGPT but is worth knowing if you migrate to that project — see the upstream issue tracker for details.

8. Developer Internals

Repository tree

nanoGPT/
  bench.py                 # micro-benchmark + torch.profiler
  configurator.py          # poor man's configurator (exec-based)
  model.py                 # LayerNorm / CausalSelfAttention / MLP / Block / GPTConfig / GPT
  sample.py                # checkpoint or HF GPT-2 -> generate
  train.py                 # the training loop
  scaling_laws.ipynb       # exploratory: parameters x compute relationships
  transformer_sizing.ipynb # exploratory: GPT-2 sizing math
  config/
    eval_gpt2.py           # eval-only against gpt2 (124M)
    eval_gpt2_medium.py    # eval-only against gpt2-medium (350M)
    eval_gpt2_large.py     # eval-only against gpt2-large (774M)
    eval_gpt2_xl.py        # eval-only against gpt2-xl (1558M)
    finetune_shakespeare.py
    train_gpt2.py          # GPT-2 124M from scratch on OpenWebText (8 x A100)
    train_shakespeare_char.py
  data/
    shakespeare/prepare.py        # tiktoken BPE, 301,966 train tokens
    shakespeare_char/prepare.py   # char-level vocab=65, ~1M tokens
    openwebtext/prepare.py        # tiktoken BPE, ~9B train tokens
  assets/                  # gpt2_124M_loss.png, nanogpt.jpg
  README.md
  LICENSE                  # MIT

model.py internals

LayerNorm (line 18)

class LayerNorm(nn.Module):
    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
    def __init__(self, ndim, bias):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
    def forward(self, input):
        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)

A drop-in for nn.LayerNorm whose only purpose is to make the bias term optional. PyTorch's built-in nn.LayerNorm(elementwise_affine=False) removes both weight and bias; nanoGPT wants weight without bias, hence this shim.

CausalSelfAttention (line 29)

The fused QKV projection: c_attn is nn.Linear(n_embd, 3 * n_embd), output is split into Q, K, V along the last dim, reshaped to (B, n_head, T, head_dim), and fed to one of two paths:

  • Flash path (scaled_dot_product_attention available, PyTorch ≥2.0):
    y = F.scaled_dot_product_attention(q, k, v, attn_mask=None,
        dropout_p=self.dropout if self.training else 0,
        is_causal=True)
    
  • Manual path (PyTorch <2.0): explicit (q @ k.transpose(-2,-1)) scaled by 1/sqrt(d_k), masked by a registered triangular bias buffer, softmaxed, dropped, and matmul'd against V.

Both paths produce the same output to numeric tolerance; the flash path is significantly faster and more memory-efficient. Output goes through a c_proj Linear and a residual dropout.

MLP (line 78)

Standard transformer MLP: n_embd → 4 * n_embd → GELU → n_embd, plus a final dropout. The c_fc and c_proj names match GPT-2's naming convention so weights map cleanly during from_pretrained.

Block (line 94)

Pre-LN residual architecture:

def forward(self, x):
    x = x + self.attn(self.ln_1(x))
    x = x + self.mlp(self.ln_2(x))
    return x

Pre-LN (LayerNorm before each sublayer) is the standard for stable training of large transformers. GPT-2 used pre-LN and so does nanoGPT.

GPTConfig (line 109, dataclass)

@dataclass
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 50304
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0
    bias: bool = True

Note bias=True here — but train.py overrides with bias=False. The class default matches GPT-2; the script default takes the perf-and-quality win of bias=False.

GPT (line 118)

Methods:

Method Lines Role
__init__ 118-148 Builds wte, wpe, drop, the h ModuleList of Blocks, ln_f, lm_head. Ties wte.weight to lm_head.weight. Applies special scaled init 0.02/sqrt(2*n_layer) to c_proj weights (per the GPT-2 paper).
get_num_params 150-156 Returns parameter count. With non_embedding=True (default) subtracts wpe's parameters but keeps wte (because of weight tying, wte.weight is also lm_head.weight).
_init_weights 158-164 Applied via self.apply. Linear weights ~ N(0, 0.02²), Linear biases zeroed, Embedding weights ~ N(0, 0.02²).
forward 170-193 Token + position embedding, dropout, N Blocks, final LayerNorm. If targets given: full lm_head + cross_entropy loss. If not: only forward lm_head on the last position (inference optimization).
crop_block_size 195-200 Model surgery: shrinks wpe (and the registered causal bias if present) to a smaller block_size. Used when loading a 1024-block GPT-2 checkpoint into a smaller-context model.
from_pretrained 207-260 Classmethod. Builds a fresh GPT whose n_layer/n_head/n_embd matches one of gpt2 / gpt2-medium / gpt2-large / gpt2-xl, then copies weights from a HuggingFace GPT2LMHeadModel. Transposes the four Conv1D-style weights (attn.c_attn.weight, attn.c_proj.weight, mlp.c_fc.weight, mlp.c_proj.weight) since GPT-2 used Conv1D and nanoGPT uses Linear.
configure_optimizers 263-286 Splits parameters into a 2D-tensor decay group and a 1D-tensor (biases, LayerNorm params) no-decay group. Detects fused AdamW availability via inspect.signature and uses it on CUDA.
estimate_mfu 289-303 PaLM Appendix B FLOPs estimator. Returns flops_achieved / 312e12.
generate 306-325 @torch.no_grad(). Iterative single-token sampling: crop context to block_size, forward, scale logits by temperature, optionally truncate to top_k, softmax, multinomial sample, append, repeat for max_new_tokens.

train.py internals

The script has three phases:

  1. Config + DDP setup (lines ~626-705): defaults → configurator.py → detect RANK → set up device, device_type, ptdtype, ctx (the autocast contextmanager).
  2. Model + optimizer init (lines ~707-805): build data loader, construct GPT from 'scratch'/'resume'/'gpt2*', call crop_block_size if needed, build optimizer, optionally torch.compile, wrap in DDP.
  3. Training loop (lines ~807-929): cosine LR per iter, eval and checkpoint every eval_interval, gradient accumulation with model.require_backward_grad_sync toggling, GradScaler-aware optimizer step, MFU EMA, terminate at max_iters.

configurator.py internals (full source)

import sys
from ast import literal_eval

for arg in sys.argv[1:]:
    if '=' not in arg:
        # assume it's the name of a config file
        assert not arg.startswith('--')
        config_file = arg
        print(f"Overriding config with {config_file}:")
        with open(config_file) as f:
            print(f.read())
        exec(open(config_file).read())
    else:
        # assume it's a --key=value argument
        assert arg.startswith('--')
        key, val = arg.split('=')
        key = key[2:]
        if key in globals():
            try:
                attempt = literal_eval(val)
            except (SyntaxError, ValueError):
                attempt = val
            assert type(attempt) == type(globals()[key])
            print(f"Overriding: {key} = {attempt}")
            globals()[key] = attempt
        else:
            raise ValueError(f"Unknown config key: {key}")

The whole thing fits on a screen. There's nothing else to it — if you don't like it, fork it. Karpathy's docstring says exactly that: "If someone comes up with a better simple Python solution I am all ears."

Note

Two consequences of this design: (1) the type of an override is inferred via ast.literal_eval and must match the existing global's type — --batch_size=12.0 will fail because the existing global is int. (2) Unknown keys raise instead of silently no-opping, which catches typos.

Testing strategy

There are no unit tests in the repo. Validation is driven by:

  • bench.py for performance regression checks on a known input shape.
  • The Shakespeare-char demo as an end-to-end smoke test: trains in ~10 minutes and produces recognizable Shakespeare on success.
  • config/eval_gpt2*.py for cross-checking against the four official HuggingFace GPT-2 sizes; the expected loss numbers are documented in Karpathy's blog posts and the README.

If you fork, the recommended additions are:

  1. A pytest that runs data/shakespeare_char/prepare.py then ten iterations of train.py and asserts loss drops below an absolute threshold (sanity check the pipeline).
  2. A pytest that loads gpt2 weights via from_pretrained and asserts the parameter count is exactly 124,439,808.

Code patterns and conventions

Pattern Reason
Module-level globals as config The configurator is the simplest way to expose every knob to the CLI without argparse or a settings library.
Data structures live in dataclasses (GPTConfig) or plain dicts (checkpoint) Avoids boilerplate. The model is the only "object" in the file; everything else is data.
print for logs No logging module. The whole codebase fits in your head; structured logging would be overkill.
Inline TODOs and warnings E.g. # TODO investigate near the weight tying. Karpathy's working notes are kept in-place rather than moved to issues.
One file per concept model.py is the model. train.py is training. sample.py is sampling. The repo deliberately avoids cross-cutting infra.

Potential improvements

  • Add weights_only=True to all torch.load calls (PyTorch 1.13+) to reduce pickle deserialization risk.
  • Replace exec(open('configurator.py').read()) with a real import so static analyzers and type checkers don't choke.
  • Add a --top_p flag to sample.py for nucleus sampling.
  • Accept gzipped train.bin/val.bin for faster network transfer of prepared datasets.
  • Add rolling checkpoint snapshots (ckpt-iter{N}.pt) gated on a --snapshot_every interval so a single corrupted save doesn't cost the whole run.

Contributing

The repo is MIT-licensed and accepts pull requests. The historic maintainer cadence is bursty — long quiet stretches followed by waves of PR review. The README explicitly notes a successor project (nanochat) that is now the more actively developed evolution.

9. Integration Points

nanoGPT does not expose any RPC, REST, or library API. Integrating with external systems happens at the filesystem boundary:

Inputs

Path Format Producer Notes
data/{dataset}/train.bin uint16 raw byte stream Your prepare.py Token IDs concatenated, no separators.
data/{dataset}/val.bin uint16 raw byte stream Your prepare.py Same format as train.bin.
data/{dataset}/meta.pkl Python pickle Your prepare.py (optional) Dict with vocab_size, stoi, itos. Read by train.py to set vocab_size.
{out_dir}/ckpt.pt torch.save pickle train.py (when resuming) Loaded by train.py (init_from='resume') and sample.py.
Config files (config/*.py) Plain Python You Run via exec() by configurator.py.
sys.argv --key=value CLI You Parsed by configurator.py.
HuggingFace cache Per HF HF SDKs Source for GPT.from_pretrained and datasets.load_dataset.

Outputs

Path Format Consumer
{out_dir}/ckpt.pt torch.save pickle train.py (resume), sample.py.
stdout Plain text Your terminal, tee, log shippers, screen/tmux buffers.
Wandb dashboard HTTP POST to wandb.ai Your wandb account, when wandb_log=True.
./bench_log/* TensorBoard traces TensorBoard / torch.profiler viewer (only when bench.py --profile=True).

External services

Service Used by When
HuggingFace Hub transformers.GPT2LMHeadModel.from_pretrained init_from='gpt2*' paths in train.py and sample.py.
HuggingFace Datasets datasets.load_dataset("openwebtext") data/openwebtext/prepare.py.
GitHub raw requests.get(...) data/shakespeare/prepare.py, data/shakespeare_char/prepare.py for tinyshakespeare/input.txt.
Wandb wandb.init, wandb.log Only when wandb_log=True.
NCCL torch.distributed DDP runs (single- and multi-node).

Embedding nanoGPT in a larger system

There is no clean library boundary — train.py is a script, not a class. The standard pattern for embedding nanoGPT components elsewhere is:

from model import GPTConfig, GPT
import torch

# build a fresh model
cfg = GPTConfig(n_layer=6, n_head=6, n_embd=384,
                block_size=256, vocab_size=65, dropout=0.0, bias=False)
model = GPT(cfg)

# or load a checkpoint
ckpt = torch.load('out/ckpt.pt', map_location='cpu')
model = GPT(GPTConfig(**ckpt['model_args']))
model.load_state_dict({k.removeprefix('_orig_mod.'): v
                        for k, v in ckpt['model'].items()})
model.eval()

# generate
import tiktoken
enc = tiktoken.get_encoding('gpt2')
ids = torch.tensor([enc.encode("Hello, world")], dtype=torch.long)
with torch.no_grad():
    out = model.generate(ids, max_new_tokens=20, temperature=0.8, top_k=200)
print(enc.decode(out[0].tolist()))

model.py has no side effects at import time, so importing it from inside your own script is safe. train.py and sample.py execute their training and sampling loops at module-level and should not be imported as libraries.

Important

If you're integrating nanoGPT into anything resembling a production system, do not import train.py or sample.py — copy out the functions you need into your own module. Both scripts run their main logic at module top-level, including torch.manual_seed, model construction, and (in sample.py) the entire generation loop.