File size: 3,286 Bytes
1375f74 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | 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,
}
|