diff --git a/docs/models/extensions/runai_model_streamer.md b/docs/models/extensions/runai_model_streamer.md index 965b2932ff..c60473bbf0 100644 --- a/docs/models/extensions/runai_model_streamer.md +++ b/docs/models/extensions/runai_model_streamer.md @@ -51,6 +51,20 @@ vllm serve s3://core-llm/Llama-3-8b \ --load-format runai_streamer ``` +## Model-specific checkpoint selection + +For models exposing `is_unused_checkpoint_weight`, `runai_streamer` skips +safetensors files containing only unused weights before streaming. It uses the +same model-owned rule as the default loader, including MTP draft models with +this hook; it does not assume a particular MTP tensor prefix. + +Object-storage checkpoints use `model.safetensors.index.json` and verify excluded +shards against their safetensors headers before streaming. Mixed shards remain +included. Missing or inconsistent metadata retains the original file list. +Models without a selection rule keep all files. Distributed streaming with +pipeline parallelism also retains the full file list. This optimization does +not change `runai_streamer_sharded`. + ## Tunable parameters You can tune parameters using `--model-loader-extra-config`: diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py index e0a083b3a8..999300e0fb 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py @@ -1,13 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json import os import types +from datetime import timedelta +from pathlib import Path from unittest.mock import patch import pytest import torch +from safetensors.torch import load_file, save_file +from tests.utils import create_new_process_for_each_test from vllm import SamplingParams from vllm.config.load import LoadConfig from vllm.model_executor.model_loader import get_model_loader @@ -39,12 +44,130 @@ def test_get_model_loader_with_runai_flag(): assert model_loader.__class__.__name__ == "RunaiModelStreamerLoader" +def test_checkpoint_selection_hook_is_optional_for_model_detection(): + from vllm.model_executor.models.interfaces import SupportsMultiModal + from vllm.model_executor.models.interfaces_base import ( + VllmModel, + is_text_generation_model, + ) + + class ModelWithoutHook(torch.nn.Module): + def __init__(self, vllm_config, prefix=""): + super().__init__() + + def embed_input_ids(self, input_ids): + return input_ids + + def forward(self, input_ids, positions): + return input_ids + + def compute_logits(self, hidden_states): + return hidden_states + + assert is_text_generation_model(ModelWithoutHook) + assert is_text_generation_model(ModelWithoutHook(None)) + + class ModelWithDefaultRule(ModelWithoutHook, VllmModel): + pass + + assert not ModelWithDefaultRule(None).is_unused_checkpoint_weight("any.weight") + + # Structural models do not inherit the protocol's optional default hook. + for names in ([], ["language_model"]): + model = torch.nn.Module() + model.language_model = ModelWithoutHook(None) + model._language_model_names = names + assert SupportsMultiModal.get_language_model(model) is model.language_model + + def test_runai_model_loader_download_files(vllm_runner): with vllm_runner(test_model, load_format=load_format) as llm: deserialized_outputs = llm.generate(prompts, sampling_params) assert deserialized_outputs +# In-process MTP compilation must not retain GPU resources in the pytest worker. +@create_new_process_for_each_test("spawn") +def test_runai_mtp_real_checkpoint(vllm_runner, tmp_path, monkeypatch): + """Skip main-only shards without changing a real MTP model's generation.""" + from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP + from vllm.platforms import current_platform + from vllm.transformers_utils.repo_utils import hf_api + from vllm.v1.metrics.reader import Counter + + if not current_platform.is_cuda_alike(): + pytest.skip("MTP serving requires CUDA or ROCm") + + checkpoint = tmp_path / "checkpoint" + hf_api().snapshot_download( + repo_id="Qwen/Qwen3.5-0.8B-Base", local_dir=str(checkpoint) + ) + original_files = list(checkpoint.glob("*.safetensors")) + tensors = {} + for path in original_files: + tensors.update(load_file(path)) + # Reshard real, unmodified weights so selection cannot pass vacuously on + # the published single mixed file. This predicate needs no model state. + selector = Qwen3_5MTP.__new__(Qwen3_5MTP) + main: dict[str, torch.Tensor] = {} + draft: dict[str, torch.Tensor] = {} + for name, tensor in tensors.items(): + (main if selector.is_unused_checkpoint_weight(name) else draft)[name] = tensor + assert main and draft + files = [checkpoint / "main.safetensors", checkpoint / "draft.safetensors"] + for path, contents in zip(files, (main, draft)): + save_file(contents, path, metadata={"format": "pt"}) + for path in original_files: + path.unlink() + (checkpoint / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + name: path.name + for path, data in zip(files, (main, draft)) + for name in data + } + } + ) + ) + del tensors, main, draft + + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + opened = [] + real_stream = rsl.runai_safetensors_weights_iterator + + def stream(paths, *args): + opened.append(list(paths)) + yield from real_stream(paths, *args) + + monkeypatch.setattr(rsl, "runai_safetensors_weights_iterator", stream) + results = [] + for load_format in ("auto", "runai_streamer"): + with vllm_runner( + str(checkpoint), + load_format=load_format, + block_size=None, + enable_chunked_prefill=True, + disable_log_stats=False, + limit_mm_per_prompt={"image": 0, "video": 0}, + speculative_config={"method": "mtp", "num_speculative_tokens": 1}, + ) as runner: + outputs = runner.generate_greedy(prompts, max_tokens=32) + results.append(outputs) + accepted = next( + m.value + for m in runner.llm.get_metrics() + if isinstance(m, Counter) + and m.name == "vllm:spec_decode_num_accepted_tokens" + ) + assert accepted > 0 + assert len(opened) == 2 + assert set(opened[0]) == {str(path) for path in files} + assert opened[1] == [str(files[1])] + matches = sum(a == b for a, b in zip(*results)) + assert matches >= 3, "RunAI MTP generation diverged from the Default loader" + + @pytest.mark.skip( reason="Temporarily disabled due to GCS access issues. " "TODO: Re-enable this test once the underlying issue is resolved." @@ -150,7 +273,7 @@ def test_runai_invalid_extra_config_leaves_environ_untouched(): ) def test_runai_get_all_weights_resolves_source(model_config, expected_source): fake_self = types.SimpleNamespace( - _get_weights_iterator=lambda path, revision: iter( + _get_weights_iterator=lambda path, revision, is_unused_weight=None: iter( [(f"{path}@{revision}", None)] ) ) @@ -164,6 +287,387 @@ def test_runai_get_all_weights_resolves_source(model_config, expected_source): assert weights == [(f"{expected_source}@myrev", None)] +@pytest.mark.parametrize("selection", ["draft", "unadapted", "reject_all"]) +def test_runai_selects_shards_before_streaming(tmp_path, monkeypatch, selection): + """Unused shards must not reach RunAI; mixed shards retain their tensors.""" + contents = [ + {"main.weight": torch.tensor([1.0])}, + {"main.bias": torch.tensor([2.0]), "draft.weight": torch.tensor([3.0])}, + {"draft.bias": torch.tensor([4.0])}, + ] + files = [] + for i, tensors in enumerate(contents): + path = str(tmp_path / f"part-{i}.safetensors") + save_file(tensors, path) + files.append(path) + monkeypatch.setattr(rsl, "list_safetensors", lambda path: files) + + opened = [] + real_stream = rsl.runai_safetensors_weights_iterator + + def stream(paths, *args): + opened.extend(paths) + yield from real_stream(paths, *args) + + monkeypatch.setattr(rsl, "runai_safetensors_weights_iterator", stream) + model = torch.nn.Module() + if selection == "draft": + model.is_unused_checkpoint_weight = lambda name: not name.startswith("draft.") + elif selection == "reject_all": + model.is_unused_checkpoint_weight = lambda name: True + config = types.SimpleNamespace(model=str(tmp_path), model_weights="", revision=None) + loaded = dict(_runai_loader({}).get_all_weights(config, model)) + + expected = files[1:] if selection == "draft" else files + assert opened == expected + assert set(loaded) == { + name + for tensors in (contents[1:] if selection == "draft" else contents) + for name in tensors + } + for name, tensor in {**contents[1], **contents[2]}.items(): + torch.testing.assert_close(loaded[name], tensor) + + +def _check_distributed_selection(rank, rendezvous, scenario): + torch.distributed.init_process_group( + "gloo", + init_method=rendezvous, + rank=rank, + world_size=2, + timeout=timedelta(seconds=30), + ) + root = "s3://bucket/weights" + if scenario == "different_roots": + root += str(rank) + files = [root + "/main.safetensors", root + "/draft.safetensors"] + + def pull_files(source, destination, **kwargs): + if scenario == "metadata_failure" and rank == 1: + raise OSError("Index temporarily unavailable") + (Path(destination) / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "main.weight": "main.safetensors", + "draft.weight": "draft.safetensors", + } + } + ) + ) + + try: + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(rsl, "list_safetensors", lambda path: files) + monkeypatch.setattr(rsl, "runai_pull_files", pull_files) + monkeypatch.setattr( + rsl, + "_get_remote_checkpoint_keys", + lambda paths: [ + ["main.weight"] + if path.endswith("main.safetensors") + else ["draft.weight"] + for path in paths + ], + ) + monkeypatch.setattr( + rsl, + "get_world_group", + lambda: types.SimpleNamespace( + world_size=2, cpu_group=torch.distributed.group.WORLD + ), + raising=False, + ) + monkeypatch.setattr( + rsl, + "get_pp_group", + lambda: types.SimpleNamespace( + world_size=2 if scenario == "pipeline_draft" else 1 + ), + raising=False, + ) + if scenario == "pipeline_draft" and rank == 0: + # Earlier PP stages never invoke the draft loader. + torch.distributed.barrier() + return + # Exercise real selection and Gloo; observe the distributed stream input. + monkeypatch.setattr( + rsl, + "runai_safetensors_weights_iterator", + lambda paths, *args: iter((path, None) for path in paths), + ) + model = torch.nn.Module() + if not (scenario == "missing_hook" and rank == 1): + needed = ( + "main." if scenario == "different_roles" and rank == 1 else "draft." + ) + model.is_unused_checkpoint_weight = lambda name: ( + not name.startswith(needed) + ) + config = types.SimpleNamespace(model=root, model_weights="", revision=None) + selected = [ + path + for path, _ in _runai_loader({"distributed": True}).get_all_weights( + config, model + ) + ] + expected = files[1:] if scenario == "healthy" else files + assert selected == expected + if scenario == "pipeline_draft": + torch.distributed.barrier() + finally: + torch.distributed.destroy_process_group() + + +@pytest.mark.parametrize( + "scenario", + [ + "healthy", + "metadata_failure", + "different_roles", + "missing_hook", + "different_roots", + "pipeline_draft", + ], +) +def test_distributed_runai_preserves_every_ranks_required_shards(tmp_path, scenario): + """All ranks must stream the same union even if one cannot filter safely.""" + torch.multiprocessing.spawn( + _check_distributed_selection, + args=((tmp_path / "rendezvous").as_uri(), scenario), + nprocs=2, + ) + + +@pytest.mark.parametrize("scheme", ["s3", "gs", "az"]) +@pytest.mark.parametrize( + "index_kind", + [ + "valid", + "missing", + "malformed", + "unavailable", + "unadapted", + "missing_shard", + "stale", + "header_unavailable", + ], +) +def test_runai_remote_selection_uses_actual_weight_source( + tmp_path, monkeypatch, scheme, index_kind +): + root = f"{scheme}://bucket/weights" + files = [ + root + "/main/part.safetensors", + root + "/draft/part.safetensors", + root + "/mixed.safetensors", + root + "/extra.safetensors", + ] + monkeypatch.setattr(rsl, "list_safetensors", lambda path: files) + pulled = [] + + def pull_files(source, destination, allow_pattern, ignore_pattern=None): + pulled.append((source.rstrip("/"), allow_pattern)) + if index_kind == "unavailable": + # Storage SDKs do not necessarily derive errors from OSError. + raise Exception("Object metadata unavailable") + index = Path(destination) / "model.safetensors.index.json" + if index_kind in ("valid", "missing_shard", "stale", "header_unavailable"): + index.write_text( + json.dumps( + { + "weight_map": { + "main.weight": "main/part.safetensors", + "draft.weight": ( + "absent.safetensors" + if index_kind == "missing_shard" + else "draft/part.safetensors" + ), + "main.bias": "mixed.safetensors", + "draft.bias": "mixed.safetensors", + } + } + ) + ) + elif index_kind == "malformed": + index.write_text("not JSON") + + monkeypatch.setattr(rsl, "runai_pull_files", pull_files, raising=False) + headers_checked = [] + + def checkpoint_keys(paths): + headers_checked.extend(paths) + if index_kind == "header_unavailable": + raise OSError("Header temporarily unavailable") + return [ + ["main.weight", "draft.unindexed_weight"] + if index_kind == "stale" + else ["main.weight"] + for _ in paths + ] + + monkeypatch.setattr( + rsl, "_get_remote_checkpoint_keys", checkpoint_keys, raising=False + ) + opened = [] + + def stream(paths, *args): + opened.extend(paths) + return iter(()) + + monkeypatch.setattr(rsl, "runai_safetensors_weights_iterator", stream) + model = torch.nn.Module() + if index_kind != "unadapted": + model.is_unused_checkpoint_weight = lambda name: not name.startswith("draft.") + # The config directory is not the checkpoint: using its index could drop + # tensors from an independent model_weights source. + (tmp_path / "model.safetensors.index.json").write_text("not the weights index") + config = types.SimpleNamespace( + model=str(tmp_path), model_weights=root, revision=None + ) + list(_runai_loader({}).get_all_weights(config, model)) + + assert opened == (files[1:] if index_kind == "valid" else files) + assert headers_checked == ( + files[:1] if index_kind in ("valid", "stale", "header_unavailable") else [] + ) + assert pulled == ( + [] + if index_kind == "unadapted" + else [ + ( + root, + [ + "weights/model.safetensors.index.json", + "model.safetensors.index.json", + ], + ) + ] + ) + + +@pytest.mark.parametrize( + "weight_map", + [None, {}, [], {"draft.weight": None}, {"draft.weight": "../part.safetensors"}], +) +def test_runai_invalid_weight_map_keeps_original_files(monkeypatch, weight_map): + files = ["s3://bucket/model/part.safetensors"] + monkeypatch.setattr(rsl, "list_safetensors", lambda path: files) + + def pull_files(source, directory, **kwargs): + (Path(directory) / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}) + ) + + monkeypatch.setattr(rsl, "runai_pull_files", pull_files) + monkeypatch.setattr( + rsl, "runai_safetensors_weights_iterator", lambda paths, *args: iter(paths) + ) + model = torch.nn.Module() + model.is_unused_checkpoint_weight = lambda name: True + config = types.SimpleNamespace( + model="s3://bucket/model", model_weights="", revision=None + ) + assert list(_runai_loader({}).get_all_weights(config, model)) == files + + +def test_runai_model_rule_error_is_not_hidden_by_metadata_fallback(monkeypatch): + files = ["s3://bucket/model/part.safetensors"] + monkeypatch.setattr(rsl, "list_safetensors", lambda path: files) + + def pull_files(source, directory, **kwargs): + (Path(directory) / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"draft.weight": "part.safetensors"}}) + ) + + monkeypatch.setattr(rsl, "runai_pull_files", pull_files) + model = torch.nn.Module() + + def invalid_rule(name): + raise ValueError("Invalid model rule") + + model.is_unused_checkpoint_weight = invalid_rule + config = types.SimpleNamespace( + model="s3://bucket/model", model_weights="", revision=None + ) + with pytest.raises(ValueError, match="Invalid model rule"): + list(_runai_loader({}).get_all_weights(config, model)) + + +def test_runai_header_reader_uses_actual_tensor_names(tmp_path): + """The RunAI header reader sees names even when an index omits them.""" + path = tmp_path / "mixed.safetensors" + save_file( + { + "main.weight": torch.tensor([1.0]), + "draft.unindexed_weight": torch.tensor([2.0]), + }, + path, + metadata={"format": "pt"}, + ) + (names,) = rsl._get_remote_checkpoint_keys([str(path)]) + assert set(names) == {"main.weight", "draft.unindexed_weight"} + + +@pytest.mark.parametrize("stale", [False, True]) +def test_runai_stale_index_keeps_actual_draft_weights(tmp_path, monkeypatch, stale): + """Verify excluded headers before dropping a shard with real draft weights.""" + root = "s3://bucket/checkpoint" + main_path = tmp_path / "main.safetensors" + draft_path = tmp_path / "draft.safetensors" + main_tensors = {"main.weight": torch.tensor([1.0])} + if stale: + main_tensors["draft.unindexed_weight"] = torch.tensor([2.0]) + save_file(main_tensors, main_path) + save_file({"draft.weight": torch.tensor([3.0])}, draft_path) + paths = { + root + "/main.safetensors": str(main_path), + root + "/draft.safetensors": str(draft_path), + } + monkeypatch.setattr(rsl, "list_safetensors", lambda path: list(paths)) + + def pull_files(source, directory, **kwargs): + (Path(directory) / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "main.weight": "main.safetensors", + "draft.weight": "draft.safetensors", + } + } + ) + ) + + monkeypatch.setattr(rsl, "runai_pull_files", pull_files) + read_headers = rsl._get_remote_checkpoint_keys + monkeypatch.setattr( + rsl, + "_get_remote_checkpoint_keys", + lambda files: read_headers([paths[path] for path in files]), + ) + stream = rsl.runai_safetensors_weights_iterator + + def stream_local(files, *args): + for name, tensor in stream([paths[path] for path in files], *args): + yield name, tensor.clone() + + monkeypatch.setattr(rsl, "runai_safetensors_weights_iterator", stream_local) + model = torch.nn.Module() + model.is_unused_checkpoint_weight = lambda name: not name.startswith("draft.") + config = types.SimpleNamespace(model=root, model_weights="", revision=None) + loaded = dict(_runai_loader({}).get_all_weights(config, model)) + assert set(loaded) == ( + {"main.weight", "draft.unindexed_weight", "draft.weight"} + if stale + else {"draft.weight"} + ) + torch.testing.assert_close(loaded["draft.weight"], torch.tensor([3.0])) + if stale: + torch.testing.assert_close( + loaded["draft.unindexed_weight"], torch.tensor([2.0]) + ) + + @pytest.mark.parametrize( "base_model,mul_model,add_model", [ diff --git a/tests/models/test_language_model_cache_is_weak.py b/tests/models/test_language_model_cache_is_weak.py index 19d93e479e..1a007275e7 100644 --- a/tests/models/test_language_model_cache_is_weak.py +++ b/tests/models/test_language_model_cache_is_weak.py @@ -12,9 +12,13 @@ import gc import weakref import pytest +import torch import torch.nn as nn -from vllm.model_executor.models.interfaces import _language_model_by_module +from vllm.model_executor.models.interfaces import ( + SupportsMultiModal, + _language_model_by_module, +) pytestmark = pytest.mark.cpu_test @@ -57,3 +61,42 @@ def test_cached_entry_is_dropped_with_its_key(): gc.collect() assert not any(isinstance(k, _MultiModalModel) for k in _language_model_by_module) + + +@pytest.mark.parametrize("marked", [False, True]) +@pytest.mark.parametrize("text_model", ["clip", "siglip"]) +def test_text_tower_without_loader_hook_is_found(text_model, marked, caplog): + from vllm.model_executor.models.clip import CLIPTextTransformer + from vllm.model_executor.models.siglip import SiglipTextTransformer + + cls = CLIPTextTransformer if text_model == "clip" else SiglipTextTransformer + # Only token embeddings are needed to exercise text tower discovery. + tower = cls.__new__(cls) + nn.Module.__init__(tower) + tower.embeddings = nn.Module() + tower.embeddings.token_embedding = nn.Embedding(4, 2) + model = nn.Module() + model.text_model = tower + model._language_model_names = ["text_model"] if marked else [] + + caplog.clear() + language_model = SupportsMultiModal.get_language_model(model) + + assert language_model is tower + input_ids = torch.tensor([0, 1]) + torch.testing.assert_close( + language_model.embed_input_ids(input_ids), + tower.embeddings.token_embedding.weight[:2], + ) + assert not caplog.records + + +@pytest.mark.parametrize("marked", [False, True]) +@pytest.mark.parametrize("method", ["embed_input_ids", "forward"]) +def test_language_model_discovery_rejects_noncallable_methods(method, marked): + model = _MultiModalModel() + model._language_model_names = ["language_model"] if marked else [] + setattr(model.language_model, method, None) + + with pytest.raises(NotImplementedError, match="No language model found"): + SupportsMultiModal.get_language_model(model) diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index e001160062..2f0f76f140 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -1,7 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json import os -from collections.abc import Generator +import posixpath +import tempfile +from collections.abc import Callable, Generator +from typing import Any +from urllib.parse import urlsplit import torch from torch import nn @@ -9,14 +14,110 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from vllm.config import ModelConfig from vllm.config.load import LoadConfig +from vllm.distributed import get_pp_group, get_world_group +from vllm.logger import init_logger from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.weight_utils import ( download_safetensors_index_file_from_hf, download_weights_from_hf, + filter_safetensors_files_by_weight_name, runai_safetensors_weights_iterator, ) from vllm.transformers_utils.repo_utils import resolve_revision -from vllm.transformers_utils.runai_utils import is_runai_obj_uri, list_safetensors +from vllm.transformers_utils.runai_utils import ( + is_runai_obj_uri, + list_safetensors, + runai_pull_files, +) + +logger = init_logger(__name__) + + +def _get_remote_checkpoint_keys(files: list[str]) -> list[list[str]]: + from runai_model_streamer import DistributedStreamer + from runai_model_streamer.safetensors_streamer.safetensors_pytorch import ( + prepare_request, + ) + + # RunAI's metadata reader uses header ranges and the streamer's backend. + with DistributedStreamer() as streamer: + metadata = prepare_request(streamer, files, None) + return [[tensor.name for tensor in tensors] for _, tensors, _ in metadata] + + +def _filter_remote_weights( + files: list[str], + source: str, + is_unused_weight: Callable[[str], bool], +) -> list[str]: + try: + with tempfile.TemporaryDirectory() as directory: + # Backends may match the full object key, not its basename. + index_key = posixpath.join( + urlsplit(source).path.strip("/"), SAFE_WEIGHTS_INDEX_NAME + ) + runai_pull_files( + source.rstrip("/") + "/", + directory, + allow_pattern=[index_key, SAFE_WEIGHTS_INDEX_NAME], + ) + with open(os.path.join(directory, SAFE_WEIGHTS_INDEX_NAME)) as index: + weight_map = json.load(index)["weight_map"] + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError("Empty or invalid weight map") + if any( + not isinstance(name, str) + or not isinstance(shard, str) + or not shard + or posixpath.isabs(shard) + or posixpath.normpath(shard) != shard + or shard in (".", "..") + or shard.startswith("../") + for name, shard in weight_map.items() + ): + raise ValueError("Invalid checkpoint index entry") + except Exception: + # Optional metadata must not make an otherwise loadable model fail. + # Model-rule and tensor-read errors remain outside this boundary. + logger.warning( + "Cannot use checkpoint index for RunAI shard selection; " + "keeping the original file list" + ) + return files + + # Unindexed files cannot be excluded using this metadata. + root = source.rstrip("/") + "/" + indexed = {root + shard for shard in weight_map.values()} + needed = { + root + shard for name, shard in weight_map.items() if not is_unused_weight(name) + } + if not needed.issubset(files): + logger.warning( + "Checkpoint index references missing required RunAI shards; " + "keeping the original file list" + ) + return files + kept = [path for path in files if path not in indexed or path in needed] + excluded = [path for path in files if path not in kept] + if excluded and kept: + try: + headers = _get_remote_checkpoint_keys(excluded) + if len(headers) != len(excluded): + raise ValueError("Incomplete checkpoint headers") + except Exception: + logger.warning( + "Cannot verify excluded RunAI shard headers; " + "keeping the original file list" + ) + return files + # A stale index can omit wanted tensors inside an indexed shard. + if any(not is_unused_weight(name) for names in headers for name in names): + logger.warning( + "Checkpoint index disagrees with RunAI shard headers; " + "keeping the original file list" + ) + return files + return kept or files class RunaiModelStreamerLoader(BaseModelLoader): @@ -122,12 +223,43 @@ class RunaiModelStreamerLoader(BaseModelLoader): return hf_weights_files def _get_weights_iterator( - self, model_or_path: str, revision: str | None + self, + model_or_path: str, + revision: str | None, + is_unused_weight: Callable[[str], bool] | None = None, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" hf_weights_files = self._prepare_weights(model_or_path, revision) + selected_files = hf_weights_files + distributed = self._is_distributed and torch.distributed.is_initialized() + # Drafts load only on the last PP stage; RunAI is invoked without a + # model-specific group, so retain its existing full-list path with PP. + can_select = not distributed or get_pp_group().world_size == 1 + if is_unused_weight is not None and can_select: + if is_runai_obj_uri(model_or_path): + selected_files = _filter_remote_weights( + hf_weights_files, model_or_path, is_unused_weight + ) + else: + selected_files = filter_safetensors_files_by_weight_name( + hf_weights_files, is_unused_weight + ) + if distributed and can_select: + # RunAI broadcasts across node-local or WORLD groups, not just TP. + # A rank that cannot filter safely keeps the group on the full list. + group = get_world_group() + selections: list[Any] = [None] * group.world_size + torch.distributed.all_gather_object( + selections, (hf_weights_files, selected_files), group=group.cpu_group + ) + if all(files == hf_weights_files for files, _ in selections): + needed = {path for _, selected in selections for path in selected} + selected_files = [path for path in hf_weights_files if path in needed] + else: + # Do not substitute another node's local checkpoint paths. + selected_files = hf_weights_files return runai_safetensors_weights_iterator( - hf_weights_files, self.load_config.use_tqdm_on_load, self._is_distributed + selected_files, self.load_config.use_tqdm_on_load, self._is_distributed ) def download_model(self, model_config: ModelConfig) -> None: @@ -142,7 +274,11 @@ class RunaiModelStreamerLoader(BaseModelLoader): model_weights = model_config.model if model_weights_override := model_config.model_weights: model_weights = model_weights_override - yield from self._get_weights_iterator(model_weights, model_config.revision) + yield from self._get_weights_iterator( + model_weights, + model_config.revision, + getattr(model, "is_unused_checkpoint_weight", None), + ) def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: """Load weights into a model.""" diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index a780cbe23c..5f51ece0b3 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -141,6 +141,13 @@ _language_model_by_module: "weakref.WeakKeyDictionary[nn.Module, VllmModel]" = ( ) +def _is_language_model(model: object) -> TypeIs[VllmModel]: + # Text towers can use config/position_ids signatures and omit loader hooks. + return callable(getattr(model, "embed_input_ids", None)) and callable( + getattr(model, "forward", None) + ) + + @runtime_checkable class SupportsMultiModalEmbeddings(Protocol): """The interface for models that can merge external multimodal embeddings.""" @@ -284,14 +291,14 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol): mod = getattr(mod, attr) candidate: object = mod - if candidate is not self_module and isinstance(candidate, VllmModel): + if candidate is not self_module and _is_language_model(candidate): _language_model_by_module[self_module] = candidate return candidate # Fallback for mod in self_module.children(): candidate = mod - if isinstance(candidate, VllmModel): + if _is_language_model(candidate): _language_model_by_module[self_module] = candidate return candidate diff --git a/vllm/model_executor/models/interfaces_base.py b/vllm/model_executor/models/interfaces_base.py index 7c7b291905..4b265173ab 100644 --- a/vllm/model_executor/models/interfaces_base.py +++ b/vllm/model_executor/models/interfaces_base.py @@ -51,6 +51,15 @@ class VllmModel(Protocol): self, input_ids: torch.Tensor, positions: torch.Tensor ) -> torch.Tensor: ... + def is_unused_checkpoint_weight(self, name: str) -> bool: + """Return True for checkpoint weights this model does not load. + + Defining this method is optional for satisfying this interface. + Override with the same rule used in ``load_weights``. Models without + a selection rule keep every checkpoint weight by default. + """ + return False + def _check_vllm_model_init(model: type[object] | object) -> bool: model_init = getattr(model, "__init__") # noqa: B009 @@ -134,10 +143,8 @@ def is_text_generation_model( if not is_vllm_model(model): return False - if isinstance(model, type): - return isinstance(model, VllmModelForTextGeneration) - - return isinstance(model, VllmModelForTextGeneration) + # Optional protocol defaults are not inherited by structural models. + return callable(getattr(model, "compute_logits", None)) @runtime_checkable