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