Container Compute Service:training-nv-pytorch 25.10

Key features

• Adds multi-architecture support for amd64 and aarch64.

• Upgrades megatron-core to 0.14.0 and transformer_engine to 2.4, incorporating the latest community features.

• Upgrades vLLM to 0.11.0, incorporating the latest community features.

Bug fixes

No bug fixes in this release.

Assets

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

Driver requirements

• The 25.10 release is based on CUDA 12.8.0 and requires NVIDIA driver version 575 or later. However, if you are running on a data center GPU (such as T4 or any other data center GPU), you can also use drivers from the R470 (470.57+), R525 (525.85+), R535 (535.86+), or R545 (545.23+) series.

• The R418, R440, R450, R460, R510, R520, R530, R545, R555, and R560 driver series are not forward-compatible with CUDA 12.8 and must be upgraded. For a complete list of supported drivers, refer to the CUDA Application Compatibility topic. For more information, refer to CUDA Compatibility and Upgrades.

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 evaluation

Using the cloud-native AI performance benchmarking tool CNP, we performed a comprehensive end-to-end performance comparison between this image and a standard base image using mainstream open-source models and framework configurations. We also conducted an ablation study to evaluate the performance contribution of each optimization component.

Image comparison: Base image and iterative evaluation

E2E performance contribution of core GPU components

The following tests are based on version 25.10 and were run on a multi-node GPU cluster to evaluate and compare end-to-end training performance. The compared configurations are:

1. Base: A standard NGC PyTorch image.

2. ACS AI Image: AC2: This image with AC2 BaseOS and no optimizations enabled.

3. ACS AI Image: AC2+CompilerOpt: This image with AC2 BaseOS and only the torch compile optimization enabled.

4. ACS AI Image: AC2+CompilerOpt+CkptOpt: This image with AC2 BaseOS and both torch compile and selective gradient checkpointing optimizations enabled.

1. Pull the image

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

2. Enable optimizations

• Enable compiler optimization

  Use the transformers Trainer API:

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

• Enable gradient checkpointing optimization

  export CHECKPOINT_OPTIMIZATION=true

3. Start the container

The model training tool ljperf is built into the image, and this topic uses it to describe the steps for starting a container and running a 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 

4. Usage recommendations

• The image contains custom modifications to libraries such as PyTorch and DeepSpeed. Do not reinstall these packages, as this will overwrite the optimizations.

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

• The built-in environment variable NCCL_SOCKET_IFNAME in this image must be dynamically adjusted based on the usage scenario:

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

  • When a single Pod requests all 16 GPUs on a machine for training or inference tasks, set NCCL_SOCKET_IFNAME=hpn0 to use the High-Performance Network (HPN).