Part V — Alignment & inference

Long context — extending the window

Lesson 14 gave you an instrument and it returned a verdict: the base model is genuinely strong — perplexity low, benchmarks solid, the knowledge demonstrably in the weights. But every one of those numbers was measured inside the window the model was trained at. Lesson 5 forced you to pretrain at a short sequence length — 8K, say — because activation memory scales with B·T·d·L and you could not afford more. The real jobs you want this model for — read a 300-page contract, reason across a whole codebase, run an agent that accumulates a long scratchpad — need 10–100× that window. Feed the model a sequence longer than it trained on and it does not gracefully reach further; it falls off a cliff. This lesson extends the window cheaply, after pretraining, instead of paying the quadratic cost of training long from the start.

The previous step left this broken

Two facts from earlier in the track collide here. First, lesson 5's activation budget (memory ∝ B·T·d·L, plus the attention term) is exactly why you trained at a modest T_train — doubling the window doubles activation memory and quadruples attention compute, so pretraining at 128K was never affordable. Second, lesson 3's RoPE was sold as "extrapolates gently" — but only gently: past the trained length the rotation angles are ones the model literally never saw during training, and quality collapses, not degrades. So the strong base model from lesson 14 is boxed into an 8K window, the KV cache (lesson 18) grows linearly with every position you add, and the jobs you actually care about need far more reach than that. You cannot just "use it at 128K"; you must extend it.

Linear position

Forced by: lesson 5's activation budget forced a short pretraining window; lesson 14 proved the base model is strong but only within that window; real long-document / codebase / agent workloads need 10–100× more, and naive RoPE (lesson 3) collapses past T_train while the KV cache (lesson 18) grows linearly with T.
New idea: extend the window cheaply after pretraining — rescale the RoPE frequencies (position interpolation → NTK-aware → YaRN) so every long position looks "in range," then run a short continued-training phase on genuinely long documents — instead of paying the quadratic cost of pretraining long from scratch.
Forces next: a long-context base model still completes text — it now has reach and knowledge but no instruction-following — which is the same wall lesson 14 named, now over a 128K window, and it forces alignment: SFT, lesson 16.

The plan

Six moves. (1) The wall: attention compute ∝ T² and activation memory ∝ T make pretraining at 128K unaffordable, so you train short and extend long. (2) Why naive extrapolation fails — RoPE encodes position as a rotation angle, and positions past T_train are angles the model never saw. (3) Position interpolation, NTK-aware scaling, and YaRN — three ways to rescale the frequencies so long positions fall back in range. (4) Continued training: a short phase at the long length on real long documents, and the data pitfalls. (5) Serving long context: the KV cache becomes the bottleneck — GQA/MLA, FlashAttention, and sequence/context parallelism. (6) Evaluating it: perplexity-at-length, needle-in-a-haystack, RULER, and "lost in the middle." Then drive the context-extension explorer.

1 · The wall: why you train short and extend long

Lesson 5 gave you two back-of-envelope models, and both punish context length directly. The attention score matrix is T × T per head, so attention compute scales as T²; the activations you must keep for the backward pass scale as B·T·d·L (linear in T for the residual stream, with the attention term itself quadratic before FlashAttention tames its memory). Put numbers on the quadratic term. Going from T_train = 8,192 to 128K is a 16× increase in sequence length, so the per-layer attention FLOPs rise by 16² = 256×:

attention FLOPs ∝ T² ⇒ (128K / 8K)² = 16² = 256×

