From 09ed2b963da7a997b9ccb40cfd9874bf4100e2d7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 16 May 2026 15:39:19 +0000 Subject: [PATCH] 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. --- unsloth/models/vision.py | 130 ++++++++++++++++++++++----------------- 1 file changed, 75 insertions(+), 55 deletions(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 69eed9743b..7d338630a4 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -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(