Project-Chronos 0.1.1__py3-none-any.whl

chronos/mlx/inference.py

@@ -0,0 +1,159 @@
+"""
+chronos/mlx/inference.py — MLX inference engine for Project Chronos.
+
+Mirrors ChronosInferenceEngine (PyTorch) but uses:
+- ChronosMLXModel for forward passes
+- MLXExpertStore (unified memory, no H2D stream)
+- mx.compile() for static-graph acceleration
+- Async prefetch via Python threading (SSD → unified memory)
+
+Usage:
+    engine = ChronosMLXInferenceEngine(model, config)
+    for token_id in engine.generate("Explain how Chronos moves expert IO into prefill.", max_new_tokens=128):
+        print(tokenizer.decode([token_id]), end="", flush=True)
+"""
+import threading
+import queue
+from typing import Iterator, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from chronos.mlx.expert_store import MLXExpertStore
+
+
+class ChronosMLXInferenceEngine:
+    """
+    Streaming token-by-token inference with async expert prefetch.
+
+    On Apple Silicon the prefetch thread reads .npy files from SSD and
+    calls mx.array() to bring weights into unified memory.  The main
+    thread calls mx.eval() (promote_to_vram) before the next forward
+    pass — Metal overlaps the IO with prior computation automatically.
+    """
+
+    def __init__(self, model, config, ssd_dir: str = "./expert_cache_mlx"):
+        self.model = model
+        self.config = config
+        self.store = MLXExpertStore(model, config, ssd_dir=ssd_dir)
+        self._prefetch_q: queue.Queue = queue.Queue(maxsize=config.prefetch_depth * 2)
+        self._stop = threading.Event()
+        self._prefetch_thread = threading.Thread(
+            target=self._prefetch_loop, daemon=True, name="mlx-prefetch"
+        )
+        self._prefetch_thread.start()
+
+    # ── Prefetch background thread ────────────────────────────────
+
+    def _prefetch_loop(self):
+        while not self._stop.is_set():
+            try:
+                expert_ids: List[int] = self._prefetch_q.get(timeout=0.05)
+            except queue.Empty:
+                continue
+            self.store.prefetch_to_ram(expert_ids)
+            self._prefetch_q.task_done()
+
+    def _schedule_prefetch(self, expert_ids: List[int]):
+        try:
+            self._prefetch_q.put_nowait(expert_ids)
+        except queue.Full:
+            pass
+
+    def stop(self):
+        self._stop.set()
+
+    # ── Generation loop ───────────────────────────────────────────
+
+    def generate(
+        self,
+        input_ids: mx.array,           # [1, S] (already tokenised)
+        max_new_tokens: int = 128,
+        temperature: float = 0.85,
+        top_p: float = 0.9,
+        scheduler=None,                # Optional[PrefillScheduler] — lazy import
+    ) -> Iterator[int]:
+        """
+        Yields token IDs one at a time.
+
+        If scheduler is provided (PrefillScheduler), runs prefill-time batch
+        loading before the generation loop (dual-layer mode).  Otherwise falls
+        back to per-token LookaheadRouter-driven prefetch.
+        """
+        # ── Prefill phase: front-load expert IO ───────────────────────────
+        if scheduler is not None:
+            import torch, numpy as np
+            ids_np = np.array(input_ids.tolist(), dtype=np.int64)
+            ids_pt = torch.from_numpy(ids_np)
+            scheduler.prepare(ids_pt, device="cpu")
+            scheduler.wait()
+
+        # Prefill forward
+        logits, lookahead_probs, cache = self.model(input_ids)
+        mx.eval(logits)
+
+        # Sample from prefill
+        next_token = self._sample(logits[:, -1, :], temperature, top_p)
+        activated_ids: List[int] = []
+        yield int(next_token.item())
+
+        for _ in range(max_new_tokens - 1):
+            if scheduler is not None:
+                # Use PrefillScheduler masks (already loaded at prefill)
+                n_layers = len(self.model.layers)
+                avail_masks = [self.store.vram_availability_mask()] * n_layers
+            else:
+                # LookaheadRouter-driven prefetch
+                if lookahead_probs is not None:
+                    future_ids = self._predict_future_experts(lookahead_probs)
+                    self._schedule_prefetch(future_ids)
+                avail_masks = self._build_avail_masks(next_token)
+
+            token_in = next_token.reshape(1, 1)
+            logits, lookahead_probs, cache = self.model(
+                token_in,
+                cache=cache,
+                available_expert_masks=avail_masks,
+            )
+            mx.eval(logits)
+
+            next_token = self._sample(logits[:, -1, :], temperature, top_p)
+            mx.eval(next_token)
+            tok_id = int(next_token.item())
+            activated_ids.append(tok_id)
+            yield tok_id
+
+        # Update frequency heuristic after generation
+        if scheduler is not None:
+            scheduler.record_activation(list(set(activated_ids)))
+
+    def _predict_future_experts(self, lookahead_probs: mx.array) -> List[int]:
+        """Extract top-1 expert IDs for t+1 and t+2 from lookahead probs."""
+        # lookahead_probs: [B, S, K+1, E] — take last token, steps 1..K
+        future = lookahead_probs[0, -1, 1:, :]         # [K, E]
+        top_ids = mx.argmax(future, axis=-1).tolist()   # [K]
+        return list(set(top_ids))
+
+    def _build_avail_masks(self, _token) -> List[mx.array]:
+        """Promote prefetched experts and return per-layer availability masks."""
+        n_layers = len(self.model.layers)
+        mask = self.store.vram_availability_mask()
+        return [mask] * n_layers
+
+    @staticmethod
+    def _sample(logits: mx.array, temperature: float, top_p: float) -> mx.array:
+        if temperature == 0:
+            return mx.argmax(logits, axis=-1)
+        scaled = logits / temperature
+        probs = mx.softmax(scaled, axis=-1)
+        # top-p nucleus sampling
+        sorted_idx = mx.argsort(-probs, axis=-1)
+        sorted_probs = mx.take_along_axis(probs, sorted_idx, axis=-1)
+        cumsum = mx.cumsum(sorted_probs, axis=-1)
+        # mask tokens beyond top-p threshold
+        cutoff = mx.array(top_p)
+        nucleus = (cumsum - sorted_probs) < cutoff
+        filtered = mx.where(nucleus, sorted_probs, mx.zeros_like(sorted_probs))
+        filtered = filtered / (filtered.sum(axis=-1, keepdims=True) + 1e-9)
+        sample_idx = mx.random.categorical(mx.log(filtered + 1e-9))
+        return mx.take(sorted_idx[0], sample_idx)

