training-xpu-pytorch 25.03, training-xpu-pytorch 25.03

Key features and bug fixes

Key features

• Upgraded the PPU SDK in the base image to 1.4.3_hotfix.

• Upgraded ACCL-P to v1.4.3-8 and EIC to 1.3.7.1-v2503, incorporating their respective features and bug fixes.

• Upgraded several components, including vLLM to 0.7.3, Ray to 2.43, FlashInfer to 0.2.1.post1, Accelerate to 1.5.2, and Transformers to 4.49.0+ali.

Bug fixes

No bug fixes in this release.

Software components

Assets

Driver requirements

• Driver version >= 1.1.0

Key features and enhancements

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.

Selective gradient checkpoint 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.

E2E performance evaluation

Using CNP, a cloud-native AI performance analysis tool, a comprehensive E2E performance comparison was conducted between this image and the standard base image across mainstream open-source models and framework configurations. An ablation study was also performed to evaluate the performance contribution of each optimization component.

Performance comparison with the base image

E2E performance contribution of core GPU components

The following tests, based on version 25.03, were conducted on a multi-node PPU cluster to evaluate E2E performance. The comparison includes:

1. Base: 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 no optimizations enabled.

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

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

Quick start

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

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

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

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

Known issues

• The PPU-compatible cuDNN currently does not support the FMA feature. Therefore, the PPU-compatible TE cannot use FusedAttention (cuDNN). You can use FlashAttention instead: export NVTE_FLASH_ATTN=1 && export NVTE_FUSED_ATTN=0.

• Issues with DeepSeek models

  • When testing the DeepSeek model, the --max-model-len setting must be < 80k. Otherwise, an OOM error may occur at runtime.

  • SGLang does not support running INT8 quantized models.

  • For vllm 0.7.3, you must set the --max-num-batched-tokens parameter to avoid the framework's default limit of 2048 tokens (see the open-source community issue [Usage]:Input prompt (2501 tokens) is too long and exceeds limit of 2048 #13370).

• Performance issues

  • For large Mixture-of-Experts (MoE) models, performance can be improved by using an autotune tool. vLLM includes an optimization feature that loads a hardware-specific fused_moe kernel configuration file to boost performance.

• Quantization issues

  • By default, vLLM uses the Marlin Kernel to accelerate quantized inference for GPTQ, AWQ, W8A8, and WOQ. You can explicitly specify the quantization method to use vLLM's original quantized inference kernels, such as Marlin, W8A8, awq, or gptq. For example, you can specify gptq_acext or awq_acext to perform quantized inference. However, acext currently does not support quantization with act_order=True or GPTQ-Int8 quantized weights. This will be optimized in the next version.

  • vLLM uses vllm-flash-attn, xformers, or flashinfer as the attention backend. The default backend, vllm-flash-attn, does not yet outperform the GPU version in terms of performance. This will be optimized in a future release.

  • In vLLM, A8W8 quantization is implemented by using the acext library, and its performance has not yet been fully optimized.

  • FP8 quantization is currently very slow because it has not yet been adapted or optimized for PPU. The support schedule is to be determined.

• Stability issues

  • When you perform multi-GPU inference with vLLM, an error related to mp_distributed_executor may occur. Use VLLM_WORKER_MULTIPROC_METHOD=spawn to prevent this issue. For more information, see the GitHub issue [Bug]: When tensor_parallel_size>1, RuntimeError: Cannot re-initialize CUDA in forked subprocess. #6152.

  • After multi-GPU inference with vLLM is complete, the process may exit with an error. This is a known community issue. For more information, see the open-source community issue [Bug]: When using tp for inference, an error occurs: Worker VllmWorkerProcess pid 3283517 died, exit code: -15. #6145.