From 5381ac86310cb6d425dad35f3c8175d881adf0f5 Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:10:26 +0800 Subject: [PATCH 1/3] Release 0.3.0: v3 recurrent-depth kernel, task-agnostic UI, speculative switch - v3 kernel: shared recurrent groups with contractive update, shared_first / per_iteration / first_only attention modes, tensor-only ACT exit, top-2 + shared-expert MoE with sorted grouped dispatch (Torch argsort, MLX gather_mm), MLA absorbed/expanded decode, int8 KV, single-dtype precision chain. - Recurrent-depth self-speculative decoding on Torch and MLX, greedy-exact, with lazy-expert masks and lookahead prefetch; Inference tab switch, draft depth and gamma on all four runner paths. - Web UI monitor and Data tab no longer assume a task: metrics, samples and dataset schemas are discovered from what a run or dataset contains. - chronos serve (Ollama/OpenAI protocol), tokenizer builder, corpus pipeline, longrun trainer, monitor, model bundles, planner. - Fixes: MLX engine forces eval mode; lazy Torch engine uses no_grad so paged experts stay usable; Inference tab uses the checkpoint's own tokenizer. - tools/bench_speculative.py and measured reports in reports/speculative/ (exact, not faster at 50M). - torch >= 2.6 (transformers 5 needs it); CI pins tested versions in .github/ci-constraints.txt and drops the dead token-based publish job. Co-Authored-By: Claude Opus 5.5 --- .github/ci-constraints.txt | 21 + .github/workflows/ci.yml | 26 +- .github/workflows/python-app.yml | 4 +- .gitignore | 24 + CHANGELOG.md | 85 ++ MANIFEST.in | 22 +- README.md | 244 +++- README_zh.md | 278 ++++- benchmark_compare.py | 35 +- benchmark_log.txt | 58 +- benchmark_results.json | 71 +- chronos/__init__.py | 2 +- chronos/app.py | 18 +- chronos/bundle/__init__.py | 22 + chronos/bundle/builder.py | 342 ++++++ chronos/bundle/cli.py | 177 +++ chronos/bundle/docs.py | 320 +++++ chronos/bundle/finetune.py | 128 ++ chronos/bundle/recommend.py | 182 +++ chronos/bundle/runtime.py | 204 ++++ chronos/cli.py | 41 + chronos/config_schema.py | 224 ++++ chronos/data/__init__.py | 11 + chronos/data/corpus/__init__.py | 7 + chronos/data/corpus/clean.py | 105 ++ chronos/data/corpus/download.py | 161 +++ chronos/data/corpus/evalsets.py | 65 ++ chronos/data/corpus/pipeline.py | 205 ++++ chronos/data/corpus/presets.py | 63 + chronos/data/corpus/readers.py | 141 +++ chronos/data/corpus/shards.py | 237 ++++ chronos/data/corpus/sources.py | 107 ++ chronos/data/flexible_dataset.py | 440 +++++++ chronos/deps.py | 8 + chronos/eval/benchmark.py | 23 +- chronos/eval/io_profiler.py | 2 +- chronos/export.py | 82 +- chronos/io/async_prefetcher.py | 57 +- chronos/io/cluster_layout.py | 37 +- chronos/io/expert_store.py | 435 ++++--- chronos/io/storage.py | 51 +- chronos/migrate.py | 51 + chronos/mlx/attention.py | 351 ++++-- chronos/mlx/expert_store.py | 100 +- chronos/mlx/inference.py | 178 ++- chronos/mlx/kv_quant.py | 87 ++ chronos/mlx/model.py | 212 +++- chronos/mlx/moe.py | 91 +- chronos/mlx/moe_kernels.py | 156 +++ chronos/mlx/recurrent.py | 89 ++ chronos/mlx/speculative.py | 227 ++++ chronos/mlx/training/inspect_checkpoint.py | 179 +++ chronos/mlx/training/io.py | 19 +- chronos/mlx/training/longrun.py | 730 ++++++++++++ chronos/mlx/training/optim.py | 45 + .../mlx/training/recipes/enzh_50m_72h.json | 89 ++ chronos/mlx/training/telemetry.py | 236 ++++ chronos/mlx/training/trainer.py | 542 +++++---- chronos/model/__init__.py | 4 +- chronos/model/checkpoint.py | 140 ++- chronos/model/config.py | 166 +++ chronos/model/factory.py | 94 ++ chronos/model/hf_io.py | 3 + chronos/model/hybrid_attention.py | 340 ++++-- chronos/model/kv_quant.py | 108 ++ chronos/model/lookahead_router.py | 13 +- chronos/model/model_chronos.py | 490 +++++++- chronos/model/moe_chronos.py | 141 ++- chronos/model/multimodal.py | 216 ++++ chronos/model/outputs.py | 62 + chronos/model/planner.py | 673 +++++++++++ chronos/model/precision_ops.py | 144 +++ chronos/model/recurrent.py | 208 ++++ chronos/model/temporal_loss.py | 60 +- chronos/monitor/__init__.py | 17 + chronos/monitor/evals.py | 110 ++ chronos/monitor/jsonl.py | 81 ++ chronos/monitor/process.py | 179 +++ chronos/monitor/server.py | 253 ++++ chronos/monitor/snapshot.py | 330 ++++++ chronos/monitor/static/app.css | 165 +++ chronos/monitor/static/app.js | 1036 +++++++++++++++++ chronos/monitor/static/index.html | 49 + chronos/precision.py | 195 ++++ chronos/runtime/cache_manager.py | 28 +- chronos/runtime/inference_engine.py | 424 ++++++- chronos/runtime/speculative.py | 225 ++++ chronos/serving/__init__.py | 11 + chronos/serving/adapter.py | 121 ++ chronos/serving/api_compat.py | 131 +++ chronos/serving/ollama_server.py | 353 ++++++ chronos/serving/runtime.py | 343 ++++++ chronos/tokenizer/__init__.py | 75 ++ chronos/tokenizer/__main__.py | 4 + chronos/tokenizer/evaluate.py | 40 + chronos/tokenizer/extract.py | 176 +++ chronos/tokenizer/model.py | 197 ++++ chronos/tokenizer/normalize.py | 81 ++ chronos/tokenizer/trainer.py | 207 ++++ chronos/trainer/__init__.py | 1 + chronos/trainer/chronos_trainer.py | 47 +- chronos/trainer/device_utils.py | 43 +- chronos/trainer/distill_trainer.py | 85 +- chronos/trainer/dpo_trainer.py | 21 +- chronos/trainer/grpo_trainer.py | 24 +- chronos/trainer/loss_mixin.py | 83 +- chronos/trainer/optim_utils.py | 253 ++++ chronos/trainer/orpo_trainer.py | 19 +- chronos/trainer/paged_store.py | 337 ++++++ chronos/trainer/precision_training.py | 170 +++ chronos/trainer/sft_trainer.py | 62 +- chronos/trainer/stage_utils.py | 103 +- chronos/tuning/chronos_auto_tuner.py | 2 +- chronos/verify.py | 8 +- pyproject.toml | 27 +- reports/speculative/mlx_gpu_d2_s13996.json | 834 +++++++++++++ reports/speculative/mlx_gpu_fp32_s13996.json | 834 +++++++++++++ reports/speculative/mlx_gpu_s13996.json | 736 ++++++++++++ reports/speculative/torch_mps_d2_s13996.json | 273 +++++ reports/speculative/torch_mps_s13996.json | 424 +++++++ reports/ui_event_contract_20260910.md | 17 + requirements.txt | 2 +- results/chronos_v2_benchmark_cpu.json | 560 +++++++++ results/chronos_v2_benchmark_mps.json | 561 +++++++++ setup.py | 25 +- tests/fixtures/v2_golden_logits.pt | Bin 0 -> 21047 bytes tests/test_bundle.py | 227 ++++ tests/test_interface_regressions.py | 230 ++++ tests/test_longrun.py | 64 + tests/test_mlx_speculative.py | 208 ++++ tests/test_moe_dispatch.py | 167 +++ tests/test_monitor.py | 301 +++++ tests/test_planner_ui.py | 303 +++++ tests/test_precision_chain.py | 458 ++++++++ tests/test_smoke.py | 23 +- tests/test_speculative.py | 123 ++ tests/test_tokenizer_builder.py | 103 ++ tests/test_ui_event_browser.py | 171 +++ tests/test_ui_event_contract.py | 95 ++ tests/test_ui_generic.py | 358 ++++++ tests/test_v2_architecture.py | 209 ++++ tests/test_v3_inference_options.py | 388 ++++++ tests/test_v3_kernel.py | 283 +++++ tools/bench_speculative.py | 325 ++++++ tools/benchmark_chronos_v2.py | 490 ++++++++ tools/check_mlx_stability.py | 80 ++ tools/compare_minimind_chronos_v3.py | 27 +- train_chronos.py | 36 +- train_chronos_distill.py | 21 +- train_chronos_dpo.py | 4 +- train_chronos_grpo.py | 4 +- train_chronos_orpo.py | 4 +- train_chronos_sft.py | 20 +- ui/estimator.py | 284 +++-- ui/event_contract.js | 48 + ui/event_contract.py | 127 ++ ui/gradio_compat.py | 11 + ui/i18n.py | 536 ++++++++- ui/planner_panel.py | 177 +++ ui/presets.py | 391 +++++-- ui/presets_v3.json | 699 +++++++++++ ui/tabs/autotune_tab.py | 10 +- ui/tabs/benchmark_tab.py | 23 +- ui/tabs/config_tab.py | 356 ++++-- ui/tabs/data_tab.py | 191 ++- ui/tabs/export_tab.py | 15 +- ui/tabs/inference_tab.py | 400 ++++--- ui/tabs/monitor_tab.py | 425 +++++++ ui/tabs/serve_tab.py | 202 ++++ ui/tabs/tokenizer_tab.py | 191 +++ ui/tabs/train_tab.py | 295 +++-- 171 files changed, 28640 insertions(+), 1854 deletions(-) create mode 100644 .github/ci-constraints.txt create mode 100644 .gitignore create mode 100644 chronos/bundle/__init__.py create mode 100644 chronos/bundle/builder.py create mode 100644 chronos/bundle/cli.py create mode 100644 chronos/bundle/docs.py create mode 100644 chronos/bundle/finetune.py create mode 100644 chronos/bundle/recommend.py create mode 100644 chronos/bundle/runtime.py create mode 100644 chronos/config_schema.py create mode 100644 chronos/data/corpus/__init__.py create mode 100644 chronos/data/corpus/clean.py create mode 100644 chronos/data/corpus/download.py create mode 100644 chronos/data/corpus/evalsets.py create mode 100644 chronos/data/corpus/pipeline.py create mode 100644 chronos/data/corpus/presets.py create mode 100644 chronos/data/corpus/readers.py create mode 100644 chronos/data/corpus/shards.py create mode 100644 chronos/data/corpus/sources.py create mode 100644 chronos/migrate.py create mode 100644 chronos/mlx/kv_quant.py create mode 100644 chronos/mlx/moe_kernels.py create mode 100644 chronos/mlx/recurrent.py create mode 100644 chronos/mlx/speculative.py create mode 100644 chronos/mlx/training/inspect_checkpoint.py create mode 100644 chronos/mlx/training/longrun.py create mode 100644 chronos/mlx/training/optim.py create mode 100644 chronos/mlx/training/recipes/enzh_50m_72h.json create mode 100644 chronos/mlx/training/telemetry.py create mode 100644 chronos/model/factory.py create mode 100644 chronos/model/kv_quant.py create mode 100644 chronos/model/multimodal.py create mode 100644 chronos/model/outputs.py create mode 100644 chronos/model/planner.py create mode 100644 chronos/model/precision_ops.py create mode 100644 chronos/model/recurrent.py create mode 100644 chronos/monitor/__init__.py create mode 100644 chronos/monitor/evals.py create mode 100644 chronos/monitor/jsonl.py create mode 100644 chronos/monitor/process.py create mode 100644 chronos/monitor/server.py create mode 100644 chronos/monitor/snapshot.py create mode 100644 chronos/monitor/static/app.css create mode 100644 chronos/monitor/static/app.js create mode 100644 chronos/monitor/static/index.html create mode 100644 chronos/precision.py create mode 100644 chronos/runtime/speculative.py create mode 100644 chronos/serving/adapter.py create mode 100644 chronos/serving/api_compat.py create mode 100644 chronos/serving/ollama_server.py create mode 100644 chronos/serving/runtime.py create mode 100644 chronos/tokenizer/__init__.py create mode 100644 chronos/tokenizer/__main__.py create mode 100644 chronos/tokenizer/evaluate.py create mode 100644 chronos/tokenizer/extract.py create mode 100644 chronos/tokenizer/model.py create mode 100644 chronos/tokenizer/normalize.py create mode 100644 chronos/tokenizer/trainer.py create mode 100644 chronos/trainer/paged_store.py create mode 100644 chronos/trainer/precision_training.py create mode 100644 reports/speculative/mlx_gpu_d2_s13996.json create mode 100644 reports/speculative/mlx_gpu_fp32_s13996.json create mode 100644 reports/speculative/mlx_gpu_s13996.json create mode 100644 reports/speculative/torch_mps_d2_s13996.json create mode 100644 reports/speculative/torch_mps_s13996.json create mode 100644 reports/ui_event_contract_20260910.md create mode 100644 results/chronos_v2_benchmark_cpu.json create mode 100644 results/chronos_v2_benchmark_mps.json create mode 100644 tests/fixtures/v2_golden_logits.pt create mode 100644 tests/test_bundle.py create mode 100644 tests/test_interface_regressions.py create mode 100644 tests/test_longrun.py create mode 100644 tests/test_mlx_speculative.py create mode 100644 tests/test_moe_dispatch.py create mode 100644 tests/test_monitor.py create mode 100644 tests/test_planner_ui.py create mode 100644 tests/test_precision_chain.py create mode 100644 tests/test_speculative.py create mode 100644 tests/test_tokenizer_builder.py create mode 100644 tests/test_ui_event_browser.py create mode 100644 tests/test_ui_event_contract.py create mode 100644 tests/test_ui_generic.py create mode 100644 tests/test_v2_architecture.py create mode 100644 tests/test_v3_inference_options.py create mode 100644 tests/test_v3_kernel.py create mode 100644 tools/bench_speculative.py create mode 100644 tools/benchmark_chronos_v2.py create mode 100644 tools/check_mlx_stability.py create mode 100644 ui/event_contract.js create mode 100644 ui/event_contract.py create mode 100644 ui/planner_panel.py create mode 100644 ui/presets_v3.json create mode 100644 ui/tabs/monitor_tab.py create mode 100644 ui/tabs/serve_tab.py create mode 100644 ui/tabs/tokenizer_tab.py diff --git a/.github/ci-constraints.txt b/.github/ci-constraints.txt new file mode 100644 index 0000000..423f029 --- /dev/null +++ b/.github/ci-constraints.txt @@ -0,0 +1,21 @@ +# Versions CI installs (resolved and tested together on Python 3.10 and 3.11). +# transformers >= 5 disables its PyTorch backend below torch 2.6, so torch and +# transformers move together. Newer transformers / numpy / flake8 dropped 3.10. +torch==2.10.0 +transformers==5.22.1; python_version >= "3.11" +transformers==5.20.2; python_version < "3.11" +numpy==2.4.6; python_version >= "3.11" +numpy==2.2.6; python_version < "3.11" +flake8==7.5.0; python_version >= "3.11" +flake8==7.4.1; python_version < "3.11" +datasets==5.1.0 +tokenizers==0.23.3 +safetensors==0.8.0 +optuna==5.0.0 +tqdm==4.70.1 +psutil==7.2.2 +fastapi==0.115.9 +uvicorn==0.34.3 +httpx==0.28.1 +pytest==9.2.0 +ruff==0.17.4 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7fc4b57..e450ff2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,10 +24,10 @@ jobs: - name: Install dependencies run: | pip install --upgrade pip - pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cpu - pip install transformers datasets safetensors optuna psutil numpy tqdm pytest ruff - # Install minimind-master as editable dependency - pip install -e ./minimind-master 2>/dev/null || true + # CPU torch wheel; versions come from .github/ci-constraints.txt + pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cpu + pip install -c .github/ci-constraints.txt transformers datasets tokenizers safetensors optuna \ + psutil numpy tqdm pytest ruff fastapi uvicorn httpx - name: Lint run: | @@ -36,21 +36,3 @@ jobs: - name: Unit tests run: | python -m pytest tests/ -v --tb=short - - publish: - if: startsWith(github.ref, 'refs/tags/v') - needs: test - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 - with: - python-version: "3.11" - - name: Build - run: | - pip install build - python -m build - - name: Publish to PyPI - uses: pypa/gh-action-pypi-publish@release/v1 - with: - password: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml index 1168bd9..6272e59 100644 --- a/.github/workflows/python-app.yml +++ b/.github/workflows/python-app.yml @@ -26,8 +26,8 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install flake8 pytest - if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cpu + pip install -c .github/ci-constraints.txt flake8 pytest psutil -r requirements.txt - name: Lint with flake8 run: | # stop the build if there are Python syntax errors or undefined names diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..484c180 --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +# Python / tooling +__pycache__/ +*.py[cod] +*.egg-info/ +build/ +dist/ +.pytest_cache/ +.ruff_cache/ +.venv/ +.DS_Store + +# Training runs, checkpoints, expert caches, bundles (weights are research-only) +runs/ +out/ +expert_cache/ +expert_cache_mlx/ +*.pth +*.pkl +*.ctsr +*.safetensors + +# Local data and archives +Translation Training Data/ +.archive/ diff --git a/CHANGELOG.md b/CHANGELOG.md index cf0412b..44eed6f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,90 @@ # Changelog +## [0.3.0] — 2026-11-06 + +v3 kernel: recurrent-depth MoE. Quality results for the 50M en→zh demo are not +included yet; they will be added after the 72-hour run and its ablations. + +### Added +- Web UI Inference tab: self-speculative switch, draft depth and γ. All four runner paths use it (Torch / MLX, full DRAM / lazy experts). The stats table shows acceptance, tokens per verify, and verify / draft forwards. Models without a shared KV stream (`per_iteration`, depth 1) decode plainly and say so. +- `ChronosInferenceEngine.generate_speculative`: self-speculative decoding with lazy experts. Every draft and verify forward runs under the current availability masks, and the lookahead prediction is prefetched before it. `chronos.runtime.speculative.speculative_supported(cfg)` tells callers whether a model can draft. +- `tools/bench_speculative.py`: plain vs speculative decoding on the real runtimes. Reports wall time, tokens/s, acceptance, identical-output count and SSD loads, with the environment and checkpoint hash. +- Recurrent-depth execution: prelude → shared recurrent groups looped `recurrent_depth` times with a contractive update → coda. +- `recurrent_attention_mode`: `shared_first` (default; attention every iteration, K/V from iteration 0), `per_iteration` (ablation), `first_only` (v2 semantics). +- Tensor-only ACT early exit, exact and batch-invariant; per-token decode exit without KV gaps. +- Recurrent-depth self-speculative decoding (`chronos/runtime/speculative.py`): greedy-exact, rejection sampling, cache rollback. +- MLX self-speculative decode loop (`chronos/mlx/speculative.py`), Torch parity: same invariant, rejection sampling with `mx.random` keys, cache truncation for MLA, SWA and int8 entries. `ChronosMLXInferenceEngine.generate(speculative=True, draft_depth=, gamma=)` runs it with the lazy-expert masks; experts each draft routes to are unioned with the lookahead prediction and prefetched before the verify forward. `SlidingWindowAttentionMLX.cache_slack` keeps `window + γ_max + 1` entries for rollback. +- `chronos bundle build|pack|check|probe|translate|eval|prepare-data|finetune|serve`: a self-contained model bundle (weights + sidecar, tokenizer, FLORES+ dev/devtest, vendored minimind, wheel, source recipe, recommended fine-tune tiers, launcher, setup script, EN/ZH README, NOTICE, SHA-256 file table). `bundle translate` takes `--backend torch|mlx` and `--speculative`. +- `CHRONOS_MINIMIND_DIR` points Chronos at a local minimind copy (bundles set it to `third_party/minimind`), so no git clone or network access is needed. +- Stage scripts take `--tokenizer_path`; `load_stage_tokenizer` loads the checkpoint's own tokenizer and refuses one that does not fit `vocab_size`. +- `chronos longrun` spec fields `init_checkpoint` (start from existing weights, checked against `model_config`) and `export_prefix`. +- Sorted grouped MoE dispatch: `mx.gather_mm` on MLX (requires mlx >= 0.32.4), `argsort` + `index_add_` on Torch; per-expert reference paths kept for parity tests. +- Single-dtype precision chain for CPU / CUDA / MPS / MLX with FP32 only in RMSNorm statistics, SDPA softmax, RoPE tables, router softmax, LM-head logits and optimizer masters; bit-exact MLX↔Torch bf16 bridge. +- `chronos serve`: Ollama/OpenAI protocol server on the native runtime (lazy experts, speculative decoding); optional `/monitor/` mount. +- `chronos tokenizer`: self-built byte-fallback BPE with language-balanced merges and HF-compatible `tokenizer.json`. +- `chronos corpus download|build|shards`: pinned downloads, cleaning, global dedup, eval-set leak removal, uint16 shards. +- `chronos longrun` (alias `run`): time-budgeted resumable MLX training with NaN rollback, validation and translation probes, and run telemetry (`status.json`, `runtime.json`, `parameters.jsonl`). +- `chronos monitor`: read-only training monitor page + JSON API (loopback by default; key login, session cookie, rate limiting off loopback); Training Monitor tab in the Web UI. +- `chronos.model.factory.fit_config`: parameter-budget configuration search. +- Optional extras `serve` and `translate` with pinned versions. +- Inference and memory options, all off by default: `kv_cache_dtype="int8"` (MLA `c_kv` and SWA k/v, 32-element groups), `mla_decode="absorbed"|"auto"` with optional `mla_cache_rope`, `router_scope="shared"`, the `draft_exit_kl_coef` / `draft_exit_depth` training term, and hard expert budgets `max_resident_experts` / `ram_budget_mb` with prediction-protected LRU and per-virtual-layer prefetch. Torch and MLX both implement them. +- Adaptive training planner (`chronos.model.planner`, Config tab → "Auto plan"). Inputs: target size, MoE on/off, experts, and either active experts or active parameters, plus a memory-saving level and datasets. Outputs: a v3 architecture, scaling-law LR/batch, a memory fit, and independent pretrain / SFT / DPO / ORPO / GRPO / distill values. Dataset profiling supplies vocabulary, sequence length, token count and tokenizer path. +- Per-stage overrides (`config["stage_overrides"]`); the Train tab merges the selected stage's values. +- Five v3 presets generated by the planner (`ui/presets_v3.json`, `python -m ui.presets`) replace the old presets; old names are aliases. +- Tokenizer and Serve tabs in the Web UI; the Train tab honours `tokenizer_path`. + +### Changed +- Routing defaults to top-2 plus a shared expert (top-1 with `norm_topk_prob=True` gave the gate no CE gradient). +- `--dtype auto` resolves to fp32 on MPS (native bf16 measured slower there). +- `lookahead_source` now comes from the prelude output. +- Repetition penalty divides positive logits and multiplies negative ones. +- `gradio` is constrained to `>=5.0,<6`; `tokenizers` is a declared dependency. +- `setup.py` is a shim; metadata lives only in `pyproject.toml`. +- `mla_decode` defaults to `expanded`. With `auto`, a bf16 prefill took the expanded path and the following decode took the absorbed path. Prefill+decode then differed from the full forward pass by 0.06 in max logit (measured). `absorbed` and `auto` are opt-in. +- Minimum torch is 2.6: transformers 5 disables its PyTorch backend below that version, which broke imports. +- CI installs torch 2.10.0 (CPU) and pins the rest in `.github/ci-constraints.txt`, tested on Python 3.10 and 3.11. The dead token-based publish job was removed; releases publish through `pypi.yml` (trusted publishing). +- The Web UI no longer assumes a task. + - The monitor (page and Gradio tab) discovers evaluation metrics from whatever numeric keys a run logs on validation / final events. Keys are grouped by metric name (`probe_en_zh_bleu`, `qa_accuracy`, `pass1` …) and charted per family, with an improvement direction for known metrics. + - Counts and durations are listed as context and not charted. + - Samples show input / reference / output for any schema (`src`/`prompt`/`question`, `ref`/`answer`, `hyp`/`response`/`prediction`). Other fields appear as tags. + - Offline inspection reports are flattened to every numeric leaf. +- Data tab: scans a file, directory or glob of jsonl / json / parquet / csv / tsv / txt. Reports record shapes (chat, preference, instruction, prompt–response, parallel pair, plain text), field presence and types, scripts, duplicates and per-record character lengths over all text fields. + +### Fixed +- The Inference tab loaded the minimind tokenizer for every checkpoint. It now uses the checkpoint's own tokenizer: the config `tokenizer_path`, the bundle's `tokenizer/` folder, or the sidecar path. Minimind is the last resort. +- MLX inference ran in training mode when the model was converted from a freshly built torch model (the Inference tab did this). MoE then masked experts at random (`fallback_mask_prob`), so repeated greedy runs disagreed. `ChronosMLXInferenceEngine` now puts its model in eval mode. +- The Torch lazy engine failed with "Inference tensors cannot be saved for backward" when a plain request followed a speculative one, because experts paged in under `torch.inference_mode` were reused. The decode paths now use `torch.no_grad`. +- The v2 every-iteration attention switch rewrote a single KV slot per iteration (prefill+decode vs full forward gap 0.94). +- CLI registered `serve`, `tokenizer` and `corpus` twice. +- Monitor CPU time on macOS: `rusage` times are Mach ticks, now scaled by `mach_timebase_info`. +- `apply_inference_overrides` validated overrides against a config that used class-default KV heads (4). Every serve/inference override on models with other head counts (e.g. the 50M demo, 6/2) raised. +- The UI estimator ignored `recurrent_groups` and `recurrent_attention_mode`; its counts now match `count_parameters`. +- Planner/dataset profiling counted GRPO prompt files (empty assistant turn) as SFT. +- Stage scripts loaded the minimind tokenizer (6,400 entries) for checkpoints trained with another vocabulary. +- A forward with `recurrent_depth` below `min_recurrent_depth` (a speculative draft) gave the last state zero ACT weight on MLX, so the coda saw zeros (max logit gap to Torch 16.6 on the 50M checkpoint, measured; draft acceptance 2 % instead of 45 %). Torch special-cased depth 1 but had the same fault for 2 ≤ depth < `min_recurrent_depth`. The minimum now clamps to the executed depth; training depths are unaffected. +- MLX ≥ 0.32 binds a lazy `mx.load` to the loading thread's stream; experts prefetched by the background thread raised "There is no Stream(cpu, N) in current thread" when promoted on the decode thread. Pages are now evaluated in the thread that loads them. + +### Known gaps +- Self-speculative decoding is exact but **not faster** for the 50M demo, on CPU, GPU or lazy experts. All figures are measured on the pretrain step-13,996 checkpoint, greedy, batch 1, FLORES+ dev en→zh, draft depth 1, γ 4. The 72-hour run was paused (SIGSTOP) for the GPU runs. Reports are in `reports/speculative/`. + + | Runtime | Plain tok/s | Speculative tok/s | Ratio | Acceptance | Identical | + |---|---|---|---|---|---| + | MLX GPU bf16, all experts in DRAM | 112.0 | 107.4 | 0.96× | 42 % | 22/24 | + | MLX GPU bf16, lazy experts (12 of 216 pages) | 43.7 | 32.0 | 0.73× | 42 % | 22/24 | + | Torch MPS fp32, all experts | 32.1 | 25.9 | 0.81× | 40 % | 12/12 | + | Torch MPS fp32, lazy experts, sync misses | 17.0 | 14.7 | 0.86× | 40 % | 12/12 | + + Draft depth 2 raised acceptance to 55–59 % but did not change the result (0.72–0.92×). + + At this size one full-depth step costs about as much as a batched verify pass, so drafting at depth 1 of 4 does not save enough. + + With lazy experts, drafts and verify passes route to different experts. Speculative decoding then read 1.4–1.6× more expert pages from SSD (12,092 vs 8,372 for MLX). + + bf16 on the GPU flips near-tie tokens (the 2/24 and 3/24 mismatches). In GPU fp32 one sentence differed, at a top-1/top-2 logit margin of 0.0006. CPU fp32 matches exactly. The switch therefore stays off by default. +- At a 1280 px window, Gradio moves the Auto-Tune and IO Monitor tabs into its "…" overflow menu. +- Planner peak-memory figures are estimates; for the 50M recipe the estimate (16.4 GB) is about twice the measured MLX peak (8.72 GB). +- Gradio 5.35 does not render the LinePlot legend; colour keys are written into plot titles instead. + ## [0.1.0] — 2026-04-20 ### Added diff --git a/MANIFEST.in b/MANIFEST.in index ad8909c..5d2942f 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,9 +1,13 @@ -include chronos/**/*.py -include train_chronos.py -include README.md -include LICENSE -include requirements.txt -exclude **/__pycache__ -exclude **/*.pyc -exclude expert_cache/ -exclude out/ +include README.md README_zh.md LICENSE THIRD_PARTY_NOTICES.md CHANGELOG.md requirements.txt +recursive-include chronos/monitor/static *.html *.js *.css +include ui/event_contract.js +global-exclude __pycache__ *.py[cod] .DS_Store +prune runs +prune out +prune expert_cache +prune expert_cache_mlx +prune results +prune reports +prune .archive +prune Translation?Training?Data +prune chronos/mlx/training/recipes diff --git a/README.md b/README.md index 6fbe058..0e443ad 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,81 @@ Under a simulated 30 ms SSD delay, the new path keeps **35 ms+ per token** of pi --- +## v3 kernel: recurrent-depth MoE (0.3.0) + +v3 reuses a few physical layers many times. A prelude runs once, each shared recurrent group loops `recurrent_depth` times with a contractive update `h ← h + dt·(block(h + e_loop) − h)`, and a coda runs once. The 50M demo has 6 physical layers that execute as 18 virtual layers (prelude 1, 4 recurrent layers × depth 4, coda 1). + +```mermaid +flowchart LR + X["tokens"] --> P["Prelude
runs once"] + P --> G["Recurrent group
MLA / SWA + top-2 MoE"] + G -->|"loop d = 1..D
contractive step, ACT halt"| G + G --> C["Coda
runs once"] + C --> H["LM head"] + P -. "lookahead source" .-> LR["LookaheadRouter
expert prefetch"] +``` + +**What changed from v2** + +| | v2 | v3 | +|---|---|---| +| Attention inside the loop | Iteration 0 only (`first_only`) | Every iteration | +| KV cache per recurrent layer | 1 slot; the every-iteration attention switch rewrote it each iteration (measured prefill+decode vs full-forward gap 0.94) | `shared_first`: Q recomputed each iteration, K/V from iteration 0, so KV slots = physical layers | +| Routing | top-1 (gate received no CE gradient) | top-2 + 1 shared expert | +| MoE dispatch | Python loop over experts with `.item()` syncs | Sorted grouped dispatch: `gather_mm` on MLX, `argsort` + `index_add_` on Torch | +| Early exit | Host sync every step | Tensor-only ACT, exact and batch-invariant; decode can exit per token without KV gaps | +| Decoding | Plain autoregressive | Recurrent-depth self-speculative decoding | + +`per_iteration` (one KV slot per virtual layer) is kept as an ablation. Old v2 checkpoints load as `first_only`; their logits match a stored golden tensor within 1e-5. + +**Self-speculative decoding (no extra parameters).** The same model drafts γ tokens at a shallow depth and verifies them with one full-depth batched forward pass. Greedy output is token-for-token identical to full-depth greedy. Sampling uses rejection sampling, which preserves the full-depth distribution. After a rejection, the cache is truncated: MLA slices `c_kv`, and the sliding-window cache keeps `window + γ_max` entries. Torch (`chronos/runtime/speculative.py`) and MLX (`chronos/mlx/speculative.py`, `ChronosMLXInferenceEngine.generate(speculative=True)`) share the algorithm. Tests cover greedy identity, the rejection distribution and rollback against a fresh prefill on both backends. The Web UI Inference tab has a switch for it, with draft depth and γ. + +For the 50M demo it is exact but not faster on any runtime we measured (pretrain step-13,996 checkpoint, greedy, batch 1, FLORES+ dev en→zh, draft depth 1, γ 4; reports in `reports/speculative/`): + +| Runtime | Plain tok/s | Speculative tok/s | Ratio | Acceptance | +|---|---|---|---|---| +| MLX GPU bf16, all experts in DRAM | 112.0 | 107.4 | 0.96× | 42 % | +| MLX GPU bf16, lazy experts (12 of 216 pages resident) | 43.7 | 32.0 | 0.73× | 42 % | +| Torch MPS fp32, all experts | 32.1 | 25.9 | 0.81× | 40 % | +| Torch MPS fp32, lazy experts | 17.0 | 14.7 | 0.86× | 40 % | + +At this size a full-depth step costs about as much as a batched verify pass. With lazy experts, drafts and verify passes also route to different experts, so speculative decoding read 1.4–1.6× more expert pages from SSD. Greedy outputs were identical in CPU fp32 and Torch MPS fp32. In GPU bf16, near-tie tokens flipped in 2 of 24 sentences. The switch therefore stays off by default. + +**Precision chain.** One compute dtype runs end to end. FP32 is used only where it pays: RMSNorm statistics, SDPA softmax, RoPE tables, router softmax, LM-head logits and optimizer master weights. The MLX↔Torch bf16 bridge is bit-exact. `auto` picks bf16 on MLX GPU and CUDA, and fp32 on CPU and MPS (MPS bf16 measured slower; see below). + +**Tooling added in v3** + +| Command | What it does | +|---|---| +| `chronos tokenizer` | Self-built byte-fallback BPE with script-aware pre-tokenization and language-balanced merges; writes an HF-compatible `tokenizer.json` + chat template | +| `chronos corpus download / build / shards` | Pinned downloads, cleaning, global dedup, eval-set leak removal, uint16 shards | +| `chronos longrun --spec recipe.json` (alias `run`) | Time-budgeted MLX run: warmup + cosine per phase wall-clock budget, NaN rollback, bit-exact resume, validation and translation probes | +| `chronos serve` | Ollama/OpenAI protocol server on the native runtime (see below) | +| `chronos monitor` | Read-only training monitor page + JSON API (see below) | + +`chronos.model.factory.fit_config` searches for a configuration that matches total and active parameter targets. + +**Status of the 50M en→zh demo (in progress).** +- Model: 49.69M total / 12.09M active parameters, 36 experts, MLA latent 128, SWA window 128, 16k self-built vocabulary. +- Run: a 72-hour MLX bf16 run (pretrain, then SFT) started on 2026-11-05. +- Interim measurement: valid CE 2.477 at pretrain step 13,996 of about 24.7k planned (it was 3.012 at step 1,656). That checkpoint is packaged as `chronos-enzh-50m-pretrain-s13996` (see Model bundle). +- BLEU/chrF on FLORES+ and WMT22, and the ablations (attention mode, depth, lazy-expert budget, speculative γ), will be published after the run finishes. No quality numbers are claimed before then. + +**Inference and memory options (all off by default).** Each option is a config field and a `chronos serve` / Inference-tab override. Tests compare each one with its reference path on Torch and MLX. + +| Field | Values (default first) | Effect | +|---|---|---| +| `kv_cache_dtype` | `auto`, `int8` | `int8` stores MLA `c_kv` and SWA k/v per token in groups of 32 with a scale; rollback and speculative decoding work on it | +| `mla_decode` | `expanded`, `absorbed`, `auto` | `absorbed` scores against `c_kv` without expanding K/V. `auto` absorbs only short cached chunks; it is opt-in because a bf16 prefill (expanded) followed by absorbed decode drifted from the full forward pass (max logit gap 0.06, measured) | +| `mla_cache_rope` | `false`, `true` | Also caches the roped key, trading memory for decode speed | +| `router_scope` | `per_iteration`, `shared` | `shared` reuses iteration 0's experts for the rest of the loop, so fewer experts are resident per token | +| `draft_exit_kl_coef`, `draft_exit_depth` | `0`, `1` | Training term: KL from the shallow exit to the detached full-depth logits, which raises speculative acceptance | +| `max_resident_experts`, `ram_budget_mb` | unset | Hard caps on the fast tier and the RAM tier. LRU eviction never evicts experts the lookahead router predicts. Prefetch covers every loop iteration of a virtual layer | + +Web UI: the Tokenizer tab (extract → train → evaluate) and the Serve tab (start/stop `chronos serve`; the API key goes through the environment, never argv). + +--- + ## Three-tier storage (M1: cluster-aware safetensors) ```mermaid @@ -189,7 +264,7 @@ flowchart LR Training dtype and resource policy: -- `--dtype auto` is the default. MPS/MLX resolve to BF16-first for training stability, CUDA/XPU resolve to FP16, and CPU resolves to FP32 unless `--dtype float16` or `--dtype bfloat16` is set explicitly. +- `--dtype auto` is the default. One compute dtype covers weights, activations, KV cache and expert pages; FP32 is kept only inside fused kernels (SDPA softmax, RMSNorm statistics), for RoPE tables, router softmax, loss logits and the optimizer's FP32 masters/moments. `auto` resolves to BF16 on MLX GPU, BF16 on CUDA (FP16 if unsupported, with loss scaling), and FP32 on CPU and MPS. On MPS, native BF16 measured slower than FP32 (Chronos 50M on an M5 Max with torch 2.10: 0.749 vs 0.619 s per training step, 49.1 vs 37.3 ms per decoded token). Pass `--dtype bfloat16` explicitly to trade speed for half the weight and KV memory. - CPU training configures PyTorch to use physical cores by default. Override with `--cpu_threads` or `--cpu_budget_percent`. - On macOS, MPS/MLX training forces DataLoader workers to `0` by default to avoid Metal command-buffer crashes from multiprocessing. CPU/CUDA still use worker processes; advanced users can override the guard with `CHRONOS_ALLOW_METAL_DATALOADER_WORKERS=1`. - Native MLX training pushes UI logs, scalar readouts, and chart points every `log_interval` steps, and Web UI Stop is checked at each batch boundary. @@ -200,7 +275,7 @@ The six-stage MLX trainer mirrors the PyTorch loss stack in `chronos.mlx.*`: masked CE, DPO/ORPO/GRPO/distillation losses, load balance, temporal locality, lookahead soft-target supervision, lookahead top-k hit loss, and router KL anchor. Numerically sensitive pieces run in FP32 even when model weights use -BF16/FP16, which is why `auto` prefers BF16 on MLX/MPS: FP16 has too little +BF16/FP16, which is why `auto` prefers BF16 over FP16 on MLX GPU: FP16 has too little exponent range for router softmax, CE, and Adam moments on small unstable training runs. @@ -371,6 +446,13 @@ pip install -e ".[dev]" pip install "Project_Chronos[mlx]" ``` +**Optional extras** (versions pinned): + +```bash +pip install "Project_Chronos[serve]" # chronos serve / chronos monitor (fastapi, uvicorn) +pip install "Project_Chronos[translate]" # corpus detokenization + BLEU/chrF probes (sacremoses, sacrebleu) +``` + **vLLM serving (optional, Linux + CUDA only):** ```bash @@ -378,16 +460,17 @@ pip install vllm ``` > **minimind dependency**: Project Chronos uses [minimind](https://github.com/jingyaogong/minimind) as its MoE kernel. -> If it is not found locally, Chronos clones it automatically into `~/.cache/chronos/minimind-master/` on first import. +> If it is not found locally, Chronos clones it automatically into `~/.cache/chronos/minimind-master/` on first import (this needs `git` and network access). +> Set `CHRONOS_MINIMIND_DIR=/path/to/minimind` to use a local copy instead; model bundles ship one under `third_party/minimind` and their launcher sets the variable. > minimind is licensed under **Apache-2.0**. See [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md) for attribution details. -**Requirements**: Python 3.10+, PyTorch 2.4+ +**Requirements**: Python 3.10+, PyTorch 2.6+ --- ## Quick start -### Web UI (M6: 8 tabs, 4 languages) +### Web UI (12 tabs, 4 languages) ```bash chronos-ui @@ -397,17 +480,127 @@ python chronos_app.py Tabs included: -- `Config` with a live parameter / memory estimator merged in from the old Designer -- `Train` with its own `data_path` +- `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. + +### Presets and the auto planner + +The Config tab ships five v3 presets. Each one is generated by the planner and stored in `ui/presets_v3.json`; regenerate the file with `python -m ui.presets`. Parameter counts are measured with a meta-device model; peak memory is an estimate. + +| Preset | Total / active params (measured) | Shape | Memory level | Est. training peak | +|---|---|---|---|---| +| Chronos-v3 Tiny (smoke / CI) | 1.51M / 1.29M | E4 top-2, vocab 6400, seq 128, depth 2 | none | 2.13 GB | +| Chronos-v3 50M (en→zh recipe) | 49.69M / 12.09M | h384, L6, E36 top-2, vocab 16384, seq 512 | none | 16.4 GB (the live run measures 8.72 GB peak) | +| Chronos-v3 200M (32 GB) — default | 200.72M / 28.20M | h576, E32 top-2 | balanced | 10.14 GB | +| Chronos-v3 500M (32 GB) | 499.53M / 61.59M | E48 top-3 | aggressive | 9.26 GB | +| Chronos-v3 1B (64 GB) | 1002.17M / 99.74M | E64 top-4, seq 2048 | aggressive | 18.92 GB | + +The old preset names (Recommended-CN, Tiny/Small/Medium/Large) still load; they map to the closest v3 tier. + +**Auto planner** (Config tab → "Auto plan"). Inputs: +- target total parameters, and whether to use MoE; +- the expert count, plus either the number of active experts or an active-parameter target; +- a memory-saving level: `none`, `balanced`, `aggressive` or `extreme`; +- optional dataset paths, RAM, token budget, vocabulary size and sequence length. + +`chronos.model.planner.plan_training` then does the following: +1. **Profiles the datasets.** It reads the shard manifest or samples the files, classifies records as pretrain, SFT, preference or RL, and reads the tokenizer report. The vocabulary size, sequence length (p95 / cap), token count and tokenizer path come from the data when available. +2. **Searches the architecture.** It looks for the hidden size, layer count, expert width and GQA heads whose meta-device count is closest to the target. +3. **Sets LR and batch size** from the DeepSeek LLM compute-optimal fits (η = 0.3118·C^−0.125, B = 0.2920·C^0.3271), capped for wide models. +4. **Fits memory.** It picks micro-batch, accumulation, int8 KV, router scope, activation checkpointing, optimizer paging and an expert RAM budget for the memory level. +5. **Writes per-stage values.** Pretrain, SFT, DPO, ORPO, GRPO and distill each get their own learning rate, batch, epochs, save interval and stage-specific terms (β, temperature, α). + +The stage table is editable; the Train tab merges the selected stage's row over the shared values, and the pretrain row stays in sync with the shared widgets. "Save plan" writes the config, stage values and report to JSON. Every number in the report is labelled measured or estimated. + Built-in i18n: `zh-Hans`, `zh-Hant`, `en`, `ja` +### Training monitor (console tab + standalone page/API) + +A `chronos longrun` run directory holds `run_spec.json`, `metrics.jsonl`, +`status.json` (heartbeat), `runtime.json` (precision policy, device, versions, +code hash), `parameters.jsonl` (per-tensor weight/grad/update statistics and +per-virtual-layer routing), `samples.jsonl`, `inspections/` and checkpoints. +The console's **Training Monitor** tab and a standalone page both read these +files and never touch the training device. Runs started before telemetry +existed still show progress; their status is derived from `metrics.jsonl` and +the process table. + +```bash +chronos monitor --runs-dir runs # http://127.0.0.1:7870/ +chronos serve --models-dir exports --monitor-runs-dir runs # same page at /monitor/ +``` + +JSON API (read-only): `GET /api/state?run=&history=1`, `/api/metrics?run=&phase=&since=`, +`/api/parameters`, `/api/tensors`, `/api/tensor?name=`, `/api/log?bytes=`, `/api/runs`. + +Security: the default bind is 127.0.0.1 with no authentication. Any other +address requires `--allow-remote` and `--api-key` (or `CHRONOS_MONITOR_KEY`). +Browsers sign in once and get an HttpOnly, SameSite=Strict session cookie; +scripts send `Authorization: Bearer `. Failed logins are rate-limited. +Off loopback, add `--ssl-certfile/--ssl-keyfile` or a TLS reverse proxy, +otherwise the key crosses the network in clear text. + +### Ollama / OpenAI protocol server + +`chronos serve` runs the native Chronos runtime behind the Ollama and OpenAI HTTP APIs. The kernel is not converted to another architecture, so recurrent depth, SSD-lazy experts with lookahead prefetch, compressed KV and self-speculative decoding stay active. + +```bash +pip install "Project_Chronos[serve]" +chronos serve --models-dir ./exports --port 11435 # 11434 would collide with a running Ollama +OLLAMA_HOST=127.0.0.1:11435 ollama list +OLLAMA_HOST=127.0.0.1:11435 ollama run "Hello" +``` + +- Endpoints: `/api/version`, `/api/tags`, `/api/show`, `/api/ps`, `/api/pull`, `/api/generate` and `/api/chat` (streaming NDJSON), `/v1/models`, `/v1/chat/completions` (SSE). +- Models are `.pth` + `.config.json` pairs found under `--models-dir`. +- Chronos options go in the request `options`: `chronos_speculative`, `chronos_draft_depth`, `chronos_gamma`, `chronos_miss_policy`. +- Stock `ollama` / llama.cpp cannot execute Chronos GGUF files (MLA+SWA hybrid layers and the contractive recurrence have no stock equivalent); use this server instead. +- Security: binds 127.0.0.1 without authentication by default. Other addresses require `--allow-remote` and `--api-key`; clients then send `Authorization: Bearer `. + +### Model bundle (ready-to-run pretrained model) + +`chronos bundle build` packages one exported checkpoint of a `chronos longrun` run into a folder that works right after unpacking. Building only reads the run directory. + +```bash +chronos bundle build --run-dir runs/enzh_50m_v1 --which best --wheel dist/project_chronos-0.3.0-py3-none-any.whl --pack +``` + +| Bundle path | Content | +|---|---| +| `model/` | weights + `.config.json` sidecar, tokenizer (`model/tokenizer/`) | +| `eval/` | FLORES+ dev/devtest (CC BY-SA 4.0) for `bundle eval` / `probe` | +| `data/sft_valid/` | the run's SFT validation shards | +| `third_party/minimind/` | vendored minimind (Apache-2.0), so no git clone is needed | +| `wheels/`, `bin/setup.sh`, `bin/chronos` | offline install into `.venv`, launcher that sets `CHRONOS_BUNDLE` / `CHRONOS_MINIMIND_DIR` | +| `recipes/source_run_spec.json`, `recommended.json` | the recipe that produced the weights; fine-tune tiers `smoke` / `quick` / `standard` / `full` | +| `bundle.json` | provenance (step, tokens, valid CE), measured quality probe, SHA-256 of every file | + +On the recipient side: + +```bash +bin/setup.sh # venv + wheel + `bundle check` +bin/chronos bundle translate "The weather is very nice today." +bin/chronos bundle translate --backend mlx --speculative --file examples/translate_en.txt +bin/chronos bundle eval --split devtest -n 200 +bin/chronos bundle prepare-data my_pairs.jsonl --out work/data # FLORES+ is removed from your data +bin/chronos bundle finetune --data work/data --run-dir work/ft --tier quick --mode sft +bin/chronos bundle serve -- --port 11435 +``` + +`--speculative` needs `--repetition-penalty 1.0` (its default), because the penalty changes the target distribution at every position. `finetune` writes a `chronos longrun` spec with `init_checkpoint` set to the bundle weights and refuses a run directory that already holds a run. Bundles built from data with unknown licences (WMT19, OPUS-100) say in `NOTICE.md` that the weights are for research only. + ### Deployment export ```bash @@ -540,7 +733,10 @@ Project_Chronos/ │ ├── deps.py # Auto-download minimind if missing │ ├── __init__.py # AutoConfig / AutoModelForCausalLM registration │ ├── model/ -│ │ ├── config.py # ChronosConfig +│ │ ├── config.py # ChronosConfig (v3 fields, v2 migration) +│ │ ├── recurrent.py # Loop/group embeddings, contractive step, tensor-only ACT +│ │ ├── precision_ops.py # RMSNorm/SDPA/RoPE with one compute dtype +│ │ ├── factory.py # Parameter-budget config search │ │ ├── hybrid_attention.py # MLAAttention + SlidingWindowAttention │ │ ├── lookahead_router.py # Per-token lookahead predictor │ │ ├── moe_chronos.py # ChronosMOEFeedForward + shared experts + soft gating @@ -559,7 +755,11 @@ Project_Chronos/ │ │ └── prefill_scheduler.py # Prefill-time expert preloader │ ├── mlx/ │ │ ├── attention.py / moe.py / model.py / expert_store.py / inference.py +│ │ ├── moe_kernels.py # Sorted gather_mm dispatch (mlx >= 0.32.4) +│ │ ├── speculative.py # MLX self-speculative decoding (Torch parity) +│ │ └── training/ # longrun.py, telemetry.py, recipes/ │ ├── runtime/ +│ │ ├── speculative.py # Recurrent-depth self-speculative decoding │ │ ├── cache_manager.py # prefetch_for_next_step / ensure_resident │ │ ├── inference_engine.py # End-to-end inference engine │ │ └── metrics.py # MetricsBus for the IO Monitor @@ -578,14 +778,19 @@ Project_Chronos/ │ │ ├── io_profiler.py # Lookahead validation │ │ └── benchmark.py # End-to-end benchmarking │ ├── data/ -│ │ └── flexible_dataset.py # Flexible JSONL dataset loader +│ │ ├── flexible_dataset.py # Flexible JSONL dataset loader +│ │ └── 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 │ ├── backend/ │ │ ├── __init__.py # BackendDispatcher (cpu/mps/cuda/mlx) │ │ ├── dispatcher.py # Capability probing + priority logic │ │ └── ext/opencl.py # Third-party OpenCL extension hook │ ├── _backend_legacy.py # Backward-compatible older APIs +│ ├── bundle/ # chronos bundle: build/check/translate/finetune │ ├── serving/ -│ │ ├── __init__.py +│ │ ├── ollama_server.py # chronos serve: Ollama/OpenAI protocol layer +│ │ ├── 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) @@ -593,7 +798,9 @@ Project_Chronos/ │ ├── estimator.py # Live parameter / memory estimator │ └── tabs/ │ ├── config_tab.py # Config + Designer merged together +│ ├── data_tab.py │ ├── train_tab.py # Owns data_path +│ ├── monitor_tab.py # Training Monitor (shares chronos.monitor) │ ├── pipeline_tab.py # Per-stage datasets across all 6 stages │ ├── inference_tab.py │ ├── benchmark_tab.py # Markdown table + gr.BarPlot @@ -611,8 +818,9 @@ Project_Chronos/ │ ├── compare_minimind_chronos_v2.py │ └── compare_minimind_chronos_v3.py ├── tests/ -│ ├── test_smoke.py -│ ├── test_smoke_cuda.py +│ ├── test_smoke.py / test_smoke_cuda.py +│ ├── test_v3_kernel.py / test_speculative.py / test_moe_dispatch.py / test_precision_chain.py +│ ├── test_tokenizer_builder.py / test_longrun.py / test_monitor.py │ └── fixtures/ ├── docs/ │ └── vllm_integration.md @@ -654,7 +862,19 @@ timeline : HF IO, vLLM adapter, multi-backend dispatch : Stage 6 distillation and pluggable rewards : Web UI v2, benchmark plots, IO Monitor + v3 (0.3.0) + : Recurrent-depth kernel with shared_first KV + : top-2 + shared expert, grouped MoE dispatch + : Self-speculative decoding (Torch and MLX) + : Ollama/OpenAI protocol server + : Ready-to-run model bundles (chronos bundle) + : Chronos tokenizer, corpus pipeline, long-run trainer + : Training monitor page, API and console tab + : int8 KV, MLA absorbed decode, shared router scope (opt-in) + : Auto planner, v3 presets, tokenizer and serve tabs Next + : 50M en-zh results and ablations + : Speculative decoding: a draft path that pays off on larger models and lazy experts : Train IntentClassifier on large activation corpora : Benchmark 7B+ checkpoints : Inject masks on the vLLM worker path diff --git a/README_zh.md b/README_zh.md index 7fff8e3..23f76ef 100644 --- a/README_zh.md +++ b/README_zh.md @@ -54,6 +54,99 @@ sequenceDiagram --- +## v3 内核：循环深度 MoE（0.3.0） + +v3 让少量物理层被多次复用。执行分三段： +- Prelude 执行一次。 +- 每个共享循环组循环 `recurrent_depth` 次，每次做收缩式更新 `h ← h + dt·(block(h + e_loop) − h)`。 +- Coda 执行一次。 + +50M demo 有 6 个物理层，实际执行 18 个虚拟层（prelude 1 + 4 个循环层 × 深度 4 + coda 1）。 + +```mermaid +flowchart LR + X["tokens"] --> P["Prelude
执行一次"] + P --> G["循环组
MLA / SWA + top-2 MoE"] + G -->|"循环 d = 1..D
收缩更新，ACT 退出"| G + G --> C["Coda
执行一次"] + C --> H["LM head"] + P -. "lookahead 输入" .-> LR["LookaheadRouter
专家预取"] +``` + +**相对 v2 的变化** + +| | v2 | v3 | +|---|---|---| +| 循环内注意力 | 只在第 0 轮执行（`first_only`） | 每轮都执行 | +| 每个循环层的 KV cache | 1 个槽位；打开"每轮注意力"开关时每轮重写（实测 prefill+decode 与 full forward 偏差 0.94） | `shared_first`：每轮重算 Q，K/V 取第 0 轮，KV 槽位数 = 物理层数 | +| 路由 | top-1（gate 拿不到 CE 梯度） | top-2 + 1 个共享专家 | +| MoE 调度 | Python 逐专家循环，含 `.item()` 同步 | 排序分组调度：MLX 用 `gather_mm`，Torch 用 `argsort` + `index_add_` | +| 提前退出 | 每步同步 host | 纯张量 ACT，结果精确且与 batch 无关；decode 可逐 token 退出，不留 KV 缺口 | +| 解码 | 普通自回归 | 循环深度自推测解码 | + +`per_iteration`（每个虚拟层一个 KV 槽位）保留，用于消融对比。v2 checkpoint 自动以 `first_only` 加载，logits 与保存的 golden 张量误差在 1e-5 以内。 + +**自推测解码（不增加参数）** +- 同一个模型先用浅深度起草 γ 个 token，再用一次全深度批量前向校验。 +- greedy 输出与全深度 greedy 逐 token 一致。 +- 采样使用 rejection sampling，保持全深度的分布不变。 +- 拒绝后截断 cache：MLA 截断 `c_kv`，滑窗 cache 保留 `窗口 + γ_max` 个条目。 +- Torch（`chronos/runtime/speculative.py`）与 MLX（`chronos/mlx/speculative.py`、`ChronosMLXInferenceEngine.generate(speculative=True)`）使用同一套算法。 +- 两个后端的测试都覆盖 greedy 一致性、拒绝采样分布，以及回滚结果与从头 prefill 一致。 +- Web UI 推理 Tab 提供自推测开关，可调起草深度和 γ。 +- 50M demo 在实测过的所有运行时上都结果一致，但都没有提速。实测条件：预训练 step 13,996 checkpoint，greedy，batch 1，FLORES+ dev 英→中，起草深度 1，γ 4。报告在 `reports/speculative/`。 + +| 运行时 | 普通 tok/s | 自推测 tok/s | 比值 | 接受率 | +|---|---|---|---|---| +| MLX GPU bf16，专家全部在内存 | 112.0 | 107.4 | 0.96× | 42 % | +| MLX GPU bf16，懒加载专家（216 页中常驻 12 页） | 43.7 | 32.0 | 0.73× | 42 % | +| Torch MPS fp32，专家全部加载 | 32.1 | 25.9 | 0.81× | 40 % | +| Torch MPS fp32，懒加载专家 | 17.0 | 14.7 | 0.86× | 40 % | + +- 这个规模下，一次全深度单步的开销与一次批量校验相当。 +- 懒加载时，起草和校验路由到的专家不同，自推测从 SSD 读取的专家页多出 1.4–1.6 倍。 +- greedy 输出在 CPU fp32 和 Torch MPS fp32 下完全一致；GPU bf16 下，24 句中有 2 句在近似平票的位置翻转了 token。 +- 因此该开关默认关闭。 + +**精度链路** +- 全链路使用同一个计算精度。 +- FP32 只用在收益明确的位置：RMSNorm 统计量、SDPA softmax、RoPE 表、路由 softmax、LM head logits、优化器 master 权重。 +- MLX↔Torch 的 bf16 互转逐位一致。 +- `auto` 在 MLX GPU 和 CUDA 上选 bf16，在 CPU 和 MPS 上选 fp32（MPS 的 bf16 实测更慢，见下文）。 + +**v3 新增工具** + +| 命令 | 作用 | +|---|---| +| `chronos tokenizer` | 自研字节回退 BPE：按文字脚本预切分，按语言平衡合并；输出 HF 兼容的 `tokenizer.json` 和 chat template | +| `chronos corpus download / build / shards` | 固定版本下载、清洗、全局去重、剔除评测集泄漏、生成 uint16 分片 | +| `chronos longrun --spec recipe.json`（别名 `run`） | 按墙钟预算训练的 MLX 任务：每个阶段 warmup + cosine，NaN 回滚，续跑逐位一致，带验证和翻译探针 | +| `chronos serve` | 原生运行时上的 Ollama/OpenAI 协议服务（见下文） | +| `chronos monitor` | 只读训练监控页面 + JSON API（见下文） | + +`chronos.model.factory.fit_config` 按总参数量和激活参数量目标搜索模型配置。 + +**50M 英译中 demo（进行中）** +- 模型：总参数 49.69M，激活 12.09M；36 个专家；MLA latent 128；SWA 窗口 128；自研 16k 词表。 +- 训练：72 小时 MLX bf16 任务（先预训练，后 SFT），2026-11-05 启动。 +- 中途实测：预训练第 13,996 步（计划约 24.7k 步）valid CE 为 2.477（第 1,656 步时为 3.012）。该 checkpoint 已打包为 `chronos-enzh-50m-pretrain-s13996`（见“模型包”）。 +- FLORES+ 和 WMT22 上的 BLEU/chrF，以及各项消融（注意力模式、深度、懒加载专家预算、推测 γ），等训练结束后公布。在此之前不给出任何质量数字。 + +**推理与内存选项（全部默认关闭）。** 每个选项既是 config 字段，也可作为 `chronos serve` 和 Inference Tab 的覆盖项；测试在 Torch 和 MLX 上把每个选项与参考路径做对比。 + +| 字段 | 取值（第一个为默认） | 作用 | +|---|---|---| +| `kv_cache_dtype` | `auto`、`int8` | `int8` 把 MLA `c_kv` 与 SWA k/v 按 token、每 32 元素一组量化并带 scale；回滚和推测解码都支持 | +| `mla_decode` | `expanded`、`absorbed`、`auto` | `absorbed` 直接对 `c_kv` 打分，不展开 K/V。`auto` 只对短的缓存块做吸收，改为显式开启，因为实测 bf16 下“prefill 走 expanded、decode 走 absorbed”会偏离 full forward（logit 最大差 0.06） | +| `mla_cache_rope` | `false`、`true` | 额外缓存已加 RoPE 的 key，用内存换 decode 速度 | +| `router_scope` | `per_iteration`、`shared` | `shared` 让循环后续各轮复用第 0 轮的专家，每个 token 驻留的专家更少 | +| `draft_exit_kl_coef`、`draft_exit_depth` | `0`、`1` | 训练项：浅层出口对 detach 后的全深度 logits 做 KL，提高推测解码接受率 | +| `max_resident_experts`、`ram_budget_mb` | 不设 | 快速层和 RAM 层的硬上限；LRU 不淘汰 lookahead 预测要用的专家；预取覆盖一个虚拟层的全部循环轮次 | + +Web UI 新增 Tokenizer Tab（抽取 → 训练 → 评估）和 Serve Tab（启停 `chronos serve`；API key 通过环境变量传入，不进 argv）。 + +--- + ## 三级存储架构（M1：cluster-aware safetensors） ```mermaid @@ -183,7 +276,7 @@ flowchart LR 训练精度与资源策略： -- 默认 `--dtype auto`。MPS/MLX 训练优先解析为 BF16 以保证稳定性，CUDA/XPU 解析为 FP16；CPU 默认 FP32，只有显式传 `--dtype float16` 或 `--dtype bfloat16` 才启用 CPU autocast。 +- 默认 `--dtype auto`。权重、激活、KV cache 和专家页统一使用一个计算精度。FP32 只保留在融合 kernel 内部（SDPA softmax、RMSNorm 统计量），以及 RoPE 表、路由 softmax、loss logits 和优化器的 FP32 master/moments。`auto` 的解析结果：MLX GPU 为 BF16；CUDA 为 BF16，不支持时改用 FP16 并开启 loss scaling；CPU 和 MPS 为 FP32。在 MPS 上，原生 BF16 实测比 FP32 慢：M5 Max、torch 2.10、Chronos 50M 每个训练步 0.749 s 对 0.619 s，decode 每个 token 49.1 ms 对 37.3 ms。如需把权重和 KV 内存减半，可以显式传 `--dtype bfloat16`，代价是速度。 - CPU 训练默认使用物理核心；可用 `--cpu_threads` 或 `--cpu_budget_percent` 覆盖。 - macOS 上 MPS/MLX 训练默认强制 DataLoader workers 为 `0`，避免 multiprocessing 触发 Metal command-buffer 崩溃。CPU/CUDA 仍可使用 worker 进程；高级用户可用 `CHRONOS_ALLOW_METAL_DATALOADER_WORKERS=1` 覆盖这个保护。 - MLX 原生训练会按 `log_interval` 同步 UI 日志、标量读数和图表点；Web UI Stop 会在每个 batch 边界检查并停止。 @@ -194,8 +287,8 @@ trainer 在 `chronos.mlx.*` 中复刻 PyTorch 训练栈：masked CE、DPO/ORPO/G 蒸馏损失、load balance、temporal locality、lookahead soft-target 监督、 lookahead top-k 命中损失、router KL anchor。即使权重使用 BF16/FP16，路由 softmax、CE、Adam moments 等数值敏感路径仍以 FP32 计算；因此 `auto` 在 -MLX/MPS 上默认优先 BF16。FP16 的指数范围太小，在小模型/小 batch/路由分布 -剧烈波动时更容易把 loss 推成 NaN。 +MLX GPU 上默认 BF16（MPS 实测 FP32 更快，默认 FP32）。FP16 的指数范围太小， +在小模型/小 batch/路由分布剧烈波动时更容易把 loss 推成 NaN。 Web UI 的 MLX 训练日志与 CPU/MPS 对齐：step、loss、steps/s、ETA、 checkpoint 保存事件、stop 事件、verify 结果都会同步到界面。Stop 在 batch @@ -337,7 +430,7 @@ D_{\mathrm{KL}} --- -## 安装 (Not Ready in PyPI Yet) +## 安装 ```bash pip install Project_Chronos @@ -356,22 +449,29 @@ pip install -e ".[dev]" pip install "Project_Chronos[mlx]" ``` +**可选依赖组**（版本已固定）： +```bash +pip install "Project_Chronos[serve]" # chronos serve / chronos monitor（fastapi、uvicorn） +pip install "Project_Chronos[translate]" # 语料反分词 + BLEU/chrF 探针（sacremoses、sacrebleu） +``` + **vLLM 服务（可选，仅 Linux+CUDA）：** ```bash pip install vllm ``` > **minimind 依赖**：Project Chronos 使用 [minimind](https://github.com/jingyaogong/minimind) 作为 MoE 内核。 -> 若本地未找到，首次 import 时自动克隆至 `~/.cache/chronos/minimind-master/`。 +> 若本地未找到，首次 import 时自动克隆至 `~/.cache/chronos/minimind-master/`（需要 `git` 和网络）。 +> 设置 `CHRONOS_MINIMIND_DIR=/path/to/minimind` 可改用本地副本；模型包自带 `third_party/minimind`，启动脚本会设置这个变量。 > minimind 采用 **Apache-2.0** 授权，完整归属见 [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md)。 -**环境要求**：Python 3.10+，PyTorch 2.4+ +**环境要求**：Python 3.10+，PyTorch 2.6+ --- ## 快速开始 -### Web UI（M6 — 8 个 Tab，4 种语言） +### Web UI（12 个 Tab，4 种语言） ```bash chronos-ui @@ -379,7 +479,116 @@ chronos-ui python chronos_app.py ``` -包含：⚙️ Config（含右侧实时参数估算面板，合并了 Designer）/ 🏋️ Train（拥有 data_path）/ 🧪 6-Stage Pipeline（每阶段独立数据路径）/ 💬 Inference（含懒加载 vs 全量 DRAM 对比）/ 📦 Export（FP16/Q8_0 safetensors 与 GGUF）/ 📊 Benchmark（Markdown 表 + 对比图）/ 🔬 Auto-Tune（持久化日志 + 一键 Apply Best → Config）/ 📡 IO Monitor。i18n 支持 zh-Hans / zh-Hant / en / ja。 +包含：⚙️ Config（v3 预设、自动规划器和实时参数估算面板）/ 🗂️ Data（数据集扫描与转换）/ 🔤 Tokenizer（从数据集抽取文本、训练 Chronos BPE、评估并保存）/ 🏋️ Train（拥有 data_path，按所选阶段合并阶段参数）/ 📈 Training Monitor（只读监控 `runs/`，见下文）/ 🧪 6-Stage Pipeline（每阶段独立数据路径）/ 💬 Inference（含懒加载 vs 全量 DRAM 对比）/ 📦 Export（FP16/Q8_0 safetensors 与 GGUF）/ 🛰️ Serve（启停 `chronos serve` 并查看日志）/ 📊 Benchmark（Markdown 表 + 对比图）/ 🔬 Auto-Tune（持久化日志 + 一键 Apply Best → Config）/ 📡 IO Monitor。i18n 支持 zh-Hans / zh-Hant / en / ja。窗口宽 1280 px 时，Gradio 会把最后几个 Tab（Auto-Tune、IO Monitor）收进 “…” 溢出菜单。 + +### 预设与自动规划器 + +Config Tab 提供 5 档 v3 预设。它们都由规划器生成，存放在 `ui/presets_v3.json`，可用 `python -m ui.presets` 重新生成。参数量在 meta device 上实测，峰值内存是估算值。 + +| 预设 | 总参数 / 激活参数（实测） | 形状 | 内存档位 | 训练峰值（估算） | +|---|---|---|---|---| +| Chronos-v3 Tiny (smoke / CI) | 1.51M / 1.29M | E4 top-2，词表 6400，seq 128，深度 2 | none | 2.13 GB | +| Chronos-v3 50M (en→zh recipe) | 49.69M / 12.09M | h384，L6，E36 top-2，词表 16384，seq 512 | none | 16.4 GB（正在跑的任务实测峰值 8.72 GB） | +| Chronos-v3 200M (32 GB)（默认） | 200.72M / 28.20M | h576，E32 top-2 | balanced | 10.14 GB | +| Chronos-v3 500M (32 GB) | 499.53M / 61.59M | E48 top-3 | aggressive | 9.26 GB | +| Chronos-v3 1B (64 GB) | 1002.17M / 99.74M | E64 top-4，seq 2048 | aggressive | 18.92 GB | + +旧预设名（Recommended-CN、Tiny/Small/Medium/Large）仍可加载，会映射到最接近的 v3 档位。 + +**自动规划器**（Config Tab →“自动规划”） + +输入： +- 目标总参数量、是否开启 MoE； +- 专家数，以及“激活专家数”或“激活参数量”二选一； +- 内存节省档位：`none`、`balanced`、`aggressive`、`extreme`； +- 可选：数据集路径、内存、token 预算、词表大小、序列长度。 + +`chronos.model.planner.plan_training` 的处理步骤： +1. **分析数据集。** 读取分片 manifest 或抽样文件，把记录分成预训练、SFT、偏好、RL 四类，并读取分词器报告。数据里有词表大小、序列长度（p95 或上限）、token 数和分词器路径时，直接采用。 +2. **搜索架构。** 找出 meta device 实测参数量最接近目标的隐藏维度、层数、专家宽度和 GQA 头数。 +3. **确定学习率和 batch。** 按 DeepSeek LLM 的计算最优拟合式（η = 0.3118·C^−0.125，B = 0.2920·C^0.3271）给出，宽模型的学习率另设上限。 +4. **匹配内存。** 按内存档位选择 micro-batch、梯度累积、int8 KV、路由范围、激活检查点、优化器分页和专家 RAM 预算。 +5. **生成各阶段参数。** pretrain、SFT、DPO、ORPO、GRPO、distill 各有独立的学习率、batch、epoch、保存间隔，以及各自的专用项（β、温度、α）。 + +阶段表可以直接编辑。Train Tab 会把所选阶段那一行合并到共享参数之上；pretrain 行与共享控件保持同步。“保存方案”把 config、各阶段参数和报告写成 JSON。报告里每个数字都标注实测或估算。 + +### 训练监控（控制台 Tab + 独立页面/API） + +`chronos longrun` 的运行目录包含以下文件： +- `run_spec.json`、`metrics.jsonl`； +- `status.json`：心跳； +- `runtime.json`：精度策略、设备、版本、代码哈希； +- `parameters.jsonl`：逐张量的权重、梯度、更新量统计，以及逐虚拟层路由； +- `samples.jsonl`、`inspections/` 和 checkpoint。 + +控制台的 **Training Monitor** Tab 和独立页面读取的是同一批文件，都不会占用训练设备。遥测功能上线前启动的运行同样能显示进度，其状态由 `metrics.jsonl` 和进程表推断。 + +```bash +chronos monitor --runs-dir runs # http://127.0.0.1:7870/ +chronos serve --models-dir exports --monitor-runs-dir runs # 同一页面挂在 /monitor/ +``` + +JSON API（只读）： +- `GET /api/state?run=&history=1` +- `/api/metrics?run=&phase=&since=` +- `/api/parameters`、`/api/tensors`、`/api/tensor?name=` +- `/api/log?bytes=`、`/api/runs` + +安全： +- 默认只监听 127.0.0.1，不带认证。 +- 绑定其他地址时，必须同时给出 `--allow-remote` 和 `--api-key`（或设置环境变量 `CHRONOS_MONITOR_KEY`）。 +- 浏览器登录一次后获得 HttpOnly、SameSite=Strict 的会话 cookie；脚本使用 `Authorization: Bearer `。 +- 登录失败有限流。 +- 非本机访问时请加 `--ssl-certfile/--ssl-keyfile` 或使用 TLS 反向代理，否则密钥会以明文在网络上传输。 + +### Ollama / OpenAI 协议服务 + +`chronos serve` 把原生 Chronos 运行时挂在 Ollama 和 OpenAI 的 HTTP API 之后。内核不会被转换成其他架构，因此循环深度、SSD 懒加载专家与 lookahead 预取、KV 压缩、自推测解码都保持生效。 + +```bash +pip install "Project_Chronos[serve]" +chronos serve --models-dir ./exports --port 11435 # 11434 会和正在运行的 Ollama 冲突 +OLLAMA_HOST=127.0.0.1:11435 ollama list +OLLAMA_HOST=127.0.0.1:11435 ollama run <模型名> "你好" +``` + +- 接口：`/api/version`、`/api/tags`、`/api/show`、`/api/ps`、`/api/pull`、`/api/generate` 和 `/api/chat`（流式 NDJSON）、`/v1/models`、`/v1/chat/completions`（SSE）。 +- 模型来源：`--models-dir` 下成对的 `.pth` 和 `.config.json`。 +- Chronos 专属参数放在请求的 `options` 里：`chronos_speculative`、`chronos_draft_depth`、`chronos_gamma`、`chronos_miss_policy`。 +- stock `ollama` / llama.cpp 无法执行 Chronos 的 GGUF 文件，因为 MLA+SWA 混合层和收缩式循环没有对应的 stock 实现；请改用本服务。 +- 安全：默认只监听 127.0.0.1，不带认证。绑定其他地址必须同时给出 `--allow-remote` 和 `--api-key`，客户端随后发送 `Authorization: Bearer `。 + +### 模型包（开箱即用的预训练模型） + +`chronos bundle build` 把 `chronos longrun` 某次运行导出的一个 checkpoint 打包成解压即用的目录。打包过程只读取运行目录。 + +```bash +chronos bundle build --run-dir runs/enzh_50m_v1 --which best --wheel dist/project_chronos-0.3.0-py3-none-any.whl --pack +``` + +| 包内路径 | 内容 | +|---|---| +| `model/` | 权重 + `.config.json` 配置文件，分词器（`model/tokenizer/`） | +| `eval/` | FLORES+ dev/devtest（CC BY-SA 4.0），供 `bundle eval` / `probe` 使用 | +| `data/sft_valid/` | 该次运行的 SFT 验证分片 | +| `third_party/minimind/` | 内置 minimind（Apache-2.0），无需 git clone | +| `wheels/`、`bin/setup.sh`、`bin/chronos` | 离线安装到 `.venv`；启动脚本设置 `CHRONOS_BUNDLE` / `CHRONOS_MINIMIND_DIR` | +| `recipes/source_run_spec.json`、`recommended.json` | 产生这份权重的配方；微调档位 `smoke` / `quick` / `standard` / `full` | +| `bundle.json` | 来源信息（步数、token 数、valid CE）、实测质量探针、每个文件的 SHA-256 | + +接收方的用法： + +```bash +bin/setup.sh # 建 venv、装 wheel、运行 `bundle check` +bin/chronos bundle translate "The weather is very nice today." +bin/chronos bundle translate --backend mlx --speculative --file examples/translate_en.txt +bin/chronos bundle eval --split devtest -n 200 +bin/chronos bundle prepare-data my_pairs.jsonl --out work/data # 会从你的数据中剔除 FLORES+ +bin/chronos bundle finetune --data work/data --run-dir work/ft --tier quick --mode sft +bin/chronos bundle serve -- --port 11435 +``` + +`--speculative` 要求 `--repetition-penalty 1.0`（此时的默认值），因为惩罚会逐位置改变目标分布。`finetune` 写出一份 `init_checkpoint` 指向包内权重的 `chronos longrun` spec；运行目录里已有训练时会拒绝写入。用许可未知的数据（WMT19、OPUS-100）训练的模型包，`NOTICE.md` 中写明权重仅供研究使用。 ### 部署导出 @@ -461,6 +670,22 @@ python tools/compare_minimind_chronos_v3.py \ 输出包括：每阶段 loss、HF roundtrip Δlogit、tokens/sec、激活专家比例、常驻专家字节、M3 流水线富余空间、当前主机后端列表。 +### v2 循环深度与分页实测 + +`tools/benchmark_chronos_v2.py` 使用固定种子的可收敛 next-token 任务，对照 v1、固定循环深度 2/4、ACT adaptive，并实际执行 `sync_on_demand`、`on_demand` 和 `fallback_diagnostic` 三种专家分页路径。报告会写入 JSON，包含训练/验证 CE、PPL、next-token accuracy、prefill/decode tokens/s、KV 长度、平均循环深度、early-exit、RSS、专家驻留比例、预测 recall 和 full-resident 精确 token 对齐率。 + +```bash +# CPU：完整四组架构 + SSD/分页路径 +python -m tools.benchmark_chronos_v2 \ + --device cpu --train_steps 60 --output results/chronos_v2_benchmark_cpu.json + +# Apple Silicon MPS：功能/质量/Metal 内存快照 +python -m tools.benchmark_chronos_v2 \ + --device mps --train_steps 20 --output results/chronos_v2_benchmark_mps.json +``` + +该基准只代表 tiny 模型的架构回归，不等同于大模型最终能力；CUDA 显存峰值、NVMe 带宽和生产级 tokens/s 必须在目标 CUDA+NVMe 机器上另行运行。 + ### 专家聚簇存储（最大化 SSD 顺序读） ```python @@ -502,7 +727,10 @@ Project_Chronos/ │ ├── deps.py # 自动下载 minimind（若本地未找到） │ ├── __init__.py # 注册 AutoConfig / AutoModelForCausalLM │ ├── model/ -│ │ ├── config.py # ChronosConfig（lookahead/temporal/anchor/storage_format 等） +│ │ ├── config.py # ChronosConfig（lookahead/temporal/anchor/storage_format、v3 字段与 v2 迁移） +│ │ ├── recurrent.py # 循环/组嵌入、收缩步长、纯张量 ACT +│ │ ├── precision_ops.py # 单一计算精度的 RMSNorm/SDPA/RoPE +│ │ ├── factory.py # 按参数预算搜索配置 │ │ ├── hybrid_attention.py # MLAAttention + SlidingWindowAttention │ │ ├── lookahead_router.py # 逐 token 前瞻预测器（第二层） │ │ ├── moe_chronos.py # ChronosMOEFeedForward + 共享专家 + Soft Gating @@ -521,7 +749,11 @@ Project_Chronos/ │ │ └── prefill_scheduler.py # 编排 prefill 阶段批量预加载 │ ├── mlx/ # Apple Silicon 原生后端 │ │ ├── attention.py / moe.py / model.py / expert_store.py / inference.py +│ │ ├── moe_kernels.py # 排序 gather_mm 调度（mlx >= 0.32.4） +│ │ ├── speculative.py # MLX 自推测解码（与 Torch 对齐） +│ │ └── training/ # longrun.py、telemetry.py、recipes/ │ ├── runtime/ +│ │ ├── speculative.py # 循环深度自推测解码 │ │ ├── cache_manager.py # prefetch_for_next_step / ensure_resident（M3） │ │ ├── inference_engine.py # 端到端推理引擎（重排为 H2D-compute 重叠） │ │ └── metrics.py # MetricsBus（IO Monitor 数据源） @@ -540,14 +772,19 @@ Project_Chronos/ │ │ ├── io_profiler.py # Phase 1 验证（前瞻准确率） │ │ └── benchmark.py # 端到端基准测试 │ ├── data/ -│ │ └── flexible_dataset.py # 自动识别任意 JSONL 字段格式 +│ │ ├── flexible_dataset.py # 自动识别任意 JSONL 字段格式 +│ │ └── corpus/ # 下载 / 清洗 / 去重 / 泄漏过滤 / 分片 +│ ├── tokenizer/ # Chronos BPE 构建器（抽取、规范化、训练、评估） +│ ├── monitor/ # 只读运行监控：快照、API 服务、静态页面 │ ├── backend/ │ │ ├── __init__.py # BackendDispatcher（cpu/mps/cuda/mlx） │ │ ├── dispatcher.py # 探针 + 优先级 + 训练能力声明 │ │ └── ext/opencl.py # 第三方 OpenCL 插件钩子（stub） │ ├── _backend_legacy.py # 向后兼容 build_model() 等旧 API +│ ├── bundle/ # chronos bundle：打包/校验/翻译/微调 │ ├── serving/ -│ │ ├── __init__.py +│ │ ├── ollama_server.py # chronos serve：Ollama/OpenAI 协议层 +│ │ ├── runtime.py # 模型发现、懒加载专家、推测生成 │ │ └── vllm_adapter.py # 可选 vLLM 注册（无 vLLM 时优雅降级） │ └── cli.py # 统一 CLI ├── ui/ # Gradio Web UI（i18n: zh-Hans/zh-Hant/en/ja） @@ -555,7 +792,9 @@ Project_Chronos/ │ ├── estimator.py # 实时参数量/内存估算（与 Config 同步） │ └── tabs/ │ ├── config_tab.py # 已合并 Designer，右侧实时估算面板 +│ ├── data_tab.py │ ├── train_tab.py # 拥有 data_path（不再属于 Config） +│ ├── monitor_tab.py # Training Monitor（与 chronos.monitor 共用数据层） │ ├── pipeline_tab.py # 6 阶段，每段独立 data_path │ ├── inference_tab.py │ ├── benchmark_tab.py # Markdown 表 + gr.BarPlot @@ -573,8 +812,9 @@ Project_Chronos/ │ ├── compare_minimind_chronos_v2.py # v2 (M3+M4) │ └── compare_minimind_chronos_v3.py # v3 (含 6 阶段 + HF roundtrip + 后端报告) ├── tests/ -│ ├── test_smoke.py # 18 个单元测试 -│ ├── test_smoke_cuda.py # 仅 CUDA 主机执行 +│ ├── test_smoke.py / test_smoke_cuda.py（后者仅 CUDA 主机执行） +│ ├── test_v3_kernel.py / test_speculative.py / test_moe_dispatch.py / test_precision_chain.py +│ ├── test_tokenizer_builder.py / test_longrun.py / test_monitor.py │ └── fixtures/ # tiny_pretrain / tiny_sft / tiny_dpo / tiny_grpo ├── docs/ │ └── vllm_integration.md @@ -616,7 +856,19 @@ timeline : HF IO、vLLM adapter、多后端调度 : Stage 6 蒸馏和可插拔 reward : Web UI v2、Benchmark 图表、IO Monitor + v3（0.3.0） + : 循环深度内核 + shared_first KV + : top-2 + 共享专家，分组 MoE 调度 + : 自推测解码（Torch 与 MLX） + : Ollama/OpenAI 协议服务 + : 开箱即用模型包（chronos bundle） + : Chronos 分词器、语料流水线、长时训练器 + : 训练监控页面、API 与控制台 Tab + : int8 KV、MLA 吸收式解码、共享路由（显式开启） + : 自动规划器、v3 预设、分词器与 serve Tab Next + : 50M 英译中结果与消融 + : 自推测解码：在更大模型与懒加载专家下真正提速的起草路径 : 大规模激活语料训练 IntentClassifier : 7B+ checkpoint 基准测试 : vLLM worker path mask 注入 diff --git a/benchmark_compare.py b/benchmark_compare.py index e48d65f..5335327 100644 --- a/benchmark_compare.py +++ b/benchmark_compare.py @@ -101,10 +101,13 @@ def bench_minimind(): train_losses.append(loss.item()) log(f' step {step+1}/5 loss={loss.item():.4f}') - ram_train = ram_used_gb() + ram_train = ram_used_gb() - ram_before + # Optimizer state is a training working set, not inference memory. + del optimizer + gc.collect() model.eval() tps = measure_tokens_per_sec(model, cfg.vocab_size, 'cpu', n_tokens=30) - ram_infer = ram_used_gb() + ram_infer = ram_used_gb() - ram_before log(f' TPS={tps:.1f} RAM_train={ram_train:.3f}GB RAM_infer={ram_infer:.3f}GB') result = { @@ -116,7 +119,7 @@ def bench_minimind(): 'tokens_per_sec': round(tps, 1), 'train_losses': [round(l, 4) for l in train_losses], } - del model, optimizer + del model gc.collect() return result @@ -164,7 +167,11 @@ def bench_chronos(): train_losses.append(loss.item()) log(f' step {step+1}/5 loss={loss.item():.4f}') - ram_train = ram_used_gb() + ram_train = ram_used_gb() - ram_before + # Remove AdamW before measuring inference: allocator high-water marks are + # retained by the process, so report the delta from a post-cleanup baseline. + del optimizer + gc.collect() # Setup SSD offload + cache manager ssd_dir = '/tmp/chronos_bench_ssd' @@ -174,7 +181,9 @@ def bench_chronos(): mgr.start() mgr.warm_up() - # Measure TPS with SSD+DRAM mode + # Measure paged mode with the non-blocking M3 pipeline. The old + # benchmark called mgr.step(), which synchronously promoted and globally + # synchronized every token, measuring I/O stalls rather than decode. model.eval() x = torch.randint(0, cfg.vocab_size, (1, 32)) t0 = time.monotonic() @@ -191,10 +200,18 @@ def bench_chronos(): from chronos.model.moe_chronos import ChronosMOEFeedForward as CMOE cur_ids = [l.mlp.last_router_probs[:,-1,:].argmax(-1).item() for l in model.model.layers if isinstance(l.mlp, CMOE)] - mgr.step(lp3, cur_ids) + # Submit lookahead only; on CPU this avoids a fake blocking + # promotion loop. GPU/MPS inference_engine performs demand + # promotion on the next step and waits only for required pages. + predicted = mgr.prefetch_for_next_step(lp3) + # Observe the same pages on the following demand boundary. This + # turns the prefetch metric into a real hit-rate measurement + # instead of the old, permanently empty counter. + for eid in predicted[:4]: + mgr.prefetcher.wait_for(eid, timeout=0.01) elapsed = time.monotonic() - t0 tps = 30 / max(elapsed, 1e-6) - ram_infer = ram_used_gb() + ram_infer = max(0.0, ram_used_gb() - ram_before) cache_stats = mgr.stats() mgr.stop() @@ -202,7 +219,7 @@ def bench_chronos(): log(f' cache_stats={cache_stats}') result = { - 'name': 'Chronos (SSD+DRAM hybrid)', + 'name': 'Chronos (SSD+DRAM hybrid, async)', 'params_m': round(params_m, 1), 'ram_model_gb': round(ram_model_gb, 3), 'ram_train_gb': round(ram_train, 3), @@ -212,7 +229,7 @@ def bench_chronos(): 'cache_stats': cache_stats, 'kv_cache_type': 'MLA(latent)+SlidingWindow', } - del model, optimizer + del model gc.collect() return result diff --git a/benchmark_log.txt b/benchmark_log.txt index ae96fcc..3f61180 100644 --- a/benchmark_log.txt +++ b/benchmark_log.txt @@ -1,30 +1,30 @@ -[05:45:46] Starting benchmark comparison... -[05:45:46] === Benchmark 1: MiniMind (pure RAM) === -[05:45:46] params=91.4M RAM delta=0.399GB -[05:45:46] step 1/5 loss=8.9834 -[05:45:46] step 2/5 loss=8.9272 -[05:45:46] step 3/5 loss=8.9045 -[05:45:47] step 4/5 loss=8.8138 -[05:45:47] step 5/5 loss=8.9289 -[05:45:47] TPS=349.9 RAM_train=2.250GB RAM_infer=2.250GB -[05:45:47] === Benchmark 2: Chronos (SSD+DRAM hybrid) === -[05:45:47] params=111.1M RAM delta=0.016GB -[05:45:47] step 1/5 loss=8.9058 -[05:45:47] step 2/5 loss=8.8842 -[05:45:48] step 3/5 loss=8.9566 -[05:45:48] step 4/5 loss=9.0308 -[05:45:48] step 5/5 loss=8.8993 -[05:45:48] TPS=224.3 RAM_train=2.443GB RAM_infer=2.631GB -[05:45:48] cache_stats={'vram_experts': 4, 'vram_capacity': 52, 'ram_experts': 4, 'ram_capacity_dynamic': 208, 'expert_size_kb': 4992, 'pinned_ram_used_gb': 0.152, 'available_ram_gb': 57.0, 'pinned_ram_fraction': 0.0027, 'h2d_stream': 'default', 'storage_format': 'safetensors', 'cluster_aware': True, 'num_clusters': 1, 'prefetch_hits': 0, 'prefetch_misses': 0, 'total_requests': 58, 'hit_rate': 0.0} -[05:45:48] +[22:49:16] Starting benchmark comparison... +[22:49:16] === Benchmark 1: MiniMind (pure RAM) === +[22:49:17] params=91.4M RAM delta=0.394GB +[22:49:17] step 1/5 loss=8.9385 +[22:49:17] step 2/5 loss=8.9239 +[22:49:17] step 3/5 loss=8.8833 +[22:49:18] step 4/5 loss=8.8745 +[22:49:18] step 5/5 loss=8.8777 +[22:49:18] TPS=167.4 RAM_train=2.156GB RAM_infer=1.907GB +[22:49:19] === Benchmark 2: Chronos (SSD+DRAM hybrid) === +[22:49:19] params=111.1M RAM delta=0.097GB +[22:49:19] step 1/5 loss=8.9743 +[22:49:19] step 2/5 loss=8.9233 +[22:49:20] step 3/5 loss=9.0412 +[22:49:20] step 4/5 loss=8.9433 +[22:49:20] step 5/5 loss=8.9354 +[22:49:21] TPS=100.4 RAM_train=0.848GB RAM_infer=0.972GB +[22:49:21] cache_stats={'device': 'cpu', 'vram_experts': 3, 'vram_capacity': 3, 'ram_experts': 4, 'ram_capacity_dynamic': 4, 'expert_size_kb': 9984, 'cache_page_bytes': 81788928, 'resident_expert_bytes': 245366784, 'ram_buffer_bytes': 327155712, 'unique_expert_bytes': 327155712, 'pinned_ram_used_gb': 0.0, 'available_ram_gb': 86.1, 'pinned_ram_fraction': 0.0, 'h2d_stream': 'default', 'storage_format': 'safetensors', 'cluster_aware': True, 'layer_aware': False, 'num_clusters': 4, 'ssd_page_reads': 32, 'expert_evictions': 0, 'prefetch_hits': 48, 'prefetch_misses': 0, 'total_requests': 59, 'prefetch_observations': 48, 'hit_rate': 1.0, 'prefetch_queue_drops': 0, 'prefetch_wait_time_s': 0.0013231200573500246} +[22:49:21] === COMPARISON SUMMARY === -[05:45:48] MiniMind (pure RAM): -[05:45:48] tokens/s : 349.9 -[05:45:48] RAM (train) : 2.25 GB -[05:45:48] RAM (infer) : 2.25 GB -[05:45:48] Chronos (SSD+DRAM hybrid): -[05:45:48] tokens/s : 224.3 -[05:45:48] RAM (train) : 2.443 GB -[05:45:48] RAM (infer) : 2.631 GB -[05:45:48] KV cache : MLA(latent)+SlidingWindow -[05:45:48] Results saved to /Users/Fona/Downloads/Hybrid_LLM/Project_Chronos/benchmark_results.json +[22:49:21] MiniMind (pure RAM): +[22:49:21] tokens/s : 167.4 +[22:49:21] RAM (train) : 2.156 GB +[22:49:21] RAM (infer) : 1.907 GB +[22:49:21] Chronos (SSD+DRAM hybrid, async): +[22:49:21] tokens/s : 100.4 +[22:49:21] RAM (train) : 0.848 GB +[22:49:21] RAM (infer) : 0.972 GB +[22:49:21] KV cache : MLA(latent)+SlidingWindow +[22:49:21] Results saved to /Users/Fona/Downloads/Project_Chronos/benchmark_results.json diff --git a/benchmark_results.json b/benchmark_results.json index 96b5a7c..65f69f0 100644 --- a/benchmark_results.json +++ b/benchmark_results.json @@ -2,49 +2,60 @@ "minimind": { "name": "MiniMind (pure RAM)", "params_m": 91.4, - "ram_model_gb": 0.399, - "ram_train_gb": 2.25, - "ram_infer_gb": 2.25, - "tokens_per_sec": 349.9, + "ram_model_gb": 0.394, + "ram_train_gb": 2.156, + "ram_infer_gb": 1.907, + "tokens_per_sec": 167.4, "train_losses": [ - 8.9834, - 8.9272, - 8.9045, - 8.8138, - 8.9289 + 8.9385, + 8.9239, + 8.8833, + 8.8745, + 8.8777 ] }, "chronos": { - "name": "Chronos (SSD+DRAM hybrid)", + "name": "Chronos (SSD+DRAM hybrid, async)", "params_m": 111.1, - "ram_model_gb": 0.016, - "ram_train_gb": 2.443, - "ram_infer_gb": 2.631, - "tokens_per_sec": 224.3, + "ram_model_gb": 0.097, + "ram_train_gb": 0.848, + "ram_infer_gb": 0.972, + "tokens_per_sec": 100.4, "train_losses": [ - 8.9058, - 8.8842, - 8.9566, - 9.0308, - 8.8993 + 8.9743, + 8.9233, + 9.0412, + 8.9433, + 8.9354 ], "cache_stats": { - "vram_experts": 4, - "vram_capacity": 52, + "device": "cpu", + "vram_experts": 3, + "vram_capacity": 3, "ram_experts": 4, - "ram_capacity_dynamic": 208, - "expert_size_kb": 4992, - "pinned_ram_used_gb": 0.152, - "available_ram_gb": 57.0, - "pinned_ram_fraction": 0.0027, + "ram_capacity_dynamic": 4, + "expert_size_kb": 9984, + "cache_page_bytes": 81788928, + "resident_expert_bytes": 245366784, + "ram_buffer_bytes": 327155712, + "unique_expert_bytes": 327155712, + "pinned_ram_used_gb": 0.0, + "available_ram_gb": 86.1, + "pinned_ram_fraction": 0.0, "h2d_stream": "default", "storage_format": "safetensors", "cluster_aware": true, - "num_clusters": 1, - "prefetch_hits": 0, + "layer_aware": false, + "num_clusters": 4, + "ssd_page_reads": 32, + "expert_evictions": 0, + "prefetch_hits": 48, "prefetch_misses": 0, - "total_requests": 58, - "hit_rate": 0.0 + "total_requests": 59, + "prefetch_observations": 48, + "hit_rate": 1.0, + "prefetch_queue_drops": 0, + "prefetch_wait_time_s": 0.0013231200573500246 }, "kv_cache_type": "MLA(latent)+SlidingWindow" } diff --git a/chronos/__init__.py b/chronos/__init__.py index 8f02bec..09bf525 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.1.0" +__version__ = "0.3.0" # Bootstrap minimind dependency before any submodule imports from chronos.deps import ensure_minimind as _ensure_minimind diff --git a/chronos/app.py b/chronos/app.py index ffbf209..0206203 100644 --- a/chronos/app.py +++ b/chronos/app.py @@ -34,6 +34,7 @@ ) from ui.gradio_compat import gr +from ui.event_contract import UIEventContract, event_contract_app_kwargs import chronos.deps # auto-bootstrap minimind from ui.i18n import ( @@ -49,6 +50,9 @@ from ui.tabs.iomon_tab import build_iomon_tab from ui.tabs.pipeline_tab import build_pipeline_tab from ui.tabs.data_tab import build_data_tab +from ui.tabs.monitor_tab import build_monitor_tab, MONITOR_CSS +from ui.tabs.tokenizer_tab import build_tokenizer_tab +from ui.tabs.serve_tab import build_serve_tab APP_CSS = """ @@ -81,10 +85,12 @@ def _hero_html() -> str: def build_app() -> gr.Blocks: + contract = UIEventContract() with gr.Blocks( title="Project Chronos", - css=APP_CSS, + css=APP_CSS + MONITOR_CSS, fill_width=True, + head=contract.head(), ) as app: gr.HTML(_hero_html()) @@ -98,10 +104,13 @@ def build_app() -> gr.Blocks: config_state, config_inputs, cfg_save_dir = build_config_tab() build_data_tab() + build_tokenizer_tab() build_train_tab(config_state, cfg_save_dir) + monitor_tab = build_monitor_tab() build_pipeline_tab() build_inference_tab(config_state) build_export_tab() + build_serve_tab() build_benchmark_tab() build_autotune_tab(config_state, config_inputs) build_iomon_tab() @@ -115,7 +124,13 @@ def on_language_change(lang): inputs=[language], outputs=get_registered_components(), ) + wiring = getattr(monitor_tab, "monitor_wiring", None) + if wiring is not None and hasattr(app, "load"): + fn, inputs, outputs = wiring + app.load(fn=fn, inputs=inputs, outputs=outputs) + contract.blocks = app + app.chronos_event_contract = contract return app @@ -142,6 +157,7 @@ def main(): server_port=args.port, share=args.share, show_error=True, + app_kwargs=event_contract_app_kwargs(app), ) diff --git a/chronos/bundle/__init__.py b/chronos/bundle/__init__.py new file mode 100644 index 0000000..97e00bb --- /dev/null +++ b/chronos/bundle/__init__.py @@ -0,0 +1,22 @@ +"""Out-of-the-box model bundles (``chronos bundle``). + +A bundle is one self-contained directory: + + / + bundle.json identity, provenance, file hashes, measured quality + README.md / README_zh.md quick start + NOTICE.md data and code licences + model/.pth weights (+ .config.json sidecar) + model/tokenizer/ the model's own tokenizer + eval/ FLORES+ dev/devtest (CC BY-SA 4.0) for smoke tests + recipes/*.json ``chronos longrun`` specs for SFT / continued pretraining + recommended.json tiered next-stage parameters with their sources + examples/ tiny en-zh pairs to prove the fine-tune path + +``chronos bundle build`` creates it from a training run, ``chronos bundle check`` +verifies hashes and runs a CPU forward smoke test, ``chronos bundle translate`` +and ``chronos bundle finetune`` are the two entry points a recipient needs. +""" +from .builder import BUNDLE_FORMAT, build_bundle, check_bundle, load_bundle_manifest + +__all__ = ["BUNDLE_FORMAT", "build_bundle", "check_bundle", "load_bundle_manifest"] diff --git a/chronos/bundle/builder.py b/chronos/bundle/builder.py new file mode 100644 index 0000000..f13116b --- /dev/null +++ b/chronos/bundle/builder.py @@ -0,0 +1,342 @@ +"""Build and verify a self-contained Chronos model bundle (layout in ``__init__``).""" +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import tarfile +import time + +BUNDLE_FORMAT = "chronos-bundle/1" +MANIFEST = "bundle.json" +TOKENIZER_FILES = ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", + "chronos_tokenizer.json", "tokenizer_report.json") +FLORES_REL = os.path.join("facebook__flores__eng-zho", "data", "pair", "eng_Latn-zho_Hans") +FLORES_FILES = ("dev-00000-of-00001.parquet", "devtest-00000-of-00001.parquet") +# minimind sub-packages Chronos imports (model.*, trainer.trainer_utils, dataset.lm_dataset). +MINIMIND_PARTS = ("model", "trainer", "dataset", "LICENSE") +# Never shipped: they hold absolute local paths, optimizer state or caches. +_SKIP_NAMES = {"__pycache__", ".DS_Store"} + + +def sha256_file(path: str) -> str: + h = hashlib.sha256() + with open(path, "rb") as fh: + for chunk in iter(lambda: fh.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def load_bundle_manifest(bundle_dir: str) -> dict: + path = os.path.join(bundle_dir, MANIFEST) + if not os.path.isfile(path): + raise FileNotFoundError(f"{bundle_dir!r} is not a Chronos bundle (no {MANIFEST})") + with open(path, encoding="utf-8") as fh: + manifest = json.load(fh) + if manifest.get("format") != BUNDLE_FORMAT: + raise ValueError(f"unsupported bundle format {manifest.get('format')!r}; expected {BUNDLE_FORMAT}") + return manifest + + +def _copy(src: str, dst: str) -> str: + os.makedirs(os.path.dirname(dst), exist_ok=True) + shutil.copy2(src, dst) + return dst + + +def _copytree(src: str, dst: str) -> None: + shutil.copytree(src, dst, ignore=shutil.ignore_patterns(*_SKIP_NAMES, "*.pyc"), dirs_exist_ok=True) + + +def _file_table(root: str) -> dict: + """``relative path → {bytes, sha256}`` for every file except the manifest.""" + table = {} + for base, dirs, files in os.walk(root): + dirs[:] = sorted(d for d in dirs if d not in _SKIP_NAMES) + for name in sorted(files): + if name in _SKIP_NAMES: + continue + full = os.path.join(base, name) + rel = os.path.relpath(full, root).replace(os.sep, "/") + if rel == MANIFEST: + continue + table[rel] = {"bytes": os.path.getsize(full), "sha256": sha256_file(full)} + return table + + +def _sanitize(value, roots: list[str]): + """Replace local absolute prefixes (repo, home) so a bundle leaks no paths.""" + if isinstance(value, str): + for root, label in roots: + if value.startswith(root): + return label + value[len(root):] + return value + if isinstance(value, list): + return [_sanitize(v, roots) for v in value] + if isinstance(value, dict): + return {k: _sanitize(v, roots) for k, v in value.items()} + return value + + +def _path_roots(run_dir: str) -> list[tuple[str, str]]: + repo = os.path.dirname(os.path.dirname(os.path.abspath(run_dir))) + return [(repo, "$CHRONOS_REPO"), (os.path.expanduser("~"), "~")] + + +def _metrics(run_dir: str) -> list[dict]: + path = os.path.join(run_dir, "metrics.jsonl") + if not os.path.exists(path): + return [] + out = [] + with open(path, encoding="utf-8") as fh: + for line in fh: + try: + out.append(json.loads(line)) + except json.JSONDecodeError: + continue # a line being written right now + return out + + +def run_provenance(run_dir: str, phase: str, which: str) -> dict: + """Training state behind ``_[.best].pth`` from ``metrics.jsonl``. + + ``best``: the validation with the lowest valid CE (the run exports + ``.best`` right after it); ``final``: the last validation.""" + records = _metrics(run_dir) + vals = [r for r in records if r.get("event") == "validation" and r.get("phase") == phase] + if not vals: + return {"phase": phase, "note": "no validation record found"} + val = min(vals, key=lambda r: r["valid_ce"]) if which == "best" else vals[-1] + steps = [r for r in records if r.get("phase") == phase and "tokens_seen" in r and r.get("step", 0) <= val["step"]] + last = steps[-1] if steps else {} + return {"phase": phase, "step": val["step"], "planned_steps_at_export": last.get("total"), + "time": val.get("time"), "valid_ce": round(float(val["valid_ce"]), 4), + "valid_ppl": round(float(val.get("valid_ppl", 0.0)), 3), "tokens_seen": last.get("tokens_seen"), + "label": "measured (held-out shard validation, 8 fixed batches)", + "validation_history": [{"step": r["step"], "valid_ce": round(float(r["valid_ce"]), 4)} for r in vals]} + + +def _stable_copy(src: str, dst: str, retries: int = 3) -> str: + """Copy ``src`` and prove it did not change mid-copy (a live run may re-export it).""" + for _ in range(retries): + before = os.stat(src) + _copy(src, dst) + after = os.stat(src) + if (before.st_mtime_ns, before.st_size) == (after.st_mtime_ns, after.st_size): + return sha256_file(dst) + time.sleep(2.0) + raise RuntimeError(f"{src} kept changing while being copied") + + +def _planner_preset(vocab_size: int, hidden_size: int) -> dict | None: + """The v3 preset with the same topology, if one exists (``ui/presets_v3.json``).""" + path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), + "ui", "presets_v3.json") + if not os.path.exists(path): + return None + with open(path, encoding="utf-8") as fh: + presets = json.load(fh).get("presets", {}) + for preset in presets.values(): + if preset.get("vocab_size") == vocab_size and preset.get("hidden_size") == hidden_size: + return preset + return None + + +def build_bundle(run_dir: str, out_dir: str | None = None, *, phase: str = "pretrain", which: str = "best", + name: str | None = None, minimind_dir: str | None = None, wheel: str | None = None, + include_sft_valid: bool = True, status: str = "interim", log=print) -> str: + """Assemble a bundle from a ``chronos longrun`` run directory (read-only on the run).""" + from chronos.model.checkpoint import checkpoint_config_path, tokenizer_metadata + from . import docs, recommend + run_dir = os.path.abspath(run_dir) + with open(os.path.join(run_dir, "run_spec.json"), encoding="utf-8") as fh: + spec = json.load(fh) + prefix = spec.get("export_prefix", "chronos_enzh") + stem = f"{prefix}_{phase}" + (".best" if which == "best" else "") + src_ckpt = os.path.join(run_dir, stem + ".pth") + if not os.path.isfile(src_ckpt): + raise FileNotFoundError(f"no exported checkpoint {src_ckpt}") + prov = run_provenance(run_dir, phase, which) + name = name or f"chronos-enzh-50m-{phase}-s{prov.get('step', 'x')}" + out_dir = os.path.abspath(out_dir or os.path.join(os.path.dirname(os.path.dirname(run_dir)), "dist", "bundles", + name)) + if os.path.exists(out_dir) and os.listdir(out_dir): + raise FileExistsError(f"{out_dir} exists and is not empty; pick another --out") + os.makedirs(out_dir, exist_ok=True) + model_name = name.replace(".", "_") + ckpt_rel = f"model/{model_name}.pth" + log(f"[bundle] weights {src_ckpt} → {ckpt_rel}") + ckpt_sha = _stable_copy(src_ckpt, os.path.join(out_dir, ckpt_rel)) + # Sidecar: same config, tokenizer path made relative (resolve_checkpoint_tokenizer + # finds /tokenizer first; "tokenizer" is also correct relative to it). + with open(checkpoint_config_path(src_ckpt), encoding="utf-8") as fh: + sidecar = json.load(fh) + tok_src = spec["tokenizer_dir"] + tok_dst = os.path.join(out_dir, "model", "tokenizer") + for f in TOKENIZER_FILES: + if os.path.exists(os.path.join(tok_src, f)): + _copy(os.path.join(tok_src, f), os.path.join(tok_dst, f)) + tok_meta = tokenizer_metadata(tok_dst) + tok_meta["path"] = "tokenizer" + sidecar.update(checkpoint=os.path.basename(ckpt_rel), tokenizer=tok_meta) + sidecar["extra"] = dict(sidecar.get("extra") or {}, bundle=name, source_checkpoint_sha256=ckpt_sha) + with open(checkpoint_config_path(os.path.join(out_dir, ckpt_rel)), "w", encoding="utf-8") as fh: + json.dump(sidecar, fh, indent=2, sort_keys=True) + # Evaluation data in the raw layout ``load_eval_sets`` reads. + raw_root = (spec.get("extra") or {}).get("raw_root") + flores = [] + for f in FLORES_FILES: + src = os.path.join(raw_root or "", FLORES_REL, f) + if raw_root and os.path.exists(src): + _copy(src, os.path.join(out_dir, "eval", FLORES_REL, f)) + flores.append(f) + if raw_root and os.path.exists(os.path.join(raw_root, "facebook__flores__eng-zho", "download_manifest.json")): + with open(os.path.join(raw_root, "facebook__flores__eng-zho", "download_manifest.json"), encoding="utf-8") as fh: + dl = json.load(fh) + for item in dl.get("files", []): + item.pop("local", None) + recommend.write_json(os.path.join(out_dir, "eval", "facebook__flores__eng-zho", "download_manifest.json"), dl) + if include_sft_valid: + valid_src = os.path.join(spec["shards_dir"], "sft", "valid") + if os.path.isdir(valid_src): + _copytree(valid_src, os.path.join(out_dir, "data", "sft_valid")) + # minimind (Apache-2.0) so the recipient needs neither git nor network. + from chronos.deps import ensure_minimind + mm = minimind_dir or str(ensure_minimind()) + for part in MINIMIND_PARTS: + src = os.path.join(mm, part) + if os.path.isdir(src): + _copytree(src, os.path.join(out_dir, "third_party", "minimind", part)) + elif os.path.isfile(src): + _copy(src, os.path.join(out_dir, "third_party", "minimind", part)) + mm_rev = _git_head(mm) + if wheel: + _copy(wheel, os.path.join(out_dir, "wheels", os.path.basename(wheel))) + # Recipes: the exact spec that produced these weights (paths made relative). + roots = _path_roots(run_dir) + source_spec = _sanitize(spec, roots) + recommend.write_json(os.path.join(out_dir, "recipes", "source_run_spec.json"), source_spec) + preset = _planner_preset(int(sidecar["config"].get("vocab_size", 0)), int(sidecar["config"].get("hidden_size", 0))) + recs = recommend.build_recommendations(spec, preset) + recommend.write_json(os.path.join(out_dir, "recommended.json"), recs) + docs.write_examples(os.path.join(out_dir, "examples")) + manifest = { + "format": BUNDLE_FORMAT, "name": name, "status": status, + "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "prompt_format": "pretrain_tags" if phase == "pretrain" else "chat_sft", + "model": {"checkpoint": ckpt_rel, "config": checkpoint_config_path(ckpt_rel), "stage": phase, + "which": which, "sha256": ckpt_sha, "params": (spec.get("extra") or {}).get("params"), + "dtype_on_disk": "see config; weights saved by the MLX exporter", + "architecture": {k: sidecar["config"].get(k) for k in ( + "architecture_version", "hidden_size", "num_hidden_layers", "prelude_layers", "coda_layers", + "recurrent_depth", "recurrent_attention_mode", "num_experts", "num_experts_per_tok", + "num_shared_experts", "vocab_size", "max_position_embeddings")}}, + "tokenizer": {"dir": "model/tokenizer", **{k: v for k, v in tok_meta.items() if k != "path"}}, + "training": prov, + "eval": {"flores_plus": flores, "licence": "CC BY-SA 4.0", "root": "eval"}, + "third_party": {"minimind": {"dir": "third_party/minimind", "licence": "Apache-2.0", "commit": mm_rev}}, + "chronos_version": _chronos_version(), + } + with open(os.path.join(out_dir, MANIFEST), "w", encoding="utf-8") as fh: + json.dump(manifest, fh, ensure_ascii=False, indent=2) + docs.write_docs(out_dir, manifest, recs) + manifest["files"] = _file_table(out_dir) + with open(os.path.join(out_dir, MANIFEST), "w", encoding="utf-8") as fh: + json.dump(manifest, fh, ensure_ascii=False, indent=2) + log(f"[bundle] wrote {out_dir} ({len(manifest['files'])} files)") + return out_dir + + +def _git_head(path: str) -> str | None: + head = os.path.join(path, ".git", "HEAD") + try: + ref = open(head, encoding="utf-8").read().strip() + if ref.startswith("ref: "): + return open(os.path.join(path, ".git", ref[5:]), encoding="utf-8").read().strip() + return ref + except OSError: + return None + + +def _chronos_version() -> str: + import chronos + return getattr(chronos, "__version__", "unknown") + + +def _rewrite_manifest(bundle_dir: str, manifest: dict) -> None: + manifest["files"] = _file_table(bundle_dir) + with open(os.path.join(bundle_dir, MANIFEST), "w", encoding="utf-8") as fh: + json.dump(manifest, fh, ensure_ascii=False, indent=2) + + +def record_quality_probe(bundle_dir: str, n: int = 50, *, repetition_penalties=(1.0, 1.2), + device: str = "cpu", log=print) -> dict: + """Measure BLEU/chrF of the bundled weights on FLORES+ dev and store it in the manifest.""" + from .runtime import BundleModel, eval_pairs, score + manifest = load_bundle_manifest(bundle_dir) + pairs = eval_pairs(bundle_dir, "dev")[:n] + if not pairs: + raise FileNotFoundError("bundle has no FLORES+ dev file under eval/") + bm = BundleModel(bundle_dir, device=device) + results = {} + for rp in repetition_penalties: + log(f"[bundle] probe: {len(pairs)} FLORES+ dev sentences, repetition_penalty={rp}") + results[f"rp_{rp}"] = score(bm, pairs, repetition_penalty=rp) + manifest["quality"] = {"set": "flores_plus_dev", "first_n": len(pairs), + "label": "measured on the bundled checkpoint (sacrebleu 2.6.0; BLEU tokenize=zh for →zh, " + "13a for →en; corpus-level)", + "results": results} + from . import docs + docs.write_docs(bundle_dir, manifest) + _rewrite_manifest(bundle_dir, manifest) + return manifest["quality"] + + +def check_bundle(bundle_dir: str, *, smoke: bool = True, device: str = "cpu", log=print) -> dict: + """Verify every file hash, then (``smoke``) load the model and translate one sentence.""" + bundle_dir = os.path.abspath(bundle_dir) + manifest = load_bundle_manifest(bundle_dir) + bad, missing = [], [] + for rel, meta in (manifest.get("files") or {}).items(): + path = os.path.join(bundle_dir, rel) + if not os.path.isfile(path): + missing.append(rel) + elif os.path.getsize(path) != meta["bytes"] or sha256_file(path) != meta["sha256"]: + bad.append(rel) + report = {"bundle": manifest["name"], "files": len(manifest.get("files") or {}), "missing": missing, + "hash_mismatch": bad, "ok": not bad and not missing} + log(f"[check] {report['files']} files, missing={len(missing)}, hash_mismatch={len(bad)}") + if smoke and report["ok"]: + from .runtime import BundleModel, eval_pairs + t0 = time.monotonic() + bm = BundleModel(bundle_dir, device=device) + pairs = eval_pairs(bundle_dir, "dev")[:1] or [("The weather is nice today.", "")] + hyp = bm.translate([pairs[0][0]], "en_zh", max_new_tokens=64)[0] + report["smoke"] = {"src": pairs[0][0], "hyp": hyp, "non_empty": bool(hyp.strip()), + "seconds": round(time.monotonic() - t0, 1), "device": bm.device} + report["ok"] = report["ok"] and bool(hyp.strip()) + log(f"[check] smoke en→zh: {hyp!r}") + return report + + +def pack_bundle(bundle_dir: str, archive: str | None = None, log=print) -> str: + """``tar.gz`` the bundle (top-level folder = bundle name) and write ``.sha256``.""" + bundle_dir = os.path.abspath(bundle_dir) + archive = os.path.abspath(archive or bundle_dir.rstrip(os.sep) + ".tar.gz") + if os.path.exists(archive): + raise FileExistsError(f"{archive} exists") + base = os.path.basename(bundle_dir.rstrip(os.sep)) + + def _filter(info): + return None if os.path.basename(info.name) in _SKIP_NAMES or info.name.endswith(".pyc") else info + + tmp = archive + ".partial" + with tarfile.open(tmp, "w:gz", compresslevel=6) as tar: + tar.add(bundle_dir, arcname=base, filter=_filter) + os.replace(tmp, archive) + digest = sha256_file(archive) + with open(archive + ".sha256", "w", encoding="utf-8") as fh: + fh.write(f"{digest} {os.path.basename(archive)}\n") + log(f"[bundle] {archive} ({os.path.getsize(archive) / 2**20:.1f} MiB) sha256={digest}") + return archive diff --git a/chronos/bundle/cli.py b/chronos/bundle/cli.py new file mode 100644 index 0000000..8fc317a --- /dev/null +++ b/chronos/bundle/cli.py @@ -0,0 +1,177 @@ +"""``chronos bundle {build|pack|check|probe|translate|eval|prepare-data|finetune|serve}``.""" +from __future__ import annotations + +import argparse +import json +import os +import sys + + +def _bundle_dir(value: str | None) -> str: + """Default: the bundle this file was installed from, else the current directory.""" + return os.path.abspath(value or os.environ.get("CHRONOS_BUNDLE") or ".") + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(prog="chronos bundle", description=__doc__) + sub = p.add_subparsers(dest="cmd", required=True) + + b = sub.add_parser("build", help="assemble a bundle from a chronos longrun run directory (read-only on the run)") + b.add_argument("--run-dir", required=True) + b.add_argument("--out", default=None) + b.add_argument("--phase", default="pretrain", choices=["pretrain", "sft"]) + b.add_argument("--which", default="best", choices=["best", "final"]) + b.add_argument("--name", default=None) + b.add_argument("--status", default="interim", choices=["interim", "final"]) + b.add_argument("--wheel", default=None, help="a project_chronos wheel to ship under wheels/") + b.add_argument("--probe-n", type=int, default=50, help="FLORES+ dev sentences for the quality probe (0 = skip)") + b.add_argument("--pack", action="store_true", help="also write .tar.gz + .sha256") + + k = sub.add_parser("pack", help="tar.gz a bundle directory") + k.add_argument("bundle") + k.add_argument("--archive", default=None) + + c = sub.add_parser("check", help="verify file hashes and run a CPU smoke translation") + c.add_argument("bundle", nargs="?") + c.add_argument("--no-smoke", action="store_true") + c.add_argument("--device", default="cpu") + + q = sub.add_parser("probe", help="measure BLEU/chrF on FLORES+ dev and store it in bundle.json") + q.add_argument("bundle", nargs="?") + q.add_argument("-n", type=int, default=50) + q.add_argument("--device", default="cpu") + + t = sub.add_parser("translate", help="translate text (argument, --file or stdin)") + t.add_argument("text", nargs="*") + t.add_argument("--bundle", default=None) + t.add_argument("--file", default=None) + t.add_argument("--direction", default="en_zh", choices=["en_zh", "zh_en"]) + t.add_argument("--repetition-penalty", type=float, default=None, + help="default 1.2; 1.0 with --speculative (required for exact speculative greedy)") + t.add_argument("--max-new-tokens", type=int, default=256) + t.add_argument("--backend", default="torch", choices=["torch", "mlx"]) + t.add_argument("--device", default="auto", help="torch: auto|cpu|mps|cuda; mlx: auto|gpu|cpu") + t.add_argument("--dtype", default="float32") + t.add_argument("--speculative", action="store_true", + 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)") + + 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("--device", default="auto") + e.add_argument("--out", default=None, help="write the report JSON here") + + 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") + d.add_argument("--bundle", default=None) + d.add_argument("--out", required=True) + d.add_argument("--zh-field", default="zh") + d.add_argument("--en-field", default="en") + d.add_argument("--unsegment-zh", action="store_true", help="remove word-segmentation spaces between CJK chars") + d.add_argument("--detokenize-en", action="store_true", help="undo Moses tokenization on the English side") + d.add_argument("--workers", type=int, default=4) + + f = sub.add_parser("finetune", help="write a chronos longrun spec that starts from the bundle weights (MLX)") + f.add_argument("--bundle", default=None) + f.add_argument("--data", required=True, help="the --out of prepare-data (or a shards dir)") + f.add_argument("--run-dir", required=True) + f.add_argument("--tier", default="quick", choices=["smoke", "quick", "standard", "full"]) + f.add_argument("--mode", default="sft", choices=["sft", "pretrain", "pretrain+sft"]) + 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") + + s = sub.add_parser("serve", help="Ollama/OpenAI-protocol server over the bundle's model/ (loopback only)") + s.add_argument("--bundle", default=None) + s.add_argument("rest", nargs=argparse.REMAINDER, help="extra `chronos serve` options") + + a = p.parse_args(argv) + return _run(a) + + +def _print(obj) -> None: + print(json.dumps(obj, ensure_ascii=False, indent=2)) + + +def _run(a) -> int: + from . import builder + if a.cmd == "build": + out = builder.build_bundle(a.run_dir, a.out, phase=a.phase, which=a.which, name=a.name, wheel=a.wheel, + status=a.status) + if a.probe_n: + _print(builder.record_quality_probe(out, a.probe_n)) + report = builder.check_bundle(out) + _print(report) + if a.pack: + builder.pack_bundle(out) + return 0 if report["ok"] else 1 + if a.cmd == "pack": + builder.pack_bundle(a.bundle, a.archive) + return 0 + if a.cmd == "check": + report = builder.check_bundle(_bundle_dir(a.bundle), smoke=not a.no_smoke, device=a.device) + _print(report) + return 0 if report["ok"] else 1 + if a.cmd == "probe": + _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: + lines = [ln.strip() for ln in fh if ln.strip()] + 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) + 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 + 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) + report["set"] = f"flores_plus_{a.split}" + _print(report) + if a.out: + with open(a.out, "w", encoding="utf-8") as fh: + json.dump(report, fh, ensure_ascii=False, indent=2) + return 0 + if a.cmd == "prepare-data": + from .finetune import prepare_data + prepare_data(_bundle_dir(a.bundle), a.inputs, a.out, zh_field=a.zh_field, en_field=a.en_field, + unsegment_zh=a.unsegment_zh, detokenize_en=a.detokenize_en, workers=a.workers) + return 0 + if a.cmd == "finetune": + from .finetune import launch, write_finetune_spec + 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) + print(f"[finetune] spec: {spec}") + if a.dry_run: + print(f"[finetune] start it with: chronos longrun --spec {spec}") + return 0 + return launch(spec) + if a.cmd == "serve": + from chronos.serving.ollama_server import main as serve_main + from .builder import load_bundle_manifest + bundle = _bundle_dir(a.bundle) + load_bundle_manifest(bundle) + rest = [r for r in a.rest if r != "--"] + return serve_main(["--models-dir", os.path.join(bundle, "model"), *rest]) + raise SystemExit(f"unknown command {a.cmd}") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/bundle/docs.py b/chronos/bundle/docs.py new file mode 100644 index 0000000..3c23c07 --- /dev/null +++ b/chronos/bundle/docs.py @@ -0,0 +1,320 @@ +"""Bundle documents: README (en/zh), NOTICE, example data and launcher scripts.""" +from __future__ import annotations + +import json +import os + +# Written for this bundle (not taken from any corpus); just enough to prove +# the prepare → fine-tune → translate path end to end. +EXAMPLE_PAIRS = [ + ("The weather is very nice today.", "今天天气很好。"), + ("Please close the window before you leave.", "离开之前请关好窗户。"), + ("I will call you when I arrive at the station.", "我到车站后会给你打电话。"), + ("This book is more interesting than I expected.", "这本书比我想象的更有趣。"), + ("We need to finish the report by Friday.", "我们需要在周五之前完成报告。"), + ("The museum is closed on Mondays.", "博物馆周一闭馆。"), + ("Could you tell me how to get to the library?", "你能告诉我怎么去图书馆吗？"), + ("He has been learning Chinese for three years.", "他学中文已经三年了。"), + ("The meeting has been moved to next Tuesday.", "会议改到了下周二。"), + ("Drink more water when the weather is hot.", "天气热的时候要多喝水。"), + ("The train was delayed by twenty minutes.", "火车晚点了二十分钟。"), + ("She opened a small coffee shop near the university.", "她在大学附近开了一家小咖啡店。"), + ("Remember to save your work before closing the program.", "关闭程序之前记得保存你的工作。"), + ("The children are playing football in the park.", "孩子们在公园里踢足球。"), + ("Prices in this market are usually lower.", "这个市场的价格通常更低。"), + ("My phone battery is almost empty.", "我的手机快没电了。"), + ("The new bridge will open to traffic next year.", "新桥将于明年通车。"), + ("Please send me the file by email.", "请通过电子邮件把文件发给我。"), + ("It took us two hours to climb the mountain.", "我们花了两个小时爬上这座山。"), + ("The doctor advised him to get more sleep.", "医生建议他多睡觉。"), + ("Our team won the first match of the season.", "我们队赢得了本赛季的第一场比赛。"), + ("The restaurant is famous for its noodles.", "这家餐厅以面条闻名。"), + ("I forgot to bring my umbrella this morning.", "我今天早上忘了带伞。"), + ("The company plans to hire fifty new engineers.", "公司计划招聘五十名新工程师。"), + ("Turn left at the second traffic light.", "在第二个红绿灯处左转。"), + ("The river flows through the center of the city.", "这条河流经城市中心。"), + ("He fixed the old bicycle by himself.", "他自己修好了那辆旧自行车。"), + ("Students must hand in their homework on time.", "学生必须按时交作业。"), + ("The flowers in the garden bloom every spring.", "花园里的花每年春天都会开。"), + ("We watched a film together last night.", "我们昨晚一起看了一部电影。"), + ("This road is very busy during rush hour.", "这条路在高峰时段非常拥堵。"), + ("She speaks three languages fluently.", "她能流利地说三种语言。"), + ("The hotel room has a view of the sea.", "酒店房间可以看到大海。"), + ("Please turn off the lights when you leave the room.", "离开房间时请关灯。"), + ("The price of vegetables rose this month.", "这个月蔬菜价格上涨了。"), + ("My grandmother tells the best stories.", "我奶奶讲的故事最好听。"), + ("The concert starts at eight in the evening.", "音乐会晚上八点开始。"), + ("They planted trees along the new road.", "他们在新路两旁种了树。"), + ("I am looking for a quiet place to study.", "我在找一个安静的地方学习。"), + ("The airport is about forty kilometres from here.", "机场离这里大约四十公里。"), +] + + +def write_examples(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) + with open(os.path.join(out_dir, "pairs_en_zh.jsonl"), "w", encoding="utf-8") as fh: + for en, zh in EXAMPLE_PAIRS: + fh.write(json.dumps({"en": en, "zh": zh}, ensure_ascii=False) + "\n") + with open(os.path.join(out_dir, "translate_en.txt"), "w", encoding="utf-8") as fh: + fh.write("The weather is very nice today.\nPlease send me the file by email.\n") + + +LAUNCHER = """#!/bin/sh +# Runs `chronos` with this bundle's settings (no git, no network needed). +HERE="$(cd "$(dirname "$0")/.." && pwd)" +export CHRONOS_BUNDLE="${CHRONOS_BUNDLE:-$HERE}" +export CHRONOS_MINIMIND_DIR="${CHRONOS_MINIMIND_DIR:-$HERE/third_party/minimind}" +if [ -x "$HERE/.venv/bin/chronos" ]; then exec "$HERE/.venv/bin/chronos" "$@"; fi +exec chronos "$@" +""" + +SETUP = """#!/bin/sh +# One-time setup: a virtualenv in .venv with Chronos and its pinned extras. +set -eu +HERE="$(cd "$(dirname "$0")/.." && pwd)" +PY="${PYTHON:-python3}" +"$PY" -c 'import sys; sys.exit(0 if sys.version_info >= (3, 10) else "Python >= 3.10 required")' +[ -d "$HERE/.venv" ] || "$PY" -m venv "$HERE/.venv" +. "$HERE/.venv/bin/activate" +python -m pip install --upgrade pip +EXTRAS="translate,serve" +if [ "$(uname -s)" = Darwin ] && [ "$(uname -m)" = arm64 ]; then EXTRAS="$EXTRAS,mlx"; fi +WHEEL="$(ls "$HERE"/wheels/*.whl 2>/dev/null | head -n 1 || true)" +if [ -n "$WHEEL" ]; then + python -m pip install "$WHEEL[$EXTRAS]" +else + echo "no wheel in wheels/: install Chronos from source with pip install -e '[$EXTRAS]'" >&2 + exit 1 +fi +"$HERE/bin/chronos" bundle check +""" + + +def _fmt(value, digits: int = 2): + return "n/a" if value is None else (f"{value:,}" if isinstance(value, int) else f"{value:.{digits}f}") + + +def _quality_rows(manifest: dict) -> list[str]: + q = manifest.get("quality") or {} + rows = [] + for key, res in (q.get("results") or {}).items(): + rp = key.split("_", 1)[1] + for direction in ("en_zh", "zh_en"): + r = res.get(direction) + if r: + rows.append(f"| {direction.replace('_', '→')} | {rp} | {r['n']} | {r['bleu']:.2f} | {r['chrf']:.2f} |") + return rows + + +def _tier_rows(recs: dict) -> list[str]: + rows = [] + for tier, t in recs["mlx_longrun"]["tiers"].items(): + sft = t["phases"]["sft"][0] + rows.append(f"| `{tier}` | {t['hours']} | {sft['peak_lr']:g} | {sft['tokens_per_step']:,} | " + f"{sft['warmup_fraction']:g} | {t['why']} |") + return rows + + +def _readme_en(m: dict, recs: dict) -> str: + tr, arch, params = m.get("training") or {}, m["model"]["architecture"], m["model"].get("params") or {} + status = ("**Interim checkpoint**: the training run was still in progress when this bundle was built. " + "Rebuild from the finished run for the final weights.") if m.get("status") == "interim" else \ + "Final checkpoint of the training run." + q = _quality_rows(m) + lines = [ + f"# {m['name']}", "", + f"Chronos v3 English↔Chinese **pretrained** model, packaged to run right after unpacking. {status}", "", + "## What is inside", "", + "| Item | Value |", "|---|---|", + f"| Stage | {m['model']['stage']} ({m['model']['which']} validation checkpoint) |", + f"| Training step | {_fmt(tr.get('step'))} of ~{_fmt(tr.get('planned_steps_at_export'))} planned |", + f"| Tokens seen | {_fmt(tr.get('tokens_seen'))} |", + f"| Held-out valid CE / PPL | {_fmt(tr.get('valid_ce'), 4)} / {_fmt(tr.get('valid_ppl'), 2)} (measured) |", + f"| Parameters | total {_fmt(params.get('total'))}, active {_fmt(params.get('active'))} per token |", + f"| Topology | hidden {arch['hidden_size']}, {arch['num_hidden_layers']} layers " + f"(prelude {arch['prelude_layers']}, coda {arch['coda_layers']}), recurrent depth {arch['recurrent_depth']} " + f"({arch['recurrent_attention_mode']}), {arch['num_experts']} experts top-{arch['num_experts_per_tok']} + " + f"{arch['num_shared_experts']} shared |", + f"| Vocabulary | {arch['vocab_size']:,} (own Chronos BPE tokenizer in `model/tokenizer/`) |", "", + "Prompt format (pretrain): `<|en|>` English `<|zh|>` → Chinese, or `<|zh|>` Chinese `<|en|>` → English. " + "The model has **not** been instruction-tuned: it translates single sentences, and long or unusual inputs " + "can loop or drop entities. A repetition penalty of 1.2 (the CLI default) suppresses most loops.", "", + ] + if q: + lines += [f"Quality on the first {m['quality']['first_n']} FLORES+ dev sentences ({m['quality']['label']}):", "", + "| Direction | Repetition penalty | n | BLEU | chrF |", "|---|---|---|---|---|", *q, ""] + lines += [ + "## Quick start", "", + "```sh", + "tar xzf " + m["name"] + ".tar.gz && cd " + m["name"], + "sh bin/setup.sh # .venv + pinned extras + `chronos bundle check`", + "bin/chronos bundle translate \"The weather is very nice today.\"", + "bin/chronos bundle translate --direction zh_en \"今天天气很好。\"", + "bin/chronos bundle translate --file examples/translate_en.txt", + "bin/chronos bundle translate --backend mlx --speculative --file examples/translate_en.txt", + "bin/chronos bundle eval --split devtest -n 100 # BLEU/chrF on FLORES+", + "```", "", + "`--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.", "", + "`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.", "", + "## Next stage: SFT (fine-tuning) on Apple silicon (MLX)", "", + "```sh", + "# 1. your pairs: .jsonl {\"en\":…, \"zh\":…} | .csv | .tsv | .parquet", + "bin/chronos bundle prepare-data my_pairs.jsonl --out work/data # clean, dedup, tokenize, shard", + "# add --unsegment-zh / --detokenize-en for word-segmented or Moses-tokenized corpora", + "# 2. fine-tune from the bundled weights (resumable; rerun the same command after an interruption)", + "bin/chronos bundle finetune --data work/data --run-dir work/run --tier quick", + "# smoke test of the whole path on the bundled examples (minutes):", + "bin/chronos bundle prepare-data examples/pairs_en_zh.jsonl --out work/smoke_data", + "bin/chronos bundle finetune --data work/smoke_data --run-dir work/smoke_run --tier smoke", + "```", "", + "`--mode sft` (default) trains the chat prompt `<|im_start|>user\\n翻译成中文，只输出译文：\\n…`; " + "`--mode pretrain` continues pretraining on new text; `--mode pretrain+sft` does both (40/60 time split). " + "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`.", "", + "Tiers (from `recommended.json`; every value is labelled with its source):", "", + "| Tier | Hours | SFT peak LR | Tokens / step | Warmup | Use |", "|---|---|---|---|---|---|", + *_tier_rows(recs), "", + "Torch (CUDA/CPU) path: `recommended.json → torch_stages` lists per-stage values for " + "`train_chronos_sft.py` and the DPO/ORPO/GRPO/distill scripts; pass `--from_weight model/.pth` " + "(the tokenizer next to it is picked up automatically, or set `--tokenizer_path model/tokenizer`).", "", + "## Serve (Ollama / OpenAI protocol)", "", + "```sh", + "bin/chronos bundle serve -- --port 11435 # loopback only, no auth", + "curl -s localhost:11435/api/generate -d '{\"model\":\"" + m["model"]["checkpoint"].split("/")[-1][:-4] + + "\",\"prompt\":\"<|en|>Good morning.<|zh|>\",\"raw\":true,\"stream\":false}'", + "```", "", + "The server binds 127.0.0.1 and has no authentication. Exposing it needs `--allow-remote` plus an API key " + "in `CHRONOS_API_KEY`.", "", + "## Files", "", + "`bundle.json` (identity, provenance, sha256 of every file) · `model/` (weights, config sidecar, tokenizer) · " + "`recipes/source_run_spec.json` (the exact training spec) · `recommended.json` (next-stage parameters) · " + "`eval/` (FLORES+) · `data/sft_valid/` (held-out SFT validation shards) · `examples/` · `third_party/` · " + "`wheels/` · `bin/`. Verify integrity any time with `bin/chronos bundle check`.", "", + "See `NOTICE.md` before redistributing: the weights are for **research use only**.", "", + ] + return "\n".join(lines) + + +def _readme_zh(m: dict, recs: dict) -> str: + tr, arch, params = m.get("training") or {}, m["model"]["architecture"], m["model"].get("params") or {} + status = "**中期检查点**：打包时训练仍在进行，训练结束后请用最终 run 重新打包。" if m.get("status") == "interim" \ + else "训练 run 的最终检查点。" + q = _quality_rows(m) + lines = [ + f"# {m['name']}", "", + f"Chronos v3 英↔中**预训练**模型，解压后即可使用。{status}", "", + "## 包内容", "", + "| 项目 | 数值 |", "|---|---|", + f"| 阶段 | {m['model']['stage']}（{m['model']['which']} 验证检查点） |", + f"| 训练步数 | {_fmt(tr.get('step'))} / 计划约 {_fmt(tr.get('planned_steps_at_export'))} |", + f"| 已见 token | {_fmt(tr.get('tokens_seen'))} |", + f"| 留出验证 CE / PPL | {_fmt(tr.get('valid_ce'), 4)} / {_fmt(tr.get('valid_ppl'), 2)}（实测） |", + f"| 参数量 | 总计 {_fmt(params.get('total'))}，每 token 激活 {_fmt(params.get('active'))} |", + f"| 结构 | hidden {arch['hidden_size']}，{arch['num_hidden_layers']} 层（prelude {arch['prelude_layers']}，" + f"coda {arch['coda_layers']}），循环深度 {arch['recurrent_depth']}（{arch['recurrent_attention_mode']}），" + f"{arch['num_experts']} 专家 top-{arch['num_experts_per_tok']} + {arch['num_shared_experts']} 共享专家 |", + f"| 词表 | {arch['vocab_size']:,}（自研 Chronos BPE 分词器，位于 `model/tokenizer/`） |", "", + "预训练格式：`<|en|>` 英文 `<|zh|>` → 中文，或 `<|zh|>` 中文 `<|en|>` → 英文。模型**未经指令微调**，" + "适合单句翻译；长句或少见输入可能出现重复或实体错误。命令行默认 repetition penalty 1.2，可抑制大部分重复。", "", + ] + if q: + lines += [f"FLORES+ dev 前 {m['quality']['first_n']} 句的质量（{m['quality']['label']}）：", "", + "| 方向 | repetition penalty | n | BLEU | chrF |", "|---|---|---|---|---|", *q, ""] + lines += [ + "## 快速开始", "", + "```sh", + "tar xzf " + m["name"] + ".tar.gz && cd " + m["name"], + "sh bin/setup.sh # 创建 .venv，安装固定版本依赖，并运行 bundle check", + "bin/chronos bundle translate \"The weather is very nice today.\"", + "bin/chronos bundle translate --direction zh_en \"今天天气很好。\"", + "bin/chronos bundle translate --file examples/translate_en.txt", + "bin/chronos bundle translate --backend mlx --speculative --file examples/translate_en.txt", + "bin/chronos bundle eval --split devtest -n 100 # FLORES+ 上的 BLEU/chrF", + "```", "", + "`--speculative`（循环深度自推测解码）的 greedy 输出与普通解码一致，接受率打印到 stderr。" + "它要求 `--repetition-penalty 1.0`（此时的默认值）。在这个规模的模型上用 CPU 实测，它比普通解码慢。", "", + "`bin/chronos` 会设置 `CHRONOS_BUNDLE` 和 `CHRONOS_MINIMIND_DIR`（指向 `third_party/` 里随包附带的 minimind），" + "全程不需要联网。不用 `bin/` 时，手动导出这两个变量后直接调用 `chronos bundle …`。", "", + "## 下一阶段：SFT 微调（Apple silicon / MLX）", "", + "```sh", + "# 1. 准备平行语料：.jsonl {\"en\":…, \"zh\":…} | .csv | .tsv | .parquet", + "bin/chronos bundle prepare-data my_pairs.jsonl --out work/data # 清洗、去重、分词、分片", + "# 中文已分词或英文经过 Moses tokenize 时，加 --unsegment-zh / --detokenize-en", + "# 2. 从包内权重开始微调（可断点续训：中断后重跑同一条命令）", + "bin/chronos bundle finetune --data work/data --run-dir work/run --tier quick", + "# 用包内示例数据把整条流程跑通（几分钟）：", + "bin/chronos bundle prepare-data examples/pairs_en_zh.jsonl --out work/smoke_data", + "bin/chronos bundle finetune --data work/smoke_data --run-dir work/smoke_run --tier smoke", + "```", "", + "`--mode sft`（默认）训练对话格式 `<|im_start|>user\\n翻译成中文，只输出译文：\\n…`；" + "`--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` 重新打包。", "", + "档位（来自 `recommended.json`，每个数值都标注了来源）：", "", + "| 档位 | 小时 | SFT 峰值 LR | 每步 token | Warmup | 用途 |", "|---|---|---|---|---|---|", + *_tier_rows(recs), "", + "Torch（CUDA/CPU）路径：`recommended.json → torch_stages` 给出 `train_chronos_sft.py` 及 DPO/ORPO/GRPO/" + "distill 脚本的分阶段参数；传入 `--from_weight model/.pth` 即可（同目录的分词器会被自动识别，" + "也可以显式指定 `--tokenizer_path model/tokenizer`）。", "", + "## 服务（Ollama / OpenAI 协议）", "", + "```sh", + "bin/chronos bundle serve -- --port 11435 # 只监听本机，无认证", + "curl -s localhost:11435/api/generate -d '{\"model\":\"" + m["model"]["checkpoint"].split("/")[-1][:-4] + + "\",\"prompt\":\"<|en|>Good morning.<|zh|>\",\"raw\":true,\"stream\":false}'", + "```", "", + "服务默认绑定 127.0.0.1，没有认证。对外开放需要 `--allow-remote`，并通过环境变量 `CHRONOS_API_KEY` 提供 API key。", "", + "## 文件", "", + "`bundle.json`（身份、来源、每个文件的 sha256）· `model/`（权重、配置 sidecar、分词器）· " + "`recipes/source_run_spec.json`（产生这份权重的训练配置）· `recommended.json`（下一阶段参数）· " + "`eval/`（FLORES+）· `data/sft_valid/`（留出的 SFT 验证分片）· `examples/` · `third_party/` · `wheels/` · " + "`bin/`。随时可用 `bin/chronos bundle check` 校验完整性。", "", + "再分发前请阅读 `NOTICE.md`：权重**仅供研究使用**。", "", + ] + return "\n".join(lines) + + +def _notice(m: dict) -> str: + mm = (m.get("third_party") or {}).get("minimind") or {} + return "\n".join([ + "# NOTICE", "", + "## Weights: research use only", "", + "The weights were trained on parallel text whose licences are not all known:", "", + "| Source | Use | Licence as recorded at download |", "|---|---|---|", + "| WMT19 zh-en (HF `wmt/wmt19`, news-commentary, UN, CWMT, wikititles) | training | unknown (dataset card); " + "research use only |", + "| OPUS-100 en-zh (HF `Helsinki-NLP/opus-100`) | training | unknown (dataset card) |", + "| iic WMT Chinese-English corpus (ModelScope) | training | not recorded |", + "| FLORES+ dev/devtest (shipped in `eval/`) | evaluation only; removed from training | CC BY-SA 4.0 |", "", + "Because the training data licences are unknown, these weights are distributed **for research and " + "evaluation only**. Do not use them commercially or redistribute them as a product without checking the " + "licences of the underlying corpora yourself.", "", + "## Code", "", + "- Chronos (`wheels/`, `bin/`): Apache-2.0.", + f"- minimind (`third_party/minimind/`, commit `{mm.get('commit') or 'unknown'}`): Apache-2.0, " + "© jingyaogong. Its `LICENSE` file is included.", + "- FLORES+ files in `eval/` keep their CC BY-SA 4.0 licence (attribution: the FLORES+ / OLDI project).", "", + "## No warranty", "", + "Provided as is. Machine translation output can be wrong; check it before relying on it.", "", + ]) + + +def write_docs(bundle_dir: str, manifest: dict, recs: dict | None = None) -> None: + if recs is None: + with open(os.path.join(bundle_dir, "recommended.json"), encoding="utf-8") as fh: + recs = json.load(fh) + for name, text in (("README.md", _readme_en(manifest, recs)), ("README_zh.md", _readme_zh(manifest, recs)), + ("NOTICE.md", _notice(manifest))): + with open(os.path.join(bundle_dir, name), "w", encoding="utf-8") as fh: + fh.write(text) + bin_dir = os.path.join(bundle_dir, "bin") + os.makedirs(bin_dir, exist_ok=True) + for name, text in (("chronos", LAUNCHER), ("setup.sh", SETUP)): + path = os.path.join(bin_dir, name) + with open(path, "w", encoding="utf-8") as fh: + fh.write(text) + os.chmod(path, 0o755) diff --git a/chronos/bundle/finetune.py b/chronos/bundle/finetune.py new file mode 100644 index 0000000..3280ea2 --- /dev/null +++ b/chronos/bundle/finetune.py @@ -0,0 +1,128 @@ +"""Recipient-side data preparation and fine-tune launch for a bundle.""" +from __future__ import annotations + +import json +import os + +from .builder import load_bundle_manifest + +KINDS = {".jsonl": "pairs_jsonl", ".json": "pairs_jsonl", ".csv": "pairs_csv", ".tsv": "pairs_tsv", + ".parquet": "pairs_parquet"} + + +def _ensure_valid_split(pairs_dir: str, min_valid: int, log=print) -> dict: + """Small user sets can hash every row into train; move ``min_valid`` rows to valid.""" + import pyarrow.parquet as pq + names = sorted(os.listdir(pairs_dir)) + train = [n for n in names if n.startswith("train.") and n.endswith(".parquet")] + valid = [n for n in names if n.startswith("valid.") and n.endswith(".parquet")] + rows = {"train": sum(pq.ParquetFile(os.path.join(pairs_dir, n)).metadata.num_rows for n in train), + "valid": sum(pq.ParquetFile(os.path.join(pairs_dir, n)).metadata.num_rows for n in valid)} + if rows["valid"] >= min_valid or not train: + return rows + need = min(min_valid - rows["valid"], max(1, rows["train"] // 10)) + name = train[-1] + table = pq.read_table(os.path.join(pairs_dir, name)) + need = min(need, max(0, table.num_rows - 1)) + if need <= 0: + return rows + keep, move = table.slice(0, table.num_rows - need), table.slice(table.num_rows - need) + pq.write_table(move, os.path.join(pairs_dir, "valid." + name[len("train."):]), compression="zstd") + pq.write_table(keep, os.path.join(pairs_dir, name), compression="zstd") + log(f"[prepare] moved {need} pairs to the validation split (set too small for the hash split)") + return {"train": rows["train"] - need, "valid": rows["valid"] + need} + + +def prepare_data(bundle_dir: str, inputs: list[str], out_dir: str, *, zh_field: str = "zh", en_field: str = "en", + unsegment_zh: bool = False, detokenize_en: bool = False, workers: int = 4, + valid_per_mille: int | None = None, min_free_gb: float = 1.0, log=print) -> dict: + """User parallel data → cleaned pairs → MLX shards tokenized with the bundle tokenizer. + + ``inputs`` are files (``.jsonl``/``.csv``/``.tsv``/``.parquet``) with a + Chinese and an English column. FLORES+ (shipped in the bundle) is removed + from the data so the bundled evaluation stays held out.""" + from chronos.data.corpus.clean import CleanConfig + from chronos.data.corpus.evalsets import load_eval_sets + from chronos.data.corpus.pipeline import PairSource, build_corpus + from chronos.data.corpus.readers import count_rows + from chronos.data.corpus.shards import build_shards + bundle_dir, out_dir = os.path.abspath(bundle_dir), os.path.abspath(out_dir) + manifest = load_bundle_manifest(bundle_dir) + if os.path.exists(os.path.join(out_dir, "shards", "shards_manifest.json")): + raise FileExistsError(f"{out_dir} already holds prepared shards; pick another --out") + sources, total = [], 0 + for i, path in enumerate(inputs): + path = os.path.abspath(os.path.expanduser(path)) + kind = KINDS.get(os.path.splitext(path)[1].lower()) + if kind is None: + raise ValueError(f"{path}: unsupported file type (use {sorted(KINDS)})") + n = count_rows(path, kind) + total += n + clean = CleanConfig(unsegment_zh=unsegment_zh, detokenize_en=detokenize_en) + sources.append(PairSource(name=f"user{i:02d}", path=path, kind=kind, fields=(zh_field, en_field), + clean=clean, priority=i, rows_per_unit=200_000)) + if valid_per_mille is None: # ≈1 % on small sets, the recipe's 0.2 % on large ones + valid_per_mille = 10 if total < 1_000_000 else 2 + eval_pairs = load_eval_sets(os.path.join(bundle_dir, "eval"), ["flores_dev", "flores_devtest"]) + corpus_dir = os.path.join(out_dir, "corpus") + corpus = build_corpus(sources, eval_pairs, corpus_dir, workers=workers, valid_per_mille=valid_per_mille, + min_free_gb=min_free_gb, log=log) + split_rows = _ensure_valid_split(os.path.join(corpus_dir, "pairs"), min_valid=min(200, max(4, total // 20)), + log=log) + tok_dir = os.path.join(bundle_dir, manifest["tokenizer"]["dir"]) + shards = build_shards(os.path.join(corpus_dir, "pairs"), tok_dir, os.path.join(out_dir, "shards"), + workers=workers, log=log) + report = {"inputs": inputs, "rows_in": total, "split_rows": split_rows, "valid_per_mille": valid_per_mille, + "corpus": {k: corpus.get(k) for k in ("rows_after_clean", "duplicates", "eval_leak_removed", + "final_train", "final_valid", "clean_reasons")}, + "shards": {k: v for k, v in shards.items() if k.endswith(("_tokens", "_records"))}} + with open(os.path.join(out_dir, "prepare_report.json"), "w", encoding="utf-8") as fh: + json.dump(report, fh, ensure_ascii=False, indent=2) + log(f"[prepare] {json.dumps(report['shards'])}") + return report + + +def _check_shards(shards_dir: str, phases: list[dict]) -> None: + path = os.path.join(shards_dir, "shards_manifest.json") + if not os.path.exists(path): + raise FileNotFoundError(f"{shards_dir} has no shards_manifest.json; run `chronos bundle prepare-data` first") + with open(path, encoding="utf-8") as fh: + m = json.load(fh) + for phase in phases: + if phase["name"] == "pretrain": + need = int(phase.get("seq_len", 512)) + 2 + for split in ("train", "valid"): + if int(m.get(f"pretrain_{split}_tokens", 0)) < need: + raise ValueError(f"pretrain/{split} has {m.get(f'pretrain_{split}_tokens', 0)} tokens, fewer " + f"than one {need}-token window; use --mode sft or add data") + else: + for split in ("train", "valid"): + if int(m.get(f"sft_{split}_records", 0)) < 1: + raise ValueError(f"sft/{split} is empty; add data") + + +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``.""" + 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) + _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") + if os.path.exists(os.path.join(spec["run_dir"], "checkpoints", "latest.state.pkl")): + # longrun resumes from latest.state.pkl; a different spec would mix two runs. + raise FileExistsError(f"{spec['run_dir']} already holds a run; resume it with " + f"`chronos longrun --spec {path}` or choose a new --run-dir") + with open(path, "w", encoding="utf-8") as fh: + json.dump(spec, fh, ensure_ascii=False, indent=2) + return path + + +def launch(spec_path: str, budget_hours: float | None = None) -> int: + from chronos.mlx.training.longrun import main as longrun_main + argv = ["--spec", spec_path] + (["--budget-hours", str(budget_hours)] if budget_hours is not None else []) + return longrun_main(argv) diff --git a/chronos/bundle/recommend.py b/chronos/bundle/recommend.py new file mode 100644 index 0000000..cb97842 --- /dev/null +++ b/chronos/bundle/recommend.py @@ -0,0 +1,182 @@ +"""Next-stage parameter recommendations shipped with a bundle. + +Every value carries a ``source``: + +* ``recipe`` – copied from the run's own RunSpec (what produced these weights); +* ``planner`` – ``chronos.model.planner`` output for the same topology + (``ui/presets_v3.json``); +* ``rule`` – a stated derivation from one of the above (the rule is + written next to the value), not a measurement. + +Nothing here is a measured optimum. The tiers trade wall time for quality +on a single Apple-silicon or CUDA machine; ``chronos bundle finetune --tier`` +turns a tier into a ``chronos longrun`` spec. +""" +from __future__ import annotations + +import copy +import json +import os + +# Wall-clock budgets per tier (hours) and the reason for each. +TIERS = { + "smoke": {"hours": 0.05, "why": "proves the pipeline end to end on the bundled examples (minutes)"}, + "quick": {"hours": 2.0, "why": "a small domain set (≈10k–100k pairs); one evening"}, + "standard": {"hours": 12.0, "why": "a medium set (≈0.1M–1M pairs); overnight"}, + "full": {"hours": 36.0, "why": "the SFT half of the original 72 h recipe"}, +} + +# Phase keys a longrun PhaseSpec accepts. +PHASE_KEYS = ("name", "budget_fraction", "peak_lr", "tokens_per_step", "micro_tokens", "seq_len", + "warmup_fraction", "min_lr_ratio", "weight_decay", "grad_clip") + + +def _phase(recipe: dict, name: str) -> dict | None: + for phase in recipe.get("phases", []): + if phase.get("name") == name: + return {k: phase[k] for k in PHASE_KEYS if k in phase} + return None + + +def _labelled(values: dict, source: str, note: str | None = None) -> dict: + out = {"values": values, "source": source} + if note: + out["note"] = note + return out + + +def _planner_stages(planner_preset: dict | None) -> dict: + stages = (planner_preset or {}).get("stage_overrides") or {} + return {k: dict(v) for k, v in stages.items()} + + +def continued_pretrain_phase(recipe: dict) -> dict: + """Re-warm the pretrain phase at 0.3× its peak LR. + + The bundled weights finished (or are partway through) a cosine decay; + restarting at the full peak would undo part of it. 0.3× is a rule of + thumb, not a tuned value.""" + phase = _phase(recipe, "pretrain") or {} + phase = dict(phase, budget_fraction=1.0, peak_lr=round(float(phase.get("peak_lr", 1.5e-3)) * 0.3, 8), + warmup_fraction=max(0.01, float(phase.get("warmup_fraction", 0.02)))) + return phase + + +def sft_phase(recipe: dict) -> dict: + phase = _phase(recipe, "sft") + if phase is None: # a pretrain-only recipe: fall back to the 72 h SFT values + phase = {"name": "sft", "peak_lr": 5e-4, "tokens_per_step": 65536, "micro_tokens": 8192, + "warmup_fraction": 0.01, "min_lr_ratio": 0.05, "weight_decay": 0.1, "grad_clip": 1.0} + return dict(phase, budget_fraction=1.0) + + +def tier_phases(recipe: dict, tier: str, mode: str = "sft") -> list[dict]: + """PhaseSpec dicts for one tier. + + ``mode``: ``sft`` (translation fine-tune), ``pretrain`` (continued + pretraining on new text) or ``pretrain+sft`` (40 % / 60 % split).""" + if tier not in TIERS: + raise ValueError(f"unknown tier {tier!r}; choose one of {sorted(TIERS)}") + sft, cpt = sft_phase(recipe), continued_pretrain_phase(recipe) + if tier == "smoke": + # Tiny steps so a CPU finishes in minutes; LR as in the recipe. + sft = dict(sft, tokens_per_step=8192, micro_tokens=2048, warmup_fraction=0.1) + cpt = dict(cpt, tokens_per_step=8192, micro_tokens=2048, seq_len=min(int(cpt.get("seq_len", 512)), 256), + warmup_fraction=0.1) + elif tier == "quick": + # Small data: halve the step size so a 2 h budget still makes enough + # updates, and warm up longer (rule: ≥ ~100 warmup steps). + sft = dict(sft, tokens_per_step=max(16384, int(sft["tokens_per_step"]) // 2), warmup_fraction=0.03) + cpt = dict(cpt, tokens_per_step=max(16384, int(cpt.get("tokens_per_step", 131072)) // 2), + warmup_fraction=0.03) + if mode == "sft": + return [sft] + if mode == "pretrain": + return [cpt] + if mode == "pretrain+sft": + return [dict(cpt, budget_fraction=0.4), dict(sft, budget_fraction=0.6)] + raise ValueError(f"unknown mode {mode!r}; choose sft, pretrain or pretrain+sft") + + +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.""" + from .builder import load_bundle_manifest + manifest = load_bundle_manifest(bundle_dir) + bundle_dir = os.path.abspath(bundle_dir) + smoke = tier == "smoke" + 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), + "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, + "validate_every_s": 60.0 if smoke else 1800.0, + "checkpoint_every_s": 120.0 if smoke else 1800.0, + "bleu_samples": 8 if smoke else 100, + "final_eval_samples": 16 if smoke else 1012, + "probe_samples": 4 if smoke else 16, + "tensor_stats_every_s": 0.0 if smoke else 1800.0, + "seed": int(seed if seed is not None else recipe.get("seed", 1234)), + "dtype": recipe.get("dtype", "bfloat16"), + "schedule": "time", + "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}, + } + return spec + + +def build_recommendations(recipe: dict, planner_preset: dict | None, decode: dict | None = None) -> dict: + """The ``recommended.json`` payload.""" + planner = _planner_stages(planner_preset) + tiers = {} + for tier, meta in TIERS.items(): + tiers[tier] = { + "hours": meta["hours"], "why": meta["why"], + "phases": {mode: tier_phases(recipe, tier, mode) for mode in ("sft", "pretrain", "pretrain+sft")}, + "source": "rule", + "rule": "SFT phase = recipe SFT phase; continued pretrain = recipe pretrain phase at 0.3x peak LR; " + "smoke uses 8k-token steps, quick halves tokens_per_step and warms up 3%", + } + torch_stages = {} + for stage, values in planner.items(): + torch_stages[stage] = _labelled(values, "planner", + "per-stage values for train_chronos_.py / the Web UI; " + "batch_size counts sequences of max_seq_len tokens") + out = { + "about": "Next-stage parameters for this bundle. 'recipe' = what produced these weights; " + "'planner' = chronos.model.planner for the same topology; 'rule' = derived as stated; " + "'measured' = measured on the bundled weights (conditions given). None is a tuned optimum.", + "mlx_longrun": { + "tiers": tiers, + "recipe_phases": _labelled({p["name"]: {k: p[k] for k in PHASE_KEYS if k in p} + 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"), + }, + "torch_stages": torch_stages, + "data": _labelled({ + "pretrain_mix": {"mono": 0.5, "en_zh": 0.25, "zh_en": 0.25}, + "sft_max_len": 512, + "clean": {"max_chars": 1024, "max_len_ratio": 3.0, "min_han_ratio": 0.3, "min_latin_ratio": 0.4, + "zh_fullwidth_punct": True}, + "min_pairs_hint": "below ~5k pairs prefer tier 'quick' with mode 'sft'; the 50M model overfits " + "small sets quickly — watch valid_ce in metrics.jsonl", + }, "recipe", "defaults of chronos.data.corpus (build_corpus / build_shards) used for the bundled weights"), + } + if decode: + out["decode"] = decode + return out + + +def write_json(path: str, payload: dict) -> str: + os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + json.dump(payload, fh, ensure_ascii=False, indent=2) + fh.write("\n") + return path diff --git a/chronos/bundle/runtime.py b/chronos/bundle/runtime.py new file mode 100644 index 0000000..818d042 --- /dev/null +++ b/chronos/bundle/runtime.py @@ -0,0 +1,204 @@ +"""Load a bundle's model and translate with it (Torch: CPU/MPS/CUDA, or MLX). + +Pretrain bundles use the ``<|en|>src<|zh|>`` (or reverse) format they were +trained on; SFT bundles use the chat prompt. Decoding is greedy with a KV +cache; ``repetition_penalty`` > 1 suppresses the loops a pretrain-only model +falls into on long inputs. ``speculative=True`` uses recurrent-depth +self-speculative decoding (same greedy output, fewer full-depth forwards); +it needs ``repetition_penalty == 1`` because the penalty changes the target +distribution position by position. +""" +from __future__ import annotations + +import os +import time + +STOP_TOKENS = ("<|endoftext|>", "<|im_end|>", "<|en|>", "<|zh|>") +LANG_TAGS = {"en_zh": ("<|en|>", "<|zh|>"), "zh_en": ("<|zh|>", "<|en|>")} + + +class BundleModel: + """Weights + tokenizer of one bundle, ready for greedy translation.""" + + def __init__(self, bundle_dir: str, *, device: str = "cpu", dtype: str = "float32", backend: str = "torch"): + import torch + from tokenizers import Tokenizer + from chronos.model.checkpoint import chronos_config_from_checkpoint, load_checkpoint_into_model + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.precision import place_torch_model + from .builder import load_bundle_manifest + self.bundle_dir = os.path.abspath(bundle_dir) + self.manifest = load_bundle_manifest(self.bundle_dir) + self.checkpoint = os.path.join(self.bundle_dir, self.manifest["model"]["checkpoint"]) + tok_dir = os.path.join(self.bundle_dir, self.manifest["tokenizer"]["dir"]) + self.tokenizer = Tokenizer.from_file(os.path.join(tok_dir, "tokenizer.json")) + 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") + model = ChronosForCausalLM(self.config) + load_checkpoint_into_model(model, self.checkpoint, map_location="cpu") + self.backend, self.last_stats = backend, {} + if backend == "mlx": + import mlx.core as mx + from chronos.mlx.model import ChronosMLXModel + from chronos.precision import resolve_mlx + 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.precision = resolve_mlx(dtype, "cpu" if device == "cpu" else "gpu") + self.model = ChronosMLXModel.from_chronos_pytorch(model.eval(), self.config, + dtype=self.precision.compute).eval() + del model + self.device = "cpu" if device == "cpu" else "gpu" + elif backend == "torch": + if device == "auto": + device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu") + self.model, self.precision = place_torch_model(model, device, dtype) + self.model.eval() + 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} + + # ---- prompts --------------------------------------------------------------- + def directions(self) -> list[str]: + return ["en_zh", "zh_en"] if self.prompt_format == "pretrain_tags" else ["en_zh"] + + def prompt_ids(self, src: str, direction: str = "en_zh") -> list[int]: + tok = self.tokenizer + if direction not in self.directions(): + raise ValueError(f"this bundle supports {self.directions()}, not {direction!r}") + if self.prompt_format == "pretrain_tags": + a, b = LANG_TAGS[direction] + return [tok.token_to_id(a)] + tok.encode(src, add_special_tokens=False).ids + [tok.token_to_id(b)] + from chronos.data.corpus.shards import SFT_PROMPT_EN_ZH + text = f"<|im_start|>user\n{SFT_PROMPT_EN_ZH}{src}<|im_end|>\n<|im_start|>assistant\n" + return tok.encode(text, add_special_tokens=False).ids + + # ---- decoding -------------------------------------------------------------- + def greedy(self, prompt: list[int], limit: int, repetition_penalty: float = 1.0, *, speculative: bool = False, + draft_depth: int = 1, gamma: int = 4) -> list[int]: + """Greedy continuation of ``prompt`` (stop tokens excluded), at most ``limit`` tokens.""" + 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: + out = self._speculative(prompt, limit, draft_depth, gamma) + elif self.backend == "mlx": + out = self._greedy_mlx(prompt, limit, repetition_penalty) + else: + out = self._greedy_torch(prompt, limit, repetition_penalty) + for i, t in enumerate(out): + if t in self.stop_ids: + return out[:i] + return out + + def _greedy_torch(self, prompt, limit, repetition_penalty): + import torch + from chronos.model.model_chronos import apply_repetition_penalty + ids = torch.tensor([prompt], dtype=torch.long, device=self.device) + seen = ids + out: list[int] = [] + cache = None + with torch.inference_mode(): + for _ in range(limit): + res = self.model(ids, past_key_values=cache, use_cache=True) + cache = res.past_key_values + logits = res.logits[:, -1, :].float() + if repetition_penalty != 1.0: + logits = apply_repetition_penalty(logits, seen, repetition_penalty) + nxt = int(torch.argmax(logits, dim=-1).item()) + if nxt in self.stop_ids: + break + out.append(nxt) + ids = torch.tensor([[nxt]], dtype=torch.long, device=self.device) + seen = torch.cat([seen, ids], dim=1) + return out + + def _greedy_mlx(self, prompt, limit, repetition_penalty): + import mlx.core as mx + import numpy as np + seen = np.zeros(int(self.config.vocab_size), dtype=bool) + seen[prompt] = True + out: list[int] = [] + logits, _, cache = self.model(mx.array([prompt], dtype=mx.int32)) + for _ in range(limit): + row = logits[0, -1].astype(mx.float32) + if repetition_penalty != 1.0: # CTRL-style, same rule as the Torch helper + p = float(repetition_penalty) + row = mx.where(mx.array(seen), mx.where(row < 0, row * p, row / p), row) + nxt = int(mx.argmax(row).item()) + if nxt in self.stop_ids: + break + out.append(nxt) + seen[nxt] = True + logits, _, cache = self.model(mx.array([[nxt]], dtype=mx.int32), cache=cache) + return out + + def _speculative(self, prompt, limit, draft_depth, gamma): + eos = sorted(self.stop_ids) + if self.backend == "mlx": + from chronos.mlx.speculative import speculative_generate + tokens, stats = speculative_generate(self.model, prompt, max_new_tokens=limit, draft_depth=draft_depth, + gamma=gamma, eos_token_id=eos) + else: + import torch + from chronos.runtime.speculative import speculative_generate + with torch.inference_mode(): + ids, stats = speculative_generate(self.model, torch.tensor([prompt], device=self.device), + max_new_tokens=limit, draft_depth=draft_depth, gamma=gamma, + eos_token_id=eos) + tokens = ids[0].tolist() + agg = self.last_stats.setdefault("speculative", {"proposed": 0, "accepted": 0, "verify_calls": 0, + "draft_calls": 0, "generated": 0}) + for k in agg: + agg[k] += int(getattr(stats, k)) + return tokens[len(prompt):] + + def translate(self, sources: list[str], direction: str = "en_zh", *, max_new_tokens: int = 256, + repetition_penalty: float = 1.0, speculative: bool = False, draft_depth: int = 1, + gamma: int = 4) -> list[str]: + out = [] + self.last_stats = {} + for src in sources: + src_len = len(self.tokenizer.encode(src, add_special_tokens=False).ids) + limit = min(int(max_new_tokens), int(1.6 * src_len) + 16) # output ≲ 1.6× source tokens + 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)) + spec = self.last_stats.get("speculative") + if spec: + spec["acceptance_rate"] = round(spec["accepted"] / spec["proposed"], 4) if spec["proposed"] else 0.0 + spec["tokens_per_verify"] = round(spec["generated"] / max(spec["verify_calls"], 1), 3) + return out + + +def eval_pairs(bundle_dir: str, split: str = "dev") -> list[tuple[str, str]]: + """FLORES+ (en, zh) pairs shipped in ``/eval``.""" + from chronos.data.corpus.evalsets import load_eval_sets + name = "flores_dev" if split == "dev" else "flores_devtest" + return load_eval_sets(os.path.join(os.path.abspath(bundle_dir), "eval"), [name]).get(name, []) + + +def score(bm: BundleModel, pairs: list[tuple[str, str]], *, directions=None, repetition_penalty: float = 1.0, + samples: int = 3) -> dict: + """Corpus BLEU/chrF with sacrebleu (BLEU tokenize=zh for →zh, 13a for →en).""" + import sacrebleu + out, t0 = {}, time.monotonic() + for direction in directions or bm.directions(): + si, ti = (0, 1) if direction == "en_zh" else (1, 0) + hyps = bm.translate([p[si] for p in pairs], direction, repetition_penalty=repetition_penalty) + refs = [p[ti] for p in pairs] + bleu = sacrebleu.corpus_bleu(hyps, [refs], tokenize="zh" if ti == 1 else "13a") + chrf = sacrebleu.corpus_chrf(hyps, [refs]) + out[direction] = {"n": len(pairs), "bleu": round(bleu.score, 2), "chrf": round(chrf.score, 2), + "samples": [{"src": p[si], "ref": p[ti], "hyp": h} + for p, h in list(zip(pairs, hyps))[:samples]]} + out["seconds"] = round(time.monotonic() - t0, 1) + out["decode"] = {"greedy": True, "kv_cache": True, "repetition_penalty": repetition_penalty, + "max_new_tokens": "min(256, 1.6*src_tokens+16)", "device": bm.device, + "dtype": getattr(bm.precision, "as_dict", lambda: str(bm.precision))()} + return out diff --git a/chronos/cli.py b/chronos/cli.py index 3548bb0..35514d9 100644 --- a/chronos/cli.py +++ b/chronos/cli.py @@ -8,6 +8,12 @@ chronos diagnose — checkpoint topology + offload/predictive-routing report chronos export — export full checkpoints or expert SSD cluster layout chronos mac — Apple Silicon backend diagnostics + chronos migrate — explicitly migrate a v1 checkpoint to v2 + chronos serve — Ollama/OpenAI-compatible server (native runtime; optional /monitor page) + chronos tokenizer — train a Chronos tokenizer from selected datasets + 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 """ import sys import os @@ -69,6 +75,25 @@ def cmd_mac(args): print(json.dumps(mac_backend_diagnostics(configure_threads=configure), indent=2, ensure_ascii=False)) +def cmd_migrate(args): + from chronos.migrate import main + main(args) + + +def _module_main(module: str, rest): + import importlib + raise SystemExit(importlib.import_module(module).main(rest)) + + +def _cmd_corpus(rest): + """chronos corpus {download|build|shards} ...""" + if not rest or rest[0] not in {"download", "build", "shards"}: + raise SystemExit("usage: chronos corpus {download|build|shards} [options]") + module = {"download": "chronos.data.corpus.sources", "build": "chronos.data.corpus.presets", + "shards": "chronos.data.corpus.shards"}[rest[0]] + _module_main(module, rest[1:]) + + def main(): parser = argparse.ArgumentParser( prog='chronos', @@ -82,8 +107,16 @@ def main(): sub.add_parser('diagnose', help='Checkpoint/offload diagnostic', add_help=False) sub.add_parser('export', help='Export checkpoint or expert cluster layout', add_help=False) sub.add_parser('mac', help='Apple Silicon backend diagnostics', add_help=False) + sub.add_parser('migrate', help='Explicitly migrate a v1 checkpoint to v2', add_help=False) + sub.add_parser('serve', help='Ollama/OpenAI-compatible server (native Chronos runtime)', add_help=False) + sub.add_parser('tokenizer', help='Train a Chronos tokenizer from selected datasets', add_help=False) + sub.add_parser('corpus', help='Download/clean/dedup/shard training corpora', add_help=False) + 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) known, rest = parser.parse_known_args() + longrun = lambda rest: _module_main('chronos.mlx.training.longrun', rest) # noqa: E731 dispatch = { 'train': cmd_train, 'eval': cmd_eval, @@ -91,6 +124,14 @@ def main(): 'diagnose': cmd_diagnose, 'export': cmd_export, 'mac': cmd_mac, + 'migrate': cmd_migrate, + 'serve': lambda rest: _module_main('chronos.serving.ollama_server', rest), + 'tokenizer': lambda rest: _module_main('chronos.tokenizer', rest), + 'corpus': _cmd_corpus, + 'longrun': longrun, + 'run': longrun, + 'monitor': lambda rest: _module_main('chronos.monitor.server', rest), + 'bundle': lambda rest: _module_main('chronos.bundle.cli', rest), } dispatch[known.command](rest) diff --git a/chronos/config_schema.py b/chronos/config_schema.py new file mode 100644 index 0000000..e399c0e --- /dev/null +++ b/chronos/config_schema.py @@ -0,0 +1,224 @@ +"""Canonical Chronos v2 configuration mapping. + +Every boundary (UI, CLI, checkpoint, HF config, GGUF metadata) must use this +module rather than maintaining another positional list of model fields. +""" +from __future__ import annotations + +import hashlib +import json +import os +import math +from typing import Any, Mapping + + +SCHEMA_VERSION = 3 +ARCHITECTURE_VERSION = 3 +LEGACY_ARCHITECTURE_VERSION = 2 + +MODEL_FIELDS = ( + "architecture_version", "hidden_size", "num_hidden_layers", "vocab_size", + "num_attention_heads", "num_key_value_heads", "head_dim", "intermediate_size", + "max_position_embeddings", "tie_word_embeddings", "dropout", "rms_norm_eps", + "rope_theta", "flash_attn", "bos_token_id", "eos_token_id", + "recurrent_enabled", "recurrent_groups", "recurrent_depth", "min_recurrent_depth", + "prelude_layers", "coda_layers", "adaptive_halting", "halting_threshold", + "halting_epsilon", "max_inference_depth", "ponder_loss_coef", + "recurrent_distill_coef", "recurrent_stability_coef", + "recurrent_contraction_coef", "recurrent_contraction_budget", + "recurrent_stability_threshold", "recurrent_dt_init", "recurrent_dt_min", + "recurrent_dt_max", "recurrent_latent_refinement", "recurrent_attention_mode", + "recurrent_strategy", + "recurrent_depth_jitter", "recurrent_parallel_refine", + "recurrent_gradient_checkpoint_groups", "depth_schedule", + "depth_schedule_warmup_fraction", + "use_moe", "num_experts", "num_experts_per_tok", "num_shared_experts", + "moe_intermediate_size", "norm_topk_prob", "router_aux_loss_coef", + "lookahead_steps", "layer_aware_lookahead", "layer_aware_cache", + "lambda_balance", "lambda_temporal", "lambda_lookahead", + "lambda_lookahead_topk", "lambda_lookahead_union", "lambda_router_locality", + "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", + "draft_exit_kl_coef", "draft_exit_depth", "max_resident_experts", "ram_budget_mb", + "vram_budget_gb", "pinned_memory_max_fraction", "prefetch_depth", + "recommended_resident_experts", "offload_working_set_experts", "storage_format", + "cluster_manifest_path", "training_paging", "optimizer_paging", + "optimizer_page_dir", "page_size_mb", "activation_checkpointing", + "modalities", "auto_detect_modalities", "image_size", "image_patch_size", + "max_image_tokens", "audio_sample_rate", "max_audio_seconds", "max_audio_tokens", + "modality_dropout", "freeze_modality_encoders", "image_encoder_type", + "audio_encoder_type", "modality_fusion_heads", "media_error_policy", + "modality_adapters", "modality_alignment_coef", +) + +RUNTIME_FIELDS = ( + "dtype", "learning_rate", "batch_size", "accumulation_steps", "epochs", + "max_steps", "save_interval", "log_interval", "save_dir", "weight_decay", + "grad_clip", "cpu_threads", "cpu_budget_percent", "num_workers", "reward_spec", + "max_gen_len", "num_generations", "temperature", "beta", "alpha", "lambda_or", +) + +CANONICAL_FIELDS = frozenset(MODEL_FIELDS + RUNTIME_FIELDS + ("max_seq_len", "extra_metadata")) + + +def _plain(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, (list, tuple)): + return [_plain(v) for v in value] + if isinstance(value, Mapping): + return {str(k): _plain(v) for k, v in value.items()} + return str(value) + + +def canonical_config_dict(source: Any, *, preserve_unknown: bool = True) -> dict: + """Return a versioned, JSON-safe configuration with no silent drops.""" + if isinstance(source, Mapping): + raw = dict(source) + elif hasattr(source, "to_dict"): + raw = dict(source.to_dict()) + else: + raw = dict(vars(source)) + if "max_seq_len" in raw and "max_position_embeddings" not in raw: + raw["max_position_embeddings"] = raw["max_seq_len"] + known = {key: _plain(raw[key]) for key in CANONICAL_FIELDS if key in raw} + extra = dict(known.pop("extra_metadata", {}) or {}) + if preserve_unknown: + for key, value in raw.items(): + if key not in CANONICAL_FIELDS and not key.startswith("_"): + extra[key] = _plain(value) + # A config without an explicit version predates v3 metadata; treat it + # as legacy so its recurrent numerics never change silently. + known["architecture_version"] = int(known.get("architecture_version", LEGACY_ARCHITECTURE_VERSION)) + known["schema_version"] = SCHEMA_VERSION + known["extra_metadata"] = extra + return known + + +def model_config_kwargs(source: Any) -> dict: + canonical = canonical_config_dict(source) + out = {key: canonical[key] for key in MODEL_FIELDS if key in canonical} + out["extra_metadata"] = canonical.get("extra_metadata", {}) + for key in ("intermediate_size", "moe_intermediate_size"): + if out.get(key) in (None, 0, ""): + out.pop(key, None) + return out + + +def normalize_model_dimensions(source: Mapping[str, Any]) -> dict: + """Resolve auto widths and the physical recurrent partition without Torch.""" + out = dict(source) + h, layers = int(out["hidden_size"]), int(out["num_hidden_layers"]) + if h < 1 or layers < 1: + raise ValueError("hidden_size and num_hidden_layers must be positive") + heads = int(out["num_attention_heads"]) + if heads < 1: + raise ValueError("num_attention_heads must be positive") + out["head_dim"] = int(out.get("head_dim") or h // heads) + out["num_key_value_heads"] = int(out.get("num_key_value_heads") or heads) + if out["head_dim"] < 1 or heads % out["num_key_value_heads"]: + raise ValueError("attention heads must divide evenly by KV heads, with head_dim >= 1") + out["intermediate_size"] = int(out.get("intermediate_size") or math.ceil(h*math.pi/64)*64) + out["moe_intermediate_size"] = int(out.get("moe_intermediate_size") or out["intermediate_size"]) + if out.get("use_moe", True) and not 1 <= int(out["num_experts_per_tok"]) <= int(out["num_experts"]): + raise ValueError("num_experts_per_tok must be between 1 and num_experts") + p = max(0, min(int(out.get("prelude_layers", 0)), layers-1)) + c = max(0, min(int(out.get("coda_layers", 0)), layers-p-1)) + g = max(1, min(int(out.get("recurrent_groups") or layers-p-c), layers-p-c)) + depth = max(1, int(out.get("recurrent_depth", 1))) + minimum = max(1, min(int(out.get("min_recurrent_depth", 1)), depth)) + out.update(prelude_layers=p, recurrent_groups=g, coda_layers=layers-p-g, + recurrent_depth=depth, min_recurrent_depth=minimum, + max_inference_depth=max(minimum, int(out.get("max_inference_depth") or depth))) + return out + + +def config_hash(source: Any) -> str: + payload = json.dumps(canonical_config_dict(source), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def file_sha256(path: str | None) -> str | None: + if not path or not os.path.isfile(path): + return None + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def ui_to_canonical(ui_state: Mapping[str, Any]) -> dict: + return canonical_config_dict(ui_state) + + +def canonical_to_ui(config: Mapping[str, Any], defaults: Mapping[str, Any]) -> dict: + out = dict(defaults) + canonical = canonical_config_dict(config) + out.update({k: v for k, v in canonical.items() if k in out or k in CANONICAL_FIELDS}) + out.pop("schema_version", None) + return out + + +RECURRENT_ATTENTION_MODES = ("shared_first", "per_iteration", "first_only") + + +def resolve_recurrent_attention_mode(source: Any) -> str: + """Return the recurrent attention contract for a config or mapping. + + v3 configs carry ``recurrent_attention_mode`` explicitly. v2 configs + predate it and are resolved from ``recurrent_latent_refinement`` so their + numerics never change: True → ``first_only`` (latent-only middle steps), + False → ``per_iteration``. + """ + get = source.get if isinstance(source, Mapping) else (lambda k, d=None: getattr(source, k, d)) + mode = get("recurrent_attention_mode", None) + if mode: + mode = str(mode).lower() + if mode not in RECURRENT_ATTENTION_MODES: + raise ValueError(f"recurrent_attention_mode must be one of {RECURRENT_ATTENTION_MODES}, got {mode!r}") + return mode + if int(get("architecture_version", LEGACY_ARCHITECTURE_VERSION) or LEGACY_ARCHITECTURE_VERSION) >= 3: + return "shared_first" + return "first_only" if bool(get("recurrent_latent_refinement", True)) else "per_iteration" + + +class RecurrentCacheLayout: + """Map (physical layer, loop iteration) to a KV-cache slot. + + ``shared_first``/``first_only`` keep one slot per physical layer. + ``per_iteration`` gives each recurrent iteration its own slot, ordered + prelude, group0·step0..stepD-1, group1·…, coda. + """ + + def __init__(self, prelude: int, groups: int, coda: int, depth: int, mode: str): + self.prelude, self.groups, self.coda = int(prelude), int(groups), int(coda) + self.depth, self.mode = max(1, int(depth)), mode + self.per_iteration = mode == "per_iteration" + physical = self.prelude + self.groups + self.coda + self.num_slots = (self.prelude + self.groups * self.depth + self.coda) if self.per_iteration else physical + + def slot(self, layer: int, step: int = 0) -> int: + if not self.per_iteration or layer < self.prelude: + return int(layer) + if layer < self.prelude + self.groups: + return self.prelude + (layer - self.prelude) * self.depth + int(step) + return self.prelude + self.groups * self.depth + (layer - self.prelude - self.groups) + + def physical_layer(self, slot: int) -> int: + if not self.per_iteration or slot < self.prelude: + return int(slot) + recurrent_end = self.prelude + self.groups * self.depth + if slot < recurrent_end: + return self.prelude + (slot - self.prelude) // self.depth + return self.prelude + self.groups + (slot - recurrent_end) + + +def recurrent_cache_layout(config: Any, depth: int | None = None) -> RecurrentCacheLayout: + get = config.get if isinstance(config, Mapping) else (lambda k, d=None: getattr(config, k, d)) + return RecurrentCacheLayout( + get("prelude_layers", 0), get("recurrent_groups", 1), get("coda_layers", 0), + depth if depth is not None else get("max_inference_depth", None) or get("recurrent_depth", 1), + resolve_recurrent_attention_mode(config), + ) diff --git a/chronos/data/__init__.py b/chronos/data/__init__.py index e69de29..3c8442b 100644 --- a/chronos/data/__init__.py +++ b/chronos/data/__init__.py @@ -0,0 +1,11 @@ +from .flexible_dataset import ( + ModalityRegistry, RecordDetector, MediaResolver, MultimodalCollator, + MultimodalStreamingDataset, MultimodalStreamingSFTDataset, + multimodal_collate, dataset_modality_summary, +) + +__all__ = [ + "ModalityRegistry", "RecordDetector", "MediaResolver", "MultimodalCollator", + "MultimodalStreamingDataset", "MultimodalStreamingSFTDataset", + "multimodal_collate", "dataset_modality_summary", +] diff --git a/chronos/data/corpus/__init__.py b/chronos/data/corpus/__init__.py new file mode 100644 index 0000000..495fd95 --- /dev/null +++ b/chronos/data/corpus/__init__.py @@ -0,0 +1,7 @@ +"""Chronos corpus pipeline: download, clean, deduplicate, mix and shard. + +The package is task-agnostic. Translation is one configuration of it; the +same stages handle monolingual pretraining text and chat/SFT records. +Modules import only the standard library, numpy and pyarrow at top level so +that data preparation never pulls in Torch or MLX. +""" diff --git a/chronos/data/corpus/clean.py b/chronos/data/corpus/clean.py new file mode 100644 index 0000000..9c37d56 --- /dev/null +++ b/chronos/data/corpus/clean.py @@ -0,0 +1,105 @@ +"""Pair cleaning: de-segmentation, detokenization, normalization and filters. + +Every rule is deterministic and recorded as a reject reason so the corpus +report can account for every dropped pair. +""" +from __future__ import annotations + +import re +from dataclasses import dataclass + +from chronos.tokenizer.normalize import normalize_text + +_HAN = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff]") +_LATIN = re.compile(r"[A-Za-z]") +_URL = re.compile(r"(https?://|www\.)\S+", re.I) +_HTML = re.compile(r"]{0,80}>|&(?:amp|lt|gt|quot|nbsp|#\d+);") +_REPEAT = re.compile(r"(.)\1{9,}") + +# Punctuation policy for Chinese text: full-width sentence punctuation after +# Han characters, matching standard simplified-Chinese typesetting. +_ZH_PUNCT = {",": "，", ";": "；", ":": "：", "?": "？", "!": "！", "(": "（", ")": "）"} +_H = r"\u3400-\u9fff\u3000-\u303f\uff00-\uffef" +# ASCII punctuation touching Han text becomes full-width; ``1,000``/``3:30`` +# inside Latin spans stay ASCII. +_ZH_PUNCT_RE = re.compile( + rf"(?<=[{_H}])([,;:?!])|([,;:?!])(?=[{_H}])" + rf"|(?<=[{_H}])([(])|([(])(?=[{_H}])|(?<=[{_H}])([)])|([)])(?=[{_H}])" +) + +# Segmenter padding around full-width punctuation, and Moses ``3 : 30``. +_FW_GAP_RE = re.compile(r"\s*([\uff08\uff09\uff0c\u3002\uff1b\uff1a\uff1f\uff01\u3001])\s*") +_EN_TIME_RE = re.compile(r"(\d)\s*([:.,])\s*(?=\d)") + + +@dataclass +class CleanConfig: + unsegment_zh: bool = False # spaces between CJK chars come from segmentation + detokenize_en: bool = False # Moses-tokenized English (`` don 't`` → ``don't``) + zh_fullwidth_punct: bool = True + min_chars: int = 1 + max_chars: int = 1024 + max_len_ratio: float = 3.0 # char-length ratio after a script-aware scale + min_han_ratio: float = 0.3 + min_latin_ratio: float = 0.4 + + +class PairCleaner: + def __init__(self, config: CleanConfig): + self.config = config + self._detok = None + if config.detokenize_en: + from sacremoses import MosesDetokenizer + self._detok = MosesDetokenizer(lang="en") + + def clean(self, zh: str, en: str) -> tuple[str, str] | str: + """Return the cleaned pair, or a reject-reason string.""" + cfg = self.config + zh = normalize_text(zh, unsegment_cjk=cfg.unsegment_zh) + en = normalize_text(en) + if self._detok is not None: + en = self._detok.detokenize(en.split()) + if en[:1].islower(): + en = en[0].upper() + en[1:] + if cfg.zh_fullwidth_punct: + convert = lambda m: _ZH_PUNCT[next(g for g in m.groups() if g)] + # Two passes: a conversion can create the Han-side context for its + # neighbour, e.g. ``Blue),`` → ``Blue)，`` → ``Blue）,``. + zh = _ZH_PUNCT_RE.sub(convert, _ZH_PUNCT_RE.sub(convert, zh)) + zh = _FW_GAP_RE.sub(r"\1", zh) + if self._detok is not None: + en = _EN_TIME_RE.sub(r"\1\2", en) + if len(zh) < cfg.min_chars or len(en) < cfg.min_chars: + return "empty" + if len(zh) > cfg.max_chars or len(en) > cfg.max_chars: + return "too_long" + if _URL.search(zh) or _URL.search(en) or _HTML.search(zh) or _HTML.search(en): + return "markup" + if _REPEAT.search(zh) or _REPEAT.search(en): + return "repetition" + han = len(_HAN.findall(zh)) + if han / max(len(zh.replace(" ", "")), 1) < cfg.min_han_ratio: + return "zh_script" + latin = len(_LATIN.findall(en)) + if latin / max(len(en.replace(" ", "")), 1) < cfg.min_latin_ratio: + return "en_script" + # One Han character carries roughly as much content as 3-4 English + # characters; compare on that scale. + zh_len, en_len = len(zh) * 3.5, len(en) + ratio = max(zh_len, en_len) / max(min(zh_len, en_len), 1) + if ratio > cfg.max_len_ratio and max(len(zh), len(en)) > 8: + return "length_ratio" + if zh == en: + return "copy" + return zh, en + + +def pair_key(zh: str, en: str) -> str: + """Normalization used for exact deduplication and leakage checks.""" + def squash(s: str) -> str: + return re.sub(r"[\s\W_]+", "", s.lower()) + return squash(zh) + "\x1f" + squash(en) + + +def side_key(text: str) -> str: + return re.sub(r"[\s\W_]+", "", text.lower()) diff --git a/chronos/data/corpus/download.py b/chronos/data/corpus/download.py new file mode 100644 index 0000000..46954cd --- /dev/null +++ b/chronos/data/corpus/download.py @@ -0,0 +1,161 @@ +"""Verified, resumable dataset downloads. + +Every source is pinned to a revision, written into its own directory and +checked against the size and SHA-256 that the hosting API reports. Results +are recorded in ``download_manifest.json`` next to the files so later stages +can prove which bytes they consumed. +""" +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import time +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import asdict, dataclass, field + +CHUNK = 1 << 20 + + +@dataclass +class RemoteFile: + url: str + rel_path: str + size: int | None = None + sha256: str | None = None + + +@dataclass +class SourceSpec: + name: str + dest_dir: str + files: list[RemoteFile] = field(default_factory=list) + revision: str | None = None + license: str | None = None + note: str = "" + + +def _get_json(url: str, timeout: float = 60.0) -> object: + req = urllib.request.Request(url, headers={"User-Agent": "chronos-corpus/1"}) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read().decode("utf-8")) + + +def hf_dataset_files(repo: str, subdir: str, revision: str | None = None) -> tuple[str, list[RemoteFile]]: + """List files under ``subdir`` of a Hugging Face dataset, pinned to a commit.""" + meta = _get_json(f"https://huggingface.co/api/datasets/{repo}" + (f"/revision/{revision}" if revision else "")) + sha = meta["sha"] + tree = _get_json(f"https://huggingface.co/api/datasets/{repo}/tree/{sha}/{urllib.parse.quote(subdir)}") + files = [] + for item in tree: + if item.get("type") != "file": + continue + lfs = item.get("lfs") or {} + files.append(RemoteFile( + url=f"https://huggingface.co/datasets/{repo}/resolve/{sha}/{urllib.parse.quote(item['path'])}", + rel_path=item["path"], + size=int(item.get("size") or 0) or None, + sha256=lfs.get("oid"), + )) + return sha, files + + +def modelscope_dataset_files(repo: str, root: str, match=None) -> list[RemoteFile]: + """List files from a ModelScope dataset tree; ``match(path)`` filters them.""" + query = urllib.parse.urlencode({"Revision": "master", "Root": root, "Recursive": "true", "PageSize": 5000}) + data = _get_json(f"https://modelscope.cn/api/v1/datasets/{repo}/repo/tree?{query}") + files = [] + for item in (data.get("Data") or {}).get("Files") or []: + if item.get("Type") != "blob" or (match and not match(item["Path"])): + continue + rev = item.get("Revision") or "master" + files.append(RemoteFile( + url=f"https://modelscope.cn/datasets/{repo}/resolve/{rev}/{urllib.parse.quote(item['Path'])}", + rel_path=item["Path"], + size=int(item.get("Size") or 0) or None, + sha256=item.get("Sha256") or None, + )) + return files + + +def sha256_file(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as handle: + for block in iter(lambda: handle.read(CHUNK), b""): + digest.update(block) + return digest.hexdigest() + + +def free_bytes(path: str) -> int: + probe = path + while not os.path.exists(probe): + probe = os.path.dirname(probe) or "." + return shutil.disk_usage(probe).free + + +def download_file(remote: RemoteFile, dest_dir: str, *, min_free_bytes: int = 8 << 30, + retries: int = 30, progress=None) -> dict: + """Download one file with HTTP range resume, then verify size and SHA-256.""" + target = os.path.join(dest_dir, remote.rel_path) + os.makedirs(os.path.dirname(target), exist_ok=True) + record = asdict(remote) | {"local": target} + if os.path.isfile(target) and remote.sha256 and os.path.getsize(target) == (remote.size or -1): + if sha256_file(target) == remote.sha256: + return record | {"status": "verified_existing"} + needed = (remote.size or 0) + min_free_bytes + if free_bytes(dest_dir) < needed: + raise OSError(f"insufficient disk for {remote.rel_path}: need {needed} bytes free") + partial = target + ".part" + for attempt in range(retries): + offset = os.path.getsize(partial) if os.path.exists(partial) else 0 + if remote.size and offset >= remote.size: + break + headers = {"User-Agent": "chronos-corpus/1"} + if offset: + headers["Range"] = f"bytes={offset}-" + try: + req = urllib.request.Request(remote.url, headers=headers) + with urllib.request.urlopen(req, timeout=120) as resp: + # A server that ignores Range restarts from byte zero. + mode = "ab" if offset and resp.status == 206 else "wb" + done = offset if mode == "ab" else 0 + with open(partial, mode) as out: + for block in iter(lambda: resp.read(CHUNK), b""): + out.write(block) + done += len(block) + if progress: + progress(remote.rel_path, done, remote.size) + except (OSError, urllib.error.URLError): # network hiccup: resume below + pass + # CDNs may close a stream early without raising; resume until the + # advertised size is reached (unknown size: accept one clean pass). + current = os.path.getsize(partial) if os.path.exists(partial) else 0 + if (remote.size and current >= remote.size) or (not remote.size and current > 0): + break + if attempt == retries - 1: + raise IOError(f"download incomplete for {remote.rel_path}: {current} of {remote.size} bytes") + time.sleep(min(2 ** attempt, 30)) + size = os.path.getsize(partial) + if remote.size and size != remote.size: + raise IOError(f"size mismatch for {remote.rel_path}: {size} != {remote.size}") + digest = sha256_file(partial) + if remote.sha256 and digest != remote.sha256: + os.remove(partial) + raise IOError(f"sha256 mismatch for {remote.rel_path}") + os.replace(partial, target) + return record | {"bytes": size, "sha256": digest, + "status": "downloaded_verified" if remote.sha256 else "downloaded_unverified_hash", + "verified_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())} + + +def download_source(spec: SourceSpec, *, progress=None) -> dict: + os.makedirs(spec.dest_dir, exist_ok=True) + records = [download_file(f, spec.dest_dir, progress=progress) for f in spec.files] + manifest = {"name": spec.name, "revision": spec.revision, "license": spec.license, + "note": spec.note, "files": records} + with open(os.path.join(spec.dest_dir, "download_manifest.json"), "w", encoding="utf-8") as handle: + json.dump(manifest, handle, ensure_ascii=False, indent=2) + return manifest diff --git a/chronos/data/corpus/evalsets.py b/chronos/data/corpus/evalsets.py new file mode 100644 index 0000000..b9158fd --- /dev/null +++ b/chronos/data/corpus/evalsets.py @@ -0,0 +1,65 @@ +"""Held-out evaluation sets (never trained on; also used for leakage filtering).""" +from __future__ import annotations + +import ast +import json +import os + + +def _flores(raw: str, split: str): + import pyarrow.parquet as pq + path = os.path.join(raw, "facebook__flores__eng-zho", "data", "pair", "eng_Latn-zho_Hans", + f"{split}-00000-of-00001.parquet") + for r in pq.read_table(path, columns=["sentence_eng_Latn", "sentence_zho_Hans"]).to_pylist(): + yield r["sentence_eng_Latn"], r["sentence_zho_Hans"] + + +def _wmt22(raw: str): + import pyarrow.parquet as pq + folder = os.path.join(raw, "haoranxu__WMT22-Test__en-zh", "en-zh") + for name in sorted(os.listdir(folder)): + if not name.endswith(".parquet"): + continue + for r in pq.read_table(os.path.join(folder, name)).to_pylist(): + value = r["en-zh"] + pair = value if isinstance(value, dict) else ast.literal_eval(value) + yield pair["en"], pair["zh"] + + +def _wmt24pp(raw: str): + path = os.path.join(raw, "google__wmt24pp__en-zh_CN", "en-zh_CN.jsonl") + with open(path, encoding="utf-8") as fh: + for line in fh: + r = json.loads(line) + if str(r.get("is_bad_source", "")).strip().lower() == "true": + continue # documented low-quality sources + yield r["source"], r["target"] + + +def _opus_test(raw: str): + import pyarrow.parquet as pq + path = os.path.join(raw, "Helsinki-NLP__opus-100", "en-zh", "test-00000-of-00001.parquet") + for r in pq.read_table(path, columns=["translation"]).to_pylist(): + yield r["translation"]["en"], r["translation"]["zh"] + + +EVAL_SETS = { + "flores_dev": lambda raw: _flores(raw, "dev"), + "flores_devtest": lambda raw: _flores(raw, "devtest"), + "wmt22_enzh": _wmt22, + "wmt24pp_enzh": _wmt24pp, + "opus100_test": _opus_test, +} + + +def load_eval_sets(raw_root: str, names=None) -> dict[str, list[tuple[str, str]]]: + """Return {name: [(en, zh), ...]} for every available set.""" + out = {} + for name, loader in EVAL_SETS.items(): + if names and name not in names: + continue + try: + out[name] = list(loader(raw_root)) + except FileNotFoundError: + continue + return out diff --git a/chronos/data/corpus/pipeline.py b/chronos/data/corpus/pipeline.py new file mode 100644 index 0000000..b6a9716 --- /dev/null +++ b/chronos/data/corpus/pipeline.py @@ -0,0 +1,205 @@ +"""Parallel corpus build: clean → dedup → leakage filter → split → manifest. + +Stage 1 (parallel): every work unit (source × row range) is read, cleaned +and written as a zstd parquet part with 64-bit keys. +Stage 2 (single process, numpy): global exact dedup on the pair key in source +priority order, cross-split / eval leakage removal on either side's key, and +a deterministic validation split by Chinese-side hash. +Outputs are append-only parquet parts plus ``corpus_manifest.json``. +""" +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import time +from collections import Counter +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import asdict, dataclass, field + +import numpy as np + +from .clean import CleanConfig, PairCleaner, pair_key, side_key +from .readers import count_rows, csv_byte_ranges, iter_csv_byte_range, iter_pairs + + +@dataclass +class PairSource: + name: str + path: str + kind: str + fields: tuple[str, str] # (zh_field, en_field) + clean: CleanConfig = field(default_factory=CleanConfig) + priority: int = 0 # lower wins on duplicates + weight: float = 1.0 # mixture weight (sampling multiplier) + rows_per_unit: int = 1_000_000 + + +def h64(text: str) -> int: + return int.from_bytes(hashlib.blake2b(text.encode("utf-8"), digest_size=8).digest(), "little", signed=True) + + +def _clean_unit(args) -> dict: + source_dict, start, stop, out_path = args + import pyarrow as pa + import pyarrow.parquet as pq + src = PairSource(**{**source_dict, "clean": CleanConfig(**source_dict["clean"])}) + cleaner = PairCleaner(src.clean) + zh_out, en_out, pk, zk, ek = [], [], [], [], [] + reasons = Counter() + if src.kind in ("pairs_csv", "pairs_tsv"): + stream = iter_csv_byte_range(src.path, src.fields, start, stop, + "\t" if src.kind == "pairs_tsv" else ",") + else: + stream = iter_pairs(src.path, src.kind, src.fields, start=start, stop=stop) + for zh, en in stream: + result = cleaner.clean(zh, en) + if isinstance(result, str): + reasons[result] += 1 + continue + zh, en = result + zh_out.append(zh) + en_out.append(en) + pk.append(h64(pair_key(zh, en))) + zk.append(h64(side_key(zh))) + ek.append(h64(side_key(en))) + reasons["kept"] = len(zh_out) + table = pa.table({"zh": zh_out, "en": en_out, + "pair_key": pa.array(pk, pa.int64()), "zh_key": pa.array(zk, pa.int64()), + "en_key": pa.array(ek, pa.int64())}) + tmp = out_path + ".tmp" + pq.write_table(table, tmp, compression="zstd") + info = {"source": src.name, "start": start, "stop": stop, "path": out_path, "reasons": dict(reasons)} + with open(out_path + ".json", "w", encoding="utf-8") as fh: + json.dump(info, fh) + os.replace(tmp, out_path) # the part exists only once complete (resume-safe) + return info + + +def _units(sources: list[PairSource], parts_dir: str): + units = [] + for src in sources: + if src.kind in ("pairs_csv", "pairs_tsv"): + # ~150 bytes per pair row: byte ranges avoid rescanning from row 0. + ranges = csv_byte_ranges(src.path, src.rows_per_unit * 150) + else: + n = count_rows(src.path, src.kind) + ranges = [(s, min(n, s + src.rows_per_unit)) for s in range(0, n, src.rows_per_unit)] + for i, (start, stop) in enumerate(ranges): + out = os.path.join(parts_dir, f"{src.name}.{i:05d}.parquet") + units.append(((asdict(src) | {"clean": asdict(src.clean)}), start, stop, out)) + return units + + +def build_corpus(sources: list[PairSource], eval_pairs: dict[str, list[tuple[str, str]]], + out_dir: str, *, workers: int = 16, valid_per_mille: int = 2, + min_free_gb: float = 8.0, keep_parts: bool = False, log=print) -> dict: + os.makedirs(out_dir, exist_ok=True) + parts_dir = os.path.join(out_dir, "clean_parts") + os.makedirs(parts_dir, exist_ok=True) + if shutil.disk_usage(out_dir).free < min_free_gb * (1 << 30): + raise OSError(f"less than {min_free_gb} GiB free at {out_dir}") + t0 = time.time() + units = [u for u in _units(sources, parts_dir) if not os.path.exists(u[3])] + done_units = [] + log(f"[corpus] stage1: {len(units)} pending units, {workers} workers") + with ProcessPoolExecutor(max_workers=workers) as pool: + futures = [pool.submit(_clean_unit, u) for u in units] + for i, fut in enumerate(as_completed(futures), 1): + info = fut.result() + done_units.append(info) + log(f"[corpus] {i}/{len(units)} {info['source']}[{info['start']}:{info['stop']}] " + f"kept={info['reasons'].get('kept', 0)} ({time.time() - t0:.0f}s)") + manifest = _finalize(sources, eval_pairs, out_dir, parts_dir, done_units, valid_per_mille, t0, log) + if not keep_parts: + # The final split files hold every kept row; drop the intermediate + # text but keep the per-unit reason sidecars for auditing. + for name in os.listdir(parts_dir): + if name.endswith(".parquet"): + os.remove(os.path.join(parts_dir, name)) + return manifest + + +def _split_of(zh_key: np.ndarray, valid_per_mille: int) -> np.ndarray: + """0 = train, 1 = valid; deterministic in the Chinese-side hash.""" + mixed = (zh_key.astype(np.uint64) * np.uint64(0x9E3779B97F4A7C15)) >> np.uint64(40) + return (mixed % np.uint64(1000) < np.uint64(valid_per_mille)).astype(np.int8) + + +def _finalize(sources, eval_pairs, out_dir, parts_dir, done_units, valid_per_mille, t0, log): + import pyarrow as pa + import pyarrow.parquet as pq + by_source = {s.name: s for s in sources} + parts = sorted((p for p in os.listdir(parts_dir) if p.endswith(".parquet")), + key=lambda p: (by_source[p.split(".")[0]].priority, p)) + # ---- gather keys (≈ 24 bytes/row) ---- + pair_keys, zh_keys, en_keys, part_id, row_id = [], [], [], [], [] + for pi, name in enumerate(parts): + t = pq.read_table(os.path.join(parts_dir, name), columns=["pair_key", "zh_key", "en_key"]) + n = t.num_rows + pair_keys.append(t["pair_key"].to_numpy()); zh_keys.append(t["zh_key"].to_numpy()) + en_keys.append(t["en_key"].to_numpy()) + part_id.append(np.full(n, pi, np.int32)); row_id.append(np.arange(n, dtype=np.int64)) + pair_keys = np.concatenate(pair_keys); zh_keys = np.concatenate(zh_keys); en_keys = np.concatenate(en_keys) + part_id = np.concatenate(part_id); row_id = np.concatenate(row_id) + total = len(pair_keys) + # ---- exact dedup: first occurrence in priority order wins ---- + _, first = np.unique(pair_keys, return_index=True) + keep = np.zeros(total, bool); keep[first] = True + dup = int(total - keep.sum()) + # ---- evaluation leakage on either side ---- + eval_zh = np.array([h64(side_key(zh)) for pairs in eval_pairs.values() for _en, zh in pairs], np.int64) + eval_en = np.array([h64(side_key(en)) for pairs in eval_pairs.values() for en, _zh in pairs], np.int64) + leak = np.isin(zh_keys, eval_zh) | np.isin(en_keys, eval_en) + leaked = int((leak & keep).sum()) + keep &= ~leak + # ---- split by zh hash; drop en-side keys shared across splits ---- + split = _split_of(zh_keys, valid_per_mille) + valid_en = np.unique(en_keys[keep & (split == 1)]) + conflict = keep & (split == 0) & np.isin(en_keys, valid_en) + keep &= ~conflict + log(f"[corpus] rows={total} duplicates={dup} eval_leak={leaked} split_conflicts={int(conflict.sum())}") + # ---- write final splits, part by part (bounded memory) ---- + final_dir = os.path.join(out_dir, "pairs") + os.makedirs(final_dir, exist_ok=True) + counts = Counter() + for pi, name in enumerate(parts): + mask = part_id == pi + rows, k, sp = row_id[mask], keep[mask], split[mask] + if not k.any(): + continue + t = pq.read_table(os.path.join(parts_dir, name), columns=["zh", "en", "pair_key"]) + src = name.split(".")[0] + for split_name, code in (("train", 0), ("valid", 1)): + sel = rows[k & (sp == code)] + if len(sel) == 0: + continue + sub = t.take(pa.array(sel)).append_column("source", pa.array([src] * len(sel))) + pq.write_table(sub, os.path.join(final_dir, f"{split_name}.{name}"), compression="zstd") + counts[(split_name, src)] += len(sel) + reasons = Counter() + for name in parts: + meta_path = os.path.join(parts_dir, name + ".json") + if os.path.exists(meta_path): + with open(meta_path, encoding="utf-8") as fh: + info = json.load(fh) + reasons.update({f"{info['source']}:{k}": v for k, v in info["reasons"].items()}) + manifest = { + "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "elapsed_s": round(time.time() - t0, 1), + "sources": [asdict(s) | {"clean": asdict(s.clean)} for s in sources], + "eval_sets": {k: len(v) for k, v in eval_pairs.items()}, + "rows_after_clean": total, "duplicates": dup, "eval_leak_removed": leaked, + "split_conflicts_removed": int(conflict.sum()), + "final": {f"{s}:{src}": n for (s, src), n in sorted(counts.items())}, + "final_train": sum(n for (s, _), n in counts.items() if s == "train"), + "final_valid": sum(n for (s, _), n in counts.items() if s == "valid"), + "clean_reasons": dict(sorted(reasons.items())), + "valid_per_mille": valid_per_mille, + "split_policy": "zh-side hash; en-side keys shared with valid removed from train; " + "eval sets removed on either side", + } + with open(os.path.join(out_dir, "corpus_manifest.json"), "w", encoding="utf-8") as fh: + json.dump(manifest, fh, ensure_ascii=False, indent=2) + return manifest diff --git a/chronos/data/corpus/presets.py b/chronos/data/corpus/presets.py new file mode 100644 index 0000000..755cfd2 --- /dev/null +++ b/chronos/data/corpus/presets.py @@ -0,0 +1,63 @@ +"""Corpus presets and the ``python -m chronos.data.corpus.presets`` CLI. + +A preset is only a list of ``PairSource`` objects plus evaluation sets; the +pipeline itself is task-agnostic. +""" +from __future__ import annotations + +import argparse +import json +import os +import sys + +from .clean import CleanConfig +from .evalsets import load_eval_sets +from .pipeline import PairSource, build_corpus + + +def translation_enzh_sources(raw_root: str) -> list[PairSource]: + sources: list[PairSource] = [] + # Priority: curated/natural sources first so duplicates keep their copy. + opus = os.path.join(raw_root, "Helsinki-NLP__opus-100", "en-zh", "train-00000-of-00001.parquet") + if os.path.exists(opus): + sources.append(PairSource("opus100", opus, "pairs_parquet", ("translation.zh", "translation.en"), + priority=0, weight=1.5, rows_per_unit=250_000)) + wmt_dir = os.path.join(raw_root, "wmt__wmt19__zh-en", "zh-en") + if os.path.isdir(wmt_dir): + for name in sorted(os.listdir(wmt_dir)): + if name.startswith("train-") and name.endswith(".parquet"): + sources.append(PairSource(f"wmt19_{name.split('-')[1]}", os.path.join(wmt_dir, name), + "pairs_parquet", ("translation.zh", "translation.en"), + priority=1, weight=1.0, rows_per_unit=1_000_000)) + iic = os.path.join(raw_root, "iic__WMT-Chinese-to-English-Machine-Translation-Training-Corpus", + "wmt_zh_en_training_corpus.csv") + if os.path.exists(iic): + sources.append(PairSource("iic_wmt", iic, "pairs_csv", ("0", "1"), + clean=CleanConfig(unsegment_zh=True, detokenize_en=True), + priority=2, weight=1.0, rows_per_unit=1_000_000)) + return sources + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description="Build the cleaned, deduplicated pair corpus") + parser.add_argument("--raw-root", required=True) + parser.add_argument("--out", required=True) + parser.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 4) - 2)) + parser.add_argument("--valid-per-mille", type=int, default=2) + parser.add_argument("--only", nargs="*", default=None, help="source-name prefixes to include") + args = parser.parse_args(argv) + sources = translation_enzh_sources(args.raw_root) + if args.only: + sources = [s for s in sources if any(s.name.startswith(p) for p in args.only)] + evals = load_eval_sets(args.raw_root) + print(f"[corpus] sources={[s.name for s in sources]} eval={ {k: len(v) for k, v in evals.items()} }", flush=True) + manifest = build_corpus(sources, evals, args.out, workers=args.workers, + valid_per_mille=args.valid_per_mille, + log=lambda m: print(m, flush=True)) + print(json.dumps({k: manifest[k] for k in ("rows_after_clean", "duplicates", "eval_leak_removed", + "final_train", "final_valid")}), flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/chronos/data/corpus/readers.py b/chronos/data/corpus/readers.py new file mode 100644 index 0000000..b71d9f8 --- /dev/null +++ b/chronos/data/corpus/readers.py @@ -0,0 +1,141 @@ +"""Streaming readers that yield (source_text, target_text) pairs or plain text. + +Every reader is lazy and bounded in memory, so multi-GB CSV/parquet sources +never load fully. Field specs use dotted paths (``translation.zh``). +""" +from __future__ import annotations + +import csv +import json +import os +import sys +from typing import Iterator + +csv.field_size_limit(min(sys.maxsize, 1 << 30)) + + +def _get(record, dotted: str): + value = record + for part in str(dotted).split("."): + if isinstance(value, dict): + value = value.get(part) + elif isinstance(value, (list, tuple)) and part.isdigit(): + value = value[int(part)] if int(part) < len(value) else None + else: + return None + return value + + +def iter_pairs(path: str, kind: str, fields: tuple[str, str], *, batch_rows: int = 65536, + start: int = 0, stop: int | None = None) -> Iterator[tuple[str, str]]: + """Yield (a, b) pairs for ``fields=(a_field, b_field)``. + + kind: pairs_parquet | pairs_csv | pairs_jsonl | pairs_tsv. + ``start``/``stop`` select a row range (for parallel shards). + """ + a_field, b_field = fields + row = 0 + if kind == "pairs_parquet": + import pyarrow.parquet as pq + pf = pq.ParquetFile(path) + top = sorted({a_field.split(".")[0], b_field.split(".")[0]}) + for batch in pf.iter_batches(batch_size=batch_rows, columns=top): + n = batch.num_rows + if stop is not None and row >= stop: + return + if row + n <= start: + row += n + continue + for record in batch.to_pylist(): + if start <= row and (stop is None or row < stop): + a, b = _get(record, a_field), _get(record, b_field) + if isinstance(a, str) and isinstance(b, str): + yield a, b + row += 1 + return + if kind in ("pairs_csv", "pairs_tsv"): + delimiter = "\t" if kind == "pairs_tsv" else "," + with open(path, newline="", encoding="utf-8", errors="replace") as fh: + reader = csv.reader(fh, delimiter=delimiter) + header = next(reader, None) or [] + ia = header.index(a_field) if a_field in header else int(a_field) + ib = header.index(b_field) if b_field in header else int(b_field) + for record in reader: + if stop is not None and row >= stop: + return + if row >= start and len(record) > max(ia, ib): + yield record[ia], record[ib] + row += 1 + return + if kind == "pairs_jsonl": + with open(path, encoding="utf-8", errors="replace") as fh: + for line in fh: + if stop is not None and row >= stop: + return + if row >= start and line.strip(): + try: + record = json.loads(line) + except json.JSONDecodeError: + row += 1 + continue + a, b = _get(record, a_field), _get(record, b_field) + if isinstance(a, str) and isinstance(b, str): + yield a, b + row += 1 + return + raise ValueError(f"unknown pair source kind {kind!r}") + + +def iter_csv_byte_range(path: str, fields: tuple[str, str], byte_start: int, byte_stop: int, + delimiter: str = ",") -> Iterator[tuple[str, str]]: + """Pairs whose line starts in [byte_start, byte_stop). Ranges must come + from ``csv_byte_ranges`` (newline aligned); the header line is skipped.""" + with open(path, "rb") as raw: + header = raw.readline().decode("utf-8", errors="replace") + names = next(csv.reader([header], delimiter=delimiter)) + a_field, b_field = fields + ia = names.index(a_field) if a_field in names else int(a_field) + ib = names.index(b_field) if b_field in names else int(b_field) + raw.seek(max(byte_start, raw.tell())) + def lines(): + while raw.tell() < byte_stop: + line = raw.readline() + if not line: + return + yield line.decode("utf-8", errors="replace") + for record in csv.reader(lines(), delimiter=delimiter): + if len(record) > max(ia, ib): + yield record[ia], record[ib] + + +def csv_byte_ranges(path: str, target_bytes: int) -> list[tuple[int, int]]: + size = os.path.getsize(path) + with open(path, "rb") as fh: + fh.readline() + cuts = [fh.tell()] + while cuts[-1] + target_bytes < size: + fh.seek(cuts[-1] + target_bytes) + fh.readline() # align to the next line start + if fh.tell() >= size: + break + cuts.append(fh.tell()) + cuts.append(size) + return list(zip(cuts, cuts[1:])) + + +def count_rows(path: str, kind: str) -> int: + if kind == "pairs_parquet": + import pyarrow.parquet as pq + return pq.ParquetFile(path).metadata.num_rows + with open(path, "rb") as fh: + lines = sum(chunk.count(b"\n") for chunk in iter(lambda: fh.read(1 << 24), b"")) + return max(0, lines - (1 if kind in ("pairs_csv", "pairs_tsv") else 0)) + + +def expand_paths(path: str, suffix: str) -> list[str]: + if os.path.isdir(path): + out = [] + for root, _dirs, files in os.walk(path): + out.extend(os.path.join(root, f) for f in files if f.endswith(suffix)) + return sorted(out) + return [path] diff --git a/chronos/data/corpus/shards.py b/chronos/data/corpus/shards.py new file mode 100644 index 0000000..b18a4db --- /dev/null +++ b/chronos/data/corpus/shards.py @@ -0,0 +1,237 @@ +"""Tokenized, memory-mapped training shards for pretraining and SFT. + +Pretrain stream: one flat uint16/uint32 token array per shard (``.bin``), +documents separated by ``<|endoftext|>``; the loader packs fixed windows. +Document mix (per pair, seeded): + * monolingual zh, monolingual en — language modelling on both sides + * ``<|en|>{en}<|zh|>{zh}`` and reverse — cross-lingual alignment +SFT stream: chat-template records ``[ids, loss_mask]`` (only the assistant +reply is supervised) in a ``.bin`` + ``.idx`` pair for O(1) random access. +""" +from __future__ import annotations + +import json +import os +import random +import time +from concurrent.futures import ProcessPoolExecutor, as_completed + +import numpy as np + +SFT_PROMPT_EN_ZH = "翻译成中文，只输出译文：\n" + + +def _dtype_for(vocab_size: int): + return np.uint16 if vocab_size <= 65535 else np.uint32 + + +def _encode_batch(tok, texts: list[str]) -> list[list[int]]: + return [e.ids for e in tok.encode_batch(texts, add_special_tokens=False)] + + +def _pretrain_part(args) -> dict: + tokenizer_dir, part_path, out_path, seed, mix = args + import pyarrow.parquet as pq + from tokenizers import Tokenizer + tok = Tokenizer.from_file(os.path.join(tokenizer_dir, "tokenizer.json")) + eot, en_tag, zh_tag = (tok.token_to_id(t) for t in ("<|endoftext|>", "<|en|>", "<|zh|>")) + table = pq.read_table(part_path, columns=["zh", "en"]) + zh, en = table["zh"].to_pylist(), table["en"].to_pylist() + rng = random.Random(seed) + zh_ids, en_ids = _encode_batch(tok, zh), _encode_batch(tok, en) + out, docs = [], 0 + for z, e in zip(zh_ids, en_ids): + r = rng.random() + if r < mix["mono"]: + out.extend(z); out.append(eot); out.extend(e); out.append(eot); docs += 2 + elif r < mix["mono"] + mix["en_zh"]: + out.append(en_tag); out.extend(e); out.append(zh_tag); out.extend(z); out.append(eot); docs += 1 + else: + out.append(zh_tag); out.extend(z); out.append(en_tag); out.extend(e); out.append(eot); docs += 1 + arr = np.asarray(out, dtype=_dtype_for(tok.get_vocab_size())) + tmp = out_path + ".tmp" + with open(tmp, "wb") as fh: + arr.tofile(fh) + os.replace(tmp, out_path) + return {"part": os.path.basename(part_path), "tokens": int(arr.size), "docs": docs, + "zh_tokens": int(sum(map(len, zh_ids))), "en_tokens": int(sum(map(len, en_ids)))} + + +def _sft_part(args) -> dict: + tokenizer_dir, part_path, out_prefix, max_len = args + import pyarrow.parquet as pq + from tokenizers import Tokenizer + tok = Tokenizer.from_file(os.path.join(tokenizer_dir, "tokenizer.json")) + table = pq.read_table(part_path, columns=["zh", "en"]) + prompts = [f"<|im_start|>user\n{SFT_PROMPT_EN_ZH}{e}<|im_end|>\n<|im_start|>assistant\n" + for e in table["en"].to_pylist()] + replies = [f"{z}<|im_end|>\n" for z in table["zh"].to_pylist()] + p_ids = [e.ids for e in tok.encode_batch(prompts, add_special_tokens=False)] + r_ids = [e.ids for e in tok.encode_batch(replies, add_special_tokens=False)] + dtype = _dtype_for(tok.get_vocab_size()) + tokens, offsets, prompt_lens, skipped = [], [0], [], 0 + for p, r in zip(p_ids, r_ids): + if len(p) + len(r) > max_len: + skipped += 1 + continue + tokens.extend(p); tokens.extend(r) + offsets.append(offsets[-1] + len(p) + len(r)) + prompt_lens.append(len(p)) + for suffix, arr in ((".bin", np.asarray(tokens, dtype)), + (".idx", np.asarray(offsets, np.int64)), + (".plen", np.asarray(prompt_lens, np.int32))): + tmp = out_prefix + suffix + ".tmp" + with open(tmp, "wb") as fh: + arr.tofile(fh) + os.replace(tmp, out_prefix + suffix) + return {"part": os.path.basename(part_path), "records": len(prompt_lens), "skipped_too_long": skipped, + "tokens": len(tokens)} + + +def build_shards(pairs_dir: str, tokenizer_dir: str, out_dir: str, *, workers: int = 16, + seed: int = 1234, mix: dict | None = None, sft_max_len: int = 512, + log=print) -> dict: + """Tokenize every ``pairs/{train,valid}.*.parquet`` into pretrain + SFT shards.""" + mix = mix or {"mono": 0.5, "en_zh": 0.25, "zh_en": 0.25} + t0 = time.time() + with open(os.path.join(tokenizer_dir, "tokenizer_report.json"), encoding="utf-8") as fh: + tok_meta = json.load(fh) + jobs = [] + for split in ("train", "valid"): + parts = sorted(p for p in os.listdir(pairs_dir) if p.startswith(split + ".") and p.endswith(".parquet")) + os.makedirs(os.path.join(out_dir, "pretrain", split), exist_ok=True) + os.makedirs(os.path.join(out_dir, "sft", split), exist_ok=True) + for i, name in enumerate(parts): + stem = name[:-len(".parquet")] + pre = os.path.join(out_dir, "pretrain", split, stem + ".bin") + sft = os.path.join(out_dir, "sft", split, stem) + if not os.path.exists(pre): + jobs.append(("pretrain", split, _pretrain_part, + (tokenizer_dir, os.path.join(pairs_dir, name), pre, seed + i, mix))) + if not os.path.exists(sft + ".plen"): + jobs.append(("sft", split, _sft_part, + (tokenizer_dir, os.path.join(pairs_dir, name), sft, sft_max_len))) + log(f"[shards] {len(jobs)} pending jobs, {workers} workers") + results = [] + with ProcessPoolExecutor(max_workers=workers) as pool: + futures = {pool.submit(fn, args): (kind, split) for kind, split, fn, args in jobs} + for i, fut in enumerate(as_completed(futures), 1): + kind, split = futures[fut] + info = fut.result() | {"kind": kind, "split": split} + results.append(info) + if i % 10 == 0 or i == len(futures): + log(f"[shards] {i}/{len(futures)} ({time.time() - t0:.0f}s)") + summary = {} + for kind in ("pretrain", "sft"): + for split in ("train", "valid"): + folder = os.path.join(out_dir, kind, split) + if kind == "pretrain": + files = [f for f in os.listdir(folder) if f.endswith(".bin")] + summary[f"{kind}_{split}_tokens"] = sum(os.path.getsize(os.path.join(folder, f)) for f in files) // 2 + else: + files = [f for f in os.listdir(folder) if f.endswith(".plen")] + summary[f"{kind}_{split}_records"] = sum(os.path.getsize(os.path.join(folder, f)) for f in files) // 4 + manifest = {"created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "elapsed_s": round(time.time() - t0, 1), "tokenizer_fingerprint": tok_meta.get("fingerprint"), + "tokenizer_dir": os.path.abspath(tokenizer_dir), "pairs_dir": os.path.abspath(pairs_dir), + "mix": mix, "seed": seed, "sft_max_len": sft_max_len, "sft_prompt_en_zh": SFT_PROMPT_EN_ZH, + "dtype": "uint16" if tok_meta.get("vocab_size", 0) <= 65535 else "uint32", + **summary, "jobs_this_run": len(results)} + with open(os.path.join(out_dir, "shards_manifest.json"), "w", encoding="utf-8") as fh: + json.dump(manifest, fh, ensure_ascii=False, indent=2) + return manifest + + +class PackedTokenStream: + """Random fixed-length windows over memory-mapped pretrain shards.""" + + def __init__(self, folder: str, seq_len: int, dtype=np.uint16, seed: int = 0): + files = sorted(os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".bin")) + self.arrays = [np.memmap(f, dtype=dtype, mode="r") for f in files] + self.arrays = [a for a in self.arrays if a.size > seq_len + 1] + sizes = np.array([a.size for a in self.arrays], dtype=np.float64) + self.weights = sizes / sizes.sum() + self.total_tokens = int(sizes.sum()) + self.seq_len = seq_len + self.rng = np.random.default_rng(seed) + + def batch(self, batch_size: int) -> np.ndarray: + which = self.rng.choice(len(self.arrays), size=batch_size, p=self.weights) + out = np.empty((batch_size, self.seq_len + 1), dtype=np.int32) + for row, i in enumerate(which): + arr = self.arrays[i] + start = int(self.rng.integers(0, arr.size - self.seq_len - 1)) + out[row] = arr[start:start + self.seq_len + 1] + return out + + +class SFTRecords: + """Length-bucketed SFT batches with assistant-only loss masks.""" + + def __init__(self, folder: str, dtype=np.uint16, seed: int = 0, pad_id: int = 0): + stems = sorted(f[:-5] for f in os.listdir(folder) if f.endswith(".plen")) + self.tokens, self.offsets, self.plens, self.index = [], [], [], [] + for si, stem in enumerate(stems): + base = os.path.join(folder, stem) + tok = np.memmap(base + ".bin", dtype=dtype, mode="r") + off = np.fromfile(base + ".idx", dtype=np.int64) + plen = np.fromfile(base + ".plen", dtype=np.int32) + self.tokens.append(tok); self.offsets.append(off); self.plens.append(plen) + self.index.append(np.stack([np.full(len(plen), si), np.arange(len(plen)), + np.diff(off)], axis=1)) + self.index = np.concatenate(self.index) if self.index else np.zeros((0, 3), np.int64) + # Sorted by length once: each batch is a contiguous similar-length + # window found by binary search (O(log n), no full scan per batch). + self.index = self.index[np.argsort(self.index[:, 2], kind="stable")] + self.lengths = self.index[:, 2] + self.rng = np.random.default_rng(seed) + self.pad_id = pad_id + + def __len__(self): + return len(self.index) + + def batch(self, token_budget: int, max_rows: int = 256, width_multiple: int = 64): + """Sample a bucket of similar-length records filling ~token_budget tokens. + + The padded width is rounded up to ``width_multiple`` so the training + graph sees a handful of shapes (Metal kernels are shape-specialized). + """ + target = int(self.lengths[int(self.rng.integers(0, len(self.index)))]) + tol = max(4, target // 8) + lo = int(np.searchsorted(self.lengths, target - tol, side="left")) + hi = int(np.searchsorted(self.lengths, target + tol, side="right")) + bucket = -(-(target + tol) // width_multiple) * width_multiple + rows = max(1, min(max_rows, token_budget // max(bucket, 1), hi - lo)) + pick = self.index[lo + self.rng.choice(hi - lo, size=rows, replace=False)] + width = int(pick[:, 2].max()) + width = -(-width // width_multiple) * width_multiple + 1 # +1: inputs/labels are shifted + ids = np.full((len(pick), width), self.pad_id, dtype=np.int32) + labels = np.full((len(pick), width), -100, dtype=np.int32) + for r, (si, ri, n) in enumerate(pick): + off = self.offsets[si][ri] + seq = self.tokens[si][off:off + n] + ids[r, :n] = seq + p = int(self.plens[si][ri]) + labels[r, p:n] = seq[p:] + return ids, labels + + +def main(argv=None) -> int: + import argparse + parser = argparse.ArgumentParser(description="Tokenize the pair corpus into training shards") + parser.add_argument("--pairs", required=True) + parser.add_argument("--tokenizer", required=True) + parser.add_argument("--out", required=True) + parser.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 4) - 2)) + parser.add_argument("--sft-max-len", type=int, default=512) + args = parser.parse_args(argv) + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + manifest = build_shards(args.pairs, args.tokenizer, args.out, workers=args.workers, + sft_max_len=args.sft_max_len, log=lambda m: print(m, flush=True)) + print(json.dumps({k: v for k, v in manifest.items() if k.endswith(("tokens", "records", "elapsed_s"))}), + flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/data/corpus/sources.py b/chronos/data/corpus/sources.py new file mode 100644 index 0000000..66e99f9 --- /dev/null +++ b/chronos/data/corpus/sources.py @@ -0,0 +1,107 @@ +"""Source registry for corpus presets. + +A preset names the remote sources it needs; each source resolves to a pinned +``SourceSpec``. Local, user-provided datasets are described by +``LocalSource`` and never re-downloaded. +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +from dataclasses import dataclass + +from .download import SourceSpec, download_source, hf_dataset_files, modelscope_dataset_files + + +@dataclass +class LocalSource: + name: str + path: str + kind: str # "pairs_csv" | "pairs_parquet" | "pairs_jsonl" | "text_jsonl" + role: str = "train" # "train" | "eval" + weight: float = 1.0 + fields: tuple[str, str] = ("zh", "en") + + +def _hf(name: str, repo: str, subdir: str, dest: str, revision: str | None = None, + license: str | None = None, note: str = "") -> SourceSpec: + sha, files = hf_dataset_files(repo, subdir, revision) + return SourceSpec(name=name, dest_dir=dest, files=files, revision=sha, license=license, note=note) + + +def _ms(name: str, repo: str, root: str, dest: str, match, license: str | None = None, + note: str = "") -> SourceSpec: + files = modelscope_dataset_files(repo, root, match) + revs = sorted({f.url.split("/resolve/")[1].split("/")[0] for f in files}) + return SourceSpec(name=name, dest_dir=dest, files=files, revision=",".join(revs), license=license, note=note) + + +def translation_enzh_remote(raw_root: str) -> dict: + """Remote sources for the en→zh validation demo (lazy: resolved on demand).""" + return { + "wmt19_zh_en": lambda: _hf( + "wmt19_zh_en", "wmt/wmt19", "zh-en", os.path.join(raw_root, "wmt__wmt19__zh-en"), + license="unknown (dataset card); research use only", + note="news-commentary, UN, CWMT, wikititles mixture"), + "flores_plus_enzh": lambda: _ms( + "flores_plus_enzh", "facebook/flores", "data/pair", + os.path.join(raw_root, "facebook__flores__eng-zho"), + match=lambda p: p.startswith("data/pair/eng_Latn-zho_Hans/"), + license="CC BY-SA 4.0", note="evaluation only: dev + devtest"), + "wmt22_test_enzh": lambda: _hf( + "wmt22_test_enzh", "haoranxu/WMT22-Test", "en-zh", + os.path.join(raw_root, "haoranxu__WMT22-Test__en-zh"), + license="unknown (dataset card); evaluation only"), + "wmt24pp_enzh": lambda: _ms( + "wmt24pp_enzh", "google/wmt24pp", "", os.path.join(raw_root, "google__wmt24pp__en-zh_CN"), + match=lambda p: p in {"en-zh_CN.jsonl", "README.md"}, + license="see README", note="evaluation only"), + } + + +def translation_enzh_local(data_root: str) -> list[LocalSource]: + raw = os.path.join(data_root, "raw") + return [ + LocalSource("iic_wmt_zh_en", os.path.join( + raw, "iic__WMT-Chinese-to-English-Machine-Translation-Training-Corpus", + "wmt_zh_en_training_corpus.csv"), "pairs_csv", fields=("0", "1")), + LocalSource("opus100_en_zh", os.path.join(raw, "Helsinki-NLP__opus-100", "en-zh", + "train-00000-of-00001.parquet"), "pairs_parquet", + fields=("translation.zh", "translation.en")), + LocalSource("opus100_en_zh_test", os.path.join(raw, "Helsinki-NLP__opus-100", "en-zh", + "test-00000-of-00001.parquet"), "pairs_parquet", + role="eval", fields=("translation.zh", "translation.en")), + ] + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description="Download pinned corpus sources") + parser.add_argument("--raw-root", required=True) + parser.add_argument("--only", nargs="*", default=None) + parser.add_argument("--list", action="store_true") + args = parser.parse_args(argv) + remote = translation_enzh_remote(args.raw_root) + names = args.only or list(remote) + for name in names: + spec = remote[name]() + total = sum(f.size or 0 for f in spec.files) + print(f"[{name}] rev={spec.revision} files={len(spec.files)} bytes={total}", flush=True) + if args.list: + continue + last = {} + + def progress(path, done, size, _last=last): + step = done // (256 << 20) + if _last.get(path) != step: + _last[path] = step + print(f" {path}: {done/1e6:.0f}/{(size or 0)/1e6:.0f} MB", flush=True) + + manifest = download_source(spec, progress=progress) + print(json.dumps({"name": name, "status": [f["status"] for f in manifest["files"]]}), flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/chronos/data/flexible_dataset.py b/chronos/data/flexible_dataset.py index b4f5d5f..d93a5c7 100644 --- a/chronos/data/flexible_dataset.py +++ b/chronos/data/flexible_dataset.py @@ -27,6 +27,9 @@ import json import os import threading +import warnings +import wave +import urllib.request import torch from torch.utils.data import Dataset @@ -45,6 +48,123 @@ ("query", "answer"), ) +_IMAGE_KEYS = ("image", "images", "image_path", "image_paths", "image_url", "image_urls") +_AUDIO_KEYS = ("audio", "audios", "audio_path", "audio_paths", "audio_url", "audio_urls", "wav") + + +class ModalityRegistry: + """Single registry used by datasets, trainers and serving adapters.""" + + fields = {"image": _IMAGE_KEYS, "audio": _AUDIO_KEYS} + + def detect(self, record: dict) -> set[str]: + return detect_modalities(record) + + def summary(self, data_path: str, max_records: int = 10000) -> dict: + return dataset_modality_summary(data_path, max_records=max_records) + + +class RecordDetector: + """Compatibility façade for callers that need schema-only detection.""" + + def __call__(self, record: dict) -> set[str]: + return detect_modalities(record) + + detect = __call__ + + +def detect_modalities(record: dict) -> set[str]: + """Detect media modalities without opening files or allocating tensors.""" + found = set() + # Accept the canonical v2 normalized record as input as well as the raw + # heterogeneous JSONL forms. This makes dataset summaries, collators and + # serving adapters true round-trip consumers of the same schema. + nested = record.get("modalities") + if isinstance(nested, dict): + if nested.get("image") not in (None, "", [], {}): + found.add("image") + if nested.get("audio") not in (None, "", [], {}): + found.add("audio") + if nested.get("image_embeds") is not None: + found.add("image") + if nested.get("audio_embeds") is not None: + found.add("audio") + if any(record.get(key) not in (None, "", [], {}) for key in _IMAGE_KEYS): + found.add("image") + if any(record.get(key) not in (None, "", [], {}) for key in _AUDIO_KEYS): + found.add("audio") + if any(key in record for key in ("image_embeds", "image_features")): + found.add("image") + if any(key in record for key in ("audio_embeds", "audio_features")): + found.add("audio") + return found + + +def normalize_multimodal_record(record: dict) -> dict: + """Normalize heterogeneous media fields to the canonical v2 schema.""" + out = { + "messages": record.get("messages") or record.get("conversations") or [], + "modalities": {"image": [], "audio": []}, + "labels": record.get("labels", {}), + "metadata": {k: v for k, v in record.items() if k not in _IMAGE_KEYS + _AUDIO_KEYS}, + } + nested = record.get("modalities") if isinstance(record.get("modalities"), dict) else {} + for key in _IMAGE_KEYS: + value = record.get(key) + if value not in (None, "", [], {}): + out["modalities"]["image"] = list(value) if isinstance(value, (list, tuple)) else [value] + break + if not out["modalities"]["image"] and nested.get("image") not in (None, "", [], {}): + value = nested["image"] + out["modalities"]["image"] = list(value) if isinstance(value, (list, tuple)) else [value] + for key in _AUDIO_KEYS: + value = record.get(key) + if value not in (None, "", [], {}): + out["modalities"]["audio"] = list(value) if isinstance(value, (list, tuple)) else [value] + break + if not out["modalities"]["audio"] and nested.get("audio") not in (None, "", [], {}): + value = nested["audio"] + out["modalities"]["audio"] = list(value) if isinstance(value, (list, tuple)) else [value] + # Preserve already-materialized features for embedding-only datasets. + for name, aliases in (("image_embeds", ("image_embeds", "image_features")), + ("audio_embeds", ("audio_embeds", "audio_features"))): + for key in aliases: + if key in record: + out["modalities"][name] = record[key] + break + if name not in out["modalities"] and nested.get(name) is not None: + out["modalities"][name] = nested[name] + return out + + +def dataset_modality_summary(data_path: str, max_records: int = 10000) -> dict: + """Stream a JSONL prefix and return deterministic modality statistics.""" + counts = {"text": 0, "image": 0, "audio": 0, "multimodal": 0, "invalid": 0} + scanned = 0 + with open(data_path, "rb") as handle: + for raw in handle: + if scanned >= max_records: + break + if not raw.strip(): + continue + scanned += 1 + try: + record = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + counts["invalid"] += 1 + continue + mods = detect_modalities(record) + if not mods: + counts["text"] += 1 + else: + for mod in mods: + counts[mod] += 1 + if len(mods) > 1: + counts["multimodal"] += 1 + counts["scanned"] = scanned + counts["detected_modalities"] = [m for m in ("image", "audio") if counts[m] > 0] + return counts + def _extract_text(record: dict) -> str: for k in _TEXT_KEYS: @@ -177,6 +297,269 @@ def __getitem__(self, index): return input_ids, labels +def _media_value(value): + if isinstance(value, (list, tuple)): + return value[0] if value else None + return value + + +class MediaResolver: + """Lazy media resolver shared by JSONL and API serving paths.""" + + def __init__(self, *, image_size: int = 224, audio_sample_rate: int = 16000, + max_audio_seconds: float = 30.0): + self.image_size = int(image_size) + self.audio_sample_rate = int(audio_sample_rate) + self.max_audio_seconds = float(max_audio_seconds) + + def image(self, value): + return _load_image(value, self.image_size) + + def audio(self, value): + return _load_audio(value, self.audio_sample_rate, self.max_audio_seconds) + + def resolve(self, modality: str, value): + if modality == "image": + return self.image(value) + if modality == "audio": + return self.audio(value) + raise ValueError(f"unsupported modality: {modality}") + + +class MultimodalCollator: + """Callable wrapper retaining the functional collator API.""" + + def __call__(self, batch): + return multimodal_collate(batch) + + +def _load_image(value, size: int) -> torch.Tensor: + try: + from PIL import Image + if isinstance(value, torch.Tensor): + image = value + if image.dim() == 3 and image.shape[0] not in (1, 3): + image = image.permute(2, 0, 1) + return torch.nn.functional.interpolate( + image.float().unsqueeze(0), size=(size, size), mode="bilinear", align_corners=False + ).squeeze(0).div(255.0 if image.max() > 1.5 else 1.0) + if isinstance(value, Image.Image): + image = value + elif str(value).startswith(("http://", "https://")): + with urllib.request.urlopen(str(value), timeout=10) as response: + import io + image = Image.open(io.BytesIO(response.read())) + else: + image = Image.open(str(value)) + image = image.convert("RGB").resize((size, size)) + import numpy as np + return torch.from_numpy(np.asarray(image)).permute(2, 0, 1).float().div(255.0) + except Exception as exc: + raise ValueError(f"unable to load image {value!r}: {exc}") from exc + + +def _load_audio(value, sample_rate: int, max_seconds: float) -> torch.Tensor: + try: + if isinstance(value, torch.Tensor): + audio = value.float().flatten() + return audio[: int(sample_rate * max_seconds)] + # wave.open accepts paths and binary file-like objects. Supporting + # bytes here makes API data-URI audio use exactly the JSONL decoder. + source = value + if isinstance(value, (bytes, bytearray)): + import io + source = io.BytesIO(value) + with wave.open(source, "rb") as handle: + n = min(handle.getnframes(), int(sample_rate * max_seconds)) + raw = handle.readframes(n) + import numpy as np + audio = torch.from_numpy(np.frombuffer(raw, dtype=np.int16).copy()).float().div(32768.0) + channels = int(handle.getnchannels()) + if channels > 1: + audio = audio.view(-1, channels).mean(dim=1) + source_rate = int(handle.getframerate()) + if source_rate != sample_rate and audio.numel() > 1: + target = int(audio.numel() * sample_rate / source_rate) + audio = torch.nn.functional.interpolate( + audio.view(1, 1, -1), size=target, mode="linear", align_corners=False + ).view(-1) + return audio[: int(sample_rate * max_seconds)] + except Exception as exc: + raise ValueError(f"unable to load audio {value!r}: {exc}") from exc + + +class MultimodalStreamingDataset(FlexibleDataset): + """Streaming LM dataset with lazy image/audio materialization. + + Text-only records retain the regular tuple contract. Media records return + a dict consumed by ``multimodal_collate`` and ChronosTrainer. + """ + def __init__(self, data_path, tokenizer, max_length=512, *, image_size=224, + audio_sample_rate=16000, max_audio_seconds=30.0, + media_error_policy="warn"): + super().__init__(data_path, tokenizer, max_length=max_length) + self.image_size = int(image_size) + self.audio_sample_rate = int(audio_sample_rate) + self.max_audio_seconds = float(max_audio_seconds) + self.media_error_policy = str(media_error_policy) + self.modality_summary = dataset_modality_summary(data_path) + + def __getitem__(self, index): + record = self._read_record(index) + base_ids, base_labels = _tokenize_record(record, self.tokenizer, self.max_length) + mods = detect_modalities(record) + if not mods: + return base_ids, base_labels + normalized = normalize_multimodal_record(record) + media = {} + # Pre-encoded features are first-class inputs. They are intentionally + # kept on CPU and only batched here; the model/projector decides how to + # move them to the execution device. + for name in ("image_embeds", "audio_embeds"): + value = normalized["modalities"].get(name) + if value is not None: + try: + media[name] = torch.as_tensor(value, dtype=torch.float32) + except Exception as exc: + if self.media_error_policy == "raise": + raise ValueError(f"invalid {name}: {exc}") from exc + warnings.warn(f"invalid {name}: {exc}", RuntimeWarning, stacklevel=2) + for modality, loader in ( + ("image", lambda v: _load_image(v, self.image_size)), + ("audio", lambda v: _load_audio(v, self.audio_sample_rate, self.max_audio_seconds)), + ): + # A pre-encoded representation takes precedence over a path/value. + if f"{modality}_embeds" in media: + continue + value = _media_value(normalized["modalities"].get(modality)) + if value is None: + continue + try: + media[modality] = loader(value) + except Exception as exc: + if self.media_error_policy == "raise": + raise + warnings.warn(str(exc), RuntimeWarning, stacklevel=2) + if not media: + return base_ids, base_labels + return { + "input_ids": base_ids, + "labels": base_labels, + "modality_inputs": media, + # Per-modality presence mask. The old scalar mask is accepted by + # the model, but a two-column mask is required for mixed batches + # so a text+image sample never accidentally attends to an audio + # placeholder (and vice versa). + "modality_mask": torch.tensor([ + float("image" in media or "image_embeds" in media), + float("audio" in media or "audio_embeds" in media), + ]), + "metadata": normalized.get("metadata", {}), + } + + +def _tokenize_record(record, tokenizer, max_length): + text = _extract_text(record) + tok = tokenizer(text, add_special_tokens=False, max_length=max_length - 2, truncation=True) + tokens = [tokenizer.bos_token_id] + tok.input_ids + [tokenizer.eos_token_id] + pad_id = tokenizer.pad_token_id or 0 + tokens = tokens + [pad_id] * (max_length - len(tokens)) + ids = torch.tensor(tokens, dtype=torch.long) + labels = ids.clone() + labels[ids == pad_id] = -100 + return ids, labels + + +def multimodal_collate(batch): + """Collate mixed text/media records without invoking absent encoders.""" + ids, labels, masks = [], [], [] + images, audios, image_embeds, audio_embeds = [], [], [], [] + has_image = has_audio = has_image_embeds = has_audio_embeds = False + for item in batch: + if isinstance(item, dict): + ids.append(item["input_ids"]); labels.append(item["labels"]) + media_item = item.get("modality_inputs", {}) + image = media_item.get("image") + audio = media_item.get("audio") + image_embed = media_item.get("image_embeds") + audio_embed = media_item.get("audio_embeds") + item_mask = item.get("modality_mask") + if isinstance(item_mask, dict): + mask = torch.tensor([ + float(item_mask.get("image", image is not None or image_embed is not None)), + float(item_mask.get("audio", audio is not None or audio_embed is not None)), + ]) + else: + mask = torch.as_tensor(item_mask if item_mask is not None else 1.0).flatten() + if mask.numel() == 1: + # Legacy scalar means "this record has media". Infer + # the exact columns from materialized inputs. + active = float(mask.item()) + mask = torch.tensor([ + active * float(image is not None or image_embed is not None), + active * float(audio is not None or audio_embed is not None), + ]) + elif mask.numel() != 2: + raise ValueError("modality_mask must have one scalar or two [image,audio] values") + masks.append(mask.float()) + has_image = has_image or image is not None + has_audio = has_audio or audio is not None + has_image_embeds = has_image_embeds or image_embed is not None + has_audio_embeds = has_audio_embeds or audio_embed is not None + images.append(image); audios.append(audio) + image_embeds.append(image_embed); audio_embeds.append(audio_embed) + else: + ids.append(item[0]); labels.append(item[1]); masks.append(torch.zeros(2)) + images.append(None); audios.append(None) + image_embeds.append(None); audio_embeds.append(None) + result = {"input_ids": torch.stack(ids), "labels": torch.stack(labels), + "modality_mask": torch.stack(masks)} + # A pre-encoded representation takes precedence for every sample. Do not + # build both branches in one batch; doing so wastes encoder compute and + # makes duplicate modality tokens dependent on batch composition. + if has_image and not has_image_embeds: + shape = next(x.shape for x in images if x is not None) + result["modality_inputs"] = {"image": torch.stack([ + x if x is not None else torch.zeros(shape) for x in images + ]), "image_mask": torch.tensor([x is not None for x in images])} + if has_audio and not has_audio_embeds: + width = max(x.numel() for x in audios if x is not None) + result.setdefault("modality_inputs", {})["audio"] = torch.stack([ + (torch.nn.functional.pad(x.flatten(), (0, width - x.numel())) + if x is not None else torch.zeros(width)) for x in audios + ]) + result["modality_inputs"]["audio_mask"] = torch.tensor([ + x is not None for x in audios + ]) + def _pad_embeddings(values): + present = [v.float() if v is not None else None for v in values] + exemplar = next(v for v in present if v is not None) + if exemplar.dim() == 1: + exemplar = exemplar.unsqueeze(0) + width = exemplar.shape[-1] + lengths = [v.shape[-2] if v is not None and v.dim() > 1 else 1 for v in present] + out = exemplar.new_zeros((len(values), max(lengths), width)) + valid = torch.zeros((len(values), max(lengths)), dtype=torch.bool) + for i, value in enumerate(present): + if value is None: + continue + value = value.unsqueeze(0) if value.dim() == 1 else value + if value.shape[-1] != width: + raise ValueError("mixed pre-encoded modality embedding widths in one batch") + out[i, :value.shape[-2]] = value + valid[i, :value.shape[-2]] = True + return out, valid + if has_image_embeds: + image_batch, image_valid = _pad_embeddings(image_embeds) + result.setdefault("modality_inputs", {})["image_embeds"] = image_batch + result["modality_inputs"]["image_embeds_mask"] = image_valid + if has_audio_embeds: + audio_batch, audio_valid = _pad_embeddings(audio_embeds) + result.setdefault("modality_inputs", {})["audio_embeds"] = audio_batch + result["modality_inputs"]["audio_embeds_mask"] = audio_valid + return result + + # ── SFT (conversations → assistant-token loss mask) ───────────────── class StreamingSFTDataset(_StreamingJSONLBase): @@ -256,6 +639,63 @@ def __getitem__(self, index): ) +class MultimodalStreamingSFTDataset(StreamingSFTDataset): + """SFT records with lazy media and optional pre-encoded features.""" + + def __init__(self, data_path, tokenizer, max_length=1024, *, image_size=224, + audio_sample_rate=16000, max_audio_seconds=30.0, + media_error_policy="warn"): + super().__init__(data_path, tokenizer, max_length=max_length) + self.image_size = int(image_size) + self.audio_sample_rate = int(audio_sample_rate) + self.max_audio_seconds = float(max_audio_seconds) + self.media_error_policy = str(media_error_policy) + self.modality_summary = dataset_modality_summary(data_path) + + def __getitem__(self, index): + input_ids, labels = super().__getitem__(index) + record = self._read_record(index) + if not detect_modalities(record): + return input_ids, labels + normalized = normalize_multimodal_record(record) + media = {} + for name in ("image_embeds", "audio_embeds"): + value = normalized["modalities"].get(name) + if value is not None: + try: + media[name] = torch.as_tensor(value, dtype=torch.float32) + except Exception as exc: + if self.media_error_policy == "raise": + raise ValueError(f"invalid {name}: {exc}") from exc + warnings.warn(f"invalid {name}: {exc}", RuntimeWarning, stacklevel=2) + for modality, loader in ( + ("image", lambda v: _load_image(v, self.image_size)), + ("audio", lambda v: _load_audio(v, self.audio_sample_rate, self.max_audio_seconds)), + ): + if f"{modality}_embeds" in media: + continue + value = _media_value(normalized["modalities"].get(modality)) + if value is None: + continue + try: + media[modality] = loader(value) + except Exception as exc: + if self.media_error_policy == "raise": + raise + warnings.warn(str(exc), RuntimeWarning, stacklevel=2) + if not media: + return input_ids, labels + return { + "input_ids": input_ids, "labels": labels, + "modality_inputs": media, + "modality_mask": torch.tensor([ + float("image" in media or "image_embeds" in media), + float("audio" in media or "audio_embeds" in media), + ]), + "metadata": normalized.get("metadata", {}), + } + + # ── DPO / ORPO (chosen + rejected pair) ───────────────────────────── class StreamingDPODataset(_StreamingJSONLBase): diff --git a/chronos/deps.py b/chronos/deps.py index ad150ad..7a01a4a 100644 --- a/chronos/deps.py +++ b/chronos/deps.py @@ -26,6 +26,14 @@ def _find_local() -> Path | None: + # A bundle ships a vendored copy (``/third_party/minimind``) and + # points here, so a recipient never needs git or network access. + override = os.environ.get("CHRONOS_MINIMIND_DIR", "").strip() + if override: + p = Path(override).expanduser() + if not (p / "model" / "model_minimind.py").exists(): + raise FileNotFoundError(f"CHRONOS_MINIMIND_DIR={override!r} has no model/model_minimind.py") + return p for p in _CANDIDATE_PATHS: if (p / "model" / "model_minimind.py").exists(): return p diff --git a/chronos/eval/benchmark.py b/chronos/eval/benchmark.py index f2073a2..75be33f 100644 --- a/chronos/eval/benchmark.py +++ b/chronos/eval/benchmark.py @@ -43,7 +43,7 @@ def compute_perplexity(model, dataloader, device, max_batches=50): if i >= max_batches: break input_ids, labels = input_ids.to(device), labels.to(device) - out, _ = model(input_ids, labels=labels) + out, _ = model(input_ids, labels=labels, legacy_output=True) n_tokens = (labels != -100).sum().item() total_loss += out.loss.item() * n_tokens total_tokens += n_tokens @@ -56,25 +56,20 @@ def measure_throughput(model, input_ids, device, max_new_tokens=64): input_ids = input_ids.to(device) t0 = time.monotonic() with torch.no_grad(): - out, lp = model(input_ids, use_cache=True) + out, lp = model(input_ids, use_cache=True, legacy_output=True) past_kv = out.past_key_values for _ in range(max_new_tokens): x = torch.randint(0, model.config.vocab_size, (1, 1), device=device) - out2, _ = model(x, past_key_values=past_kv, use_cache=True) + out2, _ = model(x, past_key_values=past_kv, use_cache=True, legacy_output=True) past_kv = out2.past_key_values elapsed = time.monotonic() - t0 return max_new_tokens / elapsed def kv_cache_memory_bytes(past_key_values, model): - """Estimate total KV cache memory in bytes (fp32).""" - total = 0 - for i, kv in enumerate(past_key_values): - if kv is None: - continue - for t in kv: - total += t.numel() * t.element_size() - return total + """Bytes held by the KV cache (any dtype, int8 entries include their scales).""" + from chronos.model.kv_quant import entry_nbytes + return sum(entry_nbytes(kv) for kv in past_key_values if kv is not None) def run_benchmark(args): @@ -132,11 +127,11 @@ def run_benchmark(args): model.eval() with torch.no_grad(): x = torch.randint(0, config.vocab_size, (1, 32), device=args.device) - out, _ = model(x, use_cache=True) + out, _ = model(x, use_cache=True, legacy_output=True) past_kv = out.past_key_values for _ in range(args.max_new_tokens): xn = torch.randint(0, config.vocab_size, (1, 1), device=args.device) - out2, _ = model(xn, past_key_values=past_kv, use_cache=True) + out2, _ = model(xn, past_key_values=past_kv, use_cache=True, legacy_output=True) past_kv = out2.past_key_values kv_bytes = kv_cache_memory_bytes(past_kv, model) @@ -149,7 +144,7 @@ def run_benchmark(args): layer_type = type(model.model.layers[i].self_attn).__name__ seq_len = kv[0].shape[1] print(f" Layer {i:2d} ({layer_type:25s}): seq_len={seq_len:4d} " - f"bytes={sum(t.numel()*t.element_size() for t in kv)//1024}KB") + f"bytes={kv_cache_memory_bytes([kv], model)//1024}KB") print(f"\n=== Summary ===") print(f" PPL={ppl:.2f} TPS={tps:.1f} KV={kv_bytes//1024}KB " diff --git a/chronos/eval/io_profiler.py b/chronos/eval/io_profiler.py index 85f2ae4..1c08343 100644 --- a/chronos/eval/io_profiler.py +++ b/chronos/eval/io_profiler.py @@ -44,7 +44,7 @@ def collect_ground_truth_routing(model, input_ids, device): model.eval() input_ids = input_ids.to(device) with torch.no_grad(): - _, _ = model(input_ids, use_cache=False) + _, _ = model(input_ids, use_cache=False, legacy_output=True) layers = [l for l in model.model.layers if isinstance(l.mlp, ChronosMOEFeedForward)] if not layers: diff --git a/chronos/export.py b/chronos/export.py index a182e5f..2ed0b64 100644 --- a/chronos/export.py +++ b/chronos/export.py @@ -32,6 +32,8 @@ "q8_0-safetensors", "fp16-gguf", "q8_0-gguf", + "chronos-native-gguf", + "ollama-compatible-gguf", ) @@ -99,6 +101,8 @@ def resolve_export_artifact(path: str) -> str: "model.q8_0.safetensors", "model.fp16.gguf", "model.q8_0.gguf", + "model.chronos-native.gguf", + "model.ollama-compatible.gguf", ): candidate = os.path.join(path, name) if os.path.exists(candidate): @@ -320,6 +324,8 @@ def _normalize_formats(formats: Iterable[str] | str | None) -> list[str]: fmt = fmt.lower().replace("fp16_safetensors", "fp16-safetensors") fmt = fmt.replace("q8_0_safetensors", "q8_0-safetensors") fmt = fmt.replace("fp16_gguf", "fp16-gguf").replace("q8_0_gguf", "q8_0-gguf") + fmt = fmt.replace("chronos_native_gguf", "chronos-native-gguf") + fmt = fmt.replace("ollama_compatible_gguf", "ollama-compatible-gguf") if fmt == "all": for known in EXPORT_FORMATS: if known not in out: @@ -428,6 +434,9 @@ def _cluster_manifest_summary(expert_cache_dir: str | None) -> dict | None: "storage_format": data.get("storage_format", "safetensors"), "num_experts": data.get("num_experts"), "num_layers": data.get("num_layers"), + "version": int(data.get("version", 1)), + "layer_aware": bool(data.get("layer_aware", False)), + "layer_to_group": data.get("layer_to_group", {}), } @@ -509,6 +518,7 @@ def _export_metadata( "config_sources": list(sources), "architecture": "chronos", "chronos": { + "architecture_version": cfg.get("architecture_version", 1), "hidden_size": cfg.get("hidden_size"), "num_hidden_layers": cfg.get("num_hidden_layers"), "num_experts": cfg.get("num_experts"), @@ -526,6 +536,21 @@ def _export_metadata( "lambda_lookahead_topk": cfg.get("lambda_lookahead_topk"), "offload_miss_policy_default": cfg.get("offload_miss_policy_default", "on_demand"), "recommended_resident_experts": cfg.get("recommended_resident_experts"), + "recurrent_enabled": cfg.get("recurrent_enabled", False), + "recurrent_groups": cfg.get("recurrent_groups"), + "recurrent_depth": cfg.get("recurrent_depth", 1), + "min_recurrent_depth": cfg.get("min_recurrent_depth", 1), + "max_inference_depth": cfg.get("max_inference_depth"), + "recurrent_latent_refinement": cfg.get("recurrent_latent_refinement", True), + "recurrent_strategy": cfg.get("recurrent_strategy", "scheduled"), + "recurrent_depth_jitter": cfg.get("recurrent_depth_jitter", 0.0), + "recurrent_parallel_refine": cfg.get("recurrent_parallel_refine", False), + "adaptive_halting": cfg.get("adaptive_halting", False), + "halting_threshold": cfg.get("halting_threshold", 0.99), + "layer_aware_cache": cfg.get("layer_aware_cache", False), + "layer_aware_lookahead": cfg.get("layer_aware_lookahead", False), + "modalities": cfg.get("modalities", []), + "modality_adapters": cfg.get("modality_adapters", {}), }, "export_contract": "predictive_lazy_routed_experts_or_shared_fallback + shared_residual", "offload_features": { @@ -538,6 +563,16 @@ def _export_metadata( "cluster_aware_safetensors": bool(_cluster_manifest_summary(expert_cache_dir)), }, } + if fmt == "ollama-compatible-gguf": + payload["compatibility_notice"] = { + "status": "manifest-only", + "stock_runtime_executable": False, + "reason": ( + "Chronos MLA/sliding attention and layer-aware MoE tensors are " + "not yet materialized into standard llama.cpp tensor names." + ), + "required_backend": "chronos-aware runtime", + } if cluster_build: payload["cluster_build"] = cluster_build payload.update(_sidecar_config_for_export(checkpoint_path, config)) @@ -677,7 +712,12 @@ def export_gguf( ) -> ExportResult: os.makedirs(output_dir, exist_ok=True) quantized = fmt == "q8_0-gguf" - filename = "model.q8_0.gguf" if quantized else "model.fp16.gguf" + if fmt == "chronos-native-gguf": + filename = "model.chronos-native.gguf" + elif fmt == "ollama-compatible-gguf": + filename = "model.ollama-compatible.gguf" + else: + filename = "model.q8_0.gguf" if quantized else "model.fp16.gguf" path = os.path.join(output_dir, filename) cfg = config_to_dict(config) @@ -691,12 +731,22 @@ def export_gguf( cluster_build=cluster_build, )["chronos_export"] ) + compatibility_mode = "expanded" if fmt == "ollama-compatible-gguf" else "native" kvs = [ _gguf_kv_string("general.architecture", "chronos"), _gguf_kv_string("general.name", "Project Chronos"), _gguf_kv_string("chronos.export.format", fmt), _gguf_kv_string("chronos.export.metadata_json", json.dumps(metadata_payload, sort_keys=True)), ] + kvs.extend([ + _gguf_kv_string("chronos.compatibility_mode", compatibility_mode), + _gguf_kv_string( + "chronos.runtime_features", + "disabled" if compatibility_mode == "expanded" else "recurrent,moe,paging,lookahead", + ), + _gguf_kv_string("chronos.stock_runtime_executable", "false"), + _gguf_kv_string("chronos.modalities", json.dumps(cfg.get("modalities", []))), + ]) for key in ( "hidden_size", "num_hidden_layers", @@ -713,6 +763,11 @@ def export_gguf( "sliding_window_size", "vocab_size", "recommended_resident_experts", + "architecture_version", + "recurrent_depth", + "recurrent_groups", + "min_recurrent_depth", + "max_inference_depth", ): value = cfg.get(key) if isinstance(value, int): @@ -792,11 +847,12 @@ def _write_deployment_notes(output_dir: str, results: Sequence[ExportResult]) -> notes = [ "# Project Chronos GGUF deployment", "", - "These files are standard GGUF containers with `general.architecture=chronos`.", - "They are intended for Chronos-aware runtimes. Stock Ollama/llama.cpp builds", - "that do not implement the Chronos architecture cannot execute the model", - "correctly because Chronos uses hybrid MLA/sliding-window attention,", - "lookahead expert prediction, and lazy MoE expert loading.", + "Chronos-native files are GGUF containers with `general.architecture=chronos`", + "and require a Chronos-aware runtime. Any `ollama-compatible` artifact in", + "this export is currently a compatibility manifest only and explicitly has", + "`chronos.stock_runtime_executable=false`; stock Ollama/llama.cpp cannot run", + "it until MLA/sliding attention, recurrent blocks, MoE and tensor names have", + "been materialized into a supported standard architecture.", "", "For an Ollama build with Chronos architecture support, use a Modelfile like:", "", @@ -812,6 +868,18 @@ def _write_deployment_notes(output_dir: str, results: Sequence[ExportResult]) -> stem = name.replace("model.", "").replace(".gguf", "") with open(os.path.join(output_dir, f"Modelfile.{stem}"), "w", encoding="utf-8") as f: f.write(f"FROM ./{name}\nPARAMETER temperature 0.7\n") + native = next((n for n in gguf_files if "chronos-native" in n), None) + compat = next((n for n in gguf_files if "ollama-compatible" in n), None) + if native: + with open(os.path.join(output_dir, "Modelfile.chronos-native"), "w", encoding="utf-8") as f: + f.write(f"FROM ./{native}\nPARAMETER temperature 0.7\n") + if compat: + with open(os.path.join(output_dir, "Modelfile.ollama-compatible"), "w", encoding="utf-8") as f: + f.write( + "# Compatibility manifest only: stock Ollama execution is not yet supported.\n" + f"# Use a Chronos-aware backend for {compat}.\n" + f"FROM ./{compat}\nPARAMETER temperature 0.7\n" + ) def export_checkpoint( @@ -921,7 +989,7 @@ def main(argv: Sequence[str] | None = None) -> None: "--formats", nargs="+", default=["fp16-safetensors"], - help="Any of: fp16-safetensors q8_0-safetensors fp16-gguf q8_0-gguf all", + help="Any of: fp16-safetensors q8_0-safetensors fp16-gguf q8_0-gguf chronos-native-gguf ollama-compatible-gguf all", ) parser.add_argument("--config_path", default=None) parser.add_argument("--expert_cache_dir", default=None) diff --git a/chronos/io/async_prefetcher.py b/chronos/io/async_prefetcher.py index ddb1828..bf35195 100644 --- a/chronos/io/async_prefetcher.py +++ b/chronos/io/async_prefetcher.py @@ -57,7 +57,10 @@ def __init__( self._done_lock = threading.Lock() self._stop_event = threading.Event() self._thread: Optional[threading.Thread] = None - self._stats = {"prefetch_hits": 0, "prefetch_misses": 0, "total_requests": 0} + self._stats = { + "prefetch_hits": 0, "prefetch_misses": 0, "total_requests": 0, + "queue_drops": 0, "wait_time_s": 0.0, + } def start(self): self._stop_event.clear() @@ -78,7 +81,7 @@ def stop(self): if self._thread: self._thread.join(timeout=2.0) - def submit(self, expert_ids: List[int]): + def submit(self, expert_ids): """ Submit a list of expert IDs to prefetch (non-blocking). Called after each decode step with LookaheadRouter predictions. @@ -93,23 +96,27 @@ def submit(self, expert_ids: List[int]): self._request_queue.get_nowait() self._request_queue.put_nowait(expert_ids) except (queue.Empty, queue.Full): - pass + with self._done_lock: + self._stats["queue_drops"] += 1 def wait_for(self, expert_id: int, timeout: float = 0.05) -> bool: """ Wait until expert_id is in RAM (or VRAM). Returns True if ready. In practice this should return immediately if prefetch worked. """ - deadline = time.monotonic() + timeout + started = time.monotonic() + deadline = started + timeout while time.monotonic() < deadline: if (self.store.ram_lru.contains(expert_id) or self.store.vram_lru.contains(expert_id)): with self._done_lock: self._stats["prefetch_hits"] += 1 + self._stats["wait_time_s"] += time.monotonic() - started return True time.sleep(0.001) with self._done_lock: self._stats["prefetch_misses"] += 1 + self._stats["wait_time_s"] += time.monotonic() - started return False def _prefetch_loop(self): @@ -142,8 +149,12 @@ def stats(self) -> dict: with self._done_lock: s = dict(self._stats) total = s["prefetch_hits"] + s["prefetch_misses"] - s["hit_rate"] = s["prefetch_hits"] / max(total, 1) - _metric("prefetch_hit_rate", s["hit_rate"]) + s["prefetch_observations"] = total + s["hit_rate"] = s["prefetch_hits"] / total if total else None + s["prefetch_queue_drops"] = int(s.pop("queue_drops", 0)) + s["prefetch_wait_time_s"] = float(s.pop("wait_time_s", 0.0)) + if total: + _metric("prefetch_hit_rate", s["hit_rate"]) return s @@ -189,13 +200,15 @@ def promote_current( (already-resident experts are skipped). """ newly = [] - for eid in current_step_expert_ids: + for eid in self.store._expand_keys(current_step_expert_ids): if not self.store.vram_lru.contains(eid): if self.store.promote_to_vram(eid, blocking=blocking): newly.append(eid) + else: + self.store.touch_expert(eid) return newly - def prefetch_only(self, lookahead_probs) -> List[int]: + def prefetch_only(self, lookahead_probs): """ Submit future expert predictions to the async prefetcher. Does NOT touch VRAM. Safe to call *before* the forward pass so the @@ -208,13 +221,35 @@ def prefetch_only(self, lookahead_probs) -> List[int]: import torch top_k = int(getattr(self.store.config, "num_experts_per_tok", 1) or 1) top_k = max(1, min(top_k, lookahead_probs.shape[-1])) - # lookahead_probs: [B, S, K+1, E] — take last token, steps 1..K. + layer_aware = lookahead_probs.dim() == 5 and bool(getattr(self.store.config, "layer_aware_cache", False)) + # v1: [B,S,K+1,E]. v2: [B,S,K+1,G,E]. # We submit t+1 first (highest priority), then t+2, etc., so the # FIFO queue naturally orders them by recency. - last_token_probs = lookahead_probs[:, -1, 1:, :] # [B, K, E] + if layer_aware: + last_token_probs = lookahead_probs[:, -1, 1:, :, :] # [B,K,G,E] + else: + if lookahead_probs.dim() == 5: + last_token_probs = lookahead_probs[:, -1, 1:, :, :].mean(dim=2) + else: + last_token_probs = lookahead_probs[:, -1, 1:, :] K = last_token_probs.shape[1] - all_submitted: List[int] = [] + all_submitted = [] for k in range(K): + if layer_aware: + for group in range(last_token_probs.shape[2]): + ids_k = torch.topk(last_token_probs[:, k, group, :], k=top_k, dim=-1).indices + ids_k = ids_k.reshape(-1).unique().cpu().tolist() + layer_ids = [ + int(li) for li, gi in getattr(self.store, "_layer_to_group", {}).items() + if int(gi) == group + ] + if not layer_ids: + layer_ids = [group] + keys = [(li, group, int(eid)) for li in layer_ids for eid in ids_k] + if keys: + self.prefetcher.submit(keys) + all_submitted.extend(keys) + continue ids_k = torch.topk(last_token_probs[:, k, :], k=top_k, dim=-1).indices ids_k = ids_k.reshape(-1).unique().cpu().tolist() if ids_k: diff --git a/chronos/io/cluster_layout.py b/chronos/io/cluster_layout.py index 273512f..8f9ef54 100644 --- a/chronos/io/cluster_layout.py +++ b/chronos/io/cluster_layout.py @@ -349,17 +349,33 @@ def repack_expert_weights_safetensors( output_dir=output_dir, num_layers=num_layers, num_experts=num_experts, - dtype=dtype if dtype is not None else _torch.float16, + dtype=dtype, + layer_to_group={ + int(li): int(li - int(getattr(model.config, "prelude_layers", 0))) + if int(getattr(model.config, "prelude_layers", 0)) <= li < ( + int(getattr(model.config, "prelude_layers", 0)) + + int(getattr(model.config, "recurrent_groups", num_layers)) + ) else -1 + for li in range(num_layers) + }, ) return manifest_path -def _parse_torch_dtype(dtype: str | torch.dtype | None) -> torch.dtype: +def _model_dtype(model) -> torch.dtype: + from chronos.precision import torch_model_dtype + return torch_model_dtype(model) + + +def _parse_torch_dtype(dtype: str | torch.dtype | None) -> torch.dtype | None: + """``None`` / ``"auto"`` -> keep the model's own dtype (no conversion).""" if isinstance(dtype, torch.dtype): return dtype if dtype is None: - return torch.float16 + return None normalized = str(dtype).lower().replace("torch.", "") + if normalized in {"auto", "native", "keep"}: + return None if normalized in {"fp16", "float16", "half"}: return torch.float16 if normalized in {"bf16", "bfloat16"}: @@ -380,7 +396,7 @@ def build_clustered_expert_cache( max_batches: int = 50, batch_size: int = 4, max_seq_len: int = 256, - dtype: str | torch.dtype | None = torch.float16, + dtype: str | torch.dtype | None = None, source_checkpoint: str | None = None, ) -> dict: """ @@ -445,7 +461,7 @@ def build_clustered_expert_cache( "max_batches": int(max_batches), "batch_size": int(batch_size), "max_seq_len": int(max_seq_len), - "dtype": str(_parse_torch_dtype(dtype)).replace("torch.", ""), + "dtype": str(_parse_torch_dtype(dtype) or _model_dtype(model)).replace("torch.", ""), "clusters": clusters, "cooccurrence_matrix": np.round(C, 6).tolist(), } @@ -467,7 +483,7 @@ def build_clustered_expert_cache_from_checkpoint( max_batches: int = 50, batch_size: int = 4, max_seq_len: int = 256, - dtype: str | torch.dtype | None = torch.float16, + dtype: str | torch.dtype | None = None, ) -> dict: from chronos.model.model_chronos import ChronosForCausalLM from chronos.model.checkpoint import ( @@ -483,6 +499,11 @@ def build_clustered_expert_cache_from_checkpoint( ) model = ChronosForCausalLM(config) weights = load_checkpoint_state_dict(model_path, map_location="cpu") + if _parse_torch_dtype(dtype) is None: + # Keep the checkpoint's stored precision end to end (BF16 stays BF16). + stored = next((v.dtype for v in weights.values() + if torch.is_tensor(v) and v.is_floating_point()), torch.float32) + model = model.to(stored) load_state_dict_controlled(model, weights) return build_clustered_expert_cache( model, @@ -514,6 +535,8 @@ def main(): parser.add_argument("--max_batches", type=int, default=50) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--max_seq_len", type=int, default=256) + parser.add_argument("--dtype", default="auto", + help="expert page dtype: auto (keep checkpoint dtype)|float32|bfloat16|float16") args = parser.parse_args() if args.format == "safetensors": @@ -527,7 +550,7 @@ def main(): max_batches=args.max_batches, batch_size=args.batch_size, max_seq_len=args.max_seq_len, - dtype=torch.float16, + dtype=args.dtype, ) print(json.dumps(summary, indent=2, ensure_ascii=False)) return diff --git a/chronos/io/expert_store.py b/chronos/io/expert_store.py index d7f942b..6145604 100644 --- a/chronos/io/expert_store.py +++ b/chronos/io/expert_store.py @@ -16,7 +16,7 @@ import os import threading import collections -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Tuple, Union import torch import torch.nn as nn @@ -31,6 +31,29 @@ def _metric(name, value): # fallback pass +def expert_budget_pages(config, page_bytes: int): + """(max_resident_pages, ram_budget_pages) from the hard budget fields. + + ``max_resident_experts`` counts fast-tier pages directly; ``ram_budget_mb`` + is converted with the page size. Either may be None (no hard cap). + """ + resident = getattr(config, "max_resident_experts", None) + budget_mb = getattr(config, "ram_budget_mb", None) + resident = None if resident in (None, "", 0) else max(1, int(resident)) + ram = None + if budget_mb not in (None, "", 0): + ram = max(1, int(float(budget_mb) * 2 ** 20 // max(int(page_bytes), 1))) + return resident, ram + + +def _process_rss_bytes() -> int: + try: + import psutil + return int(psutil.Process().memory_info().rss) + except Exception: + return 0 + + class LRUCache: """Thread-safe LRU cache tracking which expert IDs are in a tier.""" @@ -38,6 +61,20 @@ def __init__(self, capacity: int): self.capacity = capacity self._cache: collections.OrderedDict = collections.OrderedDict() self._lock = threading.Lock() + # Keys the predictor says are about to be used: evicted only when + # every resident key is protected (the hard capacity always holds). + self.protected: frozenset = frozenset() + + def _victim_locked(self): + for key in self._cache: + if key not in self.protected: + return key + return next(iter(self._cache)) + + def victim(self): + """Key the next insertion would evict (coldest unprotected), or None.""" + with self._lock: + return self._victim_locked() if self._cache else None def contains(self, key: int) -> bool: with self._lock: @@ -56,7 +93,8 @@ def put(self, key: int) -> Optional[int]: return None evicted = None if len(self._cache) >= self.capacity: - evicted, _ = self._cache.popitem(last=False) + evicted = self._victim_locked() + self._cache.pop(evicted) self._cache[key] = True return evicted @@ -73,7 +111,7 @@ def reorder_by_score(self, scores: Dict[int, float]): with self._lock: ordered = sorted( self._cache.keys(), - key=lambda key: (float(scores.get(key, 0.0)), key), + key=lambda key: (float(scores.get(key, 0.0)), repr(key)), ) self._cache = collections.OrderedDict((key, True) for key in ordered) @@ -121,14 +159,18 @@ def __init__(self, model, config, ssd_dir: str = "./expert_cache"): ] self.num_experts = config.num_experts self.num_layers = len(self.moe_layers) + self.layer_aware = bool(getattr(config, "layer_aware_cache", False)) + self._layer_to_group = self._build_layer_group_map() # Estimate VRAM capacity in number of experts + self._expert_dtype_cached = self._expert_dtype() expert_bytes = self._expert_size_bytes() self._expert_bytes = expert_bytes - self._expert_dtype_cached = self._expert_dtype() + page_bytes = expert_bytes * (1 if self.layer_aware else self.num_layers) + total_pages = self.num_experts * (self.num_layers if self.layer_aware else 1) vram_bytes = config.vram_budget_gb * (1024 ** 3) # Reserve 50% headroom for shared experts + attention KV cache - self.vram_capacity = max(1, int(vram_bytes * 0.5 / max(expert_bytes, 1))) + self.vram_capacity = min(total_pages, max(1, int(vram_bytes * 0.5 / max(page_bytes, 1)))) # Pinned RAM capacity: bounded by physical RAM safety limit # Never exceed pinned_memory_max_fraction of total physical RAM @@ -136,10 +178,18 @@ def __init__(self, model, config, ssd_dir: str = "./expert_cache"): physical_ram_bytes = self._physical_ram_bytes() safe_pinned_bytes = physical_ram_bytes * max_fraction # Each expert occupies expert_bytes * num_layers in RAM - per_expert_ram = max(expert_bytes * self.num_layers, 1) + per_expert_ram = max(expert_bytes if self.layer_aware else expert_bytes * self.num_layers, 1) ram_capacity_by_safety = max(1, int(safe_pinned_bytes / per_expert_ram)) # Also cap at 4x VRAM capacity (original heuristic) - self.ram_capacity = min(self.vram_capacity * 4, ram_capacity_by_safety) + self.ram_capacity = min(total_pages, self.vram_capacity * 4, ram_capacity_by_safety) + # Hard user budget (max_resident_experts / ram_budget_mb) on top of + # the heuristics; counted in cache pages, at least one page each. + self.max_resident_pages, self.ram_budget_pages = expert_budget_pages(config, page_bytes) + if self.max_resident_pages is not None: + self.vram_capacity = min(self.vram_capacity, self.max_resident_pages) + if self.ram_budget_pages is not None: + self.ram_capacity = min(self.ram_capacity, self.ram_budget_pages) + self._peak_rss_bytes = 0 self.vram_lru = LRUCache(capacity=self.vram_capacity) self.ram_lru = LRUCache(capacity=self.ram_capacity) @@ -147,6 +197,8 @@ def __init__(self, model, config, ssd_dir: str = "./expert_cache"): # RAM pinned buffers: {expert_id: {layer_idx: tensor_dict}} self._ram_buffers: Dict[int, Dict[int, Dict[str, torch.Tensor]]] = {} self._ram_lock = threading.Lock() + self._load_lock = threading.RLock() + self._io_stats = {"ssd_page_reads": 0, "expert_evictions": 0} self._offloaded = False # M3: per-expert H2D events. Every promote_to_vram records an event @@ -164,6 +216,40 @@ def __init__(self, model, config, ssd_dir: str = "./expert_cache"): # can run cluster-aware from step 0. self._cluster_storage = ClusterStorage(ssd_dir) + def _build_layer_group_map(self) -> Dict[int, int]: + prelude = int(getattr(self.config, "prelude_layers", 0) or 0) + groups = int(getattr(self.config, "recurrent_groups", self.num_layers) or self.num_layers) + mapping = {} + for li in range(self.num_layers): + mapping[li] = li - prelude if prelude <= li < prelude + groups else -1 + return mapping + + def expert_key(self, layer_id: int, expert_id: int) -> tuple[int, int, int]: + return (int(layer_id), int(self._layer_to_group.get(int(layer_id), -1)), int(expert_id)) + + def _expand_keys(self, items) -> list: + """Normalize legacy expert IDs and v2 page keys.""" + keys = [] + for item in items or []: + if isinstance(item, (tuple, list)): + if len(item) == 3: + keys.append((int(item[0]), int(item[1]), int(item[2]))) + elif len(item) == 2: + keys.append(self.expert_key(int(item[0]), int(item[1]))) + else: + raise ValueError(f"invalid expert page key: {item!r}") + elif self.layer_aware: + keys.extend(self.expert_key(li, int(item)) for li in range(self.num_layers)) + else: + keys.append(int(item)) + return list(dict.fromkeys(keys)) + + @staticmethod + def _page_parts(key) -> tuple[int, int]: + if not isinstance(key, tuple): + raise TypeError(f"expected a layer-aware expert key, got {key!r}") + return int(key[0]), int(key[2]) + @staticmethod def _physical_ram_bytes() -> int: """Query available (not total) RAM at call time for dynamic safety.""" @@ -180,18 +266,26 @@ def _safe_ram_capacity(self) -> int: max_fraction = getattr(self.config, 'pinned_memory_max_fraction', 0.25) available_bytes = self._physical_ram_bytes() safe_bytes = available_bytes * max_fraction - per_expert_ram = max(self._expert_size_bytes() * self.num_layers, 1) - return max(1, min( - self.ram_capacity, - self.vram_capacity * 4, - int(safe_bytes / per_expert_ram), - )) + per_expert_ram = max( + self._expert_size_bytes() if self.layer_aware else self._expert_size_bytes() * self.num_layers, + 1, + ) + cap = min(self.ram_capacity, self.vram_capacity * 4, int(safe_bytes / per_expert_ram)) + if getattr(self, "ram_budget_pages", None) is not None: + cap = min(cap, self.ram_budget_pages) + return max(1, cap) + + def set_protected_experts(self, keys) -> None: + """Pin predicted pages against eviction (expanded like prefetch keys).""" + protected = frozenset(self._expand_keys(list(keys or ()))) + self.vram_lru.protected = protected + self.ram_lru.protected = protected def _expert_size_bytes(self) -> int: if not self.moe_layers: return 0 expert = self.moe_layers[0].experts[0] - size = sum(p.numel() * 2 for p in expert.parameters()) # fp16 + size = sum(p.numel() * p.element_size() for p in expert.parameters()) if size: return size cached = getattr(self, "_expert_bytes", 0) @@ -203,20 +297,26 @@ def _expert_size_bytes(self) -> int: hidden = int(getattr(self.config, "hidden_size", 0) or 0) intermediate = int(getattr(self.config, "moe_intermediate_size", 0) or 0) if hidden and intermediate: - return 3 * hidden * intermediate * 2 # gate/up/down, fp16 bytes + dtype = self._expert_dtype() + return 3 * hidden * intermediate * torch.empty((), dtype=dtype).element_size() return 0 def _expert_dtype(self) -> torch.dtype: - if not self.moe_layers: - return torch.float16 - expert = self.moe_layers[0].experts[0] - for p in expert.parameters(): - return p.dtype - return getattr(self, "_expert_dtype_cached", torch.float16) + """Expert pages live in the trunk's compute dtype (one precision + policy end to end): a BF16 trunk gets BF16 pages, an FP32 CPU trunk + FP32 pages. SSD shards of another dtype are converted once when they + are read into RAM, never per forward.""" + from chronos.precision import torch_model_dtype + try: + return torch_model_dtype(self.model) + except StopIteration: + return torch.float32 def _build_live_expert(self): from model.model_minimind import FeedForward - return FeedForward(self.config, intermediate_size=self.config.moe_intermediate_size).to(self.device) + # No random initialization or temporary full allocation on a cache miss. + with torch.device("meta"): + return FeedForward(self.config, intermediate_size=self.config.moe_intermediate_size) def _build_placeholder_expert(self): from chronos.model.moe_chronos import LazyExpertPlaceholder @@ -239,7 +339,7 @@ def _expert_state_from_checkpoint( state = {} for key, val in state_dict.items(): if key.startswith(prefix): - state[key[len(prefix):]] = val.detach().to(dtype=torch.float16, device="cpu").contiguous() + state[key[len(prefix):]] = val.detach().cpu().contiguous() return state def offload_from_state_dict( @@ -264,8 +364,6 @@ def offload_from_state_dict( for li in range(self.num_layers): for ei in range(self.num_experts): path = self._ssd_path(ei, li) - if os.path.exists(path): - continue state = self._expert_state_from_checkpoint(state_dict, li, ei) if state: torch.save(state, path) @@ -293,8 +391,6 @@ def offload_from_state_reader( for li in range(self.num_layers): for ei in range(self.num_experts): path = self._ssd_path(ei, li) - if os.path.exists(path): - continue prefix = f"model.layers.{li}.mlp.experts.{ei}." state = {} for key in keys: @@ -302,7 +398,7 @@ def offload_from_state_reader( state[key[len(prefix):]] = ( reader.get_tensor(key) .detach() - .to(dtype=torch.float16, device="cpu") + .cpu() .contiguous() ) if state: @@ -342,23 +438,22 @@ def _offload_legacy_pt(self): for li, moe in enumerate(self.moe_layers): for ei, expert in enumerate(moe.experts): path = self._ssd_path(ei, li) - if not os.path.exists(path): - state = {k: v.half().cpu() for k, v in expert.state_dict().items()} - torch.save(state, path) + state = {k: v.detach().cpu().contiguous() for k, v in expert.state_dict().items()} + torch.save(state, path) self.replace_live_experts_with_placeholders() def _offload_safetensors(self, clusters: Optional[List[List[int]]]): if clusters is None: - # Trivial single-cluster layout — better than per-expert pickles - # because we still get one mmap-friendly contiguous file. - clusters = [list(range(self.num_experts))] + # Without calibrated co-activation data, do not couple every expert. + clusters = [[eid] for eid in range(self.num_experts)] ClusterStorage.write_clusters( moe_layers=self.moe_layers, clusters=clusters, output_dir=self.ssd_dir, num_layers=self.num_layers, num_experts=self.num_experts, - dtype=torch.float16, + dtype=self._expert_dtype_cached, + layer_to_group=self._layer_to_group if self.layer_aware else None, ) self._cluster_storage = ClusterStorage(self.ssd_dir) self.replace_live_experts_with_placeholders() @@ -370,6 +465,18 @@ def attach_cluster_manifest(self, manifest_dir: Optional[str] = None): if ClusterStorage.has_manifest(target): self._cluster_storage = ClusterStorage(target) self.storage_format = "safetensors" + issues = self._cluster_storage.manifest.validate( + num_layers=self.num_layers, num_experts=self.num_experts + ) + if issues: + import warnings + warnings.warn( + "cluster manifest topology mismatch: " + "; ".join(issues), + RuntimeWarning, + stacklevel=2, + ) + return not issues + return False def _load_from_ssd(self, expert_id: int) -> Dict[int, Dict[str, torch.Tensor]]: """Load expert weights for all layers from SSD into CPU tensors. @@ -386,7 +493,7 @@ def _load_from_ssd(self, expert_id: int) -> Dict[int, Dict[str, torch.Tensor]]: # ── RAM operations ──────────────────────────────────────────── - def prefetch_to_ram(self, expert_ids: List[int]): + def prefetch_to_ram(self, expert_ids): """ Synchronously load expert weights from SSD into pinned RAM. Call this from a background thread for async prefetch. @@ -397,87 +504,104 @@ def prefetch_to_ram(self, expert_ids: List[int]): random reads into sequential reads. Requests that land in the same cluster are deduplicated. """ + expert_ids = self._expand_keys(expert_ids) if not expert_ids: return - if self._cluster_storage is not None: - self._prefetch_clustered(expert_ids) - else: - self._prefetch_legacy(expert_ids) + # Serialize foreground demand loads and the background prefetcher. + # The RAM lock stays short so promotions can take a stable reference. + with self._load_lock: + if self._cluster_storage is not None: + self._prefetch_clustered(expert_ids) + else: + self._prefetch_legacy(expert_ids) - def _prefetch_clustered(self, expert_ids: List[int]): - """Group requests by cluster, load each cluster at most once.""" + def _cache_ram_entry(self, key, states): + cap = self._safe_ram_capacity() + with self._ram_lock: + self.ram_lru.capacity = cap + self.ram_lru.remove(key) + while len(self.ram_lru) >= cap: + victim = self.ram_lru.victim() + self.ram_lru.remove(victim) + self._ram_buffers.pop(victim, None) + self.ram_lru.put(key) + self._ram_buffers[key] = states + + def _prefetch_clustered(self, expert_ids): cs = self._cluster_storage - seen_clusters: set = set() - clusters_to_load: List[int] = [] - for eid in expert_ids: - if eid not in cs.manifest.expert_to_cluster: + # Keep each requested page last in its cluster batch, so neighboring + # pages can never evict the page the caller is about to promote. + for request in expert_ids: + if self.ram_lru.contains(request): + self.ram_lru.touch(request) continue - if self.ram_lru.contains(eid): - self.ram_lru.touch(eid) + li, eid = self._page_parts(request) if self.layer_aware else (None, int(request)) + if eid not in cs.manifest.expert_to_cluster: continue cid = cs.cluster_for_expert(eid) - if cid in seen_clusters: - continue - seen_clusters.add(cid) - clusters_to_load.append(cid) - - for cid in clusters_to_load: - # Simulate SSD read latency when CHRONOS_SIM_SSD_MS is set + cap = self._safe_ram_capacity() + neighbors = [e for e in cs.experts_in_cluster(cid) if e != eid] + candidates = neighbors[:max(0, cap - 1)] + [eid] + keys = [self.expert_key(li, e) if self.layer_aware else e for e in candidates] + keys = [key for key in keys if not self.ram_lru.contains(key)] + pages = set() + for key in keys: + if self.layer_aware: + pages.add(self._page_parts(key)) + else: + pages.update((layer, key) for layer in range(self.num_layers)) _maybe_sleep_ssd() - # One mmap read pulls the whole cluster — all experts, all layers. - nested = cs.load_cluster(cid, dtype=torch.float16) - # nested: {expert_id: {layer_id: {param_name: Tensor}}} - dynamic_cap = self._safe_ram_capacity() - self.ram_lru.capacity = dynamic_cap - for eid, layer_states in nested.items(): - if self.ram_lru.contains(eid): - self.ram_lru.touch(eid) - continue - pinned = self._pin_layer_states(layer_states) - with self._ram_lock: - evicted = self.ram_lru.put(eid) - if evicted is not None: - self._ram_buffers.pop(evicted, None) - self._ram_buffers[eid] = pinned + nested = cs.load_cluster(cid, pages=pages) + self._io_stats["ssd_page_reads"] += len(pages) + for key in keys: + if self.layer_aware: + layer, expert = self._page_parts(key) + state = nested.get(expert, {}).get(layer) + entry = {layer: state} if state is not None else {} + else: + entry = nested.get(key, {}) + if entry: + self._cache_ram_entry(key, self._pin_layer_states(entry)) def _prefetch_legacy(self, expert_ids: List[int]): - for eid in expert_ids: - if self.ram_lru.contains(eid): - self.ram_lru.touch(eid) + for request in expert_ids: + if self.ram_lru.contains(request): + self.ram_lru.touch(request) continue # Simulate SSD read latency _maybe_sleep_ssd() dynamic_cap = self._safe_ram_capacity() self.ram_lru.capacity = dynamic_cap - layer_states = self._load_from_ssd(eid) + if self.layer_aware: + li, eid = self._page_parts(request) + path = self._ssd_path(eid, li) + layer_states = {li: torch.load(path, map_location="cpu")} if os.path.exists(path) else {} + else: + eid = int(request) + layer_states = self._load_from_ssd(eid) if not layer_states: continue - pinned = self._pin_layer_states(layer_states) - with self._ram_lock: - evicted = self.ram_lru.put(eid) - if evicted is not None: - self._ram_buffers.pop(evicted, None) - self._ram_buffers[eid] = pinned + self._io_stats["ssd_page_reads"] += len(layer_states) + self._cache_ram_entry(request, self._pin_layer_states(layer_states)) - @staticmethod - def _pin_layer_states( - layer_states: Dict[int, Dict[str, torch.Tensor]] - ) -> Dict[int, Dict[str, torch.Tensor]]: - pinned: Dict[int, Dict[str, torch.Tensor]] = {} + def _pin_layer_states(self, layer_states): + pinned = {} for li, state in layer_states.items(): pinned[li] = {} for k, v in state.items(): - try: - pinned[li][k] = v.pin_memory() if not v.is_pinned() else v - except RuntimeError: - pinned[li][k] = v.cpu() + value = v.to(dtype=self._expert_dtype_cached, device="cpu") + # CPU/MPS have no CUDA DMA benefit. Some torch/macOS builds + # segfault here instead of raising an exception. + if self._h2d_stream is not None and not value.is_pinned(): + value = value.pin_memory() + pinned[li][k] = value return pinned # ── VRAM operations ─────────────────────────────────────────── - def promote_to_vram(self, expert_id: int, blocking: bool = False) -> bool: + def promote_to_vram(self, expert_id, blocking: bool = False) -> bool: """ Move expert from RAM → VRAM. Returns True if the H2D copy was issued. @@ -488,6 +612,20 @@ def promote_to_vram(self, expert_id: int, blocking: bool = False) -> bool: blocking=True (legacy path): synchronizes _h2d_stream before return, so the expert is guaranteed visible to the compute stream. """ + requests = self._expand_keys([expert_id]) + if len(requests) > 1: + changed = False + for request in requests: + changed = self.promote_to_vram(request, blocking=False) or changed + if blocking: + self.sync_h2d() + return changed + expert_id = requests[0] + if self.vram_lru.contains(expert_id): + self.touch_expert(expert_id) + if blocking: + self.wait_for_experts([expert_id]) + return False if not self.ram_lru.contains(expert_id): self.prefetch_to_ram([expert_id]) @@ -507,11 +645,17 @@ def promote_to_vram(self, expert_id: int, blocking: bool = False) -> bool: ) with ctx: for li, state in layer_states.items(): - expert = self.moe_layers[li].experts[expert_id] + eid = self._page_parts(expert_id)[1] if self.layer_aware else int(expert_id) + expert = self.moe_layers[li].experts[eid] if not isinstance(expert, nn.Module) or len(list(expert.parameters())) == 0: expert = self._build_live_expert() - self._replace_expert_module(li, expert_id, expert) - expert.load_state_dict({k: v.to(self.device, non_blocking=True) for k, v in state.items()}, strict=False) + self._replace_expert_module(li, eid, expert) + expert.load_state_dict( + {k: v.to(device=self.device, dtype=self._expert_dtype_cached, non_blocking=True) + for k, v in state.items()}, strict=True, assign=True, + ) + expert.requires_grad_(False) + expert.eval() # Record an event on the H2D stream so consumers can wait on # *just this expert* instead of the whole stream. if self._h2d_stream is not None: @@ -533,6 +677,7 @@ def wait_for_experts(self, expert_ids: List[int]) -> None: if self._h2d_stream is None: return current = torch.cuda.current_stream(self.device) + expert_ids = self._expand_keys(expert_ids) with self._events_lock: events = [self._pending_events.pop(eid, None) for eid in expert_ids] for evt in events: @@ -542,10 +687,11 @@ def wait_for_experts(self, expert_ids: List[int]) -> None: # the CPU too. current.wait_event(evt) - def touch_expert(self, expert_id: int) -> None: + def touch_expert(self, expert_id) -> None: """Mark an expert as recently useful in every resident tier.""" - self.vram_lru.touch(expert_id) - self.ram_lru.touch(expert_id) + for key in self._expand_keys([expert_id]): + self.vram_lru.touch(key) + self.ram_lru.touch(key) def reprioritize_resident_experts(self, scores: Dict[int, float]) -> None: """Use recent router scores to make low-probability experts evict first.""" @@ -559,30 +705,25 @@ def sync_h2d(self): with self._events_lock: self._pending_events.clear() - def _evict_from_vram(self, expert_id: int): - """Write VRAM expert back to RAM buffer before eviction.""" - for li, moe in enumerate(self.moe_layers): - if expert_id < len(moe.experts): - state = {} - for k, v in moe.experts[expert_id].state_dict().items(): - cpu_v = v.detach().to(dtype=torch.float16, device="cpu").contiguous() - try: - state[k] = cpu_v.pin_memory() - except RuntimeError: - # MPS/CPU-only torch builds may reject pinned storage. - # Keeping a normal CPU tensor is still correct; it only - # loses the CUDA non_blocking copy optimization. - state[k] = cpu_v - with self._ram_lock: - if expert_id not in self._ram_buffers: - self._ram_buffers[expert_id] = {} - self._ram_buffers[expert_id][li] = state - self._replace_expert_module(li, expert_id, self._build_placeholder_expert()) - with self._ram_lock: - evicted = self.ram_lru.put(expert_id) - if evicted is not None and evicted != expert_id: - self._ram_buffers.pop(evicted, None) + def _evict_from_vram(self, expert_id): + # Inference weights are immutable and already persisted. Copying them + # back on eviction wastes bandwidth and can displace the demanded RAM + # page; on CPU it also creates a duplicate of shared storage. + self.wait_for_experts([expert_id]) + if self.layer_aware: + li, eid = self._page_parts(expert_id) + candidates = [(li, eid)] + else: + candidates = [(li, int(expert_id)) for li in range(self.num_layers)] + for li, eid in candidates: + expert = self.moe_layers[li].experts[eid] + if self._h2d_stream is not None: + current = torch.cuda.current_stream(self.device) + for param in expert.parameters(): + param.record_stream(current) + self._replace_expert_module(li, eid, self._build_placeholder_expert()) self.vram_lru.remove(expert_id) + self._io_stats["expert_evictions"] += 1 # ── Mask generation ─────────────────────────────────────────── @@ -590,32 +731,66 @@ def vram_availability_mask(self) -> torch.Tensor: """Returns [num_experts] bool tensor: True if expert is in VRAM.""" mask = torch.zeros(self.num_experts, dtype=torch.bool) for eid in self.vram_lru.keys(): + if isinstance(eid, tuple): + eid = eid[2] if eid < self.num_experts: mask[eid] = True return mask + def vram_availability_masks(self) -> List[torch.Tensor]: + """Return real per-physical-layer masks for the v2 cache.""" + if not self.layer_aware: + mask = self.vram_availability_mask() + return [mask.clone() for _ in range(self.num_layers)] + masks = [torch.zeros(self.num_experts, dtype=torch.bool) for _ in range(self.num_layers)] + for key in self.vram_lru.keys(): + li, eid = self._page_parts(key) + if 0 <= li < self.num_layers and 0 <= eid < self.num_experts: + masks[li][eid] = True + return masks + def stats(self) -> dict: available_ram_gb = self._physical_ram_bytes() / (1024 ** 3) - pinned_used_gb = ( - len(self.ram_lru) * self._expert_size_bytes() * self.num_layers / (1024 ** 3) - ) - dynamic_cap = self._safe_ram_capacity() + with self._ram_lock: + tensors = [v for layers in self._ram_buffers.values() + for state in layers.values() for v in state.values()] + live = [p for moe in self.moe_layers for expert in moe.experts for p in expert.parameters()] + ram_bytes = sum(v.numel() * v.element_size() for v in tensors) + live_bytes = sum(v.numel() * v.element_size() for v in live) + pinned_bytes = sum(v.numel() * v.element_size() for v in tensors + if self._h2d_stream is not None and v.is_pinned()) + # CPU modules share RAM pages. Count storage once across the two tiers. + unique = {} + for tensor in tensors + live: + storage = tensor.untyped_storage() + unique[(str(tensor.device), storage.data_ptr())] = storage.nbytes() + rss = _process_rss_bytes() + self._peak_rss_bytes = max(self._peak_rss_bytes, rss) return { + "device": str(self.device), + "rss_mb": round(rss / 2 ** 20, 1), + "peak_rss_mb": round(self._peak_rss_bytes / 2 ** 20, 1), + "max_resident_pages": self.max_resident_pages, + "ram_budget_pages": self.ram_budget_pages, + "protected_pages": len(self.vram_lru.protected), "vram_experts": len(self.vram_lru), "vram_capacity": self.vram_capacity, "ram_experts": len(self.ram_lru), - "ram_capacity_dynamic": dynamic_cap, + "ram_capacity_dynamic": self._safe_ram_capacity(), "expert_size_kb": self._expert_size_bytes() // 1024, - "pinned_ram_used_gb": round(pinned_used_gb, 3), + "cache_page_bytes": self._expert_size_bytes() * (1 if self.layer_aware else self.num_layers), + "resident_expert_bytes": live_bytes, + "ram_buffer_bytes": ram_bytes, + "unique_expert_bytes": sum(unique.values()), + "pinned_ram_used_gb": round(pinned_bytes / (1024 ** 3), 6), "available_ram_gb": round(available_ram_gb, 1), - "pinned_ram_fraction": round(pinned_used_gb / max(available_ram_gb, 1), 4), + "pinned_ram_fraction": round(pinned_bytes / max(available_ram_gb * 1024 ** 3, 1), 4), "h2d_stream": "dedicated" if self._h2d_stream is not None else "default", "storage_format": self.storage_format, "cluster_aware": self._cluster_storage is not None, - "num_clusters": ( - len(self._cluster_storage.manifest.clusters) - if self._cluster_storage is not None else 0 - ), + "layer_aware": self.layer_aware, + "num_clusters": len(self._cluster_storage.manifest.clusters) if self._cluster_storage else 0, + **self._io_stats, } diff --git a/chronos/io/storage.py b/chronos/io/storage.py index c5baf71..166b99e 100644 --- a/chronos/io/storage.py +++ b/chronos/io/storage.py @@ -36,7 +36,7 @@ MANIFEST_FILENAME = "cluster_manifest.json" -MANIFEST_VERSION = 1 +MANIFEST_VERSION = 2 KEY_SEP = "__" # layer/expert prefix vs param-name separator @@ -69,6 +69,9 @@ class ClusterManifest: clusters: Dict[int, Tuple[str, List[int]]] # expert_id -> cluster_id expert_to_cluster: Dict[int, int] + # v2 records the recurrent group owning each physical layer. It is + # optional when reading v1 manifests and defaults to identity groups. + layer_to_group: Dict[int, int] | None = None @classmethod def from_dict(cls, data: Dict) -> "ClusterManifest": @@ -79,6 +82,11 @@ def from_dict(cls, data: Dict) -> "ClusterManifest": [int(e) for e in entry["experts"]], ) e2c = {int(k): int(v) for k, v in data["expert_to_cluster"].items()} + raw_l2g = data.get("layer_to_group") + layer_to_group = ( + {int(k): int(v) for k, v in raw_l2g.items()} + if isinstance(raw_l2g, dict) else None + ) return cls( version=int(data.get("version", MANIFEST_VERSION)), num_experts=int(data["num_experts"]), @@ -86,6 +94,7 @@ def from_dict(cls, data: Dict) -> "ClusterManifest": storage_format=data.get("storage_format", "safetensors"), clusters=clusters, expert_to_cluster=e2c, + layer_to_group=layer_to_group, ) def to_dict(self) -> Dict: @@ -99,8 +108,23 @@ def to_dict(self) -> Dict: for cid, (fname, eids) in sorted(self.clusters.items()) ], "expert_to_cluster": {str(k): v for k, v in self.expert_to_cluster.items()}, + "layer_aware": bool(self.layer_to_group), + "layer_to_group": {str(k): int(v) for k, v in (self.layer_to_group or {}).items()}, } + def validate(self, *, num_layers: int | None = None, + num_experts: int | None = None) -> list[str]: + """Return actionable manifest mismatches instead of failing in hot IO.""" + errors = [] + if num_layers is not None and int(self.num_layers) != int(num_layers): + errors.append(f"num_layers manifest={self.num_layers} runtime={num_layers}") + if num_experts is not None and int(self.num_experts) != int(num_experts): + errors.append(f"num_experts manifest={self.num_experts} runtime={num_experts}") + expected = set(range(int(self.num_experts))) + if set(self.expert_to_cluster) != expected: + errors.append("expert_to_cluster does not cover every expert") + return errors + def cluster_members(self, expert_id: int) -> List[int]: cid = self.expert_to_cluster[expert_id] return list(self.clusters[cid][1]) @@ -120,7 +144,7 @@ class ClusterStorage: Read path (runtime, hot loop): storage = ClusterStorage(output_dir) - per_layer = storage.load_cluster(cluster_id, dtype=torch.float16) + per_layer = storage.load_cluster(cluster_id) # pages keep their stored dtype # per_layer: {expert_id: {layer_id: {param_name: Tensor}}} """ @@ -146,6 +170,7 @@ def load_cluster( self, cluster_id: int, dtype: Optional[torch.dtype] = None, + pages: Optional[set] = None, ) -> Dict[int, Dict[int, Dict[str, torch.Tensor]]]: """ mmap a whole cluster file and return nested dict: @@ -164,6 +189,10 @@ def load_cluster( if parsed is None: continue layer_id, expert_id, param_name = parsed + # Select before get_tensor: a one-page miss must not materialize + # every layer/expert in a large cluster. + if pages is not None and (layer_id, expert_id) not in pages: + continue t = f.get_tensor(key) if dtype is not None and t.dtype != dtype: t = t.to(dtype) @@ -185,12 +214,17 @@ def write_clusters( output_dir: str, num_layers: int, num_experts: int, - dtype: torch.dtype = torch.float16, + dtype: Optional[torch.dtype] = None, + layer_to_group: Optional[Dict[int, int]] = None, ) -> str: """ Serialize `moe_layers` (list of ChronosMOEFeedForward) into one safetensors file per cluster and write the manifest. + ``dtype=None`` (default) stores each expert in its own dtype, i.e. the + 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. """ if not _HAS_SAFETENSORS: @@ -212,13 +246,21 @@ def write_clusters( continue expert = moe_layers[li].experts[eid] for pname, pval in expert.state_dict().items(): - tensors[_key(li, int(eid), pname)] = pval.detach().to(dtype=dtype).cpu().contiguous() + 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() file_name = f"cluster_{cid}.ctsr" file_path = os.path.join(output_dir, file_name) _st_save_file(tensors, file_path) manifest_clusters[cid] = (file_name, [int(e) for e in cluster]) + layer_to_group = ( + {int(li): int(gi) for li, gi in layer_to_group.items()} + if layer_to_group is not None + else {int(li): int(li) for li in range(int(num_layers))} + ) manifest = ClusterManifest( version=MANIFEST_VERSION, num_experts=num_experts, @@ -226,6 +268,7 @@ def write_clusters( storage_format="safetensors", clusters=manifest_clusters, expert_to_cluster=expert_to_cluster, + layer_to_group=layer_to_group, ) manifest_path = os.path.join(output_dir, MANIFEST_FILENAME) with open(manifest_path, "w") as f: diff --git a/chronos/migrate.py b/chronos/migrate.py new file mode 100644 index 0000000..83894ae --- /dev/null +++ b/chronos/migrate.py @@ -0,0 +1,51 @@ +"""Explicit, auditable Chronos checkpoint migrations.""" +from __future__ import annotations + +import argparse +import os + +from chronos.model.checkpoint import ( + chronos_config_from_checkpoint, + load_checkpoint_state_dict, + save_state_dict_with_config, +) +from chronos.model.model_chronos import ChronosForCausalLM + + +def migrate_v1_to_v2(source: str, output: str) -> str: + cfg, _ = chronos_config_from_checkpoint(source, require_unsniffable=True) + cfg.architecture_version = 2 + cfg.recurrent_enabled = True + cfg.recurrent_depth = 1 + cfg.min_recurrent_depth = 1 + cfg.max_inference_depth = 1 + cfg.layer_aware_cache = False + cfg.layer_aware_lookahead = False + cfg.extra_metadata = { + **dict(getattr(cfg, "extra_metadata", {}) or {}), + "migration": {"from": "chronos-v1", "to": "chronos-v2", "source": os.path.abspath(source)}, + } + model = ChronosForCausalLM(cfg) + state = load_checkpoint_state_dict(source, map_location="cpu") + result = model.load_state_dict(state, strict=False) + unexpected = [k for k in result.unexpected_keys if not k.startswith("model.freqs_")] + allowed_missing = ( + "model.recurrent_controller.", "model.multimodal_fusion.", + ) + bad_missing = [k for k in result.missing_keys if not k.startswith(allowed_missing)] + if unexpected or bad_missing: + raise ValueError(f"migration topology mismatch: missing={bad_missing}, unexpected={unexpected}") + save_state_dict_with_config(model, output, cfg, stage="migrated-v2") + return output + + +def main(argv=None): + parser = argparse.ArgumentParser(description="Migrate an explicit Chronos v1 checkpoint to v2") + parser.add_argument("checkpoint") + parser.add_argument("--output", required=True) + args = parser.parse_args(argv) + print(migrate_v1_to_v2(args.checkpoint, args.output)) + + +if __name__ == "__main__": + main() diff --git a/chronos/mlx/attention.py b/chronos/mlx/attention.py index fa0ef9d..02f1100 100644 --- a/chronos/mlx/attention.py +++ b/chronos/mlx/attention.py @@ -2,16 +2,23 @@ chronos/mlx/attention.py — MLX-native MLA + Sliding Window Attention. Design notes vs. the PyTorch version: -- No past_key_values list; instead we carry (c_kv_cache, k_rope_cache) as - plain mx.array and concatenate explicitly each step. -- RoPE: computed with mx operations; no separate apply_rotary_pos_emb import. +- MLA carries only the latent cache. Sliding layers carry K, V and an + absolute position counter that survives cache eviction. +- RoPE: FP32 tables applied by one fused kernel (single rounding). +- Attention: mx.fast.scaled_dot_product_attention with native GQA and a + boolean mask built here; softmax is FP32 inside the kernel, so the + low-precision chain has no score/probability casts. - SlidingWindow: slice oldest tokens with array indexing, not assignment. - mx.compile()-friendly: no Python branches on runtime tensor shapes. """ import math +from functools import partial + import mlx.core as mx import mlx.nn as nn +from chronos.mlx.kv_quant import kv_read, kv_store + def _rope_freqs( dim: int, @@ -46,27 +53,93 @@ def _rope_freqs( return cos, sin +def _rotate_half(x: mx.array) -> mx.array: + half = x.shape[-1] // 2 + return mx.concatenate([-x[..., half:], x[..., :half]], axis=-1) + + +@partial(mx.compile, shapeless=True) # decode grows T every step; no retrace +def _rope_fma(x: mx.array, rotated: mx.array, c: mx.array, s: mx.array) -> mx.array: + # One fused kernel: FP32 tables, FP32 multiply-add in registers, a single + # rounding to x.dtype. No FP32 activation tensor is materialized. + return (x * c + rotated * s).astype(x.dtype) + + def _apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array: - """ - Apply RoPE to x [B, T, nH, d]. - cos/sin: full duplicated MiniMind tables, sliced to x.shape[-1]. + """Apply RoPE to x [B, T, nH, d] with tables cos/sin [T, >=d] (FP32). + + The tables stay FP32 and are sliced to x.shape[-1]; ``_rope_fma`` fuses + the rotation so low-precision runs round once instead of three times + (measured BF16 relative error 1.96e-3 vs 2.48e-3 with BF16 tables). """ _B, T, _nH, d = x.shape c = cos[:T, :d][:, None, :] s = sin[:T, :d][:, None, :] - half = d // 2 - rotated_half = mx.concatenate([-x[..., half:], x[..., :half]], axis=-1) - return (x * c + rotated_half * s).astype(x.dtype) + return _rope_fma(x, _rotate_half(x), c, s) + + +def _sdpa(q: mx.array, k: mx.array, v: mx.array, scale: float, allowed: mx.array | None) -> mx.array: + """Fused attention on [B, nH, S, d] queries and [B, nKV, T, d] keys/values. + + The kernel keeps scores and softmax in FP32 internally, handles GQA + natively (no repeated K/V) and returns the compute dtype. ``allowed`` is + a boolean [S, T] mask; an array mask (not the "causal" string) is used + because the vector kernels cannot produce the logsumexp the VJP needs. + """ + return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=allowed) + + +def _causal_allowed(S: int, T: int, window: int | None = None, past: int | None = None) -> mx.array | None: + """Boolean [S, T] mask for queries at absolute rows ``past..past+S-1``. + + Returns None for single-token decode without a window (all keys visible). + """ + past = T - S if past is None else past + q_pos = mx.arange(S)[:, None] + past + k_pos = mx.arange(T)[None, :] + if window is None: + if S == 1: + return None + return k_pos <= q_pos + return (k_pos <= q_pos) & (k_pos > q_pos - window) + + +def _merge_mask(allowed: mx.array | None, mask: mx.array | None, S: int, T: int) -> mx.array | None: + """Combine the layer's own boolean mask with an optional caller mask. + + Callers may pass a boolean mask or a legacy additive one (0 / -inf); + both are aligned to the right edge of the key axis. + """ + if mask is None: + return allowed + m = mask if mask.dtype == mx.bool_ else (mask > -1e4) + if m.shape[-1] < T: + m = mx.pad(m, [(0, 0)] * (m.ndim - 1) + [(T - m.shape[-1], 0)], constant_values=True) + m = m[..., -S:, -T:] if m.ndim >= 2 else m + return m if allowed is None else (allowed & m) class _RopeCacheMixin: - def _init_rope_cache(self, config, dim: int) -> None: + """FP32 RoPE tables, built once and never converted per call. + + ``rope_layout="legacy_sliced"`` reproduces the v1-v3 checkpoints: MLA + slices the first ``rope_dim`` columns of the head_dim table. That is not + a norm-preserving rotation when rope_dim < head_dim, so new configs use + ``"per_head"`` (tables sized to the rotated dimension). SWA rotates the + full head_dim either way, so both layouts are identical there. + """ + + def _init_rope_cache(self, config, dim: int, rotated_dim: int | None = None) -> None: rope_len = int(getattr(config, "max_position_embeddings", 4096)) - object.__setattr__(self, "_rope_dim", int(dim)) + layout = str(getattr(config, "rope_layout", "legacy_sliced") or "legacy_sliced") + if layout == "per_head" and rotated_dim and int(rotated_dim) % 2: + raise ValueError(f"rope_layout='per_head' needs an even rotated dimension, got {rotated_dim}") + table_dim = int(rotated_dim) if (layout == "per_head" and rotated_dim) else int(dim) + object.__setattr__(self, "_rope_dim", table_dim) object.__setattr__(self, "_rope_base", float(getattr(config, "rope_theta", 1_000_000.0))) object.__setattr__(self, "_rope_scaling", getattr(config, "rope_scaling", None)) cos, sin = _rope_freqs( - int(dim), + table_dim, max_len=rope_len, base=self._rope_base, rope_scaling=self._rope_scaling, @@ -74,6 +147,19 @@ def _init_rope_cache(self, config, dim: int) -> None: object.__setattr__(self, "_cos", cos) object.__setattr__(self, "_sin", sin) + def _attend(self, q, k, v, scale: float, allowed): + if self.training and getattr(self, "dropout_p", 0.0) > 0.0: + # Attention-probability dropout needs the explicit probabilities; + # the softmax stays FP32 and is rounded once for the PV product. + if self.n_heads != self.n_kv_heads: + rep = self.n_heads // self.n_kv_heads + k, v = mx.repeat(k, rep, axis=1), mx.repeat(v, rep, axis=1) + scores = (q @ k.swapaxes(-1, -2)).astype(mx.float32) * scale + if allowed is not None: + scores = mx.where(allowed, scores, -float("inf")) + return self.attn_dropout(mx.softmax(scores, axis=-1)).astype(q.dtype) @ v + return _sdpa(q, k, v, scale, allowed) + def _ensure_rope_cache(self, needed_len: int) -> None: current_len = int(self._cos.shape[0]) if int(needed_len) <= current_len: @@ -95,8 +181,7 @@ class MLAAttentionMLX(_RopeCacheMixin, nn.Module): """ Multi-head Latent Attention for MLX. - KV cache stores (c_kv [B, T, kv_latent_dim], k_rope [B, T, nH, rope_dim]). - Both are concatenated each decode step — no copy in/out of VRAM. + KV cache stores only (c_kv [B, T, kv_latent_dim],), matching PyTorch. """ def __init__(self, config): @@ -125,74 +210,112 @@ def __init__(self, config): self.k_rope_proj = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.rope_dim, bias=False) self.v_proj = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.n_heads * self.head_dim, H, bias=False) - - self._init_rope_cache(config, self.head_dim) + self.attn_dropout = nn.Dropout(config.dropout) + self.resid_dropout = nn.Dropout(config.dropout) + self.dropout_p = float(config.dropout) + + self._init_rope_cache(config, self.head_dim, rotated_dim=self.rope_dim) + # Inference cache/decode policy (no parameters; see ChronosConfig). + object.__setattr__(self, "mla_decode", str(getattr(config, "mla_decode", "expanded") or "expanded").lower()) + object.__setattr__(self, "cache_rope", bool(getattr(config, "mla_cache_rope", False))) + object.__setattr__(self, "kv_int8", str(getattr(config, "kv_cache_dtype", "auto") or "auto").lower() == "int8") + object.__setattr__(self, "kv_group", int(getattr(config, "kv_quant_group_size", 32) or 32)) + + def _use_absorbed(self, S: int, T: int) -> bool: + if self.mla_decode == "expanded" or (self.training and self.dropout_p > 0.0): + return False + if self.mla_decode == "absorbed": + return True + return (not self.training) and S <= 8 and T > S + + def _rotated_k_rope(self, c_kv: mx.array, start: int) -> mx.array: + B, T, _ = c_kv.shape + raw = self.k_rope_proj(c_kv).reshape(B, T, self.n_kv_heads, self.rope_dim) + return _apply_rope(raw, self._cos[start:start + T], self._sin[start:start + T]) def __call__( self, x: mx.array, mask: mx.array = None, - cache=None, # (c_kv [B, T_past, kv_latent_dim], k_rope [B, T_past, nH, rope_dim]) + cache=None, # (c_kv,) or (c_kv, k_rope); arrays or QuantizedKVMLX + kv_override=None, + return_shared_kv: bool = False, ): + """``kv_override``/``return_shared_kv``: v3 shared_first contract + (see the PyTorch MLAAttention docstring). The shared view is always + ``(c_kv_full, k_rope_full)`` in the compute dtype.""" B, S, H = x.shape + dtype = x.dtype - # ── Queries ─────────────────────────────────────────────── - q_nope = self.q_nope_norm( - self.q_nope_proj(x).reshape(B, S, self.n_heads, self.nope_dim) - ) - q_rope = self.q_rope_norm( - self.q_rope_proj(x).reshape(B, S, self.n_heads, self.rope_dim) - ) - - # ── KV latent (what gets cached) ────────────────────────── - c_kv = self.kv_down_norm(self.kv_down_proj(x)) # [B, S, kv_latent_dim] + q_nope = self.q_nope_norm(self.q_nope_proj(x).reshape(B, S, self.n_heads, self.nope_dim)) + q_rope = self.q_rope_norm(self.q_rope_proj(x).reshape(B, S, self.n_heads, self.rope_dim)) - # Append to cache - if cache is not None: - c_kv_past, k_rope_past = cache - c_kv_full = mx.concatenate([c_kv_past, c_kv], axis=1) + if kv_override is not None: + c_kv_full = kv_override[0] + k_rope = kv_override[1] if len(kv_override) > 1 else None + T_full = c_kv_full.shape[1] + self._ensure_rope_cache(T_full) + new_cache = None else: - c_kv_full = c_kv - k_rope_past = None - - T_full = c_kv_full.shape[1] - self._ensure_rope_cache(T_full) - - # ── Keys / Values (recomputed from cache) ───────────────── - k_nope = self.k_nope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.nope_dim) - k_rope = self.k_rope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.rope_dim) - v = self.v_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.head_dim) + c_kv = self.kv_down_norm(self.kv_down_proj(x)) # [B, S, kv_latent_dim] + past_c = cache[0] if cache is not None else None + c_kv_full = c_kv if past_c is None else mx.concatenate([kv_read(past_c, dtype), c_kv], axis=1) + T_full = c_kv_full.shape[1] + self._ensure_rope_cache(T_full) + # MLX caches are lazy: unused (training) cache entries are never + # evaluated, so building them costs nothing there. + stored_c = kv_store(c_kv, past_c, self.kv_group) if self.kv_int8 else c_kv_full + k_rope = None + new_cache = (stored_c,) + if self.cache_rope: + past_len = T_full - S + k_rope_new = self._rotated_k_rope(c_kv, past_len) + past_r = cache[1] if cache is not None and len(cache) > 1 else None + if past_c is not None and past_r is None: + past_r = self._rotated_k_rope(kv_read(past_c, dtype), 0) + k_rope = k_rope_new if past_r is None else mx.concatenate([kv_read(past_r, dtype), k_rope_new], axis=1) + stored_r = kv_store(k_rope_new, past_r, self.kv_group) if self.kv_int8 else k_rope + new_cache = (stored_c, stored_r) + if k_rope is None: + k_rope = self._rotated_k_rope(c_kv_full, 0) # [B, T, nKV, rope_dim] - # ── RoPE ────────────────────────────────────────────────── past_len = T_full - S - q_rope = _apply_rope(q_rope, - self._cos[past_len:past_len + S], - self._sin[past_len:past_len + S]) - k_rope = _apply_rope(k_rope, self._cos[:T_full], self._sin[:T_full]) - - # Concat nope + rope - q = mx.concatenate([q_nope, q_rope], axis=-1) # [B, S, nH, head_dim] - k = mx.concatenate([k_nope, k_rope], axis=-1) # [B, T, nKV, head_dim] - - # GQA repeat - if self.n_heads != self.n_kv_heads: - rep = self.n_heads // self.n_kv_heads - k = mx.repeat(k, rep, axis=2) - v = mx.repeat(v, rep, axis=2) - - # ── Attention ───────────────────────────────────────────── + q_rope = _apply_rope(q_rope, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S]) scale = 1.0 / math.sqrt(self.head_dim) - q = q.transpose(0, 2, 1, 3) # [B, nH, S, d] - k = k.transpose(0, 2, 1, 3) # [B, nH, T, d] - v = v.transpose(0, 2, 1, 3) - attn = (q @ k.transpose(0, 1, 3, 2)) * scale # [B, nH, S, T] - if mask is not None: - attn = attn + mask[None, None, :, :] - attn = mx.softmax(attn.astype(mx.float32), axis=-1).astype(q.dtype) - out = (attn @ v).transpose(0, 2, 1, 3).reshape(B, S, -1) - - new_cache = (c_kv_full, k_rope) - return self.o_proj(out), new_cache + allowed = _merge_mask(_causal_allowed(S, T_full), mask, S, T_full) + if self._use_absorbed(S, T_full): + out = self._absorbed(q_nope, q_rope, c_kv_full, k_rope, scale, allowed) + else: + k_nope = self.k_nope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.nope_dim) + v = self.v_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.head_dim) + q = mx.concatenate([q_nope, q_rope], axis=-1).transpose(0, 2, 1, 3) # [B, nH, S, d] + k = mx.concatenate([k_nope, k_rope], axis=-1).transpose(0, 2, 1, 3) # [B, nKV, T, d] + out = self._attend(q, k, v.transpose(0, 2, 1, 3), scale, allowed) + out = out.transpose(0, 2, 1, 3) + out = self.resid_dropout(self.o_proj(out.reshape(B, S, -1))) + if return_shared_kv: + return out, new_cache, (c_kv_full, k_rope) + return out, new_cache + + def _absorbed(self, q_nope, q_rope, c_kv, k_rope, scale, allowed): + """Scores against the latent cache; W_uk folded into q, W_uv into the output. + + Fused SDPA with key width latent+rope and value width latent (MQA per + KV group, FP32 softmax inside the kernel). Returns [B, S, nH, d]. + """ + B, S, nH, _ = q_nope.shape + T, L, G, R = c_kv.shape[1], self.kv_latent_dim, self.n_kv_heads, self.n_heads // self.n_kv_heads + w_k = self.k_nope_proj.weight.reshape(G, self.nope_dim, L) + w_v = self.v_proj.weight.reshape(G, self.head_dim, L) + qn = q_nope.reshape(B, S, G, R, self.nope_dim).transpose(0, 2, 3, 1, 4) # [B, G, R, S, n] + q_lat = (qn @ w_k[None, :, None]).reshape(B, nH, S, L) # [B, nH, S, L] + q = mx.concatenate([q_lat, q_rope.transpose(0, 2, 1, 3)], axis=-1) # [B, nH, S, L+r] + c = mx.broadcast_to(c_kv[:, None], (B, G, T, L)) + k = mx.concatenate([c, k_rope.transpose(0, 2, 1, 3)], axis=-1) # [B, G, T, L+r] + lat = _sdpa(q, k, c, scale, allowed) # [B, nH, S, L] + lat = lat.reshape(B, G, R, S, L) + out = lat @ w_v.transpose(0, 2, 1)[None, :, None] # [B, G, R, S, d] + return out.transpose(0, 3, 1, 2, 4).reshape(B, S, nH, self.head_dim) class SlidingWindowAttentionMLX(_RopeCacheMixin, nn.Module): @@ -215,61 +338,73 @@ def __init__(self, config): self.o_proj = nn.Linear(self.n_heads * self.head_dim, H, bias=False) self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.attn_dropout = nn.Dropout(config.dropout) + self.resid_dropout = nn.Dropout(config.dropout) + self.dropout_p = float(config.dropout) self._init_rope_cache(config, self.head_dim) + object.__setattr__(self, "kv_int8", str(getattr(config, "kv_cache_dtype", "auto") or "auto").lower() == "int8") + object.__setattr__(self, "kv_group", int(getattr(config, "kv_quant_group_size", 32) or 32)) + # Extra retained entries so self-speculative decoding can roll back up + # to ``cache_slack`` rejected tokens without losing in-window keys. + object.__setattr__(self, "cache_slack", 0) def __call__( self, x: mx.array, mask: mx.array = None, - cache=None, # (k [B, T_past, nKV, d], v [B, T_past, nKV, d]) + cache=None, # (k [B, T_past, nKV, d], v [B, T_past, nKV, d], absolute_position) + kv_override=None, + return_shared_kv: bool = False, ): B, S, H = x.shape q = self.q_norm(self.q_proj(x).reshape(B, S, self.n_heads, self.head_dim)) - k_new = self.k_norm(self.k_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim)) - v_new = self.v_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim) - - past_len = cache[0].shape[1] if cache is not None else 0 - self._ensure_rope_cache(past_len + S) - q = _apply_rope(q, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S]) - k_new = _apply_rope(k_new, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S]) - - if cache is not None: - k_past, v_past = cache - k = mx.concatenate([k_past, k_new], axis=1) - v = mx.concatenate([v_past, v_new], axis=1) + if kv_override is not None: + k, v, end = kv_override + past_len = int(end) - S + cached_len = k.shape[1] - S + self._ensure_rope_cache(past_len + S) + q = _apply_rope(q, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S]) else: - k = k_new - v = v_new + k_new = self.k_norm(self.k_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim)) + v_new = self.v_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim) + cached_len = cache[0].shape[1] if cache is not None else 0 + past_len = (cache[2] if len(cache) > 2 else cached_len) if cache is not None else 0 + self._ensure_rope_cache(past_len + S) + q = _apply_rope(q, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S]) + k_new = _apply_rope(k_new, self._cos[past_len:past_len + S], self._sin[past_len:past_len + S]) + if cache is not None: + k_past, v_past = cache[:2] + k = mx.concatenate([kv_read(k_past, k_new.dtype), k_new], axis=1) + v = mx.concatenate([kv_read(v_past, v_new.dtype), v_new], axis=1) + else: + k = k_new + v = v_new T_full = k.shape[1] - # Slide window — keep only last window_size tokens - if T_full > self.window_size: - k = k[:, -self.window_size:, :, :] - v = v[:, -self.window_size:, :, :] - - # GQA repeat for attention (k/v expand to n_heads) - k_attn = k - v_attn = v - if self.n_heads != self.n_kv_heads: - rep = self.n_heads // self.n_kv_heads - k_attn = mx.repeat(k, rep, axis=2) - v_attn = mx.repeat(v, rep, axis=2) + # All queries need their OWN causal window. Truncating keys before + # attention leaks future tokens into early queries during long prefill. + q_positions = mx.arange(S)[:, None] + cached_len + k_positions = mx.arange(T_full)[None, :] + allowed = (k_positions <= q_positions) & (k_positions > q_positions-self.window_size) scale = 1.0 / math.sqrt(self.head_dim) - # Transpose to [B, nH, T, d] for batched matmul - qt = q.transpose(0, 2, 1, 3) - kt = k_attn.transpose(0, 2, 1, 3) - vt = v_attn.transpose(0, 2, 1, 3) - attn = (qt @ kt.transpose(0, 1, 3, 2)) * scale - if mask is not None: - # mask is [S, T_k] (2D causal); broadcast over [B, nH, S, T_k] - attn = attn + mask[None, None, :, :] - attn = mx.softmax(attn.astype(mx.float32), axis=-1).astype(q.dtype) - out = (attn @ vt).transpose(0, 2, 1, 3).reshape(B, S, -1) + out = self._attend(q.transpose(0, 2, 1, 3), k.transpose(0, 2, 1, 3), v.transpose(0, 2, 1, 3), + scale, _merge_mask(allowed, mask, S, T_full)) + out = out.transpose(0, 2, 1, 3).reshape(B, S, -1) # Cache stores [B, T, nKV, d] (pre-repeat, pre-transpose) - new_cache = (k, v) - return self.o_proj(out), new_cache + keep = self.window_size + int(self.cache_slack) + if self.kv_int8 and kv_override is None: + # Quantize only the new tokens (each token once), then evict. + qk = kv_store(k_new, cache[0] if cache is not None else None, self.kv_group) + qv = kv_store(v_new, cache[1] if cache is not None else None, self.kv_group) + start = max(0, qk.shape[1] - keep) + new_cache = (qk[:, start:], qv[:, start:], past_len + S) + else: + new_cache = (k[:, -keep:], v[:, -keep:], past_len + S) + if return_shared_kv: + return self.resid_dropout(self.o_proj(out)), new_cache, (k, v, past_len + S) + return self.resid_dropout(self.o_proj(out)), new_cache diff --git a/chronos/mlx/expert_store.py b/chronos/mlx/expert_store.py index 1d27ac2..0afc4b8 100644 --- a/chronos/mlx/expert_store.py +++ b/chronos/mlx/expert_store.py @@ -24,12 +24,40 @@ try: from safetensors import safe_open as _safe_open - from safetensors.numpy import save_file as _np_save_file _HAS_SAFETENSORS = True except ImportError: # pragma: no cover _HAS_SAFETENSORS = False +# Expert pages cross the SSD boundary through MLX's own safetensors / npy IO: +# numpy has no bfloat16, so ``np.array(bf16)`` and safetensors' ``mlx`` +# framework (which round-trips through numpy) both fail on BF16 pages. The +# file format is unchanged (.ctsr is plain safetensors; torch-written BF16/FP16 +# shards load directly). + +def save_pages(path: str, tensors: Dict[str, mx.array]) -> None: + """Atomic safetensors write in the arrays' own dtype.""" + tmp = path + ".tmp" + mx.save_safetensors(tmp, tensors) # appends ".safetensors" + os.replace(tmp + ".safetensors", path) + + +def load_pages(path: str) -> Dict[str, mx.array]: + return mx.load(path, format="safetensors") + + +def _materialize(layer_states: Dict[int, Dict[str, mx.array]]) -> None: + """Evaluate SSD pages in the thread that loaded them. + + MLX 0.32 binds a lazy ``mx.load`` to the creating thread's stream; a page + loaded by the prefetch thread and first evaluated by the decode thread + raises "There is no Stream(cpu, N) in current thread". Evaluating here + also keeps the SSD read on the prefetch thread, off the decode path.""" + arrays = [a for params in layer_states.values() for a in params.values()] + if arrays: + mx.eval(*arrays) + + class MLXExpertStore: """ Expert weight store for MLX (unified memory — no H2D stream needed). @@ -69,13 +97,32 @@ def __init__(self, model, config, ssd_dir: str = "./expert_cache_mlx"): if warm_requested in (None, "", 0): warm_requested = min(self.num_experts, max(capacity, capacity * 2)) self._warm_capacity = max(capacity, min(self.num_experts, int(warm_requested))) + # Hard budget (unified memory: hot experts also live in RAM, so the + # RAM budget bounds both tiers). One page = one expert, all layers. + from chronos.io.expert_store import expert_budget_pages + h = int(getattr(config, "hidden_size", 0) or 0) + i = int(getattr(config, "moe_intermediate_size", 0) or 0) + itemsize = int(getattr(model.embed_tokens.weight.dtype, "size", 4) or 4) + self.page_bytes = max(1, 3 * h * i * itemsize * max(1, self.num_layers)) + self.max_resident_pages, self.ram_budget_pages = expert_budget_pages(config, self.page_bytes) + if self.max_resident_pages is not None: + self._capacity = min(self._capacity, self.max_resident_pages) + if self.ram_budget_pages is not None: + self._capacity = min(self._capacity, self.ram_budget_pages) + self._warm_capacity = min(self._warm_capacity, self.ram_budget_pages) + self._warm_capacity = max(self._capacity, self._warm_capacity) + self._peak_rss_bytes = 0 self._warm_lru: collections.OrderedDict = collections.OrderedDict() self._cluster_manifest: ClusterManifest | None = None self._loaded_clusters: set[int] = set() self._protected_experts: set[int] = set() self.storage_format = "npy" - self.attach_cluster_manifest(ssd_dir) + # 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 + # read, never per forward. + self.page_dtype = model.embed_tokens.weight.dtype self._stats = collections.Counter() + self.attach_cluster_manifest(ssd_dir) # ── SSD serialisation ───────────────────────────────────────── @@ -97,7 +144,7 @@ def offload_all_to_ssd(self, clusters: Optional[List[List[int]]] = None): arr = getattr(expert, proj).weight path = self._ssd_path(ei, li, f"{proj}.weight") if not os.path.exists(path): - np.save(path, np.array(arr)) + mx.save(path, arr) # dtype-preserving (BF16 safe) def _cluster_cache_complete(self) -> bool: return self.attach_cluster_manifest(self.ssd_dir) @@ -124,9 +171,9 @@ def _write_clustered(self, clusters: List[List[int]]): expert = moe.experts[int(eid)] for proj in ("gate_proj", "up_proj", "down_proj"): key = f"l{li}_e{int(eid)}{KEY_SEP}{proj}.weight" - tensors[key] = np.array(getattr(expert, proj).weight) + tensors[key] = getattr(expert, proj).weight file_name = f"cluster_{cid}.ctsr" - _np_save_file(tensors, os.path.join(self.ssd_dir, file_name)) + save_pages(os.path.join(self.ssd_dir, file_name), tensors) manifest_clusters[cid] = (file_name, [int(e) for e in cluster]) manifest = ClusterManifest( version=1, @@ -159,9 +206,10 @@ def _load_from_ssd(self, expert_id: int) -> Dict[int, Dict[str, mx.array]]: fname = f"{name}.{subname}" path = self._ssd_path(expert_id, li, fname) if os.path.exists(path): - params[fname] = mx.array(np.load(path)) + params[fname] = self._as_page(mx.load(path)) if params: layer_states[li] = params + _materialize(layer_states) return layer_states def attach_cluster_manifest(self, manifest_dir: Optional[str] = None) -> bool: @@ -219,6 +267,12 @@ def _cluster_manifest_tensors_complete(self, manifest: ClusterManifest, target: return False return True + def _as_page(self, arr: mx.array) -> mx.array: + if arr.dtype != self.page_dtype: + self._stats["page_dtype_conversions"] += 1 + return arr.astype(self.page_dtype) + return arr + def _parse_cluster_key(self, key: str): if KEY_SEP not in key: return None @@ -246,13 +300,14 @@ def _load_cluster_for_expert(self, expert_id: int) -> None: file_name = manifest.cluster_file(cid) path = os.path.join(self.ssd_dir, file_name) nested: Dict[int, Dict[int, Dict[str, mx.array]]] = {} - with _safe_open(path, framework="mlx", device="cpu") as f: - for key in f.keys(): - parsed = self._parse_cluster_key(key) - if parsed is None: - continue - li, eid, pname = parsed - nested.setdefault(eid, {}).setdefault(li, {})[pname] = f.get_tensor(key) + for key, value in load_pages(path).items(): + parsed = self._parse_cluster_key(key) + if parsed is None: + continue + li, eid, pname = parsed + nested.setdefault(eid, {}).setdefault(li, {})[pname] = self._as_page(value) + for layer_states in nested.values(): + _materialize(layer_states) ordered = sorted(nested.items(), key=lambda item: int(item[0]) == int(expert_id)) with self._lock: for eid, layer_states in ordered: @@ -265,7 +320,7 @@ def _load_cluster_for_expert(self, expert_id: int) -> None: def prefetch_to_ram(self, expert_ids: List[int]): """ - Load expert weights from SSD into mx.array (not yet eval'd). + Load expert weights from SSD into evaluated mx.arrays (RAM tier). Call from a background thread for async prefetch. """ for eid in expert_ids: @@ -318,6 +373,13 @@ def _evict_warm_locked(self, protected_expert_id: int | None = None): if candidate != protected_expert_id and candidate not in self._protected_experts: evict_id = candidate break + if evict_id is None: + # Every other page is protected: the hard capacity still wins + # (oldest first); only the page being inserted is kept. + for candidate in self._warm_lru.keys(): + if candidate != protected_expert_id: + evict_id = candidate + break if evict_id is None: return evicted_was_hot = evict_id in self._hot_lru @@ -524,9 +586,19 @@ def warm_up(self, expert_ids: Optional[List[int]] = None): self.promote_to_vram(int(eid)) def stats(self) -> dict: + from chronos.io.expert_store import _process_rss_bytes hot_ids = self.hot_expert_ids() + rss = _process_rss_bytes() + self._peak_rss_bytes = max(self._peak_rss_bytes, rss) return { "backend": "mlx", + "rss_mb": round(rss / 2 ** 20, 1), + "peak_rss_mb": round(self._peak_rss_bytes / 2 ** 20, 1), + "page_bytes": self.page_bytes, + "max_resident_pages": self.max_resident_pages, + "ram_budget_pages": self.ram_budget_pages, + "protected_pages": len(self._protected_experts), + "resident_experts": len(hot_ids), "hot_experts": len(hot_ids), "hot_capacity": self._capacity, "warm_experts": len(self._warm), diff --git a/chronos/mlx/inference.py b/chronos/mlx/inference.py index 07714e5..0c36a02 100644 --- a/chronos/mlx/inference.py +++ b/chronos/mlx/inference.py @@ -35,6 +35,9 @@ class ChronosMLXInferenceEngine: """ def __init__(self, model, config, ssd_dir: str = "./expert_cache_mlx"): + # Inference only: training mode would randomly mask experts + # (fallback_mask_prob) and use the training ACT schedule. + model.eval() self.model = model self.config = config self.store = MLXExpertStore(model, config, ssd_dir=ssd_dir) @@ -214,6 +217,32 @@ def _all_resident_fast_path_enabled(self) -> bool: return False return len(self.store.hot_expert_ids()) >= num_experts + def _current_virtual_ids(self) -> List[int]: + """Last-token top-k ids over every MoE dispatch of the last forward. + + Recurrent groups dispatch once per loop iteration (virtual layer) + and may pick different experts each time; residency follows the + union, so the next step's identical iterations stay hot. + """ + 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 = [] + 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 ())) + if not rows: + return [] + 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() + 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]))] + def _protect_prediction_window(self, expert_ids: List[int]) -> None: capacity = max( 1, @@ -221,7 +250,8 @@ def _protect_prediction_window(self, expert_ids: List[int]) -> None: int(getattr(self.store, "_capacity", 1) or 1), ) hot_ids = sorted(self.store.hot_expert_ids()) - protected = list(dict.fromkeys([int(eid) for eid in expert_ids] + hot_ids))[:capacity] + current = list(getattr(self, "_last_current", ()) or ()) + protected = list(dict.fromkeys([int(eid) for eid in expert_ids] + current + hot_ids))[:capacity] self.store.set_protected_experts(protected) def stop(self): @@ -242,6 +272,15 @@ def generate( temperature: float = 0.85, top_p: float = 0.9, scheduler=None, # Optional[PrefillScheduler] — lazy import + *, + speculative: bool = False, # recurrent-depth self-speculative decoding + draft_depth: int = 1, + gamma: int = 4, + gamma_max: int = 8, + adaptive_gamma: bool = True, + top_k: int = 0, + seed: Optional[int] = None, + eos_token_id=None, ) -> Iterator[int]: """ Yields token IDs one at a time. @@ -249,10 +288,22 @@ def generate( If scheduler is provided (PrefillScheduler), runs prefill-time batch loading before the generation loop (dual-layer mode). Otherwise falls back to per-token LookaheadRouter-driven prefetch. + + ``speculative=True`` drafts ``gamma`` tokens at ``draft_depth`` and + 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. """ + 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) + return t0 = time.monotonic() setup_mem = self._memory_snapshot() self._last_predicted = set() + self._last_current = [] with self._stats_lock: all_resident_fast_path = self._all_resident_fast_path_enabled() self._runtime_stats = self._new_runtime_stats(all_resident_fast_path) @@ -277,6 +328,7 @@ def generate( prefill_masks = self._build_avail_masks(None) logits, lookahead_probs, cache = self.model(input_ids, available_expert_masks=prefill_masks) mx.eval(logits) + self._last_current = [] if all_resident_fast_path else self._current_virtual_ids() prefill_t1 = time.monotonic() prefill_mem = self._memory_snapshot() @@ -297,9 +349,11 @@ def generate( self._last_predicted = set(int(eid) for eid in future_ids) self._protect_prediction_window(future_ids) self._prefetch_and_promote_window(future_ids, promote_ids=immediate_ids, timeout_s=0.025) - yield int(next_token.item()) + eos = set([eos_token_id] if isinstance(eos_token_id, int) else (eos_token_id or [])) + first_id = int(next_token.item()) + yield first_id - for _ in range(max_new_tokens - 1): + for _ in range(0 if first_id in eos else max_new_tokens - 1): if all_resident_fast_path: avail_masks = None elif scheduler is not None: @@ -328,6 +382,8 @@ def generate( available_expert_masks=avail_masks, ) mx.eval(logits) + if not all_resident_fast_path: + self._last_current = self._current_virtual_ids() next_token = self._sample(logits[:, -1, :], temperature, top_p) mx.eval(next_token) @@ -335,6 +391,8 @@ def generate( activated_ids.append(tok_id) tokens += 1 yield tok_id + if tok_id in eos: + break # Update frequency heuristic after generation if scheduler is not None: @@ -354,6 +412,84 @@ def generate( **self.store.stats(), } + 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]: + """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.""" + from chronos.mlx.speculative import SpecStats, speculative_tokens + t0 = time.monotonic() + setup_mem = self._memory_snapshot() + self._last_predicted, self._last_current = set(), [] + with self._stats_lock: + fast = self._all_resident_fast_path_enabled() + self._runtime_stats = self._new_runtime_stats(fast) + if fast: + self._clear_runtime_hooks() + self.store.set_protected_experts([]) + else: + self._install_runtime_hooks() + warm_cap = int(getattr(self.store, "_warm_capacity", self.config.num_experts) or self.config.num_experts) + hot_cap = int(getattr(self.store, "_capacity", 1) or 1) + timing = {"first": None} + + def forward_fn(ids, cache, depth): + 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: + self._last_current = self._current_virtual_ids() + predicted = (self._predict_future_experts(lookahead, capacity=warm_cap) + if lookahead is not None else []) + window = list(dict.fromkeys(list(self._last_current) + predicted))[:max(warm_cap, 1)] + self._last_predicted = set(window) + self._protect_prediction_window(window) + self._prefetch_and_promote_window(window, promote_ids=window[:max(hot_cap, 1)]) + return logits, lookahead, new_cache + + stats = SpecStats() + from chronos.config_schema import resolve_recurrent_attention_mode + cfg = self.model.config + full_depth = int(getattr(cfg, "max_inference_depth", None) or cfg.recurrent_depth) + recurrent = bool(getattr(cfg, "recurrent_enabled", False)) and int(cfg.recurrent_depth) > 1 + eff_draft = max(1, min(int(draft_depth), full_depth)) + if not recurrent or resolve_recurrent_attention_mode(cfg) == "per_iteration": + eff_draft = full_depth # mirrors speculative_tokens' fallback + do_sample = float(temperature) > 0 + try: + for tok in speculative_tokens( + self.model, input_ids, stats, max_new_tokens=int(max_new_tokens), draft_depth=draft_depth, + gamma=gamma, adaptive_gamma=adaptive_gamma, gamma_max=gamma_max, do_sample=do_sample, + temperature=max(float(temperature), 1e-6), top_k=int(top_k or 0), + top_p=float(top_p) if do_sample else 1.0, eos_token_id=eos_token_id, seed=seed, + forward_fn=forward_fn): + if timing["first"] is None: + timing["first"] = time.monotonic() + yield int(tok) + finally: + end = time.monotonic() + first = timing["first"] or end + end_mem = self._memory_snapshot() + self.last_stats = { + "tokens_generated": stats.generated, + "elapsed_s": round(end - t0, 3), + "tokens_per_sec": round(stats.generated / max(end - t0, 1e-6), 2), + "prefill_time_s": round(first - t0, 3), + "decode_time_s": round(max(0.0, end - first), 3), + "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]}, + **self._memory_fields(setup_mem, "after_setup"), + **self._memory_fields(end_mem, "after_decode"), + **self._runtime_stat_fields(), + **self.store.stats(), + } + def _predict_future_experts( self, lookahead_probs: mx.array, @@ -361,28 +497,23 @@ def _predict_future_experts( max_steps: Optional[int] = None, ) -> List[int]: """Extract bounded top-k expert IDs for future steps.""" - # lookahead_probs: [B, S, K+1, E] — take last token, steps 1..K - future = lookahead_probs[0, -1, 1:, :] # [K, E] + # This store pages an expert ID across physical layers. Preserve + # each prediction head until admission, then rank the union of IDs. + future = lookahead_probs[:, -1, 1:, ...] if max_steps is not None: - future = future[:max(0, int(max_steps))] - if future.shape[0] <= 0: + future = future[:, :max(0, int(max_steps))] + if future.shape[1] <= 0: return [] - num_experts = int(getattr(self.config, "num_experts", future.shape[-1]) or future.shape[-1]) - top_k = max(1, min( - int(getattr(self.config, "num_experts_per_tok", 1) or 1), - num_experts, - )) - budget = num_experts if capacity is None else max(1, min(num_experts, int(capacity))) - scores = mx.zeros((num_experts,), dtype=mx.float32) - for offset in range(future.shape[0]): - step = future[offset].astype(mx.float32) - if top_k >= num_experts: - scores = scores + step / float(offset + 1) - continue - step_ids = mx.argpartition(-step, kth=top_k - 1, axis=-1)[:top_k] - step_vals = mx.take(step, step_ids) / float(offset + 1) - one_hot = (step_ids[:, None] == mx.arange(num_experts)).astype(mx.float32) - scores = scores + (one_hot * step_vals[:, None]).sum(axis=0) + e = int(self.config.num_experts) + 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) + ids = mx.argpartition(-step, kth=top_k-1, axis=-1)[..., :top_k] + vals = mx.take_along_axis(step, ids, axis=-1) / float(offset + 1) + 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])) @@ -439,6 +570,7 @@ def _runtime_stat_fields(self) -> dict: "on_demand_loads": int(stats.get("resident_misses", 0)), "on_demand_load_time_s": round(float(stats.get("on_demand_load_time_s", 0.0)), 4), "fallback_weight_rate": 0.0, + "virtual_dispatches": int(stats.get("virtual_dispatches", 0)), } @staticmethod diff --git a/chronos/mlx/kv_quant.py b/chronos/mlx/kv_quant.py new file mode 100644 index 0000000..d56c973 --- /dev/null +++ b/chronos/mlx/kv_quant.py @@ -0,0 +1,87 @@ +"""Inference-only int8 KV-cache storage (MLX), mirroring chronos.model.kv_quant. + +Symmetric int8 with one FP32 scale per token per ``group_size`` channels of +the last axis. Each token is quantized once when appended and dequantized +into the compute dtype when attention reads the cache. +""" +from __future__ import annotations + +import mlx.core as mx + + +class QuantizedKVMLX: + __slots__ = ("data", "scale", "group_size", "dtype") + + def __init__(self, data: mx.array, scale: mx.array, group_size: int, dtype): + self.data, self.scale, self.group_size, self.dtype = data, scale, int(group_size), dtype + + @classmethod + def quantize(cls, x: mx.array, group_size: int = 32) -> "QuantizedKVMLX": + d = int(x.shape[-1]) + g = int(group_size) if group_size and d % int(group_size) == 0 else d + groups = x.astype(mx.float32).reshape(*x.shape[:-1], d // g, g) + scale = mx.maximum(mx.max(mx.abs(groups), axis=-1, keepdims=True), 1e-12) / 127.0 + data = mx.clip(mx.round(groups / scale), -127, 127).astype(mx.int8).reshape(x.shape) + return cls(data, scale, g, x.dtype) + + def append(self, x: mx.array) -> "QuantizedKVMLX": + new = QuantizedKVMLX.quantize(x, self.group_size) + return QuantizedKVMLX(mx.concatenate([self.data, new.data], axis=1), + mx.concatenate([self.scale, new.scale], axis=1), self.group_size, self.dtype) + + def dequantize(self, dtype=None) -> mx.array: + d = int(self.data.shape[-1]) + groups = self.data.reshape(*self.data.shape[:-1], d // self.group_size, self.group_size) + return (groups.astype(mx.float32) * self.scale).reshape(self.data.shape).astype(dtype or self.dtype) + + @property + def shape(self): + return self.data.shape + + def __getitem__(self, index): + if not (isinstance(index, tuple) and len(index) == 2 and index[0] == slice(None) + and isinstance(index[1], slice) and index[1].step in (None, 1)): + raise TypeError("QuantizedKVMLX supports only [:, start:stop] token slicing") + return QuantizedKVMLX(self.data[index], self.scale[index], self.group_size, self.dtype) + + @property + def nbytes(self) -> int: + return int(self.data.nbytes) + int(self.scale.nbytes) + + def arrays(self): + return (self.data, self.scale) + + +def kv_store(x: mx.array, past, group_size: int) -> QuantizedKVMLX: + """Append ``x`` to an int8 cache (``past`` may be None, int8 or float).""" + if isinstance(past, QuantizedKVMLX): + return past.append(x) + if past is not None: + x = mx.concatenate([past.astype(x.dtype), x], axis=1) + return QuantizedKVMLX.quantize(x, group_size) + + +def kv_read(entry, dtype) -> mx.array: + return entry.dequantize(dtype) if isinstance(entry, QuantizedKVMLX) else entry + + +def entry_nbytes(entry) -> int: + total = 0 + for item in entry or (): + if isinstance(item, QuantizedKVMLX): + total += item.nbytes + elif isinstance(item, mx.array): + total += int(item.nbytes) + return total + + +def cache_arrays(cache): + """Flatten a cache list into arrays (for mx.eval / async_eval).""" + out = [] + for entry in cache or (): + for item in entry or (): + if isinstance(item, QuantizedKVMLX): + out.extend(item.arrays()) + elif isinstance(item, mx.array): + out.append(item) + return out diff --git a/chronos/mlx/model.py b/chronos/mlx/model.py index 2de1b73..62bfb85 100644 --- a/chronos/mlx/model.py +++ b/chronos/mlx/model.py @@ -17,6 +17,8 @@ from chronos.mlx.attention import MLAAttentionMLX, SlidingWindowAttentionMLX from chronos.mlx.moe import ChronosMLXMOE, FeedForwardMLX +from chronos.mlx.recurrent import RecurrentControllerMLX, auxiliary_state +from chronos.config_schema import resolve_recurrent_attention_mode, recurrent_cache_layout class RMSNormMLX(nn.Module): @@ -36,19 +38,31 @@ def __init__(self, config): super().__init__() K = getattr(config, 'lookahead_steps', 2) self.lookahead_steps = K + self.layer_aware = bool(getattr(config, "layer_aware_lookahead", False)) + self.num_groups = int(getattr(config, "recurrent_groups", 1) or 1) hidden = max(1, config.hidden_size // 4) self.proj = nn.Sequential( nn.Linear(config.hidden_size, hidden, bias=False), nn.SiLU(), - nn.Linear(hidden, config.num_experts * (K + 1), bias=False), + nn.Linear( + hidden, + config.num_experts * (K + 1) * (self.num_groups if self.layer_aware else 1), + bias=False, + ), ) self.num_experts = config.num_experts def __call__(self, x: mx.array) -> mx.array: B, S, H = x.shape out = self.proj(x).astype(mx.float32) # [B, S, E*(K+1)] - out = out.reshape(B, S, self.lookahead_steps + 1, self.num_experts) - return mx.softmax(out, axis=-1) # [B, S, K+1, E] + shape = (B, S, self.lookahead_steps + 1) + if self.layer_aware: + shape += (self.num_groups,) + expected = (self.lookahead_steps + 1)*self.num_experts*(self.num_groups if self.layer_aware else 1) + if out.shape[-1] != expected: + raise ValueError(f"MLX lookahead projection width {out.shape[-1]} != config width {expected}; load the checkpoint config with its weights") + return mx.softmax(out.reshape(*shape, self.num_experts), axis=-1) + class ChronosMLXBlock(nn.Module): @@ -65,6 +79,7 @@ def __init__(self, layer_id: int, config): self.self_attn = SlidingWindowAttentionMLX(config) else: self.self_attn = SlidingWindowAttentionMLX(config) + self.self_attn.window_size = 2**63 - 1 if getattr(config, 'use_moe', True): self.mlp = ChronosMLXMOE(config) @@ -77,14 +92,26 @@ def __call__( mask: mx.array = None, cache=None, available_expert_mask: mx.array = None, + latent_only: bool = False, + kv_override=None, + return_shared_kv: bool = False, ): - h, new_cache = self.self_attn(self.input_norm(x), mask=mask, cache=cache) - x = x + h + new_cache, shared = cache, None + if not latent_only: + out = self.self_attn(self.input_norm(x), mask=mask, cache=cache, + kv_override=kv_override, return_shared_kv=return_shared_kv) + if return_shared_kv: + h, new_cache, shared = out + else: + h, new_cache = out + x = x + h mlp_in = self.post_norm(x) if isinstance(self.mlp, ChronosMLXMOE): x = x + self.mlp(mlp_in, available_expert_mask=available_expert_mask) else: x = x + self.mlp(mlp_in) + if return_shared_kv: + return x, new_cache, shared return x, new_cache @@ -92,18 +119,23 @@ class ChronosMLXModel(nn.Module): """ Full Chronos model in MLX. - forward() returns (logits [B, S, V], lookahead_probs [B, S, K+1, E]) + Returns logits [B,S,V] and lookahead [B,S,K+1,G,E] (legacy: no G axis). Cache: list of per-layer attention caches (passed back in, None at start). """ def __init__(self, config): super().__init__() self.config = config + if config.modalities: + raise ValueError("MLX modality encoders are not implemented; use the PyTorch MPS/CUDA backend for multimodal checkpoints") + self.recurrent_controller = RecurrentControllerMLX(config) + object.__setattr__(self, "last_recurrent_state", None) + self.dropout = nn.Dropout(config.dropout) self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = [ChronosMLXBlock(i, config) for i in range(config.num_hidden_layers)] self.norm = RMSNormMLX(config.hidden_size, eps=config.rms_norm_eps) - self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.lm_head = None if config.tie_word_embeddings else nn.Linear(config.hidden_size, config.vocab_size, bias=False) # LookaheadRouter inserted after block 0 self.lookahead_router = LookaheadRouterMLX(config) @@ -113,15 +145,30 @@ def __call__( input_ids: mx.array, # [B, S] cache=None, # list of per-layer caches, or None available_expert_masks=None, # list of [E] float masks, or None + logits_dtype=None, # e.g. mx.float32 for losses + recurrent_depth=None, # draft depth override (shared_first/first_only) ): + """``logits_dtype=mx.float32`` makes the logits FP32 at birth: the + final hidden [B,S,H] and the head [V,H] are upcast (H/V the size of + the logits) instead of rounding logits to the compute dtype and + casting the large [B,S,V] tensor afterwards.""" + moes = [layer.mlp for layer in self.layers if isinstance(layer.mlp, ChronosMLXMOE)] + for moe in moes: + moe.reset_iteration_routing() + moe.open_pack_scope() + try: + return self._forward(input_ids, cache, available_expert_masks, logits_dtype, recurrent_depth) + finally: + for moe in moes: + moe.close_pack_scope() + + def _forward(self, input_ids, cache, available_expert_masks, logits_dtype, depth_override=None): B, S = input_ids.shape - x = self.embed_tokens(input_ids) # [B, S, H] + x = self.dropout(self.embed_tokens(input_ids)) # [B, S, H] - # Causal mask — only needed for prefill (S > 1); decode (S==1) needs no mask + # Each attention layer builds its own boolean causal / window mask + # (fused SDPA); no additive mask is materialized in the compute dtype. mask = None - if S > 1: - mask = nn.MultiHeadAttention.create_additive_causal_mask(S) - mask = mask.astype(x.dtype) new_caches = [] lookahead_probs = None @@ -129,34 +176,128 @@ def __call__( if cache is None: cache = [None] * len(self.layers) - for i, layer in enumerate(self.layers): - avail = (available_expert_masks[i] - if available_expert_masks is not None else None) - x, lyr_cache = layer(x, mask=mask, cache=cache[i], - available_expert_mask=avail) - new_caches.append(lyr_cache) - - # Lookahead routing after block 0 - if i == 0: - lookahead_probs = self.lookahead_router(x) + cfg = self.config + object.__setattr__(self, "last_recurrent_state", None) + def run(i, hidden, slot=None, **kw): + avail = available_expert_masks[i] if available_expert_masks is not None else None + c = cache[i if slot is None else slot] if "kv_override" not in kw else None + return self.layers[i](hidden, mask=mask, cache=c, available_expert_mask=avail, **kw) + + if cfg.recurrent_enabled and cfg.recurrent_depth > 1: + controller = self.recurrent_controller + depth = controller.execution_depth(cfg, self.training) + adaptive = controller.adaptive_active(cfg, self.training) + mode = resolve_recurrent_attention_mode(cfg) + if depth_override is not None: + # Shallow (draft) execution, Torch parity: only depth-free + # cache layouts can share a KV stream across depths. + if mode == "per_iteration": + raise ValueError("depth override needs a depth-independent cache (shared_first/first_only)") + depth = max(1, min(int(depth_override), int(cfg.recurrent_depth))) + layout = recurrent_cache_layout(cfg, depth) + if len(cache) != layout.num_slots: + if all(c is None for c in cache): + cache = [None] * layout.num_slots + else: + raise ValueError(f"MLX recurrent cache has {len(cache)} slots; {mode} depth={depth} expects {layout.num_slots}") + new_caches = [None]*layout.num_slots + for i in range(cfg.prelude_layers): + x, new_caches[i] = run(i, x) + # Lookahead observes the post-prelude state (Torch parity). + lookahead_probs = self.lookahead_router(x) + # MLX always returns a cache, so per_iteration (one slot per + # iteration) never exits early; shared_first/first_only may. + allow_exit = mode != "per_iteration" + diagnostics = [] + for group in range(cfg.recurrent_groups): + i = cfg.prelude_layers + group + states, halts = [], [] + exited, shared = False, None + halted = cumulative = None + target = min(1.0, max(1e-6, cfg.halting_threshold)) + for step in range(depth): + conditioned = controller.condition(x, step, group) + slot = layout.slot(i, step) + if mode == "per_iteration": + refined, new_caches[slot] = run(i, conditioned, slot) + elif mode == "first_only": + refined, candidate = run(i, conditioned, slot, latent_only=step > 0) + if step == 0: + new_caches[slot] = candidate + elif step == 0: + refined, new_caches[slot], shared = run(i, conditioned, slot, return_shared_kv=True) + else: + refined, _ = run(i, conditioned, kv_override=shared) + x = controller.contract(x, refined, group, cfg) + states.append(x) + prob = controller.halt_probability(x) + halts.append(prob) + # Exact per-token ACT halting (Torch parity, see model_chronos). + if step + 1 >= cfg.min_recurrent_depth: + clipped = mx.clip(prob.astype(mx.float32), 0.0, 1.0) + if halted is None: + cumulative = mx.zeros_like(clipped) + halted = mx.zeros(clipped.shape, dtype=mx.bool_) + newly = mx.logical_and(mx.logical_not(halted), cumulative + clipped >= target) + cumulative = mx.where(mx.logical_or(halted, newly), cumulative, mx.minimum(cumulative + clipped, 1.0)) + halted = mx.logical_or(halted, newly) + if not self.training and adaptive and allow_exit and halted is not None and step + 1 < depth: + if bool(mx.all(halted).item()): + exited = True + break + mixed, state = auxiliary_state(states, halts, cfg.min_recurrent_depth, cfg.halting_threshold) + if adaptive: + x = mixed + else: + x = states[-1] + state["ponder_loss"] = mx.zeros(()) + state["depth_used"] = mx.full(halts[0].shape, len(states), dtype=mx.float32) + state["exited"] = exited + diagnostics.append(state) + for i in range(cfg.prelude_layers+cfg.recurrent_groups, len(self.layers)): + slot = layout.slot(i, 0) + x, new_caches[slot] = run(i, x, slot) + max_depth = max(s["halting_probs"].shape[-1] for s in diagnostics) + halting = [mx.pad(s["halting_probs"], [(0,0), (0,0), (0,max_depth-s["halting_probs"].shape[-1])], constant_values=1) for s in diagnostics] + combined = {key: mx.stack([s[key] for s in diagnostics]).mean() for key in ("ponder_loss", "distill_loss", "stability_loss", "contraction")} + combined.update(halting_probs=mx.stack(halting, axis=-2), depth_used=mx.stack([s["depth_used"] for s in diagnostics], axis=-1), + early_exit_ratio=mx.array(sum(int(s["exited"]) for s in diagnostics)/len(diagnostics))) + object.__setattr__(self, "last_recurrent_state", combined) + else: + for i in range(len(self.layers)): + x, candidate = run(i, x) + new_caches.append(candidate) + if i == 0: + lookahead_probs = self.lookahead_router(x) x = self.norm(x) - logits = self.lm_head(x) + head = self.embed_tokens.weight if self.lm_head is None else self.lm_head.weight + if logits_dtype is not None and logits_dtype != head.dtype: + logits = x.astype(logits_dtype) @ head.astype(logits_dtype).T + else: + logits = self.embed_tokens.as_linear(x) if self.lm_head is None else self.lm_head(x) return logits, lookahead_probs, new_caches @staticmethod - def from_chronos_pytorch(pt_model, config, include_experts: bool = True): + def from_chronos_pytorch(pt_model, config, include_experts: bool = True, dtype: str | None = None): """ Convert a trained PyTorch ChronosForCausalLM to ChronosMLXModel. - Copies weights layer by layer via numpy as bridge. + + ``dtype`` is the MLX compute dtype (``chronos.precision.resolve_mlx``); + each tensor is converted once, straight from the checkpoint dtype + (BF16 crosses bit-exactly). ``None`` keeps the legacy FP32 load. """ + from chronos.precision import torch_to_mx mlx_model = ChronosMLXModel(config) + target = dtype or "float32" def _to_mx(pt_tensor): - return mx.array(pt_tensor.detach().float().cpu().numpy()) + return torch_to_mx(pt_tensor, target) def _assign_linear(linear, sd, key, missing): if key in sd: + if tuple(linear.weight.shape) != tuple(sd[key].shape): + raise ValueError(f"MLX checkpoint shape mismatch for {key}: {tuple(sd[key].shape)} != {tuple(linear.weight.shape)}") linear.weight = _to_mx(sd[key]) else: missing.append(key) @@ -174,6 +315,8 @@ def _assign_norm(norm, sd, key, missing): ("model.embed_tokens.weight", mlx_model.embed_tokens), ("lm_head.weight", mlx_model.lm_head), ]: + if obj is None: + continue if key in sd: obj.weight = _to_mx(sd[key]) else: @@ -220,6 +363,20 @@ def _assign_norm(norm, sd, key, missing): missing, ) + else: + for proj in ("gate_proj", "up_proj", "down_proj"): + _assign_linear(getattr(moe, proj), sd, f"{prefix}.mlp.{proj}.weight", missing) + + controller = mlx_model.recurrent_controller + for key in ("loop_embeddings", "group_embeddings", "step_sizes"): + source = f"model.recurrent_controller.{key}" + if source in sd: + setattr(controller, key, _to_mx(sd[source])) + else: + missing.append(source) + _assign_norm(controller.halting_norm, sd, "model.recurrent_controller.halting_norm.weight", missing) + _assign_linear(controller.halting_head, sd, "model.recurrent_controller.halting_head.weight", missing) + controller.halting_head.bias = _to_mx(sd["model.recurrent_controller.halting_head.bias"]) proj_layers = mlx_model.lookahead_router.proj.layers _assign_linear(proj_layers[0], sd, "model.lookahead_router.proj.0.weight", missing) _assign_linear(proj_layers[2], sd, "model.lookahead_router.proj.2.weight", missing) @@ -229,5 +386,8 @@ def _assign_norm(norm, sd, key, missing): "MLX conversion is missing Chronos checkpoint tensors: " + ", ".join(missing[:40]) ) + if dtype is not None: + mlx_model.set_dtype(getattr(mx, target)) # params not in the checkpoint (fresh init) + mlx_model.train(pt_model.training) mx.eval(mlx_model.parameters()) return mlx_model diff --git a/chronos/mlx/moe.py b/chronos/mlx/moe.py index 090010c..ec162e8 100644 --- a/chronos/mlx/moe.py +++ b/chronos/mlx/moe.py @@ -9,6 +9,17 @@ import mlx.core as mx import mlx.nn as nn +from chronos.mlx.moe_kernels import ( + gather_mm_supported, grouped_swiglu, sparse_dispatch_ok, stack_expert_weights, training_availability, +) + + +def _sparse_ok(dtype) -> bool: + """Sorted gather_mm dispatch is used only where it passed the one-time + self-check for this device/dtype (CPU: FP32 only). Otherwise the exact + masked-dense path runs: slower, never wrong.""" + return sparse_dispatch_ok(dtype) + class FeedForwardMLX(nn.Module): """SwiGLU FFN matching minimind's FeedForward structure.""" @@ -52,7 +63,20 @@ def __init__(self, config): FeedForwardMLX(H, I) for _ in range(config.num_shared_experts) ] object.__setattr__(self, "_last_router_probs", None) # [B, S, E] + # One entry per call: a recurrent block dispatches once per iteration. + object.__setattr__(self, "_iteration_router_probs", []) + # router_scope="shared": later iterations reuse iteration zero's ids. + object.__setattr__(self, "router_scope", str(getattr(config, "router_scope", "per_iteration") or "per_iteration")) + object.__setattr__(self, "_shared_topk_idx", None) + # Per-forward cache of the stacked [E, out, in] expert weights: a + # recurrent block runs this MoE once per iteration, but its weights + # only change between optimizer updates. Opened/closed by the model + # forward; outside a scope every call stacks fresh (always correct). + object.__setattr__(self, "_pack_scope", None) self.runtime_miss_policy = "sync_on_demand" + self.fallback_mask_prob = float(getattr(config, "fallback_mask_prob", 0.0) or 0.0) + # Sorted gather_mm dispatch (requires mlx >= 0.32.4 on M5-class GPUs). + self.sparse_dispatch = bool(getattr(config, "moe_sparse_dispatch", True)) and gather_mm_supported() self.runtime_on_demand_loader = None self.runtime_touch_expert = None @@ -60,6 +84,33 @@ def __init__(self, config): def last_router_probs(self): return self._last_router_probs + @property + def iteration_router_probs(self): + return self._iteration_router_probs + + def reset_iteration_routing(self): + object.__setattr__(self, "_iteration_router_probs", []) + object.__setattr__(self, "_shared_topk_idx", None) + + def open_pack_scope(self): + object.__setattr__(self, "_pack_scope", {}) + + def close_pack_scope(self): + object.__setattr__(self, "_pack_scope", None) + + def _expert_stacks(self): + scope = self._pack_scope + if scope is None: + return stack_expert_weights(self.experts) + # Keyed on the live weight arrays, so an expert swapped in mid-forward + # (lazy promotion) can never be served from a stale stack. + key = tuple(id(getattr(e, p).weight) for e in self.experts for p in ("gate_proj", "up_proj", "down_proj")) + hit = scope.get("w") + if hit is None or hit[0] != key: + hit = (key, stack_expert_weights(self.experts)) + scope["w"] = hit + return hit[1] + @last_router_probs.setter def last_router_probs(self, value): object.__setattr__(self, "_last_router_probs", value) @@ -82,18 +133,42 @@ def __call__( scores = mx.softmax(self.gate(x_flat).astype(mx.float32), axis=-1) # [N, E] object.__setattr__(self, "_last_router_probs", scores.reshape(B, S, self.num_experts)) + self._iteration_router_probs.append(self._last_router_probs) # top-k routing - topk_idx = mx.argpartition(-scores, kth=self.num_experts_per_tok - 1, axis=-1) - topk_idx = topk_idx[:, :self.num_experts_per_tok] # [N, K] + shared_idx = self._shared_topk_idx + if (self.router_scope == "shared" and shared_idx is not None + and shared_idx.shape[0] == scores.shape[0] and len(self._iteration_router_probs) > 1): + topk_idx = shared_idx + else: + topk_idx = mx.argpartition(-scores, kth=self.num_experts_per_tok - 1, axis=-1) + # Routing indices are discrete: gradients flow through topk_weight + # (gathered scores) only, never through the integer indices. + topk_idx = mx.stop_gradient(topk_idx[:, :self.num_experts_per_tok]) # [N, K] + if self.router_scope == "shared": + object.__setattr__(self, "_shared_topk_idx", topk_idx) topk_weight = mx.take_along_axis(scores, topk_idx, axis=-1) # [N, K] if self.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(axis=-1, keepdims=True) + 1e-20) y = mx.zeros_like(x_flat) - if available_expert_mask is None: - # Training: standard MoE scatter + if available_expert_mask is None and self.sparse_dispatch and _sparse_ok(x_flat.dtype) and not any( + isinstance(e, LazyFeedForwardMLX) for e in self.experts + ): + # Sparse dispatch: each token runs only its top-k experts. + shared = self._shared_out(x_flat) if len(self.shared_experts) > 0 else None + fallback = None + if self.training and self.fallback_mask_prob > 0.0 and shared is not None: + # Cache-miss simulation (PyTorch parity): an unavailable routed + # expert is served by the shared expert, as at offload inference. + fallback = (training_availability(self.num_experts, self.fallback_mask_prob), shared) + y = grouped_swiglu(x_flat, topk_idx, topk_weight, *self._expert_stacks(), + fallback=fallback) + if shared is not None: + y = y + shared + elif available_expert_mask is None: + # Masked-dense reference path (old mlx or explicit opt-out). for i, expert in enumerate(self.experts): # mask over tokens where expert i is selected sel = (topk_idx == i) # [N, K] bool @@ -112,6 +187,12 @@ def __call__( else: avail = available_expert_mask.astype(x_flat.dtype) # [E] float shared_out = self._shared_out(x_flat) # [N, H] + selected_host = None + if python_avail is not None: + # One host read for every expert's selection (was one + # ``.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() for i, expert in enumerate(self.experts): sel = (topk_idx == i) @@ -119,7 +200,7 @@ def __call__( w = mx.where(sel, topk_weight, 0.0).sum(axis=-1, keepdims=True) if python_avail is not None: - selected = bool(mx.any(tok_mask).item()) + selected = bool(selected_host[i]) if not selected: continue is_avail = i in python_avail diff --git a/chronos/mlx/moe_kernels.py b/chronos/mlx/moe_kernels.py new file mode 100644 index 0000000..65a082a --- /dev/null +++ b/chronos/mlx/moe_kernels.py @@ -0,0 +1,156 @@ +"""Sparse expert execution for MLX: sorted ``gather_mm`` grouped SwiGLU. + +Only the top-k selected experts run for each token (no dense E-way +compute). Known backend defects and how they are handled: + +* mlx 0.31.x on Apple M5 GPUs lacks fp32 GatherMM Metal kernels and returns + wrong bf16 results (measured 2026-11-05, relative gradient error ~1.0). +* mlx 0.31.1 on M1 Max: a ``[:, None, :]`` row view feeding a sorted + gather and a ``squeeze`` before the inverse permutation gave wrong forward + values / down-projection gradients (performance branch, 2026-11-06). + ``grouped_swiglu`` therefore uses explicit ``reshape(-1, 1, H)`` / + ``reshape(-1, H)`` layouts (measured identical to the dense oracle on mlx + 0.32.4 / M5 at N=8192, E=36, H=384, I=160). +* CPU GatherMM is float32-only. + +Instead of trusting a version number, :func:`sparse_dispatch_ok` runs a +one-time self-check per (device, dtype) against the dense oracle (value and +all gradients, > 64 sorted assignments) and falls back to the masked-dense +path if it fails. ``dense_reference_swiglu`` is that oracle. +""" +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + +MIN_MLX_GATHER_MM = (0, 32, 4) # first version validated on M5 (fp32 + bf16 Metal kernels) +_SELF_CHECK: dict[tuple[str, str], bool] = {} + + +def mlx_version_tuple() -> tuple[int, int, int]: + parts = (mx.__version__.split("+")[0].split(".") + ["0", "0"])[:3] + return tuple(int("".join(ch for ch in p if ch.isdigit()) or 0) for p in parts) + + +def gather_mm_supported() -> bool: + """``gather_mm`` exists and passes the self-check on the default device (fp32).""" + return sparse_dispatch_ok(mx.float32) + + +def sparse_dispatch_ok(dtype) -> bool: + """Whether sorted ``gather_mm`` dispatch is correct for ``dtype`` on the + current default device (memoised self-check; ~10 ms once per process).""" + if not hasattr(mx, "gather_mm"): + return False + device = mx.default_device() + if device.type == mx.DeviceType.cpu and not cpu_gather_mm_dtype_ok(dtype): + return False + key = (str(device), str(dtype)) + if key not in _SELF_CHECK: + _SELF_CHECK[key] = _self_check(dtype) + return _SELF_CHECK[key] + + +def _self_check(dtype) -> bool: + """Sorted grouped SwiGLU vs the masked-dense oracle on this device/dtype. + + 192 assignments (above any small-batch shortcut), all four weight/input + gradients plus the routing-weight gradient. Reported defects are O(0.1-1) + relative errors; TF32-class / bf16 rounding is O(1e-3-1e-2). + """ + try: + key = mx.random.key(20261106) + k = mx.random.split(key, 6) + N, K, E, H, I = 96, 2, 6, 32, 24 + x = mx.random.normal((N, H), key=k[0]).astype(dtype) + wg = (mx.random.normal((E, I, H), key=k[1]) * 0.2).astype(dtype) + wu = (mx.random.normal((E, I, H), key=k[2]) * 0.2).astype(dtype) + wd = (mx.random.normal((E, H, I), key=k[3]) * 0.2).astype(dtype) + scores = mx.softmax(mx.random.normal((N, E), key=k[4]), axis=-1) + idx = mx.argpartition(-scores, kth=K - 1, axis=-1)[:, :K] + w = mx.take_along_axis(scores, idx, axis=-1).astype(dtype) + cot = mx.random.normal((N, H), key=k[5]) + + def run(fn): + def loss(x, wg, wu, wd, w): + return (fn(x, idx, w, wg, wu, wd).astype(mx.float32) * cot).sum() + v, g = mx.value_and_grad(loss, argnums=(0, 1, 2, 3, 4))(x, wg, wu, wd, w) + mx.eval(v, g) + return [v] + list(g) + + tol = 2e-2 if dtype == mx.float32 else 8e-2 + for a, b in zip(run(grouped_swiglu), run(dense_reference_swiglu)): + a, b = a.astype(mx.float32).reshape(-1), b.astype(mx.float32).reshape(-1) + err = mx.linalg.norm(a - b) / mx.maximum(mx.linalg.norm(b), 1e-12) + if not bool(mx.isfinite(err).item()) or float(err.item()) > tol: + return False + return True + except Exception: + return False + + +def stack_expert_weights(experts) -> tuple[mx.array, mx.array, mx.array]: + """Stack per-expert SwiGLU weights into [E, out, in] arrays (differentiable).""" + return ( + mx.stack([e.gate_proj.weight for e in experts]), + mx.stack([e.up_proj.weight for e in experts]), + mx.stack([e.down_proj.weight for e in experts]), + ) + + +def cpu_gather_mm_dtype_ok(dtype) -> bool: + """MLX CPU GatherMM is float32-only (mlx 0.32.4).""" + return dtype == mx.float32 + + +def _gather_mm(a: mx.array, b: mx.array, idx: mx.array) -> mx.array: + return mx.gather_mm(a, b, rhs_indices=idx, sorted_indices=True) + + +def grouped_swiglu(x: mx.array, topk_idx: mx.array, topk_weight: mx.array, + w_gate: mx.array, w_up: mx.array, w_down: mx.array, + fallback: tuple[mx.array, mx.array] | None = None) -> mx.array: + """Σ_k w[n,k]·SwiGLU_{idx[n,k]}(x[n]) with one sorted gather_mm per projection. + + x: [N, H]; topk_idx/topk_weight: [N, K]; w_gate/w_up: [E, I, H]; w_down: [E, H, I]. + ``fallback=(available[E] bool, shared_out[N, H])`` replaces the output of + every unavailable expert with the shared-expert output (cache-miss + simulation, identical to the PyTorch masked path). + """ + N, K = topk_idx.shape + H = x.shape[-1] + flat = mx.stop_gradient(topk_idx).reshape(-1) + order = mx.stop_gradient(mx.argsort(flat)) + sorted_experts = flat[order] + # Explicit reshapes (not ``[:, None, :]`` / ``squeeze``): see module doc. + rows = x[order // K].reshape(-1, 1, H) # [N·K, 1, H] + gate = _gather_mm(rows, w_gate.swapaxes(-1, -2), sorted_experts) + up = _gather_mm(rows, w_up.swapaxes(-1, -2), sorted_experts) + out = _gather_mm(nn.silu(gate) * up, w_down.swapaxes(-1, -2), sorted_experts).reshape(-1, H) + out = out[mx.stop_gradient(mx.argsort(order))].reshape(N, K, H) + if fallback is not None: + available, shared = fallback + live = mx.stop_gradient(available[topk_idx])[..., None] # [N, K, 1] + out = mx.where(live, out, shared[:, None, :].astype(out.dtype)) + return (out * topk_weight[..., None].astype(out.dtype)).sum(axis=1) + + +def dense_reference_swiglu(x: mx.array, topk_idx: mx.array, topk_weight: mx.array, + w_gate: mx.array, w_up: mx.array, w_down: mx.array, + fallback=None) -> mx.array: + """Masked-dense oracle: runs every expert on every token (tests only).""" + out = mx.zeros(x.shape, dtype=x.dtype) + for e in range(w_gate.shape[0]): + y = (nn.silu(x @ w_gate[e].T) * (x @ w_up[e].T)) @ w_down[e].T + if fallback is not None and not bool(fallback[0][e].item()): + y = fallback[1].astype(y.dtype) + weight = mx.where(topk_idx == e, topk_weight, 0).sum(axis=-1, keepdims=True) + out = out + y * weight.astype(y.dtype) + return out + + +def training_availability(num_experts: int, miss_prob: float) -> mx.array: + """Random per-call expert availability; at least one expert stays live.""" + keep = mx.random.uniform(shape=(num_experts,)) >= miss_prob + forced = mx.arange(num_experts) == mx.random.randint(0, num_experts, ()) + return mx.where(mx.any(keep), keep, forced) diff --git a/chronos/mlx/recurrent.py b/chronos/mlx/recurrent.py new file mode 100644 index 0000000..cc8b795 --- /dev/null +++ b/chronos/mlx/recurrent.py @@ -0,0 +1,89 @@ +"""MLX execution of the Chronos recurrent controller and ACT mixture.""" +import random +import mlx.core as mx +import mlx.nn as nn + + +class RecurrentControllerMLX(nn.Module): + def __init__(self, config): + super().__init__() + h, depth, groups = config.hidden_size, config.recurrent_depth, config.recurrent_groups + self.loop_embeddings = mx.zeros((depth, h)) + self.group_embeddings = mx.zeros((groups, h)) + self.step_sizes = mx.full((groups,), config.recurrent_dt_init) + self.halting_norm = nn.RMSNorm(h, eps=config.rms_norm_eps) + self.halting_head = nn.Linear(h, 1) + self.halting_head.weight = mx.zeros((1, h)) + self.halting_head.bias = mx.full((1,), -1.0) + self.training_progress = 0.0 + + def set_training_progress(self, step, total): + self.training_progress = min(1.0, max(0.0, step / max(total, 1))) + + def adaptive_active(self, cfg, training): + if not cfg.adaptive_halting or cfg.recurrent_strategy == "fixed": + return False + if not training: + return True + if cfg.depth_schedule in {"fixed", "constant", "off"}: + return False + return (cfg.depth_schedule in {"adaptive", "act", "on"} + or self.training_progress >= cfg.depth_schedule_warmup_fraction) + + def execution_depth(self, cfg, training): + ceiling = cfg.recurrent_depth if training else cfg.max_inference_depth + if not training or cfg.recurrent_strategy == "fixed" or cfg.recurrent_depth_jitter == 0: + return ceiling + warmup = min(.99, max(0.0, cfg.depth_schedule_warmup_fraction)) + progress = max(0.0, (self.training_progress - warmup) / (1-warmup)) + high = cfg.min_recurrent_depth + round((ceiling-cfg.min_recurrent_depth)*progress) + low = max(cfg.min_recurrent_depth, high-round((high-cfg.min_recurrent_depth)*cfg.recurrent_depth_jitter)) + return random.randint(low, high) + + def condition(self, x, step, group): + return x + self.loop_embeddings[min(step, self.loop_embeddings.shape[0]-1)][None, None, :] + self.group_embeddings[group][None, None, :] + + def contract(self, before, after, group, cfg): + dt = mx.clip(self.step_sizes[group], cfg.recurrent_dt_min, cfg.recurrent_dt_max).astype(after.dtype) + return before + dt*(after-before) + + def halt_probability(self, x): + return mx.sigmoid(self.halting_head(self.halting_norm(x)).squeeze(-1).astype(mx.float32)) + + +def auxiliary_state(states, probabilities, minimum, threshold): + # A draft forward may run fewer iterations than ``min_recurrent_depth``; + # the last executed state must still receive the remaining mass (Torch parity). + minimum = min(int(minimum), len(states)) + target = min(1.0, max(1e-6, threshold)) + cumulative = mx.zeros_like(probabilities[0]) + halted = mx.zeros(cumulative.shape, dtype=mx.bool_) + mixed, depth, accumulated = mx.zeros_like(states[0]), mx.zeros_like(cumulative), [] + for idx, (state, prob) in enumerate(zip(states, probabilities)): + if idx + 1 < minimum: + weight = mx.zeros_like(prob) + elif idx == len(states)-1: + weight = mx.where(halted, 0.0, 1-cumulative) + else: + clipped = mx.clip(prob, 0.0, 1.0) + new_halt = ~halted & (cumulative+clipped >= target) + weight = mx.where(~halted, mx.where(new_halt, 1-cumulative, clipped), 0.0) + halted = halted | new_halt + weight = mx.maximum(weight, 0) + cumulative = mx.clip(cumulative+weight, 0.0, 1.0) + mixed = mixed + state*weight[..., None].astype(state.dtype) + depth = depth + weight*(idx+1) + accumulated.append(cumulative) + accumulated[-1] = mx.ones_like(cumulative) + # Shared recurrent weights require gradients through BOTH endpoints; + # a detached final state creates an expansion-rewarding semi-gradient. + final = states[-1].astype(mx.float32) + zero = mx.zeros(()) + distill = mx.stack([mx.mean((s.astype(mx.float32)-final)**2) for s in states[:-1]]).mean() if len(states)>1 else zero + deltas = [states[i]-states[i-1] for i in range(1,len(states))] + stability = mx.stack([mx.mean((deltas[i]-deltas[i-1]).astype(mx.float32)**2) for i in range(1,len(deltas))]).mean() if len(deltas)>1 else zero + # Relative step norm is a diagnostic; it is not a Lipschitz guarantee. + relative = mx.stack([mx.sqrt(mx.maximum(mx.mean(d.astype(mx.float32)**2), 1e-12)) / mx.sqrt(mx.maximum(mx.mean(s.astype(mx.float32)**2), 1e-12)) for d,s in zip(deltas, states[:-1])]).mean() if deltas else zero + return mixed, {"halting_probs": mx.stack(accumulated, axis=-1), "depth_used": depth, + "ponder_loss": mx.maximum(depth-1, 0).mean(), "distill_loss": distill, + "stability_loss": stability, "contraction": relative} diff --git a/chronos/mlx/speculative.py b/chronos/mlx/speculative.py new file mode 100644 index 0000000..9eb232e --- /dev/null +++ b/chronos/mlx/speculative.py @@ -0,0 +1,227 @@ +"""Recurrent-depth self-speculative decoding for the MLX model (Torch parity). + +Same algorithm and loop invariant as ``chronos.runtime.speculative``: + +* draft γ tokens with the same weights at ``draft_depth`` < full depth; +* verify them with one full-depth forward over ``[pending, d1..dγ]`` + (weights and experts are read once for γ+1 positions); +* greedy output is token-for-token identical to full-depth greedy decoding, + sampled output uses rejection sampling and keeps the full-depth distribution; +* the persistent cache is only extended by verify forwards and then truncated + to the accepted prefix. MLX arrays are immutable, so the draft extends a + separate list of entries and the committed cache is never modified. + +Only depth-independent cache layouts (``shared_first`` / ``first_only``) can +share one KV stream between draft and verify; ``per_iteration`` falls back to +plain decoding. +""" +from __future__ import annotations + +import time + +import mlx.core as mx + +from chronos.runtime.speculative import SpecStats + + +def cache_length(cache) -> int: + """Absolute token count of an MLX Chronos cache (MLA ``(c_kv,...)`` or SWA ``(k, v, end)``).""" + for entry in cache or (): + if entry is not None: + return int(entry[2]) if len(entry) > 2 else int(entry[0].shape[1]) + return 0 + + +def _truncate_entry(entry, length: int): + if entry is None: + return None + if len(entry) <= 2: # MLA: (c_kv,) or (c_kv, k_rope) + return tuple(t[:, :length] for t in entry) + k, v, end = entry + drop = int(end) - length + if drop < 0: + raise ValueError("cannot extend a cache by truncation") + if drop == 0: + return entry + if drop > k.shape[1]: + raise ValueError(f"rollback of {drop} exceeds retained sliding-window cache {k.shape[1]}") + return (k[:, :k.shape[1] - drop], v[:, :v.shape[1] - drop], length) + + +def truncate_cache(cache, length: int): + return [_truncate_entry(e, length) for e in cache] + + +def _cache_arrays(cache) -> list: + """Every array inside a cache (plain, int8 ``QuantizedKVMLX`` or SWA end offsets skipped).""" + out = [] + for entry in cache or (): + for t in entry or (): + if isinstance(t, mx.array): + out.append(t) + elif hasattr(t, "arrays"): + out.extend(t.arrays()) + return out + + +def set_cache_slack(model, slack: int) -> None: + """Let every sliding-window layer keep ``slack`` extra entries for rollback.""" + from chronos.mlx.attention import SlidingWindowAttentionMLX + for layer in model.layers: + attn = getattr(layer, "self_attn", None) + if isinstance(attn, SlidingWindowAttentionMLX): + object.__setattr__(attn, "cache_slack", max(int(getattr(attn, "cache_slack", 0)), int(slack))) + + +def _probs(logits: mx.array, temperature: float, top_k: int, top_p: float) -> mx.array: + logits = logits.astype(mx.float32) / max(float(temperature), 1e-6) + if top_k and top_k > 0 and top_k < logits.shape[-1]: + kth = mx.sort(logits)[-int(top_k)] + logits = mx.where(logits < kth, -mx.inf, logits) + if top_p < 1.0: + order = mx.argsort(-logits) + sorted_logits = logits[order] + cum = mx.cumsum(mx.softmax(sorted_logits), axis=-1) + drop_sorted = (cum - mx.softmax(sorted_logits)) >= top_p # keep the first token always + drop = drop_sorted[mx.argsort(order)] # back to vocabulary order + logits = mx.where(drop, -mx.inf, logits) + return mx.softmax(logits) + + +def _categorical(p: mx.array, key) -> int: + return int(mx.random.categorical(mx.log(mx.maximum(p, 1e-30)), key=key).item()) + + +def rejection_step(p: mx.array, q: mx.array, token: int, key) -> tuple[bool, int]: + """Accept ``token`` (drawn from ``q``) with prob. min(1, p/q), else draw from norm(max(p − q, 0)).""" + k1, k2 = mx.random.split(key) + ratio = min(1.0, float(p[token].item()) / max(float(q[token].item()), 1e-12)) + if float(mx.random.uniform(key=k1).item()) <= ratio: + return True, int(token) + residual = mx.maximum(p - q, 0.0) + total = float(residual.sum().item()) + return False, _categorical(residual / total if total > 0 else p, k2) + + +def speculative_tokens(model, input_ids, stats: SpecStats, *, max_new_tokens: int = 128, draft_depth: int = 1, + gamma: int = 4, adaptive_gamma: bool = True, gamma_max: int = 8, do_sample: bool = False, + temperature: float = 1.0, top_k: int = 0, top_p: float = 1.0, eos_token_id=None, + seed: int | None = None, forward_fn=None, should_stop=None): + """Yield generated token ids one by one (batch size 1), filling ``stats``. + + Same loop invariant as the Torch version: ``pending`` is emitted but not + yet cached; one full-depth forward over ``[pending, d1..dγ]`` caches it + and verifies the drafts, and the cache is truncated to the accepted + prefix. ``forward_fn(ids, cache, depth) -> (logits, lookahead, cache)`` + lets a runtime add expert masks / prefetch around every forward.""" + from chronos.config_schema import resolve_recurrent_attention_mode + ids = mx.array(input_ids, dtype=mx.int32) + if ids.ndim == 1: + ids = ids[None] + if ids.shape[0] != 1: + raise ValueError("speculative decoding supports batch size 1") + cfg = model.config + full_depth = int(getattr(cfg, "max_inference_depth", None) or cfg.recurrent_depth) + draft_depth = max(1, min(int(draft_depth), full_depth)) + recurrent = bool(getattr(cfg, "recurrent_enabled", False)) and int(cfg.recurrent_depth) > 1 + if not recurrent or resolve_recurrent_attention_mode(cfg) == "per_iteration": + draft_depth = full_depth # no shared KV stream → plain decoding + eos = set([eos_token_id] if isinstance(eos_token_id, int) else (eos_token_id or [])) + gamma_max = max(1, int(gamma_max)) + gamma = max(1, min(int(gamma), gamma_max)) + set_cache_slack(model, gamma_max + 1) + key = mx.random.key(int(seed) if seed is not None else int(time.time_ns() % (2**31))) + t0 = time.perf_counter() + if forward_fn is None: + def forward_fn(x, cache, depth): + return model(x, cache=cache, recurrent_depth=depth) + + def next_key(): + nonlocal key + key, sub = mx.random.split(key) + return sub + + def dist(row): + return _probs(row, temperature, top_k, top_p) + + def pick(row): + if not do_sample: + return int(mx.argmax(row).item()) + return _categorical(dist(row), next_key()) + + try: + logits, _, cache = forward_fn(ids, None, None) + mx.eval(logits, *_cache_arrays(cache)) + pending = pick(logits[0, -1]) + stats.generated += 1 + yield pending + while stats.generated < max_new_tokens and pending not in eos: + if should_stop is not None and should_stop(): + break + committed = cache_length(cache) + room = max_new_tokens - stats.generated + draft, draft_q, dcache, last = [pending], [], cache, pending + if draft_depth < full_depth: + for _ in range(min(gamma, room - 1) if room > 1 else 0): + dl, _, dcache = forward_fn(mx.array([[last]], dtype=mx.int32), dcache, draft_depth) + stats.draft_calls += 1 + row = dl[0, -1] + last = pick(row) # .item() evaluates the draft step + if do_sample: + draft_q.append(dist(row)) + draft.append(last) + if last in eos: + break + vlogits, _, vcache = forward_fn(mx.array([draft], dtype=mx.int32), cache, None) + stats.verify_calls += 1 + rows = vlogits[0] # rows[i] predicts the token after draft[i] + if not do_sample: + targets = [int(t) for t in mx.argmax(rows, axis=-1).tolist()] # one host sync per round + n_ok, nxt, stop = 0, None, False + for i in range(1, len(draft)): + tok = draft[i] + if not do_sample: + ok, correction = targets[i - 1] == tok, targets[i - 1] + else: + ok, correction = rejection_step(dist(rows[i - 1]), draft_q[i - 1], tok, next_key()) + if not ok: + nxt = correction + break + n_ok += 1 + stats.generated += 1 + yield tok + if tok in eos or stats.generated >= max_new_tokens: + stop = True + break + stats.proposed += len(draft) - 1 + stats.accepted += n_ok + cache = truncate_cache(vcache, committed + 1 + n_ok) + mx.eval(*_cache_arrays(cache)) + if stop: + break + if nxt is None: # every draft accepted: bonus token + nxt = targets[len(draft) - 1] if not do_sample else pick(rows[len(draft) - 1]) + pending = nxt + stats.generated += 1 + yield pending + if adaptive_gamma and len(draft) > 1: + rate = n_ok / (len(draft) - 1) + if rate >= 0.8: + gamma = min(gamma_max, gamma + 1) + elif rate < 0.4: + gamma = max(1, gamma - 1) + stats.gamma_history.append(gamma) + finally: + stats.elapsed_s = time.perf_counter() - t0 + + +def speculative_generate(model, input_ids, *, on_token=None, **kwargs): + """Return ``(prompt + generated ids as list[int], SpecStats)``; kwargs as ``speculative_tokens``.""" + stats = SpecStats() + ids = mx.array(input_ids, dtype=mx.int32) + tokens = [int(t) for t in (ids if ids.ndim == 1 else ids[0]).tolist()] + for tok in speculative_tokens(model, ids, stats, **kwargs): + tokens.append(tok) + if on_token: + on_token(tok) + return tokens, stats diff --git a/chronos/mlx/training/inspect_checkpoint.py b/chronos/mlx/training/inspect_checkpoint.py new file mode 100644 index 0000000..9494a46 --- /dev/null +++ b/chronos/mlx/training/inspect_checkpoint.py @@ -0,0 +1,179 @@ +"""Inspect a LongRun state checkpoint without disturbing the training process. + +Loads ``/checkpoints/latest.state.pkl`` on the MLX **CPU** device (so a +running GPU job keeps its throughput), then reports: + +* held-out CE / perplexity on the phase's fixed validation windows; +* routing balance and recurrent-depth statistics on those windows; +* greedy en→zh / zh→en continuations on FLORES+ dev with BLEU / chrF; + +and optionally exports a serveable ``.pth`` + ``.config.json`` (with the +run's own tokenizer recorded) for ``chronos serve``. The report is written to +``/inspections/_s.json`` so the monitor can show it. + + python -m chronos.mlx.training.inspect_checkpoint --run-dir runs/enzh_50m_v1 \ + --export serve_models/chronos_enzh_pretrain_s300.pth +""" +from __future__ import annotations + +import argparse +import json +import math +import os +import pickle +import time + +import mlx.core as mx +import mlx.nn as nn + + +def load_run_model(run_dir: str, *, device: str = "cpu", state_file: str = "latest.state.pkl"): + from chronos.model.config import ChronosConfig + from chronos.mlx.training.longrun import _restore_arrays + from chronos.mlx.training.trainer import _load_mlx_from_checkpoint + if device == "cpu": + mx.set_default_device(mx.cpu) + spec = json.load(open(os.path.join(run_dir, "run_spec.json"), encoding="utf-8")) + config = ChronosConfig(**spec["model_config"]) + model = _load_mlx_from_checkpoint(config, None, dtype_name=spec.get("dtype")) + with open(os.path.join(run_dir, "checkpoints", state_file), "rb") as fh: + payload = pickle.load(fh) # written by this project's own trainer + model.update(_restore_arrays(payload["params"], payload["dtypes"])) + # The policy picked the compute dtype at build time (CPU -> FP32: exact + # upcast of the BF16 weights; CPU GatherMM is FP32-only). Restored + # arrays keep their saved dtype, so convert them once to that dtype. + model.set_dtype(model._chronos_precision.mx_dtype()) + mx.eval(model.parameters()) + model.eval() + return spec, config, model, payload["state"] + + +def _phase(spec: dict, state: dict) -> dict: + return spec["phases"][min(int(state["phase_index"]), len(spec["phases"]) - 1)] + + +def held_out(spec, config, model, phase, n_batches: int = 8) -> dict: + from chronos.mlx.training.longrun import PretrainSource, SFTSource, recurrent_metrics, routing_metrics + from chronos.mlx.training.trainer import _collect_router_probs + source = (PretrainSource(spec["shards_dir"], phase["seq_len"], spec["seed"]) if phase["name"] == "pretrain" + else SFTSource(spec["shards_dir"], spec["seed"])) + total, count, extra = 0.0, 0, {} + for i, (ids, labels) in enumerate(source.valid_batches(n_batches, phase["micro_tokens"])): + logits, _, _ = model(mx.array(ids, dtype=mx.int32), logits_dtype=mx.float32) + lb = mx.array(labels, dtype=mx.int32).reshape(-1) + mask = lb != -100 + ce = nn.losses.cross_entropy(logits.reshape(-1, logits.shape[-1]), + mx.where(mask, lb, 0), reduction="none") + s, n = (ce * mask).sum(), mask.sum() + mx.eval(s, n) + total, count = total + float(s.item()), count + int(n.item()) + if i == 0: + extra = {**routing_metrics(_collect_router_probs(model), int(config.num_experts_per_tok)), + **recurrent_metrics(getattr(model, "last_recurrent_state", None))} + ce = total / max(count, 1) + return {"valid_ce": round(ce, 4), "valid_ppl": round(math.exp(min(ce, 30)), 2), "valid_tokens": count, **extra} + + +def greedy(model, tok, prompt_ids: list[int], stop_ids: set[int], limit: int) -> list[int]: + logits, _, cache = model(mx.array([prompt_ids], dtype=mx.int32)) + out = [] + for _ in range(limit): + nxt = int(mx.argmax(logits[0, -1]).item()) + if nxt in stop_ids: + break + out.append(nxt) + logits, _, cache = model(mx.array([[nxt]], dtype=mx.int32), cache=cache) + return out + + +def translation_probe(spec, model, phase_name: str, n: int) -> dict: + """Pretrain checkpoints are probed with the ``<|en|>…<|zh|>`` format they + were trained on; SFT checkpoints with the chat prompt.""" + import sacrebleu + from tokenizers import Tokenizer + from chronos.data.corpus.evalsets import load_eval_sets + from chronos.data.corpus.shards import SFT_PROMPT_EN_ZH + tok = Tokenizer.from_file(os.path.join(spec["tokenizer_dir"], "tokenizer.json")) + tid = tok.token_to_id + stop = {tid(t) for t in ("<|endoftext|>", "<|im_end|>", "<|en|>", "<|zh|>")} + raw = (spec.get("extra") or {}).get("raw_root") + pairs = load_eval_sets(raw, ["flores_dev"]).get("flores_dev", [])[:n] if raw else [] + if not pairs: + return {"probe": "skipped: FLORES+ dev not found"} + out, t0 = {}, time.monotonic() + directions = [("en_zh", 0, 1)] + ([("zh_en", 1, 0)] if phase_name == "pretrain" else []) + for name, si, ti in directions: + hyps, samples = [], [] + for pair in pairs: + src, ref = pair[si], pair[ti] + if phase_name == "pretrain": + tags = ("<|en|>", "<|zh|>") if name == "en_zh" else ("<|zh|>", "<|en|>") + prompt = [tid(tags[0])] + tok.encode(src, add_special_tokens=False).ids + [tid(tags[1])] + else: + text = f"<|im_start|>user\n{SFT_PROMPT_EN_ZH}{src}<|im_end|>\n<|im_start|>assistant\n" + prompt = tok.encode(text, add_special_tokens=False).ids + src_len = len(tok.encode(src, add_special_tokens=False).ids) + hyp = tok.decode(greedy(model, tok, prompt, stop, int(1.6 * src_len) + 16), skip_special_tokens=True) + hyps.append(hyp) + if len(samples) < 6: + samples.append({"src": src, "ref": ref, "hyp": hyp}) + refs = [p[ti] for p in pairs] + bleu = sacrebleu.corpus_bleu(hyps, [refs], tokenize="zh" if ti == 1 else "13a") + chrf = sacrebleu.corpus_chrf(hyps, [refs]) + out[name] = {"n": len(pairs), "bleu": round(bleu.score, 2), "chrf": round(chrf.score, 2), + "samples": samples} + out["probe_s"] = round(time.monotonic() - t0, 1) + out["probe_set"] = "flores_dev" + return out + + +def export(model, config, spec, out_path: str, stage: str) -> str: + from chronos.model.checkpoint import tokenizer_metadata + from chronos.mlx.training.trainer import _save_mlx_checkpoint + os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True) + _save_mlx_checkpoint(model, config, out_path, stage=stage) + cfg_path = out_path[:-4] + ".config.json" + meta = json.load(open(cfg_path, encoding="utf-8")) + meta["tokenizer"] = tokenizer_metadata(os.path.abspath(spec["tokenizer_dir"])) + json.dump(meta, open(cfg_path, "w", encoding="utf-8"), indent=2, sort_keys=True) + return out_path + + +def inspect_run(run_dir: str, *, samples: int = 24, n_batches: int = 8, export_path: str | None = None, + device: str = "cpu") -> dict: + t0 = time.monotonic() + spec, config, model, state = load_run_model(run_dir, device=device) + phase = _phase(spec, state) + report = {"time": time.strftime("%Y-%m-%dT%H:%M:%S"), "run_dir": os.path.abspath(run_dir), + "phase": phase["name"], "step": int(state["phase_step"]), "tokens_seen": int(state["tokens_seen"]), + "device": device, "compute_dtype": model._chronos_precision.compute, + "precision": model._chronos_precision.as_dict(), + "train_dtype": spec.get("dtype"), "load_s": round(time.monotonic() - t0, 1)} + report["held_out"] = held_out(spec, config, model, phase, n_batches) + if samples: + report["translation"] = translation_probe(spec, model, phase["name"], samples) + if export_path: + report["exported"] = export(model, config, spec, export_path, phase["name"]) + report["elapsed_s"] = round(time.monotonic() - t0, 1) + out_dir = os.path.join(run_dir, "inspections") + os.makedirs(out_dir, exist_ok=True) + with open(os.path.join(out_dir, f"{phase['name']}_s{report['step']}.json"), "w", encoding="utf-8") as fh: + json.dump(report, fh, ensure_ascii=False, indent=2) + return report + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("--run-dir", required=True) + p.add_argument("--samples", type=int, default=24, help="FLORES+ dev sentences per direction (0 = skip)") + p.add_argument("--batches", type=int, default=8, help="validation micro-batches") + p.add_argument("--export", default=None, help="write a serveable .pth here") + p.add_argument("--device", default="cpu", choices=("cpu", "gpu")) + a = p.parse_args(argv) + print(json.dumps(inspect_run(a.run_dir, samples=a.samples, n_batches=a.batches, export_path=a.export, + device=a.device), ensure_ascii=False, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/mlx/training/io.py b/chronos/mlx/training/io.py index 06e4ad4..1a1eb98 100644 --- a/chronos/mlx/training/io.py +++ b/chronos/mlx/training/io.py @@ -10,16 +10,17 @@ def _torch_tensor(arr) -> torch.Tensor: + """Exact bridge: the tensor keeps its MLX dtype (BF16 via bit-view).""" + from chronos.precision import mx_to_torch if isinstance(arr, mx.array): - arr = arr.astype(mx.float32) - mx.eval(arr) - return torch.from_numpy(np.array(arr, dtype=np.float32)).contiguous() + return mx_to_torch(arr).contiguous() + return torch.from_numpy(np.array(arr)).contiguous() def mlx_state_to_torch(mlx_model, config) -> dict[str, torch.Tensor]: sd: dict[str, torch.Tensor] = {} sd["model.embed_tokens.weight"] = _torch_tensor(mlx_model.embed_tokens.weight) - sd["lm_head.weight"] = _torch_tensor(mlx_model.lm_head.weight) + sd["lm_head.weight"] = _torch_tensor(mlx_model.embed_tokens.weight if mlx_model.lm_head is None else mlx_model.lm_head.weight) sd["model.norm.weight"] = _torch_tensor(mlx_model.norm.weight) for i, blk in enumerate(mlx_model.layers): @@ -55,6 +56,16 @@ def mlx_state_to_torch(mlx_model, config) -> dict[str, torch.Tensor]: getattr(expert, proj).weight ) + else: + for proj in ("gate_proj", "up_proj", "down_proj"): + sd[f"{prefix}.mlp.{proj}.weight"] = _torch_tensor(getattr(moe, proj).weight) + + controller = mlx_model.recurrent_controller + for key in ("loop_embeddings", "group_embeddings", "step_sizes"): + sd[f"model.recurrent_controller.{key}"] = _torch_tensor(getattr(controller, key)) + sd["model.recurrent_controller.halting_norm.weight"] = _torch_tensor(controller.halting_norm.weight) + sd["model.recurrent_controller.halting_head.weight"] = _torch_tensor(controller.halting_head.weight) + sd["model.recurrent_controller.halting_head.bias"] = _torch_tensor(controller.halting_head.bias) proj_layers = mlx_model.lookahead_router.proj.layers sd["model.lookahead_router.proj.0.weight"] = _torch_tensor(proj_layers[0].weight) sd["model.lookahead_router.proj.2.weight"] = _torch_tensor(proj_layers[2].weight) diff --git a/chronos/mlx/training/longrun.py b/chronos/mlx/training/longrun.py new file mode 100644 index 0000000..d4bc316 --- /dev/null +++ b/chronos/mlx/training/longrun.py @@ -0,0 +1,730 @@ +"""Time-budgeted, resumable MLX training runs (pretrain → SFT → evaluation). + +Designed for multi-day benchmark runs on Apple silicon: + +* The wall-clock budget is split across phases; after a short throughput + probe every phase's optimizer-step count is derived from the *measured* + tokens/s, so the LR schedule (warmup + cosine) ends when the budget does. +* Every ``checkpoint_every_s`` the full state (BF16 weights, FP32 masters and + moments, data-stream RNG, phase/step) is written atomically; ``resume`` + continues bit-for-bit from the last snapshot. +* Metrics go to ``metrics.jsonl`` (one JSON object per log interval) — + losses by component, recurrent depth/halting, routing balance/entropy, + lookahead recall, grad norm, LR, tokens/s, memory, plus periodic + validation CE and BLEU/chrF. The UI monitor tails this file. +""" +from __future__ import annotations + +import json +import math +import os +import pickle +import time +from dataclasses import asdict, dataclass, field + +import mlx.core as mx +import mlx.nn as nn +import mlx.optimizers as optim +import numpy as np +from mlx.utils import tree_flatten, tree_unflatten + + +@dataclass +class PhaseSpec: + name: str # "pretrain" | "sft" + budget_fraction: float + peak_lr: float + tokens_per_step: int # optimizer-step token budget (micro × accum) + micro_tokens: int = 8192 # tokens per forward/backward + seq_len: int = 512 # pretrain window + warmup_fraction: float = 0.02 + min_lr_ratio: float = 0.1 + weight_decay: float = 0.1 + grad_clip: float = 1.0 + + +@dataclass +class RunSpec: + run_dir: str + shards_dir: str + tokenizer_dir: str + model_config: dict + phases: list[PhaseSpec] + budget_hours: float = 72.0 + eval_reserve_fraction: float = 0.04 + log_every_steps: int = 20 + validate_every_s: float = 3600.0 + checkpoint_every_s: float = 1800.0 + bleu_samples: int = 200 # periodic FLORES devtest subset during SFT + final_eval_samples: int = 1012 # full FLORES+ devtest at the end (0 = skip) + seed: int = 1234 + dtype: str = "bfloat16" + # "time": warmup + cosine over the phase's wall-clock budget, so the LR + # reaches its floor exactly at the deadline even if throughput varies + # (shared GPU, thermal throttling). "step": classic step schedule. + schedule: str = "time" + # Telemetry (monitor pages): per-tensor weight/grad/update statistics every + # ``tensor_stats_every_s`` (0 = off) and per-virtual-layer routing. + tensor_stats_every_s: float = 1800.0 + # Greedy FLORES+ dev probe at each pretrain validation (``<|en|>…<|zh|>`` + # format, both directions); samples go to ``samples.jsonl``. 0 = off. + probe_samples: int = 16 + # Start from these weights instead of a fresh initialization (a bundle's + # ``model/*.pth``). The checkpoint's topology and vocabulary must match + # ``model_config``; optimizer state starts fresh. Ignored on resume. + init_checkpoint: str | None = None + # Exported checkpoints are ``/_[.best].pth``. + export_prefix: str = "chronos_enzh" + extra: dict = field(default_factory=dict) + + +def lr_at(step: int, total: int, peak: float, warmup_fraction: float, min_ratio: float) -> float: + warmup = max(1, int(total * warmup_fraction)) + if step < warmup: + return peak * (step + 1) / warmup + progress = min(1.0, (step - warmup) / max(1, total - warmup)) + return peak * (min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * progress))) + + +def lr_at_fraction(progress: float, peak: float, warmup_fraction: float, min_ratio: float, + first_steps: int) -> float: + """Warmup + cosine as a function of elapsed budget fraction in [0, 1].""" + progress = min(1.0, max(0.0, progress)) + wf = max(1e-6, warmup_fraction) + if progress < wf: + return peak * max(progress / wf, 1e-3 * min(first_steps + 1, 10)) + t = (progress - wf) / max(1e-6, 1 - wf) + return peak * (min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * t))) + + +def _atomic_write(path: str, write) -> None: + tmp = path + ".tmp" + write(tmp) + os.replace(tmp, path) + + +def _snapshot_arrays(tree) -> dict: + return {k: np.array(v.astype(mx.float32)) if v.dtype == mx.bfloat16 else np.array(v) + for k, v in tree_flatten(tree)} + + +def _snapshot_dtypes(tree) -> dict: + return {k: str(v.dtype) for k, v in tree_flatten(tree)} + + +def _restore_arrays(flat: dict, dtypes: dict): + items = [] + for k, v in flat.items(): + arr = mx.array(v) + if dtypes.get(k) == "mlx.core.bfloat16": + arr = arr.astype(mx.bfloat16) + items.append((k, arr)) + return tree_unflatten(items) + + +def _optimizer_snapshot(optimizer) -> dict: + flat = {k: v for k, v in tree_flatten(optimizer.state) if isinstance(v, mx.array)} + return {"arrays": {k: np.array(v) for k, v in flat.items()}, + "dtypes": {k: str(v.dtype) for k, v in flat.items()}} + + +def _optimizer_restore(optimizer, model, snap: dict) -> None: + optimizer.init(model.trainable_parameters()) # recreates python-side decay flags + current = dict(tree_flatten(optimizer.state)) + for k, v in snap["arrays"].items(): + if k in current: + current[k] = mx.array(v) + optimizer.state = tree_unflatten(list(current.items())) + + +def routing_metrics(router: mx.array | None, top_k: int) -> dict: + """Load balance and confidence over every dispatch ``[B, S, V, E]``.""" + if router is None: + return {} + E = router.shape[-1] + flat = router.reshape(-1, router.shape[2], E).astype(mx.float32) # [N, V, E] + top = mx.argpartition(-flat, kth=top_k - 1, axis=-1)[..., :top_k] + load = (top[..., None] == mx.arange(E)).astype(mx.float32).sum(-2).mean(0) / top_k # [V, E] + importance = flat.mean(0) + entropy = -(flat * mx.log(mx.maximum(flat, 1e-9))).sum(-1).mean() + cv = (load.std(-1) / mx.maximum(load.mean(-1), 1e-9)).mean() + dead = (load.max(0) < (0.05 / E)).sum() + mx.eval(load, importance, entropy, cv, dead) + return {"route_entropy": float(entropy.item()), "route_entropy_max": float(math.log(E)), + "load_cv": float(cv.item()), "dead_experts": int(dead.item()), + "max_expert_load": float(load.max().item() * E)} + + +def recurrent_metrics(state: dict | None) -> dict: + if not state: + return {} + out = {} + for key in ("ponder_loss", "distill_loss", "stability_loss", "contraction", "early_exit_ratio"): + if key in state: + out["rec_" + key] = float(mx.array(state[key]).item()) + if "depth_used" in state: + d = state["depth_used"].astype(mx.float32) + out["rec_depth_mean"] = float(d.mean().item()) + out["rec_depth_per_group"] = [round(float(v), 3) for v in np.array(d.mean(axis=(0, 1))).tolist()] + if "halting_probs" in state: + h = state["halting_probs"].astype(mx.float32) # [B, S, G, D] accumulated mass + out["rec_halt_mass_by_step"] = [round(float(v), 3) for v in np.array(h.mean(axis=(0, 1, 2))).tolist()] + return out + + +class PretrainSource: + def __init__(self, shards_dir, seq_len, seed): + from chronos.data.corpus.shards import PackedTokenStream + self.stream = PackedTokenStream(os.path.join(shards_dir, "pretrain", "train"), seq_len, seed=seed) + self.valid = PackedTokenStream(os.path.join(shards_dir, "pretrain", "valid"), seq_len, seed=seed + 1) + + def batch(self, micro_tokens: int): + rows = max(1, micro_tokens // self.stream.seq_len) + arr = self.stream.batch(rows) + return arr[:, :-1], arr[:, 1:] # inputs, next-token labels + + def valid_batches(self, n: int, micro_tokens: int): + rng_state = self.valid.rng.bit_generator.state + self.valid.rng = np.random.default_rng(0) # fixed validation windows + rows = max(1, micro_tokens // self.valid.seq_len) + out = [self.valid.batch(rows) for _ in range(n)] + self.valid.rng.bit_generator.state = rng_state + return [(a[:, :-1], a[:, 1:]) for a in out] + + def rng_state(self): + return self.stream.rng.bit_generator.state + + def set_rng_state(self, state): + self.stream.rng.bit_generator.state = state + + +class SFTSource: + def __init__(self, shards_dir, seed): + from chronos.data.corpus.shards import SFTRecords + self.records = SFTRecords(os.path.join(shards_dir, "sft", "train"), seed=seed) + self.valid = SFTRecords(os.path.join(shards_dir, "sft", "valid"), seed=seed + 1) + + @staticmethod + def _shift(ids, labels): + return ids[:, :-1], labels[:, 1:] # labels already -100 outside replies + + def batch(self, micro_tokens: int): + return self._shift(*self.records.batch(micro_tokens)) + + def valid_batches(self, n: int, micro_tokens: int): + state = self.valid.rng.bit_generator.state + self.valid.rng = np.random.default_rng(0) + out = [self._shift(*self.valid.batch(micro_tokens)) for _ in range(n)] + self.valid.rng.bit_generator.state = state + return out + + def rng_state(self): + return self.records.rng.bit_generator.state + + def set_rng_state(self, state): + self.records.rng.bit_generator.state = state + + +def check_init_checkpoint(path: str, config) -> None: + """Fail fast when ``init_checkpoint`` does not fit ``model_config``. + + ``_load_mlx_from_checkpoint`` silently starts from random weights when the + file is missing. A mismatched topology or vocabulary would instead fail + deep inside ``load_state_dict``, so both are checked here. + """ + from chronos.model.checkpoint import SNIFFABLE_TOPOLOGY_FIELDS, UNSNIFFABLE_REQUIRED_FIELDS, read_checkpoint_config + if not os.path.isfile(path): + raise FileNotFoundError(f"init_checkpoint not found: {path}") + saved, _ = read_checkpoint_config(path) + if not saved: + return # no sidecar: load_state_dict_controlled still checks shapes + keys = SNIFFABLE_TOPOLOGY_FIELDS + UNSNIFFABLE_REQUIRED_FIELDS + ( + "recurrent_depth", "prelude_layers", "coda_layers", "recurrent_attention_mode", "rope_layout") + diff = {k: (saved[k], getattr(config, k, None)) for k in keys + if k in saved and saved[k] is not None and saved[k] != getattr(config, k, None)} + if diff: + raise ValueError(f"init_checkpoint {path} does not match model_config: " + + ", ".join(f"{k}: checkpoint {a!r} vs spec {b!r}" for k, (a, b) in diff.items())) + + +class LongRun: + def __init__(self, spec: RunSpec, log=print): + from chronos.model.config import ChronosConfig + from chronos.mlx.model import ChronosMLXModel + from chronos.mlx.training.optim import MixedPrecisionAdamW + from chronos.mlx.training.trainer import _load_mlx_from_checkpoint + self.spec, self.log = spec, log + os.makedirs(spec.run_dir, exist_ok=True) + self.config = ChronosConfig(**spec.model_config) + import torch + torch.manual_seed(spec.seed) # reproducible initialization (built in Torch, then converted) + mx.random.seed(spec.seed) + init = spec.init_checkpoint + if init: + init = os.path.abspath(os.path.expanduser(init)) + check_init_checkpoint(init, self.config) + self.model = _load_mlx_from_checkpoint(self.config, init, dtype_name=spec.dtype) + self.model.train() + self.optimizer = MixedPrecisionAdamW(learning_rate=0.0, weight_decay=spec.phases[0].weight_decay) + self.optimizer.init(self.model.trainable_parameters()) + self.state = {"phase_index": 0, "phase_step": 0, "phase_total": None, "elapsed_s": 0.0, + "wall_s": 0.0, + "tokens_seen": 0, "phase_tokens_per_s": {}, "rng": {}, "best_valid": {}} + self.metrics_path = os.path.join(spec.run_dir, "metrics.jsonl") + self.ckpt_dir = os.path.join(spec.run_dir, "checkpoints") + os.makedirs(self.ckpt_dir, exist_ok=True) + with open(os.path.join(spec.run_dir, "run_spec.json"), "w", encoding="utf-8") as fh: + json.dump(asdict(spec), fh, ensure_ascii=False, indent=2) + from chronos.mlx.training import telemetry + self.telemetry = telemetry + self.heartbeat = telemetry.Heartbeat(spec.run_dir) + self.precision = getattr(self.model, "_chronos_precision", None) + self.runtime = telemetry.runtime_info(self.model, spec, self.precision) + try: + telemetry._atomic_json(os.path.join(spec.run_dir, "runtime.json"), self.runtime) + except OSError: + pass + self.parameters_path = os.path.join(spec.run_dir, "parameters.jsonl") + self.heartbeat.beat("starting", precision=self.runtime["precision"].get("compute"), + code=self.runtime["code"]["sha256"]) + + # ---- persistence ------------------------------------------------------- + def save_state(self, tag="latest"): + if getattr(self, "_run_start", None) is not None: + self.state["wall_s"] = time.monotonic() - self._run_start + source = getattr(self, "_active_source", None) + if source is not None: + phase = self.spec.phases[min(self.state["phase_index"], len(self.spec.phases) - 1)] + self.state["rng"][phase.name] = source.rng_state() + + path = os.path.join(self.ckpt_dir, f"{tag}.state.pkl") + params = self.model.parameters() + payload = {"state": self.state, "params": _snapshot_arrays(params), "dtypes": _snapshot_dtypes(params), + "optimizer": _optimizer_snapshot(self.optimizer)} + _atomic_write(path, lambda tmp: pickle.dump(payload, open(tmp, "wb"), protocol=5)) + return path + + def try_resume(self) -> bool: + path = os.path.join(self.ckpt_dir, "latest.state.pkl") + if not os.path.exists(path): + return False + with open(path, "rb") as fh: + payload = pickle.load(fh) # written by this run; not untrusted input + self.model.update(_restore_arrays(payload["params"], payload["dtypes"])) + _optimizer_restore(self.optimizer, self.model, payload["optimizer"]) + self.state = payload["state"] + mx.eval(self.model.parameters(), self.optimizer.state) + self.log(f"[run] resumed phase={self.state['phase_index']} step={self.state['phase_step']} " + f"elapsed={self.state['elapsed_s'] / 3600:.2f}h") + return True + + def export_checkpoint(self, name: str, stage: str): + from chronos.mlx.training.trainer import _save_mlx_checkpoint + out = os.path.join(self.spec.run_dir, f"{name}.pth") + _save_mlx_checkpoint(self.model, self.config, out, stage=stage) + cfg_path = out.replace(".pth", ".config.json") + if os.path.exists(cfg_path): # record the run's own tokenizer, not minimind's + from chronos.model.checkpoint import tokenizer_metadata + meta = json.load(open(cfg_path, encoding="utf-8")) + meta["tokenizer"] = tokenizer_metadata(os.path.abspath(self.spec.tokenizer_dir)) + json.dump(meta, open(cfg_path, "w", encoding="utf-8"), indent=2, sort_keys=True) + return out + + def _emit(self, record: dict): + record = {"time": time.strftime("%Y-%m-%dT%H:%M:%S"), **record} + with open(self.metrics_path, "a", encoding="utf-8") as fh: + fh.write(json.dumps(record, ensure_ascii=False) + "\n") + + def _beat(self, state: str | None = None, **fields): + st = self.state + phases = self.spec.phases + pi = min(int(st.get("phase_index", 0)), len(phases) - 1) + run_start = getattr(self, "_run_start", None) + self.heartbeat.beat( + state, phase=phases[pi].name if st.get("phase_index", 0) < len(phases) else "done", + phase_index=int(st.get("phase_index", 0)), step=int(st.get("phase_step", 0)), total=st.get("phase_total"), + phase_wall_s=st.get("phase_wall_s"), phase_budget_s=st.get("phase_budget_s"), + run_wall_s=(time.monotonic() - run_start) if run_start is not None else st.get("wall_s"), + tokens_seen=st.get("tokens_seen"), tokens_processed=st.get("tokens_processed"), + rollbacks=int(st.get("rollbacks", 0)), lr_scale=float(st.get("lr_scale", 1.0)), + best_valid=st.get("best_valid"), **fields) + + def _tensor_telemetry(self, phase, step: int, grads, before, accum: int, lr: float, norm: float): + """One ``parameters.jsonl`` row: per-tensor weight RMS/|max|, pre-clip + gradient RMS, relative update, per-group norms, expert grid, routing.""" + tel = self.telemetry + t0 = time.monotonic() + params = dict(tree_flatten(self.model.parameters())) + stats = tel.tensor_stats(params, dict(tree_flatten(grads)) if grads is not None else None, + dict(tree_flatten(before)) if before is not None else None, + grad_scale=1.0 / max(1, accum)) + row = {"time": time.strftime("%Y-%m-%dT%H:%M:%S"), "phase": phase.name, "step": step, "lr": lr, + "grad_norm": norm, **stats, "experts": tel.expert_table(stats["tensors"]), + "routing": tel.routing_snapshot(self.model)} + row["telemetry_s"] = round(time.monotonic() - t0, 3) + with open(self.parameters_path, "a", encoding="utf-8") as fh: + fh.write(json.dumps(row, ensure_ascii=False) + "\n") + return row + + # ---- training ---------------------------------------------------------- + def _loss_fn(self, model, ids, labels): + from chronos.mlx.training.trainer import _chronos_regularized_loss, draft_exit_kl, draft_exit_logits + draft = draft_exit_logits(model, ids, model.config) # None unless draft_exit_kl_coef > 0 + logits, lookahead, _ = model(ids, logits_dtype=mx.float32) + lg = logits.reshape(-1, logits.shape[-1]) + lb = labels.reshape(-1) + mask = lb != -100 + ce = nn.losses.cross_entropy(lg, mx.where(mask, lb, 0), reduction="none") + ce = (ce * mask).sum() / mx.maximum(mask.sum(), 1) + # Labels here are pre-aligned with the logits (no shift). + draft_kl = None if draft is None else draft_exit_kl(draft, logits, labels != -100) + return _chronos_regularized_loss(model, ce, lookahead, model.config, metrics={"ce": ce, "primary": ce}, + draft_kl=draft_kl) + + def _source(self, phase: PhaseSpec): + if phase.name == "pretrain": + return PretrainSource(self.spec.shards_dir, phase.seq_len, self.spec.seed) + return SFTSource(self.spec.shards_dir, self.spec.seed) + + def _micro_step(self, grad_fn, ids, labels): + loss, grads = grad_fn(self.model, mx.array(ids, dtype=mx.int32), mx.array(labels, dtype=mx.int32)) + return loss, grads + + def _apply(self, grads, n_micro, phase, lr): + from chronos.mlx.training.trainer import _scale_grads + grads = _scale_grads(grads, 1.0 / n_micro) + grads, norm = optim.clip_grad_norm(grads, phase.grad_clip) + self.optimizer.learning_rate = lr + self.optimizer.weight_decay = phase.weight_decay + self.optimizer.update(self.model, grads) + return norm + + def _validate(self, source, phase, n_batches=8) -> float: + self.model.eval() + total, count = 0.0, 0 + for ids, labels in source.valid_batches(n_batches, phase.micro_tokens): + logits, _, _ = self.model(mx.array(ids, dtype=mx.int32), logits_dtype=mx.float32) + lb = mx.array(labels, dtype=mx.int32).reshape(-1) + mask = lb != -100 + ce = nn.losses.cross_entropy(logits.reshape(-1, logits.shape[-1]), + mx.where(mask, lb, 0), reduction="none") + s, n = (ce * mask).sum(), mask.sum() + mx.eval(s, n) + total, count = total + float(s.item()), count + int(n.item()) + self.model.train() + return total / max(count, 1) + + def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): + """Train phase ``pi`` until its step total, its deadline, or ``max_steps`` + more optimizer updates (used by tests and staged runs).""" + from chronos.mlx.training.trainer import _add_grads, _collect_router_probs, _tree_all_finite + phase = self.spec.phases[pi] + source = self._source(phase) + if self.state["rng"].get(phase.name): + source.set_rng_state(self.state["rng"][phase.name]) + accum = max(1, phase.tokens_per_step // phase.micro_tokens) + grad_fn = nn.value_and_grad(self.model, self._loss_fn) + controller = self.model.recurrent_controller + top_k = int(self.config.num_experts_per_tok) + t_phase = time.monotonic() + last_ckpt = last_val = time.monotonic() + window_tokens, window_processed, window_t = 0, 0, time.monotonic() + self._active_source = source + stop_at = None if max_steps is None else self.state["phase_step"] + int(max_steps) + if self.state["phase_step"] == 0 and not self.state.get("phase_budget_s"): + self.state["phase_budget_s"] = max(1.0, deadline_s - time.monotonic()) + self.state["phase_wall_s"] = 0.0 + last_wall = time.monotonic() + last_tensors = -math.inf # first update after (re)start is sampled + tensor_every = float(getattr(self.spec, "tensor_stats_every_s", 0.0) or 0.0) + self._emit({"phase": phase.name, "event": "phase_start", "step": self.state["phase_step"], + "deadline_h": round((deadline_s - time.monotonic()) / 3600, 3)}) + self._beat("running") + while True: + step, total = self.state["phase_step"], self.state["phase_total"] + if total is not None and step >= total and self.spec.schedule != "time": + break + if stop_at is not None and step >= stop_at: + self.state["rng"][phase.name] = source.rng_state() + return + if time.monotonic() >= deadline_s: + self.log(f"[run] phase {phase.name} reached its wall-clock budget at step {step}") + break + now_wall = time.monotonic() + self.state["phase_wall_s"] = self.state.get("phase_wall_s", 0.0) + (now_wall - last_wall) + last_wall = now_wall + time_based = self.spec.schedule == "time" and self.state.get("phase_budget_s") + progress = (self.state["phase_wall_s"] / self.state["phase_budget_s"]) if time_based else \ + step / max(total or 1_000_000, 1) + controller.set_training_progress(int(progress * 1_000_000), 1_000_000) + # MLX's global RNG (cache-miss simulation, dropout) cannot be + # snapshotted, so make each step's randomness a pure function of + # (seed, phase, step, rollbacks): resume is then bit-exact. + rb = int(self.state.get("rollbacks", 0)) + mx.random.seed((self.spec.seed * 1_000_003 + pi * 100_003 + step + rb * 7_919) % (2 ** 31)) + if time_based: + lr = lr_at_fraction(progress, phase.peak_lr, phase.warmup_fraction, phase.min_lr_ratio, step) + else: + lr = lr_at(step, total or 1_000_000, phase.peak_lr, phase.warmup_fraction, phase.min_lr_ratio) + lr *= float(self.state.get("lr_scale", 1.0)) + t0 = time.monotonic() + acc, tokens, processed, loss_sum = None, 0, 0, mx.zeros(()) + for _ in range(accum): + ids, labels = source.batch(phase.micro_tokens) + loss, grads = self._micro_step(grad_fn, ids, labels) + acc = grads if acc is None else _add_grads(acc, grads) + loss_sum = loss_sum + loss + tokens += int((labels != -100).sum()) # supervised tokens + processed += int(ids.size) # forward/backward tokens + mx.eval(acc, loss_sum) # bound the lazy graph to one micro-batch + # Tensor telemetry keeps references to the pre-update weights (MLX + # arrays are immutable: no copy) and the FP32 accumulated grads. + want_tensors = tensor_every > 0 and time.monotonic() - last_tensors >= tensor_every + before = self.model.parameters() if want_tensors else None + norm = self._apply(acc, accum, phase, lr) + mx.eval(self.model.parameters(), self.optimizer.state, norm) + dt = time.monotonic() - t0 + if not math.isfinite(float(norm.item())) or not math.isfinite(float(loss_sum.item())): + # Restore the last snapshot but keep wall-clock progress, halve + # the LR and change the data/RNG stream so the failing step is + # not replayed verbatim. + keep = {k: self.state.get(k) for k in ("phase_wall_s", "phase_budget_s")} + rollbacks = int(self.state.get("rollbacks", 0)) + 1 + lr_scale = max(0.125, float(self.state.get("lr_scale", 1.0)) * 0.5) + self.log(f"[run] non-finite step {step}; rollback #{rollbacks}, lr_scale={lr_scale}") + if not self.try_resume(): + raise FloatingPointError("non-finite loss before first snapshot") + self.state.update(keep, rollbacks=rollbacks, lr_scale=lr_scale) + if self.state["rng"].get(phase.name): + source.set_rng_state(self.state["rng"][phase.name]) + source.batch(phase.micro_tokens) # skip ahead in the data stream + if rollbacks > 8: + raise FloatingPointError("too many non-finite rollbacks") + self._emit({"phase": phase.name, "event": "rollback", "step": step, "rollbacks": rollbacks, + "lr_scale": lr_scale}) + self._beat("running", last_rollback=time.time()) + last_wall = time.monotonic() + continue + self.state["phase_step"] = step + 1 + if want_tensors: + try: + self._tensor_telemetry(phase, step + 1, acc, before, accum, lr, float(norm.item())) + except Exception as exc: # noqa: BLE001 - telemetry must never stop training + self.log(f"[telemetry] tensor stats skipped: {exc!r}") + last_tensors = time.monotonic() + before = None + self.state["tokens_seen"] += tokens + self.state["tokens_processed"] = self.state.get("tokens_processed", 0) + processed + self.state["elapsed_s"] += dt + window_tokens += tokens + window_processed += processed + refresh_eta = self.spec.schedule == "time" and (step + 1) % 50 == 0 + if (total is None and step + 1 >= 6) or refresh_eta: + # Throughput probe (first time) or ETA refresh (time schedule). + tps = window_tokens / max(time.monotonic() - window_t, 1e-6) + remaining = max(0.0, deadline_s - time.monotonic()) + self.state["phase_total"] = step + 1 + int(remaining * tps / max(tokens, 1) * 0.97) + if phase.name not in self.state["phase_tokens_per_s"]: + self.log(f"[run] {phase.name}: {tps:,.0f} tok/s → ~{self.state['phase_total']} steps " + f"(~{self.state['phase_total'] * tokens / 1e9:.2f}B tokens)") + self.state["phase_tokens_per_s"][phase.name] = tps + if (step + 1) % self.spec.log_every_steps == 0 or step < 3: + comps = {k: round(float(v.item()), 5) for k, v in getattr(self.model, "_last_loss_components", {}).items()} + rec = {"phase": phase.name, "step": step + 1, "total": self.state["phase_total"], + "lr": lr, "grad_norm": float(norm.item()), "loss": float(loss_sum.item()) / accum, + "tokens_per_s": window_tokens / max(time.monotonic() - window_t, 1e-6), + "processed_tokens_per_s": window_processed / max(time.monotonic() - window_t, 1e-6), + "tokens_processed": self.state["tokens_processed"], + "tokens_seen": self.state["tokens_seen"], "step_s": dt, + "peak_mem_gb": mx.get_peak_memory() / 1e9, "active_mem_gb": mx.get_active_memory() / 1e9, + "cache_mem_gb": mx.get_cache_memory() / 1e9, + "phase_wall_s": self.state.get("phase_wall_s"), "phase_budget_s": self.state.get("phase_budget_s"), + "lr_scale": float(self.state.get("lr_scale", 1.0)), "rollbacks": int(self.state.get("rollbacks", 0)), + **comps, **recurrent_metrics(getattr(self.model, "last_recurrent_state", None)), + **routing_metrics(_collect_router_probs(self.model), top_k)} + self._emit(rec) + self.log(f"[{phase.name}] {step + 1}/{self.state['phase_total']} loss={rec['loss']:.4f} " + f"ce={comps.get('primary_loss', 0):.4f} gn={rec['grad_norm']:.2f} lr={lr:.2e} " + f"{rec['tokens_per_s']:,.0f} tok/s depth={rec.get('rec_depth_mean', 0):.2f}") + window_tokens, window_processed, window_t = 0, 0, time.monotonic() + self._beat("running", last={k: rec.get(k) for k in ( + "loss", "primary_loss", "grad_norm", "lr", "tokens_per_s", "step_s", "peak_mem_gb", + "active_mem_gb", "rec_depth_mean", "load_cv")}) + now = time.monotonic() + if now - last_val >= self.spec.validate_every_s: + self._beat("evaluating") + self._periodic_eval(source, phase) + last_val = time.monotonic() + self._beat("running") + if now - last_ckpt >= self.spec.checkpoint_every_s: + self._beat("checkpointing") + self.state["rng"][phase.name] = source.rng_state() + self.save_state() + last_ckpt = time.monotonic() + self._beat("running", last_checkpoint=time.time()) + self.state["rng"][phase.name] = source.rng_state() + self._beat("evaluating") + self._periodic_eval(source, phase) + self.export_checkpoint(f"{self.spec.export_prefix}_{phase.name}", stage=phase.name) + self._emit({"phase": phase.name, "event": "phase_end", "step": self.state["phase_step"], + "tokens_seen": self.state["tokens_seen"]}) + + def _periodic_eval(self, source, phase): + ce = self._validate(source, phase) + rec = {"phase": phase.name, "step": self.state["phase_step"], "event": "validation", "valid_ce": ce, + "valid_ppl": math.exp(min(ce, 30))} + if phase.name == "sft" and self.spec.bleu_samples: + rec.update(self.translation_eval(self.spec.bleu_samples)) + elif phase.name == "pretrain" and getattr(self.spec, "probe_samples", 0): + rec.update(self._pretrain_probe(int(self.spec.probe_samples))) + best = self.state["best_valid"].get(phase.name) + if best is None or ce < best: + self.state["best_valid"][phase.name] = ce + self.export_checkpoint(f"{self.spec.export_prefix}_{phase.name}.best", stage=phase.name) + self._emit(rec) + self.log(f"[eval] {json.dumps(rec, ensure_ascii=False)}") + + def _pretrain_probe(self, n: int) -> dict: + from chronos.mlx.training.inspect_checkpoint import translation_probe + self.model.eval() + try: + probe = translation_probe(asdict(self.spec), self.model, "pretrain", n) + except Exception as exc: # noqa: BLE001 - a probe failure must not stop training + self.log(f"[eval] pretrain probe skipped: {exc!r}") + return {} + finally: + self.model.train() + out = {} + with open(os.path.join(self.spec.run_dir, "samples.jsonl"), "a", encoding="utf-8") as fh: + for direction in ("en_zh", "zh_en"): + res = probe.get(direction) + if not isinstance(res, dict): + continue + out[f"probe_{direction}_bleu"], out[f"probe_{direction}_chrf"] = res["bleu"], res["chrf"] + for sample in res.get("samples", [])[:3]: + fh.write(json.dumps({"time": time.strftime("%Y-%m-%dT%H:%M:%S"), "phase": "pretrain", + "step": self.state["phase_step"], "direction": direction, + "set": probe.get("probe_set"), **sample}, ensure_ascii=False) + "\n") + if out: + out["probe_n"], out["probe_s"] = n, probe.get("probe_s") + return out + + # ---- translation quality (greedy, KV cache, MLX GPU) ---------------------- + def _tokenizer(self): + if not hasattr(self, "_tok"): + from tokenizers import Tokenizer + self._tok = Tokenizer.from_file(os.path.join(self.spec.tokenizer_dir, "tokenizer.json")) + return self._tok + + def translate(self, sources: list[str], max_new_tokens: int = 160) -> list[str]: + from chronos.data.corpus.shards import SFT_PROMPT_EN_ZH + tok = self._tokenizer() + end = tok.token_to_id("<|im_end|>") + self.model.eval() + out = [] + for src in sources: + prompt = f"<|im_start|>user\n{SFT_PROMPT_EN_ZH}{src}<|im_end|>\n<|im_start|>assistant\n" + ids = tok.encode(prompt, add_special_tokens=False).ids + src_len = len(tok.encode(src, add_special_tokens=False).ids) + limit = min(max_new_tokens, int(1.6 * src_len) + 16) # zh output ≲ 1.6× en tokens + logits, _, cache = self.model(mx.array([ids], dtype=mx.int32)) + gen = [] + for _ in range(limit): + nxt = int(mx.argmax(logits[0, -1]).item()) + if nxt == end: + break + gen.append(nxt) + logits, _, cache = self.model(mx.array([[nxt]], dtype=mx.int32), cache=cache) + out.append(tok.decode(gen, skip_special_tokens=True)) + self.model.train() + return out + + def translation_eval(self, n: int) -> dict: + import sacrebleu + from chronos.data.corpus.evalsets import load_eval_sets + raw = self.spec.extra.get("raw_root") + if not raw: + return {} + pairs = load_eval_sets(raw, ["flores_devtest"]).get("flores_devtest", [])[:n] + if not pairs: + return {} + t0 = time.monotonic() + hyps = self.translate([en for en, _zh in pairs]) + refs = [zh for _en, zh in pairs] + bleu = sacrebleu.corpus_bleu(hyps, [refs], tokenize="zh") + chrf = sacrebleu.corpus_chrf(hyps, [refs]) + with open(os.path.join(self.spec.run_dir, "samples.jsonl"), "a", encoding="utf-8") as fh: + for (en, zh), hyp in list(zip(pairs, hyps))[:5]: + fh.write(json.dumps({"time": time.strftime("%Y-%m-%dT%H:%M:%S"), "phase": self.spec.phases[ + min(self.state["phase_index"], len(self.spec.phases) - 1)].name, "step": self.state["phase_step"], + "direction": "en_zh", "set": "flores_devtest", "src": en, "ref": zh, "hyp": hyp}, + ensure_ascii=False) + "\n") + return {"flores_devtest_n": len(pairs), "bleu_zh": round(bleu.score, 2), "chrf": round(chrf.score, 2), + "eval_s": round(time.monotonic() - t0, 1)} + + # ---- whole run ----------------------------------------------------------- + def run(self): + # Pin every lazily-used module now so source edits during a multi-day + # run cannot change the code that evaluates it. + import sacrebleu # noqa: F401 + import tokenizers # noqa: F401 + import chronos.data.corpus.evalsets # noqa: F401 + import chronos.data.corpus.shards # noqa: F401 + import chronos.mlx.training.inspect_checkpoint # noqa: F401 + try: + return self._run() + except KeyboardInterrupt: + self._beat("stopped", error="interrupted") + raise + except BaseException as exc: + self._beat("failed", error=f"{type(exc).__name__}: {exc}"[:800]) + raise + + def _run(self): + resumed = self.try_resume() + if resumed: + self._emit({"phase": self.spec.phases[min(self.state["phase_index"], len(self.spec.phases) - 1)].name, + "event": "resume", "step": self.state["phase_step"], + "wall_h": round(self.state.get("wall_s", 0.0) / 3600, 3), "code": self.runtime["code"]["sha256"]}) + budget_s = self.spec.budget_hours * 3600 * (1 - self.spec.eval_reserve_fraction) + # Wall time already spent (including eval/checkpoint I/O) before a resume. + start = time.monotonic() - (self.state.get("wall_s", 0.0) if resumed else 0.0) + self._run_start = start + phase_end = start + for pi, phase in enumerate(self.spec.phases): + phase_end += budget_s * phase.budget_fraction + if pi < self.state["phase_index"]: + continue + self.log(f"[run] phase {phase.name}: deadline in {(phase_end - time.monotonic()) / 3600:.2f}h") + self.run_phase(pi, phase_end) + self.state.update(phase_index=pi + 1, phase_step=0, phase_total=None, + phase_budget_s=None, phase_wall_s=0.0, lr_scale=1.0) + self.save_state() + self._beat("evaluating") + final = {"event": "final", "tokens_seen": self.state["tokens_seen"], + "hours": round((time.monotonic() - start) / 3600, 2)} + if self.spec.final_eval_samples: + final.update(self.translation_eval(self.spec.final_eval_samples)) + self._emit(final) + self._beat("finished", final=final) + self.log(f"[run] done: {json.dumps(final, ensure_ascii=False)}") + return final + + +def main(argv=None) -> int: + import argparse + parser = argparse.ArgumentParser(description="Time-budgeted Chronos MLX training run") + parser.add_argument("--spec", required=True, help="RunSpec JSON (see chronos/mlx/training/recipes)") + parser.add_argument("--budget-hours", type=float, default=None) + args = parser.parse_args(argv) + raw = json.load(open(args.spec, encoding="utf-8")) + raw["phases"] = [PhaseSpec(**p) for p in raw["phases"]] + spec = RunSpec(**raw) + if args.budget_hours is not None: + spec.budget_hours = args.budget_hours + LongRun(spec, log=lambda m: print(m, flush=True)).run() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/mlx/training/optim.py b/chronos/mlx/training/optim.py new file mode 100644 index 0000000..4aadde9 --- /dev/null +++ b/chronos/mlx/training/optim.py @@ -0,0 +1,45 @@ +"""AdamW with FP32 accumulation/state and low-precision forward weights.""" +import mlx.core as mx +import mlx.optimizers as optim +from mlx.utils import tree_map_with_path + + +class MixedPrecisionAdamW(optim.Adam): + """Match the Torch AdamW recipe without rounding away small BF16 updates. + + The FP32 master copy is required only for BF16/FP16 parameters. Its memory + cost is reported separately by the UI estimator. Optimizer snapshots include + masters and moments, so restoring a snapshot restores the complete update. + """ + + def __init__(self, learning_rate, weight_decay=.01, betas=(.9, .95), eps=1e-8): + super().__init__(learning_rate, betas=betas, eps=eps, bias_correction=True) + self.weight_decay = float(weight_decay) + + def init_single(self, parameter, state): + state["m"] = mx.zeros(parameter.shape, dtype=mx.float32) + state["v"] = mx.zeros(parameter.shape, dtype=mx.float32) + if parameter.dtype in (mx.bfloat16, mx.float16): + state["master"] = parameter.astype(mx.float32) + + def init(self, parameters): + super().init(parameters) + def assign_decay(path, parameter, state): + # Same exclusion as the Torch recipe: embeddings, norms and biases. + state["decay"] = parameter.ndim >= 2 and not path.endswith("embed_tokens.weight") + tree_map_with_path(assign_decay, parameters, self.state) + + def apply_gradients(self, gradients, parameters): + # MLX's default initializes from gradients; masters must start at WEIGHTS. + if not self._initialized: + self.init(parameters.trainable_parameters() if hasattr(parameters, "trainable_parameters") else parameters) + return super().apply_gradients(gradients, parameters) + + def apply_single(self, gradient, parameter, state): + master = state.get("master", parameter.astype(mx.float32)) + if state["decay"]: + master = master * (1 - self.learning_rate * self.weight_decay) + updated = super().apply_single(gradient.astype(mx.float32), master, state) + if "master" in state: + state["master"] = updated + return updated.astype(parameter.dtype) diff --git a/chronos/mlx/training/recipes/enzh_50m_72h.json b/chronos/mlx/training/recipes/enzh_50m_72h.json new file mode 100644 index 0000000..0d8b082 --- /dev/null +++ b/chronos/mlx/training/recipes/enzh_50m_72h.json @@ -0,0 +1,89 @@ +{ + "run_dir": "runs/enzh_50m_v1", + "shards_dir": "Translation Training Data/shards_v1", + "tokenizer_dir": "Translation Training Data/tokenizer_v1", + "model_config": { + "architecture_version": 3, + "num_hidden_layers": 6, + "prelude_layers": 1, + "coda_layers": 1, + "recurrent_depth": 4, + "min_recurrent_depth": 2, + "max_inference_depth": 4, + "recurrent_attention_mode": "shared_first", + "adaptive_halting": true, + "depth_schedule": "fixed_then_adaptive", + "depth_schedule_warmup_fraction": 0.3, + "halting_threshold": 0.95, + "num_experts_per_tok": 2, + "num_shared_experts": 1, + "kv_latent_dim": 128, + "rope_dim": 32, + "sliding_window_size": 128, + "max_position_embeddings": 1024, + "layer_aware_lookahead": true, + "layer_aware_cache": true, + "lookahead_steps": 2, + "dropout": 0.0, + "fallback_mask_prob": 0.08, + "rope_theta": 1000000.0, + "bos_token_id": 1, + "eos_token_id": 2, + "tie_word_embeddings": true, + "hidden_size": 384, + "num_attention_heads": 6, + "num_key_value_heads": 2, + "num_experts": 36, + "moe_intermediate_size": 160, + "vocab_size": 16384, + "ponder_loss_coef": 0.003, + "lambda_balance": 0.005 + }, + "phases": [ + { + "name": "pretrain", + "budget_fraction": 0.5, + "peak_lr": 0.0015, + "tokens_per_step": 131072, + "micro_tokens": 8192, + "seq_len": 512, + "warmup_fraction": 0.02, + "min_lr_ratio": 0.1, + "weight_decay": 0.1, + "grad_clip": 1.0 + }, + { + "name": "sft", + "budget_fraction": 0.5, + "peak_lr": 0.0005, + "tokens_per_step": 65536, + "micro_tokens": 8192, + "warmup_fraction": 0.01, + "min_lr_ratio": 0.05, + "weight_decay": 0.1, + "grad_clip": 1.0 + } + ], + "budget_hours": 72.0, + "eval_reserve_fraction": 0.04, + "log_every_steps": 20, + "validate_every_s": 7200.0, + "checkpoint_every_s": 1800.0, + "bleu_samples": 100, + "seed": 1234, + "dtype": "bfloat16", + "extra": { + "raw_root": "Translation Training Data/raw", + "params": { + "total": 49691717, + "active": 12090437, + "active_union_upper": 16514117, + "embedding": 6291456, + "per_expert": 184320, + "virtual_layers": 18 + }, + "note": "fine-grained MoE E=36 top-2 + 1 shared; recurrent depth 4 shared_first; min depth 2, ponder 3e-3, balance 5e-3 (smoke: depth collapse to 1.5, load_cv 2.2)" + }, + "schedule": "time", + "final_eval_samples": 1012 +} diff --git a/chronos/mlx/training/telemetry.py b/chronos/mlx/training/telemetry.py new file mode 100644 index 0000000..673bd8e --- /dev/null +++ b/chronos/mlx/training/telemetry.py @@ -0,0 +1,236 @@ +"""Training telemetry for MLX runs: heartbeat, per-tensor statistics, routing. + +Design rules (the run is the product, telemetry must not cost it): +* nothing here changes training numerics — it only reads parameters, + accumulated gradients and router probabilities that already exist; +* every statistic for one telemetry step is reduced on device and brought to + the host with **one** ``mx.eval`` + ``tolist`` (no per-tensor ``.item()``); +* FP32 casts appear only inside these diagnostic reductions (sums of squares + of BF16 values must not accumulate in BF16); they run every + ``tensor_stats_every_s`` seconds, never in the training step itself. +""" +from __future__ import annotations + +import hashlib +import json +import math +import os +import platform +import socket +import sys +import time + +import mlx.core as mx +from mlx.utils import tree_flatten + +# Name fragment -> component group (first match wins). +GROUP_RULES = ( + (".mlp.shared_experts.", "shared_expert"), + (".mlp.experts.", "experts"), + (".mlp.gate.", "router"), + (".self_attn.", "attention"), + ("lookahead_router", "lookahead"), + ("recurrent_controller", "recurrent"), + ("embed_tokens", "embedding"), + ("lm_head", "head"), + ("norm", "norm"), +) + + +def component(name: str) -> str: + for frag, group in GROUP_RULES: + if frag in name: + return group + return "other" + + +def _sig(v: float, digits: int = 4) -> float: + return float(f"{v:.{digits}g}") + + +def _atomic_json(path: str, payload: dict) -> None: + tmp = f"{path}.{os.getpid()}.tmp" + with open(tmp, "w", encoding="utf-8") as fh: + json.dump(payload, fh, ensure_ascii=False, indent=1, default=str) + os.replace(tmp, path) + + +def tensor_stats(params: dict, grads: dict | None = None, before: dict | None = None, + grad_scale: float = 1.0) -> dict: + """Per-tensor RMS / mean / |max|, gradient RMS and relative update size. + + ``params``/``grads``/``before`` are flat ``{name: array}`` dicts (``grads`` + is the FP32 accumulation *before* clipping; ``grad_scale`` turns the sum + over micro-batches into the mean). ``before`` holds the parameter arrays + from before the optimizer update (MLX arrays are immutable, so keeping the + references costs no copy). + """ + names = list(params) + zero = mx.array(0.0, dtype=mx.float32) + rows = [] + for name in names: + p = params[name].astype(mx.float32) + cells = [p.sum(), (p * p).sum(), mx.abs(p).max()] + g = grads.get(name) if grads else None + cells.append(mx.square(g.astype(mx.float32)).sum() if g is not None else zero) + b = before.get(name) if before else None + cells.append(mx.square(p - b.astype(mx.float32)).sum() if b is not None else zero) + rows.append(mx.stack(cells)) + table = mx.stack(rows) if rows else mx.zeros((0, 5)) + mx.eval(table) + values = table.tolist() + + tensors, groups = [], {} + tot = {"numel": 0, "sumsq": 0.0, "grad_sumsq": 0.0, "delta_sumsq": 0.0} + for name, (s, sq, amax, gsq, dsq) in zip(names, values): + arr = params[name] + n = max(1, int(arr.size)) + mean = s / n + var = max(0.0, sq / n - mean * mean) + gsq *= grad_scale * grad_scale + entry = {"name": name, "group": component(name), "shape": list(arr.shape), "dtype": str(arr.dtype).split(".")[-1], + "rms": _sig((sq / n) ** 0.5), "mean": _sig(mean), "std": _sig(var ** 0.5), "absmax": _sig(amax)} + if grads: + entry["grad_rms"] = _sig((gsq / n) ** 0.5) + if before: + entry["update_ratio"] = _sig((dsq ** 0.5) / max(sq ** 0.5, 1e-30)) + tensors.append(entry) + acc = groups.setdefault(entry["group"], {"tensors": 0, "numel": 0, "sumsq": 0.0, "grad_sumsq": 0.0, "delta_sumsq": 0.0}) + for bucket in (acc, tot): + bucket["numel"] += n + bucket["sumsq"] += sq + bucket["grad_sumsq"] += gsq + bucket["delta_sumsq"] += dsq + acc["tensors"] += 1 + + def summarize(b): + out = {"numel": b["numel"], "rms": _sig((b["sumsq"] / max(1, b["numel"])) ** 0.5)} + if grads: + out["grad_norm"] = _sig(b["grad_sumsq"] ** 0.5) + if before: + out["update_ratio"] = _sig(b["delta_sumsq"] ** 0.5 / max(b["sumsq"] ** 0.5, 1e-30)) + if "tensors" in b: + out["tensors"] = b["tensors"] + return out + + return {"tensors": tensors, "groups": {k: summarize(v) for k, v in sorted(groups.items())}, + "stats": summarize(tot), "grad_note": "pre-clip, mean over micro-batches" if grads else None} + + +def expert_table(tensors: list[dict]) -> dict: + """``[layer][expert]`` weight RMS and gradient RMS (gate/up/down pooled).""" + cells: dict[int, dict[int, list]] = {} + for t in tensors: + parts = t["name"].split(".") + if ".mlp.experts." not in t["name"]: + continue + try: + layer = int(parts[parts.index("layers") + 1]) + expert = int(parts[parts.index("experts") + 1]) + except (ValueError, IndexError): + continue + n = 1 + for d in t["shape"]: + n *= d + slot = cells.setdefault(layer, {}).setdefault(expert, [0, 0.0, 0.0]) + slot[0] += n + slot[1] += t["rms"] ** 2 * n + slot[2] += t.get("grad_rms", 0.0) ** 2 * n + layers = sorted(cells) + width = max((max(v) + 1 for v in cells.values()), default=0) + def grid(i): + return [[_sig((cells[l].get(e, [1, 0.0, 0.0])[i] / max(1, cells[l].get(e, [1])[0])) ** 0.5) for e in range(width)] + for l in layers] + return {"layers": layers, "weight_rms": grid(1), "grad_rms": grid(2)} if layers else {} + + +def routing_snapshot(model) -> dict | None: + """Per virtual layer: load share, mean router probability and entropy of + the last forward's dispatches (``[V, E]``), with ``L·`` + labels so recurrent iterations of one physical layer sit side by side.""" + from chronos.mlx.moe import ChronosMLXMOE + probs, labels = [], [] + for li, layer in enumerate(getattr(model, "layers", [])): + moe = getattr(layer, "mlp", None) + if isinstance(moe, ChronosMLXMOE): + for it, p in enumerate(moe.iteration_router_probs): + probs.append(p) + labels.append(f"L{li}·{it}") + if not probs: + return None + top_k = int(model.config.num_experts_per_tok) + router = mx.stack(probs, axis=2) # [B, S, V, E] + E, V = router.shape[-1], router.shape[2] + flat = router.reshape(-1, V, E).astype(mx.float32) + top = mx.argpartition(-flat, kth=top_k - 1, axis=-1)[..., :top_k] + load = (top[..., None] == mx.arange(E)).astype(mx.float32).sum(-2).mean(0) / top_k + importance = flat.mean(0) + entropy = -(flat * mx.log(mx.maximum(flat, 1e-9))).sum(-1).mean(0) + mx.eval(load, importance, entropy) + r4 = lambda grid: [[round(v, 4) for v in row] for row in grid] + return {"labels": labels, "experts": E, "top_k": top_k, "load": r4(load.tolist()), + "importance": r4(importance.tolist()), "entropy": [round(v, 4) for v in entropy.tolist()], + "entropy_max": round(math.log(E), 4), "source": "last micro-batch of the step"} + + +def code_fingerprint(root: str | None = None) -> dict: + """SHA-256 over the kernel and trainer sources this process imported, so + the monitor can tell which code revision a long run is executing.""" + root = root or os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + files = [] + for sub in ("chronos/mlx", "chronos/mlx/training", "chronos/model"): + folder = os.path.join(root, sub) + if os.path.isdir(folder): + files += sorted(os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".py")) + digest = hashlib.sha256() + for path in files: + with open(path, "rb") as fh: + digest.update(os.path.relpath(path, root).encode() + b"\0" + fh.read()) + return {"sha256": digest.hexdigest()[:16], "files": len(files)} + + +def runtime_info(model, spec, precision) -> dict: + import numpy as np + params = dict(tree_flatten(model.parameters())) + by_group: dict[str, int] = {} + for name, arr in params.items(): + by_group[component(name)] = by_group.get(component(name), 0) + int(arr.size) + try: + device = mx.device_info() + except Exception: # noqa: BLE001 - device info is optional metadata + device = {} + return { + "started": time.strftime("%Y-%m-%dT%H:%M:%S"), "pid": os.getpid(), "host": socket.gethostname(), + "python": sys.version.split()[0], "platform": platform.platform(), + "versions": {"mlx": getattr(mx, "__version__", "?"), "numpy": np.__version__}, + "device": {k: device.get(k) for k in ("device_name", "architecture", "memory_size", + "max_recommended_working_set_size") if k in device}, + "default_device": str(mx.default_device()), + "precision": precision.as_dict() if precision is not None else {"compute": spec.dtype}, + "parameters": {"total": sum(by_group.values()), "by_group": dict(sorted(by_group.items()))}, + "phases": [{"name": p.name, "accum": max(1, p.tokens_per_step // p.micro_tokens), + "micro_rows": max(1, p.micro_tokens // p.seq_len) if p.name == "pretrain" else None, + "micro_tokens": p.micro_tokens, "tokens_per_step": p.tokens_per_step} for p in spec.phases], + "code": code_fingerprint(), + } + + +class Heartbeat: + """``status.json``: the run's own view of itself, rewritten atomically.""" + + def __init__(self, run_dir: str): + self.path = os.path.join(run_dir, "status.json") + self.fields: dict = {"pid": os.getpid(), "host": socket.gethostname(), "state": "starting", + "started": time.time(), "error": None} + + def beat(self, state: str | None = None, **fields) -> dict: + if state is not None: + self.fields["state"] = state + self.fields.update(fields) + self.fields["updated"] = time.time() + self.fields["updated_iso"] = time.strftime("%Y-%m-%dT%H:%M:%S") + try: + _atomic_json(self.path, self.fields) + except OSError: + pass # telemetry must never stop training + return self.fields diff --git a/chronos/mlx/training/trainer.py b/chronos/mlx/training/trainer.py index cf2dc82..d6d2b38 100644 --- a/chronos/mlx/training/trainer.py +++ b/chronos/mlx/training/trainer.py @@ -25,6 +25,7 @@ from chronos.mlx.model import ChronosMLXModel from chronos.mlx.moe import ChronosMLXMOE from chronos.mlx.training.io import mlx_state_to_torch + from chronos.mlx.training.optim import MixedPrecisionAdamW except ModuleNotFoundError: mx = None nn = None @@ -43,10 +44,15 @@ def _ids_to_mx(x) -> mx.array: return mx.array(x).astype(mx.int32) +def _f32(x: mx.array) -> mx.array: + """FP32 view for loss math; a no-op when logits were born FP32.""" + return x if x.dtype == mx.float32 else x.astype(mx.float32) + + def _masked_lm_loss(logits: mx.array, labels: mx.array) -> mx.array: shift_logits = logits[:, :-1, :] shift_labels = labels[:, 1:] - flat_logits = shift_logits.reshape(-1, shift_logits.shape[-1]).astype(mx.float32) + flat_logits = _f32(shift_logits.reshape(-1, shift_logits.shape[-1])) flat_labels = shift_labels.reshape(-1) mask = flat_labels != -100 safe_labels = mx.where(mask, flat_labels, mx.zeros_like(flat_labels)) @@ -56,7 +62,7 @@ def _masked_lm_loss(logits: mx.array, labels: mx.array) -> mx.array: def _gather_token_logprobs(logits: mx.array, labels: mx.array) -> mx.array: labels = labels.astype(mx.int32) - logits = logits.astype(mx.float32) + logits = _f32(logits) log_probs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) flat = log_probs.reshape(-1, log_probs.shape[-1]) idx = labels.reshape(-1) @@ -84,14 +90,22 @@ def _shifted_token_logprobs(logits: mx.array, ids: mx.array) -> mx.array: def _collect_router_probs(model) -> mx.array | None: + """Every dispatch of the last forward: ``[B, S, V, E]`` (V = virtual layers).""" probs = [] for layer in getattr(model, "layers", []): moe = getattr(layer, "mlp", None) - if isinstance(moe, ChronosMLXMOE) and moe.last_router_probs is not None: - probs.append(moe.last_router_probs) + if isinstance(moe, ChronosMLXMOE): + probs.extend(moe.iteration_router_probs) if not probs: return None - return mx.stack(probs, axis=2) # [B, S, L, E] + return mx.stack(probs, axis=2) + + +def _first_iteration_router_probs(model) -> mx.array | None: + """Step-0 routing per physical layer ``[B, S, L, E]`` (lookahead teacher).""" + probs = [layer.mlp.iteration_router_probs[0] for layer in getattr(model, "layers", []) + if isinstance(getattr(layer, "mlp", None), ChronosMLXMOE) and layer.mlp.iteration_router_probs] + return mx.stack(probs, axis=2) if probs else None def _temporal_locality_loss(router_mean: mx.array) -> mx.array: @@ -127,11 +141,34 @@ def _load_balance_loss(router_4d: mx.array, num_experts_per_tok: int) -> mx.arra return total +def _flatten_lookahead_groups(prediction, teacher): + """Flatten matching groups into batch; never average distinct routing heads.""" + if prediction is None or teacher is None: + return prediction, teacher + if prediction.ndim == 5: + b, seq, steps, groups, experts = prediction.shape + if teacher.shape != (b, seq, groups, experts): + raise ValueError(f"MLX lookahead teacher {teacher.shape} does not match prediction {prediction.shape}") + return (prediction.transpose(0, 3, 1, 2, 4).reshape(b*groups, seq, steps, experts), + teacher.transpose(0, 2, 1, 3).reshape(b*groups, seq, experts)) + if prediction.ndim != 4 or teacher.ndim != 3: + raise ValueError("MLX lookahead expects 4-D/3-D or 5-D/4-D prediction/teacher") + return prediction, teacher + + +def _lookahead_teacher(router, prediction, config): + if prediction is None or prediction.ndim == 4: + return mx.stop_gradient(router.mean(axis=2)) + start = int(config.prelude_layers) + return mx.stop_gradient(router[:, :, start:start + prediction.shape[3], :]) + + def _lookahead_supervision_loss( lookahead_probs: mx.array | None, teacher_router_probs: mx.array | None, lookahead_steps: int, ) -> mx.array: + lookahead_probs, teacher_router_probs = _flatten_lookahead_groups(lookahead_probs, teacher_router_probs) if lookahead_probs is None: return mx.zeros(()) if teacher_router_probs is None: @@ -159,6 +196,7 @@ def _lookahead_topk_hit_loss( lookahead_steps: int, num_experts_per_tok: int, ) -> mx.array: + lookahead_probs, teacher_router_probs = _flatten_lookahead_groups(lookahead_probs, teacher_router_probs) if lookahead_probs is None: return mx.zeros(()) if teacher_router_probs is None: @@ -190,6 +228,7 @@ def _lookahead_union_loss( lookahead_steps: int, num_experts_per_tok: int, ) -> mx.array: + lookahead_probs, teacher_router_probs = _flatten_lookahead_groups(lookahead_probs, teacher_router_probs) if lookahead_probs is None: return mx.zeros(()) if teacher_router_probs is None: @@ -217,6 +256,7 @@ def _offload_training_metrics( lookahead_steps: int, num_experts_per_tok: int, ) -> dict[str, float]: + lookahead_probs, router_mean = _flatten_lookahead_groups(lookahead_probs, router_mean) if router_mean is None: return {} B, S, E = router_mean.shape @@ -286,71 +326,97 @@ def _router_kl_anchor(current: mx.array, reference: mx.array | None, lambda_anch return float(lambda_anchor) * kl -def _chronos_regularized_loss(model, base_loss: mx.array, lookahead_probs: mx.array | None, config, router_ref=None) -> mx.array: - setattr(model, "_last_lookahead_probs", lookahead_probs) - loss = base_loss.astype(mx.float32) - router_4d = _collect_router_probs(model) - if router_4d is None: - return loss - router_4d = router_4d.astype(mx.float32) - - lambda_balance = float(getattr(config, "lambda_balance", 0.0) or 0.0) - if lambda_balance > 0.0: - loss = loss + lambda_balance * _load_balance_loss( - router_4d, - int(getattr(config, "num_experts_per_tok", 1) or 1), - ) +def draft_exit_active(model, config) -> bool: + if float(getattr(config, "draft_exit_kl_coef", 0.0) or 0.0) <= 0.0 or not getattr(model, "training", False): + return False + if not getattr(config, "recurrent_enabled", True) or int(config.recurrent_depth) <= 1: + return False + return 1 <= int(getattr(config, "draft_exit_depth", 1)) < int(config.recurrent_depth) - if router_4d.shape[1] > 1: - router_mean = router_4d.mean(axis=2) - lambda_temporal = float(getattr(config, "lambda_temporal", 0.0) or 0.0) - if lambda_temporal > 0.0: - loss = loss + lambda_temporal * _temporal_locality_loss(router_mean) - lambda_locality = float(getattr(config, "lambda_router_locality", 0.0) or 0.0) - if lambda_locality > 0.0: - loss = loss + lambda_locality * _router_locality_loss( - router_mean, - int(getattr(config, "num_experts_per_tok", 1) or 1), - ) - lookahead_steps = int(getattr(config, "lookahead_steps", 0) or 0) - lambda_lookahead = float(getattr(config, "lambda_lookahead", 0.0) or 0.0) - if lookahead_probs is not None and lambda_lookahead > 0.0 and lookahead_steps > 0: - teacher = mx.stop_gradient(router_mean) - loss = loss + lambda_lookahead * _lookahead_supervision_loss( - lookahead_probs, - teacher, - lookahead_steps, - ) - lambda_topk = float(getattr(config, "lambda_lookahead_topk", 0.0) or 0.0) - if lambda_topk > 0.0: - loss = loss + lambda_topk * _lookahead_topk_hit_loss( - lookahead_probs, - teacher, - lookahead_steps, - int(getattr(config, "num_experts_per_tok", 1) or 1), - ) - lambda_union = float(getattr(config, "lambda_lookahead_union", 0.0) or 0.0) - if lambda_union > 0.0: - loss = loss + lambda_union * _lookahead_union_loss( - lookahead_probs, - teacher, - lookahead_steps, - int(getattr(config, "num_experts_per_tok", 1) or 1), - ) +def draft_exit_logits(model, input_ids, config): + """Draft-depth FP32 logits, or None when the KL hook is off. - loss = loss + _router_kl_anchor( - router_mean, - router_ref, - float(getattr(config, "lambda_router_anchor", 0.0) or 0.0), - ) + Run *before* the full-depth forward: the full pass then owns the routing + and recurrent diagnostics the other regularizers read. + """ + if not draft_exit_active(model, config): + return None + logits, _lookahead, _cache = model(input_ids, logits_dtype=mx.float32, + recurrent_depth=int(config.draft_exit_depth)) + return logits + + +def draft_exit_kl(draft_logits: mx.array, full_logits: mx.array, keep: mx.array | None) -> mx.array: + """Mean KL(full.detach() || draft) in FP32; ``keep`` [B, S] marks supervised logits.""" + teacher = mx.stop_gradient(full_logits.astype(mx.float32)) + t_log = teacher - mx.logsumexp(teacher, axis=-1, keepdims=True) + d32 = draft_logits.astype(mx.float32) + s_log = d32 - mx.logsumexp(d32, axis=-1, keepdims=True) + kl = (mx.exp(t_log) * (t_log - s_log)).sum(axis=-1) + if keep is None: + return kl.mean() + keep = keep.astype(mx.float32) + return (kl * keep).sum() / mx.maximum(keep.sum(), 1.0) + + +def _shifted_keep(labels: mx.array) -> mx.array: + """Logit t is supervised when label t+1 is not -100 (CE shift).""" + nxt = labels[:, 1:] != -100 + return mx.concatenate([nxt, mx.zeros((labels.shape[0], 1), dtype=mx.bool_)], axis=1) + + +def _chronos_regularized_loss(model, base_loss, lookahead_probs, config, router_ref=None, metrics=None, + draft_kl=None): + """Return total loss and retain the actual weighted components for logging.""" + object.__setattr__(model, "_last_lookahead_probs", lookahead_probs) + terms = {"primary_loss": base_loss.astype(mx.float32)} + def add(name, coefficient, value): + if coefficient: + terms[name] = float(coefficient) * value.astype(mx.float32) + + recurrent = getattr(model, "last_recurrent_state", None) or {} + for name, coefficient in (("ponder_loss", "ponder_loss_coef"), ("distill_loss", "recurrent_distill_coef"), ("stability_loss", "recurrent_stability_coef")): + if name in recurrent: + add("recurrent_" + name, getattr(config, coefficient, 0), recurrent[name]) + if "contraction" in recurrent: + add("recurrent_contraction_loss", config.recurrent_contraction_coef, + mx.maximum(recurrent["contraction"]-config.recurrent_contraction_budget, 0)**2) + + router = _collect_router_probs(model) + if router is not None: + 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 router.shape[1] > 1: + mean = router.mean(axis=2) + add("temporal_loss", config.lambda_temporal, _temporal_locality_loss(mean)) + add("locality_loss", config.lambda_router_locality, _router_locality_loss(mean, top_k)) + if lookahead_probs is not None and config.lambda_lookahead > 0 and config.lookahead_steps > 0: + teacher = _lookahead_teacher(_first_iteration_router_probs(model).astype(mx.float32), + lookahead_probs, config) + add("lookahead_loss", config.lambda_lookahead, + _lookahead_supervision_loss(lookahead_probs, teacher, config.lookahead_steps)) + add("lookahead_topk_loss", config.lambda_lookahead_topk, + _lookahead_topk_hit_loss(lookahead_probs, teacher, config.lookahead_steps, top_k)) + add("lookahead_union_loss", config.lambda_lookahead_union, + _lookahead_union_loss(lookahead_probs, teacher, config.lookahead_steps, top_k)) + add("anchor_loss", 1, _router_kl_anchor(mean, router_ref, config.lambda_router_anchor)) + if draft_kl is not None: + add("draft_exit_kl", getattr(config, "draft_exit_kl_coef", 0.0), draft_kl) + loss = sum(terms.values()) + object.__setattr__(model, "_last_loss_components", terms) + object.__setattr__(model, "_last_stage_metrics", metrics or {"primary": base_loss}) return loss def _ce_stage_loss(model, input_ids, labels): - logits, lookahead, _cache = model(input_ids) + draft = draft_exit_logits(model, input_ids, model.config) + logits, lookahead, _cache = model(input_ids, logits_dtype=mx.float32) base = _masked_lm_loss(logits, labels) - return _chronos_regularized_loss(model, base, lookahead, model.config) + draft_kl = None if draft is None else draft_exit_kl(draft, logits, _shifted_keep(labels)) + return _chronos_regularized_loss(model, base, lookahead, model.config, metrics={"ce": base, "primary": base}, + draft_kl=draft_kl) def _dpo_stage_loss(model, reference, batch, beta: float = 0.1, router_ref=None): @@ -361,12 +427,12 @@ def _dpo_stage_loss(model, reference, batch, beta: float = 0.1, router_ref=None) m_c = _ids_to_mx(batch["mask_chosen"]) m_r = _ids_to_mx(batch["mask_rejected"]) x = mx.concatenate([x_c, x_r], axis=0) - logits, lookahead, _cache = model(x) + logits, lookahead, _cache = model(x, logits_dtype=mx.float32) B = x_c.shape[0] lp_c = _aligned_sequence_logprob(logits[:B], y_c, m_c) lp_r = _aligned_sequence_logprob(logits[B:], y_r, m_r) if reference is not None: - ref_logits, _ref_lp, _ref_cache = reference(x) + ref_logits, _ref_lp, _ref_cache = reference(x, logits_dtype=mx.float32) ref_c = mx.stop_gradient(_aligned_sequence_logprob(ref_logits[:B], y_c, m_c)) ref_r = mx.stop_gradient(_aligned_sequence_logprob(ref_logits[B:], y_r, m_r)) base = -mx.mean(nn.log_sigmoid(beta * ((lp_c - lp_r) - (ref_c - ref_r)))) @@ -383,7 +449,7 @@ def _orpo_stage_loss(model, batch, beta: float = 0.1, lambda_or: float = 0.1, ro m_c = _ids_to_mx(batch["mask_chosen"]) m_r = _ids_to_mx(batch["mask_rejected"]) x = mx.concatenate([x_c, x_r], axis=0) - logits, lookahead, _cache = model(x) + logits, lookahead, _cache = model(x, logits_dtype=mx.float32) B = x_c.shape[0] lp_c = _aligned_sequence_logprob(logits[:B], y_c, m_c, average=True) lp_r = _aligned_sequence_logprob(logits[B:], y_r, m_r, average=True) @@ -393,8 +459,9 @@ def _orpo_stage_loss(model, batch, beta: float = 0.1, lambda_or: float = 0.1, ro odds = (mx.log(p_c) - mx.log1p(-p_c)) - ( mx.log(p_r) - mx.log1p(-p_r) ) - base = nll - lambda_or * mx.mean(nn.log_sigmoid(beta * odds)) - return _chronos_regularized_loss(model, base, lookahead, model.config, router_ref=router_ref) + or_term = -lambda_or * mx.mean(nn.log_sigmoid(beta * odds)) + base = nll + or_term + return _chronos_regularized_loss(model, base, lookahead, model.config, router_ref=router_ref, metrics={"primary": nll, "secondary": or_term}) def _distill_stage_loss( @@ -407,14 +474,14 @@ def _distill_stage_loss( temperature: float = 4.0, router_ref=None, ): - logits, lookahead, _cache = model(input_ids) - t_logits, _tlp, _tcache = teacher(teacher_input_ids if teacher_input_ids is not None else input_ids) + logits, lookahead, _cache = model(input_ids, logits_dtype=mx.float32) + t_logits, _tlp, _tcache = teacher(teacher_input_ids if teacher_input_ids is not None else input_ids, logits_dtype=mx.float32) V = min(logits.shape[-1], t_logits.shape[-1]) S = min(logits.shape[1], t_logits.shape[1]) logits_kd = logits[:, :S, :V] t_logits = mx.stop_gradient(t_logits[:, :S, :V]) - logits_kd = logits_kd.astype(mx.float32) - t_logits = t_logits.astype(mx.float32) + logits_kd = _f32(logits_kd) + t_logits = _f32(t_logits) s_logp = logits_kd / temperature - mx.logsumexp(logits_kd / temperature, axis=-1, keepdims=True) t_prob = mx.softmax(t_logits / temperature, axis=-1) kd_tok = mx.sum(t_prob * (mx.log(mx.maximum(t_prob, 1e-12)) - s_logp), axis=-1) @@ -422,7 +489,7 @@ def _distill_stage_loss( kd = (temperature * temperature) * (kd_tok * mask).sum() / mx.maximum(mask.sum(), 1.0) ce = _masked_lm_loss(logits, labels) base = alpha * kd + (1.0 - alpha) * ce - return _chronos_regularized_loss(model, base, lookahead, model.config, router_ref=router_ref) + return _chronos_regularized_loss(model, base, lookahead, model.config, router_ref=router_ref, metrics={"primary": kd, "secondary": ce, "ce": ce}) def _fit_teacher_input(ids: mx.array, teacher_vocab_size: int) -> mx.array: @@ -433,7 +500,7 @@ def _grpo_stage_loss(model, reference, batch, beta: float = 0.04, router_ref=Non ids = _ids_to_mx(batch["ids"]) response_mask = mx.array(batch["response_mask"]).astype(mx.float32) advantages = mx.array(batch["advantages"]).astype(mx.float32) - logits, lookahead, _cache = model(ids) + logits, lookahead, _cache = model(ids, logits_dtype=mx.float32) lp_tok = _shifted_token_logprobs(logits, ids) mask = response_mask[:, 1:] seq_lp = (lp_tok * mask).sum(axis=1) / mx.maximum(mask.sum(axis=1), 1.0) @@ -441,12 +508,12 @@ def _grpo_stage_loss(model, reference, batch, beta: float = 0.04, router_ref=Non kl = mx.zeros((), dtype=logits.dtype) if reference is not None: - ref_logits, _ref_lp, _ref_cache = reference(ids) + ref_logits, _ref_lp, _ref_cache = reference(ids, logits_dtype=mx.float32) ref_lp_tok = mx.stop_gradient(_shifted_token_logprobs(ref_logits, ids)) log_ratio = ref_lp_tok - lp_tok kl = ((mx.exp(log_ratio) - log_ratio - 1.0) * mask).sum() / mx.maximum(mask.sum(), 1.0) base = pg_loss + float(beta) * kl - return _chronos_regularized_loss(model, base, lookahead, model.config, router_ref=router_ref) + return _chronos_regularized_loss(model, base, lookahead, model.config, router_ref=router_ref, metrics={"primary": pg_loss, "secondary": kl}) def _capture_router_ref(model, batch, stage: str) -> mx.array | None: @@ -550,6 +617,7 @@ class MLXStageResult: stopped: bool = False checkpoint_saved: bool = False rollbacks: int = 0 + failure_reason: str = "" def _is_mx_array(value) -> bool: @@ -564,9 +632,9 @@ def _tree_copy(tree): if isinstance(tree, tuple): return tuple(_tree_copy(v) for v in tree) if _is_mx_array(tree): - copied = mx.array(tree) - mx.eval(copied) - return copied + # MLX arrays are immutable; copy the tree without synchronizing once + # per tensor. Snapshots retain already evaluated arrays until rollback. + return mx.array(tree) return tree @@ -583,13 +651,23 @@ def _tree_arrays(tree) -> list[mx.array]: return arrays +def _tree_finite_check(tree) -> mx.array: + """Lazy device boolean: every floating array in ``tree`` is finite. + + Evaluate it together with the arrays it checks (one ``mx.eval`` for + values + guard) instead of a separate sync per check. + """ + checks = [mx.all(mx.isfinite(arr)) for arr in _tree_arrays(tree) + if mx.issubdtype(arr.dtype, mx.inexact)] + if not checks: + return mx.array(True) + return mx.all(mx.stack(checks)) if len(checks) > 1 else checks[0] + + def _tree_all_finite(tree) -> bool: - arrays = _tree_arrays(tree) - if not arrays: - return True - checks = [mx.all(mx.isfinite(arr)) for arr in arrays] - mx.eval(*checks) - return all(bool(check.item()) for check in checks) + check = _tree_finite_check(tree) + mx.eval(check) + return bool(check.item()) def _scalar_is_finite(value: mx.array) -> bool: @@ -600,8 +678,8 @@ def _scalar_is_finite(value: mx.array) -> bool: def _add_grads(a, b): if a is None: - return _tree_copy(b) - return tree_map(lambda x, y: x + y, a, b) + return tree_map(lambda g: g.astype(mx.float32), b) + return tree_map(lambda x, y: x + y.astype(mx.float32), a, b) def _scale_grads(grads, scale: float): @@ -641,22 +719,32 @@ def _mlx_dtype_from_name(dtype_name: str | None): def _load_mlx_from_checkpoint(config, checkpoint_path: str | None, dtype_name: str | None = None): + """Build the MLX model directly in the policy compute dtype. + + ``chronos.precision.resolve_mlx`` decides the dtype once (GPU: requested + or BF16; CPU: FP32 because CPU GatherMM is FP32-only). Each tensor is + converted once from the checkpoint dtype; no FP32 staging copy. + """ + from chronos.precision import resolve_mlx + policy = resolve_mlx(dtype_name or "auto") pt = ChronosForCausalLM(config) if checkpoint_path and os.path.exists(checkpoint_path): state = load_checkpoint_state_dict(checkpoint_path, map_location="cpu") load_state_dict_controlled(pt, state) - model = ChronosMLXModel.from_chronos_pytorch(pt, config) - dtype = _mlx_dtype_from_name(dtype_name) - model.set_dtype(dtype) + model = ChronosMLXModel.from_chronos_pytorch(pt, config, dtype=policy.compute) mx.eval(model.parameters()) model.config = config + model._chronos_precision = policy return model def _save_mlx_checkpoint(model, config, out_path: str, stage: str) -> None: + """Save in the model's compute dtype (BF16 stays BF16, bit-exact).""" pt = ChronosForCausalLM(config) - pt.load_state_dict(mlx_state_to_torch(model, config), strict=True) - save_state_dict_with_config(pt, out_path, config, stage=stage) + state = mlx_state_to_torch(model, config) + dtype = next(iter(state.values())).dtype + pt.to(dtype).load_state_dict(state, strict=True) + save_state_dict_with_config(pt, out_path, config, stage=stage, half=False) def run_mlx_stage( @@ -666,6 +754,7 @@ def run_mlx_stage( checkpoint_path: str | None, save_dir: str, loader=None, + validation_loader=None, prompts=None, teacher_path: str | None = None, args=None, @@ -691,13 +780,20 @@ def run_mlx_stage( teacher = _load_mlx_from_checkpoint(teacher_config, teacher_path, dtype_name=dtype_name) base_lr = float(getattr(args, "learning_rate", 5e-5) if args is not None else 5e-5) - default_eps = 1e-6 if dtype_name in {"float16", "bfloat16"} else 1e-8 - optimizer = optim.AdamW( + default_eps = 1e-8 + optimizer = MixedPrecisionAdamW( learning_rate=base_lr, weight_decay=float(getattr(args, "weight_decay", 0.01) if args is not None else 0.01), eps=float(getattr(args, "adam_eps", default_eps) if args is not None else default_eps), ) + optimizer.init(model.trainable_parameters()) + model.train() + for frozen in (reference, teacher): + if frozen is not None: + frozen.eval() max_steps = getattr(args, "steps", None) if args is not None else None + if max_steps is not None and int(max_steps) <= 0: + max_steps = None epochs = int(getattr(args, "epochs", 1) if args is not None else 1) accum = max(1, int(getattr(args, "accumulation_steps", 1) if args is not None else 1)) grad_clip = float(getattr(args, "grad_clip", 0.0) if args is not None else 0.0) @@ -710,6 +806,11 @@ def run_mlx_stage( log_interval = max(1, int(getattr(args, "log_interval", 10) if args is not None else 10)) steps = 0 last_loss = 0.0 + best_loss = float("inf") + failure_reason = "" + grad_norm = 0.0 + token_count = 0 + last_log_tokens = 0 t0 = time.monotonic() last_log_t = t0 last_log_step = 0 @@ -731,6 +832,8 @@ def _should_stop() -> bool: if data_iter is None: raise ValueError(f"MLX {stage} requires loader/prompts") total_steps = _planned_total_steps(data_iter, epochs, max_steps) + total_updates = max(1, (total_steps + accum - 1) // accum) + optimizer.learning_rate = get_lr(1, total_updates, base_lr) os.makedirs(save_dir, exist_ok=True) out_path = os.path.join( save_dir, @@ -745,6 +848,8 @@ def _should_stop() -> bool: pending_count = 0 rollback_model_state = _tree_copy(model.parameters()) rollback_optimizer_state = _tree_copy(optimizer.state) + rollback_loss = 0.0 + best_validation = float("inf") def _emit(event: dict) -> None: if progress_callback is not None: @@ -756,7 +861,9 @@ def _discard_pending() -> None: pending_count = 0 def _restore_rollback_state() -> None: + nonlocal last_loss _discard_pending() + last_loss = rollback_loss model.update(_tree_copy(rollback_model_state)) optimizer.state = _tree_copy(rollback_optimizer_state) mx.eval(model.parameters(), optimizer.state) @@ -816,8 +923,82 @@ def loss_fn(m, batch): "lr": base_lr, "accumulation_steps": accum, "save_interval": save_interval, + "optimizer": "AdamW FP32 masters/moments, betas=(0.9,0.95), bias correction", + "optimizer_updates": total_updates, }) + def rollback(reason, epoch): + nonlocal rollback_count, consecutive_rollbacks, lr_scale, failure_reason + rollback_count += 1 + consecutive_rollbacks += 1 + lr_scale *= .5 + _restore_rollback_state() + if consecutive_rollbacks >= rollback_limit: + failure_reason = reason + _emit({"event": "rollback", "stage": stage, "epoch": epoch, "epochs": epochs, + "step": steps, "steps": total_steps, "loss": last_loss, + "elapsed_s": time.monotonic()-t0, "dtype": dtype_name, + "lr": base_lr*lr_scale, "rollbacks": rollback_count, "reason": reason}) + + def apply_pending(epoch): + nonlocal rollback_model_state, rollback_optimizer_state, rollback_loss + nonlocal consecutive_rollbacks, grad_norm + if pending_grads is None: + return True + update_grads = _scale_grads(pending_grads, 1.0 / max(1, pending_count)) + # Accumulation and this norm are FP32 even when forward weights are BF16. + update_grads, norm = optim.clip_grad_norm(update_grads, grad_clip if grad_clip > 0 else float("inf")) + grads_ok = _tree_finite_check((norm, update_grads)) + optimizer.learning_rate = get_lr(int(optimizer.step.item()) + 1, total_updates, base_lr * lr_scale) + previous_snapshot = (rollback_model_state, rollback_optimizer_state, rollback_loss) + rollback_model_state = _tree_copy(model.parameters()) + rollback_optimizer_state = _tree_copy(optimizer.state) + rollback_loss = last_loss + optimizer.update(model, update_grads) + # One device sync for clip + guard + update + post-update guard. A bad + # gradient only wastes this update's compute: both guards are read + # before the result is kept, and the rollback targets are unchanged. + state_ok = _tree_finite_check((model.parameters(), optimizer.state)) + mx.eval(model.parameters(), optimizer.state, norm, grads_ok, state_ok) + grad_norm = float(norm.item()) + if not bool(grads_ok.item()): + rollback_model_state, rollback_optimizer_state, rollback_loss = previous_snapshot + rollback("non_finite_accumulated_grad", epoch) + return False + if not bool(state_ok.item()): + rollback("non_finite_parameters_or_optimizer", epoch) + return False + consecutive_rollbacks = 0 + _discard_pending() + return True + + def evaluate_validation(epoch): + nonlocal best_validation + if validation_loader is None or stage not in {"pretrain", "sft"}: + return + total_nll, valid_tokens = 0.0, 0 + was_training = model.training + model.eval() + try: + for ids, labels in validation_loader: + labels = _ids_to_mx(labels) + logits, _, _ = model(_ids_to_mx(ids), logits_dtype=mx.float32) + ce = _masked_lm_loss(logits, labels) + count = mx.sum(labels[:, 1:] != -100) + mx.eval(ce, count) + total_nll += float(ce.item()) * int(count.item()) + valid_tokens += int(count.item()) + finally: + model.train(was_training) + if valid_tokens: + value = total_nll / valid_tokens + _emit({"event": "validation", "stage": stage, "epoch": epoch, + "step": steps, "steps": total_steps, "ce": value, + "tokens": valid_tokens, "elapsed_s": time.monotonic()-t0}) + if value < best_validation: + best_validation = value + _save_mlx_checkpoint(model, config, out_path.replace(".pth", ".best.pth"), checkpoint_stage) + stop_outer = False for epoch_idx in range(epochs): for batch in data_iter: @@ -836,29 +1017,19 @@ def loss_fn(m, batch): reward_fn, args, ) + model.recurrent_controller.set_training_progress(steps + 1, total_steps) loss, grads = grad_fn(model, train_batch) - mx.eval(loss, grads) + supervised = stage in {"pretrain", "sft", "distill"} + # Loss, gradients, their finite guard, the FP32 accumulation and the + # supervised-token count are materialised by one device sync. + finite_check = _tree_finite_check((loss, grads)) + candidate = _add_grads(pending_grads, grads) + n_supervised = mx.sum(_ids_to_mx(train_batch[1])[:, 1:] != -100) if supervised else None + mx.eval(loss, finite_check, candidate, *([n_supervised] if supervised else [])) steps += 1 - finite = _scalar_is_finite(loss) and _tree_all_finite(grads) + finite = bool(finite_check.item()) if not finite: - rollback_count += 1 - consecutive_rollbacks += 1 - lr_scale *= 0.5 - _discard_pending() - _emit({ - "event": "rollback", - "stage": stage, - "epoch": epoch_idx + 1, - "epochs": epochs, - "step": steps, - "steps": total_steps, - "loss": last_loss, - "elapsed_s": time.monotonic() - t0, - "dtype": dtype_name, - "lr": base_lr * lr_scale, - "rollbacks": rollback_count, - "reason": "non_finite_loss_or_grad", - }) + rollback("non_finite_loss_or_grad", epoch_idx + 1) if consecutive_rollbacks >= rollback_limit: stop_outer = True break @@ -866,72 +1037,26 @@ def loss_fn(m, batch): break continue - last_loss = float(loss.item()) - pending_grads = _add_grads(pending_grads, grads) + observed_loss = float(loss.item()) + # A catastrophe guard, not a convergence/quality criterion. Only + # nonnegative supervised objectives have a meaningful ratio here. + if stage in {"pretrain", "sft", "distill"} and observed_loss > max(100.0, 20.0*best_loss): + rollback("finite_loss_divergence", epoch_idx + 1) + failure_reason = "finite_loss_divergence" + stop_outer = True + break + last_loss = observed_loss + best_loss = min(best_loss, observed_loss) + if supervised: + token_count += int(n_supervised.item()) + pending_grads = candidate # adopted only after the guard passed pending_count += 1 - should_update = pending_count >= accum - if max_steps is not None and steps >= int(max_steps): - should_update = True - if should_update and pending_grads is not None: - current_lr = get_lr(steps, total_steps or steps, base_lr * lr_scale) - optimizer.learning_rate = current_lr - update_grads = _scale_grads(pending_grads, 1.0 / max(1, pending_count)) - if grad_clip > 0.0: - update_grads, _norm = optim.clip_grad_norm(update_grads, grad_clip) - mx.eval(update_grads) - if not _tree_all_finite(update_grads): - rollback_count += 1 - consecutive_rollbacks += 1 - lr_scale *= 0.5 - _discard_pending() - _emit({ - "event": "rollback", - "stage": stage, - "epoch": epoch_idx + 1, - "epochs": epochs, - "step": steps, - "steps": total_steps, - "loss": last_loss, - "elapsed_s": time.monotonic() - t0, - "dtype": dtype_name, - "lr": base_lr * lr_scale, - "rollbacks": rollback_count, - "reason": "non_finite_accumulated_grad", - }) - if consecutive_rollbacks >= rollback_limit: - stop_outer = True - break - else: - rollback_model_state = _tree_copy(model.parameters()) - rollback_optimizer_state = _tree_copy(optimizer.state) - optimizer.update(model, update_grads) - mx.eval(model.parameters(), optimizer.state) - if not _tree_all_finite(model.parameters()): - rollback_count += 1 - consecutive_rollbacks += 1 - lr_scale *= 0.5 - _restore_rollback_state() - _emit({ - "event": "rollback", - "stage": stage, - "epoch": epoch_idx + 1, - "epochs": epochs, - "step": steps, - "steps": total_steps, - "loss": last_loss, - "elapsed_s": time.monotonic() - t0, - "dtype": dtype_name, - "lr": base_lr * lr_scale, - "rollbacks": rollback_count, - "reason": "non_finite_parameters", - }) - if consecutive_rollbacks >= rollback_limit: - stop_outer = True - break - else: - consecutive_rollbacks = 0 - pending_grads = None - pending_count = 0 + should_update = pending_count >= accum or (max_steps is not None and steps >= int(max_steps)) + if should_update: + apply_pending(epoch_idx + 1) + if consecutive_rollbacks >= rollback_limit: + stop_outer = True + break if steps % save_interval == 0: _save_checkpoint("interval") @@ -939,17 +1064,31 @@ def loss_fn(m, batch): if steps == 1 or steps % log_interval == 0: now = time.monotonic() delta_steps = steps - last_log_step - tps = delta_steps / max(now - last_log_t, 1e-6) + interval_s = max(now - last_log_t, 1e-6) + tps = delta_steps / interval_s + tokens_per_s = (token_count-last_log_tokens) / interval_s + last_log_tokens = token_count last_log_t = now last_log_step = steps router_4d = _collect_router_probs(model) - router_mean = router_4d.mean(axis=2) if router_4d is not None else None + prediction = getattr(model, "_last_lookahead_probs", None) + router_mean = _lookahead_teacher(router_4d, prediction, config) if router_4d is not None else None offload_metrics = _offload_training_metrics( router_mean, getattr(model, "_last_lookahead_probs", None), int(getattr(config, "lookahead_steps", 0) or 0), int(getattr(config, "num_experts_per_tok", 1) or 1), ) + components = {k: float(v.item()) for k, v in getattr(model, "_last_loss_components", {}).items()} + stage_metrics = {k: float(v.item()) for k, v in getattr(model, "_last_stage_metrics", {}).items()} + lookahead_loss = sum(v for k, v in components.items() if k.startswith("lookahead_")) + recurrent_loss = sum(v for k, v in components.items() if k.startswith("recurrent_")) + defaults = { + "pretrain": (last_loss-components.get("primary_loss", last_loss), components.get("temporal_loss", 0)), + "sft": (lookahead_loss, components.get("anchor_loss", 0)), + "dpo": (components.get("anchor_loss", 0), lookahead_loss), + } + secondary, tertiary = defaults.get(stage, (0, components.get("anchor_loss", 0))) _emit({ "event": "step", "stage": stage, @@ -958,10 +1097,18 @@ def loss_fn(m, batch): "step": steps, "steps": total_steps, "loss": last_loss, - "ce": last_loss, - "aux": 0.0, - "temporal": 0.0, + **stage_metrics, + "secondary": stage_metrics.get("secondary", secondary), + "tertiary": tertiary, + "aux": last_loss-components.get("primary_loss", last_loss), + "temporal": components.get("temporal_loss", 0), + "loss_components": components, + "recurrent_loss": recurrent_loss, + "lookahead_loss": lookahead_loss, + "grad_norm": grad_norm, + "optimizer_step": int(optimizer.step.item()), "tps": tps, + "tokens_per_s": tokens_per_s, "elapsed_s": now - t0, "dtype": dtype_name, "lr": float(optimizer.learning_rate.item()) if "learning_rate" in optimizer.state else base_lr * lr_scale, @@ -970,27 +1117,22 @@ def loss_fn(m, batch): }) if max_steps is not None and steps >= int(max_steps): break + if not stop_outer and not _should_stop(): + evaluate_validation(epoch_idx+1) if stop_outer or _should_stop() or (max_steps is not None and steps >= int(max_steps)): break stopped = _should_stop() if pending_grads is not None and pending_count > 0 and not stop_outer: - current_lr = get_lr(steps, total_steps or steps, base_lr * lr_scale) - optimizer.learning_rate = current_lr - update_grads = _scale_grads(pending_grads, 1.0 / max(1, pending_count)) - if grad_clip > 0.0: - update_grads, _norm = optim.clip_grad_norm(update_grads, grad_clip) - mx.eval(update_grads) - if _tree_all_finite(update_grads): - rollback_model_state = _tree_copy(model.parameters()) - rollback_optimizer_state = _tree_copy(optimizer.state) - optimizer.update(model, update_grads) - mx.eval(model.parameters(), optimizer.state) - if not _tree_all_finite(model.parameters()): - _restore_rollback_state() - else: - _discard_pending() - _save_checkpoint("stop" if stopped else ("rollback_limit" if stop_outer else "final")) + if not apply_pending(epochs): + stop_outer = True + failure_reason = "final_update_failed" + elif not stopped: + evaluate_validation(epochs) + if stop_outer: + # Keep a failed run's recovery artifact distinct from a successful model. + out_path = out_path.replace(".pth", ".recovery.pth") + _save_checkpoint("stop" if stopped else (failure_reason if stop_outer else "final")) elapsed = time.monotonic() - t0 _emit({ "event": "stopped" if stopped else ("failed" if stop_outer else "complete"), @@ -1003,6 +1145,7 @@ def loss_fn(m, batch): "checkpoint_path": out_path if checkpoint_saved else "", "checkpoint_saved": checkpoint_saved, "rollbacks": rollback_count, + "reason": failure_reason, }) if stopped: print( @@ -1029,4 +1172,5 @@ def loss_fn(m, batch): stopped=stopped or stop_outer, checkpoint_saved=checkpoint_saved, rollbacks=rollback_count, + failure_reason=failure_reason, ) diff --git a/chronos/model/__init__.py b/chronos/model/__init__.py index f27d51a..d577794 100644 --- a/chronos/model/__init__.py +++ b/chronos/model/__init__.py @@ -1,6 +1,8 @@ from .config import ChronosConfig from .lookahead_router import LookaheadRouter from .moe_chronos import ChronosMOEFeedForward -from .model_chronos import ChronosForCausalLM +from .model_chronos import ChronosForCausalLM, ChronosModel, RecurrentChronosModel +from .outputs import ChronosCausalLMOutput +from .recurrent import RecurrentController, RecurrentDepthState from .temporal_loss import temporal_locality_loss, total_loss from .hybrid_attention import MLAAttention, SlidingWindowAttention, make_attention diff --git a/chronos/model/checkpoint.py b/chronos/model/checkpoint.py index 7a20971..a2717e6 100644 --- a/chronos/model/checkpoint.py +++ b/chronos/model/checkpoint.py @@ -20,52 +20,11 @@ from .config import ChronosConfig -CHECKPOINT_CONFIG_VERSION = 1 +CHECKPOINT_CONFIG_VERSION = 2 -CHECKPOINT_CONFIG_FIELDS = ( - "hidden_size", - "num_hidden_layers", - "num_experts", - "num_experts_per_tok", - "num_shared_experts", - "intermediate_size", - "moe_intermediate_size", - "num_attention_heads", - "num_key_value_heads", - "head_dim", - "rope_dim", - "kv_latent_dim", - "lookahead_steps", - "vocab_size", - "max_position_embeddings", - "sliding_window_size", - "use_hybrid_attention", - "use_moe", - "tie_word_embeddings", - "norm_topk_prob", - "router_aux_loss_coef", - "lambda_balance", - "lambda_temporal", - "lambda_lookahead", - "lambda_lookahead_topk", - "lambda_lookahead_union", - "lambda_router_locality", - "lambda_router_anchor", - "fallback_mask_prob", - "offload_miss_policy_default", - "recommended_resident_experts", - "offload_working_set_experts", - "vram_budget_gb", - "pinned_memory_max_fraction", - "storage_format", - "cluster_manifest_path", - "dropout", - "rms_norm_eps", - "rope_theta", - "flash_attn", - "bos_token_id", - "eos_token_id", -) +from chronos.config_schema import MODEL_FIELDS + +CHECKPOINT_CONFIG_FIELDS = MODEL_FIELDS + ("extra_metadata",) VALID_CHRONOS_CONFIG_KEYS = set(CHECKPOINT_CONFIG_FIELDS) | {"max_seq_len"} SNIFFABLE_TOPOLOGY_FIELDS = ( @@ -136,7 +95,49 @@ def tokenizer_metadata(tokenizer_path: str | None = None) -> dict: return meta +def resolve_checkpoint_tokenizer(checkpoint_path: str | None = None, explicit: str | None = None, + *, fallback: bool = True) -> str: + """Return the tokenizer directory that belongs to ``checkpoint_path``. + + Order: + 1. ``explicit``; + 2. ``/tokenizer`` (bundle layout); + 3. the ``tokenizer.path`` recorded in the sidecar config; + 4. the minimind tokenizer, only when ``fallback`` is set. + + A path that exists but holds no ``tokenizer.json`` raises an error. + It is never silently skipped, because a wrong vocabulary corrupts training without failing. + """ + explicit = os.path.expanduser(str(explicit or "").strip()) + if explicit: + if not os.path.isfile(os.path.join(explicit, "tokenizer.json")): + raise FileNotFoundError(f"tokenizer path {explicit!r} has no tokenizer.json") + return explicit + if checkpoint_path: + candidates = [os.path.join(os.path.dirname(os.path.abspath(checkpoint_path)), "tokenizer")] + meta_path = checkpoint_config_path(checkpoint_path) + if os.path.exists(meta_path): + with open(meta_path, encoding="utf-8") as fh: + recorded = (json.load(fh).get("tokenizer") or {}).get("path") + if recorded: + candidates.append(recorded) + for candidate in candidates: + if candidate and os.path.isfile(os.path.join(candidate, "tokenizer.json")): + return candidate + if not fallback: + raise FileNotFoundError(f"no tokenizer.json next to {checkpoint_path!r} or in its config metadata") + import chronos.deps + return chronos.deps.get_tokenizer_path() + + def config_to_dict(config: ChronosConfig | dict) -> dict: + try: + from chronos.config_schema import canonical_config_dict + canonical = canonical_config_dict(config) + canonical.pop("schema_version", None) + return {k: v for k, v in canonical.items() if v is not None} + except Exception: + pass if isinstance(config, dict): src = dict(config) # UI config uses max_seq_len; model config uses max_position_embeddings. @@ -187,9 +188,17 @@ def save_state_dict_with_config( extra: dict | None = None, ) -> str: os.makedirs(os.path.dirname(os.path.abspath(checkpoint_path)) or ".", exist_ok=True) + paging = getattr(model, "_chronos_paged_training_controller", None) + if paging is not None: + # A complete checkpoint must never serialize zero-sized placeholders. + # The next paged forward reapplies the active residency budget. + paging.materialize_all() state = model.state_dict() if half: - state = {k: v.half().cpu() for k, v in state.items()} + # ``half`` compacts FP32 weights to FP16 (the legacy release format). + # Weights already in a 16-bit compute dtype keep it: BF16 -> FP16 would + # clip BF16's exponent range for no size gain. + state = {k: (v.half() if v.dtype == torch.float32 else v).cpu() for k, v in state.items()} torch.save(state, checkpoint_path) save_checkpoint_config(checkpoint_path, config, tokenizer_path, stage=stage, extra=extra) return checkpoint_path @@ -204,7 +213,25 @@ def read_checkpoint_config(checkpoint_path: str) -> tuple[dict, str | None]: cfg = meta.get("config", meta) if not isinstance(cfg, dict): return {}, path - return dict(cfg), path + cfg = dict(cfg) + version = int(meta.get("checkpoint_config_version", 1) or 1) + if version < 2 and "recurrent_enabled" not in cfg: + import warnings + warnings.warn( + "Loading a Chronos v1 checkpoint in read-only compatibility mode. " + "Run `chronos migrate --output ... CHECKPOINT` before v2 training.", + RuntimeWarning, + stacklevel=2, + ) + cfg.update({ + "architecture_version": 1, + "recurrent_enabled": False, + "recurrent_depth": 1, + "min_recurrent_depth": 1, + "layer_aware_lookahead": False, + "layer_aware_cache": False, + }) + return cfg, path def read_project_config(config_path: str | None = None) -> tuple[dict, str | None]: @@ -435,6 +462,25 @@ def load_state_dict_controlled( allow_missing_substrings: Iterable[str] = (), allow_unexpected_prefixes: Iterable[str] = (), ) -> tuple[list[str], list[str]]: + # Known optional v2 modules are absent from legacy v1 checkpoints. Keep + # this compatibility explicit and topology-gated rather than allowing + # arbitrary missing tensors: a v1 model is constructed with the v2 + # controller object for API stability, but its execution graph is still + # disabled by ``recurrent_enabled=False``. + model_cfg = getattr(model, "config", None) + if model_cfg is not None: + legacy_recurrent = ( + not bool(getattr(model_cfg, "recurrent_enabled", True)) + or int(getattr(model_cfg, "architecture_version", 2) or 2) < 2 + ) + if legacy_recurrent: + allow_missing_prefixes = tuple(allow_missing_prefixes) + ( + "model.recurrent_controller.", + ) + if not getattr(model_cfg, "modalities", None): + allow_missing_prefixes = tuple(allow_missing_prefixes) + ( + "model.multimodal_fusion.", + ) result = model.load_state_dict(state_dict, strict=False) if result is None: missing, unexpected = [], [] diff --git a/chronos/model/config.py b/chronos/model/config.py index d5e69d6..6e96663 100644 --- a/chronos/model/config.py +++ b/chronos/model/config.py @@ -16,6 +16,107 @@ class ChronosConfig(MiniMindConfig): model_type = "chronos" def __init__(self, **kwargs): + # Chronos v2 recurrent-depth kernel. ``num_hidden_layers`` remains + # the number of *physical* parameter blocks so old checkpoints stay + # shape-compatible; recurrent_depth controls how often the recurrent + # portion is applied. A depth of one is exactly the v1 execution + # graph and is used when loading legacy checkpoints. + self.architecture_version: int = kwargs.pop("architecture_version", 3) + self.recurrent_enabled: bool = kwargs.pop("recurrent_enabled", True) + self.recurrent_depth: int = kwargs.pop("recurrent_depth", 1) + self.min_recurrent_depth: int = kwargs.pop("min_recurrent_depth", 1) + self.prelude_layers: int = kwargs.pop("prelude_layers", 0) + self.coda_layers: int = kwargs.pop("coda_layers", 0) + self.recurrent_groups = kwargs.pop("recurrent_groups", None) + self.adaptive_halting: bool = kwargs.pop("adaptive_halting", False) + self.halting_threshold: float = kwargs.pop("halting_threshold", 0.99) + self.halting_epsilon: float = kwargs.pop("halting_epsilon", 0.01) + # ACT needs a visible incentive to stop once the task loss is already + # low; 1e-3 was too weak and long runs learned to keep all virtual + # steps. 1e-2 preserves quality in the convergence benchmark while + # making adaptive depth a stable operating point. + self.ponder_loss_coef: float = kwargs.pop("ponder_loss_coef", 1e-2) + self.recurrent_distill_coef: float = kwargs.pop("recurrent_distill_coef", 0.05) + self.recurrent_stability_coef: float = kwargs.pop("recurrent_stability_coef", 1e-3) + self.recurrent_contraction_coef: float = kwargs.pop("recurrent_contraction_coef", 0.0) + self.recurrent_contraction_budget: float = kwargs.pop("recurrent_contraction_budget", 1.0) + # Curvature is normalized by hidden-state RMS in the recurrent + # executor, so 0.1 is a portable relative stability budget. The old + # absolute 1e-3 default almost never allowed an adaptive exit. + self.recurrent_stability_threshold: float = kwargs.pop("recurrent_stability_threshold", 0.1) + # A full Euler step is numerically valid but is too aggressive when a + # freshly initialized block is reused 2–8 times. Starting at 0.25 + # keeps the identity-preserving migration path stable while allowing + # the learned per-group step size to grow toward dt_max. + self.recurrent_dt_init: float = kwargs.pop("recurrent_dt_init", 0.25) + self.recurrent_dt_min: float = kwargs.pop("recurrent_dt_min", 0.05) + self.recurrent_dt_max: float = kwargs.pop("recurrent_dt_max", 1.0) + # v3 recurrent attention contract (see config_schema.resolve_*): + # shared_first | per_iteration | first_only. An explicit legacy + # ``recurrent_latent_refinement`` without a mode keeps v2 meaning. + explicit_latent = "recurrent_latent_refinement" in kwargs + self.recurrent_latent_refinement: bool = kwargs.pop("recurrent_latent_refinement", True) + mode = kwargs.pop("recurrent_attention_mode", None) + if not mode: + if self.architecture_version >= 3 and not explicit_latent: + mode = "shared_first" + else: + mode = "first_only" if self.recurrent_latent_refinement else "per_iteration" + mode = str(mode).lower() + if mode not in ("shared_first", "per_iteration", "first_only"): + raise ValueError(f"unknown recurrent_attention_mode {mode!r}") + self.recurrent_attention_mode: str = mode + # Keep the legacy flag coherent for estimator/export consumers. + self.recurrent_latent_refinement = mode == "first_only" + # Execution/training policy for shared recurrent blocks. ``fixed`` + # preserves the historical graph; ``scheduled`` samples a depth + # bucket during training and ``adaptive`` enables ACT mixing once the + # warmup has elapsed. ``parallel_refine`` is an inference hint used + # by runtimes that can batch latent iterations. + self.recurrent_strategy: str = kwargs.pop("recurrent_strategy", "scheduled") + self.recurrent_depth_jitter: float = kwargs.pop("recurrent_depth_jitter", 0.0) + self.recurrent_parallel_refine: bool = kwargs.pop("recurrent_parallel_refine", False) + self.recurrent_gradient_checkpoint_groups: bool = kwargs.pop( + "recurrent_gradient_checkpoint_groups", False + ) + self.max_inference_depth = kwargs.pop("max_inference_depth", None) + # 5-D [B,S,K,L,E] prediction is enabled by v2 presets. Keeping the + # switch explicit lets depth=1 legacy checkpoints retain their public + # 4-D output contract until migrated. + self.layer_aware_lookahead: bool = kwargs.pop("layer_aware_lookahead", False) + self.layer_aware_cache: bool = kwargs.pop("layer_aware_cache", False) + self.depth_schedule: str = kwargs.pop("depth_schedule", "fixed_then_adaptive") + self.depth_schedule_warmup_fraction: float = kwargs.pop( + "depth_schedule_warmup_fraction", 0.25 + ) + self.training_paging: bool = kwargs.pop("training_paging", False) + self.optimizer_paging: bool = kwargs.pop("optimizer_paging", False) + self.optimizer_page_dir: str = kwargs.pop("optimizer_page_dir", "./optimizer_pages") + self.page_size_mb: int = kwargs.pop("page_size_mb", 64) + self.activation_checkpointing: bool = kwargs.pop("activation_checkpointing", False) + + # Multimodal encoders are native kernel inputs, but are instantiated + # only for modalities discovered/enabled by the data pipeline. + raw_modalities = kwargs.pop("modalities", []) + if isinstance(raw_modalities, str): + raw_modalities = [x.strip() for x in raw_modalities.split(",") if x.strip()] + self.modalities = sorted(set(str(x).lower() for x in raw_modalities if str(x).lower() != "text")) + self.auto_detect_modalities: bool = kwargs.pop("auto_detect_modalities", True) + self.image_size: int = kwargs.pop("image_size", 224) + self.image_patch_size: int = kwargs.pop("image_patch_size", 16) + self.max_image_tokens: int = kwargs.pop("max_image_tokens", 196) + self.audio_sample_rate: int = kwargs.pop("audio_sample_rate", 16000) + self.max_audio_seconds: float = kwargs.pop("max_audio_seconds", 30.0) + self.max_audio_tokens: int = kwargs.pop("max_audio_tokens", 128) + self.modality_dropout: float = kwargs.pop("modality_dropout", 0.05) + self.freeze_modality_encoders: bool = kwargs.pop("freeze_modality_encoders", False) + self.image_encoder_type: str = kwargs.pop("image_encoder_type", "patch") + self.audio_encoder_type: str = kwargs.pop("audio_encoder_type", "conv") + self.modality_fusion_heads: int = kwargs.pop("modality_fusion_heads", 4) + self.media_error_policy: str = kwargs.pop("media_error_policy", "warn") + self.modality_adapters: dict = dict(kwargs.pop("modality_adapters", {})) + self.modality_alignment_coef: float = kwargs.pop("modality_alignment_coef", 0.0) + self.extra_metadata: dict = dict(kwargs.pop("extra_metadata", {})) self.lookahead_steps: int = kwargs.pop("lookahead_steps", 2) self.num_shared_experts: int = kwargs.pop("num_shared_experts", 1) self.lambda_balance: float = kwargs.pop("lambda_balance", 5e-4) @@ -26,7 +127,56 @@ def __init__(self, **kwargs): self.use_hybrid_attention: bool = kwargs.pop("use_hybrid_attention", True) self.kv_latent_dim: int = kwargs.pop("kv_latent_dim", 64) # MLA compression dim self.rope_dim: int = kwargs.pop("rope_dim", 32) # RoPE portion of head_dim + # How the rotated slice of each head is built. "per_head": rotate the + # rope_dim channels with their own frequency ladder (a true rotation, + # relative-position exact). "legacy_sliced": slice the first rope_dim + # columns of the head_dim table (what checkpoints before 2026-11-06 were + # trained with; kept so they load bit-identically). + self.rope_layout: str = kwargs.pop("rope_layout", "legacy_sliced") + if self.rope_layout not in ("per_head", "legacy_sliced"): + raise ValueError(f"rope_layout must be 'per_head' or 'legacy_sliced', got {self.rope_layout!r}") self.sliding_window_size: int = kwargs.pop("sliding_window_size", 2048) + # Inference KV storage: "auto" keeps the compute dtype; "int8" stores + # MLA latents (and cached rope keys) / SWA K,V as per-token int8 with + # one FP32 scale per ``kv_quant_group_size`` channels. Training and + # no-cache forwards never quantize. + self.kv_cache_dtype: str = str(kwargs.pop("kv_cache_dtype", "auto") or "auto").lower() + if self.kv_cache_dtype in ("fp16", "bf16", "fp32", "compute", "none"): + self.kv_cache_dtype = "auto" # the cache follows the compute dtype + self.kv_quant_group_size: int = int(kwargs.pop("kv_quant_group_size", 32)) + # MLA decode path: "absorbed" folds W_uk into the query and W_uv into + # the output (scores against the latent cache, no per-step K/V + # up-projection of the whole prefix); "expanded" is the reference + # path; "auto" absorbs for cached single-chunk decode (S <= 8) only. + # Default "expanded": like every new option, absorption is opt-in. + # "auto" mixes paths (expanded prefill, absorbed decode), which in + # bf16 rounds differently from a full forward and can flip a near-tied + # top-k route; numerics then differ from the pre-v3 kernel. + self.mla_decode: str = str(kwargs.pop("mla_decode", "expanded") or "expanded").lower() + # Also cache post-RoPE rope keys [B, T, nKV, rope_dim]: memory for + # throughput (skips re-projecting + rotating the prefix each step). + self.mla_cache_rope: bool = bool(kwargs.pop("mla_cache_rope", False)) + # Expert selection across recurrent iterations of one group: + # "per_iteration" routes every iteration; "shared" reuses iteration + # zero's top-k ids (weights still re-scored), so the active expert + # union per group is at most top-k per token (lazy-loading friendly). + self.router_scope: str = str(kwargs.pop("router_scope", "per_iteration") or "per_iteration").lower() + # Self-speculative training hook: KL(full-depth || draft-depth) on the + # LM head, teacher detached. 0 disables (no extra compute). + self.draft_exit_kl_coef: float = float(kwargs.pop("draft_exit_kl_coef", kwargs.pop("draft_exit_kl", 0.0)) or 0.0) + self.draft_exit_depth: int = int(kwargs.pop("draft_exit_depth", 1)) + # Hard expert-memory budget for lazy inference (None = heuristic). + # max_resident_experts caps fast-tier (VRAM / Metal-hot) pages; + # ram_budget_mb caps the RAM tier in MiB of expert pages. + self.max_resident_experts = kwargs.pop("max_resident_experts", None) + self.ram_budget_mb = kwargs.pop("ram_budget_mb", None) + for name, value, allowed in ( + ("mla_decode", self.mla_decode, ("auto", "absorbed", "expanded")), + ("router_scope", self.router_scope, ("per_iteration", "shared")), + ("kv_cache_dtype", self.kv_cache_dtype, ("auto", "int8")), + ): + if value not in allowed: + raise ValueError(f"{name} must be one of {allowed}, got {value!r}") # Pinned memory safety: max RAM for expert buffers (fraction of physical RAM) self.pinned_memory_max_fraction: float = kwargs.pop("pinned_memory_max_fraction", 0.25) # Storage backend for expert weights on SSD. @@ -70,3 +220,19 @@ def __init__(self, **kwargs): # Force MoE on kwargs.setdefault("use_moe", True) super().__init__(**kwargs) + + from chronos.config_schema import normalize_model_dimensions + resolved = normalize_model_dimensions(vars(self)) + for key, value in resolved.items(): + setattr(self, key, value) + + @property + def virtual_num_hidden_layers(self) -> int: + """Executed attention depth, including repeated shared blocks.""" + if not self.recurrent_enabled: + return int(self.num_hidden_layers) + return ( + int(self.prelude_layers) + + int(self.recurrent_groups) * int(self.recurrent_depth) + + int(self.coda_layers) + ) diff --git a/chronos/model/factory.py b/chronos/model/factory.py new file mode 100644 index 0000000..c731987 --- /dev/null +++ b/chronos/model/factory.py @@ -0,0 +1,94 @@ +"""Parameter-budget model factory. + +``fit_config(target_total, target_active, ...)`` searches the Chronos design +space (width, expert count, expert width) for a valid configuration whose +*allocated* parameter counts are closest to the targets. Counts come from +instantiating the real model on the ``meta`` device, so they always match +what training allocates — no separate formula to drift. + +Active parameters follow the usual MoE convention: everything except the +routed experts, plus ``top_k`` routed experts per MoE layer. Because loop +iterations may route differently, the factory also reports the union upper +bound ``min(E, top_k·depth)`` experts for recurrent layers. +""" +from __future__ import annotations + +import itertools +from dataclasses import dataclass + +import torch + + +@dataclass +class FitResult: + config: dict + total: int + active: int + active_union_upper: int + embedding: int + per_expert: int + score: float + + +def count_parameters(config_kwargs: dict) -> dict: + from chronos.model.config import ChronosConfig + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.model.moe_chronos import ChronosMOEFeedForward + cfg = ChronosConfig(**config_kwargs) + with torch.device("meta"): + model = ChronosForCausalLM(cfg) + seen, total = set(), 0 + for p in model.parameters(): + if id(p) not in seen: + seen.add(id(p)); total += p.numel() + routed, per_expert, moe_layers = 0, 0, 0 + for module in model.modules(): + if isinstance(module, ChronosMOEFeedForward): + moe_layers += 1 + per_expert = sum(p.numel() for p in module.experts[0].parameters()) + routed += per_expert * len(module.experts) + k, E = int(cfg.num_experts_per_tok), int(cfg.num_experts) + active = total - routed + moe_layers * k * per_expert + union_extra = int(cfg.recurrent_groups) * (min(E, k * int(cfg.recurrent_depth)) - k) * per_expert + return {"total": total, "active": active, "active_union_upper": active + union_extra, + "embedding": int(model.model.embed_tokens.weight.numel()), "per_expert": per_expert, + "virtual_layers": int(cfg.virtual_num_hidden_layers)} + + +def fit_config(target_total: float, target_active: float, *, vocab_size: int, + base: dict | None = None, hidden_sizes=(320, 352, 384, 416, 448, 512), + expert_counts=tuple(range(8, 49, 4)), expert_widths=(160, 192, 224, 256, 288, 320), + head_dim: int = 64) -> FitResult: + """Grid search minimizing relative error on (total, active).""" + base = dict(base or {}) + best = None + for h, e, i in itertools.product(hidden_sizes, expert_counts, expert_widths): + heads = max(1, h // head_dim) + kv = 2 if heads % 2 == 0 else 1 + kw = {**base, "hidden_size": h, "num_attention_heads": heads, "num_key_value_heads": kv, + "num_experts": e, "moe_intermediate_size": i, "vocab_size": vocab_size} + try: + c = count_parameters(kw) + except (ValueError, AssertionError): + continue + score = abs(c["total"] / target_total - 1) + 2.0 * abs(c["active"] / target_active - 1) + if best is None or score < best.score: + best = FitResult(kw, c["total"], c["active"], c["active_union_upper"], c["embedding"], + c["per_expert"], score) + if best is None: + raise ValueError("no valid configuration in the search space") + return best + + +ENZH_DEMO_BASE = dict( + architecture_version=3, num_hidden_layers=6, prelude_layers=1, coda_layers=1, + recurrent_depth=4, min_recurrent_depth=1, max_inference_depth=4, + recurrent_attention_mode="shared_first", adaptive_halting=True, + depth_schedule="fixed_then_adaptive", depth_schedule_warmup_fraction=0.3, + halting_threshold=0.95, num_experts_per_tok=2, num_shared_experts=1, + kv_latent_dim=128, rope_dim=32, rope_layout="per_head", sliding_window_size=128, + max_position_embeddings=1024, + layer_aware_lookahead=True, layer_aware_cache=True, lookahead_steps=2, + dropout=0.0, fallback_mask_prob=0.08, rope_theta=1e6, + bos_token_id=1, eos_token_id=2, tie_word_embeddings=True, +) diff --git a/chronos/model/hf_io.py b/chronos/model/hf_io.py index 20126f9..e7f4409 100644 --- a/chronos/model/hf_io.py +++ b/chronos/model/hf_io.py @@ -64,6 +64,9 @@ def save_chronos_pretrained( if ``expert_cache_dir`` is provided so the offloaded layout travels with the checkpoint.""" os.makedirs(save_directory, exist_ok=True) + paging = getattr(model, "_chronos_paged_training_controller", None) + if paging is not None: + paging.materialize_all() # MiniMind/Chronos tie ``lm_head.weight`` to ``model.embed_tokens.weight`` # but doesn't declare it via ``_tied_weights_keys``, which trips diff --git a/chronos/model/hybrid_attention.py b/chronos/model/hybrid_attention.py index f275c2f..dc69dde 100644 --- a/chronos/model/hybrid_attention.py +++ b/chronos/model/hybrid_attention.py @@ -21,8 +21,38 @@ import torch.nn as nn import torch.nn.functional as F -from model.model_minimind import RMSNorm, apply_rotary_pos_emb, repeat_kv from .config import ChronosConfig +from .kv_quant import kv_read, kv_store, resolve_kv_cache_dtype +from .precision_ops import RMSNorm, apply_rope, attention, build_attention_mask, causal_allowed, rope_tables + + +class _RopeTables: + """FP32 RoPE tables owned by one attention layer. + + Stored as plain attributes, not buffers, so ``model.to(bf16)`` never + rounds them; they follow the activations' device on first use and grow + geometrically past ``max_position_embeddings``. ``rope_layout`` picks + the table width (see ``ChronosConfig.rope_layout``). + """ + + def _init_rope(self, config, full_dim: int, rotated_dim: int) -> None: + layout = str(getattr(config, "rope_layout", "legacy_sliced") or "legacy_sliced") + if layout == "per_head" and int(rotated_dim) % 2: + raise ValueError(f"rope_layout='per_head' needs an even rotated dimension, got {rotated_dim}") + self._rope_table_dim = int(rotated_dim) if layout == "per_head" else int(full_dim) + self._rope_base = float(getattr(config, "rope_theta", 1_000_000.0)) + self._rope_scaling = getattr(config, "rope_scaling", None) + self._rope_len = int(getattr(config, "max_position_embeddings", 4096)) + self._rope_cos = self._rope_sin = None + + def _rope(self, start: int, length: int, width: int, device): + """FP32 cos/sin rows ``start..start+length-1`` cut to ``width`` columns.""" + need = int(start) + int(length) + if self._rope_cos is None or self._rope_cos.device != device or self._rope_cos.shape[0] < need: + size = max(need, self._rope_len, 2 * (self._rope_cos.shape[0] if self._rope_cos is not None else 0)) + self._rope_cos, self._rope_sin = rope_tables(self._rope_table_dim, size, self._rope_base, + self._rope_scaling, device) + return self._rope_cos[start:need, :width], self._rope_sin[start:need, :width] def _key_padding_bias(attention_mask: torch.Tensor, key_len: int, dtype: torch.dtype) -> torch.Tensor: @@ -57,9 +87,18 @@ def _key_padding_bias(attention_mask: torch.Tensor, key_len: int, dtype: torch.d raise ValueError(f"attention_mask must be 2-D, 3-D, or 4-D, got shape {tuple(attention_mask.shape)}") +def _sdpa_mixed_value_width(device) -> bool: + """True when fused SDPA accepts a value width different from Q/K. + + CPU and CUDA kernels do; MPS raises a shape error (checked on torch 2.10), + so the absorbed MLA path uses explicit matmuls there. + """ + return device.type != "mps" + + # ── MLA: Multi-head Latent Attention ───────────────────────────── -class MLAAttention(nn.Module): +class MLAAttention(_RopeTables, nn.Module): """ Multi-head Latent Attention (DeepSeek MLA style). @@ -111,77 +150,152 @@ def __init__(self, config: ChronosConfig): self.o_proj = nn.Linear(self.n_heads * self.head_dim, H, bias=False) self.attn_dropout = nn.Dropout(self.dropout) self.resid_dropout = nn.Dropout(self.dropout) - self.flash = hasattr(F, 'scaled_dot_product_attention') and config.flash_attn - - def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None): + self._init_rope(config, self.head_dim, self.rope_dim) + # Inference-time cache/decode policy (no parameters; see ChronosConfig). + self.mla_decode = str(getattr(config, "mla_decode", "expanded") or "expanded").lower() + self.cache_rope = bool(getattr(config, "mla_cache_rope", False)) + self.kv_int8 = resolve_kv_cache_dtype(config) == "int8" + self.kv_group = int(getattr(config, "kv_quant_group_size", 32) or 32) + + def _use_absorbed(self, S: int, T: int) -> bool: + if self.mla_decode == "expanded": + return False + if self.mla_decode == "absorbed": + return True + # auto: cached short-chunk decode only (prefill is cheaper expanded). + return (not self.training) and S <= 8 and T > S + + def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None, + kv_override=None, return_shared_kv=False): + """MLA attention. + + Cache entries are ``(c_kv,)`` or, with ``mla_cache_rope``, + ``(c_kv, k_rope)`` (post-RoPE rope keys ``[B, T, nKV, rope_dim]``); + either tensor may be a :class:`QuantizedKV` when ``kv_cache_dtype`` + is ``int8``. ``kv_override`` is the shared-KV recurrent contract: + the same tuple (always in the compute dtype) returned by iteration + zero with ``return_shared_kv=True``; it already contains this chunk, + so queries come from ``x`` while keys/values are re-read. + """ B, S, _ = x.shape - cos, sin = position_embeddings - - # Query - q_nope = self.q_nope_norm( - self.q_nope_proj(x).view(B, S, self.n_heads, self.nope_dim) - ) - q_rope = self.q_rope_norm( - self.q_rope_proj(x).view(B, S, self.n_heads, self.rope_dim) - ) - # apply_rotary_pos_emb expects [B, S, nH, d] (NOT transposed) - # cos/sin are full-length buffers; slice to current step positions - past_len = past_key_value[0].shape[1] if past_key_value is not None else 0 - cos_r = cos[past_len:past_len + S, :self.rope_dim] # [S, rope_dim] - sin_r = sin[past_len:past_len + S, :self.rope_dim] - q_rope, _ = apply_rotary_pos_emb(q_rope, q_rope, cos_r, sin_r) # [B, S, nH, rope_dim] - q = torch.cat([q_nope, q_rope], dim=-1) # [B, S, nH, head_dim] - - # KV latent (this is what we cache) - c_kv = self.kv_down_norm(self.kv_down_proj(x)) # [B, S, kv_latent_dim] - - if past_key_value is not None: - past_c_kv = past_key_value[0] - c_kv_full = torch.cat([past_c_kv, c_kv], dim=1) + del position_embeddings # layer-owned FP32 tables (rope_layout aware) + dtype = x.dtype + + q_nope = self.q_nope_norm(self.q_nope_proj(x).view(B, S, self.n_heads, self.nope_dim)) + q_rope = self.q_rope_norm(self.q_rope_proj(x).view(B, S, self.n_heads, self.rope_dim)) + if kv_override is not None: + c_kv_full = kv_override[0] + k_rope = kv_override[1] if len(kv_override) > 1 else None + past_len = c_kv_full.shape[1] - S else: - c_kv_full = c_kv - - full_len = c_kv_full.shape[1] - - # Expand KV from latent - k_nope = self.k_nope_proj(c_kv_full).view(B, full_len, self.n_kv_heads, self.nope_dim) - k_rope_raw = self.k_rope_proj(c_kv_full).view(B, full_len, self.n_kv_heads, self.rope_dim) - v = self.v_proj(c_kv_full).view(B, full_len, self.n_kv_heads, self.head_dim) - - # RoPE on keys over full cached sequence — [B, T, nKV, rope_dim] - cos_full = cos[:full_len, :self.rope_dim] - sin_full = sin[:full_len, :self.rope_dim] - k_rope, _ = apply_rotary_pos_emb(k_rope_raw, k_rope_raw, cos_full, sin_full) - k = torch.cat([k_nope, k_rope], dim=-1) # [B, T, nKV, head_dim] - - past_kv = (c_kv_full,) if use_cache else None - - # Attention - q = q.transpose(1, 2) # [B, nH, S, dH] - k = repeat_kv(k, self.n_rep).transpose(1, 2) # [B, nH, T, dH] - v = repeat_kv(v, self.n_rep).transpose(1, 2) # [B, nH, T, dH] - - T = k.shape[2] - if self.flash and S > 1 and attention_mask is None: - out = F.scaled_dot_product_attention( - q, k, v, - dropout_p=self.dropout if self.training else 0.0, - is_causal=True, - ) + past_len = past_key_value[0].shape[1] if past_key_value is not None else 0 + cos_r, sin_r = self._rope(past_len, S, self.rope_dim, x.device) # FP32 [S, rope_dim] + q_rope = apply_rope(q_rope, cos_r, sin_r) # [B, S, nH, rope_dim] + + if kv_override is None: + # KV latent (this is what we cache). Each token is quantized + # once when appended; the prefix is read back in the compute + # dtype while this chunk stays exact. + c_kv = self.kv_down_norm(self.kv_down_proj(x)) # [B, S, latent] + quant = self.kv_int8 and use_cache + past_c = past_key_value[0] if past_key_value is not None else None + c_kv_full = c_kv if past_c is None else torch.cat([kv_read(past_c, dtype), c_kv], dim=1) + stored_c = (kv_store(c_kv, past_c, True, self.kv_group) if quant else c_kv_full) if use_cache else None + k_rope = stored_r = None + if self.cache_rope: + k_rope_new = apply_rope( + self.k_rope_proj(c_kv).view(B, S, self.n_kv_heads, self.rope_dim), cos_r, sin_r) + past_r = past_key_value[1] if past_key_value is not None and len(past_key_value) > 1 else None + if past_c is not None and past_r is None: # cache built without rope keys + past_r = self._rotated_k_rope(kv_read(past_c, dtype), 0) + k_rope = k_rope_new if past_r is None else torch.cat([kv_read(past_r, dtype), k_rope_new], 1) + stored_r = (kv_store(k_rope_new, past_r, True, self.kv_group) if quant else k_rope) if use_cache else None + if k_rope is None: # [B, T, nKV, rope_dim] + k_rope = self._rotated_k_rope(c_kv_full, 0) + past_kv = ((stored_c,) if stored_r is None else (stored_c, stored_r)) if use_cache else None + # Transient view for recurrent reader iterations (never persisted): + # carries the rotated rope keys so readers skip re-projecting them. + shared_kv = (c_kv_full, k_rope) else: - scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) - scores[:, :, :, -S:] += torch.full((S, S), float('-inf'), device=scores.device).triu(1) - if attention_mask is not None: - scores += _key_padding_bias(attention_mask, T, scores.dtype) - out = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(q)) @ v - - out = out.transpose(1, 2).reshape(B, S, -1) + past_kv = None + shared_kv = kv_override + if k_rope is None: + k_rope = self._rotated_k_rope(c_kv_full, 0) + + T = c_kv_full.shape[1] + mask, is_causal = build_attention_mask(S, T, x.device, past=T - S, window=None, + attention_mask=attention_mask) + dropout_p = self.dropout if self.training else 0.0 + if self._use_absorbed(S, T): + out = self._absorbed_attention(q_nope, q_rope, c_kv_full, k_rope, mask, is_causal, dropout_p) + else: + k_nope = self.k_nope_proj(c_kv_full).view(B, T, self.n_kv_heads, self.nope_dim) + v = self.v_proj(c_kv_full).view(B, T, self.n_kv_heads, self.head_dim) + q = torch.cat([q_nope, q_rope], dim=-1) # [B, S, nH, head_dim] + k = torch.cat([k_nope, k_rope], dim=-1) # [B, T, nKV, head_dim] + # Fused attention: native GQA, FP32 softmax inside the kernel. + out = attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), mask=mask, + is_causal=is_causal, dropout_p=dropout_p) + out = out.transpose(1, 2) + + out = out.reshape(B, S, -1) + if return_shared_kv: + return self.resid_dropout(self.o_proj(out)), past_kv, shared_kv return self.resid_dropout(self.o_proj(out)), past_kv + def _rotated_k_rope(self, c_kv: torch.Tensor, start: int) -> torch.Tensor: + B, T, _ = c_kv.shape + raw = self.k_rope_proj(c_kv).view(B, T, self.n_kv_heads, self.rope_dim) + cos, sin = self._rope(start, T, self.rope_dim, c_kv.device) + return apply_rope(raw, cos, sin) + + def _absorbed_attention(self, q_nope, q_rope, c_kv, k_rope, mask, is_causal, dropout_p): + """Attention against the latent cache without expanding K/V. + + Per KV group ``g`` (heads ``g*n_rep .. g*n_rep+n_rep-1``): + score = (q_nope · W_k_nope[g]) · c_kv + q_rope · k_rope[g] + out = W_v[g] · Σ softmax(score) · c_kv + which equals the expanded path in exact arithmetic. Returns + ``[B, S, nH, head_dim]`` in the compute dtype. + """ + B, S, nH, _ = q_nope.shape + T, L, G, R = c_kv.shape[1], self.kv_latent_dim, self.n_kv_heads, self.n_rep + scale = 1.0 / math.sqrt(self.head_dim) + w_k = self.k_nope_proj.weight.view(G, self.nope_dim, L) + w_v = self.v_proj.weight.view(G, self.head_dim, L) + q_lat = torch.einsum("bsgrn,gnl->bgrsl", q_nope.view(B, S, G, R, self.nope_dim), w_k) + q_lat = q_lat.reshape(B, nH, S, L) # head h = g*R + r + q_r = q_rope.transpose(1, 2) # [B, nH, S, rope] + k_r = k_rope.transpose(1, 2) # [B, G, T, rope] + if _sdpa_mixed_value_width(c_kv.device): + # SDPA with a latent value width (supported on CPU/CUDA): one + # fused kernel, FP32 softmax inside, native GQA over the G groups. + k = torch.cat([c_kv[:, None].expand(B, G, T, L), k_r], dim=-1) # [B, G, T, L+rope] + v = c_kv[:, None].expand(B, G, T, L) + lat = attention(torch.cat([q_lat, q_r], dim=-1), k, v, mask=mask, is_causal=is_causal, + dropout_p=dropout_p, scale=scale) # [B, nH, S, L] + else: + # MPS SDPA rejects V width != Q/K width: explicit matmuls in the + # compute dtype with an FP32 softmax. + scores = (q_lat * scale) @ c_kv[:, None].transpose(-1, -2) # [B, nH, S, T] + scores = scores + ((q_r * scale).view(B, G, R, S, -1) + @ k_r[:, :, None].transpose(-1, -2)).view(B, nH, S, T) + if mask is None and is_causal: + mask = causal_allowed(S, T, c_kv.device, past=T - S) + scores = scores.float() + if mask is not None: + scores = scores.masked_fill(~mask, float("-inf")) + probs = torch.softmax(scores, dim=-1) + if dropout_p: + probs = F.dropout(probs, p=dropout_p, training=True) + lat = probs.to(c_kv.dtype) @ c_kv[:, None] # [B, nH, S, L] + out = torch.einsum("bgrsl,gdl->bsgrd", lat.view(B, G, R, S, L), w_v) # [B, S, G, R, d] + return out.reshape(B, S, nH, self.head_dim) + # ── Sliding Window Attention ────────────────────────────────────── -class SlidingWindowAttention(nn.Module): +class SlidingWindowAttention(_RopeTables, nn.Module): """ Local sliding window attention. KV cache is capped at window_size tokens, preventing unbounded growth in long-context generation. @@ -198,6 +312,9 @@ def __init__(self, config: ChronosConfig): self.n_rep = self.n_heads // self.n_kv_heads self.dropout = config.dropout self.window_size = getattr(config, 'sliding_window_size', 2048) + # Extra retained cache entries (outside the attention window, masked + # out) so speculative decoding can roll back rejected draft tokens. + self.cache_slack = 0 H = config.hidden_size self.q_proj = nn.Linear(H, self.n_heads * self.head_dim, bias=False) @@ -208,53 +325,64 @@ def __init__(self, config: ChronosConfig): self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.attn_dropout = nn.Dropout(self.dropout) self.resid_dropout = nn.Dropout(self.dropout) - self.flash = hasattr(F, 'scaled_dot_product_attention') and config.flash_attn - - def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None): + self._init_rope(config, self.head_dim, self.head_dim) + self.kv_int8 = resolve_kv_cache_dtype(config) == "int8" + self.kv_group = int(getattr(config, "kv_quant_group_size", 32) or 32) + + def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None, + kv_override=None, return_shared_kv=False): + """Sliding-window attention. + + ``kv_override`` is ``(k_full, v_full, absolute_position)`` returned by + iteration zero with ``return_shared_kv=True``: post-RoPE keys and + values for the retained prefix plus this chunk, before eviction. + """ B, S, _ = x.shape - cos, sin = position_embeddings - # cos/sin are full-length buffers; slice to current positions - past_len = past_key_value[0].shape[1] if past_key_value is not None else 0 - cos_s = cos[past_len:past_len + S] - sin_s = sin[past_len:past_len + S] + del position_embeddings # layer-owned FP32 tables + if kv_override is not None: + xk, xv, past_len = kv_override[0], kv_override[1], int(kv_override[2]) - S + cached_len = xk.shape[1] - S + else: + # cos/sin are full-length buffers; slice to current positions + cached_len = past_key_value[0].shape[1] if past_key_value is not None else 0 + past_len = (past_key_value[2] if len(past_key_value) > 2 else cached_len) if past_key_value is not None else 0 + cos_s, sin_s = self._rope(past_len, S, self.head_dim, x.device) xq = self.q_norm(self.q_proj(x).view(B, S, self.n_heads, self.head_dim)) - xk = self.k_norm(self.k_proj(x).view(B, S, self.n_kv_heads, self.head_dim)) - xv = self.v_proj(x).view(B, S, self.n_kv_heads, self.head_dim) - - xq, xk = apply_rotary_pos_emb(xq, xk, cos_s, sin_s) # [B, S, nH/nKV, d] - - if past_key_value is not None: - xk = torch.cat([past_key_value[0], xk], dim=1) - xv = torch.cat([past_key_value[1], xv], dim=1) - - # Evict tokens beyond window_size (sliding window) - if xk.shape[1] > self.window_size: - xk = xk[:, -self.window_size:, :, :] - xv = xv[:, -self.window_size:, :, :] - - past_kv = (xk, xv) if use_cache else None - - xq = xq.transpose(1, 2) - xk = repeat_kv(xk, self.n_rep).transpose(1, 2) - xv = repeat_kv(xv, self.n_rep).transpose(1, 2) - - T = xk.shape[2] - if self.flash and S > 1 and attention_mask is None: - out = F.scaled_dot_product_attention( - xq, xk, xv, - dropout_p=self.dropout if self.training else 0.0, - is_causal=(T == S), - ) + if kv_override is None: + xk = self.k_norm(self.k_proj(x).view(B, S, self.n_kv_heads, self.head_dim)) + xv = self.v_proj(x).view(B, S, self.n_kv_heads, self.head_dim) + xq, xk = apply_rope(xq, cos_s, sin_s), apply_rope(xk, cos_s, sin_s) # [B, S, nH/nKV, d] + new_k, new_v = xk, xv + if past_key_value is not None: + xk = torch.cat([kv_read(past_key_value[0], xk.dtype), xk], dim=1) + xv = torch.cat([kv_read(past_key_value[1], xv.dtype), xv], dim=1) else: - scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) - if T == S: - scores[:, :, :, -S:] += torch.full((S, S), float('-inf'), device=scores.device).triu(1) - if attention_mask is not None: - scores += _key_padding_bias(attention_mask, T, scores.dtype) - out = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv + xq = apply_rope(xq, cos_s, sin_s) + + # Keep full keys for this chunk's queries, then bound the persistent + # cache. Track absolute positions after eviction for correct RoPE. + keep = self.window_size + int(self.cache_slack) + past_kv = None + if use_cache and self.kv_int8 and kv_override is None: + # Quantize only the new tokens (each token once); evict afterwards. + prev = past_key_value if past_key_value is not None else (None, None) + qk = kv_store(new_k, prev[0], True, self.kv_group) + qv = kv_store(new_v, prev[1], True, self.kv_group) + past_kv = (qk[:, max(0, qk.shape[1] - keep):], qv[:, max(0, qv.shape[1] - keep):], past_len + S) + elif use_cache: + past_kv = (xk[:, -keep:], xv[:, -keep:], past_len+S) + shared_kv = (xk, xv, past_len + S) + + T = xk.shape[1] + mask, is_causal = build_attention_mask(S, T, x.device, past=cached_len, window=self.window_size, + attention_mask=attention_mask) + out = attention(xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2), mask=mask, + is_causal=is_causal, dropout_p=self.dropout if self.training else 0.0) out = out.transpose(1, 2).reshape(B, S, -1) + if return_shared_kv: + return self.resid_dropout(self.o_proj(out)), past_kv, shared_kv return self.resid_dropout(self.o_proj(out)), past_kv diff --git a/chronos/model/kv_quant.py b/chronos/model/kv_quant.py new file mode 100644 index 0000000..61f3507 --- /dev/null +++ b/chronos/model/kv_quant.py @@ -0,0 +1,108 @@ +"""Inference-only int8 KV-cache storage (Torch). + +``QuantizedKV`` stores a ``[B, T, *, D]`` cache tensor as symmetric int8 +with one FP32 scale per token and ``group_size`` channels along the last +axis. Each token is quantized exactly once, when it is appended; the cache +is dequantized into the compute dtype when attention reads it (one rounding +per read, no FP32 activation copy). FP32 scales keep tiny latent norms from +underflowing and cost 4/group_size bytes per channel (12.5 % at 32). + +The object mimics the parts of the tensor interface the cache plumbing uses +(``shape``, ``[:, :n]`` slicing, ``numel``/``element_size`` byte counting), +so ``truncate_cache`` rollback and sequence-length bookkeeping stay generic. +""" +from __future__ import annotations + +import torch + +KV_CACHE_DTYPES = ("auto", "int8") + + +def resolve_kv_cache_dtype(config) -> str: + value = str(getattr(config, "kv_cache_dtype", "auto") or "auto").lower() + if value in ("fp16", "bf16", "fp32", "compute", "none"): + value = "auto" + if value not in KV_CACHE_DTYPES: + raise ValueError(f"kv_cache_dtype must be one of {KV_CACHE_DTYPES}, got {value!r}") + return value + + +class QuantizedKV: + __slots__ = ("data", "scale", "group_size", "dtype") + + def __init__(self, data: torch.Tensor, scale: torch.Tensor, group_size: int, dtype: torch.dtype): + self.data, self.scale, self.group_size, self.dtype = data, scale, int(group_size), dtype + + # ---- construction ----------------------------------------------------- + @classmethod + def quantize(cls, x: torch.Tensor, group_size: int = 32) -> "QuantizedKV": + d = int(x.shape[-1]) + g = int(group_size) if group_size and d % int(group_size) == 0 else d + groups = x.float().reshape(*x.shape[:-1], d // g, g) + scale = groups.abs().amax(dim=-1, keepdim=True).clamp_min(1e-12) / 127.0 + data = torch.round(groups / scale).clamp_(-127, 127).to(torch.int8).reshape(x.shape) + return cls(data, scale, g, x.dtype) + + def append(self, x: torch.Tensor) -> "QuantizedKV": + new = QuantizedKV.quantize(x, self.group_size) + return QuantizedKV(torch.cat([self.data, new.data], dim=1), torch.cat([self.scale, new.scale], dim=1), + self.group_size, self.dtype) + + def dequantize(self, dtype: torch.dtype | None = None) -> torch.Tensor: + d = int(self.data.shape[-1]) + groups = self.data.reshape(*self.data.shape[:-1], d // self.group_size, self.group_size) + return (groups.float() * self.scale).reshape(self.data.shape).to(dtype or self.dtype) + + # ---- tensor-like surface used by cache plumbing ----------------------- + @property + def shape(self): + return self.data.shape + + @property + def device(self): + return self.data.device + + def __getitem__(self, index): + if not (isinstance(index, tuple) and len(index) == 2 and index[0] == slice(None) + and isinstance(index[1], slice) and index[1].step in (None, 1)): + raise TypeError("QuantizedKV supports only [:, start:stop] token slicing") + return QuantizedKV(self.data[index], self.scale[index], self.group_size, self.dtype) + + def nbytes(self) -> int: + return self.data.numel() * self.data.element_size() + self.scale.numel() * self.scale.element_size() + + def numel(self) -> int: # byte accounting helpers treat numel*element_size as bytes + return self.nbytes() + + @staticmethod + def element_size() -> int: + return 1 + + +def kv_store(x: torch.Tensor, past, quantize: bool, group_size: int): + """Append ``x`` to ``past`` (tensor, QuantizedKV or None) in the cache format.""" + if quantize: + if isinstance(past, QuantizedKV): + return past.append(x) + q = QuantizedKV.quantize(x, group_size) + if past is not None: # a float prefix (e.g. a cache built before int8 was enabled) + q = QuantizedKV.quantize(torch.cat([past.to(x.dtype), x], dim=1), group_size) + return q + if isinstance(past, QuantizedKV): + past = past.dequantize(x.dtype) + return x if past is None else torch.cat([past, x], dim=1) + + +def kv_read(entry, dtype: torch.dtype) -> torch.Tensor: + return entry.dequantize(dtype) if isinstance(entry, QuantizedKV) else entry + + +def entry_nbytes(entry) -> int: + """Bytes held by one cache entry (tensors, QuantizedKV, ints).""" + total = 0 + for item in entry or (): + if isinstance(item, QuantizedKV): + total += item.nbytes() + elif isinstance(item, torch.Tensor): + total += item.numel() * item.element_size() + return total diff --git a/chronos/model/lookahead_router.py b/chronos/model/lookahead_router.py index 4ef14a3..0523943 100644 --- a/chronos/model/lookahead_router.py +++ b/chronos/model/lookahead_router.py @@ -23,22 +23,27 @@ def __init__(self, config: ChronosConfig): super().__init__() self.lookahead_steps = config.lookahead_steps self.num_experts = config.num_experts + self.layer_aware = bool(getattr(config, "layer_aware_lookahead", False)) + self.num_groups = int(getattr(config, "recurrent_groups", 1) or 1) H = config.hidden_size K = config.lookahead_steps + 1 # current + future steps # Lightweight 2-layer MLP: H -> H//4 -> E*K self.proj = nn.Sequential( - nn.Linear(H, H // 4, bias=False), + nn.Linear(H, max(1, H // 4), bias=False), nn.SiLU(), - nn.Linear(H // 4, self.num_experts * K, bias=False), + nn.Linear(max(1, H // 4), self.num_experts * K * (self.num_groups if self.layer_aware else 1), bias=False), ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # hidden_states: [B, S, H] B, S, H = hidden_states.shape K = self.lookahead_steps + 1 - logits = self.proj(hidden_states) # [B, S, E*K] - logits = logits.view(B, S, K, self.num_experts) # [B, S, K, E] + logits = self.proj(hidden_states) + if self.layer_aware: + logits = logits.view(B, S, K, self.num_groups, self.num_experts) + else: + logits = logits.view(B, S, K, self.num_experts) return F.softmax(logits, dim=-1) # [B, S, K, E] def predict_next_experts( diff --git a/chronos/model/model_chronos.py b/chronos/model/model_chronos.py index 4d85ae3..a7ec7d7 100644 --- a/chronos/model/model_chronos.py +++ b/chronos/model/model_chronos.py @@ -6,15 +6,48 @@ import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel, GenerationMixin -from transformers.modeling_outputs import MoeCausalLMOutputWithPast from model.model_minimind import ( - RMSNorm, precompute_freqs_cis, Attention, FeedForward, + precompute_freqs_cis, Attention, FeedForward, ) from .config import ChronosConfig +from .precision_ops import RMSNorm, lm_head_logits from .lookahead_router import LookaheadRouter from .moe_chronos import ChronosMOEFeedForward from .hybrid_attention import make_attention +from .recurrent import RecurrentController, recurrent_auxiliary_state +from .outputs import ChronosCausalLMOutput +from .multimodal import MultimodalFusion +from chronos.config_schema import resolve_recurrent_attention_mode, recurrent_cache_layout + + +def cache_sequence_length(past_key_values) -> int: + """Absolute number of tokens already in a Chronos KV cache. + + MLA entries are ``(c_kv,)``; sliding-window entries carry their absolute + position as the third element because their K/V tensors are evicted. + """ + for entry in past_key_values or (): + if entry is not None: + return int(entry[2]) if len(entry) > 2 else int(entry[0].shape[1]) + return 0 + + +def draft_exit_kl(draft_logits: torch.Tensor, full_logits: torch.Tensor, labels: torch.Tensor | None) -> torch.Tensor: + """Mean token KL(full.detach() || draft) in FP32 over supervised positions. + + Positions follow the CE shift (logit t predicts label t+1); without + labels every position counts. Teacher gradients are stopped so the + term only pulls the shallow exit toward the full-depth distribution. + """ + teacher = F.log_softmax(full_logits.detach().float(), dim=-1) + student = F.log_softmax(draft_logits.float(), dim=-1) + kl = (teacher.exp() * (teacher - student)).sum(-1) # [B, S] + if labels is None: + return kl.mean() + keep = torch.zeros_like(kl, dtype=torch.bool) + keep[:, :-1] = labels[:, 1:] != -100 + return (kl * keep).sum() / keep.sum().clamp_min(1) class ChronosBlock(nn.Module): @@ -43,24 +76,61 @@ def forward( use_cache=False, attention_mask=None, available_expert_mask=None, + latent_only: bool = False, + kv_override=None, + return_shared_kv: bool = False, ): + """One physical block. + + ``latent_only`` skips attention (v2 ``first_only`` recurrence). + ``return_shared_kv`` / ``kv_override`` implement the v3 + ``shared_first`` contract: iteration zero returns its full K/V view, + later iterations attend with fresh queries against that view. + """ residual = hidden_states - hidden_states, present_key_value = self.self_attn( - self.input_layernorm(hidden_states), - position_embeddings, - past_key_value, - use_cache, - attention_mask, - ) - hidden_states = hidden_states + residual + shared_kv = None + if latent_only: + # Middle recurrent iterations refine the latent state without + # rereading the prefix or appending another KV entry. + present_key_value = past_key_value + else: + extra = {} + if kv_override is not None or return_shared_kv: + extra = {"kv_override": kv_override, "return_shared_kv": return_shared_kv} + attn_out = self.self_attn( + self.input_layernorm(hidden_states), + position_embeddings, + past_key_value, + use_cache, + attention_mask, + **extra, + ) + if return_shared_kv: + hidden_states, present_key_value, shared_kv = attn_out + else: + hidden_states, present_key_value = attn_out + hidden_states = hidden_states + residual mlp_in = self.post_attention_layernorm(hidden_states) if isinstance(self.mlp, ChronosMOEFeedForward): hidden_states = hidden_states + self.mlp(mlp_in, available_expert_mask) else: hidden_states = hidden_states + self.mlp(mlp_in) + if return_shared_kv: + return hidden_states, present_key_value, shared_kv return hidden_states, present_key_value +def apply_repetition_penalty(logits, input_ids, penalty: float): + """CTRL-style penalty: shrink seen-token logits toward zero. + + Dividing a negative logit by ``penalty > 1`` would *raise* its + probability; multiply negatives and divide positives instead. + """ + seen = torch.zeros_like(logits, dtype=torch.bool).scatter_(1, input_ids, True) + penalized = torch.where(logits < 0, logits * penalty, logits / penalty) + return torch.where(seen, penalized, logits) + + class ChronosModel(nn.Module): """ Backbone: identical to MiniMindModel but: @@ -81,14 +151,233 @@ def __init__(self, config: ChronosConfig): ]) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.lookahead_router = LookaheadRouter(config) + self.recurrent_controller = RecurrentController(config) + # Pure-text models carry no image/audio encoder parameters or compute. + self.multimodal_fusion = MultimodalFusion(config) if config.modalities else None freqs_cos, freqs_sin = precompute_freqs_cis( dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling, ) - self.register_buffer("freqs_cos", freqs_cos, persistent=False) - self.register_buffer("freqs_sin", freqs_sin, persistent=False) + # Plain FP32 attributes (not buffers): ``model.to(bf16)`` must not + # round them. Attention layers own their own FP32 tables; these are + # kept for callers that read the model-level RoPE. + self.freqs_cos, self.freqs_sin = freqs_cos.float(), freqs_sin.float() + self.last_recurrent_state = None + + def _run_block( + self, layer, hidden_states, position_embeddings, *, past_key_value=None, + use_cache=False, attention_mask=None, available_expert_mask=None, + latent_only: bool = False, + kv_override=None, + return_shared_kv: bool = False, + ): + """Run one physical block with optional selective recomputation.""" + if ( + bool(getattr(self.config, "activation_checkpointing", False)) + and self.training and not use_cache and not return_shared_kv + ): + from torch.utils.checkpoint import checkpoint + + def transition(states): + return layer( + states, position_embeddings, past_key_value=None, + use_cache=False, attention_mask=attention_mask, + available_expert_mask=available_expert_mask, + latent_only=latent_only, kv_override=kv_override, + )[0] + + return checkpoint(transition, hidden_states, use_reentrant=False), None + return layer( + hidden_states, position_embeddings, + past_key_value=past_key_value, use_cache=use_cache, + attention_mask=attention_mask, + available_expert_mask=available_expert_mask, + latent_only=latent_only, + kv_override=kv_override, + return_shared_kv=return_shared_kv, + ) + + def _forward_recurrent( + self, + hidden_states, + position_embeddings, + past_key_values, + use_cache, + attention_mask, + available_expert_masks, + depth_override=None, + ): + """Execute Prelude → shared groups × depth → Coda. + + ``recurrent_attention_mode`` selects how loop iterations read context: + + * ``shared_first`` (v3 default): every iteration attends. Iteration + zero writes the group's single KV entry; later iterations recompute + queries from the refined state and attend to that same K/V view. + Training, prefill and decode therefore see identical keys, the KV + footprint equals the physical layer count, and the ACT exit never + leaves a cache hole. + * ``per_iteration``: every iteration owns a KV slot (virtual-layer + cache, ``prelude + groups·depth + coda`` entries). Exact but D× + the KV memory; kept as an ablation. Decode depth is fixed. + * ``first_only``: v2 semantics; iterations after zero skip attention. + """ + cfg = self.config + p = int(cfg.prelude_layers) + g = int(cfg.recurrent_groups) + c = int(cfg.coda_layers) + mode = resolve_recurrent_attention_mode(cfg) + depth = self.recurrent_controller.execution_depth(cfg, self.training) + if depth_override is not None: + # Self-speculative drafting runs fewer loop iterations. Only + # depth-independent cache layouts can share a KV stream. + if mode == "per_iteration": + raise ValueError("depth override needs a depth-independent cache (shared_first/first_only)") + depth = max(1, min(int(depth_override), int(cfg.recurrent_depth))) + layout = recurrent_cache_layout(cfg, depth) + if past_key_values is None: + past_key_values = [None] * layout.num_slots + elif len(past_key_values) != layout.num_slots: + raise ValueError( + f"recurrent cache has {len(past_key_values)} slots; mode={mode} depth={depth} " + f"expects {layout.num_slots} (decode must reuse the prefill depth)" + ) + presents = [None] * layout.num_slots + + def mask_for(li): + return available_expert_masks[li] if available_expert_masks else None + + # Prelude remains a normal one-pass stack. + for li in range(p): + hidden_states, presents[li] = self._run_block( + self.layers[li], hidden_states, position_embeddings, + past_key_value=past_key_values[li], use_cache=use_cache, + attention_mask=attention_mask, available_expert_mask=mask_for(li), + ) + # Lookahead observes the post-prelude state: earliest stable routing + # evidence, identical for every recurrent depth (draft or verify). + lookahead_source = hidden_states + + recurrent_states = [] + # Physical early exit never changes the KV a later token sees in + # shared_first/first_only (only iteration zero writes KV), so it is + # allowed during prefill and decode. per_iteration owns one slot per + # iteration: exiting would leave cache holes, so only no-cache + # forwards may exit there. + allow_adaptive_exit = mode != "per_iteration" or not use_cache + adaptive_active = self.recurrent_controller.adaptive_active(cfg, self.training) + threshold = float(cfg.halting_threshold) + min_depth = int(cfg.min_recurrent_depth) + for group in range(g): + li = p + group + layer = self.layers[li] + group_states, group_halts = [], [] + exited_early = False + shared_kv = None + halted = cumulative = None + target = min(1.0, max(1e-6, threshold)) + for step in range(depth): + conditioned = self.recurrent_controller.condition(hidden_states, step, group) + slot = layout.slot(li, step) + block_kwargs = dict( + use_cache=use_cache, attention_mask=attention_mask, + available_expert_mask=mask_for(li), + ) + if mode == "per_iteration": + refined, presents[slot] = self._run_block( + layer, conditioned, position_embeddings, + past_key_value=past_key_values[slot], **block_kwargs) + elif mode == "first_only": + refined, present = self._run_block( + layer, conditioned, position_embeddings, + past_key_value=past_key_values[slot], latent_only=step > 0, **block_kwargs) + if step == 0: + presents[slot] = present + elif step == 0: # shared_first, writer iteration + refined, presents[slot], shared_kv = self._run_block( + layer, conditioned, position_embeddings, + past_key_value=past_key_values[slot], return_shared_kv=True, **block_kwargs) + else: # shared_first, reader iteration + refined, _ = self._run_block( + layer, conditioned, position_embeddings, + past_key_value=None, kv_override=shared_kv, + use_cache=False, attention_mask=attention_mask, + available_expert_mask=mask_for(li)) + hidden_states = self.recurrent_controller.contract(hidden_states, refined, group) + group_states.append(hidden_states) + halt_probability = self.recurrent_controller.halt_probability(hidden_states) + group_halts.append(halt_probability) + # Exact ACT bookkeeping (same rule as recurrent_auxiliary_state): + # a token halts once its accumulated mass reaches the threshold + # after min_depth. Later iterations get zero mixture weight, so + # once *every* token has halted the remaining iterations cannot + # change the output; exiting is then exact, and each token's + # result is independent of what else is in the batch (required + # for speculative verification to equal token-by-token decode). + if step + 1 >= min_depth: + clipped = halt_probability.float().clamp(0.0, 1.0) + if halted is None: + cumulative = torch.zeros_like(clipped) + halted = torch.zeros_like(clipped, dtype=torch.bool) + newly = (~halted) & ((cumulative + clipped) >= target) + cumulative = torch.where(halted | newly, cumulative, (cumulative + clipped).clamp(max=1.0)) + halted = halted | newly + if ( + not self.training and adaptive_active and allow_adaptive_exit + and halted is not None and step + 1 < depth + ): + if bool(halted.all()): # one fused device→host read per step + exited_early = True + break + # ACT is state interpolation rather than a ragged Python loop. + mixed, state = recurrent_auxiliary_state( + group_states, group_halts, min_depth=min_depth, threshold=threshold, + ) + if adaptive_active: + hidden_states = mixed + else: + hidden_states = group_states[-1] + state.ponder_cost = hidden_states.new_zeros(()) + state.depth_used = torch.full_like(state.depth_used, float(len(group_states))) + state.exited_early = exited_early + recurrent_states.append(state) + + # Coda uses the remaining physical layers. + for ci in range(c): + li = p + g + ci + slot = layout.slot(li, 0) + hidden_states, presents[slot] = self._run_block( + self.layers[li], hidden_states, position_embeddings, + past_key_value=past_key_values[slot], use_cache=use_cache, + attention_mask=attention_mask, available_expert_mask=mask_for(li), + ) + + if recurrent_states: + max_depth = max(int(s.halting_probs.shape[-1]) for s in recurrent_states) + padded_halts = [ + F.pad(s.halting_probs, (0, max_depth - int(s.halting_probs.shape[-1])), value=1.0) + for s in recurrent_states + ] + self.last_recurrent_state = { + "halting_probs": torch.stack(padded_halts, dim=-2), + "ponder_loss": torch.stack([s.ponder_cost for s in recurrent_states]).mean(), + "stability_loss": torch.stack([s.stability_loss for s in recurrent_states]).mean(), + "distill_loss": torch.stack([s.distill_loss for s in recurrent_states]).mean(), + "contraction": torch.stack([s.contraction for s in recurrent_states]).mean(), + # Keep the group axis so layer/group-aware schedulers can + # correlate ponder depth with their cache keys. + "depth_used": torch.stack([s.depth_used for s in recurrent_states], dim=-1), + "early_exit_ratio": hidden_states.new_tensor( + sum(int(s.exited_early) for s in recurrent_states) / max(len(recurrent_states), 1) + ), + "attention_mode": mode, + "executed_depth": depth, + } + else: + self.last_recurrent_state = None + return hidden_states, presents, lookahead_source def _ensure_rope_cache(self, needed_len: int, device: torch.device) -> None: """Grow RoPE buffers when generation exceeds the checkpoint context.""" @@ -116,22 +405,32 @@ def forward( past_key_values=None, use_cache=False, available_expert_masks=None, # list[Tensor|None] per layer + modality_inputs=None, + modality_mask=None, + recurrent_depth=None, **kwargs, ): B, S = input_ids.shape + for layer in self.layers: + if isinstance(layer.mlp, ChronosMOEFeedForward): + layer.mlp.reset_iteration_routing() if hasattr(past_key_values, 'layers'): past_key_values = None - past_key_values = past_key_values or [None] * len(self.layers) - # MLA caches only c_kv tuple; SlidingWindow caches (k, v) tuple - # Determine start_pos from first non-None past_kv - start_pos = 0 - for pkv in past_key_values: - if pkv is not None: - start_pos = pkv[0].shape[1] - break + # MLA caches only c_kv; SlidingWindow adds an absolute position counter. + start_pos = cache_sequence_length(past_key_values) self._ensure_rope_cache(start_pos + S, input_ids.device) hidden_states = self.dropout(self.embed_tokens(input_ids)) + modality_states = {} + if self.multimodal_fusion is not None: + hidden_states, modality_states = self.multimodal_fusion( + hidden_states, modality_inputs, modality_mask + ) + elif modality_inputs: + raise ValueError( + "modality_inputs were provided, but this checkpoint has no enabled modalities; " + "set config.modalities before constructing the model" + ) # Pass full freqs so hybrid attention layers can slice as needed position_embeddings = ( self.freqs_cos, @@ -141,26 +440,58 @@ def forward( presents = [] lookahead_probs = None # [B, S, K, E] from block 0 - for i, (layer, past_kv) in enumerate(zip(self.layers, past_key_values)): - mask = available_expert_masks[i] if available_expert_masks else None - hidden_states, present = layer( - hidden_states, position_embeddings, - past_key_value=past_kv, - use_cache=use_cache, - attention_mask=attention_mask, - available_expert_mask=mask, + # A one-loop model is routed through the proven v1 graph. The loop + # controller is an exact no-op in this migration mode, including the + # original lookahead observation point after physical block zero. + recurrent_active = ( + bool(getattr(self.config, "recurrent_enabled", False)) + and int(getattr(self.config, "recurrent_depth", 1)) > 1 + ) + if recurrent_active: + hidden_states, presents, lookahead_source = self._forward_recurrent( + hidden_states, position_embeddings, past_key_values, use_cache, + attention_mask, available_expert_masks, depth_override=recurrent_depth, ) - presents.append(present) - # Insert LookaheadRouter after block 0 - if i == 0: - lookahead_probs = self.lookahead_router(hidden_states) + # Post-prelude state: earliest stable routing observation. + lookahead_probs = self.lookahead_router(lookahead_source) + else: + self.last_recurrent_state = None + past_key_values = past_key_values or [None] * len(self.layers) + for i, (layer, past_kv) in enumerate(zip(self.layers, past_key_values)): + mask = available_expert_masks[i] if available_expert_masks else None + hidden_states, present = self._run_block( + layer, + hidden_states, position_embeddings, + past_key_value=past_kv, + use_cache=use_cache, + attention_mask=attention_mask, + available_expert_mask=mask, + ) + presents.append(present) + if i == 0: + lookahead_probs = self.lookahead_router(hidden_states) hidden_states = self.norm(hidden_states) aux_loss = sum( - [l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, ChronosMOEFeedForward)], + [l.mlp.aux_loss for l in self.layers + if isinstance(l.mlp, ChronosMOEFeedForward) and l.mlp.aux_loss is not None], hidden_states.new_zeros(1).squeeze(), ) - return hidden_states, presents, aux_loss, lookahead_probs + return hidden_states, presents, aux_loss, lookahead_probs, modality_states + + +class RecurrentChronosModel(ChronosModel): + """Explicit v2 backbone name for runtime/checkpoint dispatch. + + ``ChronosModel`` remains importable as the v1-compatible executor. The + subclass intentionally shares parameter names so v1→depth-1 migration is + identity-preserving and does not require tensor-key rewrites. + """ + + def __init__(self, config: ChronosConfig): + super().__init__(config) + if int(getattr(config, "architecture_version", 2)) < 2: + raise ValueError("RecurrentChronosModel requires architecture_version >= 2") class ChronosForCausalLM(PreTrainedModel, GenerationMixin): @@ -184,8 +515,14 @@ def all_tied_weights_keys(self) -> dict: def __init__(self, config: ChronosConfig = None): self.config = config or ChronosConfig() super().__init__(self.config) - self.model = ChronosModel(self.config) + model_cls = ( + RecurrentChronosModel + if int(getattr(self.config, "architecture_version", 2)) >= 2 + else ChronosModel + ) + self.model = model_cls(self.config) self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False) + self.last_draft_exit_kl = None if self.config.tie_word_embeddings: self.model.embed_tokens.weight = self.lm_head.weight @@ -198,28 +535,86 @@ def forward( logits_to_keep=0, labels=None, available_expert_masks=None, + modality_inputs=None, + modality_mask=None, + legacy_output=False, + recurrent_depth=None, + logits_dtype=None, **kwargs, ): - hidden_states, past_key_values, aux_loss, lookahead_probs = self.model( + """``logits_dtype``: dtype the logits are born in. Defaults to FP32 + on loss paths (``labels`` given, or training mode: DPO/ORPO/GRPO/KD + take log-softmax of the logits) and to the compute dtype for + inference, where sampling upcasts only the kept [B, V] rows. + + With ``draft_exit_kl_coef > 0`` a training forward with labels also + runs the model at ``draft_exit_depth`` first and stores + KL(full.detach() || draft) on ``self.last_draft_exit_kl`` (consumed by + ``chronos_loss_term``); the full-depth pass runs last so routing and + recurrent diagnostics describe the full model.""" + self.last_draft_exit_kl = None + draft_logits = None + if self._draft_exit_active(labels, past_key_values, use_cache, recurrent_depth): + draft_hidden = self.model( + input_ids, attention_mask, None, False, + available_expert_masks=available_expert_masks, + modality_inputs=modality_inputs, modality_mask=modality_mask, + recurrent_depth=int(self.config.draft_exit_depth), + )[0] + draft_logits = lm_head_logits(draft_hidden, self.lm_head.weight, torch.float32) + hidden_states, past_key_values, aux_loss, lookahead_probs, modality_states = self.model( input_ids, attention_mask, past_key_values, use_cache, available_expert_masks=available_expert_masks, + modality_inputs=modality_inputs, + modality_mask=modality_mask, + recurrent_depth=recurrent_depth, ) slice_indices = slice(-logits_to_keep, None) if logits_to_keep > 0 else slice(None) - logits = self.lm_head(hidden_states[:, slice_indices, :]) + if logits_dtype is None and (labels is not None or self.training): + logits_dtype = torch.float32 + logits = lm_head_logits(hidden_states[:, slice_indices, :], self.lm_head.weight, logits_dtype) loss = None if labels is not None: x = logits[..., :-1, :].contiguous() y = labels[..., 1:].contiguous() loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100) + if draft_logits is not None: + draft_logits = draft_logits[:, slice_indices, :] + kl_labels = labels[:, -logits.shape[1]:] if labels.shape[1] != logits.shape[1] else labels + self.last_draft_exit_kl = draft_exit_kl(draft_logits, logits, kl_labels) - return MoeCausalLMOutputWithPast( + rs = self.model.last_recurrent_state or {} + result = ChronosCausalLMOutput( loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states, - ), lookahead_probs + lookahead_probs=lookahead_probs, + halting_probs=rs.get("halting_probs"), + ponder_loss=rs.get("ponder_loss"), + recurrent_stability_loss=rs.get("stability_loss"), + recurrent_distill_loss=rs.get("distill_loss"), + recurrent_contraction=rs.get("contraction"), + recurrent_depth_used=rs.get("depth_used"), + cache_hints={ + "layer_aware": bool(getattr(self.config, "layer_aware_cache", False)), + "early_exit_ratio": rs.get("early_exit_ratio"), + }, + modality_states=modality_states, + ) + return result.legacy_tuple() if legacy_output else result + + def _draft_exit_active(self, labels, past_key_values, use_cache, recurrent_depth) -> bool: + cfg = self.config + if not (self.training and labels is not None and past_key_values is None and not use_cache): + return False + if float(getattr(cfg, "draft_exit_kl_coef", 0.0) or 0.0) <= 0.0 or recurrent_depth is not None: + return False + if not getattr(cfg, "recurrent_enabled", True) or int(cfg.recurrent_depth) <= 1: + return False + return 1 <= int(getattr(cfg, "draft_exit_depth", 1)) < int(cfg.recurrent_depth) @torch.inference_mode() def generate( @@ -236,6 +631,8 @@ def generate( do_sample=True, repetition_penalty=1.0, available_expert_masks=None, + modality_inputs=None, + modality_mask=None, **kwargs, ): input_ids = kwargs.pop("input_ids", inputs) @@ -245,22 +642,25 @@ def generate( streamer.put(input_ids.cpu()) for _ in range(max_new_tokens): - past_len = past_key_values[0][0].shape[1] if past_key_values else 0 - outputs, _ = self.forward( + past_len = cache_sequence_length(past_key_values) if past_key_values else 0 + result = self.forward( input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, available_expert_masks=available_expert_masks, + modality_inputs=modality_inputs, + modality_mask=modality_mask, ) + outputs, _ = result.legacy_tuple() if attention_mask is not None: attention_mask = torch.cat( [attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1 ) - logits = outputs.logits[:, -1, :] / temperature + logits = outputs.logits[:, -1, :].float() / temperature # [B, V] sampling row only if repetition_penalty != 1.0: - for i in range(input_ids.shape[0]): - logits[i, torch.unique(input_ids[i])] /= repetition_penalty + logits = apply_repetition_penalty(logits, input_ids, repetition_penalty) if top_k > 0: - logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf') + k = min(int(top_k), int(logits.shape[-1])) + logits[logits < torch.topk(logits, k)[0][..., -1, None]] = -float('inf') if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True) mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p diff --git a/chronos/model/moe_chronos.py b/chronos/model/moe_chronos.py index f2b5ee9..b637c7f 100644 --- a/chronos/model/moe_chronos.py +++ b/chronos/model/moe_chronos.py @@ -29,6 +29,13 @@ def forward(self, _x): raise RuntimeError("lazy expert placeholder executed before materialization") +def _run_expert(expert: nn.Module, x: torch.Tensor) -> torch.Tensor: + """Run a materialized expert. Expert pages are stored and promoted in the + model's compute dtype (``ExpertStore`` follows the trunk), so there is no + cast here; a mismatch is a loading bug and fails loudly.""" + return expert(x) + + class ChronosMOEFeedForward(nn.Module): """ Extends MOEFeedForward with: @@ -78,6 +85,17 @@ def __init__(self, config: ChronosConfig): # actually shape the gate. self.last_router_probs: torch.Tensor = None self.last_router_probs_grad: torch.Tensor = None + # A recurrent block calls this layer once per loop iteration; every + # call is a real dispatch decision, so routing statistics and the + # load-balance term must cover all of them (reset per model forward). + self.iteration_router_probs: list[torch.Tensor] = [] + self.iteration_router_probs_grad: list[torch.Tensor] = [] + self._iteration_aux_raw: list[torch.Tensor] = [] + # router_scope="shared": iterations after the first reuse the first + # iteration's expert ids (weights are re-gathered from the current + # scores, so the gate still trains every iteration). + self.router_scope = str(getattr(config, "router_scope", "per_iteration") or "per_iteration") + self._shared_topk_idx: torch.Tensor | None = None self.fallback_mask_prob = float(getattr(config, "fallback_mask_prob", 0.0)) self.runtime_miss_policy = "fallback_diagnostic" self.runtime_on_demand_loader = None @@ -87,6 +105,18 @@ def __init__(self, config: ChronosConfig): self.last_on_demand_weight_mass = 0.0 self.last_fallback_weight_mass = 0.0 + def reset_iteration_routing(self) -> None: + self.iteration_router_probs = [] + self.iteration_router_probs_grad = [] + self._iteration_aux_raw = [] + self._shared_topk_idx = None + + def _record_aux(self, aux_raw: torch.Tensor) -> None: + self._iteration_aux_raw.append(aux_raw) + # Mean over loop iterations keeps the scale of a single dispatch. + self.aux_loss_raw = torch.stack(self._iteration_aux_raw).mean() + self.aux_loss = self.aux_loss_raw * self.router_aux_loss_coef + def _shared_expert_output(self, x_flat: torch.Tensor) -> torch.Tensor: """Average output of all shared experts.""" out = sum(e(x_flat) for e in self.shared_experts) @@ -109,7 +139,8 @@ def forward( """ B, S, H = x.shape x_flat = x.view(-1, H) # [N, H] - scores = F.softmax(self.gate(x_flat), dim=-1) # [N, E] + # Router softmax in FP32: an [N, E] island (E is small), mirrors MLX. + scores = F.softmax(self.gate(x_flat).float(), dim=-1) # [N, E] fp32 scores_bse = scores.view(B, S, self.num_experts) self.last_on_demand_experts = [] self.last_fallback_experts = [] @@ -120,12 +151,23 @@ def forward( # Detached view for read-only consumers (logging, cluster layout, # inference-engine prefetch). Same tensor data, no autograd edge. self.last_router_probs = scores_bse.detach() + self.iteration_router_probs.append(self.last_router_probs) + self.iteration_router_probs_grad.append(scores_bse) - topk_weight, topk_idx = torch.topk( - scores, k=self.num_experts_per_tok, dim=-1, sorted=False - ) + shared_idx = self._shared_topk_idx + if (self.router_scope == "shared" and shared_idx is not None + and shared_idx.shape[0] == scores.shape[0] and len(self.iteration_router_probs) > 1): + topk_idx = shared_idx + topk_weight = scores.gather(-1, topk_idx) + else: + topk_weight, topk_idx = torch.topk( + scores, k=self.num_experts_per_tok, dim=-1, sorted=False + ) + if self.router_scope == "shared": + self._shared_topk_idx = topk_idx if self.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20) + topk_weight = topk_weight.to(x_flat.dtype) # one [N, K] cast; expert sums stay in compute dtype y = torch.zeros_like(x_flat) @@ -145,13 +187,23 @@ def forward( if available_expert_mask is None and training_fallback_mask is None: # --- Standard path: routed experts + always-on shared residual --- + # Sorted dispatch: one argsort + one host read per layer instead + # of an ``any()`` sync per expert. Each token runs only its + # top-k experts; rows are contiguous per expert. dead_param_sum = None + flat_idx = topk_idx.reshape(-1) + order = torch.argsort(flat_idx, stable=True) + counts = torch.bincount(flat_idx, minlength=self.num_experts).tolist() + rows = order // self.num_experts_per_tok + row_weight = topk_weight.reshape(-1)[order].unsqueeze(-1) + offset = 0 for i, expert in enumerate(self.experts): - mask = (topk_idx == i) - if mask.any(): - token_idx = mask.any(dim=-1).nonzero().flatten() - weight = topk_weight[mask].view(-1, 1) - y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype)) + count = counts[i] + if count: + token_idx = rows[offset:offset + count] + weight = row_weight[offset:offset + count] + y.index_add_(0, token_idx, expert(x_flat[token_idx]) * weight) + offset += count elif self.training: # Dead-expert handling: keep the autograd edge alive # WITHOUT touching y[0,0] (the previous trick poisoned @@ -170,7 +222,7 @@ def forward( # weights, which is the dominant root cause of "loss looks # fine but generations are gibberish." if len(self.shared_experts) > 0: - y = y + self._shared_expert_output(x_flat).to(y.dtype) + y = y + self._shared_expert_output(x_flat) if self.training: # Raw, un-scaled load balance (model exposes raw; trainer scales). @@ -178,11 +230,9 @@ def forward( aux_raw = (load * scores.mean(0)).sum() * self.num_experts if dead_param_sum is not None: aux_raw = aux_raw + 0.0 * dead_param_sum - self.aux_loss_raw = aux_raw - # Backwards-compat: the legacy `aux_loss` field stays - # populated with the pre-scaled value some callers still - # read (chronos_loss_term now prefers aux_loss_raw). - self.aux_loss = aux_raw * self.router_aux_loss_coef + # ``aux_loss`` (pre-scaled) stays for legacy readers; + # chronos_loss_term prefers the iteration-mean ``aux_loss_raw``. + self._record_aux(aux_raw) else: self.aux_loss_raw = scores.new_zeros(()) self.aux_loss = scores.new_zeros(()) @@ -203,22 +253,47 @@ def forward( if len(self.shared_experts) > 0 else None ) dead_param_sum = None + # Sorted dispatch with ONE host read per call: selection counts, + # routed mass and availability travel together. The previous + # per-expert ``mask.nonzero()`` was a device sync per selected + # expert (E x virtual layers per forward; dominant on MPS/CUDA). + # Rows per expert stay in ascending token order (stable sort), so + # values and the index_add_ order are unchanged. + flat_idx = topk_idx.reshape(-1) + order = torch.argsort(flat_idx, stable=True) + sorted_rows = order // self.num_experts_per_tok + sorted_weight = topk_weight.reshape(-1)[order].unsqueeze(-1) + mass = torch.zeros(self.num_experts, dtype=torch.float32, device=x_flat.device).index_add_( + 0, flat_idx, topk_weight.detach().reshape(-1).float()) + bins = torch.bincount(flat_idx, minlength=self.num_experts) + host = torch.stack([bins.to(torch.float32), mass, (avail > 0).to(torch.float32)]).tolist() + # FP32 counts are exact below 2**24 assignments; beyond that read + # them as integers (the queue is already drained: no extra stall). + counts = [int(c) for c in host[0]] if flat_idx.numel() < 2 ** 24 else bins.tolist() + mass_host = host[1] + avail_host = [v > 0 for v in host[2]] + offset = 0 for i, expert in enumerate(self.experts): - mask = (topk_idx == i) # [N, K] bool - expert_selected = mask.any() - if not expert_selected: + if not counts[i]: if self.training: - contrib = sum(p.sum() for p in expert.parameters()) - dead_param_sum = ( - contrib if dead_param_sum is None else dead_param_sum + contrib - ) + # Paged training may keep a cold expert's parameters + # on CPU while this layer executes on CUDA/MPS. A + # cross-device 0*parameter edge is invalid and the + # cold page intentionally receives no update anyway. + live_params = [p for p in expert.parameters() if p.device == x_flat.device] + if live_params: + contrib = sum(p.sum() for p in live_params) + dead_param_sum = ( + contrib if dead_param_sum is None else dead_param_sum + contrib + ) continue - token_idx = mask.any(dim=-1).nonzero().flatten() # [T] - weight = topk_weight[mask].view(-1, 1) # [T, 1] - selected_mass = float(weight.detach().to(torch.float32).sum().item()) + token_idx = sorted_rows[offset:offset + counts[i]] # [T] ascending + weight = sorted_weight[offset:offset + counts[i]] # [T, 1] + offset += counts[i] + selected_mass = float(mass_host[i]) expert_live = not isinstance(self.experts[i], LazyExpertPlaceholder) - expert_available = float(avail[i].item()) > 0.0 + expert_available = bool(avail_host[i]) inference_stale_mask = ( available_expert_mask is not None and not self.training @@ -231,7 +306,7 @@ def forward( if not expert_available and available_expert_mask is not None and not self.training: self.last_on_demand_experts.append(int(i)) self.last_on_demand_weight_mass += selected_mass - blended = self.experts[i](x_flat[token_idx]) # [T, H] + blended = _run_expert(self.experts[i], x_flat[token_idx]) # [T, H] else: loaded = False loader = getattr(self, "runtime_on_demand_loader", None) @@ -250,8 +325,11 @@ def forward( if loaded and not isinstance(self.experts[i], LazyExpertPlaceholder): self.last_on_demand_experts.append(int(i)) self.last_on_demand_weight_mass += selected_mass - blended = self.experts[i](x_flat[token_idx]) + blended = _run_expert(self.experts[i], x_flat[token_idx]) else: + if (not self.training and + self.runtime_miss_policy == "sync_on_demand"): + raise RuntimeError(f"strict offload could not load expert {i}") self.last_fallback_experts.append(int(i)) self.last_fallback_weight_mass += selected_mass blended = ( @@ -259,21 +337,20 @@ def forward( if shared_out_full is not None else x_flat.new_zeros((token_idx.numel(), H)) ) - y.index_add_(0, token_idx, (blended * weight).to(y.dtype)) + y.index_add_(0, token_idx, blended * weight) # The shared expert is a residual component, not only the cold # fallback. This makes available_mask=all_true equivalent to the # normal path and keeps SSD/DRAM inference faithful to training. if shared_out_full is not None: - y = y + shared_out_full.to(y.dtype) + y = y + shared_out_full if self.training: load = F.one_hot(topk_idx, self.num_experts).float().mean(0) aux_raw = (load * scores.mean(0)).sum() * self.num_experts if dead_param_sum is not None: aux_raw = aux_raw + 0.0 * dead_param_sum - self.aux_loss_raw = aux_raw - self.aux_loss = aux_raw * self.router_aux_loss_coef + self._record_aux(aux_raw) else: self.aux_loss_raw = scores.new_zeros(()) self.aux_loss = scores.new_zeros(1).squeeze() diff --git a/chronos/model/multimodal.py b/chronos/model/multimodal.py new file mode 100644 index 0000000..922d2f0 --- /dev/null +++ b/chronos/model/multimodal.py @@ -0,0 +1,216 @@ +"""Dependency-light image/audio encoders and text fusion for Chronos v2.""" +from __future__ import annotations + +from typing import Dict, Optional + +import torch +import torch.nn as nn + + +class PatchImageEncoder(nn.Module): + def __init__(self, hidden_size: int, patch_size: int = 16, max_tokens: int = 196): + super().__init__() + self.patch = nn.Conv2d(3, hidden_size, patch_size, stride=patch_size, bias=False) + self.max_tokens = int(max_tokens) + self.norm = nn.LayerNorm(hidden_size) + + def forward(self, images: torch.Tensor) -> torch.Tensor: + if images.dim() == 5: # [B,N,C,H,W] -> first image axis flattened as tokens + b, n, c, h, w = images.shape + images = images.reshape(b * n, c, h, w) + x = self.patch(images).flatten(2).transpose(1, 2) + x = x.reshape(b, n * x.shape[1], x.shape[2]) + else: + x = self.patch(images).flatten(2).transpose(1, 2) + return self.norm(x[:, : self.max_tokens]) + + +class ConvAudioEncoder(nn.Module): + def __init__(self, hidden_size: int, max_tokens: int = 128): + super().__init__() + mid = max(32, hidden_size // 2) + self.net = nn.Sequential( + nn.Conv1d(1, mid, 400, stride=160, padding=120), nn.SiLU(), + nn.Conv1d(mid, hidden_size, 5, stride=2, padding=2), nn.SiLU(), + ) + self.max_tokens = int(max_tokens) + self.norm = nn.LayerNorm(hidden_size) + + def forward(self, audio: torch.Tensor) -> torch.Tensor: + if audio.dim() == 2: + audio = audio.unsqueeze(1) + x = self.net(audio.float()).transpose(1, 2) + return self.norm(x[:, : self.max_tokens]) + + +class MultimodalFusion(nn.Module): + """Cross-attend text states to lazily present modality tokens.""" + def __init__(self, config): + super().__init__() + h = int(config.hidden_size) + heads = max(1, min(int(config.modality_fusion_heads), int(config.num_attention_heads))) + while h % heads: + heads -= 1 + enabled = set(getattr(config, "modalities", []) or []) + self.image_encoder = ( + PatchImageEncoder(h, config.image_patch_size, config.max_image_tokens) + if "image" in enabled else None + ) + self.audio_encoder = ( + ConvAudioEncoder(h, config.max_audio_tokens) + if "audio" in enabled else None + ) + self.hidden_size = h + self.modality_embeddings = nn.ParameterDict({ + "image": nn.Parameter(torch.zeros(1, 1, h)), + "audio": nn.Parameter(torch.zeros(1, 1, h)), + }) + self.attn = nn.MultiheadAttention(h, heads, dropout=config.dropout, batch_first=True) + self.gate = nn.Parameter(torch.tensor(0.0)) + self.dropout = nn.Dropout(float(config.modality_dropout)) + if bool(config.freeze_modality_encoders): + for encoder in (self.image_encoder, self.audio_encoder): + if encoder is None: + continue + for p in encoder.parameters(): + p.requires_grad_(False) + + def forward( + self, + text_states: torch.Tensor, + modality_inputs: Optional[Dict[str, torch.Tensor]], + modality_mask: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, dict]: + if not modality_inputs: + return text_states, {} + tokens = [] + token_modalities = [] + token_validity = [] + states = {} + images = modality_inputs.get("image") + def _fit_width(features: torch.Tensor) -> torch.Tensor: + """Project arbitrary cached feature widths without new modules. + + Exported CLIP/Whisper features often differ from Chronos hidden + size. Deterministic pad/truncate keeps serving lazy and avoids + creating uninitialized LazyLinear parameters when only one + modality is present in a checkpoint. + """ + width = features.shape[-1] + if width == self.hidden_size: + return features + if width > self.hidden_size: + return features[..., :self.hidden_size] + return torch.nn.functional.pad(features, (0, self.hidden_size - width)) + + image_embeds = modality_inputs.get("image_embeds") + if image_embeds is not None: + image_states = image_embeds.to(text_states.device, text_states.dtype) + if image_states.dim() == 2: + image_states = image_states.unsqueeze(1) + image_states = _fit_width(image_states) + image_states = image_states + self.modality_embeddings["image"] + tokens.append(image_states) + token_modalities.append("image") + present = modality_inputs.get("image_embeds_mask") + token_validity.append(present) + states["image"] = image_states + if images is not None and image_embeds is None: + if self.image_encoder is None: + raise ValueError("image input provided but image modality is disabled") + image_states = self.image_encoder(images.to(text_states.device)) + image_states = image_states + self.modality_embeddings["image"] + tokens.append(image_states) + token_modalities.append("image") + token_validity.append(modality_inputs.get("image_mask")) + states["image"] = ( + torch.cat([states["image"], image_states], dim=1) + if "image" in states else image_states + ) + audio = modality_inputs.get("audio") + audio_embeds = modality_inputs.get("audio_embeds") + if audio_embeds is not None: + audio_states = audio_embeds.to(text_states.device, text_states.dtype) + if audio_states.dim() == 2: + audio_states = audio_states.unsqueeze(1) + audio_states = _fit_width(audio_states) + audio_states = audio_states + self.modality_embeddings["audio"] + tokens.append(audio_states) + token_modalities.append("audio") + token_validity.append(modality_inputs.get("audio_embeds_mask")) + states["audio"] = audio_states + if audio is not None and audio_embeds is None: + if self.audio_encoder is None: + raise ValueError("audio input provided but audio modality is disabled") + audio_states = self.audio_encoder(audio.to(text_states.device)) + audio_states = audio_states + self.modality_embeddings["audio"] + tokens.append(audio_states) + token_modalities.append("audio") + token_validity.append(modality_inputs.get("audio_mask")) + states["audio"] = ( + torch.cat([states["audio"], audio_states], dim=1) + if "audio" in states else audio_states + ) + if not tokens: + return text_states, {} + memory = self.dropout(torch.cat(tokens, dim=1)) + # A mixed batch contains zero-padded media for samples that do not + # carry that modality. Use a per-modality key padding mask so those + # rows cannot attend to another sample's placeholder representation. + # Scalar masks from v1 remain valid and mean "any modality present". + if modality_mask is None: + active = torch.ones( + (text_states.shape[0], 2), device=text_states.device, dtype=torch.bool + ) + active[:, 0] = "image" in token_modalities + active[:, 1] = "audio" in token_modalities + else: + raw_mask = modality_mask.to(text_states.device) + if raw_mask.dim() == 1: + active = raw_mask.view(-1, 1).expand(-1, 2) > 0 + # Preserve the exact v1 behavior for a scalar batch mask. + if raw_mask.numel() == 1 and text_states.shape[0] > 1: + active = raw_mask.reshape(1, 1).expand(text_states.shape[0], 2) > 0 + elif raw_mask.shape[-1] == 2: + active = raw_mask > 0 + else: + raise ValueError("modality_mask must be [B] or [B,2] (image,audio)") + padding_parts = [] + for modality, modality_tokens, present in zip(token_modalities, tokens, token_validity): + col = 0 if modality == "image" else 1 + valid = active[:, col] + if present is not None: + present = torch.as_tensor(present, device=text_states.device, dtype=torch.bool) + if present.dim() == 2: + if present.shape[0] != text_states.shape[0]: + raise ValueError(f"{modality} modality mask batch dimension mismatch") + # Padding mask is token-granular for pre-encoded features; + # preserve it instead of reducing all rows to one flag. + if present.shape[1] != modality_tokens.shape[1]: + present = present[:, :modality_tokens.shape[1]] + if present.shape[1] < modality_tokens.shape[1]: + present = torch.nn.functional.pad( + present, (0, modality_tokens.shape[1] - present.shape[1]) + ) + padding_parts.append(~present) + continue + present = present.flatten() + if present.numel() == 1: + present = present.expand(text_states.shape[0]) + if present.numel() != text_states.shape[0]: + raise ValueError(f"{modality} modality mask batch dimension mismatch") + valid = valid & present + padding_parts.append(~valid.unsqueeze(1).expand(-1, modality_tokens.shape[1])) + key_padding_mask = torch.cat(padding_parts, dim=1) if padding_parts else None + # MultiheadAttention does not define all-masked rows (softmax NaN). + # Keep one harmless key for those rows and zero their fused output. + represented_cols = sorted({0 if m == "image" else 1 for m in token_modalities}) + no_media = ~active[:, represented_cols].any(dim=1) + if key_padding_mask is not None and no_media.any(): + key_padding_mask[no_media, 0] = False + fused, _ = self.attn( + text_states, memory, memory, key_padding_mask=key_padding_mask, need_weights=False + ) + if no_media.any(): + fused = fused.masked_fill(no_media.view(-1, 1, 1), 0.0) + return text_states + torch.sigmoid(self.gate) * fused, states diff --git a/chronos/model/outputs.py b/chronos/model/outputs.py new file mode 100644 index 0000000..0f9bad1 --- /dev/null +++ b/chronos/model/outputs.py @@ -0,0 +1,62 @@ +"""Stable public outputs for Chronos v2.""" +from __future__ import annotations + +from dataclasses import dataclass +import os +from typing import Any, Optional + +import torch +from transformers.modeling_outputs import MoeCausalLMOutputWithPast +from transformers.utils import ModelOutput + + +@dataclass +class ChronosCausalLMOutput(ModelOutput): + loss: Optional[torch.Tensor] = None + logits: Optional[torch.Tensor] = None + past_key_values: Optional[Any] = None + hidden_states: Optional[Any] = None + attentions: Optional[Any] = None + aux_loss: Optional[torch.Tensor] = None + lookahead_probs: Optional[torch.Tensor] = None + halting_probs: Optional[torch.Tensor] = None + ponder_loss: Optional[torch.Tensor] = None + recurrent_stability_loss: Optional[torch.Tensor] = None + recurrent_distill_loss: Optional[torch.Tensor] = None + recurrent_contraction: Optional[torch.Tensor] = None + recurrent_depth_used: Optional[torch.Tensor] = None + cache_hints: Optional[dict] = None + modality_states: Optional[dict] = None + + @property + def recurrent_state(self) -> dict: + """Return recurrent diagnostics in a stable, serialization-friendly + mapping for trainers and serving adapters.""" + return { + "halting_probs": self.halting_probs, + "ponder_loss": self.ponder_loss, + "stability_loss": self.recurrent_stability_loss, + "distill_loss": self.recurrent_distill_loss, + "contraction": self.recurrent_contraction, + "depth_used": self.recurrent_depth_used, + "cache_hints": self.cache_hints, + } + + def legacy_tuple(self): + primary = MoeCausalLMOutputWithPast( + loss=self.loss, + aux_loss=self.aux_loss, + logits=self.logits, + past_key_values=self.past_key_values, + hidden_states=self.hidden_states, + attentions=self.attentions, + ) + return primary, self.lookahead_probs + + def __iter__(self): + # The canonical contract follows HuggingFace ModelOutput iteration. + # A temporary opt-in bridge exists for external v1 scripts; in-repo + # callers use fields or ``legacy_output=True``/``legacy_tuple()``. + if os.environ.get("CHRONOS_LEGACY_OUTPUT_ITER", "1") == "1": + return iter(self.legacy_tuple()) + return super().__iter__() diff --git a/chronos/model/planner.py b/chronos/model/planner.py new file mode 100644 index 0000000..b3429c3 --- /dev/null +++ b/chronos/model/planner.py @@ -0,0 +1,673 @@ +"""Adaptive training planner. + +``plan_training(PlanRequest)`` turns a handful of user targets (total +parameters, MoE on/off, expert count, active parameters *or* active experts, +memory-saving level, optional dataset paths, hardware memory) into a complete +Chronos v3 configuration plus independent hyperparameters per training stage. + +How each number is derived (the report repeats this per value): + +* **Architecture** — closed-form parameter counts (``_count``) searched over + width / depth / expert width, then the best candidate is re-counted on the + ``meta`` device with ``factory.count_parameters`` (the allocation truth). +* **Pretrain LR and batch** — DeepSeek LLM hyperparameter scaling laws + (arXiv 2401.02954, §3.1): ``η_opt = 0.3118·C^-0.125`` and + ``B_opt = 0.2920·C^0.3271`` tokens, with ``C = M·D``, where ``M`` is the + non-embedding FLOPs/token ``72·L·h² + 12·L·h·s`` (L counts executed + virtual layers) and ``D`` the planned token count. The law was fitted on + dense models over C≈1e17–2e19; below / outside that range it is an + extrapolation, so results are clamped and labelled ``estimated``. +* **Later stages** — conventional ratios to the pretrain LR (SFT ≈ 1/3, + preference stages ≈ 1/10–1/20, distill ≈ 1/3); these are heuristics. +* **Memory** — weights / optimizer state from the parameter count; the + activation term uses ``ACT_BYTES_PER_TOKEN_LAYER_H`` calibrated on the live + 50M MLX run (measured peak 8.64 GB at 8192 micro-tokens, H=384, + 18 virtual layers, bf16). Treat it as an order-of-magnitude estimate. +""" +from __future__ import annotations + +import json +import math +import os +import statistics +from dataclasses import asdict, dataclass, field + +STAGES = ("pretrain", "sft", "dpo", "orpo", "grpo", "distill") +# Keys a stage may override; the Train tab merges ``stage_overrides[mode]``. +STAGE_KEYS = ("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") +MEMORY_LEVELS = ("none", "balanced", "aggressive", "extreme") +# Activation bytes per (token · executed layer · hidden unit), MLX bf16, +# derived from the measured 50M run peak (see module docstring). +ACT_BYTES_PER_TOKEN_LAYER_H = 120.0 +DEEPSEEK_FIT_RANGE = (1e17, 2e19) + + +@dataclass +class PlanRequest: + target_total: float # parameters, e.g. 50e6 + use_moe: bool = True + num_experts: int = 32 + active_params: float | None = None # either this ... + active_experts: int | None = None # ... or this (top-k); active_params wins if both + memory_saving: str = "none" # MEMORY_LEVELS + data_paths: tuple = () # pretrain + sft + preference files/dirs + vocab_size: int | None = None # None: tokenizer report / dataset, else 16384 + seq_len: int | None = None # None: from the dataset profile + ram_gb: float | None = None # None: detected + token_budget: float | None = None # pretrain tokens; None: dataset tokens (one pass) + recurrent_depth: int = 4 + dtype: str = "bfloat16" + + +@dataclass +class DatasetProfile: + files: list = field(default_factory=list) + kinds: dict = field(default_factory=dict) # stage -> [files] + pretrain_tokens: float = 0.0 + pretrain_tokens_source: str = "none" # measured | estimated | none + sft_records: float = 0.0 + pref_records: float = 0.0 + rl_records: float = 0.0 + 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 + vocab_size: int | None = None + tokenizer_dir: str = "" + notes: list = field(default_factory=list) + + +# ── dataset profile ──────────────────────────────────────────────────── +# Fallback chars/token when no tokenizer report is available (estimated, +# typical byte-level BPE at 16–32k vocab); replaced by the report values. +_DEFAULT_CPT = {"zh": 1.4, "ja": 1.2, "ko": 1.5, "latin": 3.8, "cyrillic": 3.2, "other": 2.5} +_SAMPLE_RECORDS = 2000 + + +def _record_kind(rec) -> str: + if isinstance(rec, dict): + if "chosen" in rec and "rejected" in rec: + return "pref" + turns = rec.get("conversations") if isinstance(rec.get("conversations"), list) else rec.get("messages") + if isinstance(turns, list): + # GRPO prompt files end in an empty assistant turn to be generated. + last = turns[-1] if turns and isinstance(turns[-1], dict) else {} + if last.get("role") == "assistant" and not str(last.get("content") or "").strip(): + return "rl" + return "sft" + if "prompt" in rec and ("answer" in rec or "reward" in rec or "solution" in rec): + return "rl" + return "pretrain" + + +def _read_json(path: str) -> dict: + try: + with open(path, encoding="utf-8") as fh: + data = json.load(fh) + return data if isinstance(data, dict) else {} + except (OSError, ValueError): + return {} + + +def _tokenizer_report(paths) -> tuple[dict, str]: + """First ``tokenizer_report.json`` next to / referenced by the inputs.""" + for p in paths: + for cand in (p, os.path.dirname(p)): + rep = os.path.join(cand, "tokenizer_report.json") + if os.path.isfile(rep): + return _read_json(rep), cand + man = _read_json(os.path.join(cand, "shards_manifest.json")) + tdir = man.get("tokenizer_dir") + if not tdir: + continue + # Relative paths in a manifest are relative to the repo (cwd) or to the manifest. + for base in ("", cand): + full = os.path.abspath(os.path.join(base, tdir)) + if os.path.isfile(os.path.join(full, "tokenizer_report.json")): + return _read_json(os.path.join(full, "tokenizer_report.json")), full + return {}, "" + + +def _count_records(path: str, sampled: int, sample_bytes: int) -> tuple[float, str]: + """Exact count for parquet / small text files, size-extrapolated otherwise.""" + size = os.path.getsize(path) + if path.endswith(".parquet"): + try: + import pyarrow.parquet as pq + return float(pq.ParquetFile(path).metadata.num_rows), "measured" + except Exception: + pass + if path.endswith((".jsonl", ".txt", ".csv", ".tsv")) and size <= 512 << 20: + with open(path, "rb") as fh: + lines = sum(chunk.count(b"\n") for chunk in iter(lambda: fh.read(1 << 20), b"")) + return float(max(lines - (1 if path.endswith((".csv", ".tsv")) else 0), sampled)), "measured" + if sampled and sample_bytes: + return float(sampled) * size / sample_bytes, "estimated" + return float(sampled), "estimated" + + +def profile_datasets(paths, *, max_files: int = 64) -> DatasetProfile: + """Profile pretrain / SFT / preference files for the planner. + + Prefers exact counts from a Chronos ``shards_manifest.json`` (token + shards) and the tokenizer report; otherwise samples records, maps each + text's script bucket to chars/token, and extrapolates by file size. + """ + from chronos.tokenizer.extract import _texts_from_record, expand, iter_records, script_bucket + prof = DatasetProfile() + paths = [os.path.expanduser(str(p)) for p in (paths or ()) if str(p).strip()] + report, prof.tokenizer_dir = _tokenizer_report(paths) + cpt_by_bucket = dict(_DEFAULT_CPT) + for bucket, stats in (report.get("heldout_eval") or {}).items(): + if isinstance(stats, dict) and stats.get("chars_per_token") and int(stats.get("texts", 0)) >= 100: + cpt_by_bucket[bucket] = float(stats["chars_per_token"]) + if report.get("vocab_size"): + prof.vocab_size = int(report["vocab_size"]) + for p in paths: + for cand in (p, os.path.dirname(p)): + man = _read_json(os.path.join(cand, "shards_manifest.json")) + if man.get("pretrain_train_tokens"): + prof.pretrain_tokens = float(man["pretrain_train_tokens"]) + prof.pretrain_tokens_source = "measured" + prof.sft_records = float(man.get("sft_train_records") or 0) + if man.get("sft_max_len"): + prof.seq_cap_tokens = int(man["sft_max_len"]) + prof.notes.append(f"token counts from {os.path.join(cand, 'shards_manifest.json')}") + break + if prof.pretrain_tokens_source == "measured": + break + if prof.pretrain_tokens_source == "measured": + return prof + files = expand(paths)[:max_files] + prof.files = files + tok_lens, chars, tokens = [], 0, 0.0 + for path in files: + kind_counts, sampled, sample_bytes = {}, 0, 0 + file_tokens = 0.0 + try: + for rec in iter_records(path): + if sampled >= _SAMPLE_RECORDS: + break + sampled += 1 + kind = _record_kind(rec) + kind_counts[kind] = kind_counts.get(kind, 0) + 1 + rec_tokens = 0.0 + for text in _texts_from_record(rec, None): + n = len(text) + sample_bytes += len(text.encode("utf-8", "replace")) + 8 + chars += n + rec_tokens += n / cpt_by_bucket.get(script_bucket(text), 2.5) + file_tokens += rec_tokens + tok_lens.append(rec_tokens) + except Exception as exc: # unreadable file: report, keep going + prof.notes.append(f"skipped {os.path.basename(path)}: {type(exc).__name__}") + continue + if not sampled: + continue + kind = max(kind_counts, key=kind_counts.get) + prof.kinds.setdefault(kind, []).append(path) + records, how = _count_records(path, sampled, sample_bytes) + tokens += file_tokens / sampled * records if kind == "pretrain" else 0.0 + if kind == "sft": + prof.sft_records += records + elif kind == "pref": + prof.pref_records += records + elif kind == "rl": + prof.rl_records += records + if how == "estimated": + prof.notes.append(f"{os.path.basename(path)}: record count extrapolated from file size") + if tokens: + prof.pretrain_tokens, prof.pretrain_tokens_source = tokens, "estimated" + if tok_lens: + tok_lens.sort() + 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)") + return prof + + +# ── architecture search ──────────────────────────────────────────────── +HEAD_DIM = 64 +ROPE_DIM = 32 +LOOKAHEAD_STEPS = 2 + + +def _kv_heads(heads: int) -> int: + kv = max(1, round(heads / 3)) # GQA group of ~3 (recipe: 6 heads / 2 KV) + while heads % kv: + kv -= 1 + return kv + + +def _latent(h: int) -> int: + return int(min(512, max(32, round(h / 3 / 32) * 32))) + + +def _attn_params(h: int, heads: int, kv: int, latent: int, mla: bool) -> int: + d = HEAD_DIM + if mla: # mirrors ui/estimator.py::_attn_params + return 2 * h * heads * d + h * latent + 2 * latent * kv * d + d + latent + return 2 * h * heads * d + 2 * h * kv * d + 2 * d + + +def _base_params(h: int, layers: int, groups: int, depth: int, vocab: int, experts: int, use_moe: bool) -> int: + """Everything except FFN / expert weights (embedding tied, MLA on even layers).""" + heads = h // HEAD_DIM + kv, latent = _kv_heads(heads), _latent(h) + attn = sum(_attn_params(h, heads, kv, latent, i % 2 == 0) for i in range(layers)) + ph = max(1, h // 4) + lookahead = h * ph + ph * max(experts, 1) * (LOOKAHEAD_STEPS + 1) * groups # layer-aware predictor + controller = (depth + groups) * h + 2 * h + 1 + groups + routers = layers * h * experts if use_moe else 0 + return vocab * h + attn + (2 * layers + 1) * h + routers + lookahead + controller + + +@dataclass +class ArchCandidate: + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + kv_latent_dim: int + ffn_width: int # moe_intermediate_size (MoE) or intermediate_size (dense) + num_experts: int + num_experts_per_tok: int + total: int + active: int + union_upper: int + score: float + + +def search_architecture(target_total: float, *, use_moe: bool, num_experts: int, vocab_size: int, + depth: int, active_params: float | None = None, active_experts: int | None = None, + shared_experts: int = 1) -> ArchCandidate: + """Closed-form search: FFN width is solved from the total-parameter target. + + For each width h (multiples of 64, head_dim 64), physical depth L + (prelude 1 + groups + coda 1) and top-k, the expert width that hits + ``target_total`` is solved analytically and rounded to a multiple of 32. + Score = |log total/T| (+ |log active/A| when A is given) + a mild shape + prior (h / L in [32, 128], expert width in [h/4, 2h], k ≥ 2) and weak + tie-breakers toward h/L ≈ 64 and expert width ≈ h/2 (fine-grained). + """ + if target_total <= 0: + raise ValueError("target_total must be positive") + E = max(1, int(num_experts)) if use_moe else 1 + if use_moe and active_experts and not active_params: + k_choices = [max(1, min(E, int(active_experts)))] + elif use_moe: + k_choices = list(range(1, min(E, 8) + 1)) + else: + k_choices = [1] + best = None + for h in range(128, 4097, 64): + if vocab_size * h > 0.85 * target_total: + break # embedding alone exhausts the budget + heads = h // HEAD_DIM + kv_heads = _kv_heads(heads) + # A prime head count collapses GQA to one KV head (MQA); prefer widths + # whose head count has a divisor near heads/3. + gqa_penalty = 0.03 if heads >= 4 and heads // kv_heads > 4 else 0.0 + for layers in range(3, 49): + groups = layers - 2 + base = _base_params(h, layers, groups, depth, vocab_size, E, use_moe) + if base >= target_total: + break + for k in k_choices: + slope = layers * 3 * h * ((E + shared_experts) if use_moe else 1) + raw = (target_total - base) / slope + step = 32 if use_moe else 64 + for width in {max(step, int(raw // step) * step), max(step, int(math.ceil(raw / step)) * step)}: + total = base + slope * width + per = 3 * h * width + active = base + layers * per * (k + shared_experts) if use_moe else total + union = active + groups * (min(E, k * depth) - k) * per if use_moe else total + score = abs(math.log(total / target_total)) + if use_moe and active_params: + score += abs(math.log(active / float(active_params))) + ratio = h / layers + score += 0.05 * max(0.0, math.log(32 / ratio), math.log(ratio / 128)) + lo, hi = ((h / 4, 2 * h) if use_moe else (2 * h, 4 * h)) + score += 0.10 * max(0.0, math.log(lo / width), math.log(width / hi)) + # Tie-breakers among exact fits (heuristics, not fitted): h/L near 64, + # fine-grained experts near h/2 (E ≥ 16) or h (E < 16), dense FFN near 8/3·h. + centre = (h / 2 if E >= 16 else h) if use_moe else 8 * h / 3 + score += 0.02 * abs(math.log(ratio / 64)) + 0.02 * abs(math.log(width / centre)) + if use_moe and k == 1: + score += 0.02 # top-1 loses gate gradients under norm_topk + score += gqa_penalty + if best is None or score < best.score: + best = ArchCandidate(h, layers, heads, kv_heads, _latent(h), width, E if use_moe else 0, + k, int(total), int(active), int(union), score) + if best is None: + raise ValueError(f"no architecture fits {target_total:.3g} parameters with vocab {vocab_size}") + return best + + +# ── memory model ─────────────────────────────────────────────────────── +# memory_saving → (inference options, training options). Every option is an +# existing Chronos flag; nothing here changes numerics during training except +# activation checkpointing (recompute) and optimizer paging (SSD round trips). +MEMORY_POLICIES = { + "none": dict(kv_cache_dtype="auto", mla_cache_rope=True, router_scope="per_iteration", + resident_fraction=1.0, ram_fraction=1.0, activation_checkpointing=False, + optimizer_paging=False), + "balanced": dict(kv_cache_dtype="int8", mla_cache_rope=False, router_scope="per_iteration", + resident_fraction=0.5, ram_fraction=0.75, activation_checkpointing=False, + optimizer_paging=False), + "aggressive": dict(kv_cache_dtype="int8", mla_cache_rope=False, router_scope="shared", + resident_fraction=0.25, ram_fraction=0.5, activation_checkpointing=True, + optimizer_paging=False), + "extreme": dict(kv_cache_dtype="int8", mla_cache_rope=False, router_scope="shared", + resident_fraction=0.125, ram_fraction=0.25, activation_checkpointing=True, + optimizer_paging=True), +} +# Share of physical RAM a plan may target for the training peak (rest: OS, +# data loader, UI, other jobs). Lowered by the memory-saving level. +_RAM_SHARE = {"none": 0.6, "balanced": 0.45, "aggressive": 0.3, "extreme": 0.2} + + +def detect_ram_gb() -> float: + try: + import psutil + return psutil.virtual_memory().total / 2**30 + except Exception: + return 16.0 + + +def training_peak_bytes(total_params: int, micro_tokens: int, hidden: int, executed_layers: int, + vocab: int, *, checkpointing: bool = False, optimizer_paging: bool = False, + routed_params: int = 0) -> int: + """Estimated MLX/Torch training peak (bytes). + + * state: FP32 master + grad + Adam m, v = 16 B/param (optimizer paging + moves the routed experts' Adam state, 8 B/param, to SSD); + * activations: ``ACT_BYTES_PER_TOKEN_LAYER_H`` · tokens · layers · h + (checkpointing keeps ~1 layer boundary per layer, ≈ ×0.25 here); + * logits: FP32 logits + grad + softmax ≈ 12 B per token·vocab. + """ + state = 16 * total_params - (8 * routed_params if optimizer_paging else 0) + act = ACT_BYTES_PER_TOKEN_LAYER_H * micro_tokens * executed_layers * hidden + if checkpointing: + act *= 0.25 + return int(state + act + 12 * micro_tokens * vocab) + + +# ── hyperparameters ──────────────────────────────────────────────────── +def deepseek_hparams(compute_flops: float) -> tuple[float, float]: + """DeepSeek LLM (arXiv 2401.02954) eq. for η_opt and B_opt (tokens).""" + c = max(float(compute_flops), 1.0) + return 0.3118 * c ** -0.1250, 0.2920 * c ** 0.3271 + + +def flops_per_token(hidden: int, executed_layers: int, seq_len: int, active_ffn_ratio: float = 1.0) -> float: + """``M = 72·L·h² + 12·L·h·s`` (DeepSeek's non-embedding FLOPs/token). + + ``72·L·h² = 6 · 12·L·h²`` counts forward+backward over the 4h² attention + and 8h² FFN parameters of a standard layer (FFN 4h, 2 matrices). For a + MoE or SwiGLU layer the FFN share is scaled by ``active_ffn_ratio`` = + active FFN parameters / 8h²; L counts executed (virtual) layers. + """ + L, h = executed_layers, hidden + return (24 + 48 * active_ffn_ratio) * L * h * h + 12 * L * h * seq_len + + +def _pow2_floor(x: float) -> int: + return 1 << max(0, int(math.floor(math.log2(max(1.0, x))))) + + +# ── plan ─────────────────────────────────────────────────────────────── +def _seq_len(req: PlanRequest, prof: DatasetProfile) -> tuple[int, str]: + if req.seq_len: + return int(req.seq_len), "user" + if prof.seq_cap_tokens: + return int(prof.seq_cap_tokens), "shard manifest max length (measured)" + if prof.seq_p95_tokens: + s = 1 << max(7, int(math.ceil(math.log2(prof.seq_p95_tokens * 1.15)))) + return int(min(2048, s)), "dataset p95 × 1.15 → power of two (estimated)" + return 512, "default" + + +def _round_lr(x: float) -> float: + return float(f"{x:.2g}") + + +def plan_training(req: PlanRequest, profile: DatasetProfile | None = None, *, verify: bool = True) -> dict: + """Return ``{"config", "stage_overrides", "longrun_phases", "report"}``. + + ``config`` is a complete Config-tab dict (v3 kernel); ``stage_overrides`` + maps each stage to its own hyperparameters (the Train tab merges the + selected stage over ``config``); ``longrun_phases`` are MLX long-run + ``PhaseSpec`` dicts for pretrain + SFT. + """ + level = req.memory_saving if req.memory_saving in MEMORY_POLICIES else "none" + policy = MEMORY_POLICIES[level] + prof = profile if profile is not None else profile_datasets(req.data_paths) + notes, sources = list(prof.notes), {} + vocab = int(req.vocab_size or prof.vocab_size or 16384) + sources["vocab_size"] = "user" if req.vocab_size else ("tokenizer report" if prof.vocab_size else "default") + depth = max(1, int(req.recurrent_depth)) + arch = search_architecture(req.target_total, use_moe=req.use_moe, num_experts=req.num_experts, + vocab_size=vocab, depth=depth, active_params=req.active_params, + active_experts=req.active_experts) + h, L, groups = arch.hidden_size, arch.num_hidden_layers, arch.num_hidden_layers - 2 + seq_len, sources["max_seq_len"] = _seq_len(req, prof) + executed = L + groups * (depth - 1) # virtual layers at full depth + shared = 1 if req.use_moe else 0 + if req.use_moe: + ffn_ratio = 3 * h * arch.ffn_width * (arch.num_experts_per_tok + shared) / (8 * h * h) + else: + ffn_ratio = 3 * h * arch.ffn_width / (8 * h * h) + + # Token budget D and DeepSeek-law LR / batch. + if req.token_budget: + tokens, sources["tokens"] = float(req.token_budget), "user" + elif prof.pretrain_tokens >= 2.0 * arch.total: + tokens, sources["tokens"] = prof.pretrain_tokens, f"dataset, one pass ({prof.pretrain_tokens_source})" + else: + # Too little (or no) pretrain data to size the run: plan for 20 tokens + # per parameter and report the shortfall instead of fitting a toy set. + tokens, sources["tokens"] = 20.0 * arch.total, "default 20 tokens/parameter (estimated)" + if prof.pretrain_tokens: + notes.append(f"dataset has ≈{prof.pretrain_tokens:.3g} pretrain tokens, < 2 per parameter; " + f"the plan assumes {tokens:.3g} tokens (multiple epochs or more data)") + m = flops_per_token(h, executed, seq_len, ffn_ratio) + compute = m * tokens + eta, batch_tokens = deepseek_hparams(compute) + if not DEEPSEEK_FIT_RANGE[0] <= compute <= DEEPSEEK_FIT_RANGE[1] * 10: + notes.append(f"C = {compute:.2e} FLOPs is outside the DeepSeek fit range " + f"{DEEPSEEK_FIT_RANGE[0]:.0e}–{DEEPSEEK_FIT_RANGE[1]:.0e}; LR/batch are extrapolated and clamped") + # Chronos clamps: aux losses + recurrence make very high LR fragile; the + # ceiling shrinks with width (≈ μP-style 1/h, anchored at 3e-3 for h=384). + lr = _round_lr(min(3e-3 * min(1.0, 384 / h), max(1e-4, eta))) + batch_tokens = min(1 << 21, max(16384, _pow2_floor(batch_tokens))) + + ram_gb = float(req.ram_gb or detect_ram_gb()) + budget = ram_gb * 2**30 * _RAM_SHARE[level] + routed = L * arch.num_experts * 3 * h * arch.ffn_width if req.use_moe else 0 + # Largest power-of-two micro-batch (≤ 16k tokens) whose estimated peak + # fits the level's RAM share; accumulation makes up the optimizer batch. + micro = max(seq_len, _pow2_floor(min(batch_tokens, 16384))) + while micro > seq_len and training_peak_bytes( + arch.total, micro, h, executed, vocab, checkpointing=policy["activation_checkpointing"], + optimizer_paging=policy["optimizer_paging"], routed_params=routed) > budget: + micro //= 2 + peak = training_peak_bytes(arch.total, micro, h, executed, vocab, + checkpointing=policy["activation_checkpointing"], + optimizer_paging=policy["optimizer_paging"], routed_params=routed) + if peak > budget: + notes.append(f"estimated training peak {peak / 2**30:.1f} GiB exceeds the {budget / 2**30:.1f} GiB " + f"target even at one sequence per micro-batch; raise memory saving or lower the target") + rows = max(1, micro // seq_len) + accum = max(1, batch_tokens // (rows * seq_len)) + batch_tokens = rows * seq_len * accum + updates = max(1, int(tokens // batch_tokens)) + + # Independent per-stage hyperparameters. Ratios to the pretrain LR are + # heuristics (the 50M recipe uses SFT = pretrain/3, half the batch); + # GRPO β = 0.04 follows DeepSeekMath (arXiv 2402.03300 §4.2); other + # coefficients are the Chronos trainer defaults. + def stage(lr_value, rows_value, accum_value, **extra): + out = dict(learning_rate=_round_lr(lr_value), batch_size=int(max(1, rows_value)), + accumulation_steps=int(max(1, accum_value)), max_seq_len=seq_len, grad_clip=1.0) + out.update(extra) + return out + + epochs, pre_steps = 1, 0 + if prof.pretrain_tokens and tokens > prof.pretrain_tokens: + epochs = max(1, int(math.ceil(tokens / prof.pretrain_tokens))) + if req.token_budget or tokens != prof.pretrain_tokens: + pre_steps = updates * accum # Train tab counts micro-steps + micro_steps = pre_steps or updates * accum + + def save_interval(steps: float) -> int: + """≈20 checkpoints per stage, a multiple of 10 in [100, 100000] (the UI slider).""" + return int(min(100_000, max(100, round(steps / 20, -1)))) + + save_every = save_interval(micro_steps) + sft_epochs = 1 if prof.sft_records >= 200_000 or not prof.sft_records else (2 if prof.sft_records >= 20_000 else 3) + pair_rows = max(1, rows // 2) # chosen + rejected per pair + rl_rows = max(1, rows // 4) + num_gen = 8 if level == "none" else (4 if level in ("balanced", "aggressive") else 2) + # Post-training stages are sized by their own records (one micro-step per + # `rows` records); without records fall back to a fixed 500-step cadence. + sft_save = save_interval(prof.sft_records * sft_epochs / rows) if prof.sft_records else 500 + pref_save = save_interval(prof.pref_records / pair_rows) if prof.pref_records else 500 + rl_save = save_interval(prof.rl_records / rl_rows) if prof.rl_records else 500 + gen_len = int(min(256, max(24, (prof.seq_p95_tokens or prof.seq_cap_tokens or seq_len) // 2))) + overrides = { + "pretrain": stage(lr, rows, accum, epochs=epochs, max_steps=pre_steps, weight_decay=0.1, + save_interval=save_every, lambda_router_anchor=0.0), + "sft": stage(lr / 3, rows, max(1, accum // 2), epochs=sft_epochs, max_steps=0, weight_decay=0.1, + save_interval=sft_save, lambda_router_anchor=0.0), + "dpo": stage(max(5e-6, lr / 20), pair_rows, max(1, accum // 2), epochs=1, max_steps=0, weight_decay=0.01, + beta=0.1, lambda_router_anchor=0.01, save_interval=pref_save), + "orpo": stage(max(5e-6, lr / 10), pair_rows, max(1, accum // 2), epochs=1, max_steps=0, weight_decay=0.01, + lambda_or=0.1, lambda_router_anchor=0.01, save_interval=pref_save), + "grpo": stage(min(2e-5, max(1e-6, lr / 100)), rl_rows, 1, epochs=1, max_steps=0, weight_decay=0.01, + beta=0.04, num_generations=num_gen, temperature=1.0, max_gen_len=gen_len, + lambda_router_anchor=0.01, save_interval=rl_save), + "distill": stage(lr / 3, rows, max(1, accum // 2), epochs=1, max_steps=0, weight_decay=0.1, + alpha=0.7, temperature=4.0, lambda_router_anchor=0.0, save_interval=sft_save), + } + + p95 = prof.seq_p95_tokens or prof.seq_cap_tokens + # 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) + 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 \ + round(max(per_expert_bytes * L * max(2, arch.num_experts_per_tok) / 2**20, + expert_mb * policy["ram_fraction"]), 1) + resident_weights = (arch.total - routed * (1 - policy["resident_fraction"])) * 2 + model = dict( + architecture_version=3, recurrent_attention_mode="shared_first", rope_layout="per_head", + hidden_size=h, num_hidden_layers=L, num_attention_heads=arch.num_attention_heads, + num_key_value_heads=arch.num_key_value_heads, kv_latent_dim=arch.kv_latent_dim, rope_dim=ROPE_DIM, + sliding_window_size=int(window), use_hybrid_attention=True, vocab_size=vocab, tie_word_embeddings=True, + use_moe=bool(req.use_moe), num_experts=arch.num_experts if req.use_moe else 1, + num_experts_per_tok=arch.num_experts_per_tok, num_shared_experts=shared, + moe_intermediate_size=arch.ffn_width if req.use_moe else 0, + intermediate_size=0 if req.use_moe else arch.ffn_width, + lookahead_steps=LOOKAHEAD_STEPS, layer_aware_lookahead=True, layer_aware_cache=True, + recurrent_enabled=depth > 1, recurrent_groups=groups, recurrent_depth=depth, prelude_layers=1, + coda_layers=1, min_recurrent_depth=max(1, depth // 2), max_inference_depth=depth, + adaptive_halting=depth > 1, halting_threshold=0.95, depth_schedule="fixed_then_adaptive", + depth_schedule_warmup_fraction=0.3, recurrent_strategy="scheduled", recurrent_latent_refinement=False, + 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, + 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"], + vram_budget_gb=round(max(0.25, resident_weights / 2**30 * 1.25), 2), dtype=req.dtype, + ) + config = {**model, **overrides["pretrain"], "stage_overrides": overrides} + if prof.tokenizer_dir and os.path.isfile(os.path.join(prof.tokenizer_dir, "tokenizer.json")): + # The tokenizer the dataset was profiled with is the one the model needs. + config["tokenizer_path"] = prof.tokenizer_dir + sources["tokenizer_path"] = "dataset tokenizer report" + + longrun = [ + dict(name="pretrain", budget_fraction=0.5, peak_lr=lr, tokens_per_step=int(batch_tokens), + micro_tokens=int(micro), seq_len=seq_len, warmup_fraction=0.02, min_lr_ratio=0.1, + weight_decay=0.1, grad_clip=1.0), + dict(name="sft", budget_fraction=0.5, peak_lr=_round_lr(lr / 3), tokens_per_step=int(max(micro, batch_tokens // 2)), + micro_tokens=int(micro), seq_len=seq_len, warmup_fraction=0.01, min_lr_ratio=0.05, + weight_decay=0.1, grad_clip=1.0), + ] + counted = None + if verify: + from chronos.config_schema import model_config_kwargs + from chronos.model.factory import count_parameters + counted = count_parameters(model_config_kwargs(config)) + drift = abs(counted["total"] / arch.total - 1) + if drift > 0.01: + notes.append(f"closed-form count differs from the meta-device count by {drift:.2%}") + total = counted["total"] if counted else arch.total + active = counted["active"] if counted else arch.active + union = counted["active_union_upper"] if counted else arch.union_upper + report = { + "request": asdict(req) | {"data_paths": list(req.data_paths)}, + "architecture": { + "total_params": total, "active_params": active, "active_union_upper": union, + "params_source": "measured (meta-device count)" if counted else "estimated (closed form)", + "hidden_size": h, "physical_layers": L, "recurrent_groups": groups, "recurrent_depth": depth, + "virtual_layers": executed, "experts": arch.num_experts, "top_k": arch.num_experts_per_tok, + "expert_width": arch.ffn_width, "heads": f"{arch.num_attention_heads}/{arch.num_key_value_heads}", + }, + "dataset": {k: v for k, v in asdict(prof).items() if k not in ("files", "notes")} + | {"files": len(prof.files)}, + "scaling": { + "tokens": tokens, "tokens_source": sources["tokens"], "flops_per_token": m, + "compute_flops": compute, "deepseek_lr": eta, "deepseek_batch_tokens": deepseek_hparams(compute)[1], + "chosen_lr": lr, "chosen_batch_tokens": batch_tokens, "label": "estimated", + "law": "η=0.3118·C^-0.125, B=0.2920·C^0.3271 (DeepSeek LLM, arXiv 2401.02954)", + "tokens_per_param": round(tokens / max(1, total), 1), + }, + "memory": { + "level": level, "ram_gb": round(ram_gb, 1), "target_gb": round(budget / 2**30, 1), + "micro_tokens": micro, "rows_per_micro": rows, "accumulation_steps": accum, + "training_peak_gb": round(peak / 2**30, 2), "training_peak_label": "estimated", + "routed_expert_mb_bf16": round(expert_mb, 1), "ram_budget_mb": ram_budget, + "kv_cache_dtype": policy["kv_cache_dtype"], "router_scope": policy["router_scope"], + "activation_checkpointing": policy["activation_checkpointing"], + "optimizer_paging": policy["optimizer_paging"], + }, + "sources": sources, + "notes": notes, + } + return {"config": config, "stage_overrides": overrides, "longrun_phases": longrun, "report": report} + + +def stage_table(overrides: dict) -> tuple[list[str], list[list]]: + """(headers, rows) for the UI editor: one row per stage, one column per key.""" + headers = ["stage"] + list(STAGE_KEYS) + rows = [[name] + [overrides.get(name, {}).get(k, "") for k in STAGE_KEYS] for name in STAGES] + return headers, rows + + +def overrides_from_table(rows) -> dict: + """Inverse of :func:`stage_table`; blank cells are dropped, numbers coerced.""" + out = {} + for row in rows or (): + row = list(row) + if not row or str(row[0]).strip() not in STAGES: + continue + entry = {} + for key, value in zip(STAGE_KEYS, row[1:]): + if value is None or (isinstance(value, str) and not value.strip()): + continue + if isinstance(value, float) and math.isnan(value): + continue + number = float(value) + entry[key] = int(number) if key in _INT_KEYS else number + out[str(row[0]).strip()] = entry + return out + + +_INT_KEYS = {"batch_size", "accumulation_steps", "epochs", "max_steps", "max_seq_len", "save_interval", + "num_generations", "max_gen_len"} diff --git a/chronos/model/precision_ops.py b/chronos/model/precision_ops.py new file mode 100644 index 0000000..70e55de --- /dev/null +++ b/chronos/model/precision_ops.py @@ -0,0 +1,144 @@ +"""Torch kernels for the single-compute-dtype policy (CPU / CUDA / MPS). + +Mirrors ``chronos/mlx/attention.py``: weights, activations and KV cache stay +in one compute dtype; FP32 is obtained inside fused kernels or at a few +small, deliberate boundaries (see ``chronos/precision.py``). + +* ``RMSNorm``: FP32 statistics, one rounding. CPU/CUDA use the fused + ``F.rms_norm``; on MPS the fused low-precision kernel accumulates in the + input dtype (measured bf16 rel. error 2.34e-3 vs 1.66e-3), so MPS keeps the + explicit FP32-statistics form. +* RoPE: FP32 tables that ``model.to(dtype)`` never converts; the rotation is + computed in FP32 and rounded once (MLX measured 1.96e-3 vs 2.48e-3 with + BF16 tables). +* Attention: ``F.scaled_dot_product_attention`` with native GQA (no repeated + K/V), a boolean mask and in-kernel dropout; FP32 softmax inside the kernel + on every backend. MPS bf16 measured rel. error 1.50e-3 vs 2.87e-3 for the + explicit softmax path, and faster forward (prefill 5.6 vs 6.8 ms, decode + 0.21 vs 0.40 ms, B16 S512 GQA 6/2 d64). +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +_TORCH_VERSION = tuple(int(p) for p in torch.__version__.split("+")[0].split(".")[:2]) +SDPA_NATIVE_GQA = _TORCH_VERSION >= (2, 5) + + +class RMSNorm(nn.Module): + """RMSNorm with FP32 statistics and a single rounding to the input dtype. + + Parameter name (``weight``) and class-name suffix match minimind's + RMSNorm, so checkpoints and the no-decay optimizer grouping are unchanged. + """ + + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.device.type != "mps" or x.dtype == torch.float32: + w = self.weight if self.weight.dtype == x.dtype else self.weight.to(x.dtype) + return F.rms_norm(x, (x.shape[-1],), w, self.eps) + xf = x.float() + return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight).to(x.dtype) + + +def rope_tables(dim: int, length: int, base: float, scaling, device) -> tuple[torch.Tensor, torch.Tensor]: + """FP32 duplicated cos/sin tables [length, dim] (minimind layout).""" + import chronos.deps # noqa: F401 (minimind on sys.path) + from model.model_minimind import precompute_freqs_cis + cos, sin = precompute_freqs_cis(dim=int(dim), end=int(length), rope_base=base, rope_scaling=scaling) + return cos.to(device=device, dtype=torch.float32), sin.to(device=device, dtype=torch.float32) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """x [B, S, nH, d] (compute dtype), cos/sin [S, d] FP32 -> x.dtype, rounded once.""" + c, s = cos[:, None, :], sin[:, None, :] + return (x * c + _rotate_half(x) * s).to(x.dtype) + + +def causal_allowed(S: int, T: int, device, *, past: int | None = None, + window: int | None = None) -> torch.Tensor: + """Boolean [S, T] mask for queries at absolute key rows ``past..past+S-1``.""" + past = T - S if past is None else int(past) + q_pos = torch.arange(S, device=device)[:, None] + past + k_pos = torch.arange(T, device=device)[None, :] + allowed = k_pos <= q_pos + if window is not None: + allowed = allowed & (k_pos > q_pos - int(window)) + return allowed + + +def key_keep_mask(attention_mask: torch.Tensor, S: int, T: int) -> torch.Tensor: + """Caller padding mask -> boolean keep mask broadcastable to [B, 1, S, T]. + + Accepts the tokenizer-style 2-D mask [B, total_len], a 3-D [B, S, T] mask + or a 4-D mask that is either 0/1 or additive (0 / large negative). The + key axis is right-aligned (cached decoding, sliding-window eviction). + """ + m = attention_mask + if m.dim() == 4 and m.dtype != torch.bool and bool((m < 0).any()): + m = m > -1e4 + elif m.dtype != torch.bool: + m = m > 0.5 + if m.shape[-1] < T: + pad = m.new_ones(*m.shape[:-1], T - m.shape[-1]) + m = torch.cat([pad, m], dim=-1) + m = m[..., -T:] + if m.dim() == 2: + return m[:, None, None, :] + if m.dim() == 3: + return m[:, None, -S:, :] + return m[..., -S:, :] + + +def build_attention_mask(S: int, T: int, device, *, past: int, window: int | None, + attention_mask: torch.Tensor | None): + """Return ``(mask, is_causal)`` for :func:`attention`. + + The diagonal (each query's own key) is always kept, so left-padded query + rows never have an empty softmax row (no NaN in any dtype). + """ + if attention_mask is None: + if window is None and S == T: + return None, S > 1 # fused causal kernel (flash on CUDA) + if window is None and S == 1: + return None, False # single-token decode sees every cached key + return causal_allowed(S, T, device, past=past, window=window), False + allowed = causal_allowed(S, T, device, past=past, window=window) + diag = (torch.arange(T, device=device)[None, :] == torch.arange(S, device=device)[:, None] + past) + return (allowed & key_keep_mask(attention_mask, S, T)) | diag, False + + +def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, mask, is_causal: bool, + dropout_p: float, scale: float | None = None) -> torch.Tensor: + """Fused SDPA on q [B, nH, S, d] and k/v [B, nKV, T, d] in the compute dtype.""" + n_rep = q.shape[1] // k.shape[1] + if n_rep > 1 and not SDPA_NATIVE_GQA: + k, v = k.repeat_interleave(n_rep, dim=1), v.repeat_interleave(n_rep, dim=1) + kwargs = {"enable_gqa": True} if (n_rep > 1 and SDPA_NATIVE_GQA) else {} + return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=dropout_p, + is_causal=is_causal, scale=scale, **kwargs) + + +def lm_head_logits(hidden: torch.Tensor, weight: torch.Tensor, logits_dtype: torch.dtype | None = None): + """LM head whose logits are born in ``logits_dtype``. + + With ``logits_dtype=float32`` on a low-precision model, the hidden state + [N, H] and head weight [V, H] are upcast (H/V of the logits' size) so the + matmul accumulates and stores FP32 directly: no low-precision rounding of + the logits and no extra full-size cast before cross-entropy (measured + bf16-logit rel. error 1.4e-3 vs 1.1e-7). + """ + if logits_dtype is None or logits_dtype == weight.dtype: + return F.linear(hidden, weight) + return F.linear(hidden.to(logits_dtype), weight.to(logits_dtype)) diff --git a/chronos/model/recurrent.py b/chronos/model/recurrent.py new file mode 100644 index 0000000..c01d199 --- /dev/null +++ b/chronos/model/recurrent.py @@ -0,0 +1,208 @@ +"""Recurrent-depth primitives used by the Chronos v2 kernel. + +The recurrent controller is deliberately separate from attention/MoE. It +owns loop conditioning, contractive step sizes, ACT weights, and convergence +statistics while the shared Chronos blocks remain ordinary traceable modules. +""" +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from chronos.model.precision_ops import RMSNorm + + +@dataclass +class RecurrentDepthState: + halting_probs: torch.Tensor + ponder_cost: torch.Tensor + depth_used: torch.Tensor + stability_loss: torch.Tensor + distill_loss: torch.Tensor + contraction: torch.Tensor = None + exited_early: bool = False + + +class RecurrentController(nn.Module): + """Adaptive-compute controller shared by all recurrent groups.""" + + def __init__(self, config): + super().__init__() + hidden = int(config.hidden_size) + depth = int(config.recurrent_depth) + groups = int(config.recurrent_groups) + self.depth = depth + self.groups = groups + self.min_depth = int(config.min_recurrent_depth) + self.threshold = float(config.halting_threshold) + self.epsilon = float(config.halting_epsilon) + self.dt_min = float(config.recurrent_dt_min) + self.dt_max = float(config.recurrent_dt_max) + + self.loop_embeddings = nn.Parameter(torch.zeros(depth, hidden)) + self.group_embeddings = nn.Parameter(torch.zeros(groups, hidden)) + self.halting_norm = RMSNorm(hidden, eps=config.rms_norm_eps) + self.halting_head = nn.Linear(hidden, 1) + # A depth-one migration must initially behave exactly like v1. + nn.init.zeros_(self.halting_head.weight) + nn.init.constant_(self.halting_head.bias, -1.0) + self.step_sizes = nn.Parameter( + torch.full((groups,), float(config.recurrent_dt_init)) + ) + self.training_progress = 0.0 + + def set_training_progress(self, step: int, total_steps: int) -> None: + self.training_progress = min(1.0, max(0.0, float(step) / max(int(total_steps), 1))) + + def adaptive_active(self, config, training: bool) -> bool: + if not bool(config.adaptive_halting) or config.recurrent_strategy == "fixed": + return False + if not training: + return True + schedule = str(getattr(config, "depth_schedule", "fixed_then_adaptive")).lower() + if schedule in {"fixed", "constant", "off"}: + return False + if schedule in {"adaptive", "act", "on"}: + return True + warmup = float(getattr(config, "depth_schedule_warmup_fraction", 0.25)) + return self.training_progress >= warmup + + def execution_depth(self, config, training: bool) -> int: + """Choose one depth bucket for the complete batch/sequence. + + Batch-wide buckets keep KV and compiler shapes deterministic. Jitter + is deliberately disabled by default and never applies at inference. + The curriculum increases its ceiling linearly after warmup; the + default zero jitter retains the legacy maximum-depth training graph. + """ + ceiling = int(config.recurrent_depth if training else config.max_inference_depth) + if not training or str(getattr(config, "recurrent_strategy", "fixed")) == "fixed": + return ceiling + jitter = min(1.0, max(0.0, float(getattr(config, "recurrent_depth_jitter", 0.0)))) + if jitter == 0.0: + return ceiling + warmup = min(0.99, max(0.0, float(config.depth_schedule_warmup_fraction))) + progress = max(0.0, (self.training_progress - warmup) / (1.0 - warmup)) + curriculum_max = self.min_depth + round((ceiling - self.min_depth) * progress) + lower = max(self.min_depth, curriculum_max - round((curriculum_max - self.min_depth) * jitter)) + # CPU scalar draw does not synchronize the device and follows torch's + # seeded RNG, making resumed/seeded runs reproducible. + return int(torch.randint(lower, curriculum_max + 1, ()).item()) + + def condition(self, hidden: torch.Tensor, step: int, group: int) -> torch.Tensor: + # Reuse the last learned depth embedding for conservative test-time + # depth extrapolation. The transition itself is weight-tied. + depth_idx = min(int(step), self.loop_embeddings.shape[0] - 1) + return hidden + self.loop_embeddings[depth_idx][None, None, :] + self.group_embeddings[group][None, None, :] + + def step_size(self, group: int) -> torch.Tensor: + return self.step_sizes[group].clamp(self.dt_min, self.dt_max) + + def contract(self, before: torch.Tensor, after: torch.Tensor, group: int) -> torch.Tensor: + dt = self.step_size(group).to(dtype=after.dtype) + return before + dt * (after - before) + + def halt_probability(self, hidden: torch.Tensor) -> torch.Tensor: + # FP32 sigmoid avoids saturation in BF16/FP16 routing paths. + logits = self.halting_head(self.halting_norm(hidden)).squeeze(-1).float() + return torch.sigmoid(logits) + + +def recurrent_auxiliary_state( + states: list[torch.Tensor], + halt_probs: list[torch.Tensor], + *, + min_depth: int, + threshold: float, +) -> tuple[torch.Tensor, RecurrentDepthState]: + """ACT-style state mixture plus differentiable recurrence diagnostics.""" + if not states: + raise ValueError("recurrent_auxiliary_state requires at least one state") + if len(states) == 1: + p = halt_probs[0].new_ones(halt_probs[0].shape) + zero = states[0].new_zeros(()) + depth = p + return states[0], RecurrentDepthState( + halting_probs=p.unsqueeze(-1), + ponder_cost=zero, + depth_used=depth, + stability_loss=zero, + distill_loss=zero, + contraction=zero, + ) + + # ACT uses an accumulated halting mass. ``threshold`` is allowed to be + # zero/one from a UI or checkpoint, but the numerical target is always + # kept in [epsilon, 1] so neither setting can produce NaNs or a zero-mass + # mixture. The final pass receives the exact remainder and therefore + # always makes the total mixture weight equal to one. + epsilon = 1e-6 + target = min(1.0, max(epsilon, float(threshold))) + # Draft forwards (depth override) may stop before ``min_depth``; the last + # executed state must still receive the remaining mass, never zero. + min_depth = min(int(min_depth), len(states)) + cumulative = halt_probs[0].new_zeros(halt_probs[0].shape) + halted = torch.zeros_like(cumulative, dtype=torch.bool) + mixed = torch.zeros_like(states[0]) + depth_used = torch.zeros_like(cumulative) + accumulated = [] + for idx, (state, prob) in enumerate(zip(states, halt_probs)): + step = idx + 1 + if step < min_depth: + weight = torch.zeros_like(prob) + elif idx == len(states) - 1: + weight = torch.where(halted, torch.zeros_like(prob), 1.0 - cumulative) + else: + active = ~halted + clipped = prob.float().clamp(0.0, 1.0).to(prob.dtype) + new_halt = active & ((cumulative + clipped) >= target) + weight = torch.where(new_halt, 1.0 - cumulative, clipped) + weight = torch.where(active, weight, torch.zeros_like(weight)) + halted = halted | new_halt + weight = weight.clamp_min(0.0) + cumulative = (cumulative + weight).clamp(0.0, 1.0) + mixed = mixed + state * weight.to(state.dtype).unsqueeze(-1) + depth_used = depth_used + weight * float(step) + accumulated.append(cumulative) + + # These states share transition weights. Detaching the final state makes + # it a moving target and can REWARD expansion: for s_i = i*w, the old + # semi-gradient of mean((s_i - stop_grad(s_D))**2) points toward larger w. + # Use a symmetric depth-consistency objective, differentiating both ends. + # The legacy distill_loss/config name is retained for checkpoint compatibility. + final = states[-1].float() + distill = torch.stack([ + (s.float() - final).pow(2).mean() for s in states[:-1] + ]).mean() + if len(states) >= 3: + deltas = [states[i] - states[i - 1] for i in range(1, len(states))] + stability = torch.stack([ + (deltas[i] - deltas[i - 1]).float().pow(2).mean() + for i in range(1, len(deltas)) + ]).mean() + else: + stability = states[0].new_zeros(()) + # ACT remainder already integrates expected depth. Subtract one so a + # mandatory first pass has zero penalty and depth=1 is v1-equivalent. + ponder = (depth_used.float() - 1.0).clamp_min(0.0).mean() + # Public state reports accumulated halt mass rather than raw logits. This + # is the quantity used by ACT/recurrent-depth monitors and makes the + # contract explicit: the last recurrent pass is always fully halted. + reported_halts = list(accumulated) + reported_halts[-1] = torch.ones_like(reported_halts[-1]) + # Relative step norm; this does not establish a contractive mapping. + contraction = torch.stack([ + (states[i] - states[i - 1]).float().pow(2).mean().clamp_min(1e-12).sqrt() + / states[i - 1].float().pow(2).mean().clamp_min(1e-12).sqrt() + for i in range(1, len(states)) + ]).mean() if len(states) > 1 else states[0].new_zeros(()) + return mixed, RecurrentDepthState( + halting_probs=torch.stack(reported_halts, dim=-1), + ponder_cost=ponder, + depth_used=depth_used, + stability_loss=stability, + distill_loss=distill, + contraction=contraction, + ) diff --git a/chronos/model/temporal_loss.py b/chronos/model/temporal_loss.py index 7cf4c6a..ab013cf 100644 --- a/chronos/model/temporal_loss.py +++ b/chronos/model/temporal_loss.py @@ -2,6 +2,32 @@ import torch.nn.functional as F +def _lookahead_groups(lookahead_probs: torch.Tensor) -> torch.Tensor: + """Normalize legacy 4-D and v2 group-aware 5-D predictions.""" + if lookahead_probs.dim() == 4: + return lookahead_probs.unsqueeze(3) + if lookahead_probs.dim() == 5: + return lookahead_probs + raise ValueError( + "lookahead_probs must be [B,S,K,E] or [B,S,K,G,E], got " + f"{tuple(lookahead_probs.shape)}" + ) + + +def _teacher_groups(teacher: torch.Tensor, groups: int) -> torch.Tensor: + """Map [B,S,E] or [B,S,L,E] teachers onto recurrent groups.""" + if teacher.dim() == 3: + return teacher.unsqueeze(2).expand(-1, -1, groups, -1) + if teacher.dim() != 4: + raise ValueError(f"teacher router probs have invalid shape {tuple(teacher.shape)}") + if teacher.shape[2] == groups: + return teacher + # Deterministic nearest-layer mapping keeps gradients/metrics stable when + # physical layer count differs from recurrent group count. + indices = torch.linspace(0, teacher.shape[2] - 1, groups, device=teacher.device).round().long() + return teacher.index_select(2, indices) + + def temporal_locality_loss(expert_probs_seq: torch.Tensor) -> torch.Tensor: """ Penalizes large routing distribution shifts between adjacent time steps. @@ -68,6 +94,11 @@ def router_offload_metrics( union = (prev | cur).sum(dim=-1).float().clamp_min(1.0) out["router_adjacent_jaccard"] = float((inter / union).mean().item()) if lookahead_probs is not None and lookahead_steps > 0 and S > 1: + # Operational dashboard metrics are layer-averaged; reduce the + # v2 group head only here while all training losses retain the + # full [B,S,K,G,E] supervision. + if lookahead_probs.dim() == 5: + lookahead_probs = lookahead_probs.mean(dim=3) K = min(int(lookahead_steps), int(lookahead_probs.shape[2]) - 1) recalls = [] union_recalls = [] @@ -146,19 +177,20 @@ def lookahead_supervision_loss( return lookahead_probs.new_zeros(()) if lookahead_probs is not None else \ teacher_router_probs.new_zeros(()) - B, S, Kp1, E = lookahead_probs.shape + lookahead_probs = _lookahead_groups(lookahead_probs) + B, S, Kp1, G, E = lookahead_probs.shape K = min(lookahead_steps, Kp1 - 1) if K <= 0 or S <= 1: return lookahead_probs.new_zeros(()) - teacher = teacher_router_probs # caller is responsible for .detach() + teacher = _teacher_groups(teacher_router_probs, G) total = lookahead_probs.new_zeros(()) valid_terms = 0 for k in range(1, K + 1): if S - k <= 0: continue - teacher_k = teacher[:, k:, :] # [B, S-k, E] - pred_k = lookahead_probs[:, :-k, k, :] # [B, S-k, E] + teacher_k = teacher[:, k:, :, :] + pred_k = lookahead_probs[:, :-k, k, :, :] # Soft-target CE: −Σ_e P · log Q log_pred = pred_k.clamp_min(1e-9).log() ce = -(teacher_k * log_pred).sum(dim=-1) # [B, S-k] @@ -188,20 +220,21 @@ def lookahead_topk_hit_loss( return lookahead_probs.new_zeros(()) if lookahead_probs is not None else \ teacher_router_probs.new_zeros(()) - B, S, Kp1, E = lookahead_probs.shape + lookahead_probs = _lookahead_groups(lookahead_probs) + B, S, Kp1, G, E = lookahead_probs.shape K = min(int(lookahead_steps), Kp1 - 1) top_k = max(1, min(int(num_experts_per_tok or 1), E)) if K <= 0 or S <= 1: return lookahead_probs.new_zeros(()) - teacher = teacher_router_probs.detach() + teacher = _teacher_groups(teacher_router_probs.detach(), G) total = lookahead_probs.new_zeros(()) valid_terms = 0 for k in range(1, K + 1): if S - k <= 0: continue - teacher_k = teacher[:, k:, :] - pred_k = lookahead_probs[:, :-k, k, :] + teacher_k = teacher[:, k:, :, :] + pred_k = lookahead_probs[:, :-k, k, :, :] target_ids = torch.topk(teacher_k, k=top_k, dim=-1).indices target_mask = F.one_hot(target_ids, num_classes=E).sum(dim=-2).to(pred_k.dtype) hit_mass = (pred_k * target_mask).sum(dim=-1).clamp_min(1e-9) @@ -224,20 +257,21 @@ def lookahead_union_loss( if lookahead_probs is None or teacher_router_probs is None: return lookahead_probs.new_zeros(()) if lookahead_probs is not None else \ teacher_router_probs.new_zeros(()) - B, S, Kp1, E = lookahead_probs.shape + lookahead_probs = _lookahead_groups(lookahead_probs) + B, S, Kp1, G, E = lookahead_probs.shape K = min(int(lookahead_steps), Kp1 - 1) top_k = max(1, min(int(num_experts_per_tok or 1), E)) if K <= 0 or S <= K: return lookahead_probs.new_zeros(()) - teacher = teacher_router_probs.detach() - union_target = torch.zeros(B, S - K, E, dtype=lookahead_probs.dtype, device=lookahead_probs.device) + teacher = _teacher_groups(teacher_router_probs.detach(), G) + union_target = torch.zeros(B, S - K, G, E, dtype=lookahead_probs.dtype, device=lookahead_probs.device) for k in range(1, K + 1): - target_ids = torch.topk(teacher[:, k:k + S - K, :], k=top_k, dim=-1).indices + target_ids = torch.topk(teacher[:, k:k + S - K, :, :], k=top_k, dim=-1).indices union_target = torch.maximum( union_target, F.one_hot(target_ids, num_classes=E).sum(dim=-2).to(union_target.dtype), ) - pred_union = lookahead_probs[:, :S - K, 1:K + 1, :].amax(dim=2) + pred_union = lookahead_probs[:, :S - K, 1:K + 1, :, :].amax(dim=2) hit_mass = (pred_union * union_target).sum(dim=-1).clamp_min(1e-9) target_count = union_target.sum(dim=-1).clamp_min(1.0) return (-torch.log(hit_mass / target_count)).mean() diff --git a/chronos/monitor/__init__.py b/chronos/monitor/__init__.py new file mode 100644 index 0000000..61ee7fc --- /dev/null +++ b/chronos/monitor/__init__.py @@ -0,0 +1,17 @@ +"""Read-only training monitor shared by the Gradio console tab and the +standalone API web page. + +Everything here reads files a run writes (``run_spec.json``, ``status.json``, +``metrics.jsonl``, ``parameters.jsonl``, ``samples.jsonl``, ``train.log``, +``inspections/*.json``, ``checkpoints/``). It never imports MLX and creates +no Torch tensors (``chronos/__init__`` imports torch for HF registration, but +no device is touched), so a monitor process cannot contend for the GPU, slow a +training process down or change its numerics. +Older runs without ``status.json`` / ``parameters.jsonl`` are still shown: +their status is derived from the metrics stream and the live process table. +""" +from chronos.monitor.jsonl import JsonLines +from chronos.monitor.process import PROCESS_MEMORY, ProcessMemoryReader, find_run_pid, pid_alive +from chronos.monitor.snapshot import RunMonitor + +__all__ = ["JsonLines", "PROCESS_MEMORY", "ProcessMemoryReader", "RunMonitor", "find_run_pid", "pid_alive"] diff --git a/chronos/monitor/evals.py b/chronos/monitor/evals.py new file mode 100644 index 0000000..9f6ca18 --- /dev/null +++ b/chronos/monitor/evals.py @@ -0,0 +1,110 @@ +"""Task-agnostic evaluation view over validation/final events. + +Trainers record whatever they measure as numeric keys on ``validation`` and +``final`` events: ``valid_ce`` for every run, plus task metrics such as +``probe_en_zh_bleu``, ``accuracy``, ``pass_at_1`` or ``reward_mean``. This +module discovers them; the pages never hard-code a task or a language pair. + +A key ``_`` whose last token is a known metric name is grouped +by that metric, so ``probe_en_zh_bleu`` and ``probe_zh_en_bleu`` share one +chart with series ``probe en_zh`` / ``probe zh_en``. Everything else gets a +chart of its own. Bookkeeping keys (sample counts, durations, token counts) +are listed as context, never charted. +""" +from __future__ import annotations + +import math + +# Keys that are bookkeeping, not quality metrics. +RESERVED = frozenset({"time", "event", "phase", "step", "total", "tokens_seen", "elapsed_s", "samples"}) +CONTEXT_SUFFIXES = ("_n", "_s", "_h", "_tokens", "_count", "_samples", "_seconds", "_hours") +CONTEXT_KEYS = frozenset({"n", "hours", "valid_tokens", "wall_s", "wall_h"}) +# Last-token metric names that group across subjects (case-insensitive). +KNOWN_METRICS = frozenset({ + "bleu", "chrf", "ter", "rouge1", "rouge2", "rougel", "meteor", "comet", "bertscore", + "acc", "accuracy", "f1", "precision", "recall", "em", "exact", "auc", + "ce", "loss", "ppl", "perplexity", "nll", "bpb", "bpc", + "wer", "cer", "reward", "score", "pass1", "win", "kl", +}) +PRIMARY = ("valid_ce", "valid_ppl") +# Direction of improvement, for chart captions; unknown families get none. +LOWER_IS_BETTER = frozenset({"ce", "loss", "ppl", "perplexity", "nll", "bpb", "bpc", "wer", "cer", "ter", "kl"}) +HIGHER_IS_BETTER = frozenset({"bleu", "chrf", "rouge1", "rouge2", "rougel", "meteor", "comet", "bertscore", "acc", + "accuracy", "f1", "precision", "recall", "em", "exact", "auc", "reward", "pass1", "win"}) + + +def _num(v) -> bool: + return isinstance(v, (int, float)) and not isinstance(v, bool) and math.isfinite(v) + + +def is_context_key(key: str) -> bool: + return key in CONTEXT_KEYS or key.endswith(CONTEXT_SUFFIXES) + + +def split_key(key: str) -> tuple[str, str]: + """Map a key to (metric family, series label). + + ``probe_en_zh_bleu`` → (``bleu``, ``probe en_zh``); ``bleu_zh`` → (``bleu``, + ``bleu_zh``); ``accuracy`` → (``accuracy``, ``accuracy``); unknown keys + form a family of their own. + """ + head, sep, last = key.rpartition("_") + if sep and head and last.lower() in KNOWN_METRICS: + prefix, _, rest = head.partition("_") + return last.lower(), (f"{prefix} {rest}" if rest else head) + first, sep, _rest = key.partition("_") + if sep and first.lower() in KNOWN_METRICS: + return first.lower(), key + return key, key + + +def eval_metrics(events: list[dict]) -> dict: + """Discover metric keys on validation/final events. + + Returns ``{"families": [{"metric", "series": [{"key", "label"}], "direction"}], + "keys": [...], "context": [...]}``; ``valid_ce``/``valid_ppl`` stay in + ``keys`` but get no family of their own (the loss charts already show them). + """ + keys: dict[str, None] = {} + context: dict[str, None] = {} + for e in events or []: + if e.get("event") not in ("validation", "final"): + continue + for k, v in e.items(): + if k in RESERVED or not _num(v): + continue + (context if is_context_key(k) else keys)[k] = None + families: dict[str, list[dict]] = {} + for k in keys: + if k in PRIMARY: + continue + metric, label = split_key(k) + families.setdefault(metric, []).append({"key": k, "label": label}) + def direction(metric: str) -> str | None: + return "lower" if metric in LOWER_IS_BETTER else "higher" if metric in HIGHER_IS_BETTER else None + return {"families": [{"metric": m, "series": s, "direction": direction(m)} for m, s in families.items()], + "keys": list(keys), "context": list(context)} + + +def sample_view(sample: dict) -> dict: + """Normalise one ``samples.jsonl`` row for display, whatever the task. + + Recognised roles: input (``src``/``prompt``/``input``/``question``), + reference (``ref``/``reference``/``target``/``answer``) and output + (``hyp``/``output``/``prediction``/``response``/``completion``). Every + other scalar field becomes a tag (``direction``, ``set``, ``task``, ``score`` …). + """ + roles = {"input": ("src", "prompt", "input", "question", "source"), + "reference": ("ref", "reference", "target", "answer", "expected"), + "output": ("hyp", "output", "prediction", "response", "completion", "generated")} + used = {"time", "phase", "step"} + out = {"time": sample.get("time"), "phase": sample.get("phase"), "step": sample.get("step")} + for role, names in roles.items(): + for name in names: + if isinstance(sample.get(name), str): + out[role] = sample[name] + used.add(name) + break + out["tags"] = {k: v for k, v in sample.items() + if k not in used and isinstance(v, (str, int, float)) and not isinstance(v, bool)} + return out diff --git a/chronos/monitor/jsonl.py b/chronos/monitor/jsonl.py new file mode 100644 index 0000000..e501758 --- /dev/null +++ b/chronos/monitor/jsonl.py @@ -0,0 +1,81 @@ +"""Incremental JSON-lines reader. + +A training run appends to its JSONL files for days; re-parsing a growing file +on every refresh would cost O(file) per poll. ``JsonLines`` remembers the +(inode, byte offset) it reached and parses only the new bytes. A trailing +line without ``\\n`` is an in-flight write and is left for the next poll; a +replaced or truncated file (new inode / smaller size) is re-read from zero. +""" +from __future__ import annotations + +import json +import os +import threading + + +class JsonLines: + def __init__(self, path: str, max_rows: int | None = None): + self.path = path + self.max_rows = max_rows # keep only the newest rows (None = all) + self.rows: list[dict] = [] + self.bad_lines = 0 + self._key = None # (st_dev, st_ino) + self._offset = 0 + self._lock = threading.Lock() + + def read(self) -> list[dict]: + with self._lock: + try: + st = os.stat(self.path) + except OSError: + self.rows, self._key, self._offset = [], None, 0 + return self.rows + key = (st.st_dev, st.st_ino) + if key != self._key or st.st_size < self._offset: + self.rows, self._key, self._offset, self.bad_lines = [], key, 0, 0 + if st.st_size > self._offset: + with open(self.path, "rb") as fh: + fh.seek(self._offset) + chunk = fh.read(st.st_size - self._offset) + end = chunk.rfind(b"\n") + if end >= 0: + self._offset += end + 1 + for line in chunk[:end].splitlines(): + if not line.strip(): + continue + try: + row = json.loads(line) + except ValueError: + self.bad_lines += 1 + continue + if isinstance(row, dict): + self.rows.append(row) + if self.max_rows is not None and len(self.rows) > self.max_rows: + del self.rows[: len(self.rows) - self.max_rows] + return self.rows + + +def read_json(path: str): + """Whole-file JSON (small files: specs, status, inspections); None if absent/partial.""" + try: + with open(path, encoding="utf-8") as fh: + return json.load(fh) + except (OSError, ValueError): + return None + + +def tail_text(path: str, max_bytes: int = 16000) -> str: + """Last ``max_bytes`` of a text log, starting at a line boundary.""" + try: + size = os.path.getsize(path) + with open(path, "rb") as fh: + if size > max_bytes: + fh.seek(size - max_bytes) + data = fh.read() + cut = data.find(b"\n") + data = data[cut + 1:] if cut >= 0 else data + else: + data = fh.read() + except OSError: + return "" + return data.decode("utf-8", errors="replace") diff --git a/chronos/monitor/process.py b/chronos/monitor/process.py new file mode 100644 index 0000000..ac4f226 --- /dev/null +++ b/chronos/monitor/process.py @@ -0,0 +1,179 @@ +"""Process liveness and OS-level memory of a training process. + +CPU-only system calls; never imports MLX/Torch. On macOS the physical +footprint (``proc_pid_rusage`` RUSAGE_INFO_V0) is what Activity Monitor +reports as "Memory" and includes Metal buffers in unified memory; on Linux +``/proc//status`` gives VmRSS/VmHWM. +""" +from __future__ import annotations + +import ctypes +import os +import shlex +import subprocess +import sys +import threading +import time + + +def pid_alive(pid) -> bool: + if not isinstance(pid, int) or pid <= 0: + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: # exists, owned by another user + return True + except OSError: + return False + return True + + +class _RUsageInfoV0(ctypes.Structure): + # xnu bsd/sys/resource.h: struct rusage_info_v0 (RUSAGE_INFO_V0 = 0) + _fields_ = [("ri_uuid", ctypes.c_uint8 * 16)] + [ + (name, ctypes.c_uint64) for name in ( + "ri_user_time", "ri_system_time", "ri_pkg_idle_wkups", "ri_interrupt_wkups", + "ri_pageins", "ri_wired_size", "ri_resident_size", "ri_phys_footprint", + "ri_proc_start_abstime", "ri_proc_exit_abstime")] + + +def _mach_tick_ns() -> float: + """``ri_user_time``/``ri_system_time`` are Mach absolute-time ticks, not ns + (Apple silicon: 125/3 ns per tick).""" + if sys.platform != "darwin": + return 1.0 + + class _Timebase(ctypes.Structure): + _fields_ = [("numer", ctypes.c_uint32), ("denom", ctypes.c_uint32)] + try: + tb = _Timebase() + if ctypes.CDLL("/usr/lib/libSystem.B.dylib").mach_timebase_info(ctypes.byref(tb)) == 0 and tb.denom: + return tb.numer / tb.denom + except (OSError, AttributeError): + pass + return 1.0 + + +class ProcessMemoryReader: + """Rate-limited per-pid memory sampler (cached for ``interval_s``).""" + + def __init__(self, interval_s: float = 5.0): + self.interval_s = interval_s + self._cache: dict[int, tuple[float, dict]] = {} + self._lock = threading.Lock() + self._libproc = None + if sys.platform == "darwin": + try: + lib = ctypes.CDLL("/usr/lib/libproc.dylib", use_errno=True) + lib.proc_pid_rusage.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p] + lib.proc_pid_rusage.restype = ctypes.c_int + self._libproc = lib + except OSError: + self._libproc = None + self._tick_ns = _mach_tick_ns() + + def _sample(self, pid: int) -> dict: + row = {"pid": pid, "sampled_at": time.time()} + if self._libproc is not None: + usage = _RUsageInfoV0() + if self._libproc.proc_pid_rusage(pid, 0, ctypes.byref(usage)) != 0: + row.update(state="unavailable", reason=os.strerror(ctypes.get_errno())) + return row + row.update(state="measured", source="macOS proc_pid_rusage (phys_footprint)", + footprint_gb=usage.ri_phys_footprint / 1e9, resident_gb=usage.ri_resident_size / 1e9, + wired_gb=usage.ri_wired_size / 1e9, + cpu_s=(usage.ri_user_time + usage.ri_system_time) * self._tick_ns / 1e9) + return row + status = f"/proc/{pid}/status" + if os.path.exists(status): + fields = {} + with open(status, encoding="utf-8") as fh: + for line in fh: + key, _, value = line.partition(":") + if key in ("VmRSS", "VmHWM") and value.strip().endswith("kB"): + fields[key] = int(value.split()[0]) * 1024 / 1e9 + row.update(state="measured", source="/proc status", resident_gb=fields.get("VmRSS"), + peak_resident_gb=fields.get("VmHWM")) + return row + row.update(state="unavailable", reason="no process accounting on this platform") + return row + + def read(self, pid) -> dict: + if not pid_alive(pid): + return {"pid": pid, "state": "unavailable", "reason": "process not running"} + with self._lock: + hit = self._cache.get(pid) + if hit and time.monotonic() - hit[0] < self.interval_s: + return hit[1] + row = self._sample(pid) + if len(self._cache) > 32: + self._cache.clear() + self._cache[pid] = (time.monotonic(), row) + return row + + +def process_state(pid) -> str | None: + """``ps`` state letters (``R``/``S`` running or sleeping, ``T`` stopped by + SIGSTOP, ``Z`` zombie) or None when the process is gone.""" + if not pid_alive(pid): + return None + try: + out = subprocess.run(["ps", "-o", "stat=", "-p", str(int(pid))], capture_output=True, text=True, + timeout=5, check=False).stdout.strip() + except (OSError, subprocess.SubprocessError): + return None + return out or None + + +PROCESS_MEMORY = ProcessMemoryReader() + +_PS_CACHE: tuple[float, list[tuple[int, str]]] = (0.0, []) +_PS_LOCK = threading.Lock() + + +def _ps_rows(ttl_s: float = 10.0) -> list[tuple[int, str]]: + global _PS_CACHE + with _PS_LOCK: + if time.monotonic() - _PS_CACHE[0] < ttl_s: + return _PS_CACHE[1] + rows = [] + try: + out = subprocess.run(["ps", "-axo", "pid=,command="], capture_output=True, text=True, + timeout=5, check=False).stdout + except (OSError, subprocess.SubprocessError): + out = "" + for line in out.splitlines(): + pid, _, cmd = line.strip().partition(" ") + if pid.isdigit() and "longrun" in cmd: + rows.append((int(pid), cmd.strip())) + _PS_CACHE = (time.monotonic(), rows) + return rows + + +def find_run_pid(run_dir: str, repo_root: str | None = None) -> int | None: + """PID of a ``chronos.mlx.training.longrun --spec X`` process whose spec's + ``run_dir`` is ``run_dir`` (for runs started before ``status.json``).""" + import json + target = os.path.realpath(run_dir) + roots = [r for r in (repo_root, os.getcwd()) if r] + for pid, cmd in _ps_rows(): + try: + argv = shlex.split(cmd) + except ValueError: + continue + if not argv or "caffeinate" in os.path.basename(argv[0]) or "--spec" not in argv: + continue + spec = argv[argv.index("--spec") + 1] if argv.index("--spec") + 1 < len(argv) else "" + for root in roots: + path = spec if os.path.isabs(spec) else os.path.join(root, spec) + try: + with open(path, encoding="utf-8") as fh: + spec_dir = json.load(fh).get("run_dir", "") + except (OSError, ValueError): + continue + spec_dir = spec_dir if os.path.isabs(spec_dir) else os.path.join(root, spec_dir) + if os.path.realpath(spec_dir) == target: + return pid + return None diff --git a/chronos/monitor/server.py b/chronos/monitor/server.py new file mode 100644 index 0000000..1200fc8 --- /dev/null +++ b/chronos/monitor/server.py @@ -0,0 +1,253 @@ +"""Standalone training-monitor web page + JSON API. + + chronos monitor --runs-dir runs # http://127.0.0.1:7870/ + chronos serve --models-dir ... --monitor-runs-dir runs # same page at /monitor/ + +Security +* Default bind is 127.0.0.1 with **no authentication** (local use only); the + page says so in its header. Binding any other address requires + ``--allow-remote`` *and* an API key. +* With a key, the browser logs in once (``POST api/login``) and receives an + HttpOnly, SameSite=Strict session cookie; scripts may send + ``Authorization: Bearer `` instead. Failed logins are rate-limited per + client address. Keys are compared in constant time and never stored in the + page or the cookie. +* Read-only: no endpoint starts, stops or modifies a run. +""" +from __future__ import annotations + +import hmac +import os +import secrets +import threading +import time +from collections import deque + +try: # module level: FastAPI resolves ``Request`` annotations by name + from fastapi import FastAPI, HTTPException, Request + from fastapi.responses import FileResponse, JSONResponse, PlainTextResponse, RedirectResponse +except ImportError: # pragma: no cover - serving extra not installed + FastAPI = HTTPException = Request = FileResponse = JSONResponse = PlainTextResponse = RedirectResponse = None + +from chronos.monitor.snapshot import RunMonitor + +STATIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static") +STATIC_FILES = {"app.js": "text/javascript; charset=utf-8", "app.css": "text/css; charset=utf-8"} +COOKIE = "chronos_monitor_session" +LOOPBACK = ("127.0.0.1", "localhost", "::1") +SECURITY_HEADERS = { + "Cache-Control": "no-store", + "X-Content-Type-Options": "nosniff", + "X-Frame-Options": "DENY", + "Referrer-Policy": "no-referrer", + "Content-Security-Policy": ("default-src 'none'; script-src 'self'; style-src 'self'; img-src 'self' data:; " + "connect-src 'self'; base-uri 'none'; form-action 'self'; frame-ancestors 'none'"), +} + + +class _Sessions: + def __init__(self, ttl_s: float): + self.ttl_s = ttl_s + self._tokens: dict[str, float] = {} + self._failures: dict[str, deque] = {} + self._lock = threading.Lock() + + def issue(self) -> str: + token = secrets.token_urlsafe(32) + with self._lock: + now = time.time() + self._tokens = {t: exp for t, exp in self._tokens.items() if exp > now} + self._tokens[token] = now + self.ttl_s + return token + + def valid(self, token: str | None) -> bool: + if not token: + return False + with self._lock: + exp = self._tokens.get(token) + if exp is None or exp < time.time(): + self._tokens.pop(token, None) + return False + return True + + def revoke(self, token: str | None) -> None: + with self._lock: + self._tokens.pop(token or "", None) + + def blocked(self, client: str, limit: int = 8, window_s: float = 300.0) -> bool: + with self._lock: + q = self._failures.setdefault(client, deque()) + while q and q[0] < time.time() - window_s: + q.popleft() + return len(q) >= limit + + def fail(self, client: str) -> None: + with self._lock: + self._failures.setdefault(client, deque()).append(time.time()) + if len(self._failures) > 4096: + self._failures.clear() + + +def create_monitor_app(runs_dir: str, *, api_key: str | None = None, session_hours: float = 12.0, + monitor: RunMonitor | None = None): + """FastAPI app serving the monitor page and its JSON API (mountable).""" + if FastAPI is None: + raise ImportError("chronos monitor needs fastapi and uvicorn (pip install 'Project_Chronos[serve]')") + mon = monitor or RunMonitor(runs_dir) + sessions = _Sessions(session_hours * 3600) + app = FastAPI(title="Chronos training monitor", docs_url=None, redoc_url=None, openapi_url=None) + + @app.middleware("http") + async def headers(request: Request, call_next): + response = await call_next(request) + for key, value in SECURITY_HEADERS.items(): + response.headers.setdefault(key, value) + return response + + 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: Request) -> bool: + if not api_key: + 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 _guard(request: Request): + if not _authorized(request): + raise HTTPException(status_code=401, detail="login required") + + def _run_or_404(fn, *args, **kwargs): + try: + return fn(*args, **kwargs) + except KeyError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + @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): + if name not in STATIC_FILES: + raise HTTPException(status_code=404, detail="not found") + return FileResponse(os.path.join(STATIC_DIR, name), media_type=STATIC_FILES[name]) + + @app.get("/health") + def health(): + return {"ok": True} + + @app.get("/api/session") + def session(request: Request): + return {"auth_required": bool(api_key), "authenticated": _authorized(request), + "local_only": not api_key} + + @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") + try: + body = await request.json() + except ValueError: + body = {} + key = body.get("key") if isinstance(body, dict) else None + 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=(request.scope.get("root_path") or "") + "/") + 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=(request.scope.get("root_path") or "") + "/") + return response + + @app.get("/api/runs") + def runs(request: Request): + _guard(request) + return {"root": mon.root, "runs": mon.runs()} + + @app.get("/api/state") + def state(request: Request, run: str | None = None, history: int = 1): + _guard(request) + return _run_or_404(mon.snapshot, run, history=bool(history)) + + @app.get("/api/metrics") + def metrics(request: Request, run: str | None = None, phase: str | None = None, since: int = 0, limit: int = 5000): + _guard(request) + return {"rows": _run_or_404(mon.metrics, run, phase, since, limit)} + + @app.get("/api/parameters") + def parameters(request: Request, run: str | None = None): + """Training parameters: run spec, model config, precision/runtime, data and tokenizer manifests.""" + _guard(request) + return _run_or_404(mon.training_parameters, run) + + @app.get("/api/tensors") + def tensors(request: Request, run: str | None = None): + _guard(request) + return _run_or_404(mon.parameters, run) or {"tensors": []} + + @app.get("/api/tensor") + def tensor(request: Request, name: str, run: str | None = None): + _guard(request) + return {"name": name, "history": _run_or_404(mon.tensor_history, run, name)} + + @app.get("/api/log") + def log(request: Request, run: str | None = None, bytes: int = 16000): # noqa: A002 - query name + _guard(request) + return PlainTextResponse(_run_or_404(mon.log_tail, run, max(1000, min(int(bytes), 400_000)))) + + return app + + +def check_bind(host: str, allow_remote: bool, api_key: str | None) -> str | None: + """Error message when a bind would expose an unauthenticated monitor.""" + if host in LOOPBACK: + return None + if not (allow_remote and api_key): + return "binding a non-loopback host requires --allow-remote and --api-key" + return None + + +def main(argv=None) -> int: + import argparse + parser = argparse.ArgumentParser(prog="chronos monitor", description="Read-only Chronos training monitor page + API") + parser.add_argument("--runs-dir", default="runs", help="directory holding run folders (run_spec.json, metrics.jsonl)") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=7870) + parser.add_argument("--api-key", default=os.environ.get("CHRONOS_MONITOR_KEY") or os.environ.get("CHRONOS_API_KEY"), + help="login key (env CHRONOS_MONITOR_KEY / CHRONOS_API_KEY); required off loopback") + parser.add_argument("--allow-remote", action="store_true") + parser.add_argument("--session-hours", type=float, default=12.0) + parser.add_argument("--ssl-certfile", default=None, help="serve HTTPS (recommended with --allow-remote)") + parser.add_argument("--ssl-keyfile", default=None) + args = parser.parse_args(argv) + problem = check_bind(args.host, args.allow_remote, args.api_key) + if problem: + parser.error(problem) + import uvicorn + app = create_monitor_app(args.runs_dir, api_key=args.api_key, session_hours=args.session_hours) + scheme = "https" if args.ssl_certfile else "http" + auth = "login required" if args.api_key else "no auth (loopback only)" + print(f"[chronos monitor] {scheme}://{args.host}:{args.port}/ runs={os.path.abspath(args.runs_dir)} {auth}", flush=True) + if args.api_key and scheme == "http" and args.host not in LOOPBACK: + print("[chronos monitor] 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/monitor/snapshot.py b/chronos/monitor/snapshot.py new file mode 100644 index 0000000..135cad0 --- /dev/null +++ b/chronos/monitor/snapshot.py @@ -0,0 +1,330 @@ +"""Run snapshot: everything the monitor pages show, as plain JSON. + +``RunMonitor(root).snapshot(run)`` merges +* the full run spec (phases, LR recipe, model config, budget) — training + parameters are always visible, through the API and both pages; +* status: ``status.json`` heartbeat when the run writes one, otherwise + derived from the metrics stream plus the process table (older runs); +* progress, ETA and wall-budget accounting per phase; +* the newest metric row, a downsampled history per phase, validation, + rollback and final events; +* per-tensor parameter statistics (``parameters.jsonl``), evaluation samples + (``samples.jsonl``, any task), offline inspection reports, checkpoints on disk; +* the evaluation metrics the run actually logs, discovered from its events + (:func:`chronos.monitor.evals.eval_metrics`), so no page assumes a task; +* OS-level memory of the live training process. +""" +from __future__ import annotations + +import glob +import os +import statistics +import threading +import time +from datetime import datetime + +from chronos.monitor.evals import eval_metrics, is_context_key, sample_view +from chronos.monitor.jsonl import JsonLines, read_json, tail_text +from chronos.monitor.process import PROCESS_MEMORY, find_run_pid, pid_alive, process_state + +METRIC_GROUPS = { + "loss": ("loss", "primary_loss", "balance_loss", "lookahead_loss", "lookahead_topk_loss", + "lookahead_union_loss", "temporal_loss", "locality_loss", "anchor_loss", + "recurrent_ponder_loss", "recurrent_distill_loss", "recurrent_stability_loss"), + "optimizer": ("lr", "grad_norm", "update_ratio"), + "throughput": ("tokens_per_s", "processed_tokens_per_s", "step_s", "tokens_seen"), + "memory": ("peak_mem_gb", "active_mem_gb", "cache_mem_gb", "footprint_gb"), + "recurrent": ("rec_depth_mean", "rec_contraction", "rec_early_exit_ratio", "rec_ponder_loss", + "rec_distill_loss", "rec_stability_loss"), + "routing": ("route_entropy", "load_cv", "dead_experts", "max_expert_load"), +} +EVENT_KINDS = ("validation", "rollback", "final", "phase_start", "phase_end", "resume") + + +def downsample(rows: list[dict], limit: int) -> list[dict]: + """Evenly strided subset that always keeps the newest row.""" + if limit <= 0 or len(rows) <= limit: + return list(rows) + stride = len(rows) / limit + picked = [rows[int(i * stride)] for i in range(limit - 1)] + return picked + [rows[-1]] + + +def _age_s(stamp) -> float | None: + if not stamp: + return None + try: + if isinstance(stamp, (int, float)): + return max(0.0, time.time() - float(stamp)) + return max(0.0, time.time() - datetime.fromisoformat(str(stamp)).timestamp()) + except (TypeError, ValueError): + return None + + +def _numeric_leaves(obj, prefix: str, out: dict, depth: int = 0) -> None: + if not isinstance(obj, dict) or depth > 3: + return + for k, v in obj.items(): + key = f"{prefix}.{k}" if prefix else str(k) + if isinstance(v, dict): + _numeric_leaves(v, key, out, depth + 1) + elif (isinstance(v, (int, float)) and not isinstance(v, bool) and v == v and abs(v) != float("inf") + and not is_context_key(str(k))): + out[key] = v + + +def _safe_name(name: str) -> bool: + return bool(name) and "/" not in name and "\\" not in name and name not in (".", "..") and not name.startswith(".") + + +class RunMonitor: + """Read-only view over ``//`` directories.""" + + def __init__(self, root: str, repo_root: str | None = None, history_points: int = 600): + self.root = os.path.abspath(root) + self.repo_root = repo_root or os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + self.history_points = history_points + self._streams: dict[str, JsonLines] = {} + self._lock = threading.Lock() + + # ---- discovery --------------------------------------------------------- + def runs(self) -> list[dict]: + out = [] + try: + names = sorted(os.listdir(self.root)) + except OSError: + return out + for name in names: + path = os.path.join(self.root, name) + if not _safe_name(name) or not os.path.isdir(path): + continue + if not any(os.path.exists(os.path.join(path, f)) for f in ("run_spec.json", "metrics.jsonl", "status.json")): + continue + try: + mtime = max(os.path.getmtime(os.path.join(path, f)) for f in os.listdir(path)) + except (OSError, ValueError): + mtime = 0.0 + out.append({"name": name, "updated": mtime}) + out.sort(key=lambda r: -r["updated"]) + return out + + def run_dir(self, run: str | None) -> str | None: + runs = self.runs() + if not runs: + return None + name = run or runs[0]["name"] + if not _safe_name(name) or name not in {r["name"] for r in runs}: + raise KeyError(f"unknown run {name!r}") + return os.path.join(self.root, name) + + def _stream(self, path: str, max_rows: int | None = None) -> JsonLines: + with self._lock: + reader = self._streams.get(path) + if reader is None: + reader = self._streams[path] = JsonLines(path, max_rows=max_rows) + return reader + + # ---- status -------------------------------------------------------------- + def _status(self, path: str, spec: dict | None, rows: list[dict]) -> dict: + heartbeat = read_json(os.path.join(path, "status.json")) or {} + pid = heartbeat.get("pid") if isinstance(heartbeat.get("pid"), int) else None + source = "status.json" if heartbeat else "derived" + if pid is None or not pid_alive(pid): + found = find_run_pid(path, self.repo_root) + pid = found if found is not None else pid + alive = pid_alive(pid) + ps = process_state(pid) if alive else None + 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: + state = declared + elif finished and not alive: + state = "finished" + elif alive and ps and ps.startswith("T"): + state = "paused" # SIGSTOP (e.g. benchmark pause) + elif alive: + state = declared if declared in ("starting", "evaluating", "checkpointing") else "running" + else: + state = "stopped" if rows or heartbeat else "pending" + stamps = [r.get("time") for r in rows[-6:] if r.get("time") and "event" not in r] + intervals = [] + for a, b in zip(stamps, stamps[1:]): + try: + intervals.append(datetime.fromisoformat(b).timestamp() - datetime.fromisoformat(a).timestamp()) + except ValueError: + pass + interval = statistics.median(intervals) if intervals else None + age = _age_s(heartbeat.get("updated") or last.get("time")) + stale = bool(state == "running" and age is not None and age > max(180.0, 3 * (interval or 60.0))) + return {"state": state, "source": source, "pid": pid, "alive": alive, "ps_state": ps, + "last_update_age_s": age, "log_interval_s": interval, "stale": stale, + "heartbeat": heartbeat or None, "error": heartbeat.get("error")} + + # ---- progress / ETA ------------------------------------------------------ + @staticmethod + def _progress(spec: dict | None, status: dict, latest: dict | None, rows: list[dict]) -> dict: + phases = (spec or {}).get("phases") or [] + hb = status.get("heartbeat") or {} + budget_h = float((spec or {}).get("budget_hours") or 0.0) + reserve = float((spec or {}).get("eval_reserve_fraction") or 0.0) + train_s = budget_h * 3600 * (1 - reserve) + current = hb.get("phase") or (latest or {}).get("phase") + names = [p.get("name") for p in phases] + index = names.index(current) if current in names else 0 + out = [] + for i, phase in enumerate(phases): + prows = [r for r in rows if r.get("phase") == phase.get("name") and "event" not in r] + last = prows[-1] if prows else {} + entry = {"name": phase.get("name"), "budget_h": round(train_s * float(phase.get("budget_fraction", 0)) / 3600, 2), + "step": last.get("step", 0), "total": last.get("total"), + "state": "complete" if i < index else ("active" if i == index else "pending")} + if i == index and status.get("state") == "finished": + entry["state"] = "complete" + out.append(entry) + result = {"phases": out, "current": current, "budget_hours": budget_h, "train_budget_hours": round(train_s / 3600, 2)} + if hb.get("phase_wall_s") is not None and hb.get("phase_budget_s"): + result["phase_fraction"] = min(1.0, float(hb["phase_wall_s"]) / float(hb["phase_budget_s"])) + result["phase_remaining_s"] = max(0.0, float(hb["phase_budget_s"]) - float(hb["phase_wall_s"])) + elif latest and latest.get("total"): + result["phase_fraction"] = min(1.0, latest["step"] / max(1, latest["total"])) + result["phase_remaining_s"] = max(0, latest["total"] - latest["step"]) * float(latest.get("step_s") or 0) + if result.get("phase_remaining_s") is not None: + later = sum(e["budget_h"] for e in out[index + 1:]) * 3600 + result["run_remaining_s"] = result["phase_remaining_s"] + later + budget_h * 3600 * reserve + result["eta"] = datetime.fromtimestamp(time.time() + result["run_remaining_s"]).isoformat(timespec="minutes") + if hb.get("run_wall_s") is not None: + result["run_wall_s"] = hb["run_wall_s"] + return result + + # ---- full snapshot ------------------------------------------------------- + def snapshot(self, run: str | None = None, *, history: bool = True) -> dict: + path = self.run_dir(run) + if path is None: + return {"runs": [], "run": None, "status": {"state": "pending"}, "root": self.root} + rows = self._stream(os.path.join(path, "metrics.jsonl")).read() + spec = read_json(os.path.join(path, "run_spec.json")) + status = self._status(path, spec, rows) + train_rows = [r for r in rows if "event" not in r] + events = [r for r in rows if r.get("event") in EVENT_KINDS or (r.get("event") and r.get("event") not in EVENT_KINDS)] + latest = train_rows[-1] if train_rows else None + memory = PROCESS_MEMORY.read(status["pid"]) if status["alive"] else None + best = {} + for r in events: + if r.get("event") == "validation" and r.get("valid_ce") is not None: + key = r.get("phase", "?") + if key not in best or r["valid_ce"] < best[key]["valid_ce"]: + best[key] = {"step": r.get("step"), "valid_ce": r["valid_ce"], "time": r.get("time")} + params = self._stream(os.path.join(path, "parameters.jsonl"), max_rows=64).read() + latest_params = params[-1] if params else None + out = { + "root": self.root, "runs": self.runs(), "run": os.path.basename(path), "generated": time.time(), + "spec": spec, "runtime": read_json(os.path.join(path, "runtime.json")), "status": status, + "progress": self._progress(spec, status, latest, rows), + "latest": latest, "memory": memory, + "events": events[-200:], "best_valid": best, "evals": eval_metrics(events), + "groups": {k: list(v) for k, v in METRIC_GROUPS.items()}, + "parameters": self._parameter_summary(latest_params), + "samples": [sample_view(x) for x in + 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), + "counts": {"metric_rows": len(train_rows), "events": len(events), "parameter_rows": len(params)}, + } + if history: + out["history"] = {name: downsample([r for r in train_rows if r.get("phase") == name], self.history_points) + for name in dict.fromkeys(r.get("phase") for r in train_rows)} + return out + + @staticmethod + def _parameter_summary(row: dict | None) -> dict | None: + """Latest telemetry row without the per-tensor list (fetched separately: + ~700 tensors are ~170 KiB, too much for every status poll).""" + if not row: + return None + return {k: row.get(k) for k in ("step", "phase", "time", "lr", "grad_norm", "stats", "groups", + "experts", "routing", "grad_note", "telemetry_s")} | { + "tensor_count": len(row.get("tensors") or [])} + + def parameters(self, run: str | None) -> dict | None: + """Full latest ``parameters.jsonl`` row (every tensor).""" + path = self.run_dir(run) + if path is None: + return None + rows = self._stream(os.path.join(path, "parameters.jsonl"), max_rows=64).read() + return rows[-1] if rows else None + + def metrics(self, run: str | None, phase: str | None = None, since_step: int = 0, limit: int = 5000) -> list[dict]: + path = self.run_dir(run) + if path is None: + return [] + rows = [r for r in self._stream(os.path.join(path, "metrics.jsonl")).read() + if "event" not in r and (phase is None or r.get("phase") == phase) and int(r.get("step") or 0) > since_step] + return downsample(rows, max(1, min(int(limit), 20000))) + + def training_parameters(self, run: str | None) -> dict: + """Run spec (recipe, LR schedule, model config) + runtime facts + (precision policy, device, versions, effective batch, code hash).""" + path = self.run_dir(run) + if path is None: + return {} + spec = read_json(os.path.join(path, "run_spec.json")) or {} + out = {"run": os.path.basename(path), "spec": spec, "runtime": read_json(os.path.join(path, "runtime.json"))} + # Small manifests the run points at (paths come from the run's own spec). + for key, folder, name in (("data", spec.get("shards_dir"), "shards_manifest.json"), + ("tokenizer", spec.get("tokenizer_dir"), "tokenizer_report.json")): + if isinstance(folder, str) and folder: + file = os.path.join(folder, name) + if os.path.isfile(file) and os.path.getsize(file) < 512 * 1024: + data = read_json(file) + if isinstance(data, dict) and key == "tokenizer": + data = {k: v for k, v in data.items() if k != "special_tokens"} | { + "special_tokens_count": len(data.get("special_tokens") or [])} + out[key] = data + return out + + def tensor_history(self, run: str | None, name: str) -> list[dict]: + path = self.run_dir(run) + if path is None: + return [] + rows = self._stream(os.path.join(path, "parameters.jsonl"), max_rows=64).read() + out = [] + for row in rows: + for t in row.get("tensors") or []: + if t.get("name") == name: + out.append({"step": row.get("step"), "phase": row.get("phase"), **t}) + break + return out + + def log_tail(self, run: str | None, max_bytes: int = 16000) -> str: + path = self.run_dir(run) + return tail_text(os.path.join(path, "train.log"), max_bytes) if path else "" + + @staticmethod + def _inspections(path: str) -> list[dict]: + """Offline reports: identity fields plus every numeric leaf, flattened + (``held_out.valid_ce``, ``translation.en_zh.bleu``, ``qa.accuracy`` …).""" + meta = ("time", "phase", "step", "tokens_seen", "device", "compute_dtype", "train_dtype", "elapsed_s") + out = [] + for file in sorted(glob.glob(os.path.join(path, "inspections", "*.json")))[-12:]: + data = read_json(file) + if isinstance(data, dict): + metrics: dict[str, float] = {} + _numeric_leaves({k: v for k, v in data.items() if k not in meta + ("precision", "run_dir")}, "", metrics) + out.append({"file": os.path.basename(file), **{k: data.get(k) for k in meta}, "metrics": metrics}) + return out + + @staticmethod + def _checkpoints(path: str) -> list[dict]: + files = glob.glob(os.path.join(path, "*.pth")) + glob.glob(os.path.join(path, "checkpoints", "*.pkl")) + out = [] + for file in files: + try: + st = os.stat(file) + except OSError: + continue + out.append({"file": os.path.relpath(file, path), "size_gb": round(st.st_size / 1e9, 3), + "modified": datetime.fromtimestamp(st.st_mtime).isoformat(timespec="seconds"), + "age_s": round(time.time() - st.st_mtime, 1)}) + out.sort(key=lambda r: r["modified"], reverse=True) + return out diff --git a/chronos/monitor/static/app.css b/chronos/monitor/static/app.css new file mode 100644 index 0000000..24e7fe8 --- /dev/null +++ b/chronos/monitor/static/app.css @@ -0,0 +1,165 @@ +/* Chronos training monitor — light/dark roles from one ramp set. */ +:root { + color-scheme: light; + --page: #f9f9f7; --surface-1: #fcfcfb; --surface-2: #f3f3f0; + --text-primary: #0b0b0b; --text-secondary: #52514e; --muted: #898781; + --grid: #e1e0d9; --axis: #c3c2b7; --border: rgba(11,11,11,0.10); + --series-1: #2a78d6; --series-2: #eb6834; --series-3: #1baf7a; + --good: #0ca30c; --warning: #fab219; --serious: #ec835a; --critical: #d03b3b; + --good-text: #006300; --accent: #2a78d6; --accent-wash: #cde2fb; + --seq-0: #f0efec; --seq-1: #cde2fb; --seq-2: #9ec5f4; --seq-3: #6da7ec; --seq-4: #3987e5; --seq-5: #256abf; --seq-6: #184f95; --seq-7: #0d366b; + font-family: system-ui, -apple-system, "Segoe UI", "PingFang SC", "Hiragino Sans", sans-serif; +} +body, .viz-root { color: var(--text-primary); background: var(--page); } +.viz-root { min-height: 100vh; } +@media (prefers-color-scheme: dark) { + :root:not([data-theme="light"]) { + color-scheme: dark; + --page: #0d0d0d; --surface-1: #1a1a19; --surface-2: #232321; + --text-primary: #f0efec; --text-secondary: #c3c2b7; --muted: #898781; + --grid: #2c2c2a; --axis: #383835; --border: rgba(255,255,255,0.10); + --series-1: #3987e5; --series-2: #d95926; --series-3: #199e70; + --good-text: #0ca30c; --accent: #3987e5; --accent-wash: #1c3a5e; + --seq-0: #2a2a28; --seq-1: #0d366b; --seq-2: #104281; --seq-3: #184f95; --seq-4: #256abf; --seq-5: #3987e5; --seq-6: #6da7ec; --seq-7: #b7d3f6; + } +} +:root[data-theme="dark"] { + color-scheme: dark; + --page: #0d0d0d; --surface-1: #1a1a19; --surface-2: #232321; + --text-primary: #f0efec; --text-secondary: #c3c2b7; --muted: #898781; + --grid: #2c2c2a; --axis: #383835; --border: rgba(255,255,255,0.10); + --series-1: #3987e5; --series-2: #d95926; --series-3: #199e70; + --good-text: #0ca30c; --accent: #3987e5; --accent-wash: #1c3a5e; + --seq-0: #2a2a28; --seq-1: #0d366b; --seq-2: #104281; --seq-3: #184f95; --seq-4: #256abf; --seq-5: #3987e5; --seq-6: #6da7ec; --seq-7: #b7d3f6; +} +* { box-sizing: border-box; } +body { margin: 0; } +button, input, select { font: inherit; color: inherit; } +:focus-visible { outline: 2px solid var(--accent); outline-offset: 2px; } +.muted { color: var(--text-secondary); } +.num { font-variant-numeric: tabular-nums; } + +.bar { display: flex; justify-content: space-between; align-items: center; gap: 16px; flex-wrap: wrap; + padding: 14px max(20px, calc((100vw - 1440px) / 2)); background: var(--surface-1); border-bottom: 1px solid var(--border); + position: sticky; top: 0; z-index: 5; } +.brand h1 { font-size: 18px; font-weight: 600; margin: 2px 0 0; letter-spacing: -0.2px; } +.eyebrow { font-size: 11px; letter-spacing: .08em; text-transform: uppercase; color: var(--muted); } +.bar-controls { display: flex; align-items: center; gap: 10px; flex-wrap: wrap; } +.field { display: inline-flex; align-items: center; gap: 8px; font-size: 12px; color: var(--text-secondary); } +.field.stack { flex-direction: column; align-items: stretch; gap: 6px; } +select, input { background: var(--surface-1); border: 1px solid var(--axis); border-radius: 6px; padding: 6px 8px; min-width: 0; } +input { padding: 9px 10px; } +button { cursor: pointer; border-radius: 6px; border: 1px solid var(--axis); background: var(--surface-1); padding: 6px 10px; } +button:hover { background: var(--surface-2); } +button.ghost { border-color: transparent; color: var(--text-secondary); } +button.primary { background: var(--accent); border-color: var(--accent); color: #fff; padding: 9px 14px; } +.pill { display: inline-flex; align-items: center; gap: 6px; font-size: 12px; padding: 4px 10px; border-radius: 999px; + border: 1px solid var(--border); background: var(--surface-2); } +.pill .dot { width: 8px; height: 8px; border-radius: 50%; background: var(--muted); } +.pill.running .dot { background: var(--good); } .pill.paused .dot, .pill.evaluating .dot, .pill.checkpointing .dot { background: var(--warning); } +.pill.failed .dot, .pill.stopped .dot { background: var(--critical); } .pill.finished .dot { background: var(--accent); } +.badge { font-size: 11px; color: var(--text-secondary); border: 1px dashed var(--axis); border-radius: 6px; padding: 3px 8px; } + +main { max-width: 1440px; margin: 0 auto; padding: 18px 20px 40px; } +.card { background: var(--surface-1); border: 1px solid var(--border); border-radius: 10px; } +.notice { font-size: 13px; line-height: 1.6; padding: 10px 14px; border-radius: 8px; margin-bottom: 12px; + border: 1px solid var(--border); background: var(--surface-1); display: flex; gap: 10px; align-items: baseline; } +.notice .icon { font-weight: 700; } +.notice.warn .icon { color: #a86b00; } .notice.error .icon { color: var(--critical); } .notice.info .icon { color: var(--accent); } + +.progress { padding: 16px 18px; margin-bottom: 14px; } +.progress-head { display: flex; justify-content: space-between; gap: 12px; flex-wrap: wrap; font-size: 13px; } +.progress-head strong { font-size: 15px; font-weight: 600; } +.phases { display: grid; grid-auto-flow: column; grid-auto-columns: 1fr; gap: 4px; margin: 12px 0 8px; } +.phase { font-size: 12px; color: var(--text-secondary); } +.phase .track { height: 8px; border-radius: 4px; background: var(--seq-1); overflow: hidden; margin-bottom: 6px; } +.phase .fill { height: 100%; background: var(--accent); border-radius: 4px; } +.phase.complete .fill { width: 100%; } .phase.pending .track { background: var(--surface-2); } +.phase b { color: var(--text-primary); font-weight: 600; } + +.tiles { display: grid; grid-template-columns: repeat(auto-fill, minmax(170px, 1fr)); gap: 10px; margin-bottom: 16px; } +.tile { padding: 12px 14px; } +.tile .label { font-size: 12px; color: var(--text-secondary); } +.tile .value { font-size: 22px; font-weight: 600; margin-top: 6px; letter-spacing: -0.3px; } +.tile .sub { font-size: 11px; color: var(--muted); margin-top: 4px; line-height: 1.5; } +.tile svg { display: block; width: 100%; height: 26px; margin-top: 6px; } + +.tabs { display: flex; gap: 2px; border-bottom: 1px solid var(--border); margin-bottom: 14px; overflow-x: auto; } +.tabs button { border: 0; border-bottom: 2px solid transparent; border-radius: 0; background: none; padding: 9px 14px; + color: var(--text-secondary); white-space: nowrap; } +.tabs button[aria-selected="true"] { color: var(--text-primary); border-bottom-color: var(--accent); font-weight: 600; } +.panel.loading { opacity: .6; } + +.toolbar { display: flex; gap: 10px; align-items: center; flex-wrap: wrap; font-size: 12px; } +.box .toolbar { margin: 8px 0; } +.charts { display: grid; grid-template-columns: repeat(auto-fill, minmax(330px, 1fr)); gap: 12px; } +.chart { padding: 12px 14px 8px; min-width: 0; margin: 0; } +.chart figcaption { font-size: 11px; margin-top: 4px; } +.panel { display: flex; flex-direction: column; gap: 12px; } +.panel > .caption { margin: 0; } +.chart-head { display: flex; justify-content: space-between; align-items: baseline; gap: 8px; } +.chart-head h3 { font-size: 13px; font-weight: 600; margin: 0; } +.chart-head .last { font-size: 13px; font-variant-numeric: tabular-nums; } +.chart .legend { display: flex; gap: 12px; font-size: 11px; color: var(--text-secondary); margin-top: 4px; } +.legend i { display: inline-block; width: 14px; height: 2px; vertical-align: middle; margin-right: 5px; border-radius: 1px; } +.chart svg { display: block; width: 100%; height: 170px; touch-action: none; } +.chart .axis text { font-size: 10px; fill: var(--muted); font-variant-numeric: tabular-nums; } +.chart .grid line { stroke: var(--grid); stroke-width: 1; } +.chart .baseline { stroke: var(--axis); } +.chart .line { fill: none; stroke-width: 2; stroke-linejoin: round; stroke-linecap: round; } +.chart .s1 { stroke: var(--series-1); } .chart .s2 { stroke: var(--series-2); } .chart .s3 { stroke: var(--series-3); } +.legend .s1 { background: var(--series-1); } .legend .s2 { background: var(--series-2); } .legend .s3 { background: var(--series-3); } +.chart .event { stroke: var(--axis); stroke-dasharray: 2 3; } +.chart .event.rollback { stroke: var(--critical); } +.chart .cross { stroke: var(--muted); stroke-width: 1; } +.chart .marker { stroke: var(--surface-1); stroke-width: 2; } +.chart .marker.s1 { fill: var(--series-1); } .chart .marker.s2 { fill: var(--series-2); } .chart .marker.s3 { fill: var(--series-3); } +.chart .valid { fill: var(--series-2); stroke: var(--surface-1); stroke-width: 2; } + +.tooltip { position: fixed; z-index: 20; pointer-events: none; background: var(--surface-1); color: var(--text-primary); + border: 1px solid var(--border); border-radius: 8px; padding: 8px 10px; font-size: 12px; box-shadow: 0 4px 16px rgba(0,0,0,.12); + min-width: 140px; max-width: 320px; } +.tooltip .t-head { color: var(--text-secondary); margin-bottom: 4px; } +.tooltip .t-row { display: flex; align-items: center; gap: 8px; } +.tooltip .t-row b { font-variant-numeric: tabular-nums; font-weight: 600; } +.tooltip .t-row i { width: 12px; height: 2px; display: inline-block; } +.tooltip .s1 { background: var(--series-1); } .tooltip .s2 { background: var(--series-2); } .tooltip .s3 { background: var(--series-3); } + +.grid2 { display: grid; grid-template-columns: repeat(auto-fit, minmax(420px, 1fr)); gap: 12px; } +.box { padding: 14px 16px; min-width: 0; } +.box h3 { font-size: 14px; margin: 0 0 4px; font-weight: 600; } +.caption { font-size: 12px; color: var(--text-secondary); line-height: 1.6; margin: 0 0 10px; } +.heat { display: grid; gap: 2px; font-size: 10px; align-content: start; } +.grid2 { align-items: start; } +.heat .cell { height: 18px; border-radius: 2px; background: var(--seq-0); } +.heat .rowlab, .heat .collab { color: var(--muted); font-variant-numeric: tabular-nums; white-space: nowrap; align-self: center; } +.heat .collab { text-align: center; } +.scale { display: flex; align-items: center; gap: 6px; font-size: 11px; color: var(--muted); margin-top: 8px; } +.scale .ramp { display: flex; height: 8px; width: 160px; border-radius: 4px; overflow: hidden; } +.scale .ramp span { flex: 1; } + +table { border-collapse: collapse; width: 100%; font-size: 12px; } +th, td { text-align: left; padding: 7px 8px; border-bottom: 1px solid var(--grid); vertical-align: top; } +th { color: var(--text-secondary); font-weight: 500; position: sticky; top: 0; background: var(--surface-1); } +td.num, th.num { text-align: right; font-variant-numeric: tabular-nums; } +.table-wrap { overflow: auto; max-height: 520px; } +tr.clickable { cursor: pointer; } tr.clickable:hover, tr.selected { background: var(--surface-2); } +.kv { display: grid; grid-template-columns: minmax(110px, min(45%, max-content)) minmax(0, 1fr); gap: 4px 14px; font-size: 12px; margin: 0; } +.kv dt { color: var(--text-secondary); overflow-wrap: anywhere; } .kv dd { margin: 0; font-variant-numeric: tabular-nums; overflow-wrap: anywhere; } +details { border-top: 1px solid var(--grid); padding: 8px 0; } +details summary { cursor: pointer; font-size: 13px; font-weight: 600; } +pre { font-size: 11.5px; line-height: 1.55; background: var(--surface-2); border-radius: 8px; padding: 12px; overflow: auto; max-height: 560px; + white-space: pre-wrap; overflow-wrap: anywhere; } +.sample { padding: 10px 0; border-bottom: 1px solid var(--grid); font-size: 13px; line-height: 1.65; } +.sample .meta { font-size: 11px; color: var(--muted); } +.sample .lab { display: inline-block; min-width: 34px; color: var(--text-secondary); font-size: 12px; } +.empty { color: var(--muted); font-size: 13px; padding: 24px 4px; } +.foot { max-width: 1440px; margin: 0 auto; padding: 0 20px 24px; font-size: 11px; color: var(--muted); } + +.login { position: fixed; inset: 0; display: grid; place-items: center; background: var(--page); z-index: 30; } +.login-card { width: min(380px, calc(100vw - 32px)); padding: 24px; display: flex; flex-direction: column; gap: 12px; } +.login-card h2 { margin: 0; font-size: 18px; } +.error-text { color: var(--critical); font-size: 12px; min-height: 1em; margin: 0; } +[hidden] { display: none !important; } +@media (max-width: 640px) { .grid2 { grid-template-columns: 1fr; } .charts { grid-template-columns: 1fr; } } +@media (forced-colors: active) { .chart .line { stroke: CanvasText; } } diff --git a/chronos/monitor/static/app.js b/chronos/monitor/static/app.js new file mode 100644 index 0000000..5024df6 --- /dev/null +++ b/chronos/monitor/static/app.js @@ -0,0 +1,1036 @@ +/* Chronos training monitor page. No dependencies; strict-CSP safe (no inline + * script/style, all text through textContent, SVG built with createElementNS). */ +"use strict"; +(() => { +const SVG = "http://www.w3.org/2000/svg"; +const POLL_MS = 15000; +const $ = (id) => document.getElementById(id); + +// ---- i18n ------------------------------------------------------------------ +const STR = { + zh: { + title: "训练监控", run: "运行", logout: "退出登录", theme_light: "浅色", theme_dark: "深色", lang: "English", + login_title: "登录训练监控", login_hint: "此监控页受 API 密钥保护。会话保存在 HttpOnly Cookie 中，密钥不会写入页面。", + key: "API 密钥", login: "登录", login_bad: "密钥不正确", login_wait: "失败次数过多，请 5 分钟后再试", + auth_local: "仅本机 · 无认证", auth_key: "已登录 · API 密钥", + st_running: "训练中", st_paused: "已暂停", st_evaluating: "评测中", st_checkpointing: "保存检查点", st_starting: "启动中", + st_finished: "已完成", st_failed: "失败", st_stopped: "已停止", st_pending: "等待中", + n_stale: "最近 {age} 没有新的日志行（正常间隔约 {int}）。训练进程可能卡住或被挂起。", + n_derived: "此运行由旧版代码启动，没有 status.json / parameters.jsonl。状态由指标流和进程表推断；张量统计、路由热图和评测探针会在下次以新代码恢复（resume）后出现。", + n_error: "运行报告错误：", n_paused: "训练进程处于 SIGSTOP 挂起状态，恢复后会继续。", n_none: "没有找到运行目录。用 chronos longrun 启动训练后，这里会自动出现。", + phase: "阶段", step: "步", of: "/", budget: "预算", remaining: "剩余", eta: "预计完成", wall: "已用", elapsed: "已运行", + ph_complete: "完成", ph_active: "进行中", ph_pending: "未开始", pretrain: "预训练", sft: "SFT", + k_loss: "训练损失", k_ce: "主损失 (CE)", k_valid: "最佳验证 CE", k_tps: "吞吐", k_eta: "预计完成", k_mem: "内存", + k_grad: "梯度范数", k_depth: "平均循环深度", k_lr: "学习率", k_route: "路由熵", + sub_ppl: "困惑度 {v}", sub_at: "第 {s} 步", sub_proc: "含填充 {v} tok/s", sub_footprint: "进程物理占用 (phys_footprint)", sub_peak: "MLX 峰值 {v} · 活跃 {a}", + sub_clip: "裁剪阈值 {v}", sub_maxdepth: "最大 {v}", sub_entmax: "上限 {v}", sub_left: "剩余 {v}", + t_overview: "总览", t_loss: "损失分量", t_recurrent: "循环深度", t_routing: "MoE 路由", t_params: "参数统计", + t_evals: "评测与样例", t_config: "训练参数", t_log: "日志", t_ckpt: "检查点", + window: "窗口", w_all: "全部", w_last: "最近 {n} 点", phase_all: "全部阶段", events_on: "显示验证/回滚标记", + c_loss: "训练损失", c_primary: "主损失 (CE)", c_valid: "验证 CE", c_lr: "学习率", c_grad: "梯度范数（裁剪前）", + c_tps: "吞吐 (tok/s)", c_step: "单步耗时 (s)", c_mem: "MLX 内存 (GB)", c_tokens: "累计 token", + peak: "峰值", active: "活跃", cache: "缓存", train: "训练", valid: "验证", seen: "有效", processed: "含填充", + no_data: "暂无数据", no_telemetry: "这个运行还没有张量遥测（parameters.jsonl）。新代码恢复训练后，每 {m} 分钟写一行。", + rec_caption: "每组实际循环深度、收缩系数和提前退出比例；深度上限来自模型配置。", + route_caption: "负载份额 = 每个专家被 top-k 选中的比例；熵越接近上限，路由越均匀。", + heat_load: "专家负载（最近一次遥测）", heat_importance: "平均路由概率", heat_wrms: "专家权重 RMS", heat_grms: "专家梯度 RMS", + heat_rows: "行 = 层·迭代，列 = 专家", heat_rows_layer: "行 = 物理层，列 = 专家", + groups: "参数组", tensors: "张量", filter: "筛选张量名…", sort: "排序", name: "名称", group: "组", shape: "形状", dtype: "类型", + numel: "参数量", rms: "RMS", mean: "均值", std: "标准差", absmax: "|最大|", grad_rms: "梯度 RMS", grad_norm: "梯度范数", + update_ratio: "更新比 ‖Δ‖/‖W‖", tensors_n: "张量数", hist_of: "{n} 的历史", hist_hint: "点击表格中的张量查看历史曲线。", + tel_meta: "第 {s} 步 · {t} · 统计耗时 {ms} ms · {note}", load_tensors: "加载全部 {n} 个张量", + validations: "验证记录", probes: "评测探针", samples: "评测样例", inspections: "离线检查报告", metric: "指标", + src: "输入", ref: "参考", hyp: "输出", direction: "方向", set: "数据集", bleu: "BLEU", chrf: "chrF", ppl: "困惑度", + time: "时间", kind: "事件", detail: "详情", ev_validation: "验证", ev_rollback: "回滚", ev_final: "最终评测", + ev_phase_start: "阶段开始", ev_phase_end: "阶段结束", ev_resume: "恢复", + cfg_run: "运行配置", cfg_phases: "阶段配方", cfg_model: "模型结构", cfg_runtime: "运行时与精度", cfg_data: "数据分片", cfg_tok: "分词器", + cfg_raw: "原始 JSON", cfg_missing: "运行未记录 runtime.json（旧版代码启动）。", + log_more: "加载更多", log_refresh: "刷新", log_auto: "自动刷新", + ck_file: "文件", ck_size: "大小", ck_modified: "修改时间", ck_age: "距今", ck_none: "还没有检查点。", + foot: "只读监控 · 每 {s} 秒刷新 · 数据目录 {root} · 生成于 {t}", paused_poll: "页面不可见，已暂停刷新", + err_net: "无法连接监控服务：", generated: "更新于", + }, + en: { + title: "Training monitor", run: "Run", logout: "Sign out", theme_light: "Light", theme_dark: "Dark", lang: "中文", + login_title: "Sign in to the training monitor", login_hint: "This monitor is protected by an API key. The session lives in an HttpOnly cookie; the key is never stored in the page.", + key: "API key", login: "Sign in", login_bad: "Wrong key", login_wait: "Too many failed attempts; wait 5 minutes", + auth_local: "Local only · no auth", auth_key: "Signed in · API key", + st_running: "Running", st_paused: "Paused", st_evaluating: "Evaluating", st_checkpointing: "Checkpointing", st_starting: "Starting", + st_finished: "Finished", st_failed: "Failed", st_stopped: "Stopped", st_pending: "Pending", + n_stale: "No new log line for {age} (usual interval about {int}). The trainer may be stuck or suspended.", + n_derived: "This run was started by older code and writes no status.json / parameters.jsonl. Status is derived from the metric stream and the process table; tensor stats, routing heatmaps and evaluation probes appear after the next resume with the new code.", + n_error: "The run reported an error: ", n_paused: "The trainer is suspended (SIGSTOP) and will continue when resumed.", n_none: "No run directories found. Runs started with chronos longrun show up here.", + phase: "Phase", step: "step", of: "/", budget: "budget", remaining: "remaining", eta: "ETA", wall: "used", elapsed: "Elapsed", + ph_complete: "complete", ph_active: "active", ph_pending: "pending", pretrain: "Pretrain", sft: "SFT", + k_loss: "Training loss", k_ce: "Primary loss (CE)", k_valid: "Best valid CE", k_tps: "Throughput", k_eta: "ETA", k_mem: "Memory", + k_grad: "Grad norm", k_depth: "Mean recurrent depth", k_lr: "Learning rate", k_route: "Routing entropy", + sub_ppl: "perplexity {v}", sub_at: "step {s}", sub_proc: "{v} tok/s incl. padding", sub_footprint: "process phys_footprint", sub_peak: "MLX peak {v} · active {a}", + sub_clip: "clip at {v}", sub_maxdepth: "max {v}", sub_entmax: "max {v}", sub_left: "{v} left", + t_overview: "Overview", t_loss: "Loss terms", t_recurrent: "Recurrence", t_routing: "MoE routing", t_params: "Parameters", + t_evals: "Evals & samples", t_config: "Training parameters", t_log: "Log", t_ckpt: "Checkpoints", + window: "Window", w_all: "All", w_last: "Last {n} points", phase_all: "All phases", events_on: "Show validation/rollback marks", + c_loss: "Training loss", c_primary: "Primary loss (CE)", c_valid: "Validation CE", c_lr: "Learning rate", c_grad: "Grad norm (pre-clip)", + c_tps: "Throughput (tok/s)", c_step: "Step time (s)", c_mem: "MLX memory (GB)", c_tokens: "Tokens seen", + peak: "peak", active: "active", cache: "cache", train: "train", valid: "valid", seen: "counted", processed: "incl. padding", + no_data: "No data yet", no_telemetry: "No tensor telemetry (parameters.jsonl) for this run yet. After a resume with the new code, one row is written every {m} minutes.", + rec_caption: "Depth actually used per group, contraction and early-exit share; the depth cap comes from the model config.", + route_caption: "Load share = fraction of top-k picks per expert; entropy close to the cap means even routing.", + heat_load: "Expert load (latest telemetry)", heat_importance: "Mean router probability", heat_wrms: "Expert weight RMS", heat_grms: "Expert gradient RMS", + heat_rows: "rows = layer·iteration, columns = experts", heat_rows_layer: "rows = physical layer, columns = experts", + groups: "Parameter groups", tensors: "Tensors", filter: "Filter tensor names…", sort: "Sort", name: "Name", group: "Group", shape: "Shape", dtype: "dtype", + numel: "Params", rms: "RMS", mean: "Mean", std: "Std", absmax: "|max|", grad_rms: "Grad RMS", grad_norm: "Grad norm", + update_ratio: "Update ratio ‖Δ‖/‖W‖", tensors_n: "Tensors", hist_of: "History of {n}", hist_hint: "Click a tensor row to see its history.", + tel_meta: "step {s} · {t} · stats took {ms} ms · {note}", load_tensors: "Load all {n} tensors", + validations: "Validations", probes: "Evaluation probes", samples: "Evaluation samples", inspections: "Offline inspections", metric: "Metric", + src: "Input", ref: "Reference", hyp: "Output", direction: "Direction", set: "Set", bleu: "BLEU", chrf: "chrF", ppl: "PPL", + time: "Time", kind: "Event", detail: "Detail", ev_validation: "validation", ev_rollback: "rollback", ev_final: "final eval", + ev_phase_start: "phase start", ev_phase_end: "phase end", ev_resume: "resume", + cfg_run: "Run", cfg_phases: "Phase recipes", cfg_model: "Model architecture", cfg_runtime: "Runtime & precision", cfg_data: "Data shards", cfg_tok: "Tokenizer", + cfg_raw: "Raw JSON", cfg_missing: "No runtime.json for this run (started by older code).", + log_more: "Load more", log_refresh: "Refresh", log_auto: "Auto refresh", + ck_file: "File", ck_size: "Size", ck_modified: "Modified", ck_age: "Age", ck_none: "No checkpoints yet.", + foot: "Read-only monitor · refresh every {s} s · runs dir {root} · generated {t}", paused_poll: "Page hidden; polling paused", + err_net: "Cannot reach the monitor service: ", generated: "updated", + }, +}; +// Metric labels. recurrent_* = weighted term inside the total loss; rec_* = raw value. +const METRIC = { + zh: { loss: "总损失", primary_loss: "主损失 (CE)", balance_loss: "负载均衡（加权）", lookahead_loss: "lookahead（加权）", + lookahead_topk_loss: "lookahead top-k（加权）", lookahead_union_loss: "lookahead 并集（加权）", temporal_loss: "时序一致（加权）", + locality_loss: "路由局部性（加权）", anchor_loss: "路由锚定（加权）", recurrent_ponder_loss: "ponder（加权）", + recurrent_distill_loss: "深度蒸馏（加权）", recurrent_stability_loss: "稳定性（加权）", recurrent_contraction_loss: "收缩约束（加权）", + rec_depth_mean: "平均深度", rec_contraction: "收缩系数", rec_early_exit_ratio: "提前退出比例", rec_ponder_loss: "ponder（原始）", + rec_distill_loss: "深度蒸馏（原始）", rec_stability_loss: "稳定性（原始）", route_entropy: "路由熵", load_cv: "负载变异系数", + dead_experts: "空闲专家数", max_expert_load: "最大专家负载（均值倍数）" }, + en: { loss: "Total loss", primary_loss: "Primary loss (CE)", balance_loss: "Load balance (weighted)", lookahead_loss: "Lookahead (weighted)", + lookahead_topk_loss: "Lookahead top-k (weighted)", lookahead_union_loss: "Lookahead union (weighted)", temporal_loss: "Temporal (weighted)", + locality_loss: "Router locality (weighted)", anchor_loss: "Router anchor (weighted)", recurrent_ponder_loss: "Ponder (weighted)", + recurrent_distill_loss: "Depth distill (weighted)", recurrent_stability_loss: "Stability (weighted)", recurrent_contraction_loss: "Contraction (weighted)", + rec_depth_mean: "Mean depth", rec_contraction: "Contraction", rec_early_exit_ratio: "Early-exit share", rec_ponder_loss: "Ponder (raw)", + rec_distill_loss: "Depth distill (raw)", rec_stability_loss: "Stability (raw)", route_entropy: "Routing entropy", load_cv: "Load CV", + dead_experts: "Idle experts", max_expert_load: "Max expert load (× mean)" }, +}; + +// ---- state ------------------------------------------------------------------- +const S = { + lang: localStorage.getItem("chronos.monitor.lang") || ((navigator.language || "").startsWith("zh") ? "zh" : "en"), + theme: localStorage.getItem("chronos.monitor.theme") || "", + tab: localStorage.getItem("chronos.monitor.tab") || "overview", + run: null, snap: null, session: null, timer: null, inflight: false, error: null, + window: 0, phase: "", showEvents: true, + params: null, tensors: null, tensorFilter: "", tensorSort: "name", tensorPick: null, tensorHist: null, + log: null, logBytes: 16000, logAuto: true, +}; +const tr = (key, vars) => { + let s = (STR[S.lang] && STR[S.lang][key]) ?? STR.en[key] ?? key; + if (vars) for (const [k, v] of Object.entries(vars)) s = s.split(`{${k}}`).join(String(v)); + return s; +}; +const mlabel = (key) => (METRIC[S.lang] && METRIC[S.lang][key]) || METRIC.en[key] || key; + +// ---- DOM helpers --------------------------------------------------------------- +function h(tag, attrs, ...kids) { + const el = document.createElement(tag); + if (attrs) for (const [k, v] of Object.entries(attrs)) { + if (v == null || v === false) continue; + if (k === "class") el.className = v; + else if (k === "text") el.textContent = v; + else if (k.startsWith("on")) el.addEventListener(k.slice(2), v); + else el.setAttribute(k, v === true ? "" : v); + } + for (const kid of kids.flat()) if (kid != null && kid !== false) el.append(kid instanceof Node ? kid : document.createTextNode(String(kid))); + return el; +} +function s(tag, attrs) { + const el = document.createElementNS(SVG, tag); + if (attrs) for (const [k, v] of Object.entries(attrs)) if (v != null) el.setAttribute(k, v); + return el; +} +const clear = (el) => { while (el.firstChild) el.removeChild(el.firstChild); return el; }; + +// ---- formatting ---------------------------------------------------------------- +const isNum = (v) => typeof v === "number" && Number.isFinite(v); +function fmt(v, digits = 4) { + if (!isNum(v)) return "—"; + const a = Math.abs(v); + if (a !== 0 && (a < 1e-4 || a >= 1e6)) return v.toExponential(2); + return Number(v.toPrecision(digits)).toLocaleString(S.lang === "zh" ? "zh-CN" : "en-US", { maximumFractionDigits: 6 }); +} +function fmtInt(v) { return isNum(v) ? Math.round(v).toLocaleString(S.lang === "zh" ? "zh-CN" : "en-US") : "—"; } +function fmtCount(v) { + if (!isNum(v)) return "—"; + const a = Math.abs(v); + if (a >= 1e9) return (v / 1e9).toFixed(2) + "B"; + if (a >= 1e6) return (v / 1e6).toFixed(a >= 1e8 ? 0 : 1) + "M"; + if (a >= 1e3) return (v / 1e3).toFixed(1) + "k"; + return String(Math.round(v)); +} +function fmtGB(v) { return isNum(v) ? (v >= 100 ? v.toFixed(0) : v.toFixed(2)) + " GB" : "—"; } +function fmtDur(sec) { + if (!isNum(sec)) return "—"; + sec = Math.max(0, sec); + const d = Math.floor(sec / 86400), hh = Math.floor((sec % 86400) / 3600), mm = Math.floor((sec % 3600) / 60); + if (S.lang === "zh") return d ? `${d}天${hh}时` : hh ? `${hh}时${mm}分` : mm ? `${mm}分` : `${Math.round(sec)}秒`; + return d ? `${d}d ${hh}h` : hh ? `${hh}h ${mm}m` : mm ? `${mm}m` : `${Math.round(sec)}s`; +} +function fmtTime(stamp) { + if (!stamp) return "—"; + const d = typeof stamp === "number" ? new Date(stamp * 1000) : new Date(stamp); + if (Number.isNaN(d.getTime())) return String(stamp); + return d.toLocaleString(S.lang === "zh" ? "zh-CN" : "en-US", { month: "2-digit", day: "2-digit", hour: "2-digit", minute: "2-digit", hour12: false }); +} +const phaseName = (p) => (p === "pretrain" || p === "sft") ? tr(p) : (p || "—"); +// ---- API ------------------------------------------------------------------------- +class AuthError extends Error {} +const url = (path, params) => { + const u = new URL(path, document.baseURI); // relative: works at / and when mounted at /monitor/ + if (params) for (const [k, v] of Object.entries(params)) if (v != null && v !== "") u.searchParams.set(k, v); + return u; +}; +async function api(path, params, text = false) { + const res = await fetch(url(path, params), { credentials: "same-origin", headers: { Accept: text ? "text/plain" : "application/json" } }); + if (res.status === 401) throw new AuthError("401"); + if (!res.ok) { + let msg = `${res.status} ${res.statusText}`; + try { const j = await res.json(); if (j && j.detail) msg = String(j.detail); } catch (_) { /* not JSON */ } + throw new Error(msg); + } + return text ? res.text() : res.json(); +} + +// ---- session / login --------------------------------------------------------------- +function showLogin(message) { + stop(); + $("login").hidden = false; + $("root").setAttribute("inert", ""); + $("login-title").textContent = tr("login_title"); + $("login-hint").textContent = tr("login_hint"); + $("key-label").textContent = tr("key"); + $("login-submit").textContent = tr("login"); + $("login-error").textContent = message || ""; + $("key").focus(); +} +function hideLogin() { $("login").hidden = true; $("root").removeAttribute("inert"); } +async function submitLogin(ev) { + ev.preventDefault(); + const btn = $("login-submit"); + btn.disabled = true; + try { + const res = await fetch(url("api/login"), { method: "POST", credentials: "same-origin", + headers: { "Content-Type": "application/json" }, body: JSON.stringify({ key: $("key").value }) }); + if (res.status === 429) { $("login-error").textContent = tr("login_wait"); return; } + if (!res.ok) { $("login-error").textContent = tr("login_bad"); return; } + $("key").value = ""; + hideLogin(); + await start(); + } catch (e) { + $("login-error").textContent = tr("err_net") + e.message; + } finally { btn.disabled = false; } +} +async function logout() { + try { await fetch(url("api/logout"), { method: "POST", credentials: "same-origin" }); } catch (_) { /* offline */ } + S.snap = S.params = S.tensors = S.tensorHist = S.log = null; + showLogin(""); +} + +// ---- polling --------------------------------------------------------------------- +function stop() { clearTimeout(S.timer); S.timer = null; } +function schedule() { + stop(); + if (!document.hidden) S.timer = setTimeout(async () => { await refresh(); schedule(); }, POLL_MS); +} +async function refresh() { + if (S.inflight) return; + S.inflight = true; + try { + const snap = await api("api/state", { run: S.run, history: 1 }); + S.snap = snap; + S.error = null; + if (snap.run) S.run = snap.run; + if (S.tab === "log" && S.logAuto) await loadLog(); + renderAll(false); + } catch (e) { + if (e instanceof AuthError) showLogin(""); + else { S.error = e.message; renderNotices(); } + } finally { S.inflight = false; } +} +async function start() { + try { S.session = await api("api/session"); } catch (e) { S.error = e.message; renderNotices(); return; } + if (S.session.auth_required && !S.session.authenticated) { renderChrome(); showLogin(""); return; } + await refresh(); + schedule(); +} + +// ---- header, notices, progress ------------------------------------------------------ +const effectiveDark = () => S.theme ? S.theme === "dark" : matchMedia("(prefers-color-scheme: dark)").matches; +function renderChrome() { + document.documentElement.lang = S.lang === "zh" ? "zh-CN" : "en"; + if (S.theme) document.documentElement.dataset.theme = S.theme; else delete document.documentElement.dataset.theme; + $("title").textContent = tr("title"); + document.title = `Chronos · ${tr("title")}${S.snap && S.snap.run ? " · " + S.snap.run : ""}`; + $("run-label").textContent = tr("run"); + $("lang").textContent = tr("lang"); + $("theme").textContent = effectiveDark() ? tr("theme_light") : tr("theme_dark"); + $("logout").textContent = tr("logout"); + $("logout").hidden = !(S.session && S.session.auth_required); + $("auth-badge").textContent = S.session ? (S.session.auth_required ? tr("auth_key") : tr("auth_local")) : ""; + const sel = $("run"); + const names = ((S.snap && S.snap.runs) || []).map((r) => r.name); + if (sel.dataset.names !== names.join("\n")) { + clear(sel); + for (const n of names) sel.append(h("option", { value: n, text: n })); + sel.dataset.names = names.join("\n"); + } + if (S.snap && S.snap.run) sel.value = S.snap.run; + const st = (S.snap && S.snap.status) || { state: "pending" }; + const pill = clear($("state-pill")); + pill.className = `pill ${st.state}`; + pill.append(h("span", { class: "dot", "aria-hidden": "true" }), h("span", { text: tr("st_" + st.state) + (st.stale ? " · !" : "") })); +} +function renderNotices() { + const box = clear($("notices")); + const add = (kind, icon, text) => box.append(h("div", { class: `notice ${kind}`, role: kind === "error" ? "alert" : null }, + h("span", { class: "icon", "aria-hidden": "true", text: icon }), h("span", { text }))); + if (S.error) add("error", "✕", tr("err_net") + S.error); + const snap = S.snap; + if (!snap) return; + if (!snap.run) { add("info", "i", tr("n_none")); return; } + const st = snap.status || {}; + if (st.error) add("error", "✕", tr("n_error") + st.error); + if (st.state === "paused") add("warn", "‖", tr("n_paused")); + if (st.stale) add("warn", "!", tr("n_stale", { age: fmtDur(st.last_update_age_s), int: fmtDur(st.log_interval_s || 60) })); + if (st.source === "derived" && !snap.parameters) add("info", "i", tr("n_derived")); +} +function renderProgress() { + const box = clear($("progress")); + const snap = S.snap; + box.hidden = !(snap && snap.run); + if (box.hidden) return; + const p = snap.progress || {}; + const latest = snap.latest || {}; + const pct = isNum(p.phase_fraction) ? ` · ${(p.phase_fraction * 100).toFixed(1)}%` : ""; + const right = [p.eta ? `${tr("eta")} ${fmtTime(p.eta)}` : null, + isNum(p.run_remaining_s) ? `${tr("remaining")} ${fmtDur(p.run_remaining_s)}` : null, + isNum(p.run_wall_s) ? `${tr("wall")} ${fmtDur(p.run_wall_s)}` : null, + isNum(p.budget_hours) && p.budget_hours ? `${tr("budget")} ${p.budget_hours} h` : null].filter(Boolean).join(" · "); + box.append(h("div", { class: "progress-head" }, + h("div", null, h("strong", { text: `${phaseName(p.current)} · ${tr("step")} ${fmtInt(latest.step)} ${tr("of")} ${fmtInt(latest.total)}` }), + h("span", { class: "muted num", text: pct })), + h("div", { class: "muted num", text: right }))); + const phases = h("div", { class: "phases" }); + const list = p.phases || []; + phases.style.gridTemplateColumns = list.map((ph) => `${Math.max(0.05, ph.budget_h || 1)}fr`).join(" ") || "1fr"; + for (const ph of list) { + const frac = ph.state === "complete" ? 1 : ph.state === "active" + ? (isNum(p.phase_fraction) ? p.phase_fraction : (ph.total ? ph.step / ph.total : 0)) : 0; + const fill = h("div", { class: "fill" }); + fill.style.width = `${(Math.max(0, Math.min(1, frac)) * 100).toFixed(2)}%`; + phases.append(h("div", { class: `phase ${ph.state}` }, + h("div", { class: "track", role: "progressbar", "aria-label": phaseName(ph.name), "aria-valuemin": "0", + "aria-valuemax": "100", "aria-valuenow": String(Math.round(frac * 100)) }, fill), + h("b", { text: phaseName(ph.name) }), ` ${tr("ph_" + ph.state)} · ${ph.budget_h} h`, + ph.total ? ` · ${fmtInt(ph.step)}/${fmtInt(ph.total)}` : "")); + } + box.append(phases); +} + +// ---- KPI tiles --------------------------------------------------------------------- +function phaseRows(phase) { return ((S.snap && S.snap.history) || {})[phase] || []; } +function sparkline(rows, key) { + const pts = rows.slice(-160).map((r) => r[key]).filter(isNum); + if (pts.length < 2) return null; + let lo = Math.min(...pts), hi = Math.max(...pts); + if (lo === hi) { lo -= 1; hi += 1; } + const d = pts.map((v, i) => `${i ? "L" : "M"}${(i / (pts.length - 1) * 100).toFixed(2)} ${(24 - (v - lo) / (hi - lo) * 22).toFixed(2)}`).join(""); + const svg = s("svg", { viewBox: "0 0 100 26", preserveAspectRatio: "none", "aria-hidden": "true" }); + svg.append(s("path", { d, fill: "none", stroke: "var(--series-1)", "stroke-width": "2", "vector-effect": "non-scaling-stroke", + "stroke-linejoin": "round", "stroke-linecap": "round" })); + return svg; +} +function renderTiles() { + const box = clear($("tiles")); + const snap = S.snap; + if (!snap || !snap.run) return; + const L = snap.latest || {}; + const phase = (snap.progress || {}).current || L.phase; + const rows = phaseRows(phase); + const spec = snap.spec || {}; + const cfg = spec.model_config || {}; + const phSpec = (spec.phases || []).find((p) => p.name === phase) || {}; + const best = (snap.best_valid || {})[phase]; + const mem = snap.memory || {}; + const p = snap.progress || {}; + const tile = (label, value, sub, key) => box.append(h("div", { class: "tile card" }, + h("div", { class: "label", text: label }), h("div", { class: "value num", text: value }), + sub ? h("div", { class: "sub", text: sub }) : null, key ? sparkline(rows, key) : null)); + tile(tr("k_loss"), fmt(L.loss), tr("sub_at", { s: fmtInt(L.step) }), "loss"); + tile(tr("k_ce"), fmt(L.primary_loss), isNum(L.primary_loss) ? tr("sub_ppl", { v: fmt(Math.exp(L.primary_loss), 3) }) : null, "primary_loss"); + tile(tr("k_valid"), best ? fmt(best.valid_ce) : "—", best ? `${tr("sub_at", { s: fmtInt(best.step) })} · ${tr("sub_ppl", { v: fmt(Math.exp(best.valid_ce), 3) })}` : null); + tile(tr("k_tps"), isNum(L.tokens_per_s) ? `${fmtCount(L.tokens_per_s)} tok/s` : "—", + isNum(L.processed_tokens_per_s) ? tr("sub_proc", { v: fmtCount(L.processed_tokens_per_s) }) : null, "tokens_per_s"); + tile(tr("k_eta"), p.eta ? fmtTime(p.eta) : "—", isNum(p.run_remaining_s) ? tr("sub_left", { v: fmtDur(p.run_remaining_s) }) : null); + const osMem = isNum(mem.footprint_gb); // live process footprint has no history; MLX peak does + tile(tr("k_mem"), osMem ? fmtGB(mem.footprint_gb) : fmtGB(L.peak_mem_gb), + `${osMem ? tr("sub_footprint") + " · " : ""}${tr("sub_peak", { v: fmtGB(L.peak_mem_gb), a: fmtGB(L.active_mem_gb) })}`, osMem ? null : "peak_mem_gb"); + tile(tr("k_grad"), fmt(L.grad_norm), isNum(phSpec.grad_clip) ? tr("sub_clip", { v: phSpec.grad_clip }) : null, "grad_norm"); + tile(tr("k_lr"), fmt(L.lr, 3), isNum(phSpec.peak_lr) ? `peak ${fmt(phSpec.peak_lr, 3)}` : null, "lr"); + tile(tr("k_depth"), fmt(L.rec_depth_mean, 3), isNum(cfg.recurrent_depth) ? tr("sub_maxdepth", { v: cfg.recurrent_depth }) : null, "rec_depth_mean"); + tile(tr("k_route"), fmt(L.route_entropy, 3), isNum(L.route_entropy_max) ? `${tr("sub_entmax", { v: fmt(L.route_entropy_max, 3) })} · CV ${fmt(L.load_cv, 3)}` : null, "route_entropy"); +} + +// ---- tabs ---------------------------------------------------------------------------- +const TABS = ["overview", "loss", "recurrent", "routing", "params", "evals", "config", "log", "ckpt"]; +function renderTabs() { + const nav = clear($("tabs")); + if (!TABS.includes(S.tab)) S.tab = "overview"; + for (const id of TABS) { + nav.append(h("button", { type: "button", role: "tab", id: `tab-${id}`, "aria-selected": String(S.tab === id), + "aria-controls": "panel", tabindex: S.tab === id ? "0" : "-1", text: tr("t_" + id), + onclick: () => selectTab(id), onkeydown: (e) => tabKeys(e, id) })); + } + $("panel").setAttribute("aria-labelledby", `tab-${S.tab}`); +} +function tabKeys(e, id) { + const i = TABS.indexOf(id); + const next = e.key === "ArrowRight" ? TABS[(i + 1) % TABS.length] : e.key === "ArrowLeft" ? TABS[(i - 1 + TABS.length) % TABS.length] : null; + if (!next) return; + e.preventDefault(); + selectTab(next); + $(`tab-${next}`).focus(); +} +function selectTab(id) { + S.tab = id; + localStorage.setItem("chronos.monitor.tab", id); + history.replaceState(null, "", `#tab=${id}`); // shareable deep link + renderTabs(); + renderPanel(true); + if (id === "log") loadLog().then(() => renderPanel(true)); +} +// ---- tooltip ------------------------------------------------------------------------- +const tip = { + show(x, y, head, rows) { + const el = clear($("tooltip")); + el.append(h("div", { class: "t-head", text: head })); + for (const r of rows) { + const sw = h("i", { "aria-hidden": "true" }); + if (r.cls) sw.className = r.cls; + el.append(h("div", { class: "t-row" }, r.cls ? sw : null, h("span", { text: r.label }), h("b", { text: r.value }))); + } + el.hidden = false; + const w = el.offsetWidth, ht = el.offsetHeight; + el.style.left = `${Math.min(window.innerWidth - w - 8, x + 14)}px`; + el.style.top = `${Math.max(8, Math.min(window.innerHeight - ht - 8, y - ht - 12))}px`; + }, + hide() { $("tooltip").hidden = true; }, +}; + +// ---- scales ---------------------------------------------------------------------------- +function niceTicks(lo, hi, count = 4) { + if (!(hi > lo)) return [lo]; + const raw = (hi - lo) / count, mag = 10 ** Math.floor(Math.log10(raw)); + const step = [1, 2, 2.5, 5, 10].map((m) => m * mag).find((st) => st >= raw) || raw; + const out = []; + for (let v = Math.ceil(lo / step) * step; v <= hi + step * 1e-9; v += step) out.push(Number(v.toPrecision(12))); + return out; +} +function bisect(xs, x) { + let lo = 0, hi = xs.length - 1; + while (hi - lo > 1) { const mid = (lo + hi) >> 1; if (xs[mid] < x) lo = mid; else hi = mid; } + return Math.abs(xs[lo] - x) <= Math.abs(xs[hi] - x) ? lo : hi; +} + +// ---- line chart (small multiple; one y axis; crosshair tooltip) --------------------------- +const pendingDraws = []; +function lineChart({ title, series, events = [], fmtY = fmt, note }) { + const live = series.map((sr, i) => ({ ...sr, cls: sr.cls || `s${i + 1}`, + pts: (sr.rows || []).filter((r) => isNum(r[sr.key]) && isNum(r.step)).map((r) => ({ x: r.step, y: r[sr.key], t: r.time })) })) + .filter((sr) => sr.pts.length); + const card = h("figure", { class: "chart card", tabindex: live.length ? "0" : null, "aria-label": title }); + const lastPts = live.length ? live[0].pts[live[0].pts.length - 1] : null; + card.append(h("div", { class: "chart-head" }, h("h3", { text: title }), + h("span", { class: "last", text: lastPts ? fmtY(lastPts.y) : "" }))); + if (live.length > 1) { + card.append(h("div", { class: "legend" }, live.map((sr) => h("span", null, h("i", { class: sr.cls, "aria-hidden": "true" }), sr.label)))); + } + if (!live.length) { card.append(h("div", { class: "empty", text: tr("no_data") })); return card; } + const svg = s("svg", { role: "img", "aria-label": `${title}: ${lastPts ? fmtY(lastPts.y) : ""}` }); + card.append(svg); + if (note) card.append(h("figcaption", { class: "sub muted", text: note })); + pendingDraws.push(() => drawLines(card, svg, live, events, fmtY)); + return card; +} +function drawLines(card, svg, live, events, fmtY) { + const W = Math.max(240, Math.round(svg.getBoundingClientRect().width || card.clientWidth - 28)), H = 170; + const m = { l: 52, r: 10, t: 8, b: 22 }; + clear(svg); + svg.setAttribute("viewBox", `0 0 ${W} ${H}`); + const all = live.flatMap((sr) => sr.pts); + let x0 = Math.min(...all.map((p) => p.x)), x1 = Math.max(...all.map((p) => p.x)); + let y0 = Math.min(...all.map((p) => p.y)), y1 = Math.max(...all.map((p) => p.y)); + if (x0 === x1) { x0 -= 1; x1 += 1; } + if (y0 === y1) { const d = Math.abs(y0) * 0.05 || 1; y0 -= d; y1 += d; } + const pad = (y1 - y0) * 0.06; y0 -= pad; y1 += pad; + const X = (v) => m.l + (v - x0) / (x1 - x0) * (W - m.l - m.r); + const Y = (v) => m.t + (1 - (v - y0) / (y1 - y0)) * (H - m.t - m.b); + const grid = s("g", { class: "grid" }), axis = s("g", { class: "axis" }); + for (const v of niceTicks(y0, y1, 4)) { + grid.append(s("line", { x1: m.l, x2: W - m.r, y1: Y(v), y2: Y(v) })); + const t = s("text", { x: m.l - 6, y: Y(v) + 3, "text-anchor": "end" }); t.textContent = fmtY(v); axis.append(t); + } + for (const v of niceTicks(x0, x1, Math.max(3, Math.floor(W / 80)))) { + const t = s("text", { x: X(v), y: H - 6, "text-anchor": "middle" }); t.textContent = fmtCount(v); axis.append(t); + } + svg.append(grid, s("line", { class: "baseline", x1: m.l, x2: W - m.r, y1: H - m.b, y2: H - m.b }), axis); + for (const ev of events) { + if (!isNum(ev.step) || ev.step < x0 || ev.step > x1) continue; + svg.append(s("line", { class: `event ${ev.event}`, x1: X(ev.step), x2: X(ev.step), y1: m.t, y2: H - m.b })); + } + for (const sr of live.slice().reverse()) { + if (sr.dots) { + for (const p of sr.pts) svg.append(s("circle", { class: `valid`, cx: X(p.x), cy: Y(p.y), r: 4 })); + continue; + } + svg.append(s("path", { class: `line ${sr.cls}`, d: sr.pts.map((p, i) => `${i ? "L" : "M"}${X(p.x).toFixed(1)} ${Y(p.y).toFixed(1)}`).join("") })); + } + // crosshair layer + const cross = s("line", { class: "cross", y1: m.t, y2: H - m.b, visibility: "hidden" }); + const marks = live.map((sr) => s("circle", { class: `marker ${sr.dots ? "s2" : sr.cls}`, r: 4, visibility: "hidden" })); + const hit = s("rect", { x: m.l, y: 0, width: W - m.l - m.r, height: H, fill: "transparent" }); + svg.append(cross, ...marks, hit); + const base = live.find((sr) => !sr.dots) || live[0]; + const xs = base.pts.map((p) => p.x); + let idx = xs.length - 1; + const place = (i, clientX, clientY) => { + idx = Math.max(0, Math.min(xs.length - 1, i)); + const step = xs[idx], cx = X(step); + cross.setAttribute("x1", cx); cross.setAttribute("x2", cx); cross.setAttribute("visibility", "visible"); + const rows = []; + live.forEach((sr, k) => { + const j = bisect(sr.pts.map((p) => p.x), step), p = sr.pts[j]; + const near = sr.dots ? Math.abs(p.x - step) <= (x1 - x0) * 0.02 : true; + marks[k].setAttribute("visibility", near ? "visible" : "hidden"); + if (!near) return; + marks[k].setAttribute("cx", X(p.x)); marks[k].setAttribute("cy", Y(p.y)); + rows.push({ label: sr.label, value: fmtY(p.y), cls: live.length > 1 ? (sr.dots ? "s2" : sr.cls) : null }); + }); + const t = base.pts[idx].t; + const box = svg.getBoundingClientRect(); + tip.show(clientX ?? box.left + cx * box.width / W, clientY ?? box.top + 20, + `${tr("step")} ${fmtInt(step)}${t ? " · " + fmtTime(t) : ""}`, rows); + }; + const hide = () => { cross.setAttribute("visibility", "hidden"); marks.forEach((mk) => mk.setAttribute("visibility", "hidden")); tip.hide(); }; + hit.addEventListener("pointermove", (e) => { + const box = svg.getBoundingClientRect(); + place(bisect(xs, x0 + ((e.clientX - box.left) * W / box.width - m.l) / (W - m.l - m.r) * (x1 - x0)), e.clientX, e.clientY); + }); + hit.addEventListener("pointerleave", hide); + card.onkeydown = (e) => { + if (e.key === "ArrowLeft" || e.key === "ArrowRight") { e.preventDefault(); place(idx + (e.key === "ArrowRight" ? 1 : -1)); } + else if (e.key === "Escape") hide(); + }; + card.onfocus = () => place(idx); + card.onblur = hide; +} +function flushDraws() { const fns = pendingDraws.splice(0); requestAnimationFrame(() => fns.forEach((fn) => fn())); } + +// ---- heatmap (sequential, binned to theme ramp steps) ------------------------------------- +function heatmap({ title, caption, rowLabels, values, colLabel = (j) => String(j), fmtV = (v) => fmt(v, 3), name }) { + const box = h("div", { class: "box card" }, h("h3", { text: title }), caption ? h("p", { class: "caption", text: caption }) : null); + const flat = values.flat().filter(isNum); + if (!flat.length) { box.append(h("div", { class: "empty", text: tr("no_data") })); return box; } + const lo = Math.min(...flat), hi = Math.max(...flat); + const cols = Math.max(...values.map((r) => r.length)); + const grid = h("div", { class: "heat", role: "img", "aria-label": `${title}: ${fmtV(lo)} – ${fmtV(hi)}` }); + grid.style.gridTemplateColumns = `max-content repeat(${cols}, minmax(9px, 1fr))`; + grid.append(h("span")); + for (let j = 0; j < cols; j++) grid.append(h("span", { class: "collab", text: cols > 32 && j % 4 ? "" : colLabel(j) })); + values.forEach((row, i) => { + grid.append(h("span", { class: "rowlab", text: rowLabels[i] ?? String(i) })); + for (let j = 0; j < cols; j++) { + const v = row[j]; + const cell = h("span", { class: "cell" }); + if (isNum(v)) { + const bin = hi > lo ? 1 + Math.min(6, Math.floor((v - lo) / (hi - lo) * 7)) : 4; + cell.style.background = `var(--seq-${bin})`; + } + cell.addEventListener("pointerenter", (e) => tip.show(e.clientX, e.clientY, `${rowLabels[i] ?? i} · ${colLabel(j, true)}`, + [{ label: name || title, value: fmtV(v) }])); + cell.addEventListener("pointermove", (e) => tip.show(e.clientX, e.clientY, `${rowLabels[i] ?? i} · ${colLabel(j, true)}`, + [{ label: name || title, value: fmtV(v) }])); + cell.addEventListener("pointerleave", tip.hide); + grid.append(cell); + } + }); + const ramp = h("span", { class: "ramp", "aria-hidden": "true" }); + for (let b = 1; b <= 7; b++) { const sw = h("span"); sw.style.background = `var(--seq-${b})`; ramp.append(sw); } + box.append(grid, h("div", { class: "scale" }, h("span", { class: "num", text: fmtV(lo) }), ramp, h("span", { class: "num", text: fmtV(hi) }))); + return box; +} +// ---- panels: shared bits --------------------------------------------------------------- +function phasesInHistory() { return Object.keys((S.snap && S.snap.history) || {}); } +function chosenPhase() { + const list = phasesInHistory(); + if (S.phase && list.includes(S.phase)) return S.phase; + const cur = S.snap && S.snap.progress && S.snap.progress.current; + return list.includes(cur) ? cur : list[list.length - 1] || ""; +} +function windowed(rows) { return S.window > 0 ? rows.slice(-S.window) : rows; } +function chartEvents(phase) { + if (!S.showEvents) return []; + return ((S.snap && S.snap.events) || []).filter((e) => e.phase === phase && (e.event === "validation" || e.event === "rollback")); +} +function toolbar() { + const bar = h("div", { class: "toolbar" }); + const phases = phasesInHistory(); + if (phases.length > 1) { + const sel = h("select", { "aria-label": tr("phase"), onchange: (e) => { S.phase = e.target.value; renderPanel(true); } }, + phases.map((ph) => h("option", { value: ph, text: phaseName(ph) }))); + sel.value = chosenPhase(); + bar.append(h("label", { class: "field" }, tr("phase"), sel)); + } + const win = h("select", { "aria-label": tr("window"), onchange: (e) => { S.window = Number(e.target.value); renderPanel(true); } }, + [0, 600, 200, 50].map((n) => h("option", { value: String(n), text: n ? tr("w_last", { n }) : tr("w_all") }))); + win.value = String(S.window); + const ev = h("input", { type: "checkbox", id: "events-toggle", onchange: (e) => { S.showEvents = e.target.checked; renderPanel(true); } }); + ev.checked = S.showEvents; + bar.append(h("label", { class: "field" }, tr("window"), win), h("label", { class: "field" }, ev, tr("events_on"))); + return bar; +} +const charts = (...cards) => h("div", { class: "charts" }, cards.flat(Infinity)); + +// ---- overview ---------------------------------------------------------------------------- +function panelOverview() { + const phase = chosenPhase(); + const rows = windowed(phaseRows(phase)); + const ev = chartEvents(phase); + const valid = ((S.snap && S.snap.events) || []).filter((e) => e.event === "validation" && e.phase === phase && isNum(e.valid_ce)); + const lo = rows.length ? rows[0].step : 0; + return [toolbar(), charts( + lineChart({ title: tr("c_primary"), events: ev, series: [ + { key: "primary_loss", label: tr("train"), rows }, + { key: "valid_ce", label: tr("valid"), rows: valid.filter((e) => e.step >= lo), dots: true, cls: "s2" }] }), + lineChart({ title: tr("c_loss"), events: ev, series: [{ key: "loss", label: mlabel("loss"), rows }] }), + lineChart({ title: tr("c_lr"), series: [{ key: "lr", label: tr("k_lr"), rows }], fmtY: (v) => fmt(v, 3) }), + lineChart({ title: tr("c_grad"), events: ev, series: [{ key: "grad_norm", label: tr("k_grad"), rows }] }), + lineChart({ title: tr("c_tps"), series: [ + { key: "tokens_per_s", label: tr("seen"), rows }, { key: "processed_tokens_per_s", label: tr("processed"), rows }], fmtY: fmtCount }), + lineChart({ title: tr("c_step"), series: [{ key: "step_s", label: tr("c_step"), rows }], fmtY: (v) => fmt(v, 3) }), + lineChart({ title: tr("c_mem"), series: [ + { key: "peak_mem_gb", label: tr("peak"), rows }, { key: "active_mem_gb", label: tr("active"), rows }, + { key: "cache_mem_gb", label: tr("cache"), rows }], fmtY: (v) => fmt(v, 3) }), + lineChart({ title: tr("c_tokens"), series: [{ key: "tokens_seen", label: tr("c_tokens"), rows }], fmtY: fmtCount }), + )]; +} + +// ---- loss components --------------------------------------------------------------------------- +function panelLoss() { + const phase = chosenPhase(); + const rows = windowed(phaseRows(phase)); + const ev = chartEvents(phase); + const keys = ((S.snap && S.snap.groups && S.snap.groups.loss) || []).filter((k) => rows.some((r) => isNum(r[k]))); + for (const k of ["recurrent_contraction_loss"]) if (rows.some((r) => isNum(r[k])) && !keys.includes(k)) keys.push(k); + return [toolbar(), charts(keys.map((k) => lineChart({ title: mlabel(k), events: ev, series: [{ key: k, label: mlabel(k), rows }] })))]; +} + +// ---- recurrence ------------------------------------------------------------------------------------ +function panelRecurrent() { + const phase = chosenPhase(); + const rows = windowed(phaseRows(phase)); + const L = (S.snap && S.snap.latest) || {}; + const cfg = ((S.snap && S.snap.spec) || {}).model_config || {}; + const out = [toolbar(), h("p", { class: "caption muted", text: tr("rec_caption") })]; + out.push(charts(["rec_depth_mean", "rec_contraction", "rec_early_exit_ratio", "rec_ponder_loss", "rec_distill_loss", "rec_stability_loss"] + .filter((k) => rows.some((r) => isNum(r[k]))) + .map((k) => lineChart({ title: mlabel(k), series: [{ key: k, label: mlabel(k), rows }], fmtY: (v) => fmt(v, 3) })))); + const facts = h("dl", { class: "kv" }); + const add = (k, v) => facts.append(h("dt", { text: k }), h("dd", { class: "num", text: v })); + add("recurrent_depth", String(cfg.recurrent_depth ?? "—")); + add("min_recurrent_depth", String(cfg.min_recurrent_depth ?? "—")); + add("recurrent_attention_mode", String(cfg.recurrent_attention_mode ?? "—")); + add("adaptive_halting", String(cfg.adaptive_halting ?? "—")); + add("depth_schedule", String(cfg.depth_schedule ?? "—")); + if (Array.isArray(L.rec_depth_per_group)) add(S.lang === "zh" ? "每组深度（最新）" : "depth per group (latest)", L.rec_depth_per_group.map((v) => fmt(v, 3)).join(" · ")); + if (Array.isArray(L.rec_halt_mass_by_step)) add(S.lang === "zh" ? "累计停止质量 / 迭代" : "cumulative halt mass / iteration", L.rec_halt_mass_by_step.map((v) => fmt(v, 3)).join(" · ")); + out.push(h("div", { class: "grid2" }, h("div", { class: "box card" }, h("h3", { text: tr("cfg_model") }), facts))); + return out; +} + +// ---- routing ------------------------------------------------------------------------------------------- +function panelRouting() { + const phase = chosenPhase(); + const rows = windowed(phaseRows(phase)); + const out = [toolbar(), h("p", { class: "caption muted", text: tr("route_caption") })]; + out.push(charts(["route_entropy", "load_cv", "max_expert_load", "dead_experts"].filter((k) => rows.some((r) => isNum(r[k]))) + .map((k) => lineChart({ title: mlabel(k), series: [{ key: k, label: mlabel(k), rows }], fmtY: (v) => fmt(v, 3) })))); + const r = S.snap && S.snap.parameters && S.snap.parameters.routing; + const m = S.snap && S.snap.parameters; + if (!r) { + const every = ((S.snap && S.snap.spec) || {}).tensor_stats_every_s || 1800; + out.push(h("div", { class: "empty", text: tr("no_telemetry", { m: Math.round(every / 60) }) })); + return out; + } + const expertLabel = (j, long) => long ? `${S.lang === "zh" ? "专家" : "expert"} ${j}` : String(j); + const meta = tr("tel_meta", { s: fmtInt(m.step), t: fmtTime(m.time), ms: fmtInt((m.telemetry_s || 0) * 1000), note: r.source || "" }); + out.push(h("p", { class: "caption muted", text: meta })); + const ent = r.entropy || []; + out.push(h("div", { class: "grid2" }, + heatmap({ title: tr("heat_load"), caption: `${tr("heat_rows")} · top-${r.top_k} · ${r.experts} experts`, rowLabels: r.labels, + values: r.load, colLabel: expertLabel, fmtV: (v) => isNum(v) ? (v * 100).toFixed(1) + "%" : "—", name: tr("heat_load") }), + heatmap({ title: tr("heat_importance"), caption: tr("heat_rows"), rowLabels: r.labels.map((l, i) => isNum(ent[i]) ? `${l} H=${fmt(ent[i], 3)}` : l), + values: r.importance, colLabel: expertLabel, fmtV: (v) => fmt(v, 3), name: tr("heat_importance") }))); + return out; +} +// Extra strings used by the detail panels. +Object.assign(STR.zh, { + probe_note: "评测指标由训练器在每次验证时写入 metrics.jsonl，页面自动发现并按指标分组；探针样本量通常很小，只作趋势参考。", + lower_better: "越低越好", higher_better: "越高越好", + final_eval: "最终评测", all_events: "全部事件", no_events: "暂无事件。", no_samples: "暂无评测样例。", + tensors_loading: "正在加载张量统计…", shown: "显示 {n} / {t}", h_rms: "权重 RMS", h_grad: "梯度 RMS", h_upd: "更新比", + stats_total: "全模型", every: "每 {m} 一次", none: "无", +}); +Object.assign(STR.en, { + probe_note: "Evaluation metrics are whatever the trainer writes to metrics.jsonl at each validation; the page discovers and groups them by metric. Probes usually use a small sample, so read them as a trend.", + lower_better: "lower is better", higher_better: "higher is better", + final_eval: "Final evaluation", all_events: "All events", no_events: "No events yet.", no_samples: "No evaluation samples yet.", + tensors_loading: "Loading tensor statistics…", shown: "showing {n} of {t}", h_rms: "Weight RMS", h_grad: "Grad RMS", h_upd: "Update ratio", + stats_total: "Whole model", every: "every {m}", none: "none", +}); + +function table(columns, rows, opts = {}) { + if (opts.hideEmpty) { // drop columns with no value in any row (e.g. SFT metrics during pretrain) + const value = (c, r) => (c.get ? c.get(r) : r[c.key]); + columns = columns.filter((c) => c.keep || rows.some((r) => { const v = value(c, r); return v != null && v !== ""; })); + } + const head = h("tr", null, columns.map((c) => h("th", { class: c.num ? "num" : null, scope: "col", text: c.label }))); + const body = h("tbody"); + for (const r of rows) { + const tr_ = h("tr", { class: opts.onRow ? "clickable" : null, tabindex: opts.onRow ? "0" : null }); + if (opts.selected && opts.selected(r)) tr_.classList.add("selected"); + for (const c of columns) { + const v = c.get ? c.get(r) : r[c.key]; + tr_.append(h("td", { class: c.num ? "num" : null, text: c.fmt ? c.fmt(v, r) : (v ?? "—") })); + } + if (opts.onRow) { + tr_.addEventListener("click", () => opts.onRow(r)); + tr_.addEventListener("keydown", (e) => { if (e.key === "Enter" || e.key === " ") { e.preventDefault(); opts.onRow(r); } }); + } + body.append(tr_); + } + return h("div", { class: "table-wrap" }, h("table", null, h("thead", null, head), body)); +} + +// ---- parameters (per-tensor telemetry) ---------------------------------------------------------- +async function loadTensors() { + if (S.tensorsLoading) return; + S.tensorsLoading = true; + try { S.tensors = await api("api/tensors", { run: S.run }); } + catch (e) { if (e instanceof AuthError) showLogin(""); else S.error = e.message; } + finally { S.tensorsLoading = false; } +} +async function pickTensor(name) { + S.tensorPick = name; + try { S.tensorHist = (await api("api/tensor", { run: S.run, name })).history || []; } + catch (e) { if (e instanceof AuthError) { showLogin(""); return; } S.tensorHist = []; S.error = e.message; } + renderPanel(true); +} +function tensorTable(container) { + clear(container); + const all = (S.tensors && S.tensors.tensors) || []; + const q = S.tensorFilter.trim().toLowerCase(); + let rows = q ? all.filter((t) => t.name.toLowerCase().includes(q) || (t.group || "").includes(q)) : all.slice(); + const key = S.tensorSort; + if (key !== "name") rows.sort((a, b) => (isNum(b[key]) ? b[key] : -Infinity) - (isNum(a[key]) ? a[key] : -Infinity)); + container.append(h("p", { class: "caption muted", text: tr("shown", { n: Math.min(rows.length, 400), t: all.length }) })); + const numel = (t) => (t.shape || []).reduce((a, b) => a * b, 1); + container.append(table([ + { label: tr("name"), key: "name" }, { label: tr("group"), key: "group" }, + { label: tr("shape"), get: (t) => (t.shape || []).join("×") }, { label: tr("dtype"), key: "dtype" }, + { label: tr("numel"), get: numel, num: true, fmt: fmtCount }, + { label: tr("rms"), key: "rms", num: true, fmt: (v) => fmt(v, 3) }, { label: tr("std"), key: "std", num: true, fmt: (v) => fmt(v, 3) }, + { label: tr("absmax"), key: "absmax", num: true, fmt: (v) => fmt(v, 3) }, + { label: tr("grad_rms"), key: "grad_rms", num: true, fmt: (v) => fmt(v, 3) }, + { label: tr("update_ratio"), key: "update_ratio", num: true, fmt: (v) => fmt(v, 3) }, + ], rows.slice(0, 400), { onRow: (t) => pickTensor(t.name), selected: (t) => t.name === S.tensorPick })); +} +function panelParams() { + const m = S.snap && S.snap.parameters; + const every = ((S.snap && S.snap.spec) || {}).tensor_stats_every_s || 1800; + if (!m) return [h("div", { class: "empty", text: tr("no_telemetry", { m: Math.round(every / 60) }) })]; + const out = [h("p", { class: "caption muted", text: `${phaseName(m.phase)} · ${tr("tel_meta", { s: fmtInt(m.step), t: fmtTime(m.time), + ms: fmtInt((m.telemetry_s || 0) * 1000), note: m.grad_note || "" })} · ${tr("every", { m: fmtDur(every) })}` })]; + const groups = Object.entries(m.groups || {}).map(([name, g]) => ({ name, ...g })); + if (m.stats) groups.push({ name: tr("stats_total"), ...m.stats, tensors: m.tensor_count }); + out.push(h("div", { class: "box card" }, h("h3", { text: tr("groups") }), table([ + { label: tr("group"), key: "name" }, { label: tr("tensors_n"), key: "tensors", num: true, fmt: fmtInt }, + { label: tr("numel"), key: "numel", num: true, fmt: fmtCount }, { label: tr("rms"), key: "rms", num: true, fmt: (v) => fmt(v, 3) }, + { label: tr("grad_norm"), key: "grad_norm", num: true, fmt: (v) => fmt(v, 3) }, + { label: tr("update_ratio"), key: "update_ratio", num: true, fmt: (v) => fmt(v, 3) }], groups))); + const ex = m.experts || {}; + if (ex.layers && ex.layers.length) { + const rowsL = ex.layers.map((l) => `L${l}`); + const expertLabel = (j, long) => long ? `${S.lang === "zh" ? "专家" : "expert"} ${j}` : String(j); + out.push(h("div", { class: "grid2" }, + heatmap({ title: tr("heat_wrms"), caption: tr("heat_rows_layer"), rowLabels: rowsL, values: ex.weight_rms, colLabel: expertLabel }), + heatmap({ title: tr("heat_grms"), caption: tr("heat_rows_layer"), rowLabels: rowsL, values: ex.grad_rms, colLabel: expertLabel }))); + } + if (S.tensorPick && S.tensorHist) { + const rows = S.tensorHist; + out.push(h("div", { class: "box card" }, h("h3", { text: tr("hist_of", { n: S.tensorPick }) }), charts( + lineChart({ title: tr("h_rms"), series: [{ key: "rms", label: tr("h_rms"), rows }], fmtY: (v) => fmt(v, 3) }), + lineChart({ title: tr("h_grad"), series: [{ key: "grad_rms", label: tr("h_grad"), rows }], fmtY: (v) => fmt(v, 3) }), + lineChart({ title: tr("h_upd"), series: [{ key: "update_ratio", label: tr("h_upd"), rows }], fmtY: (v) => fmt(v, 3) })))); + } + const box = h("div", { class: "box card" }, h("h3", { text: tr("tensors") })); + const holder = h("div"); + if (!S.tensors || S.tensors.step !== m.step) { + holder.append(h("div", { class: "empty", text: tr("tensors_loading") })); + if (S.tensorsFor !== m.step) { // one attempt per telemetry step (no retry loop on errors) + S.tensorsFor = m.step; + loadTensors().then(() => renderPanel(true)); + } + } else { + const filter = h("input", { type: "search", placeholder: tr("filter"), "aria-label": tr("filter"), value: S.tensorFilter, + oninput: (e) => { S.tensorFilter = e.target.value; tensorTable(holder); } }); + const sort = h("select", { "aria-label": tr("sort"), onchange: (e) => { S.tensorSort = e.target.value; tensorTable(holder); } }, + ["name", "rms", "absmax", "grad_rms", "update_ratio"].map((k) => h("option", { value: k, text: tr(k) }))); + sort.value = S.tensorSort; + box.append(h("div", { class: "toolbar" }, filter, h("label", { class: "field" }, tr("sort"), sort)), + h("p", { class: "caption muted", text: tr("hist_hint") })); + tensorTable(holder); + } + box.append(holder); + out.push(box); + return out; +} +// ---- evaluations ------------------------------------------------------------------------------------- +function flatten(obj, prefix = "", out = [], depth = 0) { + if (obj == null) return out; + for (const [k, v] of Object.entries(obj)) { + const key = prefix ? `${prefix}.${k}` : k; + if (v && typeof v === "object" && !Array.isArray(v) && depth < 3) flatten(v, key, out, depth + 1); + else if (Array.isArray(v)) { + const prim = v.every((x) => x == null || typeof x !== "object"); + out.push([key, prim && v.length <= 24 ? v.map((x) => isNum(x) ? fmt(x, 5) : String(x)).join(", ") : `[${v.length}]`]); + } else if (v && typeof v === "object") out.push([key, JSON.stringify(v).slice(0, 240)]); + else out.push([key, isNum(v) ? (Number.isInteger(v) ? fmtInt(v) : fmt(v, 6)) : String(v)]); + } + return out; +} +function kvBox(title, obj, extra) { + const dl = h("dl", { class: "kv" }); + for (const [k, v] of flatten(obj)) dl.append(h("dt", { text: k }), h("dd", { text: v })); + return h("div", { class: "box card" }, h("h3", { text: title }), extra || null, obj ? dl : h("div", { class: "empty", text: tr("none") })); +} +// Samples come from samples.jsonl through chronos/monitor/evals.py:sample_view — +// input / reference / output roles plus free-form tags, whatever the task. +function sampleList(samples) { + if (!samples.length) return h("div", { class: "empty", text: tr("no_samples") }); + const line = (lab, v) => (v == null || v === "" ? null : h("div", null, h("span", { class: "lab", text: lab }), String(v))); + return h("div", null, samples.map((x) => { + const tags = Object.entries(x.tags || {}).map(([k, v]) => `${k}=${isNum(v) ? fmt(v, 4) : v}`); + return h("div", { class: "sample" }, + h("div", { class: "meta", text: [phaseName(x.phase), x.step != null ? `${tr("step")} ${fmtInt(x.step)}` : null, + ...tags, fmtTime(x.time)].filter(Boolean).join(" · ") }), + line(tr("src"), x.input), line(tr("ref"), x.reference), line(tr("hyp"), x.output)); + })); +} +const METRIC_NAMES = { bleu: "BLEU", chrf: "chrF", ter: "TER", rouge1: "ROUGE-1", rouge2: "ROUGE-2", rougel: "ROUGE-L", + meteor: "METEOR", comet: "COMET", bertscore: "BERTScore", acc: "Acc", accuracy: "Accuracy", f1: "F1", em: "EM", + auc: "AUC", wer: "WER", cer: "CER", ppl: "PPL", ce: "CE", nll: "NLL", bpb: "BPB", bpc: "BPC", kl: "KL" }; +const metricName = (m) => METRIC_NAMES[m] || m; +// One chart per metric family and phase (steps restart at each phase), at most three series each. +function evalCharts(valid, families) { + const cards = []; + const phases = [...new Set(valid.map((e) => e.phase))]; + for (const fam of families) { + for (const ph of phases) { + const rows = valid.filter((e) => e.phase === ph); + const live = fam.series.filter((sr) => rows.some((e) => isNum(e[sr.key]))); + for (let i = 0; i < live.length; i += 3) { + const part = live.slice(i, i + 3); + const title = `${metricName(fam.metric)} · ${phaseName(ph)}${live.length > 3 ? ` (${i / 3 + 1})` : ""}`; + cards.push(lineChart({ title, series: part.map((sr) => ({ key: sr.key, label: sr.label, rows })), + fmtY: (v) => fmt(v, 4), note: fam.direction ? tr(fam.direction + "_better") : undefined })); + } + } + } + return cards; +} +function panelEvals() { + const snap = S.snap || {}; + const events = snap.events || []; + const evals = snap.evals || { families: [], keys: [], context: [] }; + const out = []; + const final = events.filter((e) => e.event === "final").pop(); + if (final) out.push(kvBox(tr("final_eval"), Object.fromEntries(Object.entries(final).filter(([k]) => !["event", "samples"].includes(k))))); + const valid = events.filter((e) => e.event === "validation"); + const cards = evalCharts(valid, evals.families || []); + if (cards.length) out.push(h("p", { class: "caption muted", text: tr("probe_note") }), charts(cards)); + const f4 = (v) => fmt(v, 4); + const extra = (evals.keys || []).filter((k) => k !== "valid_ce" && k !== "valid_ppl"); + out.push(h("div", { class: "box card" }, h("h3", { text: tr("validations") }), valid.length ? table([ + { label: tr("time"), key: "time", fmt: fmtTime }, { label: tr("phase"), key: "phase", fmt: phaseName }, + { label: tr("step"), key: "step", num: true, fmt: fmtInt }, { label: "valid CE", key: "valid_ce", num: true, fmt: f4 }, + { label: tr("ppl"), key: "valid_ppl", num: true, fmt: f4 }, + ...extra.map((k) => ({ label: k, key: k, num: true, fmt: f4 })), + ...(evals.context || []).map((k) => ({ label: k, key: k, num: true, fmt: (v) => fmt(v, 4) })), + ], valid.slice().reverse(), { hideEmpty: true }) : h("div", { class: "empty", text: tr("no_events") }))); + const ins = snap.inspections || []; + if (ins.length) { + const keys = [...new Set(ins.flatMap((r) => Object.keys(r.metrics || {})))]; + const shown = ins.slice().reverse(); + out.push(h("div", { class: "box card" }, h("h3", { text: tr("inspections") }), + h("p", { class: "caption muted", text: shown.map((r) => `${r.file}: ${phaseName(r.phase)} ${tr("step")} ${fmtInt(r.step)} · ${r.device || "?"} · ${r.compute_dtype || "?"}`).join(" ") }), + table([{ label: tr("metric"), get: (k) => k }, + ...shown.map((r) => ({ label: `${tr("step")} ${fmtInt(r.step)}`, get: (k) => (r.metrics || {})[k], num: true, fmt: f4 }))], + keys))); + } + out.push(h("div", { class: "box card" }, h("h3", { text: tr("samples") }), sampleList((snap.samples || []).slice().reverse().slice(0, 12)))); + const other = events.filter((e) => e.event !== "validation"); + out.push(h("div", { class: "box card" }, h("h3", { text: tr("all_events") }), other.length ? table([ + { label: tr("time"), key: "time", fmt: fmtTime }, { label: tr("kind"), key: "event", fmt: (v) => tr("ev_" + v) }, + { label: tr("phase"), key: "phase", fmt: phaseName }, { label: tr("step"), key: "step", num: true, fmt: fmtInt }, + { label: tr("detail"), get: (e) => Object.entries(e).filter(([k]) => !["time", "event", "phase", "step"].includes(k)) + .map(([k, v]) => `${k}=${isNum(v) ? fmt(v, 5) : typeof v === "object" ? JSON.stringify(v) : v}`).join(" ") }, + ], other.slice().reverse()) : h("div", { class: "empty", text: tr("no_events") }))); + return out; +} +// ---- training parameters --------------------------------------------------------------------------- +async function loadParams() { + try { S.params = await api("api/parameters", { run: S.run }); } + catch (e) { if (e instanceof AuthError) { showLogin(""); return; } S.params = {}; S.error = e.message; } +} +// Every scalar top-level spec field (whatever the recipe defines), plus ``extra``. +const RUN_SKIP = new Set(["phases", "model_config", "extra"]); +function panelConfig() { + const run = S.snap && S.snap.run; + if (!S.params || S.params.run !== run) { + if (S.paramsFor !== run) { S.paramsFor = run; loadParams().then(() => renderPanel(true)); } + return [h("div", { class: "empty", text: "…" })]; + } + const spec = S.params.spec || {}; + const rt = S.params.runtime; + const phases = (spec.phases || []).map((ph) => { + const r = ((rt && rt.phases) || []).find((x) => x.name === ph.name) || {}; + return { ...ph, accum: r.accum ?? (ph.micro_tokens ? Math.max(1, Math.floor(ph.tokens_per_step / ph.micro_tokens)) : null) }; + }); + const out = []; + out.push(h("div", { class: "box card" }, h("h3", { text: tr("cfg_phases") }), table([ + { label: tr("phase"), key: "name", fmt: phaseName }, { label: "budget", key: "budget_fraction", num: true, fmt: (v) => fmt(v, 3) }, + { label: "peak_lr", key: "peak_lr", num: true, fmt: (v) => fmt(v, 3) }, { label: "tokens/step", key: "tokens_per_step", num: true, fmt: fmtInt }, + { label: "micro_tokens", key: "micro_tokens", num: true, fmt: fmtInt }, { label: "accum", key: "accum", num: true, fmt: fmtInt }, + { label: "seq_len", key: "seq_len", num: true, fmt: fmtInt }, { label: "warmup", key: "warmup_fraction", num: true, fmt: (v) => fmt(v, 3) }, + { label: "min_lr_ratio", key: "min_lr_ratio", num: true, fmt: (v) => fmt(v, 3) }, { label: "weight_decay", key: "weight_decay", num: true, fmt: (v) => fmt(v, 3) }, + { label: "grad_clip", key: "grad_clip", num: true, fmt: (v) => fmt(v, 3) }], phases))); + const runObj = Object.fromEntries(Object.entries(spec).filter(([k, v]) => !RUN_SKIP.has(k) && (v == null || typeof v !== "object"))); + if (spec.extra) runObj.extra = spec.extra; + const grid = h("div", { class: "grid2" }); + grid.append(kvBox(tr("cfg_run"), runObj), kvBox(tr("cfg_model"), spec.model_config || null)); + if (rt) { + const { phases: _p, ...rest } = rt; + grid.append(kvBox(tr("cfg_runtime"), rest)); + } else { + const dl = h("dl", { class: "kv" }, h("dt", { text: "dtype" }), h("dd", { text: String(spec.dtype ?? "—") })); + grid.append(h("div", { class: "box card" }, h("h3", { text: tr("cfg_runtime") }), h("p", { class: "caption", text: tr("cfg_missing") }), dl)); + } + if (S.params.data) grid.append(kvBox(tr("cfg_data"), S.params.data)); + if (S.params.tokenizer) grid.append(kvBox(tr("cfg_tok"), S.params.tokenizer)); + out.push(grid); + out.push(h("details", null, h("summary", { text: tr("cfg_raw") }), h("pre", { text: JSON.stringify(S.params, null, 1) }))); + return out; +} + +// ---- log ---------------------------------------------------------------------------------------------- +async function loadLog() { + try { S.log = await api("api/log", { run: S.run, bytes: S.logBytes }, true); } + catch (e) { if (e instanceof AuthError) showLogin(""); else S.log = tr("err_net") + e.message; } +} +function panelLog() { + const pre = h("pre", { class: "log", tabindex: "0", text: S.log ?? "…" }); + const auto = h("input", { type: "checkbox", onchange: (e) => { S.logAuto = e.target.checked; } }); + auto.checked = S.logAuto; + pendingDraws.push(() => { if (S.logAuto) pre.scrollTop = pre.scrollHeight; }); + return [h("div", { class: "toolbar" }, + h("button", { type: "button", text: tr("log_refresh"), onclick: () => loadLog().then(() => renderPanel(true)) }), + h("button", { type: "button", text: tr("log_more"), disabled: S.logBytes >= 400000 ? true : null, + onclick: () => { S.logBytes = Math.min(400000, S.logBytes * 4); loadLog().then(() => renderPanel(true)); } }), + h("label", { class: "field" }, auto, tr("log_auto"))), pre]; +} + +// ---- checkpoints ----------------------------------------------------------------------------------------- +function panelCkpt() { + const list = (S.snap && S.snap.checkpoints) || []; + if (!list.length) return [h("div", { class: "empty", text: tr("ck_none") })]; + return [h("div", { class: "box card" }, table([ + { label: tr("ck_file"), key: "file" }, { label: tr("ck_size"), key: "size_gb", num: true, fmt: fmtGB }, + { label: tr("ck_modified"), key: "modified", fmt: fmtTime }, { label: tr("ck_age"), key: "age_s", num: true, fmt: fmtDur }], list))]; +} + +// ---- render orchestration ----------------------------------------------------------------------------------- +const PANELS = { overview: panelOverview, loss: panelLoss, recurrent: panelRecurrent, routing: panelRouting, params: panelParams, + evals: panelEvals, config: panelConfig, log: panelLog, ckpt: panelCkpt }; +function panelKey() { + const snap = S.snap || {}; + if (S.tab === "params") return [snap.run, (snap.parameters || {}).step, S.tensors && S.tensors.step, S.tensorPick, (S.tensorHist || []).length, S.lang].join("|"); + if (S.tab === "config") return [snap.run, S.params ? S.params.run : "", S.lang].join("|"); + return null; // charts, events and the log redraw on every poll +} +function renderPanel(force) { + const panel = $("panel"); + if (!S.snap || !S.snap.run) { panel.replaceChildren(); return; } + const key = panelKey(); + if (!force && key !== null && key === S.panelKey) return; // keep filter focus / table scroll between polls + S.panelKey = key; + tip.hide(); + pendingDraws.length = 0; + let content; + try { content = PANELS[S.tab](); } + catch (e) { content = [h("div", { class: "notice error", role: "alert" }, h("span", { class: "icon", text: "✕" }), h("span", { text: String(e && e.stack || e) }))]; } + panel.replaceChildren(...[content].flat().filter(Boolean)); + flushDraws(); +} +function renderFoot() { + const snap = S.snap; + const foot = clear($("foot")); + if (!snap) return; + const st = snap.status || {}; + const bits = [tr("foot", { s: POLL_MS / 1000, root: snap.root || "—", t: fmtTime(snap.generated) })]; + if (snap.run) bits.push(`status: ${st.source || "?"}${st.pid ? ` · pid ${st.pid}` : ""}${st.ps_state ? ` (${st.ps_state})` : ""}`); + if (snap.runtime && snap.runtime.code) bits.push(`code ${snap.runtime.code.sha256}`); + foot.textContent = bits.join(" · "); +} +function renderAll(force) { renderChrome(); renderNotices(); renderProgress(); renderTiles(); renderPanel(force); renderFoot(); } + +// ---- wiring -------------------------------------------------------------------------------------------------- +function resetRunCaches() { S.params = S.tensors = S.tensorHist = S.log = null; S.tensorPick = null; S.paramsFor = S.tensorsFor = undefined; S.panelKey = null; S.phase = ""; } +function init() { + $("run").addEventListener("change", (e) => { S.run = e.target.value; resetRunCaches(); refresh(); }); + $("lang").addEventListener("click", () => { + S.lang = S.lang === "zh" ? "en" : "zh"; + localStorage.setItem("chronos.monitor.lang", S.lang); + renderTabs(); renderAll(true); + if (!$("login").hidden) showLogin($("login-error").textContent); + }); + $("theme").addEventListener("click", () => { + S.theme = effectiveDark() ? "light" : "dark"; + localStorage.setItem("chronos.monitor.theme", S.theme); + renderChrome(); + }); + $("logout").addEventListener("click", logout); + $("login-form").addEventListener("submit", submitLogin); + document.addEventListener("visibilitychange", () => { if (document.hidden) stop(); else if ($("login").hidden) refresh().then(schedule); }); + matchMedia("(prefers-color-scheme: dark)").addEventListener("change", renderChrome); + let resizeT; + window.addEventListener("resize", () => { clearTimeout(resizeT); resizeT = setTimeout(() => renderPanel(true), 200); }); + const fromHash = /(?:^|[#&])tab=([a-z]+)/.exec(location.hash); + if (fromHash && TABS.includes(fromHash[1])) S.tab = fromHash[1]; + window.addEventListener("hashchange", () => { + const m = /(?:^|[#&])tab=([a-z]+)/.exec(location.hash); + if (m && TABS.includes(m[1]) && m[1] !== S.tab) selectTab(m[1]); + }); + renderChrome(); + renderTabs(); + start(); +} +if (document.readyState === "loading") document.addEventListener("DOMContentLoaded", init); else init(); +})(); diff --git a/chronos/monitor/static/index.html b/chronos/monitor/static/index.html new file mode 100644 index 0000000..ec94642 --- /dev/null +++ b/chronos/monitor/static/index.html @@ -0,0 +1,49 @@ + + + + + +Chronos · Training monitor + + + + +
+
+
+ Project Chronos +

Training monitor

+
+
+ + + + + + +
+
+ +
+
+
+
+ +
+
+
+
+ + + + + diff --git a/chronos/precision.py b/chronos/precision.py new file mode 100644 index 0000000..e64112f --- /dev/null +++ b/chronos/precision.py @@ -0,0 +1,195 @@ +"""One precision policy for every Chronos training and inference chain. + +A run picks a single *compute dtype* when it is built. Weights, activations, +KV cache and SSD/RAM expert pages all live in that dtype, so the hot path +never converts tensors. Precision is protected only where rounding would +cost accuracy, and each of those places is either a fused kernel that +accumulates in FP32 internally (no tensor cast) or a deliberately small cast: + + site how FP32 is obtained cast size + --------------------------- --------------------------------- --------------- + attention softmax fused SDPA (MLX / torch) none + RMSNorm statistics mx.fast.rms_norm / fused torch none + RoPE fused rope / tables in compute dt none + router softmax softmax(precise=True) (MLX) none + LM head + cross-entropy upcast hidden [N,H] and W [V,H] H/V of logits + ACT halting / aux losses [B,S,G,D] scalars tiny + optimizer FP32 masters + moments one per update + gradient accumulation FP32 running sum one per micro + +Everything else stays in the compute dtype end to end. The policy is also +the single place that knows device limits: MLX's CPU backend implements +GatherMM for float32 only (mlx 0.32.4), so CPU MLX runs compute in float32 +for the whole model instead of upcasting every expert matmul. +""" +from __future__ import annotations + +from dataclasses import asdict, dataclass, field + +COMPUTE_DTYPES = ("float32", "bfloat16", "float16") +_ALIASES = {"fp32": "float32", "f32": "float32", "float": "float32", + "bf16": "bfloat16", "fp16": "float16", "half": "float16", "f16": "float16"} + +# Fused kernels (attention softmax, RMSNorm, RoPE, precise router softmax) +# that keep FP32 internally without materializing FP32 tensors. +FUSED_FP32_SITES = ("attention_softmax", "rms_norm", "rope", "router_softmax") +# The only explicit casts left in the chains, with their size bound. +BOUNDARY_CASTS = { + "lm_head_loss": "hidden [N,H] + head [V,H] -> fp32; logits are born fp32", + "halting_and_aux_losses": "[B,S,G,D] / [N,E] statistics -> fp32", + "optimizer_update": "grad -> fp32 master once per update", + "grad_accumulation": "fp32 running sum across micro-batches", +} + + +def normalize_dtype_name(name, default: str = "auto") -> str: + """'bf16' / 'torch.bfloat16' / mx.bfloat16 -> 'bfloat16'; '' -> default.""" + if name is None or name == "": + return default + text = str(name).strip().lower().replace("torch.", "").replace("mlx.core.", "") + text = _ALIASES.get(text, text) + if text not in COMPUTE_DTYPES + ("auto",): + raise ValueError(f"unsupported compute dtype {name!r}; expected one of {COMPUTE_DTYPES} or 'auto'") + return text + + +@dataclass(frozen=True) +class PrecisionPolicy: + backend: str # "mlx" | "torch" + device: str # "gpu" | "cpu" | "mps" | "cuda" + compute: str # weights, activations, KV cache, expert pages + master: str = "float32" # optimizer master weights + moments + accumulate: str = "float32" # gradient accumulation + requested: str = "auto" + reason: str = "" + fused_fp32: tuple = field(default=FUSED_FP32_SITES) + + @property + def kv_cache(self) -> str: + return self.compute + + @property + def expert_storage(self) -> str: + return self.compute + + @property + def low_precision(self) -> bool: + return self.compute != "float32" + + def as_dict(self) -> dict: + out = asdict(self) + out.update(kv_cache=self.kv_cache, expert_storage=self.expert_storage, + boundary_casts=dict(BOUNDARY_CASTS), fused_fp32=list(self.fused_fp32)) + return out + + # ---- backend dtype objects ------------------------------------------- + def mx_dtype(self, which: str = "compute"): + import mlx.core as mx + return getattr(mx, getattr(self, which)) + + def torch_dtype(self, which: str = "compute"): + import torch + return getattr(torch, getattr(self, which)) + + +def resolve_mlx(requested="auto", device: str | None = None) -> PrecisionPolicy: + """Pick the MLX compute dtype once for a whole run. + + GPU: BF16 by default (FP32 masters keep small updates). CPU: FP32, + because MLX CPU GatherMM (sparse expert dispatch) is FP32-only; a + model-wide FP32 run is exact and avoids per-matmul upcasts. + """ + import mlx.core as mx + req = normalize_dtype_name(requested) + dev = device or ("cpu" if mx.default_device().type == mx.DeviceType.cpu else "gpu") + if dev == "cpu": + reason = ("MLX CPU GatherMM supports float32 only; whole model runs in float32" + if req not in ("auto", "float32") else "MLX CPU backend has no fast low-precision GEMM") + return PrecisionPolicy("mlx", "cpu", "float32", requested=req, reason=reason) + if req == "auto": + return PrecisionPolicy("mlx", dev, "bfloat16", requested=req, reason="GPU default: BF16 compute, FP32 masters") + return PrecisionPolicy("mlx", dev, req, requested=req, reason="requested") + + +def resolve_torch(requested="auto", device: str = "cpu") -> PrecisionPolicy: + """Pick the torch compute dtype once for a whole run or served model.""" + import torch + req = normalize_dtype_name(requested) + dev = str(device).split(":")[0] + if req != "auto": + if dev == "cuda" and req == "bfloat16" and not torch.cuda.is_bf16_supported(): + return PrecisionPolicy("torch", dev, "float16", requested=req, reason="GPU lacks BF16; using FP16") + return PrecisionPolicy("torch", dev, req, requested=req, reason="requested") + if dev == "cuda": + dt = "bfloat16" if torch.cuda.is_bf16_supported() else "float16" + return PrecisionPolicy("torch", dev, dt, requested=req, reason="CUDA default") + if dev == "mps": + # Measured on M5 Max / torch 2.10 (Chronos 50M, native dtype, no + # autocast): BF16 is slower than FP32 for training fwd+bwd (0.749 vs + # 0.619 s/step) and cached decode (49.1 vs 37.3 ms/token). BF16 + # stays available on request (halves weight/KV memory). + return PrecisionPolicy("torch", dev, "float32", requested=req, + reason="MPS default: FP32 measured faster than BF16") + return PrecisionPolicy("torch", dev, "float32", requested=req, reason="CPU default: float32 is the fast CPU GEMM") + + +def place_torch_model(model, device: str, requested="auto") -> tuple[object, PrecisionPolicy]: + """Move a torch Chronos model to ``device`` in the policy compute dtype. + + One conversion at load time (weights + any placeholders); afterwards the + forward pass runs in that dtype with no per-op casts. Call this before + an ``ExpertStore`` is built so SSD/RAM expert pages adopt the same dtype. + """ + from chronos.model.moe_chronos import LazyExpertPlaceholder + policy = resolve_torch(requested, device) + dtype = policy.torch_dtype() + model = model.to(device=device, dtype=dtype) + for module in model.modules(): + if isinstance(module, LazyExpertPlaceholder): + module.dtype = dtype + model._chronos_precision = policy + return model, policy + + +def torch_model_dtype(model): + """Compute dtype of a torch Chronos model (its embedding weight).""" + inner = getattr(model, "model", model) + emb = getattr(inner, "embed_tokens", None) + if emb is not None: + return emb.weight.dtype + return next(model.parameters()).dtype + + +# ---- exact MLX <-> torch tensor bridges ------------------------------------ +# numpy has no bfloat16, so a BF16 tensor crosses as its uint16 bit pattern +# (exact, no FP32 detour and no second rounding). FP16/FP32 cross natively. + +def mx_to_torch(arr, dtype: str | None = None): + """mx.array -> CPU torch tensor, keeping (or converting once to) ``dtype``.""" + import mlx.core as mx + import numpy as np + import torch + if dtype is not None: + target = getattr(mx, normalize_dtype_name(dtype, "float32")) + if arr.dtype != target: + arr = arr.astype(target) + if arr.dtype == mx.bfloat16: + bits = arr.view(mx.uint16) + mx.eval(bits) + return torch.from_numpy(np.array(bits, copy=True)).view(torch.bfloat16) + mx.eval(arr) + return torch.from_numpy(np.array(arr, copy=True)) + + +def torch_to_mx(tensor, dtype: str | None = None): + """torch tensor -> mx.array, keeping (or converting once to) ``dtype``.""" + import mlx.core as mx + import torch + t = tensor.detach().cpu() + if dtype is not None: + target = getattr(torch, normalize_dtype_name(dtype, "float32")) + if t.dtype != target: + t = t.to(target) + if t.dtype == torch.bfloat16: + return mx.array(t.contiguous().view(torch.uint16).numpy()).view(mx.bfloat16) + return mx.array(t.contiguous().numpy()) diff --git a/chronos/runtime/cache_manager.py b/chronos/runtime/cache_manager.py index 2fabaff..86ed02e 100644 --- a/chronos/runtime/cache_manager.py +++ b/chronos/runtime/cache_manager.py @@ -35,7 +35,12 @@ class CacheManager: def __init__(self, model, config, ssd_dir: str = "./expert_cache"): self.config = config self.expert_store = ExpertStore(model, config, ssd_dir=ssd_dir) - self.prefetcher = AsyncPrefetcher(self.expert_store) + # The queue is part of the memory budget: a large lookahead queue can + # retain many decoded expert ids and defeat the RAM cap. Tie it to + # the canonical prefetch_depth setting while keeping a small minimum + # so a zero-depth diagnostic configuration remains usable. + queue_depth = max(1, int(getattr(config, "prefetch_depth", 2) or 0) * 2) + self.prefetcher = AsyncPrefetcher(self.expert_store, queue_depth=queue_depth) self.scheduler = PrefetchScheduler(self.prefetcher, self.expert_store) self._num_layers = len(self.expert_store.moe_layers) @@ -62,9 +67,9 @@ def warm_up(self, initial_expert_ids: Optional[List[int]] = None): Defaults to experts 0..vram_capacity-1. """ if initial_expert_ids is None: - initial_expert_ids = list(range( - min(self.config.num_experts, self.expert_store.vram_capacity) - )) + initial_expert_ids = self.expert_store._expand_keys( + list(range(self.config.num_experts)) + )[:self.expert_store.vram_capacity] self.expert_store.prefetch_to_ram(initial_expert_ids) for eid in initial_expert_ids: self.expert_store.promote_to_vram(eid) @@ -76,9 +81,8 @@ def availability_mask(self) -> torch.Tensor: return self.expert_store.vram_availability_mask() def availability_masks_all_layers(self) -> List[torch.Tensor]: - """Returns the same mask replicated for each MoE layer.""" - mask = self.availability_mask() - return [mask] * self._num_layers + """Return actual per-layer masks (v1 stores return cloned globals).""" + return self.expert_store.vram_availability_masks() def step(self, lookahead_probs, current_expert_ids: List[int]): """ @@ -99,7 +103,13 @@ def prefetch_for_next_step(self, lookahead_probs) -> List[int]: def prefetch_experts_to_ram(self, expert_ids: List[int]) -> List[int]: """Queue exact expert IDs for SSD->RAM prefetch without touching VRAM.""" - expert_ids = list(dict.fromkeys(int(eid) for eid in expert_ids or [])) + normalized = [] + for eid in expert_ids or []: + if isinstance(eid, (tuple, list)): + normalized.append(tuple(int(x) for x in eid)) + else: + normalized.append(int(eid)) + expert_ids = list(dict.fromkeys(normalized)) if expert_ids: self.prefetcher.submit(expert_ids) return expert_ids @@ -114,7 +124,7 @@ def ensure_resident(self, expert_ids: List[int]) -> List[int]: # The compute stream waits on each newly-promoted expert's event; # already-resident experts contribute no event, so this only blocks # on the truly necessary transfers. - self.expert_store.wait_for_experts(newly) + self.expert_store.wait_for_experts(expert_ids) return newly def stats(self) -> dict: diff --git a/chronos/runtime/inference_engine.py b/chronos/runtime/inference_engine.py index 98b2f20..fd0e1ec 100644 --- a/chronos/runtime/inference_engine.py +++ b/chronos/runtime/inference_engine.py @@ -17,11 +17,17 @@ import chronos.deps # ensure minimind on sys.path import time -from typing import Dict, List, Optional +from typing import Dict, Hashable, List, Optional import torch import torch.nn.functional as F +try: + from chronos.runtime.metrics import safe_record as _metric +except Exception: # pragma: no cover + def _metric(_name, _value): + return None + from chronos.model.config import ChronosConfig from chronos.model.model_chronos import ChronosForCausalLM from chronos.model.moe_chronos import ChronosMOEFeedForward @@ -90,7 +96,7 @@ def __init__( self.cache_manager = CacheManager(model, config, ssd_dir=ssd_dir) self.last_stats: dict = {} self._runtime_stats: dict = {} - self._route_cache_scores: Dict[int, float] = {} + self._route_cache_scores: Dict[Hashable, float] = {} self._last_prefill_ram_expert_ids: List[int] = [] def _moe_layers(self) -> List[ChronosMOEFeedForward]: @@ -118,21 +124,29 @@ def _normalize_miss_policy(miss_policy: str | None) -> str: def _install_moe_runtime_hooks(self, miss_policy: str): store = self.cache_manager.expert_store - def load_one(expert_id: int) -> dict: + def cache_key(layer_id: int, expert_id: int): + return ( + store.expert_key(layer_id, expert_id) + if bool(getattr(self.config, "layer_aware_cache", False)) + else int(expert_id) + ) + + def load_one(layer_id: int, expert_id: int) -> dict: expert_id = int(expert_id) - already_vram = store.vram_lru.contains(expert_id) - already_ram = store.ram_lru.contains(expert_id) + key = cache_key(layer_id, expert_id) + already_vram = store.vram_lru.contains(key) + already_ram = store.ram_lru.contains(key) t0 = time.monotonic() ok = False try: if already_vram or already_ram or miss_policy == "sync_on_demand": - promoted = self.cache_manager.ensure_resident([expert_id]) - ok = store.vram_lru.contains(expert_id) + promoted = self.cache_manager.ensure_resident([key]) + ok = store.vram_lru.contains(key) else: # README architecture: cold SSD misses must not enter the # decode critical path. Queue the exact missing expert for # RAM prefetch and let shared fallback cover this token. - self.cache_manager.prefetch_experts_to_ram([expert_id]) + self.cache_manager.prefetch_experts_to_ram([key]) promoted = [] ok = False finally: @@ -153,18 +167,21 @@ def load_one(expert_id: int) -> dict: int(self._runtime_stats.get("sync_ssd_loads", 0)) + 1 ) if ok: - self._runtime_stats["on_demand_loaded_experts"].append(expert_id) + self._runtime_stats["on_demand_loaded_experts"].append(key) if promoted: self._runtime_stats["on_demand_promoted_experts"].extend(promoted) return {"ok": ok, "already_vram": already_vram, "already_ram": already_ram} - def touch_one(expert_id: int) -> None: - store.touch_expert(int(expert_id)) + def touch_one(layer_id: int, expert_id: int) -> None: + store.touch_expert(cache_key(layer_id, int(expert_id))) - for moe in self._moe_layers(): + for layer_id, moe in enumerate(self._moe_layers()): moe.runtime_miss_policy = miss_policy - moe.runtime_on_demand_loader = load_one if miss_policy in {"on_demand", "sync_on_demand"} else None - moe.runtime_touch_expert = touch_one + moe.runtime_on_demand_loader = ( + (lambda expert_id, li=layer_id: load_one(li, expert_id)) + if miss_policy in {"on_demand", "sync_on_demand"} else None + ) + moe.runtime_touch_expert = lambda expert_id, li=layer_id: touch_one(li, expert_id) def _clear_moe_runtime_hooks(self): for moe in self._moe_layers(): @@ -188,28 +205,36 @@ def _consume_moe_runtime_stats(self) -> dict: def _recent_route_scores( self, - current_ids: List[int], - predicted_ids: List[int], - ram_predicted_ids: Optional[List[int]] = None, - ) -> Dict[int, float]: + current_ids: List[Hashable], + predicted_ids: List[Hashable], + ram_predicted_ids: Optional[List[Hashable]] = None, + ) -> Dict[Hashable, float]: """Decayed route scores used as a sticky eviction priority. Higher scores stay hotter in both RAM and VRAM. This intentionally spans tokens; replacing it with only the current step turns tiny budgets into a promote/evict loop. """ + def normalized(key): + if isinstance(key, (tuple, list)): + return tuple(int(value) for value in key) + return int(key) + decayed = { - int(eid): float(score) * 0.85 + normalized(eid): float(score) * 0.85 for eid, score in self._route_cache_scores.items() if float(score) * 0.85 >= 0.05 } scores = decayed for eid in current_ids: - scores[int(eid)] = scores.get(int(eid), 0.0) + 3.0 + key = normalized(eid) + scores[key] = scores.get(key, 0.0) + 3.0 for eid in predicted_ids: - scores[int(eid)] = scores.get(int(eid), 0.0) + 2.0 + key = normalized(eid) + scores[key] = scores.get(key, 0.0) + 2.0 for eid in ram_predicted_ids or []: - scores[int(eid)] = scores.get(int(eid), 0.0) + 1.0 + key = normalized(eid) + scores[key] = scores.get(key, 0.0) + 1.0 self._route_cache_scores = scores return scores @@ -229,18 +254,18 @@ def _ordered_ids_from_scores(scores: torch.Tensor, capacity: int) -> List[int]: ordered = torch.argsort(scores, descending=True).tolist() return [int(eid) for eid in ordered if float(scores[eid].item()) > 0.0][:capacity] - def _ordered_sticky_ids(self, capacity: int) -> List[int]: + def _ordered_sticky_ids(self, capacity: int) -> List[Hashable]: if capacity <= 0 or not self._route_cache_scores: return [] return [ - int(eid) + eid for eid, _score in sorted( self._route_cache_scores.items(), - key=lambda item: (-float(item[1]), int(item[0])), + key=lambda item: (-float(item[1]), repr(item[0])), ) ][:capacity] - def _promotion_ready_ids(self, expert_ids: List[int]) -> List[int]: + def _promotion_ready_ids(self, expert_ids: List[Hashable]) -> List[Hashable]: """Only return IDs already staged in RAM/VRAM. Calling ensure_resident on a cold SSD expert during decode would @@ -249,7 +274,11 @@ def _promotion_ready_ids(self, expert_ids: List[int]) -> List[int]: """ store = self.cache_manager.expert_store ready = [] - for eid in dict.fromkeys(int(e) for e in expert_ids or []): + normalized = ( + tuple(int(value) for value in eid) if isinstance(eid, (tuple, list)) else int(eid) + for eid in expert_ids or [] + ) + for eid in dict.fromkeys(normalized): if store.vram_lru.contains(eid) or store.ram_lru.contains(eid): ready.append(eid) return ready @@ -300,6 +329,25 @@ def setup_from_state_reader( def teardown(self): self.cache_manager.stop() + @staticmethod + def _virtual_router_probs(moe, last_token: bool = False) -> Optional[torch.Tensor]: + """Router probs of every dispatch of ``moe`` in the last forward. + + A recurrent group's MoE runs once per loop iteration and each + iteration may pick different experts, so prefetch/residency must + cover the union over all D iterations (virtual layers), not only + the last one. Returns ``[B, S*D, E]`` (``[B, D, E]`` for + ``last_token``) on the CPU, or None before the first forward. + """ + probs = list(getattr(moe, "iteration_router_probs", None) or ()) + if not probs: + probs = [moe.last_router_probs] if moe.last_router_probs is not None else [] + if not probs: + return None + if last_token: + probs = [p[:, -1:, :] for p in probs] + return torch.cat([p.detach() for p in probs], dim=1).cpu() + def _get_current_expert_ids(self, capacity: Optional[int] = None) -> List[int]: """Extract active top-k expert IDs from every MoE layer. @@ -312,10 +360,10 @@ def _get_current_expert_ids(self, capacity: Optional[int] = None) -> List[int]: return [] scores = torch.zeros(num_experts, dtype=torch.float32) for moe in moe_layers: - if moe.last_router_probs is None: + probs = self._virtual_router_probs(moe, last_token=True) # [B, D, E] + if probs is None: continue - probs = moe.last_router_probs[:, -1, :].detach().cpu() # [B, E] - weights, ids = self._router_topk(probs) + weights, ids = self._router_topk(probs.reshape(-1, probs.shape[-1])) scores.index_add_(0, ids.reshape(-1), weights.reshape(-1).to(scores.dtype)) if capacity is None: capacity = num_experts @@ -332,9 +380,43 @@ def _predict_expert_ids( if lookahead_probs is None: return [] probs = lookahead_probs.detach().cpu() - if probs.dim() != 4 or probs.shape[-1] == 0: + if probs.dim() not in (4, 5) or probs.shape[-1] == 0: return [] + if probs.dim() == 5 and not bool(getattr(self.config, "layer_aware_cache", False)): + probs = probs.mean(dim=3) num_experts = int(probs.shape[-1]) + layer_aware = probs.dim() == 5 and bool(getattr(self.config, "layer_aware_cache", False)) + if layer_aware: + # Return physical-layer page keys. The recurrent group is encoded + # by the prediction axis and ExpertStore validates the mapping. + group_axis = probs.shape[-2] + layer_map = getattr(self.cache_manager.expert_store, "_layer_to_group", {}) + if all_future: + future = probs[:, -1, 1:, :, :] + offsets = range(future.shape[1]) + else: + if probs.shape[2] <= future_offset: + return [] + future = probs[:, -1, future_offset:future_offset + 1, :, :] + offsets = [0] + selected = [] + for offset in offsets: + for group in range(group_axis): + step_probs = future[:, offset, group, :] + weights, ids = self._router_topk(step_probs) + scale = 1.0 / float(offset + 1) if all_future else 1.0 + for eid, weight in zip(ids.reshape(-1).tolist(), weights.reshape(-1).tolist()): + selected.append(((group, int(eid)), float(weight) * scale)) + score_map = {} + for (group, eid), score in selected: + score_map[(group, eid)] = score_map.get((group, eid), 0.0) + score + capacity = num_experts if capacity is None else int(capacity) + ordered = sorted(score_map, key=lambda key: (-score_map[key], key[0], key[1])) + out = [] + for group, eid in ordered[:capacity]: + layers = [li for li, gi in layer_map.items() if gi == group] + out.extend((int(li), int(group), int(eid)) for li in layers) + return out scores = torch.zeros(num_experts, dtype=torch.float32) if all_future: if probs.shape[2] <= 1: @@ -375,18 +457,20 @@ def _route_cache_stats( "fallback_weight_mass": 0.0, "total_route_weight_mass": 0.0, "activated_expert_ids": [], + "activated_expert_keys": [], "per_layer": [], } if not moe_layers: return totals activated: set[int] = set() + activated_keys: set[tuple[int, int, int]] = set() + store = self.cache_manager.expert_store for layer_idx, moe in enumerate(moe_layers): - if moe.last_router_probs is None: + # Every loop iteration is a real dispatch: count all of them. + probs = self._virtual_router_probs(moe, last_token=not include_all_tokens) + if probs is None: continue - probs = moe.last_router_probs.detach().cpu() - if not include_all_tokens: - probs = probs[:, -1:, :] flat_probs = probs.reshape(-1, probs.shape[-1]) weights, ids = self._router_topk(flat_probs) flat_ids = ids.reshape(-1).to(torch.long) @@ -415,6 +499,11 @@ def _route_cache_stats( total_mass = hit_mass + miss_mass activated.update(int(eid) for eid in flat_ids.unique().tolist()) + group_idx = int(getattr(store, "_layer_to_group", {}).get(layer_idx, -1)) + activated_keys.update( + (int(layer_idx), group_idx, int(eid)) + for eid in flat_ids.unique().tolist() + ) totals["cache_hits"] += hits totals["cache_misses"] += misses totals["hit_weight_mass"] += hit_mass @@ -431,12 +520,15 @@ def _route_cache_stats( }) totals["activated_expert_ids"] = sorted(activated) + totals["activated_expert_keys"] = sorted(activated_keys) return totals def _plan_prefill_experts( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, + modality_inputs: Optional[dict] = None, + modality_mask: Optional[torch.Tensor] = None, ) -> tuple[List[int], List[int]]: """Probe prompt routing with shared fallback, then load hot experts. @@ -483,6 +575,8 @@ def _plan_prefill_experts( past_key_values=None, use_cache=False, available_expert_masks=masks, + modality_inputs=modality_inputs, + modality_mask=modality_mask, ) self._runtime_stats["on_demand_loads"] = probe_start["on_demand_loads"] self._runtime_stats["on_demand_load_time_s"] = probe_start["on_demand_load_time_s"] @@ -498,10 +592,12 @@ def _plan_prefill_experts( scores = torch.zeros(num_experts, dtype=torch.float32) for layer in self.model.model.layers: moe = getattr(layer, "mlp", None) - if not isinstance(moe, ChronosMOEFeedForward) or moe.last_router_probs is None: + if not isinstance(moe, ChronosMOEFeedForward): + continue + probs = self._virtual_router_probs(moe) + if probs is None: continue - probs = moe.last_router_probs.reshape(-1, num_experts) - weights, expert_ids = self._router_topk(probs.detach().cpu()) + weights, expert_ids = self._router_topk(probs.reshape(-1, num_experts)) scores.index_add_( 0, expert_ids.reshape(-1).to(torch.long), @@ -521,7 +617,130 @@ def _plan_prefill_experts( ram_ids = self._ordered_ids_from_scores(scores, ram_capacity) return ram_ids[:vram_capacity], ram_ids - @torch.inference_mode() + # no_grad, not inference_mode: experts paged in here are reused by + # generate(); inference tensors would poison later autograd-tracked calls. + @torch.no_grad() + def generate_speculative( + self, + input_ids: torch.Tensor, + *, + max_new_tokens: int = 256, + draft_depth: int = 1, + gamma: int = 4, + gamma_max: int = 8, + adaptive_gamma: bool = True, + temperature: float = 0.0, + top_p: float = 1.0, + top_k: int = 0, + do_sample: bool = False, + eos_token_id=None, + miss_policy: str = "on_demand", + generator: Optional[torch.Generator] = None, + on_token=None, + ) -> torch.Tensor: + """Recurrent-depth self-speculative decoding with lazy experts. + + Batch size 1. Each draft and verify forward runs with the current + expert availability masks; before it, the lookahead router's + prediction from the previous forward is prefetched to RAM and the + promotable part moved to VRAM, so drafts warm the experts their own + verify pass needs. Greedy output equals :meth:`generate` greedy + output token for token. ``per_iteration`` / non-recurrent models + decode plainly (``draft_depth`` = full depth). + """ + from chronos.runtime.speculative import speculative_generate, speculative_supported + input_ids = input_ids.to(self.device) + miss_policy = self._normalize_miss_policy(miss_policy) + full_depth = int(getattr(self.config, "max_inference_depth", None) or self.config.recurrent_depth) + use_spec = speculative_supported(self.config) and int(draft_depth) < full_depth + self._route_cache_scores = {} + self._runtime_stats = { + "on_demand_loads": 0, "on_demand_load_time_s": 0.0, "sync_ssd_loads": 0, + "async_cold_miss_prefetches": 0, "resident_vram_hits": 0, "resident_ram_hits": 0, + "on_demand_loaded_experts": [], "on_demand_promoted_experts": [], + } + cm = self.cache_manager + store = cm.expert_store + state = {"lookahead": None, "hits": 0, "misses": 0, "forwards": 0, "fallback_mass": 0.0, + "on_demand": 0, "prefill_s": None} + t0 = time.monotonic() + rss0 = _rss_gb() + backend0 = _backend_memory_snapshot(self.device) + self._install_moe_runtime_hooks(miss_policy) + + def forward_fn(ids, cache, depth): + if state["lookahead"] is not None: + predicted = self._predict_expert_ids(state["lookahead"], all_future=True, + capacity=int(store.ram_capacity)) + cm.prefetch_experts_to_ram(predicted) + ready = self._promotion_ready_ids(predicted[: int(store.vram_capacity)]) + if ready: + cm.ensure_resident(ready) + masks = cm.availability_masks_all_layers() + out = self.model(ids.to(self.device), past_key_values=cache, use_cache=True, + recurrent_depth=depth, available_expert_masks=masks) + state["lookahead"] = out.lookahead_probs + route = self._route_cache_stats(masks, include_all_tokens=cache is None) + state["hits"] += int(route["cache_hits"]) + state["misses"] += int(route["cache_misses"]) + rt = self._consume_moe_runtime_stats() + state["fallback_mass"] += float(rt["fallback_executed_weight_mass"]) + state["on_demand"] += len(rt["on_demand_experts"]) + state["forwards"] += 1 + if state["prefill_s"] is None: + state["prefill_s"] = time.monotonic() - t0 + return out + + try: + prefill_ids, prefill_ram_ids = self._plan_prefill_experts(input_ids, None, None, None) + if prefill_ram_ids: + store.prefetch_to_ram(prefill_ram_ids) + if prefill_ids: + cm.ensure_resident(prefill_ids) + out, spec = speculative_generate( + self.model, input_ids, max_new_tokens=int(max_new_tokens), + draft_depth=int(draft_depth) if use_spec else full_depth, gamma=int(gamma), + gamma_max=int(gamma_max), adaptive_gamma=bool(adaptive_gamma), do_sample=bool(do_sample), + temperature=max(float(temperature), 1e-6), top_k=int(top_k or 0), top_p=float(top_p or 1.0), + eos_token_id=eos_token_id, generator=generator, on_token=on_token, forward_fn=forward_fn) + finally: + self._clear_moe_runtime_hooks() + elapsed = time.monotonic() - t0 + generated = int(out.shape[1] - input_ids.shape[1]) + total = state["hits"] + state["misses"] + runtime_stats = dict(self._runtime_stats) + prefill_s = float(state["prefill_s"] or 0.0) + self.last_stats = { + "tokens_generated": generated, + "elapsed_s": round(elapsed, 3), + "tokens_per_sec": round(generated / max(elapsed, 1e-6), 2), + "prefill_time_s": round(prefill_s, 3), + "decode_time_s": round(max(0.0, elapsed - prefill_s), 3), + "rss_before_prefill_gb": round(rss0, 3), + "rss_after_decode_gb": round(_rss_gb(), 3), + "decode_rss_delta_gb": round(_rss_gb() - rss0, 3), + **_backend_memory_fields(backend0, "before_prefill"), + **_backend_memory_fields(_backend_memory_snapshot(self.device), "after_decode"), + "cache_hits": state["hits"], "cache_misses": state["misses"], + "cache_hit_rate": round(state["hits"] / max(total, 1), 4), + "resident_hit_rate": round(state["hits"] / max(total, 1), 4), + "fallback_weight_mass": round(state["fallback_mass"], 4), + "on_demand_loads": int(runtime_stats.get("on_demand_loads", 0)), + "sync_ssd_loads": int(runtime_stats.get("sync_ssd_loads", 0)), + "async_cold_miss_prefetches": int(runtime_stats.get("async_cold_miss_prefetches", 0)), + "resident_ram_hits": int(runtime_stats.get("resident_ram_hits", 0)), + "miss_policy": miss_policy, + "quality_safe": miss_policy == "sync_on_demand", + "forwards": state["forwards"], + "prefill_expert_ids": prefill_ids, + "prefill_ram_expert_ids": prefill_ram_ids, + "speculative": spec.as_dict() | {"enabled": use_spec, "full_depth": full_depth, + "draft_depth": int(draft_depth) if use_spec else full_depth, + "gamma_final": spec.gamma_history[-1] if spec.gamma_history else int(gamma)}, + } + return out + + @torch.no_grad() def generate( self, input_ids: torch.Tensor, @@ -536,6 +755,8 @@ def generate( repetition_penalty: float = 1.0, scheduler: Optional[PrefillScheduler] = None, miss_policy: str = "on_demand", + modality_inputs: Optional[dict] = None, + modality_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: input_ids = input_ids.to(self.device) if attention_mask is not None: @@ -559,6 +780,7 @@ def generate( # ── Prefill phase: front-load expert IO before generation ────── prefill_t0 = time.monotonic() prefill_mem0 = _rss_gb() + _io0 = dict(getattr(self.cache_manager.expert_store, "_io_stats", {}) or {}) prefill_backend_mem0 = _backend_memory_snapshot(self.device) if scheduler is not None: scheduler.prepare(input_ids, device=str(self.device)) @@ -566,7 +788,9 @@ def generate( prefill_expert_ids: List[int] = [] prefill_ram_expert_ids: List[int] = [] else: - prefill_expert_ids, prefill_ram_expert_ids = self._plan_prefill_experts(input_ids, attention_mask) + prefill_expert_ids, prefill_ram_expert_ids = self._plan_prefill_experts( + input_ids, attention_mask, modality_inputs, modality_mask + ) self._last_prefill_ram_expert_ids = prefill_ram_expert_ids if prefill_ram_expert_ids: self.cache_manager.expert_store.prefetch_to_ram(prefill_ram_expert_ids) @@ -594,6 +818,8 @@ def generate( fallback_executed_weight_mass = 0.0 prediction_hits = 0 prediction_total = 0 + cluster_prediction_hits = 0 + cluster_prediction_total = 0 # M3: holds last step's lookahead predictions; we submit them as # prefetch BEFORE this step's forward so the H2D stream overlaps # with compute. None on the first iteration (no predictions yet). @@ -624,6 +850,10 @@ def generate( capacity=vram_capacity, ) predicted_current_expert_ids.extend(predicted_ids) + # Predicted pages must survive this step's evictions. + protect = getattr(store, "set_protected_experts", None) + if protect is not None: + protect(list(dict.fromkeys(predicted_ids + prev_current_ids))) submitted_prefetch_ids.extend( self.cache_manager.prefetch_experts_to_ram(ram_predicted_ids) ) @@ -653,14 +883,18 @@ def generate( else: masks = self.cache_manager.availability_masks_all_layers() - past_len = past_key_values[0][0].shape[1] if past_key_values else 0 - outputs, lookahead_probs = self.model( + from chronos.model.model_chronos import cache_sequence_length + past_len = cache_sequence_length(past_key_values) if past_key_values else 0 + result = self.model( input_ids[:, past_len:], attention_mask=attention_mask, past_key_values=past_key_values, use_cache=True, available_expert_masks=masks, + modality_inputs=modality_inputs, + modality_mask=modality_mask, ) + outputs, lookahead_probs = result.legacy_tuple() if attention_mask is not None: attention_mask = torch.cat( @@ -677,17 +911,42 @@ def generate( fallback_weight_mass += float(route_stats["fallback_weight_mass"]) total_route_weight_mass += float(route_stats["total_route_weight_mass"]) route_layer_stats = route_stats["per_layer"] - current_ids = route_stats["activated_expert_ids"] - activated_expert_ids.extend(current_ids) + current_ids = ( + route_stats["activated_expert_keys"] + if bool(getattr(self.config, "layer_aware_cache", False)) + else route_stats["activated_expert_ids"] + ) + activated_expert_ids.extend(route_stats["activated_expert_ids"]) rt = self._consume_moe_runtime_stats() on_demand_expert_ids.extend(rt["on_demand_experts"]) fallback_expert_ids.extend(rt["fallback_experts"]) on_demand_weight_mass += float(rt["on_demand_weight_mass"]) fallback_executed_weight_mass += float(rt["fallback_executed_weight_mass"]) if predicted_ids: - current_set = set(int(eid) for eid in current_ids) + current_set = set(current_ids) prediction_total += len(current_set) - prediction_hits += len(current_set.intersection(int(eid) for eid in predicted_ids)) + prediction_hits += len(current_set.intersection(predicted_ids)) + cluster_storage = getattr( + self.cache_manager.expert_store, "_cluster_storage", None + ) + if cluster_storage is not None: + manifest = cluster_storage.manifest + def _eid(key): + return int(key[2]) if isinstance(key, tuple) else int(key) + predicted_clusters = { + manifest.expert_to_cluster[_eid(key)] + for key in predicted_ids + if _eid(key) in manifest.expert_to_cluster + } + current_clusters = [ + manifest.expert_to_cluster[_eid(key)] + for key in current_set + if _eid(key) in manifest.expert_to_cluster + ] + cluster_prediction_total += len(current_clusters) + cluster_prediction_hits += sum( + int(cid in predicted_clusters) for cid in current_clusters + ) if scheduler is None: self.cache_manager.expert_store.reprioritize_resident_experts( self._recent_route_scores(current_ids, predicted_ids, ram_predicted_ids) @@ -703,7 +962,8 @@ def generate( for i in range(input_ids.shape[0]): logits[i, torch.unique(input_ids[i])] /= repetition_penalty if top_k > 0: - logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf') + k = min(int(top_k), int(logits.shape[-1])) + logits[logits < torch.topk(logits, k)[0][..., -1, None]] = -float('inf') if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True) mask_p = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) > top_p @@ -764,12 +1024,33 @@ def generate( fallback_weight_rate = fallback_weight_mass / max(total_route_weight_mass, 1e-12) executed_fallback_rate = fallback_executed_weight_mass / max(total_route_weight_mass, 1e-12) prediction_hit_rate = prediction_hits / max(prediction_total, 1) + cluster_recall = cluster_prediction_hits / max(cluster_prediction_total, 1) runtime_stats = dict(self._runtime_stats) promoted_expert_ids.extend(runtime_stats.get("on_demand_promoted_experts", [])) backend_fields = {} backend_fields.update(_backend_memory_fields(prefill_backend_mem0, "before_prefill")) backend_fields.update(_backend_memory_fields(prefill_backend_mem1, "after_prefill")) backend_fields.update(_backend_memory_fields(end_backend_mem, "after_decode")) + recurrent = getattr(self.model.model, "last_recurrent_state", None) or {} + depth_used = recurrent.get("depth_used") + average_ponder_steps = ( + float(depth_used.float().mean().item()) if depth_used is not None else 1.0 + ) + early_exit = recurrent.get("early_exit_ratio") + early_exit_ratio = ( + float(early_exit.item()) if isinstance(early_exit, torch.Tensor) + else float(early_exit or 0.0) + ) + contraction = recurrent.get("contraction") + contraction_ratio = ( + float(contraction.item()) if isinstance(contraction, torch.Tensor) + else float(contraction or 0.0) + ) + on_demand_loads = int(runtime_stats.get("on_demand_loads", 0)) + resident_ram_hits = int(runtime_stats.get("resident_ram_hits", 0)) + cold_misses = int(runtime_stats.get("async_cold_miss_prefetches", 0)) + int( + runtime_stats.get("sync_ssd_loads", 0) + ) self.last_stats = { "tokens_generated": tokens_generated, "elapsed_s": round(elapsed, 3), @@ -782,8 +1063,11 @@ def generate( "rss_after_prefill_gb": round(prefill_mem1, 3), "rss_after_decode_gb": round(end_mem, 3), **backend_fields, - "cache_hit_rate": round(1.0 if miss_policy == "on_demand" and executed_fallback_rate == 0.0 else resident_hit_rate, 4), + "cache_hit_rate": round(resident_hit_rate, 4), "resident_hit_rate": round(resident_hit_rate, 4), + "resident_vram_hit_rate": round(resident_hit_rate, 4), + "ram_promotion_hit_rate": round(resident_ram_hits / max(on_demand_loads, 1), 4), + "ssd_cold_miss_rate": round(cold_misses / max(on_demand_loads, 1), 4), "cache_hits": cache_hits, "cache_misses": cache_misses, "fallback_weight_rate": round(executed_fallback_rate, 4), @@ -803,21 +1087,55 @@ def generate( "quality_safe": miss_policy == "sync_on_demand", "miss_policy": miss_policy, "prediction_hit_rate": round(prediction_hit_rate, 4), + "prediction_topk_recall": round(prediction_hit_rate, 4), + "cluster_recall": round(cluster_recall, 4), "prediction_hits": prediction_hits, "prediction_total": prediction_total, + "cluster_prediction_hits": cluster_prediction_hits, + "cluster_prediction_total": cluster_prediction_total, + "average_ponder_steps": round(average_ponder_steps, 4), + "early_exit_ratio": round(early_exit_ratio, 4), + "recurrent_contraction": round(contraction_ratio, 6), "route_layer_stats": route_layer_stats, "prefill_expert_ids": prefill_expert_ids, "prefill_ram_expert_ids": prefill_ram_expert_ids, "activated_expert_ids": sorted(set(activated_expert_ids)), - "predicted_current_expert_ids": sorted(set(predicted_current_expert_ids)), - "prefetch_submitted_ids": sorted(set(submitted_prefetch_ids)), - "promoted_expert_ids": sorted(set(promoted_expert_ids)), - "on_demand_expert_ids": sorted(set(on_demand_expert_ids)), - "fallback_expert_ids": sorted(set(fallback_expert_ids)), + "predicted_current_expert_ids": sorted(set(predicted_current_expert_ids), key=repr), + "prefetch_submitted_ids": sorted(set(submitted_prefetch_ids), key=repr), + "promoted_expert_ids": sorted(set(promoted_expert_ids), key=repr), + "on_demand_expert_ids": sorted(set(on_demand_expert_ids), key=repr), + "fallback_expert_ids": sorted(set(fallback_expert_ids), key=repr), "prefetch_submitted_total": len(submitted_prefetch_ids), "promoted_total": len(promoted_expert_ids), **self.cache_manager.stats(), } + store = self.cache_manager.expert_store + pages_read = int(getattr(store, "_io_stats", {}).get("ssd_page_reads", 0)) - int(_io0.get("ssd_page_reads", 0)) + page_bytes = int(getattr(store, "_expert_bytes", 0) or 0) + ssd_mb = pages_read * page_bytes / 2 ** 20 + self.last_stats.update({ + "ssd_pages_read": pages_read, + "ssd_read_mb": round(ssd_mb, 3), + "ssd_read_mb_per_s": round(ssd_mb / max(elapsed, 1e-6), 3), + "resident_experts": int(self.last_stats.get("vram_experts", 0)), + }) + if hasattr(store, "set_protected_experts"): + store.set_protected_experts(()) + # Publish a compact rolling series for the IO Monitor without + # retaining per-request payloads in the hot path. + for metric_name, metric_value in ( + ("tokens_per_sec", self.last_stats["tokens_per_sec"]), + ("resident_vram_hit_rate", self.last_stats["resident_vram_hit_rate"]), + ("ram_promotion_hit_rate", self.last_stats["ram_promotion_hit_rate"]), + ("ssd_cold_miss_rate", self.last_stats["ssd_cold_miss_rate"]), + ("prediction_topk_recall", self.last_stats["prediction_topk_recall"]), + ("cluster_recall", self.last_stats["cluster_recall"]), + ("average_ponder_steps", self.last_stats["average_ponder_steps"]), + ("early_exit_ratio", self.last_stats["early_exit_ratio"]), + ("recurrent_contraction", self.last_stats["recurrent_contraction"]), + ("fallback_weight_mass", self.last_stats["fallback_weight_mass"]), + ): + _metric(metric_name, metric_value) if streamer: streamer.end() diff --git a/chronos/runtime/speculative.py b/chronos/runtime/speculative.py new file mode 100644 index 0000000..ea65bcb --- /dev/null +++ b/chronos/runtime/speculative.py @@ -0,0 +1,225 @@ +"""Recurrent-depth self-speculative decoding (no extra parameters). + +Draft: the same model run with fewer loop iterations (``draft_depth``) +proposes γ tokens. Verify: one full-depth forward over the γ+1 positions +scores them in a single batched pass (weights/experts read once). Greedy +outputs are token-for-token identical to full-depth greedy decoding; sampled +outputs use standard rejection sampling, which preserves the full-depth +distribution exactly. + +Cache contract: draft KV is written to a throw-away cache; the persistent +cache is only ever extended by verified forwards, then truncated to the +accepted prefix (MLA slices ``c_kv``; sliding-window layers keep +``window + γ_max`` entries so a rollback never loses in-window keys). +""" +from __future__ import annotations + +import time +from dataclasses import dataclass, field + +import torch + + +def _truncate_entry(entry, length: int): + if entry is None: + return None + if len(entry) <= 2: # MLA latent cache (c_kv,) or (c_kv, k_rope); tensors or QuantizedKV + return tuple(t[:, :length] for t in entry) + k, v, end = entry + drop = int(end) - length + if drop < 0: + raise ValueError("cannot extend a cache by truncation") + if drop == 0: + return entry + if drop > k.shape[1]: + raise ValueError(f"rollback of {drop} exceeds retained sliding-window cache {k.shape[1]}") + return (k[:, :k.shape[1] - drop], v[:, :v.shape[1] - drop], length) + + +def truncate_cache(cache, length: int): + return [_truncate_entry(e, length) for e in cache] + + +def set_cache_slack(model, slack: int) -> None: + from chronos.model.hybrid_attention import SlidingWindowAttention + for module in model.modules(): + if isinstance(module, SlidingWindowAttention): + module.cache_slack = max(int(getattr(module, "cache_slack", 0)), int(slack)) + + +def speculative_supported(cfg) -> bool: + """True when a shallow draft can share the KV stream with full-depth verify. + + Needs a recurrent core deeper than one iteration and a KV layout whose + prefix does not depend on depth (``shared_first`` / ``first_only``). + ``per_iteration`` writes one KV slot per virtual layer, so callers fall + back to plain decoding (``draft_depth = full depth``), as the MLX loop does. + """ + from chronos.config_schema import resolve_recurrent_attention_mode + depth = int(getattr(cfg, "max_inference_depth", None) or getattr(cfg, "recurrent_depth", 1) or 1) + recurrent = bool(getattr(cfg, "recurrent_enabled", True)) and depth > 1 + return recurrent and resolve_recurrent_attention_mode(cfg) != "per_iteration" + + +@dataclass +class SpecStats: + proposed: int = 0 + accepted: int = 0 + verify_calls: int = 0 + draft_calls: int = 0 + generated: int = 0 + elapsed_s: float = 0.0 + gamma_history: list = field(default_factory=list) + + @property + def acceptance_rate(self) -> float: + return self.accepted / self.proposed if self.proposed else 0.0 + + def as_dict(self) -> dict: + return {"proposed": self.proposed, "accepted": self.accepted, + "acceptance_rate": round(self.acceptance_rate, 4), + "verify_calls": self.verify_calls, "draft_calls": self.draft_calls, + "generated": self.generated, + "tokens_per_verify": round(self.generated / max(self.verify_calls, 1), 3), + "tokens_per_s": round(self.generated / self.elapsed_s, 2) if self.elapsed_s else None} + + +def _probs(logits: torch.Tensor, temperature: float, top_k: int, top_p: float) -> torch.Tensor: + logits = logits.float() / max(temperature, 1e-6) + if top_k and top_k > 0: + kth = torch.topk(logits, min(int(top_k), logits.shape[-1]))[0][..., -1, None] + logits = logits.masked_fill(logits < kth, float("-inf")) + if top_p < 1.0: + sorted_logits, idx = torch.sort(logits, descending=True) + cum = torch.cumsum(torch.softmax(sorted_logits, -1), -1) + mask = cum > top_p + mask[..., 1:] = mask[..., :-1].clone() + mask[..., 0] = False + logits = logits.masked_fill(mask.scatter(-1, idx, mask), float("-inf")) + return torch.softmax(logits, dim=-1) + + +def rejection_step(p: torch.Tensor, q: torch.Tensor, token: int, generator=None) -> tuple[bool, int]: + """Speculative-sampling acceptance (Leviathan et al. / Chen et al. 2023). + + ``token`` was drawn from the draft distribution ``q``. Accept with + probability min(1, p/q); otherwise draw from norm(max(p − q, 0)). The + returned token is distributed exactly as ``p``. + """ + ratio = (p[token] / q[token].clamp_min(1e-12)).clamp(max=1.0) + if bool(torch.rand((), generator=generator, device=p.device) <= ratio): + return True, int(token) + residual = (p - q).clamp_min(0) + total = residual.sum() + return False, int(torch.multinomial(residual / total if total > 0 else p, 1, generator=generator)) + + +@torch.no_grad() # not inference_mode: callers may page experts in and reuse them later +def speculative_generate(model, input_ids: torch.Tensor, *, max_new_tokens: int = 128, + draft_depth: int = 1, gamma: int = 4, adaptive_gamma: bool = True, + gamma_max: int = 8, do_sample: bool = False, temperature: float = 1.0, + top_k: int = 0, top_p: float = 1.0, eos_token_id=None, + generator: torch.Generator | None = None, on_token=None, + forward_fn=None, should_stop=None): + """Return (output_ids [1, T], SpecStats). Batch size 1 (serving path). + + Loop invariant: ``pending`` is final (already emitted) but not yet in the + persistent cache. Each round drafts from ``pending`` at shallow depth, + then one full-depth forward over ``[pending, d1..dγ]`` both writes + ``pending`` into the cache and verifies the drafts. A rejection's + correction (or the all-accepted bonus token) becomes the next pending + token, so no extra forward is ever needed. + """ + from chronos.model.model_chronos import cache_sequence_length + if input_ids.shape[0] != 1: + raise ValueError("speculative_generate supports batch size 1") + cfg = model.config + full_depth = int(getattr(cfg, "max_inference_depth", None) or cfg.recurrent_depth) + draft_depth = max(1, min(int(draft_depth), full_depth)) + eos = set([eos_token_id] if isinstance(eos_token_id, int) else (eos_token_id or [])) + gamma_max = max(1, int(gamma_max)) + gamma = max(1, min(int(gamma), gamma_max)) + set_cache_slack(model, gamma_max + 1) + device = input_ids.device + stats, t0 = SpecStats(), time.perf_counter() + tokens = input_ids[0].tolist() + # ``forward_fn(ids, cache, depth)`` lets a runtime wrap every forward with + # expert residency work (SSD lazy loading, lookahead prefetch, masks). + if forward_fn is None: + def forward_fn(ids, cache, depth): + return model(ids, past_key_values=cache, use_cache=True, recurrent_depth=depth) + + def dist(row): + return _probs(row, temperature, top_k, top_p) + + def pick(row): + if not do_sample: + return int(torch.argmax(row)) + return int(torch.multinomial(dist(row), 1, generator=generator)) + + def emit(tok): + tokens.append(tok) + stats.generated += 1 + if on_token: + on_token(tok) + + out = forward_fn(input_ids, None, None) + cache = out.past_key_values + pending = pick(out.logits[0, -1]) + emit(pending) + while stats.generated < max_new_tokens and pending not in eos: + if should_stop is not None and should_stop(): + break + committed = cache_sequence_length(cache) + room = max_new_tokens - stats.generated # tokens we may still emit + draft, draft_q, dcache, last = [pending], [], cache, pending + if draft_depth < full_depth: + for _ in range(min(gamma, room - 1) if room > 1 else 0): + d = forward_fn(torch.tensor([[last]], device=device), dcache, draft_depth) + stats.draft_calls += 1 + dcache = d.past_key_values + row = d.logits[0, -1] + last = pick(row) + if do_sample: + draft_q.append(dist(row)) + draft.append(last) + if last in eos: + break + v = forward_fn(torch.tensor([draft], device=device), cache, None) + stats.verify_calls += 1 + logits = v.logits[0] # logits[i] predicts the token after draft[i] + n_ok, nxt, stop = 0, None, False + for i in range(1, len(draft)): + row, tok = logits[i - 1], draft[i] + if not do_sample: + target = int(torch.argmax(row)) + ok = target == tok + correction = target + else: + ok, correction = rejection_step(dist(row), draft_q[i - 1], tok, generator) + if not ok: + nxt = correction + break + n_ok += 1 + emit(tok) + if tok in eos or stats.generated >= max_new_tokens: + stop = True + break + stats.proposed += len(draft) - 1 + stats.accepted += n_ok + cache = truncate_cache(v.past_key_values, committed + 1 + n_ok) + if stop: + break + if nxt is None: # every draft accepted: bonus token + nxt = pick(logits[len(draft) - 1]) + pending = nxt + emit(pending) + if adaptive_gamma and len(draft) > 1: + rate = n_ok / (len(draft) - 1) + if rate >= 0.8: + gamma = min(gamma_max, gamma + 1) + elif rate < 0.4: + gamma = max(1, gamma - 1) + stats.gamma_history.append(gamma) + stats.elapsed_s = time.perf_counter() - t0 + return torch.tensor([tokens], device=device), stats diff --git a/chronos/serving/__init__.py b/chronos/serving/__init__.py index c4b67c9..0040f92 100644 --- a/chronos/serving/__init__.py +++ b/chronos/serving/__init__.py @@ -11,3 +11,14 @@ "is_available", "HAS_VLLM", ] + +from .api_compat import ( + message_text, modality_inputs_from_messages, resolve_modality_inputs, openai_chat_response, + ollama_chat_response, ollama_generate_response, +) +from .adapter import ChronosServingAdapter + +__all__ += [ + "message_text", "modality_inputs_from_messages", "resolve_modality_inputs", "openai_chat_response", + "ollama_chat_response", "ollama_generate_response", "ChronosServingAdapter", +] diff --git a/chronos/serving/adapter.py b/chronos/serving/adapter.py new file mode 100644 index 0000000..033728b --- /dev/null +++ b/chronos/serving/adapter.py @@ -0,0 +1,121 @@ +"""Framework-neutral OpenAI/Ollama serving adapter. + +The adapter intentionally contains no FastAPI dependency. It accepts the +standard request dictionaries used by either HTTP framework and delegates +execution to :class:`ChronosInferenceEngine`; applications can mount the two +methods in FastAPI, Flask, Gradio, or a worker queue without changing model +semantics. +""" +from __future__ import annotations + +from typing import Any, Mapping, Optional + +import torch + +from .api_compat import ( + message_text, + openai_chat_response, + ollama_chat_response, + ollama_generate_response, + resolve_modality_inputs, +) + + +class ChronosServingAdapter: + """Expose a loaded Chronos engine through standard chat/generate schemas.""" + + def __init__(self, engine, tokenizer, *, model_name: str = "chronos"): + self.engine = engine + self.tokenizer = tokenizer + self.model_name = model_name + + def _encode(self, messages, prompt: Optional[str] = None): + text = prompt if prompt is not None else message_text(messages) + encoded = self.tokenizer(text, return_tensors="pt", add_special_tokens=True) + ids = encoded.input_ids + mask = encoded.get("attention_mask") if hasattr(encoded, "get") else getattr(encoded, "attention_mask", None) + modalities = resolve_modality_inputs( + messages, + image_size=getattr(self.engine.config, "image_size", 224), + audio_sample_rate=getattr(self.engine.config, "audio_sample_rate", 16000), + max_audio_seconds=getattr(self.engine.config, "max_audio_seconds", 30.0), + ) + return ids, mask, modalities + + def generate(self, *, messages=None, prompt=None, max_tokens: int = 128, + temperature: float = 0.85, top_p: float = 0.95, + top_k: int = 50, do_sample: bool = True, + miss_policy: str = "on_demand") -> tuple[str, dict]: + ids, mask, modalities = self._encode(messages or [], prompt=prompt) + output = self.engine.generate( + ids, + attention_mask=mask, + max_new_tokens=int(max_tokens), + temperature=float(temperature), + top_p=float(top_p), + top_k=int(top_k), + do_sample=bool(do_sample), + miss_policy=miss_policy, + modality_inputs=modalities or None, + modality_mask=(torch.ones(ids.shape[0]) if modalities else None), + ) + generated = output[:, ids.shape[1]:] + text = self.tokenizer.decode(generated[0].tolist(), skip_special_tokens=True) + stats = dict(getattr(self.engine, "last_stats", {}) or {}) + stats.setdefault("prompt_tokens", int(ids.shape[1])) + stats.setdefault("completion_tokens", int(generated.shape[1])) + return text, stats + + def openai_chat_completion(self, request: Mapping[str, Any]) -> dict: + messages = request.get("messages") or [] + text, stats = self.generate( + messages=messages, + max_tokens=request.get("max_tokens", 128), + temperature=request.get("temperature", 0.85), + top_p=request.get("top_p", 0.95), + top_k=request.get("top_k", 50), + do_sample=request.get("do_sample", True), + miss_policy=request.get("chronos_miss_policy", "on_demand"), + ) + return openai_chat_response( + str(request.get("model", self.model_name)), text, + prompt_tokens=stats.get("prompt_tokens", 0), + completion_tokens=stats.get("completion_tokens", 0), + chronos_metrics=stats if request.get("include_chronos_metrics") else None, + ) + + def ollama_chat(self, request: Mapping[str, Any]) -> dict: + messages = request.get("messages") or [] + options = request.get("options") or {} + text, stats = self.generate( + messages=messages, + max_tokens=request.get("num_predict", options.get("num_predict", 128)), + temperature=request.get("temperature", options.get("temperature", 0.85)), + top_p=request.get("top_p", options.get("top_p", 0.95)), + top_k=request.get("top_k", options.get("top_k", 50)), + do_sample=bool(request.get("do_sample", True)), + miss_policy=request.get("chronos_miss_policy", "on_demand"), + ) + return ollama_chat_response( + str(request.get("model", self.model_name)), text, + prompt_tokens=stats.get("prompt_tokens", 0), + completion_tokens=stats.get("completion_tokens", 0), + chronos_metrics=stats if request.get("include_chronos_metrics") else None, + ) + + def ollama_generate(self, request: Mapping[str, Any]) -> dict: + text, stats = self.generate( + prompt=request.get("prompt", ""), + max_tokens=request.get("options", {}).get("num_predict", 128), + temperature=request.get("options", {}).get("temperature", 0.85), + top_p=request.get("options", {}).get("top_p", 0.95), + top_k=request.get("options", {}).get("top_k", 50), + do_sample=bool(request.get("do_sample", True)), + miss_policy=request.get("chronos_miss_policy", "on_demand"), + ) + return ollama_generate_response( + str(request.get("model", self.model_name)), text, + prompt_tokens=stats.get("prompt_tokens", 0), + completion_tokens=stats.get("completion_tokens", 0), + chronos_metrics=stats if request.get("include_chronos_metrics") else None, + ) diff --git a/chronos/serving/api_compat.py b/chronos/serving/api_compat.py new file mode 100644 index 0000000..39a26df --- /dev/null +++ b/chronos/serving/api_compat.py @@ -0,0 +1,131 @@ +"""OpenAI/Ollama request and response compatibility helpers.""" +from __future__ import annotations + +import time +import uuid +import base64 +import io +import torch +from typing import Iterable, Optional + + +def message_text(messages: Iterable[dict]) -> str: + parts = [] + for message in messages or []: + content = message.get("content", "") + if isinstance(content, list): + content = "\n".join(str(item.get("text", "")) for item in content + if isinstance(item, dict) and item.get("type") in {"text", "input_text"}) + parts.append(f"{message.get('role', 'user')}: {content}") + return "\n".join(parts) + + +def modality_inputs_from_messages(messages: Iterable[dict]) -> dict: + modalities = {"image": [], "audio": []} + for message in messages or []: + content = message.get("content", []) + if not isinstance(content, list): + continue + for item in content: + if not isinstance(item, dict): + continue + kind = item.get("type") + if kind in {"image", "image_url", "input_image"}: + value = item.get("image") or item.get("image_url") + if isinstance(value, dict): + value = value.get("url") + if value: + modalities["image"].append(value) + elif kind in {"audio", "input_audio"}: + value = item.get("audio") or item.get("input_audio") + if isinstance(value, dict): + # OpenAI input_audio uses {data, format}; retain the + # format only as metadata and decode the base64 payload + # through the same WAV/PCM resolver as JSONL training. + value = value.get("data") or value.get("url") + if value and not str(value).startswith("data:"): + value = f"data:audio/wav;base64,{value}" + if value: + modalities["audio"].append(value) + return {key: value for key, value in modalities.items() if value} + + +def resolve_modality_inputs(messages: Iterable[dict], *, image_size: int = 224, + audio_sample_rate: int = 16000, + max_audio_seconds: float = 30.0) -> dict: + """Decode API URL/data-URI media into tensors accepted by ChronosFusion. + + Local paths and HTTP URLs are delegated to the dataset loaders so serving + and training share identical normalization. Invalid optional media is + omitted, allowing a text fallback rather than crashing a standard client. + """ + from chronos.data.flexible_dataset import _load_audio, _load_image + out = {} + raw = modality_inputs_from_messages(messages) + for kind, values in raw.items(): + decoded = [] + for value in values: + try: + if isinstance(value, str) and value.startswith("data:") and "," in value: + head, payload = value.split(",", 1) + blob = base64.b64decode(payload) + if kind == "image": + from PIL import Image + decoded.append(_load_image(Image.open(io.BytesIO(blob)), image_size)) + else: + # WAV is the portable baseline supported by the + # dependency-light training loader. Unsupported + # codecs fail closed and become a text fallback. + decoded.append(_load_audio(blob, audio_sample_rate, max_audio_seconds)) + elif kind == "image": + decoded.append(_load_image(value, image_size)) + else: + decoded.append(_load_audio(value, audio_sample_rate, max_audio_seconds)) + except Exception: + continue + if decoded: + if kind == "image": + # [B,N,C,H,W] lets PatchImageEncoder treat multiple images as + # one request; a single image remains semantically identical. + out[kind] = torch.stack(decoded).unsqueeze(0) + else: + # The baseline ConvAudioEncoder consumes one waveform per + # request. Preserve deterministic first-item semantics (the + # dataset path follows the same policy) instead of treating + # the number of clips as convolution channels. + out[kind] = decoded[0].flatten().unsqueeze(0) + return out + + +def openai_chat_response(model: str, text: str, *, prompt_tokens: int = 0, + completion_tokens: int = 0, chronos_metrics: Optional[dict] = None) -> dict: + result = { + "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion", + "created": int(time.time()), "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": int(prompt_tokens), "completion_tokens": int(completion_tokens), "total_tokens": int(prompt_tokens + completion_tokens)}, + } + if chronos_metrics is not None: + result["chronos_metrics"] = chronos_metrics + return result + + +def ollama_chat_response(model: str, text: str, *, done: bool = True, + prompt_tokens: int = 0, completion_tokens: int = 0, + chronos_metrics: Optional[dict] = None) -> dict: + result = { + "model": model, "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "message": {"role": "assistant", "content": text}, "done": bool(done), + "done_reason": "stop" if done else None, "prompt_eval_count": int(prompt_tokens), + "eval_count": int(completion_tokens), + } + if chronos_metrics is not None: + result["chronos_metrics"] = chronos_metrics + return result + + +def ollama_generate_response(model: str, text: str, **kwargs) -> dict: + result = ollama_chat_response(model, text, **kwargs) + result["response"] = text + result.pop("message", None) + return result diff --git a/chronos/serving/ollama_server.py b/chronos/serving/ollama_server.py new file mode 100644 index 0000000..8b2a8e5 --- /dev/null +++ b/chronos/serving/ollama_server.py @@ -0,0 +1,353 @@ +"""Ollama- and OpenAI-compatible HTTP server for Chronos models. + +The stock ``ollama`` CLI and any Ollama/OpenAI client can talk to the native +Chronos runtime by pointing at this server:: + + chronos serve --models-dir runs/enzh_50m_v1 --port 11435 + OLLAMA_HOST=127.0.0.1:11435 ollama run chronos_enzh_sft "Hello world" + +The kernel is never converted to another architecture: recurrent depth, +SSD-lazy experts with lookahead prefetch, compressed KV and self-speculative +decoding all stay active behind the protocol. + +Security: binds 127.0.0.1 and has no authentication by default. Binding a +non-loopback address requires ``--allow-remote`` *and* ``--api-key``. +``--monitor-runs-dir`` mounts the read-only training monitor at ``/monitor/``; +it reuses the same key through its own login page (session cookie). +""" +from __future__ import annotations + +import hashlib +import hmac +import json +import os +import time +from datetime import datetime, timezone + +try: # module level: FastAPI resolves ``Request`` annotations by name + from fastapi import FastAPI, HTTPException, Request + from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse +except ImportError: # pragma: no cover - serving extra not installed + FastAPI = HTTPException = Request = JSONResponse = PlainTextResponse = StreamingResponse = None + +OLLAMA_COMPAT_VERSION = "0.40.14" + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + + +def _digest(path: str) -> str: + st = os.stat(path) + return hashlib.sha256(f"{path}:{st.st_size}:{st.st_mtime_ns}".encode()).hexdigest() + + +def _options(body: dict) -> dict: + opts = dict(body.get("options") or {}) + return { + "max_new_tokens": int(opts.get("num_predict", body.get("max_tokens", 256)) or 256), + "temperature": float(opts.get("temperature", body.get("temperature", 0.0)) or 0.0), + "top_k": int(opts.get("top_k", body.get("top_k", 0)) or 0), + "top_p": float(opts.get("top_p", body.get("top_p", 1.0)) or 1.0), + "seed": opts.get("seed", body.get("seed")), + "stop": opts.get("stop", body.get("stop")) or [], + "speculative": bool(opts.get("chronos_speculative", True)), + "draft_depth": int(opts.get("chronos_draft_depth", 1)), + "gamma": int(opts.get("chronos_gamma", 4)), + "miss_policy": str(opts.get("chronos_miss_policy", "on_demand")), + } + + +def _details(lm_or_cfg) -> dict: + cfg = getattr(lm_or_cfg, "config", lm_or_cfg) + return {"parent_model": "", "format": "chronos", "family": "chronos", "families": ["chronos"], + "parameter_size": f"{getattr(cfg, '_total_params_m', 0) or 0:.1f}M" if hasattr(cfg, "_total_params_m") else "", + "quantization_level": "BF16"} + + +def _final_ollama(stats: dict, model: str, total_t0: float, load_s: float) -> dict: + return { + "model": model, "created_at": _now(), "done": True, + "done_reason": stats.get("done_reason", "stop"), + "total_duration": int((time.perf_counter() - total_t0) * 1e9), + "load_duration": int(load_s * 1e9), + "prompt_eval_count": stats.get("prompt_tokens", 0), "prompt_eval_duration": 0, + "eval_count": stats.get("completion_tokens", 0), + "eval_duration": stats.get("eval_duration_ns", 0), + "chronos": {k: v for k, v in stats.items() if k not in ("done", "done_reason")}, + } + + +def create_app(runtime, *, api_key: str | None = None, monitor_runs_dir: str | None = None): + if FastAPI is None: + raise ImportError("chronos serve needs fastapi and uvicorn (pip install 'Project_Chronos[serve]')") + app = FastAPI(title="Chronos Ollama-compatible server", version=OLLAMA_COMPAT_VERSION) + expected = f"Bearer {api_key}".encode("utf-8") if api_key else b"" + + @app.middleware("http") + async def auth(request: Request, call_next): + path = request.url.path + # /monitor has its own login (cookie or Bearer) with the same key. + exempt = path in ("/", "/api/version") or (monitor_runs_dir and (path == "/monitor" or path.startswith("/monitor/"))) + if api_key and not exempt: + given = request.headers.get("authorization", "").encode("utf-8") + if not hmac.compare_digest(given, expected): + return JSONResponse({"error": "unauthorized"}, status_code=401) + return await call_next(request) + + if monitor_runs_dir: + from fastapi.responses import RedirectResponse + from chronos.monitor.server import create_monitor_app + + @app.get("/monitor", include_in_schema=False) + def monitor_redirect(): + return RedirectResponse("/monitor/") + + app.mount("/monitor", create_monitor_app(monitor_runs_dir, api_key=api_key)) + + def _resolve(name: str): + try: + t0 = time.perf_counter() + lm = runtime.get(name) + return lm, time.perf_counter() - t0 + except KeyError as exc: + raise HTTPException(status_code=404, detail=f"model '{name}' not found, try pulling it first") from exc + + def _model_entry(name, path): + st = os.stat(path) + return {"name": f"{name}:latest", "model": f"{name}:latest", + "modified_at": datetime.fromtimestamp(st.st_mtime, timezone.utc).isoformat().replace("+00:00", "Z"), + "size": st.st_size, "digest": _digest(path), + "details": {"parent_model": "", "format": "chronos", "family": "chronos", + "families": ["chronos"], "parameter_size": "", "quantization_level": "BF16"}} + + @app.get("/") + def root(): + return PlainTextResponse("Ollama is running") + + @app.head("/") + def root_head(): + return PlainTextResponse("") + + @app.get("/api/version") + def version(): + return {"version": OLLAMA_COMPAT_VERSION} + + @app.get("/api/tags") + def tags(): + return {"models": [_model_entry(n, p) for n, p in sorted(runtime.available().items())]} + + @app.post("/api/show") + async def show(request: Request): + body = await request.json() + name = (body.get("model") or body.get("name") or "").split(":")[0] + path = runtime.available().get(name) + if path is None: + raise HTTPException(status_code=404, detail=f"model '{name}' not found") + meta = json.load(open(os.path.splitext(path)[0] + ".config.json", encoding="utf-8")) + cfg = meta.get("config", {}) + info = {"general.architecture": "chronos", "chronos.recurrent_depth": cfg.get("recurrent_depth"), + "chronos.recurrent_attention_mode": cfg.get("recurrent_attention_mode"), + "chronos.num_experts": cfg.get("num_experts"), "chronos.experts_per_token": cfg.get("num_experts_per_tok"), + "chronos.kv_latent_dim": cfg.get("kv_latent_dim"), "chronos.context_length": cfg.get("max_position_embeddings"), + "chronos.hidden_size": cfg.get("hidden_size"), "chronos.vocab_size": cfg.get("vocab_size")} + return {"modelfile": f"FROM {path}\n", "parameters": "temperature 0\n", "template": "{{ .Prompt }}", + "details": _model_entry(name, path)["details"], "model_info": info, + "capabilities": ["completion"], "modified_at": _model_entry(name, path)["modified_at"]} + + @app.get("/api/ps") + def ps(): + out = [] + for name, lm in list(runtime.loaded.items()): + store = getattr(getattr(lm.engine, "cache_manager", None), "expert_store", None) + st = store.stats() if store is not None else {} + entry = _model_entry(name, lm.path) + entry.update(size_vram=int(st.get("vram_bytes", 0) or 0), + expires_at=datetime.fromtimestamp(runtime.last_used.get(name, time.time()) + runtime.keep_alive_s, + timezone.utc).isoformat().replace("+00:00", "Z"), + chronos={"device": lm.device, "lazy_experts": lm.lazy_experts, "expert_store": st}) + out.append(entry) + return {"models": out} + + @app.post("/api/pull") + async def pull(request: Request): + body = await request.json() + name = (body.get("model") or body.get("name") or "").split(":")[0] + if name not in runtime.available(): + return JSONResponse({"error": f"model '{name}' is not in --models-dir; Chronos serves local checkpoints"}, + status_code=404) + return {"status": "success"} + + def _ndjson(gen): + for item in gen: + yield json.dumps(item, ensure_ascii=False) + "\n" + + def _run(body, input_ids, lm, load_s, shape): + opts = _options(body) + t0 = time.perf_counter() + name = lm.name + def events(): + for item in runtime.generate_stream(name, input_ids, **opts): + if item.get("done"): + final = _final_ollama(item, body.get("model", name), t0, load_s) + if "error" in item: + final["error"] = item["error"] + yield shape(None, final) + else: + yield shape(item["text"], None) + return events + + def _respond(body, events, collect): + if body.get("stream", True): + return StreamingResponse(_ndjson(events()), media_type="application/x-ndjson") + return JSONResponse(collect(list(events()))) + + def _unload_if_requested(body, name) -> bool: + keep = body.get("keep_alive") + if keep in (0, "0", "0s", "0m") and not body.get("prompt") and not body.get("messages"): + runtime.unload(name) + return True + return False + + @app.post("/api/generate") + async def generate(request: Request): + body = await request.json() + name = (body.get("model") or "").split(":")[0] + if _unload_if_requested(body, name): + return {"model": body.get("model"), "created_at": _now(), "response": "", "done": True, + "done_reason": "unload"} + lm, load_s = _resolve(name) + if not body.get("prompt"): # Ollama "load model" request + return {"model": body.get("model"), "created_at": _now(), "response": "", "done": True, + "done_reason": "load"} + ids = runtime.encode_prompt(lm, body["prompt"], body.get("system"), bool(body.get("raw"))) + + def shape(text, final): + if final is not None: + return {**final, "response": ""} + return {"model": body.get("model"), "created_at": _now(), "response": text, "done": False} + + def collect(items): + return {**items[-1], "response": "".join(i.get("response", "") for i in items[:-1])} + return _respond(body, _run(body, ids, lm, load_s, shape), collect) + + @app.post("/api/chat") + async def chat(request: Request): + body = await request.json() + name = (body.get("model") or "").split(":")[0] + messages = [{"role": m.get("role", "user"), "content": m.get("content", "")} + for m in body.get("messages") or []] + if _unload_if_requested(body, name): + return {"model": body.get("model"), "created_at": _now(), + "message": {"role": "assistant", "content": ""}, "done": True, "done_reason": "unload"} + lm, load_s = _resolve(name) + if not messages: + return {"model": body.get("model"), "created_at": _now(), + "message": {"role": "assistant", "content": ""}, "done": True, "done_reason": "load"} + ids = runtime.encode_chat(lm, messages) + + def shape(text, final): + if final is not None: + return {**final, "message": {"role": "assistant", "content": ""}} + return {"model": body.get("model"), "created_at": _now(), + "message": {"role": "assistant", "content": text}, "done": False} + + def collect(items): + text = "".join(i["message"]["content"] for i in items[:-1]) + return {**items[-1], "message": {"role": "assistant", "content": text}} + return _respond(body, _run(body, ids, lm, load_s, shape), collect) + + @app.get("/v1/models") + def openai_models(): + return {"object": "list", "data": [{"id": n, "object": "model", "created": int(os.stat(p).st_mtime), + "owned_by": "chronos"} for n, p in sorted(runtime.available().items())]} + + @app.post("/v1/chat/completions") + async def openai_chat(request: Request): + import uuid + body = await request.json() + name = (body.get("model") or "").split(":")[0] + lm, load_s = _resolve(name) + messages = [{"role": m.get("role", "user"), "content": m.get("content", "")} for m in body.get("messages") or []] + ids = runtime.encode_chat(lm, messages) + opts = _options({**body, "options": {"num_predict": body.get("max_tokens") or body.get("max_completion_tokens") or 256, + "temperature": body.get("temperature", 0.0), "top_p": body.get("top_p", 1.0), + "seed": body.get("seed"), "stop": body.get("stop"), + **(body.get("chronos") or {})}}) + cid, created = f"chatcmpl-{uuid.uuid4().hex}", int(time.time()) + if body.get("stream"): + def sse(): + for item in runtime.generate_stream(name, ids, **opts): + if item.get("done"): + chunk = {"id": cid, "object": "chat.completion.chunk", "created": created, "model": name, + "choices": [{"index": 0, "delta": {}, "finish_reason": item.get("done_reason", "stop")}], + "chronos": item} + yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n" + else: + chunk = {"id": cid, "object": "chat.completion.chunk", "created": created, "model": name, + "choices": [{"index": 0, "delta": {"content": item["text"]}, "finish_reason": None}]} + yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" + return StreamingResponse(sse(), media_type="text/event-stream") + items = list(runtime.generate_stream(name, ids, **opts)) + final = items[-1] + return {"id": cid, "object": "chat.completion", "created": created, "model": name, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(i["text"] for i in items[:-1])}, + "finish_reason": final.get("done_reason", "stop")}], + "usage": {"prompt_tokens": final.get("prompt_tokens", 0), "completion_tokens": final.get("completion_tokens", 0), + "total_tokens": final.get("prompt_tokens", 0) + final.get("completion_tokens", 0)}, + "chronos": final} + + return app + + +def main(argv=None) -> int: + import argparse + parser = argparse.ArgumentParser(prog="chronos serve", description=__doc__.split("\n\n")[0]) + parser.add_argument("--models-dir", required=True, help="directory searched for .pth + .config.json") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=11435, help="11434 would collide with a running Ollama") + parser.add_argument("--device", default="auto") + parser.add_argument("--dtype", default="auto", help="compute dtype: auto|float32|bfloat16|float16 " + "(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) + # 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") + parser.add_argument("--mla-decode", choices=["auto", "absorbed", "expanded"], default=None) + parser.add_argument("--mla-cache-rope", action="store_true", default=None, + help="also cache post-RoPE rope keys (memory for throughput)") + parser.add_argument("--router-scope", choices=["per_iteration", "shared"], default=None, + help="only change this for checkpoints trained with the same scope") + parser.add_argument("--max-resident-experts", type=int, default=None, + help="hard cap on fast-tier expert pages") + parser.add_argument("--ram-budget-mb", type=float, default=None, help="hard cap on RAM expert pages (MiB)") + parser.add_argument("--api-key", default=os.environ.get("CHRONOS_API_KEY")) + parser.add_argument("--allow-remote", action="store_true") + parser.add_argument("--monitor-runs-dir", default=None, + help="also serve the read-only training monitor for this runs directory at /monitor/") + args = parser.parse_args(argv) + if args.host not in ("127.0.0.1", "localhost", "::1"): + if not (args.allow_remote and args.api_key): + parser.error("binding a non-loopback host requires --allow-remote and --api-key") + import uvicorn + from chronos.serving.runtime import ChronosRuntime + overrides = {"kv_cache_dtype": args.kv_cache_dtype, "mla_decode": args.mla_decode, + "mla_cache_rope": args.mla_cache_rope, "router_scope": args.router_scope, + "max_resident_experts": args.max_resident_experts, "ram_budget_mb": args.ram_budget_mb} + 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) + 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: + print(f"[chronos serve] monitor: http://{args.host}:{args.port}/monitor/ " + f"runs={os.path.abspath(args.monitor_runs_dir)}", flush=True) + app = create_app(runtime, api_key=args.api_key, monitor_runs_dir=args.monitor_runs_dir) + uvicorn.run(app, host=args.host, port=args.port, log_level="warning") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/serving/runtime.py b/chronos/serving/runtime.py new file mode 100644 index 0000000..b765fc4 --- /dev/null +++ b/chronos/serving/runtime.py @@ -0,0 +1,343 @@ +"""Chronos serving runtime: native kernel + SSD-lazy experts + self-speculation. + +This is the engine behind the Ollama/OpenAI protocol layer. It keeps every +Chronos property — recurrent depth, shared_first KV, SSD→RAM→device expert +offload with lookahead prefetch, and recurrent-depth self-speculative +decoding — so a client talking the Ollama protocol gets the same memory and +throughput behaviour as native use. +""" +from __future__ import annotations + +import gc +import json +import os +import threading +import time +from dataclasses import dataclass, field + +import torch + + +def _rss_mb() -> float: + try: + import psutil + return psutil.Process().memory_info().rss / 2**20 + except Exception: + return 0.0 + + +@dataclass +class LoadedModel: + name: str + path: str + config: object + model: object + tokenizer: object + engine: object | None + device: str + lazy_experts: bool + loaded_at: float = field(default_factory=time.time) + meta: dict = field(default_factory=dict) + + +def _pick_device(requested: str | None) -> str: + if requested and requested != "auto": + return requested + if torch.backends.mps.is_available(): + return "mps" + if torch.cuda.is_available(): + return "cuda" + return "cpu" + + +# Inference-only options a server may override on a loaded checkpoint. +# None of them changes the trained weights; all are validated by ChronosConfig. +INFERENCE_OVERRIDE_KEYS = ("kv_cache_dtype", "kv_quant_group_size", "mla_decode", "mla_cache_rope", + "router_scope", "max_resident_experts", "ram_budget_mb") + + +def apply_inference_overrides(config, overrides: dict | None) -> dict: + """Set validated inference options on ``config``; return what was applied.""" + from chronos.model.config import ChronosConfig + applied = {k: v for k, v in (overrides or {}).items() if k in INFERENCE_OVERRIDE_KEYS and v is not None} + unknown = set(overrides or {}) - set(INFERENCE_OVERRIDE_KEYS) + if unknown: + raise ValueError(f"not an inference option: {sorted(unknown)}") + if applied: + for key in ("max_resident_experts", "ram_budget_mb"): + if key in applied and float(applied[key]) <= 0: + raise ValueError(f"{key} must be positive") + # Validate against the checkpoint's full config: a partial probe falls + # back to class defaults (e.g. 4 KV heads) that may not fit the model. + probe = {k: v for k, v in config.to_dict().items() if k not in applied} + for key in ("transformers_version", "_name_or_path", "_commit_hash", "architectures", "model_type"): + probe.pop(key, None) + checked = ChronosConfig(**probe, **applied) # raises on an invalid value, normalises aliases + for key in applied: + setattr(config, key, getattr(checked, key)) + applied[key] = getattr(checked, key) + return applied + + +def _find_tokenizer_dir(ckpt_path: str, meta: dict) -> str: + from chronos.model.checkpoint import resolve_checkpoint_tokenizer + return resolve_checkpoint_tokenizer(ckpt_path, fallback=False) + + +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: + """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`` + picks the single compute dtype (``chronos.precision.resolve_torch``): + weights, activations, KV cache and expert pages all use it. + """ + from transformers import PreTrainedTokenizerFast + from chronos.model.checkpoint import ( + chronos_config_from_checkpoint, load_checkpoint_state_dict, load_state_dict_controlled, + read_checkpoint_config, + ) + from chronos.model.model_chronos import ChronosForCausalLM + from chronos.precision import place_torch_model + meta, _ = read_checkpoint_config(path) + meta_full = {} + cfg_path = os.path.splitext(path)[0] + ".config.json" + if os.path.exists(cfg_path): + meta_full = json.load(open(cfg_path, encoding="utf-8")) + config, _sources = chronos_config_from_checkpoint(path) + applied = apply_inference_overrides(config, inference_overrides) + tok_dir = _find_tokenizer_dir(path, meta_full) + tokenizer = PreTrainedTokenizerFast.from_pretrained(tok_dir) + device = _pick_device(device) + weights = load_checkpoint_state_dict(path, map_location="cpu") + model = ChronosForCausalLM(config) + engine = None + 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, + ) + 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) + model, precision = place_torch_model(model, device, dtype) + model = model.eval() + cache_dir = expert_cache_root or _checkpoint_expert_cache_dir(path, config) + 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=[]) + else: + load_state_dict_controlled(model, weights) + model, precision = place_torch_model(model, device, dtype) + model = model.eval() + del weights + gc.collect() + return LoadedModel(name=name or os.path.splitext(os.path.basename(path))[0], path=os.path.abspath(path), + 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}) + + +class ChronosRuntime: + """Thread-safe model registry + generation for the protocol server.""" + + 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): + 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 + self.keep_alive_s = keep_alive_s + self.loaded: dict[str, LoadedModel] = {} + self.last_used: dict[str, float] = {} + self.lock = threading.RLock() + + # ---- registry ------------------------------------------------------------ + def available(self) -> dict[str, str]: + """name → checkpoint path for every ``*.pth`` with a config sidecar.""" + out = {} + if self.models_dir and os.path.isdir(self.models_dir): + for root, _dirs, files in os.walk(self.models_dir): + for f in files: + if f.endswith(".pth") and os.path.exists(os.path.join(root, f[:-4] + ".config.json")): + out[f[:-4]] = os.path.join(root, f) + for name, lm in self.loaded.items(): + out.setdefault(name, lm.path) + return out + + def get(self, name: str) -> LoadedModel: + with self.lock: + name = name.split(":")[0] + if name not in self.loaded: + path = self.available().get(name) + if path is None: + raise KeyError(f"model '{name}' not found") + self.loaded[name] = load_model(path, name=name, device=self.device, + lazy_experts=self.lazy_experts, + ram_expert_fraction=self.ram_expert_fraction, + dtype=self.dtype, + inference_overrides=self.inference_overrides) + self.last_used[name] = time.time() + return self.loaded[name] + + def unload(self, name: str) -> bool: + with self.lock: + lm = self.loaded.pop(name, None) + if lm and lm.engine is not None: + try: + lm.engine.teardown() + except Exception: + pass + gc.collect() + if torch.backends.mps.is_available(): + torch.mps.empty_cache() + return lm is not None + + # ---- generation ------------------------------------------------------------ + def _forward_fn(self, lm: LoadedModel, telemetry: dict): + engine = lm.engine + if engine is None: + def fwd(ids, cache, depth): + return lm.model(ids, past_key_values=cache, use_cache=True, recurrent_depth=depth) + return fwd + cm = engine.cache_manager + state = {"lookahead": None} + + def fwd(ids, cache, depth): + # Lookahead-guided residency: prefetch next experts to RAM, promote + # the predicted set, then run with the current availability masks. + if state["lookahead"] is not None: + predicted = engine._predict_expert_ids(state["lookahead"], all_future=True, + capacity=int(cm.expert_store.ram_capacity)) + cm.prefetch_experts_to_ram(predicted) + ready = engine._promotion_ready_ids(predicted[: int(cm.expert_store.vram_capacity)]) + if ready: + cm.ensure_resident(ready) + masks = cm.availability_masks_all_layers() + out = lm.model(ids.to(lm.device), past_key_values=cache, use_cache=True, + recurrent_depth=depth, available_expert_masks=masks) + state["lookahead"] = out.lookahead_probs + rt = engine._consume_moe_runtime_stats() + telemetry["fallback_weight_mass"] += float(rt["fallback_executed_weight_mass"]) + telemetry["on_demand_loads"] += len(rt["on_demand_experts"]) + telemetry["forwards"] += 1 + return out + return fwd + + def encode_chat(self, lm: LoadedModel, messages: list[dict]) -> list[int]: + text = lm.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + return lm.tokenizer(text, add_special_tokens=False)["input_ids"] + + def encode_prompt(self, lm: LoadedModel, prompt: str, system: str | None = None, raw: bool = False): + if raw: + return lm.tokenizer(prompt, add_special_tokens=False)["input_ids"] + msgs = ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": prompt}] + return self.encode_chat(lm, msgs) + + def generate_stream(self, name: str, input_ids: list[int], *, max_new_tokens: int = 256, + temperature: float = 0.0, top_k: int = 0, top_p: float = 1.0, seed: int | None = None, + speculative: bool = True, draft_depth: int = 1, gamma: int = 4, + miss_policy: str = "on_demand", stop: list[str] | None = None): + """Yield ``{"text": delta}`` chunks, then a final ``{"done": True, ...stats}``.""" + from chronos.runtime.speculative import speculative_generate, speculative_supported + lm = self.get(name) + tok = lm.tokenizer + eos = [i for i in (tok.convert_tokens_to_ids("<|im_end|>"), tok.convert_tokens_to_ids("<|endoftext|>")) + if isinstance(i, int) and i >= 0] + telemetry = {"fallback_weight_mass": 0.0, "on_demand_loads": 0, "forwards": 0} + import queue + chunks: queue.Queue = queue.Queue() + cancel = threading.Event() + result = {} + ids = torch.tensor([input_ids], device=lm.device) + do_sample = float(temperature) > 0 + gen = torch.Generator(device="cpu").manual_seed(int(seed)) if seed is not None else None + depth = int(getattr(lm.config, "max_inference_depth", None) or lm.config.recurrent_depth) + use_spec = bool(speculative) and speculative_supported(lm.config) and int(draft_depth) < depth + + def worker(): + with self.lock: # one generation at a time per process (single GPU) + rss0 = _rss_mb() + if lm.engine is not None: + lm.engine._runtime_stats = {"on_demand_loads": 0, "on_demand_load_time_s": 0.0, + "sync_ssd_loads": 0, "async_cold_miss_prefetches": 0, + "resident_vram_hits": 0, "resident_ram_hits": 0, + "on_demand_loaded_experts": [], "on_demand_promoted_experts": []} + lm.engine._install_moe_runtime_hooks(lm.engine._normalize_miss_policy(miss_policy)) + generated: list[int] = [] + emitted_text = "" + t0 = time.perf_counter() + + def on_token(t): + nonlocal emitted_text + if t in eos: + return + generated.append(t) + text = tok.decode(generated, skip_special_tokens=True) + if text.endswith("�"): + return # wait for the rest of a multi-byte char + delta, emitted_text = text[len(emitted_text):], text + if delta: + chunks.put({"text": delta}) + try: + _, stats = speculative_generate( + lm.model, ids, max_new_tokens=max_new_tokens, + draft_depth=draft_depth if use_spec else depth, gamma=gamma, + do_sample=do_sample, temperature=max(float(temperature), 1e-6), + top_k=int(top_k or 0), top_p=float(top_p or 1.0), eos_token_id=eos, + generator=gen, on_token=on_token, forward_fn=self._forward_fn(lm, telemetry), + should_stop=cancel.is_set) + result.update(stats.as_dict()) + except Exception as exc: # surfaced to the client as an error chunk + result["error"] = f"{type(exc).__name__}: {exc}" + finally: + if lm.engine is not None: + lm.engine._clear_moe_runtime_hooks() + elapsed = time.perf_counter() - t0 + store = getattr(getattr(lm.engine, "cache_manager", None), "expert_store", None) + result.update({ + "prompt_tokens": len(input_ids), "completion_tokens": len(generated), + "eval_duration_ns": int(elapsed * 1e9), "speculative": use_spec, + "draft_depth": draft_depth if use_spec else depth, "full_depth": depth, + "lazy_experts": lm.lazy_experts, "rss_mb": round(_rss_mb(), 1), + "rss_delta_mb": round(_rss_mb() - rss0, 1), **telemetry, + **({"expert_store": store.stats()} if store is not None else {}), + }) + chunks.put(None) + + threading.Thread(target=worker, daemon=True).start() + stop = [s for s in (stop or []) if s] + produced = "" + finished = False # True once the worker's ``None`` sentinel was consumed + try: + while True: + item = chunks.get() + if item is None: + finished = True + break + text = item["text"] + if stop: + produced += text + hits = [produced.find(s) for s in stop if s in produced] + if hits: + text = text[: max(0, len(text) - (len(produced) - min(hits)))] + if text: + yield {"text": text} + result["done_reason"] = "stop" + cancel.set() + break + yield {"text": text} + finally: + cancel.set() # client went away or stop string hit + while not finished: # drain until the worker has released the model + finished = chunks.get() is None + result.setdefault("done_reason", "length" if result.get("completion_tokens", 0) >= max_new_tokens else "stop") + yield {"done": True, **result} diff --git a/chronos/tokenizer/__init__.py b/chronos/tokenizer/__init__.py new file mode 100644 index 0000000..06d2468 --- /dev/null +++ b/chronos/tokenizer/__init__.py @@ -0,0 +1,75 @@ +"""Chronos Tokenizer Builder. + +``build_tokenizer(paths, out_dir, vocab_size)`` extracts text from any user- +selected datasets, trains the Chronos BPE (script-aware pre-tokenization, +tempered language balance, byte fallback), evaluates it and writes both the +native ``chronos_tokenizer.json`` and an HF-compatible ``tokenizer.json``. +""" +from __future__ import annotations + +import json +import os +import time + +from .evaluate import evaluate_tokenizer +from .extract import ExtractConfig, extract_texts +from .model import DEFAULT_SPECIAL_TOKENS, ChronosTokenizer +from .trainer import BPETrainerConfig, train_bpe + +__all__ = ["build_tokenizer", "ChronosTokenizer", "BPETrainerConfig", "ExtractConfig"] + + +def build_tokenizer(paths: list[str], out_dir: str, *, vocab_size: int = 16384, + extract: ExtractConfig | None = None, trainer: BPETrainerConfig | None = None, + special_tokens=DEFAULT_SPECIAL_TOKENS, heldout_fraction: float = 0.02, + log=print) -> ChronosTokenizer: + t0 = time.time() + extract = extract or ExtractConfig() + result = extract_texts(paths, extract, log=log) + train_split, heldout = {}, {} + for bucket, texts in result.texts_by_bucket.items(): + cut = max(1, int(len(texts) * heldout_fraction)) + heldout[bucket], train_split[bucket] = texts[:cut], texts[cut:] + config = trainer or BPETrainerConfig(vocab_size=vocab_size, special_tokens=tuple(special_tokens)) + config.vocab_size = vocab_size + config.special_tokens = tuple(special_tokens) + log(f"[tokenizer] training vocab={vocab_size} on " + f"{ {b: len(t) for b, t in train_split.items()} } texts ({time.time() - t0:.0f}s)") + trained = train_bpe(train_split, config, + progress=lambda n, total: log(f"[tokenizer] vocab {n}/{total}") if n % 2000 < 4 else None) + tok = ChronosTokenizer.from_trained(trained, special_tokens) + report = evaluate_tokenizer(tok, heldout) + tok.metadata.update({ + "trainer": {k: (list(v) if isinstance(v, tuple) else v) for k, v in vars(config).items() + if k != "special_tokens"}, + "train_stats": trained.stats, + "heldout_eval": report, + "sources": result.per_source, + "elapsed_s": round(time.time() - t0, 1), + }) + tok.save(out_dir) + log(f"[tokenizer] saved {out_dir}: {json.dumps(report, ensure_ascii=False)}") + return tok + + +def main(argv=None) -> int: + import argparse + parser = argparse.ArgumentParser(description="Train a Chronos tokenizer from selected datasets") + parser.add_argument("paths", nargs="+") + parser.add_argument("--out", required=True) + parser.add_argument("--vocab-size", type=int, default=16384) + parser.add_argument("--per-source", type=int, default=400_000) + parser.add_argument("--fields", nargs="*", default=None) + parser.add_argument("--temperature", type=float, default=0.5) + args = parser.parse_args(argv) + build_tokenizer( + args.paths, args.out, vocab_size=args.vocab_size, + extract=ExtractConfig(per_source_texts=args.per_source, fields=args.fields), + trainer=BPETrainerConfig(vocab_size=args.vocab_size, language_temperature=args.temperature), + log=lambda m: print(m, flush=True), + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/tokenizer/__main__.py b/chronos/tokenizer/__main__.py new file mode 100644 index 0000000..3fca214 --- /dev/null +++ b/chronos/tokenizer/__main__.py @@ -0,0 +1,4 @@ +"""``python -m chronos.tokenizer`` entry point.""" +from chronos.tokenizer import main + +raise SystemExit(main()) diff --git a/chronos/tokenizer/evaluate.py b/chronos/tokenizer/evaluate.py new file mode 100644 index 0000000..fb30662 --- /dev/null +++ b/chronos/tokenizer/evaluate.py @@ -0,0 +1,40 @@ +"""Tokenizer quality metrics: compression, fertility, fallback, round-trip.""" +from __future__ import annotations + +import re + +_HAN = re.compile(r"[㐀-䶿一-鿿豈-﫿]") + + +def evaluate_tokenizer(tokenizer, texts_by_bucket: dict[str, list[str]], limit: int = 5000) -> dict: + from .model import tokenizer_normalize + byte_ids = set(tokenizer._byte_ids) + used = set() + report = {} + for bucket, texts in texts_by_bucket.items(): + sample = texts[:limit] + if not sample: + continue + tokens = byte_tokens = chars = bytes_ = words = han = roundtrip_fail = 0 + for text in sample: + ids = tokenizer.encode(text, allow_special=False) + used.update(ids) + tokens += len(ids) + byte_tokens += sum(1 for i in ids if i in byte_ids) + chars += len(text) + bytes_ += len(text.encode("utf-8")) + words += len(text.split()) + han += len(_HAN.findall(text)) + if tokenizer.decode(ids) != tokenizer_normalize(text): + roundtrip_fail += 1 + report[bucket] = { + "texts": len(sample), + "bytes_per_token": round(bytes_ / max(tokens, 1), 3), + "chars_per_token": round(chars / max(tokens, 1), 3), + "tokens_per_word": round(tokens / max(words, 1), 3), + "tokens_per_han": round(tokens / han, 3) if han else None, + "byte_fallback_rate": round(byte_tokens / max(tokens, 1), 5), + "roundtrip_failures": roundtrip_fail, + } + report["vocab_utilization"] = round(len(used) / max(tokenizer.vocab_size, 1), 4) + return report diff --git a/chronos/tokenizer/extract.py b/chronos/tokenizer/extract.py new file mode 100644 index 0000000..eb56b77 --- /dev/null +++ b/chronos/tokenizer/extract.py @@ -0,0 +1,176 @@ +"""Automatic text extraction for tokenizer training from user-selected datasets. + +Accepts any mix of files/directories (parquet, jsonl, json, csv, tsv, txt). +For each record it finds text fields automatically (conversations, pairs, +``text``/``content``/``translation.*`` …) unless a field list is given, buckets +text by script, and keeps a seeded per-source reservoir so every source is +represented regardless of size. +""" +from __future__ import annotations + +import csv +import json +import os +import random +import re +from dataclasses import dataclass, field +from typing import Iterator + +_HAN = re.compile(r"[㐀-䶿一-鿿豈-﫿]") +_KANA = re.compile(r"[぀-ヿ]") +_HANGUL = re.compile(r"[가-힯]") +_LATIN = re.compile(r"[A-Za-zÀ-ɏ]") +_CYRILLIC = re.compile(r"[Ѐ-ӿ]") + +TEXT_KEYS = ("text", "content", "sentence", "source", "target", "prompt", "response", + "chosen", "rejected", "question", "answer", "instruction", "input", "output") +SUFFIXES = (".parquet", ".jsonl", ".json", ".csv", ".tsv", ".txt") + + +def script_bucket(text: str) -> str: + counts = {"zh": len(_HAN.findall(text)), "ja": len(_KANA.findall(text)) * 2, + "ko": len(_HANGUL.findall(text)), "latin": len(_LATIN.findall(text)), + "cyrillic": len(_CYRILLIC.findall(text))} + best = max(counts, key=counts.get) + return best if counts[best] else "other" + + +def _texts_from_record(record, fields: list[str] | None) -> Iterator[str]: + if fields: + for f in fields: + value = record + for part in f.split("."): + value = value.get(part) if isinstance(value, dict) else None + if isinstance(value, str): + yield value + return + if isinstance(record, str): + yield record + return + if not isinstance(record, dict): + return + for key in ("conversations", "messages"): + msgs = record.get(key) + if isinstance(msgs, list): + for m in msgs: + if isinstance(m, dict) and isinstance(m.get("content"), str): + yield m["content"] + return + for key, value in record.items(): + if isinstance(value, str) and (key in TEXT_KEYS or len(key) <= 3 or key.startswith("sentence")): + yield value + elif isinstance(value, dict): # translation: {en: ..., zh: ...} + for sub in value.values(): + if isinstance(sub, str): + yield sub + + +def iter_records(path: str) -> Iterator[object]: + if path.endswith(".parquet"): + import pyarrow.parquet as pq + for batch in pq.ParquetFile(path).iter_batches(batch_size=8192): + yield from batch.to_pylist() + elif path.endswith(".jsonl"): + with open(path, encoding="utf-8", errors="replace") as fh: + for line in fh: + line = line.strip() + if line: + try: + yield json.loads(line) + except json.JSONDecodeError: + continue + elif path.endswith(".json"): + with open(path, encoding="utf-8", errors="replace") as fh: + data = json.load(fh) + yield from (data if isinstance(data, list) else [data]) + elif path.endswith((".csv", ".tsv")): + csv.field_size_limit(1 << 30) + with open(path, newline="", encoding="utf-8", errors="replace") as fh: + for row in csv.DictReader(fh, delimiter="\t" if path.endswith(".tsv") else ","): + yield {k if k and not k.isdigit() else f"c{k}": v for k, v in row.items()} + elif path.endswith(".txt"): + with open(path, encoding="utf-8", errors="replace") as fh: + for line in fh: + if line.strip(): + yield line.rstrip("\n") + + +def expand(paths: list[str]) -> list[str]: + import glob + out = [] + for p in paths: + if any(ch in p for ch in "*?["): + out.extend(f for f in glob.glob(p) if f.endswith(SUFFIXES)) + elif os.path.isdir(p): + for root, _d, files in os.walk(p): + out.extend(os.path.join(root, f) for f in sorted(files) if f.endswith(SUFFIXES)) + elif p.endswith(SUFFIXES): + out.append(p) + return sorted(out) + + +@dataclass +class ExtractConfig: + per_source_texts: int = 400_000 # reservoir size per selected path + max_records_per_source: int | None = 4_000_000 + fields: list[str] | None = None # explicit dotted field list, else auto + min_chars: int = 2 + max_chars: int = 2000 + seed: int = 1234 + exclude_substrings: tuple[str, ...] = ("/valid.", "/test") + + +@dataclass +class ExtractResult: + texts_by_bucket: dict[str, list[str]] + per_source: dict[str, dict] = field(default_factory=dict) + + +def _capped(iterator, limit): + for i, item in enumerate(iterator): + if limit is not None and i >= limit: + return + yield item + + +def extract_texts(paths: list[str], config: ExtractConfig, normalize=None, log=None) -> ExtractResult: + """Seeded reservoir sample of text per source, bucketed by script.""" + rng = random.Random(config.seed) + buckets: dict[str, list[str]] = {} + per_source = {} + # Each argument is one user-selected dataset (file, directory or glob): + # one reservoir per dataset, so a corpus split into many shards is not + # over-weighted. Records are capped per file so every shard is sampled. + for source in paths: + files = [f for f in expand([source]) if not any(x in f for x in config.exclude_substrings)] + if not files: + continue + per_file = None + if config.max_records_per_source: + per_file = max(1, config.max_records_per_source // len(files)) + reservoir: list[str] = [] + seen = records = 0 + for record in (r for f in files for r in _capped(iter_records(f), per_file)): + records += 1 + for text in _texts_from_record(record, config.fields): + if normalize is not None: + text = normalize(text) + if not (config.min_chars <= len(text) <= config.max_chars): + continue + seen += 1 + if len(reservoir) < config.per_source_texts: + reservoir.append(text) + else: + j = rng.randrange(seen) + if j < config.per_source_texts: + reservoir[j] = text + counts = {} + for text in reservoir: + b = script_bucket(text) + buckets.setdefault(b, []).append(text) + counts[b] = counts.get(b, 0) + 1 + per_source[source] = {"files": len(files), "records": records, "texts_seen": seen, + "sampled": len(reservoir), "buckets": counts} + if log: + log(f"[tokenizer] {source}: files={len(files)} records={records} sampled={len(reservoir)} {counts}") + return ExtractResult(texts_by_bucket=buckets, per_source=per_source) diff --git a/chronos/tokenizer/model.py b/chronos/tokenizer/model.py new file mode 100644 index 0000000..854d27b --- /dev/null +++ b/chronos/tokenizer/model.py @@ -0,0 +1,197 @@ +"""Chronos tokenizer runtime: encoder/decoder plus HF-compatible export. + +The pure-Python encoder is the reference implementation of the trained +merges. ``save`` also writes a ``tokenizer.json`` (HF ``tokenizers`` BPE with +byte fallback and the same pre-tokenizer) so training/inference code that +expects ``AutoTokenizer`` works unchanged; tests assert both agree. +""" +from __future__ import annotations + +import hashlib +import json +import os +import re +from functools import lru_cache + +from .normalize import PRETOKEN_PATTERN, SPACE, from_model_text, pretokenize, to_model_text +from .trainer import BYTE_TOKENS, TrainedBPE + +_CONTROL_PATTERN = r"[\u0000-\u0008\u000b\u000c\u000e-\u001f\u007f\u200b\u200c\u200e\u200f\ufeff]" +_CONTROL_RE = re.compile(_CONTROL_PATTERN) +_BYTE_RE = re.compile(r"^<0x([0-9A-F]{2})>$") +_BYTE_SPLIT_RE = re.compile(r"(<0x[0-9A-F]{2}>)") + +CHAT_TEMPLATE = ( + "{% for message in messages %}" + "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n' }}" + "{% endfor %}" + "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}" +) + +DEFAULT_SPECIAL_TOKENS = ( + "<|endoftext|>", "<|im_start|>", "<|im_end|>", + "<|en|>", "<|zh|>", "<|translate|>", "<|pad|>", +) + tuple(f"<|reserved_{i}|>" for i in range(57)) + + +def tokenizer_normalize(text: str) -> str: + """Tokenizer-time normalization: NFC + control removal only. + + Whitespace collapsing, de-segmentation and punctuation policy belong to + corpus cleaning; doing them here would make ``encode(a + b)`` differ from + ``encode(a) + encode(b)`` around chat-template boundaries. + """ + import unicodedata + return _CONTROL_RE.sub("", unicodedata.normalize("NFC", text)) + + +class ChronosTokenizer: + def __init__(self, vocab: dict[str, int], merges: list[tuple[str, str]], + special_tokens: tuple[str, ...], metadata: dict | None = None): + self.vocab = dict(vocab) + self.id_to_token = {i: t for t, i in self.vocab.items()} + self.merges = [tuple(m) for m in merges] + self.ranks = {pair: r for r, pair in enumerate(self.merges)} + self.special_tokens = tuple(special_tokens) + self.metadata = dict(metadata or {}) + self._byte_ids = [self.vocab[t] for t in BYTE_TOKENS] + self._byte_of_id = {i: b for b, i in enumerate(self._byte_ids)} + escaped = sorted((re.escape(t) for t in self.special_tokens), key=len, reverse=True) + self._special_re = re.compile("(" + "|".join(escaped) + ")") if escaped else None + self._encode_word = lru_cache(maxsize=1 << 18)(self._encode_word_uncached) + + @classmethod + def from_trained(cls, trained: TrainedBPE, special_tokens, metadata=None): + return cls(trained.vocab, trained.merges, tuple(special_tokens), metadata) + + @property + def vocab_size(self) -> int: + return len(self.vocab) + + def token_id(self, token: str) -> int: + return self.vocab[token] + + def _encode_word_uncached(self, word: str) -> tuple[int, ...]: + symbols: list[str] = [] + for ch in word: + if ch in self.vocab: + symbols.append(ch) + else: + symbols.extend(BYTE_TOKENS[b] for b in ch.encode("utf-8")) + while len(symbols) > 1: + best, best_rank = None, None + for i in range(len(symbols) - 1): + rank = self.ranks.get((symbols[i], symbols[i + 1])) + if rank is not None and (best_rank is None or rank < best_rank): + best, best_rank = i, rank + if best is None: + break + a, b = symbols[best], symbols[best + 1] + merged, out, i = a + b, [], 0 + while i < len(symbols): + if i + 1 < len(symbols) and symbols[i] == a and symbols[i + 1] == b: + out.append(merged) + i += 2 + else: + out.append(symbols[i]) + i += 1 + symbols = out + return tuple(self.vocab[s] for s in symbols) + + # ---- encode / decode ------------------------------------------------- + def encode(self, text: str, *, allow_special: bool = True) -> list[int]: + text = tokenizer_normalize(text) + pieces = self._special_re.split(text) if (allow_special and self._special_re) else [text] + ids: list[int] = [] + for piece in pieces: + if not piece: + continue + if allow_special and piece in self.special_tokens: + ids.append(self.vocab[piece]) + continue + for word in pretokenize(to_model_text(piece)): + ids.extend(self._encode_word(word)) + return ids + + def decode(self, ids, *, skip_special: bool = False) -> str: + # Merged tokens may contain several ``<0xNN>`` pieces (multi-byte + # characters outside the alphabet); split every token back into raw + # bytes and text, then decode the byte runs as UTF-8. + data = bytearray() + for i in ids: + token = self.id_to_token.get(int(i), "") + if skip_special and token in self.special_tokens: + continue + for piece in _BYTE_SPLIT_RE.split(token): + if not piece: + continue + match = _BYTE_RE.match(piece) + data.extend(bytes([int(match.group(1), 16)]) if match else piece.encode("utf-8")) + return from_model_text(data.decode("utf-8", errors="replace")) + + # ---- persistence ----------------------------------------------------- + def fingerprint(self) -> str: + payload = json.dumps({"vocab": self.vocab, "merges": self.merges, + "special": self.special_tokens}, ensure_ascii=False, sort_keys=True) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def hf_tokenizer_json(self) -> dict: + added = [{"id": self.vocab[t], "content": t, "single_word": False, "lstrip": False, + "rstrip": False, "normalized": False, "special": True} for t in self.special_tokens] + return { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": added, + "normalizer": {"type": "Sequence", "normalizers": [ + {"type": "NFC"}, + {"type": "Replace", "pattern": {"Regex": _CONTROL_PATTERN}, "content": ""}, + {"type": "Replace", "pattern": {"String": " "}, "content": SPACE}, + ]}, + "pre_tokenizer": {"type": "Split", "pattern": {"Regex": PRETOKEN_PATTERN}, + "behavior": "Isolated", "invert": False}, + "post_processor": None, + "decoder": {"type": "Sequence", "decoders": [ + {"type": "Replace", "pattern": {"String": SPACE}, "content": " "}, + {"type": "ByteFallback"}, + {"type": "Fuse"}, + ]}, + "model": {"type": "BPE", "dropout": None, "unk_token": None, + "continuing_subword_prefix": None, "end_of_word_suffix": None, + "fuse_unk": False, "byte_fallback": True, "ignore_merges": False, + "vocab": self.vocab, "merges": [[a, b] for a, b in self.merges]}, + } + + def save(self, directory: str) -> str: + os.makedirs(directory, exist_ok=True) + meta = dict(self.metadata) + meta.update(vocab_size=self.vocab_size, fingerprint=self.fingerprint(), + special_tokens=list(self.special_tokens), algorithm="chronos-bpe-v1") + with open(os.path.join(directory, "chronos_tokenizer.json"), "w", encoding="utf-8") as fh: + json.dump({"vocab": self.vocab, "merges": self.merges, + "special_tokens": list(self.special_tokens), "metadata": meta}, + fh, ensure_ascii=False) + with open(os.path.join(directory, "tokenizer.json"), "w", encoding="utf-8") as fh: + json.dump(self.hf_tokenizer_json(), fh, ensure_ascii=False) + config = { + "tokenizer_class": "PreTrainedTokenizerFast", + "bos_token": "<|im_start|>", "eos_token": "<|im_end|>", + "unk_token": None, "pad_token": "<|endoftext|>", + "model_max_length": 32768, "clean_up_tokenization_spaces": False, + "chat_template": CHAT_TEMPLATE, "add_bos_token": False, "add_eos_token": False, + } + with open(os.path.join(directory, "tokenizer_config.json"), "w", encoding="utf-8") as fh: + json.dump(config, fh, ensure_ascii=False, indent=2) + with open(os.path.join(directory, "special_tokens_map.json"), "w", encoding="utf-8") as fh: + json.dump({"bos_token": "<|im_start|>", "eos_token": "<|im_end|>", + "pad_token": "<|endoftext|>"}, fh, ensure_ascii=False, indent=2) + with open(os.path.join(directory, "tokenizer_report.json"), "w", encoding="utf-8") as fh: + json.dump(meta, fh, ensure_ascii=False, indent=2) + return directory + + @classmethod + def load(cls, directory: str) -> "ChronosTokenizer": + with open(os.path.join(directory, "chronos_tokenizer.json"), encoding="utf-8") as fh: + data = json.load(fh) + return cls(data["vocab"], [tuple(m) for m in data["merges"]], + tuple(data["special_tokens"]), data.get("metadata")) diff --git a/chronos/tokenizer/normalize.py b/chronos/tokenizer/normalize.py new file mode 100644 index 0000000..828857a --- /dev/null +++ b/chronos/tokenizer/normalize.py @@ -0,0 +1,81 @@ +"""Text normalization and script-aware pre-tokenization for Chronos tokenizers. + +The same pattern is used by the trainer, the pure-Python encoder and the +exported ``tokenizer.json`` (HF ``Split`` pre-tokenizer), so all three agree +token for token. +""" +from __future__ import annotations + +import re +import unicodedata + +import regex + +SPACE = "\u2581" # marks a preceding space, SentencePiece-style. + +# Pre-token classes, in priority order: +# 1. runs of Han ideographs (BPE learns multi-character words inside them) +# 2. an optional ▁ followed by letters/marks of any non-Han script +# 3. a single digit (numbers are always split digit by digit) +# 4. an optional ▁ followed by one punctuation/symbol run +# 5. leftover ▁ runs and any other whitespace +PRETOKEN_PATTERN = ( + r"\p{Han}+" + r"|" + SPACE + r"?[^\s" + SPACE + r"\p{Han}\p{P}\p{S}\p{N}]+" + r"|\p{N}" + r"|" + SPACE + r"?[\p{P}\p{S}]+" + r"|" + SPACE + r"+" + r"|\s+" +) +_PRETOKEN_RE = regex.compile(PRETOKEN_PATTERN) + +_CONTROL_RE = re.compile(r"[\u0000-\u0008\u000b\u000c\u000e-\u001f\u007f\u200b\u200c\u200e\u200f\ufeff]") +_SPACES_RE = re.compile(r"[ \t\u00a0\u3000]+") +# Spaces touching a CJK character, or padding punctuation, come from word +# segmentation; spaces between Latin tokens (``New Day Blue``) are kept. +_CJK = r"\u3000-\u303f\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff\uff00-\uffef" +_CJK_GAP_RE = re.compile(rf"(?<=[{_CJK}])\s+|\s+(?=[{_CJK}])") +_PUNCT_GAP_RE = re.compile(r"\s+(?=[,.;:!?%)\]}])|(?<=[(\[{])\s+") + + +def normalize_text(text: str, *, form: str = "NFKC", unsegment_cjk: bool = False) -> str: + """NFKC, control-character removal, whitespace collapsing. + + ``unsegment_cjk`` removes spaces between CJK characters (for corpora that + ship word-segmented Chinese, e.g. the iic WMT CSV). + """ + if form: + text = unicodedata.normalize(form, text) + text = _CONTROL_RE.sub("", text) + text = _SPACES_RE.sub(" ", text).strip() + if unsegment_cjk: + text = _CJK_GAP_RE.sub("", text) + text = _PUNCT_GAP_RE.sub("", text) + return text + + +def to_model_text(text: str) -> str: + """Make spaces visible: ``a b`` → ``a▁b`` (no leading marker).""" + return text.replace(" ", SPACE) + + +def from_model_text(text: str) -> str: + return text.replace(SPACE, " ") + + +def pretokenize(model_text: str) -> list[str]: + return _PRETOKEN_RE.findall(model_text) + + +def looks_word_segmented_cjk(samples: list[str], threshold: float = 0.15) -> bool: + """Heuristic detector: share of CJK-CJK boundaries separated by a space.""" + gaps = joins = 0 + pair = re.compile(rf"[{_CJK}](\s?)(?=[{_CJK}])") + for text in samples: + for match in pair.finditer(text): + if match.group(1): + gaps += 1 + else: + joins += 1 + total = gaps + joins + return total > 0 and gaps / total >= threshold diff --git a/chronos/tokenizer/trainer.py b/chronos/tokenizer/trainer.py new file mode 100644 index 0000000..87f503d --- /dev/null +++ b/chronos/tokenizer/trainer.py @@ -0,0 +1,207 @@ +"""Chronos BPE trainer: script-aware, quota-balanced, byte-fallback. + +Algorithm (self-contained, no external tokenizer library): + +1. Pre-tokenize with ``normalize.PRETOKEN_PATTERN`` and count unique + pre-tokens per language bucket. +2. Base alphabet: characters covering ``char_coverage`` of the corpus mass + plus 256 ``<0xNN>`` byte tokens, so encoding is total (no UNK). +3. Merge loop with incremental pair statistics: a max-heap keyed by + language-balanced pair counts with lazy invalidation, plus an inverted + index pair → words. Each merge only updates words that contain the + pair, giving O(Σ affected-word length · log P) per merge. +4. Language balance: a pair's score is Σ_bucket count_b(pair) · w_b, where + ``w_b = (share_b ** temperature) / share_b`` re-weights buckets toward a + tempered distribution, so the high-resource language cannot crowd out the + other's subwords. +5. Optional minimum frequency stops degenerate merges on tail data. +""" +from __future__ import annotations + +import heapq +import math +from collections import Counter, defaultdict +from dataclasses import dataclass, field + +from .normalize import pretokenize, to_model_text + +BYTE_TOKENS = [f"<0x{b:02X}>" for b in range(256)] + + +@dataclass +class BPETrainerConfig: + vocab_size: int = 16384 + special_tokens: tuple[str, ...] = () + char_coverage: float = 0.9999 + # Characters beyond this share of the vocabulary fall back to bytes, + # leaving room for multi-character merges in CJK-heavy corpora. + max_alphabet_fraction: float = 0.4 + min_pair_frequency: int = 2 + language_temperature: float = 0.5 + max_token_chars: int = 16 + # Buckets below this mass share are pooled into ``other``; no bucket + # weight may exceed ``max_bucket_weight`` (stray scripts must not + # dominate merge selection). + min_bucket_share: float = 0.01 + max_bucket_weight: float = 4.0 + + +@dataclass +class TrainedBPE: + vocab: dict[str, int] + merges: list[tuple[str, str]] + alphabet: list[str] + stats: dict = field(default_factory=dict) + + +def count_pretokens(texts_by_bucket: dict[str, list[str]]) -> dict[str, Counter]: + counts = {} + for bucket, texts in texts_by_bucket.items(): + counter = Counter() + for text in texts: + counter.update(pretokenize(to_model_text(text))) + counts[bucket] = counter + return counts + + +def _pool_small_buckets(counts: dict[str, Counter], min_share: float) -> dict[str, Counter]: + mass = {b: sum(len(w) * c for w, c in ctr.items()) for b, ctr in counts.items()} + total = sum(mass.values()) or 1 + pooled: dict[str, Counter] = {} + for b, ctr in counts.items(): + key = b if mass[b] / total >= min_share else "other" + pooled.setdefault(key, Counter()).update(ctr) + return pooled + + +def _bucket_weights(counts: dict[str, Counter], temperature: float, max_weight: float = 4.0) -> dict[str, float]: + mass = {b: sum(len(w) * c for w, c in ctr.items()) for b, ctr in counts.items()} + total = sum(mass.values()) or 1 + shares = {b: m / total for b, m in mass.items() if m} + tempered = {b: s ** temperature for b, s in shares.items()} + norm = sum(tempered.values()) or 1 + return {b: min(max_weight, (tempered[b] / norm) / shares[b]) for b in shares} + + +def _alphabet(words: dict[str, float], coverage: float, limit: int | None = None) -> list[str]: + chars = Counter() + for word, freq in words.items(): + for ch in word: + chars[ch] += freq + total = sum(chars.values()) or 1 + kept, running = [], 0.0 + for ch, freq in chars.most_common(): + if running / total >= coverage or (limit is not None and len(kept) >= limit): + break + kept.append(ch) + running += freq + return sorted(kept) + + +def _initial_symbols(word: str, alphabet: set[str]) -> list[str]: + symbols = [] + for ch in word: + if ch in alphabet: + symbols.append(ch) + else: + symbols.extend(BYTE_TOKENS[b] for b in ch.encode("utf-8")) + return symbols + + +def _pairs(symbols: list[str]): + return zip(symbols, symbols[1:]) + + +def _merge_word(symbols: list[str], a: str, b: str, joined: str) -> list[str]: + out, i, n = [], 0, len(symbols) + while i < n: + if i + 1 < n and symbols[i] == a and symbols[i + 1] == b: + out.append(joined) + i += 2 + else: + out.append(symbols[i]) + i += 1 + return out + + +def train_bpe(texts_by_bucket: dict[str, list[str]], config: BPETrainerConfig, + progress=None) -> TrainedBPE: + counts = _pool_small_buckets(count_pretokens(texts_by_bucket), config.min_bucket_share) + weights = _bucket_weights(counts, config.language_temperature, config.max_bucket_weight) + word_freq: dict[str, float] = defaultdict(float) + for bucket, counter in counts.items(): + w = weights.get(bucket, 0.0) + for word, c in counter.items(): + word_freq[word] += c * w + limit = max(0, int(config.vocab_size * config.max_alphabet_fraction)) + alphabet = _alphabet(word_freq, config.char_coverage, limit) + alphabet_set = set(alphabet) + + vocab: dict[str, int] = {} + for token in (*config.special_tokens, *BYTE_TOKENS, *alphabet): + vocab.setdefault(token, len(vocab)) + if len(vocab) > config.vocab_size: + raise ValueError(f"base vocabulary ({len(vocab)}) exceeds vocab_size {config.vocab_size}; " + "lower char_coverage") + + words = list(word_freq) + freqs = [word_freq[w] for w in words] + symbols = [_initial_symbols(w, alphabet_set) for w in words] + pair_count: dict[tuple[str, str], float] = defaultdict(float) + pair_words: dict[tuple[str, str], set[int]] = defaultdict(set) + for wi, syms in enumerate(symbols): + for pair in _pairs(syms): + pair_count[pair] += freqs[wi] + pair_words[pair].add(wi) + heap = [(-c, pair) for pair, c in pair_count.items()] + heapq.heapify(heap) + + merges: list[tuple[str, str]] = [] + target = config.vocab_size + while len(vocab) < target and heap: + neg, pair = heapq.heappop(heap) + current = pair_count.get(pair, 0.0) + if current <= 0: + continue + if not math.isclose(-neg, current, rel_tol=1e-9, abs_tol=1e-9): + heapq.heappush(heap, (-current, pair)) # stale entry: re-queue true count + continue + if current < config.min_pair_frequency * min(weights.values() or [1.0]): + break + a, b = pair + joined = a + b + if len(joined) > config.max_token_chars or joined in vocab: + pair_count.pop(pair, None) + continue + merges.append(pair) + vocab[joined] = len(vocab) + touched: dict[tuple[str, str], float] = defaultdict(float) + for wi in list(pair_words.pop(pair, ())): + old = symbols[wi] + new = _merge_word(old, a, b, joined) + if len(new) == len(old): + continue + f = freqs[wi] + for p in _pairs(old): + pair_count[p] -= f + touched[p] -= f + for p in _pairs(new): + pair_count[p] += f + touched[p] += f + pair_words[p].add(wi) + symbols[wi] = new + pair_count.pop(pair, None) + for p, delta in touched.items(): + if delta > 0 and pair_count.get(p, 0.0) > 0: + heapq.heappush(heap, (-pair_count[p], p)) + if progress and len(merges) % 500 == 0: + progress(len(vocab), target) + + stats = { + "bucket_weights": weights, + "unique_pretokens": len(words), + "alphabet_size": len(alphabet), + "merges": len(merges), + "final_vocab": len(vocab), + } + return TrainedBPE(vocab=vocab, merges=merges, alphabet=alphabet, stats=stats) diff --git a/chronos/trainer/__init__.py b/chronos/trainer/__init__.py index 57b373f..1db4795 100644 --- a/chronos/trainer/__init__.py +++ b/chronos/trainer/__init__.py @@ -4,3 +4,4 @@ from .orpo_trainer import ChronosORPOTrainer from .grpo_trainer import ChronosGRPOTrainer from .distill_trainer import ChronosDistillTrainer +from .paged_store import PageKey, TieredTensorStore, PagedOptimizerState, PagedTrainingController diff --git a/chronos/trainer/chronos_trainer.py b/chronos/trainer/chronos_trainer.py index fe021d4..d80cfbc 100644 --- a/chronos/trainer/chronos_trainer.py +++ b/chronos/trainer/chronos_trainer.py @@ -21,7 +21,9 @@ from chronos.model.config import ChronosConfig from chronos.model.model_chronos import ChronosForCausalLM from chronos.model.temporal_loss import lookahead_supervision_loss -from chronos.trainer.loss_mixin import chronos_loss_term, collect_router_probs +from chronos.trainer.loss_mixin import ( + chronos_loss_term, collect_router_probs, first_iteration_router_probs, lookahead_teacher, +) from chronos.model.checkpoint import save_state_dict_with_config from chronos.trainer.device_utils import ( autocast_context, @@ -30,6 +32,7 @@ optimizer_step_with_scaler, runtime_summary, ) +from chronos.trainer.paged_store import PagedTrainingController class ChronosTrainer: @@ -63,26 +66,42 @@ def __init__(self, config: ChronosConfig, args): ) self.model = ChronosForCausalLM(config).to(self.device) - from chronos.trainer.optim_utils import build_optimizer - self.optimizer = build_optimizer( - self.model, lr=args.learning_rate, + self.paged_training = PagedTrainingController(self.model, config) + from chronos.trainer.optim_utils import build_configured_optimizer + self.optimizer = build_configured_optimizer( + self.model, config, lr=args.learning_rate, weight_decay=float(getattr(args, "weight_decay", 0.01)), ) + # One compute dtype end to end; FP32 masters live in the optimizer. + from chronos.trainer.precision_training import apply_training_precision + apply_training_precision(self, models=()) def _collect_router_probs(self) -> torch.Tensor: """Gather last_router_probs from all MoE layers → [B, S, L, E].""" return collect_router_probs(self.model) - def train_step(self, input_ids, labels, step, total_steps): + def train_step(self, input_ids, labels, step, total_steps, *, modality_inputs=None, modality_mask=None): + if hasattr(self.model.model, "recurrent_controller"): + self.model.model.recurrent_controller.set_training_progress(step, total_steps) input_ids = input_ids.to(self.device) labels = labels.to(self.device) + if modality_inputs: + modality_inputs = {k: v.to(self.device) for k, v in modality_inputs.items()} + if modality_mask is not None: + modality_mask = modality_mask.to(self.device) lr = get_lr(step, total_steps, self.args.learning_rate) for pg in self.optimizer.param_groups: pg['lr'] = lr with self.autocast_ctx: - outputs, lookahead_probs = self.model(input_ids, labels=labels) + paged_masks = self.paged_training.masks_for_forward(self.device) + result = self.model( + input_ids, labels=labels, + modality_inputs=modality_inputs, modality_mask=modality_mask, + available_expert_masks=paged_masks, + ) + outputs, lookahead_probs = result.legacy_tuple() ce_loss = outputs.loss aux_loss = outputs.aux_loss @@ -93,14 +112,14 @@ def train_step(self, input_ids, labels, step, total_steps): # Standalone lookahead loss value for logging lookahead_loss_val = 0.0 - router_4d = self._collect_router_probs() + router_4d = first_iteration_router_probs(self.model) if ( router_4d is not None and router_4d.shape[1] > 1 and lookahead_probs is not None and self.config.lookahead_steps > 0 ): - teacher = router_4d.mean(dim=2).detach() + teacher = lookahead_teacher(router_4d, lookahead_probs, self.config) la_term = lookahead_supervision_loss( lookahead_probs, teacher, self.config.lookahead_steps, ) @@ -117,6 +136,7 @@ def train_step(self, input_ids, labels, step, total_steps): self.model.parameters(), self.args.grad_clip, ) + self.paged_training.observe() return loss.item() * self.args.accumulation_steps, ce_loss.item(), aux_loss.item(), lookahead_loss_val @@ -126,14 +146,21 @@ def train_epoch(self, epoch, loader, iters, start_step=0, wandb=None, max_steps= cap = int(max_steps) if max_steps is not None else None total_steps = cap or (self.args.epochs * iters) - for local_step, (input_ids, labels) in enumerate(loader, start=1): + for local_step, batch in enumerate(loader, start=1): if local_step > iters: break step = int(start_step) + local_step if cap is not None and step > cap: 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 loss, ce_loss, aux_loss, la_loss = self.train_step( - input_ids, labels, step, total_steps + input_ids, labels, step, total_steps, + modality_inputs=media, modality_mask=media_mask, ) if local_step % self.args.log_interval == 0 or local_step == iters: diff --git a/chronos/trainer/device_utils.py b/chronos/trainer/device_utils.py index 852e732..1d72a52 100644 --- a/chronos/trainer/device_utils.py +++ b/chronos/trainer/device_utils.py @@ -61,10 +61,11 @@ def torch_dtype_from_name(name: str | None) -> TorchDType: def resolve_dtype_name(device: str | TorchDevice | None, dtype_name: str | None) -> str: """Resolve user-facing dtype policy to a concrete dtype string. - `auto` is intentionally backend-specific. Apple MPS/MLX default to bf16: - it keeps the fp32 exponent range and avoids the fp16 overflows that show up - in router softmax, CE, and Adam moments. CUDA/XPU keep the existing fp16 - default, while CPU stays fp32 unless the user explicitly requests fp16/bf16. + `auto` is backend-specific and, for torch devices, identical to the + run's precision policy (``chronos.precision.resolve_torch``) so the UI, + the trainers and serving agree: CPU and MPS FP32 (measured faster than + BF16 on MPS), CUDA BF16 (FP16 when the GPU lacks BF16). MLX GPU keeps + BF16 (FP32 masters in the optimizer); XPU keeps FP16. """ value = (dtype_name or "auto").strip().lower() if value in {"fp16", "float16", "half"}: @@ -74,20 +75,30 @@ def resolve_dtype_name(device: str | TorchDevice | None, dtype_name: str | None) if value in {"fp32", "float32", "full"}: return "float32" device_type = torch_device_type(device) - if device_type in {"mps", "mlx"}: + if device_type == "mlx": return "bfloat16" - if device_type in {"cuda", "xpu"}: + if device_type == "xpu": return "float16" + if device_type in {"cpu", "mps", "cuda"}: + try: + from chronos.precision import resolve_torch + return resolve_torch("auto", device_type).compute + except Exception: + return "float16" if device_type == "cuda" else "float32" return "float32" -def autocast_context(device: str | TorchDevice | None, dtype_name: str | None): +def autocast_context(device: str | TorchDevice | None, dtype_name: str | None, model=None): """Return an autocast context suitable for CPU/CUDA/MPS/XPU. - MPS supports torch.autocast, but GradScaler is not used there. If the user - requests float32, autocast is disabled explicitly. + When ``model`` was placed by ``prepare_torch_training`` / + ``place_torch_model`` it already computes natively in the policy dtype, + so autocast is disabled (no per-op weight casts). Otherwise this is the + legacy autocast path. float32 disables autocast explicitly. """ torch = _torch() + if model is not None and getattr(model, "_chronos_precision", None) is not None: + return contextlib.nullcontext() device_type = torch_device_type(device) dtype = torch_dtype_from_name(resolve_dtype_name(device, dtype_name)) if device_type == "mlx" or dtype == torch.float32: @@ -121,12 +132,26 @@ def optimizer_step_with_scaler( parameters, grad_clip: float, ) -> None: + """Unscale -> clip -> step -> zero. + + With FP32 masters (``prepare_torch_training``) the gradients live on the + masters, so clipping uses the optimizer's own parameters and the updated + masters are copied back to the low-precision model once. + """ torch = _torch() + masters = getattr(optimizer, "chronos_masters", None) + if masters is not None: + parameters = [p for g in optimizer.param_groups for p in g["params"] if p.grad is not None] scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(parameters, grad_clip) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) + if masters is not None: + masters.sync_to_model() + evict = getattr(optimizer, "evict_expert_moments", None) + if callable(evict): + evict() def _physical_cores() -> int: diff --git a/chronos/trainer/distill_trainer.py b/chronos/trainer/distill_trainer.py index 9a957e0..caf64ff 100644 --- a/chronos/trainer/distill_trainer.py +++ b/chronos/trainer/distill_trainer.py @@ -22,7 +22,10 @@ import chronos.deps # noqa: F401 from trainer.trainer_utils import Logger # type: ignore -from chronos.data.flexible_dataset import StreamingSFTDataset as SFTDataset +from chronos.data.flexible_dataset import ( + StreamingSFTDataset as SFTDataset, + MultimodalStreamingSFTDataset, multimodal_collate, +) from chronos.model.config import ChronosConfig from chronos.model.model_chronos import ChronosForCausalLM @@ -43,6 +46,7 @@ optimizer_step_with_scaler, runtime_summary, ) +from chronos.trainer.paged_store import PagedTrainingController def distillation_kl(student_logits: torch.Tensor, teacher_logits: torch.Tensor, @@ -83,36 +87,68 @@ def __init__(self, model: ChronosForCausalLM, teacher, config: ChronosConfig, f"device={summary.device} type={summary.device_type} dtype={summary.dtype} " f"cpu_threads={summary.cpu_threads} autocast={summary.autocast} scaler={summary.scaler}" ) - from chronos.trainer.optim_utils import build_optimizer - self.optimizer = build_optimizer( - model, lr=args.learning_rate, + from chronos.trainer.optim_utils import build_configured_optimizer + self.optimizer = build_configured_optimizer( + model, config, lr=args.learning_rate, weight_decay=float(getattr(args, "weight_decay", 0.01)), ) + # One compute dtype end to end; FP32 masters live in the optimizer. + from chronos.trainer.precision_training import apply_training_precision + apply_training_precision(self, models=(self.teacher,)) + self.paged_training = PagedTrainingController(self.model, config) self._router_ref = None - def set_calibration_batch(self, x: torch.Tensor): - self._router_ref = capture_reference_routing(self.model, x, self.device) + def set_calibration_batch(self, x: torch.Tensor, *, modality_inputs=None, modality_mask=None): + self._router_ref = capture_reference_routing( + self.model, x, self.device, + modality_inputs=modality_inputs, modality_mask=modality_mask, + ) - def _teacher_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + def _teacher_logits(self, input_ids: torch.Tensor, *, modality_inputs=None, modality_mask=None) -> torch.Tensor: with torch.no_grad(): - out = self.teacher(input_ids) + teacher_kwargs = {} + # A text-only teacher is a valid distillation source for a + # multimodal student; do not make it fail merely because the + # student batch carries media tensors. + if getattr(getattr(self.teacher, "model", None), "multimodal_fusion", None) is not None: + teacher_kwargs = { + "modality_inputs": modality_inputs, + "modality_mask": modality_mask, + } + if getattr(self.teacher, "_chronos_precision", None) is not None: + teacher_kwargs["logits_dtype"] = torch.float32 # born FP32, no cast after + out = self.teacher(input_ids, **teacher_kwargs) return (out[0].logits if isinstance(out, tuple) else out.logits).detach() - def train_step(self, input_ids, labels, step, total_steps): + def train_step(self, input_ids, labels, step, total_steps, *, modality_inputs=None, modality_mask=None): + if hasattr(self.model.model, "recurrent_controller"): + self.model.model.recurrent_controller.set_training_progress(step, total_steps) input_ids = input_ids.to(self.device) labels = labels.to(self.device) + if modality_inputs: + modality_inputs = {k: v.to(self.device) for k, v in modality_inputs.items()} + if modality_mask is not None: + modality_mask = modality_mask.to(self.device) lr = get_lr(step, total_steps, self.args.learning_rate) for pg in self.optimizer.param_groups: pg["lr"] = lr with self.autocast_ctx: - outputs, lookahead_probs = self.model(input_ids, labels=labels) + paged_masks = self.paged_training.masks_for_forward(self.device) + result = self.model( + input_ids, labels=labels, + modality_inputs=modality_inputs, modality_mask=modality_mask, + available_expert_masks=paged_masks, + ) + outputs, lookahead_probs = result.legacy_tuple() ce_loss = outputs.loss student_logits = outputs.logits # Teacher forward (no grad) - teacher_logits = self._teacher_logits(input_ids) + teacher_logits = self._teacher_logits( + input_ids, modality_inputs=modality_inputs, modality_mask=modality_mask, + ) # Align vocab dims if teacher is smaller (pad or truncate) if teacher_logits.shape[-1] != student_logits.shape[-1]: V = min(teacher_logits.shape[-1], student_logits.shape[-1]) @@ -171,6 +207,7 @@ def train_step(self, input_ids, labels, step, total_steps): self.model.parameters(), self.args.grad_clip, ) + self.paged_training.observe() return { "loss": float(loss.item() * self.args.accumulation_steps), @@ -184,13 +221,22 @@ def train_epoch(self, epoch, loader, iters, start_step=0, max_steps=None): self.model.train() cap = int(max_steps) if max_steps is not None else None total_steps = cap or (self.args.epochs * iters) - for local_step, (ids, labels) in enumerate(loader, start=1): + for local_step, batch in enumerate(loader, start=1): if local_step > iters: break step = int(start_step) + local_step if cap is not None and step > cap: break - s = self.train_step(ids, labels, step, total_steps) + if isinstance(batch, dict): + ids, labels = batch["input_ids"], batch["labels"] + s = self.train_step( + ids, labels, step, total_steps, + modality_inputs=batch.get("modality_inputs"), + modality_mask=batch.get("modality_mask"), + ) + else: + ids, labels = batch + s = self.train_step(ids, labels, step, total_steps) if local_step % self.args.log_interval == 0 or local_step == iters: lr = self.optimizer.param_groups[-1]["lr"] Logger( @@ -207,10 +253,19 @@ def _save(self, epoch, step): self.model.train() -def build_distill_loader(data_path, tokenizer, max_seq_len, batch_size, device="cpu", num_workers="auto"): - ds = SFTDataset(data_path, tokenizer, max_length=max_seq_len) +def build_distill_loader(data_path, tokenizer, max_seq_len, batch_size, device="cpu", num_workers="auto", *, + modalities=None, image_size=224, audio_sample_rate=16000, + max_audio_seconds=30.0, media_error_policy="warn"): + enabled = {str(m).lower() for m in (modalities or []) if str(m).lower() != "text"} + ds = (MultimodalStreamingSFTDataset if enabled else SFTDataset)( + data_path, tokenizer, max_length=max_seq_len, + **({"image_size": image_size, "audio_sample_rate": audio_sample_rate, + "max_audio_seconds": max_audio_seconds, "media_error_policy": media_error_policy} + if enabled else {}), + ) return DataLoader( ds, batch_size=batch_size, + collate_fn=multimodal_collate if enabled else None, **dataloader_kwargs(device=device, num_workers=num_workers, shuffle=True), ) diff --git a/chronos/trainer/dpo_trainer.py b/chronos/trainer/dpo_trainer.py index ecdee36..7d6e537 100644 --- a/chronos/trainer/dpo_trainer.py +++ b/chronos/trainer/dpo_trainer.py @@ -42,6 +42,7 @@ optimizer_step_with_scaler, runtime_summary, ) +from chronos.trainer.paged_store import PagedTrainingController def _logits_to_log_probs(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: @@ -81,9 +82,9 @@ def __init__(self, model: ChronosForCausalLM, config: ChronosConfig, args, token f"device={summary.device} type={summary.device_type} dtype={summary.dtype} " f"cpu_threads={summary.cpu_threads} autocast={summary.autocast} scaler={summary.scaler}" ) - from chronos.trainer.optim_utils import build_optimizer - self.optimizer = build_optimizer( - model, lr=args.learning_rate, + from chronos.trainer.optim_utils import build_configured_optimizer + self.optimizer = build_configured_optimizer( + model, config, lr=args.learning_rate, weight_decay=float(getattr(args, "weight_decay", 0.01)), ) @@ -95,11 +96,18 @@ def __init__(self, model: ChronosForCausalLM, config: ChronosConfig, args, token self.ref_model.eval().requires_grad_(False) self._router_ref = None + # One compute dtype end to end; FP32 masters live in the optimizer. + from chronos.trainer.precision_training import apply_training_precision + apply_training_precision(self, models=(self.ref_model,)) + self.paged_training = PagedTrainingController(self.model, config) def set_calibration_batch(self, x_calib: torch.Tensor): self._router_ref = capture_reference_routing(self.model, x_calib, self.device) def train_step(self, batch, step, total_steps): + controller = getattr(getattr(self.model, "model", None), "recurrent_controller", None) + if controller is not None: + controller.set_training_progress(step, total_steps) x_c = batch["x_chosen"].to(self.device) x_r = batch["x_rejected"].to(self.device) y_c = batch["y_chosen"].to(self.device) @@ -116,11 +124,13 @@ def train_step(self, batch, step, total_steps): with self.autocast_ctx: with torch.no_grad(): - ref_out = self.ref_model(x) + ref_out = self.ref_model(x, logits_dtype=torch.float32) ref_logits = ref_out[0].logits if isinstance(ref_out, tuple) else ref_out.logits ref_lp = _logits_to_log_probs(ref_logits, y) - outputs, lookahead_probs = self.model(x) + paged_masks = self.paged_training.masks_for_forward(self.device) + result = self.model(x, available_expert_masks=paged_masks) + outputs, lookahead_probs = result.legacy_tuple() pol_lp = _logits_to_log_probs(outputs.logits, y) base = dpo_loss(ref_lp, pol_lp, m, beta=self.beta) @@ -165,6 +175,7 @@ def train_step(self, batch, step, total_steps): self.model.parameters(), self.args.grad_clip, ) + self.paged_training.observe() return ( loss.item() * self.args.accumulation_steps, diff --git a/chronos/trainer/grpo_trainer.py b/chronos/trainer/grpo_trainer.py index 59355b8..f89cecc 100644 --- a/chronos/trainer/grpo_trainer.py +++ b/chronos/trainer/grpo_trainer.py @@ -53,6 +53,7 @@ optimizer_step_with_scaler, runtime_summary, ) +from chronos.trainer.paged_store import PagedTrainingController from chronos.trainer.reward import ToyReward, build_reward_fn # noqa: F401 @@ -107,9 +108,9 @@ def __init__( f"device={summary.device} type={summary.device_type} dtype={summary.dtype} " f"cpu_threads={summary.cpu_threads} autocast={summary.autocast} scaler={summary.scaler}" ) - from chronos.trainer.optim_utils import build_optimizer - self.optimizer = build_optimizer( - model, lr=args.learning_rate, + from chronos.trainer.optim_utils import build_configured_optimizer + self.optimizer = build_configured_optimizer( + model, config, lr=args.learning_rate, weight_decay=float(getattr(args, "weight_decay", 0.01)), ) @@ -117,6 +118,10 @@ def __init__( self.ref_model.load_state_dict(model.state_dict(), strict=False) self.ref_model.eval().requires_grad_(False) self._router_ref = None + # One compute dtype end to end; FP32 masters live in the optimizer. + from chronos.trainer.precision_training import apply_training_precision + apply_training_precision(self, models=(self.ref_model,)) + self.paged_training = PagedTrainingController(self.model, config) def set_calibration_batch(self, x: torch.Tensor): self._router_ref = capture_reference_routing(self.model, x, self.device) @@ -131,7 +136,8 @@ def _rollout(self, prompt_ids: torch.Tensor) -> Tuple[torch.Tensor, List[int]]: past = None generated: List[int] = [] for _ in range(self.max_gen_len): - past_len = past[0][0].shape[1] if past else 0 + from chronos.model.model_chronos import cache_sequence_length + past_len = cache_sequence_length(past) if past else 0 out = self.model( ids[:, past_len:] if past_len > 0 else ids, past_key_values=past, use_cache=True, @@ -151,6 +157,9 @@ def _rollout(self, prompt_ids: torch.Tensor) -> Tuple[torch.Tensor, List[int]]: # ── core step ──────────────────────────────────────────────── def train_step(self, prompt_text: str, step: int, total_steps: int): + controller = getattr(getattr(self.model, "model", None), "recurrent_controller", None) + if controller is not None: + controller.set_training_progress(step, total_steps) # Tokenize prompt prompt = self.tokenizer.apply_chat_template( [{"role": "user", "content": prompt_text}], @@ -190,14 +199,16 @@ def train_step(self, prompt_text: str, step: int, total_steps: int): pg["lr"] = lr with self.autocast_ctx: - outputs, lookahead_probs = self.model(ids_batch) + paged_masks = self.paged_training.masks_for_forward(self.device) + result = self.model(ids_batch, available_expert_masks=paged_masks) + outputs, lookahead_probs = result.legacy_tuple() log_probs = F.log_softmax(outputs.logits[:, :-1, :], dim=-1) target = ids_batch[:, 1:] lp = log_probs.gather(-1, target.unsqueeze(-1)).squeeze(-1) # [G, L-1] mask = response_mask[:, 1:] with torch.no_grad(): - ref_out = self.ref_model(ids_batch) + ref_out = self.ref_model(ids_batch, logits_dtype=torch.float32) ref_outputs = ref_out[0] if isinstance(ref_out, tuple) else ref_out ref_lp = F.log_softmax(ref_outputs.logits[:, :-1, :], dim=-1).gather( -1, target.unsqueeze(-1) @@ -251,6 +262,7 @@ def train_step(self, prompt_text: str, step: int, total_steps: int): self.model.parameters(), self.args.grad_clip, ) + self.paged_training.observe() return { "loss": float(loss.item() * self.args.accumulation_steps), diff --git a/chronos/trainer/loss_mixin.py b/chronos/trainer/loss_mixin.py index e9ca6f7..8df63fa 100644 --- a/chronos/trainer/loss_mixin.py +++ b/chronos/trainer/loss_mixin.py @@ -61,6 +61,45 @@ def collect_router_probs(model, *, with_grad: bool = False) -> Optional[torch.Te return torch.stack(probs, dim=2) +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]``. + + A recurrent block routes once per loop iteration, so V counts virtual + layers (prelude + groups·depth + coda for a full-depth forward). Load + balance and temporal locality must hold for each of them. + """ + probs: List[torch.Tensor] = [] + for layer in model.model.layers: + mlp = getattr(layer, "mlp", None) + if isinstance(mlp, ChronosMOEFeedForward): + probs.extend(mlp.iteration_router_probs_grad if with_grad else mlp.iteration_router_probs) + if not probs: + return None + return torch.stack(probs, dim=2) + + +def first_iteration_router_probs(model) -> Optional[torch.Tensor]: + """Step-0 routing per physical layer (lookahead teacher; prefetch key).""" + probs = [] + for layer in model.model.layers: + mlp = getattr(layer, "mlp", None) + if isinstance(mlp, ChronosMOEFeedForward) and mlp.iteration_router_probs: + probs.append(mlp.iteration_router_probs[0]) + return torch.stack(probs, dim=2) if probs else None + + +def lookahead_teacher(router_4d, lookahead_probs, config): + """Align prediction head g to physical layer prelude_layers + g.""" + if lookahead_probs.dim() == 4: + return router_4d.mean(dim=2).detach() + groups = lookahead_probs.shape[3] + start = int(getattr(config, "prelude_layers", 0)) + teacher = router_4d[:, :, start:start + groups, :] + if teacher.shape[2] != groups: + raise ValueError("lookahead heads do not match the recurrent physical layers") + return teacher.detach() + + def chronos_loss_term( model, base_loss: torch.Tensor, @@ -100,9 +139,11 @@ def chronos_loss_term( aux_term = raw_total if saw_raw else aux_loss loss = loss + lambda_balance * aux_term - router_4d = collect_router_probs(model, with_grad=True) + router_4d = collect_iteration_router_probs(model, with_grad=True) if router_4d is not None and router_4d.shape[1] > 1: - router_mean = router_4d.mean(dim=2) # [B, S, E] + # Locality must hold within each dispatch (virtual layer). Averaging + # layers lets opposite routing changes cancel and hides page churn. + router_mean = router_4d.permute(0, 2, 1, 3).flatten(0, 1) lambda_temporal = float(getattr(config, "lambda_temporal", 0.0)) if lambda_temporal > 0.0: loss = loss + lambda_temporal * temporal_locality_loss(router_mean) @@ -122,7 +163,7 @@ def chronos_loss_term( ): # Teacher must be detached: lookahead targets are stop-grad'd # by construction (we predict the *future* router state). - teacher = router_mean.detach() + teacher = lookahead_teacher(first_iteration_router_probs(model), lookahead_probs, config) la = lookahead_supervision_loss(lookahead_probs, teacher, lookahead_steps) loss = loss + lambda_lookahead * la lambda_lookahead_topk = float(getattr(config, "lambda_lookahead_topk", 0.0)) @@ -143,6 +184,34 @@ def chronos_loss_term( int(getattr(config, "num_experts_per_tok", 1) or 1), ) loss = loss + lambda_lookahead_union * union_loss + # Recurrent-depth regularizers are exposed by ChronosModel without + # changing the HuggingFace output tuple. This keeps all six stage + # trainers source-compatible while shaping the shared transition. + recurrent_state = getattr(getattr(model, "model", None), "last_recurrent_state", None) + if recurrent_state: + ponder_coef = float(getattr(config, "ponder_loss_coef", 0.0)) + stability_coef = float(getattr(config, "recurrent_stability_coef", 0.0)) + distill_coef = float(getattr(config, "recurrent_distill_coef", 0.0)) + if ponder_coef > 0: + loss = loss + ponder_coef * recurrent_state["ponder_loss"] + if stability_coef > 0: + loss = loss + stability_coef * recurrent_state["stability_loss"] + if distill_coef > 0: + loss = loss + distill_coef * recurrent_state["distill_loss"] + # Penalize expansive fixed points while allowing the learned + # contractive Euler step to approach one near convergence. + contraction_budget = float(getattr(config, "recurrent_contraction_budget", 1.0)) + contraction_coef = float(getattr(config, "recurrent_contraction_coef", 0.0)) + if contraction_coef > 0 and "contraction" in recurrent_state: + loss = loss + contraction_coef * ( + recurrent_state["contraction"] - contraction_budget + ).clamp_min(0.0).pow(2) + # Self-speculative draft exit: KL(full.detach() || draft-depth logits), + # computed by ChronosForCausalLM.forward when draft_exit_kl_coef > 0. + draft_kl = getattr(model, "last_draft_exit_kl", None) + draft_coef = float(getattr(config, "draft_exit_kl_coef", 0.0) or 0.0) + if draft_kl is not None and draft_coef > 0.0: + loss = loss + draft_coef * draft_kl return loss @@ -184,7 +253,7 @@ def router_kl_anchor( return lambda_anchor * kl -def capture_reference_routing(model, calibration_batch, device) -> torch.Tensor: +def capture_reference_routing(model, calibration_batch, device, *, modality_inputs=None, modality_mask=None) -> torch.Tensor: """Run `model` in eval mode on the calibration batch, return the layer-averaged routing distribution as a detached reference for the `router_kl_anchor` term. @@ -196,8 +265,12 @@ def capture_reference_routing(model, calibration_batch, device) -> torch.Tensor: model.eval() with torch.no_grad(): ids = calibration_batch.to(device) + if modality_inputs: + modality_inputs = {k: v.to(device) for k, v in modality_inputs.items()} + if modality_mask is not None: + modality_mask = modality_mask.to(device) try: - _ = model(ids, use_cache=False) + _ = model(ids, use_cache=False, modality_inputs=modality_inputs, modality_mask=modality_mask) except TypeError: _ = model(ids) # minimind models don't accept use_cache kwarg here router_4d = collect_router_probs(model) diff --git a/chronos/trainer/optim_utils.py b/chronos/trainer/optim_utils.py index 185e904..cf40c8d 100644 --- a/chronos/trainer/optim_utils.py +++ b/chronos/trainer/optim_utils.py @@ -25,6 +25,8 @@ import torch import torch.nn as nn +from chronos.trainer.paged_store import PageKey, PagedOptimizerState + def _decay_groups(model: nn.Module) -> Tuple[list, list]: """Split trainable parameters into (decay, no_decay) lists. @@ -69,6 +71,257 @@ def build_optimizer( ) +class PagedAdamW(torch.optim.AdamW): + """AdamW with crash-safe per-expert moment snapshots. + + PyTorch parameters remain optimizer-compatible; after every update, expert + moments are streamed page-by-page to SSD. This is the safe first-stage + paging implementation and provides resumable optimizer state without ever + assembling an additional full CPU state dict. Accelerator-resident moment + eviction can be layered on this contract by backend-specific trainers. + """ + def __init__(self, named_parameters, *, page_dir: str, page_capacity: int = 8, + layer_to_group: dict[int, int] | None = None, **kwargs): + named = [(name, p) for name, p in named_parameters if p.requires_grad] + self._names = {id(p): name for name, p in named} + self._param_by_name = {name: p for name, p in named} + self._dirty_expert_params = set() + self.page_store = PagedOptimizerState(page_dir, capacity=page_capacity) + self._hydrated_pages = set() + self._layer_to_group = {int(k): int(v) for k, v in (layer_to_group or {}).items()} + # Keep the canonical two-group decay policy when callers provide a + # model's named parameters. Explicit ``params`` remains supported for + # direct PyTorch use. + if "params" not in kwargs: + decay, no_decay = [], [] + for name, p in named: + if p.ndim < 2 or "norm" in name.lower() or "bias" in name.lower() or "embed" in name.lower(): + no_decay.append(p) + else: + decay.append(p) + super().__init__([ + {"params": decay, "weight_decay": float(kwargs.pop("weight_decay", 0.01))}, + {"params": no_decay, "weight_decay": 0.0}, + ], **kwargs) + else: + super().__init__([p for _, p in named], **kwargs) + + def _page_key(self, name: str): + parts = name.split(".") + try: + li = int(parts[2]) if parts[:2] == ["model", "layers"] else None + if li is None: + return None + expert_at = parts.index("experts") + eid = int(parts[expert_at + 1]) + except (ValueError, IndexError): + return None + group = self._layer_to_group.get(li, li) + return PageKey(li, group, eid, "adam") + + @torch.no_grad() + def step(self, closure=None): + # Resume expert moments lazily, one page at a time. This avoids a + # full CPU optimizer-state clone while preserving AdamW continuity. + by_page = {} + for name, p in self._param_by_name.items(): + key = self._page_key(name) + if key is None or p.grad is None: + continue + self._dirty_expert_params.add(p) + if key in self._hydrated_pages: + continue + suffix = name.split(f"experts.{key.expert}.", 1)[-1] + by_page.setdefault(key, {})[suffix] = p + for key, params in by_page.items(): + if not self.page_store.store.has(PageKey(key.layer, key.group, key.expert, "adam")): + self._hydrated_pages.add(key) + continue + loaded = self.page_store.load(key, params) + for suffix, param in params.items(): + state = self.state[param] + state["step"] = torch.tensor( + float(loaded.get("step", 0)), device=param.device + ) + state["exp_avg"] = loaded["exp_avg"][suffix].to(param.device, param.dtype) + state["exp_avg_sq"] = loaded["exp_avg_sq"][suffix].to(param.device, param.dtype) + self._hydrated_pages.add(key) + # Expert pages may have moved between CPU and accelerator residency + # since the previous microbatch. Keep live Adam tensors colocated + # with their parameter before invoking the fused/foreach update. + for param in self._dirty_expert_params: + state = self.state.get(param) + if not state or "exp_avg" not in state: + continue + if state["exp_avg"].device != param.device: + state["exp_avg"] = state["exp_avg"].to(param.device, param.dtype) + if state["exp_avg_sq"].device != param.device: + state["exp_avg_sq"] = state["exp_avg_sq"].to(param.device, param.dtype) + if torch.is_tensor(state.get("step")) and state["step"].device != param.device: + state["step"] = state["step"].to(param.device) + loss = super().step(closure) + grouped = {} + for group in self.param_groups: + for p in group["params"]: + if p not in self._dirty_expert_params: + continue + name = self._names.get(id(p), "") + key = self._page_key(name) + state = self.state.get(p) + if key is None or not state or "exp_avg" not in state: + continue + suffix = name.split(f"experts.{key.expert}.", 1)[-1] + page = grouped.setdefault(key, {"exp_avg": {}, "exp_avg_sq": {}, "step": 0}) + page["exp_avg"][suffix] = state["exp_avg"] + page["exp_avg_sq"][suffix] = state["exp_avg_sq"] + page["step"] = max(int(page["step"]), int(state.get("step", 0))) + for key, state in grouped.items(): + self.page_store.save(key, state, step=state.get("step", 0)) + self._dirty_expert_params.clear() + return loss + + def state_dict(self): + """Return the regular AdamW state plus paging metadata. + + Moments are already persisted per expert page. Omitting those entries + from the serialized PyTorch state is important: calling + ``optimizer.state_dict()`` must not create a second full CPU copy of + all expert moments, which would defeat SSD paging for large MoE + checkpoints. Non-expert state and the normal parameter-group layout + remain fully PyTorch-compatible; expert pages are hydrated lazily on + the first update after resume. + """ + payload = super().state_dict() + # ``Optimizer.state_dict`` replaces parameter objects with compact + # integer ids in param-group order. Reconstruct that mapping without + # touching tensor data and remove only expert entries. + param_ids = [] + for serialized_group, live_group in zip( + payload.get("param_groups", []), self.param_groups + ): + param_ids.extend(zip(serialized_group.get("params", []), live_group.get("params", []))) + expert_state_ids = set() + for serialized_id, param in param_ids: + name = self._names.get(id(param), "") + if self._page_key(name) is not None: + expert_state_ids.add(serialized_id) + if expert_state_ids: + payload["state"] = { + key: value for key, value in payload.get("state", {}).items() + if key not in expert_state_ids + } + payload["chronos_paging"] = { + "page_dir": self.page_store.store.root_dir, + "stats": self.page_store.store.summary(), + "expert_pages": [ + key.filename() for key in self.page_store.store._cache.keys() + if key.kind == "adam" + ], + } + return payload + + @torch.no_grad() + def evict_expert_moments(self) -> None: + """Drop accelerator/CPU Adam moments already persisted as pages. + + This is called after an optimizer step in paged training. The next + step hydrates only experts that receive gradients, so optimizer memory + follows the active working set instead of monotonically growing to all + experts. Non-expert state remains resident. + """ + expert_params = [] + for param in list(self.state): + name = self._names.get(id(param), "") + if self._page_key(name) is not None: + expert_params.append(param) + for param in expert_params: + self.state.pop(param, None) + self._hydrated_pages.clear() + self._dirty_expert_params.clear() + + def load_state_dict(self, state_dict): + # Older checkpoints may not carry the optional paging envelope. + payload = dict(state_dict) + payload.pop("chronos_paging", None) + return super().load_state_dict(payload) + + +def build_paged_optimizer(model: nn.Module, lr: float, *, page_dir: str, + weight_decay: float = 0.01, betas=(0.9, 0.95), eps=1e-8): + cfg = getattr(model, "config", None) + if cfg is None and hasattr(model, "model"): + cfg = getattr(model.model, "config", None) + prelude = int(getattr(cfg, "prelude_layers", 0) or 0) + groups = int(getattr(cfg, "recurrent_groups", 0) or 0) + layers = int(getattr(cfg, "num_hidden_layers", 0) or 0) + layer_to_group = { + li: (li - prelude if groups and prelude <= li < prelude + groups else -1) + for li in range(layers) + } + # ``page_size_mb`` is a memory budget, not merely an informational field: + # derive a bounded in-RAM page cache from the largest expert tensor set. + page_mb = max(1, int(getattr(cfg, "page_size_mb", 64) or 64)) + expert_sizes = {} + for name, param in model.named_parameters(): + if ".experts." in name: + try: + eid = int(name.split(".experts.", 1)[1].split(".", 1)[0]) + except (IndexError, ValueError): + eid = -1 + expert_sizes[eid] = expert_sizes.get(eid, 0) + int(param.numel() * 4) + expert_bytes = max(expert_sizes.values(), default=1) + page_capacity = max(1, int((page_mb * 1024 * 1024) / expert_bytes)) + # Keep a modest upper bound so an over-large page_size does not defeat + # the purpose of bounded paging on machines with ample RAM. + page_capacity = min(page_capacity, 256) + return PagedAdamW( + model.named_parameters(), page_dir=page_dir, + page_capacity=page_capacity, + layer_to_group=layer_to_group, + lr=float(lr), weight_decay=float(weight_decay), betas=betas, eps=eps, + ) + + +def build_configured_optimizer( + model: nn.Module, + config, + *, + lr: float, + weight_decay: float = 0.01, + betas=(0.9, 0.95), + eps=1e-8, +): + """Build the optimizer selected by the v2 training configuration. + + ``optimizer_paging`` is the explicit switch for SSD-persisted Adam + moments. ``training_paging`` also enables it because a paged training + run without paged optimizer state would immediately recreate a full + accelerator/CPU moment footprint at the first step. Keeping this policy + in one function prevents the six stage trainers and the WebUI from + silently diverging. + """ + paged = bool( + getattr(config, "optimizer_paging", False) + or getattr(config, "training_paging", False) + ) + if paged: + return build_paged_optimizer( + model, + lr=lr, + page_dir=str(getattr(config, "optimizer_page_dir", "./optimizer_pages")), + weight_decay=weight_decay, + betas=betas, + eps=eps, + ) + return build_optimizer( + model, + lr=lr, + weight_decay=weight_decay, + betas=betas, + eps=eps, + ) + + def get_lr( current_step: int, total_steps: int, diff --git a/chronos/trainer/orpo_trainer.py b/chronos/trainer/orpo_trainer.py index 79bdc78..7d75377 100644 --- a/chronos/trainer/orpo_trainer.py +++ b/chronos/trainer/orpo_trainer.py @@ -48,6 +48,7 @@ optimizer_step_with_scaler, runtime_summary, ) +from chronos.trainer.paged_store import PagedTrainingController def _mean_logprob(logits: torch.Tensor, labels: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: @@ -79,17 +80,24 @@ def __init__(self, model: ChronosForCausalLM, config: ChronosConfig, args, token f"device={summary.device} type={summary.device_type} dtype={summary.dtype} " f"cpu_threads={summary.cpu_threads} autocast={summary.autocast} scaler={summary.scaler}" ) - from chronos.trainer.optim_utils import build_optimizer - self.optimizer = build_optimizer( - model, lr=args.learning_rate, + from chronos.trainer.optim_utils import build_configured_optimizer + self.optimizer = build_configured_optimizer( + model, config, lr=args.learning_rate, weight_decay=float(getattr(args, "weight_decay", 0.01)), ) self._router_ref = None + # One compute dtype end to end; FP32 masters live in the optimizer. + from chronos.trainer.precision_training import apply_training_precision + apply_training_precision(self, models=()) + self.paged_training = PagedTrainingController(self.model, config) def set_calibration_batch(self, x: torch.Tensor): self._router_ref = capture_reference_routing(self.model, x, self.device) def train_step(self, batch, step, total_steps): + controller = getattr(getattr(self.model, "model", None), "recurrent_controller", None) + if controller is not None: + controller.set_training_progress(step, total_steps) x_c = batch["x_chosen"].to(self.device) x_r = batch["x_rejected"].to(self.device) y_c = batch["y_chosen"].to(self.device) @@ -104,7 +112,9 @@ def train_step(self, batch, step, total_steps): pg["lr"] = lr with self.autocast_ctx: - outputs, lookahead_probs = self.model(x) + paged_masks = self.paged_training.masks_for_forward(self.device) + result = self.model(x, available_expert_masks=paged_masks) + outputs, lookahead_probs = result.legacy_tuple() B2 = x_c.shape[0] logits_c = outputs.logits[:B2] logits_r = outputs.logits[B2:] @@ -165,6 +175,7 @@ def _log_odds(lp): self.model.parameters(), self.args.grad_clip, ) + self.paged_training.observe() return ( loss.item() * self.args.accumulation_steps, diff --git a/chronos/trainer/paged_store.py b/chronos/trainer/paged_store.py new file mode 100644 index 0000000..2056727 --- /dev/null +++ b/chronos/trainer/paged_store.py @@ -0,0 +1,337 @@ +"""Crash-safe SSD paging primitives for expert training and optimizer state. + +This module intentionally has no dependency on a particular accelerator. It +provides atomic tensor pages, bounded in-memory LRU, and an async prefetch +interface that PyTorch/MPS/CUDA trainers can share. +""" +from __future__ import annotations + +import collections +import os +import tempfile +import threading +from dataclasses import dataclass +from typing import Dict, Iterable, Optional + +import torch + + +class PagedTrainingController: + """Route-aware expert residency controller for training. + + The controller deliberately keeps scheduling separate from the optimizer: + optimizer moments are persisted by :class:`PagedOptimizerState`, while + this object selects a bounded working set for the *next* microbatch from + the previous batch's router probabilities. The first batch is all + resident (stable bootstrap); subsequent batches pass per-layer masks to + Chronos' shared-fallback MoE path. This gives a real bounded activation + and expert-compute path without cloning a full ``state_dict`` on CPU. + """ + + def __init__(self, model, config): + self.model = model + self.config = config + self.enabled = bool(getattr(config, "training_paging", False)) + self.num_experts = int(getattr(config, "num_experts", 0) or 0) + self.num_layers = int(getattr(config, "num_hidden_layers", 0) or 0) + budget = getattr(config, "offload_working_set_experts", None) + if budget is None: + budget = getattr(config, "recommended_resident_experts", None) + if budget is None: + budget = max(1, self.num_experts // 2) if self.num_experts else 0 + self.active_budget = max(1, min(int(budget), self.num_experts)) if self.num_experts else 0 + self._masks = None + try: + self.device = next(model.parameters()).device + except StopIteration: + self.device = torch.device("cpu") + self.weight_store = None + if self.enabled: + page_root = os.path.join( + str(getattr(config, "optimizer_page_dir", "./optimizer_pages")), + "expert_weights", + ) + page_mb = max(1, int(getattr(config, "page_size_mb", 64) or 64)) + # Weight pages keep the parameter's own (compute) dtype; the FP32 + # truth lives in the optimizer masters (precision_training), so an + # FP32 copy here would only double SSD traffic. + self.weight_store = TieredTensorStore( + page_root, capacity=max(1, min(64, page_mb // 4)), dtype=None + ) + # Checkpoint helpers can materialize paged weights without importing a + # specific trainer class. + setattr(model, "_chronos_paged_training_controller", self) + self._step = 0 + self._stats = {"paged_steps": 0, "active_experts": 0, + "route_tokens": 0, "cold_route_tokens": 0, + "fallback_weight_mass": 0.0} + + def masks_for_forward(self, device=None, *, allow_evict: bool = True): + if not self.enabled or self.num_experts <= 0: + return None + if self._masks is None: + # Bootstrap with a full mask so router probabilities are learned + # before any expert is made cold. + mask = torch.ones(self.num_experts, dtype=torch.bool, device=device) + masks = [mask.clone() for _ in range(self.num_layers)] + else: + masks = [mask.to(device=device) if device is not None else mask + for mask in self._masks] + self._apply_residency(masks, allow_evict=allow_evict) + return masks + + @torch.no_grad() + def _apply_residency(self, masks, *, allow_evict: bool = True): + """Move cold expert parameter pages off accelerator memory. + + Expert modules stay registered (optimizer references remain valid), + but their ``Parameter.data`` lives on CPU while the route mask ensures + they are not executed. Active pages are brought back to the training + device just before forward. CPU runs are naturally a no-op. + """ + if not self.enabled or self.device.type == "cpu": + return + backbone = getattr(self.model, "model", self.model) + for layer_idx, layer in enumerate(getattr(backbone, "layers", [])): + moe = getattr(layer, "mlp", None) + experts = getattr(moe, "experts", None) + if experts is None or layer_idx >= len(masks): + continue + mask = masks[layer_idx].detach().to("cpu", dtype=torch.bool) + for eid, expert in enumerate(experts): + active = eid < mask.numel() and bool(mask[eid]) + # Do not evict a page whose accumulated gradient has not yet + # been consumed by optimizer.step(). This preserves gradient + # accumulation semantics while still shrinking residency at + # update boundaries. + has_pending_grad = any(p.grad is not None for p in expert.parameters()) + target = ( + self.device + if active or has_pending_grad or not allow_evict + else torch.device("cpu") + ) + params = dict(expert.named_parameters()) + if not params: + continue + group = self._layer_group(layer_idx) + key = PageKey(layer_idx, group, eid, "weight") + materialized = all(param.numel() > 0 for param in params.values()) + if target != torch.device("cpu") and not materialized: + if self.weight_store is None: + raise RuntimeError("training weight paging store is not initialized") + state = self.weight_store.get(key, device=target) + if state is None: + raise RuntimeError(f"missing paged expert weight: {key.filename()}") + for name, param in params.items(): + if name not in state: + raise RuntimeError(f"missing tensor {name} in {key.filename()}") + param.data = state[name].to(device=target, dtype=param.dtype) + elif target == torch.device("cpu") and allow_evict and materialized: + # Persist updated weights before releasing accelerator and + # host storage. Empty registered Parameters keep optimizer + # identity stable while consuming no model-sized memory. + if self.weight_store is None: + raise RuntimeError("training weight paging store is not initialized") + self.weight_store.put(key, {name: p.data for name, p in params.items()}) + for param in params.values(): + param.data = torch.empty(0, device="cpu", dtype=param.dtype) + + def _layer_group(self, layer_idx: int) -> int: + prelude = int(getattr(self.config, "prelude_layers", 0) or 0) + groups = int(getattr(self.config, "recurrent_groups", 0) or 0) + return layer_idx - prelude if prelude <= layer_idx < prelude + groups else -1 + + @torch.no_grad() + def materialize_all(self) -> None: + """Restore every SSD expert page for full checkpoints/evaluation.""" + if not self.enabled: + return + masks = [torch.ones(self.num_experts, dtype=torch.bool, device=self.device) + for _ in range(self.num_layers)] + self._apply_residency(masks, allow_evict=False) + + @torch.no_grad() + def observe(self): + if not self.enabled or self.num_experts <= 0: + return + layers = [] + active_total = 0 + cold_mass = 0.0 + for layer_idx, layer in enumerate(getattr(getattr(self.model, "model", self.model), "layers", [])): + moe = getattr(layer, "mlp", None) + probs = getattr(moe, "last_router_probs", None) + if probs is None: + layers.append(torch.ones(self.num_experts, dtype=torch.bool)) + continue + scores = probs.float().reshape(-1, self.num_experts).mean(dim=0) + k = min(self.active_budget, self.num_experts) + ids = torch.topk(scores, k=k, dim=-1).indices + mask = torch.zeros(self.num_experts, dtype=torch.bool) + mask[ids.cpu()] = True + layers.append(mask) + active_total += int(mask.sum()) + self._stats["route_tokens"] += int(probs.shape[0] * probs.shape[1]) + if self._masks is not None and layer_idx < len(self._masks): + previous = self._masks[layer_idx] + top_ids = torch.topk(probs.float().reshape(-1, self.num_experts), + k=min(int(getattr(self.config, "num_experts_per_tok", 1)), self.num_experts), + dim=-1).indices + cold_mass += float((~previous[top_ids]).float().mean().item()) + if layers: + self._masks = layers + self._stats["active_experts"] = active_total + self._stats["paged_steps"] += 1 + self._step += 1 + self._stats["cold_route_tokens"] = int(cold_mass * max(self._stats["route_tokens"], 1)) + self._stats["fallback_weight_mass"] = cold_mass + try: + from chronos.runtime.metrics import safe_record + safe_record("training_active_experts", active_total) + safe_record("training_cold_route_fraction", cold_mass / max(len(layers), 1)) + except Exception: + pass + + def stats(self) -> dict: + return {**self._stats, "enabled": self.enabled, + "active_budget": self.active_budget, + "weight_store": self.weight_store.summary() if self.weight_store else None} + + +@dataclass(frozen=True) +class PageKey: + layer: int + group: int + expert: int + kind: str = "weight" + + def filename(self) -> str: + return f"{self.kind}_l{self.layer}_g{self.group}_e{self.expert}.pt" + + +class TieredTensorStore: + """Bounded CPU page cache with atomic SSD persistence.""" + + def __init__(self, root_dir: str, capacity: int = 8, *, dtype=None): + """``dtype=None`` keeps each tensor's dtype; otherwise convert once on put.""" + self.root_dir = os.path.abspath(root_dir) + os.makedirs(self.root_dir, exist_ok=True) + self.capacity = max(1, int(capacity)) + self.dtype = dtype + self._cache: "collections.OrderedDict[PageKey, dict]" = collections.OrderedDict() + self._lock = threading.RLock() + self.stats = { + "loads": 0, "stores": 0, "hits": 0, "misses": 0, + "errors": 0, "evictions": 0, + } + + def path(self, key: PageKey) -> str: + return os.path.join(self.root_dir, key.filename()) + + def has(self, key: PageKey) -> bool: + with self._lock: + return key in self._cache or os.path.exists(self.path(key)) + + def put(self, key: PageKey, state: Dict[str, torch.Tensor], *, persist: bool = True) -> None: + normalized = { + str(name): tensor.detach().to(device="cpu", dtype=self.dtype or tensor.dtype).contiguous() + for name, tensor in state.items() + } + with self._lock: + self._cache[key] = normalized + self._cache.move_to_end(key) + while len(self._cache) > self.capacity: + self._cache.popitem(last=False) + self.stats["evictions"] += 1 + if persist: + self._atomic_save(key, normalized) + self.stats["stores"] += 1 + + def get(self, key: PageKey, *, device: Optional[torch.device] = None) -> Optional[Dict[str, torch.Tensor]]: + with self._lock: + if key in self._cache: + self._cache.move_to_end(key) + self.stats["hits"] += 1 + state = self._cache[key] + else: + state = None + if state is None: + path = self.path(key) + if not os.path.exists(path): + self.stats["misses"] += 1 + return None + try: + state = torch.load(path, map_location="cpu", weights_only=True) + self.stats["loads"] += 1 + with self._lock: + self._cache[key] = state + self._cache.move_to_end(key) + while len(self._cache) > self.capacity: + self._cache.popitem(last=False) + self.stats["evictions"] += 1 + except Exception: + self.stats["errors"] += 1 + return None + if device is not None: + return {name: value.to(device=device, non_blocking=True) for name, value in state.items()} + return dict(state) + + def _atomic_save(self, key: PageKey, state: dict) -> None: + target = self.path(key) + fd, tmp = tempfile.mkstemp(prefix=".chronos-page-", suffix=".pt", dir=self.root_dir) + os.close(fd) + try: + torch.save(state, tmp) + os.replace(tmp, target) + except Exception: + self.stats["errors"] += 1 + try: + os.unlink(tmp) + except OSError: + pass + raise + + def prefetch(self, keys: Iterable[PageKey], device=None) -> Dict[PageKey, Optional[dict]]: + return {key: self.get(key, device=device) for key in keys} + + def clear_memory(self) -> None: + with self._lock: + self._cache.clear() + + def summary(self) -> dict: + with self._lock: + resident = len(self._cache) + return {**self.stats, "resident_pages": resident, "capacity": self.capacity, + "root_dir": self.root_dir} + + +class PagedOptimizerState: + """Adam moments persisted per expert page.""" + + def __init__(self, root_dir: str, capacity: int = 8): + self.store = TieredTensorStore(root_dir, capacity=capacity, dtype=torch.float32) + + def load(self, key: PageKey, params: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, torch.Tensor]]: + payload = self.store.get(PageKey(key.layer, key.group, key.expert, "adam")) + if payload is None: + return { + "exp_avg": {name: torch.zeros_like(value, dtype=torch.float32) for name, value in params.items()}, + "exp_avg_sq": {name: torch.zeros_like(value, dtype=torch.float32) for name, value in params.items()}, + } + names = list(params) + loaded = { + "exp_avg": {name: payload.get(f"avg.{name}", torch.zeros_like(params[name], dtype=torch.float32)) for name in names}, + "exp_avg_sq": {name: payload.get(f"sq.{name}", torch.zeros_like(params[name], dtype=torch.float32)) for name in names}, + } + step = payload.get("meta.step") + if step is not None: + loaded["step"] = int(step.item()) + return loaded + + def save(self, key: PageKey, state: Dict[str, Dict[str, torch.Tensor]], *, step: int = 0) -> None: + payload = {} + for name, value in state.get("exp_avg", {}).items(): + payload[f"avg.{name}"] = value + for name, value in state.get("exp_avg_sq", {}).items(): + payload[f"sq.{name}"] = value + payload["meta.step"] = torch.tensor(int(step), dtype=torch.int64) + self.store.put(PageKey(key.layer, key.group, key.expert, "adam"), payload) diff --git a/chronos/trainer/precision_training.py b/chronos/trainer/precision_training.py new file mode 100644 index 0000000..69b05f7 --- /dev/null +++ b/chronos/trainer/precision_training.py @@ -0,0 +1,170 @@ +"""Torch training under the single-compute-dtype policy (CPU / CUDA / MPS). + +The model's parameters, activations and gradients live in the compute dtype +(``chronos.precision.resolve_torch``), so forward/backward run without the +per-op weight casts that ``torch.autocast`` performs on every call. The +optimizer owns FP32 masters, mirroring ``MixedPrecisionAdamW`` on MLX: + +* each low-precision gradient is folded into an FP32 master gradient the + moment autograd finishes it (``register_post_accumulate_grad_hook``); the + model gradient is freed, so accumulation over micro-batches is exact FP32 + and no extra low-precision grad buffer stays alive; +* AdamW updates the FP32 masters (moments are FP32 as well); +* the masters are copied back to the compute dtype once per update. + +FP16 needs loss scaling (its range is too small for gradients); the scaler +operates on the FP32 master gradients. BF16/FP32 need none. A model that +is already FP32 gets no masters at all (zero overhead). +""" +from __future__ import annotations + +import torch + +from chronos.precision import PrecisionPolicy, resolve_torch + + +class _MasterWeights: + """FP32 master copies for low-precision parameters of one model.""" + + def __init__(self, params): + self.pairs: list[tuple[torch.nn.Parameter, torch.Tensor]] = [] + self.master_of: dict[int, torch.Tensor] = {} + self._handles = [] + for p in params: + if not p.requires_grad or p.dtype == torch.float32 or id(p) in self.master_of: + continue + master = p.detach().float().clone() + master.requires_grad_(True) + self.master_of[id(p)] = master + self.pairs.append((p, master)) + self._handles.append(p.register_post_accumulate_grad_hook(self._fold)) + + def _fold(self, p: torch.nn.Parameter) -> None: + master = self.master_of[id(p)] + if p.grad is None: + return + if master.shape != p.shape: # paged-out expert (zero-sized data): no update + p.grad = None + return + if master.grad is None: + master.grad = p.grad.float() + else: + master.grad.add_(p.grad) # FP32 += low precision (promoted in-kernel) + p.grad = None + + @torch.no_grad() + def sync_to_model(self) -> None: + for p, master in self.pairs: + if master.shape == p.shape: + p.copy_(master) # one rounding per update + + @torch.no_grad() + def sync_from_model(self) -> None: + """Refresh masters after the model weights were replaced (resume).""" + for p, master in self.pairs: + if master.shape == p.shape: + master.copy_(p) + + def remove(self) -> None: + for h in self._handles: + h.remove() + self._handles.clear() + + +def _swap_params_to_masters(optimizer: torch.optim.Optimizer, masters: _MasterWeights) -> None: + """Point every optimizer group at the FP32 masters (groups/decay kept).""" + if not masters.pairs: + return + for group in optimizer.param_groups: + group["params"] = [masters.master_of.get(id(p), p) for p in group["params"]] + # PagedAdamW resolves names -> params for its SSD moment pages. + by_name = getattr(optimizer, "_param_by_name", None) + if isinstance(by_name, dict): + names = getattr(optimizer, "_names", {}) + new_names = {} + for name, p in list(by_name.items()): + m = masters.master_of.get(id(p), p) + by_name[name] = m + new_names[id(m)] = name + if isinstance(names, dict): + names.update(new_names) + + +def prepare_torch_training(model, optimizer_factory, device, requested="auto"): + """Place ``model`` in the policy dtype and build an FP32-master optimizer. + + ``optimizer_factory(model)`` builds the usual optimizer (decay groups, + paging) from the *model's* parameters; they are then replaced by FP32 + masters. Returns ``(model, optimizer, policy)``; ``optimizer`` exposes + ``chronos_masters`` for :func:`optimizer_step`. + """ + policy: PrecisionPolicy = resolve_torch(requested, device) + model = model.to(device=device, dtype=policy.torch_dtype()) + optimizer = optimizer_factory(model) + masters = _MasterWeights(model.parameters()) + _swap_params_to_masters(optimizer, masters) + optimizer.chronos_masters = masters + model._chronos_precision = policy + return model, optimizer, policy + + +def grad_scaler_for(policy: PrecisionPolicy, device): + """Loss scaling only for FP16 compute (on any backend that runs FP16).""" + dev = str(device).split(":")[0] + enabled = policy.compute == "float16" and dev in ("cuda", "mps", "cpu", "xpu") + return torch.amp.GradScaler(dev, enabled=enabled) + + +def optimizer_step(scaler, optimizer, model, grad_clip: float) -> float: + """Unscale -> clip -> step on FP32 masters -> copy back. Returns grad norm.""" + masters = getattr(optimizer, "chronos_masters", None) + params = [p for g in optimizer.param_groups for p in g["params"] if p.grad is not None] + scaler.unscale_(optimizer) + norm = torch.nn.utils.clip_grad_norm_(params, grad_clip) if params else torch.zeros(()) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + model.zero_grad(set_to_none=True) + if masters is not None: + masters.sync_to_model() + evict = getattr(optimizer, "evict_expert_moments", None) + if callable(evict): + evict() + return float(norm) + + +def apply_training_precision(trainer, *, models=(), requested=None) -> PrecisionPolicy: + """Retrofit a stage trainer (``self.model`` + ``self.optimizer`` built) + to the single-dtype policy. + + Moves the trainable model and any frozen companions (reference/teacher in + ``models``) to the compute dtype, swaps the optimizer onto FP32 masters, + disables autocast (the model is native) and picks the FP16-only scaler. + Called with the trainer's ``args.dtype``; FP32 is a no-op apart from + recording the policy. + """ + requested = requested if requested is not None else getattr(trainer.args, "dtype", "auto") + policy, trainer.autocast_ctx, trainer.scaler = attach_precision( + trainer.model, trainer.optimizer, trainer.device, requested, models=models) + trainer.precision = policy + return policy + + +def attach_precision(model, optimizer, device, requested="auto", *, models=()): + """Functional form of :func:`apply_training_precision` for custom loops. + + Returns ``(policy, autocast_ctx, scaler)``. Must run before the first + optimizer step (optimizer state is keyed by the FP32 masters). + """ + import contextlib + policy = resolve_torch(requested, device) + dtype = policy.torch_dtype() + model.to(device=device, dtype=dtype) + for m in models: + if m is not None: + m.to(device=device, dtype=dtype) + masters = _MasterWeights(model.parameters()) + _swap_params_to_masters(optimizer, masters) + optimizer.chronos_masters = masters + model._chronos_precision = policy + return policy, contextlib.nullcontext(), grad_scaler_for(policy, device) diff --git a/chronos/trainer/sft_trainer.py b/chronos/trainer/sft_trainer.py index b854126..537bbb9 100644 --- a/chronos/trainer/sft_trainer.py +++ b/chronos/trainer/sft_trainer.py @@ -28,7 +28,9 @@ import chronos.deps # ensure minimind on sys.path from trainer.trainer_utils import Logger, is_main_process # type: ignore -from chronos.data.flexible_dataset import StreamingSFTDataset +from chronos.data.flexible_dataset import ( + StreamingSFTDataset, MultimodalStreamingSFTDataset, multimodal_collate, +) from chronos.model.config import ChronosConfig from chronos.model.model_chronos import ChronosForCausalLM @@ -49,6 +51,7 @@ optimizer_step_with_scaler, runtime_summary, ) +from chronos.trainer.paged_store import PagedTrainingController class ChronosSFTTrainer: @@ -73,32 +76,51 @@ def __init__(self, model: ChronosForCausalLM, config: ChronosConfig, args, token f"device={summary.device} type={summary.device_type} dtype={summary.dtype} " f"cpu_threads={summary.cpu_threads} autocast={summary.autocast} scaler={summary.scaler}" ) - from chronos.trainer.optim_utils import build_optimizer - self.optimizer = build_optimizer( - model, lr=args.learning_rate, + from chronos.trainer.optim_utils import build_configured_optimizer + self.optimizer = build_configured_optimizer( + model, config, lr=args.learning_rate, weight_decay=float(getattr(args, "weight_decay", 0.01)), ) + # One compute dtype end to end; FP32 masters live in the optimizer. + from chronos.trainer.precision_training import apply_training_precision + apply_training_precision(self, models=()) + self.paged_training = PagedTrainingController(self.model, config) # Router anchor reference — captured at the start of a stage self._router_ref = None self._calib_batch = None - def set_calibration_batch(self, input_ids: torch.Tensor): + def set_calibration_batch(self, input_ids: torch.Tensor, *, modality_inputs=None, modality_mask=None): """Provide a representative input_ids batch; at stage start the current (pre-SFT) model's router is frozen as the anchor teacher.""" self._calib_batch = input_ids - self._router_ref = capture_reference_routing(self.model, input_ids, self.device) + self._router_ref = capture_reference_routing( + self.model, input_ids, self.device, + modality_inputs=modality_inputs, modality_mask=modality_mask, + ) - def train_step(self, input_ids, labels, step, total_steps): + def train_step(self, input_ids, labels, step, total_steps, *, modality_inputs=None, modality_mask=None): + if hasattr(self.model.model, "recurrent_controller"): + self.model.model.recurrent_controller.set_training_progress(step, total_steps) input_ids = input_ids.to(self.device) labels = labels.to(self.device) + if modality_inputs: + modality_inputs = {k: v.to(self.device) for k, v in modality_inputs.items()} + if modality_mask is not None: + modality_mask = modality_mask.to(self.device) lr = get_lr(step, total_steps, self.args.learning_rate) for pg in self.optimizer.param_groups: pg['lr'] = lr with self.autocast_ctx: - outputs, lookahead_probs = self.model(input_ids, labels=labels) + paged_masks = self.paged_training.masks_for_forward(self.device) + result = self.model( + input_ids, labels=labels, + modality_inputs=modality_inputs, modality_mask=modality_mask, + available_expert_masks=paged_masks, + ) + outputs, lookahead_probs = result.legacy_tuple() ce_loss = outputs.loss aux_loss = outputs.aux_loss @@ -148,6 +170,7 @@ def train_step(self, input_ids, labels, step, total_steps): self.model.parameters(), self.args.grad_clip, ) + self.paged_training.observe() return ( loss.item() * self.args.accumulation_steps, @@ -163,14 +186,21 @@ def train_epoch(self, epoch, loader, iters, start_step=0, wandb=None, max_steps= cap = int(max_steps) if max_steps is not None else None total_steps = cap or (self.args.epochs * iters) - for local_step, (input_ids, labels) in enumerate(loader, start=1): + for local_step, batch in enumerate(loader, start=1): if local_step > iters: break step = int(start_step) + local_step if cap is not None and step > cap: 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 loss, ce, aux, la, anc = self.train_step( input_ids, labels, step, total_steps, + modality_inputs=media, modality_mask=media_mask, ) if local_step % self.args.log_interval == 0 or local_step == iters: elapsed = time.time() - start_time @@ -201,10 +231,20 @@ def _save(self, epoch, step): self.model.train() -def build_sft_loader(data_path, tokenizer, max_seq_len, batch_size, device="cpu", num_workers="auto"): - ds = StreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len) +def build_sft_loader(data_path, tokenizer, max_seq_len, batch_size, device="cpu", num_workers="auto", *, + modalities=None, image_size=224, audio_sample_rate=16000, + max_audio_seconds=30.0, media_error_policy="warn"): + enabled = {str(m).lower() for m in (modalities or []) if str(m).lower() != "text"} + ds_cls = MultimodalStreamingSFTDataset if enabled else StreamingSFTDataset + ds = ds_cls( + data_path, tokenizer, max_length=max_seq_len, + **({"image_size": image_size, "audio_sample_rate": audio_sample_rate, + "max_audio_seconds": max_audio_seconds, "media_error_policy": media_error_policy} + if enabled else {}), + ) return DataLoader( ds, batch_size=batch_size, + collate_fn=multimodal_collate if enabled else None, **dataloader_kwargs(device=device, num_workers=num_workers, shuffle=True), ) diff --git a/chronos/trainer/stage_utils.py b/chronos/trainer/stage_utils.py index d1ae15b..d4305e0 100644 --- a/chronos/trainer/stage_utils.py +++ b/chronos/trainer/stage_utils.py @@ -1,6 +1,7 @@ """Shared helpers for standalone Chronos stage scripts.""" from __future__ import annotations +import argparse import os from chronos.model.checkpoint import ( @@ -8,8 +9,10 @@ load_checkpoint_state_dict, load_state_dict_controlled, resolve_checkpoint_path, + resolve_checkpoint_tokenizer, ) from chronos.model.config import ChronosConfig +from chronos.config_schema import MODEL_FIELDS TOPOLOGY_ARG_FIELDS = ( @@ -25,6 +28,20 @@ "rope_dim", "kv_latent_dim", "vocab_size", + "recurrent_groups", + "recurrent_depth", + "min_recurrent_depth", + "prelude_layers", + "coda_layers", + "max_inference_depth", + "halting_threshold", + "adaptive_halting", + "recurrent_latent_refinement", + "training_paging", + "optimizer_paging", + "optimizer_page_dir", + "page_size_mb", + "prefetch_depth", ) @@ -48,11 +65,60 @@ def add_topology_args(parser, *, defaults: bool = False): parser.add_argument("--rope_dim", type=int, default=(32 if defaults else None)) parser.add_argument("--kv_latent_dim", type=int, default=(64 if defaults else None)) parser.add_argument("--vocab_size", type=int, default=(6400 if defaults else None)) + parser.add_argument("--tokenizer_path", default=None, + help="Tokenizer directory (tokenizer.json). Default: the checkpoint's own tokenizer " + "(/tokenizer or its config metadata), else minimind.") + parser.add_argument("--recurrent_groups", type=int, default=(2 if defaults else None)) + parser.add_argument("--recurrent_depth", type=int, default=(4 if defaults else None)) + parser.add_argument("--min_recurrent_depth", type=int, default=(2 if defaults else None)) + parser.add_argument("--prelude_layers", type=int, default=(1 if defaults else None)) + parser.add_argument("--coda_layers", type=int, default=(1 if defaults else None)) + parser.add_argument("--max_inference_depth", type=int, default=(4 if defaults else None)) + parser.add_argument("--halting_threshold", type=float, default=(0.99 if defaults else None)) + parser.add_argument( + "--adaptive_halting", action=argparse.BooleanOptionalAction, + default=(True if defaults else None), + ) + parser.add_argument( + "--recurrent_latent_refinement", action=argparse.BooleanOptionalAction, + default=(True if defaults else None), + ) + parser.add_argument( + "--training_paging", action=argparse.BooleanOptionalAction, + default=(False if defaults else None), + ) + parser.add_argument( + "--optimizer_paging", action=argparse.BooleanOptionalAction, + default=(False if defaults else None), + ) + parser.add_argument("--optimizer_page_dir", default=("./optimizer_pages" if defaults else None)) + parser.add_argument("--page_size_mb", type=int, default=(64 if defaults else None)) + parser.add_argument("--prefetch_depth", type=int, default=(2 if defaults else None)) + # Additional kernel settings use one typed registry for CLI registration; + # checkpoint-backed stages leave omitted values unset. + fields = { + "recurrent_strategy": (str, "scheduled"), + "recurrent_depth_jitter": (float, 0.), + "depth_schedule": (str, "fixed_then_adaptive"), + "depth_schedule_warmup_fraction": (float, .25), + "ponder_loss_coef": (float, .01), + "recurrent_distill_coef": (float, .05), + "recurrent_stability_coef": (float, .001), + "recurrent_stability_threshold": (float, .1), + "recurrent_dt_init": (float, .25), + "recurrent_dt_min": (float, .05), + "recurrent_dt_max": (float, 1.), + } + for name, (kind, value) in fields.items(): + parser.add_argument("--"+name, type=kind, default=value if defaults else None) + for name, value in (("recurrent_enabled", True), ("layer_aware_lookahead", True), + ("layer_aware_cache", True), ("activation_checkpointing", False)): + parser.add_argument("--"+name, action=argparse.BooleanOptionalAction, default=value if defaults else None) def topology_overrides_from_args(args) -> dict: out = {} - for key in TOPOLOGY_ARG_FIELDS: + for key in MODEL_FIELDS: value = getattr(args, key, None) if value is None: continue @@ -74,6 +140,19 @@ def build_pretrain_config(args) -> ChronosConfig: "lambda_router_locality": getattr(args, "lambda_router_locality", 0.02), "vram_budget_gb": getattr(args, "vram_budget_gb", 4.0), "use_moe": True, + "architecture_version": 2, + "recurrent_enabled": getattr(args, "recurrent_enabled", True), + "adaptive_halting": getattr(args, "adaptive_halting", True), + "halting_threshold": getattr(args, "halting_threshold", 0.99), + "max_inference_depth": getattr(args, "max_inference_depth", None), + "recurrent_latent_refinement": getattr(args, "recurrent_latent_refinement", True), + "layer_aware_lookahead": getattr(args, "layer_aware_lookahead", True), + "layer_aware_cache": getattr(args, "layer_aware_cache", True), + "training_paging": getattr(args, "training_paging", False), + "optimizer_paging": getattr(args, "optimizer_paging", False), + "optimizer_page_dir": getattr(args, "optimizer_page_dir", "./optimizer_pages"), + "page_size_mb": getattr(args, "page_size_mb", 64), + "prefetch_depth": getattr(args, "prefetch_depth", 2), }) fallback_mask_prob = getattr(args, "fallback_mask_prob", None) if fallback_mask_prob is not None: @@ -134,3 +213,25 @@ def load_required_checkpoint(model, checkpoint_path: str, device: str) -> dict: state = load_checkpoint_state_dict(checkpoint_path, map_location="cpu") load_state_dict_controlled(model, state) return state + + +def load_stage_tokenizer(args, config, checkpoint_path: str | None = None): + """Load the tokenizer that matches ``checkpoint_path`` (or ``--tokenizer_path``). + + Raises an error when the tokenizer has more entries than ``config.vocab_size``. + A larger tokenizer would emit ids the embedding cannot hold. The minimind + tokenizer (6400 entries) used on a 16k-vocabulary checkpoint fails the other + way: it trains silently on the wrong ids. + """ + from transformers import AutoTokenizer + path = resolve_checkpoint_tokenizer(checkpoint_path, getattr(args, "tokenizer_path", None)) + tokenizer = AutoTokenizer.from_pretrained(path) + vocab = int(getattr(config, "vocab_size", 0) or 0) + if vocab and len(tokenizer) > vocab: + raise ValueError(f"tokenizer at {path} has {len(tokenizer)} entries but the model vocab_size is {vocab}") + if checkpoint_path and vocab and len(tokenizer) < vocab // 2: + raise ValueError( + f"tokenizer at {path} has {len(tokenizer)} entries for a {vocab}-entry checkpoint; " + "pass the checkpoint's own tokenizer with --tokenizer_path") + print(f"[tokenizer] {path} ({len(tokenizer)} entries, model vocab {vocab})") + return tokenizer diff --git a/chronos/tuning/chronos_auto_tuner.py b/chronos/tuning/chronos_auto_tuner.py index 5edf0da..efa4ad5 100644 --- a/chronos/tuning/chronos_auto_tuner.py +++ b/chronos/tuning/chronos_auto_tuner.py @@ -197,7 +197,7 @@ def _run_probe_trial( losses = [] for step in range(probe_steps): ids = torch.randint(0, vocab_size, (batch, seq_len), device=device) - out, _ = model(ids, labels=ids) + out, _ = model(ids, labels=ids, legacy_output=True) moe_layers = [l.mlp for l in model.model.layers if isinstance(l.mlp, ChronosMOEFeedForward)] if moe_layers and moe_layers[0].last_router_probs is not None: diff --git a/chronos/verify.py b/chronos/verify.py index a37ee03..418b00f 100644 --- a/chronos/verify.py +++ b/chronos/verify.py @@ -25,7 +25,7 @@ def _torch_prefill(config, checkpoint_path: str, ids: torch.Tensor, device: str load_state_dict_controlled(model, state) model = model.to(device).eval() with torch.no_grad(): - out, _lookahead = model(ids.to(device), use_cache=False) + out, _lookahead = model(ids.to(device), use_cache=False, legacy_output=True) logits = out.logits[:, -1, :].detach().float().cpu() return { "logits": logits, @@ -44,8 +44,8 @@ def _masked_drift(config, checkpoint_path: str, ids: torch.Tensor) -> dict[str, for _ in range(int(config.num_hidden_layers)) ] with torch.no_grad(): - no_mask, _ = model(ids, use_cache=False) - all_avail, _ = model(ids, use_cache=False, available_expert_masks=masks) + no_mask, _ = model(ids, use_cache=False, legacy_output=True) + all_avail, _ = model(ids, use_cache=False, available_expert_masks=masks, legacy_output=True) a = no_mask.logits[:, -1, :].float() b = all_avail.logits[:, -1, :].float() return { @@ -62,7 +62,7 @@ def _mlx_prefill(config, checkpoint_path: str, ids: torch.Tensor) -> dict[str, A state = load_checkpoint_state_dict(checkpoint_path, map_location="cpu") pt = ChronosForCausalLM(config) load_state_dict_controlled(pt, state) - model = ChronosMLXModel.from_chronos_pytorch(pt, config) + model = ChronosMLXModel.from_chronos_pytorch(pt.eval(), config) logits, _lookahead, _cache = model(mx.array(ids.numpy()).astype(mx.int32)) mx.eval(logits) last = torch.from_numpy(__import__("numpy").array(logits[:, -1, :], dtype="float32")) diff --git a/pyproject.toml b/pyproject.toml index 6d0a940..99ad720 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,35 +4,44 @@ build-backend = "setuptools.build_meta" [project] name = "Project_Chronos" -version = "0.2.1" -description = "On-device low-latency lookahead dual-layer MoE inference — Lookahead Routing + Async Prefetch + Hybrid Attention" +version = "0.3.0" +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" requires-python = ">=3.10" -keywords = ["llm", "moe", "inference", "lookahead-routing", "hybrid-attention", "edge-ai"] +keywords = ["llm", "moe", "inference", "lookahead-routing", "hybrid-attention", "recurrent-depth", "speculative-decoding", "mlx", "edge-ai"] classifiers = [ "Development Status :: 3 - Alpha", + "Operating System :: OS Independent", "Intended Audience :: Science/Research", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "torch>=2.4.0", + "torch>=2.6.0", # transformers 5 disables its PyTorch backend below 2.6 "transformers>=4.40.0", "datasets>=2.18.0", "safetensors>=0.4.0", "optuna>=3.6.0", "numpy>=1.24.0", "psutil>=5.9.0", + "tokenizers>=0.19.0", "tqdm>=4.66.0", - "gradio>=4.0.0", + "gradio>=5.0,<6", # Web UI tested on 5.35 (gr.Timer, LinePlot color_map) ] [project.optional-dependencies] cluster = ["python-louvain>=0.16", "networkx>=3.0"] -mlx = ["mlx>=0.12.0"] -dev = ["pytest>=7.0", "black>=24.0", "ruff>=0.4.0"] +mlx = ["mlx>=0.32.4"] # gather_mm sparse MoE: 0.31.x is broken on M5 GPUs +serve = ["fastapi==0.115.9", "uvicorn==0.34.3"] # chronos serve / chronos monitor +translate = ["sacrebleu==2.6.0", "sacremoses==0.2.0"] # corpus detok + BLEU/chrF probes +dev = ["pytest>=7.0", "black>=24.0", "ruff>=0.4.0", "httpx>=0.27", "fastapi==0.115.9", "uvicorn==0.34.3"] + +[project.urls] +Homepage = "https://github.com/FonaTech/Project_Chronos" +Issues = "https://github.com/FonaTech/Project_Chronos/issues" [project.scripts] chronos = "chronos.cli:main" @@ -56,6 +65,10 @@ py-modules = [ where = ["."] include = ["chronos*", "ui*"] +[tool.setuptools.package-data] +ui = ["event_contract.js", "presets_v3.json"] +"chronos.monitor" = ["static/*.html", "static/*.js", "static/*.css"] + [tool.black] line-length = 100 target-version = ["py310"] diff --git a/reports/speculative/mlx_gpu_d2_s13996.json b/reports/speculative/mlx_gpu_d2_s13996.json new file mode 100644 index 0000000..73be40a --- /dev/null +++ b/reports/speculative/mlx_gpu_d2_s13996.json @@ -0,0 +1,834 @@ +{ + "label": "measured", + "time": "2026-11-06T19:29:10", + "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": 24, + "max_new": 128, + "configs": "mlx-full,mlx-lazy", + "mlx_device": "gpu", + "torch_device": "mps", + "mlx_dtype": "auto", + "only": "", + "ratio": 1.0, + "miss_policy": "sync_on_demand", + "draft_depth": 2, + "gamma": 4, + "warmup": 2, + "out": "reports/speculative/mlx_gpu_d2_s13996.json", + "note": "72h MLX training paused with SIGSTOP; draft_depth 2; FLORES+ dev first 24 en sentences", + "speculative": { + "draft_depth": 2, + "gamma": 4, + "gamma_max": 8, + "adaptive_gamma": true + }, + "eos_id": 0, + "decoding": "greedy, batch 1", + "n_prompts": 24, + "prompt_tokens_mean": 39.7, + "order": "plain/speculative alternate per prompt", + "prompts_source": "FLORES+ dev (eng_Latn/zho_Hans) rows 0-31, field 'en'" + }, + "model": { + "hidden_size": 384, + "num_hidden_layers": 6, + "recurrent_depth": 4, + "num_experts": 36, + "num_experts_per_tok": 2, + "recurrent_attention_mode": "shared_first", + "vocab_size": 16384 + }, + "results": [ + { + "config": "mlx-full", + "device": "gpu", + "dtype": "bfloat16", + "prompts": 24, + "identical": 21, + "plain_tokens": 713, + "spec_tokens": 759, + "plain_s": 6.334, + "spec_s": 7.312, + "plain_tok_s": 112.56, + "spec_tok_s": 103.8, + "speedup": 0.922, + "speedup_median_per_prompt": 0.882, + "acceptance_rate": 0.5876, + "tokens_per_verify": 3.163, + "plain_ssd_loads": 0, + "spec_ssd_loads": 0, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 0.4717339170019841, + "spec_s": 0.5570417080016341, + "identical": true, + "index": 0, + "first_diff": null, + "accepted": 33, + "proposed": 65, + "verify_calls": 20, + "draft_calls": 65, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 42, + "spec_tokens": 42, + "plain_s": 0.36989050000556745, + "spec_s": 0.42769991600653157, + "identical": true, + "index": 1, + "first_diff": null, + "accepted": 28, + "proposed": 52, + "verify_calls": 14, + "draft_calls": 52, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 43, + "spec_tokens": 46, + "plain_s": 0.38608349999412894, + "spec_s": 0.5144372079957975, + "identical": false, + "index": 2, + "first_diff": 22, + "accepted": 28, + "proposed": 62, + "verify_calls": 18, + "draft_calls": 62, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 17, + "spec_tokens": 17, + "plain_s": 0.15002266599913128, + "spec_s": 0.21659308299422264, + "identical": true, + "index": 3, + "first_diff": null, + "accepted": 8, + "proposed": 24, + "verify_calls": 8, + "draft_calls": 24, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.136743374998332, + "spec_s": 0.16328241598966997, + "identical": true, + "index": 4, + "first_diff": null, + "accepted": 10, + "proposed": 20, + "verify_calls": 5, + "draft_calls": 20, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 0.3225371659937082, + "spec_s": 0.36308008299965877, + "identical": true, + "index": 5, + "first_diff": null, + "accepted": 22, + "proposed": 42, + "verify_calls": 13, + "draft_calls": 42, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 24, + "spec_tokens": 24, + "plain_s": 0.21275379198777955, + "spec_s": 0.20860633300617337, + "identical": true, + "index": 6, + "first_diff": null, + "accepted": 18, + "proposed": 25, + "verify_calls": 6, + "draft_calls": 25, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 0.3303770420025103, + "spec_s": 0.40254587499657646, + "identical": true, + "index": 7, + "first_diff": null, + "accepted": 20, + "proposed": 43, + "verify_calls": 16, + "draft_calls": 43, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.19435087499732617, + "spec_s": 0.20816920799552463, + "identical": true, + "index": 8, + "first_diff": null, + "accepted": 16, + "proposed": 26, + "verify_calls": 6, + "draft_calls": 26, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 31, + "spec_tokens": 36, + "plain_s": 0.27399320898985025, + "spec_s": 0.39195091600413434, + "identical": false, + "index": 9, + "first_diff": 24, + "accepted": 25, + "proposed": 50, + "verify_calls": 11, + "draft_calls": 50, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 90, + "spec_tokens": 128, + "plain_s": 0.8118831659958232, + "spec_s": 0.9274501250038156, + "identical": false, + "index": 10, + "first_diff": 13, + "accepted": 103, + "proposed": 123, + "verify_calls": 24, + "draft_calls": 123, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.195302084001014, + "spec_s": 0.20812887500505894, + "identical": true, + "index": 11, + "first_diff": null, + "accepted": 15, + "proposed": 24, + "verify_calls": 7, + "draft_calls": 24, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 16, + "spec_tokens": 16, + "plain_s": 0.14008758300042246, + "spec_s": 0.11736845900304615, + "identical": true, + "index": 12, + "first_diff": null, + "accepted": 12, + "proposed": 14, + "verify_calls": 3, + "draft_calls": 14, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.2005982499977108, + "spec_s": 0.2242509999923641, + "identical": true, + "index": 13, + "first_diff": null, + "accepted": 14, + "proposed": 24, + "verify_calls": 9, + "draft_calls": 24, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 32, + "spec_tokens": 32, + "plain_s": 0.280891625006916, + "spec_s": 0.34949570799653884, + "identical": true, + "index": 14, + "first_diff": null, + "accepted": 20, + "proposed": 42, + "verify_calls": 12, + "draft_calls": 42, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 40, + "spec_tokens": 40, + "plain_s": 0.35443116699752863, + "spec_s": 0.37398087499605026, + "identical": true, + "index": 15, + "first_diff": null, + "accepted": 28, + "proposed": 45, + "verify_calls": 12, + "draft_calls": 45, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 18, + "spec_tokens": 18, + "plain_s": 0.15814174999832176, + "spec_s": 0.1726262919983128, + "identical": true, + "index": 16, + "first_diff": null, + "accepted": 12, + "proposed": 20, + "verify_calls": 6, + "draft_calls": 20, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 28, + "spec_tokens": 28, + "plain_s": 0.2501576669892529, + "spec_s": 0.30397249999805354, + "identical": true, + "index": 17, + "first_diff": null, + "accepted": 18, + "proposed": 36, + "verify_calls": 10, + "draft_calls": 36, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 20, + "spec_tokens": 20, + "plain_s": 0.1729286250047153, + "spec_s": 0.1395225830056006, + "identical": true, + "index": 18, + "first_diff": null, + "accepted": 16, + "proposed": 17, + "verify_calls": 4, + "draft_calls": 17, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 33, + "spec_tokens": 33, + "plain_s": 0.2965928750054445, + "spec_s": 0.3277857920038514, + "identical": true, + "index": 19, + "first_diff": null, + "accepted": 20, + "proposed": 37, + "verify_calls": 12, + "draft_calls": 37, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.129949874986778, + "spec_s": 0.11529404199973214, + "identical": true, + "index": 20, + "first_diff": null, + "accepted": 12, + "proposed": 14, + "verify_calls": 3, + "draft_calls": 14, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 20, + "spec_tokens": 20, + "plain_s": 0.17928750001010485, + "spec_s": 0.19134362500335556, + "identical": true, + "index": 21, + "first_diff": null, + "accepted": 14, + "proposed": 23, + "verify_calls": 6, + "draft_calls": 23, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.20201458400697447, + "spec_s": 0.23260941699845716, + "identical": true, + "index": 22, + "first_diff": null, + "accepted": 15, + "proposed": 27, + "verify_calls": 8, + "draft_calls": 27, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 13, + "spec_tokens": 13, + "plain_s": 0.11373162499512546, + "spec_s": 0.17506458300340455, + "identical": true, + "index": 23, + "first_diff": null, + "accepted": 6, + "proposed": 18, + "verify_calls": 7, + "draft_calls": 18, + "plain_loads": 0, + "spec_loads": 0 + } + ] + }, + { + "config": "mlx-lazy", + "device": "gpu", + "dtype": "bfloat16", + "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 + }, + "prompts": 24, + "identical": 22, + "plain_tokens": 726, + "spec_tokens": 732, + "plain_s": 16.67, + "spec_s": 23.403, + "plain_tok_s": 43.55, + "spec_tok_s": 31.28, + "speedup": 0.718, + "speedup_median_per_prompt": 0.7, + "acceptance_rate": 0.5667, + "tokens_per_verify": 2.94, + "plain_ssd_loads": 8378, + "spec_ssd_loads": 12553, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 1.0774902080011088, + "spec_s": 1.6515985839942005, + "identical": true, + "index": 0, + "first_diff": null, + "accepted": 32, + "proposed": 63, + "verify_calls": 21, + "draft_calls": 63, + "plain_loads": 416, + "spec_loads": 722 + }, + { + "plain_tokens": 42, + "spec_tokens": 45, + "plain_s": 0.9039026669925079, + "spec_s": 1.461880957998801, + "identical": false, + "index": 1, + "first_diff": 2, + "accepted": 28, + "proposed": 53, + "verify_calls": 17, + "draft_calls": 53, + "plain_loads": 388, + "spec_loads": 657 + }, + { + "plain_tokens": 38, + "spec_tokens": 38, + "plain_s": 0.8789100840076571, + "spec_s": 1.0684897090104641, + "identical": true, + "index": 2, + "first_diff": null, + "accepted": 28, + "proposed": 43, + "verify_calls": 10, + "draft_calls": 43, + "plain_loads": 410, + "spec_loads": 593 + }, + { + "plain_tokens": 16, + "spec_tokens": 19, + "plain_s": 0.39396941600716673, + "spec_s": 0.7173566249985015, + "identical": false, + "index": 3, + "first_diff": 12, + "accepted": 11, + "proposed": 25, + "verify_calls": 8, + "draft_calls": 25, + "plain_loads": 235, + "spec_loads": 406 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.4358677499985788, + "spec_s": 0.6216179170005489, + "identical": true, + "index": 4, + "first_diff": null, + "accepted": 10, + "proposed": 20, + "verify_calls": 5, + "draft_calls": 20, + "plain_loads": 293, + "spec_loads": 399 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 0.7446927080018213, + "spec_s": 1.0527858750137966, + "identical": true, + "index": 5, + "first_diff": null, + "accepted": 22, + "proposed": 42, + "verify_calls": 13, + "draft_calls": 42, + "plain_loads": 341, + "spec_loads": 462 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 0.6952921660122229, + "spec_s": 0.8811332920013228, + "identical": true, + "index": 6, + "first_diff": null, + "accepted": 21, + "proposed": 27, + "verify_calls": 6, + "draft_calls": 27, + "plain_loads": 403, + "spec_loads": 629 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 0.8665314999961993, + "spec_s": 1.3091767500009155, + "identical": true, + "index": 7, + "first_diff": null, + "accepted": 20, + "proposed": 43, + "verify_calls": 16, + "draft_calls": 43, + "plain_loads": 432, + "spec_loads": 690 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.5646691669971915, + "spec_s": 0.7218595410085982, + "identical": true, + "index": 8, + "first_diff": null, + "accepted": 16, + "proposed": 26, + "verify_calls": 6, + "draft_calls": 26, + "plain_loads": 343, + "spec_loads": 401 + }, + { + "plain_tokens": 30, + "spec_tokens": 30, + "plain_s": 0.7018145829933928, + "spec_s": 1.011187833006261, + "identical": true, + "index": 9, + "first_diff": null, + "accepted": 22, + "proposed": 37, + "verify_calls": 8, + "draft_calls": 37, + "plain_loads": 368, + "spec_loads": 606 + }, + { + "plain_tokens": 90, + "spec_tokens": 90, + "plain_s": 1.7142818330030423, + "spec_s": 2.081965917008347, + "identical": true, + "index": 10, + "first_diff": null, + "accepted": 62, + "proposed": 104, + "verify_calls": 28, + "draft_calls": 104, + "plain_loads": 514, + "spec_loads": 735 + }, + { + "plain_tokens": 19, + "spec_tokens": 19, + "plain_s": 0.4724380830011796, + "spec_s": 0.6912645830016118, + "identical": true, + "index": 11, + "first_diff": null, + "accepted": 10, + "proposed": 22, + "verify_calls": 9, + "draft_calls": 22, + "plain_loads": 297, + "spec_loads": 383 + }, + { + "plain_tokens": 16, + "spec_tokens": 16, + "plain_s": 0.4210213749902323, + "spec_s": 0.4188643749948824, + "identical": true, + "index": 12, + "first_diff": null, + "accepted": 12, + "proposed": 14, + "verify_calls": 3, + "draft_calls": 14, + "plain_loads": 267, + "spec_loads": 281 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.5316629999870202, + "spec_s": 0.6942349580058362, + "identical": true, + "index": 13, + "first_diff": null, + "accepted": 14, + "proposed": 24, + "verify_calls": 9, + "draft_calls": 24, + "plain_loads": 262, + "spec_loads": 370 + }, + { + "plain_tokens": 32, + "spec_tokens": 32, + "plain_s": 0.6988647079997463, + "spec_s": 1.1310564159939531, + "identical": true, + "index": 14, + "first_diff": null, + "accepted": 20, + "proposed": 43, + "verify_calls": 12, + "draft_calls": 43, + "plain_loads": 355, + "spec_loads": 600 + }, + { + "plain_tokens": 40, + "spec_tokens": 40, + "plain_s": 0.9208465830015484, + "spec_s": 1.2930301250016782, + "identical": true, + "index": 15, + "first_diff": null, + "accepted": 28, + "proposed": 45, + "verify_calls": 12, + "draft_calls": 45, + "plain_loads": 443, + "spec_loads": 729 + }, + { + "plain_tokens": 26, + "spec_tokens": 26, + "plain_s": 0.6660797080112388, + "spec_s": 0.9744120409886818, + "identical": true, + "index": 16, + "first_diff": null, + "accepted": 17, + "proposed": 30, + "verify_calls": 9, + "draft_calls": 30, + "plain_loads": 401, + "spec_loads": 625 + }, + { + "plain_tokens": 28, + "spec_tokens": 28, + "plain_s": 0.6445007500005886, + "spec_s": 1.009474708000198, + "identical": true, + "index": 17, + "first_diff": null, + "accepted": 18, + "proposed": 36, + "verify_calls": 10, + "draft_calls": 36, + "plain_loads": 308, + "spec_loads": 513 + }, + { + "plain_tokens": 20, + "spec_tokens": 20, + "plain_s": 0.5596654160035541, + "spec_s": 0.786746415993548, + "identical": true, + "index": 18, + "first_diff": null, + "accepted": 13, + "proposed": 22, + "verify_calls": 7, + "draft_calls": 22, + "plain_loads": 377, + "spec_loads": 509 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 0.7948367080098251, + "spec_s": 1.1376042500050971, + "identical": true, + "index": 19, + "first_diff": null, + "accepted": 24, + "proposed": 39, + "verify_calls": 13, + "draft_calls": 39, + "plain_loads": 357, + "spec_loads": 601 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.4636524579982506, + "spec_s": 0.5425763750099577, + "identical": true, + "index": 20, + "first_diff": null, + "accepted": 12, + "proposed": 14, + "verify_calls": 3, + "draft_calls": 14, + "plain_loads": 343, + "spec_loads": 452 + }, + { + "plain_tokens": 25, + "spec_tokens": 25, + "plain_s": 0.5416186669899616, + "spec_s": 0.6948067499906756, + "identical": true, + "index": 21, + "first_diff": null, + "accepted": 18, + "proposed": 26, + "verify_calls": 7, + "draft_calls": 26, + "plain_loads": 269, + "spec_loads": 387 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.6005229589936789, + "spec_s": 0.8628925410012016, + "identical": true, + "index": 22, + "first_diff": null, + "accepted": 14, + "proposed": 29, + "verify_calls": 9, + "draft_calls": 29, + "plain_loads": 355, + "spec_loads": 512 + }, + { + "plain_tokens": 16, + "spec_tokens": 16, + "plain_s": 0.3772377089917427, + "spec_s": 0.5865167909942102, + "identical": true, + "index": 23, + "first_diff": null, + "accepted": 8, + "proposed": 20, + "verify_calls": 8, + "draft_calls": 20, + "plain_loads": 201, + "spec_loads": 291 + } + ] + } + ] +} \ No newline at end of file diff --git a/reports/speculative/mlx_gpu_fp32_s13996.json b/reports/speculative/mlx_gpu_fp32_s13996.json new file mode 100644 index 0000000..a1bdf3d --- /dev/null +++ b/reports/speculative/mlx_gpu_fp32_s13996.json @@ -0,0 +1,834 @@ +{ + "label": "measured", + "time": "2026-11-06T19:27:04", + "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": 24, + "max_new": 128, + "configs": "mlx-full,mlx-lazy", + "mlx_device": "gpu", + "torch_device": "mps", + "mlx_dtype": "float32", + "only": "", + "ratio": 1.0, + "miss_policy": "sync_on_demand", + "draft_depth": 1, + "gamma": 4, + "warmup": 2, + "out": "reports/speculative/mlx_gpu_fp32_s13996.json", + "note": "72h MLX training paused with SIGSTOP; GPU fp32 exactness check; FLORES+ dev first 24 en sentences", + "speculative": { + "draft_depth": 1, + "gamma": 4, + "gamma_max": 8, + "adaptive_gamma": true + }, + "eos_id": 0, + "decoding": "greedy, batch 1", + "n_prompts": 24, + "prompt_tokens_mean": 39.7, + "order": "plain/speculative alternate per prompt", + "prompts_source": "FLORES+ dev (eng_Latn/zho_Hans) rows 0-31, field 'en'" + }, + "model": { + "hidden_size": 384, + "num_hidden_layers": 6, + "recurrent_depth": 4, + "num_experts": 36, + "num_experts_per_tok": 2, + "recurrent_attention_mode": "shared_first", + "vocab_size": 16384 + }, + "results": [ + { + "config": "mlx-full", + "device": "gpu", + "dtype": "float32", + "prompts": 24, + "identical": 23, + "plain_tokens": 703, + "spec_tokens": 700, + "plain_s": 6.223, + "spec_s": 6.994, + "plain_tok_s": 112.97, + "spec_tok_s": 100.09, + "speedup": 0.886, + "speedup_median_per_prompt": 0.891, + "acceptance_rate": 0.4101, + "tokens_per_verify": 1.983, + "plain_ssd_loads": 0, + "spec_ssd_loads": 0, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 0.47110470800544135, + "spec_s": 0.5307974999886937, + "identical": true, + "index": 0, + "first_diff": null, + "accepted": 26, + "proposed": 65, + "verify_calls": 27, + "draft_calls": 65, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 42, + "spec_tokens": 42, + "plain_s": 0.3708125419943826, + "spec_s": 0.4865471250086557, + "identical": true, + "index": 1, + "first_diff": null, + "accepted": 12, + "proposed": 47, + "verify_calls": 30, + "draft_calls": 47, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 38, + "spec_tokens": 38, + "plain_s": 0.3344719169981545, + "spec_s": 0.364178499992704, + "identical": true, + "index": 2, + "first_diff": null, + "accepted": 21, + "proposed": 50, + "verify_calls": 17, + "draft_calls": 50, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 0.23542154100141488, + "spec_s": 0.33127254201099277, + "identical": true, + "index": 3, + "first_diff": null, + "accepted": 7, + "proposed": 32, + "verify_calls": 20, + "draft_calls": 32, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.13302241599012632, + "spec_s": 0.17432904199813493, + "identical": true, + "index": 4, + "first_diff": null, + "accepted": 6, + "proposed": 19, + "verify_calls": 9, + "draft_calls": 19, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 0.32363216699741315, + "spec_s": 0.3747162500076229, + "identical": true, + "index": 5, + "first_diff": null, + "accepted": 16, + "proposed": 38, + "verify_calls": 20, + "draft_calls": 38, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 0.23844204199849628, + "spec_s": 0.19428049999987707, + "identical": true, + "index": 6, + "first_diff": null, + "accepted": 20, + "proposed": 31, + "verify_calls": 7, + "draft_calls": 31, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 0.3274757920007687, + "spec_s": 0.4292124159983359, + "identical": true, + "index": 7, + "first_diff": null, + "accepted": 13, + "proposed": 41, + "verify_calls": 24, + "draft_calls": 41, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.19240008300403133, + "spec_s": 0.2349682079948252, + "identical": true, + "index": 8, + "first_diff": null, + "accepted": 10, + "proposed": 26, + "verify_calls": 12, + "draft_calls": 26, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 30, + "spec_tokens": 30, + "plain_s": 0.26525054199737497, + "spec_s": 0.28797108300204854, + "identical": true, + "index": 9, + "first_diff": null, + "accepted": 18, + "proposed": 41, + "verify_calls": 12, + "draft_calls": 41, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 66, + "spec_tokens": 66, + "plain_s": 0.5883352499949979, + "spec_s": 0.5944269589963369, + "identical": true, + "index": 10, + "first_diff": null, + "accepted": 37, + "proposed": 76, + "verify_calls": 29, + "draft_calls": 76, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 19, + "spec_tokens": 19, + "plain_s": 0.16803979199903551, + "spec_s": 0.20387058400956448, + "identical": true, + "index": 11, + "first_diff": null, + "accepted": 8, + "proposed": 22, + "verify_calls": 11, + "draft_calls": 22, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 18, + "spec_tokens": 18, + "plain_s": 0.15573387499898672, + "spec_s": 0.17721533399890177, + "identical": true, + "index": 12, + "first_diff": null, + "accepted": 8, + "proposed": 20, + "verify_calls": 9, + "draft_calls": 20, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.203484749989002, + "spec_s": 0.21204199999920093, + "identical": true, + "index": 13, + "first_diff": null, + "accepted": 13, + "proposed": 26, + "verify_calls": 10, + "draft_calls": 26, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 32, + "spec_tokens": 32, + "plain_s": 0.2839589170034742, + "spec_s": 0.34376349998638034, + "identical": true, + "index": 14, + "first_diff": null, + "accepted": 12, + "proposed": 38, + "verify_calls": 19, + "draft_calls": 38, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 40, + "spec_tokens": 40, + "plain_s": 0.35079345799749717, + "spec_s": 0.4035519589961041, + "identical": true, + "index": 15, + "first_diff": null, + "accepted": 19, + "proposed": 50, + "verify_calls": 20, + "draft_calls": 50, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 18, + "spec_tokens": 18, + "plain_s": 0.15958762499212753, + "spec_s": 0.17199016599624883, + "identical": true, + "index": 16, + "first_diff": null, + "accepted": 10, + "proposed": 22, + "verify_calls": 8, + "draft_calls": 22, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 28, + "spec_tokens": 28, + "plain_s": 0.24673304200405255, + "spec_s": 0.23635108300368302, + "identical": true, + "index": 17, + "first_diff": null, + "accepted": 19, + "proposed": 36, + "verify_calls": 9, + "draft_calls": 36, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 20, + "spec_tokens": 20, + "plain_s": 0.17097958399972413, + "spec_s": 0.2114502500044182, + "identical": true, + "index": 18, + "first_diff": null, + "accepted": 7, + "proposed": 22, + "verify_calls": 12, + "draft_calls": 22, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 33, + "spec_tokens": 33, + "plain_s": 0.29087362499558367, + "spec_s": 0.2958425830001943, + "identical": true, + "index": 19, + "first_diff": null, + "accepted": 19, + "proposed": 38, + "verify_calls": 14, + "draft_calls": 38, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.13133412500610575, + "spec_s": 0.1342817910044687, + "identical": true, + "index": 20, + "first_diff": null, + "accepted": 9, + "proposed": 17, + "verify_calls": 6, + "draft_calls": 17, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 25, + "spec_tokens": 25, + "plain_s": 0.23259737499756739, + "spec_s": 0.21592341600626241, + "identical": true, + "index": 21, + "first_diff": null, + "accepted": 15, + "proposed": 31, + "verify_calls": 9, + "draft_calls": 31, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.20376808400033042, + "spec_s": 0.22767616700730287, + "identical": true, + "index": 22, + "first_diff": null, + "accepted": 12, + "proposed": 28, + "verify_calls": 11, + "draft_calls": 28, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 16, + "spec_tokens": 13, + "plain_s": 0.14479250001022592, + "spec_s": 0.15723979100584984, + "identical": false, + "index": 23, + "first_diff": 2, + "accepted": 5, + "proposed": 18, + "verify_calls": 8, + "draft_calls": 18, + "plain_loads": 0, + "spec_loads": 0 + } + ] + }, + { + "config": "mlx-lazy", + "device": "gpu", + "dtype": "float32", + "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 + }, + "prompts": 24, + "identical": 24, + "plain_tokens": 703, + "spec_tokens": 703, + "plain_s": 16.32, + "spec_s": 22.614, + "plain_tok_s": 43.08, + "spec_tok_s": 31.09, + "speedup": 0.722, + "speedup_median_per_prompt": 0.727, + "acceptance_rate": 0.4086, + "tokens_per_verify": 1.975, + "plain_ssd_loads": 8307, + "spec_ssd_loads": 11855, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 1.0622797079995507, + "spec_s": 1.505737749990658, + "identical": true, + "index": 0, + "first_diff": null, + "accepted": 26, + "proposed": 65, + "verify_calls": 27, + "draft_calls": 65, + "plain_loads": 409, + "spec_loads": 627 + }, + { + "plain_tokens": 42, + "spec_tokens": 42, + "plain_s": 0.8918150829995284, + "spec_s": 1.3613081660005264, + "identical": true, + "index": 1, + "first_diff": null, + "accepted": 12, + "proposed": 47, + "verify_calls": 30, + "draft_calls": 47, + "plain_loads": 386, + "spec_loads": 543 + }, + { + "plain_tokens": 38, + "spec_tokens": 38, + "plain_s": 0.8662984590046108, + "spec_s": 1.3148910000018077, + "identical": true, + "index": 2, + "first_diff": null, + "accepted": 21, + "proposed": 50, + "verify_calls": 17, + "draft_calls": 50, + "plain_loads": 401, + "spec_loads": 746 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 0.5570020829909481, + "spec_s": 0.862909083996783, + "identical": true, + "index": 3, + "first_diff": null, + "accepted": 7, + "proposed": 32, + "verify_calls": 20, + "draft_calls": 32, + "plain_loads": 253, + "spec_loads": 371 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.4427129160030745, + "spec_s": 0.6464972079993458, + "identical": true, + "index": 4, + "first_diff": null, + "accepted": 6, + "proposed": 19, + "verify_calls": 9, + "draft_calls": 19, + "plain_loads": 291, + "spec_loads": 410 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 0.7288510840007802, + "spec_s": 0.9548697910067858, + "identical": true, + "index": 5, + "first_diff": null, + "accepted": 16, + "proposed": 38, + "verify_calls": 20, + "draft_calls": 38, + "plain_loads": 297, + "spec_loads": 398 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 0.7059443750040373, + "spec_s": 0.9211508750013309, + "identical": true, + "index": 6, + "first_diff": null, + "accepted": 20, + "proposed": 31, + "verify_calls": 7, + "draft_calls": 31, + "plain_loads": 420, + "spec_loads": 691 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 0.8489263340015896, + "spec_s": 1.221835958000156, + "identical": true, + "index": 7, + "first_diff": null, + "accepted": 13, + "proposed": 41, + "verify_calls": 24, + "draft_calls": 41, + "plain_loads": 391, + "spec_loads": 542 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.569496625001193, + "spec_s": 0.7525296669919044, + "identical": true, + "index": 8, + "first_diff": null, + "accepted": 10, + "proposed": 26, + "verify_calls": 12, + "draft_calls": 26, + "plain_loads": 342, + "spec_loads": 386 + }, + { + "plain_tokens": 30, + "spec_tokens": 30, + "plain_s": 0.7071564579964615, + "spec_s": 0.9628682089969516, + "identical": true, + "index": 9, + "first_diff": null, + "accepted": 18, + "proposed": 41, + "verify_calls": 12, + "draft_calls": 41, + "plain_loads": 370, + "spec_loads": 505 + }, + { + "plain_tokens": 66, + "spec_tokens": 66, + "plain_s": 1.3164745840040268, + "spec_s": 1.5809367500041844, + "identical": true, + "index": 10, + "first_diff": null, + "accepted": 37, + "proposed": 76, + "verify_calls": 29, + "draft_calls": 76, + "plain_loads": 485, + "spec_loads": 606 + }, + { + "plain_tokens": 19, + "spec_tokens": 19, + "plain_s": 0.46918112500861753, + "spec_s": 0.7184042080043582, + "identical": true, + "index": 11, + "first_diff": null, + "accepted": 8, + "proposed": 22, + "verify_calls": 11, + "draft_calls": 22, + "plain_loads": 281, + "spec_loads": 396 + }, + { + "plain_tokens": 18, + "spec_tokens": 18, + "plain_s": 0.44788654100557324, + "spec_s": 0.579919208001229, + "identical": true, + "index": 12, + "first_diff": null, + "accepted": 8, + "proposed": 20, + "verify_calls": 9, + "draft_calls": 20, + "plain_loads": 260, + "spec_loads": 331 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.524843957988196, + "spec_s": 0.6935867919964949, + "identical": true, + "index": 13, + "first_diff": null, + "accepted": 13, + "proposed": 26, + "verify_calls": 10, + "draft_calls": 26, + "plain_loads": 263, + "spec_loads": 407 + }, + { + "plain_tokens": 32, + "spec_tokens": 32, + "plain_s": 0.7057327079965034, + "spec_s": 1.0392251659941394, + "identical": true, + "index": 14, + "first_diff": null, + "accepted": 12, + "proposed": 38, + "verify_calls": 19, + "draft_calls": 38, + "plain_loads": 355, + "spec_loads": 508 + }, + { + "plain_tokens": 40, + "spec_tokens": 40, + "plain_s": 0.9574661670048954, + "spec_s": 1.3251612919993931, + "identical": true, + "index": 15, + "first_diff": null, + "accepted": 19, + "proposed": 50, + "verify_calls": 20, + "draft_calls": 50, + "plain_loads": 484, + "spec_loads": 702 + }, + { + "plain_tokens": 18, + "spec_tokens": 18, + "plain_s": 0.5418095000059111, + "spec_s": 0.7155439589987509, + "identical": true, + "index": 16, + "first_diff": null, + "accepted": 10, + "proposed": 22, + "verify_calls": 8, + "draft_calls": 22, + "plain_loads": 382, + "spec_loads": 498 + }, + { + "plain_tokens": 28, + "spec_tokens": 28, + "plain_s": 0.6576261250011157, + "spec_s": 0.8972636670077918, + "identical": true, + "index": 17, + "first_diff": null, + "accepted": 19, + "proposed": 36, + "verify_calls": 9, + "draft_calls": 36, + "plain_loads": 311, + "spec_loads": 542 + }, + { + "plain_tokens": 20, + "spec_tokens": 20, + "plain_s": 0.5571614159998717, + "spec_s": 0.803125666003325, + "identical": true, + "index": 18, + "first_diff": null, + "accepted": 7, + "proposed": 22, + "verify_calls": 12, + "draft_calls": 22, + "plain_loads": 376, + "spec_loads": 501 + }, + { + "plain_tokens": 33, + "spec_tokens": 33, + "plain_s": 0.7482056250009919, + "spec_s": 1.023553749997518, + "identical": true, + "index": 19, + "first_diff": null, + "accepted": 19, + "proposed": 38, + "verify_calls": 14, + "draft_calls": 38, + "plain_loads": 349, + "spec_loads": 543 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.4984620420000283, + "spec_s": 0.5969373749976512, + "identical": true, + "index": 20, + "first_diff": null, + "accepted": 9, + "proposed": 17, + "verify_calls": 6, + "draft_calls": 17, + "plain_loads": 353, + "spec_loads": 422 + }, + { + "plain_tokens": 25, + "spec_tokens": 25, + "plain_s": 0.5585083330079215, + "spec_s": 0.7177489589957986, + "identical": true, + "index": 21, + "first_diff": null, + "accepted": 15, + "proposed": 31, + "verify_calls": 9, + "draft_calls": 31, + "plain_loads": 291, + "spec_loads": 396 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.5817440410028212, + "spec_s": 0.8349084999936167, + "identical": true, + "index": 22, + "first_diff": null, + "accepted": 12, + "proposed": 28, + "verify_calls": 11, + "draft_calls": 28, + "plain_loads": 356, + "spec_loads": 518 + }, + { + "plain_tokens": 16, + "spec_tokens": 16, + "plain_s": 0.3744943750061793, + "spec_s": 0.5827130420075264, + "identical": true, + "index": 23, + "first_diff": null, + "accepted": 5, + "proposed": 21, + "verify_calls": 11, + "draft_calls": 21, + "plain_loads": 201, + "spec_loads": 266 + } + ] + } + ] +} \ No newline at end of file diff --git a/reports/speculative/mlx_gpu_s13996.json b/reports/speculative/mlx_gpu_s13996.json new file mode 100644 index 0000000..f20e539 --- /dev/null +++ b/reports/speculative/mlx_gpu_s13996.json @@ -0,0 +1,736 @@ +{ + "label": "measured", + "time": "2026-11-06T19:23:05", + "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": 24, + "max_new": 128, + "configs": "mlx-full,mlx-lazy", + "mlx_device": "gpu", + "torch_device": "mps", + "ratio": 1.0, + "miss_policy": "sync_on_demand", + "draft_depth": 1, + "gamma": 4, + "warmup": 2, + "out": "reports/speculative/mlx_gpu_s13996.json", + "note": "72h MLX training paused with SIGSTOP during measurement; FLORES+ dev first 24 en sentences, en->zh pretrain tags", + "speculative": { + "draft_depth": 1, + "gamma": 4, + "gamma_max": 8, + "adaptive_gamma": true + }, + "eos_id": 0, + "decoding": "greedy, batch 1", + "n_prompts": 24, + "prompt_tokens_mean": 39.7, + "order": "plain/speculative alternate per prompt", + "prompts_source": "FLORES+ dev (eng_Latn/zho_Hans) rows 0-31, field 'en'" + }, + "model": { + "hidden_size": 384, + "num_hidden_layers": 6, + "recurrent_depth": 4, + "num_experts": 36, + "num_experts_per_tok": 2, + "recurrent_attention_mode": "shared_first", + "vocab_size": 16384 + }, + "results": [ + { + "config": "mlx-full", + "device": "gpu", + "dtype": "bfloat16", + "prompts": 24, + "identical": 22, + "plain_tokens": 713, + "spec_tokens": 756, + "plain_s": 6.364, + "spec_s": 7.037, + "plain_tok_s": 112.04, + "spec_tok_s": 107.44, + "speedup": 0.959, + "speedup_median_per_prompt": 0.919, + "acceptance_rate": 0.4251, + "tokens_per_verify": 2.088, + "plain_ssd_loads": 0, + "spec_ssd_loads": 0, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 0.46368174999952316, + "spec_s": 0.5035165000008419, + "identical": true, + "accepted": 26, + "proposed": 65, + "verify_calls": 27, + "draft_calls": 65, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 42, + "spec_tokens": 42, + "plain_s": 0.3676152080006432, + "spec_s": 0.4870770419947803, + "identical": true, + "accepted": 12, + "proposed": 47, + "verify_calls": 30, + "draft_calls": 47, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 43, + "spec_tokens": 43, + "plain_s": 0.38048129199887626, + "spec_s": 0.43994029200985096, + "identical": true, + "accepted": 21, + "proposed": 54, + "verify_calls": 22, + "draft_calls": 54, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 17, + "spec_tokens": 17, + "plain_s": 0.14746779200504534, + "spec_s": 0.19229279200953897, + "identical": true, + "accepted": 5, + "proposed": 21, + "verify_calls": 11, + "draft_calls": 21, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.1329260409984272, + "spec_s": 0.15823629198712297, + "identical": true, + "accepted": 7, + "proposed": 20, + "verify_calls": 8, + "draft_calls": 20, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 0.32108145800884813, + "spec_s": 0.34047191600257065, + "identical": true, + "accepted": 16, + "proposed": 38, + "verify_calls": 20, + "draft_calls": 38, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 24, + "spec_tokens": 24, + "plain_s": 0.21418912499211729, + "spec_s": 0.2177905419957824, + "identical": true, + "accepted": 15, + "proposed": 34, + "verify_calls": 9, + "draft_calls": 34, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 0.32615712500410154, + "spec_s": 0.3893993339879671, + "identical": true, + "accepted": 13, + "proposed": 41, + "verify_calls": 24, + "draft_calls": 41, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.29722133399627637, + "spec_s": 0.21805237499938812, + "identical": true, + "accepted": 10, + "proposed": 26, + "verify_calls": 12, + "draft_calls": 26, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 31, + "spec_tokens": 36, + "plain_s": 0.2677647079981398, + "spec_s": 0.33349391700176056, + "identical": false, + "accepted": 19, + "proposed": 46, + "verify_calls": 16, + "draft_calls": 46, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 90, + "spec_tokens": 128, + "plain_s": 0.7971258749894332, + "spec_s": 0.990837958001066, + "identical": false, + "accepted": 83, + "proposed": 150, + "verify_calls": 44, + "draft_calls": 150, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.19348674999491777, + "spec_s": 0.1813736250041984, + "identical": true, + "accepted": 14, + "proposed": 26, + "verify_calls": 8, + "draft_calls": 26, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 16, + "spec_tokens": 16, + "plain_s": 0.13883095800701994, + "spec_s": 0.1379902500048047, + "identical": true, + "accepted": 8, + "proposed": 17, + "verify_calls": 7, + "draft_calls": 17, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.20529004199488554, + "spec_s": 0.19889862500713207, + "identical": true, + "accepted": 13, + "proposed": 26, + "verify_calls": 10, + "draft_calls": 26, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 32, + "spec_tokens": 32, + "plain_s": 0.28104133300075773, + "spec_s": 0.3235602919885423, + "identical": true, + "accepted": 13, + "proposed": 39, + "verify_calls": 18, + "draft_calls": 39, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 40, + "spec_tokens": 40, + "plain_s": 0.3462434169923654, + "spec_s": 0.3820523749891436, + "identical": true, + "accepted": 19, + "proposed": 50, + "verify_calls": 20, + "draft_calls": 50, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 18, + "spec_tokens": 18, + "plain_s": 0.15613633299653884, + "spec_s": 0.16375274999882095, + "identical": true, + "accepted": 10, + "proposed": 22, + "verify_calls": 8, + "draft_calls": 22, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 28, + "spec_tokens": 28, + "plain_s": 0.24419041600776836, + "spec_s": 0.2192378329928033, + "identical": true, + "accepted": 19, + "proposed": 36, + "verify_calls": 9, + "draft_calls": 36, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 20, + "spec_tokens": 20, + "plain_s": 0.17054654099047184, + "spec_s": 0.19378608399711084, + "identical": true, + "accepted": 8, + "proposed": 23, + "verify_calls": 11, + "draft_calls": 23, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 33, + "spec_tokens": 33, + "plain_s": 0.2881738750002114, + "spec_s": 0.2779974589939229, + "identical": true, + "accepted": 19, + "proposed": 38, + "verify_calls": 14, + "draft_calls": 38, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.14033691700024065, + "spec_s": 0.13323262499761768, + "identical": true, + "accepted": 9, + "proposed": 17, + "verify_calls": 6, + "draft_calls": 17, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 20, + "spec_tokens": 20, + "plain_s": 0.17541329200321343, + "spec_s": 0.17484804200648796, + "identical": true, + "accepted": 11, + "proposed": 25, + "verify_calls": 8, + "draft_calls": 25, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.19772650000231806, + "spec_s": 0.2155040000070585, + "identical": true, + "accepted": 12, + "proposed": 28, + "verify_calls": 11, + "draft_calls": 28, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 13, + "spec_tokens": 13, + "plain_s": 0.11084666699753143, + "spec_s": 0.16339658299693838, + "identical": true, + "accepted": 4, + "proposed": 19, + "verify_calls": 9, + "draft_calls": 19, + "plain_loads": 0, + "spec_loads": 0 + } + ] + }, + { + "config": "mlx-lazy", + "device": "gpu", + "dtype": "bfloat16", + "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 + }, + "prompts": 24, + "identical": 22, + "plain_tokens": 726, + "spec_tokens": 732, + "plain_s": 16.61, + "spec_s": 22.859, + "plain_tok_s": 43.71, + "spec_tok_s": 32.02, + "speedup": 0.733, + "speedup_median_per_prompt": 0.716, + "acceptance_rate": 0.4237, + "tokens_per_verify": 2.05, + "plain_ssd_loads": 8372, + "spec_ssd_loads": 12092, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 1.0585620420024497, + "spec_s": 1.5052647500124294, + "identical": true, + "accepted": 26, + "proposed": 65, + "verify_calls": 27, + "draft_calls": 65, + "plain_loads": 417, + "spec_loads": 629 + }, + { + "plain_tokens": 42, + "spec_tokens": 45, + "plain_s": 0.8853354579914594, + "spec_s": 1.2863747919909656, + "identical": false, + "accepted": 24, + "proposed": 55, + "verify_calls": 21, + "draft_calls": 55, + "plain_loads": 389, + "spec_loads": 576 + }, + { + "plain_tokens": 38, + "spec_tokens": 38, + "plain_s": 0.8748158749949653, + "spec_s": 1.2922939999989467, + "identical": true, + "accepted": 21, + "proposed": 50, + "verify_calls": 17, + "draft_calls": 50, + "plain_loads": 410, + "spec_loads": 744 + }, + { + "plain_tokens": 16, + "spec_tokens": 19, + "plain_s": 0.39575125000556, + "spec_s": 0.6381637080048677, + "identical": false, + "accepted": 7, + "proposed": 22, + "verify_calls": 11, + "draft_calls": 22, + "plain_loads": 249, + "spec_loads": 356 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.4230492500064429, + "spec_s": 0.6264298330061138, + "identical": true, + "accepted": 6, + "proposed": 19, + "verify_calls": 9, + "draft_calls": 19, + "plain_loads": 293, + "spec_loads": 410 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 0.7194180420046905, + "spec_s": 0.9592659170011757, + "identical": true, + "accepted": 16, + "proposed": 38, + "verify_calls": 20, + "draft_calls": 38, + "plain_loads": 297, + "spec_loads": 408 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 0.6805705419974402, + "spec_s": 0.8985260000044946, + "identical": true, + "accepted": 20, + "proposed": 31, + "verify_calls": 7, + "draft_calls": 31, + "plain_loads": 403, + "spec_loads": 691 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 0.8344688329962082, + "spec_s": 1.205341540990048, + "identical": true, + "accepted": 13, + "proposed": 41, + "verify_calls": 24, + "draft_calls": 41, + "plain_loads": 391, + "spec_loads": 544 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.6573287089995574, + "spec_s": 0.7483851249999134, + "identical": true, + "accepted": 10, + "proposed": 26, + "verify_calls": 12, + "draft_calls": 26, + "plain_loads": 343, + "spec_loads": 388 + }, + { + "plain_tokens": 30, + "spec_tokens": 30, + "plain_s": 0.7049508749914821, + "spec_s": 0.9585966669983463, + "identical": true, + "accepted": 18, + "proposed": 41, + "verify_calls": 12, + "draft_calls": 41, + "plain_loads": 368, + "spec_loads": 499 + }, + { + "plain_tokens": 90, + "spec_tokens": 90, + "plain_s": 1.7142060409969417, + "spec_s": 2.0970234999986133, + "identical": true, + "accepted": 47, + "proposed": 103, + "verify_calls": 43, + "draft_calls": 103, + "plain_loads": 514, + "spec_loads": 706 + }, + { + "plain_tokens": 19, + "spec_tokens": 19, + "plain_s": 0.4721540410100715, + "spec_s": 0.6850622079946334, + "identical": true, + "accepted": 8, + "proposed": 22, + "verify_calls": 11, + "draft_calls": 22, + "plain_loads": 296, + "spec_loads": 383 + }, + { + "plain_tokens": 16, + "spec_tokens": 16, + "plain_s": 0.42573758300568443, + "spec_s": 0.4929459999984829, + "identical": true, + "accepted": 8, + "proposed": 17, + "verify_calls": 7, + "draft_calls": 17, + "plain_loads": 267, + "spec_loads": 299 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.5195040840044385, + "spec_s": 0.6826135839946801, + "identical": true, + "accepted": 13, + "proposed": 26, + "verify_calls": 10, + "draft_calls": 26, + "plain_loads": 262, + "spec_loads": 399 + }, + { + "plain_tokens": 32, + "spec_tokens": 32, + "plain_s": 0.695727625003201, + "spec_s": 1.0219530839967774, + "identical": true, + "accepted": 12, + "proposed": 38, + "verify_calls": 19, + "draft_calls": 38, + "plain_loads": 355, + "spec_loads": 505 + }, + { + "plain_tokens": 40, + "spec_tokens": 40, + "plain_s": 0.952616415990633, + "spec_s": 1.3055763749871403, + "identical": true, + "accepted": 19, + "proposed": 50, + "verify_calls": 20, + "draft_calls": 50, + "plain_loads": 484, + "spec_loads": 702 + }, + { + "plain_tokens": 26, + "spec_tokens": 26, + "plain_s": 0.6562734589970205, + "spec_s": 0.8654622500034748, + "identical": true, + "accepted": 15, + "proposed": 31, + "verify_calls": 11, + "draft_calls": 31, + "plain_loads": 401, + "spec_loads": 547 + }, + { + "plain_tokens": 28, + "spec_tokens": 28, + "plain_s": 0.6541978330060374, + "spec_s": 0.8779284170013852, + "identical": true, + "accepted": 19, + "proposed": 36, + "verify_calls": 9, + "draft_calls": 36, + "plain_loads": 310, + "spec_loads": 523 + }, + { + "plain_tokens": 20, + "spec_tokens": 20, + "plain_s": 0.5461565000005066, + "spec_s": 0.7866672499949345, + "identical": true, + "accepted": 7, + "proposed": 22, + "verify_calls": 12, + "draft_calls": 22, + "plain_loads": 377, + "spec_loads": 502 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 0.7884885420062346, + "spec_s": 1.2427802919992246, + "identical": true, + "accepted": 19, + "proposed": 41, + "verify_calls": 18, + "draft_calls": 41, + "plain_loads": 357, + "spec_loads": 688 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.4579244169872254, + "spec_s": 0.5758031250006752, + "identical": true, + "accepted": 9, + "proposed": 17, + "verify_calls": 6, + "draft_calls": 17, + "plain_loads": 343, + "spec_loads": 418 + }, + { + "plain_tokens": 25, + "spec_tokens": 25, + "plain_s": 0.5457124579988886, + "spec_s": 0.704952542000683, + "identical": true, + "accepted": 15, + "proposed": 31, + "verify_calls": 9, + "draft_calls": 31, + "plain_loads": 290, + "spec_loads": 396 + }, + { + "plain_tokens": 23, + "spec_tokens": 23, + "plain_s": 0.5727323749888455, + "spec_s": 0.8209240829892224, + "identical": true, + "accepted": 12, + "proposed": 28, + "verify_calls": 11, + "draft_calls": 28, + "plain_loads": 355, + "spec_loads": 513 + }, + { + "plain_tokens": 16, + "spec_tokens": 16, + "plain_s": 0.37382770801195875, + "spec_s": 0.5811494590016082, + "identical": true, + "accepted": 5, + "proposed": 21, + "verify_calls": 11, + "draft_calls": 21, + "plain_loads": 201, + "spec_loads": 266 + } + ] + } + ] +} \ No newline at end of file diff --git a/reports/speculative/torch_mps_d2_s13996.json b/reports/speculative/torch_mps_d2_s13996.json new file mode 100644 index 0000000..ac86159 --- /dev/null +++ b/reports/speculative/torch_mps_d2_s13996.json @@ -0,0 +1,273 @@ +{ + "label": "measured", + "time": "2026-11-06T19:30:15", + "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": 12, + "max_new": 128, + "configs": "torch-lazy", + "mlx_device": "gpu", + "torch_device": "mps", + "mlx_dtype": "auto", + "only": "", + "ratio": 1.0, + "miss_policy": "sync_on_demand", + "draft_depth": 2, + "gamma": 4, + "warmup": 1, + "out": "reports/speculative/torch_mps_d2_s13996.json", + "note": "72h MLX training paused with SIGSTOP; draft_depth 2; FLORES+ dev first 12 en sentences", + "speculative": { + "draft_depth": 2, + "gamma": 4, + "gamma_max": 8, + "adaptive_gamma": true + }, + "eos_id": 0, + "decoding": "greedy, batch 1", + "n_prompts": 12, + "prompt_tokens_mean": 45.6, + "order": "plain/speculative alternate per prompt", + "prompts_source": "FLORES+ dev (eng_Latn/zho_Hans) rows 0-31, field 'en'" + }, + "model": { + "hidden_size": 384, + "num_hidden_layers": 6, + "recurrent_depth": 4, + "num_experts": 36, + "num_experts_per_tok": 2, + "recurrent_attention_mode": "shared_first", + "vocab_size": 16384 + }, + "results": [ + { + "config": "torch-lazy", + "device": "mps", + "dtype": "float32", + "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 + }, + "prompts": 12, + "identical": 12, + "plain_tokens": 436, + "spec_tokens": 436, + "plain_s": 26.706, + "spec_s": 33.244, + "plain_tok_s": 16.33, + "spec_tok_s": 13.11, + "speedup": 0.803, + "speedup_median_per_prompt": 0.796, + "acceptance_rate": 0.5543, + "tokens_per_verify": 2.946, + "plain_ssd_loads": 8787, + "spec_ssd_loads": 14453, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 3.137656499995501, + "spec_s": 3.7229940419929335, + "identical": true, + "index": 0, + "first_diff": null, + "accepted": 34, + "proposed": 65, + "verify_calls": 19, + "draft_calls": 65, + "plain_loads": 999, + "spec_loads": 1726 + }, + { + "plain_tokens": 42, + "spec_tokens": 42, + "plain_s": 2.497561583004426, + "spec_s": 3.364663250002195, + "identical": true, + "index": 1, + "first_diff": null, + "accepted": 28, + "proposed": 51, + "verify_calls": 14, + "draft_calls": 51, + "plain_loads": 841, + "spec_loads": 1436 + }, + { + "plain_tokens": 38, + "spec_tokens": 38, + "plain_s": 2.5273539580084616, + "spec_s": 2.5479305409971857, + "identical": true, + "index": 2, + "first_diff": null, + "accepted": 28, + "proposed": 43, + "verify_calls": 10, + "draft_calls": 43, + "plain_loads": 826, + "spec_loads": 1205 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 1.5812431250087684, + "spec_s": 2.6952759589912603, + "identical": true, + "index": 3, + "first_diff": null, + "accepted": 14, + "proposed": 38, + "verify_calls": 13, + "draft_calls": 38, + "plain_loads": 464, + "spec_loads": 945 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 1.243565334007144, + "spec_s": 1.3543427080003312, + "identical": true, + "index": 4, + "first_diff": null, + "accepted": 10, + "proposed": 20, + "verify_calls": 5, + "draft_calls": 20, + "plain_loads": 458, + "spec_loads": 682 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 2.040071750001516, + "spec_s": 2.720543083996745, + "identical": true, + "index": 5, + "first_diff": null, + "accepted": 21, + "proposed": 41, + "verify_calls": 14, + "draft_calls": 41, + "plain_loads": 625, + "spec_loads": 1107 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 1.7407667080115061, + "spec_s": 1.7825224169937428, + "identical": true, + "index": 6, + "first_diff": null, + "accepted": 21, + "proposed": 27, + "verify_calls": 6, + "draft_calls": 27, + "plain_loads": 709, + "spec_loads": 955 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 2.2343357499921694, + "spec_s": 3.159669375003432, + "identical": true, + "index": 7, + "first_diff": null, + "accepted": 20, + "proposed": 43, + "verify_calls": 16, + "draft_calls": 43, + "plain_loads": 810, + "spec_loads": 1472 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 1.5214951249945443, + "spec_s": 1.6843749580002623, + "identical": true, + "index": 8, + "first_diff": null, + "accepted": 16, + "proposed": 26, + "verify_calls": 6, + "draft_calls": 26, + "plain_loads": 573, + "spec_loads": 833 + }, + { + "plain_tokens": 30, + "spec_tokens": 30, + "plain_s": 1.9180057079938706, + "spec_s": 2.556749499999569, + "identical": true, + "index": 9, + "first_diff": null, + "accepted": 22, + "proposed": 37, + "verify_calls": 8, + "draft_calls": 37, + "plain_loads": 753, + "spec_loads": 1142 + }, + { + "plain_tokens": 90, + "spec_tokens": 90, + "plain_s": 4.954227167007048, + "spec_s": 5.857818082993617, + "identical": true, + "index": 10, + "first_diff": null, + "accepted": 62, + "proposed": 103, + "verify_calls": 28, + "draft_calls": 103, + "plain_loads": 1152, + "spec_loads": 2105 + }, + { + "plain_tokens": 19, + "spec_tokens": 19, + "plain_s": 1.3099533329950646, + "spec_s": 1.7976152079936583, + "identical": true, + "index": 11, + "first_diff": null, + "accepted": 10, + "proposed": 22, + "verify_calls": 9, + "draft_calls": 22, + "plain_loads": 577, + "spec_loads": 845 + } + ] + } + ] +} \ No newline at end of file diff --git a/reports/speculative/torch_mps_s13996.json b/reports/speculative/torch_mps_s13996.json new file mode 100644 index 0000000..0a7da96 --- /dev/null +++ b/reports/speculative/torch_mps_s13996.json @@ -0,0 +1,424 @@ +{ + "label": "measured", + "time": "2026-11-06T19:24:08", + "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": 12, + "max_new": 128, + "configs": "torch-full,torch-lazy", + "mlx_device": "gpu", + "torch_device": "mps", + "ratio": 1.0, + "miss_policy": "sync_on_demand", + "draft_depth": 1, + "gamma": 4, + "warmup": 1, + "out": "reports/speculative/torch_mps_s13996.json", + "note": "72h MLX training paused with SIGSTOP during measurement; FLORES+ dev first 12 en sentences, en->zh pretrain tags", + "speculative": { + "draft_depth": 1, + "gamma": 4, + "gamma_max": 8, + "adaptive_gamma": true + }, + "eos_id": 0, + "decoding": "greedy, batch 1", + "n_prompts": 12, + "prompt_tokens_mean": 45.6, + "order": "plain/speculative alternate per prompt", + "prompts_source": "FLORES+ dev (eng_Latn/zho_Hans) rows 0-31, field 'en'" + }, + "model": { + "hidden_size": 384, + "num_hidden_layers": 6, + "recurrent_depth": 4, + "num_experts": 36, + "num_experts_per_tok": 2, + "recurrent_attention_mode": "shared_first", + "vocab_size": 16384 + }, + "results": [ + { + "config": "torch-full", + "device": "mps", + "dtype": "float32", + "prompts": 12, + "identical": 12, + "plain_tokens": 436, + "spec_tokens": 436, + "plain_s": 13.583, + "spec_s": 16.836, + "plain_tok_s": 32.1, + "spec_tok_s": 25.9, + "speedup": 0.807, + "speedup_median_per_prompt": 0.831, + "acceptance_rate": 0.4004, + "tokens_per_verify": 1.887, + "plain_ssd_loads": 0, + "spec_ssd_loads": 0, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 1.5062049579864834, + "spec_s": 1.5497601249953732, + "identical": true, + "accepted": 26, + "proposed": 65, + "verify_calls": 27, + "draft_calls": 65, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 42, + "spec_tokens": 42, + "plain_s": 1.1938127080065897, + "spec_s": 2.0981040000042412, + "identical": true, + "accepted": 12, + "proposed": 47, + "verify_calls": 30, + "draft_calls": 47, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 38, + "spec_tokens": 38, + "plain_s": 1.429154584009666, + "spec_s": 1.3815570830047363, + "identical": true, + "accepted": 21, + "proposed": 50, + "verify_calls": 17, + "draft_calls": 50, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 0.8105046669952571, + "spec_s": 1.6450333329994464, + "identical": true, + "accepted": 7, + "proposed": 32, + "verify_calls": 20, + "draft_calls": 32, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 0.6652973749878583, + "spec_s": 0.615217833998031, + "identical": true, + "accepted": 6, + "proposed": 19, + "verify_calls": 9, + "draft_calls": 19, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 1.0107927909994032, + "spec_s": 1.4642756249959348, + "identical": true, + "accepted": 16, + "proposed": 38, + "verify_calls": 20, + "draft_calls": 38, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 0.7648682500002906, + "spec_s": 0.6843811250000726, + "identical": true, + "accepted": 20, + "proposed": 31, + "verify_calls": 7, + "draft_calls": 31, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 1.0611550829926273, + "spec_s": 1.5824520830065012, + "identical": true, + "accepted": 13, + "proposed": 41, + "verify_calls": 24, + "draft_calls": 41, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 0.8924308749992633, + "spec_s": 0.792929584000376, + "identical": true, + "accepted": 10, + "proposed": 26, + "verify_calls": 12, + "draft_calls": 26, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 30, + "spec_tokens": 30, + "plain_s": 0.8431621249910677, + "spec_s": 1.3803647910099244, + "identical": true, + "accepted": 18, + "proposed": 41, + "verify_calls": 12, + "draft_calls": 41, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 90, + "spec_tokens": 90, + "plain_s": 2.869607708998956, + "spec_s": 2.8144312079966767, + "identical": true, + "accepted": 48, + "proposed": 100, + "verify_calls": 42, + "draft_calls": 100, + "plain_loads": 0, + "spec_loads": 0 + }, + { + "plain_tokens": 19, + "spec_tokens": 19, + "plain_s": 0.536420000004, + "spec_s": 0.8276020839984994, + "identical": true, + "accepted": 8, + "proposed": 22, + "verify_calls": 11, + "draft_calls": 22, + "plain_loads": 0, + "spec_loads": 0 + } + ] + }, + { + "config": "torch-lazy", + "device": "mps", + "dtype": "float32", + "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 + }, + "prompts": 12, + "identical": 12, + "plain_tokens": 436, + "spec_tokens": 436, + "plain_s": 25.707, + "spec_s": 29.709, + "plain_tok_s": 16.96, + "spec_tok_s": 14.68, + "speedup": 0.865, + "speedup_median_per_prompt": 0.851, + "acceptance_rate": 0.4004, + "tokens_per_verify": 1.887, + "plain_ssd_loads": 8790, + "spec_ssd_loads": 14344, + "rows": [ + { + "plain_tokens": 53, + "spec_tokens": 53, + "plain_s": 3.034606957997312, + "spec_s": 3.513541416992666, + "identical": true, + "accepted": 26, + "proposed": 65, + "verify_calls": 27, + "draft_calls": 65, + "plain_loads": 1000, + "spec_loads": 1630 + }, + { + "plain_tokens": 42, + "spec_tokens": 42, + "plain_s": 2.4649076250061626, + "spec_s": 3.2090206670109183, + "identical": true, + "accepted": 12, + "proposed": 47, + "verify_calls": 30, + "draft_calls": 47, + "plain_loads": 841, + "spec_loads": 1509 + }, + { + "plain_tokens": 38, + "spec_tokens": 38, + "plain_s": 2.2803186669916613, + "spec_s": 2.631822582989116, + "identical": true, + "accepted": 21, + "proposed": 50, + "verify_calls": 17, + "draft_calls": 50, + "plain_loads": 826, + "spec_loads": 1389 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 1.5297422090079635, + "spec_s": 2.1047971660009352, + "identical": true, + "accepted": 7, + "proposed": 32, + "verify_calls": 20, + "draft_calls": 32, + "plain_loads": 464, + "spec_loads": 884 + }, + { + "plain_tokens": 15, + "spec_tokens": 15, + "plain_s": 1.0554826670122566, + "spec_s": 1.292699541998445, + "identical": true, + "accepted": 6, + "proposed": 19, + "verify_calls": 9, + "draft_calls": 19, + "plain_loads": 458, + "spec_loads": 672 + }, + { + "plain_tokens": 36, + "spec_tokens": 36, + "plain_s": 2.014102125001955, + "spec_s": 2.243040415996802, + "identical": true, + "accepted": 16, + "proposed": 38, + "verify_calls": 20, + "draft_calls": 38, + "plain_loads": 625, + "spec_loads": 977 + }, + { + "plain_tokens": 27, + "spec_tokens": 27, + "plain_s": 1.7251021669944748, + "spec_s": 1.5945601669955067, + "identical": true, + "accepted": 20, + "proposed": 31, + "verify_calls": 7, + "draft_calls": 31, + "plain_loads": 709, + "spec_loads": 949 + }, + { + "plain_tokens": 37, + "spec_tokens": 37, + "plain_s": 2.2237481670017587, + "spec_s": 2.778628290994675, + "identical": true, + "accepted": 13, + "proposed": 41, + "verify_calls": 24, + "draft_calls": 41, + "plain_loads": 810, + "spec_loads": 1384 + }, + { + "plain_tokens": 22, + "spec_tokens": 22, + "plain_s": 1.417515250010183, + "spec_s": 1.6900495419977233, + "identical": true, + "accepted": 10, + "proposed": 26, + "verify_calls": 12, + "draft_calls": 26, + "plain_loads": 573, + "spec_loads": 901 + }, + { + "plain_tokens": 30, + "spec_tokens": 30, + "plain_s": 1.8801868750015274, + "spec_s": 2.253206624998711, + "identical": true, + "accepted": 18, + "proposed": 41, + "verify_calls": 12, + "draft_calls": 41, + "plain_loads": 755, + "spec_loads": 1209 + }, + { + "plain_tokens": 90, + "spec_tokens": 90, + "plain_s": 4.701089416994364, + "spec_s": 4.865764333997504, + "identical": true, + "accepted": 48, + "proposed": 100, + "verify_calls": 42, + "draft_calls": 100, + "plain_loads": 1152, + "spec_loads": 2016 + }, + { + "plain_tokens": 19, + "spec_tokens": 19, + "plain_s": 1.3797815830039326, + "spec_s": 1.5322961250058142, + "identical": true, + "accepted": 8, + "proposed": 22, + "verify_calls": 11, + "draft_calls": 22, + "plain_loads": 577, + "spec_loads": 824 + } + ] + } + ] +} \ No newline at end of file diff --git a/reports/ui_event_contract_20260910.md b/reports/ui_event_contract_20260910.md new file mode 100644 index 0000000..c832b86 --- /dev/null +++ b/reports/ui_event_contract_20260910.md @@ -0,0 +1,17 @@ +# Gradio 空输入事件修复（2026-09-10） + +新页面的 `apply_best` 按钮在真实 Chromium 中发送两个输入；自动调参轮询发送零个输入。当前事件表中两者的编号相邻，分别为 95、96。旧页面保留旧事件表后，把零输入轮询发到 `apply_best`，可以复现用户报告的错误。此前仅通过 Python `process_api` 测试，未覆盖浏览器与服务事件版本不同的场景。 + +本次增加了浏览器与服务实例的事件契约：页面为事件请求附带版本标识，服务在 Gradio 入队前校验版本和输入数量。旧页面的错位请求返回明确的刷新提示，不执行回调；合法的空输入轮询继续正常运行。未携带浏览器标识的合法 API 调用仍受支持。 + +浏览器会定期检查服务版本，并在服务更新或收到版本不匹配响应后重新加载页面。HTML 和 `/config` 禁用缓存，确保刷新取得同一版本的事件表。浏览器同步脚本已加入 wheel 的包数据。 + +验证： + +- 完整测试：`python -m pytest -q tests`，**129 passed**。 +- 8 项 HTTP 契约测试覆盖旧版本请求、零输入错调、正常按钮调用、合法轮询和缓存头。 +- 3 项真实浏览器测试覆盖按钮与轮询的实际输入、错位请求被拒绝后仅刷新一次，以及旧服务页面自动刷新。 +- 在隔离目录构建 wheel，确认包含 `ui/event_contract.py`、`ui/event_contract.js` 和启动器。 +- 已确认当前 `127.0.0.1:7880` 服务启用了校验，`/config` 返回 `Cache-Control: no-store`。 + +首次应用时，仍在旧版本上打开的页面需要手动刷新一次，以加载浏览器同步脚本。后续服务更新会自动刷新同步。无效请求不会被填入默认配置后继续执行。 diff --git a/requirements.txt b/requirements.txt index 32eabd4..dae616a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -torch>=2.4.0 +torch>=2.6.0 transformers>=4.40.0 datasets>=2.18.0 tokenizers>=0.19.0 diff --git a/results/chronos_v2_benchmark_cpu.json b/results/chronos_v2_benchmark_cpu.json new file mode 100644 index 0000000..b18b9da --- /dev/null +++ b/results/chronos_v2_benchmark_cpu.json @@ -0,0 +1,560 @@ +{ + "environment": { + "device": "cpu", + "torch_version": "2.10.0", + "cuda_available": false, + "mps_available": true, + "cpu_threads": 1 + }, + "benchmark_config": { + "device": "cpu", + "seed": 20260904, + "train_steps": 60, + "eval_batches": 4, + "batch_size": 4, + "seq_len": 32, + "prompt_len": 16, + "decode_tokens": 16, + "offload_decode_tokens": 6, + "hidden_size": 64, + "num_hidden_layers": 2, + "num_experts": 4, + "vocab_size": 128, + "lr": 0.001, + "offload_vram_budget_gb": 1e-05, + "simulated_ssd_ms": 1, + "skip_offload": false, + "output": "results/chronos_v2_benchmark_cpu.json" + }, + "variants": [ + { + "name": "v1", + "architecture_version": 1, + "recurrent_enabled": false, + "recurrent_depth": 1, + "adaptive_halting": false, + "virtual_num_hidden_layers": 2, + "train_initial_ce": 5.1463212966918945, + "train_final_ce": 0.04712142050266266, + "train_best_ce": 0.04712142050266266, + "train_tokens_per_sec": 19553.66887067549, + "eval_ce": 0.04558668937534094, + "eval_ppl": 1.0466417333993552, + "eval_next_token_accuracy": 1.0, + "prefill_ms": 0.8016669889912009, + "decode_ms_per_token": 0.68388018735277, + "decode_tokens_per_sec": 1462.244437685638, + "kv_cache_kb": 24.0, + "kv_cache_seq_len": 32, + "output_unique_ratio": 0.5, + "nan_or_inf": false, + "prefill_average_depth": 1.0, + "prefill_early_exit_ratio": 0.0, + "decode_average_depth": 1.0, + "decode_early_exit_ratio": 0.0, + "halt_final_mass": 1.0, + "ponder_loss": 0.0, + "parameter_mb": 2.156236, + "expert_parameter_mb": 1.572864, + "estimated_active_expert_mb": 0.393216, + "rss_after_model_gb": 0.423004150390625, + "rss_after_run_gb": 0.468414306640625 + }, + { + "name": "v2_fixed2", + "architecture_version": 2, + "recurrent_enabled": true, + "recurrent_depth": 2, + "adaptive_halting": false, + "virtual_num_hidden_layers": 4, + "train_initial_ce": 5.3257551193237305, + "train_final_ce": 0.0518440380692482, + "train_best_ce": 0.05169466510415077, + "train_tokens_per_sec": 13462.433041243816, + "eval_ce": 0.05116009991616011, + "eval_ppl": 1.052491383577473, + "eval_next_token_accuracy": 1.0, + "prefill_ms": 1.705917005892843, + "decode_ms_per_token": 1.2805728747480316, + "decode_tokens_per_sec": 780.9005014234448, + "kv_cache_kb": 24.0, + "kv_cache_seq_len": 32, + "output_unique_ratio": 0.5, + "nan_or_inf": false, + "prefill_average_depth": 2.0, + "prefill_early_exit_ratio": 0.0, + "decode_average_depth": 2.0, + "decode_early_exit_ratio": 0.0, + "halt_final_mass": 1.0, + "ponder_loss": 0.0, + "parameter_mb": 2.156492, + "expert_parameter_mb": 1.572864, + "estimated_active_expert_mb": 0.393216, + "rss_after_model_gb": 0.468597412109375, + "rss_after_run_gb": 0.495391845703125 + }, + { + "name": "v2_fixed4", + "architecture_version": 2, + "recurrent_enabled": true, + "recurrent_depth": 4, + "adaptive_halting": false, + "virtual_num_hidden_layers": 8, + "train_initial_ce": 5.269786834716797, + "train_final_ce": 0.052845437079668045, + "train_best_ce": 0.052845437079668045, + "train_tokens_per_sec": 8836.875510210892, + "eval_ce": 0.04947624634951353, + "eval_ppl": 1.0507206334677086, + "eval_next_token_accuracy": 1.0, + "prefill_ms": 3.25354200322181, + "decode_ms_per_token": 2.1801171878905734, + "decode_tokens_per_sec": 458.6909389800164, + "kv_cache_kb": 24.0, + "kv_cache_seq_len": 32, + "output_unique_ratio": 0.5, + "nan_or_inf": false, + "prefill_average_depth": 4.0, + "prefill_early_exit_ratio": 0.0, + "decode_average_depth": 4.0, + "decode_early_exit_ratio": 0.0, + "halt_final_mass": 1.0, + "ponder_loss": 0.0, + "parameter_mb": 2.157004, + "expert_parameter_mb": 1.572864, + "estimated_active_expert_mb": 0.393216, + "rss_after_model_gb": 0.4955596923828125, + "rss_after_run_gb": 0.541748046875 + }, + { + "name": "v2_adaptive4", + "architecture_version": 2, + "recurrent_enabled": true, + "recurrent_depth": 4, + "adaptive_halting": true, + "virtual_num_hidden_layers": 8, + "train_initial_ce": 5.3101806640625, + "train_final_ce": 0.06001603603363037, + "train_best_ce": 0.05005020275712013, + "train_tokens_per_sec": 8716.284919029247, + "eval_ce": 0.04868603590875864, + "eval_ppl": 1.04989067101851, + "eval_next_token_accuracy": 1.0, + "prefill_ms": 2.9648750060005113, + "decode_ms_per_token": 2.1126406245457474, + "decode_tokens_per_sec": 473.34127176268635, + "kv_cache_kb": 24.0, + "kv_cache_seq_len": 32, + "output_unique_ratio": 0.5, + "nan_or_inf": false, + "prefill_average_depth": 2.7917661666870117, + "prefill_early_exit_ratio": 1.0, + "decode_average_depth": 3.2012877464294434, + "decode_early_exit_ratio": 0.0, + "halt_final_mass": 1.0, + "ponder_loss": 2.2012877464294434, + "parameter_mb": 2.157004, + "expert_parameter_mb": 1.572864, + "estimated_active_expert_mb": 0.393216, + "rss_after_model_gb": 0.5419158935546875, + "rss_after_run_gb": 0.5555572509765625 + } + ], + "paged_expert_runtime": { + "sync_on_demand": { + "tokens_generated": 6, + "elapsed_s": 0.034, + "tokens_per_sec": 176.33, + "prefill_time_s": 0.019, + "decode_time_s": 0.015, + "prefill_rss_delta_gb": 0.0, + "decode_rss_delta_gb": 0.0, + "rss_before_prefill_gb": 0.558, + "rss_after_prefill_gb": 0.558, + "rss_after_decode_gb": 0.558, + "cache_hit_rate": 0.5, + "resident_hit_rate": 0.5, + "resident_vram_hit_rate": 0.5, + "ram_promotion_hit_rate": 1.0, + "ssd_cold_miss_rate": 0.0, + "cache_hits": 21, + "cache_misses": 21, + "fallback_weight_rate": 0.0, + "resident_miss_weight_rate": 0.5, + "fallback_weight_mass": 0.0, + "resident_miss_weight_mass": 21.0, + "hit_weight_mass": 21.0, + "total_route_weight_mass": 42.0, + "on_demand_weight_mass": 38.0, + "on_demand_weight_rate": 0.9048, + "on_demand_loads": 16, + "on_demand_load_time_s": 0.0113, + "sync_ssd_loads": 0, + "async_cold_miss_prefetches": 0, + "resident_vram_hits": 0, + "resident_ram_hits": 16, + "quality_safe": true, + "miss_policy": "sync_on_demand", + "prediction_hit_rate": 0.5, + "prediction_topk_recall": 0.5, + "cluster_recall": 0.0, + "prediction_hits": 3, + "prediction_total": 6, + "cluster_prediction_hits": 0, + "cluster_prediction_total": 0, + "average_ponder_steps": 2.0, + "early_exit_ratio": 0.0, + "route_layer_stats": [ + { + "layer": 0, + "hit_rate": 1.0, + "fallback_weight_rate": 0.0, + "hits": 1, + "misses": 0, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 0 + ] + }, + { + "layer": 1, + "hit_rate": 0.0, + "fallback_weight_rate": 1.0, + "hits": 0, + "misses": 1, + "activated_experts": [ + 1 + ], + "loaded_experts": [ + 0 + ] + } + ], + "prefill_expert_ids": [ + 2 + ], + "prefill_ram_expert_ids": [ + 2, + 0, + 1 + ], + "activated_expert_ids": [ + 0, + 1, + 2 + ], + "predicted_current_expert_ids": [ + 0, + 1, + 2 + ], + "prefetch_submitted_ids": [ + 0, + 1, + 2 + ], + "promoted_expert_ids": [ + 0, + 1, + 2 + ], + "on_demand_expert_ids": [ + 0, + 1, + 2 + ], + "fallback_expert_ids": [], + "prefetch_submitted_total": 10, + "promoted_total": 19, + "vram_experts": 1, + "vram_capacity": 1, + "ram_experts": 3, + "ram_capacity_dynamic": 4, + "expert_size_kb": 96, + "pinned_ram_used_gb": 0.001, + "available_ram_gb": 94.5, + "pinned_ram_fraction": 0.0, + "h2d_stream": "default", + "storage_format": "pt", + "cluster_aware": false, + "layer_aware": false, + "num_clusters": 0, + "prefetch_hits": 0, + "prefetch_misses": 0, + "total_requests": 15, + "hit_rate": 0.0, + "prefetch_queue_drops": 0, + "prefetch_wait_time_s": 0.0, + "exact_token_match_to_full_resident": 1.0, + "full_expert_weight_mb": 1.572864, + "resident_expert_weight_mb": 0.098304, + "ssd_artifact_bytes": 803952, + "estimated_expert_memory_savings": 0.9375 + }, + "on_demand": { + "tokens_generated": 6, + "elapsed_s": 0.031, + "tokens_per_sec": 190.71, + "prefill_time_s": 0.017, + "decode_time_s": 0.015, + "prefill_rss_delta_gb": 0.0, + "decode_rss_delta_gb": 0.0, + "rss_before_prefill_gb": 0.558, + "rss_after_prefill_gb": 0.558, + "rss_after_decode_gb": 0.558, + "cache_hit_rate": 1.0, + "resident_hit_rate": 0.5, + "resident_vram_hit_rate": 0.5, + "ram_promotion_hit_rate": 1.0, + "ssd_cold_miss_rate": 0.0, + "cache_hits": 21, + "cache_misses": 21, + "fallback_weight_rate": 0.0, + "resident_miss_weight_rate": 0.5, + "fallback_weight_mass": 0.0, + "resident_miss_weight_mass": 21.0, + "hit_weight_mass": 21.0, + "total_route_weight_mass": 42.0, + "on_demand_weight_mass": 38.0, + "on_demand_weight_rate": 0.9048, + "on_demand_loads": 16, + "on_demand_load_time_s": 0.0111, + "sync_ssd_loads": 0, + "async_cold_miss_prefetches": 0, + "resident_vram_hits": 0, + "resident_ram_hits": 16, + "quality_safe": false, + "miss_policy": "on_demand", + "prediction_hit_rate": 0.5, + "prediction_topk_recall": 0.5, + "cluster_recall": 0.0, + "prediction_hits": 3, + "prediction_total": 6, + "cluster_prediction_hits": 0, + "cluster_prediction_total": 0, + "average_ponder_steps": 2.0, + "early_exit_ratio": 0.0, + "route_layer_stats": [ + { + "layer": 0, + "hit_rate": 1.0, + "fallback_weight_rate": 0.0, + "hits": 1, + "misses": 0, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 0 + ] + }, + { + "layer": 1, + "hit_rate": 0.0, + "fallback_weight_rate": 1.0, + "hits": 0, + "misses": 1, + "activated_experts": [ + 1 + ], + "loaded_experts": [ + 0 + ] + } + ], + "prefill_expert_ids": [ + 2 + ], + "prefill_ram_expert_ids": [ + 2, + 0, + 1 + ], + "activated_expert_ids": [ + 0, + 1, + 2 + ], + "predicted_current_expert_ids": [ + 0, + 1, + 2 + ], + "prefetch_submitted_ids": [ + 0, + 1, + 2 + ], + "promoted_expert_ids": [ + 0, + 1, + 2 + ], + "on_demand_expert_ids": [ + 0, + 1, + 2 + ], + "fallback_expert_ids": [], + "prefetch_submitted_total": 10, + "promoted_total": 19, + "vram_experts": 1, + "vram_capacity": 1, + "ram_experts": 3, + "ram_capacity_dynamic": 4, + "expert_size_kb": 96, + "pinned_ram_used_gb": 0.001, + "available_ram_gb": 94.4, + "pinned_ram_fraction": 0.0, + "h2d_stream": "default", + "storage_format": "pt", + "cluster_aware": false, + "layer_aware": false, + "num_clusters": 0, + "prefetch_hits": 0, + "prefetch_misses": 0, + "total_requests": 15, + "hit_rate": 0.0, + "prefetch_queue_drops": 0, + "prefetch_wait_time_s": 0.0, + "exact_token_match_to_full_resident": 1.0, + "full_expert_weight_mb": 1.572864, + "resident_expert_weight_mb": 0.098304, + "ssd_artifact_bytes": 803952, + "estimated_expert_memory_savings": 0.9375 + }, + "fallback_diagnostic": { + "tokens_generated": 6, + "elapsed_s": 0.018, + "tokens_per_sec": 327.25, + "prefill_time_s": 0.009, + "decode_time_s": 0.01, + "prefill_rss_delta_gb": 0.0, + "decode_rss_delta_gb": 0.0, + "rss_before_prefill_gb": 0.558, + "rss_after_prefill_gb": 0.558, + "rss_after_decode_gb": 0.558, + "cache_hit_rate": 0.4762, + "resident_hit_rate": 0.4762, + "resident_vram_hit_rate": 0.4762, + "ram_promotion_hit_rate": 0.0, + "ssd_cold_miss_rate": 0.0, + "cache_hits": 20, + "cache_misses": 22, + "fallback_weight_rate": 0.5238, + "resident_miss_weight_rate": 0.5238, + "fallback_weight_mass": 22.0, + "resident_miss_weight_mass": 22.0, + "hit_weight_mass": 20.0, + "total_route_weight_mass": 42.0, + "on_demand_weight_mass": 0.0, + "on_demand_weight_rate": 0.0, + "on_demand_loads": 0, + "on_demand_load_time_s": 0.0, + "sync_ssd_loads": 0, + "async_cold_miss_prefetches": 0, + "resident_vram_hits": 0, + "resident_ram_hits": 0, + "quality_safe": false, + "miss_policy": "fallback_diagnostic", + "prediction_hit_rate": 0.5, + "prediction_topk_recall": 0.5, + "cluster_recall": 0.0, + "prediction_hits": 3, + "prediction_total": 6, + "cluster_prediction_hits": 0, + "cluster_prediction_total": 0, + "average_ponder_steps": 2.0, + "early_exit_ratio": 0.0, + "route_layer_stats": [ + { + "layer": 0, + "hit_rate": 1.0, + "fallback_weight_rate": 0.0, + "hits": 1, + "misses": 0, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 0 + ] + }, + { + "layer": 1, + "hit_rate": 0.0, + "fallback_weight_rate": 1.0, + "hits": 0, + "misses": 1, + "activated_experts": [ + 1 + ], + "loaded_experts": [ + 0 + ] + } + ], + "prefill_expert_ids": [ + 2 + ], + "prefill_ram_expert_ids": [ + 2, + 0, + 1 + ], + "activated_expert_ids": [ + 0, + 1, + 2 + ], + "predicted_current_expert_ids": [ + 0, + 1, + 2 + ], + "prefetch_submitted_ids": [ + 0, + 1, + 2 + ], + "promoted_expert_ids": [ + 0, + 1 + ], + "on_demand_expert_ids": [], + "fallback_expert_ids": [ + 0, + 1, + 2 + ], + "prefetch_submitted_total": 10, + "promoted_total": 2, + "vram_experts": 1, + "vram_capacity": 1, + "ram_experts": 3, + "ram_capacity_dynamic": 4, + "expert_size_kb": 96, + "pinned_ram_used_gb": 0.001, + "available_ram_gb": 94.4, + "pinned_ram_fraction": 0.0, + "h2d_stream": "default", + "storage_format": "pt", + "cluster_aware": false, + "layer_aware": false, + "num_clusters": 0, + "prefetch_hits": 0, + "prefetch_misses": 0, + "total_requests": 15, + "hit_rate": 0.0, + "prefetch_queue_drops": 0, + "prefetch_wait_time_s": 0.0, + "exact_token_match_to_full_resident": 1.0, + "full_expert_weight_mb": 1.572864, + "resident_expert_weight_mb": 0.098304, + "ssd_artifact_bytes": 803952, + "estimated_expert_memory_savings": 0.9375 + } + } +} \ No newline at end of file diff --git a/results/chronos_v2_benchmark_mps.json b/results/chronos_v2_benchmark_mps.json new file mode 100644 index 0000000..a2a2ad3 --- /dev/null +++ b/results/chronos_v2_benchmark_mps.json @@ -0,0 +1,561 @@ +{ + "environment": { + "device": "mps", + "torch_version": "2.10.0", + "cuda_available": false, + "mps_available": true, + "cpu_threads": 1 + }, + "benchmark_config": { + "device": "mps", + "seed": 20260904, + "train_steps": 20, + "eval_batches": 2, + "batch_size": 2, + "seq_len": 16, + "prompt_len": 8, + "decode_tokens": 8, + "offload_decode_tokens": 4, + "hidden_size": 32, + "num_hidden_layers": 2, + "num_experts": 2, + "vocab_size": 128, + "lr": 0.001, + "offload_vram_budget_gb": 1e-05, + "simulated_ssd_ms": 1, + "skip_offload": false, + "output": "results/chronos_v2_benchmark_mps.json" + }, + "variants": [ + { + "name": "v1", + "architecture_version": 1, + "recurrent_enabled": false, + "recurrent_depth": 1, + "adaptive_halting": false, + "virtual_num_hidden_layers": 2, + "train_initial_ce": 4.968942165374756, + "train_final_ce": 1.8350716829299927, + "train_best_ce": 1.8350716829299927, + "train_tokens_per_sec": 354.24736589101275, + "eval_ce": 1.7509395480155945, + "eval_ppl": 5.760011942265277, + "eval_next_token_accuracy": 1.0, + "prefill_ms": 15.517167004873045, + "decode_ms_per_token": 8.94931775110308, + "decode_tokens_per_sec": 111.74036142327624, + "kv_cache_kb": 8.0, + "kv_cache_seq_len": 16, + "output_unique_ratio": 1.0, + "nan_or_inf": false, + "prefill_average_depth": 1.0, + "prefill_early_exit_ratio": 0.0, + "decode_average_depth": 1.0, + "decode_early_exit_ratio": 0.0, + "halt_final_mass": 1.0, + "ponder_loss": 0.0, + "parameter_mb": 0.36382, + "expert_parameter_mb": 0.196608, + "estimated_active_expert_mb": 0.098304, + "rss_after_model_gb": 0.43060302734375, + "rss_after_run_gb": 0.867706298828125 + }, + { + "name": "v2_fixed2", + "architecture_version": 2, + "recurrent_enabled": true, + "recurrent_depth": 2, + "adaptive_halting": false, + "virtual_num_hidden_layers": 4, + "train_initial_ce": 5.005463600158691, + "train_final_ce": 1.712158441543579, + "train_best_ce": 1.712158441543579, + "train_tokens_per_sec": 541.7324814835845, + "eval_ce": 1.6443588733673096, + "eval_ppl": 5.177689288095263, + "eval_next_token_accuracy": 1.0, + "prefill_ms": 16.507457999978215, + "decode_ms_per_token": 14.298031250291388, + "decode_tokens_per_sec": 69.93969886445873, + "kv_cache_kb": 8.0, + "kv_cache_seq_len": 16, + "output_unique_ratio": 1.0, + "nan_or_inf": false, + "prefill_average_depth": 2.0, + "prefill_early_exit_ratio": 0.0, + "decode_average_depth": 2.0, + "decode_early_exit_ratio": 0.0, + "halt_final_mass": 1.0, + "ponder_loss": 0.0, + "parameter_mb": 0.363948, + "expert_parameter_mb": 0.196608, + "estimated_active_expert_mb": 0.098304, + "rss_after_model_gb": 0.8681640625, + "rss_after_run_gb": 0.9619903564453125 + }, + { + "name": "v2_fixed4", + "architecture_version": 2, + "recurrent_enabled": true, + "recurrent_depth": 4, + "adaptive_halting": false, + "virtual_num_hidden_layers": 8, + "train_initial_ce": 5.039073467254639, + "train_final_ce": 1.647542953491211, + "train_best_ce": 1.647542953491211, + "train_tokens_per_sec": 353.50717785321996, + "eval_ce": 1.5743180513381958, + "eval_ppl": 4.827448423350792, + "eval_next_token_accuracy": 1.0, + "prefill_ms": 35.65316599269863, + "decode_ms_per_token": 25.839520751105738, + "decode_tokens_per_sec": 38.7004081705814, + "kv_cache_kb": 8.0, + "kv_cache_seq_len": 16, + "output_unique_ratio": 1.0, + "nan_or_inf": false, + "prefill_average_depth": 4.0, + "prefill_early_exit_ratio": 0.0, + "decode_average_depth": 4.0, + "decode_early_exit_ratio": 0.0, + "halt_final_mass": 1.0, + "ponder_loss": 0.0, + "parameter_mb": 0.364204, + "expert_parameter_mb": 0.196608, + "estimated_active_expert_mb": 0.098304, + "rss_after_model_gb": 0.96246337890625, + "rss_after_run_gb": 1.0138702392578125 + }, + { + "name": "v2_adaptive4", + "architecture_version": 2, + "recurrent_enabled": true, + "recurrent_depth": 4, + "adaptive_halting": true, + "virtual_num_hidden_layers": 8, + "train_initial_ce": 4.99564790725708, + "train_final_ce": 1.7389887571334839, + "train_best_ce": 1.7389887571334839, + "train_tokens_per_sec": 309.9091951428779, + "eval_ce": 1.6488620042800903, + "eval_ppl": 5.201057676752897, + "eval_next_token_accuracy": 1.0, + "prefill_ms": 25.684000007458963, + "decode_ms_per_token": 26.27782812487567, + "decode_tokens_per_sec": 38.05489537597512, + "kv_cache_kb": 8.0, + "kv_cache_seq_len": 16, + "output_unique_ratio": 1.0, + "nan_or_inf": false, + "prefill_average_depth": 2.2974767684936523, + "prefill_early_exit_ratio": 1.0, + "decode_average_depth": 2.270901679992676, + "decode_early_exit_ratio": 0.0, + "halt_final_mass": 1.0, + "ponder_loss": 1.2709015607833862, + "parameter_mb": 0.364204, + "expert_parameter_mb": 0.196608, + "estimated_active_expert_mb": 0.098304, + "rss_after_model_gb": 1.0142364501953125, + "rss_after_run_gb": 1.030548095703125 + } + ], + "paged_expert_runtime": { + "sync_on_demand": { + "tokens_generated": 4, + "elapsed_s": 0.206, + "tokens_per_sec": 19.41, + "prefill_time_s": 0.11, + "decode_time_s": 0.096, + "prefill_rss_delta_gb": 0.013, + "decode_rss_delta_gb": 0.001, + "rss_before_prefill_gb": 1.034, + "rss_after_prefill_gb": 1.047, + "rss_after_decode_gb": 1.048, + "mps_allocated_before_prefill_gb": 0.003394, + "mps_driver_allocated_before_prefill_gb": 0.028473, + "mps_allocated_after_prefill_gb": 0.003414, + "mps_driver_allocated_after_prefill_gb": 0.028519, + "mps_allocated_after_decode_gb": 0.00341, + "mps_driver_allocated_after_decode_gb": 0.028519, + "cache_hit_rate": 0.8182, + "resident_hit_rate": 0.8182, + "resident_vram_hit_rate": 0.8182, + "ram_promotion_hit_rate": 1.0, + "ssd_cold_miss_rate": 0.0, + "cache_hits": 18, + "cache_misses": 4, + "fallback_weight_rate": 0.0, + "resident_miss_weight_rate": 0.1818, + "fallback_weight_mass": 0.0, + "resident_miss_weight_mass": 4.0, + "hit_weight_mass": 18.0, + "total_route_weight_mass": 22.0, + "on_demand_weight_mass": 12.0, + "on_demand_weight_rate": 0.5455, + "on_demand_loads": 9, + "on_demand_load_time_s": 0.0385, + "sync_ssd_loads": 0, + "async_cold_miss_prefetches": 0, + "resident_vram_hits": 0, + "resident_ram_hits": 9, + "quality_safe": true, + "miss_policy": "sync_on_demand", + "prediction_hit_rate": 0.6667, + "prediction_topk_recall": 0.6667, + "cluster_recall": 0.0, + "prediction_hits": 2, + "prediction_total": 3, + "cluster_prediction_hits": 0, + "cluster_prediction_total": 0, + "average_ponder_steps": 2.0, + "early_exit_ratio": 0.0, + "route_layer_stats": [ + { + "layer": 0, + "hit_rate": 1.0, + "fallback_weight_rate": 0.0, + "hits": 1, + "misses": 0, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 0 + ] + }, + { + "layer": 1, + "hit_rate": 1.0, + "fallback_weight_rate": 0.0, + "hits": 1, + "misses": 0, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 0 + ] + } + ], + "prefill_expert_ids": [ + 1 + ], + "prefill_ram_expert_ids": [ + 1, + 0 + ], + "activated_expert_ids": [ + 0, + 1 + ], + "predicted_current_expert_ids": [ + 0, + 1 + ], + "prefetch_submitted_ids": [ + 0, + 1 + ], + "promoted_expert_ids": [ + 0, + 1 + ], + "on_demand_expert_ids": [ + 0, + 1 + ], + "fallback_expert_ids": [], + "prefetch_submitted_total": 8, + "promoted_total": 11, + "vram_experts": 1, + "vram_capacity": 1, + "ram_experts": 2, + "ram_capacity_dynamic": 4, + "expert_size_kb": 24, + "pinned_ram_used_gb": 0.0, + "available_ram_gb": 94.1, + "pinned_ram_fraction": 0.0, + "h2d_stream": "default", + "storage_format": "pt", + "cluster_aware": false, + "layer_aware": false, + "num_clusters": 0, + "prefetch_hits": 0, + "prefetch_misses": 0, + "total_requests": 10, + "hit_rate": 0.0, + "prefetch_queue_drops": 0, + "prefetch_wait_time_s": 0.0, + "exact_token_match_to_full_resident": 1.0, + "full_expert_weight_mb": 0.196608, + "resident_expert_weight_mb": 0.024576, + "ssd_artifact_bytes": 107064, + "estimated_expert_memory_savings": 0.875 + }, + "on_demand": { + "tokens_generated": 4, + "elapsed_s": 0.158, + "tokens_per_sec": 25.3, + "prefill_time_s": 0.069, + "decode_time_s": 0.089, + "prefill_rss_delta_gb": 0.0, + "decode_rss_delta_gb": 0.0, + "rss_before_prefill_gb": 1.048, + "rss_after_prefill_gb": 1.048, + "rss_after_decode_gb": 1.048, + "mps_allocated_before_prefill_gb": 0.003397, + "mps_driver_allocated_before_prefill_gb": 0.028519, + "mps_allocated_after_prefill_gb": 0.003416, + "mps_driver_allocated_after_prefill_gb": 0.028519, + "mps_allocated_after_decode_gb": 0.003412, + "mps_driver_allocated_after_decode_gb": 0.028519, + "cache_hit_rate": 1.0, + "resident_hit_rate": 0.8182, + "resident_vram_hit_rate": 0.8182, + "ram_promotion_hit_rate": 1.0, + "ssd_cold_miss_rate": 0.0, + "cache_hits": 18, + "cache_misses": 4, + "fallback_weight_rate": 0.0, + "resident_miss_weight_rate": 0.1818, + "fallback_weight_mass": 0.0, + "resident_miss_weight_mass": 4.0, + "hit_weight_mass": 18.0, + "total_route_weight_mass": 22.0, + "on_demand_weight_mass": 12.0, + "on_demand_weight_rate": 0.5455, + "on_demand_loads": 9, + "on_demand_load_time_s": 0.0369, + "sync_ssd_loads": 0, + "async_cold_miss_prefetches": 0, + "resident_vram_hits": 0, + "resident_ram_hits": 9, + "quality_safe": false, + "miss_policy": "on_demand", + "prediction_hit_rate": 0.6667, + "prediction_topk_recall": 0.6667, + "cluster_recall": 0.0, + "prediction_hits": 2, + "prediction_total": 3, + "cluster_prediction_hits": 0, + "cluster_prediction_total": 0, + "average_ponder_steps": 2.0, + "early_exit_ratio": 0.0, + "route_layer_stats": [ + { + "layer": 0, + "hit_rate": 1.0, + "fallback_weight_rate": 0.0, + "hits": 1, + "misses": 0, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 0 + ] + }, + { + "layer": 1, + "hit_rate": 1.0, + "fallback_weight_rate": 0.0, + "hits": 1, + "misses": 0, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 0 + ] + } + ], + "prefill_expert_ids": [ + 1 + ], + "prefill_ram_expert_ids": [ + 1, + 0 + ], + "activated_expert_ids": [ + 0, + 1 + ], + "predicted_current_expert_ids": [ + 0, + 1 + ], + "prefetch_submitted_ids": [ + 0, + 1 + ], + "promoted_expert_ids": [ + 0, + 1 + ], + "on_demand_expert_ids": [ + 0, + 1 + ], + "fallback_expert_ids": [], + "prefetch_submitted_total": 8, + "promoted_total": 11, + "vram_experts": 1, + "vram_capacity": 1, + "ram_experts": 2, + "ram_capacity_dynamic": 4, + "expert_size_kb": 24, + "pinned_ram_used_gb": 0.0, + "available_ram_gb": 94.1, + "pinned_ram_fraction": 0.0, + "h2d_stream": "default", + "storage_format": "pt", + "cluster_aware": false, + "layer_aware": false, + "num_clusters": 0, + "prefetch_hits": 0, + "prefetch_misses": 0, + "total_requests": 10, + "hit_rate": 0.0, + "prefetch_queue_drops": 0, + "prefetch_wait_time_s": 0.0, + "exact_token_match_to_full_resident": 1.0, + "full_expert_weight_mb": 0.196608, + "resident_expert_weight_mb": 0.024576, + "ssd_artifact_bytes": 107064, + "estimated_expert_memory_savings": 0.875 + }, + "fallback_diagnostic": { + "tokens_generated": 4, + "elapsed_s": 0.123, + "tokens_per_sec": 32.51, + "prefill_time_s": 0.056, + "decode_time_s": 0.067, + "prefill_rss_delta_gb": 0.0, + "decode_rss_delta_gb": 0.0, + "rss_before_prefill_gb": 1.048, + "rss_after_prefill_gb": 1.048, + "rss_after_decode_gb": 1.049, + "mps_allocated_before_prefill_gb": 0.003394, + "mps_driver_allocated_before_prefill_gb": 0.028519, + "mps_allocated_after_prefill_gb": 0.003415, + "mps_driver_allocated_after_prefill_gb": 0.028519, + "mps_allocated_after_decode_gb": 0.00341, + "mps_driver_allocated_after_decode_gb": 0.028519, + "cache_hit_rate": 0.6818, + "resident_hit_rate": 0.6818, + "resident_vram_hit_rate": 0.6818, + "ram_promotion_hit_rate": 0.0, + "ssd_cold_miss_rate": 0.0, + "cache_hits": 15, + "cache_misses": 7, + "fallback_weight_rate": 0.3182, + "resident_miss_weight_rate": 0.3182, + "fallback_weight_mass": 7.0, + "resident_miss_weight_mass": 7.0, + "hit_weight_mass": 15.0, + "total_route_weight_mass": 22.0, + "on_demand_weight_mass": 0.0, + "on_demand_weight_rate": 0.0, + "on_demand_loads": 0, + "on_demand_load_time_s": 0.0, + "sync_ssd_loads": 0, + "async_cold_miss_prefetches": 0, + "resident_vram_hits": 0, + "resident_ram_hits": 0, + "quality_safe": false, + "miss_policy": "fallback_diagnostic", + "prediction_hit_rate": 0.3333, + "prediction_topk_recall": 0.3333, + "cluster_recall": 0.0, + "prediction_hits": 1, + "prediction_total": 3, + "cluster_prediction_hits": 0, + "cluster_prediction_total": 0, + "average_ponder_steps": 2.0, + "early_exit_ratio": 0.0, + "route_layer_stats": [ + { + "layer": 0, + "hit_rate": 0.0, + "fallback_weight_rate": 1.0, + "hits": 0, + "misses": 1, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 1 + ] + }, + { + "layer": 1, + "hit_rate": 0.0, + "fallback_weight_rate": 1.0, + "hits": 0, + "misses": 1, + "activated_experts": [ + 0 + ], + "loaded_experts": [ + 1 + ] + } + ], + "prefill_expert_ids": [ + 1 + ], + "prefill_ram_expert_ids": [ + 1, + 0 + ], + "activated_expert_ids": [ + 0, + 1 + ], + "predicted_current_expert_ids": [ + 0, + 1 + ], + "prefetch_submitted_ids": [ + 0, + 1 + ], + "promoted_expert_ids": [ + 0, + 1 + ], + "on_demand_expert_ids": [], + "fallback_expert_ids": [ + 0, + 1 + ], + "prefetch_submitted_total": 6, + "promoted_total": 2, + "vram_experts": 1, + "vram_capacity": 1, + "ram_experts": 2, + "ram_capacity_dynamic": 4, + "expert_size_kb": 24, + "pinned_ram_used_gb": 0.0, + "available_ram_gb": 94.1, + "pinned_ram_fraction": 0.0, + "h2d_stream": "default", + "storage_format": "pt", + "cluster_aware": false, + "layer_aware": false, + "num_clusters": 0, + "prefetch_hits": 0, + "prefetch_misses": 0, + "total_requests": 9, + "hit_rate": 0.0, + "prefetch_queue_drops": 0, + "prefetch_wait_time_s": 0.0, + "exact_token_match_to_full_resident": 1.0, + "full_expert_weight_mb": 0.196608, + "resident_expert_weight_mb": 0.024576, + "ssd_artifact_bytes": 107064, + "estimated_expert_memory_savings": 0.875 + } + } +} \ No newline at end of file diff --git a/setup.py b/setup.py index c6c0b23..95976c3 100644 --- a/setup.py +++ b/setup.py @@ -1,23 +1,4 @@ -from setuptools import setup, find_packages +"""Compatibility shim: all metadata lives in pyproject.toml (PEP 621).""" +from setuptools import setup -setup( - name="project-chronos", - version="0.1.0", - description="On-device low-latency lookahead dual-layer MoE inference architecture", - packages=find_packages(), - python_requires=">=3.10", - install_requires=[ - "torch>=2.1.0", - "transformers>=4.40.0", - "datasets>=2.18.0", - "safetensors>=0.4.0", - "optuna>=3.6.0", - "numpy>=1.24.0", - ], - entry_points={ - "console_scripts": [ - "chronos-train=train_chronos:main", - "chronos-eval=chronos.eval.io_profiler:main", - ], - }, -) +setup() diff --git a/tests/fixtures/v2_golden_logits.pt b/tests/fixtures/v2_golden_logits.pt new file mode 100644 index 0000000000000000000000000000000000000000..153556d868d24cdbc571fae4ef774560124ec1d7 GIT binary patch literal 21047 zcmbr`1$0#B+xPq6?(XjHn#{%_2?PQO5#sLdHpxt8;%*QixVvle1zN1di@TH}#a#;L z;$Pl(owc6#IcGiRWHp7-WM*e(@B6;4-{;!QG^tY6z@TPLga7i;%Al4(_+ST*0AE)( ze+S`MNrUy7`9$1UrNet{qglXHTmd^5r4&l};XRK~?_!mfEHP{vNac{j=IZ zRjn$?mkbQ5_Mb;~2@DIWZc;NesKzW41CvT7HFZHXO=_A{(goE@b#V#s^>uRz^$PG0 zscjbQ>K5$gI?>A|w02PKUgkk{LWAm>2i3EZpV7cQ%-1)lzNb9D=l^=n$-V(jq5u0y z8~o=<2L(0stUleKVQ5gJpU*bkpt3%wvHX$z#-x%ys7b2lkH>Br8q~}@sJT^jd2ClV zmjG9{pcel=y63--G&rc`k4LH;8r16FN2>InM`}IYpq2dd*e9^=jIaKMp_@5%jpSJo%5$ z|Na148w~%i2jCnP>K5?hw3$49@c-a>|JO12dWDAiy8SqYWW)Qdew+$FKL7h;Xkjqy zzaGPnt7^c&$>Rt8&kx|oOa9}-`sZ)X{kXK`WBR|p$Tj@$N7aV^{@UW_fB$%;eti5m z^1#Z`iXy-N@%jIBAcoT_=9ZP-$~_Ao}R#i z?!-UV&_f+Zz45-B2+F62CLRA>S&V2Ig8RiHEIfT#JbwX|EF$pSWXG*)0Zgrt#JDxN zyf%)ZU7PvL8W=^wIae$Pq~UvQ6?=AQ8GpnD)yuv-YLrAu9Y0KMf++2zBXW%;Pp`U= zuJ>SO7@yWbgVas2Q z)f|cpr+I2TtCq*I`i+LY76p8BvB1sKnym8$w5XH7lu!r8?2IG*b`Y~>>G8d)<*koB z-pf2mYMsRMXJ3rQEHP5z?0+?K)%R6q)(OSv#1!TBV4cdLPJ&VS$9G1Z>K-sMQ`#wE zm76I`8*f!*)+#k>`uj9hTitvTZYHv0xuHVFy>b(OTjkOE`5d1=T4`Td#nj00_ztb5 z7}amD9C#B;dYF#p8y3*5Wd^Hid$Gqn6}2vtz&Ft}3{0kT^L%WR1K8~A#pDSABrhw* zzLq-^<|Wg3o;yXAQuxq3fwj*y47(h_h8e4I+nYwcFIr~LF683}EroXqa4EIrsKEk? zYUOisL_S|mP9y1L8Fw3d6Sp~-b$yc=W0B0V)I`1*cw&E5&kdE1*p2CAU5Q8c$(Gqe zwbcJ7jef0N8TPS|o>@uURD5`27LD)81zeq@!TL%ZJ+|uk(KKAak{MmK(X=s+<*$NF)|?NeK~6Fg`|DZ# zdoDZeys5M+p7eIfL>x)vuV3T1WfVl!cYegj6qDD>54~v#{p5NYtBFDHsHHsLjI%?N znEEUop9?P3YNzGT*i`NtPN7AgP;w6^V|px&M?-ztV7QoW5A3LQ)tQWdaNb$_u%l%v z36Vam8@+&uBTU$IKbda@a?MuD<5tsbUflLz{R=y0Z_DBEmkicE(sDX0gtqCy^tDgp z)0G7F?8@VdR!y((TUfF(9GR_yh_kASewHb zn{c))N+Zp}oAVJ_Og_;-G5>9r;#m>O-+AFg9Gbxxzcl{17E8rWJp(>X!EBr@K_PEc z*Wc-AQE#BKt!g~+TQWIvsH^fMI*uD7^|bIP#L~@&^T)hMs2)JM^K7jDbYiQV?~|(z zRn|SSq)~Ve<>Q1wNbKK%tWKUiocnT(|3}2_)Sf$T;KESL~-(%mg;sI z9t;cNTHn#i`2GQmu}s2la0qJ-7f=^p>V2F}uWc!GYmr4C=Kyx+STe->k*dl4I40-j z)3skF^N-GA+0^YgD`{+(=kGf(2oooJs@b}dyULrQ!AbPFevVjcW2Rg!VLaCu(kZ~%v7p9NucV{O!DSC5q3e%JlTI%j<(>|*c6iX zT2sR|lAxvil$p+jEPFPK$wsBBK08www9Jw6g~1dL&BtzGI!y}4(==<2Qg%F)Z^`-G zJQs@kP7v9H(im2vrsH8*&#V(U;Tgcvznqz;O{PbWRDPBH$va#}J^XpNDw60&;oQy7 zV$W$c{fr{8oH>`VBZ?>siDBa$U%Ea_!by&8_pE3<=R7pZsb@~UZ!N{*?n0Ii&~y7q zD3f04sd+t(FALI%Y-@+UES6(YnQ+6FyP-M`G%MgxaT3?wrlM`%y<&dfQHBS1PgCX^ z`tnO&G+`r6arZmug6lgQWzhmWV7Omno}LdqO^CRIMSPE7G`AITX61y z4d+Kh(6EUk&!))ledp4sXawzIr?b@251Xc;Y;H7bth40iXk$X&SrR_Mm(x*` zdFMW!(gW@+x#~z}K^D*VxH9A79R3NiCV7h^p~oY*=^Mf7al;sPOiwMB@mwr)VB7^8 zMm?WJnR_6e+9YCsbOsZ?h0(X_1g`pyVMxCayk0u7sc|%_r-95(m1}md6OotZu-|+R zO}i@$*fp9PzG3)kTVnmdh&9UtsUL5J`i(n5mFJ>4YC~?7XaS*GlerxE7$et+nx{Utv|X*oXmK}2 zew&5GO+yZM@n+Xv2m0+WVeBkZj-*ekX!iAp%CGOT@^D@q(qIB3^MW!u3?HN?cV zd|iHA#YWpxMrUp;R=IvTWpuZ z1%1dFqmB)S@ooPw^qrTM=cMJS)~>Z8ePtB&PI;kP^=yt2O+6tPAt^H|L4GUs`$1_sJgSTt9|a zYu%Zj$Q{Ng~{j`oacl0)g{ zVI1x^gyp_d`KXw%d-Vi5|2-Ga_iFkCxKUimi}7KhJX;dWhPqbFo^QtbJXI3L;0t`m6aZM!UEiAdaI)*iwwzN1kmI1*dn0e5ds?Ds)8!3BYQ$wbgX$hNV zgOPnA!73qx-U!b*|VkTbna$55%i@aE!z6xF~f~YgetS%ccAsKG?ng)2W>9KF}FBh^>wuw z-zLss?VT{CQ{NOUdc|XBAm_q|>6AHJ@MYX2qVGAdM&rk>(kH5VcVDYaRbyE6E`YeB zeX%^`!lP^#zTQZ{pm!_g)k|RigJRY0`Tcn0Q$@M-(vbvBB1yHTD>qipWN~-dUzh0l zvNl}S#&OhG7sMlzD6H3e6E(aS?bGHd$7)z$t@Bfs%lbPlC5?&OCiC{1nmX4N*1xpG zcJok9ocdnQXG@xuELKj930CS{?Lg}57F>Q9$hv=Q*}gHEYJC;Mk1Nfno|?kFI*H1v zhLiclawe~;8dAU20D?~^FldP>>GgaG`4WfUcyrbM4t)vtwq&F66q-LbCA#_yTt<1K z?;3*f7FW(r)bhl!GnMxqHtKM50#!d}5O-oEcV?S2)I5nzkA2ztR@RJv%<=kd7IhxQ zu&t;GWm4bNcg;}JMp#oi(gxjrPadXQa>q`~<7+b+bV{eVFP_8o<7PB&k&W%QX}H`S zkM+b+xLZylY^g7EDuv>cGlLCDHguAD`mEH;cON@*M)v9}t-Q&pHInQlPK0bL*nji2 z4cc99tc_E%-aZiHhZeMJ;mNLcZIz{E3Vl0i=vE_|v(HAbCd-4Y_7m_Z-&TILbQZr? zH6o&J66(`AG;!Cn&)bfua~32IaO2r53s$d5W?ugjsv6<5FfAU#j)}F)kFIrL>+ily zRckSw zjRwwR$U%kGecdY}2l)}2@>1pTJPTEhJJ*9}ku_G$umL^r37twm*YWt<2jhCso^G-~ zRo3)n|Cz}=Hk-ixq7Z&x+Os00rP0B%elE)EMNw>gqhna2Es=EtxLzZWzrHz_V zTZ3i$B0BXdN#W( zk+_yg+?niycUBRrr5=9ZZcj{KA5Ppty$&}BV#DeKLe403$#65wub@rpJfg4vMlGs}-g*mgFIWj>@Z!;@?@tnc< zv_zKp`my4O4*UDD^gZX#gXv)$Om?MZZ#Ry7o{jp-Z0Z(jc=k^(o<_~4AR~dO&i)u& z51}B>pL1?fXGhwxvz0xb$DB#Huf-z8o(WPvtXLY%i$>Ob>oJ0oC-FS%7{s*8zFf#i z;+*s}!#DY}^Q?|aOXb|M_9K0BF2Q+Y2yW;~V$pQIRf}PZsUM%)IkI(^J&9&hE7mF% zsykobmq+T4mT!!k$F&*5RI{37s{ZjlTfVaC$coV)N{wpUtyg8l6&rQvG|#B_l&Pv) zqg_?T&64TXBM`N7xoSyKZuyG*y2{Hm8@v}^GCG%Pg5_sp3J*3fkIgSp)qUne>&`kp zo{XZ=>|i!6@}>LPNSc2!!Ny6>!*4D;8>42{A{!c}oAPd?GfAGdtgAnn0iy%CJw2LF zN$wch=dds$g)AY>1 z#-AahIULlF#8n!q3{NC)VjvaSI{t`^#Ir^eJ_ADdnv}x1i!->iT}|IPu56aNFHEwN zYfjRqN?jkmG?Kw*=h3B;Bb$1-(er62%eR_QY3)q2oQP9J5NR);FJq{`$Az+|{~D&LOc z{DxF=Ynw2A;VgCr`BUp{8q=2epj!}(cj*ioFN?;oo+l|POBUy7n9|#i+V3TgS*zvJ zcSkM^aU^?2BKN#f*!Pb$J-wo6eBO^W-tJ5t(8peJKm{Q)SiivtWjYLC(&qP5RGFZaT%ybZ#0-N+Y~z2*{r&hFrMO%)s+qR zf+$N$WkR}-vU8&g?T^HfuuC$gPFi;Cp3T$7Nz^~3p`Y~f8|$k0OLE^s7kx;0tW}&x zyDJlK=i%AYiqbx51TL}Q*VSH_WEh4(`}27)+msQq-k0yy z)3kL0e_oDY$~9-Y&T{4Me95O$d$DzV7&ogfRF04Hqp*==TTLQ}+v`P0wG6(>d9${5 zq0+FoHyLXk(VZ@#uO2tsUdS=6js%W z@og1>^Ji0@)(S=YB=g{jV`U^UMH2V zV>20GqGi6`hCO#C@jgA4QX791tGV3H(pJnlq9(Y*SJl$QB0iZ#GCtXvnRe3u|87YC z)3%%%WJ&dTYR<0o!$5wusHw1gza>^(&9KSPaK)x?#TUPW2M6^IQO32D{B>0ZCw91S zY=1OXpQ5Ro>4)`ICmuFPWoOEKY@;)iFw&TD=4=3!JESo7f{sPf!}YFHO4BV7WDNJ_wY#i$o%EEvb|veB19|@>(YBqM z7S|W?$KnLJb|gP}5W%>9i5&PXtS2Qg9Qqp0rL%gB2W7JH?^xOzXOLc=&W??dJo>Yk&`Z;(HrJo+k`0*nOa7({ zrH1U88OyyHaAtDF&ffb~xtp3*B*rW(-`c|ir$^IO2b!!iX$zfjV9mP zqpD;aV${gh*~qF@qH0wuTUCP>(rX-#W%%$ts*Yc>%4aQUpbRi{=f$hPj2b#i{}g3K z&ZPbahYwn=dXpM~*Z6c6y;0+LHja~$&t)`_ntqe?>rHe7$^LR|e=_Hrd9cjXhH>@% z=~dO01HEQ5q+S%q8q2;`7>IY>LKGLlaGpMRUy5Ie>nTYBl&VG0w2jKZy%dX zFI5aTYb6n=w`6od9{tNRu~JPVa_|^-{gcA4={gQRjb^Y-2HVyo(s6krzLI&pjLpLI zH(QpT)#5Toa-=2G30fP?>p)rSTWZise)zPBCj~{`#99P$FDafo%WQEPHIL%Uk?248 z^0cjnhm%66XXh!ol>>L?N3d>lA!QrKa89m?Pme=cHldVt9&<5N&EUS3AAb(;V|%_c zR;#54ewfLlU22SN)i@2(lUg^ErRo5-KJsPUX?gy$iOg@JA@^;dtSb?md}z++&MEY| z7{>ATI*zZ5qqWgI7G_y;?RXq}$Awe*RXVdwEjcG^&f8I8yhzQW^|nNMPEBCrq(mG% z2GKghoBYpF*ycqt`n(3s_6)YXoy*`ccC5Y`$tuY#{LVyD&?1(O8f%KXXt>%jly1_y z*{w?9ljL!0n&oh8aT?L*WuF-tvAHCzXB^GdOx+0H>bGv9$}Qaa(gRjpg1W|!pkm$SC1zRF&I_P1a@&CF9)XB0~=CGTU0NVNN0 zdH<{IQ%SZ|nXcyj11&8(2**zQtnx6o<3UXw7kXJxmKlZf;UtFrnoh*Aacn5dAmfFa7&Meq2$UBWG^-`Wp{0z zQtxdf7KNEg=MNEzL1ZyiN}N$jvT2y(MAe0%JU=s+zkVNCk@X;ur+0QJyO(7uC#H?# z##=+0ER~%8*APC|&!dL5M8O>`w?Q~ikZ}@WfeiWaor!mzq zl>>7%d>pxu!#xwI|Ja(hcjCDp@!{YXlSG>4Cb3l9=4ng{Cq{*^?2eYN=ejDE>EdEn zGNd-uW47Lw+SQZjyVR8jfltfVkCj@r%z?dPd6WAt!=gn#zmAXOmGphKqvG*xAC20! zkfdKuspj?bVVd-oH^%!L?Ol+-2+0rg2Imu&8%=JWD`6WcNyUj#8$ zJiSwho>n8`nQc6Y^c`xxzf5PZLk3&rweasO_3(KojEghaJJW`j5xE$J>2awgeWd){ z4ZmkGv0edD{ZqMY>W-1*o43O>G;xSx+Ba85f6lD9`b~#b+sewGHw#E;p2#9QUs@DP zPho7ryl!X8)TvHnEwQ9k=Tz$L9$2xo*}a3B-=dYc(j=6{ zS#+A^#@P4qJl+>i?QX#|*x^py-_yADVF3&3WZ`nz8pCSw9NwbiRvjNMm&8-HTuqNo znQSOdq}R+y5?Tb~w8xxjdqbI+nnvwc(wj$zlYMX@X-TpCxiSE2k68Y>nM(L~PfRmh z>1dP8q!~#}dr`)&ONoSRl)k@a9N%Q^^Q@LerxO|U>OY4ym5Zs=UwW^$@iZulqq5Xv zoh;L^w9n+gn_$j2&S2tWFT4`%so5umt_k6En3+v)jRw13YD{Kpupbe^fQP|U?;|~> z^!N9lr}6&9NOJGGvMe;2R^@80K9pM7CXr0bD7tS5;GT~keU)I&iY-po1n}^hD-$*) zuyBPZGo^3dbwSU?-xDykia{Nk%7)chbw}jcS`x7;!!xpN(0xSg#?)Si_g8aSXnj zgZW80Uu8aI(8B<#9Gb=7yW;rckeY#Vz8w9U#Of3uxi3I*Xeg1H{BB1TE&S;ZESHosV07VG)uZ<;j_}3agsySHHhNEq#UyDYiO5~NWV)O zeyKVfqh9{lkCtm|VmuRe$#pJ$-qzOkOe=JvmSp&Ljq=dghy`0DV4Psbo>j?Ae;i5M zFJ45MrZIhY3^_Y<>3uJQHT?tma!_jd5mJ{}=FlhAmpojM)L>b7a<^OK1jt7S#?NM&{7NUj`|OuEm0)uttV%cEm8#N14f+RK-ZGB;yn zWk%j09UE(&`u}KD)lF8k;-Qqkan2mfEuOY7Fz5CiBrhldP&!*{zq} z@Ya4+`_MU*x2>gAZz}n7?;P5%N>VCbh0x|l9yhdvgOkLEt(;jgLS_~grONlEGODH2 zh-1}g=ZF&(&of~SZ zeT(Np|6;CRIjLH{%%8-Yo=m-3V6YOJyv%l+kp_ zDprP*+9Qu6FH?ABkx1}w)7W<+iJAAZXjUs5Z&#VEDR!dhfE(@VW>dPrfjY4dcCC0E>2;(L6bj-32N9W~?W}Ae2wjN=g4CQnD={c217uu=L68 zque;T(~;jJ<#qE=6QC@n=^!;Xy@P4hF^U-{5+&14A@5BZj$LQ*C906L^L8w48pH0Q zD9Q|iWace}8e=oq*((4;>ogkFb?1uA_iQ|>;~(j-Yk6c6?wZ7#*|GF9ls)}w07)4E zOb&G>7$!^7fvz9o!|m0~BgylK%QfGxGsuqg7O z%05RLosUKL!j-kLQ<&6LM~u`oD?dh%@JZ&yY%(R=kH?@a1LbZyTl1rspe^BUEfWm> z^ryPac9^bD;*U+ClKcA-Y33vS^3;mywfCu3MYXQ*4_Z)u|F8#3&d*RCyRk~;98jm? z)Pu?uEjouAHF~^PwfdNgQDC&0(UBaTYLv-b)tNu^yibi|(c7J>^#M8Mu735E;)U+S z**!73cgsfh4Vgm~M=qTWAua#>rg3A|0P!`V%ev z|FFBV4jr+T9@wAm^<1guK8G3EQTQcmxSSZoE{8(4NDk&z;X}33B0k?sX275b`gM$C zx)RCA40#VLPo;2cGz+YgICa*NteAXy?v~dk+?c^b$I@|jD#!lO(Mk52lTzc2YbeKY zFp>A}G0c$~ICiQXfgiL~8tuoX*)#awMrPz1$I~TRLmOY2ACl|u?rU%2B@gx+kU-#u zc`TEx*`Rk6Ndx^za?oJm7m9k1WXO64mZwYp)TW5ZcgFBkY`fEdFs@84#l_o-HQQ&9 z9p}$ojX!U5To@A_L2i>Q7I#S`HAUtC3iQm3ke_#J0PFqzIA1e~Jo_Xv25P7!`R{Ns zf}G&lOq4vTOIbLNtLW&mAf8FB?N~HpF8OKkOtg!@#wP=dm6n9bdQkOb7@sfV(u`^$H5KDeu{>Uw#ch`aB|#U+jjb9kEh-{(bt)f@%8ZGu zVKXMA@YGf>^PSVVJW{XxGA5b_`YfgG{wO7WT``M;UAS{3o5y#Z>98RT*8^6J8alSZ zA~c9mw|6S0Mp??}YvZVysv>nrDxDgH%3N+9u8TGpwe|Mp+@m}O)+|(pcXVgxB(cs% zwoFZ!!n~w38vWwKmIGRBrH)M4WUcBHWX_*8LO9k}_O^CJArV#QeY=&nb6=>23dz2hWTyynQVmU3EQ*^eKo(w8L2ENJe6RPI#Pa&Eyw=8Q<-(mOLa`bwcL%MD}R@tiSc5?QAd_^%VEzi%NRL6 zpGS$2OnerKLwP($C138mwU7^mzo`iH;qxUIZr+bE>fcSx#N|33X67?4IhraH-PrGJ zDs$!6REHD8`Ba|B-Kb!UdgfppCo{m46F4sx9@{#R-}YwE{#l01E#$~Nf*m!yow4kb zNo7M@7Iw^|@SL7V<9NC|${gO-Y@VDeU{YEtN%K7DKWjGTH>KZ@y?m?u3|ANCR0MC8 zdFQ_?Db7ZvG?BHieHA}??nviYh$)Swua2nh%5LLicIqS0ywtJ&jVY*{2;6JQd&^WF z#`Le)Ua4idcT%FF{WTf?dU*_!x!g+Ml8AkfL={gnoVwY=OG&y^C%m-|~0Azm@eEQ@}}(s^+;mUg|%FqtIl z&O*ros>lqYx$Mi;0*Uz5hjB|%_$4?Qhoj4AR5_V0{bE^sP3EHHH9YB*L42DWnN72y z^v-+?Mtbn6i(~~;lTdw%Azs$#qxso%klyaIVK(po2;%z@H)CQO5(mjpA2*g<=XnV2r)A+ZK*r|p_>xmc`siRSU1uh8YjrZ$jwO&>5W&65 zp|m_6$IS!G7sYzAqV9u`F?q%I!k`ipQVL%Llv; z;>smA)p)}S)#{M$6&rW1FLx~)V`TgA7uEe@#b{^nT%$o7m#9t@PgiAo$~<_KmeMub zR2rw;@;B=nD+k91bEbbK)u;*&nQ3<8*YDE|*Kb~}8h0^~DU)*ud920BQj2Dj^fIY= zZqLVEol4JTY7S0SV>Mh&h0MfTPnpd)`$YE0e8%4&oQSWxfJ?1%nXa*7!@H@(j>zJ3 z6`2PY3)%KD7sqQl$>?=F9+^nS+ycBtd2w6j3}5GjvHq6K-AO&&`KX4i&BgovNS4{Q zU}jth<$M1aT1Y=0yv>6ZFFf!Q&vGjYW!w=x&pWH}mwaI62yYh4OuL_1Da(gWqK#Y= zO_SoN;sN-vWrxsT&ZilpH1v6x#)S*n+&tsP^i`fb2~Q-d zc_E`d$_${KgTEek!A){z`{rsM^$8(Kdek_X)%7@3z~W&!n90oa&0Qg6$EH)|eLj!> z3}OCj9SObFOwh-Wzd+5>qD%t*jzJR^&n}}FuC!vC(3*I$yFv4OZXm={r(Q`p7Vzq)Xpd z=8NuR5OwdCabdaTIG<1<1ueam`6GvIGIXjiZdy_fyOQCXpjsLW5(Q+j_4Qsz& zRpny8lRm^U?!f}()QN>$l>F9HGTBqz{IUHM&g!WyxXN6{m0hXa+gwQ*@yuWg%RYGRfYREjQ0e(V=7!v-W3fLA^V6{m*i*#Wlh=&K{}#^00Y#Wc zlqx&3f*3GN=E;}2<0Er`uU|>lcPg9>GH3G5IuB*4lWM+?1AFBC>e62F^2VWr$h{4% zC6CRK9QUt28X_L&b2xq`ts>M$D`Y0nXILp4S=Plv zkBdmzGLeg2b$G>XR+bbd@#vY@nABx=*CbG7bpg#}9L<@%-NY{`{e z&PnD?db?5{BPP2%iyKiX6ei0It<2llSIJ>)TozmVrm=hP66zUh7`VZM(h_l+k^Rau z{*?QWeobN0*?fxfvdESm=#XSpZy$D7E-IdMo|Z$4rs<4Z>cx)HdhQMi;%U=wAQa$u!g^e7Sm~;8tcFW9u>$nCo?>eAq(if>vz@o1);3_7|5bo2aFy_-ms=i z22Wi|=wcvsX!8Jebhl#A%6qC~V-nHW>sTlbxl!&@>Kc?m_ZeD}hRkB|1RY~f<#DWS zKCf&F>2k|MzDH*I2j%gwh9`qp7Bl){2Cdz-%>9%|)dq#Exw(LGnOS_2dkr2obs$Bb z!bPcl$F`Q7`%p>6+1nXhh^(qKYf;9>@jANC4d=6D-CyFYY49$T>6QVEcS_}%oJ)87 zGZ+$QMbT&<=0COL!m%s{JufPMU-f`dYW@gibie{)<+#2(<*;sbJcDFi{x&O}efJaj zm{>~wvRs+JEaKW_EuR)Gq21Yd1{kko+4W@BNCwnE`p1r^Gx>GyLfMb?j6SeZ?gi1Y zX0#%E z;H?y~u5}pw29;vwqKCu2#5jcEG*#+@n}v*y6F=FI#oDFbbd=|D$dG(RX4{7Skw(e7 zbZ+cf!~#EWnw{_>qRwJwotE`odZpzp)46h6OVsFm)}K$u(!~SYlL?%WjLv9Febu@h zQ0UM*(zos+f?hrR~&R}7pkhAm}YcG+gbV0wTbfEfCH*F=SQj* zuNIrX^ds6u7C{}idggs1NMf#JQc7$#`+kx9dqxKHX2jAVFNZp3WS{9!$n*UvM9E%N+b4@&SMsT)kvX62 zGM3-Sl=*~scK;>!*VJ6e%Q3mkTb?R4U%q5$YU(}+Wqr7Y>~m|WaV3|p-xBE?DEBue z>6yK|lwOlEsXL&Ig)c(+WojDdl2-B9P7Bv^IraBEJOjNLJ~xwQN0X?vIFm}<^BFoV z6-8_U$v>5SRM%x!Wk0RWdvE;i%;8*@f5+N#Wc0MR>`+ zyf1by_UQRGcLBA-itwEgMEqf?U(XaV{dOqZ8zrz}s;tpxTxlDU!%nkYl6Pm3(LR<2 zqq6Y-MQRjzU-#A$hdrWWLKO|w?Mv8NBZ+Bt`*SruW`!S zeWe^(k*N&aQ$WznT)G`HQVxAeBYCcl*-NJ>c5#w{PjHuWGz;U#^Qmkkbz|KEuDZ*5 zUS|PIMz&G%UpXsF*J=nV(~;%j1?5?I+)txMw_J>6p7Qe)FZ}neFl<~GDSj_h1pOn3< zu#xxZ`U1ZHq2aT*_t3^tJIa1&zB7v+&hg4qIlch}3GAqoOa1A|1eq0cQ1;B$!^bI} z$Ia=h&7$jkJ#%{$aBWl$_S1AMIjzC>$9(qC3{Gb0F!h&N%i#s|mHUN^eY54h^c0F^ zZmovQhn(B5A@&dV@=wG4Xk?tip%W>-{=dff+yLum3z@9olK;yzKrue#l-r_y*|dS$}yduv9}h8_3P0voD1;SDvZRjAC#q2vUB>3JeTVesCUQ?o*JNtQ^%>eHo%JiV)afI8 zM@|_B?#ci2k-0P%sllawt2;*S-Mezks8(M`#zxmsj#{o@-<}-WoYQh-Uk=Ml=CgKJ z2^ErM%sHDZ_Z%0KQxbrqYcA(*FX3HtJ+04Y^W*OySZj0lje%nDpTB_M|I6P$F!=w+ z-#>6%{Gq4#Lz?)*B=LuD;t$)zA9jd8_=`U*6o0t#(;sSzKlB!VcrX6YQ~Y6^_=CIn z!-SvyfcQg4@rP~V5C4ciEZ=QLOYw)j;t!|AAF_V>Lx%XnG4Y2z;t#)wKXej*I4%Bg z`lmlch(9zJf4D3D&|3UqruajGWOq5@5BA~@AH*N@;t#9EA9jd8EEIp(_0u1Ai$5sh z4~@ki{t$oIApY=G{GqnTzy9!1{Gqq_L%#TfM*Lx__(QDt!!+@SIY0fuNBp6k_(Og1 zhXc}M{r%G)s{Hhat>O=s;t${UI?_b^LHoj$)#48m#UIq-4{O99;=~`Wi9hTRf2bz@ zFj4$roA|>-@rMUL{lP%|VY>LkS@DM?@rMZUhd}X%Y2ptX#UBF1AMS}iEEaz_`O_cH zia#_Je`qTHutof#Nc`a!@dqREhtuK@pQQF)DE_cY{J};1p;Y`q`_mt$ia!h#f0!-) z;4c0!Q2aqB{%}Y9VY>K3FR35LjI$y|{Nc4^wQ=GPBgG#!h(GKRf9Ng#P*wcFLiS(Z zpZ>5%{9%LmgZEE=_#%DbVDX1x9V3}0+3qOuho<5WCE^bY#UGCU^oJhe51W4aLuc`a z6XFjY#2=1}Ka3ZD7$g4hN&LY|{2^NW!T*DnYT^%D#UH+iKWN1tdWb)Ckj$;4_`?J7 z2etTvulPf-_`?eEhsxp)I`M}@@dqpM2aWhcfANPUKmB2f_`?(Nhb}+;!BhNUx%fkt z_`?H@KOe;(M*j4NMdA->;t&0P`h%bN!^WTfP$d44BmPiZ{9&~CLxK3i6d$>tUi{%{ z73m4YAEt>vl!!kRi$6^M=?|vj524}@HN_tye)_|=N^0!HAGV4=XvH6*#UJ{JKU5Ka zSS0=s`O_axN-p@tOYWf%fA}E&(C(){v=M(W7Juj<{!m@~p}qLSUGWFRA9jd83=w~5 zD*o_M{Na-LLp$*YFY$*N;t%)59~y~2c!)oA5P!HT{%}|{Gq1!!$a|hYT^$a#2+?^KP(h~SR?*mBL47N{Na}P!`Dau`h!~h z;j{Qdql;DE@F- z{2^BSVSxBUg!sc+@dvf|Loe}%OX3fs#UEyfKP(Y{2orzMia%@@f2b|~Fjf4ahWLYt z_`@^V&u;wmhwDH6Ax!+?llX(7tYI_7AD;j8hb!U_$HgD+i$Am$e~1%*C=-9UFaGf4 zt}`9Q9~{LWhKWCzi$9DIe=zy!5A}cgLr?Jso%n;F_(MnWhx6hOL&P6OiWh7Xe|RYV zVEWS^nu|Yd7Jsl7e~1%*==#$iPWJU=n)rjM_=C6j!zJ+tJMo7Y@dqpMhgaecE9LRL zfBM55@dqFAhXnBlH}QubvxoZO0kjlqX5Puji{@^bD&{_N;Nc>@s_(PcZ zgNbB!y~Q8;i$C-fe>nWpANme(mHVMn7%Kj-7eC3V{`H6apZ;)G{9&&6!&~u(-r^7L zKmB36_(P`n!%6Xn5b+0-pZ+jj{Nc}^{xD7a;nM$A%K3ys5XEtPSQdLy$e$$A+Lb?9 zZ_UgqrKLSkvdX&o*9yr6sdaPXM7cOf4sIeYd(zsXO*#12mILJ^$yMRx`*xr0y*D%E z&unksY&-Lr-|o!Yea)x7@DEeCcfSb#Z~^~tm>7-k!aqdeA3FRS@mKhV9{gvI!#|wE zKh(iLe84{>n18sF{6iP~LpJ=w3H(C?{KGN)!)I(c{=)o2AMTy+G5@d(|8NQaPy_!k z2mf#j|BwU!Fa-Zl0RON9|FC(T_786OhlQiH_%r-N0RCYG{(&(6kk9D8H;2$pGABx}~BJdAU_=h9r9}eIjV(<^M@DFdH)ztiBe(Nl$1(d-jB|ooWr{%onX*rU4NOHKYWI7!1JAYP?$4h*wMtoX}+D25;s%u2|s9vJ0c-5kN zH4pLne40+HjGR{NyHO!+{k>VV*i7;rpP5!p-w+nd96TArC!tjfh)0}Y14fEllWBac zS*<_|qIe_CkfkOA`HZp(0eE8u%w(g<96m9uRN#bcCWWQ*&148a?H3`1i0dQHa_w~8fd}0_z3mqi(Ms3= literal 0 HcmV?d00001 diff --git a/tests/test_bundle.py b/tests/test_bundle.py new file mode 100644 index 0000000..acb4c6e --- /dev/null +++ b/tests/test_bundle.py @@ -0,0 +1,227 @@ +"""``chronos bundle``: build → check → translate → prepare-data → fine-tune spec, on a tiny model.""" +from __future__ import annotations + +import json +import os +import sys + +import pytest + +torch = pytest.importorskip("torch") +pq = pytest.importorskip("pyarrow.parquet") +pa = pytest.importorskip("pyarrow") + +_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.bundle import builder # noqa: E402 +from chronos.bundle.builder import FLORES_REL, build_bundle, check_bundle, load_bundle_manifest # noqa: E402 +from chronos.model.checkpoint import ( # noqa: E402 + checkpoint_config_path, resolve_checkpoint_tokenizer, save_state_dict_with_config, tokenizer_metadata, +) +from chronos.model.config import ChronosConfig # noqa: E402 +from chronos.model.model_chronos import ChronosForCausalLM # noqa: E402 +from chronos.tokenizer import build_tokenizer # noqa: E402 + +ZH = ["今天天气很好，我们去公园散步吧。", "经济增长放缓，央行维持利率不变。", "他说：我们明天再讨论这个问题。", + "人工智能正在改变世界。", "请把这份文件翻译成英文。"] +EN = ["The weather is nice today, let's take a walk in the park.", + "Economic growth slowed and the central bank kept rates unchanged.", + "He said: we will discuss this question again tomorrow.", + "Artificial intelligence is changing the world.", "Please translate this document into English."] + + +def _config(vocab: int) -> ChronosConfig: + return ChronosConfig(hidden_size=64, num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2, + vocab_size=vocab, num_experts=4, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=64, max_position_embeddings=256, recurrent_depth=4, + min_recurrent_depth=2, prelude_layers=1, coda_layers=1, dropout=0.0, + fallback_mask_prob=0.0, recurrent_attention_mode="shared_first", sliding_window_size=16) + + +@pytest.fixture(scope="module") +def run_dir(tmp_path_factory): + """A finished-looking ``chronos longrun`` run directory (never trained; random weights).""" + root = tmp_path_factory.mktemp("bundle_src") + corpus = root / "pairs.jsonl" + with open(corpus, "w", encoding="utf-8") as fh: + for zh, en in zip(ZH * 40, EN * 40): + fh.write(json.dumps({"translation": {"zh": zh, "en": en}}, ensure_ascii=False) + "\n") + tok_dir = root / "tokenizer" + tok = build_tokenizer([str(corpus)], str(tok_dir), vocab_size=700, log=lambda *_: None) + vocab = ((tok.vocab_size + 63) // 64) * 64 + torch.manual_seed(5) + cfg = _config(vocab) + run = root / "runs" / "tiny_run" + run.mkdir(parents=True) + ckpt = run / "chronos_enzh_pretrain.best.pth" + save_state_dict_with_config(ChronosForCausalLM(cfg).eval(), str(ckpt), cfg, stage="pretrain") + side = checkpoint_config_path(str(ckpt)) + meta = json.load(open(side, encoding="utf-8")) + meta["tokenizer"] = tokenizer_metadata(str(tok_dir)) # what the longrun exporter records + json.dump(meta, open(side, "w", encoding="utf-8"), indent=2) + raw = root / "raw" + flores = raw / FLORES_REL + flores.mkdir(parents=True) + for split in ("dev", "devtest"): + table = pa.table({"sentence_eng_Latn": EN, "sentence_zho_Hans": ZH}) + pq.write_table(table, str(flores / f"{split}-00000-of-00001.parquet")) + spec = {"run_dir": str(run), "shards_dir": str(root / "shards"), "tokenizer_dir": str(tok_dir), + "model_config": {k: getattr(cfg, k) for k in ("hidden_size", "num_hidden_layers", "vocab_size")}, + "phases": [{"name": "pretrain", "budget_fraction": 0.7, "peak_lr": 1e-3, "tokens_per_step": 65536, + "micro_tokens": 8192, "seq_len": 512, "warmup_fraction": 0.02, "min_lr_ratio": 0.1, + "weight_decay": 0.1, "grad_clip": 1.0}, + {"name": "sft", "budget_fraction": 0.3, "peak_lr": 5e-4, "tokens_per_step": 65536, + "micro_tokens": 8192, "seq_len": 512, "warmup_fraction": 0.02, "min_lr_ratio": 0.1, + "weight_decay": 0.1, "grad_clip": 1.0}], + "budget_hours": 1.0, "seed": 1234, "dtype": "bfloat16", "export_prefix": "chronos_enzh", + "extra": {"raw_root": str(raw), "params": {"total": 123456, "active": 65432}}} + json.dump(spec, open(run / "run_spec.json", "w", encoding="utf-8"), indent=2) + with open(run / "metrics.jsonl", "w", encoding="utf-8") as fh: + for rec in ({"event": "step", "phase": "pretrain", "step": 10, "total": 40, "tokens_seen": 1000}, + {"event": "validation", "phase": "pretrain", "step": 10, "valid_ce": 3.0, "valid_ppl": 20.1}, + {"event": "step", "phase": "pretrain", "step": 20, "total": 40, "tokens_seen": 2000}, + {"event": "validation", "phase": "pretrain", "step": 20, "valid_ce": 2.5, "valid_ppl": 12.2}, + {"event": "validation", "phase": "pretrain", "step": 30, "valid_ce": 2.7, "valid_ppl": 14.9}): + fh.write(json.dumps(rec) + "\n") + fh.write('{"event": "step", "phase": "pretr') # a line being written right now + return run + + +@pytest.fixture(scope="module") +def bundle(run_dir, tmp_path_factory): + out = tmp_path_factory.mktemp("bundles") / "tiny-bundle" + return build_bundle(str(run_dir), str(out), name="tiny-bundle", log=lambda *_: None) + + +def test_build_records_best_validation_and_leaks_no_paths(bundle, run_dir): + m = load_bundle_manifest(bundle) + assert m["training"]["step"] == 20 and m["training"]["valid_ce"] == 2.5 # min valid CE, not the last + assert m["training"]["tokens_seen"] == 2000 and m["prompt_format"] == "pretrain_tags" + assert m["eval"]["flores_plus"] == ["dev-00000-of-00001.parquet", "devtest-00000-of-00001.parquet"] + assert "model" in os.listdir(os.path.join(bundle, "third_party", "minimind")) + home = os.path.expanduser("~") + for rel in ("recipes/source_run_spec.json", m["model"]["config"], "bundle.json", "recommended.json"): + text = open(os.path.join(bundle, rel), encoding="utf-8").read() + assert str(run_dir) not in text and home not in text, rel + side = json.load(open(os.path.join(bundle, m["model"]["config"]), encoding="utf-8")) + assert side["tokenizer"]["path"] == "tokenizer" + assert resolve_checkpoint_tokenizer(os.path.join(bundle, m["model"]["checkpoint"])) == \ + os.path.join(bundle, "model", "tokenizer") + with pytest.raises(FileExistsError): + build_bundle(str(run_dir), bundle, name="tiny-bundle", log=lambda *_: None) + + +def test_check_detects_tampering_and_smoke_translates(bundle, tmp_path): + import shutil + report = check_bundle(bundle, log=lambda *_: None) + assert report["ok"] and report["missing"] == [] and report["smoke"]["device"] == "cpu" + copy = tmp_path / "copy" + shutil.copytree(bundle, copy) + with open(copy / "examples" / "translate_en.txt", "a", encoding="utf-8") as fh: + fh.write("tampered\n") + os.remove(copy / "recommended.json") + bad = check_bundle(str(copy), smoke=False, log=lambda *_: None) + assert not bad["ok"] and bad["hash_mismatch"] == ["examples/translate_en.txt"] + assert bad["missing"] == ["recommended.json"] + + +def test_translate_backends_and_speculative_agree(bundle): + from chronos.bundle.runtime import BundleModel + bm = BundleModel(bundle, device="cpu") + src = EN[:3] + plain = bm.translate(src, "en_zh", max_new_tokens=20) + spec = bm.translate(src, "en_zh", max_new_tokens=20, speculative=True, draft_depth=1, gamma=3) + assert spec == plain and bm.last_stats["speculative"]["verify_calls"] > 0 + with pytest.raises(ValueError, match="repetition_penalty"): + bm.translate(src, "en_zh", speculative=True, repetition_penalty=1.2) + with pytest.raises(ValueError, match="supports"): + bm.prompt_ids("x", "fr_en") + mx = pytest.importorskip("mlx.core") + prev = mx.default_device() + try: # device="cpu" sets the process-wide MLX default + ml = BundleModel(bundle, device="cpu", backend="mlx") + assert ml.translate(src, "en_zh", max_new_tokens=20) == plain + assert ml.translate(src, "en_zh", max_new_tokens=20, speculative=True, draft_depth=1) == plain + assert ml.translate(src, "zh_en", max_new_tokens=12, repetition_penalty=1.2) == \ + bm.translate(src, "zh_en", max_new_tokens=12, repetition_penalty=1.2) + 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" + rows = [{"zh": f"{zh}第{i}句。", "en": f"{en} Sentence {i}."} for i in range(30) for zh, en in zip(ZH, EN)] + rows += [{"zh": ZH[0], "en": EN[0]}] # FLORES+ pair: must be removed + with open(data, "w", encoding="utf-8") as fh: + for r in rows: + fh.write(json.dumps(r, ensure_ascii=False) + "\n") + out = tmp_path / "prepared" + report = prepare_data(bundle, [str(data)], str(out), workers=1, log=lambda *_: None) + assert report["corpus"]["eval_leak_removed"] >= 1 + assert report["split_rows"]["valid"] >= 4 + assert report["shards"]["sft_train_records"] > 0 and report["shards"]["sft_valid_records"] > 0 + with pytest.raises(FileExistsError): + prepare_data(bundle, [str(data)], str(out), workers=1, log=lambda *_: None) + run = tmp_path / "ft_run" + spec_path = write_finetune_spec(bundle, str(out / "shards"), str(run), tier="smoke", mode="sft") + spec = json.load(open(spec_path, encoding="utf-8")) + m = load_bundle_manifest(bundle) + 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" + with pytest.raises(ValueError, match="pretrain"): + write_finetune_spec(bundle, str(out / "shards"), str(tmp_path / "ft2"), tier="smoke", mode="pretrain") + (run / "checkpoints").mkdir() + (run / "checkpoints" / "latest.state.pkl").write_bytes(b"x") + with pytest.raises(FileExistsError): + write_finetune_spec(bundle, str(out / "shards"), str(run), tier="smoke", mode="sft") + + +def test_pack_writes_archive_and_digest(bundle, tmp_path): + import tarfile + archive = builder.pack_bundle(bundle, str(tmp_path / "b.tar.gz"), log=lambda *_: None) + digest = open(archive + ".sha256", encoding="utf-8").read().split()[0] + assert digest == builder.sha256_file(archive) + with tarfile.open(archive) as tar: + names = tar.getnames() + assert "tiny-bundle/bundle.json" in names and not any("__pycache__" in n for n in names) + with pytest.raises(FileExistsError): + builder.pack_bundle(bundle, archive, log=lambda *_: None) + + +def test_stage_tokenizer_guards(bundle, tmp_path): + from types import SimpleNamespace + from chronos.model.checkpoint import chronos_config_from_checkpoint + from chronos.trainer.stage_utils import load_stage_tokenizer + m = load_bundle_manifest(bundle) + ckpt = os.path.join(bundle, m["model"]["checkpoint"]) + cfg, _ = chronos_config_from_checkpoint(ckpt) + tok = load_stage_tokenizer(SimpleNamespace(tokenizer_path=None), cfg, ckpt) # the bundle's own + assert len(tok) <= cfg.vocab_size and tok.convert_tokens_to_ids("<|zh|>") is not None + small = SimpleNamespace(vocab_size=64) + with pytest.raises(ValueError, match="vocab_size"): + load_stage_tokenizer(SimpleNamespace(tokenizer_path=None), small, ckpt) + big = SimpleNamespace(vocab_size=cfg.vocab_size * 4) + with pytest.raises(ValueError, match="--tokenizer_path"): + load_stage_tokenizer(SimpleNamespace(tokenizer_path=None), big, ckpt) + with pytest.raises(FileNotFoundError): + resolve_checkpoint_tokenizer(ckpt, str(tmp_path)) # no tokenizer.json there + + +def test_init_checkpoint_guard(bundle, tmp_path): + pytest.importorskip("mlx.core") # chronos.mlx.training.longrun is the MLX trainer + from chronos.mlx.training.longrun import check_init_checkpoint + from chronos.model.checkpoint import chronos_config_from_checkpoint + m = load_bundle_manifest(bundle) + ckpt = os.path.join(bundle, m["model"]["checkpoint"]) + cfg, _ = chronos_config_from_checkpoint(ckpt) + check_init_checkpoint(ckpt, cfg) + with pytest.raises(FileNotFoundError): + check_init_checkpoint(str(tmp_path / "missing.pth"), cfg) + wrong = _config(cfg.vocab_size) + wrong.hidden_size = 96 + with pytest.raises(ValueError, match="hidden_size"): + check_init_checkpoint(ckpt, wrong) diff --git a/tests/test_interface_regressions.py b/tests/test_interface_regressions.py new file mode 100644 index 0000000..f80c311 --- /dev/null +++ b/tests/test_interface_regressions.py @@ -0,0 +1,230 @@ +"""Numerical and real frontend contracts that short forward smoke tests miss.""" +import asyncio +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +import pytest +import torch + +from chronos.model.config import ChronosConfig +from chronos.model.model_chronos import ChronosForCausalLM +from chronos.model.recurrent import recurrent_auxiliary_state + + +def tiny(**overrides): + args = dict(hidden_size=16, num_hidden_layers=4, num_attention_heads=2, + num_key_value_heads=2, head_dim=8, kv_latent_dim=8, rope_dim=4, + intermediate_size=32, moe_intermediate_size=32, num_experts=2, + num_experts_per_tok=2, num_shared_experts=1, vocab_size=32, + max_seq_len=16, recurrent_depth=3, min_recurrent_depth=2, + prelude_layers=1, recurrent_groups=2, coda_layers=1, + recurrent_enabled=True, layer_aware_lookahead=True, dropout=0.) + args.update(overrides) + return ChronosConfig(**args) + + +def test_shared_transition_consistency_gradient_penalizes_expansion(): + def objective(w): + states = [(i*w).reshape(1, 1, 1) for i in range(1, 5)] + return recurrent_auxiliary_state(states, [torch.full((1, 1), .3)]*4, + min_depth=2, threshold=.99)[1].distill_loss + w = torch.tensor(1., requires_grad=True) + value = objective(w) + value.backward() + # The old detached teacher gave -6.667, increasing the very loss it claimed + # to minimize. Check against the actual objective, not another implementation. + numeric = (objective(w.detach()+1e-3)-objective(w.detach()-1e-3))/(2e-3) + torch.testing.assert_close(w.grad, numeric, rtol=1e-3, atol=1e-3) + assert objective(w.detach()-.01*w.grad) < value + + +@pytest.mark.parametrize("overrides", [ + {}, {"recurrent_enabled": False}, {"recurrent_depth": 1}, + {"num_hidden_layers": 3}, {"num_hidden_layers": 1}, + {"use_hybrid_attention": False}, {"tie_word_embeddings": False}, + {"use_moe": False}, {"intermediate_size": 0, "moe_intermediate_size": 0}, + {"modalities": ["image"], "image_patch_size": 8}, + {"modalities": ["audio"]}, {"modalities": ["image", "audio"]}, + {"layer_aware_lookahead": False, "num_experts_per_tok": 1}, +]) +def test_estimator_matches_allocated_unique_parameters(overrides): + from ui.estimator import ArchConfig, total_params, active_param_bounds, memory_footprint + cfg = tiny(**overrides) + with torch.device("meta"): + model = ChronosForCausalLM(cfg) + arch = ArchConfig.from_config(cfg) + assert total_params(arch) == sum(p.numel() for p in model.parameters()) + low, high = active_param_bounds(arch) + assert 0 < low <= high <= total_params(arch) + mem = memory_footprint(arch) + assert mem["resident_weight_bytes"] + mem["cold_bytes"] == mem["total_bytes"] + assert mem["mlx_training_state_bytes"] == (mem["total_bytes"] + mem["optimizer_bytes"] + + mem["mlx_master_weight_bytes"] + mem["mlx_accumulated_gradient_bytes"]) + + +def test_repeated_depth_changes_work_without_duplicating_blocks(): + from ui.estimator import ArchConfig, total_params, compute_params + a, b = ArchConfig.from_config(tiny(recurrent_depth=2)), ArchConfig.from_config(tiny(recurrent_depth=5)) + assert total_params(b)-total_params(a) == 3*a.hidden_size # loop embeddings only + assert compute_params(b) > compute_params(a) + + +@pytest.mark.parametrize("depth,adaptive", [(1, False), (3, False), (3, True)]) +def test_mlx_recurrent_graph_and_checkpoint_match_torch(depth, adaptive): + mx = pytest.importorskip("mlx.core") + import numpy as np + from chronos.mlx.model import ChronosMLXModel + from chronos.mlx.training.io import mlx_state_to_torch + cfg = tiny(recurrent_depth=depth, adaptive_halting=adaptive) + torch.manual_seed(17) + pt = ChronosForCausalLM(cfg).eval() + # CPU execution isolates graph parity from Metal matmul rounding and MoE + # near-tie selection. Other tests exercise real Metal updates in BF16. + with mx.stream(mx.cpu): + model = ChronosMLXModel.from_chronos_pytorch(pt, cfg).eval() + cache_pt, cache_mx = None, None + for ids in (torch.tensor([[1, 4, 2, 5]]), torch.tensor([[6]])): + with torch.no_grad(): + out = pt(ids, past_key_values=cache_pt, use_cache=True) + logits, prediction, cache_mx = model(mx.array(ids.numpy()), cache=cache_mx) + mx.eval(logits, prediction) + np.testing.assert_allclose(np.array(logits), out.logits.numpy(), atol=2e-5, rtol=2e-4) + assert prediction.shape == (*ids.shape, cfg.lookahead_steps+1, cfg.recurrent_groups, cfg.num_experts) + cache_pt = out.past_key_values + pt.load_state_dict(mlx_state_to_torch(model, cfg), strict=True) + + +def test_mlx_fp32_masters_preserve_small_bf16_updates_and_match_adamw(): + mx = pytest.importorskip("mlx.core") + from chronos.mlx.training.optim import MixedPrecisionAdamW + from chronos.mlx.training.trainer import _add_grads + from mlx.utils import tree_flatten + opt = MixedPrecisionAdamW(learning_rate=1e-4, weight_decay=0) + parameters = {"weight": mx.ones((2,), dtype=mx.bfloat16)} + p = torch.nn.Parameter(torch.ones(2)) + reference = torch.optim.AdamW([p], lr=1e-4, betas=(.9, .95), weight_decay=0, eps=1e-8) + for _ in range(40): + grads = None + for _ in range(8): + grads = _add_grads(grads, {"weight": mx.full((2,), 1/8, dtype=mx.bfloat16)}) + parameters = opt.apply_gradients(grads, parameters) + mx.eval(parameters, opt.state) + p.grad = torch.ones_like(p) + reference.step() + assert parameters["weight"].dtype == mx.bfloat16 + assert parameters["weight"][0].item() < 1. + torch.testing.assert_close(torch.tensor(opt.state["weight"]["master"].tolist()), p, rtol=0, atol=3e-6) + assert all(v.dtype == mx.float32 for k, v in tree_flatten(opt.state) if k.endswith((".m", ".v", ".master"))) + + +def test_long_prefill_is_causal_and_matches_chunked_and_token_decode(): + mx = pytest.importorskip("mlx.core") + import numpy as np + from chronos.mlx.model import ChronosMLXModel + from ui.estimator import ArchConfig, kv_cache_bytes + cfg = tiny(sliding_window_size=3, adaptive_halting=False) + torch.manual_seed(19) + pt = ChronosForCausalLM(cfg).eval() + ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + with torch.no_grad(), mx.stream(mx.cpu): + model = ChronosMLXModel.from_chronos_pytorch(pt, cfg).eval() + full = pt(ids, use_cache=True).logits + changed = ids.clone(); changed[:, 4:] = 9 + torch.testing.assert_close(pt(changed).logits[:, :4], full[:, :4], atol=1e-6, rtol=1e-5) + for widths in ([1]*8, [2, 3, 3]): + a, b, offset = None, None, 0 + for width in widths: + chunk = ids[:, offset:offset+width] + out = pt(chunk, past_key_values=a, use_cache=True) + logits, _, b = model(mx.array(chunk.numpy()), cache=b) + mx.eval(logits) + torch.testing.assert_close(out.logits, full[:, offset:offset+width], atol=2e-5, rtol=2e-4) + np.testing.assert_allclose(np.array(logits), out.logits.numpy(), atol=2e-5, rtol=2e-4) + a = out.past_key_values + offset += width + arch = ArchConfig.from_config(cfg); arch.dtype = "float32" + assert sum(t.nbytes for entry in b for t in entry if isinstance(t, mx.array)) == kv_cache_bytes(arch, 8) + + +def test_mlx_reports_real_components_and_preserves_checkpoint_on_finite_explosion(tmp_path, monkeypatch): + mx = pytest.importorskip("mlx.core") + from torch.utils.data import DataLoader, TensorDataset + import chronos.mlx.training.trainer as trainer + cfg = tiny() + calls = 0 + original = trainer._ce_stage_loss + def inject(model, ids, labels): + nonlocal calls + calls += 1 + loss = original(model, ids, labels) + return loss if calls == 1 else loss*0 + 10000 + monkeypatch.setattr(trainer, "_ce_stage_loss", inject) + checkpoint = tmp_path / "chronos_16_moe.pth" + checkpoint.write_bytes(b"existing model must survive") + ids = torch.tensor([[1, 2, 3, 4]]*3) + events = [] + result = trainer.run_mlx_stage(stage="pretrain", config=cfg, checkpoint_path=None, + save_dir=str(tmp_path), loader=DataLoader(TensorDataset(ids, ids), batch_size=1), + args=SimpleNamespace(steps=3, epochs=1, log_interval=1, dtype="bfloat16", grad_clip=1.), + progress_callback=events.append) + assert result.failure_reason == "finite_loss_divergence" + assert result.stopped and result.checkpoint_path.endswith(".recovery.pth") + assert checkpoint.read_bytes() == b"existing model must survive" + assert not any(e["event"] == "complete" for e in events) + step = next(e for e in events if e["event"] == "step") + assert step["loss"] == pytest.approx(sum(step["loss_components"].values()), abs=1e-5) + assert step["ce"] < step["loss"] + assert step["tokens_per_s"] > 0 + + +def test_real_gradio_preset_edit_save_load_and_autotune_contract(tmp_path): + gr = pytest.importorskip("gradio") + if not hasattr(gr, "Blocks"): + pytest.skip("real Gradio Blocks required") + from gradio.state_holder import SessionState + from ui.tabs.config_tab import build_config_tab + from ui.tabs.autotune_tab import build_autotune_tab + from ui.presets import get_preset, preset_names, CONFIG_INPUT_ORDER, values_in_input_order + with gr.Blocks() as app: + state, widgets, _ = build_config_tab() + build_autotune_tab(state, widgets) + session = SessionState(app) + def handler(name): + return next(fn for fn in app.fns.values() if fn.fn.__name__ == name) + async def call(name, inputs): + return await app.process_api(handler(name), inputs, state=session) + async def exercise(): + for name in preset_names(): + result = await call("apply_preset", [name]) + assert len(result["data"]) == len(handler("apply_preset").outputs) + assert session[state._id]["hidden_size"] == get_preset(name)["hidden_size"] + await call("reset_minimind", []) + cfg = dict(session[state._id], image_patch_size=8, lambda_router_locality=.031, + modalities=["image"], recurrent_strategy="fixed", recurrent_depth_jitter=.2, + max_position_embeddings=512) + source = tmp_path / "loaded.json" + source.write_text(json.dumps(cfg)) + await call("do_load", [str(source)]) + values = values_in_input_order(cfg) + values[CONFIG_INPUT_ORDER.index("hidden_size")] = 256 + values[CONFIG_INPUT_ORDER.index("max_seq_len")] = 64 + # Save must see widget edits even if their input event is still queued. + target = tmp_path / "saved.json" + await call("do_save", [None, str(target)] + values) + saved = json.loads(target.read_text()) + assert saved["hidden_size"] == 256 + assert saved["max_position_embeddings"] == saved["max_seq_len"] == 64 + assert saved["image_patch_size"] == 8 and saved["lambda_router_locality"] == .031 + assert saved["modalities"] == ["image"] + await call("update_config", [None] + values) + assert session[state._id]["hidden_size"] == 256 + await call("apply_best", [{"learning_rate": .0002}, None]) + assert session[state._id]["learning_rate"] == .0002 + assert session[state._id]["image_patch_size"] == 8 + result = await call("do_load", [str(tmp_path / "missing.json")]) + assert len(result["data"]) == len(handler("do_load").outputs) + assert session[state._id]["image_patch_size"] == 8 + asyncio.run(exercise()) diff --git a/tests/test_longrun.py b/tests/test_longrun.py new file mode 100644 index 0000000..163ab47 --- /dev/null +++ b/tests/test_longrun.py @@ -0,0 +1,64 @@ +"""Time-budgeted MLX run: deterministic resume and data sources.""" +from __future__ import annotations + +import os +import sys + +import numpy as np +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) + +mx = pytest.importorskip("mlx.core") + + +def _tiny_shards(root): + os.makedirs(os.path.join(root, "pretrain", "train")) + os.makedirs(os.path.join(root, "pretrain", "valid")) + rng = np.random.default_rng(0) + rng.integers(3, 120, 20000).astype(np.uint16).tofile(os.path.join(root, "pretrain", "train", "a.bin")) + rng.integers(3, 120, 4000).astype(np.uint16).tofile(os.path.join(root, "pretrain", "valid", "a.bin")) + + +def _spec(run_dir, shards): + from chronos.mlx.training.longrun import PhaseSpec, RunSpec + model = dict(architecture_version=3, hidden_size=64, num_hidden_layers=4, num_attention_heads=4, + num_key_value_heads=2, vocab_size=128, num_experts=4, num_experts_per_tok=2, + num_shared_experts=1, moe_intermediate_size=64, max_position_embeddings=64, + recurrent_depth=3, prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0) + phase = PhaseSpec(name="pretrain", budget_fraction=1.0, peak_lr=1e-3, tokens_per_step=256, + micro_tokens=128, seq_len=32, warmup_fraction=0.1) + return RunSpec(run_dir=run_dir, shards_dir=shards, tokenizer_dir=shards, model_config=model, + phases=[phase], budget_hours=1.0, bleu_samples=0, validate_every_s=1e9, + checkpoint_every_s=1e9, dtype="float32") + + +def _train(run, steps): + import time + if run.state["phase_total"] is None: + run.state["phase_total"] = 12 # fixed schedule (skip the throughput probe) + run.run_phase(0, time.monotonic() + 3600, max_steps=steps) + + +def test_resume_is_bit_exact(tmp_path): + """CPU stream: bit-exact. (GPU reductions are not bitwise deterministic; + measured straight-vs-resumed 7.5e-9 == straight-vs-straight on M5.)""" + from mlx.utils import tree_flatten + from chronos.mlx.training.longrun import LongRun + shards = str(tmp_path / "shards"); _tiny_shards(shards) + with mx.stream(mx.cpu): + straight = LongRun(_spec(str(tmp_path / "a"), shards), log=lambda *_: None) + _train(straight, 6) + first = LongRun(_spec(str(tmp_path / "b"), shards), log=lambda *_: None) + _train(first, 3) + first.save_state() + resumed = LongRun(_spec(str(tmp_path / "b"), shards), log=lambda *_: None) + assert resumed.try_resume() and resumed.state["phase_step"] == 3 + _train(resumed, 3) + a = dict(tree_flatten(straight.model.parameters())) + b = dict(tree_flatten(resumed.model.parameters())) + assert a.keys() == b.keys() + worst = max(float(mx.abs(a[k] - b[k]).max().item()) for k in a) + assert worst == 0.0, worst diff --git a/tests/test_mlx_speculative.py b/tests/test_mlx_speculative.py new file mode 100644 index 0000000..720db96 --- /dev/null +++ b/tests/test_mlx_speculative.py @@ -0,0 +1,208 @@ +"""MLX recurrent-depth self-speculative decoding: exactness, rollback, sampling.""" +from __future__ import annotations + +import os +import sys + +import pytest + +mx = pytest.importorskip("mlx.core") +np = pytest.importorskip("numpy") +torch = pytest.importorskip("torch") + +_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.mlx.model import ChronosMLXModel # noqa: E402 +from chronos.mlx.speculative import ( # noqa: E402 + _probs, cache_length, rejection_step, speculative_generate, truncate_cache, +) +from chronos.model.config import ChronosConfig # noqa: E402 +from chronos.model.model_chronos import ChronosForCausalLM # noqa: E402 + + +def mlx_model(mode="shared_first", window=64, **kw): + torch.manual_seed(21) + cfg = ChronosConfig(hidden_size=64, num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2, + vocab_size=96, num_experts=4, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=64, max_position_embeddings=128, recurrent_depth=4, + prelude_layers=1, coda_layers=1, dropout=0.0, fallback_mask_prob=0.0, + recurrent_attention_mode=mode, sliding_window_size=window, **kw) + pt = ChronosForCausalLM(cfg).eval() + return ChronosMLXModel.from_chronos_pytorch(pt, cfg).eval() + + +def prompt(n=7, seed=3): + return [int(t) for t in np.random.default_rng(seed).integers(3, 96, size=n)] + + +def greedy_reference(model, ids, n): + logits, _, cache = model(mx.array([ids], dtype=mx.int32)) + seq = list(ids) + nxt = int(mx.argmax(logits[0, -1]).item()) + for _ in range(n): + seq.append(nxt) + logits, _, cache = model(mx.array([[nxt]], dtype=mx.int32), cache=cache) + nxt = int(mx.argmax(logits[0, -1]).item()) + return seq + + +@pytest.mark.parametrize("window", [64, 4]) +@pytest.mark.parametrize("gamma,draft_depth", [(1, 1), (4, 1), (6, 2), (3, 4)]) +def test_mlx_greedy_speculative_equals_full_depth_greedy(window, gamma, draft_depth): + with mx.stream(mx.cpu): + model = mlx_model(window=window) + ids = prompt() + ref = greedy_reference(model, ids, 24) + out, stats = speculative_generate(model, ids, max_new_tokens=24, gamma=gamma, + draft_depth=draft_depth, gamma_max=8) + assert out == ref + assert stats.generated == 24 + if draft_depth < 4: + assert stats.verify_calls < 24 # verification is batched + assert 0.0 <= stats.acceptance_rate <= 1.0 + else: + assert stats.draft_calls == 0 # full draft depth = plain decoding + + +@pytest.mark.parametrize("overrides", [dict(kv_cache_dtype="int8"), dict(mla_decode="absorbed", mla_cache_rope=True)]) +def test_mlx_speculative_exact_with_cache_options(overrides): + with mx.stream(mx.cpu): + model = mlx_model(window=4, **overrides) + ids = prompt(9, seed=5) + ref = greedy_reference(model, ids, 16) + out, _ = speculative_generate(model, ids, max_new_tokens=16, gamma=3, draft_depth=1) + assert out == ref + + +def test_mlx_per_iteration_falls_back_to_plain_decoding(): + with mx.stream(mx.cpu): + model = mlx_model(mode="per_iteration") + ids = prompt(5) + ref = greedy_reference(model, ids, 10) + out, stats = speculative_generate(model, ids, max_new_tokens=10, gamma=4, draft_depth=1) + assert out == ref and stats.draft_calls == 0 + + +@pytest.mark.parametrize("window", [64, 3]) +def test_mlx_cache_rollback_equals_fresh_prefill(window): + from chronos.mlx.speculative import set_cache_slack + with mx.stream(mx.cpu): + model = mlx_model(window=window) + set_cache_slack(model, 6) + ids = prompt(8) + extra = [11, 12, 13, 14, 15] + _, _, cache = model(mx.array([ids], dtype=mx.int32)) + _, _, longer = model(mx.array([extra], dtype=mx.int32), cache=cache) + rolled = truncate_cache(longer, len(ids) + 2) # keep two of the five extra tokens + assert cache_length(rolled) == len(ids) + 2 + probe = mx.array([[40]], dtype=mx.int32) + got, _, _ = model(probe, cache=rolled) + _, _, fresh = model(mx.array([ids + extra[:2]], dtype=mx.int32)) + want, _, _ = model(probe, cache=fresh) + mx.eval(got, want) + np.testing.assert_allclose(np.array(got), np.array(want), atol=2e-5, rtol=2e-4) + + +def test_mlx_sampled_speculative_is_reproducible_and_stops_at_eos(): + with mx.stream(mx.cpu): + model = mlx_model() + ids = prompt(6) + kw = dict(max_new_tokens=20, gamma=3, draft_depth=1, do_sample=True, temperature=0.9, top_k=20, top_p=0.95) + a, sa = speculative_generate(model, ids, seed=7, **kw) + b, _ = speculative_generate(model, ids, seed=7, **kw) + assert a == b and sa.generated == 20 and all(0 <= t < 96 for t in a) + eos = a[len(ids) + 4] + c, sc = speculative_generate(model, ids, seed=7, eos_token_id=eos, **kw) + assert c[-1] == eos and eos not in c[len(ids):-1] and sc.generated == len(c) - len(ids) + + +def test_mlx_probs_top_p_keeps_the_nucleus_in_vocab_order(): + with mx.stream(mx.cpu): + logits = mx.array([0.0, 3.0, 1.0, 2.9, -1.0]) + p = np.array(_probs(logits, 1.0, 0, 0.6)) + assert set(np.nonzero(p)[0]) == {1, 3} # the two largest, in their own slots + np.testing.assert_allclose(p.sum(), 1.0, atol=1e-6) + + +def test_mlx_rejection_step_preserves_the_target_distribution(): + p_np, q_np = np.array([0.5, 0.2, 0.2, 0.1]), np.array([0.1, 0.6, 0.2, 0.1]) + trials, counts = 3000, np.zeros(4) + with mx.stream(mx.cpu): + p, q = mx.array(p_np), mx.array(q_np) + key = mx.random.key(0) + for _ in range(trials): + key, k1, k2 = mx.random.split(key, 3) + draft = int(mx.random.categorical(mx.log(q), key=k1).item()) + _, tok = rejection_step(p, q, draft, k2) + counts[tok] += 1 + np.testing.assert_allclose(counts / trials, p_np, atol=0.035) # ≈4σ at n=3000 + + +def _engine(tmp_path, **kw): + from chronos.mlx.inference import ChronosMLXInferenceEngine + model = mlx_model(window=4, **kw) + return model, ChronosMLXInferenceEngine(model, model.config, ssd_dir=str(tmp_path / "experts")) + + +def test_mlx_engine_speculative_matches_plain_greedy(tmp_path): + with mx.stream(mx.cpu): + model, engine = _engine(tmp_path) + engine.store.storage_format = "full_dram" + try: + ids = mx.array([prompt(7)], dtype=mx.int32) + plain = list(engine.generate(ids, max_new_tokens=20, temperature=0.0)) + spec = list(engine.generate(ids, max_new_tokens=20, temperature=0.0, speculative=True, + draft_depth=1, gamma=4)) + stats = engine.last_stats + eos = spec[5] + stopped = list(engine.generate(ids, max_new_tokens=20, temperature=0.0, speculative=True, + eos_token_id=eos)) + stopped_plain = list(engine.generate(ids, max_new_tokens=20, temperature=0.0, eos_token_id=eos)) + finally: + engine.stop() + assert spec == plain and len(spec) == 20 + assert stats["tokens_generated"] == 20 and stats["speculative"]["verify_calls"] < 20 + assert stopped == stopped_plain == spec[:spec.index(eos) + 1] + + +def test_mlx_engine_speculative_runs_with_lazy_experts(tmp_path): + 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]) + assert engine._all_resident_fast_path_enabled() is False + toks = list(engine.generate(mx.array([prompt(6)], dtype=mx.int32), max_new_tokens=12, + temperature=0.0, speculative=True, draft_depth=1, gamma=3)) + stats = engine.last_stats + finally: + engine.stop() + assert len(toks) == 12 and all(0 <= t < 96 for t in toks) + assert stats["all_resident_fast_path"] is False + assert stats["speculative"]["draft_calls"] > 0 and stats["speculative"]["generated"] == 12 + + +@pytest.mark.parametrize("min_depth", [2, 3]) +def test_draft_depth_below_min_recurrent_depth_keeps_the_last_state(min_depth): + """A draft that stops before ``min_recurrent_depth`` must not zero the ACT mixture.""" + torch.manual_seed(21) + cfg = ChronosConfig(hidden_size=64, num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2, + vocab_size=96, num_experts=4, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=64, max_position_embeddings=128, recurrent_depth=4, + min_recurrent_depth=min_depth, prelude_layers=1, coda_layers=1, dropout=0.0, + fallback_mask_prob=0.0, recurrent_attention_mode="shared_first", sliding_window_size=64) + pt = ChronosForCausalLM(cfg).eval() + ids = np.random.default_rng(3).integers(3, 96, size=(1, 7)) + with mx.stream(mx.cpu): + model = ChronosMLXModel.from_chronos_pytorch(pt, cfg).eval() + for depth in range(1, min_depth): + with torch.no_grad(): + ref = pt(torch.tensor(ids), recurrent_depth=depth).logits.numpy() + full = pt(torch.tensor(ids)).logits.numpy() + got, _, _ = model(mx.array(ids, dtype=mx.int32), recurrent_depth=depth) + mx.eval(got) + np.testing.assert_allclose(np.array(got), ref, atol=3e-5, rtol=3e-4) + # a zeroed mixture feeds identical states to the coda: every position gets the same logits + assert np.array(got).std(axis=1).max() > 1e-3 + assert np.abs(ref - full).max() > 1e-4 # the draft really ran fewer iterations diff --git a/tests/test_moe_dispatch.py b/tests/test_moe_dispatch.py new file mode 100644 index 0000000..739bb6e --- /dev/null +++ b/tests/test_moe_dispatch.py @@ -0,0 +1,167 @@ +"""Sparse expert dispatch must match the masked-dense oracle in value and gradient.""" +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +_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.model.config import ChronosConfig +from chronos.model.moe_chronos import ChronosMOEFeedForward + + +def moe_cfg(**kw): + values = dict(hidden_size=48, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, + vocab_size=64, num_experts=6, num_experts_per_tok=2, num_shared_experts=1, + moe_intermediate_size=40, dropout=0.0, fallback_mask_prob=0.0) + values.update(kw) + return ChronosConfig(**values) + + +def torch_dense_oracle(moe, x): + """Every expert on every token, weighted by the routed top-k mask.""" + B, S, H = x.shape + flat = x.reshape(-1, H) + scores = torch.softmax(moe.gate(flat), dim=-1) + weight, idx = torch.topk(scores, moe.num_experts_per_tok, dim=-1, sorted=False) + if moe.norm_topk_prob: + weight = weight / (weight.sum(-1, keepdim=True) + 1e-20) + y = torch.zeros_like(flat) + for e, expert in enumerate(moe.experts): + y = y + expert(flat) * torch.where(idx == e, weight, 0).sum(-1, keepdim=True) + y = y + sum(s(flat) for s in moe.shared_experts) / len(moe.shared_experts) + return y.view(B, S, H) + + +@pytest.mark.parametrize("top_k", [1, 2]) +def test_torch_sorted_dispatch_matches_dense_value_and_grad(top_k): + torch.manual_seed(0) + moe = ChronosMOEFeedForward(moe_cfg(num_experts_per_tok=top_k)).train() + x = torch.randn(3, 7, 48, requires_grad=True) + sparse = moe(x) + grads_sparse = torch.autograd.grad(sparse.square().sum(), [x, *moe.parameters()], allow_unused=True) + x2 = x.detach().clone().requires_grad_(True) + dense = torch_dense_oracle(moe, x2) + grads_dense = torch.autograd.grad(dense.square().sum(), [x2, *moe.parameters()], allow_unused=True) + torch.testing.assert_close(sparse, dense, atol=1e-5, rtol=1e-5) + for a, b in zip(grads_sparse, grads_dense): + if b is None: + continue + torch.testing.assert_close(a if a is not None else torch.zeros_like(b), b, atol=1e-5, rtol=1e-4) + + +mx = pytest.importorskip("mlx.core") + + +def _mlx_case(dtype, seed=0, E=6, H=64, I=48, N=40, K=2): + mx.random.seed(seed) + with mx.stream(mx.cpu): + params = [mx.random.normal(s) * 0.1 for s in ((E, I, H), (E, I, H), (E, H, I), (N, H))] + idx = mx.random.randint(0, E, (N, K)) + weight = mx.softmax(mx.random.normal((N, K)), axis=-1) + mx.eval(params, idx, weight) + return params, idx, weight + + +def _value_and_grads(fn, params, idx, weight): + """params = [w_gate, w_up, w_down, x]; grads w.r.t. all four.""" + def loss(*p): + return fn(p[3], idx, weight, *p[:3]).astype(mx.float32).square().sum() + value = fn(params[3], idx, weight, *params[:3]) + grads = mx.grad(loss, argnums=(0, 1, 2, 3))(*params) + mx.eval(value, grads) + return value, grads + + +def _rel(a, b): + return (mx.abs(a.astype(mx.float32) - b.astype(mx.float32)).max() / + (mx.abs(b.astype(mx.float32)).max() + 1e-9)).item() + + +@pytest.mark.parametrize("dtype_name,stream_name,tol", [ + ("float32", "cpu", 1e-5), + # Apple M5 GPU fp32 matmul is TF32-class: dense and sparse both measure + # ~0.2-0.6% relative error vs CPU fp32 (2026-11-05), so budget 1e-2. + ("float32", "gpu", 1e-2), + ("bfloat16", "gpu", 2.5e-2), +]) +def test_mlx_grouped_gather_mm_matches_dense_and_fp32_reference(dtype_name, stream_name, tol): + from chronos.mlx.moe_kernels import dense_reference_swiglu, gather_mm_supported, grouped_swiglu + if not gather_mm_supported(): + pytest.skip("mlx < 0.32.4: gather_mm unreliable on this GPU; dense path is used") + dtype = getattr(mx, dtype_name) + params32, idx, weight32 = _mlx_case(mx.float32) + with mx.stream(mx.cpu): + ref_value, ref_grads = _value_and_grads(dense_reference_swiglu, params32, idx, weight32) + stream = mx.cpu if stream_name == "cpu" else mx.gpu + with mx.stream(stream): + params = [p.astype(dtype) for p in params32] + weight = weight32.astype(dtype) + sparse_value, sparse_grads = _value_and_grads(grouped_swiglu, params, idx, weight) + dense_value, dense_grads = _value_and_grads(dense_reference_swiglu, params, idx, weight) + # Both paths meet the absolute budget against the fp32 CPU reference ... + assert _rel(sparse_value, ref_value) <= tol + assert _rel(dense_value, ref_value) <= tol + for s, d, r in zip(sparse_grads, dense_grads, ref_grads): + sparse_err, dense_err = _rel(s, r), _rel(d, r) + assert sparse_err <= tol, (sparse_err, dense_err) + assert dense_err <= tol, (sparse_err, dense_err) + # ... and sparse is never materially worse than the dense oracle. + assert sparse_err <= 2.0 * dense_err + 1e-6, (sparse_err, dense_err) + + +@pytest.mark.parametrize("dtype_name,tol", [("float32", 1e-2), ("bfloat16", 3e-2)]) +def test_mlx_moe_module_sparse_matches_dense_value_and_grad(dtype_name, tol): + import mlx.nn as nn + from mlx.utils import tree_flatten, tree_map + from chronos.mlx.moe import ChronosMLXMOE + from chronos.mlx.moe_kernels import gather_mm_supported + if not gather_mm_supported(): + pytest.skip("mlx < 0.32.4") + dtype = getattr(mx, dtype_name) + mx.random.seed(3) + moe = ChronosMLXMOE(moe_cfg()) + moe.update(tree_map(lambda a: a.astype(dtype), moe.parameters())) + x = (mx.random.normal((2, 9, 48)) * 0.5).astype(dtype) + + def run(sparse): + moe.sparse_dispatch = sparse + value, grads = nn.value_and_grad(moe, lambda inp: moe(inp).astype(mx.float32).square().sum())(x) + mx.eval(value, grads) + return value, dict(tree_flatten(grads)) + + dense_value, dense_grads = run(False) + sparse_value, sparse_grads = run(True) + assert abs(sparse_value.item() - dense_value.item()) <= tol * abs(dense_value.item()) + assert sparse_grads.keys() == dense_grads.keys() + for key, dense_grad in dense_grads.items(): + assert _rel(sparse_grads[key], dense_grad) <= tol, key + + + +def test_mlx_grouped_fallback_matches_dense_oracle(): + from chronos.mlx.moe_kernels import dense_reference_swiglu, gather_mm_supported, grouped_swiglu + if not gather_mm_supported(): + pytest.skip("mlx < 0.32.4") + params, idx, weight = _mlx_case(mx.float32) + with mx.stream(mx.cpu): + available = mx.array([True, False, True, True, False, True]) + shared = mx.random.normal((params[3].shape[0], params[3].shape[1])) + def run(fn): + def loss(*p): + return fn(p[3], idx, weight, *p[:3], fallback=(available, shared)).square().sum() + return fn(params[3], idx, weight, *params[:3], fallback=(available, shared)), \ + mx.grad(loss, argnums=(0, 1, 2, 3))(*params) + sv, sg = run(grouped_swiglu) + dv, dg = run(dense_reference_swiglu) + mx.eval(sv, sg, dv, dg) + assert _rel(sv, dv) < 1e-5 + for a, b in zip(sg, dg): + assert _rel(a, b) < 1e-4 + # Unavailable experts receive no gradient (they were not executed). + assert mx.abs(sg[0][1]).max().item() == 0.0 and mx.abs(sg[0][4]).max().item() == 0.0 diff --git a/tests/test_monitor.py b/tests/test_monitor.py new file mode 100644 index 0000000..04c6209 --- /dev/null +++ b/tests/test_monitor.py @@ -0,0 +1,301 @@ +"""Read-only training monitor: JSONL reader, run snapshots, API auth, Gradio tab, telemetry.""" +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +from datetime import datetime, timedelta + +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) + + +# ---- fixtures ---------------------------------------------------------------- +def _write_jsonl(path, rows, mode="w"): + with open(path, mode, encoding="utf-8") as fh: + for row in rows: + fh.write(json.dumps(row) + "\n") + + +def _synthetic_run(root, name="demo", *, final=False, steps=10): + """A run directory as an old (pre-telemetry) longrun leaves it: spec + metrics only.""" + path = os.path.join(root, name) + os.makedirs(os.path.join(path, "checkpoints")) + spec = {"run_dir": path, "budget_hours": 10.0, "eval_reserve_fraction": 0.1, + "phases": [{"name": "pretrain", "budget_fraction": 0.7}, {"name": "sft", "budget_fraction": 0.3}]} + with open(os.path.join(path, "run_spec.json"), "w", encoding="utf-8") as fh: + json.dump(spec, fh) + t0 = datetime(2026, 11, 6, 0, 0, 0) + rows = [{"phase": "pretrain", "event": "phase_start", "step": 0, "time": t0.isoformat()}] + for i in range(1, steps + 1): + rows.append({"phase": "pretrain", "step": i, "total": 100, "loss": 5.0 - 0.1 * i, "primary_loss": 4.5 - 0.1 * i, + "lr": 1e-3, "grad_norm": 1.0, "tokens_per_s": 1000.0, "step_s": 2.0, + "time": (t0 + timedelta(seconds=2 * i)).isoformat()}) + rows.append({"phase": "pretrain", "event": "validation", "step": 4, "valid_ce": 3.2}) + rows.append({"phase": "pretrain", "event": "validation", "step": 8, "valid_ce": 2.9}) + rows.append({"phase": "pretrain", "event": "validation", "step": 9, "valid_ce": 3.0}) + if final: + rows.append({"phase": "pretrain", "event": "final", "step": steps}) + _write_jsonl(os.path.join(path, "metrics.jsonl"), rows) + return path + + +# ---- JsonLines ----------------------------------------------------------------- +def test_jsonlines_incremental_partial_and_truncate(tmp_path): + from chronos.monitor import JsonLines + path = str(tmp_path / "m.jsonl") + reader = JsonLines(path) + assert reader.read() == [] # missing file + _write_jsonl(path, [{"a": 1}, {"a": 2}]) + assert [r["a"] for r in reader.read()] == [1, 2] + with open(path, "a", encoding="utf-8") as fh: # in-flight write: no newline yet + fh.write('{"a": 3}\n{"a": 4') + assert [r["a"] for r in reader.read()] == [1, 2, 3] + with open(path, "a", encoding="utf-8") as fh: + fh.write('}\nnot json\n') + assert [r["a"] for r in reader.read()] == [1, 2, 3, 4] + assert reader.bad_lines == 1 + _write_jsonl(path, [{"a": 9}]) # truncated / rewritten + assert [r["a"] for r in reader.read()] == [9] + capped = JsonLines(path, max_rows=2) + _write_jsonl(path, [{"a": i} for i in range(5)]) + assert [r["a"] for r in capped.read()] == [3, 4] + + +# ---- snapshot ------------------------------------------------------------------ +def test_snapshot_derives_status_progress_and_best_valid(tmp_path): + from chronos.monitor import RunMonitor + _synthetic_run(str(tmp_path), "demo", steps=10) + os.makedirs(tmp_path / "not_a_run") + mon = RunMonitor(str(tmp_path)) + assert [r["name"] for r in mon.runs()] == ["demo"] + snap = mon.snapshot("demo") + status = snap["status"] + assert status["source"] == "derived" and not status["alive"] and status["state"] == "stopped" + assert status["log_interval_s"] == pytest.approx(2.0) + assert snap["latest"]["step"] == 10 + assert snap["best_valid"]["pretrain"] == {"step": 8, "valid_ce": 2.9, "time": None} + progress = snap["progress"] + assert [p["state"] for p in progress["phases"]] == ["active", "pending"] + assert progress["phase_fraction"] == pytest.approx(0.1) + assert progress["phase_remaining_s"] == pytest.approx(90 * 2.0) + # Later phase budget (0.3 · 9 h training) and the eval reserve (1 h) are added. + assert progress["run_remaining_s"] == pytest.approx(180 + 0.3 * 9 * 3600 + 3600) + assert {e["event"] for e in snap["events"]} == {"phase_start", "validation"} + assert len(snap["history"]["pretrain"]) == 10 + assert [r["step"] for r in mon.metrics("demo", "pretrain", since_step=8)] == [9, 10] + + +def test_snapshot_final_event_marks_finished(tmp_path): + from chronos.monitor import RunMonitor + _synthetic_run(str(tmp_path), "done", final=True) + snap = RunMonitor(str(tmp_path)).snapshot("done") + assert snap["status"]["state"] == "finished" + + +def test_status_json_heartbeat_wins_over_derivation(tmp_path): + from chronos.monitor import RunMonitor + path = _synthetic_run(str(tmp_path), "hb") + with open(os.path.join(path, "status.json"), "w", encoding="utf-8") as fh: + json.dump({"state": "failed", "pid": 2 ** 22 + 12345, "phase": "pretrain", "error": "boom", + "phase_wall_s": 3600.0, "phase_budget_s": 4 * 3600.0}, fh) + snap = RunMonitor(str(tmp_path)).snapshot("hb") + assert snap["status"]["source"] == "status.json" + assert snap["status"]["state"] == "failed" and snap["status"]["error"] == "boom" + assert snap["progress"]["phase_fraction"] == pytest.approx(0.25) + + +@pytest.mark.parametrize("name", ["../etc", ".hidden", "a/b", "..", "missing"]) +def test_unknown_or_unsafe_run_names_are_rejected(tmp_path, name): + from chronos.monitor import RunMonitor + _synthetic_run(str(tmp_path), "demo") + with pytest.raises(KeyError): + RunMonitor(str(tmp_path)).snapshot(name) + + +def test_mach_tick_scale_is_positive(): + from chronos.monitor.process import _mach_tick_ns + assert _mach_tick_ns() > 0 + + +# ---- standalone API ------------------------------------------------------------ +def _client(root, **kwargs): + pytest.importorskip("fastapi") + pytest.importorskip("httpx") + from fastapi.testclient import TestClient + from chronos.monitor.server import create_monitor_app + return TestClient(create_monitor_app(str(root), **kwargs)) + + +def test_api_open_on_loopback_without_key(tmp_path): + _synthetic_run(str(tmp_path), "demo") + client = _client(tmp_path) + assert client.get("/api/session").json() == {"auth_required": False, "authenticated": True, "local_only": True} + state = client.get("/api/state", params={"run": "demo"}) + assert state.status_code == 200 and state.json()["run"] == "demo" + assert client.get("/api/state", params={"run": "../demo"}).status_code == 404 + assert [r["step"] for r in client.get("/api/metrics", params={"run": "demo", "since": 8}).json()["rows"]] == [9, 10] + assert client.get("/api/parameters", params={"run": "demo"}).json()["spec"]["budget_hours"] == 10.0 + assert client.get("/api/tensors", params={"run": "demo"}).json() == {"tensors": []} + page = client.get("/") + assert page.status_code == 200 and "static/app.js" in page.text + for header in ("Content-Security-Policy", "X-Content-Type-Options", "X-Frame-Options", "Cache-Control"): + assert header in page.headers + assert "script-src 'self'" in page.headers["Content-Security-Policy"] + assert client.get("/static/app.js").headers["content-type"].startswith("text/javascript") + assert client.get("/static/app.css").status_code == 200 + for name in ("index.html", "..%2Fserver.py", "server.py"): + assert client.get(f"/static/{name}").status_code == 404 + + +def test_api_key_login_cookie_bearer_and_rate_limit(tmp_path): + from chronos.monitor.server import COOKIE + _synthetic_run(str(tmp_path), "demo") + key = "correct horse battery staple" + client = _client(tmp_path, api_key=key) + assert client.get("/").status_code == 200 # login page shell is public + assert client.get("/api/session").json()["auth_required"] is True + for path in ("/api/runs", "/api/state", "/api/metrics", "/api/parameters", "/api/tensors", "/api/log"): + assert client.get(path).status_code == 401, path + assert client.get("/api/runs", headers={"Authorization": f"Bearer {key}"}).status_code == 200 + assert client.get("/api/runs", headers={"Authorization": "Bearer nope"}).status_code == 401 + ok = client.post("/api/login", json={"key": key}) + assert ok.status_code == 200 + cookie = ok.headers["set-cookie"].lower() + assert COOKIE in cookie and "httponly" in cookie and "samesite=strict" in cookie and key not in cookie + assert client.get("/api/runs").status_code == 200 # session cookie + client.post("/api/logout") + client.cookies.clear() + assert client.get("/api/runs").status_code == 401 + other = _client(tmp_path, api_key=key) + codes = [other.post("/api/login", json={"key": f"wrong{i}"}).status_code for i in range(8)] + assert codes == [401] * 8 + assert other.post("/api/login", json={"key": key}).status_code == 429 # blocked even with the right key + assert other.post("/api/login", content=b"not json").status_code == 429 + + +def test_check_bind_requires_remote_flag_and_key(): + from chronos.monitor.server import check_bind + assert check_bind("127.0.0.1", False, None) is None + assert check_bind("::1", False, None) is None + assert check_bind("0.0.0.0", False, "k") + assert check_bind("0.0.0.0", True, None) + assert check_bind("0.0.0.0", True, "k") is None + + +def test_serve_mounts_monitor_with_same_key(tmp_path): + pytest.importorskip("fastapi") + from fastapi.testclient import TestClient + from chronos.serving.ollama_server import create_app + + class _Runtime: + loaded, last_used, keep_alive_s = {}, {}, 300 + + def available(self): + return {} + + _synthetic_run(str(tmp_path), "demo") + key = "k-123" + client = TestClient(create_app(_Runtime(), api_key=key, monitor_runs_dir=str(tmp_path))) + assert client.get("/api/tags").status_code == 401 + assert client.get("/api/tags", headers={"Authorization": f"Bearer {key}"}).status_code == 200 + assert client.get("/monitor", follow_redirects=False).status_code in (301, 302, 307, 308) + page = client.get("/monitor/") + assert page.status_code == 200 and "static/app.css" in page.text + assert client.get("/monitor/api/state").status_code == 401 + login = client.post("/monitor/api/login", json={"key": key}) + assert login.status_code == 200 and "path=/monitor/" in login.headers["set-cookie"].lower() + assert client.get("/monitor/api/state", params={"run": "demo"}).json()["run"] == "demo" + no_monitor = TestClient(create_app(_Runtime(), api_key=key)) + assert no_monitor.get("/monitor/").status_code == 401 + + +# ---- longrun telemetry -> monitor -> Gradio tab --------------------------------- +def _telemetry_run(tmp_path): + """Three CPU steps of a tiny MLX run with telemetry on (writes status/runtime/parameters).""" + mx = pytest.importorskip("mlx.core") + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from test_longrun import _spec, _tiny_shards, _train + from chronos.mlx.training.longrun import LongRun + shards = str(tmp_path / "shards"); _tiny_shards(shards) + run_dir = str(tmp_path / "runs" / "tiny") + with mx.stream(mx.cpu): + run = LongRun(_spec(run_dir, shards), log=lambda *_: None) + _train(run, 3) + return run_dir + + +def test_longrun_telemetry_feeds_monitor_and_gradio_tab(tmp_path): + run_dir = _telemetry_run(tmp_path) + for name in ("status.json", "runtime.json", "parameters.jsonl", "metrics.jsonl", "run_spec.json"): + assert os.path.isfile(os.path.join(run_dir, name)), name + with open(os.path.join(run_dir, "runtime.json"), encoding="utf-8") as fh: + runtime = json.load(fh) + assert runtime["precision"]["compute"] and runtime["code"]["sha256"] + from chronos.monitor import RunMonitor + mon = RunMonitor(os.path.dirname(run_dir)) + snap = mon.snapshot("tiny") + assert snap["status"]["source"] == "status.json" and snap["status"]["pid"] == os.getpid() + assert snap["latest"]["step"] == 3 + summary = snap["parameters"] + assert summary["step"] == 1 and summary["tensor_count"] > 0 and "tensors" not in summary + assert {"experts", "attention", "embedding"} <= set(summary["groups"]) + full = mon.parameters("tiny") + first = full["tensors"][0] + assert {"name", "group", "shape", "rms", "grad_rms", "update_ratio"} <= set(first) + assert mon.tensor_history("tiny", first["name"])[0]["step"] == 1 + assert summary["routing"] and summary["routing"]["load"] + + pytest.importorskip("gradio") + from ui.tabs.monitor_tab import refresh + out = refresh(os.path.dirname(run_dir), "tiny", "pretrain", "primary_loss") + assert len(out) == 31 + status_html = out[2] + assert "tiny" in status_html and "t", [""], [[0.1, 0.9]]) + assert " 90 + for _, code in LANGUAGE_CHOICES: + missing = keys - set(TRANSLATIONS[code]) + assert not missing, (code, sorted(missing)[:5]) + + +def test_monitor_static_assets_parse(): + static = os.path.join(_repo_root, "chronos", "monitor", "static") + for name in ("index.html", "app.js", "app.css"): + assert os.path.isfile(os.path.join(static, name)) + with open(os.path.join(static, "index.html"), encoding="utf-8") as fh: + page = fh.read() + assert "" + + +class UIEventContractMiddleware: + def __init__(self, app, contract): + self.app = app + self.contract = contract + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + return await self.app(scope, receive, send) + path = scope.get("path", "").rstrip("/") + method = scope["method"] + + async def respond(status, payload): + body = json.dumps(payload).encode() + await send({"type": "http.response.start", "status": status, "headers": [ + (b"content-type", b"application/json"), (b"cache-control", b"no-store"), + (b"content-length", str(len(body)).encode()), + ]}) + await send({"type": "http.response.body", "body": body}) + + if method == "GET" and path.endswith(CONTRACT_PATH): + return await respond(200, {"token": self.contract.token}) + + queued_event = method == "POST" and path.endswith("/queue/join") + direct_event = method == "POST" and any( + segment in path for segment in ("/run/", "/api/", "/call/") + ) + client_token = dict(scope.get("headers", [])).get(HEADER.encode(), b"").decode() + reason = None + if (queued_event or direct_event) and client_token and client_token != self.contract.token: + reason = "server_restarted" + + if queued_event and reason is None: + chunks = [] + while True: + message = await receive() + if message["type"] == "http.disconnect": + return + chunks.append(message.get("body", b"")) + if not message.get("more_body", False): + break + body = b"".join(chunks) + try: + payload = json.loads(body) + except (ValueError, UnicodeDecodeError): + payload = None # Let FastAPI produce its normal validation error. + if isinstance(payload, dict) and not payload.get("simple_format"): + blocks = self.contract.blocks + functions = blocks.fns + # Respect per-session functions if Gradio rendered them dynamically. + session = blocks.state_holder.session_data.get(payload.get("session_hash")) + if session is not None: + functions = session.blocks_config.fns + index = payload.get("fn_index") + fn = functions.get(index) if isinstance(index, int) else None + if fn is None: + reason = "unknown_event" + elif isinstance(payload.get("data"), list) and len(payload["data"]) != len(fn.inputs): + reason = "input_count_mismatch" + + delivered = False + original_receive = receive + + async def replay(): + nonlocal delivered + if not delivered: + delivered = True + return {"type": "http.request", "body": body, "more_body": False} + return await original_receive() + + receive = replay + + if reason: + outdated = client_token != self.contract.token + return await respond(409 if outdated else 422, { + "code": "chronos_ui_outdated" if outdated else "chronos_invalid_event", + "reason": reason, + "error": "页面与当前服务的事件接口不一致，请刷新页面。 / Reload the page to synchronize UI events.", + "detail": "No event was executed. Refresh the Chronos page and retry.", + }) + + async def no_cached_page(message): + if method == "GET" and message["type"] == "http.response.start": + headers = message.get("headers", []) + content_type = dict(headers).get(b"content-type", b"") + if b"text/html" in content_type or path.endswith("/config"): + message = dict(message, headers=[(k, v) for k, v in headers if k.lower() != b"cache-control"] + + [(b"cache-control", b"no-store")]) + await send(message) + + return await self.app(scope, receive, no_cached_page) + + +def event_contract_app_kwargs(blocks): + """Pass through Blocks.launch(app_kwargs=...) so Gradio keeps its lifespan.""" + contract = getattr(blocks, "chronos_event_contract", None) + if contract is None: + return {} + from starlette.middleware import Middleware + return {"middleware": [Middleware(UIEventContractMiddleware, contract=contract)]} diff --git a/ui/gradio_compat.py b/ui/gradio_compat.py index 5d20892..600ab93 100644 --- a/ui/gradio_compat.py +++ b/ui/gradio_compat.py @@ -36,9 +36,18 @@ def click(self, *args, **kwargs): def change(self, *args, **kwargs): return self + def input(self, *args, **kwargs): + return self + def tick(self, *args, **kwargs): return self + def select(self, *args, **kwargs): + return self + + def submit(self, *args, **kwargs): + return self + def launch(self, **kwargs): self.launch_kwargs = kwargs return self @@ -54,6 +63,7 @@ def _update(**kwargs): Checkbox=_Component, CheckboxGroup=_Component, Column=_Component, + Dataframe=_Component, Dropdown=_Component, File=_Component, HTML=_Component, @@ -70,4 +80,5 @@ def _update(**kwargs): Textbox=_Component, Timer=_Component, update=_update, + skip=lambda: None, ) diff --git a/ui/i18n.py b/ui/i18n.py index ccfd369..5b6626e 100644 --- a/ui/i18n.py +++ b/ui/i18n.py @@ -33,11 +33,15 @@ "tab.export": "📦 导出", "designer.title": "实时参数量与内存估算", "designer.total": "总参数量", - "designer.active": "激活参数量 (每 token)", - "designer.vram": "VRAM 估算", - "designer.ssd": "SSD 估算", + "designer.active": '激活参数范围（每 token，去重）', + "designer.compute": '权重使用量（含循环重复）', + "designer.vram": '推理驻留权重 + KV', + "designer.ram": '示例 warm 缓存预算', + "designer.ssd": '完整专家 SSD 存储', "designer.kv": "KV 缓存估算", - "designer.tps": "估算 tokens/sec", + "designer.tps": '带宽模型示例 tokens/s', + "designer.details": '参数与内存明细', + "designer.notes": '激活范围取决于各轮路由是否复用专家，不等于驻留缓存。计算量可超过总参数量。内存按 batch=1 和当前序列长度估算；auto 暂按 FP32。训练权重、梯度和 Adam 状态单列于明细，未包含激活、临时缓冲和分配器峰值。统一内存的 RAM/VRAM 不应重复相加。速度仅为 85% 命中率、3 GB/s SSD、600 GB/s 内存带宽的示例，请以实测为准。', "iomon.title": "实时 IO 与缓存指标", "iomon.refresh": "刷新", "pipeline.title": "Pretrain → SFT → DPO → ORPO → GRPO → Distill", @@ -46,6 +50,27 @@ "config.loss": "损失系数", "config.hw": "硬件预算", "config.train": "训练参数", + "config.recurrent_section": "循环深度 / 分页 / 多模态", + "config.recurrent_enabled": "启用循环深度", + "config.recurrent_groups": "循环组数", + "config.recurrent_depth": "最大循环深度", + "config.min_recurrent_depth": "最小循环深度", + "config.adaptive_halting": "自适应停机 (ACT)", + "config.halting_threshold": "停机阈值", + "config.max_inference_depth": "推理最大深度", + "config.ponder_loss_coef": "Ponder 损失系数", + "config.recurrent_distill_coef": "循环状态一致性系数", + "config.recurrent_stability_coef": "循环稳定性系数", + "config.recurrent_latent_refinement": "循环中间步仅做潜变量精炼", + "config.depth_schedule": "深度调度", + "config.depth_schedule_warmup_fraction": "自适应预热比例", + "config.layer_aware_cache": "层/组感知缓存", + "config.prefetch_depth": "预取深度", + "config.training_paging": "训练期分页", + "config.optimizer_paging": "优化器分页", + "config.page_size_mb": "分页大小 (MB)", + "config.optimizer_page_dir": "优化器页目录", + "config.modalities": "模态 (自动/图像/音频)", "config.hidden_size": "隐藏层维度", "config.num_layers": "层数", "config.num_experts": "专家数", @@ -163,12 +188,12 @@ "autotune.clear_log": "清空日志文件", "autotune.apply_best": "应用最优参数 → 配置", "autotune.status": "状态", - "data.title": "📊 数据检查器", - "data.path": "数据集路径 (.jsonl)", + "data.title": "📊 数据检查器（任意任务）", + "data.path": "数据集路径（文件、目录或通配符；jsonl / json / parquet / csv / tsv / txt）", "data.scan": "▶ 扫描", "data.summary": "数据集概要", "data.preview": "样本预览", - "data.length_hist": "Token 长度分布", + "data.length_hist": "每条记录的字符长度分布", }, "zh-Hant": { "app.title": "Project Chronos — 端側 MoE 推理工作台", @@ -185,11 +210,15 @@ "tab.export": "📦 匯出", "designer.title": "即時參數量與記憶體估算", "designer.total": "總參數量", - "designer.active": "啟用參數量 (每 token)", - "designer.vram": "VRAM 估算", - "designer.ssd": "SSD 估算", + "designer.active": '啟用參數範圍（每 token，去重）', + "designer.compute": '權重使用量（含循環重複）', + "designer.vram": '推理常駐權重 + KV', + "designer.ram": '範例 warm 快取預算', + "designer.ssd": '完整專家 SSD 儲存', "designer.kv": "KV 快取估算", - "designer.tps": "估算 tokens/sec", + "designer.tps": '頻寬模型範例 tokens/s', + "designer.details": '參數與記憶體明細', + "designer.notes": '啟用範圍取決於每輪路由是否重用專家，不等於常駐快取。記憶體以 batch=1、目前序列長度估算；auto 暫用 FP32。訓練狀態另列，不含啟用值與暫存峰值。統一記憶體勿重複加總 RAM/VRAM。速度為 85% 命中率、3 GB/s SSD、600 GB/s 記憶體頻寬範例。', "iomon.title": "即時 IO 與快取指標", "iomon.refresh": "重新整理", "pipeline.title": "Pretrain → SFT → DPO → ORPO → GRPO → Distill", @@ -198,6 +227,27 @@ "config.loss": "損失係數", "config.hw": "硬體預算", "config.train": "訓練參數", + "config.recurrent_section": "循環深度 / 分頁 / 多模態", + "config.recurrent_enabled": "啟用循環深度", + "config.recurrent_groups": "循環組數", + "config.recurrent_depth": "最大循環深度", + "config.min_recurrent_depth": "最小循環深度", + "config.adaptive_halting": "自適應停機 (ACT)", + "config.halting_threshold": "停機閾值", + "config.max_inference_depth": "推理最大深度", + "config.ponder_loss_coef": "Ponder 損失係數", + "config.recurrent_distill_coef": "循環狀態一致性係數", + "config.recurrent_stability_coef": "循環穩定性係數", + "config.recurrent_latent_refinement": "循環中間步僅做潛變量精煉", + "config.depth_schedule": "深度排程", + "config.depth_schedule_warmup_fraction": "自適應預熱比例", + "config.layer_aware_cache": "層/組感知快取", + "config.prefetch_depth": "預取深度", + "config.training_paging": "訓練期分頁", + "config.optimizer_paging": "最佳化器分頁", + "config.page_size_mb": "分頁大小 (MB)", + "config.optimizer_page_dir": "最佳化器頁目錄", + "config.modalities": "模態 (自動/影像/音訊)", "config.hidden_size": "隱藏層維度", "config.num_layers": "層數", "config.num_experts": "專家數", @@ -315,12 +365,12 @@ "autotune.clear_log": "清空日誌檔案", "autotune.apply_best": "套用最優參數 → 配置", "autotune.status": "狀態", - "data.title": "📊 資料檢查器", - "data.path": "資料集路徑 (.jsonl)", + "data.title": "📊 資料檢查器（任意任務）", + "data.path": "資料集路徑（檔案、目錄或萬用字元；jsonl / json / parquet / csv / tsv / txt）", "data.scan": "▶ 掃描", "data.summary": "資料集概要", "data.preview": "樣本預覽", - "data.length_hist": "Token 長度分布", + "data.length_hist": "每筆記錄的字元長度分布", }, "en": { "app.title": "Project Chronos — On-Device MoE Inference Workbench", @@ -337,11 +387,15 @@ "tab.export": "📦 Export", "designer.title": "Live Parameter & Memory Estimator", "designer.total": "Total Parameters", - "designer.active": "Active Parameters (per token)", - "designer.vram": "VRAM Estimate", - "designer.ssd": "SSD Estimate", + "designer.active": 'Active Unique Parameters (range/token)', + "designer.compute": 'Weight Uses (including recurrence)', + "designer.vram": 'Inference Weights + KV', + "designer.ram": 'Example Warm Cache Budget', + "designer.ssd": 'Full Expert SSD Storage', "designer.kv": "KV Cache Estimate", - "designer.tps": "Estimated tokens/sec", + "designer.tps": 'Bandwidth Scenario tokens/s', + "designer.details": 'Parameter & Memory Breakdown', + "designer.notes": 'Active range depends on expert reuse across loops; it is not resident cache size. Weight uses can exceed total unique parameters. Memory assumes batch=1 and the configured context; auto uses FP32. Training weights, gradients and Adam states are listed separately; activations, scratch buffers and allocator peaks are excluded. Do not add RAM and VRAM twice on unified memory. Speed assumes 85% hits, 3 GB/s SSD and 600 GB/s memory bandwidth; measure the actual runtime.', "iomon.title": "Live IO & Cache Metrics", "iomon.refresh": "Refresh", "pipeline.title": "Pretrain → SFT → DPO → ORPO → GRPO → Distill", @@ -350,6 +404,27 @@ "config.loss": "Loss Coefficients", "config.hw": "Hardware Budget", "config.train": "Training Parameters", + "config.recurrent_section": "Recurrent Depth / Paging / Multimodal", + "config.recurrent_enabled": "Enable recurrent depth", + "config.recurrent_groups": "Recurrent groups", + "config.recurrent_depth": "Maximum recurrent depth", + "config.min_recurrent_depth": "Minimum recurrent depth", + "config.adaptive_halting": "Adaptive halting (ACT)", + "config.halting_threshold": "Halting threshold", + "config.max_inference_depth": "Maximum inference depth", + "config.ponder_loss_coef": "Ponder loss coefficient", + "config.recurrent_distill_coef": "Recurrent state consistency coefficient", + "config.recurrent_stability_coef": "Recurrent stability coefficient", + "config.recurrent_latent_refinement": "Latent-only middle loop steps", + "config.depth_schedule": "Depth schedule", + "config.depth_schedule_warmup_fraction": "Adaptive warmup fraction", + "config.layer_aware_cache": "Layer/group-aware cache", + "config.prefetch_depth": "Prefetch depth", + "config.training_paging": "Training paging", + "config.optimizer_paging": "Optimizer paging", + "config.page_size_mb": "Page size (MB)", + "config.optimizer_page_dir": "Optimizer page directory", + "config.modalities": "Modalities (auto/image/audio)", "config.hidden_size": "Hidden Size", "config.num_layers": "Num Layers", "config.num_experts": "Num Experts", @@ -467,12 +542,12 @@ "autotune.clear_log": "Clear log file", "autotune.apply_best": "Apply Best → Config", "autotune.status": "Status", - "data.title": "📊 Data Inspector", - "data.path": "Dataset path (.jsonl)", + "data.title": "📊 Data Inspector (any task)", + "data.path": "Dataset path (file, directory or glob; jsonl / json / parquet / csv / tsv / txt)", "data.scan": "▶ Scan", "data.summary": "Dataset summary", "data.preview": "Sample preview", - "data.length_hist": "Token length histogram", + "data.length_hist": "Characters per record", }, "ja": { "app.title": "Project Chronos — オンデバイス MoE 推論ワークベンチ", @@ -489,11 +564,15 @@ "tab.export": "📦 エクスポート", "designer.title": "リアルタイム パラメータ・メモリ推定", "designer.total": "総パラメータ数", - "designer.active": "アクティブ パラメータ (トークン毎)", - "designer.vram": "VRAM 推定", - "designer.ssd": "SSD 推定", + "designer.active": '有効な固有パラメータ範囲 / token', + "designer.compute": '重み使用回数（反復込み）', + "designer.vram": '推論の常駐重み + KV', + "designer.ram": 'warm キャッシュ予算例', + "designer.ssd": '全エキスパート SSD 容量', "designer.kv": "KV キャッシュ推定", - "designer.tps": "推定 tokens/sec", + "designer.tps": '帯域モデル tokens/s', + "designer.details": 'パラメータとメモリ内訳', + "designer.notes": '有効範囲はループ間のエキスパート再利用に依存します。メモリは batch=1、指定長、auto は FP32 の仮定です。訓練状態は別計上し、活性値や一時領域のピークを含みません。統合メモリでは RAM/VRAM を重複加算しないでください。速度はヒット率85%、SSD 3 GB/s、メモリ600 GB/sの例です。', "iomon.title": "リアルタイム IO・キャッシュ指標", "iomon.refresh": "更新", "pipeline.title": "Pretrain → SFT → DPO → ORPO → GRPO → Distill", @@ -502,6 +581,27 @@ "config.loss": "損失係数", "config.hw": "ハードウェア予算", "config.train": "訓練パラメータ", + "config.recurrent_section": "循環深度 / ページング / マルチモーダル", + "config.recurrent_enabled": "循環深度を有効化", + "config.recurrent_groups": "循環グループ数", + "config.recurrent_depth": "最大循環深度", + "config.min_recurrent_depth": "最小循環深度", + "config.adaptive_halting": "適応的停止 (ACT)", + "config.halting_threshold": "停止しきい値", + "config.max_inference_depth": "推論最大深度", + "config.ponder_loss_coef": "Ponder 損失係数", + "config.recurrent_distill_coef": "循環状態整合性係数", + "config.recurrent_stability_coef": "循環安定性係数", + "config.recurrent_latent_refinement": "中間ループを潜在表現のみ精緻化", + "config.depth_schedule": "深度スケジュール", + "config.depth_schedule_warmup_fraction": "適応ウォームアップ比率", + "config.layer_aware_cache": "層/グループ対応キャッシュ", + "config.prefetch_depth": "先読み深度", + "config.training_paging": "訓練時ページング", + "config.optimizer_paging": "オプティマイザページング", + "config.page_size_mb": "ページサイズ (MB)", + "config.optimizer_page_dir": "オプティマイザページディレクトリ", + "config.modalities": "モダリティ (自動/画像/音声)", "config.hidden_size": "隠れ層次元", "config.num_layers": "層数", "config.num_experts": "エキスパート数", @@ -619,12 +719,12 @@ "autotune.clear_log": "ログファイル消去", "autotune.apply_best": "最適パラメータを設定に適用", "autotune.status": "ステータス", - "data.title": "📊 データ検査ツール", - "data.path": "データセットパス (.jsonl)", + "data.title": "📊 データ検査ツール（任意のタスク）", + "data.path": "データセットパス（ファイル・ディレクトリ・ワイルドカード；jsonl / json / parquet / csv / tsv / txt）", "data.scan": "▶ スキャン", "data.summary": "データセット概要", "data.preview": "サンプルプレビュー", - "data.length_hist": "トークン長分布", + "data.length_hist": "レコードあたりの文字数分布", }, } @@ -652,8 +752,8 @@ def build_language_update(lang: str): updates = [] for comp, key in _translatable_components: text = t(key, lang) - # gr.Button uses `value`; gr.Tab and most others use `label` - if isinstance(comp, gr.Button): + # gr.Button / gr.Markdown use `value`; gr.Tab and most others use `label` + if isinstance(comp, (gr.Button, gr.Markdown)): updates.append(gr.update(value=text)) else: updates.append(gr.update(label=text)) @@ -662,3 +762,379 @@ def build_language_update(lang: str): def get_registered_components(): return [c for c, _ in _translatable_components] + + +# ---- training monitor tab (ui/tabs/monitor_tab.py) ------------------------------- +_MONITOR_TRANSLATIONS: Dict[str, Dict[str, str]] = { + "zh-Hans": { + "tab.monitor": "📈 训练监控", + "monitor.intro": "只读监控 `runs/` 下的训练运行：配置、状态、进度、全部指标、验证与评测指标（自动发现）、张量统计与路由。与独立页面 `chronos monitor` 读取同一份数据。", + "monitor.runs_dir": "运行目录", "monitor.run": "运行", "monitor.phase": "阶段", "monitor.auto": "每 15 秒自动刷新", + "monitor.refresh": "刷新", "monitor.step": "步", "monitor.eta": "预计完成", "monitor.remaining": "剩余", "monitor.budget": "预算", + "monitor.source": "状态来源", "monitor.none": "暂无数据", "monitor.no_runs": "没有找到运行目录：", + "monitor.error": "运行报告错误", "monitor.paused": "训练进程处于挂起（SIGSTOP）状态。", + "monitor.stale": "最近 {age} 没有新的日志行，训练进程可能卡住。", + "monitor.derived": "此运行由旧版代码启动：状态由指标流和进程表推断；张量统计、路由热图和评测探针在以新代码恢复后出现。", + "monitor.no_telemetry": "暂无张量遥测（parameters.jsonl）。", + "monitor.telemetry_at": "张量遥测：第 {step} 步，统计耗时 {ms} ms；梯度为裁剪前、按微批平均。", + "monitor.state.running": "训练中", "monitor.state.paused": "已暂停", "monitor.state.evaluating": "评测中", + "monitor.state.checkpointing": "保存检查点", "monitor.state.starting": "启动中", "monitor.state.finished": "已完成", + "monitor.state.failed": "失败", "monitor.state.stopped": "已停止", "monitor.state.pending": "等待中", + "monitor.ph.complete": "完成", "monitor.ph.active": "进行中", "monitor.ph.pending": "未开始", + "monitor.k_loss": "训练损失", "monitor.k_ce": "主损失 (CE)", "monitor.k_valid": "最佳验证 CE", "monitor.k_tps": "吞吐", + "monitor.k_mem": "内存", "monitor.k_grad": "梯度范数", "monitor.k_lr": "学习率", "monitor.k_depth": "平均循环深度", + "monitor.k_route": "路由熵", + "monitor.t_overview": "总览", "monitor.t_recurrent": "循环与路由", "monitor.t_loss": "损失分量", "monitor.t_params": "参数统计", + "monitor.t_evals": "评测与样例", "monitor.t_config": "训练参数", "monitor.t_log": "日志与检查点", + "monitor.p_ce": "主损失 (CE)（蓝 = 训练，橙 = 验证）", "monitor.p_loss": "总损失", "monitor.p_lr": "学习率", "monitor.p_grad": "梯度范数（裁剪前）", + "monitor.p_tps": "吞吐 tok/s（蓝 = 有效 token，橙 = 含填充）", "monitor.p_mem": "MLX 内存 GB（蓝 = 峰值，橙 = 活跃，绿 = 缓存）", "monitor.p_depth": "平均循环深度", "monitor.p_contraction": "收缩系数", + "monitor.p_exit": "提前退出比例", "monitor.p_route_entropy": "路由熵", "monitor.p_load_cv": "专家负载变异系数", + "monitor.p_max_load": "最大专家负载（均值倍数）", + "monitor.s_train": "训练", "monitor.s_valid": "验证", "monitor.s_total": "总损失", "monitor.s_lr": "学习率", "monitor.s_grad": "梯度范数", + "monitor.s_counted": "有效 token", "monitor.s_processed": "含填充", "monitor.s_peak": "峰值", "monitor.s_active": "活跃", + "monitor.s_cache": "缓存", "monitor.s_depth": "深度", "monitor.s_contraction": "收缩", "monitor.s_exit": "提前退出", + "monitor.s_entropy": "熵", "monitor.s_load_cv": "CV", "monitor.s_max_load": "最大负载", + "monitor.loss_term": "损失分量（加权值；primary_loss 为 CE）", + "monitor.groups": "参数组", "monitor.tensors": "全部张量", "monitor.h_load": "专家负载（行 = 层·迭代）", + "monitor.h_importance": "平均路由概率", "monitor.h_wrms": "专家权重 RMS（行 = 物理层）", "monitor.h_grms": "专家梯度 RMS", + "monitor.events": "事件（验证 / 回滚 / 阶段 / 恢复 / 最终评测）", "monitor.training_params": "训练参数（配方、模型、运行时、数据、分词器）", + "monitor.log": "train.log（末尾）", "monitor.checkpoints": "检查点", + "monitor.c_group": "组", "monitor.c_tensors": "张量数", "monitor.c_numel": "参数量", "monitor.c_rms": "RMS", + "monitor.c_grad_norm": "梯度范数", "monitor.c_update": "更新比 ‖Δ‖/‖W‖", "monitor.c_name": "名称", "monitor.c_shape": "形状", + "monitor.c_dtype": "类型", "monitor.c_std": "标准差", "monitor.c_absmax": "|最大|", "monitor.c_grad_rms": "梯度 RMS", + "monitor.c_time": "时间", "monitor.c_event": "事件", "monitor.c_phase": "阶段", "monitor.c_valid_ce": "验证 CE", + "monitor.c_detail": "详情", "monitor.c_file": "文件", "monitor.c_size": "大小 (GB)", + }, + "en": { + "tab.monitor": "📈 Training Monitor", + "monitor.intro": "Read-only view of training runs under `runs/`: configuration, status, progress, every metric, validation and evaluation metrics (discovered from the logs), tensor statistics and routing. Same data as the standalone `chronos monitor` page.", + "monitor.runs_dir": "Runs directory", "monitor.run": "Run", "monitor.phase": "Phase", "monitor.auto": "Auto refresh every 15 s", + "monitor.refresh": "Refresh", "monitor.step": "step", "monitor.eta": "ETA", "monitor.remaining": "remaining", "monitor.budget": "budget", + "monitor.source": "Status source", "monitor.none": "No data yet", "monitor.no_runs": "No run directories found in", + "monitor.error": "The run reported an error", "monitor.paused": "The trainer is suspended (SIGSTOP).", + "monitor.stale": "No new log line for {age}; the trainer may be stuck.", + "monitor.derived": "Started by older code: status is derived from the metric stream and the process table; tensor stats, routing heatmaps and evaluation probes appear after a resume with the new code.", + "monitor.no_telemetry": "No tensor telemetry (parameters.jsonl) yet.", + "monitor.telemetry_at": "Tensor telemetry at step {step}, stats took {ms} ms; gradients are pre-clip, mean over micro-batches.", + "monitor.state.running": "Running", "monitor.state.paused": "Paused", "monitor.state.evaluating": "Evaluating", + "monitor.state.checkpointing": "Checkpointing", "monitor.state.starting": "Starting", "monitor.state.finished": "Finished", + "monitor.state.failed": "Failed", "monitor.state.stopped": "Stopped", "monitor.state.pending": "Pending", + "monitor.ph.complete": "complete", "monitor.ph.active": "active", "monitor.ph.pending": "pending", + "monitor.k_loss": "Training loss", "monitor.k_ce": "Primary loss (CE)", "monitor.k_valid": "Best valid CE", "monitor.k_tps": "Throughput", + "monitor.k_mem": "Memory", "monitor.k_grad": "Grad norm", "monitor.k_lr": "Learning rate", "monitor.k_depth": "Mean recurrent depth", + "monitor.k_route": "Routing entropy", + "monitor.t_overview": "Overview", "monitor.t_recurrent": "Recurrence & routing", "monitor.t_loss": "Loss terms", "monitor.t_params": "Parameters", + "monitor.t_evals": "Evals & samples", "monitor.t_config": "Training parameters", "monitor.t_log": "Log & checkpoints", + "monitor.p_ce": "Primary loss (CE) (blue = train, orange = valid)", "monitor.p_loss": "Total loss", "monitor.p_lr": "Learning rate", "monitor.p_grad": "Grad norm (pre-clip)", + "monitor.p_tps": "Throughput tok/s (blue = counted, orange = incl. padding)", "monitor.p_mem": "MLX memory GB (blue = peak, orange = active, green = cache)", "monitor.p_depth": "Mean recurrent depth", "monitor.p_contraction": "Contraction", + "monitor.p_exit": "Early-exit share", "monitor.p_route_entropy": "Routing entropy", "monitor.p_load_cv": "Expert load CV", + "monitor.p_max_load": "Max expert load (× mean)", + "monitor.s_train": "train", "monitor.s_valid": "valid", "monitor.s_total": "total", "monitor.s_lr": "lr", "monitor.s_grad": "grad norm", + "monitor.s_counted": "counted tokens", "monitor.s_processed": "incl. padding", "monitor.s_peak": "peak", "monitor.s_active": "active", + "monitor.s_cache": "cache", "monitor.s_depth": "depth", "monitor.s_contraction": "contraction", "monitor.s_exit": "early exit", + "monitor.s_entropy": "entropy", "monitor.s_load_cv": "CV", "monitor.s_max_load": "max load", + "monitor.loss_term": "Loss term (weighted; primary_loss is CE)", + "monitor.groups": "Parameter groups", "monitor.tensors": "All tensors", "monitor.h_load": "Expert load (rows = layer·iteration)", + "monitor.h_importance": "Mean router probability", "monitor.h_wrms": "Expert weight RMS (rows = physical layer)", "monitor.h_grms": "Expert gradient RMS", + "monitor.events": "Events (validation / rollback / phase / resume / final)", "monitor.training_params": "Training parameters (recipe, model, runtime, data, tokenizer)", + "monitor.log": "train.log (tail)", "monitor.checkpoints": "Checkpoints", + "monitor.c_group": "Group", "monitor.c_tensors": "Tensors", "monitor.c_numel": "Params", "monitor.c_rms": "RMS", + "monitor.c_grad_norm": "Grad norm", "monitor.c_update": "Update ratio ‖Δ‖/‖W‖", "monitor.c_name": "Name", "monitor.c_shape": "Shape", + "monitor.c_dtype": "dtype", "monitor.c_std": "Std", "monitor.c_absmax": "|max|", "monitor.c_grad_rms": "Grad RMS", + "monitor.c_time": "Time", "monitor.c_event": "Event", "monitor.c_phase": "Phase", "monitor.c_valid_ce": "Valid CE", + "monitor.c_detail": "Detail", "monitor.c_file": "File", "monitor.c_size": "Size (GB)", + }, + "zh-Hant": { + "tab.monitor": "📈 訓練監控", + "monitor.intro": "唯讀監控 `runs/` 下的訓練執行：設定、狀態、進度、全部指標、驗證與評測指標（自動發現）、張量統計與路由。與獨立頁面 `chronos monitor` 讀取同一份資料。", + "monitor.runs_dir": "執行目錄", "monitor.run": "執行", "monitor.phase": "階段", "monitor.auto": "每 15 秒自動重新整理", + "monitor.refresh": "重新整理", "monitor.step": "步", "monitor.eta": "預計完成", "monitor.remaining": "剩餘", "monitor.budget": "預算", + "monitor.source": "狀態來源", "monitor.none": "暫無資料", "monitor.no_runs": "找不到執行目錄：", + "monitor.error": "執行回報錯誤", "monitor.paused": "訓練行程處於暫停（SIGSTOP）狀態。", + "monitor.stale": "最近 {age} 沒有新的日誌行，訓練行程可能卡住。", + "monitor.derived": "此執行由舊版程式碼啟動：狀態由指標流與行程表推斷；張量統計、路由熱圖與評測探針在以新程式碼恢復後出現。", + "monitor.no_telemetry": "暫無張量遙測（parameters.jsonl）。", + "monitor.telemetry_at": "張量遙測：第 {step} 步，統計耗時 {ms} ms；梯度為裁剪前、按微批平均。", + "monitor.state.running": "訓練中", "monitor.state.paused": "已暫停", "monitor.state.evaluating": "評測中", + "monitor.state.checkpointing": "儲存檢查點", "monitor.state.starting": "啟動中", "monitor.state.finished": "已完成", + "monitor.state.failed": "失敗", "monitor.state.stopped": "已停止", "monitor.state.pending": "等待中", + "monitor.ph.complete": "完成", "monitor.ph.active": "進行中", "monitor.ph.pending": "未開始", + "monitor.k_loss": "訓練損失", "monitor.k_ce": "主損失 (CE)", "monitor.k_valid": "最佳驗證 CE", "monitor.k_tps": "吞吐", + "monitor.k_mem": "記憶體", "monitor.k_grad": "梯度範數", "monitor.k_lr": "學習率", "monitor.k_depth": "平均循環深度", + "monitor.k_route": "路由熵", + "monitor.t_overview": "總覽", "monitor.t_recurrent": "循環與路由", "monitor.t_loss": "損失分量", "monitor.t_params": "參數統計", + "monitor.t_evals": "評測與樣例", "monitor.t_config": "訓練參數", "monitor.t_log": "日誌與檢查點", + "monitor.p_ce": "主損失 (CE)（藍 = 訓練，橙 = 驗證）", "monitor.p_loss": "總損失", "monitor.p_lr": "學習率", "monitor.p_grad": "梯度範數（裁剪前）", + "monitor.p_tps": "吞吐 tok/s（藍 = 有效 token，橙 = 含填充）", "monitor.p_mem": "MLX 記憶體 GB（藍 = 峰值，橙 = 活躍，綠 = 快取）", "monitor.p_depth": "平均循環深度", "monitor.p_contraction": "收縮係數", + "monitor.p_exit": "提前退出比例", "monitor.p_route_entropy": "路由熵", "monitor.p_load_cv": "專家負載變異係數", + "monitor.p_max_load": "最大專家負載（均值倍數）", + "monitor.s_train": "訓練", "monitor.s_valid": "驗證", "monitor.s_total": "總損失", "monitor.s_lr": "學習率", "monitor.s_grad": "梯度範數", + "monitor.s_counted": "有效 token", "monitor.s_processed": "含填充", "monitor.s_peak": "峰值", "monitor.s_active": "活躍", + "monitor.s_cache": "快取", "monitor.s_depth": "深度", "monitor.s_contraction": "收縮", "monitor.s_exit": "提前退出", + "monitor.s_entropy": "熵", "monitor.s_load_cv": "CV", "monitor.s_max_load": "最大負載", + "monitor.loss_term": "損失分量（加權值；primary_loss 為 CE）", + "monitor.groups": "參數組", "monitor.tensors": "全部張量", "monitor.h_load": "專家負載（列 = 層·迭代）", + "monitor.h_importance": "平均路由機率", "monitor.h_wrms": "專家權重 RMS（列 = 物理層）", "monitor.h_grms": "專家梯度 RMS", + "monitor.events": "事件（驗證 / 回滾 / 階段 / 恢復 / 最終評測）", "monitor.training_params": "訓練參數（配方、模型、執行環境、資料、分詞器）", + "monitor.log": "train.log（末尾）", "monitor.checkpoints": "檢查點", + "monitor.c_group": "組", "monitor.c_tensors": "張量數", "monitor.c_numel": "參數量", "monitor.c_rms": "RMS", + "monitor.c_grad_norm": "梯度範數", "monitor.c_update": "更新比 ‖Δ‖/‖W‖", "monitor.c_name": "名稱", "monitor.c_shape": "形狀", + "monitor.c_dtype": "型別", "monitor.c_std": "標準差", "monitor.c_absmax": "|最大|", "monitor.c_grad_rms": "梯度 RMS", + "monitor.c_time": "時間", "monitor.c_event": "事件", "monitor.c_phase": "階段", "monitor.c_valid_ce": "驗證 CE", + "monitor.c_detail": "詳情", "monitor.c_file": "檔案", "monitor.c_size": "大小 (GB)", + }, + "ja": { + "tab.monitor": "📈 学習モニタ", + "monitor.intro": "`runs/` 配下の学習ランを読み取り専用で表示します：設定・状態・進捗・全指標・検証と評価指標（ログから自動検出）・テンソル統計とルーティング。単独ページ `chronos monitor` と同じデータです。", + "monitor.runs_dir": "ランディレクトリ", "monitor.run": "ラン", "monitor.phase": "フェーズ", "monitor.auto": "15 秒ごとに自動更新", + "monitor.refresh": "更新", "monitor.step": "ステップ", "monitor.eta": "完了予定", "monitor.remaining": "残り", "monitor.budget": "予算", + "monitor.source": "状態の取得元", "monitor.none": "データなし", "monitor.no_runs": "ランディレクトリが見つかりません：", + "monitor.error": "ランがエラーを報告しました", "monitor.paused": "学習プロセスは一時停止（SIGSTOP）中です。", + "monitor.stale": "{age} の間ログ行が更新されていません。学習プロセスが停止している可能性があります。", + "monitor.derived": "旧バージョンのコードで開始されたランです：状態は指標ストリームとプロセス表から推定しています。テンソル統計・ルーティングのヒートマップ・評価プローブは新コードで再開すると表示されます。", + "monitor.no_telemetry": "テンソルテレメトリ（parameters.jsonl）はまだありません。", + "monitor.telemetry_at": "テンソルテレメトリ：ステップ {step}、集計 {ms} ms。勾配はクリップ前・マイクロバッチ平均です。", + "monitor.state.running": "学習中", "monitor.state.paused": "一時停止", "monitor.state.evaluating": "評価中", + "monitor.state.checkpointing": "チェックポイント保存中", "monitor.state.starting": "起動中", "monitor.state.finished": "完了", + "monitor.state.failed": "失敗", "monitor.state.stopped": "停止", "monitor.state.pending": "待機中", + "monitor.ph.complete": "完了", "monitor.ph.active": "実行中", "monitor.ph.pending": "未開始", + "monitor.k_loss": "学習損失", "monitor.k_ce": "主損失 (CE)", "monitor.k_valid": "最良検証 CE", "monitor.k_tps": "スループット", + "monitor.k_mem": "メモリ", "monitor.k_grad": "勾配ノルム", "monitor.k_lr": "学習率", "monitor.k_depth": "平均再帰深さ", + "monitor.k_route": "ルーティングエントロピー", + "monitor.t_overview": "概要", "monitor.t_recurrent": "再帰とルーティング", "monitor.t_loss": "損失成分", "monitor.t_params": "パラメータ統計", + "monitor.t_evals": "評価とサンプル", "monitor.t_config": "学習パラメータ", "monitor.t_log": "ログとチェックポイント", + "monitor.p_ce": "主損失 (CE)（青 = 学習、橙 = 検証）", "monitor.p_loss": "総損失", "monitor.p_lr": "学習率", "monitor.p_grad": "勾配ノルム（クリップ前）", + "monitor.p_tps": "スループット tok/s（青 = 有効トークン、橙 = パディング込み）", "monitor.p_mem": "MLX メモリ GB（青 = ピーク、橙 = アクティブ、緑 = キャッシュ）", "monitor.p_depth": "平均再帰深さ", "monitor.p_contraction": "収縮係数", + "monitor.p_exit": "早期終了率", "monitor.p_route_entropy": "ルーティングエントロピー", "monitor.p_load_cv": "エキスパート負荷の変動係数", + "monitor.p_max_load": "最大エキスパート負荷（平均比）", + "monitor.s_train": "学習", "monitor.s_valid": "検証", "monitor.s_total": "総損失", "monitor.s_lr": "学習率", "monitor.s_grad": "勾配ノルム", + "monitor.s_counted": "有効トークン", "monitor.s_processed": "パディング込み", "monitor.s_peak": "ピーク", "monitor.s_active": "アクティブ", + "monitor.s_cache": "キャッシュ", "monitor.s_depth": "深さ", "monitor.s_contraction": "収縮", "monitor.s_exit": "早期終了", + "monitor.s_entropy": "エントロピー", "monitor.s_load_cv": "CV", "monitor.s_max_load": "最大負荷", + "monitor.loss_term": "損失成分（重み付き値；primary_loss は CE）", + "monitor.groups": "パラメータグループ", "monitor.tensors": "全テンソル", "monitor.h_load": "エキスパート負荷（行 = 層·反復）", + "monitor.h_importance": "平均ルーティング確率", "monitor.h_wrms": "エキスパート重み RMS（行 = 物理層）", "monitor.h_grms": "エキスパート勾配 RMS", + "monitor.events": "イベント（検証 / ロールバック / フェーズ / 再開 / 最終評価）", "monitor.training_params": "学習パラメータ（レシピ・モデル・実行環境・データ・トークナイザ）", + "monitor.log": "train.log（末尾）", "monitor.checkpoints": "チェックポイント", + "monitor.c_group": "グループ", "monitor.c_tensors": "テンソル数", "monitor.c_numel": "パラメータ数", "monitor.c_rms": "RMS", + "monitor.c_grad_norm": "勾配ノルム", "monitor.c_update": "更新比 ‖Δ‖/‖W‖", "monitor.c_name": "名前", "monitor.c_shape": "形状", + "monitor.c_dtype": "型", "monitor.c_std": "標準偏差", "monitor.c_absmax": "|最大|", "monitor.c_grad_rms": "勾配 RMS", + "monitor.c_time": "時刻", "monitor.c_event": "イベント", "monitor.c_phase": "フェーズ", "monitor.c_valid_ce": "検証 CE", + "monitor.c_detail": "詳細", "monitor.c_file": "ファイル", "monitor.c_size": "サイズ (GB)", + }, +} +_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 数を表示します。'}} +for _lang, _entries in _INFER_SPECULATIVE.items(): + TRANSLATIONS[_lang].update(_entries) +for _lang, _entries in _MONITOR_GENERIC_EVALS.items(): + _MONITOR_TRANSLATIONS[_lang].update(_entries) +for _lang, _entries in _MONITOR_TRANSLATIONS.items(): + TRANSLATIONS[_lang].update(_entries) + + +# ---- tokenizer builder + serve tabs (ui/tabs/tokenizer_tab.py, serve_tab.py) ------ +_TOOLS_TRANSLATIONS: Dict[str, Dict[str, str]] = { + "zh-Hans": { + "tab.tokenizer": "🔤 分词器", + "tok.intro": "从所选数据集（jsonl / parquet / csv / tsv / txt；文件、目录或通配符）自动抽样，训练 Chronos 自研字节回退 BPE 分词器，并显示留出集评估。与 `chronos tokenizer` 使用同一个构建器。", + "tok.paths": "数据路径（每行一个，或逗号分隔）", "tok.out_dir": "输出目录", + "tok.fields": "字段映射（可选，逗号分隔的点路径，如 text、messages、translation.en）", "tok.vocab_size": "词表大小", + "tok.per_source": "每个来源抽样文本数", "tok.temperature": "语言温度采样 T（越小越均衡）", + "tok.build": "🔨 训练分词器", "tok.load_report": "📂 载入已有报告", "tok.status": "状态", + "tok.log": "日志", "tok.summary": "摘要", "tok.eval": "留出集评估", + "tok.state.idle": "空闲", "tok.state.running": "训练中…", "tok.state.done": "完成", + "tok.state.failed": "失败", "tok.state.loaded": "已载入报告", + "tok.err.busy": "已有分词器任务在运行", "tok.err.no_paths": "请填写数据路径", + "tok.err.no_out": "请填写输出目录", "tok.err.no_files": "路径中没有找到可读的数据文件", + "tok.err.no_report": "输出目录中没有 tokenizer_report.json", + "tab.serve": "🛰️ 服务", + "serve.intro": "以子进程启动 `chronos serve`（Ollama / OpenAI 协议）。stock `ollama` CLI 可通过 `OLLAMA_HOST` 连接；懒加载与专家内存上限在协议调用下同样生效。", + "serve.models_dir": "模型目录", "serve.host": "监听地址", "serve.port": "端口", "serve.device": "设备", + "serve.dtype": "计算精度", "serve.lazy": "专家懒加载", "serve.ram_fraction": "RAM 专家比例", + "serve.kv_dtype": "KV cache 精度", "serve.mla_decode": "MLA 解码方式", "serve.router_scope": "路由范围", + "serve.max_resident": "驻留专家硬上限（0=不限）", "serve.ram_budget": "专家 RAM 上限 MB（0=不限）", + "serve.allow_remote": "允许非本机访问", "serve.api_key": "API key（非本机访问时必填）", + "serve.monitor_dir": "监控运行目录（可选）", + "serve.warn_noauth": "⚠️ 默认只监听本机且不带认证。监听其他地址时必须勾选“允许非本机访问”并设置 API key；key 通过环境变量传给子进程，不出现在命令行。", + "serve.start": "▶️ 启动", "serve.stop": "⏹️ 停止", "serve.status": "状态", "serve.log": "服务日志", + "serve.usage": "用法：`OLLAMA_HOST=127.0.0.1:11435 ollama run <模型名>`；OpenAI 客户端使用 `http://127.0.0.1:11435/v1`。开启认证后请求头加 `Authorization: Bearer `。", + "serve.state.running": "运行中", "serve.state.stopped": "未运行", + "serve.auth_on": "已启用认证", "serve.auth_off": "未启用认证（仅限本机）", + "serve.err.models_dir": "模型目录不存在", "serve.err.port": "端口必须在 1–65535 之间", + "serve.err.remote_needs_key": "非本机地址需要勾选“允许非本机访问”并填写 API key", + "serve.err.busy": "服务已在运行，请先停止", + }, + "en": { + "tab.tokenizer": "🔤 Tokenizer", + "tok.intro": "Samples the selected datasets (jsonl / parquet / csv / tsv / txt; files, directories or globs), trains the Chronos byte-fallback BPE tokenizer and shows the held-out evaluation. Same builder as `chronos tokenizer`.", + "tok.paths": "Data paths (one per line or comma-separated)", "tok.out_dir": "Output directory", + "tok.fields": "Field mapping (optional, comma-separated dotted paths, e.g. text, messages, translation.en)", "tok.vocab_size": "Vocabulary size", + "tok.per_source": "Texts sampled per source", "tok.temperature": "Language temperature T (lower = more balanced)", + "tok.build": "🔨 Train tokenizer", "tok.load_report": "📂 Load existing report", "tok.status": "Status", + "tok.log": "Log", "tok.summary": "Summary", "tok.eval": "Held-out evaluation", + "tok.state.idle": "Idle", "tok.state.running": "Training…", "tok.state.done": "Done", + "tok.state.failed": "Failed", "tok.state.loaded": "Report loaded", + "tok.err.busy": "A tokenizer job is already running", "tok.err.no_paths": "Enter at least one data path", + "tok.err.no_out": "Enter an output directory", "tok.err.no_files": "No readable data files found at those paths", + "tok.err.no_report": "No tokenizer_report.json in the output directory", + "tab.serve": "🛰️ Serve", + "serve.intro": "Starts `chronos serve` (Ollama / OpenAI protocol) as a child process. The stock `ollama` CLI connects through `OLLAMA_HOST`; lazy experts and the expert memory caps apply to protocol requests too.", + "serve.models_dir": "Models directory", "serve.host": "Host", "serve.port": "Port", "serve.device": "Device", + "serve.dtype": "Compute dtype", "serve.lazy": "Lazy expert loading", "serve.ram_fraction": "RAM expert fraction", + "serve.kv_dtype": "KV cache dtype", "serve.mla_decode": "MLA decode", "serve.router_scope": "Router scope", + "serve.max_resident": "Resident expert hard cap (0 = none)", "serve.ram_budget": "Expert RAM cap MB (0 = none)", + "serve.allow_remote": "Allow remote access", "serve.api_key": "API key (required for remote access)", + "serve.monitor_dir": "Monitor runs directory (optional)", + "serve.warn_noauth": "⚠️ By default the server listens on loopback only, without authentication. Any other host requires *Allow remote access* plus an API key; the key reaches the child process through an environment variable, never the command line.", + "serve.start": "▶️ Start", "serve.stop": "⏹️ Stop", "serve.status": "Status", "serve.log": "Server log", + "serve.usage": "Usage: `OLLAMA_HOST=127.0.0.1:11435 ollama run `; OpenAI clients use `http://127.0.0.1:11435/v1`. With auth on, send `Authorization: Bearer `.", + "serve.state.running": "Running", "serve.state.stopped": "Stopped", + "serve.auth_on": "auth on", "serve.auth_off": "no auth (loopback only)", + "serve.err.models_dir": "Models directory does not exist", "serve.err.port": "Port must be 1–65535", + "serve.err.remote_needs_key": "A non-loopback host needs *Allow remote access* and an API key", + "serve.err.busy": "The server is already running; stop it first", + }, + "zh-Hant": { + "tab.tokenizer": "🔤 分詞器", + "tok.intro": "從所選資料集（jsonl / parquet / csv / tsv / txt；檔案、目錄或萬用字元）自動抽樣，訓練 Chronos 自研位元組回退 BPE 分詞器，並顯示留出集評估。與 `chronos tokenizer` 使用同一個建構器。", + "tok.paths": "資料路徑（每行一個，或逗號分隔）", "tok.out_dir": "輸出目錄", + "tok.fields": "欄位對應（選填，逗號分隔的點路徑，如 text、messages、translation.en）", "tok.vocab_size": "詞表大小", + "tok.per_source": "每個來源抽樣文字數", "tok.temperature": "語言溫度取樣 T（越小越均衡）", + "tok.build": "🔨 訓練分詞器", "tok.load_report": "📂 載入已有報告", "tok.status": "狀態", + "tok.log": "日誌", "tok.summary": "摘要", "tok.eval": "留出集評估", + "tok.state.idle": "閒置", "tok.state.running": "訓練中…", "tok.state.done": "完成", + "tok.state.failed": "失敗", "tok.state.loaded": "已載入報告", + "tok.err.busy": "已有分詞器任務在執行", "tok.err.no_paths": "請填寫資料路徑", + "tok.err.no_out": "請填寫輸出目錄", "tok.err.no_files": "路徑中沒有找到可讀的資料檔", + "tok.err.no_report": "輸出目錄中沒有 tokenizer_report.json", + "tab.serve": "🛰️ 服務", + "serve.intro": "以子行程啟動 `chronos serve`（Ollama / OpenAI 協定）。stock `ollama` CLI 可透過 `OLLAMA_HOST` 連線；懶載入與專家記憶體上限在協定呼叫下同樣生效。", + "serve.models_dir": "模型目錄", "serve.host": "監聽位址", "serve.port": "連接埠", "serve.device": "裝置", + "serve.dtype": "計算精度", "serve.lazy": "專家懶載入", "serve.ram_fraction": "RAM 專家比例", + "serve.kv_dtype": "KV cache 精度", "serve.mla_decode": "MLA 解碼方式", "serve.router_scope": "路由範圍", + "serve.max_resident": "常駐專家硬上限（0=不限）", "serve.ram_budget": "專家 RAM 上限 MB（0=不限）", + "serve.allow_remote": "允許非本機存取", "serve.api_key": "API key（非本機存取時必填）", + "serve.monitor_dir": "監控執行目錄（選填）", + "serve.warn_noauth": "⚠️ 預設只監聽本機且不帶認證。監聽其他位址時必須勾選「允許非本機存取」並設定 API key；key 透過環境變數傳給子行程，不出現在命令列。", + "serve.start": "▶️ 啟動", "serve.stop": "⏹️ 停止", "serve.status": "狀態", "serve.log": "服務日誌", + "serve.usage": "用法：`OLLAMA_HOST=127.0.0.1:11435 ollama run <模型名>`；OpenAI 用戶端使用 `http://127.0.0.1:11435/v1`。啟用認證後請求標頭加 `Authorization: Bearer `。", + "serve.state.running": "執行中", "serve.state.stopped": "未執行", + "serve.auth_on": "已啟用認證", "serve.auth_off": "未啟用認證（僅限本機）", + "serve.err.models_dir": "模型目錄不存在", "serve.err.port": "連接埠必須在 1–65535 之間", + "serve.err.remote_needs_key": "非本機位址需要勾選「允許非本機存取」並填寫 API key", + "serve.err.busy": "服務已在執行，請先停止", + }, + "ja": { + "tab.tokenizer": "🔤 トークナイザー", + "tok.intro": "選択したデータセット（jsonl / parquet / csv / tsv / txt；ファイル・ディレクトリ・glob）から自動サンプリングし、Chronos 独自のバイトフォールバック BPE トークナイザーを学習して、ホールドアウト評価を表示します。`chronos tokenizer` と同じビルダーです。", + "tok.paths": "データパス（1 行に 1 つ、またはカンマ区切り）", "tok.out_dir": "出力ディレクトリ", + "tok.fields": "フィールド対応（任意、カンマ区切りのドットパス、例 text・messages・translation.en）", "tok.vocab_size": "語彙サイズ", + "tok.per_source": "ソースごとのサンプル数", "tok.temperature": "言語温度 T（小さいほど均等）", + "tok.build": "🔨 トークナイザーを学習", "tok.load_report": "📂 既存レポートを読込", "tok.status": "状態", + "tok.log": "ログ", "tok.summary": "概要", "tok.eval": "ホールドアウト評価", + "tok.state.idle": "待機中", "tok.state.running": "学習中…", "tok.state.done": "完了", + "tok.state.failed": "失敗", "tok.state.loaded": "レポート読込済み", + "tok.err.busy": "トークナイザーのジョブが実行中です", "tok.err.no_paths": "データパスを入力してください", + "tok.err.no_out": "出力ディレクトリを入力してください", "tok.err.no_files": "指定パスに読めるデータファイルがありません", + "tok.err.no_report": "出力ディレクトリに tokenizer_report.json がありません", + "tab.serve": "🛰️ サーバー", + "serve.intro": "`chronos serve`（Ollama / OpenAI プロトコル）を子プロセスとして起動します。標準の `ollama` CLI は `OLLAMA_HOST` で接続でき、遅延ロードと専門家メモリ上限はプロトコル経由でも有効です。", + "serve.models_dir": "モデルディレクトリ", "serve.host": "ホスト", "serve.port": "ポート", "serve.device": "デバイス", + "serve.dtype": "計算精度", "serve.lazy": "専門家の遅延ロード", "serve.ram_fraction": "RAM 専門家比率", + "serve.kv_dtype": "KV cache 精度", "serve.mla_decode": "MLA デコード", "serve.router_scope": "ルーター範囲", + "serve.max_resident": "常駐専門家の上限（0=なし）", "serve.ram_budget": "専門家 RAM 上限 MB（0=なし）", + "serve.allow_remote": "リモートアクセスを許可", "serve.api_key": "API key（リモート時は必須）", + "serve.monitor_dir": "モニター実行ディレクトリ（任意）", + "serve.warn_noauth": "⚠️ 既定ではループバックのみで待ち受け、認証はありません。他のホストでは「リモートアクセスを許可」と API key が必要です。key は環境変数で子プロセスに渡り、コマンドラインには出ません。", + "serve.start": "▶️ 起動", "serve.stop": "⏹️ 停止", "serve.status": "状態", "serve.log": "サーバーログ", + "serve.usage": "使い方：`OLLAMA_HOST=127.0.0.1:11435 ollama run <モデル名>`。OpenAI クライアントは `http://127.0.0.1:11435/v1`。認証有効時は `Authorization: Bearer ` を送信。", + "serve.state.running": "実行中", "serve.state.stopped": "停止中", + "serve.auth_on": "認証あり", "serve.auth_off": "認証なし（ループバックのみ）", + "serve.err.models_dir": "モデルディレクトリが存在しません", "serve.err.port": "ポートは 1–65535 の範囲です", + "serve.err.remote_needs_key": "ループバック以外では「リモートアクセスを許可」と API key が必要です", + "serve.err.busy": "サーバーは既に実行中です。先に停止してください", + }, +} +for _lang, _entries in _TOOLS_TRANSLATIONS.items(): + TRANSLATIONS[_lang].update(_entries) + + +# ---- Config tab: adaptive training planner (ui/planner_panel.py) ------------------- +_PLANNER_TRANSLATIONS: Dict[str, Dict[str, str]] = { + "zh-Hans": { + "plan.title": "🧭 自动规划训练参数", + "plan.intro": "填写目标后生成 v3 架构和各阶段独立的训练参数。学习率与 batch 按 DeepSeek LLM 缩放律（arXiv 2401.02954）估算，内存按实测标定的模型估算；数据集路径用于推断 token 数、序列长度、记录数和分词器。所有数值都标注来源。", + "plan.target": "目标总参数量（M）", "plan.use_moe": "启用 MoE", "plan.experts": "专家数", + "plan.depth": "循环深度", "plan.active_mode": "激活规模的指定方式", + "plan.mode.experts": "激活专家数（top-k）", "plan.mode.params": "激活参数量（M）", + "plan.active_value": "激活值（专家数或 M 参数；0 = 自动）", "plan.memory": "内存节省幅度", + "plan.mem.none": "不节省（最快）", "plan.mem.balanced": "均衡（int8 KV，专家 RAM 75%）", + "plan.mem.aggressive": "激进（共享路由，激活重算）", "plan.mem.extreme": "极限（优化器分页，专家 RAM 25%）", + "plan.data": "数据集路径（每行一个；可为 shards 目录、jsonl、parquet）", "plan.ram": "可用内存 GB（0 = 自动检测）", + "plan.tokens": "预训练 token 预算（B；0 = 按数据集）", "plan.vocab": "词表大小（0 = 按分词器报告）", + "plan.seq_len": "序列长度（0 = 按数据集）", "plan.run": "⚙️ 生成规划", "plan.path": "规划文件路径", + "plan.save": "💾 保存规划", "plan.table": "各阶段训练参数", + "plan.table_help": "每行对应一个训练阶段，编辑后立即生效；空单元格表示沿用共享参数。pretrain 行与下方的共享训练参数同步。训练页按所选阶段合并对应行。", + "plan.report": "规划报告", "plan.applied": "✅ 已更新各阶段参数", + "plan.err.target": "❌ 目标参数量必须大于 0", "plan.err.active_experts": "❌ 激活专家数不能超过专家总数", + "plan.err.no_plan": "❌ 请先生成规划", + }, + "en": { + "plan.title": "🧭 Auto-plan training parameters", + "plan.intro": "Describe the target and get a v3 architecture plus independent hyperparameters for every stage. LR and batch follow the DeepSeek LLM scaling law (arXiv 2401.02954); memory uses a model calibrated on measured runs. Dataset paths supply token counts, sequence length, record counts and the tokenizer. Every figure is labelled with its source.", + "plan.target": "Target total parameters (M)", "plan.use_moe": "Use MoE", "plan.experts": "Experts", + "plan.depth": "Recurrent depth", "plan.active_mode": "Specify active size by", + "plan.mode.experts": "Active experts (top-k)", "plan.mode.params": "Active parameters (M)", + "plan.active_value": "Active value (experts or M params; 0 = auto)", "plan.memory": "Memory saving", + "plan.mem.none": "None (fastest)", "plan.mem.balanced": "Balanced (int8 KV, 75% expert RAM)", + "plan.mem.aggressive": "Aggressive (shared routing, activation recompute)", + "plan.mem.extreme": "Extreme (optimizer paging, 25% expert RAM)", + "plan.data": "Dataset paths (one per line; shards dir, jsonl, parquet)", "plan.ram": "Available RAM GB (0 = detect)", + "plan.tokens": "Pretrain token budget (B; 0 = from dataset)", "plan.vocab": "Vocab size (0 = tokenizer report)", + "plan.seq_len": "Sequence length (0 = from dataset)", "plan.run": "⚙️ Generate plan", "plan.path": "Plan file", + "plan.save": "💾 Save plan", "plan.table": "Per-stage hyperparameters", + "plan.table_help": "One row per training stage; edits apply immediately and blank cells fall back to the shared values. The pretrain row mirrors the shared training widgets below. The Train tab merges the row of the selected stage.", + "plan.report": "Plan report", "plan.applied": "✅ Stage parameters updated", + "plan.err.target": "❌ Target parameters must be > 0", "plan.err.active_experts": "❌ Active experts cannot exceed the expert count", + "plan.err.no_plan": "❌ Generate a plan first", + }, + "zh-Hant": { + "plan.title": "🧭 自動規劃訓練參數", + "plan.intro": "填寫目標後產生 v3 架構和各階段獨立的訓練參數。學習率與 batch 依 DeepSeek LLM 縮放律（arXiv 2401.02954）估算，記憶體依實測標定的模型估算；資料集路徑用於推斷 token 數、序列長度、記錄數和分詞器。所有數值都標註來源。", + "plan.target": "目標總參數量（M）", "plan.use_moe": "啟用 MoE", "plan.experts": "專家數", + "plan.depth": "循環深度", "plan.active_mode": "啟用規模的指定方式", + "plan.mode.experts": "啟用專家數（top-k）", "plan.mode.params": "啟用參數量（M）", + "plan.active_value": "啟用值（專家數或 M 參數；0 = 自動）", "plan.memory": "記憶體節省幅度", + "plan.mem.none": "不節省（最快）", "plan.mem.balanced": "均衡（int8 KV，專家 RAM 75%）", + "plan.mem.aggressive": "激進（共享路由，啟動重算）", "plan.mem.extreme": "極限（最佳化器分頁，專家 RAM 25%）", + "plan.data": "資料集路徑（每行一個；可為 shards 目錄、jsonl、parquet）", "plan.ram": "可用記憶體 GB（0 = 自動偵測）", + "plan.tokens": "預訓練 token 預算（B；0 = 依資料集）", "plan.vocab": "詞表大小（0 = 依分詞器報告）", + "plan.seq_len": "序列長度（0 = 依資料集）", "plan.run": "⚙️ 產生規劃", "plan.path": "規劃檔案路徑", + "plan.save": "💾 儲存規劃", "plan.table": "各階段訓練參數", + "plan.table_help": "每列對應一個訓練階段，編輯後立即生效；空儲存格表示沿用共享參數。pretrain 列與下方的共享訓練參數同步。訓練頁依所選階段合併對應列。", + "plan.report": "規劃報告", "plan.applied": "✅ 已更新各階段參數", + "plan.err.target": "❌ 目標參數量必須大於 0", "plan.err.active_experts": "❌ 啟用專家數不能超過專家總數", + "plan.err.no_plan": "❌ 請先產生規劃", + }, + "ja": { + "plan.title": "🧭 学習パラメータの自動計画", + "plan.intro": "目標を入力すると、v3 アーキテクチャとステージごとに独立したハイパーパラメータを生成します。学習率と batch は DeepSeek LLM のスケーリング則（arXiv 2401.02954）、メモリは実測で較正したモデルで推定します。データセットのパスからトークン数・系列長・レコード数・トークナイザを推定します。すべての数値に出典を付けます。", + "plan.target": "目標総パラメータ数（M）", "plan.use_moe": "MoE を使う", "plan.experts": "エキスパート数", + "plan.depth": "再帰深さ", "plan.active_mode": "アクティブ規模の指定方法", + "plan.mode.experts": "アクティブエキスパート数（top-k）", "plan.mode.params": "アクティブパラメータ数（M）", + "plan.active_value": "アクティブ値（エキスパート数または M パラメータ；0 = 自動）", "plan.memory": "メモリ節約レベル", + "plan.mem.none": "なし（最速）", "plan.mem.balanced": "バランス（int8 KV、エキスパート RAM 75%）", + "plan.mem.aggressive": "積極（共有ルーティング、活性再計算）", "plan.mem.extreme": "最大（オプティマイザのページング、エキスパート RAM 25%）", + "plan.data": "データセットのパス（1 行 1 つ；shards ディレクトリ、jsonl、parquet）", "plan.ram": "使用可能 RAM GB（0 = 自動検出）", + "plan.tokens": "事前学習トークン予算（B；0 = データセットから）", "plan.vocab": "語彙サイズ（0 = トークナイザレポート）", + "plan.seq_len": "系列長（0 = データセットから）", "plan.run": "⚙️ 計画を生成", "plan.path": "計画ファイル", + "plan.save": "💾 計画を保存", "plan.table": "ステージ別ハイパーパラメータ", + "plan.table_help": "1 行が 1 ステージに対応し、編集はすぐに反映されます。空欄は共有値を使います。pretrain 行は下の共有学習パラメータと同期します。学習タブは選択したステージの行をマージします。", + "plan.report": "計画レポート", "plan.applied": "✅ ステージ別パラメータを更新しました", + "plan.err.target": "❌ 目標パラメータ数は 0 より大きくしてください", "plan.err.active_experts": "❌ アクティブエキスパート数はエキスパート総数以下にしてください", + "plan.err.no_plan": "❌ 先に計画を生成してください", + }, +} +for _lang, _entries in _PLANNER_TRANSLATIONS.items(): + TRANSLATIONS[_lang].update(_entries) diff --git a/ui/planner_panel.py b/ui/planner_panel.py new file mode 100644 index 0000000..785f91e --- /dev/null +++ b/ui/planner_panel.py @@ -0,0 +1,177 @@ +"""ui/planner_panel.py — Config tab "auto plan" accordion. + +The user states a target (total parameters, MoE on/off, expert count, active +parameters *or* active experts, memory-saving level, optional datasets) and +:func:`chronos.model.planner.plan_training` returns a v3 architecture plus +independent hyperparameters for every training stage. The stage table is +editable; edits are written back to ``config_state["stage_overrides"]``, +which the Train tab merges over the shared values for the selected stage. + +The handlers are plain functions so they can be tested without Gradio. +""" +from __future__ import annotations + +import json +import os + +from chronos.model.planner import ( + MEMORY_LEVELS, STAGE_KEYS, PlanRequest, overrides_from_table, plan_training, stage_table, +) +from ui.gradio_compat import gr +from ui.i18n import register_translatable, t +from ui.tabs.tokenizer_tab import parse_paths + +ACTIVE_MODES = ("experts", "params") +STAGE_HEADERS = ["stage"] + list(STAGE_KEYS) + + +def run_plan(target_m, use_moe, num_experts, active_mode, active_value, memory_saving, data_paths, + ram_gb, token_budget_b, depth, vocab_size, seq_len) -> dict: + """Build a PlanRequest from UI values (millions / billions) and run the planner.""" + target = float(target_m or 0) * 1e6 + if target <= 0: + raise ValueError("target") + use_moe = bool(use_moe) + value = float(active_value or 0) + active_params = value * 1e6 if use_moe and active_mode == "params" and value > 0 else None + active_experts = int(value) if use_moe and active_mode == "experts" and value >= 1 else None + if use_moe and active_experts and active_experts > int(num_experts or 1): + raise ValueError("active_experts") + request = PlanRequest( + target_total=target, use_moe=use_moe, num_experts=int(num_experts or 1), + active_params=active_params, active_experts=active_experts, + memory_saving=memory_saving if memory_saving in MEMORY_LEVELS else "none", + data_paths=tuple(parse_paths(data_paths or "")), + ram_gb=float(ram_gb) if ram_gb else None, + token_budget=float(token_budget_b) * 1e9 if token_budget_b else None, + recurrent_depth=int(depth or 4), + vocab_size=int(vocab_size) if vocab_size else None, + seq_len=int(seq_len) if seq_len else None, + ) + return plan_training(request) + + +def plan_summary(plan: dict) -> str: + """Short Markdown summary; every figure carries its label.""" + arch, scal, mem = plan["report"]["architecture"], plan["report"]["scaling"], plan["report"]["memory"] + lines = [ + f"**{arch['total_params'] / 1e6:.2f}M** total · **{arch['active_params'] / 1e6:.2f}M** active " + f"(union ≤ {arch['active_union_upper'] / 1e6:.2f}M) — {arch['params_source']}", + f"h={arch['hidden_size']} · L={arch['physical_layers']} (groups {arch['recurrent_groups']} × depth " + f"{arch['recurrent_depth']} → {arch['virtual_layers']} virtual) · heads {arch['heads']} · " + f"E={arch['experts']} top-{arch['top_k']} width {arch['expert_width']}", + f"LR {scal['chosen_lr']:g} · batch {scal['chosen_batch_tokens']:,} tokens · " + f"{scal['tokens']:.3g} tokens ({scal['tokens_source']}) — estimated, {scal['law']}", + f"micro {mem['micro_tokens']:,} tokens × accum {mem['accumulation_steps']} · training peak ≈ " + f"{mem['training_peak_gb']} GiB (estimated) of {mem['target_gb']} GiB target · KV {mem['kv_cache_dtype']} · " + f"router {mem['router_scope']}", + ] + lines += [f"- {note}" for note in plan["report"]["notes"]] + return "\n\n".join(lines) + + +def apply_stage_rows(cfg: dict, rows) -> dict: + """Write an edited stage table back into the config (pretrain row also sets the shared values).""" + cfg = dict(cfg or {}) + if hasattr(rows, "values") and hasattr(rows, "columns"): # pandas DataFrame + rows = rows.values.tolist() + overrides = overrides_from_table(rows) + cfg["stage_overrides"] = overrides + cfg.update(overrides.get("pretrain", {})) + return cfg + + +def stage_rows(cfg: dict) -> list[list]: + return stage_table((cfg or {}).get("stage_overrides") or {})[1] + + +def save_plan(plan: dict, path: str) -> str: + """Write the full plan (config, per-stage values, MLX long-run phases, report) as JSON.""" + path = os.path.abspath(os.path.expanduser(path or "./out/chronos_plan.json")) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + json.dump(plan, fh, indent=2, ensure_ascii=False, default=str) + return path + + +def build_planner_panel(config_state, config_result, config_outputs: list) -> None: + """Render the accordion inside the Config tab and wire it to ``config_state``. + + ``config_result(cfg)`` must return the values for ``config_outputs`` (state, + JSON view, every editor widget and the estimator panel). + """ + with gr.Accordion(t("plan.title"), open=False) as acc: + register_translatable(acc, "plan.title") + intro = gr.Markdown(t("plan.intro")) + with gr.Row(): + target = gr.Number(label=t("plan.target"), value=200, precision=2, minimum=0.1) + use_moe = gr.Checkbox(label=t("plan.use_moe"), value=True) + experts = gr.Number(label=t("plan.experts"), value=32, precision=0, minimum=1) + depth = gr.Number(label=t("plan.depth"), value=4, precision=0, minimum=1, maximum=16) + with gr.Row(): + active_mode = gr.Radio(choices=[(t("plan.mode.experts"), "experts"), (t("plan.mode.params"), "params")], + value="experts", label=t("plan.active_mode")) + active_value = gr.Number(label=t("plan.active_value"), value=2, precision=2, minimum=0) + memory = gr.Dropdown(choices=[(t(f"plan.mem.{lvl}"), lvl) for lvl in MEMORY_LEVELS], + value="balanced", label=t("plan.memory")) + with gr.Row(): + data_paths = gr.Textbox(label=t("plan.data"), lines=2, scale=3, + placeholder="./dataset/shards/\n./dataset/sft_mini_512.jsonl") + with gr.Column(scale=2): + ram_gb = gr.Number(label=t("plan.ram"), value=0, precision=1, minimum=0) + token_budget = gr.Number(label=t("plan.tokens"), value=0, precision=3, minimum=0) + with gr.Column(scale=2): + vocab = gr.Number(label=t("plan.vocab"), value=0, precision=0, minimum=0) + seq_len = gr.Number(label=t("plan.seq_len"), value=0, precision=0, minimum=0) + with gr.Row(): + plan_btn = gr.Button(t("plan.run"), variant="primary") + plan_path = gr.Textbox(value="./out/chronos_plan.json", label=t("plan.path"), scale=2) + save_btn = gr.Button(t("plan.save"), variant="secondary") + summary = gr.Markdown("") + table_md = gr.Markdown(t("plan.table_help")) + table = gr.Dataframe(headers=STAGE_HEADERS, value=[], type="array", interactive=True, + label=t("plan.table"), wrap=False, static_columns=[0]) + report = gr.JSON(label=t("plan.report"), value={}) + plan_state = gr.State({}) + for comp, key in ((intro, "plan.intro"), (target, "plan.target"), (use_moe, "plan.use_moe"), + (experts, "plan.experts"), (depth, "plan.depth"), (active_mode, "plan.active_mode"), + (active_value, "plan.active_value"), (memory, "plan.memory"), (data_paths, "plan.data"), + (ram_gb, "plan.ram"), (token_budget, "plan.tokens"), (vocab, "plan.vocab"), + (seq_len, "plan.seq_len"), (plan_btn, "plan.run"), (plan_path, "plan.path"), + (save_btn, "plan.save"), (table_md, "plan.table_help"), (table, "plan.table"), + (report, "plan.report")): + register_translatable(comp, key) + + def on_plan(*values): + try: + plan = run_plan(*values) + except (ValueError, TypeError, OSError) as exc: + key = f"plan.err.{exc}" if str(exc) in ("target", "active_experts") else None + message = t(key) if key else f"❌ {type(exc).__name__}: {exc}" + return [message, gr.update(), gr.update(), gr.update()] + [gr.update() for _ in config_outputs] + cfg = plan["config"] + return [plan_summary(plan), stage_rows(cfg), plan["report"], plan] + config_result(cfg) + + def on_apply(cfg, rows): + try: + cfg = apply_stage_rows(cfg, rows) + except (ValueError, TypeError) as exc: + return [f"❌ {exc}"] + [gr.update() for _ in config_outputs] + return [t("plan.applied")] + config_result(cfg) + + def on_save(plan, cfg, path): + if not plan: + return t("plan.err.no_plan") + plan = dict(plan, config=dict(cfg), stage_overrides=dict((cfg or {}).get("stage_overrides") or {})) + return f"✅ {save_plan(plan, path)}" + + plan_inputs = [target, use_moe, experts, active_mode, active_value, memory, data_paths, ram_gb, token_budget, + depth, vocab, seq_len] + plan_btn.click(fn=on_plan, inputs=plan_inputs, outputs=[summary, table, report, plan_state] + config_outputs) + # Cell edits apply immediately, so the table never holds un-applied values + # that a later refresh could discard. + table.input(fn=on_apply, inputs=[config_state, table], outputs=[summary] + config_outputs) + save_btn.click(fn=on_save, inputs=[plan_state, config_state, plan_path], outputs=[summary]) + # Presets, loaded configs and shared-widget edits (which keep the pretrain + # row in sync, see merge_config_inputs) refresh the table. + config_state.change(fn=stage_rows, inputs=[config_state], outputs=[table]) diff --git a/ui/presets.py b/ui/presets.py index e1c3cc6..9ac41b8 100644 --- a/ui/presets.py +++ b/ui/presets.py @@ -23,138 +23,248 @@ # These values match MiniMindConfig(hidden_size=768, num_hidden_layers=8, # use_moe=True) field-for-field, with Chronos-specific extensions left at # their conservative defaults (no router-anchor by default). -MINIMIND_MOE_DEFAULTS: Dict = { - # MiniMind core - "hidden_size": 768, - "num_hidden_layers": 8, - "num_attention_heads": 8, - "num_key_value_heads": 4, - "vocab_size": 6400, - "intermediate_size": 0, # 0 = auto: ceil(H·π/64)·64 → 2432 - "moe_intermediate_size": 0, # 0 = auto, same value - "max_seq_len": 512, # training sequence length (NOT max_position_embeddings) - "tie_word_embeddings": True, - - # MoE - "num_experts": 4, - "num_experts_per_tok": 2, # M8c: top-2 doubles activated params at fixed total params - "num_shared_experts": 1, # Chronos extension; minimind has no shared - - # Attention shape (Chronos hybrid) - "kv_latent_dim": 64, - "rope_dim": 32, - "sliding_window_size": 2048, - "use_hybrid_attention": True, - - # Chronos prediction - "lookahead_steps": 2, - - # Loss coefficients - "lambda_balance": 5e-4, - "lambda_temporal": 1e-3, - "lambda_lookahead": 0.1, - "lambda_lookahead_topk": 0.05, - "lambda_router_anchor": 0.0, - "fallback_mask_prob": 0.0, - - # Hardware - "vram_budget_gb": 4.0, - "pinned_memory_max_fraction": 0.25, - "storage_format": "safetensors", - - # Training - "learning_rate": 5e-4, - "batch_size": 16, - "accumulation_steps": 8, - "epochs": 2, - "max_steps": 0, # 0 = no cap; use epochs * dataset length - "save_interval": 1000, - "weight_decay": 0.01, - "grad_clip": 1.0, - "log_interval": 10, - "reward_spec": "toy", - "max_gen_len": 24, - "num_generations": 4, - "temperature": 1.0, - "beta": 0.1, - "alpha": 0.7, - "lambda_or": 0.1, - "save_dir": "./out", - "dtype": "auto", - "cpu_threads": "auto", - "cpu_budget_percent": 100, - "num_workers": "auto", -} - - -# Order matches ui/tabs/config_tab.py::all_inputs (used by Apply-Preset -# to push values back into the slider widgets in the right slots). -CONFIG_INPUT_ORDER: List[str] = [ - "hidden_size", "num_hidden_layers", "num_experts", "num_experts_per_tok", - "num_shared_experts", "lookahead_steps", "kv_latent_dim", "sliding_window_size", - "num_attention_heads", "num_key_value_heads", "rope_dim", "moe_intermediate_size", - "vocab_size", "dtype", "tie_word_embeddings", - "lambda_balance", "lambda_temporal", "lambda_lookahead", "lambda_lookahead_topk", "lambda_router_anchor", - "fallback_mask_prob", - "vram_budget_gb", "pinned_memory_max_fraction", "use_hybrid_attention", "storage_format", - "learning_rate", "batch_size", "accumulation_steps", "max_seq_len", - "epochs", "max_steps", "save_interval", "log_interval", "save_dir", - "weight_decay", "grad_clip", - "cpu_threads", "cpu_budget_percent", "num_workers", - "reward_spec", "max_gen_len", "num_generations", "temperature", - "beta", "alpha", "lambda_or", +MINIMIND_MOE_DEFAULTS: Dict = {'architecture_version': 2, + 'recurrent_enabled': True, + 'recurrent_groups': 2, + 'recurrent_depth': 4, + 'min_recurrent_depth': 2, + 'prelude_layers': 1, + 'coda_layers': 1, + 'adaptive_halting': True, + 'halting_threshold': 0.99, + 'halting_epsilon': 0.01, + 'max_inference_depth': 4, + 'ponder_loss_coef': 0.01, + 'recurrent_distill_coef': 0.05, + 'recurrent_stability_coef': 0.001, + 'recurrent_stability_threshold': 0.1, + 'recurrent_dt_init': 0.25, + 'recurrent_dt_min': 0.05, + 'recurrent_dt_max': 1.0, + 'recurrent_latent_refinement': True, + 'recurrent_strategy': 'scheduled', + 'recurrent_depth_jitter': 0.0, + 'recurrent_parallel_refine': False, + 'depth_schedule': 'fixed_then_adaptive', + 'depth_schedule_warmup_fraction': 0.25, + 'layer_aware_lookahead': True, + 'layer_aware_cache': True, + 'training_paging': False, + 'optimizer_paging': False, + 'optimizer_page_dir': './optimizer_pages', + 'page_size_mb': 64, + 'prefetch_depth': 2, + 'activation_checkpointing': False, + 'modalities': [], + 'auto_detect_modalities': True, + 'image_size': 224, + 'image_patch_size': 16, + 'max_image_tokens': 196, + 'audio_sample_rate': 16000, + 'max_audio_seconds': 30.0, + 'max_audio_tokens': 128, + 'modality_dropout': 0.05, + 'freeze_modality_encoders': False, + 'image_encoder_type': 'patch', + 'audio_encoder_type': 'conv', + 'modality_fusion_heads': 4, + 'media_error_policy': 'warn', + 'modality_adapters': {}, + 'modality_alignment_coef': 0.0, + 'hidden_size': 768, + 'num_hidden_layers': 8, + 'num_attention_heads': 8, + 'num_key_value_heads': 4, + 'vocab_size': 6400, + 'intermediate_size': 0, + 'moe_intermediate_size': 0, + 'max_seq_len': 512, + 'tie_word_embeddings': True, + 'num_experts': 4, + 'num_experts_per_tok': 2, + 'num_shared_experts': 1, + 'kv_latent_dim': 64, + 'rope_dim': 32, + 'sliding_window_size': 2048, + 'use_hybrid_attention': True, + 'lookahead_steps': 2, + 'lambda_balance': 0.0005, + 'lambda_temporal': 0.001, + 'lambda_lookahead': 0.1, + 'lambda_lookahead_topk': 0.05, + 'lambda_router_anchor': 0.0, + 'fallback_mask_prob': 0.0, + 'vram_budget_gb': 4.0, + 'pinned_memory_max_fraction': 0.25, + 'storage_format': 'safetensors', + 'learning_rate': 0.0005, + 'batch_size': 16, + 'accumulation_steps': 8, + 'epochs': 2, + 'max_steps': 0, + 'save_interval': 1000, + 'weight_decay': 0.01, + 'grad_clip': 1.0, + 'log_interval': 10, + 'reward_spec': 'toy', + 'max_gen_len': 24, + 'num_generations': 4, + 'temperature': 1.0, + 'beta': 0.1, + 'alpha': 0.7, + 'lambda_or': 0.1, + 'save_dir': './out', + 'dtype': 'auto', + 'cpu_threads': 'auto', + 'cpu_budget_percent': 100, + 'num_workers': 'auto'} + + +CONFIG_INPUT_ORDER = ['hidden_size', + 'num_hidden_layers', + 'num_experts', + 'num_experts_per_tok', + 'num_shared_experts', + 'lookahead_steps', + 'kv_latent_dim', + 'sliding_window_size', + 'num_attention_heads', + 'num_key_value_heads', + 'rope_dim', + 'moe_intermediate_size', + 'vocab_size', + 'dtype', + 'tie_word_embeddings', + 'lambda_balance', + 'lambda_temporal', + 'lambda_lookahead', + 'lambda_lookahead_topk', + 'lambda_router_anchor', + 'fallback_mask_prob', + 'vram_budget_gb', + 'pinned_memory_max_fraction', + 'use_hybrid_attention', + 'storage_format', + 'learning_rate', + 'batch_size', + 'accumulation_steps', + 'max_seq_len', + 'epochs', + 'max_steps', + 'save_interval', + 'log_interval', + 'save_dir', + 'weight_decay', + 'grad_clip', + 'cpu_threads', + 'cpu_budget_percent', + 'num_workers', + 'reward_spec', + 'max_gen_len', + 'num_generations', + 'temperature', + 'beta', + 'alpha', + 'lambda_or', + 'recurrent_enabled', + 'recurrent_groups', + 'recurrent_depth', + 'min_recurrent_depth', + 'adaptive_halting', + 'halting_threshold', + 'max_inference_depth', + 'ponder_loss_coef', + 'recurrent_distill_coef', + 'recurrent_stability_coef', + 'recurrent_latent_refinement', + 'depth_schedule', + 'recurrent_strategy', + 'recurrent_depth_jitter', + 'depth_schedule_warmup_fraction', + 'layer_aware_cache', + 'prefetch_depth', + 'training_paging', + 'optimizer_paging', + 'page_size_mb', + 'optimizer_page_dir', + 'modalities'] + + +# ── Chronos v3 tiers ─────────────────────────────────────────────── +# Generated by the adaptive planner (chronos.model.planner) and stored as JSON +# so importing the UI never runs the search. ``python -m ui.presets`` +# regenerates the file from V3_TIER_REQUESTS. +V3_PRESETS_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "presets_v3.json") + +V3_TIER_REQUESTS = [ + ("Chronos-v3 Tiny (smoke / CI)", "1.5M smoke test; trains in minutes on CPU.", + dict(target_total=1.5e6, num_experts=4, active_experts=2, vocab_size=6400, seq_len=128, recurrent_depth=2, + memory_saving="none", token_budget=2e6, ram_gb=16)), + ("Chronos-v3 50M (en→zh recipe)", "Exactly the enzh_50m_72h recipe (16k Chronos tokenizer, shards_v1).", + dict(target_total=50e6, num_experts=36, active_experts=2, data_paths=("Translation Training Data/shards_v1",), + memory_saving="none", ram_gb=32)), + ("Chronos-v3 200M (32 GB)", "General 200M MoE, balanced memory saving, minimind 6400 vocab.", + dict(target_total=200e6, num_experts=32, active_experts=2, vocab_size=6400, seq_len=1024, + memory_saving="balanced", ram_gb=32)), + ("Chronos-v3 500M (32 GB)", "500M MoE, aggressive memory saving (int8 KV, shared routing, checkpointing).", + dict(target_total=500e6, num_experts=48, active_experts=3, vocab_size=6400, seq_len=1024, + memory_saving="aggressive", ram_gb=32)), + ("Chronos-v3 1B (64 GB)", "1B MoE, 64 fine-grained experts, aggressive memory saving.", + dict(target_total=1e9, num_experts=64, active_experts=4, vocab_size=6400, seq_len=2048, + memory_saving="aggressive", ram_gb=64)), ] -# Named presets — keys are user-facing labels (also used in the dropdown). +def build_v3_presets() -> dict: + """Run the planner for every tier; returns the presets_v3.json document.""" + from chronos.model.planner import PlanRequest, plan_training + presets, meta = {}, {} + for name, blurb, request in V3_TIER_REQUESTS: + plan = plan_training(PlanRequest(**request)) + cfg = dict(plan["config"]) + if "tokenizer_path" in cfg: + cfg["tokenizer_path"] = "./Translation Training Data/tokenizer_v1" # repo-relative + presets[name] = {k: v for k, v in cfg.items() if MINIMIND_MOE_DEFAULTS.get(k, object()) != v} + arch, mem = plan["report"]["architecture"], plan["report"]["memory"] + meta[name] = dict( + summary=blurb, request={k: (list(v) if isinstance(v, tuple) else v) for k, v in request.items()}, + total_params=arch["total_params"], active_params=arch["active_params"], + active_union_upper=arch["active_union_upper"], params_label=arch["params_source"], + training_peak_gb=mem["training_peak_gb"], training_peak_label="estimated", + notes=plan["report"]["notes"]) + return {"_about": "Generated by chronos.model.planner.plan_training (inputs in _meta.request). Parameter " + "counts are measured on the meta device; memory figures are estimates. Regenerate with " + "`python -m ui.presets` instead of hand-editing.", + "_meta": meta, "presets": presets} + + +def _load_v3_presets() -> tuple[Dict[str, Dict], Dict[str, Dict]]: + try: + with open(V3_PRESETS_PATH, encoding="utf-8") as fh: + doc = json.load(fh) + except (OSError, ValueError): + return {}, {} + return ({name: {**MINIMIND_MOE_DEFAULTS, **cfg} for name, cfg in doc.get("presets", {}).items()}, + dict(doc.get("_meta", {}))) + + +V3_PRESETS, V3_PRESET_META = _load_v3_presets() +DEFAULT_PRESET = "Chronos-v3 200M (32 GB)" if "Chronos-v3 200M (32 GB)" in V3_PRESETS else "MiniMind-MoE (default)" + PRESETS: Dict[str, Dict] = { + **V3_PRESETS, + # Legacy v2 shape kept for minimind compatibility (first_only recurrence). "MiniMind-MoE (default)": dict(MINIMIND_MOE_DEFAULTS), +} - # M8c: "Recommended-CN" — sized for the /Users/Fona/Downloads/Hybrid_LLM/Dataset - # 1.3 GB pretrain corpus. Previously the default 15M-param model - # (H=256, L=4, E=4, topk=1 → ~6M activated) was far below the size - # needed for coherent Chinese output (empirically ~100M activated). - # This preset targets ~120M total / ~60M activated params at vocab=6400. - "Recommended-CN (≈120M)": { - **MINIMIND_MOE_DEFAULTS, - "hidden_size": 512, "num_hidden_layers": 8, "num_experts": 8, - "num_experts_per_tok": 2, "num_shared_experts": 1, - "num_attention_heads": 8, "num_key_value_heads": 2, - "kv_latent_dim": 64, "rope_dim": 32, "sliding_window_size": 2048, - "max_seq_len": 512, "batch_size": 8, "accumulation_steps": 8, - "epochs": 1, "save_interval": 500, - "learning_rate": 3e-4, - }, - - "Tiny (smoke / CI)": { - **MINIMIND_MOE_DEFAULTS, - "hidden_size": 128, "num_hidden_layers": 2, "num_experts": 4, - "num_attention_heads": 4, "num_key_value_heads": 4, - "kv_latent_dim": 16, "rope_dim": 8, "sliding_window_size": 64, - "max_seq_len": 64, "batch_size": 2, "accumulation_steps": 1, - "epochs": 1, "save_interval": 100, - }, - - "Small (256H × 4L × 4E)": { - **MINIMIND_MOE_DEFAULTS, - "hidden_size": 256, "num_hidden_layers": 4, "num_experts": 4, - "max_seq_len": 256, "batch_size": 8, - }, - - "Medium (512H × 8L × 8E)": { - **MINIMIND_MOE_DEFAULTS, - "hidden_size": 512, "num_experts": 8, - "max_seq_len": 1024, "batch_size": 8, - }, - - "Large (1024H × 16L × 16E)": { - **MINIMIND_MOE_DEFAULTS, - "hidden_size": 1024, "num_hidden_layers": 16, "num_experts": 16, - "num_attention_heads": 16, "num_key_value_heads": 4, - "kv_latent_dim": 128, "rope_dim": 64, - "max_seq_len": 2048, "batch_size": 4, "accumulation_steps": 16, - "vram_budget_gb": 16.0, - }, +# Pre-v3 preset names (saved sessions, docs) resolve to the closest v3 tier. +_PRESET_ALIASES = { + "Recommended-CN (≈120M)": "Chronos-v3 200M (32 GB)", + "Recommended-CN (≈192M)": "Chronos-v3 200M (32 GB)", + "Tiny (smoke / CI)": "Chronos-v3 Tiny (smoke / CI)", + "Small (256H × 4L × 4E)": "Chronos-v3 50M (en→zh recipe)", + "Medium (512H × 8L × 8E)": "Chronos-v3 200M (32 GB)", + "Large (1024H × 16L × 16E)": "Chronos-v3 1B (64 GB)", } @@ -163,9 +273,10 @@ def preset_names() -> List[str]: def get_preset(name: str) -> Dict: - """Return a deep-ish copy of the named preset (avoids accidental mutation).""" + """Return a copy of the named preset (avoids accidental mutation).""" + name = _PRESET_ALIASES.get(name, name) cfg = dict(MINIMIND_MOE_DEFAULTS) - cfg.update(PRESETS.get(name, {})) + cfg.update(json.loads(json.dumps(PRESETS.get(name, {})))) # nested stage_overrides copied too return cfg @@ -199,6 +310,17 @@ def values_in_input_order(cfg: Dict) -> List: "reward_spec": "toy", "max_gen_len": 24, "num_generations": 4, "temperature": 1.0, "beta": 0.1, "alpha": 0.7, "lambda_or": 0.1, + "recurrent_enabled": True, "recurrent_groups": 2, + "recurrent_depth": 4, "min_recurrent_depth": 2, + "adaptive_halting": True, "halting_threshold": 0.99, + "max_inference_depth": 4, "ponder_loss_coef": 1e-2, + "recurrent_distill_coef": 0.05, "recurrent_stability_coef": 1e-3, + "recurrent_latent_refinement": True, + "depth_schedule": "fixed_then_adaptive", "recurrent_strategy": "scheduled", + "recurrent_depth_jitter": 0.0, "depth_schedule_warmup_fraction": 0.25, + "layer_aware_cache": True, "prefetch_depth": 2, + "training_paging": False, "optimizer_paging": False, "page_size_mb": 64, + "optimizer_page_dir": "./optimizer_pages", "modalities": "", } out = [] for k in CONFIG_INPUT_ORDER: @@ -215,6 +337,8 @@ def values_in_input_order(cfg: Dict) -> List: "full": "float32", "": "auto", }.get(normalized, normalized) + if k == "modalities" and isinstance(v, (list, tuple)): + v = ",".join(v) out.append(v) return out @@ -233,5 +357,16 @@ def load_config(path: str) -> Dict: with open(os.path.abspath(path)) as f: loaded = json.load(f) cfg = dict(MINIMIND_MOE_DEFAULTS) + if not isinstance(loaded, dict): + raise ValueError("Configuration JSON must be an object") + if "max_position_embeddings" in loaded and "max_seq_len" not in loaded: + loaded["max_seq_len"] = loaded["max_position_embeddings"] cfg.update(loaded) return cfg + + +if __name__ == "__main__": # python -m ui.presets → regenerate presets_v3.json + with open(V3_PRESETS_PATH, "w", encoding="utf-8") as fh: + json.dump(build_v3_presets(), fh, indent=1, ensure_ascii=False) + fh.write("\n") + print(f"wrote {V3_PRESETS_PATH}") diff --git a/ui/presets_v3.json b/ui/presets_v3.json new file mode 100644 index 0000000..45cc2c9 --- /dev/null +++ b/ui/presets_v3.json @@ -0,0 +1,699 @@ +{ + "_about": "Generated by chronos.model.planner.plan_training (inputs in _meta.request). Parameter counts are measured on the meta device; memory figures are estimates. Regenerate with `python -m ui.presets` instead of hand-editing.", + "_meta": { + "Chronos-v3 Tiny (smoke / CI)": { + "summary": "1.5M smoke test; trains in minutes on CPU.", + "request": { + "target_total": 1500000.0, + "num_experts": 4, + "active_experts": 2, + "vocab_size": 6400, + "seq_len": 128, + "recurrent_depth": 2, + "memory_saving": "none", + "token_budget": 2000000.0, + "ram_gb": 16 + }, + "total_params": 1511106, + "active_params": 1289922, + "active_union_upper": 1363650, + "params_label": "measured (meta-device count)", + "training_peak_gb": 2.13, + "training_peak_label": "estimated", + "notes": [ + "C = 1.00e+13 FLOPs is outside the DeepSeek fit range 1e+17–2e+19; LR/batch are extrapolated and clamped" + ] + }, + "Chronos-v3 50M (en→zh recipe)": { + "summary": "Exactly the enzh_50m_72h recipe (16k Chronos tokenizer, shards_v1).", + "request": { + "target_total": 50000000.0, + "num_experts": 36, + "active_experts": 2, + "data_paths": [ + "Translation Training Data/shards_v1" + ], + "memory_saving": "none", + "ram_gb": 32 + }, + "total_params": 49691717, + "active_params": 12090437, + "active_union_upper": 16514117, + "params_label": "measured (meta-device count)", + "training_peak_gb": 16.4, + "training_peak_label": "estimated", + "notes": [ + "token counts from Translation Training Data/shards_v1/shards_manifest.json" + ] + }, + "Chronos-v3 200M (32 GB)": { + "summary": "General 200M MoE, balanced memory saving, minimind 6400 vocab.", + "request": { + "target_total": 200000000.0, + "num_experts": 32, + "active_experts": 2, + "vocab_size": 6400, + "seq_len": 1024, + "memory_saving": "balanced", + "ram_gb": 32 + }, + "total_params": 200724295, + "active_params": 28200775, + "active_union_upper": 54079303, + "params_label": "measured (meta-device count)", + "training_peak_gb": 10.14, + "training_peak_label": "estimated", + "notes": [] + }, + "Chronos-v3 500M (32 GB)": { + "summary": "500M MoE, aggressive memory saving (int8 KV, shared routing, checkpointing).", + "request": { + "target_total": 500000000.0, + "num_experts": 48, + "active_experts": 3, + "vocab_size": 6400, + "seq_len": 1024, + "memory_saving": "aggressive", + "ram_gb": 32 + }, + "total_params": 499529738, + "active_params": 61585418, + "active_union_upper": 133249034, + "params_label": "measured (meta-device count)", + "training_peak_gb": 9.26, + "training_peak_label": "estimated", + "notes": [] + }, + "Chronos-v3 1B (64 GB)": { + "summary": "1B MoE, 64 fine-grained experts, aggressive memory saving.", + "request": { + "target_total": 1000000000.0, + "num_experts": 64, + "active_experts": 4, + "vocab_size": 6400, + "seq_len": 2048, + "memory_saving": "aggressive", + "ram_gb": 64 + }, + "total_params": 1002169739, + "active_params": 99739019, + "active_union_upper": 250144139, + "params_label": "measured (meta-device count)", + "training_peak_gb": 18.92, + "training_peak_label": "estimated", + "notes": [] + } + }, + "presets": { + "Chronos-v3 Tiny (smoke / CI)": { + "architecture_version": 3, + "recurrent_attention_mode": "shared_first", + "rope_layout": "per_head", + "hidden_size": 128, + "num_hidden_layers": 3, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "kv_latent_dim": 32, + "sliding_window_size": 128, + "use_moe": true, + "moe_intermediate_size": 96, + "recurrent_groups": 1, + "recurrent_depth": 2, + "min_recurrent_depth": 1, + "max_inference_depth": 2, + "halting_threshold": 0.95, + "depth_schedule_warmup_fraction": 0.3, + "recurrent_latent_refinement": false, + "ponder_loss_coef": 0.003, + "lambda_balance": 0.005, + "fallback_mask_prob": 0.08, + "dropout": 0.0, + "kv_cache_dtype": "auto", + "mla_cache_rope": true, + "router_scope": "per_iteration", + "ram_budget_mb": null, + "vram_budget_gb": 0.25, + "dtype": "bfloat16", + "learning_rate": 0.003, + "batch_size": 128, + "accumulation_steps": 1, + "max_seq_len": 128, + "epochs": 1, + "max_steps": 122, + "weight_decay": 0.1, + "save_interval": 100, + "stage_overrides": { + "pretrain": { + "learning_rate": 0.003, + "batch_size": 128, + "accumulation_steps": 1, + "max_seq_len": 128, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 122, + "weight_decay": 0.1, + "save_interval": 100, + "lambda_router_anchor": 0.0 + }, + "sft": { + "learning_rate": 0.001, + "batch_size": 128, + "accumulation_steps": 1, + "max_seq_len": 128, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "save_interval": 500, + "lambda_router_anchor": 0.0 + }, + "dpo": { + "learning_rate": 0.00015, + "batch_size": 64, + "accumulation_steps": 1, + "max_seq_len": 128, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "orpo": { + "learning_rate": 0.0003, + "batch_size": 64, + "accumulation_steps": 1, + "max_seq_len": 128, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "lambda_or": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "grpo": { + "learning_rate": 2e-05, + "batch_size": 32, + "accumulation_steps": 1, + "max_seq_len": 128, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.04, + "num_generations": 8, + "temperature": 1.0, + "max_gen_len": 64, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "distill": { + "learning_rate": 0.001, + "batch_size": 128, + "accumulation_steps": 1, + "max_seq_len": 128, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "alpha": 0.7, + "temperature": 4.0, + "lambda_router_anchor": 0.0, + "save_interval": 500 + } + } + }, + "Chronos-v3 50M (en→zh recipe)": { + "architecture_version": 3, + "recurrent_attention_mode": "shared_first", + "rope_layout": "per_head", + "hidden_size": 384, + "num_hidden_layers": 6, + "num_attention_heads": 6, + "num_key_value_heads": 2, + "kv_latent_dim": 128, + "sliding_window_size": 128, + "vocab_size": 16384, + "use_moe": true, + "num_experts": 36, + "moe_intermediate_size": 160, + "recurrent_groups": 4, + "halting_threshold": 0.95, + "depth_schedule_warmup_fraction": 0.3, + "recurrent_latent_refinement": false, + "ponder_loss_coef": 0.003, + "lambda_balance": 0.005, + "fallback_mask_prob": 0.08, + "dropout": 0.0, + "kv_cache_dtype": "auto", + "mla_cache_rope": true, + "router_scope": "per_iteration", + "ram_budget_mb": null, + "vram_budget_gb": 0.25, + "dtype": "bfloat16", + "learning_rate": 0.0021, + "batch_size": 32, + "epochs": 1, + "weight_decay": 0.1, + "save_interval": 3810, + "stage_overrides": { + "pretrain": { + "learning_rate": 0.0021, + "batch_size": 32, + "accumulation_steps": 8, + "max_seq_len": 512, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "save_interval": 3810, + "lambda_router_anchor": 0.0 + }, + "sft": { + "learning_rate": 0.0007, + "batch_size": 32, + "accumulation_steps": 4, + "max_seq_len": 512, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "save_interval": 31560, + "lambda_router_anchor": 0.0 + }, + "dpo": { + "learning_rate": 0.0001, + "batch_size": 16, + "accumulation_steps": 4, + "max_seq_len": 512, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "orpo": { + "learning_rate": 0.00021, + "batch_size": 16, + "accumulation_steps": 4, + "max_seq_len": 512, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "lambda_or": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "grpo": { + "learning_rate": 2e-05, + "batch_size": 8, + "accumulation_steps": 1, + "max_seq_len": 512, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.04, + "num_generations": 8, + "temperature": 1.0, + "max_gen_len": 256, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "distill": { + "learning_rate": 0.0007, + "batch_size": 32, + "accumulation_steps": 4, + "max_seq_len": 512, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "alpha": 0.7, + "temperature": 4.0, + "lambda_router_anchor": 0.0, + "save_interval": 31560 + } + }, + "tokenizer_path": "./Translation Training Data/tokenizer_v1" + }, + "Chronos-v3 200M (32 GB)": { + "architecture_version": 3, + "recurrent_attention_mode": "shared_first", + "rope_layout": "per_head", + "hidden_size": 576, + "num_attention_heads": 9, + "num_key_value_heads": 3, + "kv_latent_dim": 192, + "sliding_window_size": 256, + "use_moe": true, + "num_experts": 32, + "moe_intermediate_size": 416, + "recurrent_groups": 6, + "halting_threshold": 0.95, + "depth_schedule_warmup_fraction": 0.3, + "recurrent_latent_refinement": false, + "ponder_loss_coef": 0.003, + "lambda_balance": 0.005, + "fallback_mask_prob": 0.08, + "dropout": 0.0, + "kv_cache_dtype": "int8", + "mla_cache_rope": false, + "router_scope": "per_iteration", + "ram_budget_mb": 263.2, + "vram_budget_gb": 0.25, + "dtype": "bfloat16", + "learning_rate": 0.0015, + "batch_size": 4, + "accumulation_steps": 64, + "max_seq_len": 1024, + "epochs": 1, + "max_steps": 980096, + "weight_decay": 0.1, + "save_interval": 49000, + "stage_overrides": { + "pretrain": { + "learning_rate": 0.0015, + "batch_size": 4, + "accumulation_steps": 64, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 980096, + "weight_decay": 0.1, + "save_interval": 49000, + "lambda_router_anchor": 0.0 + }, + "sft": { + "learning_rate": 0.0005, + "batch_size": 4, + "accumulation_steps": 32, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "save_interval": 500, + "lambda_router_anchor": 0.0 + }, + "dpo": { + "learning_rate": 7.5e-05, + "batch_size": 2, + "accumulation_steps": 32, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "orpo": { + "learning_rate": 0.00015, + "batch_size": 2, + "accumulation_steps": 32, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "lambda_or": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "grpo": { + "learning_rate": 1.5e-05, + "batch_size": 1, + "accumulation_steps": 1, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.04, + "num_generations": 4, + "temperature": 1.0, + "max_gen_len": 256, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "distill": { + "learning_rate": 0.0005, + "batch_size": 4, + "accumulation_steps": 32, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "alpha": 0.7, + "temperature": 4.0, + "lambda_router_anchor": 0.0, + "save_interval": 500 + } + } + }, + "Chronos-v3 500M (32 GB)": { + "architecture_version": 3, + "recurrent_attention_mode": "shared_first", + "rope_layout": "per_head", + "num_hidden_layers": 11, + "num_attention_heads": 12, + "kv_latent_dim": 256, + "sliding_window_size": 256, + "use_moe": true, + "num_experts": 48, + "num_experts_per_tok": 3, + "moe_intermediate_size": 384, + "recurrent_groups": 9, + "halting_threshold": 0.95, + "depth_schedule_warmup_fraction": 0.3, + "recurrent_latent_refinement": false, + "ponder_loss_coef": 0.003, + "lambda_balance": 0.005, + "fallback_mask_prob": 0.08, + "dropout": 0.0, + "kv_cache_dtype": "int8", + "mla_cache_rope": false, + "router_scope": "shared", + "ram_budget_mb": 445.5, + "activation_checkpointing": true, + "vram_budget_gb": 0.35, + "dtype": "bfloat16", + "learning_rate": 0.0012, + "batch_size": 2, + "accumulation_steps": 256, + "max_seq_len": 1024, + "epochs": 1, + "max_steps": 4878080, + "weight_decay": 0.1, + "save_interval": 100000, + "stage_overrides": { + "pretrain": { + "learning_rate": 0.0012, + "batch_size": 2, + "accumulation_steps": 256, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 4878080, + "weight_decay": 0.1, + "save_interval": 100000, + "lambda_router_anchor": 0.0 + }, + "sft": { + "learning_rate": 0.0004, + "batch_size": 2, + "accumulation_steps": 128, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "save_interval": 500, + "lambda_router_anchor": 0.0 + }, + "dpo": { + "learning_rate": 6e-05, + "batch_size": 1, + "accumulation_steps": 128, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "orpo": { + "learning_rate": 0.00012, + "batch_size": 1, + "accumulation_steps": 128, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "lambda_or": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "grpo": { + "learning_rate": 1.2e-05, + "batch_size": 1, + "accumulation_steps": 1, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.04, + "num_generations": 4, + "temperature": 1.0, + "max_gen_len": 256, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "distill": { + "learning_rate": 0.0004, + "batch_size": 2, + "accumulation_steps": 128, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "alpha": 0.7, + "temperature": 4.0, + "lambda_router_anchor": 0.0, + "save_interval": 500 + } + } + }, + "Chronos-v3 1B (64 GB)": { + "architecture_version": 3, + "recurrent_attention_mode": "shared_first", + "rope_layout": "per_head", + "num_hidden_layers": 12, + "num_attention_heads": 12, + "kv_latent_dim": 256, + "sliding_window_size": 512, + "use_moe": true, + "num_experts": 64, + "num_experts_per_tok": 4, + "moe_intermediate_size": 544, + "recurrent_groups": 10, + "halting_threshold": 0.95, + "depth_schedule_warmup_fraction": 0.3, + "recurrent_latent_refinement": false, + "ponder_loss_coef": 0.003, + "lambda_balance": 0.005, + "fallback_mask_prob": 0.08, + "dropout": 0.0, + "kv_cache_dtype": "int8", + "mla_cache_rope": false, + "router_scope": "shared", + "ram_budget_mb": 918.0, + "activation_checkpointing": true, + "vram_budget_gb": 0.65, + "dtype": "bfloat16", + "learning_rate": 0.0011, + "batch_size": 2, + "accumulation_steps": 128, + "max_seq_len": 2048, + "epochs": 1, + "max_steps": 4893312, + "weight_decay": 0.1, + "save_interval": 100000, + "stage_overrides": { + "pretrain": { + "learning_rate": 0.0011, + "batch_size": 2, + "accumulation_steps": 128, + "max_seq_len": 2048, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 4893312, + "weight_decay": 0.1, + "save_interval": 100000, + "lambda_router_anchor": 0.0 + }, + "sft": { + "learning_rate": 0.00037, + "batch_size": 2, + "accumulation_steps": 64, + "max_seq_len": 2048, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "save_interval": 500, + "lambda_router_anchor": 0.0 + }, + "dpo": { + "learning_rate": 5.5e-05, + "batch_size": 1, + "accumulation_steps": 64, + "max_seq_len": 2048, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "orpo": { + "learning_rate": 0.00011, + "batch_size": 1, + "accumulation_steps": 64, + "max_seq_len": 2048, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "lambda_or": 0.1, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "grpo": { + "learning_rate": 1.1e-05, + "batch_size": 1, + "accumulation_steps": 1, + "max_seq_len": 2048, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.01, + "beta": 0.04, + "num_generations": 4, + "temperature": 1.0, + "max_gen_len": 256, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "distill": { + "learning_rate": 0.00037, + "batch_size": 2, + "accumulation_steps": 64, + "max_seq_len": 2048, + "grad_clip": 1.0, + "epochs": 1, + "max_steps": 0, + "weight_decay": 0.1, + "alpha": 0.7, + "temperature": 4.0, + "lambda_router_anchor": 0.0, + "save_interval": 500 + } + } + } + } +} diff --git a/ui/tabs/autotune_tab.py b/ui/tabs/autotune_tab.py index 6651991..818c592 100644 --- a/ui/tabs/autotune_tab.py +++ b/ui/tabs/autotune_tab.py @@ -161,21 +161,23 @@ def poll_tune(): best = _tuner.result.best_params return _tuner.status, _read_log(), best - def apply_best(best): + def apply_best(best, previous): """Write best_params back into the Config tab's slider widgets. Returns a list of gr.update for every config_input — untouched inputs get gr.update() (no-op). None values are dropped so Gradio's numeric widgets don't reject them.""" updates = [gr.update() for _ in config_inputs] + cfg = dict(previous or {}) if not best: - return updates + return [cfg] + updates for k, v in best.items(): if v is None: continue idx = PARAM_TO_CONFIG_IDX.get(k) if idx is not None: updates[idx] = gr.update(value=v) - return updates + cfg[k] = v + return [cfg] + updates start_btn.click( fn=start_tune, @@ -187,7 +189,7 @@ def apply_best(best): ) stop_btn.click(fn=stop_tune, outputs=[status_box]) clear_btn.click(fn=clear_log_file, outputs=[log_box]) - apply_btn.click(fn=apply_best, inputs=[best_box], outputs=config_inputs) + apply_btn.click(fn=apply_best, inputs=[best_box, config_state], outputs=[config_state] + config_inputs) timer = gr.Timer(value=2.0) timer.tick(fn=poll_tune, outputs=[status_box, log_box, best_box]) diff --git a/ui/tabs/benchmark_tab.py b/ui/tabs/benchmark_tab.py index 711d7d5..52323e8 100644 --- a/ui/tabs/benchmark_tab.py +++ b/ui/tabs/benchmark_tab.py @@ -116,6 +116,13 @@ def _load_existing(): def build_benchmark_tab(): + # Resolve the initial state before constructing components. Assigning a + # pandas DataFrame to ``component.value`` after construction bypasses + # Gradio's postprocessing step (notably in Gradio 5.x), causing the + # nativeplot frontend to receive a string representation of the frame and + # fail during hydration (the whole UI remains stuck on "Loading..."). + initial_results, initial_markdown, initial_df, initial_log = _load_existing() + with gr.Tab(t("tab.benchmark")) as tab: register_translatable(tab, "tab.benchmark") @@ -123,11 +130,11 @@ def build_benchmark_tab(): register_translatable(run_btn, "bench.run") # Formatted summary - summary_md = gr.Markdown(value=_format_table({})) + summary_md = gr.Markdown(value=initial_markdown) # Bar chart — one BarPlot per model so units stack cleanly per metric. chart = gr.BarPlot( - value=_empty_df(), + value=initial_df, x="metric", y="value", color="model", title="MiniMind vs Chronos — comparison metrics", tooltip=["model", "metric", "value", "unit"], @@ -135,11 +142,12 @@ def build_benchmark_tab(): ) with gr.Accordion("Raw JSON", open=False): - results_box = gr.JSON(label=t("bench.results")) + results_box = gr.JSON(label=t("bench.results"), value=initial_results) register_translatable(results_box, "bench.results") log_box = gr.Textbox( - label=t("bench.log"), lines=20, interactive=False, autoscroll=True + label=t("bench.log"), lines=20, interactive=False, autoscroll=True, + value=initial_log, ) register_translatable(log_box, "bench.log") @@ -192,11 +200,4 @@ def run_benchmark(): outputs=[results_box, summary_md, chart, log_box], ) - # Render existing results on tab open - d, md, df, log = _load_existing() - results_box.value = d - summary_md.value = md - chart.value = df - log_box.value = log - return tab diff --git a/ui/tabs/config_tab.py b/ui/tabs/config_tab.py index 2488244..cc51526 100644 --- a/ui/tabs/config_tab.py +++ b/ui/tabs/config_tab.py @@ -15,10 +15,12 @@ from ui.i18n import t, register_translatable from ui.estimator import ( ArchConfig, total_params, active_params, memory_footprint, - estimated_decode_tps, fmt_bytes, fmt_params, + compute_params, active_param_bounds, parameter_breakdown, estimated_decode_tps, fmt_bytes, fmt_params, ) +from ui.planner_panel import build_planner_panel from ui.presets import ( - preset_names, get_preset, values_in_input_order, save_config, load_config, + CONFIG_INPUT_ORDER, DEFAULT_PRESET, MINIMIND_MOE_DEFAULTS, preset_names, get_preset, values_in_input_order, + save_config, load_config, ) @@ -59,38 +61,67 @@ def _estimate(cfg_dict: dict): from hidden_size", matching MiniMindConfig's default rule. This keeps the estimator consistent with what the trainer actually builds. """ - mis = int(cfg_dict.get("moe_intermediate_size", 0) or 0) - its = int(cfg_dict.get("intermediate_size", 0) or 0) - arch = ArchConfig( - hidden_size=int(cfg_dict.get("hidden_size", 256)), - num_hidden_layers=int(cfg_dict.get("num_hidden_layers", 4)), - num_attention_heads=int(cfg_dict.get("num_attention_heads", 8)), - num_key_value_heads=int(cfg_dict.get("num_key_value_heads", 4)), - vocab_size=int(cfg_dict.get("vocab_size", 6400)), - intermediate_size=its if its > 0 else None, - moe_intermediate_size=mis if mis > 0 else None, - num_experts=int(cfg_dict.get("num_experts", 4)), - num_experts_per_tok=int(cfg_dict.get("num_experts_per_tok", 1)), - num_shared_experts=int(cfg_dict.get("num_shared_experts", 1)), - lookahead_steps=int(cfg_dict.get("lookahead_steps", 2)), - kv_latent_dim=int(cfg_dict.get("kv_latent_dim", 64)), - rope_dim=int(cfg_dict.get("rope_dim", 32)), - sliding_window_size=int(cfg_dict.get("sliding_window_size", 2048)), - use_hybrid_attention=bool(cfg_dict.get("use_hybrid_attention", True)), - tie_word_embeddings=bool(cfg_dict.get("tie_word_embeddings", True)), - dtype=_normalize_dtype_name(str(cfg_dict.get("dtype", "auto"))), - ) - mem = memory_footprint(arch) + arch = ArchConfig.from_config(cfg_dict) + seq_len = int(cfg_dict.get("max_seq_len", cfg_dict.get("max_position_embeddings", 1024))) + mem = memory_footprint(arch, seq_len) + low, high = active_param_bounds(arch) + detail = { + "parameters": parameter_breakdown(arch), + "active_unique_min": low, "active_unique_max": high, + "physical_blocks": arch.num_hidden_layers, + "prelude_groups_coda": [arch.prelude_layers, arch.recurrent_groups, arch.coda_layers], + "decode_depth_ceiling": arch.executed_depth, + "weight_use_count": compute_params(arch), + "inference_batch_size": 1, "context_tokens": seq_len, + "dtype_assumption": arch.dtype if arch.dtype != "auto" else "float32 (auto upper-bound scenario)", + "memory_bytes": mem, + "training_peak_excludes": ["activations", "attention workspace", "rollback snapshots", "allocator pools"], + } return ( fmt_params(total_params(arch)), - fmt_params(active_params(arch)), - fmt_bytes(mem["vram_estimate_bytes"]), - fmt_bytes(mem["ssd_estimate_bytes"]), - fmt_bytes(mem["kv_cache_bytes"]), - f"{estimated_decode_tps(arch):.1f}", + fmt_params(low) if low == high else f"{fmt_params(low)} – {fmt_params(high)}", + fmt_params(compute_params(arch)), + fmt_bytes(mem["vram_estimate_bytes"]), fmt_bytes(mem["ram_estimate_bytes"]), + fmt_bytes(mem["ssd_estimate_bytes"]), fmt_bytes(mem["kv_cache_bytes"]), + f"{estimated_decode_tps(arch):.1f}", detail, ) +def merge_config_inputs(previous: dict, values: list) -> dict: + """Apply editor fields by name while retaining loaded fields without widgets.""" + if len(values) != len(CONFIG_INPUT_ORDER): + raise ValueError(f"Config editor expects {len(CONFIG_INPUT_ORDER)} values, received {len(values)}") + cfg = dict(previous or MINIMIND_MOE_DEFAULTS) + for key, value in zip(CONFIG_INPUT_ORDER, values): + if key == "modalities": + value = value if isinstance(value, list) else str(value or "").split(",") + value = sorted(set(m.strip().lower() for m in value if m.strip() and m.strip().lower() != "text")) + elif key == "dtype": + value = _normalize_dtype_name(value) + else: + default = MINIMIND_MOE_DEFAULTS.get(key) + if value is None: + value = default + if isinstance(default, bool): + value = bool(value) + elif isinstance(default, int): + value = int(value) + elif isinstance(default, float): + value = float(value) + cfg[key] = value + # A loaded checkpoint/HF config may contain the canonical alias too. + # Keep the visible sequence-length control authoritative when it is edited. + if "max_position_embeddings" in cfg: + cfg["max_position_embeddings"] = cfg["max_seq_len"] + # The shared training widgets are the pretrain row of the per-stage table: + # an edited widget must not be silently overridden by a stale plan value. + stages = cfg.get("stage_overrides") + if isinstance(stages, dict) and isinstance(stages.get("pretrain"), dict): + pretrain = {k: (cfg[k] if k in CONFIG_INPUT_ORDER and k in cfg else v) for k, v in stages["pretrain"].items()} + cfg["stage_overrides"] = {**stages, "pretrain": pretrain} + return cfg + + def build_config_tab(): """Returns (config_state, list_of_input_components, save_dir). @@ -98,7 +129,7 @@ def build_config_tab(): data_path has been moved to the Train tab, so the tuple now has 3 elements. Callers must be updated. """ - initial_preset = "Recommended-CN (≈120M)" + initial_preset = DEFAULT_PRESET initial_cfg = get_preset(initial_preset) config_state = gr.State(dict(initial_cfg)) @@ -130,6 +161,8 @@ def build_config_tab(): load_btn = gr.Button(t("config.load_cfg"), variant="primary", scale=1) register_translatable(load_btn, "config.load_cfg") save_status = gr.Markdown("", visible=True) + # Auto-plan accordion; filled in below once the editor outputs exist. + planner_slot = gr.Column() with gr.Row(): with gr.Column(scale=3): @@ -210,6 +243,49 @@ def build_config_tab(): register_translatable(use_hybrid_attn, "config.hybrid_attn") register_translatable(storage_format, "config.storage_format") + gr.Markdown(f"### {t('config.recurrent_section')}") + with gr.Row(): + recurrent_enabled = gr.Checkbox(value=bool(D.get("recurrent_enabled", True)), label=t("config.recurrent_enabled")) + recurrent_groups = gr.Slider(1, 128, value=int(D.get("recurrent_groups", 2)), step=1, label=t("config.recurrent_groups")) + recurrent_depth = gr.Slider(1, 16, value=int(D.get("recurrent_depth", 4)), step=1, label=t("config.recurrent_depth")) + min_recurrent_depth = gr.Slider(1, 16, value=int(D.get("min_recurrent_depth", 2)), step=1, label=t("config.min_recurrent_depth")) + adaptive_halting = gr.Checkbox(value=bool(D.get("adaptive_halting", True)), label=t("config.adaptive_halting")) + with gr.Row(): + halting_threshold = gr.Slider(0.5, 1.0, value=float(D.get("halting_threshold", .99)), step=.001, label=t("config.halting_threshold")) + max_inference_depth = gr.Slider(1, 32, value=int(D.get("max_inference_depth", D.get("recurrent_depth", 4))), step=1, label=t("config.max_inference_depth")) + ponder_loss_coef = gr.Number(value=float(D.get("ponder_loss_coef", .001)), label=t("config.ponder_loss_coef"), precision=6) + recurrent_distill_coef = gr.Number(value=float(D.get("recurrent_distill_coef", .05)), label=t("config.recurrent_distill_coef"), precision=6) + recurrent_stability_coef = gr.Number(value=float(D.get("recurrent_stability_coef", .001)), label=t("config.recurrent_stability_coef"), precision=6) + recurrent_latent_refinement = gr.Checkbox( + value=bool(D.get("recurrent_latent_refinement", True)), + label=t("config.recurrent_latent_refinement"), + ) + depth_schedule = gr.Dropdown( + choices=["fixed_then_adaptive", "fixed", "adaptive"], + value=str(D.get("depth_schedule", "fixed_then_adaptive")), label=t("config.depth_schedule"), + ) + recurrent_strategy = gr.Dropdown( + choices=["fixed", "scheduled", "adaptive"], + value=str(D.get("recurrent_strategy", "scheduled")), + label="recurrent strategy", + ) + recurrent_depth_jitter = gr.Slider( + 0.0, 1.0, value=float(D.get("recurrent_depth_jitter", 0.0)), + step=.05, label="depth jitter", + ) + with gr.Row(): + depth_schedule_warmup_fraction = gr.Slider( + 0.0, 1.0, value=float(D.get("depth_schedule_warmup_fraction", .25)), step=.01, + label=t("config.depth_schedule_warmup_fraction"), + ) + layer_aware_cache = gr.Checkbox(value=bool(D.get("layer_aware_cache", True)), label=t("config.layer_aware_cache")) + prefetch_depth = gr.Slider(0, 64, value=int(D.get("prefetch_depth", 2)), step=1, label=t("config.prefetch_depth")) + training_paging = gr.Checkbox(value=bool(D.get("training_paging", False)), label=t("config.training_paging")) + optimizer_paging = gr.Checkbox(value=bool(D.get("optimizer_paging", False)), label=t("config.optimizer_paging")) + page_size_mb = gr.Number(value=int(D.get("page_size_mb", 64)), precision=0, label=t("config.page_size_mb")) + optimizer_page_dir = gr.Textbox(value=str(D.get("optimizer_page_dir", "./optimizer_pages")), label=t("config.optimizer_page_dir")) + modalities = gr.Textbox(value=",".join(D.get("modalities", [])), label=t("config.modalities")) + gr.Markdown(f"### {t('config.train')}") with gr.Row(): learning_rate = gr.Number(value=D["learning_rate"], label=t("config.lr"), precision=6) @@ -297,13 +373,19 @@ def build_config_tab(): gr.Markdown(f"### 🧬 {t('designer.title')}") total_box = gr.Textbox(label=t("designer.total"), interactive=False) active_box = gr.Textbox(label=t("designer.active"), interactive=False) + compute_box = gr.Textbox(label=t("designer.compute"), interactive=False) vram_box = gr.Textbox(label=t("designer.vram"), interactive=False) + ram_box = gr.Textbox(label=t("designer.ram"), interactive=False) ssd_box = gr.Textbox(label=t("designer.ssd"), interactive=False) kv_box = gr.Textbox(label=t("designer.kv"), interactive=False) tps_box = gr.Textbox(label=t("designer.tps"), interactive=False) + gr.Markdown(t("designer.notes")) + detail_box = gr.JSON(label=t("designer.details")) register_translatable(total_box, "designer.total") register_translatable(active_box, "designer.active") + register_translatable(compute_box, "designer.compute") register_translatable(vram_box, "designer.vram") + register_translatable(ram_box, "designer.ram") register_translatable(ssd_box, "designer.ssd") register_translatable(kv_box, "designer.kv") register_translatable(tps_box, "designer.tps") @@ -311,111 +393,139 @@ def build_config_tab(): # IMPORTANT: this list also defines the order Autotune uses to push # best-params back via gr.update(). Don't reorder without updating # ui/tabs/autotune_tab.py. - all_inputs = [ - hidden_size, num_layers, num_experts, experts_per_tok, - num_shared, lookahead, kv_latent_dim, sliding_window, - num_heads, num_kv_heads, rope_dim, moe_intermediate, - vocab_size, dtype, tie_word_embeddings, - lambda_balance, lambda_temporal, lambda_lookahead, lambda_lookahead_topk, lambda_router_anchor, - fallback_mask_prob, - vram_budget, pinned_frac, use_hybrid_attn, storage_format, - learning_rate, batch_size, accum_steps, max_seq_len, - epochs, max_steps, save_interval, log_interval, save_dir, - weight_decay, grad_clip, - cpu_threads, cpu_budget_percent, num_workers, - reward_spec, max_gen_len, num_generations, temperature, - beta, alpha, lambda_or, - ] - estimator_outs = [total_box, active_box, vram_box, ssd_box, kv_box, tps_box] - - def update_config(*vals): - (hs, nl, ne, ept, ns, la, kv, sw, - nh, nkv, rd, moe_i, - vs, dt, tie, - lb, lt, ll, llt, lra, fmp, - vb, pf, ha, sf, - lr, bs, ac, msl, ep, max_st, si, li, sd, - wd, gc, ct, cbp, nw, - reward, mgl, ng, temp, beta, alpha, lor) = vals - cfg = { - "hidden_size": int(hs), "num_hidden_layers": int(nl), - "num_experts": int(ne), "num_experts_per_tok": int(ept), - "num_shared_experts": int(ns), "lookahead_steps": int(la), - "kv_latent_dim": int(kv), "sliding_window_size": int(sw), - "num_attention_heads": int(nh), "num_key_value_heads": int(nkv), - "rope_dim": int(rd), "moe_intermediate_size": int(moe_i), - "vocab_size": int(vs), "dtype": _normalize_dtype_name(str(dt)), - "tie_word_embeddings": bool(tie), - "lambda_balance": float(lb), "lambda_temporal": float(lt), - "lambda_lookahead": float(ll), "lambda_lookahead_topk": float(llt), - "lambda_router_anchor": float(lra), - "fallback_mask_prob": float(fmp), - "vram_budget_gb": float(vb), "pinned_memory_max_fraction": float(pf), - "use_hybrid_attention": bool(ha), - "storage_format": str(sf), - "learning_rate": float(lr), "batch_size": int(bs), - "accumulation_steps": int(ac), "max_seq_len": int(msl), - "epochs": int(ep), "max_steps": max(0, int(max_st or 0)), - "save_interval": int(si), - "weight_decay": float(wd), - "grad_clip": float(gc), - "log_interval": int(li), - "reward_spec": str(reward or "toy"), - "max_gen_len": int(mgl), - "num_generations": int(ng), - "temperature": float(temp), - "beta": float(beta), - "alpha": float(alpha), - "lambda_or": float(lor), - "save_dir": sd, - "cpu_threads": str(ct or "auto"), - "cpu_budget_percent": float(cbp or 100), - "num_workers": str(nw or "auto"), - } - est = _estimate(cfg) - return (cfg, cfg) + est + config_components = { + "hidden_size": hidden_size, + "num_hidden_layers": num_layers, + "num_experts": num_experts, + "num_experts_per_tok": experts_per_tok, + "num_shared_experts": num_shared, + "lookahead_steps": lookahead, + "kv_latent_dim": kv_latent_dim, + "sliding_window_size": sliding_window, + "num_attention_heads": num_heads, + "num_key_value_heads": num_kv_heads, + "rope_dim": rope_dim, + "moe_intermediate_size": moe_intermediate, + "vocab_size": vocab_size, + "dtype": dtype, + "tie_word_embeddings": tie_word_embeddings, + "lambda_balance": lambda_balance, + "lambda_temporal": lambda_temporal, + "lambda_lookahead": lambda_lookahead, + "lambda_lookahead_topk": lambda_lookahead_topk, + "lambda_router_anchor": lambda_router_anchor, + "fallback_mask_prob": fallback_mask_prob, + "vram_budget_gb": vram_budget, + "pinned_memory_max_fraction": pinned_frac, + "use_hybrid_attention": use_hybrid_attn, + "storage_format": storage_format, + "learning_rate": learning_rate, + "batch_size": batch_size, + "accumulation_steps": accum_steps, + "max_seq_len": max_seq_len, + "epochs": epochs, + "max_steps": max_steps, + "save_interval": save_interval, + "log_interval": log_interval, + "save_dir": save_dir, + "weight_decay": weight_decay, + "grad_clip": grad_clip, + "cpu_threads": cpu_threads, + "cpu_budget_percent": cpu_budget_percent, + "num_workers": num_workers, + "reward_spec": reward_spec, + "max_gen_len": max_gen_len, + "num_generations": num_generations, + "temperature": temperature, + "beta": beta, + "alpha": alpha, + "lambda_or": lambda_or, + "recurrent_enabled": recurrent_enabled, + "recurrent_groups": recurrent_groups, + "recurrent_depth": recurrent_depth, + "min_recurrent_depth": min_recurrent_depth, + "adaptive_halting": adaptive_halting, + "halting_threshold": halting_threshold, + "max_inference_depth": max_inference_depth, + "ponder_loss_coef": ponder_loss_coef, + "recurrent_distill_coef": recurrent_distill_coef, + "recurrent_stability_coef": recurrent_stability_coef, + "recurrent_latent_refinement": recurrent_latent_refinement, + "depth_schedule": depth_schedule, + "recurrent_strategy": recurrent_strategy, + "recurrent_depth_jitter": recurrent_depth_jitter, + "depth_schedule_warmup_fraction": depth_schedule_warmup_fraction, + "layer_aware_cache": layer_aware_cache, + "prefetch_depth": prefetch_depth, + "training_paging": training_paging, + "optimizer_paging": optimizer_paging, + "page_size_mb": page_size_mb, + "optimizer_page_dir": optimizer_page_dir, + "modalities": modalities, + } + if set(config_components) != set(CONFIG_INPUT_ORDER): + raise RuntimeError("Config component registry differs from CONFIG_INPUT_ORDER") + all_inputs = [config_components[key] for key in CONFIG_INPUT_ORDER] + estimator_outs = [total_box, active_box, compute_box, vram_box, ram_box, ssd_box, kv_box, tps_box, detail_box] + + def estimator_values(cfg): + try: + return list(_estimate(cfg)) + except (ValueError, TypeError, ZeroDivisionError) as exc: + return ["—"] * 8 + [{"configuration_error": str(exc)}] + + def update_config(previous, *values): + cfg = merge_config_inputs(previous, values) + return [cfg, cfg] + estimator_values(cfg) + # User input only: preset/load updates must not trigger 68 competing + # change handlers that overwrite each other's state. for inp in all_inputs: - inp.change( - fn=update_config, - inputs=all_inputs, - outputs=[config_state, config_display] + estimator_outs, - ) + inp.input(fn=update_config, inputs=[config_state] + all_inputs, + outputs=[config_state, config_display] + estimator_outs) + + # External producers (autotune) update the state as well as widgets. + config_state.change(fn=lambda cfg: [cfg] + estimator_values(cfg), + inputs=[config_state], outputs=[config_display] + estimator_outs) + + config_outputs = [config_state, config_display] + all_inputs + estimator_outs + + def config_result(cfg): + return [dict(cfg), dict(cfg)] + [gr.update(value=v) for v in values_in_input_order(cfg)] + estimator_values(cfg) - # ── Preset / Save / Load wiring ─────────────────────── def apply_preset(name): - cfg = get_preset(name) - return [dict(cfg), dict(cfg)] + [gr.update(value=v) for v in values_in_input_order(cfg)] + return config_result(get_preset(name)) def reset_minimind(): - cfg = get_preset("MiniMind-MoE (default)") - return [dict(cfg), dict(cfg)] + [gr.update(value=v) for v in values_in_input_order(cfg)] + return config_result(get_preset("MiniMind-MoE (default)")) - def do_save(cfg, path): + def do_save(previous, path, *values): try: - abs_path = save_config(cfg or {}, path) + # Read the widgets at click time; a pending input event must + # not cause Save to write a stale or default configuration. + cfg = merge_config_inputs(previous, values) + abs_path = save_config(cfg, path or "./out/chronos_config.json") return f"✅ Saved to `{abs_path}`" - except Exception as e: - return f"❌ Save failed: {e}" + except Exception as exc: + return f"❌ Save failed: {exc}" def do_load(path): try: cfg = load_config(path) - vals = values_in_input_order(cfg) - return [f"✅ Loaded from `{path}`", dict(cfg), dict(cfg)] + [gr.update(value=v) for v in vals] - except FileNotFoundError: - return [f"❌ Not found: `{path}`", gr.update(), gr.update()] + [gr.update() for _ in all_inputs] - except Exception as e: - return [f"❌ Load failed: {e}", gr.update(), gr.update()] + [gr.update() for _ in all_inputs] - - apply_preset_btn.click(fn=apply_preset, inputs=[preset_dd], outputs=[config_state, config_display] + all_inputs) - # Selecting a preset in the dropdown should sync immediately; - # the explicit "Load Preset" button stays as a re-apply affordance. - preset_dd.change(fn=apply_preset, inputs=[preset_dd], outputs=[config_state, config_display] + all_inputs) - reset_btn.click(fn=reset_minimind, outputs=[config_state, config_display] + all_inputs) - save_btn.click(fn=do_save, inputs=[config_state, cfg_path], outputs=[save_status]) - load_btn.click(fn=do_load, inputs=[cfg_path], outputs=[save_status, config_state, config_display] + all_inputs) - - total_box.value, active_box.value, vram_box.value, ssd_box.value, kv_box.value, tps_box.value = _estimate(initial_cfg) + return [f"✅ Loaded from `{path}`"] + config_result(cfg) + except Exception as exc: + return [f"❌ Load failed: {exc}"] + [gr.update() for _ in config_outputs] + + with planner_slot: + build_planner_panel(config_state, config_result, config_outputs) + + apply_preset_btn.click(fn=apply_preset, inputs=[preset_dd], outputs=config_outputs) + preset_dd.input(fn=apply_preset, inputs=[preset_dd], outputs=config_outputs) + reset_btn.click(fn=reset_minimind, outputs=config_outputs) + save_btn.click(fn=do_save, inputs=[config_state, cfg_path] + all_inputs, outputs=[save_status]) + load_btn.click(fn=do_load, inputs=[cfg_path], outputs=[save_status] + config_outputs) + + for box, value in zip(estimator_outs, estimator_values(initial_cfg)): + box.value = value return config_state, all_inputs, save_dir diff --git a/ui/tabs/data_tab.py b/ui/tabs/data_tab.py index 0311ef6..222600c 100644 --- a/ui/tabs/data_tab.py +++ b/ui/tabs/data_tab.py @@ -1,18 +1,20 @@ """ -ui/tabs/data_tab.py — Dataset inspector. +ui/tabs/data_tab.py — Dataset inspector for any task. -Streams a JSONL file (no full-load) and reports: - - record count +Streams a file, a directory or a glob (jsonl / json / parquet / csv / tsv / +txt; never a full load) and reports, without assuming what the data is for: + - record count and the record schemas found (chat, instruction, prompt / + response, preference, parallel pair, plain text, other), with examples + - the field inventory (dotted paths, types, how often each is present) - char-length distribution (min / mean / median / p95 / max) - - tokenized-length distribution if a tokenizer is available - - unique-record ratio (cheap dedup heuristic via line hash) - - first ~3 sample previews - - matplotlib histogram of token lengths + - script mix of the text (Han, Latin, Kana, Hangul, Cyrillic …) + - unique-record ratio (cheap dedup heuristic via record hash) + - media modalities (image / audio) and the first few previews + - a histogram of char lengths The point: before training, the user can verify the corpus is what they -expect — right field name, sane lengths, low duplication. Catches -"wrong tokenizer" / "all records empty" / "huge dup ratio" before they -burn 10 hours of training. +expect — right fields, sane lengths, low duplication — whatever the model +is meant to do. """ from __future__ import annotations @@ -20,71 +22,138 @@ import json import os import statistics +from collections import Counter from ui.gradio_compat import gr import chronos.deps # noqa: F401 from ui.i18n import t, register_translatable - -def _extract_text(rec: dict) -> str: - for k in ("text", "content", "story", "passage", "document", "article"): - if k in rec and rec[k]: - return str(rec[k]) - if "conversations" in rec and isinstance(rec["conversations"], list): - return " ".join(str(c.get("content", c.get("value", ""))) for c in rec["conversations"]) - if "messages" in rec and isinstance(rec["messages"], list): - return " ".join(str(m.get("content", "")) for m in rec["messages"]) - return " ".join(str(v) for v in rec.values() - if isinstance(v, (str, int, float))) +_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: str, max_records: int = 50_000): - """Stream up to `max_records` lines; return summary dict + length list.""" - if not path or not os.path.exists(path): - return {"error": f"file not found: {path!r}"}, [], [] - char_lens = [] - samples = [] - seen_hashes = set() - dup = 0 - n_total = 0 - with open(path, "rb") as f: - for i, line in enumerate(f): - if i >= max_records: - break - line = line.strip() - if not line: - continue - n_total += 1 - try: - rec = json.loads(line) - except json.JSONDecodeError: - continue - text = _extract_text(rec) - char_lens.append(len(text)) - h = hashlib.md5(line).hexdigest() - if h in seen_hashes: - dup += 1 - else: - seen_hashes.add(h) - if len(samples) < 3: - samples.append(text[:200]) + """Stream up to ``max_records`` records; return summary dict + length list + previews.""" + from chronos.tokenizer.extract import expand, iter_records, script_bucket + if not path or not path.strip(): + return {"error": "no path given"}, [], [] + path = os.path.expanduser(path.strip()) + files = [path] if os.path.isfile(path) else expand([path]) + if not files: + return {"error": f"no readable data files at: {path!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"}, [], samples - + return {"error": "no parseable records", "files": len(files), "read_errors": errors}, [], samples + ordered = sorted(char_lens) summary = { + "files": len(files), "scanned": n_total, - "parseable": len(char_lens), "dup_ratio": round(dup / max(n_total, 1), 4), - "char_min": min(char_lens), + "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(sorted(char_lens)[int(0.95 * len(char_lens))]), - "char_max": max(char_lens), - "file_size_mb": round(os.path.getsize(path) / (1024 * 1024), 1), - "note": (f"sampled first {max_records} records" - if n_total >= max_records else "full scan"), + "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 @@ -103,7 +172,7 @@ def _hist(lens): fig.patch.set_facecolor("#1a1a2e") ax.set_facecolor("#16213e") ax.hist(lens, bins=40, color="#39d98a", edgecolor="#1a1a2e") - ax.set_xlabel("chars per record", color="#aaa") + ax.set_xlabel("chars per record (all text fields)", color="#aaa") ax.set_ylabel("frequency", color="#aaa") ax.tick_params(colors="#aaa") for s in ax.spines.values(): @@ -119,7 +188,7 @@ def build_data_tab(): with gr.Row(): data_path = gr.Textbox( label=t("data.path"), scale=4, - placeholder="/Users/Fona/Downloads/Hybrid_LLM/Dataset/pretrain_t2t_mini.jsonl", + placeholder="./dataset/train.jsonl | ./data/ | ./data/*.parquet", ) scan_btn = gr.Button(t("data.scan"), variant="primary", scale=1) register_translatable(data_path, "data.path") diff --git a/ui/tabs/export_tab.py b/ui/tabs/export_tab.py index 49995f3..49502b9 100644 --- a/ui/tabs/export_tab.py +++ b/ui/tabs/export_tab.py @@ -59,6 +59,17 @@ def _export(model_path: str, output_dir: str, formats, config_path: str, "tensors": r.tensors, "expert_cache": _result_payload(r).get("expert_cache"), "cluster_build": _result_payload(r).get("cluster_build"), + "preserves_recurrent_depth": r.format != "ollama-compatible-gguf", + "preserves_lazy_expert_paging": r.format != "ollama-compatible-gguf", + "stock_ollama_executable": bool( + _result_payload(r).get("compatibility_notice", {}).get( + "stock_runtime_executable", False + ) + ), + "requires_chronos_backend": True, + "multimodal_encoders_embedded": bool( + _result_payload(r).get("chronos", {}).get("modalities") + ), } for r in results } @@ -66,7 +77,9 @@ def _export(model_path: str, output_dir: str, formats, config_path: str, "\n\nCompatibility note: safetensors and GGUF outputs are standard containers " "with Chronos architecture metadata. Stock Ollama/llama.cpp builds need a " "Chronos architecture adapter to execute hybrid attention, lookahead MoE " - "prediction, and lazy expert loading correctly." + "prediction, and lazy expert loading correctly. The current " + "ollama-compatible artifact is explicitly manifest-only and is not " + "reported as stock-runtime executable." ) return format_export_report(results) + compatibility, summary except Exception as exc: diff --git a/ui/tabs/inference_tab.py b/ui/tabs/inference_tab.py index c28c985..5727d87 100644 --- a/ui/tabs/inference_tab.py +++ b/ui/tabs/inference_tab.py @@ -1,8 +1,9 @@ """ ui/tabs/inference_tab.py — Real-time generation with tokens/s display """ -import time import math +import os +import time from ui.gradio_compat import gr import pandas as pd @@ -42,10 +43,18 @@ ("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"), @@ -248,6 +257,15 @@ def _apply_offload_expert_budget(engine, model_cfg, ram_load_ratio: float | str 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 @@ -263,8 +281,8 @@ def _replace_torch_experts_with_placeholders(model, model_cfg) -> None: 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. """ - import torch 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) @@ -276,7 +294,9 @@ def _replace_torch_experts_with_placeholders(model, model_cfg) -> None: if not isinstance(moe, ChronosMOEFeedForward): continue for idx in range(len(moe.experts)): - moe.experts[idx] = LazyExpertPlaceholder(hidden, intermediate, torch.float16) + # 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: @@ -297,14 +317,22 @@ def _replace_mlx_experts_with_placeholders(model) -> None: 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.numpy import save_file as np_save_file + 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)): @@ -315,7 +343,7 @@ def _populate_mlx_cache_from_state_dict(store, state_dict: dict) -> None: key = prefix + pname if key in state_dict: tensors[f"l{li}_e{eid}{KEY_SEP}{pname}"] = ( - state_dict[key].detach().to(dtype=torch.float16, device="cpu").numpy() + state_dict[key].detach().to(dtype=page_dtype, device="cpu").contiguous() ) if tensors: file_name = f"cluster_{eid}.ctsr" @@ -342,10 +370,11 @@ def _populate_mlx_cache_from_reader(store, reader) -> None: import os import json import torch - from safetensors.numpy import save_file as np_save_file + 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 = {} @@ -357,7 +386,7 @@ def _populate_mlx_cache_from_reader(store, reader) -> None: key = prefix + pname if key in keys: tensors[f"l{li}_e{eid}{KEY_SEP}{pname}"] = ( - reader.get_tensor(key).detach().to(dtype=torch.float16, device="cpu").numpy() + reader.get_tensor(key).detach().to(dtype=page_dtype, device="cpu").contiguous() ) if tensors: file_name = f"cluster_{eid}.ctsr" @@ -511,7 +540,8 @@ def _run_torch_inference(backend: str, model_cfg, model_path_val: str, token_ids 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) -> tuple[list[int], dict]: + cpu_budget_percent: int | float | str | None = 100, + speculative: dict | None = None) -> tuple[list[int], dict]: import os import importlib import gc @@ -552,7 +582,9 @@ def _run_torch_inference(backend: str, model_cfg, model_path_val: str, token_ids base_load_mem = _memory_snapshot_gb() device_map = {"cuda": "cuda", "xpu": "xpu", "mps": "mps", "cpu": "cpu"}[backend] - model = model.to(device_map).eval() + 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) @@ -604,15 +636,28 @@ def _run_torch_inference(backend: str, model_cfg, model_path_val: str, token_ids setup_stats[f"{prefix}_setup_delta_gb"] = _memory_delta(backend_mem_setup, backend_mem0, key) try: - 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, - ) + 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() @@ -655,7 +700,8 @@ def _run_torch_full_dram(backend: str, model_cfg, model_path_val: str, token_ids 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) -> tuple[list[int], dict]: + cpu_budget_percent: int | float | str | None = 100, + speculative: dict | None = None) -> tuple[list[int], dict]: import os import time as _time import gc @@ -684,7 +730,9 @@ def _run_torch_full_dram(backend: str, model_cfg, model_path_val: str, token_ids del weights gc.collect() weight_load_mem = _memory_snapshot_gb() - model = model.to(device_map).eval() + 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) @@ -700,16 +748,46 @@ def _run_torch_full_dram(backend: str, model_cfg, model_path_val: str, token_ids 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(int(max_tok)): - past_len = past_key_values[0][0].shape[1] if past_key_values else 0 - outputs, _ = model( + 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( @@ -777,6 +855,8 @@ def _run_torch_full_dram(backend: str, model_cfg, model_path_val: str, token_ids "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"), @@ -792,7 +872,9 @@ def _run_torch_full_dram(backend: str, model_cfg, model_path_val: str, token_ids def _run_mlx_full_dram(model_cfg, model_path_val: str, token_ids: list[int], - max_tok: int, temp: float, do_sample: bool = True) -> tuple[list[int], dict]: + 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 @@ -822,7 +904,9 @@ def _run_mlx_full_dram(model_cfg, model_path_val: str, token_ids: list[int], weights = {} gc.collect() weight_load_mem = _memory_snapshot_gb() - mlx_model = ChronosMLXModel.from_chronos_pytorch(pt_model, model_cfg) + 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() @@ -837,6 +921,9 @@ def _run_mlx_full_dram(model_cfg, model_path_val: str, token_ids: list[int], 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() @@ -887,7 +974,9 @@ def _run_mlx_full_dram(model_cfg, model_path_val: str, token_ids: list[int], 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) -> tuple[list[int], dict]: + quality_safe: bool = True, + eos_token_id: int | None = None, + speculative: dict | None = None) -> tuple[list[int], dict]: import os import gc import torch @@ -925,7 +1014,10 @@ def _run_mlx_lazy_inference(model_cfg, model_path_val: str, token_ids: list[int] ) _replace_torch_experts_with_placeholders(pt_model, model_cfg) model_init_mem = _memory_snapshot_gb() - mlx_model = ChronosMLXModel.from_chronos_pytorch(pt_model, model_cfg, include_experts=False) + 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() @@ -956,6 +1048,9 @@ def _run_mlx_lazy_inference(model_cfg, model_path_val: str, token_ids: list[int] 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 {}), )) stats = dict(getattr(engine, "last_stats", {})) store_stats = engine.store.stats() @@ -1089,6 +1184,43 @@ def _normalize_miss_policy(value) -> str: 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._" @@ -1128,6 +1260,14 @@ def fmt_bool(row: dict, key: str) -> str: ("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")), @@ -1146,6 +1286,13 @@ def fmt_bool(row: dict, key: str) -> str: ("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")), @@ -1261,6 +1408,14 @@ def _row_from_stats(mode: str, backend: str, tokens: list[int], elapsed: float, "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), @@ -1367,54 +1522,11 @@ def _build_model_cfg(cfg, model_path_val): source_text = "; ".join(cfg_sources) if cfg_sources else model_path_val sniff_note = f"[checkpoint config: {summary} | source={source_text}]\n" - model_cfg_kwargs = { - "hidden_size": sniffed.get("hidden_size", cfg.get("hidden_size", 768)), - "num_hidden_layers": sniffed.get("num_hidden_layers", cfg.get("num_hidden_layers", 8)), - "num_experts": sniffed.get("num_experts", cfg.get("num_experts", 4)), - "num_experts_per_tok": sniffed.get("num_experts_per_tok", cfg.get("num_experts_per_tok", 1)), - "num_shared_experts": sniffed.get("num_shared_experts", cfg.get("num_shared_experts", 1)), - "lookahead_steps": sniffed.get("lookahead_steps", cfg.get("lookahead_steps", 2)), - "kv_latent_dim": sniffed.get("kv_latent_dim", cfg.get("kv_latent_dim", 64)), - "sliding_window_size": sniffed.get("sliding_window_size", cfg.get("sliding_window_size", 2048)), - "use_hybrid_attention": sniffed.get("use_hybrid_attention", cfg.get("use_hybrid_attention", True)), - "vram_budget_gb": sniffed.get("vram_budget_gb", cfg.get("vram_budget_gb", 4.0)), - "pinned_memory_max_fraction": sniffed.get("pinned_memory_max_fraction", cfg.get("pinned_memory_max_fraction", 0.25)), - "storage_format": sniffed.get("storage_format", cfg.get("storage_format", "safetensors")), - "offload_miss_policy_default": sniffed.get("offload_miss_policy_default", cfg.get("offload_miss_policy_default", "on_demand")), - "recommended_resident_experts": sniffed.get("recommended_resident_experts", cfg.get("recommended_resident_experts", None)), - "use_moe": True, - } - if "cluster_manifest_path" in sniffed: - model_cfg_kwargs["cluster_manifest_path"] = sniffed["cluster_manifest_path"] - if "rope_dim" in sniffed: - model_cfg_kwargs["rope_dim"] = sniffed["rope_dim"] - if "num_attention_heads" in sniffed: - model_cfg_kwargs["num_attention_heads"] = sniffed["num_attention_heads"] - if "num_key_value_heads" in sniffed: - model_cfg_kwargs["num_key_value_heads"] = sniffed["num_key_value_heads"] - if "vocab_size" in sniffed: - model_cfg_kwargs["vocab_size"] = sniffed["vocab_size"] - if "moe_intermediate_size" in sniffed: - model_cfg_kwargs["intermediate_size"] = sniffed["moe_intermediate_size"] - model_cfg_kwargs["moe_intermediate_size"] = sniffed["moe_intermediate_size"] - - for opt_key, cfg_key in [ - ("num_attention_heads", "num_attention_heads"), - ("num_key_value_heads", "num_key_value_heads"), - ("rope_dim", "rope_dim"), - ("max_position_embeddings", "max_seq_len"), - ("tie_word_embeddings", "tie_word_embeddings"), - ]: - if opt_key in {"num_attention_heads", "num_key_value_heads", "rope_dim"} \ - and opt_key in model_cfg_kwargs: - continue - if cfg_key not in cfg: - continue - val = cfg[cfg_key] - if val in (None, "", 0): - continue - model_cfg_kwargs[opt_key] = val - + 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)) @@ -1430,10 +1542,24 @@ def _build_model_cfg(cfg, model_path_val): 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") -> dict: + miss_policy="on_demand", speculative=False, draft_depth=1, + gamma=4) -> dict: from transformers import AutoTokenizer from chronos.backend.mac_diagnostics import mac_backend_diagnostics @@ -1445,7 +1571,8 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, 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(chronos.deps.get_tokenizer_path()) + 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) @@ -1467,6 +1594,7 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, **_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)] @@ -1490,6 +1618,8 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, 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, ) else: lazy_tokens, lazy_stats = _run_torch_inference( @@ -1506,6 +1636,7 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, miss_policy=miss_policy, cpu_threads=cpu_threads_cfg, cpu_budget_percent=cpu_budget_cfg, + speculative=spec_opts, ) lazy_elapsed = time.monotonic() - t0 lazy_mem = _memory_snapshot_gb() @@ -1515,6 +1646,7 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, "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" @@ -1539,6 +1671,8 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, 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( @@ -1553,6 +1687,7 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, 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() @@ -1562,6 +1697,7 @@ def _run_inference_modes(cfg, selected_backend, selected_mode, model_path_val, "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 @@ -1595,12 +1731,12 @@ 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") -> dict: + miss_policy="on_demand", speculative=False, draft_depth=1, gamma=4) -> 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, + miss_policy, speculative, draft_depth, gamma, ) except Exception as e: import traceback @@ -1662,6 +1798,15 @@ def build_inference_tab(config_state: gr.State): register_translatable(miss_policy, "infer.miss_policy") register_translatable(ram_load_ratio, "infer.ram_load_ratio") register_translatable(sweep_ratios, "infer.sweep_ram_load_ratios") + with gr.Row(): + speculative = gr.Checkbox(value=False, label=t("infer.speculative")) + draft_depth = gr.Slider(1, 8, value=1, step=1, label=t("infer.draft_depth")) + spec_gamma = gr.Slider(1, 8, value=4, step=1, label=t("infer.spec_gamma")) + register_translatable(speculative, "infer.speculative") + register_translatable(draft_depth, "infer.draft_depth") + register_translatable(spec_gamma, "infer.spec_gamma") + spec_note = gr.Markdown(t("infer.spec_note")) + register_translatable(spec_note, "infer.spec_note") gen_btn = gr.Button(t("infer.generate"), variant="primary") register_translatable(gen_btn, "infer.generate") @@ -1810,98 +1955,11 @@ def _sniff_checkpoint(path: str) -> dict: except Exception: return {} - 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 {}) - 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" - - model_cfg_kwargs = { - "hidden_size": sniffed.get("hidden_size", cfg.get("hidden_size", 768)), - "num_hidden_layers": sniffed.get("num_hidden_layers", cfg.get("num_hidden_layers", 8)), - "num_experts": sniffed.get("num_experts", cfg.get("num_experts", 4)), - "num_experts_per_tok": sniffed.get("num_experts_per_tok", cfg.get("num_experts_per_tok", 1)), - "num_shared_experts": sniffed.get("num_shared_experts", cfg.get("num_shared_experts", 1)), - "lookahead_steps": sniffed.get("lookahead_steps", cfg.get("lookahead_steps", 2)), - "kv_latent_dim": sniffed.get("kv_latent_dim", cfg.get("kv_latent_dim", 64)), - "sliding_window_size": sniffed.get("sliding_window_size", cfg.get("sliding_window_size", 2048)), - "use_hybrid_attention": sniffed.get("use_hybrid_attention", cfg.get("use_hybrid_attention", True)), - "vram_budget_gb": sniffed.get("vram_budget_gb", cfg.get("vram_budget_gb", 4.0)), - "pinned_memory_max_fraction": sniffed.get("pinned_memory_max_fraction", cfg.get("pinned_memory_max_fraction", 0.25)), - "storage_format": sniffed.get("storage_format", cfg.get("storage_format", "safetensors")), - "use_moe": True, - } - if "cluster_manifest_path" in sniffed: - model_cfg_kwargs["cluster_manifest_path"] = sniffed["cluster_manifest_path"] - if "rope_dim" in sniffed: - model_cfg_kwargs["rope_dim"] = sniffed["rope_dim"] - if "num_attention_heads" in sniffed: - model_cfg_kwargs["num_attention_heads"] = sniffed["num_attention_heads"] - if "num_key_value_heads" in sniffed: - model_cfg_kwargs["num_key_value_heads"] = sniffed["num_key_value_heads"] - if "vocab_size" in sniffed: - model_cfg_kwargs["vocab_size"] = sniffed["vocab_size"] - if "moe_intermediate_size" in sniffed: - model_cfg_kwargs["intermediate_size"] = sniffed["moe_intermediate_size"] - model_cfg_kwargs["moe_intermediate_size"] = sniffed["moe_intermediate_size"] - - for opt_key, cfg_key in [ - ("num_attention_heads", "num_attention_heads"), - ("num_key_value_heads", "num_key_value_heads"), - ("rope_dim", "rope_dim"), - ("max_position_embeddings", "max_seq_len"), - ("tie_word_embeddings", "tie_word_embeddings"), - ]: - if opt_key in {"num_attention_heads", "num_key_value_heads", "rope_dim"} \ - and opt_key in model_cfg_kwargs: - continue - if cfg_key not in cfg: - continue - val = cfg[cfg_key] - if val in (None, "", 0): - continue - model_cfg_kwargs[opt_key] = val - - 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 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): + miss_policy_val, ram_load_ratio_val, sweep_ratios_val, + speculative_val=False, draft_depth_val=1, gamma_val=4): raw = generate_api( cfg, selected_backend=selected_backend, @@ -1914,6 +1972,9 @@ def generate_unified(cfg, selected_backend, selected_mode, model_path_val, ram_load_ratio=ram_load_ratio_val, sweep_ram_load_ratios=sweep_ratios_val, miss_policy=miss_policy_val, + speculative=speculative_val, + draft_depth=draft_depth_val, + gamma=gamma_val, ) if raw.get("error"): text = f"Error: {raw.get('error')}" @@ -1957,10 +2018,10 @@ 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, ], outputs=[output_box, lazy_output_box, full_output_box, stats_md, stats_chart, raw_stats], api_name=None, - api_visibility="private", ) api_btn = gr.Button("api_generate", visible=False) @@ -1970,6 +2031,7 @@ 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, ], outputs=[api_out], api_name="generate", diff --git a/ui/tabs/monitor_tab.py b/ui/tabs/monitor_tab.py new file mode 100644 index 0000000..d8e7da2 --- /dev/null +++ b/ui/tabs/monitor_tab.py @@ -0,0 +1,425 @@ +"""ui/tabs/monitor_tab.py — Training monitor inside the WebUI console. + +Reads the same run directories as the standalone monitor page +(``chronos monitor``) through :class:`chronos.monitor.RunMonitor`: run spec +and training parameters, heartbeat/derived status, progress and ETA, every +logged metric, validation events, whatever evaluation metrics the run logs +(discovered, never assumed), per-tensor statistics, MoE routing, samples of +any task, log tail and checkpoints. Read-only; it never touches the trainer. +""" +from __future__ import annotations + +import html +import json +import math +import os +import threading +import time + +from ui.gradio_compat import gr +from ui.i18n import get_current_lang, register_translatable, t + +try: # pandas ships with gradio; the slim CI shim may lack it + import pandas as pd +except ModuleNotFoundError: # pragma: no cover + pd = None + +DEFAULT_RUNS_DIR = os.environ.get("CHRONOS_RUNS_DIR", "runs") +_MONITORS: dict[str, object] = {} +_LOCK = threading.Lock() + +# Line plots: (title key, [(series key, label key)...]); one y-axis per plot. +PLOTS = { + "overview": [ + ("monitor.p_ce", [("primary_loss", "monitor.s_train"), ("valid_ce", "monitor.s_valid")]), + ("monitor.p_loss", [("loss", "monitor.s_total")]), + ("monitor.p_lr", [("lr", "monitor.s_lr")]), + ("monitor.p_grad", [("grad_norm", "monitor.s_grad")]), + ("monitor.p_tps", [("tokens_per_s", "monitor.s_counted"), ("processed_tokens_per_s", "monitor.s_processed")]), + ("monitor.p_mem", [("peak_mem_gb", "monitor.s_peak"), ("active_mem_gb", "monitor.s_active"), + ("cache_mem_gb", "monitor.s_cache")]), + ], + "recurrent": [ + ("monitor.p_depth", [("rec_depth_mean", "monitor.s_depth")]), + ("monitor.p_contraction", [("rec_contraction", "monitor.s_contraction")]), + ("monitor.p_exit", [("rec_early_exit_ratio", "monitor.s_exit")]), + ("monitor.p_route_entropy", [("route_entropy", "monitor.s_entropy")]), + ("monitor.p_load_cv", [("load_cv", "monitor.s_load_cv")]), + ("monitor.p_max_load", [("max_expert_load", "monitor.s_max_load")]), + ], +} +LOSS_TERMS = ("primary_loss", "balance_loss", "lookahead_loss", "lookahead_topk_loss", "lookahead_union_loss", + "temporal_loss", "locality_loss", "anchor_loss", "recurrent_ponder_loss", "recurrent_distill_loss", + "recurrent_stability_loss", "recurrent_contraction_loss") +STATE_ICON = {"running": "●", "paused": "‖", "evaluating": "◐", "checkpointing": "◐", "starting": "○", + "finished": "✓", "failed": "✕", "stopped": "■", "pending": "○"} +# Sequential blue ramp (light→dark), steps 150..700 of the reference palette. +SEQ = ("#cde2fb", "#9ec5f4", "#6da7ec", "#3987e5", "#256abf", "#184f95", "#0d366b") +SEQ_INK = ("#0b0b0b", "#0b0b0b", "#0b0b0b", "#ffffff", "#ffffff", "#ffffff", "#ffffff") + +MONITOR_CSS = """ +.chronos-mon{font-family:system-ui,-apple-system,"Segoe UI","PingFang SC",sans-serif;font-size:13px} +.chronos-mon .row{display:flex;gap:10px;flex-wrap:wrap;align-items:center;margin-bottom:10px} +.chronos-mon .pill{display:inline-flex;gap:6px;align-items:center;padding:3px 10px;border-radius:999px; + border:1px solid var(--border-color-primary);background:var(--background-fill-secondary);font-size:12px} +.chronos-mon .muted{color:var(--body-text-color-subdued)} +.chronos-mon .bar{height:8px;border-radius:4px;background:var(--background-fill-secondary);overflow:hidden;margin:6px 0 4px} +.chronos-mon .bar>span{display:block;height:100%;background:#2a78d6;border-radius:4px} +.chronos-mon .tiles{display:grid;grid-template-columns:repeat(auto-fill,minmax(150px,1fr));gap:8px} +.chronos-mon .tile{border:1px solid var(--border-color-primary);border-radius:8px;padding:10px 12px;background:var(--block-background-fill)} +.chronos-mon .tile b{display:block;font-size:20px;font-weight:600;margin-top:4px;font-variant-numeric:tabular-nums} +.chronos-mon .tile small{color:var(--body-text-color-subdued);font-size:11px} +.chronos-mon .note{border:1px solid var(--border-color-primary);border-radius:8px;padding:8px 12px;margin-bottom:10px} +.chronos-mon .warn{border-color:#fab219}.chronos-mon .err{border-color:#d03b3b} +.chronos-mon table.heat{border-collapse:separate;border-spacing:2px;font-size:10px;width:auto;margin:6px 0 0} +.chronos-mon table.heat,.chronos-mon table.heat tr,.chronos-mon table.heat td,.chronos-mon table.heat th{border:0!important;background-clip:padding-box} +.chronos-mon table.heat td{min-width:16px;height:18px;border-radius:2px;text-align:center;padding:0 2px;font-variant-numeric:tabular-nums} +.chronos-mon table.heat th{font-weight:400;color:var(--body-text-color-subdued);padding:0 4px;white-space:nowrap} +.chronos-mon .samp{border-bottom:1px solid var(--border-color-primary);padding:8px 0;line-height:1.6} +""" + + +# ---- data access ------------------------------------------------------------- +def _monitor(runs_dir: str): + from chronos.monitor import RunMonitor + root = os.path.abspath(os.path.expanduser(runs_dir or DEFAULT_RUNS_DIR)) + with _LOCK: + mon = _MONITORS.get(root) + if mon is None: + if len(_MONITORS) > 8: + _MONITORS.clear() + mon = _MONITORS[root] = RunMonitor(root) + return mon + + +def _f(v, digits: int = 4) -> str: + if isinstance(v, bool) or not isinstance(v, (int, float)) or not math.isfinite(v): + return "—" + a = abs(v) + if a and (a < 1e-4 or a >= 1e6): + return f"{v:.2e}" + if isinstance(v, int): + return f"{v:,}" + return f"{v:.{digits}g}" + + +def _count(v) -> str: + if not isinstance(v, (int, float)): + return "—" + for div, suf in ((1e9, "B"), (1e6, "M"), (1e3, "k")): + if abs(v) >= div: + return f"{v / div:.1f}{suf}" + return str(round(v)) + + +def _dur(sec) -> str: + if not isinstance(sec, (int, float)): + return "—" + sec = max(0, int(sec)) + d, h, m = sec // 86400, sec % 86400 // 3600, sec % 3600 // 60 + return f"{d}d {h}h" if d else f"{h}h {m}m" if h else f"{m}m" + + +def _esc(v) -> str: + return html.escape("" if v is None else str(v), quote=True) + + +def _df(rows: list[dict], columns: list[tuple[str, str]], numeric: tuple[str, ...] = ()): + """``columns``: (row key, i18n header key); ``numeric`` keys are pre-formatted (tiny values as 1.23e-06).""" + data = [[_f(r.get(k)) if k in numeric and r.get(k) is not None else r.get(k) for k, _ in columns] for r in rows] + if pd is None: + return data + return pd.DataFrame(data, columns=[t(h) for _, h in columns]) + + +def _series_df(parts: list[tuple[str, list[dict], str]]): + """Long format for gr.LinePlot: ``parts`` = (label, rows, metric key).""" + recs = [] + for label, rows, key in parts: + for r in rows: + v, step = r.get(key), r.get("step") + if isinstance(v, (int, float)) and isinstance(step, (int, float)) and math.isfinite(v): + recs.append({"step": step, "value": float(v), "series": label}) + if pd is None: + return recs + return pd.DataFrame(recs, columns=["step", "value", "series"]) + + +PALETTE = ("#2a78d6", "#eb6834", "#1baf7a") + + +def _plot_update(parts): + """Fixed colour per series slot (blue, orange, green). Multi-series plots + name their colours in the translated label; the tooltip carries the name.""" + labels = list(dict.fromkeys(label for label, _, _ in parts)) + return gr.update(value=_series_df(parts), color_map={lab: PALETTE[i % 3] for i, lab in enumerate(labels)}) + + +# ---- HTML fragments (all dynamic text escaped) ---------------------------------- +def _heat_html(title: str, labels: list, grid: list, fmt=lambda v: _f(v, 3), col_label="E") -> str: + flat = [v for row in grid or [] for v in row if isinstance(v, (int, float))] + if not flat: + return f"
{_esc(title)}

{_esc(t('monitor.none'))}

" + lo, hi = min(flat), max(flat) + cols = max(len(r) for r in grid) + head = "".join(f"{j}" for j in range(cols)) + body = [] + for i, row in enumerate(grid): + cells = [] + for j in range(cols): + v = row[j] if j < len(row) else None + if not isinstance(v, (int, float)): + cells.append("") + continue + b = min(6, int((v - lo) / (hi - lo) * 7)) if hi > lo else 3 + tip = f"{labels[i] if i < len(labels) else i} · {col_label}{j}: {fmt(v)}" + cells.append(f"") + body.append(f"{_esc(labels[i] if i < len(labels) else i)}{''.join(cells)}") + ramp = "".join(f"" for c in SEQ) + return (f"
{_esc(title)}{head}{''.join(body)}
" + f"
{_esc(fmt(lo))} {ramp} {_esc(fmt(hi))}
") + + +def _status_html(snap: dict) -> str: + if not snap.get("run"): + return f"
{_esc(t('monitor.no_runs'))} {_esc(snap.get('root'))}
" + st, p, latest = snap.get("status") or {}, snap.get("progress") or {}, snap.get("latest") or {} + state = st.get("state", "pending") + notes = [] + if st.get("error"): + notes.append(f"
✕ {_esc(t('monitor.error'))}: {_esc(st['error'])}
") + if state == "paused": + notes.append(f"
‖ {_esc(t('monitor.paused'))}
") + if st.get("stale"): + notes.append(f"
! {_esc(t('monitor.stale').format(age=_dur(st.get('last_update_age_s'))))}
") + if st.get("source") == "derived" and not snap.get("parameters"): + notes.append(f"
i {_esc(t('monitor.derived'))}
") + frac = p.get("phase_fraction") + head = (f"
{STATE_ICON.get(state, '○')} {_esc(t('monitor.state.' + state))}" + f"{_esc(snap['run'])}{_esc(p.get('current') or '—')} · {_esc(t('monitor.step'))} " + f"{_esc(_f(latest.get('step')))} / {_esc(_f(latest.get('total')))}" + f"{' · ' + format(frac * 100, '.1f') + '%' if isinstance(frac, (int, float)) else ''}" + f"{_esc(t('monitor.eta'))} {_esc((p.get('eta') or '—').replace('T', ' '))} · " + f"{_esc(t('monitor.remaining'))} {_esc(_dur(p.get('run_remaining_s')))} · {_esc(t('monitor.budget'))} " + f"{_esc(p.get('budget_hours'))} h
") + bars = [] + for ph in p.get("phases") or []: + f = 1.0 if ph["state"] == "complete" else (frac or 0.0) if ph["state"] == "active" else 0.0 + bars.append(f"
" + f"{_esc(ph['name'])} {_esc(t('monitor.ph.' + ph['state']))} · {_esc(ph.get('budget_h'))} h
") + best = (snap.get("best_valid") or {}).get(p.get("current")) or {} + mem = snap.get("memory") or {} + spec = snap.get("spec") or {} + cfg = spec.get("model_config") or {} + tiles = [ + ("monitor.k_loss", _f(latest.get("loss")), f"{t('monitor.step')} {_f(latest.get('step'))}"), + ("monitor.k_ce", _f(latest.get("primary_loss")), + f"ppl {_f(math.exp(latest['primary_loss']), 3)}" if isinstance(latest.get("primary_loss"), (int, float)) else ""), + ("monitor.k_valid", _f(best.get("valid_ce")), f"{t('monitor.step')} {_f(best.get('step'))}" if best else ""), + ("monitor.k_tps", f"{_count(latest.get('tokens_per_s'))} tok/s", f"{_count(latest.get('processed_tokens_per_s'))} {t('monitor.s_processed')}"), + ("monitor.k_mem", f"{_f(mem.get('footprint_gb') or latest.get('peak_mem_gb'), 3)} GB", + f"MLX peak {_f(latest.get('peak_mem_gb'), 3)} · active {_f(latest.get('active_mem_gb'), 3)} GB"), + ("monitor.k_grad", _f(latest.get("grad_norm")), ""), + ("monitor.k_lr", _f(latest.get("lr"), 3), ""), + ("monitor.k_depth", _f(latest.get("rec_depth_mean"), 3), f"max {_esc(cfg.get('recurrent_depth', '—'))}"), + ("monitor.k_route", _f(latest.get("route_entropy"), 3), f"max {_f(latest.get('route_entropy_max'), 3)} · CV {_f(latest.get('load_cv'), 3)}"), + ] + tiles_html = "".join(f"
{_esc(t(k))}{_esc(v)}{_esc(sub)}
" for k, v, sub in tiles) + foot = (f"

{_esc(t('monitor.source'))}: {_esc(st.get('source'))} · pid {_esc(st.get('pid'))} " + f"({_esc(st.get('ps_state') or '-')}) · {_esc(time.strftime('%H:%M:%S'))}

") + return (f"
{''.join(notes)}{head}
{''.join(bars)}
" + f"
{tiles_html}
{foot}
") + + +def _events_rows(snap: dict) -> list[dict]: + out = [] + for e in reversed(snap.get("events") or []): + detail = {k: v for k, v in e.items() if k not in ("time", "event", "phase", "step")} + out.append({"time": e.get("time"), "event": e.get("event"), "phase": e.get("phase"), "step": e.get("step"), + "valid_ce": e.get("valid_ce"), "detail": json.dumps(detail, ensure_ascii=False)[:400]}) + return out + + +def _samples_html(snap: dict) -> str: + """Samples of any task: the snapshot normalises them to input/reference/output + tags.""" + rows = list(reversed(snap.get("samples") or []))[:12] + if not rows: + return f"
{_esc(t('monitor.none'))}
" + items = [] + for r in rows: + tags = " · ".join(f"{k}={v}" for k, v in (r.get("tags") or {}).items()) + head = " · ".join(x for x in (_esc(r.get("phase")), f"{_esc(t('monitor.step'))} {_esc(r.get('step'))}", _esc(tags)) if x) + body = "".join(f"
{_esc(t(key))} {_esc(r.get(role))}" + for role, key in (("input", "monitor.sample_input"), ("reference", "monitor.sample_ref"), + ("output", "monitor.sample_output")) if r.get(role) not in (None, "")) + items.append(f"
{head}{body}
") + return f"
{''.join(items)}
" + + +def _eval_family(snap: dict, metric: str | None) -> tuple[list[str], str | None, dict | None]: + """Metric families the run logs (discovered, not assumed) and the chosen one.""" + families = (snap.get("evals") or {}).get("families") or [] + names = [f["metric"] for f in families] + chosen = metric if metric in names else (names[0] if names else None) + return names, chosen, next((f for f in families if f["metric"] == chosen), None) + + +# ---- refresh ------------------------------------------------------------------------- +def _phase_rows(snap: dict, phase: str | None) -> tuple[str | None, list[dict]]: + hist = snap.get("history") or {} + phase = phase if phase in hist else (snap.get("progress") or {}).get("current") + if phase not in hist: + phase = next(reversed(hist), None) if hist else None + return phase, hist.get(phase, []) if phase else [] + + +def _loss_parts(rows: list[dict], term: str): + term = term if term in LOSS_TERMS else "primary_loss" + return [(term, rows, term)] + + +def refresh(runs_dir: str, run: str | None, phase: str | None, loss_term: str, eval_metric: str | None = None): + """Everything the tab shows, in output order (see ``build_monitor_tab``).""" + mon = _monitor(runs_dir) + try: + snap = mon.snapshot(run or None) + except KeyError: + snap = mon.snapshot(None) + names = [r["name"] for r in snap.get("runs") or []] + run_update = gr.update(choices=names, value=snap.get("run")) + phases = list((snap.get("history") or {}).keys()) + phase, rows = _phase_rows(snap, phase) + phase_update = gr.update(choices=phases, value=phase) + valid = [e for e in snap.get("events") or [] if e.get("event") == "validation" and e.get("phase") == phase] + plots = [] + for group in ("overview", "recurrent"): + for _title, series in PLOTS[group]: + parts = [(t(label), valid if key == "valid_ce" else rows, key) for key, label in series] + plots.append(_plot_update(parts)) + present = [k for k in LOSS_TERMS if any(isinstance(r.get(k), (int, float)) for r in rows)] or ["primary_loss"] + term = loss_term if loss_term in present else present[0] + loss_update = _plot_update(_loss_parts(rows, term)) + term_update = gr.update(choices=present, value=term) + + params = snap.get("parameters") or {} + routing = params.get("routing") or {} + experts = params.get("experts") or {} + groups = [{"group": k, **v} for k, v in (params.get("groups") or {}).items()] + if params.get("stats"): + groups.append({"group": "total", **params["stats"], "tensors": params.get("tensor_count")}) + group_df = _df(groups, [("group", "monitor.c_group"), ("tensors", "monitor.c_tensors"), ("numel", "monitor.c_numel"), + ("rms", "monitor.c_rms"), ("grad_norm", "monitor.c_grad_norm"), ("update_ratio", "monitor.c_update")], + numeric=("rms", "grad_norm", "update_ratio")) + full = mon.parameters(snap.get("run")) if params else None + tensors = (full or {}).get("tensors") or [] + tensor_df = _df([{**x, "shape": "×".join(map(str, x.get("shape") or []))} for x in tensors], + [("name", "monitor.c_name"), ("group", "monitor.c_group"), ("shape", "monitor.c_shape"), ("dtype", "monitor.c_dtype"), + ("rms", "monitor.c_rms"), ("std", "monitor.c_std"), ("absmax", "monitor.c_absmax"), + ("grad_rms", "monitor.c_grad_rms"), ("update_ratio", "monitor.c_update")], + numeric=("rms", "std", "absmax", "grad_rms", "update_ratio")) + tel_note = (t("monitor.telemetry_at").format(step=params.get("step"), ms=round((params.get("telemetry_s") or 0) * 1000)) + if params else t("monitor.no_telemetry")) + load_html = _heat_html(t("monitor.h_load"), routing.get("labels") or [], routing.get("load") or [], + fmt=lambda v: f"{v * 100:.1f}%") + imp_html = _heat_html(t("monitor.h_importance"), routing.get("labels") or [], routing.get("importance") or []) + layers = [f"L{x}" for x in experts.get("layers") or []] + wrms_html = _heat_html(t("monitor.h_wrms"), layers, experts.get("weight_rms") or []) + grms_html = _heat_html(t("monitor.h_grms"), layers, experts.get("grad_rms") or []) + + events_df = _df(_events_rows(snap), [("time", "monitor.c_time"), ("event", "monitor.c_event"), ("phase", "monitor.c_phase"), + ("step", "monitor.step"), ("valid_ce", "monitor.c_valid_ce"), ("detail", "monitor.c_detail")]) + ckpt_df = _df(snap.get("checkpoints") or [], [("file", "monitor.c_file"), ("size_gb", "monitor.c_size"), + ("modified", "monitor.c_time")]) + names, chosen, family = _eval_family(snap, eval_metric) + eval_parts = [(sr["label"], valid, sr["key"]) for sr in (family or {}).get("series", [])[:len(PALETTE)]] + eval_update = gr.update(choices=names, value=chosen) + eval_plot = _plot_update(eval_parts) + training_params = mon.training_parameters(snap.get("run")) if snap.get("run") else {} + log = mon.log_tail(snap.get("run"), 24000) if snap.get("run") else "" + return (run_update, phase_update, _status_html(snap), *plots, term_update, loss_update, + tel_note, group_df, tensor_df, load_html, imp_html, wrms_html, grms_html, + eval_update, eval_plot, events_df, _samples_html(snap), training_params, log, ckpt_df) + + +# ---- layout ------------------------------------------------------------------------------ +def _labelled(component, key: str): + register_translatable(component, key) + return component + + +def build_monitor_tab(): + with gr.Tab(t("tab.monitor")) as tab: + register_translatable(tab, "tab.monitor") + gr.Markdown(t("monitor.intro")) + with gr.Row(): + runs_dir = _labelled(gr.Textbox(value=DEFAULT_RUNS_DIR, label=t("monitor.runs_dir"), scale=3), "monitor.runs_dir") + run = _labelled(gr.Dropdown(choices=[], value=None, label=t("monitor.run"), scale=2, allow_custom_value=True), "monitor.run") + phase = _labelled(gr.Dropdown(choices=[], value=None, label=t("monitor.phase"), scale=1, allow_custom_value=True), "monitor.phase") + auto = _labelled(gr.Checkbox(value=True, label=t("monitor.auto"), scale=1), "monitor.auto") + refresh_btn = _labelled(gr.Button(t("monitor.refresh"), scale=1), "monitor.refresh") + status = gr.HTML() + plots = [] + with gr.Tab(t("monitor.t_overview")) as sub: + register_translatable(sub, "monitor.t_overview") + for i in range(0, len(PLOTS["overview"]), 2): + with gr.Row(): + for title, _series in PLOTS["overview"][i:i + 2]: + plots.append(_labelled(gr.LinePlot(x="step", y="value", color="series", label=t(title), y_title="", + height=300, tooltip=["series", "step", "value"]), title)) + with gr.Tab(t("monitor.t_recurrent")) as sub: + register_translatable(sub, "monitor.t_recurrent") + for i in range(0, len(PLOTS["recurrent"]), 3): + with gr.Row(): + for title, _series in PLOTS["recurrent"][i:i + 3]: + plots.append(_labelled(gr.LinePlot(x="step", y="value", color="series", label=t(title), y_title="", + height=270, tooltip=["series", "step", "value"]), title)) + with gr.Tab(t("monitor.t_loss")) as sub: + register_translatable(sub, "monitor.t_loss") + loss_term = _labelled(gr.Dropdown(choices=list(LOSS_TERMS), value="primary_loss", label=t("monitor.loss_term")), "monitor.loss_term") + loss_plot = gr.LinePlot(x="step", y="value", color="series", y_title="", height=340, tooltip=["series", "step", "value"]) + with gr.Tab(t("monitor.t_params")) as sub: + register_translatable(sub, "monitor.t_params") + tel_note = gr.Markdown() + group_df = _labelled(gr.Dataframe(label=t("monitor.groups"), interactive=False, wrap=True), "monitor.groups") + with gr.Row(): + load_html, imp_html = gr.HTML(), gr.HTML() + with gr.Row(): + wrms_html, grms_html = gr.HTML(), gr.HTML() + tensor_df = _labelled(gr.Dataframe(label=t("monitor.tensors"), interactive=False, max_height=520, + column_widths=["30%", "10%", "8%", "8%", "8.8%", "8.8%", "8.8%", "8.8%", "8.8%"]), + "monitor.tensors") + with gr.Tab(t("monitor.t_evals")) as sub: + register_translatable(sub, "monitor.t_evals") + eval_note = gr.Markdown(t("monitor.eval_note")) + register_translatable(eval_note, "monitor.eval_note") + eval_metric = _labelled(gr.Dropdown(choices=[], value=None, label=t("monitor.eval_metric"), + allow_custom_value=True), "monitor.eval_metric") + eval_plot = gr.LinePlot(x="step", y="value", color="series", y_title="", height=300, + tooltip=["series", "step", "value"]) + events_df = _labelled(gr.Dataframe(label=t("monitor.events"), interactive=False, wrap=True, max_height=420), "monitor.events") + samples_html = gr.HTML() + with gr.Tab(t("monitor.t_config")) as sub: + register_translatable(sub, "monitor.t_config") + params_json = _labelled(gr.JSON(label=t("monitor.training_params")), "monitor.training_params") + with gr.Tab(t("monitor.t_log")) as sub: + register_translatable(sub, "monitor.t_log") + log_box = _labelled(gr.Textbox(label=t("monitor.log"), lines=24, max_lines=24, interactive=False, autoscroll=True), "monitor.log") + ckpt_df = _labelled(gr.Dataframe(label=t("monitor.checkpoints"), interactive=False), "monitor.checkpoints") + + outputs = [run, phase, status, *plots, loss_term, loss_plot, tel_note, group_df, tensor_df, + load_html, imp_html, wrms_html, grms_html, eval_metric, eval_plot, events_df, samples_html, + params_json, log_box, ckpt_df] + inputs = [runs_dir, run, phase, loss_term, eval_metric] + refresh_btn.click(fn=refresh, inputs=inputs, outputs=outputs) + run.input(fn=refresh, inputs=inputs, outputs=outputs) + phase.input(fn=refresh, inputs=inputs, outputs=outputs) + loss_term.input(fn=refresh, inputs=inputs, outputs=outputs) + eval_metric.input(fn=refresh, inputs=inputs, outputs=outputs) + runs_dir.submit(fn=refresh, inputs=inputs, outputs=outputs) + + tab.select(fn=refresh, inputs=inputs, outputs=outputs) # first paint when the tab opens + timer = gr.Timer(15.0) + timer.tick(fn=refresh, inputs=inputs, outputs=outputs) + auto.change(fn=lambda on: gr.Timer(active=bool(on)), inputs=[auto], outputs=[timer]) + tab.monitor_wiring = (refresh, inputs, outputs) # app.load() paints it on page open + return tab diff --git a/ui/tabs/serve_tab.py b/ui/tabs/serve_tab.py new file mode 100644 index 0000000..43e3f1e --- /dev/null +++ b/ui/tabs/serve_tab.py @@ -0,0 +1,202 @@ +"""ui/tabs/serve_tab.py — start/stop ``chronos serve`` from the console. + +The server runs as a child process (argument list, no shell) so it keeps its +own memory accounting and survives UI reloads only while this process lives. +Security contract (same as the CLI): loopback by default; any other host +requires *allow remote* plus an API key. The key is passed through the +``CHRONOS_API_KEY`` environment variable, never on the command line, so it +does not show up in ``ps``. +""" +from __future__ import annotations + +import atexit +import os +import subprocess +import sys +import threading +from collections import deque + +from ui.gradio_compat import gr +from ui.i18n import register_translatable, t + +LOOPBACK = ("127.0.0.1", "localhost", "::1") + + +def build_serve_command(*, models_dir: str, host: str = "127.0.0.1", port: int = 11435, + device: str = "auto", dtype: str = "auto", lazy_experts: bool = True, + ram_expert_fraction: float = 0.5, allow_remote: bool = False, + api_key: str = "", kv_cache_dtype: str = "", mla_decode: str = "", + router_scope: str = "", max_resident_experts: int | None = None, + ram_budget_mb: float | None = None, monitor_runs_dir: str = "") -> tuple[list[str], dict]: + """Return ``(argv, env_overrides)``; raises ValueError on an unsafe or invalid request.""" + host = (host or "127.0.0.1").strip() + models_dir = os.path.expanduser((models_dir or "").strip()) + if not models_dir or not os.path.isdir(models_dir): + raise ValueError("models_dir") + port = int(port) + if not 1 <= port <= 65535: + raise ValueError("port") + remote = host not in LOOPBACK + if remote and not (allow_remote and (api_key or "").strip()): + raise ValueError("remote_needs_key") + argv = [sys.executable, "-m", "chronos.serving.ollama_server", "--models-dir", models_dir, + "--host", host, "--port", str(port), "--device", device or "auto", "--dtype", dtype or "auto", + "--ram-expert-fraction", str(float(ram_expert_fraction))] + if not lazy_experts: + argv.append("--no-lazy-experts") + if allow_remote: + argv.append("--allow-remote") + for flag, value in (("--kv-cache-dtype", kv_cache_dtype), ("--mla-decode", mla_decode), + ("--router-scope", router_scope)): + if value and value != "checkpoint": + argv += [flag, value] + if max_resident_experts: + argv += ["--max-resident-experts", str(int(max_resident_experts))] + if ram_budget_mb: + argv += ["--ram-budget-mb", str(float(ram_budget_mb))] + if monitor_runs_dir and monitor_runs_dir.strip(): + argv += ["--monitor-runs-dir", os.path.expanduser(monitor_runs_dir.strip())] + env = {"CHRONOS_API_KEY": api_key.strip()} if (api_key or "").strip() else {} + return argv, env + + +def redact(argv: list[str]) -> str: + return " ".join(argv[1:]) # the key is never in argv; drop the interpreter path + + +class ServeProcess: + def __init__(self): + self.proc: subprocess.Popen | None = None + self.lines: deque[str] = deque(maxlen=500) + self.lock = threading.Lock() + self.url = "" + self.auth = False + + def running(self) -> bool: + return self.proc is not None and self.proc.poll() is None + + def _pump(self, proc: subprocess.Popen) -> None: + for line in proc.stdout: + with self.lock: + self.lines.append(line.rstrip()) + with self.lock: + self.lines.append(f"[exit {proc.wait()}]") + + def start(self, argv: list[str], env_overrides: dict, host: str, port: int) -> str: + if self.running(): + return "busy" + env = dict(os.environ) + env.pop("CHRONOS_API_KEY", None) # only an explicit key from the form enables auth + env.update(env_overrides) + env["PYTHONUNBUFFERED"] = "1" + repo = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + self.lines.clear() + self.lines.append(f"$ {redact(argv)}") + self.proc = subprocess.Popen(argv, cwd=repo, env=env, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, text=True, bufsize=1) + threading.Thread(target=self._pump, args=(self.proc,), daemon=True).start() + shown = f"[{host}]" if ":" in host else host + self.url, self.auth = f"http://{shown}:{int(port)}", bool(env_overrides.get("CHRONOS_API_KEY")) + return "started" + + def stop(self, timeout: float = 10.0) -> str: + if not self.running(): + return "idle" + self.proc.terminate() + try: + self.proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + self.proc.kill() + self.proc.wait() + return "stopped" + + def snapshot(self) -> tuple[bool, str]: + with self.lock: + return self.running(), "\n".join(self.lines) + + +_server = ServeProcess() +atexit.register(_server.stop) # never leave an orphaned server behind the UI + + +def _status(lang: str | None = None) -> str: + if not _server.running(): + return t("serve.state.stopped", lang) + auth = t("serve.auth_on", lang) if _server.auth else t("serve.auth_off", lang) + return f"{t('serve.state.running', lang)} · {_server.url} · {auth}" + + +def build_serve_tab(): + with gr.Tab(t("tab.serve")) as tab: + register_translatable(tab, "tab.serve") + intro = gr.Markdown(t("serve.intro")) + register_translatable(intro, "serve.intro") + with gr.Row(): + models_dir = gr.Textbox(label=t("serve.models_dir"), value="./out", scale=3) + host = gr.Textbox(label=t("serve.host"), value="127.0.0.1", scale=1) + port = gr.Number(label=t("serve.port"), value=11435, precision=0, scale=1) + with gr.Row(): + device = gr.Dropdown(["auto", "cpu", "mps", "cuda"], value="auto", label=t("serve.device")) + dtype = gr.Dropdown(["auto", "float32", "bfloat16", "float16"], value="auto", label=t("serve.dtype")) + lazy = gr.Checkbox(value=True, label=t("serve.lazy")) + ram_fraction = gr.Slider(0.05, 1.0, value=0.5, step=0.05, label=t("serve.ram_fraction")) + with gr.Row(): + kv_dtype = gr.Dropdown(["checkpoint", "auto", "int8"], value="checkpoint", label=t("serve.kv_dtype")) + mla_decode = gr.Dropdown(["checkpoint", "auto", "absorbed", "expanded"], value="checkpoint", + label=t("serve.mla_decode")) + router_scope = gr.Dropdown(["checkpoint", "per_iteration", "shared"], value="checkpoint", + label=t("serve.router_scope")) + max_resident = gr.Number(label=t("serve.max_resident"), value=0, precision=0) + ram_budget = gr.Number(label=t("serve.ram_budget"), value=0, precision=0) + with gr.Row(): + allow_remote = gr.Checkbox(value=False, label=t("serve.allow_remote")) + api_key = gr.Textbox(label=t("serve.api_key"), type="password", value="") + monitor_dir = gr.Textbox(label=t("serve.monitor_dir"), value="", placeholder="./runs") + warning = gr.Markdown(t("serve.warn_noauth")) + with gr.Row(): + start_btn = gr.Button(t("serve.start"), variant="primary") + stop_btn = gr.Button(t("serve.stop"), variant="stop") + status = gr.Textbox(label=t("serve.status"), interactive=False, value=_status()) + log_box = gr.Textbox(label=t("serve.log"), lines=14, interactive=False, autoscroll=True) + usage = gr.Markdown(t("serve.usage")) + for comp, key in ((models_dir, "serve.models_dir"), (host, "serve.host"), (port, "serve.port"), + (device, "serve.device"), (dtype, "serve.dtype"), (lazy, "serve.lazy"), + (ram_fraction, "serve.ram_fraction"), (kv_dtype, "serve.kv_dtype"), + (mla_decode, "serve.mla_decode"), (router_scope, "serve.router_scope"), + (max_resident, "serve.max_resident"), (ram_budget, "serve.ram_budget"), + (allow_remote, "serve.allow_remote"), (api_key, "serve.api_key"), + (monitor_dir, "serve.monitor_dir"), (warning, "serve.warn_noauth"), + (start_btn, "serve.start"), (stop_btn, "serve.stop"), (status, "serve.status"), + (log_box, "serve.log"), (usage, "serve.usage")): + register_translatable(comp, key) + + def start(md, h, p, dev, dt, lz, rf, kv, mla, rs, mr, rb, ar, key, mon): + try: + argv, env = build_serve_command( + models_dir=md, host=h, port=int(p or 11435), device=dev, dtype=dt, lazy_experts=bool(lz), + ram_expert_fraction=float(rf), allow_remote=bool(ar), api_key=key or "", + kv_cache_dtype=kv, mla_decode=mla, router_scope=rs, + max_resident_experts=int(mr or 0) or None, ram_budget_mb=float(rb or 0) or None, + monitor_runs_dir=mon or "") + except ValueError as exc: + return t(f"serve.err.{exc}"), gr.update() + outcome = _server.start(argv, env, (h or "127.0.0.1").strip(), int(p or 11435)) + if outcome == "busy": + return t("serve.err.busy"), gr.update() + return _status(), _server.snapshot()[1] + + def stop(): + _server.stop() + return _status(), _server.snapshot()[1] + + def poll(): + return _status(), _server.snapshot()[1] + + start_btn.click(fn=start, inputs=[models_dir, host, port, device, dtype, lazy, ram_fraction, kv_dtype, + mla_decode, router_scope, max_resident, ram_budget, allow_remote, + api_key, monitor_dir], + outputs=[status, log_box]) + stop_btn.click(fn=stop, outputs=[status, log_box]) + timer = gr.Timer(2.0) + timer.tick(fn=poll, outputs=[status, log_box]) + return tab diff --git a/ui/tabs/tokenizer_tab.py b/ui/tabs/tokenizer_tab.py new file mode 100644 index 0000000..6812ce6 --- /dev/null +++ b/ui/tabs/tokenizer_tab.py @@ -0,0 +1,191 @@ +"""ui/tabs/tokenizer_tab.py — Chronos Tokenizer Builder. + +Trains a tokenizer from user-selected datasets (jsonl / parquet / csv / tsv / +txt, files, directories or globs) with ``chronos.tokenizer.build_tokenizer`` +in a background thread, then shows the held-out evaluation it wrote to +``tokenizer_report.json``. The same builder backs ``chronos tokenizer``. +""" +from __future__ import annotations + +import json +import os +import threading +import time +from collections import deque + +from ui.gradio_compat import gr +from ui.i18n import register_translatable, t + +EVAL_COLUMNS = ["bucket", "texts", "bytes/token", "chars/token", "tokens/word", + "byte_fallback_rate", "roundtrip_failures"] + + +def parse_paths(text: str) -> list[str]: + """One path per line or comma; ``~`` expanded; blanks dropped.""" + out = [] + for chunk in (text or "").replace(",", "\n").splitlines(): + chunk = chunk.strip().strip('"').strip("'") + if chunk: + out.append(os.path.expanduser(chunk)) + return out + + +def report_rows(report: dict) -> list[list]: + """Per-bucket rows from a ``heldout_eval`` dict (``vocab_utilization`` skipped).""" + rows = [] + for bucket, stats in (report or {}).items(): + if not isinstance(stats, dict): + continue + rows.append([bucket, stats.get("texts"), stats.get("bytes_per_token"), stats.get("chars_per_token"), + stats.get("tokens_per_word"), stats.get("byte_fallback_rate"), + stats.get("roundtrip_failures")]) + return rows + + +def read_report(out_dir: str) -> dict: + path = os.path.join(os.path.expanduser(out_dir or ""), "tokenizer_report.json") + if not os.path.isfile(path): + return {} + with open(path, encoding="utf-8") as fh: + return json.load(fh) + + +def report_summary(report: dict) -> dict: + if not report: + return {} + keep = ("vocab_size", "fingerprint", "algorithm", "elapsed_s", "special_tokens") + out = {k: report[k] for k in keep if k in report} + heldout = report.get("heldout_eval") or {} + if "vocab_utilization" in heldout: + out["vocab_utilization"] = heldout["vocab_utilization"] + if report.get("sources"): + out["sources"] = report["sources"] + return out + + +class TokenizerJob: + """One background build at a time; log lines are kept in a bounded deque.""" + + def __init__(self): + self.lock = threading.Lock() + self.thread: threading.Thread | None = None + self.lines: deque[str] = deque(maxlen=400) + self.state = "idle" + self.out_dir = "" + self.error = "" + + def log(self, message) -> None: + with self.lock: + self.lines.append(f"{time.strftime('%H:%M:%S')} {message}") + + def running(self) -> bool: + return self.thread is not None and self.thread.is_alive() + + def start(self, paths: list[str], out_dir: str, vocab_size: int, per_source: int, + fields: list[str] | None, temperature: float) -> str: + if self.running(): + return "busy" + if not paths: + return "no_paths" + if not out_dir: + return "no_out" + from chronos.tokenizer.extract import expand + files = expand(paths) + if not files: + return "no_files" + self.lines.clear() + self.state, self.out_dir, self.error = "running", os.path.expanduser(out_dir), "" + self.log(f"{len(files)} files: " + ", ".join(os.path.basename(f) for f in files[:8]) + + (" …" if len(files) > 8 else "")) + + def work(): + from chronos.tokenizer import BPETrainerConfig, ExtractConfig, build_tokenizer + try: + build_tokenizer( + paths, self.out_dir, vocab_size=int(vocab_size), + extract=ExtractConfig(per_source_texts=int(per_source), fields=fields or None), + trainer=BPETrainerConfig(vocab_size=int(vocab_size), + language_temperature=float(temperature)), + log=self.log, + ) + self.state = "done" + except Exception as exc: # surfaced in the UI, never raised into Gradio + self.error = f"{type(exc).__name__}: {exc}" + self.log(f"error: {self.error}") + self.state = "failed" + + self.thread = threading.Thread(target=work, name="chronos-tokenizer", daemon=True) + self.thread.start() + return "started" + + def snapshot(self) -> tuple[str, str]: + with self.lock: + return self.state, "\n".join(self.lines) + + +_job = TokenizerJob() + + +def _status_text(state: str, lang: str | None = None) -> str: + return t(f"tok.state.{state}", lang) + + +def build_tokenizer_tab(): + with gr.Tab(t("tab.tokenizer")) as tab: + register_translatable(tab, "tab.tokenizer") + intro = gr.Markdown(t("tok.intro")) + register_translatable(intro, "tok.intro") + + with gr.Row(): + paths = gr.Textbox(label=t("tok.paths"), lines=4, scale=3, + placeholder="./dataset/corpus/\n./dataset/*.jsonl\n./dataset/*.parquet") + with gr.Column(scale=2): + out_dir = gr.Textbox(label=t("tok.out_dir"), value="./out/tokenizer") + fields = gr.Textbox(label=t("tok.fields"), placeholder="text, messages, prompt, response") + with gr.Row(): + vocab = gr.Number(label=t("tok.vocab_size"), value=16384, precision=0) + per_source = gr.Number(label=t("tok.per_source"), value=400_000, precision=0) + temperature = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label=t("tok.temperature")) + with gr.Row(): + build_btn = gr.Button(t("tok.build"), variant="primary") + load_btn = gr.Button(t("tok.load_report")) + status = gr.Textbox(label=t("tok.status"), interactive=False) + log_box = gr.Textbox(label=t("tok.log"), lines=12, interactive=False, autoscroll=True) + summary = gr.JSON(label=t("tok.summary"), value={}) + table = gr.Dataframe(headers=EVAL_COLUMNS, label=t("tok.eval"), interactive=False) + for comp, key in ((paths, "tok.paths"), (out_dir, "tok.out_dir"), (fields, "tok.fields"), + (vocab, "tok.vocab_size"), (per_source, "tok.per_source"), + (temperature, "tok.temperature"), (build_btn, "tok.build"), + (load_btn, "tok.load_report"), (status, "tok.status"), (log_box, "tok.log"), + (summary, "tok.summary"), (table, "tok.eval")): + register_translatable(comp, key) + + def start(p, o, v, n, f, temp): + field_list = [x.strip() for x in (f or "").split(",") if x.strip()] + outcome = _job.start(parse_paths(p), (o or "").strip(), int(v or 16384), int(n or 400_000), + field_list, float(temp or 0.5)) + if outcome != "started": + return t(f"tok.err.{outcome}") + return _status_text("running") + + def show(o): + report = read_report(o) + if not report: + return t("tok.err.no_report"), {}, [] + return t("tok.state.loaded"), report_summary(report), report_rows(report.get("heldout_eval")) + + def poll(o): + state, text = _job.snapshot() + if state in ("done",) and _job.out_dir: + report = read_report(_job.out_dir) + return (_status_text(state), text, report_summary(report), + report_rows(report.get("heldout_eval"))) + message = _status_text(state) + (f" — {_job.error}" if _job.error else "") + return message, text, gr.update(), gr.update() + + build_btn.click(fn=start, inputs=[paths, out_dir, vocab, per_source, fields, temperature], + outputs=[status]) + load_btn.click(fn=show, inputs=[out_dir], outputs=[status, summary, table]) + timer = gr.Timer(2.0) + timer.tick(fn=poll, inputs=[out_dir], outputs=[status, log_box, summary, table]) + return tab diff --git a/ui/tabs/train_tab.py b/ui/tabs/train_tab.py index 07fc397..ee3d46b 100644 --- a/ui/tabs/train_tab.py +++ b/ui/tabs/train_tab.py @@ -80,7 +80,14 @@ "kv_latent_dim", "rope_dim", "num_attention_heads", - "num_key_value_heads", + "num_key_value_heads", "head_dim", "use_moe", "use_hybrid_attention", "tie_word_embeddings", + "architecture_version", "recurrent_enabled", "recurrent_groups", "recurrent_depth", + "prelude_layers", "coda_layers", "layer_aware_lookahead", "layer_aware_cache", + "modalities", "image_patch_size", "modality_fusion_heads", + # v3 semantics that do not change tensor shapes but do change what the + # weights compute; a stage must inherit them from its init checkpoint. + "recurrent_attention_mode", "rope_layout", "rope_theta", "sliding_window_size", + "intermediate_size", "router_scope", ) STAGE_HELP_TEXT = { "pretrain": { @@ -444,6 +451,8 @@ def _effective_training_config( "teacher_path", "data_path", "data_path_source", + "lambda_router_anchor", + "stage_overrides_applied", ] out = { "mode": mode, @@ -538,6 +547,13 @@ def _planned_optimizer_updates(train_steps: int, accumulation_steps: int) -> int return (train_steps + accum - 1) // accum if train_steps else 0 +_STAGE_OVERRIDE_KEYS = frozenset(( + "learning_rate", "batch_size", "accumulation_steps", "epochs", "max_steps", "max_seq_len", + "weight_decay", "grad_clip", "save_interval", "lambda_router_anchor", "beta", "alpha", + "lambda_or", "temperature", "num_generations", "max_gen_len", +)) + + def _prepare_training_run_config( cfg: dict | None, train_mode: str, @@ -552,6 +568,11 @@ def _prepare_training_run_config( ) -> tuple[dict, dict]: merged = dict(MINIMIND_MOE_DEFAULTS) merged.update(dict(cfg) if cfg else {}) + # Per-stage hyperparameters (Config tab → auto plan) override the shared + # values for this stage only; architecture keys are never stage-specific. + stage_values = (merged.get("stage_overrides") or {}).get(train_mode) or {} + merged.update({k: v for k, v in stage_values.items() if k in _STAGE_OVERRIDE_KEYS}) + merged["stage_overrides_applied"] = sorted(k for k in stage_values if k in _STAGE_OVERRIDE_KEYS) cfg = merged requested_backend = (selected_backend or "auto").strip().lower() or "auto" resolved_backend, resolved_device = resolve_training_device(requested_backend) @@ -891,6 +912,8 @@ def __init__(self): self.last_sample: str = "" self.last_sample_step: int = 0 self.effective_config: dict = {} + self.tokenizer_path: str = "" + self.vocab_size: int = 0 def _normalise_effective_cpu_config(self, cfg: dict, resolved_backend: str) -> None: cfg["cpu_budget_percent"] = _cpu_budget_for_backend(cfg, resolved_backend) @@ -908,6 +931,8 @@ def start(self, cfg: dict, mode: str): self._metrics.clear() self.total_steps = 0 self.t_start = time.monotonic() + self.tokenizer_path = str(cfg.get("tokenizer_path") or "").strip() + self.vocab_size = int(cfg.get("vocab_size") or 0) self._thread = threading.Thread( target=self._run, args=(cfg, mode), daemon=True ) @@ -972,7 +997,7 @@ def _resolve_default_init_path(self, save_dir: str, mode: str, hidden_size: int) def _topology_mismatches(self, sniffed: dict, model_cfg_kwargs: dict) -> list[str]: mismatches = [] for k in CHECKPOINT_TOPOLOGY_KEYS: - if k in sniffed and k in model_cfg_kwargs and int(sniffed[k]) != int(model_cfg_kwargs[k]): + if k in sniffed and k in model_cfg_kwargs and sniffed[k] != model_cfg_kwargs[k]: mismatches.append(f"{k}: ckpt={sniffed[k]} != ui={model_cfg_kwargs[k]}") return mismatches @@ -983,11 +1008,9 @@ def _adopt_checkpoint_topology(self, model_cfg_kwargs: dict, sniffed: dict) -> d continue current = model_cfg_kwargs.get(key) value = sniffed[key] - if current is None or int(current) != int(value): + if current is None or current != value: adopted[key] = value model_cfg_kwargs[key] = value - if "moe_intermediate_size" in sniffed: - model_cfg_kwargs["intermediate_size"] = sniffed["moe_intermediate_size"] return adopted def _build_stage_args(self, cfg: dict, mode: str, save_dir: str, hidden_size: int): @@ -1122,8 +1145,8 @@ def _generate_live_sample(self, model, device: str, cfg: dict, mode: str): with torch.no_grad(): gen = ids.clone() for _ in range(80): - with autocast_context(device, cfg.get("dtype", "auto")): - o, _ = model(gen, use_cache=False) + with autocast_context(device, cfg.get("dtype", "auto"), model=model): + o = model(gen, use_cache=False) nxt = o.logits[:, -1, :].argmax(dim=-1, keepdim=True) gen = torch.cat([gen, nxt], dim=-1) if int(nxt.item()) == (tokenizer.eos_token_id or -1): @@ -1149,58 +1172,21 @@ def _run(self, cfg: dict, mode: str): from chronos.model.moe_chronos import ChronosMOEFeedForward self._put(f"[{mode.upper()}] Building model...") - # Pass the FULL config dict so user-set arch overrides (e.g. - # moe_intermediate_size, num_attention_heads, vocab_size) are - # actually honored. Previously this method only forwarded a - # hardcoded subset, so the trainer silently used minimind's - # auto-derived intermediate_size while the UI estimator showed - # the user's slider — causing 10x param-count mismatches. - model_cfg_kwargs = { - "hidden_size": cfg.get("hidden_size", 512), - "num_hidden_layers": cfg.get("num_hidden_layers", 8), - "num_experts": cfg.get("num_experts", 4), - "num_experts_per_tok": cfg.get("num_experts_per_tok", 1), - "num_shared_experts": cfg.get("num_shared_experts", 1), - "lookahead_steps": cfg.get("lookahead_steps", 2), - "kv_latent_dim": cfg.get("kv_latent_dim", 64), - "sliding_window_size": cfg.get("sliding_window_size", 2048), - "lambda_balance": cfg.get("lambda_balance", 5e-4), - "lambda_temporal": cfg.get("lambda_temporal", 3e-3), - "lambda_lookahead": cfg.get("lambda_lookahead", 0.1), - "lambda_lookahead_topk": cfg.get("lambda_lookahead_topk", 0.15), - "lambda_lookahead_union": cfg.get("lambda_lookahead_union", 0.05), - "lambda_router_locality": cfg.get("lambda_router_locality", 0.02), - "fallback_mask_prob": cfg.get("fallback_mask_prob", 0.08), - "vram_budget_gb": cfg.get("vram_budget_gb", 4.0), - "use_hybrid_attention": cfg.get("use_hybrid_attention", True), - "use_moe": True, - } - # Optional overrides: only forward when the user actually set them. - # For the "auto" sentinel fields (0 = let MiniMindConfig derive - # it), we skip forwarding when value is 0 so ceil(H·π/64)·64 - # takes effect instead of Linear(H, 0) crashing MPS init. - AUTO_SENTINEL_KEYS = {"intermediate_size", "moe_intermediate_size"} - for opt_key, cfg_key in [ - ("num_attention_heads", "num_attention_heads"), - ("num_key_value_heads", "num_key_value_heads"), - ("rope_dim", "rope_dim"), - ("vocab_size", "vocab_size"), - ("max_position_embeddings", "max_seq_len"), - ("intermediate_size", "moe_intermediate_size"), - ("moe_intermediate_size", "moe_intermediate_size"), - ("lambda_router_anchor", "lambda_router_anchor"), - ("pinned_memory_max_fraction", "pinned_memory_max_fraction"), - ("storage_format", "storage_format"), - ("tie_word_embeddings", "tie_word_embeddings"), - ]: - if cfg_key not in cfg: - continue - val = cfg[cfg_key] - if val in (None, ""): - continue - if opt_key in AUTO_SENTINEL_KEYS and val == 0: - continue - model_cfg_kwargs[opt_key] = val + from chronos.config_schema import model_config_kwargs + model_cfg_kwargs = model_config_kwargs(cfg) + model_cfg_kwargs.setdefault("hidden_size", 512) + model_cfg_kwargs.setdefault("num_hidden_layers", 8) + detected_modalities_pre = [] + if data_path_for_cleanup and os.path.exists(data_path_for_cleanup): + try: + from chronos.data.flexible_dataset import dataset_modality_summary + detected_modalities_pre = dataset_modality_summary(data_path_for_cleanup).get( + "detected_modalities", [] + ) + except Exception as exc: + self._put(f"[data] modality detection warning: {exc}") + if cfg.get("auto_detect_modalities", True): + model_cfg_kwargs["modalities"] = detected_modalities_pre save_dir = cfg.get("save_dir", "./out") init_weight = (cfg.get("init_weight") or "").strip() load_path = "" @@ -1276,6 +1262,13 @@ def _run(self, cfg: dict, mode: str): requested_backend = (cfg.get("train_backend") or "auto").strip().lower() or "auto" resolved_backend, device = resolve_training_device(requested_backend) + if resolved_backend == "mlx" and model_cfg.modalities: + if requested_backend != "auto": + raise ValueError("MLX 尚未实现多模态编码器训练，请选择 auto 或 PyTorch MPS/CUDA 后端。") + supported = set(training_available()) + target = next(name for name in ("cuda", "xpu", "mps", "cpu") if name in supported) + resolved_backend, device = resolve_training_device(target) + self._put(f"[data] Modalities {model_cfg.modalities}: using {resolved_backend} for encoder training.") cfg["train_backend"] = requested_backend cfg["device"] = device cfg["dtype"] = _normalize_dtype_for_trainer(cfg.get("dtype")) @@ -1324,12 +1317,34 @@ def _run(self, cfg: dict, mode: str): else: tokenizer = self._load_tokenizer() from chronos.data.flexible_dataset import ( - FlexibleDataset, StreamingSFTDataset, StreamingDPODataset, + FlexibleDataset, StreamingSFTDataset, MultimodalStreamingDataset, + MultimodalStreamingSFTDataset, StreamingDPODataset, + dataset_modality_summary, multimodal_collate, ) + modality_summary = dataset_modality_summary(data_path) + detected_modalities = modality_summary.get("detected_modalities", []) if mode == "pretrain": - full_ds = FlexibleDataset(data_path, tokenizer, max_length=max_seq_len) + if cfg.get("auto_detect_modalities", True) and detected_modalities: + model_cfg_kwargs["modalities"] = detected_modalities + full_ds = ( + MultimodalStreamingDataset(data_path, tokenizer, max_length=max_seq_len, + image_size=int(cfg.get("image_size", 224)), + audio_sample_rate=int(cfg.get("audio_sample_rate", 16000)), + max_audio_seconds=float(cfg.get("max_audio_seconds", 30.0)), + media_error_policy=cfg.get("media_error_policy", "warn")) + if detected_modalities else FlexibleDataset(data_path, tokenizer, max_length=max_seq_len) + ) elif mode in {"sft", "distill"}: - full_ds = StreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len) + if cfg.get("auto_detect_modalities", True) and detected_modalities: + model_cfg_kwargs["modalities"] = detected_modalities + full_ds = ( + MultimodalStreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len, + image_size=int(cfg.get("image_size", 224)), + audio_sample_rate=int(cfg.get("audio_sample_rate", 16000)), + max_audio_seconds=float(cfg.get("max_audio_seconds", 30.0)), + media_error_policy=cfg.get("media_error_policy", "warn")) + if detected_modalities else StreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len) + ) elif mode in {"dpo", "orpo"}: full_ds = StreamingDPODataset(data_path, tokenizer, max_length=max_seq_len) else: @@ -1352,6 +1367,7 @@ def _run(self, cfg: dict, mode: str): DataLoader( train_ds, batch_size=batch_size, + collate_fn=multimodal_collate if detected_modalities and mode in {"pretrain", "sft", "distill"} else None, **dataloader_kwargs(device=device, num_workers=cfg.get("num_workers", "auto"), shuffle=True), ) if train_ds is not None else None @@ -1360,6 +1376,7 @@ def _run(self, cfg: dict, mode: str): DataLoader( val_ds, batch_size=batch_size, + collate_fn=multimodal_collate if detected_modalities and mode in {"pretrain", "sft", "distill"} else None, **dataloader_kwargs(device=device, num_workers=cfg.get("num_workers", "auto"), shuffle=False), ) if val_ds is not None else None @@ -1413,10 +1430,16 @@ def _run(self, cfg: dict, mode: str): if mlx_checkpoint_path: self._put(f"[MLX {mode.upper()}] Loading checkpoint → {mlx_checkpoint_path}") else: - self._put(f"[MLX {mode.upper()}] Starting from current in-memory random init.") + self._put(f"[MLX {mode.upper()}] Starting from random init.") def _mlx_progress(event: dict): kind = str(event.get("event", "step")) + if kind == "validation": + value = float(event["ce"]) + self._put(f"[MLX {mode.upper()} validation] step={event['step']} CE={value:.4f} tokens={event['tokens']}") + self._put_metric({"step": event["step"], "total": value, "primary": value, + "ce": value, "aux": 0., "temporal": 0., "tps": 0., "val": True, "val_loss": value}) + return if kind == "start": total = int(event.get("steps", 0) or 0) if total > 0: @@ -1426,6 +1449,7 @@ def _mlx_progress(event: dict): f"total_steps={self.total_steps} accum={event.get('accumulation_steps', '')} " f"save_interval={event.get('save_interval', '')}" ) + self._put(f"[MLX optimizer] {event.get('optimizer', '')}") return if kind == "step": step = int(event.get("step", event.get("steps", 0)) or 0) @@ -1439,17 +1463,25 @@ def _mlx_progress(event: dict): metric = { "step": step, "total": loss, - "primary": float(event.get("ce", loss) or loss), - "secondary": float(event.get("aux", 0.0) or 0.0), - "tertiary": float(event.get("temporal", 0.0) or 0.0), - "ce": float(event.get("ce", loss) or loss), + "primary": float(event.get("primary", event.get("ce", loss))), + "secondary": float(event.get("secondary", event.get("aux", 0.0))), + "tertiary": float(event.get("tertiary", event.get("temporal", 0.0))), + "ce": float(event.get("ce", 0.0)), "aux": float(event.get("aux", 0.0) or 0.0), "temporal": float(event.get("temporal", 0.0) or 0.0), "tps": tps, + "tokens_per_s": float(event.get("tokens_per_s", 0)), + "grad_norm": float(event.get("grad_norm", 0)), + "loss_components": event.get("loss_components", {}), + "recurrent_loss": float(event.get("recurrent_loss", 0)), + "lookahead_loss": float(event.get("lookahead_loss", 0)), } self._put_metric(metric) self._put( f"[MLX {mode.upper()}] step={step}/{self.total_steps} loss={loss:.4f} " + f"primary={metric['primary']:.4f} aux={metric['aux']:.4f} " + f"recurrent={metric['recurrent_loss']:.4f} grad_norm={metric['grad_norm']:.3f} " + f"tokens/s={metric['tokens_per_s']:.0f} " f"steps/s={tps:.2f} lr={float(event.get('lr', 0.0) or 0.0):.2e} " f"elapsed={float(event.get('elapsed_s', 0.0) or 0.0):.1f}s" ) @@ -1480,7 +1512,7 @@ def _mlx_progress(event: dict): return if kind == "failed": self._put( - f"[MLX {mode.upper()}] stopped after rollback limit at " + f"[MLX {mode.upper()}] failed ({event.get('reason', 'numerical_error')}) at " f"step={int(event.get('step', 0) or 0)}/{int(event.get('steps', self.total_steps) or self.total_steps)}" ) return @@ -1496,6 +1528,7 @@ def _mlx_progress(event: dict): checkpoint_path=mlx_checkpoint_path, save_dir=save_dir, loader=loader, + validation_loader=val_loader, prompts=prompts, teacher_path=stage_args.teacher_path if mode == "distill" else None, args=stage_args, @@ -1515,7 +1548,7 @@ def _mlx_progress(event: dict): if result.stopped: self.status = "error" self._put( - f"[MLX {mode.upper()}] Stopped after rollback limit in {elapsed:.1f}s. " + f"[MLX {mode.upper()}] Failed ({getattr(result, 'failure_reason', 'numerical_error')}) in {elapsed:.1f}s. " f"Last finite checkpoint → {result.checkpoint_path}" ) if result.checkpoint_saved: @@ -1527,10 +1560,19 @@ def _mlx_progress(event: dict): return model = model.to(device) - from chronos.trainer.optim_utils import build_optimizer, get_lr, apply_lr - optimizer = build_optimizer(model, lr=base_lr, weight_decay=weight_decay) - train_autocast = autocast_context(device, cfg.get("dtype", "auto")) - scaler = grad_scaler(device, cfg.get("dtype", "auto")) + from chronos.trainer.optim_utils import build_configured_optimizer, get_lr, apply_lr + optimizer = build_configured_optimizer( + model, model_cfg, lr=base_lr, weight_decay=weight_decay, + ) + from chronos.trainer.paged_store import PagedTrainingController + paged_training = PagedTrainingController(model, model_cfg) + # One compute dtype end to end (weights, activations, KV) with + # FP32 masters in the optimizer; no autocast per-op casts. + from chronos.trainer.precision_training import attach_precision + precision, train_autocast, scaler = attach_precision( + model, optimizer, device, cfg.get("dtype", "auto")) + self._put(f"[precision] {precision.backend}/{precision.device}: compute={precision.compute} " + f"master={precision.master} ({precision.reason})") model.train() global_step = 0 @@ -1546,12 +1588,27 @@ def _eval_val() -> float: model.eval() tot, n = 0.0, 0 with torch.no_grad(): - for k, (vx, vy) in enumerate(val_loader): + for k, vbatch in enumerate(val_loader): if k >= 50: # cap eval cost break + if isinstance(vbatch, dict): + vx, vy = vbatch["input_ids"], vbatch["labels"] + vmedia, vmedia_mask = vbatch.get("modality_inputs"), vbatch.get("modality_mask") + else: + vx, vy = vbatch + vmedia = vmedia_mask = None vx = vx.to(device); vy = vy.to(device) - with autocast_context(device, cfg.get("dtype", "auto")): - vout, _ = model(vx, labels=vy) + if vmedia: + vmedia = {key: value.to(device) for key, value in vmedia.items()} + if vmedia_mask is not None: + vmedia_mask = vmedia_mask.to(device) + with autocast_context(device, cfg.get("dtype", "auto"), model=model): + eval_masks = paged_training.masks_for_forward(device, allow_evict=False) + vout = model( + vx, labels=vy, + modality_inputs=vmedia, modality_mask=vmedia_mask, + available_expert_masks=eval_masks, + ) tot += float(vout.loss.item()); n += 1 model.train() return tot / max(n, 1) @@ -1564,19 +1621,35 @@ def _eval_val() -> float: trainer = ChronosSFTTrainer(model, model_cfg, stage_args, tokenizer) if float(getattr(model_cfg, "lambda_router_anchor", 0.0)) > 0 and loader is not None: first_batch = next(iter(loader)) - trainer.set_calibration_batch(first_batch[0]) + if isinstance(first_batch, dict): + trainer.set_calibration_batch( + first_batch["input_ids"], + modality_inputs=first_batch.get("modality_inputs"), + modality_mask=first_batch.get("modality_mask"), + ) + else: + trainer.set_calibration_batch(first_batch[0]) self._put("[SFT] Router anchor reference captured.") iters = _planned_epoch_iters(loader, stage_args.steps) stage_global_step = 0 for epoch in range(stage_args.epochs): epoch_start = time.monotonic() - for step_idx, (ids, labels) in enumerate(loader, start=1): + for step_idx, batch in enumerate(loader, start=1): if self._stop.is_set(): break if step_idx > iters or (stage_args.steps is not None and stage_global_step >= stage_args.steps): break stage_global_step += 1 - result = trainer.train_step(ids, labels, stage_global_step, self.total_steps) + if isinstance(batch, dict): + ids, labels = batch["input_ids"], batch["labels"] + result = trainer.train_step( + ids, labels, stage_global_step, self.total_steps, + modality_inputs=batch.get("modality_inputs"), + modality_mask=batch.get("modality_mask"), + ) + else: + ids, labels = batch + result = trainer.train_step(ids, labels, stage_global_step, self.total_steps) if step_idx % stage_args.log_interval == 0 or step_idx == iters: elapsed = max(time.monotonic() - epoch_start, 1e-6) self._record_stage_metric(mode, stage_global_step, result, tps=step_idx / elapsed) @@ -1691,19 +1764,33 @@ def _eval_val() -> float: trainer = ChronosDistillTrainer(model, teacher, model_cfg, stage_args, tokenizer) if float(getattr(model_cfg, "lambda_router_anchor", 0.0)) > 0 and loader is not None: first = next(iter(loader)) - trainer.set_calibration_batch(first[0]) + if isinstance(first, dict): + trainer.set_calibration_batch( + first["input_ids"], modality_inputs=first.get("modality_inputs"), + modality_mask=first.get("modality_mask"), + ) + else: + trainer.set_calibration_batch(first[0]) self._put("[DISTILL] Router anchor reference captured.") iters = _planned_epoch_iters(loader, stage_args.steps) stage_global_step = 0 for epoch in range(stage_args.epochs): epoch_start = time.monotonic() - for step_idx, (ids, labels) in enumerate(loader, start=1): + for step_idx, batch in enumerate(loader, start=1): if self._stop.is_set(): break if step_idx > iters or (stage_args.steps is not None and stage_global_step >= stage_args.steps): break stage_global_step += 1 - result = trainer.train_step(ids, labels, stage_global_step, self.total_steps) + if isinstance(batch, dict): + result = trainer.train_step( + batch["input_ids"], batch["labels"], stage_global_step, self.total_steps, + modality_inputs=batch.get("modality_inputs"), + modality_mask=batch.get("modality_mask"), + ) + else: + ids, labels = batch + result = trainer.train_step(ids, labels, stage_global_step, self.total_steps) if step_idx % stage_args.log_interval == 0 or step_idx == iters: elapsed = max(time.monotonic() - epoch_start, 1e-6) self._record_stage_metric(mode, stage_global_step, result, tps=step_idx / elapsed) @@ -1724,17 +1811,33 @@ def _eval_val() -> float: for epoch in range(epochs): if self._stop.is_set() or stop_outer: break - for step, (input_ids, labels) in enumerate(loader, start=1): + for step, batch in enumerate(loader, start=1): if self._stop.is_set(): break if max_steps and global_step >= max_steps: stop_outer = True break + if isinstance(batch, dict): + input_ids, labels = batch["input_ids"], batch["labels"] + media, media_mask = batch.get("modality_inputs"), batch.get("modality_mask") + else: + input_ids, labels = batch + media = media_mask = None input_ids = input_ids.to(device) labels = labels.to(device) + if media: + media = {key: value.to(device) for key, value in media.items()} + if media_mask is not None: + media_mask = media_mask.to(device) with train_autocast: - out, lp = model(input_ids, labels=labels) + paged_masks = paged_training.masks_for_forward(device) + result = model( + input_ids, labels=labels, + modality_inputs=media, modality_mask=media_mask, + available_expert_masks=paged_masks, + ) + out, lp = result.legacy_tuple() moe_layers = [l.mlp for l in model.model.layers if isinstance(l.mlp, ChronosMOEFeedForward)] @@ -1786,6 +1889,7 @@ def _eval_val() -> float: model.parameters(), 1.0, ) + paged_training.observe() if global_step % log_interval == 0: now = time.monotonic() @@ -1865,8 +1969,8 @@ def _eval_val() -> float: with torch.no_grad(): gen = ids.clone() for _ in range(80): - with autocast_context(device, cfg.get("dtype", "auto")): - o, _ = model(gen, use_cache=False) + with autocast_context(device, cfg.get("dtype", "auto"), model=model): + o = model(gen, use_cache=False) nxt = o.logits[:, -1, :].argmax(dim=-1, keepdim=True) gen = torch.cat([gen, nxt], dim=-1) if int(nxt.item()) == (tokenizer.eos_token_id or -1): @@ -1919,7 +2023,24 @@ def _synthetic_loader(self, vocab_size, seq_len, batch_size, n=50): def _load_tokenizer(self): from transformers import AutoTokenizer - return AutoTokenizer.from_pretrained(chronos.deps.get_tokenizer_path()) + path = resolve_tokenizer_dir(self.tokenizer_path) + tok = AutoTokenizer.from_pretrained(path) + if self.vocab_size and len(tok) > self.vocab_size: + raise ValueError( + f"tokenizer at {path} has {len(tok)} tokens but the model vocab_size is " + f"{self.vocab_size}; set vocab_size >= {len(tok)} or pick the matching tokenizer" + ) + return tok + + +def resolve_tokenizer_dir(configured: str | None) -> str: + """Config ``tokenizer_path`` (a Chronos Tokenizer Builder output) wins; else minimind.""" + configured = os.path.expanduser(str(configured or "").strip()) + if configured: + if not os.path.isfile(os.path.join(configured, "tokenizer.json")): + raise FileNotFoundError(f"tokenizer_path {configured!r} has no tokenizer.json") + return configured + return chronos.deps.get_tokenizer_path() _session = TrainSession() From c224202313d475478337aebb8edbb469be1b990e Mon Sep 17 00:00:00 2001 From: Fona <62028332+FonaTech@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:11:22 +0800 Subject: [PATCH 2/3] Stop tracking compiled __pycache__ files They were uploaded with the source, are ignored by the new .gitignore and carry paths from the build machine. Co-Authored-By: Claude Opus 5.5 --- .../data/__pycache__/__init__.cpython-310.pyc | Bin 166 -> 0 bytes .../flexible_dataset.cpython-310.pyc | Bin 10770 -> 0 bytes .../__pycache__/__init__.cpython-310.pyc | Bin 500 -> 0 bytes .../__pycache__/chronos_trainer.cpython-310.pyc | Bin 5226 -> 0 bytes .../__pycache__/distill_trainer.cpython-310.pyc | Bin 7164 -> 0 bytes .../__pycache__/dpo_trainer.cpython-310.pyc | Bin 6990 -> 0 bytes .../__pycache__/grpo_trainer.cpython-310.pyc | Bin 10317 -> 0 bytes .../__pycache__/loss_mixin.cpython-310.pyc | Bin 5367 -> 0 bytes .../__pycache__/optim_utils.cpython-310.pyc | Bin 3782 -> 0 bytes .../__pycache__/orpo_trainer.cpython-310.pyc | Bin 6911 -> 0 bytes .../trainer/__pycache__/reward.cpython-310.pyc | Bin 5638 -> 0 bytes .../__pycache__/sft_trainer.cpython-310.pyc | Bin 6373 -> 0 bytes 12 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 chronos/data/__pycache__/__init__.cpython-310.pyc delete mode 100644 chronos/data/__pycache__/flexible_dataset.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/__init__.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/chronos_trainer.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/distill_trainer.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/dpo_trainer.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/grpo_trainer.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/loss_mixin.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/optim_utils.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/orpo_trainer.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/reward.cpython-310.pyc delete mode 100644 chronos/trainer/__pycache__/sft_trainer.cpython-310.pyc diff --git a/chronos/data/__pycache__/__init__.cpython-310.pyc b/chronos/data/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 0297b2f53a64f1cac9b547819a5e02af234f0da6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 166 zcmd1j<>g`k0v6k+nIQTxh(HF6K#l_t7qb9~6oz01O-8?!3`HPe1o6vLKeRZts94`E zKQB?=CBHl`CqFTzSl^>EsVFlg-p9vRKcFZ-D>blIYq;;_lhPbtkwwF6mE%mgG@7yx!qC)@x4 diff --git a/chronos/data/__pycache__/flexible_dataset.cpython-310.pyc b/chronos/data/__pycache__/flexible_dataset.cpython-310.pyc deleted file mode 100644 index 7021379ca2977922142772de97a22cac1d5f7869..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 10770 zcmbVS+jrc?ea3yU;F6*#ik2^R;Wml1mbjE-r*Rcowj)_?6NPpp*=}g75V!-9OD?gi z0ag;3W#g2xowQZmOPaQEoRGciwD}8q`Umu}uXyM~drnW)=REW%jo9xS?3EX}r^|!E z%wRB>`CY%?Z`c?e^)&p#L;qR(ml;j_cQyup78);Wnx6a-Dngsngf5J_zG}=FI_qZL zoHNzknzPj1p0m~6nR9Tr>V;Kz&c#@Jt|**2PaF{=qIBDu8+lC=qhjo~CdMLn)0``z zeMpR>Jsy?N9!2{RF@g4kYLB6PSWKclsoIBxd08t@wSJ5xdetS_h#Sd_2-9$8p&s3; z&DW!Vtt3h-%{9-voXRL%t;LJ}S1!Lkd(p2qLJ`TtzgbI{{Q0#snrM;NR=!e@MG zp;nJDdD`$Vz4Eowp6B})I_o^2FVoL||MuSZ{HBc5y#9<|597raX7*PbBC4x(z0yh> z(;`ZvYMNl01Ao=V>S`nQ)96+@y-&(RA7(-75Z@ z3CKnYrHz#+uB}J%*u%HA6sFkJYO@}u(U<&2Twn9UB&mS`F&=OdGAmAHZN8O858E7= z!4KCP320Ie=cBs+qCfrYsZ)MbPokSkk&Ma}|GaE8r)%-UCR%7nAHr3`C5|P2fA2#i zIseA%!z2Np7x0a^8u`tzCX*umWt(yADQrL`qyy?$8Yyq!d!U!cZ6zb?! z6T(sPyjR2OlHXASNN7WX)QVu&a%oRa`*jfR9mwU?YyR;UlqMuINeJpjWxrW(sik;I zc$;CmRPnC`!D5uwzIzoFM~a@iK$y;{TcsRA6<qo+f^{kfI+vYw5E zaok9AwV6e`(2<6+lUdZrtkBWO%qFkavO*Jr4b{qA(WtgoF?m)9Wm>C(Ifb&3jSOgK z=H%*GHY81Y=7O3{=x2m+4z{)0Ol67dFyttodoT}RGIbau>{<-fmLqCF=3Z|_AWb9A zoG?yqMv^1YC)ea7eAkf?I4ClUBFYwpHZx%N(-&e4JVH!7)iUb{N6p||z6 zc0-esZLO_8hI@)OYn!^E>Dp!Ox{(?fdHh{t!`Lv}h8q1++syB3w87C^=C-C4osrlroNA>S4@#ye`f>doJ|tcije9jnoW z&gfW;b~~fd+Z$aJ9=@n;6}D9$)`5B3*0v5eH8Pv&-VkjJkb_r3+8jmY+j#eCfwYlE zTYn35xNZp0;;ue>7fj2H%Bf6m%EvHN75BeB1IsCsnG20LoH^gP87rNe`L(rqSrfs< zi(i|0LpGKn%t6Oj&3u}d%!+3gpw-`vIg=hM0uBP5!NAjl%q_W zy=%xR)=bzSNb=*X>(xG0OsFqwZk)l(xi z7fi^g0V&rZ@7A`~t^c8M9ZON4E8Ce@OK1y1xRMTmVA3;t0WyIb6tmXst?|aLHD=x! zIsw{)voCY0q?)UZ>Bf$ykL|2KeGvD1ZMNPC@+_{;XLb`-F4?`z(V=xbi+2*z-qpwPXS`QJ z>FY;Po67Sr8Q?6_p!<6q({xhE$ z!r!K;Yq9&RmOOtG&-wQaTzn@vmwGU+`p4QQ(BI#tMr*-HTG|@jW^IgytZdI_hMbr5 z5LebGXB!Lir zRIg{=TQnK3O4*Q^87wSi#oUjB{F!wTzC3f$8_=smb*v<+FG%$feIYZOTmXU0Zt-)2 z00=J)g0Emb$!AcE>5e|8Te_jUow9_Qao;eX{FzhCpRPZ-k7!i7Q|&PZJzmx^f2wWi z-~c$e3|G+7l8F>95OvUZ(`@T2rhKmROqa*k;jWeqnh0HcTZ@a9mfCB_x_pZDn{am+ zQ4q!^9BKd=;rPI1Nt-sCm4c3k2%0q^e-$ez7bWdz=D^&t4`K#Zpb@62%8 z^(XcvU;Zvp3xYLCk0Ot-m_%`~e$3%1cB`Pq257LXe{_WMyJb+$+vXOahLU`fRNS!k z$h|)(_o=<*P;cNJ@COPSv+H&%U6_8M?8+CpnhPv0p`eUKE0TTyW)S!|vz4>X904FD zBA7rDO2Q3**14o-1_alX;mzHYD+QE&6og=lcq17@;X(3T5YW?~{Q0;AkGzyZsAONt(D!$6;C@GTqAb*oZz@oowRa5BZi+DhBQR{X!3m(pW#4EXO$&`Fn#WnP! zlfX|;PNE$0bV$W~s4eKCAl%!d@OKKs&eCx=at}Mm+**7}LiRmO=Lr4Z} z(Mhay;><7k;Xr~SkL-aO5%ePkQMpG!>yt0@R>NDAh((xa>T#8zv-=||BW^Ug`0t~V zYMYP+dN=wfWyN3-16+Xc=>0Bi-gf}`H>}iA8)F#YrUim@*i57Oe5c=1{cgWs;5~#S z&j~x$g|lI&MNvpS$lk~%yqa*^Hp*h#5}NSdNAU5n@rkpERVdFkGON7WkW5ntux{-+ zuAO}QHgxK0kpN$PX$a&iEmO?ypfve;ym40tfwtn~+5)CSM0l*dAs6{tV`}b%A$|Znc=AV?N0S|OjO>zNGbETu*$t;9F zshnnKg>Abj0i@--?4V6AJGu2!`oGFvPJsBlvzR%gN+9OJ65ay)ib0_8A5>gXL$Q#` z{w-Vzy_VqQ9M~RLKWa?C{TVxOeWvk0H29#@gzr=z!$ZZ_%to%Ix|BJ{g(g6ZR=ok7NNf{Fz(2IK zxrsl_8WkopF#*P_R3nlrm=SZW^NFLfh9?{jBTRPJ>f6&?o@q$S1$hlO8M2sXQDvb#&HJcjqaE3- z0y0BB`<~`H$07WH09$`t^Lw%jh2RF%#nOkdK)Zo)|$aW#cG- z(G;VqJen#fzvDtorn195{L)XRhNA!~y?!JSTgYEa#)?S$9+ZDoS<_2EOo@MOC`obJ zf9IRBQIGn0aKBQiY@B3?WN63cob(Y|)N4_mlHQf3_N#-rXbeRL5&88MWMGjDK}gp~ zYQ)7pQgcjgVt`uCK)M=?NdCo)+d`rEK?PxS3e?{d+Xv`*DBBjZ$WL z!Qkr&Md3gxrK@nQKWiuI{?dQ9vy7lwYaOY}iQUoJ_;Lb>4f)xFqL_m1MRO~vnNd*-kT#Uv!{ zTX1NLQt4p}E|t*tF&5S2C%Ba7&TX_|A~4x84C_I^_PDwJE;5huB{iCvPL6)5Ga7BA2P|R1Fe44X93`@iDmYwr+lH`(4L9N3v7m{C+03Gr zWKOb{An8wQ*{CO($6J6jU_H42Tka}jh{K!g2(3EPh*BBwpp$z3V{8@HbP_lD%Pfwu zVBja;Wbqb@msya<*yJ(AqLgI;wvL*90uAgO-C7*S<^~oI^?3{bKL5WQeKsy=3e;D%gLaBJ@R-+uYPA z$zEh!x)GIPFuW>Q3?eiCxv8&cGC?$s2()cIsu5kn%j#fzE$@Nvb?8@nqc31I+HH92 z&S(nBnb>O=g!x5naBOl4IRhuP^3j;r-bDNlQFAw3%;Mr(M%z83UH|&P{30J6ch
  • b2RV0tK>CvtAFDTR1HF3G+I|Fq*+tdGdV;WWe3>5M#>PYLv6?qGEQttg8 z!_T`;uA+C?OBb2h5=Wm=X5vgud9|)Drk|C!IYO^xMm&!zVu}-PATyZ+(wBFO!$_dq zazNYs&v-MT%#0ayy&i;M3PND1J2N?@pFeJm=}+mdQTnMhy8es%ibj9Fg9YKQ@R2?U zgQh8cCJB=#pqCT@LN979lM&GBJLa;zVaEu5p_R~fFkVAa!r9iBZKl+gEwWp|G83@u zl(4{RyIoMQ%o>Jec(d5HF^8La%Oe12-f~IUcMR2HpK2*k=FBPxeg&~d3v13X1eqdl zpy=W)rL_dfebA+pLjD>12orJ+a@%r)#c#17((SVDGCK&P0IzUfkaW-#P?|cElAC-^ z*6w1Ud>?)Pic3LeG7Wl+{1EhcN5N!s3hXW8T>K5y>Z2 zTth!v6a4gq1RsJ^fJ#GfYNC^xwmYe5#JvdK3((BN10bq@^t}F0J%N4*g6faQDJ+r3 z0SD7WD;PVo-}&=t^`h^8^A_wl5{2? z7`MwuB0l{TA9;OWAG!Gd@R7#<(MLLnV!`q~{&9X^|HxVQ`bV=P06!pUK1R_shNM+y zQN-ns`0U3lwpeVl_zM<)$)ayg6t?kjaFHa+m*M}ax+f8qKDnP_ zd)9XO<24?}%;IKDYoMv&J1^JVWtAtdgg?~;;D$J@&lh3AEFe|CfvsJ2JjTPzuMs{4;L8zzT99*eeBqXlD z3dgZR5RB#VpQo{B#&R+l17n-l50*jz-kYFr0~4H?t#jijppXLrF@%jC@PY8KXSK(D z5#YdTpASWdL#qQGi4h)IJ>+8%yy-6Pjog41=|0jb*-R5LY>^?ltD&l0+k9q}8INYz7u zMGt*6Y*0$ulHJ|V&^MqWb6EG4qg!1njb^ec&{V!uTozQL%k7aOmt zud149t>!3r{_^MlgsnwI`6o3FJ~|rjp=AF;#T2GSN@x65JDP0k9bLYSjv?P>$CPg? zE@_IY#CE5wl7tqSvD0xRtw)u(+Nny~h-z`YQ9aqwo=;ioC=Y*uI(JS%E z&dK=I&a0|&O<^@wf1t4XuGV>tsjEt>@dztbZnjimru{sOK{%KF+`pdsEZ}HWl0g)8 z(=V;b%gh`0PE~UwOxY06%bv@nK2)Jm3nHT$E;`Mpj4YEwsqS)(h2HmZG8Ya1L zvLGMyy(|bA#<4bn+>3Zi6OF4;a6ep&0-6?m#N2Am13wOvjn%8SCQL}wE^YEONwdks zs+gTHuBQBwKgj&(`VDkjj?#3?-wb@_WrKde?}k~JCSH_g877M}^=>QB;e=Y0Mt{i;!E#oKP21-X|3M@i}HE4mct}&e%4-D9X8Cbm%+7>IJZ8Mvd zA1IwNYqBbjakn9Jmu~> z$_G5{L1T9!*>+_AkGbE^TeDIwz?u> z93R_uyJ_6Vj?%1I$bzZ1H{QQ;6*hoVyXSM(nv8dSdTJlaQ*ixmcfkGbHgwqg+DC{)Vbh3&m3k4YgWmJ$h@oUf{Htva1jE0tgoMXEaN*>`~7-vWy5gA6lwXSm4h==1{&UHFa2J7n-r^LmhpW zm;y_w-4q6aM;P4!^M&5+4}`t8j)Tb0z6--QS-#y5MCmD~6eK=y#W*ZVL664TAdHwN zaUSjfvekP*xUrdgEa>{%8=uer;o9GK9{dn6(IU=|lF(s3hnFyXzL5!w1$V@gvl^76m>hbN2w0K$vnLM0tO;trG2D zv@Tb~kCSTiBv%=)%A?(r>yX3!Ug2-y>vyXwA7ug0R<5Rrzj8UfmyqFSD?i>|<012| zU%#<30i!nov#Y0i{H4P{^+#3dfL5y+P!kS-O~6B2H`DKEtR*4 zrX3qsXF-Vt??tZl0zPG=nWOETe11hQ66l)fJbjuV_R;=3EV58Xa$$Sw5NO}0wOlLo z!oZUCo&nUGpv<0iS-EWut->rUre09)EP-44@)ICX>Dllb8q*$XpJ~I=&@Kq?Wz^2l z$t&`Exv=rB77o)Nn!^e+hSfqD)%KKOt*A+#k`%RGotcLfW^NfgDeMUsQoRPVB(K(M zvJ$f&s-qc5*7I46TBmV)4d~g(=ZeN?eoxs|F*?nw_ZE&Q!^+3VDh&K~byZ31vr6{G zJ@_s9eS(61henG zX-}Js<4*hGeQ5Nu9FKPI1S>P=A>@k+X(*)@=9Qu>TaK)&>jn(0a+4p!LBZ!(k-rv1 z87QFtzKj$b)3ba;HY%OdBeDnd(ODSZ!n=Yzz{^uFNt2-U3O_~tUZRQ&d5qmRi8_Ny zRQo>nV{ZfTyPxq_iBCo;42t>uU22onjB)=qeKE65AFi@-AXUhNekP1K@Dp-#xepXK zNa!^6H$p^8F-tKi9z=dl2}MQ{{Ppp0Bc;Mg1&1>B24M2M57yp zeW;4ynWQ}Sgd>#?ka%RGEcw&}o(jVYYtU!p%gS)psg!CpCDKGWOb}bWkY)TGQuC)Y zPbE)t1oN?GwEsfuZ-^2CD*9#zmn(A|QJRGA@g5N_JF`!>gm%xn>qnjPBr$rMiBcY5 zX9C_SMZV0Qo`<|6y_{6xJ&c!KMpcDhYr>B?s;$k!&snM_y_)sdux#+^((5%f4Lyt} zRlTkq{de<76Z~cPJqN|rb`C%1DW1+F4d`aq!c5a=76e5h>H9t;C!zF|kxGV~Ya&Ys<&Fi}g{}y}D!20dN+=KJE zvnbP%g`50)*kXXy%yOUSV;qwaJ$iyi;(w=frJr^;&wu)zy0W9)KDAS7A)m{#^E-O` z?e!hQ?FQ#}Oc##fJQ|TNYffL!1>ggQM4w$3Ix^ap$!T=HN)@5HQ-9W2$1bNi2EQo{ zL>Rcz1kzy$DLg4!4S?E_yCV@H%zH?}*QP9~+WRyMp_kvJN`m>hWfInC z*F-6Zkl`Z()Uh*iwa>$^GKPgzB*A#uRaC$o5O`EG)TUZ}vSd`%w1KBqL)&@|`ygULpnDqeWJF4s{mk}8YgJM;V=!+b^M14_eD0PbA zGWO9N_Cwso>FcOL(EkEAb;PFfO^Pr!?wrVb(yg8HeulPVdvO}H`}yWGnUsz-PzpkZ zmr>prV$jfPs0%(=kA zZcbOFi(k_C6bO$iJExx8Y-woi3-X?pU2H!Uf)gL#Xj`_h#_7BSz}SWmzz~rBz*|Nsz;d6@U_-8;m58&k4%#n8?^-Gj&8I_~Xs!NCi%dlCpWr`ry{~XIw zcjjNTO~_0svA};CH(OQ8wG%%MJcL>g$IAv0(N)j8Gw`Fymy+kPwCj1i204C?DjB$~ zTS1bgoKw;uodu7GBJB=nEOn=Zms8w1ijvh)*@~go%gsZ!;k4^V<+s-BLeI7{QJtJw zn^6vtk?1|DKBG#`JSVSG?IdZh50Wk=9Jo~zn1%C{A=Im!>_s%ky=98gM#kWCvRUCw z{y!U+JkkIO949Xz8Xw?tLKl|FKRLV#ZpV50*AD)G8C&pJulD$gFOTBst?!`XlkVdy z;{LsG+~dGM8p!=9gA>cs;ioH{Ano%?sb^W8h1nXwc+fA$am-e0m6TT#^HKLIeMt1y*mff{OEO(j?l z^sX+yM%TdC2+Xk5Efsy`ZdvwK!m6gIN@#U!vd0X{VZB?Ia4DDx8{I}RW41eo8RcL; zY<8Qntr9GRt?r_PtHDyZ++CKi6|97--Bk(Kg0*nHyDs5+a3(z4JuBgv;9U4z_qovS z+Truv=fm^e^QuzS*%@}0EwSYXX7`2b3R_{T4-~d~pmkqlYi#|2(tU~PHV&HE30ao-lwsEikk4N{PTI@Ajj!Is$`@2!heWxEWfA7sNU%Gg4x=b3UE-mv< ze@S>rL|?TpT(W=qJ3m5OJ8#VQvR}3Bf-=Ado2EVw?F&I158a;cGTTqRPTR7!dcIxY zZTAz~jqck!Jl^*sJB%4{PJ3=@b3YyPD6v2Mu#+GB!w+pxZztyWT+ZwZ@8nlt#oT_c zANBDvZ0>u*QQVKx#NLk6o_#Oo!*<91PCN!GAS(9qdYB<^%XUYjpzkLqS7lwclU_Uy zSoCsgdokzZQ3`B)+an@_R`;=x#5V-72+zpAd2MUs>ifTnAyLQrjO3)Pikcfmpc~zH zvPt!590Xn*rT$)usu_>ceh9$2>Zb0E*kwLPt#l*a-IdV7&6N9Y*pGIxJdH_wnm1D% z`kgr1>F=VUSu|XW`8)17af2K0p>^$;;!ZLi`TTA_0V48BXqR0dyu{mvnE!v5C0bq_ z1iqI#T$0Qg@pwBCb46_!I5`2us^^YyKfc5L9iRJ=2M9BO@V^lW6lLyh(Rh+^Kd+}T z_j;WzKT2XQ8c%YrP`8A+6U1&RY6UNbZZZ_5oDV!wI@6~Dz{L9wDWM+$NR@$_DVeq| zRnbFjVxUFK)RUG&{YaUZXf-k;Gnvkehx)*TTG1Y=lTxNLbFA`<8K9Oz^+WSWIV`c# zkvb_yvuSxy$x4USBPA;hEW9;VeyC6TGcD6v=%G1Td0lBKP!yH@t6P2u;=4qqFDITG z_>4G^_3!$8mnilP@C=B}Nd~x|jKC}X?Z77%x9RL4X6LiVcs^By&ibJ+qB!z#gH`Uc zu}93cSy21GzS-7<$Q6p1KL0SrmO(V~P41wKMg#-XTs zfg6q-Fr6^vO^H&PUzHh+LqFi>KxzIQigtz16TCo0i;5)_LQRC43T+qw6F-W>PZi!l z`{y?|el78NvT-er+>NX8y$H9=l8s-wzs>v1xpCvYjrVyxAf_w0d86RVFHo;y_OwoQ zhQS5wl+2-M>Pza1TETBlJ*TyBt7!D_KJnQ9#LJ-bdWHwH7TGSv8+X=5-OX(;C=<*dZY zXuF=yA1F~7ZOy^Lpp}(b^&zdkShOt3+15~FHQBz*>THIr@T8(DSp}ys25SW-4K_=d zu(^kHigl)-Ji|y><~Oei(+x&FcXvlV|N2c)OHagy)5QwI8#7nv-e@c;+fr3ry4+SJ zTZ-~=U6lREg%~leViKffc6;0pm?Kqgf8Xb#e$Vgk_ELxWo_l}yw_1OG{h#+AkN{}` z=Wo!Yv-~S^L;tQDhz9rXjQbE!XP3Lo5#>U0@vq7;F5gW=HI0XUM5{{y;fvH&hN5-T zl=Ih7<3C4DdV;EyqnD}eYgBxl3gSn88HLaxgree()7W#9RMe-Ck}lz* zBy|cG>xHy*rjCU?sRLykB>Wq+YE=$^+T9Utt+Qe^3TT&bN z8IM=Ax~ie4p;jJQ#{RiyN>FFIpiuXC6YF_!!Hyq-$v8;KAw86WuV42fB&_Us946?x`wVY*qOUyi(A#&0H$VocK?aqgXPvylr8@ zO^A{d7T$*i7>@$jlHB#;FJN?A&k04y2_44%X==^F-#hJ$5=9d_lcief8TG|S`!7A4 zT2IaM$hJ?PGXRe1@&HuI)I)i}pa<;xF|4H;lp_@5Kpz+f`UJ{xU=B)GmD?q#I1}nj zeM7mkg4Xg!>JP6GRbkVy0z3fimNqF*Dp`3@om5d;6ALu2Nt)NP+Mv$#RVB0X?@;Hz zna(`jdZ10}S&bP`YV{vzJKCx;sb$s^^%>liHJK^qnpklLE8>f}4a~(2{XoO~29*Tw zU{|a%^D}y8aPxA0))<=nk8qMxD-Z^(G}I3pM+(%uGMP zrLYP?D@s;vDa*=aK5a_c%tL)Eg9Xr|1*LCJn#`Ikq>XfuT3L-*q|7I+ti`OM@sY}3 zLHiO@#wx+-GOdCB#cXlfmM!M^GP4-JI9MTb0S|16=1!JDo2BAR%g3}?DOww#7`5UH zimj%r8FtH7CCye&(=1#58SO3b#ahv zo{p2XSLBK$9j}s0^N{8@vL@k)e(cf$P0%uzV-BdAL)m(^z~({|Q`HXw_rKke|(m!l>=ZAY>PH`G^)=&TZyStxz6Zd`7yJuTnwlZDE%h zDbzG}qyuaUgQ81#noMZWRa)SRQa-i)n$Y08O66Xo1*GwMm!KOcM1ACPH-yfM$D@Q_ zr*^XR+@<2zspz5*X3}%vN$c>1r74#_q_p6hfJmOq(5WrVz z^4|a;m(m-M@~jOI1|Ui`Q-iz!vR|-ryRc!U!j8cTA(}mpXx11S9O*5x{jg|eR+cU~ zW@n{kB{Po<*t-%|*9Il%7JaBCS11OCc+7e6+B+&X|hx-&1mm8DI7 z3#W(tn{Z2bD%FA1G*bKQ&wWsYTDR?A9K~MmgU{8CeeL#z{c_tTPP`)dYG3cXzO!%G zLv}^796-;%!ab}EFmSJMMqs7L4PD`W(ZjD`>R0@>l}7`nye_FxM9HW9AYDF%#^hD= z-^DCZr*IqTC*)kfo}_X(DLYP~Ub`m6Oqg;mxmJ|+5hhdw$tf~@g{QJb6pS}U>4c7NkFAhg<{T+tN$l#o^zwvB>I*EJN@4OKpx(kxGIwy16vqFm#Q3>6l z&-y&c8LdLYk)z?l7`bUr7z1QC^T@N->oeq8oP@$oX~o{c_QHawq%_4#b4y{+5|Yir zg#R8L#?1BoM}P}HL)KKv7+CSC5Xh7W{|z0CtX@++2en2ju%Yg+f58zEY)PCbv5}if z^5WqZ$T+J>u)!&mkP)4XhjQpU)K#S@GIB%Tfi-h*r)(S$Tz8y1V>ifq`1>>+cbZ#m zDS`YCsY-%y>LBt+6r1u(C@$ko$U0XjFj=_6ddx1KDj~ zSVh8{!e^4eoE;bPEN~QccE*uRVnS2ip;^|c2Ajqc#$qlXFI}AuQACnQ)V)0a-(hg5 zM7gMmMX3(K#l^QsgOI8cXOE{~l@9J;kJ7CUtqJ#%mubttpn zSvtk=JjHjMj>$(*7Axj%kh@0G@%t6ph-|!An{G74TK;Pi_V*y;ByG~4l%8asxwuGe z#KZCqPhI^nhl|P^*cZa`7U{p_H57(YuU9nkT`KT^$RDb}AHpv$VIZ_>{c)qA{C3^? EKT_|0{r~^~ diff --git a/chronos/trainer/__pycache__/dpo_trainer.cpython-310.pyc b/chronos/trainer/__pycache__/dpo_trainer.cpython-310.pyc deleted file mode 100644 index dfde4c1114e22aa71b24cbe9b1c854d3feed87a3..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6990 zcmZ`;TaOz@cJ3RSWRpX3W;j=kMy;DA+Fp*e>v%VcZ&+*X-Hj|8mNuKvYqWd1dS=M! z3sv2eMobe~WV8taWMBhX1lTneAoFG*2@)g^39>Kw19=YWEy+v#lIOrd8eww&UemXJU3@4c>e02|FipeQ&IkdDyJVEl{XbdO$=30n9@<0%CtZYwT`Az zSr7D%E_0(}AUA?SXm(85RtSn=sZ*NvD|gDWpA}X#MODITXHK@5K{>2-YO-7m>S3eP zn2wn5EMP<_SPYw;X1LT@lJs)09IkX$W!VbW!u8I&ELVba;YMdemaD<}aI>>1%X7hn z(C*l>TnjFSmpYfi%bm;NmClv$YUe8GHd|!P$42Klc9E^Kb0}SVM`0W6{9}ckKh!$c z*(SU2Sn0HwephK-8vG4dVR?Hzj^bo1}fQ<@4KN%>2N9;kX{eNn(fH zgKlJ958`;=-Sb_xmHJ^n=5DaH?FKFl?{h!#`2(M|Fg8iC)?J&&gVg8tZ{N1jWiRIT zb>{E5gCMp3<)2=AsePkmC$XLG`L-K$cOzP&y~|zJm1|{jbS<^FgZQC6=*wPSFi2=V z&-E}Sa35j%wq?EBNXs5di!ormBcFf;$ z2Z3$5m^pE?0=8DiecMMJhYIDt?!y*LPbFLm-G&N(C+ z)7pOEQg<6 zqf%DPN~_AKoE0(?xs_^TJ=3#tX00ksWulIaiBeH;aD_WxK=&zpg%T=#K>~pz)006c z3dx=exm^cISl|xsaTs?QHt)MptHLkSs4J9QrQ|v4r+&9%dV3^~9h3L+9Cffw)+sOc zyTohWDRS@>1$nl4a*^gJk?y5y@*-+U14%<|sFr#OZwc?J%C|uD@H_L60KZAXdlLm; z>3~lijV-VOGl2m*P@&8$lnhp3Rp|c$YqA=vKUO;?Yp{7xifj&CSrYS8@xisfm*asb zNX5W9>~$n>-XX)q4e2MWK$^avO{DxiOv74|j$Vrum&k%qJiNhmM_v{%x(LnSJLR_oP! zjm%``5zRF}r7XzN#eI#HBwg|$00NJSssdTT%6iMwm4Gv}5;k{4t6jl$U(GZ~%o?p} zLbI-;wB9?1mW_!*{5!%%i8ekzQO29BHc^4Wpg0#8$yno#l#Jb9{(_!=d_x%CfVo2V z`U6qgmUrjID^Q39mOko36P!?re&lWkK9iD(9`)@(H(-ubzTA*duOItR-0Sf;n5@EP}HxcDD-uELplga!CUC-xq z^+ge?*-cZPOSF#h{1u`Xq$=gVLDPJj5)wH6Eo#!;urEsPAdNjYNkw&5a=i3FSnmLJ z?|LrK46DiWb5q6Co*5EySI!{q1_}R7nsP4iT^>P|JA{Fve6sDb<9LCaB#tArG!|vj zO;S2Vo%ErTVD-MkuxVkPDjWU-HxM=M-yd{gO*mu=IHEGc7|hZt2_x}?ofA&+*Jz8c z(|Rg7iyS&7$q%{AXYQ`!z)naV=Q|jaJdeaw^S7ifsLR^2TElC6wW?KB4J|db^rTs8 z(*4$k>(5-f_G~&)JaE9d?lxbNJ$)lR7i2HAc#5%;^*p}mci%I zo^S6BDt&+icy*V`ze9IKw+4#(|;0mi91?l3Ey)W_OH0Sc(he2P|1N$(h0 zVOL{CXafz{Kom189#blkQp(E5lxkXl1_L8|X0LdtkLFM+^~$%DkIYdmo6Bmf{Hk*Q zLr|?xzEqV|A6F*IC$Qo(Wn9fH?5Cd2W%aansPSP|XT>$8SC6!;-m8xqKZdR_@E+=_ z5|wT!$sa$2{-F2w@Jkz6BdZ*0qj}6VPjmf58!e0$L0{k)hnRR50 zmeMB5^Vw3ihms#bgFj{8S(MqbND@0{;xuzSfW~;2aZ+xn98}wDC3{)zor19nt z)zMnEHlt;0y>(VwgVvk&U+bM?b<(OM(u%92bLmF5p7+Vlk@g&-OFEI3Za?k3)U>jSKiLNx-FZ`1YAoE8kE)RTJ;@ zxuNIf@s(^D$Z>Uob0h?$vC9g+sDZy@+2c8+CBh5H+_m2IaSQeN@$>TAX77bbjx3kQ zFV6bOQQwd=k_@&9w~sJBTc%Z_72jwo;LK9KYI*kgiT)Z?NWzqUAzRrg+*MdLo0eJ& zcSH#~W-m_sNLU9jUC5C1nJ7Ivran5MR>%UJHO$G%Y=p$C+@|_1T=JGGfdo(lfH;tb z%`V9D5zW92x0yRr!ns|~$@=CkB=<0?Liz;uX%t6(Ym48eHmSS_toV1408HJK&rwzC z>mN|*htyE*yW9<->f=E_;ol=tF^Z+?=kHR@f>RO!B;aEL-v)4Ha|8SmO(6l3Of9(q zUZWD>CwC}$lM+&(yiUmrlzan;(BUz0GTVgFg(FL-G;{UfU}7CrY6Ik^@HEAO>v@B4 z5ajwb@9++FC`hP;jo^Cw=oJ7p3+W1*zG2N?)#|a47@eC>ASp2n8iv=n#Ys65i26Dutk_ z2SR-$)R5Cj3$4Eks}^t!FSlb+0roB23u!8g;?(y&%@FC{NFy5F1pG^3km1`Y;y@gN zYaQ|dgvUsqJ1;Yw^Z*ON%tF1#+w@HxUCFPL-8KI5(-3A!3qx+^$LOE@5fZYe>22a& zkZ%KcY^f#9Lb{=DJhAi|C=05lS5-qbpBN_48*S@q6{x+8_dBq}(S{N0paat}pQ&n3 z4BhzOT4D2F8_zFj1I*(zvX+ zqGaYo|4Cs~9$BNxs0yqr%6W^>T>8G2yoow;o`MPtBorp(9o15c5h|-I%4g7V`C3qd zFDC0c(sBm><1=_4nI-dTR+0H!R+TxhROWRcYU4BGSIt9cw);!cik`nCl>*d30k@T> zxsD-gM@rtX`XPLkkL>UDV{h-nFV(G~_R;lWu_dE3H-~!r<(*-{hF^6P8MM&Np@9g` z&0%R8I+Arx9lz1PAg9ig>{Eo{_1_q3Gyq2Q-?lB7aE9Arg3N@CqM5+y@7+ z?_K^U)IzuJX@RfgG9A%a50CsLP>eB2D$%AE6(qy>l|Io%R4N6H&Jc@?=$oU=27n&^_CJ+yhC77@(;R6>cso-a>1`*37&!9 z?x%YmywT%_R!ysg-K!=D6DwWJ_E81*L`9$nv5Bz)UkJGn#3!{$+k}!HPfJQoWmAx)we&i(N%}LmLAD`y$Q6dr2ps`Zw z!i0Z9J1FE<|Iew5)cU%&-9Eosn6R1ah880+@QFL;4V1LVQAF zIx~(888L#8S#)ayKiWyulI=2vB}-Rl?YYGxEv%C<&x@ktu-MB_YG)uruMQ_qkpBfGpCS>~ zH2h4?DqS=LJP?59w7gsv7xR#86t#B-kw>9xXw6s-Jw>3f8>izDjl93KciO}+$vjc@ z^0;xEk(pAQ*2I!je?Ff@7N#);(#4!&Q{w#bI84>XF&%UfWpQX?Sel&tQvUbUu{rA~ zgA46bK`(I}zQ~(o3S(IQg_6g}p6W6CLqnRs(g5V&rw@_QD%NKonk}7dAV>Zwk{g|b z{?ek%{dZ}cZV`>;VD#{;J&Hn?4UsMtR diff --git a/chronos/trainer/__pycache__/grpo_trainer.cpython-310.pyc b/chronos/trainer/__pycache__/grpo_trainer.cpython-310.pyc deleted file mode 100644 index e2cbc31f45f61a09b912a5a09092f58eedd75050..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 10317 zcma)CTaX;rS?=4+^z_W$S9_IKZC!VC*_9HvZ_|AR4uj2)pFK$pgJJi4z`9g zMO9kE)e%`z@CRC>)lnHY{jt_~bv$dar#gWaMSrq2Rh^P~mOtH^sm{oF$)9cQtzxRE zQ}*|@=BjftKH%?f*;QM{2mL2n2dW2J2df8LhpLBKhpUHW{*Zs9b+mf4HD8@ql@)~z z|A|_yu!HOnn_;tW=BmdoDQqv>_ol-3ZE4lx>@YjRo?r)1a)RmCl*-Z0cXUOuYAZYp z!e}w(ZX@vcVx6}`Crd1}H|&4C`@8nF*sXi^C+thFTz%Qj%Ix{b^OsN7!XT!Gtg>KP zmw4D|+pl=O8#mTG`)cSnY8&>;?YPluY$7iV?B$T#=hGnYWt0#ZpU7$O_e&_TS$9KXMurSI}SJOvNWa*I(}swd`>(s9^(X&9)qK znTMNRfIaitI2-I?xQs3E>?n45OfydF&SM3h9HzlM*LJbH9d2KEc)WC+m%}Jpl9evv z)I1kw(dJ>R9oy{_=@xkQCA)?kA1erh$bMtPX>@6Jw(YVt7ot)3A|#%%=eg6sMz=gS zK(tPRQmNRBkfNCFWYBT8qip#ecj9o>3mm_V_I_Bme{|RR;or<-2s^Rc6&mD2KL_nv zhht+fO(}rsZtm*N;okDYE|K_UeICTth3%7d?lRfYs=sT{XJ}41{zmx4H%^}JzW04^ zy*=-@+a~p*Jx#@3wty&CG*P@BZoDEjP7K}ZG<=3ygaXi#v$!1;k-Ob#)K*1)W#bla zFpy|$VY9LN1iA|;d0z^*M%8*w>zXuRUO(=HGR+J0mcwHuN~G_L_YR{DCAWu z26lDoEA0+CN#`oeK7I*<7B2D-0*}X{*a)&2i$cxsL|CQBL5>@76-}X!8zOg$hj#)L zr-&%Da9I@H7=w{MiJVjbQS@53JjQ4``f|9AR&IPPNrFVYcaT^hioR%^=bcjBETQ{M7J$ zm8!{9?C_BtkRYCLJ>?Tu_%j%Xml23OPzq?mOOz;jFi)ZXq^hj?RoMrE|YX;tHwo z_h_~I616v%XxnOUKh0^IdbJU=n@^zRz_!B1*`9kKxRlpBxUKXKvGHvcG!Z$@VW!}8 zCawtMmWK=ZjqyLb^e>y=epU zsCAet^jf0HU1#DpLF+1+P$W zor2F%K$^j?QgDufDgpvq4rP0dGHy_yQ}B5z(n0rx>2_i;H(aCTe$1)G>jJFTW!Gx1 z4=~I}>F@00P$F0BbmE2|@tf4;K>vmz`rz@?lvi||n(s!D;|MJdNkf*M6i6%3Dn(odVCxU}j=Qut+B&%nqLl8RLx(4~ zL(WoSRvHjm2-9&u)nmuGgjDn~1bH*B(x0i0W`CACp-rfR_%j|%X=PPI>7Z(UFrFJ% zhgD16oXwDJ7f3Ah`xWZ?{jxlVh3oRdtr3*ScCwO9aH=c1q8>tmbNDXSHa;mZb>>9IP#8dqLM_@33YS+L+1UVjf9on za!!LqoObH|oa9|Mb}#INy>r>!d}NV?lsnjWk+4r`Dm`@@=4Pz*^ljy>uR_n&o&hb- zp)9|xb~UDU^}5nE;sSEaE%kzOGwkLPozDXmwP}UE%_7sM72uc#97CT#Ep64{O9@&t z^(z|E2$xYaz@3S92Fu;kyM@?F3UO&mjIrt5m1u-kv+c*HccmI-#C53zIy!f%@eTDGM)lRtSnrSK=wYN*>!-_P#cio zta=*`w6Nnvq6ixjk_685{sVk|eD6QL*Zb)|{>9mMRiQIUy2036jx-^?s4|lZ>1}k! z83GWnP^vJ9mmk6cZBYryu*iY6fu0kLiCnbe0%t7{?f_T-<_Z8KcF%vEiu00ZMnVUg zOUM)IMLta>r54CSBM4g!fQDXkL9K`d1xx~^Fd9Jn6jRGy7-c3HsB=!TqDnbcp*N^2 z9or-fTE8gvjb$cO2YitB1iK}4kb6-;0Ky#c-0L_UE+eN(H$vy>?e2nm$SfIw4BYwE0U&uAJu z%9?|s=65OTEC__sD<-*Z^$S|p>=u!-fNA(`D%*P3O7gu@VuF?wL2ayMos{DCp?IJ< zCGt5C&qJSvG@)?!QT4$v6QvGxiZ><*&G zag35jUjxfpLg_uS*P(kxcL=G8howkOW@UPJILS4q*zmN{9f_x@J>X{!YZ?85-W}_X zV|`;;`JQBq4Nfb`P%;W?vWMX2jKW5e;bb%!>8EeM70(a=p?|G8yQOz05Zl|_hhFpD z$z&p#WTR)4+dn|=+}r9~ZBlm-qanF}Q{Cy9Xwg*j00aF2Fez+IhKMf%oY)u}$1cqz zLs|JS)gsvI&L%UxgKR9B1$`RZR=2g@VfyXvO;AVn^|0=~c%XL#VEmLE{V0vzor~v_ zxwx{Wy{+=`WR6w}D4px>mowV`%V+eVBin#EfLsQ1#nZ~$YIODyptF;`Q%d)VLdd{~l!1fEUMg)KXT#VBHQCpX?bJAwlzS)8?$EeG zQdCm9hm*tD>yydBab;5J9*Gx_+MgUrrt!OoojwKfnd=^96W#fEGCobYY?4h7f$dh3 z3Y%Cpup>V~{>PZoQ7ImOoccidvE*1kFFBS<2^-Gp9|PH*qFHv2Q`_!|_{rpWS|>S? z{%&dAlUXjw73Ja=xeM_plH-uE1v>lW+kKA8bV@uD*IE8UZ-aiFBJ^Bl$ zu0@}EuZQ2I-UZnsY1|7ywCQ`)elj^q^OWsAHI6+&&f#?Ch#lOCN!&?O$*E+X z&3rSDt9KTr2QEd9+Y&Acc%ur^T@}O^=>j$Tma^F6GPcVgA(s%GjY4b3@MVGD(hUvg z5jkQacnCOLtI`c7EV5Zr|7wly<^HEn7_ap`19dC`AVQ$6Oe z$HFP0#*g1h))^nZ2_g(<-$CWn!d5mQ2w#I1Ix}+hzi<P!x|-M9 z{Olj0G$Q_tYf{-f_5YoEzj*VU!}mI;4TGUMgVgS#u82h-mp5-@*~E5 zw{?rTeX}HW@9}q0i25aB(sAenU_6yP-;Z2TXxC*Ia54FyE$|#?J`Na$E?A574^}Rw zWMvE_#@}#iD=zwhQ=Yt3(#S~-TDbQlAAwH}OhhvC9bUzM6QiXT;}mez6x1jnYl#md0D;63M=1GqO2QD` zarhEP;scaB>egx<@JgwTDS6@zD#%HD7!HEW^CnEpl}uzqWh;+wG5RcpIXQp zUc`SNg;ldp^MwWZ&K!_J<<&y!mWrx69x%vzDS-u!OPp_mRg0j!scBdp*f9%QZnU~% z^bzY-8~kr+6a!WUQkSKE0LSY550Jfk6H_i+q_a(<`PA_V5pP4@sAJ4%*#U#Dma+&@ zY4V3I<7Gt1rqX1sT5y`fdZ)$kzog;wQtI3MyELeA%LN3diUrjnYX$P{wuHu?gdNb~ zLPwWEy)LwKEqFYF8jJ{+sy&V<*&`isWRk1s?-1m{(TSBCMj3X$rD-1+Ir0Gvs}t&s zHlvQ}$XCbJ`~yQDMQG_0@CA$$bBB+>QY~#zJ%aQUzQp&Zkg~w&^P2I&+|NyO9Q_lx(3|A5i+DrTY^KOfA^{nPMM)P4qEDjn4^ zDtxc2o` zEjlaPX#C5_Ei_GLfb8Z#Yl|_gZ`^q5r1Fl!$U^|51TnOiX^IyUlUb-=ddFa8#0RqY zAmT$=d>HWOwQ4DsIloz`;LQ*3k5Xhe^@S&wHx0X8Uy{;{=&HXYrF>Jj z`4ayt^oB|wr%JxN#J@*r(@yV*CH~ixLUO@c(b9wBmN%c{5R#4 zNhk60fJZH;0O2$L2MXw<(Ox=|W+~%;5QrSiW&$h13u!#27(|FAfk-YL7O+DZkKNGp z8q)|IIfI{}JX%6vL6pg1My@py?pFG_?Et{!2<+;Q?Hm&Mzd#a zlNYMgD>L}n--H?58`xH)d5YR95lFK)NVTP86s}w+uHdbM9Xa7@H7Ec1VAF)x(_3iA zD~}K)nl?@MTAv^_h#-}qL~wvXHME`rI|q)kZ%A4vcbORHiVMf*uc09*)1Tr;sH;L< zd;AcrGvrs51aAoHD+%75vU%b|_7bj}HE$QU+i6%C!UJA}_ds{mV+|gqqGeKbcx^ek zO5zn^G(%jebm5B9N`v7M&xz=r=o#ej+Z41ZkiH;;+DhSSMVzFylG`PD@;=mStL;>} zyvNs+lQShWipZ-t3eO2)9wFf7)JH9hZ_qILbUM^5et3B^Qw%+j;_&~ac#%BI0d9PFzAcA*-W^Z`N0c4scE|P8 z623!6jhX3If&RF;Y&$b{?2 zgOM*E8P3ajJ>r)opX`ob_S;B6KNMf-YH;V^wK0KF^w!y~kr+VsWAMFdJ&jyC@U`Z9 z@WIxVZNhwdEh-=lCk`{<%{-ClJp+DK;4~R((ziL2sQl>|-di|g)YTk+8o9;W1G|mF z1TUJ8r{M$r>?1}Zy?W+TZ>1Q3#kB^6O8;$uZ)5(vN5LN>sN^A(5#E~(ydQ&=$u~)% z;}wO!O?7@ibxb}&+dN7E{<$f`J#`s07tgyO58)k<#i+jPAm|LlhIGXqlg;)mE-3tWMOIzDA1&;Dc-?uM zN{eI&yGi~RSo%);SPe3#Li~q-G{z+1=S_?utrs!VAE-YWy0l>ztpnqvWkIT!8W8j% zBPJfNpYm<_md?ki?q?`?i2_nzdGfntYDc++c_zh)vp6FqI)%Vc%6m+WNFd#vrZy^L YCJ1AhzWBo}YMMs5G?-ICp*+Ku5%dj{WCQTx&sPl;JE_gSUg6t0-Z_tRogoW}R8 zVB#_H_-Af=PCO%?5Kp2sFWehe>)GUACCh4f0~Li)yrvcRL#ft+D2_Wr|CS%Fj>e6~ zb(zF`JCKa(S4@4M4P-EqDrWJ3D=8TF6dD;10=Ax8W5${v_F1fXU&f84k&>ovncdjf zT4tBtySB`(U8nc&@oufMbs3^V;D17+^uUkVu9TygN$w4DD~)1_)EY7#3f51!;-Qu@ z#`2F1# z+tMuF53s6W%!Ge7)LeZ0$u;|G+y;-Wo-*6c-we z71r$*-tHJLYpI4y!$`#!W=YPTE9%2Bkvin zuT4(PH0MWG6s7i7Ti7^l4-UqrI=ev!s~SXV$xD=i3p+(au^}aU_tgY^xGi`w>S(9$ zZcZ$)G-VhjVa<~*1_$s!j{zG#(un~fjsXXY167(uJ?sa7^!wou(_@GELF4>X5{5a% zB%~3EQLEhwn`C`BjXC4e~106sVENye`4c;Kn+=S`a%sQQk^%WPs zsqd+12;lfVP1BJB=s(}%W6T>3$?I6}#KD+e__0?0Hu_;!pfKN%7bG!!$^t%y3AV>$ z7LBw&^gqE69MNjyF7Zpt&SrQRM%pYhZaJB|342-X8ahXC)>ZU4U4-YAYnL~q6dRG+ z<4U0KnF&sS3fiq68!TpzO+nH*7FMQh<10I>@A-Prp)IRA){wc|Q50kqU?x%Fy)%=k zKh%-(2CG{#j3f1tXVZG|(Ekel%J^F+O6;Q0xAvXX+Oc-*OV&qE?U%HZmWU#6m-by< zN?q;VcGUB!tHP2;+Al3Sh}`TPaRPb?j3x! z%6!eStQ(qpyXH{fT7yb+C^&hk#uFAh=_Wb>33ub*82g>-#}`bAzNT)t0Ai6L+fs>K z23*MvIY0vcrexYUAV1nz`=HAKPLmU;9(Z`7BXTU0qA!6ttQJ8^i0N}J*-`>dAv!YA zNx&jTd6*$08CnD26ynUu3TS(T)fBeCz3qt@@#+FhWRpZa{YO%`C;dhd=^V z1W~_F3t3|)A^R{(YwX=kY@8_KD5^1X3OAUiX%rC)`R01=%M0L`4=*BzEVIpJcKIVv zA}n;fSEEp(?Dtq0u`AK#{H$Wu^NAV|(b%l96Y$uuHsPL)II`A~TO82g={Rzn~UfV## zuU(4vLSnXf?GNK^<%`bd=9RVgRJ0>K-6@!0tw=}@PygAX7X<-zKhEFJVW#*DUhZtu zZrTld!LH%$eCpb+J!h+@(FC%oAtB`~QlBQbd(Goikj+oeJ5P0rVW>wEi5}vINAWl1 zhu2Y)+K6^n*ogLZZD||qfY`oWx^3@Q(uyeUx@tYO4;e-mig8N;@vtY({)L5+)&1JXPo_2XMp_N+)ZL#+%V{N@S+vkv z*04&#>yESEOq*ux7in{+F6!qjQO6qSt05X<=8nBTTdcELtaC1%GwTdtaZb;qv$yTS zYQCS@qTjr3n*Gg*d9%L-uo&(8i*(+M{~}!w^JcEugV{rC|CF9fo%B??fZaU;+wGd*Q4P!v>xkD&!uZ9%uqrBjFosKJ-_FBh|hi!aD{xMKzYUhg~TR% zi)})G#Cc#*6+K3Q6sm-=Z@)$#5+=u;`+-s0 zIV;`2G(WdPY{x{RjGib|iy?Rs$vAvb7n&1GDd_8JHdhF&;$;3y@kNiJjSWL1TdIbN zyA*Y5j$ZS4Aw+X!Ht~P9P*hCUc07uqR7SBP6Rr8I+ynb)MIoHkL%G)hf}%L5vPWow zMS2;AY)$~tQGCx32&j*;ROSvP4;9ITtmKC}tKIKjJTC``1qM83i>bpd*7G7Dk$WEDq zz#!tyUr$lnse(}r^+A8D;hU(AFW}`a)*Qo?Wf#2Jv=_meE_jp5P3H{W)Vg3VP%o#3 za?PbS*H-7r^y0(RbhdE6n$BiVsA=^Q#(E{ZCJ(uMN~(Xv2NBDzt?H?DUd zekrx^4T3!cogRXCkIHAQMUX0Pn&mt8feY$!K|JoF1)8ZOXL0=hSGPydS5qx~T=2%r zEJi%#qUO7AinM<*(gURLQ39P!A!3r=3WpKgD#ei)FiPNrIktprmtKjZo*wdBMI2n8 zTp;LTsFj40!tE8NArx+*SR5-`f{HNWixiU-Z!zVefMWPCe~VI^)i~Xx`*={;SL^VC5Wo5 z6emNqiWYSE4UK=DUcaTTC4^h%><0IPaG?P2^ddG!@AxsiYBL6NggD#z+4_cQ{?HrU2t-qeP2^&BF60Z9R5+kK&;UG1Z z-?ze9*V5(j3pVq)DGKmHo(>~ZjWOuN6u6bPj9D24NnB;16SfMP^N|jzA3XjuoeO*er3RP_=%$1HSq8F zB8I@7h;AMLykd;&sYxtC-t;}W5~&qj0BJaW!VnZoP|(r=F-(zodp)MY8d4**3nT`-x`nPWu+oz*xSHWy+Yl}W^Aiw zsJeaX*lr;naFZrRr`2aoTx~nkiAvaTW_Q$UWTOP%-VgC}UViiMO#N;|D!ld=42g;M zD>c%kP^RG29Ph-c74WO`=@UP?Gc{t?E9lsAjkZ?w)18{3SCeXM%4C*jsp7dwt;sE> zJ?}}a=xGAe^fb|FN))1oeUJ5Yj+Ga(YNsP2uhYpIolY;IA2K=}e@=av8DmM7B3Pqt68wn;`Ni<-olpAqF06ZFOh$NXzDnn`B0H!)8V e$hQ$wh32=Krx%=>(?pJ1s5np0-D|uxclv*@7BUe4 diff --git a/chronos/trainer/__pycache__/optim_utils.cpython-310.pyc b/chronos/trainer/__pycache__/optim_utils.cpython-310.pyc deleted file mode 100644 index 645b9bdead61e18fde2fe8fbad4c2fba6cdd4bb7..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 3782 zcmaJ^&2Jn@74NT^X?Mm>mJQjEET}fHYfoa&*ja*5$ZmE?;@wD2vWl|`GDPE^shXK~ zdb%f7-Ek(g5t1WLzyWaq<&fOqf`o*$;)eDN|3IA}93T=x9JpX4;Pu#n%C8@ zs$Rc(AHVmU&dvrF#xH;K=O~!7tUuDr^ziWVo@LqUAtv0~u(-{g*d96?j!knnb~ju+ zz1SO8HY&sFMs>okZPd)VKb+Z^0dIlN@ZgcR5%5F2&S&}DW0yDhG5*@4>c%1N^5gt- zj~#xRzrjy|Q|D*+TYMh-XSus>H5T(Lu4M(?flQNBwK5q-iIA;yltsgKo<*^0j>bW- zJ_x1Y%&?<}Lb5Y#?dFo|4g}9*!3H8838`4Fu&Fa<;=Yh$cBvG~jL<0QGnIvX5iE?P zFswzkzItnsUH<6CBD-;u=m&=KNa4eAgM|ra*+8({qVw@pc8hdlS=ea?!L99-NfD|v zQ7lVYS0s=aLo@7;)cGJ-X|l_r8;)6Druj&nXYccHc$;lU*??_}s6WWs+^mPAQ5-?* zJc)X#9L8gk7sfFgg)$tffHB4e6z@bNUa_@sEaY|UWG!~{+PaykZWv>U)E(q9Q%npy zf^%G1HQ9OwekeIflA;@p1dEK);z2jgxnK!|sYTX_LM8AdDF^3SPo_g?fUDp}DMXSS zx+H@!(nPRP9Cig8iV&tK?QV<+vfH5?=A$K?PL{fL!wg*vn@Uq!S`(7kf#(H9#WDW9! zQ%I6DRr*y8$#0!7Ls-J_YtS{6Y};Ibr3gtAOr@0XoO_d{JrYz3p!>a0x(p12!OAji zf<#sN3_5gD2_G!MVz7p>BkK;*s4JSmbIR|Atq+Asl4fBRAvmu$9PMAtgbX{e(3M;H zh#td~i+FD)Nt)ME${$@tM&vTtd&O+{W{d<_WJK?DZIU_}KAu(lE)j7243j?o9B0Sg zaWZSy+Hni(iT%{NY~8uE<7H0aW$u=@>x@nm-d1I+`UJr7)ZVE~R{ULOr;3#s?r`@9 zHcrj&)QZ|xAkSupirTJSxc4pj_PzPSeNrzuPi@;OD(X^E#jmzC%e}jnyjnO<=877v zVXu8KzgI8p!s$7HnM(c*Tx1=omX&Zo_9#ms6N9+QxX=k@hCzLBl#%{k=c}! zh;yVh!jaO04((^95mTZFTacoU>_+}8iWMFW*|YEd_+T#-K{~`~dROqJe8gnBjU-2a z7$K4Zg`zm`s1foQODZkXoo%Oi%#lrMz)nNA(*UNs!r&c3nlqqq+GrXnx;DAsD#EXZ z$aJl=i_$@ayab#jU7?*M(KGvarf1siZX7DrZfn2YP5>1&6A4#@ zntJBS3sBR6!A9daJuAh%Jd#4S`!eL(iMVz{&b61o6uJs{g@lGLE6_$)j03g*!ok{8 zLL}OSb{n-wwaM%oJ|@Q0+8+Jfk_-p;ZnQpD2$R-onuM*(@O7MqT(v$NcVxudYirk9 zAIWqJNYkDao!0-aXnMAcFepvs;q&;0dL1P2t9IQ!f!`7PGqz*Tf!6K1iTq*(Up$KLZw|`m3CsO zD?MBwlV5~H`6ZCOSIf*d>_uiCcjTM2@=BJtw=8n2z zw2<@I+NhbhDUaKdw1TN04#I#Noa$uNuB1zPDt8FDwC?5<6Umt2O> zy#VX10AQh&S^=X=j-o-1dkMsU1x5kKsEpOFSq204BcSnaJY+C zPFTARVB0l3!2Pa`w@OhdT%2r`C*f|K-L+3zM=eOf+Hs5a1B@}-bxv9c42L5)c&(@! zc~$P?-Ft#S*{T-Rqt;XVHJsk5O(c38h2pqX)Q(!jN3hJ~8mhlsVxN8QNB>hoH>b_W zKK;ZLQ`mKL0_vO7hTSCGURgpzAszq}7lvgif4QweudykT9H9v0Tp%0}5~!JlGJ9vG zu{nXnRUoFq$GUe-KL@=G>Sb(m;;Bv44V3)NDKD}M@33VkW{QoZ!uy;^^mowbM=|T; zhw-jJ_gqf(+y_zKQF!4?iAviPcH^!K<9HRvjL;i)b4R{v;ZK*GScvk&Y Y-#KFYPTdY%=XHl_-kx)u=-i3_0BaUK82|tP diff --git a/chronos/trainer/__pycache__/orpo_trainer.cpython-310.pyc b/chronos/trainer/__pycache__/orpo_trainer.cpython-310.pyc deleted file mode 100644 index 103f696ebe89e34c45b6fba4adead95824b3ec88..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6911 zcmaJ`TWlQHd7kU;?CgcNLD5^1YSnktIy+KMj;O&gMeqU#9b;hZ74%ek!1 z%t|CXiv*=o1Sp6YZC+9%0QpoHy&^@C7I|tP`d0U?59h5Z(AEzv`qsilt$zQRU6NMh z4zd6I=bv-_bIyPM>tRx_+X|jP`I~?1{Lwi@`DbdJesnZmQ4}@(4%;+HNq8@DhmGaluJeoU{m8kKd7FWA&3WhcyYD!oZ?-L)OK&Fe zrW0h27cM(jU%Yy)<@iao4qmB~_Im+;&`AU4tgWs(N$g}>f#Yu_X%IU+NP9_~2KK{F zw&i%SlQ5QE;a=8BoL&%nVYWLO!Gmt#XW%j#bJ^*{j_0JEcry%~WSu^D9t3u_A4^O)zYr}(8-;h7WKt*>9jB*$U=xiPVd}C z56Ov!|HJ;0^YKrezx|nH_MVh0v1>UiBRp(5_nej-Ex8kTxCO}!JG8*9SOpDQddaym zY8Hc5U3vCjAMWA>=i{H#d_=bGcf)>qEVHPygQHr?f~c2pF9f$RNw&SMz+=F@J~rsZ zIM}3*6L7w@iWauUm+eo;PFt#|dU2d&q!n?xM(r4)#_43iPKKtL^s-I_#J=Ta-fH5p zfVVVJT1_@L11`oosT)DD-5yW;AWhM6>Q2T3FM{ge2x(9n5Tp9=lwV8Yjm{o6oU$+J%Q@Uy{z1@K}&~v@3F*R2Y;p>jz^$xT{c)KHYVB|)wHI=T( z45l5Zxi-QMOUyh{&nVObzNEDG*4xlwQn)L$EhpJ<-Y-`5zSEC!CTPl~zkk{B9t0fz zggFTi*j~B~yqBaMdB)41JjU~pDs+Z()%YA%E6fgf?0ljE|5yC*{O8N9vQS?aRj_r_ zWFw2b9f;q9gW;mmkJJ18AlM5;d4%X`KN1x`^rD`N!*7{<0&MscfhK^c5U1iyw6OP) z!k-28(|1?Cox(M&yq?6~%B|#K974s?l{a_Sd55{Ht8cHo%M)@x?oeG;h6n#e4pr1$ zY3`VfCcA=3QVXD=EvOe&L*>t-;p6wYhYs@Df z{iF{4Baj>#0@u+p~1uV!|)8f$>HK`A$x3EGXUzOTe(&>G#b?s#5eW$2{hP7EoNawhqdwl9#bwy*o2pNsZNMZLD!XE1qQv;2cvRqOmn$tOtmAfAC{^{`zHTh8L|f z&xNTOexG?l_j`R|t#3dVvX{OByEj?3+k>*5P|Ap4>tVn+=`4DTt@k@2bEUWH>_JCs z4};F;R_3z6_jWgbtNEvI{L|i}m*Kj!kn?LaXr8z5;$H#~hQ~KkQDVV^jvok9&Lhh3 zHi%rDzd-#iQ$HO+i@!*iDP06#p-(RnAQj@IZbF9*2+QkdiSMPEsEsh!&vu0U2KU$< z-wUzEQd+#v40U>W6#L~i+RQX-|e>{ki91pxo`4Z`0dK>p;^9;Oj@T*$YxE_0sPUH3A7>znA5(zG+`-qKT+ z`_w@`t#R~_G>+}Gi&M{4rp_yea^=#e!!A#?O?=l+wM%Flr`oU^W}dVUwIgL|}(#S3#&V`XL`uGo@N%FD-;D*OVx1f*B$miP5R4Vcxf+*0n92KBs_*IDI;a{o4{ z_6O<@k%Q;TVKujp)bD76#$YTrx;0jwR|eyGqg#i~j(;CJ8B9WT`Aj|q|1^^~_tiONFq@4ZPGIiIe730^sl(oL z`7~%#`CL<(RR;4}6Rm1KpO4`^jhSay?Z6x?u=?Okre(84h3l2xWpFk>%j(<42P$_! zhkxm-glF@_V_zMd%g>Ey`8n)uY+hkw!~W;GV{qNpfjU^Ev4iv3LcUn^$eQn?j zsgR1K;tlGB>`cB0`CRzE2IvqZ>1^Pf^YdRoKMNUM%op;D*yTc|1i2gR`$(Bn4q@A{akP=m!BXeYsim*+Gw8K>@b7G zzX+nxk%+)ek=JNfP^3_D)K8yuRu$y2n=VyoU35&z-%P4Cb??QexF{JU+%yY%%Qr%V zm@vJ$d^{YX5Rqp`rp24!A|KHlM%ZC4TqMGL0Utl1&ECbRDtW#vabv8pbwy|hZ_>?v z72o;S01%bEjMu11QGvfo0FGM`H7fohcN5%uDZfc{imbv+w>(4{nFmtVg#;Q&EI&)A zjs2#b%_!-xlkm#xG?oG~|2BbNBXEmA13>5y9Umji=%C;cW2b5bNF(J{`i3B$dH$A| z^nAY`^+Q?ONO7h773yHhsEho<^S2**oCyo*ECpg$XmKn|8Fs~(A9i}!lgA(u9=XDn zy9}t16%|RR9tgOYT9o=dCId54k3^1YJtIStFlE@`A} zD%&Ue`8Fm@e;Z(2E&iHnQ)Uq3$Qx|c)*676Nz|yVTjscq%tF)9GSt#%hKXDQJr-aI zC15k^SCN6>n?^JmX@GA@)jl(H<9{3G;@^*dc@jJq@I~^~F$eM!O*Q^!@7(`uNo$`= z#Yx-9HORE&Ux2_l!Ssi&OG6T_Ha6l1oUt6E1GDmyY!hC|EoOg+IweQNf_E*q<-49&<-3vBp5Y{%VCY@nvxb)rvJI-fd^Kpfp+G@f*eR;kq3QQ%1X&?Z2z-}74$y{!<|qn0 z{yRjXVDO|yR|}1nGaTW2O8(D=A~vP^8KlCG=k#$X8daFaBExou^fTGCS?*I`N6a%V7%A>Pk-UlTxexe`Iy#wOxVPzC?4{c~j^2pF)Vo2wr81P)C@fgPPAARWWV>z3o69}Rjgnxwmu-D6 zWa?O*iQjR^?10_9Dr`a>Xh`QZ3ZG;_zpvfb?i3P6EoN+as8*IS5HV9*eo?5~LVfV0 z{Oud8<7c?_?A`nsy2|u@QQ5ov6j8bp0)n9!0`SdoX@kPc4^jkBn`i?%R& zxY`g#H$jMz6=8L&!&nfzDV0q!oEZT*MH`lp(HK7~!lFSuq{2lh|6SUOSrjEd1YY=k z-QVs>lY&w{QKgu2#XqHOQD~i17vZwVy*1Rmxi8uviXmCL+BAaxuWK}_xoO0eBp0W|8Ef>YvVs6@P`CWmMXF5h_(bkHX>JR z!3ylArQ3IF6Y`5siOc9AeK=NMN~5T_BVAzXfeuAfQOw}Ng6YgSpi3^3t{-i(*f5AU zv#q;DZLbSkhERDDxNq9(Z)}7C-BU*H!-Q`m9z&Dtlgma<6SquJK}7akI#>QWC<56a zoI!A(*z-RoipG8^5t|s^BhW2<;Y)swo-!e8z}op=5O1r9tT#- zkoisrH!dzQMy|*$>iB2Wu{r8k-1(on;1b6-8^tFHk&w%Oi4^x?PxKgFWXl9rCUpGQ zX(BQXu`rrwv~-F)NT_iIUZjo`xMhY|Toyb}P10C-Y){i?VcjUK;wBXUDcO>t)Fv#A b@(fF(9Gu?c^hQ=fe{-sBH7w=3wetS~G00@M diff --git a/chronos/trainer/__pycache__/reward.cpython-310.pyc b/chronos/trainer/__pycache__/reward.cpython-310.pyc deleted file mode 100644 index 4bd9cbacc8cd5fa49928697fcf685c1fa1c14db0..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 5638 zcmaJ_+ix6K8K3*kt}k($lF$^PQ6x}S*z2@~0%?*aO`2XvovL+<>JrV4=bW|2o}HPT zbJmWnwN!9gcqtDEcqvFR@RGlR|KmIs2^8^y)VDsg{Jt~0>kBQjns3jX+kE%$ccz?~ zX<4|w{Ov!pf1kFjf6&9x%f-Vr%d+L$XfSKVVm5PfyYH+xHnrW{?RzU;-(T@pJ8zpky4~or4}Tx!yjWA~(GQ}nGEG(&s)5MV4ju)sxXM&k77Gic2MP_iA}P8RSK0AYhVqM^)lF~;9FdfZBk;=NKkOK zU5d?!mD@#*O**aC^S+3ox+kFwN@+i#{NlIMhdAS`) z$N`b6B;8Ebc+036yr?oQ^G@qAMMK-xQ%O;jDj_Axwxd1bM~8PR>PYgd7coow_Xb(I zse|>MRgp0rZ1ZewUFkqpiI7}3WZ6^w=%Y6%5M4>GeG`YeF>($a^m)G|9@8X+N&?1H1cYwn z3D`@~pC>z5?rd)9uoAq-4NoSjbhsuGmO%<|gVjhNDH-kP<>BEPr7^M-w^nSVGKbkK zE_;f7i-oNDz+ds$86>ACkeYn7o@A%d3Yd!|byhbg`DE}ac39z}T$rSTBPnAP@Rz8+$Fm%!!~-Mvl@hTwJu#{l(1Kw?2?u-AySGk1uK?Yg#u*qNmUp z+CwY0hE5eIYB|h$$9iNhwO#RJjDq_+@5k@mxpDnvvvu2rwRqk$gQn+^C9YuN7jG=S zE4h%1x8T!@H%IX*7hf~}9N)hE=HgqT+~BE-N5QrDjfq=_ag6+<;`n>mS&{()d)D54 za-zx3c!Cb<(AO?PU!OV=`d>6Bp^qToa5lKA3Hp>a9beS;ag$CypG9M6JtVz2L-znS z6qnHT$kx8eHIF=T8PDPxk(WVqbCOhyCTen(Mj!{JsyHUAd`@aVDMdYejGHM;Yv6{Z zZ21>7hVBOJL&ftGXU`jY5A6eM&ma2o6nU6(lsokIZRP#O+VGjT?>vJNy6@}-D)_b4 zvsu6z55l2;)7rQ9?Y+jZQO|$I+6z?x|7|?5RAcBfi#0!Tp0?2Pm_@C9XSAj{Y))Cq ztJl?U*lP`gVR+UWw$8#|+O4H!DCo4@%&J&cJk64PlzrDA0IvW1{jrU|@zYBwJ&i9? z4Eg2rB=LDPehCZQo!+yD_CxEyf?W2NCE%IwS+;c#nt+Eo;tnP*A)O*qZF5n3r=y+D zMQv|CvTGxeX9{DA=H)eBAgze2&!Nxwl=ER^?VdmMDP!Kw;;VR-?jQkMyh_mlnkH5mz4#fg3 z@nw}G8@ov0+BLC`T2Ke5D+hh;?q*ec>V#v8H)yG+07A5nGRox2OpRsY8e6P{Wf7ZX zx-xD4;u_l^i47_)^^=GBCLST2n(bw5Y}6`VLu($lNhPy(XnXjbv(MPm_67SX(A`rf z6x?>i3s|t!_Uat+6WWX7L;MuvQBVigb?SJT$PFSliM#^RZctDaNuS5DZpAUO_#mhD zbR6FsB>8Bh5yz}d<5;{z3x7t0I6b2DiD!r(VSt=Y5}FMy&z^Q#-mGi-Q*<4>C{vj*3-6Su~1KUXngd&f@9@^BPz4~Pf zrX)t!Uw*v?s6awvz+lj+Bzy*7irQU__^X!~?ofKu&8kEVl0baRgdqu*5^g2r29ZYi zuPh9Z>SaPIgXDk+jSOa+6Tq;xFKyBjoR1Xo%KZeei(dj>ClTlxsq0E~Os5ylSBY9* zROKSr0uayp&n-L~|73Nco_@Z4xIVfv+M{z6QT`8@Jc=Oex09r1i9zUjR#gyHmC*+a z&v!1)M<93Uag+1PoH?R|Z6$>wV?^6L&le=p4QKO)l|JF68Wn#n2B%|16lpiD2D)Ji z*|RVFt9H=vnr<15$0W0$03s-~l>sU?ncAs%OeL~qc#TsqY62>y!Z*IC4KJ+00fqG# z8~_;DLzmzIih#FIKmkPp6lDW6AkbjSWV<0Yu}>GWYJ8V>X{agJWr%B%>SI#Fhe$p|ycC>A)5Z;vL1+ z%I!Uyifr!m-94x8?YTD1@|la2{>;$c2nrk0GO@I8EFlh6tI7SBbn%#3c2@%uac~OZ~S&{?|>u275Iy zk=|n5gsL1N!2_TU7#OvEjKP>Yw3+uOhxugiKRIg-@R+sXJVc0Wctihy(3h_Q;HrRv zWdolth)w}h!4ZO`W$`99*B(9u2s#mpGRJ`hXKgxA;b0{tj)=$-?WP!Fq5%Q5shgAZ zF34Kq7erQwOacutcpRbDKp-50W1Is=)j%T_PUwuvBvn00I5-6A4jmZn%Q^3yebQ(q zlh8q7onBA=B{o2SSxh*a5^nN6yEb?(1C-@^8LrVlh z0phoUQ|=)S}F@34Wskt zBciI@LF2^r*~h<&CV=k3Vf{Rx_b)Aui|tW3akvF7Jr?yBX#W4*V!p0lSJDAXXrr&+ zVjf@DsLWkCZ|2oI;`2zkv+IJb$1Y5`^o2cdvadNdX#2|Uw+6*#fn(}eKcu-!Q9CfW z4`mVuYbCccl+yRH9ta#tt!jsAy_e(z{)!N#fU|!@!`^Di@TCDCYn0R;38sThe1}|< z;v$XPo7&zw5o`JseR*UfFdp{`3Jm!jkm-<$+ca(h)N>|@&e_g~bBIUr2TUOJOz#yR z6S+y`*C6ed2_6E*y#ak%0RpCaj;}&H6&I#SwsoAr$Kl(2O3kYxk diff --git a/chronos/trainer/__pycache__/sft_trainer.cpython-310.pyc b/chronos/trainer/__pycache__/sft_trainer.cpython-310.pyc deleted file mode 100644 index e7e8a8af75f436c48b3b3c5ca104099ceaa816dd..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6373 zcmZ`-TaO&ab?)m-PtRrVH{R5eZE436XGL0Z5QdfmQI`avsUTz^WNe|gx2tETXZtd& z?%|S?9>mDfc_VUhUINF+?u&suIC--k1O$lk?uR@H^_Jwt%ui4NGT*80SuV+_2UA^r z>QvRKQ|El=RHf~F3Pzx^is{Gy`#H+`IabbS1tqNv$_pu&{C!c?Y3YOM7&mFjw= z_jUO;`Uc)cWX83At!k_H>!6uYJ+}Imq}8HEZ1?S|z1eS}y&kpWPQMc`^cN)EiWcLg z{*tUWqUCs{zar~)v>LDV*JQmJt;ZYv4OwqRoAFkEOV-=bx%hnle0-sQA->qZ7+>mN zQk9Ps)?o_|6}E7s^m8$Yn^ z8)f@t@6h9clTF3}zZYf!a|U3)oKKQ4*>~n0Z2P@Q80C(~(@DadI7~wP*&CU2`@=i$ zd%2edxwAdW(j;;Y!+hxM`e|}6;F*_)X_D&)O*?D>pB;hpRJ#RdY9=QI{gH%C0 zj=Vf@-gF?uWiQLZEccT9vgCH6G|QaW%MQ9umOA+`a2O`E$Kg1LAgU8)PRN2J4}C9k z@{%JZdU7h~WpvM>gSjr`+Z{}z$c0**1r+ZmbA=YC79-9F_)|lcc$Gkfeu> zw}*jx&L@F0NICS+oG^ES`{O9|!#sjCS9_SP+Ci5T044EwcNe4I$3R%xuRnM9yvR%Z z;ET_^$$jV!6w+kMz37Y2<~8R{njKB20)O$@Y1K&+C~WMc1F6R zn+*D9>al=#HBq~n?(YX&EQFaGLv-43KgcpLuH4Rf;K330pLqeu?aH-oq{$%M2VuD) ze3a zwkpNAr`nP|Ah1t5Is2TV;=7;*N?!#sYD@(ZuCf}dKO`(%V>WA|X0jG*18-_ocS`)3zII#g$yH|?#O~ZTcmp54 zc}7$Vb^6~R4F6YD3bjyNrO;5Sg^p4y43v7IkLV2y&&@&^8As|pg>Nx60p8C{wAPq@ zpz&|y^^uimC>yg{VKNQ0%ej4|Bz4f5qt>Wh)R^&zX6aOv1=-s?&{$2<7g?QI(NerT ztE);;heY~lC11_g3Y|5WeWV|$toev!TW1Q&4Own71=?!e5{5rvp3wdAL|A)+DE0Cy z--QjBEPpT#MEx119wgpg6fmx0-QZZ*11vLFS}wt=XdMRO{xEl0;Cm1D|7`W2Kl;z< z!|%h4wTScc)M<+^4|o-=dCm3n`@;T+d+atK8A8>v@_oLtLqOlo`{`kVtzg;C??2e%A#-ou{K3v2@pMFvSDEZi zwW8C=@_dvPmg^!^=B|4M!)M<_QB%vGrLJr1YDep+#*rD~veRO?BtW~xKWMc!L6>zFGE_sNQJc1SE%+Dimt|cRR0bYv|D9CiWO@}KS~pT z*${|IS^oe{8EK)WudCBHUe@XV+v_>3As0BsnupM8p&rA3N9u_>Qw#01uH_n#2WT?_ z<{arj7iDCQYVRwb)(BI8AL@1G?k1@9zxsC|gfg@YaMf}S97dlUfPqb^O*Y+$ZFtsq^n9|T~V^%9IC29zdywv zeCLkO7tIs(_`(T&!#}9!x`J7?zgGFTB*#UR<_Y@cmqxFQoa0yVZ5_XMLjCp8H%@Sd zkbK{q^U1Go$q~sWUZ)uP2>pvD$tZhmttv%Be%Z8I%=JC2!7kTOtQM_-d0Sz|yoQLu z$)_nxh`hZZqFe+xd|&1c`)HAmSdF}Rk9l*TEPd@UMUKaW59p+`FqWa@4nCUX4teS( zNS?cI2o2yOov#Jn5<$1A_9hC^8hhM}-F^7=IOBwC(p!Z=xhnr{`fgJ}QHFP^_%;>f z$Gn3=nAy-nXfxu#ON5RTiZ2k^2$8sog)^B^9*i>zOF0mNH;HH2^ZiLY!HI!99eF(@ zc)c8x!VDrP@=nS)g*8cJ-TXybiRuTBkT*Hvb0bw7}M zR%-0~=+L*zZF28bBJO#{vExhuA*I2jsMm3@;$xqG6p{rOB&)WzYf9X52YDf_{%Ux z?Qj?6SS4r5wPOv=UzZMK6#A+%vykQ?^<(OWLcj{}{YQXn1Alwg1OyX2ZWoP$U~n5p zYb0NVa$GBHfCLgYa>5OzuuhC0nX}fcJ?qQ>j^HQ9t^)>;EM_0zo07(pa{i*esmzv& zrQ8DD*u?e0k#?S-qp_*1!e!4Xvt{&ejBI8S0Bp1$d`-8&FQ~4#1b*3jkYsFBU6o3Gd|s={Mf1Ki@baXfgRoVIH>%^e@n- zAilMKs1a;c>uFc(`DL}9U*(-6h5rirSb#1XCxEfLNK%m;tl#3lgDn6+nDRU&x9Rb9 zkRCt9Bavv@`XR#or_KlC)F1xvD|Kh8eY!oZcjejr+EnkoGng8VA6%Q74shxkJ|geh z)N)D#mfxKCt_I+MuT%ZXfWJo**ZLM^RfH`48rA%uU#B-l>bJ|L>M-87C0!mm`${zbRRhzKBWi7-hGL9oQ_4jMjSdgL_72D z1*d$Rj8>Q~5xc7r8iXk)07A=yq%eVOz!qETP`aTo550u#i3a8kX*IQ zcqGUWLI{oWbF7`a^rb_Eq8>!v71R@=zeT_777F-1y!e8)2;Vp0 z&7G&4@b62~uU`YjczXKRj3>tFANcUpFb#2Els+9i{hOX+{bN&n;rn=VNFxZjnkkv&WApQ$oWH3 zleD+#v`I&Z5TkwC1gNMp%A}D4o=8agsdiVpT_&YCZW%)lCy+JpcXgrl-WKYCQ189C zEB193`njLR<99y;Yer%Dk}`esWkpKU>Rpd;)QYlqd(ZlH7lf8(!iv3vfQ39Ot=p33 ziEm+yy?iK)QHl^N&wR~c$Z%ceW^`Vp6k8YpnL1f!9=Ie>87>VPL)vYe@qkuhmO(E; zy^MFde=wG?0NZ{(CFe1i{ssSxrYM689d&HOmLX!g@ilAg8fD&3cK~JjKomB@}et$S7U06hpN&d(*1fH(D2En@!YQK#q!diL{7Tp`W#ewQ0l<-`QOvfukybkR%mrF zp{q5hkBpQ51r-C?RNb86++E%YP>#h(D`K%cm&#PDM-obLMawU*oO_INx_b2`)+99H z^jE#B`v8K65?3#t_QOy;Xykr~I}n!!M#?Jnfg7=EY Date: Sat, 7 Nov 2026 00:12:58 +0800 Subject: [PATCH 3/3] Keep the published benchmark_compare outputs The local re-run (2026-09-09) is not part of this release and recorded a machine-specific path. Co-Authored-By: Claude Opus 5.5 --- benchmark_log.txt | 58 +++++++++++++++++----------------- benchmark_results.json | 71 ++++++++++++++++++------------------------ 2 files changed, 59 insertions(+), 70 deletions(-) diff --git a/benchmark_log.txt b/benchmark_log.txt index 3f61180..ae96fcc 100644 --- a/benchmark_log.txt +++ b/benchmark_log.txt @@ -1,30 +1,30 @@ -[22:49:16] Starting benchmark comparison... -[22:49:16] === Benchmark 1: MiniMind (pure RAM) === -[22:49:17] params=91.4M RAM delta=0.394GB -[22:49:17] step 1/5 loss=8.9385 -[22:49:17] step 2/5 loss=8.9239 -[22:49:17] step 3/5 loss=8.8833 -[22:49:18] step 4/5 loss=8.8745 -[22:49:18] step 5/5 loss=8.8777 -[22:49:18] TPS=167.4 RAM_train=2.156GB RAM_infer=1.907GB -[22:49:19] === Benchmark 2: Chronos (SSD+DRAM hybrid) === -[22:49:19] params=111.1M RAM delta=0.097GB -[22:49:19] step 1/5 loss=8.9743 -[22:49:19] step 2/5 loss=8.9233 -[22:49:20] step 3/5 loss=9.0412 -[22:49:20] step 4/5 loss=8.9433 -[22:49:20] step 5/5 loss=8.9354 -[22:49:21] TPS=100.4 RAM_train=0.848GB RAM_infer=0.972GB -[22:49:21] cache_stats={'device': 'cpu', 'vram_experts': 3, 'vram_capacity': 3, 'ram_experts': 4, 'ram_capacity_dynamic': 4, 'expert_size_kb': 9984, 'cache_page_bytes': 81788928, 'resident_expert_bytes': 245366784, 'ram_buffer_bytes': 327155712, 'unique_expert_bytes': 327155712, 'pinned_ram_used_gb': 0.0, 'available_ram_gb': 86.1, 'pinned_ram_fraction': 0.0, 'h2d_stream': 'default', 'storage_format': 'safetensors', 'cluster_aware': True, 'layer_aware': False, 'num_clusters': 4, 'ssd_page_reads': 32, 'expert_evictions': 0, 'prefetch_hits': 48, 'prefetch_misses': 0, 'total_requests': 59, 'prefetch_observations': 48, 'hit_rate': 1.0, 'prefetch_queue_drops': 0, 'prefetch_wait_time_s': 0.0013231200573500246} -[22:49:21] +[05:45:46] Starting benchmark comparison... +[05:45:46] === Benchmark 1: MiniMind (pure RAM) === +[05:45:46] params=91.4M RAM delta=0.399GB +[05:45:46] step 1/5 loss=8.9834 +[05:45:46] step 2/5 loss=8.9272 +[05:45:46] step 3/5 loss=8.9045 +[05:45:47] step 4/5 loss=8.8138 +[05:45:47] step 5/5 loss=8.9289 +[05:45:47] TPS=349.9 RAM_train=2.250GB RAM_infer=2.250GB +[05:45:47] === Benchmark 2: Chronos (SSD+DRAM hybrid) === +[05:45:47] params=111.1M RAM delta=0.016GB +[05:45:47] step 1/5 loss=8.9058 +[05:45:47] step 2/5 loss=8.8842 +[05:45:48] step 3/5 loss=8.9566 +[05:45:48] step 4/5 loss=9.0308 +[05:45:48] step 5/5 loss=8.8993 +[05:45:48] TPS=224.3 RAM_train=2.443GB RAM_infer=2.631GB +[05:45:48] cache_stats={'vram_experts': 4, 'vram_capacity': 52, 'ram_experts': 4, 'ram_capacity_dynamic': 208, 'expert_size_kb': 4992, 'pinned_ram_used_gb': 0.152, 'available_ram_gb': 57.0, 'pinned_ram_fraction': 0.0027, 'h2d_stream': 'default', 'storage_format': 'safetensors', 'cluster_aware': True, 'num_clusters': 1, 'prefetch_hits': 0, 'prefetch_misses': 0, 'total_requests': 58, 'hit_rate': 0.0} +[05:45:48] === COMPARISON SUMMARY === -[22:49:21] MiniMind (pure RAM): -[22:49:21] tokens/s : 167.4 -[22:49:21] RAM (train) : 2.156 GB -[22:49:21] RAM (infer) : 1.907 GB -[22:49:21] Chronos (SSD+DRAM hybrid, async): -[22:49:21] tokens/s : 100.4 -[22:49:21] RAM (train) : 0.848 GB -[22:49:21] RAM (infer) : 0.972 GB -[22:49:21] KV cache : MLA(latent)+SlidingWindow -[22:49:21] Results saved to /Users/Fona/Downloads/Project_Chronos/benchmark_results.json +[05:45:48] MiniMind (pure RAM): +[05:45:48] tokens/s : 349.9 +[05:45:48] RAM (train) : 2.25 GB +[05:45:48] RAM (infer) : 2.25 GB +[05:45:48] Chronos (SSD+DRAM hybrid): +[05:45:48] tokens/s : 224.3 +[05:45:48] RAM (train) : 2.443 GB +[05:45:48] RAM (infer) : 2.631 GB +[05:45:48] KV cache : MLA(latent)+SlidingWindow +[05:45:48] Results saved to /Users/Fona/Downloads/Hybrid_LLM/Project_Chronos/benchmark_results.json diff --git a/benchmark_results.json b/benchmark_results.json index 65f69f0..96b5a7c 100644 --- a/benchmark_results.json +++ b/benchmark_results.json @@ -2,60 +2,49 @@ "minimind": { "name": "MiniMind (pure RAM)", "params_m": 91.4, - "ram_model_gb": 0.394, - "ram_train_gb": 2.156, - "ram_infer_gb": 1.907, - "tokens_per_sec": 167.4, + "ram_model_gb": 0.399, + "ram_train_gb": 2.25, + "ram_infer_gb": 2.25, + "tokens_per_sec": 349.9, "train_losses": [ - 8.9385, - 8.9239, - 8.8833, - 8.8745, - 8.8777 + 8.9834, + 8.9272, + 8.9045, + 8.8138, + 8.9289 ] }, "chronos": { - "name": "Chronos (SSD+DRAM hybrid, async)", + "name": "Chronos (SSD+DRAM hybrid)", "params_m": 111.1, - "ram_model_gb": 0.097, - "ram_train_gb": 0.848, - "ram_infer_gb": 0.972, - "tokens_per_sec": 100.4, + "ram_model_gb": 0.016, + "ram_train_gb": 2.443, + "ram_infer_gb": 2.631, + "tokens_per_sec": 224.3, "train_losses": [ - 8.9743, - 8.9233, - 9.0412, - 8.9433, - 8.9354 + 8.9058, + 8.8842, + 8.9566, + 9.0308, + 8.8993 ], "cache_stats": { - "device": "cpu", - "vram_experts": 3, - "vram_capacity": 3, + "vram_experts": 4, + "vram_capacity": 52, "ram_experts": 4, - "ram_capacity_dynamic": 4, - "expert_size_kb": 9984, - "cache_page_bytes": 81788928, - "resident_expert_bytes": 245366784, - "ram_buffer_bytes": 327155712, - "unique_expert_bytes": 327155712, - "pinned_ram_used_gb": 0.0, - "available_ram_gb": 86.1, - "pinned_ram_fraction": 0.0, + "ram_capacity_dynamic": 208, + "expert_size_kb": 4992, + "pinned_ram_used_gb": 0.152, + "available_ram_gb": 57.0, + "pinned_ram_fraction": 0.0027, "h2d_stream": "default", "storage_format": "safetensors", "cluster_aware": true, - "layer_aware": false, - "num_clusters": 4, - "ssd_page_reads": 32, - "expert_evictions": 0, - "prefetch_hits": 48, + "num_clusters": 1, + "prefetch_hits": 0, "prefetch_misses": 0, - "total_requests": 59, - "prefetch_observations": 48, - "hit_rate": 1.0, - "prefetch_queue_drops": 0, - "prefetch_wait_time_s": 0.0013231200573500246 + "total_requests": 58, + "hit_rate": 0.0 }, "kv_cache_type": "MLA(latent)+SlidingWindow" }