Hand-Writing a Vision-Language Model, Part 2: SwiGLU, DecoderLayer, and a RAM Lesson

Picking Up From Part 1

Part 1 covered RMSNorm, RotaryEmbedding, and the GQA Attention block. The attention stack was verified: forward pass producing the right shape, parameter count matching the real weight shapes off the safetensors file.

This session finishes the language model side of the project: the feed-forward block, the full decoder layer wiring everything together, and the complete Model class. Then things got interesting when it came time to actually instantiate it.

The SwiGLU MLP Block

class MLP(nn.Module):
    def __init__(self, hidden_size, intermediate_size, bias=False):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=bias)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=bias)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=bias)

    def forward(self, x):
        gate = torch.nn.functional.silu(self.gate_proj(x))
        return self.down_proj(gate * self.up_proj(x))

Weight shapes confirmed before writing: gate_proj and up_proj are both [9728, 2560] (hidden to intermediate), down_proj is [2560, 9728] (intermediate back to hidden). Three matrices, not the original Transformer FFN's two. That is the whole point of SwiGLU versus a plain feed-forward network.

The Core Idea

A plain FFN is down_proj(activation(fc1(x))). One projection, activation applied to its entire output.

SwiGLU computes two independent projections of the same input in parallel: gate_proj(x) and up_proj(x). Only one of them, gate_proj, ever touches an activation function. up_proj(x) stays completely linear.

The actual work happens at the multiply: silu(gate_proj(x)) * up_proj(x). Think of up_proj(x) as content and the activated gate as a learned, per-channel, per-token valve deciding how much of that content survives. The elementwise product is the only place gate and content interact. down_proj then compresses the gated result back to hidden_size.

Why SiLU and Not ReLU

ReLU(x) = max(0, x) has a well-known failure mode: neurons stuck receiving negative inputs get zero gradient permanently (dying ReLU). SiLU(x) = x * sigmoid(x) is smooth everywhere, dips slightly below zero near small negative inputs instead of hard-clamping, so nothing goes permanently dead.

One important nuance: SiLU is not bounded to [0, 1] the way a bare sigmoid would be. For large positive inputs it behaves like a linear function. So "gate" here means learned multiplier, not learned percentage.

The GLU Family

ReGLU, GeGLU, and SwiGLU are all the same two-branch gated structure with a different activation in the gate slot. ReLU for ReGLU, GELU for GeGLU, SiLU for SwiGLU. Not competing architectures — the same skeleton with one slot swapped. Shazeer's 2020 paper "GLU Variants Improve Transformer" found SwiGLU and GeGLU consistently outperform the others, which is why Qwen, LLaMA, Mistral, and most modern open LLMs settled on SwiGLU.

What I Got Wrong First, and What Clicked

Worth being honest about the path to understanding this, because it took a few wrong turns.

First wrong model: sequential, not parallel. My first mental model was "gate then up then down" as a pipeline. Gate upscales, then "up" picks out important stuff, then down compresses. That is wrong. gate_proj and up_proj are two independent projections computed in parallel from the same input, not one feeding the other. They are siblings, not a relay.

Second wrong turn: "matrix vs single value" as the key difference. I reasoned that SwiGLU's trick was passing activations through whole matrices instead of singular values. That is not actually new — a plain ReLU FFN already applies ReLU elementwise across an entire projected vector, so that was never the distinguishing feature. The real difference is the second parallel projection and the multiplicative interaction between the two. Something a plain single-projection FFN has no equivalent of at all.

Third wrong turn: assuming the gate is a clean 0 to 1 percentage. Reasonable assumption given "gate" implies something like a probability or proportion, but SiLU is not bounded that way. It is unbounded above and dips slightly negative, so it is better described as a learned multiplier than a learned percentage.

What actually clicked: realising that only the gate branch ever touches the activation function. up_proj(x) stays completely raw and linear the entire time, and the multiply is the only point where the two branches interact. That is the one-sentence version of the whole mechanism. Gate decides how much of up_proj's content survives, channel by channel. Contrasting it directly against a plain FFN — where every value in the single projection passes through the activation with no second untouched branch at all — is what made it stick.

The DecoderLayer

class DecoderLayer(nn.Module):
    def __init__(self, hidden_size, num_heads, num_kv_heads, head_dim,
                 intermediate_size, rms_norm_eps, bias=False):
        super().__init__()
        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.self_attn = Attention(hidden_size, num_heads, num_kv_heads,
                                   head_dim, rms_norm_eps, bias=bias)
        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.mlp = MLP(hidden_size, intermediate_size, bias=bias)

    def forward(self, x, cos, sin, attn_mask=None):
        residual = x
        x = self.self_attn(self.input_layernorm(x), cos, sin, attn_mask=attn_mask)
        x = residual + x

        residual = x
        x = self.mlp(self.post_attention_layernorm(x))
        x = residual + x
        return x

The pattern is x = residual + sublayer(norm(x)), applied twice: once around attention, once around the MLP.

This is pre-norm: normalise before the sublayer runs, not after. The residual path carries the raw un-normalised signal forward untouched. The norm-attention-MLP path only ever contributes an addition on top of it. This is a significant reason why a 36-layer network is trainable at all. Gradients have a clean, unimpeded path straight back through the residual connections regardless of depth.

One detail worth noting: input_layernorm and post_attention_layernorm weights are both [2560] (full hidden_size). That is a different granularity from q_norm and k_norm in Part 1, which were [128] (per-head). The block-level norms see the full token representation; the head norms see only each head's 128-dimensional slice.

The Full Model Class

class Model(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers, num_heads, num_kv_heads,
                 head_dim, intermediate_size, rms_norm_eps, rope_theta, bias=False):
        super().__init__()
        self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList([
            DecoderLayer(hidden_size, num_heads, num_kv_heads, head_dim,
                         intermediate_size, rms_norm_eps, bias=bias)
            for _ in range(num_layers)
        ])
        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.rotary_emb = RotaryEmbedding(head_dim, theta=rope_theta)

    @classmethod
    def from_config(cls, config: dict) -> "Model":
        return cls(
            vocab_size=config["vocab_size"],
            hidden_size=config["hidden_size"],
            num_layers=config["num_hidden_layers"],
            num_heads=config["num_attention_heads"],
            num_kv_heads=config["num_key_value_heads"],
            head_dim=config["head_dim"],
            intermediate_size=config["intermediate_size"],
            rms_norm_eps=config["rms_norm_eps"],
            rope_theta=config["rope_theta"],
            bias=config.get("attention_bias", False),
        )

    def forward(self, input_ids):
        batch, seq_len = input_ids.shape
        x = self.embed_tokens(input_ids)

        position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch, -1)
        cos, sin = self.rotary_emb(x, position_ids)

        causal_mask = torch.full((seq_len, seq_len), float("-inf"), device=x.device).triu(diagonal=1)
        causal_mask = causal_mask[None, None, :, :]

        for layer in self.layers:
            x = layer(x, cos, sin, attn_mask=causal_mask)

        x = self.norm(x)
        return F.linear(x, self.embed_tokens.weight)

    def __repr__(self):
        n_params = sum(p.numel() for p in self.parameters())
        return f"Qwen3Model(layers={len(self.layers)}, hidden_size={self.embed_tokens.embedding_dim}, params={n_params:,})"

Two things worth highlighting beyond the architecture itself.

Tied embeddings made concrete. config.json has tie_word_embeddings: true and no lm_head.weight tensor exists in the checkpoint at all, confirmed back in Part 1 by scanning the safetensors index. The LM head is F.linear(x, self.embed_tokens.weight), literally reusing the token embedding table transposed as the output projection back to vocabulary logits. One matrix doing two jobs.

from_config as a classmethod factory. Rather than a constructor that expects ten positional arguments at the call site, this unpacks a raw config dict once. cls(...) rather than Model(...) is deliberate. It keeps working correctly even if Model is ever subclassed, because cls resolves to whatever class the method was actually called on.

The RAM Crash

With the full Model class written, the obvious next step was a smoke test: instantiate Model.from_config() with the real Qwen3-4B config and verify the parameter count.

Instant crash.

nn.Linear and nn.Embedding allocate fp32 by default. Qwen3-4B has approximately 4 billion parameters. At fp32, that is roughly 16GB just for parameter storage, before a forward pass, before any quantisation step. On an 8GB RAM machine, that is not a close call.

This exposed something easy to miss: quantisation via bitsandbytes (the plan for fitting the 4B model onto the GTX 1080's 8GB VRAM) only helps when weights are being loaded for actual inference or training on GPU. It does nothing for a naive fresh instantiation sitting in system RAM. The full fp32 allocation has to exist first regardless of what you plan to do with it afterward.

What Is Next

The hand-written stack is architecturally complete. The next session downloads the Qwen3-0.6B config, weights, and tokenizer, inspects the tensor names the same way Part 1 inspected the 4B checkpoint, then loads those real weights into the hand-written Model and verifies logits against the official AutoModelForCausalLM as ground truth.

Once that verification passes, the language model side is done. Attention turns to the ViT encoder and trainable projector.