Container Compute Service:training-nv-pytorch 25.02

Main Features

• The base image is updated to support NGC 25.01, CUDA is updated to 12.8.0, and cuDNN is updated to 9.7.0.66.

• ACCL-N is updated to 2.23.4.11 and supports ACCL-Barex.

• Transformers is updated to 4.48.3+ali, vLLM is updated to 0.7.2, and Ray is updated to 2.42.1 to support new features and fix bugs.

Bugs Fix

None

Public image

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

VPC image

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

  {region-id} indicates the region where your ACS is activated, such as cn-beijing.

  {image:tag} indicates the name and tag of the image.

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 benefit assessment

With the cloud-native AI performance assessment and analysis tool CNP, we can use mainstream open source models and frameworks together with standard base images to analyze E2E performance. In addition, we can use ablation study to further assess how each optimization component benefits the overall model training.

GPU core component E2E performance benefit analysis

The following test is based on 25.02. E2E performance assessment and comparison is conducted for multi-node GPU-accelerated clusters.

1. Base: The NGC PyTorch image.

2. ACS AI Image: Base+ACCL: The ACCL used by the image.

3. ACS AI Image: AC2+ACCL: The Golden image uses AC2 BaseOS, without any optimizations.

4. ACS AI Image: AC2+ACCL+CompilerOpt: The Golden image uses AC2 BaseOS, with only PyTorch compiling optimization enabled.

5. ACS AI Image: AC2+ACCL+CompilerOpt+CkptOpt: The Golden image uses AC2 BaseOS, with PyTorch compiling optimization and selective gradient checkpoint optimization enabled.

1. Select an image

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

2. Call the API to enable compiling optimization and GPU memory optimization for recomputation

• Enable compiling optimization

  Use the transformers Trainer API:

  

• Enable GPU memory optimization for recomputation

  export CHECKPOINT_OPTIMIZATION=true

3. Launch containers

The image provides a built-in model training tool named ljperf to demonstrate the procedure for launching containers and running training tasks.

# Launch a container and log on to 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. Suggestions

• Changes in the image involve the PyTorch and DeepSpeed libraries. Do not reinstall it.

• Leave zero_optimization.stage3_prefetch_bucket_size in the DeepSpeed configuration empty or set it to auto.