File size: 4,611 Bytes
9d614f1 | 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 | """System profiler (PRIVATE to Hardware_Module): CPU, RAM (cgroup aware), GPUs, model size."""
from __future__ import annotations
import json
import os
import re
from pathlib import Path
from typing import Dict, Optional
GIB = 1024 ** 3
def _read_int(path: str) -> Optional[int]:
try:
raw = Path(path).read_text().strip()
if raw == "max":
return None
value = int(raw)
return value if value < (1 << 60) else None
except Exception:
return None
def cgroup_memory() -> Dict[str, Optional[int]]:
"""Container limits (GitHub Codespaces / Docker / Kaggle) - psutil alone over-reports."""
limit = _read_int("/sys/fs/cgroup/memory.max") or _read_int("/sys/fs/cgroup/memory/memory.limit_in_bytes")
used = _read_int("/sys/fs/cgroup/memory.current") or _read_int("/sys/fs/cgroup/memory/memory.usage_in_bytes")
return {"limit": limit, "used": used}
def estimate_model(model_path: str, token: Optional[str] = None) -> Dict:
"""Parameter count (in billions) + layer geometry from config.json (local dir or hub repo)."""
cfg: Dict = {}
local = Path(str(model_path)) / "config.json"
try:
if local.exists():
cfg = json.loads(local.read_text(encoding="utf-8"))
else:
from huggingface_hub import hf_hub_download
cfg = json.loads(Path(hf_hub_download(str(model_path), "config.json", token=token or None)).read_text())
except Exception:
cfg = {}
hidden = int(cfg.get("hidden_size", 0))
layers = int(cfg.get("num_hidden_layers", 0))
vocab = int(cfg.get("vocab_size", 0))
inter = int(cfg.get("intermediate_size", 4 * hidden if hidden else 0))
heads = int(cfg.get("num_attention_heads", 1) or 1)
kv_heads = int(cfg.get("num_key_value_heads", heads) or heads)
head_dim = int(cfg.get("head_dim", hidden // heads if hidden else 0) or 0)
tied = bool(cfg.get("tie_word_embeddings", False))
if hidden and layers and vocab:
attn = hidden * heads * head_dim + 2 * hidden * kv_heads * head_dim + heads * head_dim * hidden
mlp = 3 * hidden * inter
per_layer = attn + mlp + 2 * hidden
embed = vocab * hidden
total = embed * (1 if tied else 2) + layers * per_layer + hidden
else: # fall back to the name (e.g. "Qwen2.5-0.5B-Instruct", "...-9B")
match = re.search(r"(\d+(?:\.\d+)?)\s*[bB]", str(model_path))
total = int(float(match.group(1)) * 1e9) if match else int(1e9)
layers = layers or 24
embed = int(total * 0.1)
per_layer = (total - embed) // layers
return {"params": int(total), "params_b": round(total / 1e9, 3), "num_layers": layers,
"layer_params": int(per_layer), "embed_params": int(embed), "tied_embeddings": tied,
"hidden_size": hidden, "vocab_size": vocab}
def profile(model_path: str, token: Optional[str] = None) -> Dict:
import psutil
vm = psutil.virtual_memory()
cg = cgroup_memory()
total = vm.total
avail = vm.available
if cg["limit"]:
total = min(total, cg["limit"])
if cg["used"] is not None:
avail = min(avail, max(cg["limit"] - cg["used"], 0))
gpus = []
try:
import torch
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
free, tot = torch.cuda.mem_get_info(i)
major, _minor = torch.cuda.get_device_capability(i)
gpus.append({"index": i, "name": torch.cuda.get_device_name(i), "free_bytes": int(free),
"total_bytes": int(tot), "bf16": major >= 8})
except Exception:
gpus = []
model = estimate_model(model_path, token)
return {
"cpu_cores": os.cpu_count() or 1,
"ram_total_bytes": int(total), "ram_avail_bytes": int(avail),
"gpu_count": len(gpus), "gpus": gpus,
"gpu_free_bytes": sum(g["free_bytes"] for g in gpus),
"model": model,
}
def describe(prof: Dict) -> str:
m = prof["model"]
lines = [
f" CPU cores : {prof['cpu_cores']}",
f" RAM total/free : {prof['ram_total_bytes'] / GIB:.2f} GiB / {prof['ram_avail_bytes'] / GIB:.2f} GiB",
f" GPUs : {prof['gpu_count']}"
+ "".join(f"\n cuda:{g['index']} {g['name']} free {g['free_bytes'] / GIB:.1f}/{g['total_bytes'] / GIB:.1f} GiB"
for g in prof["gpus"]),
f" Model : ~{m['params_b']}B params, {m['num_layers']} layers, hidden {m['hidden_size']}, vocab {m['vocab_size']}",
]
return "\n".join(lines) |