PaperlayerView exact revision on arXiv

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

arXiv:2402.10171v1cs.CLcs.AI

TL;DR

The paper asks how data engineering can scale models pretrained on short contexts to 128K. It studies a full-model continual-pretraining recipe and reports balanced improvements that substantially close the gap to GPT-4 128K on retrieval.

  • Problem

    The paper asks how to scale models typically pretrained below 4K to 128K through continual pretraining, given the cost concerns of full attention on longer sequences.

    Source: §1 · p. 1

  • Method

    The authors continue pretraining full-attention models on 1–5B tokens of per-source-length-upsampled data to scale context to 128K.

    Source: §1 · pp. 1–2

  • Results

    1–5 billion tokens from per-source length-upsampled data gives the most balanced performance improvements and closes the gap to GPT-4 128K on retrieval.

    Source: §7 · p. 8

  • Takeaways & Limitations

    The recipe serves as a foundation for future long-context instruction finetuning research.

    Source: §7 · p. 8

    The Huggingface-DeepSpeed configuration appears limited by memory overflow at longer contexts, leaving exploration of 200K-context training for future work.

    Source: §4 · p. 4

Abstract

from arXiv · show

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 introduction presents 128K context scaling as feasible through affordable continual pretraining, with data mixture design as its central focus. It argues that preserving domain balance while upsampling long sequences enables strong long-context performance without sacrificing short-context performance.

1. Introduction

  • A 128K context window can support tasks such as multi-document question answering, repository-level code understanding, long-history dialogue, and autonomous agents; the Needle-in-a-Haystack test probes retrieval at arbitrary document positions.The test asks a model to recite a sentence placed anywhere in a 128K-token document.Source: §1 · p. 1
  • LongLoRA and YaRN-Mistral theoretically support 100K context but fail the Needle-in-a-Haystack test at those lengths, while GPT-4 128K has demonstrated strong performance.Source: §1 · p. 1
  • The authors hypothesize that arbitrary-location retrieval is mostly acquired during pretraining, and report that 1–5B tokens of long-context data can unlock precise retrieval beyond the original context length.They contrast this with approaches using 400B tokens to inject long-context modeling capabilities.Source: §1 · pp. 1–2
  • Keeping the pretraining domain mixture while upsampling long sequences within each domain avoids the suboptimal performance of upsampling long-sequence domains alone.The introduction attributes perplexity degradations in other domains to the latter strategy and links domain-balanced upsampling to stable gains on long-context tasks while maintaining short-context performance.Source: §1 · p. 2
  • Continual pretraining LLaMA-2 7B and 13B on 1–5B tokens of per-source-length-upsampled data yields strong long-context performance and substantially closes the Needle-in-a-Haystack gap to GPT-4 128K.The recipe uses full attention; Figure 1 reports strong performance up to about 100K length.Source: §1 · pp. 1–2, Figure 1 · p. 2

2. Background

The section frames precise retrieval from arbitrary locations in 100K+ contexts as a key long-context capability and positions the paper’s data-engineering approach against prior work focused on other methods. It also contrasts the paper’s lightweight continual-pretraining hypothesis with work that uses much larger training runs.

2. Background

  • Needle-in-a-Haystack tests whether a model can recite a randomly inserted sentence from an arbitrary position in a long input, distinguishing precise retrieval at 100K+ context lengths.The section describes it as a popular testbed for long-range capabilities and notes that such retrieval had not been publicly demonstrated at this scale to the authors’ knowledge.Source: §2 · p. 2
  • The work improves long-context capability over strong open-source baselines and closes the gap to GPT-4 on Needle-in-a-Haystack, as demonstrated in Fig. 1.Source: §2 · p. 2
  • The paper’s data mixture uses 80K continual-pretraining data, balanced-domain SlimPajama, and long-sequence upsampling, unlike Together’s 32K, YaRN’s book-only PG19, and LongLoRA’s lack of upsampling.The authors identify these data choices as critical to long-range retrieval performance.Source: Table 1 · p. 3, §2 · p. 3
  • Prior approaches emphasize positional embeddings in YaRN, efficient attention in LongLoRA (Chen et al., 2023b), or a full-stack solution in Together, while this work focuses on data engineering.Source: §2 · p. 3
  • Unlike LLaMA Long (Xiong et al., 2023) and concurrent XVERSE, which continue pretraining on 32K sequences for about 500 billion tokens, this work hypothesizes that about 5B tokens can extend capabilities largely acquired in base pretraining.The section also argues that similar test loss can conceal substantially different precise-retrieval behavior.Source: §2 · p. 3

