使用 AWS Inferentia2 让你的 Llama 生成时间飞起来

更新 (2024年2月): 性能已进一步提升！请查看我们更新后的基准测试。

在Hugging Face 博客上的一篇先前文章中，我们介绍了 AWS Inferentia2，这是第二代 AWS Inferentia 加速器，并解释了如何使用 optimum-neuron 在 AWS Inferentia 2 实例上为标准的文本和视觉任务快速部署 Hugging Face 模型。

作为与 AWS Neuron SDK 集成的进一步步骤，现在可以使用 🤗 optimum-neuron 在 AWS Inferentia2 上部署用于文本生成的 LLM 模型。

还有什么模型比 Llama 2 更适合这次演示呢？它是 Hugging Face hub 上最受欢迎的模型之一。

在你的 Inferentia2 实例上设置 🤗 optimum-neuron

我们推荐使用 Hugging Face Neuron 深度学习 AMI (DLAMI)。DLAMI 预装了所有必需的库，包括 Optimum Neuron、Neuron 驱动程序、Transformers、Datasets 和 Accelerate。

或者，您可以使用 Hugging Face Neuron SDK DLC 在 Amazon SageMaker 上进行部署。

注意：敬请关注即将发布的关于 SageMaker 部署的专门文章。

最后，这些组件也可以按照 optimum-neuron 的安装说明，在一个全新的 Inferentia2 实例上手动安装。

将 Llama 2 模型导出到 Neuron

正如 optimum-neuron 文档 中所述，模型在 Neuron 设备上运行之前需要被编译并导出为序列化格式。

幸运的是，🤗 optimum-neuron 提供了一个非常简单的 API，用于将标准的 🤗 transformers 模型导出为 Neuron 格式。

>>> from optimum.neuron import NeuronModelForCausalLM

>>> compiler_args = {"num_cores": 24, "auto_cast_type": 'fp16'}
>>> input_shapes = {"batch_size": 1, "sequence_length": 2048}
>>> model = NeuronModelForCausalLM.from_pretrained(
        "meta-llama/Llama-2-7b-hf",
        export=True,
        **compiler_args,
        **input_shapes)

这值得稍作解释

• 使用 compiler_args，我们指定希望模型部署在多少个核心上（每个 neuron 设备有两个核心），以及使用哪种精度（这里是 float16），
• 使用 input_shape，我们设置模型的静态输入和输出维度。所有模型编译器都需要静态形状，neuron 也不例外。请注意，sequence_length 不仅限制了输入上下文的长度，还限制了 KV 缓存的长度，从而也限制了输出长度。

根据您选择的参数和 inferentia 主机，这可能需要几分钟到超过一个小时的时间。

幸运的是，您只需要做一次这个操作，因为您可以保存模型并在以后重新加载它。

>>> model.save_pretrained("a_local_path_for_compiled_neuron_model")

更好的是，您可以将其推送到 Hugging Face hub。

>>> model.push_to_hub(
        "a_local_path_for_compiled_neuron_model",
        repository_id="aws-neuron/Llama-2-7b-hf-neuron-latency")

在 AWS Inferentia2 上使用 Llama 2 生成文本

一旦您的模型被导出，您就可以使用 transformers 库生成文本，正如这篇先前的文章中详细描述的那样。

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

>>> model = NeuronModelForCausalLM.from_pretrained('aws-neuron/Llama-2-7b-hf-neuron-latency')
>>> tokenizer = AutoTokenizer.from_pretrained("aws-neuron/Llama-2-7b-hf-neuron-latency")

>>> inputs = tokenizer("What is deep-learning ?", return_tensors="pt")
>>> outputs = model.generate(**inputs,
                             max_new_tokens=128,
                             do_sample=True,
                             temperature=0.9,
                             top_k=50,
                             top_p=0.9)
