| ''' |
| Copyright 2020 The Microsoft DeepSpeed Team |
| ''' |
|
|
| import torch.nn.init as init |
| import torch |
| from torch import nn |
| import torch.distributed as dist |
|
|
|
|
|
|
|
|
| from .gate import TopKGate |
| import copy |
| import typing |
|
|
| from .experts import FusedExperts as Experts |
|
|
|
|
| class TaskMoE(torch.nn.Module): |
| def __init__(self, |
| hidden_size, |
| expert, |
| num_experts=1, |
| k=1, |
| capacity_factor=1., |
| eval_capacity_factor=1., |
| min_capacity=4, |
| noisy_gate_policy: typing.Optional[str] = None, |
| drop_tokens: bool = True, |
| use_rts=True, |
| use_tutel: bool = False, |
| cfg=None): |
| """Initialize an MoE layer. |
| |
| Arguments: |
| hidden_size (int): the hidden dimension of the model, importantly this is also the input and output dimension. |
| |
| expert (torch.nn.Module): the torch module that defines the expert (e.g., MLP, torch.linear). |
| |
| num_experts (int, optional): default=1, the total number of experts per layer. |
| |
| k (int, optional): default=1, top-k gating value, only supports k=1 or k=2. |
| |
| capacity_factor (float, optional): default=1.0, the capacity of the expert at training time. |
| |
| eval_capacity_factor (float, optional): default=1.0, the capacity of the expert at eval time. |
| |
| min_capacity (int, optional): default=4, the minimum capacity per expert regardless of the capacity_factor. |
| |
| noisy_gate_policy (str, optional): default=None, noisy gate policy, valid options are 'Jitter', 'RSample' or 'None'. |
| |
| 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). |
| """ |
|
|
| super().__init__() |
|
|
|
|
| self.num_experts = num_experts |
|
|
| if isinstance(expert, nn.Linear): |
| self.expert_type = 'linear' |
| elif isinstance(expert, nn.MultiheadAttention): |
| self.expert_type = 'attention' |
| else: |
| raise NotImplementedError('please check expert type') |
|
|
| experts = Experts(expert, cfg, num_experts) |
|
|
| self.gate = TopKGate(hidden_size, |
| num_experts, |
| k, |
| noisy_gate_policy, |
| cfg, |
| moe_type=self.expert_type) |
|
|
|
|
| self.experts = experts |
|
|
|
|
|
|
| def forward(self, hidden_states, gate_decision=None, **kwargs): |
| """ MoE forward |
| Arguments: |
| hidden_states (Tensor): input to the layer |
| Returns: |
| A tuple including output |
| * output (Tensor): output of the model |
| """ |
|
|
|
|
| if gate_decision is not None: |
| top_indices, gates = gate_decision |
| else: |
| top_indices, gates = self.gate(hidden_states, **kwargs) |
|
|
| expert_output = self.experts(hidden_states, top_indices, gates, **kwargs) |
|
|
| return expert_output, [top_indices, gates] |
|
|
