diff --git a/ainode/core/config.py b/ainode/core/config.py index e8e454d..27b4f67 100644 --- a/ainode/core/config.py +++ b/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 diff --git a/ainode/engine/backends/nvidia.py b/ainode/engine/backends/nvidia.py index ac70a59..74a888c 100644 --- a/ainode/engine/backends/nvidia.py +++ b/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"]) diff --git a/ainode/models/api_routes.py b/ainode/models/api_routes.py index e33fad3..37404db 100644 --- a/ainode/models/api_routes.py +++ b/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"]) diff --git a/ainode/training/api_routes.py b/ainode/training/api_routes.py index 6af0463..88b42e3 100644 --- a/ainode/training/api_routes.py +++ b/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) diff --git a/tests/test_autodata.py b/tests/test_autodata.py index d3934ed..89de4e6 100644 --- a/tests/test_autodata.py +++ b/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 diff --git a/tests/test_cluster_load.py b/tests/test_cluster_load.py index 2eb5b82..00fb0ea 100644 --- a/tests/test_cluster_load.py +++ b/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.""" diff --git a/tests/test_nvidia_backend.py b/tests/test_nvidia_backend.py index b94179b..8a9281b 100644 --- a/tests/test_nvidia_backend.py +++ b/tests/test_nvidia_backend.py @@ -8,6 +8,7 @@ from __future__ import annotations +import json import subprocess from typing import Any, List from unittest import mock @@ -1122,3 +1123,85 @@ def test_download_max_workers_env(monkeypatch): assert _download_max_workers() == 2 monkeypatch.setenv("AINODE_DOWNLOAD_MAX_WORKERS", "junk") assert _download_max_workers() == 4 + + +# --------------------------------------------------------------------------- +# Multimodal fp8-KV auto-skip — fp8 KV corrupts VLM generation on GB10 +# --------------------------------------------------------------------------- + + +def _write_model_dir(tmp_path, *, vision=False, architectures=None, config_json=True): + """Create an on-disk model dir (flat org--name slug) with a config.json so + _local_model_dir resolves and _is_multimodal_model can read it.""" + slug = "org--model" + d = tmp_path / slug + d.mkdir() + if config_json: + cfg: dict = {} + if vision: + cfg["vision_config"] = {"hidden_size": 8} + if architectures is not None: + cfg["architectures"] = architectures + (d / "config.json").write_text(json.dumps(cfg)) + else: + (d / "placeholder.txt").write_text("x") # dir non-empty, no config.json + return _make_config(model="org/model", models_dir=str(tmp_path)) + + +def _kv_dtype_arg(backend): + args = backend._build_vllm_serve_args(tp_size=1) + if "--kv-cache-dtype" not in args: + return None + return args[args.index("--kv-cache-dtype") + 1] + + +def test_multimodal_vision_config_downgrades_fp8_to_auto(tmp_path): + cfg = _write_model_dir(tmp_path, vision=True) + cfg.kv_cache_dtype = "fp8" + assert _kv_dtype_arg(NvidiaBackend(cfg)) == "auto" + + +def test_multimodal_vl_architecture_downgrades_fp8_to_auto(tmp_path): + cfg = _write_model_dir(tmp_path, architectures=["Qwen2_5_VLForConditionalGeneration"]) + cfg.kv_cache_dtype = "fp8" + assert _kv_dtype_arg(NvidiaBackend(cfg)) == "auto" + + +def test_text_model_keeps_fp8(tmp_path): + cfg = _write_model_dir(tmp_path, architectures=["LlamaForCausalLM"]) + cfg.kv_cache_dtype = "fp8" + assert _kv_dtype_arg(NvidiaBackend(cfg)) == "fp8" + + +def test_explicit_non_default_kv_wins_for_multimodal(tmp_path): + """An explicit (non-'fp8') kv_cache_dtype is never downgraded, even for a VLM.""" + cfg = _write_model_dir(tmp_path, vision=True) + cfg.kv_cache_dtype = "fp8_e5m2" + assert _kv_dtype_arg(NvidiaBackend(cfg)) == "fp8_e5m2" + + +def test_explicit_fp8_honored_for_multimodal(tmp_path): + """An EXPLICIT fp8 kv_cache_dtype (flagged kv_cache_dtype_explicit) is honored + even on a VLM — the safety downgrade only fires on the fp8 DEFAULT, so a user + who knows their VLM/vLLM combo handles fp8 KV can opt back in.""" + cfg = _write_model_dir(tmp_path, vision=True) + cfg.kv_cache_dtype = "fp8" + cfg.kv_cache_dtype_explicit = True + assert _kv_dtype_arg(NvidiaBackend(cfg)) == "fp8" + + +def test_default_fp8_still_downgrades_for_multimodal(tmp_path): + """The fp8 DEFAULT (not explicitly requested) still downgrades to auto on a + VLM — provenance is what gates the downgrade, not the literal value alone.""" + cfg = _write_model_dir(tmp_path, vision=True) + cfg.kv_cache_dtype = "fp8" + cfg.kv_cache_dtype_explicit = False + assert _kv_dtype_arg(NvidiaBackend(cfg)) == "auto" + + +def test_unreadable_config_keeps_fp8_default(tmp_path): + """No config.json (e.g. remote repo-id not yet downloaded) → keep fp8; we + never fetch remote config in the serve-args builder.""" + cfg = _write_model_dir(tmp_path, config_json=False) + cfg.kv_cache_dtype = "fp8" + assert _kv_dtype_arg(NvidiaBackend(cfg)) == "fp8"