Project-Chronos 0.1.1__py3-none-any.whl

chronos/eval/io_profiler.py

@@ -0,0 +1,232 @@
+"""
+chronos/eval/io_profiler.py
+
+Phase 1 validation tool: measures LookaheadRouter prediction accuracy
+and temporal routing smoothness offline, without running full inference.
+
+Usage:
+    python -m chronos.eval.io_profiler \
+        --model_path /path/to/chronos_checkpoint \
+        --data_path  /path/to/eval.jsonl \
+        --device     cuda
+
+Outputs:
+    - t+1 Top-1 accuracy (target: >= 85%)
+    - t+2 Top-1 accuracy (target: >= 75%)
+    - Mean L2 routing shift (lower = more cache-friendly)
+    - Cache hit rate estimate given VRAM budget
+"""
+import sys
+import chronos.deps  # ensure minimind on sys.path
+
+import argparse
+import json
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.moe_chronos import ChronosMOEFeedForward
+
+
+def collect_ground_truth_routing(model, input_ids, device):
+    """
+    Run a full forward pass and collect per-layer, per-step top-1 expert indices.
+    Returns: [B, T, num_layers] int tensor of selected expert indices.
+    """
+    model.eval()
+    input_ids = input_ids.to(device)
+    with torch.no_grad():
+        _, _ = model(input_ids, use_cache=False)
+
+    layers = [l for l in model.model.layers if isinstance(l.mlp, ChronosMOEFeedForward)]
+    if not layers:
+        return None
+
+    # last_router_probs: [B, S, E] per layer
+    probs_per_layer = [l.mlp.last_router_probs for l in layers]  # list of [B, S, E]
+    stacked = torch.stack(probs_per_layer, dim=2)                 # [B, S, L, E]
+    top1 = stacked.argmax(dim=-1)                                 # [B, S, L]
+    return top1
+
+
+def evaluate_lookahead_accuracy(model, dataloader, device, lookahead_steps=2):
+    """
+    For each sequence, compare LookaheadRouter's predicted expert at t+k
+    against the ground-truth expert selected at t+k.
+
+    Returns dict with accuracy per step and mean L2 routing shift.
+    """
+    model.eval()
+    correct = {k: 0 for k in range(1, lookahead_steps + 1)}
+    total = 0
+    l2_shifts = []
+
+    for input_ids, _ in dataloader:
+        input_ids = input_ids.to(device)
+        B, T = input_ids.shape
+        if T < lookahead_steps + 2:
+            continue
+
+        # Ground truth routing for full sequence
+        gt_routing = collect_ground_truth_routing(model, input_ids, device)
+        if gt_routing is None:
+            continue
+        # gt_routing: [B, T, L] — use layer 0 as reference
+        gt_top1 = gt_routing[:, :, 0]  # [B, T]
+
+        # LookaheadRouter predictions from block 0 hidden states
+        with torch.no_grad():
+            hidden_after_block0 = None
+
+            def _hook(module, inp, out):
+                nonlocal hidden_after_block0
+                hidden_after_block0 = out[0].detach()
+
+            handle = model.model.layers[0].register_forward_hook(_hook)
+            model(input_ids, use_cache=False)
+            handle.remove()
+
+        if hidden_after_block0 is None:
+            continue
+
+        lookahead_probs = model.model.lookahead_router(hidden_after_block0)
+        # lookahead_probs: [B, T, K+1, E], index 1..K = future predictions
+
+        for k in range(1, lookahead_steps + 1):
+            if T <= k:
+                continue
+            # Predicted expert at step t for future t+k
+            pred_top1 = lookahead_probs[:, :-k, k, :].argmax(dim=-1)  # [B, T-k]
+            gt_at_tk = gt_top1[:, k:]                                   # [B, T-k]
+            correct[k] += (pred_top1 == gt_at_tk).sum().item()
+            total_k = pred_top1.numel()
+            if k == 1:
+                total += total_k
+
+        # L2 routing shift between adjacent steps
+        # Use mean router probs across layers
+        from chronos.model.moe_chronos import ChronosMOEFeedForward
+        moe_layers = [l for l in model.model.layers if isinstance(l.mlp, ChronosMOEFeedForward)]
+        if moe_layers:
+            probs = torch.stack([l.mlp.last_router_probs for l in moe_layers], dim=2).mean(dim=2)
+            diff = probs[:, 1:, :] - probs[:, :-1, :]
+            l2_shifts.append((diff ** 2).sum(dim=-1).mean().item())
+
+    accuracies = {}
+    for k in range(1, lookahead_steps + 1):
+        denom = total if k == 1 else max(total, 1)
+        accuracies[f"top1_acc_t+{k}"] = correct[k] / max(denom, 1)
+
+    mean_l2 = sum(l2_shifts) / max(len(l2_shifts), 1)
+    return {**accuracies, "mean_l2_routing_shift": mean_l2}
+
+
+def estimate_cache_hit_rate(model, dataloader, device, vram_budget_gb=4.0):
+    """
+    Simulate a VRAM LRU cache of expert weights and estimate hit rate
+    given the LookaheadRouter's prefetch predictions.
+    """
+    import collections
+
+    # Estimate single expert weight size in GB
+    moe_layers = [l for l in model.model.layers if isinstance(l.mlp, ChronosMOEFeedForward)]
+    if not moe_layers:
+        return {"cache_hit_rate": 1.0}
+
+    expert_params = sum(p.numel() for p in moe_layers[0].mlp.experts[0].parameters())
+    expert_size_gb = expert_params * 2 / (1024 ** 3)  # fp16
+    num_experts = model.config.num_experts
+    cache_capacity = max(1, int(vram_budget_gb / expert_size_gb))
+
+    hits, total = 0, 0
+    for input_ids, _ in dataloader:
+        input_ids = input_ids.to(device)
+        gt_routing = collect_ground_truth_routing(model, input_ids, device)
+        if gt_routing is None:
+            continue
+        gt_top1 = gt_routing[:, :, 0]  # [B, T]
+
+        lru = collections.OrderedDict()
+        for t in range(gt_top1.shape[1]):
+            for b in range(gt_top1.shape[0]):
+                eid = gt_top1[b, t].item()
+                if eid in lru:
+                    lru.move_to_end(eid)
+                    hits += 1
+                else:
+                    if len(lru) >= cache_capacity:
+                        lru.popitem(last=False)
+                    lru[eid] = True
+                total += 1
+
+    return {
+        "cache_hit_rate": hits / max(total, 1),
+        "cache_capacity_experts": cache_capacity,
+        "expert_size_gb": round(expert_size_gb, 4),
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Chronos Phase 1 Validation")
+    parser.add_argument("--model_path", type=str, default=None)
+    parser.add_argument("--data_path", type=str, required=True)
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--max_seq_len", type=int, default=256)
+    parser.add_argument("--vram_budget_gb", type=float, default=4.0)
+    parser.add_argument("--hidden_size", type=int, default=512)
+    parser.add_argument("--num_hidden_layers", type=int, default=8)
+    parser.add_argument("--num_experts", type=int, default=4)
+    args = parser.parse_args()
+
+    config = ChronosConfig(
+        hidden_size=args.hidden_size,
+        num_hidden_layers=args.num_hidden_layers,
+        num_experts=args.num_experts,
+        use_moe=True,
+        vram_budget_gb=args.vram_budget_gb,
+    )
+    model = ChronosForCausalLM(config).to(args.device)
+
+    if args.model_path:
+        weights = torch.load(args.model_path, map_location=args.device)
+        model.load_state_dict(weights, strict=False)
+        print(f"Loaded weights from {args.model_path}")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        chronos.deps.get_tokenizer_path()
+    )
+
+    from dataset.lm_dataset import PretrainDataset
+    dataset = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
+
+    print("\n=== Phase 1 Validation: LookaheadRouter Accuracy ===")
+    acc_results = evaluate_lookahead_accuracy(model, loader, args.device)
+    for k, v in acc_results.items():
+        status = ""
+        if "t+1" in k:
+            status = "✓ PASS" if v >= 0.85 else "✗ FAIL (target: ≥85%)"
+        elif "t+2" in k:
+            status = "✓ PASS" if v >= 0.75 else "✗ FAIL (target: ≥75%)"
+        print(f"  {k}: {v:.4f}  {status}")
+
+    print("\n=== Cache Hit Rate Estimate ===")
+    cache_results = estimate_cache_hit_rate(model, loader, args.device, args.vram_budget_gb)
+    for k, v in cache_results.items():
+        print(f"  {k}: {v}")
+
+    print("\n=== Summary ===")
+    t1 = acc_results.get("top1_acc_t+1", 0)
+    t2 = acc_results.get("top1_acc_t+2", 0)
+    phase1_pass = t1 >= 0.85 and t2 >= 0.75
+    print(f"  Phase 1 Gate: {'PASS ✓' if phase1_pass else 'FAIL ✗'}")
+    print(f"  t+1 acc={t1:.1%}  t+2 acc={t2:.1%}  "
+          f"L2 shift={acc_results['mean_l2_routing_shift']:.4f}")
+
+
+if __name__ == "__main__":
+    main()

chronos/io/__init__.py

@@ -0,0 +1,13 @@
+from .expert_store import ExpertStore, LRUCache
+from .async_prefetcher import AsyncPrefetcher, PrefetchScheduler
+from .storage import ClusterStorage, ClusterManifest, MANIFEST_FILENAME
+from .cluster_layout import (
+    collect_activation_log,
+    build_cooccurrence_matrix,
+    cluster_experts_greedy,
+    try_louvain_clustering,
+    build_cluster_layout,
+    repack_expert_weights,
+    repack_expert_weights_safetensors,
+    load_cluster_layout,
+)

chronos/io/async_prefetcher.py

@@ -0,0 +1,220 @@
+"""
+chronos/io/async_prefetcher.py
+
+Asynchronous expert weight prefetcher driven by LookaheadRouter predictions.
+
+Architecture:
+- Background daemon thread runs a prefetch loop
+- Main thread submits prefetch requests via a queue
+- Prefetcher loads SSD → RAM (pinned) ahead of time
+- When main thread needs an expert, it's already warm in RAM → fast H2D
+
+The key insight: LookaheadRouter predicts t+1/t+2 expert needs,
+giving us 1-2 decode steps (~10-50ms) of lead time to hide SSD latency.
+"""
+import threading
+import queue
+import time
+from typing import List, Optional, Callable
+
+from chronos.io.expert_store import ExpertStore
+
+try:
+    from chronos.runtime.metrics import safe_record as _metric
+except Exception:
+    def _metric(name, value):  # fallback
+        pass
+
+
+class AsyncPrefetcher:
+    """
+    Background thread that prefetches expert weights from SSD → RAM
+    based on LookaheadRouter predictions.
+
+    Usage:
+        prefetcher = AsyncPrefetcher(expert_store)
+        prefetcher.start()
+
+        # After each decode step, submit next predictions:
+        prefetcher.submit([expert_id_t1, expert_id_t2])
+
+        # Before using expert at step t:
+        prefetcher.wait_for(expert_id_t)  # usually already done
+
+        prefetcher.stop()
+    """
+
+    def __init__(
+        self,
+        expert_store: ExpertStore,
+        queue_depth: int = 8,
+        log_fn: Optional[Callable] = None,
+    ):
+        self.store = expert_store
+        self.log_fn = log_fn or (lambda m: None)
+        self._request_queue: queue.Queue = queue.Queue(maxsize=queue_depth)
+        self._done_set: set = set()
+        self._done_lock = threading.Lock()
+        self._stop_event = threading.Event()
+        self._thread: Optional[threading.Thread] = None
+        self._stats = {"prefetch_hits": 0, "prefetch_misses": 0, "total_requests": 0}
+
+    def start(self):
+        self._stop_event.clear()
+        self._thread = threading.Thread(
+            target=self._prefetch_loop,
+            daemon=True,
+            name="chronos-prefetcher",
+        )
+        self._thread.start()
+
+    def stop(self):
+        self._stop_event.set()
+        # Unblock the queue
+        try:
+            self._request_queue.put_nowait(None)
+        except queue.Full:
+            pass
+        if self._thread:
+            self._thread.join(timeout=2.0)
+
+    def submit(self, expert_ids: List[int]):
+        """
+        Submit a list of expert IDs to prefetch (non-blocking).
+        Called after each decode step with LookaheadRouter predictions.
+        """
+        if not expert_ids:
+            return
+        try:
+            self._request_queue.put_nowait(expert_ids)
+        except queue.Full:
+            # Drop oldest request to make room
+            try:
+                self._request_queue.get_nowait()
+                self._request_queue.put_nowait(expert_ids)
+            except (queue.Empty, queue.Full):
+                pass
+
+    def wait_for(self, expert_id: int, timeout: float = 0.05) -> bool:
+        """
+        Wait until expert_id is in RAM (or VRAM). Returns True if ready.
+        In practice this should return immediately if prefetch worked.
+        """
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            if (self.store.ram_lru.contains(expert_id) or
+                    self.store.vram_lru.contains(expert_id)):
+                with self._done_lock:
+                    self._stats["prefetch_hits"] += 1
+                return True
+            time.sleep(0.001)
+        with self._done_lock:
+            self._stats["prefetch_misses"] += 1
+        return False
+
+    def _prefetch_loop(self):
+        while not self._stop_event.is_set():
+            try:
+                expert_ids = self._request_queue.get(timeout=0.01)
+            except queue.Empty:
+                continue
+
+            if expert_ids is None:  # stop sentinel
+                break
+
+            with self._done_lock:
+                self._stats["total_requests"] += len(expert_ids)
+
+            # Filter out already-warm experts
+            to_fetch = [
+                eid for eid in expert_ids
+                if not self.store.ram_lru.contains(eid)
+                and not self.store.vram_lru.contains(eid)
+            ]
+            if to_fetch:
+                try:
+                    self.store.prefetch_to_ram(to_fetch)
+                except Exception as e:
+                    self.log_fn(f"[prefetcher] error fetching {to_fetch}: {e}")
+
+    @property
+    def stats(self) -> dict:
+        with self._done_lock:
+            s = dict(self._stats)
+        total = s["prefetch_hits"] + s["prefetch_misses"]
+        s["hit_rate"] = s["prefetch_hits"] / max(total, 1)
+        _metric("prefetch_hit_rate", s["hit_rate"])
+        return s
+
+
+class PrefetchScheduler:
+    """
+    Integrates LookaheadRouter predictions with AsyncPrefetcher.
+
+    Called once per decode step to:
+    1. Extract predicted expert IDs from lookahead_probs
+    2. Submit them to the prefetcher
+    3. Promote the current-step expert from RAM → VRAM
+    """
+
+    def __init__(self, prefetcher: AsyncPrefetcher, expert_store: ExpertStore):
+        self.prefetcher = prefetcher
+        self.store = expert_store
+
+    def step(
+        self,
+        lookahead_probs,  # [B, S, K+1, E] from LookaheadRouter
+        current_step_expert_ids: List[int],
+    ):
+        """
+        Legacy single-call API: promote current experts (blocking), then
+        submit future predictions. Kept for backward compatibility.
+        Prefer `promote_current` + `prefetch_only` for the M3 pipeline.
+        """
+        self.promote_current(current_step_expert_ids, blocking=True)
+        self.prefetch_only(lookahead_probs)
+
+    def promote_current(
+        self,
+        current_step_expert_ids: List[int],
+        blocking: bool = False,
+    ) -> List[int]:
+        """
+        Move the experts needed right now from RAM → VRAM.
+
+        With ``blocking=False`` (M3 default) the H2D copies are issued on
+        the dedicated stream and return immediately; the caller is
+        responsible for waiting via ``expert_store.wait_for_experts``.
+        Returns the list of experts that were actually newly promoted
+        (already-resident experts are skipped).
+        """
+        newly = []
+        for eid in current_step_expert_ids:
+            if not self.store.vram_lru.contains(eid):
+                if self.store.promote_to_vram(eid, blocking=blocking):
+                    newly.append(eid)
+        return newly
+
+    def prefetch_only(self, lookahead_probs) -> List[int]:
+        """
+        Submit future expert predictions to the async prefetcher.
+        Does NOT touch VRAM. Safe to call *before* the forward pass so the
+        H2D pipeline can overlap with compute.
+
+        Returns the list of unique expert IDs submitted.
+        """
+        if lookahead_probs is None:
+            return []
+        import torch
+        # lookahead_probs: [B, S, K+1, E] — take last token, steps 1..K.
+        # We submit t+1 first (highest priority), then t+2, etc., so the
+        # FIFO queue naturally orders them by recency.
+        last_token_probs = lookahead_probs[:, -1, 1:, :]  # [B, K, E]
+        K = last_token_probs.shape[1]
+        all_submitted: List[int] = []
+        for k in range(K):
+            ids_k = last_token_probs[:, k, :].argmax(dim=-1).unique().cpu().tolist()
+            if ids_k:
+                self.prefetcher.submit(ids_k)
+                all_submitted.extend(ids_k)
+        return list(dict.fromkeys(all_submitted))  # preserve order, dedup