Why Kubeflow Trainer Exists: The Journey from Single GPU Training to Distributed AI on Kubernetes

Most tutorials introduce Kubeflow Trainer by showing a YAML file.

kind: TrainJob

They explain the fields, deploy the resource, and show a training job running across multiple GPUs.

The problem is that they skip the most important question.

Why does Kubeflow Trainer exist at all?

If you have already written a PyTorch Distributed Data Parallel (DDP) training script, it can seem unnecessary. PyTorch already provides torchrun, DDP handles gradient synchronization automatically, and your training loop barely changes from the single GPU version.

So why introduce another layer?

The answer lies in understanding how distributed training evolves as models grow. This post walks through that evolution stage by stage, because once you see where torchrun stops being enough, Kubeflow Trainer stops looking like an extra layer and starts looking inevitable.

Stage 1: Life is Simple with a Single GPU

Every machine learning engineer starts here.

A single Python process loads the dataset, performs the forward pass, computes the loss, runs backpropagation, and updates the model weights.

The architecture is beautifully simple.

Dataset
    |
DataLoader
    |
Forward
    |
Loss
    |
Backward
    |
Optimizer Step

One process. One model. One GPU.

There is no communication because there is nobody else to communicate with. There is no coordination because there is nothing to coordinate. If the process crashes, you restart it. If it is slow, you profile it.

As long as your model and a reasonable batch size fit into the memory of a single GPU, life is easy.

Then one of two things happens. Either your model grows past the memory of one GPU, or your dataset grows to the point where a single GPU takes days or weeks to finish an epoch. Both problems have the same first answer: more GPUs.

Stage 2: One GPU Is No Longer Enough

This is where PyTorch Distributed Data Parallel (DDP) enters.

The core idea of DDP is data parallelism. Instead of splitting the model, you replicate it. Every GPU holds a full copy of the model, and each copy trains on a different slice of the data. To make that work, PyTorch creates one Python process for every GPU.

torchrun
      |
 +----+----+
 |    |    |
GPU0 GPU1 GPU2
 |    |    |
Proc Proc Proc

This detail matters more than it looks.

Each process has:

• its own copy of the model
• its own optimizer state
• its own DataLoader
• its own memory space

These processes are completely independent operating system processes. They do not share Python objects. They do not share memory. If you set a variable in one, the others never see it.

The only thing that keeps them behaving like a single training job is DDP.

How DDP Actually Works

Each process trains on a different subset of the dataset. A DistributedSampler makes sure no two processes see the same samples in the same epoch.

GPU0 -> Samples 0 to 31
GPU1 -> Samples 32 to 63
GPU2 -> Samples 64 to 95

Each process independently performs the forward pass, computes its loss, and starts the backward pass.

During loss.backward(), something remarkable happens.

DDP intercepts every gradient as it is computed and performs an AllReduce operation across all participating GPUs. AllReduce is a collective communication operation: every GPU contributes its local gradients, the values are summed and averaged, and every GPU receives the same result back.

GPU0 Gradients
        |
GPU1 Gradients
        |
GPU2 Gradients
        |
    AllReduce
        |
 Averaged Gradients
        |
 optimizer.step()

Two properties make this efficient in practice:

1. Overlap. DDP does not wait for the full backward pass to finish. It buckets gradients and starts communicating them while later layers are still computing. Communication hides behind computation.
2. NCCL. On NVIDIA hardware, the actual data movement is handled by NCCL, which knows how to use NVLink, PCIe, and network interconnects to move tensors between GPUs as fast as the hardware allows.

The result: every process calls optimizer.step() with identical averaged gradients, so every model replica stays byte-for-byte in sync after every iteration.

And critically, your training loop barely changes compared to the single GPU version. Wrap the model in DDP, swap in a DistributedSampler, and the rest looks the same.

The Real Job of torchrun

Many engineers think torchrun performs distributed training.

It does not.

Its job is much simpler and much more boring. It is a process launcher. It starts N Python processes and hands each one the environment variables it needs to find the others:

• LOCAL_RANK: which GPU this process should use on its machine
• RANK: this process's unique global ID across the whole job
• WORLD_SIZE: the total number of processes in the job
• MASTER_ADDR and MASTER_PORT: where the rendezvous happens, the address every process connects to so the group can form

Once every process starts, your code calls:

dist.init_process_group()

Each process reads those environment variables, connects to the master address, and joins the process group. At that point torchrun has essentially completed its work. PyTorch Distributed and NCCL take over for the rest of training.

Keep this mental model: torchrun answers the question "who starts the processes and how do they find each other?" Nothing more. That framing is exactly what you need for the next stage.

Why torchrun Works Perfectly on One Machine

Imagine a server with eight GPUs.

GPU Server

GPU0 GPU1 GPU2 GPU3 GPU4 GPU5 GPU6 GPU7

Running

torchrun --nproc-per-node=8 train.py

creates eight Python processes on that server.

Because every process lives on the same machine, every hard question has a trivial answer:

• Who starts the processes? torchrun, with a single command.
• How do they find each other? MASTER_ADDR=127.0.0.1. Everyone connects over localhost.
• Do they all start together? Yes, they are children of the same command.
• What if one crashes? The whole command fails on one machine, and you rerun it.

No networking complexity exists outside the box. One command, one machine, one failure domain.

The Moment Everything Changes

Now imagine your model grows again.

One server is no longer enough. You need sixteen GPUs across two machines.

