作者:李博杰 ?

這是一個(gè)好問(wèn)題。先說(shuō)結(jié)論，大模型的訓(xùn)練用 4090 是不行的，但推理（inference/serving）用 4090 不僅可行，在性?xún)r(jià)比上還能跟 H100 打個(gè)平手。

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

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

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

把他的成本打開(kāi)，SXM 的成本不會(huì)高于 300$，封裝的 Substrate 及 CoWoS 大約也需要 $300，中間的 Logic Die 最大顆，看上去最高貴 ：） 那是 4nm 的一顆 814mm2 的 Die，TSMC 一張 12 英寸 Wafer 大致上可以制造大約 60 顆這個(gè)尺寸的 Die，Nvidia 在 Partial Good 上一向做得很好（他幾乎不賣(mài) Full Good），所以這 60 顆大致能有 50 顆可用，Nvidia 是大客戶(hù)，從 TSMC 手上拿到的價(jià)格大約是 $15000，所以這個(gè)高貴的 Die 大約只需要 $300。哦，只剩下 HBM 了，當(dāng)前 DRAM 市場(chǎng)疲軟得都快要死掉一家的鬼樣了，即使是 HBM3 大抵都是虧本在賣(mài)，差不多只需要 $15/GB，嗯，80GB 的容量成本是 $1200。TSMC 曾經(jīng)講過(guò)一個(gè)故事。臺(tái)灣同胞辛辛苦苦攢錢(qián)建廠，一張 4nm 那么先進(jìn)的工藝哦，才能賣(mài)到 $15000，但是那某個(gè)客戶(hù)拿去噢，能賣(mài)出 $1500000（$30000*50）的貨啦，機(jī)車(chē)，那樣很討厭耶。你懂我意思嗎？就如最開(kāi)始說(shuō)的，在這個(gè)世界的商業(yè)規(guī)則下，$2000 成本的東西賣(mài) $30000，只有一家，銷(xiāo)售量還很大，這是不符合邏輯的，這種金母雞得有航母才守得住。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

APNet 2023 panel 的主題是 Network, AI, and Foundational Models: Opportunties and Challenges。前面幾個(gè)問(wèn)題都中規(guī)中矩的，panelists 有點(diǎn)放不開(kāi)，我就提了一個(gè)問(wèn)題，網(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ā)言，說(shuō)我在華為肯定能知道所有這些東西；然后傳雄老師也跟了一句，要是做網(wǎng)絡(luò)的懂了太多 AI，那可能他就變成一個(gè) AI guy 了。接著主持人陳凱教授問(wèn)，你們有誰(shuí)真的訓(xùn)練過(guò)大模型？沉默了一會(huì)兒，阿里的兄弟先說(shuō)，我算是半個(gè)訓(xùn)練過(guò)大模型的，我們做的東西是支撐阿里大模型 infra 的。后面又有 panelist 說(shuō)，做 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)化來(lái)說(shuō)都很關(guān)鍵，不然根本不知道瓶頸在哪里。

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

