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, works well for small-scale single-GPU training. For large language model (LLM) training, which requires GPU memory optimization and distributed frameworks like Fully Sharded Data Parallel (FSDP) or DeepSpeed, standard torch.compile() provides limited or negative benefit.

This release addresses that limitation with two optimizations:

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

• PyTorch compiler frontend improvements: The frontend is optimized to ensure compilation even when a graph break occurs. Mode matching and dynamic shape capabilities are enhanced to optimize the compiled code.

Result: end-to-end (E2E) throughput increases by 20% when training an 8B LLM.

GPU memory optimization for activation recomputation

This feature profiles GPU memory consumption across models deployed in different clusters or with different parameters by running performance tests and collecting system metrics such as GPU memory utilization. Based on the results, it automatically determines the optimal number of activation recomputation layers and integrates the setting into PyTorch — providing GPU memory savings without manual tuning. Currently supported in the DeepSpeed framework.

ACCL

Alibaba Cloud Communication Library (ACCL) is an in-house High-Performance Network (HPN) communication library provided by Alibaba Cloud for Lingjun. ACCL-N is the GPU acceleration variant — an HPN library built on NVIDIA Collective Communications Library (NCCL) that is fully compatible with NCCL, fixes known NCCL bugs, and delivers higher performance and stability.

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.