Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum Article Swipe

MAKE IMPACT

Hadi Pouransari , Chunliang Li , Jen-Hao Rick Chang , Pavan Kumar Anasosalu Vasu , Cem Koc , Vaishaal Shankar , Oncel Tuzel · YOU? · · 2024 · Open Access · · DOI: https://doi.org/10.48550/arxiv.2405.13226

Large language models (LLMs) are commonly trained on datasets consisting of fixed-length token sequences. These datasets are created by randomly concatenating documents of various lengths and then chunking them into sequences of a predetermined target length (concat-and-chunk). Recent attention implementations mask cross-document attention, reducing the effective length of a chunk of tokens. Additionally, training on long sequences becomes computationally prohibitive due to the quadratic cost of attention. In this study, we introduce dataset decomposition, a novel variable sequence length training technique, to tackle these challenges. We decompose a dataset into a union of buckets, each containing sequences of the same size extracted from a unique document. During training, we use variable sequence length and batch-size, sampling simultaneously from all buckets with a curriculum. In contrast to the concat-and-chunk baseline, which incurs a fixed attention cost at every step of training, our proposed method incurs a computational cost proportional to the actual document lengths at each step, resulting in significant savings in training time. We train an 8k context-length 1B model at the same cost as a 2k context-length model trained with the baseline approach. Experiments on a web-scale corpus demonstrate that our approach significantly enhances performance on standard language evaluations and long-context benchmarks, reaching target accuracy with up to 6x faster training compared to the baseline. Our method not only enables efficient pretraining on long sequences but also scales effectively with dataset size. Lastly, we shed light on a critical yet less studied aspect of training large language models: the distribution and curriculum of sequence lengths, which results in a non-negligible difference in performance.

Type

preprint

Language

en

Landing Page

http://arxiv.org/abs/2405.13226

PDF

https://arxiv.org/pdf/2405.13226

OA Status

green

Cited By

1

Related Works

10

OpenAlex ID

https://openalex.org/W4398795109

Raw OpenAlex JSON

OpenAlex ID

https://openalex.org/W4398795109

Canonical identifier for this work in OpenAlex

DOI

https://doi.org/10.48550/arxiv.2405.13226

Digital Object Identifier

Title

Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum

Work title

Type

preprint

OpenAlex work type

Language

en

Primary language

Publication year

2024

Year of publication

Publication date

2024-05-21

Full publication date if available

Authors

Hadi Pouransari, Chunliang Li, Jen-Hao Rick Chang, Pavan Kumar Anasosalu Vasu, Cem Koc, Vaishaal Shankar, Oncel Tuzel

List of authors in order

Landing page

https://arxiv.org/abs/2405.13226

Publisher landing page

PDF URL

https://arxiv.org/pdf/2405.13226

Direct link to full text PDF

Open access

Yes

Whether a free full text is available

OA status

green

Open access status per OpenAlex

OA URL

https://arxiv.org/pdf/2405.13226

Direct OA link when available

Concepts

Sequence (biology), Decomposition, Curriculum, Variable (mathematics), Training (meteorology), Computer science, Mathematics education, Artificial intelligence, Psychology, Mathematics, Pedagogy, Geography, Biology, Genetics, Mathematical analysis, Meteorology, Ecology

Top concepts (fields/topics) attached by OpenAlex

Cited by

1

Total citation count in OpenAlex

Citations by year (recent)

2025: 1

Per-year citation counts (last 5 years)

Related works (count)

10

Other works algorithmically related by OpenAlex

Full payload