NaNs when training with attn_bias (f32)

[zen-d]

❓ Questions and Help

Hi, I pass in the attn_bias to xformers.ops.memory_efficient_attention, but meet the following error

NotImplementedError: No operator found for `memory_efficient_attention_forward` with inputs:
     query       : shape=(831, 43, 32, 8) (torch.float32)             
     key         : shape=(831, 43, 32, 8) (torch.float32)
     value       : shape=(831, 43, 32, 8) (torch.float32)  
     attn_bias   : <class 'torch.Tensor'>
     p           : 0.0                                                                                                                                   
`flshattF` is not supported because:
    dtype=torch.float32 (supported: {torch.bfloat16, torch.float16})                                                                                     
    attn_bias type is <class 'torch.Tensor'>
`tritonflashattF` is not supported because:                                                                                                              
    dtype=torch.float32 (supported: {torch.bfloat16, torch.float16})
    attn_bias type is <class 'torch.Tensor'>                                                                                                             
`cutlassF` is not supported because:
    attn_bias.shape[-1] % 8 != 0                                                                                                                         
`smallkF` is not supported because:                                         
    bias with non-zero stride not supported    

In my case, attn_bias is indispensable and it is hard to always satisfy that attn_bias.shape[-1] % 8 == 0, so how could I benefit from this repo? Thanks.

[danthe3rd]

Hi,
Thank for opening this issue. That's something we can work on (see #683).
What type of bias do you need? Is it a learnable bias

[zen-d]

@danthe3rd Thanks a lot for your prompt reply! #683 is highly related. In that thread I notice you may work on it #683 (comment).
First, may I know when the support for a attn_bias of torch.Tensorwith attn_bias.shape[-1] % 8 != 0 is scheduled? Would it be a very recent plan?
Second, if you could also support a learnable attn_bias, it would become more attractive.

[danthe3rd]

The bias is currently learnable :) We just need to add this padding support. Hopefully we can get that out next week

[zen-d]

Wow, fantastic! Look forward to seeing the padding support soon to relax the shape constraint.

[danthe3rd]

It's merged in b6be33a

[zen-d]

@danthe3rd Thanks! Looks good, but I don't have free GPUs temporarily. I will try on the new feature ASAP.

[zen-d]

@danthe3rd By following these hints to do padding and slicing, I'm able to run the model now. The memory burden is significantly alleviated. Thanks for your awesome job! I will continue to monitor the training process and the final accuracy.

HINT: To use an attn_bias with a sequence length that is not a multiple of 8,
you need to ensure memory is aligned by slicing a bigger tensor.
Example: use attn_bias = torch.zeros([1, 1, 5, 8])[:,:,:,:5] instead of torch.zeros([1, 1, 5, 5])

[zen-d]

Unfortunately, the training diverges in the middle (loss becomes NaN), which did not happen in the original attention-based model. Would you like to share some insights about that? Thanks.

[danthe3rd]

I don't have specific idea for this, but you can detect more precisely where the nan is coming from with the anomaly detection:

torch.autograd.set_detect_anomaly(mode=True, check_nan=True)

[zen-d]

Thanks for providing the suggestion. The only difference is the attention implementation in this controlled experiment, but I am not sure of the specific reason temporarily. I will dive deep into the issue. :)

[danthe3rd]

Also - this is running in f32 it looks like? Otherwise you might want to try to train with f32 to see if it's related to the numerical precision

[zen-d]

Yes, for safety, I am training with FP32 numerical precision now. (Similar to my experience, AMP training seems to have more chance of NaN for Transformer-based models.)

[Shannen3206]

Yes, for safety, I am training with FP32 numerical precision now. (Similar to my experience, AMP training seems to have more chance of NaN for Transformer-based models.)

I meet the same question, and i found that use fp16 can solve this problem.

[Shannen3206]

@danthe3rd By following these hints to do padding and slicing, I'm able to run the model now. The memory burden is significantly alleviated. Thanks for your awesome job! I will continue to monitor the training process and the final accuracy.

HINT: To use an attn_bias with a sequence length that is not a multiple of 8,
you need to ensure memory is aligned by slicing a bigger tensor.
Example: use attn_bias = torch.zeros([1, 1, 5, 8])[:,:,:,:5] instead of torch.zeros([1, 1, 5, 5])

Hi,
I found that use this method may cause the inference speed lower.#853
Do you have any good way?