如何透过 AWS Trainium 降低 LLM 训练成本 30%

由 Abhinandan Patni、James Park、Jerry Mannil、Ken Su 以及 RJ 于 2023 年 11 月 22 日发表
在 、、[中级 (200)](https://aws.amazon.com/blogs/machine-learning/category/learning- levels/intermediate-200/ "查看所有中级文章")

分享

文章重点

• 随著 AWS Trainium 的推出，Amazon Search M5 团队成功将模型训练成本降低了 30%。
• 使用 AWS 专门的加速器 Trainium，使得模型训练效率显著提升。
• 本文分享了实施过程中的最佳实践及技术选择，助力于 Amazon 内部多语言和多任务的模型开发。

几十年来，Amazon 在机器学习 (ML) 领域中不断创新，为客户带来愉悦的体验。从最早的书籍推荐、搜索到反欺诈，Amazon 在各种用例中利用 ML技术。随著行业内部加速硬件的进步，Amazon 团队得以使用神经网络和深度学习 (DL) 技术进一步优化其产品和服务。

Amazon Search 内的 M5计划负责整体的发现学习策略，并针对多语言、多地区、多实体、多任务以及文本、图像和视频等多模态进行大规模模型的建立。M5 团队提供通用嵌入和大型基础模型，支持 Amazon 内部多达数百个 ML 团队，同时保持对成本的严格控制。为了达成这一目标，M5 团队定期评估新技术，以降低成本。

像许多 ML 组织一样，M5 团队利用加速器加速 DL 训练和推理。2020 年，AWS 首次推出具有专用功能的加速器 ，M5团队迅速开始，从而节省成本并降低延迟。去年，AWS 推出了 加速器，针对开发和构建下一代 DL模型的性能与成本进行了优化。本文将讨论 M5 如何成功将模型训练成本降低 30%，并分享我们在过程中获得的一些最佳实践。

Trainium 实例

借助专用加速器的进步，Amazon 也提供了 AWS Inferentia 和 Trainium等出色的加速器。这些芯片被优化以满足推理和训练工作负载的需求。在大型基础模型的训练中，Trainium 是理想的选择，因为它们符合此类模型的特性。Trn1实例由最先进的 提供支持，并拥有大量的计算和记忆体资源。而 Trn1n 实例则可以选择更高的网络带宽 (1,600 Gbs)，在考虑到性能训练的同时，实现了成本优化。

要使用加速器，您需要一层支持它们的软件。借助 Trn 和 Inf 芯片， 搭配 PyTorch XLA 轻松解锁 Amazon专用的加速器。PyTorch XLA 可以将 PyTorch 的即时模式转换为基于图形的懒惰模式实现，进而提高模型运行效率。PyTorchNeuron（Neuron SDK 的一部分）使用少量代码即可实现在 Trainium NeuronCores 上训练模型。

模型与工作负载

M5 团队训练和部署基础模型及通用表示，协助 Amazon 内部各团队促进客户在 的体验。这些模型之一是一种文本编码器模型，配合多层感知器 (MLP)，模型的神经网络架构拥有数亿个可训练参数。在数十亿个标记的基础上进行训练后，该模型可在离线批推理环境中生成数百万个嵌入，这些嵌入将作为面向客户的 Amazon 服务的输入。

生产管道的基础设施使用 ，并采用 ，以使用启用 EFA 的多节点 trn1.32xlarge集群作为模型训练的计算资源。从功能上来看，该生产管道进行增量模型训练、已训练模型的评估和基于已训练模型的离线批推理，所有这些均使用 PyTorch作为基础的 DL 库。

目标

让客户满意是我们的首要信条。鉴于管道面向客户的特性，满足所有服务级别协议 (SLA) 而不出现回归是至关重要的。我们确定了两项关键的验收标准，以适应我们现有的 GPU 生产管道并将其转换为 Trainium：

• 模型质量 - 模型的质量直接影响客户体验。必须保持 GPU 和 Trainium 的模型质量差异在 0.1% 以下。
• 训练吞吐量 - 我们定期对模型进行迭代训练，以为客户提供最新体验。必须在预定的时间内（例如 1 周内）达成模型收敛，以符合我们的生产 SLA。

接下来的几个部分，我们将分享回顾这些标准的过程以及我们为支持 Amazon 机级生产工作负载所学到的经验。

训练脚本

在开始模型训练之前，我们需要针对 XLA 进行训练脚本的变更。由于模型规模较大，我们使用分布式数据并行 (DDP) 进行训练。DDP允许我们通过增大运行模型训练所使用的机器数量，实现更高的数据训练吞吐量，并且不需要变更任何代码。我们根据 的指导，在训练脚本中添加了 XLA专用结构。这些代码变更是相对简单的实践。以下是从这次实验中获得的一些重要技术经验，可以在很大程度上提高我们的模型吞吐量：

• xm.mark_step() 的位置 - xm.mark_step() 编译并运行懒惰收集的计算图。过于频繁地调用 mark_step 可能会导致产生大量的小图，而调用次数太少可能会生成较少但较大的图。根据应用情况，模型训练的吞吐量和实现会根据 xm.mark_step() 的位置而有所不同。我们的实现是在前向和反向传播之后和优化器步骤之后各放置一次 xm.mark_step()。
• 使用 XLA 多进程设备加载数据 - 这是一个关键步骤，容易被忽略。多进程设备加载器 torch_xla.distributed.parallel_loader.MpDeviceLoader 在每个 XLA 设备上加载训练数据，并支持预加载和与设备运行重叠数据加载以提高吞吐量。设备加载器也会调用 xm.mark_step()，因此能够为从主机加载数据到设备构建图。

Trainium 的编译

传统上，使用 GPU 的模型开发流程涉及对模型或训练脚本进行更改，然后直接在 GPU 设备上运行。使用 XLA 的加速器，如 Trainium，要求在模型训练运行之前执行额外步骤。XLA 计算图只能在编译后运行。一般来说，编译可以通过两种方式执行：提前编译 (AOT)，先追踪并编译所有图，然后运行，或直接编译 (JIT)，在遇到时追踪、编译并运行图。Neuron SDK 都提供这两种方式。通常，AOT编译是首先执行的。这样进行编译后才运行图。如果遇到新的图，Neuron 运行时会在运行之前调用 JIT 编译。为了执行 AOT 编译，Neuron SDK提供了 ，这是一种编译工具，能够从训练脚本的试运行中提取图并进行并行的 AOT 编译。

AOT编译的一个重要方面是确保在训练过程中未创建新的计算图。在模型训练过程中，由于训练批次的动态形状会导致新计算图的出现，从而引发重新编译的问题。我们发现使用静态形状和固定大小的批次可以消除训练时的编译，并在不影响模型准确性的情况下大幅提高训练吞吐量。通过强制执行这些限制，我们观察到仅需 4-5 步模型训练、一个模型验证步骤和一次模型检查点即可完成 AOT 编译时的所有图的追踪。需要注意的是，Neuron SDK不断发展，未来将支持动态形状。

此外，编译的图会存储在 的磁碟中或 (Amazon S3) 桶中。这对于生产工作负载尤其有用，因为模型架构和训练配置通常不会改变，避免了编译的开销。使用缓存的方式也十分简单，只需设置环境标志：

bash export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Neuron 编译器还提供三种 (O1、O2、O3)，以平衡编译时间与模型运行吞吐量。O1 启用计算图的核心优化并最小化编译时间，O3 则在增长编译时间的情况下提供改善的模型运行吞吐量，而 O2（默认选项）则兼顾两者。在我们的用例中，我们使用 O1 优化，观察到编译时间降低了 86%，而模型准确度指标没有变化，同时吞吐量与默认优化（O2）相比降低了约 5-7%。依据具体情况，您可以选择不同的优化等级。

总结来说，我们使用了以下标志进行编译：

bash NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

检查点兼容性

在编译顺利完成后，我们可以进行 Trainium上的模型训练。如前所述，我们的模型采取增量训练的方式，这意味著会加载之前训练的模型检查点并使用新数据继续训练。PyTorch 和 PyTorch XLA允许在不同加速器之间无缝过渡，这给我们带来了灵活性，使我们能够将之前的 GPU 模型加载并使用 Trainium机器进行训练。这对于确保我们可以在不造成生产停机或失去模型准确性的情况下，以最好的已训练模型进行初始化至关重要。

由于 GPU 模型是使用标准的 PyTorch 模型保存工具保存的，我们可以使用 PyTorch 检查点加载工具在 Trainium 设备上加载 GPU模型。

例如，在 GPU/CPU 上，您可以使用以下代码保存模型：

python torch.save(model.state_dict(), PATH)

然后您可以在 Trainium 上加载该模型：

python import torch_xla.core.xla_model as xm xla_device = xm.xla_device() model = MyModel(*args, **kwargs) model.load_state_dict(torch.load(PATH)) model.to(xla_device)

同理，可以使用以下代码在 Trainium 上保存模型：

# 自动将数据移至主设备的 CPU

xm.save(model.state_dict(), PATH) ```

并在 GPU/CPU 上加载模型：

`python model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH)) model.to(device) # 可以是任何设备`

实际上，由于我们使用 DDP 进行模型训练，模型加载方式不受先前检查点所用机器数量的影响。这让我们在 Trn1计算资源中进行横向扩展，而无需变更代码或对模型训练造成负面影响。这些基于 PyTorch 的检查点可以直接使用，甚至可以使用 torch-script 用于
AWS Inferentia2 或其他加速器的推理用例。

## 操作稳定性

生产运行工作负载需要满足多个 SLA，不容忽视。对于我们的用例，除了模型质量和训练吞吐量
SLA，还必须确保生产管道在运行中的稳定性，这意味著在模型训练、评估和推理期间要尽量减少停机和中断的情况。

跟现有的 GPU基础管道一样，我们添加了多项机制来确保管道的稳定性。在开始模型训练之前，我们会运行多个健康测试来评估机器的状况。这些测试通常包括简单的张量操作，以验证加速器设备的健康。我们观察到，在分布式训练中，验证实例之间的集体通信同样至关重要。我们使用了
[NCCOM 测试套件](https://awsdocs-neuron.readthedocs-
hosted.com/en/latest/tools/neuron-sys-tools/nccom-test.html