modeling_loop_lm.py

"""Looped-MoE modeling code (port of the seedvar `loop-lm` architecture).

Provenance
----------
Ported from ``modeling_loop_lm.py`` as published with
``ml-ryanlee/seedvar-looped-moe-1e18-d704-seed42..47`` (arXiv 2605.09165,
*Sparse Layers are Critical to Scaling Looped Language Models*). That file is the
architecture specification: the published checkpoints are its training product,
and the diagnostics pipeline is already validated against it.

This is a **port, not a copy**. The numerics are kept faithful (see "Faithful to
the original" below) because the baseline's whole job is to reproduce the
published architecture; the deviations are all in service of three requirements
from HANDOVER §4.3 that the original file does not meet:

  1. **Semantic parameters have no defaults.** The original ``LoopLMConfig``
     defaults ``d_model=1024``, ``num_experts=8`` and so on. Defaults are how a
     silently-wrong run happens: a typo'd key name falls back to a plausible
     number and the run looks fine. Every shape/semantic field here is required
     and a missing one raises.
  2. **Every loop step is hookable, by explicit index.** The loop counter is
     threaded down to each block, so a consumer never recovers the loop axis by
     reshaping a flattened layer axis (which silently yields transposed
     semantics).
  3. **Router logits are exposed with their semantics labelled.** Both the
     pre-softmax logits and the post-softmax probabilities are handed out, tagged
     via ``src.model.loop_trace.ROUTER_LOGITS_ARE_PRESOFTMAX``.

Scope
-----
Only the ``looped-moe`` variant is implemented. The original file carries four
variants (base / looped / moe / looped-moe); all four architectures this project
trains -- baseline and DVF-a/b/c -- are looped-moe, differing only in the (L, R, E)
triple. Porting the unused three would be dead code (project rule: no entities
beyond necessity). They remain available in the original file if ever needed.

Shape parameters, and what the experiment varies
------------------------------------------------
====================  ======  =========================================
config field          symbol  meaning
====================  ======  =========================================
num_layers_in_stack   L       physical layers in the shared stack
num_stacks            R       times the stack is called (loop count)
num_experts           E       experts per MoE layer
num_active            k       experts activated per token
====================  ======  =========================================

The DVF ("dual vector foil") series holds L*R = 16 and E*L = 64 fixed and only
moves capacity around: baseline (8,2,8), DVF-a (4,4,16), DVF-b (2,8,32),
DVF-c (1,16,64).

Faithful to the original (do not "fix" these -- they are muP, not bugs)
----------------------------------------------------------------------
* ``RMSNorm`` has **no** gain parameter.
* Attention is scaled by ``1/d_k``, not ``1/sqrt(d_k)``.
* Softmax upcasts to float32.
* Expert FFN width is ``d_ff // num_active`` -- divided by k, *not* by E. This is
  what makes E*L=64 hold total expert parameters constant across the DVF series.
* Initialisation is muP: ``std = std_base / sqrt(width_ratio)`` with
  ``std_base = sqrt(2/(fan_in_base + fan_out_base))`` against a d_base=128 proxy.

Deliberate deviation
--------------------
``RotaryPositionalEmbedding`` builds its rotation table with vectorised torch ops
instead of the original's ``max_seq_len * d_k/2`` nested Python loop, which costs
minutes at seq_len 4096. ``tests/test_rope_equivalence.py`` pins the vectorised
table against a literal transcription of the original loop.
"""

from __future__ import annotations

import math
from typing import Any, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import einsum, rearrange, reduce, repeat
from torch import Tensor
from torch.nn.functional import grouped_mm, silu
from transformers import PretrainedConfig, PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

# This file is loaded in three different ways, and the trace types have to
# resolve in all of them:
#   1. as part of this repo             -> `src.model.loop_trace`
#   2. via HuggingFace trust_remote_code -> copied into a generated package under
#      `transformers_modules/<ckpt>/`, where the sibling is a *relative* import
#   3. as a loose script with the checkpoint directory on sys.path
# Case 2 is the one that matters for HANDOVER §4.7: the diagnostics pipeline is a
# separate repository that has never heard of `pretrain`, and `save_trajectory`
# bundles `loop_trace.py` next to this file so the checkpoint stands alone.
try:
    from src.model.loop_trace import ACTIVE_TRACE_SINK, TraceMeta, TracePoint, TraceSink
except ImportError:  # pragma: no cover - covered by the cold-load test
    try:
        from .loop_trace import ACTIVE_TRACE_SINK, TraceMeta, TracePoint, TraceSink  # type: ignore[no-redef]
    except ImportError:
        from loop_trace import ACTIVE_TRACE_SINK, TraceMeta, TracePoint, TraceSink  # type: ignore[no-redef]

__all__ = ["LoopMoEConfig", "LoopMoEForCausalLM", "LoopedMoETransformer"]

# muP proxy-model widths. The initialisation std of every weight is derived from
# a d_base=128 model and rescaled by width_ratio = d_model / 128.
HEAD_TAIL_NUM_EXPERTS = 8
"""Experts in a head/tail layer, fixed across every ablation configuration.

Recipe 2026-09-12 section 2.1: head and tail are "identical in all configurations"
(1 layer, 8 experts, top-2). It is deliberately independent of the loop block's E --
S4 gives the loop 64 experts per layer and its head still has 8 -- so the loop block's
count must not be reused here.
"""

BASE_D_MODEL = 128
BASE_D_FF = 384


