feat(models): add Ant Group Ling / Bailing MoE (bailing_moe) text model support #838

Summary

Add support for the Ant Group Ling / Bailing MoE text model family (bailing_moe) to the mlxcel Rust inference runtime, porting the architecture from mlx-lm. Bailing MoE is a DeepSeek-style mixture-of-experts model with routed experts, an always-on shared expert, and optional grouped top-k routing, so much of it maps onto mlxcel's existing DeepSeek-V2/V3 MoE machinery. The routing details differ from DeepSeek in several places where a blind reuse is silently wrong; those places are called out below.

This body was revised after a pre-implementation pass that read the real checkpoint at models/ling-lite-1.5 and the upstream reference line by line. Three config fields the original text named do not exist in the checkpoint, one mapping instruction was a divergence from upstream rather than a port of it, and the shared-expert shape was described incorrectly.

Upstream reference

• mlx-lm implementation: https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/bailing_moe.py

Validation checkpoint

inclusionAI/Ling-lite-1.5, present locally at models/ling-lite-1.5. No mlx-community conversion is required: the upstream weights are raw bf16 SafeTensors and load directly.

• 16.8B total / 2.75B activated, 28 layers, hidden_size 2048, intermediate_size 1408, moe_intermediate_size 1408, 64 routed experts, num_experts_per_tok 6, num_shared_experts 2, 16 query heads / 4 KV heads (head_dim 128), vocab_size 126464, rope_theta 600000, rms_norm_eps 1e-6, tie_word_embeddings false, rope_scaling null.
• Roughly 33.6 GB across 4 shards, entirely bf16. Validation needs a machine with headroom for that plus KV cache.
• Real weight keys per layer: model.layers.N.attention.query_key_value.weight (fused QKV, [3072, 2048] = (16 + 2*4) * 128), model.layers.N.attention.dense.weight, model.layers.N.mlp.gate.weight (router, [64, 2048]), model.layers.N.mlp.experts.{0..63}.{gate_proj,up_proj,down_proj}.weight (per-expert, not pre-stacked), model.layers.N.mlp.shared_experts.{gate_proj,up_proj,down_proj}.weight, plus input_layernorm and post_attention_layernorm. Top level: model.word_embeddings.weight, model.norm.weight, lm_head.weight. Note word_embeddings, not embed_tokens.
• use_bias and use_qkv_bias are both false, so no projection carries a bias tensor, and no expert_bias tensor exists in the checkpoint.

Config fields the checkpoint does not contain

Three fields the original issue text named are absent from models/ling-lite-1.5/config.json: n_group, moe_router_enable_expert_bias, and moe_router_enable_routed_scaling. They must be optional with serde defaults matching upstream, or the config will not parse at all. The full set of upstream-optional routing fields:

score_function defaults to "softmax". DeepSeek-V3's router is sigmoid. Since this issue directs the implementer to reuse the DeepSeek-V2/V3 MoE machinery, this is exactly the place where reuse goes wrong: taking the DeepSeek gate as-is changes the routing distribution on every token while producing perfectly finite output. The score function must be read from config and both "softmax" and "sigmoid" supported.

Also absent from the checkpoint and therefore default-bearing: use_qk_norm (false), partial_rotary_factor (1.0), rotary_dim (None). With those defaults the rope dim is the full head_dim of 128 and there are no QK norm weights to load.

What the validation checkpoint actually exercises

This matters for how much a green real-model run is worth:

• n_group defaults to 1, so the grouped top-k branch never runs.
• moe_router_enable_expert_bias defaults to false, so the bias add never runs and there is no expert_bias tensor to load.
• first_k_dense_replace is 0, so all 28 layers are MoE and there is no dense prefix to exercise for this checkpoint.

Ling-lite-1.5 therefore exercises only the plain softmax top-k path with a shared expert. The grouped-routing branch, the expert-bias branch, the sigmoid score function, and the dense-prefix path all still have to be implemented, and each needs unit-test coverage against a synthetic config, because no checkpoint available here reaches them. A token-exact real-model comparison will not validate those branches. Do not read a passing real-model run as full coverage.

Router: selection uses biased scores, weights come from unbiased scores

group_expert_select computes scores = softmax(gates.astype(float32)) (or sigmoid when score_function == "sigmoid"), saves orig_scores = scores, then adds the correction bias to scores only. The grouped-group masking, when n_group > 1, also applies to that biased copy. Top-k indices come from argpartition on the biased scores, and the returned weights are then gathered from orig_scores via take_along_axis. Applying the bias to the weights as well as to the selection is a silent correctness bug: output stays finite and plausible while every routed contribution is misweighted. This is the same selection-only-bias pattern DeepSeek-V3 and ERNIE-4.5 already implement in this tree; follow those.

