GPU Infrastructure for ML: CUDA, MIG, Kubernetes Device Plugins, and Cost-Efficient Training

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

# You almost never write CUDA kernels directly for ML —
# PyTorch/JAX/TensorFlow do it for you. But understanding the model helps.

import torch

# This operation launches a CUDA kernel:
a = torch.randn(1000, 1000, device='cuda')
b = torch.randn(1000, 1000, device='cuda')
c = torch.matmul(a, b)   # Calls cuBLAS → optimized GEMM kernel

# Profile to see what's actually happening
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    with_stack=True,
) as prof:
    c = torch.matmul(a, b)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

# The most important GPU command
nvidia-smi

# Output:
# +-----------------------------------------------------------------------------+
# | NVIDIA-SMI 545.23.08    Driver Version: 545.23.08    CUDA Version: 12.3    |
# |--------------------------+----------------------+--------------------------|
# | GPU  Name        Persist | Bus-Id        Disp.A | Volatile Uncorr. ECC   |
# | Fan  Temp  Perf  Pwr:Usage/Cap| Memory-Usage      | GPU-Util  Compute M.  |
# |==========================|======================|==========================|
# |   0  NVIDIA H100...  On  | 00000000:01:00.0 Off |                      0 |
# | N/A   35C    P0   100W / 700W |  12345MiB / 81920MiB | 87%      Default  |
# +-----------------------------------------------------------------------------+

# Continuous monitoring (updates every second)
nvidia-smi dmon -s pucvmet

# Per-process GPU usage
nvidia-smi pmon -s um

# Full hardware topology
nvidia-smi topo -m

# Query specific metrics
nvidia-smi --query-gpu=index,name,utilization.gpu,utilization.memory,\
memory.used,memory.total,temperature.gpu,power.draw \
--format=csv,noheader,nounits

# Watch GPU memory in real time
watch -n 0.5 nvidia-smi --query-gpu=memory.used,memory.free --format=csv

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

# Enable MIG mode (requires GPU reset — no live processes)
sudo nvidia-smi -i 0 -mig 1

# Verify MIG is enabled
nvidia-smi -i 0 --query-gpu=mig.mode.current --format=csv,noheader
# Enabled

# Create MIG instances (A100 80GB — 7x 1g.10gb slices)
sudo nvidia-smi mig -cgi 1g.10gb,1g.10gb,1g.10gb,1g.10gb,1g.10gb,1g.10gb,1g.10gb -C

# Or a mixed configuration (for different workload sizes)
# 1x 3g.40gb + 2x 1g.10gb + 1x 2g.20gb = full GPU
sudo nvidia-smi mig -cgi 3g.40gb,1g.10gb,1g.10gb,2g.20gb -C

# List MIG instances
nvidia-smi mig -lgi
# +-------------------------------------------------------+
# | GPU instances:                                        |
# | GPU   Name          Profile  Instance   Placement     |
# |       ID                     ID         Start  Size   |
# |=======================================================|
# |   0  MIG 1g.10gb    19       1          0      1      |
# |   0  MIG 1g.10gb    19       2          1      1      |
# ...

# List compute instances within MIG instances
nvidia-smi mig -lci

# Delete all MIG instances (returns to full GPU mode)
sudo nvidia-smi mig -dci && sudo nvidia-smi mig -dgi

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

# nvidia-mig-parted for declarative MIG configuration
# Install
curl -fsSL https://nvidia.github.io/nvidia-mig-parted/gpus/install.sh | sh

# config.yaml
cat > mig-config.yaml <<'EOF'
version: v1
mig-configs:
  all-1g.10gb:
    - devices: all
      mig-enabled: true
      mig-devices:
        "1g.10gb": 7

  mixed-workloads:
    - devices: [0]
      mig-enabled: true
      mig-devices:
        "3g.40gb": 1
        "2g.20gb": 1
        "1g.10gb": 2
    - devices: [1, 2, 3]
      mig-enabled: true
      mig-devices:
        "1g.10gb": 7
EOF

# Apply a configuration
nvidia-mig-parted apply -f mig-config.yaml -c all-1g.10gb

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# Install via Helm
helm repo add nvdp https://nvidia.github.io/k8s-device-plugin
helm repo update

helm upgrade --install nvdp nvdp/nvidia-device-plugin \
  --namespace nvidia-device-plugin \
  --create-namespace \
  --version=0.14.5 \
  --set gfd.enabled=true \
  --set migStrategy=mixed    # Support both MIG and non-MIG nodes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# Request a full GPU
apiVersion: v1
kind: Pod
metadata:
  name: gpu-job
spec:
  containers:
    - name: cuda-job
      image: nvidia/cuda:12.3-runtime-ubuntu22.04
      command: ["python", "train.py"]
      resources:
        limits:
          nvidia.com/gpu: 1    # Request 1 full GPU

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# Request a MIG slice (1g.10gb = 10GB, 1/7 of an A100)
apiVersion: v1
kind: Pod
metadata:
  name: inference-pod
spec:
  containers:
    - name: inference
      image: myorg/inference:latest
      resources:
        limits:
          nvidia.com/mig-1g.10gb: 1    # Request a 1g.10gb MIG slice

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# Node labels added by GFD:
kubectl get node gpu-node-1 -o json | jq '.metadata.labels | with_entries(select(.key | startswith("nvidia")))'

# {
#   "nvidia.com/cuda.driver.major": "545",
#   "nvidia.com/cuda.runtime.major": "12",
#   "nvidia.com/gpu.count": "8",
#   "nvidia.com/gpu.memory": "81920",
#   "nvidia.com/gpu.product": "NVIDIA-H100-80GB-HBM3",
#   "nvidia.com/mig.capable": "true",
#   "nvidia.com/mig.strategy": "mixed"
# }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# Schedule only on H100 nodes with at least 4 GPUs
spec:
  nodeSelector:
    nvidia.com/gpu.product: NVIDIA-H100-80GB-HBM3
  affinity:
    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
        nodeSelectorTerms:
          - matchExpressions:
              - key: nvidia.com/gpu.count
                operator: Gt
                values: ["3"]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

# ConfigMap for the device plugin — enable time-slicing
apiVersion: v1
kind: ConfigMap
metadata:
  name: time-slicing-config
  namespace: nvidia-device-plugin
data:
  config.yaml: |
    version: v1
    flags:
      migStrategy: none
    sharing:
      timeSlicing:
        renameByDefault: false
        failRequestsGreaterThanOne: false
        resources:
          - name: nvidia.com/gpu
            replicas: 4    # Each physical GPU appears as 4 schedulable GPUs

1
2
3
4

# Apply the config
helm upgrade nvdp nvdp/nvidia-device-plugin \
  --namespace nvidia-device-plugin \
  --set config.name=time-slicing-config

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# Now 4 pods can each request 1 GPU and share the physical hardware
# Each pod gets ~25% of GPU time — no memory isolation (unlike MIG)
apiVersion: apps/v1
kind: Deployment
spec:
  replicas: 4
  template:
    spec:
      containers:
        - name: inference
          resources:
            limits:
              nvidia.com/gpu: 1   # Each gets a time-slice

1
2
3
4
5

# Install DCGM Exporter
helm repo add gpu-helm-charts https://nvidia.github.io/dcgm-exporter/helm-charts
helm install --generate-name \
  gpu-helm-charts/dcgm-exporter \
  --namespace monitoring

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

# Key metrics to alert on:
# DCGM_FI_DEV_GPU_UTIL          — GPU compute utilization %
# DCGM_FI_DEV_MEM_COPY_UTIL     — Memory bandwidth utilization %
# DCGM_FI_DEV_FB_USED           — Framebuffer (VRAM) used bytes
# DCGM_FI_DEV_POWER_USAGE       — Power draw in watts
# DCGM_FI_DEV_GPU_TEMP          — Temperature in Celsius
# DCGM_FI_DEV_XID_ERRORS        — Hardware errors (non-zero = problem)
# DCGM_FI_DEV_NVLINK_BANDWIDTH_TOTAL — NVLink utilization

# Prometheus alerting rules
apiVersion: monitoring.coreos.com/v1
kind: PrometheusRule
metadata:
  name: gpu-alerts
  namespace: monitoring
spec:
  groups:
    - name: gpu
      rules:
        - alert: GPUHighTemperature
          expr: DCGM_FI_DEV_GPU_TEMP > 85
          for: 5m
          labels:
            severity: warning
          annotations:
            summary: "GPU {{ $labels.gpu }} temperature {{ $value }}°C"

        - alert: GPULowUtilization
          expr: avg_over_time(DCGM_FI_DEV_GPU_UTIL[30m]) < 20
          for: 1h
          labels:
            severity: info
          annotations:
            summary: "GPU {{ $labels.gpu }} utilization below 20% for 1h — possible waste"

        - alert: GPUHardwareError
          expr: DCGM_FI_DEV_XID_ERRORS > 0
          for: 0m
          labels:
            severity: critical
          annotations:
            summary: "GPU {{ $labels.gpu }} hardware error detected"

        - alert: GPUMemoryPressure
          expr: |
            DCGM_FI_DEV_FB_USED / (DCGM_FI_DEV_FB_USED + DCGM_FI_DEV_FB_FREE) > 0.95
          for: 10m
          labels:
            severity: warning
          annotations:
            summary: "GPU memory > 95% used"

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

# train.py — checkpointing for spot-tolerant training
import os
import torch
import torch.distributed as dist
from pathlib import Path

CHECKPOINT_DIR = Path(os.environ.get("CHECKPOINT_DIR", "/checkpoints"))
CHECKPOINT_INTERVAL = int(os.environ.get("CHECKPOINT_INTERVAL", "500"))  # steps

def save_checkpoint(model, optimizer, step, loss, path):
    """Save training state for spot instance recovery."""
    torch.save({
        'step': step,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
    }, path)
    print(f"Checkpoint saved at step {step} → {path}")

def load_checkpoint(model, optimizer, path):
    """Resume from checkpoint after preemption."""
    if not path.exists():
        return 0
    checkpoint = torch.load(path, map_location='cuda')
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_step = checkpoint['step']
    print(f"Resumed from step {start_step}")
    return start_step

def train():
    model = MyModel().cuda()
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    checkpoint_path = CHECKPOINT_DIR / "latest.pt"
    start_step = load_checkpoint(model, optimizer, checkpoint_path)

    for step, batch in enumerate(dataloader, start=start_step):
        loss = model(batch)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if step % CHECKPOINT_INTERVAL == 0:
            save_checkpoint(model, optimizer, step, loss.item(), checkpoint_path)
            # Also save a named checkpoint periodically
            if step % (CHECKPOINT_INTERVAL * 10) == 0:
                save_checkpoint(model, optimizer, step, loss.item(),
                                CHECKPOINT_DIR / f"step_{step}.pt")

if __name__ == "__main__":
    train()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

# Kubernetes Job with spot tolerance and checkpoint volume
apiVersion: batch/v1
kind: Job
metadata:
  name: training-job
spec:
  template:
    spec:
      restartPolicy: OnFailure    # Restart on spot preemption
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
        - key: karpenter.sh/capacity-type
          value: spot
          effect: NoSchedule
      containers:
        - name: trainer
          image: myorg/trainer:latest
          command: ["python", "train.py"]
          env:
            - name: CHECKPOINT_DIR
              value: /checkpoints
          resources:
            limits:
              nvidia.com/gpu: 8
              memory: 256Gi
              cpu: 64
          volumeMounts:
            - name: checkpoints
              mountPath: /checkpoints
            - name: dshm          # Shared memory for DataLoader workers
              mountPath: /dev/shm
      volumes:
        - name: checkpoints
          persistentVolumeClaim:
            claimName: training-checkpoints   # Survives pod restarts
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 32Gi

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

# train_distributed.py — DDP (DistributedDataParallel) training
import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup():
    """Initialize distributed training."""
    dist.init_process_group(backend="nccl")   # NCCL for GPU-to-GPU
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def cleanup():
    dist.destroy_process_group()

def train():
    local_rank = setup()

    # Each process gets its own GPU
    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # Each process gets a different shard of data
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, sampler=sampler, batch_size=32)

    optimizer = torch.optim.AdamW(model.parameters())

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)   # Shuffle differently each epoch
        for batch in dataloader:
            batch = batch.cuda(local_rank)
            loss = model(batch)
            loss.backward()
            # DDP automatically all-reduces gradients across GPUs
            optimizer.step()
            optimizer.zero_grad()

    cleanup()

# Launch with torchrun (replaces the old torch.distributed.launch)
# torchrun --nproc_per_node=8 --nnodes=4 --node_rank=0 \
#   --master_addr=10.0.0.1 --master_port=29500 \
#   train_distributed.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

# Kubernetes Job for multi-node distributed training
# Using volcano or kubeflow for gang scheduling
apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata:
  name: distributed-training
spec:
  minAvailable: 4          # All 4 nodes must be available (gang scheduling)
  schedulerName: volcano
  plugins:
    env: []
    svc: []
  tasks:
    - replicas: 4          # 4 nodes × 8 GPUs = 32 GPUs total
      name: worker
      template:
        spec:
          containers:
            - name: trainer
              image: myorg/trainer:latest
              command:
                - torchrun
                - --nproc_per_node=8
                - --nnodes=4
                - --node_rank=$(VK_TASK_INDEX)
                - --master_addr=$(distributed-training-worker-0.distributed-training)
                - --master_port=29500
                - train_distributed.py
              resources:
                limits:
                  nvidia.com/gpu: 8
              env:
                - name: NCCL_DEBUG
                  value: INFO
                - name: NCCL_IB_DISABLE
                  value: "0"      # Enable InfiniBand if available

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

# Automatic Mixed Precision with PyTorch
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()   # Handles gradient scaling to prevent underflow

for batch in dataloader:
    optimizer.zero_grad()

    # Forward pass in FP16/BF16
    with autocast(dtype=torch.bfloat16):   # BF16 preferred on A100/H100
        output = model(batch)
        loss = criterion(output, labels)

    # Scale gradients, backward, unscale before clip
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

