Overview¶

This guide gives an overview for Kubeflow Trainer project.

What is Kubeflow Trainer¶

Kubeflow Trainer is a Kubernetes-native distributed AI platform for scalable large language model (LLM) fine-tuning and training of AI models across a wide range of frameworks, including PyTorch, MLX, HuggingFace, DeepSpeed, Megatron-LM, JAX, XGBoost, and more.

Kubeflow Trainer brings MPI to Kubernetes, orchestrating multi-node, multi-GPU distributed jobs efficiently across high-performance computing (HPC) clusters. This enables high-throughput communication between processes, making it ideal for large-scale AI training that requires ultra-fast synchronization between GPUs nodes.

Kubeflow Trainer seamlessly integrates with the Cloud Native AI ecosystem, including Kueue for topology-aware scheduling and multi-cluster job dispatching, Slurm Bridge for scheduling on hybrid Kubernetes and Slurm clusters, and KAI Scheduler for GPU aware scheduling.

Kubeflow Trainer reuses existing Kubernetes-native building blocks like JobSet and LeaderWorkerSet for AI workload orchestration.

Kubeflow Trainer provides a distributed data cache designed to stream large-scale data with zero-copy transfer directly to GPU nodes. This ensures memory-efficient training jobs while maximizing GPU utilization.

With the Kubeflow Python SDK, AI practitioners can effortlessly develop and fine-tune LLMs while leveraging the Kubeflow Trainer APIs: TrainJob and Runtimes.

Who is This For?¶

Kubeflow Trainer documentation is organized around three key personas:

User Personas

AI Practitioners¶

ML engineers and data scientists who use the Kubeflow Python SDK and TrainJob APIs to train and fine-tune models at scale.

What you’ll find:

Training guides for PyTorch, JAX, DeepSpeed, MLX
LLM fine-tuning blueprints with TorchTune
Local execution backends for development

Platform Administrators¶

DevOps engineers and cluster operators who deploy and manage Kubeflow Trainer on Kubernetes clusters.

What you’ll find:

Installation and configuration guides
Runtime and policy management
Integration with schedulers (Kueue, Slurm Bridge, Volcano)
Extension framework architecture

Contributors¶

Open source developers who want to contribute to the Kubeflow Trainer project.

What you’ll find:

Architecture documentation
Development workflow
Contributing guidelines
Community resources

Why Use Kubeflow Trainer?¶

Simple, Scalable, and Built for LLM Fine-Tuning¶

Train models with a single Kubernetes CRD (TrainJob) across any supported framework. Scale from single-GPU workloads to massive multi-node distributed training with minimal code changes.

Extensible and Portable¶

Run anywhere: public clouds, on-premises, or hybrid environments. The plugin-based architecture allows custom ML policies, runtimes, and schedulers to be added without modifying the core platform.

Distributed AI Data Caching¶

Optimize data loading with Apache Arrow and Apache DataFusion for high-performance, zero-copy tensor streaming. The distributed cache reduces training time by eliminating data loading bottlenecks.

LLM Fine-Tuning Blueprints¶

Pre-built templates for generative AI fine-tuning with TorchTune, supporting popular models like Llama and Qwen. Configuration-driven workflows eliminate boilerplate code.

Optimized GPU Efficiency¶

Intelligent data streaming and caching maximize GPU utilization, reducing training costs and time. Supports efficient model parallelism with PyTorch FSDP and DeepSpeed ZeRO.

Native Kubernetes Integrations¶

Achieve optimal GPU utilization and coordinated scheduling for large-scale AI workloads. Kubeflow Trainer seamlessly integrates with Kubernetes ecosystem projects like Kueue, Coscheduling, Volcano, or YuniKorn. With Slurm Bridge, TrainJobs can also be scheduled by Slurm on hybrid Kubernetes and Slurm clusters.

AI Lifecycle with Kubeflow Trainer