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
Venue: NeurIPS 2024 poster (2024)
Keywords: Large Language Models
Efficient Training
Learning Efficiency
Long Context
Primary area: natural_language_processing
TL;DR: Introducing dataset decomposition to optimize training of large language models using mixture of sequence lengths, significantly improving learning efficiency (up to 3 times faster to reach a certain accuracy).

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.

Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum
Hadi Pouransari
Chun-Liang Li
Jen-Hao Rick Chang'
Pavan Kumar Anasosalu Vasu'
Cem Koe
Vaishaal Shankar Oncel Tuzel
AApple *Anthropic
Abstract
Large language models (LLMs) are commonly trained on datasets consisting of fixed-length token sequences. These datasets are created by randomly concate- nating documents of various lengths and then chunking them into sequences of a predetermined target length (coneat-and-chunk) Recent attention implementations mask cross-document attention, reducing the effective length of chunk of tokens. Additionally, training on long sequences becomes computationally prohibitive due the quadratic cost of attention. this study, we introduce dataset decomposition, a novel variable sequence length training technique, to tackle these challenges. We decompose dataset into a union buckets, each containing sequences of the same size extracted from a unique document. During training, use variable sequence length and batch-size, sampling simultaneously from all buckets with a curriculum. In contrast to the concat-and-chuue baseline, which incurs a fixed attention cost at every step training, our proposed method incurs a computational cost propor- tional the actual document lengths at each step, resulting in significant savings in training time. We train an 8k context-length IB model at the same cost as a 2k context-length model trained with the baseline approach. Experiments on web- scale corpus demonstrate that our approach significantly enhances performance 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 non-negligible difference in performance
1 Introduction
Large language models (LLMs) are often pretrained autoregressively (i.e., predicting the next token given context) on large text corpora sourced from the web. Examples include The Pile [19], RefinedWeb [46], RedPajama [14,, and DOLMA [37]. Each of these datasets comprises multiple documents, ranging from Wikipedia articles books and code repositories. While the individual lengths of the documents vary from a few words (e.g., message) to hundreds of thousands of words (e.g.. book), the training infrastructure often supports only limited sequence length in batch. To facilitate efficient training, document chunking is necessary. In this paper, we investigate the influence document chunking, propose alternative strategies, and evaluate the proposed strategies with careful experiments.
cCorresponding author: mpouransarri@apll..omm. Work is done when at Apple. *Code to be available at ott.tp/s/ttt.a...ccoooocc.aaassp
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

+2.4
Regular eval (%)
training speed-up
Regular eval (%)
efficiency
40-
40
Baseline-Sk
Baseline-äk Dataset Decomposition
Dataset Decomposition
gts
24
2
2
2/u
2
1000
2000
3000
4000
5000
of seen tokens
GPU-Hours
(a) Data Efficiency
(b) Computational Efficiency
Figure 1: Scaling the training of the OpenL.M-410M model 1.1 trillion tokens the RefinedWeb dataset. Note that each point on the figure represents a separate training from scratch (not different checkpoints of single run). 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 the baseline. (a) Regular metrics average when training with the baseline method and the proposed method. We observe more than 4x 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 6x faster. This is the combined effect of DD accuracy and speed gains.
the same for DD and
Recent works 59 60 popularized the concat-oe concat-and-chen and-chunk approach to convert text datasets with variable document lengths into sequences with fixed target length. In this approach, during data preparation stage before training, we first randomly shuffle and concatenate all tokenized documents. Consecutive concatenated documents are separated by special token <EOT>, allowing the model to detect document boundaries. We then chunk the concatenated sequence into subsequences with target sequence length. For example, 2048 and 4096 for the Llama-1 I and Llama-2 models, respectively. The model is then pretrained on batches of sequences with fixed length.
The aoncat-and-chunk approach has several shortcomings. First, randomly concatenating documents can lead to the model attending While well-traines models learn to avoid coos-document attention, this is not explicitly enforced, leading to potential spurious modeling. Second, the cross-document attention spends unnecessary computation 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, Icading to significant under- utilization of devices with shorter documents. Third, even if 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. Ba which hinders the model's capability.
context from an unrelated document to predict the next token.
Recent and concurrent works on LLM training try to improve the concat-and-chunk approach: document-masking possible with recent implementation of attention [2] 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 decompose data based on their length and train with variable sequence length (VSL) and based curriculum to address the above issues. We obtain significant both significant accuracy improvement and straining speed-up as shown in Fig. [ 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 U,D,, where each bucket D; contains sequences of length 2 each extracted from unique document. During training with VSI., at every step of the optimization process, we sample (based on a curriculum) to form batch with 6/2 sequences from the bucket D,, which

---

D
D2
D3
original documents
concat- aed---addddaan
dataset decomposition
Figure 2: Each cell in the figure represents token. Left:
Original documents with variable lengths.
Middle: Concat-and-chun baseline to form sequences with fixed target length (here = 4). Right: Dataset decomposition method with D1, D2, and Dy buckets.
keeps the total number of tokens in batch constant (2' x 6/2' 6). regardless of which 2) is sampled.
This approach gives several advantages and resolves the aforementioned issues of the concat- and-chunk method. First, DD is simple and has negligible computational overhead during the data preparation stage, making it easy 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 used to create different mixtures/curiiclla for training. Finally, our VSL training strategy accelerates training time: the latency for one optimization step is less when sampling from D; with smaller (due to attention's 's quadratic complexity). Following is summary of our contributions:
We introduce DD, method efficiently decompose a dataset of variable-length documents into 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. D significant gains in data efficiency (> 4x) and compute efficiency (11% to 45%), resulting in combined LLM pretraining acceleration of up to 6x (time to reach certain accuracy compared to baseline).
Through careful experimentation, mixture during pretraining for different natural language and long-conten tasks. We show the effect of concatenation and chunking operations to synthetically alter sequence length (Section 3.2)
study the importance of sequence length distribution and
2 Method
2.1 Dataset decomposition
Given dataset 'D of tokenized documents {di, d2,..., [.,]) the goal of dataset decomposition (DD) is to reorganize D as a union of buckets, U,D, such that: (1) each bucket D; consists of sequences of tokens with length 4; (2) each sequence E D, is subsequence of one document D; and (3) each token in D appears in exactly one D, This decomposition produces sequences that each belong to unique document, ensuring no cross-document attention within sequence during training. Additionally, all sequences in given bucket D, have the same length 4, enabling efficient batching.
Dataset decomposition defined above is not unique. We propose specific decomposition, with 4 = 2', to optimally maintain the original document sequence length distribution while also enabling efficient batch pretraining, explained in Section 2.2 We apply decomposition the document level, which makes it very easy to integrate the method into any existing data preparation pipeline stage before model training) and scalable to large datasets. For a tokenized document d D with length I, where 1 = 2 + 2 + + 2 represents its binary decomposition, we break d into adjacent sequences , Sk. with lengths of 2 ..... 2dx, respectively. Each sequence 8, of length 26 is then assigned bucket D,,- Fig. shows schematic representation of this method.
With our proposed dataset decomposition approach, each bucket D; contains sequences extracted from original document d such that the length of d is at least 2 In Fig. 3b we show the distribution

---

10rl
3.8%
10
p.8%%
10-1
10
P0F
10
mcon-158
P06
101
BussTin-2k
Basdineedk
10m5
4483
maan--300
mese- 1064
10m9
10"
558
1344
10-7
2
git
L0 15
16
24
qu g0
Sequenon bengh
Bucket number
Contect length
(a)
(b)
(c)
Figure 3: For the RefinedWeb dataset [46] (a) Distribution of chunk lengths using different dataset preparation methods. Peaks show the percentage of chunks for each method with the same length as the target sequence length. (b) Distribution of tokens over D; 's in DD. Color/pattern shows the [log2 J,, where the length of the original document each token is extracted from. (c) Probability distribution of context length (number of tokens from the same document a token can attend to) observed during training for the concat-and-chunk baseline with target sequence length 8192 and DD) with 2 256 mixture defined in Table
of RefinedWeb dataset tokens over different buckets, where Dg (corresponding to sequences with length 512) has the maximum tokens. We also highlight the original document lengths from which tokens extracted. Most tokens in bucket D) are extracted from documents with length I such that $ 2 and some tokens are rolled over from documents with length 2 2 This demonstrates the efficacy of the method in retaining original document length, especially for long documents, which are scarce.
In Fig. Ba we show the distribution of original document lengths and chunks within 2048 and 8192 target sequence lengths formed by the concat-ande aaneat-and-chun approach. We also present the length distribution using the bin-packing approximate algorithm introduced by concurrent work [17]. Additionally, Fig. 3c] we show the distribution context length (the number of tokens from the same document token can attend to during pretraining) when using baselines with target sequence length 8192 and DD. See Appendix for additional discussion on sequence length statistics.
In contrast to the concat-and-chee sequence length distribution prior knowledge and optimize the best mixture for the target task. In Section 3.2) we show the bias of each target evaluation toward sequence length and the effect of concatenation and chunking on model performance. In Section 3.3] study the effect of different sequence mixtures for LLM pretraining, less-studied topic in LLM pretraining.
approach, which results in a static dataset, DD enables us to use
2.2 Variable sequence length training
assume set of buckets such that Di, containing sequences with length 2', 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, available choices, then pick b/2 sequences from bucket D; Since D; consists of sequences with length 2'. 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 ad advantages.
Following the setup in Section 2.1
first sample i from
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).
we show that the time to complete one optimization step (forward+backward) for a fixed (tokens per step) varies by sequence length due to the quadratic cost attention (63] With VSL training. the cost of every optimization step depends on the bucket DD sampled for that step (and hence the sequence length). Thus, the more expensive steps (corresponding to long sequences) compensated with less expensive steps (corresponding to short sequences).
Second, in Section 3.1

