優(yōu)化故事: BLOOM 模型推理

@(Engineering Practice)

經(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ù)踩過的坑吧叫挟！

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

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

在此過程中奋献，不同的人參與了不同的階段，嘗試了各種不同的優(yōu)化手段旺上，我們無法一一羅列瓶蚂，還請多多包涵。如果你發(fā)現(xiàn)本文中某些內(nèi)容可能已過時(shí)甚至完全錯(cuò)誤宣吱，這也不奇怪窃这，因?yàn)橐环矫鎸τ谌绾蝺?yōu)化超大模型性能我們?nèi)栽谂W(xué)習(xí)中，另一方面征候，市面上新硬件功能和新優(yōu)化技巧也層出不窮杭攻。

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

訓(xùn)練 BLOOM

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

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

移植至 transformers

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

有幾點(diǎn)需要注意仪召，我們后面還會提到：

小版的模型，如 bigscience/bigscience-small-testing 和 bigscience/bloom-560m 非常重要松蒜。因?yàn)槟Ｐ徒Y(jié)構(gòu)與大版的一樣但尺寸更小扔茅，所以在它們上面一切工作（如調(diào)試、測試等）都更快秸苗。

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

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

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

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

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

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

首次推理（PP + Accelerate）

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

現(xiàn)在我們有了一個(gè)能支持 BLOOM 的 transformers湃交，我們可以開始跑了熟空。

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

然后迈喉，我們轉(zhuǎn)而想使用 transformers 的 pipeline API，吃一下這個(gè) API 的狗糧温圆。然而挨摸，pipeline 不是分布式感知的（這不是它的設(shè)計(jì)目標(biāo)）。

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

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

最后诽偷，在前端架一個(gè)小型 HTTP 服務(wù)器，我們就可以開始提供 BLOOM（大模型）推理服務(wù)了７枥ぁ报慕！

起點(diǎn)

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

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

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

1. 我們可以同時(shí)服務(wù)的用戶數(shù)是多少（吞吐量）脓鹃？
2. 我們平均為每個(gè)用戶服務(wù)的時(shí)間是多少（延遲）逸尖？

我們用 locust 做了一個(gè)測試腳本，如下：

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(),
                },
            },
        )

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

我們想觀察各種實(shí)現(xiàn)方案部署時(shí)如何爬坡，并確保在熔斷時(shí)適當(dāng)?shù)亟档头?wù)器負(fù)載栋齿。熔斷意味著原本能（快速）響應(yīng)你的請求的服務(wù)不再響應(yīng)你的請求苗胀，因?yàn)橥粫r(shí)間有太多人想要使用它。避免死亡之擁（hug of death） 是極其重要的瓦堵。[譯者注：死亡之擁是一個(gè)互聯(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

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

讓我們算一下，一次前向傳播需要 17 TFlop滤港。A100 的 規(guī)格為單卡 312 TFLOPS廊蜒。這意味著單個(gè) GPU 最多能達(dá)到 17 / 312 = 54毫秒/詞元 的延遲。我們用了 16 個(gè) GPU溅漾，因此可得 3毫秒/詞元山叮。這只是個(gè)上限，我們永遠(yuǎn)不可能達(dá)到這個(gè)值添履，況且現(xiàn)實(shí)中卡的性能很少能達(dá)到其規(guī)格所宣稱的數(shù)字屁倔。此外，如果你的模型并不受限于計(jì)算[譯者注：如受限于內(nèi)存帶寬缝龄、受限于 IO 帶寬等]汰现，那么這個(gè)值你也達(dá)不到。知道理想值叔壤，只是為了讓我們對優(yōu)化目標(biāo)心里有個(gè)數(shù)瞎饲。在這里，我們到目前為止與理想值差 2 個(gè)數(shù)量級炼绘。此外嗅战，這個(gè)估計(jì)假設(shè)你將所有算力都用于延遲型服務(wù)，這意味著一次只能執(zhí)行一個(gè)請求（沒關(guān)系俺亮，因?yàn)槟阏谧畲蠡愕臋C(jī)器利用率驮捍，所以沒有太多其他事情要做；但另一個(gè)思路是脚曾，我們可以犧牲一點(diǎn)延遲东且，通過批處理方式來獲得更高的吞吐量）。

探索多條路線

注意：這里本讥，張量并行（Tensor Parallelism珊泳，TP) 意味著每個(gè) GPU 將擁有部分權(quán)重，因此所有 GPU 始終處于工作狀態(tài)拷沸，專注于分給它的部分工作色查。通常這會帶來非常輕微的開銷，因?yàn)闀幸恍┕ぷ魇侵貜?fù)的撞芍，更重要的是秧了，GPU必須定期相互通信交流它們的結(jié)果，然后再繼續(xù)計(jì)算序无。

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

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

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

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

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

• 并行方案的選擇更加容易仇轻。因此 TP 的測試會更方便京痢，這是 JAX 的設(shè)計(jì)帶來的好處之一。
• 對硬件的限制更多篷店，JAX 上 TPU 的性能可能比 GPU 更好祭椰，但 TPU 比 GPU 更難獲取（只在 GCP 上有疲陕，數(shù)量也沒有 GPU 多）方淤。
• 缺點(diǎn)：需要移植工作。但無論如何蹄殃，把它集成到我們的庫里面這件事肯定是受歡迎的携茂。

結(jié)果：

• 移植比較麻煩，因?yàn)槟承l件語句和核函數(shù)很難準(zhǔn)確復(fù)制窃爷，但尚可勉力為之邑蒋。
• 一旦移植完后，測試各種并行方案就比較方便按厘。感謝 JAX医吊，沒有食言。
• 事實(shí)證明逮京，在 Ray 集群里與 TPU worker 通信對我們來講真的太痛苦了卿堂。
  不知道是工具原因還是網(wǎng)絡(luò)的原因，或者僅僅是因?yàn)槲覀儾惶撩蓿@事實(shí)上減慢了我們的實(shí)驗(yàn)速度草描，而且需要的工作比我們預(yù)期的要多得多。
  我們啟動一個(gè)需要 5 分鐘時(shí)間運(yùn)行的實(shí)驗(yàn)策严，等了 5 分鐘沒有發(fā)生任何事情穗慕，10 分鐘之后仍然沒有任何事情發(fā)生，結(jié)果發(fā)現(xiàn)是一些 TPU worker 宕機(jī)了或者是沒有響應(yīng)妻导。我們不得不手動登進(jìn)去看逛绵，弄清楚發(fā)生了什么，修復(fù)它倔韭，重啟一些東西术浪，最后再重新啟動實(shí)驗(yàn)，就這樣半小時(shí)過去了寿酌。幾次下來胰苏，幾天就沒了。我們再強(qiáng)調(diào)一下醇疼，這未必真的是我們使用的工具的問題硕并，但我們的主觀體驗(yàn)確實(shí)如此法焰。
• 無法控制編譯
  我們運(yùn)行起來后，就嘗試了幾種設(shè)置倔毙，想找出最適合我們心目中想要的推理性能的設(shè)置壶栋，結(jié)果證明很難從這些實(shí)驗(yàn)中推測出延遲/吞吐量的規(guī)律。例如普监，在 batch_size=1 時(shí)吞吐量有 0.3 RPS（Requests Per Second, RPS）（此時(shí)每個(gè)請求/用戶都是獨(dú)立的），延遲為 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，每個(gè) batch_size 之間的延遲都差不多幌氮。
  因此缭受，我們可以以 5 倍的延遲為代價(jià)獲得 16 倍的吞吐量「没ィ看上去挺不錯(cuò)的米者，但我們更希望對延遲有更細(xì)粒度的控制，從而使得延遲能滿足 100ms, 1s, 10s, 1mn 規(guī)則中的各檔宇智。

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

• 它們應(yīng)該能處理大部分優(yōu)化工作
• 缺點(diǎn)：通常需要手動處理并行性

結(jié)果：

• 事實(shí)證明蔓搞，為了能夠 trace/jit/export 模型，我們需要重寫 PyTorch 相關(guān)的一部分代碼随橘，使其能夠很容易與純 PyTorch 方法相融合喂分。總體來講机蔗，我們發(fā)現(xiàn)我們可以通過留在 PyTorch 中獲得我們想要的大部分優(yōu)化蒲祈，使我們能夠保持靈活性而無需進(jìn)行太多編碼工作。另一件值得注意的事情是蜒车，因?yàn)槲覀冊?GPU 上運(yùn)行讳嘱，而文本生成有很多輪前向過程，所以我們需要張量留在 GPU 上酿愧，有時(shí)很難將你的張量輸給某個(gè)庫沥潭，返回結(jié)果，計(jì)算 logits（如 argmax 或采樣）嬉挡，再回輸給那個(gè)庫钝鸽。
  將循環(huán)放在外部庫里面意味著像 JAX 一樣失去靈活性汇恤，這不是我們設(shè)想的推理服務(wù)應(yīng)用場景的使用方法。

DeepSpeed

• 這是我們訓(xùn)練 BLOOM 時(shí)使用的技術(shù)拔恰，所以用它來推理也很公平
• 缺點(diǎn)：DeepSpeed 之前從未用于推理因谎，其設(shè)計(jì)也沒準(zhǔn)備用于推理

結(jié)果：

• 我們很快就得到了很不錯(cuò)的結(jié)果，這個(gè)結(jié)果與我們現(xiàn)行方案的上一版性能大致相同颜懊。
• 我們必須想出一種方法财岔，在多進(jìn)程上架設(shè)用于處理并發(fā)請求網(wǎng)絡(luò)服務(wù)，因?yàn)楝F(xiàn)在一個(gè)推理任務(wù)是由多個(gè) DeepSpeed 進(jìn)程完成的（每個(gè) GPU 一個(gè)進(jìn)程）河爹，匠璧。有一個(gè)優(yōu)秀的庫 Mii 可供使用，它雖然還達(dá)不到我們所設(shè)想的極致靈活的目標(biāo)咸这，但我們現(xiàn)在可以在它之上開始我們的工作夷恍。（當(dāng)前的解決方案稍后討論）。
• 我們在使用 DeepSpeed 時(shí)遇到的最大問題是缺乏穩(wěn)定性媳维。
  我們在 CUDA 11.4 上運(yùn)行基于 11.6 編譯的代碼時(shí)遇到了問題酿雪。而其中一個(gè)由來已久的、我們永遠(yuǎn)無法真正解決的問題是：經(jīng)常會發(fā)生核函數(shù)崩潰（CUDA 非法訪問侄刽、尺寸不匹配等）指黎。我們修復(fù)了其中一些問題，但在壓測我們的網(wǎng)絡(luò)服務(wù)時(shí)州丹，我們永遠(yuǎn)無法完全實(shí)現(xiàn)穩(wěn)定性袋励。盡管如此，我想向幫助過我們的 Microsoft 人員說当叭，感謝那些非常愉快的交流茬故，它們提高了我們對正在發(fā)生的事情的理解，并為我們的后續(xù)工作提供了真知灼見蚁鳖。
• 另一個(gè)痛點(diǎn)是我們的團(tuán)隊(duì)主要在歐洲磺芭，而微軟在加利福尼亞，所以合作時(shí)間很棘手醉箕，我們因此損失了大量時(shí)間钾腺。這與技術(shù)部分無關(guān)，但我們確實(shí)認(rèn)識到合作的組織部分也非常重要讥裤。
• 另一件需要注意的事情是放棒，DeepSpeed 依賴于 transformers 來注入其優(yōu)化，并且由于我們一直在更新我們的代碼己英，這使得 DeepSpeed 團(tuán)隊(duì)很難在我們的主分支上工作间螟。很抱歉讓它變得困難，這也可能是 transformers 被稱為技術(shù)最前沿的原因。

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

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

結(jié)果：

