update condition&comments

microsoft · Dec 25, 2024 · 8d59565 · 8d59565
1 parent 931c121
commit 8d59565
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
@@ -713,7 +713,8 @@ def wait(self, handle_dependency=True) -> None:
             param_offset += ds_tensor_numel
 
         self.complete = True
-        if not handle_dependency:
+        if not get_accelerator().is_synchronized_device() and not handle_dependency:
+            # if the device needs to handle dependencies and opts for explicit processing outside the function.
             AllGatherCoalescedHandle.data_buffer.append(partitions)
 
     @staticmethod