Source-linked AI summary
Data Engineering for Scaling Language Models to 128K Context
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, Hao Peng
TL;DR
The paper asks whether models pretrained on short contexts already possess arbitrary-location retrieval and can be extended to 128K affordably. It uses full-model continual pretraining with balanced, length-upsampled data, finding that 1-5B tokens yield balanced gains and approach GPT-4 128K retrieval performance.
Problem
It is unclear whether reliable arbitrary-location retrieval at 128K requires injecting a new capability through data-intensive continual pretraining.
Method
The paper continually pretrains full LLaMA-2 models on long contexts using SlimPajama’s balanced domain mixture with per-source length upsampling.
Results
1-5B tokens of continual pretraining produce balanced long-context improvements and close the gap to GPT-4 128K on retrieval.
Takeaways & Limitations
Careful domain balance and length upsampling provide an effective, affordable recipe for scaling language models to 128K contexts.
Takeaways & Limitations
The reported infrastructure configuration reaches the HuggingFace-DeepSpeed memory limit, so training at longer contexts such as 200K is left for future work.
Abstract
from arXiv · showhide
We study the continual pretraining recipe for scaling language models' context lengths to 128K, with a focus on data engineering. We hypothesize that long context modeling, in particular \textit{the ability to utilize information at arbitrary input locations}, is a capability that is mostly already acquired through large-scale pretraining, and that this capability can be readily extended to contexts substantially longer than seen during training~(e.g., 4K to 128K) through lightweight continual pretraining on appropriate data mixture. We investigate the \textit{quantity} and \textit{quality} of the data for continual pretraining: (1) for quantity, we show that 500 million to 5 billion tokens are enough to enable the model to retrieve information anywhere within the 128K context; (2) for quality, our results equally emphasize \textit{domain balance} and \textit{length upsampling}. Concretely, we find that naively upsampling longer data on certain domains like books, a common practice of existing work, gives suboptimal performance, and that a balanced domain mixture is important. We demonstrate that continual pretraining of the full model on 1B-5B tokens of such data is an effective and affordable strategy for scaling the context length of language models to 128K. Our recipe outperforms strong open-source long-context models and closes the gap to frontier models like GPT-4 128K.
1. Introduction
The paper studies whether lightweight continual pretraining and careful data engineering can extend models pretrained at shorter contexts to reliable information retrieval within 128K-token inputs. It proposes a full-model recipe using 1-5B tokens, balanced domain mixtures, and length upsampling.
- 128K-token contexts support applications including multi-document question answering, repository-level code understanding, long-history dialogue, and autonomous agents.
- The Needle-in-a-Haystack test evaluates whether models can recite a sentence inserted at arbitrary positions in a long document.Existing open-source models often fail as document length increases, while the recipe performs strongly to about 100K tokens.
- The method continues pretraining full-attention LLaMA-2 models on 64K-80K contexts, with only the RoPE base adjusted beyond the data recipe.The authors report that this is feasible with academic-level resources.
- 1-5B tokens of long-context continual pretraining can unlock precise retrieval over contexts much longer than the original pretraining length.This supports the hypothesis that arbitrary-location retrieval is mostly acquired during large-scale pretraining.
- Retaining the pretraining domain mixture while upsampling long sequences gives more stable gains than naively emphasizing long-sequence domains such as books.Book-focused upsampling can degrade perplexity in other domains, whereas the balanced strategy improves long-context tasks while maintaining short-context performance.
- The recipe substantially closes the gap to GPT-4 128K on Needle-in-a-Haystack retrieval while improving 7B and 13B LLaMA-2 long-context performance.
2. Background
Long-context models and benchmarks target capabilities beyond standard context windows, but existing approaches differ in their data and modeling choices. The paper positions balanced data composition and long-sequence upsampling as critical distinctions for retrieval at extended lengths.
- 100K-scale contexts motivate applications such as repository-level code understanding, long-history dialogue, and autonomous agents.The Needle-in-a-Haystack benchmark tests precise retrieval from arbitrary locations in long documents.
- The paper improves over strong open-source baselines and closes the gap to GPT-4 on the Needle-in-a-Haystack benchmark.Compared baselines include LongChat, Together AI’s LLaMA-2, YaRN Mistral, and LongLoRA.
- The method uses 80K continual-pretraining sequences, balanced SlimPajama domains, and long-sequence upsampling.These choices contrast with Together’s 32K training, YaRN’s book-only data, and LongLoRA’s lack of length upsampling.
- The paper hypothesizes that long-context retrieval is mostly acquired during large-scale pretraining rather than requiring hundreds of billions of additional tokens.This motivates lightweight continual pretraining on a small amount of long-context data.
3. Long Context Data Composition
The paper separates sequence length from domain composition when engineering continual-pretraining data for 128K context. It recommends per-source length upsampling because long-sequence exposure alone is insufficient and domain shifts can limit transfer.
- SlimPajama mirrors the LLaMA pretraining mixture, with 82% web data, 4.5% code, 4.5% Wikipedia, 4.5% books, 2.5% Arxiv, and 2.0% StackExchange.
- Length strategies: Cutting documents at 4K breaks naturally occurring long-range dependencies, while cutting at 128K preserves them without changing domain proportions.
- Length strategies: Per-source upsampling increases long-sequence exposure within each domain while retaining the overall domain mixture.
- Alternative strategies: Global upsampling slightly changes domain composition, whereas upsampling Arxiv, books, and Github changes both domain and length distributions.
- Empirical implications: Using the original mixture cut at 128K is insufficient for precise retrieval, and gains in one domain may not transfer or may hurt another.
4. Infrastructure and Engineering
The engineering approach makes full-attention continual pretraining at very long sequence lengths feasible by hiding quadratic attention cost behind largely linear input/output operations. The setup remains constrained by framework memory limits for still longer contexts.
- The recipe trains 7B models at 80K and 13B models at 64K sequence length using per-source upsampled data.
- 80K-context training is only 3x slower than 4K-context training because parallelized attention is largely hidden by linear data-transfer costs.
- The configuration is substantially cheaper than prior work using more than 400B continual-pretraining tokens.
- Limitations: The framework reaches a memory limit for longer contexts because most GPU memory is occupied by the KV cache, so 200K-context training is left for future work.
5. Experimental Results
The experiments show that lightweight continual pretraining with carefully engineered long-context data enables 7B and 13B LLaMA-2 models to perform well at 128K context while retaining short-context performance. Data quantity, per-source length upsampling, and balanced domains jointly determine retrieval and downstream results.
- Overall Performance: 5B tokens yield 88.0 retrieval accuracy, while scaling to 10B tokens lowers accuracy to 84.0.Accuracy rises from 37.8 at 100M tokens to 81.1 at 500M and 85.3 at 1B tokens before peaking at 5B.
- Overall Performance: The method outperforms strong open-source long-context baselines on Needle-in-a-Haystack and closes the gap to GPT-4 128K on retrieval and BookQA.The 13B model’s BookQA performance closes the gap to GPT-4 128K, while the method also maintains short-context performance as measured by MMLU.
- Data Quantity: 500M tokens unlock most retrieval accuracy, but 5B tokens are needed for generalization from the 0-80K training range to unseen 80K-128K lengths.At 10B tokens, length generalization begins to decrease, suggesting overfitting to the 80K training range.
- Data Mixture: Per-source length upsampling improves 4K-128K context losses while avoiding the cross-domain tradeoffs seen with other mixtures.Upsampling books, code, or Arxiv can improve in-domain performance but may fail to transfer or may harm another domain; per-source upsampling preserves the domain ratio while increasing long-sequence tokens.
6. Discussion
The work attributes strong long-context performance to careful data engineering and identifies open directions for longer-context training and finetuning. It also situates the approach within computational advances that make full-attention training feasible.
- Careful data engineering is credited with improving performance over strong open-source baselines.The paper emphasizes that data details can be as important as modeling and algorithmic innovations.
- FlashAttention reduces attention memory usage from quadratic to linear, enabling more feasible long-context training.
- RNN-like architectures are proposed as an open question for zero-shot generalization to longer contexts on Needle-in-a-Haystack.
- Instruction finetuning for 100K-context tasks remains limited, with no open-source instruction-finetuned 100K model reported.The authors position this work as a basis for future supervised finetuning research.
7. Conclusion
The paper presents a continual-pretraining recipe that scales language models to 128K tokens using a small amount of length-upsampled data. It reports balanced improvements and retrieval performance approaching GPT-4 128K, while motivating future long-context instruction finetuning.
- The recipe scales language models to 128K by continually pretraining the full model on 1-5 billion tokens of per-source length-upsampled data.
- The ability to use information at arbitrary locations is mostly acquired during large-scale pretraining, even with substantially shorter 4K contexts.
- Per-source length-upsampled data produces the most balanced performance improvements.
- The work closes the gap to GPT-4 128K on the retrieval task and provides a foundation for future long-context instruction finetuning.