fix: zero left-padded tokens before the GDN recurrence

src/mcore_bridge/model/modules/gated_delta_net.py

@@ -33,6 +33,30 @@
     _GatedDeltaNet = object
 
 
+def gdn_padding_keep(attention_mask, seq_len: int, batch: int):
+    """Return ``[seq, batch, 1]`` with 0 on padding positions, or None.
+
+    Megatron's padding mask is ``attention_mask.all(dim=(1, 2))`` on a
+    ``[batch, 1, seq, seq]`` mask, True where a key is padding. Left padding
+    has to be zero before the causal conv and the delta-rule update. The loss
+    mask does not clear that state. A mask whose length does not match the
+    hidden states is left alone rather than applied to the wrong axis.
+    """
+    if isinstance(attention_mask, dict):
+        attention_mask = attention_mask.get('full_attention')
+    if not torch.is_tensor(attention_mask):
+        return None
+    if attention_mask.dim() == 4 and attention_mask.shape[0] == batch:
+        pad = attention_mask.bool().all(dim=(1, 2))
+    elif attention_mask.dim() == 2 and tuple(attention_mask.shape) == (batch, seq_len):
+        pad = attention_mask.bool()
+    else:
+        return None
+    if tuple(pad.shape) != (batch, seq_len) or not bool(pad.any()):
+        return None
+    return (~pad).transpose(0, 1).unsqueeze(-1)
+
+
 # Code borrowed from NVIDIA/Megatron-LM
 def _unpack_sequence(x, cu_seqlens, dim=1):
     unpacked_x = []
@@ -209,8 +233,11 @@ def forward(
             (Tuple[Tensor, Tensor]) GDN output and bias.
 
         """
-        # TODO: Deal with attention_mask (There is an issue when left padding is used.)
         inference_context = deprecate_inference_params(inference_context, inference_params)
+        if packed_seq_params is None:
+            keep = gdn_padding_keep(attention_mask, hidden_states.shape[0], hidden_states.shape[1])
+            if keep is not None:
+                hidden_states = hidden_states * keep.to(dtype=hidden_states.dtype, device=hidden_states.device)
 
         use_sp = self.config.sequence_parallel and self.tp_size > 1
         tp_group = self.pg_collection.tp

tests/test_gdn_padding_keep.py

@@ -0,0 +1,32 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Left-padding positions are zeroed before the GDN recurrence."""
+import ast
+import torch
+from pathlib import Path
+
+
+def _keep():
+    path = Path(__file__).resolve().parents[1] / 'src/mcore_bridge/model/modules/gated_delta_net.py'
+    tree = ast.parse(path.read_text())
+    fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == 'gdn_padding_keep')
+    module = ast.Module(body=[fn], type_ignores=[])
+    ast.fix_missing_locations(module)
+    namespace = {'torch': torch}
+    exec(compile(module, str(path), 'exec'), namespace)
+    return namespace['gdn_padding_keep']
+
+
+def test_left_pad_column_is_dropped():
+    keep = _keep()
+    # [batch, 1, seq, seq], True = masked. Batch 0 has the first key padded.
+    mask = torch.zeros(1, 1, 3, 3, dtype=torch.bool)
+    mask[0, :, :, 0] = True
+    out = keep(mask, seq_len=3, batch=1)
+    assert out.shape == (3, 1, 1)
+    assert torch.equal(out.squeeze(), torch.tensor([0, 1, 1]))
+
+
+def test_unrelated_shape_is_ignored():
+    keep = _keep()
+    mask = torch.ones(2, 3, dtype=torch.bool)
+    assert keep(mask, seq_len=4, batch=2) is None