wehub-resource-sync
92 lines
2.8 KiB
Python
92 lines
2.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from abc import abstractmethod
|
|
from typing import Any, Dict, Optional, Type
|
|
|
|
import torch
|
|
|
|
from deepspeed.runtime.config_utils import DeepSpeedConfigModel
|
|
from ..ds_module import DSModuleBase
|
|
from ..module_registry import DSModuleRegistryBase
|
|
from ..configs import DSMoEConfig
|
|
from ...inference_parameter import InferenceParameter
|
|
|
|
|
|
class DSMoEBase(DSModuleBase):
|
|
"""
|
|
Base mixing for MoE modules. The interface represented by this module is:
|
|
|
|
expert_assignments = gate(hidden_states)
|
|
intermediate = ragged_linear(hidden_states, expert_assignments)
|
|
output = ragged_linear(intermediate, expert_assignments)
|
|
"""
|
|
|
|
@staticmethod
|
|
def config_class() -> Type[DeepSpeedConfigModel]:
|
|
return DSMoEConfig
|
|
|
|
def __init__(self, config: DSMoEConfig, implementation_config: Dict[str, Any]) -> None:
|
|
super().__init__(config, implementation_config)
|
|
|
|
@abstractmethod
|
|
def transform_gate_param(self, param: torch.Tensor) -> InferenceParameter:
|
|
"""
|
|
Perform any necessary transformations of the gate parameter.
|
|
|
|
Args:
|
|
param (torch.Tensor): gate_w (shape: [num_experts, model_dim])
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter:
|
|
"""
|
|
Perform any necessary transformations of the parameter. The specific component
|
|
being transformed should be inferred from the shape of the parameter.
|
|
|
|
Args:
|
|
param (torch.Tensor): One of either mlp_1_w, mlp_1_b
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def transform_moe_mlp_2_param(self, param: torch.Tensor) -> InferenceParameter:
|
|
"""
|
|
Perform any necessary transformations of the parameter. The specified component being
|
|
transformed should be inferred from the shape of the parameter. This interface is
|
|
separate from transform_moe_1_param because the two components may have identical
|
|
shapes.
|
|
|
|
Args:
|
|
param (torch.Tensor): One of either mlp_2_w or mlp_2_b
|
|
"""
|
|
...
|
|
|
|
def forward(self,
|
|
hidden_states: torch.Tensor,
|
|
gate_w: torch.Tensor,
|
|
mlp_1_w: torch.Tensor,
|
|
mlp_2_w: torch.Tensor,
|
|
mlp_1_b: Optional[torch.Tensor] = None,
|
|
mlp_2_b: Optional[torch.Tensor] = None) -> torch.Tensor:
|
|
raise NotImplementedError()
|
|
|
|
@property
|
|
@abstractmethod
|
|
def output(self) -> torch.Tensor:
|
|
"""
|
|
Returns the pre-allocated, padded output Tensor.
|
|
"""
|
|
...
|
|
|
|
|
|
class DSMoERegistry(DSModuleRegistryBase):
|
|
registry: Dict = {}
|
|
|
|
@staticmethod
|
|
def associated_class() -> Type[DSModuleBase]:
|
|
return DSMoEBase
|