SYNKA 27938: verified step_00028000

step_00028000/config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "LoopMoEForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_loop_lm.LoopMoEConfig",
+    "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM"
+  },
+  "context_length": 4096,
+  "d_ff": 4608,
+  "d_model": 1792,
+  "lb_loss_factor": 0.01,
+  "lz_loss_factor": 0.001,
+  "model_type": "loop-moe",
+  "model_variant": "looped-moe",
+  "num_active": 2,
+  "num_experts": 16,
+  "num_head_layers": 1,
+  "num_heads": 28,
+  "num_layers": 16,
+  "num_layers_in_stack": 4,
+  "num_stacks": 4,
+  "num_tail_layers": 1,
+  "per_pass_attention": true,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.15.0",
+  "unrolled_depth": 18,
+  "vocab_size": 49152,
+  "width_ratio": 14.0
+}

step_00028000/loop_trace.py

@@ -0,0 +1,539 @@
+"""Per-loop-step trace contract between the model and the metrics layer.
+
+This module is the *interface* between `pretrain.loopmoe` (which produces raw
+tensors) and `pretrain.metrics` (which turns them into numbers). It deliberately
+imports **nothing** from torchtitan or from any trainer, so that:
+
+  * the metrics layer can be developed and tested without a trainer, and
+  * swapping the training backend (see HANDOVER §4.1' item 1) touches nothing here.
+
+Division of labour (agreed with infra-metrics, 2026-08-11):
+
+    the model / trainer  ->  *when* to collect and *how* to persist
+    the metrics layer    ->  *what* to collect and *how* to compute it
+
+Hence the model hands over **raw tensors, never reduced scalars**. A norm
+computed in two places is two definitions of that norm; the project has already
+paid for that mistake once (HANDOVER §4.4: "same-name metrics have several
+non-equivalent definitions"). The single definition lives in the metrics layer.
+
+Tensor lifetime contract (IMPORTANT)
+------------------------------------
+`TracePoint` tensors are ``detach()``-ed **views of live activations**, not
+copies. Passing them costs ~0 extra memory, which is why the model can afford to
+hand over all R x L residual tensors instead of subsampling layers. The price is
+a rule the consumer MUST obey:
+
+  1. The collector callback is **synchronous**. When it returns, the tensors are
+     considered dead.
+  2. The collector MUST NOT store a TracePoint (or any tensor inside one) in any
+     container that outlives the call -- no ``self.foo = point``, no appending to
+     a list, no closure capture.
+  3. Anything needed across steps must first be reduced to a Python scalar.
+
+Violating this pins whole activation graphs in memory and turns a constant-memory
+probe into a leak that grows with training length. `assert_trace_released` below
+turns that failure into a loud test failure instead of a slow OOM at step 40k.
+"""
+
+from __future__ import annotations
+
+import random
+import weakref
+from dataclasses import dataclass
+from typing import Any, Protocol, runtime_checkable
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+__all__ = [
+    "ACTIVE_TRACE_SINK",
+    "ROUTER_LOGITS_ARE_PRESOFTMAX",
+    "LOOP_AXIS_PROVENANCE",
+    "TracePoint",
+    "LossComponents",
+    "TraceMeta",
+    "TraceSink",
+    "LoopTraceCollector",
+    "NullCollector",
+    "run_probe_forward",
+    "assert_trace_released",
+]
+
+
+# --- semantic markers -------------------------------------------------------
+#
+# HANDOVER §4.3: the project has been burned by mislabelling sigmoid/logits,
+# which produced results that "looked entirely reasonable but were all wrong".
+# The collector schema has a `router_logits_are_presoftmax` field that pairs with
+# this constant; it is a constant rather than a runtime flag because the model
+# has exactly one behaviour and a flag that can only take one value is a lie
+# waiting to happen.
+
+ROUTER_LOGITS_ARE_PRESOFTMAX: bool = True
+"""`TracePoint.router_logits` is the raw router linear output (pre-softmax)."""
+
+LOOP_AXIS_PROVENANCE: str = "hook_call_index"
+"""How `TracePoint.loop_step` was determined.
+
+The value means the loop index was threaded down from the Python `for` loop that
+drives the stack -- a per-module call counter. It was **not** recovered by
+reshaping a flattened layer axis, which silently yields transposed semantics
+(HANDOVER §4.3).
+
+**The exact string matters.** It is not a description; it is a value the
+diagnostics pipeline validates. `looped_diag/collect/schema.py` accepts only::
+
+    allowed = ("hook_call_index", "hook_call_index+external")
+
+and raises otherwise (schema.py:398-401), with `probe_*.py` scripts asserting the
+same. An earlier draft of this constant read ``"explicit_call_counter"`` -- the
+same claim in different words, and it would have made every §4.7 ingestion of our
+traces raise. Do not "improve" the wording.
+"""
+
+
+@dataclass(frozen=True)
+class TracePoint:
+    """Raw per-(loop_step, layer) tensors. See the lifetime contract above.
+
+    All tensors are detached views; none require grad. Shapes use
+    B=batch, S=sequence, E=num_experts, k=num_active, d=d_model.
+    """
+
+    block: str
+    """Which part of the network produced this row: "head", "loop" or "tail".
+
+    Head and tail are single MoE layers run once, outside the loop (ablation recipe
+    2026-09-12 section 2.1). `loop_step` and `layer_idx` describe a position inside the
+    shared stack and are meaningless for them; they are recorded as 0/0 and must not be
+    read as coordinates unless `block == "loop"`.
+    """
+
+    unrolled_pos: int
+    """Position in the unrolled network, 0-based: THE depth coordinate.
+
+    head = 0, the loop's layers run 1 .. num_stacks*num_layers_in_stack in execution
+    order, tail = 1 + num_stacks*num_layers_in_stack (and correspondingly higher when a
+    configuration has more than one head/tail layer). This is the sink's key, so it is
+    also the only ordering that is guaranteed unique: sorting by (loop_step, layer_idx)
+    puts head, tail and the loop's first layer on top of each other.
+    """
+
+    loop_step: int
+    """Which pass through the shared stack, 0-based. Range: [0, num_stacks).
+
+    Only meaningful when `block == "loop"`; 0 for head and tail.
+    """
+
+    layer_idx: int
+    """Which physical layer inside the stack, 0-based. Range: [0, num_layers_in_stack).
+
+    Only meaningful when `block == "loop"`; 0 for head and tail.
+    """
+
+    num_experts: int
+    """E for THIS layer, taken from the router tensor's own shape.
+
+    Per row, not per model, because head/tail layers keep 8 experts whatever the loop
+    block uses (ablation recipe 2.1) -- S4 gives the loop block 64. A consumer that reads
+    the model-level `TraceMeta.num_experts` and applies it to a head row computes, for
+    example, L2 >= 1 - 8/64 = 0.875 and reads a perfectly healthy head layer as heavily
+    collapsed. Nothing raises: the number is in range and looks plausible, and head/tail
+    behaviour is exactly what the ablation is there to measure.
+    """
+
+    top_k: int
+    """k for THIS layer, from the selection tensor's own shape. Same reasoning as above."""
+
+    router_logits: Tensor
+    """[B, S, E] raw router output, pre-softmax. See ROUTER_LOGITS_ARE_PRESOFTMAX."""
+
+    router_probs: Tensor
+    """[B, S, E] softmax over **all** E experts (not renormalised over top-k)."""
+
+    topk_idx: Tensor
+    """[B, S, k] indices of the selected experts."""
+
+    topk_weights: Tensor
+    """[B, S, k] combine weights: top-k of `router_probs`, renormalised to sum to 1.
+
+    This is the quantity the MoE combine step already computes; exposing it adds
+    no arithmetic.
+    """
+
+    residual: Tensor
+    """[B, S, d] the block's output tensor, with **no reduction applied**.
+
+    Semantically this is the diagnostics repo's KIND_RESIDUAL / CAPTURE_BLOCK_OUTPUT:
+    the return value of the decoder block's forward. The metrics layer derives
+    both the within-step and the step-boundary norms from this one tensor, so
+    that both come from a single definition.
+    """
+
+
+@dataclass(frozen=True)
+class LossComponents:
+    """Per-step loss breakdown. Cheap enough to emit on *every* step.
+
+    HANDOVER §4.4 calls the aux losses the "blood-pressure monitor" and asks for
+    them every step. The *factors* are included alongside the values because
+    milestone M3.5 has to be able to reconstruct, from the logs alone, which
+    coefficient was in force at any step (HANDOVER §3.3'').
+    """
+
+    task_loss: float
+    """Cross-entropy on next-token prediction, before any aux term."""
+
+    lb_loss: float
+    """Switch-style load-balancing loss, E * sum_i f_i * p_i, averaged over all
+    unrolled layers (num_stacks * num_layers_in_stack)."""
+
+    lb_loss_factor: float
+    """Coefficient multiplying `lb_loss` in `total_loss`."""
+
+    z_loss: float
+    """Router z-loss, mean((logsumexp logits)^2), averaged the same way."""
+
+    z_loss_factor: float
+    """Coefficient multiplying `z_loss` in `total_loss`."""
+
+    total_loss: float
+    """task_loss + lb_loss_factor * lb_loss + z_loss_factor * z_loss."""
+
+
+@dataclass(frozen=True)
+class TraceMeta:
+    """Shape/provenance context so a collector never has to infer structure."""
+
+    num_stacks: int
+    """R: how many times the shared stack is called (loop count)."""
+
+    num_layers_in_stack: int
+    """L: physical layers per stack."""
+
+    num_experts: int
+    """E: experts per MoE layer."""
+
+    num_active: int
+    """k: experts activated per token."""
+
+    router_logits_are_presoftmax: bool = ROUTER_LOGITS_ARE_PRESOFTMAX
+    loop_axis_provenance: str = LOOP_AXIS_PROVENANCE
+
+
+# A sink is just the dict the model fills in during a traced forward. The
+# trainer creates one, passes it to the model, hands it to the collector, and
+# drops it -- so no hook registration/deregistration dance, and no state living
+# on the modules between steps.
+TraceSink = dict[int, TracePoint]
+"""Keyed by ``unrolled_pos`` -- explicit, never shape-derived.
+
+It was keyed by ``(loop_step, layer_idx)`` until the head/tail layers arrived (ablation
+recipe 2026-09-12). Those run outside the loop, so they have no meaningful value for
+either coordinate; recording them as 0/0 -- which is what they are -- would have put
+head, the loop's very first layer, and tail on the same key, and a dict assignment does
+not complain. Two of the three rows would have vanished with nothing in the output
+saying so. `unrolled_pos` is the depth coordinate Research defined, and keying on it
+makes the key unique by construction rather than by a sentinel convention that a later
+reader has to know about.
+"""
+
+
+class _ActiveTraceSink:
+    """The sink the blocks write into, reached as a module-level object.
+
+    Why this exists rather than passing the dict down as a keyword argument:
+    `fully_shard` is applied per block (`src/pretrain/train_spec.py`), and FSDP2
+    repacks kwargs across that boundary -- a dict passed by keyword arrives as a
+    fresh copy at every sharded block. The blocks were writing faithfully into
+    copies that were then discarded, so multi-GPU runs produced no router rows
+    at all while every intermediate layer looked correct. Diagnosed 2026-09-05
+    by printing `id()` on both sides: identical single-process, different at
+    every block under two ranks.
+
+    A module-level object does not cross that boundary, which is exactly why
+    `AUX_LOSS_STATE` never had the problem -- the aux losses travel up through
+    return values and are published to a global. This is the same pattern.
+
+    Process-local by construction: each rank has its own interpreter and so its
+    own instance, which is what makes the collected values rank-local. That is a
+    property to label in the output, not to hide -- see `aggregation_scope` in
+    the router rows.
+    """
+
+    def __init__(self) -> None:
+        self.sink: TraceSink | None = None
+
+    def arm(self) -> TraceSink:
+        """Start collecting; returns the dict that will be filled."""
+        self.sink = {}
+        return self.sink
+
+    def disarm(self) -> TraceSink | None:
+        """Stop collecting and hand back whatever was gathered."""
+        sink, self.sink = self.sink, None
+        return sink
+
+
+ACTIVE_TRACE_SINK = _ActiveTraceSink()
+
+
+@runtime_checkable
+class LoopTraceCollector(Protocol):
+    """What the metrics layer implements; what the training loop calls."""
+
+    def on_metrics_step(
+        self, step: int, loss: LossComponents, grad_norm: float | None = None
+    ) -> dict[str, float]:
+        """Called on **every** optimizer step. Returns flat scalars to log.
+
+        `grad_norm` is passed only on the steps where the training loop
+        already has it as a host float; on all other steps it is `None` and
+        the field is omitted from the row rather than guessed at.
+        """
+        ...
+
+    def on_trace_step(
+        self, step: int, sink: TraceSink, meta: TraceMeta
+    ) -> dict[str, float]:
+        """Called every N steps with the raw tensors.
+
+        MUST be synchronous and MUST NOT retain any tensor from `sink`
+        (see the lifetime contract at the top of this module).
+        """
+        ...
+
+    def on_probe_step(
+        self, step: int, sink: TraceSink, meta: TraceMeta, per_token_loss: Tensor
+    ) -> dict[str, float]:
+        """Called every M steps with a forward over the fixed probe corpus.
+
+        Deliberately a separate method rather than `on_trace_step` with a
+        `source="probe"` flag: probe data answers different questions
+        (cross-loop-step overlap, repetition-stratified loss) and lands in a
+        different file, and a flag is something a downstream analysis can forget
+        to filter on.
+
+        `sink` is structurally identical to `on_trace_step`'s. `per_token_loss`
+        is [B, S] and obeys the same no-retention contract.
+        """
+        ...
+
+
+class NullCollector:
+    """No-op collector, so training runs with metrics switched off."""
+
+    def on_metrics_step(
+        self, step: int, loss: LossComponents, grad_norm: float | None = None
+    ) -> dict[str, float]:
+        return {}
+
+    def on_trace_step(
+        self, step: int, sink: TraceSink, meta: TraceMeta
+    ) -> dict[str, float]:
+        return {}
+
+    def on_probe_step(
+        self, step: int, sink: TraceSink, meta: TraceMeta, per_token_loss: Tensor
+    ) -> dict[str, float]:
+        return {}
+
+
+def _snapshot_rng() -> dict[str, Any]:
+    """Capture every RNG stream a probe forward could disturb."""
+    state: dict[str, Any] = {
+        "python": random.getstate(),
+        "torch": torch.get_rng_state(),
+    }
+    try:
+        import numpy as np
+
+        state["numpy"] = np.random.get_state()
+    except ImportError:  # pragma: no cover - numpy is a hard dependency in practice
+        pass
+    if torch.cuda.is_available():
+        state["cuda"] = torch.cuda.get_rng_state_all()
+    return state
+
+
+def _restore_rng(state: dict[str, Any]) -> None:
+    random.setstate(state["python"])
+    torch.set_rng_state(state["torch"])
+    if "numpy" in state:
+        import numpy as np
+
+        np.random.set_state(state["numpy"])
+    if "cuda" in state:
+        torch.cuda.set_rng_state_all(state["cuda"])
+
+
+def run_probe_forward(
+    model: Any,
+    input_ids: Tensor,
+    *,
+    collector: LoopTraceCollector,
+    step: int,
+) -> dict[str, float]:
+    """Run the fixed probe corpus through the model and hand it to the collector.
+
+    Three isolation properties, each of which exists because violating it would
+    corrupt something silently:
+
+    1. **No gradients, eval mode.** LT2 disabled probing outright because
+       activation memory multiplies by the loop count (HANDOVER §3.4); at R=16
+       this is 4x the pressure they measured, so a probe that built a graph would
+       OOM at exactly the configurations we most want to measure.
+    2. **RNG is restored on exit** (python / numpy / torch / cuda). Without this,
+       whether a probe ran would change the training trajectory, and a resumed
+       run would silently diverge from an uninterrupted one -- breaking the §4.6
+       requirement that the two be statistically indistinguishable.
+    3. **The training dataloader is never touched.** The probe corpus is a fixed
+       tensor held separately, so probing does not advance the data position that
+       checkpoints record.
+
+    The model's training/eval mode is restored even if the collector raises.
+
+    Returns whatever scalars the collector produced.
+    """
+    was_training = model.training
+    rng_state = _snapshot_rng()
+    sink: TraceSink = {}
+    try:
+        model.eval()
+        with torch.no_grad():
+            out = model(input_ids, trace_sink=sink)
+            logits = out.logits if hasattr(out, "logits") else out
+
+            # Next-token targets. The model does not shift internally (the
+            # dataloader supplies pre-shifted labels during training), so the
+            # shift is explicit here. The final position has no next token and is
+            # masked out; cross_entropy returns 0.0 there.
+            labels = input_ids.new_full(input_ids.shape, -100)
+            labels[:, :-1] = input_ids[:, 1:]
+            per_token_loss = F.cross_entropy(
+                logits.transpose(1, 2).float(), labels,
+                reduction="none", ignore_index=-100,
+            )  # [B, S]
+
+            meta = model.config.trace_meta
+            return collector.on_probe_step(step, sink, meta, per_token_loss)
+    finally:
+        sink.clear()
+        _restore_rng(rng_state)
+        if was_training:
+            model.train()
+
+
+def assert_trace_released(sink: TraceSink) -> None:
+    """Raise if a collector retained tensors from `sink` past its callback.
+
+    Call this immediately after the collector returns and after clearing local
+    references. It weak-references every tensor, drops the sink, and checks that
+    nothing kept them alive.
+
+    This is the enforcement half of the lifetime contract. It is meant to be
+    wired into the toy/CI run rather than the production hot path.
+    """
+    refs: list[weakref.ref] = []
+    for point in sink.values():
+        for tensor in (
+            point.router_logits,
+            point.router_probs,
+            point.topk_idx,
+            point.topk_weights,
+            point.residual,
+        ):
+            try:
+                refs.append(weakref.ref(tensor))
+            except TypeError:  # pragma: no cover - torch tensors are weakref-able
+                continue
+    # After a `for` loop the loop variables stay bound in this frame, so `point`
+    # and `tensor` would still reference the *last* TracePoint when we collect --
+    # and this function would report itself as the leaker. Drop them explicitly.
+    point = tensor = None  # noqa: F841 - rebinding to release references
+    sink.clear()
+
+    import gc
+
+    gc.collect()
+    leaked = sum(1 for ref in refs if ref() is not None)
+    if leaked:
+        raise RuntimeError(
+            f"{leaked}/{len(refs)} trace tensors are still alive after the collector "
+            "returned. A collector must not store TracePoint tensors beyond the "
+            "callback; reduce to scalars first. See the lifetime contract in "
+            "src/model/loop_trace.py."
+        )
+
+def unrolled_position(loop_step: int, layer_idx: int, num_layers_in_stack: int) -> int:
+    """Depth coordinate of a loop-block layer: `loop_step * n_layers + layer_idx`.
+
+    THE definition of that arithmetic. It is needed in two places -- the legacy
+    checkpoint adapter in the cross-architecture probe, and the `loop_only`
+    branch of `cross_arch_v2_offline` that reads pre-head/tail artifacts -- and
+    the two must agree, because one writes the coordinate and the other reads it
+    back. Two copies of a formula that indexes into a dump do not fail loudly
+    when they drift; they silently address different cells.
+
+    Only meaningful for `block == "loop"`. Head and tail run once and are their
+    own physical layers, so their position comes from the network's order, not
+    from this arithmetic.
+    """
+    if num_layers_in_stack <= 0:
+        raise ValueError(
+            f"num_layers_in_stack must be positive, got {num_layers_in_stack}. "
+            "A zero or negative stack size collapses every pass onto the same "
+            "position, which reads as a valid coordinate."
+        )
+    return loop_step * num_layers_in_stack + layer_idx
+
+
+def legacy_trace_point_factory(num_layers_in_stack: int):
+    """A `TracePoint` stand-in accepting the PRE-head/tail keyword set.
+
+    Checkpoints trained before the head/tail change bundle their own
+    `modeling_loop_lm.py`, which constructs `TracePoint(loop_step=..., layer_idx=...,
+    <tensors>)`. Those four fields are now required, so such a checkpoint cannot
+    be loaded at all -- the failure is a `TypeError` inside the bundled code,
+    raised before any forward pass completes.
+
+    The four fields are filled rather than defaulted, and deliberately NOT given
+    defaults on `TracePoint` itself: `num_experts` being required per cell is
+    what stops a head layer's balanced 8 experts being divided by the loop
+    block's 64, which would put a floor of 0.875 under `L2` and read as severe
+    collapse in every configuration. That protection is worth keeping for new
+    checkpoints even though old ones need this adapter.
+
+    A pre-head/tail model is all loop block by construction, so `block` is
+    `"loop"` for every row and the depth coordinate is the unrolled arithmetic.
+    """
+
+    # Bound HERE, not looked up inside `make`. The adapter is installed by
+    # REPLACING the module-level `TracePoint` name, so a lookup at call time
+    # would find this factory instead of the class -- the adapter would call
+    # itself. Capturing the class before installation is what makes the
+    # substitution safe.
+    cls = TracePoint
+
+    def make(**kwargs: Any) -> "TracePoint":
+        logits = kwargs.get("router_logits")
+        topk_idx = kwargs.get("topk_idx")
+        if logits is None or topk_idx is None:
+            raise ValueError(
+                "legacy TracePoint adapter needs router_logits and topk_idx to "
+                f"infer num_experts and top_k; got keys {sorted(kwargs)}"
+            )
+        return cls(
+            block="loop",
+            unrolled_pos=unrolled_position(
+                int(kwargs["loop_step"]), int(kwargs["layer_idx"]), num_layers_in_stack
+            ),
+            # Read off the tensors, which are the only source that exists here.
+            num_experts=int(logits.shape[-1]),
+            top_k=int(topk_idx.shape[-1]),
+            **kwargs,
+        )
+
+    return make

