Container Compute Service:training-nv-pytorch 26.05

主な機能

• Torch 関連コンポーネントをバージョン 2.11.0 にスペックアップしました。

• 組み込みのトレーニングコンポーネントである megatron-core を 0.17.0 に、transformer_engine を 2.14 にスペックアップしました。

• 推論コンポーネント vllm をバージョン 0.20.2 にスペックアップしました。

バグ修正

このリリースでは該当なし。

アセット

ドライバー要件

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

PyTorch コンパイル最適化

PyTorch 2.0 で導入された torch.compile() は、小規模な単一 GPU ワークロードにおいて多くの場合、高い効果を発揮します。しかし、LLM トレーニングは GPU メモリ最適化や FSDP や DeepSpeed などの分散フレームワークに依存するため、torch.compile() による効果が限定的であったり、逆にパフォーマンスが低下したりする可能性があります。

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

• 最適化された PyTorch ビルドを使用します：

  • PyTorch コンパイラフロントエンドが改善され、計算グラフ内でグラフ中断が発生してもコンパイルが成功するようになります。

  • パターンマッチングおよび動的形状のサポートが強化され、コンパイル後のパフォーマンスが向上します。

これらの最適化により、8B パラメーターの LLM トレーニングでは、通常、エンドツーエンドのスループットが約 20 % 向上します。

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

さまざまなモデル、クラスター、トレーニングパラメータ設定、およびベンチマーク中に収集された GPU メモリ使用率などのシステムメトリックを含む大規模なパフォーマンスデータに基づいた GPU メモリオーバーヘッドの予測モデルにより、アクティベーション再計算レイヤーの最適数が推奨されます。このアプローチは PyTorch に統合されており、最小限の手間で GPU メモリ最適化によるパフォーマンス向上を実現できます。この機能は現在、DeepSpeed フレームワークでサポートされています。

エンドツーエンドパフォーマンス評価

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

イメージ vs. ベースイメージおよび反復評価

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 を使用し、torch.compile 最適化のみを有効にしたゴールデンイメージ。

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

1. イメージの選択

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

2. API によるコンパイラおよびチェックポイントの有効化

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

  transformers Trainer API を使用します：

  training_args = TrainingArguments(
      bf16=True,
      gradient_checkpointing=True,
      torch_compile=True
  )

• 勾配チェックポイントメモリ最適化の有効化

  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 をトレーニングまたは推論タスクのためにリクエストする場合は、NCCL_SOCKET_IFNAME=hpn0 を設定して、パフォーマンス専有型ネットワーク (HPN) を使用してください。