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

Honor quantization_config.load_in_4bit for Gemma-4 MoE swap

Browse source 
The Gemma-4 MoE per-expert Linear4bit swap previously gated only on the
positional load_in_4bit argument. loader.py forwards load_in_4bit=False
to FastBaseModel.from_pretrained whenever the caller supplies a
quantization_config (BitsAndBytesConfig), so callers that opt in via
UNSLOTH_GEMMA4_MOE_4BIT=1 plus BitsAndBytesConfig(load_in_4bit=True)
silently bypassed the swap. The adjacent guardrail already normalises
load_in_4bit from quantization_config; the swap gate now does the same
and sources bnb_4bit_compute_dtype from quantization_config when no
local bnb_config is built.

The except branch around the swap also previously stated "Falling back
to BF16 experts", which misrepresents the model state when the helper
fails partway through (already-swapped Gemma4TextExperts modules stay
in 4-bit; only the remainder remain BF16). The warning now counts the
modules marked _unsloth_gemma4_moe_4bit_swapped and reports the partial
state, advising a reload to recover a uniform state.

The comment above the fused-Parameter dels in gemma4_moe_4bit.py
overstated the swap's memory bound; rephrased to describe the actual
per-module peak (fused BF16 plus accumulated per-expert nf4).
This commit is contained in:
Daniel Han 2026-05-16 14:59:43 +00:00
commit 41f8792e7c
2 changed files with 36 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

5
unsloth/models/gemma4_moe_4bit.py
View file

@ -193,8 +193,9 @@ def swap_gemma4_experts_to_per_expert_linear4bit(
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.
# The fused BF16 gate_up_proj / down_proj stay live through the loop
# above; per-module peak is fused BF16 + accumulated per-expert nf4.
# They are released here, before attaching the ModuleLists.
del module.gate_up_proj
del module.down_proj

41
unsloth/models/vision.py
View file

@ -1071,7 +1071,18 @@ class FastBaseModel:
# 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:
_user_qcfg = kwargs.get("quantization_config", None)
if isinstance(_user_qcfg, dict):
_qcfg_4bit = bool(_user_qcfg.get("load_in_4bit", False))
_qcfg_dtype = _user_qcfg.get("bnb_4bit_compute_dtype", None)
elif _user_qcfg is not None:
_qcfg_4bit = bool(getattr(_user_qcfg, "load_in_4bit", False))
_qcfg_dtype = getattr(_user_qcfg, "bnb_4bit_compute_dtype", None)
else:
_qcfg_4bit = False
_qcfg_dtype = None
_effective_load_in_4bit = bool(load_in_4bit) or _qcfg_4bit
if _effective_load_in_4bit and not full_finetuning:
try:
from unsloth.models.gemma4_moe_4bit import (
is_gemma4_moe_4bit_enabled,
@ -1079,13 +1090,15 @@ class FastBaseModel:
)
if is_gemma4_moe_4bit_enabled():
if bnb_config is not None:
_compute_dtype = bnb_config.bnb_4bit_compute_dtype
elif _qcfg_dtype is not None:
_compute_dtype = _qcfg_dtype
else:
_compute_dtype = torch.bfloat16
_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
),
compute_dtype = _compute_dtype,
)
if _swapped > 0:
print(
@ -1095,10 +1108,22 @@ class FastBaseModel:
f"https://github.com/unslothai/unsloth/issues/5344)."
)
except Exception as _e:
_partial = sum(
1 for _m in model.modules()
if getattr(_m, "_unsloth_gemma4_moe_4bit_swapped", False)
)
if _partial:
_state = (
f"{_partial} Gemma4TextExperts module(s) are "
f"already in 4-bit; remaining modules stay BF16. "
f"Reload the model to recover a uniform state."
)
else:
_state = "Falling back to BF16 experts."
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.",
f"{type(_e).__name__}: {_e}. {_state} "
f"Unset UNSLOTH_GEMMA4_MOE_4BIT to silence.",
stacklevel = 2,
)