Container Compute Service:training-nv-pytorch 25.03

Applicable scenarios

Core components

Training and inference frameworks

CV tools

Debugging and profiling

Public image

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

VPC image

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

Replace {region-id} with the region where your ACS is activated, for example, cn-beijing. Replace {image:tag} with the name and tag of the image.

Choose the right image

PyTorch compiler optimization

torch.compile() delivers clear throughput gains for single-GPU workloads, but distributed training (Fully Sharded Data Parallel (FSDP), DeepSpeed) historically prevented the compiler from seeing a complete computation graph, limiting or even negating those gains. This release addresses that with two optimizations:

• Communication granularity control in DeepSpeed: exposes a complete computation graph to the compiler, enabling broader optimization scope.

• Compiler frontend improvements: the PyTorch compiler frontend now handles graph breaks gracefully, and mode matching and dynamic shape handling are improved for better runtime performance.

Result: up to 20% higher E2E throughput in 8B LLM training.

GPU memory optimization for recomputation

This release integrates automatic activation recomputation layer recommendations directly into PyTorch. The optimal number of recomputation layers is determined by running performance tests across different cluster configurations and collecting GPU memory utilization metrics — no manual tuning required.

Currently supported in the DeepSpeed framework.

ACCL communication library

ACCL-N is Alibaba Cloud's High-Performance Networking (HPN) communication library for Lingjun, built on NCCL with full NCCL API compatibility. ACCL-N 2.23.4.12 delivers higher throughput and stability than stock NCCL and includes additional bug fixes.

E2E performance benefit assessment

The cloud-native AI performance assessment tool CNP measures E2E training performance using mainstream open-source models and standard base images, with ablation study support to isolate the contribution of each optimization.

The following chart shows the cumulative E2E benefit of each optimization layer in version 25.03, measured on a multi-node GPU cluster:

1. Base: NGC PyTorch image (baseline)

2. Base + ACCL: ACCL-N substituted for NCCL

3. AC2 + ACCL: AC2 BaseOS, no additional optimizations

4. AC2 + ACCL + CompilerOpt: AC2 BaseOS with PyTorch compiler optimization

5. AC2 + ACCL + CompilerOpt + CkptOpt: AC2 BaseOS with both PyTorch compiler optimization and selective gradient checkpointing

1. Pull the image

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

2. Enable optimizations

Compiler optimization

Call the transformers Trainer API to enable compiler optimization:

GPU memory optimization for recomputation

export CHECKPOINT_OPTIMIZATION=true

3. Launch a container and run a training demo

The image includes ljperf, a built-in model training tool. The following example launches a container and runs an LLM training demo.

# Launch 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 LLM training demo
ljperf --action train --model_name deepspeed/llama3-8b