@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 122 lines 3.75 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from typing import Any
5
6 import pytest
7 import torch
8 from transformers import AutoModelForImageTextToText
9
10 from vllm.platforms import current_platform
11 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
12
13 from ....conftest import HfRunner, ImageTestAssets, VllmRunner
14 from .vlm_utils import model_utils
15
16 MODEL = "google/gemma-3-4b-it"
17 PROMPT = (
18 "<bos><start_of_turn>user\n"
19 "<start_of_image>What is the content in the center of the image?"
20 "<end_of_turn>\n<start_of_turn>model\n"
21 )
22
23
24 def _install_prefill_hidden_capture(model):
25 model = getattr(model, "module", model)
26 model._prefill_hidden = None
27
28 language_model = model.language_model.model
29 original_forward = language_model.forward
30
31 def forward(*args, **kwargs):
32 hidden_states = original_forward(*args, **kwargs)
33 if model._prefill_hidden is None and torch.is_tensor(hidden_states):
34 # Capturing hidden states for comparison is a deliberate D2H.
35 with gpu_sync_allowed():
36 model._prefill_hidden = hidden_states.detach().float().cpu()
37 return hidden_states
38
39 language_model.forward = forward
40
41
42 def _get_prefill_hidden(model):
43 model = getattr(model, "module", model)
44 hidden = getattr(model, "_prefill_hidden", None)
45 assert hidden is not None
46 return hidden
47
48
49 def _get_hf_prefill_hidden(hf_model: HfRunner, image: Any):
50 inputs = hf_model.get_inputs([PROMPT], images=[image])[0]
51 with torch.no_grad():
52 outputs = hf_model.model.model(
53 **hf_model.wrap_device(inputs),
54 use_cache=False,
55 )
56 return outputs.last_hidden_state[0].detach().float().cpu()
57
58
59 def _get_vllm_prefill_hidden(
60 vllm_runner: type[VllmRunner],
61 image: Any,
62 vllm_runner_kwargs: dict[str, Any],
63 ):
64 with vllm_runner(
65 MODEL,
66 max_model_len=4096,
67 max_num_seqs=2,
68 enforce_eager=True,
69 limit_mm_per_prompt={"image": 1},
70 **vllm_runner_kwargs,
71 ) as vllm_model:
72 vllm_model.apply_model(_install_prefill_hidden_capture)
73 vllm_model.generate_greedy([PROMPT], max_tokens=1, images=[image])
74 return vllm_model.apply_model(_get_prefill_hidden)[0]
75
76
77 @pytest.mark.core_model
78 @pytest.mark.skipif(
79 current_platform.is_rocm(), reason="ROCm attention has accuracy issue for this test"
80 )
81 def test_mm_prefix_lm_e2e(
82 hf_runner: type[HfRunner],
83 vllm_runner: type[VllmRunner],
84 image_assets: ImageTestAssets,
85 monkeypatch: pytest.MonkeyPatch,
86 ):
87 """Regression: Gemma3 native prefill must apply image prefix-LM mask."""
88 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
89 image = image_assets[0].pil_image
90
91 vllm_runner_kwargs: dict[str, Any] = {
92 "mm_processor_cache_gb": 0,
93 "mm_processor_kwargs": {"do_pan_and_scan": True},
94 }
95 vllm_hidden = _get_vllm_prefill_hidden(vllm_runner, image, vllm_runner_kwargs)
96
97 hf_model = hf_runner(
98 MODEL,
99 auto_cls=AutoModelForImageTextToText,
100 )
101 hf_model = model_utils.gemma3_patch_hf_runner(hf_model)
102
103 with hf_model:
104 hf_hidden = _get_hf_prefill_hidden(hf_model, image)
105
106 assert vllm_hidden.shape == hf_hidden.shape
107
108 full_cos = torch.nn.functional.cosine_similarity(
109 vllm_hidden.flatten(), hf_hidden.flatten(), dim=0
110 )
111 image_cos = torch.nn.functional.cosine_similarity(
112 vllm_hidden[1:769].flatten(), hf_hidden[1:769].flatten(), dim=0
113 )
114
115 assert full_cos > 0.9, (
116 "Gemma3 mm-prefix-LM full prefill hidden states should be close to HF; "
117 f"got {full_cos=}"
118 )
119 assert image_cos > 0.9, (
120 "Gemma3 mm-prefix-LM image prefill hidden states should be close to HF; "
121 f"got {image_cos=}"
122 )