From d722f697558b2f45b1f8916eb8551f807fe3646c Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Wed, 10 Jun 2026 13:32:01 -0700 Subject: [PATCH 1/4] Refactor absorbed MLA projection handling Signed-off-by: Hollow Man --- .../absorbed_mla.py | 334 ++++++++---------- .../test_absorbed_mla.py | 148 +++++--- 2 files changed, 256 insertions(+), 226 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index 48e6c76ea2f..f8f4c37ea7c 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -18,6 +18,7 @@ import torch +from megatron.core import tensor_parallel from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.models.common.embeddings import ( RotaryEmbedding, @@ -36,7 +37,7 @@ from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import MLATransformerConfig -from megatron.core.utils import deprecate_inference_params, get_pg_size +from megatron.core.utils import deprecate_inference_params, get_pg_size, is_te_min_version try: from megatron.core.fusions.fused_mla_yarn_rope_apply import ( @@ -58,6 +59,52 @@ TEColumnParallelLinear, TELinear, Linear, set_save_original_input = None, None, None, None +def _restore_packed_thd_batch_dim( + core_attn_out: torch.Tensor, hidden_states: torch.Tensor, packed_seq_params +) -> torch.Tensor: + """Restore the singleton packed-THD batch dim only when core attention omitted it.""" + if ( + packed_seq_params is not None + and packed_seq_params.qkv_format == 'thd' + and core_attn_out.ndim == hidden_states.ndim - 1 + ): + core_attn_out = core_attn_out.unsqueeze(1) + return core_attn_out + + +def _apply_absorbed_v_up_projection( + core_attn_out: torch.Tensor, + v_up_weight: torch.Tensor, + num_attention_heads_per_partition: int, + kv_lora_rank: int, + v_head_dim: int, + core_consumed_v_up_projection: bool, +) -> torch.Tensor: + """Apply V up projection unless core attention already consumed the projection weight.""" + latent_output_size = num_attention_heads_per_partition * kv_lora_rank + projected_output_size = num_attention_heads_per_partition * v_head_dim + if core_consumed_v_up_projection: + if core_attn_out.size(-1) != projected_output_size: + raise RuntimeError( + "AbsorbedMLA core attention returned unexpected projected hidden size: " + f"{core_attn_out.size(-1)}. Expected projected={projected_output_size}." + ) + return core_attn_out + + if core_attn_out.size(-1) != latent_output_size: + raise RuntimeError( + "AbsorbedMLA core attention returned unexpected hidden size: " + f"{core_attn_out.size(-1)}. Expected latent={latent_output_size}." + ) + + core_attn_out = core_attn_out.view( + *core_attn_out.shape[:-1], num_attention_heads_per_partition, kv_lora_rank + ) + core_attn_out = torch.einsum("...nc,ndc->...nd", core_attn_out, v_up_weight) + core_attn_out = core_attn_out.contiguous() + return core_attn_out.view(*core_attn_out.shape[:-2], -1) + + @dataclass class AbsorbedMLASelfAttentionSubmodules: """ @@ -68,8 +115,7 @@ class AbsorbedMLASelfAttentionSubmodules: linear_q_down_proj: Union[ModuleSpec, type] = None linear_q_up_proj: Union[ModuleSpec, type] = None linear_kv_down_proj: Union[ModuleSpec, type] = None - linear_k_up_proj: Union[ModuleSpec, type] = None - linear_v_up_proj: Union[ModuleSpec, type] = None + linear_kv_up_proj: Union[ModuleSpec, type] = None core_attention: Union[ModuleSpec, type] = None linear_proj: Union[ModuleSpec, type] = None q_layernorm: Union[ModuleSpec, type] = None @@ -110,12 +156,16 @@ def __init__( layer_number=layer_number, attn_mask_type=attn_mask_type, attention_type="self", + cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, ) assert not config.add_bias_linear, "add_bias_linear is not supported for AbsorbedMLA" + assert not ( + config.tensor_model_parallel_size > 1 and not config.sequence_parallel + ), "AbsorbedMLA requires sequence_parallel when tensor_model_parallel_size > 1" self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads self.q_head_dim = self.config.qk_head_dim + self.config.qk_pos_emb_head_dim @@ -302,34 +352,19 @@ def __init__( **kv_down_proj_kwargs, ) - # Build separate K and V up projections - self.linear_k_up_proj = build_module( - submodules.linear_k_up_proj, + self.linear_kv_up_proj = build_module( + submodules.linear_kv_up_proj, self.config.kv_lora_rank, - self.config.num_attention_heads * self.config.qk_head_dim, + self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), config=self.config, init_method=self.config.init_method, gather_output=False, bias=False, skip_bias_add=False, is_expert=False, - tp_comm_buffer_name='k_up_proj', + tp_comm_buffer_name='kv_up_proj', tp_group=pg_collection.tp, - name=(name + ".linear_k_up_proj") if name is not None else None, - ) - self.linear_v_up_proj = build_module( - submodules.linear_v_up_proj, - self.config.kv_lora_rank, - self.config.num_attention_heads * self.config.v_head_dim, - config=self.config, - init_method=self.config.init_method, - gather_output=False, - bias=False, - skip_bias_add=False, - is_expert=False, - tp_comm_buffer_name='v_up_proj', - tp_group=pg_collection.tp, - name=(name + ".linear_v_up_proj") if name is not None else None, + name=(name + ".linear_kv_up_proj") if name is not None else None, ) if self.config.q_lora_rank is not None: @@ -511,17 +546,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = torch.unsqueeze(k_pos_emb, -2) - # Prepare k_up_weight for absorption - # k_up_weight: linear_k_up_proj.weight viewed as [n, qk_head_dim, kv_lora_rank] - assert self.linear_k_up_proj.weight.size(0) == ( - self.num_attention_heads_per_partition * self.config.qk_head_dim - ) - assert self.linear_k_up_proj.weight.size(1) == self.config.kv_lora_rank - k_up_weight = self.linear_k_up_proj.weight.view( - self.num_attention_heads_per_partition, - self.config.qk_head_dim, - self.config.kv_lora_rank, - ) + k_up_weight, _ = self._get_kv_up_weights() if self.config.apply_rope_fusion: # q_no_pe: [num_tokens, n, qk_head_dim] @@ -643,17 +668,77 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po return q_absorbed, kv_compressed, q_compressed + def _get_v_up_weight(self) -> torch.Tensor: + """Return V up-projection weight in per-head layout.""" + _, v_up_weight = self._get_kv_up_weights() + return v_up_weight + + def _get_kv_up_weights(self) -> tuple[torch.Tensor, torch.Tensor]: + """Return K and V up-projection weights from the combined per-head MLA layout.""" + expected_rows = self.num_attention_heads_per_partition * ( + self.config.qk_head_dim + self.config.v_head_dim + ) + assert self.linear_kv_up_proj.weight.size(0) == expected_rows + assert self.linear_kv_up_proj.weight.size(1) == self.config.kv_lora_rank + kv_up_weight = self.linear_kv_up_proj.weight.view( + self.num_attention_heads_per_partition, + self.config.qk_head_dim + self.config.v_head_dim, + self.config.kv_lora_rank, + ) + k_up_weight = kv_up_weight[:, : self.config.qk_head_dim, :] + v_up_weight = kv_up_weight[:, self.config.qk_head_dim :, :] + return k_up_weight, v_up_weight + + def _combine_split_kv_up_weights( + self, k_up_weight: torch.Tensor, v_up_weight: torch.Tensor + ) -> torch.Tensor: + """Combine pre-refactor split K/V up-projection weights into the new layout.""" + num_heads = self.num_attention_heads_per_partition + qk_head_dim = self.config.qk_head_dim + v_head_dim = self.config.v_head_dim + kv_lora_rank = self.config.kv_lora_rank + + k_up_weight = k_up_weight.view(num_heads, qk_head_dim, kv_lora_rank) + v_up_weight = v_up_weight.view(num_heads, v_head_dim, kv_lora_rank) + return ( + torch.cat((k_up_weight, v_up_weight), dim=1) + .contiguous() + .view(num_heads * (qk_head_dim + v_head_dim), kv_lora_rank) + ) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Load checkpoints saved with either combined or split K/V up-projection weights.""" + combined_key = f"{prefix}linear_kv_up_proj.weight" + k_up_key = f"{prefix}linear_k_up_proj.weight" + v_up_key = f"{prefix}linear_v_up_proj.weight" + if combined_key not in state_dict and k_up_key in state_dict and v_up_key in state_dict: + state_dict[combined_key] = self._combine_split_kv_up_weights( + state_dict.pop(k_up_key), state_dict.pop(v_up_key) + ) + + combined_extra_state_key = f"{prefix}linear_kv_up_proj._extra_state" + k_up_extra_state_key = f"{prefix}linear_k_up_proj._extra_state" + v_up_extra_state_key = f"{prefix}linear_v_up_proj._extra_state" + if k_up_extra_state_key in state_dict or v_up_extra_state_key in state_dict: + k_extra_state = state_dict.pop(k_up_extra_state_key, None) + v_extra_state = state_dict.pop(v_up_extra_state_key, None) + if combined_extra_state_key not in state_dict: + state_dict[combined_extra_state_key] = ( + k_extra_state if k_extra_state is not None else v_extra_state + ) + + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + def _checkpointed_attention_forward( self, q_absorbed, k_compressed, - v_compressed, hidden_states, q_compressed, attention_mask, - rotary_pos_emb=None, + up_v_weight, + position_ids=None, attn_mask_type=None, - attention_bias=None, packed_seq_params=None, ): """Forward method with selective activation checkpointing.""" @@ -661,23 +746,22 @@ def _checkpointed_attention_forward( def custom_forward(*inputs): q_absorbed = inputs[0] k_compressed = inputs[1] - v_compressed = inputs[2] - hidden_states = inputs[3] - q_compressed = inputs[4] - attention_mask = inputs[5] - attn_mask_type = inputs[7] - attention_bias = inputs[8] - packed_seq_params = inputs[9] + hidden_states = inputs[2] + q_compressed = inputs[3] + attention_mask = inputs[4] + up_v_weight = inputs[5] + attn_mask_type = inputs[6] attn_mask_type = AttnMaskType(attn_mask_type.item()) output_ = self.core_attention( q_absorbed, k_compressed, - v_compressed, - hidden_states, - q_compressed, + None, attention_mask, + x=hidden_states, + qr=q_compressed, + up_v_weight=up_v_weight, + position_ids=position_ids, attn_mask_type=attn_mask_type, - attention_bias=attention_bias, packed_seq_params=packed_seq_params, ) return output_ @@ -690,14 +774,11 @@ def custom_forward(*inputs): False, q_absorbed, k_compressed, - v_compressed, hidden_states, q_compressed, attention_mask, - rotary_pos_emb, + up_v_weight, attn_mask_type, - attention_bias, - packed_seq_params, ) return hidden_states @@ -714,6 +795,7 @@ def forward( rotary_pos_cos_sin=None, attention_bias=None, packed_seq_params=None, + position_ids=None, sequence_len_offset=None, *, inference_params=None, @@ -742,6 +824,7 @@ def forward( assert q_absorbed.is_contiguous() assert q_compressed.is_contiguous() assert kv_compressed.is_contiguous() + v_up_weight = self._get_v_up_weight() # ================================== # Core attention computation @@ -750,10 +833,11 @@ def forward( core_attn_out = self._checkpointed_attention_forward( q_absorbed, kv_compressed, - None, hidden_states, q_compressed, attention_mask, + v_up_weight, + position_ids=position_ids, packed_seq_params=packed_seq_params, ) else: @@ -761,9 +845,11 @@ def forward( q_absorbed, kv_compressed, None, - hidden_states, - q_compressed, attention_mask, + x=hidden_states, + qr=q_compressed, + up_v_weight=v_up_weight, + position_ids=position_ids, packed_seq_params=packed_seq_params, attn_mask_type=self.attn_mask_type, ) @@ -771,24 +857,21 @@ def forward( # ================================== # Apply V up projection # ================================== - assert self.linear_v_up_proj.weight.size(0) == ( - self.num_attention_heads_per_partition * self.config.v_head_dim - ) - assert self.linear_v_up_proj.weight.size(1) == self.config.kv_lora_rank - v_up_weight = self.linear_v_up_proj.weight.view( - self.num_attention_heads_per_partition, self.config.v_head_dim, self.config.kv_lora_rank + core_consumed_v_up_projection = getattr( + self.core_attention, "consumes_absorbed_v_up_projection", False ) - core_attn_out = core_attn_out.view( - *core_attn_out.shape[:-1], + core_attn_out = _apply_absorbed_v_up_projection( + core_attn_out, + v_up_weight, self.num_attention_heads_per_partition, self.config.kv_lora_rank, + self.config.v_head_dim, + core_consumed_v_up_projection, ) - core_attn_out = torch.einsum("...nc,ndc->...nd", core_attn_out, v_up_weight) - core_attn_out = core_attn_out.contiguous() - core_attn_out = core_attn_out.view(*core_attn_out.shape[:-2], -1) - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - core_attn_out = core_attn_out.unsqueeze(1) + core_attn_out = _restore_packed_thd_batch_dim( + core_attn_out, hidden_states, packed_seq_params + ) assert core_attn_out.ndim == hidden_states.ndim assert core_attn_out.shape[0] == ( @@ -823,8 +906,7 @@ def backward_dw(self) -> NoReturn: def _backward_kv_proj(self): """Computes weight gradients of KV projection layers.""" - self.linear_k_up_proj.backward_dw() - self.linear_v_up_proj.backward_dw() + self.linear_kv_up_proj.backward_dw() self.linear_kv_down_proj.backward_dw() def _backward_q_proj(self): @@ -854,115 +936,3 @@ def clip_qk(self): function after Muon optimizer step. """ raise NotImplementedError("clip_qk is not implemented for AbsorbedMLA") - - def _combine_kv_weights(self, k_weight, v_weight): - """Combine separate K and V weights into MLA's interleaved format. - - MLA's linear_kv_up_proj weight layout (per head interleaved): - [head0_K, head0_V, head1_K, head1_V, ...] - - AbsorbedMLA's separate weights layout: - K: [head0_K, head1_K, ...] - V: [head0_V, head1_V, ...] - - This method interleaves K and V per head to match MLA's format. - - Args: - k_weight: [num_heads_per_partition * qk_head_dim, kv_lora_rank] - v_weight: [num_heads_per_partition * v_head_dim, kv_lora_rank] - - Returns: - combined: [num_heads_per_partition * (qk_head_dim + v_head_dim), kv_lora_rank] - """ - n = self.num_attention_heads_per_partition - qk_dim = self.config.qk_head_dim - v_dim = self.config.v_head_dim - lora_rank = self.config.kv_lora_rank - - # Reshape to per-head format - k_per_head = k_weight.view(n, qk_dim, lora_rank) - v_per_head = v_weight.view(n, v_dim, lora_rank) - - # Concatenate K and V for each head along dim=1 - # Result: [n, qk_dim + v_dim, lora_rank] - combined_per_head = torch.cat([k_per_head, v_per_head], dim=1) - - # Reshape back to linear weight format - combined_weight = combined_per_head.view(n * (qk_dim + v_dim), lora_rank) - - return combined_weight - - def _split_kv_weights(self, combined_weight): - """Split MLA's interleaved KV weight into separate K and V weights. - - MLA's linear_kv_up_proj weight layout (per head interleaved): - [head0_K, head0_V, head1_K, head1_V, ...] - - This method extracts K and V into separate tensors: - K: [head0_K, head1_K, ...] - V: [head0_V, head1_V, ...] - - Args: - combined_weight: [num_heads_per_partition * (qk_head_dim + v_head_dim), kv_lora_rank] - - Returns: - k_weight: [num_heads_per_partition * qk_head_dim, kv_lora_rank] - v_weight: [num_heads_per_partition * v_head_dim, kv_lora_rank] - """ - n = self.num_attention_heads_per_partition - qk_dim = self.config.qk_head_dim - v_dim = self.config.v_head_dim - lora_rank = self.config.kv_lora_rank - - # Reshape to per-head format - combined_per_head = combined_weight.view(n, qk_dim + v_dim, lora_rank) - - # Split K and V for each head (slicing creates non-contiguous views) - k_per_head = combined_per_head[:, :qk_dim, :] # [n, qk_dim, lora_rank] - v_per_head = combined_per_head[:, qk_dim:, :] # [n, v_dim, lora_rank] - - # Make contiguous and reshape back to linear weight format - k_weight = k_per_head.contiguous().view(n * qk_dim, lora_rank) - v_weight = v_per_head.contiguous().view(n * v_dim, lora_rank) - - return k_weight, v_weight - - def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): - """Handle loading from checkpoints with combined KV up projection weights. - - This method splits the combined 'linear_kv_up_proj.weight' (which has per-head - interleaved K and V) into separate 'linear_k_up_proj.weight' and 'linear_v_up_proj.weight'. - """ - combined_key = f'{prefix}linear_kv_up_proj.weight' - k_up_key = f'{prefix}linear_k_up_proj.weight' - v_up_key = f'{prefix}linear_v_up_proj.weight' - - # Split combined KV weights into separate K and V - if combined_key in state_dict: - combined_weight = state_dict[combined_key] - - # Split with proper per-head de-interleaving - k_weight, v_weight = self._split_kv_weights(combined_weight) - - state_dict[k_up_key] = k_weight - state_dict[v_up_key] = v_weight - - del state_dict[combined_key] - - combined_extra_state_key = f'{prefix}linear_kv_up_proj._extra_state' - k_up_extra_state_key = f'{prefix}linear_k_up_proj._extra_state' - v_up_extra_state_key = f'{prefix}linear_v_up_proj._extra_state' - - if combined_extra_state_key in state_dict: - combined_extra_state = state_dict[combined_extra_state_key] - - assert isinstance(combined_extra_state, torch.Tensor) - # Now we can only handle the case where the extra state is empty. - assert combined_extra_state.numel() == 0 - - state_dict[k_up_extra_state_key] = combined_extra_state.clone() - state_dict[v_up_extra_state_key] = combined_extra_state.clone() - - del state_dict[combined_extra_state_key] - - super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py index eb235501ad7..aa2c0a1479f 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import random +from types import SimpleNamespace from typing import List, Optional, Tuple import pytest @@ -12,6 +13,9 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant import ( + absorbed_mla as absorbed_mla_module, +) from megatron.core.transformer.experimental_attention_variant.absorbed_mla import ( AbsorbedMLASelfAttention, AbsorbedMLASelfAttentionSubmodules, @@ -193,8 +197,7 @@ def get_absorbed_mla_submodules( linear_q_down_proj=linear_q_down_proj, linear_q_up_proj=backend.column_parallel_linear(), linear_kv_down_proj=linear_kv_down_proj, - linear_k_up_proj=backend.column_parallel_linear(), - linear_v_up_proj=backend.column_parallel_linear(), + linear_kv_up_proj=backend.column_parallel_linear(), core_attention=MockCoreAttention, linear_proj=backend.row_parallel_linear(), q_layernorm=qk_norm, @@ -227,6 +230,96 @@ def get_mla_submodules( ) +def test_checkpointed_attention_forward_captures_metadata(monkeypatch): + """Optional metadata should stay in the closure instead of checkpoint tensor args.""" + + packed_seq_params = PackedSeqParams(qkv_format='thd') + checkpoint_args = None + + def fake_checkpoint(run_function, distribute_saved_activations, *args): + nonlocal checkpoint_args + del distribute_saved_activations + checkpoint_args = args + assert all(torch.is_tensor(arg) for arg in args) + return run_function(*args) + + class CoreAttention(torch.nn.Module): + def forward(self, query, key, value, attention_mask, **kwargs): + del query, key, value, attention_mask + assert kwargs["packed_seq_params"] is packed_seq_params + assert kwargs["position_ids"] is None + return kwargs["x"] + + dummy_attention = type( + "DummyAttention", + (), + {"attn_mask_type": AttnMaskType.causal, "core_attention": CoreAttention()}, + )() + + monkeypatch.setattr(absorbed_mla_module.tensor_parallel, "checkpoint", fake_checkpoint) + + hidden_states = torch.randn(4, 1, 8) + output = AbsorbedMLASelfAttention._checkpointed_attention_forward( + dummy_attention, + q_absorbed=torch.randn(4, 1, 2, 8), + k_compressed=torch.randn(4, 1, 1, 8), + hidden_states=hidden_states, + q_compressed=torch.randn(4, 1, 8), + attention_mask=torch.empty(1), + up_v_weight=torch.randn(2, 4, 4), + position_ids=None, + packed_seq_params=packed_seq_params, + ) + + assert checkpoint_args is not None + assert all(arg is not packed_seq_params for arg in checkpoint_args) + assert all(arg is not None for arg in checkpoint_args) + assert output is hidden_states + + +def test_load_from_state_dict_combines_split_kv_up_projection(monkeypatch): + """Pre-refactor split K/V up-projection checkpoints should load into the combined layout.""" + + dummy_attention = object.__new__(AbsorbedMLASelfAttention) + dummy_attention.num_attention_heads_per_partition = 2 + dummy_attention.config = SimpleNamespace(qk_head_dim=2, v_head_dim=3, kv_lora_rank=4) + + prefix = "self_attention." + k_weight = torch.arange(2 * 2 * 4, dtype=torch.float32).view(2 * 2, 4) + v_weight = torch.arange(2 * 3 * 4, dtype=torch.float32).view(2 * 3, 4) + state_dict = { + f"{prefix}linear_k_up_proj.weight": k_weight.clone(), + f"{prefix}linear_v_up_proj.weight": v_weight.clone(), + f"{prefix}linear_k_up_proj._extra_state": torch.empty(0), + f"{prefix}linear_v_up_proj._extra_state": torch.empty(0), + } + captured_state_dict = {} + + def fake_super_load(self, state_dict, *args, **kwargs): + del self, args, kwargs + captured_state_dict.update(state_dict) + + monkeypatch.setattr(absorbed_mla_module.Attention, "_load_from_state_dict", fake_super_load) + + AbsorbedMLASelfAttention._load_from_state_dict( + dummy_attention, state_dict, prefix, {}, True, [], [], [] + ) + + expected_weight = ( + torch.cat((k_weight.view(2, 2, 4), v_weight.view(2, 3, 4)), dim=1) + .contiguous() + .view(2 * (2 + 3), 4) + ) + torch.testing.assert_close( + captured_state_dict[f"{prefix}linear_kv_up_proj.weight"], expected_weight + ) + assert f"{prefix}linear_k_up_proj.weight" not in captured_state_dict + assert f"{prefix}linear_v_up_proj.weight" not in captured_state_dict + assert f"{prefix}linear_kv_up_proj._extra_state" in captured_state_dict + assert f"{prefix}linear_k_up_proj._extra_state" not in captured_state_dict + assert f"{prefix}linear_v_up_proj._extra_state" not in captured_state_dict + + @pytest.mark.parametrize("tp_cp", [[1, 1], [2, 1], [1, 2], [2, 2]]) @pytest.mark.parametrize("qkv_format", ['sbhd', 'thd']) @pytest.mark.parametrize("down_proj_use_column_parallel", [False, True]) @@ -350,50 +443,17 @@ def _calculate_tensor_similarity(x, y): absorbed_grads = dict(absorbed_mla.named_parameters()) standard_grads = dict(standard_mla.named_parameters()) - # Map parameter names between absorbed and standard MLA - # Most parameters have the same name, except for K/V up proj for name, param in standard_grads.items(): - if 'linear_kv_up_proj' in name: - # Special handling: combine k and v up proj grads from absorbed_mla - k_name = name.replace('linear_kv_up_proj', 'linear_k_up_proj') - v_name = name.replace('linear_kv_up_proj', 'linear_v_up_proj') - - k_grad = absorbed_grads[k_name].grad - v_grad = absorbed_grads[v_name].grad - - # Combine k and v grads (interleaved by head) - # k_grad: [n * qk_head_dim, kv_lora_rank] - # v_grad: [n * v_head_dim, kv_lora_rank] - # combined: [n * (qk_head_dim + v_head_dim), kv_lora_rank] - n_heads = absorbed_mla.num_attention_heads_per_partition - qk_head_dim = absorbed_mla.config.qk_head_dim - v_head_dim = absorbed_mla.config.v_head_dim - kv_lora_rank = absorbed_mla.config.kv_lora_rank - - k_grad_3d = k_grad.view(n_heads, qk_head_dim, kv_lora_rank) - v_grad_3d = v_grad.view(n_heads, v_head_dim, kv_lora_rank) - combined_grad_3d = torch.cat([k_grad_3d, v_grad_3d], dim=1) - combined_grad = combined_grad_3d.view(-1, kv_lora_rank) - - absorbed_grad_flat = combined_grad.flatten().float() - standard_grad_flat = param.grad.flatten().float() - - cos_sim = torch.nn.functional.cosine_similarity( - absorbed_grad_flat.unsqueeze(0), standard_grad_flat.unsqueeze(0) - ).item() - assert cos_sim > 0.9999, f"name: {name}, cosine similarity = {cos_sim} < 0.9999" - assert _calculate_tensor_similarity(combined_grad, param.grad) > 0.9999 - else: - absorbed_grad = absorbed_grads[name].grad - standard_grad = param.grad + absorbed_grad = absorbed_grads[name].grad + standard_grad = param.grad - absorbed_grad_flat = absorbed_grad.flatten().float() - standard_grad_flat = standard_grad.flatten().float() + absorbed_grad_flat = absorbed_grad.flatten().float() + standard_grad_flat = standard_grad.flatten().float() - cos_sim = torch.nn.functional.cosine_similarity( - absorbed_grad_flat.unsqueeze(0), standard_grad_flat.unsqueeze(0) - ).item() - assert cos_sim > 0.9999, f"name: {name}, cosine similarity = {cos_sim} < 0.9999" - assert _calculate_tensor_similarity(absorbed_grad, standard_grad) > 0.9999 + cos_sim = torch.nn.functional.cosine_similarity( + absorbed_grad_flat.unsqueeze(0), standard_grad_flat.unsqueeze(0) + ).item() + assert cos_sim > 0.9999, f"name: {name}, cosine similarity = {cos_sim} < 0.9999" + assert _calculate_tensor_similarity(absorbed_grad, standard_grad) > 0.9999 Utils.destroy_model_parallel() From 24b6b7c9c62e1128855977ff83af644ac521ccf8 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Tue, 16 Jun 2026 10:01:50 -0700 Subject: [PATCH 2/4] address review comments Signed-off-by: Hollow Man --- .../experimental_attention_variant/absorbed_mla.py | 8 ++++---- .../experimental_attention_variant/test_absorbed_mla.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index f8f4c37ea7c..fccf674d785 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -755,8 +755,8 @@ def custom_forward(*inputs): output_ = self.core_attention( q_absorbed, k_compressed, - None, - attention_mask, + value=None, + attention_mask=attention_mask, x=hidden_states, qr=q_compressed, up_v_weight=up_v_weight, @@ -844,8 +844,8 @@ def forward( core_attn_out = self.core_attention( q_absorbed, kv_compressed, - None, - attention_mask, + value=None, + attention_mask=attention_mask, x=hidden_states, qr=q_compressed, up_v_weight=v_up_weight, diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py index aa2c0a1479f..527591c11d8 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py @@ -244,7 +244,7 @@ def fake_checkpoint(run_function, distribute_saved_activations, *args): return run_function(*args) class CoreAttention(torch.nn.Module): - def forward(self, query, key, value, attention_mask, **kwargs): + def forward(self, query, key, *, value, attention_mask, **kwargs): del query, key, value, attention_mask assert kwargs["packed_seq_params"] is packed_seq_params assert kwargs["position_ids"] is None @@ -277,8 +277,8 @@ def forward(self, query, key, value, attention_mask, **kwargs): assert output is hidden_states -def test_load_from_state_dict_combines_split_kv_up_projection(monkeypatch): - """Pre-refactor split K/V up-projection checkpoints should load into the combined layout.""" +def test_load_from_state_dict_backwards_compatible_with_split_kv_up_projection(monkeypatch): + """Pre-refactor split K/V up-projection checkpoints load into the combined layout.""" dummy_attention = object.__new__(AbsorbedMLASelfAttention) dummy_attention.num_attention_heads_per_partition = 2 From e1d16e8fda3327cf68660e86a37a801bb1388d23 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Tue, 16 Jun 2026 17:19:47 -0700 Subject: [PATCH 3/4] Fix test case Signed-off-by: Hollow Man --- .../experimental_attention_variant/test_absorbed_mla.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py index 527591c11d8..edaa7be37e0 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py @@ -41,7 +41,7 @@ def __init__(self, *args, **kwargs): self.pg_collection = kwargs.get("pg_collection") def forward( - self, q, k, v, *args, packed_seq_params: Optional[PackedSeqParams] = None, **kwargs + self, q, k, v=None, *args, packed_seq_params: Optional[PackedSeqParams] = None, **kwargs ): """Mock forward pass.""" if packed_seq_params is None: From 153658327400d620ef12d7907de08d63779124df Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Tue, 16 Jun 2026 20:53:57 -0700 Subject: [PATCH 4/4] fix test cases Signed-off-by: Hollow Man --- .../test_fine_grained_activation_offloading.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py index 4b68d4a48b5..f92485084b5 100644 --- a/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py +++ b/tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py @@ -471,6 +471,12 @@ def _run_schedule_1f1b_two_microbatches( if enable_offload_reset: off_interface.reset() + # Keep warmup-created grad buffers resident for stable peak-memory comparisons, + # but clear warmup values before capturing correctness grads. + for p in model.parameters(): + if p.grad is not None: + p.grad.zero_() + data0 = _make_schedule_inputs() data1 = _make_schedule_inputs() plan0 = model.build_schedule_plan(**data0)