一堆矩陣相乘，簡(jiǎn)單來(lái)想就是左邊若干個(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 之類(lèi)的都不是占算力的主要因素，估算的時(shí)候可以忽略。

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

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

模型的參數(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 需要多少?gòu)埧?/p>

有了模型訓(xùn)練所需的總算力，除以每個(gè) GPU 的理論算力，再除以 GPU 的有效算力利用比例，就得到了所需的 GPU-hours，這塊已經(jīng)有很多開(kāi)源數(shù)據(jù)。LLaMA 2 70B 訓(xùn)練需要 1.7M GPU hours（A100），要是用 1 個(gè) GPU，那得算 200 年。要在一個(gè)月這種比較能接受的時(shí)間周期內(nèi)訓(xùn)練出來(lái)，就得至少有 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），綜合起來(lái) 4090 單卡的訓(xùn)練速度還比 A100 稍低（參考前面 LambdaLabs 的評(píng)測(cè)）。

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

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

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

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

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

Data Parallelism 示意圖，來(lái)源：Colossal AI

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

就模型訓(xùn)練需要多少 GPU 內(nèi)存，我發(fā)現(xiàn)能算清楚的人就不多。有的人甚至以為只需要把模型的參數(shù)和反向傳播的梯度存下來(lái)就夠了。事實(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)來(lái)計(jì)算，否則單純使用 16 位浮點(diǎn)來(lái)計(jì)算的誤差太大，模型容易不收斂。因此，每個(gè)參數(shù)需要存 4 字節(jié)的 32 位版本（正向傳播時(shí)用 16 位版本，優(yōu)化時(shí)用 32 位版本，這叫做 mixed-precision），還需要存 4 字節(jié)的 momentum 和 4 字節(jié)的 variance，一共 12 字節(jié)。如果是用類(lèi)似 SGD 的優(yōu)化器，可以不存 variance，只需要 8 字節(jié)。

正向傳播的中間狀態(tài)（activation）是反向傳播時(shí)計(jì)算梯度必需的，而且跟 batch size 成正比。Batch size 越大，每次讀取模型參數(shù)內(nèi)存能做的計(jì)算就越多，這樣對(duì) GPU 內(nèi)存帶寬的壓力就越小。可是不要忘了，正向傳播的中間狀態(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í)從頭開(kāi)始重算正向傳播的中間狀態(tài)，這樣這層的正向傳播中間狀態(tài)就不用保存了。如果每一層都這么干，那么就只要 2 個(gè)字節(jié)來(lái)存這一層的梯度。但是計(jì)算中間狀態(tài)的算力開(kāi)銷(xiāo)會(huì)很大。因此實(shí)際中一般是把整個(gè) Transformer 分成若干組，一組有若干層，只保存每組第一層的中間狀態(tài)，后面的層就從該組第一層開(kāi)始重新計(jì)算，這樣就平衡了算力和內(nèi)存的開(kāi)銷(xiāo)。

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

當(dāng)然有人說(shuō)，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 示意圖，來(lái)源：Colossal AI

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

Pipeline Parallelism 示意圖，來(lái)源：GPipe

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

首先，正向傳播中間狀態(tài)（activation）的存儲(chǔ)容量會(huì)成倍增加，加劇內(nèi)存容量不足的問(wèn)題。比如流水線的第一級(jí)算出了正向傳播的中間狀態(tài)，如果有 N 個(gè)流水級(jí)，那就要正向流過(guò)后面的 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 流起來(lái)，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 才能填滿(mǎn)流水線。

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

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

中間狀態(tài)存儲(chǔ)的問(wèn)題就已經(jīng)夠大了，再看這 2048 張卡之間的通信開(kāi)銷(xiāo)。按照一張卡放一層，并且用不同的輸入數(shù)據(jù)讓它完全流水起來(lái)的做法，這 2048 張卡分別在計(jì)算自己的 mini-batch，可以認(rèn)為是獨(dú)立參與到 data parallelism 里面了。前面講過(guò)，在數(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 秒。而通過(guò) PCIe Gen4 傳輸這 4 GB 數(shù)據(jù)就已經(jīng)至少需要 0.12 秒了，還需要傳兩遍，也就是先傳梯度，再把平均梯度傳過(guò)來(lái)，這 0.24 秒的時(shí)間相比 0.6 秒來(lái)說(shuō)，是占了比較大的比例。

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

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

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

Tensor parallelism（張量并行）

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

有了張量并行，就可以緩解 GPU 放不下模型導(dǎo)致的流水級(jí)太多的問(wèn)題。分到 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ì)算空間，來(lái)源：DeepSpeed

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

Attention 的計(jì)算過(guò)程，來(lái)源：The Illustrated Transformer

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

每個(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)過(guò)，正向的計(jì)算量基本上就是每個(gè) token 過(guò)一遍所有參數(shù)的計(jì)算量，2 * 3 (Q, K, V) * batch size * token 長(zhǎng)度 * 參數(shù)個(gè)數(shù) = 2 * 3 * batch size * 4096 * 8192 * 8192 / 64。可以跟矩陣的大小對(duì)一下，看看有沒(méi)有算錯(cuò)。

那么通信量是多少呢？輸出矩陣 Z 是由每個(gè) head 拼起來(lái)的，每個(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)超過(guò) PCIe Gen4 x16 的 64 GB/s，就算 NVLink 的 900 GB/s 都撐不住。

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

還是按照 4090 的算力 / 單向通信帶寬 = 330T / (64GB/s / 2) 來(lái)算，計(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，算力就肯定跑不滿(mǎn)理論值。這讓我怎么玩？

但是，如果把 H100 的參數(shù)代入進(jìn)去，馬上就不一樣了。H100 的峰值算力是 1979 Tflops，NVLink 雙向帶寬是 900 GB/s，計(jì)算量和通信量的比例最少需要是 4400，也就是 2 * 3 * (embedding size / 張量并行 GPU 數(shù)量) / bytes per param = 2 * 3 * 8192 / 張量并行 GPU 數(shù)量 / 2 >= 4400，解得：張量并行 GPU 數(shù)量 <= 5.5，也就是單機(jī) 8 卡做張量并行，如果算力跑滿(mǎn)，網(wǎng)絡(luò)會(huì)成為瓶頸。可以看到，即使對(duì)于 900 GB/s 這么快的 NVLink，在巨大的算力面前，都容易出現(xiàn)茶壺里煮餃子倒不出來(lái)的情況。當(dāng)然，采用更優(yōu)的并行切分方式可以節(jié)約一些網(wǎng)絡(luò)通信開(kāi)銷(xiāo)。

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

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

如果在 Feed Forward Network 這里做張量并行，也是可以做類(lèi)似的推導(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ì)買(mǎi) A100/H100/H800 的 PCIe 版嗎？PCIe Gen5 雖然比 Gen 4 快一倍，但對(duì) H100 而言，計(jì)算量和通信量的比例仍然最少需要是 1979T / (128G / 2) = 30000，解出來(lái)張量并行 GPU 數(shù)量 <= 0.8，只要用了張量并行，就是損失算力的！

等到 H100 的下一代出來(lái)了，比如 GH200，算力又翻了一倍，NVLink 還是 900 GB/s，這時(shí)候 NVLink 就也開(kāi)始有點(diǎn)吃力了。所以 GH200 不失時(shí)機(jī)的推出了統(tǒng)一大內(nèi)存，號(hào)稱(chēng) 144 TB，就是為了更好的做換入換出，用內(nèi)存換網(wǎng)絡(luò)通信。如果禁令保持不變，國(guó)內(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ù)中心，還有很多其他問(wèn)題。

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

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

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

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

我們年會(huì)的時(shí)候，海哥講了個(gè)段子，我們找老婆都希望又漂亮，又能掙錢(qián)，還一心一意愛(ài)自己。可同時(shí)滿(mǎn)足這三個(gè)條件的老婆就很難找到了。類(lèi)似的，在分布式系統(tǒng)中，我們都希望性能又高，通用性又強(qiáng)，成本還低。這三個(gè)條件的交集也很小。海哥講到這里，譚博補(bǔ)充了一句，同時(shí)滿(mǎn)足這三個(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é)果是需要存下來(lái)給反向傳播用的。為了節(jié)約內(nèi)存而使用流水線并行，流水級(jí)越多，要存儲(chǔ)的中間狀態(tài)也就更多，反而加劇內(nèi)存的不足。而推理任務(wù)中的各個(gè)輸入數(shù)據(jù)之間并沒(méi)有關(guān)系，正向傳播每一層的中間狀態(tài)也不需要保存下來(lái)，因此流水線并行不需要存儲(chǔ)很多中間狀態(tài)。

首先我們需要計(jì)算一下推理需要多少算力。前面針對(duì)訓(xùn)練算力的估算，為了簡(jiǎn)單起見(jiàn)，忽略了兩個(gè)事情，首先是沒(méi)有考慮 KV Cache，其次是沒(méi)有考慮內(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 矩陣緩存起來(lái)，生成下一個(gè) token 的時(shí)候不再需要重新計(jì)算，這就是所謂的 KV Cache。Q 矩陣每次都不一樣，沒(méi)有緩存的價(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 = 4，在 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ù)的空間將超過(guò)參數(shù)本身占的 140 GB。

KV Cache 能省下來(lái)多少計(jì)算量？每一層計(jì)算 K、V 矩陣一共需要 2 (K, V) * 2 (mult, add) * embedding size * embedding size = 4 * 8192 * 8192 這么多計(jì)算量，乘以之前輸入過(guò)的 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 矩陣的過(guò)程是個(gè)典型的內(nèi)存密集型過(guò)程，它需要加載每一層的 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í)間來(lái)做重復(fù)計(jì)算。這樣，token 的輸出就會(huì)越來(lái)越慢，整個(gè)輸出時(shí)間是輸出長(zhǎng)度平方級(jí)別的，根本沒(méi)法用。

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

接下來(lái)我們就可以計(jì)算推理所需的計(jì)算量了。總的算力很好算，前面講過(guò)，大概就是?2 * 輸出 token 數(shù)量 * 參數(shù)數(shù)量 flops。如果想看細(xì)節(jié)，可以看下面這張圖，來(lái)源是這里。

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

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

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

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

雖然 LLaMA 2 的理論上限是 4096 個(gè) token，但很多輸入 prompt 用不了這么多，因此內(nèi)存訪問(wèn)是有可能成為瓶頸的。此時(shí)，就需要靠 batch size 來(lái)補(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)。

有的人問(wèn)，批量處理所需的算力跟分別單獨(dú)處理所需的算力是一樣的呀，那推理時(shí)為什么需要批量處理？答案就在訪問(wèn)內(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 推理需要多少?gòu)埧ǎ?/p>

總的存儲(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ì)來(lái)說(shuō)是很小的。

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

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

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

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

首先是推理延遲。雖然流水線的不同階段可以塞進(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 秒。這里沒(méi)有考慮 KV Cache 帶來(lái)的節(jié)約。注意，8 張卡是串行處理的，因此每個(gè) token 的時(shí)延還要乘以 8，也就是 0.14 秒。每秒只能輸出 7 個(gè) token，對(duì)于 70B 這么小的模型來(lái)說(shuō)是有點(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)存給吃滿(mǎn)了。

那么要平衡利用 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，就是 1979 / 3.35 = 590。也就是說(shuō)，對(duì) 4090 而言，batch size 小于 330 的時(shí)候 GPU 內(nèi)存帶寬是瓶頸，大于 330 的時(shí)候 GPU 算力是瓶頸。當(dāng) batch size = 330 的時(shí)候，理想情況下，內(nèi)存帶寬和算力恰好都打滿(mǎn)，每張卡處理每個(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)然，還需要考慮到通信庫(kù)本身的開(kāi)銷(xiāo)，加上 4090 不支持 GPU 之間 P2P 傳輸，需要通過(guò) 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ī)來(lái)做流水線并行。我們假設(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 來(lái)說(shuō)可以忽略，不會(huì)顯著影響系統(tǒng)的推理延遲。

當(dāng) batch size 很小時(shí)，流水線推理中的網(wǎng)絡(luò)流量是突發(fā)性（bursty）的，每過(guò) 18 ms 只會(huì)進(jìn)行 0.25 ms 數(shù)據(jù)傳輸，只有 1/72 的占空比，不用擔(dān)心流水線推理把局域網(wǎng)全部給占滿(mǎn)了，搞得沒(méi)法正常上網(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ī)間通信來(lái)做流水線并行，主機(jī)間又沒(méi)有很高的通信帶寬，就勢(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 開(kāi)源項(xiàng)目，就是利用流水線并行，把 GPU 做成了一個(gè)類(lèi)似 BitTorrent 的分布式網(wǎng)絡(luò)。雖然推理延遲確實(shí)比較高，但至少說(shuō)明了分布式 GPU 推理的可行性。

推理用張量并行怎么樣？

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

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

在推理中，如果使用張量并行，Transformer 的每一層都需要傳輸把自己負(fù)責(zé)的結(jié)果向量（大小為 batch size * embedding size / num GPUs）廣播給其他所有 GPU，并接受來(lái)自所有其他 GPU 廣播來(lái)的數(shù)據(jù)。計(jì)算 attention 的時(shí)候需要傳輸一次，計(jì)算 feed-forward network 的時(shí)候又需要傳輸一次，也就是總共需要傳輸 2 * 層數(shù)這么多次。

每次發(fā)送就是 batch size * embedding size（發(fā)送和接收是不同的方向，不能算兩次），對(duì)于 batch size = 1, embedding size = 8192，只需要傳輸 16 KB 數(shù)據(jù)，在 32 GB/s 的 PCIe Gen4 上傳輸只需要 1 us。當(dāng)然，考慮到前面討論的 CPU 中轉(zhuǎn)開(kāi)銷(xiāo)，還是需要大約 30 us 的。一共 160 次傳輸，需要 4.8 ms。

我們?cè)倏紤]計(jì)算的開(kāi)銷(xiāo)。還是考慮 batch size = 1 的情形，GPU 內(nèi)存帶寬是瓶頸，此時(shí)每張卡計(jì)算每個(gè) token 的時(shí)延就是 2 byte * 參數(shù)量 / 卡的數(shù)量 / 內(nèi)存帶寬，代入我們前面的數(shù)值，仍然是 17.5 ms。但是這里 8 張卡是并行處理的，因此總的處理時(shí)長(zhǎng)就是計(jì)算時(shí)間 + 通信時(shí)間 = 17.5 ms + 4.8 ms = 22.3 ms。這就意味著每秒可以生成 45 個(gè) token，這個(gè) token 生成速度已經(jīng)很不錯(cuò)了，至少人類(lèi)的閱讀速度是很難趕上生成的速度了。

如果 batch size 更大會(huì)怎樣？例如 batch size = 330，把 GPU 算力和內(nèi)存帶寬都充分利用起來(lái)，每次需要傳輸?shù)臄?shù)據(jù)量是 330 * 8192 * 2 = 5.4 MB，在 32 GB/s 的 PCIe Gen4 上需要 0.17 ms。一共 160 次傳輸，就是 27 ms。這下網(wǎng)絡(luò)通信開(kāi)銷(xiāo)成了延遲的大頭，總處理時(shí)長(zhǎng)為 27 + 17.5 = 44.5 ms，每秒只能生成 22 個(gè) token 了，但也不算慢。

