Amazon Search M5 如何通过使用 AWS Trainium 节省 30% 的 LLM 培训成本 | 亚马逊网络服务

几十年来，亚马逊一直在机器学习 (ML) 方面开拓创新，为客户带来愉悦的体验。 从最早开始，亚马逊就将机器学习用于各种用例，例如书籍推荐、搜索和欺诈检测。 与行业其他公司类似，加速硬件的进步使亚马逊团队能够使用神经网络和深度学习 (DL) 来追求模型架构。

亚马逊搜索中的 M5 程序拥有亚马逊的发现学习策略，并构建跨多语言、多语言环境、多实体、多任务和多模态（例如文本、图像和视频）的大规模模型。 M5 计划一直在为亚马逊数百个 ML 团队提供通用嵌入和大规模基础模型，同时保持对成本优化的严格控制。 为了实现这一目标，M5 团队定期评估新技术以降低成本。

与许多机器学习组织一样，加速器主要用于加速深度学习训练和推理。 当 AWS 通过第一个版本推出专用加速器时 AWS 推理 2020年，M5团队迅速开始 利用它们更有效地部署生产工作负载，既节省成本又减少延迟。 去年，AWS 推出了 AWS 培训 加速器，可优化开发和构建下一代深度学习模型的单位成本性能。 在这篇文章中，我们讨论 M5 如何将模型训练成本降低 30%，并分享我们在此过程中学到的一些最佳实践。

Trainium 实例

随着专用加速器的进步，Amazon 还以 AWS Inferentia 和 Trainium 的形式提供引人注目的加速器。 顾名思义，这些芯片经过优化，分别超出了推理和训练工作负载的需求。 对于达到数十亿参数大小的基础模型的大规模训练，Trainium Trn1 和 Trn1n 实例 由于其特性，它们是理想的选择。 Trn1 实例由最先进的技术提供支持 NeuronCore-v2，并拥有大量的加速器计算和内存。 还可以选择 Trn1n 实例来获得更大的网络带宽 (1,600 Gbs)，因此非常适合考虑成本优化的高性能训练。

要使用加速器，您需要一个软件层来支持它们。 使用 Trn 和 Inf 芯片， AWS 神经元开发工具包 在 PyTorch XLA 的帮助下解锁 Amazon 专用加速器。 PyTorch XLA 将 PyTorch 的 eager 模式转换为基于图的惰性模式实现。 然后使用这些图表并进一步编译以与加速器一起使用。 PyTorch Neuron（Neuron SDK 的一部分）使 PyTorch 用户能够通过几行代码在 Trainium NeuronCores 上训练他们的模型。

型号及工作量

M5 团队训练和部署基础模型和通用表示，以协助亚马逊的各个团队为 Amazon.com 顾客。 其中一种模型是文本编码器模型，后跟多层感知器 (MLP)，其具有由具有数亿个可训练参数的神经网络架构定义的显式或隐式特征交互。 该模型经过数十亿个令牌的训练，并用于在离线批量推理设置中生成数百万个嵌入。 这些嵌入是面向客户的一级亚马逊服务的输入。

生产管道的基础设施使用 AWS批处理 公平份额排队策略，使用支持 EFA 的多节点 trn1.32xlarge 集群作为模型训练的计算。 从功能上讲，生产管道执行增量模型训练、训练模型的评估以及对训练模型的离线批量推理，所有这些都使用 PyTorch 作为底层 DL 库。

理想中

让客户满意是我们的首要宗旨。 鉴于管道面向客户的性质，满足所有服务级别协议 (SLA) 且不出现倒退至关重要。 我们确定了两个关键的验收标准来调整我们现有的 GPU 生产管道并将其过渡到 Trainium：

• 模型质量 – 我们模型的质量直接影响客户体验。 我们要求 GPU 和 Trainium 之间的模型质量差异应小于 0.1%。
• 训练吞吐量 – 我们定期迭代训练我们的模型，为客户提供最新鲜的体验。 我们要求模型必须在预定的时间段（例如1周）内实现收敛，以满足我们的生产SLA。

在以下部分中，我们将分享根据此标准逆向工作的历程，以及支持亚马逊规模生产工作负载的经验教训。

训练脚本

在开始模型训练之前，我们需要对训练脚本进行更改以使其符合 XLA 要求。 考虑到模型的大小，我们使用分布式数据并行（DDP）来训练模型。 DDP 允许我们通过扩大用于运行模型训练的机器数量来提高模型训练的吞吐量，而无需更改任何代码。 我们按照中提供的说明进行操作 Neuron PyTorch MLP 训练教程 在我们的训练脚本中添加 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() 因此能够构建图表以将数据从主机加载到设备。

