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:
parent
51c08ded9b
commit
52f9c30513
1 changed files with 5 additions and 1 deletions
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue