From c9696c867b5b1948eeff8234e2c29e98937503bb Mon Sep 17 00:00:00 2001 From: amazloumi Date: Tue, 21 Jul 2026 20:15:13 -0400 Subject: [PATCH 01/11] Fix - Shift VLM captioning labels for next-token loss --- CHANGELOG.md | 2 ++ kempnerforge/data/vlm_dataset.py | 17 +++++++++-------- tests/unit/test_vlm_dataset.py | 23 +++++++++++++++-------- 3 files changed, 26 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8eb55db..c05ca40 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -117,6 +117,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **VLM captioning labels are now next-token aligned.** `_tokenize_and_mask` labeled each position with its own input token (`labels[i] == input_ids[i]`); with the no-shift loss, each text logit was scored against the current, already-visible token, so the model learned to copy it and the captioning loss collapsed to `0`. Labels are now shifted (`labels[i] = input_ids[i+1]`, last real position `-100`), matching text pretraining; prompt-predicting positions stay masked. Fixes all VLM arches. + - `kempnerforge/data/vlm_dataset.py` (+ `tests/unit/test_vlm_dataset.py`). - **Resume silently reset AdamW optimizer momentum.** `CheckpointManager` round-tripped optimizer state through raw `optimizer.state_dict()` / `optimizer.load_state_dict()`. On resume the optimizer is freshly built, so its `state_dict()` is empty — `dcp.load` then had no `exp_avg` / `exp_avg_sq` tensors to fill, and the moments were silently dropped, resetting Adam momentum to zero at every resume point. Model weights, scheduler, dataloader position, and RNG all restored correctly; only the optimizer moments were lost, so resumed runs were not bit-exact. - `kempnerforge/checkpoint/manager.py`: save and load now go through DCP's `get_model_state_dict` / `get_optimizer_state_dict` / `set_model_state_dict` / `set_optimizer_state_dict`. The getters build a load template with the optimizer moments allocated in the correct FSDP/DTensor layout, so `dcp.load` repopulates them; the setters write the loaded values back into the live optimizer. - `docs/checkpointing/dcp-model.md`: updated the save/load snippets and the "shape to fill" explanation to the DCP-aware helpers. diff --git a/kempnerforge/data/vlm_dataset.py b/kempnerforge/data/vlm_dataset.py index 4d5a746..d39b03f 100644 --- a/kempnerforge/data/vlm_dataset.py +++ b/kempnerforge/data/vlm_dataset.py @@ -132,12 +132,10 @@ def _tokenize_and_mask( max_text_len: int, prompt: str | None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Tokenize and build right-padded input_ids + labels. + """Tokenize into right-padded input_ids + next-token labels. - When ``prompt`` is provided, the prompt portion of ``labels`` is - masked with ``-100`` (loss does not backpropagate through prompt - tokens). Padding positions in both ``input_ids`` and ``labels`` are - handled via ``ignore_index=-100`` on the loss. + ``labels[i]`` is token ``i+1`` (matching text pretraining and the no-shift + loss); pad, trailing, and prompt-predicting positions are ``-100``. BPE and SentencePiece tokenizers are NOT prefix-preserving: in general ``tokenize(prompt) + tokenize(text)`` differs from @@ -166,9 +164,12 @@ def _tokenize_and_mask( if n > 0: ids_tensor = torch.tensor(full_ids, dtype=torch.long) input_ids[:n] = ids_tensor - labels[:n] = ids_tensor - if prompt_len > 0: - labels[:prompt_len] = -100 + # Next-token targets: position i predicts token i+1; the last real token + # has no successor, and positions that would predict a prompt token stay + # masked so only response tokens are supervised. + labels[: n - 1] = ids_tensor[1:] + if prompt_len > 1: + labels[: prompt_len - 1] = -100 return input_ids, labels diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index d714050..70d0288 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -95,6 +95,7 @@ def test_shape_and_pad(self): assert labels.shape == (8,) assert ids[:3].tolist() == [1, 2, 3] assert ids[3].item() == 0 # pad + assert labels[:3].tolist() == [2, 3, -100] # next-token; last has no target assert labels[3].item() == -100 # pad masked def test_prompt_masks_labels(self): @@ -102,15 +103,18 @@ def test_prompt_masks_labels(self): ids, labels = _tokenize_and_mask(tok, text="xyz", max_text_len=8, prompt="ab") # Prompt "ab" = 2 tokens; target "xyz" = 3 tokens; total 5 tokens. assert ids[:5].tolist() == [1, 2, 24, 25, 26] - assert labels[:2].tolist() == [-100, -100] # prompt masked - assert labels[2:5].tolist() == [24, 25, 26] # targets not masked + # Next-token: last prompt token predicts the first target (24); earlier + # prompt-predicting positions are masked; last real token has no target. + assert labels[0].item() == -100 + assert labels[1:4].tolist() == [24, 25, 26] + assert labels[4].item() == -100 assert labels[5].item() == -100 # pad masked def test_truncation(self): tok = _MockTokenizer() ids, labels = _tokenize_and_mask(tok, text="abcdefghij", max_text_len=4, prompt=None) assert ids.tolist() == [1, 2, 3, 4] - assert labels.tolist() == [1, 2, 3, 4] + assert labels.tolist() == [2, 3, 4, -100] # next-token labels; final position has none def test_prompt_mask_with_bpe_tokenizer(self): """Regression: BPE (gpt2) and SentencePiece tokenizers are not @@ -136,10 +140,11 @@ def test_prompt_mask_with_bpe_tokenizer(self): # input_ids exactly matches the independent-concat form. assert ids[:n].tolist() == expected_ids - # Prompt portion of labels is -100. - assert (labels[: len(prompt_ids)] == -100).all() - # Target portion is the independently-tokenized text ids, byte-for-byte. - assert labels[len(prompt_ids) : n].tolist() == text_ids + # Positions predicting a prompt token are -100 (the last prompt token + # predicts the first target, so it stays supervised). + assert (labels[: len(prompt_ids) - 1] == -100).all() + # Target labels are the independently-tokenized text ids, byte-for-byte. + assert labels[len(prompt_ids) - 1 : n - 1].tolist() == text_ids def test_prompt_mask_under_bpe_merge_keeps_target_intact(self): """Concrete BPE case: ``tokenize("foo") + tokenize("bar")`` can @@ -163,7 +168,9 @@ def test_prompt_mask_under_bpe_merge_keeps_target_intact(self): assert ids[:split_total].tolist() == prompt_ids + text_ids # Target is exactly tokenize(text) — a regression would show up # here as a drift in the label ids (not just their length). - assert labels[len(prompt_ids) : len(prompt_ids) + len(text_ids)].tolist() == text_ids + assert ( + labels[len(prompt_ids) - 1 : len(prompt_ids) - 1 + len(text_ids)].tolist() == text_ids + ) # --------------------------------------------------------------------------- From 48ee0ead2fc0f8850c0dbdf9bf63d2eb400f7f3a Mon Sep 17 00:00:00 2001 From: amazloumi Date: Tue, 21 Jul 2026 21:28:47 -0400 Subject: [PATCH 02/11] Update video-dataset tests + docstrings for next-token labels --- kempnerforge/data/vlm_dataset.py | 5 ++--- tests/unit/test_video_dataset.py | 12 +++++++----- tests/unit/test_vlm_dataset.py | 4 ++-- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/kempnerforge/data/vlm_dataset.py b/kempnerforge/data/vlm_dataset.py index d39b03f..794bce4 100644 --- a/kempnerforge/data/vlm_dataset.py +++ b/kempnerforge/data/vlm_dataset.py @@ -142,9 +142,8 @@ def _tokenize_and_mask( ``tokenize(prompt + text)`` at the boundary (tokens can merge or split). To guarantee the mask lines up with the prompt boundary we tokenize prompt and text independently, then concatenate the id - lists. The mask length is ``len(prompt_ids)``, so masking - ``labels[:prompt_len]`` cannot leak a prompt token into supervision - or erase the first target token. + lists, so the prompt/target boundary is exact and only positions + predicting prompt tokens are masked. """ if prompt is not None: prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] diff --git a/tests/unit/test_video_dataset.py b/tests/unit/test_video_dataset.py index 1085b1d..5068b5a 100644 --- a/tests/unit/test_video_dataset.py +++ b/tests/unit/test_video_dataset.py @@ -126,8 +126,8 @@ def test_caption_is_supervised_when_frames_present(self, monkeypatch): monkeypatch.setattr(vd, "decode_video_frames", lambda *a, **k: _frames(4)) ds = _StubVideoDataset(["1"], ["abc"], max_frames=8, max_text_len=8) item = ds[0] - # "abc" -> ids 1,2,3 supervised; rest -100. - assert item["labels"][:3].tolist() == [1, 2, 3] + # "abc" -> ids 1,2,3; next-token labels [2, 3, -100], rest -100. + assert item["labels"][:3].tolist() == [2, 3, -100] assert (item["labels"][3:] == -100).all() def test_decode_failure_yields_zero_clip_no_loss(self, monkeypatch): @@ -152,10 +152,12 @@ def test_prompt_is_masked(self, monkeypatch): monkeypatch.setattr(vd, "decode_video_frames", lambda *a, **k: _frames(2)) ds = _StubVideoDataset(["1"], ["xyz"], max_frames=4, max_text_len=8, prompt="ab") item = ds[0] - # prompt "ab" (2 toks) masked; "xyz" (24,25,26) supervised. + # prompt "ab": the last prompt token predicts the first target (24); + # the earlier prompt-predicting position is masked. Then 25, 26, -100. assert item["input_ids"][:5].tolist() == [1, 2, 24, 25, 26] - assert item["labels"][:2].tolist() == [-100, -100] - assert item["labels"][2:5].tolist() == [24, 25, 26] + assert item["labels"][0].item() == -100 + assert item["labels"][1:4].tolist() == [24, 25, 26] + assert item["labels"][4].item() == -100 def test_len(self): ds = _StubVideoDataset(["1", "2", "3"], ["a", "b", "c"]) diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index 70d0288..baed861 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -122,8 +122,8 @@ def test_prompt_mask_with_bpe_tokenizer(self): from ``tokenize(prompt + text)`` at the boundary (tokens can merge or split). The implementation must tokenize the prompt and text independently and concatenate the id lists, so the - ``labels[:len(prompt_ids)]`` mask lines up exactly with the - prompt boundary. Verified end-to-end on the gpt2 tokenizer. + prompt/target boundary is exact and only prompt-predicting + positions are masked. Verified end-to-end on the gpt2 tokenizer. """ from transformers import AutoTokenizer From 1f250b9f60744284845cc438de0ca2c3e81924c3 Mon Sep 17 00:00:00 2001 From: amazloumi Date: Tue, 21 Jul 2026 23:11:45 -0400 Subject: [PATCH 03/11] fix docstrings and changelog for next-token labels --- CHANGELOG.md | 2 +- kempnerforge/data/vlm_dataset.py | 11 ++++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c05ca40..16c0d63 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -118,7 +118,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - **VLM captioning labels are now next-token aligned.** `_tokenize_and_mask` labeled each position with its own input token (`labels[i] == input_ids[i]`); with the no-shift loss, each text logit was scored against the current, already-visible token, so the model learned to copy it and the captioning loss collapsed to `0`. Labels are now shifted (`labels[i] = input_ids[i+1]`, last real position `-100`), matching text pretraining; prompt-predicting positions stay masked. Fixes all VLM arches. - - `kempnerforge/data/vlm_dataset.py` (+ `tests/unit/test_vlm_dataset.py`). + - `kempnerforge/data/vlm_dataset.py` — shared by the image and video datasets (+ `tests/unit/test_vlm_dataset.py`, `tests/unit/test_video_dataset.py`). - **Resume silently reset AdamW optimizer momentum.** `CheckpointManager` round-tripped optimizer state through raw `optimizer.state_dict()` / `optimizer.load_state_dict()`. On resume the optimizer is freshly built, so its `state_dict()` is empty — `dcp.load` then had no `exp_avg` / `exp_avg_sq` tensors to fill, and the moments were silently dropped, resetting Adam momentum to zero at every resume point. Model weights, scheduler, dataloader position, and RNG all restored correctly; only the optimizer moments were lost, so resumed runs were not bit-exact. - `kempnerforge/checkpoint/manager.py`: save and load now go through DCP's `get_model_state_dict` / `get_optimizer_state_dict` / `set_model_state_dict` / `set_optimizer_state_dict`. The getters build a load template with the optimizer moments allocated in the correct FSDP/DTensor layout, so `dcp.load` repopulates them; the setters write the loaded values back into the live optimizer. - `docs/checkpointing/dcp-model.md`: updated the save/load snippets and the "shape to fill" explanation to the DCP-aware helpers. diff --git a/kempnerforge/data/vlm_dataset.py b/kempnerforge/data/vlm_dataset.py index 794bce4..a7c2949 100644 --- a/kempnerforge/data/vlm_dataset.py +++ b/kempnerforge/data/vlm_dataset.py @@ -6,9 +6,9 @@ - ``pixel_values``: ``(3, H, W)`` float tensor, resized to ``image_size`` and normalized with the provided mean/std. - ``input_ids``: ``(T,)`` int64 tensor, right-padded to ``max_text_len``. -- ``labels``: ``(T,)`` int64 tensor matching ``input_ids`` with ``-100`` - on padding positions and (optionally) on prompt positions when - ``prompt_field`` is set. +- ``labels``: ``(T,)`` int64 next-token targets (``labels[i]`` = input + token ``i+1``), right-padded to ``max_text_len``; pad, trailing, and + prompt-predicting positions are ``-100``. ``VLMCollator`` stacks a list of samples into a batch. All batches are padded to the same fixed ``max_text_len`` regardless of batch content so @@ -184,8 +184,9 @@ class HuggingFaceVLMDataset(Dataset): tokenizer_path: HF tokenizer id or local path. max_text_len: Fixed-length pad target; passed to the collator. prompt_field: Optional column name for a prompt that should NOT - receive loss (e.g. the instruction in an instruction-tuned - dataset). Prompt tokens get ``labels=-100``. + be supervised (e.g. the instruction in an instruction-tuned + dataset). Positions predicting prompt tokens are ``-100``; the + last prompt token still predicts the first caption token. image_size: Target square image size. Default 224. image_mean / image_std: Normalization stats. Defaults match SigLIP's ``(0.5, 0.5, 0.5)``. From d59807b6aae1852f879a790853a43b25553ea2a5 Mon Sep 17 00:00:00 2001 From: amazloumi Date: Tue, 21 Jul 2026 23:13:04 -0400 Subject: [PATCH 04/11] Append EOS to caption targets so captions learn to stop --- CHANGELOG.md | 2 ++ kempnerforge/data/vlm_dataset.py | 15 ++++++++++-- tests/unit/test_video_dataset.py | 15 ++++++------ tests/unit/test_vlm_dataset.py | 40 +++++++++++++++++++++----------- 4 files changed, 49 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 16c0d63..e42ef07 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -117,6 +117,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **Captions train an EOS stop token.** `_tokenize_and_mask` appends the tokenizer's EOS (when defined) so the last caption token learns to predict *stop*; previously captions (`add_special_tokens=False`) had no stop target and generation never learned to terminate. Unprompted, the first caption token stays unsupervised (standard LM) — use a prompt so its last token predicts it. + - `kempnerforge/data/vlm_dataset.py` (+ `tests/unit/test_vlm_dataset.py`, `tests/unit/test_video_dataset.py`). - **VLM captioning labels are now next-token aligned.** `_tokenize_and_mask` labeled each position with its own input token (`labels[i] == input_ids[i]`); with the no-shift loss, each text logit was scored against the current, already-visible token, so the model learned to copy it and the captioning loss collapsed to `0`. Labels are now shifted (`labels[i] = input_ids[i+1]`, last real position `-100`), matching text pretraining; prompt-predicting positions stay masked. Fixes all VLM arches. - `kempnerforge/data/vlm_dataset.py` — shared by the image and video datasets (+ `tests/unit/test_vlm_dataset.py`, `tests/unit/test_video_dataset.py`). - **Resume silently reset AdamW optimizer momentum.** `CheckpointManager` round-tripped optimizer state through raw `optimizer.state_dict()` / `optimizer.load_state_dict()`. On resume the optimizer is freshly built, so its `state_dict()` is empty — `dcp.load` then had no `exp_avg` / `exp_avg_sq` tensors to fill, and the moments were silently dropped, resetting Adam momentum to zero at every resume point. Model weights, scheduler, dataloader position, and RNG all restored correctly; only the optimizer moments were lost, so resumed runs were not bit-exact. diff --git a/kempnerforge/data/vlm_dataset.py b/kempnerforge/data/vlm_dataset.py index a7c2949..7aae322 100644 --- a/kempnerforge/data/vlm_dataset.py +++ b/kempnerforge/data/vlm_dataset.py @@ -135,7 +135,11 @@ def _tokenize_and_mask( """Tokenize into right-padded input_ids + next-token labels. ``labels[i]`` is token ``i+1`` (matching text pretraining and the no-shift - loss); pad, trailing, and prompt-predicting positions are ``-100``. + loss); pad, trailing, and prompt-predicting positions are ``-100``. An EOS + token is appended (when the tokenizer defines one) so the last caption token + is trained to stop. The first caption token is supervised only when a + ``prompt`` precedes it (the last prompt token predicts it); a plain, + unprompted caption leaves the first token unsupervised. BPE and SentencePiece tokenizers are NOT prefix-preserving: in general ``tokenize(prompt) + tokenize(text)`` differs from @@ -154,7 +158,14 @@ def _tokenize_and_mask( full_ids = list(tokenizer(text, add_special_tokens=False)["input_ids"]) prompt_len = 0 - full_ids = full_ids[:max_text_len] + # Append EOS so the last caption token is trained to predict a stop token + # (captions use add_special_tokens=False, so none is added otherwise); + # reserve its slot within max_text_len. Skip when the tokenizer has no EOS. + eos_id = tokenizer.eos_token_id + if eos_id is not None: + full_ids = full_ids[: max_text_len - 1] + [eos_id] + else: + full_ids = full_ids[:max_text_len] pad_id = resolve_pad_id(tokenizer) n = len(full_ids) diff --git a/tests/unit/test_video_dataset.py b/tests/unit/test_video_dataset.py index 5068b5a..3c572e2 100644 --- a/tests/unit/test_video_dataset.py +++ b/tests/unit/test_video_dataset.py @@ -126,9 +126,9 @@ def test_caption_is_supervised_when_frames_present(self, monkeypatch): monkeypatch.setattr(vd, "decode_video_frames", lambda *a, **k: _frames(4)) ds = _StubVideoDataset(["1"], ["abc"], max_frames=8, max_text_len=8) item = ds[0] - # "abc" -> ids 1,2,3; next-token labels [2, 3, -100], rest -100. - assert item["labels"][:3].tolist() == [2, 3, -100] - assert (item["labels"][3:] == -100).all() + # "abc" -> [1,2,3] + EOS(28); next-token labels [2,3,28,-100], rest -100. + assert item["labels"][:4].tolist() == [2, 3, 28, -100] + assert (item["labels"][4:] == -100).all() def test_decode_failure_yields_zero_clip_no_loss(self, monkeypatch): def _boom(*a, **k): @@ -152,12 +152,11 @@ def test_prompt_is_masked(self, monkeypatch): monkeypatch.setattr(vd, "decode_video_frames", lambda *a, **k: _frames(2)) ds = _StubVideoDataset(["1"], ["xyz"], max_frames=4, max_text_len=8, prompt="ab") item = ds[0] - # prompt "ab": the last prompt token predicts the first target (24); - # the earlier prompt-predicting position is masked. Then 25, 26, -100. - assert item["input_ids"][:5].tolist() == [1, 2, 24, 25, 26] + # prompt "ab": last prompt token predicts first target (24); caption + EOS(28). + assert item["input_ids"][:6].tolist() == [1, 2, 24, 25, 26, 28] assert item["labels"][0].item() == -100 - assert item["labels"][1:4].tolist() == [24, 25, 26] - assert item["labels"][4].item() == -100 + assert item["labels"][1:5].tolist() == [24, 25, 26, 28] + assert item["labels"][5].item() == -100 def test_len(self): ds = _StubVideoDataset(["1", "2", "3"], ["a", "b", "c"]) diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index baed861..ac8447c 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -93,28 +93,42 @@ def test_shape_and_pad(self): ids, labels = _tokenize_and_mask(tok, "abc", max_text_len=8, prompt=None) assert ids.shape == (8,) assert labels.shape == (8,) - assert ids[:3].tolist() == [1, 2, 3] - assert ids[3].item() == 0 # pad - assert labels[:3].tolist() == [2, 3, -100] # next-token; last has no target - assert labels[3].item() == -100 # pad masked + assert ids[:4].tolist() == [1, 2, 3, 28] # caption + appended EOS + assert ids[4].item() == 0 # pad + assert labels[:4].tolist() == [2, 3, 28, -100] # next-token; "c" predicts EOS + assert labels[4].item() == -100 # pad masked def test_prompt_masks_labels(self): tok = _MockTokenizer() ids, labels = _tokenize_and_mask(tok, text="xyz", max_text_len=8, prompt="ab") - # Prompt "ab" = 2 tokens; target "xyz" = 3 tokens; total 5 tokens. - assert ids[:5].tolist() == [1, 2, 24, 25, 26] - # Next-token: last prompt token predicts the first target (24); earlier - # prompt-predicting positions are masked; last real token has no target. + # Prompt "ab" (2 tokens) + target "xyz" (3) + appended EOS (28). + assert ids[:6].tolist() == [1, 2, 24, 25, 26, 28] + # Next-token: the last prompt token predicts the first target (24), the + # last caption token predicts EOS; earlier prompt-predicting positions masked. assert labels[0].item() == -100 - assert labels[1:4].tolist() == [24, 25, 26] - assert labels[4].item() == -100 - assert labels[5].item() == -100 # pad masked + assert labels[1:5].tolist() == [24, 25, 26, 28] + assert labels[5].item() == -100 # EOS has no successor def test_truncation(self): tok = _MockTokenizer() ids, labels = _tokenize_and_mask(tok, text="abcdefghij", max_text_len=4, prompt=None) - assert ids.tolist() == [1, 2, 3, 4] - assert labels.tolist() == [2, 3, 4, -100] # next-token labels; final position has none + assert ids.tolist() == [1, 2, 3, 28] # caption truncated to fit the appended EOS + assert labels.tolist() == [2, 3, 28, -100] + + def test_single_token_caption_supervises_eos(self): + """A one-token caption still yields a supervised target (the appended + EOS), so it is never an all--100 row.""" + tok = _MockTokenizer() + ids, labels = _tokenize_and_mask(tok, "a", max_text_len=8, prompt=None) + assert ids[:2].tolist() == [1, 28] # "a" + EOS + assert labels[0].item() == 28 # predict EOS (learn to stop) + assert (labels[1:] == -100).all() + + def test_empty_caption_all_masked(self): + """An empty caption has no target (a lone EOS has no successor).""" + tok = _MockTokenizer() + _, labels = _tokenize_and_mask(tok, "", max_text_len=8, prompt=None) + assert (labels == -100).all() def test_prompt_mask_with_bpe_tokenizer(self): """Regression: BPE (gpt2) and SentencePiece tokenizers are not From 849eaaa259b48733ec53135228402b41353b9391 Mon Sep 17 00:00:00 2001 From: amazloumi Date: Fri, 24 Jul 2026 17:50:22 -0400 Subject: [PATCH 05/11] Require vlm.max_text_len >= 2 (a caption token plus its target) --- kempnerforge/config/vlm.py | 4 ++-- tests/unit/test_vlm_config.py | 5 +++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/kempnerforge/config/vlm.py b/kempnerforge/config/vlm.py index 45e3890..3882607 100644 --- a/kempnerforge/config/vlm.py +++ b/kempnerforge/config/vlm.py @@ -128,8 +128,8 @@ def __post_init__(self) -> None: f"Registered: {sorted(registered)}. " f"Reserved (not yet implemented): {sorted(_RESERVED_ARCHS)}." ) - if self.max_text_len <= 0: - raise ValueError("vlm.max_text_len must be positive") + if self.max_text_len < 2: + raise ValueError("vlm.max_text_len must be >= 2 (a caption token plus its target)") if self.freeze_schedule: steps = [s.start_step for s in self.freeze_schedule] if steps != sorted(steps) or len(steps) != len(set(steps)): diff --git a/tests/unit/test_vlm_config.py b/tests/unit/test_vlm_config.py index 82ecb18..fc37642 100644 --- a/tests/unit/test_vlm_config.py +++ b/tests/unit/test_vlm_config.py @@ -33,6 +33,11 @@ def test_zero_max_text_len(self): with pytest.raises(ValueError, match="max_text_len"): VLMConfig(max_text_len=0) + def test_max_text_len_one_rejected(self): + # Needs >= 2: one caption token plus its next-token target. + with pytest.raises(ValueError, match="max_text_len"): + VLMConfig(max_text_len=1) + def test_freeze_schedule_monotonic(self): with pytest.raises(ValueError, match="strictly monotonic"): VLMConfig( From c1505640e779a961af93f8dc9d841dbb1279d84c Mon Sep 17 00:00:00 2001 From: amazloumi Date: Fri, 24 Jul 2026 18:19:32 -0400 Subject: [PATCH 06/11] Cover the no-EOS-tokenizer branch of _tokenize_and_mask --- tests/unit/test_vlm_dataset.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index ac8447c..6894e2c 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -130,6 +130,15 @@ def test_empty_caption_all_masked(self): _, labels = _tokenize_and_mask(tok, "", max_text_len=8, prompt=None) assert (labels == -100).all() + def test_no_eos_tokenizer_appends_nothing(self): + """A tokenizer with no EOS id appends nothing and keeps up to + max_text_len caption tokens (the no-EOS branch of the append).""" + tok = _MockTokenizer() + tok.eos_token_id = None + ids, labels = _tokenize_and_mask(tok, "abcde", max_text_len=4, prompt=None) + assert ids.tolist() == [1, 2, 3, 4] # no EOS slot reserved; kept full 4 tokens + assert labels.tolist() == [2, 3, 4, -100] # next-token; last real token unsupervised + def test_prompt_mask_with_bpe_tokenizer(self): """Regression: BPE (gpt2) and SentencePiece tokenizers are not prefix-preserving. ``tokenize(prompt) + tokenize(text)`` differs From 2aa755d4e56d98d4508c6e1017552b52d6c0560d Mon Sep 17 00:00:00 2001 From: amazloumi Date: Wed, 2 Sep 2026 13:40:42 -0400 Subject: [PATCH 07/11] Cross-check VLM labels against a derived next-token contract Adds an independently derived, index-by-index statement of the label contract and parametrized cases that pin the prompt-mask boundary in both directions, EOS survival under truncation, and cross-rank label identity. --- tests/unit/test_vlm_dataset.py | 145 +++++++++++++++++++++++++++++++++ 1 file changed, 145 insertions(+) diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index 6894e2c..ff11e94 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -2,6 +2,8 @@ from __future__ import annotations +import random + import pytest import torch from PIL import Image @@ -196,6 +198,149 @@ def test_prompt_mask_under_bpe_merge_keeps_target_intact(self): ) +def _expected_next_token_labels(ids: list[int], prompt_len: int, max_text_len: int) -> list[int]: + """Index-by-index restatement of the label contract, derived from the loss. + + ``cross_entropy_loss`` does not shift, and the VLM arches hand it logits + over the text positions, so position ``t`` scores a prediction of token + ``t+1``. Written as an explicit loop (not tensor slices) so it cannot + share an off-by-one with the implementation it checks. + """ + labels = [-100] * max_text_len + for t in range(max_text_len): + nxt = t + 1 + if nxt >= len(ids): + continue # no successor token to predict + if nxt <= prompt_len - 1: + continue # successor is a prompt token: given, not supervised + labels[t] = ids[nxt] + return labels + + +class TestNextTokenLabelContract: + """Cross-checks ``_tokenize_and_mask`` against an independently derived + statement of the contract, instead of against fitted literals.""" + + @pytest.mark.parametrize( + ("text", "prompt", "max_text_len"), + [ + ("abc", None, 8), # ordinary caption, room to spare + ("abc", "ab", 8), # prompted caption + ("abcdefg", None, 8), # caption exactly max_text_len - 1 tokens + ("abcdefgh", None, 8), # caption exactly max_text_len tokens + ("abcdefghijkl", None, 8), # caption longer than the budget + ("abcdefgh", "qrs", 8), # prompt + overlong caption + ("a", None, 8), # single-token caption + ("", None, 8), # empty caption + ("abc", "", 8), # empty prompt string + ("ab", "a", 8), # prompt_len == 1 + ("a", None, 2), # minimum legal max_text_len + ], + ) + def test_matches_derived_contract(self, text, prompt, max_text_len): + tok = _MockTokenizer() + ids, labels = _tokenize_and_mask(tok, text, max_text_len, prompt) + + # --- input_ids properties (stated, not copied from the implementation) + real = list(tok(text)["input_ids"]) + prompt_ids = list(tok(prompt)["input_ids"]) if prompt is not None else [] + budget = max_text_len - 1 # one slot reserved for the appended EOS + kept = (prompt_ids + real)[:budget] + n = len(kept) + 1 + assert ids.shape == (max_text_len,) + assert ids[: len(kept)].tolist() == kept + # EOS always lands in input_ids: truncation reserves its slot, so it can + # never be the token that falls off the end. + assert ids[len(kept)].item() == tok.eos_token_id + assert (ids[n:] == 0).all() # pad + + # --- labels match the derived contract exactly + prompt_len = min(len(prompt_ids), n) + assert labels.tolist() == _expected_next_token_labels( + ids[:n].tolist(), prompt_len, max_text_len + ) + + def test_prompt_mask_boundary_both_directions(self): + """The two indices either side of the prompt boundary, pinned. + + With shifted labels the prompt occupies ``ids[0 .. prompt_len-1]``, so + ``labels[prompt_len-1]`` predicts the *first caption token* and must stay + supervised, while ``labels[prompt_len-2]`` predicts the last prompt token + and must be masked. Widening the mask by one silently drops the first + caption token; narrowing it leaks a prompt token into the loss. + """ + tok = _MockTokenizer() + prompt, text = "abcd", "xyz" + prompt_len = len(tok(prompt)["input_ids"]) # 4 + ids, labels = _tokenize_and_mask(tok, text, max_text_len=16, prompt=prompt) + + first_caption_id = tok(text)["input_ids"][0] + assert ids[prompt_len].item() == first_caption_id + # Supervised: predicts the first caption token. + assert labels[prompt_len - 1].item() == first_caption_id + # Masked: would predict the last prompt token. + assert labels[prompt_len - 2].item() == -100 + # And every earlier position is masked too (none leak). + assert (labels[: prompt_len - 1] == -100).all() + + def test_eos_is_a_supervised_target(self): + """The last caption token predicts EOS, so captions learn to stop.""" + tok = _MockTokenizer() + _, labels = _tokenize_and_mask(tok, "hello", max_text_len=16, prompt="ab") + supervised = labels[labels != -100].tolist() + assert supervised[-1] == tok.eos_token_id + assert supervised.count(tok.eos_token_id) == 1 + + def test_eos_survives_an_overlong_caption(self): + """Truncation reserves the EOS slot, so a caption longer than the budget + still ends in a supervised stop target rather than a dropped one.""" + tok = _MockTokenizer() + ids, labels = _tokenize_and_mask(tok, "abcdefghij", max_text_len=5, prompt=None) + assert ids[-1].item() == tok.eos_token_id + assert labels[-2].item() == tok.eos_token_id + assert labels[-1].item() == -100 # EOS has no successor + + def test_prompt_filling_the_budget_is_in_range(self): + """A prompt that leaves no room for a caption must not raise or produce a + negative-index slice; it degrades to an unsupervised row. (Rejecting such + a prompt up front is a separate concern, validated at dataset init.)""" + tok = _MockTokenizer() + for prompt_chars in (7, 8, 12): + _, labels = _tokenize_and_mask(tok, "xyz", max_text_len=8, prompt="a" * prompt_chars) + assert labels.shape == (8,) + assert (labels == -100).all() or (labels != -100).sum() == 1 + + def test_labels_are_identical_across_ranks(self, monkeypatch): + """Every DP rank must build byte-identical labels for a given index. + + Label construction is pure in ``(tokenizer, text, max_text_len, prompt)``, + so a rank-dependent or RNG-dependent path here would desync the loss + across ranks. Emulates two ranks in-process: distinct rank environment + and distinct global RNG state, separate tokenizer instances, same index. + """ + from transformers import AutoTokenizer + + captions = ["a red bus", "two dogs running", "a bowl of soup"] + index = 1 + + def build(rank: int) -> tuple[bytes, bytes]: + monkeypatch.setenv("RANK", str(rank)) + monkeypatch.setenv("LOCAL_RANK", str(rank)) + monkeypatch.setenv("WORLD_SIZE", "2") + random.seed(rank) + torch.manual_seed(rank) + tok = AutoTokenizer.from_pretrained("gpt2") + ids, labels = _tokenize_and_mask( + tok, captions[index], max_text_len=16, prompt="Describe: " + ) + return ids.numpy().tobytes(), labels.numpy().tobytes() + + ids0, labels0 = build(0) + ids1, labels1 = build(1) + assert ids0 == ids1 + assert labels0 == labels1 + + # --------------------------------------------------------------------------- # VLMCollator # --------------------------------------------------------------------------- From 71171c1eeba9893b43a34cb2278eead072320d45 Mon Sep 17 00:00:00 2001 From: amazloumi Date: Wed, 2 Sep 2026 13:49:53 -0400 Subject: [PATCH 08/11] Note the objective change and the max_text_len floor in the changelog --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e42ef07..b05c509 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -121,6 +121,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `kempnerforge/data/vlm_dataset.py` (+ `tests/unit/test_vlm_dataset.py`, `tests/unit/test_video_dataset.py`). - **VLM captioning labels are now next-token aligned.** `_tokenize_and_mask` labeled each position with its own input token (`labels[i] == input_ids[i]`); with the no-shift loss, each text logit was scored against the current, already-visible token, so the model learned to copy it and the captioning loss collapsed to `0`. Labels are now shifted (`labels[i] = input_ids[i+1]`, last real position `-100`), matching text pretraining; prompt-predicting positions stay masked. Fixes all VLM arches. - `kempnerforge/data/vlm_dataset.py` — shared by the image and video datasets (+ `tests/unit/test_vlm_dataset.py`, `tests/unit/test_video_dataset.py`). + - `kempnerforge/config/vlm.py`: `max_text_len` must now be `>= 2` (one caption token plus its next-token target); `1` previously passed validation and supervised nothing. + - **Behavior note:** this changes the training objective, not a config or checkpoint key. No shim is needed and nothing fails to load, but a VLM captioning checkpoint trained before this fix learned the copy objective, so resuming such a run changes its behavior mid-flight. Restart those runs rather than resuming them. - **Resume silently reset AdamW optimizer momentum.** `CheckpointManager` round-tripped optimizer state through raw `optimizer.state_dict()` / `optimizer.load_state_dict()`. On resume the optimizer is freshly built, so its `state_dict()` is empty — `dcp.load` then had no `exp_avg` / `exp_avg_sq` tensors to fill, and the moments were silently dropped, resetting Adam momentum to zero at every resume point. Model weights, scheduler, dataloader position, and RNG all restored correctly; only the optimizer moments were lost, so resumed runs were not bit-exact. - `kempnerforge/checkpoint/manager.py`: save and load now go through DCP's `get_model_state_dict` / `get_optimizer_state_dict` / `set_model_state_dict` / `set_optimizer_state_dict`. The getters build a load template with the optimizer moments allocated in the correct FSDP/DTensor layout, so `dcp.load` repopulates them; the setters write the loaded values back into the live optimizer. - `docs/checkpointing/dcp-model.md`: updated the save/load snippets and the "shape to fill" explanation to the DCP-aware helpers. From 18c1234f962ecf384d8a397f499d54f8fbdaa2c2 Mon Sep 17 00:00:00 2001 From: amazloumi Date: Wed, 2 Sep 2026 14:20:13 -0400 Subject: [PATCH 09/11] Cover prompt lengths at and beyond max_text_len in the contract cross-check --- tests/unit/test_vlm_dataset.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index ff11e94..d88af09 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -235,6 +235,11 @@ class TestNextTokenLabelContract: ("abc", "", 8), # empty prompt string ("ab", "a", 8), # prompt_len == 1 ("a", None, 2), # minimum legal max_text_len + # prompt at and beyond the budget: prompt_len is measured before + # truncation, so these pin that it still lines up with the boundary. + ("xyz", "aaaaaaa", 8), # prompt fills max_text_len - 1 + ("xyz", "aaaaaaaa", 8), # prompt fills max_text_len + ("xyz", "aaaaaaaaaaaa", 8), # prompt longer than max_text_len ], ) def test_matches_derived_contract(self, text, prompt, max_text_len): From 02f3f4f6a51f6df81e06b8a3a97ac0c27aaf8eff Mon Sep 17 00:00:00 2001 From: amazloumi Date: Wed, 2 Sep 2026 14:21:30 -0400 Subject: [PATCH 10/11] Clarify that an over-long prompt is not rejected today --- tests/unit/test_vlm_dataset.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index d88af09..eb0742a 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -307,8 +307,9 @@ def test_eos_survives_an_overlong_caption(self): def test_prompt_filling_the_budget_is_in_range(self): """A prompt that leaves no room for a caption must not raise or produce a - negative-index slice; it degrades to an unsupervised row. (Rejecting such - a prompt up front is a separate concern, validated at dataset init.)""" + negative-index slice. It degrades to a row with nothing (or only EOS) + supervised; rejecting such a prompt up front is a separate concern and is + not validated here.""" tok = _MockTokenizer() for prompt_chars in (7, 8, 12): _, labels = _tokenize_and_mask(tok, "xyz", max_text_len=8, prompt="a" * prompt_chars) From 408ffc3886a2ab8d51a9bab8489569a0f864393c Mon Sep 17 00:00:00 2001 From: amazloumi Date: Wed, 2 Sep 2026 14:25:48 -0400 Subject: [PATCH 11/11] Run the label-contract cross-check on the no-EOS tokenizer branch too --- tests/unit/test_vlm_dataset.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_vlm_dataset.py b/tests/unit/test_vlm_dataset.py index eb0742a..b742132 100644 --- a/tests/unit/test_vlm_dataset.py +++ b/tests/unit/test_vlm_dataset.py @@ -242,21 +242,26 @@ class TestNextTokenLabelContract: ("xyz", "aaaaaaaaaaaa", 8), # prompt longer than max_text_len ], ) - def test_matches_derived_contract(self, text, prompt, max_text_len): + @pytest.mark.parametrize("with_eos", [True, False]) + def test_matches_derived_contract(self, text, prompt, max_text_len, with_eos): tok = _MockTokenizer() + if not with_eos: + tok.eos_token_id = None ids, labels = _tokenize_and_mask(tok, text, max_text_len, prompt) # --- input_ids properties (stated, not copied from the implementation) real = list(tok(text)["input_ids"]) prompt_ids = list(tok(prompt)["input_ids"]) if prompt is not None else [] - budget = max_text_len - 1 # one slot reserved for the appended EOS + # With an EOS one slot is reserved for it; without, the full budget is text. + budget = max_text_len - 1 if with_eos else max_text_len kept = (prompt_ids + real)[:budget] - n = len(kept) + 1 + n = len(kept) + (1 if with_eos else 0) assert ids.shape == (max_text_len,) assert ids[: len(kept)].tolist() == kept - # EOS always lands in input_ids: truncation reserves its slot, so it can - # never be the token that falls off the end. - assert ids[len(kept)].item() == tok.eos_token_id + if with_eos: + # EOS always lands in input_ids: truncation reserves its slot, so it + # can never be the token that falls off the end. + assert ids[len(kept)].item() == tok.eos_token_id assert (ids[n:] == 0).all() # pad # --- labels match the derived contract exactly