diff --git a/tests/test_attn_impl_honor_explicit.py b/tests/test_attn_impl_honor_explicit.py new file mode 100644 index 0000000000..3fb7a2208f --- /dev/null +++ b/tests/test_attn_impl_honor_explicit.py @@ -0,0 +1,190 @@ +"""An explicit non-flash attention request must survive the flash disable path. + +When flash attention is disabled for a model, a caller who explicitly asked for +"sdpa" or "flex_attention" should keep that choice instead of being downgraded +to whatever the conservative supports_* fallback would pick. +""" + +import pytest + +from unsloth.models._utils import ( + _disable_flash_attention_if_needed, + resolve_attention_implementation, +) + + +def test_explicit_sdpa_is_honored_even_when_not_marked_supported(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "sdpa", + supports_sdpa = False, # conservative flag would have skipped sdpa + supports_flex_attention = False, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_explicit_flex_is_honored_when_supported(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flex_attention", + supports_sdpa = True, + supports_flex_attention = True, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "flex_attention" + assert config.get("_attn_implementation") == "flex_attention" + + +def test_explicit_flex_falls_back_when_not_supported(): + # flex_attention is False for known-broken/excluded configs (e.g. gpt_oss), + # so an explicit flex request must not select that backend - it falls back. + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flex_attention", + supports_sdpa = True, + supports_flex_attention = False, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + + +def test_synthesized_config_sdpa_is_not_treated_as_explicit(): + # The language loader seeds the config with attn_implementation="sdpa"; when the + # caller passes nothing, that synthesized value must not override the flex fallback + # for a model that supports flex but not sdpa. + config = {"attn_implementation": "sdpa"} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = None, + supports_sdpa = False, + supports_flex_attention = True, + would_use_flash_attention = False, + disable_reason = "unit test forces flash disabled", + ) + assert result == "flex_attention" + + +def test_no_disable_reason_returns_request_untouched(): + result = _disable_flash_attention_if_needed( + {}, + attn_implementation = "flash_attention_2", + disable_reason = None, + ) + assert result == "flash_attention_2" + + +def test_flash_request_still_falls_back_when_disabled(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flash_attention_2", + supports_sdpa = True, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + + +def test_resolver_honors_explicit_sdpa_when_not_supported_and_flash_disabled(): + # End-to-end through the public resolver: an explicit sdpa request with a + # flash-disabled config (oversized head dim) and supports_sdpa=False must not be + # rewritten to eager by the resolver's own not-supports_sdpa guard. + config = {"model_type": "test", "head_dim": 512} # head_dim > 256 disables flash + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_resolver_downgrades_non_explicit_sdpa_when_not_supported(): + # No explicit request: the model resolution seeds sdpa/eager and the guard must + # still downgrade a synthesized sdpa to eager for a model that cannot run it. + config = {"model_type": "test", "attn_implementation": "sdpa"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = None, + supports_sdpa = False, + ) + assert result == "eager" + + +def test_resolver_downgrades_explicit_sdpa_for_sdpa_excluded_model(): + # gpt_oss is in _SDPA_EXCLUDED_MODELS (sdpa is known-broken) and _FLASH_EXCLUDED_MODELS + # (flash disabled). Honoring an explicit sdpa request must not re-enable that broken + # backend: it downgrades to eager, mirroring how an explicit flex request falls back + # for _FLEX_EXCLUDED_MODELS. supports_sdpa=True proves the exclusion overrides even a + # model that otherwise advertises SDPA support. + config = {"model_type": "gpt_oss"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = True, + ) + assert result == "eager" + assert config.get("_attn_implementation") == "eager" + + +@pytest.mark.parametrize("model_type", ["gemma3", "gemma3_text"]) +def test_resolver_downgrades_explicit_sdpa_for_disable_sdpa_model(model_type): + # gemma3 / gemma3_text are in DISABLE_SDPA_MODEL_NAMES: the loader forces + # supports_sdpa=False because their bundled SDPA modules are wrong. An explicit + # sdpa request with flash disabled must NOT re-enable that known-wrong path - it + # downgrades to eager, exactly like _SDPA_EXCLUDED_MODELS (gpt_oss). head_dim>256 + # disables flash to mirror the real flash-disabled scenario. + config = {"model_type": model_type, "head_dim": 512} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "eager" + assert config.get("_attn_implementation") == "eager" + + +def test_resolver_does_not_overmatch_gemma3n_for_explicit_sdpa(): + # The "gemma3," trailing-comma guard must not match gemma3n: gemma3n is not in + # DISABLE_SDPA_MODEL_NAMES, so it stays a conservative (not known-wrong) model and an + # explicit sdpa request is still honored. Proves the substring match neither over- nor + # under-matches. + config = {"model_type": "gemma3n", "head_dim": 512} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_resolver_downgrades_synthesized_sdpa_for_disable_sdpa_model(): + # A synthesized/default sdpa (requested is None; the value came from config) on a + # DISABLE_SDPA_MODEL_NAMES model must still downgrade to eager. + config = {"model_type": "gemma3", "attn_implementation": "sdpa"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = None, + supports_sdpa = False, + ) + assert result == "eager" + + +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, "-q"])) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b504a19f74..169b610988 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -423,6 +423,18 @@ def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_len _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2") _SDPA_EXCLUDED_MODELS = ("gpt_oss",) +# The loader (loader.py) forces supports_sdpa=False for these because their bundled +# SDPA modules are wrong. Kept here, not in loader.py, so _is_sdpa_excluded can honor +# them without a loader -> _utils import cycle (loader.py already imports from _utils +# and re-exports this name for callers like sentence_transformer.py). Entries are matched +# as substrings against a comma-joined model_types string ending in a comma, so "gemma3," +# matches a distinct "gemma3" entry but not "gemma3n", and "gemma3_text" matches the +# EmbeddingGemma text model. +DISABLE_SDPA_MODEL_NAMES = [ + "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore + "gpt_oss", +] _FLASH_EXCLUDED_MODELS = ("gpt_oss",) _EAGER_ONLY_PREFIXES = ("gemma3n",) _FLASH_ATTENTION_MAX_HEAD_DIM = 256 @@ -433,8 +445,23 @@ def _is_flex_excluded(model_type): return model_type in _FLEX_EXCLUDED_MODELS +def _is_sdpa_disabled_by_name(model_type): + # Mirror the loader's DISABLE_SDPA_MODEL_NAMES check: loader.py builds + # model_types_all = ",".join(model_types) + "," and tests `name in model_types_all`. + # Rebuild the same trailing-comma form for a single model_type so the match is + # identical (e.g. "gemma3," matches "gemma3" but not "gemma3n", and "gemma3_text" + # still matches "gemma3_text"). + model_types_all = model_type.lower() + "," + return any(name.lower() in model_types_all for name in DISABLE_SDPA_MODEL_NAMES) + + def _is_sdpa_excluded(model_type): - return model_type in _SDPA_EXCLUDED_MODELS + # SDPA is known-broken for these models, so an explicit sdpa request must not + # re-enable it. Two sources: _SDPA_EXCLUDED_MODELS (resolver-level, e.g. gpt_oss) + # and DISABLE_SDPA_MODEL_NAMES (loader-level, e.g. gemma3 / gemma3_text, which the + # loader also forces to supports_sdpa=False). + lowered = model_type.lower() + return lowered in _SDPA_EXCLUDED_MODELS or _is_sdpa_disabled_by_name(lowered) def _is_flash_excluded(model_type): @@ -610,6 +637,12 @@ def _disable_flash_attention_if_needed( if disable_reason is None: return attn_implementation + # Only an implementation passed by the caller counts as an explicit request. + # Values read from the config are synthesized by the loaders (the language path + # seeds the config with attn_implementation="sdpa") or come from Transformers + # defaults, so they must not be treated as a deliberate user choice. + explicit_request = attn_implementation + requested_attn_implementation = attn_implementation if requested_attn_implementation is None: requested_attn_implementation = _config_get(config, "_attn_implementation", None) @@ -619,6 +652,20 @@ def _disable_flash_attention_if_needed( if requested_attn_implementation == "eager": return _set_attn_impl(config, "eager") + model_type = _config_get(config, "model_type", "") + + # The disable reason is flash-specific: honor an explicit non-flash request from + # the caller instead of downgrading it. SDPA is honored unless the model's SDPA is + # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES + # (e.g. gemma3 / gemma3_text); flex_attention + # is honored only when it is actually usable, since supports_flex_attention already + # rejects the excluded/broken/unavailable configs. This keeps an explicit request + # from selecting a backend the repo marks as wrong. + if explicit_request == "sdpa" and not _is_sdpa_excluded(model_type.lower()): + return _set_attn_impl(config, "sdpa") + if explicit_request == "flex_attention" and supports_flex_attention: + return _set_attn_impl(config, "flex_attention") + if supports_sdpa: fallback_attn_implementation = "sdpa" elif supports_flex_attention: @@ -631,7 +678,6 @@ def _disable_flash_attention_if_needed( if _is_flash_attention_requested(requested_attn_implementation) else "flash_attention_2" ) - model_type = _config_get(config, "model_type", "") warning_key = ( model_type, logged_attn_implementation, @@ -845,7 +891,19 @@ def resolve_attention_implementation( final_attn_impl = requested_attn_implementation _set_attn_impl(config, final_attn_impl) - if not supports_sdpa and final_attn_impl == "sdpa": + # A caller who explicitly passes requested_attn_implementation="sdpa" keeps it even + # on a conservatively unsupported model, mirroring _disable_flash_attention_if_needed + # which honors an explicit sdpa request. The exception is a model whose SDPA is + # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES + # (e.g. gemma3 / gemma3_text, which the loader also forces to supports_sdpa=False): + # an explicit request must not re-enable it, so it still downgrades to eager, just + # like flex falls back for _FLEX_EXCLUDED_MODELS. A synthesized/default sdpa + # (requested is None, so the value came from the model resolution above or the + # config) also downgrades. + honor_explicit_sdpa = requested_attn_implementation == "sdpa" and not _is_sdpa_excluded( + model_type + ) + if not supports_sdpa and final_attn_impl == "sdpa" and not honor_explicit_sdpa: print( f"Unsloth: {(model_type_name or 'model').title()} does not support SDPA - switching to fast eager." ) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 84f808d2b5..22cb65dc4a 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -21,6 +21,10 @@ from ._utils import ( USE_MODELSCOPE, get_transformers_model_type, hf_login, + # Single source of truth is _utils.py; re-exported here so callers doing + # `from unsloth.models.loader import DISABLE_SDPA_MODEL_NAMES` keep working and so + # _is_sdpa_excluded (in _utils) can honor it without a loader -> _utils cycle. + DISABLE_SDPA_MODEL_NAMES, ) from .granite import FastGraniteModel from .llama import FastLlamaModel, logger @@ -196,14 +200,6 @@ DISABLE_COMPILE_MODEL_NAMES = [ "granite,llava_next", # Granite-vision 3 ] -global DISABLE_SDPA_MODEL_NAMES -# Disables some SDPA modules since it's wrong -DISABLE_SDPA_MODEL_NAMES = [ - "gemma3,", # Add comma bc gemma3 will match gemma3n - "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore - "gpt_oss", -] - def _fix_rope_inv_freq(model): """Fix inv_freq corruption caused by transformers v5 meta-device loading.