Container Compute Service:training-nv-pytorch 25.05

Features

• The base image CUDA is upgraded to 12.9.0.

Bug fixes

• PyTorch is upgraded to 2.6.0.7.post1, which fixes the profile crash issue in the open-source community.

25.05

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

VPC image

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

Replace the placeholders with actual values:

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.

1. Pull the image

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

2. Enable optimizations

Enable compiling optimization

Use the transformers Trainer API:

Enable GPU memory optimization for activation recomputation

export CHECKPOINT_OPTIMIZATION=true

3. Launch a container

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

LLM example:

# Launch a container and log on to 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

• This release modifies the PyTorch and DeepSpeed libraries. Do not reinstall them.

• Leave zero_optimization.stage3_prefetch_bucket_size blank or set it to auto in your DeepSpeed configuration.