Long-Context Training Recipes
How a model trained at 8k becomes a genuine 128k model, why the data mixture matters more than the token count, and what the Llama 3 and ProLong recipes agree on.
Stretching the positional encoding is the famous half of long context, and the cheap half. RoPE rescaling tricks make a 4k checkpoint mechanically accept 128k tokens (see RoPE scaling), but accepting is not using: a model that never saw genuine long-range dependency during training has no reason to have learned heads that exploit it. The other half is a training problem, and its recipes have converged in public work from three directions: an academic data-engineering study, Meta's Llama 3 report, and Princeton's ProLong.
Less data than you fear, more care than you expect
The headline from the data-engineering line is that continual pretraining on 1B to 5B tokens of long-sequence data is enough to teach 128k-context retrieval to a 7B–13B model, a rounding error next to pretraining cost (Fu et al., 2024, Data Engineering for Scaling Language Models to 128K Context, arXiv:2402.10171). The ability to attend far away appears to be mostly latent in the pretrained weights; the continual stage teaches the model to use positions it has never occupied, not to acquire a new faculty. The catch is in the mixture. Naively upsampling long documents skews the domain balance, because document length correlates hard with domain: the long tail is books and code, so "train on long documents" quietly means "train on books and code" and degrades everything else. Fu et al.'s fix is per-source length upsampling: oversample long documents within each domain while holding the cross-domain mixture fixed.
What Llama 3 and ProLong agree on
Meta ran context extension as six staged increments from 8k to 128k during the tail of Llama 3 405B's pretraining, spending roughly 800B tokens, with a gate at each stage: advance only when short-context evals have fully recovered and needle-in-a-haystack is perfect at the current length (Grattafiori et al., 2024, The Llama 3 Herd of Models, arXiv:2407.21783). ProLong, at 8B scale and 40B tokens, reached state-of-the-art 128k performance among peers and usable 512k behaviour, and published the ablations behind each choice (Gao et al., 2024, How to Train Long-Context Language Models (Effectively), arXiv:2410.02660).
Three findings recur across both. Short data is not optional: ProLong keeps a majority of high-quality short-context data in the mix specifically to prevent regression, and Llama 3's stage gates exist for the same reason. Training longer than you evaluate helps: ProLong trains at 512k to be good at 128k, suggesting the benefit is in forcing genuinely long-range attention patterns, not in matching the deployment length. And evaluation drives the recipe: perplexity at length barely discriminates, so both efforts gate on downstream behaviour, with ProLong going further and evaluating after SFT, where long-context ability actually shows (see length extrapolation for why perplexity misleads).
A last ProLong result worth keeping: SFT on ordinary short instruction data was enough; adding synthetic long instruction data did not help their models. The long-range capability comes from continual pretraining, and instruction tuning merely needs to avoid destroying it.
When it breaks
- Attention cost still gates the whole enterprise. Training at 128k-plus lengths forces sequence parallelism (ring attention or context sharding as in Llama 3), and the engineering bill lands here, not in the data pipeline.
- The stage gates are load-bearing. Extend too fast and short-context quality quietly erodes; the failure surfaces weeks later as regression on tasks nobody associated with context length.
- Long-data scarcity caps quality. Above roughly 100k tokens, natural documents thin out fast, and recipes lean on concatenation or synthetic length; dependencies that never genuinely span the distance teach positions without teaching use.
- NIAH is a gate, not a goal. Both recipes use needle tests as a floor; a model can pass them and still fail multi-fact aggregation at length, which is why ProLong's downstream-task protocol exists at all.
5 flashcards for this concept
Click a card to reveal the answer.