Title: Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum
Venue: NeurIPS 2024
Accepted: yes
Keywords: Large Language Models

Efficient Training

Learning Efficiency

Long Context
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

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 webscale corpus demonstrate that our approach significantly enhances performance on standard language evaluations and long-context benchmarks, reaching target accuracy with up to 6 × faster training compared to the baseline. Our method not only enables efficient pretraining on long sequences but also scales effectively with dataset size. Lastly, we shed light on a critical yet less studied aspect of training large language models: the distribution and curriculum of sequence lengths, which results in a non-negligible difference in performance. *

1 Introduction

Large language models (LLMs) are often pretrained autoregressively (i.e., predicting the next token given a context) on large text corpora sourced from the web. Examples include The Pile [19], RefinedWeb [46], RedPajama [14], and DOLMA [57]. Each of these datasets comprises multiple documents, ranging from Wikipedia articles to books and code repositories. While the individual lengths of the documents vary from a few words (e.g., a message) to hundreds of thousands of words (e.g., a book), the training infrastructure often supports only a limited sequence length in a batch. To facilitate efficient training, document chunking is necessary. In this paper, we investigate the influence of document chunking, propose alternative strategies, and evaluate the proposed strategies with careful experiments.

2

33

2

34

2

35

2

36

2

37

2

38

2

39

2

40

# of seen tokens

40

45

Regular eval (%)

Dataset Decomposition

>4

⨉

data

efficiency

+2.4

(a) Data Efficiency

0

1000

2000

3000

4000

5000

GPU-Hours

40

45

Regular eval (%)

Dataset Decomposition

>6

⨉

training speed-up

(b) Computational Efficiency

Recent works [43, 37, 59, 60] popularized the concat-and-chunk approach to convert text datasets with variable document lengths into sequences with a fixed target length. In this approach, during a data preparation stage before training, we first randomly shuffle and concatenate all tokenized documents. Consecutive concatenated documents are separated by a special token <EOT> , allowing the model to detect document boundaries. We then chunk the concatenated sequence into subsequences with a target sequence length . For example, 2048 and 4096 for the Llama-1 and Llama-2 models, respectively. The model is then pretrained on batches of sequences with fixed length.

The concat-and-chunk approach has several shortcomings. First, randomly concatenating documents can lead to the model attending to a context from an unrelated document to predict the next token. While well-trained models learn to avoid cross-document attention, this is not explicitly enforced, leading to potential spurious modeling. Second, the cross-document attention spends unnecessary computation on attending to unrelated tokens that do not facilitate learning. This is especially crucial due to the quadratic complexity of the attention mechanism. Even with an implementation of attention that supports cross-document attention masking, the computational cost for each optimization step would be bottlenecked by the longest document in the global batch, leading to significant underutilization of devices with shorter documents. Third, even if a document is shorter than the target sequence length, it may still be broken into two chunks when they are at the boundary of two sequences. This results in significantly smaller average chunk lengths compared to the original document length average (see Fig. 3a), which hinders the model's capability.

Recent and concurrent works on LLM training try to improve the concat-and-chunk approach: document-masking is possible with recent implementation of attention [29] as adopted in some recent pre-training recipes [39], best-fit packing [17] to reduce document chunking, and concatenating semantically related documents instead of randomly [55]. However, none of them address all three issues mentioned above together.

In this work, we introduce dataset decomposition (DD), a novel approach to decompose data based on their length and train with variable sequence length (VSL) and length-based curriculum to address the above issues. We obtain significant both significant accuracy improvement and straining speed-up as shown in Fig. 1. DD decomposes a given dataset containing documents of variable lengths into a union of datasets/buckets, each with sequences of a fixed length. Specifically, a dataset D is decomposed into buckets ∪ i D i , where each bucket D i contains sequences of length 2 i , each extracted from a unique document. During training with VSL, at every step of the optimization process, we sample i (based on a curriculum) to form a batch with b/ 2 i sequences from the bucket D i , which

<latexi sh

1_b

64="nFp3

Hz

8

O

J

j

CLZTr

GPo

>A

c

VD

S

N

U

2vq yM9gI

k

RX

Y

Q

d

mE

f

+

0

u

5

w

W

K

7

<latexi sh

1_b

64="7SO8TkJnM

K

9

5X

ANgL

v

>

c

VD

F

3U

2

q

y

Ir

R

Z

C

Y

Q

d

pmE

f

+

z

G

j

o

0

P

u

w

H

W

<latexi sh

1_b

64="yCfM

v

IQ

z

7VE

5

p8Zwc

k

>A

B/n

DLS

N

FJ3

rP

du

U

q

X

Y

j

T

0m

g

R

O

9

H+

W

2

K

o

G

concat-and-chunk

original documents

dataset decomposition

D

1

D

2

D

3

keeps the total number of tokens in a batch constant ( 2 i × b/ 2 i = b ), regardless of which D i is sampled.

This approach gives us several advantages and resolves the aforementioned issues of the concatand-chunk method. First, DD is simple and has negligible computational overhead during the data preparation stage, making it easy to scale to large datasets. Second, tokens in each sequence are ensured to be from the same document by construction, which avoids cross-document attention. Furthermore, we have access to the sequence length distribution (an auxiliary prior knowledge) which can be used to create different mixtures/curricula for training. Finally, our VSL training strategy accelerates training time: the latency for one optimization step is less when sampling from D i with smaller i (due to attention's quadratic complexity). Following is a summary of our contributions:

We introduce DD, a method to efficiently decompose a dataset of variable-length documents into a union of buckets with fixed-length sequences. DD enables efficient and robust training via VSL and length-based curriculum.

We perform large-scale experimentation using different models, datasets, and evaluation tasks to demonstrate the efficacy of the proposed method. We show (see Fig. 1) significant gains in data efficiency ( > 4 × ) and compute efficiency (11% to 45%), resulting in combined LLM pretraining acceleration of up to 6 × (time to reach certain accuracy compared to baseline).

Through careful experimentation, we study the importance of sequence length distribution and mixture during pretraining for different natural language and long-context tasks. We show the effect of concatenation and chunking operations to synthetically alter sequence length (Section 3.2).

2 Method

2.1 Dataset decomposition

Given a dataset D of tokenized documents { d 1 , d 2 , . . . , d n } , the goal of dataset decomposition (DD) is to reorganize D as a union of buckets, ∪ i D i , such that: (1) each bucket D i consists of sequences of tokens with length l i ; (2) each sequence s ∈ D i is a subsequence of one document d ∈ D ; and (3) each token in D appears in exactly one D i . This decomposition produces sequences that each belong to a unique document, ensuring no cross-document attention within a sequence during training. Additionally, all sequences in a given bucket D i have the same length l i , enabling efficient batching.

Dataset decomposition as defined above is not unique. We propose a specific decomposition, with l i = 2 i , to optimally maintain the original document sequence length distribution while also enabling efficient batch pretraining, as explained in Section 2.2. We apply decomposition at the document level, which makes it very easy to integrate the method into any existing data preparation pipeline (a stage before model training) and is scalable to large datasets. For a tokenized document d ∈ D with length l , where l = 2 i 1 +2 i 2 + . . . +2 i k represents its binary decomposition, we break d into k adjacent sequences s 1 , . . . , s k , with lengths of 2 i 1 , . . . , 2 i k , respectively. Each sequence s j of length 2 i j is then assigned to bucket D i j . Fig. 2 shows a schematic representation of this method.

With our proposed dataset decomposition approach, each bucket D i contains sequences extracted from an original document d such that the length of d is at least 2 i . In Fig. 3b, we show the distribution

2

5

2

6

2

7

2

8

2

9

2

10

2

11

2

12

2

13

2

14

Sequence length

10

-

7

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

PDF

3.8%

0.3%

0.6%

Orig. Docs

mean=598

Baseline-2k

mean =463

mean =558

Pack-8k

mean =585

(a)

0

1

2

3

4

5

6

7

8

9

10 11 12 13 14 15 16 17

Bucket number

2

4

6

8

10

12

14

16

Tokens %

0

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

(b)

2

5

2

6

2

7

2

8

2

9

2

10

2

11

2

12

2

13

Context length

10

-

5

10

-

4

10

-

3

PDF

Pack-8k

mean=1064

DD

256

mean=1344

(c)

of RefinedWeb dataset tokens over different buckets, where D 9 (corresponding to sequences with length 512) has the maximum tokens. We also highlight the original document lengths from which tokens are extracted. Most tokens in a bucket D i are extracted from documents with length l such that 2 i ≤ l < 2 i +1 , and some tokens are rolled over from documents with length l ≥ 2 i +1 . This demonstrates the efficacy of the method in retaining original document length, especially for long documents, which are scarce.

In Fig. 3a, we show the distribution of original document lengths and chunks within 2048 and 8192 target sequence lengths formed by the concat-and-chunk approach. We also present the length distribution using the bin-packing approximate algorithm introduced by a concurrent work [17]. Additionally, in Fig. 3c, we show the distribution of context length (the number of tokens from the same document a token can attend to during pretraining) when using baselines with a target sequence length of 8192 and DD. See Appendix F for additional discussion on sequence length statistics.

In contrast to the concat-and-chunk approach, which results in a static dataset, DD enables us to use sequence length distribution as prior knowledge and optimize the best mixture for the target task. In Section 3.2, we show the bias of each target evaluation toward a sequence length and the effect of concatenation and chunking on model performance. In Section 3.3, we study the effect of different sequence mixtures for LLM pretraining, a less-studied topic in LLM pretraining.

2.2 Variable sequence length training

Following the setup in Section 2.1, we assume a set of k buckets such that D i , containing sequences with length 2 i , are available. Let b be the target batch size - the number of tokens used per optimization step. In variable sequence length (VSL) training, at every step of optimization, we first sample i from available choices, then pick b/ 2 i sequences from bucket D i . Since D i consists of sequences with length 2 i , the number of seen tokens per optimization step remains b , independent of the choice of i . Training LLMs with the VSL algorithm comes with several advantages.

First, since the total number of seen tokens per optimization step does not change, VSL does not alter optimization dynamics, and the same hyperparameters as the baseline can be utilized (see Section 3).

Second, in Section 3.1, we show that the time to complete one optimization step (forward+backward) for a fixed b (tokens per step) varies by sequence length due to the quadratic cost of attention [63]. With VSL training, the cost of every optimization step depends on the bucket D i sampled for that step (and hence the sequence length). Thus, the more expensive steps (corresponding to long sequences) are compensated with less expensive steps (corresponding to short sequences).

2

6

2

7

2

8

2

9

2

10

2

11

2

12

2

13

Context length

200

400

600

800

1000

One step training latency (ms)

+35%

+88%

+23%

OpenLM-1B

OpenLM-3B

OpenLM-7B

(a)

(b)

Finally, the sampling component in VSL (which D i to choose at every optimization step) enables different curricula of sequence lengths. In Section 3.4, we show the significance of such curricula on model stability and generalization accuracy.

3 Experiments and analysis

In this section, we show the efficacy of the proposed method to train LLMs of different sizes on largescale datasets and provide additional analyses. For all experiments, except the results in Section 3.5, we use RefinedWeb [46] filtering of Common Crawl [1] with a total of ∼ 525 billion tokens using the EleutherAI/gpt-neox [9] tokenizer (vocabulary size is 50,432). Model architectures and training code are based on the OpenLM [22] † . For all experiments, other than model scaling in Section 3.5, we use the OpenLM-1B model with an 8k context length. Please refer to Appendix B for implementation details of all experiments.

Positional encoding We use Rotary Positional Embedding (RoPE) [58] to encode positions in queries and keys before the attention module. RoPE rotates the consecutive components of queries and keys with a base frequency f b = 10 , 000 . Recent studies [48, 64, 36] have suggested increasing f b to better adapt a pretrained model for longer sequences through fine-tuning. We find that using a larger f b is also beneficial when training LLMs from scratch. In Table 4, we show that increasing f b to 100,000 improves performance for both the baseline and DD methods.

Evaluation We evaluate each model on a comprehensive set of standard benchmarks, mainly using LLM Foundry [2]. We report averaged accuracies over each category, as well as the regular average , which is the average of 14 regular language modeling benchmarks detailed below:

Commonsense Reasoning (CSR) : PIQA-0-shot [8], COPA-0-shot [52], and OpenBookQA-10shots [40].

Language Understanding (LU) : Lambada-OpenAI [44], Hellaswag-0-shot [65], Winograd-3shots [30], and WinoGrande-5-shots [54].

Reading Comprehension (RC) : SQuAD-3-shots [50], BoolQ-0-shot [12], and CoQA-0-shot [51].

World Knowledge (WK) : Jeopardy-3-shots [3], ArcEasy-3-shots [13], ArcChallenge-3shots [13], and WikiDataQA-3-shots [4]

To evaluate model on longer context tasks, we adopt the following real-world benchmarks:

Multi-Document Question Answering (MDQA) : We follow the exact setup as in Liu et al. [35], where for each question from NaturalQuestions-Open [28, 27], r Wikipedia documents are retrieved such that one of them has the answer to the question, and the other r -1 documents are distractors. We report MDQA-10, MDQA-20, and MDQA-30 accuracy corresponding to r = 10 , 20 , and 30,

(a)

(b)

(c)

60

ol

CSR

o

MDQA-10

CSR

D7

40.8

50

LU

MDQA-20

{1

+

WK

RC

MDQA-30

LU

D10

48

40

RC

30

WK

D13

44.5

+2.6

20

MDQA-10

D13→10

47.1

-0.9

MDQA-20

10

D7→10

40.8

MDQA-30

0

9G

27

Pretraining sequence length

28

29

210

2i1

212

213

Sequencelength

213

38

40.5

Regularaverage

43

45.5

48

respectively. For each query, we evaluate the model by changing the location of the target document among distractors and report the averaged accuracy.

TOEFL : This dataset is a multiple-choice question answering dataset from An et al. [5]. The dataset contains QA pairs for 15 longest lectures in Tseng et al. [61], Chung et al. [11]. Only one of the choices is the correct response. We estimate the correct choice by picking the choice with the lowest mean log probability value.

QuALITY : This dataset is a multiple-choice question answering dataset from An et al. [5]. The dataset contains a long passage for context, followed by a question with multiple choices. Only one of the choices is the correct response. We estimate the correct choice by picking the choice with the lowest mean log probability value.

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 (OpenLM-1B/3B/7B) and different context lengths ( 2 6 to 2 13 ) and measure the time to train 100 batches with a fixed global batchsize of b = 8 × 8192 distributed over 8 GPUs in a single node. We repeat this 5 times and report the average time per optimization step in Fig. 4a (with STD mostly < 1 ms). See Appendix C.1 for additional results with different batchsizes b . For each model, we highlight the training time overhead (due to attention's quadratic complexity with an optimized FlashAttention2 kernel [15]) when training with 8192 context lengths compared to 64 context lengths: +35%, +88%, and +23% for OpenLM-1B, -3B ‡ , and -7B, respectively. Training overhead grows for longer context lengths (see Fig. 7 for results up to 16k context length).

The concat-and-chunk baseline method always operates at a fixed sequence length. For example, for the OpenLM-1B model, an optimization step with concat-and-chunk takes 243ms and 304ms for target context lengths of 2048 and 8192, respectively. The expected time for VSL, on the other hand, is the weighted average over different sequence lengths depending on the mixture. In Table 1, we report the training step time for different mixtures. For example, with the natural length distribution resulting from DD (Fig. 3b), training up to length 8192 sequences takes a similar time (244ms) as baseline training with length 2048 (with 243ms per step) per step-equivalent to a 20% training time reduction compared to baseline training with a fixed length of 8192 (with 304ms per step).

3.2 Sequence length bias

In this section, we study the effect of pretraining data sequence length on model performance in isolation. Using a single bucket D i as the dataset, we train an LLM from scratch on sequences with length 2 i for a total of 2 34 seen tokens. Note that the number of tokens per optimization step is fixed at 256, irrespective of sequence length. We use the same training hyperparameters for all runs. In Appendix C.2, we show that our conclusions do not depend on the choice of hyperparameters. To

reduce statistical error, we train each model twice from scratch with different random seeds and report the average metric for each benchmark (observing an average standard deviation of ∼ 0 . 3 for regular benchmarks and ∼ 1 . 6 for multi-document QA). Results are demonstrated in Fig. 5a.

We show a significant correlation between pretraining sequence length and different benchmarks. Specifically, the accuracy of commonsense reasoning, language understanding, and world knowledge shows an inverted U-shape behavior with respect to pretraining sequence length, while reading comprehension benefits from longer sequences. This behavior can be associated with training-test distribution alignment with respect to sequence length. In Fig. 5b, we show the length distribution for different benchmarks where RC demonstrates a heavier tail compared to CSR, LU, and WK. Multi-document QA benchmarks show a vivid correlation with respect to sequence length: test accuracy is ≈ 0 unless pretraining sequence length is greater than the test context length, which is ∼ 2k, 4k, and 6k for MDQA-10, -20, and -30, respectively.

It could be argued that data selection based on sequence lengths could introduce bias since the content (or source) of the documents might change based on the sequence lengths. To better understand the effect of sequence length on common metrics, we created two new buckets, D 13 → 10 and D 7 → 10 , from existing buckets D 13 and D 7 , respectively. The bucket D 13 → 10 contains sequences of length 2 10 created by chunking sequences from D 13 into 8 subsequences and then performing a global shuffle. The bucket D 7 → 10 also includes sequences of length 2 10 , each formed by concatenating 8 random sequences from D 7 .

In Fig. 5c, we compare the regular average metric of models pretrained on these buckets; for each bucket, we train two models from scratch using different random seeds and report the averaged results. D 13 → 10 gains 2.6 points compared to D 13 while including the same content. This demonstrates the pure effect of sequence length on model accuracy. Furthermore, training on D 13 → 10 underperforms D 10 by 0.9 points, even though they are of the same length, indicating that long documents (used to construct D 13 → 10 ) correlate less with our benchmarks than short documents (used to construct D 10 ). Finally, we show that concatenation , as opposed to chunking , does not mitigate length correlation. This is evident from the fact that D 7 → 10 scores the same as D 7 and still significantly worse than D 10 .

Our analysis suggests that effective base model pretraining requires a mixture of different sequence lengths to perform well on all benchmarks. Next, we systematically study the effect of dataset mixture from the sequence length perspective.

3.3 Data mixture

A key benefit of dataset decomposition is access to and control over sequence length distribution. We form datasets with different mixtures of sequence lengths and explore the performance of a model trained on each mixture. Table 1 shows the results. For all experiments, the total seen tokens and hyperparameters are fixed, and only the distribution over sequence length is changed. First, we observe that mixtures with small average context length (we provide the exact definition in Appendix F) perform poorly on MDQA, which requires long context understanding. For example, as for '1k-only', ' ≤ 2k', and 'Mid' distributions that do not include long sequences from D 12 and D 13 . Larger average context length (e.g., as in ' ≥ 1k') also correlates positively with performance on

reading comprehension tasks, consistent with our observation in Fig. 5a, but comes at the cost of a longer training step time.

Furthermore, '1k-only', that is training using only the best sequence length ( = 1024 ) from the study in Section 3.2 results in good performance on regular evaluations, especially for language understanding and world knowledge tasks, but is poor for long context tasks. Finally, we observe that 'natural' mixture, that is aligned with the distribution resulting from dataset decomposition (see Fig. 3b), obtains near-optimal performance on both regular and MDQA tasks, demonstrating the scalability of the proposed approach to large datasets without a need for intervention on the natural underlying length distribution.

3.4 Length-based curriculum

We can think of short sequences as being "easier" compared to longer ones; hence motivating a curriculum learning [7, 18] that prioritizes short sequences. A similar idea (training with image resolutions from low to high) is explored in vision to train CLIP [49] models more efficiently [33]. In VSL, we can easily implement curriculum learning through sampling designs. At every optimization step, we sample without replacement a batch with b tokens from bucket D i with probability p i . If a bucket is empty, we exclude it from sampling. We study different curricula for the " ≥ 256 " mixture (with an equal number of tokens in D 8 , . . . , D 13 ). Results are shown in Table 2. For each curriculum, we determine the odds of picking a batch from each bucket ( = p i 's when normalized). Details of our length-based sampling and curriculum are provided in Algorithm 1. We consider curricula that shift from short to long sequences at different paces controlled by p i 's changing linearly, with powers of 2, and with powers of 100 between buckets.

Due to the presence of other hyperparameter schedules during the course of training (e.g., learning rate and weight decay), a curriculum on length may result in a potential implicit bias. For example, if we only see long sequences toward the end of training, long sequence learning occurs only when the learning rate is too small. To address this potential issue, we also explore cyclic curricula, where a curriculum is applied in cycles similar to cyclic learning rate schedules [56] as shown in Fig. 6. Note that when we train on a sequence of length l , we have l next-token prediction losses (applied in parallel) with context lengths 0 , 1 , . . . , l -1 . This already implies some mixing: when training on a 'hard' example (i.e., a long sequence), we also include 'easy' examples (its shorter sub-sequences). Therefore, even towards the end of each cycle, we still have some losses with short contexts.

Our results show that the cyclic "Grow-P2" curriculum is near optimal with different metrics. An additional benefit of curriculum is training stability. Li et al. [31] noticed that long sequences contribute to extreme gradient variance, especially at the beginning of training, resulting in instability. We also observe (see Appendix E) that our proposed approach with curriculum results in more stable training dynamics, thus enabling more efficient training with larger batch sizes and learning rates.

3.5 Scaling

Dataset scaling In Fig. 1a, we show the performance of models trained with 2 34 , 2 35 , 2 36 , 2 37 , and 2 38 total tokens using DD and baseline. We use the ' ≥ 256 ' mixture and 'Grow-Linear' curriculum with 8 cycles for DD, and a fixed target sequence length 8192 for the baseline. Results show > 2 × data efficiency: our proposed method reaches the same accuracy as the baseline using less than half the tokens.

Model scaling We report results on OpenLM-1B, -3B, and -7B trained from scratch for a total of 2 37 tokens in Fig. 4b. We compare baseline training with a fixed target sequence length 8192 and VSL training with a DD ≥ 256 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 on another large-scale dataset, DataComp-LM [32]. We train models with different numbers of parameters: OpenLM-160M, -410M, and -1B, for a total of 137B tokens. We compare the baseline with a DD ≥ 256 mixture trained with the "Grow-P2" curriculum with 8 cycles. Results are reported in Table 3, demonstrating significant accuracy and training efficiency gains.

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 packing [17], and in-context pretraining (ICLM) [55]. We describe the details of our implementation of the best-fit packing in Appendix D. For ICLM, we use the official implementation § applied to the RefinedWeb dataset. The results are shown in Table 5.

Pre-training context length is an important factor in determining a model's long-context performance. We empirically validate this in the results shown in Fig. 5a, where models trained on longer sequences perform better on multi-document QA. Our proposed method has an average context length (as defined in Eq. (2)) of 1,344 for the RefinedWeb dataset, compared to 930 for the baseline (see Fig. 3c) and 1,064 when packing [17] is applied. This explains why the dataset decomposition mixture, even without any length-based curriculum (the first row in Table 2), outperforms Baseline-8k-DM and Pack-8k+DM (second and third rows in Table 5). Here, DM refers to applying document masking during training to avoid cross-document attention.

Document masking improves the baseline on regular evaluations from 51 . 5 to 52 . 4 by preventing cross-document attention. However, Xiong et al. [64] demonstrate that including concatenated unrelated documents can still enhance long-context metrics compared to training solely with shorter sequences. Therefore, DM experiences a slight decline in long-context evaluations, dropping from 27 . 5 to 27 . 1 . Baseline-8k multi-document QA performance is even slightly better than our proposed dataset decomposition mixture when used without length-based curriculum (the first row in Table 2).

In-context pre-training LMs (ICLM) [55] proposes document sorting based on content similarity. Although the benefits of ICLM with large-scale Common Crawl data (used in our experiments) are marginal in regular evaluation, we observe that ICLM results in slightly better multi-document QA performance than Baseline-8k when 30 documents are in the context compared with Baseline-8k (22.0% vs. 20.5%). The average long-context metric boosts from 27 . 5 for Baseline-8k to 28 . 7 for ICLM. However, the similarity finding step proposed by ICLM is resource-intensive at scale ¶ .

Finally, as shown in in Table 2 our proposed cyclic length-based curriculum, for example, Grow-P2 with 8 cycles, results in a significant improvement in the model's long-context capability. Our proposed method avoids cross-document attention to unrelated content, maintains coherent long sequences, and benefits from a length-based curriculum, effectively improving performance in both regular and long-context evaluations compared to all baselines. We further summarize long-context performance of different methods discussed above in Table 6.

4 Related works

Recent works have raised concerns regarding cross-document attention. For example, Llama3 [39], ICLM [55], and [17], which we discussed in Section 3.6. Similarly, [26] discuss challenges with the baseline concat-andchunk approach and propose an approximate bin-packing algorithm.

Related to our study on sequence length bias, [62] shows the importance of train-vs-test time distribution shift from a sequence length perspective on a string editing task. [6, 66, 25, 36]

highlight the challenge of generalizing to lengths beyond what the model has seen during training and discuss the importance of positional encoding. Several works [41, 67, 23, 64, 10, 47, 48, 53, 34] address enabling LLM inference with long context (see [45] for an overview). These approaches are orthogonal to our contribution and can be applied post-pretraining to adapt to longer lengths. GrowLength [24] proposes accelerating LLM pretraining by progressively growing context length using the baseline sequence formation method, but does not show results on LLMs. Similarly, increasing sequence length has been shown in BERT model training [42] to improve compute efficiency.

The idea of dynamic batching has been explored in other domains. In vision, methods like NaViT [16, 38] use images with variable resolutions (a similar concept to context length for LLMs). In seq-to-seq tasks (e.g., automatic speech recognition, text-to-speech, and neural machine translation), the inputs have different lengths. An efficient approach is to sort inputs by their length and form batches of inputs with similar lengths during training (after possible padding). Batchsize is dynamically adjusted inversely proportional to input lengths [20, 21]. Different from these works, in dataset decomposition, we do not simply put documents with similar lengths into the same bucket. Instead, we decompose each document into multiple subsequences and form multiple buckets. We form batches with different lengths during training by sampling from these buckets using a target mixture and curriculum.

5 Conclusion and limitations

In this paper, we explore the shortcomings of a popular LLM pretraining approach, concat-andchunk, and introduce dataset decomposition, a method to decompose a dataset of text documents into buckets containing fixed sequence lengths. We show results of variable sequence training using DD with different mixtures, curricula, datasets, and models, demonstrating significant LLM pretraining speedup and a final model accuracy boost on a wide range of benchmarks. Furthermore, we provide analysis on sequence length bias and attention masking. We compare our proposed method with recent works that also address concat-and-chunk shortcomings in a unified experimental setup and show gains in data preparation cost, training time, and final model accuracy.

Limitations. The training speed gains compared to the baseline are significant only when the target sequence length is long enough. Otherwise, the attention cost is not a dominant fraction of training, and hence no significant training speedup is expected.

References

Common crawl. https://commoncrawl.org .

Llm foundry v0.7.0. https://github.com/mosaicml/llm-foundry .

Jeopardy. https://huggingface.co/datasets/jeopardy . [Used custom curated version by LLM Foundry].

Big-bench qa wikidata. https://github.com/google/BIG-bench/tree/main/ bigbench/benchmark_tasks/qa_wikidata . [Used through LLM Foundry].

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.

C. Anil, Y. Wu, A. Andreassen, A. Lewkowycz, V. Misra, V. Ramasesh, A. Slone, G. Gur-Ari, E. Dyer, and B. Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems , 35:38546-38556, 2022.

Y. Bengio, J. Louradour, R. Collobert, and J. Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning , pages 41-48, 2009.

Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence , volume 34, pages 7432-7439, 2020.

S. Black, S. Biderman, E. Hallahan, Q. Anthony, L. Gao, L. Golding, H. He, C. Leahy, K. McDonell, J. Phang, et al. Gpt-neox-20b: An open-source autoregressive language model. arXiv preprint arXiv:2204.06745 , 2022.

S. Chen, S. Wong, L. Chen, and Y. Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595 , 2023.

Y.-A. Chung, H.-Y. Lee, and J. Glass. Supervised and unsupervised transfer learning for question answering. In M. Walker, H. Ji, and A. Stent, editors, Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers) . Association for Computational Linguistics, 2018.

C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv:1905.10044 , 2019.

P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457 , 2018.

T. Computer. Redpajama: An open source recipe to reproduce llama training dataset, 2023. URL https://github.com/togethercomputer/RedPajama-Data .

T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691 , 2023.

M. Dehghani, B. Mustafa, J. Djolonga, J. Heek, M. Minderer, M. Caron, A. Steiner, J. Puigcerver, R. Geirhos, I. M. Alabdulmohsin, et al. Patch n'pack: Navit, a vision transformer for any aspect ratio and resolution. Advances in Neural Information Processing Systems , 36, 2024.

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.

J. L. Elman. Learning and development in neural networks: The importance of starting small. Cognition , 48(1):71-99, 1993.

L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 , 2020.

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.

P. Gonzalez, T. S. Alstrøm, and T. May. On batching variable size inputs for training end-toend speech enhancement systems. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 1-5. IEEE, 2023.

S. Gururangan, M. Wortsman, S. Y. Gadre, A. Dave, M. Kilian, W. Shi, J. Mercat, G. Smyrnis, G. Ilharco, M. Jordan, R. Heckel, A. Dimakis, A. Farhadi, V. Shankar, and L. Schmidt. OpenLM: a minimal but performative language modeling (lm) repository, 2023. URL https://github.com/mlfoundations/open_lm/ . GitHub repository.

C. Han, Q. Wang, W. Xiong, Y. Chen, H. Ji, and S. Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv:2308.16137 , 2023.

H. Jin, X. Han, J. Yang, Z. Jiang, C.-Y. Chang, and X. Hu. Growlength: Accelerating llms pretraining by progressively growing training length. arXiv preprint arXiv:2310.00576 , 2023.

A. Kazemnejad, I. Padhi, K. Natesan Ramamurthy, P. Das, and S. Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems , 36, 2024.

M. M. Krell, M. Kosec, S. P. Perez, and A. Fitzgibbon. Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance. arXiv preprint arXiv:2107.02027 , 2021.

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.

K. Lee, M.-W. Chang, and K. Toutanova. Latent retrieval for weakly supervised open domain question answering. arXiv preprint arXiv:1906.00300 , 2019.

B. Lefaudeux, F. Massa, D. Liskovich, W. Xiong, V. Caggiano, S. Naren, M. Xu, J. Hu, M. Tintore, S. Zhang, P. Labatut, D. Haziza, L. Wehrstedt, J. Reizenstein, and G. Sizov. xformers: A modular and hackable transformer modelling library. https://github.com/ facebookresearch/xformers , 2022.

H. Levesque, E. Davis, and L. Morgenstern. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning , 2012.

C. Li, M. Zhang, and Y. He. The stability-efficiency dilemma: Investigating sequence length warmup for training gpt models. Advances in Neural Information Processing Systems , 35: 26736-26750, 2022.

J. Li, A. Fang, G. Smyrnis, M. Ivgi, M. Jordan, S. Gadre, H. Bansal, E. Guha, S. Keh, K. Arora, et al. Datacomp-lm: In search of the next generation of training sets for language models. arXiv preprint arXiv:2406.11794 , 2024.

X. Li, Z. Wang, and C. Xie. An inverse scaling law for clip training. Advances in Neural Information Processing Systems , 36, 2024.

H. Liu, W. Yan, M. Zaharia, and P. Abbeel. World model on million-length video and language with blockwise ringattention. arXiv preprint arXiv:2402.08268 , 2024.

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.

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.

Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

S. Mehta, F. Abdolhosseini, and M. Rastegari. Cvnets: High performance library for computer vision. In Proceedings of the 30th ACM International Conference on Multimedia , pages 7327-7330, 2022.

Meta. Introducing meta llama 3: The most capable openly available llm to date, 2024. URL https://ai.meta.com/blog/meta-llama-3 .

T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02789 , 2018.

A. Mohtashami and M. Jaggi. Random-access infinite context length for transformers. Advances in Neural Information Processing Systems , 36, 2024.

K. Nagatsuka, C. Broni-Bediako, and M. Atsumi. Pre-training a bert with curriculum learning by increasing block-size of input text. In Proceedings of the International Conference on Recent Advances in Natural Language Processing (RANLP 2021) , pages 989-996, 2021.

M. Ott, S. Edunov, A. Baevski, A. Fan, S. Gross, N. Ng, D. Grangier, and M. Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations , 2019.

D. Paperno, G. Kruszewski, A. Lazaridou, Q. N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031 , 2016.

S. Pawar, S. Tonmoy, S. Zaman, V. Jain, A. Chadha, and A. Das. The what, why, and how of context length extension techniques in large language models-a detailed survey. arXiv preprint arXiv:2401.07872 , 2024.

G. Penedo, Q. Malartic, D. Hesslow, R. Cojocaru, A. Cappelli, H. Alobeidli, B. Pannier, E. Almazrouei, and J. Launay. The RefinedWeb dataset for Falcon LLM: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01116 , 2023. URL https://arxiv.org/abs/2306.01116 .

B. Peng and J. Quesnelle. Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation. 2023. URL https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ ntkaware_scaled_rope_allows_llama_models_to_have .

B. Peng, J. Quesnelle, H. Fan, and E. Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071 , 2023.

A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning , pages 8748-8763. PMLR, 2021.

P. Rajpurkar, R. Jia, and P. Liang. Know what you don't know: Unanswerable questions for squad. arXiv preprint arXiv:1806.03822 , 2018.

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.

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.

A. Ruoss, G. Delétang, T. Genewein, J. Grau-Moya, R. Csordás, M. Bennani, S. Legg, and J. Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint arXiv:2305.16843 , 2023.

K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM , 64(9):99-106, 2021.

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.

L. N. Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter conference on applications of computer vision (WACV) , pages 464-472. IEEE, 2017.

L. Soldaini, R. Kinney, A. Bhagia, D. Schwenk, D. Atkinson, R. Authur, B. Bogin, K. Chandu, J. Dumas, Y. Elazar, V. Hofmann, A. H. Jha, S. Kumar, L. Lucy, X. Lyu, N. Lambert, I. Magnusson, J. Morrison, N. Muennighoff, A. Naik, C. Nam, M. E. Peters, A. Ravichander, K. Richardson, Z. Shen, E. Strubell, N. Subramani, O. Tafjord, P. Walsh, L. Zettlemoyer, N. A. Smith, H. Hajishirzi, I. Beltagy, D. Groeneveld, J. Dodge, and K. Lo. Dolma: an Open Corpus of Three Trillion Tokens for Language Model Pretraining Research. arXiv preprint , 2024.

J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing , 568:127063, 2024.

H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

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.

D. Variš and O. Bojar. Sequence length is a domain: Length-based overfitting in transformer models. arXiv preprint arXiv:2109.07276 , 2021.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

W. Xiong, J. Liu, I. Molybog, H. Zhang, P. Bhargava, R. Hou, L. Martin, R. Rungta, K. A. Sankararaman, B. Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.16039 , 2023.

R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830 , 2019.

Y. Zhou, U. Alon, X. Chen, X. Wang, R. Agarwal, and D. Zhou. Transformers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371 , 2024.

D. Zhu, N. Yang, L. Wang, Y. Song, W. Wu, F. Wei, and S. Li. Pose: Efficient context window extension of llms via positional skip-wise training. arXiv preprint arXiv:2309.10400 , 2023.

A Broader impacts

This work enables faster training of LLMs, which are among the most compute-intensive applications in the field. A positive societal/environmental impact of this work is training LLMs with a smaller carbon footprint.

Another potential societal advantage of this work is training LLMs with fewer hallucinations. While we did not directly measure this potential benefit, a concurrent work [17] shows such a benefit when cross-document attention is not allowed during LLM pretraining.

B Implementation details

B.1 Training details

Software and hardware details All experiments in this paper are conducted using the OpenLM || repository, which is based on PyTorch. We use Fully Sharded Data Parallelism (FSDP) with Bfloat16 mixed precision for all experiments. We use the Xformers [29] implementation for attention. For hardware, we use one or more nodes of 8 × NVIDIA H100 GPUs (Hopper architecture), each with 80GB memory, and 192 CPU cores with 2000GB of RAM. Nodes are connected through Elastic Fabric Adapter (EFA) for efficient inter-node communication hosted by AWS.

Model architecture details We provide details of all architectures used in the paper in Table 7 to Table 11.

Baseline hyper parameters We list our baseline hyperparameters in Table 12 and iterate over changes for each section next. Note that we did not explicitly optimize hyperparameters for any of the experiments, and we always use the same hyperparameters when using either the baseline method or ours.

Implementation details of Section 3.2 experiments Experiments in this section are done using the same hyperparameters as in Table 12 for a total of 2 34 tokens on the OpenLM-1B model. We trained each model twice with different random seeds and report the averaged results. For models in this section, we use RoPE with f b = 10 , 000 . In Table 12, we show that our results and conclusions in this section are not sensitive to hyperparameters, including the RoPE base frequency f b .

Implementation details of Section 3.3 and Section 3.4 experiments Experiments in this section are done with OpenLM-1B model, trained for total of 96 × 10 34 ≈ 103 B tokens. Hyper-parameters are the same as Table 12, except we used 20000 warmup steps for all models presented in this section. We use RoPE with f b = 10 , 000 for all models in Section 3.3 and f b = 100 , 000 for models in Section 3.4.

Implementation details of Section 3.5

Dataset scaling : Experiments in this section are trained with the OpenLM-1B model, RoPE with f b = 100 , 000 , and the baseline setup as in Table 12 except for the following changes for different dataset sizes:

total tokens = 2 34 , warmup steps = 5 , 000

total tokens = 2 35 , warmup steps = 5 , 000

total tokens = 2 36 , warmup steps = 10 , 000

total tokens = 2 37 , warmup steps = 20 , 000

total tokens = 2 38 , warmup steps = 40 , 000

Model scaling : Experiments in this section are trained with the OpenLM-1B, OpenLM-3B, and OpenLM-7B models, 2 37 ≈ 137 B total seen tokens, RoPE with f b = 100 , 000 , and the baseline setup as in Table 12 except for the following changes for different model sizes:

OpenLM-1B, warmup steps = 20 , 000 , max-lr = 3 × 10 -3 , batchsize b = 2 19 , with 32 H100 GPUs

OpenLM-3B, warmup steps = 20 , 000 , max-lr = 2 × 10 -3 , batchsize b = 2 20 , with 64 H100 GPUs

OpenLM-7B, warmup steps = 20 , 000 , max-lr = 1 × 10 -3 , batchsize b = 2 22 , with 128 H100 GPUs

Alternative dataset : Experiments in this section are trained with the OpenLM-160M, OpenLM410M, and OpenLM-1B models, 2 37 ≈ 137 B total seen tokens, RoPE with f b = 100 , 000 , and the baseline setup as in Table 12 except for the following changes for different model sizes:

OpenLM-160M, warmup steps = 20 , 000 , max-lr = 5 × 10 -3 , weight-decay = 0 . 033 , with 16 H100 GPUs

OpenLM-410M, warmup steps = 20 , 000 , max-lr = 4 × 10 -3 , weight-decay = 0 . 066 , with 16 H100 GPUs

OpenLM-1B, warmup steps = 20 , 000 , max-lr = 3 × 10 -3 , weight-decay = 0 . 1 , with 32 H100 GPUs

For the DD experiments we used 'Grow-P2' length curriculum which is visualized in Fig. 6.

Sequence length

213

0.005

0.004

0.003

0.002

0.001

0.000

0

10000

20000

30000

40000

50000

60000

Optimization step

Implementation details of Section 3.6 All experiments in this section are done with the OpenLM1B model, trained for a total of 96 × 10 34 ≈ 103 B 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 b = 100 , 000 .

B.2 Length based sampling and curriculum algorithm

We present the details of our length-based sampling and curriculum in Algorithm 1.

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 response by using only the text before the first occurrence of a newline character as the answer. We noticed that sometimes the model responds with multiple newline characters before providing any valid text. In view of this behavior, we updated the evaluation script to look for the first non-empty text output from the model instead of the first string after newline character. Apart from this change in processing the model output, the rest of the evaluation follows the official implementation [35].

TOEFL We follow the setup described in [5]. As described in Section 3, the dataset contains multiple-choice QA pairs for the 15 longest lectures in [61, 11]. To obtain a response from the model, we follow MMLU-style prompting, where the choices are appended to the original prompt individually and the mean log-probability is computed for each choice. The choice corresponding to

64

128

256

512

1024

2048

4096

8192

16384

sequence length

1

2

4

8

16

32

64

128

256

number of sequences per batch

100

99

100

100

100

107

153

304

719

99

100

100

101

106

143

265

566

100

102

100

105

138

243

487

100

100

106

137

232

445

102

105

135

227

426

105

135

225

416

136

224

411

224

409

410

(a) Average time for OpenLM-1B

64

128

256

512

1024

2048

4096

8192

16384

sequence length

1

2

4

8

16

32

64

128

256

number of sequences per batch

144

143

143

144

154

203

344

777

2218

143

143

144

151

190

292

589

1452

143

144

152

185

267

501

1095

144

150

180

254

449

924

149

180

245

428

832

180

242

411

789

246

406

758

412

751

760

(b) Average time for OpenLM-3B

64

128

256

512

1024

2048

4096

8192

16384

sequence length

1

2

4

8

16

32

64

128

256

number of sequences per batch

269

270

269

273

286

348

562

1091

269

268

274

285

344

538

993

269

274

285

341

526

938

273

284

340

520

911

285

340

516

899

340

515

891

515

888

889

(c) Average time for OpenLM-7B

64

128

256

512

1024

2048

4096

8192

16384

sequence length

1

2

4

8

16

32

64

128

256

number of sequences per batch

0.53

0.22

0.18

0.17

0.52

0.09

0.23

2.02

2.60

0.13

0.25

0.45

0.45

0.10

0.36

0.77

0.62

0.27

1.05

0.24

0.19

0.25

0.67

1.00

0.10

0.18

0.38

0.04

0.50

0.63

0.38

0.31

0.15

0.59

0.72

0.42

0.25

0.46

0.62

0.16

0.59

0.54

0.48

0.55

0.39

(d) STD for OpenLM-1B

64

128

256

512

1024

2048

4096

8192

16384

sequence length

1

2

4

8

16

32

64

128

256

number of sequences per batch

0.62

0.44

0.42

0.29

0.54

0.34

0.59

4.58

4.58

0.35

0.28

0.26

0.23

0.63

0.21

1.11

3.34

0.48

0.40

0.36

0.40

0.28

0.65

0.94

0.29

0.49

0.33

0.44

0.92

0.68

0.34

0.34

0.39

0.60

0.76

0.53

0.61

1.04

0.56

0.45

0.82

1.17

0.82

0.86

0.79

(e) STD for OpenLM-3B

64

128

256

512

1024

2048

4096

8192

16384

sequence length

1

2

4

8

16

32

64

128

256

number of sequences per batch

0.47

0.64

0.40

0.40

0.41

1.10

2.25

1.89

0.28

0.35

0.32

0.38

0.39

0.64

1.20

0.58

0.31

0.30

0.24

0.37

0.63

0.31

0.41

0.32

0.26

0.79

0.39

0.32

0.33

0.35

0.34

0.67

0.37

0.69

0.65

0.87

the argmax of mean log-probability is then chosen as the model's response. After we obtain the response, the computation of accuracy follows the official implementation [5].

QuALITY We follow the setup described in [5]. The dataset contains long documents with each document containing multiple-choice QA pairs. Sometimes the context for a QA pair can be longer than 8192 tokens. To account for the longer sequence length, we increase the base frequency of RoPE positional encoding from 100k to 200k without any fine-tuning. To obtain a response from the model, we follow MMLU-style prompting, where the choices are appended to the original prompt individually and the mean log-probability is computed for each choice. The choice corresponding to the argmax of mean log-probability is then chosen as the model's response. After we obtain the model output, the rest of the evaluation follows the official implementation [5].

C Additional results

C.1 Additional results for training efficiency

We enumerate model sizes (OpenLM-1B, OpenLM-3B, OpenLM-7B), the number of sequences in a batch (from 1 to 256), and sequence lengths ( 2 6 to 2 14 ) and measure the time to train 100 batches. We repeat this 5 times and report the average and standard deviation time per batch in Fig. 7. Notice that in the figure, each diagonal corresponds to a fixed b (number of tokens seen per optimization step).

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 a 2048 sequence length. Results are shown in Table 13,

demonstrating that the gap cannot be simply filled by choosing a different hyperparameter and is fundamental to the choice of pretraining sequence length.

C.3 Additional results for scaling experiments

In this section, we show additional results for the experiments presented in Section 3.5. Table 14 shows results for dataset scaling, Table 15 for model scaling, and Table 16 for experiments on an alternative dataset.

D Comparison to best-fit sequence packing

Some recent works have employed a bin packing-based strategy [17] which aims to reduce document cross-attention by minimizing unnecessary document truncation. To achieve this, they implement a known approximation algorithm called best-fit decreasing, which packs document chunks into sequences as tightly as possible. To compare with our method, we created a new dataset based on our implementation of the best-fit decreasing algorithm and trained a new model using this dataset. We present our implementation of the best-fit decreasing algorithm, the dataset we created, and the model we trained for comparison.

Given a dataset D , the input to the algorithm is a list of tokenized document chunks C = { c 1 , c 2 , . . . , c K } such that ⋃ K i =1 c i = D , where each chunk is at most context size n (e.g., 2048) in length. The output of the algorithm is a list of bins B = { b 1 , b 2 , . . . , b M } such that c i ∈ b j . As a pre-processing step, we first tokenize the documents and convert them into chunks. Truncation is applied during this step only when necessary. Next, we sort the chunks from largest to smallest and start from the first chunk to pack into bins of size n . We track the remaining capacities for each

bin while we iterate over the chunks. In each iteration, the algorithm finds the best bin that is both feasible and optimal for placing the chunk. Feasible bins are those that can accommodate the chunk, and optimal bins are those left with the minimum remaining capacity after placing the chunk. If such a bin is not found, we open a new bin and place the chunk inside. After all the chunks have been placed, we select the bins that have non-zero remaining capacities and fill them with pad tokens <PAD> .

We process the RefinedWeb [46] dataset using the aforementioned procedure and create training sequences by concatenating all chunks in a bin. Figure 3 shows that while best-fit packing results in a higher average context length compared to the baseline concat-and-chunk , it is still much lower compared to our method dataset decomposition . Furthermore, the best-fit packing method does not prevent tokens from different documents from appearing in training sequences, whereas our method does. The presence of padding tokens in best-fit packed sequences also means that some context is wasted during each optimization step.

E Training stability with VSL and curriculum

[31] presents the stability-efficiency dilemma : efficient LLM pretraining with massive data parallelism results in a large batch size and requires a high learning rate. However, such a setup can result in training instability, leading to poor generalization. They observe a correlation between training instability and long sequences, especially at the early stages of training, suggesting that training on long sequences when the model is not well-trained can be a main source of training instability.

Here, we show that dataset decomposition alleviates this problem when used with a curriculum: starting training by sampling more from short sequence buckets. We empirically demonstrate this by training an OpenLM-1B model from scratch with a high learning rate ( = 10 -2 ) and no gradient clipping, once with baseline-8k and once with DD using the "Grow-P100" curriculum. Training loss is shown in Fig. 8, demonstrating the stability of training with DD in comparison to the baseline. This suggests that our proposed method can also be beneficial for large-scale pretraining with large batches and high learning rates in terms of efficiency.

Loss

DD-Grow-P100

10

9

8

7

6

5

4

3

Step,

0

2k

4k

6k

8k

F Average sequence length vs average context length

We compute the mean of length (Fig. 3a) and context (Fig. 3c) distributions as follows. Assume a list of sequences with lengths l 1 , l 2 , . . . , l N , which are, for example, the chunk lengths in the concat-and-chunk approach or the sequence lengths in different buckets of the dataset decomposition approach. We define the average sequence length as follows:

In auto-regressive training on a sequence with length l , we apply l losses for next-token prediction on each token in parallel. Hence, for a sequence with length l , we see contexts with lengths equal to

0 , 1 , 2 , . . . , l -1 . We define the average context length, which is different from the average sequence length, as follows:

In Fig. 3a, Fig. 3c, and Table 1, we report the average sequence length and average context length for original documents, concat-and-chunk, and dataset decomposition with different mixtures.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

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 as it 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? Answer: [Yes]

Justification: We provide limitations in Section 5.

Guidelines:

The answer NA means that the paper has no 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 a 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 the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [NA]

Justification: Our work does not 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide all experimental details in Appendix B. Further, we will be releasing the code.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Our results are based on OpenLM repository ( https://github.com/ mlfoundations/open_lm ) and RefinedWeb [46] data, both publicly available. We will be releasing our (small) changes to the repository after acceptance.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not

including code, unless this is central to the contribution (e.g., for a 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/public/guides/CodeSubmissionPolicy ) 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 a subset of experiments are reproducible, they 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).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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, or 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 are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ? 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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 in Appendix A. Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 a 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 as 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 to monitor how a 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 no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We do our best to appropriately cite/acknowledge all external code/data used in this work.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Aside from code, this work does not intend to release any asset.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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.

15. 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 nor research with human subjects.

Depending on 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 to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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.

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\nOn 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).\n\nQUESTIONS:\nQuestion:\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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nThe same as the weaknesses\n\nLIMITATIONS:\nNo potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1: 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.\n\nWEAKNESSES:\n1. 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\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-15"}