Two further details from the same function:

• norm_topk_prob is conditional and has an epsilon. Upstream applies it only when top_k > 1, and the denominator is scores.sum(-1, keepdims=True) + 1e-20, not a bare sum.
• The score function runs in float32 and the result is cast back to the input dtype only at the very end, after normalization and scaling.

moe_router_enable_routed_scaling is dead code upstream

BailingMoeGate stores self.enable_routed_scaling = args.moe_router_enable_routed_scaling and never reads it again. __call__ passes routed_scaling_factor unconditionally, and group_expert_select ends with an unconditional scores = scores * routed_scaling_factor. The flag has no effect upstream.

The original instruction in this issue to "map moe_router_enable_routed_scaling onto routed scaling" would therefore be a deliberate divergence from the reference rather than a port of it. It is unobservable on Ling-lite-1.5, where routed_scaling_factor defaults to 1.0, but it is observable on any checkpoint that sets a non-unit factor together with the flag false.

Pick one behavior explicitly: either mirror upstream (always scale, flag parsed and ignored) or honor the flag (scale only when true). State which one was chosen and why in the PR body, and pin it with a unit test so a later reader cannot mistake it for an oversight.

Shared experts are one wide MLP, not N separate experts

shared_dim = moe_shared_expert_intermediate_size or moe_intermediate_size, and the shared MLP is constructed with intermediate_size = shared_dim * num_shared_experts. For Ling-lite-1.5 that is a single MLP of width 1408 * 2 = 2816, which the real tensors confirm: mlp.shared_experts.gate_proj.weight is [2816, 2048] and mlp.shared_experts.down_proj.weight is [2048, 2816]. There is no per-shared-expert axis anywhere in the checkpoint.

It is built only when num_shared_experts > 0 and moe_router_enable_shared_expert, it is added to the routed mixture at a fixed weight of 1.0 (out = out + self.shared_experts(x)), and it is never packed into the switch tensors. Do not stack it into SwitchLinear and do not give it a routing weight.

Vendor config keys upstream partly ignores

config.json sets norm_head: false and norm_softmax: false.

• norm_head is live upstream, contrary to what a quick read suggests. mlx-lm's ModelArgs declares it, Model stores it, and sanitize L2-normalizes lm_head.weight along axis 0 in float32 with a +1e-7 epsilon before casting back to the original dtype. The vendored modeling_bailing_moe.py does the same normalization at inference time, once, then clears the flag. Implement it in sanitize to match, and unit-test it: it is a few lines, and a checkpoint that sets it true would otherwise produce wrong logits with no error.
• norm_softmax is dead everywhere. mlx-lm declares it in ModelArgs and never reads it, and it does not appear anywhere in the vendored modeling_bailing_moe.py; configuration_bailing_moe.py does not even name it, so it survives only via **kwargs. A checkpoint setting it true is asking for behavior no released code defines. Reject that config at load rather than parsing it and silently ignoring it.

Weight remap

Upstream sanitize does two things per layer at or after first_k_dense_replace:

1. Stacks mlp.experts.{e}.{gate_proj,down_proj,up_proj}.{weight,scales,biases} over e in 0..num_experts into mlp.switch_mlp.{m}.{k}. The scales and biases variants exist only for quantized checkpoints; the raw bf16 checkpoint has weight only.
2. Renames the router mlp.gate.weight to mlp.gate.gate_proj.weight, and mlp.gate.bias to mlp.gate.gate_proj.bias when present.

Watch the naming collision in step 2: mlp.gate.* is the router, while gate_proj is the SwiGLU gate of an expert MLP, and after the rename the router itself lives at mlp.gate.gate_proj.weight. A remap rule that matches on gate_proj without anchoring the prefix will swallow the router weight.

Upstream also carries a quant_predicate that quantizes mlp.gate.gate_proj at group_size 64, bits 8 while leaving the rest at the default, and a cast_predicate that excludes expert_bias from dtype casting. Both matter only for quantized conversions, but note them if one is validated later.

Architecture notes

• DeepSeek-style MoE block:
  • num_experts routed experts (64 here) plus a single shared MLP sized moe_intermediate_size * num_shared_experts, gated on moe_router_enable_shared_expert.
  • Optional n_group grouped top-k routing (experts partitioned into n_group groups, topk_group groups kept, the rest zeroed before top-k). Inactive by default and inactive on this checkpoint.
  • Optional router expert bias (moe_router_enable_expert_bias), the same role as DeepSeek's e_score_correction_bias, applied to selection only.
  • Routed scaling by routed_scaling_factor, unconditional upstream (see above).
  • norm_topk_prob, applied only when top_k > 1, with a 1e-20 epsilon in the denominator.
  • first_k_dense_replace layers before the MoE layers use a plain BailingMoeMLP at intermediate_size. Zero on this checkpoint.
