{"id":17795,"date":"2021-07-15T09:00:47","date_gmt":"2021-07-15T16:00:47","guid":{"rendered":"https:\/\/engineering.fb.com\/?p=17795"},"modified":"2021-07-22T11:02:59","modified_gmt":"2021-07-22T18:02:59","slug":"fsdp","status":"publish","type":"post","link":"https:\/\/engineering.fb.com\/2021\/07\/15\/open-source\/fsdp\/","title":{"rendered":"Fully Sharded Data Parallel: faster AI training with fewer GPUs"},"content":{"rendered":"

Training AI models at a large scale isn\u2019t easy. Aside from the need for large amounts of computing power and resources, there is also considerable engineering complexity behind training very large models. At Facebook AI Research (FAIR) Engineering, we have been working on building tools and infrastructure to make training large AI models easier. Our recent work in areas such as <\/span>intra-layer model parallelism<\/span><\/a>, <\/span>pipeline model parallelism<\/span><\/a>, <\/span>optimizer state+gradient sharding<\/span><\/a>, and <\/span>mixture of experts<\/span><\/a> is just part of our work to make training advanced AI models for any number of tasks more efficient.<\/span><\/p>\n

Fully Sharded Data Parallel (FSDP) is the newest tool we\u2019re introducing. It shards<\/a> an AI model\u2019s parameters across data parallel workers and can optionally offload part of the training computation to the CPUs. As its name suggests, FSDP is a type of data-parallel training algorithm. Although the parameters are sharded to different GPUs<\/a>, the computation for each microbatch of data is still local to each GPU worker. This conceptual simplicity makes FSDP easier to understand and more applicable to a wide range of usage scenarios (compared with intra-layer parallelism and pipeline parallelism). Compared with optimizer state+gradient sharding data parallel methods, FSDP shards parameters more uniformly and is capable of better performance via communication and computation overlapping during training.<\/span><\/p>\n

With FSDP, it is now possible to more efficiently train models that are orders of magnitude larger using fewer GPUs. FSDP has been implemented in the <\/span>FairScale library<\/span><\/a> and allows engineers and developers to scale and optimize the training of their models with simple APIs. At Facebook, FSDP has already been integrated and tested for training some of our <\/span>NLP<\/span><\/a> and<\/span> Vision<\/span><\/a> models.<\/span><\/p>\n

The high computational cost of large-scale training<\/span><\/h2>\n

NLP research<\/span><\/a> is one particular area where we can see the importance of efficiently leveraging compute for training AI. Last year, OpenAI announced that they had trained <\/span>GPT-3<\/span><\/a>, the largest-ever neural language model, with 175 billion parameters. It is <\/span>estimated<\/span><\/a> to have taken roughly 355 GPU years to train GPT-3, or the equivalent of 1,000 GPUs working continuously for more than four months.<\/span><\/p>\n

Besides requiring a lot of compute and engineering resources, most approaches to scaling like this introduce additional communication costs and require engineers to carefully evaluate trade-offs between memory use and computational efficiency. For example, typical data parallel training requires maintaining redundant copies of the model on each GPU, and model parallel training introduces additional communication costs to move activations between workers (GPUs).<\/span><\/p>\n

FSDP is relatively free of trade-offs in comparison. It improves memory efficiency by sharding model parameters, gradients, and optimizer states across GPUs, and improves computational efficiency by decomposing the communication and overlapping it with both the forward and backward passes. FSDP produces identical results as standard distributed data parallel (DDP) training and is available in an easy-to-use interface that\u2019s a drop-in replacement for PyTorch\u2019s DistributedDataParallel module. Our early testing has shown that FSDP can enable scaling to trillions of parameters.<\/span><\/p>\n

How FSDP works<\/span><\/h2>\n

In standard DDP training, every worker processes a separate batch and the gradients are summed across workers using an <\/span>all-reduce operation<\/span><\/a>. While DDP has become very popular, it takes more GPU memory than it needs because the model weights and optimizer states are replicated across all DDP workers.<\/span><\/p>\n

One method to reduce replications is to apply a process called full parameter sharding, where only a subset of the model parameters, gradients, and optimizers needed for a local computation is made available. An implementation of this method, ZeRO-3, has already been popularized by Microsoft.\u00a0<\/span><\/p>\n

The key insight to unlock full parameter sharding is that we can decompose the <\/span>all-reduce<\/span><\/a> operations in DDP into separate <\/span>reduce-scatter<\/span><\/a> and all-gather<\/a> operations:<\/span><\/p>\n

\"Full
All-reduce as a combination of reduce-scatter and all-gather. The standard all-reduce operation to aggregate gradients can be decomposed into two separate phases: reduce-scatter and all-gather. During the reduce-scatter phase, the gradients are summed in equal blocks among ranks on each GPU based on their rank index. During the all-gather phase, the sharded portion of aggregated gradients available on each GPU are made available to all GPUs (see here for details on those operators).<\/figcaption><\/figure>\n

We can then rearrange the reduce-scatter and all-gather so that each DDP worker needs to store only a single shard of parameters and optimizer states. The figure below illustrates standard DDP training (top) and FSDP training (bottom):<\/span><\/p>\n

\"Full
A comparison of standard data parallel training and fully sharded data parallel training. In standard data parallel training methods, a copy of the model is present on each GPU and a sequence of forward and backward passes are evaluated on only a shard of the data. After these local computations, the parameters and optimizers for each local process are shared with the other GPUs in order to calculate the global weight update. In FSDP, only a shard of the model is present on a GPU. Then, locally, all weights are gathered from the other GPUs \u2014 by means of an all-gather step \u2014 to calculate the forward pass. This gathering of weights is then performed again before the backward pass. After that backward pass, the local gradients are averaged and sharded across the GPUs by means of a reduce-scatter step, which allows each GPU to update its local weight shard.<\/figcaption><\/figure>\n

To maximize memory efficiency, we can discard the full weights after each layer\u2019s forward pass, saving memory for subsequent layers. This can be implemented by applying the FSDP wrapper to every layer in the network (with <\/span>reshard_after_forward=True<\/span>). <\/span><\/p>\n

In pseudo-code:<\/span><\/p>\n

FSDP forward pass:<\/span>\r\n \u00a0\u00a0\u00a0for layer_i in layers:<\/span>\r\n \u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0all-gather full weights for layer_i<\/span>\r\n \u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0forward pass for layer_i<\/span>\r\n\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0discard full weights for layer_i<\/span>\r\n\r\nFSDP backward pass:<\/span>\r\n \u00a0\u00a0\u00a0for layer_i in layers:<\/span>\r\n \u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0all-gather full weights for layer_i<\/span>\r\n \u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0backward pass for layer_i<\/span>\r\n \u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0discard full weights for layer_i<\/span>\r\n\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0reduce-scatter gradients for layer_i<\/span><\/pre>\n
How to use FSDP<\/span><\/h2>\n
There are several ways to use FSDP in large-scale AI research.<\/span> At this time, we offer four solutions to adapt to different needs.<\/span><\/p>\n
1. Using FSDP in language models<\/span><\/h3>\n
For language models, FSDP is supported in the <\/span>fairseq<\/span><\/i> framework<\/span><\/a> via the following new arguments:<\/span><\/p>\n