Qwen3 Architecture Layer-by-Layer — Reading the HuggingFace Implementation

All permalinks in this post are pinned to huggingface/transformers tag v4.57.1. Line numbers are verified against that revision.

Concrete shape numbers are taken from Qwen/Qwen3-0.6B (config.json). Config values: hidden_size=1024, num_hidden_layers=28, num_attention_heads=16, num_key_value_heads=8, head_dim=128, intermediate_size=3072, vocab_size=151936, max_position_embeddings=40960, rms_norm_eps=1e-06.

Why a layer-by-layer pass

LLM model code in transformers looks small until you actually read it. A Qwen3 forward pass touches an embedding lookup, 28 transformer blocks, a final norm, an LM head, and a loss — roughly 250 lines of Python in a single file (modeling_qwen3.py). Once you’ve internalized that file, the differences across modern open LLMs (Llama, Mistral, Gemma, Qwen2/3, OLMo) collapse into a handful of localized variations: where the residual streams cross norms, whether they use GQA, whether RMSNorm runs on the head dim or the model dim, what the MLP looks like, and how RoPE is parameterized.

This post reads Qwen3ForCausalLM top-down, with the actual source pasted at every layer. By the end you’ll be able to point at any tensor anywhere in the model and say what its shape is, what it’ll be when the next layer touches it, and what op produced it.

The whole file lives at src/transformers/models/qwen3/modeling_qwen3.py. Note the file header — it’s auto-generated from modular_qwen3.py by HF’s modular tooling. The file we read is the flattened version Inductor / dynamo see.

The call stack at a glance

Qwen3ForCausalLM.forward                 (line 445)
  └─ self.model = Qwen3Model              (line 336)
       └─ self.embed_tokens = nn.Embedding
       └─ self.layers = ModuleList[Qwen3DecoderLayer × 28]   (line 233)
       │     └─ self.input_layernorm = Qwen3RMSNorm           (line 50)
       │     └─ self.self_attn = Qwen3Attention               (line 158)
       │     │     ├─ q_proj / k_proj / v_proj / o_proj
       │     │     ├─ q_norm = Qwen3RMSNorm(head_dim)
       │     │     ├─ k_norm = Qwen3RMSNorm(head_dim)
       │     │     └─ attention_interface ∈ {sdpa,fa2,fa3,eager}
       │     └─ self.post_attention_layernorm = Qwen3RMSNorm
       │     └─ self.mlp = Qwen3MLP                            (line 70, SwiGLU)
       └─ self.norm = Qwen3RMSNorm
       └─ self.rotary_emb = Qwen3RotaryEmbedding              (line 299)
  └─ self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
  └─ self.loss_function (set by config; default is causal LM)

Eight callable units total. We’ll walk all of them.

Data flow at a glance (ASCII)

The class tree above shows containment; this shows the tensors as they move forward, with concrete Qwen3-0.6B shapes for B=2, T=2048 in bf16.

                  input_ids  (B, T) = (2, 2048)  int64
                       │
                       ▼
                 embed_tokens                              vocab_size = 151936
                  (V → H)                                  hidden_size = H = 1024
                       │
                       ▼
              hidden_states  (B, T, H) = (2, 2048, 1024)  bf16
                       │
        rotary_emb(position_ids) ──► (cos, sin)  (B, T, head_dim)
                       │            │   shared across all 28 layers
        ┌──────────────┴────────────┴────────────────────────────────┐
        │                                                            │
        │                Qwen3DecoderLayer  × 28                     │
        │                                                            │
        │   hidden_states ──────────●─── residual ─┐                 │
        │         │                                │                 │
        │         ▼                                │                 │
        │    input_layernorm   (RMSNorm over H)    │                 │
        │         │                                │                 │
        │         ▼                                │                 │
        │    self_attn  (q_proj/k_proj/v_proj +    │                 │
        │                q_norm/k_norm + RoPE +    │                 │
        │                attn kernel + o_proj)     │                 │
        │         │                                │                 │
        │         ▼                                │                 │
        │         ⊕ ◄─────────────────────────────┘ + residual       │
        │         │                                                  │
        │   hidden_states ──────────●─── residual ─┐                 │
        │         │                                │                 │
        │         ▼                                │                 │
        │    post_attention_layernorm              │                 │
        │         │                                │                 │
        │         ▼                                │                 │
        │    mlp  (SwiGLU: gate_proj/up_proj/      │                 │
        │          silu/elementwise-mul/down_proj) │                 │
        │         │                                │                 │
        │         ▼                                │                 │
        │         ⊕ ◄─────────────────────────────┘ + residual       │
        │         │                                                  │
        └─────────┼──────────────────────────────────────────────────┘
                  ▼
              norm   (final RMSNorm over H)
                  │
                  ▼
              hidden_states  (B, T, H)
                  │
                  ▼
              lm_head  (H → V)         tied to embed_tokens.weight
                  │
                  ▼
              logits  (B, T, V) = (2, 2048, 151936)  bf16   ← ~1.2 GB
                  │
                  ▼
              cross_entropy(logits, labels, ignore_index=-100)
                  │
                  ▼
              loss  (scalar)

