Container Compute Service:training-nv-pytorch 25.07

Main features

• Upgraded vllm to v0.9.2.

• Added support for TransformerEngine v2.3.0+5de3e14 , peft v0.16.0, and diffusers v0.34.0.

Bug fixes

None for this release.

25.07

• egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:25.07-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

Using the Cloud Native Platform (CNP) AI performance analysis tool, we conducted a comprehensive end-to-end performance comparison between this image and a standard base image using mainstream open-source models and framework configurations. Through ablation studies, we also evaluated the performance contribution of each optimization component to overall model training.

E2E performance contribution analysis of GPU core components

We conducted the following tests, based on version 25.07, on a multi-node GPU cluster to evaluate and compare end-to-end training performance. The comparison configurations include:

1. Base: NGC PyTorch image.

2. ACS AI image: Base+ACCL: Base image with the ACCL communication library.

3. ACS AI image: AC2+ACCL: Golden image with AC2 Base OS and no optimizations enabled.

4. ACS AI image: AC2+ACCL+CompilerOpt: Golden image with AC2 Base OS and only the torch compile optimization enabled.

5. ACS AI image: AC2+ACCL+CompilerOpt+CkptOpt: Golden image with AC2 Base OS and both torch compile and selective gradient checkpointing optimizations enabled.

1. Pull the image

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

2. Enable compilation and gradient checkpointing optimizations (Optional)

• Enable compilation optimization

  Use the Transformers Trainer API:

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

• Enable gradient checkpointing optimization

  export CHECKPOINT_OPTIMIZATION=true

3. Start the container

The image includes ljperf, a built-in model training tool. This example shows how to start the container and run 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 a training demo
ljperf benchmark --model deepspeed/llama3-8b 

4. Usage notes

• The image includes modifications to 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 NCCL_SOCKET_IFNAME environment variable must be adjusted depending on the use case:

  • If a single pod requests 1, 2, 4, or 8 GPUs for a training/inference task, set NCCL_SOCKET_IFNAME=eth0. This is the default setting.

  • If a single pod requests all 16 GPUs on a host to use the high-performance network (HPN) for a training/inference task, set NCCL_SOCKET_IFNAME=hpn0.