最佳化故事: BLOOM 模型推理

經過“九九八十一難”，大模型終於煉成。下一步就是架設服務，準備開門營業了。真這麼簡單？恐怕未必！行百里者半九十，推理最佳化又是新的雄關漫道。如何進行延遲最佳化？如何進行成本最佳化 (別忘了 OpenAI 8K 上下文的 GPT-4 模型，提示每 1000 詞元只需 0.03 美金，補全每 1000 詞元只需 0.06 美金)？如何在延遲和吞吐量之間折衷？如何處理大模型特有的分散式推理後端和網路服務前端的協作問題……要不動手之前還是先看看 BLOOM 推理服務踩過的坑吧！

本文介紹了我們在實現 BLOOM 模型高效推理服務的過程中發生的幕後故事。

在短短數週內，我們把推理延遲降低了 5 倍 (同時，吞吐量增加了 50 倍)。我們將分享我們為達成這一效能改進而經歷的所有鬥爭和史詩般的勝利。

在此過程中，不同的人參與了不同的階段，嘗試了各種不同的最佳化手段，我們無法一一羅列，還請多多包涵。如果你發現本文中某些內容可能已過時甚至完全錯誤，這也不奇怪，因為一方面對於如何最佳化超大模型效能我們仍在努力學習中，另一方面，市面上新硬體功能和新最佳化技巧也層出不窮。

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

訓練 BLOOM

這是不言而喻的，如果不先獲取到大模型，那推理最佳化就無從談起。大模型訓練是一項由很多不同的人共同領導的超級工程。

為了最大化 GPU 的利用率，我們探索了多種訓練方案。最後，我們選擇了 Megatron-Deepspeed 來訓練最終模型。這意味著訓練程式碼與 transformers 庫並不完全相容。

移植至 transformers

由於上文提及的原因，我們第一件事是將現有模型移植到 transformers 上。我們需要從訓練程式碼中提取相關程式碼並將其實現至 transformers 裡。Younes 負責完成了這項工作。這個工作量絕對不小，我們大概花了將近一個月的時間，進行了 200 次提交 才最終完成。

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

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

首先，你必須放棄那種最終你會得到位元級一致的 logits 結果的幻想。不同的 PyTorch 版本間的運算元核函式更改都會引入細微差別，更不用說不同的硬體可能會因為體系架構不同而產生不同的結果 (而出於成本原因，你可能並不能一直在 A100 GPU 上開發)。

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

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

第一個模型 (small-testing) 和大 BLOOM 一樣，精度是 bfloat16 的。我們原以為兩者應該非常相似，但由於小模型沒有經過太多訓練或者單純只是效能差，最終表現出來的結果是它的輸出波動很大。這意味著我們用它進行生成測試會有問題。第二個模型更穩定，但模型資料精度是 float16 而不是 bfloat16，因此兩者間的誤差空間更大。

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

在此步驟中，我們發現並實現了一個重要的折衷。因為 BLOOM 是在分散式環境中訓練的，所以部分程式碼會對 Linear 層作張量並行，這意味著在單 GPU 上執行相同的操作會得到 不同的數值結果。我們花了一段時間才查明這個問題。這個問題沒辦法徹底解決，要麼我們追求 100% 的數值一致性而犧牲模型執行速度，要麼我們接受每次生成時都會出現一些小的差異但執行速度更快，程式碼更簡單。我們為此設了一個標誌位供使用者自己配置。

首次推理 (PP + Accelerate)

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

現在我們有了一個能支援 BLOOM 的 transformers，我們可以開始跑了。

BLOOM 是一個 352GB (176B bf16 引數) 的模型，我們至少需要那麼多視訊記憶體才能放下它。我們花了一點時間試了試在小視訊記憶體的 GPU 上使用 CPU 解除安裝的方式來推理，但是推理速度慢了幾個數量級，所以我們很快放棄了它。

然後，我們轉而想使用 transformers 的 pipeline API，吃一下這個 API 的狗糧。然而， pipeline 不是分散式感知的 (這不是它的設計目標)。

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

它的工作原理是將 transformer 模型按層進行切分，每個 GPU 分到一些層。真正執行時，是 GPU0 先開始工作，然後將結果交給 GPU1，依次下去。

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

起點

至此，我們甚至還沒有開始討論最佳化！

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

最佳化的第一步是測量效能。在整個最佳化過程中，效能測量貫穿始終。所以，首先需要考慮我們需要測量什麼，也即我們關心的是什麼。對於一個支援多種選項的開放式推理服務而言，使用者會向該服務傳送各種不同的查詢請求，我們關心的是:

1. 我們可以同時服務的使用者數是多少 (吞吐量)？
2. 我們平均為每個使用者服務的時間是多少 (延遲)？

我們用 locust 做了一個測試指令碼，如下:

from locust import HttpUser, between, task
from random import randrange, random

class QuickstartUser(HttpUser):
    wait_time = between(1, 5)

    @task
    def bloom_small(self):
        sentence = "Translate to chinese. EN: I like soup. CN: "
        self.client.post(
            "/generate",
            json={
                "inputs": sentence[: randrange(1, len(sentence))],
                "parameters": {"max_new_tokens": 20, "seed": random()},
            },
        )

    @task
    def bloom_small(self):
        sentence = "Translate to chinese. EN: I like soup. CN: "
        self.client.post(
            "/generate",
            json={
                "inputs": sentence[: randrange(1, len(sentence))],
                "parameters": {
                    "max_new_tokens": 20,
                    "do_sample": True,
                    "top_p": 0.9,
                    "seed": random(),
                },
            },
        )

注意: 這不是我們最佳的也不是唯一的負載測試，但始終是我們第一個執行的負載測試，因此它可用於公平地比較不同方案。在此基準測試表現最好並不意味著它絕對是最好的解決方案。我們還需要使用其他更復雜的測試場景來模擬真實場景的真實效能。

我們想觀察各種實現方案部署時如何爬坡，並確保在熔斷時適當地降低伺服器負載。熔斷意味著原本能 (快速) 響應你的請求的服務不再響應你的請求，因為同一時間有太多人想要使用它。避免 死亡之擁 (hug of death) 是極其重要的。[譯者注: 死亡之擁是一個網際網路領域的隱喻，意指由於極端峰值流量而導致網際網路服務當機]

在上述基準測試中，我們得到的初始效能是 (使用 GCP 上的 16xA100 40G 環境測得，本文後續所有測試都基於該環境):

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

這兩個值並不是很好。在正式開始工作之前，我們可以預估一下我們能得到的最好結果。BLOOM 模型所需的計算量公式為 $24Bsh^2 + 4Bs^2h * 24Bsh^2 + 4Bs^2h$，其中 B 是 batch size， s 是序列長度， h 是隱含層維度。

讓我們算一下，一次前向傳播需要 17 TFlop。A100 的 規格 為單卡 312 TFLOPS。這意味著單個 GPU 最多能達到 17 / 312 = 54 毫秒/詞元 的延遲。我們用了 16 個 GPU，因此可得 3 毫秒/詞元。這只是個上限，我們永遠不可能達到這個值，況且現實中卡的效能很少能達到其規格所宣稱的數字。此外，如果你的模型並不受限於計算 [譯者注: 如受限於記憶體頻寬、受限於 IO 頻寬等]，那麼這個值你也達不到。知道理想值，只是為了讓我們對最佳化目標心裡有個數。在這裡，我們到目前為止與理想值差 2 個數量級。此外，這個估計假設你將所有算力都用於延遲型服務，這意味著一次只能執行一個請求 (沒關係，因為你正在最大化你的機器利用率，所以沒有太多其他事情要做; 但另一個思路是，我們可以犧牲一點延遲，透過批處理方式來獲得更高的吞吐量)。

探索多條路線

注意: 這裡，張量並行 (Tensor Parallelism，TP) 意味著每個 GPU 將擁有部分權重，因此所有 GPU 始終處於工作狀態，專注於分給它的部分工作。通常這會帶來非常輕微的開銷，因為會有一些工作是重複的，更重要的是，GPU 必須定期相互通訊交流它們的結果，然後再繼續計算。

現在我們已經比較清楚地瞭解了我們的處境，是時候開始工作了。

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

每次嘗試都值得寫一篇專門的博文，由於篇幅所限，在這裡我們僅將它們列出來，並只深入解釋並研究那些最終應用到當前服務中去的技術的細節。從流水線並行 (PP) 切換到張量並行 (TP) 是延遲最佳化的一個重要一步。每個 GPU 將擁有部分引數，並且所有 GPU 將同時工作，所以延遲應該會迅速下降。但是付出的代價是通訊開銷，因為它們的中間結果需要經常互相通訊。

需要注意的是，這裡涉及的方法相當廣泛。我們會有意識地學習更多關於每個工具的知識，以及在後續最佳化中如何使用它。

將程式碼移植到 JAX/Flax 中以在 TPU 上執行

• 並行方案的選擇更加容易。因此 TP 的測試會更方便，這是 JAX 的設計帶來的好處之一。
• 對硬體的限制更多，JAX 上 TPU 的效能可能比 GPU 更好，但 TPU 比 GPU 更難獲取 (只在 GCP 上有，數量也沒有 GPU 多)。
• 缺點: 需要移植工作。但無論如何，把它整合到我們的庫裡面這件事肯定是受歡迎的。

結果:

• 移植比較麻煩，因為某些條件語句和核函式很難準確複製，但尚可勉力為之。
• 一旦移植完後，測試各種並行方案就比較方便。感謝 JAX，沒有食言。

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

  不知道是工具原因還是網路的原因，或者僅僅是因為我們不太懂，但這事實上減慢了我們的實驗速度，而且需要的工作比我們預期的要多得多。
  我們啟動一個需要 5 分鐘時間執行的實驗，等了 5 分鐘沒有發生任何事情，10 分鐘之後仍然沒有任何事情發生，結果發現是一些 TPU worker 當機了或者是沒有響應。我們不得不手動登進去看，弄清楚發生了什麼，修復它，重啟一些東西，最後再重新啟動實驗，就這樣半小時過去了。幾次下來，幾天就沒了。我們再強調一下，這未必真的是我們使用的工具的問題，但我們的主觀體驗確實如此。

• 無法控制編譯

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

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

• 它們應該能處理大部分最佳化工作
• 缺點: 通常需要手動處理並行性

結果:

• 事實證明，為了能夠 trace/jit/export 模型，我們需要重寫 PyTorch 相關的一部分程式碼，使其能夠很容易與純 PyTorch 方法相融合。總體來講，我們發現我們可以透過留在 PyTorch 中獲得我們想要的大部分最佳化，使我們能夠保持靈活性而無需進行太多編碼工作。另一件值得注意的事情是，因為我們在 GPU 上執行，而文字生成有很多輪前向過程，所以我們需要張量留在 GPU 上，有時很難將你的張量輸給某個庫，返回結果，計算 logits (如 argmax 或取樣)，再回輸給那個庫。

  將迴圈放在外部庫裡面意味著像 JAX 一樣失去靈活性，這不是我們設想的推理服務應用場景的使用方法。

DeepSpeed

• 這是我們訓練 BLOOM 時使用的技術，所以用它來推理也很公平
• 缺點: DeepSpeed 之前從未用於推理，其設計也沒準備用於推理

結果:

• 我們很快就得到了很不錯的結果，這個結果與我們現行方案的上一版效能大致相同。
• 我們必須想出一種方法，在多程式上架設用於處理併發請求網路服務，因為現在一個推理任務是由多個 DeepSpeed 程式完成的 (每個 GPU 一個程式)，。有一個優秀的庫 Mii 可供使用，它雖然還達不到我們所設想的極致靈活的目標，但我們現在可以在它之上開始我們的工作。(當前的解決方案稍後討論)。

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

  我們在 CUDA 11.4 上執行基於 11.6 編譯的程式碼時遇到了問題。而其中一個由來已久的、我們永遠無法真正解決的問題是: 經常會發生核函式崩潰 (CUDA 非法訪問、尺寸不匹配等)。我們修復了其中一些問題，但在壓測我們的網路服務時，我們永遠無法完全實現穩定性。儘管如此，我想向幫助過我們的 Microsoft 人員說，感謝那些非常愉快的交流，它們提高了我們對正在發生的事情的理解，併為我們的後續工作提供了真知灼見。

