File size: 1,159 Bytes
ab56428 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | 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))
|