本文将介绍Replit如何利用Databricks、Hugging Face和MosAIcML训练大型语言模型(LLMs)。Replit是美国一家编码平台,提供了一个可从浏览器访问的IDE,无需设置即可开始编码,此外还提供了一系列简化编码任务的其他功能。

引言

大型语言模型，如OpenAI的GPT-4或谷歌的PaLM，已经在人工智能领域掀起了狂潮。然而，大多数公司目前还没有能力训练这些模型，完全依赖于少数几个大型科技公司作为技术提供者。

Replit大力投入从零开始训练自己的大型语言模型所需的基础设施。本文中，将介绍Replit如何训练LLMs，从原始数据到部署到面向用户的生产环境。还会介绍过程中遇到的工程挑战，以及如何利用LLM服务例如：Databricks、Hugging Face和MosaicML。

虽然模型主要用于代码生成，但讨论的技术和经验教训适用于所有类型的LLMs，包括通用语言模型。后面还会陆续紧跟他们的进展，敬请关注。

为什么要训练自己的LLM？

Replit人工智能团队最常见的问题之一是：“为什么要训练自己的模型？”公司可能决定训练自己的LLMs的原因有很多，从数据隐私和安全到对更新和改进的更多控制。

在Replit，主要关注定制、降低依赖性和成本效益。

• 定制化需求：训练自定义模型允许根据特定需求和要求进行调整，包括平台特定功能、术语和上下文，这些在通用模型如GPT-4甚至代码特定模型如Codex中都不会被很好支持的。例如，模型经过训练，可以更好地处理在Replit上受欢迎Web语言，包括JAVAscript React（JSX）和Typescript React（TSX）。
• 降低依赖性：虽然需要根据手头的任务选择合适的模型，但减少对少数AI提供商的依赖是有益的。这不仅适用于Replit，也适用于更广泛的开发者社区。
• 成本效益：尽管成本会继续降低，但LLM在全球开发者社区中仍然昂贵得令人望而却步。Replit的使命是让下一个十亿软件创作者上线。为了实现这一目标，训练定制的模型，使其更小、更高效，并且可以大幅降低托管成本。

数据管道

LLM需要大量的数据进行训练。训练它们需要构建健壮的数据管道，这些管道需要高度优化，同时又足够灵活，以便轻松地包含来自公共和专有数据的新数据源。

堆栈

从作为主要数据源的堆栈开始，它可以在Hugging Face上找到。Hugging Face是数据集和预训练模型的极好资源。他们还提供了Transformers库中的许多有用工具，包括用于标记化、模型推理和代码评估的工具。

堆栈由BigCode项目提供。在去重后，数据集的1.2版包含了大约2.7 TB的许可范围内的源代码，涵盖了350多种编程语言。

Transformers库在处理许多与模型训练相关方面做得很好，包括处理大规模数据。然而，它对我们的过程还不够，因为需要对数据进行更多的控制，并能够以分布式的方式处理数据。

 

数据处理

当需要进行更高级的数据处理时，使用Databricks来构建管道。这种方法也能够轻松地将其他数据源（如Replit或Stack Overflow）引入。

第一步是从Hugging Face下载原始数据。使用Apache Spark将数据集构建器过程在每种编程语言之间并行化。然后，重新分区数据，并以parquet格式使用针对下游处理优化的设置重写数据。

接下来，清理和预处理我们的数据。通常，重要的是要对数据进行去重并修复各种编码问题，The Stack完成了这项工作。然而，将Replit数据引入管道，就需要重新运行去重过程。这时，拥有Databricks这样的工具就显得非常有价值，可以将The Stack、Stackoverflow和Replit数据视为更大数据湖中的三个源，并根据需要在下游过程中使用它们。

使用Databricks的另一个好处是可以在底层数据上运行可扩展且可追踪的分析。对数据源运行各种摘要统计信息，检查长尾分布，并诊断过程中的任何问题或不一致。所有这些都是在Databricks notebook中完成的，这些notebook还可以与MLFlow集成，以跟踪和重现在过程中进行的所有分析。这个步骤相当于对我们的数据进行定期X光检查，它还有助于了解在预处理中采取的各种步骤。

在预处理过程中，采取以下步骤：

• 通过删除任何个人身份信息（PII），包括电子邮件、IP地址和密钥，对数据进行匿名处理。
• 使用一些启发式方法来检测和删除自动生成的代码。
• 对于某些语言，删除无法编译或无法使用标准语法解析器解析的代码。
• 根据平均行长度、最大行长度和字母数字字符百分比来过滤文件。

 

分词和词汇训练

在分词之前，使用与模型训练相同的数据的随机样本来训练自己的自定义词汇表。自定义词汇表使模型更好地理解和生成代码内容。这提高了模型性能，加速了模型训练和推理。

这个步骤是整个过程中最重要的一个，因为它被用在三个阶段中（数据管道、模型训练、推理）。这凸显了为模型训练过程拥有一个强大且完全集成的基础设施的重要性。

在较高层次上，还需要考虑的一些重要事项包括词汇量大小、特殊标记和哨兵标记的保留空间。

完成自定义词汇表的训练后，对数据进行分词。最后，构建训练数据集，并将其写入一个分片格式，以便优化地输入到模型训练过程中。

模型训练

使用MosaicML来训练模型。之前已经部署了训练集群，发现MosaicML平台带来了一些关键优势。

