Refining Packing and Shuffling Strategies for Enhanced Performance in Generative Language Models
Yanbing Chen, Ruilin Wang, Zihao Yang, Lavender Yao Jiang, Eric Karl Oermann
TL;DR
This work analyzes how data packing strategies and the choice of atom size affect language model training. By pretraining GPT-2 124M on WikiText with two packing methods (concat and padding) across multiple atom sizes and $MSL$ values, and keeping a fixed parameter count via Alibi, the authors evaluate final perplexity, perplexity ranking, and efficiency. They find that matching the atom size to $MSL$ optimizes performance for both packing methods, with padding delivering lower final perplexity but requiring more training steps and lower efficiency. The results provide practical guidance for selecting packing strategies based on data availability and training time, highlighting a performance-efficiency trade-off that researchers can leverage in LM pretraining.
Abstract
Packing and shuffling tokens is a common practice in training auto-regressive language models (LMs) to prevent overfitting and improve efficiency. Typically documents are concatenated to chunks of maximum sequence length (MSL) and then shuffled. However setting the atom size, the length for each data chunk accompanied by random shuffling, to MSL may lead to contextual incoherence due to tokens from different documents being packed into the same chunk. An alternative approach is to utilize padding, another common data packing strategy, to avoid contextual incoherence by only including one document in each shuffled chunk. To optimize both packing strategies (concatenation vs padding), we investigated the optimal atom size for shuffling and compared their performance and efficiency. We found that matching atom size to MSL optimizes performance for both packing methods (concatenation and padding), and padding yields lower final perplexity (higher performance) than concatenation at the cost of more training steps and lower compute efficiency. This trade-off informs the choice of packing methods in training language models.
