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 通信点，不增加重分布开销。
正文 · 原文
Abstract:Variable-length training creates conflicting parallelism requirements across the layers of a hybrid-attention model. The longest samples require wide sequence-parallel (SP) groups of 16 or 32 ranks for full attention, whose cost grows quadratically with sequence length. Sparse-attention layers, dominated by near-linear MoE work, can instead balance most evaluated workloads within node-local SP8 groups. Reusing the wide groups for these layers pushes their all-to-all communication across nodes, where the observed effective rate is 7.0x lower. Yet switching to smaller groups changes which ranks own each sample; a separate redistribution at every full-sparse boundary would introduce another data transfer.
Layerwise Repacking keeps each SP16/32 parent for full attention and repacks its samples into node-local sparse children at SP{1, 2, 4, 8}, where hidden states persist between full layers. The layout change rides on the two existing full-layer Ulysses sites: at each, one variable-size all-to-all performs both the head exchange and the sample repacking, so no redistribution pass is added. Because the balance the children can reach depends on which samples share a parent, an auxiliary parent solver rebalances parent membership before the children are formed.
On all 22 steps with SP16/32 parents among the first 200 training steps of a 1F3S hybrid-attention model (one full-attention layer per three sparse-attention layers), repacking alone reduces the mean end-to-end time by 5.95% and raises model FLOPs utilization (MFU) by 3.75 percentage points over the baseline with repacking disabled; with the improved solver, end-to-end time falls by 7.61% and MFU rises by 4.83 percentage points.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2610.35989 [cs.LG] |
| (or arXiv:2610.35989v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.35989 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Jiahao Zhang [view email]
[v1]
Mon, 26 Oct 2026 12:07:29 UTC (148 KB)
来源：arXiv cs.AI · arxiv.org