The four RMSNorms per decoder layer (input_layernorm + q_norm + k_norm + post_attention_layernorm) are not visible at this zoom level — they’re folded inside self_attn and the layer’s pre-norm wrapper. Multiplied across 28 layers, the model has 113 RMSNorm calls per forward pass (4 per layer × 28 + 1 final). This becomes relevant when a fused-norm Triton kernel like Liger’s is in play — every one is an HOP node in the FX graph.

Inside one Qwen3Attention.forward, the data flow is:

hidden_states  (B, T, H) = (2, 2048, 1024)
        │
        ├──────────────────┬──────────────────┐
        ▼                  ▼                  ▼
    q_proj             k_proj             v_proj
   (H → Hq)           (H → Hkv)          (H → Hkv)
   1024 → 2048        1024 → 1024        1024 → 1024
        │                  │                  │
        ▼                  ▼                  ▼
  view(B,T,Hq/D,D)   view(B,T,Hkv/D,D)  view(B,T,Hkv/D,D)
   = (B,T,16,128)     = (B,T,8,128)      = (B,T,8,128)
        │                  │                  │
        ▼                  ▼                  │
    q_norm             k_norm                 │
  (RMSNorm over D)  (RMSNorm over D)          │
        │                  │                  │
        ▼                  ▼                  ▼
   transpose(1,2)    transpose(1,2)     transpose(1,2)
  (B,16,T,128)       (B,8,T,128)        (B,8,T,128)
        │                  │                  │
        └────► apply_rotary_pos_emb(q, k, cos, sin)
        │                  │
        ▼                  ▼
        Q                  K (rotated)
        │                  │                  │
        └──────────────────┴──────────────────┘
                            │
                            ▼
            attention_interface(Q, K, V, mask, ...)
              ∈ { eager / sdpa / flash_attention_2 / flash_attention_3 }
              GQA broadcast: K, V repeated num_heads/num_kv_heads = 16/8 = 2×
                            │
                            ▼
                  attn_output  (B, T, Hq) = (2, 2048, 2048)
                            │
                            ▼
                       o_proj  (Hq → H)
                       2048 → 1024
                            │
                            ▼
                  attn_output  (B, T, H) = (2, 2048, 1024)

Notation: H = hidden_size = 1024, Hq = num_heads * head_dim = 2048, Hkv = num_kv_heads * head_dim = 1024, D = head_dim = 128. The asymmetry Hq ≠ H is the consequence of head_dim=128 being explicitly larger than hidden_size / num_heads = 64 — it makes q_proj and o_proj rectangular.

Layer 0 — Qwen3ForCausalLM: head + trunk + loss

modeling_qwen3.py:429. Three fields: model (the trunk), lm_head (logit projection), and vocab_size. tie_word_embeddings=True in Qwen3-0.6B’s config means lm_head.weight is the same parameter as embed_tokens.weight — saved memory, slightly faster cold-load.

The body of forward (line 480) is small enough to paste in full:

outputs: BaseModelOutputWithPast = self.model(
    input_ids=input_ids,
    attention_mask=attention_mask,
    position_ids=position_ids,
    past_key_values=past_key_values,
    inputs_embeds=inputs_embeds,
    use_cache=use_cache,
    cache_position=cache_position,
    **kwargs,
)

hidden_states = outputs.last_hidden_state
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])

loss = None
if labels is not None:
    loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

return CausalLMOutputWithPast(
    loss=loss,
    logits=logits,
    ...
)

Three observations worth pinning down before we go deeper:

