diff --git a/src/mobius/models/deepseek.py b/src/mobius/models/deepseek.py index 329810eb..430f6aef 100644 --- a/src/mobius/models/deepseek.py +++ b/src/mobius/models/deepseek.py @@ -315,7 +315,14 @@ def __init__( super().__init__() assert config.num_local_experts is not None assert config.moe_intermediate_size is not None - self.moe = MoELayer(config, gate=gate) + # ``linear_class`` must reach the routed-expert dense-loop fallback + # too (not just the shared expert below): otherwise a quantized + # config quantizes every other linear in the model (attention, + # dense FFN, shared expert) but silently leaves the routed MoE + # experts as plain float `MatMul`, which both loses quantization + # and breaks the `fuse_dense_moe_to_qmoe` post-hoc rewrite (it + # only matches a quantized `MatMulNBits` dense-fallback pattern). + self.moe = MoELayer(config, gate=gate, linear_class=linear_class) # Shared expert uses moe_intermediate_size * n_shared_experts n_shared = config.n_shared_experts or 1 shared_intermediate = config.moe_intermediate_size * n_shared diff --git a/src/mobius/models/deepseek_test.py b/src/mobius/models/deepseek_test.py index 8cedd3e7..6bcfa131 100644 --- a/src/mobius/models/deepseek_test.py +++ b/src/mobius/models/deepseek_test.py @@ -97,6 +97,57 @@ def test_deepseek_dense_moe_fallback_matches_numpy_reference(): assert all(node.domain != "com.microsoft" for node in graph) +def test_deepseek_moe_ffn_linear_class_reaches_routed_experts(): + """``linear_class`` must reach the routed dense-loop experts, not just the shared expert. + + Regression test for a bug where ``_DeepSeekMoEFFN`` constructed its + ``MoELayer`` without forwarding ``linear_class``, so a quantized config + quantized attention/dense-FFN/shared-expert linears but silently left the + routed MoE experts as plain float ``MatMul`` -- losing quantization and + breaking the ``fuse_dense_moe_to_qmoe`` post-hoc rewrite, which only + matches a quantized ``MatMulNBits`` dense-fallback pattern. + """ + from mobius.components._common import Linear + + created: list[Linear] = [] + + class TrackingLinear(Linear): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + created.append(self) + + config = make_config( + hidden_size=4, + intermediate_size=8, + moe_intermediate_size=3, + num_local_experts=4, + num_experts_per_tok=2, + n_group=2, + topk_group=1, + n_shared_experts=1, + scoring_func="sigmoid", + topk_method="noaux_tc", + ) + module = _DeepSeekMoEFFN(config, DeepSeekMoEGate(config), linear_class=TrackingLinear) + + # The dense loop-over-experts fallback must have been built with + # TrackingLinear for every routed expert's projections. Check + # `module.moe.experts` directly (not the global `created` list) so this + # assertion can't pass merely because the shared expert below picked up + # the class -- that path already worked before this fix. + assert module.moe.experts is not None + routed_linears = [m for m in module.moe.experts.modules() if isinstance(m, Linear)] + assert len(routed_linears) > 0 + assert all(isinstance(m, TrackingLinear) for m in routed_linears) + # The shared expert must still be quantized too (it already worked + # before this fix; guard against a future regression there as well). + shared_linears = [m for m in module.shared_experts.modules() if isinstance(m, Linear)] + assert len(shared_linears) > 0 + assert all(isinstance(m, TrackingLinear) for m in shared_linears) + # Sanity: the tracking list saw both groups (routed + shared). + assert len(created) == len(routed_linears) + len(shared_linears) + + def test_noaux_tc_supports_single_expert_groups(): config = make_config( hidden_size=4,