| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | import torch.nn as nn |
| 6 | from transformers import AutoConfig, AutoModel, CLIPImageProcessor |
| 7 | |
| 8 | from vllm.distributed import cleanup_dist_env_and_memory |
| 9 | from vllm.model_executor.models.radio import RadioModel |
| 10 | from vllm.platforms import current_platform |
| 11 | from vllm.transformers_utils.configs.radio import RadioConfig |
| 12 | from vllm.transformers_utils.repo_utils import hf_api |
| 13 | from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE |
| 14 | |
| 15 | from ....conftest import ImageTestAssets |
| 16 | |
| 17 | # we use snapshot_download to prevent conflicts between |
| 18 | # dynamic_module and trust_remote_code for hf_runner |
| 19 | DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"] |
| 20 | |
| 21 | DEVICE_TYPE = current_platform.device_type |
| 22 | |
| 23 | |
| 24 | @torch.inference_mode() |
| 25 | def run_radio_test( |
| 26 | image_assets: ImageTestAssets, |
| 27 | model_id: str, |
| 28 | *, |
| 29 | dtype: str, |
| 30 | ): |
| 31 | model = hf_api().snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) |
| 32 | torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype] |
| 33 | |
| 34 | img_processor = CLIPImageProcessor.from_pretrained(model) |
| 35 | images = [asset.pil_image for asset in image_assets] |
| 36 | # Input resolution must be a multiple of `self.min_resolution_step`. |
| 37 | # Using `self.get_nearest_supported_resolution`, for assets 432x642 the |
| 38 | # nearest supported resolution is 432x640. |
| 39 | pixel_values = [ |
| 40 | img_processor(image, return_tensors="pt").pixel_values.to(torch_dtype)[ |
| 41 | :, :, :, :640 |
| 42 | ] |
| 43 | for image in images |
| 44 | ] |
| 45 | |
| 46 | hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) |
| 47 | |
| 48 | # RADIO model on HF does not properly handle torch_dtype argument |
| 49 | # And relies on args["dtype"] which we have to patch manually: |
| 50 | hf_config.args["dtype"] = torch_dtype |
| 51 | |
| 52 | hf_model = AutoModel.from_pretrained( |
| 53 | model_id, |
| 54 | config=hf_config, |
| 55 | dtype=torch_dtype, |
| 56 | trust_remote_code=True, |
| 57 | ).to(DEVICE_TYPE) |
| 58 | hf_model.eval() |
| 59 | |
| 60 | # A HF model has image normalization as a part of model's forward |
| 61 | # However in vLLM we don't make normalization a part of the model |
| 62 | # forward step since mean/std stored as model's parameters and |
| 63 | # subject to precision loss (when using fp16/bf16) which negatively |
| 64 | # affects evaluation benchmarks. |
| 65 | hf_model.make_preprocessor_external() |
| 66 | |
| 67 | hf_outputs_per_image = [ |
| 68 | hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values |
| 69 | ] |
| 70 | |
| 71 | vllm_config = RadioConfig( |
| 72 | model_name=hf_config.args["model"], |
| 73 | **hf_config.args, |
| 74 | ) |
| 75 | vllm_model = RadioModel(vllm_config) |
| 76 | vllm_model.load_weights(hf_model.state_dict()) |
| 77 | vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype) |
| 78 | |
| 79 | vllm_outputs_per_image = [ |
| 80 | vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE)) |
| 81 | for pixel_value in pixel_values |
| 82 | ] |
| 83 | del vllm_model, hf_model |
| 84 | cleanup_dist_env_and_memory() |
| 85 | |
| 86 | cos_similar = nn.CosineSimilarity(dim=-1) |
| 87 | for vllm_output, hf_output in zip(vllm_outputs_per_image, hf_outputs_per_image): |
| 88 | assert cos_similar(vllm_output[0], hf_output[0]).mean() > 0.99 |
| 89 | assert cos_similar(vllm_output[1], hf_output[1]).mean() > 0.99 |
| 90 | |
| 91 | |
| 92 | @pytest.mark.parametrize( |
| 93 | "model_id", |
| 94 | [ |
| 95 | "nvidia/C-RADIOv2-H", |
| 96 | ], |
| 97 | ) |
| 98 | @pytest.mark.parametrize("dtype", ["half", "bfloat16"]) |
| 99 | def test_radio( |
| 100 | default_vllm_config, dist_init, image_assets, model_id, dtype: str |
| 101 | ) -> None: |
| 102 | run_radio_test( |
| 103 | image_assets, |
| 104 | model_id, |
| 105 | dtype=dtype, |
| 106 | ) |