ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

gemma-4 moe: per-expert Linear4bit swap so 26B-A4B fits at 4-bit (#5344)

Browse source 
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
This commit is contained in:
Daniel Han 2026-05-15 03:49:17 +00:00
commit ddf54efa5f
3 changed files with 390 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

142
tests/test_gemma4_moe_4bit_swap.py Normal file
View file

@ -0,0 +1,142 @@
"""Unit tests for the Gemma-4 MoE per-expert Linear4bit swap (#5344).
End-to-end correctness on the real 26B-A4B checkpoint requires a GPU + the
checkpoint on disk, so this file restricts itself to fast CPU-only tests
that exercise the swap helper's shape contract, idempotence, and gating
behaviour. The full repro (resident VRAM 46 GB -> 14.27 GB, cosine sim 0.994
vs BF16) is documented in the PR description.
"""
import importlib
import os
import torch
import torch.nn as nn
def _stub_gemma4_module():
"""Construct a stub Gemma4TextExperts-like module without importing
transformers' Gemma4Config (which would force a fresh transformers
download in CPU-only CI)."""
try:
from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts
except Exception:
return None
# The class init requires a config; build a tiny synthetic one and then
# overwrite the fused weights with shapes small enough for CPU tests.
class _StubConfig:
num_experts = 4
hidden_size = 16
moe_intermediate_size = 8
hidden_activation = "gelu_pytorch_tanh"
module = Gemma4TextExperts.__new__(Gemma4TextExperts)
nn.Module.__init__(module)
module.num_experts = _StubConfig.num_experts
module.hidden_dim = _StubConfig.hidden_size
module.intermediate_dim = _StubConfig.moe_intermediate_size
module.gate_up_proj = nn.Parameter(
torch.randn(
_StubConfig.num_experts,
2 * _StubConfig.moe_intermediate_size,
_StubConfig.hidden_size,
dtype = torch.bfloat16,
),
requires_grad = False,
)
module.down_proj = nn.Parameter(
torch.randn(
_StubConfig.num_experts,
_StubConfig.hidden_size,
_StubConfig.moe_intermediate_size,
dtype = torch.bfloat16,
),
requires_grad = False,
)
from transformers.activations import ACT2FN
module.act_fn = ACT2FN[_StubConfig.hidden_activation]
return module
def test_is_enabled_reads_env_var():
from unsloth.models import gemma4_moe_4bit
old = os.environ.pop("UNSLOTH_GEMMA4_MOE_4BIT", None)
try:
assert gemma4_moe_4bit.is_gemma4_moe_4bit_enabled() is False
os.environ["UNSLOTH_GEMMA4_MOE_4BIT"] = "1"
assert gemma4_moe_4bit.is_gemma4_moe_4bit_enabled() is True
os.environ["UNSLOTH_GEMMA4_MOE_4BIT"] = "0"
assert gemma4_moe_4bit.is_gemma4_moe_4bit_enabled() is False
finally:
if old is None:
os.environ.pop("UNSLOTH_GEMMA4_MOE_4BIT", None)
else:
os.environ["UNSLOTH_GEMMA4_MOE_4BIT"] = old
def test_swap_skips_models_without_gemma4_experts():
from unsloth.models.gemma4_moe_4bit import (
swap_gemma4_experts_to_per_expert_linear4bit,
)
model = nn.Sequential(nn.Linear(8, 8), nn.Linear(8, 8))
assert swap_gemma4_experts_to_per_expert_linear4bit(model) == 0
def test_swap_skips_when_transformers_lacks_gemma4():
"""If transformers does not expose Gemma4TextExperts, the helper must
return 0 without raising. We simulate the ImportError by patching."""
import unsloth.models.gemma4_moe_4bit as g4m
real_import = importlib.import_module
def _broken_import(name, *args, **kwargs):
if name == "transformers.models.gemma4.modeling_gemma4":
raise ImportError("simulated absence")
return real_import(name, *args, **kwargs)
try:
importlib.import_module = _broken_import
# Re-exercise via the public helper. It imports Gemma4TextExperts
# inside its try/except, so the simulated ImportError must yield 0.
model = nn.Sequential(nn.Linear(8, 8))
assert g4m.swap_gemma4_experts_to_per_expert_linear4bit(model) == 0
finally:
importlib.import_module = real_import
def test_swap_idempotent_on_stub_module_without_cuda():
"""On CPU we cannot exercise bnb (Linear4bit requires CUDA). Verify the
helper at least returns 0 for the no-bnb-experts case without raising,
and is idempotent across repeated calls."""
from unsloth.models.gemma4_moe_4bit import (
swap_gemma4_experts_to_per_expert_linear4bit,
)
if not torch.cuda.is_available():
# CPU-only: bnb's Linear4bit init would fail. Validate the model-walk
# path on an empty Sequential to confirm the helper is side-effect-free.
model = nn.Sequential(nn.Linear(4, 4))
assert swap_gemma4_experts_to_per_expert_linear4bit(model) == 0
assert swap_gemma4_experts_to_per_expert_linear4bit(model) == 0
return
# GPU path: build the stub and run a real swap.
module = _stub_gemma4_module()
if module is None:
return # transformers without gemma4 module: nothing to test
model = nn.Sequential(module.to("cuda"))
n1 = swap_gemma4_experts_to_per_expert_linear4bit(model)
n2 = swap_gemma4_experts_to_per_expert_linear4bit(model)
assert n1 == 1
assert n2 == 0 # idempotent: already-swapped modules are skipped
assert hasattr(module, "gate_up_proj_4bit")
assert hasattr(module, "down_proj_4bit")
assert len(module.gate_up_proj_4bit) == module.num_experts
assert len(module.down_proj_4bit) == module.num_experts
if __name__ == "__main__":
test_is_enabled_reads_env_var()
test_swap_skips_models_without_gemma4_experts()
test_swap_skips_when_transformers_lacks_gemma4()
test_swap_idempotent_on_stub_module_without_cuda()
print("All 4 swap tests passed.")

215
unsloth/models/gemma4_moe_4bit.py Normal file
View file

@ -0,0 +1,215 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Per-expert bitsandbytes Linear4bit swap for Gemma-4 MoE experts.
Refs: https://github.com/unslothai/unsloth/issues/5344
Gemma4TextExperts stores all experts as two fused 3D Parameters
(gate_up_proj, down_proj) shaped (num_experts, out_dim, in_dim) so that
torch._grouped_mm can dispatch a single grouped matmul per layer. The
fused storage is great for forward throughput but breaks bnb 4-bit
quantization: bnb.nn.Linear4bit only swaps nn.Linear instances, so the
fused 3D Parameters stay in BF16, defeating QLoRA VRAM savings.
This module swaps each Gemma4TextExperts module's fused weights for two
nn.ModuleList[Linear4bit] of length num_experts, and overrides forward to
dispatch per-expert. The trade-off is the loss of torch._grouped_mm
throughput in exchange for a ~4x reduction in expert weight VRAM
(measured on unsloth/gemma-4-26B-A4B-it: 46 GB -> 14.27 GB resident).
Gated on UNSLOTH_GEMMA4_MOE_4BIT (default off) and on load_in_4bit=True.
Default off until the matching per-expert LoRA path lands; opt in via
the env var if you want the VRAM win without QLoRA training.
"""
from __future__ import annotations
import os
from types import MethodType
import torch
import torch.nn as nn
__all__ = [
"is_gemma4_moe_4bit_enabled",
"swap_gemma4_experts_to_per_expert_linear4bit",
]
def is_gemma4_moe_4bit_enabled() -> bool:
"""Opt-in via UNSLOTH_GEMMA4_MOE_4BIT=1."""
return os.environ.get("UNSLOTH_GEMMA4_MOE_4BIT", "0") == "1"
def _per_expert_forward(
self,
hidden_states: torch.Tensor,
top_k_index: torch.Tensor,
top_k_weights: torch.Tensor,
) -> torch.Tensor:
"""Replacement Gemma4TextExperts.forward using per-expert Linear4bit.
Mirrors the reference forward in transformers.models.gemma4.modeling_gemma4
but dispatches through swapped nn.ModuleList[Linear4bit] modules instead
of nn.functional.linear on the fused 3D Parameters.
"""
final_hidden_states = torch.zeros_like(hidden_states)
with torch.no_grad():
expert_mask = torch.nn.functional.one_hot(
top_k_index, num_classes = self.num_experts,
)
expert_mask = expert_mask.permute(2, 1, 0)
expert_hit = torch.greater(expert_mask.sum(dim = (-1, -2)), 0).nonzero()
for expert_idx in expert_hit:
expert_idx = expert_idx[0]
if expert_idx == self.num_experts:
continue
top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
current_state = hidden_states[token_idx]
gate_up = self.gate_up_proj_4bit[expert_idx](current_state)
gate, up = gate_up.chunk(2, dim = -1)
current_hidden_states = self.act_fn(gate) * up
current_hidden_states = self.down_proj_4bit[expert_idx](
current_hidden_states
)
current_hidden_states = (
current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
)
final_hidden_states.index_add_(
0,
token_idx,
current_hidden_states.to(final_hidden_states.dtype),
)
return final_hidden_states
def _quantize_one_expert_to_linear4bit(
weight_2d: torch.Tensor,
compute_dtype: torch.dtype,
quant_type: str = "nf4",
):
"""Build a bnb.nn.Linear4bit from a single (out, in) BF16 weight slice.
Params4bit triggers on-the-fly quantization on .to(device).
"""
import bitsandbytes as bnb
out_features, in_features = weight_2d.shape
layer = bnb.nn.Linear4bit(
in_features,
out_features,
bias = False,
compute_dtype = compute_dtype,
quant_type = quant_type,
quant_storage = torch.uint8,
)
layer.weight = bnb.nn.Params4bit(
data = weight_2d.detach().clone().contiguous(),
requires_grad = False,
quant_type = quant_type,
)
return layer
def swap_gemma4_experts_to_per_expert_linear4bit(
model: nn.Module,
compute_dtype: torch.dtype = torch.bfloat16,
quant_type: str = "nf4",
verbose: bool = False,
) -> int:
"""Find every Gemma4TextExperts module in `model`, replace its fused 3D
weights with two nn.ModuleList[Linear4bit] (per-expert), and patch
forward to dispatch per-expert.
Returns the count of swapped modules. Zero if the model has no Gemma-4
MoE experts or if transformers does not expose Gemma4TextExperts.
"""
try:
from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts
except Exception:
return 0
swapped = 0
for module in model.modules():
if not isinstance(module, Gemma4TextExperts):
continue
# Idempotent: once swapped, the fused 3D Parameters are gone.
if hasattr(module, "_unsloth_gemma4_moe_4bit_swapped"):
continue
if not hasattr(module, "gate_up_proj") or not hasattr(module, "down_proj"):
continue
gate_up = module.gate_up_proj
down = module.down_proj
if not isinstance(gate_up, nn.Parameter) or gate_up.ndim != 3:
continue
if not isinstance(down, nn.Parameter) or down.ndim != 3:
continue
num_experts, two_intermediate, hidden = gate_up.shape
num_experts_d, hidden_d, intermediate = down.shape
if (
num_experts != num_experts_d
or hidden != hidden_d
or two_intermediate != 2 * intermediate
):
# Unrecognised layout: skip rather than risk corrupting weights.
if verbose:
print(
f"Unsloth: skipping Gemma4TextExperts swap due to "
f"unexpected shapes gate_up={tuple(gate_up.shape)} "
f"down={tuple(down.shape)}"
)
continue
device = gate_up.device
gate_up_list = nn.ModuleList()
down_list = nn.ModuleList()
for e in range(num_experts):
gu = _quantize_one_expert_to_linear4bit(
gate_up.data[e],
compute_dtype = compute_dtype,
quant_type = quant_type,
)
dp = _quantize_one_expert_to_linear4bit(
down.data[e],
compute_dtype = compute_dtype,
quant_type = quant_type,
)
gate_up_list.append(gu.to(device))
down_list.append(dp.to(device))
# Drop the fused Parameters before attaching the ModuleLists so peak
# VRAM during the swap stays bounded by one expert at a time.
del module.gate_up_proj
del module.down_proj
module.gate_up_proj_4bit = gate_up_list
module.down_proj_4bit = down_list
# Per-instance forward bind so other Gemma4TextExperts instances
# (e.g. in a sibling model) keep the class-level method.
module.forward = MethodType(_per_expert_forward, module)
module._unsloth_gemma4_moe_4bit_swapped = True
swapped += 1
if swapped > 0 and torch.cuda.is_available():
# Free the cached fused tensors so post-swap VRAM reflects 4-bit.
torch.cuda.empty_cache()
return swapped

33
unsloth/models/vision.py
View file

@ -1051,6 +1051,39 @@ class FastBaseModel:
# attn_implementation = attn_implementation,
**kwargs,
)
# Opt-in per-expert Linear4bit swap for Gemma-4 MoE checkpoints
# whose fused 3D expert weights bnb cannot quantize (#5344).
# Off by default; users enable via UNSLOTH_GEMMA4_MOE_4BIT=1.
if load_in_4bit and not full_finetuning:
try:
from unsloth.models.gemma4_moe_4bit import (
is_gemma4_moe_4bit_enabled,
swap_gemma4_experts_to_per_expert_linear4bit,
)
if is_gemma4_moe_4bit_enabled():
_swapped = swap_gemma4_experts_to_per_expert_linear4bit(
model,
compute_dtype = (
bnb_config.bnb_4bit_compute_dtype
if bnb_config is not None
else torch.bfloat16
),
)
if _swapped > 0:
print(
f"Unsloth: swapped {_swapped} "
f"Gemma4TextExperts module(s) to per-expert "
f"Linear4bit (see "
f"https://github.com/unslothai/unsloth/issues/5344)."
)
except Exception as _e:
warnings.warn(
f"Unsloth: Gemma-4 MoE 4-bit swap failed: "
f"{type(_e).__name__}: {_e}. Falling back to BF16 "
f"experts. Unset UNSLOTH_GEMMA4_MOE_4BIT to silence.",
stacklevel = 2,
)
# Guardrail: see _warn_if_quantization_silently_dropped + #5344.
_warn_if_quantization_silently_dropped(
model,