OPT training with torch.compile decomposes SDPA into Triton kernels due to attention_mask

[msnliu]

System Info

Summary

When training OPT with torch.compile, an attention_mask=None input is converted into an all-ones mask for positional embeddings. This makes the model pass an explicit mask to SDPA instead of using is_causal=True, preventing the optimized FlashAttention/oneDNN backward kernel and causing a 20–30% performance regression.

Reproduction

Root Cause

In src/transformers/models/opt/modeling_opt.py:

if attention_mask is None:
    attention_mask = torch.ones(
        (inputs_embeds.shape[0], all_seq_len), dtype=torch.bool, device=inputs_embeds.device
    )

causal_mask = create_causal_mask(
    config=self.config,
    input_tensor=inputs_embeds,
    attention_mask=attention_mask,  # <- passing non-None torch.ones mask
)

[VimalN2005]

Hi @msnliu @zucchini-nlp @ArthurZucker,

I would love to work on this issue!

Looking at src/transformers/models/opt/modeling_opt.py (around line 347), the root cause is that attention_mask is overwritten in-place with torch.ones(...) for position embedding computation, which then leaks into create_causal_mask(...). This forces an explicit dense mask into SDPA instead of preserving is_causal=True, preventing the fused FlashAttention kernel under torch.compile.

I can fix this cleanly by scoping the fallback mask specifically to position ID computation (e.g. pos_attention_mask) while preserving attention_mask=None for create_causal_mask, along with a regression test ensuring causal_mask is None when unpadded.

Could you please assign this issue to me? I will open a clean PR with tests right away.

[msnliu]

Hi @msnliu @zucchini-nlp @ArthurZucker,

I would love to work on this issue!

Looking at src/transformers/models/opt/modeling_opt.py (around line 347), the root cause is that attention_mask is overwritten in-place with torch.ones(...) for position embedding computation, which then leaks into create_causal_mask(...). This forces an explicit dense mask into SDPA instead of preserving is_causal=True, preventing the fused FlashAttention kernel under torch.compile.

I can fix this cleanly by scoping the fallback mask specifically to position ID computation (e.g. pos_attention_mask) while preserving attention_mask=None for create_causal_mask, along with a regression test ensuring causal_mask is None when unpadded.

Could you please assign this issue to me? I will open a clean PR with tests right away.

Thank you for your effort. However, I’m afraid this is not the intended change.

The goal of this issue is me trying to communicate with the Transformers community and confirm which change is correct. To avoid any misunderstanding, please do not use AI to apply the expected change from my issue description.