Container Compute Service:training-nv-pytorch 25.07

Main Features

• Upgraded vllm to v0.9.2.

• Support for transformerEngine v2.3.0+5de3e14, peft v0.16.0, and diffusers v0.34.0.

Bugs Fix

(None at this time)

25.07

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

PyTorch compiling optimization

The compiling optimization feature introduced in PyTorch 2.0 is suitable or small-scale training on one GPU. However, LLM training requires GPU memory optimization and a distributed framework, such as FSDP or DeepSpeed. Consequently, torch.compile() cannot benefit your training or provide negative benefits.

• Controlling the communication granularity in the DeepSpeed framework helps the compiler obtain a complete compute graph for a wider scope of compiling optimization.

• Optimized PyTorch:

  • The frontend of the PyTorch compiler is optimized to ensure compiling when any graph break occurs in a compute graph.

  • The mode matching and dynamic shape capabilities are enhanced to optimize the compiled code.

After the preceding optimizations, the E2E throughput is increased by 20% when a 8B LLM is trained.

GPU memory optimization for recomputation

We forecast and analyze the consumption of GPU memory of models by running performance tests on models deployed in different clusters or configured with different parameters and collecting system metrics, such as GPU memory utilization. Based on the results, we suggest the optimal number of activation recomputation layers and integrate it into PyTorch. This allows users to easily benefit from GPU memory optimization. Currently, this feature can be used in the DeepSpeed framework.

ACCL

ACCL is an in-house HPN communication library provided by Alibaba Cloud for Lingjun. It provides ACCL-N for GPU acceleration scenarios. ACCL-N is an HPN library customized based on NCCL. It is completely compatible with NCCL and fixes some bugs in NCCL. ACCL-N also provides higher performance and stability.

E2E performance gain assessment

Using the cloud-native AI performance evaluation and analysis tool CNP, we conducted comprehensive end-to-end performance comparison analyses with mainstream open-source models and framework configurations against standard base images. Through ablation experiments, we further evaluated the contribution of each optimization component to the overall model training performance.

GPU core component E2E performance contribution analysis.

The following tests are based on 25.07, conducting training E2E performance evaluation and comparative analysis on multi-node GPU clusters, with comparisons including the following:

1. Base: NGC PyTorch Image.

2. ACS AI Image: Base+ACCL: Image using the ACCL communication library.

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

4. ACS AI Image: AC2+ACCL+CompilerOpt: Golden image using AC2 Base OS, with only torch compile optimization enabled.

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

1. Select image

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

2. Call API to enable compiler + recomputation memory optimization

• Enable compilation optimization

  Using transformers Trainer API:

  

• Enable recomputation memory optimization

  export CHECKPOINT_OPTIMIZATION=true

3. Start container

The image includes the built-in model training tool ljperf, which is used to illustrate the steps for starting a container and running training tasks.

# Start container and enter
docker run --rm -it --ipc=host --net=host  --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]

# Run training demo
ljperf benchmark --model deepspeed/llama3-8b 

4. Usage recommendations

• The changes in the image involve libraries such as Pytorch and Deepspeed. Do not reinstall them.

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

• The environment variable NCCL_SOCKET_IFNAME built into this image needs to be adjusted dynamically according to the usage scenario:

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

  • When a single Pod requests all 16 cards of a machine (at this point you can use HPN high-performance network) for training/inference tasks: you need to set NCCL_SOCKET_IFNAME=hpn0.