def softmax(logits: Tensor, dim: int) -> Tensor:
    """Max-shifted softmax in float32 (verbatim semantics from the original)."""
    logits = logits.float()
    max_values = torch.max(logits, dim=dim, keepdim=True).values
    shifted = logits - max_values
    shifted_exps = torch.exp(shifted)
    shifted_exp_sums = torch.sum(shifted_exps, dim=dim, keepdim=True)
    return shifted_exps / shifted_exp_sums


class Linear(nn.Module):
    """Bias-free linear layer with muP initialisation."""

    def __init__(self, in_features, out_features, width_ratio, std_base, device=None, dtype=None):
        super().__init__()
        # Registered before init so the shape exists under HF meta-device loading.
        self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype, device=device))
        # Kept so the init can be replayed: torchtitan builds the model on the
        # meta device and then calls `init_weights()` on materialised (but
        # uninitialised) storage, so constructor-time init alone leaves the
        # model full of garbage.
        self._init_std = std_base / math.sqrt(width_ratio)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        std = self._init_std
        nn.init.trunc_normal_(self.weight, mean=0.0, std=std, a=-3 * std, b=3 * std)

    def forward(self, x: Tensor) -> Tensor:
        return einsum(self.weight, x, "d_out d_in, ... d_in -> ... d_out")


class Embedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(num_embeddings, embedding_dim, dtype=dtype, device=device))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.trunc_normal_(self.weight, mean=0.0, std=1.0, a=-3, b=3)

    def forward(self, token_ids: Tensor) -> Tensor:
        return self.weight[token_ids]


class RMSNorm(nn.Module):
    """RMS norm **without** a gain parameter (muP convention)."""

    def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None):
        super().__init__()
        self.d_model = d_model
        self.eps = eps

    def forward(self, x: Tensor) -> Tensor:
        in_dtype = x.dtype
        x = x.to(torch.float32)
        mean_squared_sum = (1 / self.d_model) * einsum(x, x, "... seq d, ... seq d -> ... seq")
        rms = torch.sqrt(mean_squared_sum + self.eps)
        rms_norm = einsum(x, 1 / rms, "... seq d, ... seq -> ... seq d")
        return rms_norm.to(in_dtype)


class PositionwiseFeedforward(nn.Module):
    """SwiGLU: W2(SiLU(W1 x) * W3 x)."""

    def __init__(self, d_model: int, d_ff: int, width_ratio: float, device=None, dtype=None):
        super().__init__()
        w_std_base = math.sqrt(2 / (BASE_D_MODEL + BASE_D_FF))
        self.w1 = Linear(d_model, d_ff, width_ratio, w_std_base, device=device, dtype=dtype)
        self.w2 = Linear(d_ff, d_model, width_ratio, w_std_base, device=device, dtype=dtype)
        self.w3 = Linear(d_model, d_ff, width_ratio, w_std_base, device=device, dtype=dtype)

    def forward(self, x: Tensor) -> Tensor:
        return self.w2(silu(self.w1(x)) * self.w3(x))


class RotaryPositionalEmbedding(nn.Module):
    """RoPE with a precomputed [seq, d_k/2, 2, 2] rotation table.

    Vectorised rebuild of the original's nested Python loop; see the module
    docstring and ``tests/test_rope_equivalence.py``.
    """

    def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None, dtype=None):
        super().__init__()
        # Retained so the table can be rebuilt: `to_empty()` replaces buffer
        # storage with uninitialised memory just as it does for parameters, so a
        # meta-device build leaves the rotation table as garbage unless
        # `reset_parameters()` regenerates it.
        self._rope_theta, self._rope_d_k = theta, d_k
        self._rope_max_seq_len, self._rope_dtype = max_seq_len, dtype
        rotations = self._build_table(theta, d_k, max_seq_len, device, dtype)
        self.register_buffer("rotations", rotations, persistent=True)

    @staticmethod
    def _build_table(theta: float, d_k: int, max_seq_len: int, device, dtype) -> Tensor:
        """[seq, d_k/2, 2, 2] rotation table.

        Angles are built in float64 and only then cast down. At seq_len 4096 the
        largest angle is ~4096 rad, where float32 spacing is ~2.4e-4; computing
        cos/sin at float32 there loses ~4 decimal digits. The original does this
        implicitly (Python floats are float64), so float64 here is both more
        accurate and what keeps the table equal to the reference.
        """
        positions = torch.arange(max_seq_len, device=device, dtype=torch.float64)
        pair_idx = torch.arange(d_k // 2, device=device, dtype=torch.float64)
        inv_freq = theta ** (2 * pair_idx / d_k)
        angles = positions[:, None] / inv_freq[None, :]
        cos, sin = torch.cos(angles), torch.sin(angles)
        # rows of the 2x2 rotation: [[cos, -sin], [sin, cos]]
        table = torch.stack(
            [torch.stack([cos, -sin], dim=-1), torch.stack([sin, cos], dim=-1)], dim=-2
        )
        return table.to(dtype if dtype is not None else torch.float32)

    @torch.no_grad()
    def reset_parameters(self) -> None:
        """Regenerate the rotation table in place (buffers survive nothing)."""
        self.rotations.copy_(
            self._build_table(
                self._rope_theta, self._rope_d_k, self._rope_max_seq_len,
                self.rotations.device, self.rotations.dtype,
            )
        )

    def forward(self, x: Tensor, token_positions: Tensor) -> Tensor:
        rot = self.rotations[token_positions].to(dtype=x.dtype)
        x_pairs = rearrange(x, "... seq_dim (feature_dim i) -> ... seq_dim feature_dim i", i=2)
        y_pairs = einsum(
            rot,
            x_pairs,
            "... seq_dim feature_dim i j, ... seq_dim feature_dim j -> ... seq_dim feature_dim i",
        )
        return rearrange(y_pairs, "... seq_dim feature_dim i -> ... seq_dim (feature_dim i)")


class MultiheadSelfAttention(nn.Module):
    """Causal MHSA with RoPE. muP: attention logits scaled by 1/d_k."""

    def __init__(self, d_model: int, num_heads: int, max_seq_len: int, theta: float,
                 width_ratio: float, device=None, dtype=None):
        super().__init__()
        if d_model % num_heads != 0:
            raise ValueError(f"d_model ({d_model}) must be divisible by num_heads ({num_heads})")
        self.d_model = d_model
        self.num_heads = num_heads

        attn_std_base = math.sqrt(2 / (BASE_D_MODEL + BASE_D_MODEL))
        self.q_proj = Linear(d_model, d_model, width_ratio, attn_std_base, device=device, dtype=dtype)
        self.k_proj = Linear(d_model, d_model, width_ratio, attn_std_base, device=device, dtype=dtype)
        self.v_proj = Linear(d_model, d_model, width_ratio, attn_std_base, device=device, dtype=dtype)
        self.output_proj = Linear(d_model, d_model, width_ratio, attn_std_base, device=device, dtype=dtype)
        self.rope = RotaryPositionalEmbedding(theta, d_model // num_heads, max_seq_len, device, dtype)

    def forward(self, x: Tensor, token_positions: Optional[Tensor] = None) -> Tensor:
        d_k = self.d_model // self.num_heads
        q_heads = rearrange(self.q_proj(x), "... seq (heads d_k) -> ... heads seq d_k", d_k=d_k)
        k_heads = rearrange(self.k_proj(x), "... seq (heads d_k) -> ... heads seq d_k", d_k=d_k)
        v_heads = rearrange(self.v_proj(x), "... seq (heads d_v) -> ... heads seq d_v", d_v=d_k)

        if token_positions is None:
            token_positions = rearrange(torch.arange(x.shape[-2], device=x.device), "seq -> 1 seq")
        q_heads = self.rope(q_heads, token_positions)
        k_heads = self.rope(k_heads, token_positions)

        mha_heads = F.scaled_dot_product_attention(
            q_heads, k_heads, v_heads, is_causal=True, scale=1.0 / d_k
        )
        return self.output_proj(rearrange(mha_heads, "... heads seq d_v -> ... seq (heads d_v)"))


class Router(nn.Module):
    """Top-k softmax router. Returns pre-softmax logits *and* probabilities.

    The two are returned side by side, and the caller labels which is which via
    ``ROUTER_LOGITS_ARE_PRESOFTMAX``. There is no jitter noise and no temperature;
    routing is deterministic given the input (matching the original).
    """

    def __init__(self, d_model: int, num_experts: int, num_active: int, width_ratio: float,
                 device=None, dtype=None):
        super().__init__()
        std_base = math.sqrt(2 / (BASE_D_MODEL + num_experts))
        self.gate = Linear(d_model, num_experts, width_ratio, std_base, device=device, dtype=dtype)
        self.num_active = num_active

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
        logits = self.gate(x)  # [B, S, E] -- pre-softmax
        probs = softmax(logits, dim=-1)  # [B, S, E] -- over all E experts
        top_scores, top_experts = torch.topk(probs, k=self.num_active, dim=-1)
        # Renormalise within the selected set so the combine weights sum to 1.
        top_scores = top_scores / torch.sum(top_scores, dim=-1, keepdim=True)
        return logits, probs, top_scores, top_experts


class GroupedMoEPrenormBlock(nn.Module):
    """Pre-norm block whose FFN is a grouped top-k MoE.

    Layout: x -> +attn(ln1(x)) -> +moe(ln2(.)). Aux losses are returned rather
    than stashed on the module, so nothing has to be reset between loop steps.
    """

    @staticmethod
    def _init_expert_weights(num_experts, in_features, out_features, width_ratio, std_base,
                             device, dtype) -> nn.Parameter:
        w = torch.empty(num_experts, in_features, out_features, device=device, dtype=dtype)
        std_scaled = std_base / math.sqrt(width_ratio)
        nn.init.trunc_normal_(w, mean=0.0, std=std_scaled, a=-3 * std_scaled, b=3 * std_scaled)
        return nn.Parameter(w)

    @torch.no_grad()
    def reset_parameters(self) -> None:
        """Re-init the grouped expert weights (see Linear.reset_parameters)."""
        std = self._expert_init_std
        for w in (self.experts_w1, self.experts_w2, self.experts_w3):
            nn.init.trunc_normal_(w, mean=0.0, std=std, a=-3 * std, b=3 * std)

    def __init__(self, d_model: int, num_heads: int, d_ff: int, num_experts: int, num_active: int,
                 max_seq_len: int, theta: float, width_ratio: float, device=None, dtype=None,
                 num_attention_sets: int = 1):
        """`num_attention_sets` = how many independent attention parameter sets
        this block holds, one per pass over the looped stack.

        1 is the shared-attention architecture every run before 2026-09-18 used.
        It keeps `self.attn` a single module, so the state dict key stays
        `attn.q_proj.weight` and existing checkpoints load unchanged -- a
        ModuleList of one would rename every key to `attn.0.*` and silently
        invalidate every archive we hold.

        >1 makes `self.attn` a ModuleList of that many attention modules, each
        initialised independently under the same muP standard deviation, chosen
        at forward time by `loop_step`. Everything else in the block -- norms,
        router, experts -- remains shared across passes, which is the point of
        the comparison.
        """
        super().__init__()
        if num_attention_sets < 1:
            raise ValueError(f"num_attention_sets must be >= 1, got {num_attention_sets}")
        self.num_attention_sets = num_attention_sets
        self.ln1 = RMSNorm(d_model, device=device, dtype=dtype)
        _make_attn = lambda: MultiheadSelfAttention(  # noqa: E731 -- one expression, used twice
            d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype
        )
        self.attn = (
            _make_attn() if num_attention_sets == 1
            else nn.ModuleList([_make_attn() for _ in range(num_attention_sets)])
        )
        self.ln2 = RMSNorm(d_model, device=device, dtype=dtype)
        self.router = Router(d_model, num_experts, num_active, width_ratio, device=device, dtype=dtype)

        self.num_experts = num_experts
        self.num_active = num_active

        # NOTE: divided by num_active (k), not by num_experts (E). This is what
        # keeps total expert parameters constant across the DVF series.
        d_ff_expert = d_ff // num_active
        w_std_base = math.sqrt(2 / (BASE_D_MODEL + BASE_D_FF))
        self._expert_init_std = w_std_base / math.sqrt(width_ratio)
        self.experts_w1 = self._init_expert_weights(num_experts, d_model, d_ff_expert, width_ratio, w_std_base, device, dtype)
        self.experts_w2 = self._init_expert_weights(num_experts, d_ff_expert, d_model, width_ratio, w_std_base, device, dtype)
        self.experts_w3 = self._init_expert_weights(num_experts, d_model, d_ff_expert, width_ratio, w_std_base, device, dtype)

    def attention_for(self, loop_step: Optional[int]):
        """The attention module this pass uses.

        `loop_step` is the explicit loop variable of
        `LoopedMoETransformer.forward` -- the index of the current pass over the
        shared stack -- threaded down as an argument. This block keeps **no
        cross-call state**: nothing here remembers which pass ran last, and two
        calls with the same `loop_step` select the same module. That is why the
        looped model can be traced, resumed and re-entered in any order.

        There is deliberately **no loop-step embedding** (owner's ruling
        2026-09-18): the passes differ only by their attention parameters, which
        is what keeps U1-U4 a controlled comparison against S1-S4.

        Refuses rather than defaults when `loop_step` is missing on a per-pass
        block: falling back to set 0 would run every pass through the first
        pass's attention, which trains, converges and reports nothing unusual
        while being a different architecture from the one on the config.
        """
        if self.num_attention_sets == 1:
            return self.attn
        assert self.num_attention_sets == len(self.attn), (
            f"{self.num_attention_sets} sets declared but {len(self.attn)} modules held; "
            "the count and the container have drifted apart"
        )
        if loop_step is None:
            raise ValueError(
                "this block has per-pass attention "
                f"({self.num_attention_sets} sets) and needs loop_step to choose one; "
                "got None. Every caller in this file passes it -- a new caller must too."
            )
        if not 0 <= loop_step < len(self.attn):
            raise IndexError(
                f"loop_step {loop_step} is outside the {self.num_attention_sets} attention "
                "sets this block holds. The block is built with one set per pass, so a "
                "loop_step beyond that means the model and the config disagree about R."
            )
        return self.attn[loop_step]

    def forward(
        self,
        x: Tensor,
        token_positions: Optional[Tensor] = None,
        *,
        loop_step: Optional[int] = None,
        layer_idx: Optional[int] = None,
        block: Optional[str] = None,
        unrolled_pos: Optional[int] = None,
        trace_sink: Optional[TraceSink] = None,
    ) -> tuple[Tensor, Tensor, Tensor]:
        batch, seq, dim = x.shape
        total_tokens = batch * seq

        norm1_out = self.ln1(x)
        attn_out = self.attention_for(loop_step)(norm1_out, token_positions)
        assert x.shape == attn_out.shape
        resid1_out = attn_out + x

        norm2_out = self.ln2(resid1_out)
        logits, probs, top_scores, top_experts = self.router(norm2_out)

        # `softmax` computes in float32 and does not cast back, so `top_scores`
        # is float32 regardless of the activation dtype. The combine weights get
        # multiplied into bf16 expert outputs below, so they must match, or
        # einsum raises "expected m1 and m2 to have the same dtype".
        #
        # This is invisible in the original: seedvar's published checkpoints are
        # float32, where the cast is a no-op. Under the bf16 training this
        # project uses (LT2 runs pure bf16, HANDOVER §4.1) it is a hard failure
        # on the first forward.
        #
        # Only the combine weights are cast. `probs` and `logits` stay float32
        # for the aux-loss and z-loss reductions, which is where the extra
        # precision is worth having.
        top_scores = top_scores.to(x.dtype)

        # Flatten and sort by expert so grouped_mm can run one matmul per expert.
        x_flat = rearrange(norm2_out, "b s d -> (b s) d")
        flat_expert_ids = rearrange(top_experts, "b s k -> (b s k)")
        flat_scores = rearrange(top_scores, "b s k -> (b s k)")
        flat_positions = torch.arange(total_tokens, device=x.device)
        flat_token_ids = repeat(flat_positions, "n -> (n k)", k=self.num_active)

        sort_indices = flat_expert_ids.argsort(stable=True)
        sorted_expert_ids = flat_expert_ids[sort_indices]
        sorted_token_ids = flat_token_ids[sort_indices]
        sorted_scores = flat_scores[sort_indices]
        sorted_x = x_flat[sorted_token_ids]

        counts = torch.bincount(sorted_expert_ids, minlength=self.num_experts)
        offs = counts.cumsum(0).to(torch.int32)

        h1 = grouped_mm(sorted_x, self.experts_w1, offs=offs)
        h3 = grouped_mm(sorted_x, self.experts_w3, offs=offs)
        gated = silu(h1) * h3
        expert_out = grouped_mm(gated, self.experts_w2, offs=offs)

        expert_out = einsum(expert_out, sorted_scores, "n d, n -> n d")
        output_flat = torch.zeros(total_tokens, dim, device=x.device, dtype=expert_out.dtype)
        output_flat.index_add_(0, sorted_token_ids, expert_out)
        experts_out = rearrange(output_flat, "(b s) d -> b s d", b=batch, s=seq)

        # Aux losses, per HANDOVER §3.3':
        #   L_LB = E * sum_i f_i * p_i   (switch-style load balancing)
        #   L_RZ = mean( (logsumexp logits)^2 )   (router z-loss)
        # Both are computed per layer per loop step; the caller averages over the
        # unrolled depth (num_stacks * num_layers_in_stack).
        fi = counts.float() / (total_tokens * self.num_active)
        pi = reduce(probs, "b s e -> e", "mean")
        lb = self.num_experts * einsum(fi, pi, "e, e ->")

        logsumexp = torch.logsumexp(logits.float(), dim=-1)
        lz = reduce(logsumexp**2, "... -> ", "mean")

        assert experts_out.shape == resid1_out.shape
        final_out = resid1_out + experts_out

        # Write into every sink that is armed. Under FSDP2 the `trace_sink`
        # keyword arrives as a per-block COPY (see `_ActiveTraceSink`), so the
        # module-level one is the only sink the trainer can actually read back;
        # the keyword remains for callers that pass their own dict directly
        # (the toy launcher, the offline probes), where it is the same object.
        # Writing to both is harmless: each is first-write-wins.
        sinks = [s for s in (ACTIVE_TRACE_SINK.sink, trace_sink) if s is not None]
        if sinks:
            if loop_step is None or layer_idx is None or block is None or unrolled_pos is None:
                raise ValueError(
                    "trace_sink was provided but loop_step/layer_idx/block/unrolled_pos "
                    "were not. Both axes must be explicit counters, never inferred."
                )
            if block not in ("head", "loop", "tail"):
                raise ValueError(f"block must be head/loop/tail, got {block!r}")
            # Keyed by the depth coordinate: head, the loop's first layer and tail all
            # carry loop_step=layer_idx=0, so the old pair-key silently collapsed them.
            key = unrolled_pos
            # First write wins. Under selective activation checkpointing the
            # block's forward runs a second time during backward to recompute
            # activations, so every key is legitimately visited twice -- that is
            # how AC works, not a bug. The recomputed values are identical by
            # construction, so keeping the first and ignoring the rest is both
            # correct and cheap.
            #
            # An earlier version raised on the second visit. That guard was aimed
            # at double-*counting*, which first-write-wins prevents directly; as
            # written it instead killed every AC-enabled run at the first traced
            # step. `test_ac_recomputation_does_not_disturb_the_trace` pins the
            # property that actually matters: same keys, same values, with AC on.
            # Detached views, not copies -- see the lifetime contract in
            # src/model/loop_trace.py.
            point = TracePoint(
                block=block,
                unrolled_pos=unrolled_pos,
                # From the tensors themselves, never from the model config: this layer's
                # E and k are what produced these numbers, and head/tail differ from the
                # loop block.
                num_experts=int(logits.shape[-1]),
                top_k=int(top_experts.shape[-1]),
                loop_step=loop_step,
                layer_idx=layer_idx,
                router_logits=logits.detach(),
                router_probs=probs.detach(),
                topk_idx=top_experts.detach(),
                topk_weights=top_scores.detach(),
                residual=final_out.detach(),
            )
            for sink in sinks:
                sink.setdefault(key, point)

        return final_out, lb, lz


class LoopedStack(nn.Module):
    """The stack of L MoE blocks that gets called R times."""

    def __init__(self, context_length: int, d_model: int, num_layers_in_stack: int, num_heads: int,
                 d_ff: int, rope_theta: float, width_ratio: float, num_experts: int,
                 num_active: int, device=None, dtype=None, num_attention_sets: int = 1):
        super().__init__()
        #: One attention set per pass when per-pass attention is on, 1 otherwise.
        #: Only the LOOPED layers get this: head and tail run once, so "per pass"
        #: has no meaning for them and they keep a single set in every variant.
        self.num_attention_sets = num_attention_sets
        self.layers = nn.ModuleList(
            [
                GroupedMoEPrenormBlock(
                    d_model, num_heads, d_ff, num_experts, num_active,
                    context_length, rope_theta, width_ratio, device, dtype,
                    num_attention_sets=num_attention_sets,
                )
                for _ in range(num_layers_in_stack)
            ]
        )

    def forward(
        self,
        x: Tensor,
        *,
        loop_step: int,
        unrolled_pos_start: int,
        trace_sink: Optional[TraceSink] = None,
    ) -> tuple[Tensor, Tensor, Tensor]:
        """`unrolled_pos_start` is the depth coordinate this call's first layer occupies.

        Passed in rather than recomputed from `loop_step`, so the caller owns the depth
        axis in one place: the stack does not need to know how many layers ran before it.
        """
        lb_total = x.new_zeros(())
        lz_total = x.new_zeros(())
        for layer_idx, layer in enumerate(self.layers):
            x, lb, lz = layer(
                x, loop_step=loop_step, layer_idx=layer_idx, block="loop",
                unrolled_pos=unrolled_pos_start + layer_idx, trace_sink=trace_sink,
            )
            lb_total = lb_total + lb
            lz_total = lz_total + lz
        return x, lb_total, lz_total


class LoopedMoETransformer(nn.Module):
    """Looped MoE transformer: one shared stack applied ``num_stacks`` times.

    The loop is an explicit Python ``for``; ``loop_step`` is the loop variable and
    is threaded all the way down to each block. Nothing downstream ever has to
    recover it from tensor shapes.
    """

    def __init__(self, vocab_size: int, context_length: int, d_model: int, num_layers_in_stack: int,
                 num_stacks: int, num_heads: int, d_ff: int, num_experts: int, num_active: int,
                 rope_theta: float, width_ratio: float, num_head_layers: int = 0,
                 num_tail_layers: int = 0, head_tail_num_experts: int = HEAD_TAIL_NUM_EXPERTS,
                 device=None, dtype=None, per_pass_attention: bool = False):
        super().__init__()
        self.num_stacks = num_stacks
        self.num_layers_in_stack = num_layers_in_stack
        self.total_layers = num_stacks * num_layers_in_stack
        self.num_head_layers = num_head_layers
        self.num_tail_layers = num_tail_layers
        # The unrolled depth, and the denominator the aux losses are averaged over.
        # Written as head + loop + tail rather than the recipe's shorthand "2 + D*L":
        # S15 has two head and two tail layers, so a literal 2 would divide by the wrong
        # number there -- and it would not fail, it would just make the auxiliary losses
        # quietly larger than intended.
        self.unrolled_depth = num_head_layers + self.total_layers + num_tail_layers

        self.token_embeddings = Embedding(vocab_size, d_model, device=device, dtype=dtype)
        # Head and tail are ordinary MoE layers that run once. They keep E=8/top-2
        # regardless of the loop block's expert count (recipe section 2.1: "identical in
        # every configuration"), so the loop block's E is deliberately not passed here.
        make_outer = lambda: GroupedMoEPrenormBlock(
            d_model, num_heads, d_ff, head_tail_num_experts, num_active,
            context_length, rope_theta, width_ratio, device, dtype,
        )
        self.head_layers = nn.ModuleList([make_outer() for _ in range(num_head_layers)])
        self.tail_layers = nn.ModuleList([make_outer() for _ in range(num_tail_layers)])
        self.per_pass_attention = per_pass_attention
        self.stack = LoopedStack(
            context_length, d_model, num_layers_in_stack, num_heads, d_ff, rope_theta,
            width_ratio, num_experts, num_active, device=device, dtype=dtype,
            # R sets when on: the stack is entered `num_stacks` times, and each
            # entry is what "a pass" means here.
            num_attention_sets=num_stacks if per_pass_attention else 1,
        )
        self.ln_final = RMSNorm(d_model, device=device, dtype=dtype)
        std_base_lm_head = math.sqrt(2 / (BASE_D_MODEL + vocab_size))
        self.lm_head = Linear(d_model, vocab_size, width_ratio, std_base_lm_head, device=device, dtype=dtype)

    @classmethod
    def from_config(cls, config: "LoopMoEConfig", *, device=None, dtype=None) -> "LoopedMoETransformer":
        """THE way to build this model from a config. Both call sites use it.

        There were two: the HF wrapper and `pretrain/train_spec.py`, each with its own
        hand-written keyword list. When head/tail layers were added, the training path's
        list was not updated, so it silently built a model with no head or tail while its
        config said otherwise -- it trained, the loss looked plausible, and every artifact
        recorded the config's depth rather than the depth that ran. Nothing could raise,
        because a shorter model is a perfectly valid model.

        A single entry point makes that class of drift impossible rather than merely
        tested-for: a field added to the config is read here once, and both paths get it.
        """
        return cls(
            vocab_size=config.vocab_size,
            context_length=config.context_length,
            d_model=config.d_model,
            num_layers_in_stack=config.num_layers_in_stack,
            num_stacks=config.num_stacks,
            num_heads=config.num_heads,
            d_ff=config.d_ff,
            num_experts=config.num_experts,
            num_active=config.num_active,
            rope_theta=config.rope_theta,
            width_ratio=config.width_ratio,
            num_head_layers=config.num_head_layers,
            num_tail_layers=config.num_tail_layers,
            # `getattr` with the config's own default, not a bare False: a config
            # object built by older code has no such attribute, and this is the
            # one entry point, so a silent False here would be the only place the
            # variant could be lost.
            per_pass_attention=getattr(config, "per_pass_attention", False),
            device=device,
            dtype=dtype,
        )

    def forward(
        self,
        x: Tensor,
        *,
        trace_sink: Optional[TraceSink] = None,
    ) -> tuple[Tensor, Tensor, Tensor]:
        lb_total = None
        lz_total = None

        x = self.token_embeddings(x)

        def run_outer(layers, block: str, pos: int, x, lb_total, lz_total):
            for i, layer in enumerate(layers):
                x, lb, lz = layer(
                    x, loop_step=0, layer_idx=0, block=block, unrolled_pos=pos + i,
                    trace_sink=trace_sink,
                )
                lb_total = lb if lb_total is None else lb_total + lb
                lz_total = lz if lz_total is None else lz_total + lz
            return x, lb_total, lz_total

        # head -> loop x num_stacks -> tail, with one running depth coordinate. head and
        # tail record loop_step=layer_idx=0 because neither coordinate means anything
        # outside the loop; `block` and `unrolled_pos` are what identifies them.
        x, lb_total, lz_total = run_outer(self.head_layers, "head", 0, x, lb_total, lz_total)
        pos = self.num_head_layers
        for loop_step in range(self.num_stacks):
            x, lb, lz = self.stack(
                x, loop_step=loop_step, unrolled_pos_start=pos, trace_sink=trace_sink,
            )
            pos += self.num_layers_in_stack
            lb_total = lb if lb_total is None else lb_total + lb
            lz_total = lz if lz_total is None else lz_total + lz
        x, lb_total, lz_total = run_outer(self.tail_layers, "tail", pos, x, lb_total, lz_total)

        x = self.lm_head(self.ln_final(x))

        # Averaged over the *unrolled* depth: every physical layer contributes once per
        # loop step, and head/tail contribute once each (recipe section 2.1). Equals
        # `total_layers` exactly when there are no head/tail layers, which is every
        # pre-ablation configuration -- so their published losses are unchanged.
        return x, lb_total / self.unrolled_depth, lz_total / self.unrolled_depth


def _require(kwargs: dict[str, Any], name: str) -> Any:
    """Fetch a required config field or raise.

    Project rule (HANDOVER §4.3 / §5.6): semantic parameters get no defaults.
    A default is a silent-wrong-answer generator -- a mistyped or dropped key
    becomes a plausible number instead of an error.
    """
    if name not in kwargs or kwargs[name] is None:
        raise ValueError(
            f"LoopMoEConfig: required field {name!r} is missing. Semantic "
            "parameters have no defaults in this project; state it explicitly."
        )
    return kwargs.pop(name)


class LoopMoEConfig(PretrainedConfig):
    """Config for the looped-MoE architecture. **Every field is required.**

    Compatible with ``save_pretrained``/``from_pretrained``: a config.json written
    by this class round-trips, and one is rejected loudly if a field is absent.
    """

    model_type = "loop-moe"

    # Tells transformers not to introspect defaults by constructing `cls()` with
    # no arguments -- which this class deliberately rejects. Without it,
    # `save_pretrained` fails inside `_get_generation_parameters`. This is the
    # supported escape hatch for configs whose fields are all required.
    has_no_defaults_at_init = True

    def __init__(self, **kwargs: Any):
        # `from_pretrained` on a *torch-saved* config, and some HF-internal paths,
        # construct with no arguments at all; only a fully-specified call is valid.
        self.vocab_size = _require(kwargs, "vocab_size")
        self.context_length = _require(kwargs, "context_length")
        self.d_model = _require(kwargs, "d_model")
        self.num_heads = _require(kwargs, "num_heads")
        self.d_ff = _require(kwargs, "d_ff")
        self.rope_theta = _require(kwargs, "rope_theta")
        self.width_ratio = _require(kwargs, "width_ratio")
        self.num_layers_in_stack = _require(kwargs, "num_layers_in_stack")  # L
        self.num_stacks = _require(kwargs, "num_stacks")  # R
        self.num_experts = _require(kwargs, "num_experts")  # E
        self.num_active = _require(kwargs, "num_active")  # k
        self.lb_loss_factor = _require(kwargs, "lb_loss_factor")
        # Head/tail layers: MoE layers run once, outside the loop (ablation recipe 2026-09-12
        # section 2.1). 0 means the architecture has none, which is not a guess -- it is what
        # every configuration built before this recipe actually is, and
        # `test_config_registry` pins their parameter counts as unchanged. Real ablation
        # configs never rely on the fallback: `config_registry.build_ablation` states both
        # counts for every entry, and a test asserts it does.
        self.num_head_layers = int(kwargs.pop("num_head_layers", 0))
        self.num_tail_layers = int(kwargs.pop("num_tail_layers", 0))

        # Per-pass attention: each of the R passes over the shared stack gets its
        # OWN attention parameters, while the MoE, the norms and the router stay
        # shared. Default False, and defaulted rather than required precisely
        # because every config written before 2026-09-18 lacks the key: those
        # files must keep loading, and they must keep meaning what they meant.
        # False reproduces the previous model bit for bit, including the state
        # dict's key names -- see GroupedMoEPrenormBlock.
        self.per_pass_attention = bool(kwargs.pop("per_pass_attention", False))
        self.lz_loss_factor = _require(kwargs, "lz_loss_factor")

        # Stated, not inherited. `PretrainedConfig` defaults this to True, and the
        # only reason the embedding and the LM head are not already sharing storage
        # is that this model never implemented `get_output_embeddings()`. The day
        # someone adds it for tool compatibility, every configuration would start
        # tying weights -- a different model, trained to a different loss, with
        # nothing in any artifact saying so. The architecture uses untied weights
        # (the parameter counts in the recipe assume it), so the config says so.
        kwargs.pop("tie_word_embeddings", None)
        self.tie_word_embeddings = False

        self._validate()

        # Derived, for readers; never an input. The unrolled depth now includes the
        # layers that run once outside the loop.
        self.num_layers = self.num_stacks * self.num_layers_in_stack
        self.unrolled_depth = self.num_head_layers + self.num_layers + self.num_tail_layers

        # The original config mirrored `context_length` into `max_length` for
        # lm-evaluation-harness. transformers >=5 classifies `max_length` as a
        # generation parameter and refuses to serialise a config carrying one
        # (the check is `hasattr`, so even a property trips it). `context_length`
        # is therefore the single source of truth for sequence length; pass
        # `max_length` to the harness explicitly at eval time instead.
        # Popped so that loading a seedvar-era config.json cannot reintroduce it.
        kwargs.pop("max_length", None)

        super().__init__(**kwargs)

    def _validate(self) -> None:
        """Reject out-of-domain values loudly rather than failing deep in a kernel."""
        positive = (
            "vocab_size", "context_length", "d_model", "num_heads", "d_ff",
            "num_layers_in_stack", "num_stacks", "num_experts", "num_active",
        )
        for name in positive:
            value = getattr(self, name)
            if not isinstance(value, int) or value < 1:
                raise ValueError(f"LoopMoEConfig.{name} must be a positive int, got {value!r}")
        if self.d_model % self.num_heads != 0:
            raise ValueError(
                f"d_model ({self.d_model}) must be divisible by num_heads ({self.num_heads})"
            )
        if self.num_active > self.num_experts:
            raise ValueError(
                f"num_active ({self.num_active}) cannot exceed num_experts ({self.num_experts})"
            )
        if self.d_ff % self.num_active != 0:
            raise ValueError(
                f"d_ff ({self.d_ff}) must be divisible by num_active ({self.num_active}); "
                "expert width is d_ff // num_active and truncation would silently "
                "change the parameter count."
            )
        for name in ("num_head_layers", "num_tail_layers"):
            value = getattr(self, name)
            if not isinstance(value, int) or value < 0:
                raise ValueError(f"LoopMoEConfig.{name} must be a non-negative int, got {value!r}")
        if (self.d_model // self.num_heads) % 2 != 0:
            raise ValueError(
                f"head dim ({self.d_model // self.num_heads}) must be even for RoPE"
            )

    @property
    def trace_meta(self) -> TraceMeta:
        """Shape/provenance block handed to the metrics collector."""
        return TraceMeta(
            num_stacks=self.num_stacks,
            num_layers_in_stack=self.num_layers_in_stack,
            num_experts=self.num_experts,
            num_active=self.num_active,
        )


class LoopMoEForCausalLM(PreTrainedModel, GenerationMixin):
    """HF-compatible causal LM wrapper.

    Kept HF-shaped on purpose: HANDOVER §4.7 makes "the diagnostics pipeline
    ingests our checkpoints unchanged" an acceptance criterion, and that pipeline
    loads models through ``from_pretrained``.
    """

    config_class = LoopMoEConfig

    def __init__(self, config: LoopMoEConfig):
        super().__init__(config)
        self.model = LoopedMoETransformer.from_config(config)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.token_embeddings

    def set_input_embeddings(self, value):
        self.model.token_embeddings = value

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[Tensor] = None,  # unused: the mask is built in
        labels: Optional[torch.LongTensor] = None,
        trace_sink: Optional[TraceSink] = None,
        **kwargs: Any,
    ) -> CausalLMOutputWithPast:
        """Forward pass.

        Returns a ``CausalLMOutputWithPast`` whose ``loss`` is the *total* loss
        (CE + weighted aux). The unweighted components are attached as
        ``task_loss`` / ``lb_loss`` / ``z_loss`` so the training loop can log the
        breakdown without recomputing anything.

        **Label contract (documented here 2026-08-21, `ABCI_ERR_20260821_0405_
        gate1_label_shift_root_cause.md`)**: ``labels`` must already be
        next-token-shifted by the caller -- ``labels[..., t] == input_ids[..., t+1]``,
        with the last position set to ``-100`` (no target exists after it). This
        method does **not** shift internally; it passes ``labels`` to
        ``F.cross_entropy`` exactly as given. Before this date the only written
        record of this contract was a comment in
        ``src/model/loop_trace.py`` ("the dataloader supplies pre-shifted
        labels during training, so the shift is explicit here") -- not here, at
        the definition itself. That gap let four separate call sites
        (the pretraining dataloader path aside, which was correct) independently
        get this wrong the same way, rather than it being four unrelated
        mistakes. Any caller not shifting first -- e.g. ``model(input_ids=ids,
        labels=ids)`` -- silently trains/evaluates on the trivial
        copy-the-current-token target instead of next-token prediction.
        """
        logits, lb, lz = self.model(input_ids, trace_sink=trace_sink)

        loss = task_loss = None
        if labels is not None:
            task_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
            loss = (
                task_loss
                + self.config.lb_loss_factor * lb
                + self.config.lz_loss_factor * lz
            )

        out = CausalLMOutputWithPast(loss=loss, logits=logits)
        # Unweighted components; the trainer pairs them with the factors from
        # config to build LossComponents.
        out.task_loss = task_loss
        out.lb_loss = lb
        out.z_loss = lz
        return out

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}