Log rank information (#1068)

facebookresearch · Feb 27, 2025 · 4fc4c80 · 4fc4c80
1 parent 8b85568
commit 4fc4c80
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/src/fairseq2/recipes/common/_gang.py b/src/fairseq2/recipes/common/_gang.py
@@ -86,12 +86,23 @@ def setup_gangs(context: RuntimeContext, recipe_config: object) -> Gangs:
     log.info("Parallel gangs initialized.")
 
     try:
-        return _maybe_setup_fsdp_gangs(context, recipe_config, gangs)
+        gangs = _maybe_setup_fsdp_gangs(context, recipe_config, gangs)
     except GangError as ex:
         raise ProgramError(
             "The hybrid sharded data parallel gangs cannot set up. See the nested exception for details."
         ) from ex
 
+    s = (
+        f"Data: {gangs.dp.rank} | "
+        f"Data/Replicated: {gangs.rdp.rank} | "
+        f"Data/Sharded: {gangs.sdp.rank} | "
+        f"Tensor: {gangs.tp.rank}"
+    )
+
+    log.info("Process Ranks - {}", s)
+
+    return gangs
+
 
 def _maybe_setup_fsdp_gangs(
     context: RuntimeContext, recipe_config: object, gangs: Gangs