Pipeline timers (deepspeedai#15)

* update default configs * fix bug with onebitadam + p.p hanging * send tensors to cuda Co-authored-by: sid <sidney.black@aleph-alpha.de>
determined-ai · Mar 10, 2021 · e020c27 · e020c27
1 parent d10a3f8
commit e020c27
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/deepspeed/runtime/fp16/fused_optimizer.py b/deepspeed/runtime/fp16/fused_optimizer.py
@@ -220,9 +220,9 @@ def step(self, closure=None):
             if hasattr(self.optimizer.comm_backend_handle,
                        "mpu") and self.optimizer.comm_backend_handle.mpu is not None:
                 if self.overflow:
-                    bool_tensor = torch.zeros(1)
+                    bool_tensor = torch.zeros(1).cuda()
                 else:
-                    bool_tensor = torch.ones(1)
+                    bool_tensor = torch.ones(1).cuda()
                 dist.all_reduce(bool_tensor, op=dist.ReduceOp.PRODUCT)
                 if not any(bool_tensor):
                     self.overflow = True