Skip to content

fix(models/glm4-moe): fuse separate switch_mlp gate_proj/up_proj experts at load#620

Merged
inureyes merged 1 commit into
mainfrom
fix/issue-530-fuse-switch-mlp-gate-up
Jul 2, 2026
Merged

fix(models/glm4-moe): fuse separate switch_mlp gate_proj/up_proj experts at load#620
inureyes merged 1 commit into
mainfrom
fix/issue-530-fuse-switch-mlp-gate-up

Conversation

@inureyes

@inureyes inureyes commented Jul 2, 2026

Copy link
Copy Markdown
Member

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 with Failed 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_proj expert 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_proj and switch_mlp.up_proj. Confirmed by reading the checkpoint's 12 *.safetensors headers directly (its model.safetensors.index.json is stale/pre-quant and was ignored):

language_model.model.layers.1.mlp.switch_mlp.gate_proj.weight  U32  [128, 1408, 512]
language_model.model.layers.1.mlp.switch_mlp.gate_proj.scales  BF16 [128, 1408, 64]
language_model.model.layers.1.mlp.switch_mlp.gate_proj.biases  BF16 [128, 1408, 64]
language_model.model.layers.1.mlp.switch_mlp.up_proj.{weight,scales,biases}   (same shapes)
language_model.model.layers.1.mlp.switch_mlp.down_proj.weight  U32  [128, 4096, 176]
language_model.model.layers.1.mlp.shared_experts.{gate,up,down}_proj.{weight,scales,biases}
language_model.model.layers.1.mlp.gate.weight                  BF16 [128, 4096]
language_model.model.layers.1.mlp.gate.e_score_correction_bias BF16 [128]

Fix

SwitchGLU::from_weights (in src/models/glm4_moe.rs) now prefers a pre-fused gate_up_proj when present and otherwise fuses the separate tensors at load via the new SwitchLinear::from_separate_gate_up.

  • Axis / order: concatenate gate-then-up along the expert output-feature axis, i.e. axis 1 of [num_experts, out_features, in_features] ([128, 1408, ...] + [128, 1408, ...] -> [128, 2816, ...]). Gate is the first half, matching the [0:I) = gate, [I:2I) = up split already done in SwitchGLU::forward.
  • Quant-aware: 4/8-bit weights pack along the input axis (last dim: 4096 -> 512 U32), so concatenating along the output axis never splits a packed word. The matching per-output-row scales/biases concatenate on the same axis. No dequantization is needed; the packed U32 words are carried through intact.
  • down_proj, shared_experts, and the router gate (+ 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's SwitchGLU consumes a fused gate_up_proj, so this fuses rather than splits. The fusion order/axis were cross-checked against references/mlx-vlm/mlx_vlm/models/glm4v_moe/language.py, whose sanitize stacks gate_proj/up_proj/down_proj separately (the layout this checkpoint ships).

Because the fix lives in the shared Glm4Moe::from_weights path, both the standalone Glm4MoeModel and the Glm4vMoeTextModel VLM backbone accept the separate layout, and the pre-fused layout still loads (backward compatible).

What changed

  • src/models/glm4_moe.rs
    • SwitchLinear::from_separate_gate_up: quant-aware fusion of separate gate_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_ref helper (borrow without copying, for in-place fusion).
    • Two unit tests: fuses_separate_switch_mlp_gate_up_experts (separate -> fused, asserts [E, 2I, ...] shapes for weight/scales/biases and unchanged down_proj) and loads_prefused_switch_mlp_gate_up_experts (backward compat).

Test plan

  • cargo test --lib glm4 (9 passed, including the 2 new tests)
  • cargo check --lib --tests
  • cargo 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).

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.
@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions area:models Model architectures, weights, loading, metadata priority:high High priority status:review Under review labels Jul 2, 2026
@inureyes inureyes merged commit 4b117a7 into main Jul 2, 2026
5 checks passed
@inureyes inureyes deleted the fix/issue-530-fuse-switch-mlp-gate-up branch July 2, 2026 13:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:models Model architectures, weights, loading, metadata priority:high High priority status:review Under review type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant