Project-Chronos 0.1.1__py3-none-any.whl

chronos/runtime/metrics.py

@@ -0,0 +1,75 @@
+"""
+chronos/runtime/metrics.py
+
+A tiny in-process metrics bus. ExpertStore and AsyncPrefetcher push events
+into the singleton bus; the WebUI's IO Monitor tab polls ``snapshot()`` and
+plots the rolling buffers.
+
+Design constraints:
+
+- **Zero deps** — pure stdlib. No Prometheus, no opentelemetry.
+- **Thread-safe** — both the prefetcher daemon and the main thread emit.
+- **Bounded memory** — each metric is a deque(maxlen=600) so a 1-Hz UI
+  poll over a 10-minute window stays under a few KB.
+- **Optional** — emitting is wrapped in a try/except so missing
+  ``chronos.runtime.metrics`` cannot crash the inference loop.
+"""
+from __future__ import annotations
+
+import threading
+import time
+from collections import deque
+from typing import Deque, Dict, List, Tuple
+
+
+_BUFFER_SIZE = 600  # ~10 min at 1 Hz
+
+
+class MetricsBus:
+    def __init__(self):
+        self._lock = threading.Lock()
+        # metric_name -> deque[(timestamp, value)]
+        self._buffers: Dict[str, Deque[Tuple[float, float]]] = {}
+
+    def record(self, name: str, value: float, ts: float = None) -> None:
+        if ts is None:
+            ts = time.monotonic()
+        with self._lock:
+            buf = self._buffers.get(name)
+            if buf is None:
+                buf = deque(maxlen=_BUFFER_SIZE)
+                self._buffers[name] = buf
+            buf.append((float(ts), float(value)))
+
+    def latest(self, name: str) -> float:
+        with self._lock:
+            buf = self._buffers.get(name)
+            if not buf:
+                return float("nan")
+            return buf[-1][1]
+
+    def series(self, name: str) -> List[Tuple[float, float]]:
+        with self._lock:
+            buf = self._buffers.get(name)
+            return list(buf) if buf else []
+
+    def snapshot(self) -> Dict[str, List[Tuple[float, float]]]:
+        with self._lock:
+            return {k: list(v) for k, v in self._buffers.items()}
+
+    def reset(self) -> None:
+        with self._lock:
+            self._buffers.clear()
+
+
+# Module-level singleton. Importing this module is essentially free.
+bus = MetricsBus()
+
+
+def safe_record(name: str, value: float) -> None:
+    """Convenience wrapper that swallows any exception so metrics emission
+    can never break the hot path."""
+    try:
+        bus.record(name, value)
+    except Exception:
+        pass

chronos/serving/__init__.py

@@ -0,0 +1,13 @@
+from chronos.serving.vllm_adapter import (
+    register_chronos_with_vllm,
+    set_available_expert_masks,
+    is_available,
+    HAS_VLLM,
+)
+
+__all__ = [
+    "register_chronos_with_vllm",
+    "set_available_expert_masks",
+    "is_available",
+    "HAS_VLLM",
+]

chronos/serving/vllm_adapter.py

