原文标题: Libra: Taming Attention Workload Skew in Long-Context LLM Training with Bounded Sequence Pool
首次公开: 2026年7月25日
署名单位: 中国科学院大学、阿里巴巴集团、新加坡国立大学
核心术语: 长上下文训练(Long-Context)、大语言模型(LLM)
原论文: https://arxiv.org/abs/2607.23250
龙哥导读
长上下文训练有个很反直觉的问题:每张卡拿到的 token 一样多,工作量却可能差十倍。阿里巴巴 Jingren Zhou 团队参与的 Libra,没有继续扩大一个全局调度池,而是把注意力计算限制在固定大小的序列池里,再用序列放置、细粒度切片和通信流水共同削平拖尾。在通义千问训练模型(Qwen3-Turbo)的 256K 与 1M 上下文训练中,端到端吞吐相对 Ulysses 最高提升 2.54 倍,数据并行扩到 16 路时,每卡扩展效率从 27.6% 提高到 70.3%。
明明总 token 数相同,为什么有的 GPU 早早算完,只能等着最慢那张卡?这不是显存分配的小毛病,而是长上下文训练扩到几百、几千张卡后,会持续吞掉集群利用率的系统问题。
此前讨论视觉和视频模型提速时,主线都是"别让所有计算一视同仁"。Libra 把同一个工程判断推到训练集群:真正昂贵的不是 token 本身,而是 token 背后的注意力关系。如果调度器只数 token,不数注意力工作量,看起来平均,实际仍会被极少数超长样本拖住。
一、同样装满一箱货,为什么重量能差十倍?
训练系统常把多个不同长度的文本拼成固定 token 数的样本。这样做能让显存占用和线性计算更整齐:每个训练样本看起来都装满了同样大小的"箱子"。但注意力计算不是按箱子总体积收费,而是与箱子里每段序列长度的平方有关。
论文给了一个非常直观的例子。在其他内容相同的情况下,一段 4 万 token 的长序列,产生的注意力浮点计算量,大约是十段各 4000 token 序列的十倍。两边总 token 都是 4 万,显存与很多线性算子看起来接近,注意力部分却完全不是一回事。
可以把核心成本写成一个简单关系:一个拼接样本的注意力工作量,近似等于其中每段序列长度平方的总和。这里每个长度代表一段原始文本包含多少 token;先分别平方再相加,是因为每个 token 要与同一段序列里的大量 token 建立注意力关系。这个公式的用途是估算真正的计算负载,而不是只看打包后的总长度。
问题又被真实数据分布放大。论文所用长上下文语料的序列长度中位数只有 644 token,但尾部存在接近 100 万 token 的样本。大多数样本很短,少数样本极长,这就是典型的长尾分布。调度稍有不慎,某个数据并行副本就会抽到"重量级箱子"。
同步训练最怕这种拖尾。每轮更新都要等最慢的副本完成;流水线训练还会因为不同微批次耗时差异产生气泡。快卡不能提前进入下一轮,只能在原地等待。卡越多,遇到极端样本的机会越大,看似增加了硬件,实际增加的等待也更多。
这里还要区分"数据长度"和"计算长度"。两个打包样本都可能占满 256K token,但一个由许多短文档组成,另一个包含少量超长文档。前者的注意力矩阵被多个掩码区域切开,后者则包含更大的连续注意力区域。训练框架如果只记录打包后的总长度,就会把两种完全不同的算力需求当成同一种样本。
二、扩到 256 张卡,为什么只得到 4.42 倍吞吐?
论文在 Qwen3-Turbo、100 万 token 打包样本上展示了扩展问题。上下文并行度固定为 16,把数据并行度从1扩大到16,总 GPU 数量从 16 张增加到 256 张。理想情况下,吞吐应该接近 16 倍;实际只提高 4.42 倍。
换算成每卡扩展效率,只有 27.6%。也就是说,大量新增计算卡没有转化为等比例训练速度。它们并不是完全没在工作,而是在反复等待少数注意力任务最重的工作进程。

