Add XPU type for work-around -inf mask causing sdpa NaN issue in modeling files

[Liangliang-Ma]

Recently when we run transformers + qlora doing fine-tuning, we found NaN produced by torch.nn.functional.scaled_dot_product_attention.
Given that xpu has similar implement of fused sdpa, we would like to follow pytorch/transformers tmp solution to modify the mask here, which could solve the issue on XPU device too.

[Rocketknight1]

cc @muellerzr @SunMarc because I'm not sure who to ping for xpu! Feel free to ping someone else if needed

[SunMarc]

Could you check if you get the same issue as here with XPU ? Also, does this PR fixes your issue ? cc @ydshieh as you might know something

[Liangliang-Ma]

Could you check if you get the same issue as here with XPU ? Also, does this PR fixes your issue ? cc @ydshieh as you might know something

Yes we used XPU to do fine-tuning and got the same issue. With this PR the issue fixed.

[ydshieh]

I can only say that this seems reasonable to me but not more than that.

@Liangliang-Ma It would be better to provide a tiny code snippet to demonstrate the issue. Like providing a mask with a row with all places being masked, and pass it to F.scaled_dot_product_attention to show we do get NaN (on XPU).

[Liangliang-Ma]

I can only say that this seems reasonable to me but not more than that.

@Liangliang-Ma It would be better to provide a tiny code snippet to demonstrate the issue. Like providing a mask with a row with all places being masked, and pass it to F.scaled_dot_product_attention to show we do get NaN (on XPU).

import torch
import intel_extension_for_pytorch
from torch.nn import functional as F

torch.manual_seed(0)

a = 3
b = 4

q = torch.randn(size=(1, 1, a, b))
k = torch.randn(size=(1, 1, a, b))
v = torch.randn(size=(1, 1, a, b))

def check(q, k, v, device):

    q = q.to(device)
    k = k.to(device)
    v = v.to(device)

    neg_value = torch.finfo(q.dtype).min
    mask = [[neg_value, neg_value, neg_value], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]
    mask = torch.tensor([[mask]]).to(device)

    with torch.amp.autocast("xpu", dtype=torch.bfloat16):
        o = F.scaled_dot_product_attention(q, k, v, mask, 0.0, is_causal=False)
    print(o)

check(q, k, v, "cpu")
check(q, k, v, "xpu")

Thanks @ydshieh , I modified your test and get the NaN result like this:

tensor([[[[ 0.1210,  0.3627, -0.9969, -0.6149],
          [ 0.1295,  0.4572, -1.0491, -0.6166],
          [ 0.1095,  0.3819, -0.7369, -0.8267]]]])
tensor([[[[    nan,     nan,     nan,     nan],
          [ 0.1299,  0.4590, -1.0469, -0.6172],
          [ 0.1094,  0.3809, -0.7344, -0.8281]]]], device='xpu:0',
       dtype=torch.bfloat16)

I found that this issue caused from casting torch.finfo(torch.float).min to bfloat16, which result in a row of -inf.
This one can make sdpa kernel output NaN.

[Liangliang-Ma]

@Rocketknight1 Hi, may I know if the CI workflow failures are expected or not, for it seems to be generated modeling from original code different from what I modified. Thanks!

[SunMarc]

Hey @Liangliang-Ma, you need to modify the modular file for mistral and bamba as the modeling file is generated automatically from this file. This should fix the CI issue

[SunMarc]

Thanks for the snippet, LGTM with the fix I proposed for the CI

[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.

[Liangliang-Ma]

@SunMarc Thanks. The CI passed with the fix.

[SunMarc]

gentle ping @ArthurZucker as this concerns attention class

[Liangliang-Ma]

@SunMarc @ArthurZucker Soft reminder of this PR.

[Liangliang-Ma]

gentle ping @SunMarc @ArthurZucker again.

[SunMarc]

Merging this as this only concerns XPU workflow.

[ArthurZucker]

Sorry @Liangliang-Ma ! And thanks for the fix 🤗