Project-Chronos 0.3.0__tar.gz → 0.3.1a0__tar.gz

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/CHANGELOG.md

@@ -1,5 +1,46 @@
 # Changelog
 
+## [Unreleased]
+
+## [0.3.1a0] — 2026-11-07 (beta)
+
+Core half of 0.3.1: cluster-locality training on by default, inference that
+adapts its expert paging to how the checkpoint was trained, and lazy experts in
+the bundle. The rewritten Web UI follows in 0.3.1a1.
+
+### Changed
+- Cluster locality is on by default for new training: `train_chronos*.py --lambda_cluster_locality` defaults to 0.3, the planner and presets set λ = 0.3 for MoE models with more than 2k experts, and the window follows the data (`recommended_window`: power of two nearest the median SFT record, clamped to [16, seq_len]; 32 without record lengths). The `ChronosConfig` class default stays 0.0, so a stored recipe without the field keeps its original objective on resume, and checkpoints written before the field existed load as λ = 0. The λ = 0.3 default comes from the synthetic sweep below (from-scratch training); its effect on real data and on fine-tuning has not been measured.
+- Lazy-expert inference picks its page layout from the checkpoint (`paging="auto"`): weights trained with λ > 0 are paged one training cluster per SSD read, other weights one expert per read. `paging="cluster"|"expert"` forces a layout. The layout is part of the expert-cache directory name (`…_mlx_c2`, `…_c2`), so a cache written for one layout is never reused for the other. Per-expert cache names are unchanged.
+
+### Added
+- `chronos/lazy_experts.py`: one lazy-expert implementation for the Web UI, bundle and server — `LazyOptions` (paging, residency, decay, decode/verify waits, RAM ratio; validated), `page_layout`, `checkpoint_cache_dir`, the expert budget, and `build_mlx_lazy_engine` (trunk only in memory, cache built or reused).
+- Torch offload writes cluster layouts as `cluster_<c>.ctsr` files with a v2 manifest (`ClusterStorage.write_clusters_from_states`); `setup_from_state_dict/_reader(clusters=…)`.
+- MLX `generate(decode_wait_s=0.012, verify_wait_s=0.02, repetition_penalty=1.0)`. The penalty uses the CTRL rule and is rejected together with `speculative=True`, which needs penalty 1.0 to stay exact.
+- `chronos bundle translate|eval --lazy-experts` with `--paging`, `--residency`, `--decay`, `--ram-ratio`, `--decode-wait`, `--verify-wait`. The cache lives outside the bundle (`$CHRONOS_EXPERT_CACHE`, default `~/.cache/chronos/expert_cache`); `eval` also gains `--backend` and `--dtype`.
+- `chronos bundle finetune --[no-]cluster-locality --lambda-cluster --cluster-window --clusters` (on by default; dense models and models with fewer than two clusters stay off). The spec records where each value came from (`extra.cluster_locality_source`); the window comes from the median SFT record length of `--data`.
+- `chronos serve --paging auto|cluster|expert`; `/api/ps` metadata reports the resolved layout.
+- Web UI Inference tab: paging, planner decay and the two prefetch waits, with four-language labels.
+
+- Next-use residency planning for lazy MLX experts (`chronos/residency_planner.py`, `generate(residency="planned")`). The planner records the page-access trace of every forward, keyed by (layer, loop iteration). The previous forward of the same kind, and in speculative decoding the queue of draft traces since the last verify, then predicts when each page is needed next. The store evicts the page needed latest (`MLXExpertStore.set_eviction_ranker`). The running dispatch's pending experts are never evicted. Pages needed early are prefetched only into spare slots. Greedy output is unchanged. The Web UI Inference tab has a switch for it, on by default (MLX lazy only); the stats table shows the eviction policy.
+- Cluster-locality training constraint `lambda_cluster_locality`, with `cluster_locality_clusters` (0 = ⌈E/2k⌉) and `cluster_locality_window`. The loss is 1 − I(window; cluster)/log C over every MoE dispatch, recurrent iterations included. It keeps a stream inside few expert clusters and spreads different streams over different clusters. The partition (`chronos/expert_clusters.py`) is the one lazy inference pages by. Torch and MLX trainers and all stage scripts take it.
+- `tools/bench_residency.py`: LRU window vs planned residency, plain and speculative, on a real checkpoint.
+- `tools/bench_cluster_locality.py`: a tiny-model training A/B of the constraint, reporting held-out CE, cluster entropies, topic purity and lazy-decode cluster reads.
+
+### Measured (reports/residency_20261107/)
+- 50M s13996 checkpoint, 16 FLORES+ dev prompts, greedy, CPU, 12 hot / 12 RAM pages.
+  - Plain decoding: planned residency needs 4438 synchronous SSD loads instead of 6876 (−35%). Throughput is 22.47 vs 23.17 tok/s.
+  - Speculative decoding: 6258 vs 9297 loads (−33%), and 18.96 vs 16.03 tok/s.
+  - All outputs are identical to plain greedy (16/16).
+  - The checkpoint was trained without the constraint and routes near-uniformly over all 36 experts within a stream.
+- Tiny model on synthetic 4-topic streams, 300 steps, λ=0.3 vs 0.
+  - Window cluster entropy falls from 0.49 to 0.005; topic→cluster purity rises from 0.81 to 0.999.
+  - Marginal cluster entropy rises from 0.71 to 0.99, and expert-usage entropy is unchanged (0.87).
+  - Held-out CE is 1.712 vs 1.703.
+  - Cluster-paged lazy decoding needs 0.023 vs 1.48 cluster reads per token.
+
+### Fixed
+- MLX lazy store: an expert evicted by the prefetch thread between an SSD read and promotion raised "could not be loaded on demand". Reads now return the page they loaded.
+
 ## [0.3.0] — 2026-11-06
 
 v3 kernel: recurrent-depth MoE. Quality results for the 50M en→zh demo are not

