0.3.3a1: every loop knob in planner and auto-tune, probe scores cross-entropy, demo model under demos/

.github/workflows/ci.yml

@@ -36,3 +36,8 @@ jobs:
       - name: Unit tests
         run: |
           python -m pytest tests/ -v --tb=short
+
+      - name: Demo driver tests
+        run: |
+          # The translation demo's model driver; the kernel suite above stays task-neutral.
+          python -m pytest demos/translation/tests/test_model_enzh_50m_sft.py -v --tb=short

CHANGELOG.md

@@ -1,5 +1,44 @@
 # Changelog
 
+## [0.3.3a1] — 2026-11-12 (beta)
+
+This release closes three gaps in 0.3.3a0. Every loop knob now has a planner recommendation and an auto-tune switch. The auto-tune probe now scores what it was meant to score. The en-zh model folder has left the repository root. Kernel defaults are unchanged, so existing configs and checkpoints behave as in 0.3.3a0.
+
+### Added
+- Planner: a looped plan (`recurrent_depth > 1`) now also sets six more loop knobs. Each comes with a `sources` entry that states the measurement behind it (15M parameters, one seed per arm). The Web UI shows these values as planner recommendations, and `ui/presets_v3.json` gains the four keys whose value differs from the base defaults.
+  - `recurrent_later_topk=0`: routing later iterations to 1 expert raised validation CE from 4.543 to 4.600 at the same speed (71.6k against 71.5k tokens/s).
+  - `recurrent_later_groups=0`: the full block has the lowest loss per step. Running 2 of 4 layers in later iterations gives about 1.3x more steps per second at depth 4. It costs up to 0.012 more training loss per step with dt fixed at 1.0 and 0.013-0.055 with the sigmoid step size. 1 is never recommended, because one layer looping alone is no longer a block loop.
+  - `lambda_iteration_cluster=0`: with cluster locality on, 99.5-100% of the clusters that later iterations use are already used at iteration 0.
+  - `lookahead_teacher_scope="first"`: the lookahead head already picks the right cluster for 0.97 of tokens at every iteration.
+  - `recurrent_dt_min=0.05` and `recurrent_dt_max=1.0`: the sigmoid arms used the ranges 0.05-1 (depth 4) and 0.1-1 (depth 8), and their learned dt stayed between 0.14 and 0.30, inside the range.
+- Auto-tune can now search every loop knob. The list adds `recurrent_dt_param`, `recurrent_dt_min`, `recurrent_dt_max`, `recurrent_later_topk` (0 or 1), `recurrent_later_groups` (0 or 2), `lambda_iteration_cluster` (0, 0.1, 0.3 or 1.0) and `lookahead_teacher_scope`. They join `recurrent_order` and `recurrent_dt_init`.
+  - Searching any of them makes the probe a looped model with 4 block layers, one prelude layer and one coda layer. The 0.3.3a0 probe had 2 block layers, so `recurrent_later_groups=2` would have meant the full block.
+  - `recurrent_later_groups` is drawn only in block order, and 1 is not offered.
+  - If a drawn `recurrent_dt_init` falls outside the drawn `[dt_min, dt_max]`, the trial starts at the nearest bound.
+  - Each trial logs its median step time and stores it as the Optuna user attribute `step_time_s`. `recurrent_later_topk` and `recurrent_later_groups` save time, not loss, so a cross-entropy score alone favours "off".
+- Help text for these knobs and for the auto-tune search list now gives the measured reason and the recommended value, in English, 简体中文, 繁體中文 and 日本語.
+
+### Fixed
+- Auto-tune score. The probe trained on and scored `CE + lambda_balance * L_balance + lambda_temporal * L_temporal`, so a smaller weight lowered a trial's score by itself. The probe now trains on the trainers' objective (`chronos_loss_term`, every weighted term of the config) and scores cross-entropy only.
+- Auto-tune `lambda_lookahead` had no effect. The probe never passed it to the model config or to the loss, and it never trained the lookahead head, so `lookahead_steps` changed the model without supervising it. With the 0.3.3a0 probe, `lambda_lookahead` 1e-3 and 1.0 gave identical per-step scores; with this release the scores differ after the first update. The lookahead, iteration-cluster and recurrent terms of the config now act in the probe as they do in training.
+- Web UI schema payload. `Field.as_dict` dropped every value equal to `False`, and in Python `0 == False`. So `min=0` of 48 fields, and `recommended=0` of `max_steps` and `cluster_locality_clusters`, never reached the page. Those number inputs had no lower bound, and their tooltips showed the minimum as "−∞". The payload now drops only unset values.
+
+### Changed
+- The en-zh model folder moved from `models/chronos-enzh-50m-sft` to `demos/translation/models/chronos-enzh-50m-sft`, and its driver test moved from `tests/` to `demos/translation/tests/`. CI runs that test in its own step after the kernel suite, and the repository root no longer holds a task-specific model. The driver finds its files relative to itself, so commands run inside the folder are unchanged, and the weights stay on the release `enzh-50m-sft-v1`. `MANIFEST.in` no longer prunes `models/`.
+- The main READMEs now describe the speculative-decoding, native-kernel parity and MLX compile measurements as taken on the 50M demo checkpoint. They link to a new section of `demos/translation/README.md` and `README_zh.md` that states the exact conditions (FLORES+ dev en→zh prompts, en→zh shards); the numbers are unchanged. Each main README names the translation demo in one pointer paragraph only. `tests/test_ui_generic.py::test_task_models_and_tests_live_under_demos` checks this rule and the folder layout.
+- The preset label "Chronos-v3 50M (16k vocab)" stays. It names the preset's explicit 16,384-token vocabulary and no task.
+
+### Verified
+- End to end through the Web UI endpoint, with an in-process client and no open port: a 4-trial study of 10 probe steps each searched all nine loop knobs plus `lambda_lookahead`. It trained on 19,090 real dataset tokens with the bundled 32k tokenizer, completed every trial, logged each trial's cross-entropy and median step time, and returned a best-config patch with every searched knob.
+
+### Tests
+- New `tests/test_loop_knobs_wiring.py` (10 items). It checks that every loop knob is searchable and reaches the probe model, and that `recurrent_later_groups` never offers 1 and is drawn only in block order. It checks that `lambda_lookahead` 1e-3 and 1.0 give the same score at step 0 and different scores after the first update. It also covers the best-config patch, the Web UI search list with help in four languages, planner recommendations with a source at depths 2, 4 and 8, a depth-1 plan with no loop knob, and a schema payload that keeps 0.
+- New `tests/test_ui_generic.py::test_task_models_and_tests_live_under_demos`: `models/` and `tests/` hold no task-specific names, and each main README names the en-zh demo in one line, which points to `demos/translation`.
+- `tests/test_autotune_loop.py`: `recurrent_later_groups` is now expected in the search list.
+- `demos/translation/tests/test_model_enzh_50m_sft.py` (15 items, moved from `tests/`) passes from its new folder; CI runs it as the step "Demo driver tests".
+- Full suite on macOS (Apple silicon, Python 3.11, MLX on CPU): 728 passed, 4 skipped. `ruff check .` passes on the whole tree.
+- A local wheel and sdist build both carry the `general_32k` tokenizer and no file from `demos/`, no model folder and no weights.
+
 ## [0.3.3a0] — 2026-11-12 (beta)
 
 ### Added

MANIFEST.in

@@ -13,4 +13,3 @@ prune reports
 prune .archive
 prune Translation?Training?Data
 prune demos
-prune models

README.md

@@ -89,7 +89,7 @@ flowchart LR
 
 **Self-speculative decoding (no extra parameters).** The same model drafts γ tokens at a shallow depth and verifies them with one full-depth batched forward pass. Greedy output is token-for-token identical to full-depth greedy. Sampling uses rejection sampling, which preserves the full-depth distribution. After a rejection, the cache is truncated: MLA slices `c_kv`, and the sliding-window cache keeps `window + γ_max` entries. Torch (`chronos/runtime/speculative.py`) and MLX (`chronos/mlx/speculative.py`, `ChronosMLXInferenceEngine.generate(speculative=True)`) share the algorithm. Tests cover greedy identity, the rejection distribution and rollback against a fresh prefill on both backends. The Web UI Inference tab has a switch for it, with draft depth and γ.
 
-For the 50M demo it is exact but not faster on any runtime we measured (pretrain step-13,996 checkpoint, greedy, batch 1, FLORES+ dev en→zh, draft depth 1, γ 4; reports in `reports/speculative/`):
+For the 50M demo checkpoint it is exact but not faster on any runtime we measured (pretrain step-13,996 checkpoint of the demo in `demos/translation`, greedy, batch 1, the demo's dev prompts, draft depth 1, γ 4; exact conditions in [`demos/translation/README.md`](demos/translation/README.md); reports in `reports/speculative/`):
 
 | Runtime | Plain tok/s | Speculative tok/s | Ratio | Acceptance |
 |---|---|---|---|---|
@@ -113,7 +113,7 @@ At this size a full-depth step costs about as much as a batched verify pass. Wit
 
 `chronos.model.factory.fit_config` searches for a configuration that matches total and active parameter targets.
 
-**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`. The 50M en→zh status, the released model `chronos-enzh-50m-sft` with its measurements, and the model bundle are documented in [`demos/translation/README.md`](demos/translation/README.md).
+**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`. The 50M en→zh status, the released model `chronos-enzh-50m-sft` (files in `demos/translation/models/`) with its measurements, and the model bundle are documented in [`demos/translation/README.md`](demos/translation/README.md).
 
 **Inference and memory options (all off by default).** Each option is a config field and a `chronos serve` / Inference-tab override. Tests compare each one with its reference path on Torch and MLX.
 
@@ -153,11 +153,11 @@ Up to 0.3.1a3, Chronos took its config base, RoPE tables, gated FFN, dense atten
 
 **Old vs new, measured 2026-11-08**
 
-Setup: Apple M5 Max, PyTorch 2.10 on CPU, FP32, 4 threads, median of 5–7 runs, each kernel in its own process, runs alternated twice. The 50M en→zh SFT run was training on the GPU at the same time, so absolute timings carry noise. Old = 0.3.1a3 + minimind `24896cd`.
+Setup: Apple M5 Max, PyTorch 2.10 on CPU, FP32, 4 threads, median of 5–7 runs, each kernel in its own process, runs alternated twice. A 50M training run was using the GPU at the same time, so absolute timings carry noise. Old = 0.3.1a3 + minimind `24896cd`.
 
 | Check | 0.3.1a3 + minimind | 0.3.2 native |
 |---|---|---|
-| Logits, `out/chronos_256_moe.pth` and the 50M en→zh checkpoint (full, padded and prefill+decode paths) | reference | bit-identical (max abs diff 0) |
+| Logits, `out/chronos_256_moe.pth` and the 50M demo checkpoint (full, padded and prefill+decode paths) | reference | bit-identical (max abs diff 0) |
 | Random-init `state_dict` and config values (same seed) | reference | identical |
 | `import chronos` with no cache and no network | fails (`git clone` error) | works |
 | minimind on `sys.path`, top-level `model` module imported | yes | no |
@@ -181,7 +181,7 @@ Up to 0.3.1a3 the MLX trainers ran every micro step eagerly and kept every activ
 - **Compiled micro step** (`chronos/mlx/training/compiled.py`, on by default). Loss and gradients run as one `mx.compile` graph per signature. The executed recurrent depth, the ACT switch, input shapes and plain config values are part of the signature, so a step never reuses a stale graph. Weights and the RNG state are compile inputs and outputs, so dropout and the cache-miss simulation consume the RNG stream exactly as before. Multimodal batches, lazy (paged) experts and any signature that fails to trace run eagerly, and the reason is logged once. `CHRONOS_MLX_COMPILE=auto|on|off` or the recipe's `extra.mlx_compile` switches it.
 - **Activation recomputation** (`chronos/mlx/checkpointing.py`, `activation_checkpointing=true`; the planner's memory-constrained tiers turn it on). Each block call becomes one recomputed region. A plain `mx.checkpoint` is wrong here in two ways, both measured on mlx 0.32.4. Dropout masks and the cache-miss draws would be drawn again during recomputation; they are recorded in the forward pass and replayed. A region that returns the same graph node twice keeps the gradient of only one copy. A block's KV cache and its shared view are such a pair whenever the SWA window covers the sequence (a full-length slice is a new Python object over the same node), and the gradients were 37–40% off. Region outputs are therefore packed into one array per dtype.
 
-**Measured 2026-11-08.** Setup: Apple M5 Max, mlx 0.32.4, bf16, the 50M en→zh recipe's model config from random init, real en→zh shards, the recipe's step sizes (SFT 8 micro-batches per step, pretrain 16). Median of the second half of 16 (SFT) or 12 (pretrain) steps. The 72-hour run was training on the same GPU, so absolute timings carry noise.
+**Measured 2026-11-08.** Setup: Apple M5 Max, mlx 0.32.4, bf16, the 50M demo recipe's model config (`demos/translation/recipes/`) from random init, the demo's real training shards, the recipe's step sizes (SFT 8 micro-batches per step, pretrain 16). Median of the second half of 16 (SFT) or 12 (pretrain) steps. The 72-hour run was training on the same GPU, so absolute timings carry noise.
 
 | Micro step | SFT s/step | SFT peak memory | Pretrain s/step | Pretrain peak memory |
 |---|---|---|---|---|
@@ -519,7 +519,7 @@ pip install "Project_Chronos[mlx]"
 
 ```bash
 pip install "Project_Chronos[serve]"      # chronos serve / chronos monitor (fastapi, uvicorn)
-pip install -e demos/translation          # en↔zh demo extension, from a checkout only (not on PyPI): chronos-translate
+pip install -e demos/translation          # translation demo extension (chronos-translate), from a checkout only, not on PyPI
 ```
 
 **vLLM serving (optional, Linux + CUDA only):**
@@ -858,7 +858,7 @@ Project_Chronos/
 ├── docs/
 │   └── vllm_integration.md
 ├── demos/
-│   └── translation/               # en↔zh extension (chronos-translate): corpus, runs, bundles; not in the wheel
+│   └── translation/               # demo extension (chronos-translate): corpus, runs, bundles, the 50M demo model; not in the wheel
 ├── pyproject.toml
 └── README.md / README_zh.md / THIRD_PARTY_NOTICES.md
 ```
@@ -908,7 +908,7 @@ timeline
         : int8 KV, MLA absorbed decode, shared router scope (opt-in)
         : Auto planner, v3 presets, tokenizer and serve tabs
     Next
-        : 50M en-zh results and ablations
+        : Multi-seed and larger-size runs of the 0.3.3 loop settings
         : Speculative decoding: a draft path that pays off on larger models and lazy experts
         : Train IntentClassifier on large activation corpora
         : Benchmark 7B+ checkpoints