AI Engineer Exhaustive
HardPlusgpt-oss

Transformer Block

gpt-oss

Hard

Pre-norm GPT-OSS block: RMSNorm, GQA attention with sinks plus optional sliding window, residual, RMSNorm, sparse MoE, residual.

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

A gpt-oss layer is a Pre-LN residual block in the GPT-2 sense. Section 2.2: residual width 2880; RMSNorm on the activations before each attention and each MoE sublayer; Grouped-Query attention with sinks and alternating sliding-window / dense masks; then a sparse MoE instead of a dense MLP. The official TransformerBlock is just those two submodules in series. Each submodule owns its own RMSNorm and its own residual add.

There is no extra “sandwich” norm after the residual, no QK-norm, and no shared expert beside the routed MoE. Layer index matters: it selects the attention mask (window on even layers, full causal on odd layers) but does not change widths, expert counts, or RoPE parameters.

How it works

Let (x\in\mathbb{R}^{T\times 2880}) enter layer (\ell). RMSNorm in the reference is

\mathrm{RMSNorm}(x)= \frac{x}{\sqrt{\mathrm{mean}(x^{2})+\varepsilon}} \odot \gamma,

with (\varepsilon=10^{-5}) and (\gamma\in\mathbb{R}^{2880}) stored in float32, then cast back to the activation dtype (BF16). Mean is over the last axis only.

Attention sublayer (AttentionBlock):

x \leftarrow x + W_o\,\mathrm{SDPA}\!\bigl(\mathrm{RoPE}(W_{qkv}\,\mathrm{RMSNorm}(x))\bigr).

Concretely: fused QKV to 64 query heads and 8 KV heads of width 64; YaRN-RoPE on (Q,K); sdpa with scale (1/\sqrt{64}), optional (W=128) band when (\ell) is even, and a learned sink logit per query head; output projection back to 2880; residual add. Details sit in the GQA, sinks, window, and YaRN notes.

MoE sublayer (MLPBlock):

x \leftarrow x + \sum_{e\in\mathrm{TopK}} w_e\,\mathrm{Expert}_e\!\bigl(\mathrm{RMSNorm}(x)\bigr).

The second RMSNorm is a new parameter tensor, not a reuse of the attention norm. Routing is top-4 with softmax after the selection. Each expert is clamped SwiGLU (limit 7, (\alpha=1.702), (+1) on the linear branch) with inner width 2880. Mixture weights (w_e) sum to one. The skip is the post-attention residual, not the original block input.

Chaining the two sublayers is the whole TransformerBlock.forward: x = self.attn(x); x = self.mlp(x). Writing a single residual around both, or a post-norm after the pair, would be a different architecture than the one that matches the released weights.

Both model sizes share this block. They differ in how many times it is stacked (36 vs 24) and in (E) (128 vs 32), not in (d), head layout, window, or (k).

What the card specifies: Pre-LN, RMSNorm, GQA 64/8/64, sinks, alternating band 128, YaRN on dense context, MoE top-4 with selected-only softmax, unconventional SwiGLU. What is inferred from official code: two distinct norms, residuals inside each submodule, even-layer window, exact SwiGLU, RMSNorm (\varepsilon), and BF16 linears with default biases on QKV, output, and the router.

Official code

gpt_oss/torch/model.py: TransformerBlock, AttentionBlock, MLPBlock, RMSNorm. layer_idx is passed only into AttentionBlock to set sliding_window. There is no separate norm3 or final block-level scale.

Watch-outs

Sources