training-xpu-pytorch 25.07, training-xpu-pytorch 25.07

Main features

• Upgraded PPU SDK to 1.5.3.

• Upgraded vllm to v0.9.1 and optimized sailSHMEM performance.

Bug fixes

SDK 1.5.3 includes the following fixes:

• torch profiler: Fixed incorrect performance data reporting in PyTorch 2.6.

• acext: Fixed a precision issue caused by passing a bias to the w8a8-int8 gemm interface.

• acext: Fixed a crash in the int8_gemm Python API during CUDA Graph Capture.

25.07

PyTorch compilation optimizations

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.

Memory optimization through recomputation

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.

E2E performance gain assessment

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

Image comparison: Base image vs. iterative evaluation.

E2E performance contribution of PPU core components.

We conducted the following tests on a multi-node PPU cluster using release 25.07 to evaluate and compare end-to-end (E2E) training performance. The comparison items include:

1. Base: The PPU PyTorch image.

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

3. ACS AI Image (AC2 + ACCL): The golden image on AC2 BaseOS with optimizations disabled.

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

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

1. Select the image

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

2. Enable compiler and memory optimization

• Enable compilation optimization

  Using the transformers Trainer API:

  

• Enable memory optimization with recomputation

  export CHECKPOINT_OPTIMIZATION=true

3. Start a container

The following steps show how to start a container and run a training task with the built-in ljperf tool.

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

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