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 13.0.2 (driver >= 580) and CUDA 12.8.0 (driver >= 575). CUDA Application Compatibility. 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 CNP (cloud-native AI performance benchmarking tool), we ran E2E performance comparisons against a standard base image with mainstream open-source models. Ablation studies measure each optimization component's contribution to training performance.

Image comparison and iteration evaluation

E2E performance contribution of core GPU components

Configurations tested on a multi-node GPU cluster:

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 ljperf, a model training tool. Use it to start a container and run training:

# 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

• Do not reinstall PyTorch or DeepSpeed. This image includes custom patches to these libraries.

• 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.