Container Compute Service:training-nv-pytorch 25.06

New features

• This release upgrades PyTorch and related components to v2.7.1.8, Triton to v3.3.0, and vLLM to v0.9.1, adding support for the new Blackwell architecture.

Bug fixes

• Upgraded PyTorch to v2.7.1.8 to fix degraded VRAM optimization performance in older images.

25.06

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

acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}

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

Comprehensive end-to-end performance comparisons were conducted with the Cloud Native Platform (CNP) AI performance analysis tool, benchmarking this image against a standard base image with mainstream open-source models and framework configurations. In addition, ablation studies evaluated the performance contribution of each optimization component to the overall model training.

E2E performance contribution of GPU core components

The following tests were run on a multi-node GPU cluster to compare end-to-end training performance across the following configurations:

1. Base: NGC PyTorch Image

2. Base+ACCL: Base image with ACCL.

3. ACS AI Image: AC2+ACCL: AC2 BaseOS with ACCL and no other optimizations.

4. ACS AI Image: AC2+ACCL+CompilerOpt: AC2 BaseOS with ACCL and torch.compile optimization.

5. ACS AI Image: AC2+ACCL+CompilerOpt+CkptOpt: AC2 BaseOS with ACCL, torch.compile, and selective gradient checkpointing optimizations.

1. Pull the image

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

2. Enable compiler and selective gradient checkpointing

• Enable compilation optimization

  Use the Transformers Trainer API:

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

• Enable selective gradient checkpointing

  export CHECKPOINT_OPTIMIZATION=true

3. Start the container

This section uses the built-in model training tool ljperf to describe the steps for starting a container and running a training task.

# Start the container and enter its command line
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 notes

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

• 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 in this image must be dynamically adjusted based on the use case:

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

  • When a single Pod requests all 16 GPUs on an entire machine for training or inference tasks (which allows you to use HPN high-performance networking), you must set NCCL_SOCKET_IFNAME=hpn0.