大模型訓(xùn)練：為何A100更勝4090一籌？

作者：李博杰

這是一個(gè)好問題。先說結(jié)論，大模型的訓(xùn)練用 4090 是不行的，但推理（inference/serving）用 4090 不僅可行，在性價(jià)比上還能比 H100 稍高。4090 如果極致優(yōu)化，性價(jià)比甚至可以達(dá)到 H100 的 2 倍。

事實(shí)上，H100/A100 和 4090 最大的區(qū)別就在通信和內(nèi)存上，算力差距不大。

?

?

NVIDIA 的算力表里面油水很多，比如 H100 TF16 算力寫的是 1979 Tflops，但那是加了 sparsity（稀疏）的，稠密的算力只有一半；4090 官方宣傳 Tensor Core 算力高達(dá) 1321 Tflops，但那是 int8 的，F(xiàn)P16 直只有 330 Tflops。這篇文章的第一版就是用了錯(cuò)的數(shù)據(jù)，H100 和 4090 的數(shù)據(jù)都用錯(cuò)了，得到的結(jié)論非常離譜。

H100 這個(gè)售價(jià)其實(shí)是有 10 倍以上油水的。

2016 年我在 MSRA 的時(shí)候，見證了微軟給每塊服務(wù)器部署了 FPGA，把 FPGA 打到了沙子的價(jià)格，甚至成為了供應(yīng)商 Altera 被 Intel 收購的重要推手。2017 年我還自己挖過礦，知道什么顯卡最劃算。后來在華為，我也是鯤鵬、昇騰生態(tài)軟件研發(fā)的核心參與者。因此，一個(gè)芯片成本多少，我心里大概是有數(shù)的。

鯤鵬的首席架構(gòu)師夏 Core 有一篇知名文章《談一下英偉達(dá)帝國的破腚》，很好的分析了 H100 的成本：

把他的成本打開，SXM 的成本不會(huì)高于 300300，中間的 Logic Die 最大顆，看上去最高貴 ：） 那是 4nm 的一顆 814mm2 的 Die，TSMC 一張 12 英寸 Wafer 大致上可以制造大約 60 顆這個(gè)尺寸的 Die，Nvidia 在 Partial Good 上一向做得很好（他幾乎不賣 Full Good），所以這 60 顆大致能有 50 顆可用，Nvidia 是大客戶，從 TSMC 手上拿到的價(jià)格大約是?300。哦，只剩下 HBM 了，當(dāng)前 DRAM 市場(chǎng)疲軟得都快要死掉一家的鬼樣了，即使是 HBM3 大抵都是虧本在賣，差不多只需要?1200。TSMC 曾經(jīng)講過一個(gè)故事。臺(tái)灣同胞辛辛苦苦攢錢建廠，一張 4nm 那么先進(jìn)的工藝哦，才能賣到?1500000（2000 成本的東西賣 $30000，只有一家，銷售量還很大，這是不符合邏輯的，這種金母雞得有航母才守得住。

據(jù)說微軟和 OpenAI 包下了 H100 2024 年產(chǎn)能的一半，猜猜他們會(huì)不會(huì)發(fā)揮當(dāng)年跟 Altera 砍價(jià)的傳統(tǒng)藝能？會(huì)真的花 $40,000 * 500,000 = 200 億美金去買卡？

咱們?cè)俜治鱿?4090 的成本，5nm 的 609mm2 Die，大約成本是?10 算，100 吧。封裝和風(fēng)扇這些東西，算它?900，這樣的東西賣 $1600，算是良心價(jià)了，因?yàn)檠邪l(fā)成本也是錢啊，更何況 NVIDIA 的大部分研發(fā)人員可是在世界上程序員平均薪酬最高的硅谷。

可以說，H100 就像是中國一線城市的房子，本身鋼筋水泥不值多少錢，房?jī)r(jià)完全是被供求關(guān)系吹起來的。我在 LA 已經(jīng)住了兩周，公司租的房子使用面積是我北京房子的 4 倍，但售價(jià)只貴了 30%，還帶個(gè)小院，相當(dāng)于單位面積的房?jī)r(jià)是北京的 1/3。我跟本地的老外聊天，他們都很吃驚，你們的平均收入水平比 LA 低這么多，怎么買得起北京的房子的？

問題來了，如果 4090 這么香的話，為啥大家還要爭(zhēng)著買 H100，搞得 H100 都斷貨了？甚至 H100 都要對(duì)華禁售，搞出個(gè) H800 的閹割版？

大模型訓(xùn)練為什么不能用 4090

GPU 訓(xùn)練性能和成本對(duì)比

LambdaLabs 有個(gè)很好的 GPU 單機(jī)訓(xùn)練性能和成本對(duì)比，在此摘錄如下。

首先看吞吐量，看起來沒有什么違和的，在單卡能放下模型的情況下，確實(shí)是 H100 的吞吐量最高，達(dá)到 4090 的兩倍?？此懔蛢?nèi)存也能看出來，H100 的 FP16 算力大約是 4090 的 3 倍，內(nèi)存帶寬是 3.35 倍，訓(xùn)練過程中由于 batch size 比較大，大多數(shù)算子是 compute bound（計(jì)算密集型），少數(shù)算子是 memory bound（內(nèi)存密集型），這個(gè)結(jié)果是不意外的。

LambdaLabs PyTorch 單卡訓(xùn)練吞吐量對(duì)比圖

LambdaLabs PyTorch 單卡訓(xùn)練吞吐量對(duì)比表

然后看性價(jià)比，就有意思了，原來排在榜首的 H100 現(xiàn)在幾乎墊底了，而且 4090 和 H100 的差距高達(dá)接近 10 倍。這就是因?yàn)?H100 比 4090 貴太多了。

由于 H100 貨源緊張，云廠商的 H100 租用價(jià)格就更黑了，按照標(biāo)價(jià)大約 7 個(gè)月就可以回本。就算大客戶價(jià)能便宜一半，一年半也足夠回本了。

在價(jià)格戰(zhàn)中過慣了苦日子的 IaaS 云服務(wù)商看到這樣的 H100 回本速度，估計(jì)要感嘆，這真是比區(qū)塊鏈挖礦回本還快吶。

LambdaLabs PyTorch 單卡訓(xùn)練單位成本吞吐量對(duì)比圖

LambdaLabs PyTorch 單卡訓(xùn)練單位成本吞吐量對(duì)比表

大模型訓(xùn)練的算力需求

既然 4090 單卡訓(xùn)練的性價(jià)比這么高，為啥不能用來做大模型訓(xùn)練呢？拋開不允許游戲顯卡用于數(shù)據(jù)中心這樣的許可證約束不談，從技術(shù)上講，根本原因是大模型訓(xùn)練需要高性能的通信，但 4090 的通信效率太低。

