Make Gemma-4 MoE swap honor full quantization_config and cover vLLM path
- Normalize string compute_dtype values from dict-style quantization_config
(e.g. {"bnb_4bit_compute_dtype": "bfloat16"}) into torch.dtype before
forwarding to bitsandbytes. A raw string would propagate to
bnb.nn.Linear4bit and crash the first forward with
"Invalid device string: 'bfloat16'".
- Forward bnb_4bit_quant_type from the user's BitsAndBytesConfig or dict
config into swap_gemma4_experts_to_per_expert_linear4bit so swapped
experts match the quantization type used by the rest of the model
(previously always nf4 even when the caller requested fp4).
- Wrap the swap and its warning into a local closure and invoke it from
both the regular auto_model.from_pretrained branch and the
fast_inference=True / convert_vllm_to_huggingface branch. The closure
is idempotent on non-Gemma-4 models, so the vLLM call is free when the
loaded model has no Gemma4TextExperts modules.
- Escalate partial-state swap failures: if the helper raises after one
or more Gemma4TextExperts modules were already committed to 4-bit,
the wrapper re-raises a RuntimeError instructing the caller to reload
the model. Previously a warning implied a clean BF16 fallback, which
is false when partial conversion has already occurred.
The closure is multi-line (long block) because it needs to capture the
already-resolved quantization parameters and be reusable across both
load paths; the alternative is duplicating the entire block.
This commit is contained in:
Daniel Han
parent
41f8792e7c
commit
09ed2b963d
1 changed files with 75 additions and 55 deletions
|
|
@ -1044,6 +1044,79 @@ class FastBaseModel:
|
|||
|
||||
verify_fp8_support_if_applicable(model_config)
|
||||
|
||||
# Resolve 4-bit + Gemma4 swap parameters once (shared by both load paths).
|
||||
_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)
|
||||
_qcfg_quant_type = _user_qcfg.get("bnb_4bit_quant_type", 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)
|
||||
_qcfg_quant_type = getattr(_user_qcfg, "bnb_4bit_quant_type", None)
|
||||
else:
|
||||
_qcfg_4bit = False
|
||||
_qcfg_dtype = None
|
||||
_qcfg_quant_type = None
|
||||
if isinstance(_qcfg_dtype, str):
|
||||
_qcfg_dtype_str = _qcfg_dtype.removeprefix("torch.")
|
||||
_maybe_dtype = getattr(torch, _qcfg_dtype_str, None)
|
||||
_qcfg_dtype = _maybe_dtype if isinstance(_maybe_dtype, torch.dtype) else None
|
||||
_effective_load_in_4bit = bool(load_in_4bit) or _qcfg_4bit
|
||||
|
||||
def _maybe_swap_gemma4_moe_4bit(_target_model):
|
||||
if not (_effective_load_in_4bit and not full_finetuning):
|
||||
return
|
||||
try:
|
||||
from unsloth.models.gemma4_moe_4bit import (
|
||||
is_gemma4_moe_4bit_enabled,
|
||||
swap_gemma4_experts_to_per_expert_linear4bit,
|
||||
)
|
||||
|
||||
if not is_gemma4_moe_4bit_enabled():
|
||||
return
|
||||
if bnb_config is not None:
|
||||
_compute_dtype = bnb_config.bnb_4bit_compute_dtype
|
||||
_quant_type = getattr(bnb_config, "bnb_4bit_quant_type", "nf4")
|
||||
else:
|
||||
_compute_dtype = (
|
||||
_qcfg_dtype if _qcfg_dtype is not None else torch.bfloat16
|
||||
)
|
||||
_quant_type = (
|
||||
_qcfg_quant_type if _qcfg_quant_type is not None else "nf4"
|
||||
)
|
||||
_swapped = swap_gemma4_experts_to_per_expert_linear4bit(
|
||||
_target_model,
|
||||
compute_dtype = _compute_dtype,
|
||||
quant_type = _quant_type,
|
||||
)
|
||||
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:
|
||||
_partial = sum(
|
||||
1 for _m in _target_model.modules()
|
||||
if getattr(_m, "_unsloth_gemma4_moe_4bit_swapped", False)
|
||||
)
|
||||
if _partial:
|
||||
raise RuntimeError(
|
||||
f"Unsloth: Gemma-4 MoE 4-bit swap failed after "
|
||||
f"converting {_partial} module(s); model is in a "
|
||||
f"mixed 4-bit/BF16 state. Reload the model to "
|
||||
f"recover. Original error: "
|
||||
f"{type(_e).__name__}: {_e}"
|
||||
) from _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,
|
||||
)
|
||||
|
||||
raise_handler = RaiseUninitialized()
|
||||
if not fast_inference:
|
||||
# Prevent load_in_fp8 from being forwarded into HF internal model loading
|
||||
|
|
@ -1071,61 +1144,7 @@ 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.
|
||||
_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,
|
||||
swap_gemma4_experts_to_per_expert_linear4bit,
|
||||
)
|
||||
|
||||
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 = _compute_dtype,
|
||||
)
|
||||
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:
|
||||
_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}. {_state} "
|
||||
f"Unset UNSLOTH_GEMMA4_MOE_4BIT to silence.",
|
||||
stacklevel = 2,
|
||||
)
|
||||
_maybe_swap_gemma4_moe_4bit(model)
|
||||
|
||||
# Guardrail: see _warn_if_quantization_silently_dropped + #5344.
|
||||
_warn_if_quantization_silently_dropped(
|
||||
|
|
@ -1249,6 +1268,7 @@ class FastBaseModel:
|
|||
bnb_config,
|
||||
is_vision_model = is_vlm,
|
||||
)
|
||||
_maybe_swap_gemma4_moe_4bit(model)
|
||||
model.vllm_engine = llm
|
||||
model.fast_generate = model.vllm_engine.generate
|
||||
model.fast_generate_batches = functools.partial(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue