Micro Batch 是深度學(xué)習(xí)中常用的一個(gè)概念，特別是在處理大模型的訓(xùn)練和推理時(shí)。它指的是將一個(gè)大的批次(batch) 進(jìn)一步劃分為多個(gè)較小的批次來處理。這種策略可以用于提升硬件利用率、優(yōu)化內(nèi)存使用或適應(yīng)特定的硬件限制。

Micro Batch 的定義

在訓(xùn)練或推理中，數(shù)據(jù)通常會(huì)以批量（batch) 的形式喂給模型。一個(gè)batch 包含多個(gè)樣本，比如一組句子或者圖像。
Micro batch 是將一個(gè)大批量 (global batch) 切分成若干較小的批次，分別依次送入模型中進(jìn)行計(jì)算，

• Global Batch Size : 整體的批次大小（e.g. : 128)
• Micro Batch Size : 切分后的每個(gè)小批次的大小 (e.g. : 16)
  示例 : 如果 global batch size = 128, 而 micro batch size = 16, 那么整個(gè)批次會(huì)被分成 8 個(gè) micro batches, 每個(gè)包含 16 個(gè)樣本。

2. Micro Batch 的用途

(1) 訓(xùn)練階段

• 顯存優(yōu)化:
  • 處理大模型時(shí)，如果一次性加載整個(gè) batch 所需要的顯存太大，切分成 micro batches 可以逐步處理，減少顯存占用.
  • 通過 梯度累積 (Gradient Accumulation), 每次處理一個(gè) micro batch 的梯度，最后在所有 micro batches 的梯度累加后再更新模型參數(shù)。
• 硬件限制適配:
  • 當(dāng) GPU 的顯存容量不足以支持大 batch 時(shí)，通過 micro batch 的方式可以讓訓(xùn)練依然在較小的 GPU 上進(jìn)行。

(2) 推理階段

• 動(dòng)態(tài)批處理 (Dynamic Batching) ：
  • 推理任務(wù)中，多個(gè)請(qǐng)求可能是并發(fā)的，但每個(gè)請(qǐng)求的數(shù)據(jù)量可能不一致。通過 micro batch 將它們動(dòng)態(tài)合并，能高效地利用硬件資源。
  • for example：一個(gè) LLM 的推理服務(wù)器可能需要同時(shí)處理多個(gè)輸入，將這些輸入組成一個(gè) micro batch 來提高吞吐量。
• 流水線并行優(yōu)化 :
  • miro batch 可以幫助流水線并行 (Pipeline Parallelism) 更好的調(diào)度計(jì)算資源，例如讓每個(gè)模型 slice 都在不同的 micro batches 上并行計(jì)算。

3. Micro Batch 的優(yōu)勢(shì)

1. 顯存效率:

• 通過切分大的 batch，減少顯存需求，適配資源有限的硬件。

2. 吞吐量提升:

• 在推理中，通過 micro batch 動(dòng)態(tài)組合多個(gè)請(qǐng)求，可以顯著提升模型的吞吐量

3. 靈活性:

• 可以適配不規(guī)則的輸入數(shù)據(jù)( for example 長(zhǎng)度不一的序列)

4. Micro Batch 的注意事項(xiàng)

1. 梯度累積（在訓(xùn)練中）：

• 如果切分成 micro batches，需要在最后累積梯度后統(tǒng)一更新參數(shù)，否則效果等同于降低 batch size。

2. 批次內(nèi)數(shù)據(jù)依賴 (inference) :

• micro batch slice 可能導(dǎo)致數(shù)據(jù)的上下文丟失 (e.g. 序列生成任務(wù)中)，需要保證切片后每個(gè) micro batch 的數(shù)據(jù)仍然是獨(dú)立的。

3. 通訊開銷

• 在多 GPU 場(chǎng)景中，如果每一個(gè) micro batch 都需要與其他 GPU 通信，會(huì)帶來額外的通信開銷

5. 應(yīng)用場(chǎng)景

• 大模型訓(xùn)練 ：分布式訓(xùn)練中，通過 micro batch 和 梯度累積來減少顯存的占用。
• 大模型推理 ：LLM 的推理框架 （e.g. Hugging Face 、vLLM) 常用 micro batching 來提高請(qǐng)求的吞吐量。
• 混合并行 : 結(jié)合數(shù)據(jù)并行、模型并行和流水線并行時(shí)，micro batch 常用于更高效的任務(wù)分配。