Project-Chronos 0.2.0__tar.gz → 0.2.1__tar.gz

{project_chronos-0.2.0 → project_chronos-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Project_Chronos
-Version: 0.2.0
+Version: 0.2.1
 Summary: On-device low-latency lookahead dual-layer MoE inference — Lookahead Routing + Async Prefetch + Hybrid Attention
 License-Expression: Apache-2.0
 Keywords: llm,moe,inference,lookahead-routing,hybrid-attention,edge-ai

{project_chronos-0.2.0 → project_chronos-0.2.1}/Project_Chronos.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Project_Chronos
-Version: 0.2.0
+Version: 0.2.1
 Summary: On-device low-latency lookahead dual-layer MoE inference — Lookahead Routing + Async Prefetch + Hybrid Attention
 License-Expression: Apache-2.0
 Keywords: llm,moe,inference,lookahead-routing,hybrid-attention,edge-ai

{project_chronos-0.2.0 → project_chronos-0.2.1}/chronos/mlx/expert_store.py

@@ -72,6 +72,7 @@ class MLXExpertStore:
         self._warm_lru: collections.OrderedDict = collections.OrderedDict()
         self._cluster_manifest: ClusterManifest | None = None
         self._loaded_clusters: set[int] = set()
+        self._protected_experts: set[int] = set()
         self.storage_format = "npy"
         self.attach_cluster_manifest(ssd_dir)
         self._stats = collections.Counter()
@@ -296,18 +297,25 @@ class MLXExpertStore:
         self._warm_lru.move_to_end(expert_id)
         self._evict_warm_locked(protected_expert_id=protected_expert_id)
 
+    def set_protected_experts(self, expert_ids: List[int] | set[int] | tuple[int, ...]):
+        with self._lock:
+            self._protected_experts = {
+                int(eid) for eid in expert_ids
+                if 0 <= int(eid) < int(self.num_experts)
+            }
+
     def _evict_warm_locked(self, protected_expert_id: int | None = None):
         while len(self._warm_lru) > self._warm_capacity:
             evict_id = None
             for candidate in self._warm_lru.keys():
-                if candidate == protected_expert_id:
+                if candidate == protected_expert_id or candidate in self._protected_experts:
                     continue
                 if candidate not in self._hot_lru:
                     evict_id = candidate
                     break
             if evict_id is None:
                 for candidate in self._warm_lru.keys():
-                    if candidate != protected_expert_id:
+                    if candidate != protected_expert_id and candidate not in self._protected_experts:
                         evict_id = candidate
                         break
             if evict_id is None:
@@ -455,9 +463,17 @@ class MLXExpertStore:
                 self._hot_lru.move_to_end(expert_id)
             else:
                 while len(self._hot_lru) >= self._capacity:
-                    evicted, _ = self._hot_lru.popitem(last=False)
-                    if evicted == expert_id:
-                        continue
+                    evicted = None
+                    for candidate in self._hot_lru.keys():
+                        if candidate != expert_id and candidate not in self._protected_experts:
+                            evicted = candidate
+                            break
+                    if evicted is None:
+                        evicted, _ = self._hot_lru.popitem(last=False)
+                        if evicted == expert_id:
+                            continue
+                    else:
+                        self._hot_lru.pop(evicted, None)
                     # Keep warm weights in unified memory, but drop the
                     # executable live module so masked MoE cannot call an
                     # expert outside the hot execution budget.

{project_chronos-0.2.0 → project_chronos-0.2.1}/chronos/mlx/inference.py

@@ -47,7 +47,11 @@ class ChronosMLXInferenceEngine:
         self.last_stats: dict = {}
         self._last_predicted: set[int] = set()
         self._stats_lock = threading.Lock()
-        self._runtime_stats = {
+        self._runtime_stats = self._new_runtime_stats()
+
+    @staticmethod
+    def _new_runtime_stats(all_resident_fast_path: bool = False) -> dict:
+        return {
             "resident_hits": 0,
             "resident_misses": 0,
             "resident_vram_hits": 0,
@@ -60,6 +64,9 @@ class ChronosMLXInferenceEngine:
             "prefetch_wait_time_s": 0.0,
             "sync_ssd_loads": 0,
             "on_demand_load_time_s": 0.0,
+            "cold_prediction_hits": 0,
+            "cold_prediction_total": 0,
+            "all_resident_fast_path": bool(all_resident_fast_path),
         }
 
     def setup(self, warm_expert_ids: Optional[List[int]] = None):
@@ -90,6 +97,9 @@ class ChronosMLXInferenceEngine:
             if not was_hot and not was_warm:
                 with self._stats_lock:
                     self._runtime_stats["sync_ssd_loads"] += 1
+                    self._runtime_stats["cold_prediction_total"] += 1
+                    if eid in self._last_predicted:
+                        self._runtime_stats["cold_prediction_hits"] += 1
             ok = self.store.promote_to_vram(eid)
             elapsed = time.monotonic() - t0
             with self._stats_lock:
@@ -129,6 +139,14 @@ class ChronosMLXInferenceEngine:
                 moe.runtime_on_demand_loader = loader
                 moe.runtime_touch_expert = touch
 
+    def _clear_runtime_hooks(self):
+        for layer in self.model.layers:
+            moe = getattr(layer, "mlp", None)
+            if isinstance(moe, ChronosMLXMOE):
+                moe.runtime_miss_policy = "sync_on_demand"
+                moe.runtime_on_demand_loader = None
+                moe.runtime_touch_expert = None
+
     # ── Prefetch background thread ────────────────────────────────
 
     def _prefetch_loop(self):
@@ -149,37 +167,63 @@ class ChronosMLXInferenceEngine:
             with self._stats_lock:
                 self._runtime_stats["prefetch_queue_drops"] += 1
 
-    def _prefetch_and_promote_window(self, expert_ids: List[int], timeout_s: float = 0.012) -> None:
-        if not expert_ids or self.store.storage_format == "full_dram":
+    def _prefetch_and_promote_window(
+        self,
+        prefetch_ids: List[int],
+        promote_ids: Optional[List[int]] = None,
+        timeout_s: float = 0.012,
+    ) -> None:
+        if not prefetch_ids or self.store.storage_format == "full_dram":
             return
+        promote_ids = list(dict.fromkeys(int(eid) for eid in (promote_ids or prefetch_ids)))
         pending = []
-        for eid in dict.fromkeys(int(eid) for eid in expert_ids):
+        for eid in dict.fromkeys(int(eid) for eid in prefetch_ids):
             with self.store._lock:
                 if eid in self.store._hot_lru:
                     continue
                 if eid in self.store._warm and self.store._layer_states_complete(self.store._warm.get(eid)):
-                    pending.append(eid)
                     continue
             pending.append(eid)
-        if not pending:
-            return
-        self._schedule_prefetch(pending)
+        if pending:
+            self._schedule_prefetch(pending)
         deadline = time.monotonic() + max(0.0, float(timeout_s))
         while time.monotonic() < deadline:
-            self._promote_ready(pending)
+            self._promote_ready(promote_ids)
             with self.store._lock:
                 ready = all(
                     eid in self.store._hot_lru
                     or (eid in self.store._warm and self.store._layer_states_complete(self.store._warm.get(eid)))
-                    for eid in pending
+                    for eid in promote_ids
                 )
             if ready:
                 break
             time.sleep(0.001)
-        self._promote_ready(pending)
+        self._promote_ready(promote_ids)
         with self._stats_lock:
             self._runtime_stats["prefetch_wait_time_s"] += max(0.0, time.monotonic() - (deadline - max(0.0, float(timeout_s))))
 
+    def _all_resident_fast_path_enabled(self) -> bool:
+        if self.store.storage_format == "full_dram":
+            return True
+        num_experts = int(getattr(self.config, "num_experts", 0) or 0)
+        if num_experts <= 0:
+            return False
+        if int(getattr(self.store, "_capacity", 0) or 0) < num_experts:
+            return False
+        if int(getattr(self.store, "_warm_capacity", 0) or 0) < num_experts:
+            return False
+        return len(self.store.hot_expert_ids()) >= num_experts
+
+    def _protect_prediction_window(self, expert_ids: List[int]) -> None:
+        capacity = max(
+            1,
+            int(getattr(self.store, "_warm_capacity", 1) or 1),
+            int(getattr(self.store, "_capacity", 1) or 1),
+        )
+        hot_ids = sorted(self.store.hot_expert_ids())
+        protected = list(dict.fromkeys([int(eid) for eid in expert_ids] + hot_ids))[:capacity]
+        self.store.set_protected_experts(protected)
+
     def stop(self):
         self._stop.set()
         try:
@@ -210,23 +254,15 @@ class ChronosMLXInferenceEngine:
         setup_mem = self._memory_snapshot()
         self._last_predicted = set()
         with self._stats_lock:
-            self._runtime_stats = {
-                "resident_hits": 0,
-                "resident_misses": 0,
-                "resident_vram_hits": 0,
-                "resident_ram_hits": 0,
-                "selection_hits": 0,
-                "selection_misses": 0,
-                "prediction_hits": 0,
-                "prediction_total": 0,
-                "prefetch_queue_drops": 0,
-                "prefetch_wait_time_s": 0.0,
-                "sync_ssd_loads": 0,
-                "on_demand_load_time_s": 0.0,
-            }
-        self._install_runtime_hooks()
+            all_resident_fast_path = self._all_resident_fast_path_enabled()
+            self._runtime_stats = self._new_runtime_stats(all_resident_fast_path)
+        if all_resident_fast_path:
+            self._clear_runtime_hooks()
+            self.store.set_protected_experts([])
+        else:
+            self._install_runtime_hooks()
         # ── Prefill phase: front-load expert IO ───────────────────────────
-        if scheduler is not None:
+        if scheduler is not None and not all_resident_fast_path:
             import torch, numpy as np
             ids_np = np.array(input_ids.tolist(), dtype=np.int64)
             ids_pt = torch.from_numpy(ids_np)
@@ -235,7 +271,7 @@ class ChronosMLXInferenceEngine:
 
         # Prefill forward
         prefill_t0 = time.monotonic()
-        if self.store.storage_format == "full_dram":
+        if all_resident_fast_path:
             prefill_masks = None
         else:
             prefill_masks = self._build_avail_masks(None)
@@ -248,21 +284,41 @@ class ChronosMLXInferenceEngine:
         next_token = self._sample(logits[:, -1, :], temperature, top_p)
         activated_ids: List[int] = []
         tokens = 1
-        if scheduler is None and lookahead_probs is not None:
-            future_ids = self._predict_future_experts(lookahead_probs)
+        if not all_resident_fast_path and scheduler is None and lookahead_probs is not None:
+            future_ids = self._predict_future_experts(
+                lookahead_probs,
+                capacity=int(getattr(self.store, "_warm_capacity", self.config.num_experts) or self.config.num_experts),
+            )
+            immediate_ids = self._predict_future_experts(
+                lookahead_probs,
+                capacity=int(getattr(self.store, "_capacity", 1) or 1),
+                max_steps=1,
+            )
             self._last_predicted = set(int(eid) for eid in future_ids)
-            self._prefetch_and_promote_window(future_ids, timeout_s=0.025)
+            self._protect_prediction_window(future_ids)
+            self._prefetch_and_promote_window(future_ids, promote_ids=immediate_ids, timeout_s=0.025)
         yield int(next_token.item())
 
         for _ in range(max_new_tokens - 1):
-            if scheduler is not None:
+            if all_resident_fast_path:
+                avail_masks = None
+            elif scheduler is not None:
                 avail_masks = self._build_avail_masks(next_token)
             else:
                 # LookaheadRouter-driven prefetch
                 if lookahead_probs is not None:
-                    future_ids = self._predict_future_experts(lookahead_probs)
+                    future_ids = self._predict_future_experts(
+                        lookahead_probs,
+                        capacity=int(getattr(self.store, "_warm_capacity", self.config.num_experts) or self.config.num_experts),
+                    )
+                    immediate_ids = self._predict_future_experts(
+                        lookahead_probs,
+                        capacity=int(getattr(self.store, "_capacity", 1) or 1),
+                        max_steps=1,
+                    )
                     self._last_predicted = set(int(eid) for eid in future_ids)
-                    self._prefetch_and_promote_window(future_ids)
+                    self._protect_prediction_window(future_ids)
+                    self._prefetch_and_promote_window(future_ids, promote_ids=immediate_ids)
                 avail_masks = self._build_avail_masks(next_token)
 
             token_in = next_token.reshape(1, 1)
@@ -298,22 +354,39 @@ class ChronosMLXInferenceEngine:
             **self.store.stats(),
         }
 
-    def _predict_future_experts(self, lookahead_probs: mx.array) -> List[int]:
+    def _predict_future_experts(
+        self,
+        lookahead_probs: mx.array,
+        capacity: Optional[int] = None,
+        max_steps: Optional[int] = None,
+    ) -> List[int]:
         """Extract bounded top-k expert IDs for future steps."""
         # lookahead_probs: [B, S, K+1, E] — take last token, steps 1..K
         future = lookahead_probs[0, -1, 1:, :]         # [K, E]
+        if max_steps is not None:
+            future = future[:max(0, int(max_steps))]
+        if future.shape[0] <= 0:
+            return []
+        num_experts = int(getattr(self.config, "num_experts", future.shape[-1]) or future.shape[-1])
         top_k = max(1, min(
             int(getattr(self.config, "num_experts_per_tok", 1) or 1),
-            int(getattr(self.config, "num_experts", future.shape[-1]) or future.shape[-1]),
+            num_experts,
         ))
-        ids: list[int] = []
-        for k in range(future.shape[0]):
-            step = future[k]
-            if top_k == 1:
-                ids.append(int(mx.argmax(step).item()))
-            else:
-                ids.extend(int(v) for v in mx.argpartition(-step, kth=top_k - 1, axis=-1)[:top_k].tolist())
-        return list(dict.fromkeys(ids))
+        budget = num_experts if capacity is None else max(1, min(num_experts, int(capacity)))
+        scores = mx.zeros((num_experts,), dtype=mx.float32)
+        for offset in range(future.shape[0]):
+            step = future[offset].astype(mx.float32)
+            if top_k >= num_experts:
+                scores = scores + step / float(offset + 1)
+                continue
+            step_ids = mx.argpartition(-step, kth=top_k - 1, axis=-1)[:top_k]
+            step_vals = mx.take(step, step_ids) / float(offset + 1)
+            one_hot = (step_ids[:, None] == mx.arange(num_experts)).astype(mx.float32)
+            scores = scores + (one_hot * step_vals[:, None]).sum(axis=0)
+        mx.eval(scores)
+        scored = [(idx, float(value)) for idx, value in enumerate(scores.tolist()) if float(value) > 0.0]
+        scored.sort(key=lambda item: (-item[1], item[0]))
+        return [int(idx) for idx, _value in scored[:budget]]
 
     def _build_avail_masks(self, _token) -> List[set[int]] | None:
         """Promote prefetched experts and return per-layer availability masks."""
@@ -339,9 +412,14 @@ class ChronosMLXInferenceEngine:
             stats = dict(self._runtime_stats)
         total = int(stats.get("selection_hits", 0)) + int(stats.get("selection_misses", 0))
         pred_total = int(stats.get("prediction_total", 0))
+        cold_pred_total = int(stats.get("cold_prediction_total", 0))
+        all_resident_fast_path = bool(stats.get("all_resident_fast_path", False))
+        resident_hit_rate = 1.0 if all_resident_fast_path else (
+            float(stats.get("selection_hits", 0)) / max(total, 1)
+        )
         return {
-            "resident_hit_rate": round(float(stats.get("selection_hits", 0)) / max(total, 1), 4),
-            "cache_hit_rate": round(float(stats.get("selection_hits", 0)) / max(total, 1), 4),
+            "resident_hit_rate": round(resident_hit_rate, 4),
+            "cache_hit_rate": round(resident_hit_rate, 4),
             "cache_hits": int(stats.get("selection_hits", 0)),
             "cache_misses": int(stats.get("selection_misses", 0)),
             "expert_selection_hits": int(stats.get("selection_hits", 0)),
@@ -349,6 +427,10 @@ class ChronosMLXInferenceEngine:
             "prediction_hit_rate": round(float(stats.get("prediction_hits", 0)) / max(pred_total, 1), 4),
             "prediction_hits": int(stats.get("prediction_hits", 0)),
             "prediction_total": pred_total,
+            "cold_miss_predict_hit_rate": round(float(stats.get("cold_prediction_hits", 0)) / max(cold_pred_total, 1), 4),
+            "cold_prediction_hits": int(stats.get("cold_prediction_hits", 0)),
+            "cold_prediction_total": cold_pred_total,
+            "all_resident_fast_path": all_resident_fast_path,
             "resident_vram_hits": int(stats.get("resident_vram_hits", 0)),
             "resident_ram_hits": int(stats.get("resident_ram_hits", 0)),
             "prefetch_queue_drops": int(stats.get("prefetch_queue_drops", 0)),

{project_chronos-0.2.0 → project_chronos-0.2.1}/chronos/mlx/training/trainer.py

@@ -101,6 +101,19 @@ def _temporal_locality_loss(router_mean: mx.array) -> mx.array:
     return (diff * diff).sum(axis=-1).mean()
 
 
+def _router_locality_loss(router_mean: mx.array, num_experts_per_tok: int) -> mx.array:
+    if router_mean.shape[1] <= 1:
+        return mx.zeros(())
+    B, _S, E = router_mean.shape
+    top_k = max(1, min(int(num_experts_per_tok or 1), E))
+    prev = mx.stop_gradient(router_mean[:, :-1, :])
+    cur = router_mean[:, 1:, :]
+    target_ids = mx.argpartition(-prev, kth=top_k - 1, axis=-1)[:, :, :top_k]
+    target_mask = (target_ids[..., None] == mx.arange(E)).astype(cur.dtype).sum(axis=-2)
+    mass = (cur * target_mask).sum(axis=-1)
+    return (-mx.log(mx.maximum(mass / float(top_k), 1e-9))).mean()
+
+
 def _load_balance_loss(router_4d: mx.array, num_experts_per_tok: int) -> mx.array:
     B, S, L, E = router_4d.shape
     top_k = max(1, min(int(num_experts_per_tok or 1), E))
@@ -171,6 +184,96 @@ def _lookahead_topk_hit_loss(
     return total / max(terms, 1)
 
 
+def _lookahead_union_loss(
+    lookahead_probs: mx.array | None,
+    teacher_router_probs: mx.array | None,
+    lookahead_steps: int,
+    num_experts_per_tok: int,
+) -> mx.array:
+    if lookahead_probs is None:
+        return mx.zeros(())
+    if teacher_router_probs is None:
+        return mx.zeros((), dtype=lookahead_probs.dtype)
+    B, S, Kp1, E = lookahead_probs.shape
+    K = min(int(lookahead_steps or 0), Kp1 - 1)
+    top_k = max(1, min(int(num_experts_per_tok or 1), E))
+    if K <= 0 or S <= K:
+        return mx.zeros((), dtype=lookahead_probs.dtype)
+    teacher = mx.stop_gradient(teacher_router_probs.astype(mx.float32))
+    union_target = mx.zeros((B, S - K, E), dtype=mx.float32)
+    for k in range(1, K + 1):
+        target_ids = mx.argpartition(-teacher[:, k:k + S - K, :], kth=top_k - 1, axis=-1)[:, :, :top_k]
+        target_mask = (target_ids[..., None] == mx.arange(E)).astype(mx.float32).sum(axis=-2)
+        union_target = mx.maximum(union_target, target_mask)
+    pred_union = lookahead_probs[:, :S - K, 1:K + 1, :].astype(mx.float32).max(axis=2)
+    mass = (pred_union * union_target).sum(axis=-1)
+    target_count = mx.maximum(union_target.sum(axis=-1), 1.0)
+    return (-mx.log(mx.maximum(mass / target_count, 1e-9))).mean()
+
+
+def _offload_training_metrics(
+    router_mean: mx.array | None,
+    lookahead_probs: mx.array | None,
+    lookahead_steps: int,
+    num_experts_per_tok: int,
+) -> dict[str, float]:
+    if router_mean is None:
+        return {}
+    B, S, E = router_mean.shape
+    top_k = max(1, min(int(num_experts_per_tok or 1), E))
+    top_ids = mx.argpartition(-router_mean, kth=top_k - 1, axis=-1)[:, :, :top_k]
+    mx.eval(top_ids)
+    import numpy as np
+
+    ids_np = np.array(top_ids)
+    out = {
+        "expert_working_set": float(len(set(int(v) for v in ids_np.reshape(-1)))),
+        "router_adjacent_jaccard": 1.0,
+        "lookahead_topk_recall": 0.0,
+        "lookahead_union_recall": 0.0,
+    }
+    if S > 1:
+        vals = []
+        for b in range(B):
+            for t in range(S - 1):
+                a = set(int(v) for v in ids_np[b, t])
+                c = set(int(v) for v in ids_np[b, t + 1])
+                vals.append(len(a & c) / max(len(a | c), 1))
+        out["router_adjacent_jaccard"] = float(sum(vals) / max(len(vals), 1))
+    if lookahead_probs is not None and lookahead_steps > 0 and S > 1:
+        K = min(int(lookahead_steps), int(lookahead_probs.shape[2]) - 1)
+        pred = lookahead_probs.astype(mx.float32)
+        mx.eval(pred)
+        pred_np = np.array(pred)
+        recalls = []
+        union_recalls = []
+        pred_budget = max(top_k, min(E, top_k * max(1, K)))
+        for k in range(1, K + 1):
+            if S - k <= 0:
+                continue
+            pred_ids = np.argpartition(-pred_np[:, :-k, k, :], kth=top_k - 1, axis=-1)[:, :, :top_k]
+            for b in range(B):
+                for t in range(S - k):
+                    target = set(int(v) for v in ids_np[b, t + k])
+                    got = set(int(v) for v in pred_ids[b, t])
+                    recalls.append(len(target & got) / float(top_k))
+        if K > 0 and S - K > 0:
+            pred_union_ids = np.argpartition(-pred_np[:, :S - K, 1:K + 1, :].max(axis=2),
+                                             kth=pred_budget - 1, axis=-1)[:, :, :pred_budget]
+            for b in range(B):
+                for t in range(S - K):
+                    target = set()
+                    for k in range(1, K + 1):
+                        target.update(int(v) for v in ids_np[b, t + k])
+                    got = set(int(v) for v in pred_union_ids[b, t])
+                    union_recalls.append(len(target & got) / max(len(target), 1))
+        if recalls:
+            out["lookahead_topk_recall"] = float(sum(recalls) / len(recalls))
+        if union_recalls:
+            out["lookahead_union_recall"] = float(sum(union_recalls) / len(union_recalls))
+    return out
+
+
 def _router_kl_anchor(current: mx.array, reference: mx.array | None, lambda_anchor: float) -> mx.array:
     if reference is None or lambda_anchor <= 0.0:
         return mx.zeros((), dtype=current.dtype)
@@ -184,6 +287,7 @@ def _router_kl_anchor(current: mx.array, reference: mx.array | None, lambda_anch
 
 
 def _chronos_regularized_loss(model, base_loss: mx.array, lookahead_probs: mx.array | None, config, router_ref=None) -> mx.array:
+    setattr(model, "_last_lookahead_probs", lookahead_probs)
     loss = base_loss.astype(mx.float32)
     router_4d = _collect_router_probs(model)
     if router_4d is None:
@@ -202,6 +306,12 @@ def _chronos_regularized_loss(model, base_loss: mx.array, lookahead_probs: mx.ar
         lambda_temporal = float(getattr(config, "lambda_temporal", 0.0) or 0.0)
         if lambda_temporal > 0.0:
             loss = loss + lambda_temporal * _temporal_locality_loss(router_mean)
+        lambda_locality = float(getattr(config, "lambda_router_locality", 0.0) or 0.0)
+        if lambda_locality > 0.0:
+            loss = loss + lambda_locality * _router_locality_loss(
+                router_mean,
+                int(getattr(config, "num_experts_per_tok", 1) or 1),
+            )
 
         lookahead_steps = int(getattr(config, "lookahead_steps", 0) or 0)
         lambda_lookahead = float(getattr(config, "lambda_lookahead", 0.0) or 0.0)
@@ -220,6 +330,14 @@ def _chronos_regularized_loss(model, base_loss: mx.array, lookahead_probs: mx.ar
                     lookahead_steps,
                     int(getattr(config, "num_experts_per_tok", 1) or 1),
                 )
+            lambda_union = float(getattr(config, "lambda_lookahead_union", 0.0) or 0.0)
+            if lambda_union > 0.0:
+                loss = loss + lambda_union * _lookahead_union_loss(
+                    lookahead_probs,
+                    teacher,
+                    lookahead_steps,
+                    int(getattr(config, "num_experts_per_tok", 1) or 1),
+                )
 
         loss = loss + _router_kl_anchor(
             router_mean,
@@ -824,6 +942,14 @@ def run_mlx_stage(
                 tps = delta_steps / max(now - last_log_t, 1e-6)
                 last_log_t = now
                 last_log_step = steps
+                router_4d = _collect_router_probs(model)
+                router_mean = router_4d.mean(axis=2) if router_4d is not None else None
+                offload_metrics = _offload_training_metrics(
+                    router_mean,
+                    getattr(model, "_last_lookahead_probs", None),
+                    int(getattr(config, "lookahead_steps", 0) or 0),
+                    int(getattr(config, "num_experts_per_tok", 1) or 1),
+                )
                 _emit({
                     "event": "step",
                     "stage": stage,
@@ -840,6 +966,7 @@ def run_mlx_stage(
                     "dtype": dtype_name,
                     "lr": float(optimizer.learning_rate.item()) if "learning_rate" in optimizer.state else base_lr * lr_scale,
                     "rollbacks": rollback_count,
+                    **offload_metrics,
                 })
             if max_steps is not None and steps >= int(max_steps):
                 break

{project_chronos-0.2.0 → project_chronos-0.2.1}/chronos/model/checkpoint.py

@@ -48,10 +48,13 @@ CHECKPOINT_CONFIG_FIELDS = (
     "lambda_temporal",
     "lambda_lookahead",
     "lambda_lookahead_topk",
+    "lambda_lookahead_union",
+    "lambda_router_locality",
     "lambda_router_anchor",
     "fallback_mask_prob",
     "offload_miss_policy_default",
     "recommended_resident_experts",
+    "offload_working_set_experts",
     "vram_budget_gb",
     "pinned_memory_max_fraction",
     "storage_format",

{project_chronos-0.2.0 → project_chronos-0.2.1}/chronos/model/config.py

@@ -19,7 +19,7 @@ class ChronosConfig(MiniMindConfig):
         self.lookahead_steps: int = kwargs.pop("lookahead_steps", 2)
         self.num_shared_experts: int = kwargs.pop("num_shared_experts", 1)
         self.lambda_balance: float = kwargs.pop("lambda_balance", 5e-4)
-        self.lambda_temporal: float = kwargs.pop("lambda_temporal", 1e-3)
+        self.lambda_temporal: float = kwargs.pop("lambda_temporal", 3e-3)
         self.vram_budget_gb: float = kwargs.pop("vram_budget_gb", 4.0)
         self.prefetch_depth: int = kwargs.pop("prefetch_depth", 2)
         # Hybrid attention params
@@ -42,7 +42,9 @@ class ChronosConfig(MiniMindConfig):
         # Auxiliary hit-rate loss for the lookahead head. The soft CE above
         # learns distribution shape; this term directly rewards future top-k
         # expert-set recall, which is what offload prefetch needs.
-        self.lambda_lookahead_topk: float = kwargs.pop("lambda_lookahead_topk", 0.05)
+        self.lambda_lookahead_topk: float = kwargs.pop("lambda_lookahead_topk", 0.15)
+        self.lambda_lookahead_union: float = kwargs.pop("lambda_lookahead_union", 0.05)
+        self.lambda_router_locality: float = kwargs.pop("lambda_router_locality", 0.02)
         # M4: weight on the router KL anchor (alignment stages only).
         # Prevents SFT/DPO/ORPO/GRPO gradients from drifting routing away
         # from the pretrained distribution that the cluster layout was
@@ -52,7 +54,7 @@ class ChronosConfig(MiniMindConfig):
         # Optional training-time miss simulation. When >0, MoE layers randomly
         # route selected experts through the shared fallback while still adding
         # the shared residual, matching the SSD/DRAM offload inference formula.
-        self.fallback_mask_prob: float = kwargs.pop("fallback_mask_prob", 0.0)
+        self.fallback_mask_prob: float = kwargs.pop("fallback_mask_prob", 0.08)
         # Runtime offload policy. "on_demand" is the predictive low-latency
         # path: cold SSD misses are queued for RAM and covered by shared
         # fallback this token. "sync_on_demand" is the diagnostic path that
@@ -60,6 +62,7 @@ class ChronosConfig(MiniMindConfig):
         # "fallback_diagnostic" keeps the pure shared-fallback stress path.
         self.offload_miss_policy_default: str = kwargs.pop("offload_miss_policy_default", "on_demand")
         self.recommended_resident_experts: int | None = kwargs.pop("recommended_resident_experts", None)
+        self.offload_working_set_experts: int | None = kwargs.pop("offload_working_set_experts", None)
         # M7: regularization defaults. MiniMindConfig defaults dropout=0.0
         # and weight_decay isn't even a config field — both major causes
         # of overfitting on small corpora. Override here.