Container Compute Service:training-nv-pytorch 25.10

Key features

• Multi-architecture support: The image now supports both amd64 and aarch64 architectures.

• Megatron-Core upgraded to 0.14.0.

• Transformer Engine upgraded to 2.4.

• vLLM upgraded to 0.11.0.

Bug fixes

No bug fixes in this release.

PyTorch compiling optimization

torch.compile(), introduced in PyTorch 2.0, works well for single-GPU training but cannot benefit LLM training that requires GPU memory optimization and distributed frameworks such as Fully Sharded Data Parallel (FSDP) or DeepSpeed — and may provide negative benefits in these scenarios.

This release addresses the issue with two optimizations:

• Communication granularity control in DeepSpeed: Regulating communication granularity gives the compiler a more complete compute graph, enabling wider-scope optimization.

• Frontend improvements: The PyTorch compiler frontend now handles graph breaks without stopping compilation. Mode matching and dynamic shape support are also improved.

After these changes, end-to-end (E2E) throughput improves by 20% when training an 8B LLM.

GPU memory optimization for recomputation

This release integrates into PyTorch the suggested optimal number of activation recomputation layers. The recommendation is derived from performance tests across clusters and configurations, using collected system metrics such as GPU memory utilization. This allows users to easily benefit from GPU memory optimization. The 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]

2. Enable optimizations

Compiler optimization — To enable the compile optimization using the transformers trainer API:

Gradient checkpointing optimization — Set the following environment variable before starting training:

export CHECKPOINT_OPTIMIZATION=true

3. Start the container and run training

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

Configuration notes

Version 25.10

Public image:

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

VPC image:

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

Replace the placeholders with actual values: