Container Compute Service:training-nv-pytorch 26.04

主な機能

• vllm を 0.19.0 にスペックアップしました。

• transformer_engine を 2.13 にスペックアップしました。

バグ修正

なし。

アセット

ドライバーの要件

• リリース 26.04 は、NVIDIA ドライバーのバージョンに応じて CUDA 12.8.0 および CUDA 13.0.2 をサポートします。CUDA 13.0.2 には NVIDIA ドライバー リリース 580 以降が、CUDA 12.8.0 には NVIDIA ドライバー リリース 575 以降が必要です。サポート対象ドライバーの完全な一覧については、「CUDA Application Compatibility」をご参照ください。詳細については、「CUDA Compatibility and Upgrades」をご参照ください。

PyTorch コンパイル最適化

PyTorch 2.0 で導入された torch.compile() は、シングル GPU でのトレーニングには効果的ですが、GPU メモリ最適化や Fully Sharded Data Parallel (FSDP) や DeepSpeed などの分散フレームワークに依存する大規模言語モデル (LLM) のトレーニングでは、限定的またはマイナスの効果しか得られません。

本リリースでは、以下の 2 つの最適化により、分散 LLM トレーニングにおける torch.compile() の性能を向上させました。

• DeepSpeed における通信粒度制御： 通信粒度を制御することで、コンパイラに完全な計算グラフを提供し、より広範なコンパイル最適化を可能にします。

• フロントエンドの改善： PyTorch コンパイラのフロントエンドは、グラフ中断が発生した場合でもコンパイルを実行できるようになり、拡張モードマッチングおよび動的形状機能が強化されました。

結果として、80 億パラメーター規模の LLM トレーニングにおいて、エンドツーエンドのスループットが約 20 % 向上しました。

再計算時の GPU メモリ最適化

さまざまなクラスターおよびパラメーター構成での性能テストに基づき、本リリースでは最適なアクティベーション再計算レイヤー数を PyTorch に直接統合しました。環境変数を 1 つ設定するだけで有効化でき、手動でのチューニングは不要です。

この機能は DeepSpeed フレームワークで利用可能です。

エンドツーエンド性能評価

クラウドネイティブ AI 性能ベンチマークおよび分析ツールである CNP を使用して、標準的なベースイメージに対して、主流のオープンソースモデルおよびフレームワーク構成によるエンドツーエンド性能比較を実施しました。また、各最適化コンポーネントが全体のトレーニング性能にどの程度寄与しているかを測定するためのアブレーションスタディも実施しました。

イメージ比較および反復評価

コア GPU コンポーネントのエンドツーエンド性能への寄与

以下のテストでは、マルチノード GPU クラスター上で、このイメージのエンドツーエンドトレーニング性能を複数の構成と比較しています。

1. ベース： NGC PyTorch イメージ

2. ACS AI イメージ (ベース + ACCL)： ACCL 通信ライブラリを追加したベースイメージ。

3. ACS AI イメージ (AC2 + ACCL)： AC2 BaseOS を搭載したゴールデンイメージで、最適化は有効化されていません。

4. ACS AI イメージ (AC2 + ACCL + CompilerOpt)： AC2 BaseOS を搭載したゴールデンイメージで、PyTorch コンパイル最適化のみ有効化されています。

5. ACS AI イメージ (AC2 + ACCL + CompilerOpt + CkptOpt)： AC2 BaseOS を搭載したゴールデンイメージで、PyTorch コンパイル最適化および選択的勾配チェックポイント最適化の両方が有効化されています。

ステップ 1： イメージの選択

docker pull egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]

ステップ 2： コンパイルおよび再計算最適化の有効化

• コンパイル最適化の有効化

  Transformers Trainer API を使用します：

  

• 再計算 GPU メモリ最適化の有効化

  export CHECKPOINT_OPTIMIZATION=true

ステップ 3： コンテナの起動

このイメージにはモデルトレーニングツール ljperf が含まれています。以下の手順では、このツールを使用してコンテナを起動し、トレーニングタスクを実行する方法を示します。

# コンテナを起動してシェルを開く
docker run --rm -it --ipc=host --net=host --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]

# トレーニングデモを実行
ljperf benchmark --model deepspeed/llama3-8b 

ステップ 4： 使用上の注意事項

• このイメージには PyTorch や DeepSpeed などのライブラリに対する変更が含まれています。これらを再インストールしないでください。

• DeepSpeed 構成ファイルでは、zero_optimization.stage3_prefetch_bucket_size を空のままにするか、auto に設定してください。

• 組み込みの環境変数 NCCL_SOCKET_IFNAME は、ご利用のシナリオに応じて設定する必要があります。

  • 単一の Pod がトレーニングまたは推論タスクで 1、2、4、または 8 個の GPU を使用する場合は、NCCL_SOCKET_IFNAME=eth0 に設定します。これは本イメージのデフォルト設定です。

  • 単一の Pod がホスト上のすべての 16 個の GPU をトレーニングまたは推論タスクで使用する場合は、HPN を使用するために NCCL_SOCKET_IFNAME=hpn0 に設定します。