问题现象
昇腾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)
