メインコンテンツまでスキップ
バージョン: v2506

分散学習フレームワーク向けのチューニング

分散学習において広く使用されるフレームワーク向けのチューニング例について説明します。

DeepSpeedとMegatron-LMのサンプルコードとしてMDK: Model Development Kitを使用しているため、以下に示すようにfaibと同階層にクローンしてください。

$ ls
faib mdk

DeepSpeed​

CommandOutputTunerを用いて、ハイパーパラメータに依存してDeepSpeed configutationファイルを書き換え、deepspeedコマンドにより分散学習を実行し、標準出力内のtrain_runtimeを最小化します。

この例では、MDKのstep-by-step.mdを完了している事を前提としています。

# MDKディレクトリに移動
cd mdk/

# Python仮想環境の作成
python3 -m venv .venv
. .venv/bin/activate

# ZenithTuneのインストール
pip install ../faib/intelligence/zenith-tune

# DeepSpeedチューニングスクリプトをカレントディレクトリにコピー
cp ../faib/intelligence/zenith-tune/examples/deepspeed/* .

# 単一プロセスでチューニングを実行
python tune_deepspeed.py

Megatron-LM​

CommandOutputTunerを用いて、Megatron-LMの学習スクリプトの値を書き換えながら、標準出力内の100イテレーションのスループットを抽出。 初回10イテレーションのウォームアップを除いて調和平均を取った値を最適化します。 この例では、単一プロセスのコマンド文字列を生成し、チューニングスクリプト自体をmpirunで並列実行します。 なお、MDKのpretrain-megatron-lm.mdを完了している事を前提としています。

# MDKのMegatron-LM作業用ディレクトリに移動
cd mdk/outputs/Megatron-LM/

# ZenithTuneをSingularityコンテナ内にインストール
singularity exec --nv --bind ../../../faib/intelligence/zenith-tune:/zenith-tune ./pytorch.sif pip install /zenith-tune

# チューニングスクリプトをカレントディレクトリにコピー
cp ../../../faib/intelligence/zenith-tune/examples/megatronlm/* .

# mpirunによって複数プロセスでチューニングを実行
mpirun --np 8 singularity exec --nv ./pytorch.sif python tune_megatronlm.py

Megatron-LMの高速化関連オプションを参考にその他のパラメータもチューニングすることで更なる高速化を期待できます。

優先して検討すべきオプション​

特に以下のパラメーターについては、性能向上を目指したい時に変更を検討してみてください。

  • tensor-model-parallel-size：テンソル並列の並列数
  • pipeline-model-parallel-size：パイプライン並列の並列数
  • global-batch-size：学習の1イテレーションで消費されるサンプル数
  • micro-batch-size：1GPUが消費するサンプル数
  • seq-length：1サンプルで学習する文字列長

これらのパラメータは、システム性能と学習の安定性に大きな影響を与えるため、以下の点に注意して設定してください。

  • 必須要件

    • 合計GPU数（＝ノード数×8）mod（tensor-model-parallel-size×pipeline-model-parallel-size）= 0
    • global-batch-size mod（micro-batch-size×合計GPU数/tensor-model-parallel-size/pipeline-model-parallel-size）= 0

  • 調整項目

    1. メモリ使用量の許す範囲で、tensor-model-parallel-size×pipeline-model-parallel-sizeを小さくする
    2. メモリ使用量の許す範囲で、micro-batch-size×seq-lengthを大きくする
    3. 学習精度が悪くならない範囲で、global-batch-sizeを大きくする

また、以下のオプションの利用も性能向上の可能性があります。

  • --overlap-param-gather：パラメータをGPU間で交換する通信を計算と重ねることで、GPUの遊休時間を削減する。

変更を推奨しないオプション​

少なくとも以下のオプションを無効化すると、処理性能が劣化する可能性が高いため、Megatron-LMを用いた分散並列事前学習の実行方法を説明しているハウツーの既定値から変更しないことを推奨します。

  • --use-distributed-optimizer：各GPUで分割して計算を行うパラメータに対応して、optimizerのパラメータも分散して保存させます。これにより、計算の効率性が向上します。特に、bf16フォーマットと併用することで、大きな高速化が期待できます。詳しくは参考文献も参照してください。
  • --overlap-grad-reduce：勾配の集約通信を計算とオーバーラップさせます。これにより、通信時間が計算に隠蔽されることで、実効性能が向上します。詳しくは参考文献も参照してください。
  • --sequence-parallel：シーケンス方向に並列化することで、GPUの計算リソースを効率的に活用させます。特に、テンソル並列が適用できないlayer norm, attention, dropoutレイヤを並列化することができます。詳しくはMegatron-LMの論文も参照してください。

MMEngine​

CommandOutputTunerを用いて、学習スクリプトのオプション--cfg-optionsにハイパーパラメータを設定し、標準出力内の100イテレーションの実行時間を最適化します。

# MMEngine examplesディレクトリに移動
cd faib/intelligence/zenith-tune/examples/mmengine/

# Python仮想環境の作成
python3 -m venv .venv
. .venv/bin/activate
pip install torch torchvision mmengine

# ZenithTuneのインストール
pip install ../../

# CIFAR-10データセットのダウンロード
mkdir -p data/cifar10
pushd data/cifar10/
wget https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
tar -xzf cifar-10-python.tar.gz
popd

# torchrunによって複数プロセスでチューニングを実行
torchrun --nproc-per-node 4 tune_mmengine.py