{project_chronos-0.3.0/Project_Chronos.egg-info → project_chronos-0.3.1a0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Project_Chronos
-Version: 0.3.0
+Version: 0.3.1a0
 Summary: Storage-aware recurrent-depth MoE: lookahead expert prefetch, SSD-lazy experts, hybrid MLA/SWA attention, self-speculative decoding
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://github.com/FonaTech/Project_Chronos

{project_chronos-0.3.0 → project_chronos-0.3.1a0/Project_Chronos.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Project_Chronos
-Version: 0.3.0
+Version: 0.3.1a0
 Summary: Storage-aware recurrent-depth MoE: lookahead expert prefetch, SSD-lazy experts, hybrid MLA/SWA attention, self-speculative decoding
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://github.com/FonaTech/Project_Chronos

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/Project_Chronos.egg-info/SOURCES.txt

@@ -25,9 +25,12 @@ chronos/app.py
 chronos/cli.py
 chronos/config_schema.py
 chronos/deps.py
+chronos/expert_clusters.py
 chronos/export.py
+chronos/lazy_experts.py
 chronos/migrate.py
 chronos/precision.py
+chronos/residency_planner.py
 chronos/verify.py
 chronos/backend/__init__.py
 chronos/backend/dispatcher.py
@@ -144,12 +147,14 @@ chronos/tuning/_base_tuner.py
 chronos/tuning/chronos_auto_tuner.py
 tests/test_bundle.py
 tests/test_interface_regressions.py
+tests/test_lazy_experts.py
 tests/test_longrun.py
 tests/test_mlx_speculative.py
 tests/test_moe_dispatch.py
 tests/test_monitor.py
 tests/test_planner_ui.py
 tests/test_precision_chain.py
+tests/test_residency_cluster.py
 tests/test_smoke.py
 tests/test_smoke_cuda.py
 tests/test_speculative.py

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/chronos/__init__.py

@@ -1,5 +1,5 @@
 # Project Chronos: On-device low-latency lookahead dual-layer MoE inference
-__version__ = "0.3.0"
+__version__ = "0.3.1a0"
 
 # Bootstrap minimind dependency before any submodule imports
 from chronos.deps import ensure_minimind as _ensure_minimind

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/chronos/bundle/cli.py

@@ -12,6 +12,43 @@ def _bundle_dir(value: str | None) -> str:
     return os.path.abspath(value or os.environ.get("CHRONOS_BUNDLE") or ".")
 
 
+def _add_lazy_args(parser) -> None:
+    """SSD-lazy expert options (MLX backend); defaults are the engine defaults."""
+    from chronos.lazy_experts import (
+        DEFAULT_DECODE_WAIT_S, DEFAULT_RAM_RATIO, DEFAULT_RESIDENCY, DEFAULT_RESIDENCY_DECAY,
+        DEFAULT_VERIFY_WAIT_S, RESIDENCY_MODES,
+    )
+    g = parser.add_argument_group("lazy experts (MLX)")
+    g.add_argument("--lazy-experts", action="store_true",
+                   help="keep only the trunk in memory and page routed experts from an SSD cache")
+    g.add_argument("--paging", default="auto", choices=["auto", "cluster", "expert"],
+                   help="page layout; auto = by training cluster if the model was trained with cluster "
+                        "locality, else per expert (default: auto)")
+    g.add_argument("--residency", default=DEFAULT_RESIDENCY, choices=list(RESIDENCY_MODES),
+                   help=f"planned = next-use eviction, window = LRU + protected window (default: {DEFAULT_RESIDENCY})")
+    g.add_argument("--decay", type=float, default=DEFAULT_RESIDENCY_DECAY,
+                   help=f"planned residency: weight of older forwards, [0, 1) (default: {DEFAULT_RESIDENCY_DECAY})")
+    g.add_argument("--ram-ratio", type=float, default=DEFAULT_RAM_RATIO,
+                   help=f"expert pages kept loaded, × the routed working set (default: {DEFAULT_RAM_RATIO})")
+    g.add_argument("--decode-wait", type=float, default=DEFAULT_DECODE_WAIT_S,
+                   help=f"seconds a decode step waits for prefetched pages (default: {DEFAULT_DECODE_WAIT_S})")
+    g.add_argument("--verify-wait", type=float, default=DEFAULT_VERIFY_WAIT_S,
+                   help=f"seconds a speculative verify waits for prefetched pages (default: {DEFAULT_VERIFY_WAIT_S})")
+
+
+def _lazy_options(a):
+    from chronos.lazy_experts import LazyOptions
+    return LazyOptions(paging=a.paging, residency=a.residency, residency_decay=a.decay, ram_ratio=a.ram_ratio,
+                       decode_wait_s=a.decode_wait, verify_wait_s=a.verify_wait).validated()
+
+
+def _bundle_model(bundle: str, a):
+    from .runtime import BundleModel
+    backend = "mlx" if a.lazy_experts else a.backend
+    return BundleModel(bundle, device=a.device, dtype=a.dtype, backend=backend, lazy_experts=a.lazy_experts,
+                       lazy_options=_lazy_options(a) if a.lazy_experts else None)
+
+
 def main(argv=None) -> int:
     p = argparse.ArgumentParser(prog="chronos bundle", description=__doc__)
     sub = p.add_subparsers(dest="cmd", required=True)
@@ -56,14 +93,18 @@ def main(argv=None) -> int:
                    help="recurrent-depth self-speculative decoding (same greedy output, fewer full-depth forwards)")
     t.add_argument("--draft-depth", type=int, default=1)
     t.add_argument("--gamma", type=int, default=4, help="initial draft length (adapts between 1 and 8)")
+    _add_lazy_args(t)
 
     e = sub.add_parser("eval", help="BLEU/chrF on the bundled FLORES+ split")
     e.add_argument("--bundle", default=None)
     e.add_argument("--split", default="devtest", choices=["dev", "devtest"])
     e.add_argument("-n", type=int, default=0, help="first n sentences (0 = all)")
     e.add_argument("--repetition-penalty", type=float, default=1.2)
+    e.add_argument("--backend", default="torch", choices=["torch", "mlx"])
     e.add_argument("--device", default="auto")
+    e.add_argument("--dtype", default="float32")
     e.add_argument("--out", default=None, help="write the report JSON here")
+    _add_lazy_args(e)
 
     d = sub.add_parser("prepare-data", help="your parallel data → shards tokenized with the bundle tokenizer")
     d.add_argument("inputs", nargs="+", help=".jsonl/.csv/.tsv/.parquet files with zh and en columns")
@@ -84,6 +125,16 @@ def main(argv=None) -> int:
     f.add_argument("--hours", type=float, default=None, help="override the tier's wall-clock budget")
     f.add_argument("--seed", type=int, default=None)
     f.add_argument("--dry-run", action="store_true", help="write the spec, do not start training")
+    from chronos.expert_clusters import DEFAULT_LAMBDA
+    f.add_argument("--cluster-locality", action=argparse.BooleanOptionalAction, default=True,
+                   help="train with the expert-cluster locality constraint so lazy inference can page "
+                        "by cluster (default: on)")
+    f.add_argument("--lambda-cluster", type=float, default=None,
+                   help=f"constraint strength (default: {DEFAULT_LAMBDA}, from a measured sweep)")
+    f.add_argument("--cluster-window", type=int, default=None,
+                   help="tokens per locality window (default: nearest power of two to the median SFT "
+                        "record length of --data, clamped to [16, seq_len])")
+    f.add_argument("--clusters", type=int, default=None, help="number of clusters (default: ceil(E / 2k))")
 
     s = sub.add_parser("serve", help="Ollama/OpenAI-protocol server over the bundle's model/ (loopback only)")
     s.add_argument("--bundle", default=None)
@@ -120,7 +171,6 @@ def _run(a) -> int:
         _print(builder.record_quality_probe(_bundle_dir(a.bundle), a.n, device=a.device))
         return 0
     if a.cmd == "translate":
-        from .runtime import BundleModel
         lines = list(a.text) and [" ".join(a.text)]
         if a.file:
             with open(a.file, encoding="utf-8") as fh:
@@ -128,19 +178,30 @@ def _run(a) -> int:
         elif not lines:
             lines = [ln.strip() for ln in sys.stdin if ln.strip()]
         rp = a.repetition_penalty if a.repetition_penalty is not None else (1.0 if a.speculative else 1.2)
-        bm = BundleModel(_bundle_dir(a.bundle), device=a.device, dtype=a.dtype, backend=a.backend)
-        for hyp in bm.translate(lines, a.direction, max_new_tokens=a.max_new_tokens, repetition_penalty=rp,
-                                speculative=a.speculative, draft_depth=a.draft_depth, gamma=a.gamma):
-            print(hyp, flush=True)
+        bm = _bundle_model(_bundle_dir(a.bundle), a)
+        try:
+            for hyp in bm.translate(lines, a.direction, max_new_tokens=a.max_new_tokens, repetition_penalty=rp,
+                                    speculative=a.speculative, draft_depth=a.draft_depth, gamma=a.gamma):
+                print(hyp, flush=True)
+        finally:
+            bm.close()
+        if bm.last_stats.get("lazy"):
+            print(json.dumps({"lazy": bm.last_stats["lazy"]}, ensure_ascii=False), file=sys.stderr)
         if bm.last_stats.get("speculative"):
             print(f"[speculative] {json.dumps(bm.last_stats['speculative'])}", file=sys.stderr)
         return 0
     if a.cmd == "eval":
-        from .runtime import BundleModel, eval_pairs, score
+        from .runtime import eval_pairs, score
         bundle = _bundle_dir(a.bundle)
         pairs = eval_pairs(bundle, a.split)
         pairs = pairs[: a.n] if a.n else pairs
-        report = score(BundleModel(bundle, device=a.device), pairs, repetition_penalty=a.repetition_penalty)
+        bm = _bundle_model(bundle, a)
+        try:
+            report = score(bm, pairs, repetition_penalty=a.repetition_penalty)
+        finally:
+            bm.close()
+        if bm.last_stats.get("lazy"):
+            report["lazy"] = bm.last_stats["lazy"]
         report["set"] = f"flores_plus_{a.split}"
         _print(report)
         if a.out:
@@ -157,7 +218,9 @@ def _run(a) -> int:
         data = os.path.abspath(a.data)
         shards = os.path.join(data, "shards") if os.path.isdir(os.path.join(data, "shards")) else data
         spec = write_finetune_spec(_bundle_dir(a.bundle), shards, a.run_dir, tier=a.tier, mode=a.mode,
-                                   budget_hours=a.hours, seed=a.seed)
+                                   budget_hours=a.hours, seed=a.seed, cluster_locality=a.cluster_locality,
+                                   lambda_cluster=a.lambda_cluster, cluster_window=a.cluster_window,
+                                   clusters=a.clusters)
         print(f"[finetune] spec: {spec}")
         if a.dry_run:
             print(f"[finetune] start it with: chronos longrun --spec {spec}")

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/chronos/bundle/docs.py

@@ -157,6 +157,13 @@ def _readme_en(m: dict, recs: dict) -> str:
         "`--speculative` (recurrent-depth self-speculative decoding) gives the same greedy output as plain "
         "decoding and reports its acceptance rate on stderr. It needs `--repetition-penalty 1.0`, its default. "
         "On CPU at this model size it was measured slower than plain decoding.", "",
+        "`--lazy-experts` (MLX) keeps only the trunk in memory and pages the routed experts from an SSD cache "
+        "(`~/.cache/chronos/expert_cache`, or `$CHRONOS_EXPERT_CACHE`). The greedy output is the same. The page "
+        "layout adapts to the weights: one file per expert cluster when the model was trained with the "
+        "cluster-locality constraint, one per expert otherwise (`--paging auto|cluster|expert`). "
+        "`--residency planned` (default) evicts the page needed latest; `window` keeps LRU order. "
+        "`--decay`, `--ram-ratio`, `--decode-wait` and `--verify-wait` tune it; `--help` lists the defaults. "
+        "The page counts are printed to stderr.", "",
         "`bin/chronos` sets `CHRONOS_BUNDLE` and `CHRONOS_MINIMIND_DIR` (the vendored minimind in "
         "`third_party/`), so nothing is fetched from the network. Without `bin/`, export both variables and call "
         "`chronos bundle …` directly.", "",
@@ -176,6 +183,12 @@ def _readme_en(m: dict, recs: dict) -> str:
         "Monitor a run with `bin/chronos monitor --runs-dir work`. The fine-tuned checkpoint lands in "
         "`work/run/*_ft_sft.best.pth`; rebuild a bundle from it with `chronos bundle build --run-dir work/run "
         "--phase sft`.", "",
+        "Fine-tuning trains with the expert-cluster locality constraint by default, so the result can be paged "
+        "by cluster at inference (`--no-cluster-locality` turns it off). Defaults: `--lambda-cluster 0.3` (from "
+        "a measured sweep on a synthetic task with from-scratch training; its effect on a fine-tune has not "
+        "been measured) and `--cluster-window` = the power of two nearest the median SFT record length of "
+        "your data. The spec records where each value came from (`extra.cluster_locality_source`). Models "
+        "with fewer than two clusters (E ≤ 2k) train without it.", "",
         "Tiers (from `recommended.json`; every value is labelled with its source):", "",
         "| Tier | Hours | SFT peak LR | Tokens / step | Warmup | Use |", "|---|---|---|---|---|---|",
         *_tier_rows(recs), "",
@@ -238,6 +251,12 @@ def _readme_zh(m: dict, recs: dict) -> str:
         "```", "",
         "`--speculative`（循环深度自推测解码）的 greedy 输出与普通解码一致，接受率打印到 stderr。"
         "它要求 `--repetition-penalty 1.0`（此时的默认值）。在这个规模的模型上用 CPU 实测，它比普通解码慢。", "",
+        "`--lazy-experts`（MLX）只把主干留在内存里，路由专家从 SSD 缓存按需分页加载"
+        "（`~/.cache/chronos/expert_cache`，或 `$CHRONOS_EXPERT_CACHE`），greedy 输出不变。分页方式按权重自适应："
+        "模型训练时开了专家簇局部性约束就按簇读文件，否则按单个专家读（`--paging auto|cluster|expert`）。"
+        "`--residency planned`（默认）优先淘汰最晚才会用到的页，`window` 保持 LRU 顺序。"
+        "可用 `--decay`、`--ram-ratio`、`--decode-wait`、`--verify-wait` 调整，默认值见 `--help`。"
+        "页读取统计打印到 stderr。", "",
         "`bin/chronos` 会设置 `CHRONOS_BUNDLE` 和 `CHRONOS_MINIMIND_DIR`（指向 `third_party/` 里随包附带的 minimind），"
         "全程不需要联网。不用 `bin/` 时，手动导出这两个变量后直接调用 `chronos bundle …`。", "",
         "## 下一阶段：SFT 微调（Apple silicon / MLX）", "",
@@ -255,6 +274,10 @@ def _readme_zh(m: dict, recs: dict) -> str:
         "`--mode pretrain` 在新文本上继续预训练；`--mode pretrain+sft` 两者都做（时间按 40/60 分配）。"
         "用 `bin/chronos monitor --runs-dir work` 监控训练。微调后的检查点在 `work/run/*_ft_sft.best.pth`，"
         "可用 `chronos bundle build --run-dir work/run --phase sft` 重新打包。", "",
+        "微调默认带专家簇局部性约束训练，这样得到的模型在推理时可以按簇分页（`--no-cluster-locality` 关闭）。"
+        "默认值：`--lambda-cluster 0.3`（来自合成任务上从头训练的实测扫描，对微调的效果尚未实测）；"
+        "`--cluster-window` 取最接近你数据中位 SFT 记录长度的 2 的幂。spec 中记录了每个值的来源"
+        "（`extra.cluster_locality_source`）。簇数少于 2（E ≤ 2k）的模型不加这项约束。", "",
         "档位（来自 `recommended.json`，每个数值都标注了来源）：", "",
         "| 档位 | 小时 | SFT 峰值 LR | 每步 token | Warmup | 用途 |", "|---|---|---|---|---|---|",
         *_tier_rows(recs), "",

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/chronos/bundle/finetune.py

@@ -102,14 +102,20 @@ def _check_shards(shards_dir: str, phases: list[dict]) -> None:
 
 
 def write_finetune_spec(bundle_dir: str, shards_dir: str, run_dir: str, *, tier: str = "quick", mode: str = "sft",
-                        budget_hours: float | None = None, seed: int | None = None) -> str:
-    """Write ``<run_dir>/finetune_spec.json`` for ``chronos longrun --spec``."""
+                        budget_hours: float | None = None, seed: int | None = None, cluster_locality: bool = True,
+                        lambda_cluster: float | None = None, cluster_window: int | None = None,
+                        clusters: int | None = None) -> str:
+    """Write ``<run_dir>/finetune_spec.json`` for ``chronos longrun --spec``.
+
+    The expert-cluster locality constraint is on by default; see
+    ``recommend.cluster_locality_config`` for the defaults and their sources."""
     from .recommend import longrun_spec
     bundle_dir = os.path.abspath(bundle_dir)
     with open(os.path.join(bundle_dir, "recipes", "source_run_spec.json"), encoding="utf-8") as fh:
         recipe = json.load(fh)
     spec = longrun_spec(bundle_dir, recipe, run_dir=run_dir, shards_dir=shards_dir, tier=tier, mode=mode,
-                        budget_hours=budget_hours, seed=seed)
+                        budget_hours=budget_hours, seed=seed, cluster_locality=cluster_locality,
+                        lambda_cluster=lambda_cluster, cluster_window=cluster_window, clusters=clusters)
     _check_shards(spec["shards_dir"], spec["phases"])
     os.makedirs(spec["run_dir"], exist_ok=True)
     path = os.path.join(spec["run_dir"], "finetune_spec.json")

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/chronos/bundle/recommend.py

@@ -98,19 +98,100 @@ def tier_phases(recipe: dict, tier: str, mode: str = "sft") -> list[dict]:
     raise ValueError(f"unknown mode {mode!r}; choose sft, pretrain or pretrain+sft")
 
 
+def median_record_tokens(shards_dir: str, max_files: int = 8) -> float | None:
+    """Median SFT record length (tokens) from the ``.idx`` offsets under ``<shards>/sft/train``.
+
+    Reads at most ``max_files`` index files (offsets only, no token data).
+    Returns None when the shards have no SFT records (packed pretrain text
+    has no record boundaries)."""
+    import glob
+    import numpy as np
+    files = sorted(glob.glob(os.path.join(shards_dir, "sft", "train", "*.idx")))[:max_files]
+    lengths = [np.diff(np.fromfile(f, dtype=np.int64)) for f in files]
+    lengths = [x for x in lengths if x.size]
+    if not lengths:
+        return None
+    return float(np.median(np.concatenate(lengths)))
+
+
+def cluster_locality_config(model_config: dict, shards_dir: str | None, *, enabled: bool = True,
+                            lambda_cluster: float | None = None, window: int | None = None,
+                            clusters: int | None = None, seq_len: int | None = None) -> dict:
+    """Cluster-locality fields for a fine-tune and where each value comes from.
+
+    On by default at ``DEFAULT_LAMBDA`` (measured sweep on a synthetic
+    multi-topic task, from-scratch training; its effect on a fine-tune of
+    weights trained without the constraint has not been measured).  Off for
+    dense models and when there are fewer than two clusters."""
+    from chronos.expert_clusters import DEFAULT_LAMBDA, recommended_window
+    from chronos.model.config import ChronosConfig
+    # The values training will use: longrun builds ChronosConfig(**model_config).
+    cfg = ChronosConfig(**{k: v for k, v in model_config.items() if v is not None})
+    e = int(getattr(cfg, "num_experts", 0) or 0)
+    k = int(getattr(cfg, "num_experts_per_tok", 1) or 1)
+    seq = int(seq_len or getattr(cfg, "max_position_embeddings", 512) or 512)
+    usable = bool(getattr(cfg, "use_moe", True)) and e > 2 * k
+    lam = 0.0 if not (enabled and usable) else float(DEFAULT_LAMBDA if lambda_cluster is None else lambda_cluster)
+    if lam < 0:
+        raise ValueError("lambda_cluster must be >= 0")
+    median = median_record_tokens(shards_dir) if (shards_dir and window is None) else None
+    win = int(window) if window is not None else recommended_window(median, seq)
+    if win < 0:
+        raise ValueError("cluster_window must be >= 0 (0 = whole sequence)")
+    values = {"lambda_cluster_locality": lam, "cluster_locality_window": win,
+              "cluster_locality_clusters": int(clusters or 0)}
+    if not enabled:
+        why = "off (--no-cluster-locality)"
+    elif not usable:
+        why = "off: dense model or fewer than two clusters"
+    elif lambda_cluster is None:
+        why = f"default {DEFAULT_LAMBDA} (measured sweep, synthetic task; not measured for fine-tuning)"
+    else:
+        why = "user value"
+    if window is not None:
+        wsrc = "user value"
+    elif median is not None:
+        wsrc = f"rule: power of two nearest the median SFT record ({median:.0f} tokens), clamped to [16, {seq}]"
+    else:
+        wsrc = "default (no SFT record lengths in the data)"
+    return {"values": values, "source": {"lambda_cluster_locality": why, "cluster_locality_window": wsrc,
+                                         "cluster_locality_clusters": "user value" if clusters else "ceil(E / 2k)"}}
+
+
 def longrun_spec(bundle_dir: str, recipe: dict, *, run_dir: str, shards_dir: str, tier: str,
-                 mode: str = "sft", budget_hours: float | None = None, seed: int | None = None) -> dict:
-    """A ``chronos longrun`` RunSpec (as a dict) that starts from the bundle weights."""
+                 mode: str = "sft", budget_hours: float | None = None, seed: int | None = None,
+                 cluster_locality: bool = True, lambda_cluster: float | None = None,
+                 cluster_window: int | None = None, clusters: int | None = None) -> dict:
+    """A ``chronos longrun`` RunSpec (as a dict) that starts from the bundle weights.
+
+    The cluster-locality constraint is on by default (``cluster_locality_config``)."""
     from .builder import load_bundle_manifest
     manifest = load_bundle_manifest(bundle_dir)
     bundle_dir = os.path.abspath(bundle_dir)
     smoke = tier == "smoke"
+    model_config = copy.deepcopy(recipe["model_config"])
+    phases = tier_phases(recipe, tier, mode)
+    seq_len = min((int(p["seq_len"]) for p in phases if p.get("seq_len")), default=None)
+    # Decide from the weights being fine-tuned: the recipe may omit fields that
+    # took class defaults when it ran.
+    topology = dict(model_config)
+    try:
+        from chronos.model.checkpoint import read_checkpoint_config
+        saved, _ = read_checkpoint_config(os.path.join(bundle_dir, manifest["model"]["checkpoint"]))
+        topology.update({k: saved[k] for k in ("num_experts", "num_experts_per_tok", "use_moe",
+                                               "max_position_embeddings") if saved and k in saved})
+    except (OSError, ValueError):
+        pass
+    locality = cluster_locality_config(topology, shards_dir, enabled=cluster_locality,
+                                       lambda_cluster=lambda_cluster, window=cluster_window, clusters=clusters,
+                                       seq_len=seq_len)
+    model_config.update(locality["values"])
     spec = {
         "run_dir": os.path.abspath(run_dir),
         "shards_dir": os.path.abspath(shards_dir),
         "tokenizer_dir": os.path.join(bundle_dir, manifest["tokenizer"]["dir"]),
-        "model_config": copy.deepcopy(recipe["model_config"]),
-        "phases": tier_phases(recipe, tier, mode),
+        "model_config": model_config,
+        "phases": phases,
         "budget_hours": float(budget_hours if budget_hours is not None else TIERS[tier]["hours"]),
         "eval_reserve_fraction": 0.1 if smoke else 0.04,
         "log_every_steps": 2 if smoke else 20,
@@ -126,7 +207,7 @@ def longrun_spec(bundle_dir: str, recipe: dict, *, run_dir: str, shards_dir: str
         "init_checkpoint": os.path.join(bundle_dir, manifest["model"]["checkpoint"]),
         "export_prefix": manifest["name"].replace("-", "_") + "_ft",
         "extra": {"raw_root": os.path.join(bundle_dir, "eval"), "from_bundle": manifest["name"],
-                  "tier": tier, "mode": mode},
+                  "tier": tier, "mode": mode, "cluster_locality_source": locality["source"]},
     }
     return spec
 
