To optimize the training speed of the model, we use several optimizations. We use mixed-precision for training and perform the forward pass computations in FP16. Since computation happens in FP16, for the un-sharding of parameters via all-gather operation (which performs communication of parameters over the network), we exchange FP16 weights instead of FP32. This speeds-up the training by communicating model parameters faster. We note that for certain special layers such as SyncBatchNorm, we still use FP32 as otherwise the training becomes unstable. Further, we use LARC optimizer [1]} from NVIDIA Apex libraryhttps://github.com/NVIDIA/apex for large batch size training. Since the model parameters are sharded, we adapted the LARC implementation to compute the distributed norms of parameters but without all-gather of model weights. We share more details on this in Appendix . Additionally, we add the activation checkpointing in the order FSDP(checkpointing(model layer)) instead of the other way around. This is because activation checkpointing re-computes the forward pass on part of the model during back-propagation and doing a forward pass on FSDP wrapped layer requires “un-sharding” of layer which involves communication of weights across all GPUs. Hence, FSDP(checkpointing(model layer)) ensures that we do not trigger excessive “un-sharding” / communication cost across GPUs.
