training-nv-pytorch 26.06

Main features

• Upgraded vllm to 0.22.0.

• Upgraded transformer_engine to 2.15.

Bug fixes

None.

Assets

Driver requirements

• The 26.06 release supports CUDA 13.0.2, which requires NVIDIA driver version 580 or later. For the complete list of supported drivers, see CUDA Application Compatibility. For more information, see 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.

ACCL

ACCL is Alibaba Cloud’s high-performance communication library built for Lingjun. ACCL-N is the GPU-focused version. ACCL-N is a high-performance communication library customized from NVIDIA NCCL. It is fully compatible with NCCL, fixes issues in the upstream NCCL release, and includes performance and stability improvements.

End-to-end performance evaluation

Using the cloud-native AI performance benchmarking tool CNP, comprehensive end-to-end performance comparisons were run against a standard base image with mainstream open-source models and framework configurations.Ablation studies further evaluated how each optimization component contributes to overall model training performance.

Image comparison against base image and iteration evaluation

E2E performance contribution analysis of GPU core components

The following tests are based on this release image. Training E2E performance benchmarks and comparisons were run on a multi-node GPU cluster. The following configurations were compared:

1. Base: NGC PyTorch Image

2. ACS AI Image: Base+ACCL: The image uses the ACCL communication library.

3. ACS AI Image: AC2+ACCL: The golden image uses AC2 BaseOS with no optimizations enabled.

4. ACS AI Image: AC2+ACCL+CompilerOpt: Golden image uses AC2 BaseOS with only the torch compile optimization enabled.

5. ACS AI Image: AC2+ACCL+CompilerOpt+CkptOpt: Golden image uses AC2 BaseOS with both torch compile and selective gradient checkpoint optimizations enabled.

1. Select the image

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

2. Enable the compiler and recomputation memory optimization through the API

• Enable compilation optimization

  Using the transformers Trainer API:

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

• Enable recomputation memory optimization

  export CHECKPOINT_OPTIMIZATION=true

3. Start the container

The image includes the built-in model training tool ljperf, which is used in the following example to demonstrate how to start the container and run a training task.

# Start the container and enter it
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

• Changes in the image affect 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 adjusted for your scenario:

  • When a single Pod requests only 1, 2, 4, or 8 cards for a training or inference task, set NCCL_SOCKET_IFNAME=eth0 (the default configuration in this image).

  • When a single Pod requests all 16 GPUs on the machine (in which case you can use HPN high-performance networking) for a training or inference task, set NCCL_SOCKET_IFNAME=hpn0.