That 256× is per token of training, on the attention sublayer, before you count the activation memory that also balloons. Pretraining a frontier model already costs C ≈ 6ND over trillions of tokens (lesson 5's budget); multiplying the attention cost by hundreds and the activation footprint by 16× — for the entire run — is simply not a budget anyone spends. You would be paying the long-context tax on every one of the trillions of pretraining tokens, the overwhelming majority of which are short documents that never needed it.

So the field settled on a two-phase split that the rest of this lesson elaborates:

train shortpretrain the whole budget at a modest T_train (e.g. 8K), where attention is cheap and activations fit — this is where the model learns language and knowledge (lesson 14 confirms it works).

extend longa short continued-training phase (often <1% of pretraining tokens) at the long length, after rescaling RoPE so the new positions are in-distribution — this teaches reach, not knowledge.

The budget framing is the whole point: "train short, extend long" is the cheap path, and it works because long-context capability turns out to be a thin adaptation on top of a strong short-context base — not a property you must bake in from the first token. The widget at the end makes the cost gap concrete: full long-pretraining scales with T²; train-short-then-extend is nearly flat in the long length.

2 · Why naive extrapolation fails

Recall the RoPE mechanism from lesson 3. Position is injected by rotating the query and key vectors: split each head's dimensions into pairs, and pair i rotates at frequency θ_i = θ^−2i/d_head, so the query at position m is turned by angle m·θ_i. Because the attention score between two rotated vectors depends only on the difference of their angles, attention sees relative distance (m − n) — that is why RoPE works at all and why it is parameter-free.

Now look at what happens past the training length. During training at T_train = 8,192, the largest relative distance the model ever experienced is about 8,192, so the rotation angles it learned to interpret span a bounded range. Ask it about positions 8,193 — or 100,000 — and the low-frequency pairs (the ones that rotate slowly and encode long-range position) are now turning through angles larger than anything in the training distribution. The model has no learned behavior for "tokens 50,000 apart"; that input is out-of-distribution in the most literal sense. The result is not graceful decay — it is collapse: attention patterns become noise, perplexity explodes, and the model produces garbage well before you reach the nominal new length.

The intuition: a clock with too-fast hands

Think of each RoPE frequency as a clock hand. High-frequency pairs (fast hands) complete many rotations within the trained window — they encode local position and have seen every angle already, so they extrapolate fine. Low-frequency pairs (slow hands) barely complete a fraction of a turn across 8K — they encode global position, and past T_train they swing into angular territory they never saw in training. Naive extrapolation breaks precisely because those slow, global hands are asked to read a time that was never on the dial. Every extension method below is a way to keep those slow hands inside the range they were trained on.

3 · Rescale the frequencies: PI → NTK-aware → YaRN

The fix is to change the angles you feed RoPE at the long length so they fall back into the range the model trained on — then fine-tune briefly to clean up the residual mismatch. Three methods, each a refinement of the last.

Position interpolation (PI, Chen et al. 2023)

The simplest idea: linearly squash the position indices so the new range [0, L_new] maps into the trained range [0, T_train]. Multiply every position by the scale factor s = T_train / L_new before applying the rotation. Extending 8K → 32K means s = 1/4: position 32,000 is presented to RoPE as if it were position 8,000, an angle the model has seen. Now every position looks in-range — no angle exceeds what training covered.

RoPE angle (PI): (m · s) · θ_i, s = T_train / L_new ≤ 1

The cost: you have compressed the resolution. Two adjacent tokens that used to differ by one full step now differ by s of a step, so the high-frequency (local) pairs are crowded together and fine-grained local position is blurred. PI alone degrades short-range modeling — but a few hundred steps of fine-tuning at the long length recovers quality, because the model only has to re-learn to read the rescaled local angles, not learn long context from scratch. PI was the breakthrough that made cheap extension practical.

NTK-aware scaling

PI's flaw is that it stretches every frequency band by the same factor s, including the high-frequency local pairs that did not need stretching — they had already seen all their angles. NTK-aware scaling fixes this by scaling the RoPE base θ instead of the positions (e.g. 10,000 → 500,000). Increasing the base changes each frequency band by a different amount: high-frequency pairs are barely touched (local resolution preserved), while low-frequency pairs — the ones that actually went out-of-distribution — are stretched the most. It targets the stretch where the problem is, so NTK-aware extends further than PI for a given amount of fine-tuning, and a modest extension can even work with no fine-tuning at all.

YaRN (Yet another RoPE extensioN, Peng et al. 2023)

YaRN takes the NTK insight to its conclusion: don't apply one global rule — ramp the interpolation per frequency band. Partition the RoPE dimensions into three regimes by how many full rotations they complete inside T_train: high-frequency pairs (many rotations, fully seen) are left untouched — pure extrapolation; low-frequency pairs (less than one rotation, never seen at long range) are fully interpolated like PI; and a smooth ramp blends the two across the middle band. YaRN adds a small attention-temperature correction to compensate for the entropy change. The result is the current best trade-off — the biggest extension for the least fine-tuning — and it is what most open long-context models (Llama-3.1's 128K, many others) use to reach 32K–128K from an 8K base.

