LLM 推理批量大小怎样在显存和吞吐之间取舍

LLM 推理里的 batch size 不是越大越好。更稳妥的做法是先把模型权重、KV cache 和运行时工作区留出安全余量，再从 1、2、4、8 这样的序列逐步压测。吞吐仍在增长且 P95 延迟没有越过服务目标时可以继续放大；一旦显存逼近上限、P95 突然抬升或吞吐进入平台期，就应把前一个配置作为工作点。

如果使用 Hugging Face Transformers，官方推理优化文档可作为实现参考：https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one。本文只讨论推理 batch size、显存和吞吐的工程取舍，不把训练时的梯度累积结论直接套到在线服务。

一、先区分固定 batch 与连续批处理

离线生成常见的是固定 batch：程序收集一批长度接近的输入，一次调用模型。此时 batch size 直接决定同一轮并行的请求数，输入越长、输出上限越大，峰值显存越高。

在线服务还要考虑连续批处理。它不必等所有请求同时到达，而是在已有请求生成过程中让新请求加入调度。这样更容易提高 GPU 利用率，但“配置一个 batch size”不再等于“永远同时处理 N 条完整序列”：调度器还要受并发上限、最大序列长度和 KV cache 预算约束。

二、把显存预算拆成四部分

调参前先回答“显存到底被谁吃掉”。可以把峰值显存粗略拆成四类：

• 模型权重：主要由参数量和数据类型决定，batch 变大时通常不会按请求数线性增长。
• KV cache：自回归生成会为每个请求保存历史注意力键值，batch、上下文长度和生成长度都会放大它。
• 激活与临时工作区：注意力后端、矩阵计算和张量转换可能需要额外空间，峰值常出现在某个算子执行期间。
• 框架与碎片余量：分配器缓存、通信缓冲区和显存碎片不能当成可用空间全部拿来下注。

因此，“模型能加载”只说明权重放得下，不说明 batch 8 能稳定生成。尤其是输入长度不一致时，静态批次往往需要 padding，短请求也会跟着最长请求承担一部分计算和缓存成本。

三、用固定条件做第一轮压测

第一轮压测最重要的是固定变量：模型、精度、输入长度区间、最大新 token、采样参数和 GPU 都不要同时改变。下面示例用 Transformers 做离线推理骨架，注释只保留关键的显存与计时逻辑。

import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 使用可替换的公开模型名；真实项目中按自己的模型和精度调整。
MODEL_ID = "distilgpt2"
DEVICE = "cuda"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
).to(DEVICE).eval()

texts = ["请用一句话解释 batch size 的作用。"] * 4
inputs = tokenizer(texts, return_tensors="pt", padding=True).to(DEVICE)

# inference_mode 关闭 autograd；同步 CUDA 后再计时，避免异步执行造成误判。
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()
start = time.perf_counter()
with torch.inference_mode():
    outputs = model.generate(**inputs, max_new_tokens=64)
torch.cuda.synchronize()
elapsed = time.perf_counter() - start

# 峰值显存只能作为本轮固定条件下的观测值，不能脱离长度与模型比较。
peak_gb = torch.cuda.max_memory_allocated() / 1024**3
new_tokens = outputs.shape[1] - inputs["input_ids"].shape[1]
print(f"batch={len(texts)} elapsed={elapsed:.3f}s peak={peak_gb:.2f}GiB new_tokens={new_tokens}")

把 texts 的数量改成 1、2、4、8，分别记录总耗时、生成 token 数、tokens/s、峰值显存和错误信息。不要只记录平均耗时；在线服务至少还要记录 P95 或 P99 延迟，因为大 batch 往往先伤害尾延迟。

四、处理 padding 与长度差异

同样的 batch size，在“8 条短问题”和“8 条长上下文”上不是同一种负载。固定 batch 中只要有一条很长，padding 就可能把其余请求补齐到相近长度，增加无效计算；生成阶段还会因为不同请求的结束时间而出现空转。

工程上可以先按输入长度分桶，再在桶内组成 batch；在线服务则优先使用支持动态调度的 serving 方案，并为单请求输入长度、最大新 token 和总 token 预算设上限。这样调小 batch 不是唯一办法，减少长度离散度同样可能释放显存和吞吐。

# 按 token 长度分桶，减少同一批次的 padding 浪费。
def bucket_by_length(texts, tokenizer, width=128):
    buckets = {}
    for text in texts:
        # 这里只测长度，不在分桶阶段生成模型输出。
        length = len(tokenizer.encode(text, add_special_tokens=True))
        key = (length // width) * width
        buckets.setdefault(key, []).append(text)
    return list(buckets.values())

# 生产环境还应给单桶设置最大 batch 和最大总 token 数，防止长文本聚集。
batches = bucket_by_length(texts, tokenizer)
print([len(batch) for batch in batches])

五、用拐点而不是显存极限选配置

压测结果建议放在同一张表里比较：

不要把示例表中的文字当成固定测量结果。实际拐点由模型、量化方式、上下文长度、GPU、注意力实现和请求分布共同决定。若 batch 从 4 增到 8 只换来很小的吞吐增益，却让 P95 超过目标，batch 4 往往比“显存刚好打满”的 batch 8 更适合线上。

六、把安全值写进服务配置

最终配置至少应同时记录以下内容：安全 batch size、最大输入 token、最大新 token、并发上限、超时、显存水位和 OOM 降级策略。若业务更看重首 token 延迟，就把 P95 首 token 延迟放在首要指标；若是离线任务，则可以用总 tokens/s 和单批完成时间做主指标。

一种实用的发布步骤是：先把压测中低一个档位作为默认值，灰度观察显存峰值与 P95，再逐步增大；遇到 OOM 不要盲目重复提交同一请求，而是降低动态批次、缩短最大输出或将请求转入排队。新模型、新精度、新 GPU 或新的输入长度分布都应重新跑这组基准。

常见问题

batch size 越大，吞吐一定越高吗？

不一定。GPU 并行度可能在较小 batch 就已经吃满，继续增大只会增加排队、padding、KV cache 和尾延迟。

显存还有余量，为什么吞吐不再增长？

瓶颈可能转移到内存带宽、解码串行阶段、CPU 调度、通信或输入长度差异。此时应看算子与服务端指标，不要只看显存占用。

在线服务应该固定 batch size 吗？

固定 batch 适合可预测的离线任务；请求到达时间和长度变化明显的在线服务，更适合动态或连续批处理，并用最大并发、总 token 和延迟目标做约束。

结论：先以 batch 1 建立可重复基线，再按 2 的幂次逐级压测；把峰值显存、tokens/s 和 P95 放在一起看，取吞吐开始趋平或尾延迟越界前的配置，并保留足够的显存余量。