nanoGPT Rebuilt — A Complete Tested Decoder

Rebuild a specified GPT, not an invented commit history

The useful exercise in nanoGPT is tracing a compact decoder-only model end to end: configuration, embeddings, fused QKV attention, pre-norm blocks, final norm, tied output projection, optimizer groups and generation. Its public model/train files are primary implementation references. Their current branch is mutable, so this review stores dated raw hashes rather than inventing January-2023 commit messages, star counts, video views or precise speedup chronology.

The implementation below is a teaching rebuild, not a byte-compatible checkpoint loader. It uses exact GELU by default, learned absolute positions, dense causal attention, optional biases and weight tying. Loading a GPT-2 checkpoint additionally requires exact configuration, activation convention, parameter naming and transposition of Conv1D-style weights; matching a few module names does not make load_state_dict sufficient. GPT-4's public report withholds architectural specifics, and Llama/DeepSeek variants differ beyond merely tensor size. No claim to reconstruct proprietary internals is made.

Prerequisites are MHA, positions, blocks and training-loop mechanics. Goals: a complete CPU-runnable model; manual/SDPA parity; tied-parameter and optimizer coverage checks; valid sampling; and an honest comparison protocol for reference implementations.

Configuration and the complete model

Save the next fence as gpt_lab.py. Imports and defaults are complete. The small defaults intentionally avoid allocating a 124M model merely to inspect the code. block_size bounds the learned position table; vocab_size is the actual output vocabulary. Padding GPT-2's 50,257 classes to 50,304 may improve particular matrix kernels, but it adds parameters/logit work and invalid output IDs that need an explicit masking policy. It is not a zero-cost change. Our model does not pad the vocabulary.

One Linear from d to 3d is algebraically a concatenation of Q/K/V projections. Whether it is faster is a kernel/workload measurement. The manual fallback registers a nonpersistent triangular buffer; it moves with the module but is neither optimized nor saved in the state dict. Persistent buffers do normally enter the state dict, contrary to the old claim that a buffer always disappears from checkpoints.

from dataclasses import dataclass,asdict
import math, copy
import torch
from torch import nn
from torch.nn import functional as F
 
@dataclass
class GPTConfig:
    block_size:int=16
    vocab_size:int=12
    n_layer:int=2
    n_head:int=2
    n_embd:int=16
    dropout:float=0.
    bias:bool=False
    tied:bool=True
    backend:str='manual'
    def __post_init__(self):
        if min(self.block_size,self.vocab_size,self.n_layer,self.n_head,self.n_embd)<1:raise ValueError('positive dimensions required')
        if self.n_embd%self.n_head:raise ValueError('width must divide across heads')
        if not 0<=self.dropout<1:raise ValueError('dropout must be in [0,1)')
        if self.backend not in ('manual','sdpa'):raise ValueError('backend')
 
class CausalSelfAttention(nn.Module):
    def __init__(self,cfg):
        super().__init__();self.cfg=cfg;d=cfg.n_embd
        self.qkv=nn.Linear(d,3*d,bias=cfg.bias);self.proj=nn.Linear(d,d,bias=cfg.bias)
        self.attn_drop=nn.Dropout(cfg.dropout);self.resid_drop=nn.Dropout(cfg.dropout)
        self.register_buffer('allowed',torch.ones(cfg.block_size,cfg.block_size,dtype=torch.bool).tril(),persistent=False)
    def forward(self,x):
        B,T,C=x.shape;h=self.cfg.n_head;dk=C//h
        q,k,v=self.qkv(x).split(C,-1)
        q,k,v=[z.reshape(B,T,h,dk).transpose(1,2) for z in (q,k,v)]
        if self.cfg.backend=='sdpa':
            out=F.scaled_dot_product_attention(q,k,v,attn_mask=None,
                dropout_p=self.cfg.dropout if self.training else 0.,is_causal=True)
        else:
            score=(q@k.transpose(-1,-2))/math.sqrt(dk)
            weights=score.masked_fill(~self.allowed[:T,:T],-torch.inf).softmax(-1)
            out=self.attn_drop(weights)@v
        return self.resid_drop(self.proj(out.transpose(1,2).reshape(B,T,C)))
 
