| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import json |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | import torch.nn.functional as F |
| 9 | from transformers import AutoModel |
| 10 | |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.utils.mem_constants import MiB_bytes |
| 13 | |
| 14 | from ....conftest import HfRunner |
| 15 | from ....utils import VLLM_PATH |
| 16 | from ...registry import HF_EXAMPLE_MODELS |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from _typeshed import StrPath |
| 20 | |
| 21 | |
| 22 | FIXTURES_PATH = VLLM_PATH / "tests/models/fixtures" |
| 23 | assert FIXTURES_PATH.exists() |
| 24 | FIXTURE_REWARD_RESULT = { |
| 25 | "Qwen/Qwen2.5-Math-PRM-7B": FIXTURES_PATH / "qwen2_5_math_prm_reward_step.json", |
| 26 | } |
| 27 | |
| 28 | |
| 29 | @pytest.fixture |
| 30 | def math_step_prompts(): |
| 31 | # ruff: noqa: E501 |
| 32 | data = { |
| 33 | "system": "Please reason step by step, and put your final answer within \\boxed{}. ", |
| 34 | "query": "Sue lives in a fun neighborhood. One weekend, the neighbors decided to play a prank on Sue. On Friday morning, the neighbors placed 18 pink plastic flamingos out on Sue's front yard. On Saturday morning, the neighbors took back one third of the flamingos, painted them white, and put these newly painted white flamingos back out on Sue's front yard. Then, on Sunday morning, they added another 18 pink plastic flamingos to the collection. At noon on Sunday, how many more pink plastic flamingos were out than white plastic flamingos?", |
| 35 | "response": [ |
| 36 | "To find out how many more pink plastic flamingos were out than white plastic flamingos at noon on Sunday, we can break down the problem into steps. First, on Friday, the neighbors start with 18 pink plastic flamingos.", |
| 37 | "On Saturday, they take back one third of the flamingos. Since there were 18 flamingos, (1/3 \\times 18 = 6) flamingos are taken back. So, they have (18 - 6 = 12) flamingos left in their possession. Then, they paint these 6 flamingos white and put them back out on Sue's front yard. Now, Sue has the original 12 pink flamingos plus the 6 new white ones. Thus, by the end of Saturday, Sue has (12 + 6 = 18) pink flamingos and 6 white flamingos.", |
| 38 | "On Sunday, the neighbors add another 18 pink plastic flamingos to Sue's front yard. By the end of Sunday morning, Sue has (18 + 18 = 36) pink flamingos and still 6 white flamingos.", |
| 39 | "To find the difference, subtract the number of white flamingos from the number of pink flamingos: (36 - 6 = 30). Therefore, at noon on Sunday, there were 30 more pink plastic flamingos out than white plastic flamingos. The answer is (\\boxed{30}).", |
| 40 | ], |
| 41 | } |
| 42 | answer = "<extra_0>".join(data["response"]) + "<extra_0>" |
| 43 | prompt = f"<im_start>system\n{data['system']}<im_end>\n<im_start>user\n{data['query']}<im_end>\n<im_start>assistant\n{answer}<im_end><|endoftext|>" |
| 44 | return [prompt] |
| 45 | |
| 46 | |
| 47 | def step_reward_patch_hf_model(hf_model: HfRunner): |
| 48 | # Patch the hf_runner to use the step reward function |
| 49 | def make_step_rewards( |
| 50 | logits: torch.Tensor, token_masks: torch.Tensor |
| 51 | ) -> list[list[float]]: |
| 52 | probabilities = F.softmax(logits, dim=-1) |
| 53 | probabilities = probabilities * token_masks.unsqueeze(-1) |
| 54 | |
| 55 | all_scores_res: list[list[float]] = [] |
| 56 | for i in range(probabilities.size(0)): |
| 57 | sample = probabilities[i] # seq_len, num_labels |
| 58 | positive_probs = sample[sample != 0].view(-1, 2) |
| 59 | non_zero_elements_list = positive_probs.cpu().tolist() |
| 60 | all_scores_res.append(non_zero_elements_list) |
| 61 | return all_scores_res |
| 62 | |
| 63 | def reward(prompts: list[str]) -> list[list[float]]: |
| 64 | input_ids = hf_model.tokenizer(prompts, return_tensors="pt").input_ids |
| 65 | input_ids = hf_model.wrap_device(input_ids) |
| 66 | outputs = hf_model.model(input_ids=input_ids) |
| 67 | |
| 68 | step_sep_id = hf_model.tokenizer.encode("<extra_0>")[0] |
| 69 | token_masks = input_ids == step_sep_id |
| 70 | return make_step_rewards(outputs[0], token_masks) |
| 71 | |
| 72 | hf_model.reward = reward # type: ignore[attr-defined] |
| 73 | |
| 74 | return hf_model |
| 75 | |
| 76 | |
| 77 | def dump_reward_outputs(outputs: list[list[float]], filename: "StrPath"): |
| 78 | with open(filename, "w", encoding="utf-8") as f: |
| 79 | json.dump(outputs, f) |
| 80 | |
| 81 | |
| 82 | def load_reward_outputs(filename: "StrPath") -> list[list[float]]: |
| 83 | with open(filename, encoding="utf-8") as f: |
| 84 | return json.load(f) |
| 85 | |
| 86 | |
| 87 | @pytest.mark.parametrize( |
| 88 | "model", |
| 89 | [ |
| 90 | pytest.param( |
| 91 | "Qwen/Qwen2.5-Math-PRM-7B", |
| 92 | marks=[pytest.mark.core_model, pytest.mark.cpu_model], |
| 93 | ), |
| 94 | ], |
| 95 | ) |
| 96 | @pytest.mark.parametrize("dtype", ["half"]) |
| 97 | def test_prm_models( |
| 98 | hf_runner, |
| 99 | vllm_runner, |
| 100 | math_step_prompts, |
| 101 | monkeypatch, |
| 102 | model: str, |
| 103 | dtype: str, |
| 104 | ) -> None: |
| 105 | model_info = HF_EXAMPLE_MODELS.find_hf_info(model) |
| 106 | model_info.check_transformers_version(on_fail="skip") |
| 107 | |
| 108 | if current_platform.is_cpu(): |
| 109 | pytest.skip("CPU only supports V1") |
| 110 | |
| 111 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") |
| 112 | with vllm_runner( |
| 113 | model, |
| 114 | max_model_len=1024, |
| 115 | dtype=dtype, |
| 116 | kv_cache_memory_bytes=64 * MiB_bytes, |
| 117 | ) as vllm_model: |
| 118 | assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner |
| 119 | vllm_outputs = vllm_model.token_classify(math_step_prompts) |
| 120 | |
| 121 | with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model: |
| 122 | hf_model = step_reward_patch_hf_model(hf_model) |
| 123 | hf_outputs = hf_model.reward(math_step_prompts) |
| 124 | |
| 125 | dump_reward_outputs( |
| 126 | hf_outputs, |
| 127 | FIXTURE_REWARD_RESULT[model], |
| 128 | ) |
| 129 | |
| 130 | # check logits difference |
| 131 | for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): |
| 132 | hf_output = torch.tensor(hf_output).float() |
| 133 | vllm_output = torch.tensor(vllm_output).float() |
| 134 | |
| 135 | assert torch.allclose(hf_output, vllm_output, 1.5e-2) |
| 136 | |
| 137 | |
| 138 | @pytest.mark.parametrize( |
| 139 | "model", |
| 140 | [ |
| 141 | pytest.param( |
| 142 | "Qwen/Qwen2.5-Math-PRM-7B", |
| 143 | marks=[pytest.mark.core_model, pytest.mark.cpu_model], |
| 144 | ), |
| 145 | ], |
| 146 | ) |
| 147 | @pytest.mark.parametrize("dtype", ["half"]) |
| 148 | def test_prm_models_with_golden_outputs( |
| 149 | vllm_runner, |
| 150 | math_step_prompts, |
| 151 | monkeypatch, |
| 152 | model: str, |
| 153 | dtype: str, |
| 154 | ) -> None: |
| 155 | if not FIXTURE_REWARD_RESULT.get(model): |
| 156 | pytest.skip(f"No available golden outputs for {model}.") |
| 157 | |
| 158 | use_v2_model_runner = not current_platform.is_cpu() |
| 159 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_model_runner else "0") |
| 160 | with vllm_runner( |
| 161 | model, |
| 162 | max_model_len=1024, |
| 163 | dtype=dtype, |
| 164 | kv_cache_memory_bytes=64 * MiB_bytes, |
| 165 | ) as vllm_model: |
| 166 | assert ( |
| 167 | vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner |
| 168 | is use_v2_model_runner |
| 169 | ) |
| 170 | vllm_outputs = vllm_model.token_classify(math_step_prompts) |
| 171 | |
| 172 | golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model]) |
| 173 | |
| 174 | # check logits difference |
| 175 | for golden_output, vllm_output in zip(golden_outputs, vllm_outputs): |
| 176 | golden_output = torch.tensor(golden_output).float() |
| 177 | vllm_output = torch.tensor(vllm_output).float() |
| 178 | |
| 179 | assert torch.allclose(golden_output, vllm_output, 1.5e-2) |