技术博客

packing 与文档边界:序列拼接的效率收益与注意力隔离问题

第 120 / 147 篇 · 2025-10-10 · ai · 预训练数据管线注意力

预训练语料是变长的:网页 300 token、论文 8k、代码文件 2k。而训练硬件喜欢定长——注意力 kernel、并行切分、显存规划都按固定窗口优化。中间的空隙怎么填,引出本节的主角:packing(拼接)。朴素做法是 bucket 批处理(按长度分桶再 padding),浪费率能压到一成上下;packing 更进一步,直接把多篇文档首尾相接塞进一个窗口,理论上把利用率推到 100%。但拼接不是免费的午餐,它把"文档边界"这个原本隐式的结构信息变成了显式问题。这篇算效率账,再拆三个边界坑。

效率账:padding 浪费有多贵

按长度分桶 + 桶内 padding 是基线。把语料按长度排序后切成固定大小的桶,同桶内最大长度做 pad,浪费率约等于"桶内长度方差/均值"。实测在网页型语料(长度方差大)上,即使做了 50 桶精细分桶,平均浪费仍在 8%–15%——这部分是纯烧算力,pad 位置的注意力计算和损失都被 mask 掉了,但 FLOPs 和显存照付。

packing 用贪心装箱(first-fit decreasing)把文档拼进定长窗口,利用率能到 99%+。换算成训练预算:同样的算力,packing 等于多训了约 10% 的有效 token。对一个 20 token/参数的 Chinchilla 式模型,这相当于省出小半个数量级的数据采购成本。账很简单,所以 Llama、Qwen、DeepSeek 系的主流实现全都做了 packing。剩下的问题全在边界处理上。

坑一:跨文档注意力污染

packing 之前,每个窗口是"一篇文档(+少量 pad)",注意力天然局限在语义相关的上下文内。packing 之后,窗口第 1000 个位置的 token 能 attend 到前面九篇毫不相干的文档。问题分两层:

污染有多严重? 直观上很吓人,但实验结果出乎意料地温和。Wettig 等人 2024 年的对照实验(Unpacking the Differences in Packing Methods 一类工作)显示:跨文档注意力 vs 完全的文档隔离注意力,下游指标差异在小数点后一位——模型似乎学会了"当前文档之外的信息不可信",注意力权重在边界处自然衰减。这和 attention sink 的机制同构:模型把边界 token 当作天然的"段落分隔符"处理。

但温和不等于免费。 差异虽小且方向不稳定(有的任务跨文档略好、有的略差),真正确定的成本在算力:满注意力矩阵下,所有跨文档 block 的计算都是浪费。一个 4k 窗口拼 10 篇 400-token 文档,有效注意力 block 只占对角线的 10%,90% 的注意力 FLOPs 花在了语义无关的 pair 上。FlashAttention 2 的 varlen 接口(flash_attn_varlen_func)就是为解决这个问题生的:传入每篇文档的偏移量,kernel 内部按 block-diagonal 组织计算,隔离的代价降到几乎为零。

结论:能用 varlen 就用 varlen,既干净又省算力;受限于框架只能用普通 attention 时,跨文档污染的模型质量损失可以接受,但算力浪费是实打实的。

packing 的注意力矩阵:有效 block 只占对角带,跨文档 block 是纯浪费 文档 A(400 tok) 文档 B 200 文档 C(420 tok) pad 12 4k 窗口利用率 = (400+200+420)/4096 ≈ 98.9% ■ 有效 block(对角带,约 10% 面积) □ 跨文档 block:满注意力时照算 FLOPs,varlen 下整块跳过

坑二:位置编码要不要重置

拼进窗口的第二篇文档,它的 RoPE 位置从 400 开始数还是从 0 开始数?

不重置(连续计数)实现最简单,但每篇文档的位置分布都依赖它被拼在第几个——同一篇文档在不同 step 出现在不同偏移,位置相位不同。这等价于给数据注入了不必要的噪声,RoPE 的相对位置性质部分失效。

重置(每个文档从 0 起)语义最干净:位置编码回归"文档内相对位置"的本义,每篇文档的训练信号一致。代价是实现上要在 kernel 层面支持 per-document 的位置偏移,varlen 接口顺带就把这事办了。主流实现(Llama 系、DeepSeek 系)在开了 varlen 之后都选择了重置。

折中:不重置也不隔离(GPT-3 早期的做法)——文档之间放 <|endoftext|> 分隔,位置连续、注意力全开,全靠模型自己学边界。这是历史上第一个可扩展方案,但在今天既没有算力优势(不如 varlen 省)也没有质量优势,只剩实现简单的优点,新项目没有理由再选。

坑三:边界 token 的 loss 怎么算

自回归损失里,每个位置 $t$ 用前 $t-1$ 个 token 预测第 $t$ 个。packing 后,文档 B 的第一个 token 的"上文"是文档 A 的尾部——让模型预测"苹果。苹果。[SEP] 香蕉"里 [SEP] 后面接什么,是在教它一个无法学习的映射(跨文档边界没有规律),这部分 loss 是纯噪声。

处理办法是给文档首 token 的预测位做 loss mask:该位置照常参与 attention(作为上下文的一部分),但不计入损失。配合 EOS/BOS 设计——文档 A 以 EOS 收尾、文档 B 以 BOS 开头——模型能学到"EOS/BOS 意味着重置",边界信号反而变成了可利用的结构。一个细节:BOS 和 EOS 用同一个 token 还是分开,主流实现两者皆可,但必须在 packing 时保证每个文档边界恰好有这对标记,漏一个会导致后续所有位置的训练信号错位,这种 bug 在 loss 曲线上表现为全程抬高几个毫、极难排查。

实操清单

落地一个 packing 管线,核对项按重要性排:

  1. varlen kernel(FlashAttention 2/3 或等效实现)+ per-document 位置重置 + 文档首 token loss mask,三件套一起做;
  2. 装箱算法用 first-fit decreasing 并留 0.1% 以内的余量给 BOS/EOS;
  3. 统计口径切换:训练按"窗口数 × 窗口长"记名义 token,按实际文档 token 记有效 token,两者比值进监控——利用率跌破 97% 说明装箱逻辑出了回归;
  4. 去重和 packing 的顺序:先按文档去重,再装箱。反过来会把同一文档的不同碎片拼进同一窗口,制造虚假的长程依赖。

小结

packing 用接近 100% 的窗口利用率换掉分桶 padding 的 8%–15% 浪费,是预训练管线里性价比最高的优化之一。它引入的三个问题——跨文档注意力、位置重置、边界 loss——有成熟的标准解法:varlen kernel 做文档隔离(顺带省掉九成无效注意力 FLOPs)、每文档位置从 0 重置、文档首 token 不计损失。历史上"分隔符 + 全开注意力"的方案只剩考古价值。监控上把名义 token 与有效 token 的比值纳入日常看板,装箱回归能第一时间暴露。