From 0d4de47c0adcf8432f591028592cc2667defcbae Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Sat, 7 Nov 2026 10:19:02 +0800 Subject: [PATCH 01/18] 0.3.1a0: cluster-locality default-on, architecture-adaptive expert paging - Cluster-locality training is on by default at the entry points (train_chronos.py, planner/presets, bundle finetune), lambda 0.3 from a measured synthetic sweep; the window follows the median SFT record length. The ChronosConfig class default stays 0.0 so stored recipes resume unchanged, and old checkpoints load as lambda 0. - Lazy-expert inference adapts to the checkpoint: models trained with the constraint are paged one training cluster per SSD read, others one expert per read (paging auto|cluster|expert). Layouts get separate cache dirs. - chronos/lazy_experts.py: shared lazy-expert implementation for the Web UI, bundle and server (LazyOptions, page_layout, build_mlx_lazy_engine). - Next-use residency planner for MLX lazy experts, with decay and prefetch wait parameters exposed in the UI, bundle CLI and generate(). - chronos bundle translate/eval --lazy-experts; bundle finetune --[no-]cluster-locality --lambda-cluster --cluster-window --clusters. - chronos serve --paging; /api/ps reports the resolved layout. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 41 + chronos/__init__.py | 2 +- chronos/bundle/cli.py | 79 +- chronos/bundle/docs.py | 23 + chronos/bundle/finetune.py | 12 +- chronos/bundle/recommend.py | 96 +- chronos/bundle/runtime.py | 78 +- chronos/config_schema.py | 1 + chronos/expert_clusters.py | 121 +++ chronos/io/expert_store.py | 38 +- chronos/io/storage.py | 48 +- chronos/lazy_experts.py | 387 ++++++++ chronos/mlx/expert_store.py | 113 ++- chronos/mlx/inference.py | 225 ++++- chronos/mlx/moe.py | 2 + chronos/mlx/training/trainer.py | 39 + chronos/model/checkpoint.py | 9 + chronos/model/config.py | 13 + chronos/model/planner.py | 15 +- chronos/model/temporal_loss.py | 76 ++ chronos/residency_planner.py | 154 +++ chronos/runtime/inference_engine.py | 8 +- chronos/serving/ollama_server.py | 10 +- chronos/serving/runtime.py | 39 +- chronos/trainer/loss_mixin.py | 15 + chronos/trainer/stage_utils.py | 6 + pyproject.toml | 2 +- ...nch_residency_50m_s13996_flores16_cpu.json | 911 ++++++++++++++++++ .../cluster_locality_ab_tiny.json | 226 +++++ .../sweep/lam_0.03_w32.json | 226 +++++ .../residency_20261107/sweep/lam_0.03_w32.log | 32 + .../residency_20261107/sweep/lam_0.1_w32.json | 226 +++++ .../residency_20261107/sweep/lam_0.1_w32.log | 32 + .../residency_20261107/sweep/lam_0.3_w0.json | 227 +++++ .../residency_20261107/sweep/lam_0.3_w0.log | 32 + .../residency_20261107/sweep/lam_0.3_w16.json | 227 +++++ .../residency_20261107/sweep/lam_0.3_w16.log | 32 + .../residency_20261107/sweep/lam_0.3_w32.json | 226 +++++ .../residency_20261107/sweep/lam_0.3_w32.log | 32 + .../residency_20261107/sweep/lam_0.3_w64.json | 227 +++++ .../residency_20261107/sweep/lam_0.3_w64.log | 32 + .../residency_20261107/sweep/lam_1.0_w32.json | 226 +++++ .../residency_20261107/sweep/lam_1.0_w32.log | 32 + tests/test_bundle.py | 31 + tests/test_lazy_experts.py | 285 ++++++ tests/test_residency_cluster.py | 265 +++++ tests/test_smoke.py | 5 +- tools/bench_cluster_locality.py | 277 ++++++ tools/bench_residency.py | 163 ++++ train_chronos.py | 6 + train_chronos_distill.py | 2 + train_chronos_dpo.py | 2 + train_chronos_grpo.py | 2 + train_chronos_orpo.py | 2 + train_chronos_sft.py | 3 + ui/i18n.py | 32 + ui/presets.py | 7 + ui/presets_v3.json | 1 + ui/tabs/config_tab.py | 13 + ui/tabs/inference_tab.py | 458 +++------ 60 files changed, 5703 insertions(+), 449 deletions(-) create mode 100644 chronos/expert_clusters.py create mode 100644 chronos/lazy_experts.py create mode 100644 chronos/residency_planner.py create mode 100644 reports/residency_20261107/bench_residency_50m_s13996_flores16_cpu.json create mode 100644 reports/residency_20261107/cluster_locality_ab_tiny.json create mode 100644 reports/residency_20261107/sweep/lam_0.03_w32.json create mode 100644 reports/residency_20261107/sweep/lam_0.03_w32.log create mode 100644 reports/residency_20261107/sweep/lam_0.1_w32.json create mode 100644 reports/residency_20261107/sweep/lam_0.1_w32.log create mode 100644 reports/residency_20261107/sweep/lam_0.3_w0.json create mode 100644 reports/residency_20261107/sweep/lam_0.3_w0.log create mode 100644 reports/residency_20261107/sweep/lam_0.3_w16.json create mode 100644 reports/residency_20261107/sweep/lam_0.3_w16.log create mode 100644 reports/residency_20261107/sweep/lam_0.3_w32.json create mode 100644 reports/residency_20261107/sweep/lam_0.3_w32.log create mode 100644 reports/residency_20261107/sweep/lam_0.3_w64.json create mode 100644 reports/residency_20261107/sweep/lam_0.3_w64.log create mode 100644 reports/residency_20261107/sweep/lam_1.0_w32.json create mode 100644 reports/residency_20261107/sweep/lam_1.0_w32.log create mode 100644 tests/test_lazy_experts.py create mode 100644 tests/test_residency_cluster.py create mode 100644 tools/bench_cluster_locality.py create mode 100644 tools/bench_residency.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 44eed6f..3de03a2 100644 --- a/CHANGELOG.md +++ b/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_.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 diff --git a/chronos/__init__.py b/chronos/__init__.py index 09bf525..8fdf6b3 100644 --- a/chronos/__init__.py +++ b/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 diff --git a/chronos/bundle/cli.py b/chronos/bundle/cli.py index 8fc317a..5450c7a 100644 --- a/chronos/bundle/cli.py +++ b/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}") diff --git a/chronos/bundle/docs.py b/chronos/bundle/docs.py index 3c23c07..d6db830 100644 --- a/chronos/bundle/docs.py +++ b/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), "", diff --git a/chronos/bundle/finetune.py b/chronos/bundle/finetune.py index 3280ea2..38278d5 100644 --- a/chronos/bundle/finetune.py +++ b/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 ``/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 ``/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") diff --git a/chronos/bundle/recommend.py b/chronos/bundle/recommend.py index cb97842..fb6401a 100644 --- a/chronos/bundle/recommend.py +++ b/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 ``/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({ diff --git a/chronos/bundle/runtime.py b/chronos/bundle/runtime.py index 818d042..51b601f 100644 --- a/chronos/bundle/runtime.py +++ b/chronos/bundle/runtime.py @@ -7,6 +7,11 @@ 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 @@ 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 @@ def __init__(self, bundle_dir: str, *, device: str = "cpu", dtype: str = "float3 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 @@ def __init__(self, bundle_dir: str, *, device: str = "cpu", dtype: str = "float3 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 @@ def greedy(self, prompt: list[int], limit: int, repetition_penalty: float = 1.0, 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 @@ def _greedy_mlx(self, prompt, limit, repetition_penalty): 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 @@ def translate(self, sources: list[str], direction: str = "en_zh", *, max_new_tok 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 diff --git a/chronos/config_schema.py b/chronos/config_schema.py index e399c0e..33d11f2 100644 --- a/chronos/config_schema.py +++ b/chronos/config_schema.py @@ -37,6 +37,7 @@ "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", diff --git a/chronos/expert_clusters.py b/chronos/expert_clusters.py new file mode 100644 index 0000000..916a82e --- /dev/null +++ b/chronos/expert_clusters.py @@ -0,0 +1,121 @@ +"""Expert-cluster partition shared by training and lazy inference (no torch / mlx imports). + +A cluster is a fixed set of expert ids that is stored and paged together +(one ``cluster_.ctsr`` file in the MLX store). The training-time +cluster-locality constraint and the inference-time cluster prefetch use the +same partition, so routing that training concentrates on one cluster is the +routing that one SSD read serves at inference. +""" +from __future__ import annotations + +import math +from typing import List, Optional, Sequence + +# Default strength of the training constraint (on by default). Measured on +# the synthetic multi-topic A/B (tools/bench_cluster_locality.py, 300 steps, +# reports/residency_20261107/sweep/): λ = 0.03 / 0.1 / 0.3 / 1.0 cut lazy +# cluster reads per token from 1.48 to 0.51 / 0.15 / 0.023 / 0.026, with +# held-out CE within +0.3–0.6 % of λ = 0 at every value. 0.3 is the smallest +# value that reaches the plateau. +DEFAULT_LAMBDA = 0.3 +# Window used when the data gives no record length (tokens). The same sweep +# put the best window at the stream segment length (16: 0.044, 32: 0.023, +# 64: 0.11, whole sequence: 0.11 reads per token at a 32-token segment). +DEFAULT_WINDOW = 32 +PAGING_MODES = ("auto", "cluster", "expert") + + +def resolve_cluster_count(config) -> int: + """Number of clusters: ``cluster_locality_clusters`` or ``⌈E / (2k)⌉``, at least 1.""" + e = int(getattr(config, "num_experts", 0) or 0) + if e <= 0: + return 1 + n = int(getattr(config, "cluster_locality_clusters", 0) or 0) + if n <= 0: + k = max(1, int(getattr(config, "num_experts_per_tok", 1) or 1)) + n = -(-e // (2 * k)) + return max(1, min(n, e)) + + +def contiguous_clusters(num_experts: int, num_clusters: int) -> List[List[int]]: + """Split ``0..E-1`` into ``num_clusters`` contiguous, near-equal groups.""" + e, c = int(num_experts), max(1, min(int(num_clusters), int(num_experts))) + base, extra = divmod(e, c) + out, start = [], 0 + for i in range(c): + size = base + (1 if i < extra else 0) + out.append(list(range(start, start + size))) + start += size + return out + + +def expert_to_cluster(clusters: Sequence[Sequence[int]], num_experts: int) -> List[int]: + """``assign[e]`` = cluster id of expert ``e``; every expert must appear exactly once.""" + assign = [-1] * int(num_experts) + for cid, members in enumerate(clusters): + for eid in members: + eid = int(eid) + if not 0 <= eid < int(num_experts) or assign[eid] != -1: + raise ValueError(f"expert {eid} is out of range or in two clusters") + assign[eid] = cid + if -1 in assign: + raise ValueError(f"experts {[i for i, c in enumerate(assign) if c == -1]} have no cluster") + return assign + + +def training_clusters(config) -> List[List[int]]: + """The partition the training constraint uses (and that inference should page by).""" + e = int(getattr(config, "num_experts", 0) or 0) + return contiguous_clusters(e, resolve_cluster_count(config)) + + +def cluster_membership(clusters: Sequence[Sequence[int]], num_experts: int) -> List[List[float]]: + """Row-per-expert one-hot ``[E, C]`` (nested lists; callers wrap it as a tensor).""" + assign = expert_to_cluster(clusters, num_experts) + c = len(clusters) + return [[1.0 if assign[e] == j else 0.0 for j in range(c)] for e in range(int(num_experts))] + + +def trained_with_clusters(config) -> bool: + """True when the weights were trained with the cluster-locality constraint (λ > 0).""" + try: + return float(getattr(config, "lambda_cluster_locality", 0.0) or 0.0) > 0.0 + except (TypeError, ValueError): + return False + + +def inference_clusters(config, paging: str = "auto") -> Optional[List[List[int]]]: + """Page layout for lazy inference, adapted to how the model was trained. + + ``paging``: ``"auto"`` pages by the training partition when the weights + were trained with the constraint and per expert otherwise; ``"cluster"`` + forces the training partition; ``"expert"`` forces one page per expert. + Returns the cluster list (singletons for per-expert paging). + """ + if paging not in PAGING_MODES: + raise ValueError(f"paging must be one of {PAGING_MODES}, not {paging!r}") + e = int(getattr(config, "num_experts", 0) or 0) + if e <= 0: + return [] + if paging == "cluster" or (paging == "auto" and trained_with_clusters(config)): + return training_clusters(config) + return [[i] for i in range(e)] + + +def paging_tag(clusters: Sequence[Sequence[int]]) -> str: + """Short cache-directory suffix that identifies a page layout.""" + sizes = sorted({len(c) for c in clusters}) + if sizes == [1]: + return "e" + return f"c{len(clusters)}" + + +def recommended_window(median_record_tokens: Optional[float], seq_len: int) -> int: + """Cluster-locality window for a dataset: the power of two nearest the median + record length (one record ≈ one stream segment), clamped to ``[16, seq_len]``. + Falls back to ``DEFAULT_WINDOW`` when no record length is known.""" + seq_len = max(1, int(seq_len)) + if not median_record_tokens or median_record_tokens <= 0: + return min(DEFAULT_WINDOW, seq_len) + w = 1 << int(round(math.log2(float(median_record_tokens)))) + return int(max(min(16, seq_len), min(w, seq_len))) diff --git a/chronos/io/expert_store.py b/chronos/io/expert_store.py index 6145604..cb38b3a 100644 --- a/chronos/io/expert_store.py +++ b/chronos/io/expert_store.py @@ -353,7 +353,15 @@ def offload_from_state_dict( This is the real lazy-load path for CPU/MPS inference: we avoid materializing MoE expert weights into the live model before setup. """ - if self.storage_format == "safetensors" and clusters is None: + if clusters is not None: + # Paging by the training clusters (lambda_cluster_locality > 0): + # one .ctsr per cluster, so one SSD read brings a stream's experts. + self._write_clusters_from( + lambda li, ei: self._expert_state_from_checkpoint(state_dict, li, ei), clusters) + if replace_live: + self.replace_live_experts_with_placeholders() + return + if self.storage_format == "safetensors": # A raw .pth checkpoint has no co-activation cluster layout. # Falling back to one .pt shard per (layer, expert) preserves # true lazy loading; writing a single all-expert .ctsr cluster @@ -383,11 +391,20 @@ def offload_from_state_reader( materializing the full expert state dict in DRAM: each expert tensor is pulled from the export artifact only while that expert shard is written. """ - if self.storage_format == "safetensors" and clusters is None: + keys = list(reader.keys()) + if clusters is not None: + def from_reader(li, ei): + prefix = f"model.layers.{li}.mlp.experts.{ei}." + return {k[len(prefix):]: reader.get_tensor(k).detach().cpu().contiguous() + for k in keys if k.startswith(prefix)} + self._write_clusters_from(from_reader, clusters) + if replace_live: + self.replace_live_experts_with_placeholders() + return + if self.storage_format == "safetensors": self.storage_format = "pt" self._cluster_storage = None - keys = list(reader.keys()) for li in range(self.num_layers): for ei in range(self.num_experts): path = self._ssd_path(ei, li) @@ -407,6 +424,21 @@ def offload_from_state_reader( if replace_live: self.replace_live_experts_with_placeholders() + def _write_clusters_from(self, get_state, clusters: List[List[int]]): + """Shard experts into one ``.ctsr`` per cluster from ``get_state(layer, expert)``.""" + dtype = self._expert_dtype_cached + + def cast(li, ei): + return {k: (v.to(dtype=dtype) if v.is_floating_point() and v.dtype != dtype else v) + for k, v in get_state(li, ei).items()} + + ClusterStorage.write_clusters_from_states( + cast, clusters=clusters, output_dir=self.ssd_dir, num_layers=self.num_layers, + num_experts=self.num_experts, layer_to_group=self._layer_to_group if self.layer_aware else None, + ) + self.storage_format = "safetensors" + self._cluster_storage = ClusterStorage(self.ssd_dir) + def replace_live_experts_with_placeholders(self): """Drop live expert modules from the active model after SSD offload.""" for li in range(self.num_layers): diff --git a/chronos/io/storage.py b/chronos/io/storage.py index 166b99e..90dc988 100644 --- a/chronos/io/storage.py +++ b/chronos/io/storage.py @@ -225,6 +225,41 @@ def write_clusters( trunk compute dtype: BF16 stays BF16 (no FP16 detour that would clip BF16's range), FP32 stays FP32. Pass a dtype only to convert once. + Returns the path to cluster_manifest.json. + """ + def get_state(li: int, eid: int) -> Dict[str, torch.Tensor]: + if li >= len(moe_layers): + return {} + out = {} + for pname, pval in moe_layers[li].experts[eid].state_dict().items(): + t = pval.detach() + if dtype is not None and t.dtype != dtype: + t = t.to(dtype=dtype) + out[pname] = t + return out + + return ClusterStorage.write_clusters_from_states( + get_state, clusters=clusters, output_dir=output_dir, num_layers=num_layers, + num_experts=num_experts, layer_to_group=layer_to_group, + ) + + @staticmethod + def write_clusters_from_states( + get_state, + clusters: List[List[int]], + output_dir: str, + num_layers: int, + num_experts: int, + layer_to_group: Optional[Dict[int, int]] = None, + ) -> str: + """Write one safetensors file per cluster from ``get_state(layer, expert)``. + + ``get_state`` returns ``{param_name: Tensor}`` for one expert in one + layer (empty when that layer has no routed experts). Only one + cluster's tensors are held at a time, so a full checkpoint state dict + or an export reader can be sharded without materialising every + expert twice. Every expert must appear in exactly one cluster. + Returns the path to cluster_manifest.json. """ if not _HAS_SAFETENSORS: @@ -232,6 +267,9 @@ def write_clusters( "write_clusters requires the `safetensors` package. " "Install with `pip install safetensors>=0.4.0`." ) + seen = [int(e) for cl in clusters for e in cl] + if sorted(seen) != list(range(int(num_experts))): + raise ValueError("clusters must cover every expert exactly once") os.makedirs(output_dir, exist_ok=True) manifest_clusters: Dict[int, Tuple[str, List[int]]] = {} @@ -242,14 +280,8 @@ def write_clusters( for eid in cluster: expert_to_cluster[int(eid)] = cid for li in range(num_layers): - if li >= len(moe_layers): - continue - expert = moe_layers[li].experts[eid] - for pname, pval in expert.state_dict().items(): - t = pval.detach() - if dtype is not None and t.dtype != dtype: - t = t.to(dtype=dtype) - tensors[_key(li, int(eid), pname)] = t.cpu().contiguous() + for pname, t in (get_state(li, int(eid)) or {}).items(): + tensors[_key(li, int(eid), pname)] = t.detach().cpu().contiguous() file_name = f"cluster_{cid}.ctsr" file_path = os.path.join(output_dir, file_name) diff --git a/chronos/lazy_experts.py b/chronos/lazy_experts.py new file mode 100644 index 0000000..632d831 --- /dev/null +++ b/chronos/lazy_experts.py @@ -0,0 +1,387 @@ +"""SSD-lazy routed experts: one implementation for the Web UI, bundle and server. + +The page layout adapts to how the checkpoint was trained. Weights trained +with the cluster-locality constraint (``lambda_cluster_locality > 0``) keep a +stream inside one training cluster, so they are paged one cluster file per +read; other weights are paged one expert per read, which never loads +experts the stream does not use (``chronos.expert_clusters.inference_clusters``). +The layout is part of the cache-directory name, so a cache written for one +layout is never reused for the other. + +Residency (what stays loaded) is independent of the layout: ``"planned"`` +evicts the page needed latest by the next-use planner, ``"window"`` keeps the +LRU order with the lookahead window protected. Neither changes the output. +""" +from __future__ import annotations + +import gc +import hashlib +import json +import math +import os +from dataclasses import asdict, dataclass + +from chronos.expert_clusters import PAGING_MODES, inference_clusters, paging_tag, trained_with_clusters + +RESIDENCY_MODES = ("planned", "window") +# Defaults shared by every entry point (the engine's own defaults). +DEFAULT_RESIDENCY = "planned" +DEFAULT_RESIDENCY_DECAY = 0.7 +DEFAULT_DECODE_WAIT_S = 0.012 +DEFAULT_VERIFY_WAIT_S = 0.02 +DEFAULT_RAM_RATIO = 1.0 + + +@dataclass +class LazyOptions: + """User-facing lazy-expert settings; every field has a working default.""" + + paging: str = "auto" # auto | cluster | expert + residency: str = DEFAULT_RESIDENCY # planned | window + residency_decay: float = DEFAULT_RESIDENCY_DECAY + decode_wait_s: float = DEFAULT_DECODE_WAIT_S + verify_wait_s: float = DEFAULT_VERIFY_WAIT_S + ram_ratio: float = DEFAULT_RAM_RATIO + + def validated(self) -> "LazyOptions": + if self.paging not in PAGING_MODES: + raise ValueError(f"paging must be one of {PAGING_MODES}, not {self.paging!r}") + if self.residency not in RESIDENCY_MODES: + raise ValueError(f"residency must be one of {RESIDENCY_MODES}, not {self.residency!r}") + decay = float(self.residency_decay) + if not 0.0 <= decay < 1.0: + raise ValueError("residency_decay must be in [0, 1)") + waits = float(self.decode_wait_s), float(self.verify_wait_s) + if min(waits) < 0: + raise ValueError("decode_wait_s and verify_wait_s must be >= 0") + ratio = float(self.ram_ratio) + if not ratio > 0: + raise ValueError("ram_ratio must be > 0") + return LazyOptions(self.paging, self.residency, decay, waits[0], waits[1], ratio) + + def generate_kwargs(self) -> dict: + """Keyword arguments for ``ChronosMLXInferenceEngine.generate``.""" + return {"residency": self.residency, "residency_decay": self.residency_decay, + "decode_wait_s": self.decode_wait_s, "verify_wait_s": self.verify_wait_s} + + def as_dict(self) -> dict: + return asdict(self) + + +def page_layout(config, paging: str = "auto") -> dict: + """Resolved layout for a checkpoint: clusters, tag and the reason (for UI/stats).""" + clusters = inference_clusters(config, paging) or [] + tag = paging_tag(clusters) if clusters else "e" + trained = trained_with_clusters(config) + if paging == "auto": + reason = "trained with cluster locality" if trained else "trained without cluster locality" + else: + reason = f"forced {paging}" + return {"paging": paging, "resolved": "expert" if tag == "e" else "cluster", "clusters": clusters, + "tag": tag, "num_pages": len(clusters), "trained_with_clusters": trained, "reason": reason} + + +# ── cache location ──────────────────────────────────────────────── + +def checkpoint_cache_dir(model_path: str, config, *, backend: str = "torch", tag: str = "e", + root: str | None = None) -> str: + """Per-checkpoint expert cache directory. + + The key is path + size + mtime + core topology, so a changed checkpoint + gets a fresh namespace. ``tag`` is the page layout (``paging_tag``): + per-expert caches keep their historical names (``ckpt_`` for Torch, + ``ckpt__mlx`` for MLX); cluster layouts add ``_``. + """ + base = root or os.path.join(os.getcwd(), "expert_cache") + if not model_path or not os.path.exists(model_path): + name_dir = base + else: + try: + st = os.stat(model_path) + except OSError: + name_dir = base + else: + topology = tuple(getattr(config, k, "") for k in ( + "vocab_size", "hidden_size", "num_hidden_layers", "num_experts", "moe_intermediate_size")) + raw = f"{os.path.abspath(model_path)}:{st.st_size}:{st.st_mtime_ns}:{topology}" + name_dir = os.path.join(base, "ckpt_" + hashlib.sha1(raw.encode("utf-8")).hexdigest()[:16]) + if backend == "mlx": + name_dir += "_mlx" + if tag and tag != "e": + name_dir += f"_{tag}" + return name_dir + + +# ── expert budget ───────────────────────────────────────────────── + +def active_expert_count(config) -> int: + """Routed working set per token: observed/recommended size, else top-k × layers.""" + num_experts = int(getattr(config, "num_experts", 1) or 1) + for attr in ("observed_active_experts", "offload_working_set_experts", "recommended_working_set_experts"): + value = getattr(config, attr, None) + if value in (None, "", 0): + continue + try: + return max(1, min(num_experts, int(value))) + except (TypeError, ValueError): + pass + top_k = int(getattr(config, "num_experts_per_tok", 1) or 1) + num_layers = int(getattr(config, "num_hidden_layers", 1) or 1) + # top_k is per layer; the cache is keyed by global expert id, so one token + # can touch up to top_k × layers distinct ids before repeats. + return max(1, min(num_experts, top_k * max(1, num_layers))) + + +def expert_budget(config, ram_ratio: float | str | None = 1.0) -> dict: + """Expert-page budget from the routed working set and the user ratio. + + The ratio produces real lower-budget runs for sweeps; output quality is + preserved by on-demand loading, never by clamping to top-k. + """ + active = active_expert_count(config) + num_experts = int(getattr(config, "num_experts", active) or active) + try: + ratio = float(ram_ratio) + except (TypeError, ValueError): + ratio = 1.0 + ratio = max(0.01, ratio) + routing_top_k = int(getattr(config, "num_experts_per_tok", active) or active) + hard_cap = max(active, int(math.floor(active * 1.25))) + requested = max(1, int(math.ceil(active * ratio))) + effective = max(1, min(num_experts, requested, hard_cap)) + return { + "ideal_active_experts": active, + "routing_top_k": routing_top_k, + "num_moe_layers": int(getattr(config, "num_hidden_layers", 1) or 1), + "requested_ram_load_ratio": round(ratio, 3), + "requested_expert_budget": requested, + "max_allowed_expert_budget": min(num_experts, hard_cap), + "effective_expert_budget": effective, + "effective_vram_expert_budget": effective, + "effective_ram_expert_budget": effective, + "min_vram_expert_budget": 1, + "budget_clamped": effective != requested, + } + + +def apply_torch_budget(engine, config, ram_ratio: float | str | None = 1.0) -> dict: + """Set the Torch store's VRAM/RAM page capacity; a hard budget always wins.""" + budget = expert_budget(config, ram_ratio) + vram_cap = int(budget["effective_vram_expert_budget"]) + ram_cap = int(budget["effective_ram_expert_budget"]) + store = getattr(getattr(engine, "cache_manager", None), "expert_store", None) + if store is None: + return budget + hard_vram = getattr(store, "max_resident_pages", None) + hard_ram = getattr(store, "ram_budget_pages", None) + if hard_vram is not None: + vram_cap = max(1, min(vram_cap, int(hard_vram))) + if hard_ram is not None: + ram_cap = max(1, min(ram_cap, int(hard_ram))) + budget["effective_vram_expert_budget"], budget["effective_ram_expert_budget"] = vram_cap, ram_cap + store.vram_capacity = vram_cap + store.ram_capacity = ram_cap + store.vram_lru.capacity = vram_cap + store.ram_lru.capacity = ram_cap + return budget + + +# ── live-expert placeholders ────────────────────────────────────── + +def torch_experts_to_placeholders(model, config) -> None: + """Drop the constructor's random experts before ``model.to(device)``.""" + from chronos.model.moe_chronos import ChronosMOEFeedForward, LazyExpertPlaceholder + from chronos.precision import torch_model_dtype + + hidden = int(getattr(config, "hidden_size", 0) or 0) + intermediate = int(getattr(config, "moe_intermediate_size", 0) or 0) + if not hidden or not intermediate: + return + for layer in getattr(getattr(model, "model", None), "layers", []): + moe = getattr(layer, "mlp", None) + if not isinstance(moe, ChronosMOEFeedForward): + continue + for idx in range(len(moe.experts)): + # Placeholder dtype follows the trunk; place_torch_model() updates it on move. + moe.experts[idx] = LazyExpertPlaceholder(hidden, intermediate, torch_model_dtype(model)) + + +def mlx_experts_to_placeholders(model) -> None: + import mlx.core as mx + from chronos.mlx.moe import ChronosMLXMOE, LazyFeedForwardMLX + + for layer in getattr(model, "layers", []): + moe = getattr(layer, "mlp", None) + if not isinstance(moe, ChronosMLXMOE): + continue + for idx in range(len(moe.experts)): + moe.experts[idx] = LazyFeedForwardMLX() + gc.collect() + try: + mx.clear_cache() + except Exception: + pass + + +# ── MLX expert cache ────────────────────────────────────────────── + +def _store_torch_dtype(store): + """Torch dtype of the MLX store's pages (= the MLX trunk compute dtype).""" + import torch + name = str(getattr(store, "page_dtype", "float32")).rsplit(".", 1)[-1] + return getattr(torch, name) + + +def write_mlx_cache(store, get_tensor, keys, clusters=None) -> None: + """Write the MLX store's cache: one ``cluster_.ctsr`` per page + manifest. + + ``get_tensor(key)`` returns a checkpoint tensor (``model.layers.L.mlp.experts.E.P``); + ``keys`` is the set of checkpoint keys. ``clusters`` defaults to one page per expert. + Only one page's tensors are held at a time. + """ + from safetensors.torch import save_file # BF16-capable + from chronos.io.storage import ClusterManifest, KEY_SEP, MANIFEST_FILENAME + + num_experts, num_layers = int(store.num_experts), int(store.num_layers) + if clusters is None: + clusters = [[i] for i in range(num_experts)] + os.makedirs(store.ssd_dir, exist_ok=True) + page_dtype = _store_torch_dtype(store) + keys = set(keys) + manifest_clusters, expert_to_cluster = {}, {} + for cid, members in enumerate(clusters): + tensors = {} + for eid in members: + for li in range(num_layers): + for pname in store._PARAM_NAMES: + key = f"model.layers.{li}.mlp.experts.{int(eid)}.{pname}" + if key in keys: + tensors[f"l{li}_e{int(eid)}{KEY_SEP}{pname}"] = ( + get_tensor(key).detach().to(dtype=page_dtype, device="cpu").contiguous()) + if not tensors: + continue + file_name = f"cluster_{cid}.ctsr" + save_file(tensors, os.path.join(store.ssd_dir, file_name)) + manifest_clusters[cid] = (file_name, [int(e) for e in members]) + for eid in members: + expert_to_cluster[int(eid)] = cid + if len(expert_to_cluster) != num_experts: + raise RuntimeError(f"MLX lazy cache is incomplete: wrote {len(expert_to_cluster)}/{num_experts} experts") + manifest = ClusterManifest(version=1, num_experts=num_experts, num_layers=num_layers, + storage_format="safetensors", clusters=manifest_clusters, + expert_to_cluster=expert_to_cluster) + with open(os.path.join(store.ssd_dir, MANIFEST_FILENAME), "w", encoding="utf-8") as f: + json.dump(manifest.to_dict(), f, indent=2) + + +def write_mlx_cache_from_state_dict(store, state_dict: dict, clusters=None) -> None: + write_mlx_cache(store, state_dict.__getitem__, state_dict.keys(), clusters) + + +def write_mlx_cache_from_reader(store, reader, clusters=None) -> None: + write_mlx_cache(store, reader.get_tensor, reader.keys(), clusters) + + +def _manifest_matches(store, clusters) -> bool: + manifest = getattr(store, "_cluster_manifest", None) + if manifest is None: + return False + have = sorted(sorted(int(e) for e in members) for _f, members in manifest.clusters.values()) + return have == sorted(sorted(int(e) for e in c) for c in clusters) + + +def ensure_mlx_cache(store, cache_dir: str, *, reader=None, weights=None, clusters=None) -> bool: + """Attach a complete cache with this layout at ``cache_dir`` or rebuild it. + + Returns True when an existing cache was reused. + """ + import shutil + + if clusters is None: + clusters = [[i] for i in range(int(store.num_experts))] + if store.attach_cluster_manifest(cache_dir) and _manifest_matches(store, clusters): + return True + if os.path.exists(cache_dir): + shutil.rmtree(cache_dir) + os.makedirs(cache_dir, exist_ok=True) + store.ssd_dir = cache_dir + store._cluster_manifest = None + store._loaded_clusters.clear() + store._warm.clear() + store._warm_lru.clear() + if reader is not None: + write_mlx_cache_from_reader(store, reader, clusters) + elif weights is not None: + write_mlx_cache_from_state_dict(store, weights, clusters) + else: + store.offload_all_to_ssd(clusters=clusters) + if not (store.attach_cluster_manifest(cache_dir) and _manifest_matches(store, clusters)): + raise RuntimeError(f"failed to build a complete MLX lazy expert cache at {cache_dir}") + return False + + +def build_mlx_lazy_engine(model_path: str, config, options: LazyOptions | None = None, *, + cache_root: str | None = None, dtype: str = "auto"): + """Load a checkpoint or export into an MLX engine with SSD-lazy experts. + + Only the trunk is materialised; experts are written to the page cache + (layout from ``options.paging`` and the checkpoint) and stream in on demand. + Sets ``config.recommended_resident_experts`` / ``recommended_ram_experts`` + from the budget. Returns ``(engine, info)``; the caller owns ``engine.stop()``. + """ + from chronos.export import is_export_artifact, open_export_reader + from chronos.mlx.inference import ChronosMLXInferenceEngine + from chronos.mlx.model import ChronosMLXModel + from chronos.model.checkpoint import load_checkpoint_state_dict, load_state_dict_controlled + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.precision import resolve_mlx + + opts = (options or LazyOptions()).validated() + layout = page_layout(config, opts.paging) + budget = expert_budget(config, opts.ram_ratio) + config.recommended_resident_experts = int(budget["effective_vram_expert_budget"]) + config.recommended_ram_experts = int(budget["effective_ram_expert_budget"]) + cache_dir = checkpoint_cache_dir(model_path, config, backend="mlx", tag=layout["tag"], root=cache_root) + + reader = weights = None + if model_path and is_export_artifact(model_path): + reader = open_export_reader(model_path) + base_only = reader.load_state_dict(include_experts=False) + elif model_path and os.path.exists(model_path): + weights = load_checkpoint_state_dict(model_path, map_location="cpu") + base_only = {k: v for k, v in weights.items() if ".mlp.experts." not in k} + else: + raise FileNotFoundError("lazy-expert inference needs a checkpoint or export artifact") + engine = None + try: + pt_model = ChronosForCausalLM(config) + if base_only: + load_state_dict_controlled(pt_model, base_only, allow_missing_substrings=(".mlp.experts.",)) + torch_experts_to_placeholders(pt_model, config) + precision = resolve_mlx(dtype) + mlx_model = ChronosMLXModel.from_chronos_pytorch(pt_model.eval(), config, include_experts=False, + dtype=precision.compute) + mlx_experts_to_placeholders(mlx_model) + del pt_model, base_only + gc.collect() + engine = ChronosMLXInferenceEngine(mlx_model, config, ssd_dir=cache_dir) + reused = ensure_mlx_cache(engine.store, cache_dir, reader=reader, weights=weights, + clusters=layout["clusters"]) + except BaseException: + if engine is not None: + engine.stop() + raise + finally: + reader = weights = None + gc.collect() + info = {**budget, "cache_dir": cache_dir, "cache_reused": reused, + "paging": layout["paging"], "paging_resolved": layout["resolved"], "paging_reason": layout["reason"], + "num_pages": layout["num_pages"], "trained_with_clusters": layout["trained_with_clusters"], + "lazy_options": opts.as_dict(), "precision": precision.as_dict()} + return engine, info + + +def warm_ids(config, budget: dict) -> list[int]: + """Experts warmed before the first token (the execution budget, lowest ids first).""" + count = int(budget["effective_vram_expert_budget"]) + return list(range(min(int(getattr(config, "num_experts", 0) or 0), count))) diff --git a/chronos/mlx/expert_store.py b/chronos/mlx/expert_store.py index 0afc4b8..d2348f6 100644 --- a/chronos/mlx/expert_store.py +++ b/chronos/mlx/expert_store.py @@ -116,6 +116,8 @@ def __init__(self, model, config, ssd_dir: str = "./expert_cache_mlx"): self._cluster_manifest: ClusterManifest | None = None self._loaded_clusters: set[int] = set() self._protected_experts: set[int] = set() + # Optional residency planner (see set_eviction_ranker); None = LRU. + self._eviction_ranker = None self.storage_format = "npy" # Pages live in the trunk's compute dtype (one precision policy end # to end); a cache written in another dtype is converted once per SSD @@ -194,9 +196,9 @@ def _ssd_path(self, expert_id: int, layer_idx: int, param_name: str) -> str: def _load_from_ssd(self, expert_id: int) -> Dict[int, Dict[str, mx.array]]: if self._cluster_manifest is not None: - self._load_cluster_for_expert(expert_id) - with self._lock: - return self._warm.get(expert_id, {}) + # The loaded page is returned directly: another thread may evict it + # from the warm tier as soon as the lock is released. + return self._load_cluster_for_expert(expert_id) or {} layer_states = {} for li in range(self.num_layers): @@ -285,10 +287,10 @@ def _parse_cluster_key(self, key: str): except ValueError: return None - def _load_cluster_for_expert(self, expert_id: int) -> None: + def _load_cluster_for_expert(self, expert_id: int) -> Dict[int, Dict[str, mx.array]] | None: manifest = self._cluster_manifest if manifest is None or expert_id not in manifest.expert_to_cluster: - return + return None cid = manifest.expert_to_cluster[int(expert_id)] with self._lock: if ( @@ -296,7 +298,7 @@ def _load_cluster_for_expert(self, expert_id: int) -> None: and expert_id in self._warm and self._layer_states_complete(self._warm.get(expert_id)) ): - return + return self._warm[expert_id] file_name = manifest.cluster_file(cid) path = os.path.join(self.ssd_dir, file_name) nested: Dict[int, Dict[int, Dict[str, mx.array]]] = {} @@ -315,6 +317,7 @@ def _load_cluster_for_expert(self, expert_id: int) -> None: self._stats["cluster_experts_loaded"] += 1 self._stats["cluster_loads"] += 1 self._loaded_clusters.add(cid) + return nested.get(int(expert_id)) # ── Prefetch (SSD → unified memory) ────────────────────────── @@ -324,18 +327,25 @@ def prefetch_to_ram(self, expert_ids: List[int]): Call from a background thread for async prefetch. """ for eid in expert_ids: - with self._lock: - if eid in self._warm and self._layer_states_complete(self._warm.get(int(eid))): - self._touch_warm_locked(int(eid)) - self._stats["warm_hits"] += 1 - continue - layer_states = self._complete_state(self._load_from_ssd(eid)) - if not layer_states: - self._stats["prefetch_misses"] += 1 - continue - with self._lock: - self._put_warm_locked(int(eid), layer_states, protected_expert_id=int(eid)) - self._stats["prefetch_loads"] += 1 + self._fetch_to_ram(int(eid)) + + def _fetch_to_ram(self, eid: int) -> Dict[int, Dict[str, mx.array]]: + """Warm page of ``eid`` (SSD read if needed); returned even if evicted meanwhile.""" + with self._lock: + hit = self._warm.get(eid) + if self._layer_states_complete(hit): + self._touch_warm_locked(eid) + self._stats["warm_hits"] += 1 + return hit + layer_states = self._complete_state(self._load_from_ssd(eid)) + if not layer_states: + self._stats["prefetch_misses"] += 1 + return {} + with self._lock: + if self._warm.get(eid) is not layer_states: + self._put_warm_locked(eid, layer_states, protected_expert_id=eid) + self._stats["prefetch_loads"] += 1 + return layer_states def _touch_warm_locked(self, expert_id: int): if expert_id in self._warm_lru: @@ -359,7 +369,31 @@ def set_protected_experts(self, expert_ids: List[int] | set[int] | tuple[int, .. if 0 <= int(eid) < int(self.num_experts) } + def set_eviction_ranker(self, ranker=None): + """Plug a residency planner into both tiers' eviction. + + ``ranker(candidates) -> victim`` receives the evictable pages in LRU + order (oldest first; protected pages are already excluded) and + returns the one to drop, e.g. ``NextUsePlanner.victim``. ``None`` + restores plain LRU. Eviction only decides *when* a page is read + again: a selected page is always materialised before it runs. + """ + with self._lock: + self._eviction_ranker = ranker + + def _coldest_locked(self, order, skip) -> int | None: + candidates = [c for c in order if c not in skip] + if not candidates: + return None + ranker = self._eviction_ranker + if ranker is None: + return candidates[0] + victim = ranker(candidates) + return candidates[0] if victim is None else victim + def _evict_warm_locked(self, protected_expert_id: int | None = None): + if self._eviction_ranker is not None: + return self._evict_warm_scored_locked(protected_expert_id) while len(self._warm_lru) > self._warm_capacity: evict_id = None for candidate in self._warm_lru.keys(): @@ -393,6 +427,31 @@ def _evict_warm_locked(self, protected_expert_id: int | None = None): } self._stats["warm_evictions"] += 1 + def _evict_warm_scored_locked(self, protected_expert_id: int | None = None): + """Score-ordered warm eviction (same tiers and hard capacity as the LRU path).""" + keep = {protected_expert_id} if protected_expert_id is not None else set() + while len(self._warm_lru) > self._warm_capacity: + order = list(self._warm_lru.keys()) + guarded = keep | self._protected_experts + evict_id = self._coldest_locked(order, guarded | set(self._hot_lru.keys())) + if evict_id is None: + evict_id = self._coldest_locked(order, guarded) + if evict_id is None: + evict_id = self._coldest_locked(order, keep) + if evict_id is None: + return + evicted_was_hot = evict_id in self._hot_lru + self._warm_lru.pop(evict_id, None) + self._warm.pop(evict_id, None) + if not evicted_was_hot: + self._replace_expert_with_placeholder(evict_id) + self._loaded_clusters = { + cid for cid in self._loaded_clusters + if not self._cluster_evicted_locked(cid) + } + self._stats["warm_evictions"] += 1 + self._stats["planned_evictions"] += 1 + def _cluster_evicted_locked(self, cluster_id: int) -> bool: manifest = self._cluster_manifest if manifest is None: @@ -487,11 +546,7 @@ def promote_to_vram(self, expert_id: int) -> bool: self._touch_warm_locked(expert_id) if not self._layer_states_complete(layer_states): - self.prefetch_to_ram([expert_id]) - with self._lock: - layer_states = self._warm.get(expert_id) - if layer_states is not None: - self._touch_warm_locked(expert_id) + layer_states = self._fetch_to_ram(expert_id) layer_states = self._complete_state(layer_states) if not layer_states: self._stats["promote_misses"] += 1 @@ -525,11 +580,12 @@ def promote_to_vram(self, expert_id: int) -> bool: self._hot_lru.move_to_end(expert_id) else: while len(self._hot_lru) >= self._capacity: - evicted = None - for candidate in self._hot_lru.keys(): - if candidate != expert_id and candidate not in self._protected_experts: - evicted = candidate - break + evicted = self._coldest_locked( + list(self._hot_lru.keys()), {expert_id} | self._protected_experts) + if evicted is None and self._eviction_ranker is not None: + # Every other page is protected: the hard capacity + # still wins, and the lowest-priority one goes. + evicted = self._coldest_locked(list(self._hot_lru.keys()), {expert_id}) if evicted is None: evicted, _ = self._hot_lru.popitem(last=False) if evicted == expert_id: @@ -598,6 +654,7 @@ def stats(self) -> dict: "max_resident_pages": self.max_resident_pages, "ram_budget_pages": self.ram_budget_pages, "protected_pages": len(self._protected_experts), + "eviction_policy": "planned" if self._eviction_ranker is not None else "lru", "resident_experts": len(hot_ids), "hot_experts": len(hot_ids), "hot_capacity": self._capacity, diff --git a/chronos/mlx/inference.py b/chronos/mlx/inference.py index 0c36a02..43f9c2c 100644 --- a/chronos/mlx/inference.py +++ b/chronos/mlx/inference.py @@ -142,6 +142,60 @@ def touch(eid: int) -> None: moe.runtime_on_demand_loader = loader moe.runtime_touch_expert = touch + def _attach_planner(self, planner) -> None: + """Report every lazy-expert run to ``planner`` and let it rank evictions. + + Wraps the hooks of :meth:`_install_runtime_hooks`: before an expert + of dispatch ``(layer, loop iteration)`` is loaded or run, the planner + records it; the dispatch's not-yet-run experts (an MoE runs its + selection in id order) are kept from eviction. The store then evicts + the page needed latest instead of the least recent one. + """ + layer_of = {} + for li, layer in enumerate(self.model.layers): + moe = getattr(layer, "mlp", None) + if not isinstance(moe, ChronosMLXMOE): + continue + layer_of[id(moe)] = li + + def record(eid, moe=moe, li=li): + eid = int(eid) + selected = getattr(moe, "_selected_host", None) or () + pending = [j for j in range(eid + 1, len(selected)) if selected[j]] + planner.access(li, len(moe.iteration_router_probs) - 1, eid, pending) + + base_load, base_touch = moe.runtime_on_demand_loader, moe.runtime_touch_expert + + def load(eid, record=record, base=base_load): + record(eid) + return bool(base(eid)) if base is not None else False + + def touch(eid, record=record, base=base_touch): + record(eid) + if base is not None: + base(eid) + + moe.runtime_on_demand_loader = load + moe.runtime_touch_expert = touch + # Ranking replaces the protected window: a protected set the size of + # the budget would turn every eviction into a forced fallback. + self.store.set_protected_experts([]) + self.store.set_eviction_ranker(planner.victim) + + def _planned_prefetch(self, planner, timeout_s: float) -> int: + """Queue (and briefly await) the pages the next forward needs early; returns the count.""" + hot_cap = max(int(getattr(self.store, "_capacity", 1) or 1), 1) + warm_cap = max(int(getattr(self.store, "_warm_capacity", hot_cap) or hot_cap), 1) + with self.store._lock: + warm = set(self.store._warm.keys()) + hot = self.store.hot_expert_ids() + fetch = planner.prefetch_plan(warm | hot, warm_cap) + promote = planner.prefetch_plan(hot, hot_cap) + if fetch or promote: + self._prefetch_and_promote_window(list(dict.fromkeys(promote + fetch)), promote_ids=promote, + timeout_s=timeout_s) + return len(promote) + def _clear_runtime_hooks(self): for layer in self.model.layers: moe = getattr(layer, "mlp", None) @@ -224,24 +278,37 @@ def _current_virtual_ids(self) -> List[int]: and may pick different experts each time; residency follows the union, so the next step's identical iterations stay hot. """ + values = self._virtual_route_weights(last_only=True) + return [i for i, _v in sorted(((i, v) for i, v in enumerate(values) if v > 0), key=lambda t: (-t[1], t[0]))] + + def _virtual_route_weights(self, last_only: bool = True) -> List[float]: + """Per-expert top-k routing weight summed over every dispatch of the last forward. + + ``last_only`` scores the last position (one decode or draft token); + otherwise every position counts and the sum is divided by the number + of positions, so a verify or prefill forward weighs like one token. + """ from chronos.mlx.moe import ChronosMLXMOE e = int(self.config.num_experts) k = max(1, min(int(self.config.num_experts_per_tok), e)) - rows = [] + rows, positions = [], 1 for layer in self.model.layers: moe = getattr(layer, "mlp", None) if isinstance(moe, ChronosMLXMOE): - rows.extend(p[:, -1, :].reshape(-1, e) for p in (moe.iteration_router_probs or ())) + for p in (moe.iteration_router_probs or ()): + if not last_only: + positions = max(positions, int(p.shape[1])) + rows.append((p[:, -1, :] if last_only else p).reshape(-1, e)) if not rows: - return [] + return [0.0] * e probs = mx.concatenate(rows, axis=0).astype(mx.float32) ids = mx.argpartition(-probs, kth=k - 1, axis=-1)[:, :k] weights = mx.take_along_axis(probs, ids, axis=-1) scores = ((ids[..., None] == mx.arange(e)).astype(mx.float32) * weights[..., None]).sum(axis=(0, 1)) - values = scores.tolist() + values = [float(v) / positions for v in scores.tolist()] with self._stats_lock: self._runtime_stats["virtual_dispatches"] = self._runtime_stats.get("virtual_dispatches", 0) + len(rows) - return [i for i, _v in sorted(((i, v) for i, v in enumerate(values) if v > 0), key=lambda t: (-t[1], t[0]))] + return values def _protect_prediction_window(self, expert_ids: List[int]) -> None: capacity = max( @@ -281,6 +348,11 @@ def generate( top_k: int = 0, seed: Optional[int] = None, eos_token_id=None, + residency: str = "window", # lazy experts: "window" (LRU + protected window) | "planned" + residency_decay: float = 0.7, + decode_wait_s: float = 0.012, # planned: wait for prefetched pages before a decode step + verify_wait_s: float = 0.02, # planned + speculative: wait before the verify forward + repetition_penalty: float = 1.0, ) -> Iterator[int]: """ Yields token IDs one at a time. @@ -293,13 +365,51 @@ def generate( verifies them in one full-depth forward (``chronos.mlx.speculative``); greedy output equals plain greedy decoding. ``eos_token_id`` (int or list) stops generation after that token is yielded. + + ``residency="planned"`` replaces LRU eviction of lazy expert pages + with next-use planning (``chronos.residency_planner``): the previous + forward of the same kind (and, when speculative, the draft queue) + predicts the order in which pages are needed, the page needed latest + is evicted, and pages needed early are prefetched. Output is + unchanged; residency only decides what is loaded when. + ``decode_wait_s`` / ``verify_wait_s`` bound how long a step waits for + its prefetched pages before running (a page still missing is loaded + on demand, so the wait only trades latency for fewer stalls). + + ``repetition_penalty`` (plain decoding only) applies the CTRL rule to + every token already in the prompt or output; speculative decoding + needs 1.0 because the penalty changes the target distribution. """ + if residency not in ("window", "planned"): + raise ValueError("residency must be 'window' or 'planned'") + if not 0.0 <= float(residency_decay) < 1.0: + raise ValueError("residency_decay must be in [0, 1)") + if float(decode_wait_s) < 0 or float(verify_wait_s) < 0: + raise ValueError("decode_wait_s and verify_wait_s must be >= 0") + if float(repetition_penalty) <= 0: + raise ValueError("repetition_penalty must be > 0") + if speculative and float(repetition_penalty) != 1.0: + raise ValueError("speculative decoding needs repetition_penalty=1.0") + self.store.set_eviction_ranker(None) if speculative: yield from self._generate_speculative( input_ids, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, draft_depth=draft_depth, gamma=gamma, gamma_max=gamma_max, adaptive_gamma=adaptive_gamma, - seed=seed, eos_token_id=eos_token_id) + seed=seed, eos_token_id=eos_token_id, residency=residency, residency_decay=residency_decay, + verify_wait_s=verify_wait_s) return + penalty = float(repetition_penalty) + seen = None + if penalty != 1.0: + import numpy as np + seen = np.zeros(int(self.config.vocab_size), dtype=bool) + seen[np.array(input_ids.tolist(), dtype=np.int64).reshape(-1)] = True + + def penalised(row): + if seen is None: + return row + row = row.astype(mx.float32) + return mx.where(mx.array(seen), mx.where(row < 0, row * penalty, row / penalty), row) t0 = time.monotonic() setup_mem = self._memory_snapshot() self._last_predicted = set() @@ -312,6 +422,12 @@ def generate( self.store.set_protected_experts([]) else: self._install_runtime_hooks() + planner = None + if not all_resident_fast_path and scheduler is None and residency == "planned": + from chronos.residency_planner import NextUsePlanner + planner = NextUsePlanner(int(self.config.num_experts), decay=float(residency_decay)) + self._attach_planner(planner) + hot_cap = int(getattr(self.store, "_capacity", 1) or 1) # ── Prefill phase: front-load expert IO ─────────────────────────── if scheduler is not None and not all_resident_fast_path: import torch, numpy as np @@ -326,17 +442,24 @@ def generate( prefill_masks = None else: prefill_masks = self._build_avail_masks(None) + if planner is not None: + planner.begin_forward("multi") logits, lookahead_probs, cache = self.model(input_ids, available_expert_masks=prefill_masks) mx.eval(logits) + if planner is not None: + planner.end_forward() self._last_current = [] if all_resident_fast_path else self._current_virtual_ids() prefill_t1 = time.monotonic() prefill_mem = self._memory_snapshot() # Sample from prefill - next_token = self._sample(logits[:, -1, :], temperature, top_p) + next_token = self._sample(penalised(logits[:, -1, :]), temperature, top_p) activated_ids: List[int] = [] tokens = 1 - if not all_resident_fast_path and scheduler is None and lookahead_probs is not None: + if planner is not None: + if lookahead_probs is not None: + planner.hint(self._predict_future_experts(lookahead_probs, capacity=hot_cap, max_steps=1)) + elif not all_resident_fast_path and scheduler is None and lookahead_probs is not None: future_ids = self._predict_future_experts( lookahead_probs, capacity=int(getattr(self.store, "_warm_capacity", self.config.num_experts) or self.config.num_experts), @@ -351,6 +474,8 @@ def generate( self._prefetch_and_promote_window(future_ids, promote_ids=immediate_ids, timeout_s=0.025) eos = set([eos_token_id] if isinstance(eos_token_id, int) else (eos_token_id or [])) first_id = int(next_token.item()) + if seen is not None: + seen[first_id] = True yield first_id for _ in range(0 if first_id in eos else max_new_tokens - 1): @@ -358,6 +483,11 @@ def generate( avail_masks = None elif scheduler is not None: avail_masks = self._build_avail_masks(next_token) + elif planner is not None: + planner.begin_forward("single") + self._planned_prefetch(planner, float(decode_wait_s)) + self._last_predicted = set(planner.upcoming()) + avail_masks = self._build_avail_masks(next_token) else: # LookaheadRouter-driven prefetch if lookahead_probs is not None: @@ -382,12 +512,18 @@ def generate( available_expert_masks=avail_masks, ) mx.eval(logits) + if planner is not None: + planner.end_forward() + if lookahead_probs is not None: + planner.hint(self._predict_future_experts(lookahead_probs, capacity=hot_cap, max_steps=1)) if not all_resident_fast_path: self._last_current = self._current_virtual_ids() - next_token = self._sample(logits[:, -1, :], temperature, top_p) + next_token = self._sample(penalised(logits[:, -1, :]), temperature, top_p) mx.eval(next_token) tok_id = int(next_token.item()) + if seen is not None: + seen[tok_id] = True activated_ids.append(tok_id) tokens += 1 yield tok_id @@ -405,22 +541,38 @@ def generate( "tokens_per_sec": round(tokens / max(end - t0, 1e-6), 2), "prefill_time_s": round(prefill_t1 - prefill_t0, 3), "decode_time_s": round(max(0.0, end - prefill_t1), 3), + "residency": {"mode": "full" if all_resident_fast_path else ( + "planned" if planner is not None else "window"), + **({"planner": dict(planner.stats)} if planner is not None else {})}, **self._memory_fields(setup_mem, "after_setup"), **self._memory_fields(prefill_mem, "after_prefill"), **self._memory_fields(end_mem, "after_decode"), **self._runtime_stat_fields(), **self.store.stats(), } + self.store.set_eviction_ranker(None) def _generate_speculative(self, input_ids, *, max_new_tokens, temperature, top_p, top_k, draft_depth, gamma, - gamma_max, adaptive_gamma, seed, eos_token_id) -> Iterator[int]: + gamma_max, adaptive_gamma, seed, eos_token_id, residency="window", + residency_decay=0.7, verify_wait_s=0.02) -> Iterator[int]: """Self-speculative decoding with the same residency work as plain decoding. Every forward (draft or verify) runs with the current availability masks. After it, the experts its last position routed to (all virtual layers) are unioned with the LookaheadRouter prediction and prefetched, so experts a draft touches are being loaded before the verify forward - that re-reads them.""" + that re-reads them. + + ``residency="planned"`` hands eviction to a next-use planner + (``chronos.residency_planner.NextUsePlanner``). Each draft token ran + the prelude and the first ``draft_depth`` iterations, and verify + re-runs exactly those dispatches for the same tokens plus the deeper + ones, so the draft queue is the verify forward's page-access + reference: pages the queue needs early are prefetched while the + drafts run and awaited right before verify (``verify_wait_s``), and + the page needed latest (or not at all) is evicted first; the + previous forward of the same kind covers the deeper iterations. + Output is unchanged: residency only decides what is loaded when.""" from chronos.mlx.speculative import SpecStats, speculative_tokens t0 = time.monotonic() setup_mem = self._memory_snapshot() @@ -437,12 +589,38 @@ def _generate_speculative(self, input_ids, *, max_new_tokens, temperature, top_p hot_cap = int(getattr(self.store, "_capacity", 1) or 1) timing = {"first": None} + planner = None + if not fast and residency == "planned": + from chronos.residency_planner import NextUsePlanner + planner = NextUsePlanner(int(self.config.num_experts), decay=float(residency_decay)) + self._attach_planner(planner) + queue_stats = {"verify_plans": 0, "verify_planned": 0, "verify_ready": 0} + def forward_fn(ids, cache, depth): + verify = depth is None and cache is not None + if planner is not None: + kind = "draft" if depth is not None else ("multi" if int(ids.shape[1]) > 1 else "single") + planner.begin_forward(kind) + # Drafts only queue SSD reads (the next draft is not delayed); + # verify waits briefly for the draft queue's pages. + self._planned_prefetch(planner, 0.0 if depth is not None else float(verify_wait_s)) + if verify: + hot = self.store.hot_expert_ids() + upcoming = planner.upcoming(hot_cap) + queue_stats["verify_plans"] += 1 + queue_stats["verify_planned"] += len(upcoming) + queue_stats["verify_ready"] += sum(1 for eid in upcoming if eid in hot) masks = None if fast else self._build_avail_masks(None) logits, lookahead, new_cache = self.model(ids, cache=cache, available_expert_masks=masks, recurrent_depth=depth) mx.eval(logits) - if not fast: + if planner is not None: + planner.end_forward(verify=verify) + if lookahead is not None: + planner.hint(self._predict_future_experts(lookahead, capacity=hot_cap, max_steps=1)) + self._last_current = self._current_virtual_ids() + self._last_predicted = set(planner.upcoming()) + elif not fast: self._last_current = self._current_virtual_ids() predicted = (self._predict_future_experts(lookahead, capacity=warm_cap) if lookahead is not None else []) @@ -484,11 +662,16 @@ def forward_fn(ids, cache, depth): "speculative": stats.as_dict() | { "draft_depth": eff_draft, "full_depth": full_depth, "enabled": eff_draft < full_depth, "gamma_final": (stats.gamma_history or [int(gamma)])[-1]}, + "residency": {"mode": "full" if fast else residency, **queue_stats, + "verify_ready_rate": round(queue_stats["verify_ready"] + / max(queue_stats["verify_planned"], 1), 4), + **({"planner": dict(planner.stats)} if planner is not None else {})}, **self._memory_fields(setup_mem, "after_setup"), **self._memory_fields(end_mem, "after_decode"), **self._runtime_stat_fields(), **self.store.stats(), } + self.store.set_eviction_ranker(None) def _predict_future_experts( self, @@ -497,16 +680,24 @@ def _predict_future_experts( max_steps: Optional[int] = None, ) -> List[int]: """Extract bounded top-k expert IDs for future steps.""" + e = int(self.config.num_experts) + budget = e if capacity is None else max(1, min(e, int(capacity))) + values = self._future_expert_scores(lookahead_probs, max_steps=max_steps) + scored = [(idx, value) for idx, value in enumerate(values) if value > 0.0] + scored.sort(key=lambda item: (-item[1], item[0])) + return [int(idx) for idx, _value in scored[:budget]] + + def _future_expert_scores(self, lookahead_probs: mx.array, max_steps: Optional[int] = None) -> List[float]: + """Per-expert LookaheadRouter score: top-k weight of each future step, weighted 1/(offset+1).""" # This store pages an expert ID across physical layers. Preserve # each prediction head until admission, then rank the union of IDs. + e = int(self.config.num_experts) future = lookahead_probs[:, -1, 1:, ...] if max_steps is not None: future = future[:, :max(0, int(max_steps))] if future.shape[1] <= 0: - return [] - e = int(self.config.num_experts) + return [0.0] * e top_k = max(1, min(int(self.config.num_experts_per_tok), e)) - budget = e if capacity is None else max(1, min(e, int(capacity))) scores = mx.zeros((e,), dtype=mx.float32) for offset in range(future.shape[1]): step = future[:, offset, ...].reshape(-1, e).astype(mx.float32) @@ -515,9 +706,7 @@ def _predict_future_experts( one_hot = (ids[..., None] == mx.arange(e)).astype(mx.float32) scores = scores + (one_hot*vals[..., None]).sum(axis=(0, 1)) mx.eval(scores) - scored = [(idx, float(value)) for idx, value in enumerate(scores.tolist()) if float(value) > 0.0] - scored.sort(key=lambda item: (-item[1], item[0])) - return [int(idx) for idx, _value in scored[:budget]] + return [float(v) for v in scores.tolist()] def _build_avail_masks(self, _token) -> List[set[int]] | None: """Promote prefetched experts and return per-layer availability masks.""" diff --git a/chronos/mlx/moe.py b/chronos/mlx/moe.py index ec162e8..9bd1620 100644 --- a/chronos/mlx/moe.py +++ b/chronos/mlx/moe.py @@ -193,6 +193,8 @@ def __call__( # ``.item()`` sync per expert per MoE call during decode). hit = (topk_idx.reshape(-1)[:, None] == mx.arange(self.num_experts)[None, :]).any(axis=0) selected_host = hit.tolist() + # Read by residency planners from the runtime hooks below. + object.__setattr__(self, "_selected_host", selected_host) for i, expert in enumerate(self.experts): sel = (topk_idx == i) diff --git a/chronos/mlx/training/trainer.py b/chronos/mlx/training/trainer.py index d6d2b38..805d215 100644 --- a/chronos/mlx/training/trainer.py +++ b/chronos/mlx/training/trainer.py @@ -1,6 +1,7 @@ """Minimal but real MLX-native six-stage training support.""" from __future__ import annotations +import math import os import time from dataclasses import dataclass @@ -128,6 +129,40 @@ def _router_locality_loss(router_mean: mx.array, num_experts_per_tok: int) -> mx return (-mx.log(mx.maximum(mass / float(top_k), 1e-9))).mean() +_MEMBERSHIP_CACHE: dict = {} + + +def _cluster_membership(config) -> mx.array: + """``[E, C]`` one-hot of ``chronos.expert_clusters.training_clusters(config)`` (cached).""" + from chronos.expert_clusters import cluster_membership, resolve_cluster_count, training_clusters + e = int(config.num_experts) + key = (e, resolve_cluster_count(config)) + hit = _MEMBERSHIP_CACHE.get(key) + if hit is None: + hit = mx.array(cluster_membership(training_clusters(config), e), dtype=mx.float32) + _MEMBERSHIP_CACHE[key] = hit + return hit + + +def _cluster_locality_loss(router_4d: mx.array, membership: mx.array, window: int = 0) -> mx.array: + """MLX twin of ``chronos.model.temporal_loss.cluster_locality_loss``: ``1 − I(window; cluster)/log C``.""" + B, S, V, E = router_4d.shape + C = int(membership.shape[1]) + w = S if int(window or 0) <= 0 else min(int(window), S) + n = S // w + if C <= 1 or n <= 0: + return mx.zeros(()) + mass = (router_4d.astype(mx.float32).reshape(-1, E) @ membership).reshape(B, S, V, C) + pw = mass[:, : n * w].reshape(B, n, w, V, C).mean(axis=(2, 3)).reshape(B * n, C) + + def entropy(p): + p = mx.maximum(p, 1e-9) + return -(p * mx.log(p)).sum(axis=-1) + + mutual = entropy(pw.mean(axis=0)) - entropy(pw).mean() + return 1.0 - mutual / math.log(C) + + def _load_balance_loss(router_4d: mx.array, num_experts_per_tok: int) -> mx.array: B, S, L, E = router_4d.shape top_k = max(1, min(int(num_experts_per_tok or 1), E)) @@ -388,6 +423,10 @@ def add(name, coefficient, value): router = router.astype(mx.float32) top_k = config.num_experts_per_tok add("balance_loss", config.lambda_balance, _load_balance_loss(router, top_k)) + if float(getattr(config, "lambda_cluster_locality", 0.0) or 0.0) > 0.0: + add("cluster_locality_loss", config.lambda_cluster_locality, + _cluster_locality_loss(router, _cluster_membership(config), + int(getattr(config, "cluster_locality_window", 0) or 0))) if router.shape[1] > 1: mean = router.mean(axis=2) add("temporal_loss", config.lambda_temporal, _temporal_locality_loss(mean)) diff --git a/chronos/model/checkpoint.py b/chronos/model/checkpoint.py index a2717e6..3640ffa 100644 --- a/chronos/model/checkpoint.py +++ b/chronos/model/checkpoint.py @@ -231,6 +231,10 @@ def read_checkpoint_config(checkpoint_path: str) -> tuple[dict, str | None]: "layer_aware_lookahead": False, "layer_aware_cache": False, }) + if "lambda_cluster_locality" not in cfg: + # Saved before the cluster-locality field existed: those weights were + # trained without the constraint (the field defaults on since 0.3.1). + cfg["lambda_cluster_locality"] = 0.0 return cfg, path @@ -395,6 +399,9 @@ def config_dict_for_checkpoint( cfg.update(sniffed) sources.append(f"{checkpoint_path} (tensor shapes)") + # How the weights were trained is a property of the checkpoint: a project + # config must not fill it in (0.0 would otherwise count as "missing"). + trained_lambda = cfg.get("lambda_cluster_locality") project_cfg, project_path = read_project_config(project_config_path) if project_cfg: for key, value in project_cfg.items(): @@ -403,6 +410,8 @@ def config_dict_for_checkpoint( if key not in cfg or cfg.get(key) in (None, "", 0): cfg[key] = value sources.append(project_path or project_config_path or "chronos_config.json") + if trained_lambda is not None: + cfg["lambda_cluster_locality"] = trained_lambda if overrides: for key, value in overrides.items(): diff --git a/chronos/model/config.py b/chronos/model/config.py index 6e96663..bad1ccb 100644 --- a/chronos/model/config.py +++ b/chronos/model/config.py @@ -195,6 +195,19 @@ def __init__(self, **kwargs): self.lambda_lookahead_topk: float = kwargs.pop("lambda_lookahead_topk", 0.15) self.lambda_lookahead_union: float = kwargs.pop("lambda_lookahead_union", 0.05) self.lambda_router_locality: float = kwargs.pop("lambda_router_locality", 0.02) + # Cluster-locality constraint: within each ``cluster_locality_window`` + # tokens (0 = whole sequence) and each virtual layer (every loop + # iteration), push routing mass into few expert clusters. The + # partition is ``chronos.expert_clusters.training_clusters`` and is the + # one lazy inference pages by (``inference_clusters``: automatic when + # λ > 0). The class default stays 0 so a stored recipe or checkpoint + # that predates the field keeps its training objective; every entry + # point that creates a new training config (planner, presets, CLI, + # bundle finetune) turns it on at the measured + # ``chronos.expert_clusters.DEFAULT_LAMBDA``. + self.lambda_cluster_locality: float = kwargs.pop("lambda_cluster_locality", 0.0) + self.cluster_locality_clusters: int = kwargs.pop("cluster_locality_clusters", 0) + self.cluster_locality_window: int = kwargs.pop("cluster_locality_window", 32) # M4: weight on the router KL anchor (alignment stages only). # Prevents SFT/DPO/ORPO/GRPO gradients from drifting routing away # from the pretrained distribution that the cluster layout was diff --git a/chronos/model/planner.py b/chronos/model/planner.py index b3429c3..8d76e88 100644 --- a/chronos/model/planner.py +++ b/chronos/model/planner.py @@ -73,6 +73,7 @@ class DatasetProfile: chars_per_token: float | None = None seq_p95_tokens: int | None = None # sampled records (estimated) seq_cap_tokens: int | None = None # exact cap from a shard manifest + seq_median_tokens: float | None = None # sampled records (estimated) vocab_size: int | None = None tokenizer_dir: str = "" notes: list = field(default_factory=list) @@ -224,7 +225,8 @@ def profile_datasets(paths, *, max_files: int = 64) -> DatasetProfile: prof.seq_p95_tokens = int(math.ceil(tok_lens[int(0.95 * (len(tok_lens) - 1))])) total_tok = sum(tok_lens) prof.chars_per_token = round(chars / total_tok, 3) if total_tok else None - prof.notes.append(f"median record ≈ {statistics.median(tok_lens):.0f} tokens (estimated)") + prof.seq_median_tokens = float(statistics.median(tok_lens)) + prof.notes.append(f"median record ≈ {prof.seq_median_tokens:.0f} tokens (estimated)") return prof @@ -557,6 +559,16 @@ def save_interval(steps: float) -> int: # SWA window: the typical record must fit; long documents use MLA layers. window = min(seq_len, max(128, 1 << int(math.ceil(math.log2(max(2, p95 // 4)))))) if p95 \ else max(128, seq_len // 4) + # Cluster-locality training (on for MoE): measured default strength, window + # = one record (stream segment); C = ⌈E/2k⌉ is resolved from the topology. + from chronos.expert_clusters import DEFAULT_LAMBDA, recommended_window + cluster_lambda = DEFAULT_LAMBDA if req.use_moe and arch.num_experts > 2 * arch.num_experts_per_tok else 0.0 + cluster_window = recommended_window(prof.seq_median_tokens, seq_len) + sources["lambda_cluster_locality"] = ( + "measured sweep, tools/bench_cluster_locality.py (smallest λ on the plateau)" if cluster_lambda + else "off: dense model or fewer than two clusters") + sources["cluster_locality_window"] = ( + "median record length (estimated)" if prof.seq_median_tokens else "default (no record lengths sampled)") per_expert_bytes = 3 * h * arch.ffn_width * 2 # bf16 page expert_mb = routed * 2 / 2**20 ram_budget = None if policy["ram_fraction"] >= 1 or not req.use_moe else \ @@ -580,6 +592,7 @@ def save_interval(steps: float) -> int: ponder_loss_coef=3e-3, recurrent_distill_coef=0.05, recurrent_stability_coef=1e-3, lambda_balance=5e-3 if req.use_moe else 0.0, lambda_temporal=1e-3, lambda_lookahead=0.1, lambda_lookahead_topk=0.05, fallback_mask_prob=0.08 if req.use_moe else 0.0, dropout=0.0, + lambda_cluster_locality=cluster_lambda, cluster_locality_clusters=0, cluster_locality_window=cluster_window, kv_cache_dtype=policy["kv_cache_dtype"], mla_cache_rope=policy["mla_cache_rope"], router_scope=policy["router_scope"], ram_budget_mb=ram_budget, activation_checkpointing=policy["activation_checkpointing"], optimizer_paging=policy["optimizer_paging"], diff --git a/chronos/model/temporal_loss.py b/chronos/model/temporal_loss.py index ab013cf..61bdc97 100644 --- a/chronos/model/temporal_loss.py +++ b/chronos/model/temporal_loss.py @@ -1,3 +1,5 @@ +import math + import torch import torch.nn.functional as F @@ -65,6 +67,80 @@ def router_locality_loss( return (-torch.log(overlap_mass / float(top_k))).mean() +def _window_cluster_mass(router_4d: torch.Tensor, membership: torch.Tensor, window: int): + """``[W, C]``: routing mass per cluster for each window (tokens × every dispatch).""" + B, S, V, E = router_4d.shape + C = int(membership.shape[1]) + w = S if int(window or 0) <= 0 else min(int(window), S) + n = S // w + if n <= 0 or C <= 1: + return None + mass = torch.einsum("bsve,ec->bsvc", router_4d.float(), membership.to(router_4d.device, torch.float32)) + return mass[:, : n * w].reshape(B, n, w, V, C).mean(dim=(2, 3)).reshape(B * n, C) + + +def _entropy(p: torch.Tensor) -> torch.Tensor: + p = p.clamp_min(1e-9) + return -(p * p.log()).sum(dim=-1) + + +def cluster_locality_loss( + router_4d: torch.Tensor, + membership: torch.Tensor, + window: int = 0, +) -> torch.Tensor: + """Cluster-locality constraint over a stream and its recurrent iterations. + + ``router_4d`` is ``[B, S, V, E]``: every MoE dispatch (V = virtual layers, + so a recurrent group contributes one slice per loop iteration). + ``membership`` is the ``[E, C]`` one-hot expert→cluster map that the lazy + store pages by. For each window of ``window`` tokens (0 = the whole + sequence) the routing mass per cluster is averaged over its tokens *and* + all its dispatches, giving ``p_w``. The loss is + + 1 − I(window; cluster) / log C, I = H(mean_w p_w) − mean_w H(p_w) + + Low ``H(p_w)`` keeps a stream, across tokens, layers and loop + iterations, inside few clusters: the pages one cluster read serves. + High ``H(mean_w p_w)`` spreads different windows over different + clusters, so experts specialise per stream instead of collapsing onto + one cluster. Range ``[0, 1]``; one window in the batch gives no signal. + """ + if router_4d is None or router_4d.dim() != 4: + return torch.tensor(0.0) + pw = _window_cluster_mass(router_4d, membership, window) + if pw is None: + return router_4d.new_zeros(()) + log_c = math.log(int(membership.shape[1])) + mutual = _entropy(pw.mean(dim=0)) - _entropy(pw).mean() + return 1.0 - mutual / log_c + + +def cluster_routing_metrics(router_4d: torch.Tensor, membership: torch.Tensor, top_k: int, + window: int = 0) -> dict[str, float]: + """Detached diagnostics of the cluster constraint (normalised entropies, clusters touched).""" + if router_4d is None or router_4d.dim() != 4: + return {} + with torch.no_grad(): + B, S, V, E = router_4d.shape + C = int(membership.shape[1]) + w = S if int(window or 0) <= 0 else min(int(window), S) + n = S // w + pw = _window_cluster_mass(router_4d, membership, window) + if pw is None: + return {} + log_c = math.log(C) + k = max(1, min(int(top_k), E)) + cid = membership.to(router_4d.device).argmax(dim=-1)[router_4d.float().topk(k, dim=-1).indices] + hit = torch.zeros(B, S, V, C, device=router_4d.device).scatter_(-1, cid, 1.0) + touched = hit[:, : n * w].reshape(B, n, w * V, C).amax(dim=2).sum(dim=-1) + return { + "cluster_window_entropy": float(_entropy(pw).mean() / log_c), + "cluster_marginal_entropy": float(_entropy(pw.mean(dim=0)) / log_c), + "clusters_per_window": float(touched.mean()), + } + + def router_offload_metrics( router_probs: torch.Tensor, lookahead_probs: torch.Tensor = None, diff --git a/chronos/residency_planner.py b/chronos/residency_planner.py new file mode 100644 index 0000000..b1b571c --- /dev/null +++ b/chronos/residency_planner.py @@ -0,0 +1,154 @@ +"""Next-use residency planning for lazy expert pages (backend-free; no torch / mlx imports). + +A lazy store holds fewer expert pages than one forward touches, so the page +to evict is the one needed *latest*. Decoding is repetitive enough to +predict that order: + +* forwards of one kind run the same dispatch sequence (prelude, every loop + iteration of every recurrent layer, coda), and consecutive tokens route + similarly, so the previous forward of the same kind is a next-use + reference; +* in self-speculative decoding the *draft queue* is a better reference for + the verify forward: each draft token ran the prelude and the first + ``draft_depth`` iterations, and verify re-runs exactly those dispatches + for the same token, plus the deeper iterations. + +``NextUsePlanner`` records the page-access trace of every forward (fed by +the MoE runtime hooks), builds the next-use table of the reference before +each forward, and ranks pages for eviction: + + -1 selected by the dispatch running now and not executed yet → never + 0 used again later in this forward → evict the one needed latest + 1 not again in this forward, expected in the next one + 2 not expected → evict first; pages the LookaheadRouter predicts + (:meth:`hint`) go last, then the lowest decayed demand + (``score ← decay·score + uses`` per forward, the sliding window) + +It also proposes a bounded prefetch: pages the reference needs early that +are not resident, at most as many as free slots plus resident pages of +class 2, so a prefetch never displaces a page needed sooner. +""" +from __future__ import annotations + +from typing import Dict, Iterable, List, Optional, Tuple + +Key = Tuple[int, int] # (physical MoE layer, loop iteration) + +_FAR = 1 << 30 + + +class NextUsePlanner: + def __init__(self, num_experts: int, *, decay: float = 0.7): + if not 0.0 <= float(decay) < 1.0: + raise ValueError("decay must be in [0, 1)") + self.num_experts = int(num_experts) + self.decay = float(decay) + self.score = [0.0] * self.num_experts + self._order: Dict[Key, int] = {} # dispatch → execution ordinal (learned once) + self._prev: Dict[str, List[Tuple[int, int]]] = {} + self._drafts: List[List[Tuple[int, int]]] = [] + self._kind = "single" + self._trace: List[Tuple[int, int]] = [] + self._next_use: Dict[int, List[int]] = {} + self._next_forward: set[int] = set() + self._here = -1 + self._pending: frozenset = frozenset() + self._hint: frozenset = frozenset() + self.stats = {"forwards": 0, "accesses": 0, "queue_verifies": 0} + + # ── forward boundaries ──────────────────────────────────────── + + def begin_forward(self, kind: str) -> None: + """``kind``: ``"single"`` (one-token decode), ``"multi"`` (prefill / verify) or ``"draft"``.""" + self._kind = kind + self._trace = [] + self._here = -1 + self._pending = frozenset() + ref = list(self._prev.get(kind, ())) + if kind == "multi" and self._drafts: + ref += [a for d in self._drafts for a in d] + self.stats["queue_verifies"] += 1 + table: Dict[int, List[int]] = {} + for ordinal, eid in ref: + table.setdefault(eid, []).append(ordinal) + for v in table.values(): + v.sort() + self._next_use = table + nxt = "draft" if kind == "multi" and "draft" in self._prev else kind + self._next_forward = {eid for _o, eid in self._prev.get(nxt, ())} + + def access(self, layer: int, iteration: int, eid: int, pending: Iterable[int] = ()) -> None: + """Dispatch ``(layer, iteration)`` is about to run expert ``eid``. + + ``pending`` are the dispatch's other selected experts that run after + it (an MoE runs its selected experts in id order); they are never + evicted to make room for ``eid``. + """ + key = (int(layer), int(iteration)) + ordinal = self._order.get(key) + if ordinal is None: + ordinal = self._order[key] = len(self._order) + self._here = ordinal + self._pending = frozenset(int(e) for e in pending) + entry = (ordinal, int(eid)) + if not self._trace or self._trace[-1] != entry: # loader + touch report one run + self._trace.append(entry) + self.stats["accesses"] += 1 + + def end_forward(self, *, verify: bool = False) -> None: + trace = self._trace + self._prev[self._kind] = trace + if self._kind == "draft": + self._drafts.append(trace) + elif verify: + self._drafts = [] + used = {eid for _o, eid in trace} + d = self.decay + self.score = [s * d + (1.0 if i in used else 0.0) for i, s in enumerate(self.score)] + self.stats["forwards"] += 1 + + def hint(self, expert_ids: Iterable[int]) -> None: + """Pages predicted for the next token (LookaheadRouter); kept over other class-2 pages.""" + self._hint = frozenset(int(e) for e in expert_ids) + + # ── ranking ─────────────────────────────────────────────────── + + def evict_key(self, eid: int): + """Sort key; the *largest* is evicted first.""" + if int(eid) in self._pending: + return (-1, 0, 0.0) + later = self._next_use.get(int(eid)) + if later: + for ordinal in later: + if ordinal > self._here: + return (0, ordinal, 0.0) + if int(eid) in self._next_forward: + return (1, 0, -self.score[int(eid)]) + return (2, 0 if int(eid) in self._hint else 1, -self.score[int(eid)]) + + def victim(self, candidates: Iterable[int]) -> Optional[int]: + best, best_key = None, None + for eid in candidates: + key = self.evict_key(eid) + if best_key is None or key > best_key: + best, best_key = eid, key + return best + + def upcoming(self, limit: Optional[int] = None) -> List[int]: + """Pages the upcoming forward is expected to touch, earliest first.""" + first = sorted((v[0], eid) for eid, v in self._next_use.items() if v) + ids = [eid for _o, eid in first] + return ids if limit is None else ids[:max(0, int(limit))] + + def prefetch_plan(self, resident: Iterable[int], capacity: int) -> List[int]: + """Missing pages the upcoming forward needs, earliest first, that fit without + displacing a page needed sooner: free slots plus resident class-2 pages. + + Call between :meth:`begin_forward` and the forward itself. + """ + resident = set(int(e) for e in resident) + spare = sum(1 for e in resident if self.evict_key(e)[0] == 2) + budget = spare + max(0, int(capacity) - len(resident)) + if budget <= 0: + return [] + return [eid for eid in self.upcoming() if eid not in resident][:budget] diff --git a/chronos/runtime/inference_engine.py b/chronos/runtime/inference_engine.py index fd0e1ec..2b24b21 100644 --- a/chronos/runtime/inference_engine.py +++ b/chronos/runtime/inference_engine.py @@ -299,15 +299,18 @@ def setup_from_state_dict( self, state_dict: dict, warm_expert_ids: Optional[List[int]] = None, + clusters: Optional[List[List[int]]] = None, ): """Setup path for real lazy expert loading on CPU/MPS. The caller loads base/non-expert weights into the live model, then hands the full checkpoint here so ExpertStore can shard only expert tensors to SSD and replace live experts with placeholders. + ``clusters`` pages by cluster file (``chronos.lazy_experts``); None + keeps one shard per expert. """ print("[Chronos] Offloading expert weights to SSD from checkpoint...") - self.cache_manager.expert_store.offload_from_state_dict(state_dict) + self.cache_manager.expert_store.offload_from_state_dict(state_dict, clusters=clusters) self.cache_manager.start() print("[Chronos] Warming up VRAM cache...") self.cache_manager.warm_up(warm_expert_ids) @@ -317,10 +320,11 @@ def setup_from_state_reader( self, reader, warm_expert_ids: Optional[List[int]] = None, + clusters: Optional[List[List[int]]] = None, ): """Setup lazy expert loading from a safetensors/GGUF export reader.""" print("[Chronos] Offloading expert weights to SSD from export reader...") - self.cache_manager.expert_store.offload_from_state_reader(reader) + self.cache_manager.expert_store.offload_from_state_reader(reader, clusters=clusters) self.cache_manager.start() print("[Chronos] Warming up VRAM cache...") self.cache_manager.warm_up(warm_expert_ids) diff --git a/chronos/serving/ollama_server.py b/chronos/serving/ollama_server.py index 8b2a8e5..47fdf69 100644 --- a/chronos/serving/ollama_server.py +++ b/chronos/serving/ollama_server.py @@ -165,7 +165,10 @@ def ps(): 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}) + chronos={"device": lm.device, "lazy_experts": lm.lazy_experts, + **{k: lm.meta[k] for k in ("paging", "paging_resolved", "num_pages") + if k in (lm.meta or {})}, + "expert_store": st}) out.append(entry) return {"models": out} @@ -311,6 +314,9 @@ def main(argv=None) -> int: "(auto: CUDA bf16/fp16, MPS fp32, CPU fp32)") parser.add_argument("--no-lazy-experts", action="store_true", help="keep every expert resident (no SSD offload)") parser.add_argument("--ram-expert-fraction", type=float, default=0.5) + parser.add_argument("--paging", choices=["auto", "cluster", "expert"], default="auto", + help="expert page layout; auto pages by training cluster when the checkpoint " + "was trained with cluster locality, per expert otherwise") # Inference-only options (default: whatever the checkpoint config says). parser.add_argument("--kv-cache-dtype", choices=["auto", "int8"], default=None, help="int8: per-token int8 KV cache with FP32 group scales") @@ -338,7 +344,7 @@ def main(argv=None) -> int: overrides = {k: v for k, v in overrides.items() if v is not None} runtime = ChronosRuntime(args.models_dir, device=args.device, lazy_experts=not args.no_lazy_experts, ram_expert_fraction=args.ram_expert_fraction, dtype=args.dtype, - inference_overrides=overrides) + inference_overrides=overrides, paging=args.paging) print(f"[chronos serve] http://{args.host}:{args.port} models={sorted(runtime.available())} " f"auth={'on' if args.api_key else 'off (loopback only)'} overrides={overrides or 'none'}", flush=True) if args.monitor_runs_dir: diff --git a/chronos/serving/runtime.py b/chronos/serving/runtime.py index b765fc4..3aa6c37 100644 --- a/chronos/serving/runtime.py +++ b/chronos/serving/runtime.py @@ -87,12 +87,14 @@ def _find_tokenizer_dir(ckpt_path: str, meta: dict) -> str: def load_model(path: str, *, name: str | None = None, device: str | None = "auto", lazy_experts: bool = True, ram_expert_fraction: float = 0.5, expert_cache_root: str | None = None, dtype: str = "auto", - inference_overrides: dict | None = None) -> LoadedModel: + inference_overrides: dict | None = None, paging: str = "auto") -> LoadedModel: """Load a Chronos checkpoint for serving. With ``lazy_experts`` (default) the routed experts are sharded to an SSD cache and only ``ram_expert_fraction`` of them are ever resident; the - rest stream in on demand, guided by the lookahead router. ``dtype`` + rest stream in on demand, guided by the lookahead router. ``paging`` + (``chronos.lazy_experts.page_layout``) pages by training cluster when the + checkpoint was trained with cluster locality, per expert otherwise. ``dtype`` picks the single compute dtype (``chronos.precision.resolve_torch``): weights, activations, KV cache and expert pages all use it. """ @@ -116,21 +118,27 @@ def load_model(path: str, *, name: str | None = None, device: str | None = "auto weights = load_checkpoint_state_dict(path, map_location="cpu") model = ChronosForCausalLM(config) engine = None + paging_info = {} if lazy_experts and getattr(config, "use_moe", True): - from chronos.runtime.inference_engine import ChronosInferenceEngine - from ui.tabs.inference_tab import ( - _apply_offload_expert_budget, _checkpoint_expert_cache_dir, - _replace_torch_experts_with_placeholders, + from chronos.lazy_experts import ( + apply_torch_budget, checkpoint_cache_dir, page_layout, torch_experts_to_placeholders, ) + from chronos.runtime.inference_engine import ChronosInferenceEngine + layout = page_layout(config, paging) base = {k: v for k, v in weights.items() if ".mlp.experts." not in k} load_state_dict_controlled(model, base, allow_missing_substrings=(".mlp.experts.",)) - _replace_torch_experts_with_placeholders(model, config) + torch_experts_to_placeholders(model, config) model, precision = place_torch_model(model, device, dtype) model = model.eval() - cache_dir = expert_cache_root or _checkpoint_expert_cache_dir(path, config) + cache_dir = checkpoint_cache_dir(path, config, tag=layout["tag"]) + if expert_cache_root: + cache_dir = os.path.join(expert_cache_root, os.path.basename(cache_dir)) engine = ChronosInferenceEngine(model, config, ssd_dir=cache_dir) - _apply_offload_expert_budget(engine, config, ram_expert_fraction) - engine.setup_from_state_dict(weights, warm_expert_ids=[]) + apply_torch_budget(engine, config, ram_expert_fraction) + engine.setup_from_state_dict(weights, warm_expert_ids=[], + clusters=None if layout["tag"] == "e" else layout["clusters"]) + paging_info = {"paging": layout["paging"], "paging_resolved": layout["resolved"], + "num_pages": layout["num_pages"]} else: load_state_dict_controlled(model, weights) model, precision = place_torch_model(model, device, dtype) @@ -141,7 +149,7 @@ def load_model(path: str, *, name: str | None = None, device: str | None = "auto config=config, model=model, tokenizer=tokenizer, engine=engine, device=device, lazy_experts=engine is not None, meta={"tokenizer_dir": tok_dir, "precision": precision.as_dict(), - "inference_overrides": applied, **meta_full}) + "inference_overrides": applied, **paging_info, **meta_full}) class ChronosRuntime: @@ -149,12 +157,16 @@ class ChronosRuntime: def __init__(self, models_dir: str | None = None, *, device: str = "auto", lazy_experts: bool = True, ram_expert_fraction: float = 0.5, keep_alive_s: float = 300.0, dtype: str = "auto", - inference_overrides: dict | None = None): + inference_overrides: dict | None = None, paging: str = "auto"): self.models_dir = models_dir self.inference_overrides = dict(inference_overrides or {}) self.dtype = dtype self.device, self.lazy_experts = device, lazy_experts self.ram_expert_fraction = ram_expert_fraction + from chronos.expert_clusters import PAGING_MODES + if paging not in PAGING_MODES: + raise ValueError(f"paging must be one of {PAGING_MODES}") + self.paging = paging self.keep_alive_s = keep_alive_s self.loaded: dict[str, LoadedModel] = {} self.last_used: dict[str, float] = {} @@ -184,7 +196,8 @@ def get(self, name: str) -> LoadedModel: lazy_experts=self.lazy_experts, ram_expert_fraction=self.ram_expert_fraction, dtype=self.dtype, - inference_overrides=self.inference_overrides) + inference_overrides=self.inference_overrides, + paging=self.paging) self.last_used[name] = time.time() return self.loaded[name] diff --git a/chronos/trainer/loss_mixin.py b/chronos/trainer/loss_mixin.py index 8df63fa..55774f8 100644 --- a/chronos/trainer/loss_mixin.py +++ b/chronos/trainer/loss_mixin.py @@ -27,6 +27,7 @@ from chronos.model.moe_chronos import ChronosMOEFeedForward from chronos.model.temporal_loss import ( + cluster_locality_loss, temporal_locality_loss, lookahead_supervision_loss, lookahead_topk_hit_loss, @@ -61,6 +62,13 @@ def collect_router_probs(model, *, with_grad: bool = False) -> Optional[torch.Te return torch.stack(probs, dim=2) +def cluster_membership_tensor(config, device=None) -> torch.Tensor: + """``[E, C]`` one-hot of ``chronos.expert_clusters.training_clusters(config)``.""" + from chronos.expert_clusters import cluster_membership, training_clusters + e = int(config.num_experts) + return torch.tensor(cluster_membership(training_clusters(config), e), dtype=torch.float32, device=device) + + def collect_iteration_router_probs(model, *, with_grad: bool = False) -> Optional[torch.Tensor]: """Every dispatch decision of the last forward as ``[B, S, V, E]``. @@ -153,6 +161,13 @@ def chronos_loss_term( router_mean, int(getattr(config, "num_experts_per_tok", 1) or 1), ) + lambda_cluster = float(getattr(config, "lambda_cluster_locality", 0.0) or 0.0) + if lambda_cluster > 0.0: + loss = loss + lambda_cluster * cluster_locality_loss( + router_4d, + cluster_membership_tensor(config, router_4d.device), + int(getattr(config, "cluster_locality_window", 0) or 0), + ) lambda_lookahead = float(getattr(config, "lambda_lookahead", 0.0)) lookahead_steps = int(getattr(config, "lookahead_steps", 0)) diff --git a/chronos/trainer/stage_utils.py b/chronos/trainer/stage_utils.py index d4305e0..9851dca 100644 --- a/chronos/trainer/stage_utils.py +++ b/chronos/trainer/stage_utils.py @@ -12,6 +12,7 @@ resolve_checkpoint_tokenizer, ) from chronos.model.config import ChronosConfig +from chronos.expert_clusters import DEFAULT_LAMBDA, DEFAULT_WINDOW from chronos.config_schema import MODEL_FIELDS @@ -138,6 +139,9 @@ def build_pretrain_config(args) -> ChronosConfig: "lambda_lookahead_topk": getattr(args, "lambda_lookahead_topk", 0.15), "lambda_lookahead_union": getattr(args, "lambda_lookahead_union", 0.05), "lambda_router_locality": getattr(args, "lambda_router_locality", 0.02), + "lambda_cluster_locality": getattr(args, "lambda_cluster_locality", DEFAULT_LAMBDA), + "cluster_locality_clusters": getattr(args, "cluster_locality_clusters", 0), + "cluster_locality_window": getattr(args, "cluster_locality_window", DEFAULT_WINDOW), "vram_budget_gb": getattr(args, "vram_budget_gb", 4.0), "use_moe": True, "architecture_version": 2, @@ -192,6 +196,8 @@ def build_config_from_upstream( "lambda_lookahead_topk", "lambda_lookahead_union", "lambda_router_locality", + "lambda_cluster_locality", + "cluster_locality_window", "fallback_mask_prob", ): value = getattr(args, key, None) diff --git a/pyproject.toml b/pyproject.toml index 99ad720..3bdc377 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "Project_Chronos" -version = "0.3.0" +version = "0.3.1a0" description = "Storage-aware recurrent-depth MoE: lookahead expert prefetch, SSD-lazy experts, hybrid MLA/SWA attention, self-speculative decoding" readme = "README.md" license = "Apache-2.0" diff --git a/reports/residency_20261107/bench_residency_50m_s13996_flores16_cpu.json b/reports/residency_20261107/bench_residency_50m_s13996_flores16_cpu.json new file mode 100644 index 0000000..5b282a1 --- /dev/null +++ b/reports/residency_20261107/bench_residency_50m_s13996_flores16_cpu.json @@ -0,0 +1,911 @@ +{ + "label": "measured", + "time": "2026-11-07T02:26:27", + "env": { + "python": "3.10.20", + "platform": "macOS-26.4.1-arm64-arm-64bit", + "torch": "2.10.0", + "mlx.core": "0.32.4", + "chip": "Apple M5 Max" + }, + "checkpoint": { + "path": "dist/bundles/chronos-enzh-50m-pretrain-s13996/model/chronos-enzh-50m-pretrain-s13996.pth", + "sha256": "8d4b356de7a0c7b372f33fdf5adb80e44a3dd1770ed6a6a1a5b890efd8a5328f" + }, + "settings": { + "checkpoint": "dist/bundles/chronos-enzh-50m-pretrain-s13996/model/chronos-enzh-50m-pretrain-s13996.pth", + "prompts": "flores_dev32.jsonl", + "field": "en", + "format": "tags:<|en|>,<|zh|>", + "eos": "<|endoftext|>", + "n": 16, + "max_new": 96, + "device": "cpu", + "dtype": "auto", + "ratio": 1.0, + "draft_depth": 1, + "gamma": 4, + "decay": 0.7, + "modes": "plain,planned,spec-window,spec-planned", + "out": "res16.json", + "note": "", + "eos_id": 0, + "decoding": "greedy, batch 1", + "order": "modes alternate per prompt; residency reset to the same cold start" + }, + "budget": { + "ideal_active_experts": 12, + "routing_top_k": 2, + "num_moe_layers": 6, + "requested_ram_load_ratio": 1.0, + "requested_expert_budget": 12, + "max_allowed_expert_budget": 15, + "effective_expert_budget": 12, + "effective_vram_expert_budget": 12, + "effective_ram_expert_budget": 12, + "min_vram_expert_budget": 1, + "budget_clamped": false + }, + "dtype": "float32", + "summary": { + "plain": { + "tokens": 549, + "seconds": 23.698, + "tok_s": 23.17, + "sync_ssd_loads": 6876, + "ssd_loads_per_token": 12.5246, + "selection_miss_rate": 0.2175, + "prefetch_loads": 7376, + "identical_to_plain": 16 + }, + "planned": { + "tokens": 549, + "seconds": 24.433, + "tok_s": 22.47, + "sync_ssd_loads": 4438, + "ssd_loads_per_token": 8.0838, + "selection_miss_rate": 0.1394, + "prefetch_loads": 5396, + "identical_to_plain": 16 + }, + "spec-window": { + "tokens": 549, + "seconds": 34.244, + "tok_s": 16.03, + "sync_ssd_loads": 9297, + "ssd_loads_per_token": 16.9344, + "selection_miss_rate": 0.2054, + "prefetch_loads": 11394, + "identical_to_plain": 16 + }, + "spec-planned": { + "tokens": 549, + "seconds": 28.961, + "tok_s": 18.96, + "sync_ssd_loads": 6258, + "ssd_loads_per_token": 11.3989, + "selection_miss_rate": 0.1395, + "prefetch_loads": 7004, + "identical_to_plain": 16, + "verify_ready_rate_mean": 0.6714 + } + }, + "rows": [ + { + "index": 0, + "plain": { + "s": 1.7617874999996275, + "ssd": 521, + "hits": 2239, + "misses": 521, + "cluster_loads": 564, + "prefetch_loads": 564, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 53 + }, + "planned": { + "s": 1.853657250001561, + "ssd": 360, + "hits": 2424, + "misses": 360, + "cluster_loads": 431, + "prefetch_loads": 431, + "planned_evictions": 419, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 53 + }, + "spec-window": { + "s": 2.780898666998837, + "ssd": 686, + "hits": 3327, + "misses": 686, + "cluster_loads": 852, + "prefetch_loads": 852, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 53 + }, + "spec-planned": { + "s": 2.416797707992373, + "ssd": 505, + "hits": 3480, + "misses": 505, + "cluster_loads": 571, + "prefetch_loads": 571, + "planned_evictions": 557, + "eviction_policy": "planned", + "verify_ready_rate": 0.7103, + "n_tokens": 53 + } + }, + { + "index": 1, + "plain": { + "s": 1.7775090409995755, + "ssd": 496, + "hits": 1835, + "misses": 496, + "cluster_loads": 529, + "prefetch_loads": 529, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 42 + }, + "planned": { + "s": 1.9302419169980567, + "ssd": 318, + "hits": 2003, + "misses": 318, + "cluster_loads": 392, + "prefetch_loads": 392, + "planned_evictions": 380, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 42 + }, + "spec-window": { + "s": 2.7124265000020387, + "ssd": 651, + "hits": 3039, + "misses": 651, + "cluster_loads": 809, + "prefetch_loads": 809, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 42 + }, + "spec-planned": { + "s": 2.4060076670139097, + "ssd": 455, + "hits": 3192, + "misses": 455, + "cluster_loads": 499, + "prefetch_loads": 499, + "planned_evictions": 486, + "eviction_policy": "planned", + "verify_ready_rate": 0.7361, + "n_tokens": 42 + } + }, + { + "index": 2, + "plain": { + "s": 1.6499649589968612, + "ssd": 474, + "hits": 1697, + "misses": 474, + "cluster_loads": 516, + "prefetch_loads": 516, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 38 + }, + "planned": { + "s": 1.8351264579978306, + "ssd": 327, + "hits": 1898, + "misses": 327, + "cluster_loads": 396, + "prefetch_loads": 396, + "planned_evictions": 384, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 38 + }, + "spec-window": { + "s": 2.4732384579983773, + "ssd": 795, + "hits": 2577, + "misses": 795, + "cluster_loads": 984, + "prefetch_loads": 984, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 38 + }, + "spec-planned": { + "s": 2.148162457990111, + "ssd": 518, + "hits": 2799, + "misses": 518, + "cluster_loads": 589, + "prefetch_loads": 589, + "planned_evictions": 574, + "eviction_policy": "planned", + "verify_ready_rate": 0.6078, + "n_tokens": 38 + } + }, + { + "index": 3, + "plain": { + "s": 1.062335875001736, + "ssd": 326, + "hits": 1221, + "misses": 326, + "cluster_loads": 346, + "prefetch_loads": 346, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 27 + }, + "planned": { + "s": 1.0693116670008749, + "ssd": 189, + "hits": 1366, + "misses": 189, + "cluster_loads": 230, + "prefetch_loads": 230, + "planned_evictions": 218, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 27 + }, + "spec-window": { + "s": 1.8674862500047311, + "ssd": 438, + "hits": 1881, + "misses": 438, + "cluster_loads": 537, + "prefetch_loads": 537, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 27 + }, + "spec-planned": { + "s": 1.51929187499627, + "ssd": 291, + "hits": 2029, + "misses": 291, + "cluster_loads": 324, + "prefetch_loads": 324, + "planned_evictions": 310, + "eviction_policy": "planned", + "verify_ready_rate": 0.7753, + "n_tokens": 27 + } + }, + { + "index": 4, + "plain": { + "s": 0.78412991600635, + "ssd": 314, + "hits": 819, + "misses": 314, + "cluster_loads": 337, + "prefetch_loads": 337, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 15 + }, + "planned": { + "s": 0.7847320839937311, + "ssd": 179, + "hits": 963, + "misses": 179, + "cluster_loads": 214, + "prefetch_loads": 214, + "planned_evictions": 202, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 15 + }, + "spec-window": { + "s": 1.2725300839956617, + "ssd": 408, + "hits": 1221, + "misses": 408, + "cluster_loads": 493, + "prefetch_loads": 493, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 15 + }, + "spec-planned": { + "s": 1.0609587910003029, + "ssd": 246, + "hits": 1372, + "misses": 246, + "cluster_loads": 282, + "prefetch_loads": 282, + "planned_evictions": 269, + "eviction_policy": "planned", + "verify_ready_rate": 0.6667, + "n_tokens": 15 + } + }, + { + "index": 5, + "plain": { + "s": 1.4923341250105295, + "ssd": 378, + "hits": 1573, + "misses": 378, + "cluster_loads": 408, + "prefetch_loads": 408, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 36 + }, + "planned": { + "s": 1.5851479999982985, + "ssd": 252, + "hits": 1725, + "misses": 252, + "cluster_loads": 311, + "prefetch_loads": 311, + "planned_evictions": 299, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 36 + }, + "spec-window": { + "s": 1.901140916001168, + "ssd": 479, + "hits": 2134, + "misses": 479, + "cluster_loads": 602, + "prefetch_loads": 602, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 36 + }, + "spec-planned": { + "s": 1.6241188749991124, + "ssd": 333, + "hits": 2286, + "misses": 333, + "cluster_loads": 373, + "prefetch_loads": 373, + "planned_evictions": 357, + "eviction_policy": "planned", + "verify_ready_rate": 0.7292, + "n_tokens": 36 + } + }, + { + "index": 6, + "plain": { + "s": 1.3860392089991365, + "ssd": 454, + "hits": 1325, + "misses": 454, + "cluster_loads": 489, + "prefetch_loads": 489, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 27 + }, + "planned": { + "s": 1.3286414169997443, + "ssd": 296, + "hits": 1498, + "misses": 296, + "cluster_loads": 364, + "prefetch_loads": 364, + "planned_evictions": 352, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 27 + }, + "spec-window": { + "s": 1.7600588750065072, + "ssd": 737, + "hits": 1609, + "misses": 737, + "cluster_loads": 871, + "prefetch_loads": 871, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 27 + }, + "spec-planned": { + "s": 1.4574707499996293, + "ssd": 446, + "hits": 1862, + "misses": 446, + "cluster_loads": 506, + "prefetch_loads": 506, + "planned_evictions": 490, + "eviction_policy": "planned", + "verify_ready_rate": 0.5476, + "n_tokens": 27 + } + }, + { + "index": 7, + "plain": { + "s": 1.5913087079970865, + "ssd": 505, + "hits": 1734, + "misses": 505, + "cluster_loads": 545, + "prefetch_loads": 545, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 37 + }, + "planned": { + "s": 1.8959565830009524, + "ssd": 365, + "hits": 1897, + "misses": 365, + "cluster_loads": 441, + "prefetch_loads": 441, + "planned_evictions": 429, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 37 + }, + "spec-window": { + "s": 2.58848045799823, + "ssd": 676, + "hits": 2633, + "misses": 676, + "cluster_loads": 849, + "prefetch_loads": 849, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 37 + }, + "spec-planned": { + "s": 2.261977083995589, + "ssd": 500, + "hits": 2827, + "misses": 500, + "cluster_loads": 555, + "prefetch_loads": 555, + "planned_evictions": 539, + "eviction_policy": "planned", + "verify_ready_rate": 0.6979, + "n_tokens": 37 + } + }, + { + "index": 8, + "plain": { + "s": 1.0879823330033105, + "ssd": 381, + "hits": 1064, + "misses": 381, + "cluster_loads": 405, + "prefetch_loads": 405, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 22 + }, + "planned": { + "s": 0.9571430000069086, + "ssd": 207, + "hits": 1239, + "misses": 207, + "cluster_loads": 244, + "prefetch_loads": 244, + "planned_evictions": 232, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 22 + }, + "spec-window": { + "s": 1.5000025830086088, + "ssd": 455, + "hits": 1656, + "misses": 455, + "cluster_loads": 542, + "prefetch_loads": 542, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 22 + }, + "spec-planned": { + "s": 1.3046174999908544, + "ssd": 286, + "hits": 1806, + "misses": 286, + "cluster_loads": 313, + "prefetch_loads": 313, + "planned_evictions": 300, + "eviction_policy": "planned", + "verify_ready_rate": 0.7222, + "n_tokens": 22 + } + }, + { + "index": 9, + "plain": { + "s": 1.3477633750007953, + "ssd": 481, + "hits": 1427, + "misses": 481, + "cluster_loads": 508, + "prefetch_loads": 508, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 30 + }, + "planned": { + "s": 1.5166876250004862, + "ssd": 303, + "hits": 1621, + "misses": 303, + "cluster_loads": 356, + "prefetch_loads": 356, + "planned_evictions": 344, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 30 + }, + "spec-window": { + "s": 2.0585927500069374, + "ssd": 614, + "hits": 2197, + "misses": 614, + "cluster_loads": 735, + "prefetch_loads": 735, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 30 + }, + "spec-planned": { + "s": 1.565166667001904, + "ssd": 400, + "hits": 2350, + "misses": 400, + "cluster_loads": 447, + "prefetch_loads": 447, + "planned_evictions": 433, + "eviction_policy": "planned", + "verify_ready_rate": 0.6389, + "n_tokens": 30 + } + }, + { + "index": 10, + "plain": { + "s": 3.1965762080071727, + "ssd": 560, + "hits": 3602, + "misses": 560, + "cluster_loads": 599, + "prefetch_loads": 599, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 90 + }, + "planned": { + "s": 3.1805184590048157, + "ssd": 406, + "hits": 3763, + "misses": 406, + "cluster_loads": 510, + "prefetch_loads": 510, + "planned_evictions": 498, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 90 + }, + "spec-window": { + "s": 4.091535250001471, + "ssd": 711, + "hits": 4599, + "misses": 711, + "cluster_loads": 891, + "prefetch_loads": 891, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 90 + }, + "spec-planned": { + "s": 3.731865875000949, + "ssd": 540, + "hits": 4759, + "misses": 540, + "cluster_loads": 594, + "prefetch_loads": 594, + "planned_evictions": 576, + "eviction_policy": "planned", + "verify_ready_rate": 0.759, + "n_tokens": 90 + } + }, + { + "index": 11, + "plain": { + "s": 1.0975627080042614, + "ssd": 392, + "hits": 979, + "misses": 392, + "cluster_loads": 412, + "prefetch_loads": 412, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 19 + }, + "planned": { + "s": 1.0790139170130715, + "ssd": 234, + "hits": 1147, + "misses": 234, + "cluster_loads": 278, + "prefetch_loads": 278, + "planned_evictions": 266, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 19 + }, + "spec-window": { + "s": 1.487006334005855, + "ssd": 503, + "hits": 1492, + "misses": 503, + "cluster_loads": 600, + "prefetch_loads": 600, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 19 + }, + "spec-planned": { + "s": 1.2269112909998512, + "ssd": 327, + "hits": 1661, + "misses": 327, + "cluster_loads": 365, + "prefetch_loads": 365, + "planned_evictions": 350, + "eviction_policy": "planned", + "verify_ready_rate": 0.6061, + "n_tokens": 19 + } + }, + { + "index": 12, + "plain": { + "s": 0.8778314169903751, + "ssd": 280, + "hits": 888, + "misses": 280, + "cluster_loads": 304, + "prefetch_loads": 304, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 18 + }, + "planned": { + "s": 0.7966191669984255, + "ssd": 156, + "hits": 988, + "misses": 156, + "cluster_loads": 187, + "prefetch_loads": 187, + "planned_evictions": 175, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 18 + }, + "spec-window": { + "s": 1.1623969590000343, + "ssd": 342, + "hits": 1238, + "misses": 342, + "cluster_loads": 408, + "prefetch_loads": 408, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 18 + }, + "spec-planned": { + "s": 0.9529859169997508, + "ssd": 207, + "hits": 1347, + "misses": 207, + "cluster_loads": 228, + "prefetch_loads": 228, + "planned_evictions": 216, + "eviction_policy": "planned", + "verify_ready_rate": 0.713, + "n_tokens": 18 + } + }, + { + "index": 13, + "plain": { + "s": 1.1606041670020204, + "ssd": 352, + "hits": 1090, + "misses": 352, + "cluster_loads": 383, + "prefetch_loads": 383, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 23 + }, + "planned": { + "s": 1.1387737920013024, + "ssd": 215, + "hits": 1252, + "misses": 215, + "cluster_loads": 266, + "prefetch_loads": 266, + "planned_evictions": 254, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 23 + }, + "spec-window": { + "s": 1.6385499169991817, + "ssd": 457, + "hits": 1471, + "misses": 457, + "cluster_loads": 544, + "prefetch_loads": 544, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 23 + }, + "spec-planned": { + "s": 1.339740791998338, + "ssd": 259, + "hits": 1640, + "misses": 259, + "cluster_loads": 291, + "prefetch_loads": 291, + "planned_evictions": 275, + "eviction_policy": "planned", + "verify_ready_rate": 0.5917, + "n_tokens": 23 + } + }, + { + "index": 14, + "plain": { + "s": 1.5357000829972094, + "ssd": 429, + "hits": 1438, + "misses": 429, + "cluster_loads": 456, + "prefetch_loads": 456, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 32 + }, + "planned": { + "s": 1.5111392090038862, + "ssd": 269, + "hits": 1612, + "misses": 269, + "cluster_loads": 327, + "prefetch_loads": 327, + "planned_evictions": 315, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 32 + }, + "spec-window": { + "s": 2.233119500000612, + "ssd": 558, + "hits": 2179, + "misses": 558, + "cluster_loads": 690, + "prefetch_loads": 690, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 32 + }, + "spec-planned": { + "s": 1.8321951249963604, + "ssd": 388, + "hits": 2313, + "misses": 388, + "cluster_loads": 434, + "prefetch_loads": 434, + "planned_evictions": 420, + "eviction_policy": "planned", + "verify_ready_rate": 0.6872, + "n_tokens": 32 + } + }, + { + "index": 15, + "plain": { + "s": 1.8883872500009602, + "ssd": 533, + "hits": 1808, + "misses": 533, + "cluster_loads": 575, + "prefetch_loads": 575, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": null, + "n_tokens": 40 + }, + "planned": { + "s": 1.9698821249912726, + "ssd": 362, + "hits": 1997, + "misses": 362, + "cluster_loads": 449, + "prefetch_loads": 449, + "planned_evictions": 437, + "eviction_policy": "planned", + "verify_ready_rate": null, + "n_tokens": 40 + }, + "spec-window": { + "s": 2.7161081660015043, + "ssd": 787, + "hits": 2707, + "misses": 787, + "cluster_loads": 987, + "prefetch_loads": 987, + "planned_evictions": 0, + "eviction_policy": "lru", + "verify_ready_rate": 0.0, + "n_tokens": 40 + }, + "spec-planned": { + "s": 2.11298083300062, + "ssd": 557, + "hits": 2886, + "misses": 557, + "cluster_loads": 633, + "prefetch_loads": 633, + "planned_evictions": 616, + "eviction_policy": "planned", + "verify_ready_rate": 0.5542, + "n_tokens": 40 + } + } + ] +} \ No newline at end of file diff --git a/reports/residency_20261107/cluster_locality_ab_tiny.json b/reports/residency_20261107/cluster_locality_ab_tiny.json new file mode 100644 index 0000000..159fc19 --- /dev/null +++ b/reports/residency_20261107/cluster_locality_ab_tiny.json @@ -0,0 +1,226 @@ +{ + "label": "measured", + "time": "2026-11-07T02:31:36", + "settings": { + "topics": 4, + "sub_vocab": 24, + "experts": 16, + "clusters": 4, + "hidden": 64, + "seq": 128, + "segment": 32, + "batch": 8, + "steps": 300, + "lr": 0.003, + "lam": 0.3, + "hot": 4, + "warm": 8, + "decode": 48, + "prompts": 8, + "seed": 7, + "out": "/tmp/chronos_scratch/cluster_ab.json" + }, + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; segments of 32 tokens, one topic each", + "arms": { + "baseline": { + "lambda_cluster_locality": 0.0, + "train_seconds": 39.36, + "curve": [ + { + "step": 0, + "loss": 5.7093, + "ce": 4.8183, + "cluster_term": 0.0 + }, + { + "step": 30, + "loss": 2.699, + "ce": 2.2045, + "cluster_term": 0.0 + }, + { + "step": 60, + "loss": 2.2713, + "ce": 1.8351, + "cluster_term": 0.0 + }, + { + "step": 90, + "loss": 2.1471, + "ce": 1.7156, + "cluster_term": 0.0 + }, + { + "step": 120, + "loss": 2.2024, + "ce": 1.7713, + "cluster_term": 0.0 + }, + { + "step": 150, + "loss": 2.1309, + "ce": 1.7002, + "cluster_term": 0.0 + }, + { + "step": 180, + "loss": 2.1575, + "ce": 1.74, + "cluster_term": 0.0 + }, + { + "step": 210, + "loss": 2.1656, + "ce": 1.748, + "cluster_term": 0.0 + }, + { + "step": 240, + "loss": 2.1477, + "ce": 1.7286, + "cluster_term": 0.0 + }, + { + "step": 270, + "loss": 2.0949, + "ce": 1.6715, + "cluster_term": 0.0 + }, + { + "step": 299, + "loss": 2.0888, + "ce": 1.6615, + "cluster_term": 0.0 + } + ], + "routing": { + "heldout_ce": 1.7034, + "cluster_window_entropy": 0.4883, + "cluster_marginal_entropy": 0.7105, + "clusters_per_window": 4.0, + "topic_cluster_purity": 0.8074, + "expert_usage_entropy": 0.8739, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 533, + "cluster_loads": 569, + "seconds": 4.606, + "expert_misses_per_token": 1.388, + "cluster_reads_per_token": 1.4818 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 368, + "cluster_loads": 428, + "seconds": 4.011, + "expert_misses_per_token": 0.9583, + "cluster_reads_per_token": 1.1146 + }, + "planned_identical_to_window": true + } + }, + "cluster": { + "lambda_cluster_locality": 0.3, + "train_seconds": 40.53, + "curve": [ + { + "step": 0, + "loss": 6.0084, + "ce": 4.8183, + "cluster_term": 0.29905 + }, + { + "step": 30, + "loss": 2.6347, + "ce": 2.1431, + "cluster_term": 0.04581 + }, + { + "step": 60, + "loss": 2.2721, + "ce": 1.8685, + "cluster_term": 0.01446 + }, + { + "step": 90, + "loss": 2.186, + "ce": 1.7711, + "cluster_term": 0.0239 + }, + { + "step": 120, + "loss": 2.1266, + "ce": 1.7495, + "cluster_term": 0.01283 + }, + { + "step": 150, + "loss": 2.0976, + "ce": 1.7335, + "cluster_term": 0.0125 + }, + { + "step": 180, + "loss": 2.137, + "ce": 1.7468, + "cluster_term": 0.0198 + }, + { + "step": 210, + "loss": 2.1371, + "ce": 1.7592, + "cluster_term": 0.01066 + }, + { + "step": 240, + "loss": 2.1124, + "ce": 1.7393, + "cluster_term": 0.03342 + }, + { + "step": 270, + "loss": 2.0276, + "ce": 1.6778, + "cluster_term": 0.00998 + }, + { + "step": 299, + "loss": 2.0718, + "ce": 1.6619, + "cluster_term": 0.04451 + } + ], + "routing": { + "heldout_ce": 1.7116, + "cluster_window_entropy": 0.0054, + "cluster_marginal_entropy": 0.993, + "clusters_per_window": 1.422, + "topic_cluster_purity": 0.999, + "expert_usage_entropy": 0.8709, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 9, + "cluster_loads": 9, + "seconds": 3.514, + "expert_misses_per_token": 0.0234, + "cluster_reads_per_token": 0.0234 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 9, + "cluster_loads": 9, + "seconds": 3.161, + "expert_misses_per_token": 0.0234, + "cluster_reads_per_token": 0.0234 + }, + "planned_identical_to_window": true + } + } + } +} \ No newline at end of file diff --git a/reports/residency_20261107/sweep/lam_0.03_w32.json b/reports/residency_20261107/sweep/lam_0.03_w32.json new file mode 100644 index 0000000..812e054 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.03_w32.json @@ -0,0 +1,226 @@ +{ + "label": "measured", + "time": "2026-11-07T08:57:35", + "settings": { + "topics": 4, + "sub_vocab": 24, + "experts": 16, + "clusters": 4, + "hidden": 64, + "seq": 128, + "segment": 32, + "batch": 8, + "steps": 300, + "lr": 0.003, + "lam": 0.03, + "hot": 4, + "warm": 8, + "decode": 48, + "prompts": 8, + "seed": 7, + "out": "reports/residency_20261107/sweep/lam_0.03_w32.json" + }, + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; segments of 32 tokens, one topic each", + "arms": { + "baseline": { + "lambda_cluster_locality": 0.0, + "train_seconds": 36.97, + "curve": [ + { + "step": 0, + "loss": 5.7093, + "ce": 4.8183, + "cluster_term": 0.0 + }, + { + "step": 30, + "loss": 2.699, + "ce": 2.2045, + "cluster_term": 0.0 + }, + { + "step": 60, + "loss": 2.2713, + "ce": 1.8351, + "cluster_term": 0.0 + }, + { + "step": 90, + "loss": 2.1471, + "ce": 1.7156, + "cluster_term": 0.0 + }, + { + "step": 120, + "loss": 2.2024, + "ce": 1.7713, + "cluster_term": 0.0 + }, + { + "step": 150, + "loss": 2.1309, + "ce": 1.7002, + "cluster_term": 0.0 + }, + { + "step": 180, + "loss": 2.1575, + "ce": 1.74, + "cluster_term": 0.0 + }, + { + "step": 210, + "loss": 2.1656, + "ce": 1.748, + "cluster_term": 0.0 + }, + { + "step": 240, + "loss": 2.1477, + "ce": 1.7286, + "cluster_term": 0.0 + }, + { + "step": 270, + "loss": 2.0949, + "ce": 1.6715, + "cluster_term": 0.0 + }, + { + "step": 299, + "loss": 2.0888, + "ce": 1.6615, + "cluster_term": 0.0 + } + ], + "routing": { + "heldout_ce": 1.7034, + "cluster_window_entropy": 0.4883, + "cluster_marginal_entropy": 0.7105, + "clusters_per_window": 4.0, + "topic_cluster_purity": 0.8074, + "expert_usage_entropy": 0.8739, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 533, + "cluster_loads": 569, + "seconds": 4.06, + "expert_misses_per_token": 1.388, + "cluster_reads_per_token": 1.4818 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 368, + "cluster_loads": 428, + "seconds": 3.845, + "expert_misses_per_token": 0.9583, + "cluster_reads_per_token": 1.1146 + }, + "planned_identical_to_window": true + } + }, + "cluster": { + "lambda_cluster_locality": 0.03, + "train_seconds": 39.17, + "curve": [ + { + "step": 0, + "loss": 5.7392, + "ce": 4.8183, + "cluster_term": 0.0299 + }, + { + "step": 30, + "loss": 2.7107, + "ce": 2.2036, + "cluster_term": 0.01597 + }, + { + "step": 60, + "loss": 2.3892, + "ce": 1.9457, + "cluster_term": 0.01212 + }, + { + "step": 90, + "loss": 2.2087, + "ce": 1.773, + "cluster_term": 0.01174 + }, + { + "step": 120, + "loss": 2.1906, + "ce": 1.7711, + "cluster_term": 0.01078 + }, + { + "step": 150, + "loss": 2.1079, + "ce": 1.6981, + "cluster_term": 0.01051 + }, + { + "step": 180, + "loss": 2.1562, + "ce": 1.7441, + "cluster_term": 0.01042 + }, + { + "step": 210, + "loss": 2.1631, + "ce": 1.7572, + "cluster_term": 0.00952 + }, + { + "step": 240, + "loss": 2.1311, + "ce": 1.7359, + "cluster_term": 0.00991 + }, + { + "step": 270, + "loss": 2.0735, + "ce": 1.6829, + "cluster_term": 0.00883 + }, + { + "step": 299, + "loss": 2.0591, + "ce": 1.6517, + "cluster_term": 0.01303 + } + ], + "routing": { + "heldout_ce": 1.7089, + "cluster_window_entropy": 0.2782, + "cluster_marginal_entropy": 0.9959, + "clusters_per_window": 3.922, + "topic_cluster_purity": 0.8981, + "expert_usage_entropy": 0.9212, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 184, + "cluster_loads": 197, + "seconds": 3.659, + "expert_misses_per_token": 0.4792, + "cluster_reads_per_token": 0.513 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 115, + "cluster_loads": 122, + "seconds": 3.5, + "expert_misses_per_token": 0.2995, + "cluster_reads_per_token": 0.3177 + }, + "planned_identical_to_window": true + } + } + } +} \ No newline at end of file diff --git a/reports/residency_20261107/sweep/lam_0.03_w32.log b/reports/residency_20261107/sweep/lam_0.03_w32.log new file mode 100644 index 0000000..be09a52 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.03_w32.log @@ -0,0 +1,32 @@ +Skipping import of cpp extensions due to incompatible torch version. Please upgrade to torch >= 2.11.0 (found 2.10.0). +W1107 08:57:38.444000 84347 site-packages/torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs. +[baseline] lambda_cluster_locality=0.0 + step 0 loss 5.7093 ce 4.8183 cluster 0.00000 + step 30 loss 2.6990 ce 2.2045 cluster 0.00000 + step 60 loss 2.2713 ce 1.8351 cluster 0.00000 + step 90 loss 2.1471 ce 1.7156 cluster 0.00000 + step 120 loss 2.2024 ce 1.7713 cluster 0.00000 + step 150 loss 2.1309 ce 1.7002 cluster 0.00000 + step 180 loss 2.1575 ce 1.7400 cluster 0.00000 + step 210 loss 2.1656 ce 1.7480 cluster 0.00000 + step 240 loss 2.1477 ce 1.7286 cluster 0.00000 + step 270 loss 2.0949 ce 1.6715 cluster 0.00000 + step 299 loss 2.0888 ce 1.6615 cluster 0.00000 +[baseline] heldout_ce 1.7034 cluster_window_entropy 0.4883 cluster_marginal_entropy 0.7105 clusters_per_window 4.0 topic_cluster_purity 0.8074 expert_usage_entropy 0.8739 num_clusters 4 +[baseline] lazy window expert misses/token 1.388 cluster reads/token 1.482 tokens 384 +[baseline] lazy planned expert misses/token 0.958 cluster reads/token 1.115 tokens 384 +[cluster] lambda_cluster_locality=0.03 + step 0 loss 5.7392 ce 4.8183 cluster 0.02990 + step 30 loss 2.7107 ce 2.2036 cluster 0.01597 + step 60 loss 2.3892 ce 1.9457 cluster 0.01212 + step 90 loss 2.2087 ce 1.7730 cluster 0.01174 + step 120 loss 2.1906 ce 1.7711 cluster 0.01078 + step 150 loss 2.1079 ce 1.6981 cluster 0.01051 + step 180 loss 2.1562 ce 1.7441 cluster 0.01042 + step 210 loss 2.1631 ce 1.7572 cluster 0.00952 + step 240 loss 2.1311 ce 1.7359 cluster 0.00991 + step 270 loss 2.0735 ce 1.6829 cluster 0.00883 + step 299 loss 2.0591 ce 1.6517 cluster 0.01303 +[cluster] heldout_ce 1.7089 cluster_window_entropy 0.2782 cluster_marginal_entropy 0.9959 clusters_per_window 3.922 topic_cluster_purity 0.8981 expert_usage_entropy 0.9212 num_clusters 4 +[cluster] lazy window expert misses/token 0.479 cluster reads/token 0.513 tokens 384 +[cluster] lazy planned expert misses/token 0.299 cluster reads/token 0.318 tokens 384 diff --git a/reports/residency_20261107/sweep/lam_0.1_w32.json b/reports/residency_20261107/sweep/lam_0.1_w32.json new file mode 100644 index 0000000..8d06f54 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.1_w32.json @@ -0,0 +1,226 @@ +{ + "label": "measured", + "time": "2026-11-07T08:59:15", + "settings": { + "topics": 4, + "sub_vocab": 24, + "experts": 16, + "clusters": 4, + "hidden": 64, + "seq": 128, + "segment": 32, + "batch": 8, + "steps": 300, + "lr": 0.003, + "lam": 0.1, + "hot": 4, + "warm": 8, + "decode": 48, + "prompts": 8, + "seed": 7, + "out": "reports/residency_20261107/sweep/lam_0.1_w32.json" + }, + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; segments of 32 tokens, one topic each", + "arms": { + "baseline": { + "lambda_cluster_locality": 0.0, + "train_seconds": 38.17, + "curve": [ + { + "step": 0, + "loss": 5.7093, + "ce": 4.8183, + "cluster_term": 0.0 + }, + { + "step": 30, + "loss": 2.699, + "ce": 2.2045, + "cluster_term": 0.0 + }, + { + "step": 60, + "loss": 2.2713, + "ce": 1.8351, + "cluster_term": 0.0 + }, + { + "step": 90, + "loss": 2.1471, + "ce": 1.7156, + "cluster_term": 0.0 + }, + { + "step": 120, + "loss": 2.2024, + "ce": 1.7713, + "cluster_term": 0.0 + }, + { + "step": 150, + "loss": 2.1309, + "ce": 1.7002, + "cluster_term": 0.0 + }, + { + "step": 180, + "loss": 2.1575, + "ce": 1.74, + "cluster_term": 0.0 + }, + { + "step": 210, + "loss": 2.1656, + "ce": 1.748, + "cluster_term": 0.0 + }, + { + "step": 240, + "loss": 2.1477, + "ce": 1.7286, + "cluster_term": 0.0 + }, + { + "step": 270, + "loss": 2.0949, + "ce": 1.6715, + "cluster_term": 0.0 + }, + { + "step": 299, + "loss": 2.0888, + "ce": 1.6615, + "cluster_term": 0.0 + } + ], + "routing": { + "heldout_ce": 1.7034, + "cluster_window_entropy": 0.4883, + "cluster_marginal_entropy": 0.7105, + "clusters_per_window": 4.0, + "topic_cluster_purity": 0.8074, + "expert_usage_entropy": 0.8739, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 533, + "cluster_loads": 569, + "seconds": 4.633, + "expert_misses_per_token": 1.388, + "cluster_reads_per_token": 1.4818 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 368, + "cluster_loads": 428, + "seconds": 4.345, + "expert_misses_per_token": 0.9583, + "cluster_reads_per_token": 1.1146 + }, + "planned_identical_to_window": true + } + }, + "cluster": { + "lambda_cluster_locality": 0.1, + "train_seconds": 40.51, + "curve": [ + { + "step": 0, + "loss": 5.809, + "ce": 4.8183, + "cluster_term": 0.09968 + }, + { + "step": 30, + "loss": 2.6409, + "ce": 2.155, + "cluster_term": 0.04694 + }, + { + "step": 60, + "loss": 2.3039, + "ce": 1.8898, + "cluster_term": 0.03667 + }, + { + "step": 90, + "loss": 2.1818, + "ce": 1.7484, + "cluster_term": 0.04719 + }, + { + "step": 120, + "loss": 2.1637, + "ce": 1.7678, + "cluster_term": 0.04002 + }, + { + "step": 150, + "loss": 2.1021, + "ce": 1.7184, + "cluster_term": 0.03471 + }, + { + "step": 180, + "loss": 2.1399, + "ce": 1.7409, + "cluster_term": 0.03846 + }, + { + "step": 210, + "loss": 2.152, + "ce": 1.7567, + "cluster_term": 0.04088 + }, + { + "step": 240, + "loss": 2.0977, + "ce": 1.7378, + "cluster_term": 0.0378 + }, + { + "step": 270, + "loss": 2.0252, + "ce": 1.6751, + "cluster_term": 0.03228 + }, + { + "step": 299, + "loss": 2.0511, + "ce": 1.665, + "cluster_term": 0.0372 + } + ], + "routing": { + "heldout_ce": 1.7129, + "cluster_window_entropy": 0.1508, + "cluster_marginal_entropy": 0.8669, + "clusters_per_window": 3.25, + "topic_cluster_purity": 0.9338, + "expert_usage_entropy": 0.8691, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 58, + "cluster_loads": 58, + "seconds": 3.366, + "expert_misses_per_token": 0.151, + "cluster_reads_per_token": 0.151 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 48, + "cluster_loads": 50, + "seconds": 3.25, + "expert_misses_per_token": 0.125, + "cluster_reads_per_token": 0.1302 + }, + "planned_identical_to_window": true + } + } + } +} \ No newline at end of file diff --git a/reports/residency_20261107/sweep/lam_0.1_w32.log b/reports/residency_20261107/sweep/lam_0.1_w32.log new file mode 100644 index 0000000..3392bd2 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.1_w32.log @@ -0,0 +1,32 @@ +Skipping import of cpp extensions due to incompatible torch version. Please upgrade to torch >= 2.11.0 (found 2.10.0). +W1107 08:59:18.647000 84494 site-packages/torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs. +[baseline] lambda_cluster_locality=0.0 + step 0 loss 5.7093 ce 4.8183 cluster 0.00000 + step 30 loss 2.6990 ce 2.2045 cluster 0.00000 + step 60 loss 2.2713 ce 1.8351 cluster 0.00000 + step 90 loss 2.1471 ce 1.7156 cluster 0.00000 + step 120 loss 2.2024 ce 1.7713 cluster 0.00000 + step 150 loss 2.1309 ce 1.7002 cluster 0.00000 + step 180 loss 2.1575 ce 1.7400 cluster 0.00000 + step 210 loss 2.1656 ce 1.7480 cluster 0.00000 + step 240 loss 2.1477 ce 1.7286 cluster 0.00000 + step 270 loss 2.0949 ce 1.6715 cluster 0.00000 + step 299 loss 2.0888 ce 1.6615 cluster 0.00000 +[baseline] heldout_ce 1.7034 cluster_window_entropy 0.4883 cluster_marginal_entropy 0.7105 clusters_per_window 4.0 topic_cluster_purity 0.8074 expert_usage_entropy 0.8739 num_clusters 4 +[baseline] lazy window expert misses/token 1.388 cluster reads/token 1.482 tokens 384 +[baseline] lazy planned expert misses/token 0.958 cluster reads/token 1.115 tokens 384 +[cluster] lambda_cluster_locality=0.1 + step 0 loss 5.8090 ce 4.8183 cluster 0.09968 + step 30 loss 2.6409 ce 2.1550 cluster 0.04694 + step 60 loss 2.3039 ce 1.8898 cluster 0.03667 + step 90 loss 2.1818 ce 1.7484 cluster 0.04719 + step 120 loss 2.1637 ce 1.7678 cluster 0.04002 + step 150 loss 2.1021 ce 1.7184 cluster 0.03471 + step 180 loss 2.1399 ce 1.7409 cluster 0.03846 + step 210 loss 2.1520 ce 1.7567 cluster 0.04088 + step 240 loss 2.0977 ce 1.7378 cluster 0.03780 + step 270 loss 2.0252 ce 1.6751 cluster 0.03228 + step 299 loss 2.0511 ce 1.6650 cluster 0.03720 +[cluster] heldout_ce 1.7129 cluster_window_entropy 0.1508 cluster_marginal_entropy 0.8669 clusters_per_window 3.25 topic_cluster_purity 0.9338 expert_usage_entropy 0.8691 num_clusters 4 +[cluster] lazy window expert misses/token 0.151 cluster reads/token 0.151 tokens 384 +[cluster] lazy planned expert misses/token 0.125 cluster reads/token 0.130 tokens 384 diff --git a/reports/residency_20261107/sweep/lam_0.3_w0.json b/reports/residency_20261107/sweep/lam_0.3_w0.json new file mode 100644 index 0000000..c954ffe --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.3_w0.json @@ -0,0 +1,227 @@ +{ + "label": "measured", + "time": "2026-11-07T09:09:40", + "settings": { + "topics": 4, + "sub_vocab": 24, + "experts": 16, + "clusters": 4, + "hidden": 64, + "seq": 128, + "segment": 32, + "window": 0, + "batch": 8, + "steps": 300, + "lr": 0.003, + "lam": 0.3, + "hot": 4, + "warm": 8, + "decode": 48, + "prompts": 8, + "seed": 7, + "out": "reports/residency_20261107/sweep/lam_0.3_w0.json" + }, + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; segments of 32 tokens, one topic each", + "arms": { + "baseline": { + "lambda_cluster_locality": 0.0, + "train_seconds": 41.9, + "curve": [ + { + "step": 0, + "loss": 5.7093, + "ce": 4.8183, + "cluster_term": 0.0 + }, + { + "step": 30, + "loss": 2.699, + "ce": 2.2045, + "cluster_term": 0.0 + }, + { + "step": 60, + "loss": 2.2713, + "ce": 1.8351, + "cluster_term": 0.0 + }, + { + "step": 90, + "loss": 2.1471, + "ce": 1.7156, + "cluster_term": 0.0 + }, + { + "step": 120, + "loss": 2.2024, + "ce": 1.7713, + "cluster_term": 0.0 + }, + { + "step": 150, + "loss": 2.1309, + "ce": 1.7002, + "cluster_term": 0.0 + }, + { + "step": 180, + "loss": 2.1575, + "ce": 1.74, + "cluster_term": 0.0 + }, + { + "step": 210, + "loss": 2.1656, + "ce": 1.748, + "cluster_term": 0.0 + }, + { + "step": 240, + "loss": 2.1477, + "ce": 1.7286, + "cluster_term": 0.0 + }, + { + "step": 270, + "loss": 2.0949, + "ce": 1.6715, + "cluster_term": 0.0 + }, + { + "step": 299, + "loss": 2.0888, + "ce": 1.6615, + "cluster_term": 0.0 + } + ], + "routing": { + "heldout_ce": 1.7034, + "cluster_window_entropy": 0.4883, + "cluster_marginal_entropy": 0.7105, + "clusters_per_window": 4.0, + "topic_cluster_purity": 0.8074, + "expert_usage_entropy": 0.8739, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 533, + "cluster_loads": 569, + "seconds": 4.976, + "expert_misses_per_token": 1.388, + "cluster_reads_per_token": 1.4818 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 369, + "cluster_loads": 429, + "seconds": 4.712, + "expert_misses_per_token": 0.9609, + "cluster_reads_per_token": 1.1172 + }, + "planned_identical_to_window": true + } + }, + "cluster": { + "lambda_cluster_locality": 0.3, + "train_seconds": 42.92, + "curve": [ + { + "step": 0, + "loss": 6.0085, + "ce": 4.8183, + "cluster_term": 0.2992 + }, + { + "step": 30, + "loss": 2.8483, + "ce": 2.1912, + "cluster_term": 0.23548 + }, + { + "step": 60, + "loss": 2.5009, + "ce": 1.8866, + "cluster_term": 0.23418 + }, + { + "step": 90, + "loss": 2.366, + "ce": 1.7525, + "cluster_term": 0.24937 + }, + { + "step": 120, + "loss": 2.3545, + "ce": 1.7693, + "cluster_term": 0.2361 + }, + { + "step": 150, + "loss": 2.2689, + "ce": 1.7137, + "cluster_term": 0.21901 + }, + { + "step": 180, + "loss": 2.3383, + "ce": 1.7517, + "cluster_term": 0.24165 + }, + { + "step": 210, + "loss": 2.2916, + "ce": 1.7544, + "cluster_term": 0.208 + }, + { + "step": 240, + "loss": 2.3199, + "ce": 1.7452, + "cluster_term": 0.24055 + }, + { + "step": 270, + "loss": 2.1678, + "ce": 1.6755, + "cluster_term": 0.17658 + }, + { + "step": 299, + "loss": 2.1693, + "ce": 1.6665, + "cluster_term": 0.19042 + } + ], + "routing": { + "heldout_ce": 1.7173, + "cluster_window_entropy": 0.0282, + "cluster_marginal_entropy": 0.9934, + "clusters_per_window": 3.234, + "topic_cluster_purity": 0.9943, + "expert_usage_entropy": 0.9028, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 43, + "cluster_loads": 43, + "seconds": 3.812, + "expert_misses_per_token": 0.112, + "cluster_reads_per_token": 0.112 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 33, + "cluster_loads": 33, + "seconds": 3.616, + "expert_misses_per_token": 0.0859, + "cluster_reads_per_token": 0.0859 + }, + "planned_identical_to_window": true + } + } + } +} \ No newline at end of file diff --git a/reports/residency_20261107/sweep/lam_0.3_w0.log b/reports/residency_20261107/sweep/lam_0.3_w0.log new file mode 100644 index 0000000..10d952b --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.3_w0.log @@ -0,0 +1,32 @@ +Skipping import of cpp extensions due to incompatible torch version. Please upgrade to torch >= 2.11.0 (found 2.10.0). +W1107 09:09:46.508000 85889 site-packages/torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs. +[baseline] lambda_cluster_locality=0.0 + step 0 loss 5.7093 ce 4.8183 cluster 0.00000 + step 30 loss 2.6990 ce 2.2045 cluster 0.00000 + step 60 loss 2.2713 ce 1.8351 cluster 0.00000 + step 90 loss 2.1471 ce 1.7156 cluster 0.00000 + step 120 loss 2.2024 ce 1.7713 cluster 0.00000 + step 150 loss 2.1309 ce 1.7002 cluster 0.00000 + step 180 loss 2.1575 ce 1.7400 cluster 0.00000 + step 210 loss 2.1656 ce 1.7480 cluster 0.00000 + step 240 loss 2.1477 ce 1.7286 cluster 0.00000 + step 270 loss 2.0949 ce 1.6715 cluster 0.00000 + step 299 loss 2.0888 ce 1.6615 cluster 0.00000 +[baseline] heldout_ce 1.7034 cluster_window_entropy 0.4883 cluster_marginal_entropy 0.7105 clusters_per_window 4.0 topic_cluster_purity 0.8074 expert_usage_entropy 0.8739 num_clusters 4 +[baseline] lazy window expert misses/token 1.388 cluster reads/token 1.482 tokens 384 +[baseline] lazy planned expert misses/token 0.961 cluster reads/token 1.117 tokens 384 +[cluster] lambda_cluster_locality=0.3 + step 0 loss 6.0085 ce 4.8183 cluster 0.29920 + step 30 loss 2.8483 ce 2.1912 cluster 0.23548 + step 60 loss 2.5009 ce 1.8866 cluster 0.23418 + step 90 loss 2.3660 ce 1.7525 cluster 0.24937 + step 120 loss 2.3545 ce 1.7693 cluster 0.23610 + step 150 loss 2.2689 ce 1.7137 cluster 0.21901 + step 180 loss 2.3383 ce 1.7517 cluster 0.24165 + step 210 loss 2.2916 ce 1.7544 cluster 0.20800 + step 240 loss 2.3199 ce 1.7452 cluster 0.24055 + step 270 loss 2.1678 ce 1.6755 cluster 0.17658 + step 299 loss 2.1693 ce 1.6665 cluster 0.19042 +[cluster] heldout_ce 1.7173 cluster_window_entropy 0.0282 cluster_marginal_entropy 0.9934 clusters_per_window 3.234 topic_cluster_purity 0.9943 expert_usage_entropy 0.9028 num_clusters 4 +[cluster] lazy window expert misses/token 0.112 cluster reads/token 0.112 tokens 384 +[cluster] lazy planned expert misses/token 0.086 cluster reads/token 0.086 tokens 384 diff --git a/reports/residency_20261107/sweep/lam_0.3_w16.json b/reports/residency_20261107/sweep/lam_0.3_w16.json new file mode 100644 index 0000000..1a981c1 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.3_w16.json @@ -0,0 +1,227 @@ +{ + "label": "measured", + "time": "2026-11-07T09:05:49", + "settings": { + "topics": 4, + "sub_vocab": 24, + "experts": 16, + "clusters": 4, + "hidden": 64, + "seq": 128, + "segment": 32, + "window": 16, + "batch": 8, + "steps": 300, + "lr": 0.003, + "lam": 0.3, + "hot": 4, + "warm": 8, + "decode": 48, + "prompts": 8, + "seed": 7, + "out": "reports/residency_20261107/sweep/lam_0.3_w16.json" + }, + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; segments of 32 tokens, one topic each", + "arms": { + "baseline": { + "lambda_cluster_locality": 0.0, + "train_seconds": 40.64, + "curve": [ + { + "step": 0, + "loss": 5.7093, + "ce": 4.8183, + "cluster_term": 0.0 + }, + { + "step": 30, + "loss": 2.699, + "ce": 2.2045, + "cluster_term": 0.0 + }, + { + "step": 60, + "loss": 2.2713, + "ce": 1.8351, + "cluster_term": 0.0 + }, + { + "step": 90, + "loss": 2.1471, + "ce": 1.7156, + "cluster_term": 0.0 + }, + { + "step": 120, + "loss": 2.2024, + "ce": 1.7713, + "cluster_term": 0.0 + }, + { + "step": 150, + "loss": 2.1309, + "ce": 1.7002, + "cluster_term": 0.0 + }, + { + "step": 180, + "loss": 2.1575, + "ce": 1.74, + "cluster_term": 0.0 + }, + { + "step": 210, + "loss": 2.1656, + "ce": 1.748, + "cluster_term": 0.0 + }, + { + "step": 240, + "loss": 2.1477, + "ce": 1.7286, + "cluster_term": 0.0 + }, + { + "step": 270, + "loss": 2.0949, + "ce": 1.6715, + "cluster_term": 0.0 + }, + { + "step": 299, + "loss": 2.0888, + "ce": 1.6615, + "cluster_term": 0.0 + } + ], + "routing": { + "heldout_ce": 1.7034, + "cluster_window_entropy": 0.4883, + "cluster_marginal_entropy": 0.7105, + "clusters_per_window": 4.0, + "topic_cluster_purity": 0.8074, + "expert_usage_entropy": 0.8739, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 533, + "cluster_loads": 569, + "seconds": 4.67, + "expert_misses_per_token": 1.388, + "cluster_reads_per_token": 1.4818 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 369, + "cluster_loads": 429, + "seconds": 4.808, + "expert_misses_per_token": 0.9609, + "cluster_reads_per_token": 1.1172 + }, + "planned_identical_to_window": true + } + }, + "cluster": { + "lambda_cluster_locality": 0.3, + "train_seconds": 45.91, + "curve": [ + { + "step": 0, + "loss": 6.0083, + "ce": 4.8183, + "cluster_term": 0.29898 + }, + { + "step": 30, + "loss": 2.6376, + "ce": 2.1361, + "cluster_term": 0.04301 + }, + { + "step": 60, + "loss": 2.2429, + "ce": 1.8522, + "cluster_term": 0.0113 + }, + { + "step": 90, + "loss": 2.1596, + "ce": 1.7482, + "cluster_term": 0.02165 + }, + { + "step": 120, + "loss": 2.1236, + "ce": 1.7626, + "cluster_term": 0.01112 + }, + { + "step": 150, + "loss": 2.0536, + "ce": 1.7019, + "cluster_term": 0.00981 + }, + { + "step": 180, + "loss": 2.1449, + "ce": 1.7574, + "cluster_term": 0.02011 + }, + { + "step": 210, + "loss": 2.1236, + "ce": 1.7599, + "cluster_term": 0.0106 + }, + { + "step": 240, + "loss": 2.1066, + "ce": 1.734, + "cluster_term": 0.03241 + }, + { + "step": 270, + "loss": 2.0052, + "ce": 1.6671, + "cluster_term": 0.00998 + }, + { + "step": 299, + "loss": 2.0156, + "ce": 1.6502, + "cluster_term": 0.04441 + } + ], + "routing": { + "heldout_ce": 1.7041, + "cluster_window_entropy": 0.0043, + "cluster_marginal_entropy": 0.9929, + "clusters_per_window": 2.094, + "topic_cluster_purity": 0.9993, + "expert_usage_entropy": 0.892, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 17, + "cluster_loads": 17, + "seconds": 3.859, + "expert_misses_per_token": 0.0443, + "cluster_reads_per_token": 0.0443 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 17, + "cluster_loads": 17, + "seconds": 3.501, + "expert_misses_per_token": 0.0443, + "cluster_reads_per_token": 0.0443 + }, + "planned_identical_to_window": true + } + } + } +} \ No newline at end of file diff --git a/reports/residency_20261107/sweep/lam_0.3_w16.log b/reports/residency_20261107/sweep/lam_0.3_w16.log new file mode 100644 index 0000000..631607a --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.3_w16.log @@ -0,0 +1,32 @@ +Skipping import of cpp extensions due to incompatible torch version. Please upgrade to torch >= 2.11.0 (found 2.10.0). +W1107 09:05:52.285000 85098 site-packages/torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs. +[baseline] lambda_cluster_locality=0.0 + step 0 loss 5.7093 ce 4.8183 cluster 0.00000 + step 30 loss 2.6990 ce 2.2045 cluster 0.00000 + step 60 loss 2.2713 ce 1.8351 cluster 0.00000 + step 90 loss 2.1471 ce 1.7156 cluster 0.00000 + step 120 loss 2.2024 ce 1.7713 cluster 0.00000 + step 150 loss 2.1309 ce 1.7002 cluster 0.00000 + step 180 loss 2.1575 ce 1.7400 cluster 0.00000 + step 210 loss 2.1656 ce 1.7480 cluster 0.00000 + step 240 loss 2.1477 ce 1.7286 cluster 0.00000 + step 270 loss 2.0949 ce 1.6715 cluster 0.00000 + step 299 loss 2.0888 ce 1.6615 cluster 0.00000 +[baseline] heldout_ce 1.7034 cluster_window_entropy 0.4883 cluster_marginal_entropy 0.7105 clusters_per_window 4.0 topic_cluster_purity 0.8074 expert_usage_entropy 0.8739 num_clusters 4 +[baseline] lazy window expert misses/token 1.388 cluster reads/token 1.482 tokens 384 +[baseline] lazy planned expert misses/token 0.961 cluster reads/token 1.117 tokens 384 +[cluster] lambda_cluster_locality=0.3 + step 0 loss 6.0083 ce 4.8183 cluster 0.29898 + step 30 loss 2.6376 ce 2.1361 cluster 0.04301 + step 60 loss 2.2429 ce 1.8522 cluster 0.01130 + step 90 loss 2.1596 ce 1.7482 cluster 0.02165 + step 120 loss 2.1236 ce 1.7626 cluster 0.01112 + step 150 loss 2.0536 ce 1.7019 cluster 0.00981 + step 180 loss 2.1449 ce 1.7574 cluster 0.02011 + step 210 loss 2.1236 ce 1.7599 cluster 0.01060 + step 240 loss 2.1066 ce 1.7340 cluster 0.03241 + step 270 loss 2.0052 ce 1.6671 cluster 0.00998 + step 299 loss 2.0156 ce 1.6502 cluster 0.04441 +[cluster] heldout_ce 1.7041 cluster_window_entropy 0.0043 cluster_marginal_entropy 0.9929 clusters_per_window 2.094 topic_cluster_purity 0.9993 expert_usage_entropy 0.892 num_clusters 4 +[cluster] lazy window expert misses/token 0.044 cluster reads/token 0.044 tokens 384 +[cluster] lazy planned expert misses/token 0.044 cluster reads/token 0.044 tokens 384 diff --git a/reports/residency_20261107/sweep/lam_0.3_w32.json b/reports/residency_20261107/sweep/lam_0.3_w32.json new file mode 100644 index 0000000..ead5518 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.3_w32.json @@ -0,0 +1,226 @@ +{ + "label": "measured", + "time": "2026-11-07T09:00:59", + "settings": { + "topics": 4, + "sub_vocab": 24, + "experts": 16, + "clusters": 4, + "hidden": 64, + "seq": 128, + "segment": 32, + "batch": 8, + "steps": 300, + "lr": 0.003, + "lam": 0.3, + "hot": 4, + "warm": 8, + "decode": 48, + "prompts": 8, + "seed": 7, + "out": "reports/residency_20261107/sweep/lam_0.3_w32.json" + }, + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; segments of 32 tokens, one topic each", + "arms": { + "baseline": { + "lambda_cluster_locality": 0.0, + "train_seconds": 39.13, + "curve": [ + { + "step": 0, + "loss": 5.7093, + "ce": 4.8183, + "cluster_term": 0.0 + }, + { + "step": 30, + "loss": 2.699, + "ce": 2.2045, + "cluster_term": 0.0 + }, + { + "step": 60, + "loss": 2.2713, + "ce": 1.8351, + "cluster_term": 0.0 + }, + { + "step": 90, + "loss": 2.1471, + "ce": 1.7156, + "cluster_term": 0.0 + }, + { + "step": 120, + "loss": 2.2024, + "ce": 1.7713, + "cluster_term": 0.0 + }, + { + "step": 150, + "loss": 2.1309, + "ce": 1.7002, + "cluster_term": 0.0 + }, + { + "step": 180, + "loss": 2.1575, + "ce": 1.74, + "cluster_term": 0.0 + }, + { + "step": 210, + "loss": 2.1656, + "ce": 1.748, + "cluster_term": 0.0 + }, + { + "step": 240, + "loss": 2.1477, + "ce": 1.7286, + "cluster_term": 0.0 + }, + { + "step": 270, + "loss": 2.0949, + "ce": 1.6715, + "cluster_term": 0.0 + }, + { + "step": 299, + "loss": 2.0888, + "ce": 1.6615, + "cluster_term": 0.0 + } + ], + "routing": { + "heldout_ce": 1.7034, + "cluster_window_entropy": 0.4883, + "cluster_marginal_entropy": 0.7105, + "clusters_per_window": 4.0, + "topic_cluster_purity": 0.8074, + "expert_usage_entropy": 0.8739, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 533, + "cluster_loads": 569, + "seconds": 4.147, + "expert_misses_per_token": 1.388, + "cluster_reads_per_token": 1.4818 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 368, + "cluster_loads": 428, + "seconds": 3.806, + "expert_misses_per_token": 0.9583, + "cluster_reads_per_token": 1.1146 + }, + "planned_identical_to_window": true + } + }, + "cluster": { + "lambda_cluster_locality": 0.3, + "train_seconds": 38.06, + "curve": [ + { + "step": 0, + "loss": 6.0084, + "ce": 4.8183, + "cluster_term": 0.29905 + }, + { + "step": 30, + "loss": 2.6347, + "ce": 2.1431, + "cluster_term": 0.04581 + }, + { + "step": 60, + "loss": 2.2721, + "ce": 1.8685, + "cluster_term": 0.01446 + }, + { + "step": 90, + "loss": 2.186, + "ce": 1.7711, + "cluster_term": 0.0239 + }, + { + "step": 120, + "loss": 2.1266, + "ce": 1.7495, + "cluster_term": 0.01283 + }, + { + "step": 150, + "loss": 2.0976, + "ce": 1.7335, + "cluster_term": 0.0125 + }, + { + "step": 180, + "loss": 2.137, + "ce": 1.7468, + "cluster_term": 0.0198 + }, + { + "step": 210, + "loss": 2.1371, + "ce": 1.7592, + "cluster_term": 0.01066 + }, + { + "step": 240, + "loss": 2.1124, + "ce": 1.7393, + "cluster_term": 0.03342 + }, + { + "step": 270, + "loss": 2.0276, + "ce": 1.6778, + "cluster_term": 0.00998 + }, + { + "step": 299, + "loss": 2.0718, + "ce": 1.6619, + "cluster_term": 0.04451 + } + ], + "routing": { + "heldout_ce": 1.7116, + "cluster_window_entropy": 0.0054, + "cluster_marginal_entropy": 0.993, + "clusters_per_window": 1.422, + "topic_cluster_purity": 0.999, + "expert_usage_entropy": 0.8709, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 9, + "cluster_loads": 9, + "seconds": 3.136, + "expert_misses_per_token": 0.0234, + "cluster_reads_per_token": 0.0234 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 9, + "cluster_loads": 9, + "seconds": 2.993, + "expert_misses_per_token": 0.0234, + "cluster_reads_per_token": 0.0234 + }, + "planned_identical_to_window": true + } + } + } +} \ No newline at end of file diff --git a/reports/residency_20261107/sweep/lam_0.3_w32.log b/reports/residency_20261107/sweep/lam_0.3_w32.log new file mode 100644 index 0000000..b578ed5 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.3_w32.log @@ -0,0 +1,32 @@ +Skipping import of cpp extensions due to incompatible torch version. Please upgrade to torch >= 2.11.0 (found 2.10.0). +W1107 09:01:02.027000 84654 site-packages/torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs. +[baseline] lambda_cluster_locality=0.0 + step 0 loss 5.7093 ce 4.8183 cluster 0.00000 + step 30 loss 2.6990 ce 2.2045 cluster 0.00000 + step 60 loss 2.2713 ce 1.8351 cluster 0.00000 + step 90 loss 2.1471 ce 1.7156 cluster 0.00000 + step 120 loss 2.2024 ce 1.7713 cluster 0.00000 + step 150 loss 2.1309 ce 1.7002 cluster 0.00000 + step 180 loss 2.1575 ce 1.7400 cluster 0.00000 + step 210 loss 2.1656 ce 1.7480 cluster 0.00000 + step 240 loss 2.1477 ce 1.7286 cluster 0.00000 + step 270 loss 2.0949 ce 1.6715 cluster 0.00000 + step 299 loss 2.0888 ce 1.6615 cluster 0.00000 +[baseline] heldout_ce 1.7034 cluster_window_entropy 0.4883 cluster_marginal_entropy 0.7105 clusters_per_window 4.0 topic_cluster_purity 0.8074 expert_usage_entropy 0.8739 num_clusters 4 +[baseline] lazy window expert misses/token 1.388 cluster reads/token 1.482 tokens 384 +[baseline] lazy planned expert misses/token 0.958 cluster reads/token 1.115 tokens 384 +[cluster] lambda_cluster_locality=0.3 + step 0 loss 6.0084 ce 4.8183 cluster 0.29905 + step 30 loss 2.6347 ce 2.1431 cluster 0.04581 + step 60 loss 2.2721 ce 1.8685 cluster 0.01446 + step 90 loss 2.1860 ce 1.7711 cluster 0.02390 + step 120 loss 2.1266 ce 1.7495 cluster 0.01283 + step 150 loss 2.0976 ce 1.7335 cluster 0.01250 + step 180 loss 2.1370 ce 1.7468 cluster 0.01980 + step 210 loss 2.1371 ce 1.7592 cluster 0.01066 + step 240 loss 2.1124 ce 1.7393 cluster 0.03342 + step 270 loss 2.0276 ce 1.6778 cluster 0.00998 + step 299 loss 2.0718 ce 1.6619 cluster 0.04451 +[cluster] heldout_ce 1.7116 cluster_window_entropy 0.0054 cluster_marginal_entropy 0.993 clusters_per_window 1.422 topic_cluster_purity 0.999 expert_usage_entropy 0.8709 num_clusters 4 +[cluster] lazy window expert misses/token 0.023 cluster reads/token 0.023 tokens 384 +[cluster] lazy planned expert misses/token 0.023 cluster reads/token 0.023 tokens 384 diff --git a/reports/residency_20261107/sweep/lam_0.3_w64.json b/reports/residency_20261107/sweep/lam_0.3_w64.json new file mode 100644 index 0000000..2372698 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.3_w64.json @@ -0,0 +1,227 @@ +{ + "label": "measured", + "time": "2026-11-07T09:07:42", + "settings": { + "topics": 4, + "sub_vocab": 24, + "experts": 16, + "clusters": 4, + "hidden": 64, + "seq": 128, + "segment": 32, + "window": 64, + "batch": 8, + "steps": 300, + "lr": 0.003, + "lam": 0.3, + "hot": 4, + "warm": 8, + "decode": 48, + "prompts": 8, + "seed": 7, + "out": "reports/residency_20261107/sweep/lam_0.3_w64.json" + }, + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; segments of 32 tokens, one topic each", + "arms": { + "baseline": { + "lambda_cluster_locality": 0.0, + "train_seconds": 44.57, + "curve": [ + { + "step": 0, + "loss": 5.7093, + "ce": 4.8183, + "cluster_term": 0.0 + }, + { + "step": 30, + "loss": 2.699, + "ce": 2.2045, + "cluster_term": 0.0 + }, + { + "step": 60, + "loss": 2.2713, + "ce": 1.8351, + "cluster_term": 0.0 + }, + { + "step": 90, + "loss": 2.1471, + "ce": 1.7156, + "cluster_term": 0.0 + }, + { + "step": 120, + "loss": 2.2024, + "ce": 1.7713, + "cluster_term": 0.0 + }, + { + "step": 150, + "loss": 2.1309, + "ce": 1.7002, + "cluster_term": 0.0 + }, + { + "step": 180, + "loss": 2.1575, + "ce": 1.74, + "cluster_term": 0.0 + }, + { + "step": 210, + "loss": 2.1656, + "ce": 1.748, + "cluster_term": 0.0 + }, + { + "step": 240, + "loss": 2.1477, + "ce": 1.7286, + "cluster_term": 0.0 + }, + { + "step": 270, + "loss": 2.0949, + "ce": 1.6715, + "cluster_term": 0.0 + }, + { + "step": 299, + "loss": 2.0888, + "ce": 1.6615, + "cluster_term": 0.0 + } + ], + "routing": { + "heldout_ce": 1.7034, + "cluster_window_entropy": 0.4883, + "cluster_marginal_entropy": 0.7105, + "clusters_per_window": 4.0, + "topic_cluster_purity": 0.8074, + "expert_usage_entropy": 0.8739, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 533, + "cluster_loads": 569, + "seconds": 5.407, + "expert_misses_per_token": 1.388, + "cluster_reads_per_token": 1.4818 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 369, + "cluster_loads": 429, + "seconds": 4.665, + "expert_misses_per_token": 0.9609, + "cluster_reads_per_token": 1.1172 + }, + "planned_identical_to_window": true + } + }, + "cluster": { + "lambda_cluster_locality": 0.3, + "train_seconds": 42.17, + "curve": [ + { + "step": 0, + "loss": 6.0085, + "ce": 4.8183, + "cluster_term": 0.29913 + }, + { + "step": 30, + "loss": 2.737, + "ce": 2.135, + "cluster_term": 0.17106 + }, + { + "step": 60, + "loss": 2.3557, + "ce": 1.8703, + "cluster_term": 0.12535 + }, + { + "step": 90, + "loss": 2.2283, + "ce": 1.7268, + "cluster_term": 0.14113 + }, + { + "step": 120, + "loss": 2.219, + "ce": 1.7517, + "cluster_term": 0.12312 + }, + { + "step": 150, + "loss": 2.1632, + "ce": 1.6988, + "cluster_term": 0.13186 + }, + { + "step": 180, + "loss": 2.2051, + "ce": 1.7333, + "cluster_term": 0.12294 + }, + { + "step": 210, + "loss": 2.2253, + "ce": 1.7631, + "cluster_term": 0.12291 + }, + { + "step": 240, + "loss": 2.2258, + "ce": 1.7285, + "cluster_term": 0.15393 + }, + { + "step": 270, + "loss": 2.1326, + "ce": 1.673, + "cluster_term": 0.12227 + }, + { + "step": 299, + "loss": 2.1113, + "ce": 1.6562, + "cluster_term": 0.12908 + } + ], + "routing": { + "heldout_ce": 1.7092, + "cluster_window_entropy": 0.0067, + "cluster_marginal_entropy": 0.9929, + "clusters_per_window": 1.953, + "topic_cluster_purity": 0.9989, + "expert_usage_entropy": 0.8772, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 44, + "cluster_loads": 44, + "seconds": 4.307, + "expert_misses_per_token": 0.1146, + "cluster_reads_per_token": 0.1146 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 24, + "cluster_loads": 24, + "seconds": 4.575, + "expert_misses_per_token": 0.0625, + "cluster_reads_per_token": 0.0625 + }, + "planned_identical_to_window": true + } + } + } +} \ No newline at end of file diff --git a/reports/residency_20261107/sweep/lam_0.3_w64.log b/reports/residency_20261107/sweep/lam_0.3_w64.log new file mode 100644 index 0000000..cedf965 --- /dev/null +++ b/reports/residency_20261107/sweep/lam_0.3_w64.log @@ -0,0 +1,32 @@ +Skipping import of cpp extensions due to incompatible torch version. Please upgrade to torch >= 2.11.0 (found 2.10.0). +W1107 09:07:46.245000 85622 site-packages/torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs. +[baseline] lambda_cluster_locality=0.0 + step 0 loss 5.7093 ce 4.8183 cluster 0.00000 + step 30 loss 2.6990 ce 2.2045 cluster 0.00000 + step 60 loss 2.2713 ce 1.8351 cluster 0.00000 + step 90 loss 2.1471 ce 1.7156 cluster 0.00000 + step 120 loss 2.2024 ce 1.7713 cluster 0.00000 + step 150 loss 2.1309 ce 1.7002 cluster 0.00000 + step 180 loss 2.1575 ce 1.7400 cluster 0.00000 + step 210 loss 2.1656 ce 1.7480 cluster 0.00000 + step 240 loss 2.1477 ce 1.7286 cluster 0.00000 + step 270 loss 2.0949 ce 1.6715 cluster 0.00000 + step 299 loss 2.0888 ce 1.6615 cluster 0.00000 +[baseline] heldout_ce 1.7034 cluster_window_entropy 0.4883 cluster_marginal_entropy 0.7105 clusters_per_window 4.0 topic_cluster_purity 0.8074 expert_usage_entropy 0.8739 num_clusters 4 +[baseline] lazy window expert misses/token 1.388 cluster reads/token 1.482 tokens 384 +[baseline] lazy planned expert misses/token 0.961 cluster reads/token 1.117 tokens 384 +[cluster] lambda_cluster_locality=0.3 + step 0 loss 6.0085 ce 4.8183 cluster 0.29913 + step 30 loss 2.7370 ce 2.1350 cluster 0.17106 + step 60 loss 2.3557 ce 1.8703 cluster 0.12535 + step 90 loss 2.2283 ce 1.7268 cluster 0.14113 + step 120 loss 2.2190 ce 1.7517 cluster 0.12312 + step 150 loss 2.1632 ce 1.6988 cluster 0.13186 + step 180 loss 2.2051 ce 1.7333 cluster 0.12294 + step 210 loss 2.2253 ce 1.7631 cluster 0.12291 + step 240 loss 2.2258 ce 1.7285 cluster 0.15393 + step 270 loss 2.1326 ce 1.6730 cluster 0.12227 + step 299 loss 2.1113 ce 1.6562 cluster 0.12908 +[cluster] heldout_ce 1.7092 cluster_window_entropy 0.0067 cluster_marginal_entropy 0.9929 clusters_per_window 1.953 topic_cluster_purity 0.9989 expert_usage_entropy 0.8772 num_clusters 4 +[cluster] lazy window expert misses/token 0.115 cluster reads/token 0.115 tokens 384 +[cluster] lazy planned expert misses/token 0.062 cluster reads/token 0.062 tokens 384 diff --git a/reports/residency_20261107/sweep/lam_1.0_w32.json b/reports/residency_20261107/sweep/lam_1.0_w32.json new file mode 100644 index 0000000..142c96b --- /dev/null +++ b/reports/residency_20261107/sweep/lam_1.0_w32.json @@ -0,0 +1,226 @@ +{ + "label": "measured", + "time": "2026-11-07T09:02:39", + "settings": { + "topics": 4, + "sub_vocab": 24, + "experts": 16, + "clusters": 4, + "hidden": 64, + "seq": 128, + "segment": 32, + "batch": 8, + "steps": 300, + "lr": 0.003, + "lam": 1.0, + "hot": 4, + "warm": 8, + "decode": 48, + "prompts": 8, + "seed": 7, + "out": "reports/residency_20261107/sweep/lam_1.0_w32.json" + }, + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; segments of 32 tokens, one topic each", + "arms": { + "baseline": { + "lambda_cluster_locality": 0.0, + "train_seconds": 38.0, + "curve": [ + { + "step": 0, + "loss": 5.7093, + "ce": 4.8183, + "cluster_term": 0.0 + }, + { + "step": 30, + "loss": 2.699, + "ce": 2.2045, + "cluster_term": 0.0 + }, + { + "step": 60, + "loss": 2.2713, + "ce": 1.8351, + "cluster_term": 0.0 + }, + { + "step": 90, + "loss": 2.1471, + "ce": 1.7156, + "cluster_term": 0.0 + }, + { + "step": 120, + "loss": 2.2024, + "ce": 1.7713, + "cluster_term": 0.0 + }, + { + "step": 150, + "loss": 2.1309, + "ce": 1.7002, + "cluster_term": 0.0 + }, + { + "step": 180, + "loss": 2.1575, + "ce": 1.74, + "cluster_term": 0.0 + }, + { + "step": 210, + "loss": 2.1656, + "ce": 1.748, + "cluster_term": 0.0 + }, + { + "step": 240, + "loss": 2.1477, + "ce": 1.7286, + "cluster_term": 0.0 + }, + { + "step": 270, + "loss": 2.0949, + "ce": 1.6715, + "cluster_term": 0.0 + }, + { + "step": 299, + "loss": 2.0888, + "ce": 1.6615, + "cluster_term": 0.0 + } + ], + "routing": { + "heldout_ce": 1.7034, + "cluster_window_entropy": 0.4883, + "cluster_marginal_entropy": 0.7105, + "clusters_per_window": 4.0, + "topic_cluster_purity": 0.8074, + "expert_usage_entropy": 0.8739, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 533, + "cluster_loads": 569, + "seconds": 4.167, + "expert_misses_per_token": 1.388, + "cluster_reads_per_token": 1.4818 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 368, + "cluster_loads": 428, + "seconds": 4.084, + "expert_misses_per_token": 0.9583, + "cluster_reads_per_token": 1.1146 + }, + "planned_identical_to_window": true + } + }, + "cluster": { + "lambda_cluster_locality": 1.0, + "train_seconds": 40.04, + "curve": [ + { + "step": 0, + "loss": 6.7062, + "ce": 4.8183, + "cluster_term": 0.99683 + }, + { + "step": 30, + "loss": 2.8887, + "ce": 2.4338, + "cluster_term": 0.05808 + }, + { + "step": 60, + "loss": 2.3366, + "ce": 1.9381, + "cluster_term": 0.02302 + }, + { + "step": 90, + "loss": 2.1795, + "ce": 1.7506, + "cluster_term": 0.05562 + }, + { + "step": 120, + "loss": 2.1525, + "ce": 1.7667, + "cluster_term": 0.02868 + }, + { + "step": 150, + "loss": 2.0918, + "ce": 1.7111, + "cluster_term": 0.02724 + }, + { + "step": 180, + "loss": 2.1362, + "ce": 1.7263, + "cluster_term": 0.06234 + }, + { + "step": 210, + "loss": 2.1232, + "ce": 1.7469, + "cluster_term": 0.03004 + }, + { + "step": 240, + "loss": 2.1643, + "ce": 1.7269, + "cluster_term": 0.10323 + }, + { + "step": 270, + "loss": 2.0505, + "ce": 1.6843, + "cluster_term": 0.02998 + }, + { + "step": 299, + "loss": 2.1568, + "ce": 1.6689, + "cluster_term": 0.14572 + } + ], + "routing": { + "heldout_ce": 1.7087, + "cluster_window_entropy": 0.0009, + "cluster_marginal_entropy": 0.9929, + "clusters_per_window": 1.391, + "topic_cluster_purity": 0.9999, + "expert_usage_entropy": 0.897, + "num_clusters": 4 + }, + "lazy_decode": { + "window": { + "tokens": 384, + "sync_ssd_loads": 10, + "cluster_loads": 10, + "seconds": 3.54, + "expert_misses_per_token": 0.026, + "cluster_reads_per_token": 0.026 + }, + "planned": { + "tokens": 384, + "sync_ssd_loads": 10, + "cluster_loads": 10, + "seconds": 3.4, + "expert_misses_per_token": 0.026, + "cluster_reads_per_token": 0.026 + }, + "planned_identical_to_window": true + } + } + } +} \ No newline at end of file diff --git a/reports/residency_20261107/sweep/lam_1.0_w32.log b/reports/residency_20261107/sweep/lam_1.0_w32.log new file mode 100644 index 0000000..79bcb7c --- /dev/null +++ b/reports/residency_20261107/sweep/lam_1.0_w32.log @@ -0,0 +1,32 @@ +Skipping import of cpp extensions due to incompatible torch version. Please upgrade to torch >= 2.11.0 (found 2.10.0). +W1107 09:02:42.439000 84813 site-packages/torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs. +[baseline] lambda_cluster_locality=0.0 + step 0 loss 5.7093 ce 4.8183 cluster 0.00000 + step 30 loss 2.6990 ce 2.2045 cluster 0.00000 + step 60 loss 2.2713 ce 1.8351 cluster 0.00000 + step 90 loss 2.1471 ce 1.7156 cluster 0.00000 + step 120 loss 2.2024 ce 1.7713 cluster 0.00000 + step 150 loss 2.1309 ce 1.7002 cluster 0.00000 + step 180 loss 2.1575 ce 1.7400 cluster 0.00000 + step 210 loss 2.1656 ce 1.7480 cluster 0.00000 + step 240 loss 2.1477 ce 1.7286 cluster 0.00000 + step 270 loss 2.0949 ce 1.6715 cluster 0.00000 + step 299 loss 2.0888 ce 1.6615 cluster 0.00000 +[baseline] heldout_ce 1.7034 cluster_window_entropy 0.4883 cluster_marginal_entropy 0.7105 clusters_per_window 4.0 topic_cluster_purity 0.8074 expert_usage_entropy 0.8739 num_clusters 4 +[baseline] lazy window expert misses/token 1.388 cluster reads/token 1.482 tokens 384 +[baseline] lazy planned expert misses/token 0.958 cluster reads/token 1.115 tokens 384 +[cluster] lambda_cluster_locality=1.0 + step 0 loss 6.7062 ce 4.8183 cluster 0.99683 + step 30 loss 2.8887 ce 2.4338 cluster 0.05808 + step 60 loss 2.3366 ce 1.9381 cluster 0.02302 + step 90 loss 2.1795 ce 1.7506 cluster 0.05562 + step 120 loss 2.1525 ce 1.7667 cluster 0.02868 + step 150 loss 2.0918 ce 1.7111 cluster 0.02724 + step 180 loss 2.1362 ce 1.7263 cluster 0.06234 + step 210 loss 2.1232 ce 1.7469 cluster 0.03004 + step 240 loss 2.1643 ce 1.7269 cluster 0.10323 + step 270 loss 2.0505 ce 1.6843 cluster 0.02998 + step 299 loss 2.1568 ce 1.6689 cluster 0.14572 +[cluster] heldout_ce 1.7087 cluster_window_entropy 0.0009 cluster_marginal_entropy 0.9929 clusters_per_window 1.391 topic_cluster_purity 0.9999 expert_usage_entropy 0.897 num_clusters 4 +[cluster] lazy window expert misses/token 0.026 cluster reads/token 0.026 tokens 384 +[cluster] lazy planned expert misses/token 0.026 cluster reads/token 0.026 tokens 384 diff --git a/tests/test_bundle.py b/tests/test_bundle.py index acb4c6e..4120f70 100644 --- a/tests/test_bundle.py +++ b/tests/test_bundle.py @@ -150,6 +150,34 @@ def test_translate_backends_and_speculative_agree(bundle): mx.set_default_device(prev) +def test_lazy_experts_translate_like_the_full_model(bundle, tmp_path): + """SSD-lazy experts (MLX) give the same greedy output; the cache stays outside the bundle.""" + mx = pytest.importorskip("mlx.core") + from chronos.bundle.runtime import BundleModel + from chronos.lazy_experts import LazyOptions + prev = mx.default_device() + try: + full = BundleModel(bundle, device="cpu", backend="mlx") + src = EN[:2] + ref = full.translate(src, "en_zh", max_new_tokens=16) + for residency in ("planned", "window"): + lazy = BundleModel(bundle, device="cpu", backend="mlx", lazy_experts=True, + lazy_options=LazyOptions(residency=residency, ram_ratio=0.5), + cache_root=str(tmp_path / "cache")) + try: + assert lazy.translate(src, "en_zh", max_new_tokens=16) == ref + assert lazy.translate(src, "en_zh", max_new_tokens=16, speculative=True) == ref + stats = lazy.last_stats["lazy"] + assert stats["paging_resolved"] == "expert" and stats["residency"] == residency + finally: + lazy.close() + assert not os.path.exists(os.path.join(bundle, "expert_cache")) + with pytest.raises(ValueError, match="mlx"): + BundleModel(bundle, device="cpu", backend="torch", lazy_experts=True) + finally: + mx.set_default_device(prev) + + def test_prepare_data_and_finetune_spec(bundle, tmp_path): from chronos.bundle.finetune import prepare_data, write_finetune_spec data = tmp_path / "mine.jsonl" @@ -172,6 +200,9 @@ def test_prepare_data_and_finetune_spec(bundle, tmp_path): assert spec["init_checkpoint"] == os.path.join(bundle, m["model"]["checkpoint"]) assert spec["tokenizer_dir"] == os.path.join(bundle, "model", "tokenizer") assert [p["name"] for p in spec["phases"]] == ["sft"] and spec["export_prefix"] == "tiny_bundle_ft" + # the tiny model has 4 experts with k=2 (one cluster): the constraint is off and says why + assert spec["model_config"]["lambda_cluster_locality"] == 0.0 + assert "fewer than two clusters" in spec["extra"]["cluster_locality_source"]["lambda_cluster_locality"] with pytest.raises(ValueError, match="pretrain"): write_finetune_spec(bundle, str(out / "shards"), str(tmp_path / "ft2"), tier="smoke", mode="pretrain") (run / "checkpoints").mkdir() diff --git a/tests/test_lazy_experts.py b/tests/test_lazy_experts.py new file mode 100644 index 0000000..788effa --- /dev/null +++ b/tests/test_lazy_experts.py @@ -0,0 +1,285 @@ +"""Architecture-adaptive lazy experts (chronos.lazy_experts) and the cluster-locality defaults.""" +from __future__ import annotations + +import json +import os +import sys + +import pytest + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +from chronos.expert_clusters import ( # noqa: E402 + DEFAULT_LAMBDA, inference_clusters, paging_tag, recommended_window, trained_with_clusters, +) +from chronos.lazy_experts import LazyOptions, checkpoint_cache_dir, page_layout # noqa: E402 + + +class Cfg: + num_experts, num_experts_per_tok, cluster_locality_clusters = 8, 2, 0 + lambda_cluster_locality = 0.0 + + +def _cfg(lam): + c = Cfg() + c.lambda_cluster_locality = lam + return c + + +# ── layout selection ────────────────────────────────────────────── + +def test_auto_paging_follows_how_the_model_was_trained(): + trained, plain = _cfg(0.3), _cfg(0.0) + assert trained_with_clusters(trained) and not trained_with_clusters(plain) + assert inference_clusters(trained) == [[0, 1, 2, 3], [4, 5, 6, 7]] # ⌈8 / (2·2)⌉ = 2 clusters + assert inference_clusters(plain) == [[i] for i in range(8)] + assert inference_clusters(plain, "cluster") == inference_clusters(trained) + assert inference_clusters(trained, "expert") == [[i] for i in range(8)] + with pytest.raises(ValueError): + inference_clusters(trained, "louvain") + assert page_layout(trained)["resolved"] == "cluster" and page_layout(plain)["resolved"] == "expert" + assert page_layout(trained)["reason"] == "trained with cluster locality" + + +def test_paging_tag_and_cache_dirs_never_mix_layouts(tmp_path): + assert paging_tag([[i] for i in range(8)]) == "e" + assert paging_tag([[0, 1, 2, 3], [4, 5, 6, 7]]) == "c2" + ckpt = tmp_path / "m.pth" + ckpt.write_bytes(b"x") + cfg = _cfg(0.3) + root = str(tmp_path / "cache") + torch_e = checkpoint_cache_dir(str(ckpt), cfg, root=root) + mlx_e = checkpoint_cache_dir(str(ckpt), cfg, backend="mlx", root=root) + mlx_c = checkpoint_cache_dir(str(ckpt), cfg, backend="mlx", tag="c2", root=root) + assert mlx_e == torch_e + "_mlx" and mlx_c == mlx_e + "_c2" # per-expert names unchanged + assert len({torch_e, mlx_e, mlx_c, checkpoint_cache_dir(str(ckpt), cfg, tag="c2", root=root)}) == 4 + + +def test_recommended_window_from_record_lengths(): + assert recommended_window(None, 512) == 32 + assert recommended_window(91, 512) == 128 + assert recommended_window(5, 512) == 16 + assert recommended_window(4000, 512) == 512 + assert recommended_window(None, 8) == 8 + + +def test_lazy_options_validate(): + assert LazyOptions().validated().generate_kwargs()["residency"] == "planned" + for bad in (dict(paging="x"), dict(residency="queue"), dict(residency_decay=1.0), + dict(decode_wait_s=-1), dict(ram_ratio=0)): + with pytest.raises(ValueError): + LazyOptions(**bad).validated() + + +# ── checkpoints and training defaults ───────────────────────────── + +torch = pytest.importorskip("torch") + + +def _tiny(lam, num_experts=8): + from chronos.model.config import ChronosConfig + return ChronosConfig(hidden_size=64, num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2, + vocab_size=96, num_experts=num_experts, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=32, max_position_embeddings=128, recurrent_depth=2, + prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0, + kv_latent_dim=16, rope_dim=8, sliding_window_size=64, use_moe=True, + lambda_cluster_locality=lam, cluster_locality_window=16) + + +def _save(tmp_path, lam, name): + from chronos.model.checkpoint import save_state_dict_with_config + from chronos.model.model_chronos import ChronosForCausalLM + torch.manual_seed(5) + cfg = _tiny(lam) + path = str(tmp_path / f"{name}.pth") + save_state_dict_with_config(ChronosForCausalLM(cfg).eval(), path, cfg, stage="pretrain") + return path + + +def test_checkpoint_records_lambda_and_legacy_sidecars_load_as_untrained(tmp_path): + from chronos.model.checkpoint import checkpoint_config_path, chronos_config_from_checkpoint + path = _save(tmp_path, 0.3, "trained") + cfg, _ = chronos_config_from_checkpoint(path) + assert cfg.lambda_cluster_locality == pytest.approx(0.3) and page_layout(cfg)["resolved"] == "cluster" + side = checkpoint_config_path(path) + meta = json.load(open(side, encoding="utf-8")) + meta.pop("lambda_cluster_locality", None) # a sidecar written before the field existed + if isinstance(meta.get("config"), dict): + meta["config"].pop("lambda_cluster_locality", None) + json.dump(meta, open(side, "w", encoding="utf-8")) + legacy, _ = chronos_config_from_checkpoint(path) + assert legacy.lambda_cluster_locality == 0.0 and page_layout(legacy)["resolved"] == "expert" + + +def test_class_default_is_off_and_new_config_entry_points_turn_it_on(monkeypatch): + """Stored recipes without the field keep their objective; new configs get the measured default.""" + from chronos.model.config import ChronosConfig + assert ChronosConfig().lambda_cluster_locality == 0.0 + import train_chronos + monkeypatch.setattr(sys, "argv", ["train_chronos.py", "--data_path", "unused.jsonl"]) + assert train_chronos.parse_args().lambda_cluster_locality == DEFAULT_LAMBDA + from ui.presets import MINIMIND_MOE_DEFAULTS + assert MINIMIND_MOE_DEFAULTS["lambda_cluster_locality"] == DEFAULT_LAMBDA + + +def test_finetune_spec_turns_cluster_locality_on_by_default(tmp_path): + from chronos.bundle.recommend import cluster_locality_config + import numpy as np + moe = {"num_experts": 36, "num_experts_per_tok": 2, "hidden_size": 64, "num_hidden_layers": 4} + sft = tmp_path / "shards" / "sft" / "train" + sft.mkdir(parents=True) + # four records of 100 tokens → window 128 (power of two nearest the median) + np.asarray([0, 100, 200, 300, 400], np.int64).tofile(sft / "train.x.00000.idx") + on = cluster_locality_config(moe, str(tmp_path / "shards"), seq_len=512) + assert on["values"] == {"lambda_cluster_locality": DEFAULT_LAMBDA, "cluster_locality_window": 128, + "cluster_locality_clusters": 0} + assert "median SFT record (100 tokens)" in on["source"]["cluster_locality_window"] + assert cluster_locality_config(moe, None, enabled=False)["values"]["lambda_cluster_locality"] == 0.0 + assert cluster_locality_config(moe, None, lambda_cluster=0.1, window=0)["values"][ + "cluster_locality_window"] == 0 + dense = cluster_locality_config(dict(moe, num_experts=4), None) # 4 experts, k=2: one cluster + assert dense["values"]["lambda_cluster_locality"] == 0.0 + + +# ── engines: cluster paging is exact ────────────────────────────── + +PROMPT = [5, 17, 33, 2, 61, 9, 40] + + +def _mlx_cpu(): + mx = pytest.importorskip("mlx.core") + previous = mx.default_device() + mx.set_default_device(mx.cpu) + return mx, previous + + +def _mlx_reference(path, n, penalty=1.0): + """Full-memory MLX greedy (CTRL repetition penalty when != 1).""" + import mlx.core as mx + import numpy as np + from chronos.mlx.model import ChronosMLXModel + from chronos.model.checkpoint import chronos_config_from_checkpoint, load_checkpoint_into_model + from chronos.model.model_chronos import ChronosForCausalLM + cfg, _ = chronos_config_from_checkpoint(path) + pt = ChronosForCausalLM(cfg) + load_checkpoint_into_model(pt, path, map_location="cpu") + model = ChronosMLXModel.from_chronos_pytorch(pt.eval(), cfg, dtype="float32").eval() + seen = np.zeros(int(cfg.vocab_size), dtype=bool) + seen[PROMPT] = True + logits, _, cache = model(mx.array([PROMPT], dtype=mx.int32)) + out = [] + for _ in range(n): + row = logits[0, -1].astype(mx.float32) + if penalty != 1.0: + row = mx.where(mx.array(seen), mx.where(row < 0, row * penalty, row / penalty), row) + nxt = int(mx.argmax(row).item()) + out.append(nxt) + seen[nxt] = True + logits, _, cache = model(mx.array([[nxt]], dtype=mx.int32), cache=cache) + return out + + +def _lazy_generate(path, tmp_path, opts, n=16, **gen): + import mlx.core as mx + from chronos.lazy_experts import build_mlx_lazy_engine, warm_ids + from chronos.model.checkpoint import chronos_config_from_checkpoint + cfg, _ = chronos_config_from_checkpoint(path) + engine, info = build_mlx_lazy_engine(path, cfg, opts, cache_root=str(tmp_path / "cache")) + try: + engine.store.warm_up(warm_ids(cfg, info)) + toks = list(engine.generate(mx.array([PROMPT], dtype=mx.int32), max_new_tokens=n, temperature=0.0, + top_p=1.0, **opts.generate_kwargs(), **gen)) + return toks, info, engine.store.stats() + finally: + engine.stop() + + +@pytest.mark.parametrize("residency", ["planned", "window"]) +def test_mlx_lazy_engine_adapts_paging_and_matches_full_memory_greedy(tmp_path, residency): + pytest.importorskip("mlx.core") + mx, previous = _mlx_cpu() + try: + trained, plain = _save(tmp_path, 0.3, "trained"), _save(tmp_path, 0.0, "plain") + ref = _mlx_reference(trained, 16) + opts = LazyOptions(residency=residency, ram_ratio=0.5).validated() + toks, info, st = _lazy_generate(trained, tmp_path, opts) + assert toks == ref + assert info["paging_resolved"] == "cluster" and info["num_pages"] == 2 + assert info["cache_dir"].endswith("_mlx_c2") and st["num_clusters"] == 2 and st["cluster_loads"] > 0 + toks2, info2, _ = _lazy_generate(trained, tmp_path, opts) # same layout: cache reused + assert toks2 == ref and info2["cache_reused"] + forced = LazyOptions(paging="expert", residency=residency, ram_ratio=0.5).validated() + toks3, info3, st3 = _lazy_generate(trained, tmp_path, forced) + assert toks3 == ref and info3["num_pages"] == 8 and st3["num_clusters"] == 8 + assert info3["cache_dir"] != info["cache_dir"] + toks4, info4, _ = _lazy_generate(plain, tmp_path, opts) # untrained → per expert + assert info4["paging_resolved"] == "expert" and info4["cache_dir"].endswith("_mlx") + assert toks4 == _mlx_reference(plain, 16) + finally: + mx.set_default_device(previous) + + +def test_mlx_engine_repetition_penalty_and_waits(tmp_path): + pytest.importorskip("mlx.core") + mx, previous = _mlx_cpu() + try: + path = _save(tmp_path, 0.3, "trained") + opts = LazyOptions(decode_wait_s=0.0, verify_wait_s=0.0).validated() + toks, _info, _st = _lazy_generate(path, tmp_path, opts, n=12, repetition_penalty=1.3) + assert toks == _mlx_reference(path, 12, penalty=1.3) + with pytest.raises(ValueError): + _lazy_generate(path, tmp_path, opts, n=4, repetition_penalty=1.3, speculative=True) + spec, _i, _s = _lazy_generate(path, tmp_path, opts, n=12, speculative=True, draft_depth=1, gamma=3) + assert spec == _mlx_reference(path, 12) + finally: + mx.set_default_device(previous) + + +@pytest.mark.parametrize("paging", ["auto", "expert"]) +def test_torch_offload_pages_by_cluster_when_trained_and_stays_exact(tmp_path, monkeypatch, paging): + monkeypatch.chdir(tmp_path) + from chronos.model.checkpoint import chronos_config_from_checkpoint + from ui.tabs.inference_tab import _clone_model_cfg, _run_torch_full_dram, _run_torch_inference + path = _save(tmp_path, 0.3, "trained") + cfg, _ = chronos_config_from_checkpoint(path) + mask = [1] * len(PROMPT) + ref, _ = _run_torch_full_dram("cpu", _clone_model_cfg(cfg), path, PROMPT, mask, 12, 1.0, None, do_sample=False) + out, stats = _run_torch_inference("cpu", _clone_model_cfg(cfg), path, PROMPT, mask, 12, 1.0, None, + do_sample=False, ram_load_ratio=0.5, paging=paging, + miss_policy="sync_on_demand") # exact: a miss loads, never falls back + assert out == ref + cache_root = tmp_path / "expert_cache" + dirs = [d for d in os.listdir(cache_root)] + if paging == "auto": + assert stats["paging_resolved"] == "cluster" and stats["num_pages"] == 2 + assert any(d.endswith("_c2") for d in dirs) + cdir = cache_root / next(d for d in dirs if d.endswith("_c2")) + assert sorted(f for f in os.listdir(cdir) if f.endswith(".ctsr")) == ["cluster_0.ctsr", "cluster_1.ctsr"] + else: + assert stats["paging_resolved"] == "expert" and not any(d.endswith("_c2") for d in dirs) + + +@pytest.mark.parametrize("lam,resolved,pages", [(0.3, "cluster", 2), (0.0, "expert", 8)]) +def test_serve_load_model_adapts_paging_and_reports_it(tmp_path, monkeypatch, lam, resolved, pages): + pytest.importorskip("fastapi") + from fastapi.testclient import TestClient + import transformers + import chronos.serving.runtime as srt + from chronos.serving.ollama_server import create_app + path = _save(tmp_path, lam, "served") + monkeypatch.setattr(srt, "_find_tokenizer_dir", lambda *_a, **_k: str(tmp_path)) # tokenizer not under test + monkeypatch.setattr(transformers.PreTrainedTokenizerFast, "from_pretrained", classmethod(lambda cls, *_a, **_k: None)) + lm = srt.load_model(path, name="served", device="cpu", dtype="float32", + expert_cache_root=str(tmp_path / "cache")) + try: + assert lm.lazy_experts and lm.meta["paging_resolved"] == resolved and lm.meta["num_pages"] == pages + cache = os.listdir(tmp_path / "cache") + assert len(cache) == 1 and cache[0].endswith("_c2") == (resolved == "cluster") + runtime = srt.ChronosRuntime(None, device="cpu") + runtime.loaded["served"] = lm + ps = TestClient(create_app(runtime)).get("/api/ps").json()["models"][0]["chronos"] + assert ps["paging_resolved"] == resolved and ps["num_pages"] == pages and ps["paging"] == "auto" + finally: + lm.engine.teardown() diff --git a/tests/test_residency_cluster.py b/tests/test_residency_cluster.py new file mode 100644 index 0000000..7005ea0 --- /dev/null +++ b/tests/test_residency_cluster.py @@ -0,0 +1,265 @@ +"""Next-use residency planning and the cluster-locality training constraint.""" +from __future__ import annotations + +import math +import os +import sys + +import pytest + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +from chronos.expert_clusters import ( # noqa: E402 + cluster_membership, contiguous_clusters, expert_to_cluster, resolve_cluster_count, +) +from chronos.residency_planner import NextUsePlanner # noqa: E402 + + +# ── planner ─────────────────────────────────────────────────────── + +def _run(planner, kind, trace, verify=False): + planner.begin_forward(kind) + for layer, it, eid in trace: + planner.access(layer, it, eid) + planner.end_forward(verify=verify) + + +def test_planner_evicts_the_page_needed_latest(): + p = NextUsePlanner(8) + trace = [(0, 0, 1), (1, 0, 2), (1, 1, 3), (2, 0, 4)] + _run(p, "single", trace) + p.begin_forward("single") + p.access(0, 0, 1) + # 2 is needed next, 4 last, 7 never: never-used first, then the latest one + assert p.victim([2, 4, 7]) == 7 + assert p.victim([2, 3, 4]) == 4 + p.access(1, 0, 2) + assert p.victim([1, 3]) == 1 # 1 is done for this forward, 3 still ahead + + +def test_planner_never_evicts_pending_experts_of_the_running_dispatch(): + p = NextUsePlanner(8) + p.begin_forward("single") + p.access(0, 0, 2, pending=[5, 6]) + assert p.victim([5, 6]) in (5, 6) # forced choice still returns a candidate + assert p.victim([5, 6, 7]) == 7 + + +def test_planner_uses_the_draft_queue_for_verify(): + p = NextUsePlanner(8) + _run(p, "multi", [(0, 0, 0)]) # prefill + _run(p, "draft", [(0, 0, 3), (1, 0, 4)]) # two draft tokens + _run(p, "draft", [(0, 0, 5), (1, 0, 4)]) + p.begin_forward("multi") # verify + assert p.stats["queue_verifies"] == 1 + assert set(p.upcoming()) >= {3, 4, 5} + assert p.victim([3, 4, 5, 6]) == 6 + p.end_forward(verify=True) + p.begin_forward("multi") # queue cleared after verify + assert p.stats["queue_verifies"] == 1 + + +def test_planner_prefetch_is_bounded_by_spare_slots(): + p = NextUsePlanner(8) + _run(p, "single", [(0, 0, 1), (1, 0, 2), (2, 0, 3)]) + p.begin_forward("single") + # resident {1, 6, 7}: 6 and 7 are never used → two spare slots at capacity 3 + assert p.prefetch_plan({1, 6, 7}, capacity=3) == [2, 3] + assert p.prefetch_plan({1, 2, 3}, capacity=3) == [] + assert p.prefetch_plan({1}, capacity=2) == [2] + + +def test_planner_dedupes_loader_and_touch_reports(): + p = NextUsePlanner(4) + p.begin_forward("single") + p.access(0, 0, 1) + p.access(0, 0, 1) + p.end_forward() + assert p.stats["accesses"] == 1 + + +# ── cluster partition ───────────────────────────────────────────── + +def test_contiguous_clusters_cover_every_expert_once(): + for e, c in ((36, 9), (16, 4), (10, 3), (5, 8)): + clusters = contiguous_clusters(e, c) + assign = expert_to_cluster(clusters, e) + assert sorted(x for cl in clusters for x in cl) == list(range(e)) + assert len(set(assign)) == min(c, e) + rows = cluster_membership(clusters, e) + assert all(sum(r) == 1 for r in rows) + with pytest.raises(ValueError): + expert_to_cluster([[0, 1], [1, 2]], 3) + + +def test_default_cluster_count_is_e_over_2k(): + class Cfg: + num_experts, num_experts_per_tok, cluster_locality_clusters = 36, 2, 0 + assert resolve_cluster_count(Cfg) == 9 + Cfg.cluster_locality_clusters = 4 + assert resolve_cluster_count(Cfg) == 4 + + +# ── cluster-locality loss ───────────────────────────────────────── + +torch = pytest.importorskip("torch") + + +def _membership(e=8, c=4): + return torch.tensor(cluster_membership(contiguous_clusters(e, c), e), dtype=torch.float32) + + +def test_cluster_locality_loss_range_and_extremes(): + from chronos.model.temporal_loss import cluster_locality_loss + m = _membership() + B, S, V, E = 2, 16, 3, 8 + uniform = torch.full((B, S, V, E), 1.0 / E) + assert cluster_locality_loss(uniform, m, window=4).item() == pytest.approx(1.0, abs=1e-5) + # every window inside one cluster, the four windows of a row on four clusters + local = torch.zeros(B, S, V, E) + for j in range(4): + local[:, j * 4:(j + 1) * 4, :, 2 * j:2 * j + 2] = 0.5 + assert cluster_locality_loss(local, m, window=4).item() == pytest.approx(0.0, abs=1e-4) + # collapse onto one cluster everywhere: windows are pure but carry no information + collapsed = torch.zeros(B, S, V, E) + collapsed[..., 0:2] = 0.5 + assert cluster_locality_loss(collapsed, m, window=4).item() == pytest.approx(1.0, abs=1e-4) + rand = torch.softmax(torch.randn(B, S, V, E), dim=-1) + val = cluster_locality_loss(rand, m, window=4).item() + assert 0.0 <= val <= 1.0 + + +def test_cluster_locality_loss_counts_every_recurrent_iteration(): + """A window that switches cluster at a later loop iteration is not local.""" + from chronos.model.temporal_loss import cluster_locality_loss + m = _membership() + probs = torch.zeros(1, 16, 2, 8) + for j in range(4): # window j: cluster j in both iterations + probs[:, j * 4:(j + 1) * 4, :, 2 * j:2 * j + 2] = 0.5 + same = cluster_locality_loss(probs, m, window=4).item() + probs[:, :4, 1] = 0.0 + probs[:, :4, 1, 2:4] = 0.5 # iteration 1 of window 0 moves to cluster 1 + split = cluster_locality_loss(probs, m, window=4).item() + assert same == pytest.approx(0.0, abs=1e-4) and split > same + 0.1 + + +def test_cluster_locality_gradient_moves_mass_into_the_window_cluster(): + from chronos.model.temporal_loss import cluster_locality_loss + torch.manual_seed(0) + m = _membership() + logits = torch.randn(2, 16, 2, 8, requires_grad=True) + opt = torch.optim.Adam([logits], lr=0.1) + first = None + for _ in range(200): + loss = cluster_locality_loss(torch.softmax(logits, -1), m, window=4) + first = first if first is not None else loss.item() + opt.zero_grad() + loss.backward() + opt.step() + assert loss.item() < 0.5 * first + + +def test_torch_and_mlx_cluster_locality_losses_match(): + mx = pytest.importorskip("mlx.core") + import numpy as np + from chronos.mlx.training.trainer import _cluster_locality_loss + from chronos.model.temporal_loss import cluster_locality_loss + rng = np.random.default_rng(1) + logits = rng.normal(size=(2, 24, 5, 12)).astype(np.float32) + probs = np.exp(logits) / np.exp(logits).sum(-1, keepdims=True) + member = cluster_membership(contiguous_clusters(12, 3), 12) + for window in (0, 8, 5): + ref = cluster_locality_loss(torch.tensor(probs), torch.tensor(member), window=window).item() + with mx.stream(mx.cpu): + got = float(_cluster_locality_loss(mx.array(probs), mx.array(member, dtype=mx.float32), window)) + assert got == pytest.approx(ref, abs=1e-5) + + +def test_cluster_term_reaches_both_training_losses(): + """``lambda_cluster_locality`` adds the term in the Torch and MLX trainers (0 = absent).""" + mx = pytest.importorskip("mlx.core") + from chronos.mlx.model import ChronosMLXModel + from chronos.mlx.training.trainer import _ce_stage_loss + from chronos.model.config import ChronosConfig + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.trainer.loss_mixin import chronos_loss_term + + def cfg(lam): + return ChronosConfig(hidden_size=32, num_hidden_layers=3, num_attention_heads=4, num_key_value_heads=2, + vocab_size=48, num_experts=8, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=32, max_position_embeddings=64, recurrent_depth=2, + prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0, + lambda_cluster_locality=lam, cluster_locality_window=8) + ids = torch.randint(3, 48, (2, 16)) + totals = {} + for lam in (0.0, 0.5): + torch.manual_seed(3) + c = cfg(lam) + pt = ChronosForCausalLM(c) + out = pt(ids) + base = out.logits.sum() * 0.0 + totals[("torch", lam)] = chronos_loss_term(pt, base, None, c).item() + with mx.stream(mx.cpu): + m = ChronosMLXModel.from_chronos_pytorch(pt, c) + m.train() + _ce_stage_loss(m, mx.array(ids.numpy(), dtype=mx.int32), mx.array(ids.numpy(), dtype=mx.int32)) + totals[("mlx", lam)] = "cluster_locality_loss" in m._last_loss_components + assert totals[("torch", 0.5)] > totals[("torch", 0.0)] + assert totals[("mlx", 0.5)] and not totals[("mlx", 0.0)] + + +# ── lazy engine ─────────────────────────────────────────────────── + +@pytest.mark.parametrize("speculative", [False, True]) +def test_planned_residency_is_exact_and_uses_the_planner(tmp_path, speculative): + """Greedy output is unchanged; eviction runs through the planner and is removed afterwards. + + Load counts are not asserted here: a 4-expert toy with a 2-page budget is + too small for next-use prediction to beat LRU reliably; the effect is + measured on the 50M checkpoint (``tools/bench_residency.py``). + """ + mx = pytest.importorskip("mlx.core") + from test_mlx_speculative import _engine, prompt + with mx.stream(mx.cpu): + _model, engine = _engine(tmp_path, recommended_resident_experts=2, recommended_ram_experts=2) + try: + engine.setup(warm_expert_ids=[0, 1]) + ids = mx.array([prompt(7)], dtype=mx.int32) + plain = list(engine.generate(ids, max_new_tokens=24, temperature=0.0)) + kw = dict(speculative=True, draft_depth=1, gamma=4) if speculative else {} + loads, policy = {}, {} + for res in ("window", "planned"): + engine.store.replace_live_experts_with_placeholders() + with engine.store._lock: + engine.store._warm.clear() + engine.store._warm_lru.clear() + engine.store.warm_up([0, 1]) + out = list(engine.generate(ids, max_new_tokens=24, temperature=0.0, residency=res, **kw)) + assert out == plain + st = engine.last_stats + loads[res] = st["sync_ssd_loads"] + policy[res] = st["eviction_policy"] + if res == "planned": + planner = st["residency"]["planner"] + assert planner["accesses"] > 0 and planner["forwards"] > 0 + if speculative: + assert planner["queue_verifies"] > 0 + assert engine.store.stats()["eviction_policy"] == "lru" # ranker removed afterwards + finally: + engine.stop() + assert policy == {"window": "lru", "planned": "planned"} + assert loads["planned"] > 0 and loads["window"] > 0 + + +def test_unknown_residency_mode_is_rejected(tmp_path): + mx = pytest.importorskip("mlx.core") + from test_mlx_speculative import _engine, prompt + with mx.stream(mx.cpu): + _model, engine = _engine(tmp_path) + try: + with pytest.raises(ValueError): + list(engine.generate(mx.array([prompt(5)], dtype=mx.int32), max_new_tokens=2, residency="queue")) + finally: + engine.stop() diff --git a/tests/test_smoke.py b/tests/test_smoke.py index cc69b38..cb72dc4 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -891,8 +891,9 @@ class _FakeEngine: def __init__(self, _model, _cfg, ssd_dir=None): pass - def setup_from_state_dict(self, state, warm_expert_ids=None): + def setup_from_state_dict(self, state, warm_expert_ids=None, clusters=None): captured["setup_keys"] = sorted(state.keys()) + captured["clusters"] = clusters def setup(self, warm_expert_ids=None): captured["setup_called"] = True @@ -930,6 +931,8 @@ def teardown(self): assert captured["strict"] is False assert captured["loaded_keys"] == ["model.embed_tokens.weight"] assert any(".mlp.experts." in key for key in captured["setup_keys"]) + assert captured["clusters"] is None # trained without cluster locality: per-expert pages + assert stats["paging_resolved"] == "expert" assert captured["teardown"] is True diff --git a/tools/bench_cluster_locality.py b/tools/bench_cluster_locality.py new file mode 100644 index 0000000..f8821cd --- /dev/null +++ b/tools/bench_cluster_locality.py @@ -0,0 +1,277 @@ +"""Cluster-locality training A/B: does the constraint keep a stream inside one expert cluster? + +Two tiny Chronos models start from the same initialisation and train on the +same synthetic multi-topic streams (MLX, the production training loss), +one with ``lambda_cluster_locality = 0`` and one with it on. Each topic +uses its own sub-vocabulary and Markov chain, so a stream has a latent +domain the router can learn to specialise on. Reported, all measured: + +* held-out CE (the constraint must not cost quality); +* cluster window entropy / clusters touched per window / marginal entropy + over every dispatch, recurrent iterations included; +* topic → cluster purity (how much of a topic's routing mass lands in its + dominant cluster) — the specialisation the constraint should produce; +* lazy greedy decoding with experts paged per cluster (one SSD read per + cluster file): cluster reads and expert misses per token under the same + budget, LRU and next-use planned residency. + + python tools/bench_cluster_locality.py --steps 300 --out r.json +""" +from __future__ import annotations + +import argparse +import json +import math +import os +import shutil +import sys +import tempfile +import time + +_repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo not in sys.path: + sys.path.insert(0, _repo) + +SPECIAL = 4 # ids 0..3 reserved + + +def make_topics(num_topics: int, sub_vocab: int, seed: int): + """Per topic: disjoint token range and a sparse random transition table.""" + import numpy as np + rng = np.random.default_rng(seed) + topics = [] + for t in range(num_topics): + lo = SPECIAL + t * sub_vocab + succ = rng.integers(0, sub_vocab, size=(sub_vocab, 3)) # 3 likely successors per token + topics.append((lo, succ)) + return topics + + +def sample_stream(topics, length: int, segment: int, rng): + """Concatenated single-topic segments of ``segment`` tokens (topic per segment).""" + ids, tags = [], [] + while len(ids) < length: + t = int(rng.integers(len(topics))) + lo, succ = topics[t] + cur = int(rng.integers(succ.shape[0])) + for _ in range(segment): + ids.append(lo + cur) + tags.append(t) + cur = int(succ[cur, rng.integers(3)]) if rng.random() < 0.9 else int(rng.integers(succ.shape[0])) + return ids[:length], tags[:length] + + +def tiny_config(args, lam: float): + from chronos.model.config import ChronosConfig + return ChronosConfig( + hidden_size=args.hidden, num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2, + vocab_size=SPECIAL + args.topics * args.sub_vocab, num_experts=args.experts, num_experts_per_tok=2, + num_shared_experts=1, moe_intermediate_size=args.hidden, max_position_embeddings=args.seq * 2, + recurrent_depth=3, prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0, + recurrent_attention_mode="shared_first", sliding_window_size=args.seq, + lambda_cluster_locality=lam, cluster_locality_clusters=args.clusters, + cluster_locality_window=args.segment if args.window is None else args.window, lookahead_steps=2) + + +def build_pair(args): + """Two MLX models with identical initial weights (one Torch init, converted twice).""" + import torch + from chronos.mlx.model import ChronosMLXModel + from chronos.model.model_chronos import ChronosForCausalLM + models = {} + for arm, lam in (("baseline", 0.0), ("cluster", args.lam)): + cfg = tiny_config(args, lam) + torch.manual_seed(args.seed) + pt = ChronosForCausalLM(cfg) + models[arm] = (ChronosMLXModel.from_chronos_pytorch(pt, cfg), cfg) + return models + + +def train(model, cfg, batches, args, log): + import mlx.core as mx + import mlx.nn as nn + import mlx.optimizers as optim + from chronos.mlx.training.trainer import _ce_stage_loss + model.train() + opt = optim.AdamW(learning_rate=args.lr, weight_decay=0.0) + grad_fn = nn.value_and_grad(model, lambda m, x: _ce_stage_loss(m, x, x)) + curve = [] + for step, batch in enumerate(batches): + x = mx.array(batch, dtype=mx.int32) + loss, grads = grad_fn(model, x) + grads, _ = optim.clip_grad_norm(grads, 1.0) + opt.update(model, grads) + mx.eval(model.parameters(), opt.state, loss) + terms = {k: float(v) for k, v in (getattr(model, "_last_loss_components", {}) or {}).items()} + if step % max(1, args.steps // 10) == 0 or step == len(batches) - 1: + curve.append({"step": step, "loss": round(float(loss), 4), + "ce": round(terms.get("primary_loss", float("nan")), 4), + "cluster_term": round(terms.get("cluster_locality_loss", 0.0), 5)}) + log(f" step {step:4d} loss {float(loss):.4f} ce {terms.get('primary_loss', 0):.4f} " + f"cluster {terms.get('cluster_locality_loss', 0.0):.5f}") + model.eval() + return curve + + +def routing_metrics(model, cfg, eval_ids, eval_tags, args): + """Held-out CE and cluster metrics over every dispatch (recurrent iterations included).""" + import mlx.core as mx + import numpy as np + from chronos.expert_clusters import expert_to_cluster, training_clusters + from chronos.mlx.training.trainer import _collect_router_probs, _masked_lm_loss + x = mx.array(eval_ids, dtype=mx.int32) + logits, _la, _c = model(x, logits_dtype=mx.float32) + ce = float(_masked_lm_loss(logits, x)) + probs = np.array(_collect_router_probs(model).astype(mx.float32)) # [B, S, V, E] + B, S, V, E = probs.shape + assign = np.array(expert_to_cluster(training_clusters(cfg), E)) + C = int(assign.max()) + 1 + member = np.eye(C)[assign] # [E, C] + w = args.segment + n = S // w + mass = (probs @ member)[:, : n * w].reshape(B, n, w, V, C).mean(axis=(2, 3)).reshape(-1, C) + ent = lambda p: -(np.clip(p, 1e-12, None) * np.log(np.clip(p, 1e-12, None))).sum(-1) + top = np.argsort(-probs, axis=-1)[..., :cfg.num_experts_per_tok] # [B, S, V, k] + touched = np.zeros((B, n, C), dtype=bool) + for b in range(B): + for j in range(n): + touched[b, j, np.unique(assign[top[b, j * w:(j + 1) * w]])] = True + tags = np.array(eval_tags) # [B, S] + cmass = (probs @ member).mean(axis=2) # [B, S, C] + purity = [] + for t in range(args.topics): + sel = cmass[tags == t] + if len(sel): + purity.append(float(sel.mean(axis=0).max())) + usage = np.bincount(top.reshape(-1), minlength=E) / top.size + return { + "heldout_ce": round(ce, 4), + "cluster_window_entropy": round(float(ent(mass).mean() / math.log(C)), 4), + "cluster_marginal_entropy": round(float(ent(mass.mean(0)) / math.log(C)), 4), + "clusters_per_window": round(float(touched.sum(-1).mean()), 3), + "topic_cluster_purity": round(float(np.mean(purity)), 4), + "expert_usage_entropy": round(float(ent(usage) / math.log(E)), 4), + "num_clusters": C, + } + + +def lazy_decode(model, cfg, prompts, args, workdir): + """Greedy lazy decoding, experts paged per training cluster; LRU vs planned residency.""" + import mlx.core as mx + from chronos.expert_clusters import training_clusters + from chronos.mlx.inference import ChronosMLXInferenceEngine + cfg.recommended_resident_experts = args.hot + cfg.recommended_ram_experts = args.warm + engine = ChronosMLXInferenceEngine(model, cfg, ssd_dir=os.path.join(workdir, "experts")) + clusters = training_clusters(cfg) + out = {} + try: + engine.setup_with_clusters(clusters, warm_expert_ids=[]) + store = engine.store + outputs = {} + for residency in ("window", "planned"): + tot = {"tokens": 0, "sync_ssd_loads": 0, "cluster_loads": 0, "seconds": 0.0} + outputs[residency] = [] + for p in prompts: + store.set_protected_experts([]) + store.set_eviction_ranker(None) + store.replace_live_experts_with_placeholders() + with store._lock: + store._warm.clear() + store._warm_lru.clear() + store._loaded_clusters.clear() + store._stats.clear() + t0 = time.perf_counter() + toks = list(engine.generate(mx.array([p], dtype=mx.int32), max_new_tokens=args.decode, + temperature=0.0, top_p=1.0, residency=residency)) + tot["seconds"] += time.perf_counter() - t0 + st = engine.last_stats + outputs[residency].append(toks) + tot["tokens"] += len(toks) + tot["sync_ssd_loads"] += int(st.get("sync_ssd_loads", 0)) + tot["cluster_loads"] += int(st.get("cluster_loads", 0)) + n = max(tot["tokens"], 1) + out[residency] = {**tot, "seconds": round(tot["seconds"], 3), + "expert_misses_per_token": round(tot["sync_ssd_loads"] / n, 4), + "cluster_reads_per_token": round(tot["cluster_loads"] / n, 4)} + out["planned_identical_to_window"] = outputs["planned"] == outputs["window"] + finally: + engine.stop() + return out + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--topics", type=int, default=4) + ap.add_argument("--sub-vocab", type=int, default=24) + ap.add_argument("--experts", type=int, default=16) + ap.add_argument("--clusters", type=int, default=4) + ap.add_argument("--hidden", type=int, default=64) + ap.add_argument("--seq", type=int, default=128) + ap.add_argument("--segment", type=int, default=32, help="topic segment length of the synthetic streams") + ap.add_argument("--window", type=int, default=None, + help="cluster_locality_window (default: --segment; 0 = whole sequence)") + ap.add_argument("--batch", type=int, default=8) + ap.add_argument("--steps", type=int, default=300) + ap.add_argument("--lr", type=float, default=3e-3) + ap.add_argument("--lam", type=float, default=0.3, help="lambda_cluster_locality of the constrained arm") + ap.add_argument("--hot", type=int, default=4, help="hot expert pages (lazy decode budget)") + ap.add_argument("--warm", type=int, default=8, help="RAM expert pages (lazy decode budget)") + ap.add_argument("--decode", type=int, default=48) + ap.add_argument("--prompts", type=int, default=8) + ap.add_argument("--seed", type=int, default=7) + ap.add_argument("--out", default=None) + args = ap.parse_args(argv) + + import mlx.core as mx + import numpy as np + log = lambda m: print(m, flush=True) + topics = make_topics(args.topics, args.sub_vocab, args.seed) + rng = np.random.default_rng(args.seed + 1) + batches = [[sample_stream(topics, args.seq, args.segment, rng)[0] for _ in range(args.batch)] + for _ in range(args.steps)] + held = [sample_stream(topics, args.seq, args.segment, np.random.default_rng(10_000 + i)) for i in range(16)] + eval_ids, eval_tags = [h[0] for h in held], [h[1] for h in held] + prompts = [] + for i in range(args.prompts): # one-topic prompts, round-robin over topics + lo, succ = topics[i % args.topics] + r = np.random.default_rng(20_000 + i) + cur, p = int(r.integers(succ.shape[0])), [] + for _ in range(12): + p.append(lo + cur) + cur = int(succ[cur, r.integers(3)]) + prompts.append(p) + + report = {"label": "measured", "time": time.strftime("%Y-%m-%dT%H:%M:%S"), "settings": vars(args), + "data": "synthetic: per-topic disjoint sub-vocabulary + sparse Markov chain; " + f"segments of {args.segment} tokens, one topic each", "arms": {}} + with mx.stream(mx.cpu): + for arm, (model, cfg) in build_pair(args).items(): + log(f"[{arm}] lambda_cluster_locality={cfg.lambda_cluster_locality}") + mx.random.seed(args.seed) + t0 = time.perf_counter() + curve = train(model, cfg, batches, args, log) + train_s = time.perf_counter() - t0 + metrics = routing_metrics(model, cfg, eval_ids, eval_tags, args) + log(f"[{arm}] " + " ".join(f"{k} {v}" for k, v in metrics.items())) + workdir = tempfile.mkdtemp(prefix=f"chronos_cluster_{arm}_") + try: + decode = lazy_decode(model, cfg, prompts, args, workdir) + finally: + shutil.rmtree(workdir, ignore_errors=True) + for res in ("window", "planned"): + d = decode[res] + log(f"[{arm}] lazy {res:7s} expert misses/token {d['expert_misses_per_token']:.3f} " + f"cluster reads/token {d['cluster_reads_per_token']:.3f} tokens {d['tokens']}") + report["arms"][arm] = {"lambda_cluster_locality": cfg.lambda_cluster_locality, + "train_seconds": round(train_s, 2), "curve": curve, + "routing": metrics, "lazy_decode": decode} + if args.out: + os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) + with open(args.out, "w", encoding="utf-8") as fh: + json.dump(report, fh, ensure_ascii=False, indent=1) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/bench_residency.py b/tools/bench_residency.py new file mode 100644 index 0000000..2437124 --- /dev/null +++ b/tools/bench_residency.py @@ -0,0 +1,163 @@ +"""Lazy-expert residency benchmark: LRU window vs next-use planning, plain and speculative. + +One MLX lazy engine (SSD expert pages under the Inference-tab budget) decodes +the same prompts greedily in up to five modes, alternating per prompt: + +* ``plain`` — plain decoding, LRU + protected LookaheadRouter window; +* ``planned`` — plain decoding, next-use eviction + planned prefetch; +* ``spec-window`` — self-speculative, LRU + last forward's routes + lookahead; +* ``spec-planned`` — self-speculative, next-use eviction planned from the + draft queue, pre-verify prefetch (``chronos.residency_planner``). + +Each prompt starts from the same cold store (the residency state is reset +between modes), so SSD loads compare like with like. Every number in the +report is measured. + + python tools/bench_residency.py --checkpoint path.pth --prompts flores.jsonl --device cpu --out r.json +""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +import statistics +import sys +import tempfile +import time + +_repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo not in sys.path: + sys.path.insert(0, _repo) + +from tools.bench_speculative import _display_path, _sha256, _versions, encode, load_prompts # noqa: E402 + +MODES = ("plain", "planned", "spec-window", "spec-planned") + + +def _reset(engine, warm: int) -> None: + """Cold, identical starting residency for every (prompt, mode).""" + store = engine.store + store.set_protected_experts([]) + store.set_eviction_ranker(None) + store.replace_live_experts_with_placeholders() + with store._lock: + store._warm.clear() + store._warm_lru.clear() + store._loaded_clusters.clear() + store._stats.clear() # per-run store counters (cluster/prefetch loads) + store.warm_up(list(range(min(int(store.num_experts), warm)))) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--prompts", required=True) + ap.add_argument("--field", default="en") + ap.add_argument("--format", default="tags:<|en|>,<|zh|>") + ap.add_argument("--eos", default="<|endoftext|>") + ap.add_argument("--n", type=int, default=32) + ap.add_argument("--max-new", type=int, default=96) + ap.add_argument("--device", default="cpu", choices=("cpu", "gpu")) + ap.add_argument("--dtype", default="auto") + ap.add_argument("--ratio", type=float, default=1.0, help="lazy expert budget (Inference-tab RAM load ratio)") + ap.add_argument("--draft-depth", type=int, default=1) + ap.add_argument("--gamma", type=int, default=4) + ap.add_argument("--decay", type=float, default=0.7) + ap.add_argument("--modes", default=",".join(MODES)) + ap.add_argument("--out", default=None) + ap.add_argument("--note", default="") + args = ap.parse_args(argv) + + import mlx.core as mx + from tokenizers import Tokenizer + from chronos.model.checkpoint import chronos_config_from_checkpoint, resolve_checkpoint_tokenizer + from tools.bench_speculative import mlx_runner + from ui.tabs.inference_tab import _clone_model_cfg + + modes = [m.strip() for m in args.modes.split(",") if m.strip()] + cfg, _ = chronos_config_from_checkpoint(args.checkpoint) + cfg = _clone_model_cfg(cfg) + tok = Tokenizer.from_file(os.path.join(resolve_checkpoint_tokenizer(args.checkpoint, fallback=False), + "tokenizer.json")) + eos = tok.token_to_id(args.eos) + prompts = [encode(tok, x, args.format) for x in load_prompts(args.prompts, args.field, args.n)] + limits = [min(args.max_new, int(1.6 * (len(p) - 2)) + 16) for p in prompts] + workdir = tempfile.mkdtemp(prefix="chronos_residency_") + rows = [] + try: + _decode, info, close = mlx_runner(cfg, args.checkpoint, lazy=True, ratio=args.ratio, device=args.device, + workdir=workdir, dtype=args.dtype) + engine = close.__self__ # mlx_runner returns the engine's bound stop() + warm = int(info["budget"]["effective_vram_expert_budget"]) + try: + for i, (prompt, limit) in enumerate(zip(prompts, limits)): + row = {"index": i} + for mode in modes: + _reset(engine, warm) + ids = mx.array([prompt], dtype=mx.int32) + kw = dict(residency="planned" if mode.endswith("planned") else "window", + residency_decay=args.decay) + if mode.startswith("spec"): + kw.update(speculative=True, draft_depth=args.draft_depth, gamma=args.gamma, gamma_max=8) + t0 = time.perf_counter() + toks = list(engine.generate(ids, max_new_tokens=limit, temperature=0.0, top_p=1.0, + eos_token_id=eos, **kw)) + dt = time.perf_counter() - t0 + st = engine.last_stats + row[mode] = {"tokens": toks, "s": dt, "ssd": int(st.get("sync_ssd_loads", 0)), + "hits": int(st.get("cache_hits", 0)), "misses": int(st.get("cache_misses", 0)), + "cluster_loads": int(st.get("cluster_loads", 0)), + "prefetch_loads": int(st.get("prefetch_loads", 0)), + "planned_evictions": int(st.get("planned_evictions", 0)), + "eviction_policy": st.get("eviction_policy"), + "verify_ready_rate": (st.get("residency") or {}).get("verify_ready_rate")} + rows.append(row) + print(f"[{i + 1}/{len(prompts)}] " + " ".join( + f"{m} ssd {row[m]['ssd']:4d}" for m in modes), flush=True) + finally: + close() + finally: + shutil.rmtree(workdir, ignore_errors=True) + + summary = {} + for mode in modes: + toks = sum(len(r[mode]["tokens"]) for r in rows) + secs = sum(r[mode]["s"] for r in rows) + hits, misses = sum(r[mode]["hits"] for r in rows), sum(r[mode]["misses"] for r in rows) + summary[mode] = { + "tokens": toks, "seconds": round(secs, 3), "tok_s": round(toks / max(secs, 1e-9), 2), + "sync_ssd_loads": sum(r[mode]["ssd"] for r in rows), + "ssd_loads_per_token": round(sum(r[mode]["ssd"] for r in rows) / max(toks, 1), 4), + "selection_miss_rate": round(misses / max(hits + misses, 1), 4), + "prefetch_loads": sum(r[mode]["prefetch_loads"] for r in rows), + "identical_to_plain": sum(r[mode]["tokens"] == r["plain"]["tokens"] for r in rows) if "plain" in modes + else None, + } + if mode == "spec-planned": + rates = [r[mode]["verify_ready_rate"] for r in rows if r[mode]["verify_ready_rate"] is not None] + summary[mode]["verify_ready_rate_mean"] = round(statistics.mean(rates), 4) if rates else None + for mode in modes: + s = summary[mode] + print(f"{mode:12s} tok/s {s['tok_s']:8.2f} ssd loads {s['sync_ssd_loads']:6d} " + f"({s['ssd_loads_per_token']:.3f}/tok) miss rate {s['selection_miss_rate']:.4f} " + f"identical {s['identical_to_plain']}/{len(rows)}", flush=True) + if args.out: + report = {"label": "measured", "time": time.strftime("%Y-%m-%dT%H:%M:%S"), "env": _versions(), + "checkpoint": {"path": _display_path(args.checkpoint), "sha256": _sha256(args.checkpoint)}, + "settings": {**vars(args), "checkpoint": _display_path(args.checkpoint), + "prompts": _display_path(args.prompts), + "out": _display_path(args.out), "eos_id": eos, "decoding": "greedy, batch 1", + "order": "modes alternate per prompt; residency reset to the same cold start"}, + "budget": info.get("budget"), "dtype": str(info.get("dtype")), "summary": summary, + "rows": [{"index": r["index"], **{m: {k: v for k, v in r[m].items() if k != "tokens"} + | {"n_tokens": len(r[m]["tokens"])} for m in modes}} + for r in rows]} + os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) + with open(args.out, "w", encoding="utf-8") as fh: + json.dump(report, fh, ensure_ascii=False, indent=1, default=str) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/train_chronos.py b/train_chronos.py index f35f066..f4999ce 100644 --- a/train_chronos.py +++ b/train_chronos.py @@ -56,6 +56,12 @@ def parse_args(): p.add_argument("--lambda_lookahead_topk", type=float, default=0.15) p.add_argument("--lambda_lookahead_union", type=float, default=0.05) p.add_argument("--lambda_router_locality", type=float, default=0.02) + p.add_argument("--lambda_cluster_locality", type=float, default=0.3, + help="cluster-locality constraint weight (default 0.3, measured; 0 = off)") + p.add_argument("--cluster_locality_clusters", type=int, default=0, + help="expert clusters for the constraint (0 = ceil(E / 2k))") + p.add_argument("--cluster_locality_window", type=int, default=32, + help="tokens per locality window (0 = whole sequence)") p.add_argument("--fallback_mask_prob", type=float, default=0.08, help="Randomly mask routed experts during training to teach the shared offload fallback.") # VRAM budget diff --git a/train_chronos_distill.py b/train_chronos_distill.py index 9f71ffb..f4ebc4a 100644 --- a/train_chronos_distill.py +++ b/train_chronos_distill.py @@ -49,6 +49,8 @@ def main(): p.add_argument("--lambda_lookahead_topk", type=float, default=None) p.add_argument("--lambda_lookahead_union", type=float, default=None) p.add_argument("--lambda_router_locality", type=float, default=None) + p.add_argument("--lambda_cluster_locality", type=float, default=None) + p.add_argument("--cluster_locality_window", type=int, default=None) p.add_argument("--fallback_mask_prob", type=float, default=0.02) p.add_argument("--alpha", type=float, default=0.7, help="Weight on KD loss (vs. label CE).") diff --git a/train_chronos_dpo.py b/train_chronos_dpo.py index 06026a1..2a19245 100644 --- a/train_chronos_dpo.py +++ b/train_chronos_dpo.py @@ -40,6 +40,8 @@ def main(): p.add_argument("--lambda_lookahead_topk", type=float, default=None) p.add_argument("--lambda_lookahead_union", type=float, default=None) p.add_argument("--lambda_router_locality", type=float, default=None) + p.add_argument("--lambda_cluster_locality", type=float, default=None) + p.add_argument("--cluster_locality_window", type=int, default=None) p.add_argument("--fallback_mask_prob", type=float, default=0.02) p.add_argument("--beta", type=float, default=0.1) p.add_argument("--lambda_router_anchor", type=float, default=0.1, diff --git a/train_chronos_grpo.py b/train_chronos_grpo.py index 4a02013..7946e98 100644 --- a/train_chronos_grpo.py +++ b/train_chronos_grpo.py @@ -41,6 +41,8 @@ def main(): p.add_argument("--lambda_lookahead_topk", type=float, default=None) p.add_argument("--lambda_lookahead_union", type=float, default=None) p.add_argument("--lambda_router_locality", type=float, default=None) + p.add_argument("--lambda_cluster_locality", type=float, default=None) + p.add_argument("--cluster_locality_window", type=int, default=None) p.add_argument("--fallback_mask_prob", type=float, default=0.02) p.add_argument("--beta", type=float, default=0.04, help="KL penalty weight") p.add_argument("--lambda_router_anchor", type=float, default=0.1) diff --git a/train_chronos_orpo.py b/train_chronos_orpo.py index 2023604..c262b11 100644 --- a/train_chronos_orpo.py +++ b/train_chronos_orpo.py @@ -40,6 +40,8 @@ def main(): p.add_argument("--lambda_lookahead_topk", type=float, default=None) p.add_argument("--lambda_lookahead_union", type=float, default=None) p.add_argument("--lambda_router_locality", type=float, default=None) + p.add_argument("--lambda_cluster_locality", type=float, default=None) + p.add_argument("--cluster_locality_window", type=int, default=None) p.add_argument("--fallback_mask_prob", type=float, default=0.02) p.add_argument("--beta", type=float, default=0.1) p.add_argument("--lambda_or", type=float, default=0.1) diff --git a/train_chronos_sft.py b/train_chronos_sft.py index 934bdc2..4ea06ef 100644 --- a/train_chronos_sft.py +++ b/train_chronos_sft.py @@ -42,6 +42,9 @@ def main(): p.add_argument("--lambda_lookahead_topk", type=float, default=None) p.add_argument("--lambda_lookahead_union", type=float, default=None) p.add_argument("--lambda_router_locality", type=float, default=None) + p.add_argument("--lambda_cluster_locality", type=float, default=None, + help="default: inherited from the upstream checkpoint (0 = off)") + p.add_argument("--cluster_locality_window", type=int, default=None) p.add_argument("--fallback_mask_prob", type=float, default=0.04) p.add_argument("--lambda_router_anchor", type=float, default=0.01, help="SFT default: weak anchor.") diff --git a/ui/i18n.py b/ui/i18n.py index 5b6626e..8027fd3 100644 --- a/ui/i18n.py +++ b/ui/i18n.py @@ -166,6 +166,9 @@ "config.lambda_la": "λ 前瞻监督", "config.lambda_anchor": "λ 路由锚点", "config.fallback_mask": "专家冷缺失模拟概率", + "config.lambda_cluster": "λ 专家簇局部性（0 = 关闭；推荐 0.3）", + "config.cluster_count": "专家簇数（0 = 自动 ⌈E/2k⌉）", + "config.cluster_window": "簇局部性窗口 token 数（0 = 整条序列；推荐 ≈ 单条记录长度）", "config.storage_format": "存储格式", "config.cpu_threads": "CPU 线程数 (auto = CPU 后端默认 100% 物理核心)", "config.cpu_budget_percent": "CPU 资源比例 (%)", @@ -343,6 +346,9 @@ "config.lambda_la": "λ 前瞻監督", "config.lambda_anchor": "λ 路由錨點", "config.fallback_mask": "專家冷缺失模擬概率", + "config.lambda_cluster": "λ 專家簇局部性（0 = 關閉；建議 0.3）", + "config.cluster_count": "專家簇數（0 = 自動 ⌈E/2k⌉）", + "config.cluster_window": "簇局部性視窗 token 數（0 = 整條序列；建議 ≈ 單筆記錄長度）", "config.storage_format": "儲存格式", "config.cpu_threads": "CPU 執行緒數 (auto = CPU 後端預設 100% 物理核心)", "config.cpu_budget_percent": "CPU 資源比例 (%)", @@ -519,6 +525,9 @@ "config.tie_emb": "tie_word_embeddings", "config.lambda_la": "λ lookahead", "config.lambda_anchor": "λ router-anchor", + "config.lambda_cluster": "λ expert-cluster locality (0 = off; recommended 0.3)", + "config.cluster_count": "Expert clusters (0 = auto ⌈E/2k⌉)", + "config.cluster_window": "Cluster-locality window in tokens (0 = whole sequence; recommended ≈ one record)", "config.fallback_mask": "Expert miss simulation probability", "config.storage_format": "storage_format", "config.cpu_threads": "CPU threads (auto = 100% physical cores for CPU backend)", @@ -696,6 +705,9 @@ "config.tie_emb": "重み共有", "config.lambda_la": "λ 先読み監督", "config.lambda_anchor": "λ ルータアンカー", + "config.lambda_cluster": "λ エキスパートクラスタ局所性（0 = 無効、推奨 0.3）", + "config.cluster_count": "エキスパートクラスタ数（0 = 自動 ⌈E/2k⌉）", + "config.cluster_window": "クラスタ局所性ウィンドウ（トークン数、0 = 系列全体、推奨 ≈ 1 レコード長）", "config.fallback_mask": "Expert miss simulation probability", "config.storage_format": "保存形式", "config.cpu_threads": "CPU スレッド数 (auto = CPU バックエンドでは物理コアの 100%)", @@ -925,6 +937,26 @@ def get_registered_components(): } _MONITOR_GENERIC_EVALS: Dict[str, Dict[str, str]] = {'zh-Hans': {'monitor.eval_metric': '评测指标（从日志自动发现）', 'monitor.eval_note': '这里显示训练器在验证时实际记录的指标，按指标名分组（BLEU、accuracy、reward …），不预设任务类型。没有任务指标的运行只显示验证 CE。', 'monitor.sample_input': '输入', 'monitor.sample_ref': '参考', 'monitor.sample_output': '输出'}, 'en': {'monitor.eval_metric': 'Evaluation metric (discovered from the logs)', 'monitor.eval_note': 'Shows the metrics the trainer actually logged at validation, grouped by metric name (BLEU, accuracy, reward …); no task type is assumed. Runs without task metrics show validation CE only.', 'monitor.sample_input': 'Input', 'monitor.sample_ref': 'Reference', 'monitor.sample_output': 'Output'}, 'zh-Hant': {'monitor.eval_metric': '評測指標（從日誌自動發現）', 'monitor.eval_note': '這裡顯示訓練器在驗證時實際記錄的指標，按指標名分組（BLEU、accuracy、reward …），不預設任務類型。沒有任務指標的執行只顯示驗證 CE。', 'monitor.sample_input': '輸入', 'monitor.sample_ref': '參考', 'monitor.sample_output': '輸出'}, 'ja': {'monitor.eval_metric': '評価指標（ログから自動検出）', 'monitor.eval_note': '学習器が検証時に実際に記録した指標を、指標名ごと（BLEU・accuracy・reward …）にまとめて表示します。タスクの種類は仮定しません。タスク指標のないランは検証 CE のみを表示します。', 'monitor.sample_input': '入力', 'monitor.sample_ref': '参照', 'monitor.sample_output': '出力'}} _INFER_SPECULATIVE: Dict[str, Dict[str, str]] = {'zh-Hans': {'infer.speculative': '自推测解码（循环深度草稿）', 'infer.draft_depth': '草稿深度（循环次数）', 'infer.spec_gamma': '每轮草稿 token 数 γ（自适应，上限 8）', 'infer.spec_note': '自推测：用较浅的循环深度连续起草 γ 个 token，再用全深度一次前向批量校验。不需要额外权重；贪心输出与普通贪心逐 token 一致，采样用拒绝采样保持全深度分布。`per_iteration` 或非循环模型会自动退回普通解码。草稿深度 ≥ 全深度时等同普通解码。统计表会显示接受率与每次校验产出的 token 数。'}, 'zh-Hant': {'infer.speculative': '自推測解碼（循環深度草稿）', 'infer.draft_depth': '草稿深度（循環次數）', 'infer.spec_gamma': '每輪草稿 token 數 γ（自適應，上限 8）', 'infer.spec_note': '自推測：用較淺的循環深度連續起草 γ 個 token，再用全深度一次前向批次校驗。不需要額外權重；貪婪輸出與一般貪婪逐 token 一致，取樣用拒絕取樣保持全深度分布。`per_iteration` 或非循環模型會自動退回一般解碼。草稿深度 ≥ 全深度時等同一般解碼。統計表會顯示接受率與每次校驗產出的 token 數。'}, 'en': {'infer.speculative': 'Self-speculative decoding (recurrent-depth draft)', 'infer.draft_depth': 'Draft depth (loop iterations)', 'infer.spec_gamma': 'Draft tokens per round γ (adaptive, max 8)', 'infer.spec_note': 'Self-speculative: draft γ tokens at a shallow recurrent depth, then verify them in one full-depth forward. No extra weights; greedy output is token-for-token identical to plain greedy, and sampling uses rejection sampling to keep the full-depth distribution. `per_iteration` or non-recurrent models fall back to plain decoding, and so does a draft depth ≥ the full depth. The stats table shows the acceptance rate and tokens per verify.'}, 'ja': {'infer.speculative': '自己投機デコード（再帰深度ドラフト）', 'infer.draft_depth': 'ドラフト深度（ループ回数）', 'infer.spec_gamma': '1 ラウンドのドラフト token 数 γ（適応、上限 8）', 'infer.spec_note': '自己投機：浅い再帰深度で γ 個の token を連続ドラフトし、全深度の 1 回のフォワードでまとめて検証します。追加の重みは不要です。貪欲出力は通常の貪欲出力と token 単位で一致し、サンプリングは棄却サンプリングで全深度の分布を保ちます。`per_iteration` や非再帰モデル、ドラフト深度 ≥ 全深度の場合は通常デコードに戻ります。統計表に受理率と検証 1 回あたりの token 数を表示します。'}} +_INFER_SPECULATIVE['zh-Hans']['infer.residency_planner'] = '专家驻留规划（按下一次使用驱逐 + 草稿队列预取，仅 MLX 懒加载）' +_INFER_SPECULATIVE['zh-Hant']['infer.residency_planner'] = '專家駐留規劃（按下一次使用驅逐 + 草稿佇列預取，僅 MLX 懶載入）' +_INFER_SPECULATIVE['en']['infer.residency_planner'] = 'Expert residency planner (next-use eviction + draft-queue prefetch, MLX lazy only)' +_INFER_SPECULATIVE['ja']['infer.residency_planner'] = 'エキスパート常駐計画（次回使用順の退避 + ドラフトキュー先読み、MLX 遅延ロードのみ）' +_INFER_SPECULATIVE['zh-Hans']['infer.paging'] = '专家分页方式（auto = 按模型训练方式自动选择）' +_INFER_SPECULATIVE['zh-Hans']['infer.residency_decay'] = '驻留规划衰减（0–0.95，推荐 0.7）' +_INFER_SPECULATIVE['zh-Hans']['infer.decode_wait'] = '解码前等待预取（秒，推荐 0.012）' +_INFER_SPECULATIVE['zh-Hans']['infer.verify_wait'] = '校验前等待预取（秒，推荐 0.02）' +_INFER_SPECULATIVE['zh-Hant']['infer.paging'] = '專家分頁方式（auto = 依模型訓練方式自動選擇）' +_INFER_SPECULATIVE['zh-Hant']['infer.residency_decay'] = '駐留規劃衰減（0–0.95，建議 0.7）' +_INFER_SPECULATIVE['zh-Hant']['infer.decode_wait'] = '解碼前等待預取（秒，建議 0.012）' +_INFER_SPECULATIVE['zh-Hant']['infer.verify_wait'] = '校驗前等待預取（秒，建議 0.02）' +_INFER_SPECULATIVE['en']['infer.paging'] = 'Expert paging (auto = follow how the model was trained)' +_INFER_SPECULATIVE['en']['infer.residency_decay'] = 'Residency planner decay (0–0.95, recommended 0.7)' +_INFER_SPECULATIVE['en']['infer.decode_wait'] = 'Prefetch wait before a decode step (s, recommended 0.012)' +_INFER_SPECULATIVE['en']['infer.verify_wait'] = 'Prefetch wait before a verify step (s, recommended 0.02)' +_INFER_SPECULATIVE['ja']['infer.paging'] = 'エキスパートのページング（auto = 学習方法に合わせて自動選択）' +_INFER_SPECULATIVE['ja']['infer.residency_decay'] = '常駐計画の減衰（0–0.95、推奨 0.7）' +_INFER_SPECULATIVE['ja']['infer.decode_wait'] = 'デコード前の先読み待ち（秒、推奨 0.012）' +_INFER_SPECULATIVE['ja']['infer.verify_wait'] = '検証前の先読み待ち（秒、推奨 0.02）' for _lang, _entries in _INFER_SPECULATIVE.items(): TRANSLATIONS[_lang].update(_entries) for _lang, _entries in _MONITOR_GENERIC_EVALS.items(): diff --git a/ui/presets.py b/ui/presets.py index 9ac41b8..17cdd1d 100644 --- a/ui/presets.py +++ b/ui/presets.py @@ -93,6 +93,9 @@ 'lambda_lookahead': 0.1, 'lambda_lookahead_topk': 0.05, 'lambda_router_anchor': 0.0, + 'lambda_cluster_locality': 0.3, + 'cluster_locality_clusters': 0, + 'cluster_locality_window': 32, 'fallback_mask_prob': 0.0, 'vram_budget_gb': 4.0, 'pinned_memory_max_fraction': 0.25, @@ -140,6 +143,9 @@ 'lambda_lookahead', 'lambda_lookahead_topk', 'lambda_router_anchor', + 'lambda_cluster_locality', + 'cluster_locality_clusters', + 'cluster_locality_window', 'fallback_mask_prob', 'vram_budget_gb', 'pinned_memory_max_fraction', @@ -298,6 +304,7 @@ def values_in_input_order(cfg: Dict) -> List: "dtype": "auto", "tie_word_embeddings": True, "lambda_balance": 5e-4, "lambda_temporal": 1e-3, "lambda_lookahead": 0.1, "lambda_lookahead_topk": 0.05, "lambda_router_anchor": 0.0, + "lambda_cluster_locality": 0.3, "cluster_locality_clusters": 0, "cluster_locality_window": 32, "fallback_mask_prob": 0.0, "vram_budget_gb": 4.0, "pinned_memory_max_fraction": 0.25, "use_hybrid_attention": True, "storage_format": "safetensors", diff --git a/ui/presets_v3.json b/ui/presets_v3.json index 45cc2c9..4130c7e 100644 --- a/ui/presets_v3.json +++ b/ui/presets_v3.json @@ -128,6 +128,7 @@ "lambda_balance": 0.005, "fallback_mask_prob": 0.08, "dropout": 0.0, + "lambda_cluster_locality": 0.0, "kv_cache_dtype": "auto", "mla_cache_rope": true, "router_scope": "per_iteration", diff --git a/ui/tabs/config_tab.py b/ui/tabs/config_tab.py index cc51526..9150c85 100644 --- a/ui/tabs/config_tab.py +++ b/ui/tabs/config_tab.py @@ -230,6 +230,16 @@ def build_config_tab(): register_translatable(lambda_lookahead_topk, "config.lambda_la_topk") register_translatable(lambda_router_anchor, "config.lambda_anchor") register_translatable(fallback_mask_prob, "config.fallback_mask") + with gr.Row(): + lambda_cluster_locality = gr.Number(value=D.get("lambda_cluster_locality", 0.3), precision=4, + label=t("config.lambda_cluster")) + cluster_locality_clusters = gr.Number(value=D.get("cluster_locality_clusters", 0), precision=0, + label=t("config.cluster_count")) + cluster_locality_window = gr.Number(value=D.get("cluster_locality_window", 32), precision=0, + label=t("config.cluster_window")) + register_translatable(lambda_cluster_locality, "config.lambda_cluster") + register_translatable(cluster_locality_clusters, "config.cluster_count") + register_translatable(cluster_locality_window, "config.cluster_window") gr.Markdown(f"### {t('config.hw')}") with gr.Row(): @@ -414,6 +424,9 @@ def build_config_tab(): "lambda_lookahead": lambda_lookahead, "lambda_lookahead_topk": lambda_lookahead_topk, "lambda_router_anchor": lambda_router_anchor, + "lambda_cluster_locality": lambda_cluster_locality, + "cluster_locality_clusters": cluster_locality_clusters, + "cluster_locality_window": cluster_locality_window, "fallback_mask_prob": fallback_mask_prob, "vram_budget_gb": vram_budget, "pinned_memory_max_fraction": pinned_frac, diff --git a/ui/tabs/inference_tab.py b/ui/tabs/inference_tab.py index 5727d87..3963c1d 100644 --- a/ui/tabs/inference_tab.py +++ b/ui/tabs/inference_tab.py @@ -1,7 +1,6 @@ """ ui/tabs/inference_tab.py — Real-time generation with tokens/s display """ -import math import os import time @@ -106,34 +105,33 @@ def _resolve_inference_backend(requested_backend: str, model_path_val: str, snif return backend, note -def _checkpoint_expert_cache_dir(model_path_val: str, model_cfg) -> str: - """Keep lazy expert shards isolated per checkpoint. +# Lazy-expert helpers live in chronos.lazy_experts (shared with the bundle and +# the server); these names stay for existing callers. +from chronos.lazy_experts import ( # noqa: E402,F401 (re-exported names) + _store_torch_dtype, + active_expert_count as _actual_active_expert_count, + apply_torch_budget as _apply_offload_expert_budget, + expert_budget as _bounded_offload_expert_budget, + mlx_experts_to_placeholders as _replace_mlx_experts_with_placeholders, + torch_experts_to_placeholders as _replace_torch_experts_with_placeholders, + write_mlx_cache_from_reader as _populate_mlx_cache_from_reader, + write_mlx_cache_from_state_dict as _populate_mlx_cache_from_state_dict, +) - Reusing a single ./expert_cache across different .pth files can silently - load stale expert shards. The cache key uses path + stat metadata + core - topology so repeated runs reuse the same shards, while changed checkpoints - get a clean cache namespace. - """ - import hashlib - import os - base = os.path.join(os.getcwd(), "expert_cache") - if not model_path_val or not os.path.exists(model_path_val): - return base - try: - st = os.stat(model_path_val) - except OSError: - return base - topology = ( - getattr(model_cfg, "vocab_size", ""), - getattr(model_cfg, "hidden_size", ""), - getattr(model_cfg, "num_hidden_layers", ""), - getattr(model_cfg, "num_experts", ""), - getattr(model_cfg, "moe_intermediate_size", ""), - ) - raw = f"{os.path.abspath(model_path_val)}:{st.st_size}:{st.st_mtime_ns}:{topology}" - key = hashlib.sha1(raw.encode("utf-8")).hexdigest()[:16] - return os.path.join(base, f"ckpt_{key}") +def _checkpoint_expert_cache_dir(model_path_val: str, model_cfg, tag: str = "e") -> str: + from chronos.lazy_experts import checkpoint_cache_dir + return checkpoint_cache_dir(model_path_val, model_cfg, tag=tag) + + +def _rebuild_mlx_lazy_cache(store, cache_dir: str, reader, weights, clusters=None) -> None: + """Rebuild the MLX page cache at ``cache_dir`` (always rewrites).""" + import shutil + from chronos.lazy_experts import ensure_mlx_cache + if os.path.exists(cache_dir): + shutil.rmtree(cache_dir) + store._cluster_manifest = None + ensure_mlx_cache(store, cache_dir, reader=reader, weights=weights, clusters=clusters) def _memory_snapshot_gb() -> float: @@ -180,61 +178,6 @@ def _memory_delta(after: dict, before: dict, key: str) -> float: return 0.0 -def _actual_active_expert_count(model_cfg) -> int: - num_experts = int(getattr(model_cfg, "num_experts", 1) or 1) - for attr in ("observed_active_experts", "offload_working_set_experts", "recommended_working_set_experts"): - value = getattr(model_cfg, attr, None) - if value in (None, "", 0): - continue - try: - return max(1, min(num_experts, int(value))) - except (TypeError, ValueError): - pass - top_k = int(getattr(model_cfg, "num_experts_per_tok", 1) or 1) - num_layers = int(getattr(model_cfg, "num_hidden_layers", 1) or 1) - # top_k is per layer. The cache is keyed by global expert id, and a - # single generated token can therefore touch up to top_k * layers distinct - # expert ids before repeats. Capping at top_k alone makes the lazy path - # replace a large part of the routed MoE with shared fallback. - return max(1, min(num_experts, top_k * max(1, num_layers))) - - -def _bounded_offload_expert_budget(model_cfg, ram_load_ratio: float | str | None = 1.0) -> dict: - """Derive the offload cache budget from the trained MoE topology. - - The ideal active count is the routed working set, not the checkpoint's - recommended resident cache size. The user-controlled ratio must produce - real lower-budget runs for sweeps; exact quality is preserved by - on-demand materialisation rather than by silently clamping to top-k. - """ - active = _actual_active_expert_count(model_cfg) - num_experts = int(getattr(model_cfg, "num_experts", active) or active) - try: - ratio = float(ram_load_ratio) - except (TypeError, ValueError): - ratio = 1.0 - ratio = max(0.01, ratio) - routing_top_k = int(getattr(model_cfg, "num_experts_per_tok", active) or active) - min_vram = 1 - hard_cap = max(active, int(math.floor(active * 1.25))) - requested = max(1, int(math.ceil(active * ratio))) - effective = max(min_vram, min(num_experts, requested, hard_cap)) - ram_buffer = effective - return { - "ideal_active_experts": active, - "routing_top_k": routing_top_k, - "num_moe_layers": int(getattr(model_cfg, "num_hidden_layers", 1) or 1), - "requested_ram_load_ratio": round(ratio, 3), - "requested_expert_budget": requested, - "max_allowed_expert_budget": min(num_experts, hard_cap), - "effective_expert_budget": effective, - "effective_vram_expert_budget": effective, - "effective_ram_expert_budget": ram_buffer, - "min_vram_expert_budget": min_vram, - "budget_clamped": effective != requested, - } - - def _clone_model_cfg(model_cfg): from chronos.model.config import ChronosConfig @@ -249,188 +192,6 @@ def _clone_model_cfg(model_cfg): return ChronosConfig(**data) -def _apply_offload_expert_budget(engine, model_cfg, ram_load_ratio: float | str | None = 1.0) -> dict: - budget = _bounded_offload_expert_budget(model_cfg, ram_load_ratio) - vram_cap = int(budget["effective_vram_expert_budget"]) - ram_cap = int(budget["effective_ram_expert_budget"]) - cache_manager = getattr(engine, "cache_manager", None) - store = getattr(cache_manager, "expert_store", None) - if store is None: - return budget - # The hard budget (max_resident_experts / ram_budget_mb) always wins - # over the ratio-derived working set. - hard_vram = getattr(store, "max_resident_pages", None) - hard_ram = getattr(store, "ram_budget_pages", None) - if hard_vram is not None: - vram_cap = max(1, min(vram_cap, int(hard_vram))) - if hard_ram is not None: - ram_cap = max(1, min(ram_cap, int(hard_ram))) - budget["effective_vram_expert_budget"], budget["effective_ram_expert_budget"] = vram_cap, ram_cap - store.vram_capacity = vram_cap - store.ram_capacity = ram_cap - store.vram_lru.capacity = vram_cap - store.ram_lru.capacity = ram_cap - return budget - - -def _replace_torch_experts_with_placeholders(model, model_cfg) -> None: - """Remove randomly initialised live experts before model.to(device). - - For .pth inference the real expert tensors are sharded from the checkpoint - to SSD by ChronosInferenceEngine.setup_from_state_dict(). Keeping the - constructor's random experts alive until model.to(mps/cuda) defeats the - SSD+DRAM lazy-load design by moving full dummy experts onto the device. - """ - from chronos.model.moe_chronos import ChronosMOEFeedForward, LazyExpertPlaceholder - from chronos.precision import torch_model_dtype - - layers = getattr(getattr(model, "model", None), "layers", []) - hidden = int(getattr(model_cfg, "hidden_size", 0) or 0) - intermediate = int(getattr(model_cfg, "moe_intermediate_size", 0) or 0) - if not hidden or not intermediate: - return - for layer in layers: - moe = getattr(layer, "mlp", None) - if not isinstance(moe, ChronosMOEFeedForward): - continue - for idx in range(len(moe.experts)): - # Placeholder dtype follows the trunk; place_torch_model() updates - # it again when the model is moved to its compute dtype. - moe.experts[idx] = LazyExpertPlaceholder(hidden, intermediate, torch_model_dtype(model)) - - -def _replace_mlx_experts_with_placeholders(model) -> None: - import gc - import mlx.core as mx - from chronos.mlx.moe import ChronosMLXMOE, LazyFeedForwardMLX - - for layer in getattr(model, "layers", []): - moe = getattr(layer, "mlp", None) - if not isinstance(moe, ChronosMLXMOE): - continue - for idx in range(len(moe.experts)): - moe.experts[idx] = LazyFeedForwardMLX() - gc.collect() - try: - mx.clear_cache() - except Exception: - pass - - -def _store_torch_dtype(store): - """Torch dtype of the MLX store's pages (= the MLX trunk compute dtype).""" - import torch - name = str(getattr(store, "page_dtype", "float32")).rsplit(".", 1)[-1] - return getattr(torch, name) - - -def _populate_mlx_cache_from_state_dict(store, state_dict: dict) -> None: - import os - import json - import torch - from safetensors.torch import save_file as np_save_file # BF16-capable - from chronos.io.storage import ClusterManifest, KEY_SEP, MANIFEST_FILENAME - - os.makedirs(store.ssd_dir, exist_ok=True) - page_dtype = _store_torch_dtype(store) - manifest_clusters = {} - expert_to_cluster = {} - for eid in range(int(store.num_experts)): - tensors = {} - for li in range(int(store.num_layers)): - prefix = f"model.layers.{li}.mlp.experts.{eid}." - for pname in store._PARAM_NAMES: - key = prefix + pname - if key in state_dict: - tensors[f"l{li}_e{eid}{KEY_SEP}{pname}"] = ( - state_dict[key].detach().to(dtype=page_dtype, device="cpu").contiguous() - ) - if tensors: - file_name = f"cluster_{eid}.ctsr" - np_save_file(tensors, os.path.join(store.ssd_dir, file_name)) - manifest_clusters[eid] = (file_name, [eid]) - expert_to_cluster[eid] = eid - if len(expert_to_cluster) != int(store.num_experts): - raise RuntimeError( - f"MLX lazy cache is incomplete: wrote {len(expert_to_cluster)}/{int(store.num_experts)} experts" - ) - manifest = ClusterManifest( - version=1, - num_experts=int(store.num_experts), - num_layers=int(store.num_layers), - storage_format="safetensors", - clusters=manifest_clusters, - expert_to_cluster=expert_to_cluster, - ) - with open(os.path.join(store.ssd_dir, MANIFEST_FILENAME), "w", encoding="utf-8") as f: - json.dump(manifest.to_dict(), f, indent=2) - - -def _populate_mlx_cache_from_reader(store, reader) -> None: - import os - import json - import torch - from safetensors.torch import save_file as np_save_file # BF16-capable - from chronos.io.storage import ClusterManifest, KEY_SEP, MANIFEST_FILENAME - - os.makedirs(store.ssd_dir, exist_ok=True) - page_dtype = _store_torch_dtype(store) - keys = set(reader.keys()) - manifest_clusters = {} - expert_to_cluster = {} - for eid in range(int(store.num_experts)): - tensors = {} - for li in range(int(store.num_layers)): - prefix = f"model.layers.{li}.mlp.experts.{eid}." - for pname in store._PARAM_NAMES: - key = prefix + pname - if key in keys: - tensors[f"l{li}_e{eid}{KEY_SEP}{pname}"] = ( - reader.get_tensor(key).detach().to(dtype=page_dtype, device="cpu").contiguous() - ) - if tensors: - file_name = f"cluster_{eid}.ctsr" - np_save_file(tensors, os.path.join(store.ssd_dir, file_name)) - manifest_clusters[eid] = (file_name, [eid]) - expert_to_cluster[eid] = eid - if len(expert_to_cluster) != int(store.num_experts): - raise RuntimeError( - f"MLX lazy cache is incomplete: wrote {len(expert_to_cluster)}/{int(store.num_experts)} experts" - ) - manifest = ClusterManifest( - version=1, - num_experts=int(store.num_experts), - num_layers=int(store.num_layers), - storage_format="safetensors", - clusters=manifest_clusters, - expert_to_cluster=expert_to_cluster, - ) - with open(os.path.join(store.ssd_dir, MANIFEST_FILENAME), "w", encoding="utf-8") as f: - json.dump(manifest.to_dict(), f, indent=2) - - -def _rebuild_mlx_lazy_cache(store, cache_dir: str, reader, weights) -> None: - import os - import shutil - - if os.path.exists(cache_dir): - shutil.rmtree(cache_dir) - os.makedirs(cache_dir, exist_ok=True) - store.ssd_dir = cache_dir - store._cluster_manifest = None - store._loaded_clusters.clear() - store._warm.clear() - store._warm_lru.clear() - if reader is not None: - _populate_mlx_cache_from_reader(store, reader) - elif weights is not None: - _populate_mlx_cache_from_state_dict(store, weights) - else: - store.offload_all_to_ssd(clusters=[[i] for i in range(int(store.num_experts))]) - if not store.attach_cluster_manifest(cache_dir): - raise RuntimeError(f"failed to build a complete MLX lazy expert cache at {cache_dir}") - - def _encode_prompt(tokenizer, prompt_val: str, raw_prompt: bool = False) -> tuple[list[int], list[int]]: if raw_prompt: enc = tokenizer(prompt_val, add_special_tokens=False, return_attention_mask=True) @@ -541,7 +302,8 @@ def _run_torch_inference(backend: str, model_cfg, model_path_val: str, token_ids miss_policy: str = "on_demand", cpu_threads: int | str | None = "auto", cpu_budget_percent: int | float | str | None = 100, - speculative: dict | None = None) -> tuple[list[int], dict]: + speculative: dict | None = None, + paging: str = "auto") -> tuple[list[int], dict]: import os import importlib import gc @@ -588,20 +350,27 @@ def _run_torch_inference(backend: str, model_cfg, model_path_val: str, token_ids device_mem = _memory_snapshot_gb() input_ids = torch.tensor([token_ids]).to(device_map) attn = torch.tensor([attention_mask]).to(device_map) + from chronos.lazy_experts import page_layout + layout = page_layout(model_cfg, paging) + # Per-expert layout keeps the per-(layer, expert) shards; cluster layout + # writes one file per training cluster. + page_clusters = None if layout["tag"] == "e" else layout["clusters"] engine = inference_engine_mod.ChronosInferenceEngine( model, model_cfg, - ssd_dir=_checkpoint_expert_cache_dir(model_path_val, model_cfg), + ssd_dir=_checkpoint_expert_cache_dir(model_path_val, model_cfg, tag=layout["tag"]), ) expert_budget = _apply_offload_expert_budget(engine, model_cfg, ram_load_ratio) + expert_budget.update({"paging": layout["paging"], "paging_resolved": layout["resolved"], + "paging_reason": layout["reason"], "num_pages": layout["num_pages"]}) miss_policy = _normalize_miss_policy(miss_policy) if weights is not None: - engine.setup_from_state_dict(weights, warm_expert_ids=[]) + engine.setup_from_state_dict(weights, warm_expert_ids=[], clusters=page_clusters) del weights weights = None gc.collect() elif reader is not None: - engine.setup_from_state_reader(reader, warm_expert_ids=[]) + engine.setup_from_state_reader(reader, warm_expert_ids=[], clusters=page_clusters) del reader gc.collect() else: @@ -976,97 +745,57 @@ def _run_mlx_lazy_inference(model_cfg, model_path_val: str, token_ids: list[int] ram_load_ratio: float | str | None = 1.0, quality_safe: bool = True, eos_token_id: int | None = None, - speculative: dict | None = None) -> tuple[list[int], dict]: - import os - import gc - import torch + speculative: dict | None = None, + residency: str = "planned", + paging: str = "auto", + residency_decay: float | None = None, + decode_wait_s: float | None = None, + verify_wait_s: float | None = None) -> tuple[list[int], dict]: import mlx.core as mx - from chronos.model.model_chronos import ChronosForCausalLM - from chronos.model.checkpoint import load_checkpoint_state_dict, load_state_dict_controlled - from chronos.export import is_export_artifact, open_export_reader - from chronos.mlx.model import ChronosMLXModel - from chronos.mlx.inference import ChronosMLXInferenceEngine - - budget = _bounded_offload_expert_budget(model_cfg, ram_load_ratio) - model_cfg.recommended_resident_experts = int(budget["effective_vram_expert_budget"]) - model_cfg.recommended_ram_experts = int(budget["effective_ram_expert_budget"]) - cache_dir = _checkpoint_expert_cache_dir(model_path_val, model_cfg) + "_mlx" - + from chronos.lazy_experts import LazyOptions, build_mlx_lazy_engine, warm_ids + + defaults = LazyOptions() + opts = LazyOptions( + paging=paging, residency=residency, ram_ratio=_normalize_ram_load_ratio(ram_load_ratio), + residency_decay=defaults.residency_decay if residency_decay is None else float(residency_decay), + decode_wait_s=defaults.decode_wait_s if decode_wait_s is None else float(decode_wait_s), + verify_wait_s=defaults.verify_wait_s if verify_wait_s is None else float(verify_wait_s), + ).validated() setup_mem0 = _memory_snapshot_gb() backend_mem0 = _backend_memory_snapshot("mlx") - reader = None - weights = None - if model_path_val and is_export_artifact(model_path_val): - reader = open_export_reader(model_path_val) - base_only = reader.load_state_dict(include_experts=False) - elif model_path_val and os.path.exists(model_path_val): - weights = load_checkpoint_state_dict(model_path_val, map_location="cpu") - base_only = {k: v for k, v in weights.items() if ".mlp.experts." not in k} - else: - raise FileNotFoundError("MLX lazy/offload inference requires a checkpoint or export artifact.") - - pt_model = ChronosForCausalLM(model_cfg) - if base_only: - load_state_dict_controlled( - pt_model, - base_only, - allow_missing_substrings=(".mlp.experts.",), - ) - _replace_torch_experts_with_placeholders(pt_model, model_cfg) - model_init_mem = _memory_snapshot_gb() - from chronos.precision import resolve_mlx - precision = resolve_mlx("auto") - mlx_model = ChronosMLXModel.from_chronos_pytorch(pt_model.eval(), model_cfg, include_experts=False, - dtype=precision.compute) - _replace_mlx_experts_with_placeholders(mlx_model) - del pt_model - gc.collect() - - engine = ChronosMLXInferenceEngine(mlx_model, model_cfg, ssd_dir=cache_dir) + engine, info = build_mlx_lazy_engine(model_path_val, model_cfg, opts) try: - # MLX lazy mode must keep the same expert budget as the offload path: - # warming all experts hides the architecture and can leave stale LRU - # state when the budget evicts placeholders. Warm only the execution - # budget; selected misses are synchronously materialized expert-by-expert. - warm_count = int(budget["effective_vram_expert_budget"]) - warm_ids = list(range(min(int(model_cfg.num_experts), warm_count))) - if not engine.store.attach_cluster_manifest(cache_dir): - _rebuild_mlx_lazy_cache(engine.store, cache_dir, reader, weights) - if reader is not None: - del reader - reader = None - if weights is not None: - del weights - weights = None - gc.collect() - engine.store.warm_up(warm_ids) + # Warm only the execution budget: warming every expert would hide the + # lazy path; selected misses load synchronously, page by page. + engine.store.warm_up(warm_ids(model_cfg, info)) setup_mem1 = _memory_snapshot_gb() backend_mem_setup = _backend_memory_snapshot("mlx") - ids = mx.array([token_ids]) tokens = list(engine.generate( - ids, + mx.array([token_ids]), max_new_tokens=int(max_tok), temperature=float(temp) if do_sample else 0.0, top_p=0.85, eos_token_id=eos_token_id, speculative=bool(speculative), + **opts.generate_kwargs(), **(speculative or {}), )) stats = dict(getattr(engine, "last_stats", {})) store_stats = engine.store.stats() stats.update({ "setup_rss_delta_gb": round(setup_mem1 - setup_mem0, 3), - "model_init_rss_delta_gb": round(model_init_mem - setup_mem0, 3), - "expert_setup_rss_delta_gb": round(setup_mem1 - model_init_mem, 3), "rss_before_setup_gb": round(setup_mem0, 3), - "rss_after_model_init_gb": round(model_init_mem, 3), "rss_after_setup_gb": round(setup_mem1, 3), "rss_after_prefill_gb": stats.get("rss_after_prefill_gb", round(setup_mem1, 3)), "rss_after_decode_gb": stats.get("rss_after_decode_gb", round(setup_mem1, 3)), - **budget, + **{k: v for k, v in info.items() if k not in ("lazy_options", "precision")}, + "residency_decay": opts.residency_decay, + "decode_wait_s": opts.decode_wait_s, + "verify_wait_s": opts.verify_wait_s, "storage_format": store_stats.get("storage_format", "safetensors"), "cluster_aware": store_stats.get("cluster_aware", False), "num_clusters": store_stats.get("num_clusters", 0), + "cluster_loads": store_stats.get("cluster_loads", 0), "vram_experts": store_stats.get("hot_experts", 0), "vram_capacity": store_stats.get("hot_capacity", 0), "ram_experts": store_stats.get("warm_experts", 0), @@ -1074,20 +803,12 @@ def _run_mlx_lazy_inference(model_cfg, model_path_val: str, token_ids: list[int] "miss_policy": "mlx_lazy_quality", "quality_safe": quality_safe, }) - for snapshot, suffix in [ - (backend_mem0, "before_setup"), - (backend_mem_setup, "after_setup"), - ]: + for snapshot, suffix in [(backend_mem0, "before_setup"), (backend_mem_setup, "after_setup")]: for key in ("mlx_active_gb", "mlx_cache_gb", "mlx_peak_gb"): if key in snapshot: - prefix = key[:-3] if key.endswith("_gb") else key - stats[f"{prefix}_{suffix}_gb"] = snapshot[key] + stats[f"{key[:-3]}_{suffix}_gb"] = snapshot[key] return tokens, stats finally: - if reader is not None: - del reader - if weights is not None: - del weights engine.stop() @@ -1303,6 +1024,7 @@ def fmt_bool(row: dict, key: str) -> str: ("Warm evictions", lambda r: fmt_int(r, "warm_evictions")), ("Async misses", lambda r: fmt_int(r, "async_cold_miss_prefetches")), ("Sync SSD loads", lambda r: fmt_int(r, "sync_ssd_loads")), + ("Eviction policy", lambda r: str(r.get("eviction_policy") or "")), ("Expert hits/misses", fmt_hits), ("All-resident fast path", lambda r: fmt_bool(r, "all_resident_fast_path")), ("Load budget", lambda r: str(r.get("load_budget", "all"))), @@ -1428,6 +1150,7 @@ def _row_from_stats(mode: str, backend: str, tokens: list[int], elapsed: float, "prefetch_queue_drops": stats.get("prefetch_queue_drops", 0), "prefetch_wait_time_s": stats.get("prefetch_wait_time_s", 0.0), "sync_ssd_loads": stats.get("sync_ssd_loads", 0), + "eviction_policy": stats.get("eviction_policy"), "resident_vram_hits": stats.get("resident_vram_hits", 0), "resident_ram_hits": stats.get("resident_ram_hits", 0), "warm_hits": stats.get("warm_hits", 0), @@ -1559,7 +1282,8 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, prompt_val, max_tok, temp, raw_prompt_val, ram_load_ratio=1.0, sweep_ram_load_ratios=False, miss_policy="on_demand", speculative=False, draft_depth=1, - gamma=4) -> dict: + gamma=4, residency_planner=True, paging="auto", residency_decay=None, + decode_wait_s=None, verify_wait_s=None) -> dict: from transformers import AutoTokenizer from chronos.backend.mac_diagnostics import mac_backend_diagnostics @@ -1620,6 +1344,11 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, quality_safe=(mode == "compare" and not bool(sweep_ram_load_ratios)), eos_token_id=tokenizer.eos_token_id, speculative=spec_opts, + residency="planned" if residency_planner else "window", + paging=paging, + residency_decay=residency_decay, + decode_wait_s=decode_wait_s, + verify_wait_s=verify_wait_s, ) else: lazy_tokens, lazy_stats = _run_torch_inference( @@ -1637,6 +1366,7 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, cpu_threads=cpu_threads_cfg, cpu_budget_percent=cpu_budget_cfg, speculative=spec_opts, + paging=paging, ) lazy_elapsed = time.monotonic() - t0 lazy_mem = _memory_snapshot_gb() @@ -1731,12 +1461,16 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, def generate_api(cfg, selected_backend="auto", selected_mode="compare", model_path_val="", prompt_val="", max_tok=128, temp=0.85, raw_prompt_val=False, ram_load_ratio=1.0, sweep_ram_load_ratios=False, - miss_policy="on_demand", speculative=False, draft_depth=1, gamma=4) -> dict: + miss_policy="on_demand", speculative=False, draft_depth=1, gamma=4, + residency_planner=True, paging="auto", residency_decay=None, + decode_wait_s=None, verify_wait_s=None) -> dict: try: return _run_inference_modes( cfg, selected_backend, selected_mode, model_path_val, prompt_val, max_tok, temp, raw_prompt_val, ram_load_ratio, sweep_ram_load_ratios, - miss_policy, speculative, draft_depth, gamma, + miss_policy, speculative, draft_depth, gamma, residency_planner, + paging=paging, residency_decay=residency_decay, + decode_wait_s=decode_wait_s, verify_wait_s=verify_wait_s, ) except Exception as e: import traceback @@ -1805,6 +1539,21 @@ def build_inference_tab(config_state: gr.State): register_translatable(speculative, "infer.speculative") register_translatable(draft_depth, "infer.draft_depth") register_translatable(spec_gamma, "infer.spec_gamma") + residency_planner = gr.Checkbox(value=True, label=t("infer.residency_planner")) + register_translatable(residency_planner, "infer.residency_planner") + with gr.Row(): + from chronos.lazy_experts import ( + DEFAULT_DECODE_WAIT_S, DEFAULT_RESIDENCY_DECAY, DEFAULT_VERIFY_WAIT_S, + ) + paging_mode = gr.Dropdown(choices=["auto", "cluster", "expert"], value="auto", + label=t("infer.paging")) + residency_decay = gr.Slider(0.0, 0.95, value=DEFAULT_RESIDENCY_DECAY, step=0.05, + label=t("infer.residency_decay")) + decode_wait = gr.Number(value=DEFAULT_DECODE_WAIT_S, minimum=0.0, label=t("infer.decode_wait")) + verify_wait = gr.Number(value=DEFAULT_VERIFY_WAIT_S, minimum=0.0, label=t("infer.verify_wait")) + for comp, key in ((paging_mode, "infer.paging"), (residency_decay, "infer.residency_decay"), + (decode_wait, "infer.decode_wait"), (verify_wait, "infer.verify_wait")): + register_translatable(comp, key) spec_note = gr.Markdown(t("infer.spec_note")) register_translatable(spec_note, "infer.spec_note") @@ -1959,7 +1708,9 @@ def _sniff_checkpoint(path: str) -> dict: def generate_unified(cfg, selected_backend, selected_mode, model_path_val, prompt_val, max_tok, temp, raw_prompt_val, miss_policy_val, ram_load_ratio_val, sweep_ratios_val, - speculative_val=False, draft_depth_val=1, gamma_val=4): + speculative_val=False, draft_depth_val=1, gamma_val=4, + residency_planner_val=True, paging_val="auto", + residency_decay_val=None, decode_wait_val=None, verify_wait_val=None): raw = generate_api( cfg, selected_backend=selected_backend, @@ -1975,6 +1726,11 @@ def generate_unified(cfg, selected_backend, selected_mode, model_path_val, speculative=speculative_val, draft_depth=draft_depth_val, gamma=gamma_val, + residency_planner=residency_planner_val, + paging=paging_val, + residency_decay=residency_decay_val, + decode_wait_s=decode_wait_val, + verify_wait_s=verify_wait_val, ) if raw.get("error"): text = f"Error: {raw.get('error')}" @@ -2018,7 +1774,8 @@ def generate_unified(cfg, selected_backend, selected_mode, model_path_val, inputs=[ config_state, inference_backend, inference_mode, model_path, prompt, max_tokens, temperature, raw_prompt, miss_policy, ram_load_ratio, sweep_ratios, - speculative, draft_depth, spec_gamma, + speculative, draft_depth, spec_gamma, residency_planner, + paging_mode, residency_decay, decode_wait, verify_wait, ], outputs=[output_box, lazy_output_box, full_output_box, stats_md, stats_chart, raw_stats], api_name=None, @@ -2031,7 +1788,8 @@ def generate_unified(cfg, selected_backend, selected_mode, model_path_val, inputs=[ config_state, inference_backend, inference_mode, model_path, prompt, max_tokens, temperature, raw_prompt, ram_load_ratio, sweep_ratios, miss_policy, - speculative, draft_depth, spec_gamma, + speculative, draft_depth, spec_gamma, residency_planner, + paging_mode, residency_decay, decode_wait, verify_wait, ], outputs=[api_out], api_name="generate", From de469625e2ab5982ea701ae44a11201f565fd73f Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Sat, 7 Nov 2026 16:15:45 +0800 Subject: [PATCH 02/18] 0.3.1a1: new Web UI, file picking/upload, LAN access, lazy-expert probe - chronos-ui / chronos_app.py start the rewritten Web UI (FastAPI + single-page app in chronos/webui, logic in chronos/workbench); --legacy keeps Gradio. - zh-Hans / zh-Hant / en / ja, browser-language detection, hover help after 2 s; buttons, chips, options and picker places are translated. - Path fields: "Browse..." lists the server's folders; "Upload..." opens the native dialog on the server's own machine (path used in place) or uploads to /uploads/ from another machine. Several datasets merge into one JSONL. - --lan (off by default): key from CHRONOS_UI_KEY or generated; the server's own machine needs no sign-in. --max-upload-mb, --no-native-dialog. - chronos lazy-probe and per-validation lazy-expert probe in training runs. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 18 + MANIFEST.in | 2 + README.md | 31 +- README_zh.md | 18 +- chronos/__init__.py | 2 +- chronos/cli.py | 9 + chronos/lazy_experts.py | 4 +- chronos/monitor/snapshot.py | 45 +- chronos/webui/i18n/en.json | 860 ++++++++++ chronos/webui/i18n/ja.json | 860 ++++++++++ chronos/webui/i18n/zh-Hans.json | 860 ++++++++++ chronos/webui/i18n/zh-Hant.json | 860 ++++++++++ chronos/webui/jobs.py | 195 +++ chronos/webui/native_dialog.py | 128 ++ chronos/webui/schema.py | 460 ++++++ chronos/webui/server.py | 1083 +++++++++++++ chronos/webui/static/app.css | 252 +++ chronos/webui/static/app.js | 165 ++ chronos/webui/static/charts.js | 250 +++ chronos/webui/static/core.js | 369 +++++ chronos/webui/static/forms.js | 219 +++ chronos/webui/static/icon.svg | 1 + chronos/webui/static/index.html | 58 + chronos/webui/static/pages_config.js | 180 +++ chronos/webui/static/pages_data.js | 93 ++ chronos/webui/static/pages_infer.js | 94 ++ chronos/webui/static/pages_monitor.js | 162 ++ chronos/webui/static/pages_ops.js | 144 ++ chronos/webui/static/pages_probe.js | 96 ++ chronos/webui/static/pages_train.js | 84 + chronos/webui/static/picker.js | 334 ++++ chronos/workbench/__init__.py | 1 + chronos/workbench/data.py | 150 ++ chronos/workbench/files.py | 226 +++ chronos/workbench/inference.py | 1472 +++++++++++++++++ chronos/workbench/lazy_probe.py | 567 +++++++ chronos/workbench/serve.py | 116 ++ chronos/workbench/tokenizer.py | 114 ++ chronos/workbench/train_worker.py | 306 ++++ chronos/workbench/training.py | 2062 ++++++++++++++++++++++++ chronos_app.py | 19 +- pyproject.toml | 5 +- tests/test_lazy_probe.py | 241 +++ tests/test_smoke.py | 12 +- tests/test_ui_event_browser.py | 2 +- tests/test_webui.py | 503 ++++++ ui/tabs/data_tab.py | 137 +- ui/tabs/inference_tab.py | 1515 +----------------- ui/tabs/serve_tab.py | 104 +- ui/tabs/tokenizer_tab.py | 117 +- ui/tabs/train_tab.py | 2097 +------------------------ 51 files changed, 13811 insertions(+), 3891 deletions(-) create mode 100644 chronos/webui/i18n/en.json create mode 100644 chronos/webui/i18n/ja.json create mode 100644 chronos/webui/i18n/zh-Hans.json create mode 100644 chronos/webui/i18n/zh-Hant.json create mode 100644 chronos/webui/jobs.py create mode 100644 chronos/webui/native_dialog.py create mode 100644 chronos/webui/schema.py create mode 100644 chronos/webui/server.py create mode 100644 chronos/webui/static/app.css create mode 100644 chronos/webui/static/app.js create mode 100644 chronos/webui/static/charts.js create mode 100644 chronos/webui/static/core.js create mode 100644 chronos/webui/static/forms.js create mode 100644 chronos/webui/static/icon.svg create mode 100644 chronos/webui/static/index.html create mode 100644 chronos/webui/static/pages_config.js create mode 100644 chronos/webui/static/pages_data.js create mode 100644 chronos/webui/static/pages_infer.js create mode 100644 chronos/webui/static/pages_monitor.js create mode 100644 chronos/webui/static/pages_ops.js create mode 100644 chronos/webui/static/pages_probe.js create mode 100644 chronos/webui/static/pages_train.js create mode 100644 chronos/webui/static/picker.js create mode 100644 chronos/workbench/__init__.py create mode 100644 chronos/workbench/data.py create mode 100644 chronos/workbench/files.py create mode 100644 chronos/workbench/inference.py create mode 100644 chronos/workbench/lazy_probe.py create mode 100644 chronos/workbench/serve.py create mode 100644 chronos/workbench/tokenizer.py create mode 100644 chronos/workbench/train_worker.py create mode 100644 chronos/workbench/training.py create mode 100644 tests/test_lazy_probe.py create mode 100644 tests/test_webui.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3de03a2..5d8c843 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,24 @@ ## [Unreleased] +## [0.3.1a1] — 2026-11-07 (beta) + +Web UI half of 0.3.1: the rewritten Web UI, file picking and upload, LAN access, and the lazy-expert probe. + +### Changed +- `chronos-ui` and `python chronos_app.py` start the new Web UI (`chronos/webui/`: FastAPI and a hand-written single-page app, no build step, CSP `script-src 'self'`). The Gradio console stays available as `--legacy`. +- The UI is complete in zh-Hans, zh-Hant, en and ja. It follows the browser language and falls back to English. Buttons, chips, dropdown options and picker shortcuts are translated; device and format ids (cuda, float16, gguf) are shown as they are. +- Hover help appears after the pointer rests for 2 s; the "?" mark shows it after 0.4 s and Tab focus shows it at once. + +### Added +- Path fields have two ways to pick a path. + - "Browse…" lists folders on the server's machine. + - "Upload…" uses the system file manager. On the server's own machine (loopback) the server opens the native dialog (macOS `osascript`, Windows PowerShell, Linux `zenity`/`kdialog`) and the chosen path is used in place. From another machine the files are streamed to `/uploads///` and that path is used. + - Uploads refuse traversal, hidden names and executable or script types. They stop while less than max(8 GiB, 5 % of the disk) would stay free. `--max-upload-mb` caps a single file (default 20480). +- Dataset fields take several files or folders, one per line. Training stages and export calibration merge them into one JSONL file (`chronos/workbench/files.py:merge_datasets`, with a `.sources.json` manifest). An unchanged selection reuses the earlier merge. +- `--lan` opens the UI to the local network; it is off by default. Browsers on other machines sign in with `CHRONOS_UI_KEY`, or with a key printed at start when none is set. The server's own machine needs no sign-in. The key crosses plain HTTP in clear text unless `--ssl-certfile`/`--ssl-keyfile` or a TLS proxy is used. `--no-native-dialog` turns the native dialog off (for example over SSH). +- `chronos lazy-probe`: lazy-expert efficiency, hit rate and fidelity of a checkpoint, or of each validation of a run with `--watch`; the Train page can run it at every validation. + ## [0.3.1a0] — 2026-11-07 (beta) Core half of 0.3.1: cluster-locality training on by default, inference that diff --git a/MANIFEST.in b/MANIFEST.in index 5d2942f..babb05d 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,5 +1,7 @@ include README.md README_zh.md LICENSE THIRD_PARTY_NOTICES.md CHANGELOG.md requirements.txt recursive-include chronos/monitor/static *.html *.js *.css +recursive-include chronos/webui/static *.html *.js *.css *.svg +recursive-include chronos/webui/i18n *.json include ui/event_contract.js global-exclude __pycache__ *.py[cod] .DS_Store prune runs diff --git a/README.md b/README.md index 0e443ad..55dfd16 100644 --- a/README.md +++ b/README.md @@ -470,30 +470,19 @@ pip install vllm ## Quick start -### Web UI (12 tabs, 4 languages) +### Web UI (11 pages, 4 languages) ```bash -chronos-ui -# or -python chronos_app.py +chronos-ui # or: python chronos_app.py +chronos-ui --lan # also open it to the local network (off by default) +chronos-ui --legacy # the previous Gradio console ``` -Tabs included: +A FastAPI server with a single-page app (`chronos/webui/`). It listens on 127.0.0.1 without a login. Pages: Model (v3 presets, the auto planner below, a live parameter / memory estimate), Data (scan one or several datasets), Tokenizer (extract → train → evaluate the Chronos BPE), Train (stages in order, each with its own data; optional lazy-expert probe at every validation), Monitor (read-only view of any run folder, including runs started from the command line), Inference (lazy experts vs full memory on one prompt), Export, Serve (start and stop `chronos serve`), Benchmark, Auto-tune, IO metrics. -- `Config`: v3 presets, the auto planner (below) and a live parameter / memory estimator -- `Data` for dataset scanning and conversion -- `Tokenizer`: extract text from datasets, train the Chronos BPE, evaluate and save -- `Train` with its own `data_path`; applies the selected stage's overrides -- `Training Monitor` (read-only view of `runs/`; see below) -- `6-Stage Pipeline` with per-stage dataset paths -- `Inference` -- `Export` for FP16/Q8_0 safetensors and GGUF deployment artifacts -- `Serve`: start and stop `chronos serve` and view its log -- `Benchmark` with Markdown table + comparison plots -- `Auto-Tune` with persistent logs and one-click `Apply Best -> Config` -- `IO Monitor` - -At a 1280 px wide window, Gradio moves the last tabs (Auto-Tune, IO Monitor) into its "…" overflow menu. +- **Languages.** zh-Hans, zh-Hant, en and ja; the page follows the browser language and falls back to English. Every setting has hover help with its meaning and recommended value (after a 2 s rest; the "?" mark after 0.4 s). +- **Paths.** "Browse…" lists folders on the server's machine. "Upload…" uses the system file manager: on the server's own machine the server opens the native dialog and the file is used where it is; from another machine the files are uploaded to `/uploads/` and that path is used. Dataset fields take several files or folders, one per line; training and export calibration merge them into one JSONL file. +- **LAN access** (`--lan`). Other machines sign in with `CHRONOS_UI_KEY`, or with a key printed at start when none is set; the server's own machine needs no sign-in. A signed-in browser can train and read or write files with the server user's rights, so use it on a trusted network only. The key crosses plain HTTP in clear text unless `--ssl-certfile` / `--ssl-keyfile` or a TLS proxy is used. `--max-upload-mb` caps a single upload (default 20480); `--no-native-dialog` turns the native dialog off, for example over SSH. `--allow-remote` with a key in the environment binds any other host. ### Presets and the auto planner @@ -782,6 +771,8 @@ Project_Chronos/ │ │ └── corpus/ # download / clean / dedup / leak filter / shards │ ├── tokenizer/ # Chronos BPE builder (extract, normalize, train, evaluate) │ ├── monitor/ # Read-only run monitor: snapshot, API server, static page +│ ├── webui/ # Web UI: FastAPI server, field schema, i18n/*.json, static/ SPA +│ ├── workbench/ # Web UI jobs: train worker, data scan, file picking / merge, lazy probe │ ├── backend/ │ │ ├── __init__.py # BackendDispatcher (cpu/mps/cuda/mlx) │ │ ├── dispatcher.py # Capability probing + priority logic @@ -793,7 +784,7 @@ Project_Chronos/ │ │ ├── runtime.py # Model discovery, lazy experts, speculative generation │ │ └── vllm_adapter.py # Optional vLLM registration │ └── cli.py # Unified CLI -├── ui/ # Gradio Web UI (zh-Hans / zh-Hant / en / ja) +├── ui/ # Legacy Gradio console (chronos-ui --legacy) │ ├── i18n.py │ ├── estimator.py # Live parameter / memory estimator │ └── tabs/ diff --git a/README_zh.md b/README_zh.md index 23f76ef..09ecbea 100644 --- a/README_zh.md +++ b/README_zh.md @@ -471,15 +471,19 @@ pip install vllm ## 快速开始 -### Web UI（12 个 Tab，4 种语言） +### Web UI（11 个页面，4 种语言） ```bash -chronos-ui -# 或 -python chronos_app.py +chronos-ui # 或：python chronos_app.py +chronos-ui --lan # 同时开放局域网访问（默认关闭） +chronos-ui --legacy # 旧版 Gradio 控制台 ``` -包含：⚙️ 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）收进 “…” 溢出菜单。 +FastAPI 服务加单页前端（`chronos/webui/`），默认只监听 127.0.0.1，无需登录。页面包括：模型（v3 预设、下文的自动规划器、实时参数与内存估算）、数据（扫描一个或多个数据集）、分词器（抽取 → 训练 → 评估 Chronos BPE）、训练（按顺序执行各阶段，每阶段可单独指定数据；可在每次验证时探测懒加载）、监控（只读查看任意运行目录，包括命令行启动的训练）、推理（同一提示词对比懒加载专家与全内存）、导出、服务（启停 `chronos serve`）、基准测试、自动调参、IO 指标。 + +- **语言。** 支持 zh-Hans、zh-Hant、en、ja，自动跟随浏览器语言，不支持的语言回退到英文。每个设置项都有悬浮说明，给出含义和推荐值（光标停留 2 秒后显示；“?” 标记 0.4 秒后显示）。 +- **路径。** “浏览…”列出服务所在电脑上的文件夹。“上传…”使用系统文件管理器：在服务所在的电脑上由服务打开系统对话框，直接使用原路径；从其他设备访问时，文件上传到 `/uploads/` 后使用该路径。数据集字段可以填多个文件或文件夹（每行一个），训练和导出校准会把它们合并成一个 JSONL 文件。 +- **局域网访问**（`--lan`）。其他设备用 `CHRONOS_UI_KEY` 登录；未设置时启动时会生成并打印一个密钥。服务所在的电脑无需登录。登录后的浏览器能以服务用户的权限训练、读写文件，只应在可信网络中使用。不配 `--ssl-certfile` / `--ssl-keyfile` 或 TLS 反向代理时，密钥以明文经过 HTTP。`--max-upload-mb` 限制单个上传文件（默认 20480）；`--no-native-dialog` 关闭系统对话框（例如通过 SSH 运行时）。绑定其他地址需要 `--allow-remote` 并在环境变量中提供密钥。 ### 预设与自动规划器 @@ -776,6 +780,8 @@ Project_Chronos/ │ │ └── corpus/ # 下载 / 清洗 / 去重 / 泄漏过滤 / 分片 │ ├── tokenizer/ # Chronos BPE 构建器（抽取、规范化、训练、评估） │ ├── monitor/ # 只读运行监控：快照、API 服务、静态页面 +│ ├── webui/ # Web UI：FastAPI 服务、字段 schema、i18n/*.json、static/ 单页前端 +│ ├── workbench/ # Web UI 任务：训练 worker、数据扫描、文件选取与合并、懒加载探针 │ ├── backend/ │ │ ├── __init__.py # BackendDispatcher（cpu/mps/cuda/mlx） │ │ ├── dispatcher.py # 探针 + 优先级 + 训练能力声明 @@ -787,7 +793,7 @@ Project_Chronos/ │ │ ├── runtime.py # 模型发现、懒加载专家、推测生成 │ │ └── vllm_adapter.py # 可选 vLLM 注册（无 vLLM 时优雅降级） │ └── cli.py # 统一 CLI -├── ui/ # Gradio Web UI（i18n: zh-Hans/zh-Hant/en/ja） +├── ui/ # 旧版 Gradio 控制台（chronos-ui --legacy） │ ├── i18n.py │ ├── estimator.py # 实时参数量/内存估算（与 Config 同步） │ └── tabs/ diff --git a/chronos/__init__.py b/chronos/__init__.py index 8fdf6b3..6b93853 100644 --- a/chronos/__init__.py +++ b/chronos/__init__.py @@ -1,5 +1,5 @@ # Project Chronos: On-device low-latency lookahead dual-layer MoE inference -__version__ = "0.3.1a0" +__version__ = "0.3.1a1" # Bootstrap minimind dependency before any submodule imports from chronos.deps import ensure_minimind as _ensure_minimind diff --git a/chronos/cli.py b/chronos/cli.py index 35514d9..22f402a 100644 --- a/chronos/cli.py +++ b/chronos/cli.py @@ -14,6 +14,8 @@ chronos corpus — download / clean / dedup / shard training corpora chronos longrun — time-budgeted resumable MLX training run (alias: run) chronos monitor — read-only training monitor page + JSON API + chronos lazy-probe — lazy-expert efficiency / hit rate / fidelity of a checkpoint or a run + chronos webui — Web UI (alias: ui) """ import sys import os @@ -114,6 +116,10 @@ def main(): sub.add_parser('longrun', help='Time-budgeted resumable MLX training run', aliases=['run'], add_help=False) sub.add_parser('monitor', help='Read-only training monitor page + JSON API', add_help=False) sub.add_parser('bundle', help='Out-of-the-box model bundle: build/check/translate/finetune/serve', add_help=False) + sub.add_parser('lazy-probe', help='Lazy-expert efficiency/hit-rate/fidelity probe (checkpoint or --watch run)', + add_help=False) + sub.add_parser('webui', help='Web UI (FastAPI; --legacy for the Gradio console)', aliases=['ui'], + add_help=False) known, rest = parser.parse_known_args() longrun = lambda rest: _module_main('chronos.mlx.training.longrun', rest) # noqa: E731 @@ -132,6 +138,9 @@ def main(): 'run': longrun, 'monitor': lambda rest: _module_main('chronos.monitor.server', rest), 'bundle': lambda rest: _module_main('chronos.bundle.cli', rest), + 'lazy-probe': lambda rest: _module_main('chronos.workbench.lazy_probe', rest), + 'webui': lambda rest: _module_main('chronos.webui.server', rest), + 'ui': lambda rest: _module_main('chronos.webui.server', rest), } dispatch[known.command](rest) diff --git a/chronos/lazy_experts.py b/chronos/lazy_experts.py index 632d831..e16b150 100644 --- a/chronos/lazy_experts.py +++ b/chronos/lazy_experts.py @@ -92,7 +92,9 @@ def checkpoint_cache_dir(model_path: str, config, *, backend: str = "torch", tag per-expert caches keep their historical names (``ckpt_`` for Torch, ``ckpt__mlx`` for MLX); cluster layouts add ``_``. """ - base = root or os.path.join(os.getcwd(), "expert_cache") + # ``CHRONOS_EXPERT_CACHE_ROOT`` lets a dedicated process (the validation + # probe) keep its caches in a scratch directory it deletes afterwards. + base = root or os.environ.get("CHRONOS_EXPERT_CACHE_ROOT") or os.path.join(os.getcwd(), "expert_cache") if not model_path or not os.path.exists(model_path): name_dir = base else: diff --git a/chronos/monitor/snapshot.py b/chronos/monitor/snapshot.py index 135cad0..3a3a2e7 100644 --- a/chronos/monitor/snapshot.py +++ b/chronos/monitor/snapshot.py @@ -137,7 +137,9 @@ def _status(self, path: str, spec: dict | None, rows: list[dict]) -> dict: last = rows[-1] if rows else {} finished = any(r.get("event") == "final" for r in rows[-50:]) declared = heartbeat.get("state") - if declared in ("finished", "failed") and not alive: + if declared in ("finished", "failed") and (not alive or finished): + # A heartbeat that declares the end, backed by a ``final`` event, wins even when the pid + # is still alive (the writer runs inside a longer-lived process, or the pid was reused). state = declared elif finished and not alive: state = "finished" @@ -229,6 +231,7 @@ def snapshot(self, run: str | None = None, *, history: bool = True) -> dict: self._stream(os.path.join(path, "samples.jsonl"), max_rows=200).read()[-40:] if isinstance(x, dict)], "inspections": self._inspections(path), "checkpoints": self._checkpoints(path), + "lazy_probe": self._lazy_probes(path, limit=20), "counts": {"metric_rows": len(train_rows), "events": len(events), "parameter_rows": len(params)}, } if history: @@ -314,6 +317,46 @@ def _inspections(path: str) -> list[dict]: out.append({"file": os.path.basename(file), **{k: data.get(k) for k in meta}, "metrics": metrics}) return out + # Columns the monitor table shows per probe variant (the full result stays in lazy_probe.jsonl). + PROBE_MEASURED = ("tokens_per_s", "decode_tokens_per_s", "resident_hit_rate", "prediction_hit_rate", + "cold_miss_predict_hit_rate", "on_demand_loads", "prefetch_wait_time_s", "peak_rss_gb", + "spec_acceptance_rate", "speed_vs_full", "paging_resolved", "num_pages", + "effective_expert_budget", "tokens") + PROBE_FIDELITY = ("exact_match_rate", "positionwise_agreement", "min_matching_prefix") + + def _lazy_probes(self, path: str, limit: int = 20) -> list[dict]: + """Mid-training lazy-expert probe results (``lazy_probe.jsonl``), newest last. + + Each entry keeps phase / step / wall time and, per variant, the label, + settings, measured efficiency + hit rates and fidelity versus the + full-memory reference; skipped requests keep their reason.""" + out = [] + for row in self._stream(os.path.join(path, "lazy_probe.jsonl"), max_rows=500).read()[-limit:]: + if not isinstance(row, dict): + continue + entry = {k: row.get(k) for k in ("time", "phase", "step", "checkpoint", "wall_s", "skipped", "error") if k in row} + result = row.get("result") or {} + if result.get("skipped"): + entry["skipped"] = result["skipped"] + if result.get("variants") is not None: + entry.update({k: result.get(k) for k in ("trained_with_clusters", "auto_paging", "num_experts", + "top_k", "truncated", "variants_requested", "elapsed_s", + "conditions")}) + entry["references"] = {b: (r.get("measured") or {}).get("tokens_per_s") + for b, r in (result.get("references") or {}).items()} + entry["variants"] = [{ + "label": v.get("label"), "settings": v.get("settings"), + **({"error": v["error"]} if "error" in v else {}), + "measured": {k: (v.get("measured") or {}).get(k) for k in self.PROBE_MEASURED}, + "fidelity": {k: (v.get("fidelity") or {}).get(k) for k in self.PROBE_FIDELITY}, + } for v in result["variants"]] + out.append(entry) + return out + + def lazy_probes(self, run: str | None, limit: int = 200) -> list[dict]: + path = self.run_dir(run) + return self._lazy_probes(path, limit=limit) if path else [] + @staticmethod def _checkpoints(path: str) -> list[dict]: files = glob.glob(os.path.join(path, "*.pth")) + glob.glob(os.path.join(path, "checkpoints", "*.pkl")) diff --git a/chronos/webui/i18n/en.json b/chronos/webui/i18n/en.json new file mode 100644 index 0000000..8f239b5 --- /dev/null +++ b/chronos/webui/i18n/en.json @@ -0,0 +1,860 @@ +{ + "app.subtitle": "Recurrent-depth MoE · lazy experts · speculative decoding", + "app.title": "Chronos Workbench", + "bench.better": "Better", + "bench.higher": "Higher is better", + "bench.intro": "Compare MiniMind and Chronos on parameters, memory and decode speed with benchmark_compare.py.", + "bench.lower": "Lower is better", + "bench.measured_note": "Measured by benchmark_compare.py on this machine.", + "bench.metric": "Metric", + "bench.none": "No results yet. Run the benchmark to create them.", + "bench.params": "Parameters", + "bench.ram_infer": "Inference RAM", + "bench.ram_model": "Model RAM", + "bench.ram_train": "Training RAM", + "bench.run": "Run benchmark", + "bench.title": "Benchmark: MiniMind vs Chronos", + "bench.tps": "Decode tokens/s", + "bench.unavailable": "benchmark_compare.py is not in this installation.", + "choice.absorbed": "Absorbed", + "choice.adaptive": "Adaptive", + "choice.aggressive": "Aggressive", + "choice.auto": "Auto", + "choice.balanced": "Balanced", + "choice.cluster": "By cluster", + "choice.compare": "Compare lazy vs full", + "choice.expanded": "Expanded (reference)", + "choice.expert": "By expert", + "choice.experts": "Active experts", + "choice.extreme": "Extreme", + "choice.fallback_diagnostic": "Never load, diagnostic (fallback_diagnostic)", + "choice.fixed": "Fixed", + "choice.fixed_then_adaptive": "Fixed, then adaptive", + "choice.fullload": "Full memory only", + "choice.none": "None", + "choice.offload": "Lazy experts only", + "choice.on_demand": "Load, may skip late experts (on_demand)", + "choice.params": "Active parameters", + "choice.planned": "Next-use planner", + "choice.recurrent_attention_mode.first_only": "First iteration only, v2 (first_only)", + "choice.recurrent_attention_mode.per_iteration": "Every iteration, own KV each (per_iteration)", + "choice.recurrent_attention_mode.shared_first": "Every iteration, shared first KV (shared_first)", + "choice.router_scope.per_iteration": "Route every iteration (per_iteration)", + "choice.router_scope.shared": "Reuse iteration 0 routing (shared)", + "choice.scheduled": "Scheduled", + "choice.shared": "Shared", + "choice.sync_on_demand": "Load and wait (sync_on_demand)", + "choice.tune.search.lr": "Learning rate", + "choice.window": "Recent-use window", + "config.active": "Active per token", + "config.compute": "Weight uses per token (with recurrence)", + "config.estimate": "Size and memory estimate", + "config.estimate_note": "Estimated from the configuration, not measured.", + "config.file": "Config file", + "config.ignored": "Ignored unknown or invalid keys", + "config.intro": "Every model setting in one place. Hover any option for its meaning, default and recommended value.", + "config.kv": "KV cache", + "config.loaded": "Configuration loaded", + "config.path_help": "A JSON config file. Load reads it into the form; Save writes the form there. Relative paths start at the workbench folder.", + "config.preset": "Preset", + "config.preset_help": "A starting configuration. Applying it replaces every field on this page; the estimate updates at once.", + "config.ram": "Warm expert cache", + "config.ssd": "Expert storage on disk", + "config.total": "Total parameters", + "config.tps": "Rough decode speed (estimated)", + "config.train_peak": "Training peak memory", + "config.train_state": "Training state (weights, gradients, optimizer)", + "config.virtual_layers": "Virtual layers", + "config.vram": "Resident weights + KV", + "data.chars": "Characters per record", + "data.dups": "Duplicate share", + "data.fields": "Fields", + "data.files": "Files", + "data.full_scan": "full scan", + "data.hist": "Record length distribution", + "data.intro": "Scan a file or folder to see its format, fields, scripts and record lengths before training.", + "data.max_help": "Upper bound on records read for the summary (100 to 500,000).", + "data.max_records": "Records to scan", + "data.median": "Median", + "data.p95": "95th percentile", + "data.path_help": "Files or folders (jsonl, json, parquet, csv, tsv, txt), one per line. Several are scanned as one dataset, as training sees them after merging. Relative paths start at the workbench folder.", + "data.present": "Present in", + "data.records": "Records", + "data.sampled": "first {n} records", + "data.samples": "Sample records", + "data.scanned": "Scanned", + "data.schema": "Detected format", + "data.script": "Script", + "data.scripts": "Writing systems", + "data.share": "Share", + "data.title": "Dataset inspector", + "data.type": "Type", + "err.active_experts": "Active experts cannot exceed the expert count", + "err.dialog_open": "A file dialog is already open; finish or cancel it in the system window first", + "err.models_dir": "Models folder not found", + "err.native_local_only": "The system file dialog works only on the server's own machine", + "err.native_unavailable": "This machine cannot open a system file dialog", + "err.no_files": "No readable files under those paths", + "err.no_out": "Choose an output folder", + "err.no_paths": "Add at least one path", + "err.not_found": "Not found", + "err.paging": "Choose auto, cluster or expert paging", + "err.port": "Invalid port", + "err.probe_invalid": "The lazy-loading check settings are invalid", + "err.remote_needs_key": "Remote binding needs Allow remote and CHRONOS_API_KEY", + "err.target": "Enter a target size", + "err.upload_batch": "Invalid upload batch id; reload the page and try again", + "err.upload_disk_full": "Not enough free disk space on the server's machine; upload stopped", + "err.upload_failed": "Upload interrupted; check the connection and try again", + "err.upload_kind": "This field does not take uploads", + "err.upload_name": "File name or folder depth not allowed (hidden item, \"..\" or too deep)", + "err.upload_too_large": "File exceeds the upload limit (--max-upload-mb)", + "err.upload_type": "Nothing to upload: the file type does not fit this field, or it is a blocked executable", + "event.final": "Final", + "event.lazy_probe": "Lazy-expert probe", + "event.phase_end": "Stage end", + "event.phase_start": "Stage start", + "event.resume": "Resumed", + "event.rollback": "Rollback", + "event.validation": "Validation", + "export.done": "Exported", + "export.format": "Format", + "export.intro": "Write a checkpoint as safetensors or GGUF files plus the expert cache layout chronos serve loads.", + "export.note": "Ollama clients run Chronos models through chronos serve (set OLLAMA_HOST to its address); stock Ollama cannot execute the recurrent MoE.", + "export.run": "Export", + "export.size": "Size (MB)", + "export.tensors": "Tensors", + "export.title": "Export", + "field.accumulation_steps.help": "Micro-batches per optimizer update. Effective batch = batch size × accumulation.", + "field.accumulation_steps.label": "Gradient accumulation", + "field.activation_checkpointing.help": "Recompute activations in the backward pass to cut peak memory at the cost of speed.", + "field.activation_checkpointing.label": "Activation checkpointing", + "field.adaptive_halting.help": "Let each token stop looping once its halting probability passes the threshold. Saves compute on easy tokens.", + "field.adaptive_halting.label": "Adaptive halting (ACT)", + "field.alpha.help": "Distillation mix: α × teacher KL + (1 − α) × CE.", + "field.alpha.label": "α (distill)", + "field.batch_size.help": "Sequences per micro-batch. Limited by memory.", + "field.batch_size.label": "Batch size", + "field.beta.help": "Strength of the KL tie to the reference policy in preference training.", + "field.beta.label": "β (DPO/GRPO)", + "field.cluster_locality_clusters.help": "Number of expert clusters for the locality term. 0 chooses it from the expert count.", + "field.cluster_locality_clusters.label": "Clusters", + "field.cluster_locality_window.help": "Segment length (tokens) the cluster constraint applies to. Best results matched the typical record length.", + "field.cluster_locality_window.label": "Locality window", + "field.coda_layers.help": "Layers that run once after the loop, before the output head.", + "field.coda_layers.label": "Coda layers", + "field.cpu_budget_percent.help": "Share of physical cores auto may use, in percent.", + "field.cpu_budget_percent.label": "CPU budget", + "field.cpu_threads.help": "Threads for PyTorch/BLAS. auto uses the physical-core budget.", + "field.cpu_threads.label": "CPU threads", + "field.depth_schedule.help": "How training chooses the depth per batch: fixed, adaptive, or fixed for a warm-up and then adaptive.", + "field.depth_schedule.label": "Depth schedule", + "field.depth_schedule_warmup_fraction.help": "Fraction of training that runs at fixed full depth before the schedule becomes adaptive.", + "field.depth_schedule_warmup_fraction.label": "Schedule warm-up", + "field.draft_exit_depth.help": "Recurrent depth whose exit is trained as the speculative draft.", + "field.draft_exit_depth.label": "Draft depth (training)", + "field.draft_exit_kl_coef.help": "Extra KL between the draft-depth exit and the full-depth output, for self-speculative decoding. 0 disables it.", + "field.draft_exit_kl_coef.label": "Draft-exit KL", + "field.dropout.help": "Random activation dropout during training. Lower it for large datasets; raise it for tiny ones.", + "field.dropout.label": "Dropout", + "field.dtype.help": "Numeric precision for training and inference. Auto picks bfloat16 where the backend supports it, otherwise float32.", + "field.dtype.label": "Compute dtype", + "field.epochs.help": "Passes over the dataset. Max steps, when set, stops earlier.", + "field.epochs.label": "Epochs", + "field.export.auto_cluster.help": "Group co-used experts into pages from calibration data, so lazy loading reads fewer files.", + "field.export.auto_cluster.label": "Build clusters", + "field.export.calibration_data_path.help": "Text used to measure which experts are used together, one path per line; several are merged. Empty skips automatic clustering.", + "field.export.calibration_data_path.label": "Calibration data", + "field.export.cluster_batch_size.help": "Sequences per calibration batch.", + "field.export.cluster_batch_size.label": "Calibration batch size", + "field.export.cluster_device.help": "Device for the calibration pass.", + "field.export.cluster_device.label": "Calibration device", + "field.export.cluster_max_batches.help": "Number of batches observed for clustering.", + "field.export.cluster_max_batches.label": "Calibration batches", + "field.export.cluster_max_seq_len.help": "Tokens per calibration sequence.", + "field.export.cluster_max_seq_len.label": "Calibration length", + "field.export.config_path.help": "Override the sidecar config. Usually empty.", + "field.export.config_path.label": "Config file", + "field.export.copy_cache.help": "Include the lazy-expert pages in the export so serving starts without rebuilding.", + "field.export.copy_cache.label": "Copy expert cache", + "field.export.expert_cache_dir.help": "Existing lazy-expert cache to package. Empty builds one.", + "field.export.expert_cache_dir.label": "Expert cache", + "field.export.formats.help": "Containers to write. Safetensors and GGUF keep the Chronos metadata; run them with chronos serve. Ollama clients connect to chronos serve, not stock Ollama.", + "field.export.formats.label": "Formats", + "field.export.model_path.help": "Checkpoint to export.", + "field.export.model_path.label": "Checkpoint", + "field.export.output_dir.help": "Where the export files are written.", + "field.export.output_dir.label": "Output folder", + "field.fallback_mask_prob.help": "During training, randomly hide experts so the model tolerates a cache miss. Only matters for lazy inference with fallback.", + "field.fallback_mask_prob.label": "Miss simulation", + "field.grad_clip.help": "Clip the global gradient norm to this value. 0 disables clipping.", + "field.grad_clip.label": "Gradient clip", + "field.halting_threshold.help": "Cumulative halting probability at which a token stops iterating. Closer to 1 runs deeper.", + "field.halting_threshold.label": "Halting threshold", + "field.hidden_size.help": "Width of every token vector. Larger is more capable but grows parameters roughly quadratically. Must be divisible by the number of attention heads.", + "field.hidden_size.label": "Hidden size", + "field.infer.backend.help": "Runtime for inference. Auto chooses from the checkpoint and the machine.", + "field.infer.backend.label": "Backend", + "field.infer.decode_wait_s.help": "Wait for a prefetched expert before loading it on demand.", + "field.infer.decode_wait_s.label": "Decode wait (s)", + "field.infer.draft_depth.help": "Recurrent depth used for drafting.", + "field.infer.draft_depth.label": "Draft depth", + "field.infer.gamma.help": "Tokens drafted per verify pass.", + "field.infer.gamma.label": "Draft length γ", + "field.infer.max_tokens.help": "Upper bound on generated tokens.", + "field.infer.max_tokens.label": "Max new tokens", + "field.infer.miss_policy.help": "What a cache miss does. sync_on_demand gives output identical to full memory.", + "field.infer.miss_policy.label": "Miss policy", + "field.infer.mode.help": "compare runs lazy and full memory side by side; offload runs lazy experts only; fullload loads everything.", + "field.infer.mode.label": "Mode", + "field.infer.model_path.help": "The .pth checkpoint to run. Its sidecar config decides the architecture.", + "field.infer.model_path.label": "Checkpoint", + "field.infer.paging.help": "auto follows the checkpoint: cluster pages when trained with cluster locality, otherwise per expert.", + "field.infer.paging.label": "Paging layout", + "field.infer.ram_load_ratio.help": "Fraction of experts kept in RAM. Above 1.0 adds slack for prefetching.", + "field.infer.ram_load_ratio.label": "RAM ratio", + "field.infer.raw_prompt.help": "Send the text as is, without the chat template.", + "field.infer.raw_prompt.label": "Raw prompt", + "field.infer.residency_decay.help": "How fast the planner forgets older routing evidence.", + "field.infer.residency_decay.label": "Planner decay", + "field.infer.residency_planner.help": "Evict the expert needed furthest in the future instead of the least recently used one.", + "field.infer.residency_planner.label": "Next-use planner", + "field.infer.speculative.help": "Draft tokens at shallow depth and verify them at full depth. Output is unchanged; at 50M it was measured slower, so it is off by default.", + "field.infer.speculative.label": "Self-speculative decoding", + "field.infer.sweep_ram_ratios.help": "Run every standard ratio and chart speed and hit rate against memory.", + "field.infer.sweep_ram_ratios.label": "Sweep RAM ratios", + "field.infer.temperature.help": "Sampling randomness. 0 is greedy and reproducible.", + "field.infer.temperature.label": "Temperature", + "field.infer.verify_wait_s.help": "Same wait for the speculative verify pass.", + "field.infer.verify_wait_s.label": "Verify wait (s)", + "field.kv_cache_dtype.help": "auto keeps the compute dtype; int8 quantizes the inference KV cache per token (about half the memory, small error). Training is never quantized.", + "field.kv_cache_dtype.label": "KV cache dtype", + "field.kv_latent_dim.help": "Width of the compressed key/value latent cached by MLA layers. Smaller saves KV memory.", + "field.kv_latent_dim.label": "MLA latent size", + "field.kv_quant_group_size.help": "Channels that share one int8 scale. Smaller groups are more accurate and slightly larger.", + "field.kv_quant_group_size.label": "KV quant group", + "field.lambda_balance.help": "Pushes the router to spread tokens evenly across experts.", + "field.lambda_balance.label": "Load balance", + "field.lambda_cluster_locality.help": "Trains tokens in a segment to stay inside one expert cluster, so lazy inference can page whole clusters. 0 disables it.", + "field.lambda_cluster_locality.label": "Cluster locality", + "field.lambda_lookahead.help": "Trains the lookahead router to predict the experts future tokens will use.", + "field.lambda_lookahead.label": "Lookahead loss", + "field.lambda_lookahead_topk.help": "Extra weight on getting the exact top-k future experts right.", + "field.lambda_lookahead_topk.label": "Lookahead top-k", + "field.lambda_lookahead_union.help": "Weight on covering the union of experts over the lookahead window.", + "field.lambda_lookahead_union.label": "Lookahead union", + "field.lambda_or.help": "Weight of the odds-ratio term in ORPO.", + "field.lambda_or.label": "λ (ORPO)", + "field.lambda_router_anchor.help": "Keeps routing close to the upstream checkpoint during fine-tuning. 0 for pretraining.", + "field.lambda_router_anchor.label": "Router anchor", + "field.lambda_router_locality.help": "Mild penalty on routing changes between neighbouring tokens.", + "field.lambda_router_locality.label": "Router locality", + "field.lambda_temporal.help": "Rewards consecutive tokens for choosing the same experts, so fewer experts are loaded per stretch of text.", + "field.lambda_temporal.label": "Temporal locality", + "field.layer_aware_cache.help": "Track resident experts per virtual layer instead of per physical layer.", + "field.layer_aware_cache.label": "Layer-aware cache", + "field.learning_rate.help": "Peak optimizer step size. The planner suggests one from the model size; lower it for fine-tuning.", + "field.learning_rate.label": "Learning rate", + "field.log_interval.help": "Write a metrics row every N steps.", + "field.log_interval.label": "Log interval", + "field.lookahead_steps.help": "How many future tokens the lookahead router predicts experts for, so they can be prefetched before they are needed.", + "field.lookahead_steps.label": "Lookahead steps", + "field.max_gen_len.help": "Tokens generated per GRPO rollout.", + "field.max_gen_len.label": "Rollout length", + "field.max_inference_depth.help": "Maximum iterations at inference. Usually equal to the training depth; more than trained is not validated.", + "field.max_inference_depth.label": "Inference depth", + "field.max_position_embeddings.help": "Longest sequence the rotary embedding is prepared for. Only change it for very long contexts.", + "field.max_position_embeddings.label": "Max positions", + "field.max_resident_experts.help": "Hard cap on experts kept in memory per layer. Empty lets the planner decide.", + "field.max_resident_experts.label": "Max resident experts", + "field.max_seq_len.help": "Tokens per training sequence. Longer costs memory roughly linearly with hybrid attention.", + "field.max_seq_len.label": "Sequence length", + "field.max_steps.help": "Stop after this many optimizer steps. 0 means no cap (run all epochs).", + "field.max_steps.label": "Max steps", + "field.min_recurrent_depth.help": "Lowest depth the schedule or early exit may use. 1 lets easy tokens exit after one pass.", + "field.min_recurrent_depth.label": "Minimum depth", + "field.mla_cache_rope.help": "Also cache the rotary key part, trading KV memory for decode speed.", + "field.mla_cache_rope.label": "Cache RoPE keys", + "field.mla_decode.help": "expanded: reference path. absorbed: folds projections into the query (less compute per step). auto mixes paths and showed a small logit gap, so it stays opt-in.", + "field.mla_decode.label": "MLA decode path", + "field.modalities.help": "Extra input encoders, comma-separated (image, audio). Leave empty for text-only models; datasets are auto-detected.", + "field.modalities.label": "Modalities", + "field.moe_intermediate_size.help": "Hidden width of each expert's feed-forward network. 0 derives it from the hidden size.", + "field.moe_intermediate_size.label": "Expert FFN width", + "field.num_attention_heads.help": "Query heads. Hidden size must divide evenly by this number.", + "field.num_attention_heads.label": "Attention heads", + "field.num_experts.help": "Routed experts per MoE layer. More experts add capacity without adding per-token compute; they are the unit of lazy loading.", + "field.num_experts.label": "Experts", + "field.num_experts_per_tok.help": "Experts each token uses. Top-2 lets the router learn from the loss; top-1 needs norm_topk_prob off.", + "field.num_experts_per_tok.label": "Active experts (top-k)", + "field.num_generations.help": "Samples per prompt that GRPO compares.", + "field.num_generations.label": "Rollouts per prompt", + "field.num_hidden_layers.help": "Distinct transformer layers stored in the checkpoint (prelude + recurrent groups + coda). Recurrence reuses them, so virtual depth can be larger.", + "field.num_hidden_layers.label": "Physical layers", + "field.num_key_value_heads.help": "Key/value heads for grouped-query attention on sliding-window layers. Fewer heads shrink the KV cache.", + "field.num_key_value_heads.label": "KV heads", + "field.num_shared_experts.help": "Always-on experts every token passes through. They stay resident in memory.", + "field.num_shared_experts.label": "Shared experts", + "field.num_workers.help": "Data-loader worker processes. auto picks a safe number for the platform.", + "field.num_workers.label": "Data workers", + "field.optimizer_page_dir.help": "Folder for paged optimizer state. Use a fast local SSD.", + "field.optimizer_page_dir.label": "Optimizer page dir", + "field.optimizer_paging.help": "Store optimizer state for inactive experts on disk.", + "field.optimizer_paging.label": "Paged optimizer", + "field.page_size_mb.help": "Size of one optimizer page file, in MB.", + "field.page_size_mb.label": "Page size", + "field.pinned_memory_max_fraction.help": "Largest fraction of system RAM used for pinned expert buffers.", + "field.pinned_memory_max_fraction.label": "Pinned RAM share", + "field.plan.active_mode.help": "Set the per-token budget as a number of experts or as active parameters.", + "field.plan.active_mode.label": "Active budget by", + "field.plan.active_value.help": "Active experts per token, or active parameters in millions.", + "field.plan.active_value.label": "Active budget", + "field.plan.data_paths.help": "Files or folders the planner measures to size tokens, sequence length and windows, one per line.", + "field.plan.data_paths.label": "Datasets", + "field.plan.depth.help": "Loop iterations per recurrent group.", + "field.plan.depth.label": "Recurrent depth", + "field.plan.memory_saving.help": "How hard to trade speed for memory (int8 KV, shared routing, checkpointing).", + "field.plan.memory_saving.label": "Memory saving", + "field.plan.num_experts.help": "Routed experts per MoE layer.", + "field.plan.num_experts.label": "Experts", + "field.plan.ram_gb.help": "Memory to plan for. Empty detects this machine.", + "field.plan.ram_gb.label": "RAM (GB)", + "field.plan.seq_len.help": "Empty derives it from the dataset profile.", + "field.plan.seq_len.label": "Sequence length", + "field.plan.target_m.help": "Total parameters in millions.", + "field.plan.target_m.label": "Target size (M)", + "field.plan.token_budget_b.help": "Pretraining tokens in billions. Empty: one pass over the data.", + "field.plan.token_budget_b.label": "Token budget (B)", + "field.plan.use_moe.help": "Plan an MoE model; off plans a dense one.", + "field.plan.use_moe.label": "Mixture of experts", + "field.plan.vocab_size.help": "Empty reads it from the tokenizer report or uses 16384.", + "field.plan.vocab_size.label": "Vocabulary", + "field.ponder_loss_coef.help": "Penalty on the expected number of iterations. Higher pushes the model to stop earlier.", + "field.ponder_loss_coef.label": "Ponder loss", + "field.prefetch_depth.help": "How far ahead the prefetcher loads predicted experts. 0 disables prefetching.", + "field.prefetch_depth.label": "Prefetch depth", + "field.prelude_layers.help": "Layers that run once before the loop. Lookahead routing reads their output.", + "field.prelude_layers.label": "Prelude layers", + "field.probe.backends.help": "Runtimes to test. Each backend gets its own full-memory reference.", + "field.probe.backends.label": "Probe backends", + "field.probe.decode_wait_s.help": "How long a decode step waits for a prefetched expert before loading it on demand.", + "field.probe.decode_wait_s.label": "Decode wait (s)", + "field.probe.draft_depth.help": "Recurrent depth used to draft tokens in speculative variants.", + "field.probe.draft_depth.label": "Draft depth", + "field.probe.enabled.help": "After every validation, decode the fresh checkpoint with lazy experts and compare it with full memory: speed, hit rates and token agreement. Runs in a low-priority child process on a copy of the checkpoint.", + "field.probe.enabled.label": "Lazy-expert probe", + "field.probe.gamma.help": "Tokens drafted per verify pass in speculative variants.", + "field.probe.gamma.label": "Draft length γ", + "field.probe.max_new_tokens.help": "Greedy tokens decoded per prompt. More gives steadier rates but a longer probe.", + "field.probe.max_new_tokens.label": "Tokens per prompt", + "field.probe.max_variants.help": "Upper bound on tested combinations; extra ones are reported as truncated.", + "field.probe.max_variants.label": "Variant cap", + "field.probe.miss_policies.help": "Torch behaviour on a cache miss. sync_on_demand loads and waits (identical output); on_demand may skip late experts; fallback_diagnostic never loads (worst case).", + "field.probe.miss_policies.label": "Miss policies", + "field.probe.paging.help": "auto follows how the model was trained (cluster pages when the cluster-locality loss was on). Add expert to compare per-expert paging.", + "field.probe.paging.label": "Paging layouts", + "field.probe.ram_ratios.help": "Fraction of all experts allowed in RAM at once. Lower saves memory and raises misses.", + "field.probe.ram_ratios.label": "RAM ratios", + "field.probe.residency.help": "MLX eviction: planned uses the next-use planner, window is the plain recent-use window.", + "field.probe.residency.label": "Residency policies", + "field.probe.residency_decay.help": "How fast the next-use planner forgets older routing evidence.", + "field.probe.residency_decay.label": "Planner decay", + "field.probe.speculative.help": "Also test self-speculative decoding with lazy experts.", + "field.probe.speculative.label": "Speculative variants", + "field.probe.verify_wait_s.help": "Same wait for the speculative verify pass.", + "field.probe.verify_wait_s.label": "Verify wait (s)", + "field.ram_budget_mb.help": "RAM the lazy expert cache may use, in MB. Empty uses the RAM ratio instead.", + "field.ram_budget_mb.label": "RAM budget", + "field.recurrent_attention_mode.help": "How attention runs inside the loop. shared_first: every iteration attends, K/V from iteration 0 (exact cache, KV does not grow with depth). per_iteration: one KV per iteration (D× KV). first_only: v2 behaviour.", + "field.recurrent_attention_mode.label": "Recurrent attention", + "field.recurrent_depth.help": "Iterations of each recurrent group during training. Virtual layers = prelude + groups × depth + coda.", + "field.recurrent_depth.label": "Training depth", + "field.recurrent_depth_jitter.help": "Random spread around the scheduled depth, so the model learns to exit at several depths.", + "field.recurrent_depth_jitter.label": "Depth jitter", + "field.recurrent_distill_coef.help": "Teaches shallow exits to match full-depth predictions. Raises speculative acceptance and early-exit quality.", + "field.recurrent_distill_coef.label": "Depth distillation", + "field.recurrent_enabled.help": "Run the shared middle groups several times per token, so a small set of weights acts like a deeper network.", + "field.recurrent_enabled.label": "Recurrent depth", + "field.recurrent_groups.help": "How many layers between prelude and coda form looping groups. Empty derives it from layers − prelude − coda.", + "field.recurrent_groups.label": "Recurrent groups", + "field.recurrent_stability_coef.help": "Keeps the recurrent update contractive so repeated iterations do not drift or explode.", + "field.recurrent_stability_coef.label": "Stability loss", + "field.recurrent_strategy.help": "Sampling strategy for the recurrent depth bucket during training.", + "field.recurrent_strategy.label": "Depth strategy", + "field.reward_spec.help": "toy for the built-in reward, or lm:/path/to/reward-model[,device=cuda].", + "field.reward_spec.label": "Reward (GRPO)", + "field.rope_dim.help": "Dimensions per head that carry rotary position information in MLA layers.", + "field.rope_dim.label": "RoPE dims", + "field.rope_theta.help": "Rotary position base frequency. Larger values stretch positional resolution for longer contexts.", + "field.rope_theta.label": "RoPE base", + "field.router_scope.help": "per_iteration: route again on every loop iteration. shared: reuse iteration 0's routing, which touches fewer experts and suits lazy loading.", + "field.router_scope.label": "Router scope", + "field.save_dir.help": "Where checkpoints and their .config.json sidecars are written.", + "field.save_dir.label": "Checkpoint folder", + "field.save_interval.help": "Save (and validate) every N steps.", + "field.save_interval.label": "Save interval", + "field.serve.allow_remote.help": "Listen on a non-loopback address. Requires CHRONOS_API_KEY in the server's environment; clients must send it.", + "field.serve.allow_remote.label": "Allow remote access", + "field.serve.device.help": "Runtime for serving.", + "field.serve.device.label": "Device", + "field.serve.dtype.help": "Serving precision.", + "field.serve.dtype.label": "Dtype", + "field.serve.host.help": "Address to listen on. 127.0.0.1 is reachable from this machine only.", + "field.serve.host.label": "Host", + "field.serve.kv_cache_dtype.help": "Override the KV storage. int8 roughly halves KV memory; empty keeps the model setting.", + "field.serve.kv_cache_dtype.label": "KV cache dtype", + "field.serve.lazy_experts.help": "Keep experts on SSD and page them in on demand, so memory stays low through the Ollama protocol too.", + "field.serve.lazy_experts.label": "Lazy experts", + "field.serve.max_resident_experts.help": "Hard cap on experts in RAM per layer. Empty: derived from the RAM share.", + "field.serve.max_resident_experts.label": "Max resident experts", + "field.serve.mla_cache_rope.help": "Cache rotated keys: more memory, less recomputation.", + "field.serve.mla_cache_rope.label": "Cache RoPE keys", + "field.serve.mla_decode.help": "Override the MLA decode path. Empty keeps the model setting.", + "field.serve.mla_decode.label": "MLA decode", + "field.serve.models_dir.help": "Folder with exported models; each one appears as an Ollama model name.", + "field.serve.models_dir.label": "Models folder", + "field.serve.paging.help": "auto follows the checkpoint's training (cluster or expert pages).", + "field.serve.paging.label": "Paging layout", + "field.serve.port.help": "TCP port. 11435 avoids clashing with a stock Ollama on 11434.", + "field.serve.port.label": "Port", + "field.serve.ram_budget_mb.help": "Hard cap on expert RAM. Empty: no cap.", + "field.serve.ram_budget_mb.label": "RAM budget (MB)", + "field.serve.ram_expert_fraction.help": "Fraction of experts allowed in RAM.", + "field.serve.ram_expert_fraction.label": "RAM expert share", + "field.serve.router_scope.help": "Override routing reuse across iterations. Empty keeps the model setting.", + "field.serve.router_scope.label": "Router scope", + "field.sliding_window_size.help": "Tokens each sliding-window layer can see. Short windows suit sentence-level tasks and cap KV memory.", + "field.sliding_window_size.label": "Sliding window", + "field.storage_format.help": "File format of the on-disk expert store.", + "field.storage_format.label": "Expert storage", + "field.temperature.help": "Sampling temperature for GRPO rollouts.", + "field.temperature.label": "Rollout temperature", + "field.tie_word_embeddings.help": "Share the input embedding with the output head. Saves vocab × hidden parameters with little quality cost at small sizes.", + "field.tie_word_embeddings.label": "Tie embeddings", + "field.tok.fields.help": "JSON fields to read. Empty auto-detects text, conversations and pairs.", + "field.tok.fields.label": "Text fields", + "field.tok.out_dir.help": "Where tokenizer.json, the chat template and the report are written.", + "field.tok.out_dir.label": "Output folder", + "field.tok.paths.help": "Files or folders to sample text from, one per line. Each source is sampled on its own quota, not merged.", + "field.tok.paths.label": "Source files", + "field.tok.per_source.help": "Reservoir sample size per source file.", + "field.tok.per_source.label": "Texts per source", + "field.tok.temperature.help": "Balances languages by byte share. Lower evens them out, so English does not crowd out Chinese.", + "field.tok.temperature.label": "Language temperature", + "field.tok.vocab_size.help": "Number of tokens to learn. Set the model's vocabulary to the same value.", + "field.tok.vocab_size.label": "Vocabulary size", + "field.tokenizer_path.help": "Tokenizer folder (tokenizer.json) the model is trained with. Empty uses the checkpoint's own tokenizer, else the bundled fallback. The vocabulary size must match.", + "field.tokenizer_path.label": "Tokenizer", + "field.train.backend.help": "Device to train on. Auto picks MLX on Apple silicon, then CUDA, XPU, MPS, CPU.", + "field.train.backend.label": "Training backend", + "field.train.data_distill.help": "Text or chat data scored by the teacher model. Files or folders, one per line; several are merged. Empty reuses the first stage's dataset.", + "field.train.data_distill.label": "Distillation data", + "field.train.data_dpo.help": "Preference pairs with chosen and rejected replies. Files or folders, one per line; several are merged. Empty reuses the first stage's dataset.", + "field.train.data_dpo.label": "DPO data", + "field.train.data_grpo.help": "Prompts for reward rollouts. Files or folders, one per line; several are merged. Empty reuses the first stage's dataset.", + "field.train.data_grpo.label": "GRPO data", + "field.train.data_orpo.help": "Preference pairs with chosen and rejected replies, as for DPO. Files or folders, one per line; several are merged. Empty reuses the first stage's dataset.", + "field.train.data_orpo.label": "ORPO data", + "field.train.data_path.help": "Datasets for the first stage: JSONL, JSON, Parquet, CSV, TSV or TXT files, or folders of them, one per line. Several are merged into one JSONL file in the run folder. Empty runs a synthetic smoke test.", + "field.train.data_path.label": "Dataset", + "field.train.data_sft.help": "Chat data with conversations (user / assistant turns). Files or folders, one per line; several are merged. Empty reuses the first stage's dataset.", + "field.train.data_sft.label": "Supervised fine-tuning data", + "field.train.init_weight.help": "Checkpoint to start from. Empty: pretrain starts from random init, later stages from the previous stage.", + "field.train.init_weight.label": "Start from", + "field.train.max_steps.help": "Stop after this many optimizer steps. 0 runs the configured epochs.", + "field.train.max_steps.label": "Max steps", + "field.train.sample_prompt.help": "Prompt the trainer completes at intervals so you can read progress. Empty uses the stage default.", + "field.train.sample_prompt.label": "Sample prompt", + "field.train.stages.help": "Stages run in order; each starts from the previous stage's checkpoint. Pretrain → SFT is the usual path.", + "field.train.stages.label": "Stages", + "field.train.teacher_path.help": "Distillation only: the model whose outputs the student learns to match.", + "field.train.teacher_path.label": "Teacher checkpoint", + "field.train.val_ratio.help": "Share of records held out for validation. Each validation also triggers the lazy-expert probe.", + "field.train.val_ratio.label": "Validation split", + "field.training_paging.help": "Keep inactive experts out of memory during training. Slower, but trains models larger than RAM.", + "field.training_paging.label": "Paged training", + "field.tune.data_path.help": "Datasets checked before auto-tune starts, one per line. The current probe trials train on random tokens, so this data does not change the trial losses.", + "field.tune.data_path.label": "Dataset", + "field.tune.n_trials.help": "Number of configurations Optuna tries.", + "field.tune.n_trials.label": "Trials", + "field.tune.probe_steps.help": "Training steps per trial; poor trials are pruned early.", + "field.tune.probe_steps.label": "Steps per trial", + "field.tune.search.help": "Parameters the search may change.", + "field.tune.search.label": "Search over", + "field.use_hybrid_attention.help": "Alternate MLA (compressed latent KV) and sliding-window layers. Keeps the KV cache small at long context.", + "field.use_hybrid_attention.label": "Hybrid attention", + "field.vocab_size.help": "Number of token ids. It must equal the tokenizer's vocabulary, otherwise ids fall outside the embedding table.", + "field.vocab_size.label": "Vocabulary size", + "field.vram_budget_gb.help": "Accelerator memory the expert cache may use, in GB.", + "field.vram_budget_gb.label": "VRAM budget", + "field.weight_decay.help": "AdamW decoupled weight decay.", + "field.weight_decay.label": "Weight decay", + "group.attention": "Attention and KV cache", + "group.autotune": "Auto-tune", + "group.benchmark": "Benchmark", + "group.export": "Export", + "group.inference": "Generation", + "group.lazy": "Lazy experts", + "group.memory": "Memory and lazy experts", + "group.model": "Model shape", + "group.moe": "Mixture of experts", + "group.objectives": "Training objectives", + "group.optimizer": "Optimizer", + "group.planner": "Auto plan", + "group.probe": "Lazy-expert probe at validation", + "group.recurrent": "Recurrent depth", + "group.rl": "Preference and RL stages", + "group.runtime": "Runtime", + "group.serve": "Ollama / OpenAI server", + "group.speculative": "Speculative decoding", + "group.stage_data": "Data for later stages", + "group.tokenizer": "Tokenizer builder", + "group.train": "Training job", + "infer.arch_note": "The architecture is read from the checkpoint's config file; the Model page fills only what the file lacks.", + "infer.col.cache_hit_rate": "Cache hit rate", + "infer.col.cold_miss_predict_hit_rate": "Cold-miss predicted", + "infer.col.decode_time_s": "Decode (s)", + "infer.col.fallback_weight_rate": "Fallback weight", + "infer.col.prediction_hit_rate": "Prediction hit rate", + "infer.col.prefill_time_s": "Prefill (s)", + "infer.col.response_time_s": "Response (s)", + "infer.col.rss_after_decode_gb": "RSS after decode (GB)", + "infer.col.rss_delta_gb": "RSS change (GB)", + "infer.col.spec_acceptance_rate": "Draft acceptance", + "infer.col.tokens": "Tokens", + "infer.col.tokens_per_sec": "Tokens/s", + "infer.common_prefix": "first {n} characters identical", + "infer.diff": "Outputs differ", + "infer.full_output": "Full memory", + "infer.generate": "Generate", + "infer.hit_rates": "Hit rates", + "infer.intro": "Run one prompt with lazy experts and with full memory, then compare outputs, speed, hit rates and memory.", + "infer.lazy_output": "Lazy experts", + "infer.measured_note": "Measured on this machine for this prompt; a single run, not an average.", + "infer.mode": "Mode", + "infer.mode.full_dram": "Full memory", + "infer.mode.lazy_offload": "Lazy experts", + "infer.output": "Output", + "infer.prompt": "Prompt", + "infer.prompt_help": "Text sent to the model. The SFT chat template is applied unless Raw prompt is on.", + "infer.ram_ratio": "RAM load ratio", + "infer.same": "Outputs identical", + "infer.similarity": "similarity", + "infer.stats": "Run statistics", + "infer.sweep": "RAM ratio sweep", + "infer.title": "Inference", + "io.empty": "No metrics yet. They appear while inference or serving runs in this process.", + "io.intro": "Expert cache and IO counters recorded by inference and serving in this server process.", + "io.latest": "Latest", + "io.metric": "Metric", + "io.samples": "Samples", + "io.seconds": "Seconds", + "io.title": "IO and cache metrics", + "label.estimated": "estimated", + "label.measured": "measured", + "label.measured_meta": "measured (meta-device count)", + "metric.active_mem_gb": "Active", + "metric.balance_loss": "Expert balance", + "metric.bleu_zh": "BLEU (zh)", + "metric.chrf": "chrF", + "metric.error": "Error", + "metric.eval_s": "Eval time (s)", + "metric.flores_devtest_n": "FLORES sentences", + "metric.grad_norm": "Gradient norm", + "metric.load_cv": "Load imbalance (CV)", + "metric.locality_loss": "Cluster locality", + "metric.lookahead_loss": "Lookahead router", + "metric.loss": "Total loss", + "metric.lr": "Learning rate", + "metric.peak_mem_gb": "Peak", + "metric.primary_loss": "Cross-entropy", + "metric.rec_contraction": "Contraction", + "metric.rec_depth_mean": "Mean depth", + "metric.rec_early_exit_ratio": "Early-exit share", + "metric.route_entropy": "Routing entropy", + "metric.state": "State", + "metric.temporal_loss": "Temporal locality", + "metric.tokens_per_s": "Tokens/s", + "metric.valid_ce": "Validation CE", + "metric.valid_ppl": "Validation perplexity", + "mon.best": "Best", + "mon.chart.aux": "Auxiliary losses", + "mon.chart.loss": "Training loss", + "mon.chart.lr": "Learning rate", + "mon.chart.memory": "Memory (GB)", + "mon.chart.optimizer": "Gradient norm", + "mon.chart.recurrent": "Recurrent depth", + "mon.chart.routing": "Expert routing", + "mon.chart.throughput": "Tokens per second", + "mon.checkpoints": "Checkpoints", + "mon.eta": "Time left", + "mon.event": "Event", + "mon.events": "Events", + "mon.grad_norm": "Gradient norm", + "mon.input": "Input", + "mon.intro": "Live view of any run folder, including long runs started from the command line. Read-only.", + "mon.log": "Training log", + "mon.loss": "Loss", + "mon.lr": "Learning rate", + "mon.memory": "Memory", + "mon.no_runs": "No runs yet. Start one on the Train page.", + "mon.phase_note": "Steps restart in every stage, so charts show one stage at a time.", + "mon.progress": "Progress", + "mon.reference": "Reference", + "mon.run": "Run", + "mon.samples": "Samples", + "mon.size": "Size", + "mon.stale": "No update for a while", + "mon.throughput": "Tokens per second", + "mon.title": "Run monitor", + "mon.tokens_seen": "tokens seen", + "mon.updated": "Updated {s} s ago", + "mon.valid_ce": "Validation CE", + "mon.validation": "Validation", + "nav.autotune": "Auto-tune", + "nav.benchmark": "Benchmark", + "nav.config": "Model", + "nav.data": "Data", + "nav.export": "Export", + "nav.inference": "Inference", + "nav.io": "IO metrics", + "nav.monitor": "Monitor", + "nav.serve": "Serve", + "nav.tokenizer": "Tokenizer", + "nav.train": "Train", + "phase.active": "running", + "phase.complete": "complete", + "phase.pending": "waiting", + "picker.add": "Add selected", + "picker.add_button": "Add…", + "picker.browse_tip": "Browse folders on the server's machine inside the page and pick a path.", + "picker.empty": "No matching files or folders here", + "picker.entries": "Folder contents", + "picker.go": "Go", + "picker.hidden": "Show hidden", + "picker.multi_hint": "You can add several files or folders; they are merged before training.", + "picker.n_selected": "{n} selected", + "picker.native_fallback": "This machine cannot open a system file dialog, so the files will be uploaded instead.", + "picker.native_tip": "Open the system file manager. Nothing is copied; the chosen path is used as it is.", + "picker.pick_file": "Select a file", + "picker.place.desktop": "Desktop", + "picker.place.documents": "Documents", + "picker.place.downloads": "Downloads", + "picker.place.home": "Home", + "picker.place.root": "Root (/)", + "picker.place.volumes": "All volumes", + "picker.placeholder_multi": "One path per line, or use Add…", + "picker.places": "Places", + "picker.select": "Select", + "picker.this_folder": "Nothing selected: this folder will be used ({path})", + "picker.title.any": "Choose a file or folder", + "picker.title.checkpoint": "Choose a checkpoint", + "picker.title.config": "Choose a config file", + "picker.title.data": "Choose datasets (files or folders)", + "picker.title.dir": "Choose a folder", + "picker.truncated": "Showing the first {n} entries; type a path to go deeper", + "picker.up": "Up one folder", + "picker.upload_files": "Upload files…", + "picker.upload_folder": "Upload folder…", + "picker.upload_tip": "Pick files with this device's file manager. They are uploaded to the uploads folder on the server's machine, and that path is filled in.", + "picker.uploaded": "Uploaded {n} files.", + "picker.uploaded_skipped": "Uploaded {n} files; skipped {skipped} hidden or wrong-type files.", + "picker.uploading": "Uploading {done} / {total}…", + "picker.will_merge": "{n} paths will be merged into one JSONL file before use.", + "plan.apply": "Use this plan", + "plan.note": "The planner searches a configuration for the target size, data and memory; review it before training.", + "plan.run": "Plan", + "plan.stages": "Per-stage settings", + "plan.summary": "Plan", + "plan.title": "Auto plan", + "plan.tokens": "Training tokens", + "preset.b1": "1B (64 GB)", + "preset.b1.summary": "1B MoE with 64 fine-grained experts and aggressive memory saving.", + "preset.enzh50m": "50M en→zh recipe", + "preset.enzh50m.summary": "The 72-hour en→zh recipe: 16k Chronos tokenizer, shards_v1.", + "preset.m200": "200M (32 GB)", + "preset.m200.summary": "General 200M MoE with balanced memory saving and the 6400-token minimind vocabulary.", + "preset.m500": "500M (32 GB)", + "preset.m500.summary": "500M MoE with aggressive memory saving: int8 KV, shared routing, activation checkpointing.", + "preset.tiny": "Tiny (smoke test / CI)", + "preset.tiny.summary": "1.5M smoke test; trains in minutes on CPU.", + "probe.accept": "Draft acceptance", + "probe.agree": "Token agreement", + "probe.auto_paging": "Auto paging", + "probe.cold_hit": "Cold-miss predicted", + "probe.conditions": "Greedy decoding, {n} prompts, {t} tokens each; measured.", + "probe.exact": "Exact match", + "probe.history": "Hit rate and agreement over training", + "probe.hit": "Resident hit rate", + "probe.loads": "On-demand loads", + "probe.none": "No probe yet. One runs after each validation.", + "probe.pages": "Paging · pages", + "probe.pred_hit": "Prediction hit rate", + "probe.prefix": "Shortest matching prefix", + "probe.reference": "Full-memory reference", + "probe.rss": "Peak RSS (GB)", + "probe.skip.busy": "the previous probe is still running", + "probe.skip.dense": "dense model, no routed experts to page", + "probe.skip.no_checkpoint": "no checkpoint on disk yet", + "probe.skip.superseded": "replaced by the newer validation at step {step}", + "probe.skip.writing": "the checkpoint is still being written", + "probe.skipped": "Skipped: {reason}", + "probe.title": "Lazy-expert probe", + "probe.tps": "Tokens/s", + "probe.trained_clusters": "Trained with cluster locality", + "probe.variant": "Variant", + "probe.vs_full": "Speed vs full memory", + "probe.wait": "Prefetch wait (s)", + "rule.clusters_auto": "0 = derived from the expert count", + "rule.equals_recurrent_depth": "equal to the training depth", + "rule.loopback": "127.0.0.1 (this machine only)", + "rule.sync_is_exact": "sync_on_demand: output identical to full memory", + "rule.vocab_matches_tokenizer": "equal to the tokenizer vocabulary", + "rule.window_median_record": "power of two nearest the median record length (32 when unknown)", + "serve.auth_off": "No key: reachable from this machine only.", + "serve.auth_on": "Clients must send the API key.", + "serve.hint": "Connect with: OLLAMA_HOST={addr} ollama run ", + "serve.intro": "Start chronos serve: Ollama and OpenAI clients connect to it while experts stay lazy-loaded.", + "serve.remote_warning": "Binding beyond this machine needs Allow remote and a key in CHRONOS_API_KEY; without a key the server refuses to start.", + "serve.running": "Serving at {url}", + "serve.start": "Start server", + "serve.stop": "Stop server", + "serve.stopped": "Server stopped", + "serve.title": "Ollama / OpenAI-compatible server", + "source.code": "Default", + "source.measured": "Measured", + "source.planner": "Adaptive planner", + "source.rule": "Rule", + "stage.distill": "Distillation", + "stage.dpo": "DPO", + "stage.grpo": "GRPO", + "stage.orpo": "ORPO", + "stage.pretrain": "Pretraining", + "stage.sft": "Supervised fine-tuning", + "state.checkpointing": "Saving checkpoint", + "state.done": "Done", + "state.error": "Error", + "state.evaluating": "Evaluating", + "state.failed": "Failed", + "state.finished": "Finished", + "state.idle": "Idle", + "state.loading": "Loading", + "state.paused": "Paused", + "state.pending": "Pending", + "state.running": "Running", + "state.stale": "No recent update", + "state.starting": "Starting", + "state.stopped": "Stopped", + "state.stopping": "Stopping", + "tip.auto_empty": "Empty = automatic", + "tip.default": "Default", + "tip.range": "Range", + "tip.recommended": "Recommended", + "tip.source": "Source", + "tok.algorithm": "Algorithm", + "tok.bucket": "Bucket", + "tok.byte_fallback": "Byte fallback rate", + "tok.chars_per_token": "Characters per token", + "tok.compression": "Bytes per token", + "tok.elapsed": "Build time", + "tok.fertility": "Tokens per word", + "tok.fingerprint": "Fingerprint", + "tok.intro": "Train a Chronos tokenizer from your own data: sampling, normalisation, byte-fallback BPE and a held-out report.", + "tok.out_dir_now": "Output folder", + "tok.progress": "Build progress", + "tok.report": "Evaluation report", + "tok.report_note": "Measured on held-out text that was not used to train the tokenizer.", + "tok.roundtrip": "Lossless round trip", + "tok.texts": "Texts", + "tok.title": "Chronos tokenizer builder", + "tok.vocab_use": "Vocabulary used", + "train.config_note": "The architecture comes from the Model page; per-stage learning rates come from the plan when one was applied.", + "train.intro": "Start a training job with the model settings from the Model page. Progress appears in the monitor.", + "train.jobs": "Training jobs", + "train.open_monitor": "Open in monitor", + "train.probe_note": "Each validation decodes the newly saved checkpoint with lazy experts and with full memory. The probe runs as a low-priority child on a copy, so training continues.", + "train.probe_off": "The lazy-expert check is off for this job.", + "train.probe_title": "Lazy-expert check at every validation", + "train.run_dir": "Run folder", + "train.stage_data_note": "Each stage after the first reads its own data format. Leave a row empty to reuse the first stage's dataset.", + "train.stage_n": "Stage {n}", + "train.start": "Start training", + "train.started": "Training started", + "train.stop": "Stop training", + "train.title": "Train", + "train.variants": "{n} combinations will be tested", + "train.variants_capped": "{n} combinations requested, the first {cap} are tested", + "trial.complete": "complete", + "trial.failed": "failed", + "trial.pruned": "pruned", + "tune.apply": "Apply to model config", + "tune.best": "Best parameters", + "tune.duration": "Duration", + "tune.intro": "Search training settings with Optuna TPE: each trial trains a few steps and reports its loss.", + "tune.loss": "Probe loss", + "tune.note": "Probe loss after a short run; a ranking signal, not a final quality measure.", + "tune.start": "Start search", + "tune.stop": "Stop search", + "tune.title": "Auto-tune (Optuna TPE)", + "tune.trial": "Trial", + "tune.trials": "Trials", + "ui.access_tip.lan": "You are connected over the LAN. \"Upload…\" sends files to the uploads folder on the server's machine and uses that path. \"Browse…\" lists files on the server's machine.", + "ui.access_tip.lan_local": "You opened the UI on the server's own machine, so no sign-in is needed. \"Upload…\" opens the system file manager and uses the original path. Other devices on the LAN need the access key.", + "ui.access_tip.local_only": "The UI listens on this machine only (127.0.0.1); other devices cannot open it and no sign-in is needed. Start with --lan to allow LAN access.", + "ui.apply": "Apply", + "ui.backend": "Backend", + "ui.browse": "Browse…", + "ui.busy": "Another job of this kind is running", + "ui.cancel": "Cancel", + "ui.chart_view": "Chart", + "ui.clear": "Clear", + "ui.close": "Close", + "ui.copy": "Copy", + "ui.empty": "Nothing yet", + "ui.error": "Error", + "ui.from_model": "From the model file", + "ui.hide_advanced": "Hide advanced", + "ui.invalid": "Check the highlighted fields", + "ui.key": "Access key", + "ui.lan_local": "This machine · LAN open", + "ui.language": "Language", + "ui.load": "Load", + "ui.loading": "Loading…", + "ui.local_only": "Local only, no login", + "ui.log": "Log", + "ui.login": "Sign in", + "ui.login_blocked": "Too many failed attempts. Wait 5 minutes.", + "ui.login_failed": "Wrong key", + "ui.login_needed": "This workbench needs the access key set on the server.", + "ui.logout": "Sign out", + "ui.no": "No", + "ui.none": "None", + "ui.path": "Path", + "ui.phase": "Stage", + "ui.raw_json": "Raw JSON", + "ui.refresh": "Refresh", + "ui.remote_auth": "Remote access, signed in", + "ui.reset": "Reset to defaults", + "ui.result": "Result", + "ui.run": "Run", + "ui.save": "Save", + "ui.saved": "Saved", + "ui.scan": "Scan", + "ui.sections": "Sections", + "ui.show_advanced": "Show advanced", + "ui.skip": "Skip to content", + "ui.start": "Start", + "ui.started": "Started", + "ui.state": "State", + "ui.step": "Step", + "ui.stop": "Stop", + "ui.table_view": "Table", + "ui.theme": "Theme", + "ui.theme.dark": "Dark", + "ui.theme.light": "Light", + "ui.theme.system": "System", + "ui.time": "Time", + "ui.upload": "Upload…", + "ui.value": "Value", + "ui.yes": "Yes" +} diff --git a/chronos/webui/i18n/ja.json b/chronos/webui/i18n/ja.json new file mode 100644 index 0000000..f28330a --- /dev/null +++ b/chronos/webui/i18n/ja.json @@ -0,0 +1,860 @@ +{ + "app.subtitle": "再帰深さ MoE · 遅延エキスパート · 投機的デコード", + "app.title": "Chronos ワークベンチ", + "bench.better": "優位", + "bench.higher": "高いほど良い", + "bench.intro": "benchmark_compare.py で MiniMind と Chronos のパラメータ数・メモリ・デコード速度を比較します。", + "bench.lower": "低いほど良い", + "bench.measured_note": "benchmark_compare.py がこのマシンで実測。", + "bench.metric": "指標", + "bench.none": "結果はまだありません。ベンチマークを実行すると作成されます。", + "bench.params": "パラメータ数", + "bench.ram_infer": "推論時の RAM", + "bench.ram_model": "モデルの RAM", + "bench.ram_train": "学習時の RAM", + "bench.run": "ベンチマークを実行", + "bench.title": "ベンチマーク：MiniMind と Chronos", + "bench.tps": "デコード トークン/秒", + "bench.unavailable": "このインストールには benchmark_compare.py がありません。", + "choice.absorbed": "吸収", + "choice.adaptive": "適応", + "choice.aggressive": "積極的", + "choice.auto": "自動", + "choice.balanced": "バランス", + "choice.cluster": "クラスタ単位", + "choice.compare": "遅延ロードと全メモリを比較", + "choice.expanded": "展開（参照）", + "choice.expert": "エキスパート単位", + "choice.experts": "アクティブエキスパート数", + "choice.extreme": "最大限", + "choice.fallback_diagnostic": "読み込まない・診断用（fallback_diagnostic）", + "choice.fixed": "固定", + "choice.fixed_then_adaptive": "固定の後に適応", + "choice.fullload": "全メモリのみ", + "choice.none": "なし", + "choice.offload": "遅延エキスパートのみ", + "choice.on_demand": "読み込み、遅延したエキスパートは省略（on_demand）", + "choice.params": "アクティブパラメータ数", + "choice.planned": "次回使用プランナー", + "choice.recurrent_attention_mode.first_only": "初回のみ注意、v2 動作（first_only）", + "choice.recurrent_attention_mode.per_iteration": "毎反復で注意、反復ごとに KV（per_iteration）", + "choice.recurrent_attention_mode.shared_first": "毎反復で注意、初回 KV を共有（shared_first）", + "choice.router_scope.per_iteration": "反復ごとにルーティング（per_iteration）", + "choice.router_scope.shared": "初回のルーティングを再利用（shared）", + "choice.scheduled": "スケジュール", + "choice.shared": "共有", + "choice.sync_on_demand": "読み込んで待機（sync_on_demand）", + "choice.tune.search.lr": "学習率", + "choice.window": "最近使用ウィンドウ", + "config.active": "トークンあたりのアクティブパラメータ", + "config.compute": "トークンあたりの重み使用量（再帰込み）", + "config.estimate": "規模とメモリの見積もり", + "config.estimate_note": "設定から求めた推定値で、実測ではありません。", + "config.file": "設定ファイル", + "config.ignored": "不明または無効なキーを無視しました", + "config.intro": "モデルの設定はすべてここにあります。項目にカーソルを合わせると、意味・既定値・推奨値が表示されます。", + "config.kv": "KV キャッシュ", + "config.loaded": "設定を読み込みました", + "config.path_help": "JSON の設定ファイル。「読み込み」でフォームに読み込み、「保存」でフォームをここに書き出します。相対パスはワークベンチのフォルダが起点です。", + "config.preset": "プリセット", + "config.preset_help": "出発点となる設定。適用するとこのページの全項目が置き換わり、推定値もすぐ更新されます。", + "config.ram": "エキスパートのホットキャッシュ", + "config.ssd": "ディスク上のエキスパート", + "config.total": "総パラメータ数", + "config.tps": "デコード速度の概算（推定）", + "config.train_peak": "学習時のピークメモリ", + "config.train_state": "学習状態（重み・勾配・オプティマイザ）", + "config.virtual_layers": "仮想層数", + "config.vram": "常駐重み + KV", + "data.chars": "レコードあたりの文字数", + "data.dups": "重複率", + "data.fields": "フィールド", + "data.files": "ファイル数", + "data.full_scan": "全件走査", + "data.hist": "レコード長の分布", + "data.intro": "学習前にファイルやフォルダを走査し、形式・フィールド・文字体系・レコード長を確認します。", + "data.max_help": "集計時に読み込む最大レコード数（100〜500,000）。", + "data.max_records": "走査件数", + "data.median": "中央値", + "data.p95": "95 パーセンタイル", + "data.path_help": "ファイルまたはフォルダ（jsonl・json・parquet・csv・tsv・txt）を 1 行に 1 つ。複数指定は、結合後に学習が読むのと同じ 1 つのデータセットとしてスキャンします。相対パスはワークベンチのフォルダが起点です。", + "data.present": "出現率", + "data.records": "レコード数", + "data.sampled": "先頭 {n} 件", + "data.samples": "サンプルレコード", + "data.scanned": "スキャン済み", + "data.schema": "検出された形式", + "data.script": "文字", + "data.scripts": "文字体系", + "data.share": "割合", + "data.title": "データセットの確認", + "data.type": "型", + "err.active_experts": "アクティブエキスパート数はエキスパート総数を超えられません", + "err.dialog_open": "ファイルダイアログがすでに開いています。先にシステムのウィンドウで完了またはキャンセルしてください", + "err.models_dir": "モデルフォルダが見つかりません", + "err.native_local_only": "システムのファイルダイアログはサーバーと同じマシンでのみ使えます", + "err.native_unavailable": "このマシンではシステムのファイルダイアログを開けません", + "err.no_files": "指定したパスに読み込めるファイルがありません", + "err.no_out": "出力フォルダを選んでください", + "err.no_paths": "パスを 1 つ以上追加してください", + "err.not_found": "見つかりません", + "err.paging": "ページングは auto・cluster・expert から選んでください", + "err.port": "ポートが無効です", + "err.probe_invalid": "遅延読み込みチェックの設定が無効です", + "err.remote_needs_key": "リモートバインドには「リモートアクセスを許可」と CHRONOS_API_KEY が必要です", + "err.target": "目標サイズを入力してください", + "err.upload_batch": "アップロードのバッチ ID が無効です。ページを再読み込みしてやり直してください", + "err.upload_disk_full": "サーバー側のディスク空き容量が不足しているため、アップロードを中止しました", + "err.upload_failed": "アップロードが中断されました。接続を確認してやり直してください", + "err.upload_kind": "この項目はアップロードに対応していません", + "err.upload_name": "ファイル名またはフォルダー階層が許可されていません（隠し項目、「..」、階層が深すぎる）", + "err.upload_too_large": "ファイルがアップロード上限（--max-upload-mb）を超えています", + "err.upload_type": "アップロードできるファイルがありません。この項目に合わない種類か、許可されていない実行ファイルです", + "event.final": "終了", + "event.lazy_probe": "遅延読み込みプローブ", + "event.phase_end": "段階終了", + "event.phase_start": "段階開始", + "event.resume": "再開", + "event.rollback": "ロールバック", + "event.validation": "検証", + "export.done": "エクスポートしました", + "export.format": "形式", + "export.intro": "チェックポイントを safetensors または GGUF ファイルとして書き出し、chronos serve が読み込むエキスパートキャッシュ配置を添えます。", + "export.note": "Ollama クライアントは chronos serve を通じて Chronos モデルを実行します（OLLAMA_HOST をそのアドレスに設定）。標準の Ollama は再帰 MoE を実行できません。", + "export.run": "エクスポート", + "export.size": "サイズ（MB）", + "export.tensors": "テンソル数", + "export.title": "エクスポート", + "field.accumulation_steps.help": "オプティマイザ更新 1 回あたりのマイクロバッチ数。実効バッチ = バッチサイズ × 累積数。", + "field.accumulation_steps.label": "勾配累積", + "field.activation_checkpointing.help": "逆伝播で活性化を再計算し、速度と引き換えにピークメモリを減らします。", + "field.activation_checkpointing.label": "活性化チェックポイント", + "field.adaptive_halting.help": "停止確率がしきい値を超えたトークンは反復をやめます。簡単なトークンの計算を節約します。", + "field.adaptive_halting.label": "適応的停止（ACT）", + "field.alpha.help": "蒸留の混合比：α × 教師 KL + (1 − α) × CE。", + "field.alpha.label": "α（蒸留）", + "field.batch_size.help": "マイクロバッチあたりの系列数。メモリで制限されます。", + "field.batch_size.label": "バッチサイズ", + "field.beta.help": "選好学習で参照方策への KL の結びつきの強さ。", + "field.beta.label": "β（DPO/GRPO）", + "field.cluster_locality_clusters.help": "局所性項で使うエキスパートクラスタの数。0 ならエキスパート数から決めます。", + "field.cluster_locality_clusters.label": "クラスタ数", + "field.cluster_locality_window.help": "クラスタ制約をかける区間の長さ（トークン）。典型的なレコード長に合わせたときに最良の結果が出ました。", + "field.cluster_locality_window.label": "局所性ウィンドウ", + "field.coda_layers.help": "ループの後、出力ヘッドの前に 1 回だけ実行する層。", + "field.coda_layers.label": "コーダ層", + "field.cpu_budget_percent.help": "auto が使える物理コアの割合（%）。", + "field.cpu_budget_percent.label": "CPU 予算", + "field.cpu_threads.help": "PyTorch／BLAS のスレッド数。auto は物理コアの予算を使います。", + "field.cpu_threads.label": "CPU スレッド", + "field.depth_schedule.help": "学習時にバッチごとの深さを決める方法：固定、適応、またはウォームアップ中は固定でその後に適応。", + "field.depth_schedule.label": "深さスケジュール", + "field.depth_schedule_warmup_fraction.help": "スケジュールが適応に切り替わる前に、固定の最大深さで学習する割合。", + "field.depth_schedule_warmup_fraction.label": "スケジュールのウォームアップ", + "field.draft_exit_depth.help": "投機的ドラフトとして出口を学習させる再帰深さ。", + "field.draft_exit_depth.label": "ドラフト深さ（学習）", + "field.draft_exit_kl_coef.help": "自己投機的デコード用に、ドラフト深さの出口と最大深さの出力の間に追加する KL。0 で無効。", + "field.draft_exit_kl_coef.label": "ドラフト出口 KL", + "field.dropout.help": "学習中に活性化をランダムに無効化します。大規模データでは下げ、ごく小さなデータでは上げてください。", + "field.dropout.label": "ドロップアウト", + "field.dtype.help": "学習と推論の数値精度。auto はバックエンドが対応していれば bfloat16、そうでなければ float32 を選びます。", + "field.dtype.label": "計算精度", + "field.epochs.help": "データセットを通す回数。最大ステップを設定すると、そちらで早めに止まります。", + "field.epochs.label": "エポック数", + "field.export.auto_cluster.help": "キャリブレーションデータから同時に使われるエキスパートをページにまとめ、遅延ロードで読むファイルを減らします。", + "field.export.auto_cluster.label": "クラスタを構築", + "field.export.calibration_data_path.help": "どのエキスパートが一緒に使われるかを測定するテキスト。1 行に 1 つのパス。複数指定は自動で結合します。空欄なら自動クラスタリングを行いません。", + "field.export.calibration_data_path.label": "キャリブレーションデータ", + "field.export.cluster_batch_size.help": "キャリブレーションバッチあたりの系列数。", + "field.export.cluster_batch_size.label": "キャリブレーションのバッチサイズ", + "field.export.cluster_device.help": "キャリブレーションパスに使うデバイス。", + "field.export.cluster_device.label": "キャリブレーションのデバイス", + "field.export.cluster_max_batches.help": "クラスタリングのために観測するバッチの数。", + "field.export.cluster_max_batches.label": "キャリブレーションのバッチ数", + "field.export.cluster_max_seq_len.help": "キャリブレーション系列あたりのトークン数。", + "field.export.cluster_max_seq_len.label": "キャリブレーションの長さ", + "field.export.config_path.help": "サイドカーの設定を上書きします。通常は空欄です。", + "field.export.config_path.label": "設定ファイル", + "field.export.copy_cache.help": "遅延エキスパートのページをエクスポートに含め、サービス起動時に再構築しなくて済むようにします。", + "field.export.copy_cache.label": "エキスパートキャッシュをコピー", + "field.export.expert_cache_dir.help": "同梱する既存の遅延エキスパートキャッシュ。空欄なら新しく作成します。", + "field.export.expert_cache_dir.label": "エキスパートキャッシュ", + "field.export.formats.help": "書き出すコンテナ。Safetensors と GGUF は Chronos のメタデータを保持し、chronos serve で実行します。Ollama クライアントは標準の Ollama ではなく chronos serve に接続します。", + "field.export.formats.label": "形式", + "field.export.model_path.help": "エクスポートするチェックポイント。", + "field.export.model_path.label": "チェックポイント", + "field.export.output_dir.help": "エクスポートファイルの書き込み先。", + "field.export.output_dir.label": "出力フォルダ", + "field.fallback_mask_prob.help": "学習中にランダムにエキスパートを隠し、キャッシュミスに耐えられるようにします。フォールバック付きの遅延推論でのみ意味があります。", + "field.fallback_mask_prob.label": "ミスのシミュレーション", + "field.grad_clip.help": "全体の勾配ノルムをこの値でクリップします。0 で無効。", + "field.grad_clip.label": "勾配クリップ", + "field.halting_threshold.help": "トークンが反復をやめる累積停止確率。1 に近いほど深く実行します。", + "field.halting_threshold.label": "停止しきい値", + "field.hidden_size.help": "各トークンベクトルの幅。大きいほど表現力は上がりますが、パラメータ数はおおよそ二乗で増えます。アテンションヘッド数で割り切れる必要があります。", + "field.hidden_size.label": "隠れ層サイズ", + "field.infer.backend.help": "推論のランタイム。auto はチェックポイントとマシンから選びます。", + "field.infer.backend.label": "バックエンド", + "field.infer.decode_wait_s.help": "プリフェッチ中のエキスパートをオンデマンドでロードする前に待つ時間。", + "field.infer.decode_wait_s.label": "デコード待機（秒）", + "field.infer.draft_depth.help": "下書きに使う再帰深さ。", + "field.infer.draft_depth.label": "ドラフト深さ", + "field.infer.gamma.help": "検証 1 回あたりに下書きするトークン数。", + "field.infer.gamma.label": "ドラフト長 γ", + "field.infer.max_tokens.help": "生成するトークン数の上限。", + "field.infer.max_tokens.label": "最大生成トークン数", + "field.infer.miss_policy.help": "キャッシュミス時の動作。sync_on_demand なら出力は全メモリと同一です。", + "field.infer.miss_policy.label": "ミス時のポリシー", + "field.infer.mode.help": "compare は遅延ロードと全メモリを並べて実行。offload は遅延エキスパートのみ。fullload はすべてをロードします。", + "field.infer.mode.label": "モード", + "field.infer.model_path.help": "実行する .pth チェックポイント。アーキテクチャはサイドカーの設定で決まります。", + "field.infer.model_path.label": "チェックポイント", + "field.infer.paging.help": "auto はチェックポイントに従います。クラスタ局所性で学習していればクラスタ単位、そうでなければエキスパート単位。", + "field.infer.paging.label": "ページング方式", + "field.infer.ram_load_ratio.help": "RAM に保持するエキスパートの割合。1.0 を超えるとプリフェッチ用の余裕が加わります。", + "field.infer.ram_load_ratio.label": "RAM 比率", + "field.infer.raw_prompt.help": "チャットテンプレートを使わず、テキストをそのまま送ります。", + "field.infer.raw_prompt.label": "生のプロンプト", + "field.infer.residency_decay.help": "プランナーが古いルーティング情報を忘れる速さ。", + "field.infer.residency_decay.label": "プランナーの減衰", + "field.infer.residency_planner.help": "最も長く使われていないエキスパートではなく、次に必要になるのが最も遠いエキスパートを追い出します。", + "field.infer.residency_planner.label": "次回使用プランナー", + "field.infer.speculative.help": "浅い深さでトークンを下書きし、最大深さで検証します。出力は変わりません。50M では遅くなることが実測されたため、既定ではオフです。", + "field.infer.speculative.label": "自己投機的デコード", + "field.infer.sweep_ram_ratios.help": "標準の比率をすべて実行し、メモリに対する速度とヒット率をグラフにします。", + "field.infer.sweep_ram_ratios.label": "RAM 比率をスイープ", + "field.infer.temperature.help": "サンプリングのランダムさ。0 は貪欲デコードで再現可能です。", + "field.infer.temperature.label": "温度", + "field.infer.verify_wait_s.help": "投機的デコードの検証パスでの同様の待機時間。", + "field.infer.verify_wait_s.label": "検証待機（秒）", + "field.kv_cache_dtype.help": "auto は計算精度のまま。int8 は推論時の KV キャッシュをトークン単位で量子化します（メモリ約半分、誤差は小さい）。学習は量子化されません。", + "field.kv_cache_dtype.label": "KV キャッシュ精度", + "field.kv_latent_dim.help": "MLA 層がキャッシュする圧縮キー／バリュー潜在の幅。小さいほど KV メモリを節約します。", + "field.kv_latent_dim.label": "MLA 潜在サイズ", + "field.kv_quant_group_size.help": "1 つの int8 スケールを共有するチャネル数。小さいほど精度が上がり、サイズはわずかに増えます。", + "field.kv_quant_group_size.label": "KV 量子化グループ", + "field.lambda_balance.help": "トークンをエキスパートに均等に振り分けるようルーターを促します。", + "field.lambda_balance.label": "負荷分散", + "field.lambda_cluster_locality.help": "区間内のトークンを 1 つのエキスパートクラスタ内にとどめるよう学習させ、遅延推論でクラスタ単位のページングを可能にします。0 で無効。", + "field.lambda_cluster_locality.label": "クラスタ局所性", + "field.lambda_lookahead.help": "将来のトークンが使うエキスパートを予測するよう先読みルーターを学習させます。", + "field.lambda_lookahead.label": "先読み損失", + "field.lambda_lookahead_topk.help": "将来の top-k エキスパートを正確に当てることへの追加の重み。", + "field.lambda_lookahead_topk.label": "先読み top-k", + "field.lambda_lookahead_union.help": "先読みウィンドウ内のエキスパートの和集合をカバーすることへの重み。", + "field.lambda_lookahead_union.label": "先読み和集合", + "field.lambda_or.help": "ORPO のオッズ比項の重み。", + "field.lambda_or.label": "λ（ORPO）", + "field.lambda_router_anchor.help": "ファインチューニング中にルーティングを上流チェックポイントに近く保ちます。事前学習では 0。", + "field.lambda_router_anchor.label": "ルーターアンカー", + "field.lambda_router_locality.help": "隣接トークン間でルーティングが変わることへの弱いペナルティ。", + "field.lambda_router_locality.label": "ルーター局所性", + "field.lambda_temporal.help": "連続するトークンが同じエキスパートを選ぶと報酬を与え、テキストの区間ごとにロードするエキスパートを減らします。", + "field.lambda_temporal.label": "時間的局所性", + "field.layer_aware_cache.help": "常駐エキスパートを物理層ではなく仮想層ごとに管理します。", + "field.layer_aware_cache.label": "層別キャッシュ", + "field.learning_rate.help": "オプティマイザの最大ステップ幅。プランナーがモデルサイズから提案します。ファインチューニングでは下げてください。", + "field.learning_rate.label": "学習率", + "field.log_interval.help": "N ステップごとにメトリクスを 1 行書き出します。", + "field.log_interval.label": "ログ間隔", + "field.lookahead_steps.help": "先読みルーターが何トークン先までエキスパートを予測するか。必要になる前にプリフェッチできます。", + "field.lookahead_steps.label": "先読みステップ", + "field.max_gen_len.help": "GRPO のロールアウトごとに生成するトークン数。", + "field.max_gen_len.label": "ロールアウト長", + "field.max_inference_depth.help": "推論時の最大反復回数。通常は学習時の深さと同じです。学習時より深くした場合の挙動は検証されていません。", + "field.max_inference_depth.label": "推論時の深さ", + "field.max_position_embeddings.help": "回転位置埋め込みが対応する最長の系列長。非常に長いコンテキストを使う場合のみ変更してください。", + "field.max_position_embeddings.label": "最大位置数", + "field.max_resident_experts.help": "層ごとにメモリに保持するエキスパート数の上限。空欄ならプランナーが決めます。", + "field.max_resident_experts.label": "最大常駐エキスパート数", + "field.max_seq_len.help": "学習系列あたりのトークン数。ハイブリッドアテンションではメモリはおおよそ線形に増えます。", + "field.max_seq_len.label": "系列長", + "field.max_steps.help": "オプティマイザがこのステップ数に達したら停止します。0 は上限なし（全エポック実行）。", + "field.max_steps.label": "最大ステップ", + "field.min_recurrent_depth.help": "スケジュールや早期終了が使える最小の深さ。1 なら簡単なトークンは 1 回で終了できます。", + "field.min_recurrent_depth.label": "最小深さ", + "field.mla_cache_rope.help": "回転キー部分もキャッシュし、KV メモリと引き換えにデコードを速くします。", + "field.mla_cache_rope.label": "RoPE キーをキャッシュ", + "field.mla_decode.help": "expanded：参照経路。absorbed：射影をクエリに畳み込みます（1 ステップあたりの計算が少ない）。auto は経路を混在させ、わずかなロジット差が出たため、明示的に選んだ場合のみ使います。", + "field.mla_decode.label": "MLA デコード経路", + "field.modalities.help": "追加の入力エンコーダ（image、audio をカンマ区切り）。テキストのみのモデルでは空欄にします。データセットは自動判定されます。", + "field.modalities.label": "モダリティ", + "field.moe_intermediate_size.help": "各エキスパートのフィードフォワードネットワークの隠れ幅。0 なら隠れ層サイズから求めます。", + "field.moe_intermediate_size.label": "エキスパート FFN 幅", + "field.num_attention_heads.help": "クエリヘッドの数。隠れ層サイズはこの数で割り切れる必要があります。", + "field.num_attention_heads.label": "アテンションヘッド数", + "field.num_experts.help": "MoE 層ごとのルーティング対象エキスパート数。トークンあたりの計算を増やさずに容量を増やせます。遅延ロードの単位でもあります。", + "field.num_experts.label": "エキスパート数", + "field.num_experts_per_tok.help": "各トークンが使うエキスパートの数。top-2 ならルーターが損失から学習できます。top-1 では norm_topk_prob をオフにする必要があります。", + "field.num_experts_per_tok.label": "アクティブエキスパート数（top-k）", + "field.num_generations.help": "GRPO がプロンプトごとに比較するサンプル数。", + "field.num_generations.label": "プロンプトあたりのロールアウト数", + "field.num_hidden_layers.help": "チェックポイントに保存される独立した Transformer 層の数（プレリュード＋再帰グループ＋コーダ）。再帰で再利用されるため、仮想的な深さはこれより大きくなります。", + "field.num_hidden_layers.label": "物理層数", + "field.num_key_value_heads.help": "スライディングウィンドウ層のグループ化クエリアテンションで使うキー／バリューヘッドの数。少ないほど KV キャッシュが小さくなります。", + "field.num_key_value_heads.label": "KV ヘッド数", + "field.num_shared_experts.help": "すべてのトークンが常に通るエキスパート。常にメモリに常駐します。", + "field.num_shared_experts.label": "共有エキスパート", + "field.num_workers.help": "データローダーのワーカープロセス数。auto はプラットフォームに合った安全な数を選びます。", + "field.num_workers.label": "データワーカー", + "field.optimizer_page_dir.help": "ページングしたオプティマイザ状態を置くフォルダ。高速なローカル SSD を使ってください。", + "field.optimizer_page_dir.label": "オプティマイザのページフォルダ", + "field.optimizer_paging.help": "非アクティブなエキスパートのオプティマイザ状態をディスクに保存します。", + "field.optimizer_paging.label": "オプティマイザのページング", + "field.page_size_mb.help": "オプティマイザのページファイル 1 つの大きさ（MB）。", + "field.page_size_mb.label": "ページサイズ", + "field.pinned_memory_max_fraction.help": "ピン留めエキスパートバッファに使うシステム RAM の最大割合。", + "field.pinned_memory_max_fraction.label": "ピン留め RAM の割合", + "field.plan.active_mode.help": "トークンあたりの予算をエキスパート数で指定するか、アクティブパラメータ数で指定するか。", + "field.plan.active_mode.label": "アクティブ予算の指定方法", + "field.plan.active_value.help": "トークンあたりのアクティブエキスパート数、またはアクティブパラメータ数（百万単位）。", + "field.plan.active_value.label": "アクティブ予算", + "field.plan.data_paths.help": "トークン数、系列長、ウィンドウを決めるためにプランナーが測定するファイルまたはフォルダ（1 行に 1 つ）。", + "field.plan.data_paths.label": "データセット", + "field.plan.depth.help": "再帰グループあたりのループ反復回数。", + "field.plan.depth.label": "再帰的な深さ", + "field.plan.memory_saving.help": "速度と引き換えにどこまでメモリを節約するか（int8 KV、共有ルーティング、チェックポイント）。", + "field.plan.memory_saving.label": "メモリ節約", + "field.plan.num_experts.help": "MoE 層ごとのルーティング対象エキスパート数。", + "field.plan.num_experts.label": "エキスパート数", + "field.plan.ram_gb.help": "計画の前提とするメモリ。空欄ならこのマシンを検出します。", + "field.plan.ram_gb.label": "RAM（GB）", + "field.plan.seq_len.help": "空欄ならデータセットのプロファイルから求めます。", + "field.plan.seq_len.label": "系列長", + "field.plan.target_m.help": "総パラメータ数（百万単位）。", + "field.plan.target_m.label": "目標サイズ（M）", + "field.plan.token_budget_b.help": "事前学習のトークン数（十億単位）。空欄ならデータを 1 周します。", + "field.plan.token_budget_b.label": "トークン予算（B）", + "field.plan.use_moe.help": "MoE モデルを計画します。オフなら密なモデルを計画します。", + "field.plan.use_moe.label": "混合エキスパート", + "field.plan.vocab_size.help": "空欄ならトークナイザーのレポートから読むか、16384 を使います。", + "field.plan.vocab_size.label": "語彙", + "field.ponder_loss_coef.help": "期待反復回数へのペナルティ。大きいほど早く停止するよう促します。", + "field.ponder_loss_coef.label": "思考コスト損失", + "field.prefetch_depth.help": "予測したエキスパートをどこまで先読みしてロードするか。0 でプリフェッチ無効。", + "field.prefetch_depth.label": "プリフェッチ深さ", + "field.prelude_layers.help": "ループの前に 1 回だけ実行する層。先読みルーティングはこの出力を読みます。", + "field.prelude_layers.label": "プレリュード層", + "field.probe.backends.help": "テストするランタイム。バックエンドごとに全メモリの参照を取ります。", + "field.probe.backends.label": "プローブのバックエンド", + "field.probe.decode_wait_s.help": "デコードの各ステップで、プリフェッチ中のエキスパートをオンデマンドでロードする前に待つ時間。", + "field.probe.decode_wait_s.label": "デコード待機（秒）", + "field.probe.draft_depth.help": "投機的デコードの組み合わせでトークンを下書きする再帰深さ。", + "field.probe.draft_depth.label": "ドラフト深さ", + "field.probe.enabled.help": "検証のたびに、保存したばかりのチェックポイントを遅延エキスパートと全メモリでデコードして比較します（速度、ヒット率、トークン一致率）。チェックポイントのコピー上で、低優先度の子プロセスとして実行します。", + "field.probe.enabled.label": "遅延エキスパートのプローブ", + "field.probe.gamma.help": "投機的デコードの組み合わせで、検証 1 回あたりに下書きするトークン数。", + "field.probe.gamma.label": "ドラフト長 γ", + "field.probe.max_new_tokens.help": "プロンプトごとに貪欲デコードするトークン数。多いほど測定値は安定しますが、プローブに時間がかかります。", + "field.probe.max_new_tokens.label": "プロンプトあたりのトークン数", + "field.probe.max_variants.help": "テストする組み合わせ数の上限。超えた分は打ち切りとして報告されます。", + "field.probe.max_variants.label": "組み合わせの上限", + "field.probe.miss_policies.help": "キャッシュミス時の Torch の動作。sync_on_demand はロードして待つ（出力は同一）。on_demand は間に合わないエキスパートを飛ばすことがある。fallback_diagnostic はロードしない（最悪ケース）。", + "field.probe.miss_policies.label": "ミス時のポリシー", + "field.probe.paging.help": "auto はモデルの学習方法に従います（クラスタ局所性損失を使っていればクラスタ単位）。expert を加えるとエキスパート単位のページングと比較できます。", + "field.probe.paging.label": "ページング方式", + "field.probe.ram_ratios.help": "同時に RAM に置けるエキスパートの割合。低いほどメモリを節約し、ミスが増えます。", + "field.probe.ram_ratios.label": "RAM 比率", + "field.probe.residency.help": "MLX の追い出し方式：planned は次回使用プランナー、window は単純な最近使用ウィンドウ。", + "field.probe.residency.label": "常駐ポリシー", + "field.probe.residency_decay.help": "次回使用プランナーが古いルーティング情報を忘れる速さ。", + "field.probe.residency_decay.label": "プランナーの減衰", + "field.probe.speculative.help": "遅延エキスパートでの自己投機的デコードもテストします。", + "field.probe.speculative.label": "投機的デコードも測定", + "field.probe.verify_wait_s.help": "投機的デコードの検証パスでの同様の待機時間。", + "field.probe.verify_wait_s.label": "検証待機（秒）", + "field.ram_budget_mb.help": "遅延エキスパートキャッシュが使える RAM（MB）。空欄なら RAM 比率を使います。", + "field.ram_budget_mb.label": "RAM 予算", + "field.recurrent_attention_mode.help": "ループ内でのアテンションの実行方法。shared_first：毎反復でアテンションし、K/V は反復 0 のもの（キャッシュは厳密で、KV は深さで増えない）。per_iteration：反復ごとに KV を持つ（KV が D 倍）。first_only：v2 の挙動。", + "field.recurrent_attention_mode.label": "再帰内アテンション", + "field.recurrent_depth.help": "学習中に各再帰グループを反復する回数。仮想層数 = プレリュード + グループ数 × 深さ + コーダ。", + "field.recurrent_depth.label": "学習時の深さ", + "field.recurrent_depth_jitter.help": "スケジュールされた深さの周りにランダムな幅を持たせ、複数の深さで終了できるよう学習させます。", + "field.recurrent_depth_jitter.label": "深さのゆらぎ", + "field.recurrent_distill_coef.help": "浅い出口の予測を最大深さの予測に合わせるよう学習させます。投機的デコードの受理率と早期終了の品質が上がります。", + "field.recurrent_distill_coef.label": "深さ蒸留", + "field.recurrent_enabled.help": "中間の共有グループをトークンごとに複数回実行し、少ない重みでより深いネットワークのように働かせます。", + "field.recurrent_enabled.label": "再帰的な深さ", + "field.recurrent_groups.help": "プレリュードとコーダの間で反復するグループを構成する層の数。空欄なら 層数 − プレリュード − コーダ から求めます。", + "field.recurrent_groups.label": "再帰グループ数", + "field.recurrent_stability_coef.help": "再帰更新を縮小写像に保ち、反復を重ねても発散やドリフトが起きないようにします。", + "field.recurrent_stability_coef.label": "安定化損失", + "field.recurrent_strategy.help": "学習時に再帰深さのバケットをサンプリングする戦略。", + "field.recurrent_strategy.label": "深さの戦略", + "field.reward_spec.help": "組み込み報酬なら toy、または lm:/path/to/reward-model[,device=cuda]。", + "field.reward_spec.label": "報酬（GRPO）", + "field.rope_dim.help": "MLA 層で回転位置情報を持つヘッドあたりの次元数。", + "field.rope_dim.label": "RoPE 次元", + "field.rope_theta.help": "回転位置埋め込みの基底周波数。大きいほど位置分解能が引き伸ばされ、長いコンテキストに向きます。", + "field.rope_theta.label": "RoPE 基数", + "field.router_scope.help": "per_iteration：ループの反復ごとにルーティングし直す。shared：反復 0 のルーティングを再利用し、触れるエキスパートが少なく遅延ロードに向きます。", + "field.router_scope.label": "ルーター範囲", + "field.save_dir.help": "チェックポイントとその .config.json サイドカーの書き込み先。", + "field.save_dir.label": "チェックポイントフォルダ", + "field.save_interval.help": "N ステップごとに保存（と検証）を行います。", + "field.save_interval.label": "保存間隔", + "field.serve.allow_remote.help": "ループバック以外のアドレスで待ち受けます。サーバーの環境変数に CHRONOS_API_KEY が必要で、クライアントはそのキーを送る必要があります。", + "field.serve.allow_remote.label": "リモートアクセスを許可", + "field.serve.device.help": "サービスのランタイム。", + "field.serve.device.label": "デバイス", + "field.serve.dtype.help": "サービス時の数値精度。", + "field.serve.dtype.label": "精度", + "field.serve.host.help": "待ち受けるアドレス。127.0.0.1 はこのマシンからのみ接続できます。", + "field.serve.host.label": "ホスト", + "field.serve.kv_cache_dtype.help": "KV の保存形式を上書きします。int8 なら KV メモリはおよそ半分。空欄ならモデルの設定のまま。", + "field.serve.kv_cache_dtype.label": "KV キャッシュ精度", + "field.serve.lazy_experts.help": "エキスパートを SSD に置き、必要に応じてページインします。Ollama プロトコル経由でもメモリを低く保てます。", + "field.serve.lazy_experts.label": "遅延エキスパート", + "field.serve.max_resident_experts.help": "層ごとに RAM に置くエキスパート数の上限。空欄なら RAM 割合から求めます。", + "field.serve.max_resident_experts.label": "最大常駐エキスパート数", + "field.serve.mla_cache_rope.help": "回転済みキーをキャッシュします。メモリは増え、再計算は減ります。", + "field.serve.mla_cache_rope.label": "RoPE キーをキャッシュ", + "field.serve.mla_decode.help": "MLA デコード経路を上書きします。空欄ならモデルの設定のまま。", + "field.serve.mla_decode.label": "MLA デコード", + "field.serve.models_dir.help": "エクスポート済みモデルのフォルダ。各モデルが Ollama のモデル名として表示されます。", + "field.serve.models_dir.label": "モデルフォルダ", + "field.serve.paging.help": "auto はチェックポイントの学習方法に従います（クラスタ単位またはエキスパート単位）。", + "field.serve.paging.label": "ページング方式", + "field.serve.port.help": "TCP ポート。11435 なら 11434 の標準 Ollama と衝突しません。", + "field.serve.port.label": "ポート", + "field.serve.ram_budget_mb.help": "エキスパート用 RAM の上限。空欄なら上限なし。", + "field.serve.ram_budget_mb.label": "RAM 予算（MB）", + "field.serve.ram_expert_fraction.help": "RAM に置けるエキスパートの割合。", + "field.serve.ram_expert_fraction.label": "RAM 上のエキスパート割合", + "field.serve.router_scope.help": "反復間のルーティング再利用を上書きします。空欄ならモデルの設定のまま。", + "field.serve.router_scope.label": "ルーター範囲", + "field.sliding_window_size.help": "スライディングウィンドウ層が参照できるトークン数。短いウィンドウは文単位のタスクに向き、KV メモリを抑えます。", + "field.sliding_window_size.label": "スライディングウィンドウ", + "field.storage_format.help": "ディスク上のエキスパートストアのファイル形式。", + "field.storage_format.label": "エキスパート保存形式", + "field.temperature.help": "GRPO ロールアウトのサンプリング温度。", + "field.temperature.label": "ロールアウト温度", + "field.tie_word_embeddings.help": "入力埋め込みを出力ヘッドと共有します。語彙×隠れ層サイズ分のパラメータを節約でき、小規模モデルでは品質低下はわずかです。", + "field.tie_word_embeddings.label": "埋め込みの共有", + "field.tok.fields.help": "読み込む JSON フィールド。空欄なら text、conversations、pairs を自動検出します。", + "field.tok.fields.label": "テキストフィールド", + "field.tok.out_dir.help": "tokenizer.json、チャットテンプレート、レポートの書き込み先。", + "field.tok.out_dir.label": "出力フォルダ", + "field.tok.paths.help": "テキストをサンプリングするファイルまたはフォルダを 1 行に 1 つ。各ソースはそれぞれの割り当てでサンプリングし、結合しません。", + "field.tok.paths.label": "ソースファイル", + "field.tok.per_source.help": "ソースファイルごとのリザーバサンプリングの件数。", + "field.tok.per_source.label": "ソースあたりのテキスト数", + "field.tok.temperature.help": "バイト比率で言語のバランスを取ります。低いほど均等になり、英語が中国語を押し出しません。", + "field.tok.temperature.label": "言語温度", + "field.tok.vocab_size.help": "学習するトークン数。モデルの語彙サイズも同じ値にしてください。", + "field.tok.vocab_size.label": "語彙サイズ", + "field.tokenizer_path.help": "モデルの学習に使うトークナイザーのフォルダ（tokenizer.json を含む）。空欄ならチェックポイント付属のもの、なければ同梱の代替を使います。語彙サイズは一致させてください。", + "field.tokenizer_path.label": "トークナイザー", + "field.train.backend.help": "学習に使うデバイス。auto は Apple シリコンなら MLX、次に CUDA、XPU、MPS、CPU の順で選びます。", + "field.train.backend.label": "学習バックエンド", + "field.train.data_distill.help": "教師モデルが採点するテキストまたは対話データ。ファイルまたはフォルダを 1 行に 1 つ。複数指定は自動で結合します。空欄なら最初のステージのデータセットを使います。", + "field.train.data_distill.label": "蒸留のデータ", + "field.train.data_dpo.help": "chosen と rejected の応答を含む選好ペア。ファイルまたはフォルダを 1 行に 1 つ。複数指定は自動で結合します。空欄なら最初のステージのデータセットを使います。", + "field.train.data_dpo.label": "DPOのデータ", + "field.train.data_grpo.help": "報酬 rollout 用のプロンプト。ファイルまたはフォルダを 1 行に 1 つ。複数指定は自動で結合します。空欄なら最初のステージのデータセットを使います。", + "field.train.data_grpo.label": "GRPOのデータ", + "field.train.data_orpo.help": "chosen と rejected の応答を含む選好ペア（DPO と同じ形式）。ファイルまたはフォルダを 1 行に 1 つ。複数指定は自動で結合します。空欄なら最初のステージのデータセットを使います。", + "field.train.data_orpo.label": "ORPOのデータ", + "field.train.data_path.help": "最初のステージのデータセット。JSONL・JSON・Parquet・CSV・TSV・TXT ファイル、またはそれらを含むフォルダを 1 行に 1 つ。複数指定すると実行フォルダ内で 1 つの JSONL ファイルに結合します。空欄なら合成データでスモークテストを行います。", + "field.train.data_path.label": "データセット", + "field.train.data_sft.help": "conversations（user / assistant のターン）を含む対話データ。ファイルまたはフォルダを 1 行に 1 つ。複数指定は自動で結合します。空欄なら最初のステージのデータセットを使います。", + "field.train.data_sft.label": "教師ありファインチューニングのデータ", + "field.train.init_weight.help": "学習を始めるチェックポイント。空欄の場合、事前学習はランダム初期化から、以降のステージは前のステージから始めます。", + "field.train.init_weight.label": "開始チェックポイント", + "field.train.max_steps.help": "オプティマイザがこのステップ数に達したら停止します。0 なら設定したエポック数を実行します。", + "field.train.max_steps.label": "最大ステップ", + "field.train.sample_prompt.help": "学習の進み具合を読めるよう、一定間隔で補完させるプロンプト。空欄ならステージの既定値を使います。", + "field.train.sample_prompt.label": "サンプルプロンプト", + "field.train.stages.help": "順番に実行するステージ。各ステージは前のステージのチェックポイントから始まります。通常は 事前学習 → SFT です。", + "field.train.stages.label": "ステージ", + "field.train.teacher_path.help": "蒸留専用：生徒モデルが出力を合わせる対象のモデル。", + "field.train.teacher_path.label": "教師チェックポイント", + "field.train.val_ratio.help": "検証用に取り分けるレコードの割合。検証のたびに遅延エキスパートのプローブも実行されます。", + "field.train.val_ratio.label": "検証用の割合", + "field.training_paging.help": "学習中、非アクティブなエキスパートをメモリから外します。遅くなりますが、RAM より大きいモデルを学習できます。", + "field.training_paging.label": "ページング学習", + "field.tune.data_path.help": "自動チューニングの開始前に確認するデータセット（1 行に 1 つ）。現在の試験学習はランダムなトークンで行うため、このデータは試験の損失に影響しません。", + "field.tune.data_path.label": "データセット", + "field.tune.n_trials.help": "Optuna が試す設定の数。", + "field.tune.n_trials.label": "試行回数", + "field.tune.probe_steps.help": "試行ごとの学習ステップ数。成績の悪い試行は早めに打ち切られます。", + "field.tune.probe_steps.label": "試行あたりのステップ数", + "field.tune.search.help": "探索で変更してよいパラメータ。", + "field.tune.search.label": "探索対象", + "field.use_hybrid_attention.help": "MLA（圧縮潜在 KV）層とスライディングウィンドウ層を交互に配置します。長いコンテキストでも KV キャッシュを小さく保てます。", + "field.use_hybrid_attention.label": "ハイブリッドアテンション", + "field.vocab_size.help": "トークン ID の数。トークナイザーの語彙と一致させる必要があります。一致しないと ID が埋め込み表の範囲外になります。", + "field.vocab_size.label": "語彙サイズ", + "field.vram_budget_gb.help": "エキスパートキャッシュが使えるアクセラレータメモリ（GB）。", + "field.vram_budget_gb.label": "VRAM 予算", + "field.weight_decay.help": "AdamW の分離型重み減衰。", + "field.weight_decay.label": "重み減衰", + "group.attention": "アテンションと KV キャッシュ", + "group.autotune": "自動チューニング", + "group.benchmark": "ベンチマーク", + "group.export": "エクスポート", + "group.inference": "生成", + "group.lazy": "遅延エキスパート", + "group.memory": "メモリと遅延エキスパート", + "group.model": "モデル形状", + "group.moe": "混合エキスパート", + "group.objectives": "学習目標", + "group.optimizer": "オプティマイザ", + "group.planner": "自動プラン", + "group.probe": "検証時の遅延エキスパートプローブ", + "group.recurrent": "再帰的な深さ", + "group.rl": "選好学習と強化学習のステージ", + "group.runtime": "ランタイム", + "group.serve": "Ollama / OpenAI サーバー", + "group.speculative": "投機的デコード", + "group.stage_data": "後続ステージのデータ", + "group.tokenizer": "トークナイザー構築", + "group.train": "学習ジョブ", + "infer.arch_note": "アーキテクチャはチェックポイントの設定ファイルから読み込みます。「モデル」ページはファイルにない項目だけを補います。", + "infer.col.cache_hit_rate": "キャッシュヒット率", + "infer.col.cold_miss_predict_hit_rate": "コールドミス予測済み", + "infer.col.decode_time_s": "デコード（秒）", + "infer.col.fallback_weight_rate": "フォールバック重み", + "infer.col.prediction_hit_rate": "予測ヒット率", + "infer.col.prefill_time_s": "プリフィル（秒）", + "infer.col.response_time_s": "応答（秒）", + "infer.col.rss_after_decode_gb": "デコード後 RSS（GB）", + "infer.col.rss_delta_gb": "RSS 変化（GB）", + "infer.col.spec_acceptance_rate": "ドラフト受理率", + "infer.col.tokens": "トークン数", + "infer.col.tokens_per_sec": "トークン/秒", + "infer.common_prefix": "先頭 {n} 文字が一致", + "infer.diff": "出力が異なります", + "infer.full_output": "全メモリ", + "infer.generate": "生成", + "infer.hit_rates": "ヒット率", + "infer.intro": "同じプロンプトを遅延読み込みエキスパートとフルメモリで実行し、出力・速度・ヒット率・メモリを比較します。", + "infer.lazy_output": "遅延エキスパート", + "infer.measured_note": "このマシンでこのプロンプトについて実測。1 回の実行で、平均値ではありません。", + "infer.mode": "モード", + "infer.mode.full_dram": "フルメモリ", + "infer.mode.lazy_offload": "遅延読み込みエキスパート", + "infer.output": "出力", + "infer.prompt": "プロンプト", + "infer.prompt_help": "モデルに送るテキスト。「生のプロンプト」がオフのときは SFT のチャットテンプレートを適用します。", + "infer.ram_ratio": "RAM 読み込み比率", + "infer.same": "出力は同一です", + "infer.similarity": "類似度", + "infer.stats": "実行統計", + "infer.sweep": "RAM 比率スイープ", + "infer.title": "推論", + "io.empty": "まだメトリクスがありません。このプロセスで推論やサーバーを実行すると表示されます。", + "io.intro": "このサーバープロセスで推論とサービングが記録したエキスパートキャッシュと IO のカウンター。", + "io.latest": "最新値", + "io.metric": "指標", + "io.samples": "サンプル数", + "io.seconds": "秒", + "io.title": "IO とキャッシュのメトリクス", + "label.estimated": "推定", + "label.measured": "実測", + "label.measured_meta": "実測（meta デバイスで計数）", + "metric.active_mem_gb": "使用中", + "metric.balance_loss": "エキスパート均衡", + "metric.bleu_zh": "BLEU（中国語）", + "metric.chrf": "chrF", + "metric.error": "エラー", + "metric.eval_s": "評価時間（秒）", + "metric.flores_devtest_n": "FLORES 文数", + "metric.grad_norm": "勾配ノルム", + "metric.load_cv": "負荷の偏り（変動係数）", + "metric.locality_loss": "クラスタ局所性", + "metric.lookahead_loss": "先読みルーター", + "metric.loss": "総損失", + "metric.lr": "学習率", + "metric.peak_mem_gb": "ピーク", + "metric.primary_loss": "交差エントロピー", + "metric.rec_contraction": "収縮率", + "metric.rec_depth_mean": "平均深さ", + "metric.rec_early_exit_ratio": "早期終了の割合", + "metric.route_entropy": "ルーティングエントロピー", + "metric.state": "状態", + "metric.temporal_loss": "時間局所性", + "metric.tokens_per_s": "トークン/秒", + "metric.valid_ce": "検証交差エントロピー", + "metric.valid_ppl": "検証パープレキシティ", + "mon.best": "最良", + "mon.chart.aux": "補助損失", + "mon.chart.loss": "学習損失", + "mon.chart.lr": "学習率", + "mon.chart.memory": "メモリ（GB）", + "mon.chart.optimizer": "勾配ノルム", + "mon.chart.recurrent": "再帰深さ", + "mon.chart.routing": "エキスパートルーティング", + "mon.chart.throughput": "毎秒トークン数", + "mon.checkpoints": "チェックポイント", + "mon.eta": "残り時間", + "mon.event": "イベント", + "mon.events": "イベント", + "mon.grad_norm": "勾配ノルム", + "mon.input": "入力", + "mon.intro": "コマンドラインから起動した長時間学習を含め、任意の実行フォルダをライブ表示します。読み取り専用です。", + "mon.log": "学習ログ", + "mon.loss": "損失", + "mon.lr": "学習率", + "mon.memory": "メモリ", + "mon.no_runs": "まだ実行がありません。「学習」ページから開始してください。", + "mon.phase_note": "ステップは段階ごとに数え直すため、グラフは一度に 1 つの段階だけを表示します。", + "mon.progress": "進捗", + "mon.reference": "参照訳", + "mon.run": "実行", + "mon.samples": "サンプル出力", + "mon.size": "サイズ", + "mon.stale": "しばらく更新がありません", + "mon.throughput": "毎秒トークン数", + "mon.title": "実行モニター", + "mon.tokens_seen": "学習済みトークン", + "mon.updated": "{s} 秒前に更新", + "mon.valid_ce": "検証 CE", + "mon.validation": "検証", + "nav.autotune": "自動チューニング", + "nav.benchmark": "ベンチマーク", + "nav.config": "モデル", + "nav.data": "データ", + "nav.export": "エクスポート", + "nav.inference": "推論", + "nav.io": "IO メトリクス", + "nav.monitor": "モニター", + "nav.serve": "サーバー", + "nav.tokenizer": "トークナイザー", + "nav.train": "学習", + "phase.active": "実行中", + "phase.complete": "完了", + "phase.pending": "待機中", + "picker.add": "選択項目を追加", + "picker.add_button": "追加…", + "picker.browse_tip": "サーバー側のフォルダーをページ内で参照してパスを選びます。", + "picker.empty": "該当するファイルやフォルダはありません", + "picker.entries": "フォルダの内容", + "picker.go": "移動", + "picker.hidden": "隠し項目を表示", + "picker.multi_hint": "複数のファイルやフォルダを追加できます。学習前に自動で結合されます。", + "picker.n_selected": "{n} 件を選択", + "picker.native_fallback": "このマシンではシステムのファイルダイアログを開けないため、アップロードに切り替えます。", + "picker.native_tip": "システムのファイルマネージャーを開きます。コピーはせず、選んだ元のパスをそのまま使います。", + "picker.pick_file": "ファイルを選択してください", + "picker.place.desktop": "デスクトップ", + "picker.place.documents": "書類", + "picker.place.downloads": "ダウンロード", + "picker.place.home": "ホーム", + "picker.place.root": "ルート（/）", + "picker.place.volumes": "すべてのボリューム", + "picker.placeholder_multi": "1 行に 1 つのパス。または「追加…」", + "picker.places": "よく使う場所", + "picker.select": "選択", + "picker.this_folder": "未選択：このフォルダを使います（{path}）", + "picker.title.any": "ファイルまたはフォルダを選択", + "picker.title.checkpoint": "checkpoint を選択", + "picker.title.config": "設定ファイルを選択", + "picker.title.data": "データセットを選択（ファイルまたはフォルダ）", + "picker.title.dir": "フォルダを選択", + "picker.truncated": "先頭の {n} 件のみ表示しています。パスを入力すると直接移動できます", + "picker.up": "1 つ上のフォルダ", + "picker.upload_files": "ファイルをアップロード…", + "picker.upload_folder": "フォルダーをアップロード…", + "picker.upload_tip": "この端末のファイルマネージャーで選んだファイルをサーバー側の uploads フォルダーにアップロードし、そのパスを入力します。", + "picker.uploaded": "{n} 個のファイルをアップロードしました。", + "picker.uploaded_skipped": "{n} 個のファイルをアップロードしました。種類が合わないか隠しファイルの {skipped} 個はスキップしました。", + "picker.uploading": "アップロード中 {done} / {total}…", + "picker.will_merge": "{n} 個のパスを使用前に 1 つの JSONL ファイルへ結合します。", + "plan.apply": "このプランを使う", + "plan.note": "プランナーは目標サイズ、データ、メモリから設定を探索します。学習の前に確認してください。", + "plan.run": "プランを作成", + "plan.stages": "ステージごとの設定", + "plan.summary": "プラン", + "plan.title": "自動プラン", + "plan.tokens": "学習トークン数", + "preset.b1": "1B（64 GB）", + "preset.b1.summary": "1B MoE。64 個の細粒度エキスパート、積極的なメモリ節約。", + "preset.enzh50m": "50M 英→中レシピ", + "preset.enzh50m.summary": "72 時間の英→中レシピ：16k Chronos トークナイザー、shards_v1。", + "preset.m200": "200M（32 GB）", + "preset.m200.summary": "汎用 200M MoE。メモリ節約はバランス型で、minimind 6400 語彙を使用。", + "preset.m500": "500M（32 GB）", + "preset.m500.summary": "500M MoE。積極的なメモリ節約：int8 KV、共有ルーティング、活性化チェックポイント。", + "preset.tiny": "Tiny（スモークテスト / CI）", + "preset.tiny.summary": "1.5M のスモークテスト。CPU で数分で学習できます。", + "probe.accept": "ドラフト受理率", + "probe.agree": "トークン一致率", + "probe.auto_paging": "自動ページング", + "probe.cold_hit": "予測済みのコールドミス", + "probe.conditions": "貪欲デコード、プロンプト {n} 件、各 {t} トークン、実測。", + "probe.exact": "完全一致", + "probe.history": "学習の進行に伴うヒット率と一致率", + "probe.hit": "常駐ヒット率", + "probe.loads": "オンデマンドロード回数", + "probe.none": "まだプローブ結果がありません。検証のたびに 1 回実行されます。", + "probe.pages": "ページング · ページ数", + "probe.pred_hit": "予測ヒット率", + "probe.prefix": "最短の一致プレフィックス", + "probe.reference": "全メモリの参照", + "probe.rss": "ピーク RSS（GB）", + "probe.skip.busy": "前回のプローブがまだ実行中です", + "probe.skip.dense": "密モデルのため、ページングするルーティングエキスパートがありません", + "probe.skip.no_checkpoint": "ディスク上にチェックポイントがまだありません", + "probe.skip.superseded": "ステップ {step} の新しい検証に置き換えられました", + "probe.skip.writing": "チェックポイントを書き込み中です", + "probe.skipped": "スキップ：{reason}", + "probe.title": "遅延エキスパートのプローブ", + "probe.tps": "トークン/秒", + "probe.trained_clusters": "クラスタ局所性で学習済み", + "probe.variant": "組み合わせ", + "probe.vs_full": "全メモリ比の速度", + "probe.wait": "プリフェッチ待ち（秒）", + "rule.clusters_auto": "0 = エキスパート数から決定", + "rule.equals_recurrent_depth": "学習時の深さと同じ", + "rule.loopback": "127.0.0.1（このマシンのみ）", + "rule.sync_is_exact": "sync_on_demand：出力は全メモリと同一", + "rule.vocab_matches_tokenizer": "トークナイザーの語彙サイズと同じ", + "rule.window_median_record": "レコード長の中央値に最も近い 2 のべき乗（不明なら 32）", + "serve.auth_off": "キーなし：このマシンからのみ接続できます。", + "serve.auth_on": "クライアントは API キーを送る必要があります。", + "serve.hint": "接続方法：OLLAMA_HOST={addr} ollama run ", + "serve.intro": "chronos serve を起動します。Ollama と OpenAI のクライアントが接続し、エキスパートは遅延読み込みのままです。", + "serve.remote_warning": "このマシン以外のアドレスにバインドするには「リモートアクセスを許可」と CHRONOS_API_KEY のキーが必要です。キーがないとサーバーは起動しません。", + "serve.running": "{url} で稼働中", + "serve.start": "サーバーを起動", + "serve.stop": "サーバーを停止", + "serve.stopped": "サーバーは停止しました", + "serve.title": "Ollama / OpenAI 互換サーバー", + "source.code": "既定値", + "source.measured": "実測", + "source.planner": "適応プランナー", + "source.rule": "ルール", + "stage.distill": "蒸留", + "stage.dpo": "DPO", + "stage.grpo": "GRPO", + "stage.orpo": "ORPO", + "stage.pretrain": "事前学習", + "stage.sft": "教師ありファインチューニング", + "state.checkpointing": "チェックポイント保存中", + "state.done": "完了", + "state.error": "エラー", + "state.evaluating": "評価中", + "state.failed": "失敗", + "state.finished": "完了", + "state.idle": "待機中", + "state.loading": "読み込み中", + "state.paused": "一時停止", + "state.pending": "保留中", + "state.running": "実行中", + "state.stale": "最近の更新なし", + "state.starting": "起動中", + "state.stopped": "停止済み", + "state.stopping": "停止中", + "tip.auto_empty": "空欄 = 自動", + "tip.default": "既定値", + "tip.range": "範囲", + "tip.recommended": "推奨値", + "tip.source": "根拠", + "tok.algorithm": "アルゴリズム", + "tok.bucket": "バケット", + "tok.byte_fallback": "バイトフォールバック率", + "tok.chars_per_token": "トークンあたり文字数", + "tok.compression": "トークンあたりのバイト数", + "tok.elapsed": "構築時間", + "tok.fertility": "単語あたりのトークン数", + "tok.fingerprint": "フィンガープリント", + "tok.intro": "自分のデータで Chronos トークナイザーを学習します：サンプリング、正規化、バイトフォールバック BPE、保留データでの評価レポート。", + "tok.out_dir_now": "出力フォルダ", + "tok.progress": "構築の進捗", + "tok.report": "評価レポート", + "tok.report_note": "トークナイザーの学習に使っていない保留テキストで実測。", + "tok.roundtrip": "ロスレス往復", + "tok.texts": "テキスト数", + "tok.title": "Chronos トークナイザー構築", + "tok.vocab_use": "語彙の使用率", + "train.config_note": "アーキテクチャは「モデル」ページから取得します。自動計画を適用した場合、各段階の学習率は計画の値です。", + "train.intro": "「モデル」ページの設定で学習ジョブを開始します。進捗はモニターページに表示されます。", + "train.jobs": "学習ジョブ", + "train.open_monitor": "モニターで開く", + "train.probe_note": "検証のたびに、保存したばかりのチェックポイントを遅延エキスパートと全メモリでそれぞれデコードします。プローブはコピー上で低優先度の子プロセスとして動くため、学習は止まりません。", + "train.probe_off": "このジョブでは遅延読み込みチェックがオフです。", + "train.probe_title": "検証のたびに遅延エキスパートを確認", + "train.run_dir": "実行フォルダ", + "train.stage_data_note": "最初より後のステージは、それぞれ専用の形式のデータを読み込みます。空欄の行は最初のステージのデータセットを使います。", + "train.stage_n": "ステージ {n}", + "train.start": "学習を開始", + "train.started": "学習を開始しました", + "train.stop": "学習を停止", + "train.title": "学習", + "train.variants": "{n} 通りの組み合わせをテストします", + "train.variants_capped": "{n} 通りが指定されました。最初の {cap} 通りをテストします", + "trial.complete": "完了", + "trial.failed": "失敗", + "trial.pruned": "枝刈り", + "tune.apply": "モデル設定に適用", + "tune.best": "最良のパラメータ", + "tune.duration": "所要時間", + "tune.intro": "Optuna TPE で学習設定を探索します。各試行は数ステップだけ学習して損失を報告します。", + "tune.loss": "プローブ損失", + "tune.note": "短い学習後のプローブ損失です。順位付けの目安で、最終品質の指標ではありません。", + "tune.start": "探索を開始", + "tune.stop": "探索を停止", + "tune.title": "自動チューニング（Optuna TPE）", + "tune.trial": "試行", + "tune.trials": "試行", + "ui.access_tip.lan": "LAN 経由で接続しています。「アップロード…」はファイルをサーバー側の uploads フォルダーに送り、そのパスを使います。「参照…」はサーバー側のファイルを表示します。", + "ui.access_tip.lan_local": "サーバーと同じマシンで開いているため、ログインは不要です。「アップロード…」はシステムのファイルマネージャーを開き、元のパスをそのまま使います。LAN 上の他の端末にはアクセスキーが必要です。", + "ui.access_tip.local_only": "UI はこのマシン（127.0.0.1）だけで待ち受けています。他の端末からは開けず、ログインも不要です。--lan を付けて起動すると LAN からアクセスできます。", + "ui.apply": "適用", + "ui.backend": "バックエンド", + "ui.browse": "参照…", + "ui.busy": "同じ種類のジョブが実行中です", + "ui.cancel": "キャンセル", + "ui.chart_view": "グラフ", + "ui.clear": "クリア", + "ui.close": "閉じる", + "ui.copy": "コピー", + "ui.empty": "まだありません", + "ui.error": "エラー", + "ui.from_model": "モデルファイルに従う", + "ui.hide_advanced": "詳細オプションを隠す", + "ui.invalid": "マークされた項目を確認してください", + "ui.key": "アクセスキー", + "ui.lan_local": "このマシン · LAN 公開中", + "ui.language": "言語", + "ui.load": "読み込み", + "ui.loading": "読み込み中…", + "ui.local_only": "ローカル接続のみ、ログイン不要", + "ui.log": "ログ", + "ui.login": "ログイン", + "ui.login_blocked": "失敗が多すぎます。5 分後にもう一度お試しください。", + "ui.login_failed": "キーが正しくありません", + "ui.login_needed": "このワークベンチにはサーバーで設定したアクセスキーが必要です。", + "ui.logout": "ログアウト", + "ui.no": "いいえ", + "ui.none": "なし", + "ui.path": "パス", + "ui.phase": "フェーズ", + "ui.raw_json": "生の JSON", + "ui.refresh": "更新", + "ui.remote_auth": "リモート接続、ログイン済み", + "ui.reset": "既定値に戻す", + "ui.result": "結果", + "ui.run": "実行", + "ui.save": "保存", + "ui.saved": "保存しました", + "ui.scan": "スキャン", + "ui.sections": "セクション", + "ui.show_advanced": "詳細オプションを表示", + "ui.skip": "本文へ移動", + "ui.start": "開始", + "ui.started": "開始時刻", + "ui.state": "状態", + "ui.step": "ステップ", + "ui.stop": "停止", + "ui.table_view": "表", + "ui.theme": "テーマ", + "ui.theme.dark": "ダーク", + "ui.theme.light": "ライト", + "ui.theme.system": "システムに合わせる", + "ui.time": "時刻", + "ui.upload": "アップロード…", + "ui.value": "値", + "ui.yes": "はい" +} diff --git a/chronos/webui/i18n/zh-Hans.json b/chronos/webui/i18n/zh-Hans.json new file mode 100644 index 0000000..cf3e9fe --- /dev/null +++ b/chronos/webui/i18n/zh-Hans.json @@ -0,0 +1,860 @@ +{ + "app.subtitle": "循环深度 MoE · 懒加载专家 · 推测解码", + "app.title": "Chronos 工作台", + "bench.better": "更优", + "bench.higher": "越高越好", + "bench.intro": "用 benchmark_compare.py 比较 MiniMind 与 Chronos 的参数量、内存和解码速度。", + "bench.lower": "越低越好", + "bench.measured_note": "由 benchmark_compare.py 在本机实测。", + "bench.metric": "指标", + "bench.none": "还没有结果。运行基准测试后生成。", + "bench.params": "参数量", + "bench.ram_infer": "推理内存", + "bench.ram_model": "模型内存", + "bench.ram_train": "训练内存", + "bench.run": "运行基准测试", + "bench.title": "基准测试：MiniMind 与 Chronos", + "bench.tps": "解码 token/秒", + "bench.unavailable": "当前安装中没有 benchmark_compare.py。", + "choice.absorbed": "吸收", + "choice.adaptive": "自适应", + "choice.aggressive": "激进", + "choice.auto": "自动", + "choice.balanced": "均衡", + "choice.cluster": "按簇", + "choice.compare": "懒加载与全内存对比", + "choice.expanded": "展开（参考）", + "choice.expert": "按专家", + "choice.experts": "激活专家数", + "choice.extreme": "极限", + "choice.fallback_diagnostic": "不加载，仅诊断（fallback_diagnostic）", + "choice.fixed": "固定", + "choice.fixed_then_adaptive": "先固定，后自适应", + "choice.fullload": "仅全内存", + "choice.none": "无", + "choice.offload": "仅懒加载专家", + "choice.on_demand": "按需加载，可跳过迟到专家（on_demand）", + "choice.params": "激活参数量", + "choice.planned": "下次使用规划器", + "choice.recurrent_attention_mode.first_only": "仅首轮注意力，v2 行为（first_only）", + "choice.recurrent_attention_mode.per_iteration": "每轮注意力，每轮独立 KV（per_iteration）", + "choice.recurrent_attention_mode.shared_first": "每轮注意力，共享首轮 KV（shared_first）", + "choice.router_scope.per_iteration": "每轮重新路由（per_iteration）", + "choice.router_scope.shared": "复用首轮路由（shared）", + "choice.scheduled": "按调度", + "choice.shared": "共享", + "choice.sync_on_demand": "同步按需加载（sync_on_demand）", + "choice.tune.search.lr": "学习率", + "choice.window": "最近使用窗口", + "config.active": "每 token 激活参数", + "config.compute": "每 token 权重使用量（含循环）", + "config.estimate": "规模与内存估算", + "config.estimate_note": "根据配置估算，并非实测。", + "config.file": "配置文件", + "config.ignored": "已忽略未知或无效的键", + "config.intro": "所有模型设置集中在这里。光标悬停在任一选项上，可查看含义、默认值和推荐值。", + "config.kv": "KV 缓存", + "config.loaded": "配置已加载", + "config.path_help": "JSON 配置文件。“加载”把它读入表单，“保存”把表单写到这里。相对路径从工作台文件夹起算。", + "config.preset": "预设", + "config.preset_help": "起始配置。应用后会替换本页所有字段，估算随即更新。", + "config.ram": "专家热缓存", + "config.ssd": "磁盘专家存储", + "config.total": "总参数量", + "config.tps": "解码速度粗估（估算）", + "config.train_peak": "训练峰值内存", + "config.train_state": "训练状态（权重、梯度、优化器）", + "config.virtual_layers": "虚拟层数", + "config.vram": "常驻权重 + KV", + "data.chars": "每条记录字符数", + "data.dups": "重复占比", + "data.fields": "字段", + "data.files": "文件数", + "data.full_scan": "全量扫描", + "data.hist": "记录长度分布", + "data.intro": "训练前扫描文件或文件夹，查看格式、字段、文字系统和记录长度。", + "data.max_help": "汇总时最多读取的记录数（100 到 500,000）。", + "data.max_records": "扫描记录数", + "data.median": "中位数", + "data.p95": "95 分位", + "data.path_help": "文件或文件夹（jsonl、json、parquet、csv、tsv、txt），每行一个。多个路径按合并后训练看到的样子作为一个数据集扫描。相对路径从工作台目录算起。", + "data.present": "出现比例", + "data.records": "记录数", + "data.sampled": "前 {n} 条记录", + "data.samples": "样例记录", + "data.scanned": "已扫描", + "data.schema": "识别出的格式", + "data.script": "文字", + "data.scripts": "文字系统", + "data.share": "占比", + "data.title": "数据集检查", + "data.type": "类型", + "err.active_experts": "激活专家数不能超过专家总数", + "err.dialog_open": "已有一个文件对话框打开，请先在系统窗口中完成或取消", + "err.models_dir": "找不到模型目录", + "err.native_local_only": "系统文件对话框只能在服务所在的电脑上使用", + "err.native_unavailable": "这台电脑无法打开系统文件对话框", + "err.no_files": "这些路径下没有可读取的文件", + "err.no_out": "请选择输出目录", + "err.no_paths": "请至少添加一个路径", + "err.not_found": "未找到", + "err.paging": "分页方式请选择 auto、cluster 或 expert", + "err.port": "端口无效", + "err.probe_invalid": "懒加载检查的设置无效", + "err.remote_needs_key": "远程绑定需要允许远程访问并设置 CHRONOS_API_KEY", + "err.target": "请输入目标规模", + "err.upload_batch": "上传批次标识无效，请刷新页面后重试", + "err.upload_disk_full": "服务所在电脑的磁盘剩余空间不足，已停止上传", + "err.upload_failed": "上传中断，请检查网络连接后重试", + "err.upload_kind": "该字段不支持上传", + "err.upload_name": "文件名或文件夹层级不允许（隐藏项、“..”或层级过深）", + "err.upload_too_large": "文件超过上传上限（--max-upload-mb）", + "err.upload_type": "没有可上传的文件：类型与该字段不符，或为不允许的可执行文件", + "event.final": "结束", + "event.lazy_probe": "懒加载探测", + "event.phase_end": "阶段结束", + "event.phase_start": "阶段开始", + "event.resume": "已恢复", + "event.rollback": "回滚", + "event.validation": "验证", + "export.done": "已导出", + "export.format": "格式", + "export.intro": "把检查点写成 safetensors 或 GGUF 文件，并附带 chronos serve 加载所需的专家缓存布局。", + "export.note": "Ollama 客户端通过 chronos serve 运行 Chronos 模型（把 OLLAMA_HOST 设为其地址）；原版 Ollama 无法执行循环 MoE。", + "export.run": "导出", + "export.size": "大小（MB）", + "export.tensors": "张量数", + "export.title": "导出", + "field.accumulation_steps.help": "每次优化器更新累积的微批数。有效批大小 = 批大小 × 累积步数。", + "field.accumulation_steps.label": "梯度累积", + "field.activation_checkpointing.help": "反向传播时重算激活，以速度换取更低的峰值内存。", + "field.activation_checkpointing.label": "激活检查点", + "field.adaptive_halting.help": "每个 token 的停止概率超过阈值后就不再循环，为简单 token 节省计算。", + "field.adaptive_halting.label": "自适应停止（ACT）", + "field.alpha.help": "蒸馏混合比例：α × 教师 KL + (1 − α) × CE。", + "field.alpha.label": "α（蒸馏）", + "field.batch_size.help": "每个微批的序列数，受内存限制。", + "field.batch_size.label": "批大小", + "field.beta.help": "偏好训练中与参考策略 KL 约束的强度。", + "field.beta.label": "β（DPO/GRPO）", + "field.cluster_locality_clusters.help": "局部性约束使用的专家簇数。0 表示按专家数自动选择。", + "field.cluster_locality_clusters.label": "簇数", + "field.cluster_locality_window.help": "簇约束作用的片段长度（token）。与典型记录长度一致时效果最好。", + "field.cluster_locality_window.label": "局部性窗口", + "field.coda_layers.help": "循环之后、输出头之前只运行一次的层。", + "field.coda_layers.label": "收尾层", + "field.cpu_budget_percent.help": "自动模式可使用的物理核心比例（百分比）。", + "field.cpu_budget_percent.label": "CPU 预算", + "field.cpu_threads.help": "PyTorch/BLAS 线程数。自动按物理核心预算设置。", + "field.cpu_threads.label": "CPU 线程", + "field.depth_schedule.help": "训练时每个 batch 的深度选法：固定、自适应，或先固定预热再自适应。", + "field.depth_schedule.label": "深度调度", + "field.depth_schedule_warmup_fraction.help": "训练中先以固定全深度运行的比例，之后才切换为自适应。", + "field.depth_schedule_warmup_fraction.label": "调度预热比例", + "field.draft_exit_depth.help": "作为推测草稿来训练的循环深度。", + "field.draft_exit_depth.label": "草稿深度（训练）", + "field.draft_exit_kl_coef.help": "草稿深度出口与全深度输出之间的额外 KL，用于自推测解码。0 表示关闭。", + "field.draft_exit_kl_coef.label": "草稿出口 KL", + "field.dropout.help": "训练时随机丢弃激活。数据量大时调低，数据很少时调高。", + "field.dropout.label": "Dropout", + "field.dtype.help": "训练与推理的数值精度。自动：后端支持时用 bfloat16，否则 float32。", + "field.dtype.label": "计算精度", + "field.epochs.help": "遍历数据集的次数。设置了最大步数时会更早停止。", + "field.epochs.label": "轮数", + "field.export.auto_cluster.help": "用校准数据把常一起使用的专家分到同一页，懒加载时读取的文件更少。", + "field.export.auto_cluster.label": "构建专家簇", + "field.export.calibration_data_path.help": "用于统计哪些专家会一起使用的文本，每行一个路径；多个会自动合并。留空则跳过自动聚类。", + "field.export.calibration_data_path.label": "校准数据", + "field.export.cluster_batch_size.help": "每个校准批的序列数。", + "field.export.cluster_batch_size.label": "校准批大小", + "field.export.cluster_device.help": "校准运行使用的设备。", + "field.export.cluster_device.label": "校准设备", + "field.export.cluster_max_batches.help": "聚类时观察的批次数。", + "field.export.cluster_max_batches.label": "校准批数", + "field.export.cluster_max_seq_len.help": "每条校准序列的 token 数。", + "field.export.cluster_max_seq_len.label": "校准长度", + "field.export.config_path.help": "覆盖附属配置。通常留空。", + "field.export.config_path.label": "配置文件", + "field.export.copy_cache.help": "把懒加载专家页一并导出，部署时无需重建。", + "field.export.copy_cache.label": "复制专家缓存", + "field.export.expert_cache_dir.help": "要打包的现有懒加载专家缓存。留空则重新构建。", + "field.export.expert_cache_dir.label": "专家缓存", + "field.export.formats.help": "要写出的容器格式。safetensors 和 GGUF 保留 Chronos 元数据，需用 chronos serve 运行。Ollama 客户端连接 chronos serve，而不是原版 Ollama。", + "field.export.formats.label": "格式", + "field.export.model_path.help": "要导出的 checkpoint。", + "field.export.model_path.label": "Checkpoint", + "field.export.output_dir.help": "导出文件的写入位置。", + "field.export.output_dir.label": "输出目录", + "field.fallback_mask_prob.help": "训练时随机屏蔽专家，让模型能容忍缓存未命中。只影响带回退的懒加载推理。", + "field.fallback_mask_prob.label": "缺失模拟", + "field.grad_clip.help": "把全局梯度范数裁剪到该值。0 表示不裁剪。", + "field.grad_clip.label": "梯度裁剪", + "field.halting_threshold.help": "累计停止概率达到该值时 token 停止迭代。越接近 1 跑得越深。", + "field.halting_threshold.label": "停止阈值", + "field.hidden_size.help": "每个 token 向量的宽度。越大能力越强，参数量约按平方增长。必须能被注意力头数整除。", + "field.hidden_size.label": "隐藏维度", + "field.infer.backend.help": "推理运行时。自动按 checkpoint 和本机选择。", + "field.infer.backend.label": "后端", + "field.infer.decode_wait_s.help": "等待预取专家的时长，超时则按需加载。", + "field.infer.decode_wait_s.label": "解码等待（秒）", + "field.infer.draft_depth.help": "起草使用的循环深度。", + "field.infer.draft_depth.label": "草稿深度", + "field.infer.gamma.help": "每次校验前起草的 token 数。", + "field.infer.gamma.label": "草稿长度 γ", + "field.infer.max_tokens.help": "生成 token 数的上限。", + "field.infer.max_tokens.label": "最大新 token 数", + "field.infer.miss_policy.help": "缓存未命中时的处理方式。sync_on_demand 的输出与全内存完全一致。", + "field.infer.miss_policy.label": "未命中策略", + "field.infer.mode.help": "对比：懒加载与全内存并排运行；懒加载：只运行懒加载专家；全内存：全部加载。", + "field.infer.mode.label": "模式", + "field.infer.model_path.help": "要运行的 .pth checkpoint，其附属配置决定架构。", + "field.infer.model_path.label": "Checkpoint", + "field.infer.paging.help": "自动跟随 checkpoint：用簇局部性训练过则按簇分页，否则按专家分页。", + "field.infer.paging.label": "分页布局", + "field.infer.ram_load_ratio.help": "保留在内存中的专家比例。大于 1.0 时为预取留出余量。", + "field.infer.ram_load_ratio.label": "内存比例", + "field.infer.raw_prompt.help": "按原样发送文本，不套用对话模板。", + "field.infer.raw_prompt.label": "原始提示", + "field.infer.residency_decay.help": "规划器遗忘旧路由证据的速度。", + "field.infer.residency_decay.label": "规划器衰减", + "field.infer.residency_planner.help": "淘汰未来最晚才会用到的专家，而不是最久未用的专家。", + "field.infer.residency_planner.label": "下次使用规划器", + "field.infer.speculative.help": "以浅深度起草 token，再用全深度校验，输出不变。50M 模型实测更慢，因此默认关闭。", + "field.infer.speculative.label": "自推测解码", + "field.infer.sweep_ram_ratios.help": "依次运行各标准比例，绘制速度、命中率与内存的关系。", + "field.infer.sweep_ram_ratios.label": "扫描内存比例", + "field.infer.temperature.help": "采样随机度。0 为贪心解码，结果可复现。", + "field.infer.temperature.label": "温度", + "field.infer.verify_wait_s.help": "推测校验阶段的同类等待时长。", + "field.infer.verify_wait_s.label": "校验等待（秒）", + "field.kv_cache_dtype.help": "自动：沿用计算精度；int8：推理 KV 缓存按 token 量化（内存约减半，误差很小）。训练从不量化。", + "field.kv_cache_dtype.label": "KV 缓存精度", + "field.kv_latent_dim.help": "MLA 层缓存的压缩键值潜在向量宽度。越小越省 KV 内存。", + "field.kv_latent_dim.label": "MLA 潜在维度", + "field.kv_quant_group_size.help": "共用一个 int8 缩放系数的通道数。分组越小越精确，体积略大。", + "field.kv_quant_group_size.label": "KV 量化分组", + "field.lambda_balance.help": "促使路由器把 token 均匀分配到各专家。", + "field.lambda_balance.label": "负载均衡", + "field.lambda_cluster_locality.help": "训练同一片段内的 token 留在同一个专家簇里，懒加载推理就能整簇分页。0 表示关闭。", + "field.lambda_cluster_locality.label": "簇局部性", + "field.lambda_lookahead.help": "训练前瞻路由器预测未来 token 会用到的专家。", + "field.lambda_lookahead.label": "前瞻损失", + "field.lambda_lookahead_topk.help": "额外加权，要求准确预测未来的 top-k 专家。", + "field.lambda_lookahead_topk.label": "前瞻 top-k", + "field.lambda_lookahead_union.help": "要求覆盖前瞻窗口内专家并集的权重。", + "field.lambda_lookahead_union.label": "前瞻并集", + "field.lambda_or.help": "ORPO 中优势比项的权重。", + "field.lambda_or.label": "λ（ORPO）", + "field.lambda_router_anchor.help": "微调时让路由保持接近上游 checkpoint。预训练设为 0。", + "field.lambda_router_anchor.label": "路由锚定", + "field.lambda_router_locality.help": "对相邻 token 之间路由变化的轻微惩罚。", + "field.lambda_router_locality.label": "路由局部性", + "field.lambda_temporal.help": "奖励相邻 token 选择相同专家，一段文本需要加载的专家更少。", + "field.lambda_temporal.label": "时间局部性", + "field.layer_aware_cache.help": "按虚拟层而不是物理层跟踪常驻专家。", + "field.layer_aware_cache.label": "按层缓存", + "field.learning_rate.help": "优化器的峰值步长。规划器会按模型规模给出建议；微调时调低。", + "field.learning_rate.label": "学习率", + "field.log_interval.help": "每 N 步写一行指标。", + "field.log_interval.label": "日志间隔", + "field.lookahead_steps.help": "前瞻路由为未来多少个 token 预测专家，以便在用到之前预取。", + "field.lookahead_steps.label": "前瞻步数", + "field.max_gen_len.help": "每次 GRPO 采样生成的 token 数。", + "field.max_gen_len.label": "采样长度", + "field.max_inference_depth.help": "推理时的最大迭代次数。通常等于训练深度；超过训练深度未经验证。", + "field.max_inference_depth.label": "推理深度", + "field.max_position_embeddings.help": "旋转位置编码预备的最长序列。只有超长上下文才需要修改。", + "field.max_position_embeddings.label": "最大位置数", + "field.max_resident_experts.help": "每层内存中保留专家数的硬上限。留空由规划器决定。", + "field.max_resident_experts.label": "最大常驻专家", + "field.max_seq_len.help": "每条训练序列的 token 数。在混合注意力下内存约线性增长。", + "field.max_seq_len.label": "序列长度", + "field.max_steps.help": "达到该优化器步数后停止。0 表示不设上限（跑完所有轮数）。", + "field.max_steps.label": "最大步数", + "field.min_recurrent_depth.help": "深度调度或提前退出允许的最低深度。设为 1 时简单 token 一轮即可退出。", + "field.min_recurrent_depth.label": "最小深度", + "field.mla_cache_rope.help": "额外缓存旋转键部分，用 KV 内存换解码速度。", + "field.mla_cache_rope.label": "缓存 RoPE 键", + "field.mla_decode.help": "展开：参考路径。吸收：把投影折叠进查询（每步计算更少）。自动会混用两条路径，实测有微小 logit 偏差，因此需手动开启。", + "field.mla_decode.label": "MLA 解码路径", + "field.modalities.help": "额外的输入编码器，用逗号分隔（image、audio）。纯文本模型留空，数据集会自动检测。", + "field.modalities.label": "模态", + "field.moe_intermediate_size.help": "每个专家前馈网络的隐藏宽度。0 表示按隐藏维度推导。", + "field.moe_intermediate_size.label": "专家 FFN 宽度", + "field.num_attention_heads.help": "查询头数。隐藏维度必须能被它整除。", + "field.num_attention_heads.label": "注意力头数", + "field.num_experts.help": "每个 MoE 层的路由专家数。增加专家能扩容而不增加每 token 计算量；专家也是懒加载的单位。", + "field.num_experts.label": "专家数", + "field.num_experts_per_tok.help": "每个 token 使用的专家数。top-2 让路由器能从损失中学习；top-1 需关闭 norm_topk_prob。", + "field.num_experts_per_tok.label": "激活专家数（top-k）", + "field.num_generations.help": "GRPO 对每个提示比较的样本数。", + "field.num_generations.label": "每题采样数", + "field.num_hidden_layers.help": "checkpoint 中实际存储的层数（前置层 + 循环组 + 收尾层）。循环会复用这些层，所以虚拟深度可以更大。", + "field.num_hidden_layers.label": "物理层数", + "field.num_key_value_heads.help": "滑动窗口层分组查询注意力的键值头数。越少 KV 缓存越小。", + "field.num_key_value_heads.label": "KV 头数", + "field.num_shared_experts.help": "每个 token 都会经过的常开专家，始终常驻内存。", + "field.num_shared_experts.label": "共享专家", + "field.num_workers.help": "数据加载器的工作进程数。自动按平台选择安全的数量。", + "field.num_workers.label": "数据加载进程", + "field.optimizer_page_dir.help": "分页优化器状态的存放目录。请使用本地高速 SSD。", + "field.optimizer_page_dir.label": "优化器页目录", + "field.optimizer_paging.help": "不活跃专家的优化器状态存到磁盘。", + "field.optimizer_paging.label": "分页优化器", + "field.page_size_mb.help": "单个优化器页文件的大小，单位 MB。", + "field.page_size_mb.label": "页大小", + "field.pinned_memory_max_fraction.help": "用于锁页专家缓冲区的系统内存最大比例。", + "field.pinned_memory_max_fraction.label": "锁页内存占比", + "field.plan.active_mode.help": "按激活专家数或按激活参数量设定每 token 预算。", + "field.plan.active_mode.label": "激活预算方式", + "field.plan.active_value.help": "每 token 激活专家数，或激活参数量（百万）。", + "field.plan.active_value.label": "激活预算", + "field.plan.data_paths.help": "规划器测量这些文件或文件夹，以确定 token 数、序列长度和窗口，每行一个。", + "field.plan.data_paths.label": "数据集", + "field.plan.depth.help": "每个循环组的迭代次数。", + "field.plan.depth.label": "循环深度", + "field.plan.memory_saving.help": "以速度换内存的力度（int8 KV、共享路由、检查点）。", + "field.plan.memory_saving.label": "省内存程度", + "field.plan.num_experts.help": "每个 MoE 层的路由专家数。", + "field.plan.num_experts.label": "专家数", + "field.plan.ram_gb.help": "规划时假定的内存。留空则检测本机。", + "field.plan.ram_gb.label": "内存（GB）", + "field.plan.seq_len.help": "留空则从数据集分析结果推导。", + "field.plan.seq_len.label": "序列长度", + "field.plan.target_m.help": "总参数量，单位百万。", + "field.plan.target_m.label": "目标规模（M）", + "field.plan.token_budget_b.help": "预训练 token 数，单位十亿。留空：遍历一遍数据。", + "field.plan.token_budget_b.label": "Token 预算（B）", + "field.plan.use_moe.help": "规划 MoE 模型；关闭则规划稠密模型。", + "field.plan.use_moe.label": "混合专家", + "field.plan.vocab_size.help": "留空则从分词器报告读取，或使用 16384。", + "field.plan.vocab_size.label": "词表", + "field.ponder_loss_coef.help": "对期望迭代次数的惩罚。越高越促使模型提前停止。", + "field.ponder_loss_coef.label": "思考损失", + "field.prefetch_depth.help": "预取器提前加载预测专家的距离。0 表示不预取。", + "field.prefetch_depth.label": "预取深度", + "field.prelude_layers.help": "循环之前只运行一次的层。前瞻路由读取它们的输出。", + "field.prelude_layers.label": "前置层", + "field.probe.backends.help": "要测试的运行时。每个后端各有一份全内存参考。", + "field.probe.backends.label": "探针后端", + "field.probe.decode_wait_s.help": "解码步骤等待预取专家的时长，超时则按需加载。", + "field.probe.decode_wait_s.label": "解码等待（秒）", + "field.probe.draft_depth.help": "推测变体中起草 token 使用的循环深度。", + "field.probe.draft_depth.label": "草稿深度", + "field.probe.enabled.help": "每次验证后，用懒加载专家解码刚保存的 checkpoint，并与全内存对比：速度、命中率和 token 一致性。探针在低优先级子进程中对 checkpoint 副本运行。", + "field.probe.enabled.label": "懒加载探针", + "field.probe.gamma.help": "推测变体中每次校验前起草的 token 数。", + "field.probe.gamma.label": "草稿长度 γ", + "field.probe.max_new_tokens.help": "每条提示贪心解码的 token 数。越多速率越稳定，探针耗时也越长。", + "field.probe.max_new_tokens.label": "每条提示 token 数", + "field.probe.max_variants.help": "测试组合数的上限，超出部分标记为已截断。", + "field.probe.max_variants.label": "组合上限", + "field.probe.miss_policies.help": "Torch 缓存未命中时的行为。sync_on_demand 加载并等待（输出完全一致）；on_demand 可能跳过迟到的专家；fallback_diagnostic 从不加载（最差情况）。", + "field.probe.miss_policies.label": "未命中策略", + "field.probe.paging.help": "自动跟随模型的训练方式（开启簇局部性损失时按簇分页）。加入按专家可对比逐专家分页。", + "field.probe.paging.label": "分页布局", + "field.probe.ram_ratios.help": "同时允许驻留内存的专家占比。越低越省内存，未命中越多。", + "field.probe.ram_ratios.label": "内存比例", + "field.probe.residency.help": "MLX 淘汰策略：下次使用规划器，或普通的最近使用窗口。", + "field.probe.residency.label": "常驻策略", + "field.probe.residency_decay.help": "下次使用规划器遗忘旧路由证据的速度。", + "field.probe.residency_decay.label": "规划器衰减", + "field.probe.speculative.help": "同时测试带懒加载专家的自推测解码。", + "field.probe.speculative.label": "推测变体", + "field.probe.verify_wait_s.help": "推测校验阶段的同类等待时长。", + "field.probe.verify_wait_s.label": "校验等待（秒）", + "field.ram_budget_mb.help": "懒加载专家缓存可用的内存，单位 MB。留空则改用内存比例。", + "field.ram_budget_mb.label": "内存预算", + "field.recurrent_attention_mode.help": "循环内注意力的运行方式。shared_first：每轮都做注意力，K/V 取第 0 轮（缓存精确，KV 不随深度增长）。per_iteration：每轮一份 KV（D 倍 KV）。first_only：v2 行为。", + "field.recurrent_attention_mode.label": "循环注意力", + "field.recurrent_depth.help": "训练时每个循环组的迭代次数。虚拟层数 = 前置 + 组数 × 深度 + 收尾。", + "field.recurrent_depth.label": "训练深度", + "field.recurrent_depth_jitter.help": "在调度深度附近随机波动，让模型学会在多个深度退出。", + "field.recurrent_depth_jitter.label": "深度抖动", + "field.recurrent_distill_coef.help": "让浅层出口学习全深度的预测，提高推测接受率和提前退出质量。", + "field.recurrent_distill_coef.label": "深度蒸馏", + "field.recurrent_enabled.help": "让中间的共享层组每个 token 运行多次，用少量权重获得更深网络的效果。", + "field.recurrent_enabled.label": "循环深度", + "field.recurrent_groups.help": "前置层与收尾层之间组成循环组的层数。留空则按 层数 − 前置 − 收尾 推导。", + "field.recurrent_groups.label": "循环组数", + "field.recurrent_stability_coef.help": "保持循环更新收缩，避免多次迭代后漂移或发散。", + "field.recurrent_stability_coef.label": "稳定性损失", + "field.recurrent_strategy.help": "训练时循环深度分桶的采样策略。", + "field.recurrent_strategy.label": "深度策略", + "field.reward_spec.help": "toy 使用内置奖励，或 lm:/path/to/reward-model[,device=cuda]。", + "field.reward_spec.label": "奖励（GRPO）", + "field.rope_dim.help": "MLA 层每个头中携带旋转位置信息的维度数。", + "field.rope_dim.label": "RoPE 维度", + "field.rope_theta.help": "旋转位置编码的基础频率。值越大，长上下文的位置分辨率越好。", + "field.rope_theta.label": "RoPE 基数", + "field.router_scope.help": "per_iteration：每轮循环重新路由。共享：复用第 0 轮的路由，用到的专家更少，更适合懒加载。", + "field.router_scope.label": "路由范围", + "field.save_dir.help": "checkpoint 及其 .config.json 附属文件的写入位置。", + "field.save_dir.label": "Checkpoint 目录", + "field.save_interval.help": "每 N 步保存（并验证）一次。", + "field.save_interval.label": "保存间隔", + "field.serve.allow_remote.help": "监听非回环地址。服务器环境中必须设置 CHRONOS_API_KEY，客户端须携带该密钥。", + "field.serve.allow_remote.label": "允许远程访问", + "field.serve.device.help": "服务使用的运行时。", + "field.serve.device.label": "设备", + "field.serve.dtype.help": "服务使用的精度。", + "field.serve.dtype.label": "精度", + "field.serve.host.help": "监听的地址。127.0.0.1 只允许本机访问。", + "field.serve.host.label": "监听地址", + "field.serve.kv_cache_dtype.help": "覆盖 KV 存储方式。int8 约减半 KV 内存；留空保持模型设置。", + "field.serve.kv_cache_dtype.label": "KV 缓存精度", + "field.serve.lazy_experts.help": "专家保存在 SSD 上并按需换入，通过 Ollama 协议调用时内存同样保持较低。", + "field.serve.lazy_experts.label": "懒加载专家", + "field.serve.max_resident_experts.help": "每层内存专家数的硬上限。留空：按内存占比推导。", + "field.serve.max_resident_experts.label": "最大常驻专家", + "field.serve.mla_cache_rope.help": "缓存旋转后的键：内存更多，重算更少。", + "field.serve.mla_cache_rope.label": "缓存 RoPE 键", + "field.serve.mla_decode.help": "覆盖 MLA 解码路径。留空保持模型设置。", + "field.serve.mla_decode.label": "MLA 解码", + "field.serve.models_dir.help": "存放导出模型的目录，每个模型显示为一个 Ollama 模型名。", + "field.serve.models_dir.label": "模型目录", + "field.serve.paging.help": "自动跟随 checkpoint 的训练方式（按簇或按专家分页）。", + "field.serve.paging.label": "分页布局", + "field.serve.port.help": "TCP 端口。11435 可避免与原版 Ollama 的 11434 冲突。", + "field.serve.port.label": "端口", + "field.serve.ram_budget_mb.help": "专家内存的硬上限。留空：不设上限。", + "field.serve.ram_budget_mb.label": "内存预算（MB）", + "field.serve.ram_expert_fraction.help": "允许驻留内存的专家比例。", + "field.serve.ram_expert_fraction.label": "内存专家占比", + "field.serve.router_scope.help": "覆盖跨迭代的路由复用方式。留空保持模型设置。", + "field.serve.router_scope.label": "路由范围", + "field.sliding_window_size.help": "每个滑动窗口层能看到的 token 数。短窗口适合句级任务，并限制 KV 内存。", + "field.sliding_window_size.label": "滑动窗口", + "field.storage_format.help": "磁盘专家存储的文件格式。", + "field.storage_format.label": "专家存储格式", + "field.temperature.help": "GRPO 采样的温度。", + "field.temperature.label": "采样温度", + "field.tie_word_embeddings.help": "输入嵌入与输出头共用权重。可省下 词表 × 隐藏维度 个参数，小模型几乎不损失质量。", + "field.tie_word_embeddings.label": "共享嵌入", + "field.tok.fields.help": "要读取的 JSON 字段。留空自动识别 text、conversations 和 pairs。", + "field.tok.fields.label": "文本字段", + "field.tok.out_dir.help": "tokenizer.json、对话模板和报告的写入位置。", + "field.tok.out_dir.label": "输出目录", + "field.tok.paths.help": "抽样文本的文件或文件夹，每行一个。每个来源按各自配额抽样，不做合并。", + "field.tok.paths.label": "源文件", + "field.tok.per_source.help": "每个源文件的水库抽样数量。", + "field.tok.per_source.label": "每源文本数", + "field.tok.temperature.help": "按字节占比平衡各语言。越低越均衡，避免英文挤占中文。", + "field.tok.temperature.label": "语言温度", + "field.tok.vocab_size.help": "要学习的 token 数。模型词表需设为相同的值。", + "field.tok.vocab_size.label": "词表大小", + "field.tokenizer_path.help": "模型训练所用的分词器文件夹（含 tokenizer.json）。留空时使用 checkpoint 自带的分词器，否则用内置兜底分词器。词表大小必须一致。", + "field.tokenizer_path.label": "分词器", + "field.train.backend.help": "训练使用的设备。自动：Apple 芯片优先 MLX，其次 CUDA、XPU、MPS、CPU。", + "field.train.backend.label": "训练后端", + "field.train.data_distill.help": "由教师模型评分的文本或对话数据。文件或文件夹，每行一个；多个会自动合并。留空则沿用第一阶段的数据集。", + "field.train.data_distill.label": "蒸馏数据", + "field.train.data_dpo.help": "偏好对数据，含 chosen 与 rejected 回复。文件或文件夹，每行一个；多个会自动合并。留空则沿用第一阶段的数据集。", + "field.train.data_dpo.label": "DPO数据", + "field.train.data_grpo.help": "用于奖励 rollout 的 prompt 数据。文件或文件夹，每行一个；多个会自动合并。留空则沿用第一阶段的数据集。", + "field.train.data_grpo.label": "GRPO数据", + "field.train.data_orpo.help": "偏好对数据，含 chosen 与 rejected 回复，格式同 DPO。文件或文件夹，每行一个；多个会自动合并。留空则沿用第一阶段的数据集。", + "field.train.data_orpo.label": "ORPO数据", + "field.train.data_path.help": "第一阶段的数据集：JSONL、JSON、Parquet、CSV、TSV 或 TXT 文件，或包含这些文件的文件夹，每行一个。多个数据集会在运行目录中合并为一个 JSONL 文件。留空则运行合成冒烟测试。", + "field.train.data_path.label": "数据集", + "field.train.data_sft.help": "对话数据，含 conversations（user / assistant 轮次）。文件或文件夹，每行一个；多个会自动合并。留空则沿用第一阶段的数据集。", + "field.train.data_sft.label": "监督微调数据", + "field.train.init_weight.help": "起始 checkpoint。留空：预训练随机初始化，后续阶段接上一阶段。", + "field.train.init_weight.label": "起始权重", + "field.train.max_steps.help": "达到该优化器步数后停止。0 表示跑完配置的轮数。", + "field.train.max_steps.label": "最大步数", + "field.train.sample_prompt.help": "训练过程中定期续写的提示词，方便查看进展。留空使用阶段默认值。", + "field.train.sample_prompt.label": "采样提示词", + "field.train.stages.help": "按顺序运行，每个阶段从上一阶段的 checkpoint 开始。常用路径是 预训练 → SFT。", + "field.train.stages.label": "阶段", + "field.train.teacher_path.help": "仅用于蒸馏：学生模型要拟合其输出的模型。", + "field.train.teacher_path.label": "教师 checkpoint", + "field.train.val_ratio.help": "留作验证的记录比例。每次验证还会触发懒加载探针。", + "field.train.val_ratio.label": "验证集比例", + "field.training_paging.help": "训练时把不活跃的专家移出内存。更慢，但能训练超过内存容量的模型。", + "field.training_paging.label": "分页训练", + "field.tune.data_path.help": "自动调参启动前会检查这些数据集，每行一个。当前的试训使用随机 token，所以这些数据不影响试训损失。", + "field.tune.data_path.label": "数据集", + "field.tune.n_trials.help": "Optuna 尝试的配置数量。", + "field.tune.n_trials.label": "试验次数", + "field.tune.probe_steps.help": "每次试验的训练步数；表现差的试验会被提前剪枝。", + "field.tune.probe_steps.label": "每次试验步数", + "field.tune.search.help": "搜索可以调整的参数。", + "field.tune.search.label": "搜索参数", + "field.use_hybrid_attention.help": "MLA（压缩潜在 KV）与滑动窗口层交替，长上下文时 KV 缓存仍然很小。", + "field.use_hybrid_attention.label": "混合注意力", + "field.vocab_size.help": "token id 的数量，必须等于分词器词表大小，否则 id 会超出嵌入表。", + "field.vocab_size.label": "词表大小", + "field.vram_budget_gb.help": "专家缓存可用的加速器内存，单位 GB。", + "field.vram_budget_gb.label": "显存预算", + "field.weight_decay.help": "AdamW 解耦权重衰减。", + "field.weight_decay.label": "权重衰减", + "group.attention": "注意力与 KV 缓存", + "group.autotune": "自动调参", + "group.benchmark": "基准测试", + "group.export": "导出", + "group.inference": "生成", + "group.lazy": "懒加载专家", + "group.memory": "内存与懒加载专家", + "group.model": "模型形状", + "group.moe": "混合专家", + "group.objectives": "训练目标", + "group.optimizer": "优化器", + "group.planner": "自动规划", + "group.probe": "验证时的懒加载探针", + "group.recurrent": "循环深度", + "group.rl": "偏好与强化学习阶段", + "group.runtime": "运行时", + "group.serve": "Ollama / OpenAI 服务", + "group.speculative": "推测解码", + "group.stage_data": "后续阶段的数据", + "group.tokenizer": "分词器构建", + "group.train": "训练任务", + "infer.arch_note": "架构读取自检查点的配置文件；“模型”页只补充文件中缺少的项。", + "infer.col.cache_hit_rate": "缓存命中率", + "infer.col.cold_miss_predict_hit_rate": "冷未命中已预测", + "infer.col.decode_time_s": "解码（秒）", + "infer.col.fallback_weight_rate": "回退权重", + "infer.col.prediction_hit_rate": "预测命中率", + "infer.col.prefill_time_s": "预填充（秒）", + "infer.col.response_time_s": "响应（秒）", + "infer.col.rss_after_decode_gb": "解码后 RSS（GB）", + "infer.col.rss_delta_gb": "RSS 变化（GB）", + "infer.col.spec_acceptance_rate": "草稿接受率", + "infer.col.tokens": "token 数", + "infer.col.tokens_per_sec": "token/秒", + "infer.common_prefix": "前 {n} 个字符相同", + "infer.diff": "输出不同", + "infer.full_output": "全内存", + "infer.generate": "生成", + "infer.hit_rates": "命中率", + "infer.intro": "同一个提示分别用懒加载专家和全内存运行，比较输出、速度、命中率和内存。", + "infer.lazy_output": "懒加载专家", + "infer.measured_note": "在本机针对此提示实测，单次运行，不是平均值。", + "infer.mode": "模式", + "infer.mode.full_dram": "全内存", + "infer.mode.lazy_offload": "懒加载专家", + "infer.output": "输出", + "infer.prompt": "提示词", + "infer.prompt_help": "发送给模型的文本。未开启“原始提示”时会套用 SFT 对话模板。", + "infer.ram_ratio": "内存加载比例", + "infer.same": "输出完全一致", + "infer.similarity": "相似度", + "infer.stats": "运行统计", + "infer.sweep": "内存比例扫描", + "infer.title": "推理", + "io.empty": "暂无指标。本进程运行推理或服务时会出现。", + "io.intro": "本服务进程中推理和服务记录的专家缓存与 IO 计数。", + "io.latest": "最新值", + "io.metric": "指标", + "io.samples": "样本数", + "io.seconds": "秒", + "io.title": "IO 与缓存指标", + "label.estimated": "估算", + "label.measured": "实测", + "label.measured_meta": "实测（meta 设备计数）", + "metric.active_mem_gb": "活跃", + "metric.balance_loss": "专家均衡", + "metric.bleu_zh": "BLEU（中文）", + "metric.chrf": "chrF", + "metric.error": "错误", + "metric.eval_s": "评测用时（秒）", + "metric.flores_devtest_n": "FLORES 句数", + "metric.grad_norm": "梯度范数", + "metric.load_cv": "负载不均（变异系数）", + "metric.locality_loss": "簇局部性", + "metric.lookahead_loss": "前瞻路由", + "metric.loss": "总损失", + "metric.lr": "学习率", + "metric.peak_mem_gb": "峰值", + "metric.primary_loss": "交叉熵", + "metric.rec_contraction": "收缩率", + "metric.rec_depth_mean": "平均深度", + "metric.rec_early_exit_ratio": "提前退出占比", + "metric.route_entropy": "路由熵", + "metric.state": "状态", + "metric.temporal_loss": "时间局部性", + "metric.tokens_per_s": "token/秒", + "metric.valid_ce": "验证交叉熵", + "metric.valid_ppl": "验证困惑度", + "mon.best": "最佳", + "mon.chart.aux": "辅助损失", + "mon.chart.loss": "训练损失", + "mon.chart.lr": "学习率", + "mon.chart.memory": "内存（GB）", + "mon.chart.optimizer": "梯度范数", + "mon.chart.recurrent": "循环深度", + "mon.chart.routing": "专家路由", + "mon.chart.throughput": "每秒 token 数", + "mon.checkpoints": "Checkpoint", + "mon.eta": "剩余时间", + "mon.event": "事件", + "mon.events": "事件", + "mon.grad_norm": "梯度范数", + "mon.input": "输入", + "mon.intro": "实时查看任意运行目录，包括命令行启动的长时间训练。只读。", + "mon.log": "训练日志", + "mon.loss": "损失", + "mon.lr": "学习率", + "mon.memory": "内存", + "mon.no_runs": "还没有运行记录。请在训练页面启动。", + "mon.phase_note": "每个阶段的步数从头计数，因此图表一次只显示一个阶段。", + "mon.progress": "进度", + "mon.reference": "参考译文", + "mon.run": "运行", + "mon.samples": "样例", + "mon.size": "大小", + "mon.stale": "已有一段时间没有更新", + "mon.throughput": "每秒 token 数", + "mon.title": "运行监控", + "mon.tokens_seen": "已训练 token", + "mon.updated": "{s} 秒前更新", + "mon.valid_ce": "验证 CE", + "mon.validation": "验证", + "nav.autotune": "自动调参", + "nav.benchmark": "基准测试", + "nav.config": "模型", + "nav.data": "数据", + "nav.export": "导出", + "nav.inference": "推理", + "nav.io": "IO 指标", + "nav.monitor": "监控", + "nav.serve": "服务", + "nav.tokenizer": "分词器", + "nav.train": "训练", + "phase.active": "进行中", + "phase.complete": "已完成", + "phase.pending": "等待中", + "picker.add": "添加所选", + "picker.add_button": "添加…", + "picker.browse_tip": "在页面内浏览服务所在电脑上的文件夹并选取路径。", + "picker.empty": "这里没有符合类型的文件或文件夹", + "picker.entries": "文件夹内容", + "picker.go": "前往", + "picker.hidden": "显示隐藏项", + "picker.multi_hint": "可以添加多个文件或文件夹，训练前会自动合并。", + "picker.n_selected": "已选 {n} 项", + "picker.native_fallback": "这台电脑无法打开系统文件对话框，已改为上传方式。", + "picker.native_tip": "打开系统文件管理器。文件不会被复制，直接使用所选的原路径。", + "picker.pick_file": "请选择一个文件", + "picker.place.desktop": "桌面", + "picker.place.documents": "文稿", + "picker.place.downloads": "下载", + "picker.place.home": "主目录", + "picker.place.root": "根目录（/）", + "picker.place.volumes": "所有磁盘", + "picker.placeholder_multi": "每行一个路径，或点“添加…”", + "picker.places": "常用位置", + "picker.select": "选择", + "picker.this_folder": "未选择条目：将使用当前文件夹（{path}）", + "picker.title.any": "选择文件或文件夹", + "picker.title.checkpoint": "选择 checkpoint", + "picker.title.config": "选择配置文件", + "picker.title.data": "选择数据集（文件或文件夹）", + "picker.title.dir": "选择文件夹", + "picker.truncated": "只显示前 {n} 项；输入路径可直接进入更深的目录", + "picker.up": "上一级文件夹", + "picker.upload_files": "上传文件…", + "picker.upload_folder": "上传文件夹…", + "picker.upload_tip": "用本机的文件管理器选择文件，上传到服务所在电脑的 uploads 文件夹，并填入上传后的路径。", + "picker.uploaded": "已上传 {n} 个文件。", + "picker.uploaded_skipped": "已上传 {n} 个文件；{skipped} 个类型不符或为隐藏文件，已跳过。", + "picker.uploading": "正在上传 {done} / {total}…", + "picker.will_merge": "{n} 个路径将在使用前合并为一个 JSONL 文件。", + "plan.apply": "使用此方案", + "plan.note": "规划器会按目标规模、数据和内存搜索配置；训练前请先检查。", + "plan.run": "规划", + "plan.stages": "各阶段设置", + "plan.summary": "方案", + "plan.title": "自动规划", + "plan.tokens": "训练 token 数", + "preset.b1": "1B（64 GB）", + "preset.b1.summary": "1B MoE，64 个细粒度专家，激进省内存。", + "preset.enzh50m": "50M 英译中配方", + "preset.enzh50m.summary": "72 小时英译中配方：16k Chronos 分词器，shards_v1。", + "preset.m200": "200M（32 GB）", + "preset.m200.summary": "通用 200M MoE，均衡省内存，使用 minimind 6400 词表。", + "preset.m500": "500M（32 GB）", + "preset.m500.summary": "500M MoE，激进省内存：int8 KV、共享路由、激活检查点。", + "preset.tiny": "Tiny（冒烟测试 / CI）", + "preset.tiny.summary": "1.5M 冒烟测试，CPU 上几分钟即可训练完。", + "probe.accept": "草稿接受率", + "probe.agree": "token 一致率", + "probe.auto_paging": "自动分页", + "probe.cold_hit": "冷未命中预测率", + "probe.conditions": "贪心解码，{n} 条提示，每条 {t} 个 token；实测。", + "probe.exact": "完全一致", + "probe.history": "训练过程中的命中率与一致率", + "probe.hit": "常驻命中率", + "probe.loads": "按需加载次数", + "probe.none": "还没有探针结果。每次验证后会运行一次。", + "probe.pages": "分页 · 页数", + "probe.pred_hit": "预测命中率", + "probe.prefix": "最短一致前缀", + "probe.reference": "全内存参考", + "probe.rss": "峰值 RSS（GB）", + "probe.skip.busy": "上一次探测仍在运行", + "probe.skip.dense": "稠密模型，没有可分页的路由专家", + "probe.skip.no_checkpoint": "磁盘上还没有检查点", + "probe.skip.superseded": "已被第 {step} 步的新验证取代", + "probe.skip.writing": "检查点仍在写入", + "probe.skipped": "已跳过：{reason}", + "probe.title": "懒加载探针", + "probe.tps": "token/秒", + "probe.trained_clusters": "使用簇局部性训练", + "probe.variant": "组合", + "probe.vs_full": "相对全内存速度", + "probe.wait": "预取等待（秒）", + "rule.clusters_auto": "0 = 按专家数推导", + "rule.equals_recurrent_depth": "等于训练深度", + "rule.loopback": "127.0.0.1（仅本机）", + "rule.sync_is_exact": "sync_on_demand：输出与全内存完全一致", + "rule.vocab_matches_tokenizer": "等于分词器词表大小", + "rule.window_median_record": "最接近记录长度中位数的 2 的幂（未知时为 32）", + "serve.auth_off": "未设密钥：仅本机可访问。", + "serve.auth_on": "客户端必须携带 API 密钥。", + "serve.hint": "连接方式：OLLAMA_HOST={addr} ollama run <模型>", + "serve.intro": "启动 chronos serve：Ollama 和 OpenAI 客户端连接到它，专家保持懒加载。", + "serve.remote_warning": "绑定到本机以外的地址需要勾选允许远程访问，并在 CHRONOS_API_KEY 中设置密钥；没有密钥时服务拒绝启动。", + "serve.running": "服务地址 {url}", + "serve.start": "启动服务", + "serve.stop": "停止服务", + "serve.stopped": "服务已停止", + "serve.title": "Ollama / OpenAI 兼容服务", + "source.code": "默认值", + "source.measured": "实测", + "source.planner": "自适应规划器", + "source.rule": "规则", + "stage.distill": "蒸馏", + "stage.dpo": "DPO", + "stage.grpo": "GRPO", + "stage.orpo": "ORPO", + "stage.pretrain": "预训练", + "stage.sft": "监督微调", + "state.checkpointing": "正在保存 checkpoint", + "state.done": "完成", + "state.error": "出错", + "state.evaluating": "评估中", + "state.failed": "失败", + "state.finished": "已完成", + "state.idle": "空闲", + "state.loading": "加载中", + "state.paused": "已暂停", + "state.pending": "等待中", + "state.running": "运行中", + "state.stale": "近期无更新", + "state.starting": "启动中", + "state.stopped": "已停止", + "state.stopping": "正在停止", + "tip.auto_empty": "留空 = 自动", + "tip.default": "默认值", + "tip.range": "范围", + "tip.recommended": "推荐值", + "tip.source": "来源", + "tok.algorithm": "算法", + "tok.bucket": "分桶", + "tok.byte_fallback": "字节回退率", + "tok.chars_per_token": "每 token 字符数", + "tok.compression": "每 token 字节数", + "tok.elapsed": "构建用时", + "tok.fertility": "每词 token 数", + "tok.fingerprint": "指纹", + "tok.intro": "用你自己的数据训练 Chronos 分词器：抽样、规范化、字节回退 BPE，并输出留出集评估报告。", + "tok.out_dir_now": "输出文件夹", + "tok.progress": "构建进度", + "tok.report": "评估报告", + "tok.report_note": "在未参与分词器训练的留出文本上实测。", + "tok.roundtrip": "无损往返", + "tok.texts": "文本数", + "tok.title": "Chronos 分词器构建", + "tok.vocab_use": "词表利用率", + "train.config_note": "架构来自“模型”页；应用过自动规划时，各阶段学习率来自规划结果。", + "train.intro": "使用“模型”页的设置启动训练任务，进度显示在监控页。", + "train.jobs": "训练任务", + "train.open_monitor": "在监控中打开", + "train.probe_note": "每次验证都会用懒加载专家和全内存分别解码刚保存的 checkpoint。探针以低优先级子进程在副本上运行，不会中断训练。", + "train.probe_off": "此任务已关闭懒加载检查。", + "train.probe_title": "每次验证时检查懒加载", + "train.run_dir": "运行目录", + "train.stage_data_note": "第一阶段之后的每个阶段读取各自格式的数据。某一行留空时沿用第一阶段的数据集。", + "train.stage_n": "阶段 {n}", + "train.start": "开始训练", + "train.started": "训练已开始", + "train.stop": "停止训练", + "train.title": "训练", + "train.variants": "将测试 {n} 种组合", + "train.variants_capped": "请求了 {n} 种组合，将测试前 {cap} 种", + "trial.complete": "完成", + "trial.failed": "失败", + "trial.pruned": "已剪枝", + "tune.apply": "应用到模型配置", + "tune.best": "最佳参数", + "tune.duration": "用时", + "tune.intro": "用 Optuna TPE 搜索训练设置：每次试验只训练几步并报告损失。", + "tune.loss": "探测损失", + "tune.note": "短程训练后的探测损失，只用于排序，不代表最终质量。", + "tune.start": "开始搜索", + "tune.stop": "停止搜索", + "tune.title": "自动调参（Optuna TPE）", + "tune.trial": "试验", + "tune.trials": "试验", + "ui.access_tip.lan": "你通过局域网访问。“上传…”会把文件传到服务所在电脑的 uploads 文件夹，再使用那里的路径；“浏览…”列出的是服务所在电脑上的文件。", + "ui.access_tip.lan_local": "你在服务所在的电脑上打开，无需登录；“上传…”会打开系统文件管理器并直接使用原路径。局域网内其他设备需要访问密钥才能登录。", + "ui.access_tip.local_only": "界面只监听本机（127.0.0.1），其他设备无法打开，也无需登录。启动时加 --lan 可开放局域网访问。", + "ui.apply": "应用", + "ui.backend": "后端", + "ui.browse": "浏览…", + "ui.busy": "同类任务正在运行", + "ui.cancel": "取消", + "ui.chart_view": "图表", + "ui.clear": "清空", + "ui.close": "关闭", + "ui.copy": "复制", + "ui.empty": "暂无内容", + "ui.error": "错误", + "ui.from_model": "沿用模型文件", + "ui.hide_advanced": "隐藏高级选项", + "ui.invalid": "请检查标出的字段", + "ui.key": "访问密钥", + "ui.lan_local": "本机访问 · 局域网已开放", + "ui.language": "语言", + "ui.load": "加载", + "ui.loading": "加载中…", + "ui.local_only": "仅本机访问，无需登录", + "ui.log": "日志", + "ui.login": "登录", + "ui.login_blocked": "失败次数过多，请 5 分钟后再试。", + "ui.login_failed": "密钥错误", + "ui.login_needed": "此工作台需要服务器上设置的访问密钥。", + "ui.logout": "退出登录", + "ui.no": "否", + "ui.none": "无", + "ui.path": "路径", + "ui.phase": "阶段", + "ui.raw_json": "原始 JSON", + "ui.refresh": "刷新", + "ui.remote_auth": "远程访问，已登录", + "ui.reset": "恢复默认", + "ui.result": "结果", + "ui.run": "运行", + "ui.save": "保存", + "ui.saved": "已保存", + "ui.scan": "扫描", + "ui.sections": "功能分区", + "ui.show_advanced": "显示高级选项", + "ui.skip": "跳到正文", + "ui.start": "开始", + "ui.started": "开始时间", + "ui.state": "状态", + "ui.step": "步", + "ui.stop": "停止", + "ui.table_view": "表格", + "ui.theme": "主题", + "ui.theme.dark": "深色", + "ui.theme.light": "浅色", + "ui.theme.system": "跟随系统", + "ui.time": "时间", + "ui.upload": "上传…", + "ui.value": "数值", + "ui.yes": "是" +} diff --git a/chronos/webui/i18n/zh-Hant.json b/chronos/webui/i18n/zh-Hant.json new file mode 100644 index 0000000..7510d49 --- /dev/null +++ b/chronos/webui/i18n/zh-Hant.json @@ -0,0 +1,860 @@ +{ + "app.subtitle": "循環深度 MoE · 延遲載入專家 · 推測解碼", + "app.title": "Chronos 工作台", + "bench.better": "較佳者", + "bench.higher": "越高越好", + "bench.intro": "以 benchmark_compare.py 比較 MiniMind 與 Chronos 的參數量、記憶體與解碼速度。", + "bench.lower": "越低越好", + "bench.measured_note": "由 benchmark_compare.py 在本機實測。", + "bench.metric": "指標", + "bench.none": "尚無結果。執行基準測試後產生。", + "bench.params": "參數量", + "bench.ram_infer": "推論記憶體", + "bench.ram_model": "模型記憶體", + "bench.ram_train": "訓練記憶體", + "bench.run": "執行基準測試", + "bench.title": "基準測試：MiniMind 與 Chronos", + "bench.tps": "解碼 token/秒", + "bench.unavailable": "目前安裝中沒有 benchmark_compare.py。", + "choice.absorbed": "吸收", + "choice.adaptive": "自適應", + "choice.aggressive": "積極", + "choice.auto": "自動", + "choice.balanced": "平衡", + "choice.cluster": "按叢集", + "choice.compare": "延遲載入與全記憶體比較", + "choice.expanded": "展開（參考）", + "choice.expert": "按專家", + "choice.experts": "啟用專家數", + "choice.extreme": "極限", + "choice.fallback_diagnostic": "不載入，僅診斷（fallback_diagnostic）", + "choice.fixed": "固定", + "choice.fixed_then_adaptive": "先固定，後自適應", + "choice.fullload": "僅全記憶體", + "choice.none": "無", + "choice.offload": "僅延遲載入專家", + "choice.on_demand": "按需載入，可略過遲到專家（on_demand）", + "choice.params": "啟用參數量", + "choice.planned": "下次使用規劃器", + "choice.recurrent_attention_mode.first_only": "僅首輪注意力，v2 行為（first_only）", + "choice.recurrent_attention_mode.per_iteration": "每輪注意力，每輪獨立 KV（per_iteration）", + "choice.recurrent_attention_mode.shared_first": "每輪注意力，共用首輪 KV（shared_first）", + "choice.router_scope.per_iteration": "每輪重新路由（per_iteration）", + "choice.router_scope.shared": "沿用首輪路由（shared）", + "choice.scheduled": "依排程", + "choice.shared": "共用", + "choice.sync_on_demand": "同步按需載入（sync_on_demand）", + "choice.tune.search.lr": "學習率", + "choice.window": "最近使用視窗", + "config.active": "每 token 啟用參數", + "config.compute": "每 token 權重使用量（含循環）", + "config.estimate": "規模與記憶體估算", + "config.estimate_note": "依設定估算，並非實測。", + "config.file": "設定檔", + "config.ignored": "已忽略未知或無效的鍵", + "config.intro": "所有模型設定集中於此。游標停在任一選項上，可查看含義、預設值與建議值。", + "config.kv": "KV 快取", + "config.loaded": "設定已載入", + "config.path_help": "JSON 設定檔。「載入」把它讀入表單，「儲存」把表單寫到這裡。相對路徑從工作台資料夾起算。", + "config.preset": "預設組合", + "config.preset_help": "起始設定。套用後會取代本頁所有欄位，估算隨即更新。", + "config.ram": "專家熱快取", + "config.ssd": "磁碟專家儲存", + "config.total": "總參數量", + "config.tps": "解碼速度粗估（估算）", + "config.train_peak": "訓練峰值記憶體", + "config.train_state": "訓練狀態（權重、梯度、最佳化器）", + "config.virtual_layers": "虛擬層數", + "config.vram": "常駐權重 + KV", + "data.chars": "每條記錄字元數", + "data.dups": "重複比例", + "data.fields": "欄位", + "data.files": "檔案數", + "data.full_scan": "完整掃描", + "data.hist": "記錄長度分布", + "data.intro": "訓練前掃描檔案或資料夾，查看格式、欄位、文字系統與記錄長度。", + "data.max_help": "彙總時最多讀取的記錄數（100 至 500,000）。", + "data.max_records": "掃描記錄數", + "data.median": "中位數", + "data.p95": "95 百分位", + "data.path_help": "檔案或資料夾（jsonl、json、parquet、csv、tsv、txt），每行一個。多個路徑按合併後訓練看到的樣子作為一個資料集掃描。相對路徑從工作台資料夾起算。", + "data.present": "出現比例", + "data.records": "記錄數", + "data.sampled": "前 {n} 筆記錄", + "data.samples": "範例記錄", + "data.scanned": "已掃描", + "data.schema": "辨識出的格式", + "data.script": "文字", + "data.scripts": "文字系統", + "data.share": "比例", + "data.title": "資料集檢查", + "data.type": "類型", + "err.active_experts": "啟用專家數不能超過專家總數", + "err.dialog_open": "已有一個檔案對話框開啟，請先在系統視窗中完成或取消", + "err.models_dir": "找不到模型資料夾", + "err.native_local_only": "系統檔案對話框只能在服務所在的電腦上使用", + "err.native_unavailable": "這台電腦無法開啟系統檔案對話框", + "err.no_files": "這些路徑下沒有可讀取的檔案", + "err.no_out": "請選擇輸出資料夾", + "err.no_paths": "請至少新增一個路徑", + "err.not_found": "找不到", + "err.paging": "分頁方式請選擇 auto、cluster 或 expert", + "err.port": "連接埠無效", + "err.probe_invalid": "懶載入檢查的設定無效", + "err.remote_needs_key": "遠端綁定需要「允許遠端」與 CHRONOS_API_KEY", + "err.target": "請輸入目標規模", + "err.upload_batch": "上傳批次識別碼無效，請重新整理頁面後再試", + "err.upload_disk_full": "服務所在電腦的磁碟剩餘空間不足，已停止上傳", + "err.upload_failed": "上傳中斷，請檢查網路連線後再試", + "err.upload_kind": "此欄位不支援上傳", + "err.upload_name": "檔名或資料夾層級不允許（隱藏項目、「..」或層級過深）", + "err.upload_too_large": "檔案超過上傳上限（--max-upload-mb）", + "err.upload_type": "沒有可上傳的檔案：類型與此欄位不符，或為不允許的可執行檔", + "event.final": "結束", + "event.lazy_probe": "懶載入探測", + "event.phase_end": "階段結束", + "event.phase_start": "階段開始", + "event.resume": "已恢復", + "event.rollback": "回滾", + "event.validation": "驗證", + "export.done": "已匯出", + "export.format": "格式", + "export.intro": "將檢查點寫成 safetensors 或 GGUF 檔案，並附上 chronos serve 載入所需的專家快取配置。", + "export.note": "Ollama 用戶端透過 chronos serve 執行 Chronos 模型（將 OLLAMA_HOST 設為其位址）；原版 Ollama 無法執行循環 MoE。", + "export.run": "匯出", + "export.size": "大小（MB）", + "export.tensors": "張量數", + "export.title": "匯出", + "field.accumulation_steps.help": "每次最佳化器更新累積的微批次數。有效批次 = 批次大小 × 累積步數。", + "field.accumulation_steps.label": "梯度累積", + "field.activation_checkpointing.help": "反向傳播時重新計算激活值，以速度換取較低的峰值記憶體。", + "field.activation_checkpointing.label": "激活檢查點", + "field.adaptive_halting.help": "每個 token 的停止機率超過門檻後就不再循環，為簡單 token 節省運算。", + "field.adaptive_halting.label": "自適應停止（ACT）", + "field.alpha.help": "蒸餾混合比例：α × 教師 KL + (1 − α) × CE。", + "field.alpha.label": "α（蒸餾）", + "field.batch_size.help": "每個微批次的序列數，受記憶體限制。", + "field.batch_size.label": "批次大小", + "field.beta.help": "偏好訓練中與參考策略 KL 約束的強度。", + "field.beta.label": "β（DPO/GRPO）", + "field.cluster_locality_clusters.help": "區域性約束使用的專家叢集數。0 表示按專家數自動選擇。", + "field.cluster_locality_clusters.label": "叢集數", + "field.cluster_locality_window.help": "叢集約束作用的片段長度（token）。與典型記錄長度一致時效果最好。", + "field.cluster_locality_window.label": "區域性視窗", + "field.coda_layers.help": "循環之後、輸出頭之前只執行一次的層。", + "field.coda_layers.label": "收尾層", + "field.cpu_budget_percent.help": "自動模式可使用的實體核心比例（百分比）。", + "field.cpu_budget_percent.label": "CPU 預算", + "field.cpu_threads.help": "PyTorch/BLAS 執行緒數。自動按實體核心預算設定。", + "field.cpu_threads.label": "CPU 執行緒", + "field.depth_schedule.help": "訓練時每個 batch 的深度選法：固定、自適應，或先固定暖機再自適應。", + "field.depth_schedule.label": "深度排程", + "field.depth_schedule_warmup_fraction.help": "訓練中先以固定全深度執行的比例，之後才切換為自適應。", + "field.depth_schedule_warmup_fraction.label": "排程暖機比例", + "field.draft_exit_depth.help": "作為推測草稿來訓練的循環深度。", + "field.draft_exit_depth.label": "草稿深度（訓練）", + "field.draft_exit_kl_coef.help": "草稿深度出口與全深度輸出之間的額外 KL，用於自推測解碼。0 表示關閉。", + "field.draft_exit_kl_coef.label": "草稿出口 KL", + "field.dropout.help": "訓練時隨機丟棄激活值。資料量大時調低，資料很少時調高。", + "field.dropout.label": "Dropout", + "field.dtype.help": "訓練與推論的數值精度。自動：後端支援時用 bfloat16，否則用 float32。", + "field.dtype.label": "運算精度", + "field.epochs.help": "走訪資料集的次數。設定最大步數時會更早停止。", + "field.epochs.label": "輪數", + "field.export.auto_cluster.help": "以校準資料把常一起使用的專家分到同一頁，延遲載入時讀取的檔案更少。", + "field.export.auto_cluster.label": "建立專家叢集", + "field.export.calibration_data_path.help": "用於統計哪些專家會一起使用的文字，每行一個路徑；多個會自動合併。留空則跳過自動聚類。", + "field.export.calibration_data_path.label": "校準資料", + "field.export.cluster_batch_size.help": "每個校準批次的序列數。", + "field.export.cluster_batch_size.label": "校準批次大小", + "field.export.cluster_device.help": "校準執行使用的裝置。", + "field.export.cluster_device.label": "校準裝置", + "field.export.cluster_max_batches.help": "叢集分析時觀察的批次數。", + "field.export.cluster_max_batches.label": "校準批次數", + "field.export.cluster_max_seq_len.help": "每條校準序列的 token 數。", + "field.export.cluster_max_seq_len.label": "校準長度", + "field.export.config_path.help": "覆寫附屬設定。通常留空。", + "field.export.config_path.label": "設定檔", + "field.export.copy_cache.help": "把延遲載入專家頁一併匯出，部署時無需重建。", + "field.export.copy_cache.label": "複製專家快取", + "field.export.expert_cache_dir.help": "要打包的現有延遲載入專家快取。留空則重新建置。", + "field.export.expert_cache_dir.label": "專家快取", + "field.export.formats.help": "要寫出的容器格式。safetensors 與 GGUF 保留 Chronos 中繼資料，需以 chronos serve 執行。Ollama 用戶端連線到 chronos serve，而非原版 Ollama。", + "field.export.formats.label": "格式", + "field.export.model_path.help": "要匯出的 checkpoint。", + "field.export.model_path.label": "Checkpoint", + "field.export.output_dir.help": "匯出檔案的寫入位置。", + "field.export.output_dir.label": "輸出資料夾", + "field.fallback_mask_prob.help": "訓練時隨機遮蔽專家，讓模型能容忍快取未命中。只影響帶備援的延遲載入推論。", + "field.fallback_mask_prob.label": "缺失模擬", + "field.grad_clip.help": "把全域梯度範數裁剪到此值。0 表示不裁剪。", + "field.grad_clip.label": "梯度裁剪", + "field.halting_threshold.help": "累計停止機率達到此值時 token 停止迭代。越接近 1 跑得越深。", + "field.halting_threshold.label": "停止門檻", + "field.hidden_size.help": "每個 token 向量的寬度。越大能力越強，參數量約按平方成長。必須能被注意力頭數整除。", + "field.hidden_size.label": "隱藏維度", + "field.infer.backend.help": "推論執行環境。自動依 checkpoint 與本機選擇。", + "field.infer.backend.label": "後端", + "field.infer.decode_wait_s.help": "等待預先擷取專家的時間，逾時則按需載入。", + "field.infer.decode_wait_s.label": "解碼等待（秒）", + "field.infer.draft_depth.help": "起草使用的循環深度。", + "field.infer.draft_depth.label": "草稿深度", + "field.infer.gamma.help": "每次驗證前起草的 token 數。", + "field.infer.gamma.label": "草稿長度 γ", + "field.infer.max_tokens.help": "生成 token 數的上限。", + "field.infer.max_tokens.label": "最大新 token 數", + "field.infer.miss_policy.help": "快取未命中時的處理方式。sync_on_demand 的輸出與全記憶體完全一致。", + "field.infer.miss_policy.label": "未命中策略", + "field.infer.mode.help": "比較：延遲載入與全記憶體並排執行；延遲載入：只執行延遲載入專家；全記憶體：全部載入。", + "field.infer.mode.label": "模式", + "field.infer.model_path.help": "要執行的 .pth checkpoint，其附屬設定決定架構。", + "field.infer.model_path.label": "Checkpoint", + "field.infer.paging.help": "自動依 checkpoint：以叢集區域性訓練過則按叢集分頁，否則按專家分頁。", + "field.infer.paging.label": "分頁配置", + "field.infer.ram_load_ratio.help": "保留在記憶體中的專家比例。大於 1.0 時為預先擷取保留餘裕。", + "field.infer.ram_load_ratio.label": "記憶體比例", + "field.infer.raw_prompt.help": "依原樣送出文字，不套用對話範本。", + "field.infer.raw_prompt.label": "原始提示", + "field.infer.residency_decay.help": "規劃器遺忘舊路由證據的速度。", + "field.infer.residency_decay.label": "規劃器衰減", + "field.infer.residency_planner.help": "淘汰未來最晚才會用到的專家，而非最久未用的專家。", + "field.infer.residency_planner.label": "下次使用規劃器", + "field.infer.speculative.help": "以淺深度起草 token，再以全深度驗證，輸出不變。50M 模型實測較慢，因此預設關閉。", + "field.infer.speculative.label": "自推測解碼", + "field.infer.sweep_ram_ratios.help": "依序執行各標準比例，繪製速度、命中率與記憶體的關係。", + "field.infer.sweep_ram_ratios.label": "掃描記憶體比例", + "field.infer.temperature.help": "取樣隨機度。0 為貪婪解碼，結果可重現。", + "field.infer.temperature.label": "溫度", + "field.infer.verify_wait_s.help": "推測驗證階段的同類等待時間。", + "field.infer.verify_wait_s.label": "驗證等待（秒）", + "field.kv_cache_dtype.help": "自動：沿用運算精度；int8：推論 KV 快取按 token 量化（記憶體約減半，誤差很小）。訓練從不量化。", + "field.kv_cache_dtype.label": "KV 快取精度", + "field.kv_latent_dim.help": "MLA 層快取的壓縮鍵值潛在向量寬度。越小越省 KV 記憶體。", + "field.kv_latent_dim.label": "MLA 潛在維度", + "field.kv_quant_group_size.help": "共用一個 int8 縮放係數的通道數。分組越小越精確，體積略大。", + "field.kv_quant_group_size.label": "KV 量化分組", + "field.lambda_balance.help": "促使路由器把 token 平均分配到各專家。", + "field.lambda_balance.label": "負載平衡", + "field.lambda_cluster_locality.help": "訓練同一片段內的 token 留在同一個專家叢集，延遲載入推論就能整叢分頁。0 表示關閉。", + "field.lambda_cluster_locality.label": "叢集區域性", + "field.lambda_lookahead.help": "訓練前瞻路由器預測未來 token 會用到的專家。", + "field.lambda_lookahead.label": "前瞻損失", + "field.lambda_lookahead_topk.help": "額外加權，要求準確預測未來的 top-k 專家。", + "field.lambda_lookahead_topk.label": "前瞻 top-k", + "field.lambda_lookahead_union.help": "要求涵蓋前瞻視窗內專家聯集的權重。", + "field.lambda_lookahead_union.label": "前瞻聯集", + "field.lambda_or.help": "ORPO 中勝算比項的權重。", + "field.lambda_or.label": "λ（ORPO）", + "field.lambda_router_anchor.help": "微調時讓路由保持接近上游 checkpoint。預訓練設為 0。", + "field.lambda_router_anchor.label": "路由錨定", + "field.lambda_router_locality.help": "對相鄰 token 之間路由變化的輕微懲罰。", + "field.lambda_router_locality.label": "路由區域性", + "field.lambda_temporal.help": "獎勵相鄰 token 選擇相同專家，一段文字需要載入的專家更少。", + "field.lambda_temporal.label": "時間區域性", + "field.layer_aware_cache.help": "按虛擬層而非實體層追蹤常駐專家。", + "field.layer_aware_cache.label": "按層快取", + "field.learning_rate.help": "最佳化器的峰值步長。規劃器會按模型規模給出建議；微調時調低。", + "field.learning_rate.label": "學習率", + "field.log_interval.help": "每 N 步寫一行指標。", + "field.log_interval.label": "日誌間隔", + "field.lookahead_steps.help": "前瞻路由為未來多少個 token 預測專家，以便在用到之前預先擷取。", + "field.lookahead_steps.label": "前瞻步數", + "field.max_gen_len.help": "每次 GRPO 取樣生成的 token 數。", + "field.max_gen_len.label": "取樣長度", + "field.max_inference_depth.help": "推論時的最大迭代次數。通常等於訓練深度；超過訓練深度未經驗證。", + "field.max_inference_depth.label": "推論深度", + "field.max_position_embeddings.help": "旋轉位置編碼預備的最長序列。只有超長上下文才需要修改。", + "field.max_position_embeddings.label": "最大位置數", + "field.max_resident_experts.help": "每層記憶體中保留專家數的硬上限。留空由規劃器決定。", + "field.max_resident_experts.label": "最大常駐專家", + "field.max_seq_len.help": "每條訓練序列的 token 數。在混合注意力下記憶體約線性成長。", + "field.max_seq_len.label": "序列長度", + "field.max_steps.help": "達到此最佳化器步數後停止。0 表示不設上限（跑完所有輪數）。", + "field.max_steps.label": "最大步數", + "field.min_recurrent_depth.help": "深度排程或提前退出允許的最低深度。設為 1 時簡單 token 一輪即可退出。", + "field.min_recurrent_depth.label": "最小深度", + "field.mla_cache_rope.help": "額外快取旋轉鍵部分，以 KV 記憶體換取解碼速度。", + "field.mla_cache_rope.label": "快取 RoPE 鍵", + "field.mla_decode.help": "展開：參考路徑。吸收：把投影折入查詢（每步運算更少）。自動會混用兩條路徑，實測有微小 logit 偏差，因此需手動開啟。", + "field.mla_decode.label": "MLA 解碼路徑", + "field.modalities.help": "額外的輸入編碼器，以逗號分隔（image、audio）。純文字模型留空，資料集會自動偵測。", + "field.modalities.label": "模態", + "field.moe_intermediate_size.help": "每個專家前饋網路的隱藏寬度。0 表示按隱藏維度推導。", + "field.moe_intermediate_size.label": "專家 FFN 寬度", + "field.num_attention_heads.help": "查詢頭數。隱藏維度必須能被它整除。", + "field.num_attention_heads.label": "注意力頭數", + "field.num_experts.help": "每個 MoE 層的路由專家數。增加專家能擴充容量而不增加每 token 運算量；專家也是延遲載入的單位。", + "field.num_experts.label": "專家數", + "field.num_experts_per_tok.help": "每個 token 使用的專家數。top-2 讓路由器能從損失中學習；top-1 需關閉 norm_topk_prob。", + "field.num_experts_per_tok.label": "啟用專家數（top-k）", + "field.num_generations.help": "GRPO 對每個提示比較的樣本數。", + "field.num_generations.label": "每題取樣數", + "field.num_hidden_layers.help": "checkpoint 中實際儲存的層數（前置層 + 循環組 + 收尾層）。循環會重複使用這些層，所以虛擬深度可以更大。", + "field.num_hidden_layers.label": "實體層數", + "field.num_key_value_heads.help": "滑動視窗層分組查詢注意力的鍵值頭數。越少 KV 快取越小。", + "field.num_key_value_heads.label": "KV 頭數", + "field.num_shared_experts.help": "每個 token 都會經過的常開專家，始終常駐記憶體。", + "field.num_shared_experts.label": "共用專家", + "field.num_workers.help": "資料載入器的工作程序數。自動按平台選擇安全的數量。", + "field.num_workers.label": "資料載入程序", + "field.optimizer_page_dir.help": "分頁最佳化器狀態的存放資料夾。請使用本機高速 SSD。", + "field.optimizer_page_dir.label": "最佳化器頁面資料夾", + "field.optimizer_paging.help": "不活躍專家的最佳化器狀態存到磁碟。", + "field.optimizer_paging.label": "分頁最佳化器", + "field.page_size_mb.help": "單一最佳化器頁面檔的大小，單位 MB。", + "field.page_size_mb.label": "頁面大小", + "field.pinned_memory_max_fraction.help": "用於鎖頁專家緩衝區的系統記憶體最大比例。", + "field.pinned_memory_max_fraction.label": "鎖頁記憶體占比", + "field.plan.active_mode.help": "以啟用專家數或啟用參數量設定每 token 預算。", + "field.plan.active_mode.label": "啟用預算方式", + "field.plan.active_value.help": "每 token 啟用專家數，或啟用參數量（百萬）。", + "field.plan.active_value.label": "啟用預算", + "field.plan.data_paths.help": "規劃器量測這些檔案或資料夾，以決定 token 數、序列長度與視窗，每行一個。", + "field.plan.data_paths.label": "資料集", + "field.plan.depth.help": "每個循環組的迭代次數。", + "field.plan.depth.label": "循環深度", + "field.plan.memory_saving.help": "以速度換取記憶體的力道（int8 KV、共用路由、檢查點）。", + "field.plan.memory_saving.label": "省記憶體程度", + "field.plan.num_experts.help": "每個 MoE 層的路由專家數。", + "field.plan.num_experts.label": "專家數", + "field.plan.ram_gb.help": "規劃時假定的記憶體。留空則偵測本機。", + "field.plan.ram_gb.label": "記憶體（GB）", + "field.plan.seq_len.help": "留空則從資料集分析結果推導。", + "field.plan.seq_len.label": "序列長度", + "field.plan.target_m.help": "總參數量，單位百萬。", + "field.plan.target_m.label": "目標規模（M）", + "field.plan.token_budget_b.help": "預訓練 token 數，單位十億。留空：走訪一遍資料。", + "field.plan.token_budget_b.label": "Token 預算（B）", + "field.plan.use_moe.help": "規劃 MoE 模型；關閉則規劃稠密模型。", + "field.plan.use_moe.label": "混合專家", + "field.plan.vocab_size.help": "留空則從分詞器報告讀取，或使用 16384。", + "field.plan.vocab_size.label": "詞彙表", + "field.ponder_loss_coef.help": "對期望迭代次數的懲罰。越高越促使模型提早停止。", + "field.ponder_loss_coef.label": "思考損失", + "field.prefetch_depth.help": "預先擷取器提前載入預測專家的距離。0 表示不預先擷取。", + "field.prefetch_depth.label": "預先擷取深度", + "field.prelude_layers.help": "循環之前只執行一次的層。前瞻路由讀取它們的輸出。", + "field.prelude_layers.label": "前置層", + "field.probe.backends.help": "要測試的執行環境。每個後端各有一份全記憶體參考。", + "field.probe.backends.label": "探針後端", + "field.probe.decode_wait_s.help": "解碼步驟等待預先擷取專家的時間，逾時則按需載入。", + "field.probe.decode_wait_s.label": "解碼等待（秒）", + "field.probe.draft_depth.help": "推測變體中起草 token 使用的循環深度。", + "field.probe.draft_depth.label": "草稿深度", + "field.probe.enabled.help": "每次驗證後，以延遲載入專家解碼剛儲存的 checkpoint，並與全記憶體比較：速度、命中率與 token 一致性。探針在低優先序子程序中對 checkpoint 副本執行。", + "field.probe.enabled.label": "延遲載入探針", + "field.probe.gamma.help": "推測變體中每次驗證前起草的 token 數。", + "field.probe.gamma.label": "草稿長度 γ", + "field.probe.max_new_tokens.help": "每條提示貪婪解碼的 token 數。越多速率越穩定，探針耗時也越長。", + "field.probe.max_new_tokens.label": "每條提示 token 數", + "field.probe.max_variants.help": "測試組合數的上限，超出部分標記為已截斷。", + "field.probe.max_variants.label": "組合上限", + "field.probe.miss_policies.help": "Torch 快取未命中時的行為。sync_on_demand 載入並等待（輸出完全一致）；on_demand 可能略過遲到的專家；fallback_diagnostic 從不載入（最差情況）。", + "field.probe.miss_policies.label": "未命中策略", + "field.probe.paging.help": "自動依模型的訓練方式（啟用叢集區域性損失時按叢集分頁）。加入按專家可比較逐專家分頁。", + "field.probe.paging.label": "分頁配置", + "field.probe.ram_ratios.help": "同時允許常駐記憶體的專家占比。越低越省記憶體，未命中越多。", + "field.probe.ram_ratios.label": "記憶體比例", + "field.probe.residency.help": "MLX 淘汰策略：下次使用規劃器，或一般的最近使用視窗。", + "field.probe.residency.label": "常駐策略", + "field.probe.residency_decay.help": "下次使用規劃器遺忘舊路由證據的速度。", + "field.probe.residency_decay.label": "規劃器衰減", + "field.probe.speculative.help": "同時測試搭配延遲載入專家的自推測解碼。", + "field.probe.speculative.label": "推測變體", + "field.probe.verify_wait_s.help": "推測驗證階段的同類等待時間。", + "field.probe.verify_wait_s.label": "驗證等待（秒）", + "field.ram_budget_mb.help": "延遲載入專家快取可用的記憶體，單位 MB。留空則改用記憶體比例。", + "field.ram_budget_mb.label": "記憶體預算", + "field.recurrent_attention_mode.help": "循環內注意力的執行方式。shared_first：每輪都做注意力，K/V 取第 0 輪（快取精確，KV 不隨深度成長）。per_iteration：每輪一份 KV（D 倍 KV）。first_only：v2 行為。", + "field.recurrent_attention_mode.label": "循環注意力", + "field.recurrent_depth.help": "訓練時每個循環組的迭代次數。虛擬層數 = 前置 + 組數 × 深度 + 收尾。", + "field.recurrent_depth.label": "訓練深度", + "field.recurrent_depth_jitter.help": "在排程深度附近隨機波動，讓模型學會在多個深度退出。", + "field.recurrent_depth_jitter.label": "深度抖動", + "field.recurrent_distill_coef.help": "讓淺層出口學習全深度的預測，提高推測接受率與提前退出品質。", + "field.recurrent_distill_coef.label": "深度蒸餾", + "field.recurrent_enabled.help": "讓中間的共用層組對每個 token 執行多次，用少量權重得到更深網路的效果。", + "field.recurrent_enabled.label": "循環深度", + "field.recurrent_groups.help": "前置層與收尾層之間組成循環組的層數。留空則按 層數 − 前置 − 收尾 推導。", + "field.recurrent_groups.label": "循環組數", + "field.recurrent_stability_coef.help": "保持循環更新收縮，避免多次迭代後漂移或發散。", + "field.recurrent_stability_coef.label": "穩定性損失", + "field.recurrent_strategy.help": "訓練時循環深度分桶的取樣策略。", + "field.recurrent_strategy.label": "深度策略", + "field.reward_spec.help": "toy 使用內建獎勵，或 lm:/path/to/reward-model[,device=cuda]。", + "field.reward_spec.label": "獎勵（GRPO）", + "field.rope_dim.help": "MLA 層每個頭中攜帶旋轉位置資訊的維度數。", + "field.rope_dim.label": "RoPE 維度", + "field.rope_theta.help": "旋轉位置編碼的基礎頻率。值越大，長上下文的位置解析度越好。", + "field.rope_theta.label": "RoPE 基數", + "field.router_scope.help": "per_iteration：每輪循環重新路由。共用：沿用第 0 輪的路由，用到的專家更少，更適合延遲載入。", + "field.router_scope.label": "路由範圍", + "field.save_dir.help": "checkpoint 及其 .config.json 附屬檔的寫入位置。", + "field.save_dir.label": "Checkpoint 資料夾", + "field.save_interval.help": "每 N 步儲存（並驗證）一次。", + "field.save_interval.label": "儲存間隔", + "field.serve.allow_remote.help": "監聽非迴路位址。伺服器環境中必須設定 CHRONOS_API_KEY，用戶端須附上此金鑰。", + "field.serve.allow_remote.label": "允許遠端存取", + "field.serve.device.help": "服務使用的執行環境。", + "field.serve.device.label": "裝置", + "field.serve.dtype.help": "服務使用的精度。", + "field.serve.dtype.label": "精度", + "field.serve.host.help": "監聽的位址。127.0.0.1 只允許本機存取。", + "field.serve.host.label": "監聽位址", + "field.serve.kv_cache_dtype.help": "覆寫 KV 儲存方式。int8 約減半 KV 記憶體；留空保留模型設定。", + "field.serve.kv_cache_dtype.label": "KV 快取精度", + "field.serve.lazy_experts.help": "專家存放在 SSD 上並按需換入，透過 Ollama 協定呼叫時記憶體同樣維持較低。", + "field.serve.lazy_experts.label": "延遲載入專家", + "field.serve.max_resident_experts.help": "每層記憶體專家數的硬上限。留空：依記憶體占比推導。", + "field.serve.max_resident_experts.label": "最大常駐專家", + "field.serve.mla_cache_rope.help": "快取旋轉後的鍵：記憶體較多，重新計算較少。", + "field.serve.mla_cache_rope.label": "快取 RoPE 鍵", + "field.serve.mla_decode.help": "覆寫 MLA 解碼路徑。留空保留模型設定。", + "field.serve.mla_decode.label": "MLA 解碼", + "field.serve.models_dir.help": "存放匯出模型的資料夾，每個模型顯示為一個 Ollama 模型名稱。", + "field.serve.models_dir.label": "模型資料夾", + "field.serve.paging.help": "自動依 checkpoint 的訓練方式（按叢集或按專家分頁）。", + "field.serve.paging.label": "分頁配置", + "field.serve.port.help": "TCP 連接埠。11435 可避免與原版 Ollama 的 11434 衝突。", + "field.serve.port.label": "連接埠", + "field.serve.ram_budget_mb.help": "專家記憶體的硬上限。留空：不設上限。", + "field.serve.ram_budget_mb.label": "記憶體預算（MB）", + "field.serve.ram_expert_fraction.help": "允許常駐記憶體的專家比例。", + "field.serve.ram_expert_fraction.label": "記憶體專家占比", + "field.serve.router_scope.help": "覆寫跨迭代的路由沿用方式。留空保留模型設定。", + "field.serve.router_scope.label": "路由範圍", + "field.sliding_window_size.help": "每個滑動視窗層能看到的 token 數。短視窗適合句子級任務，並限制 KV 記憶體。", + "field.sliding_window_size.label": "滑動視窗", + "field.storage_format.help": "磁碟專家儲存的檔案格式。", + "field.storage_format.label": "專家儲存格式", + "field.temperature.help": "GRPO 取樣的溫度。", + "field.temperature.label": "取樣溫度", + "field.tie_word_embeddings.help": "輸入嵌入與輸出頭共用權重。可省下 詞彙表 × 隱藏維度 個參數，小模型幾乎不損失品質。", + "field.tie_word_embeddings.label": "共用嵌入", + "field.tok.fields.help": "要讀取的 JSON 欄位。留空自動辨識 text、conversations 與 pairs。", + "field.tok.fields.label": "文字欄位", + "field.tok.out_dir.help": "tokenizer.json、對話範本與報告的寫入位置。", + "field.tok.out_dir.label": "輸出資料夾", + "field.tok.paths.help": "取樣文字的檔案或資料夾，每行一個。每個來源按各自配額取樣，不做合併。", + "field.tok.paths.label": "來源檔案", + "field.tok.per_source.help": "每個來源檔案的蓄水池取樣數量。", + "field.tok.per_source.label": "每個來源文字數", + "field.tok.temperature.help": "依位元組占比平衡各語言。越低越平均，避免英文擠壓中文。", + "field.tok.temperature.label": "語言溫度", + "field.tok.vocab_size.help": "要學習的 token 數。模型詞彙表需設為相同值。", + "field.tok.vocab_size.label": "詞彙表大小", + "field.tokenizer_path.help": "模型訓練所用的分詞器資料夾（含 tokenizer.json）。留空時使用 checkpoint 內建的分詞器，否則用內建備援分詞器。詞彙表大小必須一致。", + "field.tokenizer_path.label": "分詞器", + "field.train.backend.help": "訓練使用的裝置。自動：Apple 晶片優先 MLX，其次 CUDA、XPU、MPS、CPU。", + "field.train.backend.label": "訓練後端", + "field.train.data_distill.help": "由教師模型評分的文字或對話資料。檔案或資料夾，每行一個；多個會自動合併。留空則沿用第一階段的資料集。", + "field.train.data_distill.label": "蒸餾資料", + "field.train.data_dpo.help": "偏好對資料，含 chosen 與 rejected 回覆。檔案或資料夾，每行一個；多個會自動合併。留空則沿用第一階段的資料集。", + "field.train.data_dpo.label": "DPO資料", + "field.train.data_grpo.help": "用於獎勵 rollout 的 prompt 資料。檔案或資料夾，每行一個；多個會自動合併。留空則沿用第一階段的資料集。", + "field.train.data_grpo.label": "GRPO資料", + "field.train.data_orpo.help": "偏好對資料，含 chosen 與 rejected 回覆，格式同 DPO。檔案或資料夾，每行一個；多個會自動合併。留空則沿用第一階段的資料集。", + "field.train.data_orpo.label": "ORPO資料", + "field.train.data_path.help": "第一階段的資料集：JSONL、JSON、Parquet、CSV、TSV 或 TXT 檔案，或包含這些檔案的資料夾，每行一個。多個資料集會在執行目錄中合併為一個 JSONL 檔案。留空則執行合成冒煙測試。", + "field.train.data_path.label": "資料集", + "field.train.data_sft.help": "對話資料，含 conversations（user / assistant 輪次）。檔案或資料夾，每行一個；多個會自動合併。留空則沿用第一階段的資料集。", + "field.train.data_sft.label": "監督微調資料", + "field.train.init_weight.help": "起始 checkpoint。留空：預訓練隨機初始化，後續階段接上一階段。", + "field.train.init_weight.label": "起始權重", + "field.train.max_steps.help": "達到此最佳化器步數後停止。0 表示跑完設定的輪數。", + "field.train.max_steps.label": "最大步數", + "field.train.sample_prompt.help": "訓練過程中定期續寫的提示詞，方便查看進度。留空使用階段預設值。", + "field.train.sample_prompt.label": "取樣提示詞", + "field.train.stages.help": "依序執行，每個階段從上一階段的 checkpoint 開始。常用路徑是 預訓練 → SFT。", + "field.train.stages.label": "階段", + "field.train.teacher_path.help": "僅用於蒸餾：學生模型要擬合其輸出的模型。", + "field.train.teacher_path.label": "教師 checkpoint", + "field.train.val_ratio.help": "保留作驗證的記錄比例。每次驗證也會觸發延遲載入探針。", + "field.train.val_ratio.label": "驗證集比例", + "field.training_paging.help": "訓練時把不活躍的專家移出記憶體。較慢，但能訓練超過記憶體容量的模型。", + "field.training_paging.label": "分頁訓練", + "field.tune.data_path.help": "自動調參啟動前會檢查這些資料集，每行一個。目前的試訓使用隨機 token，所以這些資料不影響試訓損失。", + "field.tune.data_path.label": "資料集", + "field.tune.n_trials.help": "Optuna 嘗試的設定數量。", + "field.tune.n_trials.label": "試驗次數", + "field.tune.probe_steps.help": "每次試驗的訓練步數；表現差的試驗會被提前剪枝。", + "field.tune.probe_steps.label": "每次試驗步數", + "field.tune.search.help": "搜尋可以調整的參數。", + "field.tune.search.label": "搜尋參數", + "field.use_hybrid_attention.help": "MLA（壓縮潛在 KV）與滑動視窗層交替，長上下文時 KV 快取仍然很小。", + "field.use_hybrid_attention.label": "混合注意力", + "field.vocab_size.help": "token id 的數量，必須等於分詞器詞彙表大小，否則 id 會超出嵌入表。", + "field.vocab_size.label": "詞彙表大小", + "field.vram_budget_gb.help": "專家快取可用的加速器記憶體，單位 GB。", + "field.vram_budget_gb.label": "顯示記憶體預算", + "field.weight_decay.help": "AdamW 解耦權重衰減。", + "field.weight_decay.label": "權重衰減", + "group.attention": "注意力與 KV 快取", + "group.autotune": "自動調參", + "group.benchmark": "基準測試", + "group.export": "匯出", + "group.inference": "生成", + "group.lazy": "延遲載入專家", + "group.memory": "記憶體與延遲載入專家", + "group.model": "模型形狀", + "group.moe": "混合專家", + "group.objectives": "訓練目標", + "group.optimizer": "最佳化器", + "group.planner": "自動規劃", + "group.probe": "驗證時的延遲載入探針", + "group.recurrent": "循環深度", + "group.rl": "偏好與強化學習階段", + "group.runtime": "執行環境", + "group.serve": "Ollama / OpenAI 服務", + "group.speculative": "推測解碼", + "group.stage_data": "後續階段的資料", + "group.tokenizer": "分詞器建置", + "group.train": "訓練任務", + "infer.arch_note": "架構讀取自檢查點的設定檔；「模型」頁只補上檔案中缺少的項目。", + "infer.col.cache_hit_rate": "快取命中率", + "infer.col.cold_miss_predict_hit_rate": "冷未命中已預測", + "infer.col.decode_time_s": "解碼（秒）", + "infer.col.fallback_weight_rate": "回退權重", + "infer.col.prediction_hit_rate": "預測命中率", + "infer.col.prefill_time_s": "預填充（秒）", + "infer.col.response_time_s": "回應（秒）", + "infer.col.rss_after_decode_gb": "解碼後 RSS（GB）", + "infer.col.rss_delta_gb": "RSS 變化（GB）", + "infer.col.spec_acceptance_rate": "草稿接受率", + "infer.col.tokens": "token 數", + "infer.col.tokens_per_sec": "token/秒", + "infer.common_prefix": "前 {n} 個字元相同", + "infer.diff": "輸出不同", + "infer.full_output": "全記憶體", + "infer.generate": "生成", + "infer.hit_rates": "命中率", + "infer.intro": "同一個提示分別以懶載入專家與完整記憶體執行，比較輸出、速度、命中率與記憶體。", + "infer.lazy_output": "延遲載入專家", + "infer.measured_note": "在本機針對此提示實測，單次執行，並非平均值。", + "infer.mode": "模式", + "infer.mode.full_dram": "完整記憶體", + "infer.mode.lazy_offload": "懶載入專家", + "infer.output": "輸出", + "infer.prompt": "提示詞", + "infer.prompt_help": "傳送給模型的文字。未開啟「原始提示」時會套用 SFT 對話範本。", + "infer.ram_ratio": "記憶體載入比例", + "infer.same": "輸出完全相同", + "infer.similarity": "相似度", + "infer.stats": "執行統計", + "infer.sweep": "記憶體比例掃描", + "infer.title": "推論", + "io.empty": "尚無指標。在本程序中執行推論或服務時才會出現。", + "io.intro": "本服務程序中推論與服務記錄的專家快取與 IO 計數。", + "io.latest": "最新值", + "io.metric": "指標", + "io.samples": "樣本數", + "io.seconds": "秒", + "io.title": "IO 與快取指標", + "label.estimated": "估算", + "label.measured": "實測", + "label.measured_meta": "實測（meta 裝置計數）", + "metric.active_mem_gb": "活躍", + "metric.balance_loss": "專家均衡", + "metric.bleu_zh": "BLEU（中文）", + "metric.chrf": "chrF", + "metric.error": "錯誤", + "metric.eval_s": "評測用時（秒）", + "metric.flores_devtest_n": "FLORES 句數", + "metric.grad_norm": "梯度範數", + "metric.load_cv": "負載不均（變異係數）", + "metric.locality_loss": "叢集局部性", + "metric.lookahead_loss": "前瞻路由", + "metric.loss": "總損失", + "metric.lr": "學習率", + "metric.peak_mem_gb": "峰值", + "metric.primary_loss": "交叉熵", + "metric.rec_contraction": "收縮率", + "metric.rec_depth_mean": "平均深度", + "metric.rec_early_exit_ratio": "提前退出比例", + "metric.route_entropy": "路由熵", + "metric.state": "狀態", + "metric.temporal_loss": "時間局部性", + "metric.tokens_per_s": "token/秒", + "metric.valid_ce": "驗證交叉熵", + "metric.valid_ppl": "驗證困惑度", + "mon.best": "最佳", + "mon.chart.aux": "輔助損失", + "mon.chart.loss": "訓練損失", + "mon.chart.lr": "學習率", + "mon.chart.memory": "記憶體（GB）", + "mon.chart.optimizer": "梯度範數", + "mon.chart.recurrent": "循環深度", + "mon.chart.routing": "專家路由", + "mon.chart.throughput": "每秒 token 數", + "mon.checkpoints": "Checkpoint", + "mon.eta": "剩餘時間", + "mon.event": "事件", + "mon.events": "事件", + "mon.grad_norm": "梯度範數", + "mon.input": "輸入", + "mon.intro": "即時檢視任一執行資料夾，包括從命令列啟動的長時間訓練。唯讀。", + "mon.log": "訓練日誌", + "mon.loss": "損失", + "mon.lr": "學習率", + "mon.memory": "記憶體", + "mon.no_runs": "尚無執行記錄。請在「訓練」頁開始一次。", + "mon.phase_note": "每個階段的步數重新計數，因此圖表一次只顯示一個階段。", + "mon.progress": "進度", + "mon.reference": "參考譯文", + "mon.run": "執行", + "mon.samples": "取樣輸出", + "mon.size": "大小", + "mon.stale": "一段時間沒有更新", + "mon.throughput": "每秒 token 數", + "mon.title": "執行監控", + "mon.tokens_seen": "已訓練 token", + "mon.updated": "{s} 秒前更新", + "mon.valid_ce": "驗證 CE", + "mon.validation": "驗證", + "nav.autotune": "自動調參", + "nav.benchmark": "基準測試", + "nav.config": "模型", + "nav.data": "資料", + "nav.export": "匯出", + "nav.inference": "推論", + "nav.io": "IO 指標", + "nav.monitor": "監控", + "nav.serve": "服務", + "nav.tokenizer": "分詞器", + "nav.train": "訓練", + "phase.active": "進行中", + "phase.complete": "已完成", + "phase.pending": "等待中", + "picker.add": "新增所選", + "picker.add_button": "新增…", + "picker.browse_tip": "在頁面內瀏覽服務所在電腦上的資料夾並選取路徑。", + "picker.empty": "這裡沒有符合類型的檔案或資料夾", + "picker.entries": "資料夾內容", + "picker.go": "前往", + "picker.hidden": "顯示隱藏項目", + "picker.multi_hint": "可以新增多個檔案或資料夾，訓練前會自動合併。", + "picker.n_selected": "已選 {n} 項", + "picker.native_fallback": "這台電腦無法開啟系統檔案對話框，已改為上傳方式。", + "picker.native_tip": "開啟系統檔案管理員。檔案不會被複製，直接使用所選的原路徑。", + "picker.pick_file": "請選擇一個檔案", + "picker.place.desktop": "桌面", + "picker.place.documents": "文件", + "picker.place.downloads": "下載", + "picker.place.home": "主目錄", + "picker.place.root": "根目錄（/）", + "picker.place.volumes": "所有磁碟", + "picker.placeholder_multi": "每行一個路徑，或點「新增…」", + "picker.places": "常用位置", + "picker.select": "選擇", + "picker.this_folder": "未選擇項目：將使用目前資料夾（{path}）", + "picker.title.any": "選擇檔案或資料夾", + "picker.title.checkpoint": "選擇 checkpoint", + "picker.title.config": "選擇設定檔", + "picker.title.data": "選擇資料集（檔案或資料夾）", + "picker.title.dir": "選擇資料夾", + "picker.truncated": "只顯示前 {n} 項；輸入路徑可直接進入更深的目錄", + "picker.up": "上一層資料夾", + "picker.upload_files": "上傳檔案…", + "picker.upload_folder": "上傳資料夾…", + "picker.upload_tip": "用本機的檔案管理員選擇檔案，上傳到服務所在電腦的 uploads 資料夾，並填入上傳後的路徑。", + "picker.uploaded": "已上傳 {n} 個檔案。", + "picker.uploaded_skipped": "已上傳 {n} 個檔案；{skipped} 個類型不符或為隱藏檔案，已略過。", + "picker.uploading": "正在上傳 {done} / {total}…", + "picker.will_merge": "{n} 個路徑將在使用前合併為一個 JSONL 檔案。", + "plan.apply": "使用此方案", + "plan.note": "規劃器會依目標規模、資料與記憶體搜尋設定；訓練前請先檢查。", + "plan.run": "規劃", + "plan.stages": "各階段設定", + "plan.summary": "方案", + "plan.title": "自動規劃", + "plan.tokens": "訓練 token 數", + "preset.b1": "1B（64 GB）", + "preset.b1.summary": "1B MoE，64 個細粒度專家，積極節省記憶體。", + "preset.enzh50m": "50M 英譯中配方", + "preset.enzh50m.summary": "72 小時英譯中配方：16k Chronos 分詞器，shards_v1。", + "preset.m200": "200M（32 GB）", + "preset.m200.summary": "通用 200M MoE，均衡節省記憶體，使用 minimind 6400 詞表。", + "preset.m500": "500M（32 GB）", + "preset.m500.summary": "500M MoE，積極節省記憶體：int8 KV、共享路由、激活檢查點。", + "preset.tiny": "Tiny（冒煙測試 / CI）", + "preset.tiny.summary": "1.5M 冒煙測試，CPU 上幾分鐘即可訓練完成。", + "probe.accept": "草稿接受率", + "probe.agree": "逐 token 一致率", + "probe.auto_paging": "自動分頁", + "probe.cold_hit": "冷未命中被預測", + "probe.conditions": "貪婪解碼，{n} 條提示，每條 {t} 個 token；實測。", + "probe.exact": "完全一致", + "probe.history": "命中率與一致率隨訓練的變化", + "probe.hit": "常駐命中率", + "probe.loads": "按需載入次數", + "probe.none": "尚無探針結果。每次驗證後會執行一次。", + "probe.pages": "分頁 · 頁數", + "probe.pred_hit": "預測命中率", + "probe.prefix": "最短一致前綴", + "probe.reference": "全記憶體參考", + "probe.rss": "峰值 RSS（GB）", + "probe.skip.busy": "上一次探測仍在執行", + "probe.skip.dense": "稠密模型，沒有可分頁的路由專家", + "probe.skip.no_checkpoint": "磁碟上還沒有檢查點", + "probe.skip.superseded": "已被第 {step} 步的新驗證取代", + "probe.skip.writing": "檢查點仍在寫入", + "probe.skipped": "已略過：{reason}", + "probe.title": "延遲載入專家探針", + "probe.tps": "Token/秒", + "probe.trained_clusters": "訓練時啟用叢集局部性", + "probe.variant": "組合", + "probe.vs_full": "相對全記憶體速度", + "probe.wait": "預先擷取等待（秒）", + "rule.clusters_auto": "0 = 依專家數推導", + "rule.equals_recurrent_depth": "等於訓練深度", + "rule.loopback": "127.0.0.1（僅本機）", + "rule.sync_is_exact": "sync_on_demand：輸出與全記憶體完全一致", + "rule.vocab_matches_tokenizer": "等於分詞器詞彙表大小", + "rule.window_median_record": "最接近記錄長度中位數的 2 的冪次（未知時為 32）", + "serve.auth_off": "未設金鑰：僅本機可存取。", + "serve.auth_on": "用戶端必須附帶 API 金鑰。", + "serve.hint": "連線方式：OLLAMA_HOST={addr} ollama run ", + "serve.intro": "啟動 chronos serve：Ollama 與 OpenAI 用戶端連線到它，專家維持懶載入。", + "serve.remote_warning": "綁定到本機以外的位址需要勾選「允許遠端」，並在 CHRONOS_API_KEY 中設定金鑰；沒有金鑰時服務拒絕啟動。", + "serve.running": "服務位址：{url}", + "serve.start": "啟動服務", + "serve.stop": "停止服務", + "serve.stopped": "服務已停止", + "serve.title": "Ollama / OpenAI 相容服務", + "source.code": "預設值", + "source.measured": "實測", + "source.planner": "自適應規劃器", + "source.rule": "規則", + "stage.distill": "蒸餾", + "stage.dpo": "DPO", + "stage.grpo": "GRPO", + "stage.orpo": "ORPO", + "stage.pretrain": "預訓練", + "stage.sft": "監督微調", + "state.checkpointing": "正在儲存 checkpoint", + "state.done": "完成", + "state.error": "錯誤", + "state.evaluating": "評估中", + "state.failed": "失敗", + "state.finished": "已完成", + "state.idle": "閒置", + "state.loading": "載入中", + "state.paused": "已暫停", + "state.pending": "等待中", + "state.running": "執行中", + "state.stale": "近期無更新", + "state.starting": "啟動中", + "state.stopped": "已停止", + "state.stopping": "正在停止", + "tip.auto_empty": "留空 = 自動", + "tip.default": "預設值", + "tip.range": "範圍", + "tip.recommended": "建議值", + "tip.source": "來源", + "tok.algorithm": "演算法", + "tok.bucket": "分桶", + "tok.byte_fallback": "位元組備援率", + "tok.chars_per_token": "每 token 字元數", + "tok.compression": "每 token 位元組數", + "tok.elapsed": "建置用時", + "tok.fertility": "每詞 token 數", + "tok.fingerprint": "指紋", + "tok.intro": "以你自己的資料訓練 Chronos 分詞器：抽樣、正規化、位元組回退 BPE，並產生保留集評估報告。", + "tok.out_dir_now": "輸出資料夾", + "tok.progress": "建置進度", + "tok.report": "評估報告", + "tok.report_note": "在未參與分詞器訓練的保留文本上實測。", + "tok.roundtrip": "無損往返", + "tok.texts": "文本數", + "tok.title": "Chronos 分詞器建置", + "tok.vocab_use": "詞彙表使用率", + "train.config_note": "架構來自「模型」頁；套用過自動規劃時，各階段學習率來自規劃結果。", + "train.intro": "使用「模型」頁的設定啟動訓練工作，進度顯示在監控頁。", + "train.jobs": "訓練任務", + "train.open_monitor": "在監控中開啟", + "train.probe_note": "每次驗證後，用延遲載入專家和全記憶體分別解碼剛儲存的 checkpoint。探針以低優先權子程序在副本上執行，不中斷訓練。", + "train.probe_off": "此工作已關閉懶載入檢查。", + "train.probe_title": "每次驗證時檢查延遲載入專家", + "train.run_dir": "執行資料夾", + "train.stage_data_note": "第一階段之後的每個階段讀取各自格式的資料。某一列留空時沿用第一階段的資料集。", + "train.stage_n": "第 {n} 階段", + "train.start": "開始訓練", + "train.started": "訓練已開始", + "train.stop": "停止訓練", + "train.title": "訓練", + "train.variants": "將測試 {n} 種組合", + "train.variants_capped": "要求 {n} 種組合，只測試前 {cap} 種", + "trial.complete": "完成", + "trial.failed": "失敗", + "trial.pruned": "已剪枝", + "tune.apply": "套用到模型設定", + "tune.best": "最佳參數", + "tune.duration": "用時", + "tune.intro": "以 Optuna TPE 搜尋訓練設定：每次試驗只訓練幾步並回報損失。", + "tune.loss": "探測損失", + "tune.note": "短程訓練後的探測損失，只用於排序，不代表最終品質。", + "tune.start": "開始搜尋", + "tune.stop": "停止搜尋", + "tune.title": "自動調參（Optuna TPE）", + "tune.trial": "試驗", + "tune.trials": "試驗", + "ui.access_tip.lan": "你透過區域網路存取。「上傳…」會把檔案傳到服務所在電腦的 uploads 資料夾，再使用那裡的路徑；「瀏覽…」列出的是服務所在電腦上的檔案。", + "ui.access_tip.lan_local": "你在服務所在的電腦上開啟，無需登入；「上傳…」會開啟系統檔案管理員並直接使用原路徑。區域網路內的其他裝置需要存取金鑰才能登入。", + "ui.access_tip.local_only": "介面只監聽本機（127.0.0.1），其他裝置無法開啟，也無需登入。啟動時加上 --lan 可開放區域網路存取。", + "ui.apply": "套用", + "ui.backend": "後端", + "ui.browse": "瀏覽…", + "ui.busy": "同類任務正在執行", + "ui.cancel": "取消", + "ui.chart_view": "圖表", + "ui.clear": "清除", + "ui.close": "關閉", + "ui.copy": "複製", + "ui.empty": "尚無內容", + "ui.error": "錯誤", + "ui.from_model": "沿用模型檔案", + "ui.hide_advanced": "隱藏進階選項", + "ui.invalid": "請檢查標示的欄位", + "ui.key": "存取金鑰", + "ui.lan_local": "本機存取 · 區域網路已開放", + "ui.language": "語言", + "ui.load": "載入", + "ui.loading": "載入中…", + "ui.local_only": "僅本機存取，無需登入", + "ui.log": "日誌", + "ui.login": "登入", + "ui.login_blocked": "失敗次數過多，請 5 分鐘後再試。", + "ui.login_failed": "金鑰錯誤", + "ui.login_needed": "此工作台需要伺服器上設定的存取金鑰。", + "ui.logout": "登出", + "ui.no": "否", + "ui.none": "無", + "ui.path": "路徑", + "ui.phase": "階段", + "ui.raw_json": "原始 JSON", + "ui.refresh": "重新整理", + "ui.remote_auth": "遠端存取，已登入", + "ui.reset": "還原預設", + "ui.result": "結果", + "ui.run": "執行", + "ui.save": "儲存", + "ui.saved": "已儲存", + "ui.scan": "掃描", + "ui.sections": "功能分區", + "ui.show_advanced": "顯示進階選項", + "ui.skip": "跳至內容", + "ui.start": "開始", + "ui.started": "開始時間", + "ui.state": "狀態", + "ui.step": "步", + "ui.stop": "停止", + "ui.table_view": "表格", + "ui.theme": "主題", + "ui.theme.dark": "深色", + "ui.theme.light": "淺色", + "ui.theme.system": "跟隨系統", + "ui.time": "時間", + "ui.upload": "上傳…", + "ui.value": "數值", + "ui.yes": "是" +} diff --git a/chronos/webui/jobs.py b/chronos/webui/jobs.py new file mode 100644 index 0000000..ce10e74 --- /dev/null +++ b/chronos/webui/jobs.py @@ -0,0 +1,195 @@ +"""Background jobs for the Web UI: one registry, two kinds of work. + +* ``SubprocessJob`` – training, benchmark: a child process whose stdout is + kept in a bounded buffer; ``stop`` sends SIGTERM (the train worker turns + that into a clean stage stop), then SIGKILL after a grace period. +* ``ThreadJob`` – inference, export, data scan, planning: a function run on a + thread so the server stays responsive; the result is kept for polling. + +Only one job of each *exclusive* kind runs at a time (two trainings would +fight for the same accelerator); finished jobs are kept for the session. +""" +from __future__ import annotations + +import itertools +import os +import signal +import subprocess +import threading +import time +from collections import deque + +EXCLUSIVE_KINDS = frozenset({"train", "benchmark", "export", "inference"}) + + +class JobBusy(RuntimeError): + """A job of an exclusive kind is already running.""" + + +class _Job: + def __init__(self, job_id: str, kind: str, title: str = ""): + self.id, self.kind, self.title = job_id, kind, title + self.state = "starting" + self.started = time.time() + self.ended: float | None = None + self.lines: deque[str] = deque(maxlen=2000) + self.result = None + self.error = "" + self.meta: dict = {} + self._lock = threading.Lock() + + def log(self, text: str) -> None: + with self._lock: + for line in str(text).splitlines() or [""]: + self.lines.append(line) + + def running(self) -> bool: + return self.state in ("starting", "running", "stopping") + + def view(self, *, log_lines: int = 200, with_result: bool = True) -> dict: + with self._lock: + tail = list(self.lines)[-log_lines:] if log_lines else [] + out = {"id": self.id, "kind": self.kind, "title": self.title, "state": self.state, + "started": self.started, "ended": self.ended, "error": self.error or None, + "log": tail, "meta": self.meta} + if with_result: + out["result"] = self.result + return out + + +class SubprocessJob(_Job): + def __init__(self, job_id: str, kind: str, argv: list[str], *, cwd: str, env: dict | None = None, + title: str = "", on_exit=None): + super().__init__(job_id, kind, title) + self.argv, self.on_exit = argv, on_exit + full_env = dict(os.environ) + full_env.update(env or {}) + full_env.setdefault("PYTHONUNBUFFERED", "1") + self.proc = subprocess.Popen(argv, cwd=cwd, env=full_env, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, text=True, bufsize=1, + start_new_session=True) + self.state = "running" + self.meta["pid"] = self.proc.pid + threading.Thread(target=self._pump, name=f"job-{job_id}", daemon=True).start() + + def _pump(self) -> None: + assert self.proc.stdout is not None + for line in self.proc.stdout: + self.log(line.rstrip("\n")) + code = self.proc.wait() + self.meta["returncode"] = code + if self.state == "stopping": + self.state = "stopped" + else: + self.state = "finished" if code == 0 else "failed" + if code != 0: + self.error = f"exit code {code}" + self.ended = time.time() + if self.on_exit is not None: + try: + self.result = self.on_exit(self) + except Exception as exc: # a broken result file must not kill the pump + self.error = self.error or f"{type(exc).__name__}: {exc}" + + def stop(self, grace_s: float = 20.0) -> None: + if self.proc.poll() is not None: + return + self.state = "stopping" + try: + self.proc.send_signal(signal.SIGTERM) + except ProcessLookupError: + return + + def _escalate(): + try: + self.proc.wait(timeout=grace_s) + except subprocess.TimeoutExpired: + try: + os.killpg(self.proc.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + self.proc.kill() + + threading.Thread(target=_escalate, daemon=True).start() + + +class ThreadJob(_Job): + def __init__(self, job_id: str, kind: str, fn, *, title: str = ""): + super().__init__(job_id, kind, title) + self.cancel = threading.Event() + self.state = "running" + threading.Thread(target=self._run, args=(fn,), name=f"job-{job_id}", daemon=True).start() + + def _run(self, fn) -> None: + try: + self.result = fn(self) + self.state = "stopped" if self.cancel.is_set() else "finished" + except Exception as exc: + import traceback + self.error = f"{type(exc).__name__}: {exc}" + self.log(traceback.format_exc()) + self.state = "failed" + self.ended = time.time() + + def stop(self, grace_s: float = 0.0) -> None: + if self.running(): + self.cancel.set() + self.state = "stopping" + + +class JobRegistry: + """Creates, finds and stops jobs; enforces one running job per exclusive kind.""" + + def __init__(self, keep: int = 64): + self._jobs: dict[str, _Job] = {} + self._ids = itertools.count(1) + self._lock = threading.Lock() + self.keep = keep + + def _new_id(self, kind: str) -> str: + return f"{kind}-{next(self._ids)}-{int(time.time())}" + + def _check(self, kind: str) -> None: + if kind in EXCLUSIVE_KINDS and any(j.kind == kind and j.running() for j in self._jobs.values()): + raise JobBusy(kind) + + def _add(self, job: _Job) -> _Job: + self._jobs[job.id] = job + done = [j for j in self._jobs.values() if not j.running()] + for old in sorted(done, key=lambda j: j.started)[:-self.keep] if len(done) > self.keep else []: + self._jobs.pop(old.id, None) + return job + + def spawn(self, kind: str, argv: list[str], *, cwd: str, env: dict | None = None, title: str = "", + on_exit=None, meta: dict | None = None) -> _Job: + with self._lock: + self._check(kind) + job = SubprocessJob(self._new_id(kind), kind, argv, cwd=cwd, env=env, title=title, on_exit=on_exit) + job.meta.update(meta or {}) + return self._add(job) + + def run(self, kind: str, fn, *, title: str = "", meta: dict | None = None) -> _Job: + with self._lock: + self._check(kind) + job = ThreadJob(self._new_id(kind), kind, fn, title=title) + job.meta.update(meta or {}) + return self._add(job) + + def get(self, job_id: str) -> _Job: + job = self._jobs.get(job_id) + if job is None: + raise KeyError(job_id) + return job + + def list(self, kind: str | None = None) -> list[dict]: + jobs = [j for j in self._jobs.values() if kind is None or j.kind == kind] + return [j.view(log_lines=0, with_result=False) for j in sorted(jobs, key=lambda j: j.started, reverse=True)] + + def stop(self, job_id: str) -> dict: + job = self.get(job_id) + job.stop() + return job.view(log_lines=0, with_result=False) + + def stop_all(self) -> None: + for job in list(self._jobs.values()): + if job.running(): + job.stop(grace_s=5.0) diff --git a/chronos/webui/native_dialog.py b/chronos/webui/native_dialog.py new file mode 100644 index 0000000..d8f7ffc --- /dev/null +++ b/chronos/webui/native_dialog.py @@ -0,0 +1,128 @@ +"""The operating system's own file dialog, opened by the Web UI server for a browser on the same machine. + +A browser never reveals the real path of a file the user picks, so when the page is opened on the +server's own machine (loopback), the server opens the native dialog itself and returns the chosen +paths. Nothing is copied. Parameters reach the dialog as process arguments or environment variables, +never spliced into script text. + +macOS uses ``osascript`` (``choose file`` / ``choose folder``); Linux uses ``zenity`` or ``kdialog`` +when a display is set; Windows uses PowerShell with ``System.Windows.Forms``. +""" +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +import threading + +TIMEOUT_S = 900 +_lock = threading.Lock() + +_MAC = r'''on run argv + set promptText to item 1 of argv + set startDir to item 2 of argv + set wantFolder to (item 3 of argv) is "folder" + set multi to (item 4 of argv) is "1" + try + set loc to (POSIX file startDir) as alias + on error + set loc to path to home folder + end try + activate + if wantFolder then + set picked to choose folder with prompt promptText default location loc multiple selections allowed multi + else + set picked to choose file with prompt promptText default location loc multiple selections allowed multi + end if + if class of picked is not list then set picked to {picked} + set out to "" + repeat with f in picked + set out to out & (POSIX path of f) & linefeed + end repeat + return out +end run +''' + +_WIN = r'''Add-Type -AssemblyName System.Windows.Forms +if ($env:CHRONOS_PICK_MODE -eq "folder") { + $d = New-Object System.Windows.Forms.FolderBrowserDialog + $d.Description = $env:CHRONOS_PICK_TITLE + if ($env:CHRONOS_PICK_START) { $d.SelectedPath = $env:CHRONOS_PICK_START } + if ($d.ShowDialog() -eq "OK") { [Console]::Out.WriteLine($d.SelectedPath) } else { exit 1 } +} else { + $d = New-Object System.Windows.Forms.OpenFileDialog + $d.Title = $env:CHRONOS_PICK_TITLE + $d.Multiselect = ($env:CHRONOS_PICK_MULTI -eq "1") + if ($env:CHRONOS_PICK_START) { $d.InitialDirectory = $env:CHRONOS_PICK_START } + if ($d.ShowDialog() -eq "OK") { $d.FileNames | ForEach-Object { [Console]::Out.WriteLine($_) } } else { exit 1 } +} +''' + + +class DialogBusy(RuntimeError): + """Another native dialog is already open.""" + + +def backend() -> str | None: + """Which native dialog this machine can show, or ``None``.""" + if sys.platform == "darwin": + return "osascript" if shutil.which("osascript") else None + if sys.platform.startswith("win"): + return "powershell" if shutil.which("powershell") else None + if not (os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")): + return None + for tool in ("zenity", "kdialog"): + if shutil.which(tool): + return tool + return None + + +def _command(tool: str, folder: bool, multiple: bool, start: str, title: str) -> tuple[list[str], str | None, dict]: + env = dict(os.environ) + if tool == "osascript": + return ["osascript", "-", title, start, "folder" if folder else "file", "1" if multiple else "0"], _MAC, env + if tool == "powershell": + env.update(CHRONOS_PICK_MODE="folder" if folder else "file", CHRONOS_PICK_MULTI="1" if multiple else "0", + CHRONOS_PICK_START=start, CHRONOS_PICK_TITLE=title) + return ["powershell", "-NoProfile", "-STA", "-Command", "-"], _WIN, env + if tool == "zenity": + cmd = ["zenity", "--file-selection", f"--title={title}", f"--filename={start.rstrip('/')}/"] + if folder: + cmd.append("--directory") + if multiple: + cmd += ["--multiple", "--separator=\n"] + return cmd, None, env + if folder: + return ["kdialog", "--title", title, "--getexistingdirectory", start], None, env + cmd = ["kdialog", "--title", title, "--getopenfilename", start] + if multiple: + cmd += ["--multiple", "--separate-output"] + return cmd, None, env + + +def choose(*, folder: bool, multiple: bool, start: str = "", title: str = "", timeout: float = TIMEOUT_S): + """Paths picked in the native dialog; ``None`` when cancelled. + + Raises ``RuntimeError`` when no native dialog is available here and :class:`DialogBusy` when one + is already open (the dialog belongs to the desktop session, so only one at a time). + """ + tool = backend() + if tool is None: + raise RuntimeError("no native file dialog on this machine") + start = start if start and os.path.isdir(start) else os.path.expanduser("~") + if not _lock.acquire(blocking=False): + raise DialogBusy("a file dialog is already open") + try: + cmd, script, env = _command(tool, folder, multiple, start, title[:200] or "Chronos") + try: + done = subprocess.run(cmd, input=script, capture_output=True, text=True, timeout=timeout, env=env, + check=False) + except subprocess.TimeoutExpired: + return None + finally: + _lock.release() + if done.returncode != 0: + return None # cancelled (osascript -128, zenity / kdialog 1) + paths = [line.strip() for line in done.stdout.splitlines() if line.strip()] + return [p.rstrip("/") if len(p) > 1 else p for p in paths] or None diff --git a/chronos/webui/schema.py b/chronos/webui/schema.py new file mode 100644 index 0000000..013bc1a --- /dev/null +++ b/chronos/webui/schema.py @@ -0,0 +1,460 @@ +"""Field schema for the Web UI: every setting the pages render, in one place. + +Each field carries its group, type, bounds / choices, the code default and — +where the project has one — a recommended value with its *source*: + +* ``code`` – the default in ``ChronosConfig`` / the runtime module; +* ``planner`` – chosen by the adaptive planner (``chronos.model.planner``); +* ``measured`` – backed by a measurement in ``reports/`` (named in ``ref``); +* ``rule`` – derived from the model / data (e.g. "= recurrent_depth"). + +Labels, hover help and choice names live in ``chronos/webui/i18n/.json`` +under ``field..label`` / ``field..help`` / ``choice..``; +``tests/test_webui.py`` checks every field here has them in all four +languages. The recommended value is rendered from this schema, never +duplicated in the catalogs, so the two cannot drift. +""" +from __future__ import annotations + +from dataclasses import asdict, dataclass + + +@dataclass(frozen=True) +class Field: + key: str + group: str + type: str # int | float | bool | choice | text | path | multi + default: object = None + min: float | None = None + max: float | None = None + step: float | None = None + choices: tuple = () + recommended: object = None + source: str | None = None # code | planner | measured | rule + ref: str | None = None # report path / rule id for the recommendation + advanced: bool = False + scope: tuple = ("config",) # which forms show it + nullable: bool = False # "auto" / empty allowed + + def as_dict(self) -> dict: + out = {k: v for k, v in asdict(self).items() if v not in (None, (), False) or k in ("default",)} + out["choices"] = list(self.choices) if self.choices else None + out["scope"] = list(self.scope) + return {k: v for k, v in out.items() if v is not None or k == "default"} + + +def _f(key, group, type_, default=None, **kw) -> Field: + for name in ("choices", "scope"): + if name in kw and isinstance(kw[name], list): + kw[name] = tuple(kw[name]) + return Field(key, group, type_, default, **kw) + + +LANGUAGES = ("zh-Hans", "zh-Hant", "en", "ja") +DEFAULT_LANGUAGE = "en" + +# Config-page groups, in page order. +CONFIG_GROUPS = ("model", "recurrent", "moe", "attention", "objectives", "memory", "optimizer", "runtime", "rl") +# Every group id (config + other pages); each needs ``group.`` in the catalogs. +GROUPS = CONFIG_GROUPS + ("train", "stage_data", "probe", "inference", "lazy", "speculative", "export", "serve", + "tokenizer", "benchmark", "autotune", "planner") + +RES = "reports/residency_20261107" +SPEC = "reports/speculative" + +CONFIG_FIELDS: tuple[Field, ...] = ( + # ---- model shape ------------------------------------------------------- + _f("hidden_size", "model", "int", 768, min=8, max=16384, step=8), + _f("num_hidden_layers", "model", "int", 8, min=1, max=128, step=1), + _f("vocab_size", "model", "int", 6400, min=32, max=200000, step=32, + recommended="tokenizer", source="rule", ref="vocab_matches_tokenizer"), + _f("max_position_embeddings", "model", "int", 32768, min=128, max=1048576, step=128, advanced=True), + _f("tie_word_embeddings", "model", "bool", True, recommended=True, source="code"), + _f("dropout", "model", "float", 0.1, min=0.0, max=0.5, step=0.01, advanced=True), + _f("rope_theta", "model", "float", 1e6, min=1e3, max=1e8, advanced=True), + _f("dtype", "model", "choice", "auto", choices=["auto", "float32", "bfloat16", "float16"], + recommended="auto", source="code"), + _f("modalities", "model", "text", "", advanced=True), + # ---- recurrent depth --------------------------------------------------- + _f("recurrent_enabled", "recurrent", "bool", True, recommended=True, source="code"), + _f("recurrent_groups", "recurrent", "int", 8, min=1, max=128, step=1, nullable=True), + _f("prelude_layers", "recurrent", "int", 0, min=0, max=16, step=1, advanced=True), + _f("coda_layers", "recurrent", "int", 0, min=0, max=16, step=1, advanced=True), + _f("recurrent_depth", "recurrent", "int", 1, min=1, max=16, step=1, recommended=4, source="planner"), + _f("min_recurrent_depth", "recurrent", "int", 1, min=1, max=16, step=1), + _f("max_inference_depth", "recurrent", "int", 1, min=1, max=32, step=1, + recommended="recurrent_depth", source="rule", ref="equals_recurrent_depth"), + _f("recurrent_attention_mode", "recurrent", "choice", "shared_first", + choices=["shared_first", "per_iteration", "first_only"], recommended="shared_first", source="code"), + _f("adaptive_halting", "recurrent", "bool", False), + _f("halting_threshold", "recurrent", "float", 0.99, min=0.5, max=1.0, step=0.001), + _f("depth_schedule", "recurrent", "choice", "fixed_then_adaptive", + choices=["fixed_then_adaptive", "fixed", "adaptive"], recommended="fixed_then_adaptive", source="code"), + _f("depth_schedule_warmup_fraction", "recurrent", "float", 0.25, min=0.0, max=1.0, step=0.01), + _f("recurrent_strategy", "recurrent", "choice", "scheduled", choices=["fixed", "scheduled", "adaptive"], + advanced=True), + _f("recurrent_depth_jitter", "recurrent", "float", 0.0, min=0.0, max=1.0, step=0.05, advanced=True), + _f("ponder_loss_coef", "recurrent", "float", 0.01, min=0.0, max=1.0, step=0.001), + _f("recurrent_distill_coef", "recurrent", "float", 0.05, min=0.0, max=1.0, step=0.01), + _f("recurrent_stability_coef", "recurrent", "float", 0.001, min=0.0, max=1.0, step=0.001), + _f("draft_exit_kl_coef", "recurrent", "float", 0.0, min=0.0, max=1.0, step=0.01, advanced=True), + _f("draft_exit_depth", "recurrent", "int", 1, min=1, max=16, step=1, advanced=True), + # ---- mixture of experts ------------------------------------------------ + _f("num_experts", "moe", "int", 4, min=1, max=1024, step=1), + _f("num_experts_per_tok", "moe", "int", 1, min=1, max=32, step=1, recommended=2, source="planner"), + _f("num_shared_experts", "moe", "int", 1, min=0, max=32, step=1, recommended=1, source="code"), + _f("moe_intermediate_size", "moe", "int", 0, min=0, max=65536, step=8, nullable=True), + _f("router_scope", "moe", "choice", "per_iteration", choices=["per_iteration", "shared"], + recommended="per_iteration", source="code"), + _f("lookahead_steps", "moe", "int", 2, min=0, max=16, step=1, recommended=2, source="code"), + _f("layer_aware_cache", "moe", "bool", False), + _f("prefetch_depth", "moe", "int", 2, min=0, max=64, step=1, recommended=2, source="code"), + _f("fallback_mask_prob", "moe", "float", 0.08, min=0.0, max=1.0, step=0.01, recommended=0.08, source="code"), + # ---- attention / KV ---------------------------------------------------- + _f("num_attention_heads", "attention", "int", 8, min=1, max=256, step=1), + _f("num_key_value_heads", "attention", "int", 4, min=1, max=256, step=1), + _f("use_hybrid_attention", "attention", "bool", True, recommended=True, source="code"), + _f("kv_latent_dim", "attention", "int", 64, min=4, max=1024, step=4), + _f("rope_dim", "attention", "int", 32, min=2, max=256, step=2), + _f("sliding_window_size", "attention", "int", 2048, min=8, max=65536, step=8), + _f("kv_cache_dtype", "attention", "choice", "auto", choices=["auto", "int8"], recommended="auto", source="code"), + _f("kv_quant_group_size", "attention", "int", 32, min=8, max=256, step=8, advanced=True), + _f("mla_decode", "attention", "choice", "expanded", choices=["expanded", "absorbed", "auto"], + recommended="expanded", source="code", ref="mla_auto_gap"), + _f("mla_cache_rope", "attention", "bool", False, advanced=True), + # ---- training objectives ----------------------------------------------- + _f("lambda_balance", "objectives", "float", 5e-4, min=0.0, max=1.0, step=1e-4, recommended=5e-4, source="code"), + _f("lambda_temporal", "objectives", "float", 3e-3, min=0.0, max=1.0, step=1e-4, recommended=3e-3, source="code"), + _f("lambda_lookahead", "objectives", "float", 0.1, min=0.0, max=10.0, step=0.01, recommended=0.1, source="code"), + _f("lambda_lookahead_topk", "objectives", "float", 0.15, min=0.0, max=10.0, step=0.01, advanced=True), + _f("lambda_lookahead_union", "objectives", "float", 0.05, min=0.0, max=10.0, step=0.01, advanced=True), + _f("lambda_router_locality", "objectives", "float", 0.02, min=0.0, max=10.0, step=0.01, advanced=True), + _f("lambda_router_anchor", "objectives", "float", 0.0, min=0.0, max=10.0, step=0.01), + _f("lambda_cluster_locality", "objectives", "float", 0.3, min=0.0, max=10.0, step=0.01, + recommended=0.3, source="measured", ref=f"{RES}/sweep"), + _f("cluster_locality_clusters", "objectives", "int", 0, min=0, max=1024, step=1, + recommended=0, source="rule", ref="clusters_auto"), + _f("cluster_locality_window", "objectives", "int", 32, min=0, max=65536, step=1, + recommended="data", source="rule", ref="window_median_record"), + # ---- memory / lazy experts --------------------------------------------- + _f("vram_budget_gb", "memory", "float", 4.0, min=0.01, max=256.0, step=0.01), + _f("pinned_memory_max_fraction", "memory", "float", 0.25, min=0.01, max=0.95, step=0.01), + _f("max_resident_experts", "memory", "int", None, min=1, max=4096, step=1, nullable=True), + _f("ram_budget_mb", "memory", "float", None, min=1.0, max=1048576.0, step=1.0, nullable=True), + _f("storage_format", "memory", "choice", "safetensors", choices=["safetensors", "pt"], + recommended="safetensors", source="code"), + _f("training_paging", "memory", "bool", False), + _f("optimizer_paging", "memory", "bool", False), + _f("page_size_mb", "memory", "int", 64, min=1, max=4096, step=1, advanced=True), + _f("optimizer_page_dir", "memory", "path", "./optimizer_pages", advanced=True), + _f("activation_checkpointing", "memory", "bool", False), + # ---- optimizer --------------------------------------------------------- + _f("learning_rate", "optimizer", "float", 5e-4, min=1e-7, max=1.0, step=1e-5, + recommended="planner", source="planner"), + _f("batch_size", "optimizer", "int", 16, min=1, max=512, step=1), + _f("accumulation_steps", "optimizer", "int", 8, min=1, max=256, step=1), + _f("max_seq_len", "optimizer", "int", 512, min=8, max=131072, step=8), + _f("epochs", "optimizer", "int", 2, min=1, max=1000, step=1), + _f("max_steps", "optimizer", "int", 0, min=0, max=10_000_000, step=1, recommended=0, source="code"), + _f("weight_decay", "optimizer", "float", 0.01, min=0.0, max=1.0, step=0.001), + _f("grad_clip", "optimizer", "float", 1.0, min=0.0, max=100.0, step=0.1), + # ---- runtime ----------------------------------------------------------- + _f("save_dir", "runtime", "path", "./out"), + _f("tokenizer_path", "runtime", "path", ""), + _f("save_interval", "runtime", "int", 1000, min=1, max=1_000_000, step=1), + _f("log_interval", "runtime", "int", 10, min=1, max=100000, step=1), + _f("cpu_threads", "runtime", "text", "auto", recommended="auto", source="code"), + _f("cpu_budget_percent", "runtime", "float", 100, min=10, max=100, step=1), + _f("num_workers", "runtime", "text", "auto", recommended="auto", source="code"), + # ---- RL / preference --------------------------------------------------- + _f("beta", "rl", "float", 0.1, min=0.0, max=10.0, step=0.01), + _f("alpha", "rl", "float", 0.7, min=0.0, max=1.0, step=0.01), + _f("lambda_or", "rl", "float", 0.1, min=0.0, max=10.0, step=0.01), + _f("reward_spec", "rl", "text", "toy"), + _f("max_gen_len", "rl", "int", 24, min=1, max=4096, step=1), + _f("num_generations", "rl", "int", 4, min=1, max=64, step=1), + _f("temperature", "rl", "float", 1.0, min=0.0, max=2.0, step=0.05), +) + +TRAIN = ("train",) +PROBE = ("train", "probe") +INFER = ("inference",) + +FORM_FIELDS: tuple[Field, ...] = ( + # ---- training job ------------------------------------------------------ + _f("train.backend", "train", "choice", "auto", choices=["auto", "mlx", "cuda", "xpu", "mps", "cpu"], + recommended="auto", source="code", scope=TRAIN), + _f("train.stages", "train", "multi", ["pretrain"], choices=["pretrain", "sft", "dpo", "orpo", "grpo", "distill"], + scope=TRAIN), + _f("train.data_path", "train", "path", "", scope=TRAIN), + # Later stages need their own data format (conversations, preference pairs, prompts); the Train + # page shows the row of each selected stage after the first. Empty = reuse the first stage's data. + *(_f(f"train.data_{m}", "stage_data", "path", "", scope=TRAIN, nullable=True) + for m in ("sft", "dpo", "orpo", "grpo", "distill")), + _f("train.init_weight", "train", "path", "", scope=TRAIN, nullable=True), + _f("train.teacher_path", "train", "path", "", scope=TRAIN, nullable=True), + _f("train.max_steps", "train", "int", 0, min=0, max=10_000_000, step=1, scope=TRAIN), + _f("train.val_ratio", "train", "float", 0.05, min=0.0, max=0.5, step=0.01, recommended=0.05, + source="code", scope=TRAIN), + _f("train.sample_prompt", "train", "text", "", scope=TRAIN), + # ---- lazy-expert probe after each validation --------------------------- + _f("probe.enabled", "probe", "bool", True, recommended=True, source="code", scope=PROBE), + _f("probe.backends", "probe", "multi", ["auto"], choices=["auto", "mlx", "cuda", "xpu", "mps", "cpu"], + recommended=["auto"], source="code", scope=PROBE), + _f("probe.paging", "probe", "multi", ["auto", "expert"], choices=["auto", "cluster", "expert"], + recommended=["auto", "expert"], source="code", scope=PROBE), + _f("probe.residency", "probe", "multi", ["planned"], choices=["planned", "window"], + recommended=["planned"], source="measured", ref=f"{RES}/bench_residency_50m_s13996_flores16_cpu.json", + scope=PROBE), + _f("probe.ram_ratios", "probe", "multi", [0.5, 1.0], + choices=[0.1, 0.2, 0.25, 0.33, 0.5, 0.67, 0.75, 0.9, 1.0], scope=PROBE), + _f("probe.miss_policies", "probe", "multi", ["sync_on_demand"], + choices=["sync_on_demand", "on_demand", "fallback_diagnostic"], recommended=["sync_on_demand"], + source="rule", ref="sync_is_exact", scope=PROBE), + _f("probe.residency_decay", "probe", "float", 0.7, min=0.0, max=0.95, step=0.05, recommended=0.7, + source="code", scope=PROBE, advanced=True), + _f("probe.decode_wait_s", "probe", "float", 0.012, min=0.0, max=1.0, step=0.001, recommended=0.012, + source="code", scope=PROBE, advanced=True), + _f("probe.verify_wait_s", "probe", "float", 0.02, min=0.0, max=1.0, step=0.001, recommended=0.02, + source="code", scope=PROBE, advanced=True), + _f("probe.speculative", "probe", "multi", [False], choices=[False, True], scope=PROBE), + _f("probe.draft_depth", "probe", "int", 1, min=1, max=16, step=1, scope=PROBE, advanced=True), + _f("probe.gamma", "probe", "int", 4, min=1, max=8, step=1, scope=PROBE, advanced=True), + _f("probe.max_new_tokens", "probe", "int", 32, min=1, max=512, step=1, recommended=32, source="code", + scope=PROBE), + _f("probe.max_variants", "probe", "int", 12, min=1, max=256, step=1, recommended=12, source="code", + scope=PROBE), + # ---- inference --------------------------------------------------------- + _f("infer.model_path", "inference", "path", "", scope=INFER), + _f("infer.backend", "inference", "choice", "auto", choices=["auto", "mlx", "cuda", "xpu", "mps", "cpu"], + recommended="auto", source="code", scope=INFER), + _f("infer.mode", "inference", "choice", "compare", choices=["compare", "offload", "fullload"], scope=INFER), + _f("infer.max_tokens", "inference", "int", 128, min=1, max=4096, step=1, scope=INFER), + _f("infer.temperature", "inference", "float", 0.85, min=0.0, max=2.0, step=0.05, scope=INFER), + _f("infer.raw_prompt", "inference", "bool", False, scope=INFER), + _f("infer.ram_load_ratio", "lazy", "float", 1.0, min=0.1, max=1.25, step=0.01, scope=INFER), + _f("infer.sweep_ram_ratios", "lazy", "bool", False, scope=INFER), + _f("infer.miss_policy", "lazy", "choice", "sync_on_demand", + choices=["on_demand", "sync_on_demand", "fallback_diagnostic"], recommended="sync_on_demand", + source="rule", ref="sync_is_exact", scope=INFER), + _f("infer.paging", "lazy", "choice", "auto", choices=["auto", "cluster", "expert"], recommended="auto", + source="code", scope=INFER), + _f("infer.residency_planner", "lazy", "bool", True, recommended=True, source="measured", + ref=f"{RES}/bench_residency_50m_s13996_flores16_cpu.json", scope=INFER), + _f("infer.residency_decay", "lazy", "float", 0.7, min=0.0, max=0.95, step=0.05, recommended=0.7, + source="code", scope=INFER, advanced=True), + _f("infer.decode_wait_s", "lazy", "float", 0.012, min=0.0, max=1.0, step=0.001, recommended=0.012, + source="code", scope=INFER, advanced=True), + _f("infer.verify_wait_s", "lazy", "float", 0.02, min=0.0, max=1.0, step=0.001, recommended=0.02, + source="code", scope=INFER, advanced=True), + _f("infer.speculative", "speculative", "bool", False, recommended=False, source="measured", + ref=SPEC, scope=INFER), + _f("infer.draft_depth", "speculative", "int", 1, min=1, max=8, step=1, scope=INFER), + _f("infer.gamma", "speculative", "int", 4, min=1, max=8, step=1, scope=INFER), + # ---- export ------------------------------------------------------------ + _f("export.model_path", "export", "path", "", scope=("export",)), + _f("export.output_dir", "export", "path", "./exports/chronos", scope=("export",)), + _f("export.formats", "export", "multi", ["fp16-safetensors", "q8_0-gguf"], + choices=["fp16-safetensors", "q8_0-safetensors", "fp16-gguf", "q8_0-gguf", "chronos-native-gguf"], + scope=("export",)), + _f("export.config_path", "export", "path", "", scope=("export",), nullable=True, advanced=True), + _f("export.expert_cache_dir", "export", "path", "", scope=("export",), nullable=True, advanced=True), + _f("export.copy_cache", "export", "bool", True, scope=("export",)), + _f("export.auto_cluster", "export", "bool", True, recommended=True, source="code", scope=("export",)), + _f("export.calibration_data_path", "export", "path", "", scope=("export",), nullable=True), + _f("export.cluster_device", "export", "choice", "cpu", choices=["cpu", "mps", "cuda"], scope=("export",), + advanced=True), + _f("export.cluster_max_batches", "export", "int", 50, min=1, max=10000, step=1, scope=("export",), + advanced=True), + _f("export.cluster_batch_size", "export", "int", 4, min=1, max=256, step=1, scope=("export",), advanced=True), + _f("export.cluster_max_seq_len", "export", "int", 256, min=8, max=8192, step=8, scope=("export",), + advanced=True), + # ---- serve (Ollama / OpenAI protocol) ---------------------------------- + _f("serve.models_dir", "serve", "path", "./exports", scope=("serve",)), + _f("serve.host", "serve", "text", "127.0.0.1", recommended="127.0.0.1", source="rule", ref="loopback", + scope=("serve",)), + _f("serve.port", "serve", "int", 11435, min=1024, max=65535, step=1, scope=("serve",)), + _f("serve.device", "serve", "choice", "auto", choices=["auto", "mlx", "cuda", "mps", "cpu"], scope=("serve",)), + _f("serve.dtype", "serve", "choice", "auto", choices=["auto", "float32", "bfloat16", "float16"], + scope=("serve",)), + _f("serve.lazy_experts", "serve", "bool", True, recommended=True, source="code", scope=("serve",)), + _f("serve.ram_expert_fraction", "serve", "float", 0.5, min=0.05, max=1.0, step=0.05, scope=("serve",)), + _f("serve.paging", "serve", "choice", "auto", choices=["auto", "cluster", "expert"], recommended="auto", + source="code", scope=("serve",)), + _f("serve.kv_cache_dtype", "serve", "choice", "", choices=["", "auto", "int8"], scope=("serve",), + advanced=True), + _f("serve.mla_decode", "serve", "choice", "", choices=["", "expanded", "absorbed", "auto"], scope=("serve",), + advanced=True), + _f("serve.router_scope", "serve", "choice", "", choices=["", "per_iteration", "shared"], scope=("serve",), + advanced=True), + _f("serve.mla_cache_rope", "serve", "bool", False, scope=("serve",), advanced=True), + _f("serve.max_resident_experts", "serve", "int", None, min=1, max=4096, step=1, nullable=True, + scope=("serve",), advanced=True), + _f("serve.ram_budget_mb", "serve", "float", None, min=1.0, max=1048576.0, step=1.0, nullable=True, + scope=("serve",), advanced=True), + _f("serve.allow_remote", "serve", "bool", False, recommended=False, source="rule", ref="loopback", + scope=("serve",)), + # ---- tokenizer builder ------------------------------------------------- + _f("tok.paths", "tokenizer", "text", "", scope=("tokenizer",)), + _f("tok.out_dir", "tokenizer", "path", "./out/tokenizer", scope=("tokenizer",)), + _f("tok.vocab_size", "tokenizer", "int", 16384, min=256, max=65535, step=64, scope=("tokenizer",)), + _f("tok.per_source", "tokenizer", "int", 400_000, min=1000, max=100_000_000, step=1000, scope=("tokenizer",)), + _f("tok.temperature", "tokenizer", "float", 0.5, min=0.1, max=1.0, step=0.05, recommended=0.5, source="code", + scope=("tokenizer",)), + _f("tok.fields", "tokenizer", "text", "", scope=("tokenizer",), advanced=True), + # ---- auto-tune --------------------------------------------------------- + _f("tune.data_path", "autotune", "path", "", scope=("autotune",)), + _f("tune.n_trials", "autotune", "int", 20, min=2, max=200, step=1, scope=("autotune",)), + _f("tune.probe_steps", "autotune", "int", 80, min=10, max=1000, step=10, scope=("autotune",)), + _f("tune.search", "autotune", "multi", ["lambda_balance", "lambda_temporal", "lr"], + choices=["lambda_balance", "lambda_temporal", "lambda_lookahead", "lookahead_steps", "lr", "hidden_size", + "num_experts", "num_shared_experts", "kv_latent_dim"], scope=("autotune",)), + # ---- planner ----------------------------------------------------------- + _f("plan.target_m", "planner", "float", 50.0, min=0.1, max=100000.0, step=0.1, scope=("planner",)), + _f("plan.use_moe", "planner", "bool", True, scope=("planner",)), + _f("plan.num_experts", "planner", "int", 32, min=1, max=1024, step=1, scope=("planner",)), + _f("plan.active_mode", "planner", "choice", "experts", choices=["experts", "params"], scope=("planner",)), + _f("plan.active_value", "planner", "float", 2, min=0, max=100000, step=1, scope=("planner",)), + _f("plan.memory_saving", "planner", "choice", "none", choices=["none", "balanced", "aggressive", "extreme"], + scope=("planner",)), + _f("plan.data_paths", "planner", "text", "", scope=("planner",)), + _f("plan.ram_gb", "planner", "float", None, min=1, max=4096, step=1, nullable=True, scope=("planner",)), + _f("plan.token_budget_b", "planner", "float", None, min=0.001, max=100000, step=0.1, nullable=True, + scope=("planner",)), + _f("plan.depth", "planner", "int", 4, min=1, max=16, step=1, scope=("planner",)), + _f("plan.vocab_size", "planner", "int", None, min=256, max=200000, step=32, nullable=True, scope=("planner",)), + _f("plan.seq_len", "planner", "int", None, min=16, max=131072, step=16, nullable=True, scope=("planner",)), +) + +ALL_FIELDS: tuple[Field, ...] = CONFIG_FIELDS + FORM_FIELDS +FIELDS = {f.key: f for f in ALL_FIELDS} +# Recommendation ids rendered as text (``rule.`` in the catalogs). +RULE_REFS = tuple(sorted({f.ref for f in ALL_FIELDS if f.source == "rule" and f.ref})) +SOURCES = ("code", "planner", "measured", "rule") + + +# Technical identifiers shown verbatim in every language (backend, dtype, file-format and +# stage names); every other choice value needs ``choice..`` in each catalog. +LITERAL_CHOICES = frozenset({ + "mlx", "cuda", "xpu", "mps", "cpu", "float32", "bfloat16", "float16", "int8", "safetensors", "pt", + "fp16-safetensors", "q8_0-safetensors", "fp16-gguf", "q8_0-gguf", "chronos-native-gguf", + "pretrain", "sft", "dpo", "orpo", "grpo", "distill", "shared_first", "per_iteration", "first_only", + "lambda_balance", "lambda_temporal", "lambda_lookahead", "lookahead_steps", "lr", "hidden_size", + "num_experts", "num_shared_experts", "kv_latent_dim", "sync_on_demand", "on_demand", "fallback_diagnostic", +}) + + +def choice_keys() -> list[str]: + """``choice.`` keys every catalog must define. Values mean the same thing in every + field that offers them (``auto``, ``cluster``, ``fixed`` …), so they are translated once; + booleans, numbers, blanks and :data:`LITERAL_CHOICES` are shown verbatim.""" + out = [] + for f in ALL_FIELDS: + if f.type in ("choice", "multi"): + for value in f.choices: + if (isinstance(value, bool) or value == "" or isinstance(value, (int, float)) + or value in LITERAL_CHOICES): + continue + out.append(f"choice.{value}") + return list(dict.fromkeys(out)) + + +def required_keys() -> set[str]: + keys = {f"group.{g}" for g in GROUPS} | {f"source.{s}" for s in SOURCES} | {f"rule.{r}" for r in RULE_REFS} + for f in ALL_FIELDS: + keys.add(f"field.{f.key}.label") + keys.add(f"field.{f.key}.help") + keys.update(choice_keys()) + return keys + + +# What the browse button next to a path field opens: (kind, multiple). ``kind`` is a filter of +# ``chronos.workbench.files.KINDS``; ``multiple`` fields hold one path per line and several datasets +# are merged into one JSONL before training or calibration reads them. +PICK = { + "optimizer_page_dir": ("dir", False), "save_dir": ("dir", False), "tokenizer_path": ("dir", False), + "train.data_path": ("data", True), + **{f"train.data_{m}": ("data", True) for m in ("sft", "dpo", "orpo", "grpo", "distill")}, + "train.init_weight": ("checkpoint", False), "train.teacher_path": ("checkpoint", False), + "infer.model_path": ("checkpoint", False), + "export.model_path": ("checkpoint", False), "export.output_dir": ("dir", False), + "export.config_path": ("config", False), "export.expert_cache_dir": ("dir", False), + "export.calibration_data_path": ("data", True), + "serve.models_dir": ("dir", False), + "tok.paths": ("data", True), "tok.out_dir": ("dir", False), + "tune.data_path": ("data", True), + "plan.data_paths": ("data", True), +} +MULTI_PATH = frozenset(k for k, (_kind, many) in PICK.items() if many) +# Multi-dataset fields whose selection is merged into one JSONL (training stages, export calibration). +# The tokenizer samples each source on its own and the planner measures each, so they are not merged. +MERGE_PATH = frozenset({"train.data_path", "export.calibration_data_path"} + | {f"train.data_{m}" for m in ("sft", "dpo", "orpo", "grpo", "distill")}) +MAX_TEXT = 4096 +MAX_MULTI_TEXT = 65536 + + +def schema_payload() -> dict: + fields = [] + for f in ALL_FIELDS: + item = f.as_dict() + if f.key in PICK: + item["pick"] = {"kind": PICK[f.key][0], "multiple": PICK[f.key][1], "merge": f.key in MERGE_PATH} + fields.append(item) + return {"version": 1, "languages": list(LANGUAGES), "default_language": DEFAULT_LANGUAGE, + "config_groups": list(CONFIG_GROUPS), "groups": list(GROUPS), "fields": fields} + + +def config_defaults() -> dict: + return {f.key: f.default for f in CONFIG_FIELDS} + + +def coerce(key: str, value): + """Validate one value against its field; raises ``ValueError(key)``.""" + f = FIELDS[key] + if value is None or value == "": + if f.nullable or f.type in ("text", "path") or f.default in (None, ""): + return None if f.type not in ("text", "path") else "" + raise ValueError(key) + if f.type == "bool": + if isinstance(value, bool): + return value + raise ValueError(key) + if f.type in ("int", "float"): + try: + num = float(value) + except (TypeError, ValueError): + raise ValueError(key) from None + if num != num or (f.min is not None and num < f.min) or (f.max is not None and num > f.max): + raise ValueError(key) + if f.type == "int": + if num != int(num): + raise ValueError(key) + return int(num) + return num + if f.type == "choice": + if value not in f.choices: + raise ValueError(key) + return value + if f.type == "multi": + items = value if isinstance(value, list) else [value] + if not items or any(v not in f.choices for v in items): + raise ValueError(key) + return list(dict.fromkeys(items)) + limit = MAX_MULTI_TEXT if key in MULTI_PATH else MAX_TEXT + if not isinstance(value, str) or len(value) > limit or "\x00" in value: + raise ValueError(key) + return value + + +def coerce_config(values: dict) -> tuple[dict, list[str]]: + """Model/runtime config from the Config page: ``(clean, bad_keys)``; unknown keys are dropped.""" + clean, bad = {}, [] + for key, value in (values or {}).items(): + if key not in FIELDS or FIELDS[key] not in CONFIG_FIELDS: + continue + try: + clean[key] = coerce(key, value) + except ValueError: + bad.append(key) + return clean, bad diff --git a/chronos/webui/server.py b/chronos/webui/server.py new file mode 100644 index 0000000..b03fcb6 --- /dev/null +++ b/chronos/webui/server.py @@ -0,0 +1,1083 @@ +"""Chronos Web UI: FastAPI JSON API + a hand-written single-page front end. + + chronos webui # http://127.0.0.1:7860/ + chronos webui --lan # also reachable from the local network (login for other machines) + chronos webui --legacy # the old Gradio console + +Security +* Default bind is 127.0.0.1 with **no authentication** (local use only); the + page says so. Any other address needs ``--allow-remote`` *and* a key in + ``CHRONOS_UI_KEY`` / ``CHRONOS_API_KEY`` (never on the command line). +* With a key, the browser logs in once and receives an HttpOnly, + SameSite=Strict session cookie; scripts may send ``Authorization: Bearer``. + Failed logins are rate-limited per client; keys are compared in constant time. +* Every request must carry an allowed ``Host`` (DNS-rebinding guard); every + mutating request must carry ``X-Chronos-Request: 1`` and, when present, an + ``Origin`` on the same allow-list (CSRF guard). The page is served with a + strict CSP (``script-src 'self'``, no inline script). +* ``--lan`` (off by default) listens on every interface and accepts this + machine's LAN addresses as Host. Browsers on other machines must log in (the key + comes from ``CHRONOS_UI_KEY`` or is generated and printed at start); a browser + on this machine (loopback peer *and* loopback Host, no proxy headers) needs no + login. Do not put ``--lan`` behind a reverse proxy; use ``--allow-remote``. +* File fields: on this machine the server opens the system file dialog and the + chosen paths are used in place. From another machine files are uploaded + (streamed, size-capped, disk reserve kept) into ``/uploads//``. +* This UI starts training, inference, export and servers and reads and writes + files the server user can access: treat access to it as shell access. +""" +from __future__ import annotations + +import hmac +import json +import os +import re +import shutil +import sys +import time + +try: # module level: FastAPI resolves ``Request`` annotations by name + from fastapi import FastAPI, HTTPException, Request + from fastapi.concurrency import run_in_threadpool + from fastapi.responses import FileResponse, JSONResponse +except ImportError: # pragma: no cover - serving extra not installed + FastAPI = HTTPException = Request = FileResponse = JSONResponse = run_in_threadpool = None + +from chronos.monitor.server import LOOPBACK, SECURITY_HEADERS, _Sessions, check_bind +from chronos.webui import schema +from chronos.webui.jobs import JobBusy, JobRegistry + +HERE = os.path.dirname(os.path.abspath(__file__)) +STATIC_DIR = os.path.join(HERE, "static") +I18N_DIR = os.path.join(HERE, "i18n") +REPO_ROOT = os.path.dirname(os.path.dirname(HERE)) +COOKIE = "chronos_ui_session" +STATIC_TYPES = {".js": "text/javascript; charset=utf-8", ".css": "text/css; charset=utf-8", + ".svg": "image/svg+xml"} +_STATIC_NAME = re.compile(r"^[a-z0-9_\-]+\.(js|css|svg)$") +_LANG_ALIASES = {"zh-tw": "zh-Hant", "zh-hk": "zh-Hant", "zh-mo": "zh-Hant", "zh-hant": "zh-Hant"} +MUTATING = frozenset({"POST", "PUT", "PATCH", "DELETE"}) +STAGE_ORDER = ("pretrain", "sft", "dpo", "orpo", "grpo", "distill") + + +def resolve_language(tag: str | None) -> str: + """Map a BCP-47 tag (or Accept-Language list) to a supported catalog; English otherwise.""" + for part in (tag or "").split(","): + code = part.split(";")[0].strip().lower().replace("_", "-") + if not code: + continue + if code in _LANG_ALIASES or code.startswith("zh-hant"): + return "zh-Hant" + if code == "zh" or code.startswith("zh-"): + return "zh-Hans" + if code == "ja" or code.startswith("ja-"): + return "ja" + if code == "en" or code.startswith("en-"): + return "en" + return schema.DEFAULT_LANGUAGE + + +def load_catalog(lang: str) -> dict: + if lang not in schema.LANGUAGES: + raise KeyError(lang) + with open(os.path.join(I18N_DIR, f"{lang}.json"), encoding="utf-8") as fh: + return json.load(fh) + + +MAX_UPLOAD_MB = 20480 # per file; --max-upload-mb changes it +UPLOAD_DISK_RESERVE = 8 * 2**30 # uploads stop while less than this (or 5 % of the disk) stays free +LOOPBACK_PEERS = frozenset({"127.0.0.1", "::1", "localhost"}) +LOOPBACK_HOSTS = frozenset({"127.0.0.1", "localhost", "::1", "[::1]"}) + + +def lan_addresses() -> list[str]: + """This machine's IPv4 addresses other than loopback (what a LAN browser types).""" + import socket + out = [] + try: + import psutil + for addrs in psutil.net_if_addrs().values(): + out += [a.address for a in addrs if a.family == socket.AF_INET and not a.address.startswith("127.")] + except Exception: # pragma: no cover - psutil is a dependency; keep a fallback anyway + try: + out = [a for a in socket.gethostbyname_ex(socket.gethostname())[2] if not a.startswith("127.")] + except OSError: + out = [] + return sorted(set(out), key=lambda a: (not a.startswith(("192.168.", "10.")), a)) + + +def allowed_hosts(bind_host: str, extra: tuple[str, ...] = ()) -> set[str]: + """Host header values (without port) the server answers to.""" + hosts = {"127.0.0.1", "localhost", "::1", "[::1]"} + if bind_host and bind_host not in ("0.0.0.0", "::"): + hosts.add(bind_host.strip("[]")) + hosts.add(f"[{bind_host.strip('[]')}]" if ":" in bind_host else bind_host) + hosts.update(h.strip().lower() for h in extra if h.strip()) + return hosts + + +def _host_only(value: str) -> str: + value = (value or "").strip().lower() + if value.startswith("["): + return value[: value.find("]") + 1] if "]" in value else value + return value.rsplit(":", 1)[0] if value.count(":") == 1 else value + + +# ── config sanitising ────────────────────────────────────────────────── +def _plain_scalar(value) -> bool: + return value is None or isinstance(value, (bool, int, float)) or (isinstance(value, str) and len(value) <= 4096) + + +def _small_json(value, limit: int = 16_384) -> bool: + try: + return len(json.dumps(value, allow_nan=False)) <= limit + except (TypeError, ValueError): + return False + + +def sanitize_config(values: dict) -> tuple[dict, list[str]]: + """Config sent by the page → ``(clean, bad_keys)``. + + Schema fields are validated by :func:`schema.coerce`; ``modalities`` is a + list (comma text in the form); other canonical Chronos fields (planner or + preset values such as ``use_moe``) pass when they are small JSON values; + ``stage_overrides`` is ``{stage: {key: scalar}}``; anything else is dropped.""" + from chronos.config_schema import CANONICAL_FIELDS + from chronos.model.planner import STAGE_KEYS, STAGES + values = dict(values or {}) + modalities = values.pop("modalities", None) + clean, bad = schema.coerce_config(values) + if modalities is not None: + items = modalities if isinstance(modalities, list) else str(modalities).split(",") + names = sorted({str(m).strip().lower() for m in items if str(m).strip()} - {"text"}) + if all(n in ("image", "audio") for n in names): + clean["modalities"] = names + else: + bad.append("modalities") + for key, value in values.items(): + if key in clean or key in bad: + continue + if key == "stage_overrides" and isinstance(value, dict): + clean[key] = {stage: {k: v for k, v in (row or {}).items() if k in STAGE_KEYS and _plain_scalar(v)} + for stage, row in value.items() if stage in STAGES and isinstance(row, dict)} + elif key in CANONICAL_FIELDS and key not in schema.FIELDS and _small_json(value): + clean[key] = value + return clean, bad + + +def form_values(body: dict, prefix: str) -> tuple[dict, list[str]]: + """``{".": value}`` from a page form → ``({name: value}, bad_keys)``; defaults fill gaps.""" + out, bad = {}, [] + given = body or {} + for f in schema.FORM_FIELDS: + if not f.key.startswith(prefix + "."): + continue + name = f.key[len(prefix) + 1:] + raw = given.get(f.key, given.get(name, f.default)) + try: + out[name] = schema.coerce(f.key, raw) + except ValueError: + bad.append(f.key) + return out, bad + + +def probe_sweep(values: dict) -> dict: + """``probe.*`` form values → a ``ProbeSweep`` dict for the train worker.""" + return { + "backends": values["backends"], "paging": values["paging"], "residency": values["residency"], + "ram_ratios": [float(r) for r in values["ram_ratios"]], "miss_policies": values["miss_policies"], + "residency_decay": [values["residency_decay"]], "decode_wait_s": [values["decode_wait_s"]], + "verify_wait_s": [values["verify_wait_s"]], "speculative": [bool(s) for s in values["speculative"]], + "draft_depth": values["draft_depth"], "gamma": values["gamma"], + "max_new_tokens": values["max_new_tokens"], "max_variants": values["max_variants"], + } + + +def config_estimate(cfg: dict) -> dict: + """Size / memory estimate for the Config page (``ui.estimator``; estimated, not measured).""" + from ui.estimator import ( + ArchConfig, active_param_bounds, compute_params, estimated_decode_tps, memory_footprint, + parameter_breakdown, total_params, + ) + arch = ArchConfig.from_config(cfg) + lo, hi = active_param_bounds(arch) + groups = int(arch.recurrent_groups or 0) + depth = int(arch.executed_depth) if arch.recurrent_enabled else 1 + virtual = (int(arch.prelude_layers) + groups * depth + int(arch.coda_layers) + if arch.recurrent_enabled and groups else int(arch.num_hidden_layers)) + return {"label": "estimated", "total_params": total_params(arch), "active_params": lo, + "active_union_upper": hi, "compute_params": compute_params(arch), + "breakdown": parameter_breakdown(arch), "memory": memory_footprint(arch), + "decode_tps_estimate": round(float(estimated_decode_tps(arch)), 1), "virtual_layers": virtual} + + +def plan_request(values: dict): + """``plan.*`` form values → ``PlanRequest`` (millions / billions in the form); ValueError(err key).""" + from chronos.model.planner import MEMORY_LEVELS, PlanRequest + from chronos.workbench.tokenizer import parse_paths + target = float(values.get("target_m") or 0) * 1e6 + if target <= 0: + raise ValueError("target") + use_moe = bool(values.get("use_moe")) + value = float(values.get("active_value") or 0) + experts = int(values.get("num_experts") or 1) + active_params = value * 1e6 if use_moe and values.get("active_mode") == "params" and value > 0 else None + active_experts = int(value) if use_moe and values.get("active_mode") == "experts" and value >= 1 else None + if active_experts and active_experts > experts: + raise ValueError("active_experts") + return PlanRequest( + target_total=target, use_moe=use_moe, num_experts=experts, active_params=active_params, + active_experts=active_experts, + memory_saving=values.get("memory_saving") if values.get("memory_saving") in MEMORY_LEVELS else "none", + data_paths=tuple(parse_paths(values.get("data_paths") or "")), + ram_gb=float(values["ram_gb"]) if values.get("ram_gb") else None, + token_budget=float(values["token_budget_b"]) * 1e9 if values.get("token_budget_b") else None, + recurrent_depth=int(values.get("depth") or 4), + vocab_size=int(values["vocab_size"]) if values.get("vocab_size") else None, + seq_len=int(values["seq_len"]) if values.get("seq_len") else None, + ) + + +def system_info() -> dict: + import platform + info = {"python": platform.python_version(), "platform": platform.platform(terse=True), + "machine": platform.machine(), "cpu_count": os.cpu_count()} + try: + from chronos import __version__ + info["chronos"] = __version__ + except Exception: # pragma: no cover - partial install + info["chronos"] = "?" + try: + import psutil + info["ram_gb"] = round(psutil.virtual_memory().total / 2**30, 1) + except Exception: + info["ram_gb"] = None + backends = ["cpu"] + try: + import torch + if torch.cuda.is_available(): + backends.append("cuda") + if getattr(torch, "xpu", None) is not None and torch.xpu.is_available(): + backends.append("xpu") + if torch.backends.mps.is_available(): + backends.append("mps") + except Exception: + pass + try: + import mlx.core # noqa: F401 + backends.append("mlx") + except Exception: + pass + info["backends"] = backends + return info + + +def read_presets() -> dict: + """Planner presets as complete config dicts (``ui.presets`` merges each over the base defaults).""" + from ui.presets import DEFAULT_PRESET, V3_PRESET_META, V3_PRESETS + presets = [{"name": name, "config": sanitize_config(cfg)[0], "meta": V3_PRESET_META.get(name, {})} + for name, cfg in V3_PRESETS.items()] + names = [p["name"] for p in presets] + return {"default": DEFAULT_PRESET if DEFAULT_PRESET in names else (names[0] if names else ""), + "presets": presets} + + +def base_config() -> dict: + """What an empty Config page starts from: the default planner preset, else the schema defaults.""" + data = read_presets() + by_name = {p["name"]: p["config"] for p in data["presets"]} + return {**schema.config_defaults(), **by_name.get(data["default"], {})} + + +def _json_safe(value): + """Plain JSON types (unknown objects → str); NaN / inf become None via :func:`_finite`.""" + return _finite(json.loads(json.dumps(value, default=str))) if value is not None else None + + +def _finite(value): + """Replace NaN / inf (not valid JSON) with None, recursively.""" + if isinstance(value, float): + return value if value == value and value not in (float("inf"), float("-inf")) else None + if isinstance(value, dict): + return {k: _finite(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_finite(v) for v in value] + return value + + +class UIState: + """Everything one server instance owns: jobs, the serve child, the tokenizer job, the tuner.""" + + def __init__(self, workdir: str, runs_dir: str): + from chronos.monitor.snapshot import RunMonitor + from chronos.workbench.serve import ServeProcess + from chronos.workbench.tokenizer import TokenizerJob + self.workdir = os.path.abspath(workdir) + self.runs_dir = os.path.abspath(os.path.join(self.workdir, runs_dir)) + os.makedirs(self.runs_dir, exist_ok=True) + self.jobs = JobRegistry() + self.serve = ServeProcess() + self.tokenizer = TokenizerJob() + self.monitor = RunMonitor(self.runs_dir) + self.tuner = None + self.tune_log: list[str] = [] + + def path(self, value: str) -> str: + """User path relative to the work directory, ``~`` expanded.""" + value = os.path.expanduser((value or "").strip()) + return value if not value or os.path.isabs(value) else os.path.join(self.workdir, value) + + def paths(self, value) -> list[str]: + """Several user paths (one per line, or a list), each resolved like :meth:`path`.""" + from chronos.workbench.files import parse_paths + return [self.path(p) for p in parse_paths(value)] + + def shutdown(self) -> None: + self.jobs.stop_all() + self.serve.stop(timeout=5.0) + if self.tuner is not None and self.tuner.is_running(): + self.tuner.stop() + + +def create_app(workdir: str = ".", *, runs_dir: str = "runs", api_key: str | None = None, + bind_host: str = "127.0.0.1", extra_hosts: tuple[str, ...] = (), session_hours: float = 12.0, + max_upload_mb: int = MAX_UPLOAD_MB, trust_local: bool = False, native_dialog: bool = True): + """The Web UI FastAPI app; ``app.state.ui`` holds the :class:`UIState`.""" + if FastAPI is None: + raise ImportError("chronos webui needs fastapi and uvicorn (pip install 'Project_Chronos[serve]')") + ui = UIState(workdir, runs_dir) + sessions = _Sessions(session_hours * 3600) + hosts = allowed_hosts(bind_host, extra_hosts) + app = FastAPI(title="Chronos Web UI", docs_url=None, redoc_url=None, openapi_url=None) + app.state.ui = ui + app.add_event_handler("shutdown", ui.shutdown) + public = {"/", "/health", "/api/session", "/api/login", "/api/logout", "/api/i18n"} + + def _key_ok(given: str) -> bool: + return bool(api_key) and hmac.compare_digest(given.encode("utf-8"), api_key.encode("utf-8")) + + def _authorized(request) -> bool: + if not api_key or (trust_local and _is_local(request)): + return True + auth = request.headers.get("authorization", "") + if auth.startswith("Bearer ") and _key_ok(auth[7:]): + return True + return sessions.valid(request.cookies.get(COOKIE)) + + def _is_local(request) -> bool: + """The browser runs on this machine: loopback peer *and* a loopback Host (not a LAN URL).""" + peer = request.client.host if request.client else "" + if any(h in request.headers for h in ("forwarded", "x-forwarded-for", "x-real-ip")): + return False # came through a proxy: the real browser is elsewhere + return peer in LOOPBACK_PEERS and _host_only(request.headers.get("host", "")) in LOOPBACK_HOSTS + + def _deny(status: int, detail: str): + response = JSONResponse({"detail": detail}, status_code=status) + for key, value in SECURITY_HEADERS.items(): + response.headers.setdefault(key, value) + return response + + @app.middleware("http") + async def guard(request: Request, call_next): + if _host_only(request.headers.get("host", "")) not in hosts: + return _deny(421, "unexpected Host header") + path = request.url.path + if request.method in MUTATING: + if request.headers.get("x-chronos-request") != "1": + return _deny(403, "missing X-Chronos-Request header") + origin = request.headers.get("origin") + if origin and _host_only(origin.split("://", 1)[-1].split("/", 1)[0]) not in hosts: + return _deny(403, "cross-origin request") + if path.startswith("/api/") and not (path in public or path.startswith("/api/i18n/")) \ + and not _authorized(request): + return _deny(401, "login required") + response = await call_next(request) + for key, value in SECURITY_HEADERS.items(): + response.headers.setdefault(key, value) + return response + + async def _body(request) -> dict: + try: + data = await request.json() + except ValueError: + raise HTTPException(status_code=400, detail="invalid JSON") from None + if not isinstance(data, dict): + raise HTTPException(status_code=400, detail="expected a JSON object") + return data + + def _check_data(paths: list[str], key: str) -> None: + """Every selected dataset path exists and at least one readable data file is behind them.""" + from chronos.workbench.files import data_files + if not paths: + return + missing = [p for p in paths if not os.path.exists(p)] + if missing: + _bad([key], "not_found", missing[0]) + if not data_files(paths): + _bad([key], "no_files") + + def _bad(keys, code: str = "invalid", message: str | None = None): + detail = {"error": code, "fields": list(keys)} + if message: + detail["message"] = message + raise HTTPException(status_code=422, detail=detail) + + # ── page, session, catalogs ──────────────────────────────────────── + @app.get("/") + def index(): + return FileResponse(os.path.join(STATIC_DIR, "index.html"), media_type="text/html; charset=utf-8") + + @app.get("/static/{name}") + def static(name: str): + ext = os.path.splitext(name)[1] + full = os.path.join(STATIC_DIR, name) + if not _STATIC_NAME.match(name) or not os.path.isfile(full): + raise HTTPException(status_code=404, detail="not found") + return FileResponse(full, media_type=STATIC_TYPES[ext]) + + @app.get("/health") + def health(): + return {"ok": True} + + @app.get("/api/session") + def session(request: Request): + from chronos.webui.native_dialog import backend + local = _is_local(request) + return {"auth_required": bool(api_key) and not (trust_local and local), "authenticated": _authorized(request), + "local_only": not api_key, "language": resolve_language(request.headers.get("accept-language")), + "access": "local" if local else "lan", "lan_enabled": bool(api_key) and bind_host not in LOOPBACK, + "native_dialog": bool(local and native_dialog and backend()), "max_upload_mb": max_upload_mb} + + @app.post("/api/login") + async def login(request: Request): + client = request.client.host if request.client else "?" + if not api_key: + return {"authenticated": True, "auth_required": False} + if sessions.blocked(client): + raise HTTPException(status_code=429, detail="too many failed logins; wait 5 minutes") + body = await _body(request) + key = body.get("key") + if not isinstance(key, str) or not _key_ok(key): + sessions.fail(client) + raise HTTPException(status_code=401, detail="invalid key") + response = JSONResponse({"authenticated": True, "auth_required": True}) + response.set_cookie(COOKIE, sessions.issue(), max_age=int(sessions.ttl_s), httponly=True, + samesite="strict", secure=request.url.scheme == "https", path="/") + return response + + @app.post("/api/logout") + def logout(request: Request): + sessions.revoke(request.cookies.get(COOKIE)) + response = JSONResponse({"authenticated": False}) + response.delete_cookie(COOKIE, path="/") + return response + + @app.get("/api/i18n/{lang}") + def i18n(lang: str): + try: + return {"language": lang, "messages": load_catalog(lang)} + except KeyError: + raise HTTPException(status_code=404, detail="unsupported language") from None + + @app.get("/api/schema") + def get_schema(): + return schema.schema_payload() + + @app.get("/api/system") + def system(): + return {**system_info(), "workdir": ui.workdir, "runs_dir": ui.runs_dir, + "auth_required": bool(api_key)} + + @app.get("/api/presets") + def presets(): + return read_presets() + + # ── config, planner, data, tokenizer ─────────────────────────────── + @app.post("/api/config/estimate") + async def estimate(request: Request): + clean, bad = sanitize_config((await _body(request)).get("config") or {}) + if bad: + _bad(bad) + try: + return _finite(config_estimate({**base_config(), **clean})) + except (ValueError, ZeroDivisionError) as exc: + raise HTTPException(status_code=422, detail={"error": "invalid", "message": str(exc)}) from exc + + @app.post("/api/config/save") + async def config_save(request: Request): + body = await _body(request) + clean, bad = sanitize_config(body.get("config") or {}) + if bad: + _bad(bad) + path = ui.path(body.get("path") or "./out/chronos_config.json") + if not path.endswith(".json"): + _bad(["path"]) + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + json.dump(clean, fh, indent=2, ensure_ascii=False) + return {"path": path} + + @app.post("/api/config/load") + async def config_load(request: Request): + path = ui.path((await _body(request)).get("path") or "") + if not path.endswith(".json") or not os.path.isfile(path): + raise HTTPException(status_code=404, detail="not_found") + with open(path, encoding="utf-8") as fh: + raw = json.load(fh) + if isinstance(raw, dict) and isinstance(raw.get("config"), dict): # a saved plan + raw = {**raw["config"], "stage_overrides": raw.get("stage_overrides") or raw["config"].get( + "stage_overrides")} + clean, bad = sanitize_config(raw if isinstance(raw, dict) else {}) + return {"config": clean, "ignored": bad, "path": path} + + @app.post("/api/plan") + async def plan(request: Request): + from chronos.model.planner import plan_training + values, bad = form_values(await _body(request), "plan") + if bad: + _bad(bad) + try: + req = plan_request(values) + except ValueError as exc: + _bad([str(exc)], code=str(exc)) + return _finite(_json_safe(plan_training(req))) + + # ── file picker: lists server folders (the UI is local; typed paths reach the same files) ── + @app.get("/api/fs/list") + def fs_list(path: str = "", kind: str = "any", hidden: bool = False): + from chronos.workbench.files import list_dir + if len(path) > schema.MAX_TEXT or "\x00" in path: + _bad(["path"]) + try: + return _json_safe(list_dir(ui.path(path) if path else ui.workdir, kind=kind, show_hidden=hidden)) + except ValueError: + _bad(["kind"]) + except FileNotFoundError: + raise HTTPException(status_code=404, detail={"error": "not_found"}) from None + + # The browser hides real file paths. On this machine the server opens the system dialog and the + # chosen paths are used in place; from another machine the files are uploaded instead. + @app.post("/api/fs/native") + async def fs_native(request: Request): + from chronos.webui.native_dialog import DialogBusy, backend, choose + if not _is_local(request): + raise HTTPException(status_code=403, detail={"error": "native_local_only"}) + if not native_dialog or backend() is None: + raise HTTPException(status_code=501, detail={"error": "native_unavailable"}) + body = await _body(request) + kind = body.get("kind") or "any" + if kind not in ("dir", "data", "checkpoint", "config", "any"): + _bad(["kind"]) + folder = kind == "dir" or bool(body.get("folder")) + start = ui.path(str(body.get("start") or "")) or ui.workdir + start = start if os.path.isdir(start) else os.path.dirname(start) + try: + paths = await run_in_threadpool(choose, folder=folder, multiple=bool(body.get("multiple")) and not + (folder and sys.platform.startswith("win")), start=start, + title=str(body.get("title") or "")[:200]) + except DialogBusy: + raise HTTPException(status_code=409, detail={"error": "dialog_open"}) from None + except RuntimeError: + raise HTTPException(status_code=501, detail={"error": "native_unavailable"}) from None + return {"paths": paths or [], "cancelled": not paths} + + @app.put("/api/fs/upload") + async def fs_upload(request: Request, kind: str, batch: str, name: str): + from chronos.workbench.files import upload_path + try: + dst = upload_path(ui.workdir, kind, batch, name) + except ValueError as exc: + raise HTTPException(status_code=422, detail={"error": f"upload_{exc}", "fields": ["upload"]}) from None + limit = max_upload_mb * 2**20 + declared = request.headers.get("content-length") + if declared and declared.isdigit() and int(declared) > limit: + raise HTTPException(status_code=413, detail={"error": "upload_too_large"}) + os.makedirs(os.path.dirname(dst), exist_ok=True) + free = shutil.disk_usage(os.path.dirname(dst)).free + reserve = max(UPLOAD_DISK_RESERVE, int(0.05 * shutil.disk_usage(os.path.dirname(dst)).total)) + expected = int(declared) if declared and declared.isdigit() else 0 + if free - expected < reserve: # keep room for checkpoints of running jobs + raise HTTPException(status_code=507, detail={"error": "upload_disk_full"}) + tmp = f"{dst}.part-{os.getpid()}-{time.monotonic_ns()}" + size, next_check = 0, 256 * 2**20 + try: + with open(tmp, "wb") as fh: + async for chunk in request.stream(): + size += len(chunk) + if size > limit: + raise HTTPException(status_code=413, detail={"error": "upload_too_large"}) + if size >= next_check: # re-check free space every 256 MiB written + next_check += 256 * 2**20 + if shutil.disk_usage(os.path.dirname(dst)).free < reserve: + raise HTTPException(status_code=507, detail={"error": "upload_disk_full"}) + fh.write(chunk) + os.replace(tmp, dst) + finally: + if os.path.exists(tmp): + os.remove(tmp) + from chronos.workbench.files import upload_root + return {"path": dst, "size": size, "root": upload_root(ui.workdir, kind, batch)} + + @app.post("/api/data/scan") + async def data_scan(request: Request): + from chronos.workbench.data import scan_dataset + body = await _body(request) + raw = body.get("path") or "" + if not isinstance(raw, (str, list)) or len(str(raw)) > schema.MAX_MULTI_TEXT: + _bad(["path"]) + paths = ui.paths(raw) + if not paths: + _bad(["path"], "no_paths") + try: + limit = max(100, min(int(body.get("max_records") or 50_000), 500_000)) + except (TypeError, ValueError): + _bad(["max_records"]) + summary, lengths, samples = await run_in_threadpool(scan_dataset, paths, max_records=limit) + if "error" in summary: + raise HTTPException(status_code=422, detail={"error": "no_files", "message": summary["error"]}) + return _finite(_json_safe({"summary": summary, "histogram": _histogram(lengths), + "samples": samples[:5]})) + + @app.post("/api/tokenizer/start") + async def tokenizer_start(request: Request): + from chronos.workbench.tokenizer import parse_paths + values, bad = form_values(await _body(request), "tok") + if bad: + _bad(bad) + fields = [f.strip() for f in (values["fields"] or "").split(",") if f.strip()] + result = ui.tokenizer.start([ui.path(p) for p in parse_paths(values["paths"])], ui.path(values["out_dir"]), + values["vocab_size"], values["per_source"], fields, values["temperature"]) + if result in ("no_paths", "no_out", "no_files"): + _bad(["tok.paths" if result != "no_out" else "tok.out_dir"], result) + if result == "busy": + raise HTTPException(status_code=409, detail="busy") + return tokenizer_state() + + @app.get("/api/tokenizer/state") + def tokenizer_state(): + from chronos.workbench.tokenizer import read_report, report_rows, report_summary + state, log = ui.tokenizer.snapshot() + report = read_report(ui.tokenizer.out_dir) if ui.tokenizer.out_dir and state == "done" else {} + return _finite(_json_safe({"state": state, "log": log, "error": ui.tokenizer.error or None, + "out_dir": ui.tokenizer.out_dir, "summary": report_summary(report), + "rows": report_rows((report or {}).get("heldout_eval") or {})})) + + # ── training jobs (subprocess worker, run dir read by the monitor) ── + @app.post("/api/train/start") + async def train_start(request: Request): + from chronos.workbench.lazy_probe import ProbeSweep, expand_variants + body = await _body(request) + cfg, bad = sanitize_config(body.get("config") or {}) + train, bad_t = form_values(body.get("train") or {}, "train") + probe, bad_p = form_values(body.get("probe") or {}, "probe") + if bad or bad_t or bad_p: + _bad(bad + bad_t + bad_p) + sweep = probe_sweep(probe) + try: + ProbeSweep.from_dict(sweep) + except ValueError as exc: + _bad(["probe"], "probe_invalid", str(exc)) + stamp = time.strftime("%Y%m%d_%H%M%S") + run_name = f"webui_{stamp}" + run_dir = os.path.join(ui.runs_dir, run_name) + os.makedirs(run_dir, exist_ok=True) + cfg = {**base_config(), **cfg} + cfg["save_dir"] = ui.path(cfg.get("save_dir") or "./out") + stages = [] + modes = [m for m in STAGE_ORDER if m in train["stages"]] # canonical order, not click order + first_data = ui.paths(train["data_path"]) + _check_data(first_data, "train.data_path") + for i, mode in enumerate(modes): + own = ui.paths(train.get(f"data_{mode}") or "") if i else [] + _check_data(own, f"train.data_{mode}") + data = own or first_data + stages.append({"mode": mode, "data_paths": data, "data_path": data[0] if len(data) == 1 else "", + "init_weight": ui.path(train["init_weight"]) if i == 0 else "", + "teacher_path": ui.path(train["teacher_path"]) if mode == "distill" else "", + "max_steps": train["max_steps"], "val_ratio": train["val_ratio"], + "sample_prompt": train["sample_prompt"]}) + job_spec = {"run_dir": run_dir, "config": cfg, "stages": stages, "backend": train["backend"], + "lazy_probe": {"enabled": bool(probe["enabled"]), "sweep": sweep}} + job_path = os.path.join(run_dir, "job.json") + with open(job_path, "w", encoding="utf-8") as fh: + json.dump(job_spec, fh, indent=2, ensure_ascii=False) + variants, total = expand_variants(ProbeSweep.from_dict(sweep), lambda p: p) + try: + job = ui.jobs.spawn("train", [sys.executable, "-m", "chronos.workbench.train_worker", job_path], + cwd=REPO_ROOT, title=" → ".join(modes), + meta={"run": run_name, "run_dir": run_dir, "stages": modes}) + except JobBusy: + raise HTTPException(status_code=409, detail="busy") from None + return {"job": job.view(log_lines=0), "run": run_name, + "probe": {"enabled": bool(probe["enabled"]), "variants": len(variants), "requested": total}} + + @app.post("/api/probe/preview") + async def probe_preview(request: Request): + """How many combinations a probe form expands to (before ``auto`` paging resolves).""" + from chronos.workbench.lazy_probe import ProbeSweep, expand_variants, variant_label + probe, bad = form_values(await _body(request), "probe") + if bad: + _bad(bad) + try: + variants, total = expand_variants(ProbeSweep.from_dict(probe_sweep(probe)), lambda p: p) + except ValueError as exc: + _bad(["probe"], "probe_invalid", str(exc)) + return {"variants": len(variants), "requested": total, "labels": [variant_label(v) for v in variants]} + + @app.get("/api/jobs") + def jobs(kind: str | None = None): + return {"jobs": ui.jobs.list(kind)} + + @app.get("/api/jobs/{job_id}") + def job(job_id: str, log: int = 200): + try: + return _finite(_json_safe(ui.jobs.get(job_id).view(log_lines=max(0, min(int(log), 2000))))) + except KeyError: + raise HTTPException(status_code=404, detail="not_found") from None + + @app.post("/api/jobs/{job_id}/stop") + def job_stop(job_id: str): + try: + return ui.jobs.stop(job_id) + except KeyError: + raise HTTPException(status_code=404, detail="not_found") from None + + # ── run monitor (any run under runs_dir, including the CLI long runs) ── + def _mon(fn, *args, **kwargs): + try: + return _finite(_json_safe(fn(*args, **kwargs))) + except KeyError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + @app.get("/api/monitor/runs") + def monitor_runs(): + return {"root": ui.monitor.root, "runs": _mon(ui.monitor.runs)} + + @app.get("/api/monitor/state") + def monitor_state(run: str | None = None, history: int = 1): + return _mon(ui.monitor.snapshot, run, history=bool(history)) + + @app.get("/api/monitor/metrics") + def monitor_metrics(run: str | None = None, phase: str | None = None, since: int = 0, limit: int = 5000): + return {"rows": _mon(ui.monitor.metrics, run, phase, since, max(1, min(int(limit), 20000)))} + + @app.get("/api/monitor/log") + def monitor_log(run: str | None = None, bytes: int = 16000): # noqa: A002 - query name + return {"text": _mon(ui.monitor.log_tail, run, max(1000, min(int(bytes), 400_000)))} + + @app.get("/api/monitor/parameters") + def monitor_parameters(run: str | None = None): + return _mon(ui.monitor.training_parameters, run) + + @app.get("/api/monitor/lazy_probes") + def monitor_probes(run: str | None = None, limit: int = 200): + return {"probes": _mon(ui.monitor.lazy_probes, run, max(1, min(int(limit), 1000)))} + + # ── inference (thread job; architecture comes from the checkpoint sidecar) ── + @app.post("/api/infer") + async def infer(request: Request): + from chronos.workbench.inference import generate_api + body = await _body(request) + cfg, bad = sanitize_config(body.get("config") or {}) + values, bad_i = form_values(body.get("infer") or {}, "infer") + prompt = body.get("prompt") + if bad or bad_i: + _bad(bad + bad_i) + if not isinstance(prompt, str) or not prompt.strip() or len(prompt) > 32_000: + _bad(["prompt"]) + model_path = ui.path(values["model_path"]) + if not model_path or not os.path.exists(model_path): + _bad(["infer.model_path"], "not_found") + + def work(job): + job.log(f"{values['mode']} · {values['backend']} · {os.path.basename(model_path)}") + out = generate_api({**base_config(), **cfg}, values["backend"], values["mode"], model_path, + prompt, values["max_tokens"], values["temperature"], values["raw_prompt"], + values["ram_load_ratio"], values["sweep_ram_ratios"], values["miss_policy"], + values["speculative"], values["draft_depth"], values["gamma"], + values["residency_planner"], values["paging"], values["residency_decay"], + values["decode_wait_s"], values["verify_wait_s"]) + if out.get("traceback"): + job.log(out.pop("traceback")) + if out.get("error"): + raise RuntimeError(out["error"]) + return _finite(_json_safe(out)) + + try: + job = ui.jobs.run("inference", work, title=os.path.basename(model_path)) + except JobBusy: + raise HTTPException(status_code=409, detail="busy") from None + return {"job": job.view(log_lines=0)} + + # ── export ────────────────────────────────────────────────────────── + @app.post("/api/export") + async def export(request: Request): + values, bad = form_values(await _body(request), "export") + if bad: + _bad(bad) + model_path = ui.path(values["model_path"]) + if not model_path or not os.path.isfile(model_path): + _bad(["export.model_path"], "not_found") + if not values["output_dir"]: + _bad(["export.output_dir"], "no_out") + + calibration = ui.paths(values["calibration_data_path"]) + if calibration and not all(os.path.exists(p) for p in calibration): + _bad(["export.calibration_data_path"], "not_found") + + def work(job): + from chronos.export import export_checkpoint, format_export_report + from chronos.workbench.files import merge_datasets + calibration_path = (merge_datasets(calibration, os.path.join(ui.workdir, "out", "merged_data"), + log=job.log) if calibration else None) + results = export_checkpoint( + model_path, ui.path(values["output_dir"]), formats=values["formats"], + config_path=ui.path(values["config_path"]) or None, + expert_cache_dir=ui.path(values["expert_cache_dir"]) or None, + copy_expert_cache=values["copy_cache"], auto_cluster=values["auto_cluster"], + calibration_data_path=calibration_path, + cluster_max_batches=values["cluster_max_batches"], cluster_batch_size=values["cluster_batch_size"], + cluster_max_seq_len=values["cluster_max_seq_len"], cluster_device=values["cluster_device"]) + job.log(format_export_report(results)) + return [{"format": r.format, "path": r.path, "size_mb": round(r.bytes / 2**20, 2), + "tensors": r.tensors} for r in results] + + try: + job = ui.jobs.run("export", work, title=os.path.basename(model_path)) + except JobBusy: + raise HTTPException(status_code=409, detail="busy") from None + return {"job": job.view(log_lines=0)} + + # ── Ollama / OpenAI-compatible server child ───────────────────────── + @app.post("/api/serve/start") + async def serve_start(request: Request): + from chronos.workbench.serve import build_serve_command + values, bad = form_values(await _body(request), "serve") + if bad: + _bad(bad) + serve_key = os.environ.get("CHRONOS_API_KEY", "") if values["allow_remote"] else "" + try: + argv, env = build_serve_command( + models_dir=ui.path(values["models_dir"]), host=values["host"], port=values["port"], + device=values["device"], dtype=values["dtype"], lazy_experts=values["lazy_experts"], + ram_expert_fraction=values["ram_expert_fraction"], allow_remote=values["allow_remote"], + api_key=serve_key, kv_cache_dtype=values["kv_cache_dtype"] or "", + mla_decode=values["mla_decode"] or "", router_scope=values["router_scope"] or "", + max_resident_experts=values["max_resident_experts"], ram_budget_mb=values["ram_budget_mb"], + monitor_runs_dir=ui.runs_dir, paging=values["paging"], mla_cache_rope=values["mla_cache_rope"]) + except ValueError as exc: + _bad([f"serve.{exc}" if str(exc) in ("models_dir", "port", "paging") else "serve.host"], str(exc)) + if ui.serve.start(argv, env, values["host"], values["port"]) == "busy": + raise HTTPException(status_code=409, detail="busy") + return serve_state() + + @app.post("/api/serve/stop") + def serve_stop(): + ui.serve.stop() + return serve_state() + + @app.get("/api/serve/state") + def serve_state(): + running, log = ui.serve.snapshot() + url = ui.serve.url if running else "" + addr = url.split("://", 1)[-1] if url else "" + return {"running": running, "url": url, "addr": addr, "auth": ui.serve.auth if running else False, + "log": log} + + # ── benchmark (MiniMind vs Chronos script in the repo) ────────────── + bench_script = os.path.join(REPO_ROOT, "benchmark_compare.py") + bench_results = os.path.join(REPO_ROOT, "benchmark_results.json") + + def _bench_read(_job=None): + if not os.path.isfile(bench_results): + return None + with open(bench_results, encoding="utf-8") as fh: + return _finite(json.load(fh)) + + @app.post("/api/benchmark/start") + def benchmark_start(): + if not os.path.isfile(bench_script): + raise HTTPException(status_code=404, detail="not_found") + try: + job = ui.jobs.spawn("benchmark", [sys.executable, bench_script], cwd=REPO_ROOT, + title="benchmark_compare", on_exit=_bench_read) + except JobBusy: + raise HTTPException(status_code=409, detail="busy") from None + return {"job": job.view(log_lines=0)} + + @app.get("/api/benchmark/results") + def benchmark_results(): + return {"available": os.path.isfile(bench_script), "results": _bench_read(), + "metrics": [{"key": k, "label": label, "lower_is_better": low} for k, label, low in BENCH_METRICS]} + + # ── auto-tune (Optuna TPE; one search per server) ─────────────────── + @app.post("/api/autotune/start") + async def autotune_start(request: Request): + from chronos.tuning.chronos_auto_tuner import ChronosAutoTuner, ChronosSearchSpaceConfig + body = await _body(request) + values, bad = form_values(body.get("tune") or {}, "tune") + cfg, bad_c = sanitize_config(body.get("config") or {}) + if bad or bad_c: + _bad(bad + bad_c) + from chronos.workbench.files import data_files + data = ui.paths(values["data_path"]) + if not data or not all(os.path.exists(p) for p in data): + _bad(["tune.data_path"], "not_found") + if not data_files(data): + _bad(["tune.data_path"], "no_files") + data_path = data[0] + if ui.tuner is not None and ui.tuner.is_running(): + raise HTTPException(status_code=409, detail="busy") + space = ChronosSearchSpaceConfig(**{f"tune_{name}": name in values["search"] + for name in schema.FIELDS["tune.search"].choices}) + cfg = {**base_config(), **cfg} + out_dir = os.path.join(ui.workdir, "auto_tune_cache") + os.makedirs(out_dir, exist_ok=True) + ui.tuner, ui.tune_log = ChronosAutoTuner(), [] + ui.tuner.start(model_id=os.path.join(ui.path(cfg.get("save_dir") or "./out"), + f"chronos_{cfg.get('hidden_size', 256)}_moe.pth"), + dataset_path=data_path, train_ratio=0.95, prompt_template="alpaca", think_mode="keep", + search_space=space, n_trials=values["n_trials"], probe_steps=values["probe_steps"], + output_dir=out_dir, seed=42) + return autotune_state() + + @app.post("/api/autotune/stop") + def autotune_stop(): + if ui.tuner is not None and ui.tuner.is_running(): + ui.tuner.stop() + return autotune_state() + + @app.get("/api/autotune/state") + def autotune_state(): + if ui.tuner is None: + return {"state": "idle", "log": [], "trials": [], "best": {}} + for event in ui.tuner.poll(): + if event.get("type") == "log": + ui.tune_log.append(str(event.get("line", ""))) + del ui.tune_log[:-500] + trials = [{"trial": tr.trial_number + 1, "params": tr.params, "loss": tr.train_loss, + "duration_s": tr.duration_s, "status": tr.status} for tr in list(ui.tuner.trials)[-200:]] + best = dict(ui.tuner.result.best_params) if ui.tuner.result and ui.tuner.result.best_params else {} + return _finite(_json_safe({"state": ui.tuner.status, "log": ui.tune_log[-200:], + "trials": trials, "best": best})) + + # ── live IO / cache metrics of this process ───────────────────────── + @app.get("/api/io") + def io_metrics(samples: int = 120): + from chronos.runtime.metrics import bus + keep = max(1, min(int(samples), 2000)) + return _finite({"series": {name: [[round(t, 3), v] for t, v in points[-keep:]] + for name, points in sorted(bus.snapshot().items())}}) + + return app + + +BENCH_METRICS = (("params_m", "bench.params", False), ("ram_model_gb", "bench.ram_model", True), + ("ram_train_gb", "bench.ram_train", True), ("ram_infer_gb", "bench.ram_infer", True), + ("tokens_per_sec", "bench.tps", False)) + + +def _histogram(lengths, bins: int = 24) -> dict: + """Record-length histogram (characters) for the Data page; log-spaced bins.""" + import math + values = [int(v) for v in lengths or () if v is not None and v >= 0] + if not values: + return {"edges": [], "counts": []} + lo, hi = max(1, min(values)), max(values) + if hi <= lo: + return {"edges": [lo, lo + 1], "counts": [len(values)]} + edges = sorted({int(round(math.exp(math.log(lo) + (math.log(hi + 1) - math.log(lo)) * i / bins))) + for i in range(bins + 1)}) + counts = [0] * (len(edges) - 1) + for v in values: + idx = min(len(counts) - 1, max(0, next((i for i, e in enumerate(edges[1:]) if v < e), len(counts) - 1))) + counts[idx] += 1 + return {"edges": edges, "counts": counts} + + +def main(argv=None) -> int: + import argparse + argv = list(sys.argv[1:] if argv is None else argv) + if "--legacy" in argv: # the Gradio console, unchanged + argv.remove("--legacy") + from chronos.app import main as legacy_main + sys.argv = ["chronos-ui"] + argv + return legacy_main() or 0 + parser = argparse.ArgumentParser(prog="chronos webui", description="Chronos Web UI (FastAPI + single-page app)") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=7860) + parser.add_argument("--workdir", default=".", help="base for relative paths (data, out/, runs/)") + parser.add_argument("--runs-dir", default="runs", help="run folders shown in the monitor (relative to --workdir)") + parser.add_argument("--lan", action="store_true", + help="also open the UI to this machine's local network (off by default). Browsers on other " + "machines log in with CHRONOS_UI_KEY, or with a key printed here when none is set; " + "this machine needs no login. Not for use behind a reverse proxy") + parser.add_argument("--allow-remote", action="store_true", + help="listen beyond loopback; needs a key in CHRONOS_UI_KEY or CHRONOS_API_KEY") + parser.add_argument("--max-upload-mb", type=int, default=MAX_UPLOAD_MB, + help="largest single file a LAN browser may upload (MiB)") + parser.add_argument("--no-native-dialog", action="store_true", + help="never open the system file dialog on this machine (e.g. when run over SSH)") + parser.add_argument("--allowed-host", action="append", default=[], + help="extra Host header value to accept (e.g. a reverse-proxy name); repeatable") + parser.add_argument("--session-hours", type=float, default=12.0) + parser.add_argument("--ssl-certfile", default=None) + parser.add_argument("--ssl-keyfile", default=None) + parser.add_argument("--legacy", action="store_true", help="start the old Gradio console instead") + parser.add_argument("--cpu-threads", default="auto", help="PyTorch/BLAS CPU threads for in-process work") + parser.add_argument("--cpu-budget-percent", type=float, default=100, + help="CPU physical-core budget used when --cpu-threads=auto") + parser.add_argument("--share", action="store_true", help=argparse.SUPPRESS) + args = parser.parse_args(argv) + if args.share: + parser.error("--share (a public Gradio tunnel) is not offered by this UI; use --allow-remote with a key " + "and TLS, or --legacy --share for the old Gradio console") + from chronos.trainer.device_utils import configure_cpu_threads + configure_cpu_threads(args.cpu_threads, budget_percent=args.cpu_budget_percent) + api_key = os.environ.get("CHRONOS_UI_KEY") or os.environ.get("CHRONOS_API_KEY") or None + extra, generated, lan = list(args.allowed_host), False, [] + if args.lan: + import secrets + import socket + if args.host in LOOPBACK: + args.host = "0.0.0.0" + args.allow_remote = True + if not api_key: + api_key, generated = secrets.token_urlsafe(18), True + lan = lan_addresses() + name = socket.gethostname().lower() + extra += lan + [name, f"{name.split('.')[0]}.local"] + if args.max_upload_mb < 1: + parser.error("--max-upload-mb must be at least 1") + problem = check_bind(args.host, args.allow_remote, api_key) + if problem: + parser.error("binding a non-loopback host requires --allow-remote and a key in CHRONOS_UI_KEY " + "or CHRONOS_API_KEY") + import uvicorn + app = create_app(args.workdir, runs_dir=args.runs_dir, api_key=api_key, + bind_host=args.host, extra_hosts=tuple(extra), session_hours=args.session_hours, + max_upload_mb=args.max_upload_mb, trust_local=args.lan, native_dialog=not args.no_native_dialog) + scheme = "https" if args.ssl_certfile else "http" + if args.lan: + print(f"[chronos webui] this machine: {scheme}://127.0.0.1:{args.port}/ (no login) " + f"workdir={os.path.abspath(args.workdir)}", flush=True) + for addr in lan or [""]: + print(f"[chronos webui] LAN: {scheme}://{addr}:{args.port}/ (login required)", flush=True) + if generated: + print(f"[chronos webui] LAN login key (generated for this run; set CHRONOS_UI_KEY to choose one): " + f"{api_key}", flush=True) + else: + shown = f"[{args.host}]" if ":" in args.host else args.host + auth = "login required" if api_key else "no auth (loopback only)" + print(f"[chronos webui] {scheme}://{shown}:{args.port}/ workdir={os.path.abspath(args.workdir)} {auth}", + flush=True) + if api_key and scheme == "http" and args.host not in LOOPBACK: + print("[chronos webui] warning: the login key crosses the network in clear text; use --ssl-certfile/" + "--ssl-keyfile or a TLS reverse proxy", flush=True) + uvicorn.run(app, host=args.host, port=args.port, log_level="warning", + ssl_certfile=args.ssl_certfile, ssl_keyfile=args.ssl_keyfile) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/webui/static/app.css b/chronos/webui/static/app.css new file mode 100644 index 0000000..3e890e8 --- /dev/null +++ b/chronos/webui/static/app.css @@ -0,0 +1,252 @@ +/* Chronos Web UI. Tokens first; light and dark are each selected, not flipped. */ +:root { + color-scheme: light; + --page: #f9f9f7; --surface: #fcfcfb; --surface-2: #f3f2ef; --line: #e1e0d9; --axis: #c3c2b7; + --ink: #0b0b0b; --ink-2: #52514e; --muted: #6f6d68; --accent: #2a78d6; --accent-ink: #ffffff; + --accent-soft: #e4effc; --focus: #2a78d6; + --s1: #2a78d6; --s2: #eb6834; --s3: #1baf7a; --s4: #eda100; --s5: #e87ba4; --s6: #008300; --s7: #6250d6; --s8: #e34948; + --good: #0ca30c; --warning: #fab219; --serious: #ec835a; --critical: #d03b3b; + --seq-lo: #cde2fb; --seq-hi: #104281; + --radius: 10px; --shadow: 0 1px 2px rgb(0 0 0 / 6%), 0 4px 16px rgb(0 0 0 / 5%); + --font: system-ui, -apple-system, "Segoe UI", "PingFang SC", "Hiragino Sans", "Noto Sans CJK SC", "Microsoft YaHei", sans-serif; + --mono: ui-monospace, "SF Mono", Menlo, Consolas, monospace; +} +@media (prefers-color-scheme: dark) { + :root:where(:not([data-theme="light"])) { + color-scheme: dark; + --page: #0d0d0d; --surface: #1a1a19; --surface-2: #222220; --line: #2c2c2a; --axis: #383835; + --ink: #f0efec; --ink-2: #c3c2b7; --muted: #a09e97; --accent: #3987e5; --accent-ink: #ffffff; + --accent-soft: #1b2b40; --focus: #6da7ec; + --s1: #3987e5; --s2: #d95926; --s3: #199e70; --s4: #c98500; --s5: #d55181; --s6: #008300; --s7: #9085e9; --s8: #e66767; + --seq-lo: #1b2b40; --seq-hi: #86b6ef; + --shadow: 0 1px 2px rgb(0 0 0 / 40%), 0 4px 16px rgb(0 0 0 / 30%); + } +} +:root[data-theme="dark"] { + color-scheme: dark; + --page: #0d0d0d; --surface: #1a1a19; --surface-2: #222220; --line: #2c2c2a; --axis: #383835; + --ink: #f0efec; --ink-2: #c3c2b7; --muted: #a09e97; --accent: #3987e5; --accent-ink: #ffffff; + --accent-soft: #1b2b40; --focus: #6da7ec; + --s1: #3987e5; --s2: #d95926; --s3: #199e70; --s4: #c98500; --s5: #d55181; --s6: #008300; --s7: #9085e9; --s8: #e66767; + --seq-lo: #1b2b40; --seq-hi: #86b6ef; + --shadow: 0 1px 2px rgb(0 0 0 / 40%), 0 4px 16px rgb(0 0 0 / 30%); +} + +* { box-sizing: border-box; } +html, body { margin: 0; } +body { font: 14px/1.5 var(--font); color: var(--ink); background: var(--page); min-height: 100vh; } +h1, h2, h3, h4 { margin: 0; line-height: 1.25; font-weight: 600; } +h1 { font-size: 17px; } h2 { font-size: 18px; } h3 { font-size: 15px; } h4 { font-size: 13px; color: var(--ink-2); } +p { margin: 0; } +a { color: var(--accent); } +code, pre, .mono { font-family: var(--mono); font-size: 12.5px; } +:focus-visible { outline: 2px solid var(--focus); outline-offset: 2px; } +.skip { position: absolute; left: -9999px; } +.skip:focus { left: 12px; top: 8px; z-index: 50; background: var(--surface); padding: 6px 10px; border-radius: 6px; } +[hidden] { display: none !important; } + +/* top bar + shell */ +.top { display: flex; align-items: center; justify-content: space-between; gap: 16px; padding: 10px 20px; + background: var(--surface); border-bottom: 1px solid var(--line); position: sticky; top: 0; z-index: 20; } +.brand { display: flex; align-items: center; gap: 10px; min-width: 0; } +.brand .sub { color: var(--muted); font-size: 12px; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; } +.top-tools { display: flex; align-items: center; gap: 12px; flex-wrap: wrap; justify-content: flex-end; } +.shell { display: grid; grid-template-columns: 200px minmax(0, 1fr); min-height: calc(100vh - 57px); } +.nav { padding: 14px 10px; border-right: 1px solid var(--line); background: var(--surface); position: sticky; + top: 57px; height: calc(100vh - 57px); overflow-y: auto; } +.nav a { display: flex; align-items: center; gap: 10px; padding: 8px 12px; border-radius: 8px; color: var(--ink-2); + text-decoration: none; font-weight: 500; } +.nav a:hover { background: var(--surface-2); color: var(--ink); } +.nav a[aria-current="page"] { background: var(--accent-soft); color: var(--ink); } +.nav a[aria-current="page"]::before { content: ""; width: 3px; height: 16px; border-radius: 2px; background: var(--accent); margin-left: -6px; } +.nav .dot { width: 7px; height: 7px; border-radius: 50%; background: var(--good); margin-left: auto; } +main { padding: 22px 28px 60px; max-width: 1360px; width: 100%; outline: none; } +@media (max-width: 860px) { + .shell { grid-template-columns: 1fr; } + .nav { position: static; height: auto; display: flex; flex-wrap: wrap; gap: 4px; border-right: 0; border-bottom: 1px solid var(--line); } + main { padding: 16px; } + .brand .sub { display: none; } +} +/* page furniture */ +.page-head { display: flex; align-items: flex-end; justify-content: space-between; gap: 16px; margin-bottom: 18px; flex-wrap: wrap; } +.page-head p { color: var(--ink-2); margin-top: 4px; max-width: 72ch; } +.card { background: var(--surface); border: 1px solid var(--line); border-radius: var(--radius); padding: 16px 18px; box-shadow: var(--shadow); } +.card + .card, .stack > * + * { margin-top: 16px; } +.card > h3 { margin-bottom: 12px; display: flex; align-items: center; gap: 8px; } +.grid { display: grid; gap: 16px; } +.grid.two { grid-template-columns: repeat(auto-fit, minmax(380px, 1fr)); } +.grid.three { grid-template-columns: repeat(auto-fit, minmax(280px, 1fr)); } +.row { display: flex; gap: 10px; align-items: center; flex-wrap: wrap; } +.row.end { justify-content: flex-end; } +.row.between { justify-content: space-between; } +.mt { margin-top: 12px; } +.note { color: var(--muted); font-size: 12.5px; } +.error { color: var(--critical); } +.badge { display: inline-flex; align-items: center; gap: 6px; padding: 3px 10px; border-radius: 999px; font-size: 12px; + background: var(--surface-2); color: var(--ink-2); border: 1px solid var(--line); white-space: nowrap; } +.badge::before { content: ""; width: 7px; height: 7px; border-radius: 50%; background: var(--muted); } +.badge.good::before, .badge.running::before, .badge.finished::before, .badge.done::before { background: var(--good); } +.badge.warning::before, .badge.paused::before, .badge.stale::before, .badge.stopping::before { background: var(--warning); } +.badge.failed::before, .badge.error::before { background: var(--critical); } +.badge.running::before { animation: pulse 1.6s ease-in-out infinite; } +@keyframes pulse { 50% { opacity: .35; } } +@media (prefers-reduced-motion: reduce) { .badge.running::before { animation: none; } } +.pill { font-size: 11.5px; padding: 1px 8px; border-radius: 999px; background: var(--accent-soft); color: var(--ink); } + +/* controls */ +button, .button { font: inherit; font-weight: 500; border-radius: 8px; padding: 7px 14px; cursor: pointer; + border: 1px solid var(--line); background: var(--surface); color: var(--ink); min-height: 34px; } +button:hover { background: var(--surface-2); } +button.primary { background: var(--accent); color: var(--accent-ink); border-color: var(--accent); } +button.primary:hover { filter: brightness(1.08); } +button.danger { color: var(--critical); } +button.ghost { background: transparent; border-color: transparent; } +button:disabled { opacity: .5; cursor: not-allowed; } +input, select, textarea { font: inherit; color: var(--ink); background: var(--surface); border: 1px solid var(--axis); + border-radius: 7px; padding: 6px 9px; min-height: 34px; width: 100%; } +textarea { min-height: 84px; resize: vertical; font-family: var(--font); } +input[type="checkbox"] { width: 18px; height: 18px; min-height: 0; accent-color: var(--accent); margin: 0; } +input[type="range"] { padding: 0; accent-color: var(--accent); } +input:invalid, .field.bad input, .field.bad select { border-color: var(--critical); } +label.compact { display: inline-flex; align-items: center; gap: 6px; font-size: 12.5px; color: var(--ink-2); } +label.compact select { width: auto; min-height: 30px; padding: 3px 6px; } + +/* schema-driven forms */ +.fields { display: grid; grid-template-columns: repeat(auto-fill, minmax(230px, 1fr)); gap: 12px 16px; } +.field { display: flex; flex-direction: column; gap: 5px; min-width: 0; } +.field > .label { display: flex; align-items: center; gap: 6px; font-size: 12.5px; font-weight: 500; color: var(--ink-2); } +.field > .label .name { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.field .help-dot { flex: none; width: 16px; height: 16px; border-radius: 50%; border: 1px solid var(--axis); color: var(--muted); + font-size: 10.5px; line-height: 14px; text-align: center; cursor: help; background: transparent; padding: 0; min-height: 0; } +.field .help-dot:hover, .field .help-dot:focus-visible { border-color: var(--accent); color: var(--accent); } +.field .rec { margin-left: auto; font-size: 10.5px; color: var(--muted); white-space: nowrap; } +.field.changed .rec::after { content: " •"; color: var(--accent); } +.field.bool { flex-direction: row; align-items: center; gap: 10px; padding-top: 18px; } +.field.bool > .label { order: 2; } +.field.wide { grid-column: 1 / -1; } +.chips { display: flex; flex-wrap: wrap; gap: 6px; } +.chips label { display: inline-flex; align-items: center; gap: 6px; padding: 4px 10px; border: 1px solid var(--axis); + border-radius: 999px; cursor: pointer; font-size: 12.5px; user-select: none; } +.chips input { position: absolute; opacity: 0; width: 1px; height: 1px; } +.chips label:has(input:checked) { background: var(--accent-soft); border-color: var(--accent); color: var(--ink); } +.chips label:has(input:focus-visible) { outline: 2px solid var(--focus); outline-offset: 2px; } +details.group { border-top: 1px solid var(--line); padding: 12px 0 4px; } +details.group:first-of-type { border-top: 0; padding-top: 0; } +details.group > summary { cursor: pointer; font-weight: 600; font-size: 14px; list-style: none; display: flex; align-items: center; gap: 8px; margin-bottom: 10px; } +details.group > summary::before { content: "▸"; color: var(--muted); transition: transform .15s; } +details.group[open] > summary::before { transform: rotate(90deg); } +.adv-toggle { margin-top: 10px; } +/* path fields: input + browse button; several paths are one per line */ +.path-box { display: flex; flex-direction: column; gap: 4px; min-width: 0; } +.path-row { display: flex; flex-wrap: wrap; gap: 6px; align-items: flex-start; min-width: 0; } +.path-row > input, .path-row > textarea { flex: 1 1 220px; min-width: 0; } +.upload-note { margin: 0; overflow-wrap: anywhere; } +.path-row > .browse, .path-row > .upload { flex: none; white-space: nowrap; } +textarea.paths { min-height: 34px; resize: vertical; font-family: var(--mono); font-size: 12.5px; line-height: 1.5; white-space: pre; overflow-x: auto; } +.path-note { margin: 0; } +.fields { container-type: inline-size; } +@container (min-width: 500px) { .field.path { grid-column: span 2; } } +/* file picker dialog */ +dialog.picker { width: min(760px, 94vw); padding: 0; display: flex; flex-direction: column; max-height: min(640px, 88vh); } +dialog.picker:not([open]) { display: none; } +.picker-head { display: flex; align-items: center; justify-content: space-between; padding: 14px 18px 8px; } +.picker-head button { font-size: 18px; line-height: 1; padding: 4px 10px; } +.picker-bar { display: flex; gap: 6px; padding: 0 18px; align-items: center; } +.picker-bar button { flex: none; white-space: nowrap; } +.picker-bar input { font-family: var(--mono); font-size: 12.5px; } +.picker-roots { display: flex; flex-wrap: wrap; gap: 6px; padding: 10px 18px 8px; } +.picker-roots button { min-height: 28px; padding: 2px 10px; font-size: 12.5px; } +.picker-list { list-style: none; margin: 0 18px; padding: 4px; border: 1px solid var(--line); border-radius: 8px; + overflow-y: auto; flex: 1; min-height: 220px; background: var(--page); } +.picker-list:focus-visible { outline: 2px solid var(--focus); outline-offset: 1px; } +.picker-item { display: grid; grid-template-columns: 18px minmax(0, 1fr) auto auto; align-items: center; gap: 4px 10px; + padding: 5px 8px; border-radius: 6px; cursor: default; user-select: none; } +.picker-item:hover { background: var(--surface-2); } +.picker-item.cursor { box-shadow: inset 0 0 0 1px var(--axis); } +.picker-item[aria-selected="true"] { background: var(--accent-soft); box-shadow: inset 0 0 0 1px var(--accent); } +.picker-item.dir { cursor: pointer; } +.picker-name { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.picker-meta { color: var(--muted); font-size: 11.5px; font-variant-numeric: tabular-nums; white-space: nowrap; } +.picker-icon { fill: none; stroke: var(--muted); stroke-width: 1.2; stroke-linejoin: round; } +.picker-item.dir .picker-icon { stroke: var(--accent); } +.picker-empty { color: var(--muted); padding: 18px; text-align: center; } +.picker-foot { display: grid; grid-template-columns: auto minmax(0, 1fr) auto; align-items: center; gap: 12px; padding: 10px 18px 16px; } +.picker-foot .note { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +dialog.picker label.compact { flex-direction: row; white-space: nowrap; } +@media (max-width: 640px) { .picker-foot { grid-template-columns: 1fr; } .picker-item { grid-template-columns: 18px minmax(0, 1fr) auto; } .picker-item .picker-meta:last-child { display: none; } } + +/* tooltip (one element, positioned by app.js; never gates information) */ +.tooltip { position: fixed; z-index: 60; max-width: 360px; background: var(--surface); color: var(--ink); + border: 1px solid var(--line); border-radius: 8px; box-shadow: var(--shadow); padding: 10px 12px; font-size: 12.5px; pointer-events: none; } +.tooltip .t-title { font-weight: 600; margin-bottom: 4px; } +.tooltip .t-help { color: var(--ink-2); } +.tooltip dl { display: grid; grid-template-columns: auto 1fr; gap: 2px 10px; margin: 8px 0 0; } +.tooltip dt { color: var(--muted); } +.tooltip dd { margin: 0; font-family: var(--mono); font-size: 12px; overflow-wrap: anywhere; } +.tooltip .t-row { display: flex; align-items: center; gap: 8px; } +.tooltip .t-row b { font-variant-numeric: tabular-nums; } +.tooltip .key { width: 12px; height: 2px; border-radius: 1px; flex: none; } + +/* stat tiles */ +.tiles { display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; } +.tile { background: var(--surface); border: 1px solid var(--line); border-radius: var(--radius); padding: 12px 14px; } +.tile .k { font-size: 12px; color: var(--muted); } +.tile .v { font-size: 22px; font-weight: 600; font-variant-numeric: tabular-nums; margin-top: 2px; } +.tile .s { font-size: 11.5px; color: var(--muted); margin-top: 2px; } + +/* tables */ +.table-wrap { overflow-x: auto; border: 1px solid var(--line); border-radius: 8px; } +table { border-collapse: collapse; width: 100%; font-size: 12.5px; } +th, td { padding: 6px 10px; text-align: left; border-bottom: 1px solid var(--line); white-space: nowrap; } +th { font-weight: 600; color: var(--ink-2); background: var(--surface-2); position: sticky; top: 0; } +td.num, th.num { text-align: right; font-variant-numeric: tabular-nums; } +tr:last-child td { border-bottom: 0; } +tbody tr:hover { background: var(--surface-2); } + +/* charts */ +.chart { position: relative; } +.chart svg { display: block; width: 100%; height: auto; overflow: visible; } +.chart .grid-line { stroke: var(--line); stroke-width: 1; } +.chart .axis-line { stroke: var(--axis); stroke-width: 1; } +.chart text { fill: var(--muted); font-size: 11px; font-family: var(--font); } +.chart .series { fill: none; stroke-width: 2; stroke-linejoin: round; stroke-linecap: round; } +.chart .end-dot { stroke: var(--surface); stroke-width: 2; } +.chart .crosshair { stroke: var(--axis); stroke-width: 1; } +.chart .bar:hover, .chart .bar:focus { filter: brightness(1.12); outline: none; } +.chart-head { display: flex; align-items: baseline; justify-content: space-between; gap: 8px; margin-bottom: 6px; flex-wrap: wrap; } +.legend { display: flex; flex-wrap: wrap; gap: 4px 14px; font-size: 12px; color: var(--ink-2); } +.legend span { display: inline-flex; align-items: center; gap: 6px; } +.legend i { display: inline-block; width: 14px; height: 2px; border-radius: 1px; } +.legend i.box { width: 10px; height: 10px; border-radius: 2px; } +.empty { color: var(--muted); padding: 24px; text-align: center; border: 1px dashed var(--line); border-radius: 8px; } + +/* logs, outputs */ +pre.log { background: var(--surface-2); border: 1px solid var(--line); border-radius: 8px; padding: 10px 12px; + max-height: 340px; overflow: auto; white-space: pre-wrap; word-break: break-word; margin: 0; } +.output { white-space: pre-wrap; background: var(--surface-2); border-radius: 8px; padding: 12px; min-height: 64px; border: 1px solid var(--line); } +.progress { height: 8px; background: var(--surface-2); border-radius: 4px; overflow: hidden; } +.progress > div { height: 100%; background: var(--accent); border-radius: 4px; transition: width .4s; } +.phase-row { display: grid; grid-template-columns: minmax(96px, 160px) 1fr minmax(120px, auto); align-items: center; gap: 12px; } +.phase-row + .phase-row { margin-top: 8px; } +.phase-name { font-weight: 500; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.phase-row .note { text-align: right; font-variant-numeric: tabular-nums; } +@media (max-width: 640px) { .phase-row { grid-template-columns: 1fr; gap: 4px; } .phase-row .note { text-align: left; } } +.chart-block { min-width: 0; } +.sample { padding: 10px 0; display: grid; gap: 4px; overflow-wrap: anywhere; } +.sample + .sample { border-top: 1px solid var(--line); } +.tabs { display: flex; gap: 4px; border-bottom: 1px solid var(--line); margin-bottom: 12px; } +.tabs button { border: 0; border-bottom: 2px solid transparent; border-radius: 0; background: transparent; color: var(--ink-2); } +.tabs button[aria-selected="true"] { border-bottom-color: var(--accent); color: var(--ink); } + +/* toasts, dialog */ +.toasts { position: fixed; right: 18px; bottom: 18px; display: flex; flex-direction: column; gap: 8px; z-index: 70; } +.toast { background: var(--surface); border: 1px solid var(--line); border-left: 3px solid var(--accent); border-radius: 8px; + padding: 10px 14px; box-shadow: var(--shadow); max-width: 420px; } +.toast.error { border-left-color: var(--critical); color: var(--ink); } +dialog { border: 1px solid var(--line); border-radius: 12px; background: var(--surface); color: var(--ink); padding: 20px; width: min(420px, 92vw); } +dialog::backdrop { background: rgb(0 0 0 / 45%); } +dialog form > * + * { margin-top: 12px; } +dialog label { display: flex; flex-direction: column; gap: 6px; } +@media (forced-colors: active) { + .chart .series { stroke: CanvasText; } .badge::before { forced-color-adjust: none; } +} diff --git a/chronos/webui/static/app.js b/chronos/webui/static/app.js new file mode 100644 index 0000000..605de0e --- /dev/null +++ b/chronos/webui/static/app.js @@ -0,0 +1,165 @@ +// Entry point: session, language, theme, navigation, page lifecycle and polling. +import { state, t, h, api, post, loadLanguage, saveLanguage, preferredLanguage, translateStatic, applyTheme, savedTheme, + setUnauthorizedHandler, reportError, hideTip, clear, setTip } from "./core.js"; +import { renderConfig } from "./pages_config.js"; +import { renderData, renderTokenizer } from "./pages_data.js"; +import { renderTrain } from "./pages_train.js"; +import { renderMonitor } from "./pages_monitor.js"; +import { renderInference } from "./pages_infer.js"; +import { renderExport, renderServe, renderBenchmark, renderAutotune, renderIO } from "./pages_ops.js"; + +const PAGES = [ + ["config", renderConfig], ["data", renderData], ["tokenizer", renderTokenizer], ["train", renderTrain], + ["monitor", renderMonitor], ["inference", renderInference], ["export", renderExport], ["serve", renderServe], + ["benchmark", renderBenchmark], ["autotune", renderAutotune], ["io", renderIO], +]; +const ROUTES = new Map(PAGES); +let current = null; // { id, timer, poll, alive } + +function route() { + const id = location.hash.replace(/^#\/?/, ""); + return ROUTES.has(id) ? id : "config"; +} + +function renderNav(active) { + const nav = document.getElementById("nav"); + nav.setAttribute("aria-label", t("ui.sections")); + nav.replaceChildren(...PAGES.map(([id]) => h("a", { href: `#${id}`, "aria-current": id === active ? "page" : null, text: t(`nav.${id}`) }))); +} + +function stopPage() { + if (!current) return; + current.alive = false; + clearTimeout(current.timer); + current = null; +} + +// Each page may register one poller: it returns the delay to the next call (0 = stop). +function makePage(id) { + const page = { id, alive: true, timer: 0, fn: null }; + const tick = async () => { + if (!page.alive || !page.fn) return; + clearTimeout(page.timer); + let delay = 0; + try { delay = await page.fn(); } catch (err) { reportError(err); delay = 10000; } + if (page.alive && delay > 0) page.timer = setTimeout(tick, document.hidden ? Math.max(delay, 15000) : delay); + }; + page.poll = (fn) => { page.fn = fn; tick(); }; + page.refresh = () => { if (page.fn) tick(); }; + return page; +} + +async function show({ focus = false } = {}) { + stopPage(); + hideTip(); + const id = route(); + renderNav(id); + const main = document.getElementById("main"); + clear(main); + current = makePage(id); + const page = current; + try { + await ROUTES.get(id)(main, page); + } catch (err) { + reportError(err); + main.append(h("div", { class: "empty", text: t("ui.error") })); + } + if (focus && page.alive) main.focus({ preventScroll: true }); +} + +document.addEventListener("chronos:refresh", () => { if (current) current.refresh(); }); +document.addEventListener("visibilitychange", () => { if (!document.hidden && current) current.refresh(); }); +window.addEventListener("hashchange", () => show({ focus: true })); +// ── session / login ──────────────────────────────────────────────────── +function paintSession() { + const s = state.session || {}; + const badgeEl = document.getElementById("auth-badge"); + // local_only: loopback, LAN off · lan_local: this machine while LAN is on · lan: another machine, signed in + const mode = s.access === "lan" ? "lan" : (s.lan_enabled ? "lan_local" : "local_only"); + badgeEl.className = `badge ${mode === "lan_local" ? "warning" : (mode === "lan" ? "good" : "")}`; + badgeEl.textContent = t({ local_only: "ui.local_only", lan_local: "ui.lan_local", lan: "ui.remote_auth" }[mode]); + setTip(badgeEl, t(`ui.access_tip.${mode}`)); + document.getElementById("logout").hidden = !(s.auth_required && s.authenticated); +} + +let loginOpen = false; +function askLogin() { + const dialog = document.getElementById("login-dialog"); + if (loginOpen) return; + loginOpen = true; + stopPage(); + document.getElementById("login-error").hidden = true; + dialog.showModal(); + document.getElementById("login-key").focus(); +} + +async function onLogin(event) { + event.preventDefault(); + const input = document.getElementById("login-key"); + const error = document.getElementById("login-error"); + try { + await post("/api/login", { key: input.value }); + input.value = ""; + loginOpen = false; + document.getElementById("login-dialog").close(); + state.session = await api("/api/session"); + paintSession(); + await start(); + } catch (err) { + error.textContent = err.status === 429 ? t("ui.login_blocked") : t("ui.login_failed"); + error.hidden = false; + } +} + +async function onLogout() { + try { await post("/api/logout"); } catch (err) { reportError(err); } + state.session = await api("/api/session"); + paintSession(); + askLogin(); +} + +// ── language / theme ─────────────────────────────────────────────────── +async function setLanguage(lang) { + await loadLanguage(lang); + saveLanguage(state.lang); + document.getElementById("lang-select").value = state.lang; + translateStatic(); + paintSession(); + if (state.schema) await show(); +} + +let started = false; +async function start() { + if (!state.schema) { + state.schema = await api("/api/schema"); + state.fields = Object.fromEntries(state.schema.fields.map((f) => [f.key, f])); + } + started = true; + await show(); +} + +async function boot() { + applyTheme(savedTheme()); + const theme = document.getElementById("theme-select"); + theme.value = savedTheme(); + theme.addEventListener("change", () => applyTheme(theme.value)); + const lang = document.getElementById("lang-select"); + lang.addEventListener("change", () => setLanguage(lang.value).catch(reportError)); + document.getElementById("login-form").addEventListener("submit", onLogin); + document.getElementById("login-dialog").addEventListener("cancel", (e) => e.preventDefault()); + document.getElementById("logout").addEventListener("click", onLogout); + setUnauthorizedHandler(askLogin); + + await loadLanguage(preferredLanguage()); + lang.value = state.lang; + translateStatic(); + state.session = await api("/api/session"); + paintSession(); + if (state.session.auth_required && !state.session.authenticated) { askLogin(); return; } + await start(); +} + +boot().catch((err) => { + reportError(err); + if (!started) document.getElementById("main").append(h("div", { class: "empty", text: String(err.message || err) })); +}); diff --git a/chronos/webui/static/charts.js b/chronos/webui/static/charts.js new file mode 100644 index 0000000..2f10a38 --- /dev/null +++ b/chronos/webui/static/charts.js @@ -0,0 +1,250 @@ +// Inline-SVG charts: line (crosshair + tooltip listing every series), horizontal +// bars, histogram. Each chart has a legend (>= 2 series) and a table view. +// Colors come from CSS custom properties, so light and dark are both selected. +import { h, t, num, clear, showTipAt, hideTip, SERIES } from "./core.js"; + +const NS = "http://www.w3.org/2000/svg"; +export function s(tag, attrs, ...kids) { + const el = document.createElementNS(NS, tag); + for (const [k, v] of Object.entries(attrs || {})) { + if (v === undefined || v === null) continue; + // var() colors go through CSSOM: presentation attributes do not resolve custom properties everywhere, + // and the CSP forbids style attributes. + if ((k === "fill" || k === "stroke") && String(v).startsWith("var(")) el.style[k] = v; + else el.setAttribute(k, String(v)); + } + for (const kid of kids) if (kid !== null && kid !== undefined) el.append(kid instanceof Node ? kid : document.createTextNode(String(kid))); + return el; +} +export const color = (i) => `var(${SERIES[((i % SERIES.length) + SERIES.length) % SERIES.length]})`; + +function niceStep(span, count) { + const raw = span / Math.max(1, count); + const mag = 10 ** Math.floor(Math.log10(raw)); + const n = raw / mag; + return (n <= 1 ? 1 : n <= 2 ? 2 : n <= 5 ? 5 : 10) * mag; +} +export function ticks(lo, hi, count = 4) { + if (!(hi > lo)) { const pad = Math.abs(lo) * 0.1 || 1; lo -= pad; hi += pad; } + const step = niceStep(hi - lo, count); + const a = Math.floor(lo / step) * step; const b = Math.ceil(hi / step) * step; + const values = []; + for (let v = a; v <= b + step * 1e-9; v += step) values.push(Math.abs(v) < step * 1e-9 ? 0 : v); + return { lo: a, hi: b, values }; +} +const extent = (vals) => [Math.min(...vals), Math.max(...vals)]; +const textWidth = (str) => [...String(str)].reduce((w, ch) => w + (ch.charCodeAt(0) > 0x2e7f ? 11 : 6.4), 0); + +// Frame shared by every chart: title, legend, chart/table toggle, resize. +function frame({ title, legend, draw, tableOf, note }) { + const body = h("div", { class: "chart" }); + const head = h("div", { class: "chart-head" }); + const view = { table: false }; + const toggle = h("button", { type: "button", class: "ghost", text: t("ui.table_view"), + onclick: () => { view.table = !view.table; toggle.textContent = t(view.table ? "ui.chart_view" : "ui.table_view"); paint(); } }); + head.append(h("h4", { text: title || "" }), + h("div", { class: "row" }, legend && legend.length >= 2 ? h("div", { class: "legend" }, legend.map((l) => { + const sw = h("i", { class: l.box ? "box" : null }); sw.style.background = l.color; + return h("span", null, sw, l.label); + })) : null, toggle)); + const wrap = h("div", { class: "chart-block" }, head, body, note ? h("p", { class: "note", text: note }) : null); + let lastWidth = 0; + const paint = () => { + clear(body); + if (view.table) { body.append(tableOf()); return; } + const width = Math.max(260, Math.floor(body.clientWidth || wrap.clientWidth || 640)); + lastWidth = width; + body.append(draw(width)); + }; + const ro = new ResizeObserver(() => { + const w = Math.floor(body.clientWidth); + if (!view.table && w && Math.abs(w - lastWidth) > 4) paint(); + }); + ro.observe(body); + requestAnimationFrame(paint); + return wrap; +} + +export function emptyState(text) { return h("div", { class: "empty", text: text || t("ui.empty") }); } +function tipBody(title, rows) { + return h("div", null, h("div", { class: "t-title", text: title }), rows.map((r) => { + const key = h("i", { class: "key" }); key.style.background = r.color; + return h("div", { class: "t-row" }, key, h("span", { text: r.label }), h("b", { text: r.value })); + })); +} + +// series: [{label, color?, points: [[x, y], ...]}]; colors follow the entity (index given by caller). +export function lineChart({ title, series, xLabel, yFmt = (v) => num(v, 3), xFmt = (v) => num(v, 0), height = 220, note, empty }) { + const live = (series || []).map((sr, i) => ({ ...sr, color: sr.color || color(i), + points: (sr.points || []).filter((p) => p && Number.isFinite(p[0]) && Number.isFinite(p[1])).sort((a, b) => a[0] - b[0]) })) + .filter((sr) => sr.points.length); + if (!live.length) return h("div", { class: "chart-block" }, h("h4", { text: title || "" }), emptyState(empty)); + const xs = [...new Set(live.flatMap((sr) => sr.points.map((p) => p[0])))].sort((a, b) => a - b); + const [y0, y1] = extent(live.flatMap((sr) => sr.points.map((p) => p[1]))); + const yt = ticks(y0, y1, 4); + const at = (sr, x) => { // nearest point of one series to x + let lo = 0; let hi = sr.points.length - 1; + while (hi - lo > 1) { const mid = (lo + hi) >> 1; if (sr.points[mid][0] < x) lo = mid; else hi = mid; } + return Math.abs(sr.points[lo][0] - x) <= Math.abs(sr.points[hi][0] - x) ? sr.points[lo] : sr.points[hi]; + }; + const draw = (width) => { + const left = Math.max(...yt.values.map((v) => textWidth(yFmt(v)))) + 12; + const m = { l: left, r: 14, t: 10, b: xLabel ? 40 : 26 }; + const W = width - m.l - m.r; const H = height - m.t - m.b; + const [xa, xb] = xs.length > 1 ? [xs[0], xs[xs.length - 1]] : [xs[0] - 1, xs[0] + 1]; + const X = (x) => m.l + ((x - xa) / (xb - xa)) * W; + const Y = (y) => m.t + H - ((y - yt.lo) / (yt.hi - yt.lo)) * H; + const svg = s("svg", { width, height, viewBox: `0 0 ${width} ${height}`, role: "img", "aria-label": title || "" }); + for (const v of yt.values) { + svg.append(s("line", { class: "grid-line", x1: m.l, x2: m.l + W, y1: Y(v), y2: Y(v) }), + s("text", { x: m.l - 8, y: Y(v) + 4, "text-anchor": "end" }, yFmt(v))); + } + svg.append(s("line", { class: "axis-line", x1: m.l, x2: m.l + W, y1: m.t + H, y2: m.t + H })); + const xt = ticks(xa, xb, Math.max(2, Math.floor(W / 110))).values.filter((v) => v >= xa && v <= xb); + for (const v of xt) svg.append(s("text", { x: X(v), y: m.t + H + 16, "text-anchor": "middle" }, xFmt(v))); + if (xLabel) svg.append(s("text", { x: m.l + W / 2, y: height - 6, "text-anchor": "middle" }, xLabel)); + for (const sr of live) { + const d = sr.points.map((p, i) => `${i ? "L" : "M"}${X(p[0]).toFixed(1)},${Y(p[1]).toFixed(1)}`).join(""); + svg.append(s("path", { class: "series", d, stroke: sr.color })); + const last = sr.points[sr.points.length - 1]; + svg.append(s("circle", { class: "end-dot", cx: X(last[0]), cy: Y(last[1]), r: 4, fill: sr.color })); + } + const cross = s("line", { class: "crosshair", y1: m.t, y2: m.t + H, visibility: "hidden" }); + const dots = live.map((sr) => s("circle", { class: "end-dot", r: 4, fill: sr.color, visibility: "hidden" })); + svg.append(cross, ...dots); + const overlay = s("rect", { x: m.l, y: m.t, width: W, height: H, fill: "transparent", tabindex: 0, + "aria-label": `${title || ""}: ${t("ui.table_view")}` }); + let index = xs.length - 1; + const show = (i, cx, cy) => { + index = Math.max(0, Math.min(xs.length - 1, i)); + const x = xs[index]; + cross.setAttribute("x1", X(x)); cross.setAttribute("x2", X(x)); cross.setAttribute("visibility", "visible"); + const rows = live.map((sr, j) => { + const p = at(sr, x); + dots[j].setAttribute("cx", X(p[0])); dots[j].setAttribute("cy", Y(p[1])); dots[j].setAttribute("visibility", "visible"); + return { color: sr.color, label: sr.label, value: yFmt(p[1]) }; + }); + const box = svg.getBoundingClientRect(); + showTipAt(cx ?? box.left + X(x), cy ?? box.top + m.t, tipBody(`${xLabel ? `${xLabel} ` : ""}${xFmt(x)}`, rows)); + }; + const hide = () => { cross.setAttribute("visibility", "hidden"); dots.forEach((d) => d.setAttribute("visibility", "hidden")); hideTip(); }; + overlay.addEventListener("pointermove", (e) => { + const box = svg.getBoundingClientRect(); + const x = xa + ((e.clientX - box.left - m.l) / W) * (xb - xa); + let best = 0; + for (let i = 1; i < xs.length; i += 1) if (Math.abs(xs[i] - x) < Math.abs(xs[best] - x)) best = i; + show(best, e.clientX, e.clientY); + }); + overlay.addEventListener("pointerleave", hide); + overlay.addEventListener("blur", hide); + overlay.addEventListener("focus", () => show(index)); + overlay.addEventListener("keydown", (e) => { + if (e.key === "ArrowLeft") { show(index - 1); e.preventDefault(); } + if (e.key === "ArrowRight") { show(index + 1); e.preventDefault(); } + }); + svg.append(overlay); + return svg; + }; + const tableOf = () => { + const rows = xs.slice(-500).map((x) => [x, ...live.map((sr) => { const p = at(sr, x); return p[0] === x ? p[1] : null; })]); + return h("div", { class: "table-wrap" }, h("table", null, + h("thead", null, h("tr", null, h("th", { class: "num", text: xLabel || "x" }), live.map((sr) => h("th", { class: "num", text: sr.label })))), + h("tbody", null, rows.map((r) => h("tr", null, h("td", { class: "num", text: xFmt(r[0]) }), + r.slice(1).map((v) => h("td", { class: "num", text: v === null ? "–" : yFmt(v) }))))))); + }; + return frame({ title, legend: live.map((sr) => ({ label: sr.label, color: sr.color })), draw, tableOf, note }); +} +// Bar with a 4px rounded data end; the baseline end stays square. +function barPath(x, y, w, hgt, horizontal) { + const r = Math.min(4, Math.abs(horizontal ? w : hgt) / 2, (horizontal ? hgt : w) / 2); + if (horizontal) { + if (w <= 0) return ""; + return `M${x},${y}h${w - r}a${r},${r} 0 0 1 ${r},${r}v${hgt - 2 * r}a${r},${r} 0 0 1 ${-r},${r}h${-(w - r)}z`; + } + if (hgt <= 0) return ""; + return `M${x},${y + hgt}v${-(hgt - r)}a${r},${r} 0 0 1 ${r},${-r}h${w - 2 * r}a${r},${r} 0 0 1 ${r},${r}v${hgt - r}z`; +} + +function markTip(el, content) { + el.setAttribute("tabindex", "0"); + const show = (e) => { + const box = el.getBoundingClientRect(); + showTipAt(e && e.clientX !== undefined ? e.clientX : box.right, e && e.clientY !== undefined ? e.clientY : box.top, content()); + }; + el.addEventListener("pointermove", show); + el.addEventListener("focus", () => show()); + el.addEventListener("pointerleave", () => hideTip()); + el.addEventListener("blur", () => hideTip()); +} + +// Horizontal grouped bars: categories down the side, one bar per series. +export function barChart({ title, categories, series, fmt = (v) => num(v, 2), note, empty }) { + const live = (series || []).map((sr, i) => ({ ...sr, color: sr.color || color(i) })); + const values = live.flatMap((sr) => sr.values.filter((v) => Number.isFinite(v))); + if (!values.length) return h("div", { class: "chart-block" }, h("h4", { text: title || "" }), emptyState(empty)); + const bar = live.length > 1 ? 14 : 18; const gap = 2; const groupGap = 14; + const groupH = live.length * bar + (live.length - 1) * gap; + const height = categories.length * (groupH + groupGap) + 24; + const max = Math.max(0, ...values); const min = Math.min(0, ...values); + const xt = ticks(min, max, 4); + const draw = (width) => { + const left = Math.min(220, Math.max(...categories.map(textWidth)) + 14); + const right = Math.max(...values.map((v) => textWidth(fmt(v)))) + 14; + const W = width - left - right; + const X = (v) => left + ((v - xt.lo) / (xt.hi - xt.lo)) * W; + const svg = s("svg", { width, height, viewBox: `0 0 ${width} ${height}`, role: "img", "aria-label": title || "" }); + for (const v of xt.values) svg.append(s("line", { class: "grid-line", x1: X(v), x2: X(v), y1: 0, y2: height - 20 }), + s("text", { x: X(v), y: height - 6, "text-anchor": "middle" }, fmt(v))); + categories.forEach((cat, ci) => { + const gy = ci * (groupH + groupGap) + 4; + svg.append(s("text", { x: left - 10, y: gy + groupH / 2 + 4, "text-anchor": "end" }, cat)); + live.forEach((sr, si) => { + const v = sr.values[ci]; + if (!Number.isFinite(v)) return; + const y = gy + si * (bar + gap); + const x0 = X(Math.max(0, xt.lo)); const x1 = X(v); + const el = s("path", { class: "bar", d: barPath(Math.min(x0, x1), y, Math.max(1, Math.abs(x1 - x0)), bar, true), fill: sr.color }); + markTip(el, () => tipBody(cat, [{ color: sr.color, label: sr.label, value: fmt(v) }])); + svg.append(el, s("text", { x: Math.max(x0, x1) + 6, y: y + bar / 2 + 4 }, fmt(v))); + }); + }); + svg.append(s("line", { class: "axis-line", x1: X(Math.max(0, xt.lo)), x2: X(Math.max(0, xt.lo)), y1: 0, y2: height - 20 })); + return svg; + }; + const tableOf = () => h("div", { class: "table-wrap" }, h("table", null, + h("thead", null, h("tr", null, h("th", { text: "" }), live.map((sr) => h("th", { class: "num", text: sr.label })))), + h("tbody", null, categories.map((cat, ci) => h("tr", null, h("td", { text: cat }), + live.map((sr) => h("td", { class: "num", text: Number.isFinite(sr.values[ci]) ? fmt(sr.values[ci]) : "–" }))))))); + return frame({ title, legend: live.map((sr) => ({ label: sr.label, color: sr.color, box: true })), draw, tableOf, note }); +} + +// Histogram over (possibly log-spaced) bin edges; one hue, per-bar tooltip. +export function histogram({ title, edges, counts, xLabel, height = 200, note, empty }) { + if (!counts || !counts.length) return h("div", { class: "chart-block" }, h("h4", { text: title || "" }), emptyState(empty)); + const yt = ticks(0, Math.max(...counts), 4); + const draw = (width) => { + const left = Math.max(...yt.values.map((v) => textWidth(num(v, 0)))) + 12; + const m = { l: left, r: 10, t: 8, b: xLabel ? 40 : 26 }; + const W = width - m.l - m.r; const H = height - m.t - m.b; + const bw = W / counts.length; + const Y = (v) => m.t + H - (v / yt.hi) * H; + const svg = s("svg", { width, height, viewBox: `0 0 ${width} ${height}`, role: "img", "aria-label": title || "" }); + for (const v of yt.values) svg.append(s("line", { class: "grid-line", x1: m.l, x2: m.l + W, y1: Y(v), y2: Y(v) }), + s("text", { x: m.l - 8, y: Y(v) + 4, "text-anchor": "end" }, num(v, 0))); + counts.forEach((c, i) => { + const x = m.l + i * bw + 1; const w = Math.max(1, bw - 2); + const el = s("path", { class: "bar", d: barPath(x, Y(c), w, m.t + H - Y(c), false) || `M${x},${m.t + H}h${w}v-1h${-w}z`, fill: "var(--s1)" }); + markTip(el, () => tipBody(`${num(edges[i], 0)} – ${num(edges[i + 1], 0)}`, [{ color: "var(--s1)", label: t("ui.value"), value: num(c, 0) }])); + svg.append(el); + }); + svg.append(s("line", { class: "axis-line", x1: m.l, x2: m.l + W, y1: m.t + H, y2: m.t + H })); + const every = Math.max(1, Math.ceil(counts.length / Math.max(2, Math.floor(W / 70)))); + edges.forEach((e, i) => { if (i % every === 0 || i === edges.length - 1) svg.append(s("text", { x: m.l + i * bw, y: m.t + H + 16, "text-anchor": "middle" }, num(e, 0))); }); + if (xLabel) svg.append(s("text", { x: m.l + W / 2, y: height - 6, "text-anchor": "middle" }, xLabel)); + return svg; + }; + const tableOf = () => h("div", { class: "table-wrap" }, h("table", null, + h("thead", null, h("tr", null, h("th", { class: "num", text: xLabel || "" }), h("th", { class: "num", text: t("ui.value") }))), + h("tbody", null, counts.map((c, i) => h("tr", null, h("td", { class: "num", text: `${num(edges[i], 0)} – ${num(edges[i + 1], 0)}` }), h("td", { class: "num", text: num(c, 0) })))))); + return frame({ title, legend: null, draw, tableOf, note }); +} diff --git a/chronos/webui/static/core.js b/chronos/webui/static/core.js new file mode 100644 index 0000000..bffe2b6 --- /dev/null +++ b/chronos/webui/static/core.js @@ -0,0 +1,369 @@ +// Shared runtime: API calls, language, translation, DOM helpers, tooltip, toasts. +// All text reaches the page through textContent; no HTML strings are ever parsed. + +export const LANGS = ["zh-Hans", "zh-Hant", "en", "ja"]; +const LANG_KEY = "chronos.ui.lang"; +const THEME_KEY = "chronos.ui.theme"; +const HANT = new Set(["zh-tw", "zh-hk", "zh-mo", "zh-hant"]); + +export const state = { + lang: "en", messages: {}, fallback: {}, schema: null, fields: {}, session: null, + values: {}, // per-form values kept across re-renders (language switch, navigation) + config: null, // the model config shared by Config / Train / Inference / Auto-tune +}; + +// Same mapping as server.resolve_language: first supported tag wins, English otherwise. +export function resolveLanguage(tags) { + for (const raw of tags || []) { + const code = String(raw || "").trim().toLowerCase().replace(/_/g, "-"); + if (!code) continue; + if (HANT.has(code) || code.startsWith("zh-hant")) return "zh-Hant"; + if (code === "zh" || code.startsWith("zh-")) return "zh-Hans"; + if (code === "ja" || code.startsWith("ja-")) return "ja"; + if (code === "en" || code.startsWith("en-")) return "en"; + } + return "en"; +} + +export function preferredLanguage() { + const saved = localStorage.getItem(LANG_KEY); + if (LANGS.includes(saved)) return saved; + return resolveLanguage(navigator.languages && navigator.languages.length ? navigator.languages : [navigator.language]); +} + +export class ApiError extends Error { + constructor(status, detail) { + super(typeof detail === "string" ? detail : (detail && (detail.message || detail.error)) || `HTTP ${status}`); + this.status = status; + this.detail = detail; + this.fields = (detail && Array.isArray(detail.fields)) ? detail.fields : []; + this.code = (detail && typeof detail === "object" && detail.error) || (typeof detail === "string" ? detail : ""); + } +} + +let onUnauthorized = () => {}; +export function setUnauthorizedHandler(fn) { onUnauthorized = fn; } + +export async function api(path, { method = "GET", body, signal } = {}) { + const headers = { "X-Chronos-Request": "1", Accept: "application/json" }; + const init = { method, headers, credentials: "same-origin", signal }; + if (body !== undefined) { + headers["Content-Type"] = "application/json"; + init.body = JSON.stringify(body); + } + const res = await fetch(path, init); + let data = null; + try { data = await res.json(); } catch { data = null; } + if (res.status === 401 && !path.startsWith("/api/login")) onUnauthorized(); + if (!res.ok) throw new ApiError(res.status, data && data.detail !== undefined ? data.detail : data); + return data; +} +export const post = (path, body = {}) => api(path, { method: "POST", body }); + +// ── translation ───────────────────────────────────────────────────────── +export async function loadLanguage(lang) { + const code = LANGS.includes(lang) ? lang : "en"; + const [main, en] = await Promise.all([ + api(`/api/i18n/${code}`), + code === "en" ? null : (Object.keys(state.fallback).length ? { messages: state.fallback } : api("/api/i18n/en")), + ]); + state.lang = code; + state.messages = main.messages || {}; + state.fallback = code === "en" ? state.messages : (en.messages || {}); + document.documentElement.lang = code; + return code; +} + +export function saveLanguage(lang) { localStorage.setItem(LANG_KEY, lang); } + +export function has(key) { return key in state.messages || key in state.fallback; } + +export function t(key, vars) { + let text = state.messages[key]; + if (text === undefined) text = state.fallback[key]; + if (text === undefined) text = key; + if (vars) text = text.replace(/\{(\w+)\}/g, (m, name) => (name in vars ? String(vars[name]) : m)); + return text; +} + +export function translateStatic(root = document) { + for (const node of root.querySelectorAll("[data-i18n]")) node.textContent = t(node.dataset.i18n); + for (const node of root.querySelectorAll("[data-i18n-label]")) node.setAttribute("aria-label", t(node.dataset.i18nLabel)); + document.title = t("app.title"); +} +// ── DOM (text only) ───────────────────────────────────────────────────── +const tips = new WeakMap(); + +export function h(tag, attrs, ...children) { + const el = tag === "svg" || attrs?.svg ? document.createElementNS("http://www.w3.org/2000/svg", tag) + : document.createElement(tag); + for (const [key, value] of Object.entries(attrs || {})) { + if (value === undefined || value === null || value === false || key === "svg") continue; + if (key === "class") el.setAttribute("class", value); + else if (key === "text") el.textContent = value; + else if (key === "tip") setTip(el, value); + else if (key.startsWith("on") && typeof value === "function") el.addEventListener(key.slice(2), value); + else if (key === "dataset") Object.assign(el.dataset, value); + else if (key in el && !(el instanceof SVGElement) && typeof value !== "string") el[key] = value; + else el.setAttribute(key, value === true ? "" : String(value)); + } + append(el, children); + return el; +} + +export function append(el, children) { + for (const child of children.flat(Infinity)) { + if (child === null || child === undefined || child === false) continue; + el.append(child instanceof Node ? child : document.createTextNode(String(child))); + } + return el; +} + +export function clear(el) { while (el.firstChild) el.firstChild.remove(); return el; } + +export function card(title, ...body) { return h("section", { class: "card" }, title ? h("h3", { text: title }) : null, ...body); } + +export function badge(stateName) { + const key = `state.${stateName}`; + return h("span", { class: `badge ${stateName || ""}`, role: "status", text: has(key) ? t(key) : String(stateName ?? "") }); +} + +export function table(columns, rows, { caption } = {}) { + // columns: [{label, num, fmt(row) -> text|Node}] + const head = h("tr", null, columns.map((c) => h("th", { class: c.num ? "num" : null, scope: "col", text: c.label }))); + const body = rows.map((row) => h("tr", null, columns.map((c) => { + const value = c.fmt ? c.fmt(row) : row[c.key]; + return h("td", { class: c.num ? "num" : null }, value === null || value === undefined || value === "" ? "–" : value); + }))); + return h("div", { class: "table-wrap" }, + h("table", null, caption ? h("caption", { class: "skip", text: caption }) : null, h("thead", null, head), h("tbody", null, body))); +} + +// ── number / time formatting in the page language ─────────────────────── +const nf = (opts) => new Intl.NumberFormat(state.lang, opts); +export function num(v, digits = 2) { + if (v === null || v === undefined || v === "" || Number.isNaN(Number(v))) return "–"; + const x = Number(v); + if (x !== 0 && Math.abs(x) < 10 ** -digits) return x.toExponential(1); + return nf({ maximumFractionDigits: digits }).format(x); +} +export const int = (v) => num(v, 0); +export const pct = (v, digits = 1) => (v === null || v === undefined ? "–" : nf({ style: "percent", maximumFractionDigits: digits }).format(Number(v))); +export function params(v) { + if (v === null || v === undefined) return "–"; + const x = Number(v); + if (Math.abs(x) >= 1e9) return `${num(x / 1e9, 2)}B`; + if (Math.abs(x) >= 1e6) return `${num(x / 1e6, 2)}M`; + if (Math.abs(x) >= 1e3) return `${num(x / 1e3, 1)}K`; + return int(x); +} +export function bytes(v) { + if (v === null || v === undefined) return "–"; + const x = Number(v); + const units = ["B", "KB", "MB", "GB", "TB"]; + let i = 0; let y = Math.abs(x); + while (y >= 1024 && i < units.length - 1) { y /= 1024; i += 1; } + return `${num(Math.sign(x) * y, y >= 100 ? 0 : 2)} ${units[i]}`; +} +export function duration(seconds) { + if (seconds === null || seconds === undefined || !Number.isFinite(Number(seconds))) return "–"; + let s = Math.max(0, Math.round(Number(seconds))); + const parts = { hours: Math.floor(s / 3600), minutes: Math.floor((s % 3600) / 60), seconds: s % 60 }; + if (parts.hours) delete parts.seconds; + if (typeof Intl.DurationFormat === "function") { + try { return new Intl.DurationFormat(state.lang, { style: "narrow" }).format(parts); } catch { /* fall through */ } + } + s = parts.hours ? `${parts.hours}h ${parts.minutes}m` : `${parts.minutes}m ${parts.seconds}s`; + return s; +} +export function dateTime(v) { + if (!v) return "–"; + const d = typeof v === "number" ? new Date(v * 1000) : new Date(v); + return Number.isNaN(d.getTime()) ? String(v) : new Intl.DateTimeFormat(state.lang, { dateStyle: "short", timeStyle: "short" }).format(d); +} +// ── tooltip: one element ──────────────────────────────────────────────── +// Hover shows a field's help after the pointer rests on it for HOVER_MS, so moving across a form +// never covers the next fields. The "?" dot is an explicit request and answers after HELP_MS. +// Keyboard focus (Tab) shows it at once; a click, typing, Escape or scrolling hides it. +const HOVER_MS = 2000; +const HELP_MS = 400; +let tipEl = null; +let tipOwner = null; +let tipTimer = 0; +let tipPending = null; +let viaKeyboard = false; +let suppressed = null; + +function cancelTip() { + clearTimeout(tipTimer); + tipTimer = 0; + tipPending = null; +} + +function scheduleTip(el, ms) { + if (tipOwner === el || suppressed === el) return; + if (tipPending === el && ms >= HOVER_MS) return; + cancelTip(); + tipPending = el; + tipTimer = setTimeout(() => { tipPending = null; if (el.isConnected) showTip(el); }, ms); +} + +export function setTip(el, content) { + tips.set(el, content); + if (el.dataset.tipBound) return; + el.dataset.tipBound = "1"; + el.addEventListener("pointerenter", (e) => { if (e.pointerType !== "touch") scheduleTip(el, HOVER_MS); }); + el.addEventListener("pointerover", (e) => { + if (e.pointerType !== "touch" && e.target instanceof Element && e.target.closest(".help-dot")) scheduleTip(el, HELP_MS); + }); + el.addEventListener("pointerleave", () => { + if (tipPending === el) cancelTip(); + if (suppressed === el) suppressed = null; + hideTip(el); + }); + el.addEventListener("pointerdown", (e) => { + cancelTip(); + if (e.target instanceof Element && e.target.closest(".help-dot")) { suppressed = null; showTip(el); return; } + suppressed = el; + hideTip(); + }); + el.addEventListener("focusin", () => { if (viaKeyboard) { cancelTip(); showTip(el); } }); + el.addEventListener("focusout", (e) => { if (!el.contains(e.relatedTarget)) hideTip(el); }); +} + +function fillTip(content) { + clear(tipEl); + const value = typeof content === "function" ? content() : content; + if (value instanceof Node) tipEl.append(value); + else tipEl.textContent = String(value ?? ""); +} + +function showTip(el) { + tipEl = tipEl || document.getElementById("tooltip"); + const content = tips.get(el); + if (!tipEl || content === undefined) return; + tipOwner = el; + fillTip(content); + tipEl.hidden = false; + el.setAttribute("aria-describedby", "tooltip"); + const r = el.getBoundingClientRect(); + const w = tipEl.offsetWidth; const hgt = tipEl.offsetHeight; + let left = Math.min(window.innerWidth - w - 8, Math.max(8, r.left)); + let top = r.bottom + 8; + if (top + hgt > window.innerHeight - 8) top = Math.max(8, r.top - hgt - 8); + if (left < 8) left = 8; + tipEl.style.left = `${left}px`; + tipEl.style.top = `${top}px`; +} + +export function hideTip(el) { + if (!el) cancelTip(); + if (!tipEl || (el && el !== tipOwner)) return; + tipEl.hidden = true; + if (tipOwner) tipOwner.removeAttribute("aria-describedby"); + tipOwner = null; +} + +// Chart crosshair tooltips are positioned at a point instead of an element. +export function showTipAt(x, y, content) { + tipEl = tipEl || document.getElementById("tooltip"); + if (!tipEl) return; + cancelTip(); + tipOwner = null; + fillTip(content); + tipEl.hidden = false; + const w = tipEl.offsetWidth; const hgt = tipEl.offsetHeight; + let left = x + 14; let top = y + 14; + if (left + w > window.innerWidth - 8) left = x - w - 14; + if (top + hgt > window.innerHeight - 8) top = y - hgt - 14; + tipEl.style.left = `${Math.max(8, left)}px`; + tipEl.style.top = `${Math.max(8, top)}px`; +} + +const MODIFIERS = new Set(["Shift", "Control", "Alt", "Meta", "CapsLock"]); +document.addEventListener("keydown", (e) => { + if (MODIFIERS.has(e.key)) return; + viaKeyboard = e.key === "Tab"; + if (!viaKeyboard) hideTip(); +}, true); +document.addEventListener("pointerdown", () => { viaKeyboard = false; }, true); +document.addEventListener("input", () => hideTip(), true); +window.addEventListener("scroll", () => hideTip(), { passive: true, capture: true }); + +// ── toasts and errors ─────────────────────────────────────────────────── +export function toast(message, kind = "info") { + const host = document.getElementById("toasts"); + if (!host) return; + const el = h("div", { class: `toast ${kind}`, role: kind === "error" ? "alert" : "status", text: message }); + host.append(el); + setTimeout(() => el.remove(), kind === "error" ? 9000 : 4500); +} + +export function errorText(err) { + if (err instanceof ApiError) { + // Backend detail messages are machine text (English); they are appended, never used as the only text. + const extra = err.detail && typeof err.detail === "object" && err.detail.message ? ` (${err.detail.message})` : ""; + if (err.code && has(`err.${err.code}`)) return t(`err.${err.code}`) + extra; + if (err.status === 409) return t("ui.busy"); + if (err.status === 422 && err.fields.length) return t("ui.invalid") + extra; + if (err.status === 404) return t("err.not_found"); + } + return `${t("ui.error")}: ${err && err.message ? err.message : String(err)}`; +} + +export function reportError(err) { + if (err instanceof ApiError && err.status === 401) return; + console.error(err); + toast(errorText(err), "error"); +} + +// ── theme ─────────────────────────────────────────────────────────────── +export function savedTheme() { + const v = localStorage.getItem(THEME_KEY); + return ["light", "dark", "system"].includes(v) ? v : "system"; +} +export function applyTheme(value) { + document.documentElement.dataset.theme = value; + localStorage.setItem(THEME_KEY, value); +} + +// CSS custom property in the active theme (charts read series colors this way). +export function cssVar(name) { return getComputedStyle(document.documentElement).getPropertyValue(name).trim(); } +export const SERIES = ["--s1", "--s2", "--s3", "--s4", "--s5", "--s6", "--s7", "--s8"]; + +// ── page furniture ────────────────────────────────────────────────────── +export function pageHead(titleKey, introKey, ...actions) { + return h("div", { class: "page-head" }, + h("div", null, h("h2", { text: t(titleKey) }), introKey ? h("p", { text: t(introKey) }) : null), + actions.length ? h("div", { class: "row" }, actions) : null); +} + +export function tile(label, value, sub, tip) { + return h("div", { class: "tile", tip }, h("div", { class: "k", text: label }), h("div", { class: "v", text: value ?? "–" }), + sub ? h("div", { class: "s", text: sub }) : null); +} + +export function fieldLabel(key) { return has(`field.${key}.label`) ? t(`field.${key}.label`) : key; } + +export function button(labelKey, onclick, { primary = false, danger = false, disabled = false } = {}) { + return h("button", { type: "button", class: primary ? "primary" : danger ? "danger" : null, disabled, text: t(labelKey), onclick }); +} + +// Run an async action with the button disabled; errors become toasts. +export async function busy(btn, fn) { + if (btn) btn.disabled = true; + try { return await fn(); } catch (err) { reportError(err); return undefined; } finally { if (btn) btn.disabled = false; } +} + +export function logBox(lines) { + const text = Array.isArray(lines) ? lines.join("\n") : String(lines || ""); + const pre = h("pre", { class: "log", tabindex: "0", text: text || t("ui.empty") }); + requestAnimationFrame(() => { pre.scrollTop = pre.scrollHeight; }); + return pre; +} + +export function rawJson(value) { + return h("details", null, h("summary", { text: t("ui.raw_json") }), h("pre", { class: "log", text: JSON.stringify(value, null, 2) })); +} + +export const clone = (v) => (v === undefined ? v : JSON.parse(JSON.stringify(v))); diff --git a/chronos/webui/static/forms.js b/chronos/webui/static/forms.js new file mode 100644 index 0000000..c628ab0 --- /dev/null +++ b/chronos/webui/static/forms.js @@ -0,0 +1,219 @@ +// Schema-driven forms. Every field carries a hover / focus tooltip in the page +// language: what it does, its default, the recommended value and where that +// recommendation comes from (code default, adaptive planner, measurement, rule). +import { state, t, has, h, num, setTip, ApiError, reportError } from "./core.js"; +import { pathControl } from "./picker.js"; + +const TEXTAREA = new Set(["tok.paths", "plan.data_paths", "train.sample_prompt", "reward_spec"]); +const WIDE = new Set([...TEXTAREA, "train.stages", "probe.ram_ratios", "export.formats", "tune.search", + "infer.model_path", "train.data_path", "serve.models_dir"]); +const showAdvanced = {}; +// Planner validation codes that name a form field. +const ERROR_ALIASES = { target: "plan.target_m", active_experts: "plan.active_value" }; + +export function fieldsFor(scope, group) { + return state.schema.fields.filter((f) => f.scope.includes(scope) && (!group || f.group === group)); +} + +export function defaultsFor(scope) { + const out = {}; + for (const f of fieldsFor(scope)) out[f.key] = Array.isArray(f.default) ? [...f.default] : f.default; + return out; +} + +// Values of one form, kept in state.values so they survive re-renders. +export function formState(id, scope) { + if (!state.values[id]) state.values[id] = defaultsFor(scope || id); + return state.values[id]; +} + +// Label of a choice value. With the field: its own wording first (the same id can mean different things +// in different fields; "serve.router_scope" shares "router_scope"), stage names for the stage chips and +// field names for the auto-tune search list. Format and device ids (float16, cuda, gguf…) stay as they are. +export function choiceLabel(value, f) { + if (typeof value === "boolean") return t(value ? "ui.yes" : "ui.no"); + if (value === "" || value === null || value === undefined) return t("ui.from_model"); + if (typeof value === "number") return num(value, 4); + const s = String(value); + if (f) { + for (const key of [f.key, f.key.replace(/^[a-z]+\./, "")]) if (has(`choice.${key}.${s}`)) return t(`choice.${key}.${s}`); + if (f.key === "train.stages" && has(`stage.${s}`)) return t(`stage.${s}`); + if (f.key === "tune.search" && has(`field.${s}.label`)) return t(`field.${s}.label`); + } + return has(`choice.${s}`) ? t(`choice.${s}`) : s; +} + +export function valueText(f, value) { + if (value === null || value === undefined || value === "") { + if (f.type === "choice") return choiceLabel(""); + return f.nullable || f.type === "text" || f.type === "path" ? t("tip.auto_empty") : "–"; + } + if (Array.isArray(value)) return value.map((x) => choiceLabel(x, f)).join(", "); + if (f.type === "bool") return t(value ? "ui.yes" : "ui.no"); + if (f.type === "choice") return choiceLabel(value, f); + if (typeof value === "number") return num(value, 6); + return String(value); +} + +export function recommendedText(f) { + if (f.recommended === undefined) return null; + if (f.source === "rule" && f.ref && has(`rule.${f.ref}`)) return t(`rule.${f.ref}`); + return valueText(f, f.recommended); +} + +function sourceText(f) { + if (!f.source) return null; + const base = t(`source.${f.source}`); + return f.source === "measured" && f.ref ? `${base} · ${f.ref.split("/").pop()}` : base; +} + +export function fieldTip(f) { + return () => { + const dl = h("dl"); + const row = (label, value) => { if (value !== null && value !== undefined) dl.append(h("dt", { text: label }), h("dd", { text: value })); }; + row(t("tip.default"), valueText(f, f.default)); + row(t("tip.recommended"), recommendedText(f)); + row(t("tip.source"), sourceText(f)); + if (f.min !== undefined || f.max !== undefined) row(t("tip.range"), `${f.min ?? "−∞"} – ${f.max ?? "∞"}`); + return h("div", null, + h("div", { class: "t-title", text: t(`field.${f.key}.label`) }), + h("div", { class: "t-help", text: t(`field.${f.key}.help`) }), + dl); + }; +} +const same = (a, b) => JSON.stringify(a ?? null) === JSON.stringify(b ?? null); +const domId = (key) => `f-${key.replace(/[^a-zA-Z0-9]/g, "-")}`; + +function control(f, values, set) { + const id = domId(f.key); + const v = values[f.key]; + if (f.type === "bool") { + const el = h("input", { type: "checkbox", id }); + el.checked = !!v; + el.addEventListener("change", () => set(el.checked)); + return el; + } + if (f.type === "choice") { + const el = h("select", { id }, f.choices.map((c, i) => h("option", { value: String(i), text: choiceLabel(c, f) }))); + const index = f.choices.findIndex((c) => same(c, v)); + el.value = String(index >= 0 ? index : 0); + el.addEventListener("change", () => set(f.choices[Number(el.value)])); + return el; + } + if (f.type === "multi") { + const picked = new Set((Array.isArray(v) ? v : []).map((x) => JSON.stringify(x))); + const box = h("div", { class: "chips", role: "group", id }); + const inputs = f.choices.map((c) => { + const input = h("input", { type: "checkbox" }); + input.checked = picked.has(JSON.stringify(c)); + input.addEventListener("change", () => set(f.choices.filter((x, j) => inputs[j].checked))); + box.append(h("label", null, input, choiceLabel(c, f))); + return input; + }); + return box; + } + if (f.type === "int" || f.type === "float") { + const el = h("input", { type: "number", id, inputmode: f.type === "int" ? "numeric" : "decimal", + min: f.min, max: f.max, step: f.step ?? (f.type === "int" ? 1 : "any"), + placeholder: f.nullable || f.default === null ? t("tip.auto_empty") : null }); + el.value = v === null || v === undefined ? "" : String(v); + el.addEventListener("input", () => set(el.value.trim() === "" ? null : Number(el.value))); + return el; + } + const paths = !!(f.pick && f.pick.multiple); + const multiline = TEXTAREA.has(f.key) || paths; + const el = h(multiline ? "textarea" : "input", { id, type: multiline ? null : "text", spellcheck: "false", + autocomplete: "off", class: paths ? "paths" : null, + placeholder: paths ? t("picker.placeholder_multi") : (f.type === "path" ? t("ui.path") : (f.nullable ? t("tip.auto_empty") : null)) }); + el.value = v ?? ""; + const fit = () => { if (paths) el.rows = Math.max(1, Math.min(8, el.value.split("\n").length)); }; + fit(); + el.addEventListener("input", () => { fit(); set(el.value); }); + return el; +} + +export function renderField(f, values, onChange) { + const wide = WIDE.has(f.key) || (f.type === "multi" && f.choices.length > 4) || !!(f.pick && f.pick.multiple); + const single = !!(f.pick && !f.pick.multiple); + const wrap = h("div", { class: `field${f.type === "bool" ? " bool" : ""}${wide ? " wide" : ""}${single && !wide ? " path" : ""}`, + dataset: { key: f.key } }); + const ctl = control(f, values, (x) => { + values[f.key] = x; + wrap.classList.toggle("changed", !same(x, f.default)); + wrap.classList.remove("bad"); + ctl.removeAttribute("aria-invalid"); + if (onChange) onChange(f.key, x); + }); + const rec = recommendedText(f); + const label = h(f.type === "multi" ? "span" : "label", { class: "label", for: f.type === "multi" ? null : domId(f.key) }, + h("span", { class: "name", text: t(`field.${f.key}.label`) }), + h("span", { class: "help-dot", "aria-hidden": "true", text: "?" }), + rec && rec.length <= 18 ? h("span", { class: "rec", text: `✓ ${rec}` }) : null); + if (f.type === "multi") ctl.setAttribute("aria-label", t(`field.${f.key}.label`)); + wrap.classList.toggle("changed", !same(values[f.key], f.default)); + wrap.append(label, f.pick ? pathControl(ctl, f.pick) : ctl); + setTip(wrap, fieldTip(f)); + return wrap; +} + +// One schema group as a collapsible section; advanced fields behind a toggle. +export function renderGroup(scope, group, values, onChange, { open = true, summary = true } = {}) { + const all = fieldsFor(scope, group); + const basic = all.filter((f) => !f.advanced); + const adv = all.filter((f) => f.advanced); + const key = `${scope}:${group}`; + const grid = h("div", { class: "fields" }, basic.map((f) => renderField(f, values, onChange))); + const advGrid = h("div", { class: "fields", hidden: !showAdvanced[key] }, adv.map((f) => renderField(f, values, onChange))); + const toggle = adv.length ? h("button", { type: "button", class: "ghost adv-toggle", "aria-expanded": String(!!showAdvanced[key]), + text: `${t(showAdvanced[key] ? "ui.hide_advanced" : "ui.show_advanced")} (${adv.length})`, + onclick: () => { + showAdvanced[key] = !showAdvanced[key]; + advGrid.hidden = !showAdvanced[key]; + toggle.setAttribute("aria-expanded", String(showAdvanced[key])); + toggle.textContent = `${t(showAdvanced[key] ? "ui.hide_advanced" : "ui.show_advanced")} (${adv.length})`; + } }) : null; + if (adv.length) advGrid.classList.add("mt"); + const body = [grid, toggle, advGrid]; + if (!summary) return h("div", null, body); + return h("details", { class: "group", open }, h("summary", { text: t(`group.${group}`) }), body); +} + +export function markErrors(root, keys, { scroll = true } = {}) { + for (const el of root.querySelectorAll(".field.bad")) el.classList.remove("bad"); + let first = null; + for (const raw of keys || []) { + const key = ERROR_ALIASES[raw] || raw; + const wrap = root.querySelector(`.field[data-key="${CSS.escape(key)}"]`); + if (!wrap) continue; + wrap.classList.add("bad"); + const ctl = wrap.querySelector("input, select, textarea, .chips"); + if (ctl) ctl.setAttribute("aria-invalid", "true"); + first = first || wrap; + } + if (first && scroll) { + const details = first.closest("details"); + if (details) details.open = true; + first.scrollIntoView({ block: "center", behavior: "smooth" }); + } +} + +// Replace a form's values (preset / loaded file / plan) without losing unknown keys. +export function assign(values, patch) { + for (const [key, value] of Object.entries(patch || {})) values[key] = Array.isArray(value) ? [...value] : value; + return values; +} + +// Run a form action: clears old marks, marks the fields a 422 names, reports other errors. +export async function submit(btn, root, fn) { + if (btn) btn.disabled = true; + try { + markErrors(root, [], { scroll: false }); + return await fn(); + } catch (err) { + if (err instanceof ApiError && err.status === 422 && err.fields.length) markErrors(root, err.fields); + reportError(err); + return undefined; + } finally { + if (btn) btn.disabled = false; + } +} diff --git a/chronos/webui/static/icon.svg b/chronos/webui/static/icon.svg new file mode 100644 index 0000000..756744e --- /dev/null +++ b/chronos/webui/static/icon.svg @@ -0,0 +1 @@ + diff --git a/chronos/webui/static/index.html b/chronos/webui/static/index.html new file mode 100644 index 0000000..c40f34b --- /dev/null +++ b/chronos/webui/static/index.html @@ -0,0 +1,58 @@ + + + + + + +Chronos + + + + + + +
+
+ +
+

Chronos

+

+
+
+
+ + + + +
+
+
+ +
+
+ +
+ +
+

Log in

+

+ + +
+
+ + + diff --git a/chronos/webui/static/pages_config.js b/chronos/webui/static/pages_config.js new file mode 100644 index 0000000..225f1b9 --- /dev/null +++ b/chronos/webui/static/pages_config.js @@ -0,0 +1,180 @@ +// Model page: presets, every config group, live size / memory estimate, auto plan. +import { state, t, has, h, api, post, card, tile, pageHead, button, busy, params, bytes, num, int, + toast, reportError, ApiError, table, fieldLabel, clone } from "./core.js"; +import { renderGroup, formState, markErrors, assign, submit } from "./forms.js"; +import { pathControl } from "./picker.js"; + +const PRESET_SLUGS = { "Chronos-v3 Tiny (smoke / CI)": "tiny", "Chronos-v3 50M (en→zh recipe)": "enzh50m", + "Chronos-v3 200M (32 GB)": "m200", "Chronos-v3 500M (32 GB)": "m500", "Chronos-v3 1B (64 GB)": "b1" }; + +let presetCache = null; +export async function presets() { + if (!presetCache) presetCache = await api("/api/presets"); + return presetCache; +} + +function schemaDefaults() { + const out = {}; + for (const f of state.schema.fields) if (f.scope.includes("config")) out[f.key] = clone(f.default); + return out; +} + +// The model config shared by every page; starts from the default preset. +export async function ensureConfig() { + if (state.config) return state.config; + const data = await presets(); + const preset = data.presets.find((p) => p.name === data.default); + state.config = { ...schemaDefaults(), ...clone(preset ? preset.config : {}) }; + state.presetName = preset ? preset.name : ""; + return state.config; +} + +export function presetLabel(p) { + const slug = PRESET_SLUGS[p.name]; + return slug && has(`preset.${slug}`) ? t(`preset.${slug}`) : p.name; +} + +let estimateTimer = 0; +function scheduleEstimate(host, form) { + clearTimeout(estimateTimer); + estimateTimer = setTimeout(() => refreshEstimate(host, form), 350); +} + +async function refreshEstimate(host, form) { + try { + const e = await post("/api/config/estimate", { config: state.config }); + markErrors(form, [], { scroll: false }); + const m = e.memory || {}; + host.replaceChildren(h("div", { class: "tiles" }, + tile(t("config.total"), params(e.total_params)), + tile(t("config.active"), params(e.active_params), `≤ ${params(e.active_union_upper)}`), + tile(t("config.compute"), params(e.compute_params)), + tile(t("config.virtual_layers"), int(e.virtual_layers)), + tile(t("config.vram"), bytes(m.vram_estimate_bytes)), + tile(t("config.ram"), bytes(m.ram_estimate_bytes)), + tile(t("config.ssd"), bytes(m.ssd_estimate_bytes)), + tile(t("config.kv"), bytes(m.kv_cache_bytes)), + tile(t("config.train_state"), bytes(m.training_state_bytes)), + tile(t("config.tps"), num(e.decode_tps_estimate, 0))), + h("p", { class: "note", text: t("config.estimate_note") })); + } catch (err) { + if (err instanceof ApiError && err.status === 422) { + markErrors(form, err.fields, { scroll: false }); + host.replaceChildren(h("p", { class: "error", text: t("ui.invalid") })); + } else reportError(err); + } +} +function presetNote(p) { + const m = (p && p.meta) || {}; + const parts = []; + if (m.total_params) parts.push(`${t("config.total")} ${params(m.total_params)}`); + if (m.active_params) parts.push(`${t("config.active")} ${params(m.active_params)}`); + if (m.training_peak_gb) parts.push(`${t("config.train_peak")} ${num(m.training_peak_gb, 1)} GB (${t("label.estimated")})`); + const slug = p && PRESET_SLUGS[p.name]; + const summary = slug && has(`preset.${slug}.summary`) ? t(`preset.${slug}.summary`) : ""; + return [summary, parts.join(" · ")].filter(Boolean).join(" — "); +} + +export async function renderConfig(main) { + const data = await presets(); + await ensureConfig(); + const estimate = h("div", null, h("p", { class: "note", text: t("ui.loading") })); + const form = h("div"); + const redraw = () => { + form.replaceChildren(...state.schema.config_groups.map((g, i) => + renderGroup("config", g, state.config, () => scheduleEstimate(estimate, form), { open: i < 3 }))); + }; + const apply = (cfg, message) => { + assign(state.config, cfg); + redraw(); + refreshEstimate(estimate, form); + if (message) toast(message); + }; + + // presets + const picker = h("select", { id: "preset" }, data.presets.map((p) => h("option", { value: p.name, text: presetLabel(p) }))); + picker.value = state.presetName || data.default; + const note = h("p", { class: "note" }); + const showNote = () => { note.textContent = presetNote(data.presets.find((p) => p.name === picker.value)); }; + picker.addEventListener("change", showNote); + showNote(); + const usePreset = button("ui.apply", () => { + const p = data.presets.find((x) => x.name === picker.value); + if (!p) return; + state.config = { ...schemaDefaults(), ...clone(p.config) }; + state.presetName = p.name; + apply({}, `${t("config.loaded")}: ${presetLabel(p)}`); + }, { primary: true }); + + // file load / save + const path = h("input", { type: "text", id: "config-path", spellcheck: "false", value: state.values.configPath || "./out/chronos_config.json" }); + path.addEventListener("input", () => { state.values.configPath = path.value; }); + const saveBtn = button("ui.save", () => busy(saveBtn, async () => { + const r = await post("/api/config/save", { config: state.config, path: path.value }); + toast(`${t("ui.saved")}: ${r.path}`); + })); + const loadBtn = button("ui.load", () => busy(loadBtn, async () => { + const r = await post("/api/config/load", { path: path.value }); + apply(r.config, `${t("config.loaded")}: ${r.path}`); + if (r.ignored && r.ignored.length) toast(`${t("config.ignored")}: ${r.ignored.map(fieldLabel).join(", ")}`, "error"); + })); + const resetBtn = button("ui.reset", () => { + state.config = null; + ensureConfig().then(() => { picker.value = state.presetName; showNote(); apply({}); }); + }); + + main.append( + pageHead("nav.config", "config.intro"), + h("div", { class: "grid two" }, + card(t("config.preset"), + h("div", { class: "row", tip: t("config.preset_help") }, h("label", { class: "compact", for: "preset", text: t("config.preset") }), picker, usePreset), + note), + card(t("config.file"), + h("div", { class: "field wide", tip: t("config.path_help") }, h("label", { class: "compact", for: "config-path", text: t("ui.path") }), + pathControl(path, { kind: "config", multiple: false })), + h("div", { class: "row mt" }, loadBtn, saveBtn, resetBtn))), + card(t("config.estimate"), estimate), + card(null, form), + planCard(apply)); + redraw(); + refreshEstimate(estimate, form); +} +// Auto plan: target size + data + memory → config and per-stage settings (adaptive planner). +function planCard(apply) { + const values = formState("plan", "planner"); + const root = h("div"); + const out = h("div", { class: "stack mt" }); + const run = button("plan.run", () => submit(run, root, async () => { + showPlan(out, await post("/api/plan", values), apply); + }), { primary: true }); + root.append(h("p", { class: "note", text: t("plan.note") }), renderGroup("planner", "planner", values, null, { summary: false }), + h("div", { class: "row mt" }, run), out); + return card(t("plan.title"), root); +} + +function showPlan(out, plan, apply) { + const a = (plan.report && plan.report.architecture) || {}; + const sc = (plan.report && plan.report.scaling) || {}; + const mem = (plan.report && plan.report.memory) || {}; + const stages = plan.stage_overrides || {}; + const keys = [...new Set(Object.values(stages).flatMap((row) => Object.keys(row || {})))]; + const use = button("plan.apply", () => apply({ ...plan.config, stage_overrides: plan.stage_overrides }, t("config.loaded")), { primary: true }); + out.replaceChildren( + h("h4", { text: t("plan.summary") }), + h("div", { class: "tiles" }, + tile(t("config.total"), params(a.total_params), t("label.measured_meta")), + tile(t("config.active"), params(a.active_params), `≤ ${params(a.active_union_upper)}`), + tile(t("config.virtual_layers"), int(a.virtual_layers), `${int(a.physical_layers)} × ${int(a.recurrent_depth)}`), + tile(fieldLabel("hidden_size"), int(a.hidden_size)), + tile(fieldLabel("num_experts"), `${int(a.experts)} / top-${int(a.top_k)}`), + tile(t("plan.tokens"), params(sc.tokens), t("label.estimated")), + tile(fieldLabel("learning_rate"), num(sc.chosen_lr, 5), t("label.estimated")), + tile(t("config.train_peak"), mem.training_peak_gb === undefined ? "–" : `${num(mem.training_peak_gb, 1)} GB`, t("label.estimated"))), + h("h4", { text: t("plan.stages") }), + table([{ label: t("ui.phase"), fmt: (r) => r.stage }, ...keys.map((k) => ({ label: fieldLabel(k), num: true, + fmt: (r) => (r[k] === undefined ? "" : (typeof r[k] === "number" ? num(r[k], 6) : String(r[k]))) }))], + Object.entries(stages).map(([stage, row]) => ({ stage, ...row }))), + h("div", { class: "row" }, use), + h("details", null, h("summary", { text: t("ui.raw_json") }), h("pre", { class: "log", text: JSON.stringify(plan, null, 2) }))); +} + diff --git a/chronos/webui/static/pages_data.js b/chronos/webui/static/pages_data.js new file mode 100644 index 0000000..e27af13 --- /dev/null +++ b/chronos/webui/static/pages_data.js @@ -0,0 +1,93 @@ +// Data inspector and Chronos tokenizer builder. +import { state, t, h, post, api, card, tile, pageHead, button, busy, int, num, pct, table, logBox, badge, fieldLabel } from "./core.js"; +import { renderGroup, formState, submit } from "./forms.js"; +import { histogram } from "./charts.js"; +import { pathControl } from "./picker.js"; + +function pathsInput(id, value, onInput) { + const el = h("textarea", { id, class: "paths", spellcheck: "false", autocomplete: "off", placeholder: t("picker.placeholder_multi") }); + el.value = value ?? ""; + const fit = () => { el.rows = Math.max(1, Math.min(8, el.value.split("\n").length)); }; + fit(); + el.addEventListener("input", () => { fit(); onInput(el.value); }); + return el; +} + +export async function renderData(main) { + const v = state.values.data || (state.values.data = { path: "", max: 50000 }); + const out = h("div", { class: "stack" }); + const path = pathsInput("data-path", v.path, (x) => { v.path = x; }); + const max = h("input", { type: "number", id: "data-max", min: 100, max: 500000, step: 100 }); + max.value = String(v.max); + max.addEventListener("input", () => { v.max = Number(max.value) || 50000; }); + const scan = button("ui.scan", () => busy(scan, async () => { + v.result = await post("/api/data/scan", { path: v.path, max_records: v.max }); + show(); + }), { primary: true }); + const show = () => { + const r = v.result; + if (!r) { out.replaceChildren(); return; } + const sm = r.summary || {}; + const fields = Object.entries(sm.fields || {}); + out.replaceChildren( + h("div", { class: "tiles" }, + tile(t("data.files"), int(sm.files), sm.size_mb !== undefined ? `${num(sm.size_mb, 1)} MB` : null), + tile(t("data.scanned"), int(sm.scanned), sm.note === "full scan" ? t("data.full_scan") : t("data.sampled", { n: int(v.max) })), + tile(t("data.median"), int(sm.char_median), t("data.chars")), + tile(t("data.p95"), int(sm.char_p95), t("data.chars")), + tile(t("data.dups"), pct(sm.dup_ratio))), + card(null, histogram({ title: t("data.hist"), edges: r.histogram.edges, counts: r.histogram.counts, xLabel: t("data.chars") })), + h("div", { class: "grid two" }, + card(t("data.schema"), table([{ label: t("data.schema"), key: "k" }, { label: t("data.records"), key: "n", num: true }], + Object.entries(sm.schemas || {}).map(([k, n]) => ({ k, n: int(n) })))), + card(t("data.fields"), table([{ label: t("data.fields"), key: "k" }, { label: t("data.type"), key: "type" }, + { label: t("data.present"), num: true, fmt: (row) => pct(row.present) }], + fields.map(([k, f]) => ({ k, type: f.type, present: f.present }))))), + Object.keys(sm.scripts || {}).length ? card(t("data.scripts"), table([{ label: t("data.script"), key: "k" }, + { label: t("data.share"), num: true, fmt: (row) => pct(row.share) }], Object.entries(sm.scripts).map(([k, share]) => ({ k, share })))) : null, + card(t("data.samples"), ...(r.samples || []).map((x) => h("div", { class: "output mt", text: x }))), + sm.read_errors ? card(t("ui.error"), logBox(sm.read_errors)) : null); + }; + main.append(pageHead("data.title", "data.intro"), + card(null, h("div", { class: "fields" }, + h("div", { class: "field wide", tip: t("data.path_help") }, h("label", { class: "label", for: "data-path", text: t("ui.path") }), + pathControl(path, { kind: "data", multiple: true })), + h("div", { class: "field", tip: t("data.max_help") }, h("label", { class: "label", for: "data-max", text: t("data.max_records") }), max)), + h("div", { class: "row mt" }, scan)), + out); + show(); +} + +// ── tokenizer builder ────────────────────────────────────────────────── +export async function renderTokenizer(main, page) { + const values = formState("tok", "tokenizer"); + const root = h("div"); + const status = h("div", { class: "stack" }); + const start = button("ui.start", () => submit(start, root, async () => { paint(await post("/api/tokenizer/start", values)); }), { primary: true }); + const paint = (s) => { + start.disabled = s.state === "running"; + const sm = s.summary || {}; + status.replaceChildren( + h("div", { class: "row between" }, h("h3", { text: t("tok.progress") }), badge(s.state)), + s.error ? h("p", { class: "error", text: s.error }) : null, + logBox(s.log), + s.state === "done" ? h("div", { class: "tiles" }, + tile(fieldLabel("tok.vocab_size"), int(sm.vocab_size)), + tile(t("tok.vocab_use"), pct(sm.vocab_utilization)), + tile(t("tok.algorithm"), sm.algorithm || "–"), + tile(t("tok.fingerprint"), sm.fingerprint ? String(sm.fingerprint).slice(0, 12) : "–"), + tile(t("tok.elapsed"), sm.elapsed_s === undefined ? "–" : `${num(sm.elapsed_s, 1)} s`)) : null, + s.rows && s.rows.length ? h("div", null, h("h4", { text: t("tok.report") }), table([ + { label: t("tok.bucket"), fmt: (r) => r[0] }, { label: t("tok.texts"), num: true, fmt: (r) => int(r[1]) }, + { label: t("tok.compression"), num: true, fmt: (r) => num(r[2], 2) }, { label: t("tok.chars_per_token"), num: true, fmt: (r) => num(r[3], 2) }, + { label: t("tok.fertility"), num: true, fmt: (r) => num(r[4], 2) }, { label: t("tok.byte_fallback"), num: true, fmt: (r) => pct(r[5], 2) }, + { label: t("tok.roundtrip"), num: true, fmt: (r) => (r[6] === 0 ? t("ui.yes") : `${t("ui.no")} (${int(r[6])})`) }], s.rows), + h("p", { class: "note", text: t("tok.report_note") })) : null, + s.out_dir ? h("p", { class: "note", text: `${t("tok.out_dir_now")}: ${s.out_dir}` }) : null); + }; + root.append(renderGroup("tokenizer", "tokenizer", values, null, { summary: false }), h("div", { class: "row mt" }, start)); + main.append(pageHead("tok.title", "tok.intro"), card(null, root), card(null, status)); + const poll = async () => { const s = await api("/api/tokenizer/state"); paint(s); return s.state === "running" ? 1500 : 6000; }; + page.poll(poll); +} + diff --git a/chronos/webui/static/pages_infer.js b/chronos/webui/static/pages_infer.js new file mode 100644 index 0000000..97c083d --- /dev/null +++ b/chronos/webui/static/pages_infer.js @@ -0,0 +1,94 @@ +// Inference: lazy experts vs full memory on one prompt; architecture comes from the checkpoint. +import { state, t, has, h, api, post, card, pageHead, button, num, int, pct, table, logBox, badge, toast } from "./core.js"; +import { renderGroup, formState, submit } from "./forms.js"; +import { lineChart, barChart } from "./charts.js"; +import { ensureConfig } from "./pages_config.js"; + +const COLS = [ + ["tokens", (v) => int(v)], ["tokens_per_sec", (v) => num(v, 1)], ["response_time_s", (v) => num(v, 2)], + ["prefill_time_s", (v) => num(v, 3)], ["decode_time_s", (v) => num(v, 2)], + ["cache_hit_rate", (v) => pct(v)], ["prediction_hit_rate", (v) => pct(v)], ["cold_miss_predict_hit_rate", (v) => pct(v)], + ["fallback_weight_rate", (v) => pct(v)], ["spec_acceptance_rate", (v) => pct(v)], + ["rss_delta_gb", (v) => num(v, 3)], ["rss_after_decode_gb", (v) => num(v, 2)], +]; +const colLabel = (key) => (has(`infer.col.${key}`) ? t(`infer.col.${key}`) : key); +function modeLabel(mode) { + const m = /^lazy_offload_([\d.]+)x$/.exec(String(mode)); + if (m) return `${t("infer.lazy_output")} · ${m[1]}×`; + return has(`infer.mode.${mode}`) ? t(`infer.mode.${mode}`) : String(mode ?? ""); +} + +export async function renderInference(main, page) { + await ensureConfig(); + const values = formState("infer", "inference"); + const v = state.values.inferPage || (state.values.inferPage = { prompt: "", job: null }); + const root = h("div"); + const out = h("div", { class: "stack" }); + const prompt = h("textarea", { id: "infer-prompt", rows: 4, placeholder: t("infer.prompt") }); + prompt.value = v.prompt; + prompt.addEventListener("input", () => { v.prompt = prompt.value; }); + const go = button("infer.generate", () => submit(go, root, async () => { + const r = await post("/api/infer", { config: state.config, infer: values, prompt: v.prompt }); + v.job = r.job.id; + page.refresh(); + }), { primary: true }); + root.append( + h("div", { class: "field wide", dataset: { key: "prompt" }, tip: t("infer.prompt_help") }, h("label", { class: "label", for: "infer-prompt", text: t("infer.prompt") }), prompt), + h("div", { class: "mt" }, ["inference", "lazy", "speculative"].map((g, i) => renderGroup("inference", g, values, null, { open: i < 2 }))), + h("p", { class: "note", text: t("infer.arch_note") }), + h("div", { class: "row end" }, go)); + main.append(pageHead("infer.title", "infer.intro"), card(null, root), out); + page.poll(async () => { + if (!v.job) return 0; + let job; + try { job = await api(`/api/jobs/${encodeURIComponent(v.job)}?log=200`); } catch (err) { if (err.status === 404) { v.job = null; return 0; } throw err; } + go.disabled = ["running", "starting", "stopping"].includes(job.state); + paint(out, job); + if (job.state === "failed" && !v.reported) { v.reported = true; toast(job.error || t("state.failed"), "error"); } + return go.disabled ? 1500 : 0; + }); +} + +function paint(out, job) { + const r = job.result || {}; + const rows = r.rows || []; + const o = r.outputs || {}; + const diff = r.output_diff; + const blocks = [h("div", { class: "row between" }, h("h3", { text: t("infer.output") }), badge(job.state))]; + if (job.state !== "finished") blocks.push(logBox(job.log)); + if (o.lazy || o.full_dram) { + blocks.push(h("div", { class: "grid two" }, + card(t("infer.lazy_output"), h("div", { class: "output", text: o.lazy || "" })), + card(t("infer.full_output"), h("div", { class: "output", text: o.full_dram || "" })))); + } else if (o.single) blocks.push(h("div", { class: "output", text: o.single })); + if (diff) { + blocks.push(h("p", { class: diff.same ? "note" : "error", + text: `${t(diff.same ? "infer.same" : "infer.diff")} · ${t("infer.similarity")} ${pct(diff.similarity)} · ${t("infer.common_prefix", { n: int(diff.common_prefix_chars) })}` })); + } + if (r.backend) blocks.push(h("p", { class: "note", text: `${t("ui.backend")}: ${r.backend}${r.backend_note ? ` — ${r.backend_note}` : ""}` })); + if (rows.length) { + const used = COLS.filter(([key]) => rows.some((row) => row[key] !== undefined && row[key] !== null)); + blocks.push(card(t("infer.stats"), table([{ label: t("infer.mode"), fmt: (row) => modeLabel(row.mode) }, + ...used.map(([key, fmt]) => ({ label: colLabel(key), num: true, fmt: (row) => (row[key] === undefined || row[key] === null ? "–" : fmt(row[key])) }))], rows), + h("p", { class: "note", text: t("infer.measured_note") }))); + const sweep = rows.filter((row) => row.ram_load_ratio !== undefined && row.ram_load_ratio !== null); + if (sweep.length > 1) { + const pts = (key) => sweep.map((row) => [row.ram_load_ratio, row[key]]).filter((p) => p[1] !== null && p[1] !== undefined); + blocks.push(card(t("infer.sweep"), h("div", { class: "grid two" }, + lineChart({ title: colLabel("tokens_per_sec"), xLabel: t("infer.ram_ratio"), xFmt: (x) => `${num(x, 2)}×`, yFmt: (y) => num(y, 1), + series: [{ label: colLabel("tokens_per_sec"), color: "var(--s1)", points: pts("tokens_per_sec") }] }), + lineChart({ title: t("infer.hit_rates"), xLabel: t("infer.ram_ratio"), xFmt: (x) => `${num(x, 2)}×`, yFmt: (y) => pct(y, 0), + series: [["cache_hit_rate", "var(--s1)"], ["prediction_hit_rate", "var(--s2)"]].map(([key, c]) => ({ label: colLabel(key), color: c, points: pts(key) })) }), + lineChart({ title: colLabel("rss_delta_gb"), xLabel: t("infer.ram_ratio"), xFmt: (x) => `${num(x, 2)}×`, yFmt: (y) => num(y, 2), + series: [{ label: colLabel("rss_delta_gb"), color: "var(--s1)", points: pts("rss_delta_gb") }] })))); + } else if (rows.length > 1) { + blocks.push(card(null, h("div", { class: "grid two" }, + barChart({ title: colLabel("tokens_per_sec"), categories: rows.map((row) => modeLabel(row.mode)), + series: [{ label: colLabel("tokens_per_sec"), values: rows.map((row) => Number(row.tokens_per_sec)) }], fmt: (x) => num(x, 1) }), + barChart({ title: colLabel("rss_delta_gb"), categories: rows.map((row) => modeLabel(row.mode)), + series: [{ label: colLabel("rss_delta_gb"), values: rows.map((row) => Number(row.rss_delta_gb)) }], fmt: (x) => num(x, 3) })))); + } + } + if (job.error) blocks.push(h("p", { class: "error", text: job.error })); + out.replaceChildren(card(null, ...blocks)); +} diff --git a/chronos/webui/static/pages_monitor.js b/chronos/webui/static/pages_monitor.js new file mode 100644 index 0000000..fefe7fd --- /dev/null +++ b/chronos/webui/static/pages_monitor.js @@ -0,0 +1,162 @@ +// Run monitor: any run folder under runs_dir (Web UI jobs and the CLI long runs). +// Reads only; never signals the training process. +import { state, t, has, h, api, card, tile, pageHead, badge, num, int, params, duration, dateTime, logBox, table, toast } from "./core.js"; +import { lineChart } from "./charts.js"; +import { probeSection } from "./pages_probe.js"; + +// Small multiples: one chart per measure family, one y-axis each. +const CHARTS = [ + { id: "loss", keys: ["loss", "primary_loss"] }, + { id: "aux", keys: ["balance_loss", "lookahead_loss", "temporal_loss", "locality_loss"] }, + { id: "throughput", keys: ["tokens_per_s"] }, + { id: "memory", keys: ["peak_mem_gb", "active_mem_gb"] }, + { id: "optimizer", keys: ["grad_norm"] }, + { id: "lr", keys: ["lr"] }, + { id: "recurrent", keys: ["rec_depth_mean", "rec_contraction", "rec_early_exit_ratio"] }, + { id: "routing", keys: ["route_entropy", "load_cv"] }, +]; +const metricLabel = (key) => (has(`metric.${key}`) ? t(`metric.${key}`) : key); +const chartFmt = { lr: (v) => num(v, 6), throughput: (v) => num(v, 0) }; + +function phaseLabel(name) { return has(`stage.${name}`) ? t(`stage.${name}`) : String(name ?? "–"); } + +export async function renderMonitor(main, page) { + const runSel = h("select", { id: "mon-run" }); + const head = h("div", { class: "row" }, h("label", { class: "compact", for: "mon-run", text: t("mon.run") }), runSel); + const body = h("div", { class: "stack" }, h("p", { class: "note", text: t("ui.loading") })); + main.append(pageHead("mon.title", "mon.intro", head), body); + runSel.addEventListener("change", () => { state.values.monitorRun = runSel.value; page.refresh(); }); + page.poll(async () => { + const run = state.values.monitorRun; + let snap; + try { + snap = await api(`/api/monitor/state?history=1${run ? `&run=${encodeURIComponent(run)}` : ""}`); + } catch (err) { + if (err.status === 404 && run) { state.values.monitorRun = ""; toast(t("err.not_found"), "error"); return 100; } + throw err; + } + const names = (snap.runs || []).map((r) => r.name); + if (runSel.options.length !== names.length || [...runSel.options].some((o, i) => o.value !== names[i])) { + runSel.replaceChildren(...names.map((n) => h("option", { value: n, text: n }))); + } + if (snap.run) runSel.value = snap.run; + if (!snap.run) { body.replaceChildren(h("div", { class: "empty", text: t("mon.no_runs") })); return 10000; } + paint(body, snap); + return snap.status && snap.status.state === "running" ? 5000 : 15000; + }); +} + +function paint(body, snap) { + const st = snap.status || {}; + const pr = snap.progress || {}; + const latest = snap.latest || {}; + const mem = snap.memory || {}; + const best = snap.best_valid || {}; + const current = pr.current || latest.phase; + const bestNow = best[current]; + const phaseBar = (pr.phases || []).map((p) => h("div", { class: "phase-row" }, + h("span", { class: "phase-name", text: phaseLabel(p.name) }), + h("div", { class: "progress", role: "progressbar", "aria-valuemin": 0, "aria-valuemax": p.total || 0, "aria-valuenow": p.step || 0, + "aria-label": phaseLabel(p.name) }, bar(p.total ? Math.min(1, p.step / p.total) : (p.state === "complete" ? 1 : 0))), + h("span", { class: "note", text: `${int(p.step)} / ${p.total ? int(p.total) : "–"} · ${t(`phase.${p.state}`)}` }))); + body.replaceChildren( + h("div", { class: "row between" }, + h("div", { class: "row" }, badge(st.stale ? "stale" : st.state), st.stale ? h("span", { class: "note", text: t("mon.stale") }) : null, + st.last_update_age_s !== null && st.last_update_age_s !== undefined ? h("span", { class: "note", text: t("mon.updated", { s: int(st.last_update_age_s) }) }) : null, + st.pid ? h("span", { class: "note mono", text: `PID ${st.pid}` }) : null), + st.error ? h("span", { class: "error", text: st.error }) : null), + h("div", { class: "tiles" }, + tile(t("ui.phase"), phaseLabel(current)), + tile(t("ui.step"), int(latest.step), latest.total ? `/ ${int(latest.total)}` : null), + tile(t("mon.loss"), num(latest.loss, 4), latest.primary_loss !== undefined ? `${metricLabel("primary_loss")} ${num(latest.primary_loss, 4)}` : null), + tile(t("mon.valid_ce"), bestNow ? num(bestNow.valid_ce, 4) : "–", bestNow ? `${t("mon.best")} · ${t("ui.step")} ${int(bestNow.step)}` : null), + tile(t("mon.throughput"), num(latest.tokens_per_s, 0), latest.tokens_seen ? `${params(latest.tokens_seen)} ${t("mon.tokens_seen")}` : null), + tile(t("mon.memory"), mem.footprint_gb !== undefined ? `${num(mem.footprint_gb, 1)} GB` : (latest.peak_mem_gb !== undefined ? `${num(latest.peak_mem_gb, 1)} GB` : "–"), + mem.state === "measured" ? t("label.measured") : null), + tile(t("mon.eta"), duration(pr.run_remaining_s), pr.eta ? dateTime(pr.eta) : null), + tile(t("mon.lr"), num(latest.lr, 6), latest.grad_norm !== undefined ? `${t("mon.grad_norm")} ${num(latest.grad_norm, 3)}` : null)), + phaseBar.length ? card(t("mon.progress"), ...phaseBar) : null, + probeSection(snap), + h("div", { class: "grid two" }, CHARTS.map((c) => chart(c, snap))), + evalCharts(snap), + eventsCard(snap), + samplesCard(snap), + checkpointsCard(snap), + logCard(snap)); +} + +function bar(fraction) { + const fill = h("div"); + fill.style.width = `${Math.round(Math.max(0, Math.min(1, fraction)) * 1000) / 10}%`; + return fill; +} +// Steps restart in every phase, so charts show one phase at a time (chosen in the card header). +function viewPhase(snap) { + const phases = Object.keys(snap.history || {}); + const want = state.values.monitorPhase; + return phases.includes(want) ? want : ((snap.progress && snap.progress.current) || phases[phases.length - 1]); +} + +function chart(c, snap) { + const rows = (snap.history || {})[viewPhase(snap)] || []; + const series = c.keys.map((key, i) => ({ label: metricLabel(key), color: `var(--s${i + 1})`, + points: rows.filter((r) => r[key] !== null && r[key] !== undefined).map((r) => [r.step, Number(r[key])]) })); + return card(null, lineChart({ title: t(`mon.chart.${c.id}`), series, xLabel: t("ui.step"), yFmt: chartFmt[c.id] || ((v) => num(v, 3)) })); +} + +function evalCharts(snap) { + const phase = viewPhase(snap); + const vals = (snap.events || []).filter((e) => e.event === "validation" && e.phase === phase); + const keys = (snap.evals && snap.evals.keys) || ["valid_ce"]; + const charts = keys.filter((k) => k !== "valid_ppl" && vals.some((e) => e[k] !== undefined && e[k] !== null)).map((k, i) => + card(null, lineChart({ title: metricLabel(k), xLabel: t("ui.step"), + series: [{ label: metricLabel(k), color: `var(--s${(i % 8) + 1})`, points: vals.filter((e) => e[k] !== null && e[k] !== undefined).map((e) => [e.step, Number(e[k])]) }] }))); + const phases = Object.keys(snap.history || {}); + const picker = h("select", { id: "mon-phase" }, phases.map((p) => h("option", { value: p, text: phaseLabel(p) }))); + picker.value = phase || ""; + picker.addEventListener("change", () => { state.values.monitorPhase = picker.value; document.dispatchEvent(new CustomEvent("chronos:refresh")); }); + return card(t("mon.validation"), h("div", { class: "row" }, h("label", { class: "compact", for: "mon-phase", text: t("ui.phase") }), picker, + h("span", { class: "note", text: t("mon.phase_note") })), + charts.length ? h("div", { class: "grid two mt" }, charts) : h("p", { class: "note mt", text: t("ui.empty") })); +} + +function eventsCard(snap) { + const events = (snap.events || []).slice(-40).reverse(); + if (!events.length) return null; + const known = ["event", "phase", "step", "time"]; + return card(t("mon.events"), table([ + { label: t("ui.time"), fmt: (e) => dateTime(e.time) }, + { label: t("ui.phase"), fmt: (e) => phaseLabel(e.phase) }, + { label: t("ui.step"), num: true, fmt: (e) => int(e.step) }, + { label: t("mon.event"), fmt: (e) => (has(`event.${e.event}`) ? t(`event.${e.event}`) : e.event) }, + { label: t("ui.value"), fmt: (e) => Object.entries(e).filter(([k, v]) => !known.includes(k) && v !== null && typeof v !== "object") + .slice(0, 6).map(([k, v]) => `${metricLabel(k)} ${typeof v === "number" ? num(v, 4) : v}`).join(" · ") }, + ], events)); +} + +function samplesCard(snap) { + const samples = (snap.samples || []).slice(-6).reverse(); + if (!samples.length) return null; + return card(t("mon.samples"), samples.map((x) => h("div", { class: "sample" }, + h("div", { class: "note", text: `${phaseLabel(x.phase)} · ${t("ui.step")} ${int(x.step)}` }), + x.input ? h("div", null, h("b", { text: `${t("mon.input")}: ` }), x.input) : null, + x.reference ? h("div", null, h("b", { text: `${t("mon.reference")}: ` }), x.reference) : null, + h("div", { class: "output" }, x.output || "")))); +} + +function checkpointsCard(snap) { + const rows = snap.checkpoints || []; + if (!rows.length) return null; + return card(t("mon.checkpoints"), table([ + { label: t("ui.path"), fmt: (r) => h("span", { class: "mono", text: r.file }) }, + { label: t("mon.size"), num: true, fmt: (r) => `${num(r.size_gb, 3)} GB` }, + { label: t("ui.time"), fmt: (r) => dateTime(r.modified) }, + ], rows.slice(0, 12))); +} + +function logCard(snap) { + const box = h("div", null, h("p", { class: "note", text: t("ui.loading") })); + api(`/api/monitor/log?run=${encodeURIComponent(snap.run)}&bytes=12000`).then((r) => box.replaceChildren(logBox(r.text))) + .catch(() => box.replaceChildren(h("p", { class: "note", text: t("ui.empty") }))); + return h("details", { class: "card" }, h("summary", { text: t("mon.log") }), box); +} diff --git a/chronos/webui/static/pages_ops.js b/chronos/webui/static/pages_ops.js new file mode 100644 index 0000000..0da89fc --- /dev/null +++ b/chronos/webui/static/pages_ops.js @@ -0,0 +1,144 @@ +// Export, Ollama/OpenAI server, benchmark, auto-tune and IO metrics pages. +import { state, t, h, api, post, card, pageHead, button, num, int, table, logBox, badge, toast, fieldLabel, dateTime } from "./core.js"; +import { renderGroup, formState, submit, assign } from "./forms.js"; +import { lineChart, barChart } from "./charts.js"; +import { ensureConfig } from "./pages_config.js"; + +const active = (s) => ["running", "starting", "stopping", "loading"].includes(s); + +// ── export ───────────────────────────────────────────────────────────── +export async function renderExport(main, page) { + const values = formState("export", "export"); + const v = state.values.exportPage || (state.values.exportPage = { job: null }); + const root = h("div"); + const out = h("div"); + const run = button("export.run", () => submit(run, root, async () => { + v.job = (await post("/api/export", values)).job.id; + page.refresh(); + }), { primary: true }); + root.append(h("p", { class: "note", text: t("export.note") }), renderGroup("export", "export", values, null, { summary: false }), + h("div", { class: "row end mt" }, run)); + main.append(pageHead("export.title", "export.intro"), card(null, root), out); + page.poll(async () => { + if (!v.job) return 0; + const job = await api(`/api/jobs/${encodeURIComponent(v.job)}?log=300`); + run.disabled = active(job.state); + out.replaceChildren(card(null, h("div", { class: "row between" }, h("h3", { text: t("ui.result") }), badge(job.state)), + Array.isArray(job.result) ? table([ + { label: t("export.format"), fmt: (r) => r.format }, { label: t("ui.path"), fmt: (r) => h("span", { class: "mono", text: r.path }) }, + { label: t("export.size"), num: true, fmt: (r) => num(r.size_mb, 2) }, { label: t("export.tensors"), num: true, fmt: (r) => int(r.tensors) }, + ], job.result) : null, + job.error ? h("p", { class: "error", text: job.error }) : null, logBox(job.log))); + return run.disabled ? 1500 : 0; + }); +} + +// ── serve ────────────────────────────────────────────────────────────── +export async function renderServe(main, page) { + const values = formState("serve", "serve"); + const root = h("div"); + const status = h("div"); + const start = button("serve.start", () => submit(start, root, async () => { paint(await post("/api/serve/start", values)); }), { primary: true }); + const stop = button("serve.stop", async () => { paint(await post("/api/serve/stop")); toast(t("serve.stopped")); }, { danger: true }); + const paint = (s) => { + start.disabled = s.running; + stop.disabled = !s.running; + status.replaceChildren(card(null, + h("div", { class: "row between" }, h("h3", { text: s.running ? t("serve.running", { url: s.url }) : t("serve.stopped") }), + badge(s.running ? "running" : "stopped")), + s.running ? h("pre", { class: "log", text: t("serve.hint", { addr: s.addr }) }) : null, + s.running ? h("p", { class: "note", text: s.auth ? t("serve.auth_on") : t("serve.auth_off") }) : null, + logBox(s.log))); + }; + root.append(h("p", { class: "note", text: t("serve.remote_warning") }), renderGroup("serve", "serve", values, null, { summary: false }), + h("div", { class: "row end mt" }, stop, start)); + main.append(pageHead("serve.title", "serve.intro"), card(null, root), status); + page.poll(async () => { const s = await api("/api/serve/state"); paint(s); return s.running ? 3000 : 10000; }); +} + +// ── benchmark ────────────────────────────────────────────────────────── +export async function renderBenchmark(main, page) { + const v = state.values.benchPage || (state.values.benchPage = { job: null }); + const out = h("div", { class: "stack" }); + const run = button("bench.run", async () => { + run.disabled = true; + try { v.job = (await post("/api/benchmark/start")).job.id; page.refresh(); } catch (err) { run.disabled = false; throw err; } + }, { primary: true }); + main.append(pageHead("bench.title", "bench.intro", run), out); + page.poll(async () => { + const res = await api("/api/benchmark/results"); + let job = null; + if (v.job) { try { job = await api(`/api/jobs/${encodeURIComponent(v.job)}?log=200`); } catch { v.job = null; } } + run.disabled = !res.available || (job && active(job.state)); + const results = res.results || {}; + const names = Object.keys(results); + const charts = res.metrics.filter((m) => names.some((n) => Number.isFinite(results[n][m.key]))).map((m) => card(null, barChart({ + title: t(m.label), categories: names.map((n) => results[n].name || n), fmt: (x) => num(x, 2), + series: [{ label: t(m.label), values: names.map((n) => Number(results[n][m.key])) }], + note: t(m.lower_is_better ? "bench.lower" : "bench.higher") }))); + out.replaceChildren( + !res.available ? h("p", { class: "note", text: t("bench.unavailable") }) : null, + names.length ? h("div", { class: "grid two" }, charts) : h("div", { class: "empty", text: t("bench.none") }), + names.length ? h("p", { class: "note", text: t("bench.measured_note") }) : null, + job ? card(null, h("div", { class: "row between" }, h("h3", { text: t("ui.log") }), badge(job.state)), logBox(job.log), + job.error ? h("p", { class: "error", text: job.error }) : null) : null); + return job && active(job.state) ? 2000 : 0; + }); +} +// ── auto-tune ────────────────────────────────────────────────────────── +export async function renderAutotune(main, page) { + await ensureConfig(); + const values = formState("tune", "autotune"); + const root = h("div"); + const out = h("div", { class: "stack" }); + const start = button("tune.start", () => submit(start, root, async () => { paint(await post("/api/autotune/start", { tune: values, config: state.config })); }), { primary: true }); + const stop = button("tune.stop", async () => paint(await post("/api/autotune/stop")), { danger: true }); + const paint = (s) => { + start.disabled = active(s.state); + stop.disabled = !active(s.state); + const trials = s.trials || []; + const best = s.best || {}; + const ok = trials.filter((tr) => tr.status === "complete" && Number.isFinite(tr.loss)); + const keys = [...new Set(trials.flatMap((tr) => Object.keys(tr.params || {})))]; + out.replaceChildren( + card(null, h("div", { class: "row between" }, h("h3", { text: t("tune.trials") }), badge(s.state)), + ok.length ? lineChart({ title: t("tune.loss"), xLabel: t("tune.trial"), yFmt: (y) => num(y, 4), + series: [{ label: t("tune.loss"), color: "var(--s1)", points: ok.map((tr) => [tr.trial, tr.loss]) }] }) : null, + trials.length ? table([{ label: t("tune.trial"), num: true, fmt: (tr) => int(tr.trial) }, + { label: t("tune.loss"), num: true, fmt: (tr) => num(tr.loss, 4) }, + { label: t("ui.state"), fmt: (tr) => (tr.status && t(`trial.${tr.status}`)) || tr.status }, + { label: t("tune.duration"), num: true, fmt: (tr) => `${num(tr.duration_s, 1)} s` }, + ...keys.map((k) => ({ label: fieldLabel(k), num: true, fmt: (tr) => ((tr.params || {})[k] === undefined ? "" : num(tr.params[k], 6)) }))], + trials.slice().reverse()) : null, + h("p", { class: "note", text: t("tune.note") })), + Object.keys(best).length ? card(t("tune.best"), table([{ label: t("ui.value"), fmt: (r) => fieldLabel(r[0]) }, { label: "", num: true, fmt: (r) => num(r[1], 6) }], + Object.entries(best)), h("div", { class: "row end mt" }, button("tune.apply", () => { + assign(state.config, best); + toast(t("config.loaded")); + }, { primary: true }))) : null, + card(t("ui.log"), logBox(s.log))); + }; + root.append(renderGroup("autotune", "autotune", values, null, { summary: false }), h("div", { class: "row end mt" }, stop, start)); + main.append(pageHead("tune.title", "tune.intro"), card(null, root), out); + page.poll(async () => { const s = await api("/api/autotune/state"); paint(s); return active(s.state) ? 2000 : 10000; }); +} + +// ── IO / cache metrics of this server process ────────────────────────── +export async function renderIO(main, page) { + const out = h("div", { class: "stack" }); + main.append(pageHead("io.title", "io.intro"), out); + page.poll(async () => { + const r = await api("/api/io?samples=240"); + const names = Object.keys(r.series || {}); + if (!names.length) { out.replaceChildren(h("div", { class: "empty", text: t("io.empty") })); return 4000; } + const t0 = Math.min(...names.map((n) => (r.series[n][0] || [Infinity])[0])); + out.replaceChildren( + card(t("io.latest"), table([{ label: t("io.metric"), fmt: (n) => h("span", { class: "mono", text: n }) }, + { label: t("io.latest"), num: true, fmt: (n) => num(r.series[n][r.series[n].length - 1][1], 4) }, + { label: t("io.samples"), num: true, fmt: (n) => int(r.series[n].length) }, + { label: t("ui.time"), fmt: (n) => dateTime(r.series[n][r.series[n].length - 1][0]) }], names)), + h("div", { class: "grid two" }, names.slice(0, 12).map((n) => card(null, lineChart({ title: n, xLabel: t("io.seconds"), + series: [{ label: n, color: "var(--s1)", points: r.series[n].map(([ts, val]) => [ts - t0, val]) }], yFmt: (y) => num(y, 3) }))))); + return 3000; + }); +} diff --git a/chronos/webui/static/pages_probe.js b/chronos/webui/static/pages_probe.js new file mode 100644 index 0000000..16d94cd --- /dev/null +++ b/chronos/webui/static/pages_probe.js @@ -0,0 +1,96 @@ +// Lazy-expert probe results in the run monitor: efficiency, hit rates and +// fidelity of every probed variant against the full-memory reference. +import { state, t, has, h, card, num, int, pct, dateTime, table } from "./core.js"; +import { lineChart } from "./charts.js"; + +const SKIPS = [ + [/^dense model/, "probe.skip.dense"], [/^no checkpoint/, "probe.skip.no_checkpoint"], + [/^previous probe/, "probe.skip.busy"], [/^checkpoint still/, "probe.skip.writing"], + [/^superseded by step (\d+)/, "probe.skip.superseded"], +]; +function skipText(reason) { + for (const [re, key] of SKIPS) { + const m = re.exec(String(reason || "")); + if (m) return t(key, { step: m[1] || "" }); + } + return String(reason || ""); +} + +const phaseLabel = (name) => (has(`stage.${name}`) ? t(`stage.${name}`) : String(name ?? "–")); +const yesNo = (v) => (v === null || v === undefined ? "–" : t(v ? "ui.yes" : "ui.no")); + +function variantColumns() { + const m = (key, fmt) => (v) => (v.error ? "" : fmt((v.measured || {})[key])); + const f = (key, fmt) => (v) => (v.error ? "" : fmt((v.fidelity || {})[key])); + return [ + { label: t("probe.variant"), fmt: (v) => h("span", { class: "mono", text: v.label }) }, + { label: t("probe.exact"), num: true, fmt: f("exact_match_rate", (x) => pct(x)) }, + { label: t("probe.agree"), num: true, fmt: f("positionwise_agreement", (x) => pct(x)) }, + { label: t("probe.prefix"), num: true, fmt: f("min_matching_prefix", (x) => int(x)) }, + { label: t("probe.hit"), num: true, fmt: m("resident_hit_rate", (x) => pct(x)) }, + { label: t("probe.pred_hit"), num: true, fmt: m("prediction_hit_rate", (x) => pct(x)) }, + { label: t("probe.cold_hit"), num: true, fmt: m("cold_miss_predict_hit_rate", (x) => pct(x)) }, + { label: t("probe.loads"), num: true, fmt: m("on_demand_loads", (x) => int(x)) }, + { label: t("probe.wait"), num: true, fmt: m("prefetch_wait_time_s", (x) => num(x, 3)) }, + { label: t("probe.tps"), num: true, fmt: m("tokens_per_s", (x) => num(x, 1)) }, + { label: t("probe.vs_full"), num: true, fmt: m("speed_vs_full", (x) => (x === null || x === undefined ? "–" : `${num(x, 2)}×`)) }, + { label: t("probe.accept"), num: true, fmt: m("spec_acceptance_rate", (x) => pct(x)) }, + { label: t("probe.rss"), num: true, fmt: m("peak_rss_gb", (x) => num(x, 2)) }, + { label: t("probe.pages"), num: true, fmt: (v) => (v.error ? "" : `${(v.measured || {}).paging_resolved ?? "–"} · ${int((v.measured || {}).num_pages)}`) }, + { label: t("ui.error"), fmt: (v) => (v.error ? h("span", { class: "error", text: v.error }) : "") }, + ]; +} + +export function probeSection(snap) { + const probes = snap.lazy_probe || []; + if (!probes.length) return card(t("probe.title"), h("p", { class: "note", text: t("probe.none") })); + const done = probes.filter((p) => Array.isArray(p.variants)); + const last = probes[probes.length - 1]; + const latest = done[done.length - 1]; + const parts = []; + if (last !== latest && (last.skipped || last.error)) { + parts.push(h("p", { class: "note", text: `${phaseLabel(last.phase)} · ${t("ui.step")} ${int(last.step)}: ${last.skipped ? t("probe.skipped", { reason: skipText(last.skipped) }) : last.error}` })); + } + if (latest) { + const c = latest.conditions || {}; + parts.push( + h("div", { class: "row" }, + h("span", { class: "pill", text: `${phaseLabel(latest.phase)} · ${t("ui.step")} ${int(latest.step)}` }), + h("span", { class: "note", text: dateTime(latest.time) }), + h("span", { class: "note", text: `${t("probe.trained_clusters")}: ${yesNo(latest.trained_with_clusters)}` }), + h("span", { class: "note", text: `${t("probe.auto_paging")}: ${latest.auto_paging ?? "–"}` }), + latest.truncated ? h("span", { class: "note", text: t("train.variants_capped", { n: int(latest.variants_requested), cap: int(latest.variants.length) }) }) : null), + h("p", { class: "note", text: t("probe.conditions", { n: int(c.prompts), t: int(c.max_new_tokens) }) }), + table(variantColumns(), latest.variants), + Object.keys(latest.references || {}).length ? h("p", { class: "note", text: `${t("probe.reference")}: ${Object.entries(latest.references) + .map(([b, tps]) => `${b} ${num(tps, 1)} ${t("probe.tps")}`).join(" · ")}` }) : null, + historyChart(done)); + } + return card(t("probe.title"), ...parts); +} + +// History of one chosen variant over training (sync_on_demand first by default). +function historyChart(all) { + // Steps restart in every phase: plot the newest phase only. + const phase = all[all.length - 1].phase; + const done = all.filter((p) => p.phase === phase); + if (done.length < 2) return null; + const labels = [...new Set(done.flatMap((p) => p.variants.filter((v) => !v.error).map((v) => v.label)))]; + if (!labels.length) return null; + const want = state.values.probeVariant; + const pick = labels.includes(want) ? want : (labels.find((l) => l.includes("sync_on_demand")) || labels[0]); + const series = [ + ["probe.hit", (v) => v.measured.resident_hit_rate], + ["probe.pred_hit", (v) => v.measured.prediction_hit_rate], + ["probe.agree", (v) => v.fidelity.positionwise_agreement], + ["probe.exact", (v) => v.fidelity.exact_match_rate], + ].map(([key, get], i) => ({ label: t(key), color: `var(--s${i + 1})`, + points: done.map((p) => { const v = p.variants.find((x) => x.label === pick && !x.error); return v ? [p.step, get(v)] : null; }) + .filter((pt) => pt && pt[1] !== null && pt[1] !== undefined) })); + const sel = h("select", { id: "probe-variant" }, labels.map((l) => h("option", { value: l, text: l }))); + sel.value = pick; + sel.addEventListener("change", () => { state.values.probeVariant = sel.value; document.dispatchEvent(new CustomEvent("chronos:refresh")); }); + return h("div", { class: "mt" }, + h("div", { class: "row" }, h("label", { class: "compact", for: "probe-variant", text: t("probe.variant") }), sel), + lineChart({ title: `${t("probe.history")} · ${phaseLabel(phase)}`, series, xLabel: t("ui.step"), yFmt: (v) => pct(v, 0) })); +} diff --git a/chronos/webui/static/pages_train.js b/chronos/webui/static/pages_train.js new file mode 100644 index 0000000..5ae2281 --- /dev/null +++ b/chronos/webui/static/pages_train.js @@ -0,0 +1,84 @@ +// Train page: job form, lazy-expert probe sweep (runs after every validation), job list. +import { state, t, h, has, api, post, card, pageHead, button, badge, dateTime, logBox, toast, table, int } from "./core.js"; +import { renderGroup, formState, submit } from "./forms.js"; +import { ensureConfig } from "./pages_config.js"; + +let previewTimer = 0; +const STAGE_ORDER = ["pretrain", "sft", "dpo", "orpo", "grpo", "distill"]; +const stageName = (m) => (has(`stage.${m}`) ? t(`stage.${m}`) : m); + +// One data row per selected stage after the first (canonical order); the first stage uses train.data_path. +function paintStageData(host, train) { + const later = STAGE_ORDER.filter((m) => train["train.stages"].includes(m)).slice(1); + host.hidden = !later.length; + for (const el of host.querySelectorAll(".field[data-key^='train.data_']")) { + el.hidden = !later.includes(el.dataset.key.slice("train.data_".length)); + } +} + +async function preview(out, values) { + try { + const p = await post("/api/probe/preview", values); + const text = p.requested > p.variants ? t("train.variants_capped", { n: int(p.requested), cap: int(p.variants) }) + : t("train.variants", { n: int(p.variants) }); + out.replaceChildren(h("p", { class: "note", text: values["probe.enabled"] ? text : t("train.probe_off") }), + values["probe.enabled"] && p.labels.length ? h("details", null, h("summary", { text: t("probe.variant") }), + h("ul", { class: "mono" }, p.labels.map((l) => h("li", { text: l })))) : null); + } catch { + out.replaceChildren(h("p", { class: "error", text: t("ui.invalid") })); + } +} + +export async function renderTrain(main, page) { + await ensureConfig(); + const train = formState("train", "train"); + const probe = formState("probe", "probe"); + const root = h("div"); + const variants = h("div"); + const jobs = h("div"); + const stageData = card(t("group.stage_data"), h("p", { class: "note", text: t("train.stage_data_note") }), + renderGroup("train", "stage_data", train, null, { summary: false })); + const onTrain = (key) => { if (key === "train.stages") paintStageData(stageData, train); }; + const onProbe = () => { clearTimeout(previewTimer); previewTimer = setTimeout(() => preview(variants, probe), 300); }; + const start = button("train.start", () => submit(start, root, async () => { + const r = await post("/api/train/start", { config: state.config, train, probe }); + state.values.monitorRun = r.run; + toast(`${t("train.started")}: ${r.run}`); + page.refresh(); + }), { primary: true }); + root.append( + card(t("group.train"), h("p", { class: "note", text: t("train.config_note") }), renderGroup("train", "train", train, onTrain, { summary: false })), + stageData, + card(t("train.probe_title"), h("p", { class: "note", text: t("train.probe_note") }), + renderGroup("probe", "probe", probe, onProbe, { summary: false }), variants), + h("div", { class: "row end mt" }, start)); + main.append(pageHead("train.title", "train.intro"), root, card(t("train.jobs"), jobs)); + paintStageData(stageData, train); + preview(variants, probe); + page.poll(async () => { + const r = await api("/api/jobs?kind=train"); + paintJobs(jobs, r.jobs || [], page); + return (r.jobs || []).some((j) => ["running", "starting", "stopping"].includes(j.state)) ? 2000 : 8000; + }); +} + +function paintJobs(host, list, page) { + if (!list.length) { host.replaceChildren(h("p", { class: "note", text: t("ui.empty") })); return; } + host.replaceChildren(table([ + { label: t("train.run_dir"), fmt: (j) => (j.meta && j.meta.run) || j.id }, + { label: t("ui.phase"), fmt: (j) => ((j.meta && j.meta.stages) || []).map(stageName).join(" → ") || j.title }, + { label: t("ui.state"), fmt: (j) => badge(j.state) }, + { label: t("ui.started"), fmt: (j) => dateTime(j.started) }, + { label: "", fmt: (j) => h("div", { class: "row" }, + h("a", { href: "#monitor", text: t("train.open_monitor"), onclick: () => { state.values.monitorRun = j.meta && j.meta.run; } }), + ["running", "starting"].includes(j.state) ? button("train.stop", async (e) => { + e.target.disabled = true; + try { await post(`/api/jobs/${encodeURIComponent(j.id)}/stop`); } finally { page.refresh(); } + }, { danger: true }) : null, + button("ui.log", async () => { + const full = await api(`/api/jobs/${encodeURIComponent(j.id)}?log=400`); + host.append(card(`${t("ui.log")} · ${(j.meta && j.meta.run) || j.id}`, logBox(full.log), + full.error ? h("p", { class: "error", text: full.error }) : null)); + })) }, + ], list.slice().reverse())); +} diff --git a/chronos/webui/static/picker.js b/chronos/webui/static/picker.js new file mode 100644 index 0000000..c2f09a8 --- /dev/null +++ b/chronos/webui/static/picker.js @@ -0,0 +1,334 @@ +// File / folder picking for path fields. Two ways in, both labelled in the user's language: +// "Browse…" lists the server's folders through /api/fs/list (in-page dialog). +// "Upload…" uses the system file manager. On the server's own machine the server opens the native +// dialog and the real path is used in place; from another machine (LAN access) the files +// are uploaded into /uploads/… and that server path is used. +// kind: dir | data | checkpoint | config | any. multiple: several files/folders, one per line. +import { state, t, has, h, api, post, clear, bytes, dateTime, reportError, toast, errorText, ApiError } from "./core.js"; + +const lastDir = {}; // per kind, for this browser session + +function splitPaths(text) { + return String(text || "").split(/\r?\n/).map((x) => x.trim()).filter(Boolean); +} + +// Opens the picker; resolves to a list of paths, or null when cancelled. +export function openPicker({ kind = "any", multiple = false, start = "", title } = {}) { + return new Promise((resolve) => { + const picked = new Set(); + let view = null; + let cursor = -1; + const dlg = h("dialog", { class: "picker", "aria-labelledby": "picker-title" }); + const where = h("input", { type: "text", spellcheck: "false", autocomplete: "off", "aria-label": t("ui.path") }); + const list = h("ul", { class: "picker-list", role: "listbox", tabindex: "0", "aria-multiselectable": multiple ? "true" : null, + "aria-label": t("picker.entries") }); + const roots = h("div", { class: "picker-roots", role: "group", "aria-label": t("picker.places") }); + const status = h("p", { class: "note", "aria-live": "polite" }); + const hidden = h("input", { type: "checkbox", id: "picker-hidden" }); + const choose = h("button", { type: "button", class: "primary" }); + const thisFolder = kind === "dir" || kind === "data" || kind === "any"; + const finish = (value) => { dlg.close(); dlg.remove(); resolve(value); }; + + let seq = 0; // only the newest folder request may repaint + const load = async (path) => { + const mine = ++seq; + let next; + try { + next = await api(`/api/fs/list?${new URLSearchParams({ path, kind, hidden: hidden.checked ? "true" : "false" })}`); + } catch (err) { + if (mine !== seq) return; + if (err.status === 404 && path) return load(""); + reportError(err); + return; + } + if (mine !== seq || !dlg.isConnected) return; + view = next; + lastDir[kind] = view.path; + where.value = view.path; + cursor = view.entries.length ? 0 : -1; + paint(); + }; + const entryAt = (i) => (view && i >= 0 ? view.entries[i] : null); + const selectable = (e) => e && (!e.dir || thisFolder); + const toggle = (e) => { + if (!selectable(e)) return; + if (!multiple) picked.clear(); + if (picked.has(e.path)) picked.delete(e.path); else picked.add(e.path); + mark(); + }; + const open = (e) => { if (e && e.dir) load(e.path); else if (e && !multiple) finish([e.path]); }; + const result = () => (picked.size ? [...picked] : (thisFolder && view ? [view.path] : [])); + + // The list is rebuilt only when the folder changes; selection updates attributes in place, so a + // double-click lands on the same element as its first click. + function paint() { + clear(list); + view.entries.forEach((e, i) => { + let before = false; // selection before the first click of a double-click + const li = h("li", { role: "option", id: `pick-${i}`, class: `picker-item${e.dir ? " dir" : ""}` }, + icon(e.dir), h("span", { class: "picker-name", text: e.name }), + h("span", { class: "picker-meta", text: e.dir ? "" : bytes(e.size) }), + h("span", { class: "picker-meta", text: e.mtime ? dateTime(e.mtime) : "" })); + li.addEventListener("click", (ev) => { + cursor = i; + if (ev.detail > 1) return; + before = picked.has(e.path); + if (selectable(e)) toggle(e); else mark(); + }); + li.addEventListener("dblclick", () => { + if (e.dir) { if (!before) picked.delete(e.path); load(e.path); return; } + if (!multiple) finish([e.path]); + }); + list.append(li); + }); + if (!view.entries.length) list.append(h("li", { class: "picker-empty", text: t("picker.empty") })); + if (view.truncated) list.append(h("li", { class: "picker-empty", text: t("picker.truncated", { n: view.entries.length }) })); + if (!roots.childElementCount) { + roots.append(...view.roots.map((r) => h("button", { type: "button", title: r.path, onclick: () => load(r.path), + text: r.id && has(`picker.place.${r.id}`) ? t(`picker.place.${r.id}`) : r.name }))); + } + up.disabled = !view.parent; + mark(); + } + function mark() { + list.querySelectorAll(".picker-item").forEach((li, i) => { + const e = view.entries[i]; + li.setAttribute("aria-selected", picked.has(e.path) ? "true" : "false"); + li.classList.toggle("cursor", i === cursor); + }); + list.setAttribute("aria-activedescendant", cursor >= 0 ? `pick-${cursor}` : ""); + list.querySelector(".cursor")?.scrollIntoView({ block: "nearest" }); + const out = result(); + choose.disabled = !out.length; + status.textContent = picked.size + ? (multiple ? t("picker.n_selected", { n: picked.size }) : [...picked][0]) + : (thisFolder && view ? t("picker.this_folder", { path: view.path }) : t("picker.pick_file")); + } + + const up = h("button", { type: "button", text: "↑", "aria-label": t("picker.up"), title: t("picker.up"), + onclick: () => view && view.parent && load(view.parent) }); + const go = h("button", { type: "button", text: t("picker.go"), onclick: () => load(where.value) }); + const cancel = h("button", { type: "button", text: t("ui.cancel"), onclick: () => finish(null) }); + choose.textContent = multiple ? t("picker.add") : t("picker.select"); + choose.addEventListener("click", () => { const out = result(); if (out.length) finish(out); }); + where.addEventListener("keydown", (e) => { if (e.key === "Enter") { e.preventDefault(); load(where.value); } }); + hidden.addEventListener("change", () => view && load(view.path)); + list.addEventListener("keydown", (ev) => { + if (!view || !view.entries.length) return; + const n = view.entries.length; + const move = { ArrowDown: 1, ArrowUp: -1, PageDown: 10, PageUp: -10 }[ev.key]; + if (move) cursor = Math.max(0, Math.min(n - 1, cursor + move)); + else if (ev.key === "Home") cursor = 0; + else if (ev.key === "End") cursor = n - 1; + else if (ev.key === " ") toggle(entryAt(cursor)); + else if (ev.key === "Enter") { const e = entryAt(cursor); if (e && e.dir) open(e); else if (e) { toggle(e); if (!multiple) open(e); } } + else if (ev.key === "Backspace") { if (view.parent) load(view.parent); } + else return; + ev.preventDefault(); + mark(); + }); + dlg.addEventListener("cancel", (e) => { e.preventDefault(); finish(null); }); + dlg.append( + h("div", { class: "picker-head" }, h("h3", { id: "picker-title", text: title || t(`picker.title.${kind}`) }), + h("button", { type: "button", class: "ghost", text: "×", "aria-label": t("ui.close"), onclick: () => finish(null) })), + h("div", { class: "picker-bar" }, up, where, go), + roots, list, + h("div", { class: "picker-foot" }, + h("label", { class: "compact", for: "picker-hidden" }, hidden, t("picker.hidden")), status, + h("div", { class: "row end" }, cancel, choose))); + document.body.append(dlg); + dlg.showModal(); + load(start || lastDir[kind] || "").then(() => list.focus()); + }); +} + +function icon(dir) { + return h("svg", { viewBox: "0 0 16 16", width: 16, height: 16, class: "picker-icon", "aria-hidden": "true" }, + h("path", { svg: true, d: dir ? "M1.5 3.5h4.5l1.5 1.5h7v8.5h-13z" : "M3.5 1.5h6l3 3v10h-9zM9.5 1.5v3h3" })); +} + +// ── shared field helpers ─────────────────────────────────────────────── +const lastLine = (field) => { const c = splitPaths(field.value); return c[c.length - 1] || ""; }; + +// Multiple fields append (no repeats); single fields take the first path. +function putPaths(field, got, multiple, onChange) { + const current = splitPaths(field.value); + field.value = multiple ? [...new Set([...current, ...got])].join("\n") : got[0]; + field.dispatchEvent(new Event("input", { bubbles: true })); + if (onChange) onChange(field.value); + field.focus(); +} + +// A "Browse…" button: the in-page list of the server's folders. +export function browseButton(field, { kind = "any", multiple = false, onChange } = {}) { + const btn = h("button", { type: "button", class: "browse", text: t(multiple ? "picker.add_button" : "ui.browse"), + tip: t("picker.browse_tip") }); + btn.addEventListener("click", async () => { + const got = await openPicker({ kind, multiple, start: lastLine(field) || lastDir[kind] || "" }); + if (got && got.length) putPaths(field, got, multiple, onChange); + }); + return btn; +} + +// ── "Upload…": system file manager ───────────────────────────────────── +// Mirrors chronos/workbench/files.py (KINDS, BLOCKED_SUFFIXES); the server checks again. +const ACCEPT = { + data: [".jsonl", ".json", ".parquet", ".csv", ".tsv", ".txt"], + checkpoint: [".pth", ".pt", ".safetensors", ".npz", ".gguf", ".json"], // + the .config.json sidecar + config: [".json"], +}; +const WEIGHTS = [".pth", ".pt", ".safetensors", ".npz", ".gguf"]; +const BLOCKED = [".py", ".pyc", ".pyo", ".so", ".dylib", ".dll", ".exe", ".bat", ".cmd", ".com", ".sh", ".command", + ".ps1", ".app", ".pkg", ".dmg", ".jar"]; +const ends = (name, list) => { const low = name.toLowerCase(); return list.some((s) => low.endsWith(s)); }; +const accepts = (kind, name) => !ends(name, BLOCKED) && (!ACCEPT[kind] || ends(name, ACCEPT[kind])); +const baseName = (p) => String(p).split(/[\\/]/).filter(Boolean).pop() || String(p); + +function uploadMode() { + const s = state.session || {}; + return s.access === "local" && s.native_dialog ? "native" : "upload"; +} + +function batchId() { // [A-Za-z0-9_-]{6,48}; works on plain-HTTP LAN pages too + const b = new Uint8Array(9); + crypto.getRandomValues(b); + return `${Date.now().toString(36)}-${Array.from(b, (x) => x.toString(16).padStart(2, "0")).join("")}`; +} +// The server opens the operating system's dialog on its own desktop; the chosen paths are used in place. +async function nativePick(field, { kind, multiple, folder }) { + const got = await post("/api/fs/native", { kind, multiple, folder, start: lastLine(field) || lastDir[kind] || "", + title: t(`picker.title.${folder ? "dir" : kind}`) }); + return got && got.paths && got.paths.length ? got.paths : null; +} + +// One file as a streamed PUT body (XHR, because fetch has no upload progress). +function putFile(file, query, onProgress) { + return new Promise((resolve, reject) => { + const xhr = new XMLHttpRequest(); + xhr.open("PUT", `/api/fs/upload?${new URLSearchParams(query)}`); + xhr.withCredentials = true; + xhr.setRequestHeader("X-Chronos-Request", "1"); + xhr.setRequestHeader("Content-Type", "application/octet-stream"); + xhr.upload.onprogress = (e) => onProgress(e.loaded); + xhr.onload = () => { + let data = null; + try { data = JSON.parse(xhr.responseText); } catch { data = null; } + if (xhr.status >= 200 && xhr.status < 300) resolve(data); + else reject(new ApiError(xhr.status, data && data.detail !== undefined ? data.detail : data)); + }; + xhr.onerror = () => reject(new ApiError(0, { error: "upload_failed" })); + xhr.onabort = () => reject(new ApiError(0, { error: "upload_failed" })); + xhr.send(file); + }); +} + +// Uploads the picked files into one server batch folder; resolves to the server paths a field should use. +async function uploadFiles(files, { kind, folder, note }) { + const all = [...files]; + const list = all.filter((f) => accepts(kind, f.webkitRelativePath || f.name) + && !(f.webkitRelativePath || f.name).split("/").some((p) => p.startsWith("."))); + if (!list.length) throw new ApiError(422, { error: "upload_type" }); + const limit = ((state.session && state.session.max_upload_mb) || 20480) * 2 ** 20; + const big = list.find((f) => f.size > limit); + if (big) throw new ApiError(413, { error: "upload_too_large", message: big.name }); + const total = list.reduce((s, f) => s + f.size, 0); + const batch = batchId(); + let done = 0; + const show = (n) => { note.textContent = t("picker.uploading", { done: bytes(n), total: bytes(total) }); }; + note.hidden = false; + show(0); + const saved = []; + for (const f of list) { + const got = await putFile(f, { kind, batch, name: f.webkitRelativePath || f.name }, (n) => show(done + n)); + done += f.size; + saved.push(got); + } + note.textContent = all.length > list.length + ? t("picker.uploaded_skipped", { n: list.length, skipped: all.length - list.length }) + : t("picker.uploaded", { n: list.length }); + if (folder) { // the uploaded folder itself: / + const top = (list[0].webkitRelativePath || "").split("/")[0]; + return [top ? `${saved[0].root}/${top}` : saved[0].root]; + } + if (kind === "checkpoint") { // weights + sidecar: the field takes the weights file + const weights = saved.filter((s) => ends(s.path, WEIGHTS)); + return (weights.length ? weights : saved).map((s) => s.path); + } + return saved.map((s) => s.path); +} +// Which "Upload" buttons a field gets: folders only for dir, files + folders for datasets. +function uploadKinds(kind) { + if (kind === "dir") return [true]; + if (kind === "data" || kind === "any") return [false, true]; + return [false]; +} + +// "Upload…" button(s) for a path field, plus the progress note they share. +export function uploadButtons(field, { kind = "any", multiple = false, onChange } = {}, note) { + const mode = uploadMode(); + const folders = uploadKinds(kind); + return folders.map((folder) => { + const label = folders.length > 1 ? (folder ? "picker.upload_folder" : "picker.upload_files") : "ui.upload"; + const btn = h("button", { type: "button", class: "upload", text: t(label), + tip: t(mode === "native" ? "picker.native_tip" : "picker.upload_tip") }); + let input = null; + const run = async (files) => { + btn.disabled = true; + try { + const got = await uploadFiles(files, { kind, folder, note }); + if (got.length) putPaths(field, got, multiple, onChange); + } catch (err) { + note.hidden = false; + note.textContent = errorText(err); + reportError(err); + } finally { + btn.disabled = false; + if (input) input.value = ""; + } + }; + btn.addEventListener("click", async () => { + if (uploadMode() === "native") { + btn.disabled = true; + try { + const got = await nativePick(field, { kind, multiple, folder }); + if (got) putPaths(field, got, multiple, onChange); + return; + } catch (err) { + if (!(err instanceof ApiError) || err.code !== "native_unavailable") { reportError(err); return; } + state.session.native_dialog = false; // no dialog on this machine: fall back to uploading + toast(t("picker.native_fallback")); + } finally { + btn.disabled = false; + } + } + if (!input) { + input = h("input", { type: "file", hidden: true, tabindex: "-1", "aria-hidden": "true", + multiple: multiple || folder || kind === "checkpoint" ? true : null, + accept: folder || !ACCEPT[kind] ? null : ACCEPT[kind].join(",") }); + if (folder) input.webkitdirectory = true; + input.addEventListener("change", () => { if (input.files && input.files.length) run(input.files); }); + btn.after(input); + } + input.click(); + }); + return btn; + }); +} + +// Input/textarea + Browse + Upload in one row. ``pick.merge`` fields say when several datasets will be +// merged into one training file; the upload note reports progress and where the files went. +export function pathControl(field, pick) { + const status = h("p", { class: "note upload-note", "aria-live": "polite", hidden: true }); + const row = h("div", { class: "path-row" }, field, browseButton(field, pick), uploadButtons(field, pick, status)); + const box = h("div", { class: "path-box" }, row, status); + if (!pick.merge) return box; + const note = h("p", { class: "note path-note", "aria-live": "polite" }); + const update = () => { + const n = splitPaths(field.value).length; + note.textContent = n > 1 ? t("picker.will_merge", { n }) : t("picker.multi_hint"); + }; + field.addEventListener("input", update); + update(); + box.append(note); + return box; +} diff --git a/chronos/workbench/__init__.py b/chronos/workbench/__init__.py new file mode 100644 index 0000000..9935d47 --- /dev/null +++ b/chronos/workbench/__init__.py @@ -0,0 +1 @@ +"""UI-independent services behind the Chronos Web UI (training sessions, inference runners, data tools).""" diff --git a/chronos/workbench/data.py b/chronos/workbench/data.py new file mode 100644 index 0000000..e2a95d2 --- /dev/null +++ b/chronos/workbench/data.py @@ -0,0 +1,150 @@ +"""Dataset inspection for any task (no UI dependency). + +Streams a file, a directory or a glob (jsonl / json / parquet / csv / tsv / +txt; never a full load) and reports the record schemas, field inventory, +char-length distribution, script mix, duplicate ratio and media modalities. +""" +from __future__ import annotations + +import hashlib +import json +import os +import statistics +from collections import Counter + +import chronos.deps # noqa: F401 + +_CHAT_KEYS = ("conversations", "messages", "dialog", "dialogue") +_PROMPT_KEYS = ("prompt", "question", "query", "input", "instruction") +_ANSWER_KEYS = ("response", "answer", "completion", "output", "target", "chosen") + + +def record_schema(rec) -> str: + """Coarse shape of one record; never task-specific beyond the field names.""" + if isinstance(rec, str): + return "plain_text" + if not isinstance(rec, dict): + return "other" + if any(isinstance(rec.get(k), list) for k in _CHAT_KEYS): + return "chat" + if "chosen" in rec and "rejected" in rec: + return "preference" + if "instruction" in rec and any(k in rec for k in ("output", "response", "answer")): + return "instruction" + if any(k in rec for k in _PROMPT_KEYS) and any(k in rec for k in _ANSWER_KEYS): + return "prompt_response" + nested = [v for v in rec.values() if isinstance(v, dict) + and len(v) >= 2 and all(isinstance(x, str) for x in v.values())] + short = [k for k, v in rec.items() if isinstance(v, str) and len(k) <= 3] + if nested or len(short) >= 2: + return "parallel_pair" + if any(isinstance(rec.get(k), str) for k in ("text", "content", "document", "passage", "article", "story")): + return "plain_text" + return "other" + + +def _extract_text(rec) -> str: + """All text a record carries, in field order (chat turns, pairs, any schema).""" + from chronos.tokenizer.extract import _texts_from_record + parts = list(_texts_from_record(rec, None)) + if not parts and isinstance(rec, dict): + parts = [str(v) for v in rec.values() if isinstance(v, str)] + return " ".join(parts) + + +def _fields(rec, prefix: str = "", out: dict | None = None, depth: int = 0) -> dict: + out = {} if out is None else out + if not isinstance(rec, dict) or depth > 2: + return out + for k, v in rec.items(): + key = f"{prefix}.{k}" if prefix else str(k) + if isinstance(v, dict): + _fields(v, key, out, depth + 1) + else: + out[key] = type(v).__name__ if v is not None else "null" + return out + + +def _scan(path, max_records: int = 50_000): + """Stream up to ``max_records`` records; return summary dict + length list + previews. + + ``path`` is one file, folder or glob, or a list of them (one per line in text): several + datasets are scanned as one, in order, as training would see them after merging. + """ + from chronos.tokenizer.extract import iter_records, script_bucket + from chronos.workbench.files import data_files, parse_paths + paths = parse_paths(path) + if not paths: + return {"error": "no path given"}, [], [] + files = data_files(paths) + if not files: + return {"error": f"no readable data files at: {', '.join(paths)!r} (jsonl, json, parquet, csv, tsv, txt)"}, [], [] + char_lens, samples = [], [] + schemas, scripts, field_hits, field_types = Counter(), Counter(), Counter(), {} + examples: dict[str, str] = {} + seen, dup, n_total, errors = set(), 0, 0, [] + modal = Counter() + try: + from chronos.data.flexible_dataset import detect_modalities + except Exception: # pragma: no cover - torch-free environments + detect_modalities = None + for file in files: + if n_total >= max_records: + break + try: + for rec in iter_records(file): + if n_total >= max_records: + break + n_total += 1 + kind = record_schema(rec) + schemas[kind] += 1 + for key, typ in _fields(rec).items(): + field_hits[key] += 1 + field_types.setdefault(key, typ) + if detect_modalities is not None and isinstance(rec, dict): + for mod in detect_modalities(rec): + modal[mod] += 1 + text = _extract_text(rec) + char_lens.append(len(text)) + if text: + scripts[script_bucket(text[:400])] += 1 + digest = hashlib.md5(json.dumps(rec, sort_keys=True, ensure_ascii=False, default=str) + .encode("utf-8")).hexdigest() + if digest in seen: + dup += 1 + else: + seen.add(digest) + if kind not in examples: + examples[kind] = json.dumps(rec, ensure_ascii=False, default=str)[:300] + if len(samples) < 3: + samples.append(text[:200]) + except Exception as exc: # unreadable file: report it, keep scanning the rest + errors.append(f"{os.path.basename(file)}: {type(exc).__name__}: {exc}"[:200]) + if not char_lens: + return {"error": "no parseable records", "files": len(files), "read_errors": errors}, [], samples + ordered = sorted(char_lens) + summary = { + "files": len(files), + "scanned": n_total, + "dup_ratio": round(dup / max(n_total, 1), 4), + "schemas": dict(schemas.most_common()), + "schema_examples": examples, + "fields": {k: {"type": field_types[k], "present": round(c / n_total, 3)} + for k, c in field_hits.most_common(24)}, + "scripts": {k: round(c / max(1, sum(scripts.values())), 3) for k, c in scripts.most_common()}, + "char_min": ordered[0], + "char_mean": round(statistics.mean(char_lens), 1), + "char_median": int(statistics.median(char_lens)), + "char_p95": int(ordered[int(0.95 * (len(ordered) - 1))]), + "char_max": ordered[-1], + "size_mb": round(sum(os.path.getsize(f) for f in files) / (1024 * 1024), 1), + "note": (f"sampled first {max_records} records" if n_total >= max_records else "full scan"), + } + if modal: + summary["modalities"] = dict(modal) + if errors: + summary["read_errors"] = errors[:8] + return summary, char_lens, samples + + +scan_dataset = _scan diff --git a/chronos/workbench/files.py b/chronos/workbench/files.py new file mode 100644 index 0000000..97f3f8e --- /dev/null +++ b/chronos/workbench/files.py @@ -0,0 +1,226 @@ +"""Server-side file picking and dataset merging for the Web UI (no UI dependency). + +The Web UI runs on the user's machine and every path it sends is a server path, so the picker lists +directories through :func:`list_dir` instead of a browser upload. ``kind`` filters what a field +accepts: ``dir`` (folders only), ``data`` (jsonl / json / parquet / csv / tsv / txt), ``checkpoint`` +(.pth / .safetensors / .npz / .gguf), ``config`` (.json) or ``any``. + +:func:`merge_datasets` streams several dataset files or folders into one JSONL file, the format every +Chronos training loader reads. Records are copied as they are (plain-text lines become +``{"text": ...}``); nothing is loaded whole into memory. + +:func:`upload_path` places files a browser on another machine uploads (LAN access) under +``/uploads///``, keeping a folder's +relative layout and refusing traversal, hidden names and executable or script types. +""" +from __future__ import annotations + +import hashlib +import json +import os +import re +import time + +DATA_SUFFIXES = (".parquet", ".jsonl", ".json", ".csv", ".tsv", ".txt") +KINDS = { + "dir": (), + "data": DATA_SUFFIXES, + "checkpoint": (".pth", ".pt", ".safetensors", ".npz", ".gguf"), + "config": (".json",), + "any": None, +} +MAX_ENTRIES = 2000 + + +def _entry(path: str, name: str, is_dir: bool) -> dict: + try: + st = os.stat(path) + size, mtime = (None if is_dir else st.st_size), st.st_mtime + except OSError: + size, mtime = None, None + return {"name": name, "path": path, "dir": is_dir, "size": size, "mtime": mtime} + + +def list_dir(path: str, kind: str = "any", show_hidden: bool = False) -> dict: + """Folders and matching files directly under ``path`` (a file path lists its parent). + + Returns ``{"path", "parent", "entries", "truncated", "roots"}``; raises ``FileNotFoundError`` + when nothing usable exists there and ``ValueError`` for an unknown ``kind``. + """ + if kind not in KINDS: + raise ValueError("kind") + path = os.path.abspath(os.path.expanduser(path or "~")) + while not os.path.isdir(path): + parent = os.path.dirname(path) + if parent == path: + raise FileNotFoundError(path) + path = parent + suffixes = KINDS[kind] + dirs, files = [], [] + try: + with os.scandir(path) as it: + for e in it: + if not show_hidden and e.name.startswith("."): + continue + try: + is_dir = e.is_dir() + except OSError: + continue + if is_dir: + dirs.append(_entry(e.path, e.name, True)) + elif suffixes is None or (suffixes and e.name.lower().endswith(suffixes)): + files.append(_entry(e.path, e.name, False)) + except PermissionError: + pass + key = lambda x: x["name"].lower() # noqa: E731 + entries = sorted(dirs, key=key) + sorted(files, key=key) + parent = os.path.dirname(path) + return {"path": path, "parent": parent if parent != path else None, + "entries": entries[:MAX_ENTRIES], "truncated": len(entries) > MAX_ENTRIES, "roots": _roots()} + + +def _roots() -> list[dict]: + """Shortcut places; ``id`` lets the page show the name in the user's language.""" + home = os.path.expanduser("~") + out = [{"id": "home", "name": "~", "path": home}] + for name in ("Downloads", "Desktop", "Documents"): + p = os.path.join(home, name) + if os.path.isdir(p): + out.append({"id": name.lower(), "name": name, "path": p}) + if os.path.isdir("/Volumes"): + out.append({"id": "volumes", "name": "/Volumes", "path": "/Volumes"}) + out.append({"id": "root", "name": "/", "path": os.path.abspath(os.sep)}) + return out + + +def parse_paths(text) -> list[str]: + """Paths from a list or from text with one path per line (``~`` expanded, blanks and repeats dropped).""" + items = text if isinstance(text, (list, tuple)) else str(text or "").splitlines() + out = [] + for item in items: + p = os.path.expanduser(str(item).strip().strip('"').strip("'")) + if p and p not in out: + out.append(p) + return out + + +def data_files(paths: list[str]) -> list[str]: + """Readable data files behind files, folders (recursive) or globs, in a stable order.""" + from chronos.tokenizer.extract import expand + out = [] + for p in paths: + for f in ([p] if os.path.isfile(p) else expand([p])): + if f.lower().endswith(DATA_SUFFIXES) and f not in out: + out.append(f) + return out + + +def merge_datasets(paths: list[str], out_dir: str, log=None) -> str: + """One JSONL training file for the given files / folders. + + A single ``.jsonl`` source is used in place. Otherwise every record of every source is streamed + into ``out_dir/merged_.jsonl`` (the hash covers paths, sizes and mtimes, so an unchanged + selection reuses the earlier merge). A ``.sources.json`` manifest records what went in. + Raises ``FileNotFoundError`` when no readable data file is found. + """ + from chronos.tokenizer.extract import iter_records + files = data_files(paths) + if not files: + raise FileNotFoundError("no_files") + if len(files) == 1 and files[0].lower().endswith(".jsonl"): + return files[0] + sig = hashlib.sha256() + for f in files: + st = os.stat(f) + sig.update(f"{os.path.abspath(f)}\0{st.st_size}\0{int(st.st_mtime)}\n".encode()) + os.makedirs(out_dir, exist_ok=True) + dst = os.path.join(out_dir, f"merged_{sig.hexdigest()[:12]}.jsonl") + if os.path.isfile(dst) and os.path.isfile(dst + ".sources.json"): + if log: + log(f"[data] reusing merged dataset {dst}") + return dst + counts, t0 = [], time.monotonic() + tmp = dst + ".part" + with open(tmp, "w", encoding="utf-8") as fh: + for f in files: + n = 0 + for rec in iter_records(f): + fh.write(json.dumps(rec if isinstance(rec, dict) else {"text": str(rec)}, ensure_ascii=False, + default=str) + "\n") + n += 1 + counts.append({"path": os.path.abspath(f), "records": n}) + if log: + log(f"[data] merged {n:,} records from {f}") + if not sum(c["records"] for c in counts): + os.remove(tmp) + raise FileNotFoundError("no_files") + os.replace(tmp, dst) + with open(dst + ".sources.json", "w", encoding="utf-8") as fh: + json.dump({"sources": counts, "records": sum(c["records"] for c in counts), + "seconds": round(time.monotonic() - t0, 1)}, fh, indent=1, ensure_ascii=False) + if log: + log(f"[data] {sum(c['records'] for c in counts):,} records from {len(files)} files → {dst}") + return dst + + +# ── uploads from a browser on another machine ────────────────────────── +UPLOAD_DIRS = {"data": "datasets", "checkpoint": "checkpoints", "config": "configs", "dir": "folders", "any": "files"} +BLOCKED_SUFFIXES = (".py", ".pyc", ".pyo", ".so", ".dylib", ".dll", ".exe", ".bat", ".cmd", ".com", ".sh", ".command", + ".ps1", ".app", ".pkg", ".dmg", ".jar") +_BATCH = re.compile(r"^[A-Za-z0-9_-]{6,48}$") +MAX_DEPTH = 8 + + +def upload_parts(name: str) -> list[str] | None: + """Safe relative path components of an uploaded file name (``folder/sub/x.jsonl``), else ``None``. + + Rejects absolute paths, ``..``, hidden components, control characters and over-long names. + """ + if not isinstance(name, str) or not name or len(name) > 1024 or "\x00" in name: + return None + parts = [p for p in name.replace("\\", "/").split("/") if p != ""] + if not parts or len(parts) > MAX_DEPTH or name.startswith("/"): + return None + for p in parts: + if p in (".", "..") or p.startswith(".") or len(p) > 200 or any(ord(c) < 32 for c in p): + return None + return parts + + +def upload_accepts(kind: str, filename: str) -> bool: + """Whether a field of ``kind`` takes an uploaded file with this name.""" + low = filename.lower() + if low.endswith(BLOCKED_SUFFIXES): + return False + if kind == "data": + return low.endswith(DATA_SUFFIXES) + if kind == "checkpoint": # the ``.config.json`` sidecar travels with the weights + return low.endswith(KINDS["checkpoint"] + (".json",)) + if kind == "config": + return low.endswith(".json") + return kind in ("dir", "any") + + +def upload_path(workdir: str, kind: str, batch: str, name: str) -> str: + """Where an upload lands: ``/uploads///``. + + Raises ``ValueError`` naming the bad part (``kind``, ``batch``, ``name`` or ``type``). + """ + if kind not in UPLOAD_DIRS: + raise ValueError("kind") + if not _BATCH.match(batch or ""): + raise ValueError("batch") + parts = upload_parts(name) + if parts is None: + raise ValueError("name") + if not upload_accepts(kind, parts[-1]): + raise ValueError("type") + root = os.path.abspath(os.path.join(workdir, "uploads", UPLOAD_DIRS[kind], batch)) + dst = os.path.abspath(os.path.join(root, *parts)) + if os.path.commonpath([root, dst]) != root: + raise ValueError("name") + return dst + + +def upload_root(workdir: str, kind: str, batch: str) -> str: + return os.path.abspath(os.path.join(workdir, "uploads", UPLOAD_DIRS[kind], batch)) diff --git a/chronos/workbench/inference.py b/chronos/workbench/inference.py new file mode 100644 index 0000000..b533aae --- /dev/null +++ b/chronos/workbench/inference.py @@ -0,0 +1,1472 @@ +"""Inference runners for the Web UI: Torch / MLX, full-memory and SSD-lazy experts. + +UI-independent (no Gradio, no i18n). ``run_inference_modes`` runs one +request in ``compare`` / ``offload`` / ``fullload`` mode and returns plain +JSON-able rows, outputs and decode diagnostics; the Web UI server, the legacy +Gradio tab and the benchmark tools all call it. +""" +import os +import time + +import chronos.deps # noqa: F401 (auto-bootstrap minimind on sys.path) +from chronos.trainer.device_utils import configure_cpu_threads, cpu_thread_snapshot + + +INFERENCE_BACKEND_CHOICES = ("mlx", "cuda", "xpu", "mps", "cpu") +TORCH_INFERENCE_PRIORITY = ("cuda", "xpu", "mps", "cpu") +INFERENCE_MODE_CHOICES = ("compare", "offload", "fullload") +MISS_POLICY_CHOICES = ("on_demand", "sync_on_demand", "fallback_diagnostic") +RAM_LOAD_RATIO_CHOICES = ( + "0.10", "0.20", "0.25", "0.33", "0.50", "0.67", + "0.75", "0.90", "1.00", "1.10", "1.25", +) +RAM_LOAD_SWEEP_RATIOS = (0.10, 0.20, 0.25, 0.33, 0.50, 0.67, 0.75, 0.90, 1.00, 1.10, 1.25) +INFERENCE_CHART_METRICS = ( + ("response_time_s", "Response time", "s"), + ("tokens_per_sec", "Decode speed", "tokens/s"), + ("rss_delta_gb", "RSS delta", "GB"), + ("setup_rss_delta_gb", "Setup RSS delta", "GB"), + ("prefill_rss_delta_gb", "Prefill RSS delta", "GB"), + ("decode_rss_delta_gb", "Decode RSS delta", "GB"), + ("rss_after_setup_gb", "Setup RSS actual", "GB"), + ("rss_after_prefill_gb", "Prefill RSS actual", "GB"), + ("rss_after_decode_gb", "Decode RSS actual", "GB"), + ("mps_allocated_after_setup_gb", "Setup MPS allocated", "GB"), + ("mps_allocated_after_prefill_gb", "Prefill MPS allocated", "GB"), + ("mps_allocated_after_decode_gb", "Decode MPS allocated", "GB"), + ("mlx_active_after_setup_gb", "Setup MLX active", "GB"), + ("mlx_active_after_prefill_gb", "Prefill MLX active", "GB"), + ("mlx_active_after_decode_gb", "Decode MLX active", "GB"), + ("prefill_time_s", "Prefill time", "s"), + ("decode_time_s", "Decode time", "s"), + ("cache_hit_rate", "Cache hit rate", "ratio"), + ("resident_hit_rate", "Resident hit rate", "ratio"), + ("prediction_hit_rate", "Prediction hit rate", "ratio"), + ("prediction_topk_recall", "Prediction top-k recall", "ratio"), + ("resident_vram_hit_rate", "Resident VRAM hit", "ratio"), + ("ram_promotion_hit_rate", "RAM promotion hit", "ratio"), + ("ssd_cold_miss_rate", "SSD cold miss", "ratio"), + ("cold_miss_predict_hit_rate", "Cold miss predict hit", "ratio"), + ("fallback_weight_rate", "Fallback weight", "ratio"), + ("on_demand_loads", "On-demand loads", "count"), + ("on_demand_load_time_s", "On-demand load time", "s"), + ("average_ponder_steps", "Average recurrent depth", "steps"), + ("early_exit_ratio", "Recurrent early exit", "ratio"), + ("recurrent_contraction", "Recurrent contraction", "ratio"), + ("cluster_recall", "Cluster recall", "ratio"), +) +INFERENCE_SWEEP_CHART_METRICS = ( + ("tokens_per_sec", "Decode speed", "tokens/s"), + ("response_time_s", "Response time", "s"), + ("cache_hit_rate", "Cache hit rate", "ratio"), + ("prediction_hit_rate", "Prediction hit rate", "ratio"), + ("cold_miss_predict_hit_rate", "Cold miss predict hit", "ratio"), + ("fallback_weight_rate", "Fallback weight", "ratio"), + ("rss_delta_gb", "RSS delta", "GB"), + ("prefill_rss_delta_gb", "Prefill RSS delta", "GB"), + ("decode_rss_delta_gb", "Decode RSS delta", "GB"), +) + + +def _best_torch_inference_backend() -> str: + from chronos.backend import BackendDispatcher + + d = BackendDispatcher() + for name in TORCH_INFERENCE_PRIORITY: + if d.info(name).available: + return name + return "cpu" + + +def _available_inference_backend_choices() -> list[str]: + from chronos.backend import available + + detected = set(available()) + return ["auto"] + [name for name in INFERENCE_BACKEND_CHOICES if name in detected] + + +def _default_inference_backend_value() -> str: + return "auto" + + +def _resolve_inference_backend(requested_backend: str, model_path_val: str, sniffed: dict) -> tuple[str, str]: + from chronos.backend import select + + requested = (requested_backend or "auto").strip().lower() or "auto" + backend = select(None if requested == "auto" else requested) + note = "" + + if requested not in {"", "auto"} and backend != requested: + note = ( + f"[backend fallback] Requested {requested}, but it is not available; " + f"using {backend}.\n" + ) + + return backend, note + + +# Lazy-expert helpers live in chronos.lazy_experts (shared with the bundle and +# the server); these names stay for existing callers. +from chronos.lazy_experts import ( # noqa: E402,F401 (re-exported names) + _store_torch_dtype, + active_expert_count as _actual_active_expert_count, + apply_torch_budget as _apply_offload_expert_budget, + expert_budget as _bounded_offload_expert_budget, + mlx_experts_to_placeholders as _replace_mlx_experts_with_placeholders, + torch_experts_to_placeholders as _replace_torch_experts_with_placeholders, + write_mlx_cache_from_reader as _populate_mlx_cache_from_reader, + write_mlx_cache_from_state_dict as _populate_mlx_cache_from_state_dict, +) + + +def _checkpoint_expert_cache_dir(model_path_val: str, model_cfg, tag: str = "e") -> str: + from chronos.lazy_experts import checkpoint_cache_dir + return checkpoint_cache_dir(model_path_val, model_cfg, tag=tag) + + +def _rebuild_mlx_lazy_cache(store, cache_dir: str, reader, weights, clusters=None) -> None: + """Rebuild the MLX page cache at ``cache_dir`` (always rewrites).""" + import shutil + from chronos.lazy_experts import ensure_mlx_cache + if os.path.exists(cache_dir): + shutil.rmtree(cache_dir) + store._cluster_manifest = None + ensure_mlx_cache(store, cache_dir, reader=reader, weights=weights, clusters=clusters) + + +def _memory_snapshot_gb() -> float: + try: + import psutil + + return psutil.Process().memory_info().rss / (1024 ** 3) + except Exception: + return 0.0 + + +def _backend_memory_snapshot(backend: str) -> dict: + try: + from chronos.backend.mac_diagnostics import mps_memory_snapshot, mlx_memory_snapshot + except Exception: + return {} + if backend == "mps": + return mps_memory_snapshot() + if backend == "mlx": + return mlx_memory_snapshot() + return {} + + +def _configure_inference_cpu_threads_if_needed( + backend: str, + requested: int | str | None = "auto", + budget_percent: int | float | str | None = 75, +) -> dict: + """Keep WebUI CPU inference from inheriting a single-thread shell env.""" + if backend != "cpu": + return {} + threads = configure_cpu_threads(requested, budget_percent=budget_percent) + snap = cpu_thread_snapshot() + return { + "cpu_threads": int(threads), + "cpu_thread_env": snap, + } + + +def _memory_delta(after: dict, before: dict, key: str) -> float: + try: + return round(float(after.get(key, 0.0) or 0.0) - float(before.get(key, 0.0) or 0.0), 6) + except Exception: + return 0.0 + + +def _clone_model_cfg(model_cfg): + from chronos.model.config import ChronosConfig + + try: + data = dict(model_cfg.to_dict()) + except Exception: + data = dict(getattr(model_cfg, "__dict__", {})) + data.pop("architectures", None) + data.pop("transformers_version", None) + data.pop("_name_or_path", None) + data.pop("_commit_hash", None) + return ChronosConfig(**data) + + +def _encode_prompt(tokenizer, prompt_val: str, raw_prompt: bool = False) -> tuple[list[int], list[int]]: + if raw_prompt: + enc = tokenizer(prompt_val, add_special_tokens=False, return_attention_mask=True) + else: + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt_val}], + tokenize=False, + add_generation_prompt=True, + ) + enc = tokenizer(rendered, add_special_tokens=False, return_attention_mask=True) + return list(enc["input_ids"]), list(enc.get("attention_mask") or [1] * len(enc["input_ids"])) + + +def _decode_generation_diagnostics(tokenizer, tokens: list[int], stats: dict | None = None) -> dict: + tokens = [int(t) for t in (tokens or [])] + stats = stats or {} + special_ids = { + "eos": getattr(tokenizer, "eos_token_id", None), + "pad": getattr(tokenizer, "pad_token_id", None), + "unk": getattr(tokenizer, "unk_token_id", None), + "bos": getattr(tokenizer, "bos_token_id", None), + } + raw_text = tokenizer.decode(tokens, skip_special_tokens=False) if tokens else "" + clean_text = tokenizer.decode(tokens, skip_special_tokens=True) if tokens else "" + first_token_id = tokens[0] if tokens else None + first_token_kind = "" + for name, token_id in special_ids.items(): + if token_id is not None and first_token_id == int(token_id): + first_token_kind = name + break + special_count = 0 + special_id_values = {int(v) for v in special_ids.values() if v is not None} + for token in tokens: + if int(token) in special_id_values: + special_count += 1 + first_topk = stats.get("first_token_topk", []) + warning = "" + if tokens and not clean_text.strip(): + warning = "Generated tokens decode to an empty clean string because they are all special/filtered tokens." + elif first_token_kind in {"eos", "pad"}: + warning = f"First generated token is {first_token_kind.upper()}, so generation may stop or decode to empty text." + return { + "generated_token_ids": tokens, + "token_count": len(tokens), + "raw_decode": raw_text, + "clean_decode": clean_text, + "first_token_id": first_token_id, + "first_token_kind": first_token_kind, + "special_token_count": special_count, + "special_token_ratio": round(special_count / max(len(tokens), 1), 6), + "all_special_or_filtered": bool(tokens and not clean_text.strip()), + "first_token_topk": first_topk, + "first_token_special_probs": stats.get("first_token_special_probs", {}), + "warning": warning, + } + + +def _stage_warning_from_checkpoint(sniffed: dict, model_path_val: str) -> str: + stage = str((sniffed or {}).get("stage") or "").strip().lower() + if stage == "chronos": + return ( + f"`{model_path_val}` is a pretrain (`stage=chronos`) checkpoint. " + "It is not SFT/chat-aligned yet, so chat prompts can produce repetition or low-quality text." + ) + return "" + + +def _checkpoint_sidecar_metadata(model_path_val: str) -> dict: + if not model_path_val: + return {} + try: + import json + import os + from chronos.model.checkpoint import checkpoint_config_path + + path = checkpoint_config_path(model_path_val) + if not path or not os.path.exists(path): + return {} + with open(path, encoding="utf-8") as f: + meta = json.load(f) + return dict(meta) if isinstance(meta, dict) else {} + except Exception: + return {} + + +def _format_offload_stats(stats: dict) -> str: + if not stats: + return "" + prefill = stats.get("prefill_expert_ids", []) + return ( + "[offload/predict] " + f"cache_hit_rate={stats.get('cache_hit_rate', 0):.2%} " + f"prefill={prefill} " + f"predicted={stats.get('prefetch_submitted_ids', [])} " + f"prefetch_requests={stats.get('total_requests', 0)} " + f"promoted={stats.get('promoted_expert_ids', [])} " + f"vram={stats.get('vram_experts', 0)}/{stats.get('vram_capacity', 0)} " + f"ram={stats.get('ram_experts', 0)}/{stats.get('ram_capacity_dynamic', 0)} " + f"storage={stats.get('storage_format', 'n/a')} " + f"cluster_aware={stats.get('cluster_aware', False)}\n" + ) + + +def _run_torch_inference(backend: str, model_cfg, model_path_val: str, token_ids: list[int], + attention_mask: list[int], max_tok: int, temp: float, + eos_token_id: int | None, do_sample: bool = True, + ram_load_ratio: float | str | None = 1.0, + miss_policy: str = "on_demand", + cpu_threads: int | str | None = "auto", + cpu_budget_percent: int | float | str | None = 100, + speculative: dict | None = None, + paging: str = "auto") -> tuple[list[int], dict]: + import os + import importlib + import gc + import torch + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.model.checkpoint import load_checkpoint_state_dict, load_state_dict_controlled + from chronos.export import is_export_artifact, open_export_reader + inference_engine_mod = importlib.import_module("chronos.runtime.inference_engine") + + cpu_thread_stats = _configure_inference_cpu_threads_if_needed( + backend, + requested=cpu_threads, + budget_percent=cpu_budget_percent, + ) + setup_t0 = time.monotonic() + setup_mem0 = _memory_snapshot_gb() + backend_mem0 = _backend_memory_snapshot(backend) + model = ChronosForCausalLM(model_cfg) + model_init_mem = _memory_snapshot_gb() + weights = None + reader = None + if model_path_val and os.path.exists(model_path_val): + if is_export_artifact(model_path_val): + reader = open_export_reader(model_path_val) + base_only = reader.load_state_dict(include_experts=False) + else: + weights = load_checkpoint_state_dict(model_path_val, map_location="cpu") + base_only = { + k: v for k, v in weights.items() + if ".mlp.experts." not in k + } + load_state_dict_controlled( + model, + base_only, + allow_missing_substrings=(".mlp.experts.",), + ) + _replace_torch_experts_with_placeholders(model, model_cfg) + base_load_mem = _memory_snapshot_gb() + + device_map = {"cuda": "cuda", "xpu": "xpu", "mps": "mps", "cpu": "cpu"}[backend] + from chronos.precision import place_torch_model + model, precision = place_torch_model(model, device_map) + model = model.eval() + device_mem = _memory_snapshot_gb() + input_ids = torch.tensor([token_ids]).to(device_map) + attn = torch.tensor([attention_mask]).to(device_map) + from chronos.lazy_experts import page_layout + layout = page_layout(model_cfg, paging) + # Per-expert layout keeps the per-(layer, expert) shards; cluster layout + # writes one file per training cluster. + page_clusters = None if layout["tag"] == "e" else layout["clusters"] + engine = inference_engine_mod.ChronosInferenceEngine( + model, + model_cfg, + ssd_dir=_checkpoint_expert_cache_dir(model_path_val, model_cfg, tag=layout["tag"]), + ) + expert_budget = _apply_offload_expert_budget(engine, model_cfg, ram_load_ratio) + expert_budget.update({"paging": layout["paging"], "paging_resolved": layout["resolved"], + "paging_reason": layout["reason"], "num_pages": layout["num_pages"]}) + miss_policy = _normalize_miss_policy(miss_policy) + if weights is not None: + engine.setup_from_state_dict(weights, warm_expert_ids=[], clusters=page_clusters) + del weights + weights = None + gc.collect() + elif reader is not None: + engine.setup_from_state_reader(reader, warm_expert_ids=[], clusters=page_clusters) + del reader + gc.collect() + else: + engine.setup(warm_expert_ids=[]) + setup_mem1 = _memory_snapshot_gb() + backend_mem_setup = _backend_memory_snapshot(backend) + setup_stats = { + "setup_time_s": round(time.monotonic() - setup_t0, 3), + "setup_rss_delta_gb": round(setup_mem1 - setup_mem0, 3), + "model_init_rss_delta_gb": round(model_init_mem - setup_mem0, 3), + "base_load_rss_delta_gb": round(base_load_mem - model_init_mem, 3), + "device_move_rss_delta_gb": round(device_mem - base_load_mem, 3), + "expert_setup_rss_delta_gb": round(setup_mem1 - device_mem, 3), + "rss_before_setup_gb": round(setup_mem0, 3), + "rss_after_model_init_gb": round(model_init_mem, 3), + "rss_after_base_load_gb": round(base_load_mem, 3), + "rss_after_device_move_gb": round(device_mem, 3), + "rss_after_setup_gb": round(setup_mem1, 3), + **cpu_thread_stats, + } + for snapshot, suffix in [ + (backend_mem0, "before_setup"), + (backend_mem_setup, "after_setup"), + ]: + for key in ("mps_allocated_gb", "mps_driver_allocated_gb", "mlx_active_gb", "mlx_cache_gb", "mlx_peak_gb"): + if key in snapshot: + prefix = key[:-3] if key.endswith("_gb") else key + setup_stats[f"{prefix}_{suffix}_gb"] = snapshot[key] + for key in ("mps_allocated_gb", "mps_driver_allocated_gb", "mlx_active_gb", "mlx_cache_gb", "mlx_peak_gb"): + if key in backend_mem_setup: + prefix = key[:-3] if key.endswith("_gb") else key + setup_stats[f"{prefix}_setup_delta_gb"] = _memory_delta(backend_mem_setup, backend_mem0, key) + + try: + if speculative: + out = engine.generate_speculative( + input_ids, + max_new_tokens=int(max_tok), + temperature=float(temp), + top_p=0.85, + top_k=50, + do_sample=do_sample, + eos_token_id=eos_token_id, + miss_policy=miss_policy, + **speculative, + ) + else: + out = engine.generate( + input_ids, + attention_mask=attn, + max_new_tokens=int(max_tok), + temperature=float(temp), + eos_token_id=eos_token_id, + do_sample=do_sample, + miss_policy=miss_policy, + ) + stats = dict(getattr(engine, "last_stats", {})) + finally: + engine.teardown() + + prompt_len = input_ids.shape[1] + generated = out[0, prompt_len:].tolist() + stats.update(setup_stats) + stats.update(expert_budget) + + return generated, stats + + +def _next_token_from_logits(logits, input_ids, temperature: float, top_p: float = 0.85, + top_k: int = 50, do_sample: bool = True, + repetition_penalty: float = 1.0): + import torch + import torch.nn.functional as F + + temperature = max(float(temperature), 1e-6) + logits = logits / temperature + if repetition_penalty != 1.0: + for i in range(input_ids.shape[0]): + logits[i, torch.unique(input_ids[i])] /= repetition_penalty + if top_k > 0: + k = min(int(top_k), 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 = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) > top_p + mask[..., 1:] = mask[..., :-1].clone() + mask[..., 0] = 0 + logits[mask.scatter(1, sorted_indices, mask)] = -float("inf") + return ( + torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) + if do_sample else torch.argmax(logits, dim=-1, keepdim=True) + ) + + +def _run_torch_full_dram(backend: str, model_cfg, model_path_val: str, token_ids: list[int], + attention_mask: list[int], max_tok: int, temp: float, + eos_token_id: int | None, do_sample: bool = True, + cpu_threads: int | str | None = "auto", + cpu_budget_percent: int | float | str | None = 100, + speculative: dict | None = None) -> tuple[list[int], dict]: + import os + import time as _time + import gc + import torch + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.model.checkpoint import load_checkpoint_state_dict, load_state_dict_controlled + from chronos.export import is_export_artifact, open_export_reader + + device_map = {"cuda": "cuda", "xpu": "xpu", "mps": "mps", "cpu": "cpu"}[backend] + cpu_thread_stats = _configure_inference_cpu_threads_if_needed( + backend, + requested=cpu_threads, + budget_percent=cpu_budget_percent, + ) + setup_t0 = _time.monotonic() + setup_mem0 = _memory_snapshot_gb() + backend_mem0 = _backend_memory_snapshot(backend) + model = ChronosForCausalLM(model_cfg) + model_init_mem = _memory_snapshot_gb() + if model_path_val and os.path.exists(model_path_val): + if is_export_artifact(model_path_val): + weights = open_export_reader(model_path_val).load_state_dict(include_experts=True) + else: + weights = load_checkpoint_state_dict(model_path_val, map_location="cpu") + load_state_dict_controlled(model, weights) + del weights + gc.collect() + weight_load_mem = _memory_snapshot_gb() + from chronos.precision import place_torch_model + model, precision = place_torch_model(model, device_map) + model = model.eval() + setup_mem1 = _memory_snapshot_gb() + backend_mem_setup = _backend_memory_snapshot(backend) + input_ids = torch.tensor([token_ids], device=device_map) + attn = torch.tensor([attention_mask], device=device_map) + prompt_len = input_ids.shape[1] + past_key_values = None + finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device) + tokens_generated = 0 + prefill_t0 = _time.monotonic() + prefill_mem0 = _memory_snapshot_gb() + backend_mem_prefill0 = _backend_memory_snapshot(backend) + prefill_time_s = 0.0 + prefill_mem1 = prefill_mem0 + backend_mem_prefill1 = backend_mem_prefill0 + prefill_done_t = prefill_t0 + spec_stats = None + if speculative: + from chronos.runtime.speculative import speculative_generate, speculative_supported + full_depth = int(getattr(model_cfg, "max_inference_depth", None) or model_cfg.recurrent_depth) + use_spec = speculative_supported(model_cfg) and int(speculative["draft_depth"]) < full_depth + first = {} + + def _fwd(ids, cache, depth): + out = model(ids, past_key_values=cache, use_cache=True, recurrent_depth=depth) + if "t" not in first: + first["t"] = _time.monotonic() + return out + gen_out, spec = speculative_generate( + model, input_ids, max_new_tokens=int(max_tok), + draft_depth=int(speculative["draft_depth"]) if use_spec else full_depth, + gamma=int(speculative["gamma"]), gamma_max=int(speculative.get("gamma_max", 8)), + adaptive_gamma=bool(speculative.get("adaptive_gamma", True)), do_sample=do_sample, + temperature=max(float(temp), 1e-6), top_k=50, top_p=0.85, eos_token_id=eos_token_id, + forward_fn=_fwd) + input_ids = gen_out + tokens_generated = int(gen_out.shape[1] - prompt_len) + prefill_done_t = first.get("t", prefill_t0) + prefill_time_s = prefill_done_t - prefill_t0 + prefill_mem1 = _memory_snapshot_gb() + backend_mem_prefill1 = _backend_memory_snapshot(backend) + spec_stats = spec.as_dict() | {"enabled": use_spec, "full_depth": full_depth, + "draft_depth": int(speculative["draft_depth"]) if use_spec else full_depth, + "gamma_final": spec.gamma_history[-1] if spec.gamma_history else int(speculative["gamma"])} + + with torch.inference_mode(): + for step in range(0 if speculative else int(max_tok)): + from chronos.model.model_chronos import cache_sequence_length + past_len = cache_sequence_length(past_key_values) if past_key_values else 0 + result = model( + input_ids[:, past_len:], + attention_mask=attn, + past_key_values=past_key_values, + use_cache=True, + ) + outputs = result + if attn is not None: + attn = torch.cat([attn, attn.new_ones(attn.shape[0], 1)], -1) + next_token = _next_token_from_logits( + outputs.logits[:, -1, :], + input_ids, + temperature=float(temp), + do_sample=do_sample, + ) + if eos_token_id is not None: + next_token = torch.where( + finished.unsqueeze(-1), + next_token.new_full((next_token.shape[0], 1), eos_token_id), + next_token, + ) + input_ids = torch.cat([input_ids, next_token], dim=-1) + past_key_values = outputs.past_key_values + tokens_generated += 1 + + if step == 0: + prefill_done_t = _time.monotonic() + prefill_mem1 = _memory_snapshot_gb() + backend_mem_prefill1 = _backend_memory_snapshot(backend) + prefill_time_s = prefill_done_t - prefill_t0 + + if eos_token_id is not None: + finished |= next_token.squeeze(-1).eq(eos_token_id) + if finished.all(): + break + + end_t = _time.monotonic() + end_mem = _memory_snapshot_gb() + backend_mem_end = _backend_memory_snapshot(backend) + if tokens_generated == 0: + prefill_done_t = end_t + prefill_mem1 = end_mem + backend_mem_prefill1 = backend_mem_end + prefill_time_s = end_t - prefill_t0 + elapsed = end_t - prefill_t0 + generated = input_ids[0, prompt_len:].tolist() + stats = { + "tokens_generated": len(generated), + "elapsed_s": round(elapsed, 3), + "tokens_per_sec": round(len(generated) / max(elapsed, 1e-6), 2), + "setup_time_s": round(_time.monotonic() - setup_t0 - elapsed, 3), + "setup_rss_delta_gb": round(setup_mem1 - setup_mem0, 3), + "model_init_rss_delta_gb": round(model_init_mem - setup_mem0, 3), + "base_load_rss_delta_gb": round(weight_load_mem - model_init_mem, 3), + "device_move_rss_delta_gb": round(setup_mem1 - weight_load_mem, 3), + "expert_setup_rss_delta_gb": 0.0, + "rss_before_setup_gb": round(setup_mem0, 3), + "rss_after_model_init_gb": round(model_init_mem, 3), + "rss_after_base_load_gb": round(weight_load_mem, 3), + "rss_after_device_move_gb": round(setup_mem1, 3), + "rss_after_setup_gb": round(setup_mem1, 3), + "prefill_time_s": round(prefill_time_s, 3), + "decode_time_s": round(max(0.0, end_t - prefill_done_t), 3), + "prefill_rss_delta_gb": round(prefill_mem1 - prefill_mem0, 3), + "decode_rss_delta_gb": round(end_mem - prefill_mem1, 3), + "rss_before_prefill_gb": round(prefill_mem0, 3), + "rss_after_prefill_gb": round(prefill_mem1, 3), + "rss_after_decode_gb": round(end_mem, 3), + "storage_format": "full_dram", + "cache_hit_rate": 1.0, + "cache_hits": 0, + "cache_misses": 0, + **cpu_thread_stats, + } + if spec_stats is not None: + stats["speculative"] = spec_stats + for snapshot, suffix in [ + (backend_mem0, "before_setup"), + (backend_mem_setup, "after_setup"), + (backend_mem_prefill0, "before_prefill"), + (backend_mem_prefill1, "after_prefill"), + (backend_mem_end, "after_decode"), + ]: + for key in ("mps_allocated_gb", "mps_driver_allocated_gb", "mlx_active_gb", "mlx_cache_gb", "mlx_peak_gb"): + if key in snapshot: + prefix = key[:-3] if key.endswith("_gb") else key + stats[f"{prefix}_{suffix}_gb"] = snapshot[key] + return generated, stats + + +def _run_mlx_full_dram(model_cfg, model_path_val: str, token_ids: list[int], + max_tok: int, temp: float, do_sample: bool = True, + eos_token_id: int | None = None, + speculative: dict | None = None) -> tuple[list[int], dict]: + import os + import time as _time + import gc + import torch + import mlx.core as mx + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.model.checkpoint import load_checkpoint_state_dict, load_state_dict_controlled + from chronos.export import is_export_artifact, open_export_reader + from chronos.mlx.model import ChronosMLXModel + from chronos.mlx.inference import ChronosMLXInferenceEngine + + setup_t0 = _time.monotonic() + setup_mem0 = _memory_snapshot_gb() + backend_mem0 = _backend_memory_snapshot("mlx") + if model_path_val and is_export_artifact(model_path_val): + weights = open_export_reader(model_path_val).load_state_dict(include_experts=True) + elif model_path_val and os.path.exists(model_path_val): + weights = load_checkpoint_state_dict(model_path_val, map_location="cpu") + else: + weights = {} + + pt_model = ChronosForCausalLM(model_cfg) + model_init_mem = _memory_snapshot_gb() + if weights: + load_state_dict_controlled(pt_model, weights) + del weights + weights = {} + gc.collect() + weight_load_mem = _memory_snapshot_gb() + from chronos.precision import resolve_mlx + precision = resolve_mlx("auto") + mlx_model = ChronosMLXModel.from_chronos_pytorch(pt_model.eval(), model_cfg, dtype=precision.compute) + del pt_model + gc.collect() + setup_mem1 = _memory_snapshot_gb() + backend_mem_setup = _backend_memory_snapshot("mlx") + engine = ChronosMLXInferenceEngine(mlx_model, model_cfg) + engine.store.storage_format = "full_dram" + try: + ids = mx.array([token_ids]) + run_t0 = _time.monotonic() + tokens = list(engine.generate( + ids, + max_new_tokens=int(max_tok), + temperature=float(temp) if do_sample else 0.0, + top_p=0.85, + eos_token_id=eos_token_id, + speculative=bool(speculative), + **(speculative or {}), + )) + run_elapsed = _time.monotonic() - run_t0 + end_mem = _memory_snapshot_gb() + backend_mem_end = _backend_memory_snapshot("mlx") + stats = dict(getattr(engine, "last_stats", {})) + prefill_mem = float(stats.get("rss_after_prefill_gb", setup_mem1) or setup_mem1) + prefill_time_s = float(stats.get("prefill_time_s", 0.0) or 0.0) + stats.update({ + "setup_time_s": round(_time.monotonic() - setup_t0 - run_elapsed, 3), + "setup_rss_delta_gb": round(setup_mem1 - setup_mem0, 3), + "model_init_rss_delta_gb": round(model_init_mem - setup_mem0, 3), + "base_load_rss_delta_gb": round(weight_load_mem - model_init_mem, 3), + "device_move_rss_delta_gb": round(setup_mem1 - weight_load_mem, 3), + "expert_setup_rss_delta_gb": 0.0, + "rss_before_setup_gb": round(setup_mem0, 3), + "rss_after_model_init_gb": round(model_init_mem, 3), + "rss_after_base_load_gb": round(weight_load_mem, 3), + "rss_after_device_move_gb": round(setup_mem1, 3), + "rss_after_setup_gb": round(setup_mem1, 3), + "prefill_rss_delta_gb": round(prefill_mem - setup_mem1, 3), + "decode_rss_delta_gb": round(end_mem - prefill_mem, 3), + "rss_before_prefill_gb": round(setup_mem1, 3), + "rss_after_prefill_gb": round(prefill_mem, 3), + "rss_after_decode_gb": round(end_mem, 3), + "prefill_time_s": round(prefill_time_s, 3), + "decode_time_s": round(max(0.0, run_elapsed - prefill_time_s), 3), + "storage_format": "mlx_full_dram", + "cache_hit_rate": 1.0, + "resident_hit_rate": 1.0, + "cache_hits": 0, + "cache_misses": 0, + "miss_policy": "mlx_full_dram", + }) + for snapshot, suffix in [ + (backend_mem0, "before_setup"), + (backend_mem_setup, "after_setup"), + (backend_mem_end, "after_decode"), + ]: + for key in ("mlx_active_gb", "mlx_cache_gb", "mlx_peak_gb"): + if key in snapshot: + prefix = key[:-3] if key.endswith("_gb") else key + stats[f"{prefix}_{suffix}_gb"] = snapshot[key] + return tokens, stats + finally: + engine.stop() + + +def _run_mlx_lazy_inference(model_cfg, model_path_val: str, token_ids: list[int], + max_tok: int, temp: float, do_sample: bool = True, + ram_load_ratio: float | str | None = 1.0, + quality_safe: bool = True, + eos_token_id: int | None = None, + speculative: dict | None = None, + residency: str = "planned", + paging: str = "auto", + residency_decay: float | None = None, + decode_wait_s: float | None = None, + verify_wait_s: float | None = None) -> tuple[list[int], dict]: + import mlx.core as mx + from chronos.lazy_experts import LazyOptions, build_mlx_lazy_engine, warm_ids + + defaults = LazyOptions() + opts = LazyOptions( + paging=paging, residency=residency, ram_ratio=_normalize_ram_load_ratio(ram_load_ratio), + residency_decay=defaults.residency_decay if residency_decay is None else float(residency_decay), + decode_wait_s=defaults.decode_wait_s if decode_wait_s is None else float(decode_wait_s), + verify_wait_s=defaults.verify_wait_s if verify_wait_s is None else float(verify_wait_s), + ).validated() + setup_mem0 = _memory_snapshot_gb() + backend_mem0 = _backend_memory_snapshot("mlx") + engine, info = build_mlx_lazy_engine(model_path_val, model_cfg, opts) + try: + # Warm only the execution budget: warming every expert would hide the + # lazy path; selected misses load synchronously, page by page. + engine.store.warm_up(warm_ids(model_cfg, info)) + setup_mem1 = _memory_snapshot_gb() + backend_mem_setup = _backend_memory_snapshot("mlx") + tokens = list(engine.generate( + mx.array([token_ids]), + max_new_tokens=int(max_tok), + temperature=float(temp) if do_sample else 0.0, + top_p=0.85, + eos_token_id=eos_token_id, + speculative=bool(speculative), + **opts.generate_kwargs(), + **(speculative or {}), + )) + stats = dict(getattr(engine, "last_stats", {})) + store_stats = engine.store.stats() + stats.update({ + "setup_rss_delta_gb": round(setup_mem1 - setup_mem0, 3), + "rss_before_setup_gb": round(setup_mem0, 3), + "rss_after_setup_gb": round(setup_mem1, 3), + "rss_after_prefill_gb": stats.get("rss_after_prefill_gb", round(setup_mem1, 3)), + "rss_after_decode_gb": stats.get("rss_after_decode_gb", round(setup_mem1, 3)), + **{k: v for k, v in info.items() if k not in ("lazy_options", "precision")}, + "residency_decay": opts.residency_decay, + "decode_wait_s": opts.decode_wait_s, + "verify_wait_s": opts.verify_wait_s, + "storage_format": store_stats.get("storage_format", "safetensors"), + "cluster_aware": store_stats.get("cluster_aware", False), + "num_clusters": store_stats.get("num_clusters", 0), + "cluster_loads": store_stats.get("cluster_loads", 0), + "vram_experts": store_stats.get("hot_experts", 0), + "vram_capacity": store_stats.get("hot_capacity", 0), + "ram_experts": store_stats.get("warm_experts", 0), + "ram_capacity_dynamic": store_stats.get("warm_capacity", int(model_cfg.num_experts)), + "miss_policy": "mlx_lazy_quality", + "quality_safe": quality_safe, + }) + for snapshot, suffix in [(backend_mem0, "before_setup"), (backend_mem_setup, "after_setup")]: + for key in ("mlx_active_gb", "mlx_cache_gb", "mlx_peak_gb"): + if key in snapshot: + stats[f"{key[:-3]}_{suffix}_gb"] = snapshot[key] + return tokens, stats + finally: + engine.stop() + + +def _text_diff_summary(a: str, b: str) -> dict: + import difflib + + ratio = difflib.SequenceMatcher(a=a, b=b).ratio() if (a or b) else 1.0 + prefix = 0 + for ca, cb in zip(a, b): + if ca != cb: + break + prefix += 1 + return { + "same": a == b, + "similarity": round(ratio, 4), + "common_prefix_chars": prefix, + "len_lazy": len(a), + "len_full_dram": len(b), + } + + +def _rows_to_chart_records(rows: list[dict]) -> list[dict]: + has_sweep = any(row.get("ram_load_ratio") is not None for row in rows) + metrics = INFERENCE_SWEEP_CHART_METRICS if has_sweep else INFERENCE_CHART_METRICS + metric_max: dict[str, float] = {} + for _key, label, _unit in metrics: + vals = [] + for row in rows: + value = row.get(_key) + if value is None: + continue + try: + vals.append(abs(float(value))) + except (TypeError, ValueError): + continue + metric_max[label] = max(vals) if vals else 0.0 + + data = [] + for row in rows: + mode = str(row.get("mode", "")) + ratio = row.get("ram_load_ratio") + x_value = f"{float(ratio):.2g}x" if ratio is not None else mode + for key, label, unit in metrics: + value = row.get(key) + if value is None: + continue + value = float(value) + denom = metric_max.get(label, 0.0) + data.append({ + "metric": label, + "mode": mode, + "x": x_value, + "value": value, + "normalized_value": value / denom if denom > 0 else 0.0, + "unit": unit, + }) + return data + + +def _normalize_inference_mode(selected_mode: str | None) -> str: + mode = (selected_mode or "compare").strip().lower() + return mode if mode in INFERENCE_MODE_CHOICES else "compare" + + +def _normalize_ram_load_ratio(value) -> float: + try: + return float(value) + except (TypeError, ValueError): + return 1.0 + + +def _normalize_miss_policy(value) -> str: + policy = (value or "on_demand").strip().lower().replace("-", "_") + aliases = { + "quality_safe": "sync_on_demand", + "quality": "sync_on_demand", + "safe": "sync_on_demand", + "strict": "sync_on_demand", + "sync": "sync_on_demand", + "blocking": "sync_on_demand", + "fallback": "fallback_diagnostic", + "diagnostic": "fallback_diagnostic", + } + policy = aliases.get(policy, policy) + return policy if policy in MISS_POLICY_CHOICES else "on_demand" + + +def _speculative_options(enabled, draft_depth=1, gamma=4) -> dict | None: + """UI controls → options for recurrent-depth self-speculative decoding. + + ``None`` means plain decoding. Speculative decoding needs no extra + weights; greedy output is token-for-token identical to plain greedy and + sampled output keeps the full-depth distribution (rejection sampling). + """ + if not bool(enabled): + return None + try: + depth = max(1, int(float(draft_depth))) + except (TypeError, ValueError): + depth = 1 + try: + g = max(1, min(8, int(float(gamma)))) + except (TypeError, ValueError): + g = 4 + return {"draft_depth": depth, "gamma": g, "gamma_max": 8, "adaptive_gamma": True} + + +def _speculative_row_fields(stats: dict) -> dict: + spec = (stats or {}).get("speculative") or {} + if not spec: + return {} + enabled = spec.get("enabled", int(spec.get("draft_calls", 0) or 0) > 0) + return { + "speculative": bool(enabled), + "spec_draft_depth": spec.get("draft_depth"), + "spec_full_depth": spec.get("full_depth"), + "spec_acceptance_rate": spec.get("acceptance_rate"), + "spec_tokens_per_verify": spec.get("tokens_per_verify"), + "spec_verify_calls": spec.get("verify_calls"), + "spec_draft_calls": spec.get("draft_calls"), + "spec_gamma_final": spec.get("gamma_final"), + } + + +def _format_inference_stats(rows: list[dict], diff: dict | None = None) -> str: + if not rows: + return "_No inference stats yet._" + mode_labels = [str(row.get("mode", f"run_{idx + 1}")) for idx, row in enumerate(rows)] + + def fmt_num(row: dict, key: str, suffix: str = "", precision: int = 3) -> str: + try: + return f"{float(row.get(key, 0.0) or 0.0):.{precision}f}{suffix}" + except (TypeError, ValueError): + return "" + + def fmt_pct(row: dict, key: str, fallback_key: str | None = None) -> str: + value = row.get(key, row.get(fallback_key, 0.0) if fallback_key else 0.0) + try: + return f"{float(value or 0.0):.2%}" + except (TypeError, ValueError): + return "" + + def fmt_int(row: dict, key: str) -> str: + try: + return str(int(row.get(key, 0) or 0)) + except (TypeError, ValueError): + return "" + + def fmt_hits(row: dict) -> str: + return f"{fmt_int(row, 'cache_hits')}/{fmt_int(row, 'cache_misses')}" + + def fmt_bool(row: dict, key: str) -> str: + value = row.get(key) + if value in (None, ""): + return "" + return "True" if bool(value) else "False" + + metric_rows = [ + ("Backend", lambda r: str(r.get("backend", ""))), + ("Policy", lambda r: str(r.get("miss_policy", "full"))), + ("CPU threads", lambda r: fmt_int(r, "cpu_threads")), + ("Tokens", lambda r: fmt_int(r, "tokens")), + ("Tokens/s", lambda r: fmt_num(r, "tokens_per_sec", precision=2)), + ("Self-speculative", lambda r: fmt_bool(r, "speculative")), + ("Draft / full depth", lambda r: (f"{r.get('spec_draft_depth')} / {r.get('spec_full_depth')}" + if r.get("spec_draft_depth") is not None else "")), + ("Draft acceptance", lambda r: fmt_pct(r, "spec_acceptance_rate") if r.get("spec_acceptance_rate") is not None else ""), + ("Tokens per verify", lambda r: fmt_num(r, "spec_tokens_per_verify", precision=2) + if r.get("spec_tokens_per_verify") is not None else ""), + ("Verify / draft forwards", lambda r: (f"{r.get('spec_verify_calls')} / {r.get('spec_draft_calls')}" + if r.get("spec_verify_calls") is not None else "")), + ("Response time", lambda r: fmt_num(r, "response_time_s", "s")), + ("Prefill time", lambda r: fmt_num(r, "prefill_time_s", "s")), + ("Decode time", lambda r: fmt_num(r, "decode_time_s", "s")), + ("RSS delta", lambda r: fmt_num(r, "rss_delta_gb", " GB")), + ("Setup RSS delta", lambda r: fmt_num(r, "setup_rss_delta_gb", " GB")), + ("Prefill RSS delta", lambda r: fmt_num(r, "prefill_rss_delta_gb", " GB")), + ("Decode RSS delta", lambda r: fmt_num(r, "decode_rss_delta_gb", " GB")), + ("Setup RSS actual", lambda r: fmt_num(r, "rss_after_setup_gb", " GB")), + ("Prefill RSS actual", lambda r: fmt_num(r, "rss_after_prefill_gb", " GB")), + ("Decode RSS actual", lambda r: fmt_num(r, "rss_after_decode_gb", " GB")), + ("MPS setup", lambda r: fmt_num(r, "mps_allocated_after_setup_gb", " GB")), + ("MPS prefill", lambda r: fmt_num(r, "mps_allocated_after_prefill_gb", " GB")), + ("MPS decode", lambda r: fmt_num(r, "mps_allocated_after_decode_gb", " GB")), + ("MLX setup", lambda r: fmt_num(r, "mlx_active_after_setup_gb", " GB")), + ("MLX prefill", lambda r: fmt_num(r, "mlx_active_after_prefill_gb", " GB")), + ("MLX decode", lambda r: fmt_num(r, "mlx_active_after_decode_gb", " GB")), + ("Resident hit", lambda r: fmt_pct(r, "resident_hit_rate", "cache_hit_rate")), + ("Predict hit", lambda r: fmt_pct(r, "prediction_hit_rate")), + ("Prediction top-k recall", lambda r: fmt_pct(r, "prediction_topk_recall")), + ("VRAM hit", lambda r: fmt_pct(r, "resident_vram_hit_rate")), + ("RAM promotion hit", lambda r: fmt_pct(r, "ram_promotion_hit_rate")), + ("SSD cold miss", lambda r: fmt_pct(r, "ssd_cold_miss_rate")), + ("Average recurrent depth", lambda r: fmt_num(r, "average_ponder_steps", precision=2)), + ("Early exit", lambda r: fmt_pct(r, "early_exit_ratio")), + ("Cluster recall", lambda r: fmt_pct(r, "cluster_recall")), + ("Cold miss predict hit", lambda r: fmt_pct(r, "cold_miss_predict_hit_rate")), + ("Fallback weight", lambda r: fmt_pct(r, "fallback_weight_rate")), + ("On-demand loads", lambda r: fmt_int(r, "on_demand_loads")), + ("Prefetch drops", lambda r: fmt_int(r, "prefetch_queue_drops")), + ("Prefetch wait", lambda r: fmt_num(r, "prefetch_wait_time_s", "s")), + ("Warm hits", lambda r: fmt_int(r, "warm_hits")), + ("Hot evictions", lambda r: fmt_int(r, "hot_evictions")), + ("Warm evictions", lambda r: fmt_int(r, "warm_evictions")), + ("Async misses", lambda r: fmt_int(r, "async_cold_miss_prefetches")), + ("Sync SSD loads", lambda r: fmt_int(r, "sync_ssd_loads")), + ("Eviction policy", lambda r: str(r.get("eviction_policy") or "")), + ("Expert hits/misses", fmt_hits), + ("All-resident fast path", lambda r: fmt_bool(r, "all_resident_fast_path")), + ("Load budget", lambda r: str(r.get("load_budget", "all"))), + ("VRAM experts", lambda r: str(r.get("vram_experts", ""))), + ("RAM experts", lambda r: str(r.get("ram_experts", ""))), + ("Cluster aware", lambda r: str(r.get("cluster_aware", False))), + ("First token", lambda r: str(r.get("first_token_id", ""))), + ("First token kind", lambda r: str(r.get("first_token_kind", ""))), + ("Special token ratio", lambda r: fmt_pct(r, "special_token_ratio")), + ] + + lines = [ + "| Metric | " + " | ".join(mode_labels) + " |", + "|---|" + "|".join("---:" for _ in mode_labels) + "|", + ] + for metric, getter in metric_rows: + values = [getter(row) for row in rows] + if all(value in {"", "0.000 GB", "0.000s", "0.00%", "0"} for value in values): + continue + lines.append("| " + metric + " | " + " | ".join(values) + " |") + if diff: + lines.append("") + lines.append( + f"Output diff: same=`{diff.get('same')}`, similarity=`{diff.get('similarity')}`, " + f"common_prefix_chars=`{diff.get('common_prefix_chars')}`, " + f"len_lazy=`{diff.get('len_lazy')}`, len_full_dram=`{diff.get('len_full_dram')}`" + ) + warnings = [str(row.get("decode_warning", "")).strip() for row in rows if row.get("decode_warning")] + if warnings: + lines.append("") + lines.append("Decode warnings:") + for warning in dict.fromkeys(warnings): + lines.append(f"- {warning}") + return "\n".join(lines) + + +def _format_inference_notices(raw: dict) -> str: + notices = [] + checkpoint_warning = (raw or {}).get("checkpoint_warning") + if checkpoint_warning: + notices.append(f"Checkpoint notice: {checkpoint_warning}") + for row in (raw or {}).get("rows", []): + warning = row.get("decode_warning") + if warning: + notices.append(f"{row.get('mode', 'run')}: {warning}") + if not notices: + return "" + return "\n\n".join(f"> {notice}" for notice in dict.fromkeys(notices)) + "\n\n" + + +def _row_from_stats(mode: str, backend: str, tokens: list[int], elapsed: float, + rss_delta: float, stats: dict) -> dict: + effective_budget = stats.get("effective_expert_budget") + ideal_budget = stats.get("ideal_active_experts") + max_budget = stats.get("max_allowed_expert_budget") + load_budget = ( + "all" if mode == "full_dram" else + f"{effective_budget}/{max_budget} (ideal {ideal_budget})" + ) + return { + "mode": mode, + "backend": backend, + "cpu_threads": stats.get("cpu_threads", 0), + "cpu_thread_env": stats.get("cpu_thread_env", {}), + "tokens": len(tokens), + "tokens_per_sec": round(len(tokens) / max(elapsed, 1e-6), 2), + "response_time_s": round(elapsed, 3), + "rss_delta_gb": round(rss_delta, 3), + "setup_time_s": stats.get("setup_time_s", 0.0), + "setup_rss_delta_gb": stats.get("setup_rss_delta_gb", 0.0), + "model_init_rss_delta_gb": stats.get("model_init_rss_delta_gb", 0.0), + "base_load_rss_delta_gb": stats.get("base_load_rss_delta_gb", 0.0), + "device_move_rss_delta_gb": stats.get("device_move_rss_delta_gb", 0.0), + "expert_setup_rss_delta_gb": stats.get("expert_setup_rss_delta_gb", 0.0), + "rss_before_setup_gb": stats.get("rss_before_setup_gb", 0.0), + "rss_after_model_init_gb": stats.get("rss_after_model_init_gb", 0.0), + "rss_after_base_load_gb": stats.get("rss_after_base_load_gb", 0.0), + "rss_after_device_move_gb": stats.get("rss_after_device_move_gb", 0.0), + "rss_after_setup_gb": stats.get("rss_after_setup_gb", 0.0), + "prefill_time_s": stats.get("prefill_time_s", 0.0), + "decode_time_s": stats.get("decode_time_s", 0.0), + "prefill_rss_delta_gb": stats.get("prefill_rss_delta_gb", 0.0), + "decode_rss_delta_gb": stats.get("decode_rss_delta_gb", 0.0), + "rss_before_prefill_gb": stats.get("rss_before_prefill_gb", 0.0), + "rss_after_prefill_gb": stats.get("rss_after_prefill_gb", 0.0), + "rss_after_decode_gb": stats.get("rss_after_decode_gb", 0.0), + "mps_allocated_before_setup_gb": stats.get("mps_allocated_before_setup_gb", 0.0), + "mps_allocated_after_setup_gb": stats.get("mps_allocated_after_setup_gb", 0.0), + "mps_allocated_before_prefill_gb": stats.get("mps_allocated_before_prefill_gb", 0.0), + "mps_allocated_after_prefill_gb": stats.get("mps_allocated_after_prefill_gb", 0.0), + "mps_allocated_after_decode_gb": stats.get("mps_allocated_after_decode_gb", 0.0), + "mps_driver_allocated_after_setup_gb": stats.get("mps_driver_allocated_after_setup_gb", 0.0), + "mps_driver_allocated_after_prefill_gb": stats.get("mps_driver_allocated_after_prefill_gb", 0.0), + "mps_driver_allocated_after_decode_gb": stats.get("mps_driver_allocated_after_decode_gb", 0.0), + "mlx_active_before_setup_gb": stats.get("mlx_active_before_setup_gb", 0.0), + "mlx_active_after_setup_gb": stats.get("mlx_active_after_setup_gb", 0.0), + "mlx_active_before_prefill_gb": stats.get("mlx_active_before_prefill_gb", 0.0), + "mlx_active_after_prefill_gb": stats.get("mlx_active_after_prefill_gb", 0.0), + "mlx_active_after_decode_gb": stats.get("mlx_active_after_decode_gb", 0.0), + "mlx_cache_after_setup_gb": stats.get("mlx_cache_after_setup_gb", 0.0), + "mlx_cache_after_prefill_gb": stats.get("mlx_cache_after_prefill_gb", 0.0), + "mlx_cache_after_decode_gb": stats.get("mlx_cache_after_decode_gb", 0.0), + "cache_hit_rate": stats.get("cache_hit_rate", 1.0 if mode == "full_dram" else 0), + "resident_hit_rate": stats.get("resident_hit_rate", stats.get("cache_hit_rate", 1.0 if mode == "full_dram" else 0)), + "prediction_hit_rate": stats.get("prediction_hit_rate", 0.0), + "prediction_topk_recall": stats.get("prediction_topk_recall", 0.0), + "resident_vram_hit_rate": stats.get("resident_vram_hit_rate", 0.0), + "ram_promotion_hit_rate": stats.get("ram_promotion_hit_rate", 0.0), + "ssd_cold_miss_rate": stats.get("ssd_cold_miss_rate", 0.0), + "average_ponder_steps": stats.get("average_ponder_steps", 1.0), + "early_exit_ratio": stats.get("early_exit_ratio", 0.0), + "recurrent_contraction": stats.get("recurrent_contraction", 0.0), + "cluster_recall": stats.get("cluster_recall", 0.0), + "prediction_hits": stats.get("prediction_hits", 0), + "prediction_total": stats.get("prediction_total", 0), + "cold_miss_predict_hit_rate": stats.get("cold_miss_predict_hit_rate", 0.0), + "cold_prediction_hits": stats.get("cold_prediction_hits", 0), + "cold_prediction_total": stats.get("cold_prediction_total", 0), + "all_resident_fast_path": stats.get("all_resident_fast_path", mode == "full_dram"), + "on_demand_loads": stats.get("on_demand_loads", 0), + "on_demand_load_time_s": stats.get("on_demand_load_time_s", 0.0), + "async_cold_miss_prefetches": stats.get("async_cold_miss_prefetches", 0), + "prefetch_queue_drops": stats.get("prefetch_queue_drops", 0), + "prefetch_wait_time_s": stats.get("prefetch_wait_time_s", 0.0), + "sync_ssd_loads": stats.get("sync_ssd_loads", 0), + "eviction_policy": stats.get("eviction_policy"), + "resident_vram_hits": stats.get("resident_vram_hits", 0), + "resident_ram_hits": stats.get("resident_ram_hits", 0), + "warm_hits": stats.get("warm_hits", 0), + "warm_evictions": stats.get("warm_evictions", 0), + "hot_evictions": stats.get("hot_evictions", 0), + "quality_safe": stats.get("quality_safe", mode != "full_dram"), + "miss_policy": stats.get("miss_policy", "full_dram" if mode == "full_dram" else "on_demand"), + "cache_hits": stats.get("cache_hits", 0), + "cache_misses": stats.get("cache_misses", 0), + "fallback_weight_rate": stats.get("fallback_weight_rate", 0.0), + "fallback_weight_mass": stats.get("fallback_weight_mass", 0.0), + "hit_weight_mass": stats.get("hit_weight_mass", 0.0), + "total_route_weight_mass": stats.get("total_route_weight_mass", 0.0), + "route_layer_stats": stats.get("route_layer_stats", []), + "load_budget": load_budget, + "ideal_active_experts": stats.get("ideal_active_experts"), + "requested_ram_load_ratio": stats.get("requested_ram_load_ratio"), + "requested_expert_budget": stats.get("requested_expert_budget"), + "max_allowed_expert_budget": stats.get("max_allowed_expert_budget"), + "effective_expert_budget": stats.get("effective_expert_budget"), + "effective_vram_expert_budget": stats.get("effective_vram_expert_budget"), + "effective_ram_expert_budget": stats.get("effective_ram_expert_budget"), + "min_vram_expert_budget": stats.get("min_vram_expert_budget"), + "budget_clamped": stats.get("budget_clamped"), + "prefill_experts": stats.get("prefill_expert_ids", []), + "prefill_ram_experts": stats.get("prefill_ram_expert_ids", []), + "activated_experts": stats.get("activated_expert_ids", []), + "predicted_current_experts": stats.get("predicted_current_expert_ids", []), + "predicted_experts": stats.get("prefetch_submitted_ids", []), + "promoted_experts": stats.get("promoted_expert_ids", []), + "on_demand_experts": stats.get("on_demand_expert_ids", []), + "fallback_experts": stats.get("fallback_expert_ids", []), + "vram_experts": ( + "all" if mode == "full_dram" + else f"{stats.get('vram_experts', 0)}/{stats.get('vram_capacity', 0)}" + ), + "ram_experts": ( + "all" if mode == "full_dram" + else f"{stats.get('ram_experts', 0)}/{stats.get('ram_capacity_dynamic', 0)}" + ), + "cluster_aware": False if mode == "full_dram" else stats.get("cluster_aware", False), + } + + +def _attach_decode_diagnostics(row: dict, diagnostics: dict) -> dict: + row["generated_token_ids"] = diagnostics.get("generated_token_ids", []) + row["raw_decode"] = diagnostics.get("raw_decode", "") + row["clean_decode"] = diagnostics.get("clean_decode", "") + row["first_token_id"] = diagnostics.get("first_token_id") + row["first_token_kind"] = diagnostics.get("first_token_kind", "") + row["special_token_count"] = diagnostics.get("special_token_count", 0) + row["special_token_ratio"] = diagnostics.get("special_token_ratio", 0.0) + row["all_special_or_filtered"] = diagnostics.get("all_special_or_filtered", False) + row["first_token_topk"] = diagnostics.get("first_token_topk", []) + row["first_token_special_probs"] = diagnostics.get("first_token_special_probs", {}) + row["decode_warning"] = diagnostics.get("warning", "") + return row + + +def _build_model_cfg(cfg, model_path_val): + from chronos.model.config import ChronosConfig + from chronos.model.checkpoint import config_dict_for_checkpoint, sniff_checkpoint_config + from chronos.export import config_dict_from_export, is_export_artifact + + cfg = cfg or {} + cfg_sources = [] + try: + if model_path_val and is_export_artifact(model_path_val): + ckpt_cfg = config_dict_from_export(model_path_val) + cfg_sources = [model_path_val] + else: + ckpt_cfg, cfg_sources = config_dict_for_checkpoint( + model_path_val if model_path_val else None, + require_unsniffable=False, + ) + except Exception: + ckpt_cfg, cfg_sources = {}, [] + sniffed = ckpt_cfg or (sniff_checkpoint_config(model_path_val) if model_path_val else {}) + sidecar_meta = _checkpoint_sidecar_metadata(model_path_val) + if sidecar_meta: + sniffed = dict(sniffed) + for meta_key in ("stage", "checkpoint_format", "checkpoint_config_version"): + if sidecar_meta.get(meta_key) not in (None, ""): + sniffed[meta_key] = sidecar_meta[meta_key] + sniff_note = "" + if sniffed: + summary = ", ".join( + f"{k}={v}" for k, v in sniffed.items() + if k in ("vocab_size", "hidden_size", "num_hidden_layers", + "num_experts", "moe_intermediate_size") + ) + source_text = "; ".join(cfg_sources) if cfg_sources else model_path_val + sniff_note = f"[checkpoint config: {summary} | source={source_text}]\n" + + from chronos.config_schema import model_config_kwargs + # Checkpoint metadata defines the trained graph. UI values fill missing + # fields for legacy checkpoints, including all new recurrent/media fields. + model_cfg_kwargs = model_config_kwargs({**cfg, **sniffed}) + model_cfg_kwargs.setdefault("hidden_size", 768) + if "num_attention_heads_total_dim" in sniffed: + h = model_cfg_kwargs["hidden_size"] + n_heads = model_cfg_kwargs.get("num_attention_heads", cfg.get("num_attention_heads", 8)) + head_dim = h // n_heads + if head_dim > 0: + correct_n_heads = sniffed["num_attention_heads_total_dim"] // head_dim + correct_n_kv = sniffed["num_kv_heads_total_dim"] // head_dim + if correct_n_heads > 0: + model_cfg_kwargs["num_attention_heads"] = correct_n_heads + if correct_n_kv > 0: + model_cfg_kwargs["num_key_value_heads"] = correct_n_kv + + return ChronosConfig(**model_cfg_kwargs), sniffed, sniff_note + + +def _inference_tokenizer_dir(cfg, model_path_val: str) -> str: + """The tokenizer that belongs to the checkpoint, not a fixed default. + + Order: config ``tokenizer_path``; the checkpoint's bundled ``tokenizer/`` + or the path in its sidecar; the minimind tokenizer only as the last resort + (older checkpoints that never recorded one). + """ + from chronos.model.checkpoint import resolve_checkpoint_tokenizer + explicit = str((cfg or {}).get("tokenizer_path") or "").strip() or None + path = model_path_val if model_path_val and os.path.exists(model_path_val) else None + return resolve_checkpoint_tokenizer(path, explicit, fallback=True) + + +def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, + prompt_val, max_tok, temp, raw_prompt_val, + ram_load_ratio=1.0, sweep_ram_load_ratios=False, + miss_policy="on_demand", speculative=False, draft_depth=1, + gamma=4, residency_planner=True, paging="auto", residency_decay=None, + decode_wait_s=None, verify_wait_s=None) -> dict: + from transformers import AutoTokenizer + from chronos.backend.mac_diagnostics import mac_backend_diagnostics + + if not str(prompt_val or "").strip(): + return {"error": "Please enter a prompt."} + + backend_diag = mac_backend_diagnostics(configure_threads=True) + model_cfg, sniffed, sniff_note = _build_model_cfg(cfg, model_path_val) + backend, backend_note = _resolve_inference_backend(selected_backend, model_path_val, sniffed) + cpu_threads_cfg = (cfg or {}).get("cpu_threads", "auto") + cpu_budget_cfg = (cfg or {}).get("cpu_budget_percent", 100) + tokenizer = AutoTokenizer.from_pretrained(_inference_tokenizer_dir(cfg, model_path_val)) + spec_opts = _speculative_options(speculative, draft_depth, gamma) + token_ids, attn_mask = _encode_prompt(tokenizer, prompt_val, bool(raw_prompt_val)) + + mode = _normalize_inference_mode(selected_mode) + miss_policy = _normalize_miss_policy(miss_policy) + rows: list[dict] = [] + outputs = {"lazy": "", "full_dram": "", "single": ""} + raw = { + "mode": mode, + "backend": backend, + "backend_note": backend_note.strip(), + "checkpoint_note": sniff_note.strip(), + "checkpoint_stage": str(sniffed.get("stage", "") or ""), + "checkpoint_warning": _stage_warning_from_checkpoint(sniffed, model_path_val), + "backend_diagnostics": backend_diag, + "expert_budget_policy": { + "source": "checkpoint num_experts_per_tok", + "cap": "floor(active_experts * 1.25)", + "miss_policy": miss_policy, + **_bounded_offload_expert_budget(model_cfg, ram_load_ratio), + }, + "rows": rows, + "speculative": spec_opts or {"enabled": False}, + } + + ratio_values = [_normalize_ram_load_ratio(ram_load_ratio)] + if mode in {"compare", "offload"} and bool(sweep_ram_load_ratios): + ratio_values = list(RAM_LOAD_SWEEP_RATIOS) + ratio_values = list(dict.fromkeys(float(r) for r in ratio_values)) + + representative_lazy_text = "" + if mode in {"compare", "offload"}: + for ratio in ratio_values: + run_model_cfg = _clone_model_cfg(model_cfg) + mem0 = _memory_snapshot_gb() + t0 = time.monotonic() + if backend == "mlx": + lazy_tokens, lazy_stats = _run_mlx_lazy_inference( + run_model_cfg, + model_path_val, + token_ids, + int(max_tok), + float(temp), + do_sample=(mode != "compare"), + ram_load_ratio=ratio, + quality_safe=(mode == "compare" and not bool(sweep_ram_load_ratios)), + eos_token_id=tokenizer.eos_token_id, + speculative=spec_opts, + residency="planned" if residency_planner else "window", + paging=paging, + residency_decay=residency_decay, + decode_wait_s=decode_wait_s, + verify_wait_s=verify_wait_s, + ) + else: + lazy_tokens, lazy_stats = _run_torch_inference( + backend, + run_model_cfg, + model_path_val, + token_ids, + attn_mask, + int(max_tok), + float(temp), + tokenizer.eos_token_id, + do_sample=(mode != "compare"), + ram_load_ratio=ratio, + miss_policy=miss_policy, + cpu_threads=cpu_threads_cfg, + cpu_budget_percent=cpu_budget_cfg, + speculative=spec_opts, + paging=paging, + ) + lazy_elapsed = time.monotonic() - t0 + lazy_mem = _memory_snapshot_gb() + lazy_diag = _decode_generation_diagnostics(tokenizer, lazy_tokens, lazy_stats) + lazy_text = lazy_diag["clean_decode"] + row = _row_from_stats( + "lazy_offload", backend, lazy_tokens, lazy_elapsed, + lazy_mem - mem0, lazy_stats, + ) + row.update(_speculative_row_fields(lazy_stats)) + _attach_decode_diagnostics(row, lazy_diag) + if len(ratio_values) > 1: + row["mode"] = f"lazy_offload_{ratio:g}x" + row["ram_load_ratio"] = round(float(ratio), 3) + rows.append(row) + if abs(float(ratio) - _normalize_ram_load_ratio(ram_load_ratio)) < 1e-9: + representative_lazy_text = lazy_text + if not representative_lazy_text and rows: + representative_lazy_text = lazy_text + outputs["lazy"] = representative_lazy_text + + full_text = "" + if mode in {"compare", "fullload"}: + full_model_cfg = _clone_model_cfg(model_cfg) + mem1 = _memory_snapshot_gb() + t1 = time.monotonic() + if backend == "mlx": + full_tokens, full_stats = _run_mlx_full_dram( + full_model_cfg, + model_path_val, + token_ids, + int(max_tok), + float(temp), + do_sample=(mode != "compare"), + eos_token_id=tokenizer.eos_token_id, + speculative=spec_opts, + ) + else: + full_tokens, full_stats = _run_torch_full_dram( + backend, + full_model_cfg, + model_path_val, + token_ids, + attn_mask, + int(max_tok), + float(temp), + tokenizer.eos_token_id, + do_sample=(mode != "compare"), + cpu_threads=cpu_threads_cfg, + cpu_budget_percent=cpu_budget_cfg, + speculative=spec_opts, + ) + full_elapsed = time.monotonic() - t1 + full_mem = _memory_snapshot_gb() + full_diag = _decode_generation_diagnostics(tokenizer, full_tokens, full_stats) + full_text = full_diag["clean_decode"] + row = _row_from_stats( + "full_dram", backend, full_tokens, full_elapsed, + full_mem - mem1, full_stats, + ) + row.update(_speculative_row_fields(full_stats)) + _attach_decode_diagnostics(row, full_diag) + rows.append(row) + outputs["full_dram"] = full_text + + diff = _text_diff_summary(representative_lazy_text, full_text) if mode == "compare" else None + if diff: + raw["output_diff"] = diff + raw["rows"] = rows + raw["outputs"] = outputs + raw["decode_diagnostics"] = { + row["mode"]: { + "generated_token_ids": row.get("generated_token_ids", []), + "raw_decode": row.get("raw_decode", ""), + "clean_decode": row.get("clean_decode", ""), + "first_token_id": row.get("first_token_id"), + "first_token_kind": row.get("first_token_kind", ""), + "special_token_ratio": row.get("special_token_ratio", 0.0), + "all_special_or_filtered": row.get("all_special_or_filtered", False), + "warning": row.get("decode_warning", ""), + } + for row in rows + } + raw["chart"] = _rows_to_chart_records(rows) + if mode == "offload": + outputs["single"] = representative_lazy_text + elif mode == "fullload": + outputs["single"] = full_text + return raw + + +def generate_api(cfg, selected_backend="auto", selected_mode="compare", model_path_val="", + prompt_val="", max_tok=128, temp=0.85, raw_prompt_val=False, + ram_load_ratio=1.0, sweep_ram_load_ratios=False, + miss_policy="on_demand", speculative=False, draft_depth=1, gamma=4, + residency_planner=True, paging="auto", residency_decay=None, + decode_wait_s=None, verify_wait_s=None) -> dict: + try: + return _run_inference_modes( + cfg, selected_backend, selected_mode, model_path_val, prompt_val, + max_tok, temp, raw_prompt_val, ram_load_ratio, sweep_ram_load_ratios, + miss_policy, speculative, draft_depth, gamma, residency_planner, + paging=paging, residency_decay=residency_decay, + decode_wait_s=decode_wait_s, verify_wait_s=verify_wait_s, + ) + except Exception as e: + import traceback + + return {"error": str(e), "traceback": traceback.format_exc()} + + +run_inference_modes = _run_inference_modes diff --git a/chronos/workbench/lazy_probe.py b/chronos/workbench/lazy_probe.py new file mode 100644 index 0000000..8842b36 --- /dev/null +++ b/chronos/workbench/lazy_probe.py @@ -0,0 +1,567 @@ +"""Lazy-expert validation probe: efficiency, hit rate and fidelity of a checkpoint. + +At a validation point the trainer has just written a checkpoint. This probe +loads that checkpoint the way inference would (``chronos.lazy_experts``): +trunk in memory, routed experts paged from an SSD cache, and decodes a few +fixed prompts under every requested setting. Each setting is compared with a +full-memory reference decode of the same prompts on the same backend. + +Measured per variant (labelled *measured*, conditions in ``conditions``): + +* efficiency – decode tokens/s, prefill/decode seconds, on-demand loads and + their load time, prefetch wait, cluster loads, peak RSS; +* hit rate – resident (selection) hit rate, lookahead prediction hit rate, + cold-miss prediction hit rate; +* fidelity – greedy token agreement with the reference (exact match, + matching prefix, position-wise agreement). + +Settings: backend (mlx / torch), paging (auto / cluster / expert), residency +(planned / window), RAM ratio, Torch miss policy, residency decay, prefetch +waits and self-speculative decoding. A sweep is the Cartesian product of the +lists given, deduplicated and capped at ``max_variants`` (the cap is reported, +never silent). +""" +from __future__ import annotations + +import itertools +import os +import time +from dataclasses import asdict, dataclass, field + +# Prompts used when the checkpoint has no tokenizer: fixed id sequences, so +# results from different validation points are comparable. +_SYNTHETIC_PROMPTS = ((5, 17, 33, 2, 61, 9, 40, 12), (7, 3, 29, 44, 18, 6, 51, 23, 11, 2)) +DEFAULT_PROMPTS = ( + "Translate to Chinese: The weather is nice today.", + "Explain what a mixture-of-experts layer does in one sentence.", +) + + +@dataclass +class ProbeSweep: + """What to measure; every list is one sweep axis.""" + + backends: list[str] = field(default_factory=lambda: ["auto"]) + paging: list[str] = field(default_factory=lambda: ["auto", "expert"]) # deduped after resolving + residency: list[str] = field(default_factory=lambda: ["planned"]) + ram_ratios: list[float] = field(default_factory=lambda: [0.5, 1.0]) + miss_policies: list[str] = field(default_factory=lambda: ["sync_on_demand"]) # torch only + residency_decay: list[float] = field(default_factory=lambda: [0.7]) # mlx planned only + decode_wait_s: list[float] = field(default_factory=lambda: [0.012]) # mlx only + verify_wait_s: list[float] = field(default_factory=lambda: [0.02]) # mlx speculative only + speculative: list[bool] = field(default_factory=lambda: [False]) + draft_depth: int = 1 + gamma: int = 4 + prompts: list[str] = field(default_factory=lambda: list(DEFAULT_PROMPTS)) + max_new_tokens: int = 32 + max_variants: int = 12 + dtype: str = "float32" + cache_root: str | None = None + + @classmethod + def from_dict(cls, data: dict | None) -> "ProbeSweep": + data = dict(data or {}) + known = {f for f in cls.__dataclass_fields__} + out = cls(**{k: v for k, v in data.items() if k in known}) + for name in ("backends", "paging", "residency", "ram_ratios", "miss_policies", "residency_decay", + "decode_wait_s", "verify_wait_s", "speculative", "prompts"): + value = getattr(out, name) + if not isinstance(value, (list, tuple)): + setattr(out, name, [value]) + return out.validated() + + def validated(self) -> "ProbeSweep": + from chronos.expert_clusters import PAGING_MODES + from chronos.lazy_experts import RESIDENCY_MODES + checks = (("paging", self.paging, PAGING_MODES), ("residency", self.residency, RESIDENCY_MODES), + ("miss_policies", self.miss_policies, ("on_demand", "sync_on_demand", "fallback_diagnostic")), + ("backends", self.backends, ("auto", "mlx", "cpu", "mps", "cuda", "xpu"))) + for name, values, allowed in checks: + bad = [v for v in values if v not in allowed] + if bad or not values: + raise ValueError(f"{name}: {bad or 'empty'} not in {allowed}") + if any(float(r) <= 0 for r in self.ram_ratios) or not self.ram_ratios: + raise ValueError("ram_ratios must be > 0") + if any(not 0.0 <= float(d) < 1.0 for d in self.residency_decay): + raise ValueError("residency_decay must be in [0, 1)") + if any(float(w) < 0 for w in (*self.decode_wait_s, *self.verify_wait_s)): + raise ValueError("waits must be >= 0") + if not 1 <= int(self.max_new_tokens) <= 512 or not 1 <= int(self.max_variants) <= 256: + raise ValueError("max_new_tokens in [1, 512], max_variants in [1, 256]") + if int(self.draft_depth) < 1 or not 1 <= int(self.gamma) <= 8: + raise ValueError("draft_depth >= 1, gamma in [1, 8]") + return self + + def as_dict(self) -> dict: + return asdict(self) + + +def resolve_backend(requested: str) -> str: + from chronos.backend import select + return select(None if requested == "auto" else requested) + + +def expand_variants(sweep: ProbeSweep, resolve_paging=None) -> tuple[list[dict], int]: + """Deduplicated settings; axes that do not apply to a backend collapse. + + ``resolve_paging(p)`` maps a paging request to the layout it produces for + the checkpoint (``auto`` → ``cluster`` or ``expert``); settings that + resolve to the same layout are measured once. Returns + ``(variants, total_before_cap)``.""" + resolve_paging = resolve_paging or (lambda p: p) + seen, out = set(), [] + for backend in dict.fromkeys(resolve_backend(b) for b in sweep.backends): + mlx = backend == "mlx" + for paging, residency, ratio, policy, decay, dwait, vwait, spec in itertools.product( + sweep.paging, sweep.residency if mlx else ["-"], sweep.ram_ratios, + ["-"] if mlx else sweep.miss_policies, sweep.residency_decay, sweep.decode_wait_s, + sweep.verify_wait_s, sweep.speculative): + v = {"backend": backend, "paging": resolve_paging(paging), "ram_ratio": round(float(ratio), 4), + "speculative": bool(spec)} + if mlx: + v.update(residency=residency, decode_wait_s=float(dwait)) + if residency == "planned": + v["residency_decay"] = float(decay) + if spec: + v["verify_wait_s"] = float(vwait) + else: + v["miss_policy"] = policy + key = tuple(sorted(v.items())) + if key not in seen: + seen.add(key) + out.append(v) + return out[: int(sweep.max_variants)], len(out) + + +def encode_prompts(model_path: str, sweep: ProbeSweep, tokenizer_dir: str | None = None, + vocab_size: int | None = None) -> tuple[list[list[int]], str]: + """Prompt ids from the checkpoint's tokenizer, else the fixed synthetic ids. + + A tokenizer whose ids do not fit the model's vocabulary (a fallback + tokenizer for a checkpoint that recorded none) is not used.""" + try: + from transformers import AutoTokenizer + from chronos.workbench.inference import _encode_prompt, _inference_tokenizer_dir + tok = AutoTokenizer.from_pretrained(tokenizer_dir or _inference_tokenizer_dir({}, model_path)) + ids = [_encode_prompt(tok, p, False)[0] for p in sweep.prompts if str(p).strip()] + if ids and (not vocab_size or max(max(x) for x in ids if x) < int(vocab_size)): + return ids, "tokenizer" + except Exception: + pass + limit = int(vocab_size or 64) + return [[t % limit for t in p] for p in _SYNTHETIC_PROMPTS], "synthetic" + + +def agreement(tokens: list[int], ref: list[int]) -> dict: + """Greedy fidelity of one decode against the full-memory reference.""" + prefix = 0 + for a, b in zip(tokens, ref): + if a != b: + break + prefix += 1 + n = max(len(ref), 1) + same = sum(1 for a, b in zip(tokens, ref) if a == b) + return {"exact": tokens == ref, "prefix": prefix, "positionwise": same / n, "ref_len": len(ref)} + + +def _decode(backend: str, cfg, path: str, ids: list[int], sweep: ProbeSweep, variant: dict | None): + """One greedy decode; ``variant`` None = full-memory reference.""" + from chronos.workbench.inference import ( + _clone_model_cfg, _run_mlx_full_dram, _run_mlx_lazy_inference, _run_torch_full_dram, + _run_torch_inference, _speculative_options, + ) + spec = _speculative_options(bool(variant and variant["speculative"]), sweep.draft_depth, sweep.gamma) + run_cfg, n = _clone_model_cfg(cfg), int(sweep.max_new_tokens) + if backend == "mlx": + if variant is None: + return _run_mlx_full_dram(run_cfg, path, ids, n, 0.0, do_sample=False, speculative=None) + return _run_mlx_lazy_inference( + run_cfg, path, ids, n, 0.0, do_sample=False, ram_load_ratio=variant["ram_ratio"], quality_safe=True, + speculative=spec, residency=variant["residency"], paging=variant["paging"], + residency_decay=variant.get("residency_decay"), decode_wait_s=variant.get("decode_wait_s"), + verify_wait_s=variant.get("verify_wait_s")) + mask = [1] * len(ids) + if variant is None: + return _run_torch_full_dram(backend, run_cfg, path, ids, mask, n, 1.0, None, do_sample=False) + return _run_torch_inference(backend, run_cfg, path, ids, mask, n, 1.0, None, do_sample=False, + ram_load_ratio=variant["ram_ratio"], miss_policy=variant["miss_policy"], + speculative=spec, paging=variant["paging"]) + + +# Stats summed over prompts, and rates recomputed from their counts. +_SUMS = ("on_demand_loads", "on_demand_load_time_s", "prefetch_wait_time_s", "sync_ssd_loads", "cluster_loads", + "cache_hits", "cache_misses", "prediction_hits", "prediction_total", "cold_prediction_hits", + "cold_prediction_total", "prefill_time_s", "decode_time_s") + + +def _merge(stats: list[dict], tokens: int, elapsed: float) -> dict: + out = {k: round(sum(float(s.get(k, 0) or 0) for s in stats), 4) for k in _SUMS} + sel = out["cache_hits"] + out["cache_misses"] + out["resident_hit_rate"] = round(out["cache_hits"] / sel, 4) if sel else ( + round(sum(float(s.get("resident_hit_rate", 0) or 0) for s in stats) / max(len(stats), 1), 4)) + if out["prediction_total"]: + out["prediction_hit_rate"] = round(out["prediction_hits"] / out["prediction_total"], 4) + elif any("prediction_total" in s for s in stats): + out["prediction_hit_rate"] = None # this path made no lookahead predictions + else: + out["prediction_hit_rate"] = round(sum(float(s.get("prediction_hit_rate", 0) or 0) for s in stats) + / max(len(stats), 1), 4) + if out["cold_prediction_total"]: + out["cold_miss_predict_hit_rate"] = round(out["cold_prediction_hits"] / out["cold_prediction_total"], 4) + out["tokens"] = tokens + out["tokens_per_s"] = round(tokens / elapsed, 2) if elapsed > 0 else 0.0 + out["decode_tokens_per_s"] = round(tokens / out["decode_time_s"], 2) if out["decode_time_s"] > 0 else None + last = stats[-1] if stats else {} + for key in ("paging_resolved", "num_pages", "effective_expert_budget", "ideal_active_experts", + "eviction_policy", "num_clusters", "storage_format"): + if key in last: + out[key] = last[key] + peaks = [float(s.get("rss_after_decode_gb") or 0) for s in stats] + out["peak_rss_gb"] = round(max(peaks), 3) if peaks else None + spec = [s.get("speculative") for s in stats if isinstance(s.get("speculative"), dict)] + if spec: + acc = [float(x.get("acceptance_rate", 0) or 0) for x in spec] + out["spec_acceptance_rate"] = round(sum(acc) / len(acc), 4) + return out + + +def variant_label(v: dict) -> str: + parts = [v["backend"], v["paging"], f"r{v['ram_ratio']:g}"] + if v["backend"] == "mlx": + parts.append(v["residency"]) + if "residency_decay" in v: + parts.append(f"d{v['residency_decay']:g}") + parts.append(f"w{v['decode_wait_s'] * 1000:g}ms") + else: + parts.append(v["miss_policy"]) + if v["speculative"]: + parts.append("spec") + return "/".join(parts) + + +def run_probe(model_path: str, sweep: ProbeSweep | dict | None = None, *, tokenizer_dir: str | None = None, + log=None) -> dict: + """Measure every variant of ``sweep`` on ``model_path``; see the module docstring. + + Lazy variants run first and the full-memory references last, so process + RSS after a lazy decode is not inflated by a fully loaded model.""" + from chronos.lazy_experts import page_layout + from chronos.workbench.inference import _build_model_cfg + log = log or (lambda _m: None) + sweep = sweep if isinstance(sweep, ProbeSweep) else ProbeSweep.from_dict(sweep) + cfg, _sniffed, _note = _build_model_cfg({"dtype": sweep.dtype}, model_path) + if not getattr(cfg, "use_moe", True) or int(getattr(cfg, "num_experts", 0) or 0) < 2: + return {"skipped": "dense model: no routed experts to page", "checkpoint": model_path} + variants, total = expand_variants(sweep, lambda p: page_layout(cfg, p)["resolved"]) + prompts, prompt_source = encode_prompts(model_path, sweep, tokenizer_dir, int(cfg.vocab_size)) + t_start = time.monotonic() + results = [] + for v in variants: + label = variant_label(v) + log(f"[lazy-probe] {label}") + stats, toks, t0 = [], [], time.monotonic() + try: + for ids in prompts: + out, st = _decode(v["backend"], cfg, model_path, ids, sweep, v) + toks.append(list(out)) + stats.append(st) + except Exception as exc: # one bad setting must not hide the others + results.append({"label": label, "settings": v, "error": f"{type(exc).__name__}: {exc}"[:400]}) + continue + results.append({"label": label, "settings": v, "tokens": toks, + "measured": _merge(stats, sum(len(t) for t in toks), time.monotonic() - t0)}) + references = {} + for backend in dict.fromkeys(r["settings"]["backend"] for r in results if "tokens" in r): + try: + refs, t0 = [], time.monotonic() + ref_stats = [] + for ids in prompts: + out, st = _decode(backend, cfg, model_path, ids, sweep, None) + refs.append(list(out)) + ref_stats.append(st) + references[backend] = {"tokens": refs, + "measured": _merge(ref_stats, sum(len(t) for t in refs), time.monotonic() - t0)} + except Exception as exc: + references[backend] = {"error": f"{type(exc).__name__}: {exc}"[:400]} + for r in results: + ref = references.get(r["settings"]["backend"], {}) + if "tokens" not in r or "tokens" not in ref: + continue + per = [agreement(a, b) for a, b in zip(r.pop("tokens"), ref["tokens"])] + r["fidelity"] = {"exact_match_rate": round(sum(p["exact"] for p in per) / len(per), 4), + "positionwise_agreement": round(sum(p["positionwise"] for p in per) / len(per), 4), + "min_matching_prefix": min(p["prefix"] for p in per), + "reference_tokens": sum(p["ref_len"] for p in per)} + full_tps = ref["measured"].get("tokens_per_s") or 0 + if full_tps: + r["measured"]["speed_vs_full"] = round(r["measured"]["tokens_per_s"] / full_tps, 4) + for ref in references.values(): + ref.pop("tokens", None) + layout = page_layout(cfg) + return { + "checkpoint": os.path.abspath(model_path), + "trained_with_clusters": layout["trained_with_clusters"], "auto_paging": layout["resolved"], + "num_experts": int(cfg.num_experts), "top_k": int(cfg.num_experts_per_tok), + "variants": results, "references": references, + "truncated": total > len(variants), "variants_requested": total, + "elapsed_s": round(time.monotonic() - t_start, 2), + "conditions": { + "label": "measured", "decode": "greedy", "max_new_tokens": int(sweep.max_new_tokens), + "prompts": len(prompts), "prompt_source": prompt_source, "dtype": sweep.dtype, + "reference": "full-memory decode of the same prompts on the same backend", + "rss_note": "process RSS after each decode; variants share one process in the listed order", + }, + "sweep": sweep.as_dict(), + } + + +def _sidecar(path: str) -> str: + return os.path.splitext(path)[0] + ".config.json" + + +def newest_checkpoint(candidates) -> str | None: + """The most recently written existing ``.pth`` among ``candidates``.""" + found = [c for c in candidates if c and os.path.isfile(c)] + return max(found, key=os.path.getmtime) if found else None + + +def _settled(path: str, wait_s: float = 1.0, tries: int = 10) -> bool: + """True once the checkpoint's size and mtime are unchanged across ``wait_s``. + + ``save_state_dict_with_config`` writes the ``.pth`` first and its + ``.config.json`` sidecar second, so a sidecar that is older than the + weights (or still missing) is waited for as well; checkpoints saved + without any sidecar are accepted once stable.""" + def sig(): + st = os.stat(path) + side = _sidecar(path) + fresh = os.path.isfile(side) and os.path.getmtime(side) >= st.st_mtime + return st.st_size, st.st_mtime_ns, fresh + try: + last = sig() + for attempt in range(tries): + time.sleep(wait_s) + now = sig() + if now == last and (now[2] or attempt >= 2): + return True + last = now + except OSError: + return False + return False + + +class ProbeScheduler: + """Runs one probe at a time in a low-priority child process. + + ``request`` copies the checkpoint (and its config sidecar) to + ``/lazy_probe/step_/`` first, so training can overwrite the + original while the probe reads the snapshot; the snapshot and the probe's + expert cache are deleted when the probe ends. Results are appended to + ``/lazy_probe.jsonl``; ``poll`` returns the finished summaries + (one line each). A request while a probe runs is recorded as + skipped (validation never waits for the probe). + """ + + def __init__(self, run_dir: str, sweep: dict | None = None, *, tokenizer_dir: str | None = None, + nice: int = 10, log=None): + self.run_dir = os.path.abspath(run_dir) + self.sweep = ProbeSweep.from_dict(sweep).as_dict() # validate up front + self.tokenizer_dir = tokenizer_dir or None + self.nice, self.log = int(nice), (log or (lambda _m: None)) + self.proc = None + self.current: dict | None = None + self.results_path = os.path.join(self.run_dir, "lazy_probe.jsonl") + os.makedirs(self.run_dir, exist_ok=True) + + def busy(self) -> bool: + return self.proc is not None and self.proc.poll() is None + + def skip(self, *, phase: str, step: int, reason: str) -> dict: + """Record a validation that will not be probed (e.g. superseded by a newer one).""" + return self._record({"phase": phase, "step": step, "skipped": reason}) + + def request(self, checkpoint: str | None, *, phase: str, step: int) -> dict: + import json + import shutil + import subprocess + import sys + if not checkpoint or not os.path.isfile(checkpoint): + return self._record({"phase": phase, "step": step, "skipped": "no checkpoint on disk yet"}) + if self.busy(): + return self._record({"phase": phase, "step": step, "skipped": "previous probe still running"}) + if not _settled(checkpoint): + return self._record({"phase": phase, "step": step, "skipped": "checkpoint still being written"}) + work = os.path.join(self.run_dir, "lazy_probe", f"{phase}_step_{int(step):08d}") + os.makedirs(work, exist_ok=True) + snap = os.path.join(work, os.path.basename(checkpoint)) + shutil.copy2(checkpoint, snap) + if os.path.isfile(_sidecar(checkpoint)): + shutil.copy2(_sidecar(checkpoint), _sidecar(snap)) + sweep_file = os.path.join(work, "sweep.json") + with open(sweep_file, "w", encoding="utf-8") as fh: + json.dump(self.sweep, fh) + argv = [sys.executable, "-m", "chronos.workbench.lazy_probe", "--checkpoint", snap, + "--sweep-file", sweep_file, "--out", os.path.join(work, "result.json")] + if self.tokenizer_dir: + argv += ["--tokenizer-dir", self.tokenizer_dir] + preexec = (lambda: os.nice(self.nice)) if hasattr(os, "nice") and self.nice > 0 else None + self.log_file = open(os.path.join(work, "probe.log"), "w", encoding="utf-8") + self.proc = subprocess.Popen(argv, stdout=self.log_file, stderr=subprocess.STDOUT, + preexec_fn=preexec, env={**os.environ, "PYTHONUNBUFFERED": "1"}) + self.current = {"phase": phase, "step": int(step), "checkpoint": os.path.abspath(checkpoint), + "work": work, "t0": time.monotonic()} + self.log(f"[lazy-probe] started for {phase} step {step} ({os.path.basename(checkpoint)})") + return {"phase": phase, "step": step, "started": True} + + def poll(self) -> list[dict]: + """Entries of probes that finished since the last call (at most one).""" + import json + import shutil + if self.proc is None or self.proc.poll() is None: + return [] + cur, code = self.current or {}, self.proc.returncode + self.proc, self.current = None, None + self.log_file.close() + result, out = None, os.path.join(cur.get("work", ""), "result.json") + try: + with open(out, encoding="utf-8") as fh: + result = json.load(fh) + except (OSError, ValueError): + pass + entry = {"phase": cur.get("phase"), "step": cur.get("step"), "checkpoint": cur.get("checkpoint"), + "wall_s": round(time.monotonic() - cur.get("t0", time.monotonic()), 1)} + if result is None: + try: + with open(os.path.join(cur["work"], "probe.log"), encoding="utf-8") as fh: + tail = fh.read()[-600:] + except (OSError, KeyError): + tail = "" + entry["error"] = f"probe exited {code}: {tail.strip()}"[:800] + else: + entry["result"] = result + shutil.rmtree(cur.get("work", ""), ignore_errors=True) # our snapshot + cache only + self._record(entry) + return [entry] + + def wait(self, timeout_s: float = 600.0) -> list[dict]: + t0 = time.monotonic() + while self.busy() and time.monotonic() - t0 < timeout_s: + time.sleep(0.5) + if self.busy(): + self.stop() + return self.poll() + + def stop(self) -> None: + if self.busy(): + self.proc.terminate() + try: + self.proc.wait(timeout=10) + except Exception: + self.proc.kill() + + def _record(self, entry: dict) -> dict: + import json + from datetime import datetime + with open(self.results_path, "a", encoding="utf-8") as fh: + fh.write(json.dumps({"time": datetime.now().isoformat(timespec="seconds"), **entry}, + ensure_ascii=False, default=str) + "\n") + return entry + + +def watch(run_dir: str, sweep: dict | None = None, *, out_dir: str | None = None, tokenizer_dir: str | None = None, + poll_s: float = 5.0, once: bool = False, log=print) -> int: + """Probe any trainer's run directory from outside the training process. + + Tails ``/metrics.jsonl``; after each ``validation`` event it + probes the newest ``.pth`` under ``run_dir`` written since the previous + probe (``chronos longrun`` exports ``*.best.pth`` at an improving + validation; the Web UI trainer saves before validating). Results go to + ``/lazy_probe.jsonl`` (default: the run directory). The run is + only read. ``once``: probe the newest checkpoint now and exit. + """ + import glob + import json + run_dir = os.path.abspath(run_dir) + sched = ProbeScheduler(out_dir or run_dir, sweep, tokenizer_dir=tokenizer_dir, log=log) + + def newest(): + pths = [p for p in glob.glob(os.path.join(run_dir, "**", "*.pth"), recursive=True) + if os.sep + "lazy_probe" + os.sep not in p] + return newest_checkpoint(pths) + + if once: + sched.request(newest(), phase="manual", step=0) + for entry in sched.wait(timeout_s=24 * 3600): + log(json.dumps({k: v for k, v in entry.items() if k != "result"}, ensure_ascii=False)) + return 0 + metrics, offset, probed = os.path.join(run_dir, "metrics.jsonl"), 0, 0.0 + if os.path.exists(metrics): + offset = os.path.getsize(metrics) # only validations from now on + try: + while True: + sched.poll() + if os.path.exists(metrics) and os.path.getsize(metrics) > offset: + with open(metrics, encoding="utf-8") as fh: + fh.seek(offset) + chunk = fh.read() + cut = chunk.rfind("\n") + 1 # keep a partial last line for later + offset += len(chunk[:cut].encode("utf-8")) + for line in chunk[:cut].splitlines(): + try: + row = json.loads(line) + except ValueError: + continue + if row.get("event") != "validation": + continue + ckpt = newest() + if ckpt and os.path.getmtime(ckpt) > probed: + probed = os.path.getmtime(ckpt) + sched.request(ckpt, phase=str(row.get("phase", "")), step=int(row.get("step", 0) or 0)) + time.sleep(poll_s) + except KeyboardInterrupt: + sched.stop() + return 0 + + +def main(argv=None) -> int: + import argparse + import json + import shutil + import tempfile + p = argparse.ArgumentParser(prog="chronos lazy-probe", description=__doc__.split("\n\n")[0]) + src = p.add_mutually_exclusive_group(required=True) + src.add_argument("--checkpoint", help="probe this checkpoint once") + src.add_argument("--watch", metavar="RUN_DIR", help="probe after every validation of a run directory") + p.add_argument("--out-dir", default=None, help="--watch: where lazy_probe.jsonl goes (default: the run dir)") + p.add_argument("--sweep", default=None, help="JSON object with ProbeSweep fields") + p.add_argument("--sweep-file", default=None) + p.add_argument("--tokenizer-dir", default=None) + p.add_argument("--out", default=None, help="write the result JSON here (default: stdout)") + p.add_argument("--keep-cache", action="store_true", help="keep the scratch expert cache") + a = p.parse_args(argv) + sweep = json.loads(a.sweep) if a.sweep else {} + if a.sweep_file: + with open(a.sweep_file, encoding="utf-8") as fh: + sweep = json.load(fh) + if a.watch: + return watch(a.watch, sweep, out_dir=a.out_dir, tokenizer_dir=a.tokenizer_dir) + scratch = tempfile.mkdtemp(prefix="chronos_probe_cache_") + os.environ["CHRONOS_EXPERT_CACHE_ROOT"] = scratch + try: + result = run_probe(a.checkpoint, sweep, tokenizer_dir=a.tokenizer_dir, + log=lambda m: print(m, flush=True)) + finally: + if not a.keep_cache: + shutil.rmtree(scratch, ignore_errors=True) + text = json.dumps(result, ensure_ascii=False, indent=2, default=str) + if a.out: + with open(a.out + ".tmp", "w", encoding="utf-8") as fh: + fh.write(text) + os.replace(a.out + ".tmp", a.out) + else: + print(text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/workbench/serve.py b/chronos/workbench/serve.py new file mode 100644 index 0000000..be8712a --- /dev/null +++ b/chronos/workbench/serve.py @@ -0,0 +1,116 @@ +"""``chronos serve`` as a child process (argument list, no shell). + +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 os +import subprocess +import sys +import threading +from collections import deque + +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 = "", + paging: str = "auto", mla_cache_rope: bool = False) -> 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 paging not in ("auto", "cluster", "expert"): + raise ValueError("paging") + if not lazy_experts: + argv.append("--no-lazy-experts") + elif paging != "auto": + argv += ["--paging", paging] + 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 mla_cache_rope: + argv.append("--mla-cache-rope") + 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) diff --git a/chronos/workbench/tokenizer.py b/chronos/workbench/tokenizer.py new file mode 100644 index 0000000..c4dffc9 --- /dev/null +++ b/chronos/workbench/tokenizer.py @@ -0,0 +1,114 @@ +"""Chronos Tokenizer Builder jobs and report helpers (no UI dependency).""" +from __future__ import annotations + +import json +import os +import threading +import time +from collections import deque + +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) diff --git a/chronos/workbench/train_worker.py b/chronos/workbench/train_worker.py new file mode 100644 index 0000000..4a9f7fd --- /dev/null +++ b/chronos/workbench/train_worker.py @@ -0,0 +1,306 @@ +"""Training job worker: ``python -m chronos.workbench.train_worker ``. + +Runs one or more ``TrainSession`` stages in sequence and writes a run +directory that ``chronos.monitor.snapshot.RunMonitor`` (and the Web UI live +monitor) can read: + +* ``run_spec.json`` – what was asked for (config, stages, backend); +* ``status.json`` – heartbeat: pid, state, current phase, step, error; +* ``metrics.jsonl`` – one row per logged step, plus ``validation`` / + ``phase_start`` / ``phase_end`` / ``final`` events; +* ``train.log`` – the session log; +* ``samples.jsonl`` – live generation samples (prompt → output); +* ``lazy_probe.jsonl`` – after each validation, the lazy-expert probe + (``chronos.workbench.lazy_probe``): efficiency, hit rate and fidelity of + the just-saved checkpoint under every requested paging / residency / + RAM-ratio / miss-policy / wait / speculative setting. + +The job file is written by the Web UI server; it is JSON with +``{"run_dir", "config", "stages": [{"mode", "data_path", "data_paths", "init_weight", +"teacher_path", "max_steps", "val_ratio", "sample_prompt"}], "backend", +"lazy_probe": {"enabled": true, "sweep": {...ProbeSweep fields}}}``. +SIGTERM (the UI stop button) stops the current stage and skips the rest. +""" +from __future__ import annotations + +import json +import os +import signal +import sys +import time +from datetime import datetime + +STAGE_MODES = ("pretrain", "sft", "dpo", "orpo", "grpo", "distill") + +# TrainSession reports three stage-specific terms; name them for the monitor. +TERM_NAMES = { + "pretrain": ("primary_loss", "balance_loss", "temporal_loss"), + "sft": ("primary_loss", "lookahead_loss", "anchor_loss"), + "dpo": ("primary_loss", "anchor_loss", "lookahead_loss"), + "orpo": ("primary_loss", "or_loss", "anchor_loss"), + "grpo": ("primary_loss", "kl_loss", "reward_mean"), + "distill": ("primary_loss", "ce_loss", "anchor_loss"), +} + + +def _now() -> str: + return datetime.now().isoformat(timespec="seconds") + + +def _write_json(path: str, payload: dict) -> None: + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as fh: + json.dump(payload, fh, ensure_ascii=False, indent=2, default=str) + os.replace(tmp, path) + + +class RunWriter: + """Appends monitor-format rows and keeps the heartbeat current.""" + + def __init__(self, run_dir: str): + self.run_dir = run_dir + os.makedirs(run_dir, exist_ok=True) + self.metrics = open(os.path.join(run_dir, "metrics.jsonl"), "a", encoding="utf-8") + self.log = open(os.path.join(run_dir, "train.log"), "a", encoding="utf-8") + self.samples_path = os.path.join(run_dir, "samples.jsonl") + self.status = {"pid": os.getpid(), "state": "starting", "started": _now()} + self.t0 = time.monotonic() + + def row(self, payload: dict) -> None: + self.metrics.write(json.dumps({"time": _now(), **payload}, ensure_ascii=False, default=float) + "\n") + self.metrics.flush() + + def lines(self, text: str) -> None: + if text: + self.log.write(text.rstrip("\n") + "\n") + self.log.flush() + + def sample(self, payload: dict) -> None: + with open(self.samples_path, "a", encoding="utf-8") as fh: + fh.write(json.dumps({"time": _now(), **payload}, ensure_ascii=False) + "\n") + + def heartbeat(self, **fields) -> None: + self.status.update(fields, updated=_now(), run_wall_s=round(time.monotonic() - self.t0, 1)) + _write_json(os.path.join(self.run_dir, "status.json"), self.status) + + def close(self) -> None: + self.metrics.close() + self.log.close() + + +def metric_row(mode: str, metric: dict, total_steps: int) -> dict: + """Map one ``TrainSession`` metric dict onto monitor keys.""" + if metric.get("val"): + return {"event": "validation", "phase": mode, "step": int(metric.get("step", 0)), + "valid_ce": float(metric.get("val_loss", metric.get("total", 0.0)))} + names = TERM_NAMES.get(mode, TERM_NAMES["pretrain"]) + row = {"phase": mode, "step": int(metric.get("step", 0)), "total": int(total_steps or 0), + "loss": float(metric.get("total", 0.0))} + for name, key in zip(names, ("primary", "secondary", "tertiary")): + if metric.get(key) is not None: + row[name] = float(metric[key]) + if metric.get("tokens_per_s"): + row["tokens_per_s"] = float(metric["tokens_per_s"]) + if metric.get("tps"): + row["steps_per_s"] = float(metric["tps"]) + row["step_s"] = 1.0 / float(metric["tps"]) + for key in ("grad_norm", "lr", "lookahead_union_recall", "router_adjacent_jaccard", "recurrent_loss"): + value = metric.get(key) + if isinstance(value, (int, float)): + row[key] = float(value) + return row + + +PROBE_GRACE_S = 3.0 + + +def _fresh_checkpoint(candidates, newer_than: float) -> str | None: + from chronos.workbench.lazy_probe import newest_checkpoint + fresh = [] + for path in candidates: + try: + if os.path.getmtime(path) > newer_than: + fresh.append(path) + except OSError: + continue + return newest_checkpoint(fresh) + + +def probe_event(entry: dict) -> dict: + """A compact ``lazy_probe`` event row for metrics.jsonl (the full result is in lazy_probe.jsonl).""" + row = {"event": "lazy_probe", "phase": entry.get("phase"), "step": entry.get("step"), + "wall_s": entry.get("wall_s")} + if "error" in entry: + row["error"] = str(entry["error"])[:200] + result = entry.get("result") or {} + if result.get("skipped"): + row["skipped"] = str(result["skipped"])[:200] + failed = [v for v in result.get("variants") or [] if "error" in v] + if failed: + row["probe_errors"] = len(failed) + variants = [v for v in result.get("variants") or [] if "fidelity" in v] + if variants: + row["probe_variants"] = len(variants) + row["probe_min_exact_match"] = min(v["fidelity"]["exact_match_rate"] for v in variants) + hits = [v["measured"]["resident_hit_rate"] for v in variants + if v.get("measured", {}).get("resident_hit_rate") is not None] + if hits: + row["probe_mean_resident_hit_rate"] = round(sum(hits) / len(hits), 4) + return row + + +def stage_checkpoints(run_cfg: dict, mode: str) -> list[str]: + """Checkpoint files this stage writes (``__moe[.best].pth`` in ``save_dir``).""" + import glob + from chronos.workbench.training import STAGE_CHECKPOINT_PREFIX + prefix = STAGE_CHECKPOINT_PREFIX.get(mode, mode) + return glob.glob(os.path.join(run_cfg.get("save_dir", "./out"), f"{prefix}_*_moe*.pth")) + + +def stage_data_path(stage: dict, writer: RunWriter, cfg: dict) -> str: + """The one JSONL file a stage trains on. + + ``data_paths`` (several files / folders, any supported format) are merged into + ``/data/merged_.jsonl``; a single ``.jsonl`` file is used in place. An empty + selection keeps the old behaviour (``data_path`` or the config's, else a synthetic smoke run). + """ + paths = [p for p in stage.get("data_paths") or [] if p] + if not paths: + return stage.get("data_path") or cfg.get("data_path", "") + from chronos.workbench.files import merge_datasets + try: + return merge_datasets(paths, os.path.join(writer.run_dir, "data"), log=writer.lines) + except FileNotFoundError: + raise FileNotFoundError(f"no readable records in {', '.join(paths)}") from None + + +def run_stage(session, writer: RunWriter, cfg: dict, stage: dict, backend: str, stop, probe=None) -> str: + from chronos.workbench.training import _prepare_training_run_config + mode = stage["mode"] + data_path = stage_data_path(stage, writer, cfg) + run_cfg, effective = _prepare_training_run_config( + cfg, mode, backend, data_path, None, + stage.get("init_weight", ""), stage.get("max_steps", cfg.get("max_steps", 0)), + stage.get("val_ratio", 0.0), stage.get("teacher_path", ""), stage.get("sample_prompt", "")) + if stage.get("max_steps") is not None: + run_cfg["max_steps"] = int(stage["max_steps"] or 0) + run_cfg["cleanup_gradio_temp_dataset"] = False # never delete a user dataset + session.effective_config = effective + writer.row({"event": "phase_start", "phase": mode, "step": 0}) + writer.heartbeat(state="running", phase=mode, phase_started=_now()) + session.start(run_cfg, mode) + seen, last_sample, last_beat = 0, "", 0.0 + pending = None # validation step whose checkpoint has not been probed yet + probed_mtime = time.time() # only checkpoints newer than this belong to this stage / probe + while True: + alive = session.is_running() + writer.lines(session.drain_log()) + rows = session.get_metrics() + for metric in rows[seen:]: + row = metric_row(mode, metric, session.total_steps) + writer.row(row) + if probe is not None and row.get("event") == "validation": + if pending is not None: # validations arrive faster than checkpoints settle + probe.skip(phase=mode, step=pending["step"], reason=f"superseded by step {row['step']}") + pending = {"step": row["step"], "t": time.monotonic()} + seen = len(rows) + if probe is not None and pending is not None: + # Torch saves before validating; MLX validates, then may save ``.best`` and the final + # checkpoint. Give the trainer a short grace period, then probe the newest checkpoint + # written since the last probe (the weights that were just validated). + ckpt = _fresh_checkpoint(stage_checkpoints(run_cfg, mode), probed_mtime) + due = time.monotonic() - pending["t"] >= PROBE_GRACE_S or not alive + if due and (ckpt is not None or not alive): + probe.request(ckpt, phase=mode, step=pending["step"]) + if ckpt is not None: + probed_mtime = os.path.getmtime(ckpt) + pending = None + if probe is not None: + for entry in probe.poll(): + writer.row(probe_event(entry)) + if session.last_sample and session.last_sample != last_sample: + last_sample = session.last_sample + writer.sample({"phase": mode, "step": session.last_sample_step, + "prompt": run_cfg.get("sample_prompt") or "", "output": last_sample}) + if time.monotonic() - last_beat > 2.0 or not alive: + writer.heartbeat(step=session.step, total=session.total_steps, loss=session.loss, + effective_config=session.get_effective_config()) + last_beat = time.monotonic() + if stop["requested"] and alive: + session.stop() + if not alive: + break + time.sleep(0.5) + state = session.status if session.status in ("finished", "stopped", "error") else "finished" + writer.row({"event": "phase_end", "phase": mode, "step": session.step, "state": state}) + return state + + +def run_job(job: dict) -> int: + from chronos.workbench.training import TrainSession + run_dir = os.path.abspath(job["run_dir"]) + stages = [s for s in job.get("stages", []) if s.get("mode") in STAGE_MODES] + if not stages: + raise ValueError("no valid stages") + writer = RunWriter(run_dir) + _write_json(os.path.join(run_dir, "run_spec.json"), { + "run_dir": run_dir, "source": "chronos-webui", "backend": job.get("backend", "auto"), + "phases": [{"name": s["mode"], "budget_fraction": round(1 / len(stages), 4), + "data_path": s.get("data_path", ""), "data_paths": s.get("data_paths") or [], + "max_steps": s.get("max_steps")} for s in stages], + "model_config": job.get("config", {}), "created": _now(), + }) + stop = {"requested": False} + probe = None + probe_cfg = job.get("lazy_probe") or {} + if probe_cfg.get("enabled", True): + from chronos.workbench.lazy_probe import ProbeScheduler + probe = ProbeScheduler(run_dir, probe_cfg.get("sweep"), tokenizer_dir=job.get("config", {}).get( + "tokenizer_path") or None, log=writer.lines) + + def _on_term(_sig, _frm): + stop["requested"] = True + + signal.signal(signal.SIGTERM, _on_term) + signal.signal(signal.SIGINT, _on_term) + writer.heartbeat(state="starting", stages=[s["mode"] for s in stages]) + final, error = "finished", None + try: + for stage in stages: + session = TrainSession() + state = run_stage(session, writer, dict(job.get("config", {})), stage, job.get("backend", "auto"), + stop, probe) + if state == "error": + final, error = "failed", f"stage {stage['mode']} failed (see train.log)" + break + if stop["requested"] or state == "stopped": + final = "stopped" + break + except Exception as exc: # recorded for the UI, then re-raised for the exit code + final, error = "failed", f"{type(exc).__name__}: {exc}" + writer.lines(error) + raise + finally: + if probe is not None: + if stop["requested"]: + probe.stop() + for entry in probe.wait(): + writer.row(probe_event(entry)) + writer.row({"event": "final", "state": final, **({"error": error} if error else {})}) + writer.heartbeat(state="finished" if final == "finished" else ("failed" if final == "failed" else "stopped"), + error=error) + writer.close() + return 0 if final == "finished" else (1 if final == "failed" else 0) + + +def main(argv=None) -> int: + argv = sys.argv[1:] if argv is None else argv + if len(argv) != 1: + print("usage: python -m chronos.workbench.train_worker ", file=sys.stderr) + return 2 + with open(argv[0], encoding="utf-8") as fh: + return run_job(json.load(fh)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/workbench/training.py b/chronos/workbench/training.py new file mode 100644 index 0000000..d152007 --- /dev/null +++ b/chronos/workbench/training.py @@ -0,0 +1,2062 @@ +"""Training sessions for the Web UI: all six stages, Torch and MLX backends. + +UI-independent (no Gradio). ``TrainSession`` runs one stage in a worker +thread and exposes status, metrics, a log queue and the effective config; +``chronos.workbench.train_worker`` runs it as a subprocess for the Web UI. +Language-dependent defaults (sample prompts, placeholders) use +``set_language_provider``; the default language is English. +""" +import os +import glob +import time +import threading +import queue +import shutil +import tempfile +from contextlib import contextmanager +from types import SimpleNamespace + + +import chronos.deps # auto-bootstrap minimind on sys.path +from ui.presets import MINIMIND_MOE_DEFAULTS +from chronos.backend import training_available, resolve_training_device +from chronos.trainer.device_utils import ( + autocast_context, + configure_cpu_threads, + cpu_thread_snapshot, + dataloader_kwargs, + grad_scaler, + optimizer_step_with_scaler, + resolve_dtype_name, + runtime_summary, +) +from chronos.model.checkpoint import ( + config_dict_for_checkpoint, + load_checkpoint_state_dict, + load_state_dict_controlled, + save_state_dict_with_config, + sniff_checkpoint_config, +) + + +_language_provider = None + + +def set_language_provider(fn) -> None: + """Callable returning the UI language code (``zh-Hans``/``zh-Hant``/``en``/``ja``).""" + global _language_provider + _language_provider = fn + + +def _language() -> str: + try: + return (_language_provider() if _language_provider else None) or "en" + except Exception: + return "en" + + +STAGE_UI_ORDER = ["pretrain", "sft", "dpo", "orpo", "grpo", "distill"] +STAGE_LABELS = { + "pretrain": "Stage 1 · Pretrain", + "sft": "Stage 2 · SFT", + "dpo": "Stage 3 · DPO", + "orpo": "Stage 4 · ORPO", + "grpo": "Stage 5 · GRPO", + "distill": "Stage 6 · Distill", +} +STAGE_CHECKPOINT_PREFIX = { + "pretrain": "chronos", + "sft": "sft", + "dpo": "dpo", + "orpo": "orpo", + "grpo": "grpo", + "distill": "distill", +} +STAGE_DEFAULT_INIT = { + "sft": "chronos", + "dpo": "sft", + "orpo": "sft", + "grpo": "orpo", + "distill": "grpo", +} +STAGE_DEFAULT_DATA = { + "pretrain": "./tests/fixtures/tiny_pretrain.jsonl", + "sft": "./tests/fixtures/tiny_sft.jsonl", + "dpo": "./tests/fixtures/tiny_dpo.jsonl", + "orpo": "./tests/fixtures/tiny_dpo.jsonl", + "grpo": "./tests/fixtures/tiny_grpo.jsonl", + "distill": "./tests/fixtures/tiny_sft.jsonl", +} +CHECKPOINT_TOPOLOGY_KEYS = ( + "hidden_size", + "num_hidden_layers", + "num_experts", + "num_experts_per_tok", + "num_shared_experts", + "moe_intermediate_size", + "vocab_size", + "lookahead_steps", + "kv_latent_dim", + "rope_dim", + "num_attention_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": { + "zh-Hans": "从通用语料继续预训练。`init_weight` 可留空；若存在同名 checkpoint，将按当前拓扑尝试恢复。", + "zh-Hant": "從通用語料繼續預訓練。`init_weight` 可留空；若存在同名 checkpoint，將按目前拓撲嘗試恢復。", + "en": "Continue base pretraining on LM-style data. `init_weight` may be left blank; if a matching checkpoint exists, Train will resume from it.", + "ja": "汎用コーパスで事前学習を続けます。`init_weight` は空で構いません。同名 checkpoint があり拓撲が一致すれば再開します。", + }, + "sft": { + "zh-Hans": "使用 `conversations` 数据做监督微调。`init_weight` 留空时，会从当前 `save_dir` 自动接续 `chronos__moe.pth`；只有手动填路径时才覆盖。", + "zh-Hant": "使用 `conversations` 資料做監督微調。`init_weight` 留空時，會從目前 `save_dir` 自動接續 `chronos__moe.pth`；只有手動填路徑時才覆蓋。", + "en": "Run supervised fine-tuning on `conversations` data. If `init_weight` is blank, Train auto-resolves `chronos__moe.pth` from the current `save_dir`; only a pasted path overrides it.", + "ja": "`conversations` データで SFT を行います。`init_weight` を空欄にすると、現在の `save_dir` から `chronos__moe.pth` を自動解決します。パスを入力した場合のみ上書きします。", + }, + "dpo": { + "zh-Hans": "使用 `chosen / rejected` 偏好对做 DPO。`init_weight` 留空时，会自动加载当前 `save_dir` 下的 `sft__moe.pth`；手动填写路径可覆盖。", + "zh-Hant": "使用 `chosen / rejected` 偏好對做 DPO。`init_weight` 留空時，會自動載入目前 `save_dir` 下的 `sft__moe.pth`；手動填寫路徑可覆蓋。", + "en": "Run DPO on `chosen / rejected` preference pairs. Leaving `init_weight` blank auto-loads `sft__moe.pth` from the current `save_dir`; paste a path only when you want to override it.", + "ja": "`chosen / rejected` の選好ペアで DPO を実行します。`init_weight` を空欄にすると、現在の `save_dir` から `sft__moe.pth` を自動読み込みします。上書きしたい場合だけパスを入力してください。", + }, + "orpo": { + "zh-Hans": "使用同样的偏好对做 ORPO。`init_weight` 留空时，会自动加载当前 `save_dir` 下的 `sft__moe.pth`；手动填写路径可覆盖。", + "zh-Hant": "使用相同的偏好對做 ORPO。`init_weight` 留空時，會自動載入目前 `save_dir` 下的 `sft__moe.pth`；手動填寫路徑可覆蓋。", + "en": "Run ORPO on the same preference-pair format. Leaving `init_weight` blank auto-loads `sft__moe.pth` from the current `save_dir`; paste a path only when you want to override it.", + "ja": "同じ選好ペア形式で ORPO を実行します。`init_weight` を空欄にすると、現在の `save_dir` から `sft__moe.pth` を自動読み込みします。上書きしたい場合だけパスを入力してください。", + }, + "grpo": { + "zh-Hans": "使用 prompt 数据做 GRPO rollout。`init_weight` 留空时，会自动加载当前 `save_dir` 下的 `orpo__moe.pth`；`reward` 可配置。", + "zh-Hant": "使用 prompt 資料做 GRPO rollout。`init_weight` 留空時，會自動載入目前 `save_dir` 下的 `orpo__moe.pth`；`reward` 可設定。", + "en": "Run GRPO rollouts on prompt-style data. Leaving `init_weight` blank auto-loads `orpo__moe.pth` from the current `save_dir`; `reward` remains configurable.", + "ja": "prompt 形式のデータで GRPO rollout を行います。`init_weight` を空欄にすると、現在の `save_dir` から `orpo__moe.pth` を自動読み込みします。`reward` は設定できます。", + }, + "distill": { + "zh-Hans": "蒸馏阶段中，学生 `init_weight` 留空时会自动接续当前 `save_dir` 下的 `grpo__moe.pth`；`teacher_path` 留空时会自动解析 `sft__moe.pth`。只有手动填路径时才覆盖。", + "zh-Hant": "蒸餾階段中，學生 `init_weight` 留空時會自動接續目前 `save_dir` 下的 `grpo__moe.pth`；`teacher_path` 留空時會自動解析 `sft__moe.pth`。只有手動填路徑時才覆蓋。", + "en": "For distillation, a blank student `init_weight` auto-resolves `grpo__moe.pth` from the current `save_dir`, and a blank `teacher_path` auto-resolves `sft__moe.pth`. Paste paths only when you want to override either one.", + "ja": "蒸留では、学生側の `init_weight` を空欄にすると現在の `save_dir` から `grpo__moe.pth` を自動解決し、`teacher_path` を空欄にすると `sft__moe.pth` を自動解決します。上書きしたい場合だけパスを入力してください。", + }, +} +STAGE_METRIC_LABELS = { + "pretrain": ("Total Loss", "CE Loss", "Aux Loss"), + "sft": ("SFT Loss", "CE Loss", "Lookahead"), + "dpo": ("DPO Loss", "Preference", "Anchor"), + "orpo": ("ORPO Loss", "Chosen NLL", "OR Term"), + "grpo": ("GRPO Loss", "Policy", "KL"), + "distill": ("Distill Loss", "KD Loss", "CE Loss"), +} +STAGE_CHART_META = { + "pretrain": { + "title": "Pretrain Metrics", + "primary": "CE", + "secondary": "Aux", + "tertiary": "Temporal", + "val": "Validation", + "speed_title": "Training Speed (steps/s)", + "speed_label": "steps/s", + }, + "sft": { + "title": "SFT Metrics", + "primary": "CE", + "secondary": "Lookahead", + "tertiary": "Anchor", + "val": "Validation", + "speed_title": "SFT Speed (steps/s)", + "speed_label": "steps/s", + }, + "dpo": { + "title": "DPO Metrics", + "primary": "Preference", + "secondary": "Anchor", + "tertiary": "Lookahead", + "val": "Validation", + "speed_title": "Preference Speed (steps/s)", + "speed_label": "steps/s", + }, + "orpo": { + "title": "ORPO Metrics", + "primary": "Chosen NLL", + "secondary": "OR Term", + "tertiary": "Anchor", + "val": "Validation", + "speed_title": "ORPO Speed (steps/s)", + "speed_label": "steps/s", + }, + "grpo": { + "title": "GRPO Metrics", + "primary": "Policy", + "secondary": "KL", + "tertiary": "Reward", + "val": "Validation", + "speed_title": "Rollout Speed + Reward", + "speed_label": "rollouts/s", + "speed_overlay": "Mean Reward", + }, + "distill": { + "title": "Distill Metrics", + "primary": "KD", + "secondary": "CE", + "tertiary": "Anchor", + "val": "Validation", + "speed_title": "Distill Speed (steps/s)", + "speed_label": "steps/s", + }, +} + + +DEFAULT_STAGE_SAMPLE_PROMPTS = { + "pretrain": { + "zh-Hans": "请用简洁中文介绍 Project Chronos 如何把 MoE 的专家加载提前到 Prefill 阶段。", + "zh-Hant": "請用精簡中文介紹 Project Chronos 如何把 MoE 的專家載入提前到 Prefill 階段。", + "en": "Explain how Project Chronos moves MoE expert loading into the prefill stage.", + "ja": "Project Chronos が MoE の専門家ロードを Prefill 段階へ前倒しする仕組みを説明してください。", + }, + "sft": { + "zh-Hans": "用户：为什么 Chronos 的双层路由比传统被动 offload 更适合消费级硬件？\n助手：", + "zh-Hant": "使用者：為什麼 Chronos 的雙層路由比傳統被動 offload 更適合消費級硬體？\n助手：", + "en": "User: Why is Chronos's dual-layer routing better suited to consumer hardware than reactive expert offload?\nAssistant:", + "ja": "ユーザー: なぜ Chronos の二層ルーティングは、従来の受動的な expert offload よりも民生向けハードウェアに向いているのですか？\nアシスタント:", + }, + "dpo": { + "zh-Hans": "用户：请比较 Chronos 和传统 reactive offload 在推理时的关键差异。\n助手：", + "zh-Hant": "使用者：請比較 Chronos 與傳統 reactive offload 在推理時的關鍵差異。\n助手：", + "en": "User: Compare Chronos with conventional reactive offload during inference.\nAssistant:", + "ja": "ユーザー: 推論時における Chronos と従来の reactive offload の主要な違いを比較してください。\nアシスタント:", + }, + "orpo": { + "zh-Hans": "用户：请写一段面向开发者的 Project Chronos 介绍。\n助手：", + "zh-Hant": "使用者：請寫一段面向開發者的 Project Chronos 介紹。\n助手：", + "en": "User: Write a developer-facing introduction to Project Chronos.\nAssistant:", + "ja": "ユーザー: 開発者向けの Project Chronos 紹介文を書いてください。\nアシスタント:", + }, + "grpo": { + "zh-Hans": "请比较 Chronos 与传统 MoE offload 在吞吐、首 token 延迟和缓存未命中行为上的权衡。", + "zh-Hant": "請比較 Chronos 與傳統 MoE offload 在吞吐、首 token 延遲和快取未命中行為上的權衡。", + "en": "Compare Chronos with conventional MoE offload in terms of throughput, first-token latency, and cache-miss behavior.", + "ja": "Chronos と従来の MoE offload を、スループット、初回トークン遅延、キャッシュミス時の挙動で比較してください。", + }, + "distill": { + "zh-Hans": "请用更现代、更有产品感的英文改写 Project Chronos 的首页介绍，保留核心技术点。", + "zh-Hant": "請用更現代、更有產品感的英文改寫 Project Chronos 的首頁介紹，保留核心技術點。", + "en": "Rewrite the Project Chronos landing-page intro in a more modern product voice while preserving the core technical claims.", + "ja": "Project Chronos のトップ紹介文を、技術的な要点を保ったまま、よりモダンでプロダクト寄りの英語に書き換えてください。", + }, +} + + +def _stage_sample_prompt(mode: str, lang: str | None = None) -> str: + lang = lang or _language() + prompts = DEFAULT_STAGE_SAMPLE_PROMPTS.get(mode, DEFAULT_STAGE_SAMPLE_PROMPTS["pretrain"]) + return prompts.get(lang, prompts["en"]) + + +def _default_sample_prompts() -> set[str]: + vals: set[str] = set() + for prompts in DEFAULT_STAGE_SAMPLE_PROMPTS.values(): + vals.update(v.strip() for v in prompts.values() if v and v.strip()) + return vals + + +DEFAULT_SAMPLE_PROMPT_VALUES = _default_sample_prompts() +DEFAULT_STAGE_DATA_VALUES = {v.strip() for v in STAGE_DEFAULT_DATA.values() if v.strip()} +DEFAULT_STAGE_INIT_VALUES = { + f"./out/{prefix}_768_moe.pth" + for prefix in STAGE_DEFAULT_INIT.values() +} + +PRETRAIN_INIT_PLACEHOLDER = { + "zh-Hans": "留空则按当前 save_dir 自动续训同阶段 checkpoint", + "zh-Hant": "留空則按目前 save_dir 自動續訓同階段 checkpoint", + "en": "Leave blank to auto-resume the current-stage checkpoint from save_dir", + "ja": "空欄なら save_dir から同一ステージ checkpoint を自動再開", +} +UPSTREAM_INIT_PLACEHOLDER = { + "zh-Hans": "留空则按当前 save_dir 自动载入 {prefix}__moe.pth，手动填写才覆盖", + "zh-Hant": "留空則按目前 save_dir 自動載入 {prefix}__moe.pth，手動填寫才覆蓋", + "en": "Leave blank to auto-load {prefix}__moe.pth from save_dir; paste a path to override", + "ja": "空欄なら save_dir から {prefix}__moe.pth を自動読み込み。上書き時だけパスを指定", +} +DISTILL_TEACHER_PLACEHOLDER = { + "zh-Hans": "留空则按当前 save_dir 自动载入 sft__moe.pth，手动填写才覆盖", + "zh-Hant": "留空則按目前 save_dir 自動載入 sft__moe.pth，手動填寫才覆蓋", + "en": "Leave blank to auto-load sft__moe.pth from save_dir; paste a path to override", + "ja": "空欄なら save_dir から sft__moe.pth を自動読み込み。上書き時だけパスを指定", +} + +TRAIN_BACKEND_CHOICES = ["auto", "cuda", "xpu", "mlx", "mps", "cpu"] + +TRAIN_BACKEND_LABELS = { + "auto": { + "zh-Hans": "自动 (默认最快)", + "zh-Hant": "自動 (預設最快)", + "en": "Auto (fastest default)", + "ja": "自動 (最速を既定)", + }, + "cuda": { + "zh-Hans": "CUDA", + "zh-Hant": "CUDA", + "en": "CUDA", + "ja": "CUDA", + }, + "mps": { + "zh-Hans": "MPS", + "zh-Hant": "MPS", + "en": "MPS", + "ja": "MPS", + }, + "xpu": { + "zh-Hans": "XPU", + "zh-Hant": "XPU", + "en": "XPU", + "ja": "XPU", + }, + "mlx": { + "zh-Hans": "MLX", + "zh-Hant": "MLX", + "en": "MLX", + "ja": "MLX", + }, + "cpu": { + "zh-Hans": "CPU", + "zh-Hant": "CPU", + "en": "CPU", + "ja": "CPU", + }, +} + + +def _stage_text(stage_map: dict[str, str], mode: str, lang: str | None = None) -> str: + lang = lang or _language() + per_stage = stage_map.get(mode, stage_map["pretrain"]) + return per_stage.get(lang, per_stage["en"]) + + +def _stage_init_placeholder(mode: str, lang: str | None = None) -> str: + lang = lang or _language() + if mode == "pretrain": + return PRETRAIN_INIT_PLACEHOLDER.get(lang, PRETRAIN_INIT_PLACEHOLDER["en"]) + upstream = STAGE_DEFAULT_INIT.get(mode, "chronos") + template = UPSTREAM_INIT_PLACEHOLDER.get(lang, UPSTREAM_INIT_PLACEHOLDER["en"]) + return template.format(prefix=upstream) + + +def _normalize_stage_init_value(current_init_weight: str | None) -> str: + current_init = (current_init_weight or "").strip() + if current_init in DEFAULT_STAGE_INIT_VALUES: + return "" + return current_init + + +def _extract_gradio_file_path(value) -> str: + """Return the current path from Gradio File values across versions. + + Gradio's `type="filepath"` normally passes a string, but depending on + version/event it may pass a small dict, a FileData-like object, or a one + item list. Keep this tolerant so Start can read the latest upload directly + instead of relying only on the textbox update event. + """ + if value in (None, ""): + return "" + if isinstance(value, (list, tuple)): + return _extract_gradio_file_path(value[0] if value else "") + if isinstance(value, dict): + for key in ("path", "name", "orig_name"): + if value.get(key): + return str(value[key]).strip() + return "" + for attr in ("path", "name"): + candidate = getattr(value, attr, None) + if candidate: + return str(candidate).strip() + return str(value).strip() + + +def _distill_teacher_placeholder(lang: str | None = None) -> str: + lang = lang or _language() + return DISTILL_TEACHER_PLACEHOLDER.get(lang, DISTILL_TEACHER_PLACEHOLDER["en"]) + + +def _available_train_backend_choices() -> list[str]: + available = set(training_available()) + return ["auto"] + [name for name in TRAIN_BACKEND_CHOICES[1:] if name in available] + + +def _train_backend_label(name: str, lang: str | None = None) -> str: + lang = lang or _language() + labels = TRAIN_BACKEND_LABELS.get(name, TRAIN_BACKEND_LABELS["auto"]) + return labels.get(lang, labels["en"]) + + +def _train_backend_dropdown_choices(lang: str | None = None) -> list[tuple[str, str]]: + return [ + (_train_backend_label(name, lang), name) + for name in _available_train_backend_choices() + ] + + +def _default_train_backend_value() -> str: + return "auto" + + +def _normalize_dtype_for_trainer(dtype: str | None) -> str: + dt = (dtype or "auto").strip().lower() + if dt in {"auto", ""}: + return "auto" + if dt in {"bf16", "bfloat16"}: + return "bfloat16" + if dt in {"fp32", "float32", "full"}: + return "float32" + return "float16" + + +def _is_auto_value(value) -> bool: + return value in (None, "", "auto", 0, "0") + + +def _cpu_budget_for_backend(cfg: dict, resolved_backend: str) -> float: + requested_threads = cfg.get("cpu_threads", "auto") + raw_budget = cfg.get("cpu_budget_percent", None) + if resolved_backend == "cpu" and _is_auto_value(requested_threads): + return 100.0 if raw_budget in (None, "", "auto") else float(raw_budget) + if raw_budget in (None, "", "auto"): + return 100.0 + return float(raw_budget) + + +def _effective_training_config( + cfg: dict, + *, + mode: str, + requested_backend: str, + resolved_backend: str, + device: str, + cpu_threads: int, + workers: int | None = None, + total_train_steps: int | None = None, + optimizer_update_steps: int | None = None, + progress_unit: str = "micro_step", +) -> dict: + keys = [ + "learning_rate", + "batch_size", + "accumulation_steps", + "epochs", + "max_steps", + "max_seq_len", + "save_interval", + "log_interval", + "weight_decay", + "grad_clip", + "dtype", + "cpu_threads", + "cpu_budget_percent", + "num_workers", + "max_gen_len", + "num_generations", + "temperature", + "beta", + "alpha", + "lambda_or", + "reward_spec", + "teacher_path", + "data_path", + "data_path_source", + "lambda_router_anchor", + "stage_overrides_applied", + ] + out = { + "mode": mode, + "requested_backend": requested_backend, + "resolved_backend": resolved_backend, + "device": device, + "effective_cpu_threads": int(cpu_threads), + } + if workers is not None: + out["effective_num_workers"] = int(workers) + if total_train_steps is not None: + out["total_train_steps"] = int(total_train_steps) + if optimizer_update_steps is not None: + out["optimizer_update_steps"] = int(optimizer_update_steps) + if total_train_steps is not None or optimizer_update_steps is not None: + out["progress_unit"] = progress_unit + for key in keys: + if key in cfg: + out[key] = cfg[key] + return out + + +def _format_effective_training_config(summary: dict) -> str: + parts = [] + for key in [ + "mode", + "requested_backend", + "resolved_backend", + "device", + "dtype", + "learning_rate", + "batch_size", + "accumulation_steps", + "epochs", + "max_steps", + "max_seq_len", + "weight_decay", + "grad_clip", + "save_interval", + "log_interval", + "effective_cpu_threads", + "cpu_budget_percent", + "effective_num_workers", + "total_train_steps", + "optimizer_update_steps", + "progress_unit", + ]: + if key in summary: + parts.append(f"{key}={summary[key]}") + return "Effective config: " + ", ".join(parts) + + +def _coerce_max_steps(value) -> int | None: + try: + steps = int(value or 0) + except Exception: + return None + return steps if steps > 0 else None + + +def _planned_train_steps(data_iter, epochs: int, max_steps=None) -> int: + """Return the real UI progress denominator in micro/batch steps. + + Dataset loaders must always see the full corpus. `max_steps` is only a + training-loop cap, so it is applied after deriving len(data_iter) * epochs. + """ + cap = _coerce_max_steps(max_steps) + planned = 0 + try: + planned = int(len(data_iter)) * max(1, int(epochs or 1)) + except Exception: + planned = 0 + if cap is not None: + return min(planned, cap) if planned > 0 else cap + return max(0, planned) + + +def _planned_epoch_iters(data_iter, max_steps=None) -> int: + cap = _coerce_max_steps(max_steps) + try: + n = int(len(data_iter)) + except Exception: + n = 0 + if cap is not None and n > 0: + return max(1, min(cap, n)) + return max(1, n) + + +def _planned_optimizer_updates(train_steps: int, accumulation_steps: int) -> int: + train_steps = max(0, int(train_steps or 0)) + accum = max(1, int(accumulation_steps or 1)) + 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, + selected_backend: str | None, + data_path: str | None, + dataset_upload, + init_weight: str | None, + max_steps, + val_ratio, + teacher_path: str | None, + sample_prompt: str | None, +) -> 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) + cfg["train_backend"] = requested_backend + cfg["device"] = resolved_device + resolved_data_path, data_path_source = _resolve_training_data_path( + data_path or cfg.get("data_path", ""), + dataset_upload, + ) + cfg["data_path"] = resolved_data_path + cfg["data_path_source"] = data_path_source + cfg["init_weight"] = _normalize_stage_init_value(init_weight) + if "max_steps" not in cfg: + cfg["max_steps"] = 0 + cfg["val_ratio"] = float(val_ratio or 0.0) + cfg["teacher_path"] = teacher_path or "" + cfg["sample_prompt"] = sample_prompt or "" + cfg["cpu_budget_percent"] = _cpu_budget_for_backend(cfg, resolved_backend) + initial_effective = _effective_training_config( + cfg, + mode=train_mode, + requested_backend=requested_backend, + resolved_backend=resolved_backend, + device=resolved_device, + cpu_threads=0, + workers=None, + ) + return cfg, initial_effective + + +def _gradio_temp_roots() -> list[str]: + roots = [] + for env_key in ("GRADIO_TEMP_DIR", "GRADIO_TEMP"): + value = os.environ.get(env_key) + if value: + roots.append(value) + roots.append(os.path.join(tempfile.gettempdir(), "gradio")) + out = [] + seen = set() + for root in roots: + try: + real = os.path.realpath(os.path.abspath(os.path.expanduser(str(root)))) + except Exception: + continue + if real not in seen: + out.append(real) + seen.add(real) + return out + + +def _path_in_gradio_temp(path: str | None) -> tuple[bool, str, str]: + if not path: + return False, "", "" + try: + real_path = os.path.realpath(os.path.abspath(os.path.expanduser(str(path)))) + except Exception: + return False, "", "" + for root in _gradio_temp_roots(): + try: + if os.path.commonpath([real_path, root]) == root and real_path != root: + return True, real_path, root + except ValueError: + continue + return False, real_path, "" + + +def _project_root() -> str: + return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def _training_upload_cache_root() -> str: + root = os.environ.get("CHRONOS_TRAIN_UPLOAD_CACHE_DIR") + if root: + return os.path.realpath(os.path.abspath(os.path.expanduser(root))) + return os.path.join(_project_root(), ".chronos_uploads", "datasets") + + +def _path_in_training_upload_cache(path: str | None) -> tuple[bool, str, str]: + if not path: + return False, "", "" + try: + real_path = os.path.realpath(os.path.abspath(os.path.expanduser(str(path)))) + root = os.path.realpath(os.path.abspath(_training_upload_cache_root())) + return os.path.commonpath([real_path, root]) == root and real_path != root, real_path, root + except Exception: + return False, "", "" + + +def _stage_gradio_upload_for_training(path: str | None) -> str: + """Move Gradio's private temp upload into a stable project cache. + + Training directly from `/private/.../T/gradio` is fragile on macOS: Gradio + owns that directory and may rotate or clean files independently of the + training thread. Manual paths are returned unchanged. + """ + path = _extract_gradio_file_path(path) + inside, real_path, _root = _path_in_gradio_temp(path) + if not inside or not os.path.exists(real_path): + return path or "" + cache_root = _training_upload_cache_root() + os.makedirs(cache_root, exist_ok=True) + base = os.path.basename(real_path) or "dataset.jsonl" + stem, ext = os.path.splitext(base) + if not ext: + ext = ".jsonl" + stamp = time.strftime("%Y%m%d_%H%M%S") + size = os.path.getsize(real_path) + target = os.path.join(cache_root, f"{stamp}_{os.getpid()}_{size}_{stem[:48]}{ext}") + shutil.copy2(real_path, target) + _cleanup_gradio_temp_dataset(real_path) + return target + + +def _resolve_training_data_path(current_path: str | None, upload_value=None) -> tuple[str, str]: + """Resolve the dataset path at Start time. + + The upload `change` callback updates the textbox in normal operation, but + Gradio's private upload/event state can lag behind Start. Resolve again + when Start is clicked so a newly uploaded file wins over the default tiny + fixture immediately. If the upload temp file was already moved into the + project cache, keep the textbox path. + """ + current = (current_path or "").strip() + upload_path = _extract_gradio_file_path(upload_value) + if upload_path: + staged = _stage_gradio_upload_for_training(upload_path) + if staged and os.path.exists(staged): + if staged != current: + return staged, "upload" + return staged, "textbox" + if current: + return current, "textbox" + return "", "empty" + + +def _close_dataset_handles(*objects) -> None: + """Close lazy JSONL file handles before deleting Gradio upload temp files.""" + seen: set[int] = set() + + def visit(obj): + if obj is None: + return + oid = id(obj) + if oid in seen: + return + seen.add(oid) + fh = getattr(obj, "_fh", None) + if fh is not None: + try: + fh.close() + except Exception: + pass + try: + obj._fh = None + except Exception: + pass + for attr in ("dataset", "datasets"): + child = getattr(obj, attr, None) + if child is not None: + visit(child) + if isinstance(obj, (list, tuple, set)): + for item in obj: + visit(item) + elif isinstance(obj, dict): + for item in obj.values(): + visit(item) + + for item in objects: + visit(item) + + +def _cleanup_gradio_temp_dataset(path: str | None) -> dict: + """Delete a dataset upload copied into Gradio's temp area. + + The guard deliberately refuses to touch any path outside Gradio's temp + root, so manual local dataset paths remain untouched. + """ + inside, real_path, root = _path_in_gradio_temp(path) + result = { + "path": real_path, + "root": root, + "deleted": False, + "removed_dirs": [], + "skipped": not inside, + "reason": "" if inside else "not_gradio_temp", + } + if not inside: + return result + if not os.path.lexists(real_path): + result["skipped"] = True + result["reason"] = "missing" + return result + try: + if os.path.isdir(real_path) and not os.path.islink(real_path): + shutil.rmtree(real_path) + else: + os.unlink(real_path) + result["deleted"] = True + result["skipped"] = False + except Exception as exc: + result["skipped"] = True + result["reason"] = f"delete_failed: {exc}" + return result + + parent = os.path.dirname(real_path) + while parent and parent != root: + try: + os.rmdir(parent) + result["removed_dirs"].append(parent) + except OSError: + break + parent = os.path.dirname(parent) + return result + + +def _cleanup_training_dataset_artifact(path: str | None) -> dict: + """Clean staged upload artifacts while preserving user-owned paths.""" + inside, real_path, root = _path_in_training_upload_cache(path) + if not inside: + return _cleanup_gradio_temp_dataset(path) + result = { + "path": real_path, + "root": root, + "deleted": False, + "removed_dirs": [], + "skipped": False, + "reason": "", + } + if not os.path.lexists(real_path): + result["skipped"] = True + result["reason"] = "missing" + return result + try: + if os.path.isdir(real_path) and not os.path.islink(real_path): + shutil.rmtree(real_path) + else: + os.unlink(real_path) + result["deleted"] = True + except Exception as exc: + result["skipped"] = True + result["reason"] = f"delete_failed: {exc}" + return result + parent = os.path.dirname(real_path) + while parent and parent != root: + try: + os.rmdir(parent) + result["removed_dirs"].append(parent) + except OSError: + break + parent = os.path.dirname(parent) + return result + + +@contextmanager +def _trainer_logger_sink(put_line): + from trainer import trainer_utils as _tu # type: ignore + + orig_logger = _tu.Logger + + def _hook(content): + line = str(content) + put_line(line) + try: + orig_logger(content) + except Exception: + pass + + _tu.Logger = _hook + try: + yield + finally: + _tu.Logger = orig_logger + + +def _sniff_checkpoint(path: str) -> dict: + try: + cfg, _sources = config_dict_for_checkpoint(path, require_unsniffable=False) + return cfg or sniff_checkpoint_config(path) + except Exception: + return {} + + +def _verify_checkpoint_warn(checkpoint_path: str, model_cfg, put_line) -> None: + try: + from chronos.verify import verify_checkpoint + + report = verify_checkpoint( + checkpoint_path, + model_cfg, + device="cpu", + include_mlx=True, + write_json=True, + ) + status = "ok" if report.get("ok") else "warning" + put_line(f"[verify] {status}; report={report.get('verify_json', '')}") + for warning in report.get("warnings", []): + put_line(f"[verify warning] {warning}") + except Exception as exc: + put_line(f"[verify warning] checkpoint verification failed: {exc}") + + +def _sync_metal_backend(device: str | None) -> None: + """Flush pending Metal work before a WebUI training thread exits.""" + if str(device or "").lower() != "mps": + return + try: + import torch + + if hasattr(torch, "mps"): + synchronize = getattr(torch.mps, "synchronize", None) + if synchronize is not None: + synchronize() + empty_cache = getattr(torch.mps, "empty_cache", None) + if empty_cache is not None: + empty_cache() + except Exception: + pass + + +# ── Background trainer thread ───────────────────────────────────── + +class TrainSession: + def __init__(self): + self._stop = threading.Event() + self._thread = None + self._log_q: queue.Queue = queue.Queue(maxsize=5000) + self.status = "idle" + self.step = 0 + self.loss = None + # Separate metric history for charts + self._metrics: list[dict] = [] # [{step, total, ce, aux, temporal, tps}] + self._metric_lock = threading.Lock() + # Progress + ETA — set when training kicks off + self.total_steps: int = 0 + self.t_start: float = 0.0 + # M8a: live generation sample (updated each save_interval). + 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) + + def is_running(self): + return self._thread is not None and self._thread.is_alive() + + def start(self, cfg: dict, mode: str): + if self.is_running(): + return + self._stop.clear() + self.status = "running" + self.step = 0 + self.loss = None + 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 + ) + self._thread.start() + + def stop(self): + self._stop.set() + self.status = "stopped" + + def _put(self, msg: str): + try: + self._log_q.put_nowait(msg) + except queue.Full: + pass + + def _put_metric(self, m: dict): + with self._metric_lock: + self._metrics.append(m) + + def drain_log(self): + """Return all pending log lines as a single string (no newline prefix).""" + lines = [] + while True: + try: + lines.append(self._log_q.get_nowait()) + except queue.Empty: + break + return "\n".join(lines) + + def get_metrics(self): + with self._metric_lock: + return list(self._metrics) + + def get_effective_config(self): + return dict(self.effective_config) + + def _stage_checkpoint_path(self, save_dir: str, mode: str, hidden_size: int) -> str: + prefix = STAGE_CHECKPOINT_PREFIX.get(mode, "chronos") + return os.path.join(save_dir, f"{prefix}_{hidden_size}_moe.pth") + + def _default_init_path(self, save_dir: str, mode: str, hidden_size: int) -> str: + upstream = STAGE_DEFAULT_INIT.get(mode) + if not upstream: + return "" + return os.path.join(save_dir, f"{upstream}_{hidden_size}_moe.pth") + + def _resolve_default_init_path(self, save_dir: str, mode: str, hidden_size: int) -> str: + exact = self._default_init_path(save_dir, mode, hidden_size) + if os.path.exists(exact): + return exact + upstream = STAGE_DEFAULT_INIT.get(mode) + if not upstream: + return exact + candidates = [ + p for p in glob.glob(os.path.join(save_dir, f"{upstream}_*_moe.pth")) + if os.path.isfile(p) + ] + if not candidates: + return exact + return max(candidates, key=os.path.getmtime) + + 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 sniffed[k] != model_cfg_kwargs[k]: + mismatches.append(f"{k}: ckpt={sniffed[k]} != ui={model_cfg_kwargs[k]}") + return mismatches + + def _adopt_checkpoint_topology(self, model_cfg_kwargs: dict, sniffed: dict) -> dict: + adopted = {} + for key in CHECKPOINT_TOPOLOGY_KEYS: + if key not in sniffed: + continue + current = model_cfg_kwargs.get(key) + value = sniffed[key] + if current is None or current != value: + adopted[key] = value + model_cfg_kwargs[key] = value + return adopted + + def _build_stage_args(self, cfg: dict, mode: str, save_dir: str, hidden_size: int): + reward_spec = (cfg.get("reward_spec") or "toy").strip() or "toy" + teacher_path = (cfg.get("teacher_path") or "").strip() + return SimpleNamespace( + device=cfg.get("device", "cpu"), + dtype=_normalize_dtype_for_trainer(cfg.get("dtype")), + learning_rate=float(cfg.get("learning_rate", 5e-4)), + accumulation_steps=max(1, int(cfg.get("accumulation_steps", 1))), + grad_clip=float(cfg.get("grad_clip", 1.0)), + cpu_threads=cfg.get("cpu_threads", "auto"), + cpu_budget_percent=cfg.get("cpu_budget_percent", 100), + num_workers=cfg.get("num_workers", "auto"), + log_interval=max(1, int(cfg.get("log_interval", 10))), + save_interval=max(1, int(cfg.get("save_interval", 500))), + save_dir=save_dir, + epochs=max(1, int(cfg.get("epochs", 1))), + batch_size=max(1, int(cfg.get("batch_size", 1))), + max_seq_len=max(8, int(cfg.get("max_seq_len", 512))), + max_gen_len=max(1, int(cfg.get("max_gen_len", 24))), + num_generations=max(1, int(cfg.get("num_generations", 4))), + temperature=float(cfg.get("temperature", 1.0)), + beta=float(cfg.get("beta", 0.1)), + alpha=float(cfg.get("alpha", 0.7)), + lambda_or=float(cfg.get("lambda_or", 0.1)), + lambda_router_anchor=float(cfg.get("lambda_router_anchor", 0.0)), + reward=reward_spec, + teacher_path=teacher_path or os.path.join(save_dir, f"sft_{hidden_size}_moe.pth"), + hidden_size=hidden_size, + num_hidden_layers=int(cfg.get("num_hidden_layers", 8)), + num_experts=int(cfg.get("num_experts", 4)), + steps=int(cfg.get("max_steps", 0)) or None, + weight_decay=float(cfg.get("weight_decay", 0.01)), + ) + + def _record_stage_metric(self, mode: str, step: int, result, tps: float = 0.0): + meta = STAGE_CHART_META.get(mode, STAGE_CHART_META["pretrain"]) + if isinstance(result, tuple): + if mode == "sft": + total, ce_loss, aux_loss, lookahead, anchor = result + metric = { + "total": total, + "primary": ce_loss, + "secondary": lookahead, + "tertiary": anchor, + "ce": ce_loss, + "aux": lookahead, + "temporal": anchor, + } + elif mode == "dpo": + total, pref_loss, _aux_loss, lookahead, anchor = result + metric = { + "total": total, + "primary": pref_loss, + "secondary": anchor, + "tertiary": lookahead, + "ce": pref_loss, + "aux": anchor, + "temporal": lookahead, + } + elif mode == "orpo": + total, chosen_nll, or_term, _lookahead, anchor = result + metric = { + "total": total, + "primary": chosen_nll, + "secondary": or_term, + "tertiary": anchor, + "ce": chosen_nll, + "aux": or_term, + "temporal": anchor, + } + elif mode == "distill": + total, kd_loss, ce_loss, _lookahead, anchor = result + metric = { + "total": total, + "primary": kd_loss, + "secondary": ce_loss, + "tertiary": anchor, + "ce": kd_loss, + "aux": ce_loss, + "temporal": anchor, + } + else: + return + elif isinstance(result, dict) and mode == "grpo": + metric = { + "total": result.get("loss", 0.0), + "primary": result.get("pg_loss", 0.0), + "secondary": result.get("kl", 0.0), + "tertiary": result.get("mean_reward", 0.0), + "ce": result.get("pg_loss", 0.0), + "aux": result.get("kl", 0.0), + "temporal": result.get("mean_reward", 0.0), + } + else: + return + + if mode == "pretrain": + metric.setdefault("primary", metric.get("ce", 0.0)) + metric.setdefault("secondary", metric.get("aux", 0.0)) + metric.setdefault("tertiary", metric.get("temporal", 0.0)) + + metric.update({"step": int(step), "tps": float(tps)}) + self.step = int(step) + self.loss = float(metric["total"]) + self._put_metric(metric) + self._put( + f"[{mode.upper()} summary] step={step} total={metric['total']:.4f} " + f"{meta['primary']}={metric.get('primary', 0.0):.4f} " + f"{meta['secondary']}={metric.get('secondary', 0.0):.4f} " + f"{meta['tertiary']}={metric.get('tertiary', 0.0):.4f} " + f"tps={tps:.2f}" + ) + + def _generate_live_sample(self, model, device: str, cfg: dict, mode: str): + import torch + + sp = (cfg.get("sample_prompt") or "").strip() + if not sp: + sp = _stage_sample_prompt(mode) + tokenizer = self._load_tokenizer() + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": sp}], + tokenize=False, + add_generation_prompt=True, + ) + enc = tokenizer(rendered, return_tensors="pt", add_special_tokens=False) + ids = enc.input_ids.to(device) + was_training = model.training + model.eval() + with torch.no_grad(): + gen = ids.clone() + for _ in range(80): + 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): + break + txt = tokenizer.decode(gen[0, ids.shape[1]:].tolist(), skip_special_tokens=True) + if was_training: + model.train() + self.last_sample = txt + self.last_sample_step = self.step or self.last_sample_step + self._put(f"[sample @ step {self.last_sample_step or 0}] {txt[:200]}") + + def _run(self, cfg: dict, mode: str): + t_start = time.monotonic() + data_path_for_cleanup = (cfg.get("data_path") or "").strip() + cleanup_after_run = bool(cfg.get("cleanup_gradio_temp_dataset", True)) + train_ds = val_ds = full_ds = loader = val_loader = None + try: + import torch + from torch.utils.data import DataLoader + + from chronos.model.config import ChronosConfig + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.model.moe_chronos import ChronosMOEFeedForward + + self._put(f"[{mode.upper()}] Building model...") + 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 = "" + init_sniffed = {} + if mode != "pretrain": + load_path = init_weight or self._resolve_default_init_path( + save_dir, mode, int(model_cfg_kwargs["hidden_size"]) + ) + if os.path.exists(load_path): + init_sniffed = _sniff_checkpoint(load_path) + adopted = self._adopt_checkpoint_topology(model_cfg_kwargs, init_sniffed) + if adopted: + summary = ", ".join(f"{k}={v}" for k, v in adopted.items()) + self._put( + f"[{mode.upper()}] Adopted topology from init checkpoint " + f"{load_path}: {summary}" + ) + + model_cfg = ChronosConfig(**model_cfg_kwargs) + model = ChronosForCausalLM(model_cfg) + params_m = sum(p.numel() for p in model.parameters()) / 1e6 + self._put(f"Model: {params_m:.1f}M params " + f"(H={model_cfg.hidden_size}, L={model_cfg.num_hidden_layers}, " + f"E={model_cfg.num_experts}, ffn={model_cfg.intermediate_size}, " + f"vocab={model_cfg.vocab_size})") + + ckp_path = self._stage_checkpoint_path(save_dir, mode, model_cfg.hidden_size) + mlx_checkpoint_path = None + + # Init-weight resolution: + # - pretrain: optionally resume from its own checkpoint if present + # - other stages: require an upstream weight (explicit init_weight + # or the stage's default predecessor checkpoint). + if mode == "pretrain": + resume_path = init_weight or ckp_path + if os.path.exists(resume_path): + sniffed = _sniff_checkpoint(resume_path) + mismatches = self._topology_mismatches(sniffed, model_cfg_kwargs) + if mismatches: + self._put( + "Resume checkpoint exists but does not match current UI config; " + "starting fresh instead.\n " + "\n ".join(mismatches) + ) + else: + weights = load_checkpoint_state_dict(resume_path, map_location="cpu") + load_state_dict_controlled(model, weights) + self._put(f"Resumed from {resume_path}") + mlx_checkpoint_path = resume_path + else: + self._put("Pretraining from random init") + else: + if not os.path.exists(load_path): + raise FileNotFoundError( + f"[{mode.upper()}] requires an upstream checkpoint to initialize from. " + f"Tried: {load_path}\n" + f"Set 'init_weight' in the Train tab to point at a valid .pth, " + f"or run the prior stage first." + ) + + sniffed = init_sniffed or _sniff_checkpoint(load_path) + mismatch_hints = self._topology_mismatches(sniffed, model_cfg_kwargs) + if mismatch_hints: + raise RuntimeError( + f"[{mode.upper()}] init checkpoint topology does not match current UI config:\n " + + "\n ".join(mismatch_hints) + + "\nLoad the matching preset/config first, or switch init_weight to the correct checkpoint." + ) + + weights = load_checkpoint_state_dict(load_path, map_location="cpu") + load_state_dict_controlled(model, weights) + self._put(f"[{mode.upper()}] Initialized from {load_path}") + mlx_checkpoint_path = load_path + + 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")) + self._normalise_effective_cpu_config(cfg, resolved_backend) + cpu_threads = configure_cpu_threads( + cfg.get("cpu_threads", "auto"), + budget_percent=cfg.get("cpu_budget_percent", 100), + ) + rt = runtime_summary(device, cfg.get("dtype", "auto")) + self._put(f"Training backend: {resolved_backend} (requested: {requested_backend})") + self._put(f"Device: {device}") + self._put( + f"Runtime: dtype={rt.dtype} cpu_threads={rt.cpu_threads} " + f"autocast={rt.autocast} scaler={rt.scaler}" + ) + thread_diag = cpu_thread_snapshot() + self._put( + "CPU thread env: " + f"torch={thread_diag.get('torch_num_threads')} " + f"interop={thread_diag.get('torch_num_interop_threads')} " + f"physical={thread_diag.get('physical_cores')} " + f"OMP={thread_diag.get('OMP_NUM_THREADS')} " + f"MKL={thread_diag.get('MKL_NUM_THREADS')} " + f"VECLIB={thread_diag.get('VECLIB_MAXIMUM_THREADS')} " + f"NUMEXPR={thread_diag.get('NUMEXPR_NUM_THREADS')}" + ) + + data_path = data_path_for_cleanup + self._put(f"Dataset: {data_path or ''} (source={cfg.get('data_path_source', 'unknown')})") + max_seq_len = cfg.get("max_seq_len", 512) + batch_size = cfg.get("batch_size", 4) + accum = max(1, int(cfg.get("accumulation_steps", 8))) + epochs = cfg.get("epochs", 1) + save_interval = cfg.get("save_interval", 500) + log_interval = max(1, cfg.get("log_interval", 10)) + val_ratio = float(cfg.get("val_ratio", 0.05)) + max_steps = _coerce_max_steps(cfg.get("max_steps", 0)) # 0 => no cap + base_lr = float(cfg.get("learning_rate", 5e-4)) + weight_decay = float(cfg.get("weight_decay", 0.01)) + optimizer = None + + if not data_path or not os.path.exists(data_path): + self._put("No dataset found — running synthetic smoke test (50 steps)") + loader = self._synthetic_loader(model_cfg.vocab_size, max_seq_len, batch_size, n=50) + val_loader = None + else: + tokenizer = self._load_tokenizer() + from chronos.data.flexible_dataset import ( + 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": + 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"}: + 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: + full_ds = None + # Deterministic val split: every 1/val_ratio-th record goes to val + # (idx-modulo split → reproducible without rescanning the file). + if full_ds is not None and val_ratio > 0 and len(full_ds) >= 20: + stride = max(2, int(round(1.0 / val_ratio))) + val_idx = [i for i in range(len(full_ds)) if i % stride == 0] + train_idx = [i for i in range(len(full_ds)) if i % stride != 0] + from torch.utils.data import Subset + train_ds = Subset(full_ds, train_idx) + val_ds = Subset(full_ds, val_idx) + self._put(f"Split: train={len(train_ds)} val={len(val_ds)}") + elif full_ds is not None: + train_ds, val_ds = full_ds, None + else: + train_ds, val_ds = None, None + loader = ( + 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 + ) + val_loader = ( + 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 + ) + if loader is not None: + workers = getattr(loader, "num_workers", 0) + self._put( + f"CPU threads configured: {cpu_threads}; DataLoader workers: {workers}" + ) + else: + workers = None + if mode in {"grpo"}: + tokenizer = self._load_tokenizer() + from chronos.trainer.grpo_trainer import load_grpo_prompts + prompts = load_grpo_prompts(data_path, max_prompts=None) + if not prompts: + raise ValueError(f"[{mode.upper()}] no prompts found in {data_path}") + else: + prompts = None + + stage_args = self._build_stage_args(cfg, mode, save_dir, model_cfg.hidden_size) + stage_args.device = device + train_iter_for_steps = prompts if mode == "grpo" else loader + planned_total_steps = _planned_train_steps( + train_iter_for_steps, + stage_args.epochs, + stage_args.steps, + ) + planned_optimizer_updates = _planned_optimizer_updates( + planned_total_steps, + stage_args.accumulation_steps, + ) + self.total_steps = planned_total_steps + self.effective_config = _effective_training_config( + cfg, + mode=mode, + requested_backend=requested_backend, + resolved_backend=resolved_backend, + device=device, + cpu_threads=cpu_threads, + workers=workers, + total_train_steps=planned_total_steps, + optimizer_update_steps=planned_optimizer_updates, + ) + self._put(_format_effective_training_config(self.effective_config)) + + if resolved_backend == "mlx": + from chronos.mlx.training import run_mlx_stage + + self._put("[MLX] Running native MLX trainer; skipping PyTorch model.to()/optimizer path.") + if mlx_checkpoint_path: + self._put(f"[MLX {mode.upper()}] Loading checkpoint → {mlx_checkpoint_path}") + else: + 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: + self.total_steps = total + self._put( + f"[MLX {mode.upper()}] started dtype={event.get('dtype', '')} " + 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) + total = int(event.get("steps", self.total_steps) or self.total_steps) + if total > 0: + self.total_steps = total + loss = float(event.get("loss", 0.0) or 0.0) + tps = float(event.get("tps", 0.0) or 0.0) + self.step = step + self.loss = loss + metric = { + "step": step, + "total": 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" + ) + return + if kind == "rollback": + step = int(event.get("step", 0) or 0) + total = int(event.get("steps", self.total_steps) or self.total_steps) + if total > 0: + self.total_steps = total + self._put( + f"[MLX {mode.upper()} rollback] step={step}/{self.total_steps} " + f"reason={event.get('reason', '')} rollbacks={int(event.get('rollbacks', 0) or 0)} " + f"next_lr={float(event.get('lr', 0.0) or 0.0):.2e}" + ) + return + if kind == "save": + self._put( + f"[MLX {mode.upper()}] saved checkpoint ({event.get('reason', '')}) → " + f"{event.get('checkpoint_path', '')}" + ) + return + if kind == "stopped": + if int(event.get("steps", 0) or 0) > 0: + self.total_steps = int(event.get("steps", 0) or 0) + self._put( + f"[MLX {mode.upper()}] stopped at step={int(event.get('step', 0) or 0)}/{self.total_steps}" + ) + return + if kind == "failed": + self._put( + 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 + if kind == "complete": + self._put( + f"[MLX {mode.upper()}] completed step={int(event.get('step', 0) or 0)}/" + f"{int(event.get('steps', self.total_steps) or self.total_steps)}" + ) + + result = run_mlx_stage( + stage=mode, + config=model_cfg, + 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, + progress_callback=_mlx_progress, + stop_event=self._stop, + ) + self.step = int(result.steps) + self.loss = float(result.last_loss) + self.total_steps = max(self.total_steps, int(result.total_steps or 0), self.step) + elapsed = time.monotonic() - t_start + if self._stop.is_set(): + self.status = "stopped" + self._put(f"[MLX {mode.upper()}] Stopped in {elapsed:.1f}s. Saved → {result.checkpoint_path}") + if result.checkpoint_saved: + _verify_checkpoint_warn(result.checkpoint_path, model_cfg, self._put) + return + if result.stopped: + self.status = "error" + self._put( + 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: + _verify_checkpoint_warn(result.checkpoint_path, model_cfg, self._put) + return + self._put(f"[MLX {mode.upper()}] Training complete in {elapsed:.1f}s. Saved → {result.checkpoint_path}") + _verify_checkpoint_warn(result.checkpoint_path, model_cfg, self._put) + self.status = "finished" + return + + model = model.to(device) + 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 + step_t = time.monotonic() + best_val = float("inf") + best_path = ckp_path.replace(".pth", ".best.pth") + if mode in {"sft", "dpo", "orpo", "grpo", "distill"}: + tokenizer = self._load_tokenizer() + + def _eval_val() -> float: + if val_loader is None: + return float("nan") + model.eval() + tot, n = 0.0, 0 + with torch.no_grad(): + 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) + 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) + + if mode in {"sft", "dpo", "orpo", "grpo", "distill"}: + with _trainer_logger_sink(self._put): + if mode == "sft": + from chronos.trainer.sft_trainer import ChronosSFTTrainer + + 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)) + 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, 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 + 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) + if step_idx % stage_args.save_interval == 0: + trainer._save(epoch, step_idx) + self._generate_live_sample(model, device, cfg, mode) + if self._stop.is_set() or (stage_args.steps is not None and stage_global_step >= stage_args.steps): + break + trainer._save(epoch=stage_args.epochs - 1, step=iters) + elif mode == "dpo": + from chronos.trainer.dpo_trainer import ChronosDPOTrainer + + trainer = ChronosDPOTrainer(model, 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)) + calib = torch.cat([first["x_chosen"], first["x_rejected"]], dim=0) + trainer.set_calibration_batch(calib) + self._put("[DPO] 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, 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(batch, 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) + if step_idx % stage_args.save_interval == 0: + trainer._save(epoch, step_idx) + self._generate_live_sample(model, device, cfg, mode) + if self._stop.is_set() or (stage_args.steps is not None and stage_global_step >= stage_args.steps): + break + trainer._save(epoch=stage_args.epochs - 1, step=iters) + elif mode == "orpo": + from chronos.trainer.orpo_trainer import ChronosORPOTrainer + + trainer = ChronosORPOTrainer(model, 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)) + calib = torch.cat([first["x_chosen"], first["x_rejected"]], dim=0) + trainer.set_calibration_batch(calib) + self._put("[ORPO] 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, 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(batch, 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) + if step_idx % stage_args.save_interval == 0: + trainer._save(epoch, step_idx) + self._generate_live_sample(model, device, cfg, mode) + if self._stop.is_set() or (stage_args.steps is not None and stage_global_step >= stage_args.steps): + break + trainer._save(epoch=stage_args.epochs - 1, step=iters) + elif mode == "grpo": + from chronos.trainer.grpo_trainer import ChronosGRPOTrainer + from chronos.trainer.reward import build_reward_fn + + reward_fn = build_reward_fn(stage_args.reward) + trainer = ChronosGRPOTrainer(model, model_cfg, stage_args, tokenizer, reward_fn=reward_fn) + if float(getattr(model_cfg, "lambda_router_anchor", 0.0)) > 0 and prompts: + calib = tokenizer( + prompts[0], return_tensors="pt", truncation=True, max_length=stage_args.max_seq_len + ).input_ids + trainer.set_calibration_batch(calib) + self._put("[GRPO] Router anchor reference captured.") + iters = _planned_epoch_iters(prompts, stage_args.steps) + stage_global_step = 0 + for epoch in range(stage_args.epochs): + epoch_start = time.monotonic() + for step_idx, prompt in enumerate(prompts, 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(prompt, 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) + if step_idx % stage_args.save_interval == 0: + trainer._save(epoch, step_idx) + self._generate_live_sample(model, device, cfg, mode) + if self._stop.is_set() or (stage_args.steps is not None and stage_global_step >= stage_args.steps): + break + trainer._save(epoch=stage_args.epochs - 1, step=iters) + elif mode == "distill": + from chronos.trainer.distill_trainer import ChronosDistillTrainer + + teacher_path = stage_args.teacher_path + if not os.path.exists(teacher_path): + raise FileNotFoundError(f"[DISTILL] teacher not found: {teacher_path}") + teacher = ChronosForCausalLM(model_cfg) + t_state = load_checkpoint_state_dict(teacher_path, map_location="cpu") + load_state_dict_controlled(teacher, t_state) + teacher = teacher.to(device) + teacher.eval().requires_grad_(False) + self._put(f"[DISTILL] Teacher: {teacher_path}") + 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)) + 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, 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 + 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) + if step_idx % stage_args.save_interval == 0: + trainer._save(epoch, step_idx) + self._generate_live_sample(model, device, cfg, mode) + if self._stop.is_set() or (stage_args.steps is not None and stage_global_step >= stage_args.steps): + break + trainer._save(epoch=stage_args.epochs - 1, step=iters) + self._generate_live_sample(model, device, cfg, mode) + elapsed = time.monotonic() - t_start + self._put(f"[{mode.upper()}] Training complete in {elapsed:.1f}s. Saved → {ckp_path}") + _verify_checkpoint_warn(ckp_path, model_cfg, self._put) + self.status = "finished" + return + + stop_outer = False + for epoch in range(epochs): + if self._stop.is_set() or stop_outer: + break + 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: + 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)] + + aux_val = out.aux_loss.item() if hasattr(out, "aux_loss") else 0.0 + + # Use the shared loss helper so lambda_temporal / + # lambda_lookahead actually flow gradient back to the + # gate. + from chronos.trainer.loss_mixin import ( + chronos_loss_term, collect_router_probs, + collect_offload_training_metrics, + ) + loss = chronos_loss_term( + model, out.loss, lp, model_cfg, aux_loss=out.aux_loss, + ) + # Logging-only: detached probs are fine here. + router_4d_det = collect_router_probs(model, with_grad=False) + if router_4d_det is not None and router_4d_det.shape[1] > 1: + from chronos.model.temporal_loss import temporal_locality_loss + temp_val = temporal_locality_loss( + router_4d_det.mean(dim=2) + ).item() + offload_metrics = collect_offload_training_metrics( + model, lp, model_cfg, + ) + else: + temp_val = 0.0 + offload_metrics = {} + + step_loss = loss / accum + + scaler.scale(step_loss).backward() + global_step += 1 + self.step = global_step + self.loss = loss.item() + + # LR schedule (warmup → cosine). Refresh every step. + sched_total = self.total_steps or (global_step + 1) + apply_lr(optimizer, get_lr(global_step, sched_total, base_lr)) + + # Optimizer.step() once per `accum` micro-batches. + # Off-by-one fix: trigger when (global_step % accum == 0). + # Previous code did the same check but had no warmup-aware + # accumulation count, occasionally stepping after step 1. + if global_step % accum == 0: + optimizer_step_with_scaler( + scaler, + optimizer, + model.parameters(), + 1.0, + ) + paged_training.observe() + + if global_step % log_interval == 0: + now = time.monotonic() + tps = log_interval / max(now - step_t, 1e-6) + step_t = now + ce_val = out.loss.item() + tot_val = loss.item() + + self._put( + f"Epoch {epoch+1}/{epochs} Step {global_step} " + f"loss={tot_val:.4f} ce={ce_val:.4f} " + f"aux={aux_val:.4f} temporal={temp_val:.4f} " + f"topk_recall={offload_metrics.get('lookahead_topk_recall', 0.0):.2%} " + f"union_recall={offload_metrics.get('lookahead_union_recall', 0.0):.2%} " + f"router_jaccard={offload_metrics.get('router_adjacent_jaccard', 0.0):.2%} " + f"steps/s={tps:.2f} lr={optimizer.param_groups[0]['lr']:.2e}" + ) + self._put_metric({ + "step": global_step, + "total": tot_val, + "ce": ce_val, + "aux": aux_val, + "temporal": temp_val, + "tps": tps, + **offload_metrics, + }) + + if global_step % save_interval == 0: + os.makedirs(save_dir, exist_ok=True) + save_state_dict_with_config( + model, + ckp_path, + model_cfg, + stage=STAGE_CHECKPOINT_PREFIX.get(mode, mode), + ) + self._put(f"Saved checkpoint → {ckp_path}") + + # Validation eval + best-ckpt tracking. + v = _eval_val() + if v == v: # not NaN + self._put(f"Validation loss: {v:.4f} (best={best_val:.4f})") + self._put_metric({ + "step": global_step, "total": v, "ce": v, + "aux": 0.0, "temporal": 0.0, "tps": 0.0, + "val": True, + "val_loss": v, + }) + if v < best_val: + best_val = v + save_state_dict_with_config( + model, + best_path, + model_cfg, + stage=STAGE_CHECKPOINT_PREFIX.get(mode, mode), + ) + self._put(f" ↑ new best → {best_path}") + + # M8a: live qualitative sample so the user sees + # whether the model is actually learning, not just + # whether the loss number went down. We greedy-decode + # ~80 tokens from the user-supplied prompt, and fall + # back to a stage-aware evaluation prompt if blank. + try: + sp = (cfg.get("sample_prompt") or "").strip() + if not sp: + sp = _stage_sample_prompt(mode) + tokenizer = self._load_tokenizer() + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": sp}], + tokenize=False, + add_generation_prompt=True, + ) + ids = tokenizer( + rendered, return_tensors="pt", add_special_tokens=False + ).input_ids.to(device) + model.eval() + with torch.no_grad(): + gen = ids.clone() + for _ in range(80): + 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): + break + txt = tokenizer.decode(gen[0, ids.shape[1]:].tolist(), skip_special_tokens=True) + model.train() + self.last_sample = txt + self.last_sample_step = global_step + self._put(f"[sample @ step {global_step}] {txt[:200]}") + except Exception as gen_err: + self._put(f"[sample failed: {gen_err}]") + + os.makedirs(save_dir, exist_ok=True) + save_state_dict_with_config( + model, + ckp_path, + model_cfg, + stage=STAGE_CHECKPOINT_PREFIX.get(mode, mode), + ) + elapsed = time.monotonic() - t_start + self._put(f"Training complete in {elapsed:.1f}s. Saved → {ckp_path}") + _verify_checkpoint_warn(ckp_path, model_cfg, self._put) + self.status = "finished" + + except Exception as e: + import traceback + self._put(f"[ERROR] {e}\n{traceback.format_exc()}") + self.status = "error" + finally: + _sync_metal_backend((cfg or {}).get("device")) + _close_dataset_handles(loader, val_loader, train_ds, val_ds, full_ds) + if cleanup_after_run: + cleanup = _cleanup_training_dataset_artifact(data_path_for_cleanup) + if cleanup.get("deleted"): + removed_dirs = len(cleanup.get("removed_dirs") or []) + suffix = f"; removed {removed_dirs} empty parent dirs" if removed_dirs else "" + self._put(f"[cleanup] removed staged upload dataset: {cleanup.get('path')}{suffix}") + elif cleanup.get("reason") not in {"not_gradio_temp", "missing"}: + self._put( + f"[cleanup] skipped staged upload dataset cleanup for " + f"{cleanup.get('path')}: {cleanup.get('reason')}" + ) + + def _synthetic_loader(self, vocab_size, seq_len, batch_size, n=50): + import torch + from torch.utils.data import TensorDataset, DataLoader + ids = torch.randint(0, vocab_size, (n * batch_size, seq_len)) + ds = TensorDataset(ids, ids.clone()) + return DataLoader(ds, batch_size=batch_size) + + def _load_tokenizer(self): + from transformers import AutoTokenizer + 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() diff --git a/chronos_app.py b/chronos_app.py index f4e7351..c4d72c8 100644 --- a/chronos_app.py +++ b/chronos_app.py @@ -1,20 +1,23 @@ """ -chronos_app.py — Backwards-compatible launcher shim. +chronos_app.py — launcher for the Chronos Web UI from a source checkout. -For installed users, prefer: chronos-ui -For repo users, you can run: python chronos_app.py [--port N] [--share] +Installed users run ``chronos-ui`` (or ``chronos webui``); both start the same server. -This file simply delegates to chronos.app.main(). + python chronos_app.py [--port N] [--host H] [--workdir DIR] [--runs-dir DIR] + python chronos_app.py --legacy [--port N] [--share] # the old Gradio console + +The default is the FastAPI + single-page UI (chronos/webui). It listens on loopback without a login; +any other host needs --allow-remote and a key in CHRONOS_UI_KEY or CHRONOS_API_KEY. """ -import sys import os +import sys -# Ensure Project_Chronos root is on sys.path when run directly +# Ensure the Project_Chronos root is on sys.path when run directly. _here = os.path.dirname(os.path.abspath(__file__)) if _here not in sys.path: sys.path.insert(0, _here) -from chronos.app import main +from chronos.webui.server import main if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 3bdc377..8703168 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "Project_Chronos" -version = "0.3.1a0" +version = "0.3.1a1" description = "Storage-aware recurrent-depth MoE: lookahead expert prefetch, SSD-lazy experts, hybrid MLA/SWA attention, self-speculative decoding" readme = "README.md" license = "Apache-2.0" @@ -45,7 +45,7 @@ Issues = "https://github.com/FonaTech/Project_Chronos/issues" [project.scripts] chronos = "chronos.cli:main" -chronos-ui = "chronos.app:main" +chronos-ui = "chronos.webui.server:main" chronos-train = "train_chronos:main" chronos-eval = "chronos.eval.io_profiler:main" chronos-bench = "chronos.eval.benchmark:main" @@ -68,6 +68,7 @@ include = ["chronos*", "ui*"] [tool.setuptools.package-data] ui = ["event_contract.js", "presets_v3.json"] "chronos.monitor" = ["static/*.html", "static/*.js", "static/*.css"] +"chronos.webui" = ["static/*.html", "static/*.js", "static/*.css", "static/*.svg", "i18n/*.json"] [tool.black] line-length = 100 diff --git a/tests/test_lazy_probe.py b/tests/test_lazy_probe.py new file mode 100644 index 0000000..364be6b --- /dev/null +++ b/tests/test_lazy_probe.py @@ -0,0 +1,241 @@ +"""Mid-training lazy-expert probe (chronos.workbench.lazy_probe) and its train-worker wiring.""" +from __future__ import annotations + +import json +import os +import sys + +import pytest + +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +from chronos.workbench.lazy_probe import ( # noqa: E402 + ProbeScheduler, ProbeSweep, agreement, expand_variants, newest_checkpoint, variant_label, +) + + +def test_sweep_axes_collapse_per_backend_and_dedupe_by_resolved_layout(): + sweep = ProbeSweep.from_dict({"backends": ["cpu", "mlx"], "paging": ["auto", "cluster", "expert"], + "residency": ["planned", "window"], "ram_ratios": [0.5], + "miss_policies": ["sync_on_demand"], "speculative": [False, True], + "max_variants": 256}) + variants, total = expand_variants(sweep, lambda p: "cluster" if p == "auto" else p) + assert total == len(variants) + torch_v = [v for v in variants if v["backend"] == "cpu"] + mlx_v = [v for v in variants if v["backend"] == "mlx"] + assert len(torch_v) == 2 * 2 # {cluster, expert} × spec; residency n/a + assert len(mlx_v) == 2 * 2 * 2 # layouts × residency × spec + assert all("residency" not in v and v["miss_policy"] == "sync_on_demand" for v in torch_v) + assert all("residency_decay" in v for v in mlx_v if v["residency"] == "planned") + assert all("residency_decay" not in v for v in mlx_v if v["residency"] == "window") + assert all(("verify_wait_s" in v) == v["speculative"] for v in mlx_v) + capped, total = expand_variants(ProbeSweep.from_dict({**sweep.as_dict(), "max_variants": 3}), + lambda p: p) + assert len(capped) == 3 and total > 3 + assert variant_label(mlx_v[0]).startswith("mlx/cluster/r0.5/planned/d0.7/w12ms") + + +@pytest.mark.parametrize("bad", [dict(paging=["louvain"]), dict(residency=["queue"]), dict(ram_ratios=[0]), + dict(residency_decay=[1.0]), dict(decode_wait_s=[-1]), + dict(miss_policies=["drop"]), dict(max_new_tokens=0), dict(gamma=9)]) +def test_sweep_rejects_invalid_settings(bad): + with pytest.raises(ValueError): + ProbeSweep.from_dict(bad) + + +def test_agreement_and_newest_checkpoint(tmp_path): + assert agreement([1, 2, 3], [1, 2, 3]) == {"exact": True, "prefix": 3, "positionwise": 1.0, "ref_len": 3} + a = agreement([1, 9, 3, 4], [1, 2, 3]) + assert not a["exact"] and a["prefix"] == 1 and a["positionwise"] == pytest.approx(2 / 3) + old, new = tmp_path / "a.pth", tmp_path / "b.pth" + old.write_bytes(b"x") + new.write_bytes(b"y") + os.utime(old, (1, 1)) + assert newest_checkpoint([str(old), str(new), str(tmp_path / "missing.pth")]) == str(new) + assert newest_checkpoint([]) is None + + +torch = pytest.importorskip("torch") + + +@pytest.fixture +def trained_ckpt(tmp_path, monkeypatch): + from test_lazy_experts import _save + monkeypatch.setenv("CHRONOS_EXPERT_CACHE_ROOT", str(tmp_path / "cache")) + return _save(tmp_path, 0.3, "trained") + + +def test_probe_torch_sync_is_exact_and_fallback_is_reported(trained_ckpt): + from chronos.workbench.lazy_probe import run_probe + res = run_probe(trained_ckpt, {"backends": ["cpu"], "paging": ["auto", "expert"], "ram_ratios": [0.5], + "miss_policies": ["sync_on_demand", "fallback_diagnostic"], + "max_new_tokens": 8}) + assert res["auto_paging"] == "cluster" and res["conditions"]["prompt_source"] == "synthetic" + by = {v["label"]: v for v in res["variants"]} + assert set(by) == {"cpu/cluster/r0.5/sync_on_demand", "cpu/cluster/r0.5/fallback_diagnostic", + "cpu/expert/r0.5/sync_on_demand", "cpu/expert/r0.5/fallback_diagnostic"} + for label, v in by.items(): + assert "error" not in v, v + m = v["measured"] + assert 0.0 <= m["resident_hit_rate"] <= 1.0 and m["tokens"] > 0 and m["tokens_per_s"] > 0 + assert m["paging_resolved"] == label.split("/")[1] + for layout in ("cluster", "expert"): + sync = by[f"cpu/{layout}/r0.5/sync_on_demand"] + assert sync["fidelity"]["exact_match_rate"] == 1.0 # a miss loads: identical to full memory + assert sync["measured"]["on_demand_loads"] > 0 + assert by[f"cpu/{layout}/r0.5/fallback_diagnostic"]["measured"]["on_demand_loads"] == 0 + assert res["references"]["cpu"]["measured"]["tokens_per_s"] > 0 + + +def test_probe_mlx_variants_are_exact(trained_ckpt): + mx = pytest.importorskip("mlx.core") + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + from chronos.workbench.lazy_probe import run_probe + res = run_probe(trained_ckpt, {"backends": ["mlx"], "paging": ["auto", "expert"], "ram_ratios": [0.5], + "residency": ["planned", "window"], "speculative": [False, True], + "max_new_tokens": 8}) + finally: + mx.set_default_device(previous) + assert len(res["variants"]) == 8 and not res["truncated"] + for v in res["variants"]: + assert "error" not in v, v + assert v["fidelity"]["exact_match_rate"] == 1.0, v["label"] + assert v["measured"]["prediction_hit_rate"] is None or 0 <= v["measured"]["prediction_hit_rate"] <= 1 + + +def test_dense_checkpoint_is_skipped(tmp_path, monkeypatch): + from chronos.model.checkpoint import save_state_dict_with_config + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.workbench.lazy_probe import run_probe + from test_lazy_experts import _tiny + cfg = _tiny(0.0) + cfg.use_moe = False + path = str(tmp_path / "dense.pth") + save_state_dict_with_config(ChronosForCausalLM(cfg).eval(), path, cfg, stage="pretrain") + assert "skipped" in run_probe(path, {"backends": ["cpu"]}) + + +def test_scheduler_snapshots_runs_in_child_and_cleans_up(trained_ckpt, tmp_path): + run_dir = tmp_path / "run" + sched = ProbeScheduler(str(run_dir), {"backends": ["cpu"], "paging": ["auto"], "ram_ratios": [1.0], + "max_new_tokens": 4}, nice=0) + assert "skipped" in sched.request(str(tmp_path / "nope.pth"), phase="pretrain", step=1) + assert sched.request(trained_ckpt, phase="pretrain", step=10)["started"] + assert "skipped" in sched.request(trained_ckpt, phase="pretrain", step=11) # one probe at a time + entries = sched.wait(timeout_s=300) + assert len(entries) == 1 and "error" not in entries[0], entries + variant = entries[0]["result"]["variants"][0] + assert variant["label"] == "cpu/cluster/r1/sync_on_demand" and variant["fidelity"]["exact_match_rate"] == 1.0 + assert not (run_dir / "lazy_probe").exists() or not os.listdir(run_dir / "lazy_probe") + lines = [json.loads(x) for x in (run_dir / "lazy_probe.jsonl").read_text().splitlines()] + assert [("skipped" in x, x["step"]) for x in lines] == [(True, 1), (True, 11), (False, 10)] + assert os.path.exists(trained_ckpt) # the original checkpoint is untouched + + +class _StubSession: + """TrainSession stand-in: writes the stage checkpoint, then reports a validation.""" + + def __init__(self, source_ckpt: str, delay_s: float = 0.0): + self.source, self.delay = source_ckpt, delay_s + self.status, self.step, self.loss, self.total_steps = "idle", 0, 0.0, 4 + self.last_sample, self.last_sample_step, self.effective_config = "", 0, {} + self._metrics, self._thread = [], None + + def start(self, cfg, mode): + import shutil + import threading + import time + + def run(): + self.status = "running" + for step in (1, 2, 3, 4): + self.step, self.loss = step, 5.0 - step + self._metrics.append({"step": step, "total": self.loss, "primary": self.loss, + "secondary": 0.1, "tertiary": 0.0, "tps": 10.0}) + time.sleep(0.6) + # MLX order: validation first, the checkpoint shortly after. + self._metrics.append({"step": 4, "total": 1.5, "val": True, "val_loss": 1.5}) + time.sleep(self.delay) + os.makedirs(cfg["save_dir"], exist_ok=True) + sidecar = os.path.splitext(self.source)[0] + ".config.json" + if os.path.exists(sidecar): # save_state_dict_with_config writes both + shutil.copy(sidecar, os.path.join(cfg["save_dir"], "chronos_32_moe.config.json")) + shutil.copy(self.source, os.path.join(cfg["save_dir"], "chronos_32_moe.pth")) + time.sleep(0.6) + self.status = "finished" + + self._thread = threading.Thread(target=run, daemon=True) + self._thread.start() + + def is_running(self): + return self._thread is not None and self._thread.is_alive() + + def stop(self): + pass + + def drain_log(self): + return "" + + def get_metrics(self): + return list(self._metrics) + + def get_effective_config(self): + return {} + + +def test_worker_probes_checkpoint_after_validation_and_monitor_shows_it(trained_ckpt, tmp_path, monkeypatch): + from chronos.monitor.snapshot import RunMonitor + from chronos.workbench import train_worker + from chronos.workbench import training + + monkeypatch.setattr(train_worker, "PROBE_GRACE_S", 0.5) + monkeypatch.setattr(training, "TrainSession", lambda: _StubSession(trained_ckpt, delay_s=1.0)) + run_dir = tmp_path / "runs" / "webui_job" + job = {"run_dir": str(run_dir), "backend": "cpu", + "config": {"save_dir": str(tmp_path / "out"), "hidden_size": 32}, + "stages": [{"mode": "pretrain", "max_steps": 4, "val_ratio": 0.1}], + "lazy_probe": {"enabled": True, "sweep": {"backends": ["cpu"], "paging": ["auto", "expert"], + "ram_ratios": [0.5], "max_new_tokens": 4}}} + assert train_worker.run_job(job) == 0 + rows = [json.loads(x) for x in (run_dir / "metrics.jsonl").read_text().splitlines()] + kinds = [r.get("event") for r in rows if r.get("event")] + assert kinds.index("validation") < kinds.index("lazy_probe") < kinds.index("final") + probe_row = next(r for r in rows if r.get("event") == "lazy_probe") + assert probe_row["step"] == 4 and probe_row["probe_variants"] == 2 and "error" not in probe_row + assert probe_row["probe_min_exact_match"] == 1.0 # sync_on_demand matches full memory + train_row = next(r for r in rows if "event" not in r) + assert train_row["primary_loss"] == 4.0 and train_row["balance_loss"] == 0.1 + + snap = RunMonitor(str(tmp_path / "runs")).snapshot("webui_job") + assert snap["status"]["state"] == "finished" + (probe,) = snap["lazy_probe"] + assert probe["step"] == 4 and probe["auto_paging"] == "cluster" + assert {v["label"] for v in probe["variants"]} == {"cpu/cluster/r0.5/sync_on_demand", + "cpu/expert/r0.5/sync_on_demand"} + for v in probe["variants"]: + assert 0 <= v["measured"]["resident_hit_rate"] <= 1 and v["fidelity"]["exact_match_rate"] == 1.0 + assert probe["conditions"]["label"] == "measured" + + +def test_worker_records_skip_when_validation_has_no_checkpoint(tmp_path, monkeypatch): + from chronos.workbench import train_worker + from chronos.workbench import training + + class NoSave(_StubSession): + def start(self, cfg, mode): + super().start({**cfg, "save_dir": str(tmp_path / "elsewhere")}, mode) + + src = tmp_path / "not_a_model.pth" + src.write_bytes(b"x") + monkeypatch.setattr(train_worker, "PROBE_GRACE_S", 0.2) + monkeypatch.setattr(training, "TrainSession", lambda: NoSave(str(src))) + run_dir = tmp_path / "run" + job = {"run_dir": str(run_dir), "backend": "cpu", "config": {"save_dir": str(tmp_path / "out")}, + "stages": [{"mode": "pretrain", "max_steps": 4}], "lazy_probe": {"sweep": {"backends": ["cpu"]}}} + assert train_worker.run_job(job) == 0 + lines = [json.loads(x) for x in (run_dir / "lazy_probe.jsonl").read_text().splitlines()] + assert len(lines) == 1 and "skipped" in lines[0] and lines[0]["step"] == 4 diff --git a/tests/test_smoke.py b/tests/test_smoke.py index cb72dc4..eb20d75 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -597,7 +597,7 @@ def test_train_backend_choice_helpers(): _train_backend_dropdown_choices, ) - with mock.patch("ui.tabs.train_tab.training_available", return_value=["mlx", "mps", "cpu"]): + with mock.patch("chronos.workbench.training.training_available", return_value=["mlx", "mps", "cpu"]): assert _available_train_backend_choices() == ["auto", "mlx", "mps", "cpu"] choices = _train_backend_dropdown_choices("en") assert choices[0][1] == "auto" @@ -1956,8 +1956,8 @@ def fake_run_mlx_stage(**kwargs): return Result() monkeypatch.setattr("chronos.mlx.training.run_mlx_stage", fake_run_mlx_stage) - monkeypatch.setattr("ui.tabs.train_tab.resolve_training_device", lambda _requested: ("mlx", "mlx")) - monkeypatch.setattr("ui.tabs.train_tab._verify_checkpoint_warn", lambda *args, **kwargs: None) + monkeypatch.setattr("chronos.workbench.training.resolve_training_device", lambda _requested: ("mlx", "mlx")) + monkeypatch.setattr("chronos.workbench.training._verify_checkpoint_warn", lambda *args, **kwargs: None) session = TrainSession() messages = [] @@ -2052,7 +2052,7 @@ def test_train_cleanup_removes_only_gradio_temp_dataset(monkeypatch, tmp_path): def test_train_start_merge_preserves_config_weight_decay(monkeypatch, tmp_path): from ui.tabs import train_tab - monkeypatch.setattr(train_tab, "resolve_training_device", lambda requested: ("cpu", "cpu")) + monkeypatch.setattr("chronos.workbench.training.resolve_training_device", lambda requested: ("cpu", "cpu")) cfg, effective = train_tab._prepare_training_run_config( {"weight_decay": 0.123, "cpu_threads": "auto"}, "pretrain", @@ -2075,7 +2075,7 @@ def test_train_start_prefers_latest_uploaded_dataset(monkeypatch, tmp_path): from types import SimpleNamespace from ui.tabs import train_tab - monkeypatch.setattr(train_tab, "resolve_training_device", lambda requested: ("cpu", "cpu")) + monkeypatch.setattr("chronos.workbench.training.resolve_training_device", lambda requested: ("cpu", "cpu")) gradio_root = tmp_path / "gradio" upload_dir = gradio_root / "upload_hash" upload_dir.mkdir(parents=True) @@ -2540,7 +2540,7 @@ def test_generate_api_returns_plain_json_with_chart_records(): "decode_diagnostics": {"lazy_offload": {"generated_token_ids": [2], "warning": "empty clean decode"}}, "chart": [{"metric": "Decode speed", "mode": "lazy_offload", "value": 1.0, "unit": "tokens/s"}], } - with mock.patch("ui.tabs.inference_tab._run_inference_modes", return_value=fake_raw): + with mock.patch("chronos.workbench.inference._run_inference_modes", return_value=fake_raw): out = mod.generate_api({}, "cpu", "offload", "", "hello", 1, 0.1, False, "1.00", False) assert isinstance(out, dict) diff --git a/tests/test_ui_event_browser.py b/tests/test_ui_event_browser.py index 541f124..a0ddfca 100644 --- a/tests/test_ui_event_browser.py +++ b/tests/test_ui_event_browser.py @@ -38,7 +38,7 @@ def browser_server(tmp_path_factory): port = sock.getsockname()[1] url = f"http://127.0.0.1:{port}" log = (tmp_path_factory.mktemp("browser-server") / "server.log").open("w+") - process = subprocess.Popen([sys.executable, "-u", "chronos_app.py", "--host", "127.0.0.1", "--port", str(port)], + process = subprocess.Popen([sys.executable, "-u", "chronos_app.py", "--legacy", "--host", "127.0.0.1", "--port", str(port)], cwd=ROOT, stdout=log, stderr=subprocess.STDOUT) try: deadline = time.monotonic() + 40 diff --git a/tests/test_webui.py b/tests/test_webui.py new file mode 100644 index 0000000..6faf85e --- /dev/null +++ b/tests/test_webui.py @@ -0,0 +1,503 @@ +"""Chronos Web UI (FastAPI + SPA): catalogs, language detection, security, forms, jobs.""" +import json +import os +import re + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from chronos.webui import schema # noqa: E402 +from chronos.webui import server # noqa: E402 + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +STATIC = os.path.join(ROOT, "chronos", "webui", "static") +PH = re.compile(r"\{(\w+)\}") +MUTATE = {"X-Chronos-Request": "1"} + + +def _catalog(lang): + with open(os.path.join(ROOT, "chronos", "webui", "i18n", f"{lang}.json"), encoding="utf-8") as fh: + return json.load(fh) + + +def _client(tmp_path, **kw): + from fastapi.testclient import TestClient + app = server.create_app(str(tmp_path), runs_dir="runs", **kw) + return TestClient(app, base_url="http://127.0.0.1"), app + + +# ── catalogs ──────────────────────────────────────────────────────────── +def test_catalogs_have_identical_keys_placeholders_and_no_empty_values(): + en = _catalog("en") + for lang in schema.LANGUAGES: + cat = _catalog(lang) + assert set(cat) == set(en), (lang, sorted(set(en) ^ set(cat))[:10]) + assert not [k for k, v in cat.items() if not str(v).strip()], lang + bad = [k for k in en if set(PH.findall(en[k])) != set(PH.findall(cat[k]))] + assert not bad, (lang, bad[:10]) + + +def test_every_schema_field_group_choice_and_rule_is_translated(): + need = {f"group.{g}" for g in schema.GROUPS} | {f"source.{s}" for s in schema.SOURCES} + need |= {f"rule.{r}" for r in schema.RULE_REFS} | set(schema.choice_keys()) + for f in schema.ALL_FIELDS: + need |= {f"field.{f.key}.label", f"field.{f.key}.help"} + for lang in schema.LANGUAGES: + missing = sorted(need - set(_catalog(lang))) + assert not missing, (lang, missing[:10]) + + +def test_static_files_use_only_catalog_keys(): + en = _catalog("en") + call = re.compile(r"""\b(?:t|button|pageHead|has)\(\s*["']([a-z0-9_.]+)["'](?:\s*,\s*["']([a-z0-9_.]+)["'])?""") + used = set() + for name in os.listdir(STATIC): + text = open(os.path.join(STATIC, name), encoding="utf-8").read() + if name.endswith(".html"): + used |= set(re.findall(r'data-i18n="([^"]+)"', text)) + elif name.endswith(".js"): + for a, b in call.findall(text): + used |= {k for k in (a, b) if k and "." in k} + used |= set(re.findall(r'"(probe\.skip\.[a-z_]+)"', text)) + assert used, "no keys found" + assert not sorted(used - set(en)) + nav = re.findall(r'\["([a-z]+)",\s*render', open(os.path.join(STATIC, "app.js"), encoding="utf-8").read()) + assert len(nav) == 11 and all(f"nav.{n}" in en for n in nav) + + +def test_catalog_endpoint_and_unknown_language(tmp_path): + client, _ = _client(tmp_path) + body = client.get("/api/i18n/ja").json() + assert body["language"] == "ja" and body["messages"] == _catalog("ja") + assert client.get("/api/i18n/fr").status_code == 404 + + +# ── language detection ────────────────────────────────────────────────── +@pytest.mark.parametrize("tag,expected", [ + ("zh-CN,zh;q=0.9", "zh-Hans"), ("zh", "zh-Hans"), ("zh-SG", "zh-Hans"), ("zh-Hans-CN", "zh-Hans"), + ("zh-TW", "zh-Hant"), ("zh-HK", "zh-Hant"), ("zh-MO", "zh-Hant"), ("zh-Hant", "zh-Hant"), + ("ja-JP", "ja"), ("en-GB", "en"), ("fr-FR,de;q=0.8", "en"), ("fr-FR,ja;q=0.5", "ja"), ("", "en"), (None, "en"), +]) +def test_resolve_language(tag, expected): + assert server.resolve_language(tag) == expected + + +def test_session_reports_browser_language(tmp_path): + client, _ = _client(tmp_path) + r = client.get("/api/session", headers={"Accept-Language": "zh-TW,zh;q=0.8"}).json() + # TestClient's peer is "testclient", not loopback: the server treats it as another machine + assert r == {"auth_required": False, "authenticated": True, "local_only": True, "language": "zh-Hant", + "access": "lan", "lan_enabled": False, "native_dialog": False, "max_upload_mb": server.MAX_UPLOAD_MB} + + +# ── security ──────────────────────────────────────────────────────────── +def test_host_header_allowlist(tmp_path): + client, _ = _client(tmp_path) + assert client.get("/health").status_code == 200 + assert client.get("/health", headers={"Host": "evil.example"}).status_code == 421 + + +def test_mutations_need_marker_header_and_same_origin(tmp_path): + client, _ = _client(tmp_path) + assert client.post("/api/logout").status_code == 403 + assert client.post("/api/logout", headers={**MUTATE, "Origin": "http://evil.example"}).status_code == 403 + assert client.post("/api/logout", headers={**MUTATE, "Origin": "http://127.0.0.1"}).status_code == 200 + + +def test_security_headers_and_csp(tmp_path): + client, _ = _client(tmp_path) + r = client.get("/") + assert r.status_code == 200 + csp = r.headers.get("content-security-policy", "") + assert "script-src 'self'" in csp and "unsafe-inline" not in csp.split("style-src")[0] + html = r.text + assert "