training-nv-pytorch 25.02

Main Features

• 基础镜像对齐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，合入对应版本功能和Bug fix。

Bugs Fix

（暂无）

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.02，在多节点GPU集群上进行E2E性能评测和对比分析，对比项包括：

1. Base：NGC PyTorch Image。

2. ACS AI Image：Base+ACCL：镜像使用ACCL通信库。

3. ACS AI Image：AC2+ACCL：Golden镜像使用AC2 BaseOS，不开启任何优化。

4. ACS AI Image：AC2+ACCL+CompilerOpt：Golden镜像使用AC2 BaseOS，只启用torch compile优化。

5. ACS AI Image：AC2+ACCL+CompilerOpt+CkptOpt：Golden镜像使用AC2 BaseOS，且同时开启torch compile和selective gradient checkpoint优化。

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，以此说明启动容器和运行训练任务的步骤。

# 启动容器并进入
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等库（后续会upstream），请勿重装。

• DeepSpeed配置中的zero_optimization.stage3_prefetch_bucket_size留空或设置为auto。