fix bug when CUDA_VISABLE_DEVICES is set

Author: yexin

checkpoint_engine/distributed/hccl.py

@@ -205,7 +205,7 @@ def init_process_group(
         self.port = port
         self.rank = rank
         self.world_size = world_size
-        self.device = torch.device("npu", rank)
+        self.device = torch.device("npu", torch.npu.current_device())
 
         self.pg = StatelessProcessGroup.create(
             host, port, rank, world_size, store_timeout=int(timeout.total_seconds())

checkpoint_engine/distributed/nccl.py

@@ -122,7 +122,7 @@ def init_process_group(
         self.port = port
         self.rank = rank
         self.world_size = world_size
-        self.device = torch.device("cuda", rank)
+        self.device = torch.device("cuda", torch.cuda.current_device())
 
         self.pg = StatelessProcessGroup.create(
             host, port, rank, world_size, store_timeout=int(timeout.total_seconds())