英文原文:How to train your own Large Language Models
大型语言模型,如OpenAI的GPT-4或Google的PaLM,已经席卷了人工智能领域。然而,大多数公司目前没有能力训练这些模型,并且完全依赖于只有少数几家大型科技公司提供技术支持。
在Replit,我们投入了大量资源来建立从头开始训练自己的大型语言模型所需的基础设施。在本文中,我们将概述我们如何训练LLM(Large Language Models),从原始数据到部署到用户面向生产环境。我们将讨论沿途遇到的工程挑战以及如何利用我们认为构成现代LLM堆栈的供应商:Databricks、Hugging Face和MosaicML。
虽然我们的模型主要是针对代码生成用例设计的,但所讨论的技术和教训适用于所有类型的LLMs,包括通用语言模型。在未来几周和月份中,我们计划深入探讨过程中繁琐细节系列博客文章。
Replit的AI团队经常被问到“为什么要训练自己的模型?”有很多原因可以解释公司决定训练自己的LLM,包括数据隐私和安全性、对更新和改进具有更大的控制力等。
在Replit,我们主要关注定制化、降低依赖性和成本效益。
LLMs需要大量的数据来进行训练。训练它们需要构建强大的数据管道,这些管道高度优化,同时又足够灵活,可以轻松地包含新的公共和专有数据来源。
在整个技术栈中,我们的主要数据源可以在 Hugging Face的The Stack 获得。Hugging Face 是一个很好的数据集和预训练模型资源。他们还提供了一系列有用的工具作为 Transformers 库的一部分,包括标记化、模型推理和代码评估等工具。
The Stack由BigCode项目提供。有关数据集构建的详细信息可在Kocetkov et al. (2022)中找到。经过去重处理,数据集的1.2版本包含约350种编程语言编写的大约2.7 TB的许可源代码。
Transformers库在抽象化模型训练中处理数据规模等多个挑战方面表现出色。然而,我们发现它对于我们的流程来说还不够,因为我们需要更多地控制数据并能够以分布式方式进行处理。
当需要进行更高级的数据处理时,我们使用Databricks来构建我们的流水线。这种方法也使得我们能够轻松地将其他数据源(如Replit或Stack Overflow)引入到我们的过程中,在未来的迭代中计划这样做。
第一步是从Hugging Face下载原始数据。我们使用Apache Spark在每种编程语言之间并行化数据集构建器进程。然后,我们重新分区数据,并以优化设置为下游处理重写成parquet格式。
接下来,我们转向清理和预处理数据。通常情况下,去重复和修复各种编码问题很重要,但The Stack已经使用Kocetkov et al.(2022年)概述的近似去重技术为我们完成了这项工作。然而,在开始引入Replit数据到我们的流水线时,必须重新运行去重复过程。这就是拥有像Databricks这样的工具所带来好处之一,在其中可以将The Stack、Stackoverflow和Replit 数据视为较大数据湖中的三个来源,并根据需要在下游过程中利用它们。
使用Databricks 的另一个好处是可以对底层数据运行可扩展且易于跟踪分析。 我们对所有类型的统计信息进行汇总统计,并检查长尾分布,并诊断任何问题或不一致性 。所有这些都在Databricks笔记本中完成,这些笔记本还可以与MLFlow集成,以跟踪和重现我们沿途的所有分析。这一步相当于对数据进行定期X光检查,也有助于为预处理采取各种步骤提供信息。
对于预处理,我们采取以下步骤:
在进行分词之前,我们使用相同数据的随机子样本训练自己的定制词汇表来进行模型训练。自定义词汇表使我们的模型更好地理解和生成代码内容。这将导致改进模型性能,并加快模型训练和推断速度。
这一步是整个过程中最重要的步骤之一,因为它在我们流程(数据管道、模型训练、推断)的所有三个阶段中都被使用到。这凸显了拥有一个强大而完全集成化基础设施对于您的模型训练过程至关重要。
我们计划在未来的博客文章中深入探讨标记化。从高层次上看,需要考虑一些重要事项,如词汇量大小、特殊token以及保留用于哨兵token的空间。
一旦我们完成了自定义词汇表培训,就会对数据进行标记化处理。最后,我们构建培训数据集并将其写入优化用于输入到模型培训过程中分片格式文件当中。
我们使用MosaicML来训练我们的模型。在此之前,我们曾经部署过自己的训练集群,但是发现MosaicML平台为我们提供了一些关键优势。
多个云服务提供商。Mosaic使我们能够利用不同云服务提供商的GPU,而无需设置帐户和所有必需的集成开销。
LLM训练配置。Composer库具有多种针对各种模型和不同类型培训目标进行调整的配置。
托管基础设施。他们管理的基础设施为我们提供了编排、效率优化和容错(即从节点故障中恢复)。
在确定模型参数时,我们考虑到模型大小、上下文窗口、推理时间、内存占用等方面之间的各种权衡取舍。较大的模型通常具有更好性能,并且更适合于迁移学习。然而这些模型需要更高计算要求进行培训和推理, 后者对于我们尤其重要. Replit是一个云原生IDE,其性能感觉就像桌面本地应用程序一样快速, 因此代码完成模型需要非常快速. 出于这个原因, 我们通常会选择具有较小内存占用量和低延迟推理的较小模型。
除了模型参数之外,我们还可以从各种训练目标中选择,每个目标都有其独特的优点和缺点。最常见的培训目标是下一个令牌预测。这通常对于代码完成效果很好,但无法考虑文档更深层次上下文。这可以通过使用“填充-中间”目标来减轻,其中在文档中掩盖一系列令牌,并且模型必须使用周围上下文来预测它们。另一种方法是UL2(无监督潜在语言学习),它将不同的客观函数作为去噪任务进行语言模型培训,在该任务中,模型必须恢复给定输入的缺失子序列。
一旦我们确定了模型配置和训练目标,就会在多节点GPU集群上启动训练运行。我们能够根据正在训练的模型大小以及希望完成训练过程的速度来调整为每个运行分配的节点数。运行大规模GPU集群非常昂贵,因此重要的是我们以最有效的方式利用它们。我们密切监视GPU利用率和内存,以确保从计算资源中获得最大可能使用。
我们使用Weights&Biases监控培训过程,包括资源利用情况以及培训进展情况。 我们监视损失曲线,以确保该模型在培训过程中每一步都有效地学习。 我们还观察损失峰值。 这些是损失值突然增加并通常表示底层培训数据或模型架构存在问题。 由于这些事件通常需要进一步调查和潜在调整,在我们的流程中强制执行数据确定性,因此可以更轻松地复制、诊断和解决任何此类损失峰值潜在来源问题。
为了测试我们的模型,我们使用了 Chen et al. (2021).所描述的 HumanEval 框架的变体。我们使用该模型生成一个 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赏金活动收集到的数据进行指导调整。