Container Compute Service:training-nv-pytorch 25.06

Updated frameworks

• PyTorch and related components upgraded to V2.7.1.8

• Triton Inference Server upgraded to V3.3.0

• vLLM compatibility extended to 0.9.1

• Added support for NVIDIA's Blackwell GPU architecture

Bug fix

Upgrading PyTorch to V2.7.1.8 resolves degraded VRAM (video random access memory) optimization efficiency in legacy container images.

Core components

V25.06

Public image

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

VPC image

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

Replace the placeholders:

PyTorch compilation optimization

torch.compile() delivers strong performance gains in single-GPU scenarios, but its impact is limited in large-scale LLM training because distributed frameworks like FSDP and DeepSpeed introduce frequent graph breaks that constrain the compiler.

To address this, three optimizations are applied:

• DeepSpeed communication granularity: Optimized to expose larger, more coherent computation graphs to the compiler.

• Compiler frontend: Enhanced to handle arbitrary graph breaks.

• Pattern matching and dynamic shape support: Improved for stable compiled performance across varied workloads.

Result: ~20% end-to-end (E2E) throughput improvement in 8B-parameter LLM training.

Gradient checkpointing optimization

Through extensive benchmarking across models, cluster configurations, and system metrics (including memory utilization), a predictive model identifies the optimal activation recomputation layers for each workload. This optimization is natively integrated into PyTorch and supported in DeepSpeed, so you can adopt advanced memory optimization with minimal configuration changes.

1. Pull the image

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

2. Enable compiler and memory optimization

Compilation optimization with Transformers Trainer API

Gradient checkpointing optimization

export CHECKPOINT_OPTIMIZATION=true

3. Launch the container

The image includes a built-in training tool: ljperf.

LLM training example

# Start 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