# Typical speedup: 2-3x vs FP32 on A100/H100
# Memory reduction: ~50% vs FP32 → larger batch sizes possible

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

from torch.utils.checkpoint import checkpoint_sequential

# Without gradient checkpointing: store all activations (high memory)
output = model(x)

# With gradient checkpointing: recompute activations during backward
# ~30% compute overhead, ~60-70% memory savings
output = checkpoint_sequential(model.layers, segments=4, input=x)

# Or per-layer with transformers
from transformers import AutoModel
model = AutoModel.from_pretrained("meta-llama/Llama-3-8b")
model.gradient_checkpointing_enable()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

# Standard attention: O(n²) memory in HBM
# FlashAttention: O(n) memory, same result, 2-4x faster

# With PyTorch 2.0+ (built-in scaled dot product attention)
with torch.backends.cuda.sdp_kernel(
    enable_flash=True,
    enable_math=False,
    enable_mem_efficient=False
):
    output = torch.nn.functional.scaled_dot_product_attention(
        query, key, value,
        attn_mask=None,
        dropout_p=0.0,
        is_causal=True   # For causal/decoder attention
    )
# PyTorch automatically uses FlashAttention when available

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

# Scale GPU node pools based on pending GPU requests
apiVersion: keda.sh/v1alpha1
kind: ScaledObject
metadata:
  name: gpu-workload-scaler
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: gpu-job-controller
  minReplicaCount: 0
  maxReplicaCount: 10
  triggers:
    - type: kubernetes-workload
      metadata:
        podSelector: "queue=gpu-training"
        value: "1"   # 1 node per pending pod

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

# Average GPU utilization across all nodes (goal: keep > 70%)
avg(DCGM_FI_DEV_GPU_UTIL)

# Per-job GPU utilization (requires pod labels)
avg by (pod_name) (
  DCGM_FI_DEV_GPU_UTIL * on(node) group_left(pod_name)
  kube_pod_info{namespace="ml-training"}
)

# GPU memory efficiency
avg(DCGM_FI_DEV_FB_USED) /
avg(DCGM_FI_DEV_FB_USED + DCGM_FI_DEV_FB_FREE)

# Training throughput (samples/second) from custom metrics
rate(training_samples_processed_total[5m])

# Cost per training sample (requires cost metadata)
(rate(training_samples_processed_total[5m])) /
  on(node) group_left() (gpu_cost_per_hour / 3600)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

# 1. Check actual GPU utilization
nvidia-smi dmon -s u -d 60 | awk '{print $5}' | \
  awk 'NR>1{sum+=$1; count++} END{print "Avg util:", sum/count "%"}'

# 2. Find idle GPU pods (running but low utilization)
kubectl get pods -A -l nvidia.com/gpu -o wide | \
  while read pod; do
    # Cross-reference with DCGM metrics to find low-utilization pods
    echo $pod
  done

# 3. Check for GPU memory fragmentation
nvidia-smi --query-compute-apps=pid,used_memory --format=csv

# 4. Profile a training job
nsys profile --trace=cuda,nvtx python train.py
nsys stats report.nsys-rep

# 5. Check NVLink utilization for multi-GPU jobs
nvidia-smi nvlink --status -i 0

# 6. Verify jobs are actually using Tensor Cores
ncu --metrics sm__pipe_tensor_op_hmma_cycles_active.avg.pct_of_peak_sustained_active \
    python -c "import torch; a=torch.randn(1000,1000,device='cuda',dtype=torch.float16); torch.matmul(a,a)"

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

# GPU status and monitoring
nvidia-smi                                    # Summary
nvidia-smi dmon -s pucvmet                   # Continuous stats
nvidia-smi --query-gpu=utilization.gpu,memory.used,temperature.gpu \
  --format=csv -l 1                          # Metrics every second

# MIG management
nvidia-smi -i 0 -mig 1                      # Enable MIG
nvidia-smi mig -cgi 1g.10gb,1g.10gb -C     # Create instances
nvidia-smi mig -lgi                          # List instances
nvidia-smi mig -dci && nvidia-smi mig -dgi  # Delete all

# Kubernetes GPU scheduling
kubectl describe node <gpu-node> | grep nvidia  # Available GPU resources
kubectl get pods -A --field-selector=spec.nodeName=<node> | \
  grep -v Running                               # Non-running GPU pods

# DCGM metrics
kubectl port-forward -n monitoring svc/dcgm-exporter 9400:9400
curl localhost:9400/metrics | grep DCGM_FI_DEV_GPU_UTIL

# PyTorch GPU diagnostics
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.device_count())"
python -c "import torch; print(torch.cuda.get_device_properties(0))"
python -c "import torch; torch.cuda.memory_summary(device=0, abbreviated=False)"