Container Compute Service:training-nv-pytorch 26.03

Highlights

• Upgraded torch to 2.10.

• Upgraded vllm to 0.17.0.

• Upgraded megatron-core to 0.16.0.

• Upgraded deepspeed to 0.18.8.

• Upgraded transformer_engine to 2.12.

Bug fixes

None.

Core components

Public images

VPC images

Replace the AI container image asset URI egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/{image:tag} specified in the YAML file of the ACS console with acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag} to quickly pull PG1 AI container images over the VPC.

• Where {region-id} is the available region where your ACS is activated, such as cn-beijing and cn-wulanchabu.

• {image:tag} is the name and tag of the image.

Step 1: Pull the image

Run the following command in the host shell:

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

Replace [tag] with the tag that matches your NVIDIA driver version and CPU architecture, either 26.03-cu130-serverless or 26.03-cu128-serverless. For the driver version and architectures that each tag supports, see the Image contents section.

Step 2: Start the container

Run the following command in the host shell. The command starts the container and opens a container shell, where the remaining steps run.

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

Step 3: Enable the optimizations

Enable both optimizations before training starts. Set the compilation optimization in your training script, and set the recomputation memory optimization in the container shell that Step 2 opens.

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

export CHECKPOINT_OPTIMIZATION=true

Step 4: Run the training demo

The image includes ljperf, a built-in model training tool. Run the demo in the container shell.

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

Usage recommendations

• This image ships modified libraries, such as PyTorch and DeepSpeed. Do not reinstall them.

• In the DeepSpeed configuration, leave zero_optimization.stage3_prefetch_bucket_size blank or set it to auto.

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.