Merge pull request #48 from dhruvdcoder/fix/mlm-seq2seq-suffix-window

Fix MLM seq2seq suffix window train/inference mismatch

Author: dhruvdcoder

docs/models/mlm.md

@@ -69,7 +69,7 @@ The packed FlexAttention variant uses `PackedFlexMLMBatch` (subset of the above)
 | Class | Input | Output batch | Special behavior |
 |---|---|---|---|
 | `DefaultMLMCollator` | `BaseCollatorInput` | `MLMBatch` | Pad-right to `block_size`, BOS/EOS optional, random MLM masking. |
-| `MLMSeq2SeqTrainCollator` | `Seq2SeqCollatorInput` | `MLMBatch` | Concatenates `[prompt][BOS][target][EOS]` with right padding; masks only suffix positions. |
+| `MLMSeq2SeqTrainCollator` | `Seq2SeqCollatorInput` | `MLMBatch` | Concatenates `[prompt][BOS][target][EOS]`; only `block_size` suffix positions after the prompt are visible and MLM-eligible (tail hidden via `attention_mask=0`). |
 | `MLMSeq2SeqCollator` | `Seq2SeqCollatorInput` | `MLMBatch` | Left-pads prompt and right-pads target separately (padding on both sides). |
 | `_MLMSeq2SeqPredCollator` | `Seq2SeqCollatorInput` | `MLMBatch` | Same as `MLMSeq2SeqCollator` but masks **all** suffix tokens (`mask_all=True`); used for exact-match eval. |
 | `MLMSeq2SeqPredCollator` | `Seq2SeqCollatorInput` | `MLMBatch` | `input_ids = left-padded prompt only`; `target_ids = right-padded target` (used for seq2seq prediction). |
@@ -140,8 +140,9 @@ Task dataset and preprocessing: [TinyGSM](../tasks/tinygsm.md). GSM8K and code-e
 | Tokenizer | Qwen2-0.5B (`Qwen/Qwen2-0.5B`) with added `<|mask|>` |
 | `block_size` | 512 |
 | `input_block_size` | 0 |
+| `model.max_length` | 2048 (prompt + `block_size` suffix; overrides default 516) |
 | Batching | Per-device 32; global 512 |
-| Collators | STAR seq2seq (`seq2seq_*` / `seq2seq_pred_*`); no BOS between question and code |
+| Collators | STAR seq2seq (`seq2seq_*` / `seq2seq_pred_*`); no BOS; train/val `truncate: null` |
 | Val / test prediction | Post-hoc `code_exec_accuracy` (`Gsm8kCodeEval`); token EM disabled |
 | Monitored metric | `val/lm/accumulated_loss` |
 | Training schedule | Up to 1M steps; validation every 50k steps; checkpoint every 2.5k steps (keep every 100k) |

tests/models/mlm/test_collator_mlm.py

@@ -7,6 +7,7 @@
     DefaultMLMCollator,
     MLMSeq2SeqPredCollator,
     prepare_prefix_ids,
+    prepare_prefix_suffix_ids,
 )
 from tests.models._base import BaseCollatorTests
 
@@ -143,3 +144,83 @@ def test_left_pads_to_max_seq_len(self, simple_tokenizer):
         assert (
             out["attention_mask"].sum(dim=1).tolist() == [3, 2]
         )
+
+
+class TestPreparePrefixSuffixIds:
+    """Suffix window cap for seq2seq training (STAR and TinyGSM layouts)."""
+
+    def test_star_hidden_tail_and_suffix_slot(self, simple_tokenizer):
+        torch.manual_seed(0)
+        pad = simple_tokenizer.pad_token_id
+        bos = simple_tokenizer.bos_token_id
+        eos = simple_tokenizer.eos_token_id
+        mask = simple_tokenizer.mask_token_id
+        input_block_size = 16
+        block_size = 8
+        max_seq_len = input_block_size + block_size
+
+        batch = prepare_prefix_suffix_ids(
+            prefix_ids=[[10, 11, 12], [13, 14]],
+            suffix_ids=[[20, 21], [22, 23, 24]],
+            pad_token_id=pad,
+            mask_token_id=mask,
+            eos_token_id=eos,
+            bos_token_id=bos,
+            max_seq_len=max_seq_len,
+            truncate="block",
+            suffix_block_size=block_size,
+        )
+
+        assert batch["input_ids"].shape == (2, max_seq_len)
+        # Example 0: P=3, BOS, suffix slot 8 -> visible_len=12
+        visible_len = 3 + 1 + block_size
+        assert batch["attention_mask"][0, :visible_len].all()
+        assert not batch["attention_mask"][0, visible_len:].any()
+        assert (batch["target_ids"][0, visible_len:] == -100).all()
+        # Prefix + BOS never MLM-masked
+        assert batch["input_ids"][0, 0] == 10
+        assert batch["input_ids"][0, 3] == bos
+        assert (batch["input_ids"][0, :4] != mask).all()
+        # Suffix slot layout in targets (input_ids may replace some with [MASK])
+        assert batch["target_ids"][0, 4] == 20
+        assert batch["target_ids"][0, 5] == 21
+        assert batch["target_ids"][0, 6] == eos
+        assert batch["target_ids"][0, 7] == pad
+        # MLM masks only in suffix slot
+        suffix_region = batch["input_ids"][0, 4:visible_len]
+        assert (suffix_region == mask).any()
+
+    def test_tinygsm_variable_prefix_no_shared_block(self, simple_tokenizer):
+        torch.manual_seed(0)
+        pad = simple_tokenizer.pad_token_id
+        eos = simple_tokenizer.eos_token_id
+        mask = simple_tokenizer.mask_token_id
+        block_size = 8
+
+        batch = prepare_prefix_suffix_ids(
+            prefix_ids=[[10, 11, 12, 13, 14], [20, 21]],
+            suffix_ids=[[30, 31], [40, 41, 42]],
+            pad_token_id=pad,
+            mask_token_id=mask,
+            eos_token_id=eos,
+            bos_token_id=None,
+            max_seq_len=None,
+            truncate=None,
+            suffix_block_size=block_size,
+        )
+
+        # Batch padded to max visible: 5 + 8 = 13 vs 2 + 8 = 10
+        assert batch["input_ids"].shape == (2, 13)
+        visible_len_0 = 5 + block_size
+        visible_len_1 = 2 + block_size
+        assert batch["attention_mask"][0, :visible_len_0].all()
+        assert not batch["attention_mask"][0, visible_len_0:].any()
+        assert batch["attention_mask"][1, :visible_len_1].all()
+        assert not batch["attention_mask"][1, visible_len_1:].any()
+        assert (batch["target_ids"][1, visible_len_1:] == -100).all()
+        # Suffix slot starts right after prompt (no BOS); check targets for layout
+        assert batch["target_ids"][0, 5] == 30
+        assert batch["target_ids"][0, 6] == 31
+        assert batch["target_ids"][0, 7] == eos
+        assert (batch["input_ids"][0, 5:visible_len_0] == mask).any()
+        assert (batch["input_ids"][0, :5] != mask).all()

xlm-models/mlm/configs/datamodule/tinygsm_mlm.yaml

@@ -14,10 +14,12 @@ datamodule:
       lm:
         collator:
           add_bos: false
+          truncate: null
     val:
       lm:
         collator:
           add_bos: false
+          truncate: null
       prediction:
         collator:
           add_bos: false

xlm-models/mlm/configs/experiment/tinygsm_mlm.yaml

@@ -14,6 +14,10 @@ block_size: 512
 input_block_size: 0
 monitored_metric: val/lm/accumulated_loss
 
+# prompt + block_size suffix can exceed default rotary max_length (516)
+model:
+  max_length: 2048
+
 global_components:
   tokenizer:
     _target_: xlm.datamodule.load_auto_tokenizer

xlm-models/mlm/datamodule_mlm.py

@@ -263,13 +263,18 @@ def prepare_prefix_suffix_ids(
     max_seq_len: Optional[int] = None,
     truncate: Literal["max", "block", None] = "block",
     loss_on_padding: bool = True,
+    suffix_block_size: Optional[int] = None,
 ) -> MLMBatch:
     """Prepare concatenated prefix and suffix ids for seq2seq tasks with padding on the right only
 
     Args:
         loss_on_padding: bool
           - If true, pad token is treated as a normal token: it has attention on it, it is predicted as a target token.
           - If false, it has no attention on it, it is not predicted as a target token (-100)
+        suffix_block_size: When set, reserve exactly this many suffix tokens after the prefix
+          (suffix + EOS padded/truncated to fit). Positions beyond ``prefix + BOS +
+          suffix_block_size`` are attention-hidden and excluded from MLM masking and loss,
+          matching inference ``max_new_tokens``.
     """
     input_ids: List[TT] = []
     attention_mask: List[TT] = []
@@ -280,10 +285,16 @@ def prepare_prefix_suffix_ids(
         bos_token_id is not None
     )  # always add bos before the suffix. Otherwise it is not needed.
     if truncate in ["max", None]:
-        max_len = max(
-            len(_prefix_ids) + len(_suffix_ids) + add_eos + add_bos
-            for _prefix_ids, _suffix_ids in zip(prefix_ids, suffix_ids)
-        )
+        if suffix_block_size is not None:
+            max_len = max(
+                len(_prefix_ids) + add_bos + suffix_block_size
+                for _prefix_ids, _suffix_ids in zip(prefix_ids, suffix_ids)
+            )
+        else:
+            max_len = max(
+                len(_prefix_ids) + len(_suffix_ids) + add_eos + add_bos
+                for _prefix_ids, _suffix_ids in zip(prefix_ids, suffix_ids)
+            )
         if truncate == "max" and max_seq_len is not None:
             max_len = max(max_len, max_seq_len)
     elif truncate == "block" and max_seq_len is not None:
@@ -296,58 +307,99 @@ def prepare_prefix_suffix_ids(
     for i, (_prefix_ids, _suffix_ids) in enumerate(
         zip(prefix_ids, suffix_ids)
     ):
-        # bos should not be masked
-        suffix_mask = pad_truncate_list(
-            [0] * (len(_prefix_ids) + add_bos)
-            + [1] * (len(_suffix_ids) + add_eos),
-            max_len,
-            1,
-            pad_left=False,
-        )
-        temp = pad_truncate_list(
-            _prefix_ids
-            + ([bos_token_id] * add_bos)
-            + _suffix_ids
-            + ([eos_token_id] * add_eos),
-            max_len,
-            pad_token_id,
-            pad_left=False,
-        )
-        _mask = (torch.rand(len(temp)) < t[i]).logical_and(
-            torch.tensor(suffix_mask, dtype=torch.bool)
-        )
-        _input_ids = torch.tensor(temp, dtype=torch.long)
-        input_ids.append(_input_ids)
-        if loss_on_padding:
-            attention_mask.append(
-                torch.tensor([1] * len(temp), dtype=torch.bool)
+        if suffix_block_size is not None:
+            suffix_slot = pad_truncate_list(
+                _suffix_ids + [eos_token_id] * add_eos,
+                suffix_block_size,
+                pad_token_id,
+                pad_left=False,
             )
-            target_ids.append(_input_ids.clone())
+            content = (
+                _prefix_ids + ([bos_token_id] * add_bos) + suffix_slot
+            )
+            visible_len = len(_prefix_ids) + add_bos + suffix_block_size
+            temp = pad_truncate_list(
+                content, max_len, pad_token_id, pad_left=False
+            )
+            effective_visible = min(visible_len, len(temp))
+            suffix_mask = pad_truncate_list(
+                [0] * (len(_prefix_ids) + add_bos)
+                + [1] * suffix_block_size,
+                max_len,
+                0,
+                pad_left=False,
+            )
+            _mask = (torch.rand(len(temp)) < t[i]).logical_and(
+                torch.tensor(suffix_mask, dtype=torch.bool)
+            )
+            _input_ids = torch.tensor(temp, dtype=torch.long)
+            input_ids.append(_input_ids)
+            _attn = torch.zeros(len(temp), dtype=torch.bool)
+            _attn[:effective_visible] = True
+            attention_mask.append(_attn)
+            _target_ids = _input_ids.clone()
+            _target_ids[effective_visible:] = -100
+            if not loss_on_padding:
+                content_len = (
+                    len(_prefix_ids) + add_bos + len(_suffix_ids) + add_eos
+                )
+                for j in range(min(content_len, effective_visible), effective_visible):
+                    _target_ids[j] = -100
+            target_ids.append(_target_ids)
             mask.append(_mask)
         else:
-            attention_mask.append(
-                torch.tensor(
-                    pad_truncate_list(
-                        [1]
-                        * (
-                            len(_prefix_ids)
-                            + len(_suffix_ids)
-                            + add_eos
-                            + add_bos
+            # bos should not be masked
+            suffix_mask = pad_truncate_list(
+                [0] * (len(_prefix_ids) + add_bos)
+                + [1] * (len(_suffix_ids) + add_eos),
+                max_len,
+                1,
+                pad_left=False,
+            )
+            temp = pad_truncate_list(
+                _prefix_ids
+                + ([bos_token_id] * add_bos)
+                + _suffix_ids
+                + ([eos_token_id] * add_eos),
+                max_len,
+                pad_token_id,
+                pad_left=False,
+            )
+            _mask = (torch.rand(len(temp)) < t[i]).logical_and(
+                torch.tensor(suffix_mask, dtype=torch.bool)
+            )
+            _input_ids = torch.tensor(temp, dtype=torch.long)
+            input_ids.append(_input_ids)
+            if loss_on_padding:
+                attention_mask.append(
+                    torch.tensor([1] * len(temp), dtype=torch.bool)
+                )
+                target_ids.append(_input_ids.clone())
+                mask.append(_mask)
+            else:
+                attention_mask.append(
+                    torch.tensor(
+                        pad_truncate_list(
+                            [1]
+                            * (
+                                len(_prefix_ids)
+                                + len(_suffix_ids)
+                                + add_eos
+                                + add_bos
+                            ),
+                            max_len,
+                            0,
+                            pad_left=False,
                         ),
-                        max_len,
-                        0,
-                        pad_left=False,
-                    ),
-                    dtype=torch.bool,
+                        dtype=torch.bool,
+                    )
                 )
-            )
-            mask.append(
-                _mask.logical_and(attention_mask[-1])
-            )  # no input masks in padding
-            _target_ids = _input_ids.clone()
-            _target_ids[~attention_mask[-1]] = -100  # no loss on padding
-            target_ids.append(_target_ids)
+                mask.append(
+                    _mask.logical_and(attention_mask[-1])
+                )  # no input masks in padding
+                _target_ids = _input_ids.clone()
+                _target_ids[~attention_mask[-1]] = -100  # no loss on padding
+                target_ids.append(_target_ids)
     target_ids = torch.stack(target_ids, dim=0)
     attention_mask = torch.stack(attention_mask, dim=0)
     input_ids = torch.stack(input_ids, dim=0)
@@ -452,16 +504,23 @@ def __call__(
         suffix_lists = [
             seq2seq_suffix_ids(e, self.target_field) for e in examples
         ]
+        if self.input_block_size == 0:
+            max_seq_len = None
+            truncate = None
+        else:
+            max_seq_len = self.input_block_size + self.block_size
+            truncate = self.truncate
         prefix_suffix = prepare_prefix_suffix_ids(
             [e["prompt_ids"] for e in examples],
             suffix_lists,
             self.tokenizer.pad_token_id,
             self.tokenizer.mask_token_id,
             eos_token_id=self.tokenizer.eos_token_id if self.add_eos else None,
             bos_token_id=self.tokenizer.bos_token_id if self.add_bos else None,
-            max_seq_len=(self.input_block_size + self.block_size),
-            truncate=self.truncate,
+            max_seq_len=max_seq_len,
+            truncate=truncate,
             loss_on_padding=self.loss_on_padding,
+            suffix_block_size=self.block_size,
         )
         return prefix_suffix