大模型訓(xùn)練需要多少算力？訓(xùn)練總算力（Flops）= 6 * 模型的參數(shù)量 * 訓(xùn)練數(shù)據(jù)的 token 數(shù)。

我今年初第一次看到有人煞有介事地講這個(gè)公式的時(shí)候，覺得這不是顯然的嗎？又看到 OpenAI 的高級(jí)工程師能拿 90 多萬美金的年薪，頓時(shí)整個(gè)人都不好了，還是 AI 香呀。之前我也面試過一些做 AI 的工程師，包括一些做 AI 系統(tǒng)優(yōu)化的專家，連 Q、K、V 是啥都說不清楚，LLaMA 每個(gè) tensor 的大小也算不出來，就這樣還能拿到 offer。

APNet 2023 panel 的主題是 Network, AI, and Foundational Models: Opportunties and Challenges。前面幾個(gè)問題都中規(guī)中矩的，panelists 有點(diǎn)放不開，我就提了一個(gè)問題，網(wǎng)絡(luò)歷史上的重要成就基本上都基于對(duì)應(yīng)用場(chǎng)景深刻的理解，但我們現(xiàn)在做網(wǎng)絡(luò)的很多都不了解 AI，甚至連每個(gè) tensor 的大小和每個(gè) step 傳輸?shù)臄?shù)據(jù)量都不知道，如何讓 network community 更了解 AI 呢？

這下熱鬧了，臺(tái)下的譚博首先發(fā)言，說我在華為肯定能知道所有這些東西；然后傳雄老師也跟了一句，要是做網(wǎng)絡(luò)的懂了太多 AI，那可能他就變成一個(gè) AI guy 了。接著主持人陳凱教授問，你們有誰真的訓(xùn)練過大模型？沉默了一會(huì)兒，阿里的兄弟先說，我算是半個(gè)訓(xùn)練過大模型的，我們做的東西是支撐阿里大模型 infra 的。后面又有 panelist 說，做 AI 系統(tǒng)的網(wǎng)絡(luò)優(yōu)化是否有必要自己懂 AI 呢，是不是只要會(huì)做 profiling 就行了？

我個(gè)人觀點(diǎn)仍然是，AI 并不難學(xué)，要想做好 AI 系統(tǒng)優(yōu)化，可以不懂 attention 的 softmax 里面為什么要除以 sqrt(d_k)，但不能不會(huì)計(jì)算模型所需的算力、內(nèi)存帶寬、內(nèi)存容量和通信數(shù)據(jù)量。Jeff Dean 就有個(gè)很有名的 Numbers Every Programmer Should Know，數(shù)量級(jí)的估算對(duì)任何系統(tǒng)優(yōu)化來說都很關(guān)鍵，不然根本不知道瓶頸在哪里。

回到大模型訓(xùn)練所需的總算力，其實(shí)很簡(jiǎn)單，6 * 模型的參數(shù)量 * 訓(xùn)練數(shù)據(jù)的 token 數(shù)就是所有訓(xùn)練數(shù)據(jù)過一遍所需的算力。這里的 6 就是每個(gè) token 在模型正向傳播和反向傳播的時(shí)候所需的乘法、加法計(jì)算次數(shù)。

一堆矩陣相乘，簡(jiǎn)單來想就是左邊若干個(gè)神經(jīng)元，右邊若干個(gè)神經(jīng)元，組成一個(gè)完全二分圖。選出其中任意一個(gè)左邊的神經(jīng)元 l 和右邊的神經(jīng)元 r，正向傳播的時(shí)候：

l 把它的輸出乘上 l 和 r 之間的權(quán)重 w，發(fā)給 r；

r 不可能只連一個(gè)神經(jīng)元吧，總要把多個(gè) l 的加到一起，這就是 reduce，需要一次加法。

反向傳播的時(shí)候：

r 把它收到的梯度乘上 l 和 r 之間的權(quán)重 w，發(fā)給 l；

l 也不可能只連一個(gè) r，需要把梯度 reduce 一下，做個(gè)加法；

別忘了權(quán)重 w 需要更新，那就要計(jì)算 w 的梯度，把 r 收到的梯度乘上 l 正向傳播的輸出（activation）；

一個(gè) batch 一般有多個(gè) sample，權(quán)重 w 的更新需要把這些 sample 的梯度加到一起。

一共 3 次乘法，3 次加法，不管 Transformer 多復(fù)雜，矩陣計(jì)算就是這么簡(jiǎn)單，其他的向量計(jì)算、softmax 之類的都不是占算力的主要因素，估算的時(shí)候可以忽略。

想起來我 2019 年剛加入 MindSpore 團(tuán)隊(duì)的時(shí)候，領(lǐng)導(dǎo)讓我開發(fā)一個(gè)正向算子的反向版本，我求導(dǎo)給求錯(cuò)了，搞得算子的計(jì)算結(jié)果總是不對(duì)，還以為是我們的編譯器出 bug 了。當(dāng)發(fā)現(xiàn)求導(dǎo)求錯(cuò)的時(shí)候，領(lǐng)導(dǎo)像以為我沒學(xué)過微積分一樣看著我，確實(shí)我的微積分學(xué)的不好，這也是我從數(shù)學(xué)專業(yè)轉(zhuǎn)到計(jì)算機(jī)專業(yè)的原因之一。

在 MindSpore 的時(shí)候，自動(dòng)微分一共就不到 1000 行代碼，按照微分公式遞歸計(jì)算下去就行了，但自動(dòng)微分作為一個(gè)重要特性被吹了半天，我都感覺不好意思了。

模型的參數(shù)量和訓(xùn)練數(shù)據(jù)的 token 數(shù)之間也有個(gè)比例關(guān)系，這也很容易理解，只要把模型想象成數(shù)據(jù)的壓縮版本就行了，壓縮比總是有極限的。模型的參數(shù)量太小，就吃不下訓(xùn)練數(shù)據(jù)里面所有的知識(shí)；模型的參數(shù)量如果大于訓(xùn)練數(shù)據(jù)的 token 數(shù)，那又浪費(fèi)，還容易導(dǎo)致 over-fitting。

訓(xùn)練 LLaMA-2 70B 需要多少張卡

有了模型訓(xùn)練所需的總算力，除以每個(gè) GPU 的理論算力，再除以 GPU 的有效算力利用比例，就得到了所需的 GPU-hours，這塊已經(jīng)有很多開源數(shù)據(jù)。LLaMA 2 70B 訓(xùn)練需要 1.7M GPU hours（A100），要是用 1 個(gè) GPU，那得算 200 年。要在一個(gè)月這種比較能接受的時(shí)間周期內(nèi)訓(xùn)練出來，就得至少有 2400 塊 A100。

