PyTorch Distributed Training

Multi-GPU Single-Node Training with torchrun

When using a pod equipped with 4 or 8 GPUs:

apiVersion: v1
kind: Pod
metadata:
  name: pytorch-ddp-run
  namespace: sci-myproject
spec:
  containers:
    - name: trainer
      image: nvcr.io/nvidia/pytorch:24.04-py3
      command:
        - "torchrun"
        - "--standalone"
        - "--nproc_per_node=4"
        - "train_distributed.py"
      resources:
        limits:
          nvidia.com/gpu: "4"
          cpu: "32"
          memory: "128Gi"
      volumeMounts:
        - name: dshm
          mountPath: /dev/shm
  volumes:
    - name: dshm
      emptyDir:
        medium: Memory
        sizeLimit: 32Gi

Shared Memory (/dev/shm)

PyTorch DataLoader workers and NCCL backends require large shared memory. Always mount an emptyDir with medium: Memory to avoid PyTorch DataLoader crashes.