`masking_utils.create_causal_mask` fails in `@torch.compile(fullgraph=True)`
Author: vadimkantorov
Pull request: no
Labels: name: bug

### System Info

4.57.3 2.9.1+cu130 13.0 91300

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
import torch
import transformers

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.llm = transformers.AutoModelForCausalLM.from_pretrained(
            'PleIAs/Monad',
            attn_implementation = 'flash_attention_2',
            dtype = torch.bfloat16,
        )

    def forward(self, x):
        outputs = self.llm.model(
            inputs_embeds=x,
            attention_mask=torch.ones(x.shape[0], x.shape[1], dtype=torch.long, device=x.device),
            return_dict=True,
        )
        return outputs.last_hidden_state.sum()

@torch.compile(fullgraph=True)
def compute_output(f, x):
    return f(x)

if __name__ == '__main__':
    print(transformers.__version__, torch.__version__, torch.version.cuda, torch.backends.cudnn.version())

    model = Model().to('cuda')
    x = torch.zeros((64, 145, 256), device = 'cuda', dtype = torch.bfloat16)
    with torch.no_grad():
        a = compute_output(model, x)
```

```
...
  File "/.../.venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 841, in compile_wrapper
    raise e.with_traceback(None) from e.__cause__  # User compiler error
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch._dynamo.exc.Unsupported: Data-dependent branching
  Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.
  Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.
  Hint: Use `torch.cond` to express dynamic control flow.

  Developer debug context: attempted to jump with TensorVariable()

 For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0170.html