class Block(nn.Module):
    def __init__(self,cfg):
        super().__init__();d=cfg.n_embd
        # PyTorch 2.9 LayerNorm CAN disable bias while retaining scale.
        self.ln1=nn.LayerNorm(d,bias=cfg.bias);self.attn=CausalSelfAttention(cfg)
        self.ln2=nn.LayerNorm(d,bias=cfg.bias)
        self.fc=nn.Linear(d,4*d,bias=cfg.bias);self.proj=nn.Linear(4*d,d,bias=cfg.bias)
        self.drop=nn.Dropout(cfg.dropout)
    def forward(self,x):
        x=x+self.attn(self.ln1(x))
        return x+self.drop(self.proj(F.gelu(self.fc(self.ln2(x)))))
 
class GPT(nn.Module):
    def __init__(self,cfg):
        super().__init__();self.config=cfg
        self.token=nn.Embedding(cfg.vocab_size,cfg.n_embd)
        self.position=nn.Embedding(cfg.block_size,cfg.n_embd)
        self.drop=nn.Dropout(cfg.dropout)
        self.blocks=nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.final=nn.LayerNorm(cfg.n_embd,bias=cfg.bias)
        self.head=nn.Linear(cfg.n_embd,cfg.vocab_size,bias=False)
        # Initialize independent tensors once, then establish shared identity.
        for module in self.modules():
            if isinstance(module,(nn.Linear,nn.Embedding)):
                nn.init.normal_(module.weight,std=.02)
                if isinstance(module,nn.Linear) and module.bias is not None:nn.init.zeros_(module.bias)
        if cfg.tied:self.head.weight=self.token.weight
        for block in self.blocks:
            nn.init.normal_(block.attn.proj.weight,std=.02/math.sqrt(2*cfg.n_layer))
            nn.init.normal_(block.proj.weight,std=.02/math.sqrt(2*cfg.n_layer))
    def forward(self,ids,targets=None,last_only=False):
        if ids.dtype!=torch.long or ids.ndim!=2:raise ValueError('long batch-by-time IDs required')
        B,T=ids.shape
        if B<1 or not 1<=T<=self.config.block_size:raise ValueError('empty/excess context')
        if (ids<0).any() or (ids>=self.config.vocab_size).any():raise ValueError('ID outside vocabulary')
        x=self.drop(self.token(ids)+self.position(torch.arange(T,device=ids.device)))
        for block in self.blocks:x=block(x)
        x=self.final(x)
        logits=self.head(x[:,-1:] if last_only and targets is None else x)
        loss=None
        if targets is not None:
            if targets.shape!=ids.shape or targets.dtype!=torch.long:raise ValueError('target shape/dtype')
            if ((targets < -1)|(targets>=self.config.vocab_size)).any() or not targets.ne(-1).any():raise ValueError('invalid/all-ignored targets')
            loss=F.cross_entropy(logits.reshape(-1,logits.size(-1)),targets.reshape(-1),ignore_index=-1)
        return logits,loss
    def configure_optimizers(self,lr=3e-4,weight_decay=.1,betas=(.9,.95)):
        # named_parameters deduplicates the tied embedding/head Parameter.
        parameters=[p for p in self.parameters() if p.requires_grad]
        matrix=[p for p in parameters if p.ndim>=2];other=[p for p in parameters if p.ndim<2]
        assert len({id(p) for p in matrix+other})==len(parameters)
        return torch.optim.AdamW([{'params':matrix,'weight_decay':weight_decay},
                                  {'params':other,'weight_decay':0.}],lr=lr,betas=betas)
    @torch.no_grad()
    def generate(self,ids,max_new_tokens,temperature=1.,top_k=None,greedy=False,generator=None):
        if max_new_tokens<0 or temperature<=0 or not math.isfinite(temperature):raise ValueError('invalid count/temperature')
        if top_k is not None and (not isinstance(top_k,int) or top_k<1):raise ValueError('positive integer top_k')
        mode=self.training;self.eval()
        try:
            # Validate in evaluation mode, even for zero new tokens.
            self(ids[:,-self.config.block_size:],last_only=True)
            for _ in range(max_new_tokens):
                logits,_=self(ids[:,-self.config.block_size:],last_only=True)
                scores=logits[:,-1]/temperature
                if greedy:nxt=scores.argmax(-1,keepdim=True)
                else:
                    if top_k is not None:
                        k=min(top_k,scores.size(-1));values,indices=scores.topk(k,-1)
                        kept=torch.full_like(scores,-torch.inf);scores=kept.scatter(-1,indices,values)
                    nxt=torch.multinomial(scores.softmax(-1),1,generator=generator)
                ids=torch.cat([ids,nxt],1)
            return ids
        finally:self.train(mode)
 
