Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum

Authors: hadi  pouransari
chun-liang  li
jen-hao rick chang
pavan  kumar
anasosalu  vasu
cem  koc
vaishaal  shankar
oncel  tuzel
  apple
  anthropic
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 02:21:08

Abstract: 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.

Main paper: 
 introduction: 
 Large language models (LLMs) are often pretrained autoregressively (i.e., predicting the next token given a context) on large text corpora sourced from the web. Examples include The Pile [19] , RefinedWeb [46] , RedPajama [14] , and DOLMA [57] . Each of these datasets comprises multiple documents, ranging from Wikipedia articles to books and code repositories. While the individual lengths of the documents vary from a few words (e.g., a message) to hundreds of thousands of words (e.g., a book), the training infrastructure often supports only a limited sequence length in a batch. To facilitate efficient training, document chunking is necessary. In this paper, we investigate the influence of document chunking, propose alternative strategies, and evaluate the proposed strategies with careful experiments. For the largest run, some tokens are seen more than once (the dataset has approximately 525 billion tokens). For dataset decomposition, we use the Grow-P2 curriculum with 8 cycles. Maximum context length is 8192 and all hyper-parameters are the same for DD and the baseline. (a) Regular metrics average when training with the baseline method and the proposed method. We observe more than 4× data efficiency compared to the baseline. Also, even at 1.1 trillion tokens, DD has a +2.4 accuracy improvement compared to the baseline (which has a plateauing accuracy curve even on a logarithmic x-axis). (b) Comparison of model average accuracy versus training cost (GPU-hours). DD reaches the best accuracy of the baseline more than 6× faster. This is the combined effect of DD accuracy and speed gains.Recent works [43, 37, 59, 60] popularized the concat-and-chunk approach to convert text datasets with variable document lengths into sequences with a fixed target length. In this approach, during a data preparation stage before training, we first randomly shuffle and concatenate all tokenized documents. Consecutive concatenated documents are separated by a special token <EOT>, allowing the model to detect document boundaries. We then chunk the concatenated sequence into subsequences with a target sequence length. For example, 2048 and 4096 for the Llama-1 and Llama-2 models, respectively. The model is then pretrained on batches of sequences with fixed length.The concat-and-chunk approach has several shortcomings. First, randomly concatenating documents can lead to the model attending to a context from an unrelated document to predict the next token. While well-trained models learn to avoid cross-document attention, this is not explicitly enforced, leading to potential spurious modeling. Second, the cross-document attention spends unnecessary computation on attending to unrelated tokens that do not facilitate learning. This is especially crucial due to the quadratic complexity of the attention mechanism. Even with an implementation of attention that supports cross-document attention masking, the computational cost for each optimization step would be bottlenecked by the longest document in the global batch, leading to significant underutilization of devices with shorter documents. Third, even if a document is shorter than the target sequence length, it may still be broken into two chunks when they are at the boundary of two sequences. This results in significantly smaller average chunk lengths compared to the original document length average (see Fig. 3a ), which hinders the model's capability.Recent and concurrent works on LLM training try to improve the concat-and-chunk approach: document-masking is possible with recent implementation of attention [29] as adopted in some recent pre-training recipes [39] , best-fit packing [17] to reduce document chunking, and concatenating semantically related documents instead of randomly [55] . However, none of them address all three issues mentioned above together.In this work, we introduce dataset decomposition (DD), a novel approach to decompose data based on their length and train with variable sequence length (VSL) and length-based curriculum to address the above issues. We obtain significant both significant accuracy improvement and straining speed-up as shown in Fig. 1 . DD decomposes a given dataset containing documents of variable lengths into a union of datasets/buckets, each with sequences of a fixed length. Specifically, a dataset D is decomposed into buckets ∪ i D i , where each bucket D i contains sequences of length 2 i , each extracted from a unique document. During training with VSL, at every step of the optimization process, we sample i (based on a curriculum) to form a batch with b/2 i sequences from the bucket D i , which concat-and-chunk original documents dataset decomposition < l a t e x i t s h a 1 _ b a s e 6 4 = " n F p 3 a H z 6 8 n 1 O h x J h j 1 C L Z T r z G P o = " > A A A B / n i c b V D L S s N A F J 3 U V 4 2 v q L h y M 9 g I r k r S R XV Z 1 I X L C v Y B b Q i T 6 a Q d O p m E m Y l Q Q s B f c e N C E b d + h z v / x k m b h b Y e G D ic c y / 3 z A k S R q V y n G + j s r a + s b l V 3 T Z 3 d v f 2 D 6 z D o 6 6 M U 4 F J B 8 c s F v 0 A S c I o J x 1 F F S P 9 R B A U B Y z 0 g u l N 4 f c e i Z A 0 5 g 9 q l h A v Q m N O Q 4 q R 0 p J v n d i 2 O Y y Q m m D E s t v c z 9 z c t G 3 f q j l 1 Z w 6 4 S t y S 1 E C J t m 9 9 D U c x T i P C F W Z I y o H r J M r L k F A U M 5 K b w 1 S S B O E p G p O B p h x F R H r Z P H 4 O z 7 U y g m E s 9 O M K z t X f G x m K p J x F g Z 4 s k s p l r x D / 8 w a p C q + 8 j P I k V Y T j x a E w Z V D F s O g C j q g g W L G Z J g g L q r N C This approach gives us several advantages and resolves the aforementioned issues of the concatand-chunk method. First, DD is simple and has negligible computational overhead during the data preparation stage, making it easy to scale to large datasets. Second, tokens in each sequence are ensured to be from the same document by construction, which avoids cross-document attention. Furthermore, we have access to the sequence length distribution (an auxiliary prior knowledge) which can be used to create different mixtures/curricula for training. Finally, our VSL training strategy accelerates training time: the latency for one optimization step is less when sampling from D i with smaller i (due to attention's quadratic complexity). Following is a summary of our contributions:P E E C Y a U b M 3 U J 7 v K X V 0 m 3 U X e b 9 e Z 9 o 9 a 6 L u u o g l N w B i 6 A C y 5 B C 9 y B N u g A D D L w D F 7 B m / F k v B j v x s d i t G K U O 8 f g D 4 z P H 9 m l l C Q = < / l a t e x i t > D 1 < l a t e x i t s h a 1 _ b a s e 6 4 = " 7 S O 8 T k J n M 4 K 8 9 i 5 X 4 t k A N g L 5 K v k = " > A A A B / n i c b V D L S s N A F J 3 U V 4 2 v q L h y M 9 g I r k r S R X V Z 1 I X L C v Y B b Q i T 6 a Q d O p m E m Y l Q Q s B f c e N C E b d + h z v / x k m b h b Y e G D i c c y / 3 z A k S R q V y n G + j s r a + s b l V 3 T Z 3 d v f 2 D 6 z D o 6 6 M U 4 F J B 8 c s F v 0 A S c I o J x 1 F F S P 9 R B A U B Y z 0 g u l N 4 f c e i Z A 0 5 g 9 q l h A v Q m N O Q 4 q R 0 p J v n d i 2 O Y y Q m m D E s t v c z x q 5 a d u + V X P q z h x w l b g l q Y E S b d / 6 G o 5 i n E a E K 8 y Q l A P X S Z S X I a E o Z i Q 3 h 6 k k C c J T N C Y D T T m K i P S y e f w c n m t l B M N Y 6 M c V n K u / N z I U S T m L A j 1 Z J J X L X i H + 5 w 1 S F V 5 5 G e V J q g j H i 0 N h y q C K Y d E F H F F B s G I z T R A W V G e F e I I E w k o 3 Z u o S 3 O U v r 5 J u o + 4 2 6 8 3 7 R q 1 1 X d Z R B a f g D F w A F 1 y C F r g D b d A B G G T g G b y C N + P J e D H e j Y / F a M U o d 4 7 B H x i f P 9 s t l C U = < / l a t e x i t > D 2 < l a t e x i t s h a 1 _ b a s e 6 4 = " y C f M h 6 v C I Q C z Q 7 V E b 6 5 t p 8 Z w c b k = " > A A A B / n i c b V D L S s N A F J 3 4 r P E V F V d u B h v B V U k q V J d F X b i s Y B / Q h j C Z T t q h k 0 m Y m Q g l B P w V N y 4 U c e t 3 u P N v n L R Z a O u B g c M 5 9 3 L P n C B h V C r H + T Z W V t f W N z Y r W + b 2 z u 7 e v n V w 2 J F x K j B p 4 5 j F o h c g S R j l p K 2 o Y q S X C I K i g J F u M L k p / O 4 j E Z L G / E F N E + J F a M R p S D F S W v K t Y 9 s 2 B x F S Y 4 x Y d p v 7 2 U V u 2 r Z v V Z 2 a M w N c J m 5 J q q B E y 7 e + B s M Y p x H h C j M k Z d 9 1 E u V l S C i K G c n N Q S p J g v A E j U h f U 4 4 i I r 1 s F j + H Z 1 o Z w j A W + n E F Z + r v j Q x F U k 6 j Q E 8 W S e W i V 4 j / e f 1 U h V d e R n m S K s L x / F C Y M q h i W H Q B h 1 Q Q r N h U E 4 Q F 1 V k h H i O B s N K N m b o E d / H L y 6 R T r 7 m N W u O + X m 1 e l 3 V U w A k 4 B e f A B Z e g C e 5 A C 7 Q B B h l 4 B q / g z X g y X o x 3 4 2 M + u m K U O 0 f g D 4 z P H 9 y 1 l C Y = < / l a t e x i t > D 3• We introduce DD, a method to efficiently decompose a dataset of variable-length documents into a union of buckets with fixed-length sequences. DD enables efficient and robust training via VSL and length-based curriculum.• We perform large-scale experimentation using different models, datasets, and evaluation tasks to demonstrate the efficacy of the proposed method. We show (see Fig. 1 ) significant gains in data efficiency (> 4×) and compute efficiency (11% to 45%), resulting in combined LLM pretraining acceleration of up to 6× (time to reach certain accuracy compared to baseline).• Through careful experimentation, we study the importance of sequence length distribution and mixture during pretraining for different natural language and long-context tasks. We show the effect of concatenation and chunking operations to synthetically alter sequence length (Section 3.2). 
dataset decomposition: 
 Given a dataset D of tokenized documents {d 1 , d 2 , . . . , d n }, the goal of dataset decomposition (DD) is to reorganize D as a union of buckets, ∪ i D i , such that: (1) each bucket D i consists of sequences of tokens with length l i ; (2) each sequence s ∈ D i is a subsequence of one document d ∈ D; and(3) each token in D appears in exactly one D i . This decomposition produces sequences that each belong to a unique document, ensuring no cross-document attention within a sequence during training. Additionally, all sequences in a given bucket D i have the same length l i , enabling efficient batching.Dataset decomposition as defined above is not unique. We propose a specific decomposition, with l i = 2 i , to optimally maintain the original document sequence length distribution while also enabling efficient batch pretraining, as explained in Section 2.2. We apply decomposition at the document level, which makes it very easy to integrate the method into any existing data preparation pipeline (a stage before model training) and is scalable to large datasets. For a tokenized document d ∈ D with length l, where l = 2 i1 + 2 i2 + . . . + 2 i k represents its binary decomposition, we break d into k adjacent sequences s 1 , . . . , s k , with lengths of 2 i1 , . . . , 2 i k , respectively. Each sequence s j of length 2 ij is then assigned to bucket D ij . Fig. 2 shows a schematic representation of this method.With our proposed dataset decomposition approach, each bucket D i contains sequences extracted from an original document d such that the length of d is at least 2 i . In Fig. 3b , we show the distribution 1 .of RefinedWeb dataset tokens over different buckets, where D 9 (corresponding to sequences with length 512) has the maximum tokens. We also highlight the original document lengths from which tokens are extracted. Most tokens in a bucket D i are extracted from documents with length l such that 2 i ≤ l < 2 i+1 , and some tokens are rolled over from documents with length l ≥ 2 i+1 . This demonstrates the efficacy of the method in retaining original document length, especially for long documents, which are scarce.In Fig. 3a , we show the distribution of original document lengths and chunks within 2048 and 8192 target sequence lengths formed by the concat-and-chunk approach. We also present the length distribution using the bin-packing approximate algorithm introduced by a concurrent work [17] . Additionally, in Fig. 3c , we show the distribution of context length (the number of tokens from the same document a token can attend to during pretraining) when using baselines with a target sequence length of 8192 and DD. See Appendix F for additional discussion on sequence length statistics.In contrast to the concat-and-chunk approach, which results in a static dataset, DD enables us to use sequence length distribution as prior knowledge and optimize the best mixture for the target task. In Section 3.2, we show the bias of each target evaluation toward a sequence length and the effect of concatenation and chunking on model performance. In Section 3.3, we study the effect of different sequence mixtures for LLM pretraining, a less-studied topic in LLM pretraining.Following the setup in Section 2.1, we assume a set of k buckets such that D i , containing sequences with length 2 i , are available. Let b be the target batch size -the number of tokens used per optimization step. In variable sequence length (VSL) training, at every step of optimization, we first sample i from available choices, then pick b/2 i sequences from bucket D i . Since D i consists of sequences with length 2 i , the number of seen tokens per optimization step remains b, independent of the choice of i.Training LLMs with the VSL algorithm comes with several advantages.First, since the total number of seen tokens per optimization step does not change, VSL does not alter optimization dynamics, and the same hyperparameters as the baseline can be utilized (see Section 3).Second, in Section 3.1, we show that the time to complete one optimization step (forward+backward) for a fixed b (tokens per step) varies by sequence length due to the quadratic cost of attention [63] .With VSL training, the cost of every optimization step depends on the bucket D i sampled for that step (and hence the sequence length). Thus, the more expensive steps (corresponding to long sequences) are compensated with less expensive steps (corresponding to short sequences). Finally, the sampling component in VSL (which D i to choose at every optimization step) enables different curricula of sequence lengths. In Section 3.4, we show the significance of such curricula on model stability and generalization accuracy. 
experiments and analysis: 
 In this section, we show the efficacy of the proposed method to train LLMs of different sizes on largescale datasets and provide additional analyses. For all experiments, except the results in Section 3.5, we use RefinedWeb [46] filtering of Common Crawl [1] with a total of ∼ 525 billion tokens using the EleutherAI/gpt-neox [9] Positional encoding We use Rotary Positional Embedding (RoPE) [58] to encode positions in queries and keys before the attention module. RoPE rotates the consecutive components of queries and keys with a base frequency f b = 10, 000. Recent studies [48, 64, 36] have suggested increasing f b to better adapt a pretrained model for longer sequences through fine-tuning. We find that using a larger f b is also beneficial when training LLMs from scratch. In Table 4 , we show that increasing f b to 100,000 improves performance for both the baseline and DD methods.Evaluation We evaluate each model on a comprehensive set of standard benchmarks, mainly using LLM Foundry [2]. We report averaged accuracies over each category, as well as the regular average, which is the average of 14 regular language modeling benchmarks detailed below:• Commonsense Reasoning (CSR): PIQA-0-shot [8] , COPA-0-shot [52] , and OpenBookQA-10shots [40] .• Language Understanding (LU): Lambada-OpenAI [44] , Hellaswag-0-shot [65] , Winograd-3shots [30] , and WinoGrande-5-shots [54] .• Reading Comprehension (RC): SQuAD-3-shots [50] , BoolQ-0-shot [12] , and CoQA-0-shot [51] .• World Knowledge (WK): Jeopardy-3-shots [3], ArcEasy-3-shots [13] , ArcChallenge-3shots [13] , and WikiDataQA-3-shots [4] To evaluate model on longer context tasks, we adopt the following real-world benchmarks:• Multi-Document Question Answering (MDQA): We follow the exact setup as in Liu et al. [35] , where for each question from NaturalQuestions-Open [28, 27] , r Wikipedia documents are retrieved such that one of them has the answer to the question, and the other r -1 documents are distractors. We report MDQA-10, MDQA-20, and MDQA-30 accuracy corresponding to r = 10, 20, and 30, respectively. For each query, we evaluate the model by changing the location of the target document among distractors and report the averaged accuracy.• TOEFL: This dataset is a multiple-choice question answering dataset from An et al. [5] . The dataset contains QA pairs for 15 longest lectures in Tseng et al. [61] , Chung et al. [11] . Only one of the choices is the correct response. We estimate the correct choice by picking the choice with the lowest mean log probability value.• QuALITY: This dataset is a multiple-choice question answering dataset from An et al. [5] . The dataset contains a long passage for context, followed by a question with multiple choices. Only one of the choices is the correct response. We estimate the correct choice by picking the choice with the lowest mean log probability value.We first verify that VSL training enables a higher throughput than the baseline concat-and-chunk method. We enumerate model sizes (OpenLM-1B/3B/7B) and different context lengths (2 6 to 2 13 ) and measure the time to train 100 batches with a fixed global batchsize of b = 8 × 8192 distributed over 8 GPUs in a single node. We repeat this 5 times and report the average time per optimization step in Fig. 4a (with STD mostly < 1ms). See Appendix C.1 for additional results with different batchsizes b. For each model, we highlight the training time overhead (due to attention's quadratic complexity with an optimized FlashAttention2 kernel [15] ) when training with 8192 context lengths compared to 64 context lengths: +35%, +88%, and +23% for OpenLM-1B, -3B ‡ , and -7B, respectively. Training overhead grows for longer context lengths (see Fig. 7 for results up to 16k context length).The concat-and-chunk baseline method always operates at a fixed sequence length. For example, for the OpenLM-1B model, an optimization step with concat-and-chunk takes 243ms and 304ms for target context lengths of 2048 and 8192, respectively. The expected time for VSL, on the other hand, is the weighted average over different sequence lengths depending on the mixture. In Table 1 , we report the training step time for different mixtures. For example, with the natural length distribution resulting from DD (Fig. 3b ), training up to length 8192 sequences takes a similar time (244ms) as baseline training with length 2048 (with 243ms per step) per step-equivalent to a 20% training time reduction compared to baseline training with a fixed length of 8192 (with 304ms per step).In this section, we study the effect of pretraining data sequence length on model performance in isolation. Using a single bucket D i as the dataset, we train an LLM from scratch on sequences with length 2 i for a total of 2 34 seen tokens. Note that the number of tokens per optimization step is fixed at 256, irrespective of sequence length. We use the same training hyperparameters for all runs. In Appendix C.2, we show that our conclusions do not depend on the choice of hyperparameters. To ‡ OpenLM-3B has 32 heads (×2 that of OpenLM-1B), and a per-head dimension of 2560/32=80, not suitable for FlashAttention. This makes attention a significant bottleneck for this model. reduce statistical error, we train each model twice from scratch with different random seeds and report the average metric for each benchmark (observing an average standard deviation of ∼ 0.3 for regular benchmarks and ∼ 1.6 for multi-document QA). Results are demonstrated in Fig. 5a .We show a significant correlation between pretraining sequence length and different benchmarks. Specifically, the accuracy of commonsense reasoning, language understanding, and world knowledge shows an inverted U-shape behavior with respect to pretraining sequence length, while reading comprehension benefits from longer sequences. This behavior can be associated with training-test distribution alignment with respect to sequence length. In Fig. 5b , we show the length distribution for different benchmarks where RC demonstrates a heavier tail compared to CSR, LU, and WK. Multi-document QA benchmarks show a vivid correlation with respect to sequence length: test accuracy is ≈ 0 unless pretraining sequence length is greater than the test context length, which is ∼ 2k, 4k, and 6k for MDQA-10, -20, and -30, respectively.It could be argued that data selection based on sequence lengths could introduce bias since the content (or source) of the documents might change based on the sequence lengths. To better understand the effect of sequence length on common metrics, we created two new buckets, D 13→10 and D 7→10 , from existing buckets D 13 and D 7 , respectively. The bucket D 13→10 contains sequences of length 2 10 created by chunking sequences from D 13 into 8 subsequences and then performing a global shuffle.The bucket D 7→10 also includes sequences of length 2 10 , each formed by concatenating 8 random sequences from D 7 .In Fig. 5c , we compare the regular average metric of models pretrained on these buckets; for each bucket, we train two models from scratch using different random seeds and report the averaged results. D 13→10 gains 2.6 points compared to D 13 while including the same content. This demonstrates the pure effect of sequence length on model accuracy. Furthermore, training on D 13→10 underperforms D 10 by 0.9 points, even though they are of the same length, indicating that long documents (used to construct D 13→10 ) correlate less with our benchmarks than short documents (used to construct D 10 ). Finally, we show that concatenation, as opposed to chunking, does not mitigate length correlation. This is evident from the fact that D 7→10 scores the same as D 7 and still significantly worse than D 10 .Our analysis suggests that effective base model pretraining requires a mixture of different sequence lengths to perform well on all benchmarks. Next, we systematically study the effect of dataset mixture from the sequence length perspective.A key benefit of dataset decomposition is access to and control over sequence length distribution. We form datasets with different mixtures of sequence lengths and explore the performance of a model trained on each mixture. Table 1 shows the results. For all experiments, the total seen tokens and hyperparameters are fixed, and only the distribution over sequence length is changed. First, we observe that mixtures with small average context length (we provide the exact definition in Appendix F) perform poorly on MDQA, which requires long context understanding. For example, as for "1k-only", "≤2k", and "Mid" distributions that do not include long sequences from D 12 and D 13 . Larger average context length (e.g., as in "≥1k") also correlates positively with performance on reading comprehension tasks, consistent with our observation in Fig. 5a , but comes at the cost of a longer training step time.Furthermore, "1k-only", that is training using only the best sequence length (= 1024) from the study in Section 3.2 results in good performance on regular evaluations, especially for language understanding and world knowledge tasks, but is poor for long context tasks. Finally, we observe that "natural" mixture, that is aligned with the distribution resulting from dataset decomposition (see Fig. 3b ), obtains near-optimal performance on both regular and MDQA tasks, demonstrating the scalability of the proposed approach to large datasets without a need for intervention on the natural underlying length distribution.We can think of short sequences as being "easier" compared to longer ones; hence motivating a curriculum learning [7, 18] that prioritizes short sequences. A similar idea (training with image resolutions from low to high) is explored in vision to train CLIP [49] models more efficiently [33] . In VSL, we can easily implement curriculum learning through sampling designs. At every optimization step, we sample without replacement a batch with b tokens from bucket D i with probability p i . If a bucket is empty, we exclude it from sampling. We study different curricula for the "≥ 256" mixture (with an equal number of tokens in D 8 , . . . , D 13 ). Results are shown in Table 2 . For each curriculum, we determine the odds of picking a batch from each bucket (= p i 's when normalized). Details of our length-based sampling and curriculum are provided in Algorithm 1. We consider curricula that shift from short to long sequences at different paces controlled by p i 's changing linearly, with powers of 2, and with powers of 100 between buckets.Due to the presence of other hyperparameter schedules during the course of training (e.g., learning rate and weight decay), a curriculum on length may result in a potential implicit bias. For example, if we only see long sequences toward the end of training, long sequence learning occurs only when the learning rate is too small. To address this potential issue, we also explore cyclic curricula, where a curriculum is applied in cycles similar to cyclic learning rate schedules [56] as shown in Fig. 6 .Note that when we train on a sequence of length l, we have l next-token prediction losses (applied in parallel) with context lengths 0, 1, . . . , l -1. This already implies some mixing: when training on a "hard" example (i.e., a long sequence), we also include "easy" examples (its shorter sub-sequences). Therefore, even towards the end of each cycle, we still have some losses with short contexts.Our results show that the cyclic "Grow-P2" curriculum is near optimal with different metrics. An additional benefit of curriculum is training stability. Li et al. [31] noticed that long sequences contribute to extreme gradient variance, especially at the beginning of training, resulting in instability.We also observe (see Appendix E) that our proposed approach with curriculum results in more stable training dynamics, thus enabling more efficient training with larger batch sizes and learning rates.Dataset scaling In Fig. 1a , we show the performance of models trained with 2 34 , 2 35 , 2 36 , 2 37 , and 2 38 total tokens using DD and baseline. We use the "≥ 256" mixture and "Grow-Linear" curriculum with 8 cycles for DD, and a fixed target sequence length 8192 for the baseline. Results show > 2× data efficiency: our proposed method reaches the same accuracy as the baseline using less than half the tokens. Alternative dataset We demonstrate the efficacy of our proposed method on another large-scale dataset, DataComp-LM [32] . We train models with different numbers of parameters: OpenLM-160M, -410M, and -1B, for a total of 137B tokens. We compare the baseline with a DD ≥256 mixture trained with the "Grow-P2" curriculum with 8 cycles. Results are reported in Table 3 , demonstrating significant accuracy and training efficiency gains.We compare our proposed method, data decomposition, with other approaches for handling various document lengths of pretraining data, including document masking (DM), best-fit sequence packing [17] , and in-context pretraining (ICLM) [55] . We describe the details of our implementation of the best-fit packing in Appendix D. For ICLM, we use the official implementation § applied to the RefinedWeb dataset. The results are shown in Table 5 .Pre-training context length is an important factor in determining a model's long-context performance. We empirically validate this in the results shown in Fig. 5a , where models trained on longer sequences perform better on multi-document QA. Our proposed method has an average context length (as defined in Eq. ( 2)) of 1,344 for the RefinedWeb dataset, compared to 930 for the baseline (see Fig. 3c ) and 1,064 when packing [17] is applied. This explains why the dataset decomposition mixture, even without any length-based curriculum (the first row in Table 2 ), outperforms Baseline-8k-DM and Pack-8k+DM (second and third rows in Table 5 ). Here, DM refers to applying document masking during training to avoid cross-document attention.Document masking improves the baseline on regular evaluations from 51.5 to 52.4 by preventing cross-document attention. However, Xiong et al. [64] demonstrate that including concatenated unrelated documents can still enhance long-context metrics compared to training solely with shorter sequences. Therefore, DM experiences a slight decline in long-context evaluations, dropping from 27.5 to 27.1. Baseline-8k multi-document QA performance is even slightly better than our proposed dataset decomposition mixture when used without length-based curriculum (the first row in Table 2 ).In-context pre-training LMs (ICLM) [55] proposes document sorting based on content similarity.Although the benefits of ICLM with large-scale Common Crawl data (used in our experiments) are marginal in regular evaluation, we observe that ICLM results in slightly better multi-document QA performance than Baseline-8k when 30 documents are in the context compared with Baseline-8k (22.0% vs. 20.5%). The average long-context metric boosts from 27.5 for Baseline-8k to 28.7 for ICLM. However, the similarity finding step proposed by ICLM is resource-intensive at scale ¶ .Finally, as shown in in Table 2 our proposed cyclic length-based curriculum, for example, Grow-P2 with 8 cycles, results in a significant improvement in the model's long-context capability. Our proposed method avoids cross-document attention to unrelated content, maintains coherent long sequences, and benefits from a length-based curriculum, effectively improving performance in both regular and long-context evaluations compared to all baselines. We further summarize long-context performance of different methods discussed above in Table 6 . 2 and Table 5 .Recent works have raised concerns regarding cross-document attention. For example, Llama-3 [39] , ICLM [55] , and [17] , which we discussed in Section 3.6. Similarly, [26] discuss challenges with the baseline concat-andchunk approach and propose an approximate bin-packing algorithm.Related to our study on sequence length bias, [62] shows the importance of train-vs-test time distribution shift from a sequence length perspective on a string editing task. [6, 66, 25, 36] highlight the challenge of generalizing to lengths beyond what the model has seen during training and discuss the importance of positional encoding. Several works [41, 67, 23, 64, 10, 47, 48, 53, 34] address enabling LLM inference with long context (see [45] for an overview). These approaches are orthogonal to our contribution and can be applied post-pretraining to adapt to longer lengths. GrowLength [24] proposes accelerating LLM pretraining by progressively growing context length using the baseline sequence formation method, but does not show results on LLMs. Similarly, increasing sequence length has been shown in BERT model training [42] to improve compute efficiency.The idea of dynamic batching has been explored in other domains. In vision, methods like NaViT [16, 38] use images with variable resolutions (a similar concept to context length for LLMs). In seq-to-seq tasks (e.g., automatic speech recognition, text-to-speech, and neural machine translation), the inputs have different lengths. An efficient approach is to sort inputs by their length and form batches of inputs with similar lengths during training (after possible padding). Batchsize is dynamically adjusted inversely proportional to input lengths [20, 21] . Different from these works, in dataset decomposition, we do not simply put documents with similar lengths into the same bucket. Instead, we decompose each document into multiple subsequences and form multiple buckets. We form batches with different lengths during training by sampling from these buckets using a target mixture and curriculum. 
