[BUG]Nan in gradients of scaled_dot_product_attention operation with mem_efficient backend about pytorch

[drisspg]

One Update
I tried locally compiling with string(APPEND CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -fmad=false")
And the error still repros

Also while using: NanInfdetect

I was able to determine that this training loops is encouraging very large values for q, k , v

4.9956e+08, -4.9956e+08, -4.9956e+08,  ..., -4.9956e+08,
          -4.9956e+08, -4.9956e+08],
         [ 2.7200e+08,  2.7200e+08,  2.7200e+08,  ...,  2.7200e+08,
           2.7200e+08,  2.7200e+08],
           

If you set norm_first = True I am unable to reproduce the error, the fused kernels are known to be more sensitive to large values, I have seen this directly for FAv2 because of its use of FMA.. I don think thats the case here from the first update though.