model_runtime.py

import os

import torch


MPS_FALLBACK_ENV = "PYTORCH_ENABLE_MPS_FALLBACK"
FAST_MAX_NEW_TOKENS = 96
FAST_TEMPERATURE = 0.0
FAST_TOP_P = 1.0


def enable_mps_fallback():
    os.environ.setdefault(MPS_FALLBACK_ENV, "1")


def choose_device(allow_cpu=False, requested_device="auto"):
    if requested_device not in {"auto", "cuda", "mps", "cpu"}:
        raise ValueError(
            "requested_device must be one of: auto, cuda, mps, cpu"
        )

    if requested_device in {"auto", "cuda"} and torch.cuda.is_available():
        return {
            "name": "cuda",
            "device": torch.device("cuda"),
            "device_map": "auto",
            "torch_dtype": torch.bfloat16,
            "move_after_load": False,
        }

    mps_available = (
        getattr(torch.backends, "mps", None)
        and torch.backends.mps.is_available()
        and torch.backends.mps.is_built()
    )
    if requested_device in {"auto", "mps"} and mps_available:
        enable_mps_fallback()
        return {
            "name": "mps",
            "device": torch.device("mps"),
            "device_map": None,
            "torch_dtype": torch.float16,
            "move_after_load": True,
        }

    if requested_device == "cpu" or allow_cpu:
        return {
            "name": "cpu",
            "device": torch.device("cpu"),
            "device_map": None,
            "torch_dtype": torch.float32,
            "move_after_load": False,
        }

    if requested_device == "cuda":
        raise RuntimeError("CUDA was requested, but no CUDA GPU was detected.")
    if requested_device == "mps":
        raise RuntimeError(
            "Apple MPS was requested, but this Python environment cannot use it. "
            "Use an Apple Silicon Mac and a PyTorch build with MPS support."
        )
    raise RuntimeError(
        "No CUDA or Apple MPS GPU detected. Re-run with CPU fallback enabled for "
        "slow CPU mode."
    )


def model_load_kwargs(runtime):
    kwargs = {
        "dtype": runtime["torch_dtype"],
        "low_cpu_mem_usage": True,
        "trust_remote_code": True,
        "attn_implementation": "sdpa",
    }
    if runtime["device_map"] is not None:
        kwargs["device_map"] = runtime["device_map"]
    return kwargs


def adapter_load_kwargs(runtime):
    if runtime["name"] == "mps":
        return {
            "torch_device": "cpu",
            "autocast_adapter_dtype": True,
        }
    return {}


def place_model(model, runtime):
    if runtime["move_after_load"]:
        model = model.to(runtime["device"])
    if hasattr(model, "generation_config"):
        model.generation_config.use_cache = True
    return model


def resolve_generation_settings(max_new_tokens, temperature, top_p, fast=False):
    if fast:
        return {
            "max_new_tokens": (
                FAST_MAX_NEW_TOKENS if max_new_tokens is None else max_new_tokens
            ),
            "temperature": FAST_TEMPERATURE if temperature is None else temperature,
            "top_p": FAST_TOP_P if top_p is None else top_p,
        }
    return {
        "max_new_tokens": 300 if max_new_tokens is None else max_new_tokens,
        "temperature": 0.2 if temperature is None else temperature,
        "top_p": 0.9 if top_p is None else top_p,
    }