Capstone Part 1: Model Loader, Weights, and Single-Sequence Forward

You have studied tokenization, attention, Transformer blocks, modern LLaMA-style architecture, and inference loops. Day 27 turns that knowledge into source code. The first milestone is intentionally narrow: one prompt, one sequence, no KV cache yet, and correctness before speed.

1. Read config.json into a typed LLMConfig.
2. Map safetensors weight names to model submodules.
3. Build embeddings, RMSNorm, RoPE, attention, SwiGLU FFN, blocks, and LM head.
4. Run a single forward pass and compare logits to HuggingFace Transformers.
5. Implement greedy decoding for a short continuation.

• hidden_size is model width, such as 2048.
• n_heads is query attention head count.
• n_kv_heads is key/value head count under GQA.
• head_dim = hidden_size / n_heads unless explicitly set.
• max_abs_diff is the largest absolute logit difference against the reference.

The course assumes TinyLlama-1.1B as the default capstone target because it is LLaMA-like, small enough for accessible hardware, and still structurally realistic. Example values are n_layers = 22, n_heads = 32, n_kv_heads = 4, head_dim = 64, hidden_size = 2048, and intermediate_size = 5632.

A professional engine begins with boring correctness: load config, list tensors, assert shapes, and fail loudly when names do not match. Safetensors is ideal because it stores tensor metadata and avoids pickle execution.

Single-sequence forward is token ids -> embeddings -> repeated transformer blocks -> final RMSNorm -> LM head -> logits. The logits shape is [B, T, vocab_size]. For generation you only need logits[:, -1, :], but for verification you compare the full output against Transformers.

Correctness means same token IDs, same dtype policy, same position IDs, same attention mask, same model mode, and no accidental dropout. Debug by comparing embeddings first, then norm, Q/K/V projections, attention, FFN, and final logits.

After logits match, greedy decode is simple: take argmax of the last logits, append the token, and run forward again. This is intentionally inefficient because Day 28 adds KV cache.

1. Download or point to a small LLaMA-style checkpoint.
2. Load config.json and assert every required field.
3. List all safetensors keys and map layer 0 shapes.
4. Run one forward pass and compare logits to Transformers.
5. Generate 20 greedy tokens and save the prompt, output, and logit diff.

1. Why verify logits instead of just generated text? Text can match by accident; logits expose numeric component bugs.
2. What shape does GQA change? K/V projection output width is n_kv_heads head_dim, smaller than Q.*
3. Why no KV cache today? Correct single-forward math must be proven before optimizing the decode loop.
4. Common source of RoPE bugs? Wrong position offset or wrong head_dim pairing.
5. Why use config-driven construction? The same engine can load model variants without hard-coded dimensions.