fix(models,engine): persist solo per-load overrides across restart; kv fp8 off for multimodal

ainode/core/config.py

@@ -51,6 +51,11 @@ class NodeConfig:
     # context (32k+) or vLLM OOMs sizing the cache at bf16 (see engine/AGENTS.md).
     # Set "" / "auto" to let vLLM choose if a model/quant ever rejects fp8.
     kv_cache_dtype: str = "fp8"
+    # Provenance of kv_cache_dtype: True only when a caller EXPLICITLY supplied it
+    # (per-load body or config). The multimodal fp8→auto safety downgrade
+    # (engine/backends/nvidia.py) fires only on the DEFAULT fp8 — an explicit
+    # fp8 request on a VLM is honored, giving the user a way to opt back in.
+    kv_cache_dtype_explicit: bool = False
     quantization: Optional[str] = None  # awq, gptq, fp8, None
     trust_remote_code: bool = False
 

ainode/engine/backends/nvidia.py

@@ -615,6 +615,52 @@ def _local_model_dir(self) -> Optional[str]:
             pass
         return None
 
+    def _is_multimodal_model(self) -> bool:
+        """Detect a vision/multimodal model from its on-disk ``config.json``.
+
+        True when config.json has a ``vision_config`` key, or any
+        ``architectures`` entry matches ``/VL|Vision|vision/``. Only the LOCAL
+        model dir is inspected: if config.json is unreadable (e.g. a remote
+        repo-id not yet downloaded), we return False so the caller keeps the fp8
+        default — we deliberately never fetch remote config here (no network in
+        the serve-args builder).
+
+        CEILING: a not-yet-downloaded VLM served by repo-id won't be detected
+        and will use the fp8 KV default until its weights are on disk.
+        """
+        local = self._local_model_dir()
+        if not local:
+            return False
+        try:
+            cfg = json.loads((Path(local) / "config.json").read_text())
+        except Exception:
+            return False
+        if "vision_config" in cfg:
+            return True
+        import re
+        archs = cfg.get("architectures") or []
+        if isinstance(archs, str):
+            archs = [archs]
+        return any(re.search(r"VL|Vision|vision", str(a)) for a in archs)
+
+    def _effective_kv_cache_dtype(self) -> str:
+        """Resolve the KV-cache dtype for serve args.
+
+        fp8 KV corrupts vision-model generation on GB10 (verified: Qwen2.5-VL
+        emits garbage on fp8, clean output on auto; text models are unaffected).
+        So the fp8 DEFAULT is downgraded to 'auto' for a multimodal model. Any
+        EXPLICIT ``kv_cache_dtype`` always wins — whether it's a non-fp8 value
+        (which isn't the default anyway) or an explicit 'fp8' flagged via
+        ``kv_cache_dtype_explicit`` (the user's opt-back-in for a VLM/vLLM combo
+        they know handles fp8 KV). A model whose config.json can't be read keeps
+        the fp8 default.
+        """
+        dtype = getattr(self.config, "kv_cache_dtype", "") or ""
+        explicit = getattr(self.config, "kv_cache_dtype_explicit", False)
+        if dtype == "fp8" and not explicit and self._is_multimodal_model():
+            return "auto"
+        return dtype
+
     def _is_nvfp4_model(self) -> bool:
         """Detect NVFP4 from the on-disk config.json quantization metadata (with
         the model id as a fallback) so the MARLIN serve env is applied only when
@@ -901,8 +947,11 @@ def _build_vllm_serve_args(self, tp_size: int) -> List[str]:
         # fp8 KV cache — the GB10 design default (engine/AGENTS.md): required for
         # long context or vLLM OOMs sizing the cache at bf16. Config-driven so a
         # model/quant that rejects fp8 can fall back via kv_cache_dtype="".
-        if getattr(self.config, "kv_cache_dtype", ""):
-            args.extend(["--kv-cache-dtype", self.config.kv_cache_dtype])
+        # _effective_kv_cache_dtype downgrades the fp8 DEFAULT to auto for
+        # multimodal models (fp8 corrupts VLM generation on GB10).
+        kv_dtype = self._effective_kv_cache_dtype()
+        if kv_dtype:
+            args.extend(["--kv-cache-dtype", kv_dtype])
         if tp_size > 1:
             args.extend(["--tensor-parallel-size", str(tp_size)])
             args.extend(["--distributed-executor-backend", "ray"])

ainode/models/api_routes.py

@@ -143,7 +143,51 @@ def load_instance_manifest() -> list:
 
 
 _OVERRIDE_KEYS = ("served_model_name", "max_model_len", "kv_cache_dtype",
-                  "quantization", "trust_remote_code")
+                  "kv_cache_dtype_explicit", "quantization", "trust_remote_code")
+
+
+def _resolved_overrides(gmu, overrides) -> dict:
+    """Resolve the FULL per-load override set to concrete values, defaulting
+    every field the caller did NOT supply to its NodeConfig class default.
+
+    Returned as a field→value dict covering ``gpu_memory_utilization`` and every
+    ``_OVERRIDE_KEYS`` entry. This is the single source of truth applied to BOTH
+    the per-instance launch config (``inst_config``) and the persisted primary
+    ``NodeConfig`` so the live backend and config.json can never diverge. Absent
+    fields RESET to their default rather than inheriting the previous load's
+    value — critical because ``inst_config`` is built from the SHARED, mutable
+    ``app["config"]`` (which still carries the prior primary load's overrides),
+    so without an explicit reset here, loading model B after model A would leak
+    A's kv_cache_dtype/quantization/served_model_name/etc. onto B (a bare
+    ``{"model": ...}`` load would silently inherit A's --quantization,
+    --trust-remote-code, and --served-model-name).
+    """
+    from ainode.core.config import NodeConfig
+    defaults = NodeConfig()
+    ov = overrides or {}
+    resolved = {
+        "gpu_memory_utilization":
+            gmu if gmu is not None else defaults.gpu_memory_utilization,
+    }
+    for k in _OVERRIDE_KEYS:
+        resolved[k] = ov[k] if k in ov else getattr(defaults, k)
+    return resolved
+
+
+def _persist_primary_overrides(config, gmu, overrides) -> None:
+    """Persist per-load overrides onto the SHARED NodeConfig for the primary.
+
+    The primary solo model boots from ``NodeConfig`` (config.json) after a
+    `systemctl restart` — NOT from the stacked-instance manifest — so every
+    per-load override (kv_cache_dtype, max_model_len, served_model_name,
+    trust_remote_code, quantization, gpu_memory_utilization) must be written
+    here or the boot engine serves the model with stale/default values (the
+    live VLM-came-back-on-fp8 bug). Uses the SAME resolved set as ``inst_config``
+    so the persisted config always matches the live backend just launched.
+    Caller saves config.
+    """
+    for k, v in _resolved_overrides(gmu, overrides).items():
+        setattr(config, k, v)
 
 
 def append_solo_instance(app, model: str, gmu=None, *, overrides=None, persist: bool = True) -> dict:
@@ -181,15 +225,16 @@ def append_solo_instance(app, model: str, gmu=None, *, overrides=None, persist:
     name_token = "" if port == config.api_port else str(port)  # primary keeps legacy names
     instance_id = f"{config.node_id or 'head'}:{model}"
 
+    # Build the launch config from the RESOLVED override set. app["config"] is a
+    # SHARED, mutable object that still carries the PREVIOUS primary's per-load
+    # overrides, so `replace(config, ...)` alone would leak A's kv_cache_dtype /
+    # quantization / served_model_name / trust_remote_code onto the next model B
+    # (a bare {"model": ...} load). _resolved_overrides resets every unsupplied
+    # field to its NodeConfig default, and is the SAME set persisted below, so the
+    # live backend and config.json can never diverge.
     inst_config = replace(config, model=model, distributed_mode="solo",
-                          peer_ips=[], api_port=port)
-    if gmu is not None:
-        inst_config = replace(inst_config, gpu_memory_utilization=gmu)
-    if overrides:
-        allowed = {k: v for k, v in overrides.items()
-                   if k in _OVERRIDE_KEYS and v is not None}
-        if allowed:
-            inst_config = replace(inst_config, **allowed)
+                          peer_ips=[], api_port=port,
+                          **_resolved_overrides(gmu, overrides))
 
     def _clear():
         # routing-truth: a failed primary launch must stop the node advertising a
@@ -220,8 +265,7 @@ def _clear():
         config.model = model
         config.distributed_mode = "solo"
         config.peer_ips = []
-        if gmu is not None:
-            config.gpu_memory_utilization = gmu
+        _persist_primary_overrides(config, gmu, overrides)
         try:
             config.save()
         except Exception:
@@ -231,6 +275,7 @@ def _clear():
         # Reloaded the primary while a stack exists: keep app["engine"] on the live
         # backend (not the stopped old one) so status/proxy don't dangle.
         config.model = model
+        _persist_primary_overrides(config, gmu, overrides)
         try:
             config.save()
         except Exception:
@@ -387,6 +432,11 @@ async def handle_model_load(request: web.Request) -> web.Response:
     for k in ("kv_cache_dtype", "quantization"):
         if body.get(k) is not None:
             overrides[k] = body[k]
+    if "kv_cache_dtype" in overrides:
+        # Mark provenance so the multimodal fp8→auto safety downgrade
+        # (nvidia.py _effective_kv_cache_dtype) is skipped: an EXPLICIT fp8 KV
+        # request on a VLM is honored, giving the user a way to opt back in.
+        overrides["kv_cache_dtype_explicit"] = True
     if body.get("trust_remote_code") is not None:
         overrides["trust_remote_code"] = bool(body["trust_remote_code"])
 

ainode/training/api_routes.py

@@ -633,6 +633,13 @@ def _blocking() -> dict:
                     res = meta_optimize(cfg, target_yield=target_yield, max_rounds=max_rounds)
                     report = {"kept": len(res["dataset"]), "best_yield": res["best_yield"],
                               "rounds": len(res["rounds"])}
+                    # valset objective: surface the held-out lift + active score the
+                    # meta loop already computed, so the run report reflects what the
+                    # optimizer actually maximized (not just the Δ=1 yield proxy).
+                    if res.get("objective") == "valset":
+                        report["best_lift"] = res.get("best_lift")
+                        report["best_score"] = res.get("best_score")
+                        report["best_significant"] = res.get("best_significant")
                     return {"report": report, "rounds": res["rounds"]}
                 from ainode.training.autodata.core import run as _run
                 res = _run(cfg)

tests/test_autodata.py

@@ -367,6 +367,44 @@ async def test_autodata_route_closes_the_loop(client, monkeypatch):
     assert job_resp.status == 201, await job_resp.text()
 
 
+@pytest.mark.asyncio
+async def test_autodata_meta_report_surfaces_best_lift_for_valset(client, monkeypatch):
+    """A meta run with objective=valset surfaces best_lift + best_score (what the
+    optimizer actually maximized) in the report — not just the Δ=1 yield proxy."""
+    c, _ = client
+
+    def _fake_meta(cfg, target_yield=30, max_rounds=4):
+        return {"best_prompt": "p", "best_yield": 42.0, "best_score": 0.25,
+                "best_lift": 0.25, "best_significant": True, "dataset": [{"a": 1}],
+                "rounds": [{}, {}], "objective": "valset", "out": ""}
+
+    monkeypatch.setattr("ainode.training.autodata.meta.meta_optimize", _fake_meta)
+
+    cfg = {
+        "task_spec": "x", "gen_prompt": "x", "n_tasks": 3, "concurrency": 1,
+        "challenger": {"url": "u", "model": "challenger"},
+        "weak": {"url": "u", "model": "weak"},
+        "strong": {"url": "u", "model": "strong"},
+        "judge": {"url": "u", "model": "judge"},
+    }
+    resp = await c.post("/api/training/autodata", json={"config": cfg, "meta": True})
+    assert resp.status == 202
+    run_id = (await resp.json())["run_id"]
+
+    data = None
+    for _ in range(200):
+        data = await (await c.get(f"/api/training/autodata/{run_id}")).json()
+        if data["status"] in ("completed", "failed"):
+            break
+        await asyncio.sleep(0.02)
+    assert data["status"] == "completed", data
+    report = data["report"]
+    assert report["best_lift"] == 0.25
+    assert report["best_score"] == 0.25
+    assert report["best_significant"] is True
+    assert report["best_yield"] == 42.0
+
+
 @pytest.mark.asyncio
 async def test_autodata_route_rejects_bad_body(client):
     c, _ = client

tests/test_cluster_load.py

@@ -274,6 +274,132 @@ async def _noop_sleep(*a, **k):
     assert models == {"model-A", "model-B"}
 
 
+def test_solo_load_persists_and_resets_primary_overrides(monkeypatch):
+    """A solo (primary) load persists EVERY per-load override onto the shared
+    NodeConfig — not just model + gmu — so the boot engine re-applies them after
+    a restart (the VLM-came-back-on-fp8 bug). Reloading the primary WITHOUT an
+    override resets that field to its NodeConfig default (no stale inheritance)."""
+    import ainode.models.api_routes as mr
+
+    _patch_backend(monkeypatch)
+    cfg = NodeConfig(node_id="spark1", api_port=8000)
+    cfg.save = lambda: None
+    app = {"engine": None, "config": cfg, "cluster_state": ClusterState(),
+           "ray_autostart_state": None}
+
+    # Load a VLM with the full set of per-load overrides.
+    asyncio.run(mr.handle_model_load(_Req(app, {
+        "model": "Qwen/Qwen2.5-VL-7B-Instruct",
+        "gpu_memory_utilization": 0.6,
+        "max_model_len": 16384,
+        "kv_cache_dtype": "auto",
+        "served_model_name": "vl-alias",
+        "trust_remote_code": True,
+    })))
+    assert cfg.model == "Qwen/Qwen2.5-VL-7B-Instruct"
+    assert cfg.kv_cache_dtype == "auto"
+    assert cfg.max_model_len == 16384
+    assert cfg.served_model_name == ["vl-alias"]
+    assert cfg.trust_remote_code is True
+    assert cfg.gpu_memory_utilization == 0.6
+
+    # Reload the SAME primary with NO overrides → every field resets to default
+    # (kv back to fp8, ctx/alias cleared) rather than inheriting the prior load.
+    asyncio.run(mr.handle_model_load(_Req(app, {"model": "Qwen/Qwen2.5-VL-7B-Instruct"})))
+    defaults = NodeConfig()
+    assert cfg.kv_cache_dtype == defaults.kv_cache_dtype == "fp8"
+    assert cfg.max_model_len is None
+    assert cfg.served_model_name is None
+    assert cfg.trust_remote_code is False
+    assert cfg.gpu_memory_utilization == defaults.gpu_memory_utilization
+
+
+def test_next_model_launch_config_does_not_inherit_primary_overrides(monkeypatch):
+    """BLOCKER regression: loading model B after primary A must NOT leak A's
+    per-load overrides onto B's ACTUAL launch config. The shared app["config"]
+    still carries A's kv/quant/alias/trust after A's load, but a bare B load must
+    launch with clean NodeConfig defaults — else B (unquantized) gets A's
+    --quantization awq, unconsented --trust-remote-code, and A's alias."""
+    import ainode.models.api_routes as mr
+
+    made = _patch_backend(monkeypatch)
+    cfg = NodeConfig(node_id="spark1", api_port=8000)
+    cfg.save = lambda: None
+    app = {"engine": None, "config": cfg, "cluster_state": ClusterState(),
+           "ray_autostart_state": None}
+
+    # Primary A with the full override set.
+    asyncio.run(mr.handle_model_load(_Req(app, {
+        "model": "Qwen/Qwen2.5-VL-7B-Instruct",
+        "gpu_memory_utilization": 0.6,
+        "max_model_len": 32768,
+        "kv_cache_dtype": "auto",
+        "quantization": "awq",
+        "served_model_name": "vl-alias",
+        "trust_remote_code": True,
+    })))
+    # Bare stacked load of a DIFFERENT model B — no overrides supplied.
+    asyncio.run(mr.handle_model_load(
+        _Req(app, {"model": "meta-llama/Llama-3.2-3B-Instruct"})))
+
+    b_cfg = made["backends"][-1].config
+    defaults = NodeConfig()
+    assert b_cfg.model == "meta-llama/Llama-3.2-3B-Instruct"
+    assert b_cfg.kv_cache_dtype == defaults.kv_cache_dtype  # NOT "auto"
+    assert b_cfg.max_model_len is None                      # NOT 32768
+    assert b_cfg.quantization is None                       # NOT "awq"
+    assert b_cfg.trust_remote_code is False                 # NOT True
+    assert b_cfg.served_model_name is None                  # NOT ["vl-alias"]
+    assert b_cfg.gpu_memory_utilization == defaults.gpu_memory_utilization  # NOT 0.6
+
+
+def test_reload_same_primary_keeps_live_and_persisted_in_sync(monkeypatch):
+    """MAJOR regression: reloading the SAME primary without re-supplying an
+    override resets the live launch config AND the persisted NodeConfig to
+    defaults in lockstep — no live-vs-config.json divergence a later restart
+    would surface (the exact VLM-came-back-on-fp8 corruption)."""
+    import ainode.models.api_routes as mr
+
+    made = _patch_backend(monkeypatch)
+    cfg = NodeConfig(node_id="spark1", api_port=8000)
+    cfg.save = lambda: None
+    app = {"engine": None, "config": cfg, "cluster_state": ClusterState(),
+           "ray_autostart_state": None}
+
+    # Load VLM A with an explicit non-default kv.
+    asyncio.run(mr.handle_model_load(_Req(app, {
+        "model": "Qwen/Qwen2.5-VL-7B-Instruct", "kv_cache_dtype": "auto"})))
+    # Reload the SAME model, only bumping gmu (kv_cache_dtype NOT re-supplied).
+    asyncio.run(mr.handle_model_load(_Req(app, {
+        "model": "Qwen/Qwen2.5-VL-7B-Instruct", "gpu_memory_utilization": 0.7})))
+
+    live = made["backends"][-1].config
+    defaults = NodeConfig()
+    # Live backend and persisted shared config agree — both reset kv to default.
+    assert live.kv_cache_dtype == cfg.kv_cache_dtype == defaults.kv_cache_dtype
+    assert live.gpu_memory_utilization == cfg.gpu_memory_utilization == 0.7
+
+
+def test_explicit_kv_request_marks_provenance_on_launch_config(monkeypatch):
+    """Supplying kv_cache_dtype in the load body flags kv_cache_dtype_explicit on
+    BOTH the launch config and the persisted config, so nvidia.py's multimodal
+    fp8→auto downgrade is skipped for a user who explicitly asked for fp8 KV."""
+    import ainode.models.api_routes as mr
+
+    made = _patch_backend(monkeypatch)
+    cfg = NodeConfig(node_id="spark1", api_port=8000)
+    cfg.save = lambda: None
+    app = {"engine": None, "config": cfg, "cluster_state": ClusterState(),
+           "ray_autostart_state": None}
+    asyncio.run(mr.handle_model_load(_Req(app, {
+        "model": "Qwen/Qwen2.5-VL-7B-Instruct", "kv_cache_dtype": "fp8"})))
+
+    b = made["backends"][-1].config
+    assert b.kv_cache_dtype == "fp8"
+    assert b.kv_cache_dtype_explicit is True
+    assert cfg.kv_cache_dtype_explicit is True  # persisted for restart too
+
+
 def test_unload_one_stacked_instance_leaves_the_other(monkeypatch):
     """Unloading one stacked model stops ONLY that instance; the co-resident one
     keeps serving and stays in the manager."""