如果用 4090，單卡 FP16 算力是跟 A100 差不多（330 vs 312 Tflops），但是內(nèi)存帶寬比 A100 低一半（1 vs 2 TB/s），內(nèi)存容量更是差好幾倍（24 vs 80 GB），計(jì)算梯度時(shí)需要使用的 TF32 算力也低一半（83 vs 156 Tflops），綜合起來 4090 單卡的訓(xùn)練速度還比 A100 稍低（參考前面 LambdaLabs 的評(píng)測(cè)）。

就按照 2048 塊 4090 算吧，這 2048 塊 4090 之間的通信就成了最大的問題。

為什么？一般有 tensor parallelism、pipeline parallelism、data parallelism 幾種并行方式，分別在模型的層內(nèi)、模型的層間、訓(xùn)練數(shù)據(jù)三個(gè)維度上對(duì) GPU 進(jìn)行劃分。三個(gè)并行度乘起來，就是這個(gè)訓(xùn)練任務(wù)總的 GPU 數(shù)量。

三種并行方式從三個(gè)維度劃分計(jì)算空間的示意圖，來源：DeepSpeed

Data parallelism（數(shù)據(jù)并行）

數(shù)據(jù)并行是最容易想到的并行方式。每個(gè) GPU 分別計(jì)算不同的輸入數(shù)據(jù)，計(jì)算各自的梯度（也就是模型參數(shù)的改變量），再把梯度匯總起來，取個(gè)平均值，廣播給各個(gè) GPU 分別更新。

Data Parallelism 示意圖，來源：Colossal AI

但只用數(shù)據(jù)并行是肯定不行的，因?yàn)橐粔K GPU 放不下整個(gè) LLaMA 70B 模型。

就模型訓(xùn)練需要多少 GPU 內(nèi)存，我發(fā)現(xiàn)能算清楚的人就不多。有的人甚至以為只需要把模型的參數(shù)和反向傳播的梯度存下來就夠了。事實(shí)上，訓(xùn)練需要的內(nèi)存包括模型參數(shù)、反向傳播的梯度、優(yōu)化器所用的內(nèi)存、正向傳播的中間狀態(tài)（activation）。

優(yōu)化器所用的內(nèi)存其實(shí)也很簡(jiǎn)單，如果用最經(jīng)典的 Adam 優(yōu)化器，它需要用 32 位浮點(diǎn)來計(jì)算，否則單純使用 16 位浮點(diǎn)來計(jì)算的誤差太大，模型容易不收斂。因此，每個(gè)參數(shù)需要存 4 字節(jié)的 32 位版本（正向傳播時(shí)用 16 位版本，優(yōu)化時(shí)用 32 位版本，這叫做 mixed-precision），還需要存 4 字節(jié)的 momentum 和 4 字節(jié)的 variance，一共 12 字節(jié)。如果是用類似 SGD 的優(yōu)化器，可以不存 variance，只需要 8 字節(jié)。

正向傳播的中間狀態(tài)（activation）是反向傳播時(shí)計(jì)算梯度必需的，而且跟 batch size 成正比。Batch size 越大，每次讀取模型參數(shù)內(nèi)存能做的計(jì)算就越多，這樣對(duì) GPU 內(nèi)存帶寬的壓力就越小?？墒遣灰耍騻鞑サ闹虚g狀態(tài)數(shù)量是跟 batch size 成正比的，GPU 內(nèi)存容量又會(huì)成為瓶頸。

大家也發(fā)現(xiàn)正向傳播中間狀態(tài)占的內(nèi)存太多了，可以玩一個(gè)用算力換內(nèi)存的把戲，就是不要存儲(chǔ)那么多梯度和每一層的正向傳播的中間狀態(tài)，而是在計(jì)算到某一層的時(shí)候再臨時(shí)從頭開始重算正向傳播的中間狀態(tài)，這樣這層的正向傳播中間狀態(tài)就不用保存了。如果每一層都這么干，那么就只要 2 個(gè)字節(jié)來存這一層的梯度。但是計(jì)算中間狀態(tài)的算力開銷會(huì)很大。因此實(shí)際中一般是把整個(gè) Transformer 分成若干組，一組有若干層，只保存每組第一層的中間狀態(tài)，后面的層就從該組第一層開始重新計(jì)算，這樣就平衡了算力和內(nèi)存的開銷。

如果還是算不清楚，可以讀讀這篇論文：Reducing Activation Recomputation in Large Transformer Models。

當(dāng)然有人說，GPU 內(nèi)存放不下可以換出到 CPU 內(nèi)存，但是就目前的 PCIe 速度，換出到 CPU 內(nèi)存的代價(jià)有時(shí)候還不如在 GPU 內(nèi)存里重算。如果是像 Grace Hopper 那種極高帶寬的統(tǒng)一內(nèi)存，那么換入換出倒是一個(gè)不錯(cuò)的主意，不管訓(xùn)練的正向傳播中間狀態(tài)還是 KV Cache，都有很多優(yōu)化的空間。

Pipeline parallelism（流水線并行）

既然一塊 GPU 放不下，用多塊 GPU 總行了吧？這就是?model parallelism（模型并行），可以大致分為 pipeline parallelism 和 tensor parallelism。

大家最容易想到的并行方式就是?pipeline parallelism，模型不是有很多層嗎，那就分成幾組，每組算連續(xù)的幾層，穿成一條鏈。

Pipeline Parallelism 示意圖，來源：Colossal AI

這樣就有個(gè)問題，一條鏈上只有一個(gè) GPU 在干活，剩下的都在干等。當(dāng)然聰明的你一定也想到了，既然叫 pipeline，那就可以流水線處理，可以把一個(gè) batch 分為若干個(gè) mini-batch，每個(gè) mini-batch 分別計(jì)算。

Pipeline Parallelism 示意圖，來源：GPipe

這可好，是不是把 pipeline 搞的越深越好，每個(gè) GPU 只算一層？

首先，正向傳播中間狀態(tài)（activation）的存儲(chǔ)容量會(huì)成倍增加，加劇內(nèi)存容量不足的問題。比如流水線的第一級(jí)算出了正向傳播的中間狀態(tài)，如果有 N 個(gè)流水級(jí)，那就要正向流過后面的 N - 1 個(gè)流水級(jí)，再等反向傳播 N - 1 個(gè)流水級(jí)，也就是 2N - 2 輪之后才能用到這個(gè)正向傳播的中間狀態(tài)。不要忘了每一輪都會(huì)產(chǎn)生這么多中間狀態(tài)，因此一共是保存了 2N - 1 個(gè)中間狀態(tài)。如果 N 比較大，這個(gè)存儲(chǔ)容量是非?？植赖?。

其次，pipeline 的相鄰流水級(jí)（pipeline stage）之間是要通信的，級(jí)數(shù)越多，通信的總數(shù)據(jù)量和總時(shí)延就越高。

最后，要讓這樣的 pipeline 流起來，batch size 需要等于 Transformer 里面的層數(shù)，一般是幾十，再乘以 data parallelism 的并行數(shù)，batch size 會(huì)很大，影響模型收斂的速度或模型收斂后的精度。

因此，在內(nèi)存容量足夠的情況下，最好還是少劃分一些流水級(jí)。

對(duì)于 LLaMA-2 70B 模型，模型參數(shù)需要 140 GB，反向傳播的梯度需要 140 GB，優(yōu)化器的狀態(tài)（如果用 Adam）需要 840 GB。

正向傳播的中間狀態(tài)跟 batch size 和選擇性重新計(jì)算的配置有關(guān)，我們?cè)谒懔蛢?nèi)存之間取一個(gè)折中，那么正向傳播的中間狀態(tài)需要 token 長(zhǎng)度 * batch size * hidden layer 的神經(jīng)元數(shù)量 * 層數(shù) * (10 + 24/張量并行度) 字節(jié)。假設(shè) batch size = 8，不用張量并行，那么 LLaMA-2 70B 模型的正向傳播中間狀態(tài)需要 4096 * 8 * 8192 * 80 * (10 + 24) byte = 730 GB，是不是很大？

總共需要 140 + 140 + 840 + 730 = 1850 GB，這可比單放模型參數(shù)的 140 GB 大多了。一張 A100/H100 卡也只有 80 GB 內(nèi)存，這就至少要 24 張卡；如果用 4090，一張卡 24 GB 內(nèi)存，就至少需要 78 張卡。

LLaMA-2 模型一共就只有 80 層，一張卡放一層，是不是正好？這樣就有 80 個(gè)流水級(jí)，單是流水線并行就有 80 個(gè)并行的 batch 才能填滿流水線。

這樣，正向傳播的中間狀態(tài)存儲(chǔ)就會(huì)大到無法忍受，這可是 80 * 2 = 160 輪的中間狀態(tài)，翻了 160 倍。就算是使用選擇性重新計(jì)算，比如把 80 層分成 8 組，每組 10 層，中間狀態(tài)存儲(chǔ)仍然是翻了 16 倍。

除非是用最極端的完全重新計(jì)算，反向傳播到每一層都重新從頭開始計(jì)算正向傳播的中間結(jié)果，但這樣計(jì)算開銷可是隨模型層數(shù)平方級(jí)別的增長(zhǎng)，第 1 層算 1 層，第 2 層算 2 層，一直到第 80 層算 80 層，一共算了 3240 層，計(jì)算開銷可是比正常算一次 80 層翻了 40 倍，這還能忍？

中間狀態(tài)存儲(chǔ)的問題就已經(jīng)夠大了，再看這 2048 張卡之間的通信開銷。按照一張卡放一層，并且用不同的輸入數(shù)據(jù)讓它完全流水起來的做法，這 2048 張卡分別在計(jì)算自己的 mini-batch，可以認(rèn)為是獨(dú)立參與到 data parallelism 里面了。前面講過，在數(shù)據(jù)并行中，每一輪需要傳輸?shù)氖撬?jì)算出的梯度和全局平均后的梯度，梯度的數(shù)據(jù)量就等于模型的參數(shù)數(shù)量。

把 70B 模型分成 80 層，每一層大約有 1B 參數(shù)，由于優(yōu)化器用的是 32 bit 浮點(diǎn)數(shù)，這就需要傳輸 4 GB 數(shù)據(jù)。那么一輪計(jì)算需要多久呢？總的計(jì)算量 = batch size * token 數(shù)量 * 6 * 參數(shù)量 = 8 * 4096 * 6 * 1B = 196 Tflops，在 4090 上如果假定算力利用率 100%，只需要 0.6 秒。而通過 PCIe Gen4 傳輸這 4 GB 數(shù)據(jù)就已經(jīng)至少需要 0.12 秒了，還需要傳兩遍，也就是先傳梯度，再把平均梯度傳過來，這 0.24 秒的時(shí)間相比 0.6 秒來說，是占了比較大的比例。

當(dāng)然我們也可以做個(gè)優(yōu)化，讓每個(gè) GPU 在 pipeline parallelism 中處理的 80 組梯度數(shù)據(jù)首先在內(nèi)部做個(gè)聚合，這樣理論上一個(gè) training step 就需要 48 秒，通信占用的時(shí)間不到 1 秒，通信開銷就可以接受了。當(dāng)然，通信占用時(shí)間不到 1 秒的前提是機(jī)器上插了足夠多的網(wǎng)卡，能夠把 PCIe Gen4 的帶寬都通過網(wǎng)絡(luò)吐出去，否則網(wǎng)卡就成了瓶頸。假如一臺(tái)機(jī)器上插了 8 塊 GPU，這基本上需要 8 塊 ConnectX-6 200 Gbps RDMA 網(wǎng)卡才能滿足我們的需求。

最后再看 batch size，整個(gè) 2048 張卡的集群跑起來，每個(gè) GPU 的 mini-batch 我們剛才設(shè)置為 8，那可真是 batch size = 16384，已經(jīng)是大規(guī)模訓(xùn)練中比較大的 batch size 了，如果再大，可能就影響模型的收斂速度或收斂后的精度了。

因此，單純使用流水線并行和數(shù)據(jù)并行訓(xùn)練大模型的最大問題在于流水線并行級(jí)數(shù)過多，導(dǎo)致正向傳播中間狀態(tài)（activation）存儲(chǔ)容量不足。

Tensor parallelism（張量并行）

那就沒辦法了嗎？我們還有最后一招，就是 Tensor parallelism（張量并行）。它也是模型并行的一種，但不像流水線并行那樣是在模型的層間劃分，而是在模型的層內(nèi)劃分，也就是把一層內(nèi)的 attention 計(jì)算和 Feed Forward Network 劃分到多個(gè) GPU 上處理。

有了張量并行，就可以緩解 GPU 放不下模型導(dǎo)致的流水級(jí)太多的問題。分到 80 個(gè) GPU 才能放下的模型，如果用單機(jī) 8 卡張量并行，就只需要?jiǎng)澐?10 個(gè)流水級(jí)。同時(shí)，張量并行還可以降低 batch size，因?yàn)閺埩坎⑿械膸讉€(gè) GPU 是在算同一個(gè)輸入數(shù)據(jù)。

Tensor、Pipeline、Data 三種并行方式從模型層內(nèi)、模型層間、訓(xùn)練數(shù)據(jù)三個(gè)維度上劃分計(jì)算空間，來源：DeepSpeed