注意，不管用多少個(gè) GPU 做并行推理，只要用的是張量并行，網(wǎng)絡(luò)傳輸?shù)目倲?shù)據(jù)量是相同的，因此增加 GPU 的數(shù)量只能加速計(jì)算，不能加速通信。

因此，A100/H100 的 NVLink 在降低推理延遲方面還是有很大作用的。如果用 A100/H100，取 batch size = 590 達(dá)到算力和帶寬的平衡利用，這 9.44 MB 數(shù)據(jù)只需要 9.44 MB / 450 GB/s = 0.02 ms。一共 160 次傳輸，也只有 3.2 ms。由于內(nèi)存帶寬大了，計(jì)算時(shí)間也可以大幅縮短，例如 H100 的計(jì)算時(shí)間為 2 * 70G / 8 / 3.35 TB/s = 5.2 ms。總處理時(shí)長(zhǎng)只有 5.2 ms + 3.2 ms = 8.4 ms，每秒可以生成 119 個(gè) token，非常棒！

可以說(shuō)，如果論單個(gè) prompt 的 token 生成速度，無(wú)論用多少塊 4090 也追不上 8 卡 H100。

用 4090 做推理的成本怎么樣？

對(duì)于推理，不管用流水線并行還是張量并行，batch size 不算高到太離譜的情況下內(nèi)存帶寬都是瓶頸。

