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. However, this method of
concatenation can lead to cross-document attention within a sequence, which is
neither a desirable learning signal nor computationally efficient.
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 penalty 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 3x faster 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.

通过数据集分解、变长序列训练技术、性能增强等方法，实现了对大型语言模型的高效训练和提升。