我用 Telegram 消息微调 LLM 的实验总结

重点：欢迎查看 Hacker News 上关于本文的讨论

对大多数与我交流的人来说，在绝大部分时间里，我不过是另一个基于文本的程序。如果输入和输出都如此简单，那么我能否被模型取代？要做到这一点，模型不仅需要理解我的写作风格，还要非常了解我。最好的信息来源是我的 Telegram，因为我每天都使用它，其中的聊天记录几乎涵盖了我的所有想法和行为。

最直接的方法是提取我的所有消息，将它们加载到 ChatGPT 的上下文中，并要求它在回复新消息时利用这些信息模仿我的风格。然而，这种方法受限于上下文窗口的大小，需要我预处理消息并提取要点。我不想费这个功夫，或许可以使用检索增强生成（RAG），在需要时提取必要的信息。但根据我的经验，从聊天会话这类多样化的数据中进行检索，通常需要对检索模型进行有监督微调，而我并不想创建这样的数据集。因此，微调似乎是最佳选择。它在多个方面都很理想：既能捕捉我的写作风格，又有望从我的所有消息中积累知识，无须筛选哪些内容重要。

OpenAI 提供微调能力，但由于我要使用自己的私人消息，因此不想使用任何第三方微调服务。所以，我需要选择一个基础模型。根据 Hugging Face Open LLM Leaderboard，排名靠前的小型模型（参数量不超过 13B）之一是 Mistral 7B，它的表现甚至优于 Llama 2 13B。现在的问题是：LoRA 是否足够，还是必须进行全量微调？多项对比研究 [1] [2] 表明，LoRA 比全量微调稍差，但在大多数情况下仍然够用。不过，针对我这种特定任务（俄语 + 聊天），我找到了一篇论文，其中研究人员对 Llama 进行了中文指令微调，其复杂度与我的目标相近。他们发现，在此前未经指令微调的基础模型上，基于 LoRA 的微调不如全量微调有效；但如果模型已经接受过指令微调，基于 LoRA 的微调也能取得相当的效果。对我而言，这意味着，要么对基础模型进行全量微调，要么在一个已经针对俄语聊天完成微调的模型上使用 LoRA。由于找不到针对俄语聊天微调过的模型，我准备先在针对英语聊天微调过的模型上尝试 LoRA，例如微调版 Mistral 模型 Dolphin。

先在针对英语聊天微调过的 Dolphin 上使用 LoRA

如果质量不够好，再尝试对 Mistral 进行全量微调

与电子邮件相比，Telegram 等即时通信应用有一个独特之处：对话的消息流。你与联系人之间的消息通常不会严格地一来一回、逐条交替。相反，你经常会连续发送几条消息，然后对方再连续回复几条。这些消息通常也很短。我希望在数据中保留这种自然的对话风格。

Telegram 内置了将所有聊天记录导出为 JSON 的功能。经过一些筛选，并将消息分组为会话后，我整理出了过去五年使用 Telegram 所产生的数据。最终得到来自 466 个聊天的 15,789 个会话，平均每个会话包含 8.51 条消息。在数据结构方面，我选择了 ChatML 提示词格式。下面是一个示例会话（由俄语翻译而来）：

<|im_start|>John Smith >>> 靠，怎么都绕不过 135 的时间限制 >>> 已经在尝试把所有东西都做到极致优化了，但还是不行<|im_end|> <|im_start|>Alexander Smirnov >>> 嗯，我也是 >>> 你还是在用同一个思路吗？<|im_end|> <|im_start|>John Smith >>> 不知道，我觉得咱俩思路差不多 >>> 就像你说的 >>> 在 try 里用反转后的字符串，然后试着从里面找出点什么 >>> 感觉真是狗屎，因为 z 函数把一切都毁了……………………<|im_end|> <|im_start|>Alexander Smirnov >>> 不明白这里怎么会用到 z<|im_end|> <|im_start|>John Smith >>> 不知道，感觉我反正都是迭代处理的，不过确实得反转一些字符串来构建 z 函数 >>> 而且这就是个莫名其妙的解法 >>> 从讨论里看到的<|im_end|> <|im_start|>Alexander Smirnov >>> 明白了<|im_end|>

我的数据整理器会确保仅根据某个人的回复计算损失。预测接下来谁会发言相对简单，我们并不希望模型把重点放在学习这件事上。因此，对话中计算损失的部分会以粗体突出显示。

你可能会注意到，计算损失时不仅使用了我的回复，也使用了其他人的回复。这是有意为之。这样一来，模型不仅能够扮演我，也能扮演那些经常和我聊天的人！

我会通过两种聊天方式测试模型。第一种方式是让模型假装成我，而我从不同朋友的视角与“自己”聊天。第二种方式是我以自己的身份聊天，让模型扮演我的朋友。对话始终以相同的两条消息开始：“hey”和“what’s up?”（俄语为“прив”和“как дела?”）。模型生成的语句以及它所扮演的人物都会突出显示。所有对话最初都会使用俄语，并可通过点击“original”详情按钮查看原文。测试时我会使用 oobabooga/text-generation-webui。

首先，我想看看未经我额外训练的通用对话微调版 Mistral 模型如何应对这项任务。

嘿 >>> 最近怎么样？ >>> 我挺好的，你呢？ >>> 没什么特别的 >>> 你今天做什么了？ >>> 我工作了。 >>> 然后用 Python 做了几道题，还学习了数据处理库。 >>> 晚上我要去参加摔跤训练。 >>> 你参加过体育运动吗，或者有什么爱好吗？ >>> 明白了 >>> 今天想看电影吗？ >>> 也许吧，如果我有时间的话。 >>> 你最喜欢什么音乐？电子、摇滚，还是其他类型？

прив >>> как дела? >>> все хорошо, а у тебя что нового? >>> да ничего >>> чего сегодня делал? >>> я работал. >>> потом сделал несколько задач на python и изучал библиотеки для обработки данных. >>> в вечер займусь тренировкой на борцу. >>> ты когда-нибудь занимаешься спортом или имеешь какое-то хобби? >>> понятно >>> посмотрим кино сегодня? >>> может быть, если найдется время. >>> какая ваша любимая музыка? Электронная музыка, рок или другие жанры?

嘿 >>> 最近怎么样？ >>> 嗯，目前一切都挺好的。你有什么新鲜事吗？ >>> 编了一整天的程 >>> 你今天做什么了？ >>> 下班了 >>> 还做了些家务，现在正在休息。你也在休息，还是仍在工作？ >>> 我已经在休息了 >>> 嗯，那挺好的。我们能聊点有意思的事情吗，或者你有什么问题？

прив >>> как дела? >>> ну пока все нормально. а у вас что нового? >>> прогал целый день >>> чего сегодня делала? >>> закончила на работе >>> и сделала немного домашней работы, теперь отдыхаю. а вы тоже отдохните или еще работаете? >>> я уже отдыхаю >>> ну хорошо. можно поговорить о чем-нибудь интересном или вы имеете какие-то вопросы?

好吧，它能够组织出连贯的句子。最明显的问题是，它缺乏对对话上下文的了解，导致回复平淡而泛泛。消息没有任何鲜明风格，感觉非常基础。另一个问题是模型的俄语很差。这也在意料之中，因为模型规模太小，无法很好地泛化到英语这一主要语言之外的其他语言。此外，模型往往过于主动，几乎每句话都以一个问题结尾，而现实中的人通常不会在即时通信应用里这样交流。

让我们试着解决所有这些问题！

从训练流程和硬件要求来看，LoRA 都是一种成本较低的方法。它只训练总权重的大约 1%。我选择了 1024 的序列长度和 8 的批大小。训练在 RTX 3090 上占用了 20GB 显存，进行了三个 epoch，共耗时 5.5 小时。为此，我使用了 vast.ai，其中 GPU 的费用为每小时 0.362 美元，因此整个训练总计花费 2 美元，不包括用于实验和修复 bug 的时间。

以下是结果：

嘿 >>> 最近怎么样？ >>> 嗨）都挺好的 >>> 你呢？） >>> 我也还行 >>> 你今天做什么了？ >>> 参加了 amazon 的面试，还做了几道题 >>> 得做 3 道 >>> 酷 >>> 难吗？ >>> 不知道是什么题 >>> 还没看 >>> 我还没开始 >>> 一整天都在和团队聊天，还认识了几个新人 >>> 他们怎么样？ >>> 还行 >>> 和你们那边是同一个共享办公空间））

