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 ACS and Lingjun multi-tenant scenarios. It is not compatible with Lingjun single-tenant scenarios.

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

PyTorch compiling optimization

torch.compile(), introduced in PyTorch 2.0, is effective for single-GPU training but provides limited or negative benefit for large language model (LLM) training, which depends on GPU memory optimization and distributed frameworks such as Fully Sharded Data Parallel (FSDP) or DeepSpeed.

This release improves torch.compile() for distributed LLM training through two optimizations:

• Communication granularity control in DeepSpeed: Controlling communication granularity gives the compiler a complete compute graph, enabling wider compiling optimization.

• Frontend improvements: The PyTorch compiler frontend now compiles even when a graph break occurs, with enhanced mode matching and dynamic shape capabilities.

Result: 20% higher end-to-end throughput when training an 8B LLM.

GPU memory optimization for recomputation

Based on performance tests across different clusters and parameter configurations, this release integrates the optimal number of activation recomputation layers directly into PyTorch. Enable it with a single environment variable—no manual tuning required.

Currently supported in the DeepSpeed framework.

ACCL

Alibaba Cloud Communication Library (ACCL) is a suite of high-performance networking (HPN) libraries designed for Lingjun.

Its key component is ACCL-N, a GPU-accelerated communication library customized from the NVIDIA Collective Communications Library (NCCL). ACCL-N is fully compatible with NCCL while fixing upstream bugs and providing additional performance and stability enhancements.

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