arXiv cs.AI· Jiahao Zhang, Junwei Feng, Dongliang Huang, Xiaoling Zhou, Yiqiao Liao, Chengru Song·· 13 小时前AI 评分18
Layerwise Repacking：无需额外重分布步骤的变长混合注意力训练
Layerwise Repacking: Variable-Length Hybrid-Attention Training without an Extra Redistribution Pass
AI 导读
Layerwise Repacking 让混合注意力模型在变长训练中保留 SP16/32 父组处理全注意力层，同时把样本重打包进节点内 SP{1,2,4,8} 稀疏子组，布局切换复用全注意力层已有的两处 Ulysses 通信点，不增加重分布开销。
正文 · AI 翻译
摘要：变长训练在混合注意力模型的各层之间产生了相互冲突的并行需求。最长的样本需要 16 或 32 个 rank 的宽序列并行（SP）组来执行完整注意力，其开销随序列长度呈二次增长。稀疏注意力层以近线性的 MoE 工作为主，反而可以在节点本地的 SP8 组内平衡大多数被评估的工作负载。若将这些层复用宽组，会使其 all-to-all 通信跨节点进行，而实测有效速率低 7.0 倍。然而切换到更小的组会改变每个样本由哪些 rank 持有；若在每个全-稀疏边界处单独重分布，则会引入额外的数据传输。
逐层重打包（Layerwise Repacking）为完整注意力保留每个 SP16/32 父组，并将其样本重打包到 SP{1, 2, 4, 8} 的节点本地稀疏子组中，隐藏状态在完整层之间保持。布局变化依托于两个已有的完整层 Ulysses 站点：在每个站点，一次可变大小的 all-to-all 同时完成头交换和样本重打包，因此不增加重分布过程。由于子组能达到的平衡取决于哪些样本共享一个父组，一个辅助父组求解器在形成子组之前重新平衡父组成员关系。
在 1F3S 混合注意力模型（每三个稀疏注意力层配一个完整注意力层）前 200 个训练步中，全部 22 个使用 SP16/32 父组的步骤上，仅重打包就使平均端到端时间相比禁用重打包的基线减少 5.95%，模型 FLOPs 利用率（MFU）提高 3.75 个百分点；采用改进的求解器后，端到端时间下降 7.61%，MFU 提高 4.83 个百分点。
| 主题： | 机器学习（cs.LG）；人工智能（cs.AI） |
| 引用为： | arXiv:2610.35989 [cs.LG] |
| （或 arXiv:2610.35989v1 [cs.LG] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2610.35989 由 DataCite 颁发的 arXiv DOI（待注册） |
提交历史
来自：Jiahao Zhang [查看邮箱]
[v1]
2026 年 10 月 26 日星期一 12:07:29 UTC（148 KB）
来源：arXiv cs.AI · arxiv.org