use dist.device instead of dist.rank

Author: yexin

checkpoint_engine/distributed_hccl.py

@@ -298,7 +298,7 @@ def barrier(group=None):
         subcomm = ctypes.c_void_p(group)
         dist.pyhccl.comm = subcomm
 
-    data = torch.zeros(1, device=dist.rank)
+    data = torch.zeros(1, device=dist.device)
     dist.pyhccl.all_reduce(data)
     current_stream().synchronize()
 

checkpoint_engine/distributed_nccl.py

@@ -270,7 +270,7 @@ def barrier(group=None):
         newcomm = ctypes.c_void_p(group)
         dist.pynccl.comm = newcomm
 
-    data = torch.zeros(1, device=dist.rank)
+    data = torch.zeros(1, device=dist.device)
     dist.pynccl.all_reduce(data)
     current_stream().synchronize()