昇腾910FlashAttention适配alibi问题

问题现象

昇腾910上不开FA和开FA后loss对不齐。如下图中的两个曲线,前几个step loss相差较大。

问题分析

1.没有开启alibi参数
2.开启了alibi参数,但是因为海思的pse定义跟mindspore的不同,所以在送到海思前要先把scale除掉。
海思实现方法:

score = (k * v + pse) * scale 

mindspore实现方法:

score = k * v * scale + pse

解决方法

目前除以scale的操作已经封装在nn.FlashAttention算子中,可以直接使用。(ms whl包中的路径mindspore/nn/layer/flash_attention.py)
如果使用的mindspore版本比较旧,还没有封装nn.FlashAttention算子,可用下面方法规避,只针对mindformers:

方法1. 修改类mindformers.modules.layers.AlibiTensorV2的__init__方法

# self.slopes = Tensor(slopes[None, :, None, None],mstype.float32) # (num_heads, 1)  
self.slopes = Tensor(slopes[None, :, None, None] * math.sqrt(head_dim), mstype.float32) # (num_heads, 1)

方法2.修改mindformers.modules.layers.build_alibi_tensor_v2方法

# slopes = np.expand_dims(np.expand_dims(slopes, 1), 1)  
slopes = np.expand_dims(np.expand_dims(slopes, 1), 1) * math.sqrt(head_dim)