Container Compute Service:training-nv-pytorch 26.05

Main features

• Upgraded Torch-related components to version 2.11.0.

• Upgraded the built-in training components, megatron-core to 0.17.0 and transformer_engine to 2.14.

• Upgraded the vllm inference component to version 0.20.2.

Bug fixes

None for this release.

Assets

Driver requirements

• The 26.05 release supports CUDA 13.0.2, which requires NVIDIA driver version 580 or later. For a complete list of supported drivers, see the CUDA Application Compatibility guide. For more information, see CUDA Compatibility and Upgrading.

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

We used CNP, a cloud-native AI performance benchmarking and analysis tool, to conduct a comprehensive end-to-end performance comparison against the standard base image using mainstream open-source models and framework configurations. We also performed ablation studies to evaluate the contribution of each optimized component to the overall training performance.

Image vs. base image & iterative evaluation

GPU core components E2E performance contribution analysis

We conducted the following tests on a multi-node GPU cluster to evaluate and compare the E2E training performance of this image release. 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 golden image using AC2 BaseOS with no optimizations enabled.

4. ACS AI image: AC2+ACCL+CompilerOpt: The golden image using AC2 BaseOS with only the torch.compile optimization enabled.

5. ACS AI image: AC2+ACCL+CompilerOpt+CkptOpt: The golden image using AC2 BaseOS with both torch.compile and selective gradient checkpointing optimizations enabled.

1. Select an image

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

2. Enable compiler and checkpointing via API

• Enable compilation optimization

  Use the transformers Trainer API:

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

• Enable gradient checkpointing memory optimization

  export CHECKPOINT_OPTIMIZATION=true

3. Start the container

The image has a built-in model training tool named ljperf. This tool is used as an example to demonstrate the steps for starting a container and running 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 needs to be dynamically adjusted based on the usage scenario:

  • If a single Pod requests only 1, 2, 4, or 8 GPUs for training or inference tasks, you need to set NCCL_SOCKET_IFNAME=eth0. This is the default configuration in this image.

  • When a single Pod requests all 16 GPUs on a machine for training or inference tasks, set NCCL_SOCKET_IFNAME=hpn0 to use the high-performance network (HPN).