PaperlayerView exact revision on arXiv

Source-linked AI summary

Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum

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

arXiv:2405.13226v2cs.CLcs.LG

TL;DR

Fixed-length concat-and-chunk training can mix documents and incur attention costs that grow quadratically with sequence length. The paper introduces dataset decomposition, organizing document sequences into length-based buckets for variable-sequence training. Experiments report faster training and improved accuracy across benchmarks, with acceleration of up to 6× versus the baseline.

  • Problem

    Fixed-length concat-and-chunk can expose tokens to unrelated-document context, waste computation on unrelated tokens, and split short documents into smaller chunks.

    Source: §1 · p. 2

  • Method

    Dataset decomposition organizes document tokens into fixed-length buckets of sequences from individual documents and trains by sampling buckets with a length-based curriculum.

    Source: §1 · pp. 2–3, §2.1 · p. 3

  • Results

    Up to 6× faster time to reach a target accuracy than the baseline, with reported gains in data efficiency, compute efficiency, and final model accuracy across benchmarks.

    Source: §1 · p. 3, §5 · p. 10

  • Takeaways & Limitations

    The paper finds that sequence-length distributions and curricula during pretraining can make a non-negligible difference to performance.

    Source: §1 · p. 3

    Training-speed gains are significant only when the target sequence length is long enough, because attention cost is otherwise not a dominant share of training.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

The paper examines shortcomings of fixed-length concat-and-chunk training and introduces dataset decomposition (DD), which organizes document sequences into length-based buckets for variable-sequence training. Experiments report improved data and compute efficiency, while also studying how sequence-length distributions and curricula affect performance.

1 Introduction

  • Up to 6× faster time to reach a target accuracy, with >4× data efficiency and 11%–45% compute efficiency gains over the baseline.Figure 1 reports more than 4× data efficiency and a +2.4 accuracy improvement at 1.1 trillion tokens for OpenLM-410M; DD reaches the baseline's best accuracy more than 6× faster.Source: Figure 1 · p. 2, §1 · p. 3
  • Concat-and-chunk can create cross-document attention, waste computation on unrelated tokens, and split short documents across chunks, reducing average chunk length.These issues remain relevant even when attention masks cross-document attention, because each optimization step is bottlenecked by the longest document in the global batch.Source: §1 · p. 2
  • DD decomposes variable-length documents into fixed-length buckets of sequences from unique documents, then trains with variable sequence length and a length-based curriculum.The batch uses b/2^i sequences from bucket D_i, keeping the total number of tokens constant.Source: §1 · pp. 2–3
  • The experiments study how sequence-length distributions and mixtures affect natural-language and long-context tasks, including effects of concatenation and chunking on performance.The paper also examines sequence-length curricula during pretraining.Source: §1 · p. 3

2 Method

Dataset decomposition splits documents into fixed-length, length-based buckets while preserving each token and keeping sequences within their source document. Variable sequence length training samples from those buckets with a fixed token count per step, making step cost depend on the sampled length and allowing length curricula.

2.1 Dataset decomposition

  • Each DD bucket contains fixed-length sequences drawn from individual documents, with every dataset token appearing in exactly one bucket.This construction prevents cross-document attention within a sequence and supports efficient batching.Source: §2.1 · p. 3
  • DD uses lengths l_i = 2^i, splitting each document according to its binary length decomposition and assigning adjacent sequences to corresponding buckets.The document-level preparation is described as easy to integrate into existing data pipelines and scalable to large datasets.Source: §2.1 · p. 3
  • Most RefinedWeb tokens in bucket D_i come from documents with lengths 2^i ≤ l < 2^(i+1), while some roll over from longer documents; D_9 has the most tokens.The authors say this demonstrates retention of original document lengths, especially for scarce long documents.Source: §2.1 · p. 4

2.2 Variable sequence length training

  • Variable sequence length training samples bucket D_i and uses b/2^i sequences, so each optimization step processes the same b tokens regardless of sequence length.The authors state that this leaves optimization dynamics unchanged and permits baseline hyperparameters.Source: §2.2 · p. 4
  • Because attention has quadratic cost, optimization-step cost varies with the sampled bucket; shorter-sequence steps are less expensive than longer-sequence steps.The sampling strategy also permits different sequence-length curricula, which the authors study for their effects on stability and generalization accuracy.Source: §2.2 · pp. 4–5
  • The experiments measure average optimization-step time across context lengths and report accuracy and training-speed gains for OpenLM-1B/3B/7B models.The timing setup fixes the global batch at 8 × 8192 tokens on an 8×H100 node.Source: Figure 4 · p. 5

3 Experiments and analysis

Experiments use RefinedWeb and evaluate regular language benchmarks alongside long-context tasks. They examine training efficiency, sequence-length effects, data mixtures, curricula, model and dataset scaling, and comparisons with alternative approaches.

3.1 Training efficiency

  • 244 ms per step for natural-length DD sequences up to 8192 matches the 243 ms 2048 baseline and is 20% faster than the 304 ms 8192 baseline.The concat-and-chunk baseline uses a fixed sequence length at every step.Source: §3.1 · p. 6

3.2 Sequence length bias

  • Pretraining sequence length correlates with benchmark performance: commonsense reasoning, language understanding, and world knowledge show an inverted U-shape, while reading comprehension benefits from longer sequences.For MDQA-10, -20, and -30, accuracy is approximately zero unless pretraining sequences exceed test contexts of about 2k, 4k, and 6k, respectively.Source: §3.2 · p. 7

3.3 Data mixture

  • Mixtures with small average context length perform poorly on MDQA, while larger average context length correlates with better reading comprehension but longer training steps.The natural mixture obtains near-optimal performance on regular and MDQA evaluations; the 1k-only mixture performs well on regular evaluations but poorly on long-context tasks.Source: §3.3 · pp. 7–8

3.4 Length-based curriculum

  • Grow-P2 is near-optimal across metrics, and cyclic curricula also produce more stable training dynamics than the approaches discussed.The curriculum shifts sampling from short to long sequences at different paces; cyclic schedules address a potential bias from other training schedules, such as learning rate.Source: §3.4 · p. 8

3.5 Scaling

  • DD reaches the same accuracy as the 8192-length baseline with less than half the tokens, and results also show gains across model sizes and on DataComp-LM.Model-scaling experiments compare OpenLM-1B, -3B, and -7B; the alternative-dataset experiments report gains for OpenLM-160M, -410M, and -1B.Source: §3.5 · pp. 8–9, Table 3 · p. 9

3.6 Comparison with state-of-the-art

  • Without a length-based curriculum, the DD mixture outperforms document-masked and packed baselines on long-context performance, while document masking improves regular evaluation but slightly lowers long-context evaluation.In-context pretraining improves multi-document QA over Baseline-8k for 30 documents, but its similarity-finding step is resource-intensive at scale.Source: §3.6 · p. 9
  • The cyclic Grow-P2 curriculum improves long-context capability, and the proposed method improves regular and long-context evaluations compared with all baselines.The comparison also covers baseline training with document masking and other approaches; Table 5 reports the baseline and state-of-the-art comparison.Source: §3.6 · p. 9, Table 5 · p. 10

4 Related works

Related work addresses parts of concat-and-chunk’s shortcomings through document masking, packing, or semantically related concatenation, while other work studies sequence-length bias, long-context methods, and dynamic batching. Dataset decomposition differs from dynamic batching by splitting each document into subsequences and sampling buckets with a target mixture and curriculum.

4 Related works

  • Recent works raised concerns about cross-document attention, including Llama-3, ICLM [55], and [17]; [26] discuss challenges with the baseline concat-and-chunk approach and propose approximate bin-packing.Source: §4 · p. 10
  • Sequence-length bias, generalization beyond training lengths, positional encoding, and post-pretraining long-context adaptation are distinct related research directions.GrowLength [24] accelerates pretraining by progressively increasing context length with baseline sequence formation, but reports no LLM results.Source: §4 · p. 10
  • Unlike dynamic batching that groups similarly sized documents, dataset decomposition splits documents into subsequences and samples buckets using a target mixture and curriculum.Dynamic batching adjusts batch size inversely with input length in several sequence-to-sequence domains.Source: §4 · p. 10

5 Conclusion and limitations

The paper concludes that dataset decomposition and variable-sequence training improve training efficiency and accuracy across the reported settings. Its speed advantage is conditional on sufficiently long target sequences; the authors also identify potential societal benefits.

5 Conclusion and limitations

  • Dataset decomposition is reported to improve data-preparation cost, training time, and final accuracy in a unified comparison with recent approaches.The paper reports results across mixtures, curricula, datasets, and model sizes, including standard and long-context evaluations.Source: §5 · p. 10
  • Training-speed gains over the baseline are significant only when the target sequence length is long enough.At shorter lengths, attention is not a dominant share of training cost, so significant speedup is not expected.Source: §5 · p. 10

A Broader impacts

  • The authors identify a smaller carbon footprint as a potential societal benefit of faster LLM training.They also mention fewer hallucinations as a possible advantage but note that they did not measure it directly; concurrent work [17] reports a benefit when cross-document attention is disallowed.Source: A Broader impacts · p. 15

B.1 Training details

The appendix documents hardware and experiment settings, then describes a bucket-sampling algorithm that selects sequence lengths according to sampling odds while keeping tokens per optimization step fixed.

B.1 Training details

  • Experiments use OpenLM with PyTorch, FSDP, Bfloat16 mixed precision, Xformers attention, and nodes of eight NVIDIA H100 GPUs.The nodes have 80GB per GPU, 192 CPU cores, 2000GB RAM, and inter-node Elastic Fabric Adapter connectivity.Source: §B.1 · p. 15
  • Baseline hyperparameters are listed in Table 12 and reused across baseline and proposed-method runs without explicit hyperparameter optimization.Architecture details are provided in Tables 7–11.Source: §B.1 · p. 15, Table 7 · p. 15, Table 9 · p. 15
  • The appendix specifies different token budgets, warmup steps, learning rates, batch sizes, and GPU counts across dataset-size, model-scale, and alternative-dataset experiments.Model-scale runs cover OpenLM-1B, -3B, and -7B; the alternative-dataset runs cover OpenLM-160M, -410M, and -1B.Source: §B.1 · p. 16
  • Figure 6 plots sampled sequence lengths from 256 to 8192 under Grow-P2 with eight cycles alongside a cosine learning-rate schedule with 4k warmup steps.Bucket choice is random with curriculum-determined sampling probabilities; the plotted sequence length is sampled every nine optimization steps.Source: §B.1 · p. 16, Figure 6 · p. 17

B.2 Length based sampling and curriculum algorithm

  • The algorithm partitions each bucket into non-overlapping cycle subsets, then samples a bucket by normalized odds and draws b/2^i sequences without replacement.It uses n_i tokens per bucket, o_i sampling odds, c cycles, and b tokens per optimization step; the selected bucket contains sequences of length 2^i.Source: §B.2 · p. 17

B.3 Evaluation details

The evaluation appendix describes response processing and scoring for long-context question-answering benchmarks. It reports a small change to the MDQA output-processing script and otherwise follows the specified evaluation setups.

B.3 Evaluation details

  • For MDQA, evaluation uses the first non-empty generated text rather than the first string after a newline, while otherwise following the official implementation.The change addresses responses that sometimes contain multiple newline characters before valid text.Source: §B.3 · p. 17
  • TOEFL responses are selected by appending each choice to the prompt and choosing the option with the greatest mean log-probability.The remaining accuracy computation follows the official implementation.Source: §B.3 · pp. 17–18
  • QuALITY uses the same choice-scoring procedure, with RoPE base frequency increased from 100k to 200k when contexts exceed 8192 tokens.The evaluation otherwise follows the official implementation.Source: §B.3 · p. 18