@@ -158,6 +239,11 @@ def build_recommendations(recipe: dict, planner_preset: dict | None, decode: dic
                                         for p in recipe.get("phases", [])}, "recipe"),
             "model_lambdas": _labelled({k: v for k, v in recipe.get("model_config", {}).items()
                                         if k.startswith(("lambda_", "ponder", "fallback_mask"))}, "recipe"),
+            "cluster_locality": _labelled(
+                cluster_locality_config(recipe.get("model_config", {}), None)["values"], "rule",
+                "on by default for `chronos bundle finetune` (--no-cluster-locality to disable); lambda from "
+                "a measured sweep on a synthetic task (from-scratch training), window from the median SFT "
+                "record length of the fine-tune data when it has one"),
         },
         "torch_stages": torch_stages,
         "data": _labelled({

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/chronos/bundle/runtime.py

@@ -7,6 +7,11 @@ falls into on long inputs.  ``speculative=True`` uses recurrent-depth
 self-speculative decoding (same greedy output, fewer full-depth forwards);
 it needs ``repetition_penalty == 1`` because the penalty changes the target
 distribution position by position.
+
+``lazy_experts=True`` (MLX) keeps only the trunk in memory and pages routed
+experts from an SSD cache (``chronos.lazy_experts``): by training cluster when
+the checkpoint was trained with cluster locality, per expert otherwise, with
+next-use planned residency by default.  The greedy output is the same.
 """
 from __future__ import annotations
 
@@ -20,7 +25,8 @@ LANG_TAGS = {"en_zh": ("<|en|>", "<|zh|>"), "zh_en": ("<|zh|>", "<|en|>")}
 class BundleModel:
     """Weights + tokenizer of one bundle, ready for greedy translation."""
 
-    def __init__(self, bundle_dir: str, *, device: str = "cpu", dtype: str = "float32", backend: str = "torch"):
+    def __init__(self, bundle_dir: str, *, device: str = "cpu", dtype: str = "float32", backend: str = "torch",
+                 lazy_experts: bool = False, lazy_options=None, cache_root: str | None = None):
         import torch
         from tokenizers import Tokenizer
         from chronos.model.checkpoint import chronos_config_from_checkpoint, load_checkpoint_into_model
@@ -35,6 +41,13 @@ class BundleModel:
         self.config, _ = chronos_config_from_checkpoint(self.checkpoint)
         if self.tokenizer.get_vocab_size() > int(self.config.vocab_size):
             raise ValueError("bundle tokenizer is larger than the model vocabulary")
+        self.prompt_format = self.manifest.get("prompt_format", "pretrain_tags")
+        tid = self.tokenizer.token_to_id
+        self.stop_ids = {tid(t) for t in STOP_TOKENS if tid(t) is not None}
+        self.engine, self.lazy_info = None, {}
+        if lazy_experts:
+            self._init_lazy(backend, device, dtype, lazy_options, cache_root)
+            return
         model = ChronosForCausalLM(self.config)
         load_checkpoint_into_model(model, self.checkpoint, map_location="cpu")
         self.backend, self.last_stats = backend, {}
@@ -59,9 +72,34 @@ class BundleModel:
             self.device = device
         else:
             raise ValueError(f"backend must be 'torch' or 'mlx', not {backend!r}")
-        self.prompt_format = self.manifest.get("prompt_format", "pretrain_tags")
-        tid = self.tokenizer.token_to_id
-        self.stop_ids = {tid(t) for t in STOP_TOKENS if tid(t) is not None}
+
+    def _init_lazy(self, backend, device, dtype, lazy_options, cache_root):
+        import mlx.core as mx
+        from chronos.lazy_experts import LazyOptions, build_mlx_lazy_engine, warm_ids
+        if backend != "mlx":
+            raise ValueError("lazy_experts needs backend='mlx'")
+        if not getattr(self.config, "use_moe", True):
+            raise ValueError("lazy_experts needs a MoE checkpoint")
+        if device == "cpu":
+            mx.set_default_device(mx.cpu)
+        elif device not in ("auto", "gpu"):
+            raise ValueError(f"MLX backend runs on 'gpu' or 'cpu', not {device!r}")
+        self.lazy_options = (lazy_options or LazyOptions()).validated()
+        # Outside the bundle (it may be read-only and its files are hashed).
+        root = cache_root or os.environ.get("CHRONOS_EXPERT_CACHE") or os.path.join(
+            os.path.expanduser("~"), ".cache", "chronos", "expert_cache")
+        self.engine, self.lazy_info = build_mlx_lazy_engine(self.checkpoint, self.config, self.lazy_options,
+                                                            cache_root=root, dtype=dtype)
+        self.engine.store.warm_up(warm_ids(self.config, self.lazy_info))
+        self.model, self.backend, self.last_stats = self.engine.model, "mlx", {}
+        self.precision = None
+        self.device = "cpu" if device == "cpu" else "gpu"
+
+    def close(self) -> None:
+        """Stop the lazy engine's prefetch thread (no-op without lazy experts)."""
+        if self.engine is not None:
+            self.engine.stop()
+            self.engine = None
 
     # ---- prompts ---------------------------------------------------------------
     def directions(self) -> list[str]:
@@ -85,7 +123,9 @@ class BundleModel:
         if speculative and repetition_penalty != 1.0:
             raise ValueError("speculative decoding needs repetition_penalty=1.0 (the penalty changes the target "
                              "distribution at every position)")
-        if speculative:
+        if self.engine is not None:
+            out = self._lazy(prompt, limit, repetition_penalty, speculative, draft_depth, gamma)
+        elif speculative:
             out = self._speculative(prompt, limit, draft_depth, gamma)
         elif self.backend == "mlx":
             out = self._greedy_mlx(prompt, limit, repetition_penalty)
@@ -138,6 +178,29 @@ class BundleModel:
             logits, _, cache = self.model(mx.array([[nxt]], dtype=mx.int32), cache=cache)
         return out
 
+    def _lazy(self, prompt, limit, repetition_penalty, speculative, draft_depth, gamma):
+        import mlx.core as mx
+        kw = dict(speculative=True, draft_depth=draft_depth, gamma=gamma) if speculative else {}
+        out = []
+        for tok in self.engine.generate(mx.array([prompt], dtype=mx.int32), max_new_tokens=limit, temperature=0.0,
+                                        top_p=1.0, eos_token_id=sorted(self.stop_ids),
+                                        repetition_penalty=float(repetition_penalty),
+                                        **self.lazy_options.generate_kwargs(), **kw):
+            if tok in self.stop_ids:
+                break
+            out.append(int(tok))
+        st = dict(self.engine.last_stats or {})
+        agg = self.last_stats.setdefault("lazy", {"tokens": 0, "sync_ssd_loads": 0, "cluster_loads": 0})
+        agg["tokens"] += len(out)
+        agg["sync_ssd_loads"] += int(st.get("sync_ssd_loads", 0))
+        agg["cluster_loads"] = int(self.engine.store.stats().get("cluster_loads", 0))
+        if speculative and st.get("speculative"):
+            spec = self.last_stats.setdefault("speculative", {"proposed": 0, "accepted": 0, "verify_calls": 0,
+                                                              "draft_calls": 0, "generated": 0})
+            for k in spec:
+                spec[k] += int(st["speculative"].get(k, 0))
+        return out
+
     def _speculative(self, prompt, limit, draft_depth, gamma):
         eos = sorted(self.stop_ids)
         if self.backend == "mlx":
@@ -169,6 +232,11 @@ class BundleModel:
             gen = self.greedy(self.prompt_ids(src, direction), limit, repetition_penalty, speculative=speculative,
                               draft_depth=draft_depth, gamma=gamma)
             out.append(self.tokenizer.decode(gen, skip_special_tokens=True))
+        lazy = self.last_stats.get("lazy")
+        if lazy:
+            lazy.update({k: self.lazy_info[k] for k in ("paging_resolved", "num_pages", "trained_with_clusters",
+                                                        "effective_expert_budget") if k in self.lazy_info})
+            lazy["residency"] = self.lazy_options.residency
         spec = self.last_stats.get("speculative")
         if spec:
             spec["acceptance_rate"] = round(spec["accepted"] / spec["proposed"], 4) if spec["proposed"] else 0.0

{project_chronos-0.3.0 → project_chronos-0.3.1a0}/chronos/config_schema.py

@@ -37,6 +37,7 @@ MODEL_FIELDS = (
     "lookahead_steps", "layer_aware_lookahead", "layer_aware_cache",
     "lambda_balance", "lambda_temporal", "lambda_lookahead",
     "lambda_lookahead_topk", "lambda_lookahead_union", "lambda_router_locality",
+    "lambda_cluster_locality", "cluster_locality_clusters", "cluster_locality_window",
     "lambda_router_anchor", "fallback_mask_prob", "offload_miss_policy_default",
     "use_hybrid_attention", "kv_latent_dim", "rope_dim", "rope_layout", "sliding_window_size",
     "kv_cache_dtype", "kv_quant_group_size", "mla_decode", "mla_cache_rope", "router_scope",