經(jīng)過“九九八十一難”，大模型終于煉成。下一步就是架設(shè)服務(wù)，準(zhǔn)備開門營業(yè)了。真這么簡單？恐怕未必！行百里者半九十，推理優(yōu)化又是新的雄關(guān)漫道。如何進(jìn)行延遲優(yōu)化？如何進(jìn)行成本優(yōu)化 (別忘了 OpenAI 8K 上下文的 GPT-4 模型，提示每 1000 詞元只需 0.03 美金，補(bǔ)全每 1000 詞元只需 0.06 美金)？如何在延遲和吞吐量之間折衷？如何處理大模型特有的分布式推理后端和網(wǎng)絡(luò)服務(wù)前端的協(xié)作問題……要不動手之前還是先看看 BLOOM 推理服務(wù)踩過的坑吧！

您可以通過點擊 閱讀原文 查看文段中提到的鏈接。

本文介紹了我們在實現(xiàn) BLOOM 模型高效推理服務(wù)的過程中發(fā)生的幕后故事。

在短短數(shù)周內(nèi)，我們把推理延遲降低了 5 倍 (同時，吞吐量增加了 50 倍)。我們將分享我們?yōu)檫_(dá)成這一性能改進(jìn)而經(jīng)歷的所有斗爭和史詩般的勝利。

在此過程中，不同的人參與了不同的階段，嘗試了各種不同的優(yōu)化手段，我們無法一一羅列，還請多多包涵。如果你發(fā)現(xiàn)本文中某些內(nèi)容可能已過時甚至完全錯誤，這也不奇怪，因為一方面對于如何優(yōu)化超大模型性能我們?nèi)栽谂W(xué)習(xí)中，另一方面，市面上新硬件功能和新優(yōu)化技巧也層出不窮。

我們很抱歉如果本文沒有討論你最中意的優(yōu)化技巧，或者我們對某些方法表述有誤。但請告訴我們，我們非常樂意嘗試新東西并糾正錯誤。

訓(xùn)練 BLOOM

這是不言而喻的，如果不先獲取到大模型，那推理優(yōu)化就無從談起。大模型訓(xùn)練是一項由很多不同的人共同領(lǐng)導(dǎo)的超級工程。

為了最大化 GPU 的利用率，我們探索了多種訓(xùn)練方案。最后，我們選擇了 Megatron-Deepspeed 來訓(xùn)練最終模型。這意味著訓(xùn)練代碼與?transformers?庫并不完全兼容。

移植至 transformers

由于上文提及的原因，我們第一件事是將現(xiàn)有模型移植到?transformers?上。我們需要從訓(xùn)練代碼中提取相關(guān)代碼并將其實現(xiàn)至?transformers?里。Younes 負(fù)責(zé)完成了這項工作。這個工作量絕對不小，我們大概花了將近一個月的時間，進(jìn)行了 200 次提交?才最終完成。

有幾點需要注意，我們后面還會提到:

小版的模型，如 bigscience/bigscience-small-testing 和 bigscience/bloom-560m 非常重要。因為模型結(jié)構(gòu)與大版的一樣但尺寸更小，所以在它們上面一切工作 (如調(diào)試、測試等) 都更快。

首先，你必須放棄那種最終你會得到比特級一致的?logits?結(jié)果的幻想。不同的 PyTorch 版本間的算子核函數(shù)更改都會引入細(xì)微差別，更不用說不同的硬件可能會因為體系架構(gòu)不同而產(chǎn)生不同的結(jié)果 (而出于成本原因，你可能并不能一直在 A100 GPU 上開發(fā))。

一個好的嚴(yán)格的測試套件對所有模型都非常重要

我們發(fā)現(xiàn)，最佳的測試方式是使用一組固定的提示。從測試角度，你知道提示 (prompt)，而且你想要為每個提示生成確定性的補(bǔ)全 (completion)，所以解碼器用貪心搜索就好了。如果兩次測試生成的補(bǔ)全是相同的，你基本上可以無視 logits 上的小差異。每當(dāng)你看到生成的補(bǔ)全發(fā)生漂移時，就需要調(diào)查原因。可能是你的代碼沒有做它應(yīng)該做的事; 也有可能是你的提示不在該模型的知識域內(nèi) [譯者注: 即模型的訓(xùn)練數(shù)據(jù)中并不包含提示所涉及的話題]，所以它對噪聲更敏感。如果你有多個提示且提示足夠長，不太可能每個提示都觸發(fā)上述不在知識域的問題。因此，提示越多越好，越長越好。

