PyTorch Distributed Training

Multi-GPU Single-Node Training with torchrun

For batch deep learning workloads, running torchrun inside a Kubernetes Job is recommended over a bare Pod. A Job provides automatic retry handling on failure and automatic cleanup after completion (ttlSecondsAfterFinished).

Strictly Prohibited: No sleep infinity in Training Jobs

Training Jobs must exit automatically as soon as torchrun finishes. Never run sleep infinity, tail -f /dev/null, or append ; sleep infinity after your training command to hold multi-GPU nodes open. Save checkpoints to a persistent ceph-rbd or cephfs PVC and let the container exit cleanly.

When running on a node equipped with multiple GPUs (such as 4x A40s via nvidia.com/gpu, 4x A100s via nvidia.com/a100, or 8x H100s via nvidia.com/h100), use torchrun --standalone with --nproc_per_node:

apiVersion: batch/v1
kind: Job
metadata:
  name: pytorch-ddp-job
  namespace: my-project
spec:
  backoffLimit: 2                   # Retry up to 2 times on unhandled error
  ttlSecondsAfterFinished: 86400    # Automatically delete completed pod after 24 hours
  template:
    spec:
      restartPolicy: OnFailure
      containers:
        - name: trainer
          image: nvcr.io/nvidia/pytorch:24.04-py3
          command:
            - "torchrun"
            - "--standalone"
            - "--nproc_per_node=4"
            - "train_distributed.py"
          resources:
            requests:
              cpu: "32"
              memory: "128Gi"
              nvidia.com/gpu: "4"
            limits:
              cpu: "32"
              memory: "128Gi"
              nvidia.com/gpu: "4"
          volumeMounts:
            - name: dshm
              mountPath: /dev/shm
      volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 32Gi

Shared Memory (/dev/shm)

PyTorch DataLoader multi-processing workers and NCCL collective communications require substantial shared memory. By default, Kubernetes mounts a 64MB /dev/shm partition, which will cause PyTorch to crash with Bus error. Always mount an emptyDir volume with medium: Memory to allocate shared memory from RAM.

---

Multi-Node Distributed Training

Single-node torchrun --standalone scales up to the total number of GPUs available on a single physical node (e.g., up to 4 GPUs on A100 nodes, or up to 8 GPUs on the flagship H100 node).

When training requires scaling across multiple physical nodes (e.g., 16+ GPUs):

• PyTorch requires master rendezvous configuration (--nnodes, --node_rank, --master_addr, --master_port) and inter-pod synchronization.
• On Berkelium, the recommended approach for multi-node workloads is using Distributed Ray Clusters (KubeRay) with RayTrain, which automates node discovery, cluster scaling, and fault tolerance across nodes.