training-nv-pytorch 25.09

Updated components

• PyTorch and its related components are upgraded to 2.8.0.

• Transformers is upgraded to 4.56.1+ali, incorporating features and bug fixes from the corresponding open-source version.

Bug fixes

• Fixed an error that occurred when torch.compile() was enabled for open-source Transformers on Qwen2-VL.

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.

Driver compatibility

The 25.09 release is based on CUDA 12.8.0 and requires NVIDIA driver version 575 or later. For data center GPUs such as the T4, the following driver versions are also supported:

• 470.57 or later (R470)

• 525.85 or later (R525)

• 535.86 or later (R535)

• 545.23 or later (R545)

The following driver versions are not forward-compatible with CUDA 12.8 and must be upgraded before using this image: R418, R440, R450, R460, R510, R520, R530, R545, R555, and R560. For the complete list of supported drivers, see CUDA Application Compatibility. For upgrade guidance, see CUDA Compatibility and Upgrades.

Public image

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

VPC image

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

Replace the placeholders with the following values:

Step 1: Pull the image

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

Step 2: Enable compile and recomputation optimizations

Enable compile optimization

Use the Transformers Trainer API:

Enable activation recomputation for GPU memory optimization

export CHECKPOINT_OPTIMIZATION=true

Step 3: Start the container and run a training task

The image includes a built-in model training tool named ljperf. The following example starts a container and runs an LLM 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 the training demo
ljperf benchmark --model deepspeed/llama3-8b