在 AWS Inferentia2 上使用 optimum-neuron 部署 Llama 2

Hugging Face 已将 optimum-neuron 与 AWS Neuron SDK 集成，使得 Llama 2 模型可以在 AWS Inferentia2 加速器上部署。此集成使用户能够利用专用硬件来提升文本生成性能，包括编码时间、延迟和吞吐量。

部署工作流程和设置

在 AWS Inferentia2 上部署 Llama 2 需要特定的环境设置和模型编译步骤，以将模型转换为与 Neuron 设备兼容的序列化格式。

环境配置

用户可以使用以下三种主要方法设置他们的 Inferentia2 实例：

  • Hugging Face Neuron Deep Learning AMI (DLAMI): 推荐的方法，提供预打包的库，包括 Optimum Neuron、Neuron 驱动、Transformers、Datasets 和 Accelerate。
  • Hugging Face Neuron SDK DLC: 用于在 Amazon SageMaker 上进行部署。
  • 手动安装： 按照全新实例上的 optimum-neuron 安装说明进行操作。

模型导出和编译

由于 Neuron 设备需要静态形状，模型在执行前必须进行编译。使用 NeuronModelForCausalLM API，用户在导出过程中指定以下参数：

  • Compiler Arguments: 定义核心数量（每个 neuron 设备有两个核心）和精度（例如 float16）。
  • Input Shapes: 设置 batch_size 和 sequence_length 的静态维度。sequence_length 至关重要，因为它限制了输入上下文、KV 缓存以及最大输出长度。

编译完成后，模型可以保存到本地或推送到 Hugging Face Hub 以供重复使用。

文本生成功能

optimum-neuron 可以使用 transformers 库进行标准文本生成，或使用简化的 optimum-neuron 管道。

生成策略

支持的生成策略包括：

  • 贪婪搜索
  • 带 top-k 和 top-p（包括温度）的多项式采样
  • 大多数 logits 预处理过滤器，例如重复惩罚

实现选项

为了简化部署，optimum-neuron 的 pipeline API 允许用户从 Hub 加载预编译模型并通过单个函数调用生成文本。

性能基准测试

基准测试使用了 Llama 2 7B 和 13B 模型，在 inf2.xlarge（预算型模型）和 inf2.48xlarge（延迟和吞吐量优化型模型）上进行了各种配置的测试。所有模型均使用最大序列长度 2048。

编码时间

编码时间——处理输入标记并生成第一个输出标记的持续时间——是感知用户延迟的关键指标。对于 256 个输入标记（典型的 Q&A），编码时间范围为 0.3s（Llama2 7B-B）至 0.9s（Llama2 7B-T）。对于 768 个输入标记（典型的 RAG），时间范围为 0.5s（Llama2 7B-B）至 5.2s（Llama2 13B-T）。

端到端延迟

端到端延迟衡量达到 1024 个标记序列长度所需的总时间。在高端 inf2.48xlarge 实例上，Llama2 7B-L 在 768 个新标记上的延迟为 6.2s，而 Llama2 13B-L 需要 10.2s。在 inf2.xlarge 上的“预算”模型表现出显著更高的延迟，达到 768 个新标记时为 47.3s。

吞吐量

吞吐量计算为每秒总标记数（端到端延迟除以 batch_size * sequence_length）。

  • 高性能： Llama2 7B-T 配置对于 256 个新标记可达到最高 750 标记/秒。
  • 预算： Llama2 7B-B 模型达到 22 到 32 标记/秒之间，Hugging Face 指出这在考虑平均人类阅读速度的情况下足以满足流式使用场景。

技术限制和未来工作

尽管性能强劲，Hugging Face 确定了两个主要的改进领域：

  1. 吞吐量扩展： 目前，吞吐量主要通过增加批次大小来提升，这受到设备内存的限制。正在将流水线作为替代方案进行集成。
  2. 上下文长度： 对静态序列长度的要求限制了处理非常长上下文的能力。团队正在探索注意力沉降作为潜在解决方案。

Sources