vLLM loader: use safetensors_load_strategy="eager" on network volumes

The previous approach patched safetensors.torch.load_file, but vLLM uses
safe_open (mmap) via safetensors_weights_iterator(). vLLM v0.17+ already
has a built-in "eager" strategy that does open().read() — we just set
load_config.safetensors_load_strategy="eager" on network FUSE volumes.

Benchmarking on RunPod (RTX 4090, MFS FUSE, Qwen2.5-1.5B) shows eager
is actually slower with warm page cache (2.5s vs 0.7s for mmap). The
eager path helps on cold reads from slow network FSes where mmap page
faults trigger expensive round-trips. Removed overlay from slow_fs
detection since overlay-backed storage has fast mmap.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: lightsofapollo

python/zerostart/integrations/vllm.py

@@ -3,10 +3,9 @@
 Provides a custom model loader that subclasses vLLM's DefaultModelLoader
 and runs inside vLLM's EngineCore subprocess where weights are actually loaded.
 
-Key optimizations:
-1. Network volume fix: eager read instead of mmap on FUSE/NFS (30-50x faster)
-2. Patched safe_open: detect network volumes and use fast path
-3. Auto-registered via vLLM's plugin system (entry_points)
+Key optimization:
+  Network volume fix: sets safetensors_load_strategy="eager" on FUSE/NFS
+  volumes where mmap is 30-50x slower than sequential read.
 
 Usage:
     # Option 1: Auto-registration via entry_points (pip install zerostart)
@@ -16,9 +15,6 @@
     from zerostart.integrations.vllm import register
     register()
     # Then: --load-format zerostart
-
-    # Option 3: Transparent hook (patches from_pretrained in parent process)
-    zerostart run --accelerate -p vllm -- python -m vllm.entrypoints.openai.api_server ...
 """
 
 from __future__ import annotations
@@ -31,9 +27,6 @@
 log = logging.getLogger("zerostart.vllm")
 
 if TYPE_CHECKING:
-    from collections.abc import Generator, Iterator
-
-    import torch
     import torch.nn as nn
     from vllm.config import ModelConfig
     from vllm.config.load import LoadConfig
@@ -53,7 +46,6 @@ def register() -> None:
         register_model_loader("zerostart")(ZerostartModelLoader)
         log.info("Registered zerostart model loader with vLLM")
     except ImportError:
-        # Fallback for older vLLM versions
         try:
             import vllm.model_executor.model_loader as ml
             registry = getattr(ml, "_LOAD_FORMAT_TO_MODEL_LOADER", None)
@@ -118,9 +110,12 @@ def _is_network_volume(path: str) -> bool:
         return _network_volume_cache[path]
 
     result = False
+    # Only truly network-backed filesystems where mmap page faults
+    # trigger network round-trips. Overlay is excluded because it's
+    # backed by local storage and mmap works fine there.
     slow_fs = frozenset({
         "fuse", "fuse.juicefs", "fuse.gcsfuse", "fuse.sshfs",
-        "nfs", "nfs4", "cifs", "smbfs", "9p", "overlay",
+        "nfs", "nfs4", "cifs", "smbfs", "9p",
     })
 
     try:
@@ -144,105 +139,56 @@ def _is_network_volume(path: str) -> bool:
     return result
 
 
-# ---------------------------------------------------------------------------
-# Fast weight iterator — replaces safetensors mmap with eager read on
-# network volumes, and patches safe_open for the same
-# ---------------------------------------------------------------------------
-
-def _fast_safetensors_weights_iterator(
-    hf_weights_files: list[str],
-) -> Generator[tuple[str, torch.Tensor], None, None]:
-    """Yield (name, tensor) pairs from safetensors files.
-
-    On network volumes: reads entire file into memory first (eager),
-    avoiding the 30-50x mmap penalty on FUSE/NFS.
-    On local NVMe: uses standard safe_open (mmap is fast).
-    """
-    import safetensors.torch
-
-    for st_file in hf_weights_files:
-        t0 = time.monotonic()
-
-        if _is_network_volume(st_file):
-            # Eager read: load entire file to avoid mmap page fault penalty
-            with open(st_file, "rb") as f:
-                data = f.read()
-            tensors = safetensors.torch.load(data)
-            elapsed = time.monotonic() - t0
-            log.info(
-                "Eager read %s (%.2fs, %d tensors, %.0f MB)",
-                Path(st_file).name, elapsed, len(tensors),
-                len(data) / 1e6,
-            )
-            yield from tensors.items()
-        else:
-            # Local NVMe: mmap is fast, use standard safe_open
-            from safetensors import safe_open
-            with safe_open(st_file, framework="pt") as f:
-                for name in f.keys():
-                    yield name, f.get_tensor(name)
-
-
 # ---------------------------------------------------------------------------
 # ZerostartModelLoader
 # ---------------------------------------------------------------------------
 
 class ZerostartModelLoader(_DefaultLoader):  # type: ignore[misc]
     """vLLM model loader with network volume acceleration.
 
-    Subclasses DefaultModelLoader and overrides the weight iteration
-    to use eager read on FUSE/NFS volumes. This runs INSIDE vLLM's
-    EngineCore subprocess where weights are actually loaded.
+    Subclasses DefaultModelLoader. On FUSE/NFS network volumes, sets
+    safetensors_load_strategy="eager" so vLLM reads entire files into
+    memory instead of using mmap (which is 30-50x slower on these FSes).
 