• 我們使用綁定庫 tch-rs 在 Rust 中重寫了所有代碼见坑。Rust 的目標(biāo)不是提高性能嚷掠，而是對并行性（線程/進(jìn)程）以及 web 服務(wù)和 PyTorch 的并發(fā)性進(jìn)行更細(xì)粒度的控制。由于 GIL的存在荞驴，Python 很難處理這些底層細(xì)節(jié)叠国。
• 結(jié)果表明，大部分的痛苦來自于移植工作戴尸，移植完后，實(shí)驗(yàn)就輕而易舉了冤狡。我們認(rèn)為孙蒙，通過對循環(huán)進(jìn)行精確的控制，即使在具有大量不同屬性的請求的場景中悲雳，我們也可以為每個(gè)請求提供出色的性能挎峦。如果你感興趣的話晴楔，可以查看代碼顿苇，但這份代碼沒有任何支持税弃，也沒有好的文檔纪岁。
• Rust web 服務(wù)投入生產(chǎn)了幾周西壮，因?yàn)樗鼘Σ⑿行缘闹С指鼘捤刹接疲覀兛梢愿行У厥褂?GPU（如使用 GPU0 處理請求 1尚粘，而 GPU1 處理請求 0）泽铛。在保持延遲不變的情況下，我們把吞吐從 0.3 RPS 提高到了 ~2.5 RPS舀透。雖然在最理想情況下灯荧，我們能將吞吐提高到 16 倍。但實(shí)際工作負(fù)載上的測出來能到 8 倍左右的話也還算不錯(cuò)盐杂。

純 PyTorch

• 純粹修改現(xiàn)有代碼逗载，通過刪除諸如 reshape 之類的操作、使用更優(yōu)化的核函數(shù)等方法來使其運(yùn)行速度更快链烈。
• 缺點(diǎn)：我們必須自己編寫 TP 代碼厉斟，并且我們還有一個(gè)限制，即修改后代碼最好仍然適合我們的庫（至少大部分）强衡。

結(jié)果

• 在下一章詳述擦秽。

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

編寫更高效的 PyTorch

第一件事是在代碼中刪除不必要的操作。可以通過代碼走查并找出明顯可被刪除的某些操作：

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

舊代碼：鏈接
新代碼：鏈接

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

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

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

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

再舉個(gè)例子嘴脾，在刪除大量 reshape/transpose 后男摧，我們在 tensorboard 中發(fā)現(xiàn)：

• 注意力是性能熱點(diǎn)（這是預(yù)期的，但能夠通過測量數(shù)據(jù)來驗(yàn)證總是好的）译打。
• 在注意力中耗拓，由于大量的reshape，很多核函數(shù)其實(shí)是顯存拷貝函數(shù)奏司。
• 我們可以通過修改權(quán)重和 past_key_values 的內(nèi)存布局來移除 reshape乔询。這個(gè)改動有點(diǎn)大，但性能確實(shí)有一定的提高韵洋！

支持 TP

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

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

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

此外千绪，吞吐量上升了很多充易，達(dá)到 10 RPS。 batch_size=1 和 batch_size=32 延遲基本相同荸型，因此盹靴，從這種意義上來講，在相同的延遲下瑞妇，吞吐量的上升基本上是免費(fèi)的稿静。

唾手可得的果實(shí)

現(xiàn)在我們有了一個(gè) TP 版本的實(shí)現(xiàn)，我們可以再次開始進(jìn)行性能剖析和優(yōu)化辕狰。因?yàn)椴⑿蟹桨赴l(fā)生了改變改备，我們有必要再從頭開始分析一遍。

首先蔓倍，同步 (ncclAllReduce) 開始成為主要熱點(diǎn)悬钳，這符合我們的預(yù)期，同步需要花時(shí)間偶翅。但我們不打算優(yōu)化這一部分默勾，因?yàn)樗呀?jīng)使用了 nccl。雖然可能還有一些改進(jìn)空間聚谁，但我們認(rèn)為我們很難做得更好母剥。

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

我們對 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 核函數(shù)融合成了一個(gè)核函數(shù)，從而節(jié)省了核函數(shù)啟動開銷和內(nèi)存拷貝開銷朵耕。

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

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

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

• 你有很多小的、element-wise 的操作
• 你的性能熱點(diǎn)里有一些難以去除的 reshape 算子筹误，這些算子一般就是拷貝
• 算子能融合時(shí)

