Scaling AI Training Pipelines Unlocks Performance Gains with Massive Datasets and Billions Invested

Large language models now excel at a diverse range of natural language processing tasks, but training these powerful systems demands enormous computational resources, with leading artificial intelligence companies investing heavily in supercomputing infrastructure. Alexander Interrante-Grant, Carla Varela-Rosa, Suhaas Narayan, and colleagues at the Massachusetts Institute of Technology address a critical gap in understanding how to efficiently scale the pretraining of these models. Their work investigates the practical challenges of training on massive datasets, distributed across hundreds of computing nodes, and offers valuable insights into maximizing the use of available computing power. By demystifying the pretraining pipeline, this research provides crucial guidance for optimising performance and reducing the costs associated with developing ever-larger language models.

Unfortunately, detailed information about the scaling performance and training considerations of these large training pipelines remains scarce in public literature. This work aims to demystify the large language model pretraining pipeline, particularly with respect to distributed training.

Large Language Model Scaling and Performance Analysis

Scaling up large language models demands significant computational resources, as demonstrated by recent advancements from companies like OpenAI, Google, Anthropic, xAI, Mistral AI, and DeepSeek. However, details regarding model architectures and training processes are largely undisclosed, leaving researchers to independently address scaling challenges. This work details the pretraining of a language model for a novel application, examining scaling performance with increasing dataset and model sizes, and reports key lessons learned. The experiments were conducted on a computing cluster at the Lincoln Laboratory Supercomputing Center, utilizing a system comprising 316 compute nodes, each featuring dual AMD EPYC 9254 24-core CPUs with 768 GB of DRAM, and dual Nvidia Hopper H100-NVL GPUs with 94 GB of HBM GPU RAM, interconnected by an NV-Link bridge.

Nodes mount a central Lustre parallel storage array, have 3. 8 TB of local SSD storage, and connect via 25-Gigabit Converged Ethernet. The team pretrained a Bidirectional Encoder Representation for Transformers (BERT)-like encoder model on a large dataset of binary code, comprising 202 million pretraining samples compiled from open-source projects using the nixpkgs package manager, totaling just under 2 TB. Initial scaling efforts faced challenges sharing this large dataset across hundreds of nodes. To address this, the team recommends preprocessing and tokenizing the entire dataset before training, storing only the necessary inputs and attention masks, reducing the dataset size to 25 GB, a 99% reduction.

Duplicating the dataset across nodes prior to training further mitigated network bottlenecks. Further optimization involved parallelizing data loading, carefully balancing the number of parallel data loaders to maintain near 100% GPU utilization. Utilizing PyTorch Lightning, the team scaled training to multiple GPUs and nodes, discovering that network bandwidth was not a significant bottleneck, with training performance scaling roughly linearly with the number of nodes, up to 128. However, increasing model size while maintaining a constant node count led to decreased training performance, as larger models require more GPU memory, reducing the training batch size. Scaling beyond this point would necessitate model parallelism and further tuning. These findings provide practical recommendations for optimizing large language model pretraining performance, offering insights for researchers attempting to scale custom models and datasets.

Scaling Data and Compute for Large Models

Large language models currently deliver best-in-class performance across numerous natural language processing applications, yet training these complex systems presents significant computational challenges. Researchers are investing billions of dollars in supercomputing infrastructure to train progressively larger models on increasingly massive datasets, but detailed information regarding the scaling performance and practical considerations of these training pipelines remains scarce. This work addresses this gap by demystifying the pretraining process, with a particular focus on distributed training, efficient dataset management across hundreds of nodes, and maximizing the utilization of available GPU compute capacity. The team’s investigations reveal critical insights into scaling data parallelism, a technique essential for distributing the workload across multiple processors.

The research demonstrates the importance of fully leveraging available GPU resources during training, a factor often overlooked in previous studies. By carefully analysing the training pipeline, scientists identified bottlenecks and optimised data handling procedures to improve efficiency. This optimisation allows for more effective scaling of data parallelism, enabling the training of larger models with reduced computational cost. The findings show that careful attention to these details is crucial for achieving optimal performance when working with extremely large datasets and distributed computing environments. Ultimately, this work provides practical recommendations for tuning training performance when scaling up, offering valuable guidance to researchers and engineers working in the field of artificial intelligence.

This work provides practical recommendations for optimising the pretraining of large language models, gained through the process of training models at scale. The research demonstrates that maintaining consistent computational resources across training runs does not necessarily guarantee optimal performance, and that increasing model size can indirectly reduce training efficiency due to limitations in batch size. Specifically, larger models require more GPU memory, which constrains the number of samples processed in each batch. The findings highlight the importance of careful consideration of these factors when scaling up language model pretraining. The authors acknowledge that further scaling beyond the models tested would likely necessitate model parallelism, requiring additional tuning to maximise performance. This research offers valuable insights for other researchers attempting to train custom large language models and contribute to a better understanding of the challenges associated with scaling these complex systems.