假如 batch size 能夠高到把算力 100% 利用起來(lái)，并且還能解決 KV Cache 不夠大的問(wèn)題，能解決中間結(jié)果占用內(nèi)存過(guò)多的問(wèn)題，那么這 8 張 4090 可以達(dá)到多少吞吐量？

當(dāng)然，這兩個(gè)問(wèn)題都不好解決，因此推理優(yōu)化才是一個(gè)熱門(mén)的研究領(lǐng)域，存在很多的 trade-off 和奇技淫巧。如果只是用標(biāo)準(zhǔn)的 PyTorch，那推理性能距離把算力 100% 利用起來(lái)還遠(yuǎn)得很吶。

假設(shè)都解決了，在張量并行的通信過(guò)程中我們可以利用 double buffer 做另外一個(gè) batch 的計(jì)算，也就是計(jì)算和通信并行，進(jìn)一步提高吞吐量。通信和計(jì)算分別是 27 ms 和 17.5 ms，傳輸?shù)?27 ms 是瓶頸，也就是每 27 ms 輸出一組 token，一個(gè) batch 330 個(gè) prompt，那這 8 張 4090 真是可以達(dá)到每秒 330 / 0.027 = 12.2K token 的吞吐量。

8 張 4090 的成本是 12800 美金，8 卡 PCIe Gen4 服務(wù)器本身要 2 萬(wàn)美金，加上網(wǎng)絡(luò)設(shè)備，平均每臺(tái) 4 萬(wàn)美金的設(shè)備成本。固定資產(chǎn)按照 3 年攤銷(xiāo)，每小時(shí) 1.52 美元。整機(jī)功耗大約 400W * 8 + 2 kW = 5 kW，按照 0.1 美元一度電算，每小時(shí) 0.5 美元。這 2 美元一小時(shí)的機(jī)器，滿(mǎn)打滿(mǎn)算能生成 12.2K * 3600 = 44M tokens，也就是說(shuō)?1 美元能生成 22M tokens。

