通过 Optimum Neuron 加速 Llama 2 和 SDXL 在 Inf2 和 Trn1 上部署推理

背景

Hugging Face 生态体系中的 Transformers 库为 state-of-the-art 预训练模型提供了非常方便的 API 和工具集，帮助用户轻松下载模型并完成训练，是工程师手边的必备神器。为了在更多加速硬件上（如亚马逊云科技的 Trainium 和 Inferentia）训练和运行 Transformers 模型，提升整体的效率，Hugging Face 推出了开源项目 Optimum（GitHub 链接）对其进行扩展，它提供了一套性能优化工具，可以最大效率地在目标硬件上训练和运行模型。

大语言模型 LLM 以及生成式 AI 全力推动了 AI 的进一步发展，同时也有越来越多的加速器供应商提供性能优化的专业硬件出现。因此，Optimum 使得开发人员能够在享受 Transformers 易用性的同时，更加有效地使用这些平台。Optimum 英语的含义是“最适宜的、最佳效果”，所以可以看出这个项目的目标是将 Transformers 在对应的加速硬件上获得性能的提升，以趋近于其理想状态。

Optimum 的功能包括但不限于以下几点：

为各种硬件平台提供优化工具
自动调整模型超参数以实现最佳性能
支持模型压缩和量化，以减少模型大小和内存占用
提供工具来导出和运行优化后的模型

今天介绍的重点是这个系列中的 Optimum Neuron 项目，它是 Transformers 和亚马逊云科技自研的 AI 加速器 Trainium 和 Inferentia 之间的桥梁。Optimum Neuron 也延续了 Optimum 整体设计目标，为 Transformers 用户提供了一组帮助工具，可以让用户轻松地在 Trainium 和 Inferentia 上实现不同模型的快速训练和推理，而无需进行复杂的转换。

Neuron SDK

为了更好地与底层硬件打交道，对 Trainium 和 Inferentia 进行更高效的管理，亚马逊云科技推出了 Neuron SDK（上面 Optimum Neuron 名字的由来）。Neuron 支持用户在其端到端机器学习开发生命周期中构建新模型、训练和优化，然后将它们部署到生产环境中。Neuron 包括深度学习模型的编译器、运行时和工具集，可以原生集成到 PyTorch 和 TensorFlow 等框架中。借助 Neuron，用户也可以通过 Amazon SageMaker、Amazon Elastic Container Service (ECS)、Amazon Elastic Kubernetes Service (EKS)、Amazon Batch 和 Amazon Parallel Cluster 等服务进行集成。本文实践的对象是 Trainium 和 Inferentia2 对应的 EC2 实例 Trn1/Trn1n（适用于高性能和高性价比的模型训练，也可以用于模型推理）和 Inf2（专为高性能深度学习推理应用程序设计）。

NeuronCore-v2

当然，Optimum Neuron 作为 Transformers 和加速硬件的接口，最终还是需要底层硬件的硬实力来真实地输出。关于 Trainium 和 Inferentia 的架构，可以参考对应的 Neuron 文档，需要特别提一下的是，Trainium 和新一代 Inferentia2 芯片使用的都是 NeuronCore 引擎的第二代产品 NeuronCore-v2。每个 NeuronCore-v2 都是一个完全独立的异构计算单元，具有 4 个主要引擎 Tensor/Vector/Scalar/GPSIMD 引擎（如下图所示）和软件管理的 on-chip SRAM 内存，以最大化地保障数据本地性（这部分是由编译器管理，以实现数据本地性和数据预取优化）。

Transformers Neuron

在 Optimum Neuron 和 Trainium/Inferentia 加速器中间还有一处比较重要的连接点 – Transformers Neuron（对应的软件包为 transformers-neuronx），它是针对 Trn1 和 Inf2 是一个软件包，可以帮助 PyTorch 用户在第二代 Neuron 硬件 NeuronCore-v2（Trainium/Inferentia 的核心）上进行高性能的大型语言模型 LLM 的训练和推理。

Transformers Neuron 主要负责 Neuron 和 Transformers 的集成，用于在 Neuron 上运行 Transformers 解码器推理（自回归采样）的工作流。Neuron 重新实现了一些 Hugging Face Transformers 模型，以最大化 Neuron 上解码器的执行效率，这些实现都是考虑到最大兼容性进行的，这意味着可以使用标准的 HuggingFace Transformers 库训练解码器模型例如 Llama系列，然后使用 transformers-neuronx 的 LlamaForSampling 类构建一个推理优化解码器模型。

以 PyTorch 为例，上述几个组件的内容的逻辑关系如下所示，注意 Inferentia 一代对应的 Inf1 使用的 Neuron 包是 torch-neuron，Trainium 和 Inferentia2 对应的 Trn1/Trn1n 和 Inf2 使用的是 torch-neuronx，多了个“x”。

Neuron 在 LLM 推理上的性能优化

LLM 文本生成的顺序性为高效部署带来了诸多挑战，首先，为了将模型加载到设备内存中，模型通常会在多个设备上分片，这会在设备之间产生额外的通信开销和复杂性。其次，一些部署具有严格的应用级延迟控制，因此需要进行大量的延迟优化。最后，一次仅处理一个请求生成一个 token 低强度不会充分利用到设备的能力，这可以通过批处理来提高。

考虑到这些因素，Neuron SDK 提供多项内置优化，让用户直接在部署 LLM 模型时获得更好的性能，这些内容包括：

KV-caching：transformer-neuronx 库实现了 KV Cache 优化，通过重用先前计算的 Self-Attention 的 KV 来节省计算资源，而不是为每个生成的 token 重新计算它们。
模型分片：Neuron 支持模型在设备上分片，以便将大模型加载在设备内存中，这可以通过 Tensor Parallelism 等技术来实现。
计算和通信融合：Neuron 编译器会自动在计算图中将集合通信原语与后续计算融合在一起，有助于最大限度地减少在多个设备上分片模型时产生的任何开销。
紧凑的数据类型：Neuron 支持 INT8 和 FP8 数据类型，这可以显着减少模型的内存带宽和容量要求，这对于 Memory Bound 的 LLM 推理非常有帮助。
批处理：批处理可以通过将多个 token 一起处理来提高模型的推理性能。

Optimum Neuron 示例

接下来，本文会通过在 Inf2 实例上使用 PyTorch 的示例来介绍下 Optimum Neuron 的使用，希望能通过例子感受到它的易用性。另外，本文关注在 Optimum Neuron 推理方向的应用，（希望）后续的博客会继续介绍 Optimum Neuron 的模型训练部分。另外，这里的代码也都可以在 Trn1 的实例上运行，因为 Trainium 和 Inferentia2 采用了相同的架构设计。

Optimum Neuron 提供了：1）命令行 Optimum CLI 和 2）类接口 NeuronModelForXXX 两种方式将模型编译（通过 Neuron 编译器 neuronx-cc 完成）导出到加速设备上执行。两种方式是等价的，模型最终都会经过编译转换成优化过的序列化 TorchScript 模块，过程如下所示，其中 NEFF（Neuron Executable File Format）是在 Neuron 设备上的可执行二进制格式。

首先，我们在实例上进行软件的安装，这里我们使用了 us-east-1 区域的 inf2.8xlarge 实例，采用了 Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04) 镜像，如果使用其他镜像，请先参考文档完成 torch-neuronx 等组件的安装。接着，通过以下命令安装 optimum：

pip install "optimum[neuronx, diffusers]"

本文使用的是 optimum-neuron 0.0.11 版本，新增支持了 SDXL（Stable Diffusion XL）、Llama 2 的训练和推理，以及 TGI（Text Generation Inference），Neuron SDK 版本为 2.13.2。

$ pip show optimum-neuron
Name: optimum-neuron
Version: 0.0.11
Summary: Optimum Neuron is the interface between the Hugging Face Transformers and Diffusers libraries and AWS Tranium and Inferentia accelerators. It provides a set of tools enabling easy model loading, training and inference on single and multiple neuron core settings for different downstream tasks.
Home-page: https://huggingface.co/hardware/aws
Author: HuggingFace Inc. Special Ops Team
Author-email: hardware@huggingface.co
License: Apache
Location: /opt/aws_neuron_venv_pytorch/lib/python3.8/site-packages
Requires: accelerate, huggingface-hub, numpy, optimum, protobuf, transformers
Required-by:

通过 Optimum CLI 命令行工具导出模型进行推理

先通过一个简单的命令来看下效果，这里选择 distilbert-base-uncased 模型，更多参数可以参考文档：

optimum-cli export neuron \
  --model distilbert-base-uncased \
  --batch_size 16 \
  --sequence_length 128 \
  distilbert-base-uncased-neuron/

执行后可以看到过程中进行了编译并将结果保存了下来，注意这里默认会进行数据校验，会对比原模型在 CPU 上执行和在 Neuron 设备上输出的差值 max diff，可以选择进行对比查看，也可以使用通过选项–disable-validation 禁用验证：

[Compilation Time] 32.38 seconds.
[Total compilation Time] 32.38 seconds.
Validating model model...
	- Validating Neuron Model output "logits":
		-[✓] (16, 128, 30522) matches (16, 128, 30522)
		-[x] values not close enough, max diff: 0.0013232231140136719 (atol: 0.0001)
The maximum absolute difference between the output of the reference model and the Neuron exported model is not within the set tolerance 0.0001:
- logits: max diff = 0.0013232231140136719
The Neuronx export succeeded and the exported model was saved at: distilbert-base-uncased-neuron

查看目录，可以看到生成了 model.neuron 模型文件：

$ ls -lh distilbert-base-uncased-neuron/
total 318M
-rw-rw-r-- 1 ubuntu ubuntu  997 Sep 14 15:23 config.json
-rw-rw-r-- 1 ubuntu ubuntu 318M Sep 14 15:23 model.neuron

通过 Optimum CLI 命令行工具编译 SDXL 并进行推理

Optimum Neuron 新支持在 Neuron 设备 Inf2/Trn1 上编译 SDXL Pipeline 组件。使用原生 Neuron SDK 需要按照示例脚本中的内容依次对 UNet 和 VAE 中的组件进行编译，需要对 SD 相关知识有一定基础。然而，Optimum Neuron 用户可以实现一键编译，以下是一个 SDXL 模型的示例，用来生成 1024*1024 的图片：

optimum-cli export neuron \
--model stabilityai/stable-diffusion-xl-base-1.0 \
--task stable-diffusion-xl \
--batch_size 1 \
--height 1024 \
--width 1024 \
--auto_cast matmul \
--auto_cast_type bf16 \
sdxl_neuron/

等待模型编译完成，系统会显示以下提示：

The Neuronx export succeeded and the exported model was saved at: sdxl_neuron

第一次编译的时间略长一些，但是这些模型“一次编译、多次使用“，不影响后续推理。在使用 SDXL 模型时也非常方便，参考原生 diffuers 中的 StableDiffusionPipeline 方式，在 Optimum Neuron 中直接替换为 NeuronStableDiffusionXLPipeline 即可。

from optimum.neuron import NeuronStableDiffusionXLPipeline

prompt = " An astronaut riding a green horse"
stable_diffusion_xl = NeuronStableDiffusionXLPipeline.from_pretrained(
    model_id="sdxl_neuron/", device_ids=[0, 1]
)

image = stable_diffusion_xl(prompt).images[0]

以下是演示示例生成的结果：

通过 NeuronModel 类导出模型进行推理

除了命令行工具外，Optimum Neuron 也提供了针对不同模型所对应的 NeuronModelForxxx 类，以 distilbert-base-uncased-finetuned-sst-2-english 为例，通过以下代码即可导出保存编译模型：

from optimum.neuron import NeuronModelForSequenceClassification

input_shapes = {"batch_size": 1, "sequence_length": 64} 
model = NeuronModelForSequenceClassification.from_pretrained(
  "distilbert-base-uncased-finetuned-sst-2-english", export=True, **input_shapes
) 

model.save_pretrained("./distilbert-base-uncased-finetuned-sst-2-english-code-neuron/")

使用编译模型也很方便，按照对应目录加载即可，也可以结合 pipeline 进行简化使用：

from transformers import AutoTokenizer
from optimum.neuron import NeuronModelForSequenceClassification, pipeline

tokenizer = AutoTokenizer.from_pretrained("./distilbert-base-uncased-finetuned-sst-2-english-code-neuron/")
model = NeuronModelForSequenceClassification.from_pretrained("./distilbert-base-uncased-finetuned-sst-2-english-code-neuron/")

text_classification_pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)

inputs = "You’re a good guy"

print(text_classification_pipe(inputs)[0]["label"])

>>> POSITIVE

通过 NeuronModel 类编译 Llama 2 并进行推理

Optimum Neuron 新支持在 Neuron 设备 Inf2/Trn1 上编译 Llama 2模型。使用原生 Neuron SDK 需要按照示例脚本中的内容 Llama 2 的组件进行切分然后编译，这些在 Optimum Neuron 中都可以实现一键完成，可以参考以下示例：

import torch
from transformers import AutoTokenizer
from optimum.neuron import NeuronModelForCausalLM

# Instantiate and convert to Neuron a PyTorch checkpoint
# Note this will take several minutes for compilation
model = NeuronModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", export=True)

# Save model to local path
saved_path = '/tmp/llama2'
model.save_pretrained(saved_path)

# Push model to Hugging Face Hub, note using your write token
# !huggingface-cli login
# neuron_repo = <your-neuron-repo>,
# model.push_to_hub(saved_path, repository_id=neuron_repo, use_auth_token=True)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer.pad_token_id = tokenizer.eos_token_id

tokens = tokenizer("Hello, I'm a language model,", return_tensors="pt")
with torch.inference_mode():
    sample_output = model.generate(
        **tokens,
        do_sample=True,
        min_length=8,
        max_length=128,
        temperature=0.7,
    )
    outputs = [tokenizer.decode(tok) for tok in sample_output]
print(outputs[0])

>>> <s> Hello, I'm a language model, and I'm here to help you write a better essay.
I've been trained on a huge dataset of essays, so I can help you with everything from grammar and spelling to structure and style. Just type your essay into the box below, and I'll do my best to improve it. Essay about the importance of the environment. The environment is the foundation of life on earth…

需要注意的是，在模型拆分之后可以保存模型到特定路径，这样后续使用的时候就会减少很多处理时间，当然也可以选择将模型上传到 Hugging Face 上，供后续其他环境使用。

TGI（Text Generation Inference）支持

TGI 是 Hugging Face 推出的一个开源工具，用于在生产环境中部署文本生成模型，使用 Rust、Python 和 gRPC 构建，可以提供高效、可靠的推理服务，更多内容可以参考 GitHub 项目。Optimum Neuron 新版也增加了对 TGI 的支持，提供了一个可以即刻使用的容器镜像，包含了对应的 Neuron SDK、TGI 和 Neuron 接口服务, 更多内容可以参考文档，（希望）后续的博客再继续展开介绍 TGI 推理部分。

小结

如今生成式 AI 发展如火如荼，为了能更高效、更快速地使用到像 SDXL（Stable Diffusion XL）、Llama 2 这样流行的大模型来生成图片和对话，并基于此构建更智能的应用程序，亚马逊云科技结合 Hugging Face 开源了 Optimum Neuron 项目，帮助用户方便快捷地将这些大模型部署到超高性价比的云上自研芯片 Inferentia 2 和 Trainium 上进行推理服务，在大幅降低推理成本的同时，也能提高吞吐减少延迟。如果对这方面内容感兴趣，快来开启使用吧！

参考资料

https://huggingface.co/docs/optimum-neuron/index

https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/appnotes/transformers-neuronx/generative-llm-inference-with-neuron.html

亚马逊AWS官方博客