Container Compute Service:training-nv-pytorch 25.03

Main features

• The base image is updated to align with NGC 25.02.

• Upgraded PyTorch and related components to 2.6.0.7, TransformerEngine (TE) to 2.1, and accelerate to 1.5.2.

• Upgraded ACCL-N to 2.23.4.12.

• Upgraded vLLM to 0.8.2.dev0 and Ray to 2.44, and added support for flash-infer 0.2.3. Upgraded Transformers to 4.49.0+ali and flash_attn to 2.7.2.

Bug fixes

The upgrade to vLLM 0.8.2.dev0 resolves the Illegal memory access for Mixture of Experts (MoE) on H20 (#13693) issue.

PyTorch compiler optimization

While the torch.compile() feature introduced in PyTorch 2.0 offers significant gains for single-GPU workloads, these benefits diminish in large language model (LLM) training due to the demands of GPU memory optimization and distributed training frameworks like FSDP and DeepSpeed. This release addresses these challenges by:

• Controlling communication granularity within the DeepSpeed framework, which helps the compiler obtain a more complete computation graph, enabling broader optimization.

• Including an optimized version of PyTorch that:

  • Improves the PyTorch compiler frontend to ensure successful compilation even when graph breaks occur in the computation graph.

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

Together, these optimizations increase E2E throughput by approximately 20% for 8B LLM training.

GPU memory optimization for recomputation

We developed a predictive model for GPU memory consumption based on extensive performance data from various models, clusters, and training configurations. This model recommends the optimal number of activation recomputation layers. This functionality is integrated into PyTorch, allowing you to easily leverage the performance benefits of GPU memory optimization. This feature is currently adapted for the DeepSpeed framework.

ACCL communication library

ACCL is a high-performance networking communication library developed by Alibaba Cloud for Lingjun products. The ACCL-N variant is designed for GPU workloads. Customized from NVIDIA NCCL, ACCL-N is fully compatible with the upstream library and includes bug fixes, performance optimizations, and stability enhancements.

E2E performance benefit assessment

Using the cloud-native AI performance evaluation tool CNP, we conducted a comprehensive E2E performance comparison between this image and 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 optimized component.

E2E performance contribution of GPU core components

The following tests, based on version 25.03, were conducted on a multi-node GPU cluster to assess and compare E2E performance. The comparison items include:

1. Base: NGC PyTorch Image

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

3. ACS AI Image: AC2+ACCL: This image using AC2 BaseOS with no optimizations enabled.

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

5. ACS AI Image: AC2+ACCL+CompilerOpt+CkptOpt: This image using AC2 BaseOS with both torch.compile and selective gradient checkpointing optimizations enabled.

Step 1: Pull the image

docker pull egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]

Step 2: Enable compiler and recomputation memory optimizations

• Enable compiler optimization

  When using the Transformers Trainer API, set torch_compile=True:

  training_args = TrainingArguments(
      bf16=True,
      gradient_checkpointing=True,
      torch_compile=True
  )

• Enable GPU memory optimization for recomputation

  export CHECKPOINT_OPTIMIZATION=true

Step 3: Launch the container

The image includes the ljperf model training tool. Use the following steps to launch the container and run a training task.

# Launch the container and enter the shell
docker run --rm -it --ipc=host --net=host  --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]

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

Step 4: Usage notes

• The image includes modified versions of libraries such as PyTorch and DeepSpeed. Do not reinstall them, as this will overwrite the included optimizations. These changes will be upstreamed in the future.

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