Diverging attention kernels due to `allow_is_bidirectional_skip` branching on torch.compile

[xmfan]

System Info

Hi, while we were updating the PyTorch transformers pin to v5.2.0, our regression tests caught a numerics issue between eager and compiled, the difference is very substantial (3.3 vs the typical e-4 accepted difference). Digging into it: pytorch/pytorch#175274 (comment), we found the cause to be in these lines (added in #41265):

transformers/src/transformers/masking_utils.py

Lines 490 to 491 in 147b7aa

	if allow_is_bidirectional_skip and _ignore_bidirectional_mask_sdpa(padding_mask, kv_length, local_size):
	return None

We set allow_is_bidirectional_skip=True in a few places:

transformers/src/transformers/masking_utils.py

Lines 996 to 997 in 147b7aa

	# Allow skipping the mask creation except we have additional masking operators (and/or masks)
	allow_is_bidirectional_skip = True

And in _ignore_bidirectional_mask_sdpa, we branch logic on whether we compile or not:

transformers/src/transformers/masking_utils.py

Lines 324 to 332 in 147b7aa

	if (
	not is_tracing(padding_mask)
	and (padding_mask is None or padding_mask.all())
	# in this case we need to add special patterns to the mask so cannot be skipped otherwise
	and (local_attention_size is None or kv_length < local_attention_size)
	):
	return True
	return False

This issue was found on BERT but it seems like it would affect other models too.

We've also verified that removing the branching fixes the numerical difference. I'm creating this issue to ask about the best way forward here. From the PR that added it, it looks like this was necessary specifically for executorch, but the algorithm difference is also affected all other APIs that fall under is_tracing . Can we restrict the check?

Who can help?

@vasqu @ArthurZucker @Cyrilvallez

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

I believe the description is enough, but I can provide a simpler repro on request

Expected behavior

transformers users probably shouldn't run into large numeric differences when compiling, at least not by default

[vasqu]

Just quickly commenting but essentially this caused by dropout rng being different for the different kernels (fa backend vs memory efficient). This could also happen in decoder-only models but dropout is just not really popular in modern models.

Wdyt about adding another condition on dropout? I think we have like 2-3 different names for dropout in the config so we could easily check for it (test_modeling_common should have a check for dropout somewhere surely). My issue is that we do want to avoid having to use the mask to opt into the fastest kernels where possible so see my tldr below.

Tl;dr: I can think of 2 options

• Make dropout not skip mask creations as another condition (forces slower training)
• The test is rewritten on torch side to create their masks themself (if you pass your own 4D mask we don't do anything)

[xmfan]

I see, the options make sense, we would have to do 1 if we wanted same numerics.

But back on restricting the check approach, if the issue was with export tracing, could we be using is_torchdynamo_exporting instead of is_tracing?

transformers/src/transformers/masking_utils.py

Lines 322 to 325 in 147b7aa

	# When using `torch.export` or `torch.onnx.dynamo_export`, we need to avoid to check the contents of the mask;
	# otherwise, we will encounter dynamic control flows
	if (
	not is_tracing(padding_mask)

[vasqu]

Tbh, personally I'd be for option 2 for the simple reason that users should run the fastest kernel whenever possible 😅 are there any references for the test? We know it works and we know the root cause is dropout behaving differently (different sets are disabled); not the mask being wrong.

Oh sorry overlooked the restriction: iirc, these are dyanmic control flows after the is_tracing check so (fullgraph) compile would complain, no? E.g. if mask.all(): ... else: ...

[gowthamr-tech]

I’m working on this issue.

[i3hz]

I tested it and padding_mask.all() fails under torch.compile(fullgraph=True) .
I'm happy to put up a PR for option 1 (dropout condition) which I guess would add an argument to _ignore_bidirectional_mask_sdpa for dropout and a conditional check with an early return. Should be a quick fix since the function is only called in one place .
What do you think @vasqu