Elamur Research⌕
☰
home / papers / 2024 / Data Engineering for Scaling Language Models to 128K Context

Data Engineering for Scaling Language Models to 128K Context

Summary

Open-source LLMs that claim support for 100K+ tokens often fail the Needle-in-a-Haystack test, which checks whether a model can recall a fact hidden anywhere in a very long document. The authors argue that the ability to use information at any position is mostly learned during ordinary pretraining, even with a 4K window, and only needs a modest continual-pretraining stage to extend it. They continue training the full LLaMA-2 7B and 13B models on 64K–80K sequences, changing only the RoPE base. The training data comes from SlimPajama. Their key data strategy is per-source length upsampling: each domain keeps its original share of the mixture, but longer documents are oversampled within every domain. With only 1–5B tokens, the 7B model matches or slightly exceeds GPT-4 128K on Needle-in-a-Haystack, holds short-context quality, and beats YaRN Mistral and LongLoRA. The study shows that careful data design, rather than costly training on hundreds of billions of tokens, can unlock long context on an academic budget.

Key points

  • Data volume: 500M–5B tokens suffice for retrieval across a 128K window. LLaMA Long and XVERSE, by comparison, spend roughly 400–500B tokens.
  • Needle-in-a-Haystack accuracy:
    • the authors' LLaMA-2 7B (trained at 80K): 88.0
    • GPT-4 128K: 87.1
    • LongLoRA 100K: 70.0
    • YaRN Mistral 128K: 57.4
    • Together LLaMA-2 32K: 27.9
    • LongChat v1.5 32K: 18.0
  • Upsampling recipe: SlimPajama keeps its domain proportions (67% CommonCrawl, 15% C4, 4.5% each Github/Wikipedia/books, 2.5% Arxiv, 2% StackExchange). Within each domain, the share of documents longer than 4K rises from about 30% to about 70%.
  • Book-only upsampling hurts: Oversampling only long sources such as books (the PG19 approach used for YaRN) worsens perplexity on other domains, including code and web text.
  • Validation loss is misleading: Models with nearly identical validation loss can differ sharply in retrieval ability. Needle-in-a-Haystack exposes gaps that loss conceals. Training was feasible thanks to FlashAttention. Code is published at github.com/FranxYao/Long-Context-Data-Engineering.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Natural language processing · AI & machine learning

Cited by

104 papers of our library cite it (of the 140884 whose bibliographies we read); 10 of them have a page here, 94 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2402.10171 · arXiv 2402.10171

@misc{fu2024data,
  title = {{Data Engineering for Scaling Language Models to 128K Context}},
  author = {Fu, Yao and Panda, Rameswar and Niu, Xinyao and Yue, Xiang and Hajishirzi, Hannaneh and Kim, Yoon and Peng, Hao},
  year = {2024},
  doi = {10.48550/arxiv.2402.10171},
  eprint = {2402.10171},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2402.10171}
}

References

