1.6x. That was the speedup I measured the first time I spread the assistant’s tune from one GPU to four across two nodes, and I had budgeted for something close to 4x. A distributed training job on a Kubernetes cluster hides the reason a run is slow as neatly as a pipeline hides a skipped step, and the reason here was not the Training Operator, not Kueue, and not the code. It was the network between the two nodes. This part takes the single GPU tune from last part and spreads it with the OpenShift AI distributed training stack, and most of the work is learning where that stack helps and where it quietly does not.

Picking up from a single GPU tune

Last part wrapped the InstructLab retrain into a pipeline whose tune step ran on one L40S for roughly 220 minutes, about ninety percent of the whole run. That is fine at night when nobody waits on it, and it stops scaling the moment two teams want their own nightly tune on the same cluster. This part replaces the single GPU tune with a distributed one that fans the work across several GPUs, and threads it through the same project namespace the pipeline already uses. Same Granite 3.1 8B, same synthetic data from Part 9, same evaluation gate. What changes is that one process on one GPU becomes several processes on several GPUs, kept in step by a gradient sync every training step.

Red Hat did not invent this layer, and it helps to name what is upstream. Distributed training on OpenShift AI runs on the Kubeflow Training Operator, an open source controller, with Kueue, another open source project, doing admission and quota. Red Hat’s contribution is the operator packaging, tested container images built on a known torch and CUDA, and a support line when a run wedges at 2am. Fine tuning itself is one of three ways to change a model’s behaviour, and if you are still deciding between tuning, retrieval and prompting, that trade off sits in fine tuning vs RAG vs prompting. We already chose tuning for this project back in Part 10, so here the question is only how to make it finish faster.

Data parallel training in one page

Distributed data parallel, DDP, is the shape almost every tuning job uses. Each GPU holds a full copy of the model and trains on a different slice of the batch. After each backward pass, every GPU exchanges its gradients with the others and averages them, an all reduce, so every copy applies the same update and the models never drift apart. NCCL, the NVIDIA collective library, is what moves those gradients between GPUs, over NVLink inside a node and over the network between nodes. That last sentence is the whole performance story of this part in miniature: NVLink is fast, the network usually is not.

When a model plus its optimizer state does not fit on one GPU, DDP is not enough, because DDP needs a full copy per GPU. That is where Fully Sharded Data Parallel, FSDP, comes in: it shards the parameters, gradients and optimizer state across GPUs so each holds only a fraction, at the cost of more communication. A tuned 8B model in bf16 needs roughly 16 GB for weights, and the Adam optimizer state adds another two copies, so full fine tuning lands near 90 GB and does not fit one 48 GB L40S. FSDP across two or four GPUs is what makes it fit at all. The operator does not care which you use; DDP and FSDP are both just your training script, and the operator’s job is to launch that script as a coordinated group of pods.

Training Operator, Kueue and how a job gets scheduled

A PyTorchJob is the custom resource the Training Operator watches. You describe a Master replica and some Worker replicas, each a pod template with an image, a command and a GPU request, and the operator creates the pods, injects the environment variables torchrun needs to find its peers, and reports the job state. Before any pod starts, Kueue decides whether the cluster can afford it. Kueue holds a ClusterQueue with a nominal GPU quota and a per project LocalQueue that draws from it, and a job that asks for more than the quota is Suspended, not rejected, until room frees up. That admission gate is what stops one team’s tuning run from starving the serving pods, and it is also the single most common reason a job appears to hang.

flowchart LR
  A[create_job in workbench] --> B[PyTorchJob resource]
  B --> C[Training Operator reconciles]
  C --> D[Kueue checks quota]
  D -->|admitted| E[Master and Worker pods]
  D -->|over quota| S[Suspended, waits]
  E --> F[torchrun forms process group]
  F --> G[NCCL all reduce each step]
  G --> H[checkpoint to PVC]

Set the quota up first, because a job submitted before its LocalQueue exists never admits. Create a ClusterQueue that caps the GPUs this cluster will hand to tuning, a LocalQueue in the project that points at it, and label the namespace for Kueue management. Nothing here holds a secret, so it goes straight into version control.