第一個模型 (small-testing) 和大 BLOOM 一樣，精度是?bfloat16?的。我們原以為兩者應(yīng)該非常相似，但由于小模型沒有經(jīng)過太多訓(xùn)練或者單純只是性能差，最終表現(xiàn)出來的結(jié)果是它的輸出波動很大。這意味著我們用它進(jìn)行生成測試會有問題。第二個模型更穩(wěn)定，但模型數(shù)據(jù)精度是?float16?而不是?bfloat16，因此兩者間的誤差空間更大。

公平地說，推理時將?bfloat16?模型轉(zhuǎn)換為?float16?似乎問題不大 (?bfloat16?的存在主要是為了處理大梯度，而推理中不存在大梯度)。

在此步驟中，我們發(fā)現(xiàn)并實現(xiàn)了一個重要的折衷。因為 BLOOM 是在分布式環(huán)境中訓(xùn)練的，所以部分代碼會對 Linear 層作張量并行，這意味著在單 GPU 上運(yùn)行相同的操作會得到?不同的數(shù)值結(jié)果。我們花了一段時間才查明這個問題。這個問題沒辦法徹底解決，要么我們追求 100% 的數(shù)值一致性而犧牲模型運(yùn)行速度，要么我們接受每次生成時都會出現(xiàn)一些小的差異但運(yùn)行速度更快，代碼更簡單。我們?yōu)榇嗽O(shè)了一個標(biāo)志位供用戶自己配置。

首次推理 (PP + Accelerate)

注意:?這里，流水線并行?(Pipeline Parallelism, PP)?意味著每個 GPU 將分得模型的一些層，因此每個 GPU 將完成一部分操作，然后再將其結(jié)果交給下一個 GPU。


現(xiàn)在我們有了一個能支持 BLOOM 的 ?transformers，我們可以開始跑了。

BLOOM 是一個 352GB (176B bf16 參數(shù)) 的模型，我們至少需要那么多顯存才能放下它。我們花了一點時間試了試在小顯存的 GPU 上使用 CPU 卸載的方式來推理，但是推理速度慢了幾個數(shù)量級，所以我們很快放棄了它。

然后，我們轉(zhuǎn)而想使用?transformers?的 pipeline API，吃一下這個 API 的狗糧。然而，?pipeline?不是分布式感知的 (這不是它的設(shè)計目標(biāo))。

經(jīng)過短暫的技術(shù)方案討論，我們最終使用了?accelerate?的新功能?device_map="auto?來管理模型的分片。我們不得不解決一些?accelerate?以及?transformers?的 bug，才使得這一方案能正常工作。

它的工作原理是將 transformer 模型按層進(jìn)行切分，每個 GPU 分到一些層。真正運(yùn)行時，是 GPU0 先開始工作，然后將結(jié)果交給 GPU1，依次下去。

最后，在前端架一個小型 HTTP 服務(wù)器，我們就可以開始提供 BLOOM (大模型) 推理服務(wù)了??！

起點

至此，我們甚至還沒有開始討論優(yōu)化！

我們其實做了不少優(yōu)化，這一切過程有點像紙牌疊城堡游戲。在優(yōu)化期間，我們將對底層代碼進(jìn)行修改，所以一定要確保我們不會以任何方式破壞模型，這一點非常重要，而且其實比想象中更容易做到。

優(yōu)化的第一步是測量性能。在整個優(yōu)化過程中，性能測量貫穿始終。所以，首先需要考慮我們需要測量什么，也即我們關(guān)心的是什么。對于一個支持多種選項的開放式推理服務(wù)而言，用戶會向該服務(wù)發(fā)送各種不同的查詢請求，我們關(guān)心的是:

1. 我們可以同時服務(wù)的用戶數(shù)是多少 (吞吐量)？

2. 我們平均為每個用戶服務(wù)的時間是多少 (延遲)？

我們用?locust?做了一個測試腳本，如下:

注意: 這不是我們最佳的也不是唯一的負(fù)載測試，但始終是我們第一個運(yùn)行的負(fù)載測試，因此它可用于公平地比較不同方案。在此基準(zhǔn)測試表現(xiàn)最好并不意味著它絕對是最好的解決方案。我們還需要使用其他更復(fù)雜的測試場景來模擬真實場景的真實性能。

