MoE, optimisers, decoding — every line, every shape.

§0How to read this page

Almost every confusing line in ML code is a shape change: a block of numbers gets split, stacked, reordered or multiplied. If you can say the shape before and after each line, you can explain the forward pass. So this page writes a shape next to every line, using the letters below.

When a shape is written (B, L, D), read it as "B sequences, each of L tokens, each token a list of D numbers". In Python terms it is a list of B lists of L lists of D floats. The shape traces in §5 were printed by running the real src/model.py with B = 2 and L = 6, so they are not guesses.

§1The tensor toolkit: every strange function name

A tensor is a rectangular block of numbers: a Python nested list where every inner list has the same length. PyTorch stores it as one flat run of numbers in memory plus a shape that says how to read it. Every function below either builds a tensor, changes how it is read (its shape), or does arithmetic on it. Each one has a tiny input and output, copied from a real run.

§1.1 · Making tensors

§1.2 · Reading a shape: dim, -1, keepdim

x.shape is a tuple like (2, 6, 512). x.size(0) is its first entry, x.dim() is how many entries it has (3), x.numel() is their product (6,144 numbers). An axis (also called a dim) is one position in that tuple. Negative numbers count from the end, exactly like Python list indexing, so dim=-1 always means the last axis.

A reduction such as sum, mean, max or softmax works along one axis. x.sum(-1) on a (2, 3) tensor adds the 3 numbers in each row and returns shape (2,): the axis you reduce over disappears. With keepdim=True it stays as a length-1 axis, giving (2, 1). You keep it so the result can later be divided back into the original, as in RMSNorm (§3.2).

x = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.]])      # shape (2, 3)
x.sum(-1)                             # tensor([ 6., 15.])        shape (2,)
x.sum(-1, keepdim=True)               # tensor([[ 6.], [15.]])    shape (2, 1)
x / x.sum(-1, keepdim=True)           # each row divided by its own sum -> rows add to 1

§1.3 · view and reshape: same numbers, new shape

The numbers sit in memory in row-major order: the last index changes fastest, like reading a book left to right, then down. view keeps that flat run and only changes how it is cut into rows. The product of the new shape must equal the old one. One entry may be -1, meaning "work this out for me".

x = torch.arange(12)                  # [0, 1, 2, ..., 11]        shape (12,)
x.view(3, 4)                          # [[ 0,  1,  2,  3],
                                      #  [ 4,  5,  6,  7],
                                      #  [ 8,  9, 10, 11]]        shape (3, 4)
x.view(2, -1).shape                   # (2, 6)  -- 12 / 2 = 6 was inferred

reshape does the same thing, but it also works when the numbers are no longer laid out in row-major order (after a transpose, below); then it quietly makes a copy. view refuses and raises an error in that case. That is the only difference.

flatten(0, 1) merges axes 0 and 1: (B, L, V) becomes (B·L, V). It is the same as reshape(B*L, V). Losses use it so that every token is one row.

§1.4 · unsqueeze, squeeze, [None]: adding and removing size-1 axes

x.unsqueeze(i) inserts a new axis of length 1 at position i. No numbers change; only the shape gains a 1. x.squeeze(i) removes a length-1 axis. Indexing with None is shorthand for the same thing: x[None] is x.unsqueeze(0) and x[:, None] is x.unsqueeze(1).

t = torch.arange(3)                   # [0, 1, 2]                 shape (3,)
t.unsqueeze(0)                        # [[0, 1, 2]]               shape (1, 3)   one row
t.unsqueeze(1)                        # [[0], [1], [2]]           shape (3, 1)   one column
t[:, None]                            # same as unsqueeze(1)

Why would anyone want a length-1 axis? Because of broadcasting, the next rule. A length-1 axis is a promise: "use this same value for every position along this axis".

§1.5 · Broadcasting: how tensors of different shapes combine

When you add or multiply two tensors of different shapes, PyTorch lines the shapes up from the right. For each pair of axes, the sizes must be equal, or one of them must be 1. A size-1 axis is stretched (without copying) to match the other. A missing axis on the left counts as 1.

§1.6 · expand vs repeat vs repeat_interleave

expand is broadcasting done on purpose: it makes a size-1 (or missing) axis look bigger without copying anything. repeat and repeat_interleave really copy.

torch.arange(4).expand(2, 4)          # [[0, 1, 2, 3],
                                      #  [0, 1, 2, 3]]   no copy: both rows are the same memory
torch.arange(2).repeat_interleave(3)  # [0, 0, 0, 1, 1, 1]   each element copied 3 times in a row

model.py uses torch.arange(start, start + L).expand(B, L) to give every sequence in the batch the same position numbers 0…L−1. Beam search uses arange(B).repeat_interleave(W) to say "row 0 of the prompt batch becomes rows 0…W−1 of the beam batch".

§1.7 · permute and transpose: reordering axes

x.transpose(1, 2) swaps axes 1 and 2. x.permute(2, 0, 3, 1, 4) is a general reorder: the new axis 0 is the old axis 2, the new axis 1 is the old axis 0, and so on. This really changes which number is at which index, unlike view, which never moves a number relative to the flat order.

a = torch.arange(6).view(2, 3)        # [[0, 1, 2],
                                      #  [3, 4, 5]]
a.transpose(0, 1)                     # [[0, 3], [1, 4], [2, 5]]   rows became columns
a.view(3, 2)                          # [[0, 1], [2, 3], [4, 5]]   NOT the same thing

§1.8 · chunk, cat, stack: splitting and joining

a = torch.arange(8).view(2, 4)        # [[0, 1, 2, 3], [4, 5, 6, 7]]
x1, x2 = a.chunk(2, dim=-1)           # x1 = [[0, 1], [4, 5]]   x2 = [[2, 3], [6, 7]]
torch.cat([x1, x2], dim=-1)           # back to [[0, 1, 2, 3], [4, 5, 6, 7]]
torch.stack([x1, x2])                 # shape (2, 2, 2): a NEW axis in front

cat joins along an axis that already exists. stack creates a new axis. RoPE uses chunk to split each head vector into two halves and cat to glue the rotated halves back.

§1.9 · Indexing with tensors

Plain slicing works like Python lists, with one slice per axis:

• logits[:, -1] on (B, L, V) takes the last position of every sequence: result (B, V). The indexed axis disappears.
• seqs[:, -1:] keeps the axis: result (B, 1).
• cache.k[layer][:, :, 5:9] = k writes into positions 5…8 of the third axis.

Indexing with an integer tensor replaces every integer with the row it names. If cos has shape (257, 32) and pos has shape (B, L) of integers, then cos[pos] has shape (B, L, 32): each position number became its row of 32 cosines. This is also exactly what an embedding table does with token ids.

Indexing with a boolean tensor keeps only the entries where it is True and flattens them into one axis. topi[sel] with topi of shape (B, L, k) and sel of shape (B, L) gives (count of True, k).

Comparisons such as topi == e or logits < kth return boolean tensors of the same shape. &, | and ~ are element-wise AND, OR and NOT on them.

§1.10 · The remaining names, alphabetically

§1.11 · The trailing underscore, and the optimiser helpers

A method ending in _ changes the tensor in place instead of returning a new one: p.mul_(0.9) means p *= 0.9. The optimisers use these because they update big weight tensors every step and do not want copies. Three of them combine several operations:

§1.12 · dtype, device, gradients

• dtype is the number format. float32 (fp32) is the normal float. float16 (fp16) and bfloat16 (bf16) use half the memory and run faster on GPUs, but fp16's largest value is 65,504 and its smallest positive normal value is about 6·10⁻⁵, so big sums and tiny gradients go wrong. .float() converts to fp32, .to(dtype) to anything, .type_as(x) to whatever x is, .long() to 64-bit integers.
• device is where the tensor lives: "cpu" or "cuda" (the GPU). Tensors must be on the same device to be combined. .to(dev) moves one.
• gradients: PyTorch records every operation so that loss.backward() can compute, for every weight p, how much the loss would change if p changed. The answer is stored in p.grad, with the same shape as p. @torch.no_grad() turns the recording off (for evaluation and inside optimisers). .detach() gives a copy of a value with no history. float(t) or t.item() turns a one-number tensor into a Python float, .tolist() into a Python list.

§1.13 · The nn building blocks

§1.14 · Worked example from A1: mask = mp.view(B*nb, 1, 1, W)

This line comes from the A1 BLT local encoder. With §1.3–§1.5 you can now read it. The code first pads the byte sequence up to a multiple of the window size W, then cuts it into windows and treats every window as its own little sequence:

xp = _pad_to_multiple(x, W)        # (B, nb*W, D)   padded bytes
nb = xp.size(1) // W               # number of windows per sequence
xb = xp.view(B * nb, W, D)         # (B*nb, W, D)   each window is now one "sequence" in the batch
mask = mp.view(B * nb, 1, 1, W)    # (B*nb, 1, 1, W) padding flag for each key in each window

Take B = 2 sequences of 300 bytes and W = 128. Padding gives 384 = 3 × 128, so nb = 3. The view turns (2, 384, D) into (6, 128, D). Because of row-major order, rows 0–2 are the three windows of sequence 0 and rows 3–5 are the three windows of sequence 1. The padding mask mp has shape (2, 384), one flag per byte, and the same view cuts it the same way.

Inside attention the score tensor has shape (B·nb, H, W, W) = (6, H, 128, 128): for each window, each head, each query byte, each key byte. The mask is (6, 1, 1, 128). Line them up from the right:

So the two 1s say: "if key byte t of window j is padding, no head and no query may look at it". In that A1 code True meant "hide", which is why the causal mask was OR-ed in with mask | causal_mask(W): a (128, 128) causal mask broadcasts against (6, 1, 1, 128) to (6, 1, 128, 128), hiding a key if it is padding or in the future. Finally out.view(B, nb * W, D)[:, :L] glues the windows back into sequences and cuts off the padding. (A2 uses the opposite convention, True = "may attend", because that is what F.scaled_dot_product_attention expects.)

§2The only maths you need

Five ideas cover everything in this repo: the dot product, matrix multiplication, softmax, cross-entropy and gradient descent. Each is short.

§2.1 · Dot product

Multiply matching entries and add. It is large and positive when the two vectors point the same way, near 0 when they are unrelated, negative when they point opposite ways. Attention is "take the dot product of a query with every key and see which keys score high".

§2.2 · Matrix multiplication and nn.Linear

A @ B with A of shape (n, m) and B of shape (m, p) gives (n, p): entry (i, j) is the dot product of row i of A with column j of B. The inner sizes must match and disappear. With more axes, @ treats all but the last two as batch axes: (B, H, L, hd) @ (B, H, hd, L) → (B, H, L, L), one L×L matrix per sequence per head.

A linear layer is one matrix multiplication. nn.Linear(512, 1536) stores W of shape (1536, 512) and computes $y=x{W}^{\top }$. Each of the 1536 outputs is a dot product of the input with one row of W, so it is a learned weighted sum of the 512 inputs. It has 512 × 1536 = 786,432 parameters. Training means finding good values for those numbers.

§2.3 · Softmax, log-softmax, logsumexp

Softmax turns any list of real numbers ("logits") into probabilities: all positive, summing to 1, and the biggest logit gets the biggest share. Adding the same constant to every logit changes nothing, so implementations subtract the maximum first to avoid overflow. log_softmax is $\log \text{softmax}(z)$ computed stably, and $\text{logsumexp}(z)=\log {\sum }_j{e}^{z_j}$ is the log of the denominator. Setting a logit to −∞ gives it probability exactly 0, which is how every mask and filter in this repo works.

Temperature divides the logits first: $\text{softmax}(z/T)$. T < 1 makes the distribution sharper, T > 1 flatter.

§2.4 · Cross-entropy and perplexity

The model outputs, at every position, logits over the V tokens. The loss for that position is minus the log of the probability it gave to the correct next token:

If the model gives the right token probability 1 the loss is 0. If it gives 0.01 the loss is 4.6. Perplexity is $e^{\text{mean loss}}$: a perplexity of 3.4 means the model is, on average, as unsure as if it were choosing uniformly among 3.4 tokens. F.cross_entropy(logits, y) takes logits of shape (N, V) and targets (N,). A target of -100 is skipped (ignore_index=-100), which is how padding and prompt tokens are left out of the loss.

§2.5 · Gradient descent

loss.backward() fills p.grad with $\partial ℒ/\partial p$ for every weight. That gradient points in the direction that increases the loss fastest, so stepping against it lowers the loss. $\eta$ is the learning rate. Every optimiser in Part 2 is a smarter version of this one line. Gradients add up across backward() calls until opt.zero_grad() clears them, which is what makes gradient accumulation over micro-batches work.

GELU, the non-linearity in the MLP, is $\text{GELU}(x)=x\Phi (x)$ where Φ is the standard normal CDF. It is about $x$ for large positive x, about 0 for large negative x, and smooth in between. Without a non-linearity, two linear layers in a row would collapse into one.

§3src/model.py, line by line

One file holds the model for Parts 1 and 2. It is a decoder-only transformer, the GPT design: token ids go in, and for every position the model outputs a score for every possible next token. Translation in Part 1 is done by writing the source sentence, a separator, and then letting the model continue with the English.

§3.1 · TransformerConfig

@dataclass
class TransformerConfig:
    vocab_size: int = 24000
    d_model: int = 512  # Dimension of the model
    n_ctx: int = 256  # Number of context tokens
    n_layers: int = 8
    n_heads: int = 8
    dropout: float = 0.0
    rope_base: float = 10000.0
    tie_embeddings: bool = True
    use_moe: bool = False  # Whether to use MoE or MLP
    d_ff: int = 2048  # hidden size of the dense MLP
    n_experts: int = 4  # routed experts
    top_k: int = 1  # active routed experts per token
    n_shared: int = 0  # always-active shared experts
    d_expert: int = 512  # hidden size of every expert (routed and shared)
    aux_loss_coef: float = 0.01  # Switch-style load-balancing loss
    z_loss_coef: float = 1e-3  # router z-loss (ST-MoE)

Plain settings, overwritten by the YAML configs. n_ctx is the longest sequence the model is built for. rope_base sets the RoPE frequencies (§4). The MoE fields only matter when use_moe is True.

§3.2 · RMSNorm

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        xf = x.float()
        xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
        return xf.type_as(x) * self.weight

It rescales each token vector so its root-mean-square is 1, then multiplies by a learned gain g. This keeps the numbers in a stable range as they pass through 8 layers.

Compared with LayerNorm it does not subtract the mean and has no bias, so it is cheaper and has D parameters instead of 2D.

§3.3 · KVCache

class KVCache:
    def __init__(self, n_layers, batch, n_heads, max_len, head_dim, device, dtype):
        shape = (batch, n_heads, max_len, head_dim)
        self.k = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(n_layers)]
        self.v = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(n_layers)]
        self.len = 0

    def update(self, layer: int, k: torch.Tensor, v: torch.Tensor):
        L = k.size(2)
        self.k[layer][:, :, self.len : self.len + L] = k
        self.v[layer][:, :, self.len : self.len + L] = v
        end = self.len + L
        return self.k[layer][:, :, :end], self.v[layer][:, :, :end]

When generating text one token at a time, the keys and values of earlier tokens never change (a token can only look backwards). So instead of re-running the whole sequence for every new token, the model stores them. Each layer gets one key buffer and one value buffer of shape (B, H, max_len, hd). update writes the new keys (B, H, L_new, hd) into positions len … len+L_new−1 and returns everything stored so far, (B, H, len+L_new, hd). self.len is advanced by Transformer.forward after all layers have written, because every layer writes to the same positions.

The saving: without a cache, generating T tokens costs 1 + 2 + … + T token-passes, about T²/2. With it, each step processes one token, so T in total (attention still reads all cached keys).

§3.4 · Attention

