training-xpu-pytorch 25.03

更新时间:
复制为 MD 格式

本文介绍training-xpu-pytorch 25.03版本发布记录。

Main Features and Bug Fix Lists

Main Features

  • 基础镜像升级PPU SDK1.4.3_hotfix

  • ACCL-P升级至v1.4.3-8、EIC升级至1.3.7.1-v2503,合入对应版本功能和Bug fix。

  • vLLM升级至0.7.3、ray相应升级至2.43,支持flashinfer 0.2.1.post1、accelerate升级至1.5.2、Transformers升级至4.49.0+ali,合入对应版本功能和Bug fix。

Bugs Fix

(暂无)

Contents

应用场景

训练/推理

框架

pytorch

Requirements

PPU SDK 1.4.3-hotfix

核心组件

  • Ubuntu 22.04

  • Python 3.10.13+gc

  • Torch 2.5.1.6.post4+ppu

  • CUDA 12.6

  • ACCL-P 1.4.3-8

  • triton 3.1.0+ppu

  • TransformerEngine 1.13+ppu1.4.2.oe

  • deepspeed 0.15.4+ali

  • flash-attn 2.5.7+ppu1.4.2.oe

  • xformers 0.0.29.post1+ppu1.4.2.oe

  • 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

  • bitsandbytes 0.43.1

  • mmengine 0.10.3

  • mmcv 2.1.0

  • MMsegmentation 1.2.2

  • mmdet3d 1.4.0

  • AutoGPTQ 0.8.0.dev0+cu126

  • opencv-python-headless 4.10.0.84

  • vllm 0.7.3+ppu1.4.2.oe

  • pytorch-dynamic-profiler 0.24.11

  • perf 5.4.30

  • gdb 12.1

  • peft 0.13.2

  • ray 2.43

  • spacy 3.7.4

  • cumm 0.5.3

  • spconv 2.3.6

  • ffmpeg 7.0.1

  • open3d 0.18.0

  • flashinfer 0.2.1.post1

Assets

说明

通过公网拉取ACS AI容器镜像需要先获取鉴权密钥。建议您使用VPC方式加速拉取AI容器镜像,减少镜像拉取的时间。

25.03

公网镜像

  • egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-xpu-pytorch:25.03

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、PPUAMD三个场景提供ACCL-N、ACCL-PACCL-R三个版本。ACCL-N是阿里云基于英伟达NCCL定制后提供的高性能通信库,在完全兼容NCCL的基础上,修复了官方NCCL版本的一些BUG,并进行了性能和稳定性相关的优化。ACCL-P是基于平头哥开源通信库pccl,进行二次开发的集合通信库。ACCL-R是基于AMD Rocm开源通信库rccl,进行二次开发的集合通信库。本版本主要将ACCL/ACCL-N上实现的主要特性移植到了pccl上,修复了一些问题,结合高网的相关自研产品组件进行了深度定制。

E2E性能收益评估

利用云原生AI性能评测分析工具CNP,我们采用主流开源模型和框架配置,与标准的基础镜像进行了全面的端到端性能比较分析,并通过消融实验分析,进一步评估了每个优化组件对整体模型训练性能的收益贡献。

镜像对比基础镜像&迭代评估

image.png

GPU核心组件E2E性能贡献分析

以下测试基于25.03, 在多节点PPU集群上进行E2E性能评测和对比分析,对比项包括:

  1. Base:PPU 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 compileselective gradient checkpoint优化

image.png

Quick Start

以下示例内容仅通过Docker方式拉取training-xpu-pytorch镜像。

说明

ACS中使用training-xpu-pytorch镜像需要通过控制台创建工作负载界面的制品中心页面选取,或者通过YAML文件指定镜像引用。

1. 选择镜像

说明

通过公网拉取ACS AI容器镜像需要先获取鉴权密钥。建议您使用VPC方式加速拉取AI容器镜像,减少镜像拉取的时间。

docker login egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com

2. 调用API开启编译器+重计算显存优化

  • 启用编译优化

    使用transformers Trainer API:

    image.png

  • 启用重计算显存优化

    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:[tag]

# 运行训练demo
ljperf --action train --model_name deepspeed/llama3-8b

4. 使用建议

  • 镜像中的改动涉及Pytorch、DeepSpeed等库(后续会upstream),请勿重装。

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

Known Issues

  • PPU兼容的cudnn目前不支持fma功能,因此PPU兼容的TE无法走FusedAttention (cudnn) ,可以优先使用FlashAttention:export NVTE_FLASH_ATTN=1 && export NVTE_FUSED_ATTN=0

  • DeepSeek问题

  • 性能问题

    • MOE类大模型使用Autotune工具可以提升性能 ,vLLM支持通过硬件相关的fused_moe kernel配置文件启动kernel以提升性能的优化功能。

  • 量化问题

    • vLLM默认使用Marlin KernelGPTQ/AWQ/W8A8/WOQ进行量化推理加速,可以通过显式指定quantization的方法来使用vllm原本的Marlin/W8A8/awq/gptq等量化推理kernel,例如可以指定使用gptq_acext/awq_acext来做量化推理,但是acext目前不支持act_order=True的量化,以及GPTQ-Int8量化权重,下个版本优化。

    • vLLM使用vllm-flash-attn/xformers/flashinfer作为attention的后端,默认使用vllm-flash-attn性能尚未超过GPU版本,下个版本优化。

    • vLLM量化A8W8对接acext库实现,性能需要持续优化。

    • 目前FP8量化下速度非常慢,还没有进行PPU适配与优化,支持计划待定。

  • 稳定性问题