Repository: microsoft/DeepSpeed

Issue:
[BUG] average_tensor coalesing does not take process_group into consideration for MoE
**Describe the bug**
When decision about coalescing adjacent tensors is made right now it is currently only based on `partition_id` [here](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage_1_and_2.py#L923) which is simply target's rank in the `range(dist.get_world_size(group=process_group))`. This means that even though expert data parallel and data parallel both have `partition_id=0` they could belong to different physical ranks in different `process_group`s.

**To Reproduce**
A unit test will be added soon.

**Expected behavior**
`stage_1_and_2`'s `average_tensor` should detect `param_group`s belong to different `process_group`s, e.g. by remembering `prev_process_group`.

**ds_report output**
```
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-devel package with yum
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/azzhipa/.local/lib/python3.9/site-packages/torch']
torch version .................... 1.13.1+cu116
deepspeed install path ........... ['/home/azzhipa/workspace/DeepSpeed/deepspeed']
deepspeed info ................... 0.9.2+f7d71ec1, f7d71ec1, HEAD
torch cuda version ............... 11.6
torch hip version ................ None
nvcc version ..................... 11.6
deepspeed wheel compiled w. ...... torch 1.12, cuda 11.6
```

**Screenshots**
N/A

**System info (please complete the following information):**
 - OS: [e.g. Ubuntu 18.04]
 - GPU count and types: 1 machine with 8x A100s
 - Interconnects (if applicable) 4x 100 Gbps AWS EFA adaptors
 - Python version: 3.9
 - Any other relevant info about your setup: N/A

**Launcher context**
Launching with `deepspeed` runner

**Docker context**
N/A

**Additional context**
N/A



Discussion:

Resolving patch:
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -915,7 +915,7 @@ def average_tensor(self, tensor):
             rank_and_offsets = []
             real_dp_process_group = []
             curr_size = 0
-            prev_id = -1
+            prev_id, prev_process_group = -1, None
 
             process_group = self.dp_process_group
             # count = 0
@@ -958,14 +958,14 @@ def average_tensor(self, tensor):
                         numel = partition_ids_w_offsets[idx + 1][1] - offset
 
                     # Merge bucket ranges if they belong to the same rank
-                    if partition_id == prev_id:
+                    if partition_id == prev_id and process_group == prev_process_group:
                         prev_pid, prev_size, prev_numel = rank_and_offsets[-1]
                         rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel)
                     else:
                         rank_and_offsets.append((partition_id, curr_size, numel))
                         real_dp_process_group.append(process_group)
                     curr_size += numel
-                    prev_id = partition_id
+                    prev_id, prev_process_group = partition_id, process_group
 
             if not self.ipg_bucket_has_moe_params:
                 tensor.div_(dist.get_world_size(group=self.dp_process_group))


Test patch:
diff --git a/tests/unit/moe/test_moe.py b/tests/unit/moe/test_moe.py
--- a/tests/unit/moe/test_moe.py
+++ b/tests/unit/moe/test_moe.py
@@ -8,33 +8,96 @@
 import pytest
 from unit.common import DistributedTest
 from unit.simple_model import SimplePRMoEModel, SimpleMoEModel, sequence_dataloader
+from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer, is_moe_param
 from unit.util import required_torch_version
 
 
 @pytest.mark.parametrize("ep_size", [2, 4])
+@pytest.mark.parametrize("zero_stage", [0, 1, 2])
 @pytest.mark.parametrize("use_residual", [True, False])
 class TestMoE(DistributedTest):
     world_size = 4
 
-    def test(self, ep_size, use_residual):
+    def test(self, ep_size, zero_stage, use_residual):
         if not required_torch_version():
             pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly")
 
-        config_dict = {"train_batch_size": 8, "steps_per_print": 1, "fp16": {"enabled": True}}
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "fp16": {
+                "enabled": True
+            },
+            "zero_optimization": {
+                "stage": zero_stage
+            }
+        }
         hidden_dim = 16
 
         # E+D -- ep_size = 2
         # E only -- ep_size = 4
         model = SimpleMoEModel(hidden_dim, ep_size=ep_size, use_residual=use_residual)
-        optimizer = torch.optim.AdamW(params=model.parameters())
-        model, _, _, _ = deepspeed.initialize(config=config_dict,
-                                              model=model,
-                                              optimizer=optimizer,
-                                              dist_init_required=False)
+        param_group = {'params': [p for p in model.parameters()], 'name': 'random-unique-name'}
+        params = split_params_into_different_moe_groups_for_optimizer(param_group)
+        optimizer = torch.optim.AdamW(params=params)
+        model, optimizer, _, _ = deepspeed.initialize(config=config_dict,
+                                                      model=model,
+                                                      optimizer=optimizer,
+                                                      dist_init_required=False)
         #dist_init_required=False -- parameterize to True/False?
 
         data_loader = sequence_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device)
 
+        def strict_average_tensor(tensor):
+            process_group = optimizer.dp_process_group
+            curr_size = 0
+            pg_offsets = []
+            for i, param, param_id in optimizer.params_in_ipg_bucket:
+                process_group = optimizer.dp_process_group
+                if optimizer.ipg_bucket_has_moe_params:
+                    process_group = optimizer.expert_dp_process_group[param.group_name] if is_moe_param(
+                        param) else optimizer.dp_process_group
+                partition_ids = optimizer.param_to_partition_ids[i][param_id]
+                # Get all partition ids + their offsets
+                partition_offsets = []
+                for partition_id in partition_ids:
+                    offset = optimizer.grad_start_offset[i][partition_id][param_id]
+                    partition_offsets.append(offset)
+                partition_offsets.sort()
+                # Calculate rank and offsets for grad slices
+                for idx, offset in enumerate(partition_offsets):
+                    # Calculate numel for grad slice depending on partition location
+                    if idx == len(partition_offsets) - 1:
+                        # Last partition_id uses its own offset
+                        numel = param.numel() - offset
+                    else:
+                        # Set numel to next partition's offset
+                        numel = partition_offsets[idx + 1] - offset
+                    pg_offsets.append((curr_size, process_group))
+                    curr_size += numel
+
+            def strict_narrow(dim, start, length):
+                lo, hi = 0, len(pg_offsets) - 1
+                while lo < hi:
+                    mi = lo + (hi - lo) // 2
+                    if pg_offsets[mi][0] >= start:
+                        hi = mi
+                    else:
+                        lo = mi + 1
+                curr_slice, reduce_process_group = lo, pg_offsets[lo][1]
+                while curr_slice < len(pg_offsets) and start + length > pg_offsets[curr_slice][0]:
+                    assert reduce_process_group == pg_offsets[curr_slice][
+                        1], "reduce process_group does not match the parameter's process_group"
+                    curr_slice += 1
+                return orig_narrow(dim, start, length)  # real call
+
+            orig_narrow, tensor.narrow = tensor.narrow, strict_narrow
+            type(optimizer).average_tensor(optimizer, tensor)  # real call
+            tensor.narrow = orig_narrow
+
+        if "average_tensor" in dir(optimizer):
+            optimizer.average_tensor = strict_average_tensor
+
         for n, batch in enumerate(data_loader):
             loss = model(batch[0], batch[1])
             model.backward(loss)
diff --git a/tests/unit/simple_model.py b/tests/unit/simple_model.py
--- a/tests/unit/simple_model.py
+++ b/tests/unit/simple_model.py
@@ -71,27 +71,33 @@ class SimpleMoEModel(torch.nn.Module):
 
     def __init__(self, hidden_dim, num_experts=4, ep_size=1, use_residual=False):
         super(SimpleMoEModel, self).__init__()
-        self.linear = torch.nn.Linear(hidden_dim, hidden_dim)
-        expert = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.linear1 = torch.nn.Linear(hidden_dim, hidden_dim)
+        expert = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.Linear(hidden_dim, hidden_dim))
         # using two MoE layers to check implications of sharing a single storage
-        self.linear2 = MoE(hidden_size=hidden_dim,
-                           expert=expert,
-                           ep_size=ep_size,
-                           use_residual=use_residual,
-                           num_experts=num_experts,
-                           k=1)
-        self.linear3 = MoE(hidden_size=hidden_dim,
-                           expert=expert,
-                           ep_size=ep_size,
-                           use_residual=use_residual,
-                           num_experts=num_experts,
-                           k=1)
+        self.moe_1 = MoE(hidden_size=hidden_dim,
+                         expert=expert,
+                         ep_size=ep_size,
+                         use_residual=use_residual,
+                         num_experts=num_experts,
+                         k=1)
+        # interleaving MoE modules with dense to create an opportunity
+        # for gradients to be merged in ZeRO stage 2 average_tensor reduce bucket
+        self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.moe_2 = MoE(hidden_size=hidden_dim,
+                         expert=expert,
+                         ep_size=ep_size,
+                         use_residual=use_residual,
+                         num_experts=num_experts,
+                         k=1)
+        self.linear3 = torch.nn.Linear(hidden_dim, hidden_dim)
         self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
 
     def forward(self, x, y):
-        hidden_dim = self.linear(x)
-        output, _, _ = self.linear2(hidden_dim)
-        output, _, _ = self.linear3(output)
+        hidden_dim = self.linear1(x)
+        output, _, _ = self.moe_1(hidden_dim)
+        output = self.linear2(output)
+        output, _, _ = self.moe_2(output)
+        output = self.linear3(output)
         hidden_dim = hidden_dim + output
         sentence_embed = hidden_dim.mean(1)
         return self.cross_entropy_loss(sentence_embed, y)