Orchestrate distributed training and model serving at scale on Kubernetes.

Distributed Training with Kubeflow

MPIJob for Distributed Training

apiVersion: kubeflow.org/v1
kind: MPIJob
metadata:
  name: distributed-training
spec:
  slotsPerWorker: 8  # GPUs per node
  runPolicy:
    cleanPodPolicy: Running
  mpiReplicaSpecs:
    Launcher:
      replicas: 1
      template:
        spec:
          containers:
          - name: mpi-launcher
            image: my-training-image:latest
            command: [
              "mpirun",
              "-n", "32",  # 4 nodes × 8 GPUs
              "--allow-run-as-root",
              "-bind-to", "none",
              "-map-by", "slot",
              "-x", "NCCL_DEBUG=INFO",
              "python", "distributed_train.py"
            ]
    Worker:
      replicas: 4  # 4 GPU nodes
      template:
        spec:
          containers:
          - name: mpi-worker
            image: my-training-image:latest
            resources:
              limits:
                nvidia.com/gpu: 8

PyTorchJob

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-distributed
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
          - name: pytorch
            image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime
            command: ["python", "-m", "torch.distributed.run",
                      "--nproc_per_node=8", "train.py"]
            resources:
              limits:
                nvidia.com/gpu: 8
    Worker:
      replicas: 7  # Total 8 nodes
      template:
        spec:
          containers:
          - name: pytorch
            image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime
            resources:
              limits:
                nvidia.com/gpu: 8

Model Serving

KServe (formerly KFServing)

apiVersion: "serving.kserve.io/v1beta1"
kind: "InferenceService"
metadata:
  name: "llama-3-8b"
spec:
  predictor:
    minReplicas: 2
    maxReplicas: 8
    scaleTarget: 5  # Requests per replica
    model:
      modelFormat:
        name: pytorch
      storageUri: "s3://my-bucket/models/llama-3-8b/"
      resources:
        limits:
          nvidia.com/gpu: 1
          memory: "20Gi"
        requests:
          nvidia.com/gpu: 1
          memory: "20Gi"