鄙人主要负责 attention residual 的训练 infra 优化,从训练 infra 的角度随便聊两句。
实际上 full attention residual 是被训练 infra 限制住的,具体来说是被 pipeline 并行限制住的。因为 attention residual 会导致 pipeline 并行通信+显存天然的不对称,这种不对称性会对大模型的训练 infra 造成非常严重的负担。
苏神在几个月前说要搞 full attention residual,我想了很久也没想到什么好的 pipeline 并行适配方式(我隐约觉得可能有一种特殊的 pipeline 调度能够对称化 full attention residual 的通信,但是感觉太复杂了)。所以希望能有一种类似 SWA 或者 linear attention 的思路去优化 full attention residual,在和 guangyu/张少 沟通了之后,大哥们几乎一晚上的时间就设计出了数学上更本质的 block attention residual。
在这个具有良好局部性的算法下,cross-stage caching 的通信优化就比较容易想到,除了第一个 vp chunk 外,后面的所有 pipeline 通信都是常数,解决了不对称的问题。当然这个东西在 megatron 这套「屎山」下实现起来没有那么简单。经过一个痛苦的周末,我发现这个 cross-stage caching 的工程实现还算优雅,大致思路就是在 pipeline 并行的通信后增加一个适配器,将收到的 block 与适配器中缓存的 block 进行拼接。反向也是类似思路,会收到所有 block 的 grad,所有的 grad 就在适配器里面做累加,需要发到下一个 stage 的就直接把累加 buffer 发出去。整个代码逻辑还是比较对称的,不影响网络内部逻辑。
这个优化方式其实不算特别完美,比如在 block attention residual 下其实很难做真正的完全重计算。所幸缓存这些 block 的开销不算特别大,而且完全重计算在当前硬件条件下基本可以被选择重计算+ activation offload 上位替代掉。
还有一点值得注意,cross-stage caching 会导致累加顺序的改变。如果按照普通的 pipeline 并行,每个 block 一层层的发下去,对应的 grad 一层层的传上来,无论 pipeline 并行配置如何变化,累加顺序应该都是一样的。但带上 cross-stage caching 之后,会导致被 cache 的 block 先在各个 rank 内累加,到了这个 block 需要通信的时候才会把他们再累加起来。这会导致一些 debug/精度对齐上的困难(比如 pipeline 并行配置变了,loss/norm 没法做到完全一致),当然有些肮脏的工程细节可以尽量让这个对齐工程变得容易一点,但还是增加很多复杂度。
有些人会比较关心训练性能。
实际上这个 cross-stage caching 的额外开销非常小,在最常用的 interleave pipeline 的调度下,send/recv 的开销在 steady 阶段非常容易掩盖,只有 warmup 和 cooldown 阶段的一点点通信会暴露出来。attention residual 的计算逻辑很简单,即使用大模型 vibe 也能搞出来一个访存性能比较好的算子(没有说 mhc 不好 vibe 的意思),相比网络中真正的 attention 和 moe 计算,attention residual 引入的计算开销微乎其微。所以如果想用 attention residual,各家 infra 参考这个优化思路,应该都能在一个合理的工程 /性能开销下 scale 起来,从这个角度来看,也算是个简洁有效的优化手段。
最后说一下 full attention residual 要怎么训起来。
天下苦 pipeline 并行久矣,但 pipeline 并行确实是个隔离高/低带宽域的最有效手段。但随着高带宽域越来越大,可能纯 zero3 也能训得起来,只要将 allgather 参数和 reducescatter 梯度的通信限制在高带宽域内就行了,结合 offload,不用 pipeline 并行应该也能以很好的性能训起来非常大的模型,这个时候 full attention residual 就没有太多难度了,甚至可以做一些更激进的算法优化工作(不构成股票购买建议)。
类似的演进其实已经在发生了,随着低精度训练的普及/ HBM 的增大/模型稀疏度的增加/ C2C 带宽的暴涨,传统 3D 并行里面的 TP 已经基本退出历史舞台,如果 PP 也能被优化掉,那一个真正优雅的大模型训练框架可能呼之欲出,毕竟 EP/CP 都是在 DP 维度上做文章。
P. S. 在写这个答案的时候,发现 attention residual 分析文档+框架代码实现基本都是古法手作,当时只习惯用 cursor 帮着看看代码。LLM Agent 发展太快了(主要是 opus 4.6),几个月之内世界像是变了样子,这些古法手作感觉已经成为「古典时代」的浪漫回忆了。