Elamur Research⌕
☰
home / papers / 2024 / Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum

Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum

Paper · 2024 · Hadi Pouransari, Chun-Liang Li, Jen-Hao Rick Chang et al.

Summary

Standard LLM pretraining uses concat-and-chunk: documents are shuffled, joined, and cut into fixed-length sequences. This lets tokens attend to unrelated documents. It also wastes quadratic attention compute and splits many short documents at chunk boundaries. The Apple authors propose dataset decomposition (DD) to address these problems. Each tokenized document is split according to the binary representation of its length into adjacent pieces of size 2^i. Each piece goes into bucket D_i, so every sequence comes from a single document. Training then uses variable sequence length (VSL): each step samples a bucket under a length-based curriculum and adjusts the batch size so the token count per batch stays constant. Short-sequence steps are cheaper, so overall cost tracks real document lengths rather than the maximum context. Experiments on RefinedWeb show DD reaches the baseline's accuracy up to 6× faster and improves both regular and long-context benchmarks. The work also shows that the distribution and ordering of sequence lengths during pretraining matter noticeably, and DD scales easily to web-sized corpora.

Key points

  • When scaling OpenLM-410M on RefinedWeb to 1.1 trillion tokens with the Grow-P2 curriculum (8 cycles, 8192 max context), DD is more than 4× more data-efficient and still leads by +2.4 average accuracy at the largest run.
  • Compute savings per step range from 11% to 45%, and an 8k-context 1B model costs about as much as a 2k-context baseline.
  • At 103B training tokens, DD outperforms Baseline-8k, Baseline-8k+DM, Pack-8k+DM, and ICLM:
    • regular average: 54.4 vs. 51.5 for Baseline-8k;
    • long-context average (MDQA, TOEFL, QuALITY): 30.9 vs. 27.5;
    • step time: 244 ms vs. 304 ms.
  • DD's data preparation is cheap. By contrast, ICLM's similarity-based document ordering took over a week to process 400B tokens on 96 CPUs and 16 GPUs.
  • The speed advantage only appears when the target context is long enough for attention to dominate training cost. Code is planned at github.com/apple/ml-dataset-decomposition (NeurIPS 2024).

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, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Hadi Pouransari, Chun-Liang Li, Jen-Hao Rick Chang, Pavan Kumar Anasosalu Vasu, Cem Koc, Vaishaal Shankar and Oncel Tuzel. 2024. Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum. arXiv:2405.13226.

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

@misc{pouransari2024dataset,
  title = {{Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum}},
  author = {Hadi Pouransari, Chun-Liang Li, Jen-Hao Rick Chang, Pavan Kumar Anasosalu Vasu, Cem Koc, Vaishaal Shankar and Oncel Tuzel},
  year = {2024},
  doi = {10.48550/arxiv.2405.13226},
  eprint = {2405.13226},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2405.13226}
}