是不是比 GPT-3.5 Turbo 的 $0.002 / 1K tokens，也就是 1 美元 0.5M tokens?便宜 44 倍？當(dāng)然，賬不能這么算。

首先，GPU 的算力利用率到不了 100%；

其次，如同所有 SaaS 服務(wù)一樣，用戶(hù)的請(qǐng)求數(shù)量有波峰有波谷，用戶(hù)是按量付費(fèi)的，平臺(tái)提供方可是不管有沒(méi)有人用都在燒錢(qián)的；

此外，每個(gè) batch 中不同 prompt 的長(zhǎng)度和響應(yīng) token 數(shù)量都不同，消耗的算力是 batch 中最大的那個(gè)，但收的錢(qián)是用戶(hù)實(shí)際用的 token 數(shù)；

再次，GPT-3.5 是 175B 的模型，比 70B 的 LLaMA 很可能推理成本更高；

最后，OpenAI 開(kāi)發(fā) GPT-3.5 是燒了不知道多少錢(qián)的，人家至少要賺回訓(xùn)練成本和研發(fā)人員的工資吧。

其實(shí) GPT-3.5 Turbo 的 $0.002 / 1K tokens 真的挺良心的，有的賣(mài) API 的，LLaMA-2 70B 都敢比 GPT-3.5 Turbo 賣(mài)得貴。

