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

要約

大規模なモデルのトレーニングではバッチ サイズを適切に選択することが重要ですが、それには本質的かつ避けられないジレンマが伴います。つまり、大規模なバッチ トレーニングはメモリ使用率の点でトレーニング効率を向上させますが、汎化パフォーマンスは少量の勾配ノイズにより低下することがよくあります。
このジレンマにもかかわらず、言語モデルのトレーニングでバッチ サイズを選択する一般的な方法では、トレーニング効率を優先することが多く、データ並列処理を備えた一定の大きなサイズを採用するか、バッチ サイズのウォームアップ スケジュールを実装します。
ただし、このようなバッチ サイズのスケジュール設計は依然としてヒューリスティックであり、トレーニングのダイナミクスに適応できないことが多く、適応的なバッチ サイズのスケジュールを設計するという課題が生じています。
利用可能なデータセットが豊富であり、言語モデルがデータを大量に消費する性質があることを考慮すると、データ並列処理は不可欠な分散トレーニング パラダイムとなり、勾配計算により大きなバッチ サイズを使用できるようになります。
ただし、バニラ データの並列処理では、各ワーカーでモデル パラメーター、勾配、オプティマイザー状態のレプリカが必要となるため、数十億のパラメーターを使用して大規模なモデルをトレーニングすることができません。
メモリ使用量を最適化するには、より高度な並列処理戦略を採用する必要があります。
この研究では、データ並列処理とモデル並列処理と互換性のある、汎用的で理論的に原理に基づいた適応バッチ サイズ スケジュールを提案します。
私たちは、PyTorch Fully Sharded Data Parallel を使用した実用的な実装を開発し、さまざまなサイズの言語モデルの事前トレーニングを容易にします。
私たちは、Llama ファミリーのモデルの事前トレーニングにおいて、私たちが提案するアプローチが一定のバッチ サイズやヒューリスティックなバッチ サイズのウォームアップ スケジュールよりも優れたパフォーマンスを発揮することを経験的に示しています。特に、最大 30 億のパラメーターを持つ小規模なモデルに重点を置いています。
また、一般的な滑らかな非凸目標に対して、Adam を使用して、このような適応バッチ サイズ スケジュールに対する理論的な収束保証を確立します。

要約(オリジナル)

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 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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google