The MoE bnb-4bit quantization shipped for unslothai/unsloth#4032 has an asymmetry:
- Quantize-coverage is generic:
moe_utils_bnb4bit.py patches transformers' Bnb4bit quantizer
to convert any module with fused gate_up_proj/down_proj nn.Parameters to packed Params4bit.
- Route-coverage is a list: the dequantize-at-forward dispatch only reaches architectures with a
per-arch patch file (qwen3*, glm4, gemma4_moe, deepseek_v3, lfm2, ernie4_5, gpt_oss…).
Any fused-experts arch outside the list loads "successfully" — experts quantize and VRAM drops —
then transformers' own experts-kernel integration receives packed (N, 1) uint8 storage and
crashes on the first forward. OLMoE is a shipping example.
Repro
unsloth 2026.6.9 / unsloth-zoo 2026.6.7 / transformers 5.5.0, torch 2.11.0+cu128 (outside
unsloth's <2.11 pin — the quantize-without-routing gap is coverage logic, version-independent by
code reading):
from unsloth import FastLanguageModel
from bitsandbytes.nn import Params4bit
model, tok = FastLanguageModel.from_pretrained(
"allenai/OLMoE-1B-7B-0924", max_seq_length=512, load_in_4bit=True,
)
print(type(model.model.layers[0].mlp.experts.gate_up_proj))
# <class 'bitsandbytes.nn.modules.Params4bit'> -- all 16 layers quantize (3.0 GiB of experts;
# with UNSLOTH_ENABLE_LOGGING=1: "Unsloth: Prepared ... for BNB 4-bit quantization (shapes: ...)" x16)
ids = tok("hello", return_tensors="pt").input_ids.cuda()
model(ids)
# File "transformers/integrations/moe.py", line 186, in _grouped_mm_fallback
# output = torch.zeros(input.size(0), weight.size(2), ...) # weight is the packed (N, 1) uint8
# IndexError: Dimension out of range (expected to be in range of [-2, 1], but got 2)Which error surfaces depends on the dispatch branch in transformers/integrations/moe.py, but the
root cause is the same — packed 2-D uint8 storage arrives where a 3-D float expert tensor is
expected. The full-model trace above comes via _grouped_mm_fallback because the experts forward
runs under torch.compile (unsloth compiles the MoE block) and _can_use_grouped_mm rejects
non-bf16 weights under dynamo, so pinned-torch users hit this branch too. A plain eager call on a
grouped_mm-capable device takes the other branch and dies with RuntimeError: Expected mat_a to be Float32, BFloat16 or Float16 matrix, got Byte — mat_a names the hidden states because
transformers casts the input to the packed weight's uint8 dtype (input.to(weight.dtype)) before
the kernel's dtype check fires.
LoRA adapters attach fine (both unsloth get_peft_model and raw-peft target_parameters wrap
the expert params), but the training forward then hits the same crash. There is also a LoRA-bypass
side of the same asymmetry worth noting: zoo's global ParamWrapper.forward patch skips PEFT's
parametrization for any MoE-shaped module and stashes LoRA weights (_unsloth_lora_*) that only
the zoo-routed backends consume — on unrouted archs the expert adapters would be silently ignored
at forward time even absent the crash.
bf16 OLMoE (no load_in_4bit) works — verified module-level at parity with stock math on a
4-layer truncation — so this is specifically the quantize-without-routing state.
Suggested fix
Either gate the quantizer match on the same arch predicate the forward routing uses (fail loud into
bf16 experts + a warning, like before the fix), or route generically (the transformers-v5 experts
interface (hidden_states, top_k_index, top_k_weights) is uniform). If OLMoE gets routed through
the zoo backend, note that its gate_up is square (2·1024 == hidden 2048) — it then lands
directly in the preprocess_weight transposition ambiguity (#849) unless that is fixed
first.
Found while auditing the #4032 fix; happy to test candidate patches on the affected hardware.
The MoE bnb-4bit quantization shipped for unslothai/unsloth#4032 has an asymmetry:
moe_utils_bnb4bit.pypatches transformers' Bnb4bit quantizerto convert any module with fused
gate_up_proj/down_projnn.Parameters to packedParams4bit.per-arch patch file (qwen3*, glm4, gemma4_moe, deepseek_v3, lfm2, ernie4_5, gpt_oss…).
Any fused-experts arch outside the list loads "successfully" — experts quantize and VRAM drops —
then transformers' own experts-kernel integration receives packed
(N, 1)uint8 storage andcrashes on the first forward. OLMoE is a shipping example.
Repro
unsloth 2026.6.9 / unsloth-zoo 2026.6.7 / transformers 5.5.0, torch 2.11.0+cu128 (outside
unsloth's
<2.11pin — the quantize-without-routing gap is coverage logic, version-independent bycode reading):
Which error surfaces depends on the dispatch branch in
transformers/integrations/moe.py, but theroot cause is the same — packed 2-D uint8 storage arrives where a 3-D float expert tensor is
expected. The full-model trace above comes via
_grouped_mm_fallbackbecause the experts forwardruns under torch.compile (unsloth compiles the MoE block) and
_can_use_grouped_mmrejectsnon-bf16 weights under dynamo, so pinned-torch users hit this branch too. A plain eager call on a
grouped_mm-capable device takes the other branch and dies with
RuntimeError: Expected mat_a to be Float32, BFloat16 or Float16 matrix, got Byte— mat_a names the hidden states becausetransformers casts the input to the packed weight's uint8 dtype (
input.to(weight.dtype)) beforethe kernel's dtype check fires.
LoRA adapters attach fine (both unsloth
get_peft_modeland raw-pefttarget_parameterswrapthe expert params), but the training forward then hits the same crash. There is also a LoRA-bypass
side of the same asymmetry worth noting: zoo's global
ParamWrapper.forwardpatch skips PEFT'sparametrization for any MoE-shaped module and stashes LoRA weights (
_unsloth_lora_*) that onlythe zoo-routed backends consume — on unrouted archs the expert adapters would be silently ignored
at forward time even absent the crash.
bf16 OLMoE (no
load_in_4bit) works — verified module-level at parity with stock math on a4-layer truncation — so this is specifically the quantize-without-routing state.
Suggested fix
Either gate the quantizer match on the same arch predicate the forward routing uses (fail loud into
bf16 experts + a warning, like before the fix), or route generically (the transformers-v5 experts
interface
(hidden_states, top_k_index, top_k_weights)is uniform). If OLMoE gets routed throughthe zoo backend, note that its gate_up is square (2·1024 == hidden 2048) — it then lands
directly in the
preprocess_weighttransposition ambiguity (#849) unless that is fixedfirst.
Found while auditing the #4032 fix; happy to test candidate patches on the affected hardware.