step_00028000/manifest.json

@@ -0,0 +1,96 @@
+{
+  "activation_checkpoint_mode": "selective",
+  "arch": "UL2",
+  "autocast": {
+    "enabled": false
+  },
+  "code_diff_sha256": null,
+  "code_dirty": false,
+  "code_revision_source": "job_start_env",
+  "compute_dtype": "bfloat16",
+  "data_position": {
+    "doc_idx_in_shard": 227593,
+    "dp_rank": 0,
+    "dp_world_size": 8,
+    "epoch": 0,
+    "global_sample_idx": 2688000,
+    "offset_in_shard_bytes": 952172140,
+    "seed": 42,
+    "shard_id": "shard_002",
+    "shard_sha256": "1f87dc91a903b6b1dd6994d5187a63d091d37384ed2be60eacc475ad282b15b6"
+  },
+  "epochs_completed": null,
+  "git_commit": "bfb53c6ae536fcc8159ac351d3ababe00eb6b4c5",
+  "kind": "trajectory",
+  "learning_rate": null,
+  "lr_schedule": {
+    "decay_ratio": 0.1,
+    "decay_type": "sqrt",
+    "min_lr_factor": 0.05,
+    "peak_lr": 0.0003,
+    "warmup_steps": 1000
+  },
+  "max_seq_len": null,
+  "model_config": {
+    "_name_or_path": "",
+    "architectures": [
+      "LoopMoEForCausalLM"
+    ],
+    "auto_map": {
+      "AutoConfig": "modeling_loop_lm.LoopMoEConfig",
+      "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM"
+    },
+    "chunk_size_feed_forward": 0,
+    "context_length": 4096,
+    "d_ff": 4608,
+    "d_model": 1792,
+    "dtype": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "lb_loss_factor": 0.01,
+    "lz_loss_factor": 0.001,
+    "model_type": "loop-moe",
+    "model_variant": "looped-moe",
+    "num_active": 2,
+    "num_experts": 16,
+    "num_head_layers": 1,
+    "num_heads": 28,
+    "num_layers": 16,
+    "num_layers_in_stack": 4,
+    "num_stacks": 4,
+    "num_tail_layers": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "per_pass_attention": true,
+    "problem_type": null,
+    "return_dict": true,
+    "rope_theta": 10000.0,
+    "tie_word_embeddings": false,
+    "transformers_version": "5.15.0",
+    "unrolled_depth": 18,
+    "vocab_size": 49152,
+    "width_ratio": 14.0
+  },
+  "perturbation_probe": {
+    "runs_offline": true
+  },
+  "resume_checkpoint_interval_steps": 500,
+  "rng_state_saved": false,
+  "run_name": "UL2",
+  "seed": 42,
+  "step": 28000,
+  "steps_overridden": false,
+  "storage_dtype": "bfloat16",
+  "tokenizer_name": "smollm2",
+  "tokens_consumed": 11010048000,
+  "torch_version": "2.12.0+cu126",
+  "training_steps": 50000,
+  "transformers_version": "5.15.0"
+}

step_00028000/modeling_loop_lm.py

@@ -0,0 +1,936 @@
+"""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}

step_00028000/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f470809cd2e0f4b5009b468ae669932b2a1090214d2727ee0459f9e2e5a25958
+size 2826236439