qwen_test_runner/vision/model_registry.py

"""
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,
    )