What breaks when you exceed the trained length

pip install torch

import math, time, torch, torch.nn as nn, torch.nn.functional as F

torch.set_num_threads(4)
VOCAB, DIM, HEADS, LAYERS, TRAIN_LEN = 24, 64, 4, 2, 24

def make_batch(n, seq, g):
    """The last token repeats one from the first half. Predict what followed it."""
    x = torch.randint(2, VOCAB, (n, seq), generator=g)
    idx = torch.randint(1, seq // 2, (n,), generator=g)
    x[torch.arange(n), -1] = x[torch.arange(n), idx]
    return x, x[torch.arange(n), idx + 1]

def apply_rope(x, base=10000.0):
    d, T = x.shape[-1], x.shape[-2]
    inv = 1.0 / (base ** (torch.arange(0, d, 2, dtype=torch.float32) / d))
    ang = torch.arange(T, dtype=torch.float32)[:, None] * inv[None, :]
    c, s = ang.cos(), ang.sin()
    even, odd = x[..., 0::2], x[..., 1::2]
    return torch.stack([even * c - odd * s, even * s + odd * c], -1).flatten(-2)

def alibi_slopes(n):
    start = 2 ** (-8 / n)
    return torch.tensor([start ** (i + 1) for i in range(n)])

class Block(nn.Module):
    def __init__(self, mode):
        super().__init__()
        self.mode = mode
        self.q, self.k = nn.Linear(DIM, DIM, bias=False), nn.Linear(DIM, DIM, bias=False)
        self.v, self.o = nn.Linear(DIM, DIM, bias=False), nn.Linear(DIM, DIM, bias=False)
        self.n1, self.n2 = nn.LayerNorm(DIM), nn.LayerNorm(DIM)
        self.mlp = nn.Sequential(nn.Linear(DIM, 4 * DIM), nn.GELU(), nn.Linear(4 * DIM, DIM))
        self.register_buffer("slopes", alibi_slopes(HEADS)[:, None, None])

    def forward(self, x):
        B, T, _ = x.shape
        h = self.n1(x)
        shape = (B, T, HEADS, DIM // HEADS)
        q = self.q(h).view(shape).transpose(1, 2)
        k = self.k(h).view(shape).transpose(1, 2)
        v = self.v(h).view(shape).transpose(1, 2)
        if self.mode == "rope":
            q, k = apply_rope(q), apply_rope(k)
        att = (q @ k.transpose(-1, -2)) / math.sqrt(DIM // HEADS)
        if self.mode == "alibi":
            pos = torch.arange(T)
            att = att + self.slopes * (pos[None, :] - pos[:, None]).float()
        att = att.masked_fill(torch.triu(torch.ones(T, T, dtype=torch.bool), 1), -torch.inf)
        y = (att.softmax(-1) @ v).transpose(1, 2).reshape(B, T, DIM)
        x = x + self.o(y)
        return x + self.mlp(self.n2(x))

def sinusoidal(T, d):
    pos = torch.arange(T, dtype=torch.float32)[:, None]
    i = torch.arange(d // 2, dtype=torch.float32)[None, :]
    f = pos / 10000.0 ** (2 * i / d)
    pe = torch.zeros(T, d)
    pe[:, 0::2], pe[:, 1::2] = f.sin(), f.cos()
    return pe

class Tiny(nn.Module):
    def __init__(self, mode):
        super().__init__()
        self.mode = mode
        self.emb = nn.Embedding(VOCAB, DIM)
        if mode == "learned":
            self.wpe = nn.Embedding(TRAIN_LEN, DIM)      # exactly the trained length
        self.blocks = nn.ModuleList([Block(mode) for _ in range(LAYERS)])
        self.nf = nn.LayerNorm(DIM)
        self.head = nn.Linear(DIM, VOCAB, bias=False)

    def forward(self, x):
        h = self.emb(x)
        if self.mode == "learned":
            h = h + self.wpe(torch.arange(x.shape[1]))
        elif self.mode == "sinusoidal":
            h = h + sinusoidal(x.shape[1], DIM)
        for b in self.blocks:
            h = b(h)
        return self.head(self.nf(h))[:, -1, :]           # only the last position matters

def train_and_test(mode, steps=2000):
    torch.manual_seed(0)
    g = torch.Generator().manual_seed(1)
    model = Tiny(mode)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
    for _ in range(steps):
        x, y = make_batch(64, TRAIN_LEN, g)
        loss = F.cross_entropy(model(x), y)
        opt.zero_grad(); loss.backward(); opt.step()
    model.eval()
    accs = []
    with torch.no_grad():
        for seq in [24, 32, 48]:
            if mode == "learned" and seq > TRAIN_LEN:
                accs.append(None); continue
            gg = torch.Generator().manual_seed(99)
            x, y = make_batch(512, seq, gg)
            accs.append((model(x).argmax(-1) == y).float().mean().item())
    return accs

print(f"trained only at length {TRAIN_LEN}. random guessing scores {1/(VOCAB-2):.3f}\n")
print(f"{'position method':<14} {'len 24':>8} {'len 32':>8} {'len 48':>8}")
t0 = time.time()
for mode in ["learned", "sinusoidal", "rope", "alibi", "nope"]:
    a = train_and_test(mode)
    cells = ["  crash " if v is None else f"{v:8.3f}" for v in a]
    print(f"{mode:<14}" + "".join(cells))
print(f"\ntotal training time on this CPU: {time.time() - t0:.0f} s")