training-nv-pytorch 25.03 - Container Compute Service

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

Main Features and Bug Fix Lists

Main Features

基礎鏡像對齊NGC 25.02。
Pytorch及系列組件升級至2.6.0.7、TE升級至2.1、accelerate升級至1.5.2，合入對應版本功能和Bug fix。
ACCL-N 升級為2.23.4.12，合入對應版本功能和Bug fix。
vLLM升級至0.8.2.dev0、ray相應升級至2.44，新增支援flash-infer 0.2.3、Transformers升級至 4.49.0+ali、flash_attn升級至2.7.2，合入對應版本功能和Bug fix。

Bugs Fix

升級vLLM版本至0.8.2.dev0，修複Illegal memory access for MoE On H20 #13693問題。

應用情境	訓練/推理
架構	pytorch
Requirements	NVIDIA Driver release >= 570
核心組件	Ubuntu 24.04 Python 3.12.7+gc Torch 2.6.0.7 CUDA 12.8.0 ACCL-N 2.23.4.12 triton 3.1.0 TransformerEngine 2.1 deepspeed 0.15.4+ali flash-attn 2.7.2 flashattn-hopper 3.0.0b1 transformers 4.49.0+ali megatron-core 0.9.0 grouped_gemm 1.1.4 accelerate 1.5.2 diffusers 0.31.0 openmim 0.3.9 mmengine 0.10.3 mmcv 2.1.0 mmdet 3.3.0 opencv-python-headless 4.10.0.84 ultralytics 8.2.74 timm 1.0.13 mmdet 3.3.0 vllm 0.8.2.dev0+g61c7a1b8.d20250325.cu128 flashinfer 0.2.3 pytorch-dynamic-profiler 0.24.11 perf 5.4.30 gdb 15.0.50 peft 0.13.2 ray 2.44.0

Assets

25.03 公網鏡像

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

VPC鏡像

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

重要

目前僅華北2（北京）地區支援使用VPC方式拉取鏡像。

說明

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

Driver Requirements

25.03 Release對齊NGC pytorch 25.02鏡像版本更新（因NGC是每月月底發布鏡像，Golden鏡像研發月只能基於上月版本的NGC），因此Golden-gpu的驅動遵循對應NGC鏡像版本的要求。該Release基於CUDA 12.8.0.38，需要NVIDIA驅動程式版本570或更高版本。但是，如果您在資料中心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驅動程式升級，這些驅動程式不向前相容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.03, 在多節點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 --action train --model_name deepspeed/llama3-8b

4. 使用建議

鏡像中的改動涉及Pytorch、DeepSpeed等庫（後續會upstream），請勿重裝。
DeepSpeed配置中的zero_optimization.stage3_prefetch_bucket_size留空或設定為auto。

Known Issues

（暫無）

Container Compute Service：training-nv-pytorch 25.03