Container Compute Service:training-nv-pytorch 26.02

Main features

• Upgraded DeepSpeed to version 0.18.5 and Transformers to version 4.57.6.

• Upgraded vLLM to version 0.15.0 and flashinfer-python to version 0.6.1.

Bug fixes

None.

Core components

Public images

VPC images

Usage notes

Review the following requirements before you run the training-nv-pytorch image:

• The image contains changes to libraries such as PyTorch and DeepSpeed. Do not reinstall these libraries.

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

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

  • If a single pod requests 1, 2, 4, or 8 GPUs for a training or inference job, you must set NCCL_SOCKET_IFNAME=eth0. This is the default configuration in this image.

  • If a single pod requests all 16 GPUs of a node for a training or inference job, you can use the HPN high-performance network. In this case, you must set NCCL_SOCKET_IFNAME=hpn0.

Step 1: Pull the image

Replace [tag] with the tag that matches the NVIDIA driver version and CPU architecture of your nodes, as listed in the Contents section.

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

Step 2: Start the container and run a training job

The training-nv-pytorch image includes a built-in model training tool named ljperf. The following example uses this tool to start a container and run a training job.

# Start the container and connect 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

Step 3: Enable compilation and GPU memory optimization

Enable the following optimizations of the training-nv-pytorch image:

• Enable compilation optimization: Use the Transformers Trainer API.

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

• Enable GPU memory optimization for recomputation: Set the following environment variable.

  export CHECKPOINT_OPTIMIZATION=true

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.