简介
大型语言模型（Large Language Models, LLMs），如 OpenAI 的 GPT-4 或 Google 的 PaLM，已经席卷了人工智能领域。然而，大多数公司目前没有能力从头训练这些模型，并且完全依赖于少数几家大型科技公司提供的 API 服务。
在构建自己的 AI 基础设施时，我们投入了大量资源来建立从原始数据到部署到生产环境所需的全流程管道。本文将概述如何训练 LLM，讨论沿途遇到的工程挑战以及如何利用现代 LLM 堆栈中的关键组件：Hugging Face、Databricks 和 MosaicML。
虽然我们的模型主要是针对代码生成用例设计的，但所讨论的技术和教训适用于所有类型的 LLMs，包括通用语言模型。我们将深入探讨过程中繁琐的细节，并分享一系列最佳实践。
为什么要训练你自己的 LLMs？
企业决定训练自己的 LLM 通常基于以下几个核心原因：数据隐私和安全性、对更新和改进具有更大的控制力、定制化需求以及成本效益。
- 定制化：训练一个定制化模型使我们能够根据特定业务需求进行调整，包括平台特定功能、术语和上下文。通用模型（如 GPT-4）可能无法涵盖垂直领域的专业知识。例如，我们的模型针对特定的编程语言（如 JavaScript React 和 TypeScript React）进行了优化，以提高代码生成的准确率。
- 减少依赖：减少对单一 AI 提供商的依赖是降低风险的关键策略。这不仅适用于技术团队，也适用于更广泛的开发者社区。通过开源部分自研模型，我们可以促进生态系统的多样性。
- 成本效率：尽管 API 调用成本持续下降，但对于大规模应用来说，LLM 托管仍然是巨大的开支。训练更小、更高效的定制化模型可以大幅降低推理成本，使更多用户能够访问 AI 能力。
数据管道
LLMs 需要海量的数据进行训练。构建强大的数据管道是成功的关键，这些管道必须高度优化，同时足够灵活以容纳新的公共和专有数据来源。
数据源与预处理
我们从 Hugging Face 上可用的 The Stack 作为主要数据源开始。The Stack 由 BigCode 项目提供，经过去重处理后，版本 1.2 的数据集包含大约 2.7 TB 以开放授权方式发布的源代码，涵盖超过 350 种编程语言。
Transformers 库在处理大规模数据方面表现出色，但在我们需要对数据进行额外控制并能够以分布式方式处理时，它显得不足。因此，我们使用 Databricks 来构建我们的管道。
第一步是下载来自 Hugging Face 的原始数据。我们使用 Apache Spark 将数据集构建过程在每种编程语言之间并行化。然后，我们重新划分数据，并以优化设置的 Parquet 格式重写出来，供下游处理。
接下来是清理和预处理数据。通常，重复数据删除和编码修复至关重要。一旦我们将自有数据引入管道，就必须重新运行重复数据删除过程。Databricks 允许我们将不同数据源视为更大的数据湖中的来源，并在下游流程中根据需要利用它们。
对于预处理，我们采取以下步骤：
- 匿名化：通过删除任何个人身份信息 (PII) 来匿名化数据，包括电子邮件、IP 地址和密钥。这通常涉及正则表达式匹配和命名实体识别 (NER) 模型。
- 去噪：使用多种启发式方法来检测和删除自动生成的代码或垃圾内容。
- 语法检查：对于一部分语言，我们删除了无法编译或无法使用标准语法解析器解析的代码，以确保训练数据的可执行性。
- 质量过滤：根据平均行长、最大行长和字母数字字符的百分比过滤掉文件，剔除低质量样本。
标记化和词汇训练
在标记化之前，我们使用用于模型训练的相同数据的随机子样本来训练我们自己的自定义词汇表。自定义词汇表使我们的模型能够更好地理解和生成代码内容，从而提高性能并加快训练和推理速度。
此步骤是该过程中最重要的步骤之一，因为它用于我们过程的所有三个阶段（数据管道、模型训练、推理）。在高层次上，我们必须考虑的一些重要事项是词汇量大小、特殊标记的选择以及标记保留空间的分配。
一旦我们训练了我们的自定义词汇表，我们就会标记我们的数据。最后，我们构建了我们的训练数据集并将其写成一种分片格式，该格式经过优化以用于模型训练过程，例如 TFRecord 或 WebDataset。
模型训练
我们使用 MosaicML 平台训练我们的模型。在尝试部署自己的训练集群后，我们发现 MosaicML 平台为我们提供了一些关键优势，包括支持多个云提供商的 GPU、预配置的 LLM 训练配置以及托管基础设施的容错能力。
在确定我们模型的参数时，我们考虑了模型大小、上下文窗口、推理时间、内存占用等之间的各种权衡。较大的模型通常提供更好的性能并且更能够进行迁移学习，但对训练和推理都有更高的计算要求。出于低延迟推理的需求，我们通常会选择具有较小内存占用和低延迟推理的较小模型。