培训班编译

传统上，使用 GPU 的模型开发周期涉及对模型或训练脚本进行更改并直接在 GPU 设备上运行。 使用 XLA 的 Trainium 等加速器需要执行额外的步骤，然后才能在加速器上运行模型训练。 XLA计算图只有在编译后才能运行。 一般来说，有两种方法可以执行此编译：提前 (AOT)，即首先跟踪并编译所有图，然后运行它们；或即时 (JIT)，即在图出现时跟踪、编译并运行图。都遇到了。 Neuron SDK 开箱即用地提供了这两种功能。 通常，首先执行 AOT 编译。 编译后运行图表。 如果遇到新图，Neuron 运行时会在运行它们之前调用 JIT 编译。 为了执行 AOT 编译，Neuron SDK 提供了 神经元并行编译，一个编译实用程序，可从训练脚本的试运行中提取图形并执行并行 AOT 编译。

AOT 编译的一个重要方面是确保在训练过程中不会创建新的计算图。 新计算图（以及因此重新编译）的来源之一是模型训练期间训练批次的动态形状。 我们发现，使用静态形状和固定大小的批次可以消除训练时间编译，并大大提高训练吞吐量，而不会对模型精度产生任何影响。 通过对训练施加这样的限制，我们发现在 AOT 编译期间跟踪所有图只需要 4-5 个模型训练步骤、一个模型验证步骤和一次模型检查点即可。 值得注意的是，Neuron SDK 正在不断发展，将来也将支持动态形状。

此外，编译后的图表存储在 神经元持久缓存 在磁盘上或在 亚马逊简单存储服务 （亚马逊 S3）存储桶。 这对于模型架构和训练配置不变的生产工作负载特别有用。 因此，编译的开销只产生一次。 使用缓存就像设置环境标志一样简单：

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

Neuron 编译器还提供了三种 编译器级优化选项 （O1、O2、O3）以平衡编译时间和模型运行吞吐量。 O1 支持计算图上的核心优化并最大限度地减少编译时间，O3 以更高的编译时间为代价提供改进的模型运行吞吐量，O2（默认选项）是两者之间的平衡。 对于我们的用例，我们使用 O1 优化，观察到编译时间减少了 86%，而模型精度指标没有变化，同时观察到与默认优化 (O5) 相比，吞吐量减少了大约 7-2%。 根据用例，您可以选择不同级别的优化。

总而言之，我们使用以下标志进行编译：

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

检查点兼容性

编译成功完成后，我们可以继续在 Trainium 上训练我们的模型。 如前所述，我们逐步训练模型，这意味着我们加载先前训练的模型检查点并继续使用新数据进行训练。 PyTorch 和 PyTorch XLA 允许通过检查点互操作性在加速器之间无缝转换。 拥有在 GPU 和 Trainium 之间移动的灵活性，使我们能够无缝加载以前的 GPU 模型并在 Trainium 机器上进行训练。 这对于确保我们可以使用之前训练的最佳模型来初始化模型，而不会造成任何生产停机或模型准确性损失至关重要。

由于 GPU 模型是使用标准 PyTorch 模型保存实用程序保存的，因此我们能够使用 PyTorch 检查点加载实用程序在 Trainium 设备上加载 GPU 模型。

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

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

然后将模型加载回 Trainium：

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上：

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

并将模型加载回 GPU/CPU：

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

事实上，因为我们使用 DDP 进行模型训练，所以模型加载与用于训练先前检查点的机器数量无关。 这使我们能够水平扩展 Trn1 队列，而不需要更改代码或对模型训练产生不利影响。 这些基于 PyTorch 的检查点可以直接使用，甚至可以编写 torch 脚本，用于 AWS Inferentia2 或其他加速器上的推理用例。

运行稳定性

必须强调的是，在生产环境中运行工作负载需要满足多个 SLA。 对于我们的用例，除了模型质量和训练吞吐量 SLA 之外，生产管道必须稳定运行，这意味着模型训练、评估和推理期间的停机时间和中断最少。

与现有的基于 GPU 的管道一样，我们添加了许多机制来使管道运行稳定。 在开始模型训练之前，我们运行多次健全性测试来评估机器的健康状况。 这些测试通常包括简单的张量操作，以验证加速器设备的运行状况。 我们观察到，对于分布式训练，运行测试来验证实例之间的集体通信也很重要。 我们使用了 NCCOM 测试套件 通过 Neuron SDK 来实现这一点，运行各种操作，例如全聚集、全归约和归约分散。

