training-xpu-pytorch 25.05, training-xpu-pytorch 25.05

更新时间:
复制 MD 格式

This document details the training-xpu-pytorch 25.05 release.

Key features and bug fixes

Key features

  • The base image is upgraded to Ubuntu 24.04, and Python is upgraded to 3.12.7+gc.

  • Upgraded PPU SDK to 1.5.1 and vLLM to v0.8.5 to support W8A8 (int8) quantized models for Qwen3 and DeepSeek-R1.

Bug fixes

Software components

Use case

Training/inference

Framework

pytorch

Requirements

PPU SDK 1.5.1

Core components

  • Ubuntu 24.04

  • Python 3.12.7+gc

  • Torch 2.6.0.7.post1+ppu1.5.1

  • CUDA 12.6

  • ACCL-P 1.5.0-9

  • EIC 1.3.8

  • triton 3.2.0

  • TransformerEngine 2.1

  • deepspeed 0.15.4+ali

  • flash-attn 2.5.7+torch2.6.ppu1.5.1

  • xformers 0.0.29.post1

  • cupy 13.3.0

  • transformers 4.51.2+ali

  • megatron-core 0.9.0

  • grouped_gemm 1.1.4

  • accelerate 1.6.0+ali

  • diffusers 0.31.0

  • openmim 0.3.9

  • bitsandbytes 0.45.3

  • mmengine 0.10.3

  • mmcv 2.1.0

  • mmdet3d 1.4.0

  • vllm 0.8.5

  • pytorch-dynamic-profiler 0.24.11

  • perf 5.4.30

  • gdb 12.1

  • peft 0.13.2

  • ray 2.43

  • spacy 2.3.6

  • cumm 0.5.3

  • spconv 2.3.6

  • ffmpeg 7.0.1

  • open3d 0.19.0

  • flashinfer 0.2.3

Container images

25.05

Public image

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

VPC images

  • acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}

    {region-id} is the ID of the region where your Container Service for Kubernetes (ACK) cluster is located, such as cn-beijing or cn-wulanchabu.
    {image:tag} is the actual image name and tag.

Driver requirements

  • Driver version >= 1.1.0

Key features and enhancements

PyTorch compiler optimization

The compilation optimization feature introduced in PyTorch 2.0 usually provides significant benefits in small-scale, single-GPU scenarios, but in LLM training, the need to incorporate memory optimization and distributed frameworks such as FSDP and Deepspeed makes it difficult for torch.compile() to provide benefits and can even cause performance degradation:

  • This optimization controls communication granularity within the DeepSpeed framework, allowing the compiler to capture a more complete computation graph and apply broader optimizations.

  • This release includes an optimized version of PyTorch with the following enhancements:

    • An optimized PyTorch compiler frontend that ensures successful compilation even when arbitrary graph breaks occur.

    • Enhanced pattern matching and dynamic shape capabilities to improve the performance of compiled code.

With these optimizations, you can typically achieve an end-to-end (E2E) throughput gain of around 20% when training an 8B LLM.

Recomputation memory optimization

A model to predict memory overhead was built using extensive benchmark data from various models, clusters, and training configurations—along with system metrics like memory utilization. This allows the system to recommend the optimal number of activation checkpointing layers. Integrated into PyTorch and compatible with the DeepSpeed framework, this feature helps users optimize memory usage.

ACCL communication library

ACCL is a high-performance network communication library developed by Alibaba for Lingjun products. It is available in three versions: ACCL-N for NVIDIA GPUs, ACCL-P for PPUs, and ACCL-R for AMD accelerators. ACCL-N is based on NVIDIA NCCL, offering full compatibility and including bug fixes and optimizations for performance and stability. ACCL-P is a collective communication library built upon T-Head's open-source pccl. ACCL-R is built upon AMD's open-source ROCm rccl. This release ports key features from ACCL and ACCL-N to pccl, resolves known issues, and is deeply customized for integration with Alibaba Cloud's proprietary networking components.

End-to-end performance assessment

We used CNP, a cloud-native AI performance analysis tool, to compare the end-to-end performance of this image against a standard base image using mainstream open-source models and framework configurations. We also performed an ablation study to evaluate the performance contribution of each optimization component.

Image comparison: base vs. optimized

image.png

PPU core component end-to-end performance contribution analysis

The following tests evaluate and compare end-to-end training performance on a multi-node PPU cluster using the 25.05 release. The comparison includes the following configurations:

  1. Base: The standard PPU PyTorch Image.

  2. ACS AI Image (Base + ACCL): The base image with the ACCL communication library.

  3. ACS AI Image (AC2 + ACCL): The optimized image using the AC2 base OS without any optimizations enabled.

  4. ACS AI Image (AC2 + ACCL + CompilerOpt): The optimized image using the AC2 base OS with only the torch.compile optimization enabled.

  5. ACS AI Image (AC2 + ACCL + CompilerOpt + CkptOpt): The optimized image using the AC2 base OS with both torch.compile and selective gradient checkpointing optimizations enabled.

image.png

Quick start

The following example shows how to pull the training-xpu-pytorch image using Docker.

Note

To use the training-xpu-pytorch image in ACS, select it from Artifacts when creating a workload in the console, or specify the image reference in a YAML file.

Step 1: Pull the image

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

Step 2: Enable optimizations

  • Enable compiler optimization

    With the transformers Trainer API:

    image.png

  • Enable recomputation memory optimization

    export CHECKPOINT_OPTIMIZATION=true

Step 3: Start the container

The model training tool ljperf is built into the image. The following steps describe how to start a container and run a training task by using this tool.

Note

For ACS products, use the image by specifying it in a YAML file.

LLM models

# Start and enter the container
docker run --rm -it --ipc=host --net=host  --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-xpu-pytorch:25.05

# Run the training demo
ljperf --action train --model_name deepspeed/llama3-8b

Usage notes

  • This image contains modified versions of libraries such as PyTorch and DeepSpeed. Do not reinstall them. These modifications will be contributed to the upstream projects in a future release.

  • In your DeepSpeed configuration, leave the zero_optimization.stage3_prefetch_bucket_size parameter empty or set it to auto.

  • This image supports the following quantization methods for running DeepSeek R1 (671B) and Qwen3 models:

    • vLLM 0.8.5: Supports AWQ (w4a16), GPTQ (w4a16, w8a16), and per-token/per-channel a8w8 (int8) quantization schemes (to run models quantized by using AWQ and GPTQ on vLLM 0.8.5, you must add the --quantization moe_wna16 flag).

  • The DeepEP and sailSHMEM features mentioned in the PG1 SDK 1.5.1 Release Note are not yet supported. Support for these features is not yet scheduled.

  • For optimal performance, we recommend using this image with driver version 1.5.1. For configuration instructions, see Specify GPU models and driver versions for ACS GPU pods and GPU driver versions.

Known issues

  • Performance for the Qwen3-235B-A22B model will be further optimized in a future release.

  • DeepSeek R1 model

    • vLLM images may produce garbled output during inference, so we recommend that you disable prefix caching by adding the --no-enable-prefix-caching option.

    • When testing the DeepSeek model, the --max-model-len parameter must be less than 80k. Otherwise, an OOM error may occur during runtime.

  • Stability issues

  • Performance issues

    • Using an autotune tool can improve performance for Mixture of Experts (MOE) models. vLLM supports a performance optimization feature that allows you to load a hardware-specific fused_moe kernel using a configuration file.