Training runtimes¶

A training runtime is a named bundle of defaults for a Kubeflow TrainJob: which container image to start, what environment variables to set, what volumes to mount, and how many nodes to use when you don't say otherwise. You reference a runtime by name in the Python SDK (runtime="torch-gh200") or in a YAML manifest (spec.runtimeRef.name: torch-gh200); everything in the runtime is merged with whatever you specify yourself.

On FLAME, the runtime you choose also chooses your GPU type. There is one runtime per GPU class, named torch-<class> (for example torch-gh200). Naming the runtime is how you select the hardware your job runs on. There is deliberately no general-purpose, GPU-agnostic runtime — every TrainJob names a specific class.

Available runtimes¶

FLAME provides one torch-<class> runtime for each GPU class currently installed in the cluster. To see exactly which runtimes exist right now:

kubectl get clustertrainingruntimes

or, from Python:

from kubeflow.trainer import TrainerClient
for rt in TrainerClient().list_runtimes():
    print(rt.name)

Today the cluster has GH200 nodes only, so the single available runtime is:

torch-gh200¶

The PyTorch runtime for FLAME's NVIDIA GH200 nodes. Use it for single-GPU jobs, multi-GPU jobs on one node, and multi-node distributed training on GH200 hardware.

The default image is URCF's own PyTorch image, built and kept up to date for FLAME. It includes PyTorch and the CUDA libraries needed for GH200s. Because it is built for both arm64 and amd64, it works on FLAME's GPU nodes without any extra steps from you.

Environment variables set by the runtime:

Pointing caches at /personal means they survive pod restarts and are shared between notebook sessions and TrainJobs in the same workspace.

Volumes mounted by the runtime:

A memory-backed /dev/shm is mounted automatically. This is required by PyTorch's multiprocessing data loaders (shared memory is how worker processes pass tensors to the training loop). Its size is bounded by the pod's memory limit.

Your home PVC and workspace PVC are also mounted automatically into every TrainJob — this is done cluster-wide, not by the runtime itself, so it applies to plain Jobs too.

How choosing a runtime selects your GPU¶

Each torch-<class> runtime carries a nodeSelector that pins your pods to nodes of that GPU class. That same selector is what Kueue uses to admit your job against the matching resource flavor.

So naming torch-gh200 does two things at once: it lands your pods on GH200 nodes, and it routes the job through the gh200 flavor of your research group's allocation. You never set a nodeSelector or a flavor name yourself — picking the runtime is the whole interface.

The GPU hardware¶

FLAME's GPU nodes are NVIDIA GH200 Grace Hopper Superchips. Each chip pairs an ARM (Aarch64) CPU with an H200 GPU connected over NVLink-C2C. There are currently two GH200 nodes, each with one GPU, for a total of 2 GH200 GPUs in the cluster.

The 480 GB memory capacity makes GH200s well-suited for large models — you can fit models in memory that would require multi-GPU setups on smaller cards.

Queue and scheduling¶

You do not need to set a Kueue queue name on your TrainJob — it is injected automatically. The GPU flavor is selected by the runtime you chose (via its nodeSelector, as described above), and Kueue then schedules your job against your research group's allocation for that flavor.

Specifying GPU resources¶

The runtime does not request any GPUs by default — you must ask for them explicitly. In the Python SDK, use the resources_per_node argument:

from kubeflow.trainer import TrainerClient, CustomTrainer

client = TrainerClient()
job_name = client.train(
    runtime="torch-gh200",
    trainer=CustomTrainer(
        func=my_train_fn,
        resources_per_node={"cpu": 8, "memory": "64Gi", "gpu": 1},  # (1)
    ),
)

1. "gpu": 1 requests one GH200. The cluster has two GPUs total across two nodes, so the maximum for a single-node job is 1; for a two-node job set num_nodes=2 and "gpu": 1 per node to use both.

In a YAML manifest, the equivalent looks like:

apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: my-job
spec:
  runtimeRef:
    name: torch-gh200
  trainer:
    resourcesPerNode:
      requests:
        cpu: "8"
        memory: "64Gi"
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

Using a custom image¶

To override the default image, pass image in the SDK or set spec.trainer.image in YAML. The image must be built for the CPU architecture of the GPU class you chose — GH200 nodes are arm64, so a torch-gh200 job needs an arm64 build. See Building a custom container image for how to produce one through FLAME's CI system.

Inspecting runtimes¶

To see the full definition of a runtime:

kubectl get clustertrainingruntime torch-gh200 -o yaml

To list all available runtimes (one per GPU class):

kubectl get clustertrainingruntimes