AWS Inferentia2 在 AWS Inferentia1 的基础上进行了升级，其吞吐量提高了4倍，延迟降低了10倍

机器学习（ML）模型的规模——大型语言模型（LLMs）和基础模型（FMs）——年复一年地快速增长，这些模型需要更快、更强大的加速器，尤其是用于生成式人工智能。AWS Inferentia2从根本上设计，以提供更高的性能，同时降低LLMs和生成式AI推理的成本。

在本文中，我们展示了AWS Inferentia的第二代如何建立在AWS Inferentia1引入的功能基础上，满足部署和运行LLMs和FMs的独特需求。

AWS Inferentia的第一代是一种专门构建的加速器，于2019年推出，专为加速深度学习推理而优化。AWS Inferentia帮助ML用户降低推理成本，并提高了他们的预测吞吐量和延迟。使用AWS Inferentia1，客户看到的吞吐量高达2.3倍，每次推理的成本比可比较的推理优化Amazon Elastic Compute Cloud（Amazon EC2）实例低70%。

AWS Inferentia2是新的Amazon EC2 Inf2实例中的特色，也受到Amazon SageMaker的支持，它针对大规模生成式AI推理进行了优化，是AWS的第一个针对分布式推理优化的推理实例，具有高速、低延迟的加速器之间或加速器与网络适配器之间的连接。

现在，您可以在单个Inf2实例上高效地部署一个包含1750亿个参数的模型，而无需昂贵的训练实例。到目前为止，有大型模型的客户只能使用为训练而构建的实例，但这是浪费资源的——因为它们更昂贵、耗费更多能源，而它们的工作负载没有利用所有可用的资源（如更快的网络和存储）。使用AWS Inferentia2，您可以实现比AWS Inferentia1高4倍的吞吐量和高达10倍的低延迟。此外，AWS Inferentia的第二代增加了对更多数据类型、自定义运算符、动态张量等的增强支持。

AWS Inferentia2的内存容量增加了4倍，内存带宽比AWS Inferentia1高16.4倍，并原生支持将大型模型分片到多个加速器上。加速器使用NeuronLink和Neuron Collective Communication来最大化它们之间或加速器与网络适配器之间数据传输的速度。AWS Inferentia2更适合需要跨多个加速器分片的较大模型，尽管AWS Inferentia1对于较小的模型仍然是一个很好的选择，因为它提供了比其他替代方案更好的价格性能。

体系结构演进

为了比较AWS Inferentia的两个世代，让我们回顾一下AWS Inferentia1的架构。每个芯片有四个NeuronCores v1，如下图所示。

每个芯片的规格：

• 计算 – 四个核心总共提供128 INT8 TOPS和64FP16/BF16 TFLOPS
• 内存 – 8 GB的DRAM（50 GB/sec的带宽），由所有四个核心共享
• NeuronLink – 用于在两个或多个核心之间分片模型的链接

现在让我们看看AWS Inferentia2的组织方式。每个AWS Inferentia2芯片都有基于NeuronCore-v2架构的两个升级核心。与AWS Inferentia1一样，您可以在每个NeuronCore上运行不同的模型，或者将多个核心组合起来分片大型模型。

每个芯片的规格：

• 计算 – 两个核心总共提供380 INT8 TOPS，190 FP16/BF16/cFP8/TF32 TFLOPS和47.5 FP32 TFLOPS
• 内存 – 32 GB的HBM，由两个核心共享
• NeuronLink – 用于在两个或多个芯片之间分片模型的链接（每个设备384 GB/sec）

NeuronCore-v2 采用模块化设计，具有四个独立引擎：

• ScalarEngine（比 v1 快 3 倍）– 对浮点数进行操作 —— 1600（BF16/FP16）FLOPS
• VectorEngine（比 v1 快 10 倍）– 对向量进行操作，使用单个操作进行归一化、池化等计算。
• TensorEngine（比 v1 快 6 倍）– 执行张量计算，如 Conv、Reshape、Transpose 等。
• GPSIMD-Engine – 具有八个完全可编程的 512 位宽通用处理器，可使用标准的 PyTorch 自定义 C++ 运算符 API 创建自定义运算符。这是 NeuronCore-v2 中引入的新功能。

AWS Inferentia2 NeuronCore-v2 更快，更优化。此外，它能够加速不同类型和大小的模型，从 ResNet 50 等简单模型到拥有数十亿个参数的大型语言模型或基础模型，例如 GPT-3（1750 亿个参数）。与 AWS Inferentia1 相比，AWS Inferentia2 的内部存储器更大、更快，如下表所示。

芯片	神经元核心	内存类型	内存大小	内存带宽
AWS Inferentia	x4（v1）	DDR4	8GB	50GB/S
AWS Inferentia 2	x2（v2）	HBM	32GB	820GB/S

AWS Inferentia2 中找到的存储器属于高带宽存储器（HBM）类型。每个 AWS Inferentia2 芯片具有 32 GB 存储器，可以与其他芯片组合使用 NeuronLink（设备到设备的互连）来分发非常大的模型。例如，inf2.48xlarge 具有 12 个 AWS Inferentia2 加速器，共计 384 GB 的加速存储器。AWS Inferentia2 存储器的速度比 AWS Inferentia1 快 16.4 倍，如上表所示。

其他特性

AWS Inferentia2 提供以下附加功能：

• 硬件支持– cFP8（新的、可配置的 FP8）、FP16、BF16、TF32、FP32、INT8、INT16 和 INT32。有关更多信息，请参阅数据类型。
• Lazy Tensor 推断– 我们稍后在本文中讨论 Lazy Tensor 推断。
• 自定义运算符– 开发人员可以使用标准的 PyTorch 自定义运算符编程接口使用自定义 C++ 运算符功能。自定义运算符由 Tensor Factory 函数中可用的低级原语组成，并由 GPSIMD-Engine 加速。
• 控制流程（即将推出）– 这是用于模型内部原生编程语言控制流程，最终对从一层到另一层的数据进行预处理和后处理。
• 动态形状（即将推出）– 当您的模型动态更改任何内部层的输出形状时，这将非常有用。例如：一个基于输入数据减小输出张量大小或形状的过滤器。

在 AWS Inferentia1 和 AWS Inferentia2 上加速模型

AWS Neuron SDK 用于编译和运行您的模型。它与 PyTorch 和 TensorFlow 本地集成。这样，您就不需要运行额外的工具。使用您在这些 ML 框架之一中编写的原始代码，通过几行代码更改，您就可以使用 AWS Inferentia。

让我们看看如何使用PyTorch在AWS Inferentia1和AWS Inferentia2上编译和运行模型。

从torchvision加载预训练模型（ResNet 50）

加载预训练模型并运行一次以预热它：

import torch
import torchvision

model = torchvision.models.resnet50(weights='IMAGENET1K_V1').eval().cpu()
x = torch.rand(1,3,224,224).float().cpu() # 虚拟输入
y = model(x) # 预热模型

在Inferentia1上追踪和部署加速模型

要将模型追踪到AWS Inferentia，请导入torch_neuron并调用追踪函数。请记住，模型需要是PyTorch Jit可追踪的才能正常工作。

在追踪过程结束时，将模型保存为普通的PyTorch模型。编译模型一次并多次加载它。 Neuron SDK运行时已经集成到PyTorch中，并负责自动将运算符发送到AWS Inferentia1芯片以加速您的模型。

在您的推理代码中，您始终需要导入torch_neuron以激活集成的运行时。

您可以传递其他参数给编译器，以自定义它优化模型的方式或启用特殊功能，例如神经元管道核心。将模型分片到多个核心以增加吞吐量。

import torch_neuron

# 使用AWS NeuronSDK追踪模型
neuron_model = torch_neuron.trace(model,x) # 将模型追踪到Inferentia
# 保存以备将来使用
neuron_model.save('neuron_resnet50.pt')

