[MoE] Fix Qwen-family MoE LoRA extractor shape mismatch

[lordx64]

Summary

Fix a shape-mismatch bug in the shared _make_qwen_moe_lora_extractor (used by Qwen3-MoE, Qwen3.5/3.6 MoE, and Qwen3-Next) that causes torch._grouped_mm to raise RuntimeError: contraction dimension of mat_a and mat_b must match during the first training step on MoE+LoRA.

Reproduction

# train.py (UV script)
from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Qwen3.6-35B-A3B",
    max_seq_length=4096,
    load_in_4bit=True,
)
model = FastLanguageModel.get_peft_model(
    model, r=16, lora_alpha=16,
    target_modules=["q_proj","k_proj","v_proj","o_proj",
                    "gate_proj","up_proj","down_proj"],  # includes expert LoRA
    use_gradient_checkpointing="unsloth",
)
# trainer.train() fails at step 0 with:
#   File ".../moe_utils.py", line 131, in _grouped_mm_with_backward_fix
#       return torch._grouped_mm(inputs, weight, offs=offsets)
#   RuntimeError: contraction dimension of mat_a and mat_b must match

Root cause

_make_qwen_moe_lora_extractor has three code paths that all construct first_weight from weight_B (PEFT's lora_B, shape (out_dim, E*R)):

# qwen3_moe.py:77-94 (broken)
if param_name == "down_proj" and ...:
    first_weight = weight_B.view(dim_B, E, R).permute(1, 0, 2)   # (E, out_dim, R)
    second_weight = weight_A.view(E, R, dim_A)                    # (E, R, in_dim)
    return ...
elif param_name == "gate_up_proj" and ...:
    first_weight = weight_B.view(dim_B, E, R).permute(1, 0, 2)   # (E, out_dim, R)
    ...
if dim_B == hidden_dim:
    first_weight = weight_B.view(dim_B, E, R).permute(1, 0, 2)   # (E, out_dim, R)
    ...

But forward_native_grouped_mm calls:

# moe_utils.py:862
lora_out = _grouped_mm_with_backward_fix(permuted_input, first_weight, offsets)

with permuted_input of shape (N, in_dim). torch._grouped_mm requires first_weight.shape[-2] == in_dim, but the extractor gives out_dim there.

On Qwen3.6-35B-A3B (gate_up_proj stored as (E, 2*I, H) in modeling_qwen3_5_moe.py:723), in_dim = H and out_dim = 2*I — and these are unequal for every MoE config, so the mismatch fires on every run with expert LoRA.

The fallback at the bottom of the function (lines 102-106 pre-fix) was already correct, but none of the earlier branches fell through to it for Qwen3-MoE-style models.

Fix

Drop the broken branches. The correct mapping — identical to:

• the default extractor in moe_utils.py::_extract_lora_from_wrapper (lines 421-426),
• the working Qwen3-VL-MoE extractor in qwen3_vl_moe.py::_qwen3_vl_lora_extractor (lines 275-279),

— is format-independent:

first_weight  = weight_A.view(E, R, in_dim).permute(0, 2, 1).contiguous()  # (E, in_dim, R)
second_weight = weight_B.view(out_dim, E, R).permute(1, 2, 0).contiguous() # (E, R, out_dim)

PEFT LoRA weights have fixed shape relative to the linear's in/out dims — they don't depend on whether the base weights are stored "standard" (E, out, in) or "transposed" (E, in, out). That distinction is handled by preprocess_weight for the base-weight path, not the LoRA path.

Test plan

• Reproduce the contraction dimension crash on unsloth/Qwen3.6-35B-A3B with target_modules that include gate_proj, up_proj, down_proj (expert LoRA)
• Verify the fix makes the LoRA forward path through torch._grouped_mm succeed — training progresses past step 0 into the expected memory-bound regime on 1× A100-80GB
• Regression check on unsloth/Qwen3-30B-A3B (128 experts, same code path) — expected to continue working
• Regression check on Qwen3-Next — same extractor, same mapping applies

Happy to add a unit test for the extractor shape contract if useful; let me know if you have a preferred fixture pattern.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request simplifies the _qwen_moe_lora_extractor function for Qwen-family MoE models by removing complex conditional logic that previously checked for specific parameter names and dimensions. The implementation now uses a unified approach to reshape and permute LoRA weights (weight_A and weight_B) into the format required for grouped_mm, regardless of the base weight storage layout. The docstring has also been updated to reflect this streamlined logic and explain the mapping from PEFT's dimension layout. I have no feedback to provide as no review comments were submitted.

[Datta0]

The changes do look good to me. But can you do a comparison run with transformers baseline vs unsloth and report_to wandb.
Once we can confirm that the losses are closely matching, we can merge the PR

[danielhanchen]

Possible duplicate of a trusted maintainer's PR. This PR looks like it solves the same underlying problem as unslothai/unsloth-zoo#495 by @Datta0 (trusted maintainer).

Directly patches qwen3_moe extractor layout for Qwen3Next/Qwen3.5, matching the same shape-mismatch fix path.

Canonical PR summary: The PR updates Unsloth Zoo compatibility patches for newer transformers v5, adding new Qwen3Next and Qwen3.5 MoE temporary patches, relaxing Gemma3/CSM/DeepSeekV3/GLM4-MoE signatures, and fixing MoE backend, 4-bit loading, and compiler behavior.

Other related PRs:

• unslothai/unsloth-zoo#453 — @danielhanchen [trusted] — Fixes MoE weight unwrapping and reshaping in the same LoRA extraction path, including grouped GEMM layout handling.

The auto-review is still running against this PR — reviewers will factor in the canonical above. If this PR is genuinely different, call out the delta in the review discussion so the maintainer can decide which to merge.

[danielhanchen]

Auto-review verdict: Approved

Replaces the branchy Qwen-family MoE LoRA extractor with the canonical PEFT shape mapping (weight_A -> (E, in_dim, R), weight_B -> (E, R, out_dim)), eliminating the torch._grouped_mm contraction-dimension crash that reproduced on every gate_up_proj and down_proj step for Qwen3-MoE, Qwen3.5/3.6, and Qwen3-Next. The fix matches the existing default extractor in moe_utils.py and qwen3_vl_moe.py, and is verified by CPU shape and numerical-equivalence tests.

Reason: Correct minimal fix; all 11 reviewers approved across 2 iterations with no blocking findings.

[danielhanchen]

Thanks so much this works!