[`0.0.18`, `memory_efficient_attention` with `attn_bias`]Getting NANs with arbitrary attn_bias mask with xformers==0.0.18

[toothacher17]

🐛 Bug

I am trying to use xformers to replace my native pytorch MHA implementation, sth like:

scale = 1 / query.shape[-1] ** 0.5
query = query * scale
attn = query @ key.transpose(-2, -1)
if attn_bias is not None:
    attn = attn + attn_bias
attn = attn.softmax(-1)
attn = F.dropout(attn, p)
return attn @ value

After switching to xformers, I am usiing xops.memory_efficient_attention(q, k ,v, attn_bias).

This works fine when I am using a lower triangular mask, either by passing in a LowerTriangularMask() or passing a torch.tensor with the same shape that built by my own.

However, when I am switching to use a arbitrary mask (supposing in pretraining stage, you opened reset_position_ids and reset_attention_mask flags, so you'll get a new start inside one sequence), I am getting NANs during evaluating (no grad forward) or training (with grad). Based on the log, the program is using the CUTLASS op.

Based on my observations, xformers saves 10-15% GPU memory and improves overall TFLOPs by 10-15%, so I really want to use it to replace with my native pytorch implementations. Could you help on this issue?

Environment

some depenedencies
trition==2.0
xformers==0.0.18
pytorch==2.0

[toothacher17]

So my question is: is arbitrary attention mask supported by xformers 0.0.18 yet?
I tried to follow several issues thread from the past and based on the doc, my understanding is that the answer to this question is YES. However, I could not make it work

[danthe3rd]

Hi @toothacher17
This should be supported by xFormers, but the behavior you report is definitively a bug. Unless a line in your mask is entirely masked-out, it shouldn't give you NaNs.
Do you have an independant minimum repro example so I can try it?

[danthe3rd]

I actually managed to repro it with this script:

import math
import torch
import xformers.ops.fmha as fmha

B, M, H, K = 1, 1024, 2, 64
dtype = torch.float16
device = "cuda"

q, k, v = [torch.randn([B, M, H, K], dtype=dtype, device=device) for _ in range(3)]
mask = torch.zeros([B, H, M, M], dtype=dtype, device=device)
mask[:, :, :256, :256] = -math.inf
out = fmha.memory_efficient_attention(q, k, v, attn_bias=mask)
print(out.sum())

It looks like it happens when the first 128 tokens of a sentence are entirely masked out.

Regarding your issue specifically, as you want to handle sequences of varying length, I recommend you use this mask, it will also save compute

[toothacher17]

hi, @danthe3rd

Thanks a lot for the reply.

Yes, I took another look at the BlockDiagonalCausalMask at flash_attention and Megatron-LM repo a few days ago. Basically it meets my need of pretraining with reset-position-id and reset-attention-mask, as resetting positions changes a lower-triangular matrix to a block diagonal matrix. To use this mask, like Megatron-LM, I'll need to change the shape, and merge all sentences in a batch into a single sentence, calculate the cumulative sequence length and pass them in to the ops.

[toothacher17]

Btw, I tested with cutlass based operators as it is the only class that supports customized attention bias. And the NANs is generated by the cutlass operator...

[danthe3rd]

To use this mask, like Megatron-LM, I'll need to change the shape, and merge all sentences in a batch into a single sentence, calculate the cumulative sequence length and pass them in to the ops.

It shouldn't be that involved, as normally everything operates at the token level, except for the attention. It might be a bit more involved if you are using some specific positional embedding tho.

Regarding the bug, I believe I understand where it comes from and should have a fix coming soon

[toothacher17]

Thanks for getting back to me. I am using sth like ROPE, but that's before it enters into attention, so as long as the cumulative sequence length for qkv is calculated correctly, the block diagonal causal matrix should be fine.

How soon do you think you can release the fix? If it is really coming soon, I'll wait for the fix before changing my code for the block diagonal causal matrix

[danthe3rd]

How soon do you think you can release the fix? If it is really coming soon, I'll wait for the fix before changing my code for the block diagonal causal matrix

Hopefully this week or next week for the xFormers development version (eg we might not necessarily push a new version tag yet)

[danthe3rd]

It should be fixed as of 540fcbf, and will be included in the next release (0.0.19). In the meantime, you can also use a development build >=0.0.19.dev516