diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e450ff2..5c6d2fd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,7 +4,7 @@ on: push: branches: [main, dev] pull_request: - branches: [main] + branches: [main, dev] jobs: test: diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d8c843..e8dce09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,39 @@ ## [Unreleased] +## [0.3.1a2] — 2026-11-07 (beta) + +Kernel/extension split: the `chronos` kernel trains general models and holds no translation-specific code. The en↔zh demo moved to `demos/translation` (`chronos-translate`), which imports the kernel and is not part of the `Project_Chronos` wheel. The extension is not published on PyPI; install it from a checkout with `pip install -e demos/translation` (see `demos/translation/README.md`). + +### Moved (to `demos/translation/chronos_translate/`) +- `chronos.data.corpus.*` → `chronos_translate.corpus.*` (WMT19/OPUS-100/FLORES+ download, cleaning, dedup, en–zh shards, doc-mix sets). +- `chronos.bundle.*` → `chronos_translate.bundle.*`; `chronos bundle …` → `chronos-translate bundle …`. A bundle's `bin/chronos bundle …` still works: the launcher sends `bundle` to `chronos-translate`. `bin/setup.sh` installs both wheels from `wheels/`; `bundle build --wheel` can be repeated. +- `chronos.mlx.training.doc_run` → `chronos_translate.doc_run` (`chronos-translate doc-run`). +- `inspect_checkpoint.translation_probe` → `chronos_translate.probe` (`chronos-translate inspect`). +- `chronos.model.factory.ENZH_DEMO_BASE` → `chronos_translate.model_presets.ENZH_50M_BASE`. +- `chronos/mlx/training/recipes/enzh_50m_72h.json` → `demos/translation/recipes/`, with paths relative to the file. + +### Added +- Monitor (`chronos monitor`): a "Lazy experts / 懒加载" tab shows every metric the lazy-expert probe records after each validation. It has four tables (efficiency with the full-memory reference row; paging and loads; hit rates; fidelity), per-variant history charts (hit rate, fidelity, speed, loads, peak RSS) and the full probe log with skip and failure reasons. New endpoint: `GET /api/lazy_probes?run=…&limit=…`. +- Web UI monitor page: the probe card shows the same four grouped tables and four history charts in all four UI languages. Previously it showed a subset of the columns and no load, timing or reference metrics. +- The run snapshot's `lazy_probe` entries carry every measured probe field plus `reference_measured` (the full-memory reference per backend). +- Probes run from Web UI training jobs. For a CLI run (`chronos longrun`, `chronos-translate run`), start `chronos lazy-probe --watch ` next to it, or the tab stays empty. + +### Changed +- `LongRun` has task hooks: `task_eval(phase, n, final=)` (returns `{}` in the kernel), `eval_samples_for`, `generate`, `write_samples`, `pinned_modules`. The kernel `RunSpec` keeps `eval_samples` / `final_eval_samples` (default 0) and `export_prefix` (default `chronos`); `bleu_samples` and `probe_samples` moved to `TranslationRunSpec`. Run translation recipes with `chronos-translate run --spec …`; the kernel runner rejects unknown keys and names the task runner. +- `chronos longrun` and `load_spec` resolve relative recipe paths against the recipe's folder (they resolved against the working directory before). `--max-steps` and `--run-dir` give a short smoke run. +- Task-eval exceptions during SFT validation and the final evaluation are logged as `task_eval_error` instead of stopping the run (pretrain probes were already logged). The smoke path still raises. +- Tokenizer special tokens: the kernel default is `<|endoftext|>`, `<|im_start|>`, `<|im_end|>`, `<|pad|>` and reserved slots up to 64. Task tokens come from `special_tokens_with(extra)` or `chronos tokenizer --extra-special …`. `special_tokens_with(["<|en|>", "<|zh|>", "<|translate|>"])` reproduces the 0.3.1a1 order, so existing tokenizers and checkpoints are unchanged. +- The shard readers (`PackedTokenStream`, `SFTRecords`, `dtype_for`, `manifest_dtype`) live in `chronos.data.shards`. `median_record_tokens` and `cluster_locality_config` moved to `chronos.expert_clusters`. +- Preset "Chronos-v3 50M (en→zh recipe)" is now "Chronos-v3 50M (16k vocab)" (the old name is an alias) and uses the default 20 tokens/param schedule instead of the production-shard measurement. Same shape: 49,691,717 total / 12,090,437 active parameters (measured, meta device). + +### Fixed +- Web UI monitor page: a card with nothing to show no longer prints a stray "null". + +### Removed +- `chronos corpus`, `chronos bundle` and the `translate` extra (`sacrebleu`, `sacremoses` are dependencies of `chronos-translate`). +- Translation metric labels (`bleu_zh`, `flores_devtest_n`) from the Web UI i18n files; the monitor still shows any numeric metric a task writes. + ## [0.3.1a1] — 2026-11-07 (beta) Web UI half of 0.3.1: the rewritten Web UI, file picking and upload, LAN access, and the lazy-expert probe. diff --git a/MANIFEST.in b/MANIFEST.in index babb05d..da4788b 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -12,4 +12,4 @@ prune results prune reports prune .archive prune Translation?Training?Data -prune chronos/mlx/training/recipes +prune demos diff --git a/README.md b/README.md index 55dfd16..7acc334 100644 --- a/README.md +++ b/README.md @@ -107,14 +107,15 @@ At this size a full-depth step costs about as much as a batched verify pass. Wit | 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 longrun --spec recipe.json` (alias `run`) | Time-budgeted MLX run: warmup + cosine per phase wall-clock budget, NaN rollback, bit-exact resume, validation CE; tasks add their own metrics by subclassing `LongRun` (`task_eval`) | | `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).** +**Kernel and extensions.** The `chronos` package trains general models and carries no task-specific code. Task tooling lives in extensions that import the kernel. The en↔zh translation demo is one: `demos/translation` (`chronos-translate`: corpus pipeline, BLEU/chrF runs, model bundles; see its README). It is not part of the `Project_Chronos` wheel and is not published on PyPI; install it from a checkout with `pip install -e demos/translation`. + +**Status of the 50M en→zh demo (in progress; `demos/translation`).** - 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). @@ -450,7 +451,7 @@ pip install "Project_Chronos[mlx]" ```bash pip install "Project_Chronos[serve]" # chronos serve / chronos monitor (fastapi, uvicorn) -pip install "Project_Chronos[translate]" # corpus detokenization + BLEU/chrF probes (sacremoses, sacrebleu) +pip install -e demos/translation # en↔zh demo extension, from a checkout only (not on PyPI): chronos-translate ``` **vLLM serving (optional, Linux + CUDA only):** @@ -491,7 +492,7 @@ The Config tab ships five v3 presets. Each one is generated by the planner and s | 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 50M (16k vocab) | 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 | @@ -558,12 +559,13 @@ OLLAMA_HOST=127.0.0.1:11435 ollama run "Hello" - 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) +### Model bundle (translation demo extension) -`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. +Bundles belong to the translation extension (`demos/translation`). Build the two wheels first (`python -m build --outdir dist .` and `python -m build --outdir dist demos/translation`). `chronos-translate bundle build` packages one exported checkpoint of a 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 +chronos-translate bundle build --run-dir runs/enzh_50m_v1 --which best --pack \ + --wheel dist/project_chronos-0.3.1a2-py3-none-any.whl --wheel dist/chronos_translate-0.3.1a2-py3-none-any.whl ``` | Bundle path | Content | @@ -572,14 +574,14 @@ chronos bundle build --run-dir runs/enzh_50m_v1 --which best --wheel dist/projec | `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` | +| `wheels/`, `bin/setup.sh`, `bin/chronos` | offline install of the kernel and extension wheels into `.venv`; launcher that sets `CHRONOS_BUNDLE` / `CHRONOS_MINIMIND_DIR` and sends `bundle …` to `chronos-translate` | | `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/setup.sh # venv + wheels + `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 @@ -588,7 +590,7 @@ bin/chronos bundle finetune --data work/data --run-dir work/ft --tier quick --mo 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. +`--speculative` needs `--repetition-penalty 1.0` (its default), because the penalty changes the target distribution at every position. `finetune` writes a `chronos-translate run` 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 @@ -768,7 +770,7 @@ Project_Chronos/ │ │ └── benchmark.py # End-to-end benchmarking │ ├── data/ │ │ ├── flexible_dataset.py # Flexible JSONL dataset loader -│ │ └── corpus/ # download / clean / dedup / leak filter / shards +│ │ └── shards.py # memory-mapped pretrain / SFT shard readers │ ├── tokenizer/ # Chronos BPE builder (extract, normalize, train, evaluate) │ ├── monitor/ # Read-only run monitor: snapshot, API server, static page │ ├── webui/ # Web UI: FastAPI server, field schema, i18n/*.json, static/ SPA @@ -778,7 +780,6 @@ Project_Chronos/ │ │ ├── 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/ │ │ ├── ollama_server.py # chronos serve: Ollama/OpenAI protocol layer │ │ ├── runtime.py # Model discovery, lazy experts, speculative generation @@ -815,6 +816,8 @@ Project_Chronos/ │ └── fixtures/ ├── docs/ │ └── vllm_integration.md +├── demos/ +│ └── translation/ # en↔zh extension (chronos-translate): corpus, runs, bundles; not in the wheel ├── pyproject.toml └── README.md / README_zh.md / THIRD_PARTY_NOTICES.md ``` diff --git a/README_zh.md b/README_zh.md index 09ecbea..d150e35 100644 --- a/README_zh.md +++ b/README_zh.md @@ -119,14 +119,15 @@ flowchart LR | 命令 | 作用 | |---|---| | `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 longrun --spec recipe.json`（别名 `run`） | 按墙钟预算训练的 MLX 任务：每个阶段 warmup + cosine，NaN 回滚，续跑逐位一致，验证 CE；任务通过继承 `LongRun`（`task_eval`）加入自己的指标 | | `chronos serve` | 原生运行时上的 Ollama/OpenAI 协议服务（见下文） | | `chronos monitor` | 只读训练监控页面 + JSON API（见下文） | `chronos.model.factory.fit_config` 按总参数量和激活参数量目标搜索模型配置。 -**50M 英译中 demo（进行中）** +**内核与扩展。** `chronos` 包用于训练通用模型，不含任何任务特化代码。任务工具放在调用内核的扩展里。英↔中翻译 demo 就是这样一个扩展：`demos/translation`（`chronos-translate`：语料流水线、BLEU/chrF 训练、模型包，详见其 README）。它不在 `Project_Chronos` wheel 里，也不发布到 PyPI，需从仓库检出后用 `pip install -e demos/translation` 安装。 + +**50M 英译中 demo（进行中；`demos/translation`）** - 模型：总参数 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`（见“模型包”）。 @@ -452,7 +453,7 @@ 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） +pip install -e demos/translation # 英↔中 demo 扩展 chronos-translate，仅限仓库检出安装（不在 PyPI） ``` **vLLM 服务（可选，仅 Linux+CUDA）：** @@ -492,7 +493,7 @@ Config Tab 提供 5 档 v3 预设。它们都由规划器生成，存放在 `ui/ | 预设 | 总参数 / 激活参数（实测） | 形状 | 内存档位 | 训练峰值（估算） | |---|---|---|---|---| | 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 50M (16k vocab) | 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 | @@ -562,12 +563,13 @@ OLLAMA_HOST=127.0.0.1:11435 ollama run <模型名> "你好" - stock `ollama` / llama.cpp 无法执行 Chronos 的 GGUF 文件，因为 MLA+SWA 混合层和收缩式循环没有对应的 stock 实现；请改用本服务。 - 安全：默认只监听 127.0.0.1，不带认证。绑定其他地址必须同时给出 `--allow-remote` 和 `--api-key`，客户端随后发送 `Authorization: Bearer `。 -### 模型包（开箱即用的预训练模型） +### 模型包（翻译 demo 扩展） -`chronos bundle build` 把 `chronos longrun` 某次运行导出的一个 checkpoint 打包成解压即用的目录。打包过程只读取运行目录。 +模型包属于翻译扩展（`demos/translation`）。先构建两个 wheel（`python -m build --outdir dist .` 和 `python -m build --outdir dist demos/translation`）。`chronos-translate bundle build` 把某次运行导出的一个 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 +chronos-translate bundle build --run-dir runs/enzh_50m_v1 --which best --pack \ + --wheel dist/project_chronos-0.3.1a2-py3-none-any.whl --wheel dist/chronos_translate-0.3.1a2-py3-none-any.whl ``` | 包内路径 | 内容 | @@ -576,14 +578,14 @@ chronos bundle build --run-dir runs/enzh_50m_v1 --which best --wheel dist/projec | `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` | +| `wheels/`、`bin/setup.sh`、`bin/chronos` | 把内核和扩展两个 wheel 离线安装到 `.venv`；启动脚本设置 `CHRONOS_BUNDLE` / `CHRONOS_MINIMIND_DIR`，并把 `bundle …` 转给 `chronos-translate` | | `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/setup.sh # 建 venv、装 wheels、运行 `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 @@ -592,7 +594,7 @@ bin/chronos bundle finetune --data work/data --run-dir work/ft --tier quick --mo bin/chronos bundle serve -- --port 11435 ``` -`--speculative` 要求 `--repetition-penalty 1.0`（此时的默认值），因为惩罚会逐位置改变目标分布。`finetune` 写出一份 `init_checkpoint` 指向包内权重的 `chronos longrun` spec；运行目录里已有训练时会拒绝写入。用许可未知的数据（WMT19、OPUS-100）训练的模型包，`NOTICE.md` 中写明权重仅供研究使用。 +`--speculative` 要求 `--repetition-penalty 1.0`（此时的默认值），因为惩罚会逐位置改变目标分布。`finetune` 写出一份 `init_checkpoint` 指向包内权重的 `chronos-translate run` spec；运行目录里已有训练时会拒绝写入。用许可未知的数据（WMT19、OPUS-100）训练的模型包，`NOTICE.md` 中写明权重仅供研究使用。 ### 部署导出 @@ -777,7 +779,7 @@ Project_Chronos/ │ │ └── benchmark.py # 端到端基准测试 │ ├── data/ │ │ ├── flexible_dataset.py # 自动识别任意 JSONL 字段格式 -│ │ └── corpus/ # 下载 / 清洗 / 去重 / 泄漏过滤 / 分片 +│ │ └── shards.py # 内存映射的预训练 / SFT 分片读取器 │ ├── tokenizer/ # Chronos BPE 构建器（抽取、规范化、训练、评估） │ ├── monitor/ # 只读运行监控：快照、API 服务、静态页面 │ ├── webui/ # Web UI：FastAPI 服务、字段 schema、i18n/*.json、static/ 单页前端 @@ -787,7 +789,6 @@ Project_Chronos/ │ │ ├── dispatcher.py # 探针 + 优先级 + 训练能力声明 │ │ └── ext/opencl.py # 第三方 OpenCL 插件钩子（stub） │ ├── _backend_legacy.py # 向后兼容 build_model() 等旧 API -│ ├── bundle/ # chronos bundle：打包/校验/翻译/微调 │ ├── serving/ │ │ ├── ollama_server.py # chronos serve：Ollama/OpenAI 协议层 │ │ ├── runtime.py # 模型发现、懒加载专家、推测生成 @@ -824,6 +825,8 @@ Project_Chronos/ │ └── fixtures/ # tiny_pretrain / tiny_sft / tiny_dpo / tiny_grpo ├── docs/ │ └── vllm_integration.md +├── demos/ +│ └── translation/ # 英↔中扩展（chronos-translate）：语料、训练、模型包；不进 wheel ├── pyproject.toml └── README.md / README_zh.md / THIRD_PARTY_NOTICES.md ``` diff --git a/chronos/__init__.py b/chronos/__init__.py index 6b93853..2c5d388 100644 --- a/chronos/__init__.py +++ b/chronos/__init__.py @@ -1,5 +1,5 @@ # Project Chronos: On-device low-latency lookahead dual-layer MoE inference -__version__ = "0.3.1a1" +__version__ = "0.3.1a2" # Bootstrap minimind dependency before any submodule imports from chronos.deps import ensure_minimind as _ensure_minimind diff --git a/chronos/cli.py b/chronos/cli.py index 22f402a..ab81410 100644 --- a/chronos/cli.py +++ b/chronos/cli.py @@ -11,11 +11,13 @@ 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 chronos lazy-probe — lazy-expert efficiency / hit rate / fidelity of a checkpoint or a run chronos webui — Web UI (alias: ui) + +Task-specific tooling (corpora, evaluation, model bundles) lives in extensions +outside the kernel, e.g. ``demos/translation`` (``chronos-translate``). """ import sys import os @@ -87,15 +89,6 @@ def _module_main(module: str, rest): 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', @@ -112,10 +105,8 @@ def main(): 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) sub.add_parser('lazy-probe', help='Lazy-expert efficiency/hit-rate/fidelity probe (checkpoint or --watch run)', add_help=False) sub.add_parser('webui', help='Web UI (FastAPI; --legacy for the Gradio console)', aliases=['ui'], @@ -133,11 +124,9 @@ def main(): '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), 'lazy-probe': lambda rest: _module_main('chronos.workbench.lazy_probe', rest), 'webui': lambda rest: _module_main('chronos.webui.server', rest), 'ui': lambda rest: _module_main('chronos.webui.server', rest), diff --git a/chronos/data/corpus/evalsets.py b/chronos/data/corpus/evalsets.py deleted file mode 100644 index b9158fd..0000000 --- a/chronos/data/corpus/evalsets.py +++ /dev/null @@ -1,65 +0,0 @@ -"""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/shards.py b/chronos/data/shards.py new file mode 100644 index 0000000..0e586de --- /dev/null +++ b/chronos/data/shards.py @@ -0,0 +1,112 @@ +"""Readers for tokenized, memory-mapped training shards (any task, any language). + +Layout under a shard folder (written by any task's own builder): + +* ``pretrain/{train,valid}/*.bin``: one flat uint16/uint32 token array per + file, documents separated by the EOS token; the reader samples fixed windows. +* ``sft/{train,valid}/.{bin,idx,plen}``: chat-template records in a + token array (``.bin``), int64 record offsets (``.idx``, n+1 entries) and + int32 prompt lengths (``.plen``). Only tokens after the prompt are + supervised, so the loss covers the assistant reply alone. + +An optional ``shards_manifest.json`` may record ``"dtype": "uint16"|"uint32"``. +""" +from __future__ import annotations + +import json +import os + +import numpy as np + + +def dtype_for(vocab_size: int): + """Smallest unsigned dtype that holds every token id of ``vocab_size``.""" + return np.uint16 if vocab_size <= 65535 else np.uint32 + + +_dtype_for = dtype_for # pre-0.3.1a2 name + + +def manifest_dtype(shards_dir: str, default=np.uint16): + """Token dtype recorded in ``/shards_manifest.json`` (``default`` if absent).""" + path = os.path.join(shards_dir, "shards_manifest.json") + try: + with open(path, encoding="utf-8") as fh: + name = json.load(fh).get("dtype") + except (OSError, ValueError): + return default + return {"uint16": np.uint16, "uint32": np.uint32}.get(name, default) + + +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 diff --git a/chronos/expert_clusters.py b/chronos/expert_clusters.py index 916a82e..7bf6aa2 100644 --- a/chronos/expert_clusters.py +++ b/chronos/expert_clusters.py @@ -119,3 +119,64 @@ def recommended_window(median_record_tokens: Optional[float], seq_len: int) -> i return min(DEFAULT_WINDOW, seq_len) w = 1 << int(round(math.log2(float(median_record_tokens)))) return int(max(min(16, seq_len), min(w, seq_len))) + + +def median_record_tokens(shards_dir: str, max_files: int = 8) -> float | None: + """Median SFT record length (tokens) from the ``.idx`` offsets under ``/sft/train``. + + Reads at most ``max_files`` index files (offsets only, no token data). + Returns None when the shards have no SFT records (packed pretrain text + has no record boundaries).""" + import glob + import os + + import numpy as np + files = sorted(glob.glob(os.path.join(shards_dir, "sft", "train", "*.idx")))[:max_files] + lengths = [np.diff(np.fromfile(f, dtype=np.int64)) for f in files] + lengths = [x for x in lengths if x.size] + if not lengths: + return None + return float(np.median(np.concatenate(lengths))) + + +def cluster_locality_config(model_config: dict, shards_dir: str | None, *, enabled: bool = True, + lambda_cluster: float | None = None, window: int | None = None, + clusters: int | None = None, seq_len: int | None = None) -> dict: + """Cluster-locality fields for a fine-tune and where each value comes from. + + On by default at ``DEFAULT_LAMBDA`` (measured sweep on a synthetic + multi-topic task, from-scratch training; its effect on a fine-tune of + weights trained without the constraint has not been measured). Off for + dense models and when there are fewer than two clusters.""" + from chronos.model.config import ChronosConfig # lazy: this module stays torch/mlx-free + # The values training will use: longrun builds ChronosConfig(**model_config). + cfg = ChronosConfig(**{k: v for k, v in model_config.items() if v is not None}) + e = int(getattr(cfg, "num_experts", 0) or 0) + k = int(getattr(cfg, "num_experts_per_tok", 1) or 1) + seq = int(seq_len or getattr(cfg, "max_position_embeddings", 512) or 512) + usable = bool(getattr(cfg, "use_moe", True)) and e > 2 * k + lam = 0.0 if not (enabled and usable) else float(DEFAULT_LAMBDA if lambda_cluster is None else lambda_cluster) + if lam < 0: + raise ValueError("lambda_cluster must be >= 0") + median = median_record_tokens(shards_dir) if (shards_dir and window is None) else None + win = int(window) if window is not None else recommended_window(median, seq) + if win < 0: + raise ValueError("cluster_window must be >= 0 (0 = whole sequence)") + values = {"lambda_cluster_locality": lam, "cluster_locality_window": win, + "cluster_locality_clusters": int(clusters or 0)} + if not enabled: + why = "off (disabled by the caller)" + elif not usable: + why = "off: dense model or fewer than two clusters" + elif lambda_cluster is None: + why = f"default {DEFAULT_LAMBDA} (measured sweep, synthetic task; not measured for fine-tuning)" + else: + why = "user value" + if window is not None: + wsrc = "user value" + elif median is not None: + wsrc = f"rule: power of two nearest the median SFT record ({median:.0f} tokens), clamped to [16, {seq}]" + else: + wsrc = "default (no SFT record lengths in the data)" + return {"values": values, "source": {"lambda_cluster_locality": why, "cluster_locality_window": wsrc, + "cluster_locality_clusters": "user value" if clusters else "ceil(E / 2k)"}} diff --git a/chronos/lazy_experts.py b/chronos/lazy_experts.py index e16b150..365b17d 100644 --- a/chronos/lazy_experts.py +++ b/chronos/lazy_experts.py @@ -1,4 +1,4 @@ -"""SSD-lazy routed experts: one implementation for the Web UI, bundle and server. +"""SSD-lazy routed experts: one implementation for the Web UI, model bundles and the server. The page layout adapts to how the checkpoint was trained. Weights trained with the cluster-locality constraint (``lambda_cluster_locality > 0``) keep a diff --git a/chronos/mlx/training/inspect_checkpoint.py b/chronos/mlx/training/inspect_checkpoint.py index 9494a46..eb4183e 100644 --- a/chronos/mlx/training/inspect_checkpoint.py +++ b/chronos/mlx/training/inspect_checkpoint.py @@ -5,14 +5,16 @@ * 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; +* optional task probes: ``inspect_run(..., probes={"name": fn})`` stores + ``fn(spec, model, phase_name, samples)`` under ``report[name]`` (a task + extension supplies them; the kernel has none); 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 + python -m chronos.mlx.training.inspect_checkpoint --run-dir runs/my_run \ + --export serve_models/my_run_pretrain.pth """ from __future__ import annotations @@ -86,47 +88,6 @@ def greedy(model, tok, prompt_ids: list[int], stop_ids: set[int], limit: int) -> 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 @@ -140,7 +101,7 @@ def export(model, config, spec, out_path: str, stage: str) -> str: def inspect_run(run_dir: str, *, samples: int = 24, n_batches: int = 8, export_path: str | None = None, - device: str = "cpu") -> dict: + device: str = "cpu", probes: dict | None = None) -> dict: t0 = time.monotonic() spec, config, model, state = load_run_model(run_dir, device=device) phase = _phase(spec, state) @@ -150,8 +111,9 @@ def inspect_run(run_dir: str, *, samples: int = 24, n_batches: int = 8, export_p "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) + for name, probe in (probes or {}).items(): + if samples: + report[name] = 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) @@ -165,7 +127,7 @@ def inspect_run(run_dir: str, *, samples: int = 24, n_batches: int = 8, export_p 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("--samples", type=int, default=24, help="samples per task probe (0 = skip; no probes here)") 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")) diff --git a/chronos/mlx/training/longrun.py b/chronos/mlx/training/longrun.py index d4bc316..5405637 100644 --- a/chronos/mlx/training/longrun.py +++ b/chronos/mlx/training/longrun.py @@ -11,7 +11,13 @@ * 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. + validation CE. The UI monitor tails this file. + +The runner is task-agnostic. A task adds its own quality metrics by +subclassing ``LongRun`` and overriding ``task_eval`` (called at every +validation and once at the end); whatever keys it returns are logged next to +``valid_ce`` and discovered by the monitor. ``demos/translation`` is one +such extension. """ from __future__ import annotations @@ -55,8 +61,10 @@ class RunSpec: 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) + # Task-metric sample counts handed to ``LongRun.task_eval`` (0 = skip): + # at each validation and once at the end of the run. + eval_samples: int = 0 + final_eval_samples: int = 0 seed: int = 1234 dtype: str = "bfloat16" # "time": warmup + cosine over the phase's wall-clock budget, so the LR @@ -66,15 +74,12 @@ class RunSpec: # 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" + export_prefix: str = "chronos" extra: dict = field(default_factory=dict) @@ -174,7 +179,7 @@ def recurrent_metrics(state: dict | None) -> dict: class PretrainSource: def __init__(self, shards_dir, seq_len, seed): - from chronos.data.corpus.shards import PackedTokenStream + from chronos.data.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) @@ -200,7 +205,7 @@ def set_rng_state(self, state): class SFTSource: def __init__(self, shards_dir, seed): - from chronos.data.corpus.shards import SFTRecords + from chronos.data.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) @@ -573,10 +578,7 @@ 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))) + rec.update(self._safe_task_eval(phase.name, final=False)) best = self.state["best_valid"].get(phase.name) if best is None or ce < best: self.state["best_valid"][phase.name] = ce @@ -584,93 +586,74 @@ def _periodic_eval(self, source, phase): 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 + # ---- task hooks (overridden by task extensions) -------------------------- + def task_eval(self, phase_name: str, n: int, *, final: bool = False) -> dict: + """Task metrics for the current weights; the kernel itself has none. + + ``n`` comes from ``eval_samples_for`` (never 0: a zero count skips the call). + Return flat ``{metric_key: number}``; the keys land in the + ``validation`` / ``final`` rows of ``metrics.jsonl``. Use + ``write_samples`` for example outputs and ``generate`` for decoding. + """ + return {} + + def eval_samples_for(self, phase_name: str, *, final: bool) -> int: + """Sample count for ``task_eval`` (per phase in task extensions).""" + return int((self.spec.final_eval_samples if final else self.spec.eval_samples) or 0) + + def _safe_task_eval(self, phase_name: str, *, final: bool, n: int | None = None) -> dict: + n = self.eval_samples_for(phase_name, final=final) if n is None else int(n) + if n <= 0: + return {} 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 {} + return dict(self.task_eval(phase_name, n, final=final) or {}) + except Exception as exc: # noqa: BLE001 - a task metric must not stop training + self.log(f"[eval] task_eval skipped: {exc!r}") + return {"task_eval_error": f"{type(exc).__name__}: {exc}"[:200]} 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() + def generate(self, prompt_ids: list[int], limit: int, stop_ids=()) -> list[int]: + """Greedy continuation with the KV cache (no sampling, no penalty).""" + stop = set(int(i) for i in stop_ids) + logits, _, cache = self.model(mx.array([prompt_ids], dtype=mx.int32)) 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() + for _ in range(int(limit)): + nxt = int(mx.argmax(logits[0, -1]).item()) + if nxt in stop: + break + out.append(nxt) + logits, _, cache = self.model(mx.array([[nxt]], dtype=mx.int32), cache=cache) 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]) + def write_samples(self, rows: list[dict]) -> None: + """Append example outputs to ``samples.jsonl`` (any keys; the monitor shows them).""" + if not rows: + return + phase = self.spec.phases[min(self.state["phase_index"], len(self.spec.phases) - 1)].name + stamp = time.strftime("%Y-%m-%dT%H:%M:%S") 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)} + for r in rows: + fh.write(json.dumps({"time": stamp, "phase": phase, "step": self.state["phase_step"], **r}, + ensure_ascii=False) + "\n") + + def pinned_modules(self) -> tuple[str, ...]: + """Modules imported at start so source edits during a long run cannot + change the code that evaluates it; extensions add theirs.""" + return ("tokenizers", "chronos.data.shards", "chronos.mlx.training.inspect_checkpoint") # ---- 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 + import importlib + for name in self.pinned_modules(): + importlib.import_module(name) try: return self._run() except KeyboardInterrupt: @@ -703,26 +686,83 @@ def _run(self): 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)) + phase = self.spec.phases[-1].name if self.spec.phases else "" + final.update(self._safe_task_eval(phase, final=True)) 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: +PATH_KEYS = ("run_dir", "shards_dir", "tokenizer_dir", "init_checkpoint") + + +def load_spec(path: str, budget_hours: float | None = None, *, spec_cls=None, extra_path_keys=()) -> RunSpec: + """Read a RunSpec JSON. + + Relative paths (``PATH_KEYS`` and ``extra.``) resolve + against the recipe's folder, so a recipe can travel with its data. Keys + starting with ``_`` are comments. Any other unknown key is an error: it + usually belongs to a task extension's own spec class (``spec_cls``). + """ + spec_cls = spec_cls or RunSpec + base = os.path.dirname(os.path.abspath(path)) + with open(path, encoding="utf-8") as fh: + raw = {k: v for k, v in json.load(fh).items() if not k.startswith("_")} + + def resolve(v): + return v if not v or os.path.isabs(v) else os.path.normpath(os.path.join(base, os.path.expanduser(v))) + + for key in PATH_KEYS: + if raw.get(key): + raw[key] = resolve(raw[key]) + extra = dict(raw.get("extra") or {}) + for key in extra_path_keys: + if extra.get(key): + extra[key] = resolve(extra[key]) + raw["extra"] = extra + unknown = sorted(set(raw) - set(spec_cls.__dataclass_fields__)) + if unknown: + raise ValueError(f"{path}: unknown RunSpec keys {unknown}; task-specific keys need the task's " + f"runner (e.g. `chronos-translate run` for demos/translation recipes)") + raw["phases"] = [PhaseSpec(**p) for p in raw["phases"]] + spec = spec_cls(**raw) + if budget_hours is not None: + spec.budget_hours = float(budget_hours) + return spec + + +def main(argv=None, *, runner_cls=None, spec_cls=None, extra_path_keys=(), description=None, + prepare=None) -> int: + """CLI for a time-budgeted run; task extensions pass their own runner/spec classes + (``prepare(spec, args)`` may adjust the spec after loading).""" 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)") + runner_cls = runner_cls or LongRun + parser = argparse.ArgumentParser(description=description or "Time-budgeted Chronos MLX training run") + parser.add_argument("--spec", required=True, help="RunSpec JSON (paths may be relative to it)") parser.add_argument("--budget-hours", type=float, default=None) + parser.add_argument("--max-steps", type=int, default=None, + help="smoke test: this many updates per phase, then the closing eval (no full budget)") + parser.add_argument("--run-dir", default=None, help="override the recipe's run_dir (e.g. a smoke-test folder)") 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() + spec = load_spec(args.spec, args.budget_hours, spec_cls=spec_cls, extra_path_keys=extra_path_keys) + if args.run_dir: + spec.run_dir = os.path.abspath(args.run_dir) + if prepare is not None: + prepare(spec, args) + run = runner_cls(spec, log=lambda m: print(m, flush=True)) + if args.max_steps is None: + run.run() + return 0 + deadline = time.monotonic() + spec.budget_hours * 3600 + for pi in range(len(spec.phases)): + run.state.update(phase_index=pi, phase_step=0, phase_total=None, phase_budget_s=None, phase_wall_s=0.0) + run.run_phase(pi, deadline, max_steps=args.max_steps) + last = spec.phases[-1].name + n = min(run.eval_samples_for(last, final=True), 4) + if n > 0: # smoke: errors surface instead of being logged + run.model.eval() + print(json.dumps(run.task_eval(last, n, final=True), ensure_ascii=False), flush=True) return 0 diff --git a/chronos/model/config.py b/chronos/model/config.py index bad1ccb..b964770 100644 --- a/chronos/model/config.py +++ b/chronos/model/config.py @@ -203,7 +203,7 @@ def __init__(self, **kwargs): # λ > 0). The class default stays 0 so a stored recipe or checkpoint # that predates the field keeps its training objective; every entry # point that creates a new training config (planner, presets, CLI, - # bundle finetune) turns it on at the measured + # task extensions) turns it on at the measured # ``chronos.expert_clusters.DEFAULT_LAMBDA``. self.lambda_cluster_locality: float = kwargs.pop("lambda_cluster_locality", 0.0) self.cluster_locality_clusters: int = kwargs.pop("cluster_locality_clusters", 0) diff --git a/chronos/model/factory.py b/chronos/model/factory.py index c731987..59d976b 100644 --- a/chronos/model/factory.py +++ b/chronos/model/factory.py @@ -78,17 +78,3 @@ def fit_config(target_total: float, target_active: float, *, vocab_size: int, 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/monitor/evals.py b/chronos/monitor/evals.py index 9f6ca18..8e6c705 100644 --- a/chronos/monitor/evals.py +++ b/chronos/monitor/evals.py @@ -2,12 +2,12 @@ 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. +``dev_bleu``, ``accuracy``, ``pass_at_1`` or ``reward_mean``. This +module discovers them; the pages never hard-code a task. 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 +by that metric, so ``probe_math_accuracy`` and ``probe_code_accuracy`` share +one chart with series ``probe math`` / ``probe code``. Everything else gets a chart of its own. Bookkeeping keys (sample counts, durations, token counts) are listed as context, never charted. """ @@ -44,8 +44,8 @@ def is_context_key(key: str) -> bool: 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 + ``probe_math_accuracy`` → (``accuracy``, ``probe math``); ``bleu_dev`` → + (``bleu``, ``bleu_dev``); ``accuracy`` → (``accuracy``, ``accuracy``); unknown keys form a family of their own. """ head, sep, last = key.rpartition("_") diff --git a/chronos/monitor/server.py b/chronos/monitor/server.py index 1200fc8..f3e7ce1 100644 --- a/chronos/monitor/server.py +++ b/chronos/monitor/server.py @@ -203,6 +203,12 @@ 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/lazy_probes") + def lazy_probes(request: Request, run: str | None = None, limit: int = 200): + """Every lazy-expert probe of a run (``lazy_probe.jsonl``), newest last; the snapshot keeps 20.""" + _guard(request) + return {"probes": _run_or_404(mon.lazy_probes, run, max(1, min(int(limit), 1000)))} + @app.get("/api/log") def log(request: Request, run: str | None = None, bytes: int = 16000): # noqa: A002 - query name _guard(request) diff --git a/chronos/monitor/snapshot.py b/chronos/monitor/snapshot.py index 3a3a2e7..3c22b13 100644 --- a/chronos/monitor/snapshot.py +++ b/chronos/monitor/snapshot.py @@ -306,7 +306,7 @@ def log_tail(self, run: str | None, max_bytes: int = 16000) -> str: @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`` …).""" + (``held_out.valid_ce``, ``probe.dev.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:]: @@ -318,11 +318,16 @@ def _inspections(path: str) -> list[dict]: return out # Columns the monitor table shows per probe variant (the full result stays in lazy_probe.jsonl). - PROBE_MEASURED = ("tokens_per_s", "decode_tokens_per_s", "resident_hit_rate", "prediction_hit_rate", - "cold_miss_predict_hit_rate", "on_demand_loads", "prefetch_wait_time_s", "peak_rss_gb", - "spec_acceptance_rate", "speed_vs_full", "paging_resolved", "num_pages", - "effective_expert_budget", "tokens") - PROBE_FIDELITY = ("exact_match_rate", "positionwise_agreement", "min_matching_prefix") + PROBE_MEASURED = ("tokens_per_s", "decode_tokens_per_s", "speed_vs_full", "tokens", + "prefill_time_s", "decode_time_s", + "resident_hit_rate", "cache_hits", "cache_misses", + "prediction_hit_rate", "prediction_hits", "prediction_total", + "cold_miss_predict_hit_rate", "cold_prediction_hits", "cold_prediction_total", + "on_demand_loads", "on_demand_load_time_s", "sync_ssd_loads", "cluster_loads", + "prefetch_wait_time_s", "peak_rss_gb", "spec_acceptance_rate", + "paging_resolved", "num_pages", "num_clusters", "storage_format", "eviction_policy", + "effective_expert_budget", "ideal_active_experts") + PROBE_FIDELITY = ("exact_match_rate", "positionwise_agreement", "min_matching_prefix", "reference_tokens") def _lazy_probes(self, path: str, limit: int = 20) -> list[dict]: """Mid-training lazy-expert probe results (``lazy_probe.jsonl``), newest last. @@ -342,8 +347,12 @@ def _lazy_probes(self, path: str, limit: int = 20) -> list[dict]: entry.update({k: result.get(k) for k in ("trained_with_clusters", "auto_paging", "num_experts", "top_k", "truncated", "variants_requested", "elapsed_s", "conditions")}) - entry["references"] = {b: (r.get("measured") or {}).get("tokens_per_s") - for b, r in (result.get("references") or {}).items()} + refs = result.get("references") or {} + entry["references"] = {b: (r.get("measured") or {}).get("tokens_per_s") for b, r in refs.items()} + entry["reference_measured"] = { + b: ({"error": r["error"]} if "error" in r else + {k: (r.get("measured") or {}).get(k) for k in self.PROBE_MEASURED}) + for b, r in refs.items()} entry["variants"] = [{ "label": v.get("label"), "settings": v.get("settings"), **({"error": v["error"]} if "error" in v else {}), diff --git a/chronos/monitor/static/app.js b/chronos/monitor/static/app.js index 5024df6..2fbd3f2 100644 --- a/chronos/monitor/static/app.js +++ b/chronos/monitor/static/app.js @@ -25,7 +25,7 @@ const STR = { 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: "检查点", + t_evals: "评测与样例", t_lazy: "懒加载", 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", @@ -67,7 +67,7 @@ const STR = { 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", + t_evals: "Evals & samples", t_lazy: "Lazy experts", 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", @@ -248,6 +248,7 @@ async function refresh() { S.error = null; if (snap.run) S.run = snap.run; if (S.tab === "log" && S.logAuto) await loadLog(); + if (S.tab === "lazy") await loadLazy(); renderAll(false); } catch (e) { if (e instanceof AuthError) showLogin(""); @@ -379,7 +380,7 @@ function renderTiles() { } // ---- tabs ---------------------------------------------------------------------------- -const TABS = ["overview", "loss", "recurrent", "routing", "params", "evals", "config", "log", "ckpt"]; +const TABS = ["overview", "loss", "recurrent", "routing", "params", "evals", "lazy", "config", "log", "ckpt"]; function renderTabs() { const nav = clear($("tabs")); if (!TABS.includes(S.tab)) S.tab = "overview"; @@ -405,6 +406,7 @@ function selectTab(id) { renderTabs(); renderPanel(true); if (id === "log") loadLog().then(() => renderPanel(true)); + if (id === "lazy") loadLazy().then(() => renderPanel(true)); } // ---- tooltip ------------------------------------------------------------------------- const tip = { @@ -441,7 +443,7 @@ function bisect(xs, x) { // ---- line chart (small multiple; one y axis; crosshair tooltip) --------------------------- const pendingDraws = []; -function lineChart({ title, series, events = [], fmtY = fmt, note }) { +function lineChart({ title, series, events = [], fmtY = fmt, note, yRange }) { 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); @@ -456,10 +458,10 @@ function lineChart({ title, series, events = [], fmtY = fmt, note }) { 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)); + pendingDraws.push(() => drawLines(card, svg, live, events, fmtY, yRange)); return card; } -function drawLines(card, svg, live, events, fmtY) { +function drawLines(card, svg, live, events, fmtY, yRange) { 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); @@ -470,6 +472,8 @@ function drawLines(card, svg, live, events, fmtY) { 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; + if (yRange) { y0 = Math.max(y0, yRange[0]); y1 = Math.min(y1, yRange[1]); } // rates stay in [0, 1] + if (yRange && y0 >= y1) { y0 = yRange[0]; y1 = yRange[1]; } 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" }); @@ -957,6 +961,162 @@ function panelLog() { h("label", { class: "field" }, auto, tr("log_auto"))), pre]; } +// ---- lazy experts (lazy_probe.jsonl) --------------------------------------------------------------------- +Object.assign(STR.zh, { + ev_lazy_probe: "懒加载探针", + lz_none: "这个运行还没有懒加载探针结果（lazy_probe.jsonl）。Web UI 启动的训练默认在每次验证后探测；命令行训练可另开 `chronos lazy-probe --watch <运行目录>`。", + lz_latest: "最近一次探针", lz_meta: "{phase} · 第 {step} 步 · {t} · 耗时 {s}", lz_cond: "贪心解码，{n} 条提示，每条最多 {k} 个 token；全部为实测值。全内存参考 = 同一后端的全内存解码。", + lz_trained: "使用簇局部性训练", lz_auto: "自动分页", lz_experts: "专家 / top-k", lz_capped: "变体数被限制：请求 {n}，实测 {c}", + lz_skipped: "已跳过：{r}", lz_failed: "失败：{r}", yes: "是", no: "否", + lz_skip_dense: "稠密模型，没有可分页的路由专家", lz_skip_no_ckpt: "磁盘上还没有检查点", + lz_skip_busy: "上一次探测仍在运行", lz_skip_writing: "检查点仍在写入", + lz_skip_superseded: "已被第 {step} 步的新验证取代", + lz_eff: "效率", lz_paging: "分页与加载", lz_hit: "命中率", lz_fid: "保真度（对照全内存参考）", lz_ref: "全内存参考", + lz_variant: "组合", lz_backend: "后端", lz_tps: "token/秒", lz_dtps: "解码 token/秒", lz_vsfull: "相对全内存速度", lz_prefill: "预填充 (s)", lz_decode: "解码 (s)", + lz_rss: "峰值 RSS", lz_wait: "预取等待 (s)", lz_loads: "按需加载次数", lz_load_s: "加载耗时 (s)", lz_sync: "同步 SSD 读", lz_cluster_loads: "簇加载", + lz_pages: "分页 · 页数", lz_budget: "专家预算", lz_ideal: "理想激活专家", lz_evict: "淘汰策略", lz_accept: "草稿接受率", + lz_res_hit: "常驻命中率", lz_hits: "命中 / 未命中", lz_pred_hit: "预测命中率", lz_pred_n: "预测命中 / 总数", lz_cold: "冷未命中预测率", lz_cold_n: "冷未命中预测 / 总数", + lz_exact: "完全一致", lz_agree: "token 一致率", lz_prefix: "最短一致前缀", lz_ref_tokens: "参照 token", lz_error: "错误", + lz_pick: "组合", lz_c_hit: "命中率 · {p}", lz_c_fid: "保真度 · {p}", lz_c_speed: "解码速度（token/秒）· {p}", + lz_c_load: "加载次数 · {p}", lz_c_rss: "峰值 RSS (GB) · {p}", lz_history_note: "每次验证后运行一次；步数在每个阶段重新计数，所以只画当前阶段。", + lz_all: "全部探针记录", lz_status: "状态", lz_variants_n: "组合数", lz_ok: "完成", +}); +Object.assign(STR.en, { + ev_lazy_probe: "Lazy-expert probe", + lz_none: "No lazy-expert probe results for this run yet (lazy_probe.jsonl). Web UI training probes after each validation by default; for a command-line run, start `chronos lazy-probe --watch `.", + lz_latest: "Latest probe", lz_meta: "{phase} · step {step} · {t} · took {s}", lz_cond: "Greedy decoding, {n} prompts, up to {k} tokens each; all values measured. Reference = full-memory decode on the same backend.", + lz_trained: "Trained with cluster locality", lz_auto: "Auto paging layout", lz_experts: "Experts / top-k", lz_capped: "Variants capped: {n} requested, {c} measured", + lz_skipped: "Skipped: {r}", lz_failed: "Failed: {r}", yes: "yes", no: "no", + lz_skip_dense: "dense model, no routed experts to page", lz_skip_no_ckpt: "no checkpoint on disk yet", + lz_skip_busy: "previous probe still running", lz_skip_writing: "checkpoint still being written", + lz_skip_superseded: "superseded by the validation at step {step}", + lz_eff: "Efficiency", lz_paging: "Paging and loads", lz_hit: "Hit rates", lz_fid: "Fidelity (vs the full-memory reference)", lz_ref: "Full-memory reference", + lz_variant: "Variant", lz_backend: "Backend", lz_tps: "Tokens/s", lz_dtps: "Decode tokens/s", lz_vsfull: "Speed vs full memory", lz_prefill: "Prefill (s)", lz_decode: "Decode (s)", + lz_rss: "Peak RSS", lz_wait: "Prefetch wait (s)", lz_loads: "On-demand loads", lz_load_s: "Load time (s)", lz_sync: "Sync SSD reads", lz_cluster_loads: "Cluster loads", + lz_pages: "Paging · pages", lz_budget: "Expert budget", lz_ideal: "Ideal active experts", lz_evict: "Eviction", lz_accept: "Draft acceptance", + lz_res_hit: "Resident hit rate", lz_hits: "Hits / misses", lz_pred_hit: "Prediction hit rate", lz_pred_n: "Predicted hits / total", lz_cold: "Cold-miss predicted", lz_cold_n: "Cold hits / total", + lz_exact: "Exact match", lz_agree: "Token agreement", lz_prefix: "Shortest matching prefix", lz_ref_tokens: "Reference tokens", lz_error: "Error", + lz_pick: "Variant", lz_c_hit: "Hit rates · {p}", lz_c_fid: "Fidelity · {p}", lz_c_speed: "Decode speed (tokens/s) · {p}", + lz_c_load: "Loads · {p}", lz_c_rss: "Peak RSS (GB) · {p}", lz_history_note: "One probe after each validation; steps restart in every phase, so only the current phase is plotted.", + lz_all: "All probes", lz_status: "Status", lz_variants_n: "Variants", lz_ok: "done", +}); +// Results of chronos/workbench/lazy_probe.py after each validation: the checkpoint is +// decoded with experts paged from SSD and compared with a full-memory decode. +const fmtPct = (v) => (isNum(v) ? (v * 100).toFixed(1) + "%" : "—"); +const fmtSec = (v) => (isNum(v) ? (Math.abs(v) >= 10 ? v.toFixed(1) : v.toFixed(3)) : "—"); +const pair = (a, b) => (isNum(a) || isNum(b) ? `${fmtInt(a)} / ${fmtInt(b)}` : "—"); +const yesNo = (v) => (v === null || v === undefined ? "—" : tr(v ? "yes" : "no")); +async function loadLazy() { + try { const r = await api("api/lazy_probes", { run: S.run, limit: 500 }); S.lazy = r.probes || []; S.lazyFor = S.run; } + catch (e) { if (e instanceof AuthError) showLogin(""); else S.lazy = null; } +} +function lazyProbes() { return (S.lazy && S.lazyFor === S.run ? S.lazy : (S.snap && S.snap.lazy_probe)) || []; } +const lzm = (label, key, f) => ({ label, get: (v) => (v.measured || {})[key], num: true, fmt: f }); +const lzf = (label, key, f) => ({ label, get: (v) => (v.fidelity || {})[key], num: true, fmt: f }); +// Skip reasons come from chronos.workbench.lazy_probe / train_worker in English; unknown ones pass through. +const LZ_SKIPS = [[/^dense model/, "lz_skip_dense"], [/^no checkpoint/, "lz_skip_no_ckpt"], + [/^previous probe/, "lz_skip_busy"], [/^checkpoint still/, "lz_skip_writing"], + [/^superseded by step (\d+)/, "lz_skip_superseded"]]; +function lazySkipText(reason) { + for (const [re, key] of LZ_SKIPS) { + const m = re.exec(String(reason || "")); + if (m) return tr(key, { step: m[1] || "" }); + } + return String(reason || ""); +} +function lazyStatus(p) { + if (p.error) return tr("lz_failed", { r: p.error }); + if (p.skipped) return tr("lz_skipped", { r: lazySkipText(p.skipped) }); + return tr("lz_ok"); +} +function lazyLatest(latest) { + const c = latest.conditions || {}; + const refs = Object.entries(latest.reference_measured || {}).map(([b, m]) => ({ + label: `${tr("lz_ref")} · ${b}`, measured: m && !m.error ? m : {}, error: m && m.error })); + const variants = latest.variants || []; + const name = { label: tr("lz_variant"), get: (v) => v.label, keep: true }; + const err = { label: tr("lz_error"), get: (v) => v.error || "" }; + const box = (title, columns, rows) => h("div", { class: "box card" }, h("h3", { text: title }), table(columns, rows, { hideEmpty: true })); + const facts = [`${tr("lz_trained")}: ${yesNo(latest.trained_with_clusters)}`, `${tr("lz_auto")}: ${latest.auto_paging ?? "—"}`, + `${tr("lz_experts")}: ${latest.num_experts ?? "—"} / ${latest.top_k ?? "—"}`, + latest.truncated ? tr("lz_capped", { n: fmtInt(latest.variants_requested), c: fmtInt(variants.length) }) : null].filter(Boolean); + return [ + h("div", { class: "box card" }, h("h3", { text: tr("lz_latest") }), + h("p", { class: "caption", text: tr("lz_meta", { phase: phaseName(latest.phase), step: fmtInt(latest.step), + t: fmtTime(latest.time), s: fmtDur(latest.elapsed_s ?? latest.wall_s) }) }), + h("p", { class: "caption muted", text: facts.join(" · ") }), + h("p", { class: "caption muted", text: tr("lz_cond", { n: fmtInt(c.prompts), k: fmtInt(c.max_new_tokens) }) })), + box(tr("lz_eff"), [name, lzm(tr("lz_tps"), "tokens_per_s", (v) => fmt(v, 4)), lzm(tr("lz_dtps"), "decode_tokens_per_s", (v) => fmt(v, 4)), + lzm(tr("lz_vsfull"), "speed_vs_full", (v) => (isNum(v) ? `${v.toFixed(2)}×` : "—")), + lzm(tr("lz_prefill"), "prefill_time_s", fmtSec), lzm(tr("lz_decode"), "decode_time_s", fmtSec), + lzm(tr("lz_rss"), "peak_rss_gb", fmtGB), lzm(tr("lz_accept"), "spec_acceptance_rate", fmtPct), err], [...variants, ...refs]), + box(tr("lz_paging"), [name, lzm(tr("lz_loads"), "on_demand_loads", fmtInt), lzm(tr("lz_load_s"), "on_demand_load_time_s", fmtSec), + lzm(tr("lz_sync"), "sync_ssd_loads", fmtInt), lzm(tr("lz_cluster_loads"), "cluster_loads", fmtInt), + lzm(tr("lz_wait"), "prefetch_wait_time_s", fmtSec), + { label: tr("lz_pages"), get: (v) => ((v.measured || {}).paging_resolved ? `${v.measured.paging_resolved} · ${fmtInt(v.measured.num_pages)}` : null) }, + lzm(tr("lz_budget"), "effective_expert_budget", fmtInt), lzm(tr("lz_ideal"), "ideal_active_experts", fmtInt), + { label: tr("lz_evict"), get: (v) => (v.measured || {}).eviction_policy }], variants), + box(tr("lz_hit"), [name, lzm(tr("lz_res_hit"), "resident_hit_rate", fmtPct), + { label: tr("lz_hits"), num: true, get: (v) => (v.measured && (isNum(v.measured.cache_hits) || isNum(v.measured.cache_misses)) ? pair(v.measured.cache_hits, v.measured.cache_misses) : null) }, + lzm(tr("lz_pred_hit"), "prediction_hit_rate", fmtPct), + { label: tr("lz_pred_n"), num: true, get: (v) => (v.measured && v.measured.prediction_total ? pair(v.measured.prediction_hits, v.measured.prediction_total) : null) }, + lzm(tr("lz_cold"), "cold_miss_predict_hit_rate", fmtPct), + { label: tr("lz_cold_n"), num: true, get: (v) => (v.measured && v.measured.cold_prediction_total ? pair(v.measured.cold_prediction_hits, v.measured.cold_prediction_total) : null) }], + variants), + box(tr("lz_fid"), [name, lzf(tr("lz_exact"), "exact_match_rate", fmtPct), lzf(tr("lz_agree"), "positionwise_agreement", fmtPct), + lzf(tr("lz_prefix"), "min_matching_prefix", fmtInt), lzf(tr("lz_ref_tokens"), "reference_tokens", fmtInt), err], variants), + ]; +} +// One chosen variant across the probes of the newest phase (steps restart per phase). +function lazyHistory(done) { + const phase = done[done.length - 1].phase; + const runs = done.filter((p) => p.phase === phase); + const labels = [...new Set(runs.flatMap((p) => (p.variants || []).filter((v) => !v.error).map((v) => v.label)))]; + if (runs.length < 2 || !labels.length) return []; + const pick = labels.includes(S.lazyPick) ? S.lazyPick : (labels.find((l) => l.includes("sync_on_demand")) || labels[0]); + const rows = runs.map((p) => { + const v = (p.variants || []).find((x) => x.label === pick && !x.error); + return v ? { step: p.step, time: p.time, ...(v.measured || {}), ...(v.fidelity || {}) } : null; + }).filter(Boolean); + const sel = h("select", { "aria-label": tr("lz_pick"), onchange: (e) => { S.lazyPick = e.target.value; renderPanel(true); } }, + labels.map((l) => h("option", { value: l, text: l }))); + sel.value = pick; + const p = phaseName(phase); + const line = (title, series, fmtY, yRange) => + lineChart({ title, fmtY, yRange, series: series.map(([key, label]) => ({ key, label, rows })) }); + return [h("div", { class: "toolbar" }, h("label", { class: "field" }, tr("lz_pick"), sel)), + h("p", { class: "caption muted", text: tr("lz_history_note") }), + charts( + line(tr("lz_c_hit", { p }), [["resident_hit_rate", tr("lz_res_hit")], ["prediction_hit_rate", tr("lz_pred_hit")], + ["cold_miss_predict_hit_rate", tr("lz_cold")]], fmtPct, [0, 1]), + line(tr("lz_c_fid", { p }), [["exact_match_rate", tr("lz_exact")], ["positionwise_agreement", tr("lz_agree")]], fmtPct, [0, 1]), + line(tr("lz_c_speed", { p }), [["tokens_per_s", tr("lz_tps")], ["decode_tokens_per_s", tr("lz_dtps")]], (v) => fmt(v, 4)), + line(tr("lz_c_load", { p }), [["on_demand_loads", tr("lz_loads")], ["sync_ssd_loads", tr("lz_sync")], ["cluster_loads", tr("lz_cluster_loads")]], fmtCount), + line(tr("lz_c_rss", { p }), [["peak_rss_gb", tr("lz_rss")]], (v) => fmt(v, 3)))]; +} +function panelLazy() { + const probes = lazyProbes(); + if (!probes.length) return [h("div", { class: "empty", text: tr("lz_none") })]; + const done = probes.filter((p) => Array.isArray(p.variants)); + const last = probes[probes.length - 1]; + const out = []; + if (last !== done[done.length - 1]) { + out.push(h("div", { class: "notice warn", role: "status" }, h("span", { class: "icon", text: "!" }), + h("span", { text: `${phaseName(last.phase)} · ${tr("step")} ${fmtInt(last.step)}: ${lazyStatus(last)}` }))); + } + if (done.length) out.push(...lazyLatest(done[done.length - 1]), ...lazyHistory(done)); + out.push(h("div", { class: "box card" }, h("h3", { text: tr("lz_all") }), 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: tr("lz_status"), get: lazyStatus }, + { label: tr("lz_variants_n"), num: true, get: (p) => (Array.isArray(p.variants) ? p.variants.length : null), fmt: fmtInt }, + { label: tr("lz_res_hit"), num: true, fmt: fmtPct, + get: (p) => { const xs = (p.variants || []).map((v) => (v.measured || {}).resident_hit_rate).filter(isNum); return xs.length ? Math.min(...xs) : null; } }, + { label: tr("lz_exact"), num: true, fmt: fmtPct, + get: (p) => { const xs = (p.variants || []).map((v) => (v.fidelity || {}).exact_match_rate).filter(isNum); return xs.length ? Math.min(...xs) : null; } }, + ], probes.slice().reverse()))); + return out; +} + // ---- checkpoints ----------------------------------------------------------------------------------------- function panelCkpt() { const list = (S.snap && S.snap.checkpoints) || []; @@ -968,7 +1128,7 @@ function panelCkpt() { // ---- render orchestration ----------------------------------------------------------------------------------- const PANELS = { overview: panelOverview, loss: panelLoss, recurrent: panelRecurrent, routing: panelRouting, params: panelParams, - evals: panelEvals, config: panelConfig, log: panelLog, ckpt: panelCkpt }; + evals: panelEvals, lazy: panelLazy, 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("|"); @@ -1002,7 +1162,7 @@ function renderFoot() { 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 resetRunCaches() { S.params = S.tensors = S.tensorHist = S.log = S.lazy = 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", () => { diff --git a/chronos/serving/ollama_server.py b/chronos/serving/ollama_server.py index 47fdf69..e5eb920 100644 --- a/chronos/serving/ollama_server.py +++ b/chronos/serving/ollama_server.py @@ -3,8 +3,8 @@ 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" + chronos serve --models-dir runs/my_run --port 11435 + OLLAMA_HOST=127.0.0.1:11435 ollama run chronos_sft "Hello world" The kernel is never converted to another architecture: recurrent depth, SSD-lazy experts with lookahead prefetch, compressed KV and self-speculative diff --git a/chronos/tokenizer/__init__.py b/chronos/tokenizer/__init__.py index 06d2468..e88831d 100644 --- a/chronos/tokenizer/__init__.py +++ b/chronos/tokenizer/__init__.py @@ -13,10 +13,10 @@ from .evaluate import evaluate_tokenizer from .extract import ExtractConfig, extract_texts -from .model import DEFAULT_SPECIAL_TOKENS, ChronosTokenizer +from .model import DEFAULT_SPECIAL_TOKENS, ChronosTokenizer, special_tokens_with from .trainer import BPETrainerConfig, train_bpe -__all__ = ["build_tokenizer", "ChronosTokenizer", "BPETrainerConfig", "ExtractConfig"] +__all__ = ["build_tokenizer", "ChronosTokenizer", "BPETrainerConfig", "ExtractConfig", "special_tokens_with"] def build_tokenizer(paths: list[str], out_dir: str, *, vocab_size: int = 16384, @@ -61,11 +61,14 @@ def main(argv=None) -> int: 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) + parser.add_argument("--extra-special", nargs="*", default=(), metavar="TOKEN", + help="task control tokens placed at ids 3, 4, … (e.g. language tags)") 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), + special_tokens=special_tokens_with(args.extra_special), log=lambda m: print(m, flush=True), ) return 0 diff --git a/chronos/tokenizer/model.py b/chronos/tokenizer/model.py index 854d27b..09051da 100644 --- a/chronos/tokenizer/model.py +++ b/chronos/tokenizer/model.py @@ -28,10 +28,31 @@ "{% 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)) +#: Ids 0–2 are fixed (EOS / chat start / chat end); the rest of the 64 +#: control slots are ``<|reserved_i|>`` until a task claims them. +CORE_SPECIAL_TOKENS = ("<|endoftext|>", "<|im_start|>", "<|im_end|>", "<|pad|>") +SPECIAL_SLOTS = 64 + + +def special_tokens_with(extra=(), *, slots: int = SPECIAL_SLOTS) -> tuple[str, ...]: + """Core control tokens, then ``extra`` task tokens, padded with reserved slots. + + ``extra`` tokens go right after ``<|im_end|>`` (ids 3, 4, …) in the given + order; a token already in the core set keeps its core position. + """ + head = list(CORE_SPECIAL_TOKENS[:3]) + for tok in extra: + if tok not in head: + head.append(tok) + for tok in CORE_SPECIAL_TOKENS[3:]: + if tok not in head: + head.append(tok) + if len(head) > slots: + raise ValueError(f"{len(head)} special tokens exceed the {slots} control slots") + return tuple(head) + tuple(f"<|reserved_{i}|>" for i in range(slots - len(head))) + + +DEFAULT_SPECIAL_TOKENS = special_tokens_with() def tokenizer_normalize(text: str) -> str: diff --git a/chronos/tokenizer/normalize.py b/chronos/tokenizer/normalize.py index 828857a..efe9689 100644 --- a/chronos/tokenizer/normalize.py +++ b/chronos/tokenizer/normalize.py @@ -42,7 +42,7 @@ def normalize_text(text: str, *, form: str = "NFKC", unsegment_cjk: bool = False """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). + ship word-segmented Chinese). """ if form: text = unicodedata.normalize(form, text) diff --git a/chronos/webui/i18n/en.json b/chronos/webui/i18n/en.json index 8f239b5..6878759 100644 --- a/chronos/webui/i18n/en.json +++ b/chronos/webui/i18n/en.json @@ -555,11 +555,9 @@ "label.measured_meta": "measured (meta-device count)", "metric.active_mem_gb": "Active", "metric.balance_loss": "Expert balance", - "metric.bleu_zh": "BLEU (zh)", "metric.chrf": "chrF", "metric.error": "Error", "metric.eval_s": "Eval time (s)", - "metric.flores_devtest_n": "FLORES sentences", "metric.grad_norm": "Gradient norm", "metric.load_cv": "Load imbalance (CV)", "metric.locality_loss": "Cluster locality", @@ -670,10 +668,10 @@ "plan.tokens": "Training tokens", "preset.b1": "1B (64 GB)", "preset.b1.summary": "1B MoE with 64 fine-grained experts and aggressive memory saving.", - "preset.enzh50m": "50M en→zh recipe", - "preset.enzh50m.summary": "The 72-hour en→zh recipe: 16k Chronos tokenizer, shards_v1.", "preset.m200": "200M (32 GB)", "preset.m200.summary": "General 200M MoE with balanced memory saving and the 6400-token minimind vocabulary.", + "preset.m50": "50M (16k vocabulary)", + "preset.m50.summary": "50M MoE with 36 experts (top-2) for a 16k Chronos tokenizer; no memory saving.", "preset.m500": "500M (32 GB)", "preset.m500.summary": "500M MoE with aggressive memory saving: int8 KV, shared routing, activation checkpointing.", "preset.tiny": "Tiny (smoke test / CI)", @@ -681,16 +679,37 @@ "probe.accept": "Draft acceptance", "probe.agree": "Token agreement", "probe.auto_paging": "Auto paging", + "probe.budget": "Expert budget", + "probe.c_loads": "Loads", + "probe.c_rss": "Peak RSS (GB)", + "probe.c_speed": "Decode speed (tokens/s)", + "probe.cluster_loads": "Cluster loads", "probe.cold_hit": "Cold-miss predicted", + "probe.cold_n": "Cold-miss predicted / total", "probe.conditions": "Greedy decoding, {n} prompts, {t} tokens each; measured.", + "probe.decode": "Decode (s)", + "probe.dtps": "Decode tokens/s", + "probe.evict": "Eviction", "probe.exact": "Exact match", + "probe.group.eff": "Efficiency", + "probe.group.fid": "Fidelity (vs the full-memory reference)", + "probe.group.hit": "Hit rates", + "probe.group.paging": "Paging and loads", "probe.history": "Hit rate and agreement over training", + "probe.history_note": "One probe after each validation; steps restart in every phase, so only the newest phase is plotted.", "probe.hit": "Resident hit rate", + "probe.hits": "Hits / misses", + "probe.ideal": "Ideal active experts", + "probe.load_s": "Load time (s)", "probe.loads": "On-demand loads", "probe.none": "No probe yet. One runs after each validation.", "probe.pages": "Paging · pages", "probe.pred_hit": "Prediction hit rate", + "probe.pred_n": "Predicted hits / total", + "probe.prefill": "Prefill (s)", "probe.prefix": "Shortest matching prefix", + "probe.ref_row": "Full-memory reference · {backend}", + "probe.ref_tokens": "Reference tokens", "probe.reference": "Full-memory reference", "probe.rss": "Peak RSS (GB)", "probe.skip.busy": "the previous probe is still running", @@ -699,6 +718,7 @@ "probe.skip.superseded": "replaced by the newer validation at step {step}", "probe.skip.writing": "the checkpoint is still being written", "probe.skipped": "Skipped: {reason}", + "probe.sync": "Sync SSD reads", "probe.title": "Lazy-expert probe", "probe.tps": "Tokens/s", "probe.trained_clusters": "Trained with cluster locality", diff --git a/chronos/webui/i18n/ja.json b/chronos/webui/i18n/ja.json index f28330a..4bfd1cb 100644 --- a/chronos/webui/i18n/ja.json +++ b/chronos/webui/i18n/ja.json @@ -555,11 +555,9 @@ "label.measured_meta": "実測（meta デバイスで計数）", "metric.active_mem_gb": "使用中", "metric.balance_loss": "エキスパート均衡", - "metric.bleu_zh": "BLEU（中国語）", "metric.chrf": "chrF", "metric.error": "エラー", "metric.eval_s": "評価時間（秒）", - "metric.flores_devtest_n": "FLORES 文数", "metric.grad_norm": "勾配ノルム", "metric.load_cv": "負荷の偏り（変動係数）", "metric.locality_loss": "クラスタ局所性", @@ -670,10 +668,10 @@ "plan.tokens": "学習トークン数", "preset.b1": "1B（64 GB）", "preset.b1.summary": "1B MoE。64 個の細粒度エキスパート、積極的なメモリ節約。", - "preset.enzh50m": "50M 英→中レシピ", - "preset.enzh50m.summary": "72 時間の英→中レシピ：16k Chronos トークナイザー、shards_v1。", "preset.m200": "200M（32 GB）", "preset.m200.summary": "汎用 200M MoE。メモリ節約はバランス型で、minimind 6400 語彙を使用。", + "preset.m50": "50M（16k 語彙）", + "preset.m50.summary": "36 エキスパート（top-2）の 50M MoE、16k Chronos トークナイザー向け。メモリ節約なし。", "preset.m500": "500M（32 GB）", "preset.m500.summary": "500M MoE。積極的なメモリ節約：int8 KV、共有ルーティング、活性化チェックポイント。", "preset.tiny": "Tiny（スモークテスト / CI）", @@ -681,16 +679,37 @@ "probe.accept": "ドラフト受理率", "probe.agree": "トークン一致率", "probe.auto_paging": "自動ページング", + "probe.budget": "エキスパート予算", + "probe.c_loads": "読み込み回数", + "probe.c_rss": "ピーク RSS（GB）", + "probe.c_speed": "デコード速度（トークン/秒）", + "probe.cluster_loads": "クラスタ読み込み", "probe.cold_hit": "予測済みのコールドミス", + "probe.cold_n": "コールドミス予測 / 総数", "probe.conditions": "貪欲デコード、プロンプト {n} 件、各 {t} トークン、実測。", + "probe.decode": "デコード（秒）", + "probe.dtps": "デコード トークン/秒", + "probe.evict": "退避方式", "probe.exact": "完全一致", + "probe.group.eff": "効率", + "probe.group.fid": "忠実度（フルメモリ参照との比較）", + "probe.group.hit": "ヒット率", + "probe.group.paging": "ページングと読み込み", "probe.history": "学習の進行に伴うヒット率と一致率", + "probe.history_note": "検証ごとに 1 回実行します。ステップは段階ごとに数え直すため、最新の段階だけを描きます。", "probe.hit": "常駐ヒット率", + "probe.hits": "ヒット / ミス", + "probe.ideal": "理想アクティブエキスパート数", + "probe.load_s": "読み込み時間（秒）", "probe.loads": "オンデマンドロード回数", "probe.none": "まだプローブ結果がありません。検証のたびに 1 回実行されます。", "probe.pages": "ページング · ページ数", "probe.pred_hit": "予測ヒット率", + "probe.pred_n": "予測ヒット / 総数", + "probe.prefill": "プリフィル（秒）", "probe.prefix": "最短の一致プレフィックス", + "probe.ref_row": "フルメモリ参照 · {backend}", + "probe.ref_tokens": "参照トークン数", "probe.reference": "全メモリの参照", "probe.rss": "ピーク RSS（GB）", "probe.skip.busy": "前回のプローブがまだ実行中です", @@ -699,6 +718,7 @@ "probe.skip.superseded": "ステップ {step} の新しい検証に置き換えられました", "probe.skip.writing": "チェックポイントを書き込み中です", "probe.skipped": "スキップ：{reason}", + "probe.sync": "同期 SSD 読み込み", "probe.title": "遅延エキスパートのプローブ", "probe.tps": "トークン/秒", "probe.trained_clusters": "クラスタ局所性で学習済み", diff --git a/chronos/webui/i18n/zh-Hans.json b/chronos/webui/i18n/zh-Hans.json index cf3e9fe..97e2cbc 100644 --- a/chronos/webui/i18n/zh-Hans.json +++ b/chronos/webui/i18n/zh-Hans.json @@ -555,11 +555,9 @@ "label.measured_meta": "实测（meta 设备计数）", "metric.active_mem_gb": "活跃", "metric.balance_loss": "专家均衡", - "metric.bleu_zh": "BLEU（中文）", "metric.chrf": "chrF", "metric.error": "错误", "metric.eval_s": "评测用时（秒）", - "metric.flores_devtest_n": "FLORES 句数", "metric.grad_norm": "梯度范数", "metric.load_cv": "负载不均（变异系数）", "metric.locality_loss": "簇局部性", @@ -670,10 +668,10 @@ "plan.tokens": "训练 token 数", "preset.b1": "1B（64 GB）", "preset.b1.summary": "1B MoE，64 个细粒度专家，激进省内存。", - "preset.enzh50m": "50M 英译中配方", - "preset.enzh50m.summary": "72 小时英译中配方：16k Chronos 分词器，shards_v1。", "preset.m200": "200M（32 GB）", "preset.m200.summary": "通用 200M MoE，均衡省内存，使用 minimind 6400 词表。", + "preset.m50": "50M（16k 词表）", + "preset.m50.summary": "50M MoE，36 个专家（top-2），配 16k Chronos 分词器；不启用省内存选项。", "preset.m500": "500M（32 GB）", "preset.m500.summary": "500M MoE，激进省内存：int8 KV、共享路由、激活检查点。", "preset.tiny": "Tiny（冒烟测试 / CI）", @@ -681,16 +679,37 @@ "probe.accept": "草稿接受率", "probe.agree": "token 一致率", "probe.auto_paging": "自动分页", + "probe.budget": "专家预算", + "probe.c_loads": "加载次数", + "probe.c_rss": "峰值 RSS（GB）", + "probe.c_speed": "解码速度（token/秒）", + "probe.cluster_loads": "簇加载次数", "probe.cold_hit": "冷未命中预测率", + "probe.cold_n": "冷未命中预测 / 总数", "probe.conditions": "贪心解码，{n} 条提示，每条 {t} 个 token；实测。", + "probe.decode": "解码（秒）", + "probe.dtps": "解码 token/秒", + "probe.evict": "淘汰策略", "probe.exact": "完全一致", + "probe.group.eff": "效率", + "probe.group.fid": "保真度（对照全内存参考）", + "probe.group.hit": "命中率", + "probe.group.paging": "分页与加载", "probe.history": "训练过程中的命中率与一致率", + "probe.history_note": "每次验证后运行一次；步数在每个阶段重新计数，所以只画最新阶段。", "probe.hit": "常驻命中率", + "probe.hits": "命中 / 未命中", + "probe.ideal": "理想激活专家数", + "probe.load_s": "加载耗时（秒）", "probe.loads": "按需加载次数", "probe.none": "还没有探针结果。每次验证后会运行一次。", "probe.pages": "分页 · 页数", "probe.pred_hit": "预测命中率", + "probe.pred_n": "预测命中 / 总数", + "probe.prefill": "预填充（秒）", "probe.prefix": "最短一致前缀", + "probe.ref_row": "全内存参考 · {backend}", + "probe.ref_tokens": "参考 token 数", "probe.reference": "全内存参考", "probe.rss": "峰值 RSS（GB）", "probe.skip.busy": "上一次探测仍在运行", @@ -699,6 +718,7 @@ "probe.skip.superseded": "已被第 {step} 步的新验证取代", "probe.skip.writing": "检查点仍在写入", "probe.skipped": "已跳过：{reason}", + "probe.sync": "同步 SSD 读取", "probe.title": "懒加载探针", "probe.tps": "token/秒", "probe.trained_clusters": "使用簇局部性训练", diff --git a/chronos/webui/i18n/zh-Hant.json b/chronos/webui/i18n/zh-Hant.json index 7510d49..1a98dc7 100644 --- a/chronos/webui/i18n/zh-Hant.json +++ b/chronos/webui/i18n/zh-Hant.json @@ -555,11 +555,9 @@ "label.measured_meta": "實測（meta 裝置計數）", "metric.active_mem_gb": "活躍", "metric.balance_loss": "專家均衡", - "metric.bleu_zh": "BLEU（中文）", "metric.chrf": "chrF", "metric.error": "錯誤", "metric.eval_s": "評測用時（秒）", - "metric.flores_devtest_n": "FLORES 句數", "metric.grad_norm": "梯度範數", "metric.load_cv": "負載不均（變異係數）", "metric.locality_loss": "叢集局部性", @@ -670,10 +668,10 @@ "plan.tokens": "訓練 token 數", "preset.b1": "1B（64 GB）", "preset.b1.summary": "1B MoE，64 個細粒度專家，積極節省記憶體。", - "preset.enzh50m": "50M 英譯中配方", - "preset.enzh50m.summary": "72 小時英譯中配方：16k Chronos 分詞器，shards_v1。", "preset.m200": "200M（32 GB）", "preset.m200.summary": "通用 200M MoE，均衡節省記憶體，使用 minimind 6400 詞表。", + "preset.m50": "50M（16k 詞表）", + "preset.m50.summary": "50M MoE，36 個專家（top-2），配 16k Chronos 分詞器；不啟用省記憶體選項。", "preset.m500": "500M（32 GB）", "preset.m500.summary": "500M MoE，積極節省記憶體：int8 KV、共享路由、激活檢查點。", "preset.tiny": "Tiny（冒煙測試 / CI）", @@ -681,16 +679,37 @@ "probe.accept": "草稿接受率", "probe.agree": "逐 token 一致率", "probe.auto_paging": "自動分頁", + "probe.budget": "專家預算", + "probe.c_loads": "載入次數", + "probe.c_rss": "峰值 RSS（GB）", + "probe.c_speed": "解碼速度（token/秒）", + "probe.cluster_loads": "簇載入次數", "probe.cold_hit": "冷未命中被預測", + "probe.cold_n": "冷未命中預測 / 總數", "probe.conditions": "貪婪解碼，{n} 條提示，每條 {t} 個 token；實測。", + "probe.decode": "解碼（秒）", + "probe.dtps": "解碼 token/秒", + "probe.evict": "淘汰策略", "probe.exact": "完全一致", + "probe.group.eff": "效率", + "probe.group.fid": "保真度（對照全記憶體參考）", + "probe.group.hit": "命中率", + "probe.group.paging": "分頁與載入", "probe.history": "命中率與一致率隨訓練的變化", + "probe.history_note": "每次驗證後執行一次；步數在每個階段重新計數，因此只繪製最新階段。", "probe.hit": "常駐命中率", + "probe.hits": "命中 / 未命中", + "probe.ideal": "理想啟用專家數", + "probe.load_s": "載入耗時（秒）", "probe.loads": "按需載入次數", "probe.none": "尚無探針結果。每次驗證後會執行一次。", "probe.pages": "分頁 · 頁數", "probe.pred_hit": "預測命中率", + "probe.pred_n": "預測命中 / 總數", + "probe.prefill": "預填充（秒）", "probe.prefix": "最短一致前綴", + "probe.ref_row": "全記憶體參考 · {backend}", + "probe.ref_tokens": "參考 token 數", "probe.reference": "全記憶體參考", "probe.rss": "峰值 RSS（GB）", "probe.skip.busy": "上一次探測仍在執行", @@ -699,6 +718,7 @@ "probe.skip.superseded": "已被第 {step} 步的新驗證取代", "probe.skip.writing": "檢查點仍在寫入", "probe.skipped": "已略過：{reason}", + "probe.sync": "同步 SSD 讀取", "probe.title": "延遲載入專家探針", "probe.tps": "Token/秒", "probe.trained_clusters": "訓練時啟用叢集局部性", diff --git a/chronos/webui/schema.py b/chronos/webui/schema.py index 013bc1a..5f041be 100644 --- a/chronos/webui/schema.py +++ b/chronos/webui/schema.py @@ -204,7 +204,7 @@ def _f(key, group, type_, default=None, **kw) -> Field: _f("probe.paging", "probe", "multi", ["auto", "expert"], choices=["auto", "cluster", "expert"], recommended=["auto", "expert"], source="code", scope=PROBE), _f("probe.residency", "probe", "multi", ["planned"], choices=["planned", "window"], - recommended=["planned"], source="measured", ref=f"{RES}/bench_residency_50m_s13996_flores16_cpu.json", + recommended=["planned"], source="measured", ref=f"{RES}/bench_residency_50m_s13996_cpu.json", scope=PROBE), _f("probe.ram_ratios", "probe", "multi", [0.5, 1.0], choices=[0.1, 0.2, 0.25, 0.33, 0.5, 0.67, 0.75, 0.9, 1.0], scope=PROBE), @@ -240,7 +240,7 @@ def _f(key, group, type_, default=None, **kw) -> Field: _f("infer.paging", "lazy", "choice", "auto", choices=["auto", "cluster", "expert"], recommended="auto", source="code", scope=INFER), _f("infer.residency_planner", "lazy", "bool", True, recommended=True, source="measured", - ref=f"{RES}/bench_residency_50m_s13996_flores16_cpu.json", scope=INFER), + ref=f"{RES}/bench_residency_50m_s13996_cpu.json", scope=INFER), _f("infer.residency_decay", "lazy", "float", 0.7, min=0.0, max=0.95, step=0.05, recommended=0.7, source="code", scope=INFER, advanced=True), _f("infer.decode_wait_s", "lazy", "float", 0.012, min=0.0, max=1.0, step=0.001, recommended=0.012, diff --git a/chronos/webui/static/pages_config.js b/chronos/webui/static/pages_config.js index 225f1b9..f1f0e48 100644 --- a/chronos/webui/static/pages_config.js +++ b/chronos/webui/static/pages_config.js @@ -4,7 +4,7 @@ import { state, t, has, h, api, post, card, tile, pageHead, button, busy, params import { renderGroup, formState, markErrors, assign, submit } from "./forms.js"; import { pathControl } from "./picker.js"; -const PRESET_SLUGS = { "Chronos-v3 Tiny (smoke / CI)": "tiny", "Chronos-v3 50M (en→zh recipe)": "enzh50m", +const PRESET_SLUGS = { "Chronos-v3 Tiny (smoke / CI)": "tiny", "Chronos-v3 50M (16k vocab)": "m50", "Chronos-v3 200M (32 GB)": "m200", "Chronos-v3 500M (32 GB)": "m500", "Chronos-v3 1B (64 GB)": "b1" }; let presetCache = null; diff --git a/chronos/webui/static/pages_monitor.js b/chronos/webui/static/pages_monitor.js index fefe7fd..1fb90da 100644 --- a/chronos/webui/static/pages_monitor.js +++ b/chronos/webui/static/pages_monitor.js @@ -59,7 +59,8 @@ function paint(body, snap) { h("div", { class: "progress", role: "progressbar", "aria-valuemin": 0, "aria-valuemax": p.total || 0, "aria-valuenow": p.step || 0, "aria-label": phaseLabel(p.name) }, bar(p.total ? Math.min(1, p.step / p.total) : (p.state === "complete" ? 1 : 0))), h("span", { class: "note", text: `${int(p.step)} / ${p.total ? int(p.total) : "–"} · ${t(`phase.${p.state}`)}` }))); - body.replaceChildren( + // Cards with nothing to show return null; replaceChildren would print it as "null". + body.replaceChildren(...[ h("div", { class: "row between" }, h("div", { class: "row" }, badge(st.stale ? "stale" : st.state), st.stale ? h("span", { class: "note", text: t("mon.stale") }) : null, st.last_update_age_s !== null && st.last_update_age_s !== undefined ? h("span", { class: "note", text: t("mon.updated", { s: int(st.last_update_age_s) }) }) : null, @@ -82,7 +83,7 @@ function paint(body, snap) { eventsCard(snap), samplesCard(snap), checkpointsCard(snap), - logCard(snap)); + logCard(snap)].filter(Boolean)); } function bar(fraction) { diff --git a/chronos/webui/static/pages_probe.js b/chronos/webui/static/pages_probe.js index 16d94cd..6168e4e 100644 --- a/chronos/webui/static/pages_probe.js +++ b/chronos/webui/static/pages_probe.js @@ -19,25 +19,50 @@ function skipText(reason) { const phaseLabel = (name) => (has(`stage.${name}`) ? t(`stage.${name}`) : String(name ?? "–")); const yesNo = (v) => (v === null || v === undefined ? "–" : t(v ? "ui.yes" : "ui.no")); -function variantColumns() { - const m = (key, fmt) => (v) => (v.error ? "" : fmt((v.measured || {})[key])); - const f = (key, fmt) => (v) => (v.error ? "" : fmt((v.fidelity || {})[key])); +const secs = (x) => (x === null || x === undefined ? "–" : num(x, Math.abs(Number(x)) >= 10 ? 1 : 3)); +const pair = (a, b) => (a == null && b == null ? "–" : `${int(a)} / ${int(b)}`); +const mOf = (v) => v.measured || {}; +const fOf = (v) => v.fidelity || {}; + +// Columns that no row fills are dropped (e.g. draft acceptance without speculative variants). +function groupTable(title, columns, rows) { + const live = columns.filter((c) => c.keep || rows.some((r) => { const x = c.value(r); return x !== null && x !== undefined && x !== ""; })); + return h("div", { class: "mt" }, h("h4", { text: title }), + table(live.map((c) => ({ label: c.label, num: c.num, fmt: (r) => (r.error && !c.keep && !c.always ? "" : c.fmt(c.value(r), r)) })), rows)); +} +const col = (key, get, fmt) => ({ label: t(key), num: true, value: get, fmt }); +function variantTables(latest) { + const variants = latest.variants || []; + const refs = Object.entries(latest.reference_measured || {}).map(([b, m]) => ({ + label: t("probe.ref_row", { backend: b }), measured: m && !m.error ? m : {}, ...(m && m.error ? { error: m.error } : {}) })); + const name = { label: t("probe.variant"), keep: true, value: (v) => v.label, fmt: (x) => h("span", { class: "mono", text: x }) }; + const errCol = { label: t("ui.error"), always: true, fmt: (x) => (x ? h("span", { class: "error", text: x }) : "") }; return [ - { label: t("probe.variant"), fmt: (v) => h("span", { class: "mono", text: v.label }) }, - { label: t("probe.exact"), num: true, fmt: f("exact_match_rate", (x) => pct(x)) }, - { label: t("probe.agree"), num: true, fmt: f("positionwise_agreement", (x) => pct(x)) }, - { label: t("probe.prefix"), num: true, fmt: f("min_matching_prefix", (x) => int(x)) }, - { label: t("probe.hit"), num: true, fmt: m("resident_hit_rate", (x) => pct(x)) }, - { label: t("probe.pred_hit"), num: true, fmt: m("prediction_hit_rate", (x) => pct(x)) }, - { label: t("probe.cold_hit"), num: true, fmt: m("cold_miss_predict_hit_rate", (x) => pct(x)) }, - { label: t("probe.loads"), num: true, fmt: m("on_demand_loads", (x) => int(x)) }, - { label: t("probe.wait"), num: true, fmt: m("prefetch_wait_time_s", (x) => num(x, 3)) }, - { label: t("probe.tps"), num: true, fmt: m("tokens_per_s", (x) => num(x, 1)) }, - { label: t("probe.vs_full"), num: true, fmt: m("speed_vs_full", (x) => (x === null || x === undefined ? "–" : `${num(x, 2)}×`)) }, - { label: t("probe.accept"), num: true, fmt: m("spec_acceptance_rate", (x) => pct(x)) }, - { label: t("probe.rss"), num: true, fmt: m("peak_rss_gb", (x) => num(x, 2)) }, - { label: t("probe.pages"), num: true, fmt: (v) => (v.error ? "" : `${(v.measured || {}).paging_resolved ?? "–"} · ${int((v.measured || {}).num_pages)}`) }, - { label: t("ui.error"), fmt: (v) => (v.error ? h("span", { class: "error", text: v.error }) : "") }, + groupTable(t("probe.group.eff"), [name, + col("probe.tps", (v) => mOf(v).tokens_per_s, (x) => num(x, 1)), col("probe.dtps", (v) => mOf(v).decode_tokens_per_s, (x) => num(x, 1)), + col("probe.vs_full", (v) => mOf(v).speed_vs_full, (x) => (x === null || x === undefined ? "–" : `${num(x, 2)}×`)), + col("probe.prefill", (v) => mOf(v).prefill_time_s, secs), col("probe.decode", (v) => mOf(v).decode_time_s, secs), + col("probe.rss", (v) => mOf(v).peak_rss_gb, (x) => num(x, 2)), col("probe.accept", (v) => mOf(v).spec_acceptance_rate, (x) => pct(x)), + { ...errCol, value: (v) => v.error }], [...variants, ...refs]), + groupTable(t("probe.group.paging"), [name, + col("probe.loads", (v) => mOf(v).on_demand_loads, int), col("probe.load_s", (v) => mOf(v).on_demand_load_time_s, secs), + col("probe.sync", (v) => mOf(v).sync_ssd_loads, int), col("probe.cluster_loads", (v) => mOf(v).cluster_loads, int), + col("probe.wait", (v) => mOf(v).prefetch_wait_time_s, secs), + { label: t("probe.pages"), value: (v) => mOf(v).paging_resolved, fmt: (x, r) => `${x ?? "–"} · ${int(mOf(r).num_pages)}` }, + col("probe.budget", (v) => mOf(v).effective_expert_budget, int), col("probe.ideal", (v) => mOf(v).ideal_active_experts, int), + { label: t("probe.evict"), value: (v) => mOf(v).eviction_policy, fmt: (x) => x ?? "–" }], variants), + groupTable(t("probe.group.hit"), [name, + col("probe.hit", (v) => mOf(v).resident_hit_rate, (x) => pct(x)), + col("probe.hits", (v) => (mOf(v).cache_hits ?? mOf(v).cache_misses ?? null), (_x, r) => pair(mOf(r).cache_hits, mOf(r).cache_misses)), + col("probe.pred_hit", (v) => mOf(v).prediction_hit_rate, (x) => pct(x)), + col("probe.pred_n", (v) => (mOf(v).prediction_total || null), (_x, r) => pair(mOf(r).prediction_hits, mOf(r).prediction_total)), + col("probe.cold_hit", (v) => mOf(v).cold_miss_predict_hit_rate, (x) => pct(x)), + col("probe.cold_n", (v) => (mOf(v).cold_prediction_total || null), (_x, r) => pair(mOf(r).cold_prediction_hits, mOf(r).cold_prediction_total))], + variants), + groupTable(t("probe.group.fid"), [name, + col("probe.exact", (v) => fOf(v).exact_match_rate, (x) => pct(x)), col("probe.agree", (v) => fOf(v).positionwise_agreement, (x) => pct(x)), + col("probe.prefix", (v) => fOf(v).min_matching_prefix, int), col("probe.ref_tokens", (v) => fOf(v).reference_tokens, int), + { ...errCol, value: (v) => v.error }], variants), ]; } @@ -61,16 +86,16 @@ export function probeSection(snap) { h("span", { class: "note", text: `${t("probe.auto_paging")}: ${latest.auto_paging ?? "–"}` }), latest.truncated ? h("span", { class: "note", text: t("train.variants_capped", { n: int(latest.variants_requested), cap: int(latest.variants.length) }) }) : null), h("p", { class: "note", text: t("probe.conditions", { n: int(c.prompts), t: int(c.max_new_tokens) }) }), - table(variantColumns(), latest.variants), - Object.keys(latest.references || {}).length ? h("p", { class: "note", text: `${t("probe.reference")}: ${Object.entries(latest.references) + ...variantTables(latest), + !latest.reference_measured && Object.keys(latest.references || {}).length ? h("p", { class: "note", text: `${t("probe.reference")}: ${Object.entries(latest.references) .map(([b, tps]) => `${b} ${num(tps, 1)} ${t("probe.tps")}`).join(" · ")}` }) : null, - historyChart(done)); + historyCharts(done)); } return card(t("probe.title"), ...parts); } // History of one chosen variant over training (sync_on_demand first by default). -function historyChart(all) { +function historyCharts(all) { // Steps restart in every phase: plot the newest phase only. const phase = all[all.length - 1].phase; const done = all.filter((p) => p.phase === phase); @@ -79,18 +104,22 @@ function historyChart(all) { if (!labels.length) return null; const want = state.values.probeVariant; const pick = labels.includes(want) ? want : (labels.find((l) => l.includes("sync_on_demand")) || labels[0]); - const series = [ - ["probe.hit", (v) => v.measured.resident_hit_rate], - ["probe.pred_hit", (v) => v.measured.prediction_hit_rate], - ["probe.agree", (v) => v.fidelity.positionwise_agreement], - ["probe.exact", (v) => v.fidelity.exact_match_rate], - ].map(([key, get], i) => ({ label: t(key), color: `var(--s${i + 1})`, - points: done.map((p) => { const v = p.variants.find((x) => x.label === pick && !x.error); return v ? [p.step, get(v)] : null; }) + const at = (p) => p.variants.find((x) => x.label === pick && !x.error); + const series = (defs) => defs.map(([key, get], i) => ({ label: t(key), color: `var(--s${i + 1})`, + points: done.map((p) => { const v = at(p); return v ? [p.step, get(v)] : null; }) .filter((pt) => pt && pt[1] !== null && pt[1] !== undefined) })); + const chart = (title, defs, yFmt) => lineChart({ title: `${title} · ${phaseLabel(phase)}`, series: series(defs), xLabel: t("ui.step"), yFmt }); const sel = h("select", { id: "probe-variant" }, labels.map((l) => h("option", { value: l, text: l }))); sel.value = pick; sel.addEventListener("change", () => { state.values.probeVariant = sel.value; document.dispatchEvent(new CustomEvent("chronos:refresh")); }); return h("div", { class: "mt" }, h("div", { class: "row" }, h("label", { class: "compact", for: "probe-variant", text: t("probe.variant") }), sel), - lineChart({ title: `${t("probe.history")} · ${phaseLabel(phase)}`, series, xLabel: t("ui.step"), yFmt: (v) => pct(v, 0) })); + h("p", { class: "note", text: t("probe.history_note") }), + h("div", { class: "grid two mt" }, + chart(t("probe.history"), [["probe.hit", (v) => mOf(v).resident_hit_rate], ["probe.pred_hit", (v) => mOf(v).prediction_hit_rate], + ["probe.agree", (v) => fOf(v).positionwise_agreement], ["probe.exact", (v) => fOf(v).exact_match_rate]], (v) => pct(v, 0)), + chart(t("probe.c_speed"), [["probe.tps", (v) => mOf(v).tokens_per_s], ["probe.dtps", (v) => mOf(v).decode_tokens_per_s]], (v) => num(v, 1)), + chart(t("probe.c_loads"), [["probe.loads", (v) => mOf(v).on_demand_loads], ["probe.sync", (v) => mOf(v).sync_ssd_loads], + ["probe.cluster_loads", (v) => mOf(v).cluster_loads]], (v) => num(v, 0)), + chart(t("probe.c_rss"), [["probe.rss", (v) => mOf(v).peak_rss_gb]], (v) => num(v, 2)))); } diff --git a/chronos/workbench/lazy_probe.py b/chronos/workbench/lazy_probe.py index 8842b36..bd132e6 100644 --- a/chronos/workbench/lazy_probe.py +++ b/chronos/workbench/lazy_probe.py @@ -32,7 +32,7 @@ # results from different validation points are comparable. _SYNTHETIC_PROMPTS = ((5, 17, 33, 2, 61, 9, 40, 12), (7, 3, 29, 44, 18, 6, 51, 23, 11, 2)) DEFAULT_PROMPTS = ( - "Translate to Chinese: The weather is nice today.", + "Summarize in one sentence: the weather is nice today and the park is busy.", "Explain what a mixture-of-experts layer does in one sentence.", ) diff --git a/demos/translation/README.md b/demos/translation/README.md new file mode 100644 index 0000000..1fc9f56 --- /dev/null +++ b/demos/translation/README.md @@ -0,0 +1,99 @@ +# chronos-translate: an en↔zh translation demo on the Chronos kernel + +The Chronos kernel (`chronos/`) trains general models and carries no +translation-specific code. Everything that builds the English↔Chinese +benchmark model lives here, as an extension that imports the kernel. + +| Path | What it is | +|---|---| +| `chronos_translate/corpus/` | download, clean, dedup and shard en–zh data (WMT19, OPUS-100, FLORES+, …); `docmix` builds merged n-in-1 sets | +| `chronos_translate/run.py` | `TranslationLongRun`: the kernel `LongRun` plus FLORES+ BLEU/chrF validation | +| `chronos_translate/doc_run.py` | doc-mix variant: micro-batch auto-fit, `target_passes`, document-level eval | +| `chronos_translate/probe.py` | offline FLORES+ probe of a run checkpoint | +| `chronos_translate/bundle/` | self-contained model bundles: build, check, translate, prepare-data, fine-tune, serve | +| `chronos_translate/model_presets.py` | the 50M topology base and the tokenizer's language tokens | +| `recipes/enzh_50m_72h.json` | the 72h benchmark recipe (paths relative to this file) | +| `recipes/docmix_plan.json` | planner output for the doc-mix packages A/B/C | +| `scripts/` | `plan_docmix.py`, `make_docmix_packages.py` | +| `tests/` | extension tests (`pytest demos/translation/tests`) | + +## Install + +This extension is a demo inside the repository. It is not published on PyPI; +install it from a checkout. The kernel comes from PyPI (`pip install --pre` +while 0.3.1 is in beta) or from the same checkout. + +```sh +git clone https://github.com/FonaTech/Project_Chronos.git && cd Project_Chronos +pip install -e '.[mlx,serve]' # the kernel (or: pip install --pre 'Project_Chronos[mlx,serve]==0.3.1a2') +pip install -e demos/translation # this extension (sacrebleu, sacremoses, pyarrow pinned) +chronos-translate --version +``` + +The extension pins `Project_Chronos==0.3.1a2`. Use the checkout's tag +(`git checkout v0.3.1a2`) so both sides match. A model bundle needs the two +as wheels; build them from the checkout: + +```sh +python -m build --outdir dist . # dist/project_chronos-0.3.1a2-py3-none-any.whl +python -m build --outdir dist demos/translation # dist/chronos_translate-0.3.1a2-py3-none-any.whl +``` + +The extension's tests run from the repo root: `pytest demos/translation/tests`. +The kernel CI does not run them, because they need sacrebleu and pyarrow. + +## Commands + +```sh +chronos tokenizer data/pairs.jsonl --out tok --vocab-size 16384 \ + --extra-special "<|en|>" "<|zh|>" "<|translate|>" # kernel tokenizer + language tags +chronos-translate corpus download --help # also: build | shards | docmix +chronos-translate run --spec demos/translation/recipes/enzh_50m_72h.json +chronos-translate doc-run --spec /recipe.json +chronos-translate inspect --run-dir runs/enzh_50m_v1 --samples 24 +chronos-translate bundle build --run-dir runs/enzh_50m_v1 --wheel dist/project_chronos-*.whl \ + --wheel dist/chronos_translate-*.whl +``` + +`chronos serve`, `chronos monitor` and `chronos tokenizer` stay kernel commands. + +## How it plugs into the kernel + +This folder is also the template for other task extensions. It uses only +public kernel hooks and never patches kernel code: + +| Kernel hook | What this extension does with it | +|---|---| +| `RunSpec` subclass | `TranslationRunSpec` adds `bleu_samples`, `probe_samples` and its own `export_prefix` | +| `LongRun.eval_samples_for` | sample count per phase (pretrain probe, SFT BLEU subset, final devtest) | +| `LongRun.task_eval` | returns `{metric: value}`; the keys land in the `validation` / `final` rows of `metrics.jsonl` and show up in the monitor | +| `LongRun.generate`, `write_samples` | greedy decoding with the KV cache; example outputs go to `samples.jsonl` | +| `LongRun.pinned_modules` | records the versions of sacrebleu and the extension modules in the run | +| `longrun.main(runner_cls=, spec_cls=, extra_path_keys=)` | the `chronos-translate run` CLI, with `extra.raw_root` resolved against the recipe folder | +| `special_tokens_with(...)` | the tokenizer's `<|en|>`, `<|zh|>`, `<|translate|>` slots | + +A failing `task_eval` is logged as `task_eval_error` and never stops training. +A new task copies this layout: its own package, a `RunSpec` subclass, a +`LongRun` subclass and a CLI that calls the kernel `main`. + +## Migrating from 0.3.1a1 + +| 0.3.1a1 (kernel) | 0.3.1a2 (this extension) | +|---|---| +| `chronos corpus …` | `chronos-translate corpus …` | +| `chronos bundle …` | `chronos-translate bundle …` (a bundle's `bin/chronos bundle …` still works) | +| `chronos longrun --spec enzh_50m_72h.json` | `chronos-translate run --spec …` | +| `python -m chronos.mlx.training.doc_run` | `chronos-translate doc-run` | +| `chronos.data.corpus.*` | `chronos_translate.corpus.*` | +| `chronos.bundle.*` | `chronos_translate.bundle.*` | +| `inspect_checkpoint.translation_probe` | `chronos_translate.probe.translation_probe` | +| `chronos.model.factory.ENZH_DEMO_BASE` | `chronos_translate.model_presets.ENZH_50M_BASE` | + +Recipes keep their format: `bleu_samples` and `probe_samples` belong to +`TranslationRunSpec`, so the kernel runner rejects them with a pointer here. +Relative paths in a recipe now resolve against the recipe's own folder. + +## Data and weights + +WMT19 and OPUS-100 list their licences as unknown, so weights trained on them +are for research use only. FLORES+ (CC BY-SA 4.0) is used for evaluation only. diff --git a/demos/translation/README_zh.md b/demos/translation/README_zh.md new file mode 100644 index 0000000..3ac1e05 --- /dev/null +++ b/demos/translation/README_zh.md @@ -0,0 +1,93 @@ +# chronos-translate：基于 Chronos 内核的英↔中翻译 demo + +Chronos 内核（`chronos/`）用于训练通用模型，不含任何翻译特化代码。 +构建英↔中标杆模型所需的全部内容都放在这里，作为调用内核的扩展。 + +| 路径 | 内容 | +|---|---| +| `chronos_translate/corpus/` | 下载、清洗、去重并分片英中数据（WMT19、OPUS-100、FLORES+ 等）；`docmix` 构建 n 合 1 合并集 | +| `chronos_translate/run.py` | `TranslationLongRun`：内核 `LongRun` 加 FLORES+ BLEU/chrF 验证 | +| `chronos_translate/doc_run.py` | doc-mix 版本：微批自适应、`target_passes`、篇章级评测 | +| `chronos_translate/probe.py` | 对运行检查点做离线 FLORES+ 探测 | +| `chronos_translate/bundle/` | 自包含模型包：构建、校验、翻译、准备数据、微调、服务 | +| `chronos_translate/model_presets.py` | 50M 拓扑基础配置和分词器语言标记 | +| `recipes/enzh_50m_72h.json` | 72 小时标杆配方（路径相对于本文件） | +| `recipes/docmix_plan.json` | doc-mix 训练包 A/B/C 的规划器输出 | +| `scripts/` | `plan_docmix.py`、`make_docmix_packages.py` | +| `tests/` | 扩展测试（`pytest demos/translation/tests`） | + +## 安装 + +本扩展是仓库里的一个 demo，不发布到 PyPI，需要从仓库检出后安装。 +内核可以从 PyPI 安装（0.3.1 处于 beta 阶段，需加 `pip install --pre`），也可以用同一份检出安装。 + +```sh +git clone https://github.com/FonaTech/Project_Chronos.git && cd Project_Chronos +pip install -e '.[mlx,serve]' # 内核（或：pip install --pre 'Project_Chronos[mlx,serve]==0.3.1a2'） +pip install -e demos/translation # 本扩展（sacrebleu、sacremoses、pyarrow 固定版本） +chronos-translate --version +``` + +扩展固定依赖 `Project_Chronos==0.3.1a2`。请检出对应标签（`git checkout v0.3.1a2`），使两边版本一致。 +模型包需要这两个 wheel，可以从检出目录构建： + +```sh +python -m build --outdir dist . # dist/project_chronos-0.3.1a2-py3-none-any.whl +python -m build --outdir dist demos/translation # dist/chronos_translate-0.3.1a2-py3-none-any.whl +``` + +扩展测试在仓库根目录运行：`pytest demos/translation/tests`。 +内核 CI 不运行这些测试，因为它们需要 sacrebleu 和 pyarrow。 + +## 命令 + +```sh +chronos tokenizer data/pairs.jsonl --out tok --vocab-size 16384 \ + --extra-special "<|en|>" "<|zh|>" "<|translate|>" # 内核分词器 + 语言标记 +chronos-translate corpus download --help # 另有 build | shards | docmix +chronos-translate run --spec demos/translation/recipes/enzh_50m_72h.json +chronos-translate doc-run --spec <训练包>/recipe.json +chronos-translate inspect --run-dir runs/enzh_50m_v1 --samples 24 +chronos-translate bundle build --run-dir runs/enzh_50m_v1 --wheel dist/project_chronos-*.whl \ + --wheel dist/chronos_translate-*.whl +``` + +`chronos serve`、`chronos monitor` 和 `chronos tokenizer` 仍是内核命令。 + +## 如何接入内核 + +这个目录也是其他任务扩展的模板。它只使用内核公开的钩子，不修改内核代码： + +| 内核钩子 | 本扩展的用法 | +|---|---| +| `RunSpec` 子类 | `TranslationRunSpec` 增加 `bleu_samples`、`probe_samples` 和自己的 `export_prefix` | +| `LongRun.eval_samples_for` | 每个阶段的样本数（预训练探测、SFT BLEU 子集、最终 devtest） | +| `LongRun.task_eval` | 返回 `{指标: 数值}`；这些键写入 `metrics.jsonl` 的 `validation` / `final` 行，并显示在 monitor 中 | +| `LongRun.generate`、`write_samples` | 带 KV cache 的贪心解码；示例输出写入 `samples.jsonl` | +| `LongRun.pinned_modules` | 在运行记录中写入 sacrebleu 和扩展模块的版本 | +| `longrun.main(runner_cls=, spec_cls=, extra_path_keys=)` | `chronos-translate run` 命令行；`extra.raw_root` 相对于配方目录解析 | +| `special_tokens_with(...)` | 分词器中的 `<|en|>`、`<|zh|>`、`<|translate|>` 位置 | + +`task_eval` 出错时记为 `task_eval_error`，不会中断训练。 +新任务照搬这个结构即可：独立的包、一个 `RunSpec` 子类、一个 `LongRun` 子类，以及调用内核 `main` 的命令行。 + +## 从 0.3.1a1 迁移 + +| 0.3.1a1（内核） | 0.3.1a2（本扩展） | +|---|---| +| `chronos corpus …` | `chronos-translate corpus …` | +| `chronos bundle …` | `chronos-translate bundle …`（模型包里的 `bin/chronos bundle …` 仍可用） | +| `chronos longrun --spec enzh_50m_72h.json` | `chronos-translate run --spec …` | +| `python -m chronos.mlx.training.doc_run` | `chronos-translate doc-run` | +| `chronos.data.corpus.*` | `chronos_translate.corpus.*` | +| `chronos.bundle.*` | `chronos_translate.bundle.*` | +| `inspect_checkpoint.translation_probe` | `chronos_translate.probe.translation_probe` | +| `chronos.model.factory.ENZH_DEMO_BASE` | `chronos_translate.model_presets.ENZH_50M_BASE` | + +配方格式不变：`bleu_samples` 和 `probe_samples` 属于 `TranslationRunSpec`， +内核运行器遇到它们会报错并指向这里。配方中的相对路径现在相对于配方所在目录解析。 + +## 数据与权重 + +WMT19 和 OPUS-100 的许可证标注为 unknown，用它们训练的权重只能用于研究。 +FLORES+（CC BY-SA 4.0）只用于评测。 diff --git a/demos/translation/chronos_translate/__init__.py b/demos/translation/chronos_translate/__init__.py new file mode 100644 index 0000000..73a99ba --- /dev/null +++ b/demos/translation/chronos_translate/__init__.py @@ -0,0 +1,19 @@ +"""English↔Chinese translation demo built on the Chronos kernel. + +Everything translation-specific lives here, outside the ``chronos`` kernel: + +* ``corpus``: download, clean, dedup and shard en–zh parallel data + (WMT19, OPUS-100, FLORES+, …) into the kernel's shard layout; + ``corpus.docmix`` builds merged multi-sentence ("n-in-1") sets; +* ``run``: ``TranslationLongRun``, the kernel ``LongRun`` with BLEU/chrF + validation (``chronos-translate run --spec …``); +* ``doc_run``: the doc-mix variant with document-level evaluation; +* ``probe``: offline FLORES+ probe for a run checkpoint; +* ``bundle``: self-contained model bundles (build, check, translate, + prepare-data, fine-tune, serve); +* ``model_presets``: the 50M en→zh model topology and tokenizer tokens. + +The kernel needs none of this to train a general model. +""" + +__version__ = "0.3.1a2" diff --git a/chronos/bundle/__init__.py b/demos/translation/chronos_translate/bundle/__init__.py similarity index 66% rename from chronos/bundle/__init__.py rename to demos/translation/chronos_translate/bundle/__init__.py index 97e00bb..e87896d 100644 --- a/chronos/bundle/__init__.py +++ b/demos/translation/chronos_translate/bundle/__init__.py @@ -1,4 +1,4 @@ -"""Out-of-the-box model bundles (``chronos bundle``). +"""Out-of-the-box translation model bundles (``chronos-translate bundle``). A bundle is one self-contained directory: @@ -9,13 +9,13 @@ 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 + recipes/*.json ``chronos-translate run`` 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. +``chronos-translate bundle build`` creates it from a training run, ``... check`` +verifies hashes and runs a CPU forward smoke test, ``... translate`` +and ``... finetune`` are the two entry points a recipient needs. """ from .builder import BUNDLE_FORMAT, build_bundle, check_bundle, load_bundle_manifest diff --git a/chronos/bundle/builder.py b/demos/translation/chronos_translate/bundle/builder.py similarity index 96% rename from chronos/bundle/builder.py rename to demos/translation/chronos_translate/bundle/builder.py index f13116b..f55498e 100644 --- a/chronos/bundle/builder.py +++ b/demos/translation/chronos_translate/bundle/builder.py @@ -144,9 +144,12 @@ def _planner_preset(vocab_size: int, hidden_size: int) -> dict | 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, + name: str | None = None, minimind_dir: str | None = None, wheel=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).""" + """Assemble a bundle from a run directory (read-only on the run). + + ``wheel``: one path or a list of paths (the kernel and the + ``chronos_translate`` wheels) copied to ``wheels/`` for ``bin/setup.sh``.""" from chronos.model.checkpoint import checkpoint_config_path, tokenizer_metadata from . import docs, recommend run_dir = os.path.abspath(run_dir) @@ -211,8 +214,8 @@ def build_bundle(run_dir: str, out_dir: str | None = None, *, phase: str = "pret 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))) + for whl in ([wheel] if isinstance(wheel, str) else list(wheel or ())): + _copy(whl, os.path.join(out_dir, "wheels", os.path.basename(whl))) # Recipes: the exact spec that produced these weights (paths made relative). roots = _path_roots(run_dir) source_spec = _sanitize(spec, roots) @@ -237,6 +240,7 @@ def build_bundle(run_dir: str, out_dir: str | None = None, *, phase: str = "pret "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(), + "chronos_translate_version": _extension_version(), } with open(os.path.join(out_dir, MANIFEST), "w", encoding="utf-8") as fh: json.dump(manifest, fh, ensure_ascii=False, indent=2) @@ -264,6 +268,11 @@ def _chronos_version() -> str: return getattr(chronos, "__version__", "unknown") +def _extension_version() -> str: + from .. import __version__ + return __version__ + + 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: diff --git a/chronos/bundle/cli.py b/demos/translation/chronos_translate/bundle/cli.py similarity index 94% rename from chronos/bundle/cli.py rename to demos/translation/chronos_translate/bundle/cli.py index 5450c7a..d4faeca 100644 --- a/chronos/bundle/cli.py +++ b/demos/translation/chronos_translate/bundle/cli.py @@ -1,4 +1,4 @@ -"""``chronos bundle {build|pack|check|probe|translate|eval|prepare-data|finetune|serve}``.""" +"""``chronos-translate bundle {build|pack|check|probe|translate|eval|prepare-data|finetune|serve}``.""" from __future__ import annotations import argparse @@ -50,17 +50,18 @@ def _bundle_model(bundle: str, a): def main(argv=None) -> int: - p = argparse.ArgumentParser(prog="chronos bundle", description=__doc__) + p = argparse.ArgumentParser(prog="chronos-translate 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 = sub.add_parser("build", help="assemble a bundle from a 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("--wheel", action="append", default=None, + help="a wheel to ship under wheels/ (repeat: the project_chronos and chronos_translate 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") @@ -116,7 +117,7 @@ def main(argv=None) -> int: 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 = sub.add_parser("finetune", help="write a run 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) @@ -223,7 +224,7 @@ def _run(a) -> int: clusters=a.clusters) print(f"[finetune] spec: {spec}") if a.dry_run: - print(f"[finetune] start it with: chronos longrun --spec {spec}") + print(f"[finetune] start it with: chronos-translate run --spec {spec}") return 0 return launch(spec) if a.cmd == "serve": diff --git a/chronos/bundle/docs.py b/demos/translation/chronos_translate/bundle/docs.py similarity index 94% rename from chronos/bundle/docs.py rename to demos/translation/chronos_translate/bundle/docs.py index d6db830..75673b1 100644 --- a/chronos/bundle/docs.py +++ b/demos/translation/chronos_translate/bundle/docs.py @@ -60,16 +60,21 @@ def write_examples(out_dir: str) -> None: LAUNCHER = """#!/bin/sh -# Runs `chronos` with this bundle's settings (no git, no network needed). +# Runs Chronos with this bundle's settings (no git, no network needed). +# `bin/chronos bundle …` goes to the translation extension (`chronos-translate +# bundle …`); every other command goes to the Chronos kernel CLI. 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 "$@" +BIN=chronos +if [ "${1:-}" = bundle ]; then BIN=chronos-translate; fi +if [ -x "$HERE/.venv/bin/$BIN" ]; then exec "$HERE/.venv/bin/$BIN" "$@"; fi +exec "$BIN" "$@" """ SETUP = """#!/bin/sh -# One-time setup: a virtualenv in .venv with Chronos and its pinned extras. +# One-time setup: a virtualenv in .venv with the Chronos kernel, the +# translation extension (both from wheels/) and their pinned extras. set -eu HERE="$(cd "$(dirname "$0")/.." && pwd)" PY="${PYTHON:-python3}" @@ -77,13 +82,15 @@ def write_examples(out_dir: str) -> None: [ -d "$HERE/.venv" ] || "$PY" -m venv "$HERE/.venv" . "$HERE/.venv/bin/activate" python -m pip install --upgrade pip -EXTRAS="translate,serve" +EXTRAS="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]" +KERNEL="$(ls "$HERE"/wheels/project_chronos-*.whl 2>/dev/null | head -n 1 || true)" +EXT="$(ls "$HERE"/wheels/chronos_translate-*.whl 2>/dev/null | head -n 1 || true)" +if [ -n "$KERNEL" ] && [ -n "$EXT" ]; then + python -m pip install "$KERNEL[$EXTRAS]" "$EXT" else - echo "no wheel in wheels/: install Chronos from source with pip install -e '[$EXTRAS]'" >&2 + echo "wheels/ needs project_chronos-*.whl and chronos_translate-*.whl; from a source tree use" >&2 + echo " pip install -e '[$EXTRAS]' -e '/demos/translation'" >&2 exit 1 fi "$HERE/bin/chronos" bundle check @@ -147,7 +154,7 @@ def _readme_en(m: dict, recs: dict) -> str: "## Quick start", "", "```sh", "tar xzf " + m["name"] + ".tar.gz && cd " + m["name"], - "sh bin/setup.sh # .venv + pinned extras + `chronos bundle check`", + "sh bin/setup.sh # .venv + pinned extras + `bin/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", @@ -166,7 +173,7 @@ def _readme_en(m: dict, recs: dict) -> str: "The page counts are printed to stderr.", "", "`bin/chronos` sets `CHRONOS_BUNDLE` and `CHRONOS_MINIMIND_DIR` (the vendored minimind in " "`third_party/`), so nothing is fetched from the network. Without `bin/`, export both variables and call " - "`chronos bundle …` directly.", "", + "`chronos-translate bundle …` directly.", "", "## Next stage: SFT (fine-tuning) on Apple silicon (MLX)", "", "```sh", "# 1. your pairs: .jsonl {\"en\":…, \"zh\":…} | .csv | .tsv | .parquet", @@ -181,7 +188,7 @@ def _readme_en(m: dict, recs: dict) -> str: "`--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 " + "`work/run/*_ft_sft.best.pth`; rebuild a bundle from it with `chronos-translate bundle build --run-dir work/run " "--phase sft`.", "", "Fine-tuning trains with the expert-cluster locality constraint by default, so the result can be paged " "by cluster at inference (`--no-cluster-locality` turns it off). Defaults: `--lambda-cluster 0.3` (from " @@ -258,7 +265,7 @@ def _readme_zh(m: dict, recs: dict) -> str: "可用 `--decay`、`--ram-ratio`、`--decode-wait`、`--verify-wait` 调整，默认值见 `--help`。" "页读取统计打印到 stderr。", "", "`bin/chronos` 会设置 `CHRONOS_BUNDLE` 和 `CHRONOS_MINIMIND_DIR`（指向 `third_party/` 里随包附带的 minimind），" - "全程不需要联网。不用 `bin/` 时，手动导出这两个变量后直接调用 `chronos bundle …`。", "", + "全程不需要联网。不用 `bin/` 时，手动导出这两个变量后直接调用 `chronos-translate bundle …`。", "", "## 下一阶段：SFT 微调（Apple silicon / MLX）", "", "```sh", "# 1. 准备平行语料：.jsonl {\"en\":…, \"zh\":…} | .csv | .tsv | .parquet", @@ -273,7 +280,7 @@ def _readme_zh(m: dict, recs: dict) -> str: "`--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` 重新打包。", "", + "可用 `chronos-translate bundle build --run-dir work/run --phase sft` 重新打包。", "", "微调默认带专家簇局部性约束训练，这样得到的模型在推理时可以按簇分页（`--no-cluster-locality` 关闭）。" "默认值：`--lambda-cluster 0.3`（来自合成任务上从头训练的实测扫描，对微调的效果尚未实测）；" "`--cluster-window` 取最接近你数据中位 SFT 记录长度的 2 的幂。spec 中记录了每个值的来源" diff --git a/chronos/bundle/finetune.py b/demos/translation/chronos_translate/bundle/finetune.py similarity index 91% rename from chronos/bundle/finetune.py rename to demos/translation/chronos_translate/bundle/finetune.py index 38278d5..f2066ce 100644 --- a/chronos/bundle/finetune.py +++ b/demos/translation/chronos_translate/bundle/finetune.py @@ -41,11 +41,11 @@ def prepare_data(bundle_dir: str, inputs: list[str], out_dir: str, *, zh_field: ``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 + from chronos_translate.corpus.clean import CleanConfig + from chronos_translate.corpus.evalsets import load_eval_sets + from chronos_translate.corpus.pipeline import PairSource, build_corpus + from chronos_translate.corpus.readers import count_rows + from chronos_translate.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")): @@ -105,7 +105,7 @@ def write_finetune_spec(bundle_dir: str, shards_dir: str, run_dir: str, *, tier: budget_hours: float | None = None, seed: int | None = None, cluster_locality: bool = True, lambda_cluster: float | None = None, cluster_window: int | None = None, clusters: int | None = None) -> str: - """Write ``/finetune_spec.json`` for ``chronos longrun --spec``. + """Write ``/finetune_spec.json`` for ``chronos-translate run --spec``. The expert-cluster locality constraint is on by default; see ``recommend.cluster_locality_config`` for the defaults and their sources.""" @@ -120,15 +120,15 @@ def write_finetune_spec(bundle_dir: str, shards_dir: str, run_dir: str, *, tier: 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. + # the run 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") + f"`chronos-translate run --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 + from ..run import main as run_main # TranslationRunSpec: bleu_samples / probe_samples argv = ["--spec", spec_path] + (["--budget-hours", str(budget_hours)] if budget_hours is not None else []) - return longrun_main(argv) + return run_main(argv) diff --git a/chronos/bundle/recommend.py b/demos/translation/chronos_translate/bundle/recommend.py similarity index 74% rename from chronos/bundle/recommend.py rename to demos/translation/chronos_translate/bundle/recommend.py index fb6401a..0d4be7d 100644 --- a/chronos/bundle/recommend.py +++ b/demos/translation/chronos_translate/bundle/recommend.py @@ -10,7 +10,7 @@ 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. +turns a tier into a ``chronos-translate run`` spec. """ from __future__ import annotations @@ -18,6 +18,9 @@ import json import os +# Generic helpers (kernel); re-exported for existing imports. +from chronos.expert_clusters import cluster_locality_config, median_record_tokens # noqa: F401 + # 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)"}, @@ -98,71 +101,11 @@ def tier_phases(recipe: dict, tier: str, mode: str = "sft") -> list[dict]: raise ValueError(f"unknown mode {mode!r}; choose sft, pretrain or pretrain+sft") -def median_record_tokens(shards_dir: str, max_files: int = 8) -> float | None: - """Median SFT record length (tokens) from the ``.idx`` offsets under ``/sft/train``. - - Reads at most ``max_files`` index files (offsets only, no token data). - Returns None when the shards have no SFT records (packed pretrain text - has no record boundaries).""" - import glob - import numpy as np - files = sorted(glob.glob(os.path.join(shards_dir, "sft", "train", "*.idx")))[:max_files] - lengths = [np.diff(np.fromfile(f, dtype=np.int64)) for f in files] - lengths = [x for x in lengths if x.size] - if not lengths: - return None - return float(np.median(np.concatenate(lengths))) - - -def cluster_locality_config(model_config: dict, shards_dir: str | None, *, enabled: bool = True, - lambda_cluster: float | None = None, window: int | None = None, - clusters: int | None = None, seq_len: int | None = None) -> dict: - """Cluster-locality fields for a fine-tune and where each value comes from. - - On by default at ``DEFAULT_LAMBDA`` (measured sweep on a synthetic - multi-topic task, from-scratch training; its effect on a fine-tune of - weights trained without the constraint has not been measured). Off for - dense models and when there are fewer than two clusters.""" - from chronos.expert_clusters import DEFAULT_LAMBDA, recommended_window - from chronos.model.config import ChronosConfig - # The values training will use: longrun builds ChronosConfig(**model_config). - cfg = ChronosConfig(**{k: v for k, v in model_config.items() if v is not None}) - e = int(getattr(cfg, "num_experts", 0) or 0) - k = int(getattr(cfg, "num_experts_per_tok", 1) or 1) - seq = int(seq_len or getattr(cfg, "max_position_embeddings", 512) or 512) - usable = bool(getattr(cfg, "use_moe", True)) and e > 2 * k - lam = 0.0 if not (enabled and usable) else float(DEFAULT_LAMBDA if lambda_cluster is None else lambda_cluster) - if lam < 0: - raise ValueError("lambda_cluster must be >= 0") - median = median_record_tokens(shards_dir) if (shards_dir and window is None) else None - win = int(window) if window is not None else recommended_window(median, seq) - if win < 0: - raise ValueError("cluster_window must be >= 0 (0 = whole sequence)") - values = {"lambda_cluster_locality": lam, "cluster_locality_window": win, - "cluster_locality_clusters": int(clusters or 0)} - if not enabled: - why = "off (--no-cluster-locality)" - elif not usable: - why = "off: dense model or fewer than two clusters" - elif lambda_cluster is None: - why = f"default {DEFAULT_LAMBDA} (measured sweep, synthetic task; not measured for fine-tuning)" - else: - why = "user value" - if window is not None: - wsrc = "user value" - elif median is not None: - wsrc = f"rule: power of two nearest the median SFT record ({median:.0f} tokens), clamped to [16, {seq}]" - else: - wsrc = "default (no SFT record lengths in the data)" - return {"values": values, "source": {"lambda_cluster_locality": why, "cluster_locality_window": wsrc, - "cluster_locality_clusters": "user value" if clusters else "ceil(E / 2k)"}} - - def longrun_spec(bundle_dir: str, recipe: dict, *, run_dir: str, shards_dir: str, tier: str, mode: str = "sft", budget_hours: float | None = None, seed: int | None = None, cluster_locality: bool = True, lambda_cluster: float | None = None, cluster_window: int | None = None, clusters: int | None = None) -> dict: - """A ``chronos longrun`` RunSpec (as a dict) that starts from the bundle weights. + """A ``TranslationRunSpec`` (as a dict) that starts from the bundle weights. The cluster-locality constraint is on by default (``cluster_locality_config``).""" from .builder import load_bundle_manifest @@ -253,7 +196,7 @@ def build_recommendations(recipe: dict, planner_preset: dict | None, decode: dic "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"), + }, "recipe", "defaults of chronos_translate.corpus (build_corpus / build_shards) used for the bundled weights"), } if decode: out["decode"] = decode diff --git a/chronos/bundle/runtime.py b/demos/translation/chronos_translate/bundle/runtime.py similarity index 99% rename from chronos/bundle/runtime.py rename to demos/translation/chronos_translate/bundle/runtime.py index 51b601f..3fbdd52 100644 --- a/chronos/bundle/runtime.py +++ b/demos/translation/chronos_translate/bundle/runtime.py @@ -112,7 +112,7 @@ def prompt_ids(self, src: str, direction: str = "en_zh") -> list[int]: 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 + from chronos_translate.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 @@ -246,7 +246,7 @@ def translate(self, sources: list[str], direction: str = "en_zh", *, max_new_tok 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 + from chronos_translate.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, []) diff --git a/demos/translation/chronos_translate/cli.py b/demos/translation/chronos_translate/cli.py new file mode 100644 index 0000000..dce2cfa --- /dev/null +++ b/demos/translation/chronos_translate/cli.py @@ -0,0 +1,59 @@ +"""``chronos-translate``: the en↔zh translation demo's command line. + + chronos-translate corpus download|build|shards|docmix … data pipeline + chronos-translate run --spec recipes/enzh_50m_72h.json time-budgeted MLX run + chronos-translate doc-run --spec doc-mix variant + chronos-translate inspect --run-dir runs/… offline FLORES+ probe + chronos-translate bundle build|check|translate|… model bundles + +Kernel commands (``chronos serve``, ``chronos monitor``, ``chronos tokenizer``, +…) stay in the ``chronos`` CLI. +""" +from __future__ import annotations + +import importlib +import sys + +CORPUS = { + "download": "chronos_translate.corpus.sources", + "build": "chronos_translate.corpus.presets", + "shards": "chronos_translate.corpus.shards", + "docmix": "chronos_translate.corpus.docmix", +} +COMMANDS = { + "run": "chronos_translate.run", + "doc-run": "chronos_translate.doc_run", + "inspect": "chronos_translate.probe", + "bundle": "chronos_translate.bundle.cli", +} + + +def _usage() -> str: + return (__doc__ or "").strip() + "\n" + + +def main(argv=None) -> int: + argv = list(sys.argv[1:] if argv is None else argv) + if not argv or argv[0] in {"-h", "--help", "help"}: + print(_usage()) + return 0 if argv else 2 + if argv[0] in {"-V", "--version"}: + from . import __version__ + print(__version__) + return 0 + cmd, rest = argv[0], argv[1:] + if cmd == "corpus": + if not rest or rest[0] not in CORPUS: + print("usage: chronos-translate corpus {" + "|".join(CORPUS) + "} [options]", file=sys.stderr) + return 2 + module, rest = CORPUS[rest[0]], rest[1:] + elif cmd in COMMANDS: + module = COMMANDS[cmd] + else: + print(f"unknown command {cmd!r}\n\n{_usage()}", file=sys.stderr) + return 2 + return int(importlib.import_module(module).main(rest) or 0) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/chronos/data/corpus/__init__.py b/demos/translation/chronos_translate/corpus/__init__.py similarity index 100% rename from chronos/data/corpus/__init__.py rename to demos/translation/chronos_translate/corpus/__init__.py diff --git a/chronos/data/corpus/clean.py b/demos/translation/chronos_translate/corpus/clean.py similarity index 100% rename from chronos/data/corpus/clean.py rename to demos/translation/chronos_translate/corpus/clean.py diff --git a/demos/translation/chronos_translate/corpus/docmix.py b/demos/translation/chronos_translate/corpus/docmix.py new file mode 100644 index 0000000..dee18ff --- /dev/null +++ b/demos/translation/chronos_translate/corpus/docmix.py @@ -0,0 +1,566 @@ +"""Document-mix translation shards: science-weighted subsets, merged n-in-1 +samples up to a context limit, and bidirectional SFT records. + +Why: ``shards_v1`` holds single sentences (SFT p50 75 / p99 225 tokens), so a +model trained on it never sees long inputs. Here consecutive rows of the +same corpus unit are merged into one sample (2-in-1 … n-in-1). + +Measured on ``corpus_v1/pairs`` (adjacent-pair vs random-pair content-word +overlap): most WMT19 units and iic keep document order (ratio 2–25), while +OPUS-100 and WMT19 unit 00012 are shuffled (ratio ≈ 1). Each block is +classified here with the same statistic; merges inside an ordered block are +coherent discourse. Merges of shuffled rows are random concatenations: they +teach length and position handling but not coherence, and the manifest +counts them separately. Document boundaries are not stored in the corpus, +so some ordered merges cross a document boundary. + +Layout of ``out_dir`` (same readers as ``shards.py`` plus one extra file): + pretrain/{train,valid}/*.bin flat tokens, ``<|endoftext|>``-separated + sft/{train,valid}/.bin/.idx/.plen as in shards.py + sft/{train,valid}/.dir uint8 per record: 0 = en→zh, 1 = zh→en + sft/{train,valid}/.n uint8 per record: sentences merged (1..255) +""" +from __future__ import annotations + +import json +import math +import os +import random +import re +import time +from concurrent.futures import ProcessPoolExecutor +from dataclasses import asdict, dataclass, field + +import numpy as np + +from chronos_translate.corpus.shards import SFT_PROMPTS, _dtype_for + +DIRECTIONS = ("en_zh", "zh_en") +DIR_CODE = {d: i for i, d in enumerate(DIRECTIONS)} +# Share of SAMPLES per merge size; "long" draws n uniformly in [6, n_max] +# where n_max is what fits the context limit. +DEFAULT_SHARES = {"1": 0.40, "2": 0.20, "3": 0.15, "4": 0.10, "5": 0.05, "long": 0.10} + +# Science score = keyword hits per row (English: a trailing "*" is a prefix, +# otherwise a whole word with an optional plural "s"; Chinese: substrings). +# Generic UN/policy vocabulary (measure, general, physical, data protection) +# is avoided on purpose. Chunks are ranked by mean score per row. +SCIENCE_EN = ( + "research*", "scien*", "experiment*", "laborator*", "clinic*", "patient", "disease", "medic*", + "vaccin*", "virus", "viruses", "bacteri*", "gene", "genes", "genetic*", "genom*", "protein", "cell", + "molecul*", "chemi*", "physics", "physicist", "quantum", "particle", "atom", "atomic", "nuclear", + "energy", "solar", "battery", "batteries", "electron*", "electric*", "climate", "carbon", "emission", + "ocean*", "geolog*", "ecolog*", "ecosystem", "species", "evolution*", "biolog*", "neuro*", "brain", + "algorithm*", "comput*", "software", "robot*", "artificial intelligence", "machine learning", + "engineer*", "technolog*", "satellite", "planet*", "astronom*", "telescope", "mathemat*", "equation", + "statistic*", "hypothes*", "sensor", "semiconductor", "polymer", "catalys*", "reactor", "fusion", + "laser", "optic*", "spectr*", "drug", "therap*", "diagnos*", "infection", "cancer", "tumou?r", "dna", + "rna", "enzyme", "microb*", "fossil", "isotope", "radiation", "theorem", "physiolog*", "pharmac*", +) + + +def _en_pattern(term: str) -> str: + if term.endswith("*"): + return re.escape(term[:-1]) + r"\w*" + body = term if "?" in term else re.escape(term) + return body + r"s?\b" + + + +SCIENCE_ZH = ( + "研究", "科学", "实验", "实验室", "临床", "患者", "疾病", "医学", "医疗", "健康", "疫苗", "病毒", "细菌", + "基因", "蛋白", "细胞", "分子", "化学", "物理", "量子", "粒子", "原子", "核能", "能源", "太阳能", "电池", + "电子", "气候", "碳", "排放", "海洋", "地质", "生态", "物种", "进化", "生物", "神经", "大脑", "算法", + "计算", "软件", "数据", "网络", "机器人", "人工智能", "机器学习", "工程", "技术", "卫星", "太空", "行星", + "天文", "望远镜", "数学", "方程", "统计", "理论", "假设", "测量", "传感", "材料", "半导体", "聚合物", + "催化", "反应堆", "聚变", "激光", "光学", "光谱", "药物", "治疗", "诊断", +) +_SCI_EN_RE = re.compile(r"\b(?:" + "|".join(_en_pattern(s) for s in SCIENCE_EN) + r")") +_SCI_ZH_RE = re.compile("|".join(re.escape(s) for s in sorted(SCIENCE_ZH, key=len, reverse=True))) +# Adjacency statistic (same as the measurement that motivated this module). +_WORD = re.compile(r"[a-z]{4,}") +_STOP = frozenset(( + "that this with from have will been were they their there which would could about into than them these " + "those what when where while also more most such only over other some very after before because being " + "between both each every however should since through under until upon your said says year years").split()) + + +def science_hits(en: str, zh: str) -> int: + """Keyword hits of one pair (English side lower-cased, Chinese side as is).""" + return len(_SCI_EN_RE.findall(en.lower())) + len(_SCI_ZH_RE.findall(zh)) + + +def _bag(text: str) -> frozenset: + return frozenset(w for w in _WORD.findall(text.lower()) if w not in _STOP) + + +@dataclass +class DocMixSpec: + fraction: float # share of train rows kept (0.15 = 15 %) + ctx_len: int # SFT record cap = pretrain window (tokens) + shares: dict = field(default_factory=lambda: dict(DEFAULT_SHARES)) + long_min: int = 6 # "long" samples: n ~ U[long_min, n_cap] + zh_en_share: float = 1 / 3 # en→zh : zh→en = 2 : 1 (SFT and tagged pretrain docs) + mono_share: float = 0.5 # pretrain: monolingual docs; the rest are tagged pairs + para_break: float = 0.15 # chance that a sentence boundary becomes a newline + science_weights: tuple = (2.0, 1.2, 1.0) # high / medium / other + science_quantiles: tuple = (0.10, 0.20) # chunk shares of the high / medium tier + chunk_rows: int = 256 # selection unit (consecutive rows of one part) + block_rows: int = 65536 # adjacency classification window + ordered_ratio: float = 2.0 # adjacent/random overlap ratio for "ordered" + ordered_min_adj: float = 0.05 # ... and at least this adjacent overlap + seed: int = 20261107 + + +def _scan_part(args) -> dict: + path, block_rows, seed = args + import pyarrow.parquet as pq + table = pq.read_table(path, columns=["zh", "en"]) + zh, en = table["zh"].to_pylist(), table["en"].to_pylist() + sci = np.fromiter((science_hits(e, z) for e, z in zip(en, zh)), dtype=np.int16, count=len(en)) + rng = random.Random(seed) + blocks = [] + for b0 in range(0, len(en), block_rows): + bags = [_bag(s) for s in en[b0:b0 + block_rows]] + n = len(bags) + adj = sum(1 for a, b in zip(bags, bags[1:]) if a & b) + k = min(20000, 2 * n) if n > 1 else 0 + rnd = sum(1 for _ in range(k) if bags[rng.randrange(n)] & bags[rng.randrange(n)]) + blocks.append((adj, max(0, n - 1), rnd, k)) + return {"name": os.path.basename(path), "rows": len(en), "sci": sci, + "blocks": np.asarray(blocks, dtype=np.int64).reshape(-1, 4)} + + +def scan_corpus(pairs_dir: str, cache_dir: str, *, block_rows: int = 65536, workers: int = 8, + log=print) -> dict: + """Science hits per row and adjacency counts per block, cached in ``cache_dir``.""" + os.makedirs(cache_dir, exist_ok=True) + meta_path = os.path.join(cache_dir, "scan.json") + names = sorted(p for p in os.listdir(pairs_dir) if p.endswith(".parquet")) + if os.path.exists(meta_path): + meta = json.load(open(meta_path, encoding="utf-8")) + if meta.get("names") == names and meta.get("block_rows") == block_rows: + data = np.load(os.path.join(cache_dir, "scan.npz")) + return {"names": names, "block_rows": block_rows, "rows": meta["rows"], + "sci": {n: data[f"sci/{n}"] for n in names}, "blocks": {n: data[f"blk/{n}"] for n in names}} + t0 = time.time() + jobs = [(os.path.join(pairs_dir, n), block_rows, 17 + i) for i, n in enumerate(names)] + with ProcessPoolExecutor(max_workers=workers) as pool: + results = {r["name"]: r for r in pool.map(_scan_part, jobs)} + out = {"names": names, "block_rows": block_rows, "rows": {n: results[n]["rows"] for n in names}, + "sci": {n: results[n]["sci"] for n in names}, "blocks": {n: results[n]["blocks"] for n in names}} + np.savez(os.path.join(cache_dir, "scan.npz"), **{f"sci/{n}": out["sci"][n] for n in names}, + **{f"blk/{n}": out["blocks"][n] for n in names}) + json.dump({"names": names, "block_rows": block_rows, "rows": out["rows"], "keywords_en": SCIENCE_EN, + "keywords_zh": SCIENCE_ZH, "elapsed_s": round(time.time() - t0, 1)}, + open(meta_path, "w", encoding="utf-8"), ensure_ascii=False, indent=1) + log(f"[docmix] scanned {len(names)} parts in {time.time() - t0:.0f}s") + return out + + +def block_ordered(blocks: np.ndarray, ratio: float, min_adj: float) -> np.ndarray: + """Per block: adjacent-pair overlap ≥ ``ratio`` × random-pair overlap and ≥ ``min_adj``.""" + adj = blocks[:, 0] / np.maximum(blocks[:, 1], 1) + rnd = blocks[:, 2] / np.maximum(blocks[:, 3], 1) + return (adj >= ratio * np.maximum(rnd, 1e-3)) & (adj >= min_adj) + + +def select_chunks(scan: dict, spec: DocMixSpec) -> dict: + """Science-weighted chunk sample of the train parts, without replacement. + + Chunks (``chunk_rows`` consecutive rows) are ranked by mean keyword hits + per row; the top ``science_quantiles[0]`` get weight ``science_weights[0]``, + the next ``science_quantiles[1]`` the medium weight, the rest 1. Keys + ``log(u)/w`` (Efraimidis–Spirakis) pick chunks until ``fraction`` of all + train rows is reached; consecutive picked chunks form one run. + """ + names = [n for n in scan["names"] if n.startswith("train.")] + rows_of, cr = scan["rows"], spec.chunk_rows + part, start, stop, mean, ordered = [], [], [], [], [] + for pi, name in enumerate(names): + r = rows_of[name] + starts = np.arange(0, r, cr) + stops = np.minimum(starts + cr, r) + sci = np.add.reduceat(scan["sci"][name].astype(np.float64), starts) if r else np.zeros(0) + blk = block_ordered(scan["blocks"][name], spec.ordered_ratio, spec.ordered_min_adj) + part.append(np.full(len(starts), pi)); start.append(starts); stop.append(stops) + mean.append(sci / (stops - starts)); ordered.append(blk[starts // scan["block_rows"]]) + part, start, stop = np.concatenate(part), np.concatenate(start), np.concatenate(stop) + mean, ordered = np.concatenate(mean), np.concatenate(ordered) + q_hi, q_mid = spec.science_quantiles + t_hi = float(np.quantile(mean, 1 - q_hi)) + t_mid = float(np.quantile(mean, 1 - q_hi - q_mid)) + tier = np.where(mean >= t_hi, 0, np.where(mean >= t_mid, 1, 2)) + weight = np.asarray(spec.science_weights, dtype=np.float64)[tier] + rng = np.random.default_rng(spec.seed) + keys = np.log(rng.random(len(part))) / weight + order = np.argsort(-keys, kind="stable") + total = int((stop - start).sum()) + target = int(round(spec.fraction * total)) + cum = np.cumsum((stop - start)[order]) + take = np.sort(order[:int(np.searchsorted(cum, target)) + 1]) + runs = {name: [] for name in names} + for ci in take: + lst = runs[names[part[ci]]] + s, e, o = int(start[ci]), int(stop[ci]), bool(ordered[ci]) + if lst and lst[-1][1] == s and lst[-1][2] == o: + lst[-1][1] = e # extend the run + else: + lst.append([s, e, o]) + rows = stop - start + stats = {"train_rows_total": total, "rows_selected": int(rows[take].sum()), + "chunks_selected": int(len(take)), "chunks_total": int(len(part)), + "tier_thresholds_mean_hits": {"high": round(t_hi, 4), "medium": round(t_mid, 4)}, + "tier_row_share_all": [round(float(rows[tier == t].sum() / total), 4) for t in range(3)], + "tier_row_share_selected": [round(float(rows[take][tier[take] == t].sum() / max(rows[take].sum(), 1)), 4) + for t in range(3)], + "mean_hits_per_row_all": round(float((mean * rows).sum() / total), 4), + "mean_hits_per_row_selected": round(float((mean[take] * rows[take]).sum() / max(rows[take].sum(), 1)), 4), + "ordered_row_share_selected": round(float(rows[take][ordered[take]].sum() / max(rows[take].sum(), 1)), 4)} + return {"runs": {k: v for k, v in runs.items() if v}, "stats": stats} + + +SIZE_KEYS = ("1", "2", "3", "4", "5", "long") +LEN_BINS = (0, 128, 256, 512, 1024, 2048, 3072, 4096, 1 << 30) + + +def merge_plan(cost: np.ndarray, cap: int, spec: DocMixSpec, rng: random.Random) -> list: + """Split consecutive rows into samples: [(start, stop, size_key), ...]. + + ``cost`` is the token cost of each row inside a merged sample; a sample + never exceeds ``cap``. The size key is drawn from ``spec.shares``; "long" + takes n ~ U[long_min, n_fit] (n_fit = rows that still fit the cap). When + fewer rows fit than requested the sample is shortened; a row that alone + exceeds the cap is skipped (counted by the caller). + """ + keys = list(SIZE_KEYS) + weights = [float(spec.shares.get(k, 0.0)) for k in keys] + cs = np.concatenate([[0], np.cumsum(cost, dtype=np.int64)]) + out, i, n_rows = [], 0, len(cost) + while i < n_rows: + n_fit = int(np.searchsorted(cs, cs[i] + cap, side="right")) - 1 - i + n_fit = min(n_fit, n_rows - i) + if n_fit <= 0: + i += 1 + continue + key = rng.choices(keys, weights)[0] + if key == "long": + n = n_fit if n_fit <= spec.long_min else rng.randint(spec.long_min, n_fit) + else: + n = min(int(key), n_fit) + out.append((i, i + n, key)) + i += n + return out + + +def _join(rows: list, sep: str, breaks: list) -> str: + parts = [rows[0]] + for text, br in zip(rows[1:], breaks): + parts.append("\n" if br else sep) + parts.append(text) + return "".join(parts) + + +def _write_arrays(prefix: str, arrays: dict) -> None: + for suffix, arr in arrays.items(): + tmp = prefix + suffix + ".tmp" + with open(tmp, "wb") as fh: + arr.tofile(fh) + os.replace(tmp, prefix + suffix) + + +def _len_bin(n: int) -> str: + for lo, hi in zip(LEN_BINS, LEN_BINS[1:]): + if n <= hi: + return f"{lo + 1}-{hi}" if hi < (1 << 30) else f">{lo}" + return "?" + + +def _sft_samples(tok, en, zh, en_len, zh_len, runs_local, spec, rng, overhead): + """Merged chat records for one part. ``runs_local``: [(start, stop, ordered)] in row space.""" + cap = spec.ctx_len - overhead + prompts, replies, meta = [], [], [] + cost = np.asarray(en_len, np.int64) + np.asarray(zh_len, np.int64) + 1 + for s, e, ordered in runs_local: + for a, b, key in merge_plan(cost[s:e], cap, spec, rng): + a, b = a + s, b + s + breaks = [rng.random() < spec.para_break for _ in range(b - a - 1)] + en_text = _join(en[a:b], " ", breaks) + zh_text = _join(zh[a:b], "", breaks) + d = "zh_en" if rng.random() < spec.zh_en_share else "en_zh" + src, tgt = (en_text, zh_text) if d == "en_zh" else (zh_text, en_text) + prompts.append(f"<|im_start|>user\n{SFT_PROMPTS[d]}{src}<|im_end|>\n<|im_start|>assistant\n") + replies.append(f"{tgt}<|im_end|>\n") + meta.append((DIR_CODE[d], b - a, key, bool(ordered))) + p_ids = [x.ids for x in tok.encode_batch(prompts, add_special_tokens=False)] if prompts else [] + r_ids = [x.ids for x in tok.encode_batch(replies, add_special_tokens=False)] if replies else [] + return p_ids, r_ids, meta + + +def _pretrain_docs(tok, en, zh, en_len, zh_len, runs_local, spec, rng, ids): + """Merged documents for the packed pretrain stream (same merge shares).""" + eot, en_tag, zh_tag = ids + cost = np.asarray(en_len, np.int64) + np.asarray(zh_len, np.int64) + 1 + plan = [] + for s, e, _ordered in runs_local: + for a, b, _key in merge_plan(cost[s:e], spec.ctx_len, spec, rng): + breaks = [rng.random() < spec.para_break for _ in range(b - a - 1)] + r = rng.random() + kind = "mono" if r < spec.mono_share else ( + "zh_en" if r < spec.mono_share + (1 - spec.mono_share) * spec.zh_en_share else "en_zh") + plan.append((kind, _join(en[a + s:b + s], " ", breaks), _join(zh[a + s:b + s], "", breaks))) + if not plan: + return [], {} + enc_en = [x.ids for x in tok.encode_batch([p[1] for p in plan], add_special_tokens=False)] + enc_zh = [x.ids for x in tok.encode_batch([p[2] for p in plan], add_special_tokens=False)] + out, counts = [], {"mono": 0, "en_zh": 0, "zh_en": 0} + for (kind, _e, _z), e_ids, z_ids in zip(plan, enc_en, enc_zh): + counts[kind] += 1 + if kind == "mono": + out.extend(z_ids); out.append(eot); out.extend(e_ids); out.append(eot) + elif kind == "en_zh": + out.append(en_tag); out.extend(e_ids); out.append(zh_tag); out.extend(z_ids); out.append(eot) + else: + out.append(zh_tag); out.extend(z_ids); out.append(en_tag); out.extend(e_ids); out.append(eot) + return out, counts + + +def _build_part(args) -> dict: + tokenizer_dir, part_path, runs, out_dir, split, spec_d, seed = args + import pyarrow.parquet as pq + from tokenizers import Tokenizer + spec = DocMixSpec(**spec_d) + tok = Tokenizer.from_file(os.path.join(tokenizer_dir, "tokenizer.json")) + dtype = _dtype_for(tok.get_vocab_size()) + special = tuple(tok.token_to_id(t) for t in ("<|endoftext|>", "<|en|>", "<|zh|>")) + table = pq.read_table(part_path, columns=["zh", "en"]) + zh_all, en_all = table["zh"].to_pylist(), table["en"].to_pylist() + en, zh, local = [], [], [] + for s, e, o in runs: # gather selected rows; runs → local spans + local.append((len(en), len(en) + (e - s), o)) + en.extend(en_all[s:e]); zh.extend(zh_all[s:e]) + del table, zh_all, en_all + en_len = [len(x.ids) for x in tok.encode_batch(en, add_special_tokens=False)] + zh_len = [len(x.ids) for x in tok.encode_batch(zh, add_special_tokens=False)] + # Chat-template overhead of one record (+8 margin: joins can re-tokenize). + overhead = 8 + max(len(tok.encode(f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n<|im_end|>\n", + add_special_tokens=False).ids) for p in SFT_PROMPTS.values()) + stem = os.path.basename(part_path)[:-len(".parquet")] + p_ids, r_ids, meta = _sft_samples(tok, en, zh, en_len, zh_len, local, spec, random.Random(seed), overhead) + tokens, offsets, plens, dirs, ns = [], [0], [], [], [] + st = {"records": 0, "skipped_too_long": 0, "by_size": {}, "by_dir": {d: 0 for d in DIRECTIONS}, + "by_len": {}, "tokens_by_len": {}, "ordered_records": 0, "random_concat_records": 0, + "sft_tokens": 0, "sft_supervised_tokens": 0, "max_len": 0} + for p, r, (d, n, key, ordered) in zip(p_ids, r_ids, meta): + total = len(p) + len(r) + if total > spec.ctx_len: + st["skipped_too_long"] += 1 + continue + tokens.extend(p); tokens.extend(r) + offsets.append(offsets[-1] + total); plens.append(len(p)); dirs.append(d); ns.append(min(n, 255)) + st["records"] += 1 + st["by_size"][key] = st["by_size"].get(key, 0) + 1 + st["by_dir"][DIRECTIONS[d]] += 1 + b = _len_bin(total) + st["by_len"][b] = st["by_len"].get(b, 0) + 1 + st["tokens_by_len"][b] = st["tokens_by_len"].get(b, 0) + total + if n > 1: + st["ordered_records" if ordered else "random_concat_records"] += 1 + st["sft_tokens"] += total + st["sft_supervised_tokens"] += len(r) + st["max_len"] = max(st["max_len"], total) + sft_dir = os.path.join(out_dir, "sft", split) + _write_arrays(os.path.join(sft_dir, stem), { + ".bin": np.asarray(tokens, dtype), ".idx": np.asarray(offsets, np.int64), + ".n": np.asarray(ns, np.uint8), ".dir": np.asarray(dirs, np.uint8), ".plen": np.asarray(plens, np.int32)}) + pre, counts = _pretrain_docs(tok, en, zh, en_len, zh_len, local, spec, random.Random(seed + 1), special) + _write_arrays(os.path.join(out_dir, "pretrain", split, stem), {".bin": np.asarray(pre, dtype)}) + return {"part": stem, "split": split, "rows": len(en), "en_tokens": int(sum(en_len)), + "zh_tokens": int(sum(zh_len)), "pretrain_tokens": len(pre), "pretrain_docs": counts, **st} + + +def _merge_stats(results: list) -> dict: + out = {} + for r in results: + acc = out.setdefault(r["split"], {}) + for k, v in r.items(): + if k in ("part", "split"): + continue + if isinstance(v, dict): + d = acc.setdefault(k, {}) + for kk, vv in v.items(): + d[kk] = d.get(kk, 0) + vv + elif k == "max_len": + acc[k] = max(acc.get(k, 0), v) + else: + acc[k] = acc.get(k, 0) + v + for acc in out.values(): + recs = max(acc.get("records", 0), 1) + acc["size_share_records"] = {k: round(acc.get("by_size", {}).get(k, 0) / recs, 4) for k in SIZE_KEYS} + acc["dir_share_records"] = {k: round(v / recs, 4) for k, v in acc.get("by_dir", {}).items()} + tot = max(acc.get("sft_tokens", 0), 1) + acc["token_share_by_len"] = {k: round(v / tot, 4) for k, v in sorted( + acc.get("tokens_by_len", {}).items(), key=lambda kv: int(kv[0].lstrip(">").split("-")[0]))} + return out + + +def build_docmix(pairs_dir: str, tokenizer_dir: str, out_dir: str, spec: DocMixSpec, *, scan_dir: str, + workers: int = 8, log=print) -> dict: + """Select, merge and tokenize one doc-mix subset into ``out_dir``.""" + t0 = time.time() + scan = scan_corpus(pairs_dir, scan_dir, block_rows=spec.block_rows, workers=workers, log=log) + sel = select_chunks(scan, spec) + for kind in ("pretrain", "sft"): + for split in ("train", "valid"): + os.makedirs(os.path.join(out_dir, kind, split), exist_ok=True) + spec_d = asdict(spec) + jobs = [] + for i, (name, runs) in enumerate(sorted(sel["runs"].items())): + jobs.append((tokenizer_dir, os.path.join(pairs_dir, name), runs, out_dir, "train", spec_d, spec.seed + 101 * i)) + for i, name in enumerate(n for n in scan["names"] if n.startswith("valid.")): + rows = scan["rows"][name] + if rows: + ordered = bool(block_ordered(scan["blocks"][name], spec.ordered_ratio, spec.ordered_min_adj).all()) + jobs.append((tokenizer_dir, os.path.join(pairs_dir, name), [[0, rows, ordered]], out_dir, "valid", + spec_d, spec.seed + 7 + 101 * i)) + log(f"[docmix] {len(jobs)} parts, {sel['stats']['rows_selected']:,} train rows selected") + with ProcessPoolExecutor(max_workers=workers) as pool: + results = list(pool.map(_build_part, jobs)) + with open(os.path.join(tokenizer_dir, "tokenizer_report.json"), encoding="utf-8") as fh: + tok_meta = json.load(fh) + manifest = {"created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "elapsed_s": round(time.time() - t0, 1), "format": "chronos-docmix-v1", + "tokenizer_fingerprint": tok_meta.get("fingerprint"), "spec": spec_d, + "prompts": SFT_PROMPTS, "directions": list(DIRECTIONS), + "selection": sel["stats"], "splits": _merge_stats(results), + "parts": sorted(results, key=lambda r: (r["split"], r["part"]))} + with open(os.path.join(out_dir, "docmix_manifest.json"), "w", encoding="utf-8") as fh: + json.dump(manifest, fh, ensure_ascii=False, indent=1) + log(f"[docmix] done in {time.time() - t0:.0f}s → {out_dir}") + return manifest + + +def main(argv=None) -> int: + import argparse + p = argparse.ArgumentParser(description="Science-weighted, merged (n-in-1), bidirectional translation shards") + p.add_argument("--pairs", required=True) + p.add_argument("--tokenizer", required=True) + p.add_argument("--out", required=True) + p.add_argument("--scan-dir", required=True, help="cache for per-row science hits and block adjacency") + p.add_argument("--fraction", type=float, required=True) + p.add_argument("--ctx", type=int, required=True) + p.add_argument("--seed", type=int, default=DocMixSpec.seed) + p.add_argument("--workers", type=int, default=8) + args = p.parse_args(argv) + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + spec = DocMixSpec(fraction=args.fraction, ctx_len=args.ctx, seed=args.seed) + m = build_docmix(args.pairs, args.tokenizer, args.out, spec, scan_dir=args.scan_dir, workers=args.workers, + log=lambda s: print(s, flush=True)) + print(json.dumps({"selection": m["selection"], "train": {k: v for k, v in m["splits"].get("train", {}).items() + if not isinstance(v, dict) or "share" in k}}, + ensure_ascii=False), flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +class DocSFTRecords: + """SFT batches that give every record the same expected passes, at any length. + + ``shards.SFTRecords`` picks the batch's target record uniformly; a batch + holds ~budget / width records, so each short record is seen ~W_long/W_short + times as often as a long one. Here a record is the target with weight + 1 / rows(its length), where rows = records that one batch of that length + holds (budget // padded width, capped at ``max_rows``): every length then + yields records ∝ its count, i.e. the same expected passes per record for + short sentences and 4k-token documents, floor and cap effects included. + Reads ``.dir`` / ``.n`` and counts what was drawn per length bin + (``seen``; stored totals in ``bin_tokens``). + """ + + def __init__(self, folder: str, dtype=np.uint16, seed: int = 0, pad_id: int = 0, tol_div: int = 16): + stems = sorted(f[:-5] for f in os.listdir(folder) if f.endswith(".plen")) + self.tokens, self.offsets, self.plens, rows = [], [], [], [] + for si, stem in enumerate(stems): + base = os.path.join(folder, stem) + off = np.fromfile(base + ".idx", dtype=np.int64) + plen = np.fromfile(base + ".plen", dtype=np.int32) + dirs = np.fromfile(base + ".dir", dtype=np.uint8) + ns = np.fromfile(base + ".n", dtype=np.uint8) + if not (len(off) == len(plen) + 1 == len(dirs) + 1 == len(ns) + 1): + raise ValueError(f"{base}: .idx/.plen/.dir/.n lengths disagree") + if not len(plen): + continue + self.tokens.append(np.memmap(base + ".bin", dtype=dtype, mode="r")) + self.offsets.append(off); self.plens.append(plen) + si = len(self.tokens) - 1 + rows.append(np.stack([np.full(len(plen), si), np.arange(len(plen)), np.diff(off), + dirs.astype(np.int64), ns.astype(np.int64)], axis=1)) + idx = np.concatenate(rows) if rows else np.zeros((0, 5), np.int64) + self.index = idx[np.argsort(idx[:, 2], kind="stable")] + self.lengths = self.index[:, 2] + self.pad_id, self.tol_div = pad_id, tol_div + self._cum_cache: dict = {} + self._total = int(self.lengths.sum()) + # Stored tokens per record-length bin (same bins as ``seen["len_bins"]``). + bins = np.searchsorted(np.asarray(LEN_BINS[1:]), self.lengths, side="left") + self.bin_tokens = {_len_bin(int(LEN_BINS[bi + 1])): int(self.lengths[bins == bi].sum()) + for bi in np.unique(bins)} + self.rng = np.random.default_rng(seed) + self.seen = {"records": 0, "tokens": 0, "dir": [0, 0], "len_bins": {}} + self.last_meta = None + + def __len__(self): + return len(self.index) + + def _bucket(self, n, width_multiple: int = 64): + """Padded width of a batch whose target length is ``n`` (as in ``batch``).""" + n = np.asarray(n, dtype=np.int64) + tol = np.maximum(4, n // self.tol_div) + return -(-(n + tol) // width_multiple) * width_multiple + + def draw_cum(self, token_budget: int, max_rows: int = 256, width_multiple: int = 64) -> np.ndarray: + """Cumulative target weights 1 / rows(length) for one batch shape (cached).""" + key = (int(token_budget), int(max_rows), int(width_multiple)) + cum = self._cum_cache.get(key) + if cum is None: + rows = np.clip(int(token_budget) // self._bucket(self.lengths, width_multiple), 1, int(max_rows)) + cum = self._cum_cache[key] = np.cumsum(1.0 / rows) + return cum + + def total_tokens(self) -> int: + return self._total + + def supervised_tokens(self) -> int: + return int(sum(int(o[-1]) - int(p.sum()) for o, p in zip(self.offsets, self.plens))) + + def batch(self, token_budget: int, max_rows: int = 256, width_multiple: int = 64): + cum = self.draw_cum(token_budget, max_rows, width_multiple) + i = min(int(np.searchsorted(cum, self.rng.random() * cum[-1], side="right")), len(cum) - 1) + target = int(self.lengths[i]) + tol = max(4, target // self.tol_div) + 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_multiple) * width_multiple + 1 + 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, _d, _k) 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:] + self.last_meta = pick[:, 2:5].copy() # length, direction, sentences merged + s = self.seen + s["records"] += len(pick); s["tokens"] += int(pick[:, 2].sum()) + for d in (0, 1): + s["dir"][d] += int((pick[:, 3] == d).sum()) + for n in pick[:, 2]: # tokens drawn per record-length bin + b = _len_bin(int(n)) + s["len_bins"][b] = s["len_bins"].get(b, 0) + int(n) + return ids, labels diff --git a/chronos/data/corpus/download.py b/demos/translation/chronos_translate/corpus/download.py similarity index 100% rename from chronos/data/corpus/download.py rename to demos/translation/chronos_translate/corpus/download.py diff --git a/demos/translation/chronos_translate/corpus/evalsets.py b/demos/translation/chronos_translate/corpus/evalsets.py new file mode 100644 index 0000000..ed3ab30 --- /dev/null +++ b/demos/translation/chronos_translate/corpus/evalsets.py @@ -0,0 +1,120 @@ +"""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 + + +# ---- document-level sets (doc-mix runs) ------------------------------------ +def _flores_docs(raw: str, split: str): + """FLORES+ sentences grouped by source article (URL), in id order.""" + 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") + rows = pq.read_table(path, columns=["id", "URL", "sentence_eng_Latn", "sentence_zho_Hans"]).to_pylist() + docs: dict = {} + for r in sorted(rows, key=lambda r: int(r["id"])): + docs.setdefault(r["URL"], []).append((r["sentence_eng_Latn"], r["sentence_zho_Hans"])) + return [{"doc_id": url, "segments": segs} for url, segs in docs.items()] + + +def _wmt24pp_docs(raw: str): + """WMT24++ en→zh_CN segments grouped by document_id, in segment order (bad sources dropped).""" + path = os.path.join(raw, "google__wmt24pp__en-zh_CN", "en-zh_CN.jsonl") + docs: dict = {} + 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 + docs.setdefault(r["document_id"], []).append((int(r["segment_id"]), r["source"], r["target"], + r.get("domain"))) + out = [] + for doc_id, segs in docs.items(): + segs.sort() + out.append({"doc_id": doc_id, "domain": segs[0][3], "segments": [(s, t) for _i, s, t, _d in segs]}) + return out + + +EVAL_DOCUMENTS = { + "flores_dev_docs": lambda raw: _flores_docs(raw, "dev"), + "flores_devtest_docs": lambda raw: _flores_docs(raw, "devtest"), + "wmt24pp_docs": _wmt24pp_docs, +} + + +def load_eval_documents(raw_root: str, names=None) -> dict[str, list[dict]]: + """Return {name: [{"doc_id", "segments": [(en, zh), ...]}, ...]} for every available set. + + A document's source is its English segments joined by " ", its reference + the Chinese segments joined without separator (same joins as doc-mix). + """ + out = {} + for name, loader in EVAL_DOCUMENTS.items(): + if names and name not in names: + continue + try: + out[name] = loader(raw_root) + except FileNotFoundError: + continue + return out diff --git a/chronos/data/corpus/pipeline.py b/demos/translation/chronos_translate/corpus/pipeline.py similarity index 100% rename from chronos/data/corpus/pipeline.py rename to demos/translation/chronos_translate/corpus/pipeline.py diff --git a/chronos/data/corpus/presets.py b/demos/translation/chronos_translate/corpus/presets.py similarity index 97% rename from chronos/data/corpus/presets.py rename to demos/translation/chronos_translate/corpus/presets.py index 755cfd2..b7e961e 100644 --- a/chronos/data/corpus/presets.py +++ b/demos/translation/chronos_translate/corpus/presets.py @@ -1,4 +1,4 @@ -"""Corpus presets and the ``python -m chronos.data.corpus.presets`` CLI. +"""Corpus presets and the ``python -m chronos_translate.corpus.presets`` CLI. A preset is only a list of ``PairSource`` objects plus evaluation sets; the pipeline itself is task-agnostic. diff --git a/chronos/data/corpus/readers.py b/demos/translation/chronos_translate/corpus/readers.py similarity index 100% rename from chronos/data/corpus/readers.py rename to demos/translation/chronos_translate/corpus/readers.py diff --git a/chronos/data/corpus/shards.py b/demos/translation/chronos_translate/corpus/shards.py similarity index 67% rename from chronos/data/corpus/shards.py rename to demos/translation/chronos_translate/corpus/shards.py index b18a4db..7c411e6 100644 --- a/chronos/data/corpus/shards.py +++ b/demos/translation/chronos_translate/corpus/shards.py @@ -7,6 +7,7 @@ * ``<|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. +The layout and its readers are generic: ``chronos.data.shards``. """ from __future__ import annotations @@ -18,11 +19,14 @@ import numpy as np -SFT_PROMPT_EN_ZH = "翻译成中文，只输出译文：\n" - +# Readers live in the kernel; re-exported so existing imports keep working. +from chronos.data.shards import PackedTokenStream, SFTRecords, _dtype_for # noqa: F401 -def _dtype_for(vocab_size: int): - return np.uint16 if vocab_size <= 65535 else np.uint32 +SFT_PROMPT_EN_ZH = "翻译成中文，只输出译文：\n" +# Reverse direction (doc-mix shards, ``chronos_translate.corpus.docmix``): same +# template, so one model serves both directions from the instruction alone. +SFT_PROMPT_ZH_EN = "翻译成英文，只输出译文：\n" +SFT_PROMPTS = {"en_zh": SFT_PROMPT_EN_ZH, "zh_en": SFT_PROMPT_ZH_EN} def _encode_batch(tok, texts: list[str]) -> list[list[int]]: @@ -142,80 +146,6 @@ def build_shards(pairs_dir: str, tokenizer_dir: str, out_dir: str, *, workers: i 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") diff --git a/chronos/data/corpus/sources.py b/demos/translation/chronos_translate/corpus/sources.py similarity index 100% rename from chronos/data/corpus/sources.py rename to demos/translation/chronos_translate/corpus/sources.py diff --git a/demos/translation/chronos_translate/doc_run.py b/demos/translation/chronos_translate/doc_run.py new file mode 100644 index 0000000..0c452f0 --- /dev/null +++ b/demos/translation/chronos_translate/doc_run.py @@ -0,0 +1,365 @@ +"""Time-budgeted MLX run on doc-mix shards (``chronos_translate.corpus.docmix``). + +Same loop, checkpoints, telemetry and resume as the kernel ``LongRun`` (via +``TranslationLongRun``); this subclass changes only what the doc-mix data needs: + +* SFT batches come from ``DocSFTRecords`` (target record drawn ∝ length, so + short sentences and 4k-token documents get the same passes per record) + and carry both directions (``.dir``: 0 = en→zh, 1 = zh→en). +* Translation eval covers both directions at sentence level (FLORES+ + devtest) and document level (FLORES+ devtest articles, WMT24++ documents + that fit the context). ``bleu_zh`` / ``chrf`` keep the en→zh sentence + scores under the live run's key names, so monitors and comparisons keep working. +* ``micro_tokens`` is fitted at launch to the GPU's recommended working set + (largest power-of-two fraction that runs one full-length step without an + allocation failure and under the memory target); ``tokens_per_step`` is + kept by raising gradient accumulation, so the optimisation recipe is unchanged. + +Paths in the recipe may be relative to the recipe file, so a package can be +copied to another machine and launched in place. +""" +from __future__ import annotations + +import json +import math +import os +import time + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from chronos.mlx.training.longrun import PhaseSpec, PretrainSource, SFTSource +from chronos.mlx.training.longrun import load_spec as _load_spec +from chronos.mlx.training.longrun import main as _longrun_main + +from .run import EXTRA_PATH_KEYS, TranslationLongRun, TranslationRunSpec + +DIR_NAMES = ("en_zh", "zh_en") + + +class _PhaseDone(Exception): + """Raised from ``_emit`` when a phase reaches ``extra.target_passes``.""" + + +class DocSFTSource(SFTSource): + """``SFTSource`` over doc-mix shards (records ∝ length, both directions).""" + + def __init__(self, shards_dir, seed): # noqa: D107 - same interface as SFTSource + from chronos_translate.corpus.docmix import DocSFTRecords + self.records = DocSFTRecords(os.path.join(shards_dir, "sft", "train"), seed=seed) + self.valid = DocSFTRecords(os.path.join(shards_dir, "sft", "valid"), seed=seed + 1) + + def seen(self) -> dict: + s = self.records.seen + n = max(1, s["records"]) + tot = max(1, s["tokens"]) + return {"sft_drawn_records": s["records"], "sft_drawn_tokens": s["tokens"], + "sft_drawn_zh_en_share": round(s["dir"][1] / n, 4), + "sft_drawn_len_share": {k: round(v / tot, 4) for k, v in sorted(s["len_bins"].items())}, + "sft_passes": round(s["tokens"] / max(1, self.records.total_tokens()), 4)} + + +def _docmix_manifest(shards_dir: str) -> dict: + path = os.path.join(shards_dir, "docmix_manifest.json") + if not os.path.isfile(path): + raise FileNotFoundError(f"{shards_dir} is not a doc-mix shard folder (no docmix_manifest.json)") + with open(path, encoding="utf-8") as fh: + return json.load(fh) + + +def _gpu_info() -> dict: + """Metal device info (unified memory: the GPU working set bounds CPU-stream runs too).""" + if not hasattr(mx, "device_info"): + return {} + try: + return dict(mx.device_info(mx.gpu)) + except Exception: # noqa: BLE001 - no Metal device + return dict(mx.device_info()) + + +class DocLongRun(TranslationLongRun): + def __init__(self, spec: TranslationRunSpec, log=print): + self.docmix = _docmix_manifest(spec.shards_dir) + self.ctx_len = int(self.docmix["spec"]["ctx_len"]) + if int(spec.model_config.get("max_position_embeddings", self.ctx_len)) < self.ctx_len: + raise ValueError("model_config.max_position_embeddings is below the shards' ctx_len") + super().__init__(spec, log=log) + self._data_tokens: dict = {} + self._projected: set = set() + + # ---- data ---------------------------------------------------------------- + def _source(self, phase: PhaseSpec): + if phase.name == "pretrain": + src = PretrainSource(self.spec.shards_dir, phase.seq_len, self.spec.seed) + self._data_tokens[phase.name] = int(src.stream.total_tokens) + else: + src = DocSFTSource(self.spec.shards_dir, self.spec.seed) + self._data_tokens[phase.name] = int(src.records.supervised_tokens()) + return src + + # ---- memory fit ------------------------------------------------------------ + def _probe_peak(self, grad_fn, rows: int, seq: int) -> float: + """Peak bytes of two accumulated micro-steps on ``rows`` full-length rows.""" + from chronos.mlx.training.trainer import _add_grads + rng = np.random.default_rng(0) + ids = mx.array(rng.integers(8, int(self.config.vocab_size), size=(rows, seq)), dtype=mx.int32) + mx.clear_cache() + mx.reset_peak_memory() + _, g1 = grad_fn(self.model, ids, ids) + mx.eval(g1) + _, g2 = grad_fn(self.model, ids, ids) + acc = _add_grads(g1, g2) + mx.eval(acc) + peak = float(mx.get_peak_memory()) + del g1, g2, acc + mx.clear_cache() + return peak + + def _fit_micro_tokens(self, phase: PhaseSpec) -> int: + """Largest ``micro_tokens`` ≤ the recipe's whose step fits the memory target. + + Peak memory is measured at 1 and 2 full-context rows and extrapolated + linearly (activations scale with rows); the chosen size is then run + once to confirm. The result is stored in the run state, so a resume + keeps the same micro-batches. + """ + fitted = self.state.setdefault("fitted_micro_tokens", {}) + if phase.name in fitted: + return int(fitted[phase.name]) + extra = self.spec.extra or {} + seq = max(int(phase.seq_len), self.ctx_len) if phase.name != "pretrain" else int(phase.seq_len) + want_rows = max(1, int(phase.micro_tokens) // seq) + info = _gpu_info() + limit = float(info.get("max_recommended_working_set_size") or info.get("memory_size") or 0) + if extra.get("micro_fit", "auto") == "off" or not limit: + fitted[phase.name] = int(phase.micro_tokens) + return int(phase.micro_tokens) + target = float(extra.get("micro_fit_fraction", 0.7)) * limit + grad_fn = nn.value_and_grad(self.model, self._loss_fn) + p1 = self._probe_peak(grad_fn, 1, seq) + p2 = self._probe_peak(grad_fn, 2, seq) if want_rows >= 2 else p1 + per_row = max(p2 - p1, 1.0) + rows = 1 + while rows * 2 <= want_rows and p1 + (rows * 2 - 1) * per_row <= target: + rows *= 2 + peak = self._probe_peak(grad_fn, rows, seq) if rows > 2 else (p2 if rows == 2 else p1) + while rows > 1 and peak > target: + rows //= 2 + peak = self._probe_peak(grad_fn, rows, seq) + micro = rows * seq if rows < want_rows else int(phase.micro_tokens) + fitted[phase.name] = micro + report = {"phase": phase.name, "seq": seq, "rows": rows, "micro_tokens": micro, + "recipe_micro_tokens": int(phase.micro_tokens), "peak_gb": round(peak / 1e9, 2), + "target_gb": round(target / 1e9, 2), "limit_gb": round(limit / 1e9, 2), + "device": info.get("device_name") or info.get("architecture")} + self.state.setdefault("micro_fit", {})[phase.name] = report + self.log(f"[doc] micro-batch fit: {json.dumps(report, ensure_ascii=False)}") + if peak > target: + self.log(f"[doc] warning: one {seq}-token row needs {peak / 1e9:.1f} GB, above the " + f"{target / 1e9:.1f} GB target; training continues at 1 row per micro-batch") + return micro + + # ---- training ------------------------------------------------------------ + def run_phase(self, pi: int, deadline_s: float, max_steps: int | None = None): + from dataclasses import replace + original = self.spec.phases[pi] + micro = self._fit_micro_tokens(original) + phase = replace(original, micro_tokens=micro) + accum = max(1, phase.tokens_per_step // micro) + data = self._data_tokens.get(phase.name) + self.log(f"[doc] {phase.name}: micro_tokens={micro} × accum {accum} = {micro * accum} tokens/step " + f"(recipe {original.tokens_per_step}); ctx={self.ctx_len}") + if self.state["phase_step"] == 0: + self.state.setdefault("phase_start_tokens", {})[phase.name] = int(self.state["tokens_seen"]) + self._phase_name = phase.name + self.spec.phases[pi] = phase + try: + super().run_phase(pi, deadline_s, max_steps=max_steps) + except _PhaseDone: + # Same closing steps as LongRun.run_phase after its loop. + source = self._active_source + self.log(f"[doc] {phase.name}: reached {self._passes(phase.name):.2f} passes " + f"(target {self._target(phase.name)}) at step {self.state['phase_step']}") + 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"], "reason": "target_passes"}) + finally: + self.spec.phases[pi] = original + self._phase_name = None + if data and max_steps is None: + self.log(f"[doc] {phase.name}: {self._passes(phase.name):.2f} passes over {data:,} " + f"{'tokens' if phase.name == 'pretrain' else 'supervised tokens'}") + + def _target(self, name: str) -> float | None: + t = ((self.spec.extra or {}).get("target_passes") or {}).get(name) + return float(t) if t else None + + def _emit(self, record: dict): + """Pass accounting on training rows; ends the phase at ``target_passes``. + + Once the cumulative rate (tokens per wall second, evals included) is + known, the time schedule's phase budget shrinks to the predicted time of + the target, so warmup + cosine end where the data passes do. Time left + over moves to the next phase (its deadline is fixed by the recipe). + """ + name = getattr(self, "_phase_name", None) + target = self._target(name) if name and "event" not in record and record.get("phase") == name else None + done = False + if target and self._data_tokens.get(name): + st = self.state + data = self._data_tokens[name] + seen = int(st["tokens_seen"]) - int((st.get("phase_start_tokens") or {}).get(name, 0)) + wall = float(st.get("phase_wall_s") or 0.0) + done = seen >= target * data + if wall >= 60 and seen > 0 and self.spec.schedule == "time" and st.get("phase_budget_s"): + need = wall * target * data / seen + orig = st.setdefault("phase_budget_orig_s", {}).setdefault(name, float(st["phase_budget_s"])) + st["phase_budget_s"] = max(wall, min(orig, need)) + step = int(record.get("step") or st["phase_step"]) + st["phase_total"] = record["total"] = max(step, int(step * st["phase_budget_s"] / wall)) + record["passes"] = round(seen / data, 4) + if name not in self._projected and wall >= 600: + self._projected.add(name) + rate = seen / wall + budget = (st.get("phase_budget_orig_s") or {}).get(name, st.get("phase_budget_s") or wall) + self.log(f"[doc] {name}: {rate:,.0f} tok/s over {wall / 60:.0f} min → " + f"{rate * budget / data:.2f} passes fit the phase budget; target {target} " + f"({'reached early' if rate * budget >= target * data else 'not reached in time'})") + super()._emit(record) + if done: + raise _PhaseDone + + def _passes(self, name: str) -> float: + start = (self.state.get("phase_start_tokens") or {}).get(name, 0) + return (int(self.state["tokens_seen"]) - int(start)) / max(1, self._data_tokens.get(name) or 1) + + def _periodic_eval(self, source, phase): + rec = {"phase": phase.name, "step": self.state["phase_step"], "event": "data_passes", + "passes": round(self._passes(phase.name), 4)} + if isinstance(source, DocSFTSource): + rec.update(source.seen()) + self._emit(rec) + super()._periodic_eval(source, phase) + + # ---- translation --------------------------------------------------------- + def translate(self, sources: list[str], max_new_tokens: int | None = None, + direction: str = "en_zh") -> list[str]: + """Greedy translation with the doc-mix prompt for ``direction``. + + The output limit scales with the source (zh output ≲ 1.6× en tokens, + en ≲ 2× zh tokens) and is capped by the context left after the prompt. + """ + from chronos_translate.corpus.shards import SFT_PROMPTS + tok = self._tokenizer() + end = tok.token_to_id("<|im_end|>") + ratio = 1.6 if direction == "en_zh" else 2.0 + self.model.eval() + out = [] + try: + for src in sources: + prompt = f"<|im_start|>user\n{SFT_PROMPTS[direction]}{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) + room = self.ctx_len - len(ids) + limit = min(room, int(ratio * src_len) + 16) + if max_new_tokens: + limit = min(limit, int(max_new_tokens)) + out.append(tok.decode(self.generate(ids, max(1, limit), (end,)), skip_special_tokens=True) + if room > 0 else "") + finally: + self.model.train() + return out + + def _score(self, hyps, refs, direction: str) -> tuple[float, float]: + import sacrebleu + bleu = sacrebleu.corpus_bleu(hyps, [refs], tokenize="zh" if direction == "en_zh" else "13a") + return round(bleu.score, 2), round(sacrebleu.corpus_chrf(hyps, [refs]).score, 2) + + def _fits(self, src: str, ref: str, direction: str) -> bool: + from chronos_translate.corpus.shards import SFT_PROMPTS + tok = self._tokenizer() + n = len(tok.encode(SFT_PROMPTS[direction] + src, add_special_tokens=False).ids) + return n + len(tok.encode(ref, add_special_tokens=False).ids) + 12 <= self.ctx_len + + def translation_eval(self, n: int) -> dict: + """Both directions on FLORES+ devtest sentences and on whole documents. + + Documents are translated in one pass each (source segments joined the + way the doc-mix records join them) and scored at document level; + ``*_len_ratio`` = hypothesis / reference characters flags early stops + and run-on output on long inputs. + """ + from chronos_translate.corpus.evalsets import load_eval_documents, load_eval_sets + raw = (self.spec.extra or {}).get("raw_root") + if not raw: + return {} + t0 = time.monotonic() + out, samples = {}, [] + pairs = load_eval_sets(raw, ["flores_devtest"]).get("flores_devtest", [])[:n] + for direction in DIR_NAMES: + if not pairs: + break + srcs = [en if direction == "en_zh" else zh for en, zh in pairs] + refs = [zh if direction == "en_zh" else en for en, zh in pairs] + hyps = self.translate(srcs, direction=direction) + bleu, chrf = self._score(hyps, refs, direction) + if direction == "en_zh": + out.update({"flores_devtest_n": len(pairs), "bleu_zh": bleu, "chrf": chrf}) + else: + out.update({"sent_zh_en_bleu": bleu, "sent_zh_en_chrf": chrf}) + samples += [{"direction": direction, "set": "flores_devtest", "src": s, "ref": r, "hyp": h} + for s, r, h in list(zip(srcs, refs, hyps))[:3]] + extra = self.spec.extra or {} + final = n >= int(self.spec.final_eval_samples or 0) > 0 + # Periodic evals use a fixed small subset; the final eval every document + # that fits the context unless the recipe caps it. + n_docs = int(extra.get("final_doc_eval_docs", 1 << 30) if final else extra.get("doc_eval_docs", 16)) + docs = load_eval_documents(raw, ["flores_devtest_docs", "wmt24pp_docs"]) + plan = (("doc", "flores_devtest_docs", DIR_NAMES), ("wmt24pp_doc", "wmt24pp_docs", ("en_zh",))) + for prefix, name, directions in plan: + for direction in directions: + srcs, refs = [], [] + for d in docs.get(name, []): + en = " ".join(e for e, _z in d["segments"]) + zh = "".join(z for _e, z in d["segments"]) + src, ref = (en, zh) if direction == "en_zh" else (zh, en) + if self._fits(src, ref, direction): + srcs.append(src); refs.append(ref) + if len(srcs) >= n_docs: + break + if not srcs: + continue + hyps = self.translate(srcs, direction=direction) + bleu, chrf = self._score(hyps, refs, direction) + ratio = sum(map(len, hyps)) / max(1, sum(map(len, refs))) + out.update({f"{prefix}_{direction}_bleu": bleu, f"{prefix}_{direction}_chrf": chrf, + f"{prefix}_{direction}_len_ratio": round(ratio, 3), f"{prefix}_{direction}_n": len(srcs)}) + samples += [{"direction": direction, "set": name, "src": s, "ref": r, "hyp": h} + for s, r, h in list(zip(srcs, refs, hyps))[:2]] + self.write_samples(samples) + if out: + out["eval_s"] = round(time.monotonic() - t0, 1) + return out + + +def load_spec(path: str, budget_hours: float | None = None) -> TranslationRunSpec: + """Read a recipe; relative paths (and ``extra.raw_root``) resolve against its folder.""" + return _load_spec(path, budget_hours, spec_cls=TranslationRunSpec, extra_path_keys=EXTRA_PATH_KEYS) + + +def _smoke(spec, args) -> None: + if args.max_steps is not None: # smoke: keep the closing eval to a few documents + spec.extra = dict(spec.extra or {}, doc_eval_docs=2) + + +def main(argv=None) -> int: + return _longrun_main(argv, runner_cls=DocLongRun, spec_cls=TranslationRunSpec, extra_path_keys=EXTRA_PATH_KEYS, + description="Time-budgeted Chronos MLX run on doc-mix shards", prepare=_smoke) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/demos/translation/chronos_translate/model_presets.py b/demos/translation/chronos_translate/model_presets.py new file mode 100644 index 0000000..543eef9 --- /dev/null +++ b/demos/translation/chronos_translate/model_presets.py @@ -0,0 +1,37 @@ +"""The en→zh demo's model topology and tokenizer special tokens. + +``ENZH_50M_BASE`` is the fixed part of the 50M benchmark topology (moved out +of ``chronos.model.factory`` in 0.3.1a2, where it was ``ENZH_DEMO_BASE``). The +searched part (hidden size, expert count and width) comes from the kernel's +``fit_config``:: + + from chronos.model.factory import fit_config + fit = fit_config(50e6, 12e6, vocab_size=16384, base=ENZH_50M_BASE) + +The 72h recipe (``recipes/enzh_50m_72h.json``) stores the full ``model_config`` +it was trained with; that file, not this search, is the reference. + +``SPECIAL_TOKENS`` is the order the 0.3.1a1 tokenizer used (language tags in +ids 3-5), so a tokenizer trained with ``chronos tokenizer --out +--extra-special "<|en|>" "<|zh|>" "<|translate|>"`` matches the 50M run. +""" +from __future__ import annotations + +from chronos.tokenizer.model import special_tokens_with + +LANG_TOKENS = ("<|en|>", "<|zh|>", "<|translate|>") +SPECIAL_TOKENS = special_tokens_with(LANG_TOKENS) + +ENZH_50M_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, +) +ENZH_DEMO_BASE = ENZH_50M_BASE # pre-0.3.1a2 name (chronos.model.factory) diff --git a/demos/translation/chronos_translate/probe.py b/demos/translation/chronos_translate/probe.py new file mode 100644 index 0000000..bfd7221 --- /dev/null +++ b/demos/translation/chronos_translate/probe.py @@ -0,0 +1,78 @@ +"""Greedy en↔zh translation probe for a model loaded from a run (FLORES+ dev, BLEU / chrF). + +Pretrain checkpoints are probed in the ``<|en|>…<|zh|>`` tag format they +were trained on (both directions); SFT checkpoints use the chat prompt. Used +by ``TranslationLongRun`` at each pretrain validation and by +``chronos-translate inspect`` (offline, on the MLX CPU device). +""" +from __future__ import annotations + +import os +import time + +LANG_TAGS = ("<|en|>", "<|zh|>") + + +def translation_probe(spec: dict, model, phase_name: str, n: int) -> dict: + import sacrebleu + from tokenizers import Tokenizer + + from chronos.mlx.training.inspect_checkpoint import greedy + + from .corpus.evalsets import load_eval_sets + from .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|>", *LANG_TAGS)} + 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 = LANG_TAGS if name == "en_zh" else LANG_TAGS[::-1] + 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 main(argv=None) -> int: + """``chronos-translate inspect``: the kernel inspector plus this probe.""" + import argparse + import json + + from chronos.mlx.training.inspect_checkpoint import inspect_run + p = argparse.ArgumentParser(description="Inspect a run checkpoint (CPU) with the en↔zh FLORES+ dev probe") + 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) + report = inspect_run(a.run_dir, samples=a.samples, n_batches=a.batches, export_path=a.export, + device=a.device, probes={"translation": translation_probe}) + print(json.dumps(report, ensure_ascii=False, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/demos/translation/chronos_translate/run.py b/demos/translation/chronos_translate/run.py new file mode 100644 index 0000000..e8dfac1 --- /dev/null +++ b/demos/translation/chronos_translate/run.py @@ -0,0 +1,115 @@ +"""en↔zh translation run: the kernel ``LongRun`` plus translation metrics. + +* every SFT validation: greedy en→zh on a FLORES+ devtest subset + (``bleu_samples``), BLEU (``tokenize=zh``) and chrF; +* every pretrain validation: the ``<|en|>…<|zh|>`` probe on FLORES+ dev in + both directions (``probe_samples``, keys ``probe_{en_zh,zh_en}_{bleu,chrf}``); +* end of run: en→zh on ``final_eval_samples`` FLORES+ devtest sentences. + +FLORES+ is read from ``extra.raw_root`` (a ``chronos-translate corpus download`` +tree). The recipe format is the kernel's ``RunSpec`` plus the three sample +counts above, so 0.3.1a1 recipes run unchanged:: + + chronos-translate run --spec demos/translation/recipes/enzh_50m_72h.json +""" +from __future__ import annotations + +import time +from dataclasses import asdict, dataclass + +from chronos.mlx.training.longrun import LongRun, RunSpec +from chronos.mlx.training.longrun import main as _longrun_main + +#: ``extra`` keys holding paths (resolved against the recipe's folder). +EXTRA_PATH_KEYS = ("raw_root",) + + +@dataclass +class TranslationRunSpec(RunSpec): + final_eval_samples: int = 1012 # full FLORES+ devtest at the end (0 = skip) + export_prefix: str = "chronos_enzh" + bleu_samples: int = 200 # periodic FLORES+ devtest subset during SFT + # Greedy FLORES+ dev probe at each pretrain validation (``<|en|>…<|zh|>`` + # format, both directions); samples go to ``samples.jsonl``. 0 = off. + probe_samples: int = 16 + + +class TranslationLongRun(LongRun): + def pinned_modules(self) -> tuple[str, ...]: + return super().pinned_modules() + ("sacrebleu", "chronos_translate.corpus.evalsets", + "chronos_translate.corpus.shards", "chronos_translate.probe") + + def eval_samples_for(self, phase_name: str, *, final: bool) -> int: + if final: + return int(self.spec.final_eval_samples or 0) + if phase_name == "sft": + return int(getattr(self.spec, "bleu_samples", 0) or 0) + if phase_name == "pretrain": + return int(getattr(self.spec, "probe_samples", 0) or 0) + return super().eval_samples_for(phase_name, final=final) + + def task_eval(self, phase_name: str, n: int, *, final: bool = False) -> dict: + if final or phase_name != "pretrain": + return self.translation_eval(n) + return self._pretrain_probe(n) + + def _pretrain_probe(self, n: int) -> dict: + from .probe import translation_probe + probe = translation_probe(asdict(self.spec), self.model, "pretrain", n) + out, rows = {}, [] + 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"] + rows += [{"direction": direction, "set": probe.get("probe_set"), **sample} + for sample in res.get("samples", [])[:3]] + self.write_samples(rows) + if out: + out["probe_n"], out["probe_s"] = n, probe.get("probe_s") + return out + + # ---- translation quality (greedy, KV cache, MLX GPU) ---------------------- + def translate(self, sources: list[str], max_new_tokens: int = 160) -> list[str]: + from .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 + out.append(tok.decode(self.generate(ids, limit, (end,)), skip_special_tokens=True)) + self.model.train() + return out + + def translation_eval(self, n: int) -> dict: + import sacrebleu + + from .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]) + self.write_samples([{"direction": "en_zh", "set": "flores_devtest", "src": en, "ref": zh, "hyp": hyp} + for (en, zh), hyp in list(zip(pairs, hyps))[:5]]) + return {"flores_devtest_n": len(pairs), "bleu_zh": round(bleu.score, 2), "chrf": round(chrf.score, 2), + "eval_s": round(time.monotonic() - t0, 1)} + + +def main(argv=None) -> int: + return _longrun_main(argv, runner_cls=TranslationLongRun, spec_cls=TranslationRunSpec, + extra_path_keys=EXTRA_PATH_KEYS, description="Time-budgeted en↔zh Chronos MLX run") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/demos/translation/pyproject.toml b/demos/translation/pyproject.toml new file mode 100644 index 0000000..4f96781 --- /dev/null +++ b/demos/translation/pyproject.toml @@ -0,0 +1,28 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "chronos-translate" +version = "0.3.1a2" +description = "English-Chinese translation demo for the Chronos kernel: corpus pipeline, BLEU/chrF runs, model bundles" +readme = "README.md" +license = "Apache-2.0" +requires-python = ">=3.10" +dependencies = [ + "Project_Chronos==0.3.1a2", + "sacrebleu==2.6.0", + "sacremoses==0.2.0", + "pyarrow==23.0.1", +] + +[project.optional-dependencies] +mlx = ["Project_Chronos[mlx]==0.3.1a2"] +serve = ["Project_Chronos[serve]==0.3.1a2"] + +[project.scripts] +chronos-translate = "chronos_translate.cli:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["chronos_translate*"] diff --git a/demos/translation/recipes/docmix_plan.json b/demos/translation/recipes/docmix_plan.json new file mode 100644 index 0000000..6986b72 --- /dev/null +++ b/demos/translation/recipes/docmix_plan.json @@ -0,0 +1,614 @@ +{ + "A": { + "model": { + "architecture_version": 3, + "hidden_size": 384, + "num_hidden_layers": 5, + "vocab_size": 16384, + "num_attention_heads": 6, + "num_key_value_heads": 2, + "max_position_embeddings": 4096, + "tie_word_embeddings": true, + "dropout": 0.0, + "recurrent_enabled": true, + "recurrent_groups": 3, + "recurrent_depth": 4, + "min_recurrent_depth": 2, + "prelude_layers": 1, + "coda_layers": 1, + "adaptive_halting": true, + "halting_threshold": 0.95, + "max_inference_depth": 4, + "ponder_loss_coef": 0.003, + "recurrent_distill_coef": 0.05, + "recurrent_stability_coef": 0.001, + "recurrent_latent_refinement": false, + "recurrent_attention_mode": "shared_first", + "recurrent_strategy": "scheduled", + "depth_schedule": "fixed_then_adaptive", + "depth_schedule_warmup_fraction": 0.3, + "use_moe": true, + "num_experts": 18, + "num_experts_per_tok": 2, + "num_shared_experts": 1, + "moe_intermediate_size": 384, + "lookahead_steps": 2, + "layer_aware_lookahead": true, + "layer_aware_cache": true, + "lambda_balance": 0.005, + "lambda_temporal": 0.001, + "lambda_lookahead": 0.1, + "lambda_lookahead_topk": 0.05, + "lambda_cluster_locality": 0.3, + "cluster_locality_clusters": 0, + "cluster_locality_window": 128, + "lambda_router_anchor": 0.0, + "fallback_mask_prob": 0.08, + "use_hybrid_attention": true, + "kv_latent_dim": 128, + "rope_dim": 32, + "rope_layout": "per_head", + "sliding_window_size": 512, + "kv_cache_dtype": "auto", + "mla_cache_rope": true, + "router_scope": "per_iteration", + "ram_budget_mb": null, + "vram_budget_gb": 0.25, + "optimizer_paging": false, + "activation_checkpointing": false, + "extra_metadata": { + "stage_overrides": { + "pretrain": { + "learning_rate": 0.002, + "batch_size": 4, + "accumulation_steps": 8, + "max_seq_len": 4096, + "grad_clip": 1.0, + "epochs": 4, + "max_steps": 45288, + "weight_decay": 0.1, + "save_interval": 2260, + "lambda_router_anchor": 0.0 + }, + "sft": { + "learning_rate": 0.00067, + "batch_size": 4, + "accumulation_steps": 4, + "max_seq_len": 4096, + "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.0001, + "batch_size": 2, + "accumulation_steps": 4, + "max_seq_len": 4096, + "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.0002, + "batch_size": 2, + "accumulation_steps": 4, + "max_seq_len": 4096, + "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": 1, + "accumulation_steps": 1, + "max_seq_len": 4096, + "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.00067, + "batch_size": 4, + "accumulation_steps": 4, + "max_seq_len": 4096, + "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 + } + } + } + }, + "phases": [ + { + "name": "pretrain", + "budget_fraction": 0.5, + "peak_lr": 0.002, + "tokens_per_step": 131072, + "micro_tokens": 16384, + "seq_len": 4096, + "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.00067, + "tokens_per_step": 65536, + "micro_tokens": 16384, + "seq_len": 4096, + "warmup_fraction": 0.01, + "min_lr_ratio": 0.05, + "weight_decay": 0.1, + "grad_clip": 1.0 + } + ], + "arch": { + "total_params": 50328836, + "active_params": 14939396, + "active_union_upper": 22902020, + "params_source": "measured (meta-device count)", + "hidden_size": 384, + "physical_layers": 5, + "recurrent_groups": 3, + "recurrent_depth": 4, + "virtual_layers": 14, + "experts": 18, + "top_k": 2, + "expert_width": 384, + "heads": "6/2" + }, + "scaling": { + "chosen_lr": 0.002, + "chosen_batch_tokens": 131072 + }, + "memory": { + "level": "none", + "ram_gb": 128.0, + "target_gb": 76.8, + "micro_tokens": 16384, + "rows_per_micro": 4, + "accumulation_steps": 8, + "training_peak_gb": 13.59, + "training_peak_label": "estimated", + "routed_expert_mb_bf16": 75.9, + "ram_budget_mb": null, + "kv_cache_dtype": "auto", + "router_scope": "per_iteration", + "activation_checkpointing": false, + "optimizer_paging": false + } + }, + "B": { + "model": { + "architecture_version": 3, + "hidden_size": 384, + "num_hidden_layers": 5, + "vocab_size": 16384, + "num_attention_heads": 6, + "num_key_value_heads": 2, + "max_position_embeddings": 1024, + "tie_word_embeddings": true, + "dropout": 0.0, + "recurrent_enabled": true, + "recurrent_groups": 3, + "recurrent_depth": 4, + "min_recurrent_depth": 2, + "prelude_layers": 1, + "coda_layers": 1, + "adaptive_halting": true, + "halting_threshold": 0.95, + "max_inference_depth": 4, + "ponder_loss_coef": 0.003, + "recurrent_distill_coef": 0.05, + "recurrent_stability_coef": 0.001, + "recurrent_latent_refinement": false, + "recurrent_attention_mode": "shared_first", + "recurrent_strategy": "scheduled", + "depth_schedule": "fixed_then_adaptive", + "depth_schedule_warmup_fraction": 0.3, + "use_moe": true, + "num_experts": 18, + "num_experts_per_tok": 2, + "num_shared_experts": 1, + "moe_intermediate_size": 384, + "lookahead_steps": 2, + "layer_aware_lookahead": true, + "layer_aware_cache": true, + "lambda_balance": 0.005, + "lambda_temporal": 0.001, + "lambda_lookahead": 0.1, + "lambda_lookahead_topk": 0.05, + "lambda_cluster_locality": 0.3, + "cluster_locality_clusters": 0, + "cluster_locality_window": 128, + "lambda_router_anchor": 0.0, + "fallback_mask_prob": 0.08, + "use_hybrid_attention": true, + "kv_latent_dim": 128, + "rope_dim": 32, + "rope_layout": "per_head", + "sliding_window_size": 256, + "kv_cache_dtype": "auto", + "mla_cache_rope": true, + "router_scope": "per_iteration", + "ram_budget_mb": null, + "vram_budget_gb": 0.25, + "optimizer_paging": false, + "activation_checkpointing": false, + "extra_metadata": { + "stage_overrides": { + "pretrain": { + "learning_rate": 0.0025, + "batch_size": 16, + "accumulation_steps": 4, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 4, + "max_steps": 15256, + "weight_decay": 0.1, + "save_interval": 760, + "lambda_router_anchor": 0.0 + }, + "sft": { + "learning_rate": 0.00083, + "batch_size": 16, + "accumulation_steps": 2, + "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": 0.00013, + "batch_size": 8, + "accumulation_steps": 2, + "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.00025, + "batch_size": 8, + "accumulation_steps": 2, + "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": 2e-05, + "batch_size": 4, + "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": 8, + "temperature": 1.0, + "max_gen_len": 256, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "distill": { + "learning_rate": 0.00083, + "batch_size": 16, + "accumulation_steps": 2, + "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 + } + } + } + }, + "phases": [ + { + "name": "pretrain", + "budget_fraction": 0.5, + "peak_lr": 0.0025, + "tokens_per_step": 65536, + "micro_tokens": 16384, + "seq_len": 1024, + "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.00083, + "tokens_per_step": 32768, + "micro_tokens": 16384, + "seq_len": 1024, + "warmup_fraction": 0.01, + "min_lr_ratio": 0.05, + "weight_decay": 0.1, + "grad_clip": 1.0 + } + ], + "arch": { + "total_params": 50328836, + "active_params": 14939396, + "active_union_upper": 22902020, + "params_source": "measured (meta-device count)", + "hidden_size": 384, + "physical_layers": 5, + "recurrent_groups": 3, + "recurrent_depth": 4, + "virtual_layers": 14, + "experts": 18, + "top_k": 2, + "expert_width": 384, + "heads": "6/2" + }, + "scaling": { + "chosen_lr": 0.0025, + "chosen_batch_tokens": 65536 + }, + "memory": { + "level": "none", + "ram_gb": 128.0, + "target_gb": 76.8, + "micro_tokens": 16384, + "rows_per_micro": 16, + "accumulation_steps": 4, + "training_peak_gb": 13.59, + "training_peak_label": "estimated", + "routed_expert_mb_bf16": 75.9, + "ram_budget_mb": null, + "kv_cache_dtype": "auto", + "router_scope": "per_iteration", + "activation_checkpointing": false, + "optimizer_paging": false + } + }, + "C": { + "model": { + "architecture_version": 3, + "hidden_size": 384, + "num_hidden_layers": 5, + "vocab_size": 16384, + "num_attention_heads": 6, + "num_key_value_heads": 2, + "max_position_embeddings": 1024, + "tie_word_embeddings": true, + "dropout": 0.0, + "recurrent_enabled": true, + "recurrent_groups": 3, + "recurrent_depth": 4, + "min_recurrent_depth": 2, + "prelude_layers": 1, + "coda_layers": 1, + "adaptive_halting": true, + "halting_threshold": 0.95, + "max_inference_depth": 4, + "ponder_loss_coef": 0.003, + "recurrent_distill_coef": 0.05, + "recurrent_stability_coef": 0.001, + "recurrent_latent_refinement": false, + "recurrent_attention_mode": "shared_first", + "recurrent_strategy": "scheduled", + "depth_schedule": "fixed_then_adaptive", + "depth_schedule_warmup_fraction": 0.3, + "use_moe": true, + "num_experts": 18, + "num_experts_per_tok": 2, + "num_shared_experts": 1, + "moe_intermediate_size": 384, + "lookahead_steps": 2, + "layer_aware_lookahead": true, + "layer_aware_cache": true, + "lambda_balance": 0.005, + "lambda_temporal": 0.001, + "lambda_lookahead": 0.1, + "lambda_lookahead_topk": 0.05, + "lambda_cluster_locality": 0.3, + "cluster_locality_clusters": 0, + "cluster_locality_window": 128, + "lambda_router_anchor": 0.0, + "fallback_mask_prob": 0.08, + "use_hybrid_attention": true, + "kv_latent_dim": 128, + "rope_dim": 32, + "rope_layout": "per_head", + "sliding_window_size": 256, + "kv_cache_dtype": "auto", + "mla_cache_rope": true, + "router_scope": "per_iteration", + "ram_budget_mb": null, + "vram_budget_gb": 0.25, + "optimizer_paging": false, + "activation_checkpointing": false, + "extra_metadata": { + "stage_overrides": { + "pretrain": { + "learning_rate": 0.003, + "batch_size": 16, + "accumulation_steps": 2, + "max_seq_len": 1024, + "grad_clip": 1.0, + "epochs": 4, + "max_steps": 3100, + "weight_decay": 0.1, + "save_interval": 160, + "lambda_router_anchor": 0.0 + }, + "sft": { + "learning_rate": 0.001, + "batch_size": 16, + "accumulation_steps": 1, + "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": 0.00015, + "batch_size": 8, + "accumulation_steps": 1, + "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.0003, + "batch_size": 8, + "accumulation_steps": 1, + "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": 2e-05, + "batch_size": 4, + "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": 8, + "temperature": 1.0, + "max_gen_len": 256, + "lambda_router_anchor": 0.01, + "save_interval": 500 + }, + "distill": { + "learning_rate": 0.001, + "batch_size": 16, + "accumulation_steps": 1, + "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 + } + } + } + }, + "phases": [ + { + "name": "pretrain", + "budget_fraction": 0.5, + "peak_lr": 0.003, + "tokens_per_step": 32768, + "micro_tokens": 16384, + "seq_len": 1024, + "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.001, + "tokens_per_step": 16384, + "micro_tokens": 16384, + "seq_len": 1024, + "warmup_fraction": 0.01, + "min_lr_ratio": 0.05, + "weight_decay": 0.1, + "grad_clip": 1.0 + } + ], + "arch": { + "total_params": 50328836, + "active_params": 14939396, + "active_union_upper": 22902020, + "params_source": "measured (meta-device count)", + "hidden_size": 384, + "physical_layers": 5, + "recurrent_groups": 3, + "recurrent_depth": 4, + "virtual_layers": 14, + "experts": 18, + "top_k": 2, + "expert_width": 384, + "heads": "6/2" + }, + "scaling": { + "chosen_lr": 0.003, + "chosen_batch_tokens": 32768 + }, + "memory": { + "level": "none", + "ram_gb": 128.0, + "target_gb": 76.8, + "micro_tokens": 16384, + "rows_per_micro": 16, + "accumulation_steps": 2, + "training_peak_gb": 13.59, + "training_peak_label": "estimated", + "routed_expert_mb_bf16": 75.9, + "ram_budget_mb": null, + "kv_cache_dtype": "auto", + "router_scope": "per_iteration", + "activation_checkpointing": false, + "optimizer_paging": false + } + } +} \ No newline at end of file diff --git a/chronos/mlx/training/recipes/enzh_50m_72h.json b/demos/translation/recipes/enzh_50m_72h.json similarity index 90% rename from chronos/mlx/training/recipes/enzh_50m_72h.json rename to demos/translation/recipes/enzh_50m_72h.json index 0d8b082..dc54cdd 100644 --- a/chronos/mlx/training/recipes/enzh_50m_72h.json +++ b/demos/translation/recipes/enzh_50m_72h.json @@ -1,7 +1,7 @@ { - "run_dir": "runs/enzh_50m_v1", - "shards_dir": "Translation Training Data/shards_v1", - "tokenizer_dir": "Translation Training Data/tokenizer_v1", + "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, @@ -73,7 +73,7 @@ "seed": 1234, "dtype": "bfloat16", "extra": { - "raw_root": "Translation Training Data/raw", + "raw_root": "../../../Translation Training Data/raw", "params": { "total": 49691717, "active": 12090437, diff --git a/demos/translation/scripts/make_docmix_packages.py b/demos/translation/scripts/make_docmix_packages.py new file mode 100644 index 0000000..0fe74b6 --- /dev/null +++ b/demos/translation/scripts/make_docmix_packages.py @@ -0,0 +1,411 @@ +"""Self-contained MLX training packages for the doc-mix sets (A: 15%/4k/24h, B: 5%/1k/8h, C: 1%/1k/1h). + +Each package carries its data, tokenizer, eval sets, Chronos source, vendored +minimind, a recipe with relative paths, pinned requirements, launch/monitor +scripts and SHA256SUMS. Usage:: + + python demos/translation/scripts/make_docmix_packages.py OUT_DIR [A B C] + +The repository root defaults to three folders above this script (override with +``CHRONOS_REPO``); the planner output is ``recipes/docmix_plan.json`` (written by +``scripts/plan_docmix.py``). +""" +import json +import os +import shutil +import sys +import time + +DEMO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +REPO = os.path.abspath(os.environ.get("CHRONOS_REPO") or os.path.dirname(os.path.dirname(DEMO))) +sys.path[:0] = [REPO, DEMO] +from chronos_translate.bundle.builder import TOKENIZER_FILES, FLORES_REL, FLORES_FILES, MINIMIND_PARTS, sha256_file # noqa: E402 +from chronos.config_schema import model_config_kwargs # noqa: E402 + +DATA = os.path.join(REPO, "Translation Training Data") +DOCMIX = os.path.join(DATA, "docmix_v1") +RAW = os.path.join(DATA, "raw") +TOKENIZER = os.path.join(DATA, "tokenizer_v1") +MINIMIND = os.path.expanduser("~/.cache/chronos/minimind-master") +PLAN = os.path.join(DEMO, "recipes", "docmix_plan.json") +SKIP = {"__pycache__", ".DS_Store"} +# The extension's own recipes and tests are not needed inside a package. +SKIP_SRC = {os.path.join("demos", "translation", "recipes"), os.path.join("demos", "translation", "tests")} + +# Measured on the built sets (docmix_manifest.json, trial_lengths.py); live run: 2.79 pretrain / ~1.67 SFT passes. +PACKAGES = { + "A": dict(name="chronos_docmix_A_sci15_ctx4k_24h", set="sci15_ctx4096", hours=24.0, ctx=4096, + validate_every_s=7200.0, checkpoint_every_s=1800.0, bleu_samples=100, final_eval_samples=1012, + eval_reserve_fraction=0.06, doc_eval_docs=16, final_doc_eval_docs=1 << 30, tensor_stats_every_s=1800.0, + probe_samples=16, lengths={"p50": 139, "p95": 1826, "p99": 3568, "max": 4071}), + "B": dict(name="chronos_docmix_B_sci5_ctx1k_8h", set="sci5_ctx1024", hours=8.0, ctx=1024, + validate_every_s=3600.0, checkpoint_every_s=1200.0, bleu_samples=100, final_eval_samples=1012, + eval_reserve_fraction=0.08, doc_eval_docs=16, final_doc_eval_docs=1 << 30, tensor_stats_every_s=1800.0, + probe_samples=16, lengths={"p50": 138, "p95": 695, "p99": 943, "max": 1014}), + "C": dict(name="chronos_docmix_C_sci1_ctx1k_1h", set="sci1_ctx1024", hours=1.0, ctx=1024, + validate_every_s=1200.0, checkpoint_every_s=600.0, bleu_samples=50, final_eval_samples=200, + eval_reserve_fraction=0.12, doc_eval_docs=8, final_doc_eval_docs=40, tensor_stats_every_s=900.0, + probe_samples=8, lengths={"p50": 140, "p95": 696, "p99": 944, "max": 1013}), +} +TARGET_PASSES = {"pretrain": 4.0, "sft": 3.0} +# Exact versions resolved from the environment the packages were tested in (pip metadata walk). +REQUIREMENTS = """\ +annotated-types==0.7.0 +anyio==4.15.1 +certifi==2026.4.22 +click==8.2.1 +exceptiongroup==1.3.0 +fastapi==0.115.9 +filelock==3.18.0 +fsspec==2025.3.0 +h11==0.16.0 +hf-xet==1.4.3 +httpcore==1.0.9 +httpx==0.28.1 +huggingface-hub==1.10.2 +idna==3.20 +jinja2==3.1.6 +lxml==5.4.0 +markdown-it-py==3.0.0 +markupsafe==3.0.2 +mdurl==0.1.2 +mlx==0.32.4 +mlx-metal==0.32.4 +mpmath==1.3.0 +networkx==3.4.2 +numpy==1.26.4 +packaging==24.2 +portalocker==4.4.2 +psutil==7.2.2 +pyarrow==23.0.1 +pydantic==2.11.5 +pydantic-core==2.33.2 +pygments==2.19.1 +pyyaml==6.0 +regex==2026.4.4 +rich==14.0.0 +sacrebleu==2.6.0 +safetensors==0.5.3 +shellingham==1.5.4 +starlette==0.45.3 +sympy==1.14.0 +tabulate==0.10.0 +tokenizers==0.22.2 +torch==2.10.0 +tqdm==4.67.1 +transformers==5.5.0 +typer==0.16.0 +typing-extensions==4.16.0 +typing-inspection==0.4.1 +uvicorn==0.34.3 +""" + + +def _ignore(folder, names): + out = {n for n in names if n in SKIP or n.endswith((".pyc", ".egg-info"))} + for n in names: + if os.path.relpath(os.path.join(folder, n), REPO) in SKIP_SRC: + out.add(n) + return out + + +def _version(module): + import importlib + return getattr(importlib.import_module(module), "__version__", "unknown") + + +def _copytree(src, dst): + shutil.copytree(src, dst, ignore=_ignore, dirs_exist_ok=False) + + +def _copy(src, dst): + os.makedirs(os.path.dirname(dst), exist_ok=True) + shutil.copy2(src, dst) + + +def _write(path, text, mode=0o644): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w", encoding="utf-8") as fh: + fh.write(text) + os.chmod(path, mode) + + +def _json(path, obj): + _write(path, json.dumps(obj, ensure_ascii=False, indent=2) + "\n") + + +def _download_manifest(src, dst): + m = json.load(open(src, encoding="utf-8")) + for f in m.get("files", []): + f.pop("local", None) # local absolute/relative paths of this machine + _json(dst, m) + + +def recipe(key, cfg, plan, docmix): + model = dict(plan["model"]) + model.update(max_position_embeddings=cfg["ctx"], rope_theta=1e6, bos_token_id=1, eos_token_id=2, + extra_metadata={"docmix_package": cfg["name"]}) + model = model_config_kwargs(model) + assert model["num_experts"] == 18 and model["max_position_embeddings"] == cfg["ctx"] + assert model["rope_theta"] == 1e6 and model["eos_token_id"] == 2 + phases = [dict(p) for p in plan["phases"]] + for p in phases: + p["seq_len"] = cfg["ctx"] + train = docmix["splits"]["train"] + extra = { + "raw_root": "data/raw", "micro_fit": "auto", "micro_fit_fraction": 0.7, + "target_passes": TARGET_PASSES, "doc_eval_docs": cfg["doc_eval_docs"], + "package": cfg["name"], + "params": {k: plan["arch"][k] for k in ("total_params", "active_params", "active_union_upper", + "virtual_layers", "experts", "top_k", "expert_width")}, + "data": {"set": cfg["set"], "ctx_len": cfg["ctx"], "train_rows": train["rows"], "sft_records": train["records"], + "pretrain_tokens": train["pretrain_tokens"], "sft_tokens": train["sft_tokens"], + "sft_supervised_tokens": train["sft_supervised_tokens"], "record_len_tokens": cfg["lengths"], + "selection": docmix["selection"]}, + "planner": {"memory_estimate": plan["memory"], "scaling": plan["scaling"], + "windows": {"sliding_window_size": model["sliding_window_size"], + "cluster_locality_window": model["cluster_locality_window"], + "rule": "SWA = pow2 ≥ p95/4 (min 128, ≤ ctx); cluster = pow2 nearest p50"}}, + "note": ("target_passes: live 72h run reached 2.79 pretrain / ~1.67 SFT passes; +1 pass each, rounded. " + "A phase ends at its target or its time budget, whichever comes first."), + } + if cfg["final_doc_eval_docs"] < (1 << 30): + extra["final_doc_eval_docs"] = cfg["final_doc_eval_docs"] + return { + "_comment": f"Chronos doc-mix package {key}: paths are relative to this file.", + "run_dir": "run", "shards_dir": "data/shards", "tokenizer_dir": "data/tokenizer", + "model_config": model, "phases": phases, "budget_hours": cfg["hours"], + "eval_reserve_fraction": cfg["eval_reserve_fraction"], "log_every_steps": 20, + "validate_every_s": cfg["validate_every_s"], "checkpoint_every_s": cfg["checkpoint_every_s"], + "bleu_samples": cfg["bleu_samples"], "final_eval_samples": cfg["final_eval_samples"], "seed": 1234, + "dtype": "bfloat16", "schedule": "time", "tensor_stats_every_s": cfg["tensor_stats_every_s"], + "probe_samples": cfg["probe_samples"], "export_prefix": "chronos_docmix", "extra": extra, + } + + +def shards_manifest(cfg, docmix): + """Small summary the monitor shows under "data" (``shards_manifest.json`` in shards_dir).""" + tr, va = docmix["splits"]["train"], docmix["splits"]["valid"] + return {"format": "chronos-docmix-v1", "set": cfg["set"], "tokenizer_dir": "../tokenizer", + "tokenizer_fingerprint": docmix.get("tokenizer_fingerprint"), "created_utc": docmix.get("created_utc"), + "sft_max_len": cfg["ctx"], "pretrain_train_tokens": tr["pretrain_tokens"], + "pretrain_valid_tokens": va["pretrain_tokens"], "sft_train_records": tr["records"], + "sft_valid_records": va["records"], "sft_train_tokens": tr["sft_tokens"], + "sft_train_supervised_tokens": tr["sft_supervised_tokens"], "prompts": docmix.get("prompts"), + "size_share_records": tr.get("size_share_records"), "dir_share_records": tr.get("dir_share_records"), + "token_share_by_len": tr.get("token_share_by_len"), "ordered_records": tr.get("ordered_records"), + "random_concat_records": tr.get("random_concat_records"), "spec": docmix.get("spec")} + + +LAUNCH = r"""#!/bin/bash +# Chronos doc-mix training package: venv + pinned deps + resumable MLX run. +# ./launch.sh full run (rerun the same command to resume from run/checkpoints/latest.state.pkl) +# ./launch.sh --smoke 2 updates per phase + short eval in run_smoke/