training-nv-pytorch 26.01

Key features

• Upgraded the built-in training component megatron-core to 0.15.0, the inference component vLLM to 0.13.0, and flashinfer-python to 0.5.3.

• Updated health_check to be compatible with shuttle 1.5.3.

Fixed issues

None in this release.

Assets

Driver requirements

• The 26.01 release supports both CUDA 12.8.0 and CUDA 13.0.2, which have different driver requirements. CUDA 13.0.2 requires NVIDIA driver version 580 or later. CUDA 12.8.0 requires NVIDIA driver version 575 or later. For a full 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.

End-to-end performance evaluation

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

Image comparison: Base vs. iterative improvements

End-to-end performance contribution of core GPU components

We conducted the following tests on a multi-node GPU cluster using version 26.01 to compare end-to-end training performance. The configurations compared are:

1. Base: The standard NGC PyTorch image.

2. ACS AI Image (Base + ACCL): The base image integrated with the Alibaba Cloud Communication Library (ACCL).

3. ACS AI Image (AC2 + ACCL): The ACS AI image with the AC2 BaseOS with ACCL, but with no other optimizations enabled.

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

5. ACS AI Image (AC2 + ACCL + CompilerOpt + CkptOpt): The ACS AI image with the AC2 BaseOS with 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 compiler and gradient checkpointing optimizations

• Enable compilation optimization

  Using 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 image has a built-in model training tool named ljperf. This topic uses ljperf to explain the steps for starting a container and running a training task.

# Start the container and enter its shell
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 modified versions of 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 in this image, NCCL_SOCKET_IFNAME, needs to be dynamically adjusted based on the usage scenario:

  • When a single Pod requests only 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 an entire machine for training or inference tasks, you must set NCCL_SOCKET_IFNAME=hpn0 to use the High-Performance Network (HPN).