diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 742f5f57e0..8978e3be34 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -82,6 +82,7 @@ __all__ = [ "_get_inference_mode_context_manager", "hf_login", "is_moe_model", + "is_gemma4_shared_kv_model", "get_moe_target_parameters", "make_fast_generate_wrapper", ] @@ -3280,6 +3281,30 @@ def is_moe_model(model) -> bool: return num_experts is not None and num_experts > 0 +def is_gemma4_shared_kv_model(model) -> bool: + """ + Detect Gemma 4 E-series models (E2B / E4B) that share KV across layers. + + These have num_kv_shared_layers > 0, and in transformers 5.5.0 their + KV cache generation path is broken (huggingface/transformers#45242): + generating with a KV cache (use_cache=True / cache_implementation set) + yields different, garbled logits to the cache-free path. Gemma 4 31B / + 26B-A4B have num_kv_shared_layers == 0 and are unaffected. + """ + config = getattr(model, "config", model) + # Use _config_get so a dict config (some serialization / init paths) is read + # correctly; plain getattr would silently miss the field on a dict. + text_config = _config_get(config, "text_config", None) or config + model_type = _config_get(config, "model_type", "") or "" + text_model_type = _config_get(text_config, "model_type", "") or "" + if not ( + (isinstance(model_type, str) and model_type.startswith("gemma4")) + or (isinstance(text_model_type, str) and text_model_type.startswith("gemma4")) + ): + return False + return (_config_get(text_config, "num_kv_shared_layers", 0) or 0) > 0 + + def _resolve_moe_parameter_name(model, default_name: str, alternate_name: str) -> str: """ Resolve the actual parameter path for MoE expert weights. diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 08802f030e..e80e7c5fe4 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2156,6 +2156,14 @@ def unsloth_fast_generate(self, *args, **kwargs): # For newer HF kwargs["cache_implementation"] = "dynamic" + # Gemma 4 E-series (E2B / E4B) share KV across layers and their KV cache + # generation path is broken in transformers 5.5.0 + # (huggingface/transformers#45242): a cached forward diverges from the + # cache-free forward, so default greedy decode emits garbage. Default to the + # correct cache-free path unless the caller explicitly opts into a cache. + if "use_cache" not in kwargs and is_gemma4_shared_kv_model(self): + kwargs["use_cache"] = False + kwargs["cache_implementation"] = None # transformers 4.50 renamed num_logits_to_keep -> logits_to_keep; pop both, # re-emit under the spelling forward() accepts. _provided_num = kwargs.pop("num_logits_to_keep", None) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 3f8d9e6ce0..d802db6888 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -405,6 +405,15 @@ def unsloth_base_fast_generate(self, *args, **kwargs): if do_bfloat16_mixed_precision: cache_implementation = None + # Gemma 4 E-series (E2B / E4B) share KV across layers and their KV cache + # generation path is broken in transformers 5.5.0 + # (huggingface/transformers#45242): a cached forward diverges from the + # cache-free forward, so default greedy decode emits garbage. Default to the + # correct cache-free path unless the caller explicitly opts into a cache. + if "use_cache" not in kwargs and is_gemma4_shared_kv_model(self): + kwargs["use_cache"] = False + cache_implementation = None + if "generation_config" in kwargs: kwargs["generation_config"].cache_implementation = cache_implementation if cache_implementation is not None: