`Lfm2MoeSparseMoeBlock` routing diverges from HF: softmax instead of sigmoid, `expert_bias` in the weights, and `routed_scaling_factor` dropped

[hijohnnylin]

• AI Usage Notice: We used Opus to find this bug. The issue is written up by Opus. Code to repro is copy-pastable.

---

mlx_lm/models/lfm2_moe.py routes its experts three ways differently from transformers, so every
MoE layer of an LFM2-MoE checkpoint returns a different tensor than the reference implementation.
There is no error and the model generates plausible text, which is why this went unnoticed.

It is not only the weighting: because softmax compresses the gate's range while expert_bias is a
fixed offset, adding the bias afterwards reorders the top-k. In the repro below half the tokens
are routed to a different set of experts than transformers selects, and the weights differ by up
to 34% even where both implementations pick the same one.

Versions: mlx-lm 0.31.3, mlx 0.32.2, transformers 5.17.0, Apple M3 Max, fp32 both sides.
Checkpoint: LiquidAI/LFM2-8B-A1B (use_expert_bias=true, norm_topk_prob=true,
routed_scaling_factor=1.0, num_experts=32, num_experts_per_tok=4, num_dense_layers=2).

The three divergences

Lfm2MoeSparseMoeBlock.__call__ in mlx_lm/models/lfm2_moe.py:

gates = self.gate(x).astype(mx.float32)
gates = mx.softmax(gates, axis=-1)          # (1) HF applies sigmoid

if self.use_expert_bias:
    gates += self.expert_bias               # (2) folded into the weights, not only the selection

k = self.top_k
inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
scores = mx.take_along_axis(gates, inds, axis=-1)
if self.norm_topk_prob:
    scores /= mx.sum(scores, axis=-1, keepdims=True) + 1e-20
scores = scores.astype(x.dtype)             # (3) no routed_scaling_factor

Lfm2MoeTopKRouter.forward in transformers/models/lfm2_moe/modeling_lfm2_moe.py:

router_logits = F.linear(hidden_states, self.weight)
routing_weights = router_logits.sigmoid()                     # (1) sigmoid, per-expert
if self.use_expert_bias:
    scores_for_routing = routing_weights + expert_bias        # (2) bias selects...
    _, selected_experts = torch.topk(scores_for_routing, k=self.top_k, dim=-1)
    routing_weights = torch.gather(routing_weights, dim=1, index=selected_experts)
    #                 ^ ...but the weights gathered are the UNBIASED sigmoid scores
else:
    routing_weights, selected_experts = torch.topk(routing_weights, k=self.top_k, dim=-1)

if self.norm_topk_prob:
    routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-6)
routing_weights = routing_weights * self.routed_scaling_factor  # (3)

1. softmax vs sigmoid. The dominant difference. sigmoid is a per-expert gate whose top-k
   values do not sum to 1 before normalization; softmax couples all 32 experts, so normalizing the
   top 4 gives a different distribution even when the selection is identical.
2. expert_bias reaches the weights. HF uses the bias only to pick experts and then gathers the
   unbiased scores. mlx-lm adds it into gates before both the argpartition and the
   take_along_axis, so the returned weights carry the bias.
3. routed_scaling_factor is not read. Harmless at 1.0, as on this checkpoint, but silently
   wrong on any LFM2-MoE config that sets it.

(The normalization epsilon also differs, 1e-20 against HF's 1e-6. Too small to matter; noted for
completeness.)

Runnable repro

No checkpoint download: it builds both libraries' own classes at a small hidden size, copies the
same random gate weights and expert_bias into each, and feeds both the same input. So the only
thing that differs is the routing arithmetic.

pip install "mlx-lm>=0.31.3" "transformers>=5" torch numpy
python repro_lfm2.py

"""Repro: mlx-lm's LFM2-MoE routing disagrees with transformers'.

Runs both libraries' own classes over the same weights and the same input, so the
difference is the routing arithmetic and nothing else. No checkpoint download.

    pip install "mlx-lm>=0.31.3" "transformers>=5" torch numpy
    python repro_lfm2.py
"""

import mlx.core as mx
import numpy as np
import torch
from mlx_lm.models.lfm2_moe import Lfm2MoeSparseMoeBlock
from mlx_lm.models.lfm2_moe import ModelArgs as MlxArgs
from transformers.models.lfm2_moe.configuration_lfm2_moe import Lfm2MoeConfig
from transformers.models.lfm2_moe.modeling_lfm2_moe import Lfm2MoeTopKRouter

