Address 10-reviewer P1 findings: vision cache split, PEFT offline, retry OOM

Split the Studio vision-detection cache by local_files_only, mirroring the audio
cache fix. is_vision_model / _is_vision_model_uncached / _raw_config_has_vision_config
/ load_model_config now thread local_files_only, the cache key includes it, and the
exporter passes it. Offline detection also skips the transformers-5 network
subprocess and stays on the local cache, so an offline negative can no longer be
keyed under the online entry and poison a later online probe. Adds a regression
test mirroring the audio poison test.

Forward local_files_only to both PeftModel.from_pretrained adapter-attach sites in
loader.py so a cached remote LoRA adapter resolves from the local cache under
explicit local-only / offline loads (defence-in-depth alongside the forced-offline
window).

_offline_aware_load: run the forced-offline retry OUTSIDE the except block and
collect + empty the device cache first. An except-scoped exception keeps its
__traceback__, which pins the failed attempt's frame locals (a partially loaded
model) until the block exits; loading the model again while that copy is still
alive could OOM a large VLM. Letting the except block close drops the traceback so
the partial load is freed before the retry reallocates.

studio/backend/core/export/export.py

@@ -280,7 +280,9 @@ class ExportBackend:
                 self._audio_type = detect_audio_type(
                     model_id, hf_token = token, local_files_only = local_files_only
                 )
-                self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token)
+                self.is_vision = not self._audio_type and is_vision_model(
+                    model_id, hf_token = token, local_files_only = local_files_only
+                )
 
             if self._audio_type == "csm":
                 from unsloth import FastModel

studio/backend/tests/test_vision_cache.py

@@ -71,7 +71,7 @@ class TestVisionCacheHitMiss:
         """Two calls for the same model invoke the uncached fn once."""
         assert is_vision_model("org/my-vlm") is True
         assert is_vision_model("org/my-vlm") is True
-        mock_uncached.assert_called_once_with("org/my-vlm", None)
+        mock_uncached.assert_called_once_with("org/my-vlm", None, local_files_only = False)
 
     @patch("utils.models.model_config._is_vision_model_uncached", return_value = False)
     def test_different_models_each_detected(self, mock_uncached):
@@ -97,7 +97,7 @@ class TestVisionCacheStoresFalse:
         assert is_vision_model("org/text-only") is False
         assert is_vision_model("org/text-only") is False
         mock_uncached.assert_called_once()
-        assert _vision_detection_cache[("org/text-only", None)] is False
+        assert _vision_detection_cache[("org/text-only", None, False)] is False
 
 
 # Subprocess path (transformers 5.x) caching
@@ -120,7 +120,7 @@ class TestVisionCacheSubprocessPath:
         assert is_vision_model("unsloth/Qwen3.5-2B") is True
 
         mock_subprocess.assert_called_once()
-        assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None)] is True
+        assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None, False)] is True
 
     @patch("utils.models.model_config._raw_config_has_vision_config", return_value = True)
     @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None)
@@ -133,7 +133,9 @@ class TestVisionCacheSubprocessPath:
         assert is_vision_model("unsloth/gemma-4-E4B-it") is True
         assert is_vision_model("unsloth/gemma-4-E4B-it") is True
 
-        mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None)
+        mock_raw_config.assert_called_once_with(
+            "unsloth/gemma-4-E4B-it", hf_token = None, local_files_only = False
+        )
         mock_subprocess.assert_not_called()
 
 
@@ -405,6 +407,40 @@ class TestVisionCacheTokenHandling:
         mock_uncached.assert_called_once()
 
 
+class TestVisionCacheLocalOnly:
+    """local_files_only is part of the cache key (mirrors the audio cache). An
+    offline probe only sees the on-disk cache, so its negative result must never
+    be reused by a later online probe that can reach the Hub -- otherwise a VLM
+    is routed through the text loader until restart."""
+
+    def test_local_only_negative_does_not_poison_online(self, monkeypatch):
+        import utils.models.model_config as mc
+
+        mc._vision_detection_cache.clear()
+        monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False)
+        monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n)
+
+        seen = []
+
+        def _probe(name, hf_token = None, local_files_only = False):
+            seen.append(local_files_only)
+            # Offline: cannot run the transformers-5 subprocess / fetch -> not a VLM.
+            # Online: the remote config reveals a VLM.
+            return False if local_files_only else True
+
+        monkeypatch.setattr(mc, "_is_vision_model_uncached", _probe)
+
+        # Offline probe caches False under a local-only key.
+        assert mc.is_vision_model("some/vlm", local_files_only = True) is False
+        # A later online probe must re-run (different key) and detect the VLM.
+        assert mc.is_vision_model("some/vlm", local_files_only = False) is True
+        assert seen == [True, False]
+        # The online positive is then cached for subsequent online callers.
+        assert mc.is_vision_model("some/vlm", local_files_only = False) is True
+        assert seen == [True, False]
+        mc._vision_detection_cache.clear()
+
+
 # ---------------------------------------------------------------------------
 # Direct unit tests for _raw_config_has_vision_config
 # ---------------------------------------------------------------------------

studio/backend/utils/models/model_config.py

@@ -471,6 +471,7 @@ def load_model_config(
     use_auth: bool = False,
     token: Optional[str] = None,
     trust_remote_code: bool = False,
+    local_files_only: bool = False,
 ):
     """Load model config with optional authentication control.
 
@@ -478,12 +479,18 @@ def load_model_config(
     metadata lookups must never execute a model repo's ``auto_map`` Python.
     Deliberate remote-code loads pass the flag explicitly through
     ``FastLanguageModel.from_pretrained`` with the user's own consent.
+
+    ``local_files_only`` keeps the config read on the local HF cache (offline
+    export), so an offline probe never blocks on the network.
     """
     from transformers import AutoConfig
 
     if token:
         return AutoConfig.from_pretrained(
-            model_name, trust_remote_code = trust_remote_code, token = token
+            model_name,
+            trust_remote_code = trust_remote_code,
+            token = token,
+            local_files_only = local_files_only,
         )
 
     if not use_auth:
@@ -493,12 +500,14 @@ def load_model_config(
                 model_name,
                 trust_remote_code = trust_remote_code,
                 token = None,
+                local_files_only = local_files_only,
             )
 
     # Default auth (cached tokens)
     return AutoConfig.from_pretrained(
         model_name,
         trust_remote_code = trust_remote_code,
+        local_files_only = local_files_only,
     )
 
 
@@ -598,7 +607,7 @@ def _is_vlm(config) -> bool:
 
 
 def _raw_config_has_vision_config(
-    model_name: str, hf_token: Optional[str] = None
+    model_name: str, hf_token: Optional[str] = None, local_files_only: bool = False
 ) -> Optional[bool]:
     try:
         if is_local_path(model_name):
@@ -610,6 +619,7 @@ def _raw_config_has_vision_config(
                     repo_id = model_name,
                     filename = "config.json",
                     token = hf_token,
+                    local_files_only = local_files_only,
                 )
             )
         config = json.loads(config_path.read_text())
@@ -780,22 +790,29 @@ def _token_fingerprint(token: Optional[str]) -> Optional[str]:
 # Keyed by (normalized_model_name, token_fingerprint) to handle gated models.
 # Only definitive results are cached; transient failures (network, timeouts)
 # are NOT cached so they can be retried.
-_vision_detection_cache: Dict[Tuple[str, Optional[str]], bool] = {}
+_vision_detection_cache: Dict[Tuple[str, Optional[str], bool], bool] = {}
 _vision_cache_lock = threading.Lock()
 
 
-def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
+def is_vision_model(
+    model_name: str, hf_token: Optional[str] = None, local_files_only: bool = False
+) -> bool:
     """
     Detect vision-language models (VLMs) via architecture in config. Works for
     fine-tuned models since they inherit the base architecture.
 
     Models needing transformers 5.x are checked in a .venv_t5/ subprocess.
-    Results are cached per (model_name, token_fingerprint) for the process
-    lifetime; transient failures are not cached so they can be retried.
+    Results are cached per (model_name, token_fingerprint, local_files_only) for
+    the process lifetime; transient failures are not cached so they can be
+    retried. local_files_only is part of the key because an offline probe reads
+    only the on-disk cache and can differ from an online probe (e.g. a
+    transformers-5 VLM the offline path cannot run the subprocess for), so the
+    two results must never share a cache entry.
 
     Args:
         model_name: Model identifier (HF repo or local path)
         hf_token: Optional HF token for gated/private models
+        local_files_only: Keep detection on the local HF cache (offline export)
     """
     # Local GGUF models are served by llama-server. Their multimodal
     # capability comes from a companion mmproj, not a Transformers config.
@@ -829,7 +846,7 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
             exc,
         )
         resolved_name = model_name
-    cache_key = (resolved_name, _token_fingerprint(hf_token))
+    cache_key = (resolved_name, _token_fingerprint(hf_token), bool(local_files_only))
 
     # Lock-free fast path for cache hits. Sentinel distinguishes "key not found"
     # from "value is False" in a single atomic dict.get() call.
@@ -840,7 +857,9 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
 
     # Compute outside the lock so long-running detection isn't serialized across
     # models. Two concurrent calls may both run, but produce the same result.
-    result = _is_vision_model_uncached(resolved_name, hf_token)
+    result = _is_vision_model_uncached(
+        resolved_name, hf_token, local_files_only = local_files_only
+    )
     # Only cache definitive results; None is a transient failure, retry later.
     if result is not None:
         with _vision_cache_lock:
@@ -849,7 +868,9 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
     return False
 
 
-def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
+def _is_vision_model_uncached(
+    model_name: str, hf_token: Optional[str] = None, local_files_only: bool = False
+) -> Optional[bool]:
     """Uncached vision detection; use is_vision_model() instead.
 
     Returns True/False for definitive results, or None on transient errors
@@ -858,15 +879,19 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -
     # Try the raw-config reader FIRST (code-free, version-independent): it classifies
     # repo-code VLMs like DeepSeek-OCR via declarative vision_config with no remote-code
     # execution or transformers-5.x subprocess.
-    raw = _raw_config_has_vision_config(model_name, hf_token = hf_token)
+    raw = _raw_config_has_vision_config(
+        model_name, hf_token = hf_token, local_files_only = local_files_only
+    )
     if raw is not None:
         return raw
 
     # Raw read failed transiently: fall back to AutoConfig with remote code DISABLED
     # (in a transformers-5.x subprocess when the main process can't parse the arch).
+    # Skip the subprocess offline: it does its own network probe and would diverge
+    # from the online path, so offline we stay on the local cache via load_model_config.
     from utils.transformers_version import needs_transformers_5
 
-    if needs_transformers_5(model_name):
+    if not local_files_only and needs_transformers_5(model_name):
         logger.info(
             "Model '%s' needs transformers 5.x -- checking vision via subprocess",
             model_name,
@@ -874,7 +899,12 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -
         return _is_vision_model_subprocess(model_name, hf_token = hf_token)
 
     try:
-        config = load_model_config(model_name, use_auth = True, token = hf_token)
+        config = load_model_config(
+            model_name,
+            use_auth = True,
+            token = hf_token,
+            local_files_only = local_files_only,
+        )
 
         if _is_vlm(config):
             model_type = getattr(config, "model_type", None)

unsloth/models/loader.py

@@ -860,6 +860,7 @@ class FastLanguageModel(FastLlamaModel):
                 old_model_name,
                 token = token,
                 revision = revision,
+                local_files_only = local_files_only,
                 is_trainable = True,
                 trust_remote_code = trust_remote_code,
             )
@@ -1767,6 +1768,7 @@ class FastModel(FastBaseModel):
                     old_model_name,
                     token = token,
                     revision = revision,
+                    local_files_only = local_files_only,
                     is_trainable = True,
                     trust_remote_code = trust_remote_code,
                 )

unsloth/models/loader_utils.py

@@ -737,12 +737,25 @@ def _offline_aware_load(fn):
         try:
             return fn(*args, **kwargs)
         except Exception as e:
-            if _is_offline_related_error(e):
-                kwargs["local_files_only"] = True
-                with _force_hf_offline():
-                    return fn(*args, **kwargs)
-            raise
-
+            if not _is_offline_related_error(e):
+                raise
+        # Retry OUTSIDE the except block: an except-scoped exception still holds its
+        # __traceback__, which pins the failed attempt's frame locals (e.g. a
+        # partially loaded model) until the block exits. Loading the model a second
+        # time while the first copy is still alive can OOM a large VLM. Letting the
+        # except block close drops the traceback; collect + empty the device cache
+        # so the partial load is freed before the forced-offline retry reallocates.
+        try:
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            if hasattr(torch, "xpu") and torch.xpu.is_available():
+                torch.xpu.empty_cache()
+        except Exception:
+            pass
+        kwargs["local_files_only"] = True
+        with _force_hf_offline():
+            return fn(*args, **kwargs)
     return _wrapper