简介
像OpenAI的GPT-4和谷歌的PaLM这样的大型语言模型已经席卷了人工智能世界。然而，大多数公司目前还没有能力训练这些模型，完全依赖于少数几个大型科技公司提供技术。
在Replit，我们大力投资建设训练自己的大型语言模型所需的基础设施。在这篇博客文章中，我们将概述如何从原始数据到部署在面向用户的生产环境中训练LLM。我们将讨论沿途遇到的工程挑战，以及我们如何利用我们认为构成现代LLM技术栈的供应商：Databricks、Hugging Face和MosaicML。
虽然我们的模型主要用于代码生成，但这里讨论的技术和经验教训也适用于所有类型的LLM，包括通用语言模型。我们计划在接下来的几周和几个月里深入探讨我们过程中的具体细节。
为什么要训练自己的LLM？
Replit的AI团队最常见的问题之一就是“为什么你们要训练自己的模型？”公司决定训练自己的LLM的原因有很多，从数据隐私和安全到对更新和改进的更多控制。
在Replit，我们主要关注定制化、降低依赖和成本效益。
- 定制化。定制模型训练可以让我们根据特定需求和要求进行定制，包括平台特定功能、术语和上下文，这些在GPT-4等通用模型或者像Codex这样的代码特定模型中都不会涵盖得很好。比如，我们的模型经过训练，能更好地处理在Replit上受欢迎的特定基于网络的语言，包括Javascript React（JSX）和Typescript React（TSX）。
- 降低依赖。虽然我们总是根据手头的任务选择合适的模型，但我们认为减少对少数AI供应商的依赖有好处。这不仅对Replit有好处，对更广泛的开发者社区也是如此。这就是为什么我们计划开源我们的一些模型，而如果没有训练它们的方法，我们是无法做到这一点的。
- 成本效益。尽管成本会继续降低，但LLM对全球开发者社区来说仍然昂贵得令人望而却步。在Replit，我们的使命是让下一代十亿软件创作者上线。我们认为，印度一名在手机上编程的学生应该能够使用与硅谷专业开发者相同的AI。为了实现这一目标，我们训练定制模型，这些模型更小、更高效，并且可以以大大降低的成本进行托管。
数据管道
LLM需要大量数据进行训练。训练它们需要构建健壮的数据管道，这些管道需要高度优化，同时足够灵活，以便轻松地包含新的公共和专有数据来源。
技术栈
我们从Hugging Face上的The Stack开始，作为我们的主要数据来源。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地址和密钥，对数据进行匿名处理。
- 我们使用许多启发式方法来检测并删除自动生成的代码。
- 对于部分语言，我们会移除无法编译或无法使用标准语法解析器解析的代码。
- 我们根据平均行长度、最大行长度和字母数字字符百分比来过滤文件。
分词和词汇训练
在分词之前，我们使用与模型训练相同的随机子样本数据训练我们自己的定制词汇。定制词汇使我们的模型能够更好地理解和生成代码内容。这可以提高模型性能，并加速模型训练和推理。
这个步骤是整个过程中最重要的一个，因为它被用于我们过程的三个阶段（数据管道、模型训练、推理）。这凸显了拥有一个健壮且完全集成的模型训练基础设施的重要性。
我们计划在未来的博客文章中更深入地探讨分词。从高层次来看，我们需要考虑的一些重要事项包括词汇大小、特殊令牌和为哨兵令牌保留的空间。
一旦我们训练了自定义词汇表，我们对数据进行分词。最后，我们构建训练数据集，并将其写入一个分片格式，以便优化喂入模型训练过程。
模型训练
我们使用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赏金的数据进行指令调优。
下一步计划 虽然我们取得了很大的进展，但在训练LLM方面，我们还处于起步阶段。我们还有很多改进要做，还有很多困难问题需要解决。随着语言模型的不断发展，这一趋势只会加速。数据、算法和模型评估方面将不断出现新的挑战。
如果您对训练LLM的许多工程挑战感兴趣，我们很愿意与您交流。我们非常欢迎反馈，希望听到您对我们可能遗漏的内容以及您会采取哪些不同做法的看法。