experimental result reproducibility: 
 Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We provide all experimental details in Appendix B. Further, we will be releasing the code. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
conclusion and limitations: 
 In this paper, we explore the shortcomings of a popular LLM pretraining approach, concat-andchunk, and introduce dataset decomposition, a method to decompose a dataset of text documents into buckets containing fixed sequence lengths. We show results of variable sequence training using DD with different mixtures, curricula, datasets, and models, demonstrating significant LLM pretraining speedup and a final model accuracy boost on a wide range of benchmarks. Furthermore, we provide analysis on sequence length bias and attention masking. We compare our proposed method with recent works that also address concat-and-chunk shortcomings in a unified experimental setup and show gains in data preparation cost, training time, and final model accuracy.Limitations. The training speed gains compared to the baseline are significant only when the target sequence length is long enough. Otherwise, the attention cost is not a dominant fraction of training, and hence no significant training speedup is expected.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: Our results are based on OpenLM repository (https://github.com/ mlfoundations/open_lm) and RefinedWeb [46] data, both publicly available. We will be releasing our (small) changes to the repository after acceptance. Guidelines:• The answer NA means that paper does not include experiments requiring code. • The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We do our best to appropriately cite/acknowledge all external code/data used in this work. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: Aside from code, this work does not intend to release any asset. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: This work enables faster training of LLMs, which are among the most compute-intensive applications in the field. A positive societal/environmental impact of this work is training LLMs with a smaller carbon footprint.Another potential societal advantage of this work is training LLMs with fewer hallucinations. While we did not directly measure this potential benefit, a concurrent work [17] shows such a benefit when cross-document attention is not allowed during LLM pretraining. Baseline hyper parameters We list our baseline hyperparameters in Table 12 and iterate over changes for each section next. Note that we did not explicitly optimize hyperparameters for any of the experiments, and we always use the same hyperparameters when using either the baseline method or ours. Implementation details of Section 3.2 experiments Experiments in this section are done using the same hyperparameters as in Table 12 for a total of 2 34 tokens on the OpenLM-1B model. We each model twice with different random seeds and report the averaged results. For models in this section, we use RoPE with f b = 10, 000. In Table 12 , we show that our results and conclusions in this section are not sensitive to hyperparameters, including the RoPE base frequency f b .Implementation details of Section 3.3 and Section 3.4 experiments Experiments in this section are done with OpenLM-1B model, trained for total of 96 × 10 34 ≈ 103B tokens. Hyper-parameters are the same as Table 12 , except we used 20000 warmup steps for all models presented in this section. We use RoPE with f b = 10, 000 for all models in Section 3.3 and f b = 100, 000 for models in Section 3.4.Implementation details of Section 3.5 Dataset scaling: Experiments in this section are trained with the OpenLM-1B model, RoPE with f b = 100, 000, and the baseline setup as in Table 12 except for the following changes for different dataset sizes:• total tokens = 2 34 , warmup steps = 5, 000 • total tokens = 2 35 , warmup steps = 5, 000 • total tokens = 2 36 , warmup steps = 10, 000 • total tokens = 2 37 , warmup steps = 20, 000 • total tokens = 2 38 , warmup steps = 40, 000Model scaling: Experiments in this section are trained with the OpenLM-1B, OpenLM-3B, and OpenLM-7B models, 2 37 ≈ 137B total seen tokens, RoPE with f b = 100, 000, and the baseline setup as in Table 12 except for the following changes for different model sizes:• OpenLM-1B, warmup steps = 20, 000, max-lr = 3 × 10 -3 , batchsize b = 2 19 , with 32 H100 GPUs • OpenLM-3B, warmup steps = 20, 000, max-lr = 2 × 10 -3 , batchsize b = 2 20 , with 64 H100 GPUs • OpenLM-7B, warmup steps = 20, 000, max-lr = 1 × 10 -3 , batchsize b = 2 22 , with 128 H100 GPUsAlternative dataset: Experiments in this section are trained with the OpenLM-160M, OpenLM-410M, and OpenLM-1B models, 2 37 ≈ 137B total seen tokens, RoPE with f b = 100, 000, and the baseline setup as in Table 12 except for the following changes for different model sizes:• OpenLM-160M, warmup steps = 20, 000, max-lr = 5 × 10 -3 , weight-decay = 0.033, with 16 H100 GPUs • OpenLM-410M, warmup steps = 20, 000, max-lr = 4 × 10 -3 , weight-decay = 0.066, with 16 H100 GPUs • OpenLM-1B, warmup steps = 20, 000, max-lr = 3 × 10 -3 , weight-decay = 0.1, with 32 H100 GPUsFor the DD experiments we used "Grow-P2" length curriculum which is visualized in Fig. 6 . 12 , except we used 20,000 warmup steps for all models presented in this section. We use RoPE with f b = 100, 000.We present the details of our length-based sampling and curriculum in Algorithm 1.Require:• D i : list of buckets such that D i includes sequences with length 2 i • n i : total number of tokens to be picked from each bucket (see We follow the open-book evaluation setup described in [35] . The document containing the answer is part of the context. The evaluation script provided by the official repository processes the model's response by using only the text before the first occurrence of a newline character as the answer. We noticed that sometimes the model responds with multiple newline characters before providing any valid text. In view of this behavior, we updated the evaluation script to look for the first non-empty text output from the model instead of the first string after newline character. Apart from this change in processing the model output, the rest of the evaluation follows the official implementation [35] .TOEFL We follow the setup described in [5] . As described in Section 3, the dataset contains multiple-choice QA pairs for the 15 longest lectures in [61, 11] . To obtain a response from the model, we follow MMLU-style prompting, where the choices are appended to the original prompt individually and the mean log-probability is computed for each choice. The choice corresponding to the argmax of mean log-probability is then chosen as the model's response. After we obtain the response, the computation of accuracy follows the official implementation [5] .QuALITY We follow the setup described in [5] . The dataset contains long documents with each document containing multiple-choice QA pairs. Sometimes the context for a QA pair can be longer than 8192 tokens. To account for the longer sequence length, we increase the base frequency of RoPE positional encoding from 100k to 200k without any fine-tuning. To obtain a response from the model, we follow MMLU-style prompting, where the choices are appended to the original prompt individually and the mean log-probability is computed for each choice. The choice corresponding to the argmax of mean log-probability is then chosen as the model's response. After we obtain the model output, the rest of the evaluation follows the official implementation [5] .We enumerate model sizes (OpenLM-1B, OpenLM-3B, OpenLM-7B), the number of sequences in a batch (from 1 to 256), and sequence lengths (2 6 to 2 14 ) and measure the time to train 100 batches. We repeat this 5 times and report the average and standard deviation time per batch in Fig. 7 . Notice that in the figure, each diagonal corresponds to a fixed b (number of tokens seen per optimization step).In this section, we show that changing hyperparameters does not alter our conclusions in Section 3.2. We observed that pretraining on a sequence length of 1024 results in optimal performance with respect to regular metrics, compared to both longer and shorter lengths. For example, the regular average metric is 48.0 when pretraining with a 1024 sequence length, but it is 47.0 when pretraining with a 2048 sequence length. We explore whether this gap can be filled by using potentially better hyperparameters when training with a 2048 sequence length. Results are shown in Table 13 , demonstrating that the gap cannot be simply filled by choosing a different hyperparameter and is fundamental to the choice of pretraining sequence length.Maximum Learning Rate RoPE f b Regular Average 3 × 10 -3 10,000 47.0 3 × 10 -3 100,000 47.1 10 -3 10,000 45.9 10 -2 10,000 46.5Table 13 : Sensitivity to hyperparameters for Section 3.2 experiments. All models are trained twice with different random seeds, and averaged results are reported.In this section, we show additional results for the experiments presented in Section 3.5. Table 14 shows results for dataset scaling, Table 15 for model scaling, and Table 16 for experiments on an alternative dataset. Table 16 : Small model performance trained on an improved refined-web pipeline applied to Common Crawl. All models are trained for a total of 2 37 tokens.Some recent works have employed a bin packing-based strategy [17] which aims to reduce document cross-attention by minimizing unnecessary document truncation. To achieve this, they implement a known approximation algorithm called best-fit decreasing, which packs document chunks into sequences as tightly as possible. To compare with our method, we created a new dataset based on our implementation of the best-fit decreasing algorithm and trained a new model using this dataset. We present our implementation of the best-fit decreasing algorithm, the dataset we created, and the model we trained for comparison.Given a dataset D, the input to the algorithm is a list of tokenized document chunks C = {c 1 , c 2 , . . . , c K } such that K i=1 c i = D, where each chunk is at most context size n (e.g., 2048) in length. The output of the algorithm is a list of bins B = {b 1 , b 2 , . . . , b M } such that c i ∈ b j . As a pre-processing step, we first tokenize the documents and convert them into chunks. Truncation is applied during this step only when necessary. Next, we sort the chunks from largest to smallest and start from the first chunk to pack into bins of size n. We track the remaining capacities for each bin while we iterate over the chunks. In each iteration, the algorithm finds the best bin that is both feasible and optimal for placing the chunk. Feasible bins are those that can accommodate the chunk, and optimal bins are those left with the minimum remaining capacity after placing the chunk. If such a bin is not found, we open a new bin and place the chunk inside. After all the chunks have been placed, we select the bins that have non-zero remaining capacities and fill them with pad tokens <PAD>.We process the RefinedWeb [46] dataset using the aforementioned procedure and create training sequences by concatenating all chunks in a bin. Figure 3 shows that while best-fit packing results in a higher average context length compared to the baseline concat-and-chunk, it is still much lower compared to our method dataset decomposition. Furthermore, the best-fit packing method does not prevent tokens from different documents from appearing in training sequences, whereas our method does. The presence of padding tokens in best-fit packed sequences also means that some context is wasted during each optimization step.E Training stability with VSL and curriculum [31] presents the stability-efficiency dilemma: efficient LLM pretraining with massive data parallelism results in a large batch size and requires a high learning rate. However, such a setup can result in training instability, leading to poor generalization. They observe a correlation between training instability and long sequences, especially at the early stages of training, suggesting that training on long sequences when the model is not well-trained can be a main source of training instability.Here, we show that dataset decomposition alleviates this problem when used with a curriculum: starting training by sampling more from short sequence buckets. We empirically demonstrate this by training an OpenLM-1B model from scratch with a high learning rate (= 10 -2 ) and no gradient clipping, once with baseline-8k and once with DD using the "Grow-P100" curriculum. Training loss is shown in Fig. 8 , demonstrating the stability of training with DD in comparison to the baseline. This suggests that our proposed method can also be beneficial for large-scale pretraining with large batches and high learning rates in terms of efficiency. We compute the mean of length (Fig. 3a ) and context (Fig. 3c ) distributions as follows. Assume a list of sequences with lengths l 1 , l 2 , . . . , l N , which are, for example, the chunk lengths in the concat-and-chunk approach or the sequence lengths in different buckets of the dataset decomposition approach. We define the average sequence length as follows:EQUATIONIn auto-regressive training on a sequence with length l, we apply l losses for next-token prediction on each token in parallel. Hence, for a sequence with length l, we see contexts with lengths equal to 0, 1, 2, . . . , l -1. We define the average context length, which is different from the average sequence length, as follows:Average context length =   N i=1 li-1 j=0 j   / N i=1 l i = N i=1 l i (l i -1) / 2 N i=1 l i . (2)In Fig. 3a , Fig. 3c , and Table 1 , we report the average sequence length and average context length for original documents, concat-and-chunk, and dataset decomposition with different mixtures.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced.† https://github.com/mlfoundations/open_lm§ https://github.com/swj0419/in-context-pretraining ¶ Processing 400B tokens with the official repository required over a week using 96 CPUs and 16 GPUs.|| https://github.com/mlfoundations/open_lm 


Decision: This paper proposes a data strategy for LLM pre-training, specifically, it buckets training sequences into variable lengths and then sample a batch from one bucket based on some pre-defined curriculum learning during training. This training data strategy can reduce training overhead as shown in the paper and speed up training of LLMs effectively.

Although some reviewers question about the novelty and originality of the idea as there are already public research showing the effectiveness of multi-stage pre-training on documents of increasing sequence lengths, I found the experiments and analysis conducted in this paper quite insightful for LLM training community. 

One issue authors should pay attention to is one of their motivations in the paper "this method of concatenation can lead to cross-document attention within a sequence, which is neither a desirable learning signal nor computationally efficient." Actually, modern data concatenation approach already eliminates cross-document attention by specifying document boundaries in efficient attention implementations like FlashAttention2. Thus, I think this is not a valid point in claiming the advantages of the proposed method.

Overall, I recommend acceptance of this paper as it presents a clear, simple and effective approach with solid empirical results.

Reviews: review: paper_summary: paper_summary: The paper introduces a method called dataset decomposition for training large language models (LLMs) more efficiently. Traditional LLM training processes use fixed-length token sequences, leading to inefficiencies such as unnecessary computational costs from cross-document attention. The proposed method tackles this by organizing the training dataset into various "buckets," each containing sequences of a fixed length from a unique document. This allows for variable sequence length training, where different buckets can be sampled during training based on a curriculum that adjusts for sequence length. The approach significantly reduces the attention computation overhead, leading to faster training times and improved model performance across various language understanding benchmarks. This method enables efficient and scalable LLM pretraining on large datasets, with experimental results showing up to three times faster attainment of target accuracies compared to traditional methods.
strength_weakness: strength_weakness: 1. It might be more accurate to consider that "the cross-document attention allocates significant computational resources to attending to unrelated tokens that may not directly contribute to learning." However, it's also valuable for models to develop the ability to discern and disregard irrelevant information.

2. The concept of a length-based curriculum, while insightful, isn't entirely novel. It has been explored in previous studies, such as those detailed in references [1] and [2]. Furthermore, numerous models, including BERT, employ similar curriculum learning strategies, albeit without specific emphasis on this aspect.

[1] World Model on Million-Length Video and Language with RingAttention
[2] GrowLength: Accelerating LLMs Pretraining by Progressively Growing Training Length
questions: questions: 1. How can we ascertain the outcomes when models are trained with sufficient data? Will there remain a gap between the proposed method and the baseline?

In a short summary, my major concern is the originality of the proposed strategy. Nonetheless, the empirical findings presented in this paper certainly add valuable contributions to the community.
limitations: limitations: Yes
paper_summary: paper_summary: The paper introduces 'Dataset Decomposition' (DD), a method to enhance the pre-training of Large Language Models (LLMs). Contrary to the traditional 'concat-and-chunk' approach, which can lead to unwanted cross-document attention and computational inefficiency, DD organizes datasets into buckets with sequences of uniform length from individual documents, enabling variable sequence length training. This method is demonstrated to reduce training time, improve model performance, and scale effectively with the size of the dataset.
strength_weakness: strength_weakness: The paper's writing requires improvement, as exemplified by the following:

Line 14: "Our proposed method incurs a penalty proportional to the actual document lengths at each step, resulting in significant savings in training time." Even after careful reading, it remains unclear what is meant by "penalty" here.

Line 158 (and many other instances): "We follow the exact setup as in [32]." When using citations within sentences, it is customary to use the \citet command, resulting in "We follow the exact setup as in Liu et al. (2024)," rather than using \cite or \citep.
Table 1: It would be beneficial to clearly define the meanings of the elements in the first column within the main body of the paper to prevent confusion.

Additionally, while it adds value to the paper to provide extensive discussions of various aspects through experiments, the overall organization of the paper should be enhanced to make the main focus of the paper clearer.

In Section 3.6, DD shows a significant advantage over other methods in long-context scenarios. However, the authors do not offer a detailed analysis of why this is the case. In my understanding, with a fixed number of training tokens per gradient update, DD increases efficiency by enlarging the batch size and reducing sequence length (<8k). In contrast, the baseline and ICLM maintain an 8k sequence length. It is unclear why DD performs better, as most of its training samples are actually shorter in length.
questions: questions: The same as the weaknesses
limitations: limitations: No potential negative societal impact.
paper_summary: paper_summary: The paper explores dataset decomposition for LLM pre-training. 

The method decomposes documents into subsequences and organizes them into buckets. Sequences of similar lengths are grouped in the same bucket, and different buckets have different lengths. This amounts to more efficient training. The paper investigates various mixtures of lengths separately for their impact on performance. 

The paper also explores a length-based cyclic curriculum learning - treating smaller length buckets as "easy" examples and larger ones as "hard" examples. Curriculum learning can improve the results a bit.
strength_weakness: strength_weakness: On the one hand, the proposals can be seen as important in exploring some unorthodox training structures for the specific context of LLMs and informing future pre-training. However, on the other hand, the paper seems to be mainly an exploration of hyperparameter tuning. The main proposed techniques seem like extensions of existing strategies (bucketing and curriculum learning) for LLMs. Bucketing is already understood to make things efficient, and curriculum learning has some positive results in NLP in general in earlier papers. Cyclic curriculum learning was also used in earlier works [1].

[1] CYCLICAL CURRICULUM LEARNING - Kesgin et al. ArXiv 2022

Also, if I understand correctly, the training speed gain may be less significant with flash-attention-based LLMs and alternative models like (Mamba, Linear Transformers).
questions: questions: Question:

1. Can cross-document attention be restricted with an attention mask? Can that be explored? 
2. Wouldn't a cyclic curriculum still be a problem if the learning rate starts to degrade before the first curriculum cycle? If I understand correctly, you have to run the cycles quicker than the learning rate decay? More discussion on this can be helpful. 
3. Can you elaborate more on the pacing strategy (see more on my confusion below in the suggestions)?


Suggestions:

* I did not quite understand the exact details of the pacing mechanism for the cyclic curriculum. It would be helpful to provide a pseudo-code or explain it through mathematical formalisms. The table shows a static distribution of sampling odds, but if the model is shifting from easy to hard examples, then the sampling odds should be changing, no? It's not clear to me how exactly that is being done.
limitations: limitations: Yes.
paper_summary: paper_summary: This paper proposes to carefully segment the training data so that different documents won't be mixed together. They also propose a Grow-P2 curriculum that increases training efficiency and stability.
strength_weakness: strength_weakness: This paper presents only empirical results.
questions: questions: 1. The authors only explore the training curriculum regarding lengths. There are other dimensions when it comes to curriculum design. For example, should we train the model on simpler topics first?
limitations: limitations: None.

Comments: title: Appreciate your feedback
Appreciate your response
Appreciate your response
Appreciate your response
I have read the rebuttal.
Appreciate your response
Response
comment: We would like to thank the Reviewer *z2w6* for their positive response.
Thank you for your detailed response. I'm glad to see that the writing has been improved. I have changed the score accordingly.
We are glad that Reviewer *bcif*’s concerns have been addressed and would like to thank them again for their positive feedback.
We would like to thank Reviewer *uip1* again for their positive feedback.
Thanks for the response. It solved my concerns. Although, the novelty is still questionable. I believe some empirical results from this paper is valuable to the community. Hope more discussion about those related works can be added to the revised version.

I raised my score to 5.
Thank you for providing additional details. I'll keep my positive score at 5.
We would like to thank Reviewer H2jF again for their time and positive feedback.

As suggested by the reviewer, we will include more discussion in the revision on our investigation of length curriculum in relation to other schedules, such as learning rates.

Regarding the method details: The reviewer’s understanding is correct. We do not explicitly control the pace; instead, we sample without replacement using fixed probabilities (that determine the curriculum). When a bucket is empty for the current cycle, it is excluded from sampling. Consequently, as the reviewer mentioned, mostly long sequences remain toward the end. Note that even when we train on a sequence of length $n$, we have $n$ next-token prediction losses (applied in parallel) with context lengths $0, 1, 2, …, n-1$. This implies some mixing: when training on a hard example (i.e., a long sequence), we also have easy examples (its shorter sub-sequences). Therefore, even toward the end of each cycle, we still have some losses with short contexts. As suggested, we will augment the text with additional discussion in the revision.
Thank you for the rebuttal. Overall, I increased the score to 6. 

> We respectfully disagree with this comment. We do not tune any hyperparameters in this work. The dataset (RefinedWeb), models, and training hyperparameters are all based on publicly available OpenLM repo, and are used for both the baselines and our method without any tuning.

I meant in the sense that training scheduling strategies and data sampling can be seen as a form of hyperparameter, and exploring small changes in the details of established strategies (bucketing, curriculum learning) can be in itself seen as hyperparameter tuning. I am not saying that there is any technical issue here, just that technical novelty may appear limited as a result.

> All results presented in this paper are with FlashAttentionV2, and we achieve up to 45% faster training and more accurate models (see Fig 1b) compared to the baseline, which also uses FlashAttentionV2 for sequence length up to 8192. Training speed gains will be even more when training on longer sequences.

Thank you for the clarification. 

> In all our experiments, we use a learning rate schedule with a short initial warmup followed by a one-cycle cosine learning rate decay, as in the OpenLM repository. For the length curriculum, we analyze both one-cycle and multi-cycle curricula. In all cases, the warmup period for the learning schedule is shorter than a single length curriculum. We did not observe any stability or convergence problems in any of the setups. In fact, in Appendix E, we show that our proposed curriculum significantly improves training stability compared to the baseline. To further clarify our cyclic length curricula with respect to the learning rate schedule, please see Figure 2 in the rebuttal PDF for a visualization overlapping both schedules (to be also included in the revised paper).

My point is more related to the theoretical motivation you supplied for cyclic curriculum - that due to learning rate decay, by the time the model starts encountering hard examples. 

"Due to the presence of other hyperparameter schedules during the course of training (e.g., learning
259 rate and weight decay), a curriculum on length may result in a potential implicit bias. For example, if
260 we only see long sequences toward the end of training, long sequence learning occurs only when the
261 learning rate is too small. To address this potential issue, we also explore cyclic curricula, where a
262 curriculum is applied in cycles (similar to cyclic learning rate schedules [52])."

My point is that the cyclic curriculum itself doesn't seem to be a complete solution, but it has to be balanced properly with learning rate scheduling if this motivation applies at all. Following this motivation, if you have a monotonic decay of learning rates , and it's faster than the first cycle - the same problem arises. Also, on the other hand, if you have cyclic learning rates, the theoretical motivation for a cyclic curriculum seems to diminish because it can allow harder examples to experience higher learning rates near the end without cycles. 

More can be discussed here in these regards. 

> We apologize for the confusion. We provide further clarification on our mixture and curriculum implementation below with a pseudo-code, as suggested (to be also included in the revised paper). We will also release the full code (a small patch on top of the OpenLM repo), which should clarify the details further.

Thank you. That brought some much-needed clarity on the method. 

It seems you don't strictly have any explicit pacing function to directly control the odds of sampling easier samples (starting from high odds) as time goes on. Instead, you do sampling without replacement. So once the easy samples are made early on, near the end, mostly harder samples will remain to sample from. I think this should be discussed more in the text.