mirror of https://github.com/razor-ai/soup.git
82 lines
2.8 KiB
Python
82 lines
2.8 KiB
Python
"""Auto mixed-precision picker (v0.32.0 Part C).
|
|
|
|
Maps a model + GPU compute capability to the best mixed-precision setting:
|
|
``bf16`` (Ampere+, most modern models), ``fp16`` (Turing or known
|
|
fp16-stable models), or ``no`` (Pascal and older).
|
|
|
|
Quirk map is keyed on a lower-cased substring of the model id so we catch
|
|
both ``Qwen/Qwen2-7B-Instruct`` and ``alibaba/Qwen2.5-3B`` with one entry.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Literal
|
|
|
|
# Compute capability thresholds.
|
|
BF16_MIN_CC = 8.0 # Ampere
|
|
FP16_MIN_CC = 6.0 # Pascal
|
|
MAX_MODEL_NAME_LEN = 200
|
|
|
|
# Known fp16-stable / bf16-unstable model families.
|
|
# Value: preferred precision when CC supports it.
|
|
KNOWN_PRECISION_QUIRKS: dict[str, str] = {
|
|
"qwen2": "fp16",
|
|
"qwen2.5": "fp16",
|
|
"phi-3": "fp16",
|
|
"phi-3.5": "fp16",
|
|
"phi-4": "bf16",
|
|
"gemma-2": "bf16",
|
|
"mistral": "bf16",
|
|
"llama-3": "bf16",
|
|
"llama-2": "bf16",
|
|
}
|
|
|
|
|
|
def pick_mixed_precision(
|
|
model_name: str, compute_capability: float,
|
|
) -> Literal["bf16", "fp16", "no"]:
|
|
"""Pick mixed-precision mode for a model + GPU.
|
|
|
|
- cc < 6.0 → ``"no"`` (Pascal lacks reliable fp16 tensor cores)
|
|
- cc < 8.0 → ``"fp16"`` (no bf16 support)
|
|
- cc >= 8.0 → ``"bf16"`` unless model is in the fp16-stable quirks map
|
|
"""
|
|
if not isinstance(model_name, str) or not model_name:
|
|
raise ValueError("model_name must be a non-empty string")
|
|
if "\x00" in model_name:
|
|
raise ValueError("model_name must not contain null bytes")
|
|
if len(model_name) > MAX_MODEL_NAME_LEN:
|
|
raise ValueError(
|
|
f"model_name must be <= {MAX_MODEL_NAME_LEN} chars, got {len(model_name)}"
|
|
)
|
|
if not isinstance(compute_capability, (int, float)):
|
|
raise ValueError(
|
|
f"compute_capability must be a number, got {type(compute_capability)}"
|
|
)
|
|
if compute_capability < 0:
|
|
raise ValueError(
|
|
f"compute_capability must be non-negative, got {compute_capability}"
|
|
)
|
|
|
|
if compute_capability < FP16_MIN_CC:
|
|
return "no"
|
|
|
|
name_lc = model_name.lower()
|
|
quirk: str | None = None
|
|
# Sort longer substrings first so multi-version pairs work correctly:
|
|
# ``qwen2.5`` wins over ``qwen2``, ``phi-3.5`` wins over ``phi-3``.
|
|
# If you add a new family with multiple versions, the longest substring
|
|
# always wins — no manual ordering of the dict is required.
|
|
for substring in sorted(KNOWN_PRECISION_QUIRKS, key=len, reverse=True):
|
|
if substring in name_lc:
|
|
quirk = KNOWN_PRECISION_QUIRKS[substring]
|
|
break
|
|
|
|
if compute_capability < BF16_MIN_CC:
|
|
# Turing / Volta cannot do bf16 — fall back to fp16 regardless.
|
|
return "fp16"
|
|
|
|
if quirk is None:
|
|
return "bf16"
|
|
return "bf16" if quirk == "bf16" else "fp16"
|