PyTorch推理输入长度分桶与Padding优化实战：减少无效计算

在GPU算力平台进行大模型推理部署时，动态批处理虽然能提高吞吐，但长短请求混在同一批次，会把短文本补齐到最长序列。Padding越多，无效计算与显存占用越高，用户可能看到并发增加后延迟突然上升。输入长度分桶可以在吞吐和等待时间之间取得更稳定的平衡。

一、问题背景

批量推理需要把不同长度的输入整理成规则张量。若一个批次包含极长请求，其余短请求也会扩展到相同长度。大模型训练同样会遇到Padding浪费，但在线服务还要考虑排队时间：桶划分太细，凑批变慢；桶太粗，无效Token增多。

GPU服务器租用资源的显存和计算能力有限，不能只靠增加batch size提高吞吐。选择AI算力平台后，还应根据真实流量分布设计批处理策略，而不是只用固定样例压测。

二、环境准备

准备分词器、模型服务和一批脱敏请求样例，记录输入Token长度、首字延迟、总延迟、吞吐与峰值显存。润云智算提供GPU云服务器、开发镜像与模型资源，官网为https://www.smoothcloud.com.cn/。先统计长度分布：

python 复制代码

lengths = [len(tokenizer(x).input_ids) for x in samples]
for p in [50, 75, 90, 95, 99]:
    print(p, np.percentile(lengths, p))

桶边界应来自实际请求，不要照搬其他项目参数。

三、实操步骤

1. 记录每个请求的Token长度

在进入队列前完成分词，并保存长度与请求ID：

python 复制代码

encoded = tokenizer(text, add_special_tokens=True)
request = {
    "id": request_id,
    "input_ids": encoded["input_ids"],
    "length": len(encoded["input_ids"]),
}

日志中不要记录原始敏感文本，只保留长度、时间和匿名标识。

2. 设计长度桶

例如按业务分布设置若干上界：

python 复制代码

BUCKETS = [128, 256, 512, 1024]

def bucket_id(length):
    for upper in BUCKETS:
        if length <= upper:
            return upper
    return "long"

超长请求应进入独立队列，避免拖慢常规请求。边界数量不宜过多，否则每个队列难以及时凑成批次。

3. 设置"批量上限+等待上限"

每个桶同时设置最大批次和最大等待时间。达到任一条件就提交：

python 复制代码

if len(queue) >= max_batch or oldest_wait_ms >= max_wait_ms:
    batch = queue.pop_many(max_batch)
    run_inference(batch)

低流量时主要依赖等待上限保证延迟，高流量时主要依赖批量上限控制显存。推理部署不能只追求大批次，否则队列等待会抵消吞吐收益。

4. 在桶内动态Padding

python 复制代码

batch = tokenizer.pad(
    features,
    padding=True,
    return_tensors="pt",
)
batch = {k: v.to("cuda") for k, v in batch.items()}

只补齐到当前批次最长输入，而不是统一补到模型最大长度。注意模型的Padding方向、注意力掩码和位置编码要求，调整后要核对输出一致性。

5. 计算Padding浪费率

python 复制代码

real_tokens = sum(item["length"] for item in items)
padded_tokens = len(items) * max(item["length"] for item in items)
waste = 1 - real_tokens / padded_tokens
print("padding_waste:", waste)

同时记录队列等待、模型执行时间和显存峰值。浪费率下降并不代表体验一定改善，若等待时间上升，仍需合并部分桶或缩短凑批窗口。

6. 用真实流量回放验收

按脱敏后的长度分布进行回放，比较优化前后的P50、P95延迟、吞吐、错误率和峰值显存。深度学习服务应同时测试短请求突发、长请求集中与混合流量，避免只优化平均值。

四、常见问题与解决方案

1. 分桶后吞吐反而下降

可能是桶过细导致小批次增多。合并相近桶，或在低流量时允许相邻桶合批。

2. 长请求一直排队

为长请求设置独立等待上限和并发配额，避免因凑不满批次而饥饿。

3. Padding减少但显存变化不明显

模型权重、KV Cache和框架缓存也占显存，应分项观察，不能只看总占用。

4. 为什么首字延迟仍然高

首字延迟还包含排队、分词、Prefill与调度，需要拆分各阶段计时。

五、总结

输入长度分桶的核心是用真实长度分布确定桶边界，再通过批量上限、等待上限和动态Padding共同控制。润云智算等GPU计算环境可提供运行资源，服务团队仍应持续记录Token、延迟和显存指标，让推理性能优化建立在可验证数据上。