Transformers DataCollatorWithPadding 怎样按 batch 动态补齐

训练文本分类或批量推理时，样本的 token 数通常不一样。推荐把 padding 留到 DataLoader 组 batch 的时刻：DataCollatorWithPadding 默认只补到当前 batch 的最长序列，而不是把整个数据集都填到全局最大长度，通常能减少无效 token 和显存占用。

官方文档：https://huggingface.co/docs/transformers/

先把截断和补齐拆成两个动作

tokenizer 的 truncation 决定“太长的样本如何变短”，collator 的 padding 决定“同一个 batch 如何变成同形状张量”。前者是长度上限，后者是组 batch 时的对齐。若把整个数据集预先补到 512，短文本也会携带大量无效位置；按 batch 补齐则只承担当前批次的最大长度。

from transformers import AutoTokenizer, DataCollatorWithPadding

# 这里只限制过长文本；补齐交给每个 batch 的 collator
tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
def tokenize(batch):
    # truncation 防止单条样本超过模型上限，不在这里 padding
    return tokenizer(batch["text"], truncation=True, max_length=256)

# padding=True 等价于按当前 batch 的最长 input_ids 补齐
data_collator = DataCollatorWithPadding(
    tokenizer=tokenizer,
    padding=True,
    return_tensors="pt",
)

放入 Trainer 或 PyTorch DataLoader(collate_fn=data_collator) 后，collator 会接收一组字典，调用 tokenizer 的 pad()，并把 label 或 label_ids 统一整理为 labels。

三种 padding 选择对应三种长度策略

需要硬件友好形状时，可以叠加 pad_to_multiple_of：

# 先按 batch 最长长度决定基准，再向上补到 8 的倍数
data_collator = DataCollatorWithPadding(
    tokenizer=tokenizer,
    padding=True,
    pad_to_multiple_of=8,
    return_tensors="pt",
)

例如当前批次最长序列是 61，倍数为 8 时最终长度会向上到 64。官方文档说明该参数可用于更好地利用 NVIDIA Tensor Core；它不是“减少 token”的开关，样本很短或长度差异很大时，额外补齐反而可能抵消收益。

显存异常时按这张清单排查

• 先看最长样本：动态补齐的形状由 batch 内最长序列决定，一个异常长文本会让整批都变长。
• 再看截断：只写 padding=True 不会把超长文本截短；需要 truncation=True 与明确的 max_length。
• 检查 pad token：tokenizer 必须有可用的 pad_token；部分因果语言模型没有默认 pad token，不能直接照搬分类模型配置。
• 确认左右侧：padding_side 由 tokenizer 决定。分类任务通常关注 mask 的正确性，生成任务还要结合模型的注意力与缓存约定。
• 核对字段：自定义的 score、元数据等字段不要直接交给 tokenizer.pad；需要先取出，父类完成输入补齐后再放回 batch。

常见问题

padding=True 会把所有样本补到 512 吗？

不会。它默认只补到当前 batch 的最长序列；只有使用 padding="max_length" 并给出 max_length 时，才按固定长度补齐。

pad_to_multiple_of=8 一定更省显存吗？

不一定。它有利于某些硬件的计算对齐，但会增加 padding 位置；应结合 batch 长度分布和实际吞吐量测量。

为什么 DataLoader 报长度不一致？

常见原因是没有把 collator 传给 collate_fn，或显式设置了 padding=False。先确认 DataLoader 走的是同一个 DataCollatorWithPadding 实例。