training-nv-pytorch 26.02

Main features

• Upgraded DeepSpeed to version 0.18.5 and Transformers to version 4.57.6.

• Upgraded vLLM to version 0.15.0 and flashinfer-python to version 0.6.1.

Bug fixes

None.

Assets

Driver requirements

• The 26.02 release 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 complete 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

We used the CNP cloud-native AI performance analysis tool to conduct a comprehensive end-to-end performance comparison between this image and a standard base image using mainstream open-source models and framework configurations. We also performed an ablation study to evaluate the contribution of each optimization component to the overall model training performance.

Base image comparison and iterative evaluation

E2E performance contribution of core GPU components

The following tests, conducted on a multi-node GPU cluster with the 26.02 image, evaluate and compare end-to-end training performance. The comparison items include:

1. Base: 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 on AC2 BaseOS with no optimizations enabled.

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

5. ACS AI Image: AC2+ACCL+CompilerOpt+CkptOpt: The golden image on AC2 BaseOS with both torch 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 compilation and memory optimization

• Enable compilation optimization

  Use the Transformers Trainer API:

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

• Enable recomputation memory optimization

  export CHECKPOINT_OPTIMIZATION=true

Step 3: Start the container

The image has a built-in model training tool, ljperf, which is used to explain 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 

Step 4: Usage recommendations

• The image contains modifications to libraries such as PyTorch and DeepSpeed. Do not reinstall them.

• In the DeepSpeed configuration, leavezero_optimization.stage3_prefetch_bucket_size empty or set it toauto.

• Adjust the built-in environment variableNCCL_SOCKET_IFNAME in this image based on your use case:

  • When a single pod requests 1, 2, 4, or 8 GPUs for a training or inference task, you must setNCCL_SOCKET_IFNAME=eth0. This is the default configuration in the image.

  • When a single pod requests all 16 GPUs on a node for a training or inference task, you can use the HPN high-speed network. In this case, you must setNCCL_SOCKET_IFNAME=hpn0.