Adaptive Batch Size Schedules for Distributed Training of Language Models with Data and Model Parallelism

• URL: http://arxiv.org/abs/2412.21124v2
• Date: Sun, 16 Mar 2025 21:10:15 GMT
• Title: Adaptive Batch Size Schedules for Distributed Training of Language Models with Data and Model Parallelism
• Authors: Tim Tsz-Kit Lau, Weijian Li, Chenwei Xu, Han Liu, Mladen Kolar,
• Abstract summary: We develop a practical implementation with PyTorch Fully Shard Data Parallel.<n>We demonstrate that our proposed approaches outperform constant batch sizes.<n>We also theoretical guarantees for such adaptive batch size schedules.
• Score: 17.006352664497122
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: An appropriate choice of batch sizes in large-scale model training is crucial, yet it involves an intrinsic yet inevitable dilemma: large-batch training improves training efficiency in terms of memory utilization, while generalization performance often deteriorates due to small amounts of gradient noise. Despite this dilemma, the common practice of choosing batch sizes in language model training often prioritizes training efficiency -- employing either constant large sizes with data parallelism or implementing batch size warmup schedules. However, such batch size schedule designs remain heuristic and often fail to adapt to training dynamics, presenting the challenge of designing adaptive batch size schedules. Given the abundance of available datasets and the data-hungry nature of language models, data parallelism has become an indispensable distributed training paradigm, enabling the use of larger batch sizes for gradient computation. However, vanilla data parallelism requires replicas of model parameters, gradients, and optimizer states at each worker, which prohibits training larger models with billions of parameters. To optimize memory usage, more advanced parallelism strategies must be employed. In this work, we propose general-purpose and theoretically principled adaptive batch size schedules compatible with data parallelism and model parallelism. We develop a practical implementation with PyTorch Fully Sharded Data Parallel, facilitating the pretraining of language models of different sizes. We empirically demonstrate that our proposed approaches outperform constant batch sizes and heuristic batch size warmup schedules in the pretraining of models in the Llama 2 family, with particular focus on smaller models with up to 3 billion parameters. We also establish theoretical convergence guarantees for such adaptive batch size schedules with Adam for general smooth nonconvex objectives.

Related papers

• Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo [22.7130140114906]
  We study the scaling law behavior of DiLoCo when training LLMs under a fixed compute budget. We find that DiLoCo scales both predictably and robustly with model size. When well-tuned, DiLoCo scales better than data-parallel training with model size, and can outperform data-parallel training even at small model sizes.
  arXiv  Detail & Related papers  (2025-03-12T20:04:38Z)
• The interplay between domain specialization and model size [8.653321928148547]
  We investigate the interplay between domain and model size during continued pretraining under compute-constrained scenarios. Our goal is to identify an optimal training regime for this scenario and detect patterns in this interplay that can be generalized across different model sizes and domains.
  arXiv  Detail & Related papers  (2025-01-03T19:28:53Z)
• Test-Time Alignment via Hypothesis Reweighting [56.71167047381817]
  Large pretrained models often struggle with underspecified tasks.<n>We propose a novel framework to address the challenge of aligning models to test-time user intent.
  arXiv  Detail & Related papers  (2024-12-11T23:02:26Z)
• More Compute Is What You Need [3.184416958830696]
  We propose a new scaling law that suggests model performance depends mostly on the amount of compute spent for transformer-based models. We predict that (a) for inference efficiency, training should prioritize smaller model sizes and larger training datasets, and (b) assuming the exhaustion of available web datasets, scaling the model size might be the only way to further improve model performance.
  arXiv  Detail & Related papers  (2024-04-30T12:05:48Z)
• Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance [55.872926690722714]
  We study the predictability of model performance regarding the mixture proportions in function forms. We propose nested use of the scaling laws of training steps, model sizes, and our data mixing law. Our method effectively optimize the training mixture of a 1B model trained for 100B tokens in RedPajama.
  arXiv  Detail & Related papers  (2024-03-25T17:14:00Z)
• Efficient Parallelization Layouts for Large-Scale Distributed Model Training [17.16249954009967]
  We conduct a comprehensive study of possible training configurations for large language models. We find that using a micro-batch size of 1 usually enables the most efficient training layouts. Our most efficient configurations enable us to achieve state-of-the-art training efficiency results over a range of model sizes.
  arXiv  Detail & Related papers  (2023-11-09T18:59:38Z)
• SWARM Parallelism: Training Large Models Can Be Surprisingly Communication-Efficient [69.61083127540776]
  Deep learning applications benefit from using large models with billions of parameters. Training these models is notoriously expensive due to the need for specialized HPC clusters. We consider alternative setups for training large models: using cheap "preemptible" instances or pooling existing resources from multiple regions.
  arXiv  Detail & Related papers  (2023-01-27T18:55:19Z)
• Does compressing activations help model parallel training? [64.59298055364336]
  We present the first empirical study on the effectiveness of compression methods for model parallelism. We implement and evaluate three common classes of compression algorithms. We evaluate these methods across more than 160 settings and 8 popular datasets.
  arXiv  Detail & Related papers  (2023-01-06T18:58:09Z)
• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models [60.23234205219347]
  TeraPipe is a high-performance token-level pipeline parallel algorithm for synchronous model-parallel training of Transformer-based language models. We show that TeraPipe can speed up the training by 5.0x for the largest GPT-3 model with 175 billion parameters on an AWS cluster.
  arXiv  Detail & Related papers  (2021-02-16T07:34:32Z)
• Parallel Training of Deep Networks with Local Updates [84.30918922367442]
  Local parallelism is a framework which parallelizes training of individual layers in deep networks by replacing global backpropagation with truncated layer-wise backpropagation. We show results in both vision and language domains across a diverse set of architectures, and find that local parallelism is particularly effective in the high-compute regime.
  arXiv  Detail & Related papers  (2020-12-07T16:38:45Z)
• Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training [12.36664837965624]
  This paper presents an approach to automatically shard the weight update across replicas. We show this technique achieves substantial speedups on typical image and language models on Cloud TPUs.
  arXiv  Detail & Related papers  (2020-04-28T07:13:50Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.