# LiquidAI/LFM2-8B-A1B's own routing settings, at a small hidden size.
HIDDEN, N_EXPERTS, TOP_K = 64, 32, 4
SEED, N_TOKENS = 0, 64

# 0.005 because that is the scale the checkpoint actually ships: its 22 `expert_bias` tensors run
# about +/-0.01. The bias magnitude decides how often the *selection* moves, so a larger one here
# would overstate the effect -- at bias 0 the selection is identical (softmax and sigmoid are both
# monotonic in the logits) and only the weights differ.
rng = np.random.default_rng(SEED)
gate_w = rng.normal(0, 0.05, (N_EXPERTS, HIDDEN)).astype(np.float32)
expert_bias = rng.normal(0, 0.005, (N_EXPERTS,)).astype(np.float32)
x = rng.normal(0, 1.0, (N_TOKENS, HIDDEN)).astype(np.float32)


def mlx_routing() -> tuple:
    """mlx-lm's routing weights, read out of its own block."""
    args = MlxArgs(
        model_type="lfm2_moe",
        vocab_size=32,
        hidden_size=HIDDEN,
        intermediate_size=8,
        moe_intermediate_size=8,
        num_hidden_layers=1,
        num_experts=N_EXPERTS,
        num_experts_per_tok=TOP_K,
        norm_topk_prob=True,
        num_attention_heads=4,
        num_key_value_heads=4,
        max_position_embeddings=128,
        use_expert_bias=True,
        num_dense_layers=0,
        norm_eps=1e-5,
        conv_bias=False,
        conv_L_cache=3,
        layer_types=["full_attention"],
    )
    block = Lfm2MoeSparseMoeBlock(args)
    block.gate.weight = mx.array(gate_w)
    block.expert_bias = mx.array(expert_bias)

    # The block returns the *combined* expert output, so the weights it applied are read back by
    # re-running the lines above `switch_mlp` -- copied verbatim from mlx_lm/models/lfm2_moe.py.
    gates = block.gate(mx.array(x)).astype(mx.float32)
    gates = mx.softmax(gates, axis=-1)
    gates += block.expert_bias
    inds = mx.argpartition(gates, kth=-TOP_K, axis=-1)[..., -TOP_K:]
    scores = mx.take_along_axis(gates, inds, axis=-1)
    scores /= mx.sum(scores, axis=-1, keepdims=True) + 1e-20
    return inds, scores


def hf_routing() -> tuple:
    """transformers' routing weights, straight out of `Lfm2MoeTopKRouter.forward`."""
    config = Lfm2MoeConfig(
        hidden_size=HIDDEN,
        num_experts=N_EXPERTS,
        num_experts_per_tok=TOP_K,
        norm_topk_prob=True,
        use_expert_bias=True,
        routed_scaling_factor=1.0,
    )
    router = Lfm2MoeTopKRouter(config)
    with torch.no_grad():
        router.weight.copy_(torch.from_numpy(gate_w))
        _, weights, experts = router(torch.from_numpy(x), torch.from_numpy(expert_bias))
    return experts, weights


mlx_inds, mlx_scores = mlx_routing()
hf_inds, hf_scores = hf_routing()

mlx_i = np.asarray(mlx_inds, dtype=np.int64)
mlx_s = np.asarray(mlx_scores, dtype=np.float32)
hf_i = hf_inds.numpy()
hf_s = hf_scores.numpy().astype(np.float32)

# Sort each row by expert id so nothing depends on top-k ordering.
order_mlx, order_hf = np.argsort(mlx_i, axis=-1), np.argsort(hf_i, axis=-1)
mlx_i, mlx_s = np.take_along_axis(mlx_i, order_mlx, -1), np.take_along_axis(mlx_s, order_mlx, -1)
hf_i, hf_s = np.take_along_axis(hf_i, order_hf, -1), np.take_along_axis(hf_s, order_hf, -1)

same = int((mlx_i == hf_i).all(axis=-1).sum())
print(f"1. SELECTION: tokens routed to the same {TOP_K} experts: {same}/{N_TOKENS}")
print()
print("   token 0:")
print(f"     experts (mlx-lm)      {mlx_i[0].tolist()}")
print(f"     experts (transformers) {hf_i[0].tolist()}")
print(f"     weights (mlx-lm)      {np.round(mlx_s[0], 5).tolist()}")
print(f"     weights (transformers) {np.round(hf_s[0], 5).tolist()}")
print()