1. logits_to_keep is an inference-time micro-opt. At training, it’s 0 (default), the slice is slice(0, None) (the whole sequence), and logits has shape (B, T, V) = (B, T, 151936). At decode, you can pass logits_to_keep=1 to project only the last position — saves a ~0.5 GB GEMM on Qwen3-0.6B at B=1.
2. self.loss_function is not defined in this class — it’s set on the parent PreTrainedModel by config, default loss_utils.ForCausalLMLoss which does the standard shift-and-mean cross-entropy. Liger’s lce_forward (used in this MP’s training pipeline) replaces this whole forward and folds lm_head + cross_entropy into one fused Triton kernel — see the Liger callout at the end.
3. The (B·T, V) logit tensor is what makes a 0.6B model surprisingly memory-hungry. With bs=64, seq=4096, V=151936, you’re materializing 64×4096×151936 ≈ 40 G elements. In fp32 that’s 160 GB, well past any single GPU. In bf16 it’s 80 GB — still a lot. Hence the existence of fused-CE kernels.

Layer 1 — Qwen3Model: the trunk

modeling_qwen3.py:336. Five fields: embed_tokens, layers, norm, rotary_emb, plus a couple of small flags. The forward (line 356) is mostly bookkeeping — preparing position_ids, the causal mask, and the cache — and then a single layer loop.

Stripped to the essentials:

inputs_embeds = self.embed_tokens(input_ids)         # (B, T, hidden_size)
hidden_states = inputs_embeds

position_embeddings = self.rotary_emb(hidden_states, position_ids)  # (cos, sin)

for decoder_layer in self.layers[: self.config.num_hidden_layers]:
    hidden_states = decoder_layer(
        hidden_states,
        attention_mask=causal_mask_mapping[decoder_layer.attention_type],
        position_ids=position_ids,
        position_embeddings=position_embeddings,
        ...
    )

hidden_states = self.norm(hidden_states)
return BaseModelOutputWithPast(last_hidden_state=hidden_states, ...)

Two design choices to call out:

RoPE is computed once. self.rotary_emb(hidden_states, position_ids) returns a (cos, sin) tuple computed at the trunk level and threaded through every layer’s attention as position_embeddings. Old transformer codebases used to recompute RoPE inside each attention block — Qwen3 (and modern HF in general) shares it, since it depends only on position_ids, not on layer state.

Per-layer attention type. causal_mask_mapping is a dict with two keys: full_attention and (optionally) sliding_attention. Each layer has an attention_type field set from config.layer_types[layer_idx]. Qwen3-0.6B has sliding_window: null so all 28 layers are full_attention. Larger Qwen3 variants (and Qwen2.5) interleave full and windowed layers — this is where that machinery lives. The mask itself is built once via create_causal_mask / create_sliding_window_causal_mask.

Layer 2 — Qwen3DecoderLayer: the pre-norm transformer block

modeling_qwen3.py:233. Four submodules: input_layernorm, self_attn, post_attention_layernorm, mlp. The forward is the canonical pre-norm transformer dance:

def forward(self, hidden_states, ...):
    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    hidden_states, _ = self.self_attn(
        hidden_states=hidden_states, ...
    )
    hidden_states = residual + hidden_states

    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    hidden_states = self.mlp(hidden_states)
    hidden_states = residual + hidden_states
    return hidden_states

Pre-norm = the layernorm runs before attention/MLP, the residual stream stays in the original dtype, and the norm output feeds the heavy block. Compare to the original “Attention Is All You Need” post-norm, which puts the norm after the residual add. Pre-norm is universal in modern open LLMs — it’s the variant that trains stably without warmup / gradient clipping at trillion-token scale.

Counting the RMSNorm calls inside one layer: input_layernorm + post_attention_layernorm + (inside attention) q_norm + k_norm = 4 RMSNorms per decoder layer. Across 28 layers that’s 112 RMSNorm calls per forward pass, plus the trunk’s final norm = 113. This matters for graph-break analysis when a fused-norm kernel like Liger’s is involved (each one becomes a triton_kernel_wrapper_functional HOP node).

Layer 3 — Qwen3Attention: GQA + q_norm/k_norm + RoPE + dispatch

This is where most of the model-specific identity lives. modeling_qwen3.py:158.

Init: shapes that don’t equal hidden_size / num_heads

self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
self.scaling = self.head_dim**-0.5

self.q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=False)
self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False)
self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False)
self.o_proj = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=False)

self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None

For Qwen3-0.6B:

