"""Tests for Lab B: tiny decoder (lessons 1-3, 7-10). Covers: * forward output shape [B, T, V] * parameter count matches hand-computed formula * RMSNorm / decoder block shape preservation * causal safety (future positions cannot affect earlier outputs) * label alignment: correct shift pairs logits[t] with labels[t+1] * wrong shift pairs logits[t+1] with labels[t] (the alignment bug) * correct-shift loss decreases (measured data, not proof of alignment) Run: python -m pytest test_tiny_decoder.py -v """ from __future__ import annotations import torch import torch.nn as nn from model import ( CausalSelfAttention, DecoderBlock, RMSNorm, SwiGLU, TinyDecoder, count_parameters, expected_parameter_count, ) from train import loss_correct, loss_wrong V, D, N_HEADS, N_BLOCKS, MAX_POS = 32, 16, 4, 2, 16 B, T = 2, 6 PAD_ID = 0 IGNORE_ID = -100 def make_model() -> TinyDecoder: torch.manual_seed(0) return TinyDecoder(V, D, N_HEADS, N_BLOCKS, MAX_POS, tie_weights=True) def make_batch() -> tuple[torch.Tensor, torch.Tensor]: input_ids = torch.tensor( [[1, 2, 3, 0, 0, 0], [4, 5, 6, 7, 0, 0]], dtype=torch.long ) labels = input_ids.clone() labels[labels == PAD_ID] = -100 return input_ids, labels # --------------------------------------------------------------------------- # Shape tests # --------------------------------------------------------------------------- def test_forward_output_shape(): model = make_model() input_ids, _ = make_batch() logits = model(input_ids) assert logits.shape == (B, T, V), f"expected {(B, T, V)} got {tuple(logits.shape)}" def test_rmsnorm_preserves_shape(): norm = RMSNorm(D) x = torch.randn(B, T, D) assert norm(x).shape == x.shape def test_decoder_block_preserves_shape(): block = DecoderBlock(D, N_HEADS) x = torch.randn(B, T, D) out = block(x) assert out.shape == x.shape == (B, T, D) def test_swiglu_intermediate_shapes(): mlp = SwiGLU(D) x = torch.randn(B, T, D) proj = mlp.w(x) # [B, T, 4D] assert proj.shape == (B, T, 4 * D) gate, val = proj.chunk(2, dim=-1) assert gate.shape == (B, T, 2 * D) assert val.shape == (B, T, 2 * D) out = mlp(x) assert out.shape == (B, T, D) # --------------------------------------------------------------------------- # Parameter count test # --------------------------------------------------------------------------- def test_parameter_count_matches_formula(): model = make_model() actual = count_parameters(model) expected = expected_parameter_count( V, D, N_HEADS, N_BLOCKS, MAX_POS, tie_weights=True ) assert actual == expected, f"expected {expected:,} params, got {actual:,}" # --------------------------------------------------------------------------- # Causal safety test (future tokens must not affect earlier outputs) # --------------------------------------------------------------------------- def test_causal_safety(): torch.manual_seed(7) block = DecoderBlock(D, N_HEADS) base = torch.randn(1, T, D) corrupt = base.clone() corrupt[0, 2, :] = 999.0 # corrupt a future position (index 2) out_base = block(base) out_corrupt = block(corrupt) # positions 0 and 1 must be unaffected by position 2 (it is in their future) assert torch.allclose(out_base[0, 0], out_corrupt[0, 0], atol=1e-5) assert torch.allclose(out_base[0, 1], out_corrupt[0, 1], atol=1e-5) # --------------------------------------------------------------------------- # Label alignment / loss reduction tests # --------------------------------------------------------------------------- def test_correct_shift_loss_decreases(): """Correct-shift loss decreases on the fixed batch. Reported as measured data only. A decreasing loss does NOT by itself prove that logits[t] is paired with labels[t+1]; that is verified separately in test_specific_pairing_inspection. """ torch.manual_seed(3) model = TinyDecoder(V, D, N_HEADS, N_BLOCKS, MAX_POS, tie_weights=True) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2) input_ids, labels = make_batch() first = None last = None for _ in range(20): optimizer.zero_grad() logits = model(input_ids) loss = loss_correct(logits, labels) if first is None: first = loss.item() loss.backward() optimizer.step() last = loss.item() assert first is not None and last is not None assert last < first, ( f"correct shift: loss {first:.4f} -> {last:.4f} did not decrease" ) def test_wrong_vs_correct_loss_on_crafted_logits(): """Static diagnostic: wrong and correct shifts give different losses. Craft logits where batch 0 position 0 confidently predicts token 5 (logit 100 for token 5, -100 elsewhere), and set labels[0, 1] = 5 so token 5 is the NEXT token after position 0. Correct shift pairs logits[0,0] with labels[0,1]=5 (aligned). Wrong shift pairs logits[0,1] with labels[0,0] (misaligned — same cell of the sliced view pairs the logit at position 1 with the label at position 0). Same inputs and labels; only the shift differs. This detects the misalignment independent of training dynamics / memorisation. Actual measured values are reported; no "near-zero" claim is made because the loss averages over all positions, not just the inspected pairing. """ input_ids, labels = make_batch() logits = torch.full((B, T, V), -100.0) logits[0, 0, 5] = 100.0 labels[0, 1] = 5 loss_c = loss_correct(logits, labels).item() loss_w = loss_wrong(logits, labels).item() print(f"\n crafted-logit losses: correct={loss_c:.4f}, wrong={loss_w:.4f}") assert loss_c != loss_w, ( f"diagnostic failed: correct {loss_c:.4f} == wrong {loss_w:.4f}; " f"shifts should differ" ) def test_specific_pairing_inspection(): """Verify the concrete logit/label pairing for each shift. For batch=0: * Correct shift: logits[:, :-1] vs labels[:, 1:] sliced cell (b=0, t=0) pairs logits[0,0] with labels[0,1]. * Wrong shift: logits[:, 1:] vs labels[:, :-1] sliced cell (b=0, t=0) pairs logits[0,1] with labels[0,0] (the logit at position 1 is scored against the label at position 0 — the alignment bug). """ input_ids, labels = make_batch() print(f"\n input_ids[0] = {input_ids[0].tolist()}") print(f" labels [0] = {labels[0].tolist()}") # Correct shift pairing: logits[pos] -> labels[pos + 1] correct_logit_pos = 0 correct_label_pos = correct_logit_pos + 1 assert labels[0, correct_label_pos].item() != IGNORE_ID, ( "correct pairing target must not be padded" ) print(f" CORRECT: logits[0,{correct_logit_pos}] -> labels[0,{correct_label_pos}] " f"= {labels[0, correct_label_pos].item()} (next token)") # Wrong shift pairing: sliced cell t pairs logits[t+1] with labels[t] # logits[0,1] should predict labels[0,2] (next token), but wrong shift # pairs it with labels[0,0] (previous position's label) — the bug. wrong_logit_pos = 1 wrong_label_pos = wrong_logit_pos - 1 assert labels[0, wrong_label_pos].item() != IGNORE_ID, ( "wrong pairing target must not be padded" ) print(f" WRONG: logits[0,{wrong_logit_pos}] -> labels[0,{wrong_label_pos}] " f"= {labels[0, wrong_label_pos].item()} (previous-position label — bug)") # The two shifts must pair different (logit, label) cells for the same # sliced position, proving they are not the same computation. assert (correct_logit_pos, correct_label_pos) != ( wrong_logit_pos, wrong_label_pos, )