def parameter_count(cfg):
    d=cfg.n_embd;L=cfg.n_layer
    # 12d² matrices; 2d LN scales; bias adds 9d linear +2d LN.
    block=12*d*d+2*d+(11*d if cfg.bias else 0)
    return cfg.vocab_size*d+cfg.block_size*d+L*block+d+(d if cfg.bias else 0)+(0 if cfg.tied else cfg.vocab_size*d)

The final normalization is part of this model's specification, not a claim all Transformer variants require it. The scaled residual initializer shrinks attention/FFN output projection standard deviations by square root of 2L under a residual-variance heuristic; it does not guarantee arbitrary-depth stability or predict “NaN at step 500” when omitted. Bias-free is an option, not a universal speed or quality law.

Prove parity, causality, tying and optimizer coverage

The functional SDPA call scales scores internally; do not scale q again unless changing the contract intentionally. Backend eligibility depends on device, dtype, shape, mask and installed implementation; is_causal=True expresses semantics, not a guarantee no mask-like storage exists in the backend. Our CPU parity is not a GPU FlashAttention benchmark. Tiling exact attention reduces IO/materialization, not dense pairwise arithmetic to linear time.

if __name__=='__main__':
    torch.manual_seed(40)
    cfg=GPTConfig();model=GPT(cfg).double()
    ids=torch.tensor([[1,2,1,3],[3,2,4,5]]);targets=torch.tensor([[2,1,3,4],[2,4,5,6]])
    a,loss=model(ids,targets)
    other=GPT(GPTConfig(**{**asdict(cfg),'backend':'sdpa'})).double();other.load_state_dict(model.state_dict())
    b,loss2=other(ids,targets)
    assert torch.allclose(a,b,atol=1e-12) and torch.allclose(loss,loss2,atol=1e-12)
    loss.backward();loss2.backward()
    assert all(torch.allclose(p.grad,q.grad,atol=1e-11) for p,q in zip(model.parameters(),other.parameters()))
    changed=ids.clone();changed[:,-1]=7
    assert torch.allclose(model(changed)[0][:,:-1],a.detach()[:,:-1],atol=1e-12)
    assert model.head.weight is model.token.weight
    assert sum(p.numel() for p in model.parameters())==parameter_count(cfg)
    untied=GPT(GPTConfig(**{**asdict(cfg),'tied':False}))
    assert sum(p.numel() for p in untied.parameters())-parameter_count(cfg)==cfg.vocab_size*cfg.n_embd
    opt=model.configure_optimizers()
    ids_in_groups=[id(p) for g in opt.param_groups for p in g['params']]
    assert len(ids_in_groups)==len(set(ids_in_groups))==len(list(model.parameters()))
    assert id(model.token.weight) in {id(p) for p in opt.param_groups[0]['params']}
    assert not any('allowed' in key for key in model.state_dict())
    model.train();g=model.generate(ids[:1],20,greedy=True)
    assert g.shape==(1,24) and model.training
    sampled=model.generate(ids[:1],5,top_k=1,generator=torch.Generator().manual_seed(1))
    assert torch.equal(sampled,model.generate(ids[:1],5,greedy=True))
    for bad in (0.,-1.,float('nan')):
        try:model.generate(ids[:1],1,temperature=bad)
        except ValueError:pass
        else:raise AssertionError('temperature validation')
    print('GPT output/gradient/backend/causal/tied/optimizer/generation contracts: pass',parameter_count(cfg))

The mask is a buffer, not a Parameter. The optimizer sees every trainable parameter once, including a tied embedding matrix in the decayed matrix group. The old warning “this rule excludes embeddings” contradicted its own dimension-based code. AdamW's decoupled decay is not generally equivalent to adding an L2 objective penalty under adaptive preconditioning. Excluding gains/biases is a common empirical policy, not a shape theorem; biases can overfit and nonzero gain is not automatically annihilated by a modest decay step. New modules with other parameter ranks need reviewed grouping.

Weight tying and exact budget arithmetic

With no biases, one block has 12d² matrix parameters and 2d LN gains. A tied model totals Vd+Pd+L(12d²+2d)+d. With biases, add 11d per block and d for the final LN bias. Untying adds exactly Vd; it does not necessarily double the model.

For L=6,h=6,d=384,P=256,V=65 and no biases, the exact total is 10,745,088. Untying adds 24,960, about 0.23% of that total—not half. With all ordinary biases, the count is 10,770,816. Count these configurations arithmetically or on a meta-device instead of allocating them just to count.

