mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
"""Pure contracts shared by model loading and runtime selection paths."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Collection
|
|
from typing import Literal
|
|
|
|
|
|
VALID_MODEL_TASKS = ("causal_lm", "classification")
|
|
VALID_QUANTIZATIONS = (None, "4bit", "8bit")
|
|
VALID_DTYPES = ("float32", "float16", "bfloat16")
|
|
|
|
|
|
def validate_model_load_request(
|
|
model_name: object,
|
|
task: str,
|
|
quantization: str | None,
|
|
dtype: str,
|
|
*,
|
|
valid_tasks: Collection[str] = VALID_MODEL_TASKS,
|
|
) -> None:
|
|
"""Reject malformed loader requests before any provider access."""
|
|
if not isinstance(model_name, str) or not model_name.strip():
|
|
raise ValueError("model_name must be a non-empty HuggingFace identifier or local path")
|
|
if task not in valid_tasks:
|
|
raise ValueError(f"Unknown task {task!r}. Choose from {list(valid_tasks)}")
|
|
if quantization not in VALID_QUANTIZATIONS:
|
|
raise ValueError(
|
|
"Unknown quantization {!r}. Choose None, '4bit', or '8bit'".format(quantization),
|
|
)
|
|
if dtype not in VALID_DTYPES:
|
|
raise ValueError(f"Unknown dtype {dtype!r}. Choose from {list(VALID_DTYPES)}")
|
|
|
|
|
|
def effective_model_memory_gb(estimate_gb: float, quantization: str | None) -> float:
|
|
"""Adjust a full-precision weight estimate for runtime quantization."""
|
|
factor = {"4bit": 4, "8bit": 2}.get(quantization, 1)
|
|
return estimate_gb / factor
|
|
|
|
|
|
def quantized_model_fits_gpu(
|
|
estimate_gb: float,
|
|
quantization: str | None,
|
|
available_gpu_gb: float,
|
|
) -> bool:
|
|
"""Return whether a quantized estimate fits with 30 percent headroom."""
|
|
effective_gb = effective_model_memory_gb(estimate_gb, quantization)
|
|
return (
|
|
quantization in ("4bit", "8bit")
|
|
and effective_gb > 0
|
|
and effective_gb < available_gpu_gb * 0.7
|
|
)
|
|
|
|
|
|
def should_snapshot_model(
|
|
*,
|
|
skip_snapshot: bool | None,
|
|
initial_gpu_free_gb: float,
|
|
remaining_gpu_free_gb: float,
|
|
has_native_quantization: bool,
|
|
estimate_gb: float,
|
|
quantization: str | None,
|
|
) -> bool:
|
|
"""Decide whether a restorable state snapshot fits the memory policy."""
|
|
if skip_snapshot is True:
|
|
return False
|
|
if skip_snapshot is False:
|
|
return True
|
|
if initial_gpu_free_gb > 0 and has_native_quantization:
|
|
return remaining_gpu_free_gb >= initial_gpu_free_gb * 0.4
|
|
if initial_gpu_free_gb > 0:
|
|
effective_gb = effective_model_memory_gb(estimate_gb, quantization)
|
|
return effective_gb <= initial_gpu_free_gb * 0.5
|
|
return True
|
|
|
|
|
|
def classify_architecture_size(
|
|
*,
|
|
is_moe: bool,
|
|
total_params_b: float,
|
|
num_experts: int,
|
|
model_name: str,
|
|
large_moe_name_patterns: Collection[str],
|
|
) -> Literal["dense", "small_moe", "large_moe"]:
|
|
"""Classify dense and MoE scale using the documented precedence rules."""
|
|
if not is_moe:
|
|
return "dense"
|
|
if total_params_b > 0:
|
|
is_small = total_params_b < 100
|
|
elif num_experts > 0:
|
|
is_small = num_experts <= 16
|
|
else:
|
|
name_lower = model_name.lower()
|
|
is_small = not any(pattern.lower() in name_lower for pattern in large_moe_name_patterns)
|
|
return "small_moe" if is_small else "large_moe"
|
|
|
|
|
|
def supports_bfloat16_target(
|
|
device: str,
|
|
*,
|
|
cuda_available: bool,
|
|
cuda_major: int | None,
|
|
torch_version: str,
|
|
) -> bool:
|
|
"""Evaluate the bfloat16 capability contract from deterministic inputs."""
|
|
if device.startswith("cuda"):
|
|
return cuda_available and cuda_major is not None and cuda_major >= 8
|
|
if device == "mps":
|
|
major, minor = (int(value) for value in torch_version.split(".")[:2])
|
|
return (major, minor) >= (2, 3)
|
|
return True
|