Container Compute Service:training-nv-pytorch 25.04

25.04

Public image:

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

VPC image:

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

• {region-id}: the region where your ACS is activated, such as cn-beijing or cn-wulanchabu.

• {image:tag}: the name and tag of the image.

training-nv-pytorch:25.04-serverless is compatible with the ACS product form and the Lingjun multi-tenant product form. It is not compatible with the Lingjun single-tenant product form.

training-nv-pytorch:25.04 (no suffix) is for Lingjun single-tenant scenarios.

PyTorch compilation optimization

torch.compile(), introduced in PyTorch 2.0, works well for single-GPU training but provides limited benefit—or even negative benefit—for LLM training, which requires GPU memory optimization and distributed frameworks such as FSDP or DeepSpeed.

This release addresses that gap with two compiler improvements:

• Communication granularity control in DeepSpeed: controls the communication granularity in the DeepSpeed framework, giving the compiler a wider scope of computation graph for optimization.

• Improved PyTorch compiler frontend: ensures compilation proceeds even when graph breaks occur in a computation graph, with enhanced mode matching and dynamic shapes support.

After these optimizations, end-to-end throughput increases by 20% when training an 8B LLM.

GPU memory optimization for recomputation

This release analyzes GPU memory consumption across different cluster configurations and parameter settings, then derives the optimal number of activation recomputation layers and integrates them into PyTorch. This makes selective activation checkpointing available with minimal configuration. Currently supported in the DeepSpeed framework.

ACCL

ACCL is Alibaba Cloud's in-house High-Performance Network (HPN) communication library for Lingjun. ACCL-N, the GPU acceleration variant, is built on NCCL with full NCCL compatibility, additional bug fixes, and improved performance and stability.

End-to-end performance contribution analysis

The following tests use Golden-25.04 on multi-node GPU clusters, comparing the contribution of each optimization component to end-to-end throughput. The cloud-native AI performance tool CNP runs evaluations with mainstream open-source models and frameworks alongside standard base images.

Test configurations:

1. Base: NGC PyTorch image

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

3. ACS AI image: AC2 + ACCL: golden image with AC2 BaseOS, no optimizations enabled

4. ACS AI image: AC2 + ACCL + CompilerOpt: golden image with AC2 BaseOS, PyTorch compilation optimization enabled

5. ACS AI image: AC2 + ACCL + CompilerOpt + CkptOpt: golden image with AC2 BaseOS, both compilation optimization and selective activation checkpointing enabled

Prerequisites

Before you begin, make sure that you have:

• Docker installed on your host machine

• An NVIDIA driver that meets the driver requirements

• Access to the image registry (public image or VPC image, depending on your network)

1. Pull the image

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

2. Enable compilation optimization and GPU memory optimization for recomputation

Enable compilation optimization — use the Transformers Trainer API:

Enable GPU memory optimization for recomputation:

export CHECKPOINT_OPTIMIZATION=true

3. Launch a container and run a training task

The image includes a built-in model training tool, ljperf, for launching containers and running training tasks.

# Launch a container and log in.
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