Container Compute Service:training-nv-pytorch 26.01

主な機能

• 組み込みのトレーニングコンポーネント megatron-core が 0.15.0 にスペックアップされ、推論コンポーネント vLLM が 0.13.0 に、flashinfer-python が 0.5.3 にスペックアップされました。

• health_check が shuttle 1.5.3 との互換性を持つようにスペックアップされました。

バグ修正

なし。

アセット

ドライバー要件

• 26.01 リリースでは、異なるドライバーバージョンで CUDA 12.8.0 および CUDA 13.0.2 がサポートされています。CUDA 13.0.2 には NVIDIA ドライバーバージョン 580 以降が必要です。CUDA 12.8.0 には NVIDIA ドライバーバージョン 575 以降が必要です。サポートされているドライバーの完全な一覧については、「CUDA アプリケーション互換性」トピックをご参照ください。詳細については、「CUDA 互換性とアップグレード」をご参照ください。

PyTorch コンパイル最適化

PyTorch 2.0 で導入されたコンパイル最適化機能は、1 台の GPU での小規模トレーニングに適しています。しかし、LLM トレーニングには GPU メモリ最適化と FSDP や DeepSpeed などの分散フレームワークが必要です。そのため、torch.compile() はトレーニングにメリットをもたらさず、むしろ悪影響を与える可能性があります。

• DeepSpeed フレームワークにおける通信粒度を制御することで、コンパイラがより広範囲のコンパイル最適化のために完全な計算グラフを取得できるようになります。

• 最適化された PyTorch：

  • PyTorch コンパイラのフロントエンドが最適化され、計算グラフ内でグラフ中断が発生した場合でもコンパイルが確実に行われます。

  • モードマッチングおよび動的形状機能が強化され、コンパイル済みコードが最適化されます。

上記の最適化により、8B LLM のトレーニング時の E2E スループットが 20 % 向上しました。

再計算による GPU メモリ最適化

異なるクラスタにデプロイされたモデルや異なるパラメーターで構成されたモデルに対してパフォーマンステストを実行し、GPU メモリ使用率などのシステムメトリックを収集することで、モデルの GPU メモリ消費量を予測・分析します。その結果に基づき、アクティベーション再計算レイヤーの最適数を提案し、それを PyTorch に統合しています。これにより、ユーザーは簡単に GPU メモリ最適化の恩恵を受けられます。現在、この機能は DeepSpeed フレームワークで使用可能です。

E2E パフォーマンス向上効果の評価

クラウドネイティブ AI パフォーマンス評価・分析ツール CNP を使用して、包括的な E2E パフォーマンス比較を実施しました。標準ベースイメージに対して、主流のオープンソースモデルおよびフレームワーク構成を用いて比較を行い、さらにアブレーション実験を実施して、各最適化コンポーネントがモデル全体のトレーニングパフォーマンスに与える貢献度を評価しました。

イメージ比較：ベースイメージと反復評価

コア GPU コンポーネントの E2E パフォーマンスへの貢献度分析

以下のテストはバージョン 26.01 に基づいています。マルチノード GPU クラスター上で E2E トレーニングパフォーマンスの評価および比較分析を実施しました。比較項目は以下のとおりです。

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

2. ACS AI イメージ：ベース + ACCL： ACCL 通信ライブラリを使用したイメージ

3. ACS AI イメージ：AC2 + ACCL： 最適化を有効にしていない AC2 BaseOS を使用したゴールデンイメージ

4. ACS AI イメージ：AC2 + ACCL + CompilerOpt： AC2 BaseOS を使用し、torch コンパイル最適化のみを有効にしたゴールデンイメージ

5. ACS AI イメージ：AC2 + ACCL + CompilerOpt + CkptOpt： AC2 BaseOS を使用し、torch コンパイル最適化と選択的勾配チェックポイント最適化の両方を有効にしたゴールデンイメージ

1. イメージの選択

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

2. コンパイラおよび GPU メモリ最適化のための再計算を有効にする API の呼び出し

• コンパイル最適化を有効にする

  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 枚のカードをリクエストする場合は、NCCL_SOCKET_IFNAME=eth0 に設定します。これはこのイメージのデフォルト構成です。

  • 単一の Pod がマシンの全 16 枚のカードをトレーニングまたは推論タスクでリクエストする場合（このケースでは HPN 高速ネットワークを使用できます）、NCCL_SOCKET_IFNAME=hpn0 に設定します。