Add per-frame timestamp embedding to the VLM video path

[amazloumi]

Summary

• Add FrameTimeEmbedding (kempnerforge/model/frame_time.py): sinusoidal features of a frame's timestamp (seconds) at log-spaced periods → a zero-initialized projection (identity at step 0).
• decode_video_frames returns (frames, times); WebVidVideoDataset emits frame_times (F,), VideoCollator stacks to (B, F).
• Applied per frame in _project_visual_features as a VLMWrapper sibling submodule (video only; None for the image path); built + FSDP-sharded + meta-materialized at both build sites.
• scripts/train.py threads frame_times; docs + CHANGELOG updated.
• Make the time embedding registry-driven: [time_embedding] selects the implementation (type = "sinusoidal" default, "none" disables) via @registry.register_time_embedding; new techniques drop in as small additions. Sequence-modifying encodings (Molmo2-style text time-tokens) are flagged as a separate future hook (needs interleaved-sequence support).

Testing

• uv run ruff check kempnerforge/ tests/ passes
• uv run ruff format --check kempnerforge/ tests/ scripts/ passes
• uv run pyright kempnerforge/ passes (0 errors)
• uv run pytest tests/unit/ -v --timeout=60 passes (1527 passed, 2 skipped)
• Distributed (parallel.py changed): uv run torchrun --nproc_per_node=4 -m pytest tests/distributed/ -v ← running this now
• 2-GPU FSDP smoke on vlm_video_webvid.toml (random encoder): trains, +33,792 params confirms the module is sharded/trainable

Closes #127
refs: KEM-544

[codecov]

Codecov Report

❌ Patch coverage is 94.00000% with 9 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
kempnerforge/distributed/parallel.py	57.14%	6 Missing ⚠️
kempnerforge/model/vlm.py	92.00%	1 Missing and 1 partial ⚠️
kempnerforge/data/video_io.py	94.44%	1 Missing ⚠️

Files with missing lines	Coverage Δ
kempnerforge/config/job.py	89.34% <100.00%> (+0.74%)	⬆️
kempnerforge/config/registry.py	100.00% <100.00%> (ø)
kempnerforge/config/schema.py	100.00% <100.00%> (ø)
kempnerforge/config/time_embedding.py	100.00% <100.00%> (ø)
kempnerforge/config/video.py	100.00% <ø> (ø)
kempnerforge/config/vlm.py	100.00% <ø> (ø)
kempnerforge/data/video_dataset.py	93.39% <100.00%> (+0.25%)	⬆️
kempnerforge/eval/vlm/adapter.py	94.00% <100.00%> (ø)
kempnerforge/model/frame_time.py	100.00% <100.00%> (ø)
kempnerforge/data/video_io.py	83.82% <94.44%> (+3.82%)	⬆️
... and 2 more

🚀 New features to boost your workflow:

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

[Copilot]

Pull request overview

This PR adds absolute per-frame timestamp conditioning to the VLM video pathway by propagating decoded frame presentation times through the data pipeline and injecting a registry-configurable, zero-initialized time embedding into each frame’s visual tokens.

Changes:

• decode_video_frames now returns (frames, times); datasets/collator propagate frame_times as (F,) / (B, F) and training threads it into the model forward.
• Introduces a registry-driven time-embedding module (FrameTimeEmbedding default; "none" disables) and applies it per-frame in the VLM visual-token projection (video-only).
• Ensures distributed/FSDP build paths materialize and shard the new submodule; adds unit coverage for config, embedding behavior, and build wiring.

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/unit/test_vlm.py	Adds unit tests ensuring video wrappers attach the module, image wrappers do not, and forward wiring/shape checks behave as expected.
tests/unit/test_video_io.py	Updates tests for (frames, times) return and basic timestamp properties.
tests/unit/test_video_dataset.py	Updates dataset/collator tests to validate frame_times padding/stacking behavior.
tests/unit/test_time_embedding_config.py	Adds coverage for TimeEmbeddingConfig defaults, validation, and kwargs.
tests/unit/test_frame_time.py	Adds coverage for embedding shape, zero-init behavior, gradient flow, dtype behavior, and registry builder.
tests/unit/test_distributed.py	Verifies distributed build attaches/casts frame_time_embed for video and omits it for images.
scripts/train.py	Threads time_embedding_config into model build and passes frame_times into VLM forward.
kempnerforge/model/vlm.py	Adds frame_times plumbing and applies per-frame timestamp embeddings in _project_visual_features.
kempnerforge/model/frame_time.py	Introduces TimeEmbedding interface, FrameTimeEmbedding, and registry-driven build_time_embedding.
kempnerforge/distributed/parallel.py	Builds/materializes/casts frame_time_embed in meta/CPU paths and FSDP-shards it when present.
kempnerforge/data/video_io.py	Changes decode_video_frames to also return matched presentation timestamps.
kempnerforge/data/video_dataset.py	Emits frame_times per sample and stacks it in VideoCollator.
kempnerforge/config/time_embedding.py	Adds [time_embedding] config with validation and builder kwargs.
kempnerforge/config/schema.py	Exposes TimeEmbeddingConfig in the config schema surface.
kempnerforge/config/registry.py	Adds time-embedding registry hooks (register/get/list_time_embedding).
kempnerforge/config/job.py	Adds optional time_embedding field to JobConfig.
docs/how-to/train-on-video.md	Documents per-frame timestamp embedding and registry semantics.
CHANGELOG.md	Records the new per-frame timestamp feature and affected components.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 26 out of 26 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

kempnerforge/eval/vlm/adapter.py:291

• decode_video_frames now returns (frames, times), but eval drops times here and never threads per-frame timestamps into the model. As a result, a video checkpoint trained with [time_embedding] will (a) rebuild without frame_time_embed in _build_model(...) because build_vlm_wrapper(...) is called without time_embedding_config=config.time_embedding, and (b) run inference without frame_times=... (see _generate_batch), silently disabling the learned temporal signal. This can also cause strict checkpoint loads to fail due to missing/unexpected frame_time_embed.* keys.

Recommendation: carry frame_times alongside frames out of _render_request, pad/stack it analogously to frame_mask, pass time_embedding_config=config.time_embedding into build_vlm_wrapper in _build_model, and pass frame_times=... into model(...) calls for video.

        frames, _frame_times = decode_video_frames(
            path,
            fps=video_config.fps,
            min_frames=video_config.min_frames,
            max_frames=video_config.max_frames,
            sampling_policy=video_config.sampling_policy,
        )