Container Compute Service:training-nv-pytorch 26.04

Main features

• Upgraded vllm to 0.19.0.

• Upgraded transformer_engine to 2.13.

Bug fixes

None.

Assets

Driver requirements

• Release 26.04 supports CUDA 12.8.0 and CUDA 13.0.2, depending on the driver version. CUDA 13.0.2 requires NVIDIA driver release 580 or later, and CUDA 12.8.0 requires NVIDIA driver release 575 or later. For a full list of supported drivers, see CUDA Application Compatibility. 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 evaluation

Using CNP, a cloud-native AI performance benchmarking and analysis tool, we ran end-to-end performance comparisons against a standard base image with mainstream open-source models and framework configurations. We also ran ablation studies to measure how much each optimization component contributes to overall training performance.

Image comparison and iteration evaluation

E2E performance contribution of core GPU components

The following tests benchmark this image's E2E training performance on a multi-node GPU cluster against several configurations:

1. Base: NGC PyTorch image

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

3. ACS AI Image (AC2+ACCL): Golden image with AC2 BaseOS, with no optimizations enabled.

4. ACS AI Image (AC2+ACCL+CompilerOpt): Golden image with AC2 BaseOS, with only the PyTorch compile optimization enabled.

5. ACS AI Image (AC2+ACCL+CompilerOpt+CkptOpt): Golden image with AC2 BaseOS, with both PyTorch compile and selective gradient checkpointing optimizations enabled.

Step 1: Select an image

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

Step 2: Enable compiler and recomputation optimizations

• Enable compile optimization

  Use the Transformers Trainer API:

  

• Enable re-computation GPU memory optimization

  export CHECKPOINT_OPTIMIZATION=true

Step 3: Start the container

The image includes the model training tool ljperf. The following steps use this tool to show how to start a container and run a training task.

# Start the container and open a 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 

Step 4: Usage notes

• This image includes changes to libraries such as PyTorch and DeepSpeed. Do not reinstall them.

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

• The built-in environment variable NCCL_SOCKET_IFNAME must be set based on your scenario:

  • If a single Pod uses 1, 2, 4, or 8 GPUs for a training or inference task, set NCCL_SOCKET_IFNAME=eth0. This is the default setting in this image.

  • If a single Pod uses all 16 GPUs on a host for a training or inference task, set NCCL_SOCKET_IFNAME=hpn0 to use HPN.