Distributed training on OpenShift AI 3.4 with Kubeflow Trainer v2

Summary

Red Hat OpenShift AI 3.4 introduces Kubeflow Trainer v2, a unified training API that replaces framework-specific CRDs like PyTorchJob and TFJob with a single TrainJob resource. Under the hood, each TrainJob is translated into a JobSet, a Kubernetes-native API for managing groups of Jobs as a single unit, which handles headless services for pod discovery and automatic failure recovery. The post validates the stack through two phases: manual validation with a minimal NCCL all-reduce test across 2 nodes, and real-world training of an embedding-based neural network on a fraud detection dataset using 1, 2, and 4 NVIDIA A10G GPUs. Benchmarking shows near-linear Distributed Data Parallel scaling at 95% efficiency, successful checkpoint recovery after pod kills, and GPU utilization confirmed via DCGM Exporter. The TrainJob API simplifies distributed training by automatically configuring torchrun, environment variables, and pod coordination, while raw JobSets offer more control for multi-stage pipelines and differentiated failure handling.

Takeaways

1. 1

   The TrainJob API automatically configures torchrun, environment variables, and pod coordination, eliminating manual MASTER_ADDR and entrypoint scripts.

2. 2

   Both TrainJob and raw JobSet paths delivered near-linear DDP scaling at 95% efficiency on a fraud detection dataset with 1, 2, and 4 NVIDIA A10G GPUs.

3. 3

   JobSet failure policies allow different actions per stage: FailJobSet for data preparation failures to avoid wasted GPU time, and RestartJobSet for training failures with checkpoint recovery.

FAQ