Enable flash attention

[s-ryosky]

Thank you for sharing your codes!

I am trying to apply flash attention to MHA as well as StreamPETR, but a loss becomes NaN during the training.
Have you ever encountered this phenomenon?
And do you know of a solution to this problem?

[AlmoonYsl]

@s-ryosky Hi,
This problem is caused by the fp16 calculation used in FlashAttention .You can try to use memory-efficient attention implemented by the xformers for fp32 calculation in attention.

[s-ryosky]

@AlmoonYsl
Thank you for your reply. I'll check it.

However, I don't think that object-wise position embedding itself is incompatible with flash-attention.
Compared to StreamPETR, there seems to be differences regarding inverse sigmoid and pos2posemb3d used during position embedding.
Do you think this is one of the reasons for the training instability?

[AlmoonYsl]

@AlmoonYsl Thank you for your reply. I'll check it.

However, I don't think that object-wise position embedding itself is incompatible with flash-attention. Compared to StreamPETR, there seems to be differences regarding inverse sigmoid and pos2posemb3d used during position embedding. Do you think this is one of the reasons for the training instability?

I think there may be a numerical overflow problem when calculating the position embedding and its gradient in fp16.