如果換成用 H100 做推理，重新算一下這筆賬。一張 H100 至少要 3 萬(wàn)美金，一臺(tái) 8 卡 H100 高配服務(wù)器加上配套的 IB 網(wǎng)絡(luò)，起碼要 30 萬(wàn)美金，同樣按照 3 年攤銷(xiāo)，每小時(shí) 11.4 美元。10 kW 功耗，電費(fèi)每小時(shí) 1 美元。一共 12.4 美元一小時(shí)。

這其實(shí)已經(jīng)是非常良心的價(jià)格了，你在任何云服務(wù)商都不可能租得到這么便宜的 8 卡 H100。所以說(shuō)從云服務(wù)商租卡賣(mài)沒(méi)有護(hù)城河的 SaaS 服務(wù)，比如開(kāi)源模型的推理 API，除非有一種提高推理性能的獨(dú)門(mén)絕技，很難賺得了什么大錢(qián)，二房東的生意不是這么好做的。

再算算這臺(tái) 8 卡 H100 機(jī)器的吞吐量，張量并行也采用傳輸和計(jì)算并行，H100 的通信比較快，因此計(jì)算是瓶頸，每 5.2 ms 可以輸出一組 token，一個(gè) batch 590 個(gè) prompt，滿(mǎn)打滿(mǎn)算可以達(dá)到每秒 590 / 0.0052 = 113K token 的吞吐量。理想情況下，一小時(shí)能生成 407M tokens，也就是 1 美元能生成 33M tokens，H100 這單位 token 的成本比 4090 還要低 30%。

為什么 8 卡 H100 機(jī)器是 4090 機(jī)器價(jià)格的 6 倍，性?xún)r(jià)比卻比 4090 高？因?yàn)橐粡?H100 的算力是 4090 的 6 倍，內(nèi)存帶寬是 4090 的 3.35 倍，當(dāng) batch size 夠大，算力達(dá)到瓶頸的時(shí)候，單卡的性能就是 6 倍。而且，H100 比 4090 的網(wǎng)絡(luò)帶寬強(qiáng)太多了，導(dǎo)致 4090 在張量并行中網(wǎng)絡(luò)通信成了瓶頸，浪費(fèi)了有效算力。因此，同樣的 8 卡機(jī)器吞吐量幾乎可以達(dá)到 4090 的 10 倍。雖然一張 H100 卡的價(jià)格是 4090 的 20 倍以上，但算上服務(wù)器本身的成本和電費(fèi)，整機(jī)的成本只是 6 倍左右。

用最便宜的設(shè)備搞出最高的推理性能

我們發(fā)現(xiàn)在 8 卡 4090 機(jī)器中，3 萬(wàn)美金的設(shè)備成本，GPU 卡只占了 1.28 萬(wàn)美金，不像 H100 機(jī)器那樣 GPU 成本占了大頭。還有辦法進(jìn)一步降低嗎？

如果我們可以忍受 5 token/s 的輸出速度，甚至可以利用流水線并行，用家用臺(tái)式機(jī)和 4090 攢出個(gè)推理集群來(lái)。

遙想我當(dāng)年在 MSRA 的時(shí)候，在一臺(tái)只用 1000 美金攢出來(lái)的機(jī)器上插了 10 塊 FPGA，做出個(gè)世界最快的 Key-Value Store。其實(shí)如果讓我去設(shè)計(jì)一個(gè)性?xún)r(jià)比最高的 4090 推理集群，有很多種方案可以嘗試：

用流水線并行，臺(tái)式機(jī) + 10 Gbps 網(wǎng)卡，足夠在 5 ms 內(nèi)傳輸 batch size = 330 的 5.28 MB 數(shù)據(jù)了，通信 40 ms，計(jì)算 140 ms，達(dá)到 5 token/s 的單 prompt 輸出速度，同時(shí)又能充分利用 4090 的算力。10 Gbps 的網(wǎng)卡和交換機(jī)都很便宜，Intel X710 網(wǎng)卡只要 150 美金，20 口交換機(jī)只要 1500 美金（每 8 個(gè)口 750 美金），一臺(tái)家用臺(tái)式機(jī) 700 美金，這只要 2 萬(wàn)美金就可以搞定原本需要 4 萬(wàn)美金的設(shè)備。

用張量并行，臺(tái)式機(jī) + 200 Gbps ConnectX-6 網(wǎng)卡，上 RoCE，可以把 batch size = 330 的 5.28 MB 數(shù)據(jù)在 0.22 ms 內(nèi)傳完，160 次傳輸是 35 ms，加上計(jì)算的 17.5 ms，一個(gè) token 52.5 ms，可以達(dá)到 19 token/s 的單 prompt 輸出速度，這個(gè)速度已經(jīng)不錯(cuò)了。網(wǎng)卡 1000 美金，200G 交換機(jī) 2 萬(wàn)美金 40 個(gè)端口，平均每 8 個(gè)端口 4000 美金，一臺(tái)家用臺(tái)式機(jī) 700 美金，這只要 3 萬(wàn)美金就能搞定原本 4 萬(wàn)美金的設(shè)備。

主機(jī)內(nèi)用張量并行，主機(jī)間用流水線并行，4 卡 PCIe Gen4 服務(wù)器主板只要 1000 美金而且能跑滿(mǎn) PCIe 帶寬（因?yàn)?8 卡就需要 PCIe switch 了，價(jià)格會(huì)貴很多），兩臺(tái)主機(jī)之間用 25 Gbps 網(wǎng)卡直連，主機(jī)內(nèi)張量并行的時(shí)延是 27 ms，主機(jī)間流水線并行只需 2 次 8 ms 的傳輸（注意 25G 的網(wǎng)絡(luò)帶寬是 4 張 GPU 卡共享的），加上兩次流水線計(jì)算各 17.5 ms，總共 78 ms，可以達(dá)到 13 token/s 的單 prompt 輸出速度。網(wǎng)卡 300 美金 * 2，服務(wù)器 3000 美金 * 2，這只要 1.95 萬(wàn)美金就可以搞定原本需要 4 萬(wàn)美金的設(shè)備。

2 萬(wàn)美金按照 3 年攤銷(xiāo)是每小時(shí) 0.76 美元。按照 0.1 美元/度的電價(jià)，每小時(shí)的電費(fèi)都要 0.5 美元，接近設(shè)備成本了，這有點(diǎn)挖礦的味道了。這 1.26 美元一小時(shí)的機(jī)器如果跑滿(mǎn)了 44M tokens 的吞吐量，1 美元能生成 35M tokens，終于趕上 8 卡 H100 的 33M token per dollar 了。

為什么 H100 以 20 倍于 4090 的 GPU 價(jià)格，9 倍的性能，卻仍然能在系統(tǒng)性?xún)r(jià)比上打個(gè)平手，首先是因?yàn)槟芎某杀靖停? 卡 H100 的功耗是 10 kW，但 9 臺(tái) 8 卡 4090 的功耗是 45 kW；其次是因?yàn)橹鳈C(jī)和網(wǎng)絡(luò)設(shè)備成本更低，一臺(tái) 8 卡 H100 準(zhǔn)系統(tǒng)雖然貴，但只占整機(jī)價(jià)格的 20% 左右；但 4090 因?yàn)榭ǘ啵窍?GPU 礦機(jī)那樣壓成本，只要還是用數(shù)據(jù)中心級(jí)的設(shè)備，準(zhǔn)系統(tǒng)價(jià)格就要占到 35% 以上。

其實(shí)，這個(gè)世界上不止有 A100/H100 和 4090，還有 A10 等計(jì)算卡和 3090 等游戲卡，還有 AMD 的 GPU 和很多其他廠商的 AI 芯片。H100 和 4090 大概率都不是性?xún)r(jià)比的最優(yōu)解，例如 A10 和 AMD GPU 的性?xún)r(jià)比有可能就更高。

我都想搞一個(gè)推理性?xún)r(jià)比挑戰(zhàn)賽，看誰(shuí)能用最便宜的設(shè)備搞出最強(qiáng)的推理吞吐量，同時(shí)延遲不能太高；或者用最便宜的設(shè)備搞出最低的推理延遲，同時(shí)吞吐量不能太低。

