Tensor parallelism for Mixture of Experts

deepspeed/moe/layer.py

@@ -26,7 +26,8 @@ def __init__(self,
                  noisy_gate_policy: typing.Optional[str] = None,
                  drop_tokens: bool = True,
                  use_rts=True,
-                 use_tutel: bool = False):
+                 use_tutel: bool = False,
+                 enable_expert_tensor_parallelism: bool = False):
         """Initialize an MoE layer.
 
         Arguments:
@@ -43,20 +44,21 @@ def __init__(self,
             drop_tokens (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent to infinite capacity).
             use_rts (bool, optional): default=True, whether to use Random Token Selection.
             use_tutel (bool, optional): default=False, whether to use Tutel optimizations (if installed).
+            enable_expert_tensor_parallelism (bool, optional): default=False, whether to use tensor parallelism for experts
         """
 
         super(MoE, self).__init__()
 
         self.use_residual = use_residual
-        self.ep_size = min(
-            ep_size,
-            num_experts)  # the ep size should be less than the number of experts
+        self.enable_expert_tensor_parallelism = enable_expert_tensor_parallelism
+        assert num_experts % ep_size == 0, f"Number of experts ({num_experts}) should be divisible by expert parallel size ({ep_size})"
+        self.ep_size = ep_size
         self.expert_group_name = f"ep_size_{self.ep_size}"
         self.num_experts = num_experts
-        self.num_local_experts = 1 if num_experts < ep_size else num_experts // ep_size
+        self.num_local_experts = num_experts // self.ep_size
 
         log_dist(
-            f'Creating MoE layer with num_experts: {num_experts} | num_local_experts: {self.num_local_experts} | expert_parallel_size: {ep_size}',
+            f'Creating MoE layer with num_experts: {num_experts} | num_local_experts: {self.num_local_experts} | expert_parallel_size: {self.ep_size}',
             [0])
 
         assert noisy_gate_policy is None or noisy_gate_policy in ['None', 'Jitter', 'RSample'], \
@@ -91,9 +93,12 @@ def _create_process_groups(self):
             print(
                 f"No existing process group found, creating a new group named: {self.expert_group_name}"
             )
-            if groups.mpu is None:
+            if (groups.mpu is None) or (not self.enable_expert_tensor_parallelism):
+                # Condition 1 - no groups.mpu means no tensor parallelism
+                # Condition 2 - disabling expert tensor parallelism on purpose
                 groups._create_expert_and_data_parallel(self.ep_size)
             else:
+                # expert tensor parallelism is enabled
                 groups._create_expert_data_and_model_parallel(self.ep_size,
                                                               mpu=groups.mpu)
         # Set the group handle for the MOELayer (deepspeed_moe) object

deepspeed/moe/sharded_moe.py

@@ -20,6 +20,7 @@
 from torch import Tensor
 from torch.nn import Module
 import torch.nn.functional as F
+from deepspeed.utils import groups
 
 if TYPE_CHECKING:
     Base = Module[Tensor]
@@ -520,6 +521,15 @@ def forward(self, *input: Tensor, **kwargs: Any) -> Tensor:
         if self.wall_clock_breakdown:
             self.timers('falltoall').start()
 
+        if groups._get_expert_model_parallel_world_size() == 1:
+            # If the non-expert is tensor-parallel, it will create
+            # duplicate tokens on the tensor-parallel ranks.
+            # Since our experts are not tensor-parallel, these duplicates
+            # need to be dropped to ensure correctness.
+            # this also doubles up as a communication optimization as we are
+            # reducing the all-to-all communication volume.
+            dispatched_input = groups.drop_tokens(dispatched_input, dim=1)
+
         dispatched_input = _AllToAll.apply(self.ep_group, dispatched_input)
 
         if self.wall_clock_breakdown:
@@ -548,6 +558,12 @@ def forward(self, *input: Tensor, **kwargs: Any) -> Tensor:
                                               -1,
                                               d_model)
 
+        if groups._get_expert_model_parallel_world_size() == 1:
+            # the dropped duplicate tokens need to be gathered on each
+            # tensor parallel rank again for the tensor-parallel
+            # non-expert of the next layer.
+            expert_output = groups.gather_tokens(expert_output, dim=1)
+
         if self.use_tutel:
             combined_output = self._tutel_dispatcher.decode(expert_output.view(E * C, M))
         else:

deepspeed/utils/groups.py

@@ -28,6 +28,7 @@
 
 from deepspeed.utils import log_dist
 from deepspeed.utils.exceptions import DeprecatedException
+from deepspeed.utils.mappings import _GatherTokens, _DropTokens
 
 # Expert parallel group that the current rank belongs to.
 _EXPERT_PARALLEL_GROUP = {}
@@ -37,6 +38,8 @@
 _WORLD_GROUP = None
 # global object to maintain mpu object if passed by a Megatron client
 mpu = None
+# global object that stores tensor parallel world size for experts
+expert_tensor_parallel_world_size = 1
 
 
 # Deprecated groups initialize function.
@@ -104,7 +107,7 @@ def _create_model_parallel(model_parallel_size_):
     return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP
 
 
-def _create_expert_and_data_parallel(ep_size):
+def _create_expert_and_data_parallel(expert_parallel_size_):
     """
         Create expert and data parallel groups.
 
@@ -119,11 +122,12 @@ def _create_expert_and_data_parallel(ep_size):
     """
     assert dist.is_initialized()
 
-    log_dist(f'Creating expert and data parallel groups with size {ep_size}', ranks=[0])
+    log_dist(
+        f'Creating expert and data parallel groups with size {expert_parallel_size_}',
+        ranks=[0])
     world_size = dist.get_world_size()
     rank = dist.get_rank()
 
-    expert_parallel_size_ = min(ep_size, world_size)
     _ensure_divisibility(world_size, expert_parallel_size_)
 
     group_name = f"ep_size_{expert_parallel_size_}"
@@ -221,11 +225,17 @@ def _create_expert_data_and_model_parallel(expert_parallel_size_, mpu):
     assert dist.is_initialized(), "dist is not initialized"
     model_parallel_size_ = mpu.get_model_parallel_world_size()
 
+    global expert_tensor_parallel_world_size
+    expert_tensor_parallel_world_size = model_parallel_size_
+
     world_size = dist.get_world_size()
     rank = dist.get_rank()
     dp_world_size = mpu.get_data_parallel_world_size()
     dp_rank = mpu.get_data_parallel_rank()
 
+    _ensure_divisibility(world_size, model_parallel_size_)
+    _ensure_divisibility(dp_world_size, expert_parallel_size_)
+
     log_dist(
         f"Creating deepspeed groups with model parallel size {model_parallel_size_}, expert parallel size {expert_parallel_size_}, world size {world_size}, dp world size {dp_world_size}",
         [0])
@@ -236,9 +246,6 @@ def _create_expert_data_and_model_parallel(expert_parallel_size_, mpu):
     _DATA_PARALLEL_GROUP = mpu.get_data_parallel_group()
     _MODEL_PARALLEL_GROUP = mpu.get_model_parallel_group()
 
-    expert_parallel_size_ = min(expert_parallel_size_, dp_world_size)
-    _ensure_divisibility(world_size, expert_parallel_size_)
-
     group_name = f"ep_size_{expert_parallel_size_}"
 
     # Only create groups if they don't already exist
@@ -386,3 +393,22 @@ def _get_data_parallel_rank():
     if mpu is not None:
         return mpu.get_data_parallel_rank()
     return dist.get_rank(group=_get_data_parallel_group())
+
+
+def _get_expert_model_parallel_world_size():
+    global expert_tensor_parallel_world_size
+    return expert_tensor_parallel_world_size
+
+
+def gather_tokens(input_, dim=0):
+    global mpu
+    if mpu is None:  # no tensor parallelism for non-experts
+        return input_
+    return _GatherTokens.apply(input_, dim)
+
+
+def drop_tokens(input_, dim=0):
+    global mpu
+    if mpu is None:  # no tensor parallelism for non-experts
+        return input_
+    return _DropTokens.apply(input_, dim)

deepspeed/utils/mappings.py

@@ -0,0 +1,95 @@
+'''
+Copyright 2022 The Microsoft DeepSpeed Team
+'''
+
+# The file has been adapted from the following Megatron-LM file:
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/mpu/mappings.py
+# Git commit hash: 9dc3c42a84aa656f583703cf8b6b4f79f712b796
+# We retain the following copyright from the original files:
+
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import deepspeed
+
+
+def _gather_tokens(input_, dim=0):
+    """Gather tensors and concatenate them along a dimension"""
+    mpu = deepspeed.utils.groups.mpu
+    if mpu.get_tensor_model_parallel_world_size() == 1:
+        return input_
+
+    input_ = input_.contiguous()
+    # Size and dimension.
+    rank = mpu.get_tensor_model_parallel_rank()
+
+    tensor_list = [
+        torch.empty_like(input_)
+        for _ in range(mpu.get_tensor_model_parallel_world_size())
+    ]
+    tensor_list[rank] = input_
+    deepspeed.comm.all_gather(tensor_list,
+                              input_,
+                              group=mpu.get_tensor_model_parallel_group())
+
+    # Note: torch.cat already creates a contiguous tensor.
+    output = torch.cat(tensor_list, dim=dim).contiguous()
+
+    return output
+
+
+def _drop_tokens(input_, dim=0):
+    """Divide a tensor among the tensor parallel ranks"""
+    mpu = deepspeed.utils.groups.mpu
+    if mpu.get_tensor_model_parallel_world_size() == 1:
+        return input_
+    total_chunks = mpu.get_tensor_model_parallel_world_size()
+    this_chunk = mpu.get_tensor_model_parallel_rank()
+    assert input_.shape[dim] % total_chunks == 0, f"input dimension {dim} ({input_.shape[dim]}) is not divisible by tensor parallel world size ({total_chunks})"
+    chunk_size = input_.shape[dim] // total_chunks
+
+    return torch.narrow(input_, dim, this_chunk * chunk_size, chunk_size)
+
+
+class _GatherTokens(torch.autograd.Function):
+    """All gather tokens among the tensor parallel ranks"""
+    @staticmethod
+    def symbolic(graph, input_, dim):
+        return _gather_tokens(input_, dim)
+
+    @staticmethod
+    def forward(ctx, input_, dim):
+        ctx.dim = dim
+        return _gather_tokens(input_, dim)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return _drop_tokens(grad_output, ctx.dim), None
+
+
+class _DropTokens(torch.autograd.Function):
+    "Divide tokens equally among the tensor parallel ranks"
+
+    @staticmethod
+    def symbolic(graph, input_, dim):
+        return _drop_tokens(input_, dim)
+
+    @staticmethod
+    def forward(ctx, input_, dim):
+        ctx.dim = dim
+        return _drop_tokens(input_, dim)
+
+    @staticmethod
+    def backward(ctx, input_):
+        return _gather_tokens(input_, ctx.dim), None