| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class SwiGLU(nn.Module): | |
| """SwiGLU feed-forward block (Section 4.6), as used in Llama/PaLM. | |
| Standard formulation: down_proj(silu(gate_proj(x)) * up_proj(x)) | |
| The inner dim is scaled down from the naive 4x so that SwiGLU's extra | |
| gate_proj matrix doesn't blow the parameter budget relative to a plain MLP | |
| of the same nominal "4x" size -- this matches how Llama-style models size it. | |
| """ | |
| def __init__(self, hidden_dim: int, mult: float = 4.0): | |
| super().__init__() | |
| # standard correction: 4 * hidden * (2/3) keeps param count comparable | |
| # to a plain (non-gated) 4x MLP, rounded to a clean multiple of 8. | |
| inner_dim = int(hidden_dim * mult * 2 / 3) | |
| inner_dim = ((inner_dim + 7) // 8) * 8 | |
| self.gate_proj = nn.Linear(hidden_dim, inner_dim, bias=False) | |
| self.up_proj = nn.Linear(hidden_dim, inner_dim, bias=False) | |
| self.down_proj = nn.Linear(inner_dim, hidden_dim, bias=False) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) | |