使用deepspeed，transformers，safetensor中常见的训练精度，共享权重问题

大路诗人路小果

已于 2024-04-03 20:03:10 修改

阅读量1.4k

点赞数 5

CC 4.0 BY-SA版权

文章标签：深度学习人工智能

于 2024-04-03 19:36:53 首次发布

本文链接：https://blog.csdn.net/a14285700/article/details/137349438

使用deepspeed可能需要注意精度问题

在这里插入图片描述

混合精度，LayerNorm

虽然deepspeed有混合精度训练的功能，但是对于网络上各种奇奇怪怪的代码的DIY转化中，他还是很弱小的。它的精度问题，使用deepspeed如果模型中有部分模型使用的是half精度，那么整个模型都会使用half精度，即使是nn.LayerNorm这样新创立的层。因为我们通常可能在计算权重的时候使用half，在LayerNorm的时候使用float32这样更好的归一化，防止梯度因为精度的问题消失或者爆炸。所以通常建议使用float32精度进行计算。但是这样的强制数据类型转化，在deepspeed中就会因为将模型的全部精度都降低而难以实现。
这个难题可以解决，但是需要特别设置，尤其是原来代码没有这样设置的时候。

class LayerNorm(nn.LayerNorm):
    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
        super(LayerNorm, self).__init__(normalized_shape, eps=eps, elementwise_affine=elementwise_affine)
        # 确保权重和偏置初始化为float32，即使之后模型转换为fp16
        if self.elementwise_affine:
            self.weight.data = self.weight.data.float()
            self.bias.data = self.bias.data.float()
    
    def forward(self, x: torch.Tensor):
        # print(f"这是是layernorm")
        # embed()
        orig_type = x.dtype
        # ret = super().forward<