通过 Arcee 和 AWS Trainium 改变大型语言模型训练方式

作者：Mark McQuade, Malikeh Ehghaghi, Shamane Siri
日期：2024年4月29日
发布于 ： | |
**** | ****

关键要点

• Arcee创新 ：通过持续预训练（CPT）与模型合并，提高大型语言模型（LLM）的领域适应性。
• AWS Trainium协作 ：通过AWS Trainium的平台，加速了模型训练并降低了成本，同时保持数据的合规性和完整性。
• 模型训练效率 ：采用AWS Trainium芯片为LLM模型提供高性能计算，加快训练时间并减低支出。

近年来，大型语言模型（LLMs）因其有效性而受到广泛关注，促使各行各业采用这些通用LLM以提升数据表现。而在这种情况下，提升训练效率和硬件可用性显得尤为重要。在的努力下，我们专注于在客户中心化的方式下提升LLM的领域适应性。Arcee创新的持续预训练（CPT）和模型合并技术实现了LLM训练的显著提升，尤其在医疗、法律及金融领域取得了优异的评估结果。与的密切合作也极大提升了Arcee平台的表现，不仅加速了模型训练，还减少了总体花费，并在安全的AWS环境中强化了合规性和数据完整性。本文将重点展示我们如何通过使用Trainium芯片来提高持续预训练的效率。

理解持续预训练

Arcee认识到持续CPT的重要性，通过诸如PMC- LLaMA和ChipNeMo的研究证明了模型对特定领域的定制性。这些项目展示了通过领域适应预训练提高模型在各领域表现的潜力，例如医疗应用和工业芯片设计。基于这些启发，我们的CPT方法旨在使用特定领域数据集扩展基础模型的训练，例如使用PubMed数据集对Llama2模型进行调整，从而使模型适应专业领域的细微差别。为进一步提高我们的CPT效率，我们与Trainium团队合作，利用其尖端技术对Llama2模型进行增强，通过处理包含880亿个标记的PubMed数据集，这标志着我们推动创新的一个重要里程碑。通过这篇文章，我们期待分享所获得的变革性见解，欢迎大家与我们一起探索领域特定模型适配的未来，以及CPT与Trainium结合所带来的优化潜力。

数据集收集

我们遵循PMC- LLaMA论文中描述的方法，收集了包括来自语义学者API的PubMed论文及论文中引用的各种医学文本，最终汇集了880亿个标记的全面数据集。有关数据集的进一步细节，原始论文提供了深入的信息。

为准备该数据集进行训练，我们在管道中使用Llama2的标记器以实现高效处理。我们将数据组织为每行包含4096个标记，遵循的建议。

为什么选择Trainium？

像本篇文章中所描述的持续预训练技术需要高性能计算实例，但随著更多开发者使用生成式人工智能（AI）和LLM，这种需求变得更加迫切。传统上，这些工作负载通常部署在GPU上，但近年来GPU的成本和可用性显著影响了模型构建的创新。随着Trainium的推出，我们得以解锁新的技术，使我们能以更高效、更低成本继续推动模型创新。Trainium是AWS专门为帮助开发者访问高性能模型训练加速器而构建的第二代机器学习（ML）加速器，可以帮助降低多达50%的训练成本，相较于同类（AmazonEC2）实例。而Trainium在全球AWS区域的可用性使得开发者无需为了获得GPU集群来构建他们的模型而进行昂贵的长期计算预定。Trainium实例为开发者提供所需的性能，同时在训练效率与降低模型构建成本之间提供弹性优化。

设置Trainium集群

我们使用构建一个高性能计算（HPC）环境，使用Trn1计算节点运行我们的分布式机器学习训练任务（请参见）。您还可以使用诸如、（AmazonEKS）或Ray等开发者流程（了解更多，请参见）。节点启动后，我们运行一个训练任务以确认节点的工作状态，并使用slurm命令检查作业状态。在这一部分，我们使用AWS的pcluster命令运行.yaml文件来生成集群。我们的集群由16个节点组成，每个节点配置了一台trn1n.32xlarge实例，具备32GB的VRAM。

我们按照下图所示设置了ParallelCluster基础设施（）。

删除)

如上图所示，VPC内部有两个子网，一个公共子网和一个私有子网。头节点位于公共子网，计算集群（在此案例中为Trn1实例）位于私有子网。为了使私有子网中的节点能够连接到VPC外部的客户端，还需要一个NAT网关。在接下来的部分中，我们将描述如何为Trn1 ParallelCluster设置所需的基础设施。

设置环境

要设置您的环境，请完成以下步骤：

1. 安装VPC及其ParallelCluster所需的组件。有关说明，请参见。
2. 在VPC中创建并启动ParallelCluster。有关说明，请参见。

现在，您可以，将模型训练脚本提交为slurm作业。

部署到Trainium

基于Trainium的使用 SDK，支持PyTorch和TensorFlow等常见的ML框架。Neuron简化了分布式训练，并与Megatron Nemo和NeuronDistributed进行了集成。

在使用Trainium时，理解几个关键参数至关重要：

• 张量并行大小 – 这决定了的级别，特别是在变换器的自注意力计算过程中，对于优化模型加载时的内存使用至关重要（非计算时间效率）。
• NeuronCores – 每个Trainium设备有两个，八节点的配置意味着总共提供256个核心。
• 迷你批次 – 反映了每个批次处理的示例数量，这是由数据加载器决定的。
• 世界大小 – 参与训练操作的节点总数。

