MOE: Update global norm calculation for pipeline

When using MoE with MoE-TP disabled, use pipeline parallel group to max or sum
MoE gradients.

This also fixes the behavior for following configuration:
No pipeline, TP enabled, MoE TP disabled.

Signed-off-by: Moshe Island <[email protected]>

Author: mosheisland

deepspeed/runtime/utils.py

@@ -25,7 +25,8 @@
     from torch import inf
 
 from deepspeed.utils import groups, logger
-from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank
+from deepspeed.utils.bwc import (bwc_tensor_model_parallel_rank, bwc_pipeline_parallel_world_size,
+                                 bwc_pipeline_parallel_group)
 from deepspeed.runtime.constants import PIPE_REPLICATED
 from numpy import prod
 from deepspeed.accelerator import get_accelerator
@@ -857,8 +858,16 @@ def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, use_graph=F
             all_norms.append(t.data.abs().max().float())
         total_norm = torch.stack(all_norms).max()
         device_total_norm = total_norm.to(get_accelerator().current_device_name())
+        # Max across model parallel
         if mpu is not None:
-            dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group())
+            # For MoE grads, max over model parallel only if MoE-TP is enabled
+            if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1:
+                dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group())
+            # If MoE grads and MoE-TP disabled, max over pipeline parallel
+            elif bwc_pipeline_parallel_world_size(mpu) > 1:
+                dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=bwc_pipeline_parallel_group(mpu))
+
+        # MoE grads: max across expert parallel group
         if moe_ep_group is not None:
             dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=moe_ep_group)
         total_norm = device_total_norm.to(input_tensors[0].device)
@@ -885,8 +894,16 @@ def _norm_tensors(tensor_list, _compute_buffer, _norm_type):
 
         device_total_norm = compute_buffer[0].float().detach()
 
+        # Sum across model parallel
         if mpu is not None:
-            dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
+            # For MoE grads, sum over model parallel only if MoE-TP is enabled
+            if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1:
+                dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
+            # If MoE grads and MoE-TP disabled, sum over pipeline parallel
+            elif bwc_pipeline_parallel_world_size(mpu) > 1:
+                dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=bwc_pipeline_parallel_group(mpu))
+
+        # MoE grads: sum across expert parallel group
         if moe_ep_group is not None:
             dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=moe_ep_group)
         total_norm = device_total_norm.to(input_tensors[0].device).pow(1. / norm_type)

deepspeed/utils/bwc.py

@@ -76,3 +76,29 @@ def bwc_tensor_model_parallel_group(mpu=None):
     else:
         # Deprecated Megatron and DeepSpeed convention
         return mpu.get_model_parallel_group()
+
+
+def bwc_pipeline_parallel_world_size(mpu=None):
+    """Backwards-compatible way of querying the pipeline parallel world size."""
+    world_size = 1
+    if mpu is not None:
+        if hasattr(mpu, 'get_pipeline_model_parallel_world_size'):
+            # New Megatron and DeepSpeed convention (post pipeline-parallelism release)
+            world_size = mpu.get_pipeline_model_parallel_world_size()
+        elif hasattr(mpu, 'get_pipe_parallel_world_size'):
+            # DeepSpeed Topology
+            world_size = mpu.get_pipe_parallel_world_size()
+    return world_size
+
+
+def bwc_pipeline_parallel_group(mpu=None):
+    """Backwards-compatible way of querying the pipeline parallel group."""
+    if mpu is None:
+        return None
+    if hasattr(mpu, 'get_pipeline_model_parallel_group'):
+        # Megatron
+        return mpu.get_pipeline_model_parallel_group()
+    elif hasattr(mpu, 'get_pipe_parallel_group'):
+        # DeepSpeed Topology
+        return mpu.get_pipe_parallel_group()
+    assert False, 'mpu does not support pipeline parallel group'