The first module of the residency is to write a GPT-style transformer by hand — not to train anything useful, but to get to the point where the architecture is something I can reason about instead of recite. What follows is the thing I built, at the level of detail I wish someone had handed me on day one.

It is small on purpose: 2 layers, 4 heads, a 128-wide residual stream, a 256-token vocabulary, and a 64-token context. That comes out to 470,528 parameters, which is nothing — and that is the point. Every part is legible at this size, and the same parts scaled up are what a frontier model is made of.

The shape of the whole thing

A transformer takes a sequence of token ids and produces, for every position, a score over every token that could come next. Everything else — chat, agents, reasoning traces — is that operation in a loop.

The forward pass, from token ids through the blocks to logitstransformer block × 2token ids(B, T)token embedding + position embedding(B, T, 128)dropout(B, T, 128)layer norm → causal self-attention → + x(B, T, 128)layer norm → feed-forward → + x(B, T, 128)final layer norm(B, T, 128)lm head (no bias)(B, T, 256)logits → softmax → sample(B, 256)
The whole model, at the default config: two blocks, four heads, a 128-wide residual stream, and a 256-token vocabulary. B is batch, T is sequence length.

Read it top to bottom and you have the forward pass. Ids come in, become vectors, pass through blocks that repeatedly mix context in and think about it, and come out the other side as logits. The only genuinely interesting part is the attention step, so most of what follows is about that.

Embeddings, and why addition works

Two lookup tables. One maps a token id to a vector; the other maps a position — 0, 1, 2 — to a vector of the same width. Then you add them.

positions = torch.arange(0, seq_len, device=idx.device).unsqueeze(0)
x = self.token_embedding(idx) + self.position_embedding(positions)

Adding two unrelated meanings into one vector sounds like it should destroy both, and the first time I read it I assumed I was misunderstanding. You are not: 128 dimensions is a lot of room, the two tables are learned jointly, and nothing forces them to use the same directions. The model has every incentive to keep “which token” and “which position” in subspaces it can pull apart again. It works because it is trained to work, not because addition is principled.

Attention, which is the whole trick

Every position produces three vectors. A query: what am I looking for. A key: what do I have to offer. A value: what I hand over if you pick me. Compare each query against every key, turn those comparisons into weights, and take a weighted average of the values. That is it.

One attention head: query, key, and value projections through scores, mask, softmax, and the weighted sum of valuesxone Linear128 → 384Qwhat am I looking forKwhat do I offerVwhat I pass alongeach split into 4 heads → (B, 4, T, 32)
Q, K, and V come out of a single fused projection, then get reshaped into four heads of 32 dimensions each. The head split is a view, not a computation — the same numbers, read in a different shape.

In the code all three come out of a single linear layer that is three times as wide, then get split. That is not a conceptual choice, it is a performance one — one big matmul beats three small ones — and it is the first place the shapes get confusing, because384 is not a meaningful number, it is 3 × 128 wearing a disguise.

qkv = self.qkv(x)                    # (B, T, 384)
q, k, v = qkv.split(d_model, dim=2)  # three of (B, T, 128)

q = q.view(batch, seq_len, self.n_heads, self.head_dim)
q = q.transpose(1, 2)                # (B, 4, T, 32)

The head split is the other place people get lost, so it is worth saying plainly: it is a reshape. The 128 numbers become 4 groups of 32, and the transpose moves the head axis next to the batch axis so the matmuls treat heads as independent problems. No arithmetic happens. The same values, read in a different shape.

Then the comparison itself, scaled down by the square root of the head dimension. Without that scale, the dot products grow with dimension, the softmax saturates, and gradients go flat — you get a model that attends to exactly one thing and cannot be talked out of it.

scores = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)

The mask is where I lost an hour

A language model must not see the future. Position 3 predicting position 4 cannot be allowed to look at position 4, or training becomes a copying exercise that collapses the moment you ask it to generate something.

A six-by-six causal mask: each row can attend to itself and everything before it, never afterthecatsatonthematthecatsatonthemat−∞−∞−∞−∞−∞−∞−∞−∞−∞−∞−∞−∞−∞−∞−∞rows: the token doingthe lookingcolumns: what it isallowed to look atthe first token attendsonly to itself
The mask is filled with negative infinity before the softmax, not zeroed after. Softmax of −∞ is exactly zero; zeroing afterwards would leave the surviving weights not summing to one.
scores = scores.masked_fill(
    self.causal_mask[:, :, :seq_len, :seq_len] == 0, float("-inf")
)
weights = self.attn_dropout(F.softmax(scores, dim=-1))

Two details that cost me time. The mask is negative infinity applied before the softmax, not zeros applied after — softmax of −∞ is exactly zero, and the remaining weights still sum to one, which is not true if you zero things out afterwards. And the mask is registered as a buffer at full context length, then sliced to the current sequence: [:, :, :seq_len, :seq_len]. Forget the slice and short sequences silently attend against a mask built for long ones.

The block: residuals and where the norm goes

A block is attention, then a feed-forward network, each wrapped in a residual connection. Two lines:

x = x + self.attn(self.ln_1(x))
return x + self.ff(self.ln_2(x))

Note where the layer norm sits. It normalizes the input to the sublayer, not the output of the addition — pre-norm rather than the post-norm of the original 2017 paper. The practical consequence is that x travels from the embedding to the final norm along a path with nothing multiplicative on it. That unbroken path is why deep stacks train at all; post-norm puts a normalization in the middle of it and needs a learning-rate warmup to survive.

The feed-forward part is the least glamorous and roughly two thirds of the parameters: widen 128 to 512, apply GELU, narrow back to 128. If attention is where positions talk to each other, this is where each position thinks by itself.

Logits and the loss

A final norm, then one linear layer from 128 to the vocabulary. Cross entropy wants a flat list of predictions and a flat list of answers, so both get reshaped — the batch and time axes collapse into one:

logits = self.lm_head(self.ln_f(x))            # (B, T, 256)
loss = F.cross_entropy(
    logits.view(-1, logits.size(-1)),         # (B*T, 256)
    targets.view(-1),                         # (B*T,)
)

Every position is a training example. An eight-token sequence is eight predictions, not one, which is the quiet reason language models are so sample-efficient relative to how simple the objective looks.

Generation is a loop around all of it

Crop to the context window, run the forward pass, take the last position’s logits, divide by temperature, sample, append, repeat.

idx_cond = idx[:, -self.config.context_length :]
logits, _ = self(idx_cond)
next_logits = logits[:, -1, :] / max(temperature, 1e-6)
probs = F.softmax(next_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, next_token), dim=1)

The [:, -1, :] is easy to skim past and worth sitting with. The model computed a prediction for every position, and generation throws all but the last one away. That is pure waste, and eliminating it is exactly what a KV cache does — which is the next thing I want to build, because you cannot appreciate the cache until you have written the version that recomputes everything.

The shapes, in one place

Most of the real work was not conceptual. It was a shape mismatch and twenty minutes of figuring out which axis I had lost. Here is the whole journey for one sequence of eight tokens:

input token ids(1, 8)eight positions, integers
after embeddings(1, 8, 128)every token now a vector
fused qkv projection(1, 8, 384)Q, K, V stacked
split into heads(1, 4, 8, 32)a view, not a computation
attention scores(1, 4, 8, 8)every token against every token
back to the stream(1, 8, 128)heads concatenated
logits(1, 8, 256)one score per vocabulary entry

Print those at every step and the debugging stops being mysterious. Nearly every bug I hit was a transpose I had skipped or a view I had taken on the wrong axis, and every one of them showed up here first.

What it was actually worth

Nothing in here is novel. The value was not the artifact, it was the calibration: having debugged my own attention mask, I read claims about what models “know” differently, and I read claims about what post-training can fix differently too. The gap between following an architecture diagram and having built the thing is much wider than it looks from the diagram side.

The implementation is in the residency repo, along with the smoke tests that catch the shape errors. The rest of the sequence — policy gradients, GRPO, evals — is on the residency page.