Kubeflow Trainer: Distributed Training on Kubernetes

In the Kubeflow Pipelines guide you saw how Kubeflow turns an ML workflow into a series of containerized steps. One of those steps — training — is often the hardest to scale. This guide covers the Kubeflow subproject that handles it: Kubeflow Trainer.

A Look Into Distributed Training

When you train a machine learning model, it processes large amounts of data, runs many math operations, and updates its parameters repeatedly. For small models — like the employee attrition model from earlier guides, a scikit-learn model trained on 500,000 rows — a single CPU pod is enough.

When models get larger, training involves matrix operations that need to run in parallel, which is where GPUs come in: unlike CPUs, they're designed to perform thousands of calculations simultaneously. For many deep learning models, a single GPU is enough.

Large language models change the picture again. They contain billions of parameters — too large to fit in a single GPU's memory, and even if they did fit, training on one GPU could take months. So ML teams split the model and the data across multiple GPUs running on multiple servers. This is distributed training.

At a high level:

• The training job is split across multiple worker processes.
• Each worker runs on a different GPU.
• Workers process different parts of the data, then calculate updates and sync with each other so every worker stays consistent with the same model.
• Together, all workers behave like one large training system.

Distributed training is not just about running ML code on multiple GPUs — behind the scenes you need to handle scheduling, networking between workers, shared storage, and GPU allocation. This is the complexity Kubeflow Trainer solves.

What is Kubeflow Trainer?

When ML engineers train large models or fine-tune LLMs, the process often needs multiple GPUs running across multiple Kubernetes nodes. That requires creating worker pods, assigning GPU resources, configuring communication between workers, and tracking job status.

Kubeflow Trainer does all of this in a Kubernetes-native way and abstracts away the complexity. In simple terms, ML engineers define a TrainJob custom resource, and Kubeflow Trainer takes care of creating and managing the infrastructure for distributed training.

Kubeflow Trainer Custom Resources

Kubeflow Trainer is implemented as a Kubernetes Operator, so it ships with specific CRDs for platform engineers, MLOps engineers, and ML engineers to manage and deploy training jobs. There are three key custom resources.

How Kubeflow Trainer Works Internally

When a user creates a TrainJob, Kubeflow Trainer does not directly create training pods. Instead, it converts the TrainJob into a JobSet — a Kubernetes-native API designed for distributed workloads that manages a group of related Jobs that need to run together (for example, one coordinator node, multiple worker nodes, and GPU-based training pods).

Hands-on: Run a Distributed TrainJob

Let's set up Kubeflow Trainer and run a distributed PyTorch training job on CPU using the torch-distributed ClusterTrainingRuntime.

1. Install Kubeflow Trainer

Deploy the standalone Kubeflow Trainer using the official Helm chart, enabling all default training runtimes:

helm install kubeflow-trainer oci://ghcr.io/kubeflow/charts/kubeflow-trainer \
  --namespace kubeflow-system \
  --create-namespace \
  --version 2.2.1 \
  --set runtimes.defaultEnabled=true

After installation, confirm the trainer and JobSet controller pods are running:

$ kubectl get po -n kubeflow-system
NAME                                                   READY   STATUS
jobset-controller-67f6757844-4gffm                     1/1     Running
kubeflow-trainer-controller-manager-d68b55dd4-5fbkm    1/1     Running

Then verify the available default runtimes:

$ kubectl get clustertrainingruntimes
NAME                     AGE
deepspeed-distributed    6m30s
jax-distributed          6m30s
mlx-distributed          6m30s
torch-distributed        6m30s
torchtune-llama3.2-1b    6m30s
torchtune-llama3.2-3b    6m30s
torchtune-qwen2.5-1.5b   6m30s
xgboost-distributed      6m30s

torch-distributed, jax-distributed, and xgboost-distributed can run training on CPU nodes — that's what we'll use for this exercise.

2. Create a TrainJob

We'll run distributed training with two worker pods, each processing a portion of the training data and syncing with the other worker. The trainjob.yaml manifest defines a training script inline that creates sample data, trains a simple neural network across the workers, and saves the model on the rank-0 worker pod:

$ kubectl apply -f trainjob.yaml
trainjob.trainer.kubeflow.org/distributed-training created

List the training pods — you'll see two pods created across nodes for distributed training:

$ kubectl get pods
NAME                                    READY   STATUS
distributed-training-node-0-0-p57fq     0/1     ContainerCreating
distributed-training-node-0-1-7lh46     0/1     ContainerCreating

Since the pods are created as Kubernetes Jobs, they move to the completed state once the training script finishes. There's no PersistentVolume configured here, so the saved model is lost once the pod is deleted — for this exercise, /tmp/model.pt is only used to confirm the run succeeded.

3. Trigger training from the SDK

ML engineers can also trigger a TrainJob directly from a workstation with the Kubeflow SDK, which uses ~/.kube/config to connect to the cluster:

python3 -m venv venv
source venv/bin/activate
pip install kubeflow

python3 train.py

The script submits the training job to Kubernetes and prints the training details in the terminal.

4. Clean up

helm uninstall kubeflow-trainer -n kubeflow-system

Running Training Jobs on GPU Nodes

To run training workloads on GPUs, the cluster needs GPU-enabled nodes with the required drivers and device plugins configured. Default runtimes usually won't work as-is, because GPU nodes typically carry taints and the default runtimes don't include matching tolerations.

The fix is to create a custom runtime with GPU tolerations and node selectors, then reference it from your TrainJob. For example:

spec:
  tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: "NoSchedule"
  nodeSelector:
    doks.digitalocean.com/gpu-model: h200
  containers:
    - name: node
      image: pytorch/pytorch:2.5.1-cuda12.1-cudnn9-runtime

Any TrainJob that references this runtime is scheduled on GPU nodes and can access GPUs based on the configured resource requests.

Which Models Actually Need GPUs?

A common misconception is that ML workload always means GPU. It doesn't: