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