Method	What it rescales	Local resolution	Fine-tune needed	Reach
Naive RoPE	nothing	perfect	—	collapses past T_train
PI (Chen 2023)	positions, uniform s	blurred everywhere	~hundreds of steps	~2–4×
NTK-aware	base θ (per-band)	local mostly kept	little / sometimes none	~4–8×
YaRN	per-band ramp + temperature	local kept, global stretched	least	~8–32×

All four share the same skeleton: keep the local (high-frequency) information the model already reads well, and bring the global (low-frequency) angles back inside the trained range so they stop being out-of-distribution. They differ only in how surgically they do it — and surgery buys you reach per fine-tuning dollar.

4 · Continued training at the long length

Rescaling the angles makes long positions look in-range, but the model still has a small mismatch to absorb (the squashed local angles, the new temperature), and — more importantly — it has never actually practiced using 100K tokens of context to predict the next one. The fix is a short continued-training (long-context fine-tuning) phase: a few billion tokens — typically well under 1% of the pretraining budget — at the new sequence length, with the rescaled RoPE in place. This is cheap precisely because of §1: you pay the T² attention tax only over this tiny tail, not the whole run.

The non-obvious hard part is the data. To teach long-range dependency you need documents that genuinely are long and whose distant tokens genuinely depend on each other:

Naturally long docs

Books, long-form articles, legal/scientific documents, and whole code repositories (files concatenated in dependency order) — sources where understanding token 90,000 really does require token 100.

Concatenation

Stitching shorter documents to fill the window. Cheap and common, but risky: unrelated docs glued together teach no long-range dependency — the model learns it can safely ignore far-away tokens, the opposite of the goal.

Up-weight the long tail

Raw web skews short, so naively sampling gives few long examples. Up-weight or bucket long documents so a meaningful fraction of each batch actually exercises the new range.

Keep the short mix

Train only on long docs and short-context quality regresses. Blend in the original short distribution so the model extends its reach without forgetting how to handle ordinary prompts.

The recurring pitfall — and a genuine failure mode — is the concatenation trap: if your "long" data is just unrelated short docs packed to length, the loss falls (predicting each short doc is easy) but the model never learns to attend across the window. It scores fine on perplexity and fails every retrieval test in §6. Long-context data must contain real long-range structure, not just real length.

5 · Serving long context: the KV cache strikes back

Suppose extension worked and you have a 128K-capable model. Now the bottleneck moves to inference, and it is exactly the object lesson 18 warned about: the KV cache grows linearly with context T. From lesson 18's derivation, the cache is 2 · L · n_kv · d_head · T · batch · bytes. For the 7B-class model there (L = 32, d_head = 128, bf16, MHA) one token costs 512 KB, so a single 128K-context request needs 512 KB × 128K ≈ 69 GB of cache — on its own nearly the whole 80 GB GPU, before the 14 GB of weights. Long context is a memory problem at serving time, and the levers are all callbacks:

GQA / MLA shrink the cache

Grouped-query attention (lesson 3) cuts the cache by the group factor g — that 69 GB drops to ~17 GB at g = 4. Multi-head latent attention (MLA, DeepSeek) compresses K/V into a low-rank latent for an even smaller footprint. This is what makes 128K servable at all.

FlashAttention keeps memory O(T)

