Merge branch 'fix/mlm-seq2seq-suffix-window'

Author: dhruvdcoder

docs/models/mlm.md

@@ -138,11 +138,10 @@ Task dataset and preprocessing: [TinyGSM](../tasks/tinygsm.md). GSM8K and code-e
 | Setting | Value |
 |---------|--------|
 | 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) |
+| `block_size` | 256 |
+| `input_block_size` | 256 |
 | Batching | Per-device 32; global 512 |
-| Collators | STAR seq2seq (`seq2seq_*` / `seq2seq_pred_*`); no BOS; train/val `truncate: null` |
+| Collators | STAR seq2seq (`seq2seq_*` / `seq2seq_pred_*`); no BOS |
 | 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

@@ -189,38 +189,3 @@ def test_star_hidden_tail_and_suffix_slot(self, simple_tokenizer):
         # 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()

tests/models/mlm/test_predictor_mlm.py

@@ -100,3 +100,36 @@ def test_both_top_k_and_top_p_raises(
                 top_k=5,
                 top_p=0.9,
             )
+
+
+class TestMLMPredictorSeq2SeqMaskInjection:
+    """Seq2seq inference appends exactly ``max_new_tokens`` [MASK] slots."""
+
+    def test_appends_block_size_masks_after_prefix(
+        self, tiny_mlm_model, simple_tokenizer, dummy_noise_schedule
+    ):
+        block_size = 8
+        input_block_size = 16
+        predictor = MLMPredictor(
+            max_steps=block_size,
+            max_new_tokens=block_size,
+            tokenizer=simple_tokenizer,
+            model=tiny_mlm_model,
+            noise_schedule=dummy_noise_schedule,
+            top_k=1,
+        )
+        batch = {
+            "input_ids": torch.tensor(
+                [[0, 0, 0, 10, 11, 12, 13, 14, 15, 16, 0, 0, 0, 0, 0, 0]],
+                dtype=torch.long,
+            ),
+            "attention_mask": torch.tensor(
+                [[0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]],
+                dtype=torch.bool,
+            ),
+        }
+        with torch.no_grad():
+            preds = predictor.predict(batch)
+        assert preds["output_start_idx"] == input_block_size
+        assert preds["ids"].shape == (1, input_block_size + block_size)
+        assert preds["ids"].shape[-1] - preds["output_start_idx"] == block_size

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

@@ -17,12 +17,10 @@ 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

@@ -10,14 +10,10 @@ add_eos: true
 per_device_batch_size: 32
 global_batch_size: 512
 num_dataloader_workers: 5
-block_size: 512
-input_block_size: 0
+block_size: 256
+input_block_size: 256
 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

@@ -504,21 +504,15 @@ 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=max_seq_len,
-            truncate=truncate,
+            max_seq_len=(self.input_block_size + self.block_size),
+            truncate=self.truncate,
             loss_on_padding=self.loss_on_padding,
             suffix_block_size=self.block_size,
         )
@@ -677,12 +671,13 @@ def __call__(
         examples: List[Seq2SeqCollatorInput],
     ) -> MLMBatch:
         pf = self.prompt_field
+        prefix_ids_list = [e[pf] for e in examples]  # type: ignore[misc]
+        if self.add_bos:
+            prefix_ids_list = [
+                p + [self.tokenizer.bos_token_id] for p in prefix_ids_list
+            ]
         prefix = prepare_prefix_ids(
-            [
-                e[pf]  # type: ignore[misc]
-                + [self.tokenizer.bos_token_id] * int(self.add_bos)
-                for e in examples
-            ],
+            prefix_ids_list,
             self.tokenizer.pad_token_id,
             max_seq_len=self.input_block_size,
         )