A tied token row receives input-lookup and output-softmax gradients. An untied input row is sparse in data usage while the output row may receive gradient through the full softmax; that observation helps explain different training signals. It does not prove every output row always has a nonzero gradient or that rare untied vectors must have zero cosine similarity. Tying is a regularization/representation choice with empirical support, not a guaranteed quality improvement. Different input/output vocabularies or embedding widths require a mapping or prevent direct tying; even matching vocabularies may be left untied deliberately.

The full GPT-2-compatible 12×768,50,257-vocabulary,1024-position configuration with biases has 124,439,808 unique parameters. A rounded advertised historical label and a tensor inventory are different things. Tokenizer identity and padded output rows change the inventory and likelihood baseline.

These larger counts can be checked without allocating their weights. Meta-device tensors retain shapes but no ordinary data storage; do not attempt a forward pass on this counting model.

if __name__=='__main__':
    for bias,expected in [(False,10745088),(True,10770816)]:
        large=GPTConfig(n_layer=6,n_head=6,n_embd=384,block_size=256,vocab_size=65,bias=bias)
        with torch.device('meta'):
            counted=GPT(large)
        assert sum(p.numel() for p in counted.parameters())==parameter_count(large)==expected
    gpt2=GPTConfig(n_layer=12,n_head=12,n_embd=768,block_size=1024,vocab_size=50257,bias=True)
    assert parameter_count(gpt2)==124439808
    print('exact tiny-Shakespeare and GPT-2-shaped parameter arithmetic: pass')

Sampling and the context boundary

Temperature divides logits by a positive number. Smaller positive temperatures sharpen; they are not literally deterministic or “safe.” Greedy argmax is a separate option; temperature zero is rejected rather than divided by. Top-k above retains exactly k selected indices (or all V when k exceeds V); ties are resolved by the selected topk indices, unlike threshold masking that can retain more than k. Sampling needs a recorded generator seed for repeatability under the same runtime. Low temperature cannot repair a bad model, and high temperature does not guarantee creativity.

Only the final hidden position needs vocabulary logits during generation. Earlier hidden states still matter for attention and are recomputed in this simple model. Cropping to the last block_size tokens resets learned position indices within the cropped window: this is an explicit bounded-context policy, not a KV cache. A real cache needs per-layer K/V states and position-consistency tests. Registering a causal mask avoids repeated explicit mask allocation; SDPA can make the user-level mask unnecessary.

eval() disables module dropout, and functional SDPA receives dropout_p=0 while not training. no_grad() alone does not do either. Generation restores the previous training mode after completion so samples do not silently alter the surrounding training loop. No checkpoint or data downloads occur in the module.

Training comparison and performance experiments

The next lesson provides a complete local-text training/validation/checkpoint recipe. To compare this rebuild against a saved nanoGPT revision, align vocabulary/order, splits, biases, tied weights, GELU approximation, dropout locations, initialization, context, batch/accumulation, loss reduction, optimizer groups, LR clock and random seeds. Compare forward logits/gradients under copied tensors before comparing training trajectories. A “within noise” result requires repeated runs and a defined noise estimate; neither model.py alone nor a loss near 1.48 establishes equivalence.

Optional performance mechanisms retain distinct meanings:

• torch.compile specializes/captures computation; compile overhead, graph breaks and shapes affect amortization. No fixed 20–40% gain is claimed or tested here.
• AMP BF16/FP16 changes selected computation precision; unsupported devices or unstable operations require different choices. It is not automatically 2× throughput.
• DDP replicates parameters and synchronizes gradients; global batch/accumulation, loss denominators and communication costs must be reconciled.
• Activation checkpointing trades recomputation for stored activations and needs RNG care with dropout. It is not assumed to exist in every upstream nanoGPT revision merely because it is useful.
• Warmup/cosine schedules use an explicit optimizer-update clock; skipping an update or resuming changes that clock unless state is saved.

These are unexecuted optimization extensions, not dependencies of the bounded CPU recipe. The old fused=True AdamW assumption is removed from the portable default; inspect actual optimizer/device support before enabling fusion.

Read a speedrun as an experiment, not a timeless leaderboard

The dated modded-nanoGPT README is an experiment log with hardware, data, stopping metric and changing implementations. Muon orthogonalizes matrix updates using an approximate polar/Newton–Schulz procedure and commonly uses a different optimizer for other parameters. FP8 kernels, value residuals, changed activations, soft caps and U-shaped skip paths each change numerical/training contracts. Combining several in one PR prevents attributing a percentage of the gain to a single change without ablation.

