Assembling a GPT-Style Language Model

A decoder block is a complete unit — it takes [B, T, d] and returns [B, T, d]. Stack N of them, add an embedding table at the front and an output projection at the back, and you have a complete GPT-style language model. This lesson assembles the components from lessons 4–8 into a runnable model, counts its parameters, and verifies the forward pass.

The one-sentence version

A GPT-style model = token embeddings + N stacked decoder blocks (each preserving [B, T, d]) + final normalisation + output projection to [B, T, V] logits, with the output projection either tied to the embedding table or learned separately.

Everything else in this lesson makes that concrete.

Lab: Run the companion CPU exercise at /labs/track1/02-tiny-decoder/ (stack blocks, final norm, weight tying, parameter count).

A consistent shape trace

Use these dimensions throughout the lesson:

Symbol	Meaning	Value in this trace
B	Batch size (sequences processed together)	2
T	Sequence length (tokens per sequence)	5
d	Model width (hidden dimension)	32
h	Number of attention heads	4
d_head	Width per head (d / h)	8
V	Vocabulary size	100
N	Number of decoder blocks	2

The numbers are deliberately small. A real model uses the same pattern with larger values; the pattern is what matters.

1. Model assembly

A GPT-style decoder-only model has four parts:

1. Token embedding table E of shape [V, d] — maps token IDs to d-dim vectors.
2. Position information — learned position table P of shape [max_position, d] (GPT-2 style) or RoPE rotation (0 params, applied inside attention).
3. N decoder blocks — each preserves [B, T, d] through pre-norm attention + SwiGLU MLP with residual connections.
4. Final normalisation + output projection — RMSNorm/LayerNorm, then project [B, T, d] → [B, T, V] logits.

Forward pass

input_ids [B, T]              → look up in E → [B, T, d]
+ position info [B, T, d]     → [B, T, d]
DecoderBlock 1                → [B, T, d]
DecoderBlock 2                → [B, T, d]
... (N blocks total)           → [B, T, d]
Final norm                    → [B, T, d]
Output projection             → [B, T, V] logits

Worked example (B=2, T=5, d=32, V=100, N=2)

input_ids:        [2, 5]
token embeddings: [2, 5, 32]   E lookup
position info:    [2, 5, 32]   added to embeddings
block 1:          [2, 5, 32]
block 2:          [2, 5, 32]
final norm:       [2, 5, 32]
output projection:[2, 5, 100]  logits

2. Weight tying vs untied output projection

Tied (GPT-2 style)

The output projection matrix is the transpose of the embedding table: W_out = E^T. No separate output matrix is stored.

logits = hidden @ E^T        # [B, T, d] × [d, V] → [B, T, V]

• Saves V × d parameters
• No bias term
• Default in Hugging Face GPT2Config (tie_word_embeddings = True)

Untied (Llama 2 style)

A separate output matrix W_out of shape [d, V] is learned independently.

logits = hidden @ W_out      # [B, T, d] × [d, V] → [B, T, V]

• Adds d × V parameters
• No bias term
• Default in Hugging Face LlamaConfig (tie_word_embeddings = False)

Staff/Principal engineering judgment

Question	Recommendation	Rationale
Tie or untie for a new decoder-only model?	Tie (default)	Saves V × d params with negligible quality cost at small-to-medium scale; simpler implementation.
When to untie?	When the output task differs significantly from pre-training (e.g., classification head, specialised vocabulary)	Untied output has extra expressive capacity for the output layer.
Bias in output projection?	No bias	GPT-2 and Llama 2 both omit bias; adding one adds V params with no demonstrated benefit.

Staff/Principal judgment: For a small-to-medium GPT-style model (100M–1B parameters), tie the output projection by default. The parameter savings are real (V × d, e.g., ~205M params for V=50k, d=4k), the implementation is simpler, and the quality difference is typically negligible. Untie only if you have a concrete reason (different output task, vocabulary skew).

3. Parameter counting

For a model with vocabulary V, width d, N decoder blocks:

