get_fs_local_rank

allenai · Feb 18, 2024 · 5de468b · 5de468b
1 parent 6765317
commit 5de468b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/olmo/torch_util.py b/olmo/torch_util.py
@@ -55,7 +55,7 @@ def get_fs_local_rank() -> int:
     if all ranks share the same filesystem then `get_fs_local_rank()` will be equivalent to `get_global_rank()`,
     but if nodes do not share the same filesystem then `get_fs_local_rank()` will be equivalent to `get_local_rank()`.
     """
-    return int(os.environ.get("FS_LOCAL_RANK") or get_local_rank())
+    return int(os.environ.get("FS_LOCAL_RANK") or get_global_rank())
 
 
 def move_to_device(o: T, device: torch.device) -> T: