Replace VMap that Fails Torch.Export

[YifanShenSZ]

The specific vmap used in causal mask construction would fail torch.export by

torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function <built-in function getitem>(*(FakeTensor(..., size=(1, 8), dtype=torch.int64), (BatchedTensor(lvl=1, bdim=0, value=
    FakeTensor(..., size=(1,), dtype=torch.int64)
), BatchedTensor(lvl=3, bdim=0, value=
    FakeTensor(..., size=(8,), dtype=torch.int64)
))), **{}): got RuntimeError("vmap: It looks like you're calling .item() on a Tensor. We don't support vmap over calling .item() on a Tensor, please try to rewrite what you're doing with other operations. If error is occurring somewhere inside PyTorch internals, please file a bug report.")

The solution is to use a simple causal mask construction

    q_indices = torch.arange(kv_length - q_length, kv_length)
    k_indices = torch.arange(kv_length)
    causal_mask = q_indices[:, None] >= k_indices[None, :]

cc @ArthurZucker @Cyrilvallez

[vasqu]

We have #41852 which should be available in the v5.rc0. I don't think we will backport for v4, as we keep this to patches only (I think).

For v4 you can check the integrations/executorch.py on that version. Back then we had a different implementation that avoided this by overwriting the attention/mask fns.

[YifanShenSZ]

Thx Anton! 🙏 So by v5 you mean at release 4.57.5 we will default to the non-vmap mask creation and torch.export can work out of the box? 😆

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: llama4

[vasqu]

Haha no, I really meant v5 as in major version :D see https://github.com/huggingface/transformers/releases/tag/v5.0.0rc0

[YifanShenSZ]

Close as we will get the fix in v5