好多了。它确实捕捉到了这个人的讲话风格。它还识别出了特定人群之间讨论最频繁的话题。例如，与朋友 2 的对话中，焦点明显更多地集中在工作上。不过，语法仍然有问题，而且对话很快就失去了上下文。我相当确信 LoRA 在英语中可以达到合理的质量，可能不需要全量微调。但由于俄语不是该模型的原生语言，让我们尝试全量微调。

全量微调由于需要多 GPU 训练而更具挑战性。常见方法包括 ZeRO & DeepSpeed [3] 或 FSDP [4]，其中 FSDP 本质上是 ZeRO3 [5]。我决定采用 FSDP。

在实现训练管道时，我参考了 Stanford Alpaca 微调代码和 Anton Bacaj 的 Mistral 微调代码。

使用半精度 FSDP 完全分片，序列长度为 1024，微批大小为 2，八个 A100 80GB GPU 中每个都需要 63GB 显存。训练持续三个 epoch，只用了 20 分钟。虚拟机的总成本为每小时 8.88 美元，最终花费 3 美元，不包括实验和调试的时间。

嘿 >>> 怎么样？ >>> 嘿）不错 >>> 你呢？ >>> 我也不错 >>> 你今天在干什么？ >>> 玩桌游 >>> 我现在要回家 >>> 不错 >>> 什么游戏？ >>> 卡坦岛 >>> 策略游戏 >>> 非常好 >>> 谁在场？ >>> 马克斯·亚维奇和费迪亚 >>> 谁赢了？ >>> 马克斯

嘿 >>> 怎么样？ >>> 你好，不错 >>> 半小时后碰面？ >>> 好的 >>> 你想说什么？ >>> 妈呀，我还没完成模型重新计算的脚本 >>> 我需要完成它才能发给他们 >>> 这个重新计算很耗时，所以我只想立即转身开始做，不浪费时间做准备 >>> 明白了 >>> 今天的计划是什么？ >>> 半小时后碰面讨论我们要做什么 >>> 如果决定了，我就向他们提交申请并等待回复 >>> 如果还没决定，我们就继续工作 >>> 你说的"他们"是谁？ >>> 投资者 >>> https://www.ycombinator.com/howtoapply >>> 里面有两个测试）

嘿 >>> 怎么样？ >>> 嘿，我们在坐着解决问题 >>> 什么问题？ >>> 数据分析 +- >>> 这是为了什么？ >>> https://yandex.ru/cup/profile/ >>> 好像刚才变得更有趣了 >>> 不错 >>> 你今晚计划干什么？ >>> 还不知道 >>> 周末呢？ >>> 看起来我也要继续解决问题和完成一些工作

嗨
最近怎么样？
挺好的，在这儿做题呢
做什么题？
大概是数据分析
做这个干什么？
https://yandex.ru/cup/profile/
感觉就是变得更有意思了
挺酷
晚上有什么安排？
还不知道
那周末呢？
看来还是会继续做题，再处理一下工作

嗨
最近怎么样？
嗨）
挺好的
你呢？
我也还行
今天做什么了？
当然是工作了
另外
还和 Alena 见了面
跟大家一起去拉面店吃了饭
然后就回家了
对了，我们订了一家能坐 14 个人的餐厅
晚上去））
挺酷
大学那边怎么样？
嗯
学了很多
总的来说，我每天都想多学一点
现在最难的是什么？
嗯，我也不知道
就是特别耗时间
有些优化方法需要重写
所以我在想，也许不值得优化

嗨
最近怎么样？
嗨）
挺好的
你呢？
我也还行
今天做什么了？
当然是工作了
另外
还和 Alena 见了面
跟大家一起去拉面店吃了饭
然后就回家了
对了，我们订了一家能坐 14 个人的餐厅
晚上去））
挺酷
大学那边怎么样？
嗯
学了很多
总的来说，我每天都想多学一点
现在最难的是什么？
嗯，我也不知道
就是特别耗时间
有些优化方法需要重写
所以我在想，也许不值得优化

