feat: MTP loading improvements

modelexpress_client/python/modelexpress/engines/vllm/adapter.py

@@ -40,6 +40,12 @@
 
 _SAFETENSORS_INDEX_NAME = "model.safetensors.index.json"
 
+# Registries on compilation_config that vLLM keys by layer name.
+_LAYER_REGISTRY_FIELDS: tuple[str, ...] = (
+    "static_forward_context",
+    "static_all_moe_layers",
+)
+
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
 
@@ -85,13 +91,21 @@ def _read_safetensors_index(model_uri: str) -> dict | None:
     from runai_model_streamer import pull_files
 
     with tempfile.TemporaryDirectory() as tmp:
-        pull_files(model_uri, tmp, allow_pattern=[_SAFETENSORS_INDEX_NAME])
+        # runai's allow_pattern is a glob matched against the full object key,
+        # so a bare filename never matches; anchor it with a leading wildcard.
+        pull_files(model_uri, tmp, allow_pattern=[f"*{_SAFETENSORS_INDEX_NAME}"])
         for root, _dirs, files in os.walk(tmp):
             if _SAFETENSORS_INDEX_NAME in files:
                 with open(
                     os.path.join(root, _SAFETENSORS_INDEX_NAME), encoding="utf-8"
                 ) as handle:
                     return json.load(handle)
+    logger.warning(
+        "safetensors index %s not found under %s; draft-shard selection will "
+        "fall back to streaming all shards",
+        _SAFETENSORS_INDEX_NAME,
+        model_uri,
+    )
     return None
 
 
@@ -323,12 +337,16 @@ def after_native_load(self, result: LoadResult) -> LoadResult:
     def reinit_for_retry(self, result: LoadResult) -> LoadResult:
         from vllm.model_executor.model_loader.utils import initialize_model
 
-        old_value = result.value
+        stale_value = result.value
+        stale_model = result.model
         result.value = None
         result.model = None
-        del old_value
+        # Unregister before dropping the model: its registrations identify it,
+        # and clearing them frees its parameters before the rebuild allocates.
+        self._unregister_model_layers(stale_model)
+        del stale_value
+        del stale_model
         self.accelerator_backend.empty_cache()
-        self._reset_compilation_state()
         logger.info(
             "[Worker %s] Re-initializing vLLM model after failed strategy",
             self.get_global_rank(),
@@ -404,21 +422,35 @@ def _resolve_target_device(self) -> torch.device:
         )
         return torch.device(load_device)
 
-    def _reset_compilation_state(self) -> None:
-        compilation_config = self.vllm_config.compilation_config
-        # vLLM registers each attention / MLA / Mamba / FusedMoE layer into
-        # fields on vllm_config.compilation_config during initialize_model().
-        # Those fields live on the config object, not the model, so they survive
-        # del model and trip duplicate registration on the next initialize_model().
-        # Clear them so re-init starts from a clean slate. Audited against vLLM
-        # 0.17.1; other versions may add init=False fields that need similar
-        # treatment.
-        compilation_config.static_forward_context.clear()
-        compilation_config.static_all_moe_layers.clear()
-        compilation_config.enabled_custom_ops.clear()
-        compilation_config.disabled_custom_ops.clear()
-        compilation_config.traced_files.clear()
-        compilation_config.compilation_time = 0.0
+    def _unregister_model_layers(self, stale_model: torch.nn.Module | None) -> None:
+        """Remove `stale_model`'s layers from vLLM's layer registries.
+
+        The registries live on compilation_config, so dropping the model leaves
+        its entries behind and the rebuild fails vLLM's duplicate-name check.
+        Clearing them is wrong: one compilation_config is shared by every model
+        built from a VllmConfig, so under MTP they also hold the live target's
+        layers. Remove only what this model registered, matched by its own
+        modules or the `layer_name` they registered under.
+
+        Args:
+            stale_model: Model being discarded, or None to clear outright.
+        """
+        owned_ids: set[int] = set()
+        owned_names: set[str] = set()
+        for module in stale_model.modules() if stale_model is not None else ():
+            owned_ids.add(id(module))
+            layer_name = getattr(module, "layer_name", None)
+            if isinstance(layer_name, str):
+                owned_names.add(layer_name)
+
+        for attr in _LAYER_REGISTRY_FIELDS:
+            registry = getattr(self.vllm_config.compilation_config, attr, None)
+            if registry is None:
+                continue
+            if stale_model is None:
+                registry.clear()
+            else:
+                _drop_owned_entries(attr, registry, owned_ids, owned_names)
 
     def _model_streamer_distributed_enabled(self) -> bool:
         tp_size = getattr(self.vllm_config.parallel_config, "tensor_parallel_size", 1)
@@ -439,6 +471,35 @@ def _is_same_or_descendant(name: str, prefix: str) -> bool:
     return prefix == "" or name == prefix or name.startswith(f"{prefix}.")
 
 
+def _drop_owned_entries(
+    attr: str,
+    registry,
+    owned_ids: set[int],
+    owned_names: set[str],
+) -> None:
+    """Remove one compilation registry's entries belonging to a single model."""
+
+    def is_owned(entry) -> bool:
+        return id(entry) in owned_ids or (isinstance(entry, str) and entry in owned_names)
+
+    if isinstance(registry, dict):
+        for key in [k for k, v in registry.items() if is_owned(k) or is_owned(v)]:
+            del registry[key]
+    elif isinstance(registry, set):
+        registry.difference_update({e for e in registry if is_owned(e)})
+    elif isinstance(registry, list):
+        registry[:] = [e for e in registry if not is_owned(e)]
+    else:
+        # A leftover entry only fails the rebuild's duplicate-name check, while
+        # clearing blind could unregister a co-owner's layers.
+        logger.warning(
+            "compilation_config.%s is a %s, which cannot be filtered by owner; "
+            "leaving it untouched",
+            attr,
+            type(registry).__name__,
+        )
+
+
 def _get_vllm_worker_rank(
     vllm_config: VllmConfig, target_device: torch.device
 ) -> int:

modelexpress_client/python/tests/test_vllm_adapter.py

@@ -16,8 +16,10 @@
 from modelexpress.engines.vllm.adapter import (
     DraftShardSelection,
     VllmAdapter,
+    _SAFETENSORS_INDEX_NAME,
     _get_vllm_device_id,
     _get_vllm_worker_rank,
+    _read_safetensors_index,
     _select_draft_weight_files,
     build_vllm_load_context,
 )
@@ -395,3 +397,77 @@ def test_unreadable_index_reports_unresolved(self, tmp_path):
                 DraftShardSelection.UNRESOLVED,
                 [],
             )
+
+
+def _stub_runai(monkeypatch, available: dict[str, str]) -> list:
+    """Install a fake runai_model_streamer whose pull_files mirrors runai's
+    real semantics: allow_pattern is fnmatched against the full object key, so
+    an unanchored bare filename matches nothing. Returns the list of
+    allow_pattern values it was called with.
+
+    ``available`` maps object basenames (relative to model_uri) to file content.
+    """
+    import fnmatch
+
+    calls: list = []
+
+    def pull_files(model_uri, dest, allow_pattern=None):
+        calls.append(allow_pattern)
+        for key, content in available.items():
+            full_key = f"{model_uri.rstrip('/')}/{key}"
+            if any(fnmatch.fnmatch(full_key, pat) for pat in (allow_pattern or ["*"])):
+                with open(os.path.join(dest, key), "w", encoding="utf-8") as handle:
+                    handle.write(content)
+
+    module = ModuleType("runai_model_streamer")
+    module.pull_files = pull_files
+    monkeypatch.setitem(sys.modules, "runai_model_streamer", module)
+    return calls
+
+
+class TestReadSafetensorsIndexObjectStore:
+    """Reading the index from an object store depends on runai's glob matching
+    the full object key; the bare filename this once used matched nothing."""
+
+    def test_reads_index_via_anchored_glob(self, monkeypatch):
+        index = {"weight_map": {"mtp.fc.weight": "model-mtp.safetensors"}}
+        calls = _stub_runai(
+            monkeypatch,
+            {
+                _SAFETENSORS_INDEX_NAME: json.dumps(index),
+                "model-00001-of-00001.safetensors": "weights",
+            },
+        )
+        # Reverting to a bare, unanchored pattern makes the fake fnmatch miss,
+        # so this returns None and the assertion fails, as it should.
+        assert _read_safetensors_index("s3://bucket/model") == index
+        assert calls == [[f"*{_SAFETENSORS_INDEX_NAME}"]]
+
+    def test_returns_none_and_warns_when_index_absent(self, monkeypatch, caplog):
+        _stub_runai(monkeypatch, {"model-00001-of-00001.safetensors": "weights"})
+        with caplog.at_level("WARNING", logger="modelexpress.engines.vllm.adapter"):
+            assert _read_safetensors_index("s3://bucket/model") is None
+        assert any("not found under" in rec.message for rec in caplog.records)
+
+    def test_selects_mtp_shard_from_object_store(self, monkeypatch):
+        index = {
+            "weight_map": {
+                "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
+                "lm_head.weight": "model-00002-of-00002.safetensors",
+                "mtp.fc.weight": "model-mtp.safetensors",
+                "mtp.layers.0.input_layernorm.weight": "model-mtp.safetensors",
+            }
+        }
+        _stub_runai(monkeypatch, {_SAFETENSORS_INDEX_NAME: json.dumps(index)})
+        files = [
+            f"s3://bucket/model/{name}"
+            for name in (
+                "model-00001-of-00002.safetensors",
+                "model-00002-of-00002.safetensors",
+                "model-mtp.safetensors",
+            )
+        ]
+        assert _select_draft_weight_files("s3://bucket/model", files) == (
+            DraftShardSelection.SELECTED,
+            ["s3://bucket/model/model-mtp.safetensors"],
+        )

modelexpress_client/python/tests/test_vllm_loader.py

@@ -1837,3 +1837,78 @@ def test_falls_back_to_vllm_level_when_env_unset(self):
             assert mx_root.level == logging.DEBUG
         finally:
             self._cleanup(vllm_logger)
+
+
+# ---------------------------------------------------------------------------
+# Compilation-state coherence across re-init (MTP target/drafter co-ownership)
+# ---------------------------------------------------------------------------
+
+
+class _RegisteredLayer(nn.Module):
+    """Stand-in for a vLLM layer that records the prefix it registered under."""
+
+    def __init__(self, layer_name):
+        super().__init__()
+        self.layer_name = layer_name
+
+
+def _make_compilation_config():
+    from collections import Counter
+    from types import SimpleNamespace
+
+    return SimpleNamespace(
+        static_forward_context={},
+        static_all_moe_layers=[],
+        enabled_custom_ops=Counter(),
+        compilation_time=0.0,
+    )
+
+
+def _initialize_model(cc, prefix):
+    """Build a model under `prefix`, registering its layers the way vLLM does."""
+    from types import SimpleNamespace
+    from modelexpress.engines.vllm.adapter import VllmAdapter
+
+    model = nn.Module()
+    model.add_module("self_attn", _RegisteredLayer(f"{prefix}.layers.0.self_attn"))
+    model.add_module("mlp", _RegisteredLayer(f"{prefix}.layers.0.mlp"))
+    assert model.self_attn.layer_name not in cc.static_forward_context  # vLLM's check
+    cc.static_forward_context[model.self_attn.layer_name] = model.self_attn
+    cc.static_all_moe_layers.append(model.mlp.layer_name)
+    cc.enabled_custom_ops["rms_norm"] += 1
+
+    adapter = object.__new__(VllmAdapter)  # __init__ touches devices
+    adapter.vllm_config = SimpleNamespace(compilation_config=cc)
+    return model, adapter
+
+
+def test_unregister_leaves_co_owned_target_registrations():
+    """MTP: unregistering the drafter must not drop the live target's layers."""
+    cc = _make_compilation_config()
+    target, _ = _initialize_model(cc, "language_model.model")
+    drafter, adapter = _initialize_model(cc, "mtp")
+
+    adapter._unregister_model_layers(drafter)
+
+    assert cc.static_forward_context == {
+        "language_model.model.layers.0.self_attn": target.self_attn
+    }
+    assert cc.static_all_moe_layers == ["language_model.model.layers.0.mlp"]
+    assert cc.enabled_custom_ops["rms_norm"] == 2  # accumulating field untouched
+    _initialize_model(cc, "mtp")  # the drafter's rebuild re-registers cleanly
+
+
+def test_unregister_releases_the_discarded_model():
+    """The registries must not pin the stale model across the rebuild."""
+    import gc
+    import weakref
+
+    cc = _make_compilation_config()
+    stale, adapter = _initialize_model(cc, "language_model.model")
+    layer_ref = weakref.ref(stale.self_attn)
+
+    adapter._unregister_model_layers(stale)
+    del stale
+    gc.collect()
+
+    assert layer_ref() is None