clef_runtime/runtime.py

"""Keep the official joint head while moving only vocabulary rows over PCIe."""

from __future__ import annotations

import gc
import importlib.util
import json
import sys
from contextlib import contextmanager
from functools import partial

import torch

from .bundle import bundle_manifest, read_bundle
from .core import PREPROCESSING_VERSION, PROFILES


class LexicalRows:
    def __init__(self, weight, *, device, dtype):
        if weight.device.type != "cpu" or weight.is_meta:
            raise ValueError("語彙行列には実体のあるCPUテンソルが必要です。")
        self.weight, self.device, self.dtype = weight, device, dtype

    def __getitem__(self, token_ids):
        shape = (*token_ids.shape, self.weight.shape[1])
        return self.weight.index_select(0, token_ids.reshape(-1).cpu()).reshape(shape).to(self.device, dtype=self.dtype)


class CPUEmbedding(torch.nn.Module):
    def __init__(self, weight, *, device, dtype):
        super().__init__()
        self.rows = LexicalRows(weight, device=device, dtype=dtype)
        self.weight = torch.nn.Parameter(weight, requires_grad=False)

    def forward(self, tokens):
        return self.rows[tokens]


def load_options(profile):
    settings = PROFILES[profile]
    result = {"device_map": {"": 0}}
    if settings["cpu_embeddings"]:
        result["device_map"] = {
            "visual": 0,
            "language_model.embed_tokens": "cpu",
            "language_model.layers": 0,
            "language_model.norm": 0,
            "language_model.rotary_emb": 0,
        }
    if settings["precision"] != "bf16":
        quantization = {
            "llm_int8_skip_modules": ["visual", "lm_head"],
            "llm_int8_enable_fp32_cpu_offload": settings["cpu_embeddings"],
        }
        if settings["precision"] == "int8":
            quantization["load_in_8bit"] = True
        else:
            quantization.update(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.bfloat16,
            )
        result["quantization"] = quantization
    return result


def image_options(processor, max_pixels):
    # Transformers 5 applies per-call pixel bounds through size, rather than
    # the legacy max_pixels constructor argument. Never change processor defaults.
    return {
        "images_kwargs": {
            "size": {
                "shortest_edge": min(processor.image_processor.size["shortest_edge"], max_pixels),
                "longest_edge": max_pixels,
            }
        }
    }


def image_measurement(processor, encoded):
    grid = (encoded.media or {}).get("image_grid_thw")
    if grid is None:
        return None, 0
    frames, height, width = map(int, grid[0])
    image_processor = processor.image_processor
    patch, merge = image_processor.patch_size, image_processor.merge_size
    return [width * patch, height * patch], frames * height * width // merge**2


def encode_complete(official, processor, record, max_length):
    # The release silently truncates state. Encode once without truncation and
    # refuse overflow before collating or allocating decoder activations.
    encoded = official.encode_record(processor.tokenizer, record, processor=processor, max_length=1_000_000)
    if len(encoded.input_ids) > max_length:
        raise ValueError(
            f"入力は{len(encoded.input_ids)}トークンです。上限{max_length}を超えます。補足や項目を短くするか処理解像度を下げてください。"
        )
    return encoded


def tensor_from_source(source, key):
    from safetensors import safe_open

    index = json.loads((source / "model.safetensors.index.json").read_text())
    with safe_open(source / index["weight_map"][key], framework="pt", device="cpu", backend="pread") as stream:
        return stream.get_tensor(key)


@contextmanager
def stream_weights():
    import transformers.modeling_utils as loading

    # Windows reserves commit for every copy-on-write mmap, including unread
    # pages. Keep TF's lazy tensor loading without retaining all shard maps.
    original = loading.safe_open
    loading.safe_open = partial(original, backend="pread")
    try:
        yield
    finally:
        loading.safe_open = original


class Runner:
    @classmethod
    def from_directory(cls, directory, profile):
        return cls(None, profile, directory=directory)

    def __init__(self, root, profile, *, directory=None):
        from safetensors.torch import load_file
        from transformers import AutoProcessor, BitsAndBytesConfig, Qwen3_5Model

        if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
            raise RuntimeError("ClefにはBF16対応のNVIDIA CUDA GPUが必要です。")
        self.settings, self.profile = PROFILES[profile], profile
        total = torch.cuda.get_device_properties(0).total_memory / 2**30
        if total < self.settings["min_vram_gib"]:
            raise RuntimeError(
                f"{self.settings['label']}のGPUメモリが不足しています。16GB設定またはFlashを選んでください。"
            )
        source, _ = read_bundle(directory, profile) if directory is not None else bundle_manifest(root, profile)
        spec = importlib.util.spec_from_file_location(
            "clef_release_" + self.settings["model"].replace("-", "_"), source / "joint_schema_model.py"
        )
        self.official = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = self.official
        spec.loader.exec_module(self.official)
        options = load_options(profile)
        saved = source if self.settings["precision"] != "bf16" else None
        if "quantization" in options:
            options["quantization_config"] = BitsAndBytesConfig(**options.pop("quantization"))
        print(f"{self.settings['label']}: {'保存済み' if saved else '公式重み'}を読み込み中", flush=True)  # noqa: T201
        with stream_weights():
            self.backbone, loading = Qwen3_5Model.from_pretrained(
                saved or source,
                dtype=torch.bfloat16,
                local_files_only=True,
                attn_implementation="sdpa",
                output_loading_info=True,
                **options,
            )
        expected_extra = {"lm_head.weight"}
        if loading["missing_keys"] or set(loading["unexpected_keys"]) != expected_extra:
            raise RuntimeError(f"Clef本体のキーが一致しません: {loading}")
        self.backbone.eval()
        self.backbone.config.use_cache = False
        if self.settings["cpu_embeddings"]:
            # Read the bundle's BF16 rows directly; CPU offload may cast to FP32.
            # Discard the old embedding together with its hook. Detaching a
            # hook can temporarily restore the entire vocabulary onto CUDA.
            embedding = tensor_from_source(source, "language_model.embed_tokens.weight")
            self.backbone.language_model.embed_tokens = CPUEmbedding(embedding, device="cuda:0", dtype=torch.bfloat16)
            gc.collect()
        self.lexical = LexicalRows(tensor_from_source(source, "lm_head.weight"), device="cuda:0", dtype=torch.bfloat16)
        self.head = self.official.JointSchemaHead(**json.loads((source / "joint_head_config.json").read_text()))
        self.head.load_state_dict(load_file(source / "joint_head.safetensors", backend="pread"), strict=True)
        self.head = self.head.to(device="cuda:0", dtype=torch.bfloat16).eval()
        self.processor = AutoProcessor.from_pretrained(source, local_files_only=True)
        self.source = source
        self.root, self.saved = root, saved

    @torch.inference_mode()
    def decide(self, request, image=None):
        record = {
            "model": self.settings["model"],
            "state": request["state"],
            "questions": request["questions"],
            "media_kwargs": image_options(self.processor, request["max_pixels"]),
        }
        if image is not None:
            record["images"] = [image]
        encoded = encode_complete(self.official, self.processor, record, request["max_length"])
        processing_size, image_tokens = image_measurement(self.processor, encoded)
        batch = self.official.collate_records([encoded], self.processor.tokenizer.pad_token_id, torch.device("cuda:0"))
        torch.cuda.reset_peak_memory_stats()
        media = batch.get("media") or {}
        model = self.backbone if media else self.backbone.language_model
        outputs = model(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
            use_cache=False,
            return_dict=True,
            **media,
        )
        logits = self.head(
            outputs.last_hidden_state, batch["input_ids"], batch["attention_mask"], batch["records"], self.lexical
        )[0]
        answers = {
            q.question_id: self.official.systemone_answer(
                request["questions"][q.question_id],
                dict(zip(q.option_ids, logit.float().softmax(-1).tolist(), strict=True)),
            )
            for q, logit in zip(encoded.questions, logits, strict=True)
        }
        torch.cuda.synchronize()
        return {
            "model": self.settings["model"],
            "answers": answers,
            "preprocessing": PREPROCESSING_VERSION,
            "processing_size": processing_size,
            "usage": {"input_tokens": len(encoded.input_ids), "image_tokens": image_tokens, "output_tokens": 0},
            "peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30,
            "peak_reserved_gib": torch.cuda.max_memory_reserved() / 2**30,
        }

    def close(self):
        self.backbone = self.head = self.lexical = self.processor = None
        gc.collect()
        torch.cuda.empty_cache()