chore: import upstream snapshot with attribution

python/paddle/incubate/distributed/models/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023 PaddlePaddle Authors. 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.

python/paddle/incubate/distributed/models/moe/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2022 PaddlePaddle Authors. 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.
+
+from .gate import BaseGate, GShardGate, NaiveGate, SwitchGate  # noqa: F401
+from .grad_clip import ClipGradForMOEByGlobalNorm
+from .moe_layer import MoELayer  # noqa: F401
+
+ClipGradByGlobalNorm = ClipGradForMOEByGlobalNorm

python/paddle/incubate/distributed/models/moe/gate/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2022 PaddlePaddle Authors. 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.
+
+from .base_gate import BaseGate  # noqa: F401
+from .gshard_gate import GShardGate  # noqa: F401
+from .naive_gate import NaiveGate  # noqa: F401
+from .switch_gate import SwitchGate  # noqa: F401

python/paddle/incubate/distributed/models/moe/gate/base_gate.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2022 PaddlePaddle Authors. 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.
+#
+# The file has been adapted from the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/gates/base_gate.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+from paddle import nn
+
+
+class BaseGate(nn.Layer):
+    def __init__(self, num_expert, world_size):
+        super().__init__()
+        self.world_size = world_size
+        self.num_expert = num_expert
+        self.tot_expert = world_size * num_expert
+        self.loss = None
+
+    def forward(self, x):
+        raise NotImplementedError("Please implement the forward function.")
+
+    def set_loss(self, loss):
+        self.loss = loss
+
+    def get_loss(self, clear=True):
+        loss = self.loss
+        if clear:
+            self.loss = None
+        return loss

python/paddle/incubate/distributed/models/moe/gate/gshard_gate.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2022 PaddlePaddle Authors. 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.
+#
+# The file has been adapted from the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/gates/gshard_gate.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import math
+
+import paddle
+import paddle.nn.functional as F
+
+from ..utils import limit_by_capacity
+from .naive_gate import NaiveGate
+
+
+class GShardGate(NaiveGate):
+    def __init__(
+        self,
+        d_model,
+        num_expert,
+        world_size,
+        topk=2,
+        capacity=(1.2, 2.4),
+        random_routing=True,
+        group=None,
+    ):
+        assert topk == 2, "topk should be 2 in gshard"
+        super().__init__(d_model, num_expert, world_size)
+        self.capacity = capacity
+        self.random_routing = random_routing
+        self.group = group
+
+    def forward(self, x):
+        topk_val, topk_idx, gate_score = super().forward(
+            x, return_all_scores=True
+        )
+        s = gate_score.shape[0]
+        top1_idx = topk_idx.flatten()
+        c_e = (
+            paddle.scatter(
+                paddle.zeros(shape=[self.tot_expert]),
+                top1_idx,
+                paddle.ones_like(top1_idx, dtype="float32"),
+                overwrite=False,
+            )
+            / s
+        )
+        m_e = paddle.mean(F.softmax(gate_score, axis=1), axis=0)
+        loss = paddle.mean(c_e * m_e) * (self.num_expert**2)
+        self.set_loss(loss)
+
+        cap_rate = self.capacity[0 if self.training else 1]
+        capacity = math.ceil(cap_rate * x.shape[0])
+        _new_lec, _new_gec, topk_idx = limit_by_capacity(
+            topk_idx,
+            self.num_expert,
+            self.world_size,
+            capacity,
+            group=self.group,
+        )
+
+        if self.random_routing:
+            rand_routing_prob = paddle.rand(
+                shape=[gate_score.shape[0]], dtype="float32"
+            )
+            topk_idx = paddle.distributed.models.moe.utils._random_routing(
+                topk_idx, topk_val, rand_routing_prob
+            )
+        return topk_val, topk_idx

python/paddle/incubate/distributed/models/moe/gate/naive_gate.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2022 PaddlePaddle Authors. 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.
+#
+# The file has been adapted from the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/gates/naive_gate.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import paddle
+from paddle import nn
+
+from .base_gate import BaseGate
+
+
+class NaiveGate(BaseGate):
+    def __init__(self, d_model, num_expert, world_size, topk=2):
+        super().__init__(num_expert, world_size)
+        self.gate = nn.Linear(d_model, self.tot_expert)
+        self.gate.weight.name = "gate_" + self.gate.weight.name
+        self.gate.bias.name = "gate_" + self.gate.bias.name
+        self.top_k = topk
+
+    def forward(self, inp, return_all_scores=False):
+        gate = self.gate(inp)
+        gate_top_k_val, gate_top_k_idx = paddle.topk(
+            gate, k=self.top_k, axis=-1, largest=True, sorted=False
+        )
+
+        if return_all_scores:
+            return gate_top_k_val, gate_top_k_idx, gate
+        return gate_top_k_val, gate_top_k_idx

python/paddle/incubate/distributed/models/moe/gate/switch_gate.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2022 PaddlePaddle Authors. 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.
+#
+# The file has been adapted from the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/gates/switch_gate.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import math
+
+import paddle
+import paddle.nn.functional as F
+
+from ..utils import limit_by_capacity
+from .naive_gate import NaiveGate
+
+
+class SwitchGate(NaiveGate):
+    def __init__(
+        self,
+        d_model,
+        num_expert,
+        world_size,
+        topk=1,
+        switch_eps=0.1,
+        capacity=(1.2, 2.4),
+        group=None,
+    ):
+        assert topk == 1, "topk should be 1 in switch"
+        super().__init__(d_model, num_expert, world_size, topk=1)
+        self.switch_eps = switch_eps
+        self.capacity = capacity
+        self.group = group
+
+    def forward(self, inp):
+        score = self.gate(inp)
+
+        if self.training:
+            noise = paddle.rand(shape=score.shape)
+            noise = noise * 2 * self.switch_eps + 1.0 - self.switch_eps
+            score += noise
+
+        score = F.softmax(score, axis=-1)
+        top1_score, top1_idx = paddle.topk(score, k=1, axis=-1, largest=True)
+
+        cap_rate = self.capacity[0 if self.training else 1]
+        capacity = math.ceil(cap_rate * inp.shape[0])
+        _new_lec, _new_gec, top1_idx = limit_by_capacity(
+            top1_idx,
+            self.num_expert,
+            self.world_size,
+            capacity,
+            group=self.group,
+        )
+        valid_idx = top1_idx[top1_idx > -1]
+        valid_idx_tmp = paddle.reshape(valid_idx, shape=[len(valid_idx), 1])
+        fraction_expert = (
+            paddle.scatter_nd_add(
+                x=paddle.zeros(shape=[self.tot_expert]),
+                index=valid_idx_tmp,
+                updates=paddle.ones_like(
+                    valid_idx, dtype=paddle.float32
+                ).reshape(shape=[len(valid_idx)]),
+            )
+            / valid_idx.numel()
+        )
+        prob_expert = score.sum(axis=0) / valid_idx.numel()
+        loss = (fraction_expert * prob_expert).sum() * self.tot_expert
+        self.set_loss(loss)
+
+        return top1_score, top1_idx

python/paddle/incubate/distributed/models/moe/grad_clip.py

@@ -0,0 +1,238 @@
+#   Copyright (c) 2018 PaddlePaddle Authors. 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 paddle
+import paddle.distributed as dist
+from paddle.autograd import no_grad
+from paddle.framework import core
+from paddle.nn import clip
+from paddle.nn.clip import ClipGradBase, _squared_l2_norm
+
+
+class ClipGradForMOEByGlobalNorm(ClipGradBase):
+    r"""
+    The Algorithm is the same as paddle.nn.ClipGradByGlobalNorm
+    Given a list of Tensor :math:`t\_list` , calculate the global norm for the elements of all tensors in
+    :math:`t\_list` , and limit it to ``clip_norm`` .
+
+    - If the global norm is greater than ``clip_norm`` , all elements of :math:`t\_list` will be compressed by a ratio.
+
+    - If the global norm is less than or equal to ``clip_norm`` , nothing will be done.
+
+    The list of Tensor :math:`t\_list` is not passed from this class, but the gradients of all parameters set in ``optimizer``.
+    If ``need_clip`` of specific param is ``False`` in its ``ParamAttr``, then the gradients of this param will not be clipped.
+
+    Gradient clip will takes effect after being set in ``optimizer`` , see the document ``optimizer``
+    (for example: :ref:`api_paddle_optimizer_SGD`).
+
+    The clipping formula is:
+
+    .. math::
+
+        t\_list[i] = t\_list[i] * \frac{clip\_norm}{\max(global\_norm, clip\_norm)}
+
+    where:
+
+    .. math::
+
+        global\_norm = \sqrt{\sum_{i=0}^{N-1}(l2norm(t\_list[i]))^2}
+
+    Note:
+        ``need_clip`` of ``ClipGradyGlobalNorm`` HAS BEEN DEPRECATED since 2.0.
+        Please use ``need_clip`` in ``ParamAttr`` to specify the clip scope.
+
+    Reference:
+        https://github.com/laekov/fastmoe/blob/master/examples/megatron/clip-grad-v2.2.patch
+        Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+
+
+    Args:
+        clip_norm (float): The maximum norm value.
+        is_expert_param_func (function): a function to decide whether a param should be put into moe_params_grads
+        moe_group (Group): group for moe experts communication.
+        group_name (str, optional): The group name for this clip. Default value is ``default_moe_group``.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x = paddle.uniform([10, 10], min=-1.0, max=1.0, dtype='float32')
+            >>> linear = paddle.nn.Linear(
+            ...     in_features=10,
+            ...     out_features=10,
+            ...     weight_attr=paddle.ParamAttr(need_clip=True),
+            ...     bias_attr=paddle.ParamAttr(need_clip=False),
+            ... )
+            >>> out = linear(x)
+            >>> loss = paddle.mean(out)
+            >>> loss.backward()
+
+            >>> clip = paddle.nn.ClipGradByGlobalNorm(
+            ...     clip_norm=1.0
+            ... )  # Cause paddle.nn hasn't this interface, so we use ClipGradByGlobalNorm here.
+            >>> sdg = paddle.optimizer.SGD(learning_rate=0.1, parameters=linear.parameters(), grad_clip=clip)
+            >>> sdg.step()
+    """
+
+    def __init__(
+        self,
+        clip_norm,
+        is_expert_param_func=None,
+        moe_group=None,
+        group_name="default_moe_group",
+    ):
+        super().__init__()
+        self.clip_norm = float(clip_norm)
+        self.group_name = group_name
+        self.moe_group = moe_group
+        if moe_group is not None and moe_group.nranks > 1:
+            assert is_expert_param_func is not None, (
+                "When moe group size > 1, a function for selecting expert params must be specified."
+            )
+        self.is_expert_param_func = is_expert_param_func
+
+    def __str__(self):
+        return f"Gradient Clip By GlobalNorm, global_norm={self.clip_norm:f}"
+
+    @staticmethod
+    def get_l2_norm_pow(params_grads, sum_dtype=None):
+        sum_square_list = []
+        sum_square_list_fp16 = []
+        sum_square_list_fp32 = []
+        for p, g in params_grads:
+            if g is None:
+                continue
+            if getattr(p, 'need_clip', True) is False:
+                continue
+            merge_grad = g
+            if g.type == core.VarDesc.VarType.SELECTED_ROWS:
+                merge_grad = clip.merge_selected_rows(g)
+                merge_grad = clip.get_tensor_from_selected_rows(merge_grad)
+            sum_square = _squared_l2_norm(merge_grad)
+            if sum_square.dtype == paddle.float16:
+                sum_square_list_fp16.append(sum_square)
+            elif sum_square.dtype == paddle.float32:
+                sum_square_list_fp32.append(sum_square)
+            else:
+                sum_square_list.append(sum_square)
+
+        # all parameters have been filtered out
+        if (
+            len(sum_square_list)
+            + len(sum_square_list_fp16)
+            + len(sum_square_list_fp32)
+            == 0
+        ):
+            return None, None
+        assert sum_dtype in [
+            "float64",
+            "float32",
+            None,
+        ], "sum's type must be float64/ float32 / None"
+        if sum_dtype != "float64":
+            sum_dtype = 'float64' if len(sum_square_list) > 0 else "float32"
+
+        global_norm_var = []
+        if len(sum_square_list_fp16) > 0:
+            global_norm_var_fp16 = paddle.add_n(sum_square_list_fp16)
+            global_norm_var.append(global_norm_var_fp16.astype(sum_dtype))
+        if len(sum_square_list_fp32) > 0:
+            global_norm_var_fp32 = paddle.add_n(sum_square_list_fp32)
+            if sum_dtype == 'float32':
+                global_norm_var.append(global_norm_var_fp32)
+            else:
+                global_norm_var.append(global_norm_var_fp32.astype(sum_dtype))
+        if len(sum_square_list) > 0:
+            global_norm_var_fp64 = paddle.add_n(sum_square_list)
+            global_norm_var.append(global_norm_var_fp64)
+        global_norm_var = paddle.add_n(global_norm_var)
+        return global_norm_var, sum_dtype
+
+    @no_grad()
+    def _dygraph_clip(self, params_grads):
+        normal_params_grads = []
+        moe_params_grads = []
+
+        # separate moe params from normal params
+        if self.moe_group is not None and self.moe_group.nranks > 1:
+            for p, g in params_grads:
+                if self.is_expert_param_func(p):
+                    moe_params_grads.append((p, g))
+                else:
+                    normal_params_grads.append((p, g))
+        else:
+            normal_params_grads = params_grads
+
+        # why to return sum_dtype?
+        # we will call `get_l2_norm_pow` twice and the precisions may be different.
+        # For convenience and simplification, we use sum_dtype directly instead of global_norm_var_normal.dtype
+        global_norm_var_normal, sum_dtype = self.get_l2_norm_pow(
+            normal_params_grads
+        )
+        global_norm_var_moe = None
+        if len(moe_params_grads) > 0:
+            global_norm_var_moe, _ = self.get_l2_norm_pow(
+                moe_params_grads, sum_dtype
+            )
+            if global_norm_var_moe is not None:
+                dist.all_reduce(
+                    global_norm_var_moe,
+                    op=dist.ReduceOp.SUM,
+                    group=self.moe_group,
+                )
+
+        if global_norm_var_normal is None and global_norm_var_moe is None:
+            return params_grads
+        elif global_norm_var_normal is None:
+            global_norm_var = global_norm_var_moe
+        elif global_norm_var_moe is None:
+            global_norm_var = global_norm_var_normal
+        else:
+            if global_norm_var_normal.dtype != global_norm_var_moe.dtype:
+                # compared with normal norm, moe norm is the later one,
+                # so its precision is no lower than normal norm
+                global_norm_var_normal = global_norm_var_normal.astype(
+                    global_norm_var_moe.dtype
+                )
+            global_norm_var = global_norm_var_normal + global_norm_var_moe
+
+        params_and_grads = []
+        global_norm_var = paddle.sqrt(global_norm_var)
+        max_global_norm = paddle.full(
+            shape=[1], dtype=global_norm_var.dtype, fill_value=self.clip_norm
+        )
+        clip_var = paddle.divide(
+            x=max_global_norm,
+            y=paddle.maximum(x=global_norm_var, y=max_global_norm),
+        )
+        for p, g in params_grads:
+            if g is None:
+                continue
+            if getattr(p, 'need_clip', True) is False:
+                params_and_grads.append((p, g))
+                continue
+            # TODO(wangxi): use inplace elementwise_mul
+            clip_input = (
+                clip_var.astype('float16')
+                if g.dtype == paddle.float16
+                else clip_var
+            )
+            new_grad = paddle.multiply(x=g, y=clip_input)
+            params_and_grads.append((p, new_grad))
+        return params_and_grads
+
+
+ClipGradByGlobalNorm = ClipGradForMOEByGlobalNorm

python/paddle/incubate/distributed/models/moe/moe_layer.py

@@ -0,0 +1,503 @@
+# Copyright (c) 2021 PaddlePaddle Authors. 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.
+#
+# The file has been adapted from the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/layers.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import os
+
+import numpy as np
+
+import paddle
+from paddle import nn
+from paddle.autograd import PyLayer
+from paddle.distributed.utils.moe_utils import global_gather, global_scatter
+from paddle.distributed.utils.nccl_utils import check_nccl_version_for_p2p
+from paddle.framework import in_dynamic_mode
+from paddle.incubate.distributed.fleet import recompute_hybrid
+
+from .gate import BaseGate, GShardGate, NaiveGate, SwitchGate
+from .utils import count_by_gate
+
+
+def _local_scatter(inp, pos):
+    if pos.shape != [0]:
+        inp_buf = paddle.index_select(inp, pos, 0)
+    else:
+        inp_buf = paddle.empty([0, inp.shape[1]], dtype=inp.dtype)
+    return inp_buf
+
+
+def _local_gather(inp, pos, out_batch_size, maybe_overlap=True):
+    if pos.shape != [0]:
+        origin_dtype = inp.dtype
+        inp = paddle.cast(inp, dtype="float32")
+        inp_buf = paddle.scatter(
+            paddle.zeros(
+                shape=[out_batch_size, inp.shape[-1]], dtype="float32"
+            ),
+            pos,
+            inp,
+            overwrite=True,
+        )
+        inp_buf = paddle.cast(inp_buf, dtype=origin_dtype)
+    else:
+        inp_buf = paddle.zeros([out_batch_size, inp.shape[-1]], dtype=inp.dtype)
+    return inp_buf
+
+
+def _all_gather(tensor, group=None, use_calc_stream=True):
+    if group is not None and not group.is_member():
+        return
+
+    if in_dynamic_mode():
+        group = (
+            paddle.distributed.collective._get_default_group()
+            if group is None
+            else group
+        )
+        tensor_shape = list(tensor.shape)
+        tensor_shape[0] *= group.nranks
+        out = paddle.empty(tensor_shape, tensor.dtype)
+
+        task = group.process_group.all_gather(tensor, out)
+        task.wait()
+        return out
+    else:
+        ring_id = 0 if group is None else group.id
+        nranks = (
+            paddle.distributed.collective._get_global_group().nranks
+            if group is None
+            else group.nranks
+        )
+        return paddle._C_ops.all_gather(
+            tensor,
+            ring_id,
+            nranks,
+        )
+
+
+class MoEScatter(PyLayer):
+    r"""
+    Scatter input samples from [batch x sequences] to contiguous alone experts.
+    If `world_size` is greater than 1, the samples will first be locally
+    scattered, and then exchanged across workers.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        inp,
+        pos,
+        local_expert_count,
+        global_expert_count,
+        fwd_batch_size,
+        world_size,
+        group=None,
+    ):
+        local_input_buf = _local_scatter(inp, pos)
+        if world_size > 1:
+            global_input_buf = global_scatter(
+                local_input_buf,
+                local_expert_count,
+                global_expert_count,
+                group=group,
+            )
+        else:
+            global_input_buf = local_input_buf
+
+        ctx.moe_args = inp.shape[0], world_size, group
+
+        variables = (pos, local_expert_count, global_expert_count)
+        ctx.save_for_backward(*variables)
+        return global_input_buf
+
+    @staticmethod
+    def backward(ctx, grad):
+        (pos, local_expert_count, global_expert_count) = ctx.saved_tensor()
+        (inp_batch_size, world_size, group) = ctx.moe_args
+
+        if world_size > 1:
+            local_grad_in = global_gather(
+                grad, local_expert_count, global_expert_count, group=group
+            )
+        else:
+            local_grad_in = grad
+        grad_in = _local_gather(local_grad_in, pos, inp_batch_size)
+        return grad_in, None, None, None
+
+
+class MoEGather(PyLayer):
+    r"""
+    Gather output samples from contiguous alone experts back to [batch x
+    sequences]. Works symmetrically with MoEScatter.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        global_output_buf,
+        pos,
+        local_expert_count,
+        global_expert_count,
+        local_batch_size,
+        world_size,
+        group=None,
+    ):
+        if world_size > 1:
+            local_output_buf = global_gather(
+                global_output_buf,
+                local_expert_count,
+                global_expert_count,
+                group=group,
+            )
+        else:
+            local_output_buf = global_output_buf
+        output = _local_gather(
+            local_output_buf, pos, local_batch_size, maybe_overlap=False
+        )
+
+        ctx.moe_args = (global_output_buf.shape[0], world_size, group)
+        variables = (pos, local_expert_count, global_expert_count)
+        ctx.save_for_backward(*variables)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        pos, local_expert_count, global_expert_count = ctx.saved_tensor()
+        fwd_batch_size, world_size, group = ctx.moe_args
+        grad_out_buf = _local_scatter(grad_out, pos)
+        if world_size > 1:
+            global_grad_out_buf = global_scatter(
+                grad_out_buf,
+                local_expert_count,
+                global_expert_count,
+                group=group,
+            )
+        else:
+            global_grad_out_buf = grad_out_buf
+        return global_grad_out_buf, None, None, None
+
+
+class AllGather(PyLayer):
+    r"""
+    A wrapper for the All-Gather function to support auto-differentiation.
+    """
+
+    @staticmethod
+    def forward(ctx, inp, rank, world_size, group):
+        tensor_list = []
+        paddle.distributed.all_gather(tensor_list, inp, group=group)
+        output = paddle.concat(tensor_list, axis=0)
+        ctx.args = rank, inp.shape[0]
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        rank, dim0 = ctx.args
+        return paddle.slice(
+            grad_out, axes=[0], starts=[rank * dim0], ends=[(rank + 1) * dim0]
+        )
+
+
+class Slice(PyLayer):
+    r"""
+    A wrapper for the Slice function to support auto-differentiation.
+    """
+
+    @staticmethod
+    def forward(ctx, inp, rank, world_size, group):
+        B = inp.shape[0]
+        local_batch_size = B // world_size
+        batch_start = local_batch_size * rank
+        batch_end = min(batch_start + local_batch_size, B)
+        inp = paddle.slice(
+            inp, axes=[0], starts=[batch_start], ends=[batch_end]
+        )
+        ctx.args = world_size, group
+        return inp
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        world_size, group = ctx.args
+        return _all_gather(grad_out, group=group)
+
+
+def prepare_forward(gate, num_expert, world_size, moe_group):
+    pos, local_expert_count, global_expert_count = count_by_gate(
+        gate, num_expert, world_size, group=moe_group
+    )
+    with paddle.no_grad():
+        fwd_expert_count = global_expert_count.reshape_(
+            [world_size, num_expert]
+        ).sum(axis=0)
+        fwd_batch_size = int(fwd_expert_count.sum().item())
+    return (
+        pos,
+        local_expert_count,
+        global_expert_count,
+        fwd_expert_count,
+        fwd_batch_size,
+    )
+
+
+class MoELayer(nn.Layer):
+    """MoE Layer
+    Args:
+        d_model (int): Model dimension.
+        experts (nn.LayerList): Expert networks list.
+        gate (dict|NaiveGate|SwitchGate|NaiveGate):
+
+            - If gate is a dict:
+              gate is a gate network config, containing 2 keys:
+              `type` (str) value can be: "naive", "gshard", "switch" or None, default is "gshard".
+              `top_k` (int) Default value is 2.
+            else gate is an instance of NaiveGate|SwitchGate|NaiveGate:
+
+        moe_group: moe group for experts communication.
+        mp_group: mp group for mp communication.
+        recompute_interval (int, optional): Whether to use recompute, default 0, means to disable recompute.
+        recompute_ctx (dict, optional): The context for recompute, if recompute_interval > 1, recompute_ctx must be given.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('Until Distributed move successfully, just skip it')
+            >>> from paddle.nn import layer, LayerList
+            >>> from paddle.distributed.moe import MoElayer
+            >>> from paddle.distributed.collective import Group
+            >>> from paddle.distributed import fleet
+
+            >>> moe_group = Group(
+            ...     fleet.worker_index(),
+            ...     0,
+            ...     list(range(fleet.worker_num())),
+            ... )
+            >>> mp_group = None
+
+            >>> num_experts = 8
+            >>> dim_feedforward = 512
+            >>> d_model = 8
+            >>> top_k = 2
+
+            >>> class ExpertLayer(Layer):
+            ...     def __init__(self, d_model, d_hidden, name=None, rank=0, windex=0, num_expert=1):
+            ...         super().__init__()
+            ...         self.htoh4 = nn.Linear(d_model, d_hidden)
+            ...         self.h4toh = nn.Linear(d_hidden, d_model)
+
+            ...     def forward(self, x):
+            ...         x = self.htoh4(x)
+            ...         x = self.h4toh(x)
+            ...         return x
+
+            >>> gate_config = {
+            ...     "type": "gshard",
+            ...     "top_k": top_k,
+            ... }
+
+            >>> experts_list = LayerList()
+            >>> for expi in range(num_experts):
+            ...     exp_layer = ExpertLayer(d_model, dim_feedforward // top_k, windex=expi, num_expert=num_experts)
+            ...     experts_list.append(exp_layer)
+
+            >>> moeLayer = MoELayer(
+            ...     d_model=d_model,
+            ...     experts=experts_list,
+            ...     gate=gate_config,
+            ...     moe_group=moe_group,
+            ...     mp_group=mp_group,
+            ...     recompute_interval=0,
+            ... )
+
+    """
+
+    def __init__(
+        self,
+        d_model,
+        experts,
+        gate=None,
+        moe_group=None,
+        mp_group=None,
+        recompute_interval=0,
+        recompute_ctx=None,
+    ):
+        super().__init__()
+
+        self.recompute_ctx = recompute_ctx
+
+        if gate is None:
+            gate = {}
+
+        assert isinstance(gate, (dict, BaseGate)), (
+            "gate config' type must be dict or an instance of BaseGate"
+        )
+        # only support mp/dp
+        self.group = moe_group
+
+        self.world_size = 1
+        if self.group is not None:
+            self.world_size = self.group.nranks
+        self.num_expert = len(experts)
+        self.recompute_interval = recompute_interval
+        assert experts is not None
+        self.experts = experts
+
+        if (
+            self.world_size > 1
+            and os.getenv("PADDLE_DISTRI_BACKEND", None) != "xccl"
+        ):
+            check_nccl_version_for_p2p()
+
+        self.mp_group = mp_group
+        self.d_model = d_model
+        if isinstance(gate, dict):
+            self.top_k = gate.get("top_k", 2)
+            gate = gate.get("type", "gshard")
+            if gate == "naive" or gate is None:
+                gate = NaiveGate(
+                    self.d_model,
+                    num_expert=len(experts),
+                    world_size=self.world_size,
+                    topk=self.top_k,
+                )
+            elif gate == "gshard":
+                gate = GShardGate(
+                    self.d_model,
+                    num_expert=len(experts),
+                    world_size=self.world_size,
+                    topk=self.top_k,
+                    group=self.group,
+                )
+            elif gate == "switch":
+                gate = SwitchGate(
+                    self.d_model,
+                    num_expert=len(experts),
+                    world_size=self.world_size,
+                    topk=self.top_k,
+                    group=self.group,
+                )
+            else:
+                raise AssertionError(
+                    f"We only support naive gate,                                 gshard gate and switch gate,                                 but you choose {gate} gate."
+                )
+        elif isinstance(gate, NaiveGate):
+            self.top_k = gate.top_k
+        elif isinstance(gate, BaseGate):
+            raise TypeError(f"Unimplemented gate type: {type(gate)}")
+        else:
+            raise TypeError("gate's type must be either dict or moe.BaseGate")
+        self.gate = gate
+
+    def forward(self, inp):
+        # inp shape: b * s * m
+        assert len(inp.shape) == 3
+        origin_shape = inp.shape
+        inp = inp.reshape_([-1, origin_shape[2]])
+
+        mp_rank = 0
+        mp_size = 1
+        if self.mp_group is not None:
+            mp_rank = self.mp_group.rank
+            mp_size = self.mp_group.nranks
+        if mp_size > 1:
+            inp = Slice.apply(inp, mp_rank, mp_size, self.mp_group)
+        value, gate = self.gate(inp)
+
+        (
+            pos,
+            local_expert_count,
+            global_expert_count,
+            fwd_expert_count,
+            fwd_batch_size,
+        ) = prepare_forward(gate, self.num_expert, self.world_size, self.group)
+
+        topk = 1
+        if len(gate.shape) == 2:
+            topk = gate.shape[1]
+
+        if pos.shape != [0]:
+            temp_pos = pos // topk
+        else:
+            temp_pos = pos
+        assert topk == self.top_k
+
+        x = MoEScatter.apply(
+            inp,
+            temp_pos,
+            local_expert_count,
+            global_expert_count,
+            fwd_batch_size,
+            self.world_size,
+            self.group,
+        )
+
+        d_model = self.d_model
+
+        def experts_fwd(x, fwd_expert_count, experts):
+            if x.shape[0] == 0:
+                return x
+            y = []
+            last_index = 0
+            assert isinstance(fwd_expert_count, np.ndarray)
+            assert len(experts) == len(fwd_expert_count)
+            for idx, expert_count in enumerate(fwd_expert_count):
+                if expert_count <= 0:
+                    continue
+                y.append(
+                    experts[idx](x[last_index : expert_count + last_index])
+                )
+                last_index = expert_count + last_index
+            return paddle.concat(y, axis=0)
+
+        if self.recompute_interval <= 0 or x.shape[0] == 0:
+            x = experts_fwd(x, fwd_expert_count.numpy(), self.experts)
+        else:
+            x = recompute_hybrid(
+                self.recompute_ctx,
+                experts_fwd,
+                x,
+                fwd_expert_count.numpy(),
+                self.experts,
+            )
+
+        out_batch_size = inp.shape[0]
+        if len(gate.shape) == 2:
+            out_batch_size *= gate.shape[1]
+
+        x = MoEGather.apply(
+            x,
+            pos,
+            local_expert_count,
+            global_expert_count,
+            out_batch_size,
+            self.world_size,
+            self.group,
+        )
+
+        x = x.reshape([-1, self.top_k, d_model])
+        value = value.reshape([x.shape[0], 1, self.top_k])
+        x = paddle.bmm(value, x).reshape([-1, d_model])
+
+        if mp_size > 1:
+            x = AllGather.apply(x, mp_rank, mp_size, self.mp_group)
+
+        x = paddle.reshape_(x, origin_shape)
+
+        return x

