通过使用AWS Trainium，亚马逊搜索M5使LLM培训成本节省了30%
借助AWS Trainium，亚马逊搜索M5使LLM培训成本降低30%
几十年来，亚马逊一直在机器学习（ML）方面处于领先地位，为其客户带来愉悦的体验。从最早的日子起，亚马逊就在书籍推荐、搜索和欺诈检测等各种用例中使用了ML。与该行业的其他公司类似，加速硬件的进步使得亚马逊团队能够使用神经网络和深度学习（DL）进行模型架构的研究。
亚马逊搜索中的M5项目拥有亚马逊的发现学习战略，并构建跨多语言、多地区、多实体、多任务和多模态（例如文本、图像和视频）的大规模模型。M5项目一直为亚马逊的数百个ML团队提供通用嵌入和大规模基础模型，同时严格控制成本优化。为了实现这一点，M5团队定期评估新的技术来降低成本。
与许多ML组织一样，加速器在加速DL训练和推理方面发挥了重要作用。当AWS在2020年首次发布AWS Inferentia时，M5团队迅速开始利用它们更高效地部署生产工作负载，既节省成本又减少延迟。去年，AWS推出了AWS Trainium加速器，该加速器对于开发和构建下一代DL模型的性能成本优化至关重要。在本文中，我们将讨论M5团队如何将训练模型的成本降低30％，并分享我们沿途学到的一些最佳实践。
Trainium实例
随着专用加速器的进步，亚马逊还提供了AWS Inferentia和Trainium等引人注目的加速器。正如它们的名称所暗示的，这些芯片分别优化了推理和训练工作负载的需求。对于规模庞大、具有数十亿参数的基础模型的大规模训练，Trainium Trn1和Trn1n实例是理想选择，因为它们的特性。Trn1实例由领先的NeuronCore-v2驱动，具有大量的加速器计算和内存。Trn1n实例还可以选择更大的网络带宽（1,600 Gb/s），因此非常适合以成本优化为目标的高性能训练。
要使用加速器，您需要一个软件层来支持它们。使用Trn和Inf芯片，AWS Neuron SDK借助PyTorch XLA解锁了亚马逊特定的加速器。PyTorch XLA将PyTorch的即时模式转换为基于图形的懒惰模式实现。然后使用这些图形并进一步编译以与加速器一起使用。PyTorch Neuron（Neuron SDK的一部分）使PyTorch用户可以使用几行代码在Trainium NeuronCores上训练他们的模型。
模型和工作负载
M5团队训练并部署了基础模型和通用表示，以帮助亚马逊的各个团队为Amazon.com的客户带来愉悦体验。其中一个模型是一个文本编码器模型，后面跟着一个由神经网络架构定义的具有数亿可训练参数的多层感知机（MLP），其中包括显式或隐式特征交互。该模型在数十亿个标记上进行训练，并用于在离线批处理推理环境中生成数百万个嵌入。这些嵌入是面向客户的一级亚马逊服务的输入。
用于生产流水线的基础设施使用AWS Batch和公平共享排队策略，使用支持EFA的多节点trn1.32xlarge集群作为模型训练的计算资源。从功能上讲，生产流水线使用PyTorch作为底层DL库，执行增量模型训练、已训练模型的评估和离线批处理推理。
目标
让我们的客户满意是我们的首要准则。考虑到与客户直接相关的流程性质，务必满足所有服务级别协议（SLAs）而不产生任何倒退。我们确定了两个关键的验收准则，以适应我们现有的 GPU 生产流程并过渡到 Trainium：
- 模型质量 – 我们模型的质量直接影响客户体验。我们要求 GPU 和 Trainium 之间的模型质量差异应小于 0.1%。
- 训练吞吐量 – 我们周期性地迭代训练模型，以向客户提供最新的体验。我们要求模型的收敛必须在预定的时间段内实现，以满足我们的生产 SLAs。
在接下来的章节中，我们将分享从这些准则开始反向工作的过程和我们的经验，以支持亚马逊规模的生产工作量。
训练脚本
在开始模型训练之前，我们需要对训练脚本进行更改，使其符合 XLA 标准。考虑到模型的大小，我们使用分布式数据并行（DDP）来训练模型。DDP 允许我们增加用于运行模型训练的机器数量，从而增加模型训练的吞吐量，而无需进行任何代码更改。我们按照 Neuron PyTorch MLP training tutorial 提供的说明，在我们的训练脚本中添加了特定于 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 提供了 neuron_parallel_compile，这是一个编译工具，可以从训练脚本的试运行中提取图，并进行并行的 AOT 编译。
AOT 编译的一个重要方面是确保在训练过程中不会创建任何新的计算图。训练过程中的训练批次的动态形状是产生新计算图（因此需要重新编译）的一个源头。我们发现，使用静态形状和固定大小的批次可以消除训练时间的编译，并极大地提高训练吞吐量，而不会对模型准确性产生任何影响。通过对训练施加这样的约束，我们观察到只需要 4-5 步模型训练、1 步模型验证和一次模型检查点，就可以在 AOT 编译期间跟踪所有的图。值得注意的是，Neuron SDK 在不断发展，并将来支持动态形状。
此外，编译的图表存储在磁盘上的Neuron持久缓存或Amazon Simple Storage Service（Amazon S3）桶中。这在模型架构和训练配置不会发生变化的生产工作负载中特别有用。因此，只会产生一次编译的开销。使用缓存很简单，只需设置一个环境标志：
export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"Neuron编译器还提供了三个编译器级别的优化选项（O1，O2，O3），以平衡编译时间和模型运行吞吐量。O1在计算图上启用核心优化，最小化编译时间；O3提供了更高的模型运行吞吐量，但编译时间也更长；而O2（默认选项）则在两者之间取得平衡。对于我们的用例，我们使用了O1优化，并观察到编译时间减少了86%，而模型准确性指标没有改变，与默认优化（O2）相比，吞吐量减少了约5-7%。根据使用情况，可以选择不同级别的优化。
总结起来，我们在编译时使用了以下标志：
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 xmxla_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#自动将数据移动到CPU以用于主设备xm.save(model.state_dict(), PATH) 然后在GPU/CPU上加载模型：
model = MyModel(*args, **kwargs)model.load_state_dict(torch.load(PATH))model.to(device) #可以是任何设备事实上，由于我们使用DDP进行模型训练，模型加载不依赖于训练先前检查点时使用的机器数量。这使我们能够在Trn1集群上水平扩展，而无需更改代码或对模型训练产生不利影响。这些基于PyTorch的检查点可以直接用于AWS Inferentia2或其他加速器上的推断用例，甚至可以进行torch脚本编写。
运营稳定性
无法过分强调在生产中运行工作负载需要满足多个SLA的重要性。对于我们的用例，除了模型质量和训练吞吐量SLA之外，使生产流水线运行稳定是至关重要的，即在模型训练、评估和推断过程中减少停机时间和干扰。
与现有的基于GPU的流水线类似，我们添加了多种机制来使流水线运行稳定。在开始模型训练之前，我们运行多个健康测试来评估机器的健康状况。这些测试通常包括简单的张量运算，以验证加速器设备的健康状况。我们观察到，对于分布式训练，重要的是运行测试来验证实例之间的集体通信。我们使用Neuron SDK的NCCOM测试套件来实现这一点，运行各种操作，例如全部收集、全部归约和减少分散。
即使在我们提到的建议之后，我们发现无论使用的加速器如何，在任何流水线中都无法避免暂时性的问题。为了在任何训练流程中建立弹性，我们建议建立重试机制以解决这些潜在问题。我们使用AWS 批量自动重试来重试在模型训练期间遇到暂时性故障的作业。如果在训练结束前遇到故障，这些重启可能很昂贵。为了解决这个问题，我们修改了我们的训练脚本，加载先前训练过的模型检查点，并从那个点继续训练。有了这个功能，我们能够以最小的开销积极地重新启动失败的训练作业。
有了这些弹性机制，我们在Trn1上的工作负载达到了98.5％的成功率，与我们现有的GPU流水线成功率相当。
结果
为了验证我们模型的准确性，我们从同一个GPU检查点初始化了两个模型，并分别在 Trainium 和相似的GPU上训练。两个模型使用相同的训练超参数进行训练。用于计算指标的数据集是一个保留数据集，并且我们每隔N个全局步骤评估模型的准确性。图中X轴是全局步骤，Y轴是模型准确性。我们观察到在下图的每个点上，模型准确性之间的差异小于0.1％。
此外，为了评估模型训练的成本效益，我们更倾向于比较达到模型收敛所需的挂钟时间。相比于单位代币成本、达到的每秒浮点运算次数/美元和其他因素，我们认为这提供了更实用的成本节约视角。考虑到 Trn1.32xl 的训练时间和相似的亚马逊弹性计算云（Amazon EC2）实例，我们观察到 Trainium 的模型收敛成本更低，可节省高达30％的成本。
结论
在评估不同加速器用于深度学习工作负载时，有许多因素需要考虑。其中最重要的因素包括模型质量、吞吐量、成本和可用性。确保在选择加速器时，不会牺牲模型质量和吞吐量是至关重要的。
多亏了我们与 Annapurna Neuron 团队的合作与协作，亚马逊搜索 M5 团队通过转向 Trainium，成功节省了高达30％的成本。团队能够使用 Trainium 并在市场上与可比加速器达到模型质量和吞吐量的平衡。检查点的互操作性和对 XLA 的最小代码更改使 M5 能够在其工作负载中选择多个加速器。这使得 M5 团队能够充分利用 Trainium 的强大计算能力，并构建不受加速器限制的解决方案，为亚马逊网顾客提供愉快的体验。从运营的角度来看，Trainium 已被证明能够支持亚马逊规模的一流服务。M5 团队将继续将更多工作负载迁移到 Trainium，以以最低的成本为亚马逊提供最好的模型。
总之，通过将 Trainium 添加到加速器的车队中，M5 团队能够进行具有成本效益且符合生产标准的机器学习训练。我们鼓励您去了解 Trainium 和其他 Neuron 设备，例如 AWS Inferentia，以利用为机器学习工作负载量身定制的亚马逊硅片的好处。通过参考包含不同模型（例如Llama 2，在 Trainium 上可用）的许多教程之一，您可以轻松入门。
