AI・機械学習
上級

マイクロバッチング(マイクロバッチング)

GPU メモリに収まらない大きなバッチを小さなサブバッチ(マイクロバッチ)に分割し、逐次処理する手法。パイプライン並列やモデル並列と組み合わせて GPU の稼働率を最大化しつつ、論理的には大バッチと等価な学習を実現する。LLM の学習基盤である Megatron-LM や DeepSpeed の中核機能として実装されている。

0 回閲覧
0 いいね

マイクロバッチングとは

マイクロバッチング(Micro-Batching)は、1 回のパラメータ更新に使用するグローバルバッチを、GPU メモリに収まるサイズの小さなサブバッチ(マイクロバッチ)に分割して逐次的にフォワード・バックワードパスを実行する手法である。勾配累積と密接に関連するが、マイクロバッチングは特にパイプライン並列との組み合わせにおいて「パイプラインバブル」を最小化する目的で体系化された概念である。

LLM の学習では、70B パラメータモデルの場合パラメータだけで FP16 で 140GB、オプティマイザ状態を含めると 560GB のメモリが必要となる。H100 80GB の単一 GPU にはパラメータすら載らないため、モデル並列やテンソル並列でモデルを分割した上で、各 GPU が処理するデータ量をマイクロバッチで制御する。

マイクロバッチとグローバルバッチの関係

マイクロバッチングにおけるバッチサイズの階層構造は以下の通りである:

  • マイクロバッチサイズ(mbs): 1 回のフォワードパスで 1 つの GPU が処理するサンプル数。GPU メモリに収まる最大値に設定
  • ミニバッチサイズ: 1 つの GPU が勾配累積ステップ全体で処理するサンプル数。mbs × 累積ステップ数
  • グローバルバッチサイズ(gbs): 全 GPU が 1 回のパラメータ更新で処理する合計サンプル数。mbs × 累積ステップ数 × データ並列度

具体例として、マイクロバッチサイズ 2、勾配累積 8 ステップ、データ並列 256 GPU の場合:

gbs = 2 × 8 × 256 = 4,096 サンプル

シーケンス長 4,096 トークンなら、グローバルバッチサイズは約 1,680 万トークンとなる。

パイプライン並列との統合

マイクロバッチングが最も威力を発揮するのはパイプライン並列(Pipeline Parallelism)との組み合わせである。パイプライン並列ではモデルをレイヤーグループ(ステージ)に分割し、異なる GPU に配置する。マイクロバッチがなければ、ステージ 1 の処理完了を待ってステージ 2 が動くため、大半の GPU がアイドル状態(パイプラインバブル)になる。

マイクロバッチングにより、ステージ 1 がマイクロバッチ 1 を処理し終えた時点でマイクロバッチ 2 の処理を開始し、同時にステージ 2 がマイクロバッチ 1 を処理できる。マイクロバッチ数が多いほどパイプラインの充填率が高まり、バブルの割合が減少する。

マイクロバッチ数パイプラインバブル率GPU 稼働率適用場面
1最大(ステージ数依存)低(25〜50%)テスト用途のみ
4中程度中(60〜75%)小規模実験
8小高(80〜90%)標準的な学習
16 以上極小非常に高(90%+)大規模事前学習
32 以上ほぼゼロ理論上限に近いフロンティアモデル

NVIDIA の Megatron-LM では、パイプラインステージ数の 2〜4 倍のマイクロバッチ数を推奨している。

スケジューリング方式

マイクロバッチのスケジューリング(どの順番でフォワード・バックワードを実行するか)には複数の方式がある:

GPipe 方式:すべてのマイクロバッチのフォワードパスを先に完了し、その後すべてのバックワードパスを実行する。実装がシンプルだが、すべてのマイクロバッチの活性化値を同時にメモリに保持するため、メモリ消費が大きい。

1F1B(One Forward One Backward)方式:フォワードとバックワードを交互に実行する。パイプラインが充填された後は、1 マイクロバッチのフォワードと 1 マイクロバッチのバックワードを交互に処理するため、同時にメモリに保持する活性化値の数がパイプラインステージ数に制限される。Megatron-LM のデフォルト方式。