即使遵循了我们提到的建议，我们仍然发现，无论底层加速器如何，瞬态问题在任何管道中都是不可避免的。 为了在任何培训管道中建立弹性，我们建议建立重试机制来解决这些潜在问题。 我们用 AWS Batch 自动重试 重试在模型训练期间遇到暂时故障的作业。 如果在训练快要结束时遇到故障，这些重新启动的成本可能会很高。 为了解决这个问题，我们调整了训练脚本来加载先前训练的模型检查点并从该点继续训练。 借助此功能，我们能够以最小的开销积极地重新启动失败的训练作业。

有了这些弹性机制，我们能够在 Trn98.5 上实现 1% 的工作负载成功率，与我们现有的 GPU 管道成功率相当。

成果

为了验证模型的准确性，我们从同一 GPU 检查点初始化了两个模型，并在 Trainium 上训练一个模型，在类似的 GPU 上训练另一个模型。 两个模型都使用相同的训练超参数进行训练。 用于指标计算的数据集是一个保留数据集，我们每 N 个全局步骤评估模型在此数据集上的准确性。 X 轴是全局步长，Y 轴是模型精度。 我们观察到下图中每个点的模型精度差异小于 0.1%。

此外，为了评估模型训练的成本效益，我们更喜欢比较达到模型收敛所需的挂钟时间。 我们相信，与每个代币的成本、实现的 FLOPS/美元和其他因素等指标相比，这提供了更实际的成本节省观点。 考虑到trn1.32xl的训练时间和可比较的 亚马逊弹性计算云 (Amazon EC2) 实例中，我们观察到 Trainium 的模型收敛成本降低了 30%。

结论

在评估适合深度学习工作负载的不同加速器时，需要考虑许多因素。 其中最重要的是模型质量、吞吐量、成本和可用性。 确保您的模型质量和吞吐量不会因您选择的加速器而受到影响至关重要。

由于我们与 Annapurna Neuron 团队的合作和协作，Amazon Search M5 团队通过迁移到 Trainium 节省了高达 30% 的成本。 该团队能够使用 Trainium 并实现与市场上同类加速器同等的模型质量和吞吐量。 检查点互操作性和对 XLA 支持的最小代码更改使 M5 能够针对其工作负载在多个加速器之间进行选择。 这使得 M5 团队能够利用 Trainium 的强大计算能力，构建与加速器无关的解决方案来满足 Amazon.com 客户的需求。 从运营角度来看，Trainium 已被证明能够支持亚马逊规模的一级服务。 M1 团队继续将更多工作负载转移到 Trainium，以最低的成本为 Amazon 提供最佳模型。

总之，M5 团队通过将 Trainium 添加到加速器系列中，能够执行经济高效的生产级 ML 训练。 我们鼓励您查看 Trainium 和其他 Neuron 设备（例如 AWS Inferentia），以获得专门为 ML 工作负载构建的 Amazon 芯片的优势。 从具有不同模型的众多教程之一轻松开始，例如 Llama 2，可在 Trainium 上找到.

作者简介

阿比南丹·帕特尼 是亚马逊搜索的高级软件工程师。 他专注于为可扩展的分布式深度学习训练和实时推理构建系统和工具。

詹姆斯公园 是 Amazon Web Services 的解决方案架构师。 他与 Amazon.com 合作，在 AWS 上设计、构建和部署技术解决方案，并且对人工智能和机器学习特别感兴趣。 在业余时间，他喜欢寻找新文化、新体验，并了解最新的技术趋势。 你可以在以下位置找到他 LinkedIn.

杰里·曼尼尔 是亚马逊搜索的软件工程师。 他致力于提高分布式培训基础设施的效率、稳健性和可扩展性。

苏肯 是亚马逊搜索的软件工程师。 他致力于提高培训效率和可扩展的分布式培训工作流程。 工作之余，他喜欢徒步旅行和网球。

RJ 是亚马逊的一名工程师。 他为分布式系统构建和优化系统进行训练，并致力于优化采用系统以减少 ML 推理的延迟。 工作之余，他正在探索使用生成式人工智能来构建食物食谱。

• SEO 支持的内容和 PR 分发。 今天得到放大。
• PlatoData.Network 垂直生成人工智能。 赋予自己力量。 访问这里。
• 柏拉图爱流。 Web3 智能。 知识放大。 访问这里。
• 柏拉图ESG。 碳， 清洁科技, 能源， 环境， 太阳能， 废物管理。 访问这里。
• 柏拉图健康。 生物技术和临床试验情报。 访问这里。
• Sumber: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/