Amazon Search 如何使用 AWS Inferentia 将 ML 推理成本降低 85%

亚马逊的产品搜索引擎索引了数十亿种产品，为全球数亿客户提供服务，是世界上使用最频繁的服务之一。 Amazon Search 团队开发了机器学习 (ML) 技术，为 Amazon.com 搜索引擎，帮助客户轻松搜索。 提供出色的客户体验并以客户所需的大规模运营 Amazon.com 搜索引擎，该团队一直在寻找方法来构建具有实时延迟和吞吐量要求的更具成本效益的系统。 该团队不断探索针对深度学习优化的硬件和编译器，以加速模型训练和推理，同时全面降低运营成本。

在这篇文章中，我们描述了亚马逊搜索如何使用 AWS 推理，由 AWS 构建的高性能加速器，用于加速深度学习推理工作负载。 该团队在基于 AWS Inferentia 的基于 Transformer 的 NLP 模型上运行低延迟 ML 推理 亚马逊弹性计算云 (Amazon EC2) Inf1 实例，可节省高达 85% 的基础设施成本，同时保持强大的吞吐量和延迟性能。

用于重复和查询意图预测的深度学习

搜索亚马逊商城是一个多任务、多模式的问题，需要处理多个输入，例如 ASIN（亚马逊标准识别号，一个唯一标识产品的 10 位字母数字编号）、产品图片、文本描述和查询。 为了创建定制的用户体验，来自许多模型的预测用于搜索的不同方面。 这是一个挑战，因为搜索系统有数千个模型，在峰值负载下每秒处理数万个事务 (TPS)。 我们专注于这种体验的两个组成部分：

• 客户感知的重复预测 – 要显示与用户查询匹配的最相关产品列表，识别客户难以区分的产品非常重要
• 查询意图预测 – 为了调整搜索页面和产品布局以更好地满足客户的需求，预测用户查询的意图和类型（例如，与媒体相关的查询、帮助查询和其他查询类型）很重要

这两个预测都是使用 Transformer 模型架构进行的，即基于 BERT 的模型。 事实上，两者都共享相同的基于 BERT 的模型作为基础，并且每个都在这个主干之上堆叠了一个分类/回归头。

重复预测将一对评估产品的各种文本特征作为输入（例如产品类型、标题、描述等），并针对大型数据集定期计算。 该模型以多任务方式进行端到端训练。 Amazon SageMaker 处理作业 用于定期运行这些批处理工作负载以自动启动它们，并且只需为使用的处理时间付费。 对于这个批处理工作负载用例，推理吞吐量的要求是 8,800 总 TPS。

意图预测将用户的文本查询作为输入，并且需要实时动态地为日常流量提供服务并增强亚马逊市场上的用户体验。 该模型是针对多类分类目标进行训练的。 然后将该模型部署在 亚马逊弹性容器服务 (Amazon ECS)，它支持快速自动扩展和 易于部署定义和管理. 由于这是一个实时用例，它要求 P99 延迟低于 10 毫秒，以确保令人愉悦的用户体验。

AWS Inferentia 和 AWS Neuron 开发工具包

EC2 Inf1 实例 由 AWS Inferentia 提供支持，这是 AWS 为加速深度学习推理工作负载而构建的第一个 ML 加速器。 与基于 GPU 的 EC1 实例相比，Inf2.3 实例可提供高达 70 倍的吞吐量和高达 2% 的每次推理成本。 您可以使用您选择的框架（PyTorch、TensorFlow、MXNet）继续训练您的模型，然后轻松地将它们部署在 AWS Inferentia 上，以从内置的性能优化中受益。 您可以使用 Inf1 实例部署各种模型类型，包括图像识别、对象检测、自然语言处理 (NLP) 和现代推荐模型。

AWS 神经元 是一个软件开发工具包 (SDK)，由编译器、运行时和分析工具组成，可优化 EC2 Inf1 实例的 ML 推理性能。 Neuron 与流行的 ML 框架（如 TensorFlow 和 PyTorch）原生集成。 因此，您可以使用您选择的框架提供的相同熟悉的 API 在 AWS Inferentia 上部署深度学习模型，并从云中性能的提升和最低的每次推理成本中受益。

自推出以来，Neuron SDK 不断增加其支持的模型的广度，同时不断提高性能并降低推理成本。 这包括 NLP 模型（BERT）、图像分类模型（ResNet、VGG）和对象检测模型（OpenPose 和 SSD）。

在 Inf1 实例上部署以实现低延迟、高吞吐量和成本节约

Amazon Search 团队希望在满足重复预测的高吞吐量要求和查询意图预测的低延迟要求的同时节省成本。 他们选择部署在基于 AWS Inferentia 的 Inf1 实例上，不仅满足了高性能要求，还节省了高达 85% 的推理成本。

客户感知的重复预测

在使用 Inf1 之前，一个专用的 亚马逊电子病历 集群使用基于 CPU 的实例运行。 在不依赖硬件加速的情况下，需要大量实例才能满足每秒 8,800 个总事务的高吞吐量要求。 该团队切换到 inf1.6xlarge 实例，每个实例具有 4 个 AWS Inferentia 加速器和 16 个 NeuronCore（每个 AWS Inferentia 芯片 4 个内核）。 他们跟踪了单个 NeuronCore 的基于 Transformer 的模型，并为每个 NeuronCore 加载了一种模式，以最大限度地提高吞吐量。 通过利用 16 个可用的 NeuronCore，他们将推理成本降低了 85%（基于当前公开的 Amazon EC2 按需定价）。