• 多个云提供商：Mosaic能够利用不同云提供商的GPU，而无需设置帐户和所有必需的集成。
• LLM训练配置：Composer库有许多为训练各种模型以及不同类型的训练目标进行了优化的配置。
• 托管基础设施：他们的托管基础设施提供了编排、效率优化和容错（即从节点故障中恢复）。

在确定模型参数时，考虑模型大小、上下文窗口、推理时间、内存占用等之间的各种权衡。更大的模型通常具有更好的性能，且更适于迁移学习。然而，这些模型对训练和推理都有更高的计算要求。后者尤为重要。Replit是一个云原生IDE，其性能感觉像一个桌面原生应用程序，因此代码补全模型需要非常快速。因此，通常倾向于选择具有较小内存占用和低延迟推理的较小模型。

除了模型参数外，还从各种训练目标中进行选择，每个目标都具有其独特的优点和缺点。最常见的训练目标是下一个token预测。这对于代码补全通常效果很好，但无法考虑到文档后续上下文中的信息。这可以通过使用“完形填空”目标来弥补，其中文档中的一系列tokens被屏蔽，模型必须使用周围的上下文来预测它们。另一种方法是 UL2（无监督潜在语言学习），将训练语言模型的不同目标函数构建为去噪任务，其中模型必须恢复给定输入的缺失子序列。

 

一旦确定了模型配置和训练目标，就会在多节点GPU集群上启动训练。根据正在训练的模型的大小以及希望完成训练过程的速度，可以调整为每次运行分配的节点数量。运行大型GPU集群的成本很高，因此需要以尽可能高效的方式利用它们。同时密切监控GPU利用率和内存，以确保充分利用了计算资源。

使用Weights & Biases来监控训练过程，包括资源利用情况以及训练进度。监控损失曲线，以确保模型在训练过程的每个阶段都能有效地学习。另外还要关注损失值的突然增加，这通常表示训练数据或模型架构存在问题。因为这些情况通常需要进一步调查和可能的调整，在过程中实施数据确定性，以便更容易地重现、诊断和解决任何这种损失激增的潜在来源。

评估

为了测试模型，使用了HumanEval框架的一个变体。使用模型根据函数签名和文档字符串生成Python/ target=_blank class=infotextkey>Python代码块。然后，对生成的函数运行一个测试用例，以确定生成的代码块是否按预期工作。运行多个样本并分析相应的Pass@K值。

这种方法对于Python效果最好，有现成的评估器和测试用例。但是，因为Replit支持许多编程语言，需要评估更多其他语言的模型性能。这很难做到，而且没有广泛采用的工具或框架提供全面的解决方案。两个具体的挑战包括在任何编程语言中制作可复现的运行时环境，以及在没有广泛使用的测试用例标准的编程语言中存在的歧义（例如html、css等）。万幸的是，一个“可在任何编程语言中复现的运行时环境”正是Replit的特长！目前正在构建一个评估框架，允许任何研究人员接入并测试他们的多语言基准。

 

部署到生产环境

在训练和评估模型之后，是时候将其部署到生产环境了。如前所述，代码补全模型应具有快速的感觉，在请求之间具有非常低的延迟。使用NVIDIA的FasterTransformer和Triton服务器加速推理过程。FasterTransformer是一个实现加速基于Transformer的神经网络推理引擎的库，而Triton是一个稳定且快速的推理服务器，配置简单。这种组合提供了一个高度优化的层，位于Transformer模型和底层GPU硬件之间，并允许大型模型的超快分布式推理。

在将模型部署到生产环境后，能够使用Kube.NETes基础设施根据需求自动扩展它。托管推理服务器还面临着一系列独特的挑战。这些挑战包括大型工件（即模型权重）和特殊硬件要求（即不同的GPU大小/数量）。通过设计部署和集群配置，以便能够快速且可靠地发布。例如，集群设计为解决单个区域中的GPU短缺问题，并寻找最便宜的可用节点。

在将模型置于实际用户之前，还需要测试它，并了解模型的“氛围”。之前计算的HumanEval测试结果很有用，不过只有亲自使用模型才能真正了解它，包括其延迟、建议的一致性和总体帮助程度。将模型放在开发人员面前就像拨动一个开关一样简单。一旦对其感到满意，就拨动另一个开关，将其推广给其他用户。

 

持续监控模型性能和使用指标。对于模型性能，监控诸如请求延迟和GPU利用率之类的指标。对于使用情况，跟踪代码建议的接受率，并根据编程语言等多个维度进行细分。这还允许我们对不同模型进行A/B测试，并获得一个模型与另一个模型进行比较的定量指标。

反馈与迭代

模型训练平台使开发人员能够在不到一天的时间里，从原始数据到部署到生产环境的模型。但更重要的是，它使训练和部署模型，收集反馈，然后根据这些反馈快速迭代。

这个过程要保持对底层数据来源、模型训练目标或服务器架构的稳定性。这可以快速发展出领域中的新能力，每天似乎都有新改进。

接下来，是扩展平台，使开发人员能够使用Replit本身来改进模型。这包括基于人类反馈的强化学习（RLHF）等技术，以及使用来自Replit赏金任务的数据进行指令调整。

下一步

尽管Replit已经取得了很大的进展，但在训练LLMs方面仍处于非常初期阶段。还有很多改进要做，很多困难问题需要解决。随着语言模型的不断发展，这种趋势将只会加速。与数据、算法和模型评估相关的新挑战将持续出现。