[DCP][Draft] Checkpoint daemon process fixes

meetv18 · facebook-github-bot · commit 86081b8592a2 · 2025-03-17T12:53:59.000-07:00
Differential Revision: D71336180
diff --git a/torch/distributed/checkpoint/_async_process_executor.py b/torch/distributed/checkpoint/_async_process_executor.py
@@ -8,6 +8,8 @@
 from typing import Any, Optional, Union
 from uuid import uuid4
 
+import torch
+
 import torch.distributed as dist
 import torch.multiprocessing as mp
 from torch.distributed.checkpoint._async_executor import _AsyncCheckpointExecutor
@@ -55,7 +57,7 @@ class _ProcessGroupInitInfo:
     tcp_store_master_port: int
 
     def __init__(self, process_group: Optional[dist.ProcessGroup] = None):
-        self.local_rank = dist.get_node_local_rank(fallback_rank=0)
+        self.local_rank = dist.get_node_local_rank(fallback_rank=dist.get_rank(process_group)%8)
         self.global_rank = dist.get_rank(process_group)
         self.world_size = dist.get_world_size(process_group)
 
@@ -176,13 +178,12 @@ def _checkpointing_subprocess(
             os.environ["LOCAL_RANK"] = str(pg_init_info.local_rank)
             os.environ["RANK"] = str(pg_init_info.global_rank)
             os.environ["WORLD_SIZE"] = str(pg_init_info.world_size)
+            torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
             logger.info(
                 "Initializing dist.ProcessGroup in checkpoint background process"
             )
-            # NOTE: GLOO backend is enforced here.
-            dist.init_process_group(backend=dist.Backend.GLOO)
-            dist.barrier()
+            dist.init_process_group()
 
             logger.info("Checkpoint background process is running...")
             send.put(_CheckpointSaveProcessControlOpts.INIT_COMPLETE)
