baichuan2-13b算子溢出 loss跑飞问题和定位

经验分享 Tech Blogs 模型训练-Model Training
1

1. 环境信息

硬件环境(Ascend/GPU/CPU): Ascend 910
MindSpore版本: 2.2.11
MindFormers版本: MindFormers=1.0
执行模式（PyNative/ Graph）:Graph
Python版本: 3.9.13
操作系统平台: linux

2. 问题现象

baichaun2-13b多机开pipeline后loss先下降，后续持续溢出，loss跑飞

3. 根因分析

dump溢出算子

溢出检测设置：https://www.mindspore.cn/tutorials/experts/zh-CN/r2.2/debug/dump.html#异步dump
设置只dump溢出算子
op_debug_mode：该属性用于算子溢出调试，设置成0，表示不开启溢出；设置成1，表示开启AiCore溢出检测；设置成2，表示开启Atomic溢出检测；设置成3，表示开启全部溢出检测功能。
dump成功后分析第一个溢出的算子：

SqrtGrad.Gradients_Default_network-MFTrainOneStepCell_network-_VirtualDatasetCell_ backbone-Baichuan13BV2ForCausalLM_Im head-NormHead..  
SqrtGrad.Gradients_Default_network-MFTrainOneStepCell_network-_VirtualDatasetCell_ backbone-Baichuan13BV2ForCausalLM_Im head-NormHead..  
SqrtGrad.Gradients_Default_network-MFTrainOneStepCell_network-_VirtualDatasetCell_ backbone-Baichuan13BV2ForCausalLM_Im head-NormHead..  
SqrtGrad.Gradients_Default_network-MFTrainOneStepCell_network-_VirtualDatasetCell_ backbone-Baichuan13BV2ForCausalLM_Im head-NormHead..  
SqrtGrad.Gradients_Default_network-MFTrainOneStepCell_network-_VirtualDatasetCell_ backbone-Baichuan13BV2ForCausalLM_Im head-NormHead..  
SqrtGrad.Gradients_Default_network-MFTrainOneStepCell_network-_VirtualDatasetCell_ backbone-Baichuan13BV2ForCausalLM_Im head-NormHead..  
Opdebug.Node_OpDebug.221.19.1705388945730580.output.0.json  
Opdebug.Node_OpDebug.221.19.1705388725219277.output.0.json

通过dump溢出算子，识别到是sqrt反向溢出
单算子复现溢出场景：

@Plugins.golden.register([ "sqrt_grad"])  
def _sqrt_grad(conte xt: UniversalTe stcasestructure):  
    inputx = context.input_arrays[ø]  
    inputy = context.input_arrays[1]  
    div_val = np.divide( inputy, inputx)  
    res = div_val * 0.5  
    return [res]

>>> input1  
array([[ 6.265e+02],  
        [-2.404e+00],  
        [-1.394e+00],  
        ...,  
        [-5.176e-01],  
        [-3.635e+00],  
        [-1.166e+00]], dtype=float16)  
    
>>> input1.max()  
47580.0

算子溢出导致loss跑飞。

4. 解决方案

针对溢出算子cast成float32或者bfloat16，转flaot32性能可能会劣化，推荐转bfloat16。