Container Compute Service:training-nv-pytorch 26.01

Core components

Public network

CUDA 13.0.2 (driver ≥ 580, amd64 & aarch64)

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

CUDA 12.8 (driver ≥ 575, amd64)

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

VPC

To pull the image over a virtual private cloud (VPC), replace the registry prefix:

Where {region-id} is the ID of your Alibaba Cloud Container Compute Service (ACS) region, for example, cn-beijing or cn-wulanchabu. For a full list of supported regions, see Regions and zones.

PyTorch compiling optimization

torch.compile(), introduced in PyTorch 2.0, is effective for single-GPU training but provides limited or negative benefit for large language model (LLM) training, which depends on GPU memory optimization and distributed frameworks such as Fully Sharded Data Parallel (FSDP) or DeepSpeed.

This release improves torch.compile() for distributed LLM training through two optimizations:

• Communication granularity control in DeepSpeed: Controlling communication granularity gives the compiler a complete compute graph, enabling wider compiling optimization.

• Frontend improvements: The PyTorch compiler frontend now compiles even when a graph break occurs, with enhanced mode matching and dynamic shape capabilities.

Result: ~20% higher end-to-end throughput in 8B-parameter LLM training.

GPU memory optimization for recomputation

Based on performance tests across different clusters and parameter configurations, this release integrates the optimal number of activation recomputation layers directly into PyTorch. Enable it with a single environment variable—no manual tuning required.

This feature is available in the DeepSpeed framework.

1. Pull the image

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

Replace [tag] with the tag for your CUDA version: 26.01-cu130-serverless or 26.01-cu128-serverless.

2. Enable compiler and recomputation optimizations

Compiler optimization — enable via the Transformers Trainer API:

GPU memory optimization via activation recomputation — set the following environment variable before starting your training job:

export CHECKPOINT_OPTIMIZATION=true

3. Start the container and run a training job

The image includes ljperf, a built-in model training tool. The following example starts a container and runs an LLM training benchmark:

# 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