• head_dim = 128 (config-set, not hidden_size // num_heads = 64). This is unusual. It means q_proj projects to num_heads × head_dim = 16 × 128 = 2048 (twice hidden_size). o_proj projects back from 2048 → 1024. So q_proj and o_proj aren’t square — they’re 1024×2048 and 2048×1024.
• k_proj, v_proj project to num_kv_heads × head_dim = 8 × 128 = 1024. K and V have half as many heads as Q — that’s GQA.
• num_key_value_groups = 16 / 8 = 2. Each KV head is shared by 2 query heads.

• q_norm and k_norm are RMSNorm with hidden_size = head_dim = 128 (not the model dim). They’re applied per-head. This is the line:

  # unlike olmo, only on the head dim!

  pointing out that OLMo applies a similar “QK norm” but at model_dim. Qwen3’s variant is per-head — cheaper and arguably more consistent with the per-head attention computation.

Forward: the four-step dance

def forward(self, hidden_states, position_embeddings, attention_mask, ...):
    input_shape = hidden_states.shape[:-1]              # (B, T)
    hidden_shape = (*input_shape, -1, self.head_dim)    # (B, T, ?, 128)

    # 1. Project + reshape + per-head norm + transpose for attention
    query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
    # -> (B, num_heads,    T, head_dim) = (B, 16, T, 128)

    key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
    # -> (B, num_kv_heads, T, head_dim) = (B,  8, T, 128)

    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    # -> (B, num_kv_heads, T, head_dim) = (B,  8, T, 128)

    # 2. Apply RoPE to Q and K (V is rotation-invariant)
    cos, sin = position_embeddings
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

    # 3. (Optional) update KV cache for autoregressive decoding
    if past_key_values is not None:
        key_states, value_states = past_key_values.update(...)

    # 4. Dispatch to attention kernel (eager / sdpa / fa2 / fa3)
    attention_interface = eager_attention_forward
    if self.config._attn_implementation != "eager":
        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

    attn_output, attn_weights = attention_interface(
        self, query_states, key_states, value_states, attention_mask,
        dropout=0.0 if not self.training else self.attention_dropout,
        scaling=self.scaling,
        sliding_window=self.sliding_window,
        **kwargs,
    )

    # 5. Reshape + output projection
    attn_output = attn_output.reshape(*input_shape, -1).contiguous()  # (B, T, num_heads*head_dim) = (B, T, 2048)
    attn_output = self.o_proj(attn_output)                            # (B, T, hidden_size)        = (B, T, 1024)
    return attn_output, attn_weights

The key shape transformations, with concrete Qwen3-0.6B numbers and B=2, T=2048:

hidden_states                       : (2, 2048, 1024)   bf16   16 MB
q_proj(hidden_states)               : (2, 2048, 2048)   bf16   16 MB    ← projects up
view(... -1, 128)                   : (2, 2048,   16, 128)
q_norm                              : (2, 2048,   16, 128)             ← per-head RMS
.transpose(1, 2)                    : (2,   16, 2048, 128)             ← (B, H, T, D)

k_proj(hidden_states)               : (2, 2048, 1024)   bf16    8 MB    ← K projects to half
view(... -1, 128)                   : (2, 2048,    8, 128)
k_norm                              : (2, 2048,    8, 128)
.transpose(1, 2)                    : (2,    8, 2048, 128)

apply_rotary_pos_emb(q, k, cos, sin)
  q                                 : (2,   16, 2048, 128)
  k                                 : (2,    8, 2048, 128)

attention_interface(...)
  attn_output                       : (2, 2048, 16, 128) (varies by backend; SDPA returns (B, T, H, D))

attn_output.reshape(*input_shape, -1): (2, 2048, 2048)
o_proj                              : (2, 2048, 1024)   ← back to hidden_size

apply_rotary_pos_emb — what RoPE actually does

modeling_qwen3.py:93. The implementation is short:

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

cos, sin are produced by Qwen3RotaryEmbedding.forward (line 321) and have shape (B, T, head_dim). unsqueeze(1) makes them (B, 1, T, head_dim) so they broadcast against q shaped (B, num_heads, T, head_dim). The math is the standard “complex multiplication via the half-rotation trick” — pair adjacent dims as $(x_{2i}, x_{2i+1})$, apply a 2D rotation by angle $\theta_{i,t} = t \cdot 10000^{-2i/D}$.

In Qwen3, the base frequency is rope_theta=1000000 (note: 1M, not the usual 10K from the original LLaMA). This expands the effective context window — concretely, Qwen3-0.6B’s max_position_embeddings=40960 is supported by setting θ this high so the lowest-frequency embedding still rotates ≪ 1 full turn over 40K positions.

The attention dispatch

ALL_ATTENTION_FUNCTIONS is a registry: {"eager", "sdpa", "flash_attention_2", "flash_attention_3", "flex_attention", ...} keyed by string. The chosen one is set at from_pretrained time via attn_implementation=, and stored on config._attn_implementation. The dispatch is literally a dict lookup:

attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

The dispatched function receives module, query, key, value, attention_mask, scaling, dropout, sliding_window, **kwargs. Each backend handles GQA differently:

• eager (line 132) calls repeat_kv(k, num_kv_groups) to materialize K and V at full head-count and runs softmax(QK^T / √d) V in pure aten.
• sdpa uses torch.nn.functional.scaled_dot_product_attention, which has a enable_gqa=True flag and handles the K/V repeat internally.
• flash_attention_2 and flash_attention_3 are real torch.library.custom_ops that take Q/K/V at GQA shapes and handle the broadcast in CUDA kernels.

For our compile path, this matters because the attention kernel is the only place where the model interacts with code outside torch._inductor’s direct codegen reach: SDPA stays in aten land, FA goes through a custom op. That structural seam is the source of the FA × Liger interaction documented in the previous post on Inductor codegen.

Layer 4 — Qwen3MLP: SwiGLU

modeling_qwen3.py:70. The whole class:

class Qwen3MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]   # "silu" for Qwen3 → torch.nn.functional.silu

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

