unsloth/gemma-4-26B-A4B-it loads at ~46 GB even with load_in_4bit=True
because Gemma4TextExperts stores experts as fused 3D nn.Parameter tensors
(gate_up_proj of shape (128, 1408, 2816), down_proj of (128, 2816, 704))
so torch._grouped_mm can dispatch a single grouped matmul per layer.
bitsandbytes' replace_with_bnb_linear only swaps nn.Linear instances, so
the fused expert weights stay BF16 and dominate the VRAM footprint.
This adds an opt-in helper that walks the loaded model, finds every
Gemma4TextExperts module, slices each fused (E, O, I) Parameter into E
individual bnb.nn.Linear4bit modules (per-expert), and patches forward
to dispatch per-expert instead of via torch._grouped_mm.
Trade-off:
- VRAM win: 46 GB -> 14.27 GB resident on unsloth/gemma-4-26B-A4B-it
(B200, transformers 5.5.0, single GPU). Linear4bit count 206 -> 7886.
Forward-pass cosine similarity vs BF16 reference is 0.994 on a fixed
prompt, i.e. standard QLoRA fidelity.
- Throughput loss: per-expert dispatch loses the grouped_mm speedup.
Acceptable for "model fits at 4-bit on a single GPU"; QLoRA training
still needs the matching per-expert LoRA path which is not in this PR.
Gated on UNSLOTH_GEMMA4_MOE_4BIT=1, default off until the per-expert
LoRA path lands (the swap renames gate_up_proj -> gate_up_proj_4bit
which would break unsloth_zoo's grouped_mm LoRA extractor as-is).
The renamed attributes also make the helper idempotent: re-entering it
sees `_unsloth_gemma4_moe_4bit_swapped` and no-ops, so multiple calls
across nested loaders are safe.
No regression on non-MoE checkpoints: the helper only touches modules
that are isinstance(Gemma4TextExperts) with the expected 3D shape.
Tests cover env-var gating, no-op behaviour on non-Gemma4 models, the
transformers-without-gemma4 ImportError path, and idempotence on a stub
Gemma4TextExperts module.
Refs #5344
ddf54efa5f