File size: 4,263 Bytes
5b39202 | 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 | """Portable native WiggleGPT loader and exact GPT-2 tokenizer.
The model class is the byte-preserved source in native/model_bio.py.
Model construction and all forwards are CUDA only. Tokenization is CPU only.
"""
from __future__ import annotations
import base64
import hashlib
import os
from pathlib import Path
import subprocess
import sys
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / 'native'))
DESKTOP_GPU_APPS = {
"shellhost.exe", "crossdeviceresume.exe", "lghub_system_tray.exe", "dcv2.exe",
"snippingtool.exe", "steelseriesggclient.exe", "explorer.exe", "textinputhost.exe",
"startmenuexperiencehost.exe", "searchhost.exe", "shellexperiencehost.exe",
"msedgewebview2.exe", "windowsterminal.exe", "chrome.exe", "msedge.exe", "firefox.exe",
"chatgpt.exe", "codex-computer-use-swift.exe", "razerappengine.exe", "notepad++.exe",
"obsidian.exe", "dwm.exe",
}
def get_encoder():
import tiktoken
from tiktoken_ext.openai_public import r50k_pat_str
raw = (ROOT / "tokenizer" / "gpt2.tiktoken").read_bytes()
# Exact canonical GPT-2/r50k merge table, independent of network/cache paths.
if hashlib.sha256(raw).hexdigest() != "306cd27f03c1a714eca7108e03d66b7dc042abe8c258b44c199a7ed9838dd930":
raise ValueError("Bundled GPT-2 merge table is damaged.")
ranks = {base64.b64decode(token): int(rank) for token, rank in
(line.split() for line in raw.splitlines() if line)}
return tiktoken.Encoding(name="gpt2", pat_str=r50k_pat_str, mergeable_ranks=ranks,
special_tokens={"<|endoftext|>": 50256}, explicit_n_vocab=50257)
def check_gpu():
result = subprocess.run(
["nvidia-smi", "--query-compute-apps=pid,process_name", "--format=csv,noheader,nounits"],
capture_output=True, text=True, check=True,
)
busy = []
for line in result.stdout.splitlines():
fields = line.split(",", 1)
# Current Windows WDDM reports ordinary graphics clients as C+G too.
# Desktop rendering is covered by the measured VRAM cap; other jobs wait.
desktop = len(fields) == 2 and Path(fields[1].strip()).name.lower() in DESKTOP_GPU_APPS
if fields[0].strip().isdigit() and int(fields[0]) != os.getpid() and not desktop:
busy.append(line.strip())
if busy:
raise RuntimeError("Another GPU job is running. Finish it before loading chat: " + "; ".join(busy))
def load_model(checkpoint_path):
check_gpu()
import torch
from model_bio import GPT, GPTConfig
if not torch.cuda.is_available():
raise RuntimeError("WiggleGPT needs CUDA here; CPU model execution is disabled.")
torch.set_num_threads(2)
torch.cuda.set_device(0)
free, total = torch.cuda.mem_get_info()
budget = min(total - int(2.5 * 1024**3), free - int(0.5 * 1024**3))
if budget < 2 * 1024**3:
raise RuntimeError("Less than 2 GiB is available after leaving desktop room.")
torch.cuda.set_per_process_memory_fraction(budget / total, 0)
torch.set_float32_matmul_precision("high")
checkpoint_path = Path(checkpoint_path)
checkpoint = torch.load(checkpoint_path, map_location="cpu", mmap=True, weights_only=False)
config = GPTConfig(**checkpoint["model_args"])
# Construct directly on CUDA. In particular, never run a CPU model forward.
with torch.device("cuda:0"):
model = GPT(config)
state = {key.removeprefix("_orig_mod."): value for key, value in checkpoint["model"].items()}
model.load_state_dict(state, strict=True)
model.eval()
metadata = {"checkpoint": str(checkpoint_path), "saved_step": checkpoint.get("iter_num", checkpoint.get("step")),
"parameters": sum(parameter.numel() for parameter in model.parameters()),
"model_args": checkpoint["model_args"], "torch": torch.__version__}
del checkpoint, state
encoder = get_encoder()
print(f"CUDA: {torch.cuda.get_device_name(0)}; {metadata['parameters']:,} parameters; "
f"{budget / 1024**3:.2f} GiB memory cap", flush=True)
print(f"Checkpoint: {checkpoint_path}", flush=True)
return model, config, encoder, metadata
|