Attention 的計(jì)算過程是比較容易并行的，因?yàn)橛卸鄠€(gè) head，用來關(guān)注輸入序列中的不同位置的，那么把這些 head 分別拆開就行了。

Attention 的計(jì)算過程，來源：The Illustrated Transformer

但是我們做任何并行計(jì)算的時(shí)候都不要忘記通信開銷。

每個(gè) head 里面的 Q、K 兩個(gè)矩陣的大小是 batch size * token 長(zhǎng)度 * key 的大小，V 矩陣的大小是 batch size * token 長(zhǎng)度 * value 的大小。key/value 的大小一般等于 embedding size / heads 數(shù)量，例如在 LLaMA-2 70B 中就是 8192 / 64 = 128，矩陣大小是 batch size * 4096 * 8192 / 64（注意，這只是一個(gè) head 的）。而 Q、K、V 參數(shù)矩陣在每個(gè) head 上的大小是 embedding size * embedding size / heads num = 8192 * 8192 / 64。

我們前面推導(dǎo)過，正向的計(jì)算量基本上就是每個(gè) token 過一遍所有參數(shù)的計(jì)算量，2 * 3 (Q, K, V) * batch size * token 長(zhǎng)度 * 參數(shù)個(gè)數(shù) = 2 * 3 * batch size * 4096 * 8192 * 8192 / 64?？梢愿仃嚨拇笮?duì)一下，看看有沒有算錯(cuò)。

那么通信量是多少呢？輸出矩陣 Z 是由每個(gè) head 拼起來的，每個(gè) head 的大小是 batch size * token 長(zhǎng)度 * embedding size / heads num = batch size * 4096 * 8192 / 64。輸入矩陣 X 的大小是 batch size * token 長(zhǎng)度 * embedding size = batch size * 4096 * 8192。注意這里的 X 大小跟所有 heads 合并在一起后的 Z 大小是一致的，而我們?cè)谶@里算的是每個(gè) head 的 Z 大小。這里的單位是參數(shù)數(shù)量，如果按照字節(jié)算，還要乘以每個(gè)參數(shù)的大小。

如果我們采用最極端的方式，每個(gè) head 交給一個(gè) GPU 去算，那么計(jì)算量和通信量的比例是多少？大概是 2 * 3 * embedding size / heads num / bytes per param = 2 * 3 * 8192 / 64 / 2 = 384。代入 4090 的 330 Tflops，如果想讓通信不成為瓶頸，那么通信帶寬至少需要是 330T / 384 = 859 GB/s，發(fā)送接收雙向還得乘以 2，就是 1.7 TB/s。太大了，遠(yuǎn)遠(yuǎn)超過 PCIe Gen4 x16 的 64 GB/s，就算 NVLink 的 900 GB/s 都撐不住。

所以，tensor parallelism 不能切得太細(xì)，每個(gè) GPU 需要多算幾個(gè) heads。如果每個(gè) GPU 多算幾個(gè) attention heads，輸入矩陣 X 就是這些 heads 共享的了，因此輸入矩陣的通信開銷就被多個(gè) heads 平攤了，計(jì)算量和通信量的比例就可以提高。

還是按照 4090 的算力 / 單向通信帶寬 = 330T / (64GB/s / 2) 來算，計(jì)算量和通信量的比例最少需要是 10000，也就是 2 * 3 * (embedding size / 張量并行 GPU 數(shù)量) / bytes per param = 2 * 3 * 8192 / 張量并行 GPU 數(shù)量 / 2 >= 10000，解得：張量并行 GPU 數(shù)量 <= 2.4。也就是告訴你，要是用了張量并行，最多用 2 個(gè) GPU，如果用更多的 GPU，算力就肯定跑不滿理論值。這讓我怎么玩？

但是，如果把 H100 的參數(shù)代入進(jìn)去，馬上就不一樣了。H100 的峰值算力是 989 Tflops，NVLink 雙向帶寬是 900 GB/s，計(jì)算量和通信量的比例最少需要是 1100，也就是 2 * 3 * (embedding size / 張量并行 GPU 數(shù)量) / bytes per param = 2 * 3 * 8192 / 張量并行 GPU 數(shù)量 / 2 >= 1100，解得：張量并行 GPU 數(shù)量 <= 11，也就是單機(jī) 8 卡做張量并行，對(duì)于 embedding size = 8192 的模型，剛剛好，通信不會(huì)成為瓶頸！

閹割版的 H800 相比 H100 卡的就是網(wǎng)絡(luò)帶寬，把網(wǎng)絡(luò)帶寬從 900 GB/s 降到 400 GB/s 了。我們?cè)俅胍淮?，?jì)算量和通信量比例最少需要是 5000，那么張量并行 GPU 數(shù)量 <= 4.8。這樣單機(jī) 8 卡做張量并行，就會(huì)導(dǎo)致網(wǎng)絡(luò)成為瓶頸。當(dāng)然，計(jì)算量 989 Tflops 是理論值，并行切分方式也可以優(yōu)化，因此實(shí)際訓(xùn)練 70B 的模型 8 卡 H800 網(wǎng)絡(luò)不一定真的是瓶頸。這就是 H800 精準(zhǔn)打擊大模型訓(xùn)練，讓張量并行過得不舒服。

Feed Forward Network 的計(jì)算過程，雖然這是 encoder 的，但 decoder 也差不多，來源：Step-by-Step Illustrated Explanations of Transformer

如果在 Feed Forward Network 這里做張量并行，也是可以做類似的推導(dǎo)，在這里就不贅述了。大凡神經(jīng)網(wǎng)絡(luò)里的矩陣乘法，M*N 的矩陣乘上 N*K 的矩陣，總的計(jì)算量是 M*N*K，輸入輸出的總大小是 (M*N + N*K)，多摞幾個(gè)矩陣那也是常數(shù)（就像 Q、K、V），也就是計(jì)算和通信的比例跟矩陣的邊長(zhǎng)（dimension）是一個(gè)量級(jí)的。

這么分析完了，如果你是要做大規(guī)模大模型訓(xùn)練，你還會(huì)買 A100/H100/H800 的 PCIe 版嗎？PCIe Gen5 雖然比 Gen 4 快一倍，但對(duì) H100 而言，計(jì)算量和通信量的比例仍然最少需要是 989T / (128G / 2) = 15000，解出來張量并行 GPU 數(shù)量 <= 1.6，也就是只要用了張量并行，就是損失算力的！

等到 H100 的下一代出來了，比如 GH200，算力又翻了一倍，NVLink 還是 900 GB/s，這時(shí)候 NVLink 就也開始有點(diǎn)吃力了。所以 GH200 不失時(shí)機(jī)的推出了統(tǒng)一大內(nèi)存，號(hào)稱 144 TB，就是為了更好的做換入換出，用內(nèi)存換網(wǎng)絡(luò)通信。如果禁令保持不變，國內(nèi)版本還是卡住 400 GB/s 的通信，那性能差距會(huì)有多大？

上面的推導(dǎo)當(dāng)然都是簡(jiǎn)化的，實(shí)際上可能不會(huì)這么夸張，但數(shù)量級(jí)是差不多的。

訓(xùn)練部分小結(jié)

4090 不容易做大模型訓(xùn)練的原因除了前面分析的內(nèi)存小，通信慢，license 不支持?jǐn)?shù)據(jù)中心，還有很多其他問題。

比如，A100/H100 支持 ECC 顯存容錯(cuò)，據(jù)說 4090 也支持 ECC，但是不知道故障率會(huì)不會(huì)比 A100/H100 更高。不要小看了容錯(cuò)，2048 張卡的集群就算每張卡 1 個(gè)月出一次故障，平均 20 分鐘就會(huì)有一張卡出故障！要是沒有自動(dòng)化的故障恢復(fù)方式，煉丹師就別想睡覺了。

就算是自動(dòng)從上一個(gè) checkpoint 恢復(fù)，這可是要時(shí)間的，如果不考慮丟棄故障 GPU 梯度這種比較暴力的方式，當(dāng)前這個(gè) step 就算是白算了，還要從上一個(gè) checkpoint 加載梯度，一般需要 10 來分鐘的時(shí)間才能搞定。這樣，每 20 分鐘就浪費(fèi) 10 分鐘，這 10 分鐘恢復(fù)過程中可能又有新的卡故障，總的算下來要浪費(fèi)掉一半的有效算力。

因此，保持大規(guī)模訓(xùn)練集群的低故障率是非常重要的，這些 GPU 卡都非常金貴，可不能像挖礦機(jī)房那樣，動(dòng)不動(dòng)就過熱死機(jī)了。

據(jù)說 3090 是支持 NVLink 的，但 4090 就把 NVLink 給砍掉了。更老的卡，甚至還有支持 PCIe P2P 的，現(xiàn)在也都被砍掉了。誰感興趣可以測(cè)一測(cè) 3090 的 NVLink 性能怎么樣，是不是真的能達(dá)到標(biāo)稱的 600 GB/s，如果真的能達(dá)到的話，是否又可以用來做大模型訓(xùn)練了呢。

我們年會(huì)的時(shí)候，海哥講了個(gè)段子，我們找老婆都希望又漂亮，又能掙錢，還一心一意愛自己?？赏瑫r(shí)滿足這三個(gè)條件的老婆就很難找到了。類似的，在分布式系統(tǒng)中，我們都希望性能又高，通用性又強(qiáng)，成本還低。這三個(gè)條件的交集也很小。海哥講到這里，譚博補(bǔ)充了一句，同時(shí)滿足這三個(gè)條件的分布式系統(tǒng)根本就不存在。

Tensor、Pipeline、Data Parallelism 就像是這樣的不可能三角，相互牽制，只要集群規(guī)模夠大，模型結(jié)構(gòu)仍然是 Transformer，就很難逃出內(nèi)存容量和網(wǎng)絡(luò)帶寬的魔爪。

大模型推理為什么 4090 很香

推理和訓(xùn)練有什么區(qū)別？

首先，訓(xùn)練不僅需要存儲(chǔ)模型參數(shù)，還需要存儲(chǔ)梯度、優(yōu)化器狀態(tài)、正向傳播每一層的中間狀態(tài)（activation），后面幾個(gè)比參數(shù)更大，對(duì)模型內(nèi)存的需求量也更大。

其次，訓(xùn)練任務(wù)是一個(gè)整體，流水線并行的正向傳播中間結(jié)果是需要存下來給反向傳播用的。為了節(jié)約內(nèi)存而使用流水線并行，流水級(jí)越多，要存儲(chǔ)的中間狀態(tài)也就更多，反而加劇內(nèi)存的不足。而推理任務(wù)中的各個(gè)輸入數(shù)據(jù)之間并沒有關(guān)系，正向傳播每一層的中間狀態(tài)也不需要保存下來，因此流水線并行不需要存儲(chǔ)很多中間狀態(tài)。

首先我們需要計(jì)算一下推理需要多少算力。前面針對(duì)訓(xùn)練算力的估算，為了簡(jiǎn)單起見，忽略了兩個(gè)事情，首先是沒有考慮 KV Cache，其次是沒有考慮內(nèi)存帶寬。

KV Cache

什么是 KV Cache？對(duì)于每個(gè)輸入的 prompt，在計(jì)算第一個(gè) token 輸出的時(shí)候，每個(gè) token 的 attention 肯定是都要從頭計(jì)算。但是在后續(xù) token 的生成中，都需要計(jì)算 self-attention，也就是輸入 prompt 以及前面輸出的 token 的 attention。這是就需要用到前面每一個(gè) token 的 K 和 V，由于每一層的參數(shù)矩陣是不變的，此時(shí)只有剛生成的那個(gè) token 的 K 和 V 需要從頭計(jì)算，輸入 prompt 和之前生成的 token 的 K 和 V 其實(shí)是跟上一輪一樣的。

這時(shí)，我們就可以把每一層的 K、V 矩陣緩存起來，生成下一個(gè) token 的時(shí)候不再需要重新計(jì)算，這就是所謂的 KV Cache。Q 矩陣每次都不一樣，沒有緩存的價(jià)值。前面講的訓(xùn)練中的選擇性保存正向 activation 是個(gè)拿計(jì)算換內(nèi)存的把戲，這里的 KV Cache 就是一個(gè)拿內(nèi)存換計(jì)算的把戲。

KV Cache 需要多少存儲(chǔ)容量呢？每一層，每個(gè) token 的 K、V 矩陣都是 embedding size 這么大，再乘上 token 數(shù)量和 batch size，就是這一層的 KV Cache 所需的存儲(chǔ)容量了。一定要記住 batch size，在正向和反向傳播的幾乎所有階段，都不會(huì)涉及到對(duì) batch size 中各個(gè) sample 的合并處理，因此它始終是存儲(chǔ)量和計(jì)算量計(jì)算中的一個(gè)系數(shù)。

例如，如果 batch size = 8，在 LLaMA 2 70B 中，假設(shè)輸入和輸出的 token 數(shù)量達(dá)到了模型的極限 4096，80 層的 KV Cache 一共需要 2 (K, V) * 80 * 8192 * 4096 * 8 * 2B = 80 GB。如果 batch size 更大，那么 KV Cache 占據(jù)的空間將超過參數(shù)本身占的 140 GB。

