Container Compute Service:training-nv-pytorch 25.02

Main features

• The base image is now aligned with NGC 25.01, CUDA is upgraded to 12.8.0, and cuDNN is upgraded to 9.7.0.66.

• ACCL-N is upgraded to 2.23.4.11 and now supports ACCL-Barex.

• Transformers is upgraded to 4.48.3+ali, vLLM is upgraded to 0.7.2, and Ray is upgraded to 2.42.1. This release contains the features and bug fixes from these upstream versions.

Bug fixes

None for this release.

Public image

• egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:25.02-serverless

VPC image

acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}

Replace {region-id} with the region where your Alibaba Cloud Container Compute Service (ACS) is activated (for example, cn-beijing or cn-wulanchabu).

Replace {image:tag} with the name and tag of the image.

PyTorch compilation optimization

torch.compile(), introduced in PyTorch 2.0, often delivers strong gains for small-scale, single-GPU workloads. But LLM training depends on GPU memory optimization and distributed frameworks such as FSDP or DeepSpeed, so torch.compile() may offer limited benefits or even degrade performance.

• Control communication granularity in the DeepSpeed framework. This helps the compiler capture a more complete compute graph and apply broader compilation optimizations.

• Use an optimized PyTorch build:

  • The PyTorch compiler frontend is improved to ensure that compilation succeeds even if a graph break occurs in the compute graph.

  • Pattern matching and dynamic shape support are strengthened to improve post-compilation performance.

With these optimizations, 8B-parameter LLM training typically achieves an end-to-end throughput gain of about 20%.

GPU memory optimization for recomputation

A predictive model for GPU memory overhead, built on large-scale performance data—including different models, clusters, and training parameter settings, as well as system metrics such as GPU memory utilization collected during benchmarking—recommends the optimal number of activation recomputation layers. This approach is integrated into PyTorch, allowing you to achieve the performance gains of GPU memory optimization with minimal effort. This feature is now supported in the DeepSpeed framework.

ACCL

ACCL is Alibaba Cloud’s high-performance communication library built for Lingjun. ACCL-N is the GPU-focused version. ACCL-N is a high-performance communication library customized from NVIDIA NCCL. It is fully compatible with NCCL, fixes issues in the upstream NCCL release, and includes performance and stability improvements.

End-to-end performance assessment

Using CNP, a cloud-native AI performance benchmarking and analysis tool, we conducted a full end-to-end performance comparison against a standard base image with mainstream open-source models and framework configurations. We also conducted ablation studies to measure how much each optimization component contributes to overall model training performance.

End-to-end performance contribution analysis of GPU key components

The following tests are based on 25.02 and run end-to-end performance benchmarks and comparisons on a multi-node GPU cluster. The comparison items include:

1. Base: NGC PyTorch image.

2. ACS AI Image: Base+ACCL: Base image with the ACCL communication library.

3. ACS AI Image: AC2+ACCL: Base image on AC2 BaseOS with ACCL.

4. ACS AI Image: AC2+ACCL+CompilerOpt: Base image on AC2 BaseOS with ACCL and torch.compile() optimization enabled.

5. ACS AI Image: AC2+ACCL+CompilerOpt+CkptOpt: Base image on AC2 BaseOS with ACCL, torch.compile(), and selective gradient checkpointing optimizations enabled.

1. Pull the image

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

2. Enable compiler and memory optimization

• Enable compilation optimization

  Use the Transformers Trainer API:

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

• Enable GPU memory optimization for recomputation

  export CHECKPOINT_OPTIMIZATION=true

3. Launch the container

The image includes the model training tool ljperf. The following steps show how to start a container and run a training task.

# Launch and enter the container.
docker run --rm -it --ipc=host --net=host  --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]
# Run the training demo.
ljperf benchmark --model deepspeed/llama3-8b 

4. Usage notes

• Do not reinstall PyTorch or DeepSpeed. The image includes changes to these libraries (which will be contributed upstream later).

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