From e5a2d39d57c7e7602da8595e69ee810e18bd45b1 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 11 Mar 2026 06:10:29 +0000 Subject: [PATCH 1/7] refactor attn impl replacement --- unsloth/models/_utils.py | 81 ++++++++++++++++++++++++++-------------- unsloth/models/llama.py | 4 +- unsloth/models/vision.py | 22 ++++------- 3 files changed, 61 insertions(+), 46 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 19b5fe0574..822b93a465 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -64,7 +64,7 @@ __all__ = [ "patch_compiled_autograd", "process_vision_info", "unsloth_compile_transformers", - "prefer_flex_attn_if_supported", + "determine_attention_implementation", "patch_fast_lora", "validate_loftq_config", "RaiseUninitialized", @@ -222,36 +222,61 @@ def apply_unsloth_gradient_checkpointing( return use_gradient_checkpointing -def prefer_flex_attn_if_supported(model_class, config): - if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0": - return None - try: - from transformers.utils.import_utils import is_torch_flex_attn_available +def determine_attention_implementation(model_class, config): + model_type = getattr(config, "model_type", "").lower() - if not is_torch_flex_attn_available(): - return None - if model_class is None or not getattr( - model_class, "_supports_flex_attn", False - ): - return None - # 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") or str(model_type).startswith("gemma3n"): - return None + # 1. Flash Attention 2 + if HAS_FLASH_ATTENTION and model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3"): + supports_fa2 = False + if model_class is not None: + supports_fa2 = getattr(model_class, "_supports_flash_attn_2", False) or getattr(model_class, "_supports_flash_attn", False) + + if supports_fa2: + if config is not None: + setattr(config, "_attn_implementation", "flash_attention_2") + if hasattr(config, "attn_implementation"): + setattr(config, "attn_implementation", "flash_attention_2") + return "flash_attention_2" + + # 2. Flex Attention + if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") != "0": + try: + from transformers.utils.import_utils import is_torch_flex_attn_available + + if is_torch_flex_attn_available() and (model_class is not None) and getattr( + model_class, "_supports_flex_attn", False + ): + # GPT-OSS, Mllama and Gemma3 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. + if model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3"): + if config is not None: + setattr(config, "_attn_implementation", "flex_attention") + if hasattr(config, "attn_implementation"): + setattr(config, "attn_implementation", "flex_attention") + return "flex_attention" + except Exception: + pass + + # 3. SDPA + if model_class is not None and getattr(model_class, "_supports_sdpa", False): if config is not None: - setattr(config, "_attn_implementation", "flex_attention") + setattr(config, "_attn_implementation", "sdpa") if hasattr(config, "attn_implementation"): - setattr(config, "attn_implementation", "flex_attention") - return "flex_attention" - except Exception: - return None + setattr(config, "attn_implementation", "sdpa") + return "sdpa" + + # 4. Eager + if config is not None: + setattr(config, "_attn_implementation", "eager") + if hasattr(config, "attn_implementation"): + setattr(config, "attn_implementation", "eager") + return "eager" def _run_temporary_patches(phase): diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 93d93e26d6..ee6fc7021a 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2341,9 +2341,7 @@ class FastLlamaModel: model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__] IS_FALCON_H1 = model_config.model_type.startswith("falcon_h1") - preferred_attn_impl = ( - prefer_flex_attn_if_supported(model_function, model_config) or "eager" - ) + preferred_attn_impl = determine_attention_implementation(model_function, model_config) has_rope_scaling = False try: diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index a8adba99e7..94186da2f0 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -597,8 +597,7 @@ class FastBaseModel: custom_datatype = None correct_dtype = None - # Stop SDPA for some archs like Pixtral / Mistral3 - flex_attn_impl = None + # Unified hierarchical attention fallback: Flash > Flex > SDPA > Eager if auto_config is None: auto_config = AutoConfig.from_pretrained( model_name, @@ -609,7 +608,8 @@ class FastBaseModel: model_class = auto_model._model_mapping[auto_config.__class__] except Exception: model_class = None - flex_attn_impl = prefer_flex_attn_if_supported(model_class, auto_config) + + attn_impl = determine_attention_implementation(model_class, auto_config) # Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with # FP8 weights. We just need to update it here for sanity. @@ -620,20 +620,12 @@ class FastBaseModel: except Exception: model_class = None - 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 + kwargs["attn_implementation"] = attn_impl if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa": - if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") == "0": - print( - f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." - ) + print( + f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." + ) del kwargs["attn_implementation"] bnb_config = None From b85e7d04888105b26c9f17e966366b3d2b86d3d4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 12 Mar 2026 10:01:35 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/_utils.py | 21 ++++++++++++++++----- unsloth/models/llama.py | 4 +++- unsloth/models/vision.py | 2 +- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 822b93a465..b7fa8d5dec 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -226,10 +226,16 @@ def determine_attention_implementation(model_class, config): model_type = getattr(config, "model_type", "").lower() # 1. Flash Attention 2 - if HAS_FLASH_ATTENTION and model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3"): + if ( + HAS_FLASH_ATTENTION + and model_type not in ("gpt_oss", "mllama") + and not model_type.startswith("gemma3") + ): supports_fa2 = False if model_class is not None: - supports_fa2 = getattr(model_class, "_supports_flash_attn_2", False) or getattr(model_class, "_supports_flash_attn", False) + supports_fa2 = getattr( + model_class, "_supports_flash_attn_2", False + ) or getattr(model_class, "_supports_flash_attn", False) if supports_fa2: if config is not None: @@ -243,8 +249,10 @@ def determine_attention_implementation(model_class, config): try: from transformers.utils.import_utils import is_torch_flex_attn_available - if is_torch_flex_attn_available() and (model_class is not None) and getattr( - model_class, "_supports_flex_attn", False + if ( + is_torch_flex_attn_available() + and (model_class is not None) + and getattr(model_class, "_supports_flex_attn", False) ): # GPT-OSS, Mllama and Gemma3 use eager/sdpa attention during # inference since flex attention returns incorrect results or errors out. @@ -254,7 +262,10 @@ def determine_attention_implementation(model_class, config): # decode q_len=1, causing ValueError. Needs transformers update. # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not # support flex_attention. - if model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3"): + if model_type not in ( + "gpt_oss", + "mllama", + ) and not model_type.startswith("gemma3"): if config is not None: setattr(config, "_attn_implementation", "flex_attention") if hasattr(config, "attn_implementation"): diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index ee6fc7021a..8e3892c712 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2341,7 +2341,9 @@ class FastLlamaModel: model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__] IS_FALCON_H1 = model_config.model_type.startswith("falcon_h1") - preferred_attn_impl = determine_attention_implementation(model_function, model_config) + preferred_attn_impl = determine_attention_implementation( + model_function, model_config + ) has_rope_scaling = False try: diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 94186da2f0..c44a77eff0 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -608,7 +608,7 @@ class FastBaseModel: model_class = auto_model._model_mapping[auto_config.__class__] except Exception: model_class = None - + attn_impl = determine_attention_implementation(model_class, auto_config) # Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with From c230eb525c6a59adb3408aaa42021bfdf2d461a1 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Thu, 12 Mar 2026 10:51:57 +0000 Subject: [PATCH 3/7] Gemma3 vs gemma3n issue resolve --- unsloth/models/_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 822b93a465..8ad68a3bbc 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -226,7 +226,7 @@ def determine_attention_implementation(model_class, config): model_type = getattr(config, "model_type", "").lower() # 1. Flash Attention 2 - if HAS_FLASH_ATTENTION and model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3"): + if HAS_FLASH_ATTENTION and model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3n"): supports_fa2 = False if model_class is not None: supports_fa2 = getattr(model_class, "_supports_flash_attn_2", False) or getattr(model_class, "_supports_flash_attn", False) @@ -254,7 +254,7 @@ def determine_attention_implementation(model_class, config): # decode q_len=1, causing ValueError. Needs transformers update. # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not # support flex_attention. - if model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3"): + if model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3n"): if config is not None: setattr(config, "_attn_implementation", "flex_attention") if hasattr(config, "attn_implementation"): From c4f01056aea3efb52556ccd14f0ba50853ef413a Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Fri, 13 Mar 2026 04:50:06 +0000 Subject: [PATCH 4/7] fix comment and match old behaviour --- unsloth/models/_utils.py | 2 +- unsloth/models/vision.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 7af1c336fc..b047781e11 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -250,7 +250,7 @@ def determine_attention_implementation(model_class, config): and (model_class is not None) and getattr(model_class, "_supports_flex_attn", False) ): - # GPT-OSS, Mllama and Gemma3 use eager/sdpa attention during + # 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 diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index c44a77eff0..210f298e65 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -623,9 +623,10 @@ class FastBaseModel: if not ("attn_implementation" in kwargs): kwargs["attn_implementation"] = attn_impl if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa": - print( - f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." - ) + if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0": + print( + f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." + ) del kwargs["attn_implementation"] bnb_config = None From f77a42e0ab7d83dd09757c4e66adda9fe577220f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 16:47:48 +0000 Subject: [PATCH 5/7] Fix attention selection regressions in PR #4250 - Extract _set_attn_impl helper to reduce config-stamping boilerplate - Add nemotron_h to FA2 and flex exclusion lists (from main) - Add attention_dropout > 0 check for flex attention (from main) - Restore gemma3n eager default in vision.py (timm vision towers) - Preserve sdpa fallback for unmapped/remote-code vision configs - Restore original SDPA warning default (print when env var unset) - Make config=None safe in determine_attention_implementation --- unsloth/models/_utils.py | 55 ++++++++++++++++++++++------------------ unsloth/models/vision.py | 16 ++++++++++-- 2 files changed, 44 insertions(+), 27 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b047781e11..bd442f78ed 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -222,11 +222,24 @@ def apply_unsloth_gradient_checkpointing( return use_gradient_checkpointing +def _set_attn_impl(config, impl): + """Stamp the chosen attention implementation onto the config object.""" + if config is not None: + setattr(config, "_attn_implementation", impl) + if hasattr(config, "attn_implementation"): + setattr(config, "attn_implementation", impl) + return impl + + def determine_attention_implementation(model_class, config): - model_type = getattr(config, "model_type", "").lower() + model_type = getattr(config, "model_type", "").lower() if config else "" # 1. Flash Attention 2 - if HAS_FLASH_ATTENTION and model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3n"): + if ( + HAS_FLASH_ATTENTION + and model_type not in ("gpt_oss", "mllama", "nemotron_h") + and not model_type.startswith("gemma3n") + ): supports_fa2 = False if model_class is not None: supports_fa2 = getattr( @@ -234,11 +247,7 @@ def determine_attention_implementation(model_class, config): ) or getattr(model_class, "_supports_flash_attn", False) if supports_fa2: - if config is not None: - setattr(config, "_attn_implementation", "flash_attention_2") - if hasattr(config, "attn_implementation"): - setattr(config, "attn_implementation", "flash_attention_2") - return "flash_attention_2" + return _set_attn_impl(config, "flash_attention_2") # 2. Flex Attention if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") != "0": @@ -250,37 +259,33 @@ def determine_attention_implementation(model_class, config): and (model_class is not None) and getattr(model_class, "_supports_flex_attn", False) ): - # GPT-OSS, Mllama and Gemma3N use eager/sdpa attention during - # inference since flex attention returns incorrect results or errors out. + attention_dropout = getattr(config, "attention_dropout", 0) or 0 + # GPT-OSS, Mllama, Gemma3N and NemotronH 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. - if model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3n"): - if config is not None: - setattr(config, "_attn_implementation", "flex_attention") - if hasattr(config, "attn_implementation"): - setattr(config, "attn_implementation", "flex_attention") - return "flex_attention" + # NemotronH: hybrid Mamba-2 + Transformer model that does not + # support flex_attention (raises NotImplementedError from transformers). + is_excluded = ( + model_type in ("gpt_oss", "mllama", "nemotron_h") + or model_type.startswith("gemma3n") + ) + if attention_dropout == 0 and not is_excluded: + return _set_attn_impl(config, "flex_attention") except Exception: pass # 3. SDPA if model_class is not None and getattr(model_class, "_supports_sdpa", False): - if config is not None: - setattr(config, "_attn_implementation", "sdpa") - if hasattr(config, "attn_implementation"): - setattr(config, "attn_implementation", "sdpa") - return "sdpa" + return _set_attn_impl(config, "sdpa") # 4. Eager - if config is not None: - setattr(config, "_attn_implementation", "eager") - if hasattr(config, "attn_implementation"): - setattr(config, "attn_implementation", "eager") - return "eager" + return _set_attn_impl(config, "eager") def _run_temporary_patches(phase): diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 210f298e65..72d99c2d41 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -609,7 +609,19 @@ class FastBaseModel: except Exception: model_class = None - attn_impl = determine_attention_implementation(model_class, auto_config) + model_type = str(getattr(auto_config, "model_type", "")).lower() + if model_type.startswith("gemma3n"): + # Gemma3N variants use timm-based vision towers which do not support + # flex_attention. The old code defaulted gemma3n to eager; preserve + # that behavior rather than letting the hierarchy pick sdpa. + attn_impl = "eager" + elif model_class is None and supports_sdpa: + # When model_class cannot be resolved (remote-code or unmapped + # configs), the old code defaulted to sdpa. Preserve that fallback + # instead of falling through to eager. + attn_impl = "sdpa" + else: + attn_impl = determine_attention_implementation(model_class, auto_config) # Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with # FP8 weights. We just need to update it here for sanity. @@ -623,7 +635,7 @@ class FastBaseModel: if not ("attn_implementation" in kwargs): kwargs["attn_implementation"] = attn_impl if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa": - if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0": + if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") == "0": print( f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." ) From 3b8860c0393f43a05e355024e4e8145db8a32fab Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 17:17:37 +0000 Subject: [PATCH 6/7] Consistent config stamping in vision.py attention fallbacks Use _set_attn_impl for gemma3n eager and unmapped sdpa fallbacks in vision.py so config._attn_implementation is consistent with what gets passed to from_pretrained via kwargs. --- unsloth/models/vision.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 72d99c2d41..c1a605cfc8 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -29,7 +29,7 @@ except: from ..kernels import ( post_patch_loss_function, ) -from ._utils import __version__, importlib_version, _prepare_model_for_qat +from ._utils import __version__, importlib_version, _prepare_model_for_qat, _set_attn_impl from ._utils import * from .loader_utils import _get_fp8_mode_and_check_settings from ..save import patch_saving_functions @@ -614,12 +614,12 @@ class FastBaseModel: # Gemma3N variants use timm-based vision towers which do not support # flex_attention. The old code defaulted gemma3n to eager; preserve # that behavior rather than letting the hierarchy pick sdpa. - attn_impl = "eager" + attn_impl = _set_attn_impl(auto_config, "eager") elif model_class is None and supports_sdpa: # When model_class cannot be resolved (remote-code or unmapped # configs), the old code defaulted to sdpa. Preserve that fallback # instead of falling through to eager. - attn_impl = "sdpa" + attn_impl = _set_attn_impl(auto_config, "sdpa") else: attn_impl = determine_attention_implementation(model_class, auto_config) From 405bb58284d1afe262183d23e8ab2c5baa69dd68 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 17:29:07 +0000 Subject: [PATCH 7/7] Fix config stamping when SDPA is not supported When supports_sdpa=False and the attention hierarchy selected sdpa, the kwarg was deleted but config._attn_implementation remained as "sdpa". Stamp it to "eager" for consistency before deleting the kwarg. --- unsloth/models/vision.py | 1 + 1 file changed, 1 insertion(+) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index c1a605cfc8..0dae787e22 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -639,6 +639,7 @@ class FastBaseModel: print( f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." ) + _set_attn_impl(auto_config, "eager") del kwargs["attn_implementation"] bnb_config = None