## Understanding Distributed Training Strategies and Hardware Effects
Modern deep learning workflows often require distributing training across multiple GPUs. At a high level, adding GPUs seems simple: replicate the model and data, and wait for faster completion. While this works for many scenarios—particularly when the model fits on a single GPU—training large foundation models introduces new challenges around memory constraints and communication overhead.
When models exceed single-GPU memory capacity, simply adding more GPUs with full replication is insufficient. The model, its gradients, and optimizer states must be split across devices. This leads to a spectrum of strategies, from Distributed Data Parallel (DDP) to Fully Sharded Data Parallel (FSDP), with intermediate options like ZeRO from DeepSpeed. Each strategy balances memory usage against communication frequency.
Equally important is the underlying hardware. Data must move between GPUs, and the speed of this communication depends heavily on the physical interconnects—whether GPUs are connected via PCIe, NVLink, or NVSwitch. The same distributed strategy can perform drastically differently depending on how GPUs are wired together.
This article explores both the software strategies and hardware considerations for efficient distributed training.
—
### Part 1: The Software – What Each GPU Holds
Distributed training strategies differ in how they partition model state across GPUs. The core question is: what does each GPU keep locally?
**Example Setup:** Fine-tuning a Mistral-7B model with Adam in BF16 requires approximately:
– Parameters: 14 GB
– Gradients: 14 GB
– Optimizer states (Adam): ~58 GB
Total: ~87 GB per GPU
An A100 GPU has only 80 GB, so full replication is impossible. Distributed strategies address this by splitting state.
#### Distributed Data Parallel (DDP)
DDP keeps a full copy of the model on each GPU. Each GPU processes a different batch slice, computes gradients independently, and then averages them via an **all-reduce** operation before taking an optimizer step.
DDP is fast and communication-efficient—only one all-reduce per training step. However, it does not reduce memory usage per GPU. Adding more GPUs increases throughput but does not solve memory constraints.
#### Fully Sharded Data Parallel (FSDP)
FSDP eliminates full model replication by partitioning parameters, gradients, and optimizer states across GPUs. Before executing a layer, FSDP gathers the relevant parameters, computes the forward and backward passes, and then scatters the results.
This approach drastically reduces per-GPU memory usage but increases communication. FSDP performs an **all-gather** before each layer and a **reduce-scatter** after each backward pass. While memory-friendly, this constant communication adds overhead.
#### ZeRO Stages: A Dial Between Strategies
Rather than choosing only between DDP and FSDP, DeepSpeed’s ZeRO offers intermediate options:
– **ZeRO-1:** Shards only optimizer states
– **ZeRO-2:** Shards optimizer states and gradients
– **ZeRO-3:** Shards parameters, gradients, and optimizer states
Each stage reduces memory at the cost of increased communication. FSDP aligns closely with ZeRO-3, while DDP resembles full replication.
—
### Part 2: The Hardware – How Fast GPUs Can Talk
Communication cost depends on the physical interconnect between GPUs. The same distributed strategy can run dramatically faster or slower depending on hardware topology.
#### Common Interconnects
– **PCIe:** Default host bus connection. Slower and shared among devices.
– **NVLink:** Dedicated GPU-to-GPU link offering ~450 GB/s per H100/H200.
– **NVL:** Groups of GPUs connected internally with NVLink, separated by PCIe between groups.
– **NVSwitch:** Full NVLink mesh, allowing every GPU to communicate directly with every other at full speed.
#### Performance Differences
Measuring collective operations reveals significant differences:
– **NVLink vs. forced PCIe:** Up to 10× bandwidth difference
– **All-reduce scaling:** NVSwitch maintains consistent bandwidth as GPU count increases. NVL performs well within a group but degrades when crossing groups. PCIe-based systems show the steepest decline.
In real training, these differences are equally important. Communication-heavy strategies like FSDP suffer most on slow interconnects, while DDP remains more resilient.
—
### Combined Insights: Strategy + Hardware
Experiments show that:
– On **fast fabrics (NVSwitch)**, communication is cheap. FSDP closely matches DDP in speed while using far less memory.
– On **slow fabrics (cross-group NVL or PCIe)**, communication dominates runtime. Sharding strategies lose much of their advantage, and DDP becomes more attractive when memory permits.
Memory savings decrease predictably across ZeRO stages, but performance varies based on hardware. Choosing the right configuration requires understanding both model size and GPU interconnect topology.
—
### Practical Recommendations
1. **Check your hardware topology** using:
“`bash
nvidia-smi topo -m
“`
2. **If using NVSwitch or full NVL:** Prefer FSDP when memory-bound; otherwise DDP is simpler and slightly faster.
3. **If using NVL with small jobs:** Pin workloads to a single group for near-NVSwitch speeds.
4. **If spanning multiple NVL groups or using PCIe:** Expect significant slowdowns for sharded strategies. DDP or ZeRO-1 may be preferable when memory allows.
—
### Conclusion
Effective distributed training requires balancing software strategies with hardware capabilities. While model sharding reduces memory pressure, its performance depends heavily on how quickly GPUs can exchange data. By understanding both the algorithmic choices and the underlying interconnects, practitioners can make informed decisions that optimize speed, memory, and stability.
—
**References and Further Reading**
– PyTorch Distributed and NCCL documentation
– DeepSpeed Zero and FSDP whitepapers
– NVIDIA H200 and NVLink technical guides
*This article is based on practical measurements and real-world training runs described in the original piece “Distributed Training Is Easier Said Than Done.”*
