File size: 6,913 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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """Tier allocation planner (PRIVATE to Hardware_Module): picks the execution tier and the exact RAM/Disk split.
Tier 1 GPU (Kaggle 2x T4 / Colab) 4-bit NF4 (bitsandbytes) + device_map="auto" across all GPUs
Tier 2 CPU medium (Codespaces) safetensors memory-map, low_cpu_mem_usage=True, offload_folder=./offload_disk
Tier 3 ultra-low RAM (1-2 GB VPS) SequentialRingLoader: Layer 0 -> 1 -> ... -> N streamed from disk,
only `window` decoder layers resident at any moment
"""
from __future__ import annotations
import importlib.util
from pathlib import Path
from typing import Dict, Optional
GIB = 1024 ** 3
TIER_NAMES = {1: "GPU 4-bit NF4 (device_map=auto)", 2: "CPU memory-mapped safetensors",
3: "Sequential ring loader (disk streaming)"}
def _bytes_per_param(dtype: str) -> float:
return {"float32": 4.0, "float16": 2.0, "bfloat16": 2.0}[dtype]
def _bnb_available() -> bool:
return importlib.util.find_spec("bitsandbytes") is not None
def make_plan(prof: Dict, force_tier: Optional[int] = None) -> Dict:
model = prof["model"]
params = int(model["params"])
layers = max(1, int(model["num_layers"]))
layer_bytes_bf16 = int(model["layer_params"]) * 2
embed_bytes_bf16 = int(model["embed_params"]) * 2
ram_avail = int(prof["ram_avail_bytes"])
gpus = prof["gpus"]
vram_free = int(prof["gpu_free_bytes"])
# ---------------------------------------------------------------- tier choice
nf4_bytes = int(params * 0.56) + 300 * 1024 ** 2 # NF4 weights + double-quant constants + CUDA ctx
fp16_bytes = int(params * 2.1)
gpu_ok_nf4 = bool(gpus) and _bnb_available() and vram_free > nf4_bytes * 1.15
gpu_ok_fp16 = bool(gpus) and vram_free > fp16_bytes * 1.25
fp32_bytes = params * 4
cpu_fp32_ok = ram_avail > fp32_bytes * 1.45 + 1.0 * GIB # weights + grads of 2 layers + activations
cpu_bf16_ok = ram_avail > params * 2 * 1.45 + 1.0 * GIB
if force_tier in (1, 2, 3):
tier = int(force_tier)
if tier == 1 and not gpus:
print(" [plan] --tier 1 requested but no CUDA GPU was found -> falling back to tier 2")
tier = 2
elif gpu_ok_nf4 or gpu_ok_fp16:
tier = 1
elif cpu_fp32_ok or cpu_bf16_ok:
tier = 2
else:
tier = 3
plan: Dict = {"tier": tier, "name": TIER_NAMES[tier], "quant": None, "device_map": None,
"max_memory": None, "window": None, "evict_ends": None, "offload_folder": "./offload_disk",
"low_cpu_mem_usage": True}
# ---------------------------------------------------------------- per tier details
if tier == 1:
bf16_ok = all(g["bf16"] for g in gpus)
plan["dtype"] = "bfloat16" if bf16_ok else "float16" # T4 has no bf16 -> fp16 compute
use_nf4 = gpu_ok_nf4 or not gpu_ok_fp16
plan["quant"] = "nf4" if use_nf4 else None
if use_nf4 and not _bnb_available():
plan["quant"] = None # forced tier 1 without bitsandbytes
print(" [plan] bitsandbytes missing -> loading GPU model without quantisation")
plan["device_map"] = "auto"
plan["max_memory"] = {g["index"]: f"{max(1, int(g['free_bytes'] * 0.88 / GIB))}GiB" for g in gpus}
plan["max_memory"]["cpu"] = f"{max(1, int(ram_avail * 0.6 / GIB))}GiB"
weights = nf4_bytes if plan["quant"] else fp16_bytes
plan.update({"vram_bytes": weights, "ram_bytes": 0, "disk_bytes": 0})
elif tier == 2:
plan["dtype"] = "float32" if cpu_fp32_ok else "bfloat16"
weights = int(params * _bytes_per_param(plan["dtype"]))
resident = min(weights, int(ram_avail * 0.80))
plan.update({"device_map": {"": "cpu"}, "ram_bytes": resident, "disk_bytes": max(0, weights - resident),
"vram_bytes": 0, "max_memory": {"cpu": f"{max(1, int(ram_avail * 0.80 / GIB))}GiB"}})
else:
plan["dtype"] = "bfloat16"
window = 1 if ram_avail < 1.2 * GIB else 2
evict_ends = ram_avail < 3.0 * GIB or embed_bytes_bf16 > 0.25 * ram_avail
# embed + lm_head are each a vocab x hidden matrix (a tied checkpoint still materialises two copies):
# evict_ends -> only one of them is resident at a time, otherwise both stay in RAM
resident = window * layer_bytes_bf16 + (1 if evict_ends else 2) * embed_bytes_bf16
total = layers * layer_bytes_bf16 + embed_bytes_bf16
plan.update({"window": window, "evict_ends": evict_ends, "ram_bytes": min(resident, total),
"disk_bytes": max(0, total - resident), "vram_bytes": 0})
plan["weights_bytes"] = plan["ram_bytes"] + plan["disk_bytes"] + plan["vram_bytes"]
return plan
def describe_plan(plan: Dict, prof: Dict) -> str:
def gib(v: int) -> str:
return f"{v / GIB:.2f} GiB"
lines = [f" Tier {plan['tier']}: {plan['name']}",
f" dtype={plan['dtype']} quant={plan['quant']} device_map={plan['device_map']}",
f" split -> RAM {gib(plan['ram_bytes'])} | Disk {gib(plan['disk_bytes'])} | VRAM {gib(plan['vram_bytes'])}"]
if plan["tier"] == 3:
lines.append(f" ring window={plan['window']} layer(s), evict_ends={plan['evict_ends']} "
f"({prof['model']['num_layers']} layers streamed from disk, one forward pass = layer 0 -> N)")
return "\n".join(lines)
# ---------------------------------------------------------------- materialisation (tier 1 + 2)
def _from_pretrained(model_dir: str, dtype, **kwargs):
from transformers import AutoModelForCausalLM
try:
return AutoModelForCausalLM.from_pretrained(model_dir, dtype=dtype, **kwargs)
except TypeError:
return AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=dtype, **kwargs)
def materialize_model(model_dir: str, plan: Dict, offload_dir: Path):
import torch
dtype = getattr(torch, plan["dtype"])
offload_dir = Path(offload_dir)
offload_dir.mkdir(parents=True, exist_ok=True)
common = {"low_cpu_mem_usage": True}
if plan["tier"] == 1:
kwargs = dict(common, device_map="auto", max_memory=plan["max_memory"], offload_folder=str(offload_dir))
if plan["quant"] == "nf4":
from transformers import BitsAndBytesConfig
kwargs["quantization_config"] = BitsAndBytesConfig(
load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=dtype)
model = _from_pretrained(model_dir, dtype, **kwargs)
else: # tier 2: safetensors mmap -> CPU
model = _from_pretrained(model_dir, dtype, device_map={"": "cpu"}, offload_folder=str(offload_dir),
**common)
model.config.use_cache = True
return model |