From 925d4b067563657f5905e3c9256f02e63a185440 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 20 Jun 2026 07:43:52 +0000 Subject: [PATCH 1/4] Gemma 4 E-series: default generation to use_cache=False Gemma 4 E-series models (E2B / E4B) share KV across their last layers (num_kv_shared_layers > 0). Under transformers 5.5.0 their KV-cache generation path is broken (huggingface/transformers#45242): a cached forward diverges from the cache-free forward, so default greedy decode emits garbage even when the model has clearly learned the target (teacher-forced loss 0 on the cache-free path). Add is_gemma4_shared_kv_model and, in the text and vision fast-generate wrappers, default these models to use_cache=False (and skip the forced cache_implementation) unless the caller passed use_cache explicitly. This is gated to model_type starting with gemma4 and num_kv_shared_layers > 0, so Gemma 4 31B / 26B-A4B (num_kv_shared_layers == 0), Gemma 3, Gemma 3n and all other models are untouched. Users keep the cached path via an explicit generate(..., use_cache=True). --- unsloth/models/_utils.py | 25 +++++++++++++++++++++++++ unsloth/models/llama.py | 8 ++++++++ unsloth/models/vision.py | 9 +++++++++ 3 files changed, 42 insertions(+) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 742f5f57e0..b28f76b33e 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", ] @@ -3278,6 +3279,30 @@ def is_moe_model(model) -> bool: break return num_experts is not None and num_experts > 0 +pass + + +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) + text_config = getattr(config, "text_config", None) or config + model_type = getattr(config, "model_type", "") or "" + text_model_type = getattr(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 (getattr(text_config, "num_kv_shared_layers", 0) or 0) > 0 +pass def _resolve_moe_parameter_name(model, default_name: str, alternate_name: str) -> str: 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: From 952fa1916eaf96601c303fd734fa074f25658c0f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 20 Jun 2026 07:45:04 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/_utils.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b28f76b33e..0f03241187 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3279,7 +3279,8 @@ def is_moe_model(model) -> bool: break return num_experts is not None and num_experts > 0 -pass + + def is_gemma4_shared_kv_model(model) -> bool: @@ -3297,12 +3298,13 @@ def is_gemma4_shared_kv_model(model) -> bool: model_type = getattr(config, "model_type", "") or "" text_model_type = getattr(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")) + (isinstance(model_type, str) and model_type.startswith("gemma4")) + or (isinstance(text_model_type, str) and text_model_type.startswith("gemma4")) ): return False return (getattr(text_config, "num_kv_shared_layers", 0) or 0) > 0 -pass + + def _resolve_moe_parameter_name(model, default_name: str, alternate_name: str) -> str: From 2ef4039911e1c3175c58195db53381c2917a7cc4 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 21 Jun 2026 13:22:17 +0000 Subject: [PATCH 3/4] is_gemma4_shared_kv_model: read config via _config_get for dict configs getattr on a dict config silently misses model_type / text_config / num_kv_shared_layers (some serialization and init paths pass a dict). Use the existing _config_get helper so the Gemma 4 E-series detection is correct whether the config is an object or a dict. Behavior is unchanged for object configs. --- unsloth/models/_utils.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 0f03241187..5595766360 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3294,15 +3294,17 @@ def is_gemma4_shared_kv_model(model) -> bool: 26B-A4B have num_kv_shared_layers == 0 and are unaffected. """ config = getattr(model, "config", model) - text_config = getattr(config, "text_config", None) or config - model_type = getattr(config, "model_type", "") or "" - text_model_type = getattr(text_config, "model_type", "") or "" + # 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 (getattr(text_config, "num_kv_shared_layers", 0) or 0) > 0 + return (_config_get(text_config, "num_kv_shared_layers", 0) or 0) > 0 From b715eaad6de326ac15ca730477abfa2deececc51 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 21 Jun 2026 13:22:55 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/_utils.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 5595766360..8978e3be34 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3281,8 +3281,6 @@ 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. @@ -3307,8 +3305,6 @@ def is_gemma4_shared_kv_model(model) -> bool: 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.