如何使用Databricks、Hugging Face和MosaicML训练大型语言模型
大型语言模型，如OpenAI的GPT-4或谷歌的PaLM，已经席卷了人工智能世界。然而，大多数公司目前没有能力培训这些模型，完全依赖少数大型科技公司作为技术提供商。在本篇文章当中，我们将用Replit团队的实际经验，来概述如何训练LLM，从原始数据到面向用户的生产环境中的部署，在这一过程中面临的工程挑战，以及我们如何利用构成现代LLM堆栈的供应商：Databricks、Hugging face和MosaicML。虽然我们的模型主要用于代码生成的用例，但所讨论的技术和经验教训适用于所有类型的LLM，包括通用语言模型。为什么要培训自己的LLM？Replit人工智能团队最常见的问题之一是“为什么要训练自己的模型？”一家公司可能决定训练自己的LLM有很多原因，从数据隐私和安全，到加强对更新和改进的控制。
- 定制
训练自定义模型使我们能够根据自己的特定需求和要求进行定制，包括GPT-4等通用模型甚至Codex等代码特定模型无法很好地涵盖的平台特定功能、术语和上下文。例如，我们的模型经过培训，可以更好地使用Replit上流行的特定基于web的语言，包括Javascript React（JSX）和Typescript React（TSX）。
- 减少依赖性
虽然我们总是会根据手头任务情况，而选用对应的模型，但我们相信，减少对少数人工智能提供商的依赖是有好处的。这不仅适用于Replit，也适用于更广泛的开发人员社区。这就是为什么我们计划开源一些模型，如果没有训练它们的手段，我们就无法做到这一点。
- 成本效益
尽管成本将继续下降，但LLM在全球开发者群体中的使用成本仍然高得令人望而却步。在Replit，我们的使命是让下十亿软件创作者在线。我们认为，在印度用手机编码的学生应该能够像硅谷的专业开发人员一样使用相同的人工智能。为了实现这一点，我们训练了更小、更高效的自定义模型，并且可以大幅降低成本。数据管道LLM需要大量的数据进行训练。培训他们需要建立强大的数据管道，这些管道经过高度优化，但足够灵活，可以轻松地包括公共和专有数据的新来源。我们从The Stack开始，它是主要数据源，可在Hugging Face上获得。Hugging Face是数据集和预训练模型的一个很好的资源。作为Transformers库的一部分，它们还提供了各种有用的工具，包括用于标记化、模型推理和代码评估的工具。该堆栈由BigCode项目提供。数据集构建的细节可在Kocetkov等人（2022）中获得。重复数据消除后，数据集的1.2版包含约2.7 TB的许可源代码，这些代码是用350多种编程语言编写的。Transformers库在抽象与模型训练相关的许多挑战方面做得很好，包括大规模处理数据。然而，我们发现这对流程来说是不够的，因为需要对数据进行额外的控制，以及以分布式方式处理数据的能力。
当需要更高级的数据处理时，我们使用Databricks来构建我们的管道。这种方法还可以很容易地将额外的数据源（如Replit或Stack Overflow）引入到我们的流程中，计划在未来的迭代中这样做。第一步是从Hugging Face下载原始数据。使用ApacheSpark在每种编程语言中并行化数据集生成器过程。然后，重新划分数据，并将其改写为镶木地板格式，并为下游处理进行优化设置。接下来，我们转向数据的清理和预处理。通常，消除数据重复并解决各种编码问题很重要，但the Stack已经使用Kocetkov等人（2022）中概述的近乎重复的技术做到了这一点。但是，一旦开始将Replit数据引入我们的管道，就必须重新运行重复数据消除过程。这就是拥有Databricks这样的工具的回报所在，在这里，可以将堆栈、堆栈溢出和Replit数据视为一个更大的数据湖中的三个源，并根据需要在下游流程中使用它们。使用Databricks的另一个好处是，可以对底层数据运行可扩展和可处理的分析。在数据源上运行所有类型的汇总统计数据，检查长尾分布，并诊断过程中的任何问题或不一致。所有这些都是在Databricks笔记本电脑中完成的，它还可以与MLFlow集成，以跟踪和复制一路上的所有分析。这一步骤相当于对数据进行周期性的x射线检查，也有助于为我们采取的各种预处理步骤提供信息。对于预处理，采取以下步骤：
- 通过删除任何个人身份信息（PII）来匿名化数据，包括电子邮件、IP地址和密钥。
- 使用许多启发式方法来检测和删除自动生成的代码。
- 对于一个子集的语言，删除无法编译或无法使用标准语法解析器解析的代码。
- 根据平均行长度、最大行长度和字母数字字符的百分比来过滤文件。
标记化与词汇训练
在标记化之前，使用用于模型训练的相同数据的随机子样本，来训练我们自己的自定义词汇表。自定义词汇表使模型能够更好地理解和生成代码内容。这导致了模型性能的提高，并加快了模型训练和推理的速度。这一步骤是流程中最重要的步骤之一，因为它用于流程的所有三个阶段（数据管道、模型训练、推理）。它强调了为您的模型培训过程提供强大且完全集成的基础设施的重要性。一旦训练了我们的自定义词汇，就能够对数据进行标记。最后，构建训练数据集，并将其写入经过优化的分片格式，以供模型训练过程使用。模型培训使用MosaicML训练我们的模型。在之前部署了训练集群之后，发现MosaicML平台带来了一些关键的好处。
- 多个云提供商。Mosaic使我们能够利用来自不同云提供商的GPU，而无需设置帐户和所有所需集成的开销。
- LLM培训配置。Composer库有许多经过良好调整的配置，用于训练各种模型和不同类型的训练目标。
- 托管基础设施。他们的托管基础设施为我们提供了协调、效率优化和容错（即从节点故障中恢复）。
在确定模型的参数时，我们考虑了模型大小、上下文窗口、推理时间、内存占用等之间的各种权衡。较大的模型通常提供更好的性能，并且更有能力进行迁移学习。然而，这些模型对训练和推理都有更高的计算要求。后者尤其重要。Replit是一个云原生IDE，其性能就像桌面原生应用程序，因此我们的代码完成模型需要闪电般快。因此，我们通常会选择内存占用较小、推理延迟较低的较小模型。除了模型参数，我们还从各种训练目标中进行选择，每个目标都有自己独特的优点和缺点。最常见的训练目标是下一个令牌预测。这通常适用于代码完成，但没有考虑文档下游的上下文。这可以通过使用“中间填充”目标来缓解，在该目标中，文档中的一系列标记被屏蔽，并且模型必须使用周围的上下文来预测它们。另一种方法是UL2（无监督潜在语言学习），它将用于训练语言模型的不同目标函数框定为去噪任务，其中模型必须恢复给定输入的缺失子序列。
一旦决定了模型配置和训练目标，我们就在GPU的多节点集群上启动训练运行。我们能够根据正在训练的模型的大小，以及希望以多快的速度完成训练过程，从而来调整为每次跑步分配的节点数量。运行一个大型GPU集群是昂贵的，所以要以最有效的方式利用它们是很重要的。并且密切监控GPU的利用率和内存，以确保计算资源得到最大可能的利用。此外，使用权重和偏差来监控培训过程，包括资源利用率和培训进度。在监控损失曲线，以确保模型在训练过程的每一步都能有效学习。同样，也注意损失的峰值。这些是损失值的突然增加，通常表明底层训练数据或模型架构存在问题。由于这些事件通常需要进一步的调查和潜在的调整，我们在流程中强制执行数据决定论，因此可以更容易地复制、诊断和解决任何此类损失峰值的潜在来源。评价为了测试我们的模型，使用了Chen等人（2021）中描述的HumanEval框架的变体。我们使用该模型生成一个给定函数签名和文档字符串的Python代码块。然后，对生成的函数运行一个测试用例，以确定生成的代码块是否按预期工作。这种方法最适合Python，有现成的评估器和测试用例。但由于Replit支持许多编程语言，需要评估各种其他语言的模型性能。然而这很难做到，而且没有广泛采用的工具或框架可以提供全面的解决方案。两个具体的挑战包括在任何编程语言中创造一个可复制的运行时环境，以及在没有广泛使用的测试用例标准（如HTML、CSS等）的编程语言中的模糊性。部署到生产模型一旦训练并评估完毕，就到了将其部署到生产中的时候了。正如我们前面提到的，我们的代码完成模型应该感觉很快，请求之间的延迟非常低。使用NVIDIA的FasterTransformer和Triton Server加速我们的推理过程。FasterTransformer是一个实现基于transformer的神经网络推理加速引擎的库，Triton是一个稳定快速的推理服务器，配置简单。这种组合在转换器模型和底层GPU硬件之间为我们提供了一个高度优化的层，并允许对大型模型进行超快速的分布式推理。在将模型部署到生产中后，能够使用Kubernetes基础设施自动扩展它以满足需求。托管推理服务器带来了一系列独特的挑战。这些包括大的伪影（即，模型权重）和特殊的硬件要求（即，不同的GPU大小/计数）。我们设计了部署和集群配置，以便能够快速可靠地发货。例如，集群旨在解决单个区域的GPU短缺问题，并寻找最便宜的可用节点。在把一个模型放在实际用户面前之前，通常会先自己测试它，并了解模型的“氛围”。在之前计算的HumanEval测试结果很有用，但没有什么比使用模型来了解它更好的了，包括它的延迟、建议的一致性和一般的帮助性。
而后继续监控模型性能和使用指标。对于模型性能，监控请求延迟和GPU利用率等指标。对于使用，跟踪代码建议的接受率，并在包括编程语言在内的多个维度上进行分解。这也允许我们对不同的模型进行A/B test，并获得一个模型与另一个模型比较的定量度量。反馈和迭代我们的模型训练平台使我们能够在不到一天的时间内从原始数据到部署在生产中的模型。但更重要的是，它能够去训练和部署模型，收集反馈，然后根据反馈快速迭代。同样重要的是，流程要对底层数据源、模型训练目标或服务器架构的任何变化保持稳健。这使我们能够在一个快速发展的领域利用新的进步和能力。虽然目前已经取得了很大的进步，但仍处于LLM培训的早期阶段，还有很多改进要做，还有很多难题要解决。随着语言模型的不断发展，这一趋势只会加速。数据、算法和模型评估将面临一系列新的挑战。