# team-quota.yaml, the GPUs this cluster will lend to tuning jobs
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: support-cq
spec:
  namespaceSelector: {}
  resourceGroups:
  - coveredResources: ['cpu', 'memory', 'nvidia.com/gpu']
    flavors:
    - name: default-flavor
      resources:
      - name: 'cpu'
        nominalQuota: '32'
      - name: 'memory'
        nominalQuota: 128Gi
      - name: 'nvidia.com/gpu'
        nominalQuota: '4'
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  name: team-support
  namespace: support-assistant
spec:
  clusterQueue: support-cq

$ oc apply -f team-quota.yaml
clusterqueue.kueue.x-k8s.io/support-cq created
localqueue.kueue.x-k8s.io/team-support created
$ oc label namespace support-assistant kueue.openshift.io/managed=true --overwrite
namespace/support-assistant labeled

Running a PyTorchJob with the Training Operator SDK

You can write the PyTorchJob YAML by hand, and the docs show it, but from a workbench the SDK is cleaner. The create_job call takes your training function, a base image, the number of workers and the GPUs each needs, and the Kueue queue name, and it builds and submits the resource for you. Structuring that training function so it runs unchanged on the cluster is the same discipline as turning a notebook into a package, which is worth reading in from notebook to Python package. The credentials come from the workbench service account, not a typed token.

# Tested on Red Hat OpenShift AI 3.4 self managed, Training Operator v1 PyTorchJob
# API kubeflow.org/v1, Kueue on the cluster, kubeflow-training SDK
# base image odh-training-cuda128-torch28-py312-rhel9:v3.0 (torch 2.8, CUDA 12.8, py3.12)
from kubeflow.training import TrainingClient

def train_func():
    import os, torch, torch.distributed as dist
    dist.init_process_group(backend='nccl')
    rank, world = dist.get_rank(), dist.get_world_size()
    torch.cuda.set_device(int(os.environ['LOCAL_RANK']))
    if rank == 0:
        print(f'world size {world}, tuning granite-3.1-8b')
    # FSDP wrap, load the SDG dataset from Part 9, run the multi phase tune here
    dist.barrier()
    dist.destroy_process_group()

client = TrainingClient()   # in-cluster config from the workbench service account
client.create_job(
    name='granite-retrain',
    train_func=train_func,
    base_image='registry.redhat.io/rhoai/odh-training-cuda128-torch28-py312-rhel9:v3.0',
    num_workers=2,
    resources_per_worker={'nvidia.com/gpu': '2'},
    env_vars={'NCCL_DEBUG': 'INFO', 'TORCH_DISTRIBUTED_DEBUG': 'DETAIL'},
    labels={'kueue.x-k8s.io/queue-name': 'team-support'},
)

# the operator creates the resource, torchrun sets world size across the pods
INFO:kubeflow.training:PyTorchJob support-assistant/granite-retrain created
$ oc get pytorchjob granite-retrain -n support-assistant
NAME              STATE     AGE
granite-retrain   Running   3m
# rank 0 log, once pods are admitted and the group forms
world size 4, tuning granite-3.1-8b

Watching the run is one more call. get_job_logs streams the rank 0 pod by default, and you pass a worker rank when you are chasing an imbalance between GPUs. Where the checkpoint lands matters more than people expect: a pod filesystem is ephemeral, so a run that trains for an hour and then writes its checkpoint to local disk loses it the moment the pod exits. Mount a PVC or write to the same S3 bucket the pipeline already uses, and gate the write on rank 0 so four processes do not race to save the same file.

# follow the rank 0 pod while the tune runs
logs = client.get_job_logs('granite-retrain', follow=True)
# [rank0] step  400  loss 0.912  tok/s 1180
# [rank0] step  800  loss 0.741  tok/s 1187
# [rank0] saved checkpoint to /mnt/models/granite-retrain/epoch-1

One quiet gotcha lives in that call. Setting num_workers=2 does not mean two GPUs; it means two Worker pods, and the operator also creates one Master replica. With two GPUs requested per replica, the real ask is three pods times two GPUs, six GPUs, not four. That off by one on the Master is what turned my first overnight run into nothing.