Lesson 7's FlashAttention never materializes the T × T score matrix — it streams K/V tiles through SRAM with the online-softmax recurrence, so attention activation memory is O(T) not O(T²). Without it, long-context attention OOMs on the score matrix alone; with it, the compute is still O(T²) but the memory scales.

Sequence / context parallelism

Lesson 9's sequence/context parallelism (and ring attention) splits the sequence dimension T across GPUs, so no single device holds the whole cache or all activations — the way you fit a context too large for one GPU, at both train and serve time.

The clean way to hold it: extension is what makes the model capable of long context; GQA/MLA, FlashAttention, and context parallelism are what make that capability affordable to run. The first half of this lesson buys the reach; the second half (and lessons 7, 9, 18) pays the serving bill. The production machinery — paged KV, prefix caching, the long-context scheduler — lives in the vLLM track and the SGLang track.

6 · Evaluating it: perplexity is necessary, not sufficient

How do you know the extension actually worked, rather than merely not crashing? Lesson 14's instrument warned that an easy number can lie, and long context is a textbook case.

Perplexity-at-length — measure held-out perplexity as a function of position deep into the window — is the first check and a necessary one: if PPL rises past T_train, extension failed outright. But it is not sufficient, and this is the trap. A model can keep perplexity flat at 100K simply by using local context to predict each token well — the vast majority of next-token predictions only need the last few hundred tokens — while being completely unable to retrieve a specific fact planted 90K tokens back. Low perplexity at length proves the model isn't broken; it does not prove the model is using the distance.

needle in a haystackplant a single specific fact ("the secret code is 7741") at a random depth in a long filler context, then ask for it. Sweep depth × context length and chart retrieval accuracy — directly tests whether far-away tokens are actually attended to.

RULERa synthetic benchmark suite that goes beyond single-needle: multiple needles, multi-hop tracing, aggregation, and long-context QA. It exposes the gap between a model's claimed window and its effective window — many "128K" models degrade sharply well before 128K on RULER.

lost in the middlethe documented effect (Liu et al. 2023) that models retrieve facts placed near the start or end of the context far better than facts in the middle — a U-shaped accuracy curve. A model can have a 128K window and still systematically miss what's at position 64K.

So the honest evaluation of a long-context model is a basket (just like lesson 14): perplexity-at-length to confirm it isn't broken, plus needle/RULER to confirm it actually retrieves across the window, read with the knowledge that the middle is the weak spot. A model that aces perplexity and flunks the needle has the reach on paper and not in practice.

7 · Drive the context-extension explorer

The widget puts the whole lesson on three knobs. The target context slider sweeps from T_train out to 128K. The method toggle picks naive RoPE, PI, or YaRN. The GQA toggle halves-or-quarters the KV cache. Watch three readouts. (a) A quality proxy that holds high inside T_train for every method but, past it, collapses under naive RoPE while PI and YaRN stay high (YaRN highest, PI slightly lower from its uniform blur) — each at a small fine-tune cost shown in the KPI. (b) The KV-cache GB, rising linearly with T (callback lesson 18), with GQA dividing it. (c) The relative training cost: full long-pretraining scales as T² and explodes, while train-short-then-extend stays nearly flat. The configuration that "breaks": push to 128K with naive RoPE and the quality proxy craters; switch the method to YaRN and it recovers.

Context-extension explorer — reach vs the cost of getting it

Drag target context past T_train = 8K. With naive RoPE the quality proxy collapses the moment you exceed the trained length; switch the method to PI or YaRN and it stays high (for a small fine-tune cost). The KV-cache bar grows linearly with T (lesson 18); toggle GQA to divide it. The training-cost bars contrast full long-pretraining (∝ T², explodes) with train-short-then-extend (cheap, flat) — the reason the cheap path exists.

effective quality @ length T_train = 8K current target

naive RoPE position interpolation YaRN GQA (g = 4)

target context T: 32,768

Effective quality @ length

—

KV cache (1 seq)

—

Fine-tune cost

—

Extension factor

—

extend (this path)

—

pretrain long (∝T²)

—