3. Long Context Data Composition

The section treats document length and source domain as coupled factors in long-context data design, and recommends balancing them rather than relying on naturally long documents or selectively upsampling particular domains. Its recipes use SlimPajama to compare truncation and upsampling strategies.

3. Long Context Data Composition

  • The continual-pretraining data is SlimPajama, an open-source reproduction of LLaMA’s mixture, dominated by web data at 82%.Its other shares are 4.5% each Github code, Wikipedia, and books; 2.5% Arxiv; and 2.0% StackExchange. The authors choose it because its similarity to LLaMA’s pretraining mixture reduces concern about distribution shift.Source: §3 · p. 3
  • Document length and source domain are confounded: books and Github code are longest, while C4 and StackExchange tend to be shorter.Upsampling long sequences can therefore shift domain proportions, while changing domain proportions can shift the length distribution.Source: §3 · p. 3, Figure 2 · p. 4
  • Cutting documents at 4K preserves the original mixture but breaks dependencies in the roughly 30% of documents naturally longer than 4K; cutting at 128K retains most such dependencies without changing the mixture.The figure contrasts these length distributions with upsampling strategies.Source: §3 · p. 3, Figure 2 · p. 4
  • Per-source upsampling preserves the domain mixture while increasing long documents within each domain, and the authors recommend it for the most balanced performance gain.This strategy separates length upsampling from changes to domain proportions.Source: §3 · pp. 3–4
  • Global upsampling ignores source domains and slightly changes the mixture, while selectively upsampling Arxiv, books, or Github changes both domain and length distributions.YaRN Mistral and MPT-storyteller (Team, 2023) upsample books.Source: §3 · p. 4, Figure 2 · p. 4
  • Cutting at 128K alone is insufficient for precise long-context retrieval; long-sequence upsampling is needed, and gains in one domain may not transfer and can hurt another.The authors conclude that domain ratios must be balanced across domains.Source: §3 · p. 4

4. Infrastructure and Engineering

The engineering recipe makes long-context continual pretraining feasible: training at 80K is only 3x slower than at 4K. The Huggingface-DeepSpeed configuration is memory-limited at these lengths, and longer-context training is left for future work.

Infrastructure and Engineering

  • Training at 80K is only 3x slower than at 4K because linear or constant data-transfer costs largely hide the parallelized quadratic attention computation.Most time is spent transferring data from CPU to GPU, between GPUs via NVLink, and from HBM to SMs; these operations are constant or linear in sequence length.Source: §4 · p. 4
  • The Huggingface-DeepSpeed configuration appears to reach its framework limit: even longer contexts lead to memory overflow, since Zero Optimizer is data parallel and more GPUs cannot increase context length.Most GPU memory is taken by the KV cache; training on 200K context is left for future work.Source: §4 · p. 4
  • The recipe uses a constant 2e-5 learning rate, RoPE base adjustment, packed 80K chunks, a 4M-token batch, and 5B tokens over 2000 optimization steps.The 4M-token batch matches training at 4K context by increasing chunk length while decreasing the number of chunks per batch.Source: §4 · p. 4

5. Experimental Results

The experiments show that per-source length-upsampled continual pretraining can extend retrieval to 128K and outperform leading open-source long-context models while maintaining short-context performance.

5. Experimental Results

  • Continual pretraining a 7B LLaMA 2 took five days on 8×80G A100s, about 1% the budget of Xiong et al. (2023), which trains on 400B tokens.The reported overall configuration uses 5B training tokens.Source: §5 · p. 5, §5.2 · p. 6

5.2. Data Quantity

  • At 5B tokens, retrieval works across trained 0–80K and unseen 80K–128K lengths; 500M tokens unlock most retrieval accuracy, while scaling to 10B does not improve length generalization.At 500M–1B tokens, performance is relatively good within the 80K training context but does not generalize to 80K–128K; at 10B, generalization starts to decrease.Source: Figure 3 · p. 6, §5.2 · p. 6

5.1. Overall Performance

  • On 128K BookQA, our method outperforms LongLoRA and YaRN Mistral, and the 13B model closes the gap to GPT-4; retrieval generalizes from 80K to 128K.Mistral 7B is described as a stronger base model than the LLaMA 2 7B used here.Source: §5.1 · pp. 5–6, Table 4 · p. 6

5.3. Data Mixture

  • Per-source length upsampling preserves domain proportions while increasing sequences longer than 4K from about 30% to about 70%, yielding the most balanced mixture across domains.Upsampling a domain such as books or code can improve its in-domain performance without transferring to other domains, and may harm them; per-source upsampling improves 4K–128K losses without much sacrificing short-context losses.Source: §5.3 · pp. 7–8, §5.2 · p. 7, Table 5 · p. 7

5.1. Overall Performance

  • Similar validation loss can conceal retrieval differences: the original mixture performs worse on precise Needle-in-a-Haystack retrieval than per-source length upsampling.The evaluation also uses BookQA and MMLU; the paper notes that prior long-context benchmarks were mostly around 10K, shorter than its 128K focus.Source: §5 · p. 5, Table 3 · p. 5, Figure 4 · p. 8

6. Discussion

The discussion attributes strong long-context performance to careful data engineering and enabling systems innovations, while noting that position embeddings motivate continual pretraining and that 100K-context instruction fine-tuning remains limited.

6. Discussion

  • The authors attribute improvements over strong open-source baselines to careful data engineering, which they argue is as important as modeling and algorithmic innovations.They note that many data details crucial to long-context performance can be easily overlooked.Source: §6 · p. 8
  • FlashAttention reduces attention memory usage from quadratic to linear, while sequence parallelism could enable brute-force training of models with even longer contexts, such as 200K.The discussion identifies FlashAttention (Dao, 2023) and sequence parallelism (Li et al., 2021; Jacobs et al., 2023) as relevant systems innovations.Source: §6 · p. 8
  • Transformer position embeddings make substantial context-length generalization difficult, even with relative encodings, motivating at least some continual pretraining.The authors suggest testing whether RNN-like models can generalize zero-shot to longer contexts on Needle-in-a-Haystack.Source: §6 · p. 8
  • The work studies continual pretraining only; instruction fine-tuning for 100K-context tasks remains limited, and the authors hope this work supports future research on it.They note that repo-level code understanding is an example of such a task and that no open-source instruction-finetuned 100K-context models appear to exist so far.Source: §6 · p. 8

7. Conclusion

Continual pretraining on per-source length-upsampled data extends language models to 128K context, leveraging capabilities largely acquired during shorter-context pretraining and improving retrieval performance.

7. Conclusion

  • Full-model continual pretraining on 1–5 billion tokens of per-source length-upsampled data yields the most balanced performance improvements and closes the gap to GPT-4 128K on retrieval.The recipe also serves as a foundation for future long-context instruction finetuning research.Source: §7 · p. 8
  • The ability to use information at arbitrary locations in 128K inputs is mostly acquired during large-scale pretraining, even with a 4K pretraining context.Source: §7 · p. 8

Related papers

Loading 2402.10171v1…