| """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"]) |
|
|
| |
| nf4_bytes = int(params * 0.56) + 300 * 1024 ** 2 |
| 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 |
| 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} |
|
|
| |
| if tier == 1: |
| bf16_ok = all(g["bf16"] for g in gpus) |
| plan["dtype"] = "bfloat16" if bf16_ok else "float16" |
| 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 |
| 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 |
| |
| |
| 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) |
|
|
|
|
| |
| 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: |
| model = _from_pretrained(model_dir, dtype, device_map={"": "cpu"}, offload_folder=str(offload_dir), |
| **common) |
| model.config.use_cache = True |
| return model |