chronos/mlx/model.py

@@ -0,0 +1,181 @@
+"""
+chronos/mlx/model.py — Full MLX Chronos model.
+
+Architecture mirrors the PyTorch version (ChronosForCausalLM) but uses:
+- mlx.nn.Module throughout
+- MLAAttentionMLX (even layers) / SlidingWindowAttentionMLX (odd layers)
+- ChronosMLXMOE for feed-forward
+- LookaheadRouter as a plain Linear projection
+- mx.compile()-friendly: no Python control flow on runtime shapes
+
+On Apple Silicon, mx.eval() materialises arrays into Metal shared memory.
+Prefetch = eagerly eval expert weights during the prefetch window.
+"""
+import math
+import mlx.core as mx
+import mlx.nn as nn
+
+from chronos.mlx.attention import MLAAttentionMLX, SlidingWindowAttentionMLX
+from chronos.mlx.moe import ChronosMLXMOE, FeedForwardMLX
+
+
+class RMSNormMLX(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = mx.ones((dim,))
+        self.eps = eps
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return nn.RMSNorm(x.shape[-1], eps=self.eps)(x) * self.weight
+
+
+class LookaheadRouterMLX(nn.Module):
+    """Predicts expert probabilities for t, t+1, ..., t+K steps."""
+
+    def __init__(self, config):
+        super().__init__()
+        K = getattr(config, 'lookahead_steps', 2)
+        self.lookahead_steps = K
+        self.proj = nn.Linear(
+            config.hidden_size,
+            config.num_experts * (K + 1),
+            bias=False,
+        )
+        self.num_experts = config.num_experts
+
+    def __call__(self, x: mx.array) -> mx.array:
+        B, S, H = x.shape
+        out = self.proj(x)                                             # [B, S, E*(K+1)]
+        out = out.reshape(B, S, self.lookahead_steps + 1, self.num_experts)
+        return mx.softmax(out, axis=-1)                                # [B, S, K+1, E]
+
+
+class ChronosMLXBlock(nn.Module):
+    def __init__(self, layer_id: int, config):
+        super().__init__()
+        self.input_norm  = RMSNormMLX(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_norm   = RMSNormMLX(config.hidden_size, eps=config.rms_norm_eps)
+
+        use_hybrid = getattr(config, 'use_hybrid_attention', True)
+        if use_hybrid:
+            if layer_id % 2 == 0:
+                self.self_attn = MLAAttentionMLX(config)
+            else:
+                self.self_attn = SlidingWindowAttentionMLX(config)
+        else:
+            self.self_attn = SlidingWindowAttentionMLX(config)
+
+        if getattr(config, 'use_moe', True):
+            self.mlp = ChronosMLXMOE(config)
+        else:
+            self.mlp = FeedForwardMLX(config.hidden_size, config.intermediate_size)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: mx.array = None,
+        cache=None,
+        available_expert_mask: mx.array = None,
+    ):
+        h, new_cache = self.self_attn(self.input_norm(x), mask=mask, cache=cache)
+        x = x + h
+        mlp_in = self.post_norm(x)
+        if isinstance(self.mlp, ChronosMLXMOE):
+            x = x + self.mlp(mlp_in, available_expert_mask=available_expert_mask)
+        else:
+            x = x + self.mlp(mlp_in)
+        return x, new_cache
+
+
+class ChronosMLXModel(nn.Module):
+    """
+    Full Chronos model in MLX.
+
+    forward() returns (logits [B, S, V], lookahead_probs [B, S, K+1, E])
+    Cache: list of per-layer attention caches (passed back in, None at start).
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [ChronosMLXBlock(i, config)
+                       for i in range(config.num_hidden_layers)]
+        self.norm = RMSNormMLX(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # LookaheadRouter inserted after block 0
+        self.lookahead_router = LookaheadRouterMLX(config)
+
+    def __call__(
+        self,
+        input_ids: mx.array,               # [B, S]
+        cache=None,                         # list of per-layer caches, or None
+        available_expert_masks=None,        # list of [E] float masks, or None
+    ):
+        B, S = input_ids.shape
+        x = self.embed_tokens(input_ids)   # [B, S, H]
+
+        # Causal mask — only needed for prefill (S > 1); decode (S==1) needs no mask
+        mask = None
+        if S > 1:
+            mask = nn.MultiHeadAttention.create_additive_causal_mask(S)
+            mask = mask.astype(x.dtype)
+
+        new_caches = []
+        lookahead_probs = None
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for i, layer in enumerate(self.layers):
+            avail = (available_expert_masks[i]
+                     if available_expert_masks is not None else None)
+            x, lyr_cache = layer(x, mask=mask, cache=cache[i],
+                                 available_expert_mask=avail)
+            new_caches.append(lyr_cache)
+
+            # Lookahead routing after block 0
+            if i == 0:
+                lookahead_probs = self.lookahead_router(x)
+
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        return logits, lookahead_probs, new_caches
+
+    @staticmethod
+    def from_chronos_pytorch(pt_model, config):
+        """
+        Convert a trained PyTorch ChronosForCausalLM to ChronosMLXModel.
+        Copies weights layer by layer via numpy as bridge.
+        """
+        import numpy as np
+        import torch
+
+        mlx_model = ChronosMLXModel(config)
+
+        def _to_mx(pt_tensor):
+            return mx.array(pt_tensor.detach().float().cpu().numpy())
+
+        # Embedding + head
+        sd = pt_model.state_dict()
+        mlx_model.embed_tokens.weight = _to_mx(sd["model.embed_tokens.weight"])
+        mlx_model.lm_head.weight = _to_mx(sd["lm_head.weight"])
+        mlx_model.norm.weight = _to_mx(sd["model.norm.weight"])
+
+        for i in range(config.num_hidden_layers):
+            prefix = f"model.layers.{i}"
+            blk = mlx_model.layers[i]
+            blk.input_norm.weight = _to_mx(sd[f"{prefix}.input_layernorm.weight"])
+            blk.post_norm.weight  = _to_mx(sd[f"{prefix}.post_feedforward_layernorm.weight"])
+            # Attention weights are mapped by name if they match
+            attn = blk.self_attn
+            for attr in vars(attn):
+                pt_key = f"{prefix}.self_attn.{attr}.weight"
+                if pt_key in sd and hasattr(getattr(attn, attr, None), '__class__'):
+                    try:
+                        setattr(attn, attr,
+                                type(getattr(attn, attr))(_to_mx(sd[pt_key])))
+                    except Exception:
+                        pass
+        return mlx_model

chronos/mlx/moe.py

@@ -0,0 +1,102 @@
+"""
+chronos/mlx/moe.py — MLX-native ChronosMOEFeedForward.
+
+Soft gating is expressed as weighted sum (compile-safe):
+    out = avail[i] * expert_out + (1 - avail[i]) * shared_out
+
+mx.compile() traces a single graph — no Python branches on mask values.
+"""
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class FeedForwardMLX(nn.Module):
+    """SwiGLU FFN matching minimind's FeedForward structure."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj   = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class ChronosMLXMOE(nn.Module):
+    """
+    MLX-native Mixture-of-Experts with:
+    - LookaheadRouter output stored as last_router_probs
+    - Compile-safe soft gating (avail mask → float multiply)
+    - Shared expert fallback at zero I/O cost (unified memory)
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        H = config.hidden_size
+        I = config.moe_intermediate_size
+
+        self.gate = nn.Linear(H, config.num_experts, bias=False)
+        self.experts = [FeedForwardMLX(H, I) for _ in range(config.num_experts)]
+        self.shared_experts = [
+            FeedForwardMLX(H, I) for _ in range(config.num_shared_experts)
+        ]
+        self.last_router_probs = None  # [B, S, E] stored after each forward
+
+    def _shared_out(self, x: mx.array) -> mx.array:
+        return sum(e(x) for e in self.shared_experts) / len(self.shared_experts)
+
+    def __call__(
+        self,
+        x: mx.array,
+        available_expert_mask: mx.array = None,
+    ) -> mx.array:
+        """
+        x: [B, S, H]
+        available_expert_mask: [num_experts] float (1.0 = in GPU, 0.0 = miss)
+                                or None for training (all available).
+        """
+        B, S, H = x.shape
+        x_flat = x.reshape(-1, H)                        # [N, H]
+
+        scores = mx.softmax(self.gate(x_flat), axis=-1)  # [N, E]
+        self.last_router_probs = scores.reshape(B, S, self.num_experts)
+
+        # top-k routing
+        topk_idx = mx.argpartition(-scores, kth=self.num_experts_per_tok - 1, axis=-1)
+        topk_idx = topk_idx[:, :self.num_experts_per_tok]              # [N, K]
+        topk_weight = mx.take_along_axis(scores, topk_idx, axis=-1)   # [N, K]
+        if self.norm_topk_prob:
+            topk_weight = topk_weight / (topk_weight.sum(axis=-1, keepdims=True) + 1e-20)
+
+        y = mx.zeros_like(x_flat)
+
+        if available_expert_mask is None:
+            # Training: standard MoE scatter
+            for i, expert in enumerate(self.experts):
+                # mask over tokens where expert i is selected
+                sel = (topk_idx == i)            # [N, K] bool
+                tok_mask = sel.any(axis=-1)       # [N] bool
+                w = mx.where(sel, topk_weight, 0.0).sum(axis=-1, keepdims=True)  # [N, 1]
+                contrib = expert(x_flat) * w * tok_mask[:, None]
+                y = y + contrib
+        else:
+            # Inference: compile-safe soft gating
+            avail = available_expert_mask.astype(x_flat.dtype)  # [E] float
+            shared_out = self._shared_out(x_flat)               # [N, H]
+
+            for i, expert in enumerate(self.experts):
+                sel = (topk_idx == i)
+                tok_mask = sel.any(axis=-1)
+                w = mx.where(sel, topk_weight, 0.0).sum(axis=-1, keepdims=True)
+
+                expert_out  = expert(x_flat)
+                # avail[i]==1 → expert; avail[i]==0 → shared fallback
+                blended = avail[i] * expert_out + (1.0 - avail[i]) * shared_out
+                contrib = blended * w * tok_mask[:, None]
+                y = y + contrib
+
+        return y.reshape(B, S, H)

chronos/model/__init__.py

@@ -0,0 +1,6 @@
+from .config import ChronosConfig
+from .lookahead_router import LookaheadRouter
+from .moe_chronos import ChronosMOEFeedForward
+from .model_chronos import ChronosForCausalLM
+from .temporal_loss import temporal_locality_loss, total_loss
+from .hybrid_attention import MLAAttention, SlidingWindowAttention, make_attention

chronos/model/config.py

@@ -0,0 +1,54 @@
+import sys
+import chronos.deps  # ensure minimind on sys.path
+
+from model.model_minimind import MiniMindConfig
+
+
+class ChronosConfig(MiniMindConfig):
+    """
+    Extends MiniMindConfig with Project Chronos-specific fields:
+    - lookahead_steps: how many future token steps to predict expert routing for
+    - num_shared_experts: experts always resident in VRAM as fallback
+    - lambda_balance / lambda_temporal: loss coefficients
+    - vram_budget_gb: VRAM cap for expert cache
+    - prefetch_depth: async prefetch queue depth
+    """
+    model_type = "chronos"
+
+    def __init__(self, **kwargs):
+        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.vram_budget_gb: float = kwargs.pop("vram_budget_gb", 4.0)
+        self.prefetch_depth: int = kwargs.pop("prefetch_depth", 2)
+        # Hybrid attention params
+        self.use_hybrid_attention: bool = kwargs.pop("use_hybrid_attention", True)
+        self.kv_latent_dim: int = kwargs.pop("kv_latent_dim", 64)   # MLA compression dim
+        self.rope_dim: int = kwargs.pop("rope_dim", 32)              # RoPE portion of head_dim
+        self.sliding_window_size: int = kwargs.pop("sliding_window_size", 2048)
+        # Pinned memory safety: max RAM for expert buffers (fraction of physical RAM)
+        self.pinned_memory_max_fraction: float = kwargs.pop("pinned_memory_max_fraction", 0.25)
+        # Storage backend for expert weights on SSD.
+        # "safetensors" packs each Louvain cluster into one .ctsr file +
+        # cluster_manifest.json (mmap-friendly, sequential reads).
+        # "pt" keeps the legacy one-pickle-per-expert layout.
+        self.storage_format: str = kwargs.pop("storage_format", "safetensors")
+        # Optional path to a precomputed cluster_manifest.json. When set,
+        # CacheManager will attach this layout instead of the one in ssd_dir.
+        self.cluster_manifest_path = kwargs.pop("cluster_manifest_path", None)
+        # M2: weight on the lookahead supervision loss
+        self.lambda_lookahead: float = kwargs.pop("lambda_lookahead", 0.1)
+        # 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
+        # optimized for. Typical values:
+        #   pretrain: 0.0,  SFT: 0.01,  DPO/ORPO/GRPO: 0.1
+        self.lambda_router_anchor: float = kwargs.pop("lambda_router_anchor", 0.0)
+        # 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.
+        kwargs.setdefault("dropout", 0.1)
+        # Force MoE on
+        kwargs.setdefault("use_moe", True)
+        super().__init__(**kwargs)

chronos/model/hf_io.py

@@ -0,0 +1,149 @@
+"""
+chronos/model/hf_io.py
+
+HuggingFace save/load helpers for ChronosForCausalLM.
+
+``ChronosForCausalLM`` already inherits ``transformers.PreTrainedModel`` so
+``.from_pretrained()`` / ``.save_pretrained()`` technically work. The pieces
+this module adds:
+
+1. Register ``ChronosConfig`` with ``AutoConfig`` and ``ChronosForCausalLM``
+   with ``AutoModelForCausalLM`` so users can do
+   ``AutoModelForCausalLM.from_pretrained("path/to/chronos")`` without
+   ``trust_remote_code=True``.
+
+2. ``save_chronos_pretrained`` / ``load_chronos_pretrained`` helpers that
+   also preserve the Chronos expert-cache layout (``cluster_manifest.json``
+   alongside the ``.ctsr`` cluster files) so a round-trip keeps the
+   offloaded SSD setup intact.
+
+Importing this module is side-effect-ful: registration happens at import
+time. ``chronos/__init__.py`` imports it so ordinary usage just works.
+"""
+from __future__ import annotations
+
+import os
+import shutil
+from typing import Optional
+
+import torch
+
+
+def _register_auto_classes():
+    """Register Chronos classes with Transformers AutoXXX factories."""
+    try:
+        from transformers import AutoConfig, AutoModelForCausalLM
+    except Exception:
+        return
+    from chronos.model.config import ChronosConfig
+    from chronos.model.model_chronos import ChronosForCausalLM
+
+    # Idempotent: Transformers raises if we try to register the same config
+    # twice. We catch and ignore.
+    try:
+        AutoConfig.register("chronos", ChronosConfig)
+    except (ValueError, KeyError):
+        pass
+    try:
+        AutoModelForCausalLM.register(ChronosConfig, ChronosForCausalLM)
+    except (ValueError, KeyError):
+        pass
+
+
+_register_auto_classes()
+
+
+def save_chronos_pretrained(
+    model,
+    save_directory: str,
+    expert_cache_dir: Optional[str] = None,
+    safe_serialization: bool = True,
+) -> None:
+    """Save a Chronos model in HF format (``model.safetensors`` + ``config.json``),
+    and copy the cluster manifest + .ctsr files into ``save_directory/expert_cache/``
+    if ``expert_cache_dir`` is provided so the offloaded layout travels with
+    the checkpoint."""
+    os.makedirs(save_directory, exist_ok=True)
+
+    # MiniMind/Chronos tie ``lm_head.weight`` to ``model.embed_tokens.weight``
+    # but doesn't declare it via ``_tied_weights_keys``, which trips
+    # safetensors' shared-tensor detector. Clone lm_head once before saving;
+    # restore the tie after.
+    lm_head = getattr(model, "lm_head", None)
+    embed = getattr(getattr(model, "model", model), "embed_tokens", None)
+    re_tie = (
+        lm_head is not None and embed is not None
+        and lm_head.weight.data_ptr() == embed.weight.data_ptr()
+    )
+    if re_tie:
+        original = lm_head.weight
+        lm_head.weight = torch.nn.Parameter(original.detach().clone())
+
+    try:
+        model.save_pretrained(save_directory, safe_serialization=safe_serialization)
+    finally:
+        if re_tie:
+            lm_head.weight = embed.weight  # restore the tie
+
+    if expert_cache_dir and os.path.isdir(expert_cache_dir):
+        dst = os.path.join(save_directory, "expert_cache")
+        os.makedirs(dst, exist_ok=True)
+        for name in os.listdir(expert_cache_dir):
+            if name.endswith((".ctsr", ".json")):
+                shutil.copy2(
+                    os.path.join(expert_cache_dir, name),
+                    os.path.join(dst, name),
+                )
+
+
+def load_chronos_pretrained(
+    save_directory: str,
+    device: Optional[str] = None,
+    **from_pretrained_kwargs,
+):
+    """Load a Chronos model saved via ``save_chronos_pretrained``.
+
+    If ``save_directory/expert_cache/cluster_manifest.json`` is present,
+    the path is returned alongside the model so callers can wire it into
+    ``CacheManager(ssd_dir=...)`` or ``ExpertStore.attach_cluster_manifest(...)``.
+
+    Returns ``(model, expert_cache_dir_or_None)``.
+    """
+    from chronos.model.model_chronos import ChronosForCausalLM
+    model = ChronosForCausalLM.from_pretrained(save_directory, **from_pretrained_kwargs)
+
+    # MiniMind/Chronos precompute RoPE frequencies as *persistent* buffers.
+    # When the checkpoint is saved as fp16 and reloaded, these deterministic
+    # buffers get truncated to zeros/NaNs. Recompute them from scratch —
+    # they are a pure function of config.
+    _rehydrate_rope_buffers(model)
+
+    if device is not None:
+        model = model.to(device)
+
+    cache_dir = os.path.join(save_directory, "expert_cache")
+    if os.path.exists(os.path.join(cache_dir, "cluster_manifest.json")):
+        return model, cache_dir
+    return model, None
+
+
+def _rehydrate_rope_buffers(model) -> None:
+    """Recompute ``freqs_cos`` and ``freqs_sin`` on the inner ChronosModel.
+    These are deterministic from config; the precompute logic lives inside
+    ``ChronosModel.__init__``. We invoke it via a fresh ChronosModel and
+    copy the buffers across. Cheap — only runs once per load."""
+    try:
+        inner = getattr(model, "model", None)
+        if inner is None:
+            return
+        cfg = model.config
+        from chronos.model.model_chronos import ChronosModel
+        fresh = ChronosModel(cfg)
+        if hasattr(inner, "freqs_cos") and hasattr(fresh, "freqs_cos"):
+            inner.freqs_cos = fresh.freqs_cos
+        if hasattr(inner, "freqs_sin") and hasattr(fresh, "freqs_sin"):
+            inner.freqs_sin = fresh.freqs_sin
+    except Exception:
+        # Rehydration is best-effort; fall through silently so an unexpected
+        # config shape doesn't block the load.
+        pass