training-xpu-pytorch 25.09

Main features

• Upgraded the PPU SDK to 1.6.1.

• Upgraded vllm to 0.10.0, which includes community features and bug fixes.

• Upgraded ffmpeg to 7.1.1 to use its stable software VVC decoder.

Bug fixes

• Upgraded the SDK to 1.6.1 to resolve the following issues:

  • Fixed an out-of-memory (OOM) error that could occur during asys python sampling.

  • Resolved a bug where calls to the hgptiGetDeviceId API failed to retrieve the device ID.

  • Addressed an issue in Multi-Process Service (MPS) scenarios where certain metrics could exceed 100%.

  • Fixed a crash in DeepEP low latency mode that occurred when the number of experts exceeded 256.

• Upgraded Triton to version 3.3.0+ppu1.6.1.cu128 to support int8mm optimizations.

Assets

Driver requirements

• Driver version >= 1.1.0

PyTorch compilation 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 evaluation

Using CNP, our cloud native AI performance evaluation tool, we conducted a comprehensive end-to-end performance comparison against a standard base image using mainstream open-source models and framework configurations. We also performed an ablation study to assess the performance impact of each optimization on overall model training.

Image comparison against base image and iterative evaluation

End-to-end performance contribution analysis of PPU core components

The following tests, based on version 25.09, evaluate and compare end-to-end training performance on a multi-node PPU cluster. The comparison includes:

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 golden image with AC2 BaseOS and ACCL, with no optimizations enabled.

4. ACS AI Image: AC2+ACCL+CompilerOpt: The golden image with AC2 BaseOS, ACCL, and the torch compile optimization enabled.

5. ACS AI Image: AC2+ACCL+CompilerOpt+CkptOpt: The golden image with AC2 BaseOS, ACCL, and both torch compile and selective gradient checkpoint optimizations enabled.

1. Log in to the image registry

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

2. Enable compiler and recomputation optimization

• Enable compilation optimization

  Use the transformers Trainer API:

  

• Enable recomputation-based GPU memory optimization

  export CHECKPOINT_OPTIMIZATION=true

3. Start the container

This topic uses the model training tool ljperf, which is built into the image, to describe the steps for starting a container and running a training task.

# 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.09

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