图1:论文的数据并行扩展结果。横轴是数据并行规模,纵轴是相对单副本的归一化集群吞吐;虚线表示理想线性扩展。256K 与 1M 两组中,Libra 都比 Ulysses 更接近理想线,1M、数据并行度 16 时由 4.42 倍提升到 11.25 倍。
图里最值得看的不是某一个最高点,而是曲线分叉越来越大。数据并行规模较小时,负载偏差还有机会被平均;规模增大后,Ulysses 曲线越来越偏离理想线。Libra 没有达到完美线性扩展,但把 1M 任务的数据并行度 16 吞吐从 4.42 倍推到 11.25 倍,对应每卡效率达到 70.3%。
这项工作的价值不是"让一张卡算得更快",而是让已经买来的大量计算卡少等一会儿。 在数百张乃至上千张卡的训练任务里,等待时间下降几个百分点,都可能对应可观的训练周期和算力成本。
三、Libra 的关键选择:池子不要无限变大
面对负载不均,直觉做法是把更多计算卡放进同一个全局池,让所有任务在更大范围里重新分配。池子越大,越容易找到空闲算力,统计上也更容易把重任务摊开。但注意力任务迁移需要交换查询、键、值等张量,调度范围越大,通信跨越的节点越多。
如果计算省下来的时间,被跨机通信重新吃掉,调度就只是把瓶颈从计算搬到网络。全局池还更容易碰到低带宽链路、硬件速度差异和运行时抖动。规模越大,这些问题越难控制。
Libra 的反转在这里:它不让负载均衡池随着数据并行规模一起变大,而是固定每个序列池的规模。扩容时,系统增加更多固定大小的池,而不是把现有池不断扩张。这样每次注意力交换的范围保持有界,调度器也更容易把一个池放在局部高速互联区域。
论文借用了大数定律的直觉:多个随机样本聚在一起时,平均负载会逐渐靠近总体均值。池子不必覆盖整个集群,只要规模适中,就有机会平滑长尾波动。真正决定池大小的,应当是数据负载分布和允许的不均衡程度,而不是集群到底扩到多少张卡。
不过,纯靠随机平均还不够。长尾数据里,一个百万 token 级样本就可能破坏整个池的平衡。Libra 因此不是简单地"多抓几条序列放一起",而是又加了两层主动调度:先在池之间安排互补负载,再在池内部拆分注意力工作。