インターリーブド 1F1B 方式:各 GPU に複数の非連続なステージを割り当て、パイプラインバブルをさらに削減する。メモリ消費は 1F1B と同等だが、通信量が増加する。

スケジューリング方式メモリ消費パイプラインバブル通信量実装フレームワーク
GPipe高(全活性化保持)中低GPipe, FairScale
1F1B中(ステージ数分)中低Megatron-LM
インターリーブド 1F1B中小中Megatron-LM v3+
Zero Bubble低ほぼゼロ高研究段階

マイクロバッチサイズの最適化

マイクロバッチサイズの決定は以下の要因を考慮して行う:

GPU メモリ制約:パラメータ(モデル並列後の分割サイズ)、オプティマイザ状態(AdamW なら fp32 で 12 bytes/param)、活性化値(バッチサイズとシーケンス長に比例)、一時バッファを合計し、GPU メモリ容量(H100 で 80GB)に収まる最大の mbs を算出する。活性化チェックポインティングを併用すれば活性化メモリを大幅に削減できる。

計算効率:GPU の Tensor Core は行列演算の次元が 8 の倍数(FP16)または 16 の倍数(INT8/FP8)のとき最高効率を発揮する。マイクロバッチサイズもこれに合わせて 2 のべき乗や 8 の倍数にすることで、TFLOPS を最大化できる。

メモリの断片化:PyTorch のメモリアロケータは動的にメモリを確保・解放するため、マイクロバッチサイズが大きいほど大きな連続メモリ領域が必要になり、断片化の影響を受けやすい。torch.cuda.memory_stats() で断片化率を監視し、必要に応じて max_split_size_mb を調整する。

実装例とフレームワーク対応

主要フレームワークのマイクロバッチング対応状況:

  • Megatron-LM:--micro-batch-size と --global-batch-size パラメータで直接制御。パイプライン並列との連携が最も成熟
  • DeepSpeed:train_micro_batch_size_per_gpu で設定。ZeRO Stage 3 との組み合わせで自動的にメモリ最適化
  • PyTorch FSDP:auto_wrap_policy でシャーディング単位を設定し、勾配累積と組み合わせてマイクロバッチングを実現
  • Colossal-AI:Booster クラスの microbatch_size パラメータで設定。自動並列化機能により最適なマイクロバッチサイズを探索可能

FAQ

Q: マイクロバッチサイズ 1 でも学習できますか?

A: 技術的には可能だが、GPU の Tensor Core 利用効率が極端に低下し、理論 TFLOPS の 10〜20% 程度しか達成できないことが多い。行列積の一方の次元がバッチサイズに対応するため、mbs=1 では GEMM(General Matrix Multiply)の並列度が不足する。最低でも mbs=2、理想的には mbs=4 以上を推奨。ただし、メモリ制約が厳しい 405B クラスのモデルではパラメータだけで GPU メモリの大部分を消費するため、mbs=1 で勾配累積を多用するケースもある。

Q: マイクロバッチングと勾配累積の違いは何ですか?

A: 概念的に重複する部分が大きいが、文脈が異なる。勾配累積は「メモリに収まらないバッチサイズを分割して論理的に等価な学習を行う」汎用テクニックで、単一 GPU でも使える。マイクロバッチングは「パイプライン並列やモデル並列の文脈で、パイプラインバブルを最小化しつつバッチを分割処理する」手法であり、分散学習のスケジューリング最適化に焦点がある。結果として両方とも「小バッチの勾配を蓄積して大バッチ更新」を行うが、マイクロバッチングはパイプラインの充填率やステージ間通信まで考慮した包括的な概念である。

Q: マイクロバッチサイズを動的に変更できますか?

A: 一般的には学習中に固定するが、研究レベルでは動的調整の試みがある。たとえば、シーケンス長が可変(パッキング手法)の場合、短いシーケンスが集中するバッチではマイクロバッチサイズを増やし、長いシーケンスが集中するバッチでは減らすことで GPU メモリの利用効率を安定化できる。ただし、パイプライン並列ではマイクロバッチサイズの変動がスケジューリングを複雑化させるため、実運用では固定サイズが推奨される。