Fix DeepSeek MoE routed experts losing linear_class quantization - #562
Merged
Merged
Conversation
_DeepSeekMoEFFN.__init__ threaded linear_class into _SharedExpertMLP but not into the MoELayer that builds the routed-expert dense-loop fallback. A quantized config therefore quantized attention, dense FFN, and the shared expert, but silently left every routed MoE expert's gate/up/down projections as plain float MatMul -- losing quantization and breaking the fuse_dense_moe_to_qmoe post-hoc rewrite, which only pattern-matches a quantized MatMulNBits dense-fallback shape. Fix: pass linear_class through to MoELayer's routed-expert construction (one line). DeepSeek-V4's MoE FFN is unaffected -- it recomputes its own quantized class per-expert independently of MoELayer's linear_class and was already correct. Added test_deepseek_moe_ffn_linear_class_reaches_routed_experts, which isolates and asserts on module.moe.experts vs module.shared_experts separately (not just an aggregate creation count) so the regression is caught precisely. Verified: fails on pre-fix code, passes post-fix. Full suite (3249 passed, 9 skipped) and ruff check/format clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes a DeepSeek MoE quantization gap by ensuring the linear_class parameter is forwarded into the routed-expert dense fallback (MoELayer) within _DeepSeekMoEFFN, preventing routed experts from silently remaining unquantized and enabling downstream dense→QMoE fusion patterns to match as intended.
Changes:
- Thread
linear_classthrough_DeepSeekMoEFFNintoMoELayer(..., linear_class=...)so routed expert MLPs use the intended linear implementation. - Add a focused regression unit test that asserts the routed experts (and shared expert) are constructed using the provided
linear_class.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| src/mobius/models/deepseek.py | Forwards linear_class to the routed-expert MoELayer to avoid losing quantization in the dense fallback path. |
| src/mobius/models/deepseek_test.py | Adds a regression test verifying linear_class reaches both routed experts and the shared expert. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Root cause
_DeepSeekMoEFFN.__init__(src/mobius/models/deepseek.py) acceptedlinear_classbut only forwarded it to
_SharedExpertMLP, not to theMoELayerthat buildsthe routed-expert dense loop-over-experts fallback:
Consequence: a quantized config correctly quantized attention, dense FFN, and
the shared expert, but every routed MoE expert's gate/up/down projections
stayed plain float
MatMul. This both loses quantization for the routedexperts and breaks the
fuse_dense_moe_to_qmoepost-hoc rewrite, which onlypattern-matches a quantized
MatMulNBitsdense-fallback shape — so aquantized DeepSeek-V2/V3 config would silently fail to fuse into QMoE for its
routed experts.
DeepSeek-V4 is unaffected: its
DeepSeekV4MoEusesexpert_factoryto build_DeepSeekV4Expert, which independently recomputes its quantized class fromconfig.quantizationper-expert rather than relying onMoELayer'slinear_classforwarding — verified via review, no matching fix needed there.Fix
One line: thread
linear_classthrough to the routed-expertMoELayer:Tests
Added
test_deepseek_moe_ffn_linear_class_reaches_routed_expertsindeepseek_test.py, which isolates the assertion onmodule.moe.experts(routed) separately from
module.shared_experts(already worked before thisfix), rather than relying on an aggregate creation count — so it can't pass
for the wrong reason.
assert Falseat therouted-experts isinstance check) and passes post-fix.
PYTHONPATH=src python3 -m pytest src/→ 3249 passed, 9skipped, no failures.
ruff check/ruff format --checkon changed files: clean.Approve, no findings.
Scope
Deliberately minimal and focused: does not touch GLM-5.2 registration,
DeepSeek-V4 QMoE export, YaRN, or any other follow-up — those are tracked/
landed separately (#548, #550, #555, #559, #560).
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com