The one-line forward is SwiGLU: down_proj(silu(gate_proj(x)) * up_proj(x)).

Three matmuls, one elementwise SiLU, one elementwise multiply. Shapes for Qwen3-0.6B with B=2, T=2048:

x                           : (2, 2048, 1024)        ← hidden_size
gate_proj(x)                : (2, 2048, 3072)        ← intermediate_size
silu(gate_proj(x))          : (2, 2048, 3072)
up_proj(x)                  : (2, 2048, 3072)
elementwise multiply        : (2, 2048, 3072)
down_proj                   : (2, 2048, 1024)

Compute is dominated by the three GEMMs: 2 × 2048 × 1024 → 3072 (twice), then 2 × 2048 × 3072 → 1024. About 78M params per layer in MLP (3 * 1024 * 3072 ≈ 9.4M, but that’s per-layer); summed across 28 layers, the MLP is ~63 % of total params (296M of 463M non-embedding). Embeddings (tied with lm_head) are another ~155M.

Why two input projections (gate and up) and not one? The “gating” view: silu(gate_proj(x)) outputs a soft mask over the intermediate_size channels, and up_proj(x) is the value being masked. This is the GLU family — silu(g(x)) * f(x) instead of just silu(f(x)). Empirically gives a small quality bump for the same parameter count vs. plain SiLU MLP, and is now standard.

Primitives

Qwen3RMSNorm

modeling_qwen3.py:50:

class Qwen3RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

Standard RMSNorm: divide by RMS along the last dim, multiply by a learned per-channel scale. Two things to note:

1. It’s not LayerNorm. No mean subtraction, no learned bias. This is the simpler-and-often-better variant from Zhang & Sennrich 2019.
2. The fp32 upcast on entry, downcast on exit. Both the .pow(2) and the running sum are unsafe in bf16 — mean over head_dim=128 of bf16 squared values loses precision fast. The cast bumps it to fp32 for the reduction, then casts back.

Qwen3RotaryEmbedding

modeling_qwen3.py:299. The interesting part is the forward:

@torch.no_grad()
@dynamic_rope_update
def forward(self, x, position_ids):
    inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
    position_ids_expanded = position_ids[:, None, :].float()

    device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
    with torch.autocast(device_type=device_type, enabled=False):  # Force float32
        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos() * self.attention_scaling
        sin = emb.sin() * self.attention_scaling

    return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

Outputs (cos, sin) of shape (B, T, head_dim). Two notable choices:

• @torch.no_grad(): no learnable parameters in RoPE — it’s pure math from position_ids and a fixed inv_freq buffer.
• with torch.autocast(..., enabled=False): forces fp32 computation even under bf16 autocast. The trig functions and the long-position outer product are sensitive enough to dtype that the model degrades visibly without the upcast.

inv_freq itself is computed once at init (via ROPE_INIT_FUNCTIONS[rope_type]), based on config.rope_theta and the head dim. For Qwen3 with rope_theta=1000000 and head_dim=128:

\[\text{inv\_freq}_i = 10^{6\, \cdot\, -2i/128}, \quad i = 0, 1, \ldots, 63\]

Note: 64 frequencies for 128 dims because each frequency rotates a 2D pair.

Three Qwen3-specific things worth memorizing

Most of modeling_qwen3.py looks like LLaMA. Three deltas you’ll keep tripping over:

1. q_norm / k_norm per-head, applied between the q_proj/k_proj and the RoPE step. Adds ~64K params per layer (negligible) but visibly improves stability at long context. This is also the “QK-norm” you may have seen in OLMo and Gemma2 — Qwen3 is the per-head variant.

2. head_dim is configurable, not derived. Qwen3-0.6B sets head_dim=128 even though hidden_size / num_attention_heads = 64. This makes q_proj rectangular (1024 → 2048). Surface effect: the Q dimension grows but the K/V dimensions don’t (because of GQA), so the FLOPs split is heavier on the Q-side projection than in vanilla MHA.

3. Per-layer attention_type — read from config.layer_types[layer_idx], with values "full_attention" or "sliding_attention". The trunk’s mask construction emits one mask per type, and each layer picks via dict lookup. Qwen3-0.6B has all-full_attention, but variants with sliding window use this machinery.

Loss calculation in detail

Loss is the part of the model where you can blow up an H200 GPU with a single line. Worth understanding it carefully.

The cross-entropy math

Given logits $z \in \mathbb{R}^V$ for one token and a target class $y \in {0, \ldots, V-1}$, cross-entropy is

\[L(z, y) \;=\; -\log p(y \mid z) \;=\; -\log \frac{\exp(z_y)}{\sum_{j=0}^{V-1} \exp(z_j)} \;=\; \mathrm{logsumexp}(z) - z_y\]

For a batch of $N$ tokens, with reduction "mean":

\[L \;=\; \frac{1}{N_{\text{valid}}} \sum_{i \,:\, y_i \neq -100} \big(\mathrm{logsumexp}(z_i) - z_{i,\,y_i}\big)\]

Two add-ons HF and Liger both support:

• Label smoothing $(\alpha)$ — replaces the one-hot target with a soft distribution $q_j = (1{-}\alpha)\,\delta_{j,y} + \alpha/V$ uniform on the rest. Loss becomes $L_\alpha = (1{-}\alpha) \cdot L + \alpha \cdot (\mathrm{logsumexp}(z) - \bar z)$. Set in our config via label_smoothing_factor: 0.1.
• z-loss $(\beta)$ — adds $\beta \cdot \mathrm{logsumexp}(z)^2$ as a regularizer that pulls logits toward smaller magnitudes. Liger calls this lse_square_scale. Default 0; useful at long context where unconstrained logits drift.

The thing that matters for memory: the loss formula evaluates logsumexp over V — you cannot compute it without something logically the size of $z$, which is $V$ floats per token. We’ll come back to this.

Finetuning with labels on only some tokens

Causal-LM finetuning rarely cares about the prompt — you want gradient only on the response tokens. The standard recipe:

1. Tokenize prompt + response as one sequence.
2. Build labels as a copy of input_ids, then mask the prompt positions to -100. The remaining positions are real token ids.
3. PyTorch’s nn.functional.cross_entropy(..., ignore_index=-100) skips both the numerator and the denominator at masked positions — they don’t enter the loss or the gradient.

In this MP’s data.py (line 654-672 of training/data.py):

labels = input_ids.copy()

# Find the assistant marker (Qwen chat template), mask everything before it
response_template = "<|im_start|>assistant\n<think>\n\n</think>\n\n"
response_template_ids = self.tokenizer.encode(response_template, add_special_tokens=False)
template_len = len(response_template_ids)
for i in range(len(input_ids) - template_len + 1):
    if input_ids[i:i+template_len] == response_template_ids:
        labels[:i+template_len] = [-100] * (i + template_len)  # mask prompt
        break

# (also mask pad positions if padding="max_length" was used)
labels = [t if m == 1 else -100 for t, m in zip(labels, attention_mask)]

Effect: a 2048-token packed sequence with ~200 prompt tokens + ~50 response tokens + ~1798 pad tokens contributes loss from only ~50 tokens. The other 1998 are -100 and contribute zero gradient.

Shift-by-one alignment

A causal LM at position $t$ predicts the token at position $t+1$. So the loss at position $t$ uses logits[t] against labels[t+1]. HF’s ForCausalLMLoss (loss_utils.py:45-66) does this with a single pad+slice:

def ForCausalLMLoss(logits, labels, vocab_size, num_items_in_batch=None, ignore_index=-100, ...):
    # Upcast to float to avoid bf16 precision loss in logsumexp
    logits = logits.float()                                      # ← memory bomb here

    if shift_labels is None:
        # Shift so that tokens < n predict n
        labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
        shift_labels = labels[..., 1:].contiguous()              # left-shift labels by 1

    # Flatten the tokens
    logits = logits.view(-1, vocab_size)                         # (B*T, V)
    shift_labels = shift_labels.view(-1)                         # (B*T,)
    shift_labels = shift_labels.to(logits.device)
    loss = fixed_cross_entropy(logits, shift_labels, num_items_in_batch, ignore_index, **kwargs)
    return loss

