diff --git a/src/mcore_bridge/model/modules/gated_delta_net.py b/src/mcore_bridge/model/modules/gated_delta_net.py index 965e54b..e7e4a7a 100644 --- a/src/mcore_bridge/model/modules/gated_delta_net.py +++ b/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 diff --git a/tests/test_gdn_padding_keep.py b/tests/test_gdn_padding_keep.py new file mode 100644 index 0000000..5a75ba2 --- /dev/null +++ b/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