Project-Chronos 0.1.1__py3-none-any.whl

chronos/model/moe_chronos.py

@@ -0,0 +1,189 @@
+import sys
+import chronos.deps  # ensure minimind on sys.path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.activations import ACT2FN
+
+from model.model_minimind import FeedForward, MOEFeedForward
+from .config import ChronosConfig
+
+
+class LazyExpertPlaceholder(nn.Module):
+    """Inference-only placeholder for experts that live on SSD/RAM.
+
+    If this module executes directly, Chronos failed to materialize the
+    expert before use or failed to route the token through the shared
+    expert fallback. That is a correctness bug, so we raise loudly.
+    """
+
+    def __init__(self, hidden_size: int, intermediate_size: int, dtype: torch.dtype):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.dtype = dtype
+        self.loaded = False
+
+    def forward(self, _x):
+        raise RuntimeError("lazy expert placeholder executed before materialization")
+
+
+class ChronosMOEFeedForward(nn.Module):
+    """
+    Extends MOEFeedForward with:
+    1. shared_experts: always resident in VRAM, used as soft fallback
+    2. soft_gate: when target experts are unavailable (cache miss), blend
+       their routing weight into shared experts instead of blocking on I/O
+
+    Training (available_expert_mask=None): identical to MOEFeedForward,
+    plus stores per-token routing probs for temporal_locality_loss.
+
+    Inference (available_expert_mask provided): routes cache-miss tokens
+    to shared experts with zero latency.
+    """
+
+    def __init__(self, config: ChronosConfig):
+        super().__init__()
+        self.config = config
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+        self.experts = nn.ModuleList([
+            FeedForward(config, intermediate_size=config.moe_intermediate_size)
+            for _ in range(config.num_experts)
+        ])
+        self.shared_experts = nn.ModuleList([
+            FeedForward(config, intermediate_size=config.moe_intermediate_size)
+            for _ in range(config.num_shared_experts)
+        ])
+        self.aux_loss = torch.zeros(1)
+        # raw, *unscaled* load-balance value computed in forward(). The
+        # trainer-side λ_balance lives in chronos_loss_term so we don't
+        # double-scale (router_aux_loss_coef × λ_balance gave a 2 000×
+        # under-weighting and let the gate degenerate). Set to None until
+        # the first training forward populates it.
+        self.aux_loss_raw: torch.Tensor = None
+        # Stores routing probs for downstream loss / metrics.
+        # We keep TWO refs:
+        #   last_router_probs       — detached, safe for logging / cluster
+        #                              layout / inference engine prefetch
+        #                              (consumers that must NOT change grad).
+        #   last_router_probs_grad  — live tensor with autograd attached,
+        #                              consumed by chronos_loss_term so the
+        #                              λ_temporal / λ_lookahead penalties
+        #                              actually shape the gate.
+        self.last_router_probs: torch.Tensor = None
+        self.last_router_probs_grad: torch.Tensor = None
+
+    def _shared_expert_output(self, x_flat: torch.Tensor) -> torch.Tensor:
+        """Average output of all shared experts."""
+        out = sum(e(x_flat) for e in self.shared_experts)
+        return out / len(self.shared_experts)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        available_expert_mask: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: [B, S, H]
+            available_expert_mask: [num_experts] bool or float — which experts
+                are in VRAM. None during training (all experts available).
+
+        When available_expert_mask is provided the forward pass is written
+        as pure tensor arithmetic (no Python if/else per expert) so that
+        torch.compile can trace a single static graph without graph breaks.
+        """
+        B, S, H = x.shape
+        x_flat = x.view(-1, H)                                   # [N, H]
+        scores = F.softmax(self.gate(x_flat), dim=-1)            # [N, E]
+        scores_bse = scores.view(B, S, self.num_experts)
+        # Live (with-grad) view drives the temporal / lookahead losses.
+        self.last_router_probs_grad = scores_bse
+        # Detached view for read-only consumers (logging, cluster layout,
+        # inference-engine prefetch). Same tensor data, no autograd edge.
+        self.last_router_probs = scores_bse.detach()
+
+        topk_weight, topk_idx = torch.topk(
+            scores, k=self.num_experts_per_tok, dim=-1, sorted=False
+        )
+        if self.norm_topk_prob:
+            topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
+
+        y = torch.zeros_like(x_flat)
+
+        if available_expert_mask is None:
+            # --- Training path: standard MoE dispatch + always-on shared expert ---
+            dead_param_sum = None
+            for i, expert in enumerate(self.experts):
+                mask = (topk_idx == i)
+                if mask.any():
+                    token_idx = mask.any(dim=-1).nonzero().flatten()
+                    weight = topk_weight[mask].view(-1, 1)
+                    y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
+                elif self.training:
+                    # Dead-expert handling: keep the autograd edge alive
+                    # WITHOUT touching y[0,0] (the previous trick poisoned
+                    # the BOS position every step). We accumulate a
+                    # 0·sum(params) scalar that gets routed into aux_loss
+                    # below, which the trainer always backprops on.
+                    contrib = sum(p.sum() for p in expert.parameters())
+                    dead_param_sum = (
+                        contrib if dead_param_sum is None else dead_param_sum + contrib
+                    )
+
+            # Always-on shared expert (mirrors the inference branch). The
+            # previous code only ran shared_experts under the inference
+            # mask path, so during training they received zero gradient
+            # — and at inference, the cache-miss fallback fired untrained
+            # weights, which is the dominant root cause of "loss looks
+            # fine but generations are gibberish."
+            if len(self.shared_experts) > 0:
+                y = y + self._shared_expert_output(x_flat).to(y.dtype)
+
+            if self.training:
+                # Raw, un-scaled load balance (model exposes raw; trainer scales).
+                load = F.one_hot(topk_idx, self.num_experts).float().mean(0)
+                aux_raw = (load * scores.mean(0)).sum() * self.num_experts
+                if dead_param_sum is not None:
+                    aux_raw = aux_raw + 0.0 * dead_param_sum
+                self.aux_loss_raw = aux_raw
+                # Backwards-compat: the legacy `aux_loss` field stays
+                # populated with the pre-scaled value some callers still
+                # read (chronos_loss_term now prefers aux_loss_raw).
+                self.aux_loss = aux_raw * self.router_aux_loss_coef
+            else:
+                self.aux_loss_raw = scores.new_zeros(())
+                self.aux_loss = scores.new_zeros(())
+
+        else:
+            # --- Inference path: real lazy experts + shared fallback ---
+            # We intentionally branch at the Python level here: when an
+            # expert is cold we must not execute it at all, because the
+            # live module may have been replaced by a placeholder after
+            # expert offload. A multiply-by-zero path would still call the
+            # unloaded expert and defeat lazy loading entirely.
+            avail = available_expert_mask.to(dtype=x_flat.dtype, device=x_flat.device)  # [E]
+            shared_out_full = self._shared_expert_output(x_flat)  # [N, H] — computed once
+
+            for i, expert in enumerate(self.experts):
+                mask = (topk_idx == i)                          # [N, K] bool
+                if not mask.any():
+                    continue
+                token_idx = mask.any(dim=-1).nonzero().flatten()   # [T]
+                weight = topk_weight[mask].view(-1, 1)             # [T, 1]
+                fallback_out = shared_out_full[token_idx]          # [T, H]
+                if float(avail[i].item()) > 0.0:
+                    blended = expert(x_flat[token_idx])            # [T, H]
+                else:
+                    blended = fallback_out
+                y.index_add_(0, token_idx, (blended * weight).to(y.dtype))
+
+            self.aux_loss = scores.new_zeros(1).squeeze()
+
+        return y.view(B, S, H)

chronos/model/temporal_loss.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn.functional as F
+
+
+def temporal_locality_loss(expert_probs_seq: torch.Tensor) -> torch.Tensor:
+    """
+    Penalizes large routing distribution shifts between adjacent time steps.
+
+    L_temporal = (1/T) * sum_{t=2}^{T} ||E_t - E_{t-1}||_2^2
+
+    Args:
+        expert_probs_seq: [B, T, num_experts]  per-step routing probabilities
+    Returns:
+        scalar loss
+    """
+    diff = expert_probs_seq[:, 1:, :] - expert_probs_seq[:, :-1, :]
+    return (diff ** 2).sum(dim=-1).mean()
+
+
+def load_balance_loss(
+    router_logits: torch.Tensor,
+    num_experts: int,
+    num_experts_per_tok: int,
+) -> torch.Tensor:
+    """
+    Standard auxiliary load-balancing loss (same formulation as minimind).
+
+    Args:
+        router_logits: [N, num_experts]  (N = B*S flattened)
+    Returns:
+        scalar loss
+    """
+    scores = F.softmax(router_logits, dim=-1)
+    _, topk_idx = torch.topk(scores, k=num_experts_per_tok, dim=-1, sorted=False)
+    load = F.one_hot(topk_idx, num_experts).float().mean(0)
+    return (load * scores.mean(0)).sum() * num_experts
+
+
+def lookahead_supervision_loss(
+    lookahead_probs: torch.Tensor,
+    teacher_router_probs: torch.Tensor,
+    lookahead_steps: int,
+) -> torch.Tensor:
+    """
+    Teacher-forcing loss for the LookaheadRouter.
+
+    At time t the router predicts the routing distribution at t+k for
+    k ∈ {1..lookahead_steps}. We supervise those predictions with the
+    *actual* layer-averaged router probabilities from step t+k (stop-grad).
+
+    Formulation (soft-target cross-entropy, equivalent to forward KL
+    up to a teacher entropy constant):
+
+        L = (1 / K) · Σ_{k=1..K} [ - mean_{t} Σ_e P_{t+k}[e] · log Q_t^(k)[e] ]
+
+    Args:
+        lookahead_probs:       [B, S, K+1, E] — softmax'd output of
+                                LookaheadRouter. Index 0 = current step,
+                                indices 1..K = future-step predictions.
+        teacher_router_probs:  [B, S, E]       — layer-averaged real router
+                                probabilities. Caller MUST pass a detached
+                                (stop-grad) tensor.
+        lookahead_steps:       K, how many future steps to supervise.
+
+    Returns:
+        scalar loss. Zero when the sequence is too short to form a pair.
+    """
+    if lookahead_probs is None or teacher_router_probs is None:
+        return lookahead_probs.new_zeros(()) if lookahead_probs is not None else \
+               teacher_router_probs.new_zeros(())
+
+    B, S, Kp1, E = lookahead_probs.shape
+    K = min(lookahead_steps, Kp1 - 1)
+    if K <= 0 or S <= 1:
+        return lookahead_probs.new_zeros(())
+
+    teacher = teacher_router_probs  # caller is responsible for .detach()
+    total = lookahead_probs.new_zeros(())
+    valid_terms = 0
+    for k in range(1, K + 1):
+        if S - k <= 0:
+            continue
+        teacher_k = teacher[:, k:, :]                  # [B, S-k, E]
+        pred_k = lookahead_probs[:, :-k, k, :]         # [B, S-k, E]
+        # Soft-target CE: −Σ_e P · log Q
+        log_pred = pred_k.clamp_min(1e-9).log()
+        ce = -(teacher_k * log_pred).sum(dim=-1)        # [B, S-k]
+        total = total + ce.mean()
+        valid_terms += 1
+    if valid_terms == 0:
+        return lookahead_probs.new_zeros(())
+    return total / valid_terms
+
+
+def total_loss(
+    ce_loss: torch.Tensor,
+    balance_loss: torch.Tensor,
+    expert_probs_seq: torch.Tensor,
+    lambda1: float,
+    lambda2: float,
+    *,
+    lookahead_probs: torch.Tensor = None,
+    teacher_probs: torch.Tensor = None,
+    lookahead_steps: int = 0,
+    lambda_lookahead: float = 0.0,
+) -> torch.Tensor:
+    """
+    L_total = L_CE + λ1·L_balance + λ2·L_temporal + λ_lookahead·L_lookahead
+
+    The lookahead term is only applied when `lookahead_probs`, `teacher_probs`
+    and `lambda_lookahead > 0` are all provided — preserves the M1-era
+    signature for legacy callers.
+    """
+    loss = ce_loss + lambda1 * balance_loss + lambda2 * temporal_locality_loss(expert_probs_seq)
+    if (
+        lookahead_probs is not None
+        and teacher_probs is not None
+        and lambda_lookahead > 0.0
+        and lookahead_steps > 0
+    ):
+        loss = loss + lambda_lookahead * lookahead_supervision_loss(
+            lookahead_probs, teacher_probs, lookahead_steps,
+        )
+    return loss

chronos/router/__init__.py

@@ -0,0 +1,12 @@
+"""
+chronos/router/__init__.py
+"""
+from chronos.router.intent_classifier import IntentClassifier, IntentVector
+from chronos.router.expert_predictor import ExpertPredictor, ExpertSet
+from chronos.router.prefill_scheduler import PrefillScheduler
+
+__all__ = [
+    "IntentClassifier", "IntentVector",
+    "ExpertPredictor", "ExpertSet",
+    "PrefillScheduler",
+]

chronos/router/expert_predictor.py

@@ -0,0 +1,163 @@
+"""
+chronos/router/expert_predictor.py
+
+Layer 2 of the routing pipeline: converts IntentVector probabilities into
+a concrete expert set to pre-load before generation begins.
+
+Responsibilities:
+  1. Threshold-based selection: keep experts with prob > threshold
+  2. Budget enforcement: cap at vram_capacity experts
+  3. Per-layer masks for soft gating (avail_masks)
+  4. Confidence reporting: if max_prob < min_confidence, fall back to
+     top-K heuristic (always load the K most common experts)
+
+The ExpertPredictor is stateless — it takes an IntentVector and returns
+an ExpertSet. No weights, no training needed.
+"""
+from dataclasses import dataclass, field
+from typing import List, Optional, Set, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from chronos.router.intent_classifier import IntentVector
+
+
+@dataclass
+class ExpertSet:
+    """
+    The concrete set of experts to pre-load for a generation session.
+
+    expert_ids: flat set of expert IDs to load (union across all layers)
+    per_layer_masks: List[Tensor([num_experts], bool)] — per-layer avail
+    confidence: float — how confident the classifier was (0–1)
+    fallback: bool — True if we fell back to heuristic selection
+    """
+    expert_ids: Set[int]
+    per_layer_masks: List[torch.Tensor]
+    confidence: float
+    fallback: bool = False
+
+    def __repr__(self):
+        return (f"ExpertSet(experts={sorted(self.expert_ids)}, "
+                f"confidence={self.confidence:.2f}, fallback={self.fallback})")
+
+
+class ExpertPredictor:
+    """
+    Converts IntentVector → ExpertSet.
+
+    Args:
+        num_experts: total number of experts
+        num_moe_layers: number of MoE layers
+        vram_capacity: max experts that fit in VRAM simultaneously
+        threshold: minimum probability to include an expert
+        min_confidence: if best prob < this, use fallback heuristic
+        fallback_top_k: number of experts to always include as fallback
+    """
+
+    def __init__(
+        self,
+        num_experts: int,
+        num_moe_layers: int,
+        vram_capacity: int = 4,
+        threshold: float = 0.25,
+        min_confidence: float = 0.15,
+        fallback_top_k: int = 2,
+    ):
+        self.num_experts = num_experts
+        self.num_moe_layers = num_moe_layers
+        self.vram_capacity = vram_capacity
+        self.threshold = threshold
+        self.min_confidence = min_confidence
+        self.fallback_top_k = fallback_top_k
+
+        # Running frequency table: updated from actual activations over time
+        self._frequency: torch.Tensor = torch.ones(num_experts) / num_experts
+
+    def predict(self, intent: "IntentVector") -> ExpertSet:
+        """
+        Convert IntentVector to a concrete ExpertSet.
+
+        Strategy:
+          1. Check confidence (max prob across all layers)
+          2. If confident: threshold + budget cap
+          3. If not confident: fallback to global frequency heuristic
+        """
+        from chronos.router.intent_classifier import IntentVector
+
+        global_p = intent.global_probs  # [E]
+        confidence = float(global_p.max().item())
+        fallback = confidence < self.min_confidence
+
+        if fallback:
+            # Heuristic: always load the historically most-activated experts
+            top_ids = self._frequency.topk(
+                min(self.fallback_top_k, self.vram_capacity)
+            ).indices.tolist()
+            expert_ids = set(top_ids)
+        else:
+            # Per-layer selection
+            selected_per_layer = []
+            all_ids: Set[int] = set()
+            for layer_p in intent.per_layer_probs:
+                above = (layer_p >= self.threshold).nonzero().flatten().tolist()
+                selected_per_layer.append(set(above))
+                all_ids.update(above)
+
+            # Enforce budget: keep highest-prob experts if over capacity
+            if len(all_ids) > self.vram_capacity:
+                top_ids = global_p.topk(self.vram_capacity).indices.tolist()
+                all_ids = set(top_ids)
+
+            expert_ids = all_ids
+
+        # Build per-layer boolean masks
+        per_layer_masks = []
+        for li in range(self.num_moe_layers):
+            mask = torch.zeros(self.num_experts, dtype=torch.bool)
+            layer_p = intent.per_layer_probs[li] if li < len(intent.per_layer_probs) else global_p
+            if fallback:
+                for eid in expert_ids:
+                    mask[eid] = True
+            else:
+                for eid in expert_ids:
+                    if layer_p[eid] >= self.threshold or eid in expert_ids:
+                        mask[eid] = True
+            per_layer_masks.append(mask)
+
+        return ExpertSet(
+            expert_ids=expert_ids,
+            per_layer_masks=per_layer_masks,
+            confidence=confidence,
+            fallback=fallback,
+        )
+
+    def update_frequency(self, actual_expert_ids: List[int]):
+        """
+        Update running activation frequency from observed generation.
+        Called after each generation session with the experts that were
+        actually activated.
+        """
+        alpha = 0.05  # EMA smoothing
+        observed = torch.zeros(self.num_experts)
+        for eid in actual_expert_ids:
+            if 0 <= eid < self.num_experts:
+                observed[eid] += 1
+        if observed.sum() > 0:
+            observed = observed / observed.sum()
+        self._frequency = (1 - alpha) * self._frequency + alpha * observed
+
+    def avail_masks_float(self, expert_set: ExpertSet, device=None) -> List[torch.Tensor]:
+        """
+        Convert per_layer_masks (bool) to float tensors for soft gating.
+        0.0 = cache miss → shared expert fallback
+        1.0 = in VRAM → use directly
+        """
+        result = []
+        for mask in expert_set.per_layer_masks:
+            f = mask.float()
+            if device is not None:
+                f = f.to(device)
+            result.append(f)
+        return result

chronos/router/intent_classifier.py

@@ -0,0 +1,193 @@
+"""
+chronos/router/intent_classifier.py
+
+Layer 1 of the dual-layer architecture: a small dense model that reads
+the full prompt during prefill and outputs an "intent vector" —
+a distribution over expert-activation patterns.
+
+Design:
+  - 2-layer Transformer encoder (~8-15M params, much smaller than main MoE)
+  - Input: tokenised prompt (first min(512, seq_len) tokens)
+  - Output: IntentVector — a [num_experts] float tensor (per layer) that
+    predicts which experts the main MoE will need during generation
+
+Why separate from LookaheadRouter:
+  LookaheadRouter predicts 1-2 steps ahead within the generation loop.
+  IntentClassifier reads the ENTIRE prompt ONCE before generation starts,
+  producing a global expert-set prediction that covers the whole reply.
+  This moves all IO uncertainty to the prefill phase.
+
+Training:
+  Supervised on (prompt, expert_activation_log) pairs collected from
+  the main model's forward passes. Loss = MSE between predicted and
+  actual per-layer expert activation frequencies.
+"""
+import os
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class IntentVector:
+    """
+    Output of IntentClassifier for a single prompt.
+
+    per_layer_probs: List[Tensor([num_experts])] — one per MoE layer.
+        Values are probabilities; higher means more likely to be needed.
+    global_probs: Tensor([num_experts]) — averaged across layers.
+    """
+    per_layer_probs: List[torch.Tensor]
+    global_probs: torch.Tensor
+
+
+class IntentClassifier(nn.Module):
+    """
+    Small dense encoder that maps a prompt to an expert activation prediction.
+
+    Architecture: embedding → 2 Transformer encoder layers → mean-pool →
+                  linear head → sigmoid per-layer expert probabilities
+
+    Parameters: ~10-15M (hidden=256, heads=4, layers=2)
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        num_experts: int,
+        num_moe_layers: int,
+        hidden_size: int = 256,
+        num_heads: int = 4,
+        num_encoder_layers: int = 2,
+        max_seq_len: int = 512,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.num_moe_layers = num_moe_layers
+        self.hidden_size = hidden_size
+        self.max_seq_len = max_seq_len
+
+        self.embed = nn.Embedding(vocab_size, hidden_size, padding_idx=0)
+        self.pos_embed = nn.Embedding(max_seq_len, hidden_size)
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_size,
+            nhead=num_heads,
+            dim_feedforward=hidden_size * 4,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
+        self.norm = nn.LayerNorm(hidden_size)
+
+        # Per-layer expert probability heads
+        # Output: [num_moe_layers, num_experts] per sample
+        self.head = nn.Linear(hidden_size, num_moe_layers * num_experts)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, input_ids: torch.Tensor) -> IntentVector:
+        """
+        Args:
+            input_ids: [B, S] token ids (truncated to max_seq_len)
+        Returns:
+            IntentVector with per-layer and global expert probabilities
+        """
+        B, S = input_ids.shape
+        S = min(S, self.max_seq_len)
+        ids = input_ids[:, :S]
+
+        pos = torch.arange(S, device=ids.device).unsqueeze(0)  # [1, S]
+        x = self.embed(ids) + self.pos_embed(pos)               # [B, S, H]
+        x = self.dropout(x)
+
+        # Padding mask for attention (0 = pad token)
+        pad_mask = (ids == 0)  # [B, S]
+
+        x = self.encoder(x, src_key_padding_mask=pad_mask)
+        x = self.norm(x)
+
+        # Mean-pool over non-padding tokens
+        lengths = (~pad_mask).float().sum(dim=1, keepdim=True).clamp(min=1)  # [B, 1]
+        x_masked = x * (~pad_mask).float().unsqueeze(-1)
+        pooled = x_masked.sum(dim=1) / lengths                 # [B, H]
+
+        logits = self.head(pooled)                              # [B, L*E]
+        probs = torch.sigmoid(logits).view(B, self.num_moe_layers, self.num_experts)
+
+        # Build IntentVector (for batch size 1 during inference)
+        per_layer = [probs[0, i] for i in range(self.num_moe_layers)]
+        global_p = probs[0].mean(dim=0)
+        return IntentVector(per_layer_probs=per_layer, global_probs=global_p)
+
+    # ── Training helpers ─────────────────────────────────────────
+
+    @staticmethod
+    def collect_activation_targets(
+        model,
+        input_ids: torch.Tensor,
+        device: str,
+    ) -> torch.Tensor:
+        """
+        Run main model forward pass and collect actual expert activation
+        frequencies as training labels for IntentClassifier.
+
+        Returns: [num_moe_layers, num_experts] float tensor (frequencies)
+        """
+        from chronos.model.moe_chronos import ChronosMOEFeedForward
+        model.eval()
+        with torch.no_grad():
+            model(input_ids.to(device))
+
+        moe_layers = [l.mlp for l in model.model.layers
+                      if isinstance(l.mlp, ChronosMOEFeedForward)]
+        targets = []
+        for moe in moe_layers:
+            if moe.last_router_probs is not None:
+                # Mean activation probability across tokens
+                targets.append(moe.last_router_probs.mean(dim=(0, 1)))  # [E]
+            else:
+                targets.append(torch.zeros(model.config.num_experts))
+        return torch.stack(targets)  # [L, E]
+
+    def train_step(
+        self,
+        input_ids: torch.Tensor,
+        targets: torch.Tensor,
+        optimizer: torch.optim.Optimizer,
+    ) -> float:
+        """
+        Single training step.
+        targets: [B, num_moe_layers, num_experts] float
+        """
+        self.train()
+        optimizer.zero_grad(set_to_none=True)
+        B, S = input_ids.shape
+        logits = self.head(
+            self.norm(self.encoder(
+                self.dropout(self.embed(input_ids[:, :self.max_seq_len])
+                             + self.pos_embed(
+                                 torch.arange(min(S, self.max_seq_len),
+                                              device=input_ids.device).unsqueeze(0)))
+            ).mean(dim=1))
+        ).view(B, self.num_moe_layers, self.num_experts)
+        probs = torch.sigmoid(logits)
+        loss = F.mse_loss(probs, targets.to(probs.device))
+        loss.backward()
+        nn.utils.clip_grad_norm_(self.parameters(), 1.0)
+        optimizer.step()
+        return loss.item()
+
+    def save(self, path: str):
+        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+        torch.save(self.state_dict(), path)
+
+    @classmethod
+    def load(cls, path: str, **kwargs) -> "IntentClassifier":
+        m = cls(**kwargs)
+        m.load_state_dict(torch.load(path, map_location="cpu"))
+        return m