A Training-Memory Regression in MLA Sequence Parallelism: Why Megatron-Core Forbids Absorption, and LAGA -- a Communication-Efficient Fix
2026-07-20 • Distributed, Parallel, and Cluster Computing
Distributed, Parallel, and Cluster Computing
AI summaryⓘ
The authors studied two ways Megatron-Core implements Multi-head Latent Attention (MLA) for training and inference. They discovered that the 'absorbed' form, which reduces communication in inference, causes very high memory use during training, explaining why it is not allowed there. To fix this, they designed LAGA, a method that keeps low communication but reconstructs needed data efficiently for training. Their tests on large models show LAGA uses similar memory and speeds up training compared to the original method, making it practical for real-world large-scale use.
Multi-head Latent AttentionMegatron-Coreactivation memorycollective communicationDeepSeek-V3Latent All-Gather Attentionattention kernelAscend 910Bfused kernelsequence length
Authors
Changzheng Ma
Abstract
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.