Skip gradient averaging if there are no other peers (#440)

Optimizer will now skip grad averaging if there are no peers to average gradients with. Previously, it would invoke grad_averager.step and wait for averaging_timeout seconds. Co-authored-by: Qidong Su <[email protected]> Co-authored-by: Alexander Borzunov <[email protected]>
learning-at-home · Jan 4, 2022 · c868989 · c868989
1 parent cfc5200
commit c868989
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/hivemind/optim/optimizer.py b/hivemind/optim/optimizer.py
@@ -524,8 +524,12 @@ def _begin_averaging_gradients(self, grad_scaler: Optional[GradScaler]) -> bool:
                 logger.exception(e)
 
         if not began_averaging_gradients and self.scheduled_grads is not None and not self.scheduled_grads.done():
-            logger.log(self.status_loglevel, f"Tagging along for a pre-scheduled gradient averaging round")
-            self._tag_along_with_zero_weight(self.scheduled_grads)
+            if self.tracker.global_progress.num_peers > 1:
+                logger.log(self.status_loglevel, f"Tagging along for a pre-scheduled gradient averaging round")
+                self._tag_along_with_zero_weight(self.scheduled_grads)
+            else:
+                logger.log(self.status_loglevel, f"Skipping pre-scheduled averaging round: there are no other peers")
+                self.scheduled_grads.cancel()
             self.scheduled_grads = None
         return began_averaging_gradients