# nanoGPT Administrator & Developer Guide

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

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

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

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

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

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

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

```bash
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`.

```bash
# 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

```text
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:

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

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

```python
{
    '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:

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

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

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

```text
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)

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

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

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

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

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

---

_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. Unofficial; not affiliated with or reviewed by the project._