diff --git a/trajectory/ABCI-2319212/step_00001000/config.json b/trajectory/ABCI-2319212/step_00001000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00001000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00001000/loop_trace.py b/trajectory/ABCI-2319212/step_00001000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00001000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00001000/manifest.json b/trajectory/ABCI-2319212/step_00001000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..2da5e3f4a54e0b5123b1df4d049420c69b99d841 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00001000/manifest.json @@ -0,0 +1,88 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 125737, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 96000, + "offset_in_shard_bytes": 501831084, + "seed": 42, + "shard_id": "shard_056", + "shard_sha256": "3194a22c5cd0073af081d9ee4a029454f0d1520adddb07e8fba14d510651731a" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": null, + "chunk_size_feed_forward": 0, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "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", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 1000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 393216000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00001000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00001000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00001000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00001000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00001000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..0fc20a933fe0fb8b65f98c1ae19dcb2ca3ef6771 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00001000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:326179e1248c477bd4bd16f41c1493a419126c8cfd17b0fc8f458d292ffb304e +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00002000/config.json b/trajectory/ABCI-2319212/step_00002000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00002000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00002000/loop_trace.py b/trajectory/ABCI-2319212/step_00002000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00002000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00002000/manifest.json b/trajectory/ABCI-2319212/step_00002000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..084c8f370ce6d570fbeb15e2dd37d0a70f9077f1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00002000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 593173, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 192000, + "offset_in_shard_bytes": 2445073916, + "seed": 42, + "shard_id": "shard_044", + "shard_sha256": "86d64197bf39f0d8c224184b1de89bdd7e39e2a0a7db65131b9614779638d318" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 2000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 786432000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00002000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00002000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00002000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00002000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00002000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..40120b746c57bd411cbdec7fb5e9a838589879fc --- /dev/null +++ b/trajectory/ABCI-2319212/step_00002000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c48235a035786d155ec5fda301bf1f8d12c93c967d91dc9d97c078ef694376b5 +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00004000/config.json b/trajectory/ABCI-2319212/step_00004000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00004000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00004000/loop_trace.py b/trajectory/ABCI-2319212/step_00004000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00004000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00004000/manifest.json b/trajectory/ABCI-2319212/step_00004000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..ff3abfe20b170a55eb5a2b535f42df4269dfb037 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00004000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 245202, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 384000, + "offset_in_shard_bytes": 1014364924, + "seed": 42, + "shard_id": "shard_051", + "shard_sha256": "3964fc6509eee4cf7b83ded9524728c85be173acbb21aefb3f79170c75e0e8ca" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 4000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 1572864000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00004000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00004000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00004000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00004000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00004000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..e297d008977fc8255412523699f931424929c68c --- /dev/null +++ b/trajectory/ABCI-2319212/step_00004000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c716577c5b8f3665a67bf15f227e33138952f71d922c13bc4698a705e7f7014c +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00008000/config.json b/trajectory/ABCI-2319212/step_00008000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00008000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00008000/loop_trace.py b/trajectory/ABCI-2319212/step_00008000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00008000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00008000/manifest.json b/trajectory/ABCI-2319212/step_00008000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..7b7a829872dfd3b559c0668eaa172b9811815abd --- /dev/null +++ b/trajectory/ABCI-2319212/step_00008000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 118247, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 768000, + "offset_in_shard_bytes": 500121520, + "seed": 42, + "shard_id": "shard_113", + "shard_sha256": "59e10fd01888d744676168ed9ff15059b7808ee3fe11d9558114903f40eaed22" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 8000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 3145728000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00008000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00008000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00008000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00008000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00008000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..d23e0719a5969186330f8b66ed30b80cb55abddb --- /dev/null +++ b/trajectory/ABCI-2319212/step_00008000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:95e6291e3b54aae63e1073db365f9df81240830501f11541dfa44a29d20a44aa +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00012000/config.json b/trajectory/ABCI-2319212/step_00012000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00012000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00012000/loop_trace.py b/trajectory/ABCI-2319212/step_00012000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00012000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00012000/manifest.json b/trajectory/ABCI-2319212/step_00012000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..36ee28d2b19f42f4feced2ac300d8c4bd1db0463 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00012000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 115741, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 1152000, + "offset_in_shard_bytes": 467544940, + "seed": 42, + "shard_id": "shard_063", + "shard_sha256": "e6c9f9e89e05bc157e19b9b46f3869cd968ef84b8f539e6621e91f984f2e28dd" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 12000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 4718592000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00012000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00012000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00012000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00012000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00012000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..5f19e84fd1609c7b7c556069bc1931cbe7aa4fb8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00012000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:47839f5178b62a76329fe6f74bafbfe352fcb66306b1452e5073f85a7923a356 +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00016000/config.json b/trajectory/ABCI-2319212/step_00016000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00016000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00016000/loop_trace.py b/trajectory/ABCI-2319212/step_00016000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00016000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00016000/manifest.json b/trajectory/ABCI-2319212/step_00016000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..6af95dbca10ce6e4d42515e99e7ad0892033b3e7 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00016000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 123528, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 1536000, + "offset_in_shard_bytes": 507795720, + "seed": 42, + "shard_id": "shard_108", + "shard_sha256": "e12798b18a2785830ef068c128ff3ff16fc681f1c4caf84591dea9dafd780729" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 16000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 6291456000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00016000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00016000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00016000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00016000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00016000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..3da8a73910c74e6de82e72f92328513dc97098d9 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00016000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e519255eb2bc595a6d37f85597a8265f94f88e574d77a26fcaa65740a6143d92 +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00020000/config.json b/trajectory/ABCI-2319212/step_00020000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00020000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00020000/loop_trace.py b/trajectory/ABCI-2319212/step_00020000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00020000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00020000/manifest.json b/trajectory/ABCI-2319212/step_00020000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..c741500b1290c5725f4aec545c68477bbb165002 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00020000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 11894, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 1920000, + "offset_in_shard_bytes": 49171444, + "seed": 42, + "shard_id": "shard_105", + "shard_sha256": "de19a9baf187d99a983a6af8c8b25b94deda47839acdf59376f129ca795c69af" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 20000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 7864320000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00020000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00020000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00020000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00020000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00020000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..35df3b770c5b035181400408395169df21daa5ec --- /dev/null +++ b/trajectory/ABCI-2319212/step_00020000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:93656112b5d242e2d783e98f54773b513eac0449b683d4bca377056b37cf34a9 +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00024000/config.json b/trajectory/ABCI-2319212/step_00024000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00024000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00024000/loop_trace.py b/trajectory/ABCI-2319212/step_00024000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00024000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00024000/manifest.json b/trajectory/ABCI-2319212/step_00024000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..471481cd794b5cf1736df0bde05d52832f84a913 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00024000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 459806, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 2304000, + "offset_in_shard_bytes": 1892836108, + "seed": 42, + "shard_id": "shard_079", + "shard_sha256": "752943f357db102e08c4866818ad373b846b1b78099c8cb5861e86f151c5f01e" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 24000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 9437184000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00024000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00024000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00024000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00024000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00024000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..52c86de1a82c89b7f7bac694428a7008ee97ffee --- /dev/null +++ b/trajectory/ABCI-2319212/step_00024000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f197264b12482bdb170fd0c350f7235836681d70063933b9646e74ef7e532a7a +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00028000/config.json b/trajectory/ABCI-2319212/step_00028000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00028000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00028000/loop_trace.py b/trajectory/ABCI-2319212/step_00028000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00028000/manifest.json b/trajectory/ABCI-2319212/step_00028000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..d8e56c16717439eee23eb5e66e3cbbfe5399fe23 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00028000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "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" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00028000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00028000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00028000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00028000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00028000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..1a6f94df8b534852a98206f6f0013a5e75b0e803 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00028000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:41ab0e4cfd9b69f0211e60581f90eb69074c6026d596d0e316c22367d994bb1c +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00032000/config.json b/trajectory/ABCI-2319212/step_00032000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00032000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00032000/loop_trace.py b/trajectory/ABCI-2319212/step_00032000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00032000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00032000/manifest.json b/trajectory/ABCI-2319212/step_00032000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..6db1f645c0568b30204140b688c47a9f8a53a117 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00032000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 719298, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 3072000, + "offset_in_shard_bytes": 3013920776, + "seed": 42, + "shard_id": "shard_102", + "shard_sha256": "cf3fc5873e86b47adb1884270fd8feb9c44540a8051c79e8f2ebeae03d80538e" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 32000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 12582912000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00032000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00032000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00032000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00032000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00032000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..402c1a7666e118d971133ac63d72cd3cd7a5e7dd --- /dev/null +++ b/trajectory/ABCI-2319212/step_00032000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2401f0ae81144f10e3ae80931cc4c6df235f08be8bf5cc05985d0b32c785e82e +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00036000/config.json b/trajectory/ABCI-2319212/step_00036000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00036000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00036000/loop_trace.py b/trajectory/ABCI-2319212/step_00036000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00036000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00036000/manifest.json b/trajectory/ABCI-2319212/step_00036000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..bd71dc87aa9148d1949f6b5c86cb2bbb559cedfa --- /dev/null +++ b/trajectory/ABCI-2319212/step_00036000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 235318, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 3456000, + "offset_in_shard_bytes": 987424636, + "seed": 42, + "shard_id": "shard_107", + "shard_sha256": "5625ab74f34fc22997fb607bcba48bc96e7372c010da0b424504d6654eab181a" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 36000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 14155776000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00036000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00036000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00036000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00036000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00036000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..9c1a8941b7bc05410321363e03da16827bb23fb3 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00036000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:257d5592bf34c12b877f3fc1e207553ea505cf3894f67d4306d5d862ad3ac100 +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00040000/config.json b/trajectory/ABCI-2319212/step_00040000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00040000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00040000/loop_trace.py b/trajectory/ABCI-2319212/step_00040000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00040000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00040000/manifest.json b/trajectory/ABCI-2319212/step_00040000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..53c12030c8ad8bfedf294e42c841c7cdb57230f4 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00040000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 579217, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 3840000, + "offset_in_shard_bytes": 2435256788, + "seed": 42, + "shard_id": "shard_086", + "shard_sha256": "550379161ea9e544f6cdbce0f5109066c5a88a38c0385c754ef40ec28edc259a" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 40000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 15728640000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00040000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00040000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00040000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00040000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00040000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..a8c44d1eb542655e6696e1f04ea0b3fa02d4b608 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00040000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fa197a8b9d77ef5d4233e9047a15c560bf673d54d0e05a07d8ab50bfb36b5d54 +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00044000/config.json b/trajectory/ABCI-2319212/step_00044000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00044000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00044000/loop_trace.py b/trajectory/ABCI-2319212/step_00044000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00044000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00044000/manifest.json b/trajectory/ABCI-2319212/step_00044000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..5426931cea170696e5ac86534d11f497798844d4 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00044000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 302579, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 4224000, + "offset_in_shard_bytes": 1241113108, + "seed": 42, + "shard_id": "shard_069", + "shard_sha256": "35b030d1f9ec67e0accebb86628a7bc22042f012f4e65a2c906e01d2d60f4e6d" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 44000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 17301504000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00044000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00044000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00044000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00044000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00044000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..69a46d07dd00a8703f8eaa285d4d7d5d8a0285df --- /dev/null +++ b/trajectory/ABCI-2319212/step_00044000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e7dee8a224e1a1993929f288740ea987cb2419616762e48c47755ced16cd7b2b +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00048000/config.json b/trajectory/ABCI-2319212/step_00048000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00048000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00048000/loop_trace.py b/trajectory/ABCI-2319212/step_00048000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00048000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00048000/manifest.json b/trajectory/ABCI-2319212/step_00048000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..fb24c2426b3421a286437168d1b5d8dd49a8d22d --- /dev/null +++ b/trajectory/ABCI-2319212/step_00048000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 434963, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 4608000, + "offset_in_shard_bytes": 1809124784, + "seed": 42, + "shard_id": "shard_019", + "shard_sha256": "7d2f1826c9e099298d88d59d69e992c7f4d9e7bf12dc0ee6e2a3867e3d65489c" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 48000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 18874368000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00048000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00048000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00048000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00048000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00048000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..8b4fd0be81dcc96da15aee98786f2000ddafd5e3 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00048000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07df55548e0a42f29458f7a2b9d85f3edb30db67ba9fac89d1df61122be0e8b0 +size 1749372279 diff --git a/trajectory/ABCI-2319212/step_00050000/config.json b/trajectory/ABCI-2319212/step_00050000/config.json new file mode 100644 index 0000000000000000000000000000000000000000..a5509f0ea4b71e643bac565503cf4df9d78bfaf2 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00050000/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "LoopMoEForCausalLM" + ], + "auto_map": { + "AutoConfig": "modeling_loop_lm.LoopMoEConfig", + "AutoModelForCausalLM": "modeling_loop_lm.LoopMoEForCausalLM" + }, + "context_length": 4096, + "d_ff": 3072, + "d_model": 1024, + "lb_loss_factor": 0.01, + "lz_loss_factor": 0.001, + "model_type": "loop-moe", + "model_variant": "looped-moe", + "num_active": 2, + "num_experts": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "rope_theta": 10000.0, + "tie_word_embeddings": false, + "transformers_version": "5.15.0", + "unrolled_depth": 18, + "vocab_size": 49152, + "width_ratio": 8.0 +} diff --git a/trajectory/ABCI-2319212/step_00050000/loop_trace.py b/trajectory/ABCI-2319212/step_00050000/loop_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..19be60aeaa2f2abd5402d0a8a684d9555a2e22f8 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00050000/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=..., + )`. 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 diff --git a/trajectory/ABCI-2319212/step_00050000/manifest.json b/trajectory/ABCI-2319212/step_00050000/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..bc1a436a84717d2bedc309a99e09c1154f7e5cfc --- /dev/null +++ b/trajectory/ABCI-2319212/step_00050000/manifest.json @@ -0,0 +1,95 @@ +{ + "activation_checkpoint_mode": "none", + "arch": "S14", + "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": 577881, + "dp_rank": 0, + "dp_world_size": 8, + "epoch": 0, + "global_sample_idx": 4800000, + "offset_in_shard_bytes": 2479948288, + "seed": 42, + "shard_id": "shard_094", + "shard_sha256": "672249fc72d8c6ddca86620e0a32962461ae0da4f0e5652b21405a2642611f92" + }, + "epochs_completed": null, + "git_commit": "036829168106b6215871f9bb7e873cba970e9bbd", + "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": 3072, + "d_model": 1024, + "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": 8, + "num_head_layers": 1, + "num_heads": 16, + "num_layers": 16, + "num_layers_in_stack": 16, + "num_stacks": 1, + "num_tail_layers": 1, + "output_attentions": false, + "output_hidden_states": false, + "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": 8.0 + }, + "perturbation_probe": { + "runs_offline": true + }, + "resume_checkpoint_interval_steps": 500, + "rng_state_saved": false, + "run_name": "S14", + "seed": 42, + "step": 50000, + "steps_overridden": false, + "storage_dtype": "bfloat16", + "tokenizer_name": "smollm2", + "tokens_consumed": 19660800000, + "torch_version": "2.12.0+cu126", + "training_steps": 50000, + "transformers_version": "5.15.0" +} \ No newline at end of file diff --git a/trajectory/ABCI-2319212/step_00050000/modeling_loop_lm.py b/trajectory/ABCI-2319212/step_00050000/modeling_loop_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..d55fa5527688fdd53eb4a05b142e7d076bfb90c1 --- /dev/null +++ b/trajectory/ABCI-2319212/step_00050000/modeling_loop_lm.py @@ -0,0 +1,849 @@ +"""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//`, 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): + super().__init__() + self.ln1 = RMSNorm(d_model, device=device, dtype=dtype) + self.attn = MultiheadSelfAttention(d_model, num_heads, max_seq_len, theta, width_ratio, device, dtype) + 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 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.attn(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): + super().__init__() + self.layers = nn.ModuleList( + [ + GroupedMoEPrenormBlock( + d_model, num_heads, d_ff, num_experts, num_active, + context_length, rope_theta, width_ratio, device, dtype, + ) + 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): + 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.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, + ) + 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, + 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)) + 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} diff --git a/trajectory/ABCI-2319212/step_00050000/pytorch_model.bin b/trajectory/ABCI-2319212/step_00050000/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..68764ef2b03a7048d66ca0dee35e32019479b69e --- /dev/null +++ b/trajectory/ABCI-2319212/step_00050000/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0546a408dcee2c2d2f1f455e75b5caa4a07d2f4088bb5fdba0d23e958d659fde +size 1749372279