对这些参数有深刻了解，是希望有效利用Trainium设备的任何人所必需的。

训练模型

在本篇文章中，我们将。为了确保流畅有效的训练过程，我们遵循了以下步骤：

1. 从下载Llama 2的完整检查点（模型权重和标记器）。
2. 将这些检查点转换为与Neuron Distributed设置兼容的格式，以便在我们的训练基础设施中高效使用。
3. 计算每个时期所需的步骤数，结合有效的批大小和数据集大小，以便量身定制训练过程。
   删除)
4. 启动训练任务，仔细监控其进展和性能。
5. 定期保存训练检查点。最初，该过程可能因同步性而变得缓慢，但预计随着NeuronX团队的改进，情况会有所改善。
6. 最后，将保存的检查点转换回标准格式以供后续使用，借助脚本实现无缝转换。

有关详细信息，您可以在以下中找到完整的训练步骤实现。

清理

别忘了拆除您在本文中设置的任何资源。

结果

我们的研究重点是评估CPT增强检查点的质量。我们监控不同训练期间的PubMed验证数据集的困惑度，从而获得模型性能随时间变化的宝贵见解。

在这一过程中，我们提升了模型的能力，并希望为更广泛的社区提供有效的模型适配策略。

以下图表展示了在PMC测试数据集上，基线Llama 27B检查点与其CPT增强检查点的困惑度对比。根据这些结果，在我们研究的PubMed原始数据上进行持续预训练，导致Llama 27B检查点的提升，从而提高了模型在PMC测试集上的困惑度。

删除)

以下图表展示了Llama 2 7B模型的CPT增强检查点在不同训练标记数量下的困惑度变化。训练标记数量的增加与模型性能提升相吻合，且通过困惑度指标得以量化。

删除)

以下图表展示了基线Llama 27B模型及其CPT增强检查点的困惑度对比，分别包含和不包含数据混合。这强调了数据混合的重要性，我们添加了1%的通用标记到领域特定数据集中。CPT增强检查点结合数据混合的使用，表现在性能上优于基线Llama2 7B模型及仅使用PubMed数据训练的CPT增强检查点。

删除)

结论

Arcee通过持续预训练和模型合并的创新方式，通过与Trainium的合作，标志着在LLM，尤其是在医疗研究等专业领域中的训练进步。利用Trainium的广泛能力，我们不仅加速了模型训练过程，而且显著降低了成本，同时重视安全性和合规性，以提供安全AWS环境中的数据完整性。

我们训练实验的结果，如领域特定模型的困惑度得分提升，强调了我们方法在提升LLM在各领域表现及适用性方面的有效性。这一点尤其体现在Trainium与传统GPU设置之间时间- 训练指标的直接比较中，充分彰显了Trainium的效率和成本优势。

此外，使用PubMed数据进行领域特定训练的案例研究突显了Arcee的CPT策略在细调模型以适应高度专业化数据集方面的潜力，从而为这些领域的专业人士创造出更准确、更可靠的工具。

随着我们继续推动LLM训练的可能性边界，我们鼓励研究人员、开发者和企业抓住Trainium及Arcee方法论的可扩展性、效率及增强的安全特性。这些技术不仅促进了更有效的模型训练，也为AI驱动行业内的创新与实际应用开启了新的可能性。

Trainium的先进ML能力与Arcee在模型训练及适配中的先锋策略的结合，注定将彻底改变LLM开发的格局，使其更易获取、更经济，并将其量身定制以满足各行业不断演变的需求。

如需了解更多信息，请访问或联系我们的。

额外资源

• Arcee白皮书：
• Arcee在上的论文：
• Arcee的Mergekit仓库

参考文献

1. Gupta, Kshitij等。“大型语言模型的持续预训练：如何（重新）加热你的模型？”arXiv预印本arXiv:2308.04014 (2023)。
2. Wu, Chaoyi等。“Pmc-LLaMA：希望为医学构建开源语言模型。”arXiv预印本arXiv:2305.10415 (2023)。
3. Liu, Mingjie等。“Chipnemo：用于芯片设计的领域适应LLM。”arXiv预印本arXiv:2311.00176 (2023)。
4. Touvron, Hugo等。“Llama 2：开放基础和微调聊天模型。”arXiv预印本arXiv:2307.09288 (2023)。
5. https://aws.amazon.com/ec2/instance-types/trn1/
6. Wu, C., Zhang, X., Zhang, Y., Wang, Y., & Xie, W. (2023). Pmc-llama: 根据医学论文进一步微调llama。 arXiv预印本arXiv:2304.14454 。

关于作者

 Mark McQuade 是Arcee的首席执行官及联合创始人。Mark与他人共同创立了Arcee，旨在为企业提供行业特定的AI解决方案。这一理念源自于他在HuggingFace的经历，在那里他帮助推动了变现团队，与多家知名企业进行合作。这个一线经历使他对行业中关键痛点有了深刻的认识：不愿依赖封闭源代码API以及在不妥协数据安全的情况下训练开源模型的挑战。

 Shamane Siri博士 是Arcee的应用NLP研究负责人。在加入Arcee之前，Shamane曾在工业界和学术界工作，