[checkpoint] use gather_tensor in checkpoint and update its unit test (#1339)

2026-05-08 23:05:41 +00:00 · 2022-07-19 14:15:28 +08:00
parent f3ce7b8336
commit f92c100ddd
6 changed files with 209 additions and 91 deletions
--- a/colossalai/tensor/process_group.py
+++ b/colossalai/tensor/process_group.py
@@ -141,9 +141,18 @@ class ProcessGroup:
    def rank(self):
        return self._rank

+    def ranks_in_group(self):
+        return self._rank_list
+
    def world_size(self):
        return self._world_size

+    def tp_rank_list(self):
+        return self._tp_rank_list
+
+    def dp_rank_list(self):
+        return self._dp_rank_list
+
    def tp_local_rank(self):
        return self._rank % self._tp_degree