[AutoEP] Ensure consistent order of partition groups in all-gather when ZeRO-3

deepspeed/runtime/zero/partition_parameters.py

@@ -1540,7 +1540,14 @@ def _all_gather_coalesced(params, world_size, rank_in_group, use_secondary_tenso
         def all_gather_coalesced(params: Iterable[Parameter],
                                  safe_mode: bool = False,
                                  quantize: bool = False) -> AllGatherCoalescedHandle:
-            params = list(params)
+            # ensure that each rank has params in same order. callers may pass a set, whose
+            # order differs across ranks. the allgather is done by flattening the parameter
+            # list into a single tensor that can be allgathered in a single call - this means
+            # that if each rank gives a list of the same parameters in a different order we
+            # will silently get incorrect parameter values. params of different partition
+            # groups are gathered one group at a time, so a different order would also issue
+            # the per-group collectives in a different order and deadlock.
+            params = sorted(params, key=lambda p: p.ds_id)
             if len(params) == 0:
                 return _no_gather_coalesced(params)
 
@@ -1575,13 +1582,6 @@ def all_gather_coalesced(params: Iterable[Parameter],
                 world_size = self.num_ranks_in_param_group
 
             #pprint(dir(ds_process_group))
-            # ensure that each rank has params in same order. the allgather
-            # is done by flattening the parameter list into a single tensor that
-            # can be allgathered in a single call - this means that if each rank
-            # gives a list of the same parameters in a different order we will
-            # silently get incorrect parameter values, and have very difficult
-            # to debug correctness issues.
-            params = sorted(params, key=lambda p: p.ds_id)
 
             if logger.isEnabledFor(logging.DEBUG):
                 debug_rank0(f"-allgather_coalesced: {[p.ds_id for p in params]}")

tests/unit/runtime/zero/test_zero_stale_handle_wait.py

@@ -84,3 +84,41 @@ def __init__(self):
         victim = _stale_wait_cycle(solo.w, dense.weight)
         assert _released(victim), ("a second wait() on a released parameter marked it "
                                    f"{victim.ds_status.name} over {victim.data.numel()} elements")
+
+
+class TestMultiGroupGatherOrder(DistributedTest):
+    world_size = 2
+
+    def test_rank_dependent_param_order(self):
+        # Expert-parallel parameters are partitioned over their own group while dense parameters
+        # use the whole world, so all_gather_coalesced issues one all-gather per group. The ZeRO-3
+        # prefetcher passes a set, whose order differs across ranks; every rank must still issue
+        # the per-group all-gathers in the same order, or the collectives deadlock. A separate
+        # group over the same two ranks is enough to get two groups.
+        rank = dist.get_rank()
+        pair_group = dist.new_group([0, 1])
+
+        class PairGroupParam(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.w = torch.nn.Parameter(torch.full((HIDDEN, HIDDEN), 2.0))
+                self.w.ds_zero_partition_process_group = pair_group
+
+        class DenseParam(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.w = torch.nn.Parameter(torch.full((HIDDEN, HIDDEN), 3.0))
+
+        with deepspeed.zero.Init(dtype=preferred_dtype()):
+            dense = DenseParam()
+            pair = PairGroupParam()
+
+        params = [pair.w, dense.w]
+        if rank % 2:
+            params.reverse()
+        params[0].all_gather_coalesced(params).wait()
+
+        assert torch.equal(dense.w, torch.full_like(dense.w, 3.0))
+        assert torch.equal(pair.w, torch.full_like(pair.w, 2.0))