图2:Libra 方法总览。左侧把不同注意力工作量的打包序列分到固定大小的序列池;中间通过方差缩减放置降低池与池之间的负载差异;右侧把序列与注意力头进一步切成小块,在池内重排,并通过流水线让张量交换与注意力计算重叠。
四、第一层:先把"重箱子"分散开
方差缩减序列放置算法(VRSP)负责池与池之间的平衡。它先估算每个打包序列的注意力工作量,再把重序列与轻序列组合,尽量让每个池承接的总计算量接近。
这和搬家装车很像。只按箱子数量分车,每辆车都装十箱,看起来公平;但如果一辆车恰好全是书,另一辆全是枕头,真正负载仍然差很大。VRSP 做的是先称重,再把重箱和轻箱搭配,而不是只数箱子。
论文采用固定基数的最长处理时间优先放置思路:先处理最重的序列,并把它放到当前总负载较轻的池里,同时满足每个池包含固定数量序列的约束。计算顺序可以理解为"按重量从大到小排序—查看各池当前负载—放入最合适的池—更新负载"。
固定基数很重要。训练系统不能只追求浮点计算完全相等,还要维持样本数量、优化器步语义和数据顺序约束。VRSP 没有随意删除、复制或改写原始样本,而是在一个优化器更新窗口内重新安排它们进入哪个池。
表格实验显示,在 256K 任务、全局批量 128、池大小 8 时,池间浮点计算不均衡度从生产采样顺序下的 1.08 降到 0.005;1M 任务、全局批量 32 时,从 0.52 降到 0.007。它说明中等大小的池并非只能被动等待大数定律,显式放置可以把有限样本窗口里的残余偏差继续压低。
五、第二层:把不可分的长序列继续切小
只在池之间平衡,还会碰到一个硬问题:单个超长序列可能比池里其他序列重很多。序列如果被视作不可拆分的整体,无论怎样换位置,负责它的上下文并行组仍可能成为最慢的一组。
分块注意力池化方法(TAP)把工作粒度从完整序列继续细化到"序列片段×注意力头"的小块。序列片段代表一段 token 区间,注意力头代表模型内部不同关系子空间。把两维结合,系统就能把一个巨大注意力任务拆成更多可以独立安排的小任务。
TAP 会同时考虑两类成本。第一类是计算:某个小块需要多少注意力浮点运算;第二类是搬运:为了让另一张计算卡处理它,需要交换多少查询、键、值与输出张量。只看计算,会把任务分得很均匀,却可能让网络流量爆炸;只看通信,又会把重任务留在原处。
因此,TAP 寻找的是计算平衡与数据移动之间的折中。它不追求把每个小块都发到全局最空闲的计算卡,而是在有界池内选择合适位置。池大小固定后,通信域也固定,工程上更容易预测网络开销和拓扑影响。
这一层补上了工作量平衡四维并行等序列或微批次级方法仍可能留下的缺口。Libra 论文在 256K 实验中指出,这类方法能改善平均延迟,但不可分割的异常长打包序列仍会控制最慢步骤。分块注意力池化把这类大块继续拆开,才能真正压低"最大值"而不只是平均值。
六、第三层:通信躲到计算后面,但不能凭空消失
把注意力小块调到其他 GPU,会增加张量交换。Libra 用分块注意力池化流水线把交换切成多个分块,让中间分块的通信尽量与闪存注意力(FlashAttention)计算重叠。简单说,一边计算当前块,一边提前搬下一块,不让网络和计算完全串行。
但重叠不是魔法。第一个分块开始计算前,没有上一段计算可以遮住它的输入传输;最后一个分块算完后,也没有下一段计算可以遮住输出返回。池越大、跨计算卡移动越多,暴露在外面的通信仍会增长。
池大小扫描把这个代价讲得很清楚。256K 任务里,池大小从 4 增到 8、16、32 时,归一化通信成本从 0.07 升到 0.17、0.47、0.77。计算负载可能更平,通信却快速上升。1M 任务也呈现类似趋势。
论文大多数实验选用池大小 8,不是因为 8 在每次实测里都恰好最快,而是离线模拟器综合估计了平衡收益和通信代价。事实上,部分配置中池大小 4 的测量延迟略低。这是一个很诚实的工程细节:系统参数不是越大越高级,适中往往比全局最优想象更可靠。
Libra 真正成立的前提,是"局部调度已经足够平",而不是"通信可以免费"。 当网络拓扑、序列分布或并行布局变化时,池大小和重叠策略仍需要重新估计。
七、实验结果:赢的不是一个平均数
论文在通义千问混合专家模型(Qwen3-30B-A3B)训练上测试 256K 与 1M 两套生产工作负载。相对 Ulysses,端到端吞吐最高提升 2.54 倍;注意力微基准中,最慢步骤最高加速 3.14 倍。系统还在 32K 到 1M 上下文的通义千问系列任务上累计运行数十万 GPU 小时。
数据并行扩展更能说明方法针对的是系统瓶颈。数据并行度 16 时,256K 任务的相对吞吐由 7.63 倍提高到 13.66 倍,对应端到端 1.79 倍加速;1M 任务由 4.42 倍提高到 11.25 倍,对应 2.54 倍加速。上下文越长、负载尾部越重,平衡收益越明显。
在核心注意力层的直接比较中,Libra 同时对比 Ulysses、工作量平衡四维并行系统(WLB-LLM)和核心注意力解耦方法(DistCA)。256K 任务上,Libra 完整版本相对 WLB-LLM 的平均延迟降低 15.9%,最大延迟降低 68.4%;1M 任务上分别降低 42.1% 与 65.6%。
平均值与最大值必须分开看。平均值决定没有流水线并行时的整体吞吐,最大值更接近流水线里最慢微批次造成的气泡。WLB-LLM 已经能平衡很多常规负载,但异常长样本仍可能控制最大值;Libra 继续拆分并重排注意力小块,因此最大延迟下降更明显。
DistCA 采用全局注意力服务池,调度范围更大。在 1M、较小数据并行规模下,它与 Libra 的差距相对小;随着集群扩大,全局池通信域持续增长,而 Libra 保持固定池大小。论文在 1M 比较中报告 Libra 平均延迟仍低 7.7%、最大延迟低 7.9%。
这些结果支持的结论很明确:Libra 不仅让通常步骤更快,也重点压低了会卡住整轮训练的最慢步骤。但数值来自作者的 Qwen 训练栈、指定硬件拓扑和生产数据分布,不能直接外推为所有模型、所有网络环境都能获得 2.54 倍。
消融实验也给出了每一层设计的作用。只使用池内切片,能改善最重注意力任务;加入序列放置后,池与池之间更均匀;再加入流水重叠,部分通信被隐藏。三者不是互相替代,而是分别处理池间偏差、池内偏差和新增通信。少任何一层,另一个瓶颈都可能重新成为最长木板。
从评价设计看,论文还做对了一件事:没有只拿最终端到端吞吐说故事。它分别报告数据并行扩展、核心注意力平均延迟、最慢步骤延迟、池间不均衡、池大小扫描和通信—计算拆分。这样读者能看到收益到底来自更公平的放置、更细的任务切分,还是仅仅来自某个特定并行参数。系统论文最怕"总分赢了,但不知道为什么赢",这组实验至少把主要因果链拆得比较清楚。
八、把它放回两条真实技术路线里看
第一条路线是微软 Ulysses 长序列训练方法。Ulysses 把输入沿序列维切分,并通过高效全互联通信完成注意力计算。其论文报告,在序列长度扩大 4 倍时,相比当时基线训练速度提高 2.5 倍。它解决的是极长序列怎样跨 GPU 扩展,是 Libra 端到端实验中的基础对照。
Libra 并没有否定 Ulysses。相反,它在上下文并行已经存在的前提下追问:不同数据并行副本拿到的注意力任务不一样重时,怎样避免最慢副本拖住所有人?因此,Ulysses 更像是"把一条长序列铺开",Libra 则是在此基础上"把不同长序列带来的工作量重新摆平"。
第二条路线是工作量平衡四维并行系统。它从四维并行训练出发,在流水线并行层采用工作量感知的可变长文档打包,在上下文并行层对单个文档做细粒度分片。其正式论文在内部训练框架中报告平均 1.23 倍加速,并展示了 8000 张 H100、4050 亿参数模型、128K 上下文训练里最慢 GPU 计算延迟可高出 1.44 倍的真实问题。
两者共同点是都不再满足于"token 一样多"。区别在于,WLB-LLM 从文档打包和并行层级入手,Libra 强调有界序列池,并把不可分割的异常长序列继续拆成序列片段与注意力头的小块。一个更接近训练输入与四维并行调度,一个更聚焦注意力核心负载及其局部交换。
把这三篇工作连起来,长上下文训练的演进路径很清楚:先解决序列如何跨卡,再解决文档与并行层级如何平衡,最后继续处理注意力核心里最顽固的长尾拖尾。
九、落地时最值得问的四个问题
第一,数据分布是否真的足够长尾。 如果训练样本长度很整齐,或者上下文远没有达到几十万 token,注意力负载差异可能不大。此时引入复杂调度、张量迁移和流水重叠,收益未必覆盖工程成本。
第二,网络拓扑能否支撑池内交换。 Libra 强调把池限制在局部通信域,但真实集群的节点布局、网卡带宽、超节点结构和任务混部都会影响结果。如果调度器无法稳定把一个池放进高速互联区域,论文里的通信估计可能失真。
第三,训练栈是否允许替换注意力接口和数据采样器。 论文把 Libra 实现成可插拔数据采样器与上下文并行注意力算子,不需要改模型层、优化器和流水线计划。这比重写整个训练框架友好,但仍要求底层并行栈能暴露足够的调度与通信控制。
第四,离线估算能否跟上数据变化。 池大小、序列放置和分块策略依赖工作量估计。训练数据来源改变、课程学习阶段切换、过滤规则调整,都可能改变长度与文档组成分布。一次调好的参数,不应默认在整个训练周期里永远最优。
论文没有公开完整生产实现,外部团队很难直接复现数十万 GPU 小时级别的部署结论。公开文本足以理解设计与主要实验,但要迁移到英伟达 Megatron 训练框架、微软分布式训练框架或其他内部系统,仍需要较强的分布式系统、注意力内核和集群调度能力。
对准备评估这条路线的团队,最小可行步骤不是立刻重写注意力内核,而是先把每个训练步骤的序列组成、估算浮点计算量、平均延迟和最大延迟记录下来。若最大值长期显著偏离平均值,再用离线模拟比较简单分桶、序列重排和固定池方案。只有当拖尾确实稳定存在,才值得进入更复杂的池内切片与通信流水开发。
还有一个容易被忽略的组织成本。负载估算需要数据侧统计,序列放置涉及采样器,池内调度涉及注意力运行时,局部部署又涉及集群调度器。它跨越数据、训练框架、内核和基础设施四个团队。论文里的"可插拔"主要指不改模型数学语义,并不等于接入工作只需改几行配置。
十、龙哥点评:这不是又一个"堆卡提速"故事
龙哥最认可的地方,是 Libra 把负载均衡范围本身当成了设计变量。很多系统优化默认"能调度的范围越大越好",这篇论文却明确展示:更大的池一边带来更平的计算,一边也制造更重的通信。真正的最优解不是无限扩张,而是在统计平滑、局部通信和拖尾控制之间找一个能长期运行的平衡点。
第二个价值,是它盯住了最大延迟。平均吞吐当然重要,但同步训练和流水线训练最容易被极端步骤卡死。只把平均数做漂亮,常常掩盖真正昂贵的等待。Libra 从池间方差、池内细粒度切片到通信重叠,整条链都围绕"别让最慢那一步拖住所有人"。
第三个判断更现实:这项工作最适合已经在做超长上下文大规模预训练、并且能够修改底层并行栈的团队。对普通应用开发者,它不会直接让一次模型调用更快;对只有几十张卡、序列长度相对整齐的训练团队,部署复杂度也可能高于收益。
它的不足同样明确。核心生产实现没有完整开源,工作负载主要来自 Qwen 训练栈,数值与阿里集群拓扑、数据组成和注意力实现深度绑定。论文证明了有界池路线在这些条件下很强,但外部复现仍要回答:换网络、换模型、换数据后,最优池大小和通信重叠还能否保持。
对研究者而言,下一步最有价值的并不是再报一个更大的峰值,而是公开不同长度分布、不同互联拓扑和不同注意力实现下的稳定收益,并给出可移植的调度接口。只有外部团队能够复现实质性的拖尾下降,这套方法才会从优秀的生产经验变成更通用的训练基础设施。类似的经验分享和工程讨论,在云栈社区的技术论坛里也经常能引发不少有价值的交锋。
龙哥的最终判断是:Libra 真正值得记住的,不是 2.54 倍这个数字,而是"按真实工作量调度,并给调度范围设上限"这条原则。 当模型越来越大、上下文越来越长,系统优化的核心会从单个算子速度,转向如何避免少数极端任务绑架整个集群。
龙迷三问
如果未来注意力不再是二次复杂度,Libra 还有价值吗? 有,但价值会变化。只要样本工作量仍有长尾、同步训练仍受最慢副本限制,按真实成本而非 token 数调度就仍然重要;只是负载估算公式和最值得切分的算子可能改变。
为什么不直接丢掉特别长的样本? 因为超长样本往往正是长上下文能力需要学习的材料。简单截断或过滤虽然能让训练更整齐,却可能改变数据目标。系统优化的意义,就是尽量在保留原始样本集合和训练语义的前提下减少等待。
这会不会只是大厂专属优化? 当前实现门槛确实偏高,但原则并不专属。中小团队同样可以先统计样本真实计算成本、观察最慢步骤、避免只按 token 平均,并用更简单的分桶或批次重排减少拖尾。未必需要完整复刻 Libra,也能先拿到一部分收益。
主要参考资料
Libra 论文:
https://arxiv.org/abs/2607.23250
Ulysses 正式论文:
https://arxiv.org/abs/2309.14509
WLB-LLM 正式论文页面:
https://www.usenix.org/conference/osdi25/presentation/wang-zheng
本文基于龙哥读论文 PaperDaily 数据库及 PaperMiner 的 MCP 进行汇总整理。本文为论文解读与个人学习笔记,不构成任何论文审核意见,及对产品能力、安全性或商业可用性的保证。