Multi-Node Distributed Training on Kubernetes

A single DGX H100 node has 8 H100 GPUs and 640 GB of HBM3 memory.

That sounds like a lot.

But a Llama 3 405B model requires roughly 800 GB just to hold the weights in BF16.

It doesn't fit.

Even if it did fit, training on 8 GPUs alone would take months.

The solution is to spread the job across multiple nodes.

flowchart LR

Model["Large Model 🔢"]

-->Node1["DGX Node 1\n8x H100"]

Model

-->Node2["DGX Node 2\n8x H100"]

Model

-->Node3["DGX Node 3\n8x H100"]

Model

-->Node4["DGX Node 4\n8x H100"]

This is multi-node distributed training.

---

Why Vanilla Kubernetes Is Not Enough

Kubernetes knows how to schedule Pods.

But it doesn't know how to run a distributed training job.

The problem has three parts.

---

Problem 1: Pods Start Independently

Normally, Kubernetes starts Pods one at a time as nodes become available.

flowchart TD

PodA["Pod A ✅ — starts immediately"]

PodB["Pod B ✅ — starts when node is ready"]

PodC["Pod C ⏳ — waits for resources"]

PodD["Pod D ⏳ — waits"]

Training jobs cannot work this way.

All workers must start at the same time.

If worker 0 starts and waits for worker 3 which never comes, the job hangs forever.

---

Problem 2: Workers Need to Find Each Other

Each worker needs to know

• How many total workers exist
• What is its own rank (worker index)
• Where is the master (rank 0)?

Kubernetes doesn't set up this coordination automatically.

---

Problem 3: No Retry Logic for Job Failures

If one Pod crashes, Kubernetes restarts it — but with no awareness that the entire training job needs to restart or resume from a checkpoint.

---

The Solution: Kubeflow Training Operator

The Kubeflow Training Operator is a Kubernetes Operator that adds distributed training primitives to Kubernetes.

flowchart LR

User["kubectl apply PyTorchJob"]

-->TrainingOperator["Training Operator"]

TrainingOperator

-->MasterPod["Master Pod\nrank=0"]

TrainingOperator

-->Worker1["Worker Pod\nrank=1"]

TrainingOperator

-->Worker2["Worker Pod\nrank=2"]

TrainingOperator

-->Worker3["Worker Pod\nrank=3"]

It introduces custom resource types:

Resource	Framework
PyTorchJob	PyTorch DDP, FSDP, torchrun
TFJob	TensorFlow distributed
MPIJob	MPI-based training (Horovod)
PaddleJob	PaddlePaddle
JAXJob	JAX distributed

For DGX workloads, PyTorchJob is the most common.

---

PyTorchJob Anatomy

A PyTorchJob defines two roles.

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: llama-training
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      template:
        spec:
          containers:
          - name: trainer
            image: nvcr.io/nvidia/pytorch:24.01-py3
            resources:
              limits:
                nvidia.com/gpu: 8
    Worker:
      replicas: 7
      template:
        spec:
          containers:
          - name: trainer
            image: nvcr.io/nvidia/pytorch:24.01-py3
            resources:
              limits:
                nvidia.com/gpu: 8

This creates

1 Master Pod  × 8 GPUs  =  8 GPUs
7 Worker Pods × 8 GPUs  = 56 GPUs
                          --------
                          64 GPUs total (8 DGX nodes)

The operator automatically injects environment variables into every Pod:

Variable	Value
MASTER_ADDR	Master Pod's IP address
MASTER_PORT	23456
WORLD_SIZE	64 (total GPUs)
RANK	0–63 (this GPU's global index)
LOCAL_RANK	0–7 (this GPU's index within the node)

The training script reads these and initializes the process group:

import torch.distributed as dist

dist.init_process_group(backend="nccl")
# Every process now knows its rank, world_size, and master address

---

Gang Scheduling

The biggest operational challenge with multi-node training is resource fragmentation.

Suppose a cluster has 16 DGX nodes and you need 8 for training.

Two jobs are waiting:

• Job A: needs 8 nodes
• Job B: needs 8 nodes

The cluster has 6 nodes free.

Without gang scheduling:

flowchart TD

JobA["Job A (needs 8 nodes)"]

-->6Nodes["Gets 6 nodes\n⏳ waiting for 2 more"]

JobB["Job B (needs 8 nodes)"]

-->Waiting["Gets 0 nodes\n⏳ stuck behind Job A"]

Job A takes 6 nodes and waits.

Job B gets nothing.

Both are stuck.

---

With gang scheduling (Kueue + Kubernetes):

flowchart TD

JobA["Job A (needs 8)"]

-->Hold["Held in queue"]

JobB["Job B (needs 8)"]

-->Hold

Hold

-->Check["8 nodes free?"]

Check

-->|"Yes"| AllStart["All 8 pods start together"]

Check

-->|"No"| Wait["Wait — don't start partial job"]

The rule is simple:

All pods start together, or none start.

This prevents deadlocks and wasted GPU-hours on partial jobs.

Kueue enforces this through its admission control — a job either gets all requested resources or waits in the queue.

---

How Workers Communicate

Once all pods are running, they need to exchange gradients constantly.

The communication pattern depends on the parallelism strategy.

---

AllReduce — Data Parallelism

In Data Parallel training (PyTorch DDP):

• Each GPU holds a complete copy of the model
• Each GPU processes a different batch of data
• After each forward+backward pass, gradients are averaged across all GPUs

flowchart LR

GPU0["GPU 0\ngrad: 0.3"]

GPU1["GPU 1\ngrad: 0.5"]

GPU2["GPU 2\ngrad: 0.4"]

GPU3["GPU 3\ngrad: 0.2"]

GPU0 & GPU1 & GPU2 & GPU3

-->AllReduce["AllReduce\navg = 0.35"]

AllReduce

-->GPU0 & GPU1 & GPU2 & GPU3

Every GPU ends up with the same averaged gradient.

Every GPU applies the same weight update.

All copies of the model stay synchronized.

---

Send/Receive — Model Parallelism

When the model doesn't fit on one GPU, layers are split across GPUs.

flowchart LR

GPU0["GPU 0\nLayers 1–8"]

-->GPU1["GPU 1\nLayers 9–16"]

-->GPU2["GPU 2\nLayers 17–24"]

-->GPU3["GPU 3\nLayers 25–32"]

Each GPU sends its output to the next GPU (pipeline parallelism).

This requires point-to-point sends, not AllReduce.

---

NCCL

All GPU communication in PyTorch is handled by NCCL (NVIDIA Collective Communications Library).

NCCL chooses the fastest available transport automatically:

flowchart TD

NCCL["NCCL"]

-->IntraNode["Intra-node\nSame DGX node"]

-->NVLink["NVLink / NVSwitch\n900 GB/s"]

NCCL

-->InterNode["Inter-node\nDifferent DGX nodes"]

-->InfiniBand["InfiniBand NDR\n400 Gb/s per link"]

---

Intra-Node: NVLink and NVSwitch

Within a single DGX H100 node, all 8 GPUs are connected via NVSwitch.

flowchart TB

GPU0 & GPU1 & GPU2 & GPU3

-->NVSwitch["NVSwitch\n900 GB/s bisection bandwidth"]

GPU4 & GPU5 & GPU6 & GPU7

-->NVSwitch

Any GPU can communicate with any other GPU at full bandwidth simultaneously.

NVSwitch is the reason DGX nodes can run tensor-parallel training across 8 GPUs efficiently.

A single NVLink connection delivers 900 GB/s total bisection bandwidth.

For comparison, a PCIe Gen5 x16 slot delivers ~64 GB/s — about 14× slower.

---

Inter-Node: InfiniBand

Between DGX nodes, communication travels over InfiniBand (or RoCE for Ethernet-based clusters).

flowchart LR

DGX1["DGX Node 1\n8x H100"]

-->IB["InfiniBand NDR\n400 Gb/s"]

DGX2["DGX Node 2\n8x H100"]

-->IB

DGX3["DGX Node 3\n8x H100"]

-->IB

IB

-->Switch["InfiniBand Switch"]

Each DGX H100 node has 8 InfiniBand NDR ports — one per GPU.

GPUDirect RDMA allows GPUs to read and write from remote GPUs without involving the CPU:

sequenceDiagram

participant GPU0 as GPU 0 (Node 1)
participant RDMA as RDMA NIC
participant Network as InfiniBand
participant GPU4 as GPU 4 (Node 2)

GPU0->>RDMA: DMA transfer (CPU not involved)
RDMA->>Network: Send packet
Network->>GPU4: Direct write to GPU memory

Without GPUDirect RDMA:

GPU → CPU memory → NIC → Network → CPU memory → GPU

With GPUDirect RDMA:

GPU → NIC → Network → NIC → GPU

CPU is bypassed entirely — latency drops by 3–5×.

---

Parallelism Strategies

Large model training combines multiple parallelism types simultaneously.

Strategy	What is split	Communication	Example
Data Parallel (DDP)	Data batches	AllReduce (gradients)	8-GPU single-node
Tensor Parallel (TP)	Individual layers	AllReduce within layer	Across 8 GPUs via NVLink
Pipeline Parallel (PP)	Groups of layers	Send/Recv between stages	Across nodes
ZeRO (FSDP)	Optimizer state + weights	AllGather + ReduceScatter	Scales to 1000s of GPUs
Expert Parallel (EP)	MoE expert layers	AlltoAll	Mixtral, Grok

---

3D Parallelism

DGX SuperPOD deployments often combine TP + PP + DP together.

flowchart TB

subgraph Node1["DGX Node 1"]
GPU0["GPU 0"]
GPU1["GPU 1"]
GPU2["GPU 2"]
GPU3["GPU 3"]
end

subgraph Node2["DGX Node 2"]
GPU4["GPU 4"]
GPU5["GPU 5"]
GPU6["GPU 6"]
GPU7["GPU 7"]
end

GPU0 & GPU1 & GPU2 & GPU3

<-->|"Tensor Parallel\n(NVLink)"| GPU0

Node1

<-->|"Pipeline Parallel\n(InfiniBand)"| Node2

• Tensor Parallel across 4 GPUs within a node (uses NVLink — very fast)
• Pipeline Parallel across nodes (uses InfiniBand — still fast enough)
• Data Parallel across replicas of the above (uses AllReduce over InfiniBand)

The key principle: put the most communication-intensive parallelism on the fastest link.

---

Fault Tolerance and Checkpointing

A training job running for 2 weeks across 64 GPUs has a high probability of hitting at least one failure.

Without fault tolerance:

2 weeks of training → GPU fails on day 13 → restart from scratch

With checkpointing:

Every N steps → save model weights + optimizer state → restart from last checkpoint

---

What Gets Checkpointed

flowchart LR

Checkpoint["Checkpoint"]

-->Weights["Model Weights"]

Checkpoint

-->OptimizerState["Optimizer State\n(Adam momentum, variance)"]

Checkpoint

-->Step["Current Step"]

Checkpoint

-->RNG["RNG State\n(for reproducibility)"]

The optimizer state is often 3× the size of the weights (Adam stores first and second moments + weights).

A 70B parameter model checkpoint in BF16:

Weights:         140 GB
Optimizer state: 420 GB
Total:           560 GB per checkpoint

---

Checkpoint Storage

Checkpoints cannot be stored on a single node's local disk — that node might be the one that fails.

They need a shared distributed filesystem:

flowchart LR

GPU0 & GPU1 & GPU2 & GPU3

-->SharedFS["Shared Storage\n(Lustre / WEKA / NFS)"]

GPU4 & GPU5 & GPU6 & GPU7

-->SharedFS

Each GPU writes its shard of the model to shared storage simultaneously.

---

Elastic Training (torchrun)

PyTorch Elastic (torchrun) allows a job to continue even if a node fails, by dynamically adjusting the world size.

flowchart TD

NormalRun["8 nodes running\nworld_size=64"]

-->NodeFails["Node 3 fails"]

NodeFails

-->Elastic["torchrun detects failure\nre-initializes with 7 nodes\nworld_size=56"]

Elastic

-->Resume["Training resumes\nfrom last checkpoint"]

Without elastic training, one failed node kills the entire job.

With elastic training, the job continues at reduced scale.

---

Complete Training Flow

Putting it all together:

flowchart TB

User["kubectl apply PyTorchJob"]

-->Operator["Training Operator\ncreates Pods"]

Operator

-->Kueue["Kueue\nadmission check"]

Kueue

-->|"All nodes available"| GangStart["Gang schedule\nall pods start together"]

GangStart

-->Init["torchrun initializes\nprocess group via NCCL"]

Init

-->Training["Training loop\nforward + backward"]

Training

-->|"Intra-node gradient sync"| NVLink["NVLink / NVSwitch\n900 GB/s"]

Training

-->|"Inter-node gradient sync"| IB["InfiniBand NDR\n400 Gb/s + GPUDirect RDMA"]

NVLink & IB

-->AllReduce["AllReduce\ngradient averaging"]

AllReduce

-->WeightUpdate["Weight update\nnext step"]

WeightUpdate

-->|"Every N steps"| Checkpoint["Checkpoint\nshared storage"]

WeightUpdate

-->Training

---

Key Takeaways

Component	Role
Kubeflow Training Operator	Manages PyTorchJob lifecycle — creates pods, injects env vars, handles failures
Gang Scheduling (Kueue)	Ensures all pods start together — prevents deadlocks and wasted GPU-hours
NCCL	Handles all GPU communication — automatically uses NVLink or InfiniBand
NVLink / NVSwitch	900 GB/s intra-node bandwidth — enables tensor parallelism within a DGX node
InfiniBand + GPUDirect RDMA	400 Gb/s inter-node bandwidth — CPU bypassed, direct GPU-to-GPU memory transfers
Checkpointing	Saves model + optimizer state to shared storage — enables restart after failure
torchrun / Elastic	Adjusts world size dynamically — job survives node failures without full restart

Multi-node training on Kubernetes is not just about scheduling more Pods. It requires coordinated startup, high-speed communication at every layer of the network, and fault tolerance that spans the entire training run.