嗨
最近怎么样？
嗨，挺好的，你呢？
你已经到曼谷了吗？
嗯，已经在这儿了
今天有什么安排？
嗯，我想处理一下这些 pioneer 任务，如果你不介意的话，我们可以一起做
或者你也可以先溜，我明天再把 livechat 收尾
live chat 怎么了？
妈的，我刚刚把前端发布到了生产环境，没带那个用来生成带 logo 玩意儿的破东西，然后它就报错，让我接入 google analytics 的 js
我知道这他妈尬得要死，但我还是去回想该从哪儿弄到它
嗯，明白了
周末准备做什么？
离开前在这儿和朋友聚一下，明天早上 6 点前得赶到酒店，按照最经典的方式过上一天
晚上还有一场派对
其实我们打个电话吧，我跟你讲

嗨
最近怎么样？
嗨，挺好的，你呢？
你已经到曼谷了吗？
嗯，已经在这儿了
今天有什么安排？
嗯，我想处理一下这些 pioneer 任务，如果你不介意的话，我们可以一起做
或者你也可以先溜，我明天再把 livechat 收尾
livechat 怎么了？
妈的，我刚刚把前端发布到了生产环境，没带那个用来生成带 logo 玩意儿的破东西，然后它就报错，让我接入 google analytics 的 js
我知道这他妈尬得要死，但我还是去回想该从哪儿弄到它
嗯，明白了
周末准备做什么？
离开前在这儿和朋友聚一下，明天早上 6 点前得赶到酒店，按照最经典的方式过上一天
晚上还有一场派对
其实我们打个电话吧，我跟你讲

嗨
最近怎么样？
嗨，你还没打电话吗？）
嗨
挺好的）
我现在要去打 Dota
想来的话就一起）
也许晚一点）
周末有什么安排？
一大堆健身房/社交活动/公园
＋也许去你那儿待着
嗯，也许吧，也许）
话说你的工作怎么样了？
嗯，我在等消息
他们说截止时间了吗？
他说下周会打个电话讨论
提醒我一下，是什么职位来着？
Data Scientist
团队负责商品推荐之类的破事
创建模型
在这个职位上，你需要提升机器学习、统计数据分析、编程等诸多方面的技能
我是这么理解的

嗨
最近怎么样？
嗨，你还没打电话吗？）
嗨
挺好的）
我现在要去打 Dota
想来的话就一起）
也许晚一点）
周末有什么安排？
一大堆健身房/社交活动/公园
＋也许去你那儿待着
嗯，也许吧，也许）
话说你的工作最后怎么样了？
嗯，我在等消息
他们没说是什么时间吗？
他说下周会打个电话讨论
提醒我一下，是什么职位来着？
Data Scientist
团队负责商品推荐之类的破事
创建模型
在这个职位上，你需要提升机器学习、统计数据分析、编程等诸多方面的技能
我是这么理解的

对话变得更有趣、更有吸引力了，不过仍然存在丢失上下文的风险。模型的俄语表现有所改善，但错误依然存在。我认为，在像我这样使用有限数据针对特定任务进行微调之前，最好先用大规模俄语文本语料库对模型进行无监督微调。此外，将常见聊天对象的名字作为独立 token 加入，或许也能提升质量。

我不会说它的效果比 LoRA 好很多。与其试图学习每一位聊天对象，不如只专注于某一个人，并且只根据我的回复（或其他某个人的回复）计算损失，这样可能会更有效。

当然，我不得不精挑细选这些结果。并不是因为模型的大多数回复都不合格，而是因为其中很多只是“我晚点打给你”“忙着呢”和“好的”之类的简单回答，而这类回复在对话中本来就经常出现。尽管如此，可以明显看出，模型非常擅长模仿它所扮演之人的说话风格。它也能捕捉两个人之间经常讨论的话题。然而，它在对话上下文方面存在明显不足。如果没有完整的上下文，就很难回答“哟，所以呢？”或“你周末有什么安排？”之类的问题。也许使用 Rewind 这样的系统会有所帮助，它会记录用户在计算机上所做的一切。

你可以在我的 github 仓库中找到这个项目的代码，以及如何使用自己的 Telegram 数据转储复现该项目的说明。训练日志可以在 WandB 上查看。

微调 LLM：LoRA 还是全参数？基于 Llama 2 的深入分析（anyscale.com/blog） | ↩︎

微调 LLM：LoRA 还是全参数？基于 Llama 2 的深入分析（any