如何训练你自己的大语言模型?——来自Replit一线工程师的亲身经验
本文是Replit工程师发表的训练自己的大语言模型的过程的经验和步骤总结。Replit是一家IDE提供商,它们训练LLM的主要目的是解决编程过程的问题。Replit在训练自己的大语言模型时候使用了Databricks、Hugging Face和MosaicML等提供的技术栈。这篇文章提供的都是一线的实际经验,适合ML/AI架构师以及算法工程师学习。
本文原文地址:https://blog.replit.com/llm-training
本文的翻译来自ChatGPT,感谢ChatGPT的支持!
概述
像OpenAI的GPT-4或Google的PaLM这样的大型语言模型已经在人工智能领域掀起了一股风潮。然而,大多数公司目前并没有能力训练这些模型,完全依赖于只有少数几家大型科技公司提供的技术。
在Replit,我们投入了大量资源来建立训练自己的大型语言模型所需的基础设施。在本篇博客文章中,我们将概述我们如何训练LLM,从原始数据到在用户界面生产环境中部署。我们将讨论我们在这个过程中面临的工程挑战,以及我们如何利用我们认为构成现代LLM堆栈的供应商:Databricks、Hugging Face和MosaicML。
虽然我们的模型主要用于代码生成的用例,但所讨论的技术和经验教训适用于所有类型的LLM,包括通用语言模型。我们计划在未来几周和几个月内深入探讨我们的流程细节。
为什么要训练自己的LLMs?
Replit团队经常被问到的一个问题是,“为什么要训练自己的模型?”许多公司决定训练自己的LLMs的原因有很多,从数据隐私和安全性到对更新和改进的更多控制等方面。
在Replit,我们主要关心的是定制化、减少依赖和成本效益。
定制化。训练自定义模型使我们能够根据我们的特定需求和要求进行定制,包括平台特定的功能、术语和上下文,在通用模型如GPT-4甚至代码特定模型如Codex中不会得到很好的覆盖。例如,我们的模型经过训练,能更好地处理在Replit上流行的特定基于Web的语言,包括Javascript React(JSX)和Typescript React(TSX)。
减少依赖性。虽然我们将根据具体任务使用合适的模型,但我们认为减少对仅有少数几个AI提供商的依赖是有益的。这不仅适用于Replit,也适用于更广泛的开发者社区。这就是为什么我们计划开源一些模型,而这是没有训练它们的手段是无法实现的。
成本效益。尽管成本将继续降低,LLMs对于全球开发者社区仍然是难以承受的昂贵。在Replit,我们的使命是将下一个十亿个软件创作者带入在线世界。我们相信,在印度使用手机编程的学生应该能够访问与硅谷专业开发者相同的人工智能。为了实现这一点,我们训练定制的模型,这些模型更小、更高效,可以以大大降低的成本进行托管。
数据Pipeline
训练LLMs需要大量的数据。为了训练它们,需要建立强大的数据管道,这些管道高度优化,同时又足够灵活,能够轻松地包含新的公共和专有数据源。
The Stack
我们的主要数据来源是The Stack,该数据源可以在Hugging Face上获得。Hugging Face是一个很好的数据集和预训练模型资源。作为Transformers库的一部分,它们还提供了各种有用的工具,包括分词、模型推断和代码评估工具。
The Stack是由BigCode项目提供的。有关数据集构建的详细信息可参见Kocetkov等人(2022年)的文献。经过去重后,数据集的1.2版本包含约2.7 TB的在超过350种编程语言中编写的受许可的源代码。
Transformers库在抽象出与模型训练相关的许多挑战方面做得很好,包括处理大规模数据。然而,我们发现它对我们的过程来说还不够,因为我们需要对数据具有更多的控制,并能够以分布式方式处理数据。

数据处理
当需要进行更高级的数据处理时,我们使用Databricks构建我们的数据管道。这种方法还使我们能够轻松地将其他数据源(如Replit或Stack Overflow)引入我们的过程中,这是我们计划在未来的迭代中完成的。
第一步是从Hugging Face下载原始数据。我们使用Apache Spark将数据集构建过程并行化到每种编程语言中。然后,我们重新分区数据,并使用针对下游处理进行优化的parquet格式重新编写它。
接下来,我们转向清理和预处理数据。通常,去重和修复各种编码问题很重要,但The Stack已经使用Kocetkov等人(2022年)概述的近去重技术为我们完成了这项工作。但是,一旦我们开始引入Replit数据到我们的管道中,我们将需要重新运行去重过程。这就是使用Databricks这样的工具的好处,我们可以将The Stack、Stackoverflow和Replit数据视为更大的数据湖中的三个源,并根据需要在下游处理中使用它们。
使用Databricks的另一个好处是,我们可以对基础数据运行可扩展和可跟踪的分析。我们对数据源运行各种摘要统计信息,检查长尾分布,并诊断过程中的任何问题或不一致性。所有这些都在Databricks笔记本中完成,这些笔记本还可以与MLFlow集成,以跟踪和重现我们所有的分析过程。这一步相当于对我们的数据进行定期X光检查,也有助于指导我们进行预处理的各个步骤。
对于预处理,我们采取以下步骤:
- 我们通过删除任何个人可识别信息(PII),包括电子邮件、IP地址和秘密密钥,来对数据进行匿名化处理。
- 我们使用许多启发式方法来检测和删除自动生成的代码。
- 对于某些语言的子集,我们删除无法编译或不可使用标准语法解析器解析的代码。
- 我们基于平均行长度、最大行长度和包含字母数字字符的百分比来过滤文件。

Tokenization和Vocabulary Training
在进行分词之前,我们使用模型训练时使用的相同数据的随机子样本来训练我们自己的自定义词汇表。自定义词汇表使我们的模型能够更好地理解和生成代码内容。这导致了模型性能的提高,并加速了模型训练和推理。
这一步是整个过程中最重要的步骤之一,因为它在我们的过程的所有三个阶段(数据管道、模型训练、推理)中都会用到。这凸显了为模型训练过程建立健壮且完全集成的基础设施的重要性。
我们计划在未来的博客文章中深入探讨分词。在高层次上,我们需要考虑的一些重要因素是词汇表大小、特殊标记和用于哨兵标记的保留空间。
一旦我们训练好了自定义词汇表,我们就对数据进行分词。最后,我们构建训练数据集,并将其编写成一种分片格式,以便于馈入模型训练过程。
模型训练
我们使用MosaicML进行模型训练。之前,我们部署了自己的训练集群,但我们发现MosaicML平台为我们提供了一些关键的好处。
- 多个云提供商。Mosaic使我们能够利用来自不同云提供商的GPU,而无需设置帐户和所有必需的集成。
- LLM训练配置。Composer库具有许多针对训练各种模型和不同类型训练目标进行调优的配置。
- 托管基础设施。他们的托管基础设施为我们提供了编排、效率优化和容错能力(即从节点故障中恢复)。
在确定模型参数时,我们考虑了模型大小、上下文窗口、推理时间、内存占用等多种权衡。更大的模型通常具有更好的性能,并且更能进行迁移学习。但是这些模型对训练和推理的计算要求更高,后者对我们尤其重要。Replit是一个云本地化IDE,其性能感觉像桌面本地应用程序,因此我们的代码完成模型需要非常快。因此,我们通常倾向于使用具有较小内存占用和低延迟推理的较小模型。
除了模型参数之外,我们还可以选择各种训练目标,每种目标都有其独特的优点和缺点。最常见的训练目标是下一个标记预测。这通常对于代码完成很有效,但无法考虑到文档中下游的上下文。这可以通过使用“填充中间”目标来缓解,其中文档中的一系列标记被掩盖,模型必须使用周围的上下文来预测它们。另一种方法是UL2(无监督潜在语言学习),它将用于训练语言模型的不同目标函数框架化为去噪任务,其中模型必须恢复给定输入的丢失子序列。

一旦我们确定了模型配置和训练目标,我们就在多节点GPU集群上启动训练运行。我们能够根据我们正在训练的模型的大小和希望完成训练过程的速度来调整分配给每个运行的节点数量。运行大型GPU集群很昂贵,因此我们必须尽可能有效地利用它们。我们密切监控GPU利用率和内存,以确保我们能够最大程度地利用计算资源。
我们使用Weights&Biases来监控训练过程,包括资源利用率以及训练进度。我们监控我们的损失曲线,以确保模型在训练过程的每个步骤中都能有效学习。我们还观察损失峰值。这些是损失值突然增加的情况,通常表示底层训练数据或模型架构存在问题。由于这些情况通常需要进一步调查和潜在调整,因此我们在我们的过程中强制执行数据确定性,以便更轻松地复制、诊断和解决任何此类损失峰值的潜在源。
评估
为了测试我们的模型,我们使用了Chen等人(2021)中描述的HumanEval框架的变体。我们使用模型生成给定函数签名和文档字符串的Python代码块。然后,我们对生成的函数运行测试用例,以确定生成的代码块是否按预期工作。我们运行多个样本并分析相应的Pass@K数字。
这种方法对Python最有效,具有可用的评估器和测试用例。但是,由于Replit支持许多编程语言,我们需要对各种其他语言的模型性能进行评估。我们发现这很难做到,并且没有被广泛采用的工具或框架提供完全综合的解决方案。两个具体的挑战包括在任何编程语言中产生可重现的运行时环境,并且对于没有广泛使用的测试用例标准的编程语言存在歧义(例如HTML、CSS等)。幸运的是,“在任何编程语言中产生可重现的运行时环境”是我们在Replit中的特长!我们正在构建一个评估框架,允许任何研究人员插入并测试他们的多语言基准。我们将在未来的博客文章中讨论这个问题。

部署到生产环境
一旦我们训练和评估了模型,就可以将其部署到生产环境中。正如我们之前提到的,我们的代码补全模型应该感觉非常快,请求之间的延迟应该非常低。我们使用NVIDIA的FasterTransformer和Triton Server加速推理过程。FasterTransformer是一个实现基于Transformer的神经网络推理加速引擎的库,而Triton是一个稳定且具有易于配置的快速推理服务器。这种组合为我们提供了一个高度优化的层,位于Transformer模型和底层GPU硬件之间,并允许对大型模型进行超快速的分布式推理。
将我们的模型部署到生产环境后,我们可以使用我们的Kubernetes基础架构自动缩放以满足需求。尽管我们在以前的博客文章中已经讨论过自动缩放,但值得一提的是,托管推理服务器有一套独特的挑战。其中包括大型的工件(即模型权重)和特殊的硬件要求(即不同的GPU大小/数量)。我们设计了我们的部署和集群配置,以便我们能够快速、可靠地发货。例如,我们的集群被设计为在单个区域中的GPU短缺时工作,并寻找最便宜的可用节点。
在将模型放在实际用户面前之前,我们喜欢先自己测试一下,了解一下模型的“感觉”,包括延迟、建议的一致性和总体有用性。我们之前计算的HumanEval测试结果很有用,但与模型一起工作才能更好地感受它,包括其延迟、建议的一致性和总体有用性。将模型放在Replit员工面前就像翻开一个开关一样容易。一旦我们对它感到满意,我们会翻开另一个开关,将其推出到我们的其他用户。

我们继续监控模型性能和使用指标。对于模型性能,我们监测请求延迟和GPU利用率等指标。对于使用情况,我们跟踪代码建议的接受率,并将其分解为多个维度,包括编程语言。这也允许我们进行A/B测试不同的模型,并获得一个定量的比较一个模型与另一个模型的度量。
反馈与迭代
我们的模型训练平台使我们能够在不到一天的时间内从原始数据到部署在生产环境的模型。但更重要的是,它允许我们训练和部署模型,收集反馈,然后根据反馈快速迭代。
对于我们的流程来说,保持对底层数据源、模型训练目标或服务器架构的任何更改具有鲁棒性也很重要。这使我们能够利用这个快速发展的领域中的新进展和能力,在这个领域中,每一天似乎都带来了新的令人兴奋的消息。
接下来,我们将扩展我们的平台,使我们能够使用Replit本身来改进我们的模型。这包括使用基于人类反馈的强化学习(RLHF)等技术,以及使用从Replit赏金收集的数据进行指令调整。
欢迎大家关注DataLearner官方微信,接受最新的AI技术推送
