diff --git a/docs/models/mlm.md b/docs/models/mlm.md index f4a56f00..63d18881 100644 --- a/docs/models/mlm.md +++ b/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) | diff --git a/tests/models/mlm/test_collator_mlm.py b/tests/models/mlm/test_collator_mlm.py index b81d4102..dbb504ab 100644 --- a/tests/models/mlm/test_collator_mlm.py +++ b/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() diff --git a/xlm-models/mlm/configs/datamodule/tinygsm_mlm.yaml b/xlm-models/mlm/configs/datamodule/tinygsm_mlm.yaml index f2acde52..539cb156 100644 --- a/xlm-models/mlm/configs/datamodule/tinygsm_mlm.yaml +++ b/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 diff --git a/xlm-models/mlm/configs/experiment/tinygsm_mlm.yaml b/xlm-models/mlm/configs/experiment/tinygsm_mlm.yaml index e9488f3e..8c1efafd 100644 --- a/xlm-models/mlm/configs/experiment/tinygsm_mlm.yaml +++ b/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 diff --git a/xlm-models/mlm/datamodule_mlm.py b/xlm-models/mlm/datamodule_mlm.py index bcc236f8..349f52c0 100644 --- a/xlm-models/mlm/datamodule_mlm.py +++ b/xlm-models/mlm/datamodule_mlm.py @@ -263,6 +263,7 @@ 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 @@ -270,6 +271,10 @@ def prepare_prefix_suffix_ids( 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,6 +504,12 @@ 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, @@ -459,9 +517,10 @@ def __call__( 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