"""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