浅谈领域模型训练
共 8055字,需浏览 17分钟
·
2024-08-07 06:00
作者:ybq
原文地址:https://www.zhihu.com/question/599396505/answer/3583853852
这篇文章对如何进行领域模型训练进行一个简单的探讨,主要内容是对 post-pretrain 阶段进行分析,后续的 Alignment 阶段就先不提了,注意好老生常谈的“数据质量”和“数据多样性”即可。
背景
什么是领域模型?
领域模型为什么需要专门做 post-pretrain 训练,alignment 阶段不够吗?
简单来说,领域模型 / 专家模型 就是在某一个专业领域性能特别好的模型,可能包括法律、医学、教育、role-play 等等。一般来说,领域模型比较重要的环节是 RAG,我们需要有一个特别高精的检索库,来辅助模型做一些专业的回答。这也就是说,做好 sft 和 ppo 似乎就可以了?
其实,这么想也基本正确,因为大部分领域模型所处理的任务场景,80% 都是模型的通用能力能 cover 的。以法律大模型为例,“判断是否是法律问题、总结律师发言重点、提取法官判决结果等等?” 类似的问题基本任何一个开源模型 + 几百条 sft 语料都能做的不错。
然而,领域模型的要求的准确率是远远大于 80% 的,而剩下的那 20% case 恰恰又是 sft 无论如何也做不好的。
“张三犯抢劫罪,张三买了苹果,张三杀了个人,张三睡觉,张三挪用公款,……,张三寻衅滋事李四。” 请概括张三触犯的法条?
以上面这个 case 为例,我们的通用模型大概率是会把“寻衅滋事”当做一个动作来看待,而不会把它视为一个违法行为,进而导致概括错误。可如果连这种简单 case 都调用 RAG 的话,那么显然成本高的有些过分了,何况这种情况还很难检索准确。
因此,post-pretrain 的目的便是让模型尽可能的去认识这个领域的专有名词,知道某些词汇就是这个领域的专有名词,进而让 attention 给到这些 token 一些更大的权重。法律模型需要见过所有的法律法规、医学模型需要见过所有的症状和药品名词,以此类推。
然而,大量的 paper 已经证明:续训模型的过程,大概率是“学了新的,忘了旧的”的过程。这也就是说,你提高模型在领域知识上的认知能力的同时,往往它也在丢失通用能力。前面也说了,我们有 80% 的场景时需要通用能力来覆盖的。因此,我们更加靠谱的目标是:在 post-train 阶段学习领域知识的同时,尽最大可能去避免模型的通用能力损失。(贪不了一点,大模型有太多的工作需要 trade-off)
Post-Train
如果你的 base_model 是自己训的,那后面不用看了。使用退火前的 checkpoint,沿用 pretrain 阶段的训练数据,使用类似于“91开”的数据配比去混合领域数据续训,训完再退火,然后这个工作就完成了!
pretrain 知识回顾
emm,还往下读,应该都是没有自己 model 的同学了,咱们继续探讨!
pretrain 最重要的几个东西:数据,学习率,优化器!
数据就不多说了,质量为王,记得去重!
学习率:模型的更新幅度,size越大的模型,特征空间越大、表达能力和学习能力越强,因此学习率也应该小一点(做个假设,模型 size 无限大,有无数的神经元,那么它完全可以启用没用到的神经元来学习新知识,这样就避免了遗忘旧知识这个现象的发生)。
优化器:Adam 的基础知识我就不谈了,这里只强调一点,模型的优化方向是“历史动量”和“当前数据 grad”共同决定的。也就是说,不管当前数据多 bad,优化器都会限制你做出太大幅度的更新,梯度裁剪/梯度正则类似。因此,基本可以认为我们的模型具有一定的抗噪能力。
目前,大家基本都默认使用如下三个步骤进行 pretrain:
warmup:在训练过程中,将学习率慢慢提高。(可以这么理解,你的模型还没有积攒足够的动量去抗噪,太大的学习率容易造成不可逆的影响)
linear / constant / cosine decay:维持稳定的学习率,或者缓慢衰减的学习率。
Anneal:用小学习率去学高精数据,IFT数据,逻辑数据,去提高通用逻辑能力能力和打榜能力。
(llama 和 面壁 都明确提出了退火阶段带来的能力提升)
post-pretrain 数据储备
基础知识我们回顾完了,现在开始准备数据。说句丑话,如果你没有领域模型的高精数据,也没打算去爬数据和洗数据,那神仙难救,个人建议换个方向去研究 。
好,我们已经有了领域高精数据。那 common 数据和数据配比怎么搞呢?
先说数据质量,post-pretrain 不用那么精细,我们的目标是通用能力不下降,而不是通用能力大幅度提升。qwen2 的技术报告明确指出,训了 12T 数据的模型与训了 7T 数据的模型,基本没有提升。也就说,额外的 5T 数据仅仅是因为质量稍有下降(论文里说卡的阈值更小),就没有带来任何收益。
我们大概率拿不到比 qwen2、llama3 的 pretrain 阶段质量更好的数据,因此我个人觉着不要太执着于做一份特别干净的 pretrain 数据了,你怎么洗数据都很难带来明显收益。
这里推荐几个开源数据集,感觉基本够用了!
英文 Fineweb:https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu
英文 pile:https://huggingface.co/datasets/EleutherAI/pile
中文 SkyPile:https://huggingface.co/datasets/Skywork/SkyPile-150B
中文 CCI:https://huggingface.co/datasets/BAAI/CCI2-Data
代码 the-stack-v2:https://huggingface.co/datasets/bigcode/the-stack-v2-train-smol-ids
再说数据配比,llama3 和面壁智能明确给出了他们的数据配比,基本就是一个结论:代码很重要,英文很重要(即使是中文模型也应该保证英文语料的比例,有些 paper 认为模型的 general knowledge 基本来自于英文语料,中文更多的是对齐作用)。
这里给出不权威的个人建议:中英五五开,代码不能少,领域占比看算力。(根据个人需求和个人喜好,可以提高英文比例,如果有质量较好的 math 数据或逻辑数据,也添加一些)
果粉专区
果粉的话,建议直接使用 中文 + 领域 + 下面的数据配比,相信 apple!
苹果论文:https://https://arxiv.org/pdf/2406.11794v34v3
苹果训练数据:
- https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0
https://huggingface.co/datasets/bigcode/starcoderdata
https://huggingface.co/datasets/EleutherAI/proof-pile-2/EleutherAI/proof-pile-2
实验细节
Channel loss
就一句话:做 domain post-pretrain 不看 channel loss,你不如别开 tensorboard。
你就算随机拉一个数据集过来训,大概率也是 loss 缓慢下降的现象,你能得到啥信息呢?你难道要等训了一周,才去做实验验证数据配比和学习率配置吗?
channel loss:不同数据 channel 各自的 loss。也就是说假设 1 个 batch 有 100 条数据:40条 en,30 条 cn, 20条 code, 10 条 domain,那么就绘制四条不同 channel 的 loss 曲线和一条总的 total loss 曲线。
(题外话,我本来以为 channel_loss 需要在 dataloader 侧做很复杂的操作才能实现,后来经大佬同事指点,发现只要给每条数据加一个 channel 字段,再通过 all_gather_object 去通讯下就行,代码如下)
channel_loss = {}
for step, batch in enumerate(train_dataloader):
batch = to_device(batch, device)
channel = batch['channel'][0]
del batch['channel']
outputs = model(**batch)
loss = outputs.loss
# Update channel loss
if channel in channel_loss:
channel_loss[channel][0] += loss.item()
channel_loss[channel][1] += 1
else:
channel_loss[channel] = [loss.item(), 1]
all_channel_loss = [None for _ in range(world_size)]
torch.distributed.all_gather_object(all_channel_loss, channel_loss)
merged_channel_loss = {}
for lst in all_channel_loss:
for k, v in lst.items():
if k in merged_channel_loss:
merged_channel_loss[k][0] += v[0]
merged_channel_loss[k][1] += v[1]
else:
merged_channel_loss[k] = [v[0], v[1]]
for k,v in merged_channel_loss.items():
avg_loss = v[0] / v[1] if v[1] != 0 else 0.0
print_rank_0("The Channel {} loss is {}".format(k, avg_loss), args.global_rank)
# Log channel loss to TensorBoard
if dist.get_rank() == 0:
writer.add_scalar(f'Loss/channel_{k}', avg_loss, epoch * num_batches + step)
channel_loss = {}
Loss 分析
前面提到过,pretrain 阶段有 warmup,那么 post-pretrain 当然也要有了,原因也很简单啊。我们用的开源模型并没有提供给我们 checkpoint 对应的“优化器参数”,我们无法获得以前积攒的动量啊。
continue pretrain:https://arxiv.org/pdf/2406.01375
这篇论文讨论了“post-pretrain 模型时,warmup 应该使用的数据比例”。同时它也指出,warmup 在训练充分的时候是不太重要的。但因为我们无法判断模型是不是训练充分了,所以还是老老实实的做个 warmup 吧。
敲定 warmup 的数据比例后,选择一个顺眼的学习率和数据配比,就去开始训练和观察 channel loss 吧,在最理想情况下,我们期待得到一个这样的曲线:
domain_channel 的 loss 明显下降(新知识好学)
common_channel 的 loss 基本持平,极缓慢下降(理论上会选用作为底座的 model,通用能力已经很强了,这时候很难再让他的通用能力再进步一提升了,上文提到过 qwen2 多训了 5T 通用数据但毫无收益)
结合 loss 曲线,我们再回过头来谈谈数据配比:post-pretrain 阶段最好的数据配比,就是沿用 pretrain 阶段的数据配比,很可惜,我们不可能获取到 qwen、llama 的 pretrain数据。因此,我们也别纠结数据去重了,大概率我们使用的 common 数据是人家已经训过的,我们尽可能去找质量最高的 common 数据喂给模型就可以了。
不过从 channel loss 上,我们大概率能观察和反推一些东西:
初始 loss 低:任务简单,或者模型已经训过这份数据。如果你使用的底座模型效果巨强,比如是 qwen2-72B,llama3-70B,你甚至可以断言这个数据的质量很高(能力差的小模型不能随便下定论)。当然,loss 低也有可能存在一种情况,那就是数据十分的脏,全都是重复 token 或者 固定 pattern;
初始 loss 高:好现象,说明模型没有见过这个数据。但也有数据质量很差的风险,最好再清洗下这个数据源;
loss 持平或缓慢下降:好现象,没有比这更好的现象了,基本就是我们蒙对了底座模型 pretrain 阶段使用的数据配比才会有的现象;
loss 快速下降:说明这个数据很容易学习,有可能是 domain 数据的特点比较显著,也有可能是数据比较脏,都是固定 pattern 或者具有明显的格式(提一句,llama 说任何 markdown 数据都对模型性能有损失,所以有明显格式的数据要慎重使用);
common channel loss 下降明显:你的 common 数据显然不够 common,它相对模型来说有可能更像是 domain 数据,说明当前数据配比和 pretrain 的配比偏离有点远;
domain channel loss 下降明显:好事,鼓掌欢呼;
domain channel loss 不下降:初始 loss 低说明模型大概率已经训过这份 domain 数据了,初始 loss 高还不下降,可能是数据不够干净,也可能是数据比较难学,再多训会吧;
loss 上升:和导师或领导汇报就说学习率设置的不合适,自己私下再顺带 check 一下训练代码;
综上,通过观察 loss,多做几组实验,基本能试探出哪个数据配比和哪个开源数据最适合拿来 post-pretrain。
Scaling law
真的勇士,就应该去研究 sacling law,这也就是除了llama、qwen,我还特别推崇“面壁智能”的原因,它似乎是国内唯一一家不执着于size,而是执着于“sacaling law”的公司。
这篇 domain scaling law 的论文明确指出“domain能力“和”general 能力“是相互冲突的,也就回归到了我一开始说的:我们的目标不是提高通用能力,而是去损失尽量少的通用能力。
D-CPT:https://https://arxiv.org/pdf/2406.013755
这篇论文的结论都是比较 make sense 的:
小学习率,domain 学得快,通用忘得慢;
大学习率,domain 学得快,但到一定地步后就震荡,毕竟学习能力有限;
不同 size 的模型适合不同的学习率。
文章再多的内容我就不谈了,感兴趣的读者自己拜读一下即可,scaling law 的文章都相对晦涩一些,我还没有完全读懂,不敢班门弄斧。我引用这篇 sacaling law 论文的主要原因是,一是讴歌一下做 scaling law 的大佬们,二是想表达“学习率真的很重要”这一观点,不要因为大家都在强调数据质量的重要性,就忽略了炼丹的老本行。
这里引用我的大佬同事跟我说过的一句话:“你把学习率设成 0 ,那是不是模型怎么训效果都不下降。那根据夹逼准则,你只要找到一个好学习率,你数据再烂也能训出一个通用能力只下降一丢丢的模型。”
退火
退火本身怎么做,我就不多说了,小学习率 + 高精数据。基本每一个开源模型的技术报告,都会详细指出自己的退火数据配比。
我在这里提到退火,是想强调几个观点:
退火直接能提高刷榜能力!
我们 post-pretrain 的模型,都是做过退火的,也就是说这个模型就像是刚高考完的高三学生,考试能力是人生巅峰!现在不管教他什么知识,他的考试能力都会下降。
不管你怎么训,模型的打榜能力基本都会下降,所以大家不要太过焦虑这个现象。但做 post-pretrain 之前,一定要构建好 domain 能力的评估集,证明自己的 domain 能力在提升。要不然 common 也降、domain 也降,是钱多闲得慌吗?
参考资料
面壁技术报告:https://https://shengdinghu.notion.site/MiniCPM-c805a17c5c8046398914e47f0542095a/MiniCPM-c805a17c5c8046398914e47f0542095a
Qwen2技术报告:https://arxiv.org/abs/2407.10671
llama3技术报告:https://ai.meta.com/research/publications/the-llama-3-herd-of-models/blications/the-llama-3-herd-of-models/