@@ -0,0 +1,85 @@
+"""
+chronos/serving/vllm_adapter.py
+
+Optional vLLM serving adapter. ``vllm`` is NOT a requirement of Chronos
+and the module does the right thing whether or not it's installed.
+
+Usage::
+
+    from chronos.serving import register_chronos_with_vllm
+    register_chronos_with_vllm()  # no-op if vllm not installed
+
+When ``vllm`` is available, ``ChronosForCausalLM`` is registered in
+``vllm.ModelRegistry`` so vLLM can find the architecture by name. The
+actual worker-side glue (paged KV cache interop, per-layer expert mask
+hook) is version-sensitive; see ``docs/vllm_integration.md`` for the
+full integration steps and known gotchas.
+"""
+from __future__ import annotations
+
+import importlib.util
+from typing import Any, Optional
+
+
+HAS_VLLM = importlib.util.find_spec("vllm") is not None
+
+
+def register_chronos_with_vllm(verbose: bool = True) -> bool:
+    """Register ``ChronosForCausalLM`` with vLLM's ModelRegistry.
+
+    Returns True on success, False if vLLM is not installed or the API
+    shape has changed. This function is idempotent.
+    """
+    if not HAS_VLLM:
+        if verbose:
+            print(
+                "[chronos.serving] vllm not installed — skipping registration.\n"
+                "  To enable vLLM serving: `pip install vllm`\n"
+                "  (Linux + CUDA required; see docs/vllm_integration.md)"
+            )
+        return False
+
+    try:
+        import vllm  # noqa: F401
+        from vllm import ModelRegistry
+
+        from chronos.model.model_chronos import ChronosForCausalLM
+
+        # ModelRegistry.register_model signature changes across vLLM
+        # versions. Try the common shapes.
+        arch = "ChronosForCausalLM"
+        if hasattr(ModelRegistry, "register_model"):
+            try:
+                ModelRegistry.register_model(arch, ChronosForCausalLM)
+            except TypeError:
+                # Older API: register_model(name, module_path, class_name)
+                ModelRegistry.register_model(
+                    arch, "chronos.model.model_chronos", "ChronosForCausalLM",
+                )
+        else:
+            if verbose:
+                print("[chronos.serving] vllm.ModelRegistry API not recognized; skipping.")
+            return False
+
+        if verbose:
+            print(f"[chronos.serving] Registered {arch} with vLLM ModelRegistry.")
+        return True
+    except Exception as e:
+        if verbose:
+            print(f"[chronos.serving] vLLM registration failed: {e}")
+        return False
+
+
+def set_available_expert_masks(model, masks) -> None:
+    """Hook point for vLLM worker to inject per-step expert availability
+    masks. Sets ``model._chronos_available_expert_masks`` so the next
+    forward picks it up (if the model is wired to check there).
+
+    This is a best-effort placeholder; the actual decode-loop glue depends
+    on the vLLM version and is documented in docs/vllm_integration.md.
+    """
+    setattr(model, "_chronos_available_expert_masks", masks)
+
+
+def is_available() -> bool:
+    return HAS_VLLM

chronos/trainer/__init__.py

@@ -0,0 +1,6 @@
+from .chronos_trainer import ChronosTrainer
+from .sft_trainer import ChronosSFTTrainer
+from .dpo_trainer import ChronosDPOTrainer
+from .orpo_trainer import ChronosORPOTrainer
+from .grpo_trainer import ChronosGRPOTrainer
+from .distill_trainer import ChronosDistillTrainer

chronos/trainer/chronos_trainer.py

@@ -0,0 +1,151 @@
+import sys
+import chronos.deps  # ensure minimind on sys.path
+
+import os
+import time
+import torch
+import torch.nn.functional as F
+from torch import optim
+from torch.utils.data import DataLoader
+from contextlib import nullcontext
+
+from model.model_minimind import MiniMindConfig
+from trainer.trainer_utils import (
+    Logger, is_main_process, lm_checkpoint, setup_seed,
+)
+from chronos.trainer.optim_utils import get_lr  # warmup→cosine override
+from chronos.data.flexible_dataset import (
+    FlexibleDataset as PretrainDataset,
+    StreamingSFTDataset as SFTDataset,
+)
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.temporal_loss import lookahead_supervision_loss
+from chronos.trainer.loss_mixin import chronos_loss_term, collect_router_probs
+
+
+class ChronosTrainer:
+    """
+    Training orchestrator for Project Chronos.
+
+    Wraps minimind's training loop and injects:
+    - temporal_locality_loss via λ2 * ||E_t - E_{t-1}||^2
+    - load_balance_loss via λ1 (already computed inside ChronosMOEFeedForward)
+    - LookaheadRouter accuracy logging
+    """
+
+    def __init__(self, config: ChronosConfig, args):
+        self.config = config
+        self.args = args
+        self.device = args.device
+        seed = getattr(args, "seed", None)
+        if seed is not None:
+            setup_seed(int(seed))
+        device_type = "cuda" if "cuda" in self.device else "cpu"
+        dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        self.autocast_ctx = (
+            nullcontext() if device_type == "cpu"
+            else torch.cuda.amp.autocast(dtype=dtype)
+        )
+        self.scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == "float16"))
+
+        self.model = ChronosForCausalLM(config).to(self.device)
+        from chronos.trainer.optim_utils import build_optimizer
+        self.optimizer = build_optimizer(
+            self.model, lr=args.learning_rate,
+            weight_decay=float(getattr(args, "weight_decay", 0.01)),
+        )
+
+    def _collect_router_probs(self) -> torch.Tensor:
+        """Gather last_router_probs from all MoE layers → [B, S, L, E]."""
+        return collect_router_probs(self.model)
+
+    def train_step(self, input_ids, labels, step, total_steps):
+        input_ids = input_ids.to(self.device)
+        labels = labels.to(self.device)
+
+        lr = get_lr(step, total_steps, self.args.learning_rate)
+        for pg in self.optimizer.param_groups:
+            pg['lr'] = lr
+
+        with self.autocast_ctx:
+            outputs, lookahead_probs = self.model(input_ids, labels=labels)
+            ce_loss = outputs.loss
+            aux_loss = outputs.aux_loss
+
+            loss = chronos_loss_term(
+                self.model, ce_loss, lookahead_probs, self.config,
+                aux_loss=aux_loss,
+            )
+
+            # Standalone lookahead loss value for logging
+            lookahead_loss_val = 0.0
+            router_4d = self._collect_router_probs()
+            if (
+                router_4d is not None
+                and router_4d.shape[1] > 1
+                and lookahead_probs is not None
+                and self.config.lookahead_steps > 0
+            ):
+                teacher = router_4d.mean(dim=2).detach()
+                la_term = lookahead_supervision_loss(
+                    lookahead_probs, teacher, self.config.lookahead_steps,
+                )
+                lookahead_loss_val = float(la_term.item())
+
+            loss = loss / self.args.accumulation_steps
+
+        self.scaler.scale(loss).backward()
+
+        if step % self.args.accumulation_steps == 0:
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.grad_clip)
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            self.optimizer.zero_grad(set_to_none=True)
+
+        return loss.item() * self.args.accumulation_steps, ce_loss.item(), aux_loss.item(), lookahead_loss_val
+
+    def train_epoch(self, epoch, loader, iters, start_step=0, wandb=None):
+        self.model.train()
+        start_time = time.time()
+        total_steps = self.args.epochs * iters
+
+        for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+            loss, ce_loss, aux_loss, la_loss = self.train_step(
+                input_ids, labels, epoch * iters + step, total_steps
+            )
+
+            if step % self.args.log_interval == 0 or step == iters:
+                elapsed = time.time() - start_time
+                eta = elapsed / max(step - start_step, 1) * (iters - step) // 60
+                lr = self.optimizer.param_groups[-1]['lr']
+                Logger(
+                    f'Epoch:[{epoch+1}/{self.args.epochs}]({step}/{iters}) '
+                    f'loss:{loss:.4f} ce:{ce_loss:.4f} aux:{aux_loss:.4f} '
+                    f'la:{la_loss:.4f} lr:{lr:.2e} eta:{eta:.1f}min'
+                )
+                if wandb:
+                    wandb.log({
+                        "loss": loss, "ce_loss": ce_loss, "aux_loss": aux_loss,
+                        "lookahead_loss": la_loss, "lr": lr,
+                    })
+
+            if (step % self.args.save_interval == 0 or step == iters) and is_main_process():
+                self._save(epoch, step, wandb)
+
+    def _save(self, epoch, step, wandb=None):
+        self.model.eval()
+        os.makedirs(self.args.save_dir, exist_ok=True)
+        ckp = f'{self.args.save_dir}/chronos_{self.config.hidden_size}_moe.pth'
+        state = self.model.state_dict()
+        torch.save({k: v.half().cpu() for k, v in state.items()}, ckp)
+        lm_checkpoint(
+            self.config, weight='chronos', model=self.model,
+            optimizer=self.optimizer, scaler=self.scaler,
+            epoch=epoch, step=step, wandb=wandb,
+            save_dir=self.args.save_dir,
+        )
+        self.model.train()
+        del state

chronos/trainer/distill_trainer.py

@@ -0,0 +1,193 @@
+"""
+chronos/trainer/distill_trainer.py — Stage 6 distillation trainer.
+
+Loss:
+    L_KD    = T² · KL( softmax(student_logits/T) || softmax(teacher_logits/T) )
+    L_label = cross_entropy(student_logits, labels, ignore_index=-100)
+    L_total = α · L_KD + (1 − α) · L_label  +  chronos_loss_term (lookahead etc.)
+
+The teacher is any frozen model that returns ``.logits`` on forward (Chronos
+or MiniMind checkpoints both work). Router KL anchor still applies — Stage
+6 shouldn't corrupt the cache layout any more than SFT/DPO can.
+"""
+from __future__ import annotations
+
+import os
+import time
+from contextlib import nullcontext
+
+import torch
+import torch.nn.functional as F
+from torch import optim
+from torch.utils.data import DataLoader
+
+import chronos.deps  # noqa: F401
+from trainer.trainer_utils import Logger  # type: ignore
+from chronos.data.flexible_dataset import StreamingSFTDataset as SFTDataset
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.temporal_loss import lookahead_supervision_loss
+from chronos.trainer.optim_utils import get_lr
+from chronos.trainer.loss_mixin import (
+    chronos_loss_term,
+    collect_router_probs,
+    router_kl_anchor,
+    capture_reference_routing,
+)
+
+
+def distillation_kl(student_logits: torch.Tensor, teacher_logits: torch.Tensor,
+                    T: float = 4.0, ignore_mask: torch.Tensor = None) -> torch.Tensor:
+    """Temperature-scaled KL divergence; caller is responsible for masking."""
+    s = F.log_softmax(student_logits / T, dim=-1)
+    t = F.softmax(teacher_logits / T, dim=-1).detach()
+    kl = F.kl_div(s, t, reduction="none").sum(dim=-1)  # [B, S]
+    if ignore_mask is not None:
+        kl = kl * ignore_mask
+        denom = ignore_mask.sum().clamp_min(1.0)
+    else:
+        denom = torch.tensor(kl.numel(), dtype=kl.dtype, device=kl.device).clamp_min(1.0)
+    return (T * T) * (kl.sum() / denom)
+
+
+class ChronosDistillTrainer:
+    def __init__(self, model: ChronosForCausalLM, teacher, config: ChronosConfig,
+                 args, tokenizer):
+        self.model = model
+        self.teacher = teacher.eval().requires_grad_(False)
+        self.config = config
+        self.args = args
+        self.tokenizer = tokenizer
+        self.device = args.device
+        self.alpha = float(getattr(args, "alpha", 0.7))
+        self.temperature = float(getattr(args, "temperature", 4.0))
+
+        device_type = "cuda" if "cuda" in str(self.device) else "cpu"
+        dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        self.autocast_ctx = nullcontext() if device_type == "cpu" \
+            else torch.cuda.amp.autocast(dtype=dtype)
+        self.scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == "float16"))
+        from chronos.trainer.optim_utils import build_optimizer
+        self.optimizer = build_optimizer(
+            model, lr=args.learning_rate,
+            weight_decay=float(getattr(args, "weight_decay", 0.01)),
+        )
+        self._router_ref = None
+
+    def set_calibration_batch(self, x: torch.Tensor):
+        self._router_ref = capture_reference_routing(self.model, x, self.device)
+
+    def _teacher_logits(self, input_ids: torch.Tensor) -> torch.Tensor:
+        with torch.no_grad():
+            out = self.teacher(input_ids)
+            return (out[0].logits if isinstance(out, tuple) else out.logits).detach()
+
+    def train_step(self, input_ids, labels, step, total_steps):
+        input_ids = input_ids.to(self.device)
+        labels = labels.to(self.device)
+
+        lr = get_lr(step, total_steps, self.args.learning_rate)
+        for pg in self.optimizer.param_groups:
+            pg["lr"] = lr
+
+        with self.autocast_ctx:
+            outputs, lookahead_probs = self.model(input_ids, labels=labels)
+            ce_loss = outputs.loss
+            student_logits = outputs.logits
+
+            # Teacher forward (no grad)
+            teacher_logits = self._teacher_logits(input_ids)
+            # Align vocab dims if teacher is smaller (pad or truncate)
+            if teacher_logits.shape[-1] != student_logits.shape[-1]:
+                V = min(teacher_logits.shape[-1], student_logits.shape[-1])
+                teacher_logits = teacher_logits[..., :V]
+                student_logits_kd = student_logits[..., :V]
+            else:
+                student_logits_kd = student_logits
+            # Align seq len
+            S = min(teacher_logits.shape[1], student_logits_kd.shape[1])
+            teacher_logits = teacher_logits[:, :S, :]
+            student_logits_kd = student_logits_kd[:, :S, :]
+
+            # Loss mask = positions where label != -100
+            ignore_mask = (labels[:, :S] != -100).float() if labels.shape[1] >= S else None
+            kd = distillation_kl(
+                student_logits_kd, teacher_logits,
+                T=self.temperature, ignore_mask=ignore_mask,
+            )
+
+            base = self.alpha * kd + (1.0 - self.alpha) * ce_loss
+            loss = chronos_loss_term(
+                self.model, base, lookahead_probs, self.config,
+                aux_loss=outputs.aux_loss,
+            )
+
+            anc_val = 0.0
+            lam = float(getattr(self.config, "lambda_router_anchor", 0.0))
+            if lam > 0.0 and self._router_ref is not None:
+                r4 = collect_router_probs(self.model)
+                if r4 is not None and r4.shape[1] > 1:
+                    cur = r4.mean(dim=2)
+                    ref = self._router_ref
+                    T_ax = min(cur.shape[1], ref.shape[1])
+                    B = min(cur.shape[0], ref.shape[0])
+                    anc = router_kl_anchor(cur[:B, :T_ax, :], ref[:B, :T_ax, :], lam)
+                    loss = loss + anc
+                    anc_val = float(anc.item())
+
+            la_val = 0.0
+            r4 = collect_router_probs(self.model)
+            if (r4 is not None and r4.shape[1] > 1
+                    and lookahead_probs is not None
+                    and self.config.lookahead_steps > 0):
+                la_val = float(lookahead_supervision_loss(
+                    lookahead_probs, r4.mean(dim=2).detach(),
+                    self.config.lookahead_steps,
+                ).item())
+
+            loss = loss / self.args.accumulation_steps
+
+        self.scaler.scale(loss).backward()
+        if step % self.args.accumulation_steps == 0:
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.grad_clip)
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            self.optimizer.zero_grad(set_to_none=True)
+
+        return {
+            "loss": float(loss.item() * self.args.accumulation_steps),
+            "kd": float(kd.item()),
+            "ce": float(ce_loss.item()),
+            "la": la_val,
+            "anchor": anc_val,
+        }
+
+    def train_epoch(self, epoch, loader, iters, max_steps=None):
+        self.model.train()
+        total_steps = self.args.epochs * iters
+        for step, (ids, labels) in enumerate(loader, start=1):
+            if max_steps is not None and step > max_steps:
+                break
+            s = self.train_step(ids, labels, epoch * iters + step, total_steps)
+            if step % self.args.log_interval == 0 or step == iters:
+                lr = self.optimizer.param_groups[-1]["lr"]
+                Logger(
+                    f"[Distill] Epoch[{epoch+1}/{self.args.epochs}]({step}/{iters}) "
+                    f"loss:{s['loss']:.4f} kd:{s['kd']:.4f} ce:{s['ce']:.4f} "
+                    f"la:{s['la']:.4f} anchor:{s['anchor']:.4f} lr:{lr:.2e}"
+                )
+
+    def _save(self, epoch, step):
+        os.makedirs(self.args.save_dir, exist_ok=True)
+        ckp = os.path.join(self.args.save_dir, f"distill_{self.config.hidden_size}_moe.pth")
+        self.model.eval()
+        state = self.model.state_dict()
+        torch.save({k: v.half().cpu() for k, v in state.items()}, ckp)
+        self.model.train()
+
+
+def build_distill_loader(data_path, tokenizer, max_seq_len, batch_size):
+    ds = SFTDataset(data_path, tokenizer, max_length=max_seq_len)
+    return DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=0)