在 AWS Trainium 上使用超过 100 个节点的实例集群进行 LLM 训练

关键要点

在此文章中，我们展示了如何在 AWS Trainium 平台上，利用超过 100 个节点的集群进行大型语言模型LLMLlama 27B 的训练。我们分享了最佳实践，以提高训练效率、稳定性和恢复能力，最终展示了训练成果与开放源代码版本的质量对比。

Llama 是 Meta AI 的大型语言模型LLM，其参数从 70 亿到 700 亿不等。Llama 使用基于变换器的解码器模型架构，专注于语言令牌的生成。要从零开始训练一个模型，需要一个包含数万亿令牌的数据集。Llama 系列是目前最受欢迎的 LLM 之一。然而，训练 Llama 模型在技术上具有挑战性，且过程漫长且成本昂贵。

在本文章中，我们将展示如何通过扩展到 128 个 trn132xlarge 节点，加速 LLM 模型的全预训练，以 Llama 27B 模型为例。我们将分享在 AWS Trainium 上训练 LLM 的最佳实践，包括在超过 100 个节点的集群上扩展训练、提高从系统和硬件故障恢复的效率、提升训练稳定性和实现模型收敛。我们展示了在 Trainium 上训练的 Llama 27B 的质量，与多个任务的开放源码版本相比，包括多任务语言理解、数学推理和代码生成等，都具有相当的效果。我们还展示了 Trainium 的扩展优势。

为什么分布式训练超过 100 个节点如此具有挑战性？

训练大规模 LLM 需要跨越超过 100 个节点的分布式训练，获取弹性的大规模高性能计算集群的访问权限是非常困难的。即便你获取了所需的加速计算能力，管理一个超过 100 个节点的集群、保持硬件稳定性以及实现模型训练的稳定性与收敛性也是具挑战性的。让我们逐一看看这些挑战以及如何通过 Trainium 集群解决它们，确保端到端训练的顺利进行。

分布式训练基础设施的效率和可扩展性 LLM 的训练既需要计算也需大量内存。在本篇文章中，我们将展示如何在 Trainium 上启用不同的并行训练算法，并选择最佳超参数以实现 Llama 27B 在 Trainium 集群上的最高吞吐量。我们还展示了其他内存和计算优化技术的实现，比如 合并层 和数据类型选择。根据实证数据，我们证明 Trainium 集群的成本可以相比于同类 Amazon Elastic Compute Cloud 实例降低多达 46。高效的硬件和系统恢复 在这种规模的 LLM 训练中，难免会遇到硬件或系统故障。我们展示了如何利用 NeuronX 分布式库有效启用检查点保存和自动恢复。实证结果表明，借助于自动故障恢复，硬件的有效利用率达到了 9881，相比之下，手动恢复法的利用率为 7783。训练稳定性与收敛性 细节上，像 Llama 2 这样的深度神经网络在预训练中，损失函数频繁出现尖峰可能会导致 灾难性发散。由于训练 LLM 所需的高计算成本，我们希望降低损失函数尖峰的出现，提高训练的稳定性，并实现收敛。我们展示了在 Trainium 集群上实现的最佳实践和技术，比如 scaled initialization、梯度裁剪和缓存管理。同时我们还展示了如何监控和调试以确保训练的稳定性。

Llama 27B 预训练设置

在本节中，我们将讨论设置 Llama 27B 预训练的步骤。

基础设施

设置 Llama 27B 的基础设施包括以下组件：

EC2 集群 训练集群由 128 个 trn132xlarge 实例节点组成，总共配备 2048 个 Trainium 加速器。实例之间的网络通过 8100 Gbps EFAs 连接。我们挂载了 56 TB 的 Amazon FSx 存储，以便用于数据存储、检查点保存和加载。原始训练数据被保存在 Amazon Simple Storage ServiceAmazon S3桶内。调度 我们首先使用 trn132xlarge 集群从零开始训练 Llama 27B，并通过 Amazon Elastic Kubernetes ServiceAmazon EKS进行管理。关于设置过程的详细信息，请参考 Train Llama2 with AWS Trainium on Amazon EKS。我们遵循相同的过程，但以大规模设置配置了 128 个 trn132xlarge 实例。容器构建 我们使用了自定义的 Docker 镜像，该镜像是基于以下的 训练容器 构建的，并包括 Llama 27B 的 训练源文件。我们将自定义 Docker 镜像存储在 Amazon Elastic Container RegistryAmazon ECR注册表中，并在 EKS Pods 中部署。下图展示了集群和容器设置的架构。

数据准备

训练数据集的原始格式包含大量压缩文件。为了使用该数据集，我们首先将其转换为与 Hugging Face 数据集包兼容的格式。我们使用 Apache Arrow 格式数据集的默认存储格式将所有数据合并为一个单独的文件和一个单独的数据块。此方法显著降低了 TB 级数据集的加载时间，相比于默认需要加载多个单独文件的方法。

我们首先使用一种具有 2030 TB 内存的特殊 EC2 实例下载了预处理过的训练数据集，这是一小部分完整数据集，包含 12 万亿令牌。数据下载脚本如下所示：

蓝快加速器电脑版

pythonimport os

Cache 和 tmpdir 可能会很大。确保 / 目录有足够的磁盘空间。

osenviron[HFDATASETSCACHE] = /dataset/cacheosenviron[TMPDIR] = /dataset/tmpdir

import datasetsfrom datasets import loaddataset

savepath = //arrowsavepath = ospathexpanduser(savepath)osmakedirs(savepath existok=True)

rawdatasets = loaddataset(togethercomputer/lt1T data file namegt default numproc=448)rawdatasets[train]savetodisk( savepath numshards=1 numproc=448)