python/paddle/incubate/distributed/models/moe/utils.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2022 PaddlePaddle Authors. 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.
+#
+# The file has been adapted from the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/functions.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import paddle
+from paddle.distributed.models.moe.utils import (
+    _assign_pos,
+    _limit_by_capacity,
+    _number_count,
+    _prune_gate_by_capacity,
+)
+from paddle.framework import in_dynamic_mode
+
+
+def _alltoall(in_tensor_list, group=None, use_calc_stream=True):
+    if group is not None and not group.is_member():
+        return
+
+    if in_dynamic_mode():
+        group = (
+            paddle.distributed.collective._get_default_group()
+            if group is None
+            else group
+        )
+        out = paddle.empty(in_tensor_list.shape, in_tensor_list.dtype)
+        task = group.process_group.alltoall(out, in_tensor_list)
+        task.wait()
+        return out
+    else:
+        ring_id = 0 if group is None else group.id
+        return paddle._C_ops.all_to_all(in_tensor_list, ring_id)
+
+
+def count_by_gate(gate, num_expert, world_size, require_pos=True, group=None):
+    total_expert_count = num_expert * world_size
+    with paddle.no_grad():
+        local_expert_count = _number_count(gate, total_expert_count)
+
+        if world_size > 1:
+            global_expert_count = _alltoall(local_expert_count, group=group)
+        else:
+            global_expert_count = local_expert_count
+        if not require_pos:
+            pos = None
+        else:
+            lec_cum = paddle.cumsum(local_expert_count, axis=0)
+            pos = _assign_pos(gate, lec_cum)
+    return pos, local_expert_count, global_expert_count
+
+
+def limit_by_capacity(topk_idx, num_expert, world_size, capacity, group=None):
+    with paddle.no_grad():
+        capacity = (
+            paddle.ones(shape=[num_expert], dtype=paddle.int64) * capacity
+        )
+        pos, lec, gec = count_by_gate(
+            topk_idx, num_expert, world_size, require_pos=False, group=group
+        )
+        new_gec = _limit_by_capacity(gec, capacity, world_size)
+        if world_size > 1:
+            assert group.nranks == world_size
+            new_lec = _alltoall(new_gec, group=group)
+        else:
+            new_lec = new_gec
+
+        topk_idx = _prune_gate_by_capacity(
+            topk_idx, new_lec, num_expert, world_size
+        )
+
+    return new_lec, new_gec, topk_idx