I am primarily responsible for training infrastructure optimization for attention residual. From the perspective of training infra, let me casually chat about a few things.

In fact, full attention residual is constrained by training infrastructure—specifically, by pipeline parallelism. This is because attention residual causes inherent asymmetry in communication and GPU memory under pipeline parallelism, and this asymmetry creates a severe burden on large model training infrastructure. A few months ago, when Su Shen mentioned he wanted to implement full attention residual, I thought about it for a long time but couldn't come up with any good adaptation method for pipeline parallelism (I vaguely felt there might be some special pipeline scheduling that could symmetrize the communication of full attention residual, but it seemed too complex). Therefore, I hoped for an approach similar to SWA (Sliding Window Attention) or linear attention to optimize full attention residual. After communicating with Guangyu and Shaozhang, these big brothers designed the mathematically more fundamental block attention residual in almost one night. Under this algorithm with good locality, the communication optimization of cross-stage caching becomes relatively easy to conceive.

Of course, implementing this in Megatron's "shit mountain" codebase isn't that simple. After a painful weekend, I found that the engineering implementation of this cross-stage caching is actually quite elegant. The general idea is to add an adapter after pipeline parallel communication to concatenate the received blocks with blocks cached in the adapter. The backward pass follows a similar logic: it receives gradients from all blocks, accumulates all these gradients inside the adapter, and when a block needs to be sent to the next stage, it simply sends out the accumulation buffer. The entire code logic is relatively symmetric and doesn't affect the internal network logic.

This optimization method isn't particularly perfect. For example, under block attention residual, it's actually difficult to perform true full recomputation. Fortunately, the overhead of caching these blocks isn't too large, and full recomputation can basically be replaced by selective recomputation plus activation offload under current hardware conditions.

Another point worth noting is that cross-stage caching changes the accumulation order. Under ordinary pipeline parallelism, where each block is sent down layer by layer and corresponding gradients are passed up layer by layer, regardless of how the pipeline parallel configuration changes, the accumulation order should remain the same. But with cross-stage caching, it causes cached blocks to accumulate first within each rank, and only when the block needs to communicate are they accumulated again. This leads to some difficulties in debugging and precision alignment (for example, if the pipeline parallel configuration changes, loss/norm cannot be made completely identical). Of course, some dirty engineering details can make this alignment work easier, but it still adds a lot of complexity.

Some people might be concerned about training performance. In fact, the additional overhead of cross-stage caching is very small. Under the most commonly used interleaved pipeline scheduling, send/recv overhead in the steady stage is very easily hidden, with only a little communication exposed during warmup and cooldown stages. The computation logic of attention residual is very simple—even using the "large model vibe" approach, you can come up with an operator with good memory access performance (no disrespect to mhc intended). Compared to the real attention and MoE computation in the network, the computational overhead introduced by attention residual is negligible. So if anyone wants to use attention residual, following this optimization approach, various infra teams should be able to scale it up with reasonable engineering and performance overhead. From this perspective, it's also a concise and effective optimization method.

Finally, let's talk about how to train full attention residual. The world has suffered long under pipeline parallelism, but it is indeed the most effective means to isolate high/low bandwidth domains. However, as high-bandwidth domains become larger and larger, perhaps pure Zero3 will also be viable—simply restrict allgather parameters and reduce-scatter gradients communication within high-bandwidth domains. Combined with offload, without using pipeline parallelism, one should still be able to train very large models with good performance. At that point, full attention residual won't have many difficulties, and one could even do more aggressive algorithmic optimization work (this does not constitute investment advice).

Similar evolution is already happening. With the popularization of low-precision training, the increase in HBM, the increase in model sparsity, and the explosion of C2C bandwidth, traditional TP in 3D parallelism has basically exited the historical stage. If PP can also be optimized away, then a truly elegant large model training framework may emerge, after all, EP/CP are just variations on the DP dimension.

P.S. While writing this answer, I realized that the attention residual analysis documentation and framework code implementation were basically all done through ancient manual craftsmanship. At that time, I was only used to using Cursor to help review code. LLM Agent development has been too fast (mainly opus 4.6), and within a few months, the world seems to have changed. These ancient handcrafted methods feel like they have become romantic memories of the "classical era."