Source
Zhao et al. (2022), LIT-083 — ARXIV-2304.11277.
What is being overlapped, and with what
Sharding costs an all-gather of each unit’s parameters before it runs. Left serial, the GPU idles through every one of them. Backward prefetch issues the gather for unit i−1 while unit i’s gradients are still being computed, so the collective runs against compute that is already scheduled.
The forward direction has the same shape and is easier — the execution order is known — which is why the backward case is the one with a name: the order is discovered by autograd, so the prefetch has to predict it from the forward’s recorded sequence.
Backward prefetching is LIT-083’s, and §3.3.2 gives the specific reason for
it. FSDP runs both collectives on one process group, so in the backward pass
the ReduceScatter for the current FlatParameter blocks the AllGather
for the next, putting two exposed collectives on the critical path; FSDP
issues the next AllGather first, using the reverse of the forward order,
re-recorded every iteration, as its guess at the backward order. On
GPT-175B, where communication weighs most, prefetching gave about an 18%
speedup that held across cluster sizes, and the authors leave it on for
every later experiment (§5.2, Figure 6(b)). The record’s reading,
NOTE-003, gives the same ablation.
Cost
Prefetching means two units’ parameters are resident at once instead of one, so the peak footprint rises by roughly one unit. On a job that chose FSDP because it was out of memory (SOTA-116) that is not free, and the knob is the first thing to turn off when the run OOMs near the peak rather than at the start.
It also depends on the predicted order being right. A model whose backward order differs from the reverse of its forward — conditional branches, shared modules used twice — prefetches the wrong unit and pays the gather twice.