我們想觀察各種實現(xiàn)方案部署時如何爬坡，并確保在熔斷時適當(dāng)?shù)亟档头?wù)器負(fù)載。熔斷意味著原本能 (快速) 響應(yīng)你的請求的服務(wù)不再響應(yīng)你的請求，因為同一時間有太多人想要使用它。避免?死亡之擁 (hug of death)?是極其重要的。[譯者注: 死亡之擁是一個互聯(lián)網(wǎng)領(lǐng)域的隱喻，意指由于極端峰值流量而導(dǎo)致互聯(lián)網(wǎng)服務(wù)宕機(jī)]

在上述基準(zhǔn)測試中，我們得到的初始性能是 (使用 GCP 上的 16xA100 40G 環(huán)境測得，本文后續(xù)所有測試都基于該環(huán)境):

每秒處理請求數(shù) (吞吐量): 0.3 每詞元延遲: 350ms

這兩個值并不是很好。在正式開始工作之前，我們可以預(yù)估一下我們能得到的最好結(jié)果。BLOOM 模型所需的計算量公式為?，其中?B?是 batch size，?s?是序列長度，?h?是隱含層維度。

讓我們算一下，一次前向傳播需要?17 TFlop。A100 的?規(guī)格?為單卡?312 TFLOPS。這意味著單個 GPU 最多能達(dá)到?17 / 312 = 54 毫秒/詞元?的延遲。我們用了 16 個 GPU，因此可得?3 毫秒/詞元。這只是個上限，我們永遠(yuǎn)不可能達(dá)到這個值，況且現(xiàn)實中卡的性能很少能達(dá)到其規(guī)格所宣稱的數(shù)字。此外，如果你的模型并不受限于計算 [譯者注: 如受限于內(nèi)存帶寬、受限于 IO 帶寬等]，那么這個值你也達(dá)不到。知道理想值，只是為了讓我們對優(yōu)化目標(biāo)心里有個數(shù)。在這里，我們到目前為止與理想值差 2 個數(shù)量級。此外，這個估計假設(shè)你將所有算力都用于延遲型服務(wù)，這意味著一次只能執(zhí)行一個請求 (沒關(guān)系，因為你正在最大化你的機(jī)器利用率，所以沒有太多其他事情要做; 但另一個思路是，我們可以犧牲一點延遲，通過批處理方式來獲得更高的吞吐量)。

探索多條路線

注意:?這里，張量并行?(Tensor Parallelism，TP)?意味著每個 GPU 將擁有部分權(quán)重，因此所有 GPU 始終處于工作狀態(tài)，專注于分給它的部分工作。通常這會帶來非常輕微的開銷，因為會有一些工作是重復(fù)的，更重要的是，GPU 必須定期相互通信交流它們的結(jié)果，然后再繼續(xù)計算。


現(xiàn)在我們已經(jīng)比較清楚地了解了我們的處境，是時候開始工作了。

我們根據(jù)我們自己及其他人的各種經(jīng)驗和知識嘗試了各種方法。

每次嘗試都值得寫一篇專門的博文，由于篇幅所限，在這里我們僅將它們列出來，并只深入解釋并研究那些最終應(yīng)用到當(dāng)前服務(wù)中去的技術(shù)的細(xì)節(jié)。從流水線并行 (PP) 切換到張量并行 (TP) 是延遲優(yōu)化的一個重要一步。每個 GPU 將擁有部分參數(shù)，并且所有 GPU 將同時工作，所以延遲應(yīng)該會迅速下降。但是付出的代價是通信開銷，因為它們的中間結(jié)果需要經(jīng)?；ハ嗤ㄐ?。

需要注意的是，這里涉及的方法相當(dāng)廣泛。我們會有意識地學(xué)習(xí)更多關(guān)于每個工具的知識，以及在后續(xù)優(yōu)化中如何使用它。

將代碼移植到 JAX/Flax 中以在 TPU 上運(yùn)行

• 并行方案的選擇更加容易。因此 TP 的測試會更方便，這是 JAX 的設(shè)計帶來的好處之一。

• 對硬件的限制更多，JAX 上 TPU 的性能可能比 GPU 更好，但 TPU 比 GPU 更難獲取 (只在 GCP 上有，數(shù)量也沒有 GPU 多)。

• 缺點: 需要移植工作。但無論如何，把它集成到我們的庫里面這件事肯定是受歡迎的。

結(jié)果:

• 移植比較麻煩，因為某些條件語句和核函數(shù)很難準(zhǔn)確復(fù)制，但尚可勉力為之。

