大模型推理中的批处理大小选择

在大模型推理过程中，批处理大小（batch size）的选择对性能和资源利用率有着至关重要的影响。本文将结合实际测试案例，分享如何科学地选择批处理大小。

批处理大小的影响因素

1. 显存限制：这是最直接的约束条件。较大的batch size需要更多的GPU显存。例如，在使用Transformer模型时，batch size为64时可能占用24GB显存，而增大到128时可能超出显存上限。

2. 吞吐量优化：在某些场景下，适当增大批处理大小可以提高整体吞吐量。但超过某个临界点后，继续增加batch size反而会降低效率。

实验验证步骤

1. 环境准备：使用NVIDIA RTX 4090（24GB显存）进行测试，模型为LLaMA-7B。

# 安装依赖
pip install torch transformers accelerate

2. 推理代码示例：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 测试不同batch size
for batch_size in [1, 4, 8, 16, 32]:
    inputs = tokenizer(["Hello world"] * batch_size, return_tensors="pt", padding=True)
    inputs = {k: v.to("cuda") for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=10)
    print(f"Batch size {batch_size}: {outputs.shape}")

3. 性能测试：使用torch.cuda.memory_summary()和time.time()记录显存占用和推理时间。

最佳实践建议

• 初步测试可从较小batch size开始（如1、4），逐步增加
• 在资源充足时，可以将batch size设置为2的幂次方，以优化内存对齐
• 对于实时推理场景，优先保证响应时间而非最大吞吐量

在实际项目中，建议通过自动化脚本进行多轮测试，记录不同配置下的性能表现，从而找到最适合特定硬件和业务需求的批处理大小。

参考链接：PyTorch Batch Size Tuning