• 另一個痛點是我們的團隊主要在歐洲，而微軟在加利福尼亞，所以合作時間很棘手，我們因此損失了大量時間。這與技術部分無關，但我們確實認識到合作的組織部分也非常重要。
• 另一件需要注意的事情是，DeepSpeed 依賴於 transformers 來注入其最佳化，並且由於我們一直在更新我們的程式碼，這使得 DeepSpeed 團隊很難在我們的主分支上工作。很抱歉讓它變得困難，這也可能是 transformers 被稱為技術最前沿的原因。

有關 Web 服務的想法

• 鑑於我們準備執行一個免費服務，支援使用者向該服務傳送長短不一的文字，並要求獲取短至幾個詞，長至如整個食譜那麼長的回應，每個請求的引數也可以各不相同，web 服務需要做點什麼來支援這個需求。

結果:

• 我們使用繫結庫 tch-rs 在 Rust 中重寫了所有程式碼。Rust 的目標不是提高效能，而是對並行性 (執行緒/程式) 以及 web 服務和 PyTorch 的併發性進行更細粒度的控制。由於 GIL 的存在，Python 很難處理這些底層細節。
• 結果表明，大部分的痛苦來自於移植工作，移植完後，實驗就輕而易舉了。我們認為，透過對迴圈進行精確的控制，即使在具有大量不同屬性的請求的場景中，我們也可以為每個請求提供出色的效能。如果你感興趣的話，可以檢視 程式碼，但這份程式碼沒有任何支援，也沒有好的文件。
• Rust web 服務投入生產了幾周，因為它對並行性的支援更寬鬆，我們可以更有效地使用 GPU (如使用 GPU0 處理請求 1，而 GPU1 處理請求 0)。在保持延遲不變的情況下，我們把吞吐從 0.3 RPS 提高到了 ~2.5 RPS。雖然在最理想情況下，我們能將吞吐提高到 16 倍。但實際工作負載上的測出來能到 8 倍左右的話也還算不錯。

純 PyTorch

• 純粹修改現有程式碼，透過刪除諸如 reshape 之類的操作、使用更最佳化的核函式等方法來使其執行速度更快。
• 缺點: 我們必須自己編寫 TP 程式碼，並且我們還有一個限制，即修改後程式碼最好仍然適合我們的庫 (至少大部分)。

結果

• 在下一章詳述。

最終路線: PyTorch + TP + 1 個自定義核心 + torch.jit.script

編寫更高效的 PyTorch

第一件事是在程式碼中刪除不必要的操作。可以透過程式碼走查並找出明顯可被刪除的某些操作:

• Alibi 在 BLOOM 中用於新增位置嵌入 (position embeddings)，原始碼中計算 Alibi 的地方太多，每次都重新計算一次，我們最佳化成只計算一次，這樣效率更高。

舊程式碼: 連結

新程式碼: 連結

這個改動獲得了 10 倍的加速，最新版本還增加了對填充 (padding) 的支援！
由於此步驟僅計算一次，因此在這裡，運算本身實際速度並不重要，而總體上減少操作和張量建立的次數更重要。

當你開始 剖析 程式碼效能時，其他部分會越來越清晰，我們大量地使用了 tensorboard 來幫助我們進行效能剖析。它提供瞭如下圖所示的這類影像，可以提供有關效能的洞見:

<img src="https://devrel.andfun.cn/devrel/posts/2023/04/17/SaHLyP.jpg">

注意力層佔用了很多時間，注意這是一個 CPU 檢視，所以條形很長並不意味著核函式執行時間很長，它只意味著 CPU 正在等待上一步的 GPU 結果。

<img src="https://devrel.andfun.cn/devrel/posts/2023/04/17/hmf54c.jpg">

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

再舉個例子，在刪除大量 reshape / transpose 後，我們在 tensorboard 中發現:

• 注意力是效能熱點 (這是預期的，但能夠透過測量資料來驗證總是好的)。
• 在注意力中，由於大量的 reshape，很多核函式其實是視訊記憶體複製函式。
• 我們 可以 透過修改權重和 past_key_values 的記憶體佈局來移除 reshape。這個改動有點大，但效能確實有一定的提高！

支援 TP

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

