動かざることバグの如し

近づきたいよ 君の理想に

per_device_train_batch_sizeとgradient_accumulation_stepsの違い

llm

機械学習モデルのトレーニングにおけるバッチサイズと勾配蓄積について

per_device_train_batch_sizeとは

per_device_train_batch_sizeとは、GPUで一度に処理するデータのサンプル数を指定するパラメータである。 例えば、per_device_train_batch_size=8と設定するとモデルは8個のサンプルを同時に処理する。

がVRAMが足りない

バッチサイズは大きいほど学習効率は高くなる。 じゃあデカけりゃいいのかっていうとそうでもなくて、GPUメモリ使用量が増えてしまうので限界がある。 この問題を解決するのがgradient_accumulation_steps

VRAMが足りない状況で無理にバッチサイズを大きくすると、OutOfMemoryエラーが発生して学習が停止してしまう。 そこでgradient_accumulation_stepsを使って、擬似的にバッチサイズを大きくすることで、少ないVRAMでも効率的に学習を進めることができる。

勾配蓄積の仕組み

勾配蓄積は、大きなバッチサイズを小さな単位に分割して処理する手法。 例えば

設定例

  • per_device_train_batch_size=4, gradient_accumulation_steps=16
    • →実効バッチサイズ: 64
  • per_device_train_batch_size=1, gradient_accumulation_steps=64
    • →実効バッチサイズ: 64

例えば、実効バッチサイズを64にしたいとする。 per_device_train_batch_size=4、gradient_accumulation_steps=16と設定した場合、 モデルはまず4つのサンプルを処理して勾配を計算する。 しかし、この時点では勾配の更新は行われず、代わりに蓄積される。 このプロセスを16回繰り返し、合計64個のサンプル分の勾配が蓄積された後、パラメータが更新される。

per_device_train_batch_size=1、gradient_accumulation_steps=64の場合も同様だ。 一度に処理するサンプル数は1つだが、64回分の勾配が蓄積された後にパラメータが更新されるので、実効バッチサイズは64になる。

最適な設定方法

  1. GPUメモリの限界までper_device_train_batch_sizeを設定する
  2. 必要な実効バッチサイズに達するようにgradient_accumulation_stepsを調整する

例

  • OK: per_device_train_batch_size=4, gradient_accumulation_steps=16
  • NG: per_device_train_batch_size=1, gradient_accumulation_steps=64

勾配蓄積のデメリットとして、トレーニング速度が低下する。per_device_train_batch_sizeで事足りるならそっちに割り振るべき

まとめ

per_device_train_batch_sizeとgradient_accumulation_stepsは、GPUメモリ使用量と学習効率のバランスをとるための重要なパラメータだ。 GPUメモリが限られている場合は、勾配蓄積を使って実質的なバッチサイズを大きくすることで、学習時間を短縮し、モデルの精度を向上させることができる。