42 references, 13 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Fuzhao Xue, Zian Zheng, Y. F. J. N. Z. Z. W. Z. and You, Y. Openmoe: Open mixture-of-experts language models. https://github.com/XueFuzhao/OpenMoE, 2023. ↗
  2. An, C., Gong, S., Zhong, M., Li, M., Zhang, J., Kong, L., and Qiu, X. L-eval: Instituting standardized evaluation for long context language models. arXiv preprint arXiv:2307.11088, 2023. find our card
  3. Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. 2023. URL https://api.semanticscholar.org/CorpusID:265551773. ↗
  4. Anthropic. Model card and evaluations for claude models, July 2023. URL https://www.anthropic.com/ product. ↗
  5. Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020. find our card
  6. Azerbayev, Z., Schoelkopf, H., Paster, K., Santos, M. D., McAleer, S., Jiang, A. Q., Deng, J., Biderman, S., and Welleck, S. Llemma: An open language model for mathematics. arXiv preprint arXiv:2310.10631, 2023. find our card
  7. Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022. find
  8. Bai, J., Bai, S., Chu, Y., Cui, Z., Dang, K., Deng, X., Fan, Y., Ge, W., Han, Y., Huang, F., et al. Qwen technical report. arXiv preprint arXiv:2309.16609, 2023a. find our card
  9. Jacobs, S. A., Tanaka, M., Zhang, C., Zhang, M., Song, L., Rajbhandari, S., and He, Y. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023. find our card
  10. Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508, 2023b. find our card
  11. Bairi, R., Sonwane, A., Kanade, A., Iyer, A., Parthasarathy, S., Rajamani, S., Ashok, B., Shet, S., et al. Codeplan: Repository-level coding using llms and planning. arXiv preprint arXiv:2309.12499, 2023. — [arXiv:2309.12499](https://arxiv.org/abs/2309.12499) find
  12. Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. d. l., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023. find
  13. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877–1901, 2020. find
  14. Kamradt, G. Needle in a haystack - pressure testing llms. https://github.com/gkamradt/LLMTest NeedleInAHaystack, 2023. ↗
  15. Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. find
  16. Caciularu, A., Peters, M. E., Goldberger, J., Dagan, I., and Cohan, A. Peek across: Improving multi-document modeling via cross-document question-answering. arXiv preprint arXiv:2305.15387, 2023. find
  17. Chen, M., Tworek, J., Jun, H., Yuan, Q., Pinto, H. P. d. O., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021. find our card
  18. Li, D., Shao, R., Xie, A., Sheng, Y., Zheng, L., Gonzalez, J. E., Stoica, I., Ma, X., and Zhang, H. How long can open-source llms truly promise on context length?, June 2023a. URL https://lmsys.org/blog/ 2023-06-29-longchat. ↗
  19. Chen, S., Wong, S., Chen, L., and Tian, Y. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595, 2023a. find
  20. Li, S., Xue, F., Baranwal, C., Li, Y., and You, Y. Sequence parallelism: Long sequence training from system perspective. arXiv preprint arXiv:2105.13120, 2021. find our card
  21. Chen, Y., Qian, S., Tang, H., Lai, X., Liu, Z., Han, S., and Jia, J. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309.12307, 2023b. find our card
  22. Li, S., You, C., Guruganesh, G., Ainslie, J., Ontanon, S., Zaheer, M., Sanghai, S., Yang, Y., Kumar, S., and Bhojanapalli, S. Functional interpolation for relative positions improves long context transformers. arXiv preprint arXiv:2310.04418, 2023b. find
  23. Dao, T. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023. find
  24. Mazumder, S. and Liu, B. Lifelong and Continual Learning Dialogue Systems. Springer Nature, 2024. find
  25. Peng, B., Quesnelle, J., Fan, H., and Shippole, E. Yarn: Efficient context window extension of large language models, 2023. find
  26. Together. Llama-2-7b-32k-instruct — and fine-tuning for llama-2 models with together api, August 2023. URL https://www.together.ai/blog/ llama-2-7b-32k-instruct. ↗
  27. Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021. find
  28. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023a. find
  29. Qin, Z., Yang, S., and Zhong, Y. Hierarchically gated recurrent neural network for sequence modeling. CoRR, abs/2311.04823, 2023. doi: 10.48550/ARXIV. 2311.04823. URL https://doi.org/10.48550/arXiv. 2311.04823. ↗
  30. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023b. find
  31. Rae, J. W., Potapenko, A., Jayakumar, S. M., Hillier, C., and Lillicrap, T. P. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL https://arxiv.org/abs/1911.05507. ↗ our card
  32. Weng, L. Llm-powered autonomous agents. lilian-weng.github.io, Jun 2023. URL https://lilianweng.github.io/posts/2023-06-23-agent/. ↗
  33. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485–5551, 2020. find
  34. Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023. find
  35. Xiong, W., Liu, J., Molybog, I., Zhang, H., Bhargava, P., Hou, R., Martin, L., Rungta, R., Sankararaman, K. A., Oguz, B., et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.16039, 2023. find our card
  36. Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1–16. IEEE, 2020. find
  37. XVerse. Xverse-13b, January 2024. URL https://github.com/xverse-ai/XVERSE-13B/ blob/main/README EN.md. ↗
  38. Shaham, U., Ivgi, M., Efrat, A., Berant, J., and Levy, O. Zeroscrolls: A zero-shot benchmark for long text understanding, 2023. find our card
  39. Yang, A., Xiao, B., Wang, B., Zhang, B., Bian, C., Yin, C., Lv, C., Pan, D., Wang, D., Yan, D., et al. Baichuan 2: Open large-scale language models. arXiv preprint arXiv:2309.10305, 2023a. find our card
  40. Soboleva, D., Al-Khateeb, F., Myers, R., Steeves, J. R., Hestness, J., and Dey, N. SlimPajama: A Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated 627B token cleaned and deduplicated version of linear attention transformers with hardware-efficient train-RedPajama. https://www.cerebras.net/b… ↗
  41. Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023. find
  42. Team, M. N. Introducing mpt-7b: A new standard for open-source, commercially usable llms, 2023. URL www. mosaicml.com/blog/mpt-7b. Accessed: 2023-03-28. find

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/fu-2024-data-engineering-scaling-language-models/