01-ai/Yi项目微调过程中Loss Scale问题的分析与解决

问题背景

在使用01-ai/Yi项目进行模型微调时，部分用户遇到了"loss scale overflow"问题，特别是在Yi-6B基础模型上表现尤为明显，而Yi-6B-Chat版本则相对稳定。该问题表现为训练过程中不断出现loss scale溢出的警告信息，最终导致训练失败并报错"Current loss scale already at minimum - cannot decrease scale anymore. Exiting run"。

技术原理分析

这个问题本质上与混合精度训练中的梯度缩放机制有关。在混合精度训练中，为了保持数值稳定性，通常会使用动态loss scaling技术。当梯度值过大时，系统会自动降低loss scale；反之则会适当提高。当梯度持续过大时，loss scale会不断降低直至达到预设的最小值，此时系统会判定训练无法继续而终止。

问题根源探究

经过深入分析，发现该问题的核心原因在于Yi-6B基础模型的embedding层中某些特殊token(如<|im_start|>和<|im_end|>)的向量初始化值异常小。这些token在基础模型训练阶段未被充分训练，导致其对应的embedding向量数值范围与其他token差异较大。

在微调过程中，这些异常小的向量值会导致梯度计算出现数值不稳定，进而引发loss scale的持续降低。相比之下，Chat版本模型由于已经经过对话数据的微调，这些特殊token的embedding向量已经得到了适当调整，因此不会出现同样的问题。

解决方案

针对这一问题，可以采取以下几种解决方案：

1. 重新初始化特殊token的embedding向量：手动调整这些特殊token对应的embedding向量值，使其与其他token保持相近的数值范围。

2. 禁用混合精度训练：在DeepSpeed配置中将fp16设置为False，改用fp32精度进行训练。虽然这会增加显存消耗和计算开销，但能有效避免数值不稳定问题。

3. 调整训练参数：适当降低学习率或使用梯度裁剪技术，控制梯度更新的幅度。

4. 使用预训练的Chat模型：如果任务允许，直接使用已经调优过的Chat版本模型进行微调。

实践建议

对于使用V100等计算硬件的用户，如果显存允许，建议优先考虑方案2(禁用混合精度)。虽然这会降低训练速度，但能保证训练稳定性。如果必须使用混合精度训练，则应仔细检查并调整特殊token的embedding初始化值。

此外，在模型微调前，建议先对embedding层的数值分布进行简单分析，特别关注那些新增或特殊用途的token，确保它们的初始化值处于合理范围内。这可以有效预防类似问题的发生。

总结

数值稳定性是深度学习训练中的关键问题，特别是在大模型微调场景下。通过理解loss scale机制的工作原理，并针对性地调整模型参数和训练配置，可以有效解决这类问题，确保训练过程的顺利进行。