GELU / SwiGLU from Scratch
import torch.nn as nn
import torch.nn.functional as F
class SwiGLU(nn.Module):
def __init__(self, d_model, d_ff):
super().__init__()
self.w_gate = nn.Linear(d_model, d_ff, bias=False)
self.w_up = nn.Linear(d_model, d_ff, bias=False)
self.w_down = nn.Linear(d_ff, d_model, bias=False)
def forward(self, x):
return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
Why it matters: The gate is a content-dependent multiplier on each hidden unit — the FFN learns which features to let through, not just how to mix them. Trap: vanilla FFN uses d_ff = 4 * d_model (2 matrices); SwiGLU has 3, so to match parameter count use d_ff ≈ (8/3) * d_model — Llama, Qwen and Mistral all do this. Copying 4x into a SwiGLU silently inflates the block by ~50%.
Full implementation in rlvr-from-scratch/model/ffn.py.