chore: import upstream snapshot with attribution

deepspeed/inference/v2/model_implementations/common_parameters/moe_parameters.py

@@ -0,0 +1,78 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase, ParamList
+"""
+Moe Parameters
+
+These parameters are compatible with any model inheriting from ``DSMoETransformerModelBase``.
+"""
+
+
+class MoEGatingWeightParameter(ParameterBase):
+    """
+    Gating weight matrix.
+    """
+
+    params: torch.Tensor
+    """
+    Projection matrix from the input activations to the gate logits.
+    """
+
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform_moe_gate_param(self.params)
+
+
+class UnfusedMoEMLP1Parameter(ParameterBase):
+    """
+    This container should be used when the experts are held in separate parameters
+    and need to be joined into a single group.
+    """
+
+    experts: ParamList("n_experts")  # noqa: F821
+
+    def finalize(self) -> torch.Tensor:
+        stacked_experts = torch.stack([p for p in self.experts], dim=0)
+        return self.inference_model.transform_moe_mlp_1_param(stacked_experts)
+
+
+class UnfusedMoEMLP2Parameter(ParameterBase):
+    """
+    This container should be used when the experts are held in separate parameters
+    and need to be joined into a single group.
+    """
+
+    experts: ParamList("n_experts")  # noqa: F821
+
+    def finalize(self) -> torch.Tensor:
+        stacked_experts = torch.stack([p for p in self.experts], dim=0)
+        return self.inference_model.transform_moe_mlp_2_param(stacked_experts)
+
+
+class UnfusedMoEGatedMLPParameter(ParameterBase):
+    """
+    MoE Parameter for a gated activation function in which the gating matrix is not
+    fused in the same parameter as the non-gating matrix.
+
+    This is a stacked version of the ``GatedMLPParameter``. Please see that class for more
+    documentation on the layout of the parameters.
+    """
+
+    gating_experts: ParamList("n_experts")  # noqa: F821
+
+    up_experts: ParamList("n_experts")  # noqa: F821
+
+    def finalize(self) -> torch.Tensor:
+        transposed_experts = []
+        for gate, up in zip(self.gating_experts, self.up_experts):
+            assert gate.shape[0] == up.shape[0], "Gated MLP parameters must have the same number of neurons."
+            total_neurons = gate.shape[0] + up.shape[0]
+            fused_expert = torch.cat([gate, up], dim=-1).reshape(total_neurons, -1)
+            transposed_experts.append(fused_expert)
+
+        stacked_experts = torch.stack(transposed_experts, dim=0)
+        return self.inference_model.transform_moe_mlp_1_param(stacked_experts)