fix: pass DSA hidden states through selective attention recompute

src/mcore_bridge/model/modules/absorbed_mla.py

@@ -226,6 +226,49 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
 
         return q_absorbed, kv_compressed, q_compressed
 
+    def _checkpoint_absorbed_core_attention(self, q_absorbed, kv_compressed, absorbed_kwargs):
+        """Recompute absorbed DSA with the same arguments as the non-checkpoint path.
+
+        Tensor inputs go through ``tensor_parallel.checkpoint`` so the backward
+        recompute sees detached copies. ``packed_seq_params`` and ``attn_mask_type``
+        are not tensors; the base checkpoint closes over them the same way.
+        """
+        hidden_states = absorbed_kwargs['x']
+        q_compressed = absorbed_kwargs['qr']
+        attention_mask = absorbed_kwargs['attention_mask']
+        v_up_weight = absorbed_kwargs['up_v_weight']
+        position_ids = absorbed_kwargs['position_ids']
+        packed_seq_params = absorbed_kwargs['packed_seq_params']
+        attn_mask_type = absorbed_kwargs['attn_mask_type']
+        checkpoint_args = [q_absorbed, kv_compressed, hidden_states, q_compressed, v_up_weight]
+        include_mask = torch.is_tensor(attention_mask)
+        include_pos = torch.is_tensor(position_ids)
+        if include_mask:
+            checkpoint_args.append(attention_mask)
+        if include_pos:
+            checkpoint_args.append(position_ids)
+
+        def custom_forward(*inputs):
+            q, kv, x, qr, up_v, *optional = inputs
+            cursor = 0
+            mask = optional[cursor] if include_mask else attention_mask
+            cursor += int(include_mask)
+            pos = optional[cursor] if include_pos else position_ids
+            return self.core_attention(
+                q,
+                kv,
+                value=None,
+                attention_mask=mask,
+                x=x,
+                qr=qr,
+                up_v_weight=up_v,
+                position_ids=pos,
+                packed_seq_params=packed_seq_params,
+                attn_mask_type=attn_mask_type,
+            )
+
+        return tensor_parallel.checkpoint(custom_forward, False, *checkpoint_args)
+
     def forward(
         self,
         hidden_states,
@@ -284,30 +327,24 @@ def forward(
                                      f'"full" layer precedes it in this PP stage. Please adjust '
                                      f'`--pipeline_model_parallel_layout` to ensure each PP stage starts with '
                                      f'a "full" indexer layer. indexer_types: {self.config.hf_config.indexer_types}.')
+        absorbed_kwargs = dict(
+            value=None,
+            attention_mask=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,
+        )
         if self.checkpoint_core_attention and self.training:
-            core_attn_out = self._checkpointed_attention_forward(
-                q_absorbed,
-                kv_compressed,
-                hidden_states,
-                q_compressed,
-                attention_mask,
-                v_up_weight,
-                position_ids=position_ids,
-                packed_seq_params=packed_seq_params,
-            )
+            # The base checkpoint signature is (query, key, value, attention_mask, ...).
+            # Passing hidden_states/q_compressed positionally binds them to value and
+            # attention_mask, and `position_ids` is not a parameter, so selective
+            # core-attention recompute either raises or runs DSA without x/qr.
+            core_attn_out = self._checkpoint_absorbed_core_attention(q_absorbed, kv_compressed, absorbed_kwargs)
         else:
-            core_attn_out = self.core_attention(
-                q_absorbed,
-                kv_compressed,
-                value=None,
-                attention_mask=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,
-            )
+            core_attn_out = self.core_attention(q_absorbed, kv_compressed, **absorbed_kwargs)
 
         # ==================================
         # Apply V up projection

src/mcore_bridge/model/modules/multi_latent_attention.py

@@ -1,3 +1,4 @@
+import inspect
 import torch
 import torch.nn.functional as F
 from megatron.core import parallel_state, tensor_parallel
@@ -11,6 +12,41 @@
 
 class MLASelfAttention(McoreMLASelfAttention):
 
+    def _checkpoint_accepts_extra_kwargs(self) -> bool:
+        cached = getattr(self, '_ckpt_accepts_extra', None)
+        if cached is None:
+            params = inspect.signature(super()._checkpointed_attention_forward).parameters
+            cached = 'core_attention_extra_kwargs' in params
+            self._ckpt_accepts_extra = cached
+        return cached
+
+    def _checkpoint_mla_core_attention(self, query, key, value, attention_mask, packed_seq_params, attn_mask_type,
+                                       extra_kwargs):
+        if not extra_kwargs or self._checkpoint_accepts_extra_kwargs():
+            kwargs = {'packed_seq_params': packed_seq_params, 'attn_mask_type': attn_mask_type}
+            if extra_kwargs:
+                kwargs['core_attention_extra_kwargs'] = extra_kwargs
+            return self._checkpointed_attention_forward(query, key, value, attention_mask, **kwargs)
+
+        # mcore <= 0.16 checkpoints only q/k/v/mask and drops DSA's x/qr.
+        hidden_states = extra_kwargs['x']
+        q_compressed = extra_kwargs['qr']
+
+        def custom_forward(query, key, value, attention_mask, hidden_states, q_compressed):
+            return self.core_attention(
+                query,
+                key,
+                value,
+                attention_mask,
+                packed_seq_params=packed_seq_params,
+                attn_mask_type=attn_mask_type,
+                x=hidden_states,
+                qr=q_compressed,
+            )
+
+        return tensor_parallel.checkpoint(custom_forward, False, query, key, value, attention_mask, hidden_states,
+                                          q_compressed)
+
     def get_query_key_value_tensors(
         self,
         hidden_states,
@@ -247,22 +283,30 @@ def forward(
         if thd_qkv_format and query.shape[-1] != v_dim:
             value = F.pad(value, [0, query.shape[-1] - v_dim])
             self.core_attention.hidden_size_per_attention_head_v = value.shape[-1]
+        extra_kwargs = {}
+        if self.config.experimental_attention_variant == 'dsa':
+            if packed_seq_params is not None or self.config.context_parallel_size > 1:
+                raise ImportError('Please install the megatron-core main branch to support `DSAttention` '
+                                  'padding_free/context parallelism: '
+                                  '`pip install git+https://github.com/NVIDIA/Megatron-LM.git`')
+            # For dsa we need to pass in the original hidden states and the compressed
+            # query representation. Selective core-attention recompute has to see the
+            # same arguments; the base checkpoint call does not forward them.
+            extra_kwargs['x'] = hidden_states
+            extra_kwargs['qr'] = q_compressed
+            # for easy injection of rotary_pos_emb (patch)
+            packed_seq_params = (packed_seq_params, rotary_pos_emb)
         if self.checkpoint_core_attention and self.training:
-            core_attn_out = self._checkpointed_attention_forward(
-                query, key, value, attention_mask, packed_seq_params=packed_seq_params)
+            core_attn_out = self._checkpoint_mla_core_attention(
+                query,
+                key,
+                value,
+                attention_mask,
+                packed_seq_params,
+                attn_mask_type,
+                extra_kwargs,
+            )
         else:
-            extra_kwargs = {}
-            if self.config.experimental_attention_variant == 'dsa':
-                if packed_seq_params is not None or self.config.context_parallel_size > 1:
-                    raise ImportError('Please install the megatron-core main branch to support `DSAttention` '
-                                      'padding_free/context parallelism: '
-                                      '`pip install git+https://github.com/NVIDIA/Megatron-LM.git`')
-                # For dsa we need to pass in the original hidden states and the compressed
-                # query representation.
-                extra_kwargs['x'] = hidden_states
-                extra_kwargs['qr'] = q_compressed
-                # for easy injection of rotary_pos_emb (patch)
-                packed_seq_params = (packed_seq_params, rotary_pos_emb)
             core_attn_out = self.core_attention(
                 query,
                 key,

tests/test_absorbed_dsa_checkpoint.py

@@ -0,0 +1,99 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Selective recompute of absorbed DSA must forward x and qr as keywords."""
+import importlib.util
+import sys
+import torch
+import types
+from pathlib import Path
+
+
+def _load_absorbed():
+    root = Path(__file__).resolve().parents[1] / 'src' / 'mcore_bridge' / 'model' / 'modules' / 'absorbed_mla.py'
+
+    def module(name):
+        mod = types.ModuleType(name)
+        sys.modules[name] = mod
+        return mod
+
+    megatron = module('megatron')
+    core = module('megatron.core')
+    megatron.core = core
+
+    def checkpoint(function, distribute_saved_activations, *args):
+        return function(*args)
+
+    core.tensor_parallel = types.SimpleNamespace(checkpoint=checkpoint)
+    rope = module('megatron.core.models.common.embeddings.rope_utils')
+    rope.apply_rotary_pos_emb = lambda *args, **kwargs: None
+    for name in (
+            'megatron.core.models',
+            'megatron.core.models.common',
+            'megatron.core.models.common.embeddings',
+            'megatron.core.tensor_parallel',
+            'megatron.core.transformer',
+            'megatron.core.transformer.experimental_attention_variant',
+    ):
+        module(name)
+    mappings = module('megatron.core.tensor_parallel.mappings')
+    mappings.gather_from_sequence_parallel_region = lambda *args, **kwargs: None
+    mappings.gather_from_tensor_model_parallel_region = lambda *args, **kwargs: None
+    mappings.scatter_to_sequence_parallel_region = lambda *args, **kwargs: None
+    packed = module('megatron.core.packed_seq_params')
+    packed.PackedSeqParams = type('PackedSeqParams', (), {})
+    utils = module('megatron.core.utils')
+    utils.deprecate_inference_params = lambda *args, **kwargs: None
+
+    spec = importlib.util.spec_from_file_location('absorbed_mla_under_test', root)
+    loaded = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(loaded)
+    return loaded
+
+
+def test_checkpoint_passes_hidden_and_compressed_query_as_keywords():
+    saved = sys.modules.copy()
+    try:
+        absorbed = _load_absorbed()
+        seen = {}
+
+        class Core:
+
+            def __call__(self, query, key, **kwargs):
+                seen['query'] = query
+                seen['key'] = key
+                seen['kwargs'] = kwargs
+                return query
+
+        layer = absorbed.AbsorbedMLASelfAttention()
+        layer.core_attention = Core()
+        layer.attn_mask_type = 'causal'
+        q = torch.tensor([1.0])
+        kv = torch.tensor([2.0])
+        hidden = torch.tensor([3.0])
+        qr = torch.tensor([4.0])
+        up = torch.tensor([5.0])
+        mask = torch.tensor([6.0])
+        pos = torch.tensor([7])
+        packed = object()
+        absorbed.tensor_parallel.checkpoint = lambda function, distribute, *args: function(*args)
+        result = layer._checkpoint_absorbed_core_attention(
+            q, kv, {
+                'x': hidden,
+                'qr': qr,
+                'attention_mask': mask,
+                'up_v_weight': up,
+                'position_ids': pos,
+                'packed_seq_params': packed,
+                'attn_mask_type': 'causal',
+            })
+    finally:
+        sys.modules.clear()
+        sys.modules.update(saved)
+    assert torch.equal(result, q)
+    assert torch.equal(seen['query'], q)
+    assert torch.equal(seen['key'], kv)
+    assert seen['kwargs']['value'] is None
+    assert torch.equal(seen['kwargs']['x'], hidden)
+    assert torch.equal(seen['kwargs']['qr'], qr)
+    assert torch.equal(seen['kwargs']['up_v_weight'], up)
+    assert torch.equal(seen['kwargs']['attention_mask'], mask)
+    assert seen['kwargs']['packed_seq_params'] is packed