---

63
1000
Baseline 8k Ours
+2.7
88%
Openl
+1.5
One training
M-38
OpenLM-78
56
800
35%
Regator
52.5
+3.0
20% fonter
45%
11%
200
25
21
212
2.
19
Conteot longth
Number of model parameters
(a)
(b)
Figure 4: (a) Average time for one optimization step (b = x 8192 tokens) on an 8xH100 node with FSDP and FlashAttentio22 for different context lengths. (b) OpenLLM-IB/3B/1B models trained on 137B tokens. Accuracy and training speed gains are shown.
Finally, the sampling component in VSL (which D, to choose at every optimization step) enables different curricula of sequence lengths. In Section 3.4 show the significance of such curricula on model stability and generalization accuracy.
3 Experiments and analysis
In this section, show the efficacy of the proposed method train LLMs of different sizes large- scale datasets and provide additional analyses. For all experiments, except the results Section use RefinedWa [46] filtering of Common Crawl [ with total of ~ 525 billion tokens using the ElggtterA//gtttneoo 9 tokenizer (vocabulary size is 50,432). Model architectures and training code based the OpenLM [22 For all experiments, other than model scaling in Section 3.5
the OpenI.M18 model with an 8k context length. Please refer to Appendix B for implementation details of all experiments.
Positional encoding We 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 base frequency fo 10, 000. Recent studies 48 64 36) have suggested increasing fo to better adapt a pretrained model for longer sequences through fine-tuning. We find that using a larger fo is also beneficial when training LLMs from scratch. In Table 4 we show that increasing fb to 100,000 improves performance for both the baseline and DD methods.
Evaluation We evaluate each model on comprehensive set standard benchmarks, mainly using LLM Foundry (2). We report averaged accuracies over each category, well as the regular average, which is the average of 14 regular language modeling benchmarks detailed below:
Commonsense Reasoning (CSR): PIQA-O-shot
8 COPA-0-shot
[52] and OpenBookQA-10-
shots (40]
Language Understanding (L.U): Lambada-Opent
44), Hellaswag-O-sho 65
Winograd-3-
shots 30), and WinoGrande-5-shon 54
Reading Comprehension (RC): SQuAD--3shotts 50
BoolQ-O-shot 12 and CoQA-0-shot
[1
World Knowledge (WK): Jeopardy-3-shee
3 ArcEasy-3-shots
13
ArcChallenge-3-
shots [13, and WikisataQA---shoos
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 Liu et al. B5
where for each question from NaturalQQusttonn- [28 (2)) Wikipedia documents retrieved
such that one of them has the answer to the question, and the other r documents are distractors. We report MDQA-10, MDQA-20, and MDQA-30 accuracy corresponding to 10, 20, and 30,
5

---

CSR
MDOM-1
-0-
MDON-2
M000-13
LU
WX
Dio
WK
D1a
44.5
+2.6
MDQA-10
D55-10
0.9
MDQA-20
D7 40
MDQA-30
2
Pretuieng accemnos kmggh
211
233
123
21
Sequence length
29
20
git gi
38
40.5
Regular avenge
43
45.5
48
(a)
(b)
(c)
Figure 5: (a) Performance OpenLM-1BB model trained on 2 tokens from buckets with different sequence lengths. (b) distribution of lengths of documents for different benchmarks. (c) Effect of chunking (D13-s10) and concatenating (D7-13) sequences during pretraining on model performance.
respectively. For each query, we evaluate the model by changing the location of the target document among distractors and report the averaged accuracy.
This dataset dataset contains QA pairs for 15 longest lectures in Tseng et al. [6]], Chung et al. [TT. Only one of
TOEFL:
multiple-choice question answering dataset from An al. [5]. The
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 a multiple-choice question answering dataset from An et al. [5]. The dataset contains long passage for context, followed by 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.
3.1 Training efficiency
We first verify that VSL training enables a higher throughput than the baseline concat-and-chunk. method. We enumerate model sizes (0penLLM-1B/3B/7B and different context lengths (2 to 213) and measure the time to train 100 batches with a fixed global batchsize of b = 8192 distributed over GPUs in 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. C.1 for additional results with different batchsizes b. For each model, highlight the training time overhead (due attention's quadratic complexity with an optimized Flaetonte kernel [15)) when training with 8192 context lengths compared to
64 context lengths: +35%, +88%, and +23% for OpenLM-IB, 3 and -BB, 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 fixed sequence length. For example, for the OpenL-M1B model, an optimization step with concat-and-chume takes 243ms and 304ms for target context lengths of 2048 and 8192, respectively. The expected time for VSL, the other hand, the weighted average over different sequence lengths depending on the mixture, In Table [ we report the training step time for different mixtures. For example, with the natural length distribution resulting from DD (Fig. (b). training up to length 8192 sequences takes a similar time (244ms) as baseline training with length 2048 (with 243ms per step) per hep-eqquvaleen reduction compared to baseline training with a fixed length 8192 (with 304ms per step).
20% training time
3.2 Sequence length bias
In this section, we study the effect of pretraining data sequence length on model performance in isolation. Using single bucket D, as the dataset, we train an LLM from scratch sequences with length 2' for a total of 2 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 (x2 that of OpenLM- for FlashAttention. This makes attention significant bottleneck for this model.
-BL. 1B). and hea-head dimension of 2560/32=80, not suitable

---

tokens
Avg.
Avg.
Step
Regular
MDQA
Name
D2
Dy
Dz
Dg
Di
Pi
Pia
Dia
Len
Seq.
Ctx.
Len
Time
(186)
CSR
LU
RC
WK
Avg.
10
20 30 Avg.
Natural
21.9
Equal
12
12
12
12
12
12
12
257
1020
244
61.9
64.3
43.1
43.5
53.3
25.1
21.4
17.4
21.3
Ik-only
0
0
0
96
0
0
0
1024
512
234
60.8
66.4
43.2
44.7
54.0
0.2
0.1
0.2
0.2
52k
16
16
0
195
336
231
62.8
63.7
43.5
53.1
23.5
0.4
0.4
16
53.8
18.4
17.2
Mid
24 24
24 0
546 480
233
61.2
65.5
42.5
43.8
53.6
19.1
0.0
0.1
6.4
211k
24
24
24
24
2185
1920
263
61.9
65.0
45.8
43.3
54.0
26.7
21.6
18.1
22.1
Table I: Effecc of pretraining dataset mixture on model performance. Each row corresponds to model trained on a specific mixture of dataset decomposition buckets. All models are OpenL have seen a total of 96 200 tokens, use RoPE with base frequency of 10k, and are trained with the same hyperparameters. The definition of average context length is given in Appendix FF
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-tee 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 = O 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 metries, we created two new buckets, D13-10 and D7-10- from existing buckets D13 and D7, respectively. The bucket P13-10 contains sequences of length 2 created by chunking sequences from D13 into 8 subsequences and then performing a global shuffle. The bucket D7-110 also includes sequences of length 2 each formed by concatenating 8 random sequences from D7.
In Fig- [c] 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. D13--10 gains 2.6 points compared Dia while including the same content. This demonstrates the pure effect of sequence length on model accuracy. Furthermore, training on D13-10 underperforms D10 by 0.9 points, even though they are of the same length. indicating that long documents (used to construct D13-100 correlate less with our benchmarks than short documents (used to construct Djo). Finally, we show that concatenation, as opposed to chunking, does not mitigate length correlation. This is evident from the fact that D7>+10 scores the same as D77 and still significantly worse than D10-
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.
3.3 Data mixture
key benefit of dataset decomposition is access and control over sequence length distribution. We form datasets with different mixtures of sequence lengths and explore the performance of model trained on each mixture. Table[]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", "s2k", and "Mid" distributions that do not include long sequences from D12 and D13. Larger average context length (e.g., as in "> Ik") also correlates positively with performance on
7

---

Sampling Odds
Num. Cycles
Regular RC
MDQA
Name
Dg
Dg
Dio
Du
D12
CSR
LU
WK
Avg.
10
20
30
Avg.
Uniform
1
1
62.2
65.2
43.4
44.0
53.8
27.3
22.0
19.6
23.0
60.9
64.2
46.6
42.9
53.6
30.9
26.0
23.9
Grow-Linear
6
5
4
3
2
1
8
62.7
65.0
45.4
44.7
54.5
30.1
25.3
22.8
26.
Grow-P2 32
16
8
4
2
1
64.3
44.1
54,0
25.0
8
62.8
65.2
45,3
44.2
54.4
32.3
26.9
24.6
28.0
Grow-P100
100*
100*
1004
100
100
1
65.0
43.4
54.1
30.6
26.9
23.5
27.0
8
63.2
65.4
46.3
44.6
54,9
30.2
23.2
18.9
24.1
Shrink-P100 1 100 1002
100
100
100
60.0
62.2
37.6
40.7
50.3
24.5
18.7
15.6
19.6
Table 2: Effect of length- based curriculum. All models are OpenI.--11 1B and have seen total of 96 x 230 tokens, with exactly 2 tokens from each D, for i = .., 13. We use RoPE with a base frequency of 100k and the same default hyperparameters.
reading comprehension tasks, consistent with our observation in Fig. 5a but comes at the cost of a longer training step time.
Furthermore. "Ik-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 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 need for intervention on the natural underlying length distribution.
3.4 Length- based curriculum
We can think of short sequences as being "easier" compared longer ones; hence motivating curriculum learning 77 [8] that prioritizes short sequences. A similar idea (training with image resolutions from low high) is explored in vision 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 tokens from bucket D; with probability ps. If bucket is empty, we exclude it from sampling. We study different curricula for the "2 256" mixture (with an equal number of tokens in De, Dia). Results are shown in Table For each curriculum, we determine the odds of picking batch from each bucket (= p.'s ps's when normalized). Details of our length- based sampling and curriculum are provided Algorithm I We consider curricula that shift from short to long sequences at different paces controlled by p, '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 (c.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 curriculum is applied in cycles similar to cyclic learning rate schedules [56] as shown in Fig. 6 Note that when we train on sequence of length we have / next-token prediction losses (applied in parallel) with context lengths 0, 1, 1. This already implies some mixing: when training "hard" example (i.e., a long sequence), we also include "easy" examples (its shorter sub- 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. [3]] noticed that long sequences contribute to extreme gradient variance, especially at the beginning of training, resulting in instability. We also observe (see Appendix(Ei 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.
3.5 Scaling
Dataset scaling In Fig. la we show the performance of models trained with 234 2 2,77 and 23 total tokens using DD and baseline. We use the "> 256" mixture and "Grow-Linear" curriculum with 8 cycles for DD, and fixed target sequence length 8192 for the baseline. Results show > 2x data efficiency: our proposed method reaches the same accuracy as the baseline using less than half the tokens.
8

---

Method
(bours)
Regular Avg.
MDQA
Avg.
Method
fo
Regular Avg.
MDQA
Avg.
160M
Hascline-Sk
DD
16
18.3
15.7
-14%
40,0
39.3
+0.7
11.4
9.7
+1.7
Buseline-8k
51.3
410M
Haseline-S Sk
DD
16
29.6
58.9
-24%
49.4
48.3
+1.1
14.8
18.8
+4.0
100k
10k
51.5
53.8
24.4
20.1
Maseline
44.4
56.7
25.6
DD
100k
53.8
24,9
DD
-20%
25.6
Table Comparing baseline training with DD on an alter- native pretraining dataset and model sizes.
Table 4: Effect of RoPE base fir-- quency, ft. in pretraining.
Model scaling We report results on OpenLM-IB. -3B, and -7B trained from scratch for a total of 237 tokens in Fig. 4bb We compare baseline training with a fixed target sequence length 8192 and VSL training with a DD2256 2256 mixture and the "Grow-Linear" curriculum with 8 cycles. Training with DD results in significant accuracy gains and reductions in training wall-clock time at different scales,
Alternative dataset We demonstrate the efficacy of our proposed method another large-scale dataset, Data [32). We train models with different numbers of parameters: OpenLM- -410M, and -BB, for a total of 137B tokens. We compare the baseline with a DD2220 mixture trained with the "Grow-P2" curriculum with 8 cycles. Results are reported in Table 3 demonstrating significant accuracy and training efficiency gains.
3.6 Comparison with state-of-the-art
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 pack- ing [17], and in in-context pretraining (ICLM) [55]. We describe the details of 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
Pre-traine 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
QA. Our proposed method has an average context length (as defined in Eq. ( of 1,344 for the Web dataset, compared to 930 for the baseline (see Fig. 3cp and 1, 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 ( outperforms Baseline-8k-DD and Pack-8k+DM (sccond and third rows in Table 55. 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 52.4 by preventing cross-document attention. However, Xiong et al. 64 demonstrate that including concatenated unrelated documents can still enhance context metrics compared to training solely with shorter sequences. Therefore, DM experiences a slight decline in long-come long-context evaluations, dropping from 27.5 to 27.1. Baseline-8K 8k multi-document QA performance 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 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-88 8k to 28.7 for ICLM. However, the similarity finding step proposed by ICLM is essource-intensig at scale
Finally, shown in in Table our proposed cyclic length-based curriculum, for example, Grow-P2 with 8 cycles, results in significant improvement in the model's long-context capability. Our proposed method avoids cross-document attention unrelated content, maintains coherent long sequences, and benefits from a length- curriculum, effectively improving performance both regular and long-context evaluations compared to all baselines. We further summarize long-context performance of different methods discussed above in Table
fhttps: / //gi ////gttalt thub. com sw j 04 9/ www.-oottteeeeeeeeeeeeee Processing 400B tokens with the official repository required over week using 96 CPUs and 16 GPUs.

---

Regular
Long Context
Step Time
Data
Method
CSR
LU
RC
WK
Avg.
10
MDQA
20
30
TOFFL
QuALITY
Avg.
(ms)
Prep.
Cost
Buseline-S Sk
60.6
62.5
41.5
41.3
51.5
29.0
23.8
20.5
26.2
32.0
27.5
304
S
Baseline-8k+DMM
60.2
64.1
42.8
41.8
52.4
24.4
20.0
16.0
29.2
32.0
27.1
304
$
Pack- Buck-SkDDM [7
60.3
64.0
44.6
41.8
52.7
25.6
19.8
16.9
29.2
33.1
27.7
304
S$
ICI.M (55
60.6
62.1
44.7
40.0
51.7
26.7
20.0
28.7
28.7
$SS
DD (ours)
62.8
65.2
45.3
44.2
54.4
32.3
26.9
24.6
30.7
34.2
30.9
244
S
Table 5: Comparison with baseline and state-of-the-art methods. All models are trained with the same hyperparameters, RoPE with fo 100k, and for 103B tokens. DM denotes training with document masking. DD uses the "Grow-P2" curriculum with 8 cycles. Dataset preparation cost is symbolic to compare methods and does not reflect the wall-clock time.
4 Related works
Recent works have raised concerns regarding cross-document attention. For example, Llama- 3 139, ICLM (55). and [17]. which dis- cussed in Section 3.6 Similarly, [26] dis- cuss challenges with the baseline concat-and- chunk approach and propose an approximate bin-packing algorithm.
Method
Doc Masking
Average Cuntext
Does in Sequence
Corz.
MIQA- (0%)
Mult-rankmm
16.0
8L+DM [77
Mullspacking
169
DI-Uniform
1544
Single
X
19.6
Baseline
4006
Mult-random
X
20.5
ICLM 55
4096
Mult- semantic
X
22.0
DD-Gros-P2
N/A
1544
Single
24.6
Related to our study on sequence length bias, 62] shows the importance of train-vs-test time
Table 6: Summary of long-contexs performance for different methods from Table 2 and Table 5
distribution shift from a sequence length per-
spective on string editing task. [6] 66 [6]
highlight the challenge of generalizing to lengths beyond what the model has seen during training
and discuss the importance of positional encoding. Several works [4] 67 23 64
10
47
48
53
34
address enabling LLM inference with long context (see [45] for an overview). These approaches adapt to longer lengths. Growl BrrwhLnntth [24] proposes accelerating LLM pretraining by progressively growing context length us-
orthogonal to contribution and be applied post-pretraining
ing the baseline sequence formation method, but does not show results LLMss Similarly, increasing sequence length has been shown in BERT model training
to improve compute efficiency.
The idea of dynamic batching has been explored in other domains. In vision, methods like Na [16 38 use images with variable resolutions (a similar concept to context length for LLMs). In seq-to-seq tasks (c.g., automatic speech recognition, text-to-speech, and neural machine translation). the inputs have different lengths. An efficient approach 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 Different from these works, in dataset decomposition. 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 target mixture and curriculum.
5 Conclusion and limitations
In this paper, we explore the shortcomings of popular LLM pretraining approach, concat-and- chunk, and introduce dataset dccomposition. a method to decompose a dataset 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 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-chun shortcomings in unified experimental setup and show gains in data preparation cost, training time, and final model accuracy.
The training speed gains compared to the baseline are significant only when the target sequence length is long enough. Otherwise, the attention cost not a dominant fraction of training, and hence significant training speedup is expected.
Limitations.
10

---

Acknowledgements
We thank Tatiana Likhomanenko, Jason Ramapuram, Alexander Toshev, Barry Theobald, and Fartash Faghri from Apple for their valuable feedback and suggestions.
References
[1] Common crawl. Rttps://ccomonccawl.oorggg
[2] Lim foundry v0.7.0.
[3] Jeopardy, version by LLM Foundry).
https://uugggineeecc/////ttt//// [Used custom curated
[4] Big-bench qa wikidata. Sag.berch/benccmmarttaaaaa//////
[Used through LLM Foundry).
[5] C. An, S. Gong, M. Zhong, M. Li, J. Zhang, L. Kong, and X. Qiu. L-eval: Instituting standardized evaluation for long context language models, 2023.
[6] C. Anil, Y. Wu, A. Andreassen, Lewkowyez, Misra, V. Ramasesh, Slone, G. Gur-Ari, E. Dyer, and B. Neyshabur. Exploring length generalization in large language models. Advances Neural Information Processing Systems, 09385566385555 2022.
[7] Y. Bengio, J. Louradour, R. Collobert, and J. Weston. Curriculum learning. In Proceedings of the 26th annual international conference machine learning, pages 41-48, 2009.
[8] Y. Bisk, R. Zellers, Gao, Y. Choi. et Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
[9] S. Black, S. Biderman, E. Hallahan, Q. Anthony, L. Gao, L. Golding, H. He, C. Leahy, K. McDonell, J. Phang, et al. npt-neox-20b: An open-source autoregressive language model. arXiv preprint arXiv:220..06744
2022.
[10] S. Chen, S. Wong. L. Chen, and Y. Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306. 2023.
[11] Y.-A. Chung, H.-Y. Lee, and Glass. Supervised and unsupervised transfer learning for question answering. In M. Walker, H. Ji, and A. Stent, editors, Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long Papers). Association for Computational Linguistics, 2018.
[12] C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski. M. Collins, and K. Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint 1901v:1905.10044, 2019.
[13] Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018.
[14] T. Computer. Redpajama: An open source recipe to reproduce llama training dataset, 2023. URL htlltttts//gghhuii
http://tbb.ooooeoteeeeeeerooeeee ReddPa -aaa-Datal
[15] T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08661, 2023.
[16] M. Dehghani, B. Mustafa, J. Djolonga, J. Heek, M. Minderer, M. Caron, A. Steiner, J. Puigcerver, R. Geirhos, I. M. Alabdulmohsin, et al. Patch pack: Navit, vision transformer for any aspect ratio and resolution. Advances in Neural Information Processing Systems, 36, 2024.
[17] H. Ding, Z.. Wang, G. Paolini, V. Kumar, A. Deoras, D. Roth, and S. Soatto. Fewer truncations improve language modeling. arXiv preprint arXiv:2404.10830, 2024.

---

[18] J. L. Elman. Learning and development in neural networks: The importance of starting small. Cognition, 48(1):71-99 1993.
[19] L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.000 001000027 2020.
[20] Z. Ge, L. Kaushik, M. Omote, and S. Kumar. Speed up training with variable length inputs by efficient batching strategies. In Interspeech, pages 156-160, 2021.
[21] P. Gonzalez, T. S. Alström, and T. May. On batching variable size inputs for training end-to- end speech enhancement systems. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 1-5. IEFE, 2023.
[22] S. Gururangan, M. Wortsman, S. Y. Gadre, A. Dave, M. Kilian, W. Shi, J. Mercat, G. Smyr- nis, G. Ilharco, M. Jordan, R. Heckel, A. Dimakis, A. Farhadi, V. Shankar, and L. Schmidt. OpenIM: a minimal but performative language modeling (Im) repository, 2023. URL https://gg //g thub. www..ccooo/..... m Foundations//peaaeeeetoo
1m/ GitHub repository.
[23] C. Han, Q. Wang, W. Xiong, Y. Chen, H. Ji, and S. Wang. Lm-infinite: Simple on-the-fly fly length generalization for large language models. arviv preprint arXiv:2308. 2023.
[24] H. Jin, X. Han, J. Yang, Z. Jiang, C.-Y. Chang, and X. Hu. Growlength: Accelerating Ilms pretraining by progressively growing training length. arXiv preprint arXiv:2310.00566. 2023.
[25] A. Kazemnejad, I. Padhi, K. Natesan Ramamurthy, P. Das, and S. Reddy. The impact of positional encoding on length generalization transformers. Advances in Neural Information Processing Systems, 36, 2024.
[26] M. M. Krell, M. Kosec, S. P. Perez, and A. Fitzgibbon. Efficient sequence packing without cros-coonamination:: Accelerating large language models without impacting performance. arXiv preprini urXiv:2107.02002 2021.
[27] T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein,
I. Polosukhin, J. Devlin, K. Lee, et al. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453 466, 2019.
[28] K. Lee, M.-W. Chang, and K. Toutanova. Latent retrieval for weakly supervised open domain question answering. arXiv preprint arXiv:1906 00300, 2019.
[29] B. Lefaudeux, F. Massa, D. Liskovich, W. Xiong, V. Caggiano, S. Naren, M. Xu, Hu, M. Tintore, S. Zhang, P. Labatut, D. Haziza, L. Wehrstedt, J. Reizenstein, and G. Sizov. xformers: A modular and hackable transformer modelling library. matbbookkesearrc/xxEEooee
2022.
[30] H. Levesque, E. Davis, and L. Morgenstern. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning, 2012.
[31] C. Li, M. Zhang, and Y. He. The stability-efficienny dilemma: Investigating sequence length warmup for training gpt models. Advances in Neural Information Processing Systems, 35: 26736-26750, 2022.
[32] J. Li, A. Fang, G. Smymiss M. Ivgi, M. Jordan, S. Gadre, H. Bansal, E. Guha, S. Keh, K. Arora, et al. Datacomp-Im: In search of the next generation of training sets for language models. arXiv preprint arXiv:2406. 1794, 2024.
[33] X. Li, Z. Wang, and C. Xie. An inverse scaling law for clip training. Advances in Neural Information Processing Systems, 36, 2024.
[34] H. Liu, W. Yan, M. Zaharia, and P. Abbeel. World model on million-length video and language with blockwise ringattention. arXiv preprint arXiv:2402.08266 2024.
[35] N. F. Liu, K. Lin, J. Hewitt A. Paranjape, M. Bevilacqua, F. Petroni, and P. Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024.
12

---

[36] X. Liu, H. Yan, S. Zhang, C. An, X. Qiu, and D. Lin. Scaling laws of rope-based extrapolation. arXiv preprint arXiv:2310.05209, 2023.
[37] Y. Liu, M. Ou, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zeulemoyer, and V. Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692, 2019.
[38] Mehta, Abdolhosscini, and M. Rastegari. Cvnets: High performance library for computer vision. In Proceedings of the 30th ACM International Conference on Multimedia, pages 7327-7330, 2022.
[39] Meta. Introducing meta llama 3: The most capable openly available Ilm to date, 2024. URL https://a: meta. cww..oo////o...... www.boooo///oooooooooooooooooooo www.g//tttaaaaaaaaaaaaaaaaaaaaaa
[40] T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:/809.02789, 2018.
[41] A. Mohtashami and M. Jaggi. Random-access infinite context length for transformers. Advances in Neural Information Processing Systems, 36, 2024.
[42] K. Nagatsuka, C. Broni-Bediako, and M. Atsumi. Pre-training a bert with curriculum learning by increasing block-size of input text. In Proceedings of the International Conference Recent Advances in Natural Language Processing (RANLP 2. pages 989-996, 2021.
[43] M. Ott, S. Edunov, A. Baevski, A. Fan, S. Gross, N. Ng, D. Grangier, and M. Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-IILT 2019: Demonstrations, 2019.
[44] D. Paperno, G. Kruszewski, Lazaridou, Q. Pham, R. Bernardi, Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv: 160.00366 2016.
[45] S. Pawar, S. Tonmoy, S. Zaman, V. Jain, A. Chadha, and A. Das. The what, why, and how of context length extension techniques in large language models-a detailed survey. arXiv preprint arXiv:2401.0rm 007872, 2024.
[46] G. Penedo, Q. Malartic, D. Hesslow, R. Cojocaru, A. Cappelli, H. Alobeidli, B. Pannier, E. Almazrouei, and J. Launay. The RefinedWeb dataset for Falcon LLM: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01166, 2023, URL https: /wwiivvvvvvvv/////vvv///vv///vvv org/ abs / 2306 01116
[47] B. Peng and J. Quesnelle. Ntk-aware scaled rope allows Ilama models to have ex- tended (8k+) context size without any fine-tuning and minimal perplexity degradation. 2023. URL https: www. reddit. com/ com///LocalLLaM. http//ttttttt///////// otkaware_scaled rope allows 11ama aallaaammddes_______aaa_aa__
models to have
[48] B. Peng, Quesnelle, H. Fan, and Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint 2300023309.0077, 2023.
[49] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748-8763. PMLR, 2021.
[50] P. Rajpurkar, R. Jia, and P. Liang. Know what you don't 't know: Unanswerable questions for squad. arXiv preprint arXiv: /o0o::000000222 2018.
[51] S. Reddy, D. Chen, and C. D. Manning. Coqa: A conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7:249-266, 2019.
[52] M. Roemmele, C. A. Bejan, and A. S. Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series. 2011.
[53] A. Ruoss, G. Delétang, T. Genewein, J. Grau-Moya, R. Csordás, M. Bennani, S. Legg, and J. Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint 2305..33005.6443 2023.
13

---

[54] K. Sakaguchi, R. L. Bras, C. Bhagavatull, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
[55] W. Shi, S. Min, M. Lomeli, C. Zhou, M. Li, V. Lin, N. A. Smith, L. Zettlemoyer, S. Yih, and M. Lewis. In-context pretraining: Language modeling beyond document boundaries. arXiv preprint arXiv:2310.10638, 2023.
[56] L. N. Smith, Cyclical learning rates for training neural networks. In 2017 IEEE winter conference applications of computer vision (WACV), pages 464 464-472. IEEE, 2017.
[57] L. Soldaini, R. Kinney, A. Bhagia, D. Schwenk, D. Atkinson, R. Authur, B. Bogin, K. Chandu, J. Dumas, Y. Elazar, V. Hofmann, A. H. Jha, S. Kumar, L. Lucy, X. Lyu, N. Lambert, I. Magnus- son, Morrison, N. Muennighoff, A. Naik, C. Nam, M. E. Peters, A. Ravichander, K. Richard- son, Z. Shen, E. Strubell, N. Subramani, O. Tafjord, P. Walsh, L. Zettlemoyer, N. A. Smith, H. Hajishirzi, I. Beltagy, D. Groeneveld, Dodge, and Lo. Dolma: an Open Corpus of Three Trillion Tokens for Language Model Pretraining Research. arXiv preprint, 2024.
[58] J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
[59] H. Touvron, Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv: 1, 2023.
[60] H. Touvron, L.. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, Bhargava, S. Bhosale, al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.092288, 2023.
[61] B.-H. Tseng, S.-S. Shen, H.-Y. Lee, and L.-S. Lee. Towards machine comprehension of spoken content: Initial toefl listening comprehension test by machine, 2016.
[62] D. Variš and O. Bojar. Sequence length is domain: Length-based overfitting in transformer models. arXiv preprint arXiv:2109.07276, 2021.
[63] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, E. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[64] W. Xiong, J. Liu, Molybog, H. Zhang, P. Bhargava, R. Hou, L. Martin, R. Rungta, K. A. Sankararaman, B. Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309. 16039, 2023.
[65] R. Zellers, Holtzman, Y. Bisk, A. Farhadi, and Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv: 9055007300, 2019.
[66] Y. Zhou, U. Alon, X. Chen, X. Wang, R. Agarwal, and D. Zhou. Transformers achieve length generalization but not robustly. arXiv preprint 2XXv:2402.003771 2024.
[67] D. Zhu, N. Yang, L. Wang, Y. Song, W. Wu, F. Wei, and S. Li. Pose: Efficient context window extension of Ilms via positional skip-wise training. arXiv preprint ar Xiv::2309. 2023.
14

---

A Broader impacts
This work enables faster training of LLMs, which are among the most compute-intensie applications in the field. A positive socictal/environmmnnta impact of this work is training LLMs with 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 benefit when cross-document attention is not allowed during LLM pretraining.
B
Implementation details
B.1 Training details
Software and hardware details All experiments in this paper conducted using the OpenI.M repository, which is based on PyTorch. We use Fully Sharded Data Parallelism (FSDP) with Bfloat 16 mixed precision for all experiments. We use the Xformers [29] implementation for attention. For hardware, we use one or more nodes of 8x NVIDIA H100 GPUs (Hopper architecture), each with 80GB memory, and 192 CPU cores with 2000GB RAM. Nodes are connected through Elastic Fabric Adapter (EFA) for efficient noder-node communication hosted by AWS.
Model architecture details Table |
We provide details of all architectures used the paper Table to
Model Name
OpenI MM-60M
Model Name Hidden dimension
OpenI MM-nLMM410MM
Hidden dimension
768
1024
Number of Layers
12
Number of Layers
24
Number of Heads
12
Number of Heads
16
Number of Parameters
162,435,840
Number of Parameters
41 411,665,408
Table 7: OpenI.M-160M.
Table 8: OpenI.M-410M.
Model Name Hidden dimension Number of Layers Number of Heads
OpenLM-IB
Model Name Hidden dimension
OpenLMM3B
2048
2560
24
Number Layers Number of Heads
16
32
Number of Parameters
1,439,893,504
Number of Parameters 2,796,096,000
Table 9: OpenLM-IB.
Table 10: OpenL
Model Name Hidden dimension Number of Layers
OpenLM-7B
4096
32
Number of Heads
32
Number of Parameters 6,889,672,704
Table 11: OpenLM-7B.
Baseline hyper parameters We list our baseline hyperparameters in Table 12 and iterate over changes for each section next. Note that did not explicitly optimize hyperparameters for any of the experiments, and we always use the same hyperparameters when using either the baseline method ours.
15

---

Optimizer
Adam
AdamW-31
0.9
AdamW W-32
0.95
learning-rate schedule
cosine+warmup
Maximum learning rate
3 x 10-3
cooldown learning rate
x 10
Warm-up steps
5000
Grad Norm Clipping
1.0
Global batchsize (num tokens per step)
2
Weight Decay
0.1
Z-Loss Coefficient
10-4
Table 12: Baseline hyper- paeeer-arametess...
Implementation details of Section experiments Experiments this section are done using the same hyperparameters as in Table for total of 234 tokens the OpenLM-1B model. We trained each model twice with different random seeds and report the averaged results. For models in RoPE with 10, 000. In Table show that results and conclusions in this section are not sensitive to hyperparameters, including the RoPE base frequency fo-
this section,
Implementation details of Section 333 and Section 3.4 experiments Experiments in this section are done with OpenL.M1B B model, trained for total of 96 x 10 & 103B tokens. Hyper-parameters the same as Table [2] except used 20000 warmup steps for all models presented in this section. We use RoPE with /b = 10, 000 for all models in Section 3.3 and fo 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 fa 100, 000, and the baseline setup as in Table T2 except for the following changes for different dataset sizes:
total tokens = 231
warmup steps = 5, 000
total tokens -335 warmup steps 5. 000
total tokens = 236
warmup steps = 10, 000
total tokens -237, warmup steps 20, 000
total tokens = 238, warmup steps = 40, 000
Model scaling:
Experiments in this section trained with the OpenLM-IB, OpenLM-3B, and
OpenLM-7B models, 237 x 137B total seen tokens, RoPE with fa = 100, 000, and the baseline setup in Table [2 except for the following changes for different model sizes:
OpenL IB, warmup steps = 20, 000, max-Ir = 3 x 10 batchsize = 219, with 32
H100 GPUs
Op3nL.M-3B, warmup steps = 20, 000, max-Ir = 2 x 10 batchsize 220 with 64
H100 GPUs
-OpenLM-7BB warmup steps = 20, 000, max-Ir x 10 batchsize 222, with 128
H100 GPUs
Experiments in this section are trained with the OpenLM-160M, OpenLM- 410M, and OpenLM-IB models, 237 137B total seen tokens, RoPE with fo = 100, 000. and the baseline setup as in Table [2] for the following changes for different model sizes:
Alternative dataset:
OpenLM-1600 160M, warmup steps 20, 000, max-Ir 5 10 weight-decay 0.033, with H100 GPUs
-OpenLM-410M. warmup steps 20, 000, max-Ir 4 10-3, weight-decay 0.066, with
H100 GPUs
-ppenLM-1B3 IB, warmup steps 20, 000, max-Ir 3 10 weight-decay 0.1, with 32 H100 GPUs
For the DD experiments used "Grow-P2" length curriculum which is visualized in Fig.
16

---

Sequence length
0.0n5
212
0.001
Learning rate
2
0.003
210
0.002
29
0.001
0.000
30000
40000
50000 60000
Optimization step
Figure 6: Comparison of length-bar based curriculum schedule with learning rate schedule. Sequence length varies between 256 and 8192 based the Grow-P2 currieulum with 8 cycles. Note that the choice of bucket (and hence the sequence length) is random, with sampling probabilities determined by the curriculum. In the figure, we show the length of the sampled sequence at every optimization steps. For the learning rate, we use a cosine learning rate with warm-up for 4k steps. The job corresponds to training for total of 236 tokens, with 220 tokens seen per optimization step.
All experiments in this section are done with the OpenLM- IB model, trained for a total of 96 10 34 se 103B tokens. Hyperparameters are the same as in Table 12 except we used 20,000 warmup steps for all models presented in this section. We use RoPE with f = 100, 000.
Implementation details of Section 3.6
B.2 Length based sampling and curriculum algorithm
We present the details of our length-based sampling and curriculum Algorithm
Algorithm Length based sampling and curriculum
Require:
Dj: list of buckets such that D; includes sequences with length 2
ny: total number tokens to be picked from each bucket (see Table T
0g: sampling odd for each bucket (see Table 2
c: number of cycles
b: number of tokens per optimization step
S. < random subset of D, with n tokes for j E [1. ......... c do
non-overlapping subsets of D;
loop over cycles
while at least one sig is non-empty do odds + [o; if 8t, is not empty else for 3,2, probs + odds/ hadd//dds..mm
randomly sample index probability probs[i]
sample b/2 sequences from 8ij w/o replacement for training
end while
end for
B.3 Evaluation details
Multi Document Question Answering (MDQA) 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 's response by using only the text before the first occurrence 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, the evaluation script to look for the first non 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
updated
35
TOEFL
We follow the setup described [5] As described Section the dataset contains multiple-choice QA pairs for the 15 longest lectures in [61 [1. To obtain response from the model, we follow MMI.U-style prompting, where the choices are appended to the original prompt individually and the mean log-probability computed for each choice. The choice corresponding to
17

---

ICC
31L
Y13
144
245
143
544
2/1
cimnle dl SECUSSASS pic batok
zamber di SEQUISAASI pic Tatee
30
144
Fe
240
425
H022
3-3F
TC1
%
1890
1091
aequenee kngth
sequence kngth
sequenee kngth
(a) Average time for OpenLM-1B I-
(b) Average time for M-3BLM-3B
(c) Average time for OpenLMM7B
mumber of oequrtins pet batch
per batth
number of oocurroons per batch
4
EUR
2:26
1:20
0:32
1:31
27
1/C
CUI
0.25
EXIT
1.30
0:45
0.35
E-40
0.25
EUVO
1.:4
0.55
0:33
C.4
0.37
EVD
92
81
EXG
22
0:5
1:54
92-
12-034
0.37
9:38
8:50
munber
EU5
1:17
+4-
ExD
290
sequence kngth
Sequenoo lengt
sequence knuth
(d) STD for OpenLM- IB
(e) STD for OpenLM-3B
(I) STD for OpenLM-77
Figure 7: Top row: Average time (ms) for each node to train one batch 8 8xH100 machine using FSDP. Bottom row: measured standard deviation for each setup.
the argmax of mean log-probabilne 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 (5). The dataset contains long documents with each document containing multiple-choice QA pairs. Sometimes the context for 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-probat probability computed for each choice. The choice corresponding the a argmax of mean log- log-probability is then chosen as the model's s response. After we obtain the model output, the rest of the evaluation follows the official implementation (5].
C
Additional results
C.1 Additional results for training efficiency
We enumerate model sizes (OpenLM-1BB O-enLM-3B, OpenLM-7B), the number of sequences in a batch (from 1 to 256), and sequence lengths (2 to 214) 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 fixed b (number of tokens seen per optimization step).
C.2 Additional results for sequence length bias experiments
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 2048 sequence length. Results are shown in Table I3
18

---

demonstrating that the gap cannot be simply filled by choosing different hyperparameter and is fundamental the choice of pretraining sequence length.
Maximum Learning Rate
RoPE
Regular Average
3 x 10
10,000
47.0
3 x 10-3
100,000
47.1
J
10,000
45.9
10 2
10,000
46.5
Table 13: Sensitivity to hyperparameters for Sec- tion experiments. All models are trained twice with different random seeds, and averaged results reported.
C.3
Additional results for scaling experiments
In this section, we show additional results for the experiments presented in Section 3.5 shows results for dataset scaling, Table 13 for model scaling, and Table 16 for experiments on an alternative dataset.
Table [4]
Seen tokens
Method
Regular uverage
MDQA
2
Baseline-8k
7.8
DD
47.0
16.0
Model size
Method
Regular average
MDQA average
2
Baseline-Ke 8k
DD
47.6
23.3
15.4
IB
Baseline-Sk
51.9
24.2
23,1
2
Baseline-8k 8k
DD
50.2
22.3
19.9
3B
Baseline-&k
57.5
17.8
250
Baseline-8k
DD
54.9
$1.9
23.2
25.9
7B
Baseline &k
59.8
31,7
2
Baseline-Kk DD
25.8
56.0
29.4
Table 15: Model scaling for total of 137B tokens.
Table 14: Dataset scaling for OpenLM- IB.
Method
PIQA
COPA
OBQA
LanOAl
SQuAAD
CuQA
WikaQA
XDQQA
160M
Baseline
UD
66.5
66.4
65
61
30.2
29.3
43.6
405
37.7
37.2
66.3
63.6
52.
51.9
14.3
12.9
556
90.7
19.1
182
233
43
51.3
49.5
25.9
24
762
54
16.2
12.8
9.3
99
K2
7.1
DD
71.3
70
34
55.8
74.7
96.1
69.4
10.
53.2
28.
1811
14
Bactiie 74.9
24
43.4
63
622
802
53.4
418
66.
35.3
20.2
65.7
18.4
56.7
31.3
248
208
DD
76.7
64.7
36.3
56.7
204
Table 16: Small model performance trained on an improved refined-web pipeline applied to Common Crawl. All models are trained for total of 237 tokens.
D Comparison to best-fit sequence packing
Some recent works have employed a bin packing-based strategy [17 which aims reduce document cross-attention by minimizing unnecessary document truncation. To achieve this, they implement 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 fit decreasing algorithm and trained 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 = {c1, ....... cK) such that U c; D. where each chunk is at most context size (e.g., 2048) in length. The output the algorithm a list of bins B {b), b by) such that & E by. As 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 We track the remaining capacities for each
19

---

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 chunks in bin. 3 shows that while best-fit packing results in higher average context length compared to the baseline concat-and- 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 best-fit packed sequences also means that some context wasted during each optimization step.
E
Training stability with VSL and curriculum
presents the stability-effcienny 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 poor generalization. They observe a correlation between training instability and long sequences, especially at the early stages of training, suggesting that training on can be main source of training instability.
31
long sequences when the model not
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 OpenL.M1BB 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 shown in Fig. g demonstrating the stability training with DD comparison to the baseline. This suggests that proposed method can also be beneficial for large-scale pretraining with large batches and high learning rates in terms of efficiency.
Reseline-dt
DD-Grow-PPie
Figure 8: We compare the training loss when training with Baseline-8k DD with the "Grow- P100" curriculum. Both models are trained with identical hyperparameters. a high learning rate (= 10-2) 2), and no gradient clipping. is evident that DD results in greater stability.
Average sequence length vs average context length
We compute the mean of length (Fig. 3aj and context (Fig, 3e) distributions as follows. Assume a list of sequences with lengths 1, le..... IN. which are, for example, the chunk lengths in the concat-and-chume approach the sequence lengths in different buckets the dataset decomposition approach. We define the average sequence length as follows:
Average sequence length
Σ
4
In rego-segeeseive training on a sequence with length 1, we apply losses for next-token prediction on each token in parallel. Hence, for sequence with length we see contexts with lengths equal to
20

---

1.
We define the average context length, which is different from the average sequence
length. as follows:
ΣΣ
Σ (((--1)
Average context length-
2
(2)
we report the average sequence length and average context length for original documents, concat-and-chunk, and dataset decomposition with different mixtures.
In Fig. 3a Fig. 3c
and Table 1
21

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's s contributions and scope? Answer: [Yes]
Question:
Justification: Method proposed in Section 2.1 and experiments provided in Section 3 support all claims in the abstract.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims
made in the paper.
The abstract and/or introduction should clearly state the claims made, including the
contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? [Yes]
Answer:
Justification:
We provide limitations in Section 5
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what
implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: Our work does have any theorem/proof.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
22

---

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 provide 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.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that 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 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 specific model and empirical evaluation, it may be necessary to either make 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 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 NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., 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 welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our results are based on OpenLM repository (attps:/ com/ m1 woundat innttatiios//// peoos//ppen.. 1m) 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.
Please see the NeurlPS code and data submission guidelines (https://nipss cc/
pub 1 ic/ / gui des SCodeSuomiiieee ssi i onPol ccy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not
23

---

including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https)
//nips. cc/ publ c/ gu des CodeeuddeeSasooee Summi i onPol icy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why. At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
reproducible, they
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all implementation details in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Training LLM at scale is computationally expensive. We repeat all experiments at 17B total tokens scale twice (with different random seeds) and report mean and variance in Section 3.2 For efficiency benchmarks we repeat 5 times, report standard deviation in
Fig. 7
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence 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, 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 eтror 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 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 стror rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
24

---

Answer: [Yes]
Justification: We provide software/hardware details in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:p/:/w..rrrppp...ccoo.l.ll/a
Answer:
[Yes]
Justification: All experiments fully respect NeurIPS code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid-
eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss broader impacts of this work Appendix A
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal
impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), fairness considerations (e.g., deploymeen of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA]
Justification: We do not believe the models/code related to this work poses any risk that requires safeguarding.
Guidelines:
The answer NA means that the paper poses such risks.
25

---

Released models that have high risk for misuse or dual-use 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 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 dataset.
The authors should state which version of the asset used and, possible, include 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, paperr www.rrswithocooo com/ dat aset s has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is not available online, the authors encouraged to reach out to the asset's '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 the dataset/code/model as part 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.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This work does not include any crowdsourcing or research with Human subjects.
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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
26

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Not applicable to this work. Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
27


Reviews:
[{"review_id": "Ub7AW9ybXJ", "reviewer": "Reviewer_bcif", "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.", "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?\n\nIn 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": "Yes", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1: The paper is well organized and easy to follow. \n\n2: The motivation is clear and the proposed method looks simple yet effective. \n\n3: This paper conducted massive experiments and provided many valuable empirical results, which can support both the claim of this paper as well as many long existing guesses in the community.", "weaknesses": "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.\n\n2. 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.\n\n[1] World Model on Million-Length Video and Language with RingAttention\n[2] GrowLength: Accelerating LLMs Pretraining by Progressively Growing Training Length", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer, bcif, for their time and feedback. We are glad the reviewer finds the paper well-organized and motivated, the method effective, and the experiments valuable to the community. In the following, we address the concerns and questions raised by the reviewer and kindly request to be informed if further clarification is needed.\n\n---\n> It's also valuable for models to develop the ability to discern and disregard irrelevant information\n\nWe agree with the reviewer that, in principle, the model may benefit from attending to irrelevant documents during pre-training to develop a discerning capability. We mentioned this in the paper (line 291). Our results also confirm this. In Table 5, when comparing baseline-8k to baseline-8k with document masking (which stops cross-document attention), adding document masking improves the regular metric (51.5% → 52.4%). However, it results in weaker discerning capability, as seen in multi-document QA results with 30 documents in the context (performance drops from 20.5% → 16%). We will clarify this point further in the revision.\n\nWe would also like to point out that our proposed method surpasses all baselines in discerning capability, as seen from its superior performance on multi-document QA. This capability emerges when a length-based curriculum is deployed (compare the multi-document QA performance with a uniform curriculum and, for example, the Grow-P2 curriculum in Table 2). Please **see the summary in Table 1 of the rebuttal PDF**.\n\n---\n> Length-based curriculum, while insightful, isn't entirely novel\n\nWe thank the reviewer for pointing out two relevant works. Both are **concurrent** with ours **and not published** in a peer-reviewed venue at the time of this rebuttal. We also appreciate the reference to the length-based curriculum in BERT [3]. We will include them in the revised paper. These works highlight the computational benefits of length-based batching; however, they differ significantly from our work, as explained below:\n\n* We are the first to show length-based curriculum for autoregressive LLM pre-training. Given the huge cost of LLM pre-training, the savings from our proposed method are significant (more than 6x speed up to reach the best accuracy of baseline as shown in Figure 1 of the rebuttal PDF). The related work [1] is a continual learning setup (starting from pre-trained Llama2) only for context-length extension using book data. GrowLength [2] does not show any results on large language models. BERT [3] is only for masked-language modeling.\n* Unlike [1-3], we introduce binary decomposition, a novel method to preserve document length, form buckets with fixed sequence lengths without using pad tokens, and avoid forming multi-document sequences (thus achieving no cross-document attention without attention masking).\n* Unlike [1-3], we analyze different forms of curriculum (not just a simple multi-stage training from short to long sequences) in Table 2, and show that mixing (i.e., a mixture of long and short sequences with a changing mixture) is important for the best results. Our analysis further shows that our introduced cyclic schedule is important for best results.\n* Unlike [1-3], we systematically show the effect of pre-training sequence length on model performance for different tasks, including those requiring long context (Sections 3.2 and 3.3).\n\n\n[1] World Model on Million-Length Video and Language with RingAttention, 2024\n\n[2] GrowLength: Accelerating LLMs Pretraining by Progressively Growing Training Length, 2024\n\n[3] Pre-training a BERT with curriculum learning by increasing block-size of input text, RANLP 2021\n\n---\n> The outcomes when models are trained with sufficient data\n\nNote that the benefits of using the proposed method are twofold:\n\n* **Computational benefit**: Using variable sequence length training (Section 2.2), the training cost per seen token is reduced, i.e., we reach a certain performance faster compared to the baseline.\n* **Performance benefit**: Using DD with a curriculum for the same total number of seen tokens, compared to the baseline, DD results in a more accurate model (for both regular and long context evaluations).\n\nTo further demonstrate the above facts and the scalability of our results, we trained a 410M model with our proposed method and the baseline up to 1.1 trillion tokens. This is 128 times more tokens than the “optimal” number recommended based on the number of parameters of the model in Chinchilla paper [4]. Furthermore, 1.1 trillion tokens exceed the number used in recent state-of-the-art open LLMs (e.g., smolLM [5] uses 600 billion tokens to train their 360M model). We show **+2.4 accuracy improvement compared to baseline even at 1.1 trillion tokens**, where the baseline show a plateau in accuracy (indicative of sufficient data). Further we show **more than 4x data efficiency** and **more than 6x speed-up** to reach to baseline's best accuracy. Please **see Figure 1 of the rebuttal PDF**.\n\n[4] Training compute-optimal large language models, NeurIPS 2022\n\n[5] SmolLM - blazingly fast and remarkably powerful, 2024\n\n---\n> The originality of the proposed strategy\n\nWe would like to emphasize again that this is the first work demonstrating the computational and performance benefits of length-based curricula (and beyond simple multi-stage training) for LLM pre-training. Furthermore, the proposed binary dataset decomposition method is entirely new. As mentioned by the reviewer, we provide extensive experiments and insights on the importance of length in the mixture, training stability, and generalization to different model sizes and datasets. We believe the combination of our proposed binary decomposition method, extensive empirical results, and open-source code and models will be a useful contribution to the community.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1llqNeP3FR", "reviewer": "Reviewer_z2w6", "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.", "questions": "The same as the weaknesses", "limitations": "No potential negative societal impact.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Comprehensive Experiments: The authors conduct extensive experiments across various datasets and LLMs, with observations and analyses that are noteworthy.\n\n2. Simplicity and Effectiveness: The proposed method is straightforward and efficacious.", "weaknesses": "The paper's writing requires improvement, as exemplified by the following:\n\nLine 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.\n\nLine 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.\nTable 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.\n\nAdditionally, 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.\n\nIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer z2w6 for their time and feedback. We are happy that the reviewer finds our experiments comprehensive, analyses noteworthy, and the method simple and effective. We address the concerns and questions raised by the reviewer below and kindly request to be informed if further clarification is needed.\n\n---\n> Paper's writing requires improvement\n\nWe thank the reviewer for the editorial points. We have applied all of them in the revised manuscript including additional discussion and apologize for any potential inconvenience during the review.\n\n---\n> It remains unclear what is meant by penalty\n\nWe apologize for the confusion. By \"penalty,\" we meant \"computational cost.\" We fix this in the revised manuscript.\n\n---\n> Significant advantage over other methods in long-context scenarios\n\nWe provide a more detailed discussion here, and revise the paper accordingly.\n\nAs pointed out by the reviewer, 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 of the paper, where models trained on longer sequences perform better on multi-document QA. For the context length from the same document (i.e., the number of tokens from the same document a token can attend to), our proposed method has an average context length of 1,344 for the RefinedWeb dataset (as defined in equation 2 in Appendix F), compared to 930 for the baseline (see Figure 3c of the paper) and 1,064 when bin-packing [1] is applied. This explains why the dataset decomposition mixture, even without any length-based curriculum (the first row in Table 2 of the paper), outperforms Baseline-8k-DM and Pack-8k+DM (second and third rows in Table 5 of the paper). Here, DM refers to applying document masking to avoid cross-document attention.\n\nOne can increase context length by concatenating different documents and putting them in the context, as in the Baseline-8k result. However, simply increasing the context length by filling it with multiple documents does not necessarily lead to long-range attention during training (and hence improved long-context capability of the model). Nevertheless, as discussed in the response to reviewer bcif, a multi-document context encourages the model to learn to discern and disregard irrelevant information. Comparing Baseline-8k and Baseline-8k-DM multi-document QA results in Table 5 of the paper shows such benefit. Baseline-8k multi-document QA performance is even slightly better than our proposed dataset decomposition mixture when used without length-based curriculum (first row of Table 2 of the paper).\n\nIn-context pre-training LMs (ICLM [2]) proposes to put semantically relevant documents into the context. We observe that ICLM results in slightly better multi-document QA performance when 30 documents are in the context compared with Baseline-8k (22.0% vs. 20.5%). However, we do not observe such gains in shorter multi-document QAs (i.e., with fewer distractor documents in the context) and regular evaluations.\n\nFinally, in Table 2 of the paper, we show the importance of length-based curriculum. Note that the data mixture (and hence the average context length) is the same for all rows in Table 2 of the paper, differing only by the curriculum (i.e., the order in which different examples are seen during training). We show that using 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. For instance, multi-document QA with 30 documents improves from 19.6% with no curriculum (first row in Table 2 of the paper) to 24.6%. It is worth noting that the effect of length-based curriculum on regular metrics is less significant, with the average metric improving from 53.8% with no curriculum to 54.4% for Grow-P2 with 8 cycles curriculum.\n\nPlease **see Table 1 of the rebuttal PDF** for a  summary of all of the above contributing factors.\n\n[1] Fewer truncations improve language modeling, ICML 2024\n\n[2] In-context pretraining: Language modeling beyond document boundaries, ICLR 2024", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AFEdPuU6L2", "reviewer": "Reviewer_H2jF", "summary": "The paper explores dataset decomposition for LLM pre-training. \n\nThe 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. \n\nThe 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.", "questions": "Question:\n\n1. Can cross-document attention be restricted with an attention mask? Can that be explored? \n2. 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. \n3. Can you elaborate more on the pacing strategy (see more on my confusion below in the suggestions)?\n\n\nSuggestions:\n\n* 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": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The high-level ideas are reasonably motivated. \n* Reasonably extensive experimental analyses of the methods are provided along with baseline dataset structuring methods for pre-training.\n* Results are generally promising. Generally, it results in better task performance while taking less training time.", "weaknesses": "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].\n\n[1] CYCLICAL CURRICULUM LEARNING - Kesgin et al. ArXiv 2022\n\nAlso, if I understand correctly, the training speed gain may be less significant with flash-attention-based LLMs and alternative models like (Mamba, Linear Transformers).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer H2jF for their time and feedback. We are glad that the reviewer finds our work to be reasonably motivated, and the results to be extensive and promising. We respond to the reviewer’s concerns and questions below, address all of them in the revised paper, and kindly request that you let us know if further clarification is needed.\n\n---\n> The paper seems to be mainly an exploration of hyperparameter tuning\n\nWe 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.\n\n---\n> The main proposed techniques seem like extensions of existing strategies\n\nWhile some components introduced in this work (multi-stage training based on sequence length, length-based bucketization, and cyclic schedules) have been used by previous works in different contexts, this is **the first work to show the efficacy of length-based curriculum for autoregressive LLM pre-training**, both in terms of training efficiency (faster training) and training performance (final model accuracy, especially for long-context metrics). We would like to emphasize that given the significant cost associated with LLM pretraining, the savings provided by the proposed method are very substantial (up to 45% faster training for models considered in the paper).\n\nIn addition, while our bucketing and curriculum (e.g., [1]) may look similar to prior works for different domains and setups, they are novel and have major differences as explained below:\n\n* **Bucketing**: We introduce binary decomposition (Section 2.1), a novel method to preserve document length, form buckets with fixed sequence lengths without using pad tokens, and avoid forming multi-document sequences (thus achieving no cross-document attention without attention masking). This is different from existing bucketing methods in different domains, where cross-document attention may still occur (albeit with reduced chances) without padding.\n* **Sequence length-based curriculum**: Our analysis shows the importance of mixing sequences with different lengths during training and curriculum (Table 2), beyond the simple multi-stage training considered by previous works.\n\nPlease also see [our response to reviewer bcif on the novelty](https://openreview.net/forum?id=r8M9SfYMDi&noteId=TtzwitydGc).\n\n[1] CYCLICAL CURRICULUM LEARNING - Kesgin et al. ArXiv 2022\n\n---\n> The training speed gain may be less significant with flash-attention-based LLMs and alternative models\n\nAll 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.\n\nPlease note that the proposed method is designed to speed up the training of transformers and further increase their performance without any architectural changes. Alternative architectures, such as state-space based models (e.g., Mamba) or approximations to attention (e.g., linear attention), do not suffer from quadratic attention costs but come with performance limitations. We would like to emphasize that transformers with multi-head attention are still the predominant architecture in most large-scale language models in the community.\n\n---\n> Can cross-document attention be restricted with an attention mask? Can that be explored?\n\nYes, we have already included results with attention masking to avoid cross-document attention. Baseline-8k-DM and Pack-8k+DM in Table 5 refer to models trained with Document-Masking (DM). Applying document masking mainly improves the regular evaluation, making it closer to our results. However, document masking does not provide the computational benefits (training speed) of our proposed method. Please **see Tables 1 and 2 in the rebuttal PDF**.\n\n---\n> Cyclic curriculum vs learning rate schedule\n\nIn 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).\n\n---\n> Details of the pacing mechanism\n\nWe 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 further clarify the details.\n\n```\n# {D_i}: list of buckets such that D_i includes sequences with length 2^i\n# {n_i}: total number of tokens to be picked from each bucket (see Table 1 of paper)\n# {o_i}: sampling odd for each bucket (see Table 2 of paper)\n# c: number of cycles\n# b: number of tokens per optimization step\n\n# Form c non-overlapping random subsets from each bucket D_i:\ns_{i,j} = random subset of D_i with n_i/c tokens # for j = 1, 2, ..., c\n\nfor j in [1, 2, ..., c]: # loop over cycles\n    while at least one s_{i,j} is not empty:\n        odds = [o_i if s_{i,j} not empty else 0 for i=1,2,3,...]\n        probs = odds / odds.sum()\n        randomly sample index i with probability probs[i]\n        sample b/2^i sequences from s_{i,j} without replacement and use for training\n```", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xDu0rMfVdg", "reviewer": "Reviewer_uip1", "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.", "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": "None.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed Grow-P2 curriculum is useful to practitioners if they want to pretrain large language models.", "weaknesses": "This paper presents only empirical results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer uip1 for their feedback. We respond to the reviewer’s concerns and questions below and kindly request that you let us know if further clarification is needed.\n\n---\n> This paper presents only empirical results.\n\nWe appreciate that our work is considered an empirical contribution, yet valuable to the community. We respectfully do not consider this a weakness.\n\n---\n> other dimensions for curriculum\n\nAs suggested by the reviewer, different notions of difficulty can be explored when designing a curriculum. In this work, our focus is on sequence length-based curriculum. We show that such a curriculum results in both training efficiency (through variable sequence length training) and performance gain.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all the reviewers for their time and feedback. We are pleased with the positive comments from the reviewers: Reviewer *bcif* finds the paper **well-organized and motivated**, the **method effective**, and the **experiments valuable to the community**; Reviewer *z2w6* finds our **experiments comprehensive**, our **analyses noteworthy**, and the **method simple and effective**; Reviewer *H2jF* finds our work **reasonably motivated**, and the **results extensive and promising**, and Reviewer *uip1* find our work **useful to practitioners**.

We provide multiple tables and figures in the **one-page PDF rebuttal** to further support the contributions of the paper and clarify questions and concerns raised by the reviewers. We show **more than 4x data efficiency** and up to **more than 6x training speed up** compared to the baseline for **large-scale trainings up to 1.1 trillion tokens**. We would like to emphasize that all contributions of this paper, including code and all model checkpoints, will be **open-sourced** to facilitate follow up works. We address individual reviews below and kindly request that you let us know if any further questions or concerns remain unaddressed.

Meta-review: 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.
Decision: Accept (poster)