Where a distributed job stalls

The overnight run that did not run. I submitted the job above, saw no error, and went home. Next morning the PyTorchJob still read Suspended, eight hours gone, and I very nearly filed a bug against the Training Operator. It was Kueue doing its job correctly. My real request was six GPUs, counting the Master, against a ClusterQueue nominal quota of four, so Kueue held the job rather than overcommit the cluster. One command told me what a night of reading operator logs had not.

$ oc get pytorchjob granite-retrain -n support-assistant
NAME              STATE       AGE
granite-retrain   Suspended   8h
$ oc get workloads -n support-assistant
NAME                          QUEUE          ADMITTED   AGE
pytorchjob-granite-retrain-x  team-support   False      8h
$ oc get workload pytorchjob-granite-retrain-x -n support-assistant 
    -o jsonpath='{.status.conditions[0].message}'
couldn't assign flavors to pod set worker: insufficient quota
for nvidia.com/gpu in flavor default-flavor, request > maximum
capacity (6 > 4)

Two fixes, and the choice matters. Raising the ClusterQueue nominal quota to six lets the six GPU job admit, which is right if the cluster has them to spare. Dropping to num_workers=1 with two GPUs per replica, giving four GPUs total counting the Master, keeps the job inside a four GPU quota and, as the next section shows, is often faster anyway. I raised the quota that morning, then walked it back a week later once I had measured the scaling. Keep this lookup next to the workbench; these six symptoms cover almost every stalled run I have hit on this stack.

When multi node is the wrong call

Here is where the product docs and most tutorials will lead you astray, not by being wrong but by being silent. Multi node training is the headline feature, so the natural instinct is to reach for it, and for an 8B tune that instinct usually costs you time. Gradient sync happens every step, and its cost is set by the slowest link the gradients cross. Inside one node, GPUs talk over NVLink at hundreds of GB per second. Between two nodes on ordinary 10 or 25 GbE, that same all reduce runs one to two orders of magnitude slower, and for a small model where compute per step is modest, communication starts to dominate. The chart below is the assistant’s tune measured four ways, all against the 220 minute single GPU baseline from last part.

Before you add a single GPU, there is a cheaper lever worth pulling. Gradient accumulation lets one GPU stand in for a larger batch by summing gradients over several micro batches before it steps, trading a little wall clock for memory and needing no second node, no extra Kueue quota and no NCCL at all. For the assistant it took the effective batch from 8 to 32 on the same L40S and closed most of the quality gap I had been reaching for more GPUs to fix. Scale out when accumulation runs out of road, not before it.

One forward looking note, because the API is moving under you. OpenShift AI 3.3 added Kubeflow Trainer v2, which replaces the per framework PyTorchJob with a single TrainJob resource and pre built ClusterTrainingRuntime templates you reference instead of hand rolling pod specs. It is the cleaner long term shape and worth reading before you standardise. As of 3.4 the Training Operator v1 PyTorchJob shown here is still the supported default for most tuning work, so this part targets it, but if you are starting fresh, read the TrainJob chapter first so you do not build on the older API by habit.

Keep the tune on one node until it does not fit

References

• Red Hat OpenShift AI, running Training Operator PyTorchJob distributed workloads
• Red Hat OpenShift AI, Kubeflow Trainer v2 TrainJob and ClusterTrainingRuntime
• Red Hat Developer, fine tune LLMs with Kubeflow Trainer on OpenShift AI
• Kubeflow, PyTorchJob user guide, upstream Training Operator

About The Author

Dr. Pranay Jha

Dr. Pranay Jha is a Cloud and AI Consultant with 18+ years of experience in hybrid cloud, virtualization, and enterprise infrastructure transformation. He specializes in VMware technologies, multi-cloud strategy, and Generative AI solutions. He holds a PhD in Computer Applications with research focused on Cloud and AI, has published multiple research papers, and has been a VMware vExpert since 2016 and a VMUG Community Leader.

See author's posts