Pretrained Weights and Text Generation

Lab: Run the companion CPU exercise at /labs/track1/03-cache-and-gqa/ (weight mapping, decoding strategies, cache verification).

Mapping a pretrained checkpoint into your model

Before a model can generate text, its weights must be compatible with the implementation. Hugging Face checkpoints (e.g. GPT-2, GPT-Neo, GPT-J) store weights in a well-known state_dict format, but the key names and tensor shapes may not match a custom decoder-only architecture directly.

Weight-mapping recipe

1. Rename prefixes. Hugging Face uses transformer.h.{i}. while a custom model may use blocks.{i}. (or layers.{i.}). Replace the prefix systematically.

2. Transpose Conv1D weights. GPT-2 stores attention and MLP weights as Conv1D matrices. Hugging Face transformers.pytorch_utils.Conv1D(nf, nx) stores weight with shape [nx, nf] ([in_features, out_features]) and evaluates via x @ weight + bias. PyTorch nn.Linear(in_features, out_features) stores weight with shape [out_features, in_features] and evaluates via x @ weight.T + bias.

   Therefore, copying weights from an HF GPT-2 Conv1D layer into a PyTorch nn.Linear layer requires a transpose (.t()):

   • On non-square matrices (e.g. c_fc where nx=768, nf=3072), nn.Linear(768, 3072) expects [3072, 768], but Conv1D provides [768, 3072]. load_state_dict crashes with a shape mismatch!
   • On square matrices (e.g. c_proj where nx=768, nf=768, or sliced QKV projections [768, 768]), shapes match [768, 768] so no shape error is thrown, but the weight matrix is transposed silently, computing x @ W^T instead of x @ W and destroying numerical parity and model behavior.

   For the fused QKV projection c_attn.weight of shape [nx, 3 * nx] = [768, 2304]: split along dim=-1 into three [768, 768] column slices, then transpose each slice (.t()) before assigning to separate nn.Linear(768, 768) layers (W_q, W_k, W_v).

3. Handle tied output projection. In GPT-2 the lm_head weight is explicitly tied to wte (the token embedding matrix). In PyTorch:

   • nn.Embedding(50257, 768) stores weight of shape [num_embeddings, embedding_dim] = [50257, 768].
   • nn.Linear(768, 50257, bias=False) stores weight of shape [out_features, in_features] = [50257, 768].
   • The forward pass computes x @ weight.T = hidden @ E^T, yielding [B, T, V].
   • Therefore lm_head.weight = wte.weight directly without .t() (both are [50257, 768]).

4. Account for RoPE vs. learned positional encoding. GPT-2 checkpoints include wpe (learned positional embeddings). A model that uses Rotary Positional Embeddings (RoPE) will not have a wpe key; those parameters are simply omitted and the RoPE logic is applied during the forward pass instead. Omitting wpe in favor of RoPE, or changing LayerNorm to RMSNorm, changes the model architecture. You cannot drop wpe and expect logit parity against a standard GPT-2 reference; logit parity requires architectural identity.

5. Verify shape consistency. After mapping, run a forward pass with dummy input and compare output logits (or per-layer hidden states) against a reference implementation. Shape mismatches are the most common failure mode.

Common weight-mapping failure cases

Failure mode	Symptom	Root cause	Fix
Shape mismatch	Runtime error: expected tensor of size [d, d] but got [d, 3d]	Custom hidden dimension d differs from checkpoint, or c_attn was split incorrectly	Verify d matches; if splitting c_attn, slice columns weight[:, :d], weight[:, d:2d], weight[:, 2d:] into three [d, d] matrices
Missing keys	load_state_dict warns missing_key: transformer.wpe.weight	Custom model uses RoPE and has no wpe	Omit RoPE-keys; do not error on missing positional embeddings if using RoPE
Tied/untied mismatch	Dimension error: tied lm_head.weight [V, d] vs. a custom output_proj Linear weight [out, in]	Checkpoint ties lm_head to wte but custom model expects separate matrices	Either tie output_proj to tok_emb.weight or copy the weight and transpose
QKV layout mismatch	Attention produces incorrect results	HF c_attn is fused [d, 3d]; custom model expects three separate [d, d] matrices	Slice column thirds: c_attn.weight[:, :d], c_attn.weight[:, d:2d], c_attn.weight[:, 2d:]
Norm type mismatch	Numerical NaN during forward pass	HF uses LayerNorm with weight+biases; custom uses RMSNorm (weight only)	Not a drop-in fix: LayerNorm and RMSNorm are different architectures, so dropping biases does not preserve GPT-2 behaviour; validate any conversion separately

Traceable example: mapping GPT-2 small into a minimal decoder-only model

HF checkpoint key	Custom model key	Shape transformation
transformer.wte.weight	tok_emb.weight	No change: [50257, 768]
transformer.wpe.weight	(omitted)	RoPE model does not use this key
transformer.h.0.ln_1.weight	blocks.0.norm1.weight	No change: [768]
transformer.h.0.attn.c_attn.weight	blocks.0.attn.q_proj.weight, k_proj.weight, v_proj.weight	Split columns [768, 2304] → weight[:, :768].t(), weight[:, 768:1536].t(), weight[:, 1536:].t() → each [768, 768]
transformer.h.0.attn.c_proj.weight	blocks.0.attn.proj.weight	Transpose: [768, 768] → .t() → [768, 768]
transformer.h.0.mlp.c_fc.weight	blocks.0.mlp.fc_in.weight	Transpose: [768, 3072] → .t() → [3072, 768]
transformer.h.0.mlp.c_proj.weight	blocks.0.mlp.fc_out.weight	Transpose: [3072, 768] → .t() → [768, 3072]
lm_head.weight	lm_head.weight = tok_emb.weight	Tied: nn.Linear(768, 50257).weight = nn.Embedding(50257, 768).weight, both [50257, 768], no .t()

(Assumes GPT-2 small: vocab size V=50257, hidden dimension d=768, 12 layers, 12 heads, head dim=64.)

CPU parity check (run this to verify your mapping)

The following 15-line script loads GPT-2 small from Hugging Face, maps weights into a 1-block replica using the rules above, and asserts logit parity on a fixed prompt. Run it to confirm your mapping is correct before scaling up.

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GPT2LMHeadModel

torch.manual_seed(42)
torch.set_grad_enabled(False)

# 1-block replica matching GPT-2 small (d=768, h=12, d_head=64)
class OneBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(768)
        self.q = nn.Linear(768, 768, bias=True)
        self.k = nn.Linear(768, 768, bias=True)
        self.v = nn.Linear(768, 768, bias=True)
        self.o = nn.Linear(768, 768, bias=True)
        self.ln2 = nn.LayerNorm(768)
        self.fc1 = nn.Linear(768, 3072, bias=True)
        self.fc2 = nn.Linear(3072, 768, bias=True)

    def forward(self, x):
        x = x + self.o(F.scaled_dot_product_attention(
            self.q(self.ln1(x)), self.k(self.ln1(x)), self.v(self.ln1(x)),
            is_causal=True
        ))
        x = x + self.fc2(F.gelu(self.fc1(self.ln2(x))))
        return x

# Load HF GPT-2 small
hf_model = GPT2LMHeadModel.from_pretrained("gpt2")
hf_model.eval()

# Build replica and map weights using the CORRECTED rules
replica = OneBlock()
with torch.no_grad():
    # LayerNorm
    replica.ln1.weight.copy_(hf_model.transformer.h[0].ln_1.weight)
    replica.ln1.bias.copy_(hf_model.transformer.h[0].ln_1.bias)
    replica.ln2.weight.copy_(hf_model.transformer.h[0].ln_2.weight)
    replica.ln2.bias.copy_(hf_model.transformer.h[0].ln_2.bias)
    # Attention: c_attn [768, 2304] split + transpose
    c_attn = hf_model.transformer.h[0].attn.c_attn.weight
    replica.q.weight.copy_(c_attn[:, :768].t())
    replica.k.weight.copy_(c_attn[:, 768:1536].t())
    replica.v.weight.copy_(c_attn[:, 1536:].t())
    replica.q.bias.copy_(hf_model.transformer.h[0].attn.c_attn.bias[:768])
    replica.k.bias.copy_(hf_model.transformer.h[0].attn.c_attn.bias[768:1536])
    replica.v.bias.copy_(hf_model.transformer.h[0].attn.c_attn.bias[1536:])
    # c_proj transpose
    replica.o.weight.copy_(hf_model.transformer.h[0].attn.c_proj.weight.t())
    replica.o.bias.copy_(hf_model.transformer.h[0].attn.c_proj.bias)
    # MLP transpose
    replica.fc1.weight.copy_(hf_model.transformer.h[0].mlp.c_fc.weight.t())
    replica.fc1.bias.copy_(hf_model.transformer.h[0].mlp.c_fc.bias)
    replica.fc2.weight.copy_(hf_model.transformer.h[0].mlp.c_proj.weight.t())
    replica.fc2.bias.copy_(hf_model.transformer.h[0].mlp.c_proj.bias)

