diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 116352c504..d6255d8ed9 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -242,8 +242,12 @@ def prefer_flex_attn_if_supported(model_class, config): # decode q_len=1, causing ValueError. Needs transformers update. # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not # support flex_attention. + # NemotronH: hybrid Mamba-2 + Transformer model that does not + # support flex_attention (raises NotImplementedError from transformers). model_type = getattr(config, "model_type", "") if config else "" - if model_type in ("gpt_oss", "mllama") or str(model_type).startswith("gemma3n"): + if model_type in ("gpt_oss", "mllama", "nemotron_h") or str( + model_type + ).startswith("gemma3n"): return None if config is not None: setattr(config, "_attn_implementation", "flex_attention")