• Standard RMSNorm + RoPE attention, no exotic attention variant, but the QKV projection is fused into one query_key_value tensor split at [q_size, q_size + kv_size], and the output projection is named dense.
• Optional QK RMSNorm per head (use_qk_norm), applied after the head reshape and before RoPE. Off on this checkpoint.

Implementation plan

Much of this maps onto mlxcel's existing DeepSeek-V2/V3 MoE path (src/models/deepseek_v2.rs, src/models/deepseek_v3.rs), which already implements grouped top-k routing, shared experts, the e_score_correction_bias router bias, and routed scaling. Reuse that machinery where the semantics actually match, and do not assume they match by default:

• Experts: reuse src/models/switch_layers.rs (SwitchLinear / fused experts) as the DeepSeek path already does.
• Routing: reuse the grouped top-k router and shared-expert combination, but parameterize the score function (softmax default, not DeepSeek's sigmoid), keep the bias selection-only, keep the top_k > 1 guard and the 1e-20 epsilon on norm_topk_prob, and decide the moe_router_enable_routed_scaling question explicitly.
• Shared expert: one MLP at shared_dim * num_shared_experts, added at weight 1.0, outside the switch tensors.
• Attention: RMSNorm + RoPE with a fused query_key_value split and a dense output projection, reusing existing shared layers.
• Main new work is config field mapping (including the optional-with-defaults set above), the Bailing weight-key naming in sanitize (expert stacking plus the router rename), and the norm_head lm_head normalization.
• New model file under src/models/ (for example src/models/bailing_moe.rs), following docs/adding-models.md.

Validation

A token-exact comparison against the mlx-lm reference requires --no-chat-template on both sides. The reference harness feeds the raw prompt, while this checkpoint does ship a chat template (embedded in tokenizer_config.json as a <role>HUMAN</role> style template; there is no separate chat_template.json). Without the flag the two sides tokenize different strings and diverge on the first token for reasons unrelated to the port.

Captured reference for the raw prompt What is the capital of France?: generation stops at EOS after 11 token ids.

• text: \n\n\nThe capital of France is Paris.
• ids: 198 198 198 678 7706 300 11406 341 13997 13 126081

The trailing 126081 is <|endoftext|>, which config.json sets as both eos_token_id and pad_token_id.

Then run a second time without --no-chat-template. That run is the real-usage sanity check; it is expected to produce different text, and it should be coherent and terminate.

Defensive requirements

Every port already merged in this series (#833, #834, #836, #837) produced findings of the same class, so these are requirements rather than suggestions. The underlying hazard: a config-derived value that passes Rust-side validation but violates an undocumented precondition of an MLX C++ entry point terminates the process at the first inference rather than failing at load, because those bridge functions return UniquePtr rather than Result and a C++ throw crossing the cxx bridge is an uncatchable std::terminate. Rejecting at load is the only containment.

• Gather and slice bounds must be taken from real tensor shapes, not from config fields. MLX take wraps negative indices but does not range-check positive ones, so an out-of-range expert index or slice offset silently reads adjacent memory and the result reaches the logits as plausible numbers.
• Magnitude validation on the architecture scalars (hidden_size, num_hidden_layers, num_attention_heads, num_key_value_heads, num_experts, num_experts_per_tok, num_shared_experts, intermediate_size, moe_intermediate_size, moe_shared_expert_intermediate_size, vocab_size, max_position_embeddings, n_group, topk_group, first_k_dense_replace) before any Vec::with_capacity, truncating as i32 cast, or use as a divisor. Zero checks go before divisibility checks: 0.is_multiple_of(0) is true in Rust, so a divisibility check on its own lets a zero divisor through. src/models/gpt_neox.rs has the established shape of this (MAX_* ceilings plus a validate() returning Result<(), String>).
• Routing-specific bounds beyond magnitude: num_experts_per_tok at least 1 and no greater than num_experts; when n_group > 1, num_experts divisible by n_group, and topk_group between 1 and n_group. Upstream computes k = n_group - topk_group and calls argpartition(kth=k - 1), which is out of range as soon as topk_group >= n_group.
• RoPE dims validated even, positive, and no larger than the last axis before reaching fast_rope. partial_rotary_factor is a float, so a hostile value can be NaN or infinite; the saturating as i32 cast turns those into 0 and i32::MAX respectively, and both must be rejected with a message naming the actual field. rope_theta must be finite and positive, since RoPE takes its logarithm.
• rms_norm_eps validated finite and positive.
• Reject norm_softmax: true rather than ignoring it, per the section above.
• The tree-wide unvalidated quantization block is tracked separately as fix(core): an unvalidated config.json quantization block reaches MLX and terminates the process #929. This port only needs to avoid adding new exposure, which the per-model validate_quantization pattern already present in src/models/gpt_neox.rs achieves at negligible cost.

Touchpoints & acceptance criteria

Delivered and closed. Implemented by PR #946, merged as a3c823611; mlxcel arch now reports the family. Test coverage: 42 unit tests in src/models/bailing_moe_tests.rs, mutation-checked by reverting five behaviors and observing exactly five paired failures. Real-checkpoint validation: models/ling-lite-1.5 reproduces the captured reference ids 198 198 198 678 7706 300 11406 341 13997 13 and stops at EOS 126081, run with --no-chat-template on both sides so the reference and mlxcel see the same raw prompt. Note the gate does NOT cover grouped routing or the expert-bias branch, because n_group defaults to 1 and moe_router_enable_expert_bias defaults false, so those two rest entirely on the unit tests. The unticked boxes below are the original pre-implementation plan and were not maintained during the work; treat the merged PR and its review thread as the record of what shipped.

• Config struct + serde parsing for the Bailing MoE config, with n_group, topk_group, moe_router_enable_expert_bias, moe_router_enable_routed_scaling, routed_scaling_factor, score_function, moe_shared_expert_intermediate_size, moe_router_enable_shared_expert, use_qk_norm, partial_rotary_factor, and rotary_dim all optional and defaulted to the upstream values in the table above. A parse test over the verbatim models/ling-lite-1.5/config.json proves none of them is required.
• score_function read from config and both "softmax" and "sigmoid" implemented, with a unit test that the default is softmax and that an unknown value is rejected rather than silently falling back.
• from_weights constructor wiring the MoE block (routed experts, single wide shared MLP added at weight 1.0, grouped top-k router) and RMSNorm/RoPE attention with the fused query_key_value split.
• Router implemented with selection on the biased scores and weights gathered from the unbiased scores, norm_topk_prob applied only when top_k > 1 with the 1e-20 epsilon, and the float32 score computation cast back at the end. Unit tests pin each of these three separately.
• moe_router_enable_routed_scaling behavior chosen explicitly (mirror upstream or honor the flag), stated in the PR body, and pinned by a unit test.
• Grouped routing (n_group > 1) and the expert-bias path implemented and unit-tested against synthetic configs, since the validation checkpoint reaches neither.
• first_k_dense_replace dense-prefix path implemented and unit-tested against a synthetic config, since the validation checkpoint sets it to 0.
• sanitize stacking mlp.experts.{e}.{gate_proj,down_proj,up_proj}.{weight,scales,biases} into the fused SwitchLinear layout and renaming mlp.gate.weight / mlp.gate.bias to mlp.gate.gate_proj.*, with a test that the router weight is not captured by the expert gate_proj rule.
• norm_head implemented in sanitize (axis-0 L2 normalization of lm_head.weight in float32 with +1e-7, cast back) and unit-tested.
• norm_softmax: true rejected at load with a message explaining that no released implementation defines it.
• Config validation covering the magnitude, routing, RoPE, and norm-eps requirements in the Defensive requirements section, with a rejection test per rule.
• Arm added to src/models/detection.rs for the bailing_moe architecture / model_type.
• Registration in src/model_metadata.rs via for_each_model_registration!.
• TP / distributed arch-string entry if applicable.
• Unit tests in the model's _tests.rs (config parse, sanitize expert remap and router rename, router selection, forward shape) in addition to the targeted tests above.
• docs/supported-models.md updated to list Bailing MoE / Ling.
• Real-checkpoint validation: mlxcel arch reports the architecture, and ./target/release/mlxcel generate on models/ling-lite-1.5 reproduces the captured reference token ids exactly with --no-chat-template.
• Second real-checkpoint run without --no-chat-template produces coherent output and terminates.

Correction: the real-checkpoint validation step above originally read mlxcel list, which lists downloaded checkpoints in the local model store. The architecture registry is mlxcel arch, and that is what confirms the binary knows the family.

Done only when the model loads and generates from the real checkpoint (per docs/adding-models.md).

Effort

MEDIUM. The expert machinery and the general MoE shape already exist on the DeepSeek-V2/V3 path, but the reuse is not free: the score function default differs from DeepSeek, the shared expert is one wide MLP rather than N experts, several config fields are absent from real checkpoints and must be defaulted, and two of the routing branches cannot be validated against any available checkpoint and need synthetic-config tests instead. Add the usual registration, detection, validation, test, and doc touchpoints.