Container Compute Service:training-nv-pytorch 25.02

主な機能

• ベースイメージは NGC 25.01 に準拠し、CUDA は 12.8.0 に、cuDNN は 9.7.0.66 にアップグレードされました。

• ACCL-N が 2.23.4.11 にアップグレードされ、ACCL-Barex をサポートするようになりました。

• Transformers は 4.48.3+ali に、vLLM は 0.7.2 に、Ray は 2.42.1 にアップグレードされました。このリリースには、これらのアップストリームバージョンの機能とバグ修正が含まれます。

バグ修正

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

パブリックイメージ

• egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:25.02-serverless

VPC イメージ

• acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}

  {region-id} は、ACS がアクティブ化されているリージョン (cn-beijing、cn-wulanchabu など) を示します。

  {image:tag} は、イメージの名前とタグを示します。

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 フレームワークでサポートされるようになりました。

ACCL

ACCL は、Lingjun 向けに構築された Alibaba Cloud の高性能な通信ライブラリです。ACCL-N は GPU に特化したバージョンです。ACCL-N は、NVIDIA NCCL をカスタマイズした高性能な通信ライブラリです。NCCL と完全に互換性があり、アップストリームの NCCL リリースの問題を修正し、パフォーマンスと安定性を向上させています。

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

クラウドネイティブな AI パフォーマンスのベンチマークおよび分析ツールである CNP を使用し、主流のオープンソースモデルとフレームワーク構成を備えた標準のベースイメージとの、完全なエンドツーエンドのパフォーマンス比較を実施しました。また、各最適化コンポーネントがモデルトレーニング全体のパフォーマンスにどれだけ貢献するかを測定するために、アブレーションスタディも実施しました。

GPU 主要コンポーネントのエンドツーエンドパフォーマンス貢献度分析

以下のテストは 25.02 に基づいており、マルチノード GPU クラスターでエンドツーエンドのパフォーマンスベンチマークと比較を実行します。比較項目は次のとおりです。

1. Base： NGC PyTorch イメージ。

2. ACS AI Image: Base+ACCL： ACCL 通信ライブラリを備えたベースイメージ。

3. ACS AI Image: AC2+ACCL： AC2 BaseOS 上で ACCL を備えたベースイメージ。

4. ACS AI Image: AC2+ACCL+CompilerOpt： AC2 BaseOS 上で ACCL と torch.compile() の最適化を有効にしたベースイメージ。

5. ACS AI Image: AC2+ACCL+CompilerOpt+CkptOpt： AC2 BaseOS 上で ACCL、torch.compile()、および選択的勾配チェックポイントの最適化を有効にしたベースイメージ。

1. イメージのプル

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

2. コンパイラとメモリ最適化の有効化

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

  Transformers Trainer API を使用します：

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

• 再計算のための 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 に設定してください。