Project-Chronos 0.1.1__py3-none-any.whl

chronos/trainer/dpo_trainer.py

@@ -0,0 +1,190 @@
+"""
+chronos/trainer/dpo_trainer.py — Stage 3 DPO trainer.
+
+Implements DPO with a frozen reference Chronos model + Chronos loss mixin
+(lookahead/temporal/balance preserved) + strong router KL anchor (default
+λ=0.1) so the alignment gradients don't blow up the cluster cache layout.
+
+Mirrors minimind's train_dpo.py loss but uses the Chronos forward signature
+``(outputs, lookahead_probs) = model(x)`` and tolerates its extra return.
+"""
+from __future__ import annotations
+
+import os
+import time
+from contextlib import nullcontext
+from copy import deepcopy
+
+import torch
+import torch.nn.functional as F
+from torch import optim
+from torch.utils.data import DataLoader
+
+import chronos.deps  # noqa
+from trainer.trainer_utils import Logger, is_main_process  # type: ignore
+from chronos.data.flexible_dataset import StreamingDPODataset as DPODataset
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.temporal_loss import lookahead_supervision_loss
+from chronos.trainer.optim_utils import get_lr
+from chronos.trainer.loss_mixin import (
+    chronos_loss_term,
+    collect_router_probs,
+    router_kl_anchor,
+    capture_reference_routing,
+)
+
+
+def _logits_to_log_probs(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+    log_probs = F.log_softmax(logits, dim=-1)
+    return log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)
+
+
+def dpo_loss(ref_lp: torch.Tensor, pol_lp: torch.Tensor, mask: torch.Tensor, beta: float) -> torch.Tensor:
+    ref_lp = (ref_lp * mask).sum(dim=1)
+    pol_lp = (pol_lp * mask).sum(dim=1)
+    B = ref_lp.shape[0]
+    ch_ref, rj_ref = ref_lp[:B // 2], ref_lp[B // 2:]
+    ch_pol, rj_pol = pol_lp[:B // 2], pol_lp[B // 2:]
+    pi_logr = ch_pol - rj_pol
+    rf_logr = ch_ref - rj_ref
+    return -F.logsigmoid(beta * (pi_logr - rf_logr)).mean()
+
+
+class ChronosDPOTrainer:
+    def __init__(self, model: ChronosForCausalLM, config: ChronosConfig, args, tokenizer):
+        self.model = model
+        self.config = config
+        self.args = args
+        self.tokenizer = tokenizer
+        self.device = args.device
+        self.beta = float(getattr(args, "beta", 0.1))
+
+        device_type = "cuda" if "cuda" in str(self.device) else "cpu"
+        dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        self.autocast_ctx = nullcontext() if device_type == "cpu" \
+            else torch.cuda.amp.autocast(dtype=dtype)
+        self.scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == "float16"))
+        from chronos.trainer.optim_utils import build_optimizer
+        self.optimizer = build_optimizer(
+            model, lr=args.learning_rate,
+            weight_decay=float(getattr(args, "weight_decay", 0.01)),
+        )
+
+        # Frozen reference model — DPO needs π_ref(y|x). Build a fresh
+        # instance and load the current state_dict (deepcopy fails on
+        # autograd-tracked buffers in some torch versions).
+        self.ref_model = model.__class__(config).to(self.device)
+        self.ref_model.load_state_dict(model.state_dict(), strict=False)
+        self.ref_model.eval().requires_grad_(False)
+
+        self._router_ref = None
+
+    def set_calibration_batch(self, x_calib: torch.Tensor):
+        self._router_ref = capture_reference_routing(self.model, x_calib, self.device)
+
+    def train_step(self, batch, step, total_steps):
+        x_c = batch["x_chosen"].to(self.device)
+        x_r = batch["x_rejected"].to(self.device)
+        y_c = batch["y_chosen"].to(self.device)
+        y_r = batch["y_rejected"].to(self.device)
+        m_c = batch["mask_chosen"].to(self.device).float()
+        m_r = batch["mask_rejected"].to(self.device).float()
+        x = torch.cat([x_c, x_r], dim=0)
+        y = torch.cat([y_c, y_r], dim=0)
+        m = torch.cat([m_c, m_r], dim=0)
+
+        lr = get_lr(step, total_steps, self.args.learning_rate)
+        for pg in self.optimizer.param_groups:
+            pg["lr"] = lr
+
+        with self.autocast_ctx:
+            with torch.no_grad():
+                ref_out = self.ref_model(x)
+                ref_logits = ref_out[0].logits if isinstance(ref_out, tuple) else ref_out.logits
+            ref_lp = _logits_to_log_probs(ref_logits, y)
+
+            outputs, lookahead_probs = self.model(x)
+            pol_lp = _logits_to_log_probs(outputs.logits, y)
+            base = dpo_loss(ref_lp, pol_lp, m, beta=self.beta)
+
+            # Mix in Chronos regularizers
+            loss = chronos_loss_term(
+                self.model, base, lookahead_probs, self.config,
+                aux_loss=outputs.aux_loss,
+            )
+
+            # Router KL anchor (DPO default: strong)
+            anc_val = 0.0
+            lam = float(getattr(self.config, "lambda_router_anchor", 0.0))
+            if lam > 0.0 and self._router_ref is not None:
+                router_4d = collect_router_probs(self.model)
+                if router_4d is not None and router_4d.shape[1] > 1:
+                    cur = router_4d.mean(dim=2)
+                    ref = self._router_ref
+                    T = min(cur.shape[1], ref.shape[1])
+                    B = min(cur.shape[0], ref.shape[0])
+                    anc = router_kl_anchor(cur[:B, :T, :], ref[:B, :T, :], lam)
+                    loss = loss + anc
+                    anc_val = float(anc.item())
+
+            la_val = 0.0
+            r4 = collect_router_probs(self.model)
+            if (r4 is not None and r4.shape[1] > 1
+                    and lookahead_probs is not None
+                    and self.config.lookahead_steps > 0):
+                la_val = float(lookahead_supervision_loss(
+                    lookahead_probs, r4.mean(dim=2).detach(),
+                    self.config.lookahead_steps,
+                ).item())
+
+            loss = loss / self.args.accumulation_steps
+
+        self.scaler.scale(loss).backward()
+
+        if step % self.args.accumulation_steps == 0:
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.grad_clip)
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            self.optimizer.zero_grad(set_to_none=True)
+
+        return (
+            loss.item() * self.args.accumulation_steps,
+            float(base.item()),
+            float(outputs.aux_loss.item()) if outputs.aux_loss is not None else 0.0,
+            la_val,
+            anc_val,
+        )
+
+    def train_epoch(self, epoch, loader, iters, max_steps=None):
+        self.model.train()
+        start_time = time.time()
+        total_steps = self.args.epochs * iters
+        for step, batch in enumerate(loader, start=1):
+            if max_steps is not None and step > max_steps:
+                break
+            loss, dpo_v, aux, la, anc = self.train_step(
+                batch, epoch * iters + step, total_steps,
+            )
+            if step % self.args.log_interval == 0 or step == iters:
+                lr = self.optimizer.param_groups[-1]["lr"]
+                Logger(
+                    f"[DPO] Epoch[{epoch+1}/{self.args.epochs}]({step}/{iters}) "
+                    f"loss:{loss:.4f} dpo:{dpo_v:.4f} aux:{aux:.4f} "
+                    f"la:{la:.4f} anchor:{anc:.4f} lr:{lr:.2e}"
+                )
+
+    def _save(self, epoch, step):
+        os.makedirs(self.args.save_dir, exist_ok=True)
+        ckp = os.path.join(self.args.save_dir, f"dpo_{self.config.hidden_size}_moe.pth")
+        self.model.eval()
+        state = self.model.state_dict()
+        torch.save({k: v.half().cpu() for k, v in state.items()}, ckp)
+        self.model.train()
+
+
+def build_dpo_loader(data_path, tokenizer, max_seq_len, batch_size):
+    ds = DPODataset(data_path, tokenizer, max_length=max_seq_len)
+    return DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=0)

chronos/trainer/grpo_trainer.py

@@ -0,0 +1,288 @@
+"""
+chronos/trainer/grpo_trainer.py — Stage 5 GRPO trainer (self-contained).
+
+Group Relative Policy Optimization for Chronos. Unlike minimind's
+train_grpo.py which depends on a rollout engine + external reward model
++ sglang, this trainer is self-contained:
+
+  - Rollouts use ``model.generate`` (HF-style) with sampling temperature.
+  - Reward is a toy function (length + vocab purity + repetition penalty)
+    pluggable via ``reward_fn``.
+  - Reference model is a frozen deep-copy of the starting model (same
+    idea as our DPO trainer).
+
+Loss:
+    For each prompt p, sample G completions {y_i}.
+    advantages   = (r_i - mean(r_group)) / (std(r_group) + eps)
+    per_token_lp = log π_θ(y_i | p)                         # current policy
+    per_token_lp_ref = log π_ref(y_i | p)                   # stop-grad
+    kl           = π_ratio^{-1}·exp(lp_ref - lp) - (lp_ref - lp) - 1   (k3)
+    L = -E[ advantage · lp_π ] + β · E[kl]
+"""
+from __future__ import annotations
+
+import os
+import time
+from contextlib import nullcontext
+from copy import deepcopy
+from typing import Callable, List, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import optim
+from torch.utils.data import DataLoader
+
+import chronos.deps  # noqa
+from trainer.trainer_utils import Logger  # type: ignore
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.temporal_loss import lookahead_supervision_loss
+from chronos.trainer.optim_utils import get_lr
+from chronos.trainer.loss_mixin import (
+    chronos_loss_term,
+    collect_router_probs,
+    router_kl_anchor,
+    capture_reference_routing,
+)
+
+
+from chronos.trainer.reward import ToyReward, build_reward_fn  # noqa: F401
+
+
+# Toy vocabulary mirrors tests/fixtures/generate_tiny_alignment_data.py
+TOY_VOCAB = {
+    "the", "quick", "brown", "fox", "machine", "learning", "model", "expert",
+    "router", "cluster", "memory", "cache", "prefetch", "latency", "throughput",
+    "tensor", "training", "inference", "decode", "weight", "gradient",
+    "softmax", "attention", "layer", "embedding", "vocab", "sequence",
+    "loss", "optimizer", "epoch", "batch", "step", "token", "logit",
+    "minimind", "chronos", "hybrid", "mixture", "of", "experts", "sparse",
+    "dense", "neural", "network", "transformer", "rope", "rotary", "kv",
+}
+
+
+def default_reward(prompt: str, response_tokens: List[int], tokenizer, target_len: int = 24) -> float:
+    """Backward-compat wrapper; prefer chronos.trainer.reward.ToyReward."""
+    return ToyReward(target_len=target_len, vocab=TOY_VOCAB)(prompt, response_tokens, tokenizer)
+
+
+class ChronosGRPOTrainer:
+    def __init__(
+        self,
+        model: ChronosForCausalLM,
+        config: ChronosConfig,
+        args,
+        tokenizer,
+        reward_fn: Callable = None,
+    ):
+        self.model = model
+        self.config = config
+        self.args = args
+        self.tokenizer = tokenizer
+        self.device = args.device
+        self.beta = float(getattr(args, "beta", 0.04))
+        self.num_generations = int(getattr(args, "num_generations", 4))
+        self.max_gen_len = int(getattr(args, "max_gen_len", 24))
+        self.temperature = float(getattr(args, "temperature", 1.0))
+        self.reward_fn = reward_fn or default_reward
+
+        device_type = "cuda" if "cuda" in str(self.device) else "cpu"
+        dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        self.autocast_ctx = nullcontext() if device_type == "cpu" \
+            else torch.cuda.amp.autocast(dtype=dtype)
+        self.scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == "float16"))
+        from chronos.trainer.optim_utils import build_optimizer
+        self.optimizer = build_optimizer(
+            model, lr=args.learning_rate,
+            weight_decay=float(getattr(args, "weight_decay", 0.01)),
+        )
+
+        self.ref_model = model.__class__(config).to(self.device)
+        self.ref_model.load_state_dict(model.state_dict(), strict=False)
+        self.ref_model.eval().requires_grad_(False)
+        self._router_ref = None
+
+    def set_calibration_batch(self, x: torch.Tensor):
+        self._router_ref = capture_reference_routing(self.model, x, self.device)
+
+    # ── rollout ──────────────────────────────────────────────────
+
+    def _rollout(self, prompt_ids: torch.Tensor) -> Tuple[torch.Tensor, List[int]]:
+        """Sample a single completion. Returns (full_ids [1, P+L], response_len)."""
+        self.model.eval()
+        with torch.no_grad():
+            ids = prompt_ids.clone().to(self.device)
+            past = None
+            generated: List[int] = []
+            for _ in range(self.max_gen_len):
+                past_len = past[0][0].shape[1] if past else 0
+                out = self.model(
+                    ids[:, past_len:] if past_len > 0 else ids,
+                    past_key_values=past, use_cache=True,
+                )
+                outputs = out[0] if isinstance(out, tuple) else out
+                past = outputs.past_key_values
+                logits = outputs.logits[:, -1, :] / max(self.temperature, 1e-6)
+                probs = F.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+                ids = torch.cat([ids, next_id], dim=-1)
+                generated.append(int(next_id.item()))
+                if next_id.item() == self.tokenizer.eos_token_id:
+                    break
+        self.model.train()
+        return ids, generated
+
+    # ── core step ────────────────────────────────────────────────
+
+    def train_step(self, prompt_text: str, step: int, total_steps: int):
+        # Tokenize prompt
+        prompt = self.tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt_text}],
+            tokenize=False, add_generation_prompt=True,
+        )
+        prompt_ids = self.tokenizer(
+            prompt, return_tensors="pt", truncation=True, max_length=self.args.max_seq_len,
+        ).input_ids.to(self.device)
+
+        # Generate G completions, compute rewards
+        trajectories: List[Tuple[torch.Tensor, int, float]] = []
+        for _ in range(self.num_generations):
+            full_ids, gen_tokens = self._rollout(prompt_ids)
+            r = self.reward_fn(prompt_text, gen_tokens, self.tokenizer)
+            trajectories.append((full_ids, len(gen_tokens), r))
+
+        rewards = torch.tensor([t[2] for t in trajectories], device=self.device)
+        adv = (rewards - rewards.mean()) / (rewards.std().clamp_min(1e-6))
+
+        # Build a padded batch of full sequences for the policy update
+        max_len = max(t[0].shape[1] for t in trajectories)
+        pad_id = self.tokenizer.pad_token_id or 0
+        ids_batch = torch.full(
+            (len(trajectories), max_len), pad_id, dtype=torch.long, device=self.device,
+        )
+        response_mask = torch.zeros(
+            (len(trajectories), max_len), dtype=torch.float, device=self.device,
+        )
+        prompt_len = prompt_ids.shape[1]
+        for i, (full, glen, _) in enumerate(trajectories):
+            L = full.shape[1]
+            ids_batch[i, :L] = full[0]
+            response_mask[i, prompt_len:L] = 1.0
+
+        lr = get_lr(step, total_steps, self.args.learning_rate)
+        for pg in self.optimizer.param_groups:
+            pg["lr"] = lr
+
+        with self.autocast_ctx:
+            outputs, lookahead_probs = self.model(ids_batch)
+            log_probs = F.log_softmax(outputs.logits[:, :-1, :], dim=-1)
+            target = ids_batch[:, 1:]
+            lp = log_probs.gather(-1, target.unsqueeze(-1)).squeeze(-1)  # [G, L-1]
+            mask = response_mask[:, 1:]
+
+            with torch.no_grad():
+                ref_out = self.ref_model(ids_batch)
+                ref_outputs = ref_out[0] if isinstance(ref_out, tuple) else ref_out
+                ref_lp = F.log_softmax(ref_outputs.logits[:, :-1, :], dim=-1).gather(
+                    -1, target.unsqueeze(-1)
+                ).squeeze(-1)
+
+            # k3 KL estimator (Schulman 2020)
+            log_ratio = ref_lp - lp
+            kl_per_tok = (log_ratio.exp() - log_ratio - 1.0) * mask
+            kl = kl_per_tok.sum() / mask.sum().clamp_min(1.0)
+
+            # Advantage-weighted policy gradient (per-sequence)
+            seq_lp = (lp * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0)
+            pg_loss = -(adv.detach() * seq_lp).mean()
+            base = pg_loss + self.beta * kl
+
+            loss = chronos_loss_term(
+                self.model, base, lookahead_probs, self.config,
+                aux_loss=outputs.aux_loss,
+            )
+
+            anc_val = 0.0
+            lam = float(getattr(self.config, "lambda_router_anchor", 0.0))
+            if lam > 0.0 and self._router_ref is not None:
+                r4 = collect_router_probs(self.model)
+                if r4 is not None and r4.shape[1] > 1:
+                    cur = r4.mean(dim=2)
+                    ref = self._router_ref
+                    T = min(cur.shape[1], ref.shape[1])
+                    Bm = min(cur.shape[0], ref.shape[0])
+                    anc = router_kl_anchor(cur[:Bm, :T, :], ref[:Bm, :T, :], lam)
+                    loss = loss + anc
+                    anc_val = float(anc.item())
+
+            la_val = 0.0
+            r4 = collect_router_probs(self.model)
+            if (r4 is not None and r4.shape[1] > 1
+                    and lookahead_probs is not None
+                    and self.config.lookahead_steps > 0):
+                la_val = float(lookahead_supervision_loss(
+                    lookahead_probs, r4.mean(dim=2).detach(),
+                    self.config.lookahead_steps,
+                ).item())
+
+            loss = loss / self.args.accumulation_steps
+
+        self.scaler.scale(loss).backward()
+        if step % self.args.accumulation_steps == 0:
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.grad_clip)
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            self.optimizer.zero_grad(set_to_none=True)
+
+        return {
+            "loss": float(loss.item() * self.args.accumulation_steps),
+            "pg_loss": float(pg_loss.item()),
+            "kl": float(kl.item()),
+            "mean_reward": float(rewards.mean().item()),
+            "la": la_val,
+            "anchor": anc_val,
+        }
+
+    def train_epoch(self, epoch, prompts: List[str], iters, max_steps=None):
+        self.model.train()
+        total_steps = self.args.epochs * iters
+        for step, prompt in enumerate(prompts, start=1):
+            if max_steps is not None and step > max_steps:
+                break
+            if step > iters:
+                break
+            stats = self.train_step(prompt, epoch * iters + step, total_steps)
+            if step % self.args.log_interval == 0 or step == iters:
+                lr = self.optimizer.param_groups[-1]["lr"]
+                Logger(
+                    f"[GRPO] Epoch[{epoch+1}/{self.args.epochs}]({step}/{iters}) "
+                    f"loss:{stats['loss']:.4f} pg:{stats['pg_loss']:.4f} "
+                    f"kl:{stats['kl']:.4f} r:{stats['mean_reward']:.3f} "
+                    f"la:{stats['la']:.4f} anchor:{stats['anchor']:.4f} lr:{lr:.2e}"
+                )
+
+    def _save(self, epoch, step):
+        os.makedirs(self.args.save_dir, exist_ok=True)
+        ckp = os.path.join(self.args.save_dir, f"grpo_{self.config.hidden_size}_moe.pth")
+        self.model.eval()
+        state = self.model.state_dict()
+        torch.save({k: v.half().cpu() for k, v in state.items()}, ckp)
+        self.model.train()
+
+
+def load_grpo_prompts(data_path: str, max_prompts: int = None) -> List[str]:
+    import json
+    prompts = []
+    with open(data_path) as f:
+        for line in f:
+            row = json.loads(line)
+            # Extract first user message
+            for msg in row.get("conversations", []):
+                if msg.get("role") == "user":
+                    prompts.append(msg["content"])
+                    break
+            if max_prompts and len(prompts) >= max_prompts:
+                break
+    return prompts

chronos/trainer/loss_mixin.py

@@ -0,0 +1,165 @@
+"""
+chronos/trainer/loss_mixin.py
+
+Reusable loss-mixing helpers shared across all Chronos training stages
+(pretrain, SFT, DPO, ORPO, GRPO). Centralizing this keeps each stage
+trainer thin and guarantees lookahead supervision is preserved through
+the whole 6-stage chain rather than being silently dropped by a stage
+that forgot to add it.
+
+Two helpers are exposed:
+
+- ``chronos_loss_term(model, base_loss, lookahead_probs, config)`` —
+  appends λ_balance·L_aux + λ_temporal·L_temporal + λ_lookahead·L_lookahead
+  to any base loss (CE for pretrain/SFT, DPO loss, ORPO loss, etc.).
+
+- ``router_kl_anchor(current_router_probs, ref_router_probs, lambda_anchor)`` —
+  KL(current || reference) penalty. Captured at the start of an alignment
+  stage by running the freshly-loaded model in eval mode on a calibration
+  batch; prevents the alignment gradients from drifting routing away from
+  the pretrained distribution that the cluster layout was optimized for.
+"""
+from __future__ import annotations
+
+from typing import List, Optional
+
+import torch
+
+from chronos.model.moe_chronos import ChronosMOEFeedForward
+from chronos.model.temporal_loss import (
+    temporal_locality_loss,
+    lookahead_supervision_loss,
+)
+
+
+def collect_router_probs(model, *, with_grad: bool = False) -> Optional[torch.Tensor]:
+    """Stack router probs from every ChronosMOEFeedForward layer.
+
+    Args:
+        with_grad: when True, returns the live tensor that still has an
+                   autograd edge back to the gate (used by training loss
+                   terms that need to shape the gate). When False, returns
+                   the detached tensor (safe for logging / cluster layout
+                   / KL anchor reference / lookahead teacher).
+
+    Returns: ``[B, S, L, E]`` or ``None`` if no MoE layers fired in the
+    last forward pass (e.g. before the first forward).
+    """
+    probs: List[torch.Tensor] = []
+    for layer in model.model.layers:
+        mlp = getattr(layer, "mlp", None)
+        if isinstance(mlp, ChronosMOEFeedForward):
+            p = mlp.last_router_probs_grad if with_grad else mlp.last_router_probs
+            if p is not None:
+                probs.append(p)
+    if not probs:
+        return None
+    return torch.stack(probs, dim=2)
+
+
+def chronos_loss_term(
+    model,
+    base_loss: torch.Tensor,
+    lookahead_probs: Optional[torch.Tensor],
+    config,
+    *,
+    aux_loss: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Combine `base_loss` with the Chronos-specific regularizers:
+
+        L = L_base
+            + λ_balance  · L_aux            (if aux_loss provided)
+            + λ_temporal · L_temporal       (if seq has ≥2 tokens)
+            + λ_lookahead · L_lookahead     (if lookahead_probs + λ_la > 0)
+
+    The function is safe to call from any stage: it gracefully degrades
+    when the necessary tensors are unavailable (e.g. lookahead_probs is
+    None during a model that doesn't expose it).
+    """
+    loss = base_loss
+
+    if aux_loss is not None:
+        lambda_balance = float(getattr(config, "lambda_balance", 0.0))
+        if lambda_balance > 0.0:
+            # Prefer the un-scaled raw aux summed across MoE layers so we
+            # don't double-multiply by router_aux_loss_coef. Falls back to
+            # whatever the model handed us if aux_loss_raw isn't populated
+            # (e.g. eval-mode forward).
+            raw_total = base_loss.new_zeros(())
+            saw_raw = False
+            for layer in model.model.layers:
+                mlp = getattr(layer, "mlp", None)
+                if isinstance(mlp, ChronosMOEFeedForward) and mlp.aux_loss_raw is not None:
+                    raw_total = raw_total + mlp.aux_loss_raw
+                    saw_raw = True
+            aux_term = raw_total if saw_raw else aux_loss
+            loss = loss + lambda_balance * aux_term
+
+    router_4d = collect_router_probs(model, with_grad=True)
+    if router_4d is not None and router_4d.shape[1] > 1:
+        router_mean = router_4d.mean(dim=2)  # [B, S, E]
+        lambda_temporal = float(getattr(config, "lambda_temporal", 0.0))
+        if lambda_temporal > 0.0:
+            loss = loss + lambda_temporal * temporal_locality_loss(router_mean)
+
+        lambda_lookahead = float(getattr(config, "lambda_lookahead", 0.0))
+        lookahead_steps = int(getattr(config, "lookahead_steps", 0))
+        if (
+            lookahead_probs is not None
+            and lambda_lookahead > 0.0
+            and lookahead_steps > 0
+        ):
+            # Teacher must be detached: lookahead targets are stop-grad'd
+            # by construction (we predict the *future* router state).
+            teacher = router_mean.detach()
+            la = lookahead_supervision_loss(lookahead_probs, teacher, lookahead_steps)
+            loss = loss + lambda_lookahead * la
+    return loss
+
+
+def router_kl_anchor(
+    current_router_probs: torch.Tensor,
+    ref_router_probs: torch.Tensor,
+    lambda_anchor: float,
+) -> torch.Tensor:
+    """KL(current || reference) over routing distributions, scaled by λ.
+
+    Both inputs are layer-averaged ``[B, S, E]`` distributions (post-softmax).
+    Reference is treated as a fixed teacher (no grad needed; caller passes
+    a detached tensor). Reduces to mean over ``B*S``.
+    """
+    if lambda_anchor <= 0.0:
+        return current_router_probs.new_zeros(())
+    eps = 1e-9
+    kl = (current_router_probs * (
+        current_router_probs.clamp_min(eps).log()
+        - ref_router_probs.clamp_min(eps).log()
+    )).sum(dim=-1).mean()
+    return lambda_anchor * kl
+
+
+def capture_reference_routing(model, calibration_batch, device) -> torch.Tensor:
+    """Run `model` in eval mode on the calibration batch, return the
+    layer-averaged routing distribution as a detached reference for the
+    `router_kl_anchor` term.
+
+    The calibration batch is a single ``input_ids`` tensor (already on
+    device-or-not is handled here). Returns ``[B, S, E]`` detached.
+    """
+    was_training = model.training
+    model.eval()
+    with torch.no_grad():
+        ids = calibration_batch.to(device)
+        try:
+            _ = model(ids, use_cache=False)
+        except TypeError:
+            _ = model(ids)  # minimind models don't accept use_cache kwarg here
+        router_4d = collect_router_probs(model)
+        if router_4d is None:
+            ref = None
+        else:
+            ref = router_4d.mean(dim=2).detach()
+    if was_training:
+        model.train()
+    return ref