Container Compute Service:training-nv-pytorch 25.09

Updated components

• PyTorch and its related components are upgraded to 2.8.0.

• Transformers is upgraded to 4.56.1+ali, incorporating features and bug fixes from the corresponding open-source version.

Bug fixes

• Fixed an error that occurred when torch.compile() was enabled for open-source Transformers on Qwen2-VL.

PyTorch compile optimization

torch.compile(), introduced in PyTorch 2.0, works well for small-scale single-GPU training but provides limited or negative benefit for large language model (LLM) training, which requires GPU memory optimization and a distributed framework such as Fully Sharded Data Parallel (FSDP) or DeepSpeed.

This release addresses that gap with two compiler improvements:

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

• PyTorch compiler frontend improvements: The frontend is optimized to continue compiling when a graph break occurs. Mode matching and dynamic shape capabilities are also enhanced.

Together, these optimizations deliver a 20% increase in end-to-end (E2E) throughput when training an 8B LLM.

GPU memory optimization for activation recomputation

This release forecasts and analyzes GPU memory consumption across models deployed in different clusters and parameter configurations, collecting system metrics such as GPU memory utilization. Based on the results, it suggests the optimal number of activation recomputation layers and integrates this into PyTorch—so you get GPU memory savings without manual tuning. This feature is currently available in the DeepSpeed framework.

ACCL

ACCL is Alibaba Cloud's in-house High-Performance Network (HPN) communication library for Lingjun. It includes ACCL-N, an HPN library built on top of NVIDIA Collective Communications Library (NCCL) that is fully compatible with NCCL, fixes several NCCL bugs, and delivers higher performance and stability.

Driver compatibility

The 25.09 release is based on CUDA 12.8.0 and requires NVIDIA driver version 575 or later. For data center GPUs such as the T4, the following driver versions are also supported:

• 470.57 or later (R470)

• 525.85 or later (R525)

• 535.86 or later (R535)

• 545.23 or later (R545)

The following driver versions are not forward-compatible with CUDA 12.8 and must be upgraded before using this image: R418, R440, R450, R460, R510, R520, R530, R545, R555, and R560. For the complete list of supported drivers, see CUDA Application Compatibility. For upgrade guidance, see CUDA Compatibility and Upgrades.

Public image

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

VPC image

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

Replace the placeholders with the following values:

Step 1: Pull the image

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

Step 2: Enable compile and recomputation optimizations

Enable compile optimization

Use the Transformers Trainer API:

Enable activation recomputation for GPU memory optimization

export CHECKPOINT_OPTIMIZATION=true

Step 3: Start the container and run a training task

The image includes a built-in model training tool named ljperf. The following example starts a container and runs an LLM training task:

# Start and enter 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