简介
大型语言模型,如OpenAI的GPT-4或谷歌的PaLM,已经在人工智能领域掀起了一场风暴。然而,大多数公司目前没有能力训练这些模型,而且完全依赖少数几家大型科技公司作为技术提供者。
在Replit,我们已经大量投资于所需的基础设施,以从头开始训练我们自己的大型语言模型。在这篇博文中,我们将概述我们如何训练LLM,从原始数据到面向用户的生产环境中的部署。我们将讨论我们在这一过程中面临的工程挑战,以及我们如何利用我们认为构成现代LLM栈的供应商: Databricks, Hugging Face, and MosaicML.
虽然我们的模型主要是针对代码生成的用例,但所讨论的技术和教训适用于所有类型的LLM,包括一般语言模型。我们计划在未来几周和几个月内,在一系列的博客文章中深入探讨我们过程中的细节。
为什么要训练你自己的LLMs?
对于Replit的人工智能团队来说,最常见的问题之一是 "为什么你们要训练自己的模型?" 一个公司可能决定训练自己的LLM的原因有很多,从数据隐私和安全到增加对更新和改进的控制。
在Replit,我们主要关心的是定制化、减少依赖性和成本效率。
- 定制化。训练一个定制的模型使我们能够根据我们的具体需要和要求进行定制,包括平台特定的能力、术语和背景,这些在通用模型(如GPT-4)或甚至特定代码模型(如Codex)中都不能很好地涵盖的。例如,我们的模型经过训练,可以更好地处理Replit上流行的特定网络语言,包括Javascript React(JJSX)和Typescript React(TSX)。
- 减少依赖性。虽然我们总是根据手头的任务使用正确的模型,但我们相信减少对少数人工智能供应商的依赖是有好处的。这不仅对Replit是如此,对更广泛的开发者社区也是如此。这就是为什么我们计划开源我们的一些模型,如果没有训练它们的方法,我们就无法做到这一点。
- 成本效率。尽管成本将继续下降,但LLMs在全球开发者社区中的使用仍然是非常昂贵的。在Replit,我们的使命是将下一个10亿软件创造者带到网上。我们相信,一个在印度用手机编码的学生应该有机会获得与硅谷的专业开发者相同的人工智能。为了使这成为可能,我们训练定制的模型,这些模型更小、更有效,并且可以以大幅降低的成本进行托管。
数据管道
LLM 需要大量数据进行训练。培训他们需要构建强大的数据管道,这些管道经过高度优化,但又足够灵活,可以轻松包含公共数据和专有数据的新来源。
堆栈
我们从Hugging Face上可用的 The Stack 作为我们的主要数据源开始。Hugging Face 是数据集和预训练模型的重要资源。它们还提供各种有用的工具作为 Transformers 库的一部分,包括用于标记化、模型推理和代码评估的工具。
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 地址和密钥。
- 我们使用多种启发式方法来检测和删除自动生成的代码。
- 对于一部分语言,我们删除了无法编译或无法使用标准语法解析器解析的代码。
- 我们根据平均行长、最大行长和字母数字字符的百分比过滤掉文件。
标记化和词汇训练
在标记化之前,我们使用我们用于模型训练的相同数据的随机子样本来训练我们自己的自定义词汇表。自定义词汇表使我们的模型能够更好地理解和生成代码内容。这会提高模型性能,并加快模型训练和推理。
此步骤是该过程中最重要的步骤之一,因为它用于我们过程的所有三个阶段(数据管道、模型训练、推理)。它强调了为模型训练过程提供强大且完全集成的基础架构的重要性。
我们计划在未来的博文中更深入地探讨代币化。在高层次上,我们必须考虑的一些重要事项是词汇量、特殊标记和标记标记的保留空间。
一旦我们训练了我们的自定义词汇表,我们就会标记我们的数据。最后,我们构建了我们的训练数据集并将其写成一种分片格式,该格式经过优化以用于模型训练过程。
模型训练
我们使用MosaicML训练我们的模型。在之前部署了我们自己的训练集群后,我们发现 MosaicML 平台为我们提供了一些关键优势。
- 多个云提供商。Mosaic 使我们能够利用来自不同云提供商的 GPU,而无需设置帐户和所有必需的集成的开销。
- 法学硕士培训配置。Composer 库有许多调整良好的配置,用于训练各种模型和不同类型的训练目标。
- 托管基础设施。他们的托管基础架构为我们提供了编排、效率优化和容错(即从节点故障中恢复)。
在确定我们模型的参数时,我们考虑了模型大小、上下文窗口、推理时间、内存占用等之间的各种权衡。较大的模型通常提供更好的性能并且更能够进行迁移学习。然而,这些模型对训练和推理都有更高的计算要求。后者对我们尤为重要。Replit 是一种云原生 IDE,其性能感觉就像桌面原生应用程序,因此我们的代码完成模型需要快如闪电。出于这个原因,我们通常会选择具有较小内存占用和低延迟推理的较小模型。
除了模型参数外,我们还从各种训练目标中进行选择,每个训练目标都有其独特的优点和缺点。最常见的训练目标是下一个标记预测。这通常适用于代码完成,但未能考虑到文档下游的上下文。这可以通过使用“中间填充”目标来缓解,其中文档中的一系列标记被屏蔽,并且模型必须使用周围的上下文来预测它们。另一种方法是 UL2(无监督潜在语言学习),它将训练语言模型的不同目标函数构建为去噪任务,其中模型必须恢复给定输入的缺失子序列。
一旦我们决定了我们的模型配置和训练目标,我们就会在 GPU 的多节点集群上启动我们的训练运行。我们能够根据我们正在训练的模型的大小以及我们希望多快完成训练过程来调整为每次运行分配的节点数。运行大型 GPU 集群的成本很高,因此以尽可能最有效的方式利用它们非常重要。我们密切监控 GPU 利用率和内存,以确保我们从计算资源中获得最大可能的使用率。
我们使用 Weights & Biases 来监控训练过程,包括资源利用率和训练进度。我们监控我们的损失曲线,以确保模型在训练过程的每个步骤中都能有效地学习。我们还关注损失峰值。这些是损失值的突然增加,通常表明底层训练数据或模型架构存在问题。因为这些事件通常需要进一步调查和潜在的调整,我们在我们的流程中强制执行数据确定性,因此我们可以更轻松地重现、诊断和解决任何此类损失峰值的潜在来源。
评估
为了测试我们的模型,我们使用 Chen 等人中描述的 HumanEval 框架的变体。(2021)。给定函数签名和文档字符串,我们使用该模型生成一段 Python 代码。然后,我们对生成的函数运行测试用例,以确定生成的代码块是否按预期工作。我们运行多个样本并分析相应的Pass@K数字。
这种方法对Python来说效果最好,因为它有现成的评估器和测试案例。但由于Replit支持许多编程语言,我们需要对各种额外语言的模型性能进行评估。我们发现这很难做到,而且没有广泛采用的工具或框架来提供全面的解决方案。两个具体的挑战包括在任何编程语言中建立一个可重复的运行环境,以及对没有广泛使用的测试案例标准的编程语言(如HTML、CSS等)的模糊性。幸运的是,"任何编程语言的可重现的运行环境 "是我们在Replit的工作!我们目前正在建立一个评估框架,允许任何研究人员插入并测试他们的多语言基准。我们将在未来的一篇博文中讨论这个问题。
部署到生产
一旦我们训练和评估了我们的模型,就可以将其部署到生产环境中了。正如我们之前提到的,我们的代码完成模型应该感觉很快,请求之间的延迟非常低。我们使用 NVIDIA 的 FasterTransformer 和 Triton Server 加速我们的推理过程。FasterTransformer 是一个为基于 transformer 的神经网络的推理实现加速引擎的库,而 Triton 是一个稳定且快速的推理服务器,易于配置。这种组合为我们在转换器模型和底层 GPU 硬件之间提供了一个高度优化的层,并允许对大型模型进行超快速分布式推理。
在将我们的模型部署到生产环境中后,我们能够使用我们的 Kubernetes 基础设施对其进行自动缩放以满足需求。尽管我们在之前的博文中讨论过自动缩放,但值得一提的是,托管推理服务器会带来一系列独特的挑战。这些包括大型工件(即模型权重)和特殊硬件要求(即不同的 GPU 大小/数量)。我们设计了我们的部署和集群配置,以便我们能够快速可靠地交付。例如,我们的集群旨在解决个别区域的 GPU 短缺问题,并寻找最便宜的可用节点。
在我们将模型放在实际用户面前之前,我们喜欢自己测试它并了解模型的“氛围”。我们之前计算的 HumanEval 测试结果很有用,但没有什么比使用模型来感受它更好的了,包括它的延迟、建议的一致性和一般帮助。将模型放在 Replit 工作人员面前就像拨动开关一样简单。一旦我们对它感到满意,我们就会翻转另一个开关并将其推广给我们的其他用户。
我们将继续监控模型性能和使用指标。对于模型性能,我们监控请求延迟和 GPU 利用率等指标。对于使用情况,我们跟踪代码建议的接受率,并将其分解到包括编程语言在内的多个维度。这也允许我们对不同的模型进行 A/B 测试,并获得一个模型与另一个模型比较的定量度量。
反馈与迭代
我们的模型训练平台使我们能够在不到一天的时间内将原始数据转化为部署在生产环境中的模型。但更重要的是,它允许我们训练和部署模型、收集反馈,然后根据该反馈快速迭代。
对于我们的流程来说,保持对底层数据源、模型训练目标或服务器架构的任何变化的鲁棒性也很重要。这使我们能够在快速发展的领域中利用新的进步和功能,在这个领域中,似乎每天都有新的令人兴奋的公告。
接下来,我们将扩展我们的平台,使我们能够使用 Replit 本身来改进我们的模型。这包括基于人类反馈的强化学习 (RLHF) 等技术,以及使用从 Replit Bounties 收集的数据进行指令调整。
下一步
虽然我们取得了很大进步,但我们仍处于训练LLMs的早期阶段。我们有很多改进要做,还有很多难题需要解决。随着语言模型的不断进步,这种趋势只会加速。将会有一系列与数据、算法和模型评估相关的新挑战。
如果您对培训 LLM 的许多工程挑战感到兴奋,我们很乐意与您交谈。我们喜欢反馈,也很乐意听取您的意见,了解我们缺少什么以及您会采取哪些不同的做法。
我们一直在 Replit AI 团队中寻找有才华的工程师、研究人员和建设者。请务必在我们的职业页面上查看空缺职位。如果您没有找到合适的角色,但认为自己可以做出贡献,请与我们联系;我们很乐意听取您的意见。
Replit - How to train your own Large Language Models