GPT-OSS-120B

OpenAI's open-weight MoE: alternating sliding/full attention with learned sinks, MXFP4 experts.

params=117Bactive=5.1Blayers=36ctx=128Kvocab=201,088

source: config.json ↗ · tech report ↗ snapshot 2026-08-24

attentionFFN / MoEnorm / residualembedding / headauxiliarydata flowresidualauxiliary

Reading this diagram

Data flows top to bottom. Solid edges carry activations (tensor shapes on the labels), dashed gray edges are residual paths, and dotted edges attach auxiliary modules. The layer stack draws one representative layer — the bracketed badge carries the ×36 multiplicity, and labeled variants show which layers differ.

Hover any box for a summary; click it for the details and a PyTorch reference implementation. Selections are shareable — the URL keeps them.

highlights: alternating sliding-window (128) / full attention · learned attention sinks + attention bias · 128 experts, top-4, clipped SwiGLU · MXFP4 expert weights

embedding / head

Token embedding

Maps each of the 201,088 vocabulary ids to a learned 2880-dim vector — the model's only lookup table.

The embedding is a learned matrix W_E ∈ ℝ^{201088 × 2880}. A token id indexes a row, so the "computation" is a gather, not a matmul. The output is the initial residual stream every layer will read from and write to.

Subtleties worth knowing:

No positional information is added here in modern decoders — position enters inside attention via RoPE, not as an additive embedding (the original transformer added sinusoids at this point).
Scale: at d=2880 and V=201,088, this table alone is 0.58B parameters.
Weight tying is OFF in this model: the LM head has its own independent matrix, which tends to help at large scale at the cost of a second V×d table.

reference implementation

import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size=201088, d_model=2880):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)

    def forward(self, input_ids):          # [B, T] int64
        return self.embed(input_ids)       # [B, T, 2880]

structure

Transformer layer

One pre-norm residual block: normalize → attention → add, then normalize → FFN → add.

The repeating unit of the whole model. Each layer computes two deltas and adds them to the d=2880 residual stream:

h = x + Attention(Norm(x)) — tokens exchange information (the only place sequence positions interact);
y = h + FFN(Norm(h)) — each token is transformed independently.

Subtleties:

Pre-norm (norm on the branch, not the trunk) is what makes 36-layer stacks trainable without fragile warmup schedules — the identity path is never scaled or normalized.
Attention mixes across time, the FFN mixes across features; a transformer is just these two mixers alternated 36 times.
The diagram draws one representative layer — the ×36 badge carries the multiplicity. Where variants exist, the layer keeps this exact skeleton and only swaps the slot's contents:
- attention: layers 0–34 (alternating) — sliding window
- attention: layers 1–35 (alternating) — full attention
- FFN: all layers

reference implementation

class DecoderLayer(nn.Module):
    def __init__(self, d_model=2880):
        super().__init__()
        self.norm1, self.norm2 = RMSNorm(d_model), RMSNorm(d_model)
        self.attn = Attention(...)
        self.ffn  = FFN(...)          # dense SwiGLU or MoE, per layer index

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

norm / residual

RMSNorm

Rescales each token vector to unit RMS over its 2880 features (eps=0.00001), then applies a learned per-feature gain.

y = x / sqrt(mean(x²) + 0.00001) · g with a learned gain g ∈ ℝ^{2880}.

Compared to LayerNorm, RMSNorm drops mean subtraction and the bias term — it only controls the vector's scale, not its center. Empirically this loses nothing in quality and is faster (one reduction instead of two), which is why every model in this tool uses it.

Subtleties:

Pre-norm placement: the norm sits before attention/FFN, not after. The residual stream itself is never normalized in place — that's what keeps very deep stacks trainable (the identity path stays untouched).
The reduction is usually done in fp32 even in bf16 models; eps=0.00001 guards the near-zero-norm case.
The final RMSNorm before the LM head is the only place the residual stream's growing magnitude gets tamed before scoring tokens.

reference implementation

class RMSNorm(nn.Module):
    def __init__(self, d_model=2880, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):                     # [B, T, 2880]
        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * rms).type_as(x) * self.g

attention

Sliding-window attention

GQA restricted to the last 128 tokens — O(T·128) instead of O(T²), used on alternating layers.

Identical to GQA except the causal mask is banded: each query may only attend to the previous 128 positions. Cost and KV-cache reads become linear in window size rather than context length.

Why alternating works: a sliding layer propagates information 128 tokens; a stack of them extends reach layer by layer (receptive field grows ~linearly with depth), while the interleaved full-attention layers provide direct long-range hops. GPT-OSS alternates sliding/full 1:1 — half the layers cheap, half global.

In this model:

window = 128 tokens, 64/8 Q/KV heads, d_h=64
learned sink logits matter more here: with a hard window, heads often have nothing relevant to attend to, and the sink absorbs that probability mass
the KV cache for these layers can be a ring buffer of 128 entries — it stops growing with context, a large serving win at 128K.

reference implementation

def sliding_mask(T, window=128, device=None):
    i = torch.arange(T, device=device)
    # causal AND within the last 128 positions
    return (i[:, None] >= i[None, :]) & (i[:, None] - i[None, :] < window)

# inside forward(): same GQA math, banded mask instead of is_causal=True
out = F.scaled_dot_product_attention(q, k, v, attn_mask=sliding_mask(T))

attention

Grouped-Query Attention

64 query heads share 8 KV heads (d_h=64) — full attention quality with a 8× smaller KV cache.

Standard causal self-attention with grouped queries: the 64 query heads are split into 8 groups, and each group shares one K/V head. Per token the KV cache stores 8 × 64 keys and values instead of 64 × 64 — at inference, KV-cache size (not FLOPs) is what limits batch size and context, so this is nearly free quality-wise and hugely valuable.

The spectrum: MHA (kvHeads = heads, original), GQA (1 < kvHeads < heads), MQA (kvHeads = 1, maximum sharing, some quality cost).

In this model:

64 query heads, 8 KV heads, d_h=64
RoPE base θ=150000 with YaRN ×32 context extension
no Q/K normalization
projection biases enabled (uncommon in modern decoders)
learned sink logits: each head has a trainable logit competing in the softmax, letting heads attend to 'nothing' instead of being forced to distribute weight over real tokens — replaces the first-token 'attention sink' hack

reference implementation

class GQA(nn.Module):
    def __init__(self, d_model=2880, n_heads=64,
                 n_kv_heads=8, head_dim=64):
        super().__init__()
        self.q = nn.Linear(d_model, n_heads * head_dim, bias=True)
        self.k = nn.Linear(d_model, n_kv_heads * head_dim, bias=True)
        self.v = nn.Linear(d_model, n_kv_heads * head_dim, bias=True)
        self.o = nn.Linear(n_heads * head_dim, d_model, bias=False)
        self.groups = n_heads // n_kv_heads
        self.sinks = nn.Parameter(torch.zeros(n_heads))  # learned sink logits

    def forward(self, x, pos):
        B, T, _ = x.shape
        q = split_heads(self.q(x), 64)        # [B, 64, T, 64]
        k = split_heads(self.k(x), 8)      # [B, 8, T, 64]
        v = split_heads(self.v(x), 8)
        q, k = apply_rope(q, pos), apply_rope(k, pos)
        k, v = repeat_kv(k, self.groups), repeat_kv(v, self.groups)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o(merge_heads(out))               # [B, T, 2880]

norm / residual

Residual add

Adds the attention sublayer's output back onto the untouched input stream: y = x + f(norm(x)).

The residual connection is the transformer's gradient highway: every layer computes a delta f(norm(x)) and adds it to the stream, so the identity path from embedding to LM head is completely unobstructed.

Subtleties:

In pre-norm blocks (all models here), the add receives the raw stream x — normalization happened on the branch, not the trunk. This is why 60–90-layer models train without warmup tricks that post-norm (original transformer) needed.
The stream's magnitude grows roughly with depth since every layer adds to it; the final RMSNorm absorbs that before the head.
Residual adds are where activation dtype matters: the add is typically done in the trunk dtype (bf16) while branches may compute internally in other precisions.

reference implementation

# inside a decoder layer (pre-norm)
h = x + attention(norm1(x))     # ← this add
y = h + ffn(norm2(h))           # ffn sublayer

norm / residual

RMSNorm

Rescales each token vector to unit RMS over its 2880 features (eps=0.00001), then applies a learned per-feature gain.

y = x / sqrt(mean(x²) + 0.00001) · g with a learned gain g ∈ ℝ^{2880}.

Subtleties:

Pre-norm placement: the norm sits before attention/FFN, not after. The residual stream itself is never normalized in place — that's what keeps very deep stacks trainable (the identity path stays untouched).
The reduction is usually done in fp32 even in bf16 models; eps=0.00001 guards the near-zero-norm case.
The final RMSNorm before the LM head is the only place the residual stream's growing magnitude gets tamed before scoring tokens.

reference implementation

class RMSNorm(nn.Module):
    def __init__(self, d_model=2880, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):                     # [B, T, 2880]
        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * rms).type_as(x) * self.g

structure

Mixture-of-Experts FFN

The FFN split into 128 small experts; each token runs through only 4 — most parameters idle per token.

MoE decouples parameter count from per-token compute: the dense FFN is replaced by 128 narrow SwiGLU experts (d=2880 → 2880), and a router sends each token to its top-4. Total capacity scales with expert count while FLOPs stay near a small dense model's — the mechanism behind "128-expert, 4-active" models.

Design axes visible in this tool's models:

Granularity: many small experts (d_expert=2880) rather than few big ones — finer routing granularity consistently wins at fixed compute.
Shared expert: absent — all capacity is routed (Qwen3/GPT-OSS style).
Load balancing: routers collapse onto favorite experts without countermeasures — see the router box for this model's approach.
Systems reality: experts shard across GPUs (expert parallelism); tokens are all-to-all'd to their experts and back, which is why balanced routing matters for wall-clock, not just quality.

reference implementation

class MoEFFN(nn.Module):
    def __init__(self, d_model=2880, n_experts=128, top_k=4):
        super().__init__()
        self.router = Router(d_model, n_experts, top_k)
        self.experts = nn.ModuleList(
            SwiGLU(d_model, 2880) for _ in range(n_experts))

    def forward(self, x):
        weights, idx = self.router(x)            # [N, 4] each
        out = torch.zeros_like(x)
        for slot in range(4):            # gather/scatter in real kernels
            e = idx[:, slot]
            out += weights[:, slot, None] * dispatch(self.experts, e, x)
        return out

FFN / MoE

MoE router

Scores all 128 experts per token (softmax) and dispatches to the top-4.

A linear layer ℝ^{d} → ℝ^{128} producing per-expert affinities, turned into a top-4 choice and combine weights. Small module, most consequential knobs in the architecture:

Scoring: softmax. Classic — affinities compete (weights sum to 1 across experts).
Top-k selection: standard top-k over scores.
No group restriction — any token may use any expert.
Renormalization ON: the selected top-k weights are rescaled to sum to 1.
The router usually runs in fp32 even in quantized models — tiny cost, and routing decisions are too consequential for low precision.

reference implementation

class Router(nn.Module):
    def __init__(self, d_model, n_experts=128, top_k=4):
        super().__init__()
        self.gate = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):                               # [N, d]
        logits = self.gate(x.float())
        scores = logits.softmax(dim=-1)
        w, sel = scores.topk(4, dim=-1)
        w = w / w.sum(dim=-1, keepdim=True)             # renormalize
        return w, sel

FFN / MoE

Routed experts

128 independent SwiGLU MLPs (d=2880 → 2880); each token executes only its 4.

Each expert is a standard SwiGLU MLP, just narrow: d_ff=2880 vs the thousands a dense FFN would use. A token's routed output is the weighted sum of its 4 experts' outputs.

Subtleties:

Specialization is emergent: experts aren't assigned topics; clustering arises from routing pressure (and is famously hard to interpret cleanly).
Batched execution: kernels group all tokens per expert and run one GEMM per expert (grouped GEMM) — the loop in reference code is purely pedagogical.
MXFP4 weights: expert weights dominate parameter count, so this model stores them in MXFP4; the SwiGLU clip at ±7 keeps activations inside the format's range — attention/router/embeddings stay higher precision.
Capacity note: with 128 experts × 4 active, the ratio of stored to exercised FFN parameters per token is 32:1.

reference implementation

self.experts = nn.ModuleList(
    SwiGLU(d_model=2880, d_ff=2880)   # narrow expert
    for _ in range(128)
)

# grouped-GEMM style dispatch (conceptual):
for e in selected_experts.unique():
    rows = (selected_experts == e).nonzero()
    out[rows] += gate_w[rows] * self.experts[e](x[rows])

FFN / MoE

Expert combine

Weighted sum of expert outputs.

The reduction that reassembles a token's result: y = Σ wᵢ·eᵢ(x).

Subtleties:

The weights wᵢ come from the router's scores renormalized over the selected top-k (they sum to 1, so the combine is a convex average).
In distributed serving this sum happens after the all-to-all returns expert outputs to their home devices — it's the merge point of the scatter/gather.

reference implementation

# w: [N, k] gate weights, e_out: [N, k, d] expert outputs
w = w / w.sum(-1, keepdim=True)      # convex combination
y = (w.unsqueeze(-1) * e_out).sum(1)  # [N, d]

norm / residual

Residual add

Adds the ffn sublayer's output back onto the untouched input stream: y = x + f(norm(x)).

Subtleties:

In pre-norm blocks (all models here), the add receives the raw stream x — normalization happened on the branch, not the trunk. This is why 60–90-layer models train without warmup tricks that post-norm (original transformer) needed.
The stream's magnitude grows roughly with depth since every layer adds to it; the final RMSNorm absorbs that before the head.
Residual adds are where activation dtype matters: the add is typically done in the trunk dtype (bf16) while branches may compute internally in other precisions.

reference implementation

# inside a decoder layer (pre-norm)
h = x + attention(norm1(x))     # attention sublayer
y = h + ffn(norm2(h))           # ← this add

norm / residual

RMSNorm

Rescales each token vector to unit RMS over its 2880 features (eps=0.00001), then applies a learned per-feature gain.

y = x / sqrt(mean(x²) + 0.00001) · g with a learned gain g ∈ ℝ^{2880}.

Subtleties:

Pre-norm placement: the norm sits before attention/FFN, not after. The residual stream itself is never normalized in place — that's what keeps very deep stacks trainable (the identity path stays untouched).
The reduction is usually done in fp32 even in bf16 models; eps=0.00001 guards the near-zero-norm case.
The final RMSNorm before the LM head is the only place the residual stream's growing magnitude gets tamed before scoring tokens.

reference implementation

class RMSNorm(nn.Module):
    def __init__(self, d_model=2880, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):                     # [B, T, 2880]
        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * rms).type_as(x) * self.g

embedding / head

LM head

Projects the final 2880-dim hidden state to 201,088 logits — one score per vocabulary entry.

A single linear layer ℝ^{2880} → ℝ^{201088} applied at every position. During training all T positions produce logits (next-token loss everywhere); during generation only the last position's logits are needed.

Subtleties:

Untied from the embedding in this model — separate matrices for reading tokens in and scoring tokens out, standard at this scale.
The softmax over 201,088 entries is usually computed in fp32 even when the trunk runs bf16 — logit precision matters for sampling quality.
This matmul is a significant fraction of per-token FLOPs at small batch sizes, which is one motivation for speculative decoding (verify many drafted tokens in one pass).

reference implementation

class LMHead(nn.Module):
    def __init__(self, d_model=2880, vocab_size=201088, tie_to=None):
        super().__init__()
        if tie_to is not None:                     # weight tying
            self.weight = tie_to.embed.weight      # [V, d]
        else:
            self.weight = nn.Parameter(torch.empty(vocab_size, d_model))

    def forward(self, h):                          # [B, T, 2880]
        return h @ self.weight.T                   # [B, T, 201088] logits