,

Main features and bug fixes

Main features

• Upgraded the PPU SDK in the base image to 1.4.1.

• Upgraded ACCL-P to 1.4.0-7 and EIC to 1.3.7.cuda12.ppu.

• Upgraded TransformerEngine to 1.12 and VLLM to 0.6.4.post1.

Bug fixes

• Upgraded the SDK to 1.4.1, which fixes a hang that occurred when the HGGC_AUTO_DISPATCH_BARRIER environment variable was enabled.

• Upgraded VLLM to 0.6.4.post1, which fixes a custom all-reduce error caused by an open-source issue.

Contents

Assets

25.01

[Preview] Nightly-20250115

[Preview] Nightly-20250109

Driver requirements

Driver version >= 1.1.0

Key features and enhancements

PyTorch compilation optimizations

The compilation optimization capability introduced in PyTorch 2.0 typically delivers significant gains for small-scale, single-card workloads. However, in LLM training scenarios that require memory optimization and distributed frameworks such as FSDP or DeepSpeed, torch.compile() may not improve, and can even degrade, performance. Our optimizations include the following:

• Controlling communication granularity within the DeepSpeed framework allows the compiler to obtain a more complete computation graph and apply broader-scope optimizations.

• An optimized version of PyTorch:

  • An optimized compiler frontend that ensures successful compilation despite arbitrary graph breaks.

  • Enhanced pattern matching and dynamic shape capabilities for improved compiled code performance.

These optimizations typically provide an end-to-end throughput gain of approximately 20% for 8B LLM training.

Activation checkpointing for memory optimization

We built a predictive model for memory overhead based on extensive performance evaluation data from various models, clusters, and training configurations, along with system metrics such as memory utilization. This model recommends the optimal number of layers for activation checkpointing. Integrated into PyTorch and adapted for the DeepSpeed framework, this feature enables users to easily apply memory optimization.

ACCL collective communication library

ACCL is Alibaba's proprietary high-performance network communication library developed for Lingjun products. It provides three versions for different hardware scenarios: ACCL-N for NVIDIA GPUs, ACCL-P for PPUs, and ACCL-R for AMD GPUs. ACCL-N is a high-performance communication library customized by Alibaba Cloud based on NVIDIA NCCL. While maintaining full compatibility with NCCL, it fixes some bugs in the official NCCL release and includes performance and stability optimizations. ACCL-P is a collective communication library derived from T-Head's open-source pccl library, with additional enhancements. ACCL-R is a collective communication library derived from AMD's ROCm open-source communication library, rccl, with additional enhancements. This release primarily ports major features from ACCL/ACCL-N to pccl, fixes several issues, and integrates deeply with Alibaba Cloud's proprietary high-speed networking components.

End-to-end performance assessment

Using CNP, a cloud-native AI performance evaluation and analysis tool, we compared the end-to-end performance against the standard base image using mainstream open-source models and framework configurations. We also performed ablation studies to evaluate the performance contribution of each optimization.

Image comparison and iterative evaluation

E2E performance contribution of PPU core components

The following tests evaluate the end-to-end performance of version 25.01 on multi-node GPU/PPU clusters. The comparison includes:

1. Alibaba Cloud AI image: Base+ACCL: The base image using the ACCL communication library.

2. Alibaba Cloud AI image: AC2+ACCL: The golden image using the AC2 BaseOS without any optimizations enabled.

3. Alibaba Cloud AI image: AC2+ACCL+CompilerOpt: The golden image using the AC2 BaseOS with only the torch.compile optimization enabled.

4. Alibaba Cloud AI image: AC2+ACCL+CompilerOpt+CkptOpt: The golden image using the AC2 BaseOS with both torch.compile and selective activation checkpointing optimizations enabled.

Quick start

The following examples show how to pull the training-xpu-pytorch image by using Docker.

Step 1: Log in to the registry

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

Step 2: Enable optimizations

• Enable compilation optimization

  Use the Transformers Trainer API:

  

• Enable memory optimization with activation checkpointing

  export CHECKPOINT_OPTIMIZATION=true

Step 3: Start the container

The image includes a built-in model training tool, ljperf. The following steps describe how to start the container and run a training task by using this tool.

LLM workloads

# Start the container and go to the command line
docker run --rm -it --ipc=host --net=host  --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-xpu-pytorch:[tag]

# Run a training demo
ljperf benchmark --model deepspeed/llama3-8b 

Step 4: Usage recommendations

• The image contains modifications to libraries such as PyTorch and DeepSpeed that will be upstreamed. Do not reinstall these libraries.

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

Known issues

The PPU-compatible cuDNN version currently does not support the FMA feature. Consequently, the PPU-compatible TransformerEngine (TE) cannot use FusedAttention (cudnn). Use FlashAttention instead:export NVTE_FLASH_ATTN=1 && export NVTE_FUSED_ATTN=0.