From 52f9c30513d1a239e27987f0ebbeff218bcad34e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 18 Mar 2026 06:11:11 -0700 Subject: [PATCH] fix: exclude nemotron_h from flex_attention (#4424) * fix: exclude nemotron_h from flex_attention NemotronHForCausalLM does not support flex_attention and raises: NotImplementedError: NemotronHForCausalLM does not support an attention implementation through torch's flex_attention. Add nemotron_h to the exclusion list alongside gpt_oss and mllama so Unsloth falls back to the default attention implementation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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")