diff --git a/deepspeed/moe/layer.py b/deepspeed/moe/layer.py index fd847308b21e..fe1010d2e1a8 100644 --- a/deepspeed/moe/layer.py +++ b/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 diff --git a/deepspeed/moe/mappings.py b/deepspeed/moe/mappings.py new file mode 100644 index 000000000000..38f1630a6703 --- /dev/null +++ b/deepspeed/moe/mappings.py @@ -0,0 +1,108 @@ +''' +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 + + 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 + + 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 + + +def gather_tokens(input_, dim=0): + mpu = deepspeed.utils.groups.mpu + if mpu is None or mpu.get_tensor_model_parallel_world_size() == 1: + # no tensor parallelism for non-experts + return input_ + return _GatherTokens.apply(input_, dim) + + +def drop_tokens(input_, dim=0): + mpu = deepspeed.utils.groups.mpu + if mpu is None or mpu.get_tensor_model_parallel_world_size() == 1: + # no tensor parallelism for non-experts + return input_ + return _DropTokens.apply(input_, dim) diff --git a/deepspeed/moe/sharded_moe.py b/deepspeed/moe/sharded_moe.py index 091ca2b54a24..727b2baec0a6 100644 --- a/deepspeed/moe/sharded_moe.py +++ b/deepspeed/moe/sharded_moe.py @@ -20,6 +20,8 @@ from torch import Tensor from torch.nn import Module import torch.nn.functional as F +from deepspeed.utils import groups +from .mappings import drop_tokens, gather_tokens if TYPE_CHECKING: Base = Module[Tensor] @@ -520,6 +522,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 = drop_tokens(dispatched_input, dim=1) + dispatched_input = _AllToAll.apply(self.ep_group, dispatched_input) if self.wall_clock_breakdown: @@ -548,6 +559,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 = gather_tokens(expert_output, dim=1) + if self.use_tutel: combined_output = self._tutel_dispatcher.decode(expert_output.view(E * C, M)) else: diff --git a/deepspeed/utils/groups.py b/deepspeed/utils/groups.py index c4d4d4d7cfe5..a4d403a7dd0a 100644 --- a/deepspeed/utils/groups.py +++ b/deepspeed/utils/groups.py @@ -37,6 +37,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 +106,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 +121,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 +224,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 +245,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 +392,8 @@ 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 diff --git a/tests/unit/test_moe_tp.py b/tests/unit/test_moe_tp.py new file mode 100644 index 000000000000..60fbe9697da9 --- /dev/null +++ b/tests/unit/test_moe_tp.py @@ -0,0 +1,153 @@ +import torch +import deepspeed +import pytest +from .common import distributed_test +from .simple_model import args_from_dict +from .util import required_torch_version +from deepspeed.moe.layer import MoE + + +@pytest.mark.parametrize("ep_size, tp_size, enable_expert_tp, use_residual", + [ + (1, + 2, + False, + False), + (1, + 2, + True, + False), + (1, + 2, + False, + True), + (1, + 2, + True, + True), + (1, + 4, + False, + False), + (1, + 4, + True, + False), + (1, + 4, + False, + True), + (1, + 4, + True, + True), + (2, + 2, + False, + False), + (2, + 2, + True, + False), + (2, + 2, + False, + True), + (2, + 2, + True, + True), + ]) +def test_moe_tensor_parallel(tmpdir, ep_size, tp_size, enable_expert_tp, 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 + } + } + args = args_from_dict(tmpdir, config_dict) + hidden_dim = 16 + + class MPU(): + def __init__(self, tp_world_size): + self.rank = deepspeed.comm.get_rank() + self.world_size = deepspeed.comm.get_world_size() + self.tp_world_size = tp_world_size + + for i in range(0, self.world_size, tp_world_size): + ranks = range(i, i + tp_world_size) + group = deepspeed.comm.new_group(ranks) + if self.rank in ranks: + self.tp_group = group + + for i in range(0, tp_world_size): + ranks = range(i, self.world_size, tp_world_size) + group = deepspeed.comm.new_group(ranks) + if self.rank in ranks: + self.dp_group = group + + def get_model_parallel_rank(self): + return self.rank % self.tp_world_size + + def get_model_parallel_world_size(self): + return self.tp_world_size + + def get_data_parallel_rank(self): + return self.rank // self.tp_world_size + + def get_data_parallel_world_size(self): + return self.world_size // self.tp_world_size + + def get_data_parallel_group(self): + return self.dp_group + + def get_model_parallel_group(self): + return self.tp_group + + @distributed_test(world_size=[4]) + def _test_moe(args, hidden_dim, ep_size, tp_size, enable_expert_tp, use_residual): + + # TODO: replace this with a true parallel mlp in the future + # and run convergence tests + + tensor_parallel_expert = torch.nn.Sequential( + torch.nn.Linear(hidden_dim, + 4 * hidden_dim // tp_size), + torch.nn.ReLU(), + torch.nn.Linear(4 * hidden_dim // tp_size, + hidden_dim)) + + # set num experts to world size + world_size = deepspeed.comm.get_world_size() + model = MoE( + hidden_size=hidden_dim, + expert=tensor_parallel_expert, + num_experts=world_size, + ep_size=ep_size, + use_residual=use_residual, + enable_expert_tensor_parallelism=enable_expert_tp, + ) + optimizer = torch.optim.AdamW(params=model.parameters()) + model, _, _, _ = deepspeed.initialize(args=args, + model=model, + optimizer=optimizer, + dist_init_required=False, + mpu=MPU(tp_size)) + + assert model.num_local_experts == world_size // ep_size + if enable_expert_tp: + assert deepspeed.utils.groups._get_expert_model_parallel_world_size( + ) == tp_size + else: + assert deepspeed.utils.groups._get_expert_model_parallel_world_size() == 1 + + _test_moe(args=args, + hidden_dim=hidden_dim, + ep_size=ep_size, + tp_size=tp_size, + enable_expert_tp=enable_expert_tp, + use_residual=use_residual)