使用 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 是否是解決這個問題的有效選項將會很有趣。