| 223 | |
| 224 | class LlamaMLP(nn.Module): |
| 225 | def __init__(self, args): |
| 226 | super().__init__() |
| 227 | self.hidden_size = args.hidden_size |
| 228 | self.intermediate_size = args.intermediate_size |
| 229 | |
| 230 | self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) |
| 231 | self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) |
| 232 | self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) |
| 233 | |
| 234 | def forward(self, x): |
| 235 | return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |