All-reduce in FP16

fairseq/fp16_trainer.py

@@ -11,7 +11,7 @@ Train a network on multiple GPUs.
 
 import torch
 
-from fairseq import optim
+from fairseq import optim, utils
 from fairseq.meters import AverageMeter
 from fairseq.optim import lr_scheduler
 from fairseq.trainer import Trainer
@@ -105,8 +105,26 @@ class FP16Trainer(Trainer):
         # undo effect of dynamic loss scaling on gradients
         grad_denom *= self.scaler.loss_scale
 
-        # all-reduce and rescale gradients
-        grad_norm = super()._all_reduce_and_rescale(grad_denom)
+        if self.args.distributed_world_size > 1:
+            # flatten grads into a single buffer
+            flat_grads = self._flat_grads = self._get_flat_grads(self._flat_grads)
+
+            # scale gradients to avoid overflow in all-reduce
+            flat_grads.div_(self.args.distributed_world_size)
+            grad_denom /= self.args.distributed_world_size
+
+            # all-reduce flat grads
+            torch.distributed.all_reduce(flat_grads)
+
+            # copy grads back to FP32
+            self.fp32_params.grad.data.copy_(flat_grads)
+        else:
+            # single worker: copy grads directly to FP32
+            self._get_flat_grads(out=self.fp32_params.grad.data)
+
+        # rescale and clip grads
+        self.fp32_params.grad.data.div_(grad_denom)
+        grad_norm = utils.clip_grad_norm_(self.fp32_params.grad.data, self.args.clip_norm)
 
         # detect overflow and adjust loss scale
         overflow = DynamicLossScaler.has_overflow(grad_norm)
@@ -116,15 +134,6 @@ class FP16Trainer(Trainer):
 
         return grad_norm
 
-    def _get_flat_grads(self, out=None):
-        if out is None:
-            out = self.fp32_params.grad
-        return super()._get_flat_grads(out)
-
-    def _set_flat_grads(self, new_grads):
-        # no-op
-        assert new_grads.data_ptr() == self.fp32_params.grad.data.data_ptr()
-
     def _opt(self):
         # take an optimization step using the FP32 params and grads
         super()._opt()