| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | from torch import nn |
| 6 | |
| 7 | from vllm.config import ModelConfig |
| 8 | from vllm.config.load import LoadConfig |
| 9 | from vllm.model_executor.model_loader import get_model_loader, register_model_loader |
| 10 | from vllm.model_executor.model_loader.base_loader import BaseModelLoader |
| 11 | from vllm.model_executor.model_loader.default_loader import DefaultModelLoader |
| 12 | |
| 13 | |
| 14 | @register_model_loader("custom_load_format") |
| 15 | class CustomModelLoader(BaseModelLoader): |
| 16 | def __init__(self, load_config: LoadConfig) -> None: |
| 17 | super().__init__(load_config) |
| 18 | |
| 19 | def download_model(self, model_config: ModelConfig) -> None: |
| 20 | pass |
| 21 | |
| 22 | def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: |
| 23 | pass |
| 24 | |
| 25 | |
| 26 | def test_register_model_loader(): |
| 27 | load_config = LoadConfig(load_format="custom_load_format") |
| 28 | assert isinstance(get_model_loader(load_config), CustomModelLoader) |
| 29 | |
| 30 | |
| 31 | def test_invalid_model_loader(): |
| 32 | with pytest.raises(ValueError): |
| 33 | |
| 34 | @register_model_loader("invalid_load_format") |
| 35 | class InValidModelLoader: |
| 36 | pass |
| 37 | |
| 38 | |
| 39 | def test_default_loader_rejects_zero_num_threads(): |
| 40 | # num_threads=0 used to fail late in ThreadPoolExecutor ("max_workers must be > 0"). |
| 41 | with pytest.raises(ValueError, match="num_threads"): |
| 42 | DefaultModelLoader( |
| 43 | LoadConfig( |
| 44 | model_loader_extra_config={ |
| 45 | "enable_multithread_load": True, |
| 46 | "num_threads": 0, |
| 47 | } |
| 48 | ) |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | def test_default_loader_rejects_multithread_with_non_lazy_strategy(): |
| 53 | # The multi-thread loader ignores safetensors_load_strategy; reject the |
| 54 | # combination instead of silently dropping the requested strategy. |
| 55 | with pytest.raises(ValueError, match="does not support"): |
| 56 | DefaultModelLoader( |
| 57 | LoadConfig( |
| 58 | safetensors_load_strategy="torchao", |
| 59 | model_loader_extra_config={"enable_multithread_load": True}, |
| 60 | ) |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def test_default_loader_explicit_safetensors_does_not_misread_pt(tmp_path): |
| 65 | # Explicit safetensors must not fall back to a .pt and open it as safetensors. |
| 66 | (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00") |
| 67 | loader = DefaultModelLoader(LoadConfig(load_format="safetensors")) |
| 68 | with pytest.raises(RuntimeError, match="Cannot find any model weights"): |
| 69 | loader._prepare_weights( |
| 70 | str(tmp_path), |
| 71 | None, |
| 72 | None, |
| 73 | fall_back_to_pt=True, |
| 74 | allow_patterns_overrides=None, |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | def test_default_loader_hf_still_falls_back_to_pt(tmp_path): |
| 79 | # Control: load_format="hf" still picks up .pt weights via fallback. |
| 80 | (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00") |
| 81 | loader = DefaultModelLoader(LoadConfig(load_format="hf")) |
| 82 | _, files, use_safetensors, _ = loader._prepare_weights( |
| 83 | str(tmp_path), |
| 84 | None, |
| 85 | None, |
| 86 | fall_back_to_pt=True, |
| 87 | allow_patterns_overrides=None, |
| 88 | ) |
| 89 | assert use_safetensors is False |
| 90 | assert any(f.endswith("model.pt") for f in files) |