Container Compute Service:training-nv-pytorch 26.03

主な特徴

• torch をバージョン 2.10 にアップグレードしました。

• vllm をバージョン 0.17.0 にアップグレードしました。

• megatron-core をバージョン 0.16.0 にアップグレードしました。

• deepspeed をバージョン 0.18.8 にアップグレードしました。

• transformer_engine をバージョン 2.12 にアップグレードしました。

バグ修正

このリリースではバグ修正はありません。

アセット

ドライバー要件

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

PyTorch コンパイルの最適化

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

このリリースでは、2 つの最適化を通じて、分散 LLM トレーニングにおける torch.compile() を改善しています：

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

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

結果：8B LLM のトレーニング時にエンドツーエンドのスループットが 20% 向上します。

再計算のための GPU メモリ最適化

さまざまなクラスターとパラメーター構成での性能テストに基づき、このリリースでは最適な活性化再計算レイヤー数を PyTorch に直接統合しています。手動でのチューニングは不要で、単一の環境変数で有効にできます。

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

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

クラウドネイティブ AI 性能分析ツール CNP を使用して、主流のオープンソースモデルとフレームワーク構成を持つ標準ベースイメージとの包括的なエンドツーエンド性能比較を実施しました。また、各最適化コンポーネントがモデルトレーニング全体の性能に与える貢献度を評価するために、アブレーションスタディも行いました。

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

コア GPU コンポーネントの E2E 性能貢献度分析

以下のテストでは、このイメージリリースを使用して、マルチノード GPU クラスターでのエンドツーエンドのトレーニング性能を評価・比較します。比較項目は以下の通りです：

1. Base：NGC PyTorch イメージ

2. ACS AI Image (Base + ACCL)：ACCL 通信ライブラリを使用したイメージ。

3. ACS AI Image (AC2 + ACCL)：最適化を有効にしていない AC2 BaseOS を使用したイメージ。

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

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

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

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

ステップ 2：コンパイラとメモリ最適化の有効化

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

  transformers Trainer API を使用します：

  

• 選択的勾配チェックポイントの有効化

  export CHECKPOINT_OPTIMIZATION=true

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

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

# コンテナを起動して中に入る
docker run --rm -it --ipc=host --net=host --privileged 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 に設定してください。