👋 Hi! I’m Bibin Wilson. In each edition, I share practical tips, guides, and the latest trends in DevOps and MLOps. If someone forwarded this email to you, you can subscribe here to never miss out!

In the last edition, we learned Kubeflow Pipelines. Today we cover the second Kubeflow subproject: Kubeflow Trainer.

It covers,

By the end, you will know exactly how enterprise teams run distributed training on Kubernetes and what your job is as the MLOps engineer behind it.

In the last edition, I explained how Kubeflow Pipelines turns an ML workflow into a series of containerized steps.

Now, let’s look at one important step in that pipeline. The training step. In that step, the pipeline delegates the training job to Kubeflow Trainer as shown in the image below.

To understand the need for the Kubeflow trainer, first you need to understand why we need Distributed Training.

When we train a machine learning model, the model has to process huge amounts of data, run mathematical operations, calculate errors, and update millions or even billions of parameters again and again.

For small models, a CPU is enough. For example, our employee attrition model, training runs comfortably in a single pod. Because a scikit-learn model on 500,000 rows of data does not need much.

However, for deep learning models and LLMs, CPUs become too slow because the amount of computation is huge. This is where GPUs come in.

Unlike CPUs, GPUs are designed to perform thousands of calculations in parallel. This makes them ideal for training large models.

But there is another challenge that many engineers discover when working with large models. Adding a GPU does not magically solve every training problem. Even GPUs have limits.

For example, when training a large model, the model and data may not even fit inside one GPU memory. So ML teams distribute the training workload across multiple GPUs and sometimes across multiple servers.

This is called distributed training. The following image illustrates it better.

Here is what happens at a high level.

And this is where things become interesting from an infrastructure point of view because 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, GPU allocation, etc..

This is the complexity Kubeflow Trainer helps solve.

When ML engineers train large models or fine-tune LLMs, the training process often needs multiple GPUs running across multiple Kubernetes nodes.

For this, you need the following.

Kubeflow Trainer helps in doing all these in Kubernetes-native way and abstracts away all the complexities.

In simple terms, with Kubeflow Trainer, ML engineers define a TrainJob custom resource, and Kubeflow takes care of creating and managing the required infrastructure for distributed training.

Kubeflow Trainer is basically a Kubernetes Operator implementation.

So it comes with specific custom resource definitions for ML engineers, Platform engineers and MLOps engineers to manage and deploy Training Jobs.

It has three key CRD’s. Let's look at each custom resource in detail and who uses it.

It is managed by Platform engineers / AI infrastructure teams. It defines the following

In short, it defines cluster-scoped reusable default training runtime shared across teams with standard training environments and infrastructure policies.

Here is an example CRD.

Think of it like a golden template created by the platform team. ML engineers focus on training code, while infrastructure standards remain consistent.

It is managed by MLOps engineers or individual teams.

TrainingRuntime is the same concept as ClusterTrainingRuntime. The main difference is scope. Namespace-scoped runtime template.

It is used by ML engineers / AI engineers. It defines the actual training workload, including training code, dataset configuration, model parameters, and a reference to a TrainingRuntime or ClusterTrainingRuntime.

Here is an example TrainJob CRD.

Using this CRD, ML engineers can customize workload-specific configurations such as container image, CPU, memory, and GPU requirements without modifying the default TrainingRuntime or ClusterTrainingRuntime managed by the platform team.

The following image illustrates the high-level overview of how Kubeflow Trainer works.

When a user creates a TrainJob, Kubeflow Trainer does not directly create training pods.

Instead, it converts the TrainJob into a JobSet. A JobSet is a Kubernetes-native API designed for distributed workloads. It manages a group of related Kubernetes Jobs that need to run together.

For example, in distributed AI training, you may need,

JobSet ensures these components are created and managed as a single training workload.

Now lets look at the three key components of Kubeflow trainer.

This is the main controller of Kubeflow Trainer. It watches for TrainJob resources created by users. When a TrainJob is submitted, the controller combines TrainJob, TrainingRuntime & ClusterTrainingRuntime and generates a JobSet resource.

The JobSet controller takes care of the actual distributed job orchestration. It watches JobSet resources and creates the required Kubernetes Jobs and training pods.

It also manages the lifecycle of these workloads.

LeaderWorkerSet (LWS) controller is an optional component used for Kubeflow Trainer’s Distributed Data Cache feature. It is not required for normal TrainJob execution.

When Distributed Data Cache is enabled, the LWS controller manages the leader-worker cache workloads that help stream large datasets efficiently to GPU nodes during training.

Now let’s put everything we learned into practice with a hands-on setup. Here is what we will do.

Inside the repo, the phase-2-enterprise-setup/kubeflow-trainer folder contains the YAML manifest used to create the TrainJob custom resource.

We will set up a standalone Kubeflow Trainer using the official Helm chart.

Run the following command to deploy Kubeflow Trainer.

In the above command, I have added the --set runtimes.defaultEnabled=true flag to enable all the default training runtimes.

After the installation, ensure the trainer pod and jobset controller are in running state.

To verify the available default runtimes, run the following command.

In the above list, torch-distributed, jax-distributed, and xgboost-distributed can run training on CPU nodes.

For testing, we will use the torch-distributed runtime, which runs the training workload on CPUs.

In this example, we will perform distributed training with two worker pods. Each pod processes a portion of the training data and synchronizes with the other worker to train the model, as shown below.

You can find the trainjob.yaml manifest under phase-2-enterprise-setup/kubeflow-trainer folder

Lets deploy the TrainJob.

Once it is applied, use the following command to list the training pods. You will see two pods created by TrainJobs in two different nodes for distributed training.

Since the deployed pods are created as Kubernetes Jobs, they will move to the completed state once the training script executes successfully.

Since there is no PersistentVolume (PV/PVC) or external storage configured, the saved model in the pod will be lost when the pod gets deleted after the training job completes.

For this hands-on example, /tmp/model.pt is mainly used to validate that the distributed training completed successfully.

When ML engineers work on training workflows, they can use the Kubeflow SDK to trigger training jobs directly from their workstation.

You can try this using the train.py script available in the kubeflow-trainer folder.

Let’s set up the Python environment and install the Kubeflow SDK.

Then, run the train.py script to create the TrainJob.

The script will submit the training job to Kubernetes and print the training details in the terminal output.

If you no longer need Kubeflow Trainer, run the following command to uninstall it.

You might be wondering how we can run training jobs on GPU nodes, since we used a simple CPU-based training example.

To run training workloads on GPUs, the Kubernetes cluster should have GPU-enabled nodes with the required GPU drivers and device plugins configured.

Also, default runtimes may not work because GPU nodes usually have taints, and the default runtimes do not include the required tolerations.

To solve this, you need to create a custom runtime with GPU tolerations and node selectors, then use it in your TrainJob to schedule workloads on GPU nodes.

For example, the following ClusterTrainingRuntime spec configures GPU scheduling by adding the required tolerations, node selector, and GPU-enabled training image.

Now, any TrainJob that references this runtime will be scheduled on the GPU nodes and can access the GPUs based on the configured resource requests.

You now know both Kubeflow subprojects (pipelines & trainer) our final project uses.

Now, we need a way to track every training run, log parameters and metrics, and properly version trained models. That is done using MLflow.

In the next edition, I will cover MLflow.

The TrainJobs you learned today create the training outputs, and MLflow helps track, version, and manage them.

When I started learning MLOps, one misconception I had was: ML workload = GPU. But that is not true.

Not every model needs a GPU.