Lesson 23 of 23 · Capstone

Walk one block without losing a dimension.

Your win: narrate every operation in a pre-norm causal Transformer block with RoPE, then map that readable block onto MiniMind's cache-aware implementation.

25 minutesNeeds: Lessons 1–22Outcome: teach and extend the block

Final mission: start with \(X^{(\ell)}\), finish with \(X^{(\ell+1)}\), and explain exactly where position, comparison, probability, content, and residual updates enter.

MINIMINDBLOCK.FORWARD · THE BIG PICTURE

Two updates. One continuous state.

Attention shares information across tokens. The MLP reshapes features within each token.

Toy walkthrough: \(B=1,T=2,d_{\text{model}}=2,h=1\)

Use RMSNorm with unit scale, identity projection matrices, positions \(m=0,1\), RoPE angle step \(\theta=\pi/2\), and

\[ X^{(\ell)}=\begin{bmatrix}1&1\\1&-1\end{bmatrix}. \]

Each row has RMS \(1\), so \(\bar X=X^{(\ell)}\). Identity projections give \(Q=K=V=\bar X\).

At position 0, \((1,1)\) stays \((1,1)\). At position 1,

\[ R_{\pi/2}\begin{bmatrix}1\\-1\end{bmatrix} =\begin{bmatrix}0&-1\\1&0\end{bmatrix} \begin{bmatrix}1\\-1\end{bmatrix} =\begin{bmatrix}1\\1\end{bmatrix}. \]

Thus \(Q'=K'=\begin{bmatrix}1&1\\1&1\end{bmatrix}\). The scaled score and masked score matrices are

\[ S=\frac{Q'K'^{\mathsf T}}{\sqrt2} =\begin{bmatrix}\sqrt2&\sqrt2\\\sqrt2&\sqrt2\end{bmatrix}, \qquad A=\begin{bmatrix}\sqrt2&-\infty\\\sqrt2&\sqrt2\end{bmatrix}. \]

Row-wise softmax gives \(P=\begin{bmatrix}1&0\\0.5&0.5\end{bmatrix}\). Therefore

\[ O=PV =\begin{bmatrix}1&0\\0.5&0.5\end{bmatrix} \begin{bmatrix}1&1\\1&-1\end{bmatrix} =\begin{bmatrix}1&1\\1&0\end{bmatrix}. \] \[ Y=X^{(\ell)}+O =\begin{bmatrix}2&2\\2&-1\end{bmatrix}. \]

After the second normalization, suppose the learned SwiGLU produces the toy update \(F=\begin{bmatrix}0.1&-0.2\\0.3&0.1\end{bmatrix}\). The block output is

\[ X^{(\ell+1)}=Y+F =\begin{bmatrix}2.1&1.8\\2.3&-0.9\end{bmatrix}. \]

The chosen \(F\) stands in for the SwiGLU calculation from Lesson 22; every attention-stage number above is calculated explicitly.

Code checkpoint · assemble the TransformerModel

The block pieces from Lessons 15–22 now form a decoder-only language model. Every block preserves \([B,T,D]\); the language-model head converts the final features into one logit per vocabulary item.

Show the final assembly

class TransformerModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            TransformerBlock(config) for _ in range(config.num_layers)
        )
        self.final_norm = RMSNorm(config.hidden_size, config.norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight

    def forward(self, input_ids):
        x = self.token_embedding(input_ids)  # [B, T] -> [B, T, D]
        for layer in self.layers:
            x = layer(x)                    # [B, T, D] -> [B, T, D]
        return self.lm_head(self.final_norm(x))  # [B, T, vocab]

The complete assembly stays in this lesson so the shape ledger, block trace, and implementation can be read together. You can also download transformer_model.py.

Trace it: Starting from input_ids, name the final axis at each of the three comments.

Engineering bridge · grow this into MiniMind

The course TransformerModel is intentionally readable. The full target at model/model_minimind.py in the MiniMind checkout keeps the same pre-norm block equation, then adds runtime contracts around it. Use the table as an implementation order; each row has a shape or numerical check that catches regressions early.

Stage	MiniMind detail to add	Check before moving on
1 · Config	`MiniMindConfig` owns \(H\), \(H_{\mathrm{kv}}\), \(D_h\), RoPE limits, optional YaRN `rope_scaling`, dropout, and the `use_moe` switch.	Reject incompatible head counts and odd \(D_h\); print a shape ledger.
2 · Attention	Use GQA projections, per-head `q_norm`/`k_norm`, RoPE, and `repeat_kv`. Precompute cosine/sine tables with `precompute_freqs_cis`.	Expanded K/V agree with Q on \(H\); output returns \([B,T,D]\).
3 · Decode	Concatenate `past_key_value` along tokens and slice RoPE tables from `start_pos`.	One-token cached logits match full-prefix logits within tolerance.
4 · Runtime	Select Flash Attention for eligible prefixes; otherwise apply causal and padding masks explicitly with FP32 softmax.	Test both branches with a padding mask and with a non-empty cache.
5 · Feed-forward	Keep dense SwiGLU; optionally swap in top-k `MOEFeedForward` and accumulate its auxiliary loss.	Every route returns \([B,T,D]\); the dense path remains the baseline.
6 · Model wrapper	Add final RMSNorm, tied embedding/lm-head weights, shifted label loss, and generation sampling.	Logits are \([B,T,V]\); tied weights are the same parameter object.

RoPE lifecycle: MiniMind registers the precomputed tables as non-persistent buffers so they move with the model without bloating checkpoints. Its forward path also rebuilds them on the hidden-state device if meta-device initialization left zeroed tables.

Show the MiniMind block contract

residual = hidden_states
hidden_states, present_key_value = self.self_attn(
    self.input_layernorm(hidden_states), position_embeddings,
    past_key_value, use_cache, attention_mask
)
hidden_states += residual
hidden_states = hidden_states + self.mlp(
    self.post_attention_layernorm(hidden_states)
)
return hidden_states, present_key_value

Wrapper contract: the Hugging Face model returns logits, shifted-label loss, auxiliary MoE loss, and past_key_values. Its generation loop feeds only the uncached suffix, then applies temperature, top-k, and top-p sampling.

Suggested build sequence: port the dense block, add GQA, add cache-aware decoding, verify Flash/fallback parity, then add MoE and the Hugging Face wrapper. Keep one small deterministic input fixed for every comparison.

Retrieval check

Answer once before reopening the code; use the feedback to correct your explanation.

Where does RoPE occur in this block?

STAGE 4 · PROVE IT

Capstone studio · narrate → build → verify

Without looking above, write the eight stages of the block in order.
For \(B=2,T=5,d_{\text{model}}=12,h=3\), make a shape ledger for split Q, one head's score matrix, concatenated output, and final block output.
Repeat the toy multiplication \(PV\) and residual addition \(X^{(\ell)}+O\).
Explain why \(R_m^{\mathsf T}R_n=R_{n-m}\) matters specifically inside Step 4.
Name the two places where information is mixed across features and the place where it is mixed across tokens.
Teach the block aloud in under two minutes without saying “it just learns it.” Name the input and output of every stage.

BUILD CHALLENGE · RUN IT LOCALLY

Engineering challenge · add a shape ledger to the block

Instrument `MiniMindBlock.forward` with assertions or logging for the input state, attention update, first residual, MLP update, and returned state. Then add one cached decode case where `Tq=1`.

Standard answer · compare after your attempt

assert hidden_states.ndim == 3
residual = hidden_states
hidden_states, present_key_value = self.self_attn(...)
assert hidden_states.shape == residual.shape
hidden_states = hidden_states + residual
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
assert hidden_states.shape == residual.shape
return hidden_states, present_key_value

The block preserves `[B, T, D]`; only the cache is allowed to grow along the key/value token axis during decoding.

Check solutions

Normalize; project Q/K/V; apply RoPE to Q/K; score and causally mask; softmax and blend V; concatenate heads and apply \(W_O\); first residual; normalize, SwiGLU, and second residual.
Split Q: \(2\times3\times5\times4\); one head score per batch: \(5\times5\); concatenated output: \(2\times5\times12\); final output: \(2\times5\times12\).
\(O=\begin{bmatrix}1&1\\1&0\end{bmatrix}\); \(Y=\begin{bmatrix}2&2\\2&-1\end{bmatrix}\).
It turns the rotated query-key dot product into a comparison that depends on relative displacement \(n-m\).
Projection/output matrices and the FFN mix features; attention's \(PV\) mixes information across allowed token positions.
Use the reference sheet to check omissions after speaking, not during the first attempt.

Primary sources: Vaswani et al.; Su et al., RoFormer; Touvron et al., LLaMA.