有了模型訓(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é)。如果是用類似 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)銷會(huì)很大。因此實(shí)際中一般是把整個(gè) Transformer 分成若干組，一組有若干層，只保存每組第一層的中間狀態(tài)，后面的層就從該組第一層開(kāi)始重新計(jì)算，這樣就平衡了算力和內(nèi)存的開(kāi)銷。

如果還是算不清楚，可以讀讀這篇論文：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 才能填滿流水線。

這樣，正向傳播的中間狀態(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)銷可是隨模型層數(shù)平方級(jí)別的增長(zhǎng)，第 1 層算 1 層，第 2 層算 2 層，一直到第 80 層算 80 層，一共算了 3240 層，計(jì)算開(kāi)銷可是比正常算一次 80 層翻了 40 倍，這還能忍？

中間狀態(tài)存儲(chǔ)的問(wèn)題就已經(jīng)夠大了，再看這 2048 張卡之間的通信開(kāi)銷。按照一張卡放一層，并且用不同的輸入數(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)銷就可以接受了。當(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)卡才能滿足我們的需求。

最后再看 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)銷。

每個(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)銷就被多個(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，算力就肯定跑不滿理論值。這讓我怎么玩？

但是，如果把 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)練，讓張量并行過(guò)得不舒服。

Feed Forward Network 的計(jì)算過(guò)程，雖然這是 encoder 的，但 decoder 也差不多，來(lái)源：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ì)買(mǎi) A100/H100/H800 的 PCIe 版嗎？PCIe Gen5 雖然比 Gen 4 快一倍，但對(duì) H100 而言，計(jì)算量和通信量的比例仍然最少需要是 989T / (128G / 2) = 15000，解出來(lái)張量并行 GPU 數(shù)量 <= 1.6，也就是只要用了張量并行，就是損失算力的！

等到 H100 的下一代出來(lái)了，比如 GH200，算力又翻了一倍，NVLink 還是 900 GB/s，這時(shí)候 NVLink 就也開(kāi)始有點(diǎn)吃力了。所以 GH200 不失時(shí)機(jī)的推出了統(tǒng)一大內(nèi)存，號(hào)稱 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)稱的 600 GB/s，如果真的能達(dá)到的話，是否又可以用來(lái)做大模型訓(xùn)練了呢。

我們年會(huì)的時(shí)候，海哥講了個(gè)段子，我們找老婆都希望又漂亮，又能掙錢(qián)，還一心一意愛(ài)自己。可同時(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é)果是需要存下來(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 = 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ù)的空間將超過(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)源：https：//le.qun.ch/en/blog/2023/05/13/transformer-batching/]。

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)存帶寬，我們先不算。

如果不做任何批量輸入，也就是模型專門(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)存帶寬之比是 989 / 3.35 = 295。也就是說(shuō)，如果 context 中的 token 數(shù)量小于 330 或者 295，那么內(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)埧ǎ?/strong>