• 一旦移植完后，測試各種并行方案就比較方便。感謝 JAX，沒有食言。

• 事實證明，在 Ray 集群里與 TPU worker 通信對我們來講真的太痛苦了。

不知道是工具原因還是網(wǎng)絡(luò)的原因，或者僅僅是因為我們不太懂，但這事實上減慢了我們的實驗速度，而且需要的工作比我們預(yù)期的要多得多。我們啟動一個需要 5 分鐘時間運(yùn)行的實驗，等了 5 分鐘沒有發(fā)生任何事情，10 分鐘之后仍然沒有任何事情發(fā)生，結(jié)果發(fā)現(xiàn)是一些 TPU worker 宕機(jī)了或者是沒有響應(yīng)。我們不得不手動登進(jìn)去看，弄清楚發(fā)生了什么，修復(fù)它，重啟一些東西，最后再重新啟動實驗，就這樣半小時過去了。幾次下來，幾天就沒了。我們再強(qiáng)調(diào)一下，這未必真的是我們使用的工具的問題，但我們的主觀體驗確實如此。

• 無法控制編譯

我們運(yùn)行起來后，就嘗試了幾種設(shè)置，想找出最適合我們心目中想要的推理性能的設(shè)置，結(jié)果證明很難從這些實驗中推測出延遲/吞吐量的規(guī)律。例如，在 batch_size=1 時吞吐量有 0.3 RPS (Requests Per Second, RPS) (此時每個請求/用戶都是獨立的)，延遲為 15 毫秒/詞元 (不要與本文中的其他數(shù)字進(jìn)行太多比較，TPU 機(jī)器與 GPU 機(jī)器大不相同)，延遲很好，但是總吞吐量跟之前差不多。所以我們決定引入批處理，在 batch_size=2 的情況下，延遲增加到原來的 5 倍，而吞吐量只提高到原來的 2 倍…… 經(jīng)過進(jìn)一步調(diào)查，我們發(fā)現(xiàn)一直到 batch_size=16，每個 batch_size 之間的延遲都差不多。因此，我們可以以 5 倍的延遲為代價獲得 16 倍的吞吐量?？瓷先ネΣ诲e的，但我們更希望對延遲有更細(xì)粒度的控制，從而使得延遲能滿足 100ms, 1s, 10s, 1mn 規(guī)則中的各檔。

使用 ONNX/TRT 或其他編譯方法

• 它們應(yīng)該能處理大部分優(yōu)化工作

• 缺點: 通常需要手動處理并行性

結(jié)果:

• 事實證明，為了能夠 trace/jit/export 模型，我們需要重寫 PyTorch 相關(guān)的一部分代碼，使其能夠很容易與純 PyTorch 方法相融合?？傮w來講，我們發(fā)現(xiàn)我們可以通過留在 PyTorch 中獲得我們想要的大部分優(yōu)化，使我們能夠保持靈活性而無需進(jìn)行太多編碼工作。另一件值得注意的事情是，因為我們在 GPU 上運(yùn)行，而文本生成有很多輪前向過程，所以我們需要張量留在 GPU 上，有時很難將你的張量輸給某個庫，返回結(jié)果，計算 logits (如 argmax 或采樣)，再回輸給那個庫。

將循環(huán)放在外部庫里面意味著像 JAX 一樣失去靈活性，這不是我們設(shè)想的推理服務(wù)應(yīng)用場景的使用方法。

DeepSpeed

• 這是我們訓(xùn)練 BLOOM 時使用的技術(shù)，所以用它來推理也很公平

• 缺點: DeepSpeed 之前從未用于推理，其設(shè)計也沒準(zhǔn)備用于推理

結(jié)果:

• 我們很快就得到了很不錯的結(jié)果，這個結(jié)果與我們現(xiàn)行方案的上一版性能大致相同。

• 我們必須想出一種方法，在多進(jìn)程上架設(shè)用于處理并發(fā)請求網(wǎng)絡(luò)服務(wù)，因為現(xiàn)在一個推理任務(wù)是由多個 DeepSpeed 進(jìn)程完成的 (每個 GPU 一個進(jìn)程)，。有一個優(yōu)秀的庫 Mii 可供使用，它雖然還達(dá)不到我們所設(shè)想的極致靈活的目標(biāo)，但我們現(xiàn)在可以在它之上開始我們的工作。(當(dāng)前的解決方案稍后討論)。

• 我們在使用 DeepSpeed 時遇到的最大問題是缺乏穩(wěn)定性。

我們在 CUDA 11.4 上運(yùn)行基于 11.6 編譯的代碼時遇到了問題。而其中一個由來已久的、我們永遠(yuǎn)無法真正解決的問題是: 經(jīng)常會發(fā)生核函數(shù)崩潰 (CUDA 非法訪問、尺寸不匹配等)。我們修復(fù)了其中一些問題，但在壓測我們的網(wǎng)絡(luò)服務(wù)時，我們永遠(yuǎn)無法完全實現(xiàn)穩(wěn)定性。盡管如此，我想向幫助過我們的 Microsoft 人員說，感謝那些非常愉快的交流，它們提高了我們對正在發(fā)生的事情的理解，并為我們的后續(xù)工作提供了真知灼見。

• 另一個痛點是我們的團(tuán)隊主要在歐洲，而微軟在加利福尼亞，所以合作時間很棘手，我們因此損失了大量時間。這與技術(shù)部分無關(guān)，但我們確實認(rèn)識到合作的組織部分也非常重要。

• 另一件需要注意的事情是，DeepSpeed 依賴于?transformers?來注入其優(yōu)化，并且由于我們一直在更新我們的代碼，這使得 DeepSpeed 團(tuán)隊很難在我們的主分支上工作。很抱歉讓它變得困難，這也可能是?transformers?被稱為技術(shù)最前沿的原因。

有關(guān) Web 服務(wù)的想法

• 鑒于我們準(zhǔn)備運(yùn)行一個免費(fèi)服務(wù)，支持用戶向該服務(wù)發(fā)送長短不一的文本，并要求獲取短至幾個詞，長至如整個食譜那么長的回應(yīng)，每個請求的參數(shù)也可以各不相同，web 服務(wù)需要做點什么來支持這個需求。

結(jié)果:

• 我們使用綁定庫?tch-rs?在 ?Rust?中重寫了所有代碼。Rust 的目標(biāo)不是提高性能，而是對并行性 (線程/進(jìn)程) 以及 web 服務(wù)和 PyTorch 的并發(fā)性進(jìn)行更細(xì)粒度的控制。由于 GIL 的存在，Python 很難處理這些底層細(xì)節(jié)。

• 結(jié)果表明，大部分的痛苦來自于移植工作，移植完后，實驗就輕而易舉了。我們認(rèn)為，通過對循環(huán)進(jìn)行精確的控制，即使在具有大量不同屬性的請求的場景中，我們也可以為每個請求提供出色的性能。如果你感興趣的話，可以查看?代碼，但這份代碼沒有任何支持，也沒有好的文檔。

• Rust web 服務(wù)投入生產(chǎn)了幾周，因為它對并行性的支持更寬松，我們可以更有效地使用 GPU (如使用 GPU0 處理請求 1，而 GPU1 處理請求 0)。在保持延遲不變的情況下，我們把吞吐從 0.3 RPS 提高到了 ~2.5 RPS。雖然在最理想情況下，我們能將吞吐提高到 16 倍。但實際工作負(fù)載上的測出來能到 8 倍左右的話也還算不錯。

純 PyTorch

• 純粹修改現(xiàn)有代碼，通過刪除諸如?reshape?之類的操作、使用更優(yōu)化的核函數(shù)等方法來使其運(yùn)行速度更快。

• 缺點: 我們必須自己編寫 TP 代碼，并且我們還有一個限制，即修改后代碼最好仍然適合我們的庫 (至少大部分)。

結(jié)果

• 在下一章詳述。

最終路線: PyTorch + TP + 1 個自定義內(nèi)核 + torch.jit.script

編寫更高效的 PyTorch

第一件事是在代碼中刪除不必要的操作?？梢酝ㄟ^代碼走查并找出明顯可被刪除的某些操作:

• Alibi 在 BLOOM 中用于添加位置嵌入 (position embeddings)，源代碼中計算 Alibi 的地方太多，每次都重新計算一次，我們優(yōu)化成只計算一次，這樣效率更高。

舊代碼:?鏈接新代碼:?鏈接

這個改動獲得了 10 倍的加速，最新版本還增加了對填充 (padding) 的支持！由于此步驟僅計算一次，因此在這里，運(yùn)算本身實際速度并不重要，而總體上減少操作和張量創(chuàng)建的次數(shù)更重要。

當(dāng)你開始?剖析?代碼性能時，其他部分會越來越清晰，我們大量地使用了 tensorboard 來幫助我們進(jìn)行性能剖析。它提供了如下圖所示的這類圖像，可以提供有關(guān)性能的洞見:

