Fix torch.compile with fullgraph=True when attention_mask input is used

[fxmarty]

As per title.

Fixes #29190

[fxmarty]

Let's consider using pytorch/pytorch#120400 if this is accepted and released.

[ArthurZucker]

I guess we don't have a choice?

[ArthurZucker]

Kind of related to #29210

[fxmarty]

@ArthurZucker this would conflict but is unrelated

[ArthurZucker]

🤢

[fxmarty]

agree

[fxmarty]

I guess we don't have a choice?

The other choice would be:

            is_tracing = (
                torch.jit.is_tracing()
                or isinstance(input_tensor, torch.fx.Proxy)
                or torch._dynamo.is_fullgraph_tracing())
            )

but torch._dynamo.is_fullgraph_tracing does not exist in PyTorch.

One other possibility is to always do causal_mask = causal_mask.mul(~torch.all(causal_mask == causal_mask.min(), dim=-1, keepdim=True)) no matter what.

One other possibility is to not use causal_mask = causal_mask.mul(~torch.all(causal_mask == causal_mask.min(), dim=-1, keepdim=True)) at all, but that means that we drop support for memory-efficient attention backend in Transformers cc @drisspg

One other possibility is to move the causal mask logic outside of the modeling code.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Alright, let's make sure CIs are green en bench are not slower!

[kwen2501]

Thanks for the fix!
Just wanted to share new API offerings from PyTorch (as of a couple days ago):
pytorch/pytorch#119602

Summary:

• A more general flag is named torch.compiler.is_compiling(). This flag indicates whether a graph is traced/compiled via torch.export() or torch.compile(). The flag works even in non-strict mode (i.e. even without TorchDynamo).
• A more specific flag is named torch.compiler.is_dynamo_compiling(), it's stricter, because it's only set to True when TorchDynamo is used, so, in non-strict mode it would be False.

Cc: @mreso @khabinov

[fxmarty]

@kwen2501 Thank you! torch.export is not always using dynamo? Reading https://pytorch.org/docs/stable/export.html I thought so!