anthology-of-the-sota

LIT-083: PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

2 min read

Properties6
StatusActive
Version2
Tagsdistributed-optimization
Cited as “introduced by” bySOTA-116 — Use FSDP over DDP when model size exceeds single GPU memory, SOTA-117 — Overlap communication with computation using backward prefetch, SOTA-118 — Employ mixed precision to reduce memory usage, SOTA-119 — Choose sharding factor based on model and GPU memory size
Cited as “paper” byNOTE-003 — PyTorch FSDP
Cited as “source” bySOTA-116 — Use FSDP over DDP when model size exceeds single GPU memory, SOTA-117 — Overlap communication with computation using backward prefetch, SOTA-118 — Employ mixed precision to reduce memory usage, SOTA-119 — Choose sharding factor based on model and GPU memory size

Graph View

Zhao et al. (2022) — ARXIV-2304.11277

Key takeaways

  • The sharding factor F — the number of ranks parameters are sharded over — generalises the whole strategy space. F=1 is full replication and reduces to DDP; F=W is full sharding, each device holding 1/W; 1 < F < W is hybrid sharding
  • FSDP is co-designed with PyTorch internals — Tensor, the dispatcher, the CUDA caching allocator — so that using it does not mean restructuring the model. That is the contribution over ZeRO, which established sharding
  • Backward prefetching: with one process group serving both collectives, the ReduceScatter for the current FlatParameter blocks the AllGather for the next, putting two collectives on the critical path. FSDP issues the next AllGather first
  • Forward overlap works on CUDA streams rather than Work objects; the outermost unit’s parameters are kept resident to avoid a free-then-regather
  • A rate limiter bounds the caching allocator, which runs on the CPU thread and can over-allocate under several CUDA streams
  • Near-linear TFLOPS scalability, comparable to DDP, supporting substantially larger models

Standing in the anthology

The source of SOTA-116, SOTA-117, SOTA-118 and SOTA-119.

Read in full for #114; the reading is NOTE-003. All four practices are confirmed, which is the first such result in that issue — the three clusters read before it each inverted their source. The generic bullet shape is a signal that a note was never read, not evidence that it is wrong.

Backlinks