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

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

為什么要研究推理能力？

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

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

Llama2–70B 型號(hào)

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

理論與實(shí)踐

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

我們還將其與理論計(jì)算進(jìn)行了比較，看看它們?cè)诂F(xiàn)實(shí)世界中的表現(xiàn)如何。

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

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

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

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

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

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

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

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

因此，我們可以使用以下公式來(lái)估計(jì)中間變量所需的內(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 * 輸入長(zhǎng)度 * 輸入長(zhǎng)度) / 8 / 1000 / 1000 / 1000

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

= 5.12 * 10e-7 * 輸入長(zhǎng)度**2

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

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

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

≈ 35 + 5.12 * 10e-7 * 輸入長(zhǎng)度**2

結(jié)果和啟示

我們隨后進(jìn)行了實(shí)驗(yàn)來(lái)評(píng)估真實(shí)場(chǎng)景中 GPU 的使用情況。在我們的實(shí)驗(yàn)中，我們使用 HuggingFace 框架對(duì) 4 位量化的 Llama2-70b 模型進(jìn)行推理。下表展示了不同負(fù)載下的GPU使用情況：

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

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

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

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

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

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

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

結(jié)論

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

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

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

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