Architecture choices: derive and test a 100M decoder

Architecture is a set of budgets, not a universal recipe

Goal: choose and implement normalization, positions, attention heads, feedforward width and weight sharing; defend their parameter, cache and numerical costs. Prerequisites are decoder attention, residual connections and next-token cross-entropy. This lesson supplies its own runnable decoder instead of assuming an undefined GroupedQueryAttention from another page.

GPT-2 is already pre-LayerNorm, uses learned absolute positions, a four-times-width GELU FFN and tied input/output embeddings. The 2017 Transformer is post-LN and uses sinusoidal positions. Confusing those baselines produces a false list of seven historical changes. Original LLaMA uses pre-RMSNorm, RoPE, SwiGLU and untied embeddings; GQA is model/version-dependent, not a property of every Llama checkpoint. Llama 2 7B uses MHA; its larger GQA models, Llama 3 and Mistral illustrate different choices. A model card/config, not the family nickname, decides the implementation.

These are seven decision axes, not seven improvements shared by every modern model. Qwen, Gemma, DeepSeek, Mistral and Llama have version-specific exceptions, including bias, sharing, local attention, soft-capping, MoE and latent attention. Do not infer unpublished GPT-4 internals from a trend.

Pre-norm and the residual highway

Post-norm computes LN(x + F(x)); pre-norm computes x + F(LN(x)). Their Jacobians are respectively J_LN (I + J_F) and I + J_F J_LN. Pre-norm provides an additive identity term, but the full gradient is not unchanged: branch terms can amplify or cancel it. Post-norm is not intrinsically untrainable above 12 or 24 layers. Xiong et al. analyze initialization and warmup for specified architectures; DeepNorm deliberately revisits deep normalized residual networks.

The highway analogy helps only if it includes the branches: the residual stream is repeatedly added to, not frozen. Final normalization constrains the scale presented to the logit head; omitting it changes the architecture and can affect stability, but no fabricated personal loss improvement is evidence. Compare matched initializations, data and optimizer schedules, not just one first/last gradient ratio. A shallow successful run does not settle behavior at depth.

RMSNorm: derive the reduction before timing it

For a token vector of width d, LayerNorm subtracts its feature mean and divides by sqrt(mean((x-mean(x))**2)+eps), then applies scale and usually bias. RMSNorm uses

It removes re-centering; it does not require the residual stream to be zero-centered. For x=(3,4), epsilon zero and unit scale, r=sqrt(12.5) and the output is approximately (0.848528,1.131371). LayerNorm instead gives approximately (-1,1). Thus the functions are not interchangeable even before training. The derivative is diag(gamma)/r - (gamma*x) xᵀ/(d*r³); epsilon prevents division by zero but also changes exact scale invariance.

The RMSNorm paper reports savings in its experiments, not a guaranteed 15% whole-model speedup or equal quality on every GPU. Reduction fusion, memory traffic, dtype and the fraction of time spent in norms matter. For low-precision input, compute square/mean in FP32 explicitly; a large individual square can overflow FP16 regardless of width. FP64 tests should retain FP64 rather than silently downcasting. Our implementation does both. BF16 has a wider exponent range than FP16, not immunity to numerical failure.

RoPE: relative rotations, not automatic extrapolation

For a two-coordinate pair, use the rotation R(pθ). Orthogonality gives

The positional part depends on the offset; the score still depends on the content vectors q and k, so it is not a function of distance alone. Apply rotations to Q and K, not V. Even head dimension and consistent pair ordering are part of the checkpoint contract. At cached decode, the position offset must advance; resetting it to zero silently changes scores.

Being mathematically defined beyond the training context does not prove useful extrapolation. Position interpolation, YaRN and LongRoPE alter the frequency/position contract and require evaluation, often continued training. Sinusoidal positions are also defined beyond the training length. Neither method removes attention's compute or memory costs.

GQA: count the cache separately from the projections

With H query heads, K KV heads and head width h=d/H, Q has d² weights, K and V each have d*K*h, and the output projection has d². Total attention weights are 2d²+2dKh. Require H % K == 0. MHA is K=H; MQA is K=1; GQA is intermediate. MQA reduces KV cache by H-fold, not by 66% (the latter confuses projection-parameter arithmetic with cache storage).

At batch B, sequence T, L layers and b bytes per cache element:

For H=32, K=8, h=128 and BF16, each token/layer needs 4096 bytes. Across 32 layers that is 131072 bytes = 128 KiB, not 131 KiB. MHA would require four times as many cache bytes. It does not follow that user concurrency or throughput rises fourfold: weights, workspaces, activations, scheduling and kernels also constrain serving. Low-batch decoding can be bandwidth-bound; batching/prefill can change the bottleneck.

Our reference repeats KV heads explicitly for transparency. repeat_interleave materializes data here; it is not promised to be a view. Store compact K/V in a real cache and use a supported GQA kernel when optimizing. PyTorch 2.9 SDPA documents enable_gqa constraints; no fused-kernel/GPU performance is demonstrated in this lesson. Reducing only Q heads while retaining K/V does not shrink this cache, but changing the whole head layout can.

SwiGLU, biases and embedding sharing

A bias-free GELU FFN with hidden width 4d has d*4d + 4d*d = 8d² weights. SwiGLU computes W_down(silu(W_gate x) * W_up x) and has 3dm weights. Equating them gives m=8d/3; rounding to a hardware multiple changes the exact count. Equal weight count gives approximately equal dense-matmul FLOPs, not an automatic 5% extra training compute merely because there are three matrices.

The multiplicative gate gives input-dependent interactions; this is not a proof that ordinary MLPs cannot approximate those functions. Shazeer's GLU paper supports its reported comparisons, not a universal one-point WikiText improvement. Bias removal saves the sum of output widths of affected layers; tying saves exactly Vd when both matrices would otherwise be independent. Neither guarantees unchanged quality. Tying is a constraint coupling input and output roles, not a mathematical inverse operation.

Complete local reference decoder

Save this exact block as architecture.py for the later local recipes. Python 3.12 / torch 2.9 CPU is the tested runtime. It implements one full causal pass, not a KV-cache serving engine. Batched generation marks each row finished at its first newly generated EOS and pads that row with EOS while others continue; trim each output at that first EOS when decoding text. Prompt-contained EOS does not mark a row finished. Context exhaustion and new_tokens remain independent stopping bounds. Top-k sampling uses a score cutoff, so tied scores can admit more than k tokens. No pretrained weights or external data are required.

import math
from dataclasses import dataclass, asdict
import torch
from torch import nn
import torch.nn.functional as F
 
