116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
# Copyright (c) 2019 Shigeki Karita
|
|
# 2020 Mobvoi Inc (Binbin Zhang)
|
|
#
|
|
# 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.
|
|
"""Positionwise feed forward layer definition."""
|
|
|
|
import torch
|
|
|
|
|
|
class PositionwiseFeedForward(torch.nn.Module):
|
|
"""Positionwise feed forward layer.
|
|
|
|
FeedForward are appied on each position of the sequence.
|
|
The output dim is same with the input dim.
|
|
|
|
Args:
|
|
idim (int): Input dimenstion.
|
|
hidden_units (int): The number of hidden units.
|
|
dropout_rate (float): Dropout rate.
|
|
activation (torch.nn.Module): Activation function
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
idim: int,
|
|
hidden_units: int,
|
|
dropout_rate: float,
|
|
activation: torch.nn.Module = torch.nn.ReLU(),
|
|
):
|
|
"""Construct a PositionwiseFeedForward object."""
|
|
super(PositionwiseFeedForward, self).__init__()
|
|
self.w_1 = torch.nn.Linear(idim, hidden_units)
|
|
self.activation = activation
|
|
self.dropout = torch.nn.Dropout(dropout_rate)
|
|
self.w_2 = torch.nn.Linear(hidden_units, idim)
|
|
|
|
def forward(self, xs: torch.Tensor) -> torch.Tensor:
|
|
"""Forward function.
|
|
|
|
Args:
|
|
xs: input tensor (B, L, D)
|
|
Returns:
|
|
output tensor, (B, L, D)
|
|
"""
|
|
return self.w_2(self.dropout(self.activation(self.w_1(xs))))
|
|
|
|
|
|
class MoEFFNLayer(torch.nn.Module):
|
|
"""
|
|
Mixture of expert with Positionwise feed forward layer
|
|
See also figure 1 in https://arxiv.org/pdf/2305.15663.pdf
|
|
The output dim is same with the input dim.
|
|
|
|
Modified from https://github.com/Lightning-AI/lit-gpt/pull/823
|
|
https://github.com/mistralai/mistral-src/blob/b46d6/moe_one_file_ref.py#L203-L219
|
|
Args:
|
|
n_expert: number of expert.
|
|
n_expert_per_token: The actual number of experts used for each frame
|
|
idim (int): Input dimenstion.
|
|
hidden_units (int): The number of hidden units.
|
|
dropout_rate (float): Dropout rate.
|
|
activation (torch.nn.Module): Activation function
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
n_expert: int,
|
|
n_expert_per_token: int,
|
|
idim: int,
|
|
hidden_units: int,
|
|
dropout_rate: float,
|
|
activation: torch.nn.Module = torch.nn.ReLU(),
|
|
):
|
|
super(MoEFFNLayer, self).__init__()
|
|
self.gate = torch.nn.Linear(idim, n_expert, bias=False)
|
|
self.experts = torch.nn.ModuleList(
|
|
PositionwiseFeedForward(idim, hidden_units, dropout_rate,
|
|
activation) for _ in range(n_expert))
|
|
self.n_expert_per_token = n_expert_per_token
|
|
|
|
def forward(self, xs: torch.Tensor) -> torch.Tensor:
|
|
"""Foward function.
|
|
Args:
|
|
xs: input tensor (B, L, D)
|
|
Returns:
|
|
output tensor, (B, L, D)
|
|
|
|
"""
|
|
B, L, D = xs.size(
|
|
) # batch size, sequence length, embedding dimension (idim)
|
|
xs = xs.view(-1, D) # (B*L, D)
|
|
router = self.gate(xs) # (B*L, n_expert)
|
|
logits, indices = torch.topk(
|
|
router, self.n_expert_per_token
|
|
) # probs:(B*L, n_expert), indices: (B*L, n_expert)
|
|
weights = torch.nn.functional.softmax(
|
|
logits, dim=1,
|
|
dtype=torch.float).to(dtype=xs.dtype) # (B*L, n_expert_per_token)
|
|
output = torch.zeros_like(xs) # (B*L, D)
|
|
for i, expert in enumerate(self.experts):
|
|
mask = indices == i
|
|
batch_idx, ith_expert = torch.where(mask)
|
|
output[batch_idx] += weights[batch_idx, ith_expert, None] * expert(
|
|
xs[batch_idx])
|
|
return output.view(B, L, D)
|