Vol. XII · Mid-Training Plate 4 of 10long context, context extension, staged training, RoPE scaling, attention cost, long documentsCL

The Long-Context Phase.

Train short where it is cheap, then stretch the context in stages at the end — paying the quadratic price only briefly.

Training at a context of 128 thousand tokens from the start would be wasteful: attention’s work grows with the square of the length (plate IV.7), most documents are short, and short-range skills are learned just as well on short sequences. So models are pretrained at a few thousand tokens and extended at the end. Llama 3 increased its context in six stages from 8 thousand to 128 thousand tokens, using about 800 billion tokens, moving to the next length only once short-context performance had recovered and long-range retrieval worked.

Each stage rescales the positions — the RoPE frequencies of plate VI.6 — and continues training on data with long documents. The course has the miniature version recorded: a model trained at 64 tokens, the same model with positions interpolated to cover 256 but not retrained, and that model fine-tuned briefly at 256. The workshop compares them, and computes what attention costs at each of Llama 3’s lengths.

Below: left, held-out loss against evaluation length for the three recorded models. Right, the share of a forward pass’s FLOPs spent on attention scores for a Llama 3 8B-sized model, as the context grows.

What you will learn
Workshop
Left: recorded runs of micro’s shape with RoPE: trained at 64 tokens (grey); the same weights with positions interpolated by 4, not retrained (blue); and that model fine-tuned for 400 steps (3.3 million tokens) at 256 tokens (red); held-out loss at each evaluation length. Right: attention-score FLOPs as a share of the forward pass for the published Llama 3 8B configuration, against context length.

Stretching the context cheaply

Trained at 64 tokens, the recorded model degrades past its training length: 4.702 nats at 256 and 5.141 at 512. Interpolating its positions without retraining gives 4.804 at 512 but costs accuracy at every length, because the rescaled positions are unfamiliar. A brief fine-tune at 256 — 10% of the original training tokens — repairs that and makes the longer context useful: 4.067 at 256, lower than the original model achieved anywhere, though 512 is still beyond what it learned. That is a long-context stage in miniature: rescale, then train briefly on long sequences.

The cost explains the staging. For a model the size of Llama 3 8B, attention scores are 22% of the forward FLOPs at 8 thousand tokens and 82% at 128 thousand; a whole run at the long length would multiply the cost of pretraining for little gain on the short documents that make up most of the data. Llama 3 spent about 800 billion tokens across its six stages — about 5% of its pretraining — and checked at each stage that short-context performance had recovered before going longer. Effective context, not the nominal length, is what such a phase must deliver (plate VI.8).

Long-context data is itself scarce: most documents are short, so long-context stages upsample books, long code repositories and concatenated related documents, and some generate synthetic long-range tasks.

Terms

Context extension
Training a model to use sequences longer than its pretraining length.
Staged extension
Increasing the context in several steps, each with its own short training.
Position rescaling
Changing RoPE frequencies so that longer positions map into the trained range.
Effective context
The length over which a model actually uses information (plate VI.8).
Context is extended late, in stages, with rescaled positions and long data — paying attention’s quadratic cost only briefly.

Sources

  1. Llama Team, AI @ Meta, “The Llama 3 Herd of Models”, arXiv:2407.21783, 2024, section 3.4.2.
  2. Chen et al., “Extending Context Window of Large Language Models via Positional Interpolation”, arXiv:2306.15595, 2023.