AI Engineer Exhaustive
HardPlusgpt-oss

MoE Forward

gpt-oss

Hard

Sparse MoE forward: route, run clamped-SwiGLU experts, weighted-sum the outputs.

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

The feed-forward half of every gpt-oss block is a sparse MoE. After a pre-norm RMSNorm, a linear router picks 4 of (E) experts ((E=128) on 120b, (32) on 20b). Each chosen expert is a two-projection gated SwiGLU MLP with the same hidden size as the residual, (d=n_{\mathrm{inner}}=2880). The expert outputs are mixed by the post-top-k softmax weights and added back to the residual.

Footnote 1 of the model card warns that the SwiGLU is “unconventional, including clamping and a residual connection.” The reference swiglu function is that implementation: interleaved split, clamps at 7, SiLU-like gate with (\alpha=1.702), and a (+1) on the linear branch. The paper does not write the algebra; the official file does.

Expert weights are stored as MXFP4 in the checkpoint. The educational PyTorch path dequantizes them to BF16 at load time (gpt_oss/torch/weights.py) and then runs the math below in BF16. This note is the floating-point forward, not the packing format.

How it works

Route. (\tilde{x}=\mathrm{RMSNorm}(x)). Router logits (g=W_r\tilde{x}+b_r\in\mathbb{R}^{T\times E}). Top-4 logits (z_t) and indices (\mathcal{T}_t); (w_t=\mathrm{softmax}(z_t)) over those four values only.

Expert 1 (up + gate). For each selected expert (e\in\mathcal{T}_t),

u_{t,e} = W^{(1)}_e \tilde{x}_t + b^{(1)}_e \in \mathbb{R}^{2n_{\mathrm{inner}}}.

In the reference shapes, mlp1_weight is [E, 2n_inner, d] (tensor-parallel shards split the 2n_inner axis). The einsum is "beck,bk->bec" after gathering rows by expert_indices.

Clamped SwiGLU. Split (u) by interleaving, not by halves:

u_{\mathrm{glu}}=u[\ldots,0::2],\qquad u_{\mathrm{lin}}=u[\ldots,1::2].

Clamp (u_{\mathrm{glu}}\le 7) and (u_{\mathrm{lin}}\in[-7,7]). Then, with (\alpha=1.702),

\mathrm{SwiGLU}(u)=\bigl(u_{\mathrm{glu}}\cdot\sigma(\alpha\,u_{\mathrm{glu}})\bigr)\odot\bigl(u_{\mathrm{lin}}+1\bigr).

The (+1) is the “residual connection” the footnote mentions: a zero linear branch still passes the gated activation rather than wiping it. (\sigma) is the logistic sigmoid. This is not the common silu(W_g x) * (W_u x) pair of independent projections; the two streams are even/odd channels of one matrix.

Expert 2 (down).

y_{t,e}=W^{(2)}_e\,\mathrm{SwiGLU}(u_{t,e})+b^{(2)}_e \in\mathbb{R}^{d}.

mlp2_weight is [E, d, n_inner]. Under tensor parallelism the inner axis is sharded and the partial products are all_reduced before adding mlp2_bias (the bias is full-width and must not be reduced). Single-GPU inference skips the reduce.

Mix and residual.

x \leftarrow x + \sum_{i=1}^{k} w_{t,i}\, y_{t,\mathcal{T}_{t,i}}.

The official einsum is "bec,be->bc". Unselected experts contribute nothing; there is no shared dense expert. Biases are part of the checkpoint and are applied. The skip connection uses the pre-norm (x), not (\tilde{x}).

Active parameters per token are the 4 experts plus attention and the unembedding, matching the card’s 5.1B / 3.6B figures. Intermediate size equals hidden size, so each expert is narrower than a typical 4d FFN of a dense 3B model.

Official code

gpt_oss/torch/model.py: swiglu, MLPBlock.forward. Config: intermediate_size=2880, swiglu_limit=7.0, experts_per_token=4. Weight load and MXFP4 upcast: gpt_oss/torch/weights.py. The Triton backend runs the same SwiGLU against still-quantized expert weights; numerically it is intended to match this BF16 graph.

Watch-outs

Sources