Machine A          Machine B

GPU0 ... GPU7      GPU0 ... GPU7

Suddenly, questions appear that have nothing to do with machine learning:

• Process launch. Who runs torchrun on Machine B? You cannot SSH into every node by hand for every experiment.
• Discovery. 127.0.0.1 no longer works. Which real IP is the master? What if that machine changes between runs?
• Rank assignment. Who guarantees that the sixteen processes get global ranks 0 through 15 with no duplicates and no gaps?
• Gang start. init_process_group() blocks until all sixteen processes join. If Machine B is delayed, everyone on Machine A sits waiting, holding eight GPUs hostage.
• Failure handling. If one worker dies at 3am, fifteen processes hang on their next collective operation. Who notices? Who restarts the job?
• Capacity. What if only twelve of the sixteen GPUs are free right now? Do you wait? Who queues the job?

None of these are training problems. Your loss function does not care about any of this.

They are distributed systems problems. And solving them with shell scripts and SSH loops is how teams end up with fragile, unmaintainable training infrastructure.

Kubernetes Solves Half the Problem

Kubernetes already exists to schedule workloads across a fleet of machines. It understands:

• nodes and their available resources, including GPUs via device plugins
• pods as the unit of scheduling
• restarts when things fail
• queueing when the cluster is full

That covers a lot of the list above. Kubernetes can place your workers, restart them, and account for GPU capacity.

But Kubernetes knows absolutely nothing about distributed PyTorch training.

To Kubernetes, your training script is just another container. It has no concept of:

• DDP or NCCL
• RANK, WORLD_SIZE, MASTER_ADDR
• the requirement that all workers must start together or none should
• the fact that if one worker dies, the whole job is dead, not just one pod

If you deploy sixteen training pods with a plain Deployment, Kubernetes will happily start twelve of them, restart individual pods in isolation, and schedule them whenever capacity appears. Every one of those behaviors is correct for a web service and wrong for distributed training.

That gap, between what Kubernetes provides and what distributed training requires, is exactly why Kubeflow Trainer exists.

What Kubeflow Trainer Actually Does

This is the biggest misconception to clear up: Kubeflow Trainer does not replace DDP.

Your training script remains almost identical. You still write:

dist.init_process_group()
 
model = DDP(model)
 
loss.backward()
 
optimizer.step()

Kubeflow Trainer changes none of that. Instead, it automates everything around it. When you submit a TrainJob, the Trainer controller:

• creates the worker pods, one per node, each requesting the right number of GPUs
• sets up networking so workers can reach each other by stable DNS names, solving the discovery problem
• injects the distributed environment, assigning RANK, WORLD_SIZE, MASTER_ADDR, and MASTER_PORT into every pod so init_process_group() just works
• coordinates startup so workers rendezvous as a group rather than initializing into a half-formed job
• treats the job as a unit for failure handling: if a worker dies, the controller can restart the job according to its policy instead of leaving fifteen processes hanging forever
• integrates with Kubernetes scheduling, including gang scheduling, so a job that needs sixteen GPUs either gets all sixteen or waits, rather than deadlocking the cluster by grabbing twelve

In other words, Kubeflow Trainer does across a cluster what torchrun does on a single machine: it launches the processes and prepares them to communicate. The difference is that doing this across machines requires a controller that understands both Kubernetes and the shape of a distributed training job.

Your DDP code executes unmodified. It just wakes up in an environment where all the hard questions from the previous section have already been answered.

Thinking in Layers

The easiest way to understand a production AI training platform is to separate responsibilities into layers.

Layer 3: Kubeflow Trainer
----------------------------
Pod creation
GPU scheduling
Worker coordination
Failure recovery
Environment setup

            |

Layer 2: PyTorch Distributed
----------------------------
DDP
NCCL
Gradient synchronization
DistributedSampler

            |

Layer 1: Your Training Code
----------------------------
Forward
Loss
Backward
Optimizer

Each layer has a completely different responsibility, and each layer only needs to trust the one below it:

• Layer 1 performs machine learning. It assumes a process group exists.
• Layer 2 performs distributed training. It assumes processes have been launched and can find each other.
• Layer 3 performs distributed orchestration. It makes Layer 2's assumptions true on a Kubernetes cluster.

When something breaks, this layering also tells you where to look. NaN loss is Layer 1. Gradient desync or NCCL timeouts are Layer 2. Pods stuck in Pending or workers that never rendezvous are Layer 3.

Understanding these boundaries is the key to understanding modern AI infrastructure.

Final Thoughts

Kubeflow Trainer was not created because DDP was insufficient. DDP already solves distributed training extremely well.

Kubeflow Trainer exists because the moment distributed training leaves a single machine, someone has to launch, connect, coordinate, and babysit an entire fleet of workers. On one machine, that someone is torchrun. On a cluster, it has to be Kubernetes, and Kubernetes needs to be taught what a distributed training job actually is.

Kubeflow Trainer is that teacher. It is the bridge between a scheduler that understands pods and a framework that understands gradients.

Once you understand that, the YAML becomes the easy part.

The architecture is what really matters.

Related Articles

If you're following this GPU fundamentals series, these articles build on one another:

• Understanding GPU Work Units: Threads, Warps, Thread Blocks and Grids
• Understanding the GPU Memory Hierarchy
• Why LLM Decoding Is Memory-Bound
• Tensor Cores Explained (Coming Soon)