為什么某些 batch size 會(huì)突然導(dǎo)致性能下降？ 原創(chuàng)
編者按：你是否曾在優(yōu)化深度學(xué)習(xí)模型時(shí)感到困惑，明明增加了 batch size，GPU 利用率卻沒(méi)有如預(yù)期提升？在實(shí)際項(xiàng)目中，這個(gè)問(wèn)題可能導(dǎo)致資源浪費(fèi)、訓(xùn)練效率低下，甚至影響整個(gè) AI 產(chǎn)品的交付周期。
本文作者深入剖析了現(xiàn)代 GPU 批處理的工作原理，揭示了內(nèi)存帶寬與計(jì)算能力之間的微妙關(guān)系。通過(guò)建立理論模型并結(jié)合實(shí)際實(shí)驗(yàn)，作者不僅解釋了為什么某些 batch size 會(huì)突然導(dǎo)致性能下降，還提供了如何找到最佳 batch size 的方法。
作者 | Finbarr Timbers
編譯 | 岳揚(yáng)
一般來(lái)說(shuō)，對(duì)于現(xiàn)代深度學(xué)習(xí)系統(tǒng)而言，你能做的第一個(gè)也是最重要的優(yōu)化措施就是實(shí)現(xiàn)批處理（batching）。在進(jìn)行推理時(shí)，不是單獨(dú)處理單個(gè)輸入，而是同時(shí)處理包含 N 個(gè)輸入的一批數(shù)據(jù)。大多數(shù)情況下，這個(gè)操作是無(wú)需額外成本的 —— 無(wú)論是處理單個(gè)輸入還是 N 個(gè)輸入，推理所需的時(shí)間幾乎相同。這又是為何呢？表面上看，批量處理數(shù)據(jù)似乎應(yīng)該消耗更多資源，畢竟，工作量增加了 N 倍。
然而，如果我們使用一個(gè)簡(jiǎn)單或者不成熟的模型來(lái)理解神經(jīng)網(wǎng)絡(luò)的工作方式，那么批處理（batching）的計(jì)算并不是沒(méi)有成本的。實(shí)際上，批處理確實(shí)需要 N 倍的計(jì)算能力。如果你在 CPU 上運(yùn)行某個(gè)特定的計(jì)算任務(wù)，你會(huì)發(fā)現(xiàn)前文提到的這一點(diǎn)是成立的。
然而，在現(xiàn)代 GPU 上運(yùn)行相同的計(jì)算任務(wù)時(shí)，情況卻并非如此。以下是我們?cè)谝豢?T4 GPU 上觀(guān)察到的情況：
從圖中可以看到，batch size 從 1 到 3 時(shí)，所消耗的時(shí)間并不會(huì)增加。但是，一旦 batch size 超過(guò) 3，時(shí)間消耗就會(huì)呈線(xiàn)性增長(zhǎng)。
這是什么原因呢？關(guān)鍵在于并發(fā)處理能力。現(xiàn)代 GPU 能夠同時(shí)執(zhí)行多次運(yùn)算（盡管在單線(xiàn)程處理時(shí)，它們其實(shí)比 CPU 要慢）。
通常，當(dāng)我們談?wù)摗坝媚Ｐ蛯?duì)單個(gè)數(shù)據(jù)樣本進(jìn)行推理”時(shí)，容易把模型看作一個(gè)整體塊（single block）。但實(shí)際上，模型是由眾多矩陣組成的。推理過(guò)程中，每個(gè)矩陣都會(huì)被加載到內(nèi)存中。具體來(lái)說(shuō)，矩陣的每個(gè)塊都會(huì)被加載到設(shè)備的共享內(nèi)存單元（在 A100 顯卡上僅有 192 kb）。這個(gè)塊隨后用于計(jì)算 batch 中每個(gè)元素的結(jié)果。需要注意的是，這與 GPU RAM（即 HBM）不同。A100 顯卡根據(jù)型號(hào)不同，配備了 40 GB 或 80 GB 的 HBM，但設(shè)備內(nèi)存僅有 192 kb。這導(dǎo)致在執(zhí)行數(shù)學(xué)運(yùn)算時(shí)，內(nèi)存帶寬成為了性能瓶頸，因?yàn)閿?shù)據(jù)需要不斷地在設(shè)備內(nèi)存中讀寫(xiě)。我們可以通過(guò)模型大小除以?xún)?nèi)存帶寬來(lái)估算傳輸權(quán)重所需的時(shí)間，通過(guò)模型的浮點(diǎn)運(yùn)算次數(shù)（FLOPS）除以 GPU 的 FLOPS 來(lái)估算計(jì)算所需的時(shí)間。
使用多層感知機(jī)（MLP），浮點(diǎn)運(yùn)算次數(shù)（FLOPS）大約是參數(shù)數(shù)量的兩倍乘以 batch 中元素的數(shù)量[1]（即為 2 * m * n * b，數(shù)據(jù)批次大小（batch size）為 b ，矩陣為 m x n ）。因此，當(dāng)傳輸時(shí)間等于計(jì)算時(shí)間時(shí)，意味著：
在此，我們可以觀(guān)察到左右兩邊的參數(shù)數(shù)量是可以相互抵消的：
同時(shí)，我們可以根據(jù) batch size 來(lái)重新排列：
當(dāng) batch size 小于 FLOPS 與內(nèi)存帶寬的比值時(shí)，內(nèi)存帶寬將成為性能瓶頸。而一旦 batch size 超過(guò)了這個(gè)比值，計(jì)算能力（FLOPS）則成為新的瓶頸。 請(qǐng)注意，這一分析僅適用于多層感知機(jī)（MLP），對(duì)于像 ResNet50 這樣的卷積神經(jīng)網(wǎng)絡(luò)來(lái)說(shuō)，情況會(huì)更為復(fù)雜。
在 T4 GPU（產(chǎn)品規(guī)格表[2]）上，其浮點(diǎn)運(yùn)算能力達(dá)到 65 TFLOPS（32位浮點(diǎn)數(shù)），內(nèi)存帶寬則是 300 GB/s，按照這個(gè)數(shù)據(jù)，理想的運(yùn)算效率比（magic ratio）應(yīng)該是 216。實(shí)際運(yùn)行一個(gè)深度為 8、寬度為 1024 的多層感知機(jī)（MLP）模型時(shí)，我們得到的結(jié)果與預(yù)期相吻合。
盡管數(shù)據(jù)中存在一些噪聲干擾，但總體趨勢(shì)與我們的預(yù)測(cè)一致：推理時(shí)間在接近 128 的閾值時(shí)開(kāi)始急劇增加（在此，我們采取逐步加倍的方式來(lái)觀(guān)察和記錄不同 batch size 對(duì)推理時(shí)間（inference time）的影響）。如果我們改變 MLP 層的寬度，會(huì)發(fā)現(xiàn)這一現(xiàn)象在多種架構(gòu)中都存在（下面是一張對(duì)數(shù)-對(duì)數(shù)（log-log）坐標(biāo)圖，以便所有的數(shù)據(jù)點(diǎn)都能在圖表中清晰地顯示）。
這真是太酷??了！我們可以看到在多種不同的模型架構(gòu)中，都存在一個(gè)關(guān)鍵的閾值。有趣的是，較小的網(wǎng)絡(luò)在處理速度上并沒(méi)有隨著 batch sizes（從 1 到 512）的增加而變化，基本保持恒定。 我對(duì)此的初步解釋是，這是因?yàn)?GPU 在執(zhí)行數(shù)學(xué)運(yùn)算時(shí)速度極快，而其他硬件（如 CPU）則相對(duì)較慢。在實(shí)驗(yàn)初期，我們觀(guān)察到了大量的噪聲干擾，對(duì)于這一現(xiàn)象，我暫時(shí)只能歸咎于“系統(tǒng)開(kāi)銷(xiāo)（overhead）”。
對(duì)于許多機(jī)器學(xué)習(xí)工程師而言，他們的時(shí)間往往沒(méi)有花在機(jī)器學(xué)習(xí)本身，而是花在消除這些系統(tǒng)開(kāi)銷(xiāo)上，這些開(kāi)銷(xiāo)大多出現(xiàn)在非機(jī)器學(xué)習(xí)相關(guān)的代碼中。在強(qiáng)化學(xué)習(xí)（RL）研究領(lǐng)域，尤其是那些專(zhuān)注于持續(xù)學(xué)習(xí)問(wèn)題（continual learning problems）的研究者，除非1）他們擁有一個(gè)非常大的神經(jīng)網(wǎng)絡(luò)，或者2）對(duì)整個(gè)技術(shù)棧進(jìn)行了極致優(yōu)化，否則在實(shí)驗(yàn)中使用 GPU 往往并不劃算。如果你想讓一位曾在 DeepMind 工作過(guò)的工程師感到尷尬，可以問(wèn)他們關(guān)于“內(nèi)置計(jì)算圖環(huán)境”（in-graph environments）的問(wèn)題——在某個(gè)階段，我們甚至是在 TensorFlow 的計(jì)算圖中實(shí)現(xiàn) RL 環(huán)境。
那么，卷積神經(jīng)網(wǎng)絡(luò)的情況又是如何呢？
在卷積神經(jīng)網(wǎng)絡(luò)中，權(quán)重的總數(shù)是濾波器數(shù)量與濾波器尺寸的乘積。以 torch.nn.Conv2d 為例，權(quán)重的計(jì)算方式是 kernel_size^2 乘以 out_channels。假設(shè)我們處理的是一張分辨率為 (224, 224) 的圖像，步長(zhǎng)為 1，卷積核大小為 3，那么每個(gè)濾波器會(huì)被重復(fù)使用 224 次。這就意味著，在卷積層中，批處理的優(yōu)勢(shì)并不明顯，因?yàn)槲覀儠?huì)反復(fù)使用相同的權(quán)重。至于池化層，其處理計(jì)算量與像素?cái)?shù)量呈線(xiàn)性關(guān)系，這一點(diǎn)與你所想的相符。
Transformers 的情況又是怎么樣呢？
Transformers 本質(zhì)上就是多層感知機(jī)（MLPs），我們可以將它們視為相同的東西。它們具有注意力機(jī)制，但是，由于有了 KV 緩存（能夠?qū)⒂?jì)算數(shù)據(jù)保留在內(nèi)存中），注意力機(jī)制所消耗的時(shí)間被大幅減少。我之前已經(jīng)撰寫(xiě)文章對(duì)此進(jìn)行了深入的探討[3]。
這一觀(guān)點(diǎn)同樣適用于混合專(zhuān)家模型（Mixture of Experts model）。在許多 Transformers 的實(shí)現(xiàn)中，KV 緩存是內(nèi)置于注意力類(lèi)中的（例如，MaxText[4] 就是一個(gè)典型案例[5]）。由于 MoE 模型與普通解碼器之間的差異僅在于，某些前饋網(wǎng)絡(luò)層被替換為了 MoE 層，因此 KV 緩存的表現(xiàn)將保持一致，推理過(guò)程也是如此，但有一點(diǎn)不同。
MoE 層中的門(mén)控機(jī)制會(huì)將數(shù)據(jù)批次（batch）分配到不同的專(zhuān)家上。如果門(mén)控沒(méi)有均勻分配數(shù)據(jù)批次，就可能會(huì)引發(fā)一些問(wèn)題。雖然有避免這種情況的路由機(jī)制（如“expert’s choice”），但在自回歸解碼器中，我們通常只能采用“token’s choice”，這可能會(huì)導(dǎo)致門(mén)控出現(xiàn)偏差。強(qiáng)制門(mén)控均勻分配 tokens 是1）當(dāng)前研究的焦點(diǎn)，并且是2）在訓(xùn)練過(guò)程中需要優(yōu)化的一個(gè)重要目標(biāo)。
Thanks for reading!
Hope you have enjoyed and learned new things from this blog!
About the authors
Finbarr Timbers
empiricist. ml researcher. previously: engineering at deepmind ??
END
本期互動(dòng)內(nèi)容 ??
?你在實(shí)際項(xiàng)目中是如何選擇 batch size 的？有沒(méi)有遇到過(guò)意外的性能瓶頸？
??文中鏈接??
[1]??https://www.stat.cmu.edu/~ryantibs/convexopt-F18/scribes/Lecture_19.pdf??
[3]??https://www.artfintel.com/p/where-do-llms-spend-their-flops??
[4]??https://github.com/google/maxtext??
[5]??https://github.com/google/maxtext/blob/main/MaxText/layers/attentions.py#L91??
原文鏈接：
??https://www.artfintel.com/p/how-does-batching-work-on-modern??
