training-xpu-pytorch 25.04
本文介绍training-xpu-pytorch 25.04版本发布记录。
Main Features and Bug Fix Lists
Main Features
基础镜像升级PPU SDK至1.5.0。
ACCL-P升级至v1.5.0-9、EIC升级至1.3.7.1-v2504,合入对应版本功能和Bug fix。
Pytorch及系列组件升级至2.6.0.7、Triton适配升级到3.2.0、accelerate升级至1.6.0+ali;
vLLM升级至v0.7.3+ppu1.5.0.oe、支持flashMLA、flashinfer_python升级至0.2.2.post1+ppu1.5.0.oe、Transformers升级至 4.51.2+ali,合入对应版本功能和Bug fix;
accelerate升级至1.6.0、datasets升级至3.1.0、bitsandbytes升级至0.44.1+ppu1.5.0.ce,合入对应版本功能和Bug fix。
Bugs Fix
(暂无)
Contents
镜像Tag | 25.04 |
应用场景 | 训练/推理 |
框架 | pytorch |
Requirements | PPU SDK 1.5.0 |
核心组件 |
|
Assets
通过公网拉取ACS AI容器镜像需要先获取鉴权密钥。建议您使用VPC方式加速拉取AI容器镜像,减少镜像拉取的时间。
25.04
公网镜像
egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-xpu-pytorch:25.04
VPC镜像
acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}
{region-id}为您使用的ACS产品所在的开服地域,比如:cn-beijing、cn-wulanchabu等。{image:tag}为实际镜像的名称和Tag。
Driver Requirements
Driver version >= 1.1.0
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、PPU和AMD三个场景提供ACCL-N、ACCL-P和ACCL-R三个版本。ACCL-N是阿里云基于英伟达NCCL定制后提供的高性能通信库,在完全兼容NCCL的基础上,修复了官方NCCL版本的一些BUG,并进行了性能和稳定性相关的优化。ACCL-P是基于平头哥开源通信库pccl,进行二次开发的集合通信库。ACCL-R是基于AMD Rocm开源通信库rccl,进行二次开发的集合通信库。本版本主要将ACCL/ACCL-N上实现的主要特性移植到了pccl上,修复了一些问题,结合高网的相关自研产品组件进行了深度定制。
E2E性能收益评估
利用云原生AI性能评测分析工具CNP,我们采用主流开源模型和框架配置,与标准的基础镜像进行了全面的端到端性能比较分析,并通过消融实验分析,进一步评估了每个优化组件对整体模型训练性能的收益贡献。
镜像对比基础镜像&迭代评估

PPU核心组件E2E性能贡献分析
以下测试基于Golden-25.04, 在多节点PPU/GPU集群上进行E2E性能评测和对比分析,对比项包括:
Base:PPU 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-xpu-pytorch镜像。
在ACS中使用training-xpu-pytorch镜像需要通过控制台创建工作负载界面的制品中心页面选取,或者通过YAML文件指定镜像引用。
1. 选择镜像
通过公网拉取ACS AI容器镜像需要先获取鉴权密钥。建议您使用VPC方式加速拉取AI容器镜像,减少镜像拉取的时间。
docker login egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com2. 调用API开启编译器+重计算显存优化
启用编译优化
使用transformers Trainer API:

启用重计算显存优化
export CHECKPOINT_OPTIMIZATION=true
3. 启动容器
镜像中内置了模型训练工具ljperf,以此说明启动容器和运行训练任务的步骤。
ACS形态产品请通过YAML方式使用镜像。
LLM类
# 启动容器并进入
docker run --rm -it --ipc=host --net=host --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-xpu-pytorch:25.04
# 运行训练demo
ljperf --action train --model_name deepspeed/llama3-8bKnown Issues
VLLM
VLLM量化问题
VLLM默认使用Marlin Kernel对GPTQ/AWQ/W8A8/WOQ进行量化推理加速,可以通过显示指定quantization的方法来使用vllm原本的Marlin/W8A8/awq/gptq等量化推理kernel,例如可以指定使用gptq_acext/awq_acext来做量化推理,但是acext目前不支持
act_order=True的量化,以及GPTQ-Int8量化权重,下个版本优化。VLLM FP8量化依赖SM89/90架构,没有进行PPU适配与优化,支持计划待定。
VLLM量化Dense A8W8对接acext库实现,性能需要持续优化,MOE W8A8量化暂时不支持,预计下一版支持。
支持DPSK FP8权重load并按照W8A16量化推理,性能与BF16相当。
VLLM默认使用vllm-flash-attn作为attention的后端,该性能在PPU上尚未超过GPU版本,下个版本优化。
VLLM稳定性问题
VLLM多卡推理时,可能会出现使用mp_distributed_executor错误的情况,需使用
VLLM_WORKER_MULTIPROC_METHOD=spawn避免该问题,参考开源社区[Bug]: When tensor_parallel_size>1, RuntimeError: Cannot re-initialize CUDA in forked subprocess. #6152。VLLM多卡推理结束后会随机出现进程退出错误的社区已知问题,详见开源社区[Bug]: When using tp for inference, an error occurs: Worker VllmWorkerProcess pid 3283517 died, exit code: -15. #6145。
flashinfer
兼容性Unit Test测试Fail Case分类:
与GPU版本相比性能需要继续优化。
部分UT存在精度问题。
PPU 升级pytorch2.6,针对LLM类模型重计算显存优化性能收益效果不及历史镜像、持续优化中;
PPU兼容的cudnn目前不支持fma功能,因此PPU兼容的TE无法走FusedAttention (cudnn) ,可以优先使用FlashAttention:export NVTE_FLASH_ATTN=1 && export NVTE_FUSED_ATTN=0;