We train our models using PyTorch on NVIDIA A100 GPUs and our biggest model with 10B parameters requires 40GB of GPU memory (with additional 40GB required for optimizer state during pre-training). On a single V100_32G GPU or more recent 40GB A100, such a model can not fit and hence the DDP (Distributed Data Parallel) training can not be used. We instead resort to model sharding and use the Fully Sharded Data Parallel FSDP [1]}, [2]} training which shards the model such that each layer of the model is sharded across different data parallel workers (GPUs). The computation of minibatch is still local to each GPU worker. FSDP decomposes the all-reduce operations in DDP into separate reduce-scatter and all-gather operations. 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. During the forward pass, the parameters of the layer to be computed are temporarily assembled before they are re-sharded. For training efficiency, the communication and computation are overlapped: un-sharding the next layer parameters (via all-gather) while computing the current layer. We illustrate the communication and compute optimizations in Fig. REF . For our model training, we leverage the FSDP implementation from Fairscale https://github.com/facebookresearch/fairscale and adapt it for our model.