# Weights compared per (token, expert id), over the intersection of the two selections -- matched by
# expert rather than by position, so this measures the weight and never the ordering.
rels = []
for t in range(N_TOKENS):
    a = dict(zip(mlx_i[t].tolist(), mlx_s[t].tolist()))
    b = dict(zip(hf_i[t].tolist(), hf_s[t].tolist()))
    rels += [abs(a[e] - b[e]) / max(abs(b[e]), 1e-12) for e in set(a) & set(b)]

rels = np.array(rels)
print(f"2. WEIGHTS: over the {len(rels)} (token, expert) pairs both implementations chose,")
print("   matched by expert id so ordering cannot flatter the result:")
print(f"     max relative difference  {rels.max():.4f}  ({rels.max() * 100:.1f}%)")
print(f"     mean                     {rels.mean():.4f}  ({rels.mean() * 100:.1f}%)")
print()
print("Both follow from the same three lines: mlx-lm softmaxes the gate where transformers")
print("applies sigmoid, and folds `expert_bias` into the gathered weights where transformers")
print("uses it only to select.")
print()
print("The weights differ from the softmax/sigmoid swap alone. The *selection* moves because")
print("softmax compresses the gate's range while the bias is a fixed offset, so adding it after")
print("a softmax reorders the top-k where adding it after a sigmoid does not. Set expert_bias to")
print("zero and the selection agrees on every token while the weights still differ.")

Output on mlx-lm 0.31.3 / mlx 0.32.2 / transformers 5.17.0, Apple M3 Max:

1. SELECTION: tokens routed to the same 4 experts: 32/64

   token 0:
     experts (mlx-lm)      [1, 2, 7, 11]
     experts (transformers) [1, 2, 7, 15]
     weights (mlx-lm)      [0.225490003824234, 0.26295000314712524, 0.28485000133514404, 0.2267100065946579]
     weights (transformers) [0.2439900040626526, 0.25641998648643494, 0.2579300105571747, 0.2416599988937378]

2. WEIGHTS: over the 223 (token, expert) pairs both implementations chose,
   matched by expert id so ordering cannot flatter the result:
     max relative difference  0.3438  (34.4%)
     mean                     0.0765  (7.7%)

The selection itself moves, not only the weights -- half the tokens are routed to a different
set of four experts. softmax compresses the gate's range while expert_bias is a fixed offset, so
adding the bias after a softmax reorders the top-k where adding it after a sigmoid does not.

The two divergences are separable, and the repro shows which does what. Set expert_bias to zero
and the selection agrees on every token while the weights still differ by up to 31% -- that is
divergence (1) on its own. The bias magnitude then decides how often the selection moves: at the
checkpoint's own scale (its 22 expert_bias tensors run about +/-0.01) it is about half the tokens,
which is what the repro is calibrated to.

Measured effect end to end

Capturing activations from both implementations over the same token ids and comparing them
point-by-point, the first MoE layer is where they part. num_dense_layers=2, so layers 0–1 are
dense MLPs and layer 2 is the first MoE layer:

layer	point	agreement
2	router_logits	passes — the gate Linear is correct, and its input is still clean
2	mlp_out	fails — same input, same logits, different block output
12	mlp_out	rel 2.89e-01, cos 0.9599
23	mlp_out	rel 2.82e-01, cos 0.9664
—	final_norm	rel 3.12e-01, cos 0.9512

The pattern is what pins it to the routing rather than to anything else: at layer 2 the router's
input and its logits both agree and only the block's output differs. That error then enters the
residual stream, so by layer 12 the hidden state feeding the router is itself drifted and
router_logits starts failing too (rel 2.2e-02 to 3.0e-02) — downstream accumulation, not a second
bug.

23 of 903 compared activations disagree; all 23 are on this checkpoint. Every other MoE family
tested (Granite-MoE, DeepSeek-V2-Lite, Phi-3) agrees to rel < 7.1e-03.

Suggested fix

Drop-in replacement for the body of Lfm2MoeSparseMoeBlock.__call__ above switch_mlp:

k = self.top_k
routing_weights = mx.sigmoid(self.gate(x).astype(mx.float32))

# The bias selects; the weights gathered are the unbiased sigmoid scores.
selector = routing_weights + self.expert_bias if self.use_expert_bias else routing_weights
inds = mx.argpartition(selector, kth=-k, axis=-1)[..., -k:]

scores = mx.take_along_axis(routing_weights, inds, axis=-1)
if self.norm_topk_prob:
    scores = scores / (mx.sum(scores, axis=-1, keepdims=True) + 1e-6)
scores = (scores * self.routed_scaling_factor).astype(x.dtype)

routed_scaling_factor needs adding to ModelArgs (HF defaults it to 1.0).

I am happy to open a PR if that is useful.