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 にアップグレードしました。

バグ修正

なし。

コアコンポーネント

パブリックイメージ

VPC イメージ

ACS コンソールの YAML ファイルで指定されている AI コンテナイメージアセット URI egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/{image:tag} を acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag} に置き換えることで、VPC 経由で AI コンテナイメージを迅速にプルできます。

• ここで、{region-id} は、ACS が有効化されている利用可能なリージョン (cn-beijing や cn-wulanchabu など) です。

• {image:tag} は、イメージの名前とタグです。

ステップ 1：イメージのプル

ホストシェルで次のコマンドを実行します。

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

[tag] を、お使いの NVIDIA ドライバーのバージョンと CPU アーキテクチャに一致するタグ (26.03-cu130-serverless または 26.03-cu128-serverless のいずれか) に置き換えてください。各タグがサポートするドライバーのバージョンとアーキテクチャについては、「イメージの内容」セクションをご参照ください。

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

ホストシェルで次のコマンドを実行します。このコマンドはコンテナを起動し、コンテナシェルを開きます。残りの手順はこのシェルで実行します。

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

ステップ 3：最適化の有効化

トレーニングを開始する前に、両方の最適化を有効にします。トレーニングスクリプトでコンパイル最適化を設定し、ステップ 2 で開いたコンテナシェルで再計算によるメモリ最適化を設定します。

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

export CHECKPOINT_OPTIMIZATION=true

ステップ 4：トレーニングデモの実行

このイメージには、組み込みのモデルトレーニングツールである ljperf が含まれています。コンテナシェルでデモを実行します。

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

使用上の推奨事項

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

• DeepSpeed の設定で、zero_optimization.stage3_prefetch_bucket_size は空白のままにするか、auto に設定してください。

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