Ornith-1.5-35B-A3B-BigBang-MTP-zh-mlx-4bit

A Chinese MTP head retrained on the deployment's own traffic

Identical to KaedeTai/Ornith-1.5-35B-A3B-BigBang-MTP-mlx-4bit except for the 44 tensors of the MTP draft head. The trunk and the vision tower are byte-identical; only the drafting changes, and speculative decoding is verify-then-accept, so a draft head can only affect speed.

The head now accepts 3.4 points more of its own drafts than the stock one, on held-out replies from the deployment it serves.

110 held-out conversations, split by question so no question appears in both halves. "Accepted" means the head's token equals the trunk's own next token — the thing the runtime checks — measured teacher-forced with no server involved, so thermals and GPU contention cannot move it.

An earlier revision of this repo was a negative result. This is what changed.

That head was trained on Traditional Chinese web and news corpora, gained 3.5 points there, and roughly nothing where it mattered. The correction seemed obvious — train on the model's own replies — but replies to invented prompts still measured ten points above real traffic. Three changes, each forced by a measurement:

1. The corpus is the deployment's own shape. Real customer questions and real member contexts (built by the assistant's own context builder against actual profiles, plus 75 synthetic members in exactly that layout for prompt diversity), under the two real system prompts, answered by the model itself at temperature 0.7 — the setting production actually calls with. 1,624 conversations, 235k assistant tokens.

2. The loss is masked to assistant positions. The head drafts only during decode; system and user tokens are prefilled and never drafted. They stay in the sample because the hidden states are built from them, but they contribute no gradient. A Dr. Willy system block is ~1,000 tokens, so without the mask most of the capacity went to positions that can never pay.

3. The expert stack is frozen. The head is 844.6M parameters, overwhelmingly a 256-expert MoE, and only 8 experts see gradient per token. A corpus this size cannot converge it. Training fc, attention, the shared expert and the norms — 39.3M parameters, 4.7% of the head — is where the gain came from.

What did not work, and one that is worth knowing

Learning rate dominates everything. Every rate at or above 3e-5 destroyed the head; the customer-service stratum fell to 55%. Training the whole head at 1e-5, which is what the previous attempt did, made it worse within 80 steps. The working range is 3e-6 to 5e-6 with the experts frozen, and the run plateaus by step 600.

The theoretically correct objective was not better. oMLX's MTP patch documents Leviathan-style stochastic acceptance — min(1, p_target / p_draft) with residual sampling — and production calls at temperature 0.7. Under that rule the expected accepted drafts per position is the overlap sum_x min(p_t(x), p_d(x)), not argmax agreement, so a loss on that overlap should beat cross-entropy. Given the same budget the two finished 0.2 points apart, which this sample cannot resolve. The gain came from what was trained, not from the loss it was trained with.

4-bit rounds away a third of it. The same head packaged at 4-bit measures 74.22% overall against 75.26% at 8-bit — the quantisation noise is comparable to the improvement, so the improvement is what disappears. The head in this repo is 8-bit while the trunk stays 4-bit, which is why model-00004.safetensors is 3.11 GB rather than 2.69 GB and config.json carries 12 per-module overrides. The head is 4-6% of decode compute, so it does not have to share the trunk's width.

What this is not

It is not yet a measured speed number. Draft accuracy at depth 1 is not throughput: the previous head in this repo gained 2.6 points at depth 1 and gave all of it back at depths 2 and 3, finishing at 0.999x tokens per backbone forward. Until the same paired runtime measurement is repeated here, the honest claim is the one above — the head drafts better — and nothing about tokens per second.

If you serve this in oMLX, read this first

Adding a model directory is not enough, and the failure is silent. oMLX applies its own defaults for any setting a model's entry does not carry, and one of those defaults is enable_thinking: true. A hand-written entry in ~/.omlx/model_settings.json is not picked up — with the entry written by hand, mtp_enabled: false still logged three MTP activations in the same request, and enable_thinking: false never took effect. The model then answers in English chain-of-thought instead of Chinese and runs to the token cap, and nothing in the logs says why.

Let oMLX register the model itself (open it once in the app) before comparing anything. Two model directories that differ only in this setting are not a comparison of their heads: chain-of-thought is far more predictable than counselling prose, and the arm in thinking mode will look 26% better on tokens per backbone forward for reasons that have nothing to do with drafting.

Then set mtp_enabled: true and expect:

Speculative backend selected: Lightning MTP (model_type=qwen3_5_moe, active)

Do not set vlm_mtp_enabled — it is a different feature (an external Gemma4 assistant drafter) and oMLX treats it as mutually exclusive with mtp_enabled, silently disabling this head.

One tensor, one file

If you build a variant by hard-linking the trunk shards and writing a new head to a separate file, the head will exist twice — it also lives inside one of those shards. Loaders glob *.safetensors and weights.update(), so whichever sorts last wins, silently, when the shapes match. swap_head.py here rewrites the shard that holds the head and asserts afterwards that no tensor name appears in two files. The build in this repo passes that check: 2,134 tensors, no duplicates, all 44 head tensors in model-00004.safetensors.

Method, for anyone repeating this

• targets are the trunk's own argmax, so the objective is what speculative decoding pays for
• trunk frozen and 4-bit throughout; only the head moves, and only 39.3M of it
• AdamW, lr 5e-6 cosine with 60 warmup steps, batch 4 by accumulation
• checkpoint selected on held-out real replies, not on training loss
• ~3,400 tokens/s on an M5 Max; the run is under 30 minutes

gen_deploy2.py builds the corpus, prep_deploy.py turns it into masked samples, train_deploy.py trains, swap_head.py packages, head_scale.py checks that the trained head's distribution is still as peaked as the trunk it drafts for — all in this repo.

The single most useful thing learned across three attempts: measure acceptance on what your deployment actually emits before deciding there is a gap. Web corpora said Chinese was 25 points behind code. Replies to invented prompts said 10. Real traffic says something different again, and only the last one predicts anything.

Limitations

• One machine (M5 Max, 128 GB), MLX 0.32.0, one model family, one deployment.
• The corpus is one product's traffic — a health consultation assistant and a Traditional Chinese customer-service line. A head trained this way is tuned to that; on unrelated traffic the stock head may well be the better choice.
• 9,618 scored positions is enough to resolve three points and not enough to resolve a few tenths.
• No end-to-end throughput measurement yet. See "What this is not".

Credits

As the base repo, plus oMLX for Lightning MTP and for a patch source readable enough to find the acceptance rule in.