training-nv-pytorch 25.08 - Container Compute Service

本文介紹training-nv-pytorch 25.08版本發布記錄。

Main Features and Bug Fix Lists

Main Features

升級transformers至4.53.3+ali。
升級vllm至0.10.0，ray相應升級到2.48.0。

Bugs Fix

（無）

應用情境	訓練/推理
架構	pytorch
Requirements	NVIDIA Driver release >= 575
核心組件	Ubuntu : 24.04 Python : 3.12.7+gc CUDA : 12.8 perf : 5.4.30 gdb : 15.0.50.20240403-git torch : 2.7.1.8+nv25.3 triton : 3.3.0 transformer_engine : 2.3.0+5de3e14 deepspeed : 0.16.9+ali flash_attn : 2.7.2 flashattn-hopper : 3.0.0b1 transformers : 4.53.3+ali grouped_gemm : 1.1.4 accelerate : 1.7.0+ali diffusers : 0.34.0 mmengine : 0.10.3 mmcv : 2.1.0 mmdet : 3.3.0 opencv-python-headless : 4.11.0.86 ultralytics : 8.3.96 timm : 1.0.19 vllm : 0.10.0 flashinfer-python : 0.2.5 pytorch-dynamic-profiler : 0.24.11 peft : 0.16.0 ray : 2.48.0 accl-n : 2.27.5.14 megatron-core : 0.12.1

Assets

25.08

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

VPC鏡像

acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}
{region-id}為您使用的ACS產品所在的開服地區，比如：cn-beijing、cn-wulanchabu等。
{image:tag}為實際鏡像的名稱和Tag。

說明

egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:25.08-serverless鏡像適用於ACS產品形態、靈駿多租產品形態；該鏡像不適用於靈駿單租產品形態，請勿在靈駿單租情境使用。

Driver Requirements

25.08 Release 基於CUDA 12.8.0，需要NVIDIA驅動程式版本575或更高版本。但是，如果您在資料中心GPU（例如T4或任何其他資料中心GPU）上運行，則可以使用NVIDIA驅動程式版本470.57（或更高版本R470）、525.85（或更高版本 R525）、535.86（或更高版本 R535），或545.23（或更高版本的 R545）。
CUDA驅動程式的相容性包只支援特定的驅動程式。因此，使用者應該從所有R418、R440、R450、R460、R510、R520、R530、R545 、R555和R560驅動程式升級，這些驅動程式不向前相容CUDA 12.8。有關支援的驅動程式的完整列表，請參閱CUDA 應用程式相容性主題。有關更多資訊，請參閱 CUDA 相容性和升級。

Key Features and Enhancements

PyTorch編譯最佳化

PyTorch 2.0引入的編譯最佳化能力在單卡小規模下通常可以獲得顯著的收益，但是在LLM訓練中需要引入顯存最佳化、FSDP/DeepSpeed等分布式架構，導致torch.compile()無法簡單地獲得收益或者存在負收益：

在DeepSpeed架構下控制通訊的顆粒度，協助編譯器擷取更完整的計算圖，做更大範圍的編譯最佳化
最佳化版本的PyTorch：
- 最佳化PyTorch編譯器前端，確保在計算圖中出現任意graph break的情況下仍能正常編譯
- 強化模式比對以及dynamic shape能力，提高編譯後代碼效能

結合上述最佳化，在8B LLM訓練下通常可以獲得20%左右的E2E吞吐收益。

重計算顯存最佳化

基於大量效能評測資料，包括不同模型在不同叢集以及不同訓練參數配置，以及評測過程中採集的相關顯存利用率等系統指標資料，我們進行模型顯存開銷的預測建模分析，並推薦出最佳的啟用值重算層數，並整合到PyTorch中，讓使用者可以低門檻的使用顯存最佳化帶來的效能收益。當前已支援該特性在DeepSpeed架構中的適配。

ACCL通訊庫

ACCL是阿里針對靈駿產品自研的高效能網路通訊庫，針對GPU情境提供ACCL-N版本。ACCL-N是阿里雲基於英偉達NCCL定製後提供的高效能通訊庫，在完全相容NCCL的基礎上，修複了官方NCCL版本的一些BUG，並進行了效能和穩定性相關的最佳化。

E2E效能收益評估

利用雲原生AI效能評測分析工具CNP，我們採用主流開源模型和架構配置，與標準的基礎鏡像進行了全面的端到端效能比較分析，並通過消融實驗分析，進一步評估了每個最佳化組件對整體模型訓練效能的收益貢獻。

GPU核心組件E2E效能貢獻分析

以下測試基於25.08, 在多節點GPU叢集上進行訓練E2E效能評測和對比分析，對比項包括：

Base：NGC PyTorch Image。
ACS AI Image：Base+ACCL： 鏡像使用ACCL通訊庫。
ACS AI Image：AC2+ACCL： Golden鏡像使用AC2 BaseOS，不開啟任何最佳化。
ACS AI Image：AC2+ACCL+CompilerOpt： Golden鏡像使用AC2 BaseOS，只啟用torch compile最佳化。
ACS AI Image：AC2+ACCL+CompilerOpt+CkptOpt：Golden鏡像使用AC2 BaseOS，且同時開啟torch compile和selective gradient checkpoint最佳化。

Quick Start

以下樣本內容僅通過Docker方式拉取training-nv-pytorch鏡像。

說明

在ACS中使用training-nv-pytorch鏡像需要通過控制台建立工作負載介面的製品中心頁面選取，或者通過YAML檔案指定鏡像引用。

1. 選擇鏡像

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

2. 調用API開啟編譯器+重計算顯存最佳化

啟用編譯最佳化
使用transformers Trainer API：
啟用重計算顯存最佳化
```
export CHECKPOINT_OPTIMIZATION=true
```

3. 啟動容器

鏡像中內建了模型訓練工具ljperf，以此說明啟動容器和運行訓練任務的步驟。

LLM類

# 啟動容器並進入
docker run --rm -it --ipc=host --net=host  --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]

# 運行訓練demo
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。

Known Issues

（暫無）

Container Compute Service：training-nv-pytorch 25.08