KV Cache 能省下來多少計(jì)算量？每一層計(jì)算 K、V 矩陣一共需要 2 (K, V) * 2 (mult, add) * embedding size * embedding size = 4 * 8192 * 8192 這么多計(jì)算量，乘以之前輸入過的 token 數(shù)量、層數(shù)和 batch size，就是 4096 * 80 * 8 * 4 * 8192 * 8192 = 640 Tflops。相當(dāng)于每存儲(chǔ) 1 個(gè)字節(jié)，節(jié)約了 16K 次計(jì)算，還是很劃算的。

事實(shí)上，KV Cache 節(jié)約的遠(yuǎn)遠(yuǎn)不止這些。計(jì)算 K、V 矩陣的過程是個(gè)典型的內(nèi)存密集型過程，它需要加載每一層的 K、V 參數(shù)矩陣。也就是如果不做任何緩存，假設(shè) prompt 長(zhǎng)度很短而輸出長(zhǎng)度接近 token 的最大長(zhǎng)度 4096，到了最后一個(gè) token 的時(shí)候，單是重復(fù)計(jì)算前面每個(gè) token 的 K、V 矩陣，就需要讀取內(nèi)存 4096 * 80 * 2 * 8192 * 8192 = 40T 次，每次 2 個(gè)字節(jié)，要知道 H100 的內(nèi)存帶寬只有 3.35 TB/s，4090 更是只有 1 TB/s，這單是最后一個(gè) token 就得耗掉一張卡幾十秒的時(shí)間來做重復(fù)計(jì)算。這樣，token 的輸出就會(huì)越來越慢，整個(gè)輸出時(shí)間是輸出長(zhǎng)度平方級(jí)別的，根本沒法用。

推理是計(jì)算密集還是存儲(chǔ)密集

接下來我們就可以計(jì)算推理所需的計(jì)算量了。總的算力很好算，前面講過，大概就是?2 * 輸出 token 數(shù)量 * 參數(shù)數(shù)量 flops。如果想看細(xì)節(jié)，可以看下面這張圖，[來源：https：//le.qun.ch/en/blog/2023/05/13/transformer-batching/]。

Transformer 推理過程中每一步的矩陣形狀、所需算力和內(nèi)存訪問量，來源：Lequn Chen，Dissecting Batching Effects in GPT Inference

但算力并不能說明一切，模型還需要訪問 GPU 內(nèi)存，內(nèi)存帶寬也可能成為瓶頸。至少需要把參數(shù)從內(nèi)存里面讀出來吧？事實(shí)上，內(nèi)存帶寬的估算就這么簡(jiǎn)單，內(nèi)存訪問量 = 參數(shù)數(shù)量 * 2 bytes。中間結(jié)果有一部分是可以放在緩存里面的，緩存放不下的部分也需要占內(nèi)存帶寬，我們先不算。

如果不做任何批量輸入，也就是模型專門服務(wù)一個(gè) prompt，batch size = 1，整個(gè) context 的長(zhǎng)度很短（例如只有 128），那么整個(gè)推理過程中，每載入一個(gè)參數(shù)（2 字節(jié)），就只進(jìn)行 128 次乘法和加法計(jì)算，那么計(jì)算 flops 和訪問內(nèi)存 bytes 的比例就只有 128?；旧先魏?GPU 在這種情況下都會(huì)變成 memory bound，時(shí)間都耗在加載內(nèi)存上了。

對(duì)于 4090 來說，計(jì)算 flops 和內(nèi)存帶寬之比是 330 / 1 = 330；對(duì)于 H100 來說，計(jì)算 flops 和內(nèi)存帶寬之比是 989 / 3.35 = 295。也就是說，如果 context 中的 token 數(shù)量小于 330 或者 295，那么內(nèi)存訪問就會(huì)成為瓶頸。

雖然 LLaMA 2 的理論上限是 4096 個(gè) token，但很多輸入 prompt 用不了這么多，因此內(nèi)存訪問是有可能成為瓶頸的。此時(shí)，就需要靠 batch size 來補(bǔ)足了。推理中的批量處理，就是把幾乎同時(shí)到達(dá)后端服務(wù)的 prompt 放到一起處理。不用擔(dān)心，batch 里面的不同 prompt 的處理是完全獨(dú)立的，不用擔(dān)心會(huì)互相干擾。但這些 prompt 的輸出是步調(diào)整齊劃一的，每一輪整個(gè) batch 中的每個(gè) prompt 都會(huì)輸出一個(gè) token，因此如果有的 prompt 先輸出完了，那就只能等其他的輸出結(jié)束，造成一定的算力浪費(fèi)。

有的人問，批量處理所需的算力跟分別單獨(dú)處理所需的算力是一樣的呀，那推理時(shí)為什么需要批量處理？答案就在訪問內(nèi)存的帶寬上。

如果同時(shí)到達(dá)服務(wù)器的 prompt 很多，是不是 batch size 越大越好？也不是，因?yàn)?KV Cache 的大小可是正比于 batch size 的，batch size 大了，KV Cache 占據(jù)的 GPU 內(nèi)存容量就很可觀，比如在 LLaMA-2 70B 中，每個(gè) prompt 都要占據(jù) 5 GB 的 KV Cache，如果 batch size 搞到 32，那么 KV Cache 就會(huì)占掉 160 GB 的 GPU 內(nèi)存，比參數(shù)都大了。

70B 推理需要多少張卡？

總的存儲(chǔ)容量也很好算，推理的時(shí)候最主要占內(nèi)存的就是參數(shù)、KV Cache 和當(dāng)前層的中間結(jié)果。當(dāng) batch size = 8 時(shí)，中間結(jié)果所需的大小是 batch size * token length * embedding size = 8 * 4096 * 8192 * 2B = 0.5 GB，相對(duì)來說是很小的。

70B 模型的參數(shù)是 140 GB，不管 A100/H100 還是 4090 都是單卡放不下的。那么 2 張 H100 夠嗎？看起來 160 GB 是夠了，但是剩下的 20 GB 如果用來放 KV Cache，要么把 batch size 壓縮一半，要么把 token 最大長(zhǎng)度壓縮一半，聽起來是不太明智。因此，至少需要 3 張 H100。

對(duì)于 4090，140 GB 參數(shù) + 40 GB KV Cache = 180 GB，每張卡 24 GB，8 張卡剛好可以放下。

推理用流水線并行可以嗎？

推理使用流水線并行，最主要的問題是串行處理的推理延遲，網(wǎng)絡(luò)延遲倒是小問題。

