diff --git a/examples/vlm-evaluation/adapter.py b/examples/vlm-evaluation/adapter.py index 068b210..d05fbcd 100644 --- a/examples/vlm-evaluation/adapter.py +++ b/examples/vlm-evaluation/adapter.py @@ -15,8 +15,8 @@ v1 scope and deliberate choices (see README.md in this directory): - **Generation: no transformer KV cache, batched, data-parallel.** The decode loop - re-runs the transformer (including the vision encoder + adapter) over the - growing sequence each step. There is no transformer KV cache + re-runs the transformer over the growing sequence each step; the vision encoder + + adapter run **once** per request (``encode_visual``). There is no transformer KV cache (``Transformer.forward`` forbids combining ``kv_caches`` with any image-conditioning route), and KempnerForge has no image-conditioned KV-cache decode path. Requests are decoded in batches @@ -453,16 +453,16 @@ def _generate_batch( 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 + cache; the vision tower is encoded **once** per request and reused each step, + while ``model(...)`` re-runs the transformer over the growing **right-padded** + batch. Right-padding matches the training layout: the image prefix stays at positions ``0..n-1`` and text is contiguous from ``n`` for every row (so image/text RoPE distances are consistent across rows), and the trailing pads are causally masked, so a batched forward gives each row the same real-position logits as decoding it alone. Each row's next token is read at its own last real position; EOS / ``max_new_tokens`` / first - ``until`` match are tracked per row. ``B == 1`` reproduces the single-request - path exactly. + ``until`` match are tracked per row, and a row that finishes is dropped from + every later forward. ``B == 1`` reproduces the single-request path exactly. """ until: list[str] = resolved["until"] max_new_tokens: int = resolved["max_new_tokens"] @@ -498,43 +498,76 @@ def _generate_batch( prompts.append(ids) generated: list[list[int]] = [[] for _ in range(batch_size)] - done = [False] * batch_size - row_index = torch.arange(batch_size, device=device) + # No decode steps to run: return before touching the vision tower, whose output + # would be discarded by the zero-iteration loop below. + if max_new_tokens == 0: + return [""] * batch_size, [0] * batch_size + + # Encode the clip once and reuse it every step; text-only batches skip this. + visual_embeds = None if pixel_values is None else model.encode_visual(pixel_values) + + # Token buffer, allocated once: prompt + room for the whole generation, pad-filled. + # A step writes its sampled token in place, so no per-row tensor is rebuilt or + # copied to the device each iteration. ``lengths`` stays a Python list: it indexes + # the buffer and bounds the slice, and reading it off a device tensor would force + # a sync every step. + max_prompt_len = max(p.shape[0] for p in prompts) + token_buf = torch.full( + (batch_size, max_prompt_len + max_new_tokens), pad_id, dtype=torch.long, device=device + ) + lengths = [p.shape[0] for p in prompts] + for i, p in enumerate(prompts): + token_buf[i, : lengths[i]] = p + + # Rows still generating. Finished rows are dropped from the forward batch: their + # logits were discarded anyway, and with right-padding + causal attention a row's + # logits at its own real positions depend on neither the other rows nor the batch + # width, so survivors decode bit-identically to the uncompacted loop. + active = list(range(batch_size)) for _ in range(max_new_tokens): - # Rebuild the right-padded batch from prompt + tokens generated so far. - seqs = [ - torch.cat([prompts[i], torch.tensor(generated[i], dtype=torch.long, device=device)]) - for i in range(batch_size) - ] - real_len = torch.tensor([s.shape[0] for s in seqs], device=device) - cur_max = int(real_len.max().item()) - input_ids = torch.full((batch_size, cur_max), pad_id, dtype=torch.long, device=device) - for i, s in enumerate(seqs): - input_ids[i, : s.shape[0]] = s - - logits, _ = model(pixel_values, input_ids, frame_mask=frame_mask) + cur_max = max(lengths[i] for i in active) + if len(active) == batch_size: + # Full batch: slice the buffer in place and pass the visual tensors + # through untouched (no index_select, no copy). + input_ids = token_buf[:, :cur_max] + step_embeds, step_frame_mask = visual_embeds, frame_mask + else: + # Compacted: the visual tensors are row-aligned with input_ids, and + # VLMWrapper.forward rejects a cache whose batch differs from input_ids. + active_index = torch.tensor(active, device=device) + input_ids = token_buf[active_index, :cur_max] + step_embeds = None if visual_embeds is None else visual_embeds[active_index] + step_frame_mask = None if frame_mask is None else frame_mask[active_index] + real_len = torch.tensor([lengths[i] for i in active], device=device) + + # pixel_values is never passed: the clip is already encoded into + # visual_embeds, and forward rejects both together (they could disagree). + # Text-only batches have no cache and pass None for both. + logits, _ = model(None, input_ids, frame_mask=step_frame_mask, visual_embeds=step_embeds) # Each row's next-token logits sit at its own last real position (the # output is already trimmed to text positions for JD/MoT; CA has no # image prefix), not at [-1] (a pad for shorter rows). - next_logits = logits[row_index, real_len - 1] + next_logits = logits[torch.arange(len(active), device=device), real_len - 1] next_tokens = sample(next_logits, temperature, top_k, top_p) - for i in range(batch_size): - if done[i]: - continue - token_id = int(next_tokens[i].item()) + still_active: list[int] = [] + for j, i in enumerate(active): + token_id = int(next_tokens[j].item()) if eos_id is not None and token_id == eos_id: - done[i] = True continue generated[i].append(token_id) + token_buf[i, lengths[i]] = token_id + lengths[i] += 1 if len(generated[i]) >= max_new_tokens: - done[i] = True - elif until: + continue + if until: text = tokenizer.decode(generated[i], skip_special_tokens=True) if _first_stop(text, until) is not None: - done[i] = True - if all(done): + continue + still_active.append(i) + active = still_active + if not active: break outputs: list[str] = [] diff --git a/examples/vlm-evaluation/tests/unit/test_adapter.py b/examples/vlm-evaluation/tests/unit/test_adapter.py index d5ce4f3..36befb2 100644 --- a/examples/vlm-evaluation/tests/unit/test_adapter.py +++ b/examples/vlm-evaluation/tests/unit/test_adapter.py @@ -9,6 +9,7 @@ from __future__ import annotations +import inspect import json import sys import types @@ -590,18 +591,50 @@ def test_per_row_eos_independent(self, arch_wrapper): class _CaptureModel: - """Minimal ``VLMWrapper`` stand-in: records the ``frame_mask`` each forward - receives and returns deterministic zero logits (greedy -> token 0), so the - decode-loop plumbing can be asserted without a real transformer.""" + """Minimal ``VLMWrapper`` stand-in: records what each forward receives + (``pixel_values``, ``visual_embeds``, ``frame_mask``, batch width), counts + ``encode_visual`` / forward calls, and returns deterministic zero logits + (greedy -> token 0), so the decode-loop plumbing can be asserted without a + real transformer. + + ``__call__`` mirrors ``VLMWrapper.forward``'s signature exactly, including + ``labels`` and parameter order; ``test_capture_model_matches_forward_signature`` + pins that so a positional call site cannot bind arguments differently here + than in production. + """ def __init__(self, num_image_tokens: int, vocab_size: int = 256) -> None: self.num_image_tokens = num_image_tokens self._vocab = vocab_size self.seen_frame_masks: list[torch.Tensor | None] = [] + self.seen_pixel_values: list[torch.Tensor | None] = [] + self.seen_visual_embeds: list[torch.Tensor | None] = [] + self.seen_batch_sizes: list[int] = [] + self.encode_calls = 0 + self.forward_calls = 0 + self.returned_embeds: torch.Tensor | None = None + + def encode_visual(self, pixel_values): + self.encode_calls += 1 + # Identifiable object whose rows are distinguishable (row i is filled with i), + # so a forward can be asserted to have received *this* cache, and a compacted + # forward to have received the surviving rows of it. + batch = pixel_values.shape[0] + self.returned_embeds = ( + torch.arange(batch, dtype=torch.float32) + .view(batch, 1, 1) + .expand(-1, self.num_image_tokens, 1) + .clone() + ) + return self.returned_embeds - def __call__(self, pixel_values, input_ids, frame_mask=None): - del pixel_values + def __call__(self, pixel_values, input_ids, labels=None, frame_mask=None, visual_embeds=None): + del labels + self.forward_calls += 1 self.seen_frame_masks.append(frame_mask) + self.seen_pixel_values.append(pixel_values) + self.seen_visual_embeds.append(visual_embeds) + self.seen_batch_sizes.append(input_ids.shape[0]) b, t = input_ids.shape return torch.zeros(b, t, self._vocab), None @@ -653,6 +686,153 @@ def test_generate_batch_text_only_budget_excludes_image_tokens(): assert len(out) == 1 +def test_encode_visual_called_once_per_generate_batch(): + """Vision encoded once per request, not once per decode step -- and the cache that + encode_visual returned is what every forward actually conditions on. + + Counting encode calls alone is not enough: dropping ``visual_embeds=`` from the + model call would leave the count at 1 while the real wrapper re-projected the clip + on every step. The identity assertion is what makes that revert fail. + """ + model = _CaptureModel(num_image_tokens=8) + pixel_values = torch.randn(2, 3, 16, 16) + prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7, 3], dtype=torch.long)] + max_new = 5 + r = _resolve_gen_kwargs({"max_new_tokens": max_new}, 128) + _gen_texts(model, _MockTokenizer(), pixel_values, prompt_ids, r, 64) + assert model.encode_calls == 1 + assert model.forward_calls == max_new + # No row finishes early here, so every step takes the uncompacted fast path and + # forwards the cache object itself. + assert model.returned_embeds is not None + assert all(ve is model.returned_embeds for ve in model.seen_visual_embeds) + + +def test_generate_batch_never_passes_pixels_with_cache(): + """VLMWrapper.forward rejects pixel_values and visual_embeds together, so the decode + loop must pass only the cache once the clip is encoded.""" + model = _CaptureModel(num_image_tokens=8) + prompt_ids = [torch.tensor([5, 9], dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128) + _gen_texts(model, _MockTokenizer(), torch.randn(1, 3, 16, 16), prompt_ids, r, 64) + assert model.forward_calls == 3 + assert all(pv is None for pv in model.seen_pixel_values) + assert all(ve is not None for ve in model.seen_visual_embeds) + + +def test_capture_model_matches_forward_signature(): + """_CaptureModel must accept arguments exactly as VLMWrapper.forward does. + + A drifting mock (missing ``labels``, or ``visual_embeds`` ahead of ``frame_mask``) + would keep every decode-loop test green while a positional call site in production + bound the cache to ``labels`` and dropped image conditioning. + """ + real = list(inspect.signature(VLMWrapper.forward).parameters) + mock = list(inspect.signature(_CaptureModel.__call__).parameters) + assert real[0] == "self" and mock[0] == "self" + assert real[1:] == mock[1:] + + +def test_encode_visual_not_called_for_text_only(): + """A text-only batch (pixel_values=None) never touches the vision tower.""" + 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) + _gen_texts(model, _MockTokenizer(), None, prompt_ids, r, 64) + assert model.encode_calls == 0 + assert model.forward_calls == 3 + assert all(ve is None for ve in model.seen_visual_embeds) + + +def test_zero_max_new_tokens_skips_vision_and_decode(): + """max_new_tokens=0 is a supported task value; with no decode steps to run, the + vision tower must not be run only for its output to be discarded.""" + model = _CaptureModel(num_image_tokens=8) + prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7], dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": 0}, 128) + texts, counts = _gb(model, _MockTokenizer(), torch.randn(2, 3, 16, 16), prompt_ids, r, 64) + assert texts == ["", ""] and counts == [0, 0] + assert model.encode_calls == 0 and model.forward_calls == 0 + + +class _RowTokenModel(_CaptureModel): + """_CaptureModel that makes greedy decoding row-dependent: row j's argmax is its + own first prompt token, which is stable across steps and survives compaction (the + row is identified by its content, not its position). Lets one row hit EOS while + another keeps generating.""" + + def __call__(self, pixel_values, input_ids, labels=None, frame_mask=None, visual_embeds=None): + logits, _ = super().__call__(pixel_values, input_ids, labels, frame_mask, visual_embeds) + logits = logits.clone() + for j in range(input_ids.shape[0]): + logits[j, :, int(input_ids[j, 0])] = 1.0 + return logits, None + + +def test_finished_rows_dropped_from_forward_batch(): + """A row that hits EOS stops being forwarded: its logits were discarded anyway, so + keeping it in the batch is pure waste (and skews the reported tok/s).""" + model = _RowTokenModel(num_image_tokens=8) + # Row 0 generates token 5 forever; row 1 generates token 7 == eos and finishes + # on the first step. + prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7, 3], dtype=torch.long)] + max_new = 4 + r = _resolve_gen_kwargs({"max_new_tokens": max_new}, 128) + texts, counts = _gb( + model, _MockTokenizer(eos_token_id=7), torch.randn(2, 3, 16, 16), prompt_ids, r, 64 + ) + assert texts == ["5 5 5 5", ""] and counts == [4, 0] + # Step 1 sees both rows; every later step sees only the survivor. + assert model.seen_batch_sizes == [2, 1, 1, 1] + # The compacted cache is row 0 of the original (encode_visual fills row i with i), + # so the survivor stays conditioned on its own image. + assert model.seen_visual_embeds[-1].shape[0] == 1 + assert torch.equal(model.seen_visual_embeds[-1][0], model.returned_embeds[0]) + + +def test_compaction_preserves_survivor_output(): + """Dropping finished rows must not change what the remaining rows generate.""" + prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7, 3], dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": 4}, 128) + batched = _gen_texts( + _RowTokenModel(num_image_tokens=8), + _MockTokenizer(eos_token_id=7), + torch.randn(2, 3, 16, 16), + prompt_ids, + r, + 64, + ) + alone = _gen_texts( + _RowTokenModel(num_image_tokens=8), + _MockTokenizer(eos_token_id=7), + torch.randn(1, 3, 16, 16), + [prompt_ids[0]], + r, + 64, + ) + assert batched[0] == alone[0] + + +def test_frame_mask_compacted_with_finished_rows(): + """frame_mask is row-aligned with the forward batch, so it must be compacted too -- + otherwise the survivor would be masked with another row's padded-frame pattern.""" + model = _RowTokenModel(num_image_tokens=8) + frame_mask = torch.tensor([[True, True], [True, False]]) + prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7, 3], dtype=torch.long)] + r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128) + _gen_texts( + model, + _MockTokenizer(eos_token_id=7), + torch.randn(2, 2, 3, 16, 16), + prompt_ids, + r, + 64, + frame_mask=frame_mask, + ) + assert model.seen_frame_masks[0] is frame_mask # full batch: passed through as-is + assert torch.equal(model.seen_frame_masks[-1], frame_mask[:1]) + + # --------------------------------------------------------------------------- # Guards: arch + not-implemented methods # ---------------------------------------------------------------------------