training-nv-pytorch 25.12

Main features

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

Bug fixes

None.

Assets

Driver requirements

• The 25.12 release supports two CUDA versions, depending on your NVIDIA driver: CUDA 13.0.2 requires NVIDIA driver version 580 or later, and CUDA 12.8.0 requires NVIDIA driver version 575 or later. For a complete list of supported drivers, see the CUDA Application Compatibility topic. For more information, see CUDA Compatibility and Upgrades.

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 in 8B-parameter LLM training.

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.

This feature is available in the DeepSpeed framework.

E2E performance gain evaluation

Using the cloud-native AI performance evaluation and analysis tool (CNP), we compared the E2E performance against a standard base image using mainstream open-source models and framework configurations. An ablation study further assessed the contribution of each optimized component to the overall model training performance.

Image comparison: Base vs. iterative evaluation

E2E performance contribution analysis of core GPU components

The following tests, conducted on a multi-node GPU cluster using version 25.12, evaluate and compare E2E training performance. The comparison configurations include:

1. Base: The NGC PyTorch image.

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

3. ACS AI image (AC2 + ACCL): The reference image with AC2 BaseOS and ACCL, but with no optimizations enabled.

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

5. ACS AI image (AC2 + ACCL + CompilerOpt + CkptOpt): The reference image with AC2 BaseOS and ACCL, with both torch compile and selective gradient checkpoint optimizations enabled.

1. Select an image

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

2. Enable compiler and recomputation optimizations

• Enable compilation optimization

  Use the transformers Trainer API:

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

• Enable activation recomputation optimization

  export CHECKPOINT_OPTIMIZATION=true

3. Start the container

The following steps use the ljperf model training tool, which is built into the image, to demonstrate how to start a container and run 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

• This image contains modifications to libraries such as PyTorch and DeepSpeed. Do not reinstall them.

• In the DeepSpeed configuration, leave zero_optimization.stage3_prefetch_bucket_size 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 a training or inference task, 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 task, you can use the HPN (High Performance Network). In this case, set NCCL_SOCKET_IFNAME=hpn0.