Kubeflow Trainer
A Kubernetes-native platform for distributed AI model training and large language model fine-tuning, orchestrating multi-node, multi-GPU jobs across frameworks such as PyTorch, JAX, HuggingFace and Megatron-LM.
Kubeflow Trainer is a Kubernetes-native platform for distributed training and fine-tuning of AI models, including large language models, across a wide range of frameworks such as PyTorch, MLX, Hugging Face, DeepSpeed, Megatron-LM, JAX and XGBoost. It brings MPI to Kubernetes, orchestrating multi-node, multi-GPU jobs for high-throughput, low-latency communication between GPUs across high-performance computing clusters.
It integrates with the broader Kubernetes AI ecosystem: Kueue for topology-aware and multi-cluster scheduling, a Slurm Bridge for hybrid Kubernetes and Slurm clusters, and KAI Scheduler for GPU-aware scheduling, while reusing existing Kubernetes primitives such as JobSet and LeaderWorkerSet for workload orchestration. A built-in distributed data cache streams large-scale training data to GPU nodes with zero-copy transfer.
Recent releases have added a runtime snapshot mechanism for decoupled runtime lifecycles, JAX and XGBoost training runtimes, and Flux Framework integration for HPC and MPI workloads. It is written in Go and licensed under Apache-2.0.