Container Compute Service:training-nv-pytorch 25.02

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 ACS is activated (for example, cn-beijing), and {image:tag} with the image name and tag.

Image compatibility

The 25.02-serverless image is not suitable for Lingjun single-tenant products.

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.

This feature is currently available in the DeepSpeed framework only.

ACCL

ACCL (Alibaba Cloud Communication Library) is an in-house high-performance networking (HPN) communication library for Lingjun. ACCL-N, the GPU acceleration variant, is built on top of NCCL (NVIDIA Collective Communications Library) and is fully compatible with it, while fixing known NCCL bugs and delivering higher performance and stability.

This release updates ACCL-N to 2.23.4.11 and adds support for ACCL-Barex.

End-to-end performance assessment

The following benchmark compares end-to-end performance across five configuration levels for multi-node GPU-accelerated clusters. All results are based on 25.02, using mainstream open-source models and frameworks through the CNP (Cloud-Native Performance) tool.

Step 1: Pull the image

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

Step 2: Enable optional optimizations

Compiling optimization — use the Transformers Trainer API:

GPU memory optimization for recomputation:

export CHECKPOINT_OPTIMIZATION=true

Step 3: Launch a container and run a training task

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 training demo.
ljperf benchmark --model deepspeed/llama3-8b

Usage notes

• Do not reinstall PyTorch or DeepSpeed. Changes in the image involve these libraries.

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