滑鐵盧

在測試期間桐早，有一段時(shí)間，我們觀察到 Rust 服務(wù)的延遲比 Python 服務(wù)低 25%厨剪。這很奇怪哄酝，但因?yàn)樗鼈兊臏y試環(huán)境是一致的，而且去除了核函數(shù)后我們還是能測到這個(gè)速度增益祷膳，我們開始感覺陶衅，也許降低 Python 開銷可以帶來不錯(cuò)的性能提升。

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

我們遭遇了滑鐵盧门烂，刪除 profiler 代碼后延遲降低了 25%，兩份代碼延遲一樣了兄淫。其實(shí)我們最初也是這么想的屯远，Python 一定不會影響性能，因?yàn)槟Ｐ瓦\(yùn)行時(shí)運(yùn)行的主要還是 torch cpp 的代碼拖叙。雖然 3 天其實(shí)也不算啥氓润，但發(fā)生這樣的事還是挺糟糕的赂乐。

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

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

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

定制核函數(shù)

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

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

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)

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

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

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

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

Web 服務(wù)部分

就像我們在 Rust 里做的一樣，我們必須實(shí)現(xiàn)對具有不同參數(shù)的請求的批處理怠晴。由于我們處于 PyTorch 世界中遥金，我們幾乎可以完全控制正在發(fā)生的事情。
而又由于我們處于 Python 世界中蒜田，我們有一個(gè)限制因素稿械，即 torch.distributed 需要多進(jìn)程而不是多線程運(yùn)行，這意味著進(jìn)程之間的通信有點(diǎn)麻煩冲粤。最后美莫，我們選擇通過 Redis 發(fā)布/訂閱來傳遞原始字符串页眯，以便同時(shí)將請求分發(fā)給所有進(jìn)程。因?yàn)槲覀兲幱诓煌倪M(jìn)程中厢呵，所以這樣做比進(jìn)行張量通信更容易餐茵、通信量也很小。

然后我們不得不放棄使用 generate 函數(shù)述吸，因?yàn)檫@會將參數(shù)應(yīng)用于batch中所有的序列忿族，而實(shí)際上每個(gè)序列的參數(shù)可能各不相同。值得慶幸的是蝌矛，我們可以重用較底層的 API 道批，如 LogitsProcessor，以節(jié)省大量工作入撒。因此隆豹，我們重構(gòu)了一個(gè) generate 函數(shù)，它接受一個(gè)參數(shù)列表并將列表中的參數(shù)分別應(yīng)用于 batch 中的各個(gè)序列茅逮。

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

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

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

最后的筆記和瘋狂的想法

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

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

你試過 ...... 嗎糊啡？

由于各種原因，有些方法我們知道但我們沒使用的吁津∨镄睿可能原因有：感覺它不適合我們的場景、工作量太大碍脏、收益潛力不夠大梭依、或者甚至僅僅是因?yàn)槲覀冇刑嗟倪x擇要試而時(shí)間不夠所以就放棄了一些。以下排名不分先后：

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

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

Flash attention

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

OpenAI Triton

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

填充和 reshape

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

理想情況下坷随，我們可以永遠(yuǎn)不做這些計(jì)算房铭，永遠(yuǎn)不做 reshape驻龟。
TensorFlow 有 RaggedTensor 而 PyTorch 也有嵌套張量 的概念。這兩者似乎都不像常規(guī)張量那樣精簡缸匪，但能使我們的計(jì)算更好翁狐，這對我們有好處。

理想的情況下凌蔬，整個(gè)推理過程都可以用 CUDA 或純 GPU 代碼來實(shí)現(xiàn)露懒。考慮到我們在融合算子時(shí)看到性能改進(jìn)砂心，這種方法看起來很誘人隐锭。但我們不知道性能提升能到什么程度。如果有更聰明的GPU專家知道计贰，我們洗耳恭聽钦睡！

致謝

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

英文原文: <url> https://huggingface.co/blog/bloom-inference-optimization </url>
原文作者：Nicolas Patry
譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師秧秉，工作方向?yàn)?transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理褐桌。