class Attention(nn.Module):
    def __init__(self, cfg: TransformerConfig, rope: Rotary):
        super().__init__()
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
        self.o_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        self.rope = rope
        self.dropout = cfg.dropout

    def forward(self, x, pos, mask=None, cache: KVCache | None = None, layer: int = 0):
        B, L, D = x.shape
        q, k, v = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k = self.rope(q, pos), self.rope(k, pos)
        if cache is not None:
            k, v = cache.update(layer, k, v)
        if mask is None:
            out = F.scaled_dot_product_attention(
                q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0
            )
        else:
            out = F.scaled_dot_product_attention(
                q, k, v, attn_mask=mask, dropout_p=self.dropout if self.training else 0.0
            )
        return self.o_proj(out.transpose(1, 2).reshape(B, L, D))

Each token makes three vectors from its own vector: a query ("what am I looking for"), a key ("what do I contain") and a value ("what I hand over if chosen"). Each token's output is a weighted average of the values of the tokens it may see, weighted by how well its query matches their keys:

M is 0 where attending is allowed and −∞ where it is not (future tokens, padding). With $d_h=64$, the division by 8 matters: the dot product of two random 64-long vectors with unit variance entries has variance 64, and without scaling the softmax would put almost all weight on one key, which gives tiny gradients. The work is split across H = 8 heads, each using its own 64 of the 512 numbers, so each head can look for a different kind of relation.

F.scaled_dot_product_attention (SDPA) is PyTorch's fused version of the formula above. The assignment allows it. is_causal=True builds the "no looking at the future" mask itself (used in training, where there is no padding on the left). When a mask is passed it must be boolean with True = "may attend" and broadcastable to (B, H, L, T). The attention layer has 512·1536 + 512·512 = 1,048,576 parameters.

§3.5 · Block: one transformer layer

class Block(nn.Module):
    def __init__(self, cfg: TransformerConfig, rope: Rotary, layer: int):
        super().__init__()
        from src.part1.ffn import MLP, MoE  # local import keeps model.py standalone

        self.layer = layer
        self.norm1 = RMSNorm(cfg.d_model)
        self.attn = Attention(cfg, rope)
        self.norm2 = RMSNorm(cfg.d_model)
        self.ffn = MoE(cfg) if cfg.use_moe else MLP(cfg.d_model, cfg.d_ff)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x, pos, mask=None, cache=None, token_mask=None):
        x = x + self.drop(self.attn(self.norm1(x), pos, mask, cache, self.layer))
        return x + self.drop(self.ffn(self.norm2(x), token_mask))

This is a pre-norm residual block. Each sub-layer reads a normalised copy of the running vector and adds its result to it instead of replacing it. Two consequences: the original information is never lost, and during backpropagation the gradient flows straight through the x + path to early layers. Attention moves information between tokens. The FFN (dense MLP or MoE) then processes each token on its own. All shapes stay (B, L, 512). token_mask is only used by the MoE to leave padding out of its statistics.

§3.6 · Transformer

class Transformer(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = cfg = config
        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        rope = Rotary(cfg.d_model // cfg.n_heads, cfg.n_ctx + 1, cfg.rope_base)
        self.blocks = nn.ModuleList([Block(cfg, rope, i) for i in range(cfg.n_layers)])
        self.norm = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight
        self.apply(self._init)
        # scaled init for residual projections (GPT-2)
        for n, p in self.named_parameters():
            if n.endswith("o_proj.weight") or n.endswith("fc2.weight"):
                nn.init.normal_(p, 0.0, 0.02 / (2 * cfg.n_layers) ** 0.5)

def forward(self, idx, pos=None, mask=None, cache=None, token_mask=None):
    # Input shape: (B, L) token ids -> Output shape: (B, L, N_vocab) logits
    B, L = idx.shape
    if pos is None:
        start = cache.len if cache is not None else 0
        pos = torch.arange(start, start + L, device=idx.device).expand(B, L)
    x = self.embed(idx)
    for blk in self.blocks:
        x = blk(x, pos, mask, cache, token_mask)
    if cache is not None:
        cache.len += L
    return self.lm_head(self.norm(x))

The output at position t is the model's guess for token t+1, having seen tokens 0…t only (the causal mask). So one forward pass over a sequence gives L predictions at once, which is why training is parallel.

def moe_layers(self):
    return [b.ffn for b in self.blocks if hasattr(b.ffn, "router")]

def aux_loss(self) -> torch.Tensor:
    losses = [m.last_aux_loss for m in self.moe_layers() if m.last_aux_loss is not None]
    if not losses:
        return torch.zeros((), device=self.embed.weight.device)
    return torch.stack(losses).sum()

def param_counts(self) -> dict:
    total = sum(p.numel() for p in self.parameters())
    active = total
    for m in self.moe_layers():
        per_expert = sum(p.numel() for p in m.experts[0].parameters())
        active -= per_expert * (len(m.experts) - m.top_k)
    ffn = sum(p.numel() for b in self.blocks for p in b.ffn.parameters())
    return {"total": total, "active": active, "ffn_total": ffn}

aux_loss adds the MoE balancing losses from all 8 layers into one number (a zero-dimensional tensor, shape ()) that training adds to the main loss. param_counts counts every weight once (parameters() yields the tied matrix only once) and, for the "active" count, subtracts the experts a token does not use: E − k of them per layer. For the dense model: embedding 12,288,000 + 8 × (attention 1,048,576 + MLP 2,097,152 + two norms 1,024) + final norm 512 = 37,462,528.

§3.7 · greedy_generate: batched decoding with a cache

Used for BLEU in Parts 1 and 2. Prompts have different lengths, and the batch must be a rectangle, so short prompts are padded. They are padded on the left, so that the last column is the last real token of every prompt, and that is where the next-token prediction is read from.

B = len(prompts)
P = max(len(p) for p in prompts)
ids = torch.full((B, P), pad_id, dtype=torch.long, device=device)
keep = torch.zeros((B, P), dtype=torch.bool, device=device)
for i, p in enumerate(prompts):
    ids[i, P - len(p) :] = torch.tensor(p, device=device)
    keep[i, P - len(p) :] = True
total = P + max_new
key_keep = torch.zeros((B, total), dtype=torch.bool, device=device)
key_keep[:, :P] = keep
pos = (keep.long().cumsum(-1) - 1).clamp(min=0)
cache = model.new_cache(B, total, dtype=amp_dtype() if device.type == "cuda" else None)
causal = torch.ones(P, P, dtype=torch.bool, device=device).tril()
mask = (causal[None] & keep[:, None, :]).unsqueeze(1)  # (B,1,P,P)
mask = mask | torch.eye(P, dtype=torch.bool, device=device)[None, None]
logits = model(ids, pos, mask, cache)[:, -1]
next_pos = pos[:, -1] + 1

Worked example with two prompts, [7, 8] and [1, 2, 3, 4], so B = 2 and P = 4:

For prompt 0 the mask before and after the eye fix is (rows = queries, columns = keys):

before            after | eye
[[0, 0, 0, 0],    [[1, 0, 0, 0],     row 0 (a pad) had nothing to look at
 [0, 0, 0, 0],     [0, 1, 0, 0],     row 1 (a pad) likewise
 [0, 0, 1, 0],     [0, 0, 1, 0],     token 7 sees itself
 [0, 0, 1, 1]]     [0, 0, 1, 1]]     token 8 sees 7 and itself

A row with no allowed key would be a softmax over all −∞, which is 0/0 = NaN, and the NaN would spread. Letting pad rows see themselves avoids that. Their outputs are never read. The prefill call returns (B, P, V) logits and [:, -1] keeps the last column, (B, V).

out = torch.full((B, max_new), pad_id, dtype=torch.long, device=device)
done = torch.zeros(B, dtype=torch.bool, device=device)
for t in range(max_new):
    tok = logits.float().argmax(-1)
    tok = torch.where(done, torch.full_like(tok, pad_id), tok)
    out[:, t] = tok
    done |= tok == eos_id
    if done.all():
        break
    key_keep[:, P + t] = True
    m = key_keep[:, : P + t + 1][:, None, None, :]
    logits = model(tok[:, None], next_pos[:, None], m, cache)[:, -1]
    next_pos = next_pos + 1

No causal mask is needed in the loop: the single new query is the latest token, so every cached key is in its past. At the end, each row is cut at its first EOS.

§4Rotary: what RoPE does to q and k

Attention has no sense of word order on its own. RoPE adds it by rotating each query and key vector by an angle that grows with the token's position. It splits each 64-number head vector into 32 pairs, feature i with feature i + 32, and rotates pair i by the angle m·θ_i at position m, where θ_i = 10000^−2i/64. Values are not rotated, and nothing is learned.

Those two formulas are exactly the last line of forward below. The table follows the shapes.

class Rotary(nn.Module):
    def __init__(self, head_dim: int, max_len: int, base: float = 10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
        t = torch.arange(max_len).float()
        freqs = torch.outer(t, inv_freq)  # (max_len, hd/2)
        self.register_buffer("cos", freqs.cos(), persistent=False)
        self.register_buffer("sin", freqs.sin(), persistent=False)

    def forward(self, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
        # x: (B, H, L, hd); pos: (B, L) integer positions
        cos = self.cos[pos].unsqueeze(1).to(x.dtype)  # (B, 1, L, hd/2)
        sin = self.sin[pos].unsqueeze(1).to(x.dtype)
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)

The last two lines are literally the two formulas above, done for all 32 pairs, all heads, all positions and all sequences in one go. pos is usually just 0…L−1, but in left-padded generation it comes from the cumsum trick (§3.7), so the first real token is always position 0.

§5One forward pass, every shape

This is the answer to "explain the forward pass". The shapes were printed by hooks on the real model (Part 1 config, top-2 MoE, B = 2 sequences of L = 6 tokens). Learn this table and you can explain Parts 1 and 2.

§6Part 1 · Mixture of Experts

A Mixture of Experts replaces the one big feed-forward network in each block with several smaller ones ("experts") plus a tiny router that picks, for each token, which k experts to use. The model can hold many parameters while each token only pays for k experts' worth of compute.

§6.1 · MLP: the dense baseline (and every expert)

class MLP(nn.Module):
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff, bias=False)
        self.fc2 = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x, token_mask=None):
        return self.fc2(F.gelu(self.fc1(x)))

Shapes: (B, L, 512) → fc1 → (B, L, 2048) → GELU → fc2 → (B, L, 512). Each token is processed independently. Parameters: 2 × 512 × 2048 = 2,097,152 per layer. token_mask is accepted and ignored so that MLP and MoE can be swapped without changing the caller.

§6.2 · MoE.forward, line by line

def forward(self, x, token_mask=None):
    B, L, D = x.shape
    xf = x.reshape(-1, D)
    logits = self.router(xf).float()  # (N, E)
    probs = logits.softmax(-1)
    topv, topi = probs.topk(self.top_k, dim=-1)  # (N, k)
    if self.top_k > 1:
        topv = topv / topv.sum(-1, keepdim=True)
    gates = topv.to(x.dtype)

    out = torch.zeros_like(xf)
    for e, expert in enumerate(self.experts):
        tok, slot = (topi == e).nonzero(as_tuple=True)
        if tok.numel() == 0:
            continue
        y = expert(xf[tok]) * gates[tok, slot, None]
        out.index_add_(0, tok, y)
    for s in self.shared:
        out = out + s(xf)

A worked routing example, with 3 tokens, top-2 and 4 experts:

topi = [[2, 0],        # token 0 chose experts 2 and 0
        [1, 2],        # token 1 chose 1 and 2
        [0, 3]]        # token 2 chose 0 and 3
(topi == 2)                     -> [[T, F], [F, T], [F, F]]
.nonzero(as_tuple=True)         -> tok = [0, 1], slot = [0, 1]
# expert 2 runs on xf[[0, 1]]; token 0 is scaled by gates[0, 0], token 1 by gates[1, 1]

There is no capacity limit here: every token always reaches its k experts, none are dropped. The loop over experts runs each expert once on a variable-size group, which is simple and correct but not as fast as a fused kernel.

§6.3 · The auxiliary losses

m = token_mask.reshape(-1) if token_mask is not None else torch.ones(B * L, dtype=torch.bool, device=x.device)
mf = m.float()
n = mf.sum().clamp(min=1.0)
assign = F.one_hot(topi, self.n_experts).sum(1).float()  # (N, E) 0/1 per selected expert
f = (assign * mf[:, None]).sum(0) / (n * self.top_k)  # fraction of routed slots per expert
P = (probs * mf[:, None]).sum(0) / n  # mean router prob per expert
lb = self.n_experts * (f * P).sum()
z = ((torch.logsumexp(logits, -1) ** 2) * mf).sum() / n
self.last_aux_loss = self.aux_coef * lb + self.z_coef * z
self.last_stats = {"load": f.detach(), "lb_loss": lb.detach(), "z_loss": z.detach()}
self.last_topi = topi.view(B, L, self.top_k).detach() if self.record_routing else None
return out.view(B, L, D)

Left alone, a router tends to collapse: one expert gets slightly better, so it gets more tokens, so it gets better still, and the other experts starve. The load-balancing loss pushes against that:

$f_e$ is the share of routing slots expert e actually received, $P_e$ the average probability the router gave it. When both are uniform (1/E each) the loss is E · E · (1/E²) = 1, its minimum. f is a count, so it has no gradient. P does, and the product means "lower the probability of experts that are already overloaded".

The z-loss (from ST-MoE) penalises large router logits. Softmax does not care about the size of its inputs, only their differences, so without this the logits can drift to large values where fp16/bf16 round-off becomes significant.

§6.4 · Parameter matching: where the numbers come from

One expert with d_expert = 512 has 2 × 512 × 512 = 524,288 weights, exactly a quarter of the dense MLP. So 4 such experts have the same total as the dense MLP, plus a tiny router (512 × E). The counts below were printed by param_counts().

Variants 1–4 match on total parameters (the differences are the routers: 8 layers × 512 × 4 = 16,384). Variant 5 doubles each expert, so two active experts equal one dense MLP: active parameters match instead. Active = total − 8 × (E − k) × one expert's size.

§6.5 · data.py: turning translation into next-token prediction

A decoder-only model can only continue text. So each training example is one sequence: <vi> source tokens <sep> English tokens <eos> (or <ja>). At test time the model is given everything up to <sep> and continues. The tokenizer is one byte-level BPE with 24,000 tokens trained on all three languages. Byte-level means any Unicode text, including Japanese, can be encoded without an "unknown" token. The special tokens get ids 0–4: PAD, EOS, SEP, VI, JA.

def collate(self, idx):
    L = int(self.seq_len[idx].max()) - 1
    B = len(idx)
    x = torch.full((B, L), PAD, dtype=torch.long)
    y = torch.full((B, L), -100, dtype=torch.long)
    role = torch.full((B, L), -1, dtype=torch.int8)
    lang = torch.empty(B, dtype=torch.long)
    for r, i in enumerate(idx):
        lid, s, t = self.sample(i)
        seq = [lid, *s, SEP, *t, EOS]
        n = len(seq) - 1
        x[r, :n] = torch.tensor(seq[:-1])
        ns = len(s) + 2  # tokens up to and including <sep>
        y[r, ns - 1 : n] = torch.tensor(seq[ns:])  # predict tgt + eos
        role[r, : ns] = 0
        role[r, ns:n] = 1
        lang[r] = lid
    return x, y, role >= 0, lang, role

Worked example: Vietnamese (id 3), source tokens [10, 11], English tokens [20, 21]:

seq   = [3, 10, 11, 2, 20, 21, 1]     # <vi> src src <sep> tgt tgt <eos>      n = 6
x     = [3, 10, 11, 2, 20, 21]        # input: everything but the last token
ns    = 4                              # <vi>, 10, 11, <sep>
y     = [-100, -100, -100, 20, 21, 1] # y[3:6] = seq[4:]; position t is trained to predict x's token t+1
role  = [0, 0, 0, 0, 1, 1]            # 0 = source side, 1 = target side, -1 = padding

Input position 3 is <sep>, and its label is 20, the first English token. The source side has label −100, so the loss ignores it: the model is not graded on predicting Vietnamese. Padding columns keep x = PAD, y = −100 and role = −1. role >= 0 is the token mask passed to the MoE.

batches builds batches by a token budget instead of a fixed count: it shuffles, sorts by length inside large chunks, and packs sequences until longest × count would exceed 32,768. Sequences of similar length end up together, so little space is wasted on padding.

§6.6 · train.py: one training step

decay = [p for n, p in model.named_parameters() if p.dim() >= 2]
no_decay = [p for n, p in model.named_parameters() if p.dim() < 2]
opt = torch.optim.AdamW(
    [{"params": decay, "weight_decay": t["weight_decay"]}, {"params": no_decay, "weight_decay": 0.0}],
    lr=t["lr"], betas=tuple(t.get("betas", (0.9, 0.95))), fused=dev.type == "cuda")

Weight decay (shrinking weights a little every step) is applied to matrices only, not to the 1-D RMSNorm gains. Part 1 may use PyTorch's AdamW; Part 2 writes its own. The step itself:

prog = int(1e6 * seen / total_tokens)
lr_mult = cosine_lr(prog, 1_000_000, int(1e6 * t.get("warmup_frac", 0.02)), t.get("min_lr_ratio", 0.1))
for g in opt.param_groups:
    g["lr"] = t["lr"] * lr_mult
opt.zero_grad(set_to_none=True)
chunks = [c for c in np.array_split(np.asarray(bidx), min(micro, len(bidx))) if len(c)]
n_tgt = int(sum(train_set.tgt_len[c].sum() + len(c) for c in chunks))
for c in chunks:
    x, y, tm, lang, role = train_set.collate(c)
    with torch.autocast(dev.type, dtype=amp, enabled=dev.type == "cuda"):
        logits = fwd(x, token_mask=tm)
        ce = torch.nn.functional.cross_entropy(logits.flatten(0, 1).float(), y.flatten(), ignore_index=-100, reduction="sum") / n_tgt
        aux = model.aux_loss() / len(chunks)
    scaler.backward(ce + aux)
gn, finite = scaler.unscale_clip(model.parameters(), t.get("grad_clip", 1.0))
if finite:
    opt.step()
scaler.update(finite)

Training stops after 150M non-pad tokens for every variant, so the comparison is fair on data seen. moe_log records each expert's load and a normalised entropy $-{\sum }_e{f}_e\log f_e/\log E$ (1 = perfectly even, 0 = one expert takes everything).

§6.7 · evaluate.py: perplexity, BLEU and routing counts

tl = torch.nn.functional.cross_entropy(
    logits.float().transpose(1, 2), y, ignore_index=-100, reduction="none"
).sum(1)  # (B,)
nt = (y != -100).sum(1)

F.cross_entropy expects the class axis in position 1, so (B, L, V) is transposed to (B, V, L); with reduction="none" it returns one loss per position, (B, L), and .sum(1) gives one total per sentence (B,). Sentences are then summed separately by language, and perplexity is $\exp (\sum \text{nll}/\sum \text{tokens})$ per language.

for li, m in enumerate(moe):
    topi = m.last_topi  # (B, L, k)
    for li_lang, lid in enumerate(LANGS):
        for r in (0, 1):
            sel = (role_d == r) & (lang_d[:, None] == lid)
            ids = topi[sel].flatten()
            counts[li, li_lang, r] += torch.bincount(ids, minlength=m.n_experts).cpu().numpy()

translate sorts test sentences by source length (less padding), builds prompts [lang, *src, SEP], runs greedy_generate, and decodes the ids back to text. BLEU (from sacrebleu) compares the translations with the references by counting matching 1- to 4-word sequences, $\text{BLEU}=\text{BP}·\exp \left(\frac{1}{4}{\sum }_{n=1}^4\log p_n\right)$, where $p_n$ is n-gram precision and BP penalises translations that are too short.

In analyze.py, _fractions does counts.transpose(0, 2, 1, 3).reshape(L, 4, E): swap the language and side axes, then merge them, giving four rows per layer in the order vi-src, ja-src, vi-tgt, ja-tgt, each normalised to sum to 1. Those rows are the heatmaps. The "routing separation" plot uses the total variation distance $\frac{1}{2}{\sum }_e|p_e-q_e|$ between two rows (0 = identical routing, 1 = no expert in common). The rest of the file is matplotlib.

§7Part 2 · Five optimisers from scratch

An optimiser is the rule that turns gradients into weight changes. Plain gradient descent ($w\leftarrow w-\eta g$) is noisy and treats every weight the same. Each optimiser below fixes that differently. All of them subclass torch.optim.Optimizer only for its bookkeeping: self.param_groups (lists of weights with their own lr and weight decay) and self.state[p] (a dict per weight for its running averages). Every step is under @torch.no_grad() so the updates themselves are not recorded for backpropagation.

§7.1 · The data: fixed windows of a token stream

def window_batch(self, split: str, idx: np.ndarray):
    a = self.streams[split]
    L = self.n_ctx
    x = np.stack([a[i * L : i * L + L + 1] for i in idx]).astype(np.int64)
    x = torch.from_numpy(x)
    return x[:, :-1], x[:, 1:]

All training documents are tokenised and joined into one long array, separated by <eos>. Window i is tokens 512·i … 512·i + 512, 513 tokens. np.stack makes (B, 513). The input is the first 512, x[:, :-1], and the target is the same window shifted by one, x[:, 1:], both (B, 512): for window [a, b, c, d] the model sees a, b, c and must predict b, c, d. Documents are split 90/5/5 by their source id so that the human text and every LLM rewrite of it land in the same split (no leakage). For BLEU, the prompt is the last 128 tokens of a test document's first human chunk and the reference is the first 64 tokens of its second chunk.

train.py is the same shape as Part 1 with a fixed batch: 32 windows × 512 = 16,384 tokens per step, run as 2 micro-batches of 16, each loss divided by 2. set_lr multiplies each group's base_lr by the cosine schedule. It evaluates every 0.1× of the dataset. The LR sweep trains on 0.2× of the data at several learning rates, and the final run takes the one with the lowest validation loss (lr: best).

§7.2 · AdamW

g = p.grad
st = self.state[p]
if not st:
    st["step"] = 0
    st["m"] = torch.zeros_like(p)
    st["v"] = torch.zeros_like(p)
st["step"] += 1
t = st["step"]
m, v = st["m"], st["v"]
m.lerp_(g, 1 - b1)
v.mul_(b2).addcmul_(g, g, value=1 - b2)
denom = (v / (1 - b2**t)).sqrt_().add_(eps)
p.mul_(1 - lr * wd)
p.addcdiv_(m, denom, value=-lr / (1 - b1**t))

Dividing by $\sqrt{\widehat{v}}$ gives each weight its own step size: weights with consistently large gradients take smaller steps, and weights with tiny gradients take larger ones. Every update is roughly of size η whatever the gradient's scale.

§7.3 · MARS (variance-reduced AdamW)

then AdamW on $c_t$ in place of $g_t$.

if p.dim() >= 2 or group["optimize_1d"]:
    c = g + gamma * (b1 / (1 - b1)) * (g - st["prev_g"]) if t > 1 else g.clone()
    norm = c.norm()
    if norm > 1:
        c.div_(norm)
    st["prev_g"].copy_(g)
else:
    c = g
m.lerp_(c, 1 - b1)
v.mul_(b2).addcmul_(c, c, value=1 - b2)
...   # identical to AdamW from here

Momentum averages gradients from old weights, so it lags behind. $g_t-g_{t-1}$ estimates how the gradient has changed since last step, and adding a scaled copy of it corrects that lag, which reduces the variance of the update. With γ = 0.025 and β₁ = 0.95 the factor is 0.025 × 19 = 0.475. The corrected gradient is then clipped to length 1 per tensor (c.norm() is its L2 length). This "approximate" MARS reuses last step's gradient instead of computing a second one. The correction only applies to matrices; the 1-D norm gains use c = g. The extra prev_g buffer means 3× model memory.

§7.4 · Lion (memory-efficient)

u = torch.lerp(m, g, 1 - b1).sign_()
p.mul_(1 - lr * wd)
p.add_(u, alpha=-lr)
m.lerp_(g, 1 - b2)

Only one buffer (m), so 1× model memory instead of 2×. torch.lerp (without underscore) makes a new tensor mixing momentum and the current gradient, and sign_ replaces each entry with +1, −1 or 0. So every weight moves by exactly ±η. The direction comes from a mix with β₁ = 0.9, but the stored momentum is updated with the slower β₂ = 0.99. Because every step has full size, Lion needs a smaller learning rate than AdamW (3·10⁻⁴ here) and a larger weight decay (0.3).

§7.5 · Muon (matrix-based)

The other optimisers treat a weight matrix as a bag of separate numbers. Muon treats it as a matrix. Any matrix can be written as $US{V}^{\top }$ (the singular value decomposition): U and V are rotations and S is a diagonal of "stretch amounts". Gradient matrices tend to have a few huge singular values and many tiny ones, so a normal step mostly moves the weights along a few directions. Muon replaces S with the identity: the update $U{V}^{\top }$ has the same directions but moves equally along all of them.

def newton_schulz5(G: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor:
    a, b, c = 3.4445, -4.7750, 2.0315
    half = torch.bfloat16 if (not G.is_cuda or torch.cuda.get_device_capability()[0] >= 8) else torch.float16
    X = (G / (G.norm() + eps)).to(half)
    transposed = X.size(0) > X.size(1)
    if transposed:
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X
    if transposed:
        X = X.T
    return X.to(G.dtype)

Computing an SVD every step is slow, so Newton–Schulz does it with matrix multiplications. Why it works: if $X=US{V}^{\top }$ then $X{X}^{\top }=U{S}^2{U}^{\top }$, so

U and V are untouched. Each singular value σ just goes through the polynomial $\phi (\sigma )=a\sigma +b{\sigma }^3+c{\sigma }^5$. The constants are chosen so that five applications push every σ in (0, 1] to roughly 0.7–1.1, close enough to 1 in practice.

        buf.mul_(mu).add_(g)
        u = g.add(buf, alpha=mu) if group["nesterov"] else buf
        o = newton_schulz5(u, group["ns_steps"])
        scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
        p.mul_(1 - lr * wd)
        p.add_(o, alpha=-lr * scale)
self.adamw.step()

buf is momentum ($b=\mu b+g$, μ = 0.95). Nesterov momentum uses $g+\mu b$, a look-ahead version. The scale √(rows/cols) (√3 for the 1536 × 512 qkv matrix) keeps tall matrices' updates the right size. Only 2-D hidden matrices use Muon. build_optimizer sends the embedding (which is also the tied LM head) and all 1-D gains to an internal copy of the from-scratch AdamW with its own learning rate (3·10⁻³), because an embedding row is updated only when its token appears and orthogonalising the whole table makes no sense. zero_grad is overridden to clear both.

§7.6 · Sophia-G (Hessian approximation, bonus)

def update_hessian(self):
    ...
            st["h"].mul_(b2).addcmul_(p.grad, p.grad, value=1 - b2)

def step(self, closure=None, bs: int = 1):
    ...
            m.lerp_(p.grad, 1 - b1)
            p.mul_(1 - lr * wd)
            ratio = (m.abs() / (rho * bs * h + eps)).clamp_(max=1.0)
            p.addcmul_(m.sign(), ratio, value=-lr)

Newton's method divides the gradient by the curvature (the second derivative, h): take big steps where the loss surface is flat and small steps where it is sharply curved. Sophia keeps a running estimate of the diagonal of the curvature matrix, one number per weight, and steps by $m/h$. Where h is tiny that ratio could be huge, so it is clipped: clamp_(max=1) means no weight ever moves more than η. m.sign() * ratio is the same as $m/(\rho bh)$ clipped to [−1, 1].

The curvature estimate comes from the training loop, every 10 steps:

    with torch.no_grad():
        y_hat = torch.distributions.Categorical(logits=logits.float()).sample()
    loss_s = F.cross_entropy(logits.float().flatten(0, 1), y_hat.flatten()) / micro
    scaler.backward(loss_s)
if scaler.unscale(model.parameters()):
    opt.update_hessian()

This is the Gauss–Newton–Bartlett estimator. Instead of the real next tokens, sample fake labels y_hat from the model's own predicted distribution (Categorical(...).sample() draws one token per position, shape (B, L)), compute the loss on those, and backpropagate. The resulting gradient $\widehat{g}$ satisfies $𝔼[\widehat{g}{\widehat{g}}^{\top }]$ = the Gauss–Newton matrix (a positive approximation of the curvature), so $\widehat{g}\odot \widehat{g}$ estimates its diagonal. The loss is a mean over bs tokens, which shrinks ${\widehat{g}}^2$ by a factor of about bs, so step multiplies h by bs = y.numel() to undo that.

§8Part 3 · Decoding strategies

A language model only gives a probability distribution over the next token. Decoding is the rule for picking a token from it, appending it, and repeating. Here the model is the pretrained Pythia-160M (12 layers, D = 768, 12 heads, V = 50,304), called through Hugging Face's model(...). generate() is not allowed, so everything around the forward call is written by hand.

§8.1 · Talking to the Hugging Face model

def _left_pad(prompts: list[list[int]], pad_id: int, device):
    B, P = len(prompts), max(len(p) for p in prompts)
    ids = torch.full((B, P), pad_id, dtype=torch.long, device=device)
    mask = torch.zeros((B, P), dtype=torch.long, device=device)
    for i, p in enumerate(prompts):
        ids[i, P - len(p) :] = torch.tensor(p, device=device)
        mask[i, P - len(p) :] = 1
    return ids, mask

def _forward(model, ids, mask, pos, past):
    out = model(input_ids=ids, attention_mask=mask, position_ids=pos, past_key_values=past, use_cache=True)
    return out.logits[:, -1].float(), out.past_key_values

def _prefill(model, prompts, pad_id):
    device = next(model.parameters()).device
    ids, mask = _left_pad(prompts, pad_id, device)
    pos = (mask.cumsum(-1) - 1).clamp(min=0)
    logits, past = _forward(model, ids, mask, pos, None)
    return logits, past, mask, pos[:, -1] + 1

The same ideas as greedy_generate (§3.7), in Hugging Face's vocabulary. attention_mask is a (B, T) tensor of 1s for real tokens and 0s for padding; the model combines it with its own causal mask. position_ids are the RoPE positions, computed with the same cumsum trick. past_key_values is Hugging Face's KV cache object. out.logits is (B, L, V) and [:, -1] keeps the last position, (B, V). _prefill runs the whole prompt once and returns the logits, the cache, the mask, and the next position (B,). Pythia has no pad token, so EOS is used as padding; the mask makes that harmless.

§8.2 · Top-k and top-p filters

def filter_top_k(logits: torch.Tensor, k: int) -> torch.Tensor:
    kth = logits.topk(k, dim=-1).values[:, -1:]
    return logits.masked_fill(logits < kth, float("-inf"))

def filter_top_p(logits: torch.Tensor, p: float) -> torch.Tensor:
    sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True)
    probs = sorted_logits.softmax(-1)
    cum = probs.cumsum(-1)
    remove = (cum - probs) >= p  # mass before this token already >= p
    remove[:, 0] = False
    sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
    return torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits)