這一切都是在假設(shè)使用 LLaMA-2 70B 模型，沒(méi)有做量化壓縮的前提下。如果做了量化壓縮，那性能就更高，甚至在 Unified Memory 夠大的 MacBook Pro 上都能單機(jī)跑了。

License 問(wèn)題怎么辦？

我把這個(gè)問(wèn)題放到最后。NVIDIA Geforce driver 的 License 里寫(xiě)道：

No Datacenter Deployment. The SOFTWARE is not licensed for datacenter deployment, except that blockchain processing in a datacenter is permitted.

既然機(jī)器都是用臺(tái)式機(jī)攢起來(lái)的，這能叫 data center 嗎？還是叫礦場(chǎng)比較合適吧。人家也說(shuō)了，4090 用來(lái)做區(qū)塊鏈?zhǔn)窃试S的。

我有一個(gè)大膽的想法，如果未來(lái)的區(qū)塊鏈不再用挖礦來(lái)做 proof of work，而是用大模型推理來(lái)做 proof of work，這是不是很有意思？每個(gè)人買(mǎi)幾塊顯卡，接到礦池上，既可以自己用來(lái)玩游戲，閑時(shí)又可以貢獻(xiàn)算力。礦池直接就是個(gè)賣(mài)大模型推理 SaaS 服務(wù)的公司，提供前所未有的低價(jià) API。甚至需要大模型推理服務(wù)的人可以在區(qū)塊鏈里自己 P2P 玩起來(lái)，誰(shuí)要用大模型就付點(diǎn) gas。

當(dāng)然，目前的 proof of work 都是計(jì)算很復(fù)雜，驗(yàn)證很簡(jiǎn)單的。如果真用大模型推理做 proof of work，必須防止用戶(hù)隨意編造一個(gè)結(jié)果交上去。當(dāng)然這也是有解決方案的，就像 BitTorrent 和其他一些去中心化網(wǎng)絡(luò)一樣，采用信用機(jī)制，新人只能做驗(yàn)證別人計(jì)算結(jié)果的工作，積攢信用；老人每次算錯(cuò)了，都有比較嚴(yán)厲的懲罰。

從另一個(gè)角度看，家庭局域網(wǎng)絡(luò)的速度也越來(lái)越快，比如我家就自己部署了 10 Gbps 的網(wǎng)絡(luò)。家中的智能設(shè)備越來(lái)越多，算力越來(lái)越強(qiáng)。光纖入戶(hù)也越來(lái)越普遍，小區(qū)和城市的運(yùn)營(yíng)商機(jī)房里部署了越來(lái)越多的邊緣計(jì)算節(jié)點(diǎn)。前面我們用 1 Gbps 的網(wǎng)絡(luò)就足以把多臺(tái)主機(jī)上的 GPU 組成流水線并行，那么在未來(lái)的家庭高速網(wǎng)絡(luò)中，流水線并行甚至張量并行都將成為可能。

大多數(shù)搞 AI 推理的都只關(guān)心數(shù)據(jù)中心，忽略了家中的分布式算力。只要解決了安全、隱私和經(jīng)濟(jì)動(dòng)機(jī)問(wèn)題，我家的 Siri，也許就跑在鄰居家里的 GPU 上。

很多人都在說(shuō)要 democratize AI。但現(xiàn)在大模型平民化的最大障礙就是成本，而成本最大的來(lái)源又是 GPU 市場(chǎng)上計(jì)算卡和游戲卡價(jià)格的剪刀差。這并不是指責(zé)某家公司，其他做 AI 芯片的公司，AI 芯片的算力也并不便宜。畢竟芯片、軟件和生態(tài)的研發(fā)都是白花花的銀子。

就像本文開(kāi)頭提到的微軟給每臺(tái)服務(wù)器部署 FPGA 一樣，大規(guī)模量產(chǎn)的芯片價(jià)格就像沙子一樣。到時(shí)候，能限制大模型推理算力的就只有能源了，就像區(qū)塊鏈挖礦和通用 CPU 的云計(jì)算一樣，都在找最便宜的電力供應(yīng)。我在之前的一個(gè)采訪中就表示，長(zhǎng)期來(lái)看，能源和材料可能是制約大模型發(fā)展的關(guān)鍵。讓我們期待廉價(jià)的大模型走進(jìn)千家萬(wàn)戶(hù)，真正改變?nèi)藗兊纳睢?/p>

編輯：黃飛

?