add is_first_rank and is_last_rank

xrsrke · Sep 4, 2023 · 9ebeb47 · 9ebeb47
1 parent 79bb372
commit 9ebeb47
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/pipegoose/distributed/parallel_context.py b/pipegoose/distributed/parallel_context.py
@@ -266,6 +266,15 @@ def get_prev_local_rank(self, rank, parallel_mode: ParallelMode) -> int:
         world_size = self.get_world_size(parallel_mode)
         return (rank - 1) % world_size
 
+    def is_first_rank(self, parallel_mode: ParallelMode) -> bool:
+        local_rank = self.get_local_rank(parallel_mode)
+        return local_rank == 0
+
+    def is_last_rank(self, parallel_mode: ParallelMode) -> bool:
+        local_rank = self.get_local_rank(parallel_mode)
+        world_size = self.get_world_size(parallel_mode)
+        return local_rank == world_size - 1
+
     def destroy(self):
         assert self.is_initialized(ParallelMode.GLOBAL), "Global group must be initialized before destroying."
         for mode, group in self._groups.items():

diff --git a/tests/distributed/test_parallel_context.py b/tests/distributed/test_parallel_context.py
@@ -103,6 +103,9 @@ def init_parallel_context(
         prev_local_rank = parallel_context.get_prev_local_rank(local_rank, parallel_mode)
         assert prev_local_rank == LOCAL_RANK_TO_PREV_RANK[world_size][parallel_mode][local_rank]
 
+        assert parallel_context.is_first_rank(parallel_mode) == (local_rank == 0)
+        assert parallel_context.is_last_rank(parallel_mode) == (local_rank == parallel_context.get_world_size(parallel_mode) - 1)
+
     parallel_context.destroy()
 
     if pipeline_parallel_size > 1: