Decoder-only transformer in PyTorch

A GPT-style decoder written line by line in PyTorch, to own every tensor shape between input_ids and logits.

Things I implemented:

• Tokenization plus token and positional embeddings
• Causal multi-head self-attention built by hand
• Residual stream, layer norm and MLP blocks in the pre-norm layout
• The full training loop, trained on Apple MPS
• Sampling from the trained model

The causal mask is where most first attempts go wrong: it has to be applied before the softmax, as additive negative infinity:

mask = torch.tril(torch.ones(T, T, device=device))
att = att.masked_fill(mask == 0, float("-inf"))
att = att.softmax(dim=-1)