首先是推理延遲。雖然流水線的不同階段可以塞進(jìn)不同的 prompt，但同一個(gè) prompt 的處理仍然永遠(yuǎn)在單個(gè) GPU 上輪轉(zhuǎn)，這樣相比 Tensor parallelism 而言，單個(gè) prompt 的延遲就增大了。

對(duì)于很小的 batch size，GPU 內(nèi)存帶寬是瓶頸，此時(shí)每張卡計(jì)算每個(gè) token 的時(shí)延就是 2 byte * 參數(shù)量 / 卡的數(shù)量 / 內(nèi)存帶寬，例如 8 卡 4090 跑 LLaMA-2 70B，就是 2 * 70G / 8 / 1 TB/s = 0.0175 秒。這里沒有考慮 KV Cache 帶來的節(jié)約。注意，8 張卡是串行處理的，因此每個(gè) token 的時(shí)延還要乘以 8，也就是 0.14 秒。每秒只能輸出 7 個(gè) token，對(duì)于 70B 這么小的模型來說是有點(diǎn)慢了。

對(duì)于很大的 batch size，GPU 算力是瓶頸，此時(shí)每張卡計(jì)算每個(gè) token 的時(shí)延就是 batch size * 2 * 參數(shù)量 / 卡的數(shù)量 / 算力，例如 batch size = 1024，同樣的 8 卡例子，就是 1024 * 2 * 70G / 8 / 330 Tflops = 0.0543 秒。事實(shí)上，對(duì)于這么大的 batch size，KV Cache 和正向傳播的中間結(jié)果先把 GPU 內(nèi)存給吃滿了。

那么要平衡利用 GPU 算力和內(nèi)存帶寬，batch size 需要是多少呢？這就是 2 byte * 參數(shù)量 / 卡的數(shù)量 / 內(nèi)存帶寬 = batch size * 2 * 參數(shù)量 / 卡的數(shù)量 / 算力，左右兩邊參數(shù)量和卡的數(shù)量互相抵消，得到 batch size = 算力 / 內(nèi)存帶寬。對(duì)于 4090，就是 330 / 1 = 330；對(duì)于 H100，就是 989 / 3.35 = 295。也就是說，對(duì) 4090 而言，batch size 小于 330 的時(shí)候 GPU 內(nèi)存帶寬是瓶頸，大于 330 的時(shí)候 GPU 算力是瓶頸。當(dāng) batch size = 330 的時(shí)候，理想情況下，內(nèi)存帶寬和算力恰好都打滿，每張卡處理每個(gè) token 的時(shí)間就是 17.5 ms。

其次是網(wǎng)絡(luò)延遲。流水線并行相比張量并行的優(yōu)點(diǎn)就是網(wǎng)絡(luò)傳輸量小，流水級(jí)之間只需要傳輸 batch size * embedding size 這么多數(shù)據(jù)。例如 batch size = 8，embedding size = 8192，只需要傳輸 128 KB 數(shù)據(jù)，在 32 GB/s 的 PCIe Gen4 x16 上，只需要 4 us 就可以傳輸完成。當(dāng)然，還需要考慮到通信庫本身的開銷，加上 4090 不支持 GPU 之間 P2P 傳輸，需要通過 CPU 中轉(zhuǎn)，實(shí)際上需要幾十 us 的時(shí)間，相比計(jì)算部分動(dòng)輒幾十 ms 的時(shí)延，可以忽略不計(jì)。

即使 batch size = 330，這 5.28 MB 數(shù)據(jù)在 PCIe 上也只需要傳輸 0.16 ms，相比計(jì)算部分的 17.5 ms 仍然可以忽略不計(jì)。

如果可以忍受流水線并行的推理延遲，甚至可以用多臺(tái)主機(jī)來做流水線并行。我們假設(shè)主機(jī)間只有 1 Gbps 的普通以太網(wǎng)絡(luò)，每臺(tái)主機(jī)只有一張 4090。對(duì)于 batch size = 1，16 KB 數(shù)據(jù)需要 0.25 ms 才能傳輸完成，再加上 0.25 ms 兩端網(wǎng)絡(luò)協(xié)議棧的處理時(shí)間，每個(gè)流水級(jí)就需要 0.5 ms 的時(shí)延，8 張卡花在通信上的時(shí)間只有 4 ms，相比整體計(jì)算時(shí)延 140 ms 來說可以忽略，不會(huì)顯著影響系統(tǒng)的推理延遲。

當(dāng) batch size 很小時(shí)，流水線推理中的網(wǎng)絡(luò)流量是突發(fā)性（bursty）的，每過 18 ms 只會(huì)進(jìn)行 0.25 ms 數(shù)據(jù)傳輸，只有 1/72 的占空比，不用擔(dān)心流水線推理把局域網(wǎng)全部給占滿了，搞得沒法正常上網(wǎng)了。

如果為了充分利用算力，把 batch size 設(shè)置得很大，比如 330，那么 16 KB * 330 = 5.28 MB 數(shù)據(jù)需要傳輸 41 ms，8 張卡花在通信上的時(shí)間高達(dá) 0.33 秒，這樣就只有 3 token/s 的輸出速度了，難以忍受。因此，如果用主機(jī)間通信來做流水線并行，主機(jī)間又沒有很高的通信帶寬，就勢(shì)必需要犧牲一定的吞吐量。

例如，我們?cè)O(shè)置輸出速度不小于 5 token/s，這時(shí)留給通信的時(shí)間是 60 ms，每個(gè)流水級(jí)至多 7.5 ms，1 Gbps 網(wǎng)絡(luò)可以傳輸 960 KB 數(shù)據(jù)，這時(shí) batch size 至多設(shè)置為 60，也就是這 8 張 4090 的總吞吐量是 2400 token/s。此時(shí)的有效算力利用率只有不到 20%。

最近有一個(gè)比較火的 [Petals 開源項(xiàng)目：https://github.com/bigscience-workshop/petals]，就是利用流水線并行，把 GPU 做成了一個(gè)類似 BitTorrent 的分布式網(wǎng)絡(luò)。雖然推理延遲確實(shí)比較高，但至少說明了分布式 GPU 推理的可行性。

推理用張量并行怎么樣？

前面講到，流水線并行的最大缺點(diǎn)是 GPU 串行處理，延遲較高，導(dǎo)致輸出 token 比較慢。而張量并行的最大缺點(diǎn)是傳輸數(shù)據(jù)量大，網(wǎng)絡(luò)帶寬低的設(shè)備不一定 hold 得住。

但是推理要傳輸?shù)臄?shù)據(jù)量跟訓(xùn)練要傳輸?shù)臄?shù)據(jù)量可不是一回事?。?/strong>推理只需要傳輸正向傳播的中間結(jié)果（activation），而訓(xùn)練還需要傳輸所有參數(shù)的梯度，梯度才是數(shù)據(jù)量的大頭。