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>
This commit is contained in:
Daniel Han 2026-03-18 06:11:11 -07:00 committed by GitHub
commit 52f9c30513
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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")