-    Key difference from transparent accelerate() hook:
-    - accelerate() patches from_pretrained in the parent process
-    - This loader patches weight loading in the EngineCore subprocess
-    - vLLM loads weights via safe_open, not from_pretrained
+    On local NVMe, delegates entirely to DefaultModelLoader (mmap is fast).
     """
 
     def __init__(self, load_config: LoadConfig):
-        # Rewrite load_format to "safetensors" BEFORE super().__init__
-        # so DefaultModelLoader._prepare_weights() doesn't reject "zerostart".
-        # We store the original to know we were invoked as zerostart.
+        import os
+
+        # Rewrite load_format from "zerostart" to "safetensors" so
+        # DefaultModelLoader._prepare_weights() doesn't reject it.
         self._zerostart_requested = getattr(load_config, "load_format", None) == "zerostart"
         if self._zerostart_requested:
             load_config.load_format = "safetensors"
 
-        if _DefaultLoader is not object:
-            super().__init__(load_config)
-        else:
-            self.load_config = load_config
-
-        # Detect if we're on a network volume
-        self._on_network_volume = any(
+        # Switch to eager loading if explicitly requested or on a network FS
+        # where mmap page faults trigger expensive network round-trips.
+        #
+        # Note: on FUSE mounts with warm page cache (e.g. RunPod MFS), mmap
+        # is actually faster than eager because it avoids copying data.
+        # Eager only helps on cold reads from slow network FSes (NFS, JuiceFS).
+        # Use ZEROSTART_EAGER=1 to force eager loading.
+        force_eager = os.environ.get("ZEROSTART_EAGER", "").lower() in ("1", "true")
+        on_network_volume = any(
             _is_network_volume(p)
             for p in ["/volume", "/gpu-cli-workspaces", "/workspace"]
             if Path(p).exists()
         )
 
-        if self._on_network_volume:
-            log.info("Network volume detected — using eager read for safetensors")
-            self._patch_safe_open()
-
-    def _patch_safe_open(self) -> None:
-        """Patch safetensors in this subprocess for eager read on network volumes."""
-        try:
-            import safetensors.torch as st
-
-            original_load_file = st.load_file
-
-            def patched_load_file(filename: str, device: str = "cpu") -> dict[str, Any]:
-                if _is_network_volume(str(filename)):
-                    with open(filename, "rb") as f:
-                        data = f.read()
-                    return st.load(data, device=device)
-                return original_load_file(filename, device=device)
+        if force_eager or on_network_volume:
+            current = getattr(load_config, "safetensors_load_strategy", "lazy")
+            if current != "eager":
+                load_config.safetensors_load_strategy = "eager"
+                reason = "ZEROSTART_EAGER=1" if force_eager else "network volume detected"
+                log.info(
+                    "Switched safetensors_load_strategy to 'eager' (%s)", reason,
+                )
 
-            st.load_file = patched_load_file
-            self._original_load_file = original_load_file
-            log.debug("Patched safetensors.torch.load_file in subprocess")
-        except ImportError:
-            pass
+        if _DefaultLoader is not object:
+            super().__init__(load_config)
+        else:
+            self.load_config = load_config
 
     def download_model(self, model_config: ModelConfig) -> None:
         """Download model via HF hub (standard path)."""
@@ -259,26 +205,18 @@ def download_model(self, model_config: ModelConfig) -> None:
                 log.warning("HF download failed: %s", e)
 
     def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
-        """Load weights with network volume optimization.
-
-        On network volumes: uses eager read (30-50x faster than mmap).
-        On local NVMe: delegates to DefaultModelLoader (mmap is fast).
-        """
+        """Load weights, delegating to DefaultModelLoader."""
         t0 = time.monotonic()
 
         if _DefaultLoader is not object and hasattr(super(), "load_weights"):
-            # Let DefaultModelLoader handle it — our safe_open patch
-            # is already installed and will intercept the reads
             super().load_weights(model, model_config)
         else:
             log.warning("DefaultModelLoader not available — basic weight loading")
             self._fallback_load_weights(model, model_config)
 
         elapsed = time.monotonic() - t0
-        log.info(
-            "Weight loading complete (%.2fs, network_volume=%s)",
-            elapsed, self._on_network_volume,
-        )
+        strategy = getattr(self.load_config, "safetensors_load_strategy", "unknown")
+        log.info("Weight loading complete (%.2fs, strategy=%s)", elapsed, strategy)
 
     def _fallback_load_weights(
         self, model: nn.Module, model_config: ModelConfig,
@@ -288,10 +226,13 @@ def _fallback_load_weights(
 
         model_path = Path(model_config.model)
         if not model_path.is_dir():
-            from zerostart.snapshot import _find_hf_cache_dir
-            cache_dir = _find_hf_cache_dir(model_config.model)
-            if cache_dir:
-                model_path = cache_dir
+            try:
+                from zerostart.snapshot import _find_hf_cache_dir
+                cache_dir = _find_hf_cache_dir(model_config.model)
+                if cache_dir:
+                    model_path = cache_dir
+            except ImportError:
+                pass
 
         sf_files = sorted(model_path.glob("*.safetensors"))
         if not sf_files:

tests/test_vllm_integration.sh

@@ -162,7 +162,7 @@ r = subprocess.run(
 elapsed = time.monotonic() - t0
 print(r.stdout.strip())
 if r.returncode != 0:
-    print("STDERR:", r.stderr[-1000:])
+    print("STDERR:", r.stderr[-1500:])
 print(f"Wall clock: {elapsed:.2f}s")
 PYEOF