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

由 Abhinandan Patni、James Park、Jerry Mannil、Ken Su 以及 RJ 于 2023 年 11 月 22 日发表在 Amazon SageMaker、最佳实践、中级 (200) 持久链接 评论 分享

文章重点

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

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

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

Trainium 实例

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

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

模型与工作负载

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

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

目标

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

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

训练脚本

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

Trainium 的编译

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

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

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

bashexport NEURONCOMPILECACHEURL=s3//BUCKET/KEY

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

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

bashNEURONCCFLAGS=target trn1 autocast all autocasttype bf16 modeltype transformer optlevel O1

检查点兼容性

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

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

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

pythontorchsave(modelstatedict() PATH)

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

pythonimport torchxlacorexlamodel as xmxladevice = xmxladevice()model = MyModel(args kwargs)modelloadstatedict(torchload(PATH))modelto(xladevice)

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

pythonimport torchxlacorexlamodel as xm

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

xmsave(modelstatedict() PATH)

并在 GPU/CPU 上加载模型：

pythonmodel = MyModel(args kwargs)modelloadstatedict(torchload(PATH))modelto(device) # 可以是任何设备

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

操作稳定性

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

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