Distributed training across multiple GPU nodes
Generated by Claude through GigAI 0.1.11.10 on 2026-10-11.
Concept: Distributed training · also covers Foundation models · Step 13 of 35 in MLOps / ML platform engineer · version 1
Practice thisOn this page
When a model or dataset is too large for one GPU, training is spread across several GPUs and several nodes that must start together, find each other and stay in sync. The training framework handles gradient synchronisation; the platform handles launching the worker group, wiring up networking and recovering from failures. Ray Train and Kubeflow Trainer are the two common ways to run this on Kubernetes.
Read more
- docRay Train: Scalable Model Training (docs.ray.io)
- docKubeflow Trainer (www.kubeflow.org)
How it is done in practice
Teams start by making a training script work on one GPU, then on several GPUs in one node, and only then across nodes, measuring throughput at each step to see whether scaling is worth the cost. The platform provides a blessed launch path, either a Ray cluster template or a set of Kubeflow training runtimes, with base images that pin compatible CUDA, driver and framework versions. Jobs write checkpoints to object storage, log metrics to the experiment tracker from the rank zero worker only, and are submitted through a queue so the whole worker group is admitted together. GPU utilisation and data loader wait time are watched on dashboards because a starved input pipeline is the most common reason added GPUs do not help.
Technologies
Diagram
flowchart TB U[User submits job: workers=4, gpu=1 each] --> C[Trainer controller or Ray driver] C --> W0[Worker rank 0 on node A] C --> W1[Worker rank 1 on node A] C --> W2[Worker rank 2 on node B] C --> W3[Worker rank 3 on node B] W0 <-->|gradient all-reduce| W1 W1 <-->|gradient all-reduce| W2 W2 <-->|gradient all-reduce| W3 DS[(Sharded dataset)] --> W0 DS --> W1 DS --> W2 DS --> W3 W0 -->|checkpoints| S[(Shared storage)] W0 -->|metrics| T[Experiment tracker]
flowchart TB U[User submits job: workers=4, gpu=1 each] --> C[Trainer controller or Ray driver] C --> W0[Worker rank 0 on node A] C --> W1[Worker rank 1 on node A] C --> W2[Worker rank 2 on node B] C --> W3[Worker rank 3 on node B] W0 <-->|gradient all-reduce| W1 W1 <-->|gradient all-reduce| W2 W2 <-->|gradient all-reduce| W3 DS[(Sharded dataset)] --> W0 DS --> W1 DS --> W2 DS --> W3 W0 -->|checkpoints| S[(Shared storage)] W0 -->|metrics| T[Experiment tracker]
Code
pythonimport torch, torch.nn as nn
import ray.train
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer, prepare_model
def train_loop_per_worker(config):
model = prepare_model(nn.Linear(10, 1)) # wraps in DDP, moves to GPU
opt = torch.optim.SGD(model.parameters(), lr=config["lr"])
device = ray.train.torch.get_device()
for epoch in range(config["epochs"]):
x = torch.randn(64, 10, device=device)
loss = model(x).pow(2).mean()
opt.zero_grad(); loss.backward(); opt.step()
ray.train.report({"epoch": epoch, "loss": loss.item()})
trainer = TorchTrainer(
train_loop_per_worker,
train_loop_config={"lr": 0.01, "epochs": 3},
scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
)
result = trainer.fit()Subtopics
How data-parallel training works
Workers, ranks and gradient sync
The dominant pattern is data parallelism. Each worker process holds a full copy of the model on its own GPU and receives a different shard of every batch. After the backward pass, gradients are averaged across all workers with a collective all-reduce operation, so every copy applies the same update. In PyTorch this is DistributedDataParallel: each process has a rank, the group has a world size, and communication goes through a backend such as NCCL over the node network. A distributed sampler ensures workers see disjoint data. The effective batch size grows with the number of workers, which often means the learning rate has to be scaled too.
What limits scaling
Speedup is rarely linear. Gradient synchronisation costs network bandwidth, so interconnect speed between nodes matters, and one slow worker holds back the whole group because each step waits for everyone. Data loading is the other common bottleneck: if CPU workers or storage cannot feed the GPUs, utilisation drops no matter how many GPUs are added. Models too large for one GPU's memory need additional techniques such as sharding parameters or optimizer state across devices.
Ray Train: scaling a training loop from Python
Trainer, scaling config, training function
Ray Train wraps an existing training loop with little change. The user writes a per-worker training function and passes it to a trainer such as TorchTrainer along with a ScalingConfig that states the number of workers and whether each uses a GPU. Ray starts that many worker processes across the cluster, sets up the distributed process group, and runs the function on each. Helper calls prepare the model and data loader for distributed use, moving them to the right device and adding the sampler. Inside the loop the function reports metrics and checkpoints back to the driver. It has integrations for PyTorch, Lightning, Hugging Face Transformers, TensorFlow and XGBoost. On Kubernetes, Ray clusters are usually managed by the KubeRay operator.
Kubeflow Trainer: distributed training as a Kubernetes resource
TrainJob and runtimes
Kubeflow Trainer is the Kubernetes-native approach. A platform team defines reusable runtimes, as ClusterTrainingRuntime or namespaced TrainingRuntime objects, that describe how a framework is launched: images, the number and roles of nodes, and the startup policy. Users submit a TrainJob that references a runtime and supplies their code, arguments and resource needs, either as YAML or through the Python SDK without writing manifests. The controller creates the group of pods, gives them stable network identities so workers can discover each other, and injects the environment the framework needs. It supports PyTorch, MPI-based workloads and others, and it is the successor to the older Training Operator with its per-framework resources such as PyTorchJob.
Separation of concerns
The split between runtime and job mirrors the split between platform and data science teams. Platform engineers own runtimes, base images, node selection and integration with a queueing system such as Kueue. Model developers only choose a runtime, a node count and resources per node. That keeps infrastructure detail out of user code and lets the platform change it centrally.
Checkpointing and fault tolerance
Surviving worker loss
A multi-node job that runs for hours or days will see failures: a node is reclaimed, a GPU errors, a pod is preempted. Because all workers move in lockstep, losing one normally stops the whole group. The defence is periodic checkpointing of model weights, optimizer state, scheduler state and the current step to shared durable storage, usually object storage or a shared filesystem, written by one worker to avoid races. On restart the job loads the latest checkpoint and resumes. Ray Train can restart failed workers and restore from the last reported checkpoint up to a configured failure limit; on Kubernetes the job controller's restart policy does the same. Checkpoint frequency trades lost work against write overhead.
Sources
- docRay Train: Scalable Model Training (docs.ray.io)
- docKubeflow Trainer (www.kubeflow.org)