解鎖最大推理能力：深入研究80GB A100 GPU 上的Llama2–70B

在不斷發(fā)展的自然語言處理領(lǐng)域，模型的性能是最重要的問題。研究人員、工程師和數(shù)據(jù)科學家不斷利用日益強大的硬件和復雜的算法突破可能的界限。其中一項追求是確定 Llama2-70B 等模型在專用硬件（如 80GB A100 GPU）上運行時的最大推理能力。在本文中，我們深入探討了此類調(diào)查的理論、實踐和結(jié)果。

為什么要研究推理能力？

出于多種原因，了解模型的最大推理能力至關(guān)重要。首先，它可以幫助研究人員和開發(fā)人員就模型部署做出明智的決策。了解模型在特定硬件上的限制對于高效資源分配和避免意外瓶頸至關(guān)重要。

其次，這些知識對于在實時應(yīng)用程序中依賴自然語言處理模型的組織至關(guān)重要。聊天機器人、虛擬助理和自動化客戶支持系統(tǒng)等應(yīng)用程序需要模型能夠快速、準確地做出響應(yīng)。確定推理能力可確保這些應(yīng)用程序能夠有效地處理不同的工作負載。

Llama2–70B 型號

在本次研究中，Llama2-70B 模型占據(jù)了中心位置。Llama2–70B 是 Llama2 模型的變體，具有 4 位量化功能。術(shù)語“4 位量化”是指一種用于降低模型權(quán)重精度的技術(shù)，使模型的存儲和計算效率更高，同時犧牲最小的精度。

理論與實踐

為了測試 Llama2-70B 模型在 80GB A100 GPU 上的最大推理能力，我們要求一位研究人員部署 Llama2 模型并將其推向極限，以查看它到底可以處理多少個令牌。不同的輸入標記長度對應(yīng)于不同的近似字數(shù)。每次測試運行時都會監(jiān)控并記錄 GPU 使用情況。

我們還將其與理論計算進行了比較，看看它們在現(xiàn)實世界中的表現(xiàn)如何。

為了計算 Llama2 推理所需的內(nèi)存，我們需要考慮以下兩點：

1.模型權(quán)重所需的內(nèi)存。

2. 推理過程中中間變量所需的內(nèi)存。

模型權(quán)重所需的內(nèi)存

在計算部署模型所需的 GPU 使用率時，我們主要考慮的是模型的參數(shù)大小。Llama2–70B 模型是一個擁有 700 億個參數(shù)的大型語言模型。當我們使用 4 位量化時，這意味著每個參數(shù)將使用 4 位來存儲。因此，對于 700 億個參數(shù)，我們需要等于 70 * 1000000000 * 4 位的 GPU 內(nèi)存。這大約相當于 35 GB。

中間變量所需的內(nèi)存

在計算模型中推理所需的中間變量時，最重要的中間變量包括注意力機制中的Query、Key和Value矩陣，以及Query矩陣與Key矩陣相乘得到的Attention Weights矩陣。下表顯示了這四個矩陣的形狀及其在 Llama2-70b 模型中的實際大小。

由于本實驗中輸入的token長度超過2000，因此四個矩陣中內(nèi)存密集的部分是注意力權(quán)重矩陣。在調(diào)試 HuggingFace 的代碼來推斷 Llama2-70b 模型時，觀察到通過將 Query 矩陣與 Key 矩陣相乘來計算 Attention Weights 矩陣時，其 dtype 為 16 位，但在實際計算過程中使用了兩次空間。此外，在 Softmax 計算期間，dtype 擴展到 32 位。因此，可以理解，模型實際推理時，需要四倍的16位Attention Weights矩陣。

因此，我們可以使用以下公式來估計中間變量所需的內(nèi)存：

內(nèi)存（以 GB 為單位）：

= 4 * 16位注意力權(quán)重矩陣 / 8 / 1000 / 1000 / 1000

= 4 * 16 * 注意力權(quán)重矩陣的大小 / 8 / 1000 / 1000 / 1000

= 4 * 16 * (1 * 64 * 輸入長度 * 輸入長度) / 8 / 1000 / 1000 / 1000

= (4 * 16 * 64 / 8 / 1000 / 1000 / 1000) * 輸入長度 * 輸入長度

= 5.12 * 10e-7 * 輸入長度**2

用于推理的總內(nèi)存使用量

根據(jù)上面提到的模型權(quán)重和中間變量兩部分，我們可以估計 GPU 的總使用量（以 GB 為單位）：

總內(nèi)存使用量 = 模型權(quán)重內(nèi)存使用量 + 中間變量內(nèi)存使用量

≈ 35 + 5.12 * 10e-7 * 輸入長度**2

結(jié)果和啟示

我們隨后進行了實驗來評估真實場景中 GPU 的使用情況。在我們的實驗中，我們使用 HuggingFace 框架對 4 位量化的 Llama2-70b 模型進行推理。下表展示了不同負載下的GPU使用情況：

然后，我們將實際實驗數(shù)據(jù)點（藍色）與我們計算的理論使用曲線（紅色）繪制在一張圖表上，如下所示：

上圖的結(jié)果可以清楚地了解 Llama2–70B 模型在 80GB A100 GPU 上的行為：

· 當不提供輸入標記時，GPU 使用量約為 35.0 GB，這與我們計算的模型參數(shù)值相符。

· 我們可以看到我們測試的實際GPU使用率與我們計算的理論曲線非常吻合。隨著輸入令牌計數(shù)的增加，GPU 使用率呈二次方上升。這與我們理論計算公式中輸入長度的二次項一致。

· 根據(jù) GPU 的容量限制，最大推理能力達到約 8500 個令牌。

8500 個令牌的推理能力大致相當于 5100 個單詞，代表 GPU 得到充分利用的時刻，輸入令牌長度的進一步增加會導致 CUDA 內(nèi)存不足錯誤。此外，該觀察結(jié)果與理論理解相一致，即 Transformers 中的注意力機制是 Llama2-70B 等模型的關(guān)鍵組成部分，隨著輸入長度的增加，GPU 使用率會呈二次方增長。

基礎(chǔ) Llama2–70B 模型的上下文長度為 4096 個標記，因此該實驗表明該模型可以輕松安裝到一張 80Gb 卡上（使用 4 位量化）。然而，由于存在擴展可用上下文窗口的技術(shù)，因此了解一張卡可以處理的最大值極其重要，即使它超出了理論最大值。

結(jié)論

為了最大化自然語言處理模型的推理能力，理解模型架構(gòu)和硬件之間的相互作用至關(guān)重要。在這項研究中，4 位量化的 Llama2-70B 模型在 80GB A100 GPU 上展示了大約 8500 個令牌的最大推理能力。

這一發(fā)現(xiàn)為旨在優(yōu)化自然語言處理管道并確保在專用硬件上高效部署模型的從業(yè)者和組織提供了寶貴的見解。它還強調(diào)了仔細平衡模型復雜性和硬件資源以在實際應(yīng)用中實現(xiàn)最佳性能的重要性。

信息來源：https://nigelcannings.medium.com/unlocking-maximum-inference-capability-a-deep-dive-into-llama2-70b-on-an-80gb-a100-gpu-2ab1158d6b0b

2023年人工智能訓練與推理工作站、服務(wù)器、集群硬件配置推薦