# Test parity
x = torch.randn(1, 10, 768)
hf_out = hf_model.transformer.h[0](x)[0]
replica_out = replica(x)
diff = (hf_out - replica_out).abs().max().item()
print(f"Max logit diff: {diff:.3e}")
assert diff < 1e-4, f"Mapping error: diff {diff} exceeds tolerance"
print("PASS: Mapped weights match HF reference within 1e-4")

Note: This tests projection parity for one block, not full checkpoint parity. The full model includes 12 blocks, embeddings, final norm, and tied output projection; parity there requires mapping all components identically.

Decoding: from logits to sampled tokens

Once the model’s weights are mapped and a forward pass produces a vector of logits over the vocabulary, the decoding step turns those scores into a single next token. Four common strategies are greedy sampling, temperature scaling, top-k truncation, and top-p (nucleus) filtering. Each changes the probability distribution differently, and each has reproducible stopping rules.

The traceable example

Use this concrete example to follow each method. The vocabulary has V=6 tokens (ids 0–5), and the model outputs the following raw logits for the next token:

logits = [1.2, 0.8, 0.5, 0.3, 0.1, -0.2]

The softmax probability distribution (at temperature τ=1.0) is approximately:

id	0	1	2	3	4	5
p	0.32	0.21	0.16	0.13	0.11	0.08

We will walk through each decoding method using this distribution.

Greedy decoding

Mechanism. Select the token with the highest probability:

greedy_id = argmax(probs) = 0

Result. Token id 0 is selected deterministically. No randomness is involved, so given the same logits the output is always the same token.

When to use. When you need the most likely next token (e.g. for classification, function calling, or deterministic completion) and randomness is undesired.

Temperature scaling

Mechanism. Divide the logits by a scalar τ > 0 before applying softmax. Lower τ sharpens the distribution; higher τ flattens it.

logits_τ = logits / τ
probs_τ   = softmax(logits_τ)

Example with τ = 0.5. Dividing logits by 0.5 is equivalent to multiplying by 2:

logits / 0.5 = [2.4, 1.6, 1.0, 0.6, 0.2, -0.4]
softmax → probs ≈ [0.49, 0.22, 0.12, 0.08, 0.05, 0.03]

The distribution becomes sharper: token 0 is even more dominant, but greedy still selects id 0.

Example with τ = 1.5. Dividing by 1.5 flattens the distribution:

logits / 1.5 = [0.80, 0.53, 0.33, 0.20, 0.07, -0.13]
softmax → probs ≈ [0.26, 0.20, 0.16, 0.14, 0.13, 0.10]

Now token 0 is less dominant; sampling introduces more variety.

Stopping rule. Temperature scaling does not change the stopping rule; it remains based on the sampled token and sequence length limits (see Stopping rules below).

Top-k decoding

Mechanism. Keep only the k highest-probability tokens, set all others to zero, then renormalize so the remaining probabilities sum to 1.

Example with k = 3. The top-3 ids by logit are 2. After zeroing ids 3–5 and applying softmax (or renormalizing the original softmax):

keep ids {0, 1, 2}, renormalize → probs ≈ [0.46, 0.31, 0.23, 0, 0, 0]

Tokens 3, 4, 5 have zero probability; sampling is weighted among 2.

Stopping rule. Same as greedy: the sampled token is checked against the EOS token and the maximum new-token limit.

Top-p (nucleus) decoding

Mechanism. Sort tokens by probability, then keep the smallest prefix whose cumulative probability reaches or exceeds p. The remaining tokens are zeroed and the kept tokens are renormalized.

Example with p = 0.7. Compute cumulative probability from the softmax distribution:

id 0: 0.32 → cumulative 0.32
id 1: 0.32 + 0.21 = 0.53 → cumulative 0.53
id 2: 0.53 + 0.16 = 0.69
id 3: 0.69 + 0.13 = 0.82 ≥ 0.7, stop

Keep ids 3, renormalize:

renorm sum = 0.32 + 0.21 + 0.16 + 0.13 = 0.82
probs → [0.32/0.82, 0.21/0.82, 0.16/0.82, 0.13/0.82, 0, 0] ≈ [0.39, 0.26, 0.19, 0.16, 0, 0]

Tokens 0, 1, 2, and 3 remain possible; sampling is random among those four.

Stopping rule. Same as the other methods: check the sampled token.

Reproducible stopping rules

After a token is sampled (via any of the above methods), apply these checks before appending it to the generated sequence and continuing the loop:

