Fix VLM captioning label alignment and add an EOS stop target

CHANGELOG.md

@@ -117,6 +117,12 @@ 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`).
+  - `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.

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)):

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
@@ -132,21 +132,22 @@ 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``. 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
     ``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"]
@@ -157,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)
@@ -166,9 +174,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
 
 
@@ -184,8 +195,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)``.

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 supervised; rest -100.
-        assert item["labels"][:3].tolist() == [1, 2, 3]
-        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,10 +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" (2 toks) masked; "xyz" (24,25,26) supervised.
-        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]
+        # 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:5].tolist() == [24, 25, 26, 28]
+        assert item["labels"][5].item() == -100
 
     def test_len(self):
         ds = _StubVideoDataset(["1", "2", "3"], ["a", "b", "c"])

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(