PyTorch Training with Kubeflow Training Operator for GB200 NVL72 on RKE2 Cluster

Last Updated: Sept 24, 2025

Introduction

This article demonstrates how to run PyTorch distributed training using the Kubeflow Training Operator on an RKE2 cluster configured for the NVIDIA GB200 NVL72 platform. A full GB200 NVL72 rack consists of 18 interconnected nodes. For demonstration purposes, this guide uses a two-node RKE2 cluster to illustrate the setup and training workflow.

Prerequisites

• RKE2 cluster setup (Kubernetes >= 1.31)
• Access to RKE2 cluster kubeconfig
• Python3 

Step-by-Step Instructions

Kubeflow Trainer is a Kubernetes-native component designed for scalable, distributed AI model training, with built-in support for fine-tuning large language models (LLMs).

1. Deploy the Kubeflow Trainer control plane

   • Apply the control plane manifests to your Kubernetes cluster:

     $ export VERSION=v2.0.0
     $ kubectl apply --server-side -k "https://github.com/kubeflow/trainer.git/manifests/overlays/manager?ref=${VERSION}"

   • Verify that the controller pods are running:

     $ kubectl get pods -n kubeflow-system

   • Expected output:

     NAME                                                   READY   STATUS    RESTARTS   AGE
     jobset-controller-manager-78bcbf6455-c5m8b             1/1     Running   0          2m46s
     kubeflow-trainer-controller-manager-84db68bdff-qmhgm   1/1     Running   0          2m46s

2. Deploy the Kubeflow Trainer Runtimes

   • Install the preconfigured training runtimes:

     $ kubectl apply --server-side -k "https://github.com/kubeflow/trainer.git/manifests/overlays/runtimes?ref=${VERSION}"

   • List available runtimes:

     $ kubectl get clustertrainingruntime
     
     NAME                    AGE
     deepspeed-distributed   84s
     mlx-distributed         83s
     mpi-distributed         83s
     torch-distributed       83s
     torchtune-llama3.2-1b   83s
     torchtune-llama3.2-3b   83s

3. Install Kubeflow SDK

   • Create a virtual environment:

     $ python3 -m venv ~/kubeflow-venv

   • Activate the virtual environment:

     $ source ~/kubeflow-venv/bin/activate

   • Install the Kubeflow SDK from GitHub:

     $ pip install git+https://github.com/kubeflow/sdk.git@64d74db2b6c9a0854e39450d8d1c0201e1e9b3f7#subdirectory=python

4. Create a basic PyTorch training function script 

   • Save the following as train.py:

     def train_pytorch():
         import os
         import torch
         import torch.distributed as dist
         from torch.utils.data import DataLoader, DistributedSampler
         from torchvision import datasets, transforms, models
     
         # [1] Configure CPU/GPU device and distributed backend.
         device, backend = ("cuda", "nccl") if torch.cuda.is_available() else ("cpu", "gloo")
         dist.init_process_group(backend=backend)
         local_rank = int(os.getenv("LOCAL_RANK", 0))
         device = torch.device(f"{device}:{local_rank}")
         
         # [2] Get the pre-defined model.
         model = models.shufflenet_v2_x0_5(num_classes=10)
         model.conv1 = torch.nn.Conv2d(1, 24, kernel_size=3, stride=2, padding=1, bias=False)
         model = torch.nn.parallel.DistributedDataParallel(model.to(device))
         optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
        
         # [3] Get the FashionMNIST dataset and distribute it across all available devices.
         if local_rank == 0: # Download dataset only on local_rank=0 process.
             dataset = datasets.FashionMNIST("./data", train=True, download=True, transform=transforms.Compose([transforms.ToTensor()]))
         dist.barrier()
         dataset = datasets.FashionMNIST("./data", train=True, download=False, transform=transforms.Compose([transforms.ToTensor()]))
         train_loader = DataLoader(dataset, batch_size=100, sampler=DistributedSampler(dataset))
     
         # [4] Define the PyTorch training loop.
         for epoch in range(3):
             for batch_idx, (inputs, labels) in enumerate(train_loader):
                 inputs, labels = inputs.to(device), labels.to(device)
                 # Forward and Backward pass
                 outputs = model(inputs)
                 loss = torch.nn.functional.cross_entropy(outputs, labels)
                 optimizer.zero_grad()
                 loss.backward()
                 optimizer.step()
                 if batch_idx % 10 == 0 and dist.get_rank() == 0:
                     print(f"Epoch {epoch} [{batch_idx * len(inputs)}/{len(train_loader.dataset)}] "
                         f"Loss: {loss.item():.4f}"
                     )

5. (Optional) Adjust the runtime image for GB200 SKU

   • For ARM-based architectures, update the PyTorch runtime image to nvcr.io/nvidia/pytorch:24.01-py3  to ensure compatibility:

     $ kubectl get clustertrainingruntime torch-distributed -o yaml | grep image
     
     image: nvcr.io/nvidia/pytorch:24.01-py3

6. Create and submit a training job submission script 

   • Save the following as submit.py:

     from kubeflow.trainer import TrainerClient, CustomTrainer
     from train import train_pytorch
     
     job_id = TrainerClient().train(
         trainer=CustomTrainer(
             func=train_pytorch,
             num_nodes=2,
             resources_per_node={
                 "cpu": 3,
                 "memory": "64Gi",
                 "gpu": 4,  # Use all 4 GPUs per node
             },
         ),
         runtime=TrainerClient().get_runtime("torch-distributed"),
     )

   • Submit the Job

     $ python3 submit.py

7. Monitor Job Progress 

   • Check the job state:

     $ kubectl get trainjobs -A                                                 
     NAMESPACE   NAME           STATE   AGE
     default     v056cae10f5e           4m26s
     
     $ kubectl get pods -o wide
     NAME                          READY   STATUS    RESTARTS   AGE     IP           NODE                  NOMINATED NODE   READINESS GATES
     v056cae10f5e-node-0-0-jmsm2   1/1     Running   0          5m32s   10.42.2.25   crusoe-rke-worker-1   <none>           <none>
     v056cae10f5e-node-0-1-gsfgk   1/1     Running   0          5m32s   10.42.3.22   crusoe-rke-worker-0   <none>           <none>
     
     $ kubectl get pods -o wide
     NAME                          READY   STATUS      RESTARTS   AGE    IP           NODE                  NOMINATED NODE   READINESS GATES
     v056cae10f5e-node-0-0-jmsm2   0/1     Completed   0          9m4s   10.42.2.25   crusoe-rke-worker-1   <none>           <none>
     v056cae10f5e-node-0-1-gsfgk   0/1     Completed   0          9m4s   10.42.3.22   crusoe-rke-worker-0   <none>           <none>
     
     $ kubectl get trainjobs -A                                                 
     NAMESPACE   NAME           STATE      AGE
     default     v056cae10f5e   Complete   9m35s

8. View Training Logs

   • Retrieve logs from a completed job pod:

     $ kubectl logs v056cae10f5e-node-0-0-jmsm2
     ...
     Epoch 0 [0/60000] Loss: 2.4013
     Epoch 1 [0/60000] Loss: 2.1618
     Epoch 2 [0/60000] Loss: 1.4538

Additional Information

• Crusoe RKE2 GitHub repository
• Official docs - Kubeflow Trainer
• PyTorch on Kubernetes: Kubeflow Training Operator
• NCCL Performance Validation for GB200 NVL72 Rack on RKE2 Cluster