Train and Load a GPT from Scratch

Build the 124M GPT, train it with AdamW and cosine warmup, load GPT-2 weights, gate the pipeline.

Run the lessons#

Command	What it does
`python3 code/main.py` (35)	Counts 124M params, generates with a tiny config
`python3 code/main.py` (36)	Trains a tiny model, writes `outputs/losses.jsonl`
`python3 code/main.py` (37)	Writes `outputs/gpt2-stub.safetensors`, loads it
`python3 code/main.py` (44)	20 steps, per step lr and grad norm, text plot
`python code/main.py plan`	Validates manifest, cost estimate, prints DAG
`python code/main.py run`	Executes stages, writes `manifest.out.yaml`
`python code/main.py gate`	Applies eval gates. Prints `SHIP` or `HOLD`
`build_param_groups(model, weight_decay)`	AdamW groups, decay split

Setting	Default	Note
`vocab_size`	50257	Rows of the token table
`context_length`	1024	Rows of the position table
`d_model`	768	Embedding width
`num_heads`	12	`d_model` must divide by it
`num_layers`	12	Twelve `TransformerBlock` modules
`mlp_expansion`	4	MLP hidden is 4 times `d_model`
`dropout`	0.1	Tiny demo sets 0.0
`weight_tying`	True	`lm_head` shares `tok_embed`

ref_model = GPTModel(GPTConfig())124M reference

count_parameters(ref_model)unique params, tied once

GPTModel(GPTConfig(weight_tying=False))adds vocab_size * d_model

self.lm_head.weight = self.tok_embed.weightshare, never copy

generate(tiny_model, prompt, max_new_tokens=12,

temperature=0.8, top_k=20, seed=42)temperature must be positive

top_k_filter(logits, top_k)others to -inf

window = tokens[:, -model.cfg.context_length:]sliding window

Why transpose on load?: c_attn, c_proj and c_fc use conv1d layout. nn.Linear wants the transpose.
Why print a sample during training?: Loss is a scalar. A sample shows a model stuck on one token.
Why slice eval batches once?: Identical batches make eval loss comparable across runs.
Where does the schedule read its step?: From global_step in the checkpoint, not a local counter.

optimizer = torch.optim.AdamW(

build_param_groups(model, cfg.weight_decay),decay split

lr=cfg.max_lr, betas=(0.9, 0.95))

for step in range(cfg.num_steps):

lr = cosine_with_warmup(step, cfg.warmup_steps, cfg.num_steps, cfg.max_lr, cfg.min_lr)

for group in optimizer.param_groups:

group["lr"] = lr

inputs, targets = next(train_loader)

optimizer.zero_grad(set_to_none=True)

loss = calc_loss_batch(model, inputs, targets)flat cross entropy

loss.backward()

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=cfg.grad_clip)

optimizer.step()

Warmup: lr_max * step / warmup_steps up to warmup_steps.
Cosine: lr_min + 0.5 * (lr_max - lr_min) * (1 + cos(pi * progress)).
Floor: past total_steps the rate stays at lr_min.
Warmup size: a fraction of total steps. Typical is 1 to 3 percent.
Floor size: lr_min at 10 percent of lr_max keeps the optimizer learning.
Weight tying: saves 50257 times 768 = 38 million parameters.
Residual init: scale out projections by 1 / sqrt(2 * num_layers).
Pre-training FLOPs: 6 x params x tokens.

make_stub_safetensors(STUB_PATH, cfg, seed=42)offline fixture

model = GPTModel(cfg)

report = load_safetensors(model, STUB_PATH, verbose=False)

report.summary()loaded= missing= unexpected= shape_mismatch=

report.ok()no missing, no shape_mismatch

model.lm_head.weight.data_ptr() == model.tok_embed.weight.data_ptr()tied

Symptom	Cause	Fix
Loss near zero, model learns nothing	No shift by one	Target at positions 1..T
LayerNorm scale goes to zero	Weight decay on the scale	Scale and bias in the no decay group
One bad batch wipes hours	Huge gradient	`clip_grad_norm_` before `step`
Same samples before and after load	Name map missed every tensor	Read the `LoadReport`
`shape_mismatch` in the report	Dims differ from the config	Loader assigns nothing. Fix the config
Off by one LR late in the run	Glued three part schedule	One function over `step`, test the edges