查询意图预测

鉴于 P99 延迟要求为 10 毫秒或更短，该团队将模型加载到 inf1.6xlarge 实例上的每个可用 NeuronCore。 您可以使用 PyTorch Neuron 轻松做到这一点 火炬.神经元.DataParallel API。 在 Inf1 部署中，模型延迟为 3 毫秒，端到端延迟约为 10 毫秒，峰值负载时的最大吞吐量达到 16,000 TPS。

开始使用示例编译和部署代码

以下是一些示例代码，可帮助您开始使用 Inf1 实例并实现与 Amazon Search 团队一样的性能和成本优势。 我们展示了如何使用 PyTorch 模型编译和执行推理，使用 PyTorch神经元.

首先，模型编译为 torch.neuron.trace():

m = torch.jit.load(f="./cpu_model.pt", map_location=torch.device('cpu'))
m.eval()
model_neuron = torch.neuron.trace(
    m,
    inputs,
    compiler_workdir="work_" + str(cores) + "_" + str(batch_size),
    compiler_args=[
        '--fp32-cast=all', '--neuroncore-pipeline-cores=' + str(cores)
    ])
model_neuron.save("m5_batch" + str(batch_size) + "_cores" + str(cores) +
                  "_with_extra_op_and_fp32cast.pt")

有关可能参数的完整列表 trace 方法，参考 PyTorch-Neuron 跟踪 Python API。 如你看到的， 编译器参数 可以传递给 torch.neuron 直接API。 所有 FP32 运算符都强制转换为 BF16 with --fp32-cast=all，在保持动态范围的同时提供最高性能。 更多铸造选项可让您控制性能以进行模型精度权衡。 用于两个用例的模型都是为单个 NeuronCore 编译的（没有 流水线).

然后我们将模型加载到 Inferentia 上 torch.jit.load，并将其用于预测。 这 神经元运行时 自动将模型加载到 NeuronCores。

cm_cpd_preprocessing_jit = torch.jit.load(f=CM_CPD_PROC,
                                          map_location=torch.device('cpu'))
cm_cpd_preprocessing_jit.eval()
m5_model = torch.jit.load(f=CM_CPD_M5)
m5_model.eval()

input = get_input()
with torch.no_grad():
    batch_cm_cpd = cm_cpd_preprocessing_jit(input)
    input_ids, attention_mask, position_ids, valid_length, token_type_ids = (
        batch_cm_cpd['input_ids'].type(torch.IntTensor),
        batch_cm_cpd['attention_mask'].type(torch.HalfTensor),
        batch_cm_cpd['position_ids'].type(torch.IntTensor),
        batch_cm_cpd['valid_length'].type(torch.IntTensor),
        batch_cm_cpd['token_type_ids'].type(torch.IntTensor))
    model_res = m5_model(input_ids, attention_mask, position_ids, valid_length,
                         token_type_ids)

结论

在高流量和苛刻的性能要求下，Amazon Search 团队能够使用基于 AWS Inferentia 的 Inf85 实例将推理成本降低 1%。 AWS Inferentia 和 Neuron SDK 为团队提供了与训练分开优化部署过程的灵活性，并通过全面的工具和熟悉的框架 API 提出了一个浅薄的学习曲线。

通过开始使用本文中提供的示例代码，您可以解锁性能和成本优势。 另外，检查端到端 教程 在 Inferentia 上运行 ML 模型 PyTorch 和 TensorFlow.

关于作者

若昂·莫拉 是 Amazon Web Services 的 AI/ML 专家解决方案架构师。 他主要专注于 NLP 用例并帮助客户优化深度学习模型训练和部署。 他还是 ML 专用硬件和低代码 ML 解决方案的积极支持者。

张伟奇 是 Search M5 的软件工程经理，负责为 Amazon 机器学习应用程序生产大型模型。 他的兴趣包括信息检索和机器学习基础设施。

杰森·卡尔森 是一名软件工程师，负责开发机器学习管道，以帮助减少由于客户感知的重复而导致的被盗搜索印象数。 他主要与 Apache Spark、AWS 和 PyTorch 合作，帮助为 ML 模型部署和提供/处理数据。 在空闲时间，他喜欢阅读和跑步。

席少辉 是 Search Query Understanding Infra 团队的 SDE。 他领导了构建具有低延迟和高可用性的大规模深度学习在线推理服务的工作。 工作之余，他喜欢滑雪和探索美食。

张卓奇 是 Search Query Understanding Infra 团队的软件开发工程师。 他致力于构建模型服务框架，以改善深度学习在线推理服务的延迟和吞吐量。 工作之余，他喜欢打篮球、滑雪板和开车。

孙浩伟 是 Search Query Understanding Infra 团队的软件工程师。 她致力于设计支持深度学习在线推理服务的 API 和基础设施。 她的兴趣包括服务 API 设计、基础设施设置和维护。 工作之余，她喜欢跑步、远足和旅行。

贾斯普雷特·辛格 是 M5 团队的一名应用科学家，他在那里研究大型基础模型以改善客户购物体验。 他的研究兴趣包括多任务学习、信息检索和表示学习。

斯鲁蒂·科帕卡 是 AWS 的高级产品营销经理。 她帮助客户探索、评估和采用 EC2 加速计算基础架构来满足他们的机器学习需求。