From e967277c54fa6a60823c67e82e209e9997db4943 Mon Sep 17 00:00:00 2001 From: camilobrownpinilla Date: Mon, 13 Jul 2026 16:12:34 -0400 Subject: [PATCH 1/4] Support multi-image eval requests on video checkpoints Pack N images from a chat request as an ordered clip (a single image is the 1-frame case), reusing the existing video frame/clip machinery. Images past frames_per_clip are truncated downstream, with a warning surfaced at render. Image checkpoints still reject multiple images; text-only stays out of scope. Flips the video-path multi-image test to expect frames and adds coverage for the truncation warning. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01STRsAc9nZjMhkLzYkkYp67 --- examples/vlm-evaluation/README.md | 26 ++++----- examples/vlm-evaluation/adapter.py | 58 ++++++++++++------- .../vlm-evaluation/tests/unit/test_adapter.py | 20 ++++++- 3 files changed, 66 insertions(+), 38 deletions(-) diff --git a/examples/vlm-evaluation/README.md b/examples/vlm-evaluation/README.md index 67d998e..b4f7023 100644 --- a/examples/vlm-evaluation/README.md +++ b/examples/vlm-evaluation/README.md @@ -119,12 +119,12 @@ uv run python examples/vlm-evaluation/vlm_eval_harness.py \ on the checkpoint's frame budget matching the reference's, which is a training choice rather than a knob here. - **Scope.** One video per request, single-turn, zero-shot, generative arches - (`joint_decoder` / `cross_attention` / `mot`). A single **image** task also runs - on a video checkpoint — the image is treated as a 1-frame clip, zero-padded to - `frames_per_clip`. Multiple videos, mixed image+video, multiple images, audio, - and multi-turn / few-shot raise a clear error; MoMa still fails fast. An - **image** checkpoint cannot evaluate video and raises a clear error if handed a - video task. + (`joint_decoder` / `cross_attention` / `mot`). **Image** tasks also run on a + video checkpoint — one image is a 1-frame clip, and multiple images are packed + as an ordered clip (zero-padded, and truncated with a warning past + `frames_per_clip`). Multiple videos, mixed image+video, audio, and multi-turn / + few-shot raise a clear error; MoMa still fails fast. An **image** checkpoint + cannot evaluate video and raises a clear error if handed a video task. ## Limitations @@ -139,13 +139,13 @@ Several are tracked follow-ups. generation-only. A MoMa checkpoint fails fast with a clear error. Joint-Decoder (`joint_decoder`), Cross-Attention (`cross_attention`), and MoT (`mot`) are supported. -- **One visual per request; no multi-turn / few-shot / multi-image.** A request - carries exactly one image (image checkpoint) or one video (video checkpoint — - see [Video evaluation](#video-evaluation)). Audio, multiple images, multiple - videos, mixed image+video, and multi-turn / few-shot requests raise a clear - error. Multi-image and multi-turn/few-shot are tracked follow-ups (for chat - tasks lmms-eval delivers few-shot as extra content blocks/turns, so it reduces - to multi-image + multi-turn support). +- **One visual per request on image checkpoints; no multi-turn / few-shot.** An + image checkpoint carries exactly one image per request (multiple images raise); a + video checkpoint carries one video, or one or more images packed as an ordered + clip (see [Video evaluation](#video-evaluation)). Audio, multiple videos, mixed + image+video, and multi-turn / few-shot requests raise a clear error. Multi-turn / + few-shot is a tracked follow-up (for chat tasks lmms-eval delivers few-shot as + extra content blocks/turns, so it reduces to multi-turn support). - **Prompt flattening discards structure.** Flattening drops role/turn structure and any model-specific chat template. KempnerForge pre-training uses no chat template; once a post-training format exists, repo-wide chat-template support diff --git a/examples/vlm-evaluation/adapter.py b/examples/vlm-evaluation/adapter.py index 9736859..351beae 100644 --- a/examples/vlm-evaluation/adapter.py +++ b/examples/vlm-evaluation/adapter.py @@ -45,8 +45,10 @@ - **Image and video.** An image checkpoint evaluates exactly one image per request; a video checkpoint (a ``[video]`` config) evaluates one video per request — decoded to a fixed ``frames_per_clip`` clip via the training - frame-sampling policy — and also accepts a single image (a 1-frame clip). - Audio, multi-image, multiple videos, mixed image+video, and multi-turn/few-shot + frame-sampling policy — or one or more images packed as an ordered clip (a + single image is the 1-frame case; multiple images fill successive frame slots, + zero-padded and truncated to ``frames_per_clip``). Audio, multiple videos, + mixed image+video, multi-image on an image checkpoint, and multi-turn/few-shot requests raise ``NotImplementedError``; ``loglikelihood`` and ``generate_until_multi_round`` are not implemented (chat tasks are generation-only). Visual input is modeled as an ordered list of frames (a @@ -215,16 +217,19 @@ def _render_request( - **Image checkpoint** (``video_config is None``): exactly one image per request; video content raises (an image model cannot evaluate video). - - **Video checkpoint**: exactly one video — decoded to frames via + - **Video checkpoint**: one video — decoded to frames via ``video_io.decode_video_frames`` using the checkpoint's frame-sampling - policy — or, when no video is present, a single image treated as a - 1-frame clip (zero-padded to ``frames_per_clip`` downstream). - - Out-of-scope content (audio, multi-turn/few-shot, multi-image, multiple - videos, mixed image+video) raises ``NotImplementedError`` so the offending - task is surfaced rather than silently mishandled. Text content blocks are - concatenated in message order (newline-joined); role/turn structure is - intentionally discarded (see the module docstring on flattening). + policy — or, when no video is present, one or more images treated as an + ordered clip of frames (zero-padded to ``frames_per_clip`` downstream; + frames beyond ``frames_per_clip`` are truncated with a warning). A single + image is the length-1 clip. + + Out-of-scope content (audio, multi-turn/few-shot, multiple videos, mixed + image+video, and — on an image checkpoint — multiple images) raises + ``NotImplementedError`` so the offending task is surfaced rather than + silently mishandled. Text content blocks are concatenated in message order + (newline-joined); role/turn structure is intentionally discarded (see the + module docstring on flattening). """ images, videos, audios = messages.extract_media() if audios: @@ -262,8 +267,8 @@ def _render_request( ) return images, prompt - # Video checkpoint: exactly one visual — a video (decoded to frames) or a - # single image (treated as a 1-frame clip, zero-padded downstream). + # Video checkpoint: a video (decoded to frames), or one or more images packed + # as an ordered clip (a single image is the 1-frame case). if len(videos) > 1: raise NotImplementedError( f"Multiple videos per request are not supported, got {len(videos)}. " @@ -294,15 +299,24 @@ def _render_request( f"No frames decoded from {path}; evaluating a zero clip (result unreliable)." ) return frames, prompt - if len(images) == 1: - # A single image on a video checkpoint: a 1-frame clip (zero-padded to - # frames_per_clip downstream), consistent with how training pads short clips. - return images, prompt - raise NotImplementedError( - f"A video-checkpoint request must carry exactly one video or one image, got " - f"{len(images)} images and no video. Multi-image and text-only requests are out " - "of scope; report the task to the project owner." - ) + if not images: + raise NotImplementedError( + "A video-checkpoint request must carry a video or at least one image, got no " + "visual content. Text-only requests are out of scope; report the task to the " + "project owner." + ) + # One or more images on a video checkpoint: treat them as an ordered clip of frames + # (zero-padded to frames_per_clip downstream, extra frames truncated), mirroring how + # training packs short clips. A single image is the length-1 case. lmms-eval delivers + # a multi-image request as several image content blocks (e.g. MMMU); extract_media + # gives them to us in document order, which we keep as frame order. + if len(images) > video_config.max_frames: + logger.warning( + f"Request carries {len(images)} images but the checkpoint's clip length is " + f"frames_per_clip={video_config.max_frames}; the extra frames will be truncated " + "(result may be unreliable)." + ) + return images, prompt def _to_pil(frame: Any) -> Any: diff --git a/examples/vlm-evaluation/tests/unit/test_adapter.py b/examples/vlm-evaluation/tests/unit/test_adapter.py index a4fd7f5..7a21ccf 100644 --- a/examples/vlm-evaluation/tests/unit/test_adapter.py +++ b/examples/vlm-evaluation/tests/unit/test_adapter.py @@ -235,9 +235,23 @@ def test_audio_still_raises(self, vcfg): with pytest.raises(NotImplementedError, match="[Aa]udio"): _render_request(msg, vcfg) - def test_multi_image_raises(self, vcfg): - with pytest.raises(NotImplementedError, match="exactly one"): - _render_request(_chat([_image(_img()), _image(_img())]), vcfg) + def test_multi_image_returns_frames(self, vcfg): + # A multi-image request on a video checkpoint is packed as an ordered clip: + # the N images become N frames (a single image is the length-1 case). + frames, prompt = _render_request( + _chat([_text("compare"), _image(_img()), _image(_img())]), vcfg + ) + assert len(frames) == 2 + assert prompt == "compare" + + def test_multi_image_over_max_frames_warns(self, monkeypatch, vcfg): + # More images than the clip length (max_frames=4): all are returned here and + # truncated downstream by frames_to_clip_tensor, so render warns the drop is visible. + rec = _RecordingLogger() + monkeypatch.setattr("adapter.logger", rec) + frames, _ = _render_request(_chat([_image(_img()) for _ in range(5)]), vcfg) + assert len(frames) == 5 + assert any("truncated" in m for m in rec.warnings) def test_multi_turn_raises(self, vcfg): msg = ChatMessages( From c686ed137265ee10934859b1b80f5cb988818f75 Mon Sep 17 00:00:00 2001 From: camilobrownpinilla Date: Mon, 13 Jul 2026 16:41:46 -0400 Subject: [PATCH 2/4] Support text-only evaluation on image and video checkpoints Add a text-only (no-visual) forward path for the generative VLM arches so generate_until text benchmarks (GSM8K, IFEval, ...) can measure text-backbone drift on image and video checkpoints, in the same lmms-eval harness as the multimodal tasks. Model core (backwards-compatible, no new parameters or state-dict keys): - vlm.py: widen pixel_values to `Tensor | None` across the ModalityStrategy protocol, all four strategies, and VLMWrapper.forward. A None (text-only) request yields an empty ModalityContext for joint_decoder/cross_attention/mot and is rejected for the non-generative moma arch. - transformer.py: skip cross-attention blocks when image_features is None; run the MoT branch with n_image=0 when there is neither a prefix nor modality_ids. - mot.py: reshape the MoT projections with explicit head counts so a zero-length image stream flows through (bit-identical for the non-empty path). Only previously-erroring paths change behavior; existing image/video forwards, checkpoints, and configs are untouched. Adapter (examples/vlm-evaluation): - _render_request renders a no-visual request as empty frames on both image and video checkpoints. - _generate_batch accepts pixel_values=None (skip vision, 0 image-token budget, device from the prompts). - generate_until segregates a gen_kwargs chunk into visual and text-only sub-batches, decodes each, and scatters results back in original order. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01STRsAc9nZjMhkLzYkkYp67 --- examples/vlm-evaluation/README.md | 25 +++ examples/vlm-evaluation/adapter.py | 208 +++++++++++------- .../vlm-evaluation/tests/unit/test_adapter.py | 86 +++++++- kempnerforge/model/mot.py | 13 +- kempnerforge/model/transformer.py | 49 +++-- kempnerforge/model/vlm.py | 51 ++++- tests/unit/test_model.py | 80 ++++++- tests/unit/test_mot.py | 31 +++ tests/unit/test_vlm.py | 70 ++++++ 9 files changed, 486 insertions(+), 127 deletions(-) diff --git a/examples/vlm-evaluation/README.md b/examples/vlm-evaluation/README.md index b4f7023..21c2f66 100644 --- a/examples/vlm-evaluation/README.md +++ b/examples/vlm-evaluation/README.md @@ -126,6 +126,31 @@ uv run python examples/vlm-evaluation/vlm_eval_harness.py \ few-shot raise a clear error; MoMa still fails fast. An **image** checkpoint cannot evaluate video and raises a clear error if handed a video task. +## Text-only evaluation + +Text-only `generate_until` benchmarks (e.g. GSM8K, IFEval) run on **both image and +video checkpoints**, for the generative arches (`joint_decoder` / `cross_attention` / +`mot`). A request with no image or video renders as an empty-frame prompt and runs +the arch's **pure-text forward** — no vision encoder, no image prefix (JD/MoT), and +cross-attention blocks skipped (CA) — so the number reflects the text backbone. This +is how you measure how much VLM training drifted the base LM, in the same harness as +the multimodal tasks. + +```bash +uv run python examples/vlm-evaluation/vlm_eval_harness.py \ + --config configs/train/vlm_jd.toml \ + --checkpoint checkpoints/vlm/step_10000 \ + --tasks gsm8k \ + --limit 8 +``` + +- **Scope.** `generate_until` tasks only (generation / answer-extraction). + `loglikelihood`-scored multiple-choice suites (ARC, HellaSwag, MMLU-style) are not + supported — the adapter is generation-only. MoMa is excluded (non-generative). + Text-only, image, and video requests may be freely mixed across a task suite; each + request is decoded by its own modality path (text-only and visual requests within a + batch are decoded as separate sub-batches). + ## Limitations Several are tracked follow-ups. diff --git a/examples/vlm-evaluation/adapter.py b/examples/vlm-evaluation/adapter.py index 351beae..ae42cff 100644 --- a/examples/vlm-evaluation/adapter.py +++ b/examples/vlm-evaluation/adapter.py @@ -42,17 +42,20 @@ autoregressively generate, and chat tasks are generation-only. A MoMa checkpoint fails fast in ``__init__``. -- **Image and video.** An image checkpoint evaluates exactly one image per - request; a video checkpoint (a ``[video]`` config) evaluates one video per - request — decoded to a fixed ``frames_per_clip`` clip via the training - frame-sampling policy — or one or more images packed as an ordered clip (a - single image is the 1-frame case; multiple images fill successive frame slots, - zero-padded and truncated to ``frames_per_clip``). Audio, multiple videos, - mixed image+video, multi-image on an image checkpoint, and multi-turn/few-shot - requests raise ``NotImplementedError``; ``loglikelihood`` and - ``generate_until_multi_round`` are not implemented (chat tasks are - generation-only). Visual input is modeled as an ordered list of frames (a - single image is the length-1 case). +- **Text, image, and video.** An image checkpoint evaluates a text-only request + (no visual) or exactly one image per request; a video checkpoint (a + ``[video]`` config) evaluates a text-only request, one video per request — + decoded to a fixed ``frames_per_clip`` clip via the training frame-sampling + policy — or one or more images packed as an ordered clip (a single image is + the 1-frame case; multiple images fill successive frame slots, zero-padded and + truncated to ``frames_per_clip``). A text-only request runs the generative + arch's pure-text forward (no visual encoder), measuring text-backbone + behavior. Audio, multiple videos, mixed image+video, multi-image on an image + checkpoint, and multi-turn/few-shot requests raise ``NotImplementedError``; + ``loglikelihood`` and ``generate_until_multi_round`` are not implemented (chat + tasks are generation-only). Visual input is modeled as an ordered list of + frames (a single image is the length-1 case; a text-only request is the empty + list). """ from __future__ import annotations @@ -210,16 +213,18 @@ def _render_request( ) -> tuple[list[Any], str]: """Flatten one chat request into ``(frames, prompt_text)``. - ``frames`` is an ordered list of visual frames; a single image is the - length-1 case and a decoded video is the multi-frame case. ``video_config`` - is the checkpoint's ``[video]`` config (``None`` for an image checkpoint) - and selects the mode: - - - **Image checkpoint** (``video_config is None``): exactly one image per - request; video content raises (an image model cannot evaluate video). - - **Video checkpoint**: one video — decoded to frames via - ``video_io.decode_video_frames`` using the checkpoint's frame-sampling - policy — or, when no video is present, one or more images treated as an + ``frames`` is an ordered list of visual frames; an empty list is a text-only + request, a single image is the length-1 case, and a decoded video is the + multi-frame case. ``video_config`` is the checkpoint's ``[video]`` config + (``None`` for an image checkpoint) and selects the mode: + + - **Image checkpoint** (``video_config is None``): a text-only request (no + image, empty frames) or exactly one image; multiple images and video raise + (an image model cannot evaluate video, and multi-image needs a video + checkpoint). + - **Video checkpoint**: a text-only request (empty frames), one video — + decoded to frames via ``video_io.decode_video_frames`` using the + checkpoint's frame-sampling policy — or one or more images treated as an ordered clip of frames (zero-padded to ``frames_per_clip`` downstream; frames beyond ``frames_per_clip`` are truncated with a warning). A single image is the length-1 clip. @@ -253,18 +258,19 @@ def _render_request( ) if video_config is None: - # Image checkpoint: image-only, exactly one image per request. + # Image checkpoint: a text-only request (no image) or exactly one image. if videos: raise NotImplementedError( "This is an image checkpoint (no [video] config) and cannot evaluate video. " "Use a video checkpoint, or report the task to the project owner." ) - if len(images) != 1: + if len(images) > 1: raise NotImplementedError( - f"This adapter supports exactly one image per request, got {len(images)}. " - "Multi-image and text-only requests are out of scope; report the task to " - "the project owner." + f"This image checkpoint supports at most one image per request, got " + f"{len(images)}. Multi-image is only supported on a video checkpoint; report " + "the task to the project owner." ) + # len(images) == 0 is a text-only request (empty frames); 1 is a single image. return images, prompt # Video checkpoint: a video (decoded to frames), or one or more images packed @@ -300,11 +306,8 @@ def _render_request( ) return frames, prompt if not images: - raise NotImplementedError( - "A video-checkpoint request must carry a video or at least one image, got no " - "visual content. Text-only requests are out of scope; report the task to the " - "project owner." - ) + # Text-only request on a video checkpoint (no video, no image): empty frames. + return [], prompt # One or more images on a video checkpoint: treat them as an ordered clip of frames # (zero-padded to frames_per_clip downstream, extra frames truncated), mirroring how # training packs short clips. A single image is the length-1 case. lmms-eval delivers @@ -417,7 +420,7 @@ def _first_stop(text: str, until: list[str]) -> int | None: def _generate_batch( model: VLMWrapper, tokenizer: Any, - pixel_values: torch.Tensor, + pixel_values: torch.Tensor | None, prompt_ids: list[torch.Tensor], resolved: dict[str, Any], max_seq_len: int, @@ -428,7 +431,9 @@ def _generate_batch( Decodes ``B`` requests together (``pixel_values`` is ``(B, 3, H, W)`` or a ``(B, F, 3, H, W)`` video clip, ``prompt_ids`` a list of ``B`` 1-D token tensors; ``frame_mask`` is ``(B, F)`` bool for video, masking padded-frame - visual tokens from attention as in training). There is no transformer KV + visual tokens from attention as in training). ``pixel_values is None`` is a + text-only batch: the vision tower is skipped, ``num_image_tokens`` is 0, and + ``model(None, ...)`` runs the pure-text forward. There is no transformer KV cache and no vision cache: ``model(...)`` re-runs over the growing **right-padded** batch each step, re-encoding the vision tower each time. Right-padding matches the training @@ -447,13 +452,15 @@ def _generate_batch( top_p: float = resolved["top_p"] eos_id = tokenizer.eos_token_id pad_id = resolve_pad_id(tokenizer) - device = pixel_values.device + # Derive device from the prompts (present for every request), not pixel_values, + # which is None on a text-only batch. + device = prompt_ids[0].device batch_size = len(prompt_ids) - # Length bound: image tokens (in-residual for JD/MoT; 0 for CA) + prompt + - # generated must fit the context. Reserve room for generation and left- - # truncate any over-budget prompt (per row). - num_image_tokens = model.num_image_tokens + # Length bound: image tokens (in-residual for JD/MoT; 0 for CA, and 0 for a + # text-only batch with no visual) + prompt + generated must fit the context. + # Reserve room for generation and left-truncate any over-budget prompt (per row). + num_image_tokens = 0 if pixel_values is None else model.num_image_tokens prompt_budget = max_seq_len - num_image_tokens - max_new_tokens if prompt_budget <= 0: raise _ContextBudgetError( @@ -595,6 +602,35 @@ def __init__( f"dtype={self._dtype}, max_seq_len={self._max_seq_len}" ) + def _decode_subbatch( + self, + pixel_values: torch.Tensor | None, + prompts: list[torch.Tensor], + resolved: dict[str, Any], + frame_mask: torch.Tensor | None, + task_name: str, + ) -> list[str]: + """Decode one homogeneous sub-batch (all-visual, or all-text-only with + ``pixel_values=None``), returning one continuation per prompt. + + A ``_ContextBudgetError`` (the task's gen_kwargs over-budget the context; + every request here shares gen_kwargs) is isolated to this sub-batch as empty + strings rather than aborting the whole run. + """ + try: + return _generate_batch( + self._model, + self._tokenizer, + pixel_values, + prompts, + resolved, + self._max_seq_len, + frame_mask=frame_mask, + ) + except _ContextBudgetError as exc: + logger.warning(f"Skipping {len(prompts)} request(s) for task {task_name}: {exc}") + return [""] * len(prompts) + def generate_until(self, requests: list[Instance]) -> list[str]: # Group requests by gen_kwargs (a batch must share decode params) and, # within a group, sort by context length so similar-length prompts batch @@ -617,9 +653,15 @@ def _collate(args: tuple[Any, ...]) -> int: # below; ``""`` = a request that failed to render/preprocess, isolated # with a warning so one bad doc does not abort the whole run. chunk_outputs: list[str | None] = [None] * len(chunk) - frames_batch: list[torch.Tensor] = [] - masks_batch: list[torch.Tensor] = [] - prompt_ids: list[torch.Tensor] = [] + # A chunk shares gen_kwargs but may now mix visual and text-only requests, + # which cannot share one pixel_values tensor — decode them as two + # sub-batches, each tagged by its original slot for order-preserving scatter. + vis_frames: list[torch.Tensor] = [] + vis_masks: list[torch.Tensor] = [] + vis_prompts: list[torch.Tensor] = [] + vis_slots: list[int] = [] + txt_prompts: list[torch.Tensor] = [] + txt_slots: list[int] = [] for slot, args in enumerate(chunk): try: # Chat 6-tuple: (context, doc_to_messages, gen_kwargs, doc_id, task, split). @@ -628,7 +670,7 @@ def _collate(args: tuple[Any, ...]) -> int: frames, prompt = _render_request(messages, self._config.video) # lmms-eval may deliver image content as a path/URL string; # normalize to PIL so both the image and video packers (strict - # pil_to_tensor) accept it. + # pil_to_tensor) accept it. Empty frames => text-only request. frames = [_to_pil(f) for f in frames] # Mirror training tokenization: no chat template, no # placeholder, add_special_tokens=False (images go via pixel_values). @@ -638,7 +680,12 @@ def _collate(args: tuple[Any, ...]) -> int: # with -100 labels (and trimmed by output_slice), so there is # no valid position to predict an image-only first token. raise ValueError("empty prompt after flattening (no text content)") - if self._is_video: + pixels: torch.Tensor | None + if not frames: + # Text-only request: no visual to pack. + pixels = None + mask = None + elif self._is_video: # Fixed (frames_per_clip, 3, H, W) clip + per-frame validity # mask, zero-padded — identical to training. The mask hides # padded-frame visual tokens from attention. @@ -662,49 +709,44 @@ def _collate(args: tuple[Any, ...]) -> int: ) chunk_outputs[slot] = "" continue - # Commit atomically: only reached when every step above succeeded, so a - # mid-request failure never leaves a partial entry in these lists. - frames_batch.append(pixels) - if mask is not None: - masks_batch.append(mask) - prompt_ids.append(prompt_tensor) - - if prompt_ids: - # Video: (B, F, 3, H, W) via stack (each request is one F-frame clip), - # with a (B, F) frame mask. Image: (B, 3, H, W) via cat. cat on video - # would fold frames into the batch and trip the frames-per-clip check. + # Commit atomically (only reached on full success), routed by modality: + # a text-only request (pixels is None) goes to the text-only sub-batch. + if pixels is None: + txt_prompts.append(prompt_tensor) + txt_slots.append(slot) + else: + vis_frames.append(pixels) + if mask is not None: + vis_masks.append(mask) + vis_prompts.append(prompt_tensor) + vis_slots.append(slot) + + # Visual sub-batch: Video stacks (B, F, 3, H, W) clips with a (B, F) frame + # mask; Image cats to (B, 3, H, W). cat on video would fold frames into the + # batch and trip the frames-per-clip check. Text-only sub-batch: no pixels. + if vis_prompts: if self._is_video: - pixel_values = torch.stack(frames_batch, dim=0) - frame_mask = torch.stack(masks_batch, dim=0) + pixel_values = torch.stack(vis_frames, dim=0) + frame_mask = torch.stack(vis_masks, dim=0) else: - pixel_values = torch.cat(frames_batch, dim=0) + pixel_values = torch.cat(vis_frames, dim=0) frame_mask = None - try: - gen_outputs = _generate_batch( - self._model, - self._tokenizer, - pixel_values, - prompt_ids, - resolved, - self._max_seq_len, - frame_mask=frame_mask, - ) - except _ContextBudgetError as exc: - # One task's gen_kwargs over-budgets the context; skip its requests - # (they all share gen_kwargs) rather than aborting the whole run. - logger.warning( - f"Skipping {len(prompt_ids)} request(s) for task {chunk[0][4]}: {exc}" - ) - gen_outputs = [""] * len(prompt_ids) - else: - gen_outputs = [] - # Scatter generated continuations back into the surviving (None) slots, - # preserving alignment with ``chunk`` (skipped slots keep their ""). - gen_iter = iter(gen_outputs) - outputs = [o if o is not None else next(gen_iter) for o in chunk_outputs] - for args, output in zip(chunk, outputs, strict=True): - results.append(output) - self.cache_hook.add_partial("generate_until", (args[0], args[2]), output) + vis_outputs = self._decode_subbatch( + pixel_values, vis_prompts, resolved, frame_mask, chunk[0][4] + ) + for s, o in zip(vis_slots, vis_outputs, strict=True): + chunk_outputs[s] = o + if txt_prompts: + txt_outputs = self._decode_subbatch(None, txt_prompts, resolved, None, chunk[0][4]) + for s, o in zip(txt_slots, txt_outputs, strict=True): + chunk_outputs[s] = o + + # Every surviving slot is now filled (a continuation, or "" for a skipped + # or over-budget request), preserving alignment with ``chunk``. + for args, output in zip(chunk, chunk_outputs, strict=True): + out = output if output is not None else "" + results.append(out) + self.cache_hook.add_partial("generate_until", (args[0], args[2]), out) pbar.update(len(chunk)) pbar.close() return re_ords.get_original(results) diff --git a/examples/vlm-evaluation/tests/unit/test_adapter.py b/examples/vlm-evaluation/tests/unit/test_adapter.py index 7a21ccf..4bd297d 100644 --- a/examples/vlm-evaluation/tests/unit/test_adapter.py +++ b/examples/vlm-evaluation/tests/unit/test_adapter.py @@ -165,10 +165,11 @@ def test_multi_image_raises(self): with pytest.raises(NotImplementedError, match="one image"): _render_request(messages, None) - def test_no_image_raises(self): - messages = _chat([_text("text only question")]) - with pytest.raises(NotImplementedError, match="one image"): - _render_request(messages, None) + def test_no_image_is_text_only(self): + """A text-only request on an image checkpoint renders as empty frames.""" + images, prompt = _render_request(_chat([_text("text only question")]), None) + assert images == [] + assert prompt == "text only question" def test_multi_turn_assistant_raises(self): messages = ChatMessages( @@ -230,6 +231,12 @@ def test_single_image_is_one_frame_clip(self, vcfg): assert len(frames) == 1 assert prompt == "q" + def test_no_visual_is_text_only(self, vcfg): + """A text-only request on a video checkpoint renders as empty frames.""" + frames, prompt = _render_request(_chat([_text("text only")]), vcfg) + assert frames == [] + assert prompt == "text only" + def test_audio_still_raises(self, vcfg): msg = _chat([_text("x"), {"type": "audio", "url": "a.wav"}]) with pytest.raises(NotImplementedError, match="[Aa]udio"): @@ -493,6 +500,14 @@ def test_overlong_prompt_is_left_truncated(self, arch_wrapper, monkeypatch): assert len(out) == 1 and len(out[0].split()) == 2 assert any("left-truncating" in m for m in rec.warnings) + def test_text_only_decode_no_pixels(self, arch_wrapper): + """pixel_values=None runs the pure-text forward for every generative arch + (JD/CA/MoT) end to end and returns a continuation of max_new_tokens.""" + r = _resolve_gen_kwargs({"max_new_tokens": 6}, 128) + out = _generate_batch(arch_wrapper, _MockTokenizer(), None, self._prompt(), r, 64) + assert len(out) == 1 and isinstance(out[0], str) + assert len(out[0].split()) == 6 + @pytest.mark.parametrize("arch", GENERATIVE_ARCHES) class TestGenerateBatchMulti: @@ -595,6 +610,30 @@ def test_generate_batch_no_frame_mask_for_image(): assert model.seen_frame_masks and all(fm is None for fm in model.seen_frame_masks) +def test_generate_batch_text_only_runs_with_none(): + """pixel_values=None decodes a text-only batch (device from prompts, no vision, + frame_mask=None each step).""" + model = _CaptureModel(num_image_tokens=8) + prompt_ids = [torch.tensor([5, 9, 12], dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128) + out = _generate_batch(model, _MockTokenizer(), None, prompt_ids, r, 64) + assert len(out) == 1 + assert model.seen_frame_masks and all(fm is None for fm in model.seen_frame_masks) + + +def test_generate_batch_text_only_budget_excludes_image_tokens(): + """A text-only batch reserves no image-token budget: a max_new_tokens that would + over-budget a visual batch (num_image_tokens=16) still fits when pixel_values is None.""" + model = _CaptureModel(num_image_tokens=16) + prompt_ids = [torch.tensor([5, 9], dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": 50}, 128) + # Visual budget: 64 - 16 - 50 <= 0 -> raises. Text-only: 64 - 0 - 50 = 14 -> fits. + with pytest.raises(_ContextBudgetError): + _generate_batch(model, _MockTokenizer(), torch.randn(1, 3, 16, 16), prompt_ids, r, 64) + out = _generate_batch(model, _MockTokenizer(), None, prompt_ids, r, 64) + assert len(out) == 1 + + # --------------------------------------------------------------------------- # Guards: arch + not-implemented methods # --------------------------------------------------------------------------- @@ -883,6 +922,45 @@ def doc_to_messages(doc): assert all(isinstance(o, str) for o in batched) assert all(len(o.split()) == 3 for o in batched) # greedy emits exactly max_new_tokens + @pytest.mark.parametrize("arch", GENERATIVE_ARCHES) + def test_mixed_text_only_and_image_batch(self, monkeypatch, tiny_vlm_configs, arch): + """A chunk mixing a text-only and an image request (same gen_kwargs) decodes + them as separate sub-batches and scatters results back in original order.""" + _patch_loaders( + monkeypatch, _vlm_job_config(tiny_vlm_configs, arch=arch), _vlm_wrapper(arch) + ) + vlm = KempnerForgeVLM( + config="x", checkpoint="y", device="cpu", dtype="float32", batch_size=2 + ) + img = _img() + + def doc_to_messages(doc): + content = [_text(doc["q"])] + if doc.get("img"): + content.append(_image(img)) + return [{"role": "user", "content": content}] + + vlm.task_dict = { + "t": {"test": {"d0": {"q": "text only?"}, "d1": {"q": "and this", "img": True}}} + } + specs = [("d0", "c"), ("d1", "cc")] + instances = [ + Instance( + request_type="generate_until", + arguments=(ctx, doc_to_messages, {"max_new_tokens": 3}, doc_id, "t", "test"), + idx=i, + metadata={"task": "t", "doc_id": doc_id, "repeats": 1}, + ) + for i, (doc_id, ctx) in enumerate(specs) + ] + + # d0 (text-only) and d1 (image) land in one chunk (shared gen_kwargs); each is + # decoded in its own sub-batch, so the mixed batch matches per-request singles. + singles = [vlm.generate_until([inst])[0] for inst in instances] + batched = vlm.generate_until(instances) + assert batched == singles + assert all(isinstance(o, str) and len(o.split()) == 3 for o in batched) + class TestGenerateUntilVideo: """End-to-end video decode: stack -> (B, F, 3, H, W) -> forward -> strings.""" diff --git a/kempnerforge/model/mot.py b/kempnerforge/model/mot.py index d16ef6d..bd5537e 100644 --- a/kempnerforge/model/mot.py +++ b/kempnerforge/model/mot.py @@ -151,9 +151,12 @@ def forward( t_m = x_m.shape[1] lengths[m] = t_m - q_m = self.q_proj[m](x_m).view(batch, t_m, -1, self.head_dim) - k_m = self.k_proj[m](x_m).view(batch, t_m, -1, self.head_dim) - v_m = self.v_proj[m](x_m).view(batch, t_m, -1, self.head_dim) + # Explicit head counts (not -1) so a zero-length modality stream + # (t_m == 0, e.g. a text-only forward's empty image stream) reshapes + # unambiguously; identical to -1 inference whenever t_m > 0. + q_m = self.q_proj[m](x_m).view(batch, t_m, self.n_heads, self.head_dim) + k_m = self.k_proj[m](x_m).view(batch, t_m, self.n_kv_heads, self.head_dim) + v_m = self.v_proj[m](x_m).view(batch, t_m, self.n_kv_heads, self.head_dim) if self.q_norm is not None: q_m = self.q_norm[m](q_m) @@ -201,7 +204,9 @@ def forward( offset = 0 for m in self.modalities: t_m = lengths[m] - o_m = out[:, offset : offset + t_m, :, :].reshape(batch, t_m, -1) + o_m = out[:, offset : offset + t_m, :, :].reshape( + batch, t_m, self.n_heads * self.head_dim + ) out_streams[m] = self.o_proj[m](o_m) offset += t_m return out_streams diff --git a/kempnerforge/model/transformer.py b/kempnerforge/model/transformer.py index 525ce1b..b5e6554 100644 --- a/kempnerforge/model/transformer.py +++ b/kempnerforge/model/transformer.py @@ -450,23 +450,31 @@ def forward( h = self.norm(h) # MoT path: position-based image-then-text split, per-modality # streams through the MoTBlock stack, single global SDPA per - # layer. modality_ids is required (presence + shape checked - # against the residual). v1 uses position-based routing; the - # tags are validated for shape but not value-matched against - # positions, so a future per-token scatter/gather can land - # without changing the public interface. + # layer. modality_ids is required when an image prefix is present + # (presence + shape checked against the residual); a text-only + # forward (no prefix, no ids) runs with n_image=0 (empty image + # stream). v1 uses position-based routing; the tags are validated + # for shape but not value-matched against positions, so a future + # per-token scatter/gather can land without changing the public + # interface. elif self._mot_modalities: if modality_ids is None: - raise ValueError( - "MoT model requires modality.modality_ids (got None). Build the " - "ModalityContext via MoTStrategy or set modality_ids explicitly." - ) - if modality_ids.shape != h.shape[:2]: - raise ValueError( - f"modality.modality_ids shape {tuple(modality_ids.shape)} does not " - f"match residual shape {tuple(h.shape[:2])}" - ) - n_image = self._mot_n_image + if prefix_embeds is not None: + # Image present (prefix) but routing tags missing: a real misuse. + raise ValueError( + "MoT model requires modality.modality_ids (got None). Build the " + "ModalityContext via MoTStrategy or set modality_ids explicitly." + ) + # Text-only forward (no image prefix): route every position through the + # text modality with an empty image stream (n_image=0). + n_image = 0 + else: + if modality_ids.shape != h.shape[:2]: + raise ValueError( + f"modality.modality_ids shape {tuple(modality_ids.shape)} does not " + f"match residual shape {tuple(h.shape[:2])}" + ) + n_image = self._mot_n_image t_image = n_image t_text = h.shape[1] - n_image streams: dict[str, torch.Tensor] = { @@ -500,13 +508,10 @@ def forward( ) if ca_iter is not None and (i + 1) % self._ca_cadence == 0: ca = next(ca_iter, None) - if ca is not None: - if image_features is None: - raise ValueError( - "Cross-Attention block fired but modality.image_features is None. " - "Cross-Attention models require image_features in the " - "ModalityContext." - ) + # image_features is None on a text-only request: skip the + # cross-attention block, leaving the pure text backbone. Existing + # image forwards always pass image_features, so they are unchanged. + if ca is not None and image_features is not None: if image_features.dtype != h.dtype: image_features = image_features.to(h.dtype) h = ca(h, image_features, image_mask) diff --git a/kempnerforge/model/vlm.py b/kempnerforge/model/vlm.py index d6dfd74..660f431 100644 --- a/kempnerforge/model/vlm.py +++ b/kempnerforge/model/vlm.py @@ -66,10 +66,18 @@ class ModalityStrategy(Protocol): def prepare( self, wrapper: VLMWrapper, - pixel_values: torch.Tensor, + pixel_values: torch.Tensor | None, input_ids: torch.Tensor, frame_mask: torch.Tensor | None = None, - ) -> ModalityContext: ... + ) -> ModalityContext: + """Compose a ``ModalityContext`` from raw VLM inputs. + + ``pixel_values is None`` is a text-only request (no visual content): + the generative arches return a context that drives the pure-text + forward — an empty ``ModalityContext()`` for the image-prefix and + cross-attention arches — while a non-generative arch may reject it. + """ + ... def num_image_tokens(self, wrapper: VLMWrapper) -> int: ... @@ -179,10 +187,15 @@ class JointDecoderStrategy: def prepare( self, wrapper: VLMWrapper, - pixel_values: torch.Tensor, + pixel_values: torch.Tensor | None, input_ids: torch.Tensor, # noqa: ARG002 frame_mask: torch.Tensor | None = None, ) -> ModalityContext: + if pixel_values is None: + # Text-only request: no image prefix, so the residual carries text + # alone and the LM head runs over every position (empty context -> + # Transformer.forward's pure-text path). + return ModalityContext() img_embeds = _project_visual_features(wrapper, pixel_values) n = img_embeds.shape[1] # pooling-aware: the adapter's actual visual-token count return ModalityContext( @@ -213,10 +226,15 @@ class CrossAttentionStrategy: def prepare( self, wrapper: VLMWrapper, - pixel_values: torch.Tensor, + pixel_values: torch.Tensor | None, input_ids: torch.Tensor, # noqa: ARG002 frame_mask: torch.Tensor | None = None, ) -> ModalityContext: + if pixel_values is None: + # Text-only request: no image K/V. The cross-attention blocks are + # skipped in Transformer.forward when image_features is None, leaving + # the pure text backbone. + return ModalityContext() img_embeds = _project_visual_features(wrapper, pixel_values) return ModalityContext( image_features=img_embeds, @@ -251,10 +269,16 @@ class MoTStrategy: def prepare( self, wrapper: VLMWrapper, - pixel_values: torch.Tensor, + pixel_values: torch.Tensor | None, input_ids: torch.Tensor, frame_mask: torch.Tensor | None = None, ) -> ModalityContext: + if pixel_values is None: + # Text-only request: no image prefix. The MoT forward runs with + # n_image=0 (an empty image stream), routing every position through + # the text-modality projections/FFN; no modality_ids are needed (see + # Transformer.forward's MoT branch). + return ModalityContext() img_embeds = _project_visual_features(wrapper, pixel_values) n = img_embeds.shape[1] # pooling-aware: the adapter's actual visual-token count b, t_text = input_ids.shape @@ -297,10 +321,18 @@ class MoMaStrategy: def prepare( self, wrapper: VLMWrapper, - pixel_values: torch.Tensor, + pixel_values: torch.Tensor | None, input_ids: torch.Tensor, frame_mask: torch.Tensor | None = None, ) -> ModalityContext: + if pixel_values is None: + # MoMa's expert-choice routing is non-causal (is_generative=False), so it + # is excluded from generative/text-only evaluation; fail fast rather than + # emit an unusable context. + raise NotImplementedError( + "Text-only forward is not supported for the 'moma' arch " + "(non-causal expert-choice routing; excluded from generative evaluation)." + ) img_embeds = _project_visual_features(wrapper, pixel_values) n = img_embeds.shape[1] # pooling-aware: the adapter's actual visual-token count b, t_text = input_ids.shape @@ -335,8 +367,9 @@ class VLMWrapper(nn.Module): Forward: ``(pixel_values, input_ids, labels) -> (logits, labels)``. The strategy composes a ``ModalityContext`` from the raw inputs and the wrapper's submodules; ``Transformer.forward`` consumes the - context. ``num_image_tokens`` is arch-aware and delegates to the - strategy. + context. ``pixel_values`` may be ``None`` for a text-only (no visual) + request — the generative arches route it to the pure-text forward. + ``num_image_tokens`` is arch-aware and delegates to the strategy. """ def __init__( @@ -370,7 +403,7 @@ def num_image_tokens(self) -> int: def forward( self, - pixel_values: torch.Tensor, + pixel_values: torch.Tensor | None, input_ids: torch.Tensor, labels: torch.Tensor | None = None, frame_mask: torch.Tensor | None = None, diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py index 93a49a6..cad9751 100644 --- a/tests/unit/test_model.py +++ b/tests/unit/test_model.py @@ -693,6 +693,44 @@ def test_modality_ids_required_when_mot_active(self): with pytest.raises(ValueError, match="MoT model requires modality.modality_ids"): model(tokens, modality=ModalityContext(prefix_embeds=prefix)) + def test_text_only_forward_no_image(self): + """A text-only MoT forward (no prefix, no modality_ids) runs with + n_image=0: every position routes through the text modality over an empty + image stream. modality=None and an empty ModalityContext agree.""" + mc, lc, n_image, n_text = _mot_setup() + model = Transformer(mc, vlm_config=lc, num_image_tokens=n_image).to(DEVICE).eval() + tokens = torch.randint(0, 256, (1, n_text), device=DEVICE) + with torch.no_grad(): + out_ctx = model(tokens, modality=ModalityContext()) + out_none = model(tokens) + assert out_ctx.shape == (1, n_text, mc.vocab_size) + assert torch.isfinite(out_ctx).all() + torch.testing.assert_close(out_ctx, out_none) + + def test_text_only_forward_with_moe(self): + """MoT + MoE text-only: the per-modality image FFN (an MoE router) runs + over the empty image stream (0 tokens) without error; output is finite.""" + from kempnerforge.config.vlm import MoTConfig + + mc = ModelConfig( + dim=64, + n_layers=2, + n_heads=4, + vocab_size=256, + max_seq_len=64, + ffn_hidden_dim=128, + num_experts=4, + moe_top_k=2, + moe_frequency=1, + ) + model = Transformer(mc, vlm_config=MoTConfig(max_text_len=32), num_image_tokens=8) + model = model.to(DEVICE).eval() + tokens = torch.randint(0, 256, (1, 16), device=DEVICE) + with torch.no_grad(): + out = model(tokens, modality=ModalityContext()) + assert out.shape == (1, 16, 256) + assert torch.isfinite(out).all() + def test_modality_ids_shape_mismatch_raises(self): mc, lc, n_image, max_text_len = _mot_setup() model = Transformer(mc, vlm_config=lc, num_image_tokens=n_image).to(DEVICE).eval() @@ -941,13 +979,45 @@ def test_ca_zero_init_residual_at_construction(self): out_text = text_model(tokens) torch.testing.assert_close(out_ca, out_text, atol=1e-5, rtol=1e-5) - def test_image_features_required_when_ca_layers_present(self): - """Forward without image_features raises a clear error when - CA layers are configured.""" + def test_text_only_skips_cross_attention(self): + """A text-only forward (no image_features) skips the CA blocks and + returns text-backbone logits instead of raising.""" model = _ca_transformer(n_layers=4, cadence=2).to(DEVICE).eval() tokens = torch.randint(0, 256, (1, 8), device=DEVICE) - with pytest.raises(ValueError, match="image_features is None"): - model(tokens) + with torch.no_grad(): + out = model(tokens) + assert out.shape == (1, 8, 256) + assert torch.isfinite(out).all() + + def test_text_only_equals_text_backbone_with_perturbed_ca(self): + """A CA model's text-only forward equals a plain text Transformer with + the same backbone weights even when the CA blocks are non-identity: the + CA blocks are skipped, not applied.""" + torch.manual_seed(0) + mc_ca, lc_ca, n_image = _ca_setup(n_layers=4, cadence=2) + text_config = ModelConfig( + dim=mc_ca.dim, + n_layers=mc_ca.n_layers, + n_heads=mc_ca.n_heads, + vocab_size=mc_ca.vocab_size, + max_seq_len=mc_ca.max_seq_len, + ) + ca_model = Transformer(mc_ca, vlm_config=lc_ca, num_image_tokens=n_image).to(DEVICE).eval() + text_model = Transformer(text_config).to(DEVICE).eval() + text_state = { + k: v for k, v in ca_model.state_dict().items() if "cross_attention_layers" not in k + } + text_model.load_state_dict(text_state, strict=True) + # Perturb CA weights so they are NOT the zero-init identity; a text-only + # forward must still skip them (equal the text backbone), not apply them. + with torch.no_grad(): + for p in ca_model.cross_attention_layers.parameters(): + p.copy_(torch.randn_like(p)) + tokens = torch.randint(0, 256, (1, 8), device=DEVICE) + with torch.no_grad(): + out_ca = ca_model(tokens) # text-only: CA skipped + out_text = text_model(tokens) + torch.testing.assert_close(out_ca, out_text, atol=1e-5, rtol=1e-5) def test_image_features_with_kv_caches_raises(self): """Cross-arg invariant: image_features is training-only.""" diff --git a/tests/unit/test_mot.py b/tests/unit/test_mot.py index 5855262..bc43575 100644 --- a/tests/unit/test_mot.py +++ b/tests/unit/test_mot.py @@ -150,6 +150,21 @@ def test_output_shape_per_modality(self, batch: int, t_image: int, t_text: int): assert out["image"].shape == (batch, t_image, 64) assert out["text"].shape == (batch, t_text, 64) + def test_empty_image_stream_runs(self): + """A zero-length image stream (a text-only forward's n_image=0) flows + through MoTAttention: the per-modality projections, the concat for the + global SDPA, and the split-back all handle the 0-length modality.""" + attn = MoTAttention(dim=64, n_heads=4, n_kv_heads=2, modalities=("image", "text")) + for m in ("image", "text"): + nn.init.normal_(attn.o_proj[m].weight) + cos_t, sin_t = _rope_for(10, 16) + streams = {"image": torch.randn(2, 0, 64), "text": torch.randn(2, 10, 64)} + rope = {"image": (cos_t[:0], sin_t[:0]), "text": (cos_t, sin_t)} + out = attn(streams, rope) + assert out["image"].shape == (2, 0, 64) + assert out["text"].shape == (2, 10, 64) + assert torch.isfinite(out["text"]).all() + def test_output_dtype_matches_input(self): attn = MoTAttention(dim=64, n_heads=4, n_kv_heads=4, modalities=("image", "text")).to( dtype=torch.float32 @@ -355,6 +370,22 @@ def test_block_forward_shape_per_modality(self): assert out["image"].shape == (2, 6, 64) assert out["text"].shape == (2, 10, 64) + def test_block_empty_image_stream_runs(self): + """MoTBlock with a zero-length image stream (text-only n_image=0) runs + end to end (per-modality norms + attention + per-modality FFN).""" + cfg = _config(dim=64, n_heads=4) + block = MoTBlock(cfg, modalities=("image", "text"), layer_idx=0) + for m in ("image", "text"): + nn.init.normal_(block.attn.o_proj[m].weight) + nn.init.normal_(block.mlp[m].down_proj.weight) # type: ignore[union-attr] + cos_t, sin_t = _rope_for(10, 16) + streams = {"image": torch.randn(1, 0, 64), "text": torch.randn(1, 10, 64)} + rope = {"image": (cos_t[:0], sin_t[:0]), "text": (cos_t, sin_t)} + out = block(streams, rope) + assert out["image"].shape == (1, 0, 64) + assert out["text"].shape == (1, 10, 64) + assert torch.isfinite(out["text"]).all() + def test_block_zero_init_residual_identity(self): """At construction, MoTBlock(streams) bit-equal to streams.""" cfg = _config(dim=64, n_heads=4) diff --git a/tests/unit/test_vlm.py b/tests/unit/test_vlm.py index e9fd04c..6556fa7 100644 --- a/tests/unit/test_vlm.py +++ b/tests/unit/test_vlm.py @@ -25,6 +25,7 @@ from kempnerforge.model.vlm import ( CrossAttentionStrategy, JointDecoderStrategy, + MoMaStrategy, MoTStrategy, VLMWrapper, _is_encoder_frozen, @@ -80,6 +81,35 @@ def test_labels_none_passthrough(self): assert logits.shape == (1, 10, 256) assert labels_out is None + def test_text_only_forward_jd(self): + """pixel_values=None routes JD to a pure-text forward (no image prefix).""" + wrapper = _build_tiny_wrapper(num_image_tokens=8).to(DEVICE) + input_ids = torch.randint(0, 256, (2, 12), device=DEVICE) + logits, labels_out = wrapper(None, input_ids) + assert logits.shape == (2, 12, 256) + assert labels_out is None + assert torch.isfinite(logits).all() + + def test_text_only_forward_ca(self): + wrapper = _build_ca_tiny_wrapper(num_image_tokens=8).to(DEVICE) + input_ids = torch.randint(0, 256, (2, 12), device=DEVICE) + logits, _ = wrapper(None, input_ids) + assert logits.shape == (2, 12, 256) + assert torch.isfinite(logits).all() + + def test_text_only_forward_mot(self): + wrapper = _build_mot_tiny_wrapper(num_image_tokens=8).to(DEVICE) + input_ids = torch.randint(0, 256, (2, 12), device=DEVICE) + logits, _ = wrapper(None, input_ids) + assert logits.shape == (2, 12, 256) + assert torch.isfinite(logits).all() + + def test_text_only_forward_moma_raises(self): + wrapper = _build_moma_tiny_wrapper(num_image_tokens=8).to(DEVICE) + input_ids = torch.randint(0, 256, (1, 12), device=DEVICE) + with pytest.raises(NotImplementedError, match="moma"): + wrapper(None, input_ids) + def test_dtype_mismatch_cast(self): """Vision encoder output in fp32, transformer in bf16 -> forward still works; the cast happens inside VLMWrapper before concat.""" @@ -287,6 +317,15 @@ def _build_mot_tiny_wrapper(num_image_tokens: int = 8, feature_dim: int = 96) -> return build_vlm_wrapper(mc, vc, ac, lc) +def _build_moma_tiny_wrapper(num_image_tokens: int = 8, feature_dim: int = 96) -> VLMWrapper: + mc = ModelConfig( + dim=64, n_layers=2, n_heads=4, vocab_size=256, max_seq_len=64, ffn_hidden_dim=128 + ) + vc = VisionEncoderConfig(type="random", feature_dim=feature_dim, num_tokens=num_image_tokens) + lc = MoMaConfig(max_text_len=32, moma_experts_per_modality={"image": 2, "text": 2}) + return build_vlm_wrapper(mc, vc, AdapterConfig(), lc) + + class TestModalityStrategies: def test_joint_decoder_strategy_fills_prefix_and_slice(self): wrapper = _build_tiny_wrapper(num_image_tokens=8) @@ -356,6 +395,37 @@ def test_strategy_num_image_tokens_arch_specific(self): # CA: residual stream is text-only, so no extension. assert ca_wrapper.num_image_tokens == 0 + def test_joint_decoder_strategy_text_only_returns_empty_context(self): + """pixel_values=None (text-only): empty context -> pure-text forward.""" + wrapper = _build_tiny_wrapper(num_image_tokens=8) + input_ids = torch.randint(0, 256, (1, 16)) + ctx = JointDecoderStrategy().prepare(wrapper, None, input_ids) + assert ctx.prefix_embeds is None + assert ctx.output_slice is None + assert ctx.key_padding_mask is None + + def test_cross_attention_strategy_text_only_returns_empty_context(self): + wrapper = _build_ca_tiny_wrapper(num_image_tokens=8) + input_ids = torch.randint(0, 256, (1, 16)) + ctx = CrossAttentionStrategy().prepare(wrapper, None, input_ids) + assert ctx.image_features is None + assert ctx.image_mask is None + + def test_mot_strategy_text_only_returns_empty_context(self): + """Text-only MoT carries no modality_ids: the forward runs n_image=0.""" + wrapper = _build_mot_tiny_wrapper(num_image_tokens=8) + input_ids = torch.randint(0, 256, (1, 16)) + ctx = MoTStrategy().prepare(wrapper, None, input_ids) + assert ctx.prefix_embeds is None + assert ctx.modality_ids is None + + def test_moma_strategy_text_only_raises(self): + """MoMa is non-generative and excluded from text-only evaluation.""" + wrapper = _build_moma_tiny_wrapper(num_image_tokens=8) + input_ids = torch.randint(0, 256, (1, 16)) + with pytest.raises(NotImplementedError, match="moma"): + MoMaStrategy().prepare(wrapper, None, input_ids) + class TestVLMWrapperDispatch: def test_build_modality_strategy_joint_decoder(self): From b293b5f390cc62c8dc41cef6c81d6d2c3ff416ea Mon Sep 17 00:00:00 2001 From: camilobrownpinilla Date: Wed, 15 Jul 2026 14:21:15 -0400 Subject: [PATCH 3/4] Fix llm-judged benchmarks + clean comments --- examples/vlm-evaluation/vlm_eval_harness.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/examples/vlm-evaluation/vlm_eval_harness.py b/examples/vlm-evaluation/vlm_eval_harness.py index 2add3ea..d42481b 100644 --- a/examples/vlm-evaluation/vlm_eval_harness.py +++ b/examples/vlm-evaluation/vlm_eval_harness.py @@ -38,6 +38,7 @@ import json import logging import sys +import tempfile from pathlib import Path logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") @@ -69,8 +70,6 @@ def main() -> None: parser.add_argument( "--checkpoint", type=str, required=True, help="DCP checkpoint dir (run dir or step_N dir)" ) - # No default task suite: the representative default benchmark set is an open - # decision; --tasks is required until one is provided. parser.add_argument( "--tasks", type=str, required=True, help="Comma-separated lmms-eval task names" ) @@ -133,10 +132,22 @@ def main() -> None: max_new_tokens=args.max_new_tokens, **dtype_kwargs, ) + # lmms-eval's evaluate() reads a few attributes off a `cli_args` namespace, + # and some tasks dereference them directly rather than via getattr: judge- + # scored tasks cache their GPT responses under `cli_args.output_path`, and + # hallusion_bench's aggregation crashes on `args.output_path` when + # `task.args` is None (which is what a missing cli_args leaves it as). + judge_output_dir = Path(args.output).parent if args.output else Path(tempfile.mkdtemp()) + judge_output_dir.mkdir(parents=True, exist_ok=True) + cli_args = argparse.Namespace( + output_path=str(judge_output_dir), + process_with_media=True, + ) results = simple_evaluate( model=model, tasks=args.tasks.split(","), limit=args.limit, + cli_args=cli_args, ) # --- Print results --- From c803dd04b5212c8ef2d83221a111469a95fb9fba Mon Sep 17 00:00:00 2001 From: camilobrownpinilla Date: Wed, 15 Jul 2026 21:34:21 -0400 Subject: [PATCH 4/4] Wire lmms throughput metrics through adapter --- examples/vlm-evaluation/adapter.py | 62 +++++--- .../integration/test_lmms_eval_contract.py | 43 +++++- .../tests/integration/test_vlm_eval.py | 12 +- .../tests/unit/_fake_lmms_eval.py | 93 ++++++++++- .../vlm-evaluation/tests/unit/test_adapter.py | 144 +++++++++++++----- 5 files changed, 295 insertions(+), 59 deletions(-) diff --git a/examples/vlm-evaluation/adapter.py b/examples/vlm-evaluation/adapter.py index ae42cff..350da30 100644 --- a/examples/vlm-evaluation/adapter.py +++ b/examples/vlm-evaluation/adapter.py @@ -61,13 +61,15 @@ from __future__ import annotations import json +import time from pathlib import Path from typing import Any import torch import torch.distributed.checkpoint as dcp -from lmms_eval.api.instance import Instance +from lmms_eval.api.instance import GenerationResult, Instance, TokenCounts from lmms_eval.api.model import lmms +from lmms_eval.models.model_utils.gen_metrics import log_metrics from lmms_eval.protocol import ChatMessages from lmms_eval.utils import Collator from tqdm import tqdm @@ -425,8 +427,9 @@ def _generate_batch( resolved: dict[str, Any], max_seq_len: int, frame_mask: torch.Tensor | None = None, -) -> list[str]: - """Batched decode (no transformer KV cache); returns one continuation per request. +) -> tuple[list[str], list[int]]: + """Batched decode (no transformer KV cache); returns one continuation per request + plus its generated token count. Decodes ``B`` requests together (``pixel_values`` is ``(B, 3, H, W)`` or a ``(B, F, 3, H, W)`` video clip, ``prompt_ids`` a list of ``B`` 1-D token @@ -523,7 +526,8 @@ def _generate_batch( text = tokenizer.decode(tokens, skip_special_tokens=True) cut = _first_stop(text, until) outputs.append(text[:cut] if cut is not None else text) - return outputs + gen_counts = [len(g) for g in generated] + return outputs, gen_counts # --------------------------------------------------------------------------- # @@ -609,9 +613,10 @@ def _decode_subbatch( resolved: dict[str, Any], frame_mask: torch.Tensor | None, task_name: str, - ) -> list[str]: + ) -> tuple[list[str], list[int]]: """Decode one homogeneous sub-batch (all-visual, or all-text-only with - ``pixel_values=None``), returning one continuation per prompt. + ``pixel_values=None``), returning one continuation per prompt and its + generated token count. A ``_ContextBudgetError`` (the task's gen_kwargs over-budget the context; every request here shares gen_kwargs) is isolated to this sub-batch as empty @@ -629,9 +634,9 @@ def _decode_subbatch( ) except _ContextBudgetError as exc: logger.warning(f"Skipping {len(prompts)} request(s) for task {task_name}: {exc}") - return [""] * len(prompts) + return [""] * len(prompts), [0] * len(prompts) - def generate_until(self, requests: list[Instance]) -> list[str]: + def generate_until(self, requests: list[Instance]) -> list[GenerationResult]: # Group requests by gen_kwargs (a batch must share decode params) and, # within a group, sort by context length so similar-length prompts batch # together (less padding). Collator.get_original restores request order. @@ -644,7 +649,9 @@ def _collate(args: tuple[Any, ...]) -> int: group_fn=lambda args: args[2], # args[2] == gen_kwargs grouping=True, ) - results: list[str] = [] + results: list[GenerationResult] = [] + total_gen_tokens = 0 + decode_elapsed = 0.0 pbar = tqdm(total=len(requests), disable=(self.rank != 0), desc="KempnerForge VLM") for chunk in re_ords.get_batched(n=self._batch_size, batch_fn=None): # Every request in the chunk shares gen_kwargs (index 2); resolve once. @@ -653,6 +660,8 @@ def _collate(args: tuple[Any, ...]) -> int: # below; ``""`` = a request that failed to render/preprocess, isolated # with a warning so one bad doc does not abort the whole run. chunk_outputs: list[str | None] = [None] * len(chunk) + # Parallel per-slot generated token counts (``None`` until filled). + chunk_counts: list[int | None] = [None] * len(chunk) # A chunk shares gen_kwargs but may now mix visual and text-only requests, # which cannot share one pixel_values tensor — decode them as two # sub-batches, each tagged by its original slot for order-preserving scatter. @@ -708,6 +717,7 @@ def _collate(args: tuple[Any, ...]) -> int: f"{type(exc).__name__}: {exc}" ) chunk_outputs[slot] = "" + chunk_counts[slot] = 0 continue # Commit atomically (only reached on full success), routed by modality: # a text-only request (pixels is None) goes to the text-only sub-batch. @@ -731,24 +741,42 @@ def _collate(args: tuple[Any, ...]) -> int: else: pixel_values = torch.cat(vis_frames, dim=0) frame_mask = None - vis_outputs = self._decode_subbatch( + _t = time.perf_counter() + vis_outputs, vis_counts = self._decode_subbatch( pixel_values, vis_prompts, resolved, frame_mask, chunk[0][4] ) - for s, o in zip(vis_slots, vis_outputs, strict=True): + decode_elapsed += time.perf_counter() - _t + for s, o, c in zip(vis_slots, vis_outputs, vis_counts, strict=True): chunk_outputs[s] = o + chunk_counts[s] = c if txt_prompts: - txt_outputs = self._decode_subbatch(None, txt_prompts, resolved, None, chunk[0][4]) - for s, o in zip(txt_slots, txt_outputs, strict=True): + _t = time.perf_counter() + txt_outputs, txt_counts = self._decode_subbatch( + None, txt_prompts, resolved, None, chunk[0][4] + ) + decode_elapsed += time.perf_counter() - _t + for s, o, c in zip(txt_slots, txt_outputs, txt_counts, strict=True): chunk_outputs[s] = o + chunk_counts[s] = c # Every surviving slot is now filled (a continuation, or "" for a skipped # or over-budget request), preserving alignment with ``chunk``. - for args, output in zip(chunk, chunk_outputs, strict=True): - out = output if output is not None else "" - results.append(out) - self.cache_hook.add_partial("generate_until", (args[0], args[2]), out) + for args, output, count in zip(chunk, chunk_outputs, chunk_counts, strict=True): + text = output if output is not None else "" + n_tok = count if count is not None else 0 + total_gen_tokens += n_tok + results.append( + GenerationResult(text=text, token_counts=TokenCounts(output_tokens=n_tok)) + ) + self.cache_hook.add_partial("generate_until", (args[0], args[2]), text) pbar.update(len(chunk)) pbar.close() + avg_speed = total_gen_tokens / decode_elapsed if decode_elapsed > 0 else 0.0 + log_metrics( + total_elapsed_time=decode_elapsed, + total_gen_tokens=total_gen_tokens, + avg_speed=avg_speed, + ) return re_ords.get_original(results) def loglikelihood(self, requests: list[Instance]) -> list[tuple[float, bool]]: diff --git a/examples/vlm-evaluation/tests/integration/test_lmms_eval_contract.py b/examples/vlm-evaluation/tests/integration/test_lmms_eval_contract.py index ee52cfe..1c04778 100644 --- a/examples/vlm-evaluation/tests/integration/test_lmms_eval_contract.py +++ b/examples/vlm-evaluation/tests/integration/test_lmms_eval_contract.py @@ -19,8 +19,18 @@ "fake lmms_eval is active; skipping real-package contract tests", allow_module_level=True ) -from lmms_eval.api.instance import Instance # noqa: E402 +from lmms_eval.api.instance import ( # noqa: E402 + GenerationResult, + Instance, + TokenCounts, + unwrap_generation_output, +) from lmms_eval.api.model import lmms # noqa: E402 +from lmms_eval.models.model_utils.gen_metrics import ( # noqa: E402 + log_metrics, + reset_logged_metrics, + summarize_logged_metrics, +) from lmms_eval.protocol import ChatMessages # noqa: E402 from lmms_eval.utils import Collator # noqa: E402 @@ -90,3 +100,34 @@ def test_args_returns_arguments_tuple(self): metadata={"task": "t", "doc_id": "d0", "repeats": 1}, ) assert inst.args == ("ctx", None, {}, "d0", "t", "test") + + +class TestGenerationResultContract: + """Pins the typed generate_until return + per-sample token counters the adapter now emits.""" + + def test_generation_result_and_token_counts_fields(self): + gr = GenerationResult(text="a b c", token_counts=TokenCounts(output_tokens=3)) + assert gr.text == "a b c" + assert gr.token_counts.output_tokens == 3 + # to_dict drops None fields — the shape build_efficiency_summary consumes. + assert gr.token_counts.to_dict() == {"output_tokens": 3} + + def test_unwrap_generation_output_handles_str_and_wrapper(self): + text, tc = unwrap_generation_output( + GenerationResult(text="x", token_counts=TokenCounts(output_tokens=1)) + ) + assert text == "x" and tc.output_tokens == 1 + # A bare string (the pre-instrumentation return) unwraps to (text, None). + assert unwrap_generation_output("plain") == ("plain", None) + + +class TestGenMetricsContract: + """Pins the throughput sink the adapter calls; the evaluator resets/summarizes this.""" + + def test_log_metrics_feeds_summary(self): + reset_logged_metrics() + log_metrics(total_elapsed_time=2.0, total_gen_tokens=10, avg_speed=5.0) + summary = summarize_logged_metrics() + assert summary["total_gen_tokens"] == 10 + assert summary["total_elapsed_time"] == 2.0 + reset_logged_metrics() diff --git a/examples/vlm-evaluation/tests/integration/test_vlm_eval.py b/examples/vlm-evaluation/tests/integration/test_vlm_eval.py index e3ae8c8..fc23485 100644 --- a/examples/vlm-evaluation/tests/integration/test_vlm_eval.py +++ b/examples/vlm-evaluation/tests/integration/test_vlm_eval.py @@ -36,7 +36,7 @@ pytest.skip("fake lmms_eval is active; skipping real-package tests", allow_module_level=True) from adapter import KempnerForgeVLM # noqa: E402 -from lmms_eval.api.instance import Instance # noqa: E402 +from lmms_eval.api.instance import GenerationResult, Instance # noqa: E402 from kempnerforge.config.data import DataConfig # noqa: E402 from kempnerforge.config.schema import JobConfig # noqa: E402 @@ -110,8 +110,9 @@ def doc_to_messages(doc): outputs = vlm.generate_until(instances) assert isinstance(outputs, list) and len(outputs) == 2 - assert all(isinstance(o, str) for o in outputs) - assert all(len(o.split()) == 3 for o in outputs) # greedy emits exactly max_new_tokens + assert all(isinstance(o, GenerationResult) for o in outputs) + assert all(len(o.text.split()) == 3 for o in outputs) # greedy emits exactly max_new_tokens + assert all(o.token_counts.output_tokens == 3 for o in outputs) # per-sample output count def test_dcp_roundtrip_video_generate_until(tmp_path, tiny_video_configs, monkeypatch): @@ -177,8 +178,9 @@ def doc_to_messages(doc): outputs = vlm.generate_until(instances) assert isinstance(outputs, list) and len(outputs) == 2 - assert all(isinstance(o, str) for o in outputs) - assert all(len(o.split()) == 3 for o in outputs) # greedy emits exactly max_new_tokens + assert all(isinstance(o, GenerationResult) for o in outputs) + assert all(len(o.text.split()) == 3 for o in outputs) # greedy emits exactly max_new_tokens + assert all(o.token_counts.output_tokens == 3 for o in outputs) # per-sample output count @pytest.mark.skipif( diff --git a/examples/vlm-evaluation/tests/unit/_fake_lmms_eval.py b/examples/vlm-evaluation/tests/unit/_fake_lmms_eval.py index ae0bedf..458bfcf 100644 --- a/examples/vlm-evaluation/tests/unit/_fake_lmms_eval.py +++ b/examples/vlm-evaluation/tests/unit/_fake_lmms_eval.py @@ -18,6 +18,10 @@ - ``api.model.lmms``: base sets ``_rank=0/_world_size=1/cache_hook/task_dict`` and exposes ``rank``/``world_size`` properties. - ``api.instance.Instance``: dataclass exposing ``.args`` (the arguments tuple). +- ``api.instance.GenerationResult`` / ``TokenCounts``: the typed ``generate_until`` + return and per-request token counters (``.text`` / ``.token_counts.output_tokens``). +- ``models.model_utils.gen_metrics.log_metrics`` (+ ``reset_logged_metrics`` / + ``summarize_logged_metrics``): the throughput history the evaluator resets/summarizes. """ from __future__ import annotations @@ -172,6 +176,74 @@ def args(self) -> tuple: return self.arguments if isinstance(self.arguments, tuple) else (self.arguments,) +# --------------------------------------------------------------------------- # +# lmms_eval.api.instance.TokenCounts / GenerationResult +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass +class TokenCounts: + input_tokens: int | None = None + output_tokens: int | None = None + reasoning_tokens: int | None = None + + def to_dict(self) -> dict[str, int | None]: + d: dict[str, int | None] = {} + if self.input_tokens is not None: + d["input_tokens"] = self.input_tokens + if self.output_tokens is not None: + d["output_tokens"] = self.output_tokens + if self.reasoning_tokens is not None: + d["reasoning_tokens"] = self.reasoning_tokens + return d + + +@dataclasses.dataclass +class GenerationResult: + text: str + token_counts: TokenCounts | None = None + + +# --------------------------------------------------------------------------- # +# lmms_eval.models.model_utils.gen_metrics (throughput history) +# --------------------------------------------------------------------------- # + +_THROUGHPUT_METRICS_HISTORY: list[dict[str, Any]] = [] + + +def reset_logged_metrics() -> None: + _THROUGHPUT_METRICS_HISTORY.clear() + + +def log_metrics( + total_elapsed_time: float, + total_gen_tokens: int, + avg_speed: float, + additional_metrics: dict[str, Any] | None = None, +) -> None: + payload: dict[str, Any] = { + "total_elapsed_time": total_elapsed_time, + "total_gen_tokens": total_gen_tokens, + "avg_speed": avg_speed, + } + if additional_metrics: + payload.update(additional_metrics) + _THROUGHPUT_METRICS_HISTORY.append(payload) + + +def summarize_logged_metrics() -> dict[str, Any]: + if not _THROUGHPUT_METRICS_HISTORY: + return {} + total_gen_tokens = sum(m.get("total_gen_tokens", 0) for m in _THROUGHPUT_METRICS_HISTORY) + total_elapsed_time = sum(m.get("total_elapsed_time", 0.0) for m in _THROUGHPUT_METRICS_HISTORY) + avg_speed = (total_gen_tokens / total_elapsed_time) if total_elapsed_time > 0 else 0.0 + return { + "total_gen_tokens": total_gen_tokens, + "total_elapsed_time": total_elapsed_time, + "avg_speed": avg_speed, + } + + # --------------------------------------------------------------------------- # # Module tree assembly # --------------------------------------------------------------------------- # @@ -193,15 +265,31 @@ def _mod(name: str, **attrs: Any) -> types.ModuleType: root = _mod("lmms_eval") api = _mod("lmms_eval.api") api_model = _mod("lmms_eval.api.model", lmms=lmms, CacheHook=_CacheHook) - api_instance = _mod("lmms_eval.api.instance", Instance=Instance) + api_instance = _mod( + "lmms_eval.api.instance", + Instance=Instance, + GenerationResult=GenerationResult, + TokenCounts=TokenCounts, + ) protocol = _mod("lmms_eval.protocol", ChatMessages=ChatMessages) utils = _mod("lmms_eval.utils", Collator=Collator) + models = _mod("lmms_eval.models") + model_utils = _mod("lmms_eval.models.model_utils") + gen_metrics = _mod( + "lmms_eval.models.model_utils.gen_metrics", + log_metrics=log_metrics, + reset_logged_metrics=reset_logged_metrics, + summarize_logged_metrics=summarize_logged_metrics, + ) root.api = api root.protocol = protocol root.utils = utils + root.models = models api.model = api_model api.instance = api_instance + models.model_utils = model_utils + model_utils.gen_metrics = gen_metrics return { "lmms_eval": root, @@ -210,4 +298,7 @@ def _mod(name: str, **attrs: Any) -> types.ModuleType: "lmms_eval.api.instance": api_instance, "lmms_eval.protocol": protocol, "lmms_eval.utils": utils, + "lmms_eval.models": models, + "lmms_eval.models.model_utils": model_utils, + "lmms_eval.models.model_utils.gen_metrics": gen_metrics, } diff --git a/examples/vlm-evaluation/tests/unit/test_adapter.py b/examples/vlm-evaluation/tests/unit/test_adapter.py index 4bd297d..3f42973 100644 --- a/examples/vlm-evaluation/tests/unit/test_adapter.py +++ b/examples/vlm-evaluation/tests/unit/test_adapter.py @@ -29,7 +29,7 @@ _resolve_gen_kwargs, _to_pil, ) -from lmms_eval.api.instance import Instance +from lmms_eval.api.instance import GenerationResult, Instance, TokenCounts from lmms_eval.protocol import ChatMessages from PIL import Image @@ -56,6 +56,23 @@ GENERATIVE_ARCHES = tuple(a for a in _ALL_VLM_ARCHS if VLMConfig.for_arch(a).is_generative) NON_GENERATIVE_ARCHES = tuple(a for a in _ALL_VLM_ARCHS if not VLMConfig.for_arch(a).is_generative) + +# ``_generate_batch`` now returns ``(texts, token_counts)`` and ``generate_until`` returns +# ``list[GenerationResult]`` (the native lmms-eval efficiency/throughput surface). These thin +# wrappers recover just the continuation strings so the behavioral tests below keep asserting +# on text; the new (text, count) contract is covered by ``TestGenerationTokenCounts``. +_gb = _generate_batch + + +def _gen_texts(*args, **kwargs) -> list[str]: + texts, _counts = _gb(*args, **kwargs) + return texts + + +def _gu_texts(model, requests) -> list[str]: + return [r.text for r in model.generate_until(requests)] + + # Per-arch BUILD sizing for a tiny CPU wrapper (sizing only, NOT generativity policy): # CA needs a cross-attention cadence that fits the tiny layer count. Arches without an # entry build from defaults; a future arch needing knobs fails the build loudly here. @@ -358,7 +375,7 @@ def doc_to_messages(doc): idx=0, metadata={"task": "t", "doc_id": "d0", "repeats": 1}, ) - out = vlm.generate_until([inst]) + out = _gu_texts(vlm, [inst]) assert len(out) == 1 and len(out[0].split()) == 3 # processed, not skipped @@ -427,18 +444,18 @@ def _prompt(self) -> list[torch.Tensor]: def test_greedy_is_deterministic(self, arch_wrapper): pv, pid = _pixels(), self._prompt() r = _resolve_gen_kwargs({"max_new_tokens": 6}, 128) - out1 = _generate_batch(arch_wrapper, _MockTokenizer(), pv, pid, r, 64) - out2 = _generate_batch(arch_wrapper, _MockTokenizer(), pv, pid, r, 64) + out1 = _gen_texts(arch_wrapper, _MockTokenizer(), pv, pid, r, 64) + out2 = _gen_texts(arch_wrapper, _MockTokenizer(), pv, pid, r, 64) assert out1 == out2 and len(out1) == 1 and isinstance(out1[0], str) def test_respects_max_new_tokens(self, arch_wrapper): r = _resolve_gen_kwargs({"max_new_tokens": 6}, 128) - out = _generate_batch(arch_wrapper, _MockTokenizer(), _pixels(), self._prompt(), r, 64) + out = _gen_texts(arch_wrapper, _MockTokenizer(), _pixels(), self._prompt(), r, 64) assert len(out[0].split()) == 6 def test_until_trims_continuation(self, arch_wrapper): pv, pid = _pixels(), self._prompt() - one = _generate_batch( + one = _gen_texts( arch_wrapper, _MockTokenizer(), pv, @@ -448,7 +465,7 @@ def test_until_trims_continuation(self, arch_wrapper): )[0] # decode = space-joined ids, so the first space follows the first token: # until=[" "] trims to exactly the first generated token. - trimmed = _generate_batch( + trimmed = _gen_texts( arch_wrapper, _MockTokenizer(), pv, @@ -460,7 +477,7 @@ def test_until_trims_continuation(self, arch_wrapper): def test_eos_stops_generation(self, arch_wrapper): pv, pid = _pixels(), self._prompt() - first = _generate_batch( + first = _gen_texts( arch_wrapper, _MockTokenizer(), pv, @@ -468,7 +485,7 @@ def test_eos_stops_generation(self, arch_wrapper): _resolve_gen_kwargs({"max_new_tokens": 1}, 128), 64, )[0] - out = _generate_batch( + out = _gen_texts( arch_wrapper, _MockTokenizer(eos_token_id=int(first)), pv, @@ -483,7 +500,7 @@ def test_no_room_raises_context_budget_error(self, arch_wrapper): # per-task skip (see TestGenerateUntilFaultTolerance). r = _resolve_gen_kwargs({"max_new_tokens": 100}, 128) with pytest.raises(_ContextBudgetError, match="max_new_tokens"): - _generate_batch(arch_wrapper, _MockTokenizer(), _pixels(), self._prompt(), r, 64) + _gen_texts(arch_wrapper, _MockTokenizer(), _pixels(), self._prompt(), r, 64) def test_overlong_prompt_is_left_truncated(self, arch_wrapper, monkeypatch): """A prompt that exceeds the budget (but leaves room) is left-truncated with a warning.""" @@ -496,7 +513,7 @@ def test_overlong_prompt_is_left_truncated(self, arch_wrapper, monkeypatch): budget = 64 - arch_wrapper.num_image_tokens - max_new long_prompt = [torch.arange(1, budget + 7, dtype=torch.long)] r = _resolve_gen_kwargs({"max_new_tokens": max_new}, 128) - out = _generate_batch(arch_wrapper, _MockTokenizer(), _pixels(), long_prompt, r, 64) + out = _gen_texts(arch_wrapper, _MockTokenizer(), _pixels(), long_prompt, r, 64) assert len(out) == 1 and len(out[0].split()) == 2 assert any("left-truncating" in m for m in rec.warnings) @@ -504,7 +521,7 @@ def test_text_only_decode_no_pixels(self, arch_wrapper): """pixel_values=None runs the pure-text forward for every generative arch (JD/CA/MoT) end to end and returns a continuation of max_new_tokens.""" r = _resolve_gen_kwargs({"max_new_tokens": 6}, 128) - out = _generate_batch(arch_wrapper, _MockTokenizer(), None, self._prompt(), r, 64) + out = _gen_texts(arch_wrapper, _MockTokenizer(), None, self._prompt(), r, 64) assert len(out) == 1 and isinstance(out[0], str) assert len(out[0].split()) == 6 @@ -529,24 +546,24 @@ def test_batch_equals_sequential(self, arch_wrapper): pv = _pixels(len(prompts)) # (3, 3, 16, 16) — one image per request r = _resolve_gen_kwargs({"max_new_tokens": 5}, 128) sequential = [ - _generate_batch(arch_wrapper, _MockTokenizer(), pv[i : i + 1], [prompts[i]], r, 64)[0] + _gen_texts(arch_wrapper, _MockTokenizer(), pv[i : i + 1], [prompts[i]], r, 64)[0] for i in range(len(prompts)) ] - batched = _generate_batch(arch_wrapper, _MockTokenizer(), pv, prompts, r, 64) + batched = _gen_texts(arch_wrapper, _MockTokenizer(), pv, prompts, r, 64) assert batched == sequential def test_per_row_max_new_tokens(self, arch_wrapper): prompts = self._prompts() pv = _pixels(len(prompts)) r = _resolve_gen_kwargs({"max_new_tokens": 4}, 128) - outs = _generate_batch(arch_wrapper, _MockTokenizer(), pv, prompts, r, 64) + outs = _gen_texts(arch_wrapper, _MockTokenizer(), pv, prompts, r, 64) assert len(outs) == 3 and all(len(o.split()) == 4 for o in outs) def test_per_row_eos_independent(self, arch_wrapper): """EOS on one row stops only that row; the batch still returns all rows.""" prompts = self._prompts() pv = _pixels(len(prompts)) - first0 = _generate_batch( + first0 = _gen_texts( arch_wrapper, _MockTokenizer(), pv[:1], @@ -554,7 +571,7 @@ def test_per_row_eos_independent(self, arch_wrapper): _resolve_gen_kwargs({"max_new_tokens": 1}, 128), 64, )[0] - outs = _generate_batch( + outs = _gen_texts( arch_wrapper, _MockTokenizer(eos_token_id=int(first0)), pv, @@ -595,7 +612,7 @@ def test_generate_batch_threads_frame_mask_for_video(): frame_mask = torch.tensor([[True, True], [True, False]]) # row 1: frame 2 padded prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7], dtype=torch.long)] r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128) - _generate_batch(model, _MockTokenizer(), pixel_values, prompt_ids, r, 64, frame_mask=frame_mask) + _gen_texts(model, _MockTokenizer(), pixel_values, prompt_ids, r, 64, frame_mask=frame_mask) assert model.seen_frame_masks # decode actually ran assert all(fm is frame_mask for fm in model.seen_frame_masks) @@ -606,7 +623,7 @@ def test_generate_batch_no_frame_mask_for_image(): pixel_values = torch.randn(1, 3, 16, 16) prompt_ids = [torch.tensor([5, 9, 12], dtype=torch.long)] r = _resolve_gen_kwargs({"max_new_tokens": 2}, 128) - _generate_batch(model, _MockTokenizer(), pixel_values, prompt_ids, r, 64) + _gen_texts(model, _MockTokenizer(), pixel_values, prompt_ids, r, 64) assert model.seen_frame_masks and all(fm is None for fm in model.seen_frame_masks) @@ -616,7 +633,7 @@ def test_generate_batch_text_only_runs_with_none(): model = _CaptureModel(num_image_tokens=8) prompt_ids = [torch.tensor([5, 9, 12], dtype=torch.long)] r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128) - out = _generate_batch(model, _MockTokenizer(), None, prompt_ids, r, 64) + out = _gen_texts(model, _MockTokenizer(), None, prompt_ids, r, 64) assert len(out) == 1 assert model.seen_frame_masks and all(fm is None for fm in model.seen_frame_masks) @@ -629,8 +646,8 @@ def test_generate_batch_text_only_budget_excludes_image_tokens(): r = _resolve_gen_kwargs({"max_new_tokens": 50}, 128) # Visual budget: 64 - 16 - 50 <= 0 -> raises. Text-only: 64 - 0 - 50 = 14 -> fits. with pytest.raises(_ContextBudgetError): - _generate_batch(model, _MockTokenizer(), torch.randn(1, 3, 16, 16), prompt_ids, r, 64) - out = _generate_batch(model, _MockTokenizer(), None, prompt_ids, r, 64) + _gen_texts(model, _MockTokenizer(), torch.randn(1, 3, 16, 16), prompt_ids, r, 64) + out = _gen_texts(model, _MockTokenizer(), None, prompt_ids, r, 64) assert len(out) == 1 @@ -915,8 +932,8 @@ def doc_to_messages(doc): ] # Each request decoded alone (greedy → deterministic), in original order. - singles = [vlm.generate_until([inst])[0] for inst in instances] - batched = vlm.generate_until(instances) + singles = [_gu_texts(vlm, [inst])[0] for inst in instances] + batched = _gu_texts(vlm, instances) assert batched == singles # batching + reorder + get_original preserve per-request results assert all(isinstance(o, str) for o in batched) @@ -956,8 +973,8 @@ def doc_to_messages(doc): # d0 (text-only) and d1 (image) land in one chunk (shared gen_kwargs); each is # decoded in its own sub-batch, so the mixed batch matches per-request singles. - singles = [vlm.generate_until([inst])[0] for inst in instances] - batched = vlm.generate_until(instances) + singles = [_gu_texts(vlm, [inst])[0] for inst in instances] + batched = _gu_texts(vlm, instances) assert batched == singles assert all(isinstance(o, str) and len(o.split()) == 3 for o in batched) @@ -1004,8 +1021,8 @@ def doc_to_messages(doc): ] # Decoding each request alone must match the batched (stacked, 5-D) result; if the # adapter folded frames with cat, the forward would trip the frames-per-clip check. - singles = [vlm.generate_until([inst])[0] for inst in instances] - batched = vlm.generate_until(instances) + singles = [_gu_texts(vlm, [inst])[0] for inst in instances] + batched = _gu_texts(vlm, instances) assert batched == singles assert all(isinstance(o, str) and len(o.split()) == 3 for o in batched) @@ -1070,7 +1087,7 @@ def doc_to_messages(doc): ) for i, (doc_id, ctx) in enumerate([("bad", "c"), ("good", "cc")]) ] - outs = vlm.generate_until(instances) + outs = _gu_texts(vlm, instances) assert outs[0] == "" # bad request isolated (order restored by get_original) assert len(outs[1].split()) == 3 # good request completes normally assert any("Skipping request" in m and "doc_id=bad" in m for m in rec.warnings) @@ -1092,7 +1109,7 @@ def doc_to_messages(doc): idx=0, metadata={"task": "t", "doc_id": "d0", "repeats": 1}, ) - assert vlm.generate_until([inst]) == [""] + assert _gu_texts(vlm, [inst]) == [""] assert any("empty prompt" in m for m in rec.warnings) def test_all_requests_bad_returns_empties(self, monkeypatch, tiny_vlm_configs): @@ -1115,7 +1132,7 @@ def doc_to_messages(doc): ) for i, (doc_id, ctx) in enumerate([("d0", "c"), ("d1", "cc")]) ] - assert vlm.generate_until(instances) == ["", ""] + assert _gu_texts(vlm, instances) == ["", ""] def test_unexpected_error_isolated(self, monkeypatch, tiny_vlm_configs): # The per-request handler catches Exception (broadened from the original @@ -1142,7 +1159,7 @@ def doc_to_messages(doc): idx=0, metadata={"task": "t", "doc_id": "d0", "repeats": 1}, ) - assert vlm.generate_until([inst]) == [""] + assert _gu_texts(vlm, [inst]) == [""] assert any("Skipping request" in m and "doc_id=d0" in m for m in rec.warnings) def test_fatal_baseexception_propagates(self, monkeypatch, tiny_vlm_configs): @@ -1168,7 +1185,7 @@ def doc_to_messages(doc): metadata={"task": "t", "doc_id": "d0", "repeats": 1}, ) with pytest.raises(KeyboardInterrupt): - vlm.generate_until([inst]) + _gu_texts(vlm, [inst]) def test_missing_image_path_skipped(self, monkeypatch, tiny_vlm_configs): # A bad image path raises FileNotFoundError (subclass of OSError) inside _to_pil; @@ -1188,7 +1205,7 @@ def doc_to_messages(doc): idx=0, metadata={"task": "t", "doc_id": "d0", "repeats": 1}, ) - assert vlm.generate_until([inst]) == [""] + assert _gu_texts(vlm, [inst]) == [""] assert any("Skipping request" in m and "doc_id=d0" in m for m in rec.warnings) def test_over_budget_chunk_skipped_not_aborted(self, monkeypatch, tiny_vlm_configs): @@ -1213,5 +1230,62 @@ def doc_to_messages(doc): ) for i, (doc_id, ctx) in enumerate([("d0", "c"), ("d1", "cc")]) ] - assert vlm.generate_until(instances) == ["", ""] + assert _gu_texts(vlm, instances) == ["", ""] assert any("Skipping" in m and "max_new_tokens" in m for m in rec.warnings) + + +# --------------------------------------------------------------------------- +# Native lmms-eval throughput/efficiency wiring: the (texts, counts) tuple from +# _generate_batch and the GenerationResult(text, token_counts) from generate_until, +# plus the generation-throughput metrics fed to gen_metrics.log_metrics. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("arch", GENERATIVE_ARCHES) +class TestGenerationTokenCounts: + def test_generate_batch_returns_per_row_token_counts(self, arch_wrapper): + r = _resolve_gen_kwargs({"max_new_tokens": 6}, 128) + texts, counts = _generate_batch( + arch_wrapper, + _MockTokenizer(), + _pixels(), + [torch.tensor([5, 9, 12, 3], dtype=torch.long)], + r, + 64, + ) + # _MockTokenizer decodes ids space-joined, so word count == generated token count. + assert len(counts) == len(texts) == 1 + assert counts[0] == len(texts[0].split()) == 6 + + def test_generate_until_reports_token_counts_and_throughput( + self, monkeypatch, tiny_vlm_configs, arch + ): + from lmms_eval.models.model_utils.gen_metrics import ( + reset_logged_metrics, + summarize_logged_metrics, + ) + + _patch_loaders( + monkeypatch, _vlm_job_config(tiny_vlm_configs, arch=arch), _vlm_wrapper(arch) + ) + vlm = KempnerForgeVLM(config="x", checkpoint="y", device="cpu", dtype="float32") + + def doc_to_messages(doc): + del doc + return [{"role": "user", "content": [_text("hi?"), _image(_img())]}] + + vlm.task_dict = {"t": {"test": {"d0": {}}}} + inst = Instance( + request_type="generate_until", + arguments=("ctx", doc_to_messages, {"max_new_tokens": 3}, "d0", "t", "test"), + idx=0, + metadata={"task": "t", "doc_id": "d0", "repeats": 1}, + ) + reset_logged_metrics() + out = vlm.generate_until([inst]) + # Efficiency surface: a typed result carrying the per-sample output-token count. + assert len(out) == 1 and isinstance(out[0], GenerationResult) + assert isinstance(out[0].token_counts, TokenCounts) + assert out[0].token_counts.output_tokens == len(out[0].text.split()) == 3 + # Throughput surface: log_metrics was fed, so the evaluator's summary is non-empty. + assert summarize_logged_metrics()["total_gen_tokens"] == 3