Chapter 28

distribution training

FP16混合精度已经成为主流大规模模型训练框架的默认选项，用于训练十亿到百亿规模的模型。但其仍太容易遇到精度问题。作为补救措施，NVIDIA Ampere GPU提供了BF16浮点格式（被BLOOM采用）来缓解这个问题。然而，BF16在其他平台上不被支持，这大大缩小了它在更广泛的应用中的潜力。

为了让更多开发者使用，GLM-130B仍然选择FP16作为其训练浮点格式。同时，这意味着GLM-130B将面临着更多的稳定性挑战。幸运的是，经过多次尝试，我们发现以下的训练策略最终有助于稳定GLM-130B的训练。

我们观察到，在训练的早期阶段，嵌入层的梯度范数明显比其他层大。根据经验，我们发现大多数训练崩溃都发生在其梯度范数激增之后。为了解决这个问题，BLOOM汇报了使用嵌入归一化（我们也发现它能稳定训练），但同时，其牺牲了相对较大的下游性能。

由于根本问题是输入嵌入层的急剧梯度，我们建议缩小输入嵌入层的梯度。实现起来相当简单。

word_embedding = word_embedding * α + word_embedding.detach() * (1 - α) 这就把梯度缩小到α。在我们的实践中，我们发现α=0.1对GLM-130B是最好的。

在我们的初步实验中，我们观察到，对于早期阶段的训练来说，缩小嵌入梯度并没有减缓收敛速度；相反，没有缩小梯度的模型会出现意外的尖峰，并在5k步左右出现训练崩溃的情况。

梯度收缩是一种避免训练崩溃的事后技术。从本质上讲，崩溃是由异常的损失"梯度"形成的，要么是由于噪声数据，要么是正向计算中的精度上溢或者下溢。

我们观察到，在大型语言模型中，注意力的计算操作是最容易上溢或下溢的。

CogView显示，不同的注意力头对其注意力分数有非常不同的数值范围，有些注意力头计算出的平均分数可以达到+1e4或-1e-3。

这种不同的数值范围会导致在softmax计算中FP16下的频繁上溢或下溢。

CogView提出了精度瓶颈放松（PB-Relax）来缓解这个问题，它在做softmax之前扣除了每个头的注意力得分矩阵中的最大绝对值。

然而，事实证明，PB-Relax在GLM-130B的训练中很慢，可能是因为在96个大小为2048*2048的注意分数矩阵中寻找最大值和操作标量对CUDA内核不友好。

最后，经过几周的艰苦探索，我们发现避免这一问题的最快和最简单的方法是在softmax计算中使用FP32。与完全的FP16计算相比，它几乎没有任何速度上的损失，但明显提高了训练的稳定性。