Skip to content

v0.3.3a0

Pre-release
Pre-release

Choose a tag to compare

@FonaTech FonaTech released this 11 Nov 23:55

0.3.3a0 (pre-release). Block-order loop, sigmoid step size, general 32k tokenizer, task-neutral kernel. Kernel defaults are unchanged from 0.3.2a3. The notes below are the CHANGELOG section for this version.

Mechanism in short: a shared block applied repeatedly adds nearly parallel vectors, so contributions add constructively; each iteration multiplies the residual by g = 1 + dt * lambda_eff with lambda_eff learned, and depth R raises it to the power R - 1. dt alone is a gauge; the fixed horizon T = depth * dt = 2 is the planner default, a safe start and not a fitted optimum.

Open: all loop measurements are one seed at 15M parameters; lazy serving of a block-order checkpoint is unmeasured; one full-suite run failed once in test_worker_probes_checkpoint_after_validation_and_monitor_shows_it (cause not found; 6 later full runs and 450 isolated spawns passed).


Added

  • recurrent_order ("layer" default, "block"): the order in which the recurrent layers repeat. layer runs each recurrent layer recurrent_depth times in a row (A A B B); block runs the whole group of recurrent layers and then repeats it (A B A B), so iteration 0 is an ordinary depth-1 forward pass and later iterations refine it. In block order there is one state trajectory and one halting head per block iteration (the halting_probs and depth_used group axis has length 1); every physical layer keeps its own dt, loop and group embeddings and its own KV slot. The default path is bit-identical to 0.3.2a3 (Torch and MLX logits compared on four attention modes, max difference 0). Available in the Torch and MLX models, the web UI (four languages) and MODEL_FIELDS. A block checkpoint is not interchangeable with a layer checkpoint.
  • recurrent_later_topk (int, default 0 = off): loop iterations after the first route each token to this many experts (capped at num_experts_per_tok); the shared expert still runs. With router_scope="shared" the kept experts are the best-scoring ones of the first dispatch, so no new expert page is touched. The count is an explicit argument of the MoE, block and recompute calls in Torch and MLX, so activation recomputation replays it. The default path is bit-identical (max difference 0 on both backends, four attention modes).
  • recurrent_later_groups (int, default 0 = all): in block order, loop iterations after the first run only the first N layers of the block. Cost in this kernel scales with the number of layer calls, so this removes calls and the expert pages those layers would load. Iteration 0 still runs every layer and writes every KV slot. It needs recurrent_attention_mode shared_first or first_only, because a skipped layer has no per-iteration KV slot, and the config raises a ValueError otherwise. The default path is bit-identical.
  • lambda_iteration_cluster (float, default 0) and lookahead_teacher_scope ("first" default, "iterations"): the first adds a total-variation penalty between a token's per-iteration cluster-mass distribution and its mean over iterations, so the loop layers keep a token in one expert cluster across depth without locking experts. The second makes the lookahead router's teacher the router probabilities averaged over the layer's iterations. Both exist in Torch and MLX; the MLX loss is checked against the Torch loss and against mx.compile.
  • tools/bench_residency.py --paging expert|cluster|auto (default expert, the historical layout) and tools/bench_speculative.py mlx_runner(..., paging=...): the lazy cache can be paged by training cluster, as production serving does for checkpoints trained with lambda_cluster_locality > 0. Each row now also records cluster_experts_loaded, cluster_aware, num_clusters, warmup_cluster_loads and prefill_cluster_loads, and the summary decode_cluster_loads.
  • recurrent_dt_param ("clamp" default, "sigmoid"): how the learned loop step size is stored. clamp keeps dt = clamp(w, dt_min, dt_max), whose gradient is zero once w leaves the interval. sigmoid uses dt = dt_min + (dt_max - dt_min) * sigmoid(w), whose gradient is positive everywhere, so a step size that has become small can still grow. The stored logit is initialised so that step 0 gives exactly recurrent_dt_init (clamp and sigmoid logits agree to 1.1e-6 at init; Torch and MLX agree to 1.3e-6). The default path is unchanged.
  • chronos.model.recurrent.recommended_dt_init(depth, horizon=2.0, cap=0.5) returns min(cap, horizon / depth): 0.5 at depth 4, 0.333 at depth 6, 0.25 at depth 8. One loop iteration is an explicit Euler step x <- x + dt * f(x), so depth iterations integrate over the horizon T = depth * dt; the rule holds T fixed and lets depth refine the integration. It is a starting value for recurrent_dt_init, not a fitted optimum (see Physical mechanism and Measured). An earlier draft used 0.5 * (4 / depth) ** 3; the depth-8 run below refuted it and it was removed before release.
  • recurrent_dt_param, recurrent_dt_init, recurrent_dt_min and recurrent_dt_max appear in the Web UI recurrent section with labels and help in four languages. The planner recommends block order, sigmoid and the horizon-based recurrent_dt_init whenever recurrent_depth > 1 (config, sources report entries, and the schema's recommended value).
  • Auto-tune: tune_recurrent_order and tune_recurrent_dt_init join the search space (Web UI "Search over" list, ChronosSearchSpaceConfig, best-config patch). When either is searched the probe model runs probe_recurrent_depth iterations (the Web UI passes the configured recurrent_depth, clamped to 2-8) with one prelude and one coda layer. recurrent_later_groups is deliberately not searched: it trades quality for cost, so a loss-only search would always choose the full block. The probe loss comes from a few steps and cannot resolve differences as small as dt 0.2-0.5 at depth 4; treat the two loop knobs as a coarse screen.
  • Auto-tune probe data: chronos/tuning/probe_data.py encodes a bounded sample of the user's dataset (at most 400k tokens and 20k records) with the configured tokenizer and trains the probe on it. The probe used to train on random token ids, so no hyperparameter could change its loss in a way that carries over to real data. It falls back to random ids, with a log line, when no text can be read or the checkpoint's vocabulary is smaller than the tokenizer's. On Apple silicon the probe now picks a torch device itself when the training backend is mlx, which is not a torch device.
  • chronos/assets/tokenizers/general_32k (32,768 tokens, 1.9 MB) ships in the wheel, with chronos.deps.get_general_tokenizer_path() and GENERAL_TOKENIZER_VOCAB. It is the vocabulary built for the pretraining mix, with its control tokens relabelled to the kernel's task-neutral layout (<|endoftext|>, <|im_start|>, <|im_end|>, <|pad|>, then reserved ids 4-63); ordinary tokens, merges and ids are unchanged (1,004 of 1,004 test texts encode to the same ids as with the original tokenizer, on the native and on the Hugging Face tokenizer). Because the fingerprint hashes the special tokens, relabelling changed it from 169668b1… to 6d4b3b52…; the old value is recorded under derived_from, so a checkpoint trained with the original relabel-free tokenizer reports a different fingerprint than this asset. It carries no task tokens. NOTICE lists the corpus licences: FineWeb-Edu and Cosmopedia v2 (ODC-By), Wikipedia (CC-BY-SA 3.0 and GFDL), and Fineweb-Edu-Chinese V2.1, whose commercial use needs the owner's permission.
  • Planner default vocabulary: a new plan whose request and dataset profile name no vocabulary now uses 32,768 and writes tokenizer_path pointing at the bundled tokenizer, with sources entries (it was 16,384 with no path). An explicit vocab_size keeps the user's choice and adds no path. get_tokenizer_path() still returns the legacy 6,400 tokenizer, so checkpoints that never recorded a tokenizer are not changed.
  • Web UI: recurrent_later_topk and recurrent_later_groups appear in the recurrent section, lambda_iteration_cluster and lookahead_teacher_scope in the objectives section, each with labels and help in English, 简体中文, 繁體中文 and 日本語.

Fixed

  • MLX trainer: a recurrent step size (recurrent_dt_*) that sat exactly on recurrent_dt_min or recurrent_dt_max received no gradient, because mx.clip has gradient 0 on its bounds while the Torch trainer's clamp gives 1. A run with recurrent_dt_init equal to a bound therefore kept its step sizes frozen on MLX only (a 30-minute arm with init 1.0 = max ended with raw step sizes exactly 1.0). RecurrentControllerMLX.contract now returns the same value with the Torch gradient on the bounds. Values strictly inside the interval, and the forward pass, are unchanged, so default configs train as before; values strictly outside the interval still get no gradient on either backend.
  • MLX trainer: with recurrent_depth_jitter > 0 and gradient accumulation, each micro-batch drew its own loop depth, so one optimizer step mixed graphs of different depth and the router statistics accumulated across micro-batches failed with a shape error (broadcast_shapes (14,24) and (10,24)). The depth is now drawn once per optimizer step and shared by all its micro-batches, in eager and compiled steps. Runs without jitter are unchanged.

Changed

  • tools/bench_residency.py cluster_loads no longer includes the page reads of the store reset's warm-up (they are reported separately as warmup_cluster_loads); earlier reports included them, because the store's counters are cleared before the warm-up. Every mode starts from the same cold warm-up, so absolute counts drop by that constant per (prompt, mode).
  • The kernel is now task-neutral in code and text. chronos/cli.py, chronos/mlx/training/longrun.py, chronos/mlx/training/compiled.py, chronos/tokenizer/extract.py, chronos/data/chat_records.py and the examples in ui/i18n.py no longer name a translation task, and the tokenizer asset's provenance metadata no longer names task tokens. tests/test_ui_generic.py::test_kernel_has_no_translation_task_code walks chronos/ and ui/ and fails on any task-specific term. The en-zh translation work stays in demos/translation (chronos-translate), which imports the kernel; its status, released-model and bundle sections moved verbatim from README.md and README_zh.md into demos/translation/README.md and demos/translation/README_zh.md. The main READMEs keep one pointer sentence and still quote en-zh numbers where they are the measurement context (speculative decoding, the 0.3.2 native-kernel benchmarks, the 50M preset row); those are measurements on one checkpoint, not defaults. models/chronos-enzh-50m-sft stays in the repository as the demo's model.
  • Planner and Web UI recommendations for a looped model changed (recurrent_order="block", recurrent_dt_param="sigmoid", recurrent_dt_init = min(0.5, 2 / depth)), and ui/presets_v3.json was regenerated: all five tiers gain those three keys and nothing else changed. The kernel defaults did not change (recurrent_order="layer", recurrent_dt_param="clamp", recurrent_dt_init=0.25), so a config that does not set them, and every existing checkpoint, behaves as in 0.3.2a3.

Physical mechanism

A shared block applied again and again to the residual stream acts as one operator, so each iteration adds a vector that points along nearly the same direction as the one before. The contributions add constructively instead of as independent, orthogonal increments. Measured on the trained checkpoints, the cosine between the outputs of consecutive block calls rises with iteration, for example from 0.92 to 0.985 (depth 8, dt 0.25) and from 0.98 to 0.996 (depth 8, dt 0.5); independent layers would give a cosine near 0.

  • One iteration is h <- h + dt * (block(h + e_loop) - h). Near the dominant direction the block update is about lambda_eff * h, so each iteration multiplies the residual norm by g = 1 + dt * lambda_eff. lambda_eff is learned, not set: the optimiser rescales the block's output when dt changes. Across step sizes of 0.146 (learned by a sigmoid arm), 0.25 and 0.5 (both fixed) at depth 8, g stayed between 1.067 and 1.075 while lambda_eff = (g - 1) / dt fell from 0.48 to 0.30 to 0.13.
  • The R - 1 iterations after iteration 0 each apply the gain, so depth R raises it to the power R - 1: total amplification g^(R-1), total log-gain (R - 1) * ln g. In every arm whose step size could stay large that total was 0.36-0.51 (1.4-1.7x) at depth 4 and at depth 8 (depth 4: 0.36-0.44; depth 8: 0.38-0.51). Arms whose step size collapsed onto its floor or stayed tiny had a lower total (0.09-0.24); at equal step they were worse than the best arm of their depth, but not always worse than every arm with a larger total (see the limits below).
  • dt alone is therefore not the physical quantity; kappa = dt * lambda_eff and the total log-gain are. This is why the loss is flat for dt between about 0.2 and 0.5 at depth 4, and why the horizon T = depth * dt (the length of the Euler integration of dh/dtau = block(h) - h) is a better handle than dt: more depth at fixed T refines the integration without raising the total gain.
  • Limits of the claim: the total gain is not a loss predictor (the depth-8 arm with a collapsed step size, total 1.27x, beat the fixed-0.5 arm, 1.57x, at equal step by 0.02-0.05), the measurements are one seed per arm at 15M parameters, and T = 2 is a safe starting point inside a flat region, not a fitted optimum: the sigmoid arm that started at T = 2 settled at T = 1.16.

Measured

15M-parameter config, recurrent depth 4, 24 experts with top-2, one Mac, 30 minutes per arm (about 1,850-1,960 steps), one seed per arm unless stated, same data order in every arm. Paired differences are training loss in 10-step windows; a repeat of the same block-order arm differs from the first by +0.002 to +0.008, so differences below about 0.01 are not separable from run-to-run noise.

  • Block against layer order at the default step size (init 0.25, clip [0.05, 1]): block order has the lower paired loss by 0.025-0.075 in every block of steps after step 100, at equal speed (about 71.5k tokens/s). Validation CE at 30 minutes: 4.543 (block, 1,960 steps) against 4.586 (layer, 1,855 steps).
  • The default step sizes collapse: the learned values end at 0.045-0.09 against a floor of 0.05, and the gradient below the floor is exactly zero on both backends, so a collapsed step size never recovers. Fixing every step size at 1.0 raises the paired loss of block order by 0.050 in steps 10-100, then lowers it by 0.022, 0.025, 0.101, 0.072 and 0.026 in the blocks 100-300, 300-600, 600-1,000, 1,000-1,500 and 1,500-1,960 (largest at steps 600-1,000). Validation CE at 30 minutes is slightly higher, 4.559 at 1,850 steps against 4.543 at 1,960, because the fixed-step arm completed fewer steps.
  • With step sizes fixed at 1.0, layer order has the lower loss after about step 600 (block minus layer +0.015 to +0.034) and block order only learns faster during the first 600 steps. The advantage of block order therefore holds at the default, collapsing step size only; the default is not changed in this release.
  • Depth 8 does not follow the depth-4 result. With step sizes fixed at 1.0, block order at depth 8 (1,255 steps in 30 minutes, 1.43 s per step, final validation CE 4.821) has a higher paired loss than the default-step-size depth-8 arm in every one of 82 windows (+0.199, +0.075, +0.039 and +0.077 in steps 10-100, 100-300, 300-600 and 600-1,000), and its gradient norm is 1.4-1.8 times higher (mean 1.27 in the first 100 steps, so the 1.0 clip was active there). Against fixed-step depth 4 it has no equal-step gain (validation CE 5.328 against 5.295 at step 800, 4.794 against 4.773 at step 1,250) at 1.5 times the time per step. The default step-size rule is therefore left as it is.
  • Truncating the fixed-step depth-8 model to 1 to 8 iterations gives validation CE 5.615, 5.201, 4.995, 4.884, 4.823, 4.794, 4.779 and 4.773, so iterations 5 to 8 add 0.11 nats.
  • Cost of the loop: a depth-4 step takes 0.945-0.976 s against 0.569 s at depth 1 (1.7x), for roughly 5-13% fewer steps to reach the same validation CE, so the loop loses on wall time at this size. Cost follows the number of layer calls, not the number of experts or tokens.
  • Where the gain sits (validation CE with the loop truncated to 1, 2, 3, 4 iterations): 4.977, 4.637, 4.527, 4.506. The top-1 token equals the full-depth top-1 for 52%, 71%, 86% and 100% of tokens. The fourth iteration is worth about 0.02 nats.
  • recurrent_later_groups=2 (block order, step sizes fixed at 1.0, 2 of the 4 block layers in iterations 2-4): 2,450 steps in 30 minutes against 1,850 for the full block (0.748 against 1.001 s per step, 89.2k against 68.4k tokens/s, peak memory 6.57 against 7.43 GB) and a final validation CE of 4.456 against 4.559. At equal step the paired loss differs from the full block by -0.005, +0.007, -0.019, +0.001, +0.005 and +0.012 across the six blocks of steps, so the saving costs at most about 0.01 per step (the noise floor is 0.005-0.008). One seed. A depth-1 model trained for the same 30 minutes still reaches 4.379 in 3,170 steps, so the loop does not yet pay for itself in wall time at this size. Truncated to 1, 2, 3 and 4 iterations this model gives validation CE 4.624, 4.494, 4.439 and 4.418, against 4.987, 4.659, 4.542 and 4.514 for the full block: its first iteration is much better and the loop adds 0.206 nats in total.
  • recurrent_later_groups=2 at depth 8 (step sizes fixed at 1.0): 1,825 steps in 30 minutes against 1,255 for the full block (1.004 against 1.458 s per step) and a final validation CE of 4.586 against 4.821. At equal step the paired loss is higher by 0.007, 0.029, 0.029 and 0.022 after step 100, and validation CE at step 1,250 is 4.809 against 4.794, so skipping half the layers in seven later iterations costs about 0.02-0.03 per step, more than at depth 4. Depth 4 with the same setting is better at every point (validation CE 4.775 against 4.809 at step 1,250, 4.456 against 4.586 at 30 minutes).
  • recurrent_later_topk=1 raised validation CE from 4.543 to 4.600 and was no faster (71.6k against 71.5k tokens/s): at this width fewer routed experts saves no time.
  • Cluster level: with lambda_cluster_locality=0.3 a token touches 1.003-1.024 of the 6 expert clusters in the first iteration, and 99.5-100% of the clusters touched in the later three iterations were already touched in the first. lambda_iteration_cluster therefore has almost nothing to act on; with step sizes fixed at 1.0 its arm had paired loss lower by 0.007, 0.000, 0.005, 0.024, 0.018 and 0.008 across the six blocks of steps (one seed; all but the 600-1,500 blocks are within the noise floor), and the top-2 overlap with the first iteration differed by at most 0.011.
  • Lookahead head: its next-token prediction recovers the top-2 experts of loop iterations 1-4 with recall 0.68-0.78 (top-2 slots) and 0.97-0.98 (top-4 slots), and picks the right cluster for 0.97 of tokens at every iteration, the same as repeating the previous token's cluster. An iteration-aware head was not built.
  • Step size at depth 4 (block order, fixed step size, one seed per arm): fixing dt = 0.5 against dt = 1.0 lowers the paired loss by 0.052, 0.016, 0.032, 0.066, 0.047 and 0.026 across the six blocks of steps (every window lower in every block), and validation CE is 4.505 at 1,975 steps against 4.559 at 1,850. A sigmoid arm started at 0.5 (range 0.05-1) moved to 0.182-0.203 in all four layers with no change in loss: its paired difference from the fixed-0.5 arm is +0.002, -0.003, -0.002, +0.007, +0.001 and -0.002, inside the noise floor. At depth 4 the loss is therefore flat for dt between about 0.2 and 0.5, and worse at 1.0 or at a dt stuck at 0.05.
  • Block against layer order with a healthy step size (both arms sigmoid, init 0.5, range 0.05-1): block order has the lower paired loss in every window block, by 0.051, 0.093, 0.081, 0.053, 0.039 and 0.027 (layer minus block; 0 of 20, 0 of 30, 0 of 40, 0 of 50 and 0 of 27 windows lower for layer order after step 100). Validation CE at 30 minutes is 4.542 for block (1,830 steps) and 4.572 for layer (1,775 steps). This supersedes the reading above that block order is only better at a collapsing step size: with step sizes fixed at 1.0 layer order wins late, with a trainable step size block order wins in every block of steps measured (in steps 10-100, 7 of 9 windows).
  • Depth 8 with a fixed integration horizon (block order, one seed per arm, 30 minutes, 1,185-1,265 steps, same data order): dt = 0.25 (T = depth * dt = 2) has a lower paired loss than dt = 0.5 (T = 4) by 0.046, 0.018, 0.031, 0.079 and 0.075 across the five blocks of steps from step 10 to the end of the run (every window lower in every block); validation CE is 4.776 at 1,215 steps against 4.794 at 1,265. Against the default clamp arm, whose step size collapsed onto the 0.05 floor, it is lower by 0.003, 0.013, 0.015 and 0.002 (about the noise floor). A sigmoid arm started at 0.25 (range 0.1-1) moved to 0.141-0.147 (T about 1.16), tied with the fixed arm until step 600 (+0.002 in three blocks) and then lower by 0.023 and 0.020 (58 of 61 windows), with validation CE 4.761 at 1,220 steps against 4.776 at 1,215. At equal step the depth-8 dt = 0.25 arm is lower than the depth-4 dt = 0.5 arm by 0.041 and 0.049 in steps 100-300 and 300-600 and level or slightly higher afterwards (+0.002 and +0.008, at the noise floor).
  • The cubic depth rule dt0 = 0.5 * (4 / depth)^3 was an earlier draft and is refuted at depth 8: a sigmoid arm started at its prediction 0.0625 and ended at 0.031-0.037, and its paired loss is higher than the fixed dt = 0.5 arm by 0.066, 0.079, 0.122 and 0.256 after step 100, and higher than the fixed dt = 0.25 arm by 0.085, 0.110, 0.201 and 0.333. Its total log-gain was 0.088, the lowest of any arm. The rule was replaced by T / depth.
  • Progressive loop curriculum (recurrent_depth rising 1 to 4 over 1,600 steps, fixed dt = 0.5, depth jitter 0.01, block order): median time per step is 0.56, 0.725, 0.899 and 1.08 s at depth 1, 2, 3 and 4. At equal step it trails the fixed-depth-4 arm until the depth reaches 3-4: validation CE 5.596 against 5.215 at step 800 and 5.168 against 4.920 at step 1,000. Its early steps are cheaper, and it gives up equal-step quality until the full depth is reached; the loop's advantage is equal-token sample efficiency, not wall clock.
  • recurrent_later_groups=1 leaves three of the four block layers dormant in every later iteration, so the loop reduces to one layer applied repeatedly. With step sizes fixed at 1.0 its paired loss is higher than recurrent_later_groups=2 by 0.062 and 0.050 in steps 100-300 and 300-600 and by 0.004, 0.004 and 0.012 afterwards, and the two arms end 0.049 apart in validation CE at 30 minutes (4.407 at 2,810 steps against 4.456 at 2,450) only because 1 layer is cheaper. Use recurrent_later_groups of at least 2.
  • recurrent_later_groups=2 with the default clamp step size (block order, init 0.25, clip [0.05, 1]) was much worse than with a step size that can move. Against the full-block arm with the same default step size (block order, a repeat run), it trailed by +0.007, +0.017, +0.030, +0.046, +0.200 and +0.097 across the six blocks of steps (about +0.235 in steps 1,200-1,400), at a median gradient norm of 0.54 against 0.51 after step 1,000 and the same learning rate; it had no rollbacks. The learned step sizes ended at 0.056, 0.050, 0.133 and 0.139, so the two layers that still run in later iterations scaled their updates by 0.056 and 0.050. With a sigmoid step size from 0.5 (range 0.05-1) and the same recurrent_later_groups=2, the paired loss is lower than the clamp arm by 0.037, 0.067, 0.150, 0.284 and 0.128 after step 100, validation CE at step 1,000 is 4.980 against 5.254, and the two later-iteration layers end at 0.185 and 0.170. The step-size collapse therefore accounts for most of the regression but not all of it. The prediction that the sigmoid arm would stay within about 0.01 of the full block was not met: it is above the full-block sigmoid arm by 0.013, 0.021, 0.022, 0.055, 0.034 and 0.019 (2 of 9, 0 of 20, 2 of 30, 0 of 40, 1 of 50 and 2 of 33 windows lower), more than the 0.01 that skipping layers cost with the step size fixed at 1.0. The two arms also differ in initial step size (0.5 against 0.25), so the effect of the parameterisation and of the initial value are not separated, and the loss gap that opens at steps 1,000-1,500 and then closes is not explained. At 30 minutes the sigmoid arm reaches validation CE 4.440 in 2,365 steps, against 4.491 in 2,460 for the clamp arm and 4.542 in 1,830 for the full-block sigmoid arm; at equal step the full block is ahead (4.928 against 4.980 at step 1,000).
  • Step time of this release's kernel: the same depth-4 configuration takes a median 0.997 s per step with the clamp and 1.017 s with the sigmoid (steps 50-400, one 6-minute arm each), against 0.958 s on the development kernel that had block order and recurrent_later_groups but not the clip-gradient fix. The 4% and 6% are not attributed to a cause (the differences are the where-based clip gradient, the per-step depth draw and the sigmoid branch), and one run each cannot separate them from run-to-run noise.

Known limits

  • Block order with lazy experts: tests/test_recurrent_order.py::test_mlx_engine_block_order_runs_with_lazy_experts checks that the MLX engine with a two-expert resident set generates 10 valid tokens speculatively, and the all-resident engine is checked for exact speculative decoding. Nothing compares lazy block-order output with all-resident block-order output token by token, and the residency planner, the cluster-paged store and the async prefetcher have not been benchmarked on a block-order checkpoint (expert_store.py and cluster_layout.py index pages by physical layer, which is the same page set in both orders; only the order in which layers are used differs). Treat lazy serving of a block-order checkpoint as unmeasured.
  • All loop measurements are one seed per arm at 15M parameters (block-order depth 4 and 8, 24 experts, 30 minutes). They are not evidence for larger models, and the planner's block-order and T = 2 recommendations are extrapolated to every size it plans.
  • The Web UI probe for auto-tune sees a few steps of a 4-layer model; it cannot rank step sizes that differ by less than the measured noise.

Tests

  • tests/test_recurrent_order.py, tests/test_later_topk.py, tests/test_later_groups.py and tests/test_iteration_cluster.py cover config validation, the MODEL_FIELDS round trip, Torch against MLX, cache decode against a full forward, exact self-speculative decoding, and gradient reach into every recurrent layer, router and the controller. tests/test_mlx_checkpointing.py adds block-order and graded-top-k cases to the recompute and compile matrices.
  • tests/test_dt_clip.py compares the MLX and Torch step-size gradient at 0.4999, 0.5, 0.5001, 0.75, 0.9999, 1.0 and 1.0001; on the 0.3.2a3 kernel it fails at 0.5 and 1.0.
  • tests/test_bench_paging.py checks that mlx_runner(paging=...) builds one page per expert for expert and the training clusters for cluster and auto. The prefill/decode split and warm-up exclusion in tools/bench_residency.py run inside its main() and need a tokenizer, so they are verified by running the tool on the benchmark checkpoints, not by a unit test.
  • tests/test_dt_sigmoid.py covers config validation of recurrent_dt_param, the sigmoid logit initialisation (step 0 gives recurrent_dt_init at 0.05, 0.0625, 0.25, 0.5 and 1.0), the nonzero sigmoid gradient where the clamp gradient is zero, Torch against MLX at four logits, and recommended_dt_init (depth * dt = 2, the 0.5 cap at depths 1-3). tests/test_depth_jitter.py covers one depth draw per optimizer step under gradient accumulation, with and without mx.compile.
  • tests/test_general_tokenizer.py (7 tests) checks the bundled tokenizer's files, its 32,768 vocabulary and control ids 0-63, native against Hugging Face encodings and the round trip, that the report names no local path and records derived_from, the planner's new default (32,768 and the bundled path, with sources), that an explicit vocabulary adds no path, and that the legacy default is still 6,400.
  • tests/test_autotune_loop.py (8 tests) covers the probe's real tokens (bundled vocabulary, None for unreadable or tiny data, deterministic batches that advance with the step), the loop knobs being sampled only when searched, the best-config patch, a probe trial that builds a looped block-order model and draws its batches from real tokens, a probe without loop knobs staying a plain depth-1 model on random ids, and the Web UI search choices and four-language labels. tests/test_planner_loop.py (7 tests) covers the planner's block order, sigmoid and horizon init at depths 2, 4, 6 and 8, the plain defaults at depth 1, that the planned config builds a valid model, and that the schema's recommended values equal what the planner emits.
  • Full suite on the development tree: 732 passed, 4 skipped (three browser tests need playwright, which is not installed here; one needs the v2 reference checkpoint, which is not in the repository). ruff check . is clean.