@dataclass(frozen=True)
class Config:
    vocab: int = 32000
    width: int = 768
    layers: int = 12
    heads: int = 12
    kv_heads: int = 4
    hidden: int = 2048
    context: int = 2048
    theta: float = 10000.0
    eps: float = 1e-6
    tied: bool = True
 
    def validate(self):
        assert min(self.vocab, self.width, self.layers, self.heads,
                   self.kv_heads, self.hidden, self.context) > 0
        assert self.width % self.heads == 0
        assert self.heads % self.kv_heads == 0
        assert (self.width // self.heads) % 2 == 0
        assert self.theta > 0 and self.eps > 0
 
def parameter_count(c):
    c.validate()
    h = c.width // c.heads
    block = 2*c.width**2 + 2*c.width*c.kv_heads*h
    block += 3*c.width*c.hidden + 2*c.width
    return c.vocab*c.width*(1 if c.tied else 2) + c.layers*block + c.width
 
class RMSNorm(nn.Module):
    def __init__(self, width, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.eps = eps
    def forward(self, x):
        z = x.float() if x.dtype in (torch.float16, torch.bfloat16) else x
        z = z * torch.rsqrt(z.square().mean(-1, keepdim=True) + self.eps)
        return z.to(x.dtype) * self.weight.to(x.dtype)
 
def rope(x, theta=10000.0, offset=0):
    # x: batch, heads, time, head-width; adjacent coordinate pairing
    d = x.shape[-1]
    assert d % 2 == 0 and offset >= 0
    dtype = torch.float64 if x.dtype == torch.float64 else torch.float32
    freq = theta ** (-torch.arange(0, d, 2, device=x.device, dtype=dtype)/d)
    pos = torch.arange(offset, offset+x.shape[-2], device=x.device, dtype=dtype)
    angle = pos[:, None]*freq[None, :]
    a, b = x[..., 0::2], x[..., 1::2]
    co, si = angle.cos().to(x.dtype), angle.sin().to(x.dtype)
    return torch.stack((a*co-b*si, a*si+b*co), dim=-1).flatten(-2)
 
class Attention(nn.Module):
    def __init__(self, c):
        super().__init__(); self.c = c
        h = c.width//c.heads
        self.q = nn.Linear(c.width, c.width, bias=False)
        self.k = nn.Linear(c.width, c.kv_heads*h, bias=False)
        self.v = nn.Linear(c.width, c.kv_heads*h, bias=False)
        self.o = nn.Linear(c.width, c.width, bias=False)
    def forward(self, x):
        c = self.c; b, t, _ = x.shape; h = c.width//c.heads
        q = self.q(x).view(b,t,c.heads,h).transpose(1,2)
        k = self.k(x).view(b,t,c.kv_heads,h).transpose(1,2)
        v = self.v(x).view(b,t,c.kv_heads,h).transpose(1,2)
        q, k = rope(q,c.theta), rope(k,c.theta)
        k = k.repeat_interleave(c.heads//c.kv_heads, dim=1)
        v = v.repeat_interleave(c.heads//c.kv_heads, dim=1)
        z = F.scaled_dot_product_attention(q,k,v,is_causal=True,dropout_p=0.0)
        return self.o(z.transpose(1,2).contiguous().view(b,t,c.width))
 
class Block(nn.Module):
    def __init__(self,c):
        super().__init__()
        self.n1=RMSNorm(c.width,c.eps); self.n2=RMSNorm(c.width,c.eps)
        self.attn=Attention(c)
        self.gate=nn.Linear(c.width,c.hidden,bias=False)
        self.up=nn.Linear(c.width,c.hidden,bias=False)
        self.down=nn.Linear(c.hidden,c.width,bias=False)
    def forward(self,x):
        x=x+self.attn(self.n1(x)); z=self.n2(x)
        return x+self.down(F.silu(self.gate(z))*self.up(z))
 
class Decoder(nn.Module):
    def __init__(self,c):
        super().__init__(); c.validate(); self.cfg=c
        self.embed=nn.Embedding(c.vocab,c.width)
        self.blocks=nn.ModuleList([Block(c) for _ in range(c.layers)])
        self.norm=RMSNorm(c.width,c.eps)
        self.head=nn.Linear(c.width,c.vocab,bias=False)
        self.apply(self._init)
        if c.tied: self.head.weight=self.embed.weight
        for block in self.blocks:
            nn.init.normal_(block.attn.o.weight,std=0.02/math.sqrt(2*c.layers))
            nn.init.normal_(block.down.weight,std=0.02/math.sqrt(2*c.layers))
    @staticmethod
    def _init(m):
        if isinstance(m,(nn.Linear,nn.Embedding)): nn.init.normal_(m.weight,std=0.02)
    def forward(self,ids):
        if ids.ndim!=2 or ids.dtype!=torch.long: raise ValueError('B,T long IDs required')
        if not 0 < ids.shape[1] <= self.cfg.context: raise ValueError('context length')
        if torch.any(ids<0) or torch.any(ids>=self.cfg.vocab): raise ValueError('token ID')
        x=self.embed(ids)
        for block in self.blocks: x=block(x)
        return self.head(self.norm(x))
    @torch.no_grad()
    def generate(self,ids,new_tokens,temperature=0.0,top_k=None,eos=None):
        if new_tokens<0 or temperature<0: raise ValueError('generation arguments')
        if eos is not None and (not isinstance(eos,int) or not 0<=eos<self.cfg.vocab):
            raise ValueError('eos must be a vocabulary ID')
        old=self.training; self.eval()
        finished=torch.zeros(ids.shape[0],dtype=torch.bool,device=ids.device)
        try:
            for _ in range(new_tokens):
                if ids.shape[1]>=self.cfg.context: break
                logits=self(ids)[:,-1,:].float()
                if temperature==0: nxt=logits.argmax(-1,keepdim=True)
                else:
                    logits=logits/temperature
                    if top_k is not None:
                        if not 1<=top_k<=self.cfg.vocab: raise ValueError('top_k')
                        cut=logits.topk(top_k).values[:,-1:]
                        logits=logits.masked_fill(logits<cut,float('-inf'))
                    nxt=torch.multinomial(logits.softmax(-1),1)
                if eos is not None:
                    nxt=torch.where(finished[:,None],torch.full_like(nxt,eos),nxt)
                    finished=finished | (nxt[:,0]==eos)
                ids=torch.cat((ids,nxt),dim=1)
                if eos is not None and bool(finished.all()): break
            return ids
        finally: self.train(old)
 
if __name__ == '__main__':
    torch.set_num_threads(1); torch.manual_seed(7)
    c=Config(vocab=17,width=16,layers=2,heads=4,kv_heads=2,hidden=40,context=16)
    model=Decoder(c).double(); x=torch.tensor([[1,2,3,4],[4,3,2,1]])
    logits=model(x); changed=x.clone(); changed[:,-1]=6
    torch.testing.assert_close(logits[:,:-1],model(changed)[:,:-1],rtol=0,atol=0)
    F.cross_entropy(logits[:,:-1].reshape(-1,c.vocab),x[:,1:].reshape(-1)).backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())
    assert sum(p.numel() for p in model.parameters())==parameter_count(c)
    assert parameter_count(Config())==100092672
    assert model.head.weight is model.embed.weight
    assert model.generate(x[:1,:2],3).shape==(1,5)
    print('decoder causal/gradient/count checks passed; 100M design:',parameter_count(Config()))

The 100M arithmetic and ablation plan

For V=32000, d=768, L=12, H=12, K=4 and m=2048: embedding 24576000; QKV 983040 per layer; attention output 589824; FFN 4718592; two norm scales 1536; final norm 768. Total 100092672 unique parameters. Decayed 2D tensors total 100073472 and norm scales total 19200 under the later dimension-based optimizer policy. Tied weights must appear once, not twice.

Vary K through divisors of 12 (1,2,3,4,6,12), not arbitrary invalid head counts. Cache is linear in K; attention projection count also changes, while attention score work still has H query heads. Increasing hidden width from 2048 to 3072 adds 12*3*768*1024=28311552 weights. To compare 8/12/16 layers at a similar total budget, solve the full count including Vd and rounding; blindly scaling width by square-root depth misses embeddings and KV layout.

An unexecuted hardware ablation should use identical batches and token counts, separately time warmed-up forward/backward/update with device synchronization, record actual peak memory and several seeds, and compare held-out NLL. A 200-step norm timing run measures timing, not equal converged quality. Do not pre-fill a 12–18% speedup into the results.

Alternatives and the modern architecture survey

Sliding-window and sparse attention: Mistral-style local windows bound each layer's direct attention range; stacking layers can propagate further. Longformer and BigBird change the attention graph. RoPE scaling changes positions, not the graph or quadratic dense attention work, so it does not replace sparsity. A full 2K-context teaching model avoids these extra policies for simplicity, not because they are obsolete.

MLA: DeepSeek caches a compressed KV latent plus a separate rotary key component. The decoupled position channel avoids forcing all rotary structure through the low-rank content projection; matrix absorption can avoid reconstructing an expanded cache during decode. Do not transfer DeepSeek-V2's reported 93.3% cache reduction to V3 with an invented denominator. A hypothetical V3 cache of 512 latent + 64 rotary values is 576 values/token/layer; comparison to any MHA baseline must specify both K and V dimensions, head count, dtype and replication. The arithmetic and the paper's benchmark are different claims.

MoE: a shared dense path plus top-k routed FFNs separates total from active parameters. DeepSeek-V3 describes 256 routed experts, one shared expert and eight routed selections, alongside 671B total/37B activated parameters. Its balancing bias affects selection without being the same as a differentiable balancing penalty; the report still includes a complementary sequence-wise balance loss. It does not prove routing collapse impossible. Its reported 2.788M H800-hours and USD5.576M estimate exclude prior research/ablation costs and use an assumed hourly price; they are not a current rental quote or a like-for-like H100 comparison. No universal 7B threshold determines whether small MoE is worthwhile.

MTP: predicting future tokens adds auxiliary supervised objectives and extra training work. Meta's parallel future heads and DeepSeek's sequential MTP modules are not identical. DeepSeek conditions each depth on appropriate prior representations/token embeddings; one extra linear head is an educational alternative, not its exact implementation. Optional speculative decoding must verify drafts with the target distribution; acceptance and speed depend on workload. Auxiliary labels are true future tokens, so calling every MTP objective self-distillation is misleading.

Long context: YaRN modifies RoPE scaling; LongRoPE searches nonuniform position rescaling; LongRoPE2 addresses its own long-context training/evaluation setting. The LongRoPE2 paper reports extension to 128K with 10B tokens in its described setting, not 800B tokens for a universal million-token guarantee. Ring Attention distributes blockwise attention work and communication; it is not evidence that a particular commercial 10M window uses it. Llama 4 Scout's advertised window, Gemini and Claude product limits belong to exact dated model/API cards, not a common architectural maximum. Retrieval/needle scores alone do not establish reliable reasoning across the full window.

Beyond pure attention: Mamba-2 uses structured state-space duality; Jamba and Zamba interleave state-space and attention modules; Griffin combines gated linear recurrence, local attention and MLPs. These are distinct mechanisms, not a single unshipped future category. BitNet b1.58 uses ternary weights with higher-precision activation/training details; the public 2B4T release is a concrete model, not proof that every future checkpoint will be ternary. Titans studies test-time neural memory, distinct from both a fixed KV cache and externally retrieved documents. Open weights aid inspection, but are neither necessary nor sufficient for sound engineering.

Worked recall and design answers

1. Seven changes? Correct the premise: placement and tying already exist in GPT-2; use the decision table, and verify each actual checkpoint rather than reciting a false universal recipe.

2. RMSNorm formula and speed? Derive r and y above. Fewer reductions can help, but whole-step speed and quality require a specified experiment; there is no transferable 15% answer.

3. KV bytes? 2*8*128*2=4096 bytes/layer/token; 32 layers gives 128 KiB. Fourfold cache savings do not imply fourfold user capacity.

4. Why 8d/3? Solve 3dm=8d². Width rounding and bias choices must be included in the exact count. No guaranteed quality gain follows from equality.

5. MLA compression? State what is cached and name the baseline before quoting a percentage. A 512-dimensional latent is not the entire cache if a rotary key is also stored.

6. Design a 1B-active MoE: write N_active=N_shared+k*N_expert+N_nonexpert and N_total=N_shared+E*N_expert+N_nonexpert. For an illustrative 200M nonexpert/shared budget and eight 100M active experts, active count is 1B; E=32 makes total 3.4B. Then budget router/all-to-all, capacity, expert utilization and shared paths. This is a count example, not evidence that 32/top-8 is optimal.

7. Legacy pre-LN versus post-LN? Pre-LN uses x+F(LN(x)), preserving an additive identity term in its Jacobian. Post-LN uses LN(x+F(x)). Branch terms still change gradients, and depth alone does not prove one is trainable or the other impossible.

8. Legacy RMSNorm formula and speed? gamma*x/sqrt(mean(x*x)+eps) removes re-centering. Compare actual reduction precision and measured whole-step time; the claimed universal 15% gain and equal quality are not established.

9. Legacy GQA definition and quality? Share K/V heads among H/K query heads; require divisibility and keep head width fixed. Cache reduction is H/K relative to MHA under the same B,T,L,dtype. Quality and throughput changes require measurement, not a fixed loss penalty.

10. Legacy SwiGLU versus GELU? Multiply the activated gate and up projections before the down projection. Three d*m matrices match the two GELU matrices at m=8d/3; this is a parameter-budget equality, not a promised one-point quality improvement.

11. Legacy tying justification? One shared V*d matrix saves exactly Vd independent parameters. It couples token representations and output scores; they are not mathematical inverse functions, and unchanged quality is not guaranteed.

For the norm/first-last-gradient exercise, compare all trainable gradient norms and subsequent held-out behavior under matched runs rather than requiring a predetermined collapse.

Resources and execution boundary

Primary references: GPT-2 source, LLaMA, RMSNorm, GLU variants, GQA, RoFormer, DeepSeek-V3, LongRoPE2. Exact passages, access times and raw hashes are in the batch evidence, including the other named survey papers.

The bounded decoder/count/causality/gradient checks execute on CPU. No 100M training, GPU timing, large-model quality, KV-cache kernel or long-context benchmark was run. Optional original videos and course navigation are learning resources, not evidence of technical correctness. Review tomorrow: redraw the residual block, recompute its parameters and cache, and justify one deviation numerically.

Inherited optional resources and navigation

These pre-read links are optional further reading; their timestamps, video claims and optional pacing are not accepted technical evidence.

• LLaMA explained — architecture walkthrough
• Karpathy — 1hr Intro to LLMs
• LLaMA: Open and Efficient Foundation Language Models
• GLU Variants Improve Transformer

Previous: ← DL S049 · Next: DL S051 →