fully sharded data parallel
In plain wordsA feature of the popular PyTorch software that lets many chips train a model too big for any one of them, by spreading the model and its training bookkeeping evenly across all the chips.
A data-parallel training method built into PyTorchPyTorch, prop. n.: The most widely used free software toolkit for building and training neural networks, especially in research; programmers write ordinary Python and it handles the calculus and the graphics chips. that applies the ZeROZeRO, prop. n.: A memory-saving technique for training big models on many chips: instead of every chip keeping its own complete copy of all the training bookkeeping, the chips divide it among themselves and share pieces only when needed. stage-3 idea: each worker permanently keeps only its own slice of the parameters, gradients and optimizer states, and complete layers exist only briefly while they are being computed. Before each wrapped unit's forward and backward computation, workers all-gather its full parameters and then free them; after the backward pass, gradients are reduce-scattered so that each worker updates only its own shard. Communication overlaps with computation by prefetching the next unit's parameters. FSDP keeps the programming model of ordinary data parallelismdata parallelism, n.: Splitting the training data among many chips that each hold a full copy of the model; every chip learns from its own share, and they pool what they learned after each step so all copies stay identical. while allowing models far larger than one accelerator's memory, and it combines with tensortensor parallelism, n.: Splitting the individual giant tables of numbers inside each layer of a model across several chips, so that every chip does a slice of every calculation at the same time. and pipeline parallelismpipeline parallelism, n.: Splitting a model's layers into consecutive stages on different chips, like stations on an assembly line, and streaming small batches of data through so that every station stays busy. in multi-dimensional layouts. A hybrid mode shards within a node and replicates across nodes to reduce cross-node traffic. FSDP2 (2024) re-implemented it with per-parameter sharding on PyTorch's distributed tensor (DTensor) abstraction.
Wrap each transformer block in FSDP so only one block's full weights are materialized at a time.
Origin
Developed at Facebook AI Research (now Meta) in the FairScale library, released in early 2021 and building on Microsoft's ZeRO and DeepSpeed; upstreamed to PyTorch 1.11 as a prototype in March 2022. Yanli Zhao and colleagues described the PyTorch design in 2023.
Related
- ZeROZeRO, prop. n.: A memory-saving technique for training big models on many chips: instead of every chip keeping its own complete copy of all the training bookkeeping, the chips divide it among themselves and share pieces only when needed.
- data parallelismdata parallelism, n.: Splitting the training data among many chips that each hold a full copy of the model; every chip learns from its own share, and they pool what they learned after each step so all copies stay identical.
- PyTorchPyTorch, prop. n.: The most widely used free software toolkit for building and training neural networks, especially in research; programmers write ordinary Python and it handles the calculus and the graphics chips.
- distributed trainingdistributed training, n.: Training one AI model on many computer chips at once, splitting up the data or the model itself among them and keeping them in sync, so that jobs that would take a single chip thousands of years finish in weeks.
- 3D parallelism3D parallelism, n.: Combining three ways of dividing up training work — by data, by layer stage and by slices within layers — so that thousands of chips can train one giant model efficiently.
- collective communicationcollective communication, n.: The standard patterns by which a whole group of chips exchanges data at once — for example, everyone adding up their numbers and all getting the total — which large training jobs use constantly.
Sources
- Zhao, Y., Gu, A., Varma, R., Luo, L., Huang, C.-C. et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. arXiv:2304.11277. arxiv.org/abs/2304.11277
- PyTorch Team (2022). Introducing PyTorch Fully Sharded Data Parallel (FSDP) API. PyTorch blog. pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api
Mentioned in
data parallelismdata parallelism, n.: Splitting the training data among many chips that each hold a full copy of the model; every chip learns from its own share, and they pool what they learned after each step so all copies stay identical., distributed trainingdistributed training, n.: Training one AI model on many computer chips at once, splitting up the data or the model itself among them and keeping them in sync, so that jobs that would take a single chip thousands of years finish in weeks., 3D parallelism3D parallelism, n.: Combining three ways of dividing up training work — by data, by layer stage and by slices within layers — so that thousands of chips can train one giant model efficiently., collective communicationcollective communication, n.: The standard patterns by which a whole group of chips exchanges data at once — for example, everyone adding up their numbers and all getting the total — which large training jobs use constantly., DeepSpeedDeepSpeed, prop. n.: Free software from Microsoft for training and running AI models too big for ordinary tools, spread over many chips; it is best known for a memory-saving technique called ZeRO., ZeROZeRO, prop. n.: A memory-saving technique for training big models on many chips: instead of every chip keeping its own complete copy of all the training bookkeeping, the chips divide it among themselves and share pieces only when needed.