AI Engineer Exhaustive
HardPlusgpt-oss

Full Forward

gpt-oss

Hard

Embedding plus stacked GPT-OSS blocks (alternating attention pattern), final RMSNorm, lm head logits.

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

gpt-oss is a decoder-only MoE transformer. The full score function is the usual embed → (L) residual blocks → final RMSNorm → untied linear unembedding. The model card (Table 1, §2.2–2.3) fixes the sizes: residual 2880; 36 layers / 128 experts (120b) or 24 layers / 32 experts (20b); vocabulary (201{,}088) from the o200k_harmony BPE; context 131,072 on dense attention layers via YaRN. Unembedding parameters count toward the “active” total; input embeddings do not.

The public Transformer.forward is exactly that stack. It does not implement the harmony chat template, tool channels, or sampling; TokenGenerator softmaxes the last-position logits separately.

How it works

Embed. Token ids (x\in{0,\ldots,V-1}^{T}) index nn.Embedding(V, d) in BF16:

h^{(0)}_t = E_{x_t}\in\mathbb{R}^{2880}.

There is no additive learned position table. Position enters only through YaRN-RoPE inside attention. Embeddings are not tied to the unembedding (unembedding is a separate Linear with bias=False).

Stack. For (\ell=0,\ldots,L-1),

h^{(\ell+1)}=\mathrm{Block}_\ell\!\bigl(h^{(\ell)}\bigr),

where (\mathrm{Block}_\ell) is the pre-norm unit of the block note: RMSNorm → GQA (sinks, RoPE, window iff (\ell) even) → residual → RMSNorm → top-4 MoE (clamped SwiGLU) → residual. Alternation of local ((W=128)) and dense causal attention is the only (\ell)-dependent structural change. Expert count and (k=4) are constant across depth.

Because local layers cannot see past 128 tokens, long-range mixing happens on the odd layers. Those are also the layers whose KV cache grows with the 131k window. Even-layer caches stay bounded by 128.

Final norm and head. After the last block,

\mathrm{logits}_t = W_U\,\mathrm{RMSNorm}\!\bigl(h^{(L)}_t\bigr)\in\mathbb{R}^{V},

with (W_U\in\mathbb{R}^{V\times d}) and the same RMSNorm recipe as in the blocks ((\varepsilon=10^{-5}), float32 scale). No bias on (W_U). Softmax over (V) is not inside Transformer.forward; the module returns raw logits.

Sizes that matter for a single forward. Per token, 4 experts fire in every layer, so active MLP compute is (L\times 4\times) (one SwiGLU expert of width 2880). Attention is 64 query heads, 8 KV heads. The card’s active-parameter counts (5.13B / 3.61B) include this sparse MLP plus attention plus the unembedding, and exclude idle experts and the input embedding table.

Precision. The educational graph upcasts MXFP4 MoE weights to BF16 at load (weights.py) and runs BF16 activations. Checkpoints on disk keep mlp1_weight / mlp2_weight as MXFP4 blocks+scales; everything else is BF16. That dequant is a load-time detail, not a third stage of the forward equations.

The paper specifies architecture and counts. Untied embeddings, the exact final-norm placement, even-layer windowing, and “logits not probabilities” are inferred from Transformer in the official repo.

Official code

gpt_oss/torch/model.py: Transformer (embedding, block ModuleList, norm, unembedding); Transformer.forward is the four-step graph above. from_checkpoint reads config.json into ModelConfig (36/128 vs 24/32 live there, not hardcoded in forward). TokenGenerator.generate calls the model on the growing token list and reads logits[-1]. Tokenizer and harmony formatting live outside this file (TikToken o200k_harmony, separate openai-harmony package).

Watch-outs

Sources