LLM Engineering

With d = 4,096, a classic feed-forward layer with a 4d hidden size has 134,217,728 parameters. SwiGLU uses three matrices instead of two, so it shrinks the hidden size to two thirds of 4d to keep the count equal: 134,221,824. LLaMA rounds that hidden size to 11,008 for hardware efficiency, giving 135,266,304.

• Gated variants win in experiments, but the explanation is thin; Shazeer (2020) attributes the success to “divine benevolence”.
• Most of a dense model’s parameters sit here, and so does much of its factual recall (Geva et al. 2021) — which is why editing a fact can have side-effects.
• Feed-forward size and depth trade off; parameter count alone does not say which shape is better.

d = 4_096
classic = 2 * d * (4 * d)                       # W1 and W2 with a 4d hidden layer
swiglu_ideal = 3 * d * round(2 / 3 * 4 * d)     # three matrices at 2/3 of 4d keep the count equal
swiglu_llama = 3 * d * 11_008                   # LLaMA-7B rounds the hidden size to 11,008

Assumes import numpy as np. Configurations named after real models use their published shapes; other numbers are made up or simulated, as the comments say.

import torch, torch.nn as nn, torch.nn.functional as F

class SwiGLU_FFN(nn.Module):
    def __init__(self, d_model, expansion=8/3):
        super().__init__()
        hidden = int(d_model * expansion)
        self.w1 = nn.Linear(d_model, hidden, bias=False)
        self.w_gate = nn.Linear(d_model, hidden, bias=False)
        self.w2 = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        return self.w2(self.w1(x) * F.silu(self.w_gate(x)))Copy