training-nv-pytorch 26.03

更新时间:
复制 MD 格式

These are the release notes for version 26.03 of training-nv-pytorch.

What's new

Highlights

  • Upgraded torch to version 2.10.

  • Upgraded vllm to version 0.17.0.

  • Upgraded megatron-core to version 0.16.0.

  • Upgraded deepspeed to version 0.18.8.

  • Upgraded transformer_engine to version 2.12.

Bug fixes

None for this release.

Contents

Image name

training-nv-pytorch

Tag

26.03-cu130-serverless

26.03-cu128-serverless

Use case

Training/Inference

Framework

pytorch

Requirements

NVIDIA Driver release >= 580

NVIDIA Driver release >= 575

Supported architectures

amd64 & aarch64

amd64

Core components

  • Ubuntu: 24.04

  • Python: 3.12.7+gc

  • CUDA: 13.0

  • perf: 5.4.30

  • gdb: 15.1

  • torch: 2.10.0+ali.10.nv25.10

  • triton: 3.6.0

  • transformer_engine: 2.12.0+5671fd36

  • deepspeed: 0.18.8+ali

  • flash_attn: 2.8.3

  • transformers: 4.57.6+ali

  • grouped_gemm: 1.1.4

  • accelerate: 1.11.0+ali

  • diffusers: 0.34.0

  • mmengine: 0.10.3

  • mmcv: 2.1.0

  • mmdet: 3.3.0

  • opencv-python-headless: 4.11.0.86

  • ultralytics: 8.3.96

  • timm: 1.0.26

  • vllm: 0.17.0+cu130

  • flashinfer-python: 0.6.4

  • pytorch-dynamic-profiler: 0.24.11

  • peft: 0.16.0

  • ray: 2.54.1

  • megatron-core: 0.16.0

  • Ubuntu: 24.04

  • Python: 3.12.7+gc

  • CUDA: 12.8

  • perf: 5.4.30

  • gdb: 15.1

  • torch: 2.10.0+ali.10.nv25.3.pgo

  • triton: 3.6.0

  • transformer_engine: 2.12.0+5671fd36

  • deepspeed: 0.18.8+ali

  • flash_attn: 2.8.3

  • flash_attn_3: 3.0.0b1

  • transformers: 4.57.6+ali

  • grouped_gemm: 1.1.4

  • accelerate: 1.11.0+ali

  • diffusers: 0.34.0

  • mmengine: 0.10.3

  • mmcv: 2.1.0

  • mmdet: 3.3.0

  • opencv-python-headless: 4.11.0.86

  • ultralytics: 8.3.96

  • timm: 1.0.26

  • vllm: 0.17.0+cu128

  • flashinfer-python: 0.6.4

  • pytorch-dynamic-profiler: 0.24.11

  • peft: 0.16.0

  • ray: 2.54.1

  • megatron-core: 0.16.0

Assets

Public images

CUDA 13.0.2 (Driver >=580, amd64 and aarch64)

  • egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:26.03-cu130-serverless

CUDA 12.8 (Driver >= 575, amd64)

  • egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:26.03-cu128-serverless

VPC images

For faster access when pulling an ACS AI container image from within a VPC, replace the public URI egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/{image:tag} with the corresponding VPC URI acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}.

  • {region-id}: The ID of the ACS region, such as cn-beijing or cn-wulanchabu.

  • {image:tag}: The name and tag of the AI container image. Examples include inference-nv-pytorch:25.10-vllm0.11.0-pytorch2.8-cu128-20251028-serverless and training-nv-pytorch:25.10-serverless.

Note

This image is compatible with ACS and multi-tenant Lingjun environments, but not with single-tenant Lingjun environments.

Driver requirements

  • The 26.03 release supports both CUDA 12.8.0 and CUDA 13.0.2, which have different driver requirements. CUDA 13.0.2 requires NVIDIA Driver release 580 or later, while 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.

Key features and enhancements

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 analysis tool, we conducted a comprehensive end-to-end performance comparison between our image and a standard base image. We used popular open-source models and framework configurations. We also performed an ablation study to assess the performance contribution of each optimization component.

Image comparison and iterative evaluation

image.png

E2E performance contribution of core GPU components

We compared the following configurations to evaluate the end-to-end training performance of this image on a multi-node GPU cluster:

  1. Base: The official NGC PyTorch image.

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

  3. ACS AI Image (AC2 + ACCL): The image on AC2 BaseOS with ACCL, but no other optimizations.

  4. ACS AI Image (AC2 + ACCL + CompilerOpt): The image on AC2 BaseOS with ACCL and the torch.compile optimization enabled.

  5. ACS AI Image (AC2 + ACCL + CompilerOpt + CkptOpt): The image on AC2 BaseOS with ACCL, torch.compile, and selective gradient checkpointing enabled.

image.png

Quick start

The following examples show how to pull the training-nv-pytorch image using Docker.

Note

To use the training-nv-pytorch image in ACS, select it from the Artifact Center page when you create a workload in the console, or specify the image reference in a YAML file.

Step 1: Pull the image

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

Step 2: Enable optimizations

  • Enable compile 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

Step 3: Start the container

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

For LLMs

# 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 recommendations

  • The image contains modified versions of libraries such as PyTorch and DeepSpeed. Do not reinstall them.

  • In the DeepSpeed configuration, leave zero_optimization.stage3_prefetch_bucket_size blank or set it to auto.

Known issues

There are no known issues in this release.