然後我們繼續實現一個 TP 版。進度比我們預期的要快得多，一個 (有經驗的) 開發人員僅花了半天時間就實現出來了，程式碼見 此處。在此過程中，我們還重用了一些其他專案的程式碼，這對我們很有幫助。

延遲從 300 毫秒/詞元直接變為 91 毫秒/詞元，這是使用者體驗的巨大改進。
一個簡單的 20 個詞元的請求延遲從 6 秒變成了 2 秒，使用者體驗直接從“慢”變成了輕微延遲。

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

唾手可得的果實

現在我們有了一個 TP 版本的實現，我們可以再次開始進行效能剖析和最佳化。因為並行方案發生了改變，我們有必要再從頭開始分析一遍。

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

第二個是 Gelu 運算元，我們可以看到它啟動了許多 element-wise 類的核函式，總體而言它佔用的計算份額比我們預期的要大。

我們對 Gelu 作了如下修改:

從

def bloom_gelu_forward(x):
    return x * 0.5 *(1.0 + torch.tanh(0.79788456 * x *(1 + 0.044715 * x * x)))

改成了

@torch.jit.script
def bloom_gelu_forward(x):
    return x * 0.5 *(1.0 + torch.tanh(0.79788456 * x *(1 + 0.044715 * x * x)))

我們使用 jit 將許多小的 element-wise 核函式融合成了一個核函式，從而節省了核函式啟動開銷和記憶體複製開銷。

該最佳化降低了 10% 的延遲，從 91 毫秒/詞元到 81 毫秒/詞元，搞定！

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

我們發現它在下面幾個場合有用:

• 你有很多小的、element-wise 的操作
• 你的效能熱點裡有一些難以去除的 reshape 運算元，這些運算元一般就是複製
• 運算元能融合時

滑鐵盧

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

我們開始了為期 3 天的重新實現 torch.distributed 部分程式碼的工作，以便在 Rust 裡執行 nccl-rs。程式碼能工作，但生成的句子與 Python 版有些不一樣，於是我們開始調查這些問題，就在這個過程中，我們發現 …… 在測量 PyTorch 版效能時，我們忘記刪除 PyTorch 裡的 profiler 程式碼了 ……

我們遭遇了滑鐵盧，刪除 profiler 程式碼後延遲降低了 25%，兩份程式碼延遲一樣了。其實我們最初也是這麼想的，Python 一定不會影響效能，因為模型執行時執行的主要還是 torch cpp 的程式碼。雖然 3 天其實也不算啥，但發生這樣的事還是挺糟糕的。

針對錯誤的或不具代表性的測量資料進行最佳化，這很常見，最佳化結果最終會令人失望甚至對整個產品帶來反效果。這就是為什麼 小步快走 以及 設立正確預期有助於控制這種風險。

另一個我們必須格外小心的地方是產生第一個新詞的前向過程 [譯者注: 第一個新詞 past_key_values 為 None ] 和產生後續新詞的前向過程 [譯者注: 此時 past_key_values 不為空] 是不一樣的。如果你只針對第一個詞最佳化，你反而會拖慢後續的那些更重要並且佔大部分執行時間的詞的生成時間。

另一個很常見的罪魁禍首是測量時間，它測量的是 CPU 時間，而不是實際的 CUDA 時間，因此執行時需要用 torch.cuda.synchronize() 來確保 GPU 執行完成。

定製核函式

到目前為止，我們已經實現了接近 DeepSpeed 的效能，而無需任何自定義程式碼！很簡約。我們也不必在推理 batch size 的靈活性上做出任何妥協！

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

attn_weights = attention_scores.masked_fill_(attention_mask, torch.finfo(attention_scores.dtype).min)
attention_probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(input_dtype)

第一個 masked_fill_ 是建立一個新的張量，這裡只是告訴 softmax 運算子忽略這些值。此外，softmax 需要在 float32 上計算 (為了數值穩定性)，但在自定義核函式中，我們可以減少向上資料型別轉換的次數，僅在求和及累加時轉換。

你可以在 此處 找到我們的程式碼。
請記住，我們的最佳化只針對一個特定的 GPU 架構 (即 A100)，所以該核函式不適用於其他 GPU 架構; 同時我們也不是編寫核函式的專家，因此很有可能有更好的實現方法。

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

在那之後，我們調查、探索了更多最佳化手段，比如融合更多的運算元來刪除剩下的 reshape 等等。但還沒有哪個手段能產生足夠大的提升而值得被放入最終版本。

Web 服務部分

就像我們在 Rust 裡做的一樣，我們必須實現對具有不同引數的請求的批處理。由於我們處於 PyTorch 世界中，我們幾乎可以完全控制正在發生的事情。
而又由於我們處於 Python 世界中，我們有一個限制因素，即 torch.distributed 需要多程式而不是多執行緒執行，這意味著程式之間的通訊有點麻煩。最後，我們選擇透過 Redis 釋出/訂閱來傳遞原始字串，以便同時將請求分發給所有程式。因為我們處於不同的程式中，所以這樣做比進行張量通訊更容易、通訊量也很小。

然後我們不得不放棄使用 generate 函式，因為這會將引數應用於 batch 中所有的序列，而實際上每個序列的引數可能各不相同。值得慶幸的是，我們可以重用較底層的 API ，如 LogitsProcessor，以節省大量工作。因此，我們重構了一個 generate 函式，它接受一個引數列表並將列表中的引數分別應用於 batch 中的各個序列。

終端使用者體驗主要還是看延遲。由於我們支援不同的請求有不同的引數，因此可能出現這樣的情況: 一個請求想要生成 20 個詞元，而另一個請求想要生成 250 個詞元。由於每個詞元需要 75 毫秒的延遲，因此一個請求需要 1.5 秒，而另一個需要 18 秒。如果我們一直進行批處理的話，我們會讓第一個使用者等待 18 秒，因此看起來好像我們正在以 900 毫秒/詞元的速度執行，太慢了！

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

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

最後的筆記和瘋狂的想法

最佳化是一項永無止境的工作，與任何其他專案一樣，20% 的工作通常會產生 80% 的結果。

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

你試過……嗎？

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

• CUDA graphs
• nvFuser (它是 torch.jit.script 的後端，所以從這個角度來講，我們也算用了它。)
• FasterTransformer
• Nvidia’s Triton
• XLA (JAX 也使用 XLA!)
• torch.fx
• TensorRT

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

Flash attention

我們簡單整合過 flash attention，雖然它在生成第一個詞元 (沒有 past_key_values) 時表現非常好，但在有了 past_key_values 後，它並沒有產生太大的改進。而且如果我們要用上它，我們需要對其進行調整以支援 alibi 張量的計算。因此我們決定暫時不做這項工作。

OpenAI Triton

Triton 是一個用於在 Python 中構建定製核函式的出色框架。我們後面打算多用它，但到目前為止我們還沒有。我們很想知道它的效能是否優於我們手寫的 CUDA 核函式。當時，在做方案選擇時，我們認為直接用 CUDA 編寫似乎是實現目標的最短路徑。

填充和 reshape

正如本文通篇所提到的，每次張量複製都有成本，而生產環境中執行時的另一個隱藏成本是填充。當兩個查詢的長度不同時，你必須使用填充 (使用虛擬標記) 以使它們等長。這可能會導致很多不必要的計算。更多資訊。

理想情況下，我們可以永遠 不 做這些計算，永遠不做 reshape。
TensorFlow 有 RaggedTensor 而 PyTorch 也有 巢狀張量 的概念。這兩者似乎都不像常規張量那樣精簡，但能使我們的計算更好，這對我們有好處。
理想的情況下，整個推理過程都可以用 CUDA 或純 GPU 程式碼來實現。考慮到我們在融合運算元時看到效能改進，這種方法看起來很誘人。但我們不知道效能提升能到什麼程度。如果有更聰明的 GPU 專家知道，我們洗耳恭聽！

致謝

所有這些工作都是許多 HF 團隊成員合作的結果。以下排名不分先後， @ThomasWang @stas
@Nouamane @Suraj
@Sanchit @Patrick
@Younes @Sylvain
@Jeff (Microsoft) @Reza
以及 BigScience 專案中的所有人。

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

作者: Nicolas Patry

譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

排版/審校: zhongdongy (阿東)