数据集经过优化存储和访问的处理如下：

pythonimport pyarrow as paimport time

a = timetime()stream = pamemorymap(//arrow/trainarrow)stream = paipcopenstream(stream)table = streamreadall()print(完成步骤 1 的时间秒： timetime() a)

ca = table[text]l = catopylist()schema = paschema({text palargestring()})arr = paarray(l type=palargestring())

with paOSFile(//arrow/trainarrow wb) as sink with paipcnewstream(sink schema=schema) as writer batch = parecordbatch([arr] schema=schema) writerwrite(batch)print(完成步骤 2 的时间秒： timetime() a)

在同一实例上，我们清理了数据集并将清理后的数据集上传到 S3 桶中。然后我们利用 128 个 trn132xlarge 集群在线进行令牌化和打包例如动态填充序列和应用掩码机制。与离线打包方法相比，此在线方法显著节省了开发时间和计算资源，尤其是在使用不同大规模数据集和分词器的多个实验中。

模型超参数

我们采用了与 Llama 模型相同的训练超参数。具体而言，我们使用了余弦学习率调度器，最大学习率为 34，最小学习率为 35。训练开始的 2000 步采用线性预热。以下图展示了整体学习率调度器的图像。

我们使用了 AdamW 优化器，设置1 = 09和2 = 095。所有参数，包括归一化权重使用的权重衰减值为 01。为确保训练稳定性，我们应用了 10 的梯度范数裁剪。对于 Llama 3 等不同模型设置，这些参数需要进行调优以获得最佳性能。

分布式训练基础设施的效率和可扩展性

在训练期间，我们通过 Neuron SDK 应用了一般优化技术，如激活检查点、模型和数据并行，及计算和通信重叠，同时在 Trainium 上进行了一些特殊增强，例如使用 BF16 进行随机舍入。在本节中，我们列出了在模型预训练中采用的关键功能和配置，以提高训练效率。

模型和数据并行

Neuron 支持张量并行TP、流水线并行PP、序列并行SP和数据并行DP。对于 7B 模型，序列长度为 4096，我们发现 TP 度数为 8，PP 度数为 1，SP 度数为 8，DP 度数为 512 给出了最高的训练吞吐量。在一个 trn132xlarge 实例集群上，这意味着每个实例有四个模型副本。

我们使用了全局批量大小为 1024 个序列，最大序列长度为 4096 个令牌。每个步骤大约覆盖 400 万个令牌。梯度积累步骤为 2，最终使每个 Neuron 核的实际批量大小为 1。以下图展示了我们在训练中应用的数据并行和张量并行。

Neuron 分布式库

AWS Neuron 是用于在 AWS Inferentia 和基于 Trainium 的实例上运行深度学习工作负载的 SDK。它包括编译器、运行时和分析工具。它支持多种 数据类型，包括 FP32、BF16、FP16 和 随机舍入。Neuron SDK 通过 NeuronX 分布式库，支持 张量并行、流水线并行和 数据并行 的分布式策略。这些策略允许在保持训练模型的高精度与提高吞吐量和内存使用效率之间进行权衡。在训练过程中，我们应用了以下的功能：

选择性激活检查点 我们使用 选择性激活检查点 来提高训练效率。与完整激活检查点相比，它的内存开销略高，但能够提高整体训练吞吐量。BF16 结合随机舍入 我们比较了三种精度设置：BF16、带有随机舍入的 BF16，以及混合精度训练。根据实证，BF16 结合随机舍入显示出与 混合精度训练 相同的收敛行为，同时具有更高的训练吞吐量和更低的内存占用；而 BF16 的训练损失则发散。因此，在我们的预训练实验中选择了 BF16 结合随机舍入。合并相同输入的层 我们合并了具有相同输入的线性层，以减少张量和序列并行中的通信，提高矩阵操作的效率。具体而言，注意力块中的 Q、K 和 V 层被合并，而 SwiGLU 中的两个线性投影层也被合并。此优化技术对 LLM 是通用的。以下是示例代码片段：

qproj、kproj 和 vproj 被合并为 qkvproj：

pythonif not selfconfigseparateqkv and selfnumheads == selfnumkeyvalueheads and selfconfigkvsharedgroupsize == 1 qkvstates = selfqkvproj(hiddenstates) querystates keystates valuestates = qkvstatessplit(selfsplitsize dim=2)elif selfconfigqkvlinear querystates keystates valuestates = selfqkvproj(hiddenstates)else querystates = selfqproj(hiddenstates) keystates = selfkproj(hiddenstates) valuestates = selfvproj(hiddenstates)

gateproj和 upproj 被合并为 gateupproj：

pythongateproj upproj = selfgateupproj(x)split(selfsplitsize dim=2)

编译器优化 我们使用编译标志 distributionstrategy=llmtraining 来启用编译器对 LLM 训练运行的优化，允许在数据并行工作者之间分片参数、梯度和优化器状态。此外，我们还使用 modeltype=transformer，进行特定于变换器模型的优化。我们设置 Neuron 环境变量 NEURONFUSESOFTMAX=1，以启用对 Softmax 操作的自定义降级的编译器优化。最后，我们使用 NEURONRTASYNCEXECMAXINFLIGHTREQUESTS=3，通过异步执行来减少训练延迟。这使得加速器的某些运行与主机CPU重叠。

以下表格总结了我们在预训练过程中使用的所有超参数。

Trn NxD优化参数Seqlen4096Precisionbf16GBS1024学习率300E04minlr300E05权重衰减01梯度裁剪1学习率调度器cosine预热步骤2000常量步骤0AdamW (bete1 beta2)(09 095)AdamW eps100E05分布式参数节点数128TP8