Container Compute Service:training-nv-pytorch 26.01

Component upgrades

• The built-in training component megatron-core is upgraded to 0.15.0.

• The built-in inference component vLLM is upgraded to 0.13.0.

• The built-in component flashinfer-python is upgraded to 0.5.3.

• The built-in component health_check is upgraded to work with shuttle 1.5.3.

Fixed issues

This release has no fixed issues.

Core components

Public image

VPC image

To pull ACS AI container images quickly within a Virtual Private Cloud (VPC), replace the Asset URI egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/{image:tag} of the specified AI container image with acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}.

• {region-id}: the ID of an ACS region. For example, cn-beijing or cn-wulanchabu. For more information, see Regions.

• {image:tag}: the name and tag of the AI container image. For example, inference-nv-pytorch:25.10-vllm0.11.0-pytorch2.8-cu128-20251028-serverless or training-nv-pytorch:25.10-serverless.

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 recommends the optimal number of activation recomputation layers. The model is 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. This approach is integrated into PyTorch, allowing you to achieve the performance gains of GPU memory optimization with minimal effort. This feature is supported in the DeepSpeed framework.

1. Pull the image

Select the tag whose NVIDIA driver requirement and supported architectures match your nodes, and then replace [tag] in the following command with that tag. For the requirements of each tag, see Image contents.

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

2. Enable the compiler and gradient checkpointing optimizations

training_args = TrainingArguments(
    bf16=True,
    gradient_checkpointing=True,
    torch_compile=True
)

export CHECKPOINT_OPTIMIZATION=true

3. Start the container

The image includes the built-in model training tool ljperf, which the following steps use to start a container and run a training job.

# Start the container and log on to it
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

Usage recommendations

• Do not reinstall the libraries that the image modifies, such as PyTorch and DeepSpeed.

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

• Adjust the built-in environment variable NCCL_SOCKET_IFNAME in this image based on your scenario:

  • When a single pod requests only 1, 2, 4, or 8 GPUs for a training or inference job, set NCCL_SOCKET_IFNAME=eth0 (the default configuration in this image).When a single pod requests only 1, 2, 4, or 8 GPUs for a training or inference job, set NCCL_SOCKET_IFNAME=eth0 (the default configuration in this image). (one of the default configurations in this image).

  • When a single pod requests all 16 GPUs of a machine for a training or inference job, and you can use the High-Performance Network (HPN), set NCCL_SOCKET_IFNAME=hpn0.