wehub-resource-sync
79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
# 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)
|