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
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.