Container Compute Service:training-nv-pytorch 25.12

Main features

• Upgraded vLLM to 0.12.0 and flashinfer-python to 0.5.3.

Bug fixes

None.

Core components

The following table lists the core components and component versions in each image.

Public images

Each CUDA version of training-nv-pytorch 25.12 is published at its own image address. For the driver version and CPU architectures that each image requires, see the Image specifications section.

VPC images

Replace the AI container image asset URI egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/{image:tag} specified in the YAML file of the ACS console with acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag} to quickly pull PG1 AI container images over the VPC.

• Where {region-id} is the available region where your ACS is activated, such as cn-beijing and cn-wulanchabu.

• {image:tag} is the name and tag of the image.

PyTorch compilation optimization

torch.compile(), introduced in PyTorch 2.0, often delivers strong gains for small-scale, single-GPU workloads. But LLM training depends on GPU memory optimization and distributed frameworks such as FSDP or DeepSpeed, so torch.compile() may offer limited benefits or even degrade performance.

• Control communication granularity in the DeepSpeed framework. This helps the compiler capture a more complete compute graph and apply broader compilation optimizations.

• Use an optimized PyTorch build:

  • The PyTorch compiler frontend is improved to ensure that compilation succeeds even if a graph break occurs in the compute graph.

  • Pattern matching and dynamic shape support are strengthened to improve post-compilation performance.

With these optimizations, 8B-parameter LLM training typically achieves an end-to-end throughput gain of about 20%.

GPU memory optimization for recomputation

A predictive model for GPU memory overhead, built on large-scale performance data—including different models, clusters, and training parameter settings, as well as system metrics such as GPU memory utilization collected during benchmarking—recommends the optimal number of activation recomputation layers. This approach is integrated into PyTorch, allowing you to achieve the performance gains of GPU memory optimization with minimal effort. This feature is now supported in the DeepSpeed framework.

E2E performance gain evaluation

A comprehensive end-to-end (E2E) performance comparison against a standard base image was conducted by using the cloud-native AI performance evaluation and analysis tool (CNP), with mainstream open-source models and framework configurations. An ablation study further assessed how much each optimized component contributes to overall model training performance.

Comparison of this image with the base image across iterations

E2E performance contribution analysis of core GPU components

E2E training performance was evaluated on a multi-node GPU cluster running version 25.12 across the following configurations:

• Base: The NGC PyTorch image.

• ACS AI image (Base + ACCL): The base image with the ACCL communication library.

• ACS AI image (AC2 + ACCL): AC2 BaseOS with ACCL and no optimizations enabled.

• ACS AI image (AC2 + ACCL + CompilerOpt): AC2 BaseOS with ACCL and only the torch compile optimization enabled.

• ACS AI image (AC2 + ACCL + CompilerOpt + CkptOpt): AC2 BaseOS with ACCL and both the torch compile and selective gradient checkpoint optimizations enabled.

  

Usage recommendations

The following restrictions and settings apply whenever you run the training-nv-pytorch image:

• Bundled libraries — Do not reinstall libraries such as PyTorch and DeepSpeed. This image contains modified versions of these libraries.

• DeepSpeed configuration — In the DeepSpeed configuration, leave zero_optimization.stage3_prefetch_bucket_size empty or set it to auto.

• Network interface — Adjust the built-in NCCL_SOCKET_IFNAME environment variable in this image based on your scenario:

  • When a single Pod requests 1, 2, 4, or 8 GPUs for a training or inference job, set NCCL_SOCKET_IFNAME=eth0. This is the default configuration in this image.

  • When a single Pod requests all 16 GPUs on a node for a training or inference job, you can use the HPN (High Performance Network). In this case, set NCCL_SOCKET_IFNAME=hpn0.

1. Pull the image

Confirm the NVIDIA driver version and CPU architecture of your environment, and then run the following command to pull the image:

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

Replace [tag] with one of the published tags:

• 25.12-cu130-serverless: CUDA 13.0.2, requires NVIDIA driver 580 or later, and supports amd64 and aarch64.

• 25.12-cu128-serverless: CUDA 12.8.0, requires NVIDIA driver 575 or later, and supports amd64 only.

2. Enable compilation and recomputation optimizations

Enable compilation optimization

Set the following fields in the transformers TrainingArguments class:

training_args = TrainingArguments(
    bf16=True,
    gradient_checkpointing=True,
    torch_compile=True
)

Enable GPU memory optimization for recomputation

Set the following environment variable:

export CHECKPOINT_OPTIMIZATION=true

3. Start the container

The ljperf model training tool is built into the image. The following commands use ljperf to start a container and run a training job.

LLMs

# 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