fix(models/glm4-moe): fuse separate switch_mlp gate_proj/up_proj experts at load#620
Merged
Merged
Conversation
Follow-up hardening for #530: the GLM-4V MoE port loaded its routed experts from a fused `switch_mlp.gate_up_proj`, but real GLM-4(.5)V MoE checkpoints (e.g. glm-4.5v-4bit) ship the already expert-stacked experts as separate `switch_mlp.gate_proj` and `switch_mlp.up_proj` tensors, so loading failed with `Weight not found: model.layers.1.mlp.switch_mlp.gate_up_proj.weight`. `SwitchGLU::from_weights` now prefers a pre-fused `gate_up_proj` when present and otherwise fuses the separate tensors at load. Fusion concatenates gate-then-up along the expert output-feature axis (axis 1 of `[num_experts, out_features, in_features]`), matching the `[0:I) = gate, [I:2I) = up` split in `SwitchGLU::forward`. It is quant-aware: 4/8-bit weights pack along the input axis, so concatenating along the output axis leaves the packed words intact and the matching per-output-row `scales`/`biases` concatenate on the same axis, with no dequantization. `down_proj`, `shared_experts`, and the router `gate` (+ `e_score_correction_bias`) already match the checkpoint and are loaded unchanged. The fix flows through the shared `Glm4Moe::from_weights`, so both the standalone `Glm4MoeModel` and the `Glm4vMoeTextModel` VLM backbone accept the separate layout. Adds unit tests feeding the real separate-expert key set through the fusion and asserting the produced `gate_up_proj` weight/scales/biases shapes ([E, 2I, ...]), plus a backward-compat test for the pre-fused layout.
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.
Follow-up hardening for #530: the merged GLM-4V MoE port fails to load its real checkpoint.
mlxcel generate -m /path/to/glm-4.5v-4bit ...aborts withFailed to load GLM-4V MoE text model: Weight not found: model.layers.1.mlp.switch_mlp.gate_up_proj.weight.Root cause
The GLM-4 MoE text loader expected a single fused
switch_mlp.gate_up_projexpert stack. Real GLM-4(.5)V MoE checkpoints (converted by mlx-vlm) instead ship the routed experts as two already expert-stacked tensors,switch_mlp.gate_projandswitch_mlp.up_proj. Confirmed by reading the checkpoint's 12*.safetensorsheaders directly (itsmodel.safetensors.index.jsonis stale/pre-quant and was ignored):Fix
SwitchGLU::from_weights(insrc/models/glm4_moe.rs) now prefers a pre-fusedgate_up_projwhen present and otherwise fuses the separate tensors at load via the newSwitchLinear::from_separate_gate_up.[num_experts, out_features, in_features]([128, 1408, ...]+[128, 1408, ...]->[128, 2816, ...]). Gate is the first half, matching the[0:I) = gate, [I:2I) = upsplit already done inSwitchGLU::forward.scales/biasesconcatenate on the same axis. No dequantization is needed; the packed U32 words are carried through intact.down_proj,shared_experts, and the routergate(+e_score_correction_bias) already match the checkpoint and load unchanged.Analogous in spirit to the qwen3.5-moe per-expert -> fused fix in #519 / PR #587 (
src/models/qwen3_5.rs), but in the opposite direction: GLM'sSwitchGLUconsumes a fusedgate_up_proj, so this fuses rather than splits. The fusion order/axis were cross-checked againstreferences/mlx-vlm/mlx_vlm/models/glm4v_moe/language.py, whosesanitizestacksgate_proj/up_proj/down_projseparately (the layout this checkpoint ships).Because the fix lives in the shared
Glm4Moe::from_weightspath, both the standaloneGlm4MoeModeland theGlm4vMoeTextModelVLM backbone accept the separate layout, and the pre-fused layout still loads (backward compatible).What changed
src/models/glm4_moe.rsSwitchLinear::from_separate_gate_up: quant-aware fusion of separategate_proj/up_proj(weight + scales + biases) along axis 1.SwitchGLU::from_weights: use the pre-fused stack if present, else fuse the separate tensors.get_weight_refhelper (borrow without copying, for in-place fusion).fuses_separate_switch_mlp_gate_up_experts(separate -> fused, asserts[E, 2I, ...]shapes for weight/scales/biases and unchangeddown_proj) andloads_prefused_switch_mlp_gate_up_experts(backward compat).Test plan
cargo test --lib glm4(9 passed, including the 2 new tests)cargo check --lib --testscargo clippy --lib --tests -- -D warnings(clean)cargo fmt --check(clean)Real-model GPU validation is intentionally deferred to the orchestrator (this is a 106B MoE; checkpoint was inspected read-only, not loaded here).