The lesson of the widget is the lesson of the chain: the expensive axis (context) is not bought by spending the expensive resource (long pretraining). It is bought by a cheap rescale-and-finetune that exploits a strong short-context base — and then paid for, at serving time, with the same KV-cache levers (GQA, FlashAttention, context parallelism) that every other part of the inference stack leans on.

Failure modes & checklist

Failure modes

Running a base model past T_train with no rescaling. Feeding 32K to an 8K-trained model and expecting RoPE to "extrapolate." Signal: output is fluent up to ~8K then becomes incoherent; perplexity-at-length spikes the moment you cross the trained length.
Uniform PI when you needed per-band. Squashing all frequencies kills local resolution. Signal: long-range retrieval is okay but short-context quality regressed; switch to NTK-aware / YaRN to spare the high-frequency bands.
The concatenation trap. "Long" data that is unrelated short docs packed to length. Signal: perplexity-at-length looks fine but needle-in-a-haystack and RULER fail — the model learned to ignore distant tokens.
Trusting perplexity-at-length alone. Declaring 128K success because PPL stayed flat. Signal: great PPL, but the model can't retrieve a fact planted at 64K — it was using local context the whole time.
Forgetting the KV cache at serving. Shipping a 128K model on MHA. Signal: a single long request OOMs the GPU (≈69 GB cache on a 7B-MHA at 128K); you needed GQA/MLA, paged KV, or context parallelism.
Training only on long docs. Dropping the short-context mix during the extension phase. Signal: the model handles 100K documents but its short-prompt quality and benchmark scores slip.

Checklist

Pretrain short, extend long. Spend the budget at a modest T_train; reserve a small tail for long-context continued training.
Rescale RoPE before extending — prefer YaRN (or NTK-aware) over uniform PI so high-frequency local bands are spared.
Fine-tune at the long length on genuinely long, dependency-rich documents; up-weight the long tail and keep a short-context blend.
Evaluate a basket: perplexity-at-length to confirm it isn't broken, plus needle-in-a-haystack and RULER to confirm real retrieval; expect a "lost in the middle" dip.
Budget the KV cache at 2·L·n_kv·d_head·T·batch·bytes for your worst-case context; pick GQA g (or MLA) and use FlashAttention + context parallelism.
Quote the effective window, not the nominal one — the length where retrieval still holds, not the length you trained the angles for.

Checkpoint

Try it

Open the context-extension explorer. Set the method to naive RoPE and drag the target context to 128K: the quality proxy collapses the instant you pass T_train = 8K, and the "pretrain long" cost bar pins to the top (∝ T²). Now switch the method to YaRN at the same 128K — watch the quality proxy recover to near-baseline while the "extend" cost bar stays a sliver, for the small fine-tune cost shown in the KPI. Then toggle GQA off and read the KV-cache KPI quadruple. In one sentence each: why does naive RoPE collapse exactly at T_train and not before, and why is the "extend" cost bar nearly flat while "pretrain long" explodes?

Where this points next

You now have a model that reaches 128K cheaply: rescale RoPE with YaRN, fine-tune briefly on real long documents, serve it with GQA + FlashAttention + context parallelism. But point lesson 14's instrument at it and the verdict is the same one you got at 8K, just over a far bigger window. The model has more reach and the knowledge is still in there — but it is still a completer. Hand it a 100K-token contract and ask "what are the termination clauses?" and a base model's most-likely continuation is more contract-like text, or another question, not an answer — exactly the behavior gap lesson 14 exposed, now with 16× the context. Extension changed how far the model can see; it did nothing for how it behaves. Installing the assistant behavior — turning the long-reach completer into something that follows instructions — is supervised fine-tuning, lesson 16, SFT — from completer to assistant.

Takeaway

Lesson 5's activation budget (∝ B·T·d·L) and attention's T² compute forced you to pretrain at a short window, and lesson 14 proved the base model is strong only inside it — yet real long-document, codebase, and agent workloads need 10–100× more reach, while naive RoPE (lesson 3) collapses past T_train because long positions are rotation angles it never saw. The fix is train short, extend long: rescale the RoPE frequencies so long positions fall back in range — position interpolation squashes all positions uniformly (cheap, blurs local resolution, needs a few hundred fine-tune steps); NTK-aware scales the base θ per-band to spare local pairs; and YaRN ramps the interpolation per frequency band (untouched local, fully interpolated global, smooth middle) for the biggest reach at the least fine-tuning — then a short continued-training phase (well under 1% of the budget) at the long length on genuinely long, dependency-rich documents (beware the concatenation trap and keep a short-context blend). At serving time the KV cache grows linearly with T (lesson 18) — a 128K request is ~69 GB on a 7B-MHA — so GQA/MLA shrink it, FlashAttention (lesson 7) keeps attention memory O(T), and sequence/context parallelism (lesson 9) splits T across GPUs. Evaluate a basket: perplexity-at-length is necessary but not sufficient, so add needle-in-a-haystack and RULER and expect "lost in the middle." And the punchline that forces the next lesson: a long-context base model still just completes text — the reach is bigger, the behavior gap is unchanged — which forces SFT in lesson 16.

Interview prompts

Why pretrain at a short window and extend afterward instead of pretraining long from the start? (§1 — attention compute is ∝ T² and activation memory ∝ B·T·d·L, so 8K→128K multiplies attention FLOPs by 16²=256× and activations by 16× over the entire trillion-token run — unaffordable. Long-context ability is a thin adaptation on a strong short base, so a short tail at the long length is enough.)
Why does a RoPE model collapse past its training length rather than degrade gently? (§2 — RoPE encodes position as rotation angle m·θ_i; past T_train the low-frequency (global) pairs swing through angles never seen in training — literally out-of-distribution — so attention becomes noise. The high-frequency local pairs extrapolate fine; the slow global ones break.)
Contrast position interpolation, NTK-aware scaling, and YaRN. (§3 — PI multiplies all positions by s = T_train/L_new (uniform squash; blurs local resolution; ~hundreds of fine-tune steps). NTK-aware scales the base θ so each band stretches differently, sparing high-frequency pairs. YaRN ramps interpolation per band — untouched local, fully interpolated global, smooth middle, plus a temperature fix — biggest reach for least fine-tuning.)
What is the "concatenation trap" in long-context data and how would you detect it? (§4 — filling the window with unrelated short docs teaches the model it can ignore distant tokens, so it never learns long-range dependency. Detection: perplexity-at-length looks fine (each short doc is easy) but needle-in-a-haystack and RULER fail. Fix: use genuinely long, dependency-rich documents.)
Why is perplexity-at-length necessary but not sufficient to validate a long-context model? (§6 — most next-token predictions only need local context, so a model can hold PPL flat at 100K using the last few hundred tokens while being unable to retrieve a fact planted 90K back. PPL confirms it isn't broken; needle/RULER confirm it actually attends across the window. Expect "lost in the middle" — a U-shaped retrieval curve.)
What is the dominant serving cost of long context and which levers attack it? (§5 — the KV cache grows linearly with T (2·L·n_kv·d_head·T·batch·bytes); a 128K request is ~69 GB on a 7B-MHA. GQA/MLA divide the cache (lesson 3), FlashAttention keeps attention memory O(T) (lesson 7), and sequence/context parallelism / ring attention split T across GPUs (lesson 9).)
Your "128K" model aces perplexity-at-length but flunks RULER beyond 32K — what's the diagnosis? (§6, §4 — the nominal window exceeds the effective window: extension made positions in-range (PPL fine) but the model isn't using the far context, often from weak long-range data (concatenation trap) or insufficient long fine-tuning. Quote the effective window and retrain on dependency-rich long docs.)
You've extended a base model to 128K — why isn't it ready to deploy as an assistant? (Where this points next — extension changed reach, not behavior; the model still maximizes p(text) and completes rather than answers, the same gap lesson 14 found, now over a bigger window. SFT (lesson 16) installs the instruction-following behavior.)