深入理解LLM中的Batch Size:从原理到实践

📅 2026/8/10 12:08:11
深入理解LLM中的Batch Size:从原理到实践
深入理解LLM中的Batch Size从原理到实践本文从GPU执行机制出发系统梳理Batch Size在大语言模型训练与推理两个阶段中的真实影响并给出工程实践中的权衡建议。一、什么是 Batch SizeBatch Size批大小指的是一次前向/反向传播训练或一次前向计算推理中同时处理的样本token序列数量。在LLM的语境下它有两种常见形态阶段Batch的构成典型取值范围训练Pre-training / Fine-tuning多条独立的文本样本拼成的tensor[B, seq_len]几十 ~ 几千Global Batch推理Inference / Serving同时服务的多条请求continuous batching几 ~ 数百受KV Cache限制很多人把这两者混为一谈但它们背后的瓶颈完全不同下面分开讲。二、底层原理为什么Batch Size会产生影响2.1 GPU的吃大锅饭特性GPU的核心并行单元CUDA Core / Tensor Core数量巨大如H100有16896个CUDA Core但它的启动开销kernel launch overhead是固定的。Batch1 时矩阵乘法变成矩阵×向量GEMV算力利用率极低GPU大部分时间在等数据。Batch64 时变成矩阵×矩阵GEMMTensor Core可以满负荷运转。这引出了一个关键指标算术强度Arithmetic Intensity, FLOPs/Byte。AI计算量 (FLOPs)访存量 (Bytes) \text{AI} \frac{\text{计算量 (FLOPs)}}{\text{访存量 (Bytes)}}AI访存量(Bytes)计算量(FLOPs)​AI 低于某个阈值 →Memory-Bound显存带宽是瓶颈AI 高于阈值 →Compute-Bound算力是瓶颈增大Batch Size最直接的作用就是提升算术强度把任务从Memory-Bound推向Compute-Bound。2.2 Roofline模型视角性能 ↑ │ _______________ Compute Roof算力上限 │ / │ / ← 斜线区域Memory Roof带宽上限 │ / │ / └──────────────────────────→ 算术强度 (FLOPs/Byte) ↑ Batch增大 → 向右移动 → 性能提升直到撞上算力墙三、训练阶段Batch Size 的三重影响3.1 梯度估计的方差统计层面SGD的本质是用mini-batch的梯度去无偏估计全量数据的梯度g1B∑i1B∇θL(xi) g \frac{1}{B}\sum_{i1}^{B} \nabla_\theta \mathcal{L}(x_i)gB1​i1∑B​∇θ​L(xi​)Batch小梯度噪声大 → 更新方向抖动 → 有正则化效果有助于跳出尖锐极小值sharp minima往往泛化更好。Batch大梯度更接近真实梯度 → 训练曲线平滑 → 但容易收敛到尖锐极小值泛化可能变差。这就是著名的“Generalization Gap”现象Keskar et al., 2017。3.2 训练速度 vs. 样本效率优化层面增大Batch → 每个step的GPU利用率↑ →墙钟时间wall-clock变快。但同样的epoch数下step数变少模型更新次数减少。经验规律Batch扩大k倍为保持收敛质量通常需要线性缩放学习率Linear Scaling Rule, Goyal 2017lr × k或使用平方根缩放lr × √k对Adam类优化器更稳配合Warmup避免初期梯度爆炸3.3 显存占用工程层面训练显存 ≈ 模型参数 梯度 优化器状态 激活值Activation其中激活值与Batch Size近似线性相关。这也是为什么显存不够时第一反应是减小micro-batch用**梯度累积Gradient Accumulation**模拟大batch# 等效 Global Batch Size micro_batch × accum_steps × GPU数fori,batchinenumerate(dataloader):lossmodel(batch)/accum_steps loss.backward()if(i1)%accum_steps0:optimizer.step()optimizer.zero_grad()⚠️ 注意梯度累积能省显存但省不了时间前向反向还是要跑那么多次。四、推理阶段Batch Size 的影响重点推理阶段的故事和训练完全不同这里有一个核心概念4.1 自回归解码的两个PhaseLLM推理分为Prefill预填充一次性处理整个promptCompute-Bound算力利用率高。Decode逐token生成每步只生成1个token本质是GEMV严重Memory-Bound。Decode阶段每生成一个token都要把几十GB的模型权重从HBM读一遍但只做一次矩阵×向量运算——算术强度极低。4.2 Batching 是推理提速的核心手段把多个请求凑在一起权重只读一次同时服务B条序列算术强度∝B \text{算术强度} \propto B算术强度∝BBatch Size从1提到32吞吐量往往能提升20~30倍而单请求延迟增加不多。这就是为什么vLLM、TGI、SGLang等推理框架都极度依赖Batching。4.3 Continuous Batching连续批处理传统Static Batching的问题一个batch里必须等最长的序列生成完才能释放短请求被迫陪跑。Continuous Batching迭代级调度每生成一个token就检查一次谁生成完谁退出空位立刻塞入新请求配合PagedAttentionvLLM解决KV Cache碎片化。这让GPU利用率从~30%提升到80%。4.4 Batch Size 在推理中的权衡维度Batch小Batch大吞吐量Throughput低高 ✅单请求延迟Latency/TTFT低 ✅高要等凑batchTPOT每token时延低略增Decode阶段近似不变直到饱和KV Cache显存占用低高且是硬约束⚠️KV Cache显存估算以LLaMA-2 70B, FP16为例KV/Token2×nlayers×nkv_heads×dhead×2bytes \text{KV/Token} 2 \times n_{layers} \times n_{kv\_heads} \times d_{head} \times 2\text{bytes}KV/Token2×nlayers​×nkv_heads​×dhead​×2bytes2×80×8×128×2≈320KB/token 2 \times 80 \times 8 \times 128 \times 2 \approx 320\text{KB/token}2×80×8×128×2≈320KB/token一条4096 tokens的序列 ≈ 1.3GB。这意味着80GB显存最多同时容纳几十条长序列——Batch Size的上限往往不是算力而是KV Cache。五、一张图总结Batch Size 的影响 ┌──────────────────┴──────────────────┐ 训练阶段 推理阶段 │ │ ┌────┼─────┐ ┌────────┼────────┐ │ │ │ │ │ │ 梯度 收敛 显存 吞吐量 延迟 KV Cache 噪声 速度 占用 ↑↑↑ ↑ 硬瓶颈 ↓ ↑ ↑ (主收益) 正则 需调 线性 效果 LR 增长六、实践建议Cheat Sheet训练时✅ 显存允许的前提下尽量用大micro-batch提高GPU利用率✅ 显存不够用梯度累积凑Global Batch Size✅ Batch翻倍记得同步调整学习率 Warmup⚠️ 不要盲目追求超大batch32k样本效率会显著下降参考 Chinchilla / LLaMA 的batch schedule训练后期才逐步增大batch。推理时✅ 永远开启Continuous BatchingvLLM / SGLang / TensorRT-LLM✅ 在线服务关注P99延迟给max_num_seqs设上限✅ 离线批处理任务把batch拉满吞吐优先⚠️ 长上下文场景Batch Size上限由KV Cache决定可考虑FP8 KV Cache或MQA/GQA模型。七、一句话总结Batch Size的本质是在并行度 / 算术强度与显存 / 延迟 / 梯度噪声之间做权衡。训练时它影响的是优化动力学推理时它影响的是系统吞吐量——两者都源于同一个物理事实GPU喜欢一次吃一大口。参考Keskar et al. 2017 (Large Batch Training); Goyal et al. 2017 (ImageNet in 1 Hour); Kwon et al. 2023 (vLLM/PagedAttention); NVIDIA Roofline Model 文档。