>>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
['What is deep-learning ?\nThe term “deep-learning” refers to a type of machine-learning
that aims to model high-level abstractions of the data in the form of a hierarchy of multiple
layers of increasingly complex processing nodes.']

注意：当向模型传递多个输入提示时，生成的 token 序列必须在左侧用流结束 token 进行填充。与导出模型一起保存的 tokenizers 已相应配置。

支持以下生成策略

• 贪心搜索，
• 带 top-k 和 top-p 的多项式采样（带温度参数）。

大多数 logits 预处理/过滤器（如重复惩罚）都受支持。

使用 optimum-neuron pipelines 实现一体化

对于喜欢简单的人来说，还有一种更简单的方法来在 AWS inferentia 2 上使用 LLM 模型，即使用 optimum-neuron pipelines。

使用它们就像这样简单

>>> from optimum.neuron import pipeline

>>> p = pipeline('text-generation', 'aws-neuron/Llama-2-7b-hf-neuron-budget')
>>> p("My favorite place on earth is", max_new_tokens=64, do_sample=True, top_k=50)
[{'generated_text': 'My favorite place on earth is the ocean. It is where I feel most
at peace. I love to travel and see new places. I have a'}]

基准测试

但是在 Inferentia2 上生成文本的效率究竟有多高呢？让我们来看看吧！

我们已经在 hub 上上传了 LLama 2 7B 和 13B 模型的预编译版本，它们具有不同的配置

注意：所有模型都以 2048 的最大序列长度进行编译。

llama2 7B "预算型" 模型旨在部署在 inf2.xlarge 实例上，该实例只有一个 neuron 设备，并有足够的 cpu 内存来加载模型。

所有其他模型都经过编译，以充分利用 inf2.48xlarge 实例上可用的全部核心。

注意：有关可用实例的详细信息，请参阅 inferentia2 产品页面。

我们为 llama2 7B 和 llama2 13B 模型创建了两个“延迟”导向的配置，它们一次只能处理一个请求，但速度最快。

我们还创建了两个“吞吐量”导向的配置，最多可以并行处理四个请求。

为了评估模型，我们从 256 个输入 token 开始，生成 tokens 直到总序列长度达到 1024（即我们生成 256、512 和 768 个 tokens）。

注意：“预算型”模型的数据会被报告，但为了图表的可读性，不会包含在图中。

编码时间

编码时间是处理输入 tokens 并生成第一个输出 token 所需的时间。这是一个非常重要的指标，因为它对应于用户在流式传输生成的 tokens 时直接感知的延迟。

我们针对不断增加的上下文大小测试编码时间，256 个输入 token 大致对应于典型的问答（Q/A）用法，而 768 个输入 token 则更典型于检索增强生成（RAG）用例。

"预算型"模型 (Llama2 7B-B) 部署在 inf2.xlarge 实例上，而其他模型则部署在 inf2.48xlarge 实例上。

编码时间以秒为单位表示。

我们可以看到，所有部署的模型都表现出优异的响应时间，即使是在长上下文的情况下。

端到端延迟

端到端延迟对应于达到 1024 个 token 序列长度所需的总时间。

因此，它包括编码和生成时间。

"预算型"模型 (Llama2 7B-B) 部署在 inf2.xlarge 实例上，而其他模型则部署在 inf2.48xlarge 实例上。

延迟以秒为单位。

所有部署在高端实例上的模型都表现出良好的延迟，即使是那些实际上配置为优化吞吐量的模型也是如此。

“预算型”部署模型的延迟明显更高，但仍然可以接受。

吞吐量

我们采用与其他基准测试相同的惯例来评估吞吐量，即将端到端延迟除以输入和输出 token 的总和。换句话说，我们将端到端延迟除以 batch_size * sequence_length，以获得每秒生成的 token 数。

"预算型"模型 (Llama2 7B-B) 部署在 inf2.xlarge 实例上，而其他模型则部署在 inf2.48xlarge 实例上。

吞吐量以**令牌/秒（tokens/second）**为单位。

同样，部署在高端实例上的模型具有非常好的吞吐量，即使是那些为延迟优化的模型也是如此。

“预算型”模型的吞吐量要低得多，但对于流式使用场景来说仍然可以接受，考虑到普通读者的阅读速度大约是每秒 5 个单词。

结论

我们已经展示了使用 🤗 optimum-neuron 在 AWS Inferentia2 上部署来自 Hugging Face hub 的 llama2 模型是多么容易。

部署的模型在编码时间、延迟和吞吐量方面表现出非常好的性能。

有趣的是，已部署模型的延迟对批量大小（batch size）不太敏感，这为它们在并行处理多个请求的推理端点上的部署开辟了道路。

不过，仍有很大的改进空间

• 在当前的实现中，增加吞吐量的唯一方法是增加批处理大小，但目前这受限于设备内存。诸如流水线（pipelining）之类的替代方案目前正在集成中，
• 静态序列长度限制了模型编码长上下文的能力。看看 attention sinks 是否是解决这个问题的有效选项将会很有趣。