Paper notes
ZeRO
(0) Basic: foundational concepts and knowledge
- Inter-GPU bandwidth: communication speed depends on whether GPUs sit in the same server. Intra-node (within one server), GPUs talk directly over NVLink/NVSwitch at hundreds of GB/s — the fast path:
GPU → NVLink → GPU. Inter-node (across servers), traffic crosses the network through each node's NIC:GPU → NIC → network → NIC → GPU, using either InfiniBand (common in HPC clusters, low latency / high bandwidth) or plain Ethernet (cheaper but slower). The takeaway: inter-node links are far slower than intra-node, so sharding strategies try to keep heavy communication inside a node. - the memory: this paper refers to is GPU memory that GPU memory cannot host the full model and usually host memory is way larger than per-GPU memory
(1) Generic: one MEMORY-OPTIMIZATION mental model
- MEMORY FOOTPRINT identification: like in LLM training, the model state includes weights, gradients, optimizer states (for Adam optimizer they are 2-3 times of weights) and activation states. This analysis is domain/project specific.
- Strategies to reduce MEMORY REDUNDANCY: like in ZeRO → full sharding of all states, and all-gather when computing one layer.
- Find pros and CONS for all strategies & try to make a COMBO-solution: like ZeRO still has communication overhead and modern practice uses ZeRO + DP (FSDP in PyTorch).
(2) LLM Domain: different types of Parallelism in pretraining: DP, TP, PP
Family
Method
What is split / what each GPU stores
Computation & communication
Main drawback
Data parallelism
Data Parallel (DP)
Splits: batch (data). Stores: full model + a different mini-batch.
Each GPU runs full forward/backward independently; gradient all-reduce each step
Model must fit on one GPU; full model duplicated everywhere
Model parallelism (MP)
Tensor Parallel (TP)
Splits: model tensors (matrix rows/columns, attention heads) — within a layer. Stores: shards of weight matrices (e.g., W0, W1) + full activations.
Each GPU computes partial matmul outputs per layer; frequent all-gather / all-reduce of intermediate activations every layer
High communication overhead; tight synchronization inside every layer
Pipeline Parallel (PP)
Splits: model depth — across layers. Stores: different contiguous layers per GPU.
Micro-batches flow stage-by-stage through GPUs; activations passed between pipeline stages
Pipeline bubbles (GPU idle time), limited utilization, microbatch complexity
Data parallelism (sharded)
ZeRO / FSDP (ZeRO-3)
Splits: training states (parameters, gradients, optimizer states). Stores: sharded parameters (W0, W1, etc.).
All-gather full weights per layer → full-layer compute → discard/shard again; all-gather params + reduce-scatter gradients per layer
Communication overhead for parameter reconstruction; bandwidth-bound at scale
(3) LLM Domain: ZeRO vs.TP
| Family | Method | What is split / what each GPU stores | Computation & communication | Main drawback |
|---|---|---|---|---|
| Data parallelism | Data Parallel (DP) | Splits: batch (data). Stores: full model + a different mini-batch. | Each GPU runs full forward/backward independently; gradient all-reduce each step | Model must fit on one GPU; full model duplicated everywhere |
| Model parallelism (MP) | Tensor Parallel (TP) | Splits: model tensors (matrix rows/columns, attention heads) — within a layer. Stores: shards of weight matrices (e.g., W0, W1) + full activations. | Each GPU computes partial matmul outputs per layer; frequent all-gather / all-reduce of intermediate activations every layer | High communication overhead; tight synchronization inside every layer |
| Pipeline Parallel (PP) | Splits: model depth — across layers. Stores: different contiguous layers per GPU. | Micro-batches flow stage-by-stage through GPUs; activations passed between pipeline stages | Pipeline bubbles (GPU idle time), limited utilization, microbatch complexity | |
| Data parallelism (sharded) | ZeRO / FSDP (ZeRO-3) | Splits: training states (parameters, gradients, optimizer states). Stores: sharded parameters (W0, W1, etc.). | All-gather full weights per layer → full-layer compute → discard/shard again; all-gather params + reduce-scatter gradients per layer | Communication overhead for parameter reconstruction; bandwidth-bound at scale |
Both ZeRO and TP use intra-layer sharding, but they split different things. TP partitions the work — each GPU only ever computes part of the layer and they talk mid-computation to combine. ZeRO partitions the memory — each GPU all-gathers the full layer just-in-time, does the whole computation on its own batch, then frees the gathered weights.