References

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

  1. Common crawl. https://commoncrawl.org. ↗
  2. Llm foundry v0.7.0. https://github.com/mosaicml/llm-foundry. ↗
  3. Jeopardy. https://huggingface.co/datasets/jeopardy. [Used custom curated version by LLM Foundry]. ↗
  4. Big-bench qa wikidata. https://github.com/google/BIG-bench/tree/main/ bigbench/benchmark_tasks/qa_wikidata. [Used through LLM Foundry]. ↗
  5. C. An, S. Gong, M. Zhong, M. Li, J. Zhang, L. Kong, and X. Qiu. L-eval: Instituting standardized evaluation for long context language models, 2023.
  6. C. Anil, Y. Wu, A. Andreassen, A. Lewkowycz, V. Misra, V. Ramasesh, A. Slone, G. Gur-Ari, E. Dyer, and B. Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems, 35:38546–38556, 2022.
  7. Y. Bengio, J. Louradour, R. Collobert, and J. Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pages 41–48, 2009.
  8. Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432–7439, 2020.
  9. S. Black, S. Biderman, E. Hallahan, Q. Anthony, L. Gao, L. Golding, H. He, C. Leahy, K. McDonell, J. Phang, et al. Gpt-neox-20b: An open-source autoregressive language model. arXiv preprint arXiv:2204.06745, 2022.
  10. S. Chen, S. Wong, L. Chen, and Y. Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595, 2023. find
  11. Y.-A. Chung, H.-Y. Lee, and J. Glass. Supervised and unsupervised transfer learning for question answering. In M. Walker, H. Ji, and A. Stent, editors, Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volu… find
  12. C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv:1905.10044, 2019.
  13. P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018.
  14. T. Computer. Redpajama: An open source recipe to reproduce llama training dataset, 2023. URL https://github.com/togethercomputer/RedPajama-Data. ↗
  15. T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023.
  16. M. Dehghani, B. Mustafa, J. Djolonga, J. Heek, M. Minderer, M. Caron, A. Steiner, J. Puigcerver, R. Geirhos, I. M. Alabdulmohsin, et al. Patch n’pack: Navit, a vision transformer for any aspect ratio and resolution. Advances in Neural Information Processing Systems, 36, 2024. find
  17. H. Ding, Z. Wang, G. Paolini, V. Kumar, A. Deoras, D. Roth, and S. Soatto. Fewer truncations improve language modeling. arXiv preprint arXiv:2404.10830, 2024. — [arXiv:2404.10830](https://arxiv.org/abs/2404.10830) find
  18. J. L. Elman. Learning and development in neural networks: The importance of starting small. Cognition, 48(1):71–99, 1993. find
  19. L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. find
  20. Z. Ge, L. Kaushik, M. Omote, and S. Kumar. Speed up training with variable length inputs by efficient batching strategies. In Interspeech, pages 156–160, 2021. find
  21. P. Gonzalez, T. S. Alstrøm, and T. May. On batching variable size inputs for training end-to-end speech enhancement systems. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 1–5. IEEE, 2023. find
  22. S. Gururangan, M. Wortsman, S. Y. Gadre, A. Dave, M. Kilian, W. Shi, J. Mercat, G. Smyrnis, G. Ilharco, M. Jordan, R. Heckel, A. Dimakis, A. Farhadi, V. Shankar, and L. Schmidt. OpenLM: a minimal but performative language modeling (lm) repository, 2023. URL https://github.com/mlfoundations/open_lm/.… ↗
  23. C. Han, Q. Wang, W. Xiong, Y. Chen, H. Ji, and S. Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv:2308.16137, 2023. — [arXiv:2308.16137](https://arxiv.org/abs/2308.16137) find
  24. H. Jin, X. Han, J. Yang, Z. Jiang, C.-Y. Chang, and X. Hu. Growlength: Accelerating llms pretraining by progressively growing training length. arXiv preprint arXiv:2310.00576, 2023. — [arXiv:2310.00576](https://arxiv.org/abs/2310.00576) find
  25. A. Kazemnejad, I. Padhi, K. Natesan Ramamurthy, P. Das, and S. Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024. find
  26. M. M. Krell, M. Kosec, S. P. Perez, and A. Fitzgibbon. Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance. arXiv preprint arXiv:2107.02027, 2021. — [arXiv:2107.02027](https://arxiv.org/abs/2107.02027) find
  27. T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, J. Devlin, K. Lee, et al. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453–466, 2019.
  28. K. Lee, M.-W. Chang, and K. Toutanova. Latent retrieval for weakly supervised open domain question answering. arXiv preprint arXiv:1906.00300, 2019.
  29. B. Lefaudeux, F. Massa, D. Liskovich, W. Xiong, V. Caggiano, S. Naren, M. Xu, J. Hu, M. Tintore, S. Zhang, P. Labatut, D. Haziza, L. Wehrstedt, J. Reizenstein, and G. Sizov. xformers: A modular and hackable transformer modelling library. https://github.com/ facebookresearch/xformers, 2022. ↗
  30. H. Levesque, E. Davis, and L. Morgenstern. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning, 2012. find
  31. C. Li, M. Zhang, and Y. He. The stability-efficiency dilemma: Investigating sequence length warmup for training gpt models. Advances in Neural Information Processing Systems, 35: 26736–26750, 2022. find
  32. J. Li, A. Fang, G. Smyrnis, M. Ivgi, M. Jordan, S. Gadre, H. Bansal, E. Guha, S. Keh, K. Arora, et al. Datacomp-lm: In search of the next generation of training sets for language models. arXiv preprint arXiv:2406.11794, 2024.
  33. X. Li, Z. Wang, and C. Xie. An inverse scaling law for clip training. Advances in Neural Information Processing Systems, 36, 2024. find
  34. H. Liu, W. Yan, M. Zaharia, and P. Abbeel. World model on million-length video and language with blockwise ringattention. arXiv preprint arXiv:2402.08268, 2024. find
  35. N. F. Liu, K. Lin, J. Hewitt, A. Paranjape, M. Bevilacqua, F. Petroni, and P. Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157–173, 2024.
  36. X. Liu, H. Yan, S. Zhang, C. An, X. Qiu, and D. Lin. Scaling laws of rope-based extrapolation. arXiv preprint arXiv:2310.05209, 2023.
  37. Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692, 2019.
  38. S. Mehta, F. Abdolhosseini, and M. Rastegari. Cvnets: High performance library for computer vision. In Proceedings of the 30th ACM International Conference on Multimedia, pages 7327–7330, 2022. find
  39. Meta. Introducing meta llama 3: The most capable openly available llm to date, 2024. URL https://ai.meta.com/blog/meta-llama-3. ↗
  40. T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02789, 2018. find
  41. A. Mohtashami and M. Jaggi. Random-access infinite context length for transformers. Advances in Neural Information Processing Systems, 36, 2024. find
  42. K. Nagatsuka, C. Broni-Bediako, and M. Atsumi. Pre-training a bert with curriculum learning by increasing block-size of input text. In Proceedings of the International Conference on Recent Advances in Natural Language Processing (RANLP 2021), pages 989–996, 2021. find
  43. M. Ott, S. Edunov, A. Baevski, A. Fan, S. Gross, N. Ng, D. Grangier, and M. Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
  44. D. Paperno, G. Kruszewski, A. Lazaridou, Q. N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031, 2016.
  45. S. Pawar, S. Tonmoy, S. Zaman, V. Jain, A. Chadha, and A. Das. The what, why, and how of context length extension techniques in large language models–a detailed survey. arXiv preprint arXiv:2401.07872, 2024. — [arXiv:2401.07872](https://arxiv.org/abs/2401.07872) find
  46. G. Penedo, Q. Malartic, D. Hesslow, R. Cojocaru, A. Cappelli, H. Alobeidli, B. Pannier, E. Almazrouei, and J. Launay. The RefinedWeb dataset for Falcon LLM: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01116, 2023. URL https://arxiv.org/abs/2306.01116. ↗
  47. B. Peng and J. Quesnelle. Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation. 2023. URL https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ ntkaware_scaled_rope_allows_llama_models_to_have. ↗
  48. B. Peng, J. Quesnelle, H. Fan, and E. Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071, 2023. find
  49. A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748–8763. PMLR, 2021.
  50. P. Rajpurkar, R. Jia, and P. Liang. Know what you don’t know: Unanswerable questions for squad. arXiv preprint arXiv:1806.03822, 2018.
  51. S. Reddy, D. Chen, and C. D. Manning. Coqa: A conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7:249–266, 2019.
  52. M. Roemmele, C. A. Bejan, and A. S. Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series, 2011. find
  53. A. Ruoss, G. Delétang, T. Genewein, J. Grau-Moya, R. Csordás, M. Bennani, S. Legg, and J. Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint arXiv:2305.16843, 2023. find
  54. K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106, 2021. find
  55. W. Shi, S. Min, M. Lomeli, C. Zhou, M. Li, V. Lin, N. A. Smith, L. Zettlemoyer, S. Yih, and M. Lewis. In-context pretraining: Language modeling beyond document boundaries. arXiv preprint arXiv:2310.10638, 2023.
  56. L. N. Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter conference on applications of computer vision (WACV), pages 464–472. IEEE, 2017. find
  57. L. Soldaini, R. Kinney, A. Bhagia, D. Schwenk, D. Atkinson, R. Authur, B. Bogin, K. Chandu, J. Dumas, Y. Elazar, V. Hofmann, A. H. Jha, S. Kumar, L. Lucy, X. Lyu, N. Lambert, I. Magnusson, J. Morrison, N. Muennighoff, A. Naik, C. Nam, M. E. Peters, A. Ravichander, K. Richardson, Z. Shen, E. Strubell… find
  58. J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
  59. H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
  60. H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
  61. B.-H. Tseng, S.-S. Shen, H.-Y. Lee, and L.-S. Lee. Towards machine comprehension of spoken content: Initial toefl listening comprehension test by machine, 2016. find
  62. D. Variš and O. Bojar. Sequence length is a domain: Length-based overfitting in transformer models. arXiv preprint arXiv:2109.07276, 2021. find
  63. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  64. W. Xiong, J. Liu, I. Molybog, H. Zhang, P. Bhargava, R. Hou, L. Martin, R. Rungta, K. A. Sankararaman, B. Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.16039, 2023.
  65. R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019. find
  66. Y. Zhou, U. Alon, X. Chen, X. Wang, R. Agarwal, and D. Zhou. Transformers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371, 2024. find
  67. D. Zhu, N. Yang, L. Wang, Y. Song, W. Wu, F. Wei, and S. Li. Pose: Efficient context window extension of llms via positional skip-wise training. arXiv preprint arXiv:2309.10400, 2023.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/pouransari-2024-dataset-decomposition-faster-llm-training/