For the original three-PR exercise, record exact base/head revisions, unchanged dataset/tokenizer/evaluation target, hardware and precision, warmup/compilation timing inclusion, seeds and time-to-target. Compute relative time reduction as (old−new)/old and throughput ratio as old/new; these are not the same percentage. Report quality failures/timeouts, not only fastest successful runs. No historical “45→13→6→3 minutes” table is retained without primary row-level validation.

A worked reading of the saved track-1 record table illustrates why the original “first three PRs, one change each” premise is unsafe. Its baseline row is llm.c, 45 minutes, not a freshly reproduced nanoGPT run. The next three entries are “Tuned learning rate & rotary embeddings” (31.4 minutes), “Introduced the Muon optimizer” (24.9 minutes, with no log linked in that row), and “Muon improvements” (22.3 minutes). Relative to each preceding row, the reported elapsed-time reductions are 30.22%, 20.70% and 10.44%. Those are arithmetic comparisons of a dated author's table, not isolated causal contributions, proof that these are the first three PRs, or independently validated benchmarks. In particular, the first combines at least two changes and the second lacks a linked run log. Reject a numerical attribution rather than fill the missing ablations with guesses.

record_minutes=[45.,31.4,24.9,22.3]
record_reductions=[100*(old-new)/old for old,new in zip(record_minutes,record_minutes[1:])]
assert [round(x,2) for x in record_reductions]==[30.22,20.70,10.44]
assert math.isclose(45/22.3,2.017937219730942)
print('reported table arithmetic only; not benchmark reproduction',record_reductions)

llm.c is a separate C/CUDA training implementation with its own requirements; it is not a promise to run on any CUDA GPU. JAX/TPU speedrun ports are separate backends and hardware protocols. Community nanoDeepSeek-style projects aim to teach MLA/MoE and other mechanisms, but their names do not prove checkpoint fidelity, benchmark parity or an identical 300-line architecture. llama.cpp is an inference project, not established here as a nanoGPT fork. The value of these projects is inspectable design, not invented genealogies.

Depth/width under fixed budget is also an experiment. Matrix parameters grow roughly as 12Ld²; layers create sequential dependencies and width changes matrix sizes and sharding. More depth may aid composition; more width may fit hardware better, but neither universally wins latency or quality. Use a documented comparable baseline and profile under the actual batch/context budget before spending on a search.

Practice and worked answers

1. Five implementation choices. Fused QKV, tied token/head identity, scaled residual initialization, reviewed AdamW parameter groups and a causal mask buffer or semantically equivalent SDPA call. Each has a separate test and scope; none guarantees a particular quality gain.
2. Tiny configuration count. Vd=24,960; Pd=98,304; six blocks contain 6×(1,769,472+768); final gain has384. Total10,745,088 without biases. Untying adds24,960, not another10M.
3. Failure fingerprints. Allocating a CPU mask inside a GPU forward can cause device mismatch and overhead; generation beyond a learned position table fails unless the context policy is explicit. The implementation crops the retained suffix and recomputes its positions.
4. SDPA versus manual. Both implement scaled masked attention under aligned dropout/mask conventions. SDPA may dispatch to optimized kernels on eligible hardware; this CPU parity test proves outputs/gradients, not GPU backend identity or speed.
5. When not tie. Nonmatching vocabularies/dimensions prevent direct sharing; a separately validated architecture may choose untied weights. Sharing equal-sized tensors is a modeling choice, not universal semantic equality of two independently learned maps.
6. Speedrun review. The saved track-1 baseline/next rows report 45→31.4→24.9→22.3 minutes: 30.22%, 20.70%, 10.44% reductions relative to the immediately prior row. LR plus RoPE is a combined change; the Muon-introduction row links no log. Thus these are reported-record arithmetic, not three verified single-change PR contributions. Pin revisions/protocol and require ablations before attributing causal percentages.

Type the config/block/model path from memory, then run the tests. If a shape error appears, trace it to a configuration field, but remember not every bug is a shape mismatch: wrong masks, target shifts, normalization axes and optimizer membership can all produce plausible tensors.

Sources and optional continuation

Primary dated implementations: nanoGPT model, train, weight tying study, decoupled weight decay, PyTorch2.9 LayerNorm, SDPA, modded-nanoGPT, Muon, llm.c. Exact raw hashes and passages are in the source ledger.

Optional original Let's build GPT and reproduce GPT-2 are learning resources, not evidence of this run's quality. Previous · Next: local text training.

Executed scope: small exact CPU model, manual/SDPA forward/backward parity, causality, parameter/optimizer contracts and bounded generation. No GPT-2/TinyShakespeare benchmark, pretrained loader, compile, GPU AMP/kernel speed, distributed run, application build or deployment is claimed.