序列长度分桶为什么能提升 RNN 训练效率?
简化版
长度分桶把相近序列放入同一 batch,减少 padding 和无效计算,提高吞吐;它不改变模型目标,但会影响数据随机性、batch 大小和分布式负载,需要桶内打乱与 token budget。
详细版
-
先按长度边界分桶,再在桶内/桶间随机采样。
-
固定样本数 batch 在长桶显存更高。
-
按总 token 数组 batch 可稳定显存。
-
分布式各 rank 要获得近似相同 token 负载。
-
评估可排序提速,但最终预测需恢复原索引。
完整版教学
一、分桶减少矩形 batch 中的空白面积
普通 batch 的计算形状由最长序列决定,短序列 padding 也会经过许多算子。
相近长度成批可提高有效 token 比例。
完全全局排序会让训练先全短后全长,破坏 SGD 随机性;实践在较大 shuffle buffer 内排序或随机选择桶。
二、底层机制与公式
padding_waste=1-sum_i L_i/(B*max_i L_i)
packed/token-sparse: sum_i L_i<=budget
dense padded batch: B*max_i(L_i)<=budget
三、带数字的推演
长度 [10,12,100,110] 同批,padding 到 110,有效比例 232/440≈52.7%;分成 [10,12] 与 [100,110],有效比例分别约 91.7% 和 95.5%。
四、方案对比
| 方案/对象 | 核心特点 | 代价或边界 |
|---|---|---|
| 随机固定 batch | 随机性好 | padding 多 |
| 长度 bucketing | 吞吐高 | 采样实现复杂 |
| Token budget batch | 显存稳定 | 每批样本数变化 |
五、执行流程
统计长度分布 -> 设桶边界/token budget -> shuffle 桶与桶内样本
-> pack/pad -> 记录有效 token 比 -> 分布式平衡 -> 还原输出顺序
六、边界条件与工程代价
BatchNorm 等依赖样本数的组件会因动态 batch size 改变统计;序列模型通常用 LN,但仍需关注梯度尺度。
长尾极长序列可能单独成批并造成 straggler,应截断、单独队列或设最大 token。
记忆钩子:Bucketing 的目标是少算 padding,但不能牺牲随机性。
七、常见误区与追问
-
误区:按长度全排序即可。 会显著降低训练随机性。
-
追问:为什么按 token 数组 batch? 让每批总计算和显存更稳定。
-
误区:padding 有 mask 就没有成本。 多数算子仍会处理补齐位置。
-
追问:分布式有什么坑? 各 rank token 数差异造成等待。
-
追问:如何衡量收益? 有效 token 比、tokens/s、显存与收敛一起看。
八、加强记忆
Bucketing 的目标是少算 padding,但不能牺牲随机性。
Bucketing 通过让同批长度接近来减少 padding,但必须保留桶内和桶间随机性。若算子仍处理稠密补齐张量,预算应按 B×max(L_i) 而非只按真实长度和估算;只有 packed 或真正跳过 padding 的实现才更接近 ΣL_i。分布式训练还要平衡各 rank 的 token 负载,并在评估后恢复原样本顺序。