Container Compute Service:training-nv-pytorch 25.06

Updated frameworks

• PyTorch and related components upgraded to V2.7.1.8

• Triton Inference Server upgraded to V3.3.0

• vLLM version compatibility enhanced to support 0.9.1

• Added support for NVIDIA's Blackwell GPU architecture, enabling forward-looking development on next-generation hardware

Bug fix

• Resolved degraded video random access memory (VRAM) optimization efficiency in legacy container images by upgrading PyTorch to V2.7.1.8.

V25.06

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

PyTorch compilation optimization

While torch.compile() delivers strong performance gains in single-GPU scenarios, its benefits are limited in large-scale LLM training due to distributed frameworks like FSDP and DeepSpeed.

• To unlock broader compiler optimizations:

  • We optimized communication granularity within DeepSpeed, exposing larger, more coherent computation graphs to the compiler.

  • Enhanced the compiler frontend to handle arbitrary graph breaks.

  • Improved pattern matching and dynamic shape support for stable compiled performance.

Result: Consistent ~20% end-to-end (E2E) throughput improvement in 8B-parameter LLM training.

Gradient checkpointing optimization

Through extensive benchmarking across models, cluster configurations, and system metrics (including memory utilization), we developed a predictive model to identify optimal activation recomputation layers. This optimization is now natively integrated into PyTorch and supported in DeepSpeed, enabling low-effort adoption of advanced memory optimization techniques.

ACCL

The Alibaba Cloud Communication Library (ACCL) is a suite of high-performance networking (HPN) libraries designed for Lingjun.

One of its key components is ACCL-N, a GPU-accelerated communication library customized from the NVIDIA Collective Communications Library (NCCL). While maintaining full API compatibility with NCCL, ACCL-N provides several enhancements:

• Improved performance: Delivers significantly higher throughput and greater stability, especially in large-scale, multi-node training environments.

• Enhanced stability: Includes targeted bug fixes not yet available in the standard NCCL versions.

E2E performance gain evaluation

Using the Cloud Native Platform (CNP) AI performance analysis tool, we conducted comprehensive E2E comparisons against standard base images (such as NGC PyTorch). Tests used mainstream open-source models and frameworks, with ablation studies to quantify each optimization's contribution.

Test configuration (multi-node GPU clusters)

1. Pull the image

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

2. Enable compiler and memory optimization

• Compilation optimization with Transformers Trainer API

  

• Enable gradient checkpointing optimization

  export CHECKPOINT_OPTIMIZATION=true

3. Launch the container

The image includes a built-in training tool: ljperf.

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

• Do not reinstall PyTorch, DeepSpeed, or related libraries—the image includes pre-optimized binaries.

• In the deepspeed configuration, leave zero_optimization.stage3_prefetch_bucket_size empty or set it to auto.

• The image pre-sets NCCL_SOCKET_IFNAME:

  • When a single pod requests 1/2/4/8 cards for training/inference tasks, set NCCL_SOCKET_IFNAME=eth0. This is the default configuration in the image.

  • For 16-GPU node training: Manually set NCCL_SOCKET_IFNAME=hpn0 to leverage HPN.