Use cuda tensors for allgather (#1548)

2025-10-20 23:53:48 +08:00 · 2021-11-10 17:27:00 -08:00
parent af443f63f4
commit bd3ebddf36
1 changed files with 1 additions and 1 deletions
--- a/deepspeed/runtime/zero/partition_parameters.py
+++ b/deepspeed/runtime/zero/partition_parameters.py
@ -889,7 +889,7 @@ class Init(InsertPostInitMethodToModuleSubClasses):
        if self.use_all_gather_base:
            # try the _all_gather_base on PyTorch master branch
            handle = dist._all_gather_base(flat_tensor,
-                                           param.ds_tensor,
+                                           param.ds_tensor.cuda(),
                                           group=self.ds_process_group,
                                           async_op=async_op)
        else: