Fix dist lamb (#1185)

1. remove the weight broadcast in the constructor
2. disable unnecessary allreduces for clip-after-ar

apex/contrib/optimizers/distributed_fused_lamb.py

@@ -270,7 +270,6 @@ class DistributedFusedLAMB(torch.optim.Optimizer):
             eps = group['eps']
             weight_decay = group['weight_decay']
             for p in group['params']:
-                torch.distributed.broadcast(p, 0)
                 if not p.requires_grad:
                     continue
                 self._model_params.append(p)
@@ -729,6 +728,8 @@ class DistributedFusedLAMB(torch.optim.Optimizer):
         # check global_grad_norm and fill overflow_buf
         is_finite = (global_grad_norm + 1 > global_grad_norm).int()
         self._overflow_buf = self._one * (is_finite ^ self._one) # toggle between 0 and 1
+
+        if not self._clip_after_ar:
             torch.distributed.all_reduce(is_finite,
                                          op=torch.distributed.ReduceOp.MIN,
                                          group=self._current_process_group)