| from dataclasses import dataclass |
| from typing import Union |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from utils import bias_gelu_impl |
| from mamba_config import MambaConfig |
|
|
| class MLP(nn.Module): |
| def __init__( |
| self, config: MambaConfig, is_expert: bool = False, layer_idx=None |
| ): |
| super().__init__() |
| |
| self.config: MambaConfig = config |
| self.layer = layer_idx |
| ffn_hidden_size_1 = self.config.ffn_hidden_size |
| ffn_hidden_size_2 = self.config.ffn_hidden_size |
| |
| |
| if self.config.gated_linear_unit: |
| ffn_hidden_size_1 *= 2 |
| |
| self.linear_fc1 = nn.Linear(self.config.hidden_size, ffn_hidden_size_1, bias = self.config.add_bias_linear, device = self.config.device) |
| self.linear_fc1.is_expert = is_expert |
|
|
| if self.config.gated_linear_unit: |
|
|
| def glu(x): |
| x = torch.chunk(x, 2, dim=-1) |
| return self.config.activation_func(x[0]) * x[1] |
|
|
| self.activation_func = glu |
| else: |
| self.activation_func = self.config.activation_func |
|
|
| self.linear_fc2 = nn.Linear(ffn_hidden_size_2, self.config.hidden_size, bias = self.config.add_bias_linear, device = self.config.device) |
|
|
| def forward(self, hidden_states, inference_params=None): |
| intermediate = self.linear_fc1(hidden_states) |
| intermediate = self.activation_func(intermediate) |
| output = self.linear_fc2(intermediate) |
| return output |
