| 220 | |
| 221 | |
| 222 | class MPTMLP(nn.Module): |
| 223 | def __init__(self, d_model: int, expansion_ratio: int): |
| 224 | super().__init__() |
| 225 | self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, bias=False) |
| 226 | self.act = nn.GELU(approximate="none") |
| 227 | self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, bias=False) |
| 228 | self.down_proj._is_residual = True |
| 229 | |
| 230 | def forward(self, x): |
| 231 | return self.down_proj(self.act(self.up_proj(x))) |
| 232 | |
| 233 | |
| 234 | class MPTBlock(nn.Module): |