[Feature][Model Loader] Apply model-owned MTP shard selection to RunAI ModelStreamer (#60759)

Signed-off-by: 浩韵 <zuohaiyu.zhy@alibaba-inc.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>

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`:

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",
     [

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)

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."""

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
 

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