简介
大型语言模型，如OpenAI的GPT-4或Google的PaLM，已经席卷了人工智能领域。然而，大多数公司目前没有能力训练这些模型，并且完全依赖于只有少数几家大型科技公司提供技术支持。
在Replit，我们投入了大量资源来建立从头开始训练自己的大型语言模型所需的基础设施。在本文中，我们将概述我们如何训练LLM（Large Language Models），从原始数据到部署到用户面向生产环境。我们将讨论沿途遇到的工程挑战以及如何利用我们认为构成现代LLM堆栈的供应商：Databricks、HuggingFace和MosaicML。
虽然我们的模型主要是针对代码生成用例设计的，但所讨论的技术和教训适用于所有类型的LLMs，包括通用语言模型。在未来几周和月份中，我们计划深入探索过程中繁琐细节并发布一系列博客文章。
为什么要训练你自己的LLMs？
Replit的AI团队经常被问到“为什么要训练自己的模型？”有很多原因可以解释公司决定训练自己的LLM，包括数据隐私和安全性、对更新和改进具有更大的控制力等。
在Replit，我们主要关注定制化、减少依赖以及成本效益。
- 定制化。训练一个定制化模型使我们能够根据特定需求和要求进行调整，包括平台特定功能、术语和上下文，在通用模型（如GPT-4）或甚至代码专用模型（如Codex）中可能无法涵盖。例如，我们的模型经过了针对Replit上流行的某些基于Web的编程语言（包括Javascript React(JSX) 和Typescript React(TSX)）进行了优化。
- 减少依赖。虽然我们会根据任务选择正确的模型，但是我们认为减少对只有一小部分AI提供商依赖是有好处的。这不仅适用于Replit还适用于更广泛开发者社区。这就是为什么我们计划开源一些自己训练出来的模型，并且如果没有训练它们将无法实现。
- 成本效率。尽管成本将持续下降，但LLMs对于全球开发者社区来说仍然是难以承受的。在Replit，我们的使命是将下一个十亿软件创作者带到线上。我们相信，在印度用手机编码的学生应该能够访问与硅谷专业开发人员相同的AI。为了实现这一目标，我们训练定制化模型，这些模型更小、更高效，并且可以大幅降低成本进行托管。
数据管道
LLMs需要大量的数据进行训练。为了训练它们，需要构建强大的数据管道，这些管道高度优化，同时又足够灵活，可以轻松地包含新的公共和专有数据来源。
堆栈
我们从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，而无需设置帐户和所有必需的集成的开销。
- LLM 训练配置。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的早期阶段。我们有很多改进要做，还有很多难题需要解决。随着语言模型的不断进步，这种趋势只会加速。将会有一系列与数据、算法和模型评估相关的新挑战。
文章出处登录后可见！