注意力層占用了很多時間，注意這是一個 CPU 視圖，所以條形很長并不意味著核函數(shù)執(zhí)行時間很長，它只意味著 CPU 正在等待上一步的 GPU 結(jié)果。

我們還在?baddbmm?操作之前看到許多?cat?操作。

再舉個例子，在刪除大量?reshape?/?transpose?后，我們在 tensorboard 中發(fā)現(xiàn):

• 注意力是性能熱點 (這是預(yù)期的，但能夠通過測量數(shù)據(jù)來驗證總是好的)。

• 在注意力中，由于大量的 reshape，很多核函數(shù)其實是顯存拷貝函數(shù)。

• 我們?可以?通過修改權(quán)重和?past_key_values?的內(nèi)存布局來移除?reshape。這個改動有點大，但性能確實有一定的提高！

支持 TP

好了，我們已經(jīng)拿到了大部分唾手可得的成果，現(xiàn)在我們的 PP 版本的延遲從大約 350 毫秒/詞元降低到 300 毫秒/詞元。延遲降低了 15%，實際情況收益更大，但由于我們最初的測量并不是非常嚴(yán)格，所以就用這個數(shù)吧。

然后我們繼續(xù)實現(xiàn)一個 TP 版。進(jìn)度比我們預(yù)期的要快得多，一個 (有經(jīng)驗的) 開發(fā)人員僅花了半天時間就實現(xiàn)出來了，代碼見?此處。在此過程中，我們還重用了一些其他項目的代碼，這對我們很有幫助。

延遲從 300 毫秒/詞元直接變?yōu)?91 毫秒/詞元，這是用戶體驗的巨大改進(jìn)。一個簡單的 20 個詞元的請求延遲從 6 秒變成了 2 秒，用戶體驗直接從“慢”變成了輕微延遲。

此外，吞吐量上升了很多，達(dá)到 10 RPS。batch_size=1 和 batch_size=32 延遲基本相同，因此，從這種意義上來講，在相同的延遲下，吞吐量的上升基本上是?免費(fèi)?的。

唾手可得的果實

現(xiàn)在我們有了一個 TP 版本的實現(xiàn)，我們可以再次開始進(jìn)行性能剖析和優(yōu)化。因為并行方案發(fā)生了改變，我們有必要再從頭開始分析一遍。

首先，同步 (?ncclAllReduce) 開始成為主要熱點，這符合我們的預(yù)期，同步需要花時間。但我們不打算優(yōu)化這一部分，因為它已經(jīng)使用了?nccl。雖然可能還有一些改進(jìn)空間，但我們認(rèn)為我們很難做得更好。

第二個是?Gelu?算子，我們可以看到它啟動了許多?element-wise?類的核函數(shù)，總體而言它占用的計算份額比我們預(yù)期的要大。

我們對?Gelu?作了如下修改:

從

改成了

我們使用?jit?將許多小的?element-wise?核函數(shù)融合成了一個核函數(shù)，從而節(jié)省了核函數(shù)啟動開銷和內(nèi)存拷貝開銷。

該優(yōu)化降低了 10% 的延遲，從 91 毫秒/詞元到 81 毫秒/詞元，搞定！

不過要小心，這種方法可不是任何時候都有效，算子融合不一定每次都會發(fā)生。另外如果原來的算子實現(xiàn)已經(jīng)非常高效了，就算融合了也不能帶來很多的增益。

我們發(fā)現(xiàn)它在下面幾個場合有用:

• 你有很多小的、element-wise?的操作

• 你的性能熱點里有一些難以去除的?reshape?算子，這些算子一般就是拷貝

• 算子能融合時

滑鐵盧

在測試期間，有一段時間，我們觀察到 Rust 服務(wù)的延遲比 Python 服務(wù)低 25%。這很奇怪，但因為它們的測試環(huán)境是一致的，而且去除了核函數(shù)后我們還是能測到這個速度增益，我們開始感覺，也許降低 Python 開銷可以帶來不錯的性能提升。

我們開始了為期 3 天的重新實現(xiàn)?torch.distributed?部分代碼的工作，以便在 Rust 里運(yùn)行 nccl-rs。代碼能工作，但生成的句子與 Python 版有些不一樣，于是我們開始調(diào)查這些問題，就在這個過程中，我們發(fā)現(xiàn) ……?在測量 PyTorch 版性能時，我們忘記刪除 PyTorch 里的 profiler 代碼了?……