1. EOS check. If sampled_id == eos_id (e.g. GPT-2’s 50256), halt generation. The EOS token is the universal “stop” signal.
2. Length check. If len(generated) ≥ max_new_tokens, halt generation. The max_new_tokens limit prevents runaway generation and is typically set by the serving system or the user prompt.
3. NaN/Inf guard. If the probability distribution contains NaN or Inf values (e.g. from a weight-mapping error), halt generation and surface an error.

These two checks (EOS + max length) are the only stopping conditions needed for safe, predictable generation.

One visual summary

Decoding pipeline: greedy, temperature, top-k, and top-p

Four decoding methods applied to traceable example logits [1.2, 0.8, 0.5, 0.3, 0.1, -0.2] with V=6.

Read this diagram as text

Top panel: weight mapping from a Hugging Face GPT-2 checkpoint to a custom decoder-only model, showing key renames, shape transposes for Conv1D-to-Linear weights, and tied output projection handling. Bottom panel: the decoding pipeline where raw logits [V] pass through temperature scaling (τ), optional top-k truncation (k=3), optional top-p (nucleus) filtering with p=0.7, multinomial sampling, and a stop-condition check (EOS token or max length). A traceable example uses V=6, logits=[1.2, 0.8, 0.5, 0.3, 0.1, -0.2].

Key takeaways

Key takeaways

• A Hugging Face checkpoint must be mapped (prefix renames, Conv1D transposes, tied-projection handling) before it can be loaded into a custom decoder-only implementation.
• Weight-mapping failure cases include shape mismatches, missing RoPE keys, tied/untied mismatches, QKV layout differences, and norm-type mismatches — each has a concrete fix.
• Greedy decoding selects argmax; temperature scaling sharpens or flattens the distribution; top-k keeps only the k most likely tokens; top-p (nucleus) keeps the smallest prefix whose cumulative probability reaches p.
• Reproducible stopping rules are: (1) halt if sampled_id == eos_id, and (2) halt if len(generated) ≥ max_new_tokens. These apply regardless of which decoding method is used.
• The published lessons how-llm-inference-works.mdx and understanding-the-kv-cache.mdx provide the prerequisite vocabulary for prefill, decode, and the KV cache; this lesson builds on top of that foundation.

Exercises

Try it yourself

1. 01Using the traceable example logits [1.2, 0.8, 0.5, 0.3, 0.1, -0.2] with V=6, compute the softmax probabilities at τ=1.0, τ=0.5, top-k=3, and top-p=0.7. Verify the greedy id, the top-k kept ids, and the top-p nucleus ids match the trace in the diagram.

   Show a hint

   Apply softmax first, then temperature, then top-k truncation+renormalization, then top-p cumulative+renormalization. Use a calculator or short Python script.

2. 02You have a model checkpoint with `transformer.wpe.weight [1024, 768]` but your custom implementation uses RoPE and has no positional embedding key. Describe exactly which keys you would omit, which you would rename, and how you would verify the forward pass produces matching logits.

   Show a hint

   Rename transformer.h.{i} → blocks.{i}. For each Conv1D weight (c_attn slices, c_proj, c_fc, mlp.c_proj), transpose with .t() before loading into nn.Linear. Omit wpe entirely. Note that omitting wpe changes the model architecture; you cannot expect logit parity against a standard GPT-2 reference without architectural identity. Verify forward pass outputs carefully against a reference implementation with the same architecture.

Related lessons

• LLM Internals & Architecture · Inference foundationsHow LLM inference works9 min
• LLM Internals & Architecture · Inference foundationsUnderstanding the KV Cache12 min

Sources and scope

• Hugging Face Transformers GPT-2 source code defines the state_dict key names, Conv1D weight layout, and tied lm_head → wte convention.
• Radford et al., “Language Models are Unsupervised Multitask Learners” (OpenAI Technical Report, 2019) describes the architecture, weight tying, and byte-level BPE tokenisation. PDF | GitHub
• Bubeck et al., “Sparks of AGI” (2023) surveys decoding strategies including greedy, temperature, top-k, and top-p (nucleus) sampling.
• Hugging Face docs: model.generate() implements the decoding loop with temperature, top-k, and top-p support.
• Vaswani et al., “Attention Is All You Need” (2017) introduces the decoder-only autoregressive formulation that these decoding methods operate on.

The code in this lesson is illustrated with a minimal example (V=6, 6 logits) and does not implement a full training loop or production serving logic. Real serving systems may batch decoding, overlap prefill and decode, or use speculative decoding to accelerate generation, but those topics are left to planned lessons orders 12–14 (prefill and decode, KV cache, and context windows).