Component	Parameters (tied)	Parameters (untied)
Token embeddings	V × d	V × d
Position embeddings (learned)	max_position × d	max_position × d
Position embeddings (RoPE)	0	0
Per decoder block	10 × d² + 6d	10 × d² + 6d
All N blocks	N × (10 × d² + 6d)	N × (10 × d² + 6d)
Final norm	d	d
Output projection	0 (tied)	d × V
Total	V×d + N×(10d²+6d) + d	2V×d + N×(10d²+6d) + d

Per-block breakdown (~10 × d² + 6d)

Sub-component	Parameters
Attention Q/K/V/O (bias=True)	4 × d² + 4d
MLP combined gate+value (bias=False)	d × 2d = 2d²
MLP output projection (bias=False)	2d × d = 2d²
RMSNorm × 2	2 × d
Total per block	10d² + 6d

Worked parameter count (V=100, d=32, N=2, RoPE, tied)

Component	Calculation	Parameters
Token embeddings	100 × 32	3,200
Position info (RoPE)	0	0
Decoder blocks (2 × (10 × 32² + 6 × 32))	2 × (10,240 + 192)	20,864
Final norm	32	32
Output projection (tied)	0	0
Total		24,096

Same model untied: 24,096 + (32 × 100) = 27,296 (3,200 more params).

4. Runnable model

import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x):
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight


class SwiGLU(nn.Module):
    def __init__(self, d, hidden_mult=4):
        super().__init__()
        hidden = d * hidden_mult
        self.w = nn.Linear(d, hidden, bias=False)
        self.w_out = nn.Linear(hidden // 2, d, bias=False)

    def forward(self, x):
        x_proj = self.w(x)              # [B, T, 4d]
        gate, val = x_proj.chunk(2, dim=-1)  # each [B, T, 2d]
        x_swish = F.silu(gate) * val    # [B, T, 2d]
        return self.w_out(x_swish)      # [B, T, d]


class DecoderBlock(nn.Module):
    def __init__(self, d, n_heads=4):
        super().__init__()
        self.norm1 = RMSNorm(d)
        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
        self.norm2 = RMSNorm(d)
        self.mlp = SwiGLU(d)

    def forward(self, x, attn_mask=None):
        x_norm = self.norm1(x)
        T = x_norm.size(1)
        if attn_mask is None:
            attn_mask = torch.zeros(T, T, device=x.device, dtype=x.dtype)
        attn_out, _ = self.attn(x_norm, x_norm, x_norm,
                                attn_mask=attn_mask, need_weights=False,
                                is_causal=True)
        x = x + attn_out
        x = x + self.mlp(self.norm2(x))
        return x


class GPTStyleModel(nn.Module):
    def __init__(self, vocab_size, d, n_heads, n_blocks, max_position=1024,
                 tie_weights=True, use_rope=True):
        super().__init__()
        self.d = d
        self.vocab_size = vocab_size
        self.tie_weights = tie_weights

        # Token embeddings
        self.tok_emb = nn.Embedding(vocab_size, d)

        # Position info: learned table or RoPE (0 params)
        if use_rope:
            self.pos_emb = None  # RoPE applied inside attention, 0 params
        else:
            self.pos_emb = nn.Embedding(max_position, d)

        # Decoder blocks
        self.blocks = nn.ModuleList([
            DecoderBlock(d, n_heads) for _ in range(n_blocks)
        ])

        # Final norm
        self.final_norm = RMSNorm(d)

        # Output projection
        if tie_weights:
            # Reuse embedding table transposed: no separate params
            self.output_proj = None
        else:
            self.output_proj = nn.Linear(d, vocab_size, bias=False)

    def forward(self, input_ids):
        B, T = input_ids.shape

        # Token embeddings [B, T, d]
        h = self.tok_emb(input_ids)

        # Position info: RoPE applied inside attention blocks (0 params)
        # Learned positions would be added here: h = h + self.pos_emb(positions)

        # Decoder blocks (RoPE applied inside each block's attention)
        for block in self.blocks:
            h = block(h)  # [B, T, d] → [B, T, d]

        # Final norm [B, T, d]
        h = self.final_norm(h)

        # Output projection [B, T, V]
        if self.tie_weights:
            logits = h @ self.tok_emb.weight.T  # [B, T, d] × [d, V]
        else:
            logits = self.output_proj(h)  # [B, T, d] × [d, V]

        return logits


# Test
torch.manual_seed(42)
model = GPTStyleModel(vocab_size=100, d=32, n_heads=4, n_blocks=2,
                      tie_weights=True, use_rope=True)
input_ids = torch.randint(0, 100, (2, 5))  # [B=2, T=5]
logits = model(input_ids)
print(f"Input: {input_ids.shape}, Logits: {logits.shape}")  # [2, 5, 100]
assert logits.shape == (2, 5, 100), "Output shape mismatch"

# Count parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params:,}")
# Expected: 100×32 + 2×(10×32² + 6×32) + 32 = 3,200 + 20,864 + 32 = 24,096
assert total_params == 24_096, f"Expected 24,096 params, got {total_params}"
print("PASS: Parameter count matches formula (tied, RoPE)")

5. Component verification

Before assembling the full model, verify each component independently:

Component	Test	Expected
RMSNorm	Match nn.RMSNorm on random input [2, 16, 512]	max_diff < 1e-6
SwiGLU	Check intermediate shapes (gate/val split, output)	gate/val each [B, T, 2d], output [B, T, d]
DecoderBlock	Shape preservation [B, T, d] → [B, T, d]	Output shape == input shape
DecoderBlock	Causal safety: corrupt future position, earlier outputs unchanged	allclose(atol=1e-5)
Full model	Forward pass [B, T] → [B, T, V]	Logits shape matches vocab
Full model	Parameter count matches formula	Exact match

6. Visual and text equivalent

GPT-style model assembly

Token embeddings → N stacked decoder blocks → final norm → output projection (tied/untied), with tensor shapes at each stage.

Read this diagram as text

Text equivalent. Panel A: input token IDs [B, T] look up in embedding table E [V, d] → [B, T, d]; position info added (RoPE or learned table). Panel B: N stacked decoder blocks, each preserving [B, T, d] through pre-norm attention + SwiGLU MLP with residuals. Panel C: final RMSNorm [B, T, d], then output projection — tied (reuses E^T, 0 extra params) or untied (separate W [d, V]). Panel D: parameter count formula with breakdown by component (embeddings, blocks, output).

Text equivalent: The diagram has four panels. Panel A shows token IDs [B, T] looked up in the embedding table E [V, d] to produce [B, T, d], with position information added (either RoPE rotation applied inside attention, or a learned position table). Panel B shows N stacked decoder blocks, each taking [B, T, d] and returning [B, T, d] through pre-norm attention + SwiGLU MLP with residual connections. Panel C shows final RMSNorm [B, T, d], then the output projection — either tied (reuses E^T, zero extra parameters) or untied (separate matrix W [d, V]). Panel D shows the parameter counting formula broken down by component: token embeddings V×d, position info (0 for RoPE), N blocks N×(10d²+6d), final norm d, output projection (0 tied or d×V untied).

7. Sources

Source	Relevance
Radford et al., “Language Models are Unsupervised Multitask Learners” (GPT-2 report, 2019)	GPT-2 architecture: token embeddings, N decoder blocks, final norm, tied output projection
GPT-2 model source: src/model.py	Tied output projection (wte reused), no bias, learned positional embeddings, pre-norm
Llama 2 model source: llama/model.py	Untied output projection, no bias, RoPE, RMSNorm, SwiGLU, pre-norm
Vaswani et al., “Attention Is All You Need” (arXiv:1706.03762)	Original Transformer architecture
Lesson 8 (transformer-blocks.mdx): decoder block implementation	Pre-norm, RMSNorm, SwiGLU, residual connections, is_causal=True
Lesson 7 (positional-encoding-and-rope.mdx): RoPE mechanism	RoPE applied to Q/K inside attention, 0 params, shape preserved
Lesson 5 (causal-attention-masks.mdx): causal mask construction	Causal masking in decoder blocks

---