nnterp · encyclopedia gpt_oss · nnterp.families.gpt_oss

family gpt_oss

GPT-OSS

Llama's pre-norm block with a mixture of experts on every block, sliding-window and full-attention blocks in turn, and a learned sink logit per head in the softmax, so a pattern row sums to less than one.

Attention sinkSliding windowMixture of expertsBiased q, k, vGrouped-query attention

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block; this family's blocks differ by layer_types.

embed_tokenstoken_embeddings
→
layers× 24
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model GptOssForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] GptOssDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn GptOssAttentionlayoutwhat it isread at
attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the output of attn_weights_1 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output
attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp GptOssMLPlayoutwhat it isread at
⚠︎mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
⚠︎mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output
router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

inside router's forward: the output of F_linear_0
expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2
expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1
expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1
routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output

Printout

What print(model) shows on openai/gpt-oss-20b: the native modules, standard names and every standard value with its layout.

GptOssForCausalLM(
  (lm_head): Linear(in_features=2880, out_features=201088, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(201088, 2880, padding_idx=199999)
  (layers): ModuleList(
    (0-23): 24 x GptOssDecoderLayer(
      (self_attn): GptOssAttention(
        (q_proj): Linear(in_features=2880, out_features=4096, bias=True)
        (k_proj): Linear(in_features=2880, out_features=512, bias=True)
        (v_proj): Linear(in_features=2880, out_features=512, bias=True)
        (o_proj): Linear(in_features=4096, out_features=2880, bias=True)
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): GptOssMLP(
        (router): GptOssTopKRouter()
        (experts): GptOssExperts()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): GptOssRMSNorm((2880,), eps=1e-05)
      (post_attention_layernorm): GptOssRMSNorm((2880,), eps=1e-05)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): GptOssRMSNorm((2880,), eps=1e-05)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers24
model.hidden_size2880
model.num_heads64
model.num_kv_heads8
model.head_dim64
model.qk_head_dim64
model.vocab_size201088
model.intermediate_size2880
architecturesGptOssForCausalLM
num_hidden_layers24
hidden_size2880
intermediate_size2880
num_attention_heads64
num_key_value_heads8
head_dim64
vocab_size201088
max_position_embeddings131072
sliding_window128
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
num_experts_per_tok4
num_local_experts32
attention_biasTrue

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))       # a sink column in the softmax
out = h + mlp(post_attention_layernorm(h))    # 4 experts per token

Llama's pre-norm block and Llama's names. The block returns a tensor and adds each sublayer's output to the stream as it is, so the identity is the plain sum. The MLP module returns (hidden_states, router_scores); mlp_output is the first element, and mlp.output the tuple.

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

torch.testing.assert_close(x + attn + mlp, out)

Every head has a sink, so pattern rows sum to less than one

Each attention module holds sinks, one learned logit per query head ([num_heads], 64 on 20b and 120b), on sliding and full blocks alike. The sink's share of the softmax mixes no value: attention_probabilities has seq key columns whose rows sum to less than one, and 1 - row sum is the sink's share.

attention_scores is read at the masked scores just before the sink column joins, so it has the pattern's shape and softmax(attention_scores) is not the pattern. Appending the column rebuilds it:

attn = model.layers[0].self_attn
with model.trace(prompt):
    scores = attn.attention_scores.save()       # masked, before the sink joins
    pattern = attn.attention_probabilities.save()

assert (pattern.float().sum(-1) < 1).all()      # the sink keeps the rest

sinks = attn._module.sinks.to(scores.dtype)
column = sinks.view(1, -1, 1, 1).expand(scores.shape[0], -1, scores.shape[2], 1)
combined = torch.cat([scores, column], dim=-1)
combined = combined - combined.max(dim=-1, keepdim=True).values
full = combined.softmax(-1)
torch.testing.assert_close(full[..., :-1], pattern)
sink_share = full[..., -1]                      # [batch, num_heads, query]

The softmax runs in the scores' dtype, not in float32. A head's output is the values weighted by these rows, so a row's sum is the share of that query's attention that reaches the tokens; the sink takes the rest. An edit to attention_scores is renormalized against the sink. An edit to attention_probabilities sets the weights directly, and its rows need not sum to anything.

Sliding-window and full blocks alternate

layer_types alternates sliding_attention and full_attention, starting with a sliding block: the even blocks attend over the latest sliding_window (128) positions and the odd blocks over the whole prefix. On a prompt longer than the window, a sliding block's pattern is zero outside the window, so each row has at most 128 nonzero key columns:

with model.trace(long_prompt):                        # more than 128 tokens
    pattern = model.layers[0].self_attn.attention_probabilities.save()

assert ((pattern[0, :, -1] > 0).sum(-1) <= 128).all()   # block 0 slides

Loading: the default load is eager

transformers offers no SDPA path for GPT-OSS (_supports_sdpa = False), so a load without attn_implementation runs the eager forward and every attention value is available. The 20b and 120b configs carry an MXFP4 quantization_config for the experts only; attention, router, embeddings and lm_head are not quantized.

Attention: grouped heads, biased projections, YaRN rotary

64 query heads of width 64 share 8 key/value heads, so attention_keys and attention_values are [batch, 8, seq, 64] and query head h reads key/value head h // 8: an edit to one key/value head reaches 8 query heads. q_proj, k_proj, v_proj and o_proj all add a bias (attention_bias). The rotary embedding is YaRN (rope_parameters: factor 32 over an original 4,096 positions, rope_theta 150,000), and attention_queries and attention_keys are read after it. YaRN multiplies the rotary cosines and sines by 0.1 * ln(32) + 1 = 1.347, so both are 1.347 times the rotated projections.

The router takes a softmax over its top 4 logits

The router is a biased linear map to num_experts logits (32 on 20b, 128 on 120b); it keeps the 4 largest and takes a softmax over those 4 alone, so expert_weights sum to one at every token and router_logits outside the top 4 do not enter the weights. There is no shared expert: routed_output is mlp_output, and expert_outputs sums to it.

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()             # [batch, seq, num_experts]
    w = moe.expert_weights.save()                 # [batch, seq, top_k]
    idx = moe.expert_indices.save()
    each = moe.expert_outputs.save()              # [batch, seq, top_k, hidden_size]
    routed = moe.routed_output.save()

top = logits.topk(moe.top_k, dim=-1)
torch.testing.assert_close(top.values.softmax(-1).to(w.dtype), w)
assert torch.equal(top.indices, idx)
torch.testing.assert_close(each.sum(2), routed)

The experts are clamped SwiGLUs with biases

The experts hold their weights as stacked parameters, experts.gate_up_proj ([num_experts, hidden_size, 2 × intermediate]) and experts.down_proj, each with a bias. The gate and up halves are interleaved, gate on the even columns and up on the odd ones. The gate is clamped above at 7 and the up half to [−7, 7] (swiglu_limit, 7.0 in every released config), and the activation is (up + 1) · gate · sigmoid(1.702 · gate) (swiglu_alpha, whose 1.702 default no released config overrides). An expert's output includes its down_proj bias, so every slot adds a weighted bias, and expert_outputs is each slot's output times its weight.

Ablating an expert touches only the tokens routed to it

Zeroing expert e's weight where expert_indices == e removes that slot's term and leaves the token's other 3 weights as they were (they are not renormalized). mlp_output changes on exactly the tokens that chose e, and an expert no token chose has no effect:

moe, e = model.layers[1].mlp, 3
with model.trace(prompt):
    idx = moe.expert_indices.save()
    clean = moe.mlp_output.save()
with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == e, 0)
    ablated = moe.mlp_output.save()

changed = (ablated != clean).any(-1)                       # [batch, seq]
assert torch.equal(changed, (idx == e).any(-1))            # only the tokens that chose e

The readout and the tokenizer

logits is lm_head.output, with no cap or scale, and project_on_vocab on the last block's layer_output equals logits. The final norm is a plain RMSNorm whose gain is norm.weight. The tokenizer prepends nothing, so model.input_ids is the prompt's tokens alone; <|startoftext|> is id 199998 and the end-of-sequence token is <|return|>, id 200002.

From the family module

04 · nnterp/families/gpt_oss.py

GPT-OSS (GptOssForCausalLM).

Llama's tree with a mixture-of-experts MLP that returns (hidden_states, router_scores), sliding-window layers, and an attention sink: each head carries a learned logit that joins the softmax as one extra key column and is dropped afterwards. So the pattern's rows sum to less than one, and the softmax's input is one key wider than the pattern and shifted by its row max; the standard attention_scores is therefore read one step earlier, at the masked scores just before the sink joins them (a binding that exists only when an attention mask is passed, as it is on every prompt).