Top-k keeps only the k most likely tokens. topk(k).values is (B, k) sorted high to low, and [:, -1:] keeps the k-th largest as (B, 1) so it broadcasts against (B, V). Every logit below it becomes −∞, so the later softmax gives it probability 0 and the rest are renormalised automatically.

Top-p (nucleus) keeps the smallest set of top tokens whose probabilities add up to at least p:

Worked example with 4 tokens and p = 0.7:

probs (vocab order)     [0.50, 0.05, 0.30, 0.15]
sort descending         [0.50, 0.30, 0.15, 0.05]     sorted_idx = [0, 2, 3, 1]
cum                     [0.50, 0.80, 0.95, 1.00]
cum - probs             [0.00, 0.50, 0.80, 0.95]     probability mass BEFORE each token
remove = (...) >= 0.7   [F,    F,    T,    T   ]     keep 0.50 and 0.30
scatter back            token 0 and token 2 kept, tokens 1 and 3 set to -inf
after softmax           [0.625, 0, 0.375, 0]

cum - probs >= p means "the tokens ranked above me already reach p, so I am not needed". It always keeps the token that crosses p. remove[:, 0] = False guarantees at least one token. scatter(-1, sorted_idx, sorted_logits) writes each sorted value back to its original vocabulary position: out[b, sorted_idx[b, j]] = sorted_logits[b, j]. Unlike top-k, the number of kept tokens adapts: few when the model is confident, many when it is unsure.

§8.3 · sample_decode: greedy, top-k, top-p

logits, past, mask, pos = _prefill(model, prompts, pad_id)
for t in range(max_new):
    if strategy == "greedy":
        tok = logits.argmax(-1)
    else:
        lg = logits / temperature
        lg = filter_top_k(lg, k) if strategy == "top_k" else filter_top_p(lg, p)
        tok = torch.multinomial(lg.softmax(-1), 1, generator=generator).squeeze(-1)
    tok = torch.where(done, torch.full_like(tok, pad_id), tok)
    out[:, t] = tok
    done |= tok == eos_id
    if done.all():
        break
    mask = torch.cat([mask, torch.ones((B, 1), dtype=mask.dtype, device=mask.device)], 1)
    logits, past = _forward(model, tok[:, None], mask, pos[:, None], past)
    pos = pos + 1

§8.4 · beam_search, line by line

Greedy picks the best token at each step, but the best sentence may start with a second-best token. Beam search keeps the W best partial sequences ("beams") and scores each by its total log-probability:

logits, past, mask, pos = _prefill(model, prompts, pad_id)
B, W, dev = len(prompts), width, logits.device
logp = logits.log_softmax(-1)  # (B, V)
V = logp.size(-1)
scores, tok = logp.topk(W, dim=-1)  # (B, W)
rep = torch.arange(B, device=dev).repeat_interleave(W)
past = _reorder_cache(past, rep)
mask, pos = mask[rep], pos[rep]
seqs = tok.reshape(B * W, 1)
scores = scores.reshape(B * W)
finished = seqs[:, 0] == eos_id
lengths = torch.ones(B * W, device=dev)

for t in range(1, max_new):
    if finished.all():
        break
    mask = torch.cat([mask, torch.ones((B * W, 1), dtype=mask.dtype, device=dev)], 1)
    logits, past = _forward(model, seqs[:, -1:], mask, pos[:, None], past)
    pos = pos + 1
    logp = logits.log_softmax(-1)
    frozen = torch.full_like(logp, float("-inf"))
    frozen[:, eos_id] = 0.0
    logp = torch.where(finished[:, None], frozen, logp)
    cand = (scores[:, None] + logp).view(B, W * V)
    scores, flat = cand.topk(W, dim=-1)  # (B, W)
    origin = (flat // V + torch.arange(B, device=dev)[:, None] * W).reshape(-1)
    tok = (flat % V).reshape(-1)
    scores = scores.reshape(-1)
    past = _reorder_cache(past, origin)
    mask, pos = mask[origin], pos[origin]
    was_finished = finished[origin]
    seqs = torch.cat([seqs[origin], tok[:, None]], 1)
    lengths = lengths[origin] + (~was_finished).float()
    finished = was_finished | (tok == eos_id)
norm = (scores / lengths**length_penalty).view(B, W)
best = norm.argmax(-1) + torch.arange(B, device=dev) * W
return _strip(seqs[best], eos_id, pad_id)

Complexity. Each of the T steps runs the model on B·W sequences of one new token, with attention over up to P + T cached keys, and a topk over W·V candidates. So time is O(T · B · W · (cost of one token + V)), linear in W, and the KV cache memory is proportional to B · W · (P + T). Measured on 1,000 prompts: widths 1, 2 and 4 took 23.4 s, 29.5 s and 50.0 s. That is less than linear because a GPU running a small batch is under-used, so more beams partly come for free until it fills up.

§8.5 · continuation_nll and the metrics

logits = model(input_ids=ids, attention_mask=mask).logits[:, :-1].float()
tgt = labels[:, 1:]
nll = F.cross_entropy(logits.transpose(1, 2), tgt, ignore_index=-100, reduction="none").sum(-1)
valid = tgt != -100
correct = ((logits.argmax(-1) == tgt) & valid).sum(-1)

This scores a given continuation (the model's own generation, or the reference) under Pythia by teacher forcing: feed prompt + continuation in one pass. Here padding is on the right, since nothing is generated. labels are −100 except over the continuation. Logits at position j predict token j+1, so the last logit ([:, :-1]) and the first label ([:, 1:]) are dropped to line them up, both becoming length L−1. The transpose puts V on axis 1 for cross_entropy; result (B, L−1), summed to (B,).

The results show the classic trade-off. Beam-4 has the lowest perplexity (1.80) but the most repetition (65% repeated 4-grams): the highest-probability text is dull and loops. Top-p 0.95 is the most diverse (distinct-2 of 0.62) with perplexity 34, far above the reference's. Sampling methods land between.

§9utils.py, main.py, tests

Most of this is ordinary Python (YAML config merging, argparse, WandB logging with retries, JSON and checkpoint writing). Only three pieces involve ML.

§9.1 · cosine_lr: warmup then cosine decay

The learning rate climbs linearly from almost 0 for $T_w$ steps (large early steps on random weights, with Adam's statistics still empty, are destabilising), then follows half a cosine from 1 down to $r$ = 0.1 of the peak. Small steps at the end let the weights settle.

§9.2 · LossScaler: making fp16 gradients survive

def backward(self, loss):
    (loss * self.scale).backward()          # scale starts at 2**15 = 32768

def unscale_clip(self, params, max_norm):
    for p in params: p.grad.div_(self.scale)
    gn = torch.nn.utils.clip_grad_norm_(params, max_norm)
    return gn, bool(torch.isfinite(gn))

def update(self, finite):
    if finite:
        self.good += 1
        if self.good % self.growth_interval == 0: self.scale *= 2
    else:
        self.scale /= 2; self.good = 0

(Lightly condensed.) Many gradients are smaller than fp16 can represent and would round to 0. Multiplying the loss by 32,768 multiplies every gradient by the same amount (the chain rule is linear), lifting them into range. They are divided back before the optimiser sees them. If any gradient overflowed to infinity, the step is skipped and the scale halved. After 1,000 good steps it doubles again. With bf16 (same range as fp32) it is switched off.

§9.3 · The rest

• amp_dtype(): bf16 on GPUs with compute capability ≥ 8 (Ampere and newer), else fp16. The RTX 2070 used for the runs is Turing, so fp16.
• torch.compile is used only if triton is installed, which it is not on Windows.
• main.py caps PyTorch at 85% of an 8 GB GPU, because Windows otherwise spills into much slower shared system memory.
• tests/test_sanity.py checks that the parameter counts match, that a 1-expert MoE equals an MLP (the gate is exactly 1), that cached greedy decoding equals recomputing the whole sequence each step, that every optimiser lowers the loss on a toy problem, and that beam width 1 equals greedy.

§10Viva scripts

Each answer below is written to be said out loud or drawn on a whiteboard. Practise the "walk me through" ones with the shape table of §5 open.

1. Walk me through the forward pass of your Part 1 model.

   "The input is a (B, L) batch of token ids: language tag, source, separator, English. The embedding gives (B, L, 512). Then 8 pre-norm blocks. In each block, RMSNorm, then one linear layer makes q, k, v as (B, L, 1536), which I view as (B, L, 3, 8, 64) and permute to three (B, 8, L, 64) tensors. RoPE rotates q and k by position. Causal SDPA computes softmax(QKᵀ/8)V per head, giving (B, 8, L, 64), which I transpose and reshape back to (B, L, 512), project with o_proj and add to the residual. Then RMSNorm and the FFN, dense or MoE, also added back. After the last block, a final norm and the LM head, tied to the embedding, give (B, L, 24000) logits. The loss is cross-entropy on the English tokens only; source positions have label −100."

2. Walk me through the forward pass of your MoE layer.

   "I flatten (B, L, 512) to N = B·L tokens. A linear router gives (N, E) logits; I take the softmax in fp32 and pick the top-k experts per token, giving gates and indices of shape (N, k). For k = 2 I renormalise the two gates to sum to 1; for k = 1 I keep the raw probability so the router still gets gradient. For each expert I find the tokens that chose it with nonzero, run the expert MLP on just those rows, multiply by their gates and index_add the result back into an (N, 512) output. Shared experts run on every token and are added. I reshape back to (B, L, 512). On the side I compute the Switch load-balancing loss E·Σ f_e·P_e and the router z-loss over non-pad tokens, and add 0.01 and 0.001 times them to the training loss."

3. Did your experts specialise by language?

   "Only weakly. The routing heatmaps count which expert each token picks, split by language (vi, ja) and by side (source or English target). I measure how different two rows are with the total variation distance, ½·Σ|p − q|, where 0 means identical routing. Between Vietnamese and Japanese source tokens it never goes above 0.19 in any layer of any variant, and for top-2 it stays at or below 0.09. The bigger split is source versus target: for top-1 it reaches 0.55 in the last layer. So experts divided the work more by what the token is doing, reading the input or writing English, than by which language the input is in. The load-balancing loss also pushes every expert towards an equal share, which works against strong specialisation."

4. Why did the active-matched MoE do best and top-1 worst?

   "All four total-matched variants have 37.5M parameters, but a top-1 token only uses 24.9M of them and top-2 or shared+top-1 uses 29.1M, while the dense model uses all 37.5M. BLEU follows active parameters: top-1 is lowest (vi 43.97, ja 34.34), top-2 and shared are in the middle, dense is at 45.21 and 35.54. The active-matched variant uses the same 37.5M per token as dense but has 54M in total, and it is the only one to beat dense (45.40, 35.76). So at this scale, quality is set mostly by compute per token, and MoE's extra total parameters add a little on top."

5. Walk me through one top-p sampling step.

   "I take the last position's logits, (B, V), divide by the temperature, and sort each row from high to low. Softmax gives the sorted probabilities, and cumsum the running total. A token is removed if the mass before it, cum − prob, is already at least p, so the token that crosses p is kept and the first token is always kept. Removed logits become −∞, and scatter puts the sorted values back in vocabulary order. Softmax again renormalises the kept tokens, and multinomial draws one per row. For probabilities 0.5, 0.3, 0.15, 0.05 and p = 0.7 it keeps the first two, which become 0.625 and 0.375."

6. Walk me through a Sophia step, including where the Hessian comes from.

   "Every step: momentum m = β₁m + (1−β₁)g, weight decay, then each weight moves by −η·sign(m)·min(|m| / (ρ·bs·h + ε), 1). That is a Newton-style step, momentum divided by curvature, clipped so no weight moves more than η. Every 10 steps I refresh h. I sample fake labels from the model's own predicted distribution, compute the cross-entropy against them and backpropagate. The squared gradient ĝ⊙ĝ estimates the diagonal of the Gauss–Newton matrix, and h = β₂h + (1−β₂)ĝ². The loss is a mean over bs tokens, so h is multiplied by bs in the step to undo that scaling."

7. Which optimiser won, and why might that be?

   "At 1× the dataset the validation losses were Muon 3.917, MARS 4.068, AdamW 4.105, Sophia 4.383 and Lion 4.496, and Muon also had the highest BLEU (1.28). Muon orthogonalises each weight matrix's update, so every direction of the matrix moves by a similar amount instead of a few dominant directions taking the whole step. That seems to make faster progress per token in this short, small-model regime. MARS beat AdamW slightly through its gradient-difference correction. Lion and Sophia were behind at the learning rates the sweep picked. The sweep only ran on 0.2× of the data, so their best settings for the full run may differ."

8. Walk me through one step of beam search.

   "I keep B·W beams, each with a token sequence, a cumulative log-probability and its own KV cache row. I feed each beam's last token and get (B·W, V) log-probs. I add each beam's score to get the score of every possible extension, then view that as (B, W·V) so all candidates of one prompt are in one row, and take the top W. For each winner, flat // V tells me which beam it extends and flat % V which token it adds. I reorder the cache, masks and sequences to follow the parents and append the tokens. A finished beam can only add EOS at zero cost, so its score stays fixed. At the end I divide scores by length and take the best beam per prompt. Each step processes B·W sequences, so time and cache memory are linear in W."

9. Walk me through one AdamW step, then say how Muon differs.

   "For each weight: m = β₁m + (1−β₁)g is the momentum, v = β₂v + (1−β₂)g² the running mean square. Both start at zero, so I divide by 1−βᵗ to correct the bias. Then decay the weight by (1 − η·λ), and subtract η·m̂/(√v̂ + ε). So each weight gets its own step size."

   "Muon is for 2-D matrices. It keeps Nesterov momentum and then orthogonalises the update: it replaces the update's singular values with 1 using five Newton–Schulz iterations, X ← aX + b(XXᵀ)X + c(XXᵀ)²X, which apply a polynomial to each singular value and leave the singular vectors alone. So the update moves equally in all directions of the matrix instead of being dominated by a few. Embeddings and norm gains go to an internal AdamW."

10. What are arange, unsqueeze, expand and reshape?

    arange(n) is range(n) as a tensor. unsqueeze(i) inserts a length-1 axis at position i so a tensor can broadcast (RoPE's (B, L, 32) → (B, 1, L, 32) to share angles across heads). expand makes a size-1 axis look larger without copying (arange(L).expand(B, L) gives every sequence the same positions). reshape reads the same numbers in a new shape in row-major order, copying only if the memory is not contiguous; view is the same but never copies.

11. Why left padding for generation, and why the cumsum for positions?

    With left padding the last column is every prompt's last real token, so the next-token logits are always [:, -1] and new tokens are appended in one column for the whole batch. The cumsum − 1 of the padding mask numbers the real tokens 0, 1, 2… so each prompt's first real token has RoPE position 0, as in training. The attention mask hides the pad keys.

12. Why the load-balancing loss? What happens without it?

    Routing is self-reinforcing: an expert that gets more tokens trains more, gets better and attracts even more. Without a penalty one or two experts can take almost all tokens and the rest are wasted parameters. E·Σ f_e·P_e is minimised (value 1) when load is uniform; f has no gradient, so the push goes through the router probabilities P.

13. Why is attention scaled by √d_h?

    If q and k have independent entries of variance 1, their dot product over 64 dimensions has variance 64 (std 8). Softmax of values that far apart is nearly one-hot, and its gradient is nearly 0. Dividing by √64 = 8 brings the standard deviation back to 1.

14. What does the KV cache save, and what does it cost?

    Without it, each new token re-runs the whole sequence, so generating T tokens is about T²/2 token-passes. With it each step processes one token, because earlier keys and values never change under a causal mask. The cost is memory: 2 × layers × B × heads × length × head_dim numbers.

15. Why tie the embedding and the LM head?

    Both are (24000, 512). Tying saves 12.3M parameters, a third of the model, and makes the logit of token v the dot product of the final hidden vector with v's embedding, so "predicting v" and "reading v" use the same representation.

16. Why does Lion need a smaller learning rate than AdamW?

    Lion's update is a sign, so every weight moves by exactly η each step. AdamW's normalised update is usually smaller than 1 per weight. To get a similar overall step size, Lion's η must be smaller (3·10⁻⁴ vs 10⁻³), and its weight decay larger to keep the same effective decay η·λ.

17. Why does beam search give low perplexity but repetitive text?

    It searches for high-probability sequences, and a model assigns high probability to safe, generic and repeated phrases: once a phrase appears, repeating it becomes even more likely. So beam-4 has perplexity 1.80, below the human reference, but 65% repeated 4-grams. Sampling with top-k or top-p gives up some likelihood for diversity closer to human text.

18. Why transpose(1, 2).reshape and not view at the end of attention?

    After the transpose the tensor's memory is no longer in row-major order for the new shape, so view raises an error. reshape copies into the right order. The transpose itself is needed to bring the head axis next to the feature axis so the 8 × 64 can merge into 512.