这篇论文解释了为什么Megatron-Core的MLA吸收模式在训练时会导致内存暴增，还给出了一个叫LAGA的修复方法，通信更省，速度和内存几乎没损失。
论文揭示Megatron-Core中MLA吸收模式被禁止训练的原因：在DeepSeek-V3规模（n_h=128, seq=16384, SP=8）下，其激活内存比显式形式膨胀20-34%，最高达9.2 GB（eager kernel）或19.2 GB（fused kernel）。该限制在A100上交叉验证，原因在于中间变量维度为n_h×d_kv而非每头K/V。论文提出LAGA方法，在8x Ascend 910B上保持吸收模式的潜变量通信，但本地重建每头K/V：通信削减1.98倍，内存仅增加0.5%，SP=1时位精确，SP=2-8时误差≤1e-3，注意力块吞吐提升1.04-1.06x（单节点）和1.07-1.24x（跨节点）。
A Training-Memory Regression in MLA Sequence Parallelism: Why Megatron-Core Forbids Absorption, and LAGA -- a Communication-Efficient Fix
Multi-head Latent Attention (MLA) ships two implementations in Megatron-Core: an explicit form used for training and an absorbed form -- which slashes collective communication by gathering only the compressed latent -- that is fully implemented but hard-asserted out of training (the forward opens with "assert not (self.training and self.cache_mla_latents)"), allowed only in inference decode. The library documents no reason. We show the restriction is well-founded and quantify why: ported to training, the absorbed form is a memory trap -- its intermediates live in n_h x d_kv dimensions per token, larger than the per-head K/V they replace -- inflating activation memory by 20-34%, up to 9.2 GB at DeepSeek-V3 scale (n_h=128, seq=16384, SP=8, eager kernel; the gap widens to 19.2 GB under a fused kernel), enough to change device-fit. This measurement, validated on two axes (linear in seq and n_h) and cross-verified on NVIDIA A100, explains the otherwise-undocumented restriction and leaves practitioners with no low-communication MLA training path. We then provide one. LAGA (Latent All-Gather Attention) keeps the absorbed form's latent-gather communication but rejects the absorb reformulation, instead reconstructing per-head K/V locally from the gathered latent. On 8x Ascend 910B at real DeepSeek-V3 dimensions, LAGA cuts collective communication 1.98x, matches explicit memory within 0.5%, is bit-identical to explicit at SP=1 and equivalent to within 1e-3 at SP=2-8, and under a fused attention kernel improves attention-block throughput 1.04-1.06x single-node and 1.07-1.24x cross-node -- leading at all sequence lengths in the cross-node regime MLA is deployed for.