""" model_registry.py — The Qwen VLM model universe (analogous to SLOT_REGISTRY). One ModelSpec per checkpoint family. The benchmark iterates this registry to decide what to load on the 96GB RTX 6000 Pro. Both generations are natively multimodal (every checkpoint has its own ViT): * Qwen3.5 (qwen3_5 / qwen3_5_moe, AutoModelForMultimodalLM) — `enable_thinking` toggle on a single checkpoint. * Qwen3-VL (qwen3_vl / qwen3_vl_moe, AutoModelForImageTextToText) — separate -Instruct / -Thinking checkpoints. VRAM rule of thumb (weights only): bf16 ≈ 2·B, fp8 ≈ B, int4 ≈ 0.55·B GB. For MoE, ALL experts are resident, so total params drive memory; active params drive decode speed (and thus throughput / fleet score). Nothing hardcodes a repo id outside this file — it is the single source of truth for the model ladder, exactly as registry.py is for the schema. """ from __future__ import annotations from dataclasses import dataclass, field from typing import Literal, Optional Family = Literal["qwen3_5", "qwen3_5_moe", "qwen3_vl", "qwen3_vl_moe", "joycaption"] LoaderKind = Literal["multimodal_lm", "image_text_to_text", "llava_conditional"] ReasoningMode = Literal["toggle", "separate_ckpt", "none"] Precision = Literal["bf16", "fp8", "int4"] Reasoning = Literal["instruct", "thinking"] # weights-only bytes-per-parameter, in GB-per-billion-params _BYTES_PER_B: dict[str, float] = {"bf16": 2.0, "fp8": 1.0, "int4": 0.55} # Headroom reserved for the vision encoder, activations, and KV cache. VRAM_HEADROOM_GB = 12.0 GPU_VRAM_GB = 96.0 @dataclass(frozen=True) class ModelSpec: key: str # stable short id used on the CLI + in filenames repo_id: str # the Instruct / default checkpoint family: Family params_b: float # total parameters (billions) loader_kind: LoaderKind reasoning_mode: ReasoningMode active_b: Optional[float] = None # active params for MoE (decode speed) thinking_repo_id: Optional[str] = None # separate -Thinking ckpt (Qwen3-VL) quant_repo_ids: dict[Precision, str] = field(default_factory=dict) # fp8/int4 ckpts is_moe: bool = False is_baseline: bool = False # the user's fine-tuned reference notes: str = "" # ── VRAM math ──────────────────────────────────────────────────────────── def est_vram_gb(self, precision: Precision = "bf16") -> float: return self.params_b * _BYTES_PER_B[precision] def fits_on(self, precision: Precision = "bf16", gpu_gb: float = GPU_VRAM_GB, headroom: float = VRAM_HEADROOM_GB) -> bool: return self.est_vram_gb(precision) + headroom <= gpu_gb def available_precisions(self) -> list[Precision]: """bf16 is always available (base repo); fp8/int4 only if a quant ckpt exists.""" return ["bf16"] + [p for p in ("fp8", "int4") if p in self.quant_repo_ids] def best_fitting_precision(self, gpu_gb: float = GPU_VRAM_GB) -> Optional[Precision]: """Smallest-footprint precision that fits, preferring higher fidelity first (bf16 > fp8 > int4). Returns None if nothing fits without CPU offload.""" for prec in ("bf16", "fp8", "int4"): if prec in self.available_precisions() and self.fits_on(prec, gpu_gb): return prec return None def needs_offload(self, gpu_gb: float = GPU_VRAM_GB) -> bool: return self.best_fitting_precision(gpu_gb) is None def repo_for(self, precision: Precision, reasoning: Reasoning = "instruct") -> str: """Resolve the concrete HF repo id for a (precision, reasoning) request.""" if reasoning == "thinking" and self.reasoning_mode == "separate_ckpt": if self.thinking_repo_id is None: raise ValueError(f"{self.key}: no separate thinking checkpoint") base = self.thinking_repo_id else: base = self.repo_id if precision != "bf16" and precision in self.quant_repo_ids: # Quant repos are published for the instruct line; thinking quants are # less common, so fall back to the instruct quant id if needed. return self.quant_repo_ids[precision] return base def supports_thinking(self) -> bool: return self.reasoning_mode == "toggle" or self.thinking_repo_id is not None def enable_thinking_flag(self, reasoning: Reasoning) -> bool: """For toggle models, thinking is a generation flag, not a separate repo.""" return reasoning == "thinking" and self.reasoning_mode == "toggle" # ────────────────────────────────────────────────────────────────────────────── # THE MODEL UNIVERSE. Bench everything that fits on 96GB, both reasoning variants. # Stretch rungs (397B / 235B) are kept but flagged needs_offload by the VRAM math. # ────────────────────────────────────────────────────────────────────────────── def _q(*pairs: tuple[Precision, str]) -> dict[Precision, str]: return dict(pairs) MODEL_REGISTRY: dict[str, ModelSpec] = { # ── Qwen3.5 dense (native multimodal; enable_thinking toggle) ────────────── "qwen3.5-0.8b": ModelSpec( key="qwen3.5-0.8b", repo_id="Qwen/Qwen3.5-0.8B", family="qwen3_5", params_b=0.873, loader_kind="multimodal_lm", reasoning_mode="toggle", notes="smallest native VLM; floor of the ladder", ), "qwen3.5-2b": ModelSpec( key="qwen3.5-2b", repo_id="Qwen/Qwen3.5-2B", family="qwen3_5", params_b=2.0, loader_kind="multimodal_lm", reasoning_mode="toggle", ), "qwen3.5-4b": ModelSpec( key="qwen3.5-4b", repo_id="Qwen/Qwen3.5-4B", family="qwen3_5", params_b=4.0, loader_kind="multimodal_lm", reasoning_mode="toggle", ), "qwen3.5-9b": ModelSpec( key="qwen3.5-9b", repo_id="Qwen/Qwen3.5-9B", family="qwen3_5", params_b=9.0, loader_kind="multimodal_lm", reasoning_mode="toggle", ), "qwen3.5-27b": ModelSpec( key="qwen3.5-27b", repo_id="Qwen/Qwen3.5-27B", family="qwen3_5", params_b=27.0, loader_kind="multimodal_lm", reasoning_mode="toggle", quant_repo_ids=_q(("fp8", "Qwen/Qwen3.5-27B-FP8"), ("int4", "Qwen/Qwen3.5-27B-GPTQ-Int4")), ), # ── Qwen3.5 MoE ──────────────────────────────────────────────────────────── "qwen3.5-35b-a3b": ModelSpec( key="qwen3.5-35b-a3b", repo_id="Qwen/Qwen3.5-35B-A3B", family="qwen3_5_moe", params_b=35.0, active_b=3.0, loader_kind="multimodal_lm", reasoning_mode="toggle", is_moe=True, quant_repo_ids=_q(("fp8", "Qwen/Qwen3.5-35B-A3B-FP8"), ("int4", "Qwen/Qwen3.5-35B-A3B-GPTQ-Int4")), notes="MoE: ~3B active → decodes near a 3B dense", ), "qwen3.5-122b-a10b": ModelSpec( key="qwen3.5-122b-a10b", repo_id="Qwen/Qwen3.5-122B-A10B", family="qwen3_5_moe", params_b=122.0, active_b=10.0, loader_kind="multimodal_lm", reasoning_mode="toggle", is_moe=True, quant_repo_ids=_q(("fp8", "Qwen/Qwen3.5-122B-A10B-FP8"), ("int4", "Qwen/Qwen3.5-122B-A10B-GPTQ-Int4")), notes="fits 96GB only at GPTQ-Int4 (~67GB)", ), "qwen3.5-397b-a17b": ModelSpec( key="qwen3.5-397b-a17b", repo_id="Qwen/Qwen3.5-397B-A17B", family="qwen3_5_moe", params_b=397.0, active_b=17.0, loader_kind="multimodal_lm", reasoning_mode="toggle", is_moe=True, quant_repo_ids=_q(("fp8", "Qwen/Qwen3.5-397B-A17B-FP8"), ("int4", "Qwen/Qwen3.5-397B-A17B-GPTQ-Int4")), notes="STRETCH: needs CPU offload even at Int4", ), # ── Qwen3-VL dense (separate -Instruct / -Thinking) ──────────────────────── "qwen3vl-2b": ModelSpec( key="qwen3vl-2b", repo_id="Qwen/Qwen3-VL-2B-Instruct", thinking_repo_id="Qwen/Qwen3-VL-2B-Thinking", family="qwen3_vl", params_b=2.0, loader_kind="image_text_to_text", reasoning_mode="separate_ckpt", quant_repo_ids=_q(("fp8", "Qwen/Qwen3-VL-2B-Instruct-FP8")), ), "qwen3vl-4b": ModelSpec( key="qwen3vl-4b", repo_id="Qwen/Qwen3-VL-4B-Instruct", thinking_repo_id="Qwen/Qwen3-VL-4B-Thinking", family="qwen3_vl", params_b=4.0, loader_kind="image_text_to_text", reasoning_mode="separate_ckpt", quant_repo_ids=_q(("fp8", "Qwen/Qwen3-VL-4B-Instruct-FP8")), ), "qwen3vl-8b": ModelSpec( key="qwen3vl-8b", repo_id="Qwen/Qwen3-VL-8B-Instruct", thinking_repo_id="Qwen/Qwen3-VL-8B-Thinking", family="qwen3_vl", params_b=8.0, loader_kind="image_text_to_text", reasoning_mode="separate_ckpt", quant_repo_ids=_q(("fp8", "Qwen/Qwen3-VL-8B-Instruct-FP8")), ), "qwen3vl-32b": ModelSpec( key="qwen3vl-32b", repo_id="Qwen/Qwen3-VL-32B-Instruct", thinking_repo_id="Qwen/Qwen3-VL-32B-Thinking", family="qwen3_vl", params_b=32.0, loader_kind="image_text_to_text", reasoning_mode="separate_ckpt", quant_repo_ids=_q(("fp8", "Qwen/Qwen3-VL-32B-Instruct-FP8")), notes="top dense that fits bf16 (~64GB)", ), # ── Qwen3-VL MoE ─────────────────────────────────────────────────────────── "qwen3vl-30b-a3b": ModelSpec( key="qwen3vl-30b-a3b", repo_id="Qwen/Qwen3-VL-30B-A3B-Instruct", thinking_repo_id="Qwen/Qwen3-VL-30B-A3B-Thinking", family="qwen3_vl_moe", params_b=30.0, active_b=3.0, loader_kind="image_text_to_text", reasoning_mode="separate_ckpt", is_moe=True, quant_repo_ids=_q(("fp8", "Qwen/Qwen3-VL-30B-A3B-Instruct-FP8")), notes="MoE: ~3B active, fits bf16 (~60GB)", ), "qwen3vl-235b-a22b": ModelSpec( key="qwen3vl-235b-a22b", repo_id="Qwen/Qwen3-VL-235B-A22B-Instruct", thinking_repo_id="Qwen/Qwen3-VL-235B-A22B-Thinking", family="qwen3_vl_moe", params_b=235.0, active_b=22.0, loader_kind="image_text_to_text", reasoning_mode="separate_ckpt", is_moe=True, quant_repo_ids=_q(("fp8", "Qwen/Qwen3-VL-235B-A22B-Instruct-FP8")), notes="STRETCH: needs CPU offload", ), # ── User's fine-tuned reference (combined ViT-classification + JSON) ──────── "qwen3.5-0.8b-json-captioner": ModelSpec( key="qwen3.5-0.8b-json-captioner", repo_id="AbstractPhil/Qwen3.5-0.8B-json-captioner", family="qwen3_5", params_b=0.873, loader_kind="multimodal_lm", reasoning_mode="toggle", is_baseline=True, notes="LoRA-merged baseline: existence proof of native ViT-class + JSON; throughput ceiling", ), # ── JoyCaption (LLaVA: SigLIP2/SigLIP + Llama 3.1 8B). Captioner JSON-capacity ─ # baseline — not grounding-trained, so treat coordinate tasks as exploratory. "joycaption-beta-one": ModelSpec( key="joycaption-beta-one", repo_id="fancyfeast/llama-joycaption-beta-one-hf-llava", family="joycaption", params_b=8.48, loader_kind="llava_conditional", reasoning_mode="none", notes="SigLIP2 + Llama 3.1 8B captioner; latest JoyCaption; not grounding-trained", ), "joycaption-alpha-two": ModelSpec( key="joycaption-alpha-two", repo_id="fancyfeast/llama-joycaption-alpha-two-hf-llava", family="joycaption", params_b=8.48, loader_kind="llava_conditional", reasoning_mode="none", notes="prior JoyCaption (SigLIP v1); version-over-version JSON comparison", ), } # ────────────────────────────────────────────────────────────────────────────── # Query helpers # ────────────────────────────────────────────────────────────────────────────── def get_model(key: str) -> ModelSpec: if key not in MODEL_REGISTRY: raise KeyError(f"unknown model: {key!r}. known: {list(MODEL_REGISTRY)}") return MODEL_REGISTRY[key] def model_keys() -> list[str]: return list(MODEL_REGISTRY.keys()) def models_that_fit(gpu_gb: float = GPU_VRAM_GB, include_offload: bool = False) -> list[ModelSpec]: """Models that fit at some precision on `gpu_gb` (plus offload stretch if asked).""" out = [] for spec in MODEL_REGISTRY.values(): if spec.best_fitting_precision(gpu_gb) is not None: out.append(spec) elif include_offload: out.append(spec) return out def reasoning_variants(spec: ModelSpec, both: bool = True) -> list[Reasoning]: """The reasoning variants to run for a model.""" if not both or not spec.supports_thinking(): return ["instruct"] return ["instruct", "thinking"] def get_runner(model_key: str, precision: Precision = "bf16", reasoning: Reasoning = "instruct", device_map: Optional[str] = None, **kwargs): """Factory: build a real VLMRunner for a (model, precision, reasoning) request. Imports torch lazily (via runners), so importing this registry stays cheap. CPU-offload stretch rungs get device_map='auto' automatically. """ from .runners import VLMRunner spec = get_model(model_key) if precision == "bf16" and not spec.fits_on("bf16") and spec.best_fitting_precision() is not None: precision = spec.best_fitting_precision() # auto-downshift to a fitting quant repo = spec.repo_for(precision, reasoning) if device_map is None: device_map = "auto" if spec.needs_offload() else "cuda" return VLMRunner( model_id=repo, loader_kind=spec.loader_kind, precision=precision, enable_thinking=spec.enable_thinking_flag(reasoning), device_map=device_map, **kwargs, )