How do I manage multi-GPU setups for deep learning?

Multi-GPU deep learning works when the hardware, software stack, and training strategy line up. Miss one and you’ll waste cycles debugging throughput instead of training models.

1. Choose the Right Hardware

GPU Selection:
- Pick GPUs built for deep learning: NVIDIA A100, RTX 4090, or RTX 3090/3080.
- VRAM must fit your largest model plus activations and optimizer state.
Interconnect:
- High-speed links — NVLink, InfiniBand, or PCIe Gen4 — keep GPUs fed.
- NVLink shines when GPUs exchange gradients or activations every step.
Power and Cooling:
- Size the PSU for peak draw across all GPUs plus CPU, drives, and fans.
- Thermal throttling kills clock speeds; plan airflow or liquid cooling before you rack the box.

2. Optimize the Software Stack

GPU Drivers: Install the latest NVIDIA driver branch that supports your GPU generation.
CUDA Toolkit: Match the CUDA version to your framework and GPU architecture. Mismatches cause silent kernel falls back to slow paths.
cuDNN: Install the cuDNN version pinned by your framework release notes.
Deep Learning Frameworks:
- TensorFlow, PyTorch, and JAX all support multi-GPU training.
- Verify the framework build matches your CUDA/cuDNN versions; pip/conda usually handles this, but container images are safer.

3. Leverage Multi-GPU Training

Data Parallelism:
- Split each batch across GPUs; each replica runs forward/backward on its shard.
- Use framework built-ins: torch.nn.parallel.DistributedDataParallel (preferred) or torch.nn.DataParallel (legacy, single-process) in PyTorch; tf.distribute.MirroredStrategy in TensorFlow.
Model Parallelism:
- Partition the model itself when it exceeds one GPU’s memory.
- Requires manual sharding or pipeline libraries (e.g., Megatron-LM, DeepSpeed, PyTorch RPC); adds engineering overhead.
Mixed Precision Training:
- FP16 or BF16 cuts memory traffic and boosts throughput on Tensor Cores.
- Enable via PyTorch AMP (torch.cuda.amp.autocast) or TensorFlow mixed precision policy; watch for gradient underflow and use loss scaling.

4. Set Up Distributed Training

Single Node, Multi-GPU: Launch with torchrun (PyTorch ≥1.10) or python -m torch.distributed.launch (older). Set NCCL_DEBUG=INFO first run to catch topology issues.
Multi-Node, Multi-GPU:
- Use torch.distributed, Horovod, or tf.distribute.MultiWorkerMirroredStrategy.
- Network must be low-latency: InfiniBand or 100 GbE with RoCE. TCP over 10 GbE will bottleneck at scale.
- Synchronize clocks (PTP or chrony) and disable CPU frequency scaling on worker nodes.

5. Optimize GPU Utilization

Batch Size: Scale global batch size with GPU count. Per-GPU batch too small underutilizes compute; too large OOMs. Profile.
Asynchronous Data Loading: Use DataLoader(num_workers=N, prefetch_factor=2, persistent_workers=True) in PyTorch; tf.data.Dataset.prefetch(AUTOTUNE) in TensorFlow. Keep the GPU fed.
GPU Profiling:
- nsys profile (Nsight Systems) for kernel timeline and CPU-GPU gaps.
- ncu (Nsight Compute) for single-kernel occupancy and memory throughput.
- Framework profilers (TensorBoard Profiler, PyTorch Profiler) for operator-level view.

6. Monitor and Maintain

Monitoring:
- nvidia-smi -l 1 or dcgmi dmon for per-GPU utilization, memory, temperature, power.
- Export metrics to Prometheus via dcgm-exporter or nvidia-gpu-prometheus-exporter; visualize in Grafana.
Error Handling: Watch for OOM, CUDA kernel launch failures, NCCL timeout/heartbeat errors, and ECC correctable/uncorrectable counts. Each points to a different fix.
Updates: Pin driver/CUDA/framework versions per project. Upgrade on a schedule, not reactively; test in staging before rolling to training clusters.

7. GPU Resource Scheduling

For shared clusters, run workloads in containers with GPU isolation:
- Kubernetes + NVIDIA GPU Operator (device plugin, GFD, validator, monitoring).
- Slurm with gres/gpu and srun --gpus for batch scheduling.
- Set compute mode (DEFAULT, EXCLUSIVE_PROCESS, PROHIBITED) via nvidia-smi -c to prevent accidental sharing.

8. Plan for Scale

Horizontal Scaling: Add nodes; ensure the interconnect and scheduler handle the new topology.
Cloud GPUs: AWS p4d/p5, Azure NDv4/NDv5, GCP A3 VMs give on-demand A100/H100. Use spot/preemptible for fault-tolerant jobs to cut cost.
Hybrid: Burst to cloud when on-prem queue backs up; keep long-running or data-sensitive jobs local.

9. Test and Benchmark

Run scaling sweeps: 1, 2, 4, 8 GPUs. Plot throughput (samples/sec) vs. GPU count. Linear scaling is rare; 80%+ efficiency is good.
Compare FP32 vs. mixed precision. Measure memory at each batch size to find the ceiling.
Log NCCL all-reduce latency and bandwidth (NCCL_TESTS=1) to verify fabric health.