AI 文摘

增量预训练baichuan-13b-chat遇到的那些坑





作者: AINLP 来源: AINLP

前言

资源

单机两4090,如图

在这里插入图片描述

单卡24G,baichuan-13b-chat单卡推理需要至少26G,因此仅用一张卡,我们是无法加载百川13B的模型,所以,无论是推理还是训练,我们都必须并行!

deepspeed

核心思想:GPU显存不够,CPU内存来凑

虽然我们两张卡加起来有48G,按理说显存是足够的,实则不是。

就两张卡而言,分别为GPU0和GPU1,两块GPU上分别有一半模型参数,即6.5B,占用13G,在使用deepspeed的参数并行进行前向传播时,GPU0需要把自己身上的参数传给GPU1临时保存起来,参与前向传播,这时,GPU1身上的参数即为整个模型的参数,即13B,占用26G,超出了单卡显存上限,因此,当只有两块卡时,不使用offload策略的前提下,单卡的显存必须大于模型大小 * 2 ,即13B的模型,双卡并行时,必须保证单卡显存大于26G。在使用三卡并行 时,单卡显存大小理论上应大于:2 * 13 * 2 / 3 = 17.3G

第一个2表示:每块GPU上都需要保存一份自己独有的参数和其他卡发送过来的临时参数,共2份 第二个2表示:参数的数据类型是fp16,占两个字节,所以乘2 13表示:13B的模型 3表示:3块GPU

下面的视频展示了 ZeRO(包含所有三个阶段)如何执行训练步骤,包括前向传递、后向传递和参数更新。要用到多卡并行的同学建议一定要看! ZeRO-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters[1] 从上面的分析可以看出,我的硬件资源在不使用offload的情况下,是不足的,我的deepspeed参数配置:


{
    "bfloat16": {
        "enabled": false
    },
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_init": 0,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 1e5,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

可以看到,我这里使用offload_optimizer 到cpu 的策略,这就导致,训练速度慢了很多。

在这里插入图片描述

总共307w数据,batch size为1,梯度累积步为16,每个batch的输入token为1024,两卡4090,offload优化器参数到cpu,需要训练50天。

怕大家不理解什么是offload optimize,给大家贴张图

在这里插入图片描述

一、训练的坑

通过上面一顿操作,总算不OOM了!

我这里是使用Lora的增量预训练方式,给模型注入领域知识,成功启动训练之后,下面是训练了300个step后的模型文件。

在这里插入图片描述

理论上这个adapter_model.bin 就是lora的参数,实际并不安全是

可以看下每个文件的大小

在这里插入图片描述

细心的同学可以发现,adapter_model.bin 的文件非常小,看起来不像是真的lora模型参数,的确,使用deepspeed stage3策略训练得到的模型参数,还需要恢复权重才能用!如果你不恢复,直接使用PeftModel.from_pretrained 合并模型的话,会报错,如下图:

在这里插入图片描述

用以上报错各种Google 后都没找到正确答案,就问你绝不绝望!

细心的同学又发现了,保存的checkpoint 文件夹里有个 zero_to_fp32.py ,没错,这个代码就是用来恢复权重的,在命令行执行:

python zero_to_fp32.py . adapter_model.bin

输出:


[2023-07-30 10:51:11,819] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Processing zero checkpoint './global_step300'
Detected checkpoint of type zero stage ZeroStageEnum.weights, world_size: 2
Parsing checkpoint created by deepspeed==0.10.0
Reconstructed Frozen fp32 state dict with 283 params 13264901120 elements
Reconstructed Trainable fp32 state dict with 400 params 27893760 elements
Saving fp32 state dict to adapter_model.bin

再看下模型大小:

在这里插入图片描述

哇,25G,不太妙啊

妙不妙合并下就知道。

二、推理的坑

好吧,合并试下呗。显存大的同学,合并时可能不会有任何问题,也许就是速度慢了些。

不过,咱只有48G的显存,拿个13B去和这25G大小的模型合并,差不多需要50G的显存,实在是顶不住。

原因:之所以adapter_model.bin 达到25G,是因为这个参数包含了非lora层的参数,我们可以将非lora层的参数删除。

解决办法:


import torch 
ckpt = torch.load('adapter_model.bin')
for key in list(ckpt.keys()):
    if key.find('lora_') == -1:
        del ckpt[key]
torch.save(ckpt, 'adapter_model.bin')

移除其他层的权重之后,大概就100M多些,这下可以正常推理了

三、继续训练的坑

大模型的并行训练,中途完全有可能需要中断,然后继续训练,这应该很好理解,你除了需要加载最新的模型外,之前已训练过的数据,最好也需要有标记,标记出哪些数据已训练过,尽可能保证每条数据只训练过一次。

不过,考虑到时间上的因素,我这里只做了随机采样,在一定程度上保证了每条数据尽可能只参与一次训练。


class MyDataset(Dataset):
  

    def __init__(self, dataset):
        self.dataset = dataset
        self.total_len = len(dataset)
  

    def __getitem__(self, idx):
        # 随机采样
        random.seed(int(time.time()))
        sample_idx = random.sample(range(self.total_len), k=1)[0]  
        return self.dataset[sample_idx]
  

    def __len__(self):
        return self.total_len

注意:如果这里你固定seed值,每次取到的数据还是一样,一定要保证每次取数据时seed值都不一样,所以我用了时间戳。

总结

卡少是非多

训练代码:https://github.com/seanzhang-zhichen/train_baichuan_llm.git

References

[1] ZeRO-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters: https://www.microsoft.com/en-us/research/blog/ZeRO-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/

进技术交流群请添加AINLP小助手微信(id: ainlp2)

请备注具体方向+所用到的相关技术点

![](https://api.allorigins.win/raw?url=https://mmbiz.qpic.cn/mmbiz_jpg/nW2ZPfuYqSJADkmZ2IX6Z23znAibuEevotDMq9iaMxiapK7jfMibiauGFkycicAJEs6x5U9SGyDJZ0S1tRed9TPNUUDQ/640?wx_fmt=jpeg&wxfrom=5&wx_lazy=1&wx_co=1)

关于AINLP

AINLP 是一个有趣有AI的自然语言处理社区,专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享,主题包括LLM、预训练模型、自动生成、文本摘要、智能问答、聊天机器人、机器翻译、知识图谱、推荐系统、计算广告、招聘信息、求职经验分享等,欢迎关注!加技术交流群请添加AINLP小助手微信(id:ainlp2),备注工作/研究方向+加群目的。

  


![](https://api.allorigins.win/raw?url=https://mmbiz.qpic.cn/mmbiz_jpg/nW2ZPfuYqSKABHCqVVQkVYPrM4XY1vsd0iaeuXzyJnoFc8cibd5mYb4wdA3WMQtiaPVmr0XLZHMuVibqWncibpnTSnQ/640?wx_fmt=jpeg&wxfrom=5&wx_lazy=1&wx_co=1)

更多AI工具,参考Github-AiBard123国内AiBard123

可关注我们的公众号:每天AI新工具