From d9089de0f7682408368abb7bfa4de3d9dd002189 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 26 Feb 2026 17:48:38 -0800 Subject: [PATCH] Guard Gemma3N variants from flex attention defaults (#4116) --- unsloth/models/_utils.py | 8 +++++--- unsloth/models/vision.py | 8 +++++++- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index dafdec4e1e..f883d466f0 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -235,14 +235,16 @@ def prefer_flex_attn_if_supported(model_class, config): model_class, "_supports_flex_attn", False ): return None - # GPT-OSS and Mllama use eager/sdpa attention during inference since - # flex attention returns incorrect results or errors out. + # GPT-OSS, Mllama and Gemma3N use eager/sdpa attention during + # inference since flex attention returns incorrect results or errors out. # GPT-OSS: left padding issues cause incorrect outputs. # Mllama: _update_causal_mask uses make_flex_block_causal_mask which # creates BlockMask with Q_LEN=KV_LEN=total_seq_len, but during # decode q_len=1, causing ValueError. Needs transformers update. + # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not + # support flex_attention. model_type = getattr(config, "model_type", "") if config else "" - if model_type in ("gpt_oss", "mllama"): + if model_type in ("gpt_oss", "mllama") or str(model_type).startswith("gemma3n"): return None if config is not None: setattr(config, "_attn_implementation", "flex_attention") diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index e6c859d44e..1f6b240a7d 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -630,7 +630,13 @@ class FastBaseModel: except KeyError: pass - default_attn_impl = "flex_attention" if flex_attn_impl else "sdpa" + model_type = str(getattr(auto_config, "model_type", "")).lower() + if model_type.startswith("gemma3n"): + # Gemma3N variants initialize timm-based vision towers which do + # not support flex_attention, so default to eager unless overridden. + default_attn_impl = "eager" + else: + default_attn_impl = "flex_attention" if flex_attn_impl else "sdpa" if not ("attn_implementation" in kwargs): kwargs["attn_implementation"] = default_attn_impl if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa":