在 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 Drivers、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.3 秒（Llama2 7B-B）與 0.9 秒（Llama2 7B-T）之間。對於 768 個輸入標記（典型的 RAG），時間介於 0.5 秒（Llama2 7B-B）與 5.2 秒（Llama2 13B-T）之間。

端到端延遲

端到端延遲衡量達到 1024 個標記序列長度所需的總時間。在高端 inf2.48xlarge 實例上，Llama2 7B-L 對於 768 個新標記達成了 6.2 秒的延遲，而 Llama2 13B-L 則花費 10.2 秒。在 inf2.xlarge 上的「預算」模型表現出顯著更高的延遲，達到 47.3 秒的 768 個新標記。

吞吐量

吞吐量計算為每秒總標記數（端到端延遲除以 batch_size * sequence_length）。

  • 高效能： Llama2 7B-T 配置在 256 個新標記時可達到每秒 750 個標記。
  • 預算： Llama2 7B-B 模型達到每秒 22 到 32 個標記，Hugging Face 指出這足以應付串流使用案例，考慮到平均人類閱讀速度。

技術限制與未來工作

雖然效能強勁，Hugging Face 指出兩個主要可改進的領域：

  1. 吞吐量擴展： 目前，吞吐量主要透過增加 batch size 來提升，但這受到裝置記憶體的限制。正在將管線化作為替代方案進行整合。
  2. 上下文長度： 對靜態序列長度的需求限制了處理非常長上下文的能力。團隊正在探索注意力沉澱作為潛在的解決方案。

Sources