# 下一次您不需要再次追踪模型
# 只需加载它，AWS NeuronSDK将自动将其发送到Inferentia
neuron_model = torch.jit.load('neuron_resnet50.pt')

# 在Inferentia上加速推理
y = neuron_model(x)

在Inferentia2上追踪和部署加速模型

对于AWS Inferentia2，该过程类似。唯一的区别是您导入的软件包以x结尾：torch_neuronx。 Neuron SDK会自动为您处理模型的编译和运行。您还可以传递其他参数给编译器，以微调操作或激活特定功能。

import torch_neuronx

# 使用NeuronSDK追踪模型
neuron_model = torch_neuronx.trace(model,x) # 将模型追踪到Inferentia
# 保存以备将来使用
neuron_model.save('neuron_resnet50.pt')

# 下一次您不需要再次追踪模型
# 只需加载它，NeuronSDK将自动将其发送到Inferentia
neuron_model = torch.jit.load('neuron_resnet50.pt')

# 在Inferentia上加速推理
y = neuron_model(x)

AWS Inferentia2还提供了第二种运行模型的方法，称为Lazy Tensor推理。在此模式下，您不需要事先追踪或编译模型；相反，编译器每次运行代码时都会动态编译。鉴于追踪模式具有许多优点，不建议在生产中使用Lazy Tensor推理。但是，如果您仍在开发模型并需要更快地测试它，则Lazy Tensor推理可以是一个很好的选择。以下是如何使用Lazy Tensor编译和运行模型的方法：

import torch
import torchvision
import torch_neuronx
import torch_xla.core.xla_model as xm

device = xm.xla_device() # 创建XLA设备
model = torchvision.models.resnet50(weights='IMAGENET1K_V1').eval().cpu()
model.to(device)

x = torch.rand((1,3,224,224), device=device) # 虚拟输入
with torch.no_grad():
  y = model(x)
  xm.mark_step() # 编译发生在这里

现在您熟悉了AWS Inferentia2，下一步是开始使用PyTorch或Tensorflow，并学习如何设置开发环境并运行教程和示例。此外，请查看AWS Neuron Samples GitHub存储库，在那里您可以找到多个有关如何准备模型在Inf2、Inf1和Trn1上运行的示例。

AWS Inferentia1和AWS Inferentia2之间功能比较摘要

AWS Inferentia2编译器基于XLA，AWS是OpenXLA倡议的一部分。这是与AWS Inferentia1最大的区别，因为PyTorch、TensorFlow和JAX都有本地的XLA集成。由于它优化了图形以在单个内核启动中计算结果，XLA带来了许多性能改进。它将连续的张量操作融合在一起，并输出加速AWS Inferentia2模型运行的最佳机器代码。AWS Inferentia2的其他部分也得到了改进，同时保持用户体验尽可能简单，同时跟踪和运行模型。下表显示了编译器和运行时的两个版本中可用的功能。

功能	torch-neuron	torch-neuronx
Tensorboard	是	是
支持的实例	Inf1	Inf2＆Trn1
推理支持	是	是
训练支持	否	是
架构	NeuronCore-v1	NeuronCore-v2
跟踪API	torch_neuron.trace()	torch_neuronx.trace()
分布式推理	NeuronCore Pipeline	Collective Communications
IR	GraphDef	HLO
编译器	neuron-cc	neuronx-cc
监控	neuron-monitor / monitor-top	neuron-monitor / monitor-top

有关torch-neuron（Inf1）和torch-neuronx（Inf2＆Trn1）之间的更详细比较，请参阅推理的torch-neuron（Inf1）与torch-neuronx（Inf2＆Trn1）比较。

模型服务

跟踪要部署到Inf2的模型后，您有许多部署选项。您可以以不同的方式运行实时预测或批处理预测。Inf2可用是因为EC2实例本地集成了其他使用Deep Learning Containers（DLCs）的AWS服务，例如Amazon Elastic Container Service（Amazon ECS）、Amazon Elastic Kubernetes Service（Amazon EKS）和SageMaker。

AWS Inferentia2与最流行的部署技术兼容。以下是使用AWS Inferentia2部署模型的一些选项：

• SageMaker-全面托管的服务，用于准备数据和构建、训练和部署ML模型
• TorchServe-PyTorch集成式部署机制
• TensorFlow Serving-TensorFlow集成式部署机制
• Deep Java Library-用于模型部署和训练的开源Java机制
• Triton-NVIDIA开源的模型部署服务

基准测试

下表突出 AWS Inferentia2 相对于 AWS Inferentia1 带来的改进。具体而言，我们测量延迟（模型使用每个加速器进行预测的速度）、吞吐量（每秒推理次数）和每次推理的成本（以美元计算的每次推理费用）。延迟毫秒数和美元成本越低，吞吐量越高，则越好。

在此过程中使用了两个模型——大型语言模型：ELECTRA 大型判别器和 BERT 大写未编码。PyTorch（1.13.1）和 Hugging Face transformers（v4.7.0），在此实验中使用的主要库，运行在 Python 3.8 上。在为批处理大小 = 1 和 10 编译模型之后（使用上一节中的代码作为参考），每个模型都被预热（调用一次以初始化上下文），然后连续调用 10 次。以下表格显示了在此简单基准测试中收集的平均数字。

• Electra 大型判别器（334,092,288 个参数 ~593 MB）
• Bert 大写未编码（335,143,938 个参数 ~580 MB）
• OPT-66B（660亿个参数 ~124 GB）

模型名称	批处理大小	句子长度	延迟（毫秒）	Inf2 相对于 Inf1 改进（x 倍）	吞吐量（每秒推理次数）	每次推理的成本（EC2 美国东部1区）**
Inf1	Inf2	Inf1	Inf2	Inf1	Inf2
ElectraLargeDiscriminator	1	256	35.7	8.31	4.30	28.01	120.34	$0.0000023	$0.0000018
ElectraLargeDiscriminator	10	256	343.7	72.9	4.71	2.91	13.72	$0.0000022	$0.0000015
BertLargeUncased	1	128	28.2	3.1	9.10	35.46	322.58	$0.0000018	$0.0000007
BertLargeUncased	10	128	121.1	23.6	5.13	8.26	42.37	$0.0000008	$0.0000005

* 在此基准测试中使用了 32 个 AMD Epyc 7313 CPU 的 c6a.8xlarge。

** 2021 年 4 月 20 日，us-east-1 的 EC2 公共定价为：inf2.xlarge：0.7582 美元/小时；inf1.xlarge：0.228 美元/小时。每次推理的成本考虑批处理中每个元素的成本。（每次推理的成本等于模型调用/批处理大小的总成本。）

有关培训和推理性能的其他信息，请参阅 Trn1/Trn1n 性能。

结论

AWS Inferentia2 是一项强大的技术，旨在改善深度学习模型推理的性能并降低成本。比 AWS Inferentia1 更具性能，提供高达 4 倍的吞吐量、高达 10 倍的较低延迟和高达其他可比较推理优化 EC2 实例 50% 的性能/瓦特。最终，您可以支付更少的费用，拥有更快的应用程序，并达到可持续性目标。

将推理代码迁移到 AWS Inferentia2 很简单明了，它还支持更广泛的模型，包括用于生成 AI 的大型语言模型和基础模型。

您可以通过遵循 AWS Neuron SDK 文档来设置开发环境并开始您的加速深度学习项目。为了帮助您入门，Hugging Face 已将 Neuron 支持添加到他们的 Optimum 库中，该库优化模型以进行更快的培训和推理，并且他们有许多已准备好在 Inf2 上运行的示例任务。此外，请查看我们的“使用大型模型推理容器在 AWS Inferentia2 上部署大型语言模型”，了解如何使用模型推理容器将 LLM 部署到 AWS Inferentia2。有关其他示例，请参阅 AWS Neuron Samples GitHub 存储库。