huggingface/transformers · error · ValueError

`flex_attention` does not support `dropout`. Please use it w

Error message

`flex_attention` does not support `dropout`. Please use it with inference only (`model.eval()`) or turn off the attention dropout in the respective config.

What it means

The flex_attention integration does not implement dropout: torch's flex_attention has no dropout parameter (unlike SDPA), so any model config or call that requests attention dropout > 0 while using attn_implementation="flex_attention" is rejected immediately with ValueError. The message points to either inference-only usage (model.eval() typically sets dropout to 0 at call time) or disabling dropout in the config.

Source

Thrown at src/transformers/integrations/flex_attention.py:275

        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


def flex_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Union[torch.Tensor, "BlockMask"],
    scaling: float | None = None,
    softcap: float | None = None,
    s_aux: torch.Tensor | None = None,
    position_bias: torch.Tensor | None = None,
    **kwargs,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    if kwargs.get("dropout", 0.0) > 0:
        raise ValueError(
            "`flex_attention` does not support `dropout`. Please use it with inference"
            " only (`model.eval()`) or turn off the attention dropout in the respective config."
        )

    block_mask = None
    score_mask = None
    if isinstance(attention_mask, BlockMask):
        block_mask = attention_mask
    else:
        score_mask = attention_mask

    if score_mask is not None:
        score_mask = score_mask[:, :, :, : key.shape[-2]]

    def score_mod(score, batch_idx, head_idx, q_idx, kv_idx):
        if softcap is not None:
            score = softcap * torch.tanh(score / softcap)
        if score_mask is not None:

View on GitHub (pinned to a597f97485)

Solutions

  1. Set config.attention_dropout = 0.0 before/at model load when using flex_attention
  2. Only use flex_attention for inference (model.eval()), where dropout resolves to 0
  3. For training with dropout, use sdpa or eager attention instead

Example fix

# before
config = AutoConfig.from_pretrained(model_id)  # attention_dropout=0.1
model = AutoModelForCausalLM.from_pretrained(model_id, config=config, attn_implementation="flex_attention")
model.train(); model(**batch)  # ValueError

# after
config.attention_dropout = 0.0
model = AutoModelForCausalLM.from_pretrained(model_id, config=config, attn_implementation="flex_attention")
Defensive patterns

Strategy: validation

Validate before calling

config = AutoConfig.from_pretrained(model_id)
if config.attention_dropout:
    config.attention_dropout = 0.0  # flex_attention has no dropout
model = AutoModelForCausalLM.from_pretrained(model_id, config=config, attn_implementation="flex_attention")

Prevention

When it happens

Trigger: Instantiating a model with attn_implementation="flex_attention" whose config has attention_dropout > 0 (or passing dropout via the attention interface), then running a forward while the effective dropout value is > 0 — typically in train mode.

Common situations: Fine-tuning a model (train mode) whose config carries attention_dropout=0.1 (very common in Llama/Gemma-family configs) after switching the attention implementation to flex_attention for block-sparse masks; sharing a training config with an inference-only implementation.

Related errors


AI-assisted analysis of huggingface/transformers@a597f97485 (2026-08-14). Data as JSON: /api/errors/2622af7fe6b6078b. Report an issue: GitHub.