ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Guard Gemma3N variants from flex attention defaults (#4116)

Browse source
This commit is contained in:
Daniel Han 2026-02-26 17:48:38 -08:00 • committed by GitHub
commit d9089de0f7
2 changed files with 12 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/models/_utils.py
View file

@ -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")

8
unsloth/models/vision.py
View file

@ -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":