Note line 53 — logits = logits.float(). This is where Stage A used to OOM. Read on.

Why GPU memory explodes at the loss

For Qwen3-0.6B at production scale (bs=64, seq=4096, V=151936):

\[B \cdot T \cdot V \;=\; 64 \cdot 4096 \cdot 151936 \;\approx\; 39.8 \times 10^9 \text{ elements}\]

In bf16: $39.8\,\text{Gel} \times 2\,\text{B} = 79.6\,\text{GB}$. After the .float() upcast: 159 GB. A single H200 has 144 GB HBM.

The OOM doesn’t happen during the actual cross-entropy computation. It happens at the logits.float() line above — PyTorch allocates 79.6 GB for the upcasted tensor while the bf16 logits are still alive (another 39.8 GB). The peak is around 120 GB just for the logit tensor and its fp32 copy, before any gradient is even computed.

Why the upcast is unavoidable in naive CE: logsumexp over V=151936 values in bf16 loses too much precision. bf16 has 8 mantissa bits — exp(x) for x>17 overflows in bf16. Real logit magnitudes routinely reach 30+ during training. fp32 buys 23 mantissa bits + 8 exponent bits, comfortable headroom.

So you’re stuck: keep bf16 → numerical garbage. Cast to fp32 → blow up memory.

How Liger’s chunked fused-CE solves it

Source: liger_kernel/ops/fused_linear_cross_entropy.py:17-200. The trick is to fuse the lm_head projection and the cross-entropy into one operation, processed in chunks of $BT$, so the full $(BT, V)$ logit tensor is never materialized.

The chunking math is documented in the source comment (line 45-51):

# inputs have shape: BT x H
# materialized activations will have shape: BT x V
# the increase in memory = BT x V
# reduction can be achieved by partitioning the number of tokens BT into smaller chunks.
# for ex: if we were to achieve the same memory consumption as BT x H, then the chunk size should be:
# inc_factor = (V+H-1)//H, chunk_size = (BT + inc_factor - 1)//inc_factor
# for ex: BT = 4096*4, V = 32000, H = 4096 ==> inc_factor = 8, chunk_size = 2048

Translation: pick chunk_size so that one chunk’s (chunk_size, V) logit tile is the same memory as the input’s (BT, H) hidden tile. For Qwen3-0.6B at production scale:

Quantity	Value
$BT$ (total tokens)	$64 \cdot 4096 = 262{,}144$
$H$ (hidden_size)	$1024$
$V$ (vocab_size)	$151{,}936$
$V/H$ inc_factor	$\lceil 151936/1024 \rceil = 149$
chunk_size	$\mathrm{nextpow2}(\lceil BT / 149 \rceil) = \mathrm{nextpow2}(1759) = 2048$
num_chunks	$\lceil 262144 / 2048 \rceil = 128$

Each chunk materializes a (2048, 151936) fp32 logit slab — that’s 1.25 GB instead of 160 GB, 128× smaller. Peak memory for the loss path drops below the residual stream, no longer the bottleneck.

The chunk loop in the implementation (line 96-205):

for chunk_id in range(num_chunks):
    start_idx = chunk_id * chunk_size
    end_idx = min((chunk_id + 1) * chunk_size, BT)
    _input_chunk = _input[start_idx:end_idx]                 # chunk_size x H

    # GEMM in original precision (bf16) — this is the lm_head projection
    logits_chunk = _input_chunk @ weight.t()                 # chunk_size x V

    target_chunk = target[start_idx:end_idx]                 # chunk_size,

    # Triton kernel computes:
    #   (a) logsumexp over V, in fp32 with online numerical stability
    #   (b) per-token loss = logsumexp - z_target
    #   (c) gradient w.r.t. logits, IN-PLACE in logits_chunk (now overwritten)
    liger_cross_entropy_kernel[(n_rows,)](
        X_ptr=logits_chunk,    # used as both input AND output (gets overwritten with grad)
        Y_ptr=target_chunk,
        loss_ptr=loss_1d_slice,
        ...
    )
    grad_logits_chunk = logits_chunk    # alias — same memory, now contains dL/dlogits

    # Backprop the chunk's logit-grad to input-grad and weight-grad accumulators
    grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
    if grad_weight is not None:
        torch.addmm(grad_weight, grad_logits_chunk.t(), _input_chunk.to(grad_weight.dtype), ...)

Three things make this work:

1. The Triton kernel writes the gradient back into the same memory as the logits. Look for # Here we calculate the gradient of logits_chunk in place so we can save memory. at line 150. This means we never need to allocate a separate grad_logits tensor — the input slot is reused.

2. Backprop is done chunk-locally. grad_input[start_idx:end_idx] = grad_logits_chunk @ weight produces the input gradient for this chunk. Weight gradient is accumulated across chunks via addmm.

3. The cross-entropy kernel itself is online-stable. It computes logsumexp in a single pass over $V$ with running max + exp-sum bookkeeping (the standard “online softmax” trick from FlashAttention’s softmax). Numerical fp32 only inside the kernel; outputs are accumulated to a fp32 loss_1d buffer.

Net effect at production scale: peak HBM during loss drops from ~160 GB to ~1.25 GB. That’s the entire reason this MP can train Stage A at bs=64, seq=4096, V=151936 on a single H200 node.

The cost: the chunk loop is sequential (each chunk depends on having received its grad from the kernel before launching the matmul-back-to-input). That’s a small SM occupancy hit vs. one giant fused op, but irrelevant compared to the memory win.

The Liger entry point that the model’s forward actually calls is LigerForCausalLMLoss, which does the same pad+slice shift as HF’s ForCausalLMLoss and then dispatches to liger_fused_linear_cross_entropy instead of the dense cross-entropy.

What this means for compile

The chunked fused-CE has one wart that breaks torch.compile cleanliness: the Triton kernel needs total_n_non_ignore (the count of non--100 targets) to compute the mean correctly, and the way Liger gets it is

total_n_non_ignore = target_mask.sum().item()    # line 81

.item() extracts a Python int. Without torch._dynamo.config.capture_scalar_outputs = True, dynamo graph-breaks at this line and then specializes on the int value, baking it into the compiled graph. Each batch’s count is different → fresh fx_graph_cache key per batch → thousands of recompiles.

That’s the recompile cascade we documented in the previous post on torch.compile internals — and the reason this MP currently ships with enable_torch_compile: false whenever use_liger_kernel: true. The fix is the one-line config change above. Memory savings stay; cache thrash goes away.

Liger overlay (the version you’ll actually train)

When training with use_liger_kernel=True on Qwen3, liger_kernel.transformers.monkey_patch.apply_liger_kernel_to_qwen3 does four substitutions:

Liger swaps	What gets replaced
Qwen3RMSNorm → LigerRMSNorm	All 113 RMSNorm calls (input_ln, post_attn_ln, q_norm, k_norm per layer + final norm)
Qwen3MLP → LigerSwiGLUMLP	The 28 MLP forwards
apply_rotary_pos_emb → liger_rotary_pos_emb	RoPE per layer
Qwen3ForCausalLM.forward → lce_forward	The whole top-level forward (replaces lm_head + cross_entropy with a fused linear-CE Triton kernel that never materializes the (B·T, V) logit tensor)

That last swap is the big memory win — at bs=64, seq=4096, V=151936 the logit tensor would be ~80 GB in bf16, well past any single GPU. Fused linear-CE keeps things in tile-sized chunks and never spills the full logits.

Where to go from here

• Qwen3Config — src/transformers/models/qwen3/configuration_qwen3.py — the canonical list of all knobs (head_dim, sliding_window, rope_scaling, layer_types, attention_bias, etc.).
• The modular source — src/transformers/models/qwen3/modular_qwen3.py — the file modeling_qwen3.py is generated from. Instructive if you want to see what Qwen3 inherits from Llama unchanged vs. what’s overridden.
• Mask construction — src/transformers/masking_utils.py — create_causal_mask and create_sliding_window_causal_mask are surprisingly involved (they handle padding-free, varlen, sliding, packed, prefix-cache).
• RoPE init functions — src/transformers/modeling_rope_utils.py — ROPE_INIT_FUNCTIONS registry, including default, linear, dynamic, yarn, longrope, llama3. Different scaling laws for different long-context strategies.
• The eager attention kernel — modeling_qwen3.py:132 — eager_attention_forward. Twenty lines, the textbook QK^T / softmax / attn @ V dance, including the repeat_kv for GQA.

Once these click, every other LLaMA-family model in transformers reads in 30 minutes. The variation surface is a small handful of axes and the file structure is the same.