Hyperparameter Tuning (Katib)

Katib is the automated hyperparameter tuning system built into Kubeflow. It runs multiple trials in parallel, guided by a search algorithm (Bayesian optimisation, random search, grid search), and identifies the best hyperparameter configuration.

Supported algorithms

Algorithm	Best for
random	Fast exploration, many parameters
bayesianoptimization	Expensive trials, continuous params
grid	Small discrete spaces
hyperband	Early stopping to reduce wasted compute
cmaes	Continuous numerical optimisation

Define an Experiment

# katib-experiment.yaml
apiVersion: kubeflow.org/v1beta1
kind: Experiment
metadata:
  name: bert-lr-sweep
  namespace: rg-compsci
spec:
  objective:
    type: maximize
    goal: 0.92
    objectiveMetricName: val_accuracy
    additionalMetricNames:
      - val_loss
  algorithm:
    algorithmName: bayesianoptimization
  parallelTrialCount: 2          # concurrent trials (GPU constrained)
  maxTrialCount: 20
  maxFailedTrialCount: 3
  parameters:
    - name: learning_rate
      parameterType: double
      feasibleSpace:
        min: "1e-5"
        max: "1e-3"
    - name: batch_size
      parameterType: int
      feasibleSpace:
        list: ["16", "32", "64"]
    - name: lora_r
      parameterType: int
      feasibleSpace:
        list: ["8", "16", "32"]
  trialTemplate:
    primaryContainerName: training-container
    trialParameters:
      - name: learning_rate
        description: Learning rate
        reference: learning_rate
      - name: batch_size
        reference: batch_size
      - name: lora_r
        reference: lora_r
    trialSpec:
      apiVersion: batch/v1
      kind: Job
      spec:
        template:
          spec:
            containers:
              - name: training-container
                image: nvcr.io/nvidia/pytorch:24.08-py3
                command:
                  - python
                  - /workspace/train.py
                  - --learning-rate=${trialParameters.learning_rate}
                  - --batch-size=${trialParameters.batch_size}
                  - --lora-r=${trialParameters.lora_r}
                resources:
                  limits:
                    nvidia.com/gpu: "1"
            restartPolicy: Never

Training script for Katib

Your training script must print metrics in the format Katib expects:

# train.py
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--learning-rate", type=float)
parser.add_argument("--batch-size", type=int)
parser.add_argument("--lora-r", type=int)
args = parser.parse_args()

# ... training code ...

val_accuracy = 0.87  # replace with actual metric
val_loss = 0.32

# Katib reads these from stdout
print(f"val_accuracy={val_accuracy}")
print(f"val_loss={val_loss}")

Submit the experiment

kubectl apply -f katib-experiment.yaml -n rg-compsci

Monitor in the Kubeflow UI: Experiments (AutoML) → select your experiment → view trial status and metrics.

Retrieve best parameters

from kubeflow.katib import KatibClient

client = KatibClient(namespace="rg-compsci")
experiment = client.get_experiment("bert-lr-sweep")

best = experiment.status.current_optimal_trial
print("Best parameters:")
for param in best.parameter_assignments:
    print(f"  {param.name}: {param.value}")
print(f"Best val_accuracy: {best.observation.metrics[0].latest}")