C.1 Additional results for training efficiency

The efficiency experiment compares training time across model sizes, batch sizes, and sequence lengths, with each diagonal in Figure 7 holding tokens per optimization step fixed.

C.1 Additional results for training efficiency

  • Training time for 100 batches is measured across OpenLM-1B, OpenLM-3B, and OpenLM-7B, batch sizes from 1 to 256, and sequence lengths from 2^6 to 2^14.The figure reports average and standard deviation time per batch over five repetitions.Source: §C.1 · p. 18
  • Five repetitions are averaged to report the mean and standard deviation of time per batch.Source: §C.1 · p. 18
  • Each diagonal in Figure 7 holds the number of tokens seen per optimization step fixed.Source: §C.1 · p. 18

C.2 Additional results for sequence length bias experiments

The sequence-length experiments find that 1024-token pretraining performs best on regular metrics among the tested lengths, and the gap to 2048 is not closed by changing hyperparameters.

C.2 Additional results for sequence length bias experiments

  • 48.0 is the regular average metric for 1024-token pretraining, versus 47.0 for 2048-token pretraining.The section reports 1024 as optimal on regular metrics compared with both longer and shorter sequence lengths.Source: §C.2 · p. 18
  • Changing hyperparameters does not close the performance gap between 1024- and 2048-token pretraining, which the authors describe as fundamental to sequence-length choice.Source: §C.2 · pp. 18–19
  • Table 13 reports results averaged over two training runs with different random seeds.Source: Table 13 · p. 19

C.3 Additional results for scaling experiments

The additional results cover dataset and model scaling and an alternative dataset, alongside a comparison showing that best-fit packing retains less context than dataset decomposition and does not prevent cross-document sequences.

C.3 Additional results for scaling experiments

  • The additional scaling results include dataset scaling for OpenLM-1B, model scaling over 137B tokens, and small-model performance on an improved RefinedWeb pipeline applied to Common Crawl.Tables 14, 15, and 16 report these respective settings.Source: §C.3 · p. 19, Table 14 · p. 19, Table 15 · p. 19

D Comparison to best-fit sequence packing

  • Best-fit decreasing sorts document chunks from largest to smallest, places each in a feasible bin with the least remaining capacity, and pads unused bin capacity.If no feasible bin is found, the algorithm opens a new bin.Source: D Comparison to best-fit sequence… · pp. 19–20
  • The best-fit comparison uses tokenized document chunks of at most context size n, packed into bins of size n after truncation only when necessary.The resulting sequences concatenate the chunks in each bin.Source: D Comparison to best-fit sequence… · pp. 19–20
  • Best-fit packing yields a higher average context length than concat-and-chunk, but a lower one than dataset decomposition.The comparison uses RefinedWeb [46]; best-fit packing also permits tokens from different documents in a sequence and pads bins with unused capacity.Source: D Comparison to best-fit sequence… · p. 20

E Training stability with VSL and curriculum

With a curriculum that starts by sampling more from short-sequence buckets, dataset decomposition shows more stable training than the baseline in a high-learning-rate, unclipped setup. The section also distinguishes average sequence length from average context length.

E Training stability with VSL and curriculum

  • Prior work [31] links training instability with long sequences, especially early in training, and notes that large-batch, high-learning-rate training can be unstable.Source: E Training stability with VSL and… · p. 20
  • Dataset decomposition with the Grow-P100 curriculum shows more stable training loss than baseline-8k for OpenLM-1B at learning rate 10−2 without gradient clipping.The curriculum starts by sampling more from short-sequence buckets; the authors suggest the approach may benefit large-batch, high-learning-rate pretraining in efficiency.Source: E Training stability with VSL and… · p. 20, Figure 8 · p. 20

F Average sequence length vs average context length

Related papers

Loading 2405.13226v2…