from user code:
...
  File "/.../.venv/lib/python3.12/site-packages/torch/_functorch/apis.py", line 406, in wrapper
    return eager_transforms.grad_impl(func, argnums, has_aux, args, kwargs)
  File "/.../.venv/lib/python3.12/site-packages/torch/_functorch/eager_transforms.py", line 1406, in grad_impl
    results = grad_and_value_impl(func, argnums, has_aux, args, kwargs)
  File "/.../.venv/lib/python3.12/site-packages/torch/_functorch/vmap.py", line 47, in fn
    return f(*args, **kwargs)
  File "/.../.venv/lib/python3.12/site-packages/torch/_functorch/eager_transforms.py", line 1364, in grad_and_value_impl
    output = func(*args, **kwargs)
  File "/.../bug.py", line 14, in forward
    outputs = self.llm.model(
  File "/.../.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 1072, in wrapper
    outputs = func(self, *args, **kwargs)
  File "/.../.venv/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 382, in forward
    causal_mask = create_causal_mask(
  File "/.../.venv/lib/python3.12/site-packages/transformers/masking_utils.py", line 825, in create_causal_mask
    causal_mask = mask_interface(
  File "/.../.venv/lib/python3.12/site-packages/transformers/masking_utils.py", line 558, in flash_attention_mask
    if attention_mask.all():

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"
```

### Expected behavior

`create_causal_mask` should be fullgraph'able. If I use `attention_impl='sdpa'`, it all works okay

Problem seems to be coming from attempting an optimization:
```python
        # We only return an actual mask if there is at least 1 padding token, otherwise we return `None` and use `is_causal` in FA2
        # (note that the attention_mask is a boolean dtype here)
        if attention_mask.all():
            attention_mask = None
```

But I think there needs to be a way to disable this optimization and gain fullgraph-ability

Commenting out these two lines leads to another error:
```

    raise e.with_traceback(None) from e.__cause__  # User compiler error
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch._dynamo.exc.Unsupported: Dynamic shape operator
  Explanation: Operator `aten.nonzero.default`'s output shape depends on input Tensor data.
  Hint: Enable tracing of dynamic shape operators with `torch._dynamo.config.capture_dynamic_output_shape_ops = True`

  Developer debug context: aten.nonzero.default

 For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0036.html

from user code:
   File "/.../bug2.py", line 23, in compute_output
    return f(x)
  File "/...bug2.py", line 14, in forward
    outputs = self.llm.model(
  File "/.../.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 1072, in wrapper
    outputs = func(self, *args, **kwargs)
  File "/.../.venv/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 395, in forward
    hidden_states = decoder_layer(
  File "/.../.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 94, in __call__
    return super().__call__(*args, **kwargs)
  File "/.../.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/.../.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/.../.venv/lib/python3.12/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
  File "/.../.venv/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 294, in forward
    hidden_states, _ = self.self_attn(
  File "/h.../.venv/lib/python3.12/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
  File "/.../.venv/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 252, in forward
    attn_output, attn_weights = attention_interface(
  File "/.../.venv/lib/python3.12/site-packages/transformers/integrations/flash_attention.py", line 66, in flash_attention_forward
    attn_output = _flash_attention_forward(
  File "/.../.venv/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py", line 607, in _flash_attention_forward
    q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
  File "/.../.venv/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py", line 278, in _upad_input
    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
  File "/.../.venv/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py", line 225, in _get_unpad_data
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"
```

It is mentioned in-there:
```
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
    # NOTE: Similar to the `.item()` in prepare_fa2_from_position_ids, with torch compile,
    # this might cause a graph break
    max_seqlen_in_batch = seqlens_in_batch.max().item()
```

Could all this be done without a graph-break? E.g. allow the user to pass-in seqlens and this padded info

Comments:
Probably after fixing this, `attention_impl = 'flash_attention_2'` can also be fullgraph-able, just like `attention_impl = 'sdpa'`

A variant without graph breaks would be very useful for using `torch.func.grad`:
- https://github.com/pytorch/pytorch/issues/167729#issuecomment-3677205916

cc @cyrilvallez since he wrote that masking block, but replies may be slow until after Christmas!

It's masking module and also padding/unpadding stuff for FA. I propose to have non-graph-breaking alternatives for FA (by default, or at least as an option: so maybe the user would have to pass in some extra meta-info like max-seq-len instead of having to graph-break for in-model computation of these), just like it is for SDPA. This will enable more simple/faithful comparisons between FA and SDPA.

IMO it would also be very useful to have a test harness and a table at HF website: listing torch.compile-ability with different torch.compile's options for some popular models, and some benchmark results.

This is already fixed on main/v5 rc0/1!
Let me know if you think something still needs to be addressed!

Including both masking and pad/unpad?

Is there anywhere a per-nightly or per-release table of smoke tests of popular models under different torch.compile options? To check for errors / graph breaks / cudagraph-ability etc 

Thanks!

@Cyrilvallez In main, unpad still uses non-static nonzero and item() - these should still produce graph breaks :(

https://github.com/huggingface/transformers/blob/9971e410516146a78972cb9034eec1a408a29ede/src/transformers/modeling_flash_attention_utils.py#L195

I think, it is needed a way to do pad/unpad without graph breaks - in the worst case, to allow the user to pass-through this seqlen meta-information as model input

@Cyrilvallez I've upgraded to 5.0.0rc1, and still got the same error in `create_causal_mask` (and I assume the same errors in un/pad downstream):
```
5.0.0rc1 2.9.1+cu130 13.0 91300

  File "/.../.venv/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 411, in forward
    causal_mask = create_causal_mask(
  File "/.../.venv/lib/python3.12/site-packages/transformers/masking_utils.py", line 858, in create_causal_mask
    causal_mask = mask_interface(
  File "/.../.venv/lib/python3.12/site-packages/transformers/masking_utils.py", line 568, in flash_attention_mask
    if attention_mask.all():

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"
```

Could you please reopen the issue?

Here is the repro code:
```python
import torch
import transformers

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.llm = transformers.AutoModelForCausalLM.from_pretrained(
            'PleIAs/Monad',
            attn_implementation = 'flash_attention_2',
            dtype = torch.bfloat16,
        )

    def forward(self, x):
        outputs = self.llm.model(
            inputs_embeds=x,
            attention_mask=torch.ones(x.shape[0], x.shape[1], dtype=torch.long, device=x.device),
            return_dict=True,
        )
        return outputs.last_hidden_state.sum()

@torch.compile(fullgraph=True)
def compute_output(f, x):
    return f(x)

if __name__ == '__main__':
    print(transformers.__version__, torch.__version__, torch.version.cuda, torch.backends.cudnn.version())

    model = Model().to('cuda')
    x = torch.zeros((64, 145, 256), device = 'cuda', dtype = torch.bfloat16)
    with torch.no_grad():
        a = compute_output(model, x)
```

> [@Cyrilvallez](https://github.com/Cyrilvallez) In main, unpad still uses non-static nonzero and item() - these should still produce graph breaks :(

@Cyrilvallez Could you please reopen this issue to track all this un/pad stuff in HF (besides purely FA's problems with `torch.compile(fullgraph=True)`)?

I suggest to allow the user to pass in un/pad information directly as model input (maybe as underscored arguments if wishing to do it in experimental/temporary way), at least this should eliminate graph breaks on the HF side

Could we use sequence packing arguments like `position_ids` to bypass the HF's un/pad computations (incompatible with fullgraph)? Could this cu_seqlen code path actually be exercised even when padding is present? One can imagine that it seqs inds/lens could be not contig - and this still can bypass issue that padding info must be computed within the model
- https://github.com/huggingface/transformers/pull/31629

In your specific case, I don't see why we even create an attention mask, e.g. `attention_mask=torch.ones(x.shape[0], x.shape[1], dtype=torch.long, device=x.device),` should just be `attention_mask=None`. Note that your input needs to have bsz>1 for this to work (see #41827). Yes this is unideal :( and is due to FA packing having to check for proper packed sequences. I've added more details in the other issue https://github.com/huggingface/transformers/issues/43075#issuecomment-3711704293

This goes back to the issue with position ids and control flows cc @Cyrilvallez, another one for the environment var solution 

@vasqu Isn't there also an issue with un/pad? (in addition to create_causal_mask)

See the second point of my comment in the other issue https://github.com/huggingface/transformers/issues/43075#issuecomment-3711704293

> - Second, If you want fullgraph compatibility you will need to prepare the FA kwargs outside of the forward, e.g. via a data collator. You can take look at my [gist](https://gist.github.com/vasqu/c264d14268fbc7f1eb725e09dfb87331) - note the return_flash_attn_kwargs=True for the collator, without you fall into the pos ids case only again. I also have this PR https://github.com/huggingface/transformers/pull/41827 which checked compile compatibility.