执行一次"] + P --> G["循环组
MLA / SWA + top-2 MoE"] + G -->|"循环 d = 1..D
收缩更新，ACT 退出"| G + G --> C["Coda
执行一次"] + C --> H["LM head"] + P -. "lookahead 输入" .-> LR["LookaheadRouter
专家预取"] +``` + +**相对 v2 的变化** + +| | v2 | v3 | +|---|---|---| +| 循环内注意力 | 只在第 0 轮执行（`first_only`） | 每轮都执行 | +| 每个循环层的 KV cache | 1 个槽位；打开"每轮注意力"开关时每轮重写（实测 prefill+decode 与 full forward 偏差 0.94） | `shared_first`：每轮重算 Q，K/V 取第 0 轮，KV 槽位数 = 物理层数 | +| 路由 | top-1（gate 拿不到 CE 梯度） | top-2 + 1 个共享专家 | +| MoE 调度 | Python 逐专家循环，含 `.item()` 同步 | 排序分组调度：MLX 用 `gather_mm`，Torch 用 `argsort` + `index_add_` | +| 提前退出 | 每步同步 host | 纯张量 ACT，结果精确且与 batch 无关；decode 可逐 token 退出，不留 KV 缺口 | +| 解码 | 普通自回归 | 循环深度自推测解码 | + +`per_iteration`（每个虚拟层一个 KV 槽位）保留，用于消融对比。v2 checkpoint 自动以 `first_only` 加载，logits 与保存的 golden 张量误差在 1e-5 以内。 + +**自推测解码（不增加参数）** +- 同一个模型先用浅深度起草 γ 个 token，再用一次全深度批量前向校验。 +- greedy 输出与全深度 greedy 逐 token 一致。 +- 采样使用 rejection sampling，保持全深度的分布不变。 +- 拒绝后截断 cache：MLA 截断 `c_kv`，滑窗 cache 保留 `窗口 + γ_max` 个条目。 +- Torch（`chronos/runtime/speculative.py`）与 MLX（`chronos/mlx/speculative.py`、`ChronosMLXInferenceEngine.generate(speculative=True)`）使用同一套算法。 +- 两个后端的测试都覆盖 greedy 一致性、拒绝采样分布，以及回滚结果与从头 prefill 一致。 +- Web UI 推理 Tab 提供自推测开关，可调起草深度和 γ。 +- 50M demo 在实测过的所有运行时上都结果一致，但都没有提速。实测条件：预训练 step 13,996 checkpoint，greedy，batch 1，FLORES+ dev 英→中，起草深度 1，γ 4。报告在 `reports/speculative/`。 + +| 运行时 | 普通 tok/s | 自推测 tok/s | 比值 | 接受率 | +|---|---|---|---|---| +| MLX GPU bf16，专家全部在内存 | 112.0 | 107.4 | 0.96× | 42 % | +| MLX GPU bf16，懒加载专家（216 页中常驻 12 页） | 43.7 | 32.0 | 0.73× | 42 % | +| Torch MPS fp32，专家全部加载 | 32.1 | 25.9 | 0.81× | 40 % | +| Torch MPS fp32，懒加载专家 | 17.0 | 14.7 | 0.86× | 40 % | + +- 这个规模下，一次全深度单步的开销与一次批量校验相当。 +- 懒加载时，起草和校验路由到的专家不同，自推测从 SSD 读取的专家页多出 1.4–1.6 倍。 +- greedy 输出在 CPU fp32 和 Torch MPS fp32 下完全一致；GPU bf16 下，24 句中有 2 句在近似平票的位置翻转了 token。 +- 因此该开关默认关闭。 + +**精度链路** +- 全链路使用同一个计算精度。 +- FP32 只用在收益明确的位置：RMSNorm 统计量、SDPA softmax、RoPE 表、路由 softmax、LM head logits、优化器 master 权重。 +- MLX↔Torch 的 bf16 互转逐位一致。 +- `auto` 在 MLX GPU 和 CUDA 上选 bf16，在 CPU 和 MPS 上选 fp32（MPS 的 bf16 实测更慢，见下文）。 + +**v3 新增工具** + +| 命令 | 作用 | +|---|---| +| `chronos tokenizer` | 自研字节回退 BPE：按文字脚本预切分，按语言平衡合并；输出 HF 兼容的 `tokenizer.json` 和 chat template | +| `chronos corpus download / build / shards` | 固定版本下载、清洗、全局去重、剔除评测集泄漏、生成 uint16 分片 | +| `chronos longrun --spec recipe.json`（别名 `run`） | 按墙钟预算训练的 MLX 任务：每个阶段 warmup + cosine，NaN 回滚，续跑逐位一致，带验证和翻译探针 | +| `chronos serve` | 原生运行时上的 Ollama/OpenAI 协议服务（见下文） | +| `chronos monitor` | 只读训练监控页面 + JSON API（见下文） | + +`chronos.model.factory.fit_config` 按总参数量和激活参数量目标搜索模型配置。 + +**50M 英译中 demo（进行中）** +- 模型：总参数 49.69M，激活 12.09M；36 个专家；MLA latent 128；SWA 窗口 128；自研 16k 词表。 +- 训练：72 小时 MLX bf16 任务（先预训练，后 SFT），2026-11-05 启动。 +- 中途实测：预训练第 13,996 步（计划约 24.7k 步）valid CE 为 2.477（第 1,656 步时为 3.012）。该 checkpoint 已打包为 `chronos-enzh-50m-pretrain-s13996`（见“模型包”）。 +- FLORES+ 和 WMT22 上的 BLEU/chrF，以及各项消融（注意力模式、深度、懒加载专家预算、推测 γ），等训练结束后公布。在此之前不给出任何质量数字。 + +**推理与内存选项（全部默认关闭）。** 每个选项既是 config 字段，也可作为 `chronos serve` 和 Inference Tab 的覆盖项；测试在 Torch 和 MLX 上把每个选项与参考路径做对比。 + +| 字段 | 取值（第一个为默认） | 作用 | +|---|---|---| +| `kv_cache_dtype` | `auto`、`int8` | `int8` 把 MLA `c_kv` 与 SWA k/v 按 token、每 32 元素一组量化并带 scale；回滚和推测解码都支持 | +| `mla_decode` | `expanded`、`absorbed`、`auto` | `absorbed` 直接对 `c_kv` 打分，不展开 K/V。`auto` 只对短的缓存块做吸收，改为显式开启，因为实测 bf16 下“prefill 走 expanded、decode 走 absorbed”会偏离 full forward（logit 最大差 0.06） | +| `mla_cache_rope` | `false`、`true` | 额外缓存已加 RoPE 的 key，用内存换 decode 速度 | +| `router_scope` | `per_iteration`、`shared` | `shared` 让循环后续各轮复用第 0 轮的专家，每个 token 驻留的专家更少 | +| `draft_exit_kl_coef`、`draft_exit_depth` | `0`、`1` | 训练项：浅层出口对 detach 后的全深度 logits 做 KL，提高推测解码接受率 | +| `max_resident_experts`、`ram_budget_mb` | 不设 | 快速层和 RAM 层的硬上限；LRU 不淘汰 lookahead 预测要用的专家；预取覆盖一个虚拟层的全部循环轮次 | + +Web UI 新增 Tokenizer Tab（抽取 → 训练 → 评估）和 Serve Tab（启停 `chronos serve`；API key 通过环境变量传入，不进 argv）。 + +--- + ## 三级存储架构（M1：cluster-aware safetensors） ```mermaid @@ -183,7 +276,7 @@ flowchart LR 训练精度与资源策略： -- 默认 `--dtype auto`。MPS/MLX 训练优先解析为 BF16 以保证稳定性，CUDA/XPU 解析为 FP16；CPU 默认 FP32，只有显式传 `--dtype float16` 或 `--dtype bfloat16` 才启用 CPU autocast。 +- 默认 `--dtype auto`。权重、激活、KV cache 和专家页统一使用一个计算精度。FP32 只保留在融合 kernel 内部（SDPA softmax、RMSNorm 统计量），以及 RoPE 表、路由 softmax、loss logits 和优化器的 FP32 master/moments。`auto` 的解析结果：MLX GPU 为 BF16；CUDA 为 BF16，不支持时改用 FP16 并开启 loss scaling；CPU 和 MPS 为 FP32。在 MPS 上，原生 BF16 实测比 FP32 慢：M5 Max、torch 2.10、Chronos 50M 每个训练步 0.749 s 对 0.619 s，decode 每个 token 49.1 ms 对 37.3 ms。如需把权重和 KV 内存减半，可以显式传 `--dtype bfloat16`，代价是速度。 - CPU 训练默认使用物理核心；可用 `--cpu_threads` 或 `--cpu_budget_percent` 覆盖。 - macOS 上 MPS/MLX 训练默认强制 DataLoader workers 为 `0`，避免 multiprocessing 触发 Metal command-buffer 崩溃。CPU/CUDA 仍可使用 worker 进程；高级用户可用 `CHRONOS_ALLOW_METAL_DATALOADER_WORKERS=1` 覆盖这个保护。 - MLX 原生训练会按 `log_interval` 同步 UI 日志、标量读数和图表点；Web UI Stop 会在每个 batch 边界检查并停止。 @@ -194,8 +287,8 @@ trainer 在 `chronos.mlx.*` 中复刻 PyTorch 训练栈：masked CE、DPO/ORPO/G 蒸馏损失、load balance、temporal locality、lookahead soft-target 监督、 lookahead top-k 命中损失、router KL anchor。即使权重使用 BF16/FP16，路由 softmax、CE、Adam moments 等数值敏感路径仍以 FP32 计算；因此 `auto` 在 -MLX/MPS 上默认优先 BF16。FP16 的指数范围太小，在小模型/小 batch/路由分布 -剧烈波动时更容易把 loss 推成 NaN。 +MLX GPU 上默认 BF16（MPS 实测 FP32 更快，默认 FP32）。FP16 的指数范围太小， +在小模型/小 batch/路由分布剧烈波动时更容易把 loss 推成 NaN。 Web UI 的 MLX 训练日志与 CPU/MPS 对齐：step、loss、steps/s、ETA、 checkpoint 保存事件、stop 事件、verify 结果都会同步到界面。Stop 在 batch @@ -337,7 +430,7 @@ D_{\mathrm{KL}} --- -## 安装 (Not Ready in PyPI Yet) +## 安装 ```bash pip install Project_Chronos @@ -356,22 +449,29 @@ pip install -e ".[dev]" pip install "Project_Chronos[mlx]" ``` +**可选依赖组**（版本已固定）： +```bash +pip install "Project_Chronos[serve]" # chronos serve / chronos monitor（fastapi、uvicorn） +pip install "Project_Chronos[translate]" # 语料反分词 + BLEU/chrF 探针（sacremoses、sacrebleu） +``` + **vLLM 服务（可选，仅 Linux+CUDA）：** ```bash pip install vllm ``` > **minimind 依赖**：Project Chronos 使用 [minimind](https://github.com/jingyaogong/minimind) 作为 MoE 内核。 -> 若本地未找到，首次 import 时自动克隆至 `~/.cache/chronos/minimind-master/`。 +> 若本地未找到，首次 import 时自动克隆至 `~/.cache/chronos/minimind-master/`（需要 `git` 和网络）。 +> 设置 `CHRONOS_MINIMIND_DIR=/path/to/minimind` 可改用本地副本；模型包自带 `third_party/minimind`，启动脚本会设置这个变量。 > minimind 采用 **Apache-2.0** 授权，完整归属见 [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md)。 -**环境要求**：Python 3.10+，PyTorch 2.4+ +**环境要求**：Python 3.10+，PyTorch 2.6+ --- ## 快速开始 -### Web UI（M6 — 8 个 Tab，4 种语言） +### Web UI（12 个 Tab，4 种语言） ```bash chronos-ui @@ -379,7 +479,116 @@ chronos-ui python chronos_app.py ``` -包含：⚙️ Config（含右侧实时参数估算面板，合并了 Designer）/ 🏋️ Train（拥有 data_path）/ 🧪 6-Stage Pipeline（每阶段独立数据路径）/ 💬 Inference（含懒加载 vs 全量 DRAM 对比）/ 📦 Export（FP16/Q8_0 safetensors 与 GGUF）/ 📊 Benchmark（Markdown 表 + 对比图）/ 🔬 Auto-Tune（持久化日志 + 一键 Apply Best → Config）/ 📡 IO Monitor。i18n 支持 zh-Hans / zh-Hant / en / ja。 +包含：⚙️ Config（v3 预设、自动规划器和实时参数估算面板）/ 🗂️ Data（数据集扫描与转换）/ 🔤 Tokenizer（从数据集抽取文本、训练 Chronos BPE、评估并保存）/ 🏋️ Train（拥有 data_path，按所选阶段合并阶段参数）/ 📈 Training Monitor（只读监控 `runs/`，见下文）/ 🧪 6-Stage Pipeline（每阶段独立数据路径）/ 💬 Inference（含懒加载 vs 全量 DRAM 对比）/ 📦 Export（FP16/Q8_0 safetensors 与 GGUF）/ 🛰️ Serve（启停 `chronos serve` 并查看日志）/ 📊 Benchmark（Markdown 表 + 对比图）/ 🔬 Auto-Tune（持久化日志 + 一键 Apply Best → Config）/ 📡 IO Monitor。i18n 支持 zh-Hans / zh-Hant / en / ja。窗口宽 1280 px 时，Gradio 会把最后几个 Tab（Auto-Tune、IO Monitor）收进 “…” 溢出菜单。 + +### 预设与自动规划器 + +Config Tab 提供 5 档 v3 预设。它们都由规划器生成，存放在 `ui/presets_v3.json`，可用 `python -m ui.presets` 重新生成。参数量在 meta device 上实测，峰值内存是估算值。 + +| 预设 | 总参数 / 激活参数（实测） | 形状 | 内存档位 | 训练峰值（估算） | +|---|---|---|---|---| +| Chronos-v3 Tiny (smoke / CI) | 1.51M / 1.29M | E4 top-2，词表 6400，seq 128，深度 2 | none | 2.13 GB | +| Chronos-v3 50M (en→zh recipe) | 49.69M / 12.09M | h384，L6，E36 top-2，词表 16384，seq 512 | none | 16.4 GB（正在跑的任务实测峰值 8.72 GB） | +| Chronos-v3 200M (32 GB)（默认） | 200.72M / 28.20M | h576，E32 top-2 | balanced | 10.14 GB | +| Chronos-v3 500M (32 GB) | 499.53M / 61.59M | E48 top-3 | aggressive | 9.26 GB | +| Chronos-v3 1B (64 GB) | 1002.17M / 99.74M | E64 top-4，seq 2048 | aggressive | 18.92 GB | + +旧预设名（Recommended-CN、Tiny/Small/Medium/Large）仍可加载，会映射到最接近的 v3 档位。 + +**自动规划器**（Config Tab →“自动规划”） + +输入： +- 目标总参数量、是否开启 MoE； +- 专家数，以及“激活专家数”或“激活参数量”二选一； +- 内存节省档位：`none`、`balanced`、`aggressive`、`extreme`； +- 可选：数据集路径、内存、token 预算、词表大小、序列长度。 + +`chronos.model.planner.plan_training` 的处理步骤： +1. **分析数据集。** 读取分片 manifest 或抽样文件，把记录分成预训练、SFT、偏好、RL 四类，并读取分词器报告。数据里有词表大小、序列长度（p95 或上限）、token 数和分词器路径时，直接采用。 +2. **搜索架构。** 找出 meta device 实测参数量最接近目标的隐藏维度、层数、专家宽度和 GQA 头数。 +3. **确定学习率和 batch。** 按 DeepSeek LLM 的计算最优拟合式（η = 0.3118·C^−0.125，B = 0.2920·C^0.3271）给出，宽模型的学习率另设上限。 +4. **匹配内存。** 按内存档位选择 micro-batch、梯度累积、int8 KV、路由范围、激活检查点、优化器分页和专家 RAM 预算。 +5. **生成各阶段参数。** pretrain、SFT、DPO、ORPO、GRPO、distill 各有独立的学习率、batch、epoch、保存间隔，以及各自的专用项（β、温度、α）。 + +阶段表可以直接编辑。Train Tab 会把所选阶段那一行合并到共享参数之上；pretrain 行与共享控件保持同步。“保存方案”把 config、各阶段参数和报告写成 JSON。报告里每个数字都标注实测或估算。 + +### 训练监控（控制台 Tab + 独立页面/API） + +`chronos longrun` 的运行目录包含以下文件： +- `run_spec.json`、`metrics.jsonl`； +- `status.json`：心跳； +- `runtime.json`：精度策略、设备、版本、代码哈希； +- `parameters.jsonl`：逐张量的权重、梯度、更新量统计，以及逐虚拟层路由； +- `samples.jsonl`、`inspections/` 和 checkpoint。 + +控制台的 **Training Monitor** Tab 和独立页面读取的是同一批文件，都不会占用训练设备。遥测功能上线前启动的运行同样能显示进度，其状态由 `metrics.jsonl` 和进程表推断。 + +```bash +chronos monitor --runs-dir runs # http://127.0.0.1:7870/ +chronos serve --models-dir exports --monitor-runs-dir runs # 同一页面挂在 /monitor/ +``` + +JSON API（只读）： +- `GET /api/state?run=&history=1` +- `/api/metrics?run=&phase=&since=` +- `/api/parameters`、`/api/tensors`、`/api/tensor?name=` +- `/api/log?bytes=`、`/api/runs` + +安全： +- 默认只监听 127.0.0.1，不带认证。 +- 绑定其他地址时，必须同时给出 `--allow-remote` 和 `--api-key`（或设置环境变量 `CHRONOS_MONITOR_KEY`）。 +- 浏览器登录一次后获得 HttpOnly、SameSite=Strict 的会话 cookie；脚本使用 `Authorization: Bearer
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Project Chronos
+
+
+ Training monitor
+
+
+
+
+
+
diff --git a/chronos/precision.py b/chronos/precision.py
new file mode 100644
index 0000000..e64112f
--- /dev/null
+++ b/chronos/precision.py
@@ -0,0 +1,195 @@
+"""One precision policy for every Chronos training and inference chain.
+
+A run picks a single *compute dtype* when it is built. Weights, activations,
+KV cache and SSD/RAM expert pages all live in that dtype, so the hot path
+never converts tensors. Precision is protected only where rounding would
+cost accuracy, and each of those places is either a fused kernel that
+accumulates in FP32 internally (no tensor cast) or a deliberately small cast:
+
+ site how FP32 is obtained cast size
+ --------------------------- --------------------------------- ---------------
+ attention softmax fused SDPA (MLX / torch) none
+ RMSNorm statistics mx.fast.rms_norm / fused torch none
+ RoPE fused rope / tables in compute dt none
+ router softmax softmax(precise=True) (MLX) none
+ LM head + cross-entropy upcast hidden [N,H] and W [V,H] H/V of logits
+ ACT halting / aux losses [B,S,G,D] scalars tiny
+ optimizer FP32 masters + moments one per update
+ gradient accumulation FP32 running sum one per micro
+
+Everything else stays in the compute dtype end to end. The policy is also
+the single place that knows device limits: MLX's CPU backend implements
+GatherMM for float32 only (mlx 0.32.4), so CPU MLX runs compute in float32
+for the whole model instead of upcasting every expert matmul.
+"""
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+
+COMPUTE_DTYPES = ("float32", "bfloat16", "float16")
+_ALIASES = {"fp32": "float32", "f32": "float32", "float": "float32",
+ "bf16": "bfloat16", "fp16": "float16", "half": "float16", "f16": "float16"}
+
+# Fused kernels (attention softmax, RMSNorm, RoPE, precise router softmax)
+# that keep FP32 internally without materializing FP32 tensors.
+FUSED_FP32_SITES = ("attention_softmax", "rms_norm", "rope", "router_softmax")
+# The only explicit casts left in the chains, with their size bound.
+BOUNDARY_CASTS = {
+ "lm_head_loss": "hidden [N,H] + head [V,H] -> fp32; logits are born fp32",
+ "halting_and_aux_losses": "[B,S,G,D] / [N,E] statistics -> fp32",
+ "optimizer_update": "grad -> fp32 master once per update",
+ "grad_accumulation": "fp32 running sum across micro-batches",
+}
+
+
+def normalize_dtype_name(name, default: str = "auto") -> str:
+ """'bf16' / 'torch.bfloat16' / mx.bfloat16 -> 'bfloat16'; '' -> default."""
+ if name is None or name == "":
+ return default
+ text = str(name).strip().lower().replace("torch.", "").replace("mlx.core.", "")
+ text = _ALIASES.get(text, text)
+ if text not in COMPUTE_DTYPES + ("auto",):
+ raise ValueError(f"unsupported compute dtype {name!r}; expected one of {COMPUTE_DTYPES} or 'auto'")
+ return text
+
+
+@dataclass(frozen=True)
+class PrecisionPolicy:
+ backend: str # "mlx" | "torch"
+ device: str # "gpu" | "cpu" | "mps" | "cuda"
+ compute: str # weights, activations, KV cache, expert pages
+ master: str = "float32" # optimizer master weights + moments
+ accumulate: str = "float32" # gradient accumulation
+ requested: str = "auto"
+ reason: str = ""
+ fused_fp32: tuple = field(default=FUSED_FP32_SITES)
+
+ @property
+ def kv_cache(self) -> str:
+ return self.compute
+
+ @property
+ def expert_storage(self) -> str:
+ return self.compute
+
+ @property
+ def low_precision(self) -> bool:
+ return self.compute != "float32"
+
+ def as_dict(self) -> dict:
+ out = asdict(self)
+ out.update(kv_cache=self.kv_cache, expert_storage=self.expert_storage,
+ boundary_casts=dict(BOUNDARY_CASTS), fused_fp32=list(self.fused_fp32))
+ return out
+
+ # ---- backend dtype objects -------------------------------------------
+ def mx_dtype(self, which: str = "compute"):
+ import mlx.core as mx
+ return getattr(mx, getattr(self, which))
+
+ def torch_dtype(self, which: str = "compute"):
+ import torch
+ return getattr(torch, getattr(self, which))
+
+
+def resolve_mlx(requested="auto", device: str | None = None) -> PrecisionPolicy:
+ """Pick the MLX compute dtype once for a whole run.
+
+ GPU: BF16 by default (FP32 masters keep small updates). CPU: FP32,
+ because MLX CPU GatherMM (sparse expert dispatch) is FP32-only; a
+ model-wide FP32 run is exact and avoids per-matmul upcasts.
+ """
+ import mlx.core as mx
+ req = normalize_dtype_name(requested)
+ dev = device or ("cpu" if mx.default_device().type == mx.DeviceType.cpu else "gpu")
+ if dev == "cpu":
+ reason = ("MLX CPU GatherMM supports float32 only; whole model runs in float32"
+ if req not in ("auto", "float32") else "MLX CPU backend has no fast low-precision GEMM")
+ return PrecisionPolicy("mlx", "cpu", "float32", requested=req, reason=reason)
+ if req == "auto":
+ return PrecisionPolicy("mlx", dev, "bfloat16", requested=req, reason="GPU default: BF16 compute, FP32 masters")
+ return PrecisionPolicy("mlx", dev, req, requested=req, reason="requested")
+
+
+def resolve_torch(requested="auto", device: str = "cpu") -> PrecisionPolicy:
+ """Pick the torch compute dtype once for a whole run or served model."""
+ import torch
+ req = normalize_dtype_name(requested)
+ dev = str(device).split(":")[0]
+ if req != "auto":
+ if dev == "cuda" and req == "bfloat16" and not torch.cuda.is_bf16_supported():
+ return PrecisionPolicy("torch", dev, "float16", requested=req, reason="GPU lacks BF16; using FP16")
+ return PrecisionPolicy("torch", dev, req, requested=req, reason="requested")
+ if dev == "cuda":
+ dt = "bfloat16" if torch.cuda.is_bf16_supported() else "float16"
+ return PrecisionPolicy("torch", dev, dt, requested=req, reason="CUDA default")
+ if dev == "mps":
+ # Measured on M5 Max / torch 2.10 (Chronos 50M, native dtype, no
+ # autocast): BF16 is slower than FP32 for training fwd+bwd (0.749 vs
+ # 0.619 s/step) and cached decode (49.1 vs 37.3 ms/token). BF16
+ # stays available on request (halves weight/KV memory).
+ return PrecisionPolicy("torch", dev, "float32", requested=req,
+ reason="MPS default: FP32 measured faster than BF16")
+ return PrecisionPolicy("torch", dev, "float32", requested=req, reason="CPU default: float32 is the fast CPU GEMM")
+
+
+def place_torch_model(model, device: str, requested="auto") -> tuple[object, PrecisionPolicy]:
+ """Move a torch Chronos model to ``device`` in the policy compute dtype.
+
+ One conversion at load time (weights + any placeholders); afterwards the
+ forward pass runs in that dtype with no per-op casts. Call this before
+ an ``ExpertStore`` is built so SSD/RAM expert pages adopt the same dtype.
+ """
+ from chronos.model.moe_chronos import LazyExpertPlaceholder
+ policy = resolve_torch(requested, device)
+ dtype = policy.torch_dtype()
+ model = model.to(device=device, dtype=dtype)
+ for module in model.modules():
+ if isinstance(module, LazyExpertPlaceholder):
+ module.dtype = dtype
+ model._chronos_precision = policy
+ return model, policy
+
+
+def torch_model_dtype(model):
+ """Compute dtype of a torch Chronos model (its embedding weight)."""
+ inner = getattr(model, "model", model)
+ emb = getattr(inner, "embed_tokens", None)
+ if emb is not None:
+ return emb.weight.dtype
+ return next(model.parameters()).dtype
+
+
+# ---- exact MLX <-> torch tensor bridges ------------------------------------
+# numpy has no bfloat16, so a BF16 tensor crosses as its uint16 bit pattern
+# (exact, no FP32 detour and no second rounding). FP16/FP32 cross natively.
+
+def mx_to_torch(arr, dtype: str | None = None):
+ """mx.array -> CPU torch tensor, keeping (or converting once to) ``dtype``."""
+ import mlx.core as mx
+ import numpy as np
+ import torch
+ if dtype is not None:
+ target = getattr(mx, normalize_dtype_name(dtype, "float32"))
+ if arr.dtype != target:
+ arr = arr.astype(target)
+ if arr.dtype == mx.bfloat16:
+ bits = arr.view(mx.uint16)
+ mx.eval(bits)
+ return torch.from_numpy(np.array(bits, copy=True)).view(torch.bfloat16)
+ mx.eval(arr)
+ return torch.from_numpy(np.array(arr, copy=True))
+
+
+def torch_to_mx(tensor, dtype: str | None = None):
+ """torch tensor -> mx.array, keeping (or converting once to) ``dtype``."""
+ import mlx.core as mx
+ import torch
+ t = tensor.detach().cpu()
+ if dtype is not None:
+ target = getattr(torch, normalize_dtype_name(dtype, "float32"))
+ if t.dtype != target:
+ t = t.to(target)
+ if t.dtype == torch.bfloat16:
+ return mx.array(t.contiguous().view(torch.uint16).numpy()).view(mx.bfloat16)
+ return mx.array(t.contiguous().numpy())
diff --git a/chronos/runtime/cache_manager.py b/chronos/runtime/cache_manager.py
index 2fabaff..86ed02e 100644
--- a/chronos/runtime/cache_manager.py
+++ b/chronos/runtime/cache_manager.py
@@ -35,7 +35,12 @@ class CacheManager:
def __init__(self, model, config, ssd_dir: str = "./expert_cache"):
self.config = config
self.expert_store = ExpertStore(model, config, ssd_dir=ssd_dir)
- self.prefetcher = AsyncPrefetcher(self.expert_store)
+ # The queue is part of the memory budget: a large lookahead queue can
+ # retain many decoded expert ids and defeat the RAM cap. Tie it to
+ # the canonical prefetch_depth setting while keeping a small minimum
+ # so a zero-depth diagnostic configuration remains usable.
+ queue_depth = max(1, int(getattr(config, "prefetch_depth", 2) or 0) * 2)
+ self.prefetcher = AsyncPrefetcher(self.expert_store, queue_depth=queue_depth)
self.scheduler = PrefetchScheduler(self.prefetcher, self.expert_store)
self._num_layers = len(self.expert_store.moe_layers)
@@ -62,9 +67,9 @@ def warm_up(self, initial_expert_ids: Optional[List[int]] = None):
Defaults to experts 0..vram_capacity-1.
"""
if initial_expert_ids is None:
- initial_expert_ids = list(range(
- min(self.config.num_experts, self.expert_store.vram_capacity)
- ))
+ initial_expert_ids = self.expert_store._expand_keys(
+ list(range(self.config.num_experts))
+ )[:self.expert_store.vram_capacity]
self.expert_store.prefetch_to_ram(initial_expert_ids)
for eid in initial_expert_ids:
self.expert_store.promote_to_vram(eid)
@@ -76,9 +81,8 @@ def availability_mask(self) -> torch.Tensor:
return self.expert_store.vram_availability_mask()
def availability_masks_all_layers(self) -> List[torch.Tensor]:
- """Returns the same mask replicated for each MoE layer."""
- mask = self.availability_mask()
- return [mask] * self._num_layers
+ """Return actual per-layer masks (v1 stores return cloned globals)."""
+ return self.expert_store.vram_availability_masks()
def step(self, lookahead_probs, current_expert_ids: List[int]):
"""
@@ -99,7 +103,13 @@ def prefetch_for_next_step(self, lookahead_probs) -> List[int]:
def prefetch_experts_to_ram(self, expert_ids: List[int]) -> List[int]:
"""Queue exact expert IDs for SSD->RAM prefetch without touching VRAM."""
- expert_ids = list(dict.fromkeys(int(eid) for eid in expert_ids or []))
+ normalized = []
+ for eid in expert_ids or []:
+ if isinstance(eid, (tuple, list)):
+ normalized.append(tuple(int(x) for x in eid))
+ else:
+ normalized.append(int(eid))
+ expert_ids = list(dict.fromkeys(normalized))
if expert_ids:
self.prefetcher.submit(expert_ids)
return expert_ids
@@ -114,7 +124,7 @@ def ensure_resident(self, expert_ids: List[int]) -> List[int]:
# The compute stream waits on each newly-promoted expert's event;
# already-resident experts contribute no event, so this only blocks
# on the truly necessary transfers.
- self.expert_store.wait_for_experts(newly)
+ self.expert_store.wait_for_experts(expert_ids)
return newly
def stats(self) -> dict:
diff --git a/chronos/runtime/inference_engine.py b/chronos/runtime/inference_engine.py
index 98b2f20..fd0e1ec 100644
--- a/chronos/runtime/inference_engine.py
+++ b/chronos/runtime/inference_engine.py
@@ -17,11 +17,17 @@
import chronos.deps # ensure minimind on sys.path
import time
-from typing import Dict, List, Optional
+from typing import Dict, Hashable, List, Optional
import torch
import torch.nn.functional as F
+try:
+ from chronos.runtime.metrics import safe_record as _metric
+except Exception: # pragma: no cover
+ def _metric(_name, _value):
+ return None
+
from chronos.model.config import ChronosConfig
from chronos.model.model_chronos import ChronosForCausalLM
from chronos.model.moe_chronos import ChronosMOEFeedForward
@@ -90,7 +96,7 @@ def __init__(
self.cache_manager = CacheManager(model, config, ssd_dir=ssd_dir)
self.last_stats: dict = {}
self._runtime_stats: dict = {}
- self._route_cache_scores: Dict[int, float] = {}
+ self._route_cache_scores: Dict[Hashable, float] = {}
self._last_prefill_ram_expert_ids: List[int] = []
def _moe_layers(self) -> List[ChronosMOEFeedForward]:
@@ -118,21 +124,29 @@ def _normalize_miss_policy(miss_policy: str | None) -> str:
def _install_moe_runtime_hooks(self, miss_policy: str):
store = self.cache_manager.expert_store
- def load_one(expert_id: int) -> dict:
+ def cache_key(layer_id: int, expert_id: int):
+ return (
+ store.expert_key(layer_id, expert_id)
+ if bool(getattr(self.config, "layer_aware_cache", False))
+ else int(expert_id)
+ )
+
+ def load_one(layer_id: int, expert_id: int) -> dict:
expert_id = int(expert_id)
- already_vram = store.vram_lru.contains(expert_id)
- already_ram = store.ram_lru.contains(expert_id)
+ key = cache_key(layer_id, expert_id)
+ already_vram = store.vram_lru.contains(key)
+ already_ram = store.ram_lru.contains(key)
t0 = time.monotonic()
ok = False
try:
if already_vram or already_ram or miss_policy == "sync_on_demand":
- promoted = self.cache_manager.ensure_resident([expert_id])
- ok = store.vram_lru.contains(expert_id)
+ promoted = self.cache_manager.ensure_resident([key])
+ ok = store.vram_lru.contains(key)
else:
# README architecture: cold SSD misses must not enter the
# decode critical path. Queue the exact missing expert for
# RAM prefetch and let shared fallback cover this token.
- self.cache_manager.prefetch_experts_to_ram([expert_id])
+ self.cache_manager.prefetch_experts_to_ram([key])
promoted = []
ok = False
finally:
@@ -153,18 +167,21 @@ def load_one(expert_id: int) -> dict:
int(self._runtime_stats.get("sync_ssd_loads", 0)) + 1
)
if ok:
- self._runtime_stats["on_demand_loaded_experts"].append(expert_id)
+ self._runtime_stats["on_demand_loaded_experts"].append(key)
if promoted:
self._runtime_stats["on_demand_promoted_experts"].extend(promoted)
return {"ok": ok, "already_vram": already_vram, "already_ram": already_ram}
- def touch_one(expert_id: int) -> None:
- store.touch_expert(int(expert_id))
+ def touch_one(layer_id: int, expert_id: int) -> None:
+ store.touch_expert(cache_key(layer_id, int(expert_id)))
- for moe in self._moe_layers():
+ for layer_id, moe in enumerate(self._moe_layers()):
moe.runtime_miss_policy = miss_policy
- moe.runtime_on_demand_loader = load_one if miss_policy in {"on_demand", "sync_on_demand"} else None
- moe.runtime_touch_expert = touch_one
+ moe.runtime_on_demand_loader = (
+ (lambda expert_id, li=layer_id: load_one(li, expert_id))
+ if miss_policy in {"on_demand", "sync_on_demand"} else None
+ )
+ moe.runtime_touch_expert = lambda expert_id, li=layer_id: touch_one(li, expert_id)
def _clear_moe_runtime_hooks(self):
for moe in self._moe_layers():
@@ -188,28 +205,36 @@ def _consume_moe_runtime_stats(self) -> dict:
def _recent_route_scores(
self,
- current_ids: List[int],
- predicted_ids: List[int],
- ram_predicted_ids: Optional[List[int]] = None,
- ) -> Dict[int, float]:
+ current_ids: List[Hashable],
+ predicted_ids: List[Hashable],
+ ram_predicted_ids: Optional[List[Hashable]] = None,
+ ) -> Dict[Hashable, float]:
"""Decayed route scores used as a sticky eviction priority.
Higher scores stay hotter in both RAM and VRAM. This intentionally
spans tokens; replacing it with only the current step turns tiny
budgets into a promote/evict loop.
"""
+ def normalized(key):
+ if isinstance(key, (tuple, list)):
+ return tuple(int(value) for value in key)
+ return int(key)
+
decayed = {
- int(eid): float(score) * 0.85
+ normalized(eid): float(score) * 0.85
for eid, score in self._route_cache_scores.items()
if float(score) * 0.85 >= 0.05
}
scores = decayed
for eid in current_ids:
- scores[int(eid)] = scores.get(int(eid), 0.0) + 3.0
+ key = normalized(eid)
+ scores[key] = scores.get(key, 0.0) + 3.0
for eid in predicted_ids:
- scores[int(eid)] = scores.get(int(eid), 0.0) + 2.0
+ key = normalized(eid)
+ scores[key] = scores.get(key, 0.0) + 2.0
for eid in ram_predicted_ids or []:
- scores[int(eid)] = scores.get(int(eid), 0.0) + 1.0
+ key = normalized(eid)
+ scores[key] = scores.get(key, 0.0) + 1.0
self._route_cache_scores = scores
return scores
@@ -229,18 +254,18 @@ def _ordered_ids_from_scores(scores: torch.Tensor, capacity: int) -> List[int]:
ordered = torch.argsort(scores, descending=True).tolist()
return [int(eid) for eid in ordered if float(scores[eid].item()) > 0.0][:capacity]
- def _ordered_sticky_ids(self, capacity: int) -> List[int]:
+ def _ordered_sticky_ids(self, capacity: int) -> List[Hashable]:
if capacity <= 0 or not self._route_cache_scores:
return []
return [
- int(eid)
+ eid
for eid, _score in sorted(
self._route_cache_scores.items(),
- key=lambda item: (-float(item[1]), int(item[0])),
+ key=lambda item: (-float(item[1]), repr(item[0])),
)
][:capacity]
- def _promotion_ready_ids(self, expert_ids: List[int]) -> List[int]:
+ def _promotion_ready_ids(self, expert_ids: List[Hashable]) -> List[Hashable]:
"""Only return IDs already staged in RAM/VRAM.
Calling ensure_resident on a cold SSD expert during decode would
@@ -249,7 +274,11 @@ def _promotion_ready_ids(self, expert_ids: List[int]) -> List[int]:
"""
store = self.cache_manager.expert_store
ready = []
- for eid in dict.fromkeys(int(e) for e in expert_ids or []):
+ normalized = (
+ tuple(int(value) for value in eid) if isinstance(eid, (tuple, list)) else int(eid)
+ for eid in expert_ids or []
+ )
+ for eid in dict.fromkeys(normalized):
if store.vram_lru.contains(eid) or store.ram_lru.contains(eid):
ready.append(eid)
return ready
@@ -300,6 +329,25 @@ def setup_from_state_reader(
def teardown(self):
self.cache_manager.stop()
+ @staticmethod
+ def _virtual_router_probs(moe, last_token: bool = False) -> Optional[torch.Tensor]:
+ """Router probs of every dispatch of ``moe`` in the last forward.
+
+ A recurrent group's MoE runs once per loop iteration and each
+ iteration may pick different experts, so prefetch/residency must
+ cover the union over all D iterations (virtual layers), not only
+ the last one. Returns ``[B, S*D, E]`` (``[B, D, E]`` for
+ ``last_token``) on the CPU, or None before the first forward.
+ """
+ probs = list(getattr(moe, "iteration_router_probs", None) or ())
+ if not probs:
+ probs = [moe.last_router_probs] if moe.last_router_probs is not None else []
+ if not probs:
+ return None
+ if last_token:
+ probs = [p[:, -1:, :] for p in probs]
+ return torch.cat([p.detach() for p in probs], dim=1).cpu()
+
def _get_current_expert_ids(self, capacity: Optional[int] = None) -> List[int]:
"""Extract active top-k expert IDs from every MoE layer.
@@ -312,10 +360,10 @@ def _get_current_expert_ids(self, capacity: Optional[int] = None) -> List[int]:
return []
scores = torch.zeros(num_experts, dtype=torch.float32)
for moe in moe_layers:
- if moe.last_router_probs is None:
+ probs = self._virtual_router_probs(moe, last_token=True) # [B, D, E]
+ if probs is None:
continue
- probs = moe.last_router_probs[:, -1, :].detach().cpu() # [B, E]
- weights, ids = self._router_topk(probs)
+ weights, ids = self._router_topk(probs.reshape(-1, probs.shape[-1]))
scores.index_add_(0, ids.reshape(-1), weights.reshape(-1).to(scores.dtype))
if capacity is None:
capacity = num_experts
@@ -332,9 +380,43 @@ def _predict_expert_ids(
if lookahead_probs is None:
return []
probs = lookahead_probs.detach().cpu()
- if probs.dim() != 4 or probs.shape[-1] == 0:
+ if probs.dim() not in (4, 5) or probs.shape[-1] == 0:
return []
+ if probs.dim() == 5 and not bool(getattr(self.config, "layer_aware_cache", False)):
+ probs = probs.mean(dim=3)
num_experts = int(probs.shape[-1])
+ layer_aware = probs.dim() == 5 and bool(getattr(self.config, "layer_aware_cache", False))
+ if layer_aware:
+ # Return physical-layer page keys. The recurrent group is encoded
+ # by the prediction axis and ExpertStore validates the mapping.
+ group_axis = probs.shape[-2]
+ layer_map = getattr(self.cache_manager.expert_store, "_layer_to_group", {})
+ if all_future:
+ future = probs[:, -1, 1:, :, :]
+ offsets = range(future.shape[1])
+ else:
+ if probs.shape[2] <= future_offset:
+ return []
+ future = probs[:, -1, future_offset:future_offset + 1, :, :]
+ offsets = [0]
+ selected = []
+ for offset in offsets:
+ for group in range(group_axis):
+ step_probs = future[:, offset, group, :]
+ weights, ids = self._router_topk(step_probs)
+ scale = 1.0 / float(offset + 1) if all_future else 1.0
+ for eid, weight in zip(ids.reshape(-1).tolist(), weights.reshape(-1).tolist()):
+ selected.append(((group, int(eid)), float(weight) * scale))
+ score_map = {}
+ for (group, eid), score in selected:
+ score_map[(group, eid)] = score_map.get((group, eid), 0.0) + score
+ capacity = num_experts if capacity is None else int(capacity)
+ ordered = sorted(score_map, key=lambda key: (-score_map[key], key[0], key[1]))
+ out = []
+ for group, eid in ordered[:capacity]:
+ layers = [li for li, gi in layer_map.items() if gi == group]
+ out.extend((int(li), int(group), int(eid)) for li in layers)
+ return out
scores = torch.zeros(num_experts, dtype=torch.float32)
if all_future:
if probs.shape[2] <= 1:
@@ -375,18 +457,20 @@ def _route_cache_stats(
"fallback_weight_mass": 0.0,
"total_route_weight_mass": 0.0,
"activated_expert_ids": [],
+ "activated_expert_keys": [],
"per_layer": [],
}
if not moe_layers:
return totals
activated: set[int] = set()
+ activated_keys: set[tuple[int, int, int]] = set()
+ store = self.cache_manager.expert_store
for layer_idx, moe in enumerate(moe_layers):
- if moe.last_router_probs is None:
+ # Every loop iteration is a real dispatch: count all of them.
+ probs = self._virtual_router_probs(moe, last_token=not include_all_tokens)
+ if probs is None:
continue
- probs = moe.last_router_probs.detach().cpu()
- if not include_all_tokens:
- probs = probs[:, -1:, :]
flat_probs = probs.reshape(-1, probs.shape[-1])
weights, ids = self._router_topk(flat_probs)
flat_ids = ids.reshape(-1).to(torch.long)
@@ -415,6 +499,11 @@ def _route_cache_stats(
total_mass = hit_mass + miss_mass
activated.update(int(eid) for eid in flat_ids.unique().tolist())
+ group_idx = int(getattr(store, "_layer_to_group", {}).get(layer_idx, -1))
+ activated_keys.update(
+ (int(layer_idx), group_idx, int(eid))
+ for eid in flat_ids.unique().tolist()
+ )
totals["cache_hits"] += hits
totals["cache_misses"] += misses
totals["hit_weight_mass"] += hit_mass
@@ -431,12 +520,15 @@ def _route_cache_stats(
})
totals["activated_expert_ids"] = sorted(activated)
+ totals["activated_expert_keys"] = sorted(activated_keys)
return totals
def _plan_prefill_experts(
self,
input_ids: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
+ modality_inputs: Optional[dict] = None,
+ modality_mask: Optional[torch.Tensor] = None,
) -> tuple[List[int], List[int]]:
"""Probe prompt routing with shared fallback, then load hot experts.
@@ -483,6 +575,8 @@ def _plan_prefill_experts(
past_key_values=None,
use_cache=False,
available_expert_masks=masks,
+ modality_inputs=modality_inputs,
+ modality_mask=modality_mask,
)
self._runtime_stats["on_demand_loads"] = probe_start["on_demand_loads"]
self._runtime_stats["on_demand_load_time_s"] = probe_start["on_demand_load_time_s"]
@@ -498,10 +592,12 @@ def _plan_prefill_experts(
scores = torch.zeros(num_experts, dtype=torch.float32)
for layer in self.model.model.layers:
moe = getattr(layer, "mlp", None)
- if not isinstance(moe, ChronosMOEFeedForward) or moe.last_router_probs is None:
+ if not isinstance(moe, ChronosMOEFeedForward):
+ continue
+ probs = self._virtual_router_probs(moe)
+ if probs is None:
continue
- probs = moe.last_router_probs.reshape(-1, num_experts)
- weights, expert_ids = self._router_topk(probs.detach().cpu())
+ weights, expert_ids = self._router_topk(probs.reshape(-1, num_experts))
scores.index_add_(
0,
expert_ids.reshape(-1).to(torch.long),
@@ -521,7 +617,130 @@ def _plan_prefill_experts(
ram_ids = self._ordered_ids_from_scores(scores, ram_capacity)
return ram_ids[:vram_capacity], ram_ids
- @torch.inference_mode()
+ # no_grad, not inference_mode: experts paged in here are reused by
+ # generate(); inference tensors would poison later autograd-tracked calls.
+ @torch.no_grad()
+ def generate_speculative(
+ self,
+ input_ids: torch.Tensor,
+ *,
+ max_new_tokens: int = 256,
+ draft_depth: int = 1,
+ gamma: int = 4,
+ gamma_max: int = 8,
+ adaptive_gamma: bool = True,
+ temperature: float = 0.0,
+ top_p: float = 1.0,
+ top_k: int = 0,
+ do_sample: bool = False,
+ eos_token_id=None,
+ miss_policy: str = "on_demand",
+ generator: Optional[torch.Generator] = None,
+ on_token=None,
+ ) -> torch.Tensor:
+ """Recurrent-depth self-speculative decoding with lazy experts.
+
+ Batch size 1. Each draft and verify forward runs with the current
+ expert availability masks; before it, the lookahead router's
+ prediction from the previous forward is prefetched to RAM and the
+ promotable part moved to VRAM, so drafts warm the experts their own
+ verify pass needs. Greedy output equals :meth:`generate` greedy
+ output token for token. ``per_iteration`` / non-recurrent models
+ decode plainly (``draft_depth`` = full depth).
+ """
+ from chronos.runtime.speculative import speculative_generate, speculative_supported
+ input_ids = input_ids.to(self.device)
+ miss_policy = self._normalize_miss_policy(miss_policy)
+ full_depth = int(getattr(self.config, "max_inference_depth", None) or self.config.recurrent_depth)
+ use_spec = speculative_supported(self.config) and int(draft_depth) < full_depth
+ self._route_cache_scores = {}
+ self._runtime_stats = {
+ "on_demand_loads": 0, "on_demand_load_time_s": 0.0, "sync_ssd_loads": 0,
+ "async_cold_miss_prefetches": 0, "resident_vram_hits": 0, "resident_ram_hits": 0,
+ "on_demand_loaded_experts": [], "on_demand_promoted_experts": [],
+ }
+ cm = self.cache_manager
+ store = cm.expert_store
+ state = {"lookahead": None, "hits": 0, "misses": 0, "forwards": 0, "fallback_mass": 0.0,
+ "on_demand": 0, "prefill_s": None}
+ t0 = time.monotonic()
+ rss0 = _rss_gb()
+ backend0 = _backend_memory_snapshot(self.device)
+ self._install_moe_runtime_hooks(miss_policy)
+
+ def forward_fn(ids, cache, depth):
+ if state["lookahead"] is not None:
+ predicted = self._predict_expert_ids(state["lookahead"], all_future=True,
+ capacity=int(store.ram_capacity))
+ cm.prefetch_experts_to_ram(predicted)
+ ready = self._promotion_ready_ids(predicted[: int(store.vram_capacity)])
+ if ready:
+ cm.ensure_resident(ready)
+ masks = cm.availability_masks_all_layers()
+ out = self.model(ids.to(self.device), past_key_values=cache, use_cache=True,
+ recurrent_depth=depth, available_expert_masks=masks)
+ state["lookahead"] = out.lookahead_probs
+ route = self._route_cache_stats(masks, include_all_tokens=cache is None)
+ state["hits"] += int(route["cache_hits"])
+ state["misses"] += int(route["cache_misses"])
+ rt = self._consume_moe_runtime_stats()
+ state["fallback_mass"] += float(rt["fallback_executed_weight_mass"])
+ state["on_demand"] += len(rt["on_demand_experts"])
+ state["forwards"] += 1
+ if state["prefill_s"] is None:
+ state["prefill_s"] = time.monotonic() - t0
+ return out
+
+ try:
+ prefill_ids, prefill_ram_ids = self._plan_prefill_experts(input_ids, None, None, None)
+ if prefill_ram_ids:
+ store.prefetch_to_ram(prefill_ram_ids)
+ if prefill_ids:
+ cm.ensure_resident(prefill_ids)
+ out, spec = speculative_generate(
+ self.model, input_ids, max_new_tokens=int(max_new_tokens),
+ draft_depth=int(draft_depth) if use_spec else full_depth, gamma=int(gamma),
+ gamma_max=int(gamma_max), adaptive_gamma=bool(adaptive_gamma), do_sample=bool(do_sample),
+ temperature=max(float(temperature), 1e-6), top_k=int(top_k or 0), top_p=float(top_p or 1.0),
+ eos_token_id=eos_token_id, generator=generator, on_token=on_token, forward_fn=forward_fn)
+ finally:
+ self._clear_moe_runtime_hooks()
+ elapsed = time.monotonic() - t0
+ generated = int(out.shape[1] - input_ids.shape[1])
+ total = state["hits"] + state["misses"]
+ runtime_stats = dict(self._runtime_stats)
+ prefill_s = float(state["prefill_s"] or 0.0)
+ self.last_stats = {
+ "tokens_generated": generated,
+ "elapsed_s": round(elapsed, 3),
+ "tokens_per_sec": round(generated / max(elapsed, 1e-6), 2),
+ "prefill_time_s": round(prefill_s, 3),
+ "decode_time_s": round(max(0.0, elapsed - prefill_s), 3),
+ "rss_before_prefill_gb": round(rss0, 3),
+ "rss_after_decode_gb": round(_rss_gb(), 3),
+ "decode_rss_delta_gb": round(_rss_gb() - rss0, 3),
+ **_backend_memory_fields(backend0, "before_prefill"),
+ **_backend_memory_fields(_backend_memory_snapshot(self.device), "after_decode"),
+ "cache_hits": state["hits"], "cache_misses": state["misses"],
+ "cache_hit_rate": round(state["hits"] / max(total, 1), 4),
+ "resident_hit_rate": round(state["hits"] / max(total, 1), 4),
+ "fallback_weight_mass": round(state["fallback_mass"], 4),
+ "on_demand_loads": int(runtime_stats.get("on_demand_loads", 0)),
+ "sync_ssd_loads": int(runtime_stats.get("sync_ssd_loads", 0)),
+ "async_cold_miss_prefetches": int(runtime_stats.get("async_cold_miss_prefetches", 0)),
+ "resident_ram_hits": int(runtime_stats.get("resident_ram_hits", 0)),
+ "miss_policy": miss_policy,
+ "quality_safe": miss_policy == "sync_on_demand",
+ "forwards": state["forwards"],
+ "prefill_expert_ids": prefill_ids,
+ "prefill_ram_expert_ids": prefill_ram_ids,
+ "speculative": spec.as_dict() | {"enabled": use_spec, "full_depth": full_depth,
+ "draft_depth": int(draft_depth) if use_spec else full_depth,
+ "gamma_final": spec.gamma_history[-1] if spec.gamma_history else int(gamma)},
+ }
+ return out
+
+ @torch.no_grad()
def generate(
self,
input_ids: torch.Tensor,
@@ -536,6 +755,8 @@ def generate(
repetition_penalty: float = 1.0,
scheduler: Optional[PrefillScheduler] = None,
miss_policy: str = "on_demand",
+ modality_inputs: Optional[dict] = None,
+ modality_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
input_ids = input_ids.to(self.device)
if attention_mask is not None:
@@ -559,6 +780,7 @@ def generate(
# ── Prefill phase: front-load expert IO before generation ──────
prefill_t0 = time.monotonic()
prefill_mem0 = _rss_gb()
+ _io0 = dict(getattr(self.cache_manager.expert_store, "_io_stats", {}) or {})
prefill_backend_mem0 = _backend_memory_snapshot(self.device)
if scheduler is not None:
scheduler.prepare(input_ids, device=str(self.device))
@@ -566,7 +788,9 @@ def generate(
prefill_expert_ids: List[int] = []
prefill_ram_expert_ids: List[int] = []
else:
- prefill_expert_ids, prefill_ram_expert_ids = self._plan_prefill_experts(input_ids, attention_mask)
+ prefill_expert_ids, prefill_ram_expert_ids = self._plan_prefill_experts(
+ input_ids, attention_mask, modality_inputs, modality_mask
+ )
self._last_prefill_ram_expert_ids = prefill_ram_expert_ids
if prefill_ram_expert_ids:
self.cache_manager.expert_store.prefetch_to_ram(prefill_ram_expert_ids)
@@ -594,6 +818,8 @@ def generate(
fallback_executed_weight_mass = 0.0
prediction_hits = 0
prediction_total = 0
+ cluster_prediction_hits = 0
+ cluster_prediction_total = 0
# M3: holds last step's lookahead predictions; we submit them as
# prefetch BEFORE this step's forward so the H2D stream overlaps
# with compute. None on the first iteration (no predictions yet).
@@ -624,6 +850,10 @@ def generate(
capacity=vram_capacity,
)
predicted_current_expert_ids.extend(predicted_ids)
+ # Predicted pages must survive this step's evictions.
+ protect = getattr(store, "set_protected_experts", None)
+ if protect is not None:
+ protect(list(dict.fromkeys(predicted_ids + prev_current_ids)))
submitted_prefetch_ids.extend(
self.cache_manager.prefetch_experts_to_ram(ram_predicted_ids)
)
@@ -653,14 +883,18 @@ def generate(
else:
masks = self.cache_manager.availability_masks_all_layers()
- past_len = past_key_values[0][0].shape[1] if past_key_values else 0
- outputs, lookahead_probs = self.model(
+ from chronos.model.model_chronos import cache_sequence_length
+ past_len = cache_sequence_length(past_key_values) if past_key_values else 0
+ result = self.model(
input_ids[:, past_len:],
attention_mask=attention_mask,
past_key_values=past_key_values,
use_cache=True,
available_expert_masks=masks,
+ modality_inputs=modality_inputs,
+ modality_mask=modality_mask,
)
+ outputs, lookahead_probs = result.legacy_tuple()
if attention_mask is not None:
attention_mask = torch.cat(
@@ -677,17 +911,42 @@ def generate(
fallback_weight_mass += float(route_stats["fallback_weight_mass"])
total_route_weight_mass += float(route_stats["total_route_weight_mass"])
route_layer_stats = route_stats["per_layer"]
- current_ids = route_stats["activated_expert_ids"]
- activated_expert_ids.extend(current_ids)
+ current_ids = (
+ route_stats["activated_expert_keys"]
+ if bool(getattr(self.config, "layer_aware_cache", False))
+ else route_stats["activated_expert_ids"]
+ )
+ activated_expert_ids.extend(route_stats["activated_expert_ids"])
rt = self._consume_moe_runtime_stats()
on_demand_expert_ids.extend(rt["on_demand_experts"])
fallback_expert_ids.extend(rt["fallback_experts"])
on_demand_weight_mass += float(rt["on_demand_weight_mass"])
fallback_executed_weight_mass += float(rt["fallback_executed_weight_mass"])
if predicted_ids:
- current_set = set(int(eid) for eid in current_ids)
+ current_set = set(current_ids)
prediction_total += len(current_set)
- prediction_hits += len(current_set.intersection(int(eid) for eid in predicted_ids))
+ prediction_hits += len(current_set.intersection(predicted_ids))
+ cluster_storage = getattr(
+ self.cache_manager.expert_store, "_cluster_storage", None
+ )
+ if cluster_storage is not None:
+ manifest = cluster_storage.manifest
+ def _eid(key):
+ return int(key[2]) if isinstance(key, tuple) else int(key)
+ predicted_clusters = {
+ manifest.expert_to_cluster[_eid(key)]
+ for key in predicted_ids
+ if _eid(key) in manifest.expert_to_cluster
+ }
+ current_clusters = [
+ manifest.expert_to_cluster[_eid(key)]
+ for key in current_set
+ if _eid(key) in manifest.expert_to_cluster
+ ]
+ cluster_prediction_total += len(current_clusters)
+ cluster_prediction_hits += sum(
+ int(cid in predicted_clusters) for cid in current_clusters
+ )
if scheduler is None:
self.cache_manager.expert_store.reprioritize_resident_experts(
self._recent_route_scores(current_ids, predicted_ids, ram_predicted_ids)
@@ -703,7 +962,8 @@ def generate(
for i in range(input_ids.shape[0]):
logits[i, torch.unique(input_ids[i])] /= repetition_penalty
if top_k > 0:
- logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
+ k = min(int(top_k), int(logits.shape[-1]))
+ logits[logits < torch.topk(logits, k)[0][..., -1, None]] = -float('inf')
if top_p < 1.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
mask_p = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) > top_p
@@ -764,12 +1024,33 @@ def generate(
fallback_weight_rate = fallback_weight_mass / max(total_route_weight_mass, 1e-12)
executed_fallback_rate = fallback_executed_weight_mass / max(total_route_weight_mass, 1e-12)
prediction_hit_rate = prediction_hits / max(prediction_total, 1)
+ cluster_recall = cluster_prediction_hits / max(cluster_prediction_total, 1)
runtime_stats = dict(self._runtime_stats)
promoted_expert_ids.extend(runtime_stats.get("on_demand_promoted_experts", []))
backend_fields = {}
backend_fields.update(_backend_memory_fields(prefill_backend_mem0, "before_prefill"))
backend_fields.update(_backend_memory_fields(prefill_backend_mem1, "after_prefill"))
backend_fields.update(_backend_memory_fields(end_backend_mem, "after_decode"))
+ recurrent = getattr(self.model.model, "last_recurrent_state", None) or {}
+ depth_used = recurrent.get("depth_used")
+ average_ponder_steps = (
+ float(depth_used.float().mean().item()) if depth_used is not None else 1.0
+ )
+ early_exit = recurrent.get("early_exit_ratio")
+ early_exit_ratio = (
+ float(early_exit.item()) if isinstance(early_exit, torch.Tensor)
+ else float(early_exit or 0.0)
+ )
+ contraction = recurrent.get("contraction")
+ contraction_ratio = (
+ float(contraction.item()) if isinstance(contraction, torch.Tensor)
+ else float(contraction or 0.0)
+ )
+ on_demand_loads = int(runtime_stats.get("on_demand_loads", 0))
+ resident_ram_hits = int(runtime_stats.get("resident_ram_hits", 0))
+ cold_misses = int(runtime_stats.get("async_cold_miss_prefetches", 0)) + int(
+ runtime_stats.get("sync_ssd_loads", 0)
+ )
self.last_stats = {
"tokens_generated": tokens_generated,
"elapsed_s": round(elapsed, 3),
@@ -782,8 +1063,11 @@ def generate(
"rss_after_prefill_gb": round(prefill_mem1, 3),
"rss_after_decode_gb": round(end_mem, 3),
**backend_fields,
- "cache_hit_rate": round(1.0 if miss_policy == "on_demand" and executed_fallback_rate == 0.0 else resident_hit_rate, 4),
+ "cache_hit_rate": round(resident_hit_rate, 4),
"resident_hit_rate": round(resident_hit_rate, 4),
+ "resident_vram_hit_rate": round(resident_hit_rate, 4),
+ "ram_promotion_hit_rate": round(resident_ram_hits / max(on_demand_loads, 1), 4),
+ "ssd_cold_miss_rate": round(cold_misses / max(on_demand_loads, 1), 4),
"cache_hits": cache_hits,
"cache_misses": cache_misses,
"fallback_weight_rate": round(executed_fallback_rate, 4),
@@ -803,21 +1087,55 @@ def generate(
"quality_safe": miss_policy == "sync_on_demand",
"miss_policy": miss_policy,
"prediction_hit_rate": round(prediction_hit_rate, 4),
+ "prediction_topk_recall": round(prediction_hit_rate, 4),
+ "cluster_recall": round(cluster_recall, 4),
"prediction_hits": prediction_hits,
"prediction_total": prediction_total,
+ "cluster_prediction_hits": cluster_prediction_hits,
+ "cluster_prediction_total": cluster_prediction_total,
+ "average_ponder_steps": round(average_ponder_steps, 4),
+ "early_exit_ratio": round(early_exit_ratio, 4),
+ "recurrent_contraction": round(contraction_ratio, 6),
"route_layer_stats": route_layer_stats,
"prefill_expert_ids": prefill_expert_ids,
"prefill_ram_expert_ids": prefill_ram_expert_ids,
"activated_expert_ids": sorted(set(activated_expert_ids)),
- "predicted_current_expert_ids": sorted(set(predicted_current_expert_ids)),
- "prefetch_submitted_ids": sorted(set(submitted_prefetch_ids)),
- "promoted_expert_ids": sorted(set(promoted_expert_ids)),
- "on_demand_expert_ids": sorted(set(on_demand_expert_ids)),
- "fallback_expert_ids": sorted(set(fallback_expert_ids)),
+ "predicted_current_expert_ids": sorted(set(predicted_current_expert_ids), key=repr),
+ "prefetch_submitted_ids": sorted(set(submitted_prefetch_ids), key=repr),
+ "promoted_expert_ids": sorted(set(promoted_expert_ids), key=repr),
+ "on_demand_expert_ids": sorted(set(on_demand_expert_ids), key=repr),
+ "fallback_expert_ids": sorted(set(fallback_expert_ids), key=repr),
"prefetch_submitted_total": len(submitted_prefetch_ids),
"promoted_total": len(promoted_expert_ids),
**self.cache_manager.stats(),
}
+ store = self.cache_manager.expert_store
+ pages_read = int(getattr(store, "_io_stats", {}).get("ssd_page_reads", 0)) - int(_io0.get("ssd_page_reads", 0))
+ page_bytes = int(getattr(store, "_expert_bytes", 0) or 0)
+ ssd_mb = pages_read * page_bytes / 2 ** 20
+ self.last_stats.update({
+ "ssd_pages_read": pages_read,
+ "ssd_read_mb": round(ssd_mb, 3),
+ "ssd_read_mb_per_s": round(ssd_mb / max(elapsed, 1e-6), 3),
+ "resident_experts": int(self.last_stats.get("vram_experts", 0)),
+ })
+ if hasattr(store, "set_protected_experts"):
+ store.set_protected_experts(())
+ # Publish a compact rolling series for the IO Monitor without
+ # retaining per-request payloads in the hot path.
+ for metric_name, metric_value in (
+ ("tokens_per_sec", self.last_stats["tokens_per_sec"]),
+ ("resident_vram_hit_rate", self.last_stats["resident_vram_hit_rate"]),
+ ("ram_promotion_hit_rate", self.last_stats["ram_promotion_hit_rate"]),
+ ("ssd_cold_miss_rate", self.last_stats["ssd_cold_miss_rate"]),
+ ("prediction_topk_recall", self.last_stats["prediction_topk_recall"]),
+ ("cluster_recall", self.last_stats["cluster_recall"]),
+ ("average_ponder_steps", self.last_stats["average_ponder_steps"]),
+ ("early_exit_ratio", self.last_stats["early_exit_ratio"]),
+ ("recurrent_contraction", self.last_stats["recurrent_contraction"]),
+ ("fallback_weight_mass", self.last_stats["fallback_weight_mass"]),
+ ):
+ _metric(metric_name, metric_value)
if streamer:
streamer.end()
diff --git a/chronos/runtime/speculative.py b/chronos/runtime/speculative.py
new file mode 100644
index 0000000..ea65bcb
--- /dev/null
+++ b/chronos/runtime/speculative.py
@@ -0,0 +1,225 @@
+"""Recurrent-depth self-speculative decoding (no extra parameters).
+
+Draft: the same model run with fewer loop iterations (``draft_depth``)
+proposes γ tokens. Verify: one full-depth forward over the γ+1 positions
+scores them in a single batched pass (weights/experts read once). Greedy
+outputs are token-for-token identical to full-depth greedy decoding; sampled
+outputs use standard rejection sampling, which preserves the full-depth
+distribution exactly.
+
+Cache contract: draft KV is written to a throw-away cache; the persistent
+cache is only ever extended by verified forwards, then truncated to the
+accepted prefix (MLA slices ``c_kv``; sliding-window layers keep
+``window + γ_max`` entries so a rollback never loses in-window keys).
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+
+import torch
+
+
+def _truncate_entry(entry, length: int):
+ if entry is None:
+ return None
+ if len(entry) <= 2: # MLA latent cache (c_kv,) or (c_kv, k_rope); tensors or QuantizedKV
+ return tuple(t[:, :length] for t in entry)
+ k, v, end = entry
+ drop = int(end) - length
+ if drop < 0:
+ raise ValueError("cannot extend a cache by truncation")
+ if drop == 0:
+ return entry
+ if drop > k.shape[1]:
+ raise ValueError(f"rollback of {drop} exceeds retained sliding-window cache {k.shape[1]}")
+ return (k[:, :k.shape[1] - drop], v[:, :v.shape[1] - drop], length)
+
+
+def truncate_cache(cache, length: int):
+ return [_truncate_entry(e, length) for e in cache]
+
+
+def set_cache_slack(model, slack: int) -> None:
+ from chronos.model.hybrid_attention import SlidingWindowAttention
+ for module in model.modules():
+ if isinstance(module, SlidingWindowAttention):
+ module.cache_slack = max(int(getattr(module, "cache_slack", 0)), int(slack))
+
+
+def speculative_supported(cfg) -> bool:
+ """True when a shallow draft can share the KV stream with full-depth verify.
+
+ Needs a recurrent core deeper than one iteration and a KV layout whose
+ prefix does not depend on depth (``shared_first`` / ``first_only``).
+ ``per_iteration`` writes one KV slot per virtual layer, so callers fall
+ back to plain decoding (``draft_depth = full depth``), as the MLX loop does.
+ """
+ from chronos.config_schema import resolve_recurrent_attention_mode
+ depth = int(getattr(cfg, "max_inference_depth", None) or getattr(cfg, "recurrent_depth", 1) or 1)
+ recurrent = bool(getattr(cfg, "recurrent_enabled", True)) and depth > 1
+ return recurrent and resolve_recurrent_attention_mode(cfg) != "per_iteration"
+
+
+@dataclass
+class SpecStats:
+ proposed: int = 0
+ accepted: int = 0
+ verify_calls: int = 0
+ draft_calls: int = 0
+ generated: int = 0
+ elapsed_s: float = 0.0
+ gamma_history: list = field(default_factory=list)
+
+ @property
+ def acceptance_rate(self) -> float:
+ return self.accepted / self.proposed if self.proposed else 0.0
+
+ def as_dict(self) -> dict:
+ return {"proposed": self.proposed, "accepted": self.accepted,
+ "acceptance_rate": round(self.acceptance_rate, 4),
+ "verify_calls": self.verify_calls, "draft_calls": self.draft_calls,
+ "generated": self.generated,
+ "tokens_per_verify": round(self.generated / max(self.verify_calls, 1), 3),
+ "tokens_per_s": round(self.generated / self.elapsed_s, 2) if self.elapsed_s else None}
+
+
+def _probs(logits: torch.Tensor, temperature: float, top_k: int, top_p: float) -> torch.Tensor:
+ logits = logits.float() / max(temperature, 1e-6)
+ if top_k and top_k > 0:
+ kth = torch.topk(logits, min(int(top_k), logits.shape[-1]))[0][..., -1, None]
+ logits = logits.masked_fill(logits < kth, float("-inf"))
+ if top_p < 1.0:
+ sorted_logits, idx = torch.sort(logits, descending=True)
+ cum = torch.cumsum(torch.softmax(sorted_logits, -1), -1)
+ mask = cum > top_p
+ mask[..., 1:] = mask[..., :-1].clone()
+ mask[..., 0] = False
+ logits = logits.masked_fill(mask.scatter(-1, idx, mask), float("-inf"))
+ return torch.softmax(logits, dim=-1)
+
+
+def rejection_step(p: torch.Tensor, q: torch.Tensor, token: int, generator=None) -> tuple[bool, int]:
+ """Speculative-sampling acceptance (Leviathan et al. / Chen et al. 2023).
+
+ ``token`` was drawn from the draft distribution ``q``. Accept with
+ probability min(1, p/q); otherwise draw from norm(max(p − q, 0)). The
+ returned token is distributed exactly as ``p``.
+ """
+ ratio = (p[token] / q[token].clamp_min(1e-12)).clamp(max=1.0)
+ if bool(torch.rand((), generator=generator, device=p.device) <= ratio):
+ return True, int(token)
+ residual = (p - q).clamp_min(0)
+ total = residual.sum()
+ return False, int(torch.multinomial(residual / total if total > 0 else p, 1, generator=generator))
+
+
+@torch.no_grad() # not inference_mode: callers may page experts in and reuse them later
+def speculative_generate(model, input_ids: torch.Tensor, *, max_new_tokens: int = 128,
+ draft_depth: int = 1, gamma: int = 4, adaptive_gamma: bool = True,
+ gamma_max: int = 8, do_sample: bool = False, temperature: float = 1.0,
+ top_k: int = 0, top_p: float = 1.0, eos_token_id=None,
+ generator: torch.Generator | None = None, on_token=None,
+ forward_fn=None, should_stop=None):
+ """Return (output_ids [1, T], SpecStats). Batch size 1 (serving path).
+
+ Loop invariant: ``pending`` is final (already emitted) but not yet in the
+ persistent cache. Each round drafts from ``pending`` at shallow depth,
+ then one full-depth forward over ``[pending, d1..dγ]`` both writes
+ ``pending`` into the cache and verifies the drafts. A rejection's
+ correction (or the all-accepted bonus token) becomes the next pending
+ token, so no extra forward is ever needed.
+ """
+ from chronos.model.model_chronos import cache_sequence_length
+ if input_ids.shape[0] != 1:
+ raise ValueError("speculative_generate supports batch size 1")
+ cfg = model.config
+ full_depth = int(getattr(cfg, "max_inference_depth", None) or cfg.recurrent_depth)
+ draft_depth = max(1, min(int(draft_depth), full_depth))
+ eos = set([eos_token_id] if isinstance(eos_token_id, int) else (eos_token_id or []))
+ gamma_max = max(1, int(gamma_max))
+ gamma = max(1, min(int(gamma), gamma_max))
+ set_cache_slack(model, gamma_max + 1)
+ device = input_ids.device
+ stats, t0 = SpecStats(), time.perf_counter()
+ tokens = input_ids[0].tolist()
+ # ``forward_fn(ids, cache, depth)`` lets a runtime wrap every forward with
+ # expert residency work (SSD lazy loading, lookahead prefetch, masks).
+ if forward_fn is None:
+ def forward_fn(ids, cache, depth):
+ return model(ids, past_key_values=cache, use_cache=True, recurrent_depth=depth)
+
+ def dist(row):
+ return _probs(row, temperature, top_k, top_p)
+
+ def pick(row):
+ if not do_sample:
+ return int(torch.argmax(row))
+ return int(torch.multinomial(dist(row), 1, generator=generator))
+
+ def emit(tok):
+ tokens.append(tok)
+ stats.generated += 1
+ if on_token:
+ on_token(tok)
+
+ out = forward_fn(input_ids, None, None)
+ cache = out.past_key_values
+ pending = pick(out.logits[0, -1])
+ emit(pending)
+ while stats.generated < max_new_tokens and pending not in eos:
+ if should_stop is not None and should_stop():
+ break
+ committed = cache_sequence_length(cache)
+ room = max_new_tokens - stats.generated # tokens we may still emit
+ draft, draft_q, dcache, last = [pending], [], cache, pending
+ if draft_depth < full_depth:
+ for _ in range(min(gamma, room - 1) if room > 1 else 0):
+ d = forward_fn(torch.tensor([[last]], device=device), dcache, draft_depth)
+ stats.draft_calls += 1
+ dcache = d.past_key_values
+ row = d.logits[0, -1]
+ last = pick(row)
+ if do_sample:
+ draft_q.append(dist(row))
+ draft.append(last)
+ if last in eos:
+ break
+ v = forward_fn(torch.tensor([draft], device=device), cache, None)
+ stats.verify_calls += 1
+ logits = v.logits[0] # logits[i] predicts the token after draft[i]
+ n_ok, nxt, stop = 0, None, False
+ for i in range(1, len(draft)):
+ row, tok = logits[i - 1], draft[i]
+ if not do_sample:
+ target = int(torch.argmax(row))
+ ok = target == tok
+ correction = target
+ else:
+ ok, correction = rejection_step(dist(row), draft_q[i - 1], tok, generator)
+ if not ok:
+ nxt = correction
+ break
+ n_ok += 1
+ emit(tok)
+ if tok in eos or stats.generated >= max_new_tokens:
+ stop = True
+ break
+ stats.proposed += len(draft) - 1
+ stats.accepted += n_ok
+ cache = truncate_cache(v.past_key_values, committed + 1 + n_ok)
+ if stop:
+ break
+ if nxt is None: # every draft accepted: bonus token
+ nxt = pick(logits[len(draft) - 1])
+ pending = nxt
+ emit(pending)
+ if adaptive_gamma and len(draft) > 1:
+ rate = n_ok / (len(draft) - 1)
+ if rate >= 0.8:
+ gamma = min(gamma_max, gamma + 1)
+ elif rate < 0.4:
+ gamma = max(1, gamma - 1)
+ stats.gamma_history.append(gamma)
+ stats.elapsed_s = time.perf_counter() - t0
+ return torch.tensor([tokens], device=device), stats
diff --git a/chronos/serving/__init__.py b/chronos/serving/__init__.py
index c4b67c9..0040f92 100644
--- a/chronos/serving/__init__.py
+++ b/chronos/serving/__init__.py
@@ -11,3 +11,14 @@
"is_available",
"HAS_VLLM",
]
+
+from .api_compat import (
+ message_text, modality_inputs_from_messages, resolve_modality_inputs, openai_chat_response,
+ ollama_chat_response, ollama_generate_response,
+)
+from .adapter import ChronosServingAdapter
+
+__all__ += [
+ "message_text", "modality_inputs_from_messages", "resolve_modality_inputs", "openai_chat_response",
+ "ollama_chat_response", "ollama_generate_response", "ChronosServingAdapter",
+]
diff --git a/chronos/serving/adapter.py b/chronos/serving/adapter.py
new file mode 100644
index 0000000..033728b
--- /dev/null
+++ b/chronos/serving/adapter.py
@@ -0,0 +1,121 @@
+"""Framework-neutral OpenAI/Ollama serving adapter.
+
+The adapter intentionally contains no FastAPI dependency. It accepts the
+standard request dictionaries used by either HTTP framework and delegates
+execution to :class:`ChronosInferenceEngine`; applications can mount the two
+methods in FastAPI, Flask, Gradio, or a worker queue without changing model
+semantics.
+"""
+from __future__ import annotations
+
+from typing import Any, Mapping, Optional
+
+import torch
+
+from .api_compat import (
+ message_text,
+ openai_chat_response,
+ ollama_chat_response,
+ ollama_generate_response,
+ resolve_modality_inputs,
+)
+
+
+class ChronosServingAdapter:
+ """Expose a loaded Chronos engine through standard chat/generate schemas."""
+
+ def __init__(self, engine, tokenizer, *, model_name: str = "chronos"):
+ self.engine = engine
+ self.tokenizer = tokenizer
+ self.model_name = model_name
+
+ def _encode(self, messages, prompt: Optional[str] = None):
+ text = prompt if prompt is not None else message_text(messages)
+ encoded = self.tokenizer(text, return_tensors="pt", add_special_tokens=True)
+ ids = encoded.input_ids
+ mask = encoded.get("attention_mask") if hasattr(encoded, "get") else getattr(encoded, "attention_mask", None)
+ modalities = resolve_modality_inputs(
+ messages,
+ image_size=getattr(self.engine.config, "image_size", 224),
+ audio_sample_rate=getattr(self.engine.config, "audio_sample_rate", 16000),
+ max_audio_seconds=getattr(self.engine.config, "max_audio_seconds", 30.0),
+ )
+ return ids, mask, modalities
+
+ def generate(self, *, messages=None, prompt=None, max_tokens: int = 128,
+ temperature: float = 0.85, top_p: float = 0.95,
+ top_k: int = 50, do_sample: bool = True,
+ miss_policy: str = "on_demand") -> tuple[str, dict]:
+ ids, mask, modalities = self._encode(messages or [], prompt=prompt)
+ output = self.engine.generate(
+ ids,
+ attention_mask=mask,
+ max_new_tokens=int(max_tokens),
+ temperature=float(temperature),
+ top_p=float(top_p),
+ top_k=int(top_k),
+ do_sample=bool(do_sample),
+ miss_policy=miss_policy,
+ modality_inputs=modalities or None,
+ modality_mask=(torch.ones(ids.shape[0]) if modalities else None),
+ )
+ generated = output[:, ids.shape[1]:]
+ text = self.tokenizer.decode(generated[0].tolist(), skip_special_tokens=True)
+ stats = dict(getattr(self.engine, "last_stats", {}) or {})
+ stats.setdefault("prompt_tokens", int(ids.shape[1]))
+ stats.setdefault("completion_tokens", int(generated.shape[1]))
+ return text, stats
+
+ def openai_chat_completion(self, request: Mapping[str, Any]) -> dict:
+ messages = request.get("messages") or []
+ text, stats = self.generate(
+ messages=messages,
+ max_tokens=request.get("max_tokens", 128),
+ temperature=request.get("temperature", 0.85),
+ top_p=request.get("top_p", 0.95),
+ top_k=request.get("top_k", 50),
+ do_sample=request.get("do_sample", True),
+ miss_policy=request.get("chronos_miss_policy", "on_demand"),
+ )
+ return openai_chat_response(
+ str(request.get("model", self.model_name)), text,
+ prompt_tokens=stats.get("prompt_tokens", 0),
+ completion_tokens=stats.get("completion_tokens", 0),
+ chronos_metrics=stats if request.get("include_chronos_metrics") else None,
+ )
+
+ def ollama_chat(self, request: Mapping[str, Any]) -> dict:
+ messages = request.get("messages") or []
+ options = request.get("options") or {}
+ text, stats = self.generate(
+ messages=messages,
+ max_tokens=request.get("num_predict", options.get("num_predict", 128)),
+ temperature=request.get("temperature", options.get("temperature", 0.85)),
+ top_p=request.get("top_p", options.get("top_p", 0.95)),
+ top_k=request.get("top_k", options.get("top_k", 50)),
+ do_sample=bool(request.get("do_sample", True)),
+ miss_policy=request.get("chronos_miss_policy", "on_demand"),
+ )
+ return ollama_chat_response(
+ str(request.get("model", self.model_name)), text,
+ prompt_tokens=stats.get("prompt_tokens", 0),
+ completion_tokens=stats.get("completion_tokens", 0),
+ chronos_metrics=stats if request.get("include_chronos_metrics") else None,
+ )
+
+ def ollama_generate(self, request: Mapping[str, Any]) -> dict:
+ text, stats = self.generate(
+ prompt=request.get("prompt", ""),
+ max_tokens=request.get("options", {}).get("num_predict", 128),
+ temperature=request.get("options", {}).get("temperature", 0.85),
+ top_p=request.get("options", {}).get("top_p", 0.95),
+ top_k=request.get("options", {}).get("top_k", 50),
+ do_sample=bool(request.get("do_sample", True)),
+ miss_policy=request.get("chronos_miss_policy", "on_demand"),
+ )
+ return ollama_generate_response(
+ str(request.get("model", self.model_name)), text,
+ prompt_tokens=stats.get("prompt_tokens", 0),
+ completion_tokens=stats.get("completion_tokens", 0),
+ chronos_metrics=stats if request.get("include_chronos_metrics") else None,
+ )
diff --git a/chronos/serving/api_compat.py b/chronos/serving/api_compat.py
new file mode 100644
index 0000000..39a26df
--- /dev/null
+++ b/chronos/serving/api_compat.py
@@ -0,0 +1,131 @@
+"""OpenAI/Ollama request and response compatibility helpers."""
+from __future__ import annotations
+
+import time
+import uuid
+import base64
+import io
+import torch
+from typing import Iterable, Optional
+
+
+def message_text(messages: Iterable[dict]) -> str:
+ parts = []
+ for message in messages or []:
+ content = message.get("content", "")
+ if isinstance(content, list):
+ content = "\n".join(str(item.get("text", "")) for item in content
+ if isinstance(item, dict) and item.get("type") in {"text", "input_text"})
+ parts.append(f"{message.get('role', 'user')}: {content}")
+ return "\n".join(parts)
+
+
+def modality_inputs_from_messages(messages: Iterable[dict]) -> dict:
+ modalities = {"image": [], "audio": []}
+ for message in messages or []:
+ content = message.get("content", [])
+ if not isinstance(content, list):
+ continue
+ for item in content:
+ if not isinstance(item, dict):
+ continue
+ kind = item.get("type")
+ if kind in {"image", "image_url", "input_image"}:
+ value = item.get("image") or item.get("image_url")
+ if isinstance(value, dict):
+ value = value.get("url")
+ if value:
+ modalities["image"].append(value)
+ elif kind in {"audio", "input_audio"}:
+ value = item.get("audio") or item.get("input_audio")
+ if isinstance(value, dict):
+ # OpenAI input_audio uses {data, format}; retain the
+ # format only as metadata and decode the base64 payload
+ # through the same WAV/PCM resolver as JSONL training.
+ value = value.get("data") or value.get("url")
+ if value and not str(value).startswith("data:"):
+ value = f"data:audio/wav;base64,{value}"
+ if value:
+ modalities["audio"].append(value)
+ return {key: value for key, value in modalities.items() if value}
+
+
+def resolve_modality_inputs(messages: Iterable[dict], *, image_size: int = 224,
+ audio_sample_rate: int = 16000,
+ max_audio_seconds: float = 30.0) -> dict:
+ """Decode API URL/data-URI media into tensors accepted by ChronosFusion.
+
+ Local paths and HTTP URLs are delegated to the dataset loaders so serving
+ and training share identical normalization. Invalid optional media is
+ omitted, allowing a text fallback rather than crashing a standard client.
+ """
+ from chronos.data.flexible_dataset import _load_audio, _load_image
+ out = {}
+ raw = modality_inputs_from_messages(messages)
+ for kind, values in raw.items():
+ decoded = []
+ for value in values:
+ try:
+ if isinstance(value, str) and value.startswith("data:") and "," in value:
+ head, payload = value.split(",", 1)
+ blob = base64.b64decode(payload)
+ if kind == "image":
+ from PIL import Image
+ decoded.append(_load_image(Image.open(io.BytesIO(blob)), image_size))
+ else:
+ # WAV is the portable baseline supported by the
+ # dependency-light training loader. Unsupported
+ # codecs fail closed and become a text fallback.
+ decoded.append(_load_audio(blob, audio_sample_rate, max_audio_seconds))
+ elif kind == "image":
+ decoded.append(_load_image(value, image_size))
+ else:
+ decoded.append(_load_audio(value, audio_sample_rate, max_audio_seconds))
+ except Exception:
+ continue
+ if decoded:
+ if kind == "image":
+ # [B,N,C,H,W] lets PatchImageEncoder treat multiple images as
+ # one request; a single image remains semantically identical.
+ out[kind] = torch.stack(decoded).unsqueeze(0)
+ else:
+ # The baseline ConvAudioEncoder consumes one waveform per
+ # request. Preserve deterministic first-item semantics (the
+ # dataset path follows the same policy) instead of treating
+ # the number of clips as convolution channels.
+ out[kind] = decoded[0].flatten().unsqueeze(0)
+ return out
+
+
+def openai_chat_response(model: str, text: str, *, prompt_tokens: int = 0,
+ completion_tokens: int = 0, chronos_metrics: Optional[dict] = None) -> dict:
+ result = {
+ "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion",
+ "created": int(time.time()), "model": model,
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": int(prompt_tokens), "completion_tokens": int(completion_tokens), "total_tokens": int(prompt_tokens + completion_tokens)},
+ }
+ if chronos_metrics is not None:
+ result["chronos_metrics"] = chronos_metrics
+ return result
+
+
+def ollama_chat_response(model: str, text: str, *, done: bool = True,
+ prompt_tokens: int = 0, completion_tokens: int = 0,
+ chronos_metrics: Optional[dict] = None) -> dict:
+ result = {
+ "model": model, "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+ "message": {"role": "assistant", "content": text}, "done": bool(done),
+ "done_reason": "stop" if done else None, "prompt_eval_count": int(prompt_tokens),
+ "eval_count": int(completion_tokens),
+ }
+ if chronos_metrics is not None:
+ result["chronos_metrics"] = chronos_metrics
+ return result
+
+
+def ollama_generate_response(model: str, text: str, **kwargs) -> dict:
+ result = ollama_chat_response(model, text, **kwargs)
+ result["response"] = text
+ result.pop("message", None)
+ return result
diff --git a/chronos/serving/ollama_server.py b/chronos/serving/ollama_server.py
new file mode 100644
index 0000000..8b2a8e5
--- /dev/null
+++ b/chronos/serving/ollama_server.py
@@ -0,0 +1,353 @@
+"""Ollama- and OpenAI-compatible HTTP server for Chronos models.
+
+The stock ``ollama`` CLI and any Ollama/OpenAI client can talk to the native
+Chronos runtime by pointing at this server::
+
+ chronos serve --models-dir runs/enzh_50m_v1 --port 11435
+ OLLAMA_HOST=127.0.0.1:11435 ollama run chronos_enzh_sft "Hello world"
+
+The kernel is never converted to another architecture: recurrent depth,
+SSD-lazy experts with lookahead prefetch, compressed KV and self-speculative
+decoding all stay active behind the protocol.
+
+Security: binds 127.0.0.1 and has no authentication by default. Binding a
+non-loopback address requires ``--allow-remote`` *and* ``--api-key``.
+``--monitor-runs-dir`` mounts the read-only training monitor at ``/monitor/``;
+it reuses the same key through its own login page (session cookie).
+"""
+from __future__ import annotations
+
+import hashlib
+import hmac
+import json
+import os
+import time
+from datetime import datetime, timezone
+
+try: # module level: FastAPI resolves ``Request`` annotations by name
+ from fastapi import FastAPI, HTTPException, Request
+ from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
+except ImportError: # pragma: no cover - serving extra not installed
+ FastAPI = HTTPException = Request = JSONResponse = PlainTextResponse = StreamingResponse = None
+
+OLLAMA_COMPAT_VERSION = "0.40.14"
+
+
+def _now() -> str:
+ return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+
+
+def _digest(path: str) -> str:
+ st = os.stat(path)
+ return hashlib.sha256(f"{path}:{st.st_size}:{st.st_mtime_ns}".encode()).hexdigest()
+
+
+def _options(body: dict) -> dict:
+ opts = dict(body.get("options") or {})
+ return {
+ "max_new_tokens": int(opts.get("num_predict", body.get("max_tokens", 256)) or 256),
+ "temperature": float(opts.get("temperature", body.get("temperature", 0.0)) or 0.0),
+ "top_k": int(opts.get("top_k", body.get("top_k", 0)) or 0),
+ "top_p": float(opts.get("top_p", body.get("top_p", 1.0)) or 1.0),
+ "seed": opts.get("seed", body.get("seed")),
+ "stop": opts.get("stop", body.get("stop")) or [],
+ "speculative": bool(opts.get("chronos_speculative", True)),
+ "draft_depth": int(opts.get("chronos_draft_depth", 1)),
+ "gamma": int(opts.get("chronos_gamma", 4)),
+ "miss_policy": str(opts.get("chronos_miss_policy", "on_demand")),
+ }
+
+
+def _details(lm_or_cfg) -> dict:
+ cfg = getattr(lm_or_cfg, "config", lm_or_cfg)
+ return {"parent_model": "", "format": "chronos", "family": "chronos", "families": ["chronos"],
+ "parameter_size": f"{getattr(cfg, '_total_params_m', 0) or 0:.1f}M" if hasattr(cfg, "_total_params_m") else "",
+ "quantization_level": "BF16"}
+
+
+def _final_ollama(stats: dict, model: str, total_t0: float, load_s: float) -> dict:
+ return {
+ "model": model, "created_at": _now(), "done": True,
+ "done_reason": stats.get("done_reason", "stop"),
+ "total_duration": int((time.perf_counter() - total_t0) * 1e9),
+ "load_duration": int(load_s * 1e9),
+ "prompt_eval_count": stats.get("prompt_tokens", 0), "prompt_eval_duration": 0,
+ "eval_count": stats.get("completion_tokens", 0),
+ "eval_duration": stats.get("eval_duration_ns", 0),
+ "chronos": {k: v for k, v in stats.items() if k not in ("done", "done_reason")},
+ }
+
+
+def create_app(runtime, *, api_key: str | None = None, monitor_runs_dir: str | None = None):
+ if FastAPI is None:
+ raise ImportError("chronos serve needs fastapi and uvicorn (pip install 'Project_Chronos[serve]')")
+ app = FastAPI(title="Chronos Ollama-compatible server", version=OLLAMA_COMPAT_VERSION)
+ expected = f"Bearer {api_key}".encode("utf-8") if api_key else b""
+
+ @app.middleware("http")
+ async def auth(request: Request, call_next):
+ path = request.url.path
+ # /monitor has its own login (cookie or Bearer) with the same key.
+ exempt = path in ("/", "/api/version") or (monitor_runs_dir and (path == "/monitor" or path.startswith("/monitor/")))
+ if api_key and not exempt:
+ given = request.headers.get("authorization", "").encode("utf-8")
+ if not hmac.compare_digest(given, expected):
+ return JSONResponse({"error": "unauthorized"}, status_code=401)
+ return await call_next(request)
+
+ if monitor_runs_dir:
+ from fastapi.responses import RedirectResponse
+ from chronos.monitor.server import create_monitor_app
+
+ @app.get("/monitor", include_in_schema=False)
+ def monitor_redirect():
+ return RedirectResponse("/monitor/")
+
+ app.mount("/monitor", create_monitor_app(monitor_runs_dir, api_key=api_key))
+
+ def _resolve(name: str):
+ try:
+ t0 = time.perf_counter()
+ lm = runtime.get(name)
+ return lm, time.perf_counter() - t0
+ except KeyError as exc:
+ raise HTTPException(status_code=404, detail=f"model '{name}' not found, try pulling it first") from exc
+
+ def _model_entry(name, path):
+ st = os.stat(path)
+ return {"name": f"{name}:latest", "model": f"{name}:latest",
+ "modified_at": datetime.fromtimestamp(st.st_mtime, timezone.utc).isoformat().replace("+00:00", "Z"),
+ "size": st.st_size, "digest": _digest(path),
+ "details": {"parent_model": "", "format": "chronos", "family": "chronos",
+ "families": ["chronos"], "parameter_size": "", "quantization_level": "BF16"}}
+
+ @app.get("/")
+ def root():
+ return PlainTextResponse("Ollama is running")
+
+ @app.head("/")
+ def root_head():
+ return PlainTextResponse("")
+
+ @app.get("/api/version")
+ def version():
+ return {"version": OLLAMA_COMPAT_VERSION}
+
+ @app.get("/api/tags")
+ def tags():
+ return {"models": [_model_entry(n, p) for n, p in sorted(runtime.available().items())]}
+
+ @app.post("/api/show")
+ async def show(request: Request):
+ body = await request.json()
+ name = (body.get("model") or body.get("name") or "").split(":")[0]
+ path = runtime.available().get(name)
+ if path is None:
+ raise HTTPException(status_code=404, detail=f"model '{name}' not found")
+ meta = json.load(open(os.path.splitext(path)[0] + ".config.json", encoding="utf-8"))
+ cfg = meta.get("config", {})
+ info = {"general.architecture": "chronos", "chronos.recurrent_depth": cfg.get("recurrent_depth"),
+ "chronos.recurrent_attention_mode": cfg.get("recurrent_attention_mode"),
+ "chronos.num_experts": cfg.get("num_experts"), "chronos.experts_per_token": cfg.get("num_experts_per_tok"),
+ "chronos.kv_latent_dim": cfg.get("kv_latent_dim"), "chronos.context_length": cfg.get("max_position_embeddings"),
+ "chronos.hidden_size": cfg.get("hidden_size"), "chronos.vocab_size": cfg.get("vocab_size")}
+ return {"modelfile": f"FROM {path}\n", "parameters": "temperature 0\n", "template": "{{ .Prompt }}",
+ "details": _model_entry(name, path)["details"], "model_info": info,
+ "capabilities": ["completion"], "modified_at": _model_entry(name, path)["modified_at"]}
+
+ @app.get("/api/ps")
+ def ps():
+ out = []
+ for name, lm in list(runtime.loaded.items()):
+ store = getattr(getattr(lm.engine, "cache_manager", None), "expert_store", None)
+ st = store.stats() if store is not None else {}
+ entry = _model_entry(name, lm.path)
+ entry.update(size_vram=int(st.get("vram_bytes", 0) or 0),
+ expires_at=datetime.fromtimestamp(runtime.last_used.get(name, time.time()) + runtime.keep_alive_s,
+ timezone.utc).isoformat().replace("+00:00", "Z"),
+ chronos={"device": lm.device, "lazy_experts": lm.lazy_experts, "expert_store": st})
+ out.append(entry)
+ return {"models": out}
+
+ @app.post("/api/pull")
+ async def pull(request: Request):
+ body = await request.json()
+ name = (body.get("model") or body.get("name") or "").split(":")[0]
+ if name not in runtime.available():
+ return JSONResponse({"error": f"model '{name}' is not in --models-dir; Chronos serves local checkpoints"},
+ status_code=404)
+ return {"status": "success"}
+
+ def _ndjson(gen):
+ for item in gen:
+ yield json.dumps(item, ensure_ascii=False) + "\n"
+
+ def _run(body, input_ids, lm, load_s, shape):
+ opts = _options(body)
+ t0 = time.perf_counter()
+ name = lm.name
+ def events():
+ for item in runtime.generate_stream(name, input_ids, **opts):
+ if item.get("done"):
+ final = _final_ollama(item, body.get("model", name), t0, load_s)
+ if "error" in item:
+ final["error"] = item["error"]
+ yield shape(None, final)
+ else:
+ yield shape(item["text"], None)
+ return events
+
+ def _respond(body, events, collect):
+ if body.get("stream", True):
+ return StreamingResponse(_ndjson(events()), media_type="application/x-ndjson")
+ return JSONResponse(collect(list(events())))
+
+ def _unload_if_requested(body, name) -> bool:
+ keep = body.get("keep_alive")
+ if keep in (0, "0", "0s", "0m") and not body.get("prompt") and not body.get("messages"):
+ runtime.unload(name)
+ return True
+ return False
+
+ @app.post("/api/generate")
+ async def generate(request: Request):
+ body = await request.json()
+ name = (body.get("model") or "").split(":")[0]
+ if _unload_if_requested(body, name):
+ return {"model": body.get("model"), "created_at": _now(), "response": "", "done": True,
+ "done_reason": "unload"}
+ lm, load_s = _resolve(name)
+ if not body.get("prompt"): # Ollama "load model" request
+ return {"model": body.get("model"), "created_at": _now(), "response": "", "done": True,
+ "done_reason": "load"}
+ ids = runtime.encode_prompt(lm, body["prompt"], body.get("system"), bool(body.get("raw")))
+
+ def shape(text, final):
+ if final is not None:
+ return {**final, "response": ""}
+ return {"model": body.get("model"), "created_at": _now(), "response": text, "done": False}
+
+ def collect(items):
+ return {**items[-1], "response": "".join(i.get("response", "") for i in items[:-1])}
+ return _respond(body, _run(body, ids, lm, load_s, shape), collect)
+
+ @app.post("/api/chat")
+ async def chat(request: Request):
+ body = await request.json()
+ name = (body.get("model") or "").split(":")[0]
+ messages = [{"role": m.get("role", "user"), "content": m.get("content", "")}
+ for m in body.get("messages") or []]
+ if _unload_if_requested(body, name):
+ return {"model": body.get("model"), "created_at": _now(),
+ "message": {"role": "assistant", "content": ""}, "done": True, "done_reason": "unload"}
+ lm, load_s = _resolve(name)
+ if not messages:
+ return {"model": body.get("model"), "created_at": _now(),
+ "message": {"role": "assistant", "content": ""}, "done": True, "done_reason": "load"}
+ ids = runtime.encode_chat(lm, messages)
+
+ def shape(text, final):
+ if final is not None:
+ return {**final, "message": {"role": "assistant", "content": ""}}
+ return {"model": body.get("model"), "created_at": _now(),
+ "message": {"role": "assistant", "content": text}, "done": False}
+
+ def collect(items):
+ text = "".join(i["message"]["content"] for i in items[:-1])
+ return {**items[-1], "message": {"role": "assistant", "content": text}}
+ return _respond(body, _run(body, ids, lm, load_s, shape), collect)
+
+ @app.get("/v1/models")
+ def openai_models():
+ return {"object": "list", "data": [{"id": n, "object": "model", "created": int(os.stat(p).st_mtime),
+ "owned_by": "chronos"} for n, p in sorted(runtime.available().items())]}
+
+ @app.post("/v1/chat/completions")
+ async def openai_chat(request: Request):
+ import uuid
+ body = await request.json()
+ name = (body.get("model") or "").split(":")[0]
+ lm, load_s = _resolve(name)
+ messages = [{"role": m.get("role", "user"), "content": m.get("content", "")} for m in body.get("messages") or []]
+ ids = runtime.encode_chat(lm, messages)
+ opts = _options({**body, "options": {"num_predict": body.get("max_tokens") or body.get("max_completion_tokens") or 256,
+ "temperature": body.get("temperature", 0.0), "top_p": body.get("top_p", 1.0),
+ "seed": body.get("seed"), "stop": body.get("stop"),
+ **(body.get("chronos") or {})}})
+ cid, created = f"chatcmpl-{uuid.uuid4().hex}", int(time.time())
+ if body.get("stream"):
+ def sse():
+ for item in runtime.generate_stream(name, ids, **opts):
+ if item.get("done"):
+ chunk = {"id": cid, "object": "chat.completion.chunk", "created": created, "model": name,
+ "choices": [{"index": 0, "delta": {}, "finish_reason": item.get("done_reason", "stop")}],
+ "chronos": item}
+ yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n"
+ else:
+ chunk = {"id": cid, "object": "chat.completion.chunk", "created": created, "model": name,
+ "choices": [{"index": 0, "delta": {"content": item["text"]}, "finish_reason": None}]}
+ yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
+ return StreamingResponse(sse(), media_type="text/event-stream")
+ items = list(runtime.generate_stream(name, ids, **opts))
+ final = items[-1]
+ return {"id": cid, "object": "chat.completion", "created": created, "model": name,
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(i["text"] for i in items[:-1])},
+ "finish_reason": final.get("done_reason", "stop")}],
+ "usage": {"prompt_tokens": final.get("prompt_tokens", 0), "completion_tokens": final.get("completion_tokens", 0),
+ "total_tokens": final.get("prompt_tokens", 0) + final.get("completion_tokens", 0)},
+ "chronos": final}
+
+ return app
+
+
+def main(argv=None) -> int:
+ import argparse
+ parser = argparse.ArgumentParser(prog="chronos serve", description=__doc__.split("\n\n")[0])
+ parser.add_argument("--models-dir", required=True, help="directory searched for {_esc(title)}
"
+ lo, hi = min(flat), max(flat)
+ cols = max(len(r) for r in grid)
+ head = "".join(f"{_esc(t('monitor.none'))}
{_esc(title)}
"
+ f"
")
+
+
+def _status_html(snap: dict) -> str:
+ if not snap.get("run"):
+ return f"| {head} |
|---|
{_esc(fmt(lo))} {ramp} {_esc(fmt(hi))}
{_esc(t('monitor.no_runs'))}
{_esc(snap.get('root'))}✕ {_esc(t('monitor.error'))}: {_esc(st['error'])}
")
+ if state == "paused":
+ notes.append(f"‖ {_esc(t('monitor.paused'))}
")
+ if st.get("stale"):
+ notes.append(f"! {_esc(t('monitor.stale').format(age=_dur(st.get('last_update_age_s'))))}
")
+ if st.get("source") == "derived" and not snap.get("parameters"):
+ notes.append(f"i {_esc(t('monitor.derived'))}
")
+ frac = p.get("phase_fraction")
+ head = (f"{STATE_ICON.get(state, '○')} {_esc(t('monitor.state.' + state))}"
+ f"{_esc(snap['run'])}{_esc(p.get('current') or '—')} · {_esc(t('monitor.step'))} "
+ f"{_esc(_f(latest.get('step')))} / {_esc(_f(latest.get('total')))}"
+ f"{' · ' + format(frac * 100, '.1f') + '%' if isinstance(frac, (int, float)) else ''}"
+ f"{_esc(t('monitor.eta'))} {_esc((p.get('eta') or '—').replace('T', ' '))} · "
+ f"{_esc(t('monitor.remaining'))} {_esc(_dur(p.get('run_remaining_s')))} · {_esc(t('monitor.budget'))} "
+ f"{_esc(p.get('budget_hours'))} h
")
+ bars = []
+ for ph in p.get("phases") or []:
+ f = 1.0 if ph["state"] == "complete" else (frac or 0.0) if ph["state"] == "active" else 0.0
+ bars.append(f""
+ f"{_esc(ph['name'])} {_esc(t('monitor.ph.' + ph['state']))} · {_esc(ph.get('budget_h'))} h
")
+ best = (snap.get("best_valid") or {}).get(p.get("current")) or {}
+ mem = snap.get("memory") or {}
+ spec = snap.get("spec") or {}
+ cfg = spec.get("model_config") or {}
+ tiles = [
+ ("monitor.k_loss", _f(latest.get("loss")), f"{t('monitor.step')} {_f(latest.get('step'))}"),
+ ("monitor.k_ce", _f(latest.get("primary_loss")),
+ f"ppl {_f(math.exp(latest['primary_loss']), 3)}" if isinstance(latest.get("primary_loss"), (int, float)) else ""),
+ ("monitor.k_valid", _f(best.get("valid_ce")), f"{t('monitor.step')} {_f(best.get('step'))}" if best else ""),
+ ("monitor.k_tps", f"{_count(latest.get('tokens_per_s'))} tok/s", f"{_count(latest.get('processed_tokens_per_s'))} {t('monitor.s_processed')}"),
+ ("monitor.k_mem", f"{_f(mem.get('footprint_gb') or latest.get('peak_mem_gb'), 3)} GB",
+ f"MLX peak {_f(latest.get('peak_mem_gb'), 3)} · active {_f(latest.get('active_mem_gb'), 3)} GB"),
+ ("monitor.k_grad", _f(latest.get("grad_norm")), ""),
+ ("monitor.k_lr", _f(latest.get("lr"), 3), ""),
+ ("monitor.k_depth", _f(latest.get("rec_depth_mean"), 3), f"max {_esc(cfg.get('recurrent_depth', '—'))}"),
+ ("monitor.k_route", _f(latest.get("route_entropy"), 3), f"max {_f(latest.get('route_entropy_max'), 3)} · CV {_f(latest.get('load_cv'), 3)}"),
+ ]
+ tiles_html = "".join(f"{_esc(t(k))}{_esc(v)}{_esc(sub)}
" for k, v, sub in tiles)
+ foot = (f"{_esc(t('monitor.source'))}: {_esc(st.get('source'))} · pid {_esc(st.get('pid'))} " + f"({_esc(st.get('ps_state') or '-')}) · {_esc(time.strftime('%H:%M:%S'))}
") + return (f"{''.join(notes)}{head}
")
+
+
+def _events_rows(snap: dict) -> list[dict]:
+ out = []
+ for e in reversed(snap.get("events") or []):
+ detail = {k: v for k, v in e.items() if k not in ("time", "event", "phase", "step")}
+ out.append({"time": e.get("time"), "event": e.get("event"), "phase": e.get("phase"), "step": e.get("step"),
+ "valid_ce": e.get("valid_ce"), "detail": json.dumps(detail, ensure_ascii=False)[:400]})
+ return out
+
+
+def _samples_html(snap: dict) -> str:
+ """Samples of any task: the snapshot normalises them to input/reference/output + tags."""
+ rows = list(reversed(snap.get("samples") or []))[:12]
+ if not rows:
+ return f"{''.join(bars)}
"
+ f"{tiles_html}
{foot}{_esc(t('monitor.none'))}
"
+ items = []
+ for r in rows:
+ tags = " · ".join(f"{k}={v}" for k, v in (r.get("tags") or {}).items())
+ head = " · ".join(x for x in (_esc(r.get("phase")), f"{_esc(t('monitor.step'))} {_esc(r.get('step'))}", _esc(tags)) if x)
+ body = "".join(f"{_esc(t(key))} {_esc(r.get(role))}" + for role, key in (("input", "monitor.sample_input"), ("reference", "monitor.sample_ref"), + ("output", "monitor.sample_output")) if r.get(role) not in (None, "")) + items.append(f"
{head}{body}
")
+ return f"{''.join(items)}
"
+
+
+def _eval_family(snap: dict, metric: str | None) -> tuple[list[str], str | None, dict | None]:
+ """Metric families the run logs (discovered, not assumed) and the chosen one."""
+ families = (snap.get("evals") or {}).get("families") or []
+ names = [f["metric"] for f in families]
+ chosen = metric if metric in names else (names[0] if names else None)
+ return names, chosen, next((f for f in families if f["metric"] == chosen), None)
+
+
+# ---- refresh -------------------------------------------------------------------------
+def _phase_rows(snap: dict, phase: str | None) -> tuple[str | None, list[dict]]:
+ hist = snap.get("history") or {}
+ phase = phase if phase in hist else (snap.get("progress") or {}).get("current")
+ if phase not in hist:
+ phase = next(reversed(hist), None) if hist else None
+ return phase, hist.get(phase, []) if phase else []
+
+
+def _loss_parts(rows: list[dict], term: str):
+ term = term if term in LOSS_TERMS else "primary_loss"
+ return [(term, rows, term)]
+
+
+def refresh(runs_dir: str, run: str | None, phase: str | None, loss_term: str, eval_metric: str | None = None):
+ """Everything the tab shows, in output order (see ``build_monitor_tab``)."""
+ mon = _monitor(runs_dir)
+ try:
+ snap = mon.snapshot(run or None)
+ except KeyError:
+ snap = mon.snapshot(None)
+ names = [r["name"] for r in snap.get("runs") or []]
+ run_update = gr.update(choices=names, value=snap.get("run"))
+ phases = list((snap.get("history") or {}).keys())
+ phase, rows = _phase_rows(snap, phase)
+ phase_update = gr.update(choices=phases, value=phase)
+ valid = [e for e in snap.get("events") or [] if e.get("event") == "validation" and e.get("phase") == phase]
+ plots = []
+ for group in ("overview", "recurrent"):
+ for _title, series in PLOTS[group]:
+ parts = [(t(label), valid if key == "valid_ce" else rows, key) for key, label in series]
+ plots.append(_plot_update(parts))
+ present = [k for k in LOSS_TERMS if any(isinstance(r.get(k), (int, float)) for r in rows)] or ["primary_loss"]
+ term = loss_term if loss_term in present else present[0]
+ loss_update = _plot_update(_loss_parts(rows, term))
+ term_update = gr.update(choices=present, value=term)
+
+ params = snap.get("parameters") or {}
+ routing = params.get("routing") or {}
+ experts = params.get("experts") or {}
+ groups = [{"group": k, **v} for k, v in (params.get("groups") or {}).items()]
+ if params.get("stats"):
+ groups.append({"group": "total", **params["stats"], "tensors": params.get("tensor_count")})
+ group_df = _df(groups, [("group", "monitor.c_group"), ("tensors", "monitor.c_tensors"), ("numel", "monitor.c_numel"),
+ ("rms", "monitor.c_rms"), ("grad_norm", "monitor.c_grad_norm"), ("update_ratio", "monitor.c_update")],
+ numeric=("rms", "grad_norm", "update_ratio"))
+ full = mon.parameters(snap.get("run")) if params else None
+ tensors = (full or {}).get("tensors") or []
+ tensor_df = _df([{**x, "shape": "×".join(map(str, x.get("shape") or []))} for x in tensors],
+ [("name", "monitor.c_name"), ("group", "monitor.c_group"), ("shape", "monitor.c_shape"), ("dtype", "monitor.c_dtype"),
+ ("rms", "monitor.c_rms"), ("std", "monitor.c_std"), ("absmax", "monitor.c_absmax"),
+ ("grad_rms", "monitor.c_grad_rms"), ("update_ratio", "monitor.c_update")],
+ numeric=("rms", "std", "absmax", "grad_rms", "update_ratio"))
+ tel_note = (t("monitor.telemetry_at").format(step=params.get("step"), ms=round((params.get("telemetry_s") or 0) * 1000))
+ if params else t("monitor.no_telemetry"))
+ load_html = _heat_html(t("monitor.h_load"), routing.get("labels") or [], routing.get("load") or [],
+ fmt=lambda v: f"{v * 100:.1f}%")
+ imp_html = _heat_html(t("monitor.h_importance"), routing.get("labels") or [], routing.get("importance") or [])
+ layers = [f"L{x}" for x in experts.get("layers") or []]
+ wrms_html = _heat_html(t("monitor.h_wrms"), layers, experts.get("weight_rms") or [])
+ grms_html = _heat_html(t("monitor.h_grms"), layers, experts.get("grad_rms") or [])
+
+ events_df = _df(_events_rows(snap), [("time", "monitor.c_time"), ("event", "monitor.c_event"), ("phase", "monitor.c_phase"),
+ ("step", "monitor.step"), ("valid_ce", "monitor.c_valid_ce"), ("detail", "monitor.c_detail")])
+ ckpt_df = _df(snap.get("checkpoints") or [], [("file", "monitor.c_file"), ("size_gb", "monitor.c_size"),
+ ("modified", "monitor.c_time")])
+ names, chosen, family = _eval_family(snap, eval_metric)
+ eval_parts = [(sr["label"], valid, sr["key"]) for sr in (family or {}).get("series", [])[:len(PALETTE)]]
+ eval_update = gr.update(choices=names, value=chosen)
+ eval_plot = _plot_update(eval_parts)
+ training_params = mon.training_parameters(snap.get("run")) if snap.get("run") else {}
+ log = mon.log_tail(snap.get("run"), 24000) if snap.get("run") else ""
+ return (run_update, phase_update, _status_html(snap), *plots, term_update, loss_update,
+ tel_note, group_df, tensor_df, load_html, imp_html, wrms_html, grms_html,
+ eval_update, eval_plot, events_df, _samples_html(snap), training_params, log, ckpt_df)
+
+
+# ---- layout ------------------------------------------------------------------------------
+def _labelled(component, key: str):
+ register_translatable(component, key)
+ return component
+
+
+def build_monitor_tab():
+ with gr.Tab(t("tab.monitor")) as tab:
+ register_translatable(tab, "tab.monitor")
+ gr.Markdown(t("monitor.intro"))
+ with gr.Row():
+ runs_dir = _labelled(gr.Textbox(value=DEFAULT_RUNS_DIR, label=t("monitor.runs_dir"), scale=3), "monitor.runs_dir")
+ run = _labelled(gr.Dropdown(choices=[], value=None, label=t("monitor.run"), scale=2, allow_custom_value=True), "monitor.run")
+ phase = _labelled(gr.Dropdown(choices=[], value=None, label=t("monitor.phase"), scale=1, allow_custom_value=True), "monitor.phase")
+ auto = _labelled(gr.Checkbox(value=True, label=t("monitor.auto"), scale=1), "monitor.auto")
+ refresh_btn = _labelled(gr.Button(t("monitor.refresh"), scale=1), "monitor.refresh")
+ status = gr.HTML()
+ plots = []
+ with gr.Tab(t("monitor.t_overview")) as sub:
+ register_translatable(sub, "monitor.t_overview")
+ for i in range(0, len(PLOTS["overview"]), 2):
+ with gr.Row():
+ for title, _series in PLOTS["overview"][i:i + 2]:
+ plots.append(_labelled(gr.LinePlot(x="step", y="value", color="series", label=t(title), y_title="",
+ height=300, tooltip=["series", "step", "value"]), title))
+ with gr.Tab(t("monitor.t_recurrent")) as sub:
+ register_translatable(sub, "monitor.t_recurrent")
+ for i in range(0, len(PLOTS["recurrent"]), 3):
+ with gr.Row():
+ for title, _series in PLOTS["recurrent"][i:i + 3]:
+ plots.append(_labelled(gr.LinePlot(x="step", y="value", color="series", label=t(title), y_title="",
+ height=270, tooltip=["series", "step", "value"]), title))
+ with gr.Tab(t("monitor.t_loss")) as sub:
+ register_translatable(sub, "monitor.t_loss")
+ loss_term = _labelled(gr.Dropdown(choices=list(LOSS_TERMS), value="primary_loss", label=t("monitor.loss_term")), "monitor.loss_term")
+ loss_plot = gr.LinePlot(x="step", y="value", color="series", y_title="", height=340, tooltip=["series", "step", "value"])
+ with gr.Tab(t("monitor.t_params")) as sub:
+ register_translatable(sub, "monitor.t_params")
+ tel_note = gr.Markdown()
+ group_df = _labelled(gr.Dataframe(label=t("monitor.groups"), interactive=False, wrap=True), "monitor.groups")
+ with gr.Row():
+ load_html, imp_html = gr.HTML(), gr.HTML()
+ with gr.Row():
+ wrms_html, grms_html = gr.HTML(), gr.HTML()
+ tensor_df = _labelled(gr.Dataframe(label=t("monitor.tensors"), interactive=False, max_height=520,
+ column_widths=["30%", "10%", "8%", "8%", "8.8%", "8.8%", "8.8%", "8.8%", "8.8%"]),
+ "monitor.tensors")
+ with gr.Tab(t("monitor.t_evals")) as sub:
+ register_translatable(sub, "monitor.t_evals")
+ eval_note = gr.Markdown(t("monitor.eval_note"))
+ register_translatable(eval_note, "monitor.eval_note")
+ eval_metric = _labelled(gr.Dropdown(choices=[], value=None, label=t("monitor.eval_metric"),
+ allow_custom_value=True), "monitor.eval_metric")
+ eval_plot = gr.LinePlot(x="step", y="value", color="series", y_title="", height=300,
+ tooltip=["series", "step", "value"])
+ events_df = _labelled(gr.Dataframe(label=t("monitor.events"), interactive=False, wrap=True, max_height=420), "monitor.events")
+ samples_html = gr.HTML()
+ with gr.Tab(t("monitor.t_config")) as sub:
+ register_translatable(sub, "monitor.t_config")
+ params_json = _labelled(gr.JSON(label=t("monitor.training_params")), "monitor.training_params")
+ with gr.Tab(t("monitor.t_log")) as sub:
+ register_translatable(sub, "monitor.t_log")
+ log_box = _labelled(gr.Textbox(label=t("monitor.log"), lines=24, max_lines=24, interactive=False, autoscroll=True), "monitor.log")
+ ckpt_df = _labelled(gr.Dataframe(label=t("monitor.checkpoints"), interactive=False), "monitor.checkpoints")
+
+ outputs = [run, phase, status, *plots, loss_term, loss_plot, tel_note, group_df, tensor_df,
+ load_html, imp_html, wrms_html, grms_html, eval_metric, eval_plot, events_df, samples_html,
+ params_json, log_box, ckpt_df]
+ inputs = [runs_dir, run, phase, loss_term, eval_metric]
+ refresh_btn.click(fn=refresh, inputs=inputs, outputs=outputs)
+ run.input(fn=refresh, inputs=inputs, outputs=outputs)
+ phase.input(fn=refresh, inputs=inputs, outputs=outputs)
+ loss_term.input(fn=refresh, inputs=inputs, outputs=outputs)
+ eval_metric.input(fn=refresh, inputs=inputs, outputs=outputs)
+ runs_dir.submit(fn=refresh, inputs=inputs, outputs=outputs)
+
+ tab.select(fn=refresh, inputs=inputs, outputs=outputs) # first paint when the tab opens
+ timer = gr.Timer(15.0)
+ timer.tick(fn=refresh, inputs=inputs, outputs=outputs)
+ auto.change(fn=lambda on: gr.Timer(active=bool(on)), inputs=[auto], outputs=[timer])
+ tab.monitor_wiring = (refresh, inputs, outputs) # app.load() paints it on page open
+ return tab
diff --git a/ui/tabs/serve_tab.py b/ui/tabs/serve_tab.py
new file mode 100644
index 0000000..43e3f1e
--- /dev/null
+++ b/ui/tabs/serve_tab.py
@@ -0,0 +1,202 @@
+"""ui/tabs/serve_tab.py — start/stop ``chronos serve`` from the console.
+
+The server runs as a child process (argument list, no shell) so it keeps its
+own memory accounting and survives UI reloads only while this process lives.
+Security contract (same as the CLI): loopback by default; any other host
+requires *allow remote* plus an API key. The key is passed through the
+``CHRONOS_API_KEY`` environment variable, never on the command line, so it
+does not show up in ``ps``.
+"""
+from __future__ import annotations
+
+import atexit
+import os
+import subprocess
+import sys
+import threading
+from collections import deque
+
+from ui.gradio_compat import gr
+from ui.i18n import register_translatable, t
+
+LOOPBACK = ("127.0.0.1", "localhost", "::1")
+
+
+def build_serve_command(*, models_dir: str, host: str = "127.0.0.1", port: int = 11435,
+ device: str = "auto", dtype: str = "auto", lazy_experts: bool = True,
+ ram_expert_fraction: float = 0.5, allow_remote: bool = False,
+ api_key: str = "", kv_cache_dtype: str = "", mla_decode: str = "",
+ router_scope: str = "", max_resident_experts: int | None = None,
+ ram_budget_mb: float | None = None, monitor_runs_dir: str = "") -> tuple[list[str], dict]:
+ """Return ``(argv, env_overrides)``; raises ValueError on an unsafe or invalid request."""
+ host = (host or "127.0.0.1").strip()
+ models_dir = os.path.expanduser((models_dir or "").strip())
+ if not models_dir or not os.path.isdir(models_dir):
+ raise ValueError("models_dir")
+ port = int(port)
+ if not 1 <= port <= 65535:
+ raise ValueError("port")
+ remote = host not in LOOPBACK
+ if remote and not (allow_remote and (api_key or "").strip()):
+ raise ValueError("remote_needs_key")
+ argv = [sys.executable, "-m", "chronos.serving.ollama_server", "--models-dir", models_dir,
+ "--host", host, "--port", str(port), "--device", device or "auto", "--dtype", dtype or "auto",
+ "--ram-expert-fraction", str(float(ram_expert_fraction))]
+ if not lazy_experts:
+ argv.append("--no-lazy-experts")
+ if allow_remote:
+ argv.append("--allow-remote")
+ for flag, value in (("--kv-cache-dtype", kv_cache_dtype), ("--mla-decode", mla_decode),
+ ("--router-scope", router_scope)):
+ if value and value != "checkpoint":
+ argv += [flag, value]
+ if max_resident_experts:
+ argv += ["--max-resident-experts", str(int(max_resident_experts))]
+ if ram_budget_mb:
+ argv += ["--ram-budget-mb", str(float(ram_budget_mb))]
+ if monitor_runs_dir and monitor_runs_dir.strip():
+ argv += ["--monitor-runs-dir", os.path.expanduser(monitor_runs_dir.strip())]
+ env = {"CHRONOS_API_KEY": api_key.strip()} if (api_key or "").strip() else {}
+ return argv, env
+
+
+def redact(argv: list[str]) -> str:
+ return " ".join(argv[1:]) # the key is never in argv; drop the interpreter path
+
+
+class ServeProcess:
+ def __init__(self):
+ self.proc: subprocess.Popen | None = None
+ self.lines: deque[str] = deque(maxlen=500)
+ self.lock = threading.Lock()
+ self.url = ""
+ self.auth = False
+
+ def running(self) -> bool:
+ return self.proc is not None and self.proc.poll() is None
+
+ def _pump(self, proc: subprocess.Popen) -> None:
+ for line in proc.stdout:
+ with self.lock:
+ self.lines.append(line.rstrip())
+ with self.lock:
+ self.lines.append(f"[exit {proc.wait()}]")
+
+ def start(self, argv: list[str], env_overrides: dict, host: str, port: int) -> str:
+ if self.running():
+ return "busy"
+ env = dict(os.environ)
+ env.pop("CHRONOS_API_KEY", None) # only an explicit key from the form enables auth
+ env.update(env_overrides)
+ env["PYTHONUNBUFFERED"] = "1"
+ repo = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+ self.lines.clear()
+ self.lines.append(f"$ {redact(argv)}")
+ self.proc = subprocess.Popen(argv, cwd=repo, env=env, stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT, text=True, bufsize=1)
+ threading.Thread(target=self._pump, args=(self.proc,), daemon=True).start()
+ shown = f"[{host}]" if ":" in host else host
+ self.url, self.auth = f"http://{shown}:{int(port)}", bool(env_overrides.get("CHRONOS_API_KEY"))
+ return "started"
+
+ def stop(self, timeout: float = 10.0) -> str:
+ if not self.running():
+ return "idle"
+ self.proc.terminate()
+ try:
+ self.proc.wait(timeout=timeout)
+ except subprocess.TimeoutExpired:
+ self.proc.kill()
+ self.proc.wait()
+ return "stopped"
+
+ def snapshot(self) -> tuple[bool, str]:
+ with self.lock:
+ return self.running(), "\n".join(self.lines)
+
+
+_server = ServeProcess()
+atexit.register(_server.stop) # never leave an orphaned server behind the UI
+
+
+def _status(lang: str | None = None) -> str:
+ if not _server.running():
+ return t("serve.state.stopped", lang)
+ auth = t("serve.auth_on", lang) if _server.auth else t("serve.auth_off", lang)
+ return f"{t('serve.state.running', lang)} · {_server.url} · {auth}"
+
+
+def build_serve_tab():
+ with gr.Tab(t("tab.serve")) as tab:
+ register_translatable(tab, "tab.serve")
+ intro = gr.Markdown(t("serve.intro"))
+ register_translatable(intro, "serve.intro")
+ with gr.Row():
+ models_dir = gr.Textbox(label=t("serve.models_dir"), value="./out", scale=3)
+ host = gr.Textbox(label=t("serve.host"), value="127.0.0.1", scale=1)
+ port = gr.Number(label=t("serve.port"), value=11435, precision=0, scale=1)
+ with gr.Row():
+ device = gr.Dropdown(["auto", "cpu", "mps", "cuda"], value="auto", label=t("serve.device"))
+ dtype = gr.Dropdown(["auto", "float32", "bfloat16", "float16"], value="auto", label=t("serve.dtype"))
+ lazy = gr.Checkbox(value=True, label=t("serve.lazy"))
+ ram_fraction = gr.Slider(0.05, 1.0, value=0.5, step=0.05, label=t("serve.ram_fraction"))
+ with gr.Row():
+ kv_dtype = gr.Dropdown(["checkpoint", "auto", "int8"], value="checkpoint", label=t("serve.kv_dtype"))
+ mla_decode = gr.Dropdown(["checkpoint", "auto", "absorbed", "expanded"], value="checkpoint",
+ label=t("serve.mla_decode"))
+ router_scope = gr.Dropdown(["checkpoint", "per_iteration", "shared"], value="checkpoint",
+ label=t("serve.router_scope"))
+ max_resident = gr.Number(label=t("serve.max_resident"), value=0, precision=0)
+ ram_budget = gr.Number(label=t("serve.ram_budget"), value=0, precision=0)
+ with gr.Row():
+ allow_remote = gr.Checkbox(value=False, label=t("serve.allow_remote"))
+ api_key = gr.Textbox(label=t("serve.api_key"), type="password", value="")
+ monitor_dir = gr.Textbox(label=t("serve.monitor_dir"), value="", placeholder="./runs")
+ warning = gr.Markdown(t("serve.warn_noauth"))
+ with gr.Row():
+ start_btn = gr.Button(t("serve.start"), variant="primary")
+ stop_btn = gr.Button(t("serve.stop"), variant="stop")
+ status = gr.Textbox(label=t("serve.status"), interactive=False, value=_status())
+ log_box = gr.Textbox(label=t("serve.log"), lines=14, interactive=False, autoscroll=True)
+ usage = gr.Markdown(t("serve.usage"))
+ for comp, key in ((models_dir, "serve.models_dir"), (host, "serve.host"), (port, "serve.port"),
+ (device, "serve.device"), (dtype, "serve.dtype"), (lazy, "serve.lazy"),
+ (ram_fraction, "serve.ram_fraction"), (kv_dtype, "serve.kv_dtype"),
+ (mla_decode, "serve.mla_decode"), (router_scope, "serve.router_scope"),
+ (max_resident, "serve.max_resident"), (ram_budget, "serve.ram_budget"),
+ (allow_remote, "serve.allow_remote"), (api_key, "serve.api_key"),
+ (monitor_dir, "serve.monitor_dir"), (warning, "serve.warn_noauth"),
+ (start_btn, "serve.start"), (stop_btn, "serve.stop"), (status, "serve.status"),
+ (log_box, "serve.log"), (usage, "serve.usage")):
+ register_translatable(comp, key)
+
+ def start(md, h, p, dev, dt, lz, rf, kv, mla, rs, mr, rb, ar, key, mon):
+ try:
+ argv, env = build_serve_command(
+ models_dir=md, host=h, port=int(p or 11435), device=dev, dtype=dt, lazy_experts=bool(lz),
+ ram_expert_fraction=float(rf), allow_remote=bool(ar), api_key=key or "",
+ kv_cache_dtype=kv, mla_decode=mla, router_scope=rs,
+ max_resident_experts=int(mr or 0) or None, ram_budget_mb=float(rb or 0) or None,
+ monitor_runs_dir=mon or "")
+ except ValueError as exc:
+ return t(f"serve.err.{exc}"), gr.update()
+ outcome = _server.start(argv, env, (h or "127.0.0.1").strip(), int(p or 11435))
+ if outcome == "busy":
+ return t("serve.err.busy"), gr.update()
+ return _status(), _server.snapshot()[1]
+
+ def stop():
+ _server.stop()
+ return _status(), _server.snapshot()[1]
+
+ def poll():
+ return _status(), _server.snapshot()[1]
+
+ start_btn.click(fn=start, inputs=[models_dir, host, port, device, dtype, lazy, ram_fraction, kv_dtype,
+ mla_decode, router_scope, max_resident, ram_budget, allow_remote,
+ api_key, monitor_dir],
+ outputs=[status, log_box])
+ stop_btn.click(fn=stop, outputs=[status, log_box])
+ timer = gr.Timer(2.0)
+ timer.tick(fn=poll, outputs=[status, log_box])
+ return tab
diff --git a/ui/tabs/tokenizer_tab.py b/ui/tabs/tokenizer_tab.py
new file mode 100644
index 0000000..6812ce6
--- /dev/null
+++ b/ui/tabs/tokenizer_tab.py
@@ -0,0 +1,191 @@
+"""ui/tabs/tokenizer_tab.py — Chronos Tokenizer Builder.
+
+Trains a tokenizer from user-selected datasets (jsonl / parquet / csv / tsv /
+txt, files, directories or globs) with ``chronos.tokenizer.build_tokenizer``
+in a background thread, then shows the held-out evaluation it wrote to
+``tokenizer_report.json``. The same builder backs ``chronos tokenizer``.
+"""
+from __future__ import annotations
+
+import json
+import os
+import threading
+import time
+from collections import deque
+
+from ui.gradio_compat import gr
+from ui.i18n import register_translatable, t
+
+EVAL_COLUMNS = ["bucket", "texts", "bytes/token", "chars/token", "tokens/word",
+ "byte_fallback_rate", "roundtrip_failures"]
+
+
+def parse_paths(text: str) -> list[str]:
+ """One path per line or comma; ``~`` expanded; blanks dropped."""
+ out = []
+ for chunk in (text or "").replace(",", "\n").splitlines():
+ chunk = chunk.strip().strip('"').strip("'")
+ if chunk:
+ out.append(os.path.expanduser(chunk))
+ return out
+
+
+def report_rows(report: dict) -> list[list]:
+ """Per-bucket rows from a ``heldout_eval`` dict (``vocab_utilization`` skipped)."""
+ rows = []
+ for bucket, stats in (report or {}).items():
+ if not isinstance(stats, dict):
+ continue
+ rows.append([bucket, stats.get("texts"), stats.get("bytes_per_token"), stats.get("chars_per_token"),
+ stats.get("tokens_per_word"), stats.get("byte_fallback_rate"),
+ stats.get("roundtrip_failures")])
+ return rows
+
+
+def read_report(out_dir: str) -> dict:
+ path = os.path.join(os.path.expanduser(out_dir or ""), "tokenizer_report.json")
+ if not os.path.isfile(path):
+ return {}
+ with open(path, encoding="utf-8") as fh:
+ return json.load(fh)
+
+
+def report_summary(report: dict) -> dict:
+ if not report:
+ return {}
+ keep = ("vocab_size", "fingerprint", "algorithm", "elapsed_s", "special_tokens")
+ out = {k: report[k] for k in keep if k in report}
+ heldout = report.get("heldout_eval") or {}
+ if "vocab_utilization" in heldout:
+ out["vocab_utilization"] = heldout["vocab_utilization"]
+ if report.get("sources"):
+ out["sources"] = report["sources"]
+ return out
+
+
+class TokenizerJob:
+ """One background build at a time; log lines are kept in a bounded deque."""
+
+ def __init__(self):
+ self.lock = threading.Lock()
+ self.thread: threading.Thread | None = None
+ self.lines: deque[str] = deque(maxlen=400)
+ self.state = "idle"
+ self.out_dir = ""
+ self.error = ""
+
+ def log(self, message) -> None:
+ with self.lock:
+ self.lines.append(f"{time.strftime('%H:%M:%S')} {message}")
+
+ def running(self) -> bool:
+ return self.thread is not None and self.thread.is_alive()
+
+ def start(self, paths: list[str], out_dir: str, vocab_size: int, per_source: int,
+ fields: list[str] | None, temperature: float) -> str:
+ if self.running():
+ return "busy"
+ if not paths:
+ return "no_paths"
+ if not out_dir:
+ return "no_out"
+ from chronos.tokenizer.extract import expand
+ files = expand(paths)
+ if not files:
+ return "no_files"
+ self.lines.clear()
+ self.state, self.out_dir, self.error = "running", os.path.expanduser(out_dir), ""
+ self.log(f"{len(files)} files: " + ", ".join(os.path.basename(f) for f in files[:8])
+ + (" …" if len(files) > 8 else ""))
+
+ def work():
+ from chronos.tokenizer import BPETrainerConfig, ExtractConfig, build_tokenizer
+ try:
+ build_tokenizer(
+ paths, self.out_dir, vocab_size=int(vocab_size),
+ extract=ExtractConfig(per_source_texts=int(per_source), fields=fields or None),
+ trainer=BPETrainerConfig(vocab_size=int(vocab_size),
+ language_temperature=float(temperature)),
+ log=self.log,
+ )
+ self.state = "done"
+ except Exception as exc: # surfaced in the UI, never raised into Gradio
+ self.error = f"{type(exc).__name__}: {exc}"
+ self.log(f"error: {self.error}")
+ self.state = "failed"
+
+ self.thread = threading.Thread(target=work, name="chronos-tokenizer", daemon=True)
+ self.thread.start()
+ return "started"
+
+ def snapshot(self) -> tuple[str, str]:
+ with self.lock:
+ return self.state, "\n".join(self.lines)
+
+
+_job = TokenizerJob()
+
+
+def _status_text(state: str, lang: str | None = None) -> str:
+ return t(f"tok.state.{state}", lang)
+
+
+def build_tokenizer_tab():
+ with gr.Tab(t("tab.tokenizer")) as tab:
+ register_translatable(tab, "tab.tokenizer")
+ intro = gr.Markdown(t("tok.intro"))
+ register_translatable(intro, "tok.intro")
+
+ with gr.Row():
+ paths = gr.Textbox(label=t("tok.paths"), lines=4, scale=3,
+ placeholder="./dataset/corpus/\n./dataset/*.jsonl\n./dataset/*.parquet")
+ with gr.Column(scale=2):
+ out_dir = gr.Textbox(label=t("tok.out_dir"), value="./out/tokenizer")
+ fields = gr.Textbox(label=t("tok.fields"), placeholder="text, messages, prompt, response")
+ with gr.Row():
+ vocab = gr.Number(label=t("tok.vocab_size"), value=16384, precision=0)
+ per_source = gr.Number(label=t("tok.per_source"), value=400_000, precision=0)
+ temperature = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label=t("tok.temperature"))
+ with gr.Row():
+ build_btn = gr.Button(t("tok.build"), variant="primary")
+ load_btn = gr.Button(t("tok.load_report"))
+ status = gr.Textbox(label=t("tok.status"), interactive=False)
+ log_box = gr.Textbox(label=t("tok.log"), lines=12, interactive=False, autoscroll=True)
+ summary = gr.JSON(label=t("tok.summary"), value={})
+ table = gr.Dataframe(headers=EVAL_COLUMNS, label=t("tok.eval"), interactive=False)
+ for comp, key in ((paths, "tok.paths"), (out_dir, "tok.out_dir"), (fields, "tok.fields"),
+ (vocab, "tok.vocab_size"), (per_source, "tok.per_source"),
+ (temperature, "tok.temperature"), (build_btn, "tok.build"),
+ (load_btn, "tok.load_report"), (status, "tok.status"), (log_box, "tok.log"),
+ (summary, "tok.summary"), (table, "tok.eval")):
+ register_translatable(comp, key)
+
+ def start(p, o, v, n, f, temp):
+ field_list = [x.strip() for x in (f or "").split(",") if x.strip()]
+ outcome = _job.start(parse_paths(p), (o or "").strip(), int(v or 16384), int(n or 400_000),
+ field_list, float(temp or 0.5))
+ if outcome != "started":
+ return t(f"tok.err.{outcome}")
+ return _status_text("running")
+
+ def show(o):
+ report = read_report(o)
+ if not report:
+ return t("tok.err.no_report"), {}, []
+ return t("tok.state.loaded"), report_summary(report), report_rows(report.get("heldout_eval"))
+
+ def poll(o):
+ state, text = _job.snapshot()
+ if state in ("done",) and _job.out_dir:
+ report = read_report(_job.out_dir)
+ return (_status_text(state), text, report_summary(report),
+ report_rows(report.get("heldout_eval")))
+ message = _status_text(state) + (f" — {_job.error}" if _job.error else "")
+ return message, text, gr.update(), gr.update()
+
+ build_btn.click(fn=start, inputs=[paths, out_dir, vocab, per_source, fields, temperature],
+ outputs=[status])
+ load_btn.click(fn=show, inputs=[out_dir], outputs=[status, summary, table])
+ timer = gr.Timer(2.0)
+ timer.tick(fn=poll, inputs=[out_dir], outputs=[status, log_box, summary, table])
+ return tab
diff --git a/ui/tabs/train_tab.py b/ui/tabs/train_tab.py
index 07fc397..ee3d46b 100644
--- a/ui/tabs/train_tab.py
+++ b/ui/tabs/train_tab.py
@@ -80,7 +80,14 @@
"kv_latent_dim",
"rope_dim",
"num_attention_heads",
- "num_key_value_heads",
+ "num_key_value_heads", "head_dim", "use_moe", "use_hybrid_attention", "tie_word_embeddings",
+ "architecture_version", "recurrent_enabled", "recurrent_groups", "recurrent_depth",
+ "prelude_layers", "coda_layers", "layer_aware_lookahead", "layer_aware_cache",
+ "modalities", "image_patch_size", "modality_fusion_heads",
+ # v3 semantics that do not change tensor shapes but do change what the
+ # weights compute; a stage must inherit them from its init checkpoint.
+ "recurrent_attention_mode", "rope_layout", "rope_theta", "sliding_window_size",
+ "intermediate_size", "router_scope",
)
STAGE_HELP_TEXT = {
"pretrain": {
@@ -444,6 +451,8 @@ def _effective_training_config(
"teacher_path",
"data_path",
"data_path_source",
+ "lambda_router_anchor",
+ "stage_overrides_applied",
]
out = {
"mode": mode,
@@ -538,6 +547,13 @@ def _planned_optimizer_updates(train_steps: int, accumulation_steps: int) -> int
return (train_steps + accum - 1) // accum if train_steps else 0
+_STAGE_OVERRIDE_KEYS = frozenset((
+ "learning_rate", "batch_size", "accumulation_steps", "epochs", "max_steps", "max_seq_len",
+ "weight_decay", "grad_clip", "save_interval", "lambda_router_anchor", "beta", "alpha",
+ "lambda_or", "temperature", "num_generations", "max_gen_len",
+))
+
+
def _prepare_training_run_config(
cfg: dict | None,
train_mode: str,
@@ -552,6 +568,11 @@ def _prepare_training_run_config(
) -> tuple[dict, dict]:
merged = dict(MINIMIND_MOE_DEFAULTS)
merged.update(dict(cfg) if cfg else {})
+ # Per-stage hyperparameters (Config tab → auto plan) override the shared
+ # values for this stage only; architecture keys are never stage-specific.
+ stage_values = (merged.get("stage_overrides") or {}).get(train_mode) or {}
+ merged.update({k: v for k, v in stage_values.items() if k in _STAGE_OVERRIDE_KEYS})
+ merged["stage_overrides_applied"] = sorted(k for k in stage_values if k in _STAGE_OVERRIDE_KEYS)
cfg = merged
requested_backend = (selected_backend or "auto").strip().lower() or "auto"
resolved_backend, resolved_device = resolve_training_device(requested_backend)
@@ -891,6 +912,8 @@ def __init__(self):
self.last_sample: str = ""
self.last_sample_step: int = 0
self.effective_config: dict = {}
+ self.tokenizer_path: str = ""
+ self.vocab_size: int = 0
def _normalise_effective_cpu_config(self, cfg: dict, resolved_backend: str) -> None:
cfg["cpu_budget_percent"] = _cpu_budget_for_backend(cfg, resolved_backend)
@@ -908,6 +931,8 @@ def start(self, cfg: dict, mode: str):
self._metrics.clear()
self.total_steps = 0
self.t_start = time.monotonic()
+ self.tokenizer_path = str(cfg.get("tokenizer_path") or "").strip()
+ self.vocab_size = int(cfg.get("vocab_size") or 0)
self._thread = threading.Thread(
target=self._run, args=(cfg, mode), daemon=True
)
@@ -972,7 +997,7 @@ def _resolve_default_init_path(self, save_dir: str, mode: str, hidden_size: int)
def _topology_mismatches(self, sniffed: dict, model_cfg_kwargs: dict) -> list[str]:
mismatches = []
for k in CHECKPOINT_TOPOLOGY_KEYS:
- if k in sniffed and k in model_cfg_kwargs and int(sniffed[k]) != int(model_cfg_kwargs[k]):
+ if k in sniffed and k in model_cfg_kwargs and sniffed[k] != model_cfg_kwargs[k]:
mismatches.append(f"{k}: ckpt={sniffed[k]} != ui={model_cfg_kwargs[k]}")
return mismatches
@@ -983,11 +1008,9 @@ def _adopt_checkpoint_topology(self, model_cfg_kwargs: dict, sniffed: dict) -> d
continue
current = model_cfg_kwargs.get(key)
value = sniffed[key]
- if current is None or int(current) != int(value):
+ if current is None or current != value:
adopted[key] = value
model_cfg_kwargs[key] = value
- if "moe_intermediate_size" in sniffed:
- model_cfg_kwargs["intermediate_size"] = sniffed["moe_intermediate_size"]
return adopted
def _build_stage_args(self, cfg: dict, mode: str, save_dir: str, hidden_size: int):
@@ -1122,8 +1145,8 @@ def _generate_live_sample(self, model, device: str, cfg: dict, mode: str):
with torch.no_grad():
gen = ids.clone()
for _ in range(80):
- with autocast_context(device, cfg.get("dtype", "auto")):
- o, _ = model(gen, use_cache=False)
+ with autocast_context(device, cfg.get("dtype", "auto"), model=model):
+ o = model(gen, use_cache=False)
nxt = o.logits[:, -1, :].argmax(dim=-1, keepdim=True)
gen = torch.cat([gen, nxt], dim=-1)
if int(nxt.item()) == (tokenizer.eos_token_id or -1):
@@ -1149,58 +1172,21 @@ def _run(self, cfg: dict, mode: str):
from chronos.model.moe_chronos import ChronosMOEFeedForward
self._put(f"[{mode.upper()}] Building model...")
- # Pass the FULL config dict so user-set arch overrides (e.g.
- # moe_intermediate_size, num_attention_heads, vocab_size) are
- # actually honored. Previously this method only forwarded a
- # hardcoded subset, so the trainer silently used minimind's
- # auto-derived intermediate_size while the UI estimator showed
- # the user's slider — causing 10x param-count mismatches.
- model_cfg_kwargs = {
- "hidden_size": cfg.get("hidden_size", 512),
- "num_hidden_layers": cfg.get("num_hidden_layers", 8),
- "num_experts": cfg.get("num_experts", 4),
- "num_experts_per_tok": cfg.get("num_experts_per_tok", 1),
- "num_shared_experts": cfg.get("num_shared_experts", 1),
- "lookahead_steps": cfg.get("lookahead_steps", 2),
- "kv_latent_dim": cfg.get("kv_latent_dim", 64),
- "sliding_window_size": cfg.get("sliding_window_size", 2048),
- "lambda_balance": cfg.get("lambda_balance", 5e-4),
- "lambda_temporal": cfg.get("lambda_temporal", 3e-3),
- "lambda_lookahead": cfg.get("lambda_lookahead", 0.1),
- "lambda_lookahead_topk": cfg.get("lambda_lookahead_topk", 0.15),
- "lambda_lookahead_union": cfg.get("lambda_lookahead_union", 0.05),
- "lambda_router_locality": cfg.get("lambda_router_locality", 0.02),
- "fallback_mask_prob": cfg.get("fallback_mask_prob", 0.08),
- "vram_budget_gb": cfg.get("vram_budget_gb", 4.0),
- "use_hybrid_attention": cfg.get("use_hybrid_attention", True),
- "use_moe": True,
- }
- # Optional overrides: only forward when the user actually set them.
- # For the "auto" sentinel fields (0 = let MiniMindConfig derive
- # it), we skip forwarding when value is 0 so ceil(H·π/64)·64
- # takes effect instead of Linear(H, 0) crashing MPS init.
- AUTO_SENTINEL_KEYS = {"intermediate_size", "moe_intermediate_size"}
- for opt_key, cfg_key in [
- ("num_attention_heads", "num_attention_heads"),
- ("num_key_value_heads", "num_key_value_heads"),
- ("rope_dim", "rope_dim"),
- ("vocab_size", "vocab_size"),
- ("max_position_embeddings", "max_seq_len"),
- ("intermediate_size", "moe_intermediate_size"),
- ("moe_intermediate_size", "moe_intermediate_size"),
- ("lambda_router_anchor", "lambda_router_anchor"),
- ("pinned_memory_max_fraction", "pinned_memory_max_fraction"),
- ("storage_format", "storage_format"),
- ("tie_word_embeddings", "tie_word_embeddings"),
- ]:
- if cfg_key not in cfg:
- continue
- val = cfg[cfg_key]
- if val in (None, ""):
- continue
- if opt_key in AUTO_SENTINEL_KEYS and val == 0:
- continue
- model_cfg_kwargs[opt_key] = val
+ from chronos.config_schema import model_config_kwargs
+ model_cfg_kwargs = model_config_kwargs(cfg)
+ model_cfg_kwargs.setdefault("hidden_size", 512)
+ model_cfg_kwargs.setdefault("num_hidden_layers", 8)
+ detected_modalities_pre = []
+ if data_path_for_cleanup and os.path.exists(data_path_for_cleanup):
+ try:
+ from chronos.data.flexible_dataset import dataset_modality_summary
+ detected_modalities_pre = dataset_modality_summary(data_path_for_cleanup).get(
+ "detected_modalities", []
+ )
+ except Exception as exc:
+ self._put(f"[data] modality detection warning: {exc}")
+ if cfg.get("auto_detect_modalities", True):
+ model_cfg_kwargs["modalities"] = detected_modalities_pre
save_dir = cfg.get("save_dir", "./out")
init_weight = (cfg.get("init_weight") or "").strip()
load_path = ""
@@ -1276,6 +1262,13 @@ def _run(self, cfg: dict, mode: str):
requested_backend = (cfg.get("train_backend") or "auto").strip().lower() or "auto"
resolved_backend, device = resolve_training_device(requested_backend)
+ if resolved_backend == "mlx" and model_cfg.modalities:
+ if requested_backend != "auto":
+ raise ValueError("MLX 尚未实现多模态编码器训练，请选择 auto 或 PyTorch MPS/CUDA 后端。")
+ supported = set(training_available())
+ target = next(name for name in ("cuda", "xpu", "mps", "cpu") if name in supported)
+ resolved_backend, device = resolve_training_device(target)
+ self._put(f"[data] Modalities {model_cfg.modalities}: using {resolved_backend} for encoder training.")
cfg["train_backend"] = requested_backend
cfg["device"] = device
cfg["dtype"] = _normalize_dtype_for_trainer(cfg.get("dtype"))
@@ -1324,12 +1317,34 @@ def _run(self, cfg: dict, mode: str):
else:
tokenizer = self._load_tokenizer()
from chronos.data.flexible_dataset import (
- FlexibleDataset, StreamingSFTDataset, StreamingDPODataset,
+ FlexibleDataset, StreamingSFTDataset, MultimodalStreamingDataset,
+ MultimodalStreamingSFTDataset, StreamingDPODataset,
+ dataset_modality_summary, multimodal_collate,
)
+ modality_summary = dataset_modality_summary(data_path)
+ detected_modalities = modality_summary.get("detected_modalities", [])
if mode == "pretrain":
- full_ds = FlexibleDataset(data_path, tokenizer, max_length=max_seq_len)
+ if cfg.get("auto_detect_modalities", True) and detected_modalities:
+ model_cfg_kwargs["modalities"] = detected_modalities
+ full_ds = (
+ MultimodalStreamingDataset(data_path, tokenizer, max_length=max_seq_len,
+ image_size=int(cfg.get("image_size", 224)),
+ audio_sample_rate=int(cfg.get("audio_sample_rate", 16000)),
+ max_audio_seconds=float(cfg.get("max_audio_seconds", 30.0)),
+ media_error_policy=cfg.get("media_error_policy", "warn"))
+ if detected_modalities else FlexibleDataset(data_path, tokenizer, max_length=max_seq_len)
+ )
elif mode in {"sft", "distill"}:
- full_ds = StreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len)
+ if cfg.get("auto_detect_modalities", True) and detected_modalities:
+ model_cfg_kwargs["modalities"] = detected_modalities
+ full_ds = (
+ MultimodalStreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len,
+ image_size=int(cfg.get("image_size", 224)),
+ audio_sample_rate=int(cfg.get("audio_sample_rate", 16000)),
+ max_audio_seconds=float(cfg.get("max_audio_seconds", 30.0)),
+ media_error_policy=cfg.get("media_error_policy", "warn"))
+ if detected_modalities else StreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len)
+ )
elif mode in {"dpo", "orpo"}:
full_ds = StreamingDPODataset(data_path, tokenizer, max_length=max_seq_len)
else:
@@ -1352,6 +1367,7 @@ def _run(self, cfg: dict, mode: str):
DataLoader(
train_ds,
batch_size=batch_size,
+ collate_fn=multimodal_collate if detected_modalities and mode in {"pretrain", "sft", "distill"} else None,
**dataloader_kwargs(device=device, num_workers=cfg.get("num_workers", "auto"), shuffle=True),
)
if train_ds is not None else None
@@ -1360,6 +1376,7 @@ def _run(self, cfg: dict, mode: str):
DataLoader(
val_ds,
batch_size=batch_size,
+ collate_fn=multimodal_collate if detected_modalities and mode in {"pretrain", "sft", "distill"} else None,
**dataloader_kwargs(device=device, num_workers=cfg.get("num_workers", "auto"), shuffle=False),
)
if val_ds is not None else None
@@ -1413,10 +1430,16 @@ def _run(self, cfg: dict, mode: str):
if mlx_checkpoint_path:
self._put(f"[MLX {mode.upper()}] Loading checkpoint → {mlx_checkpoint_path}")
else:
- self._put(f"[MLX {mode.upper()}] Starting from current in-memory random init.")
+ self._put(f"[MLX {mode.upper()}] Starting from random init.")
def _mlx_progress(event: dict):
kind = str(event.get("event", "step"))
+ if kind == "validation":
+ value = float(event["ce"])
+ self._put(f"[MLX {mode.upper()} validation] step={event['step']} CE={value:.4f} tokens={event['tokens']}")
+ self._put_metric({"step": event["step"], "total": value, "primary": value,
+ "ce": value, "aux": 0., "temporal": 0., "tps": 0., "val": True, "val_loss": value})
+ return
if kind == "start":
total = int(event.get("steps", 0) or 0)
if total > 0:
@@ -1426,6 +1449,7 @@ def _mlx_progress(event: dict):
f"total_steps={self.total_steps} accum={event.get('accumulation_steps', '')} "
f"save_interval={event.get('save_interval', '')}"
)
+ self._put(f"[MLX optimizer] {event.get('optimizer', '')}")
return
if kind == "step":
step = int(event.get("step", event.get("steps", 0)) or 0)
@@ -1439,17 +1463,25 @@ def _mlx_progress(event: dict):
metric = {
"step": step,
"total": loss,
- "primary": float(event.get("ce", loss) or loss),
- "secondary": float(event.get("aux", 0.0) or 0.0),
- "tertiary": float(event.get("temporal", 0.0) or 0.0),
- "ce": float(event.get("ce", loss) or loss),
+ "primary": float(event.get("primary", event.get("ce", loss))),
+ "secondary": float(event.get("secondary", event.get("aux", 0.0))),
+ "tertiary": float(event.get("tertiary", event.get("temporal", 0.0))),
+ "ce": float(event.get("ce", 0.0)),
"aux": float(event.get("aux", 0.0) or 0.0),
"temporal": float(event.get("temporal", 0.0) or 0.0),
"tps": tps,
+ "tokens_per_s": float(event.get("tokens_per_s", 0)),
+ "grad_norm": float(event.get("grad_norm", 0)),
+ "loss_components": event.get("loss_components", {}),
+ "recurrent_loss": float(event.get("recurrent_loss", 0)),
+ "lookahead_loss": float(event.get("lookahead_loss", 0)),
}
self._put_metric(metric)
self._put(
f"[MLX {mode.upper()}] step={step}/{self.total_steps} loss={loss:.4f} "
+ f"primary={metric['primary']:.4f} aux={metric['aux']:.4f} "
+ f"recurrent={metric['recurrent_loss']:.4f} grad_norm={metric['grad_norm']:.3f} "
+ f"tokens/s={metric['tokens_per_s']:.0f} "
f"steps/s={tps:.2f} lr={float(event.get('lr', 0.0) or 0.0):.2e} "
f"elapsed={float(event.get('elapsed_s', 0.0) or 0.0):.1f}s"
)
@@ -1480,7 +1512,7 @@ def _mlx_progress(event: dict):
return
if kind == "failed":
self._put(
- f"[MLX {mode.upper()}] stopped after rollback limit at "
+ f"[MLX {mode.upper()}] failed ({event.get('reason', 'numerical_error')}) at "
f"step={int(event.get('step', 0) or 0)}/{int(event.get('steps', self.total_steps) or self.total_steps)}"
)
return
@@ -1496,6 +1528,7 @@ def _mlx_progress(event: dict):
checkpoint_path=mlx_checkpoint_path,
save_dir=save_dir,
loader=loader,
+ validation_loader=val_loader,
prompts=prompts,
teacher_path=stage_args.teacher_path if mode == "distill" else None,
args=stage_args,
@@ -1515,7 +1548,7 @@ def _mlx_progress(event: dict):
if result.stopped:
self.status = "error"
self._put(
- f"[MLX {mode.upper()}] Stopped after rollback limit in {elapsed:.1f}s. "
+ f"[MLX {mode.upper()}] Failed ({getattr(result, 'failure_reason', 'numerical_error')}) in {elapsed:.1f}s. "
f"Last finite checkpoint → {result.checkpoint_path}"
)
if result.checkpoint_saved:
@@ -1527,10 +1560,19 @@ def _mlx_progress(event: dict):
return
model = model.to(device)
- from chronos.trainer.optim_utils import build_optimizer, get_lr, apply_lr
- optimizer = build_optimizer(model, lr=base_lr, weight_decay=weight_decay)
- train_autocast = autocast_context(device, cfg.get("dtype", "auto"))
- scaler = grad_scaler(device, cfg.get("dtype", "auto"))
+ from chronos.trainer.optim_utils import build_configured_optimizer, get_lr, apply_lr
+ optimizer = build_configured_optimizer(
+ model, model_cfg, lr=base_lr, weight_decay=weight_decay,
+ )
+ from chronos.trainer.paged_store import PagedTrainingController
+ paged_training = PagedTrainingController(model, model_cfg)
+ # One compute dtype end to end (weights, activations, KV) with
+ # FP32 masters in the optimizer; no autocast per-op casts.
+ from chronos.trainer.precision_training import attach_precision
+ precision, train_autocast, scaler = attach_precision(
+ model, optimizer, device, cfg.get("dtype", "auto"))
+ self._put(f"[precision] {precision.backend}/{precision.device}: compute={precision.compute} "
+ f"master={precision.master} ({precision.reason})")
model.train()
global_step = 0
@@ -1546,12 +1588,27 @@ def _eval_val() -> float:
model.eval()
tot, n = 0.0, 0
with torch.no_grad():
- for k, (vx, vy) in enumerate(val_loader):
+ for k, vbatch in enumerate(val_loader):
if k >= 50: # cap eval cost
break
+ if isinstance(vbatch, dict):
+ vx, vy = vbatch["input_ids"], vbatch["labels"]
+ vmedia, vmedia_mask = vbatch.get("modality_inputs"), vbatch.get("modality_mask")
+ else:
+ vx, vy = vbatch
+ vmedia = vmedia_mask = None
vx = vx.to(device); vy = vy.to(device)
- with autocast_context(device, cfg.get("dtype", "auto")):
- vout, _ = model(vx, labels=vy)
+ if vmedia:
+ vmedia = {key: value.to(device) for key, value in vmedia.items()}
+ if vmedia_mask is not None:
+ vmedia_mask = vmedia_mask.to(device)
+ with autocast_context(device, cfg.get("dtype", "auto"), model=model):
+ eval_masks = paged_training.masks_for_forward(device, allow_evict=False)
+ vout = model(
+ vx, labels=vy,
+ modality_inputs=vmedia, modality_mask=vmedia_mask,
+ available_expert_masks=eval_masks,
+ )
tot += float(vout.loss.item()); n += 1
model.train()
return tot / max(n, 1)
@@ -1564,19 +1621,35 @@ def _eval_val() -> float:
trainer = ChronosSFTTrainer(model, model_cfg, stage_args, tokenizer)
if float(getattr(model_cfg, "lambda_router_anchor", 0.0)) > 0 and loader is not None:
first_batch = next(iter(loader))
- trainer.set_calibration_batch(first_batch[0])
+ if isinstance(first_batch, dict):
+ trainer.set_calibration_batch(
+ first_batch["input_ids"],
+ modality_inputs=first_batch.get("modality_inputs"),
+ modality_mask=first_batch.get("modality_mask"),
+ )
+ else:
+ trainer.set_calibration_batch(first_batch[0])
self._put("[SFT] Router anchor reference captured.")
iters = _planned_epoch_iters(loader, stage_args.steps)
stage_global_step = 0
for epoch in range(stage_args.epochs):
epoch_start = time.monotonic()
- for step_idx, (ids, labels) in enumerate(loader, start=1):
+ for step_idx, batch in enumerate(loader, start=1):
if self._stop.is_set():
break
if step_idx > iters or (stage_args.steps is not None and stage_global_step >= stage_args.steps):
break
stage_global_step += 1
- result = trainer.train_step(ids, labels, stage_global_step, self.total_steps)
+ if isinstance(batch, dict):
+ ids, labels = batch["input_ids"], batch["labels"]
+ result = trainer.train_step(
+ ids, labels, stage_global_step, self.total_steps,
+ modality_inputs=batch.get("modality_inputs"),
+ modality_mask=batch.get("modality_mask"),
+ )
+ else:
+ ids, labels = batch
+ result = trainer.train_step(ids, labels, stage_global_step, self.total_steps)
if step_idx % stage_args.log_interval == 0 or step_idx == iters:
elapsed = max(time.monotonic() - epoch_start, 1e-6)
self._record_stage_metric(mode, stage_global_step, result, tps=step_idx / elapsed)
@@ -1691,19 +1764,33 @@ def _eval_val() -> float:
trainer = ChronosDistillTrainer(model, teacher, model_cfg, stage_args, tokenizer)
if float(getattr(model_cfg, "lambda_router_anchor", 0.0)) > 0 and loader is not None:
first = next(iter(loader))
- trainer.set_calibration_batch(first[0])
+ if isinstance(first, dict):
+ trainer.set_calibration_batch(
+ first["input_ids"], modality_inputs=first.get("modality_inputs"),
+ modality_mask=first.get("modality_mask"),
+ )
+ else:
+ trainer.set_calibration_batch(first[0])
self._put("[DISTILL] Router anchor reference captured.")
iters = _planned_epoch_iters(loader, stage_args.steps)
stage_global_step = 0
for epoch in range(stage_args.epochs):
epoch_start = time.monotonic()
- for step_idx, (ids, labels) in enumerate(loader, start=1):
+ for step_idx, batch in enumerate(loader, start=1):
if self._stop.is_set():
break
if step_idx > iters or (stage_args.steps is not None and stage_global_step >= stage_args.steps):
break
stage_global_step += 1
- result = trainer.train_step(ids, labels, stage_global_step, self.total_steps)
+ if isinstance(batch, dict):
+ result = trainer.train_step(
+ batch["input_ids"], batch["labels"], stage_global_step, self.total_steps,
+ modality_inputs=batch.get("modality_inputs"),
+ modality_mask=batch.get("modality_mask"),
+ )
+ else:
+ ids, labels = batch
+ result = trainer.train_step(ids, labels, stage_global_step, self.total_steps)
if step_idx % stage_args.log_interval == 0 or step_idx == iters:
elapsed = max(time.monotonic() - epoch_start, 1e-6)
self._record_stage_metric(mode, stage_global_step, result, tps=step_idx / elapsed)
@@ -1724,17 +1811,33 @@ def _eval_val() -> float:
for epoch in range(epochs):
if self._stop.is_set() or stop_outer:
break
- for step, (input_ids, labels) in enumerate(loader, start=1):
+ for step, batch in enumerate(loader, start=1):
if self._stop.is_set():
break
if max_steps and global_step >= max_steps:
stop_outer = True
break
+ if isinstance(batch, dict):
+ input_ids, labels = batch["input_ids"], batch["labels"]
+ media, media_mask = batch.get("modality_inputs"), batch.get("modality_mask")
+ else:
+ input_ids, labels = batch
+ media = media_mask = None
input_ids = input_ids.to(device)
labels = labels.to(device)
+ if media:
+ media = {key: value.to(device) for key, value in media.items()}
+ if media_mask is not None:
+ media_mask = media_mask.to(device)
with train_autocast:
- out, lp = model(input_ids, labels=labels)
+ paged_masks = paged_training.masks_for_forward(device)
+ result = model(
+ input_ids, labels=labels,
+ modality_inputs=media, modality_mask=media_mask,
+ available_expert_masks=paged_masks,
+ )
+ out, lp = result.legacy_tuple()
moe_layers = [l.mlp for l in model.model.layers
if isinstance(l.mlp, ChronosMOEFeedForward)]
@@ -1786,6 +1889,7 @@ def _eval_val() -> float:
model.parameters(),
1.0,
)
+ paged_training.observe()
if global_step % log_interval == 0:
now = time.monotonic()
@@ -1865,8 +1969,8 @@ def _eval_val() -> float:
with torch.no_grad():
gen = ids.clone()
for _ in range(80):
- with autocast_context(device, cfg.get("dtype", "auto")):
- o, _ = model(gen, use_cache=False)
+ with autocast_context(device, cfg.get("dtype", "auto"), model=model):
+ o = model(gen, use_cache=False)
nxt = o.logits[:, -1, :].argmax(dim=-1, keepdim=True)
gen = torch.cat([gen, nxt], dim=-1)
if int(nxt.item()) == (tokenizer.eos_token_id or -1):
@@ -1919,7 +2023,24 @@ def _synthetic_loader(self, vocab_size, seq_len, batch_size, n=50):
def _load_tokenizer(self):
from transformers import AutoTokenizer
- return AutoTokenizer.from_pretrained(chronos.deps.get_tokenizer_path())
+ path = resolve_tokenizer_dir(self.tokenizer_path)
+ tok = AutoTokenizer.from_pretrained(path)
+ if self.vocab_size and len(tok) > self.vocab_size:
+ raise ValueError(
+ f"tokenizer at {path} has {len(tok)} tokens but the model vocab_size is "
+ f"{self.vocab_size}; set vocab_size >= {len(tok)} or pick the matching tokenizer"
+ )
+ return tok
+
+
+def resolve_tokenizer_dir(configured: str | None) -> str:
+ """Config ``tokenizer_path`` (a Chronos Tokenizer Builder output) wins; else minimind."""
+ configured = os.path.expanduser(str(configured or "").strip())
+ if configured:
+ if not os.path.isfile(os.path.join(configured, "tokenizer.json")):
+ raise FileNotFoundError(f"tokenizer_path {configured!r} has no tokenizer.json")
+ return configured
+ return chronos.deps.get_tokenizer_path()
_session = TrainSession()