批量推理长度分桶怎么配置或排查

我第一次给一批长短差异明显的文本做推理时，batch size 明明没有变，显存却随着某几个长样本突然冲高。原因通常不是“模型随机变大”，而是一个 batch 会按最长输入补齐：短文本也占用了最长文本的 token 空间。长度分桶的正确做法，是先在预处理阶段取得每条样本的 input_ids 长度，把相近长度放进同一批，再用动态 padding 只补到该批最长值。

先分清训练参数和推理分桶

Transformers 的 LengthGroupedSampler 会把长度相近的样本放在相邻位置，主要用于 Trainer 的数据加载；新版本文档也把 group_by_length 定义为减少 padding 的训练采样策略。它不会自动接管一个自定义推理循环。推理服务或离线脚本若直接写 for batch in DataLoader(dataset, batch_size=...)，默认仍可能把完全不同长度的样本混在一起。

因此配置要拆成两层：第一层是“哪些样本属于一批”，由长度桶和 batch_sampler 决定；第二层是“批内补多少”，由 tokenizer 或 DataCollatorWithPadding 决定。两层混在一起，是最常见的排查误区。

按 token 长度建立可控的批次

边界不要照搬模型的最大上下文长度。先统计业务样本的 token 长度，再用少量有意义的断点，例如 128、256、512、1024；超出最后边界的样本进入单独的长样本桶。下面的实现只负责产生批次索引，模型仍由你的推理代码调用。

import torch
from torch.utils.data import DataLoader
from transformers import DataCollatorWithPadding

def make_batches(lengths, batch_size, boundaries=(128, 256, 512, 1024)):
    # 先按 token 长度入桶，避免短文本被最长文本大量补齐
    buckets = [[] for _ in range(len(boundaries) + 1)]
    for index, length in enumerate(lengths):
        bucket_id = next((i for i, limit in enumerate(boundaries) if length 

这里的边界是“软约束”：最后一批可能不足 batch_size，长样本也可能把桶撑高。若模型是生成式模型，还要把输入长度和 max_new_tokens 分开预算；后者变大时，即使输入分桶合理，KV cache 仍会增加。

用可观测指标排查配置

不要只盯着平均耗时。给每个 batch 记录 batch_size、max_input_tokens、sum_input_tokens 和 padding 比例：

padding_ratio = 1 - sum_input_tokens / (batch_size * max_input_tokens)。如果比例长期很高，先调整桶边界或减小单批样本数；如果比例不高但显存仍冲高，重点检查 max_new_tokens、模型 dtype、KV cache 和是否混入超长样本。

常见问题

分桶是不是越细越快？

不一定。桶太细会让尾批变多、批次变小，调度和设备利用率可能变差。先保证批次 token 数稳定，再用指标调整。

可以直接把 max_length 当分桶边界吗？

可以作为最后一道上限，但它不等于业务分布。超过上限的输入要明确截断或拒绝，不能静默丢掉关键信息。

为什么开启 padding=True 仍然浪费很多 token？

padding=True 通常只表示补到当前批次最长序列；若同批长度差异本身很大，仍需先改批次组成，而不是继续改 padding 开关。

我的经验是先把分桶当成数据装载问题处理，再谈模型参数优化：先确认每批 token 形状稳定，再逐步调 batch size、对齐倍数和生成上限。这样出现显存峰值时，通常能很快判断是输入分布、padding，还是生成阶段的缓存。