diff --git a/tests/test_attention_implementation.py b/tests/test_attention_implementation.py new file mode 100644 index 0000000000..839d504b27 --- /dev/null +++ b/tests/test_attention_implementation.py @@ -0,0 +1,56 @@ +from types import SimpleNamespace + +import unsloth # noqa: F401 +from transformers.utils import import_utils + +from unsloth.models import _utils + + +class SupportsFlexAndSdpa: + _supports_flash_attn_2 = True + _supports_flex_attn = True + _supports_sdpa = True + + +def _config(model_type, **kwargs): + values = {"model_type": model_type, "attention_dropout": 0} + values.update(kwargs) + return SimpleNamespace(**values) + + +def _set_flex_available(monkeypatch, available): + monkeypatch.setenv("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") + monkeypatch.setattr( + import_utils, + "is_torch_flex_attn_available", + lambda: available, + raising = False, + ) + + +def test_gpt_oss_uses_eager_instead_of_flash_flex_or_sdpa(monkeypatch): + _set_flex_available(monkeypatch, True) + config = _config("gpt_oss") + + impl = _utils.resolve_attention_implementation( + SupportsFlexAndSdpa, + config, + supports_sdpa = True, + ) + + assert impl == "eager" + assert config._attn_implementation == "eager" + + +def test_gpt_oss_falls_back_to_eager_when_flex_unavailable(monkeypatch): + _set_flex_available(monkeypatch, False) + config = _config("gpt_oss") + + impl = _utils.resolve_attention_implementation( + SupportsFlexAndSdpa, + config, + supports_sdpa = True, + ) + + assert impl == "eager" + assert config._attn_implementation == "eager" diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 572ec8842a..5293943c77 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -226,8 +226,10 @@ def apply_unsloth_gradient_checkpointing( return use_gradient_checkpointing -# Models that don't work with flex_attention: -# GPT-OSS: left padding issues cause incorrect outputs. +# Models that don't work with flex_attention as the global Transformers +# attention implementation: +# GPT-OSS: training uses the custom flex sink patch, but inference intentionally +# falls back to eager because flex decoding gives incorrect outputs. # Mllama: BlockMask Q_LEN!=KV_LEN ValueError on decode. # NemotronH: hybrid Mamba-2 + Transformer, raises NotImplementedError. # Gemma3N: timm vision wrappers don't support flex_attention. @@ -235,6 +237,8 @@ def apply_unsloth_gradient_checkpointing( # access on some GPU architectures (B200). Falls back to eager safely. _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2") +_SDPA_EXCLUDED_MODELS = ("gpt_oss",) +_FLASH_EXCLUDED_MODELS = ("gpt_oss",) _EAGER_ONLY_PREFIXES = ("gemma3n",) _FLASH_ATTENTION_MAX_HEAD_DIM = 256 _FLASH_ATTENTION_DISABLED_WARNED = set() @@ -244,6 +248,14 @@ def _is_flex_excluded(model_type): return model_type in _FLEX_EXCLUDED_MODELS +def _is_sdpa_excluded(model_type): + return model_type in _SDPA_EXCLUDED_MODELS + + +def _is_flash_excluded(model_type): + return model_type in _FLASH_EXCLUDED_MODELS + + def _config_prefers_flex_attention(config): return any( _config_get(attention_config, "model_type", "").lower() @@ -360,6 +372,9 @@ def _get_max_attention_head_dim(config): def _get_flash_attention_disable_reason(config): + model_type = _config_get(config, "model_type", "").lower() + if _is_flash_excluded(model_type): + return f"{model_type} uses custom sink attention kernels" max_head_dim = _get_max_attention_head_dim(config) if max_head_dim is not None and max_head_dim > _FLASH_ATTENTION_MAX_HEAD_DIM: return ( @@ -489,9 +504,15 @@ def resolve_attention_implementation( supports_sdpa = model_class is not None and getattr( model_class, "_supports_sdpa", False ) - supports_flash_attention = model_class is not None and ( - getattr(model_class, "_supports_flash_attn_2", False) - or getattr(model_class, "_supports_flash_attn", False) + if _is_sdpa_excluded(model_type): + supports_sdpa = False + supports_flash_attention = ( + model_class is not None + and ( + getattr(model_class, "_supports_flash_attn_2", False) + or getattr(model_class, "_supports_flash_attn", False) + ) + and not _is_flash_excluded(model_type) ) supports_flex_attention = _supports_flex_attention(model_class, config, model_type) disable_reason = _get_flash_attention_disable_reason(config) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index a7dcaa4d88..0808138ac9 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -140,6 +140,7 @@ global DISABLE_SDPA_MODEL_NAMES DISABLE_SDPA_MODEL_NAMES = [ "gemma3,", # Add comma bc gemma3 will match gemma3n "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore + "gpt_oss", ]