我們遭遇了滑鐵盧，刪除 profiler 代碼后延遲降低了 25%，兩份代碼延遲一樣了。其實我們最初也是這么想的，Python 一定不會影響性能，因為模型運(yùn)行時運(yùn)行的主要還是 torch cpp 的代碼。雖然 3 天其實也不算啥，但發(fā)生這樣的事還是挺糟糕的。

針對錯誤的或不具代表性的測量數(shù)據(jù)進(jìn)行優(yōu)化，這很常見，優(yōu)化結(jié)果最終會令人失望甚至對整個產(chǎn)品帶來反效果。這就是為什么?小步快走?以及?設(shè)立正確預(yù)期有助于控制這種風(fēng)險。

另一個我們必須格外小心的地方是產(chǎn)生第一個新詞的前向過程 [譯者注: 第一個新詞?past_key_values?為?None?] 和產(chǎn)生后續(xù)新詞的前向過程 [譯者注: 此時?past_key_values?不為空] 是不一樣的。如果你只針對第一個詞優(yōu)化，你反而會拖慢后續(xù)的那些更重要并且占大部分運(yùn)行時間的詞的生成時間。

另一個很常見的罪魁禍?zhǔn)资菧y量時間，它測量的是 CPU 時間，而不是實際的 CUDA 時間，因此運(yùn)行時需要用?torch.cuda.synchronize()?來確保 GPU 執(zhí)行完成。

定制核函數(shù)

到目前為止，我們已經(jīng)實現(xiàn)了接近 DeepSpeed 的性能，而無需任何自定義代碼！很簡約。我們也不必在推理 batch size 的靈活性上做出任何妥協(xié)！

但根據(jù) DeepSpeed 的經(jīng)驗，我們也想嘗試編寫一個自定義核函數(shù)，以對?torch.jit.script?無法完成融合的一些操作進(jìn)行融合。主要就是下面兩行:

第一個?masked_fill_?是創(chuàng)建一個新的張量，這里只是告訴 softmax 運(yùn)算符忽略這些值。此外，softmax 需要在 float32 上計算 (為了數(shù)值穩(wěn)定性)，但在自定義核函數(shù)中，我們可以減少向上數(shù)據(jù)類型轉(zhuǎn)換的次數(shù)，僅在求和及累加時轉(zhuǎn)換。

你可以在?此處?找到我們的代碼。請記住，我們的優(yōu)化只針對一個特定的 GPU 架構(gòu) (即 A100)，所以該核函數(shù)不適用于其他 GPU 架構(gòu); 同時我們也不是編寫核函數(shù)的專家，因此很有可能有更好的實現(xiàn)方法。

這個自定義核函數(shù)又提供了 10% 的延遲提升，延遲從 81 毫秒/詞元降低到 71 毫秒/詞元。同時，我們繼續(xù)保持了靈活性。

在那之后，我們調(diào)查、探索了更多優(yōu)化手段，比如融合更多的算子來刪除剩下的?reshape?等等。但還沒有哪個手段能產(chǎn)生足夠大的提升而值得被放入最終版本。

Web 服務(wù)部分

就像我們在 Rust 里做的一樣，我們必須實現(xiàn)對具有不同參數(shù)的請求的批處理。由于我們處于?PyTorch?世界中，我們幾乎可以完全控制正在發(fā)生的事情。而又由于我們處于?Python?世界中，我們有一個限制因素，即?torch.distributed?需要多進(jìn)程而不是多線程運(yùn)行，這意味著進(jìn)程之間的通信有點麻煩。最后，我們選擇通過 Redis 發(fā)布/訂閱來傳遞原始字符串，以便同時將請求分發(fā)給所有進(jìn)程。因為我們處于不同的進(jìn)程中，所以這樣做比進(jìn)行張量通信更容易、通信量也很小。

然后我們不得不放棄使用 generate 函數(shù)，因為這會將參數(shù)應(yīng)用于 batch 中所有的序列，而實際上每個序列的參數(shù)可能各不相同。值得慶幸的是，我們可以重用較底層的 API ，如 LogitsProcessor，以節(jié)省大量工作。因此，我們重構(gòu)了一個?generate?函數(shù)，它接受一個參數(shù)列表并將列表中的參數(shù)分別應(yīng)用于 batch 中的各個序列。

最終用戶體驗主要還是看延遲。由于我們支持不同的請求有不同的參數(shù)，因此可能出現(xiàn)這樣的情況: 一個請求想要生成 20 個詞元，而另一個請求想要生成 250 個詞元。由于每個詞元需要 75 毫秒的延遲，因此一個請求需要 1.5 秒，而另一個需要 18 秒。如果我們一直進(jìn)行批處理的話，我們會讓第一個用戶等待 18 秒，因此看起來好像我們正在以 900 毫秒/詞元的速度運(yùn)行，太慢了！

由于我們處于具有極大靈活性的 PyTorch 世界中，我們可以做的是在生成前 20 個詞元后立即從批處理中提取第一個請求，并在 1.5 秒內(nèi)返回給該用戶！這同時也節(jié)省了 230 個詞元的計算量。

因此，靈活性對于獲得最佳延遲非常重要。

最后的筆記和瘋狂的想法

優(yōu)化是一項永無止境的工作，與任何其他項目一樣，20% 的工作通常會產(chǎn)生 80% 的結(jié)果。

從某個時間點開始，我們開始制定一個小的測試策略來確定我們的某個想法的潛在收益，如果測試沒有產(chǎn)生顯著的結(jié)果，我們就會放棄這個想法。1 天增加 10% 足夠有價值，2 周增加 10 倍也足夠有價值。2 周提高 10% 就算了吧。

你試過……嗎？

由于各種原因，有些方法我們知道但我們沒使用的。可能原因有: 感覺它不適合我們的場景、工作量太大、收益潛力不夠大、或者甚至僅僅是因為我們有太多的選擇要試而時間不夠所以就放棄了一些。以下排名不分先后:

• CUDA graphs

• nvFuser?(它是?torch.jit.script?的后端，所以從這個角度來講，我們也算用了它。)

• FasterTransformer

• Nvidia’s Triton

• XLA?(JAX 也使用 XLA!)

• torch.fx

• TensorRT

如果你最喜歡的工具沒有列在這兒，或者你認(rèn)為我們錯過了一些可能有用的重要工具，請隨時與我們聯(lián)系！

Flash attention

我們簡單集成過 flash attention，雖然它在生成第一個詞元 (沒有?past_key_values) 時表現(xiàn)非常好，但在有了?past_key_values?后，它并沒有產(chǎn)生太大的改進(jìn)。而且如果我們要用上它，我們需要對其進(jìn)行調(diào)整以支持?alibi?張量的計算。因此我們決定暫時不做這項工作。

OpenAI Triton

Triton 是一個用于在 Python 中構(gòu)建定制核函數(shù)的出色框架。我們后面打算多用它，但到目前為止我們還沒有。我們很想知道它的性能是否優(yōu)于我們手寫的 CUDA 核函數(shù)。當(dāng)時，在做方案選擇時，我們認(rèn)為直接用 CUDA 編寫似乎是實現(xiàn)目標(biāo)的最短路徑。

填充和?reshape

正如本文通篇所提到的，每次張量拷貝都有成本，而生產(chǎn)環(huán)境中運(yùn)行時的另一個隱藏成本是填充。當(dāng)兩個查詢的長度不同時，你必須使用填充 (使用虛擬標(biāo)記) 以使它們等長。這可能會導(dǎo)致很多不必要的計算。更多信息。

理想情況下，我們可以永遠(yuǎn)?不?做這些計算，永遠(yuǎn)不做?reshape。TensorFlow 有 RaggedTensor 而 PyTorch 也有?嵌套張量?的概念。這兩者似乎都不像常規(guī)張量那樣精簡，但能使我們的計算更好，這對我們有好處。理想的情況下，整個推理過程都可以用 CUDA 或純 GPU 代碼來實現(xiàn)?？紤]到我們在融合算子時看到性能改進(jìn)，這種方法看起來很誘人。但我們不知道性能提升能到什么程度。如果有更聰明的 GPU 專家知道，我們洗耳恭聽！

致謝

所有這些工作都是許多 HF 團(tuán)隊成員合作的結(jié)果。以下排名不分先后，?@ThomasWang @stas@Nouamane @Suraj@Sanchit @Patrick@Younes @Sylvain@Jeff (Microsoft)?@Reza以及 BigScience 項目中的所有人。

您可以通過打開下方英文原文鏈接查看文段中提到的鏈接。

英文原文:?https://hf.co/blog/bloom-inference-optimization

作者: Nicolas Patry

譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師，工作方向為 transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。

排版/審校: zhongdongy (阿東)