Fix VLM captioning label alignment and add an EOS stop target

[amazloumi]

Summary

• Shift VLM captioning labels for next-token loss: labels[i] is input token i+1, last real position -100. cross_entropy_loss does no internal shift and the VLM arches return logits over the text positions, so the previous labels[i] == input_ids[i] scored each text logit against the already-visible token.
• Prompt mask moves by one to match: labels[:prompt_len - 1] = -100. labels[prompt_len - 1] predicts the first caption token and stays supervised.
• Append the tokenizer's EOS to the caption (when defined), reserving its slot inside max_text_len, so the last caption token is trained to stop. Tokenizers without an EOS keep the full budget.
• vlm.max_text_len must now be >= 2 — one caption token plus its target.
• data/video_dataset.py imports _tokenize_and_mask, so this covers the image and video captioning paths and all four arches.
• Changes the training objective, not a config or checkpoint key. Nothing fails to load, but a captioning checkpoint trained before this fix learned the copy objective, so resuming such a run changes its behavior mid-flight; restart rather than resume.

Testing

• uv run ruff check kempnerforge/ tests/ scripts/ passes — All checks passed

• uv run ruff format --check kempnerforge/ tests/ scripts/ passes — 166 files already formatted

• uv run pyright kempnerforge/ passes — 0 errors, 0 warnings

• uv run pytest tests/unit/ -v --timeout=120 — 1765 passed, 3 skipped

• uv run pytest examples/vlm/tests -q — 33 passed; uv run pytest examples/vlm/eval/tests/unit -q — 104 passed

• Before/after loss trajectory, 1 GPU on kempner_eng. examples/vlm/configs/vlm_debug.toml, seed 42, batch 16, 600 steps, warmup 100, 30k-sample COCO captions (9.6k seen, under one epoch). Same worktree and venv; only data/vlm_dataset.py and config/vlm.py differ between sides, verified by recording each side's file hashes in its log.

  step	main loss	fix loss	main grad_norm	fix grad_norm
  1	10.8750	10.8125	7.688	7.156
  100	4.9375	6.3125	2.984	1.859
  300	1.1719	4.2812	1.656	2.375
  600	1.2812	4.6562	1.273	2.344

  Both start at ln(50257) = 10.82. main tail mean 1.1711 ± 0.2100, still descending (mid-third 1.7086) with grad_norm decaying toward 0 — the copy objective has no entropy floor. fix tail mean 4.3875 ± 0.2225, levelled off against its mid-third 4.9687. Separation 3.2164 = 10.5σ from the two tails' own scatter. An identical curve would have meant the fix did not take.

• Same A/B repeated with --model.tie_embeddings=true (as in the Qwen3-0.6B config the issue reproduced on) for 2000 steps, log_interval 100. main starts at 8.3125, i.e. 2.51 nats below chance, before any training: with a tied output head the copy target is already partly satisfied at initialization. fix starts at 10.8750, on chance. Tails: main 0.6791 ± 0.1641 (grad_norm 0.41) vs fix 4.0268 ± 0.1889 (grad_norm 2.34) = 13.4σ separation, main still descending.

• Rewritten expectations were re-derived from labels[i] = ids[i+1] on paper before running the code, not adjusted until green. TestNextTokenLabelContract encodes that derivation as an index-by-index loop — not tensor slices, so it cannot share an off-by-one with the implementation — and cross-checks it over 28 parametrized cases (14 input shapes x tokenizer with and without an EOS): caption exactly max_text_len - 1, exactly max_text_len, longer than the budget, empty, single-token, empty prompt, prompt_len == 1, max_text_len == 2, and prompts filling max_text_len - 1, max_text_len, and beyond.

• Five deliberate regressions injected; the suite caught all five, no survivors: unshifted labels (the original bug, 35 failures), prompt mask one too wide (10), one too narrow (15), EOS not appended (22), shift reversed (37).

• Prompt-mask boundary in both directions: test_prompt_mask_boundary_both_directions pins labels[prompt_len - 1] supervised (equal to the first caption id) and labels[prompt_len - 2] masked.

• EOS in targets: test_eos_is_a_supervised_target, test_eos_survives_an_overlong_caption, test_single_token_caption_supervises_eos; no-EOS-tokenizer branch in test_no_eos_tokenizer_appends_nothing.

• Cross-rank label identity: test_labels_are_identical_across_ranks builds labels for one index under two rank environments and RNG seeds with separate tokenizer instances, then byte-compares input_ids and labels.

• uv run torchrun --nproc_per_node=4 -m pytest tests/distributed/test_checkpoint.py tests/distributed/test_resilience.py -v — 13 passed, 0 skipped, all 4 ranks. Collateral-damage guard only: no file under tests/distributed/ imports vlm_dataset or _tokenize_and_mask and the VLM distributed tests synthesize their own labels, so this is evidence of no breakage, not evidence about label construction.

• Re-verified against .external_docs/kf-codebase-design-new.html after implementation. Part 3 admits a correctness fix, and the plan's own live case study lists "VLM label alignment" among the general fixes to promote to core. The finished diff adds no registry category and no new config field (it tightens an existing max_text_len validator), carries tests, and helps runs beyond its origin via the shared helper. No core-to-examples/ dependency introduced.

• 2-node run via scripts/slurm/multinode.sh — deliberately not run. Its detection power is collective ordering, rank divergence, and a wedged job. This change adds no collective, no checkpoint key and no rank-dependent code, and the output tensor shape is unchanged, so none of those apply. The residual risk is rank divergence in label construction, which the cross-rank label-identity test detects directly and more cheaply.

Closes #154
Closes #156

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
kempnerforge/config/vlm.py	100.00% <100.00%> (ø)
kempnerforge/data/vlm_dataset.py	100.00% <100.00%> (+0.90%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Naeemkh]

LGTM