| """Deterministic sequential ring loader (PRIVATE to Hardware_Module) - Tier 3 (1-2 GB RAM). |
| |
| The model skeleton is built with EMPTY parameters (zero-size tensors, real CPU buffers such as the |
| rotary inv_freq). Forward pre/post hooks then stream weights from the safetensors files in the exact |
| order embed -> layer 0 -> layer 1 -> ... -> layer N-1 -> norm -> lm_head (and wrap around to layer 0 for |
| the next token). At most `window` (default 2) decoder layers are resident at any time. |
| |
| Parameters are always REPLACED by fresh nn.Parameter objects (never mutated in place), so autograd |
| graphs / checkpoint recomputation that still reference an evicted weight stay valid. |
| |
| Extra capabilities used by the Elevator trainer: |
| pin(idxs)/unpin(idxs, save=True) keep trainable layers resident, persist them as overlay shards |
| export(dest) write the full (base + overlay) model as safetensors shards |
| """ |
| from __future__ import annotations |
|
|
| import gc |
| import json |
| import shutil |
| import time |
| from collections import deque |
| from pathlib import Path |
| from typing import Dict, Iterator, List, Optional, Tuple |
|
|
| import torch |
| from torch import nn |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
|
|
|
|
| def _from_config(config, dtype): |
| from transformers import AutoModelForCausalLM |
| try: |
| return AutoModelForCausalLM.from_config(config, dtype=dtype) |
| except TypeError: |
| return AutoModelForCausalLM.from_config(config, torch_dtype=dtype) |
|
|
|
|
| class SequentialRingLoader: |
| EMBED_KEY = "model.embed_tokens.weight" |
| HEAD_KEY = "lm_head.weight" |
|
|
| def __init__(self, model_dir, dtype=torch.bfloat16, window: int = 2, evict_ends: bool = True, |
| overlay_dir: Optional[str] = None, row_recipe: Optional[Dict] = None, verbose: bool = True): |
| self.model_dir = Path(model_dir) |
| self.dtype = dtype |
| self.window = max(1, int(window)) |
| self.evict_ends = bool(evict_ends) |
| self.overlay_dir = Path(overlay_dir) if overlay_dir else self.model_dir / ".rss_overlay" |
| self.row_recipe = {int(k): [int(x) for x in v] for k, v in (row_recipe or {}).items()} |
| self.verbose = verbose |
| self.trace: deque = deque(maxlen=256) |
| self.pinned: set = set() |
| self.loaded: set = set() |
| self.overlay_map: Dict[str, Path] = {} |
| self.weight_map: Dict[str, str] = {} |
| self._handles: Dict[str, object] = {} |
| self.model = None |
| self.segments: List[Dict] = [] |
| self.by_name: Dict[str, Dict] = {} |
| self.num_layers = 0 |
| self.original_tied = False |
|
|
| |
| def _index_weights(self) -> None: |
| index = self.model_dir / "model.safetensors.index.json" |
| if index.exists(): |
| self.weight_map = dict(json.loads(index.read_text(encoding="utf-8"))["weight_map"]) |
| else: |
| single = self.model_dir / "model.safetensors" |
| if not single.exists(): |
| raise FileNotFoundError(f"no model.safetensors / index in {self.model_dir}") |
| with safe_open(str(single), framework="pt", device="cpu") as handle: |
| self.weight_map = {k: "model.safetensors" for k in handle.keys()} |
|
|
| def _handle(self, filename: str): |
| if filename not in self._handles: |
| self._handles[filename] = safe_open(str(self.model_dir / filename), framework="pt", device="cpu") |
| return self._handles[filename] |
|
|
| def _raw_get(self, key: str) -> torch.Tensor: |
| if key in self.overlay_map: |
| with safe_open(str(self.overlay_map[key]), framework="pt", device="cpu") as handle: |
| return handle.get_tensor(key) |
| return self._handle(self.weight_map[key]).get_tensor(key) |
|
|
| def _get(self, key: str) -> torch.Tensor: |
| tensor = self._raw_get(key) |
| if self.row_recipe and key in (self.EMBED_KEY, self.HEAD_KEY): |
| tensor = tensor.clone() |
| for new_id, sources in self.row_recipe.items(): |
| if new_id < tensor.shape[0] and sources: |
| idx = torch.tensor([s for s in sources if s < tensor.shape[0]], dtype=torch.long) |
| if len(idx): |
| tensor[new_id] = tensor[idx].float().mean(0).to(tensor.dtype) |
| return tensor |
|
|
| |
| def build(self): |
| from accelerate import init_empty_weights |
| from transformers import AutoConfig |
|
|
| self._index_weights() |
| config = AutoConfig.from_pretrained(str(self.model_dir)) |
| self.original_tied = bool(getattr(config, "tie_word_embeddings", False)) |
| config.tie_word_embeddings = False |
| with init_empty_weights(): |
| model = _from_config(config, self.dtype) |
| for module in model.modules(): |
| for name, param in list(module._parameters.items()): |
| if param is not None: |
| module._parameters[name] = nn.Parameter(torch.empty(0, dtype=self.dtype), requires_grad=False) |
| model.eval() |
| model.config.use_cache = True |
|
|
| inner = model.model |
| self.num_layers = len(inner.layers) |
| self.segments = [{"name": "embed", "module": inner.embed_tokens, "prefix": "model.embed_tokens", "idx": None}] |
| for i, layer in enumerate(inner.layers): |
| self.segments.append({"name": f"layer{i}", "module": layer, "prefix": f"model.layers.{i}", "idx": i}) |
| self.segments.append({"name": "norm", "module": inner.norm, "prefix": "model.norm", "idx": None}) |
| self.segments.append({"name": "head", "module": model.lm_head, "prefix": "lm_head", "idx": None}) |
| self.by_name = {s["name"]: s for s in self.segments} |
| self.model = model |
|
|
| for seg in self.segments: |
| self._install_hooks(seg) |
| if not self.evict_ends: |
| for name in ("embed", "norm", "head"): |
| self._ensure(name) |
| if self.verbose: |
| print(f" [ring] built skeleton: {self.num_layers} layers on disk, window={self.window}, " |
| f"evict_ends={self.evict_ends}, dtype={str(self.dtype).replace('torch.', '')}") |
| return model |
|
|
| |
| def _key_for(self, seg: Dict, rel: str) -> Optional[str]: |
| rel = rel.replace(".base_layer", "") |
| if "lora_" in rel or "modules_to_save" in rel or "original_module" in rel: |
| return None |
| key = f"{seg['prefix']}.{rel}" |
| if key in self.weight_map: |
| return key |
| if seg["name"] == "head" and rel == "weight" and self.EMBED_KEY in self.weight_map: |
| return self.EMBED_KEY |
| return None |
|
|
| def _managed(self, seg: Dict) -> Iterator[Tuple[str, nn.Module, str]]: |
| for mod_name, mod in seg["module"].named_modules(): |
| for pname, param in list(mod._parameters.items()): |
| if param is None: |
| continue |
| rel = f"{mod_name}.{pname}" if mod_name else pname |
| key = self._key_for(seg, rel) |
| if key: |
| yield key, mod, pname |
|
|
| |
| def _ensure(self, name: str) -> None: |
| if name in self.loaded: |
| return |
| started = time.time() |
| for key, mod, pname in self._managed(self.by_name[name]): |
| old = mod._parameters[pname] |
| tensor = self._get(key).to(old.dtype) |
| mod._parameters[pname] = nn.Parameter(tensor, requires_grad=old.requires_grad) |
| self.loaded.add(name) |
| self.trace.append(("load", name, round(time.time() - started, 3))) |
|
|
| def _evict(self, name: str) -> None: |
| if name not in self.loaded or name in self.pinned: |
| return |
| for _key, mod, pname in self._managed(self.by_name[name]): |
| old = mod._parameters[pname] |
| mod._parameters[pname] = nn.Parameter(torch.empty(0, dtype=self.dtype), requires_grad=old.requires_grad) |
| self.loaded.discard(name) |
| self.trace.append(("evict", name, 0.0)) |
|
|
| |
| def _install_hooks(self, seg: Dict) -> None: |
| name, idx = seg["name"], seg["idx"] |
| module = seg["module"] |
| if idx is not None: |
| def pre(_m, _args, i=idx): |
| self._enter_layer(i) |
|
|
| def post(_m, _args, _out, i=idx): |
| if self.window == 1: |
| self._evict(f"layer{i}") |
| module.register_forward_pre_hook(pre) |
| module.register_forward_hook(post) |
| else: |
| def pre_end(_m, _args, n=name): |
| self._ensure(n) |
|
|
| def post_end(_m, _args, _out, n=name): |
| if self.evict_ends: |
| self._evict(n) |
| module.register_forward_pre_hook(pre_end) |
| module.register_forward_hook(post_end) |
|
|
| def _enter_layer(self, i: int) -> None: |
| keep = {f"layer{i}"} |
| self._ensure(f"layer{i}") |
| if self.window >= 2 and self.num_layers > 1: |
| nxt = f"layer{(i + 1) % self.num_layers}" |
| self._ensure(nxt) |
| keep.add(nxt) |
| for name in list(self.loaded): |
| if name.startswith("layer") and name not in keep and name not in self.pinned: |
| self._evict(name) |
|
|
| |
| def pin(self, idxs) -> None: |
| for i in idxs: |
| self._ensure(f"layer{i}") |
| self.pinned.add(f"layer{i}") |
|
|
| def _save_overlay(self, i: int) -> None: |
| seg = self.by_name[f"layer{i}"] |
| tensors = {} |
| for key, mod, pname in self._managed(seg): |
| tensors[key] = mod._parameters[pname].detach().to(self.dtype).contiguous().clone() |
| self.overlay_dir.mkdir(parents=True, exist_ok=True) |
| path = self.overlay_dir / f"layer_{i}_{int(time.time() * 1000)}.safetensors" |
| save_file(tensors, str(path), metadata={"format": "pt"}) |
| for key in tensors: |
| old = self.overlay_map.get(key) |
| self.overlay_map[key] = path |
| if old and old != path and old.exists() and old not in set(self.overlay_map.values()): |
| old.unlink() |
|
|
| def unpin(self, idxs, save: bool = True) -> None: |
| for i in idxs: |
| name = f"layer{i}" |
| if save: |
| self._save_overlay(i) |
| self.pinned.discard(name) |
| for _key, mod, pname in self._managed(self.by_name[name]): |
| mod._parameters[pname].requires_grad_(False) |
| self._evict(name) |
| gc.collect() |
|
|
| def export(self, dest) -> Path: |
| """Write base + overlay weights as per-segment safetensors shards (1 segment in RAM at a time).""" |
| dest = Path(dest) |
| dest.mkdir(parents=True, exist_ok=True) |
| covered = set() |
| order = [s["name"] for s in self.segments] |
| keys_by_segment: Dict[str, List[str]] = {} |
| for name in order: |
| seg = self.by_name[name] |
| keys = sorted({key for key, _m, _p in self._managed(seg)}) |
| if name == "head" and self.original_tied: |
| keys = [] |
| keys_by_segment[name] = keys |
| covered.update(keys) |
| leftovers = sorted(k for k in self.weight_map if k not in covered |
| and not (self.original_tied and k == self.HEAD_KEY) |
| and not (k == self.EMBED_KEY)) |
| total_files = len([n for n in order if keys_by_segment[n]]) or 1 |
| weight_map_out: Dict[str, str] = {} |
| total_size = 0 |
| shard_no = 0 |
| for name in order: |
| keys = keys_by_segment[name] |
| if name == order[-1]: |
| keys = keys + leftovers |
| if not keys: |
| continue |
| shard_no += 1 |
| fname = f"model-{shard_no:05d}-of-{total_files:05d}.safetensors" |
| tensors = {k: self._get(k).to(self.dtype).contiguous() for k in keys} |
| save_file(tensors, str(dest / fname), metadata={"format": "pt"}) |
| for k, t in tensors.items(): |
| weight_map_out[k] = fname |
| total_size += t.numel() * t.element_size() |
| del tensors |
| gc.collect() |
| (dest / "model.safetensors.index.json").write_text( |
| json.dumps({"metadata": {"total_size": total_size}, "weight_map": weight_map_out}, indent=1), encoding="utf-8") |
| for fname in ("config.json", "generation_config.json"): |
| if (self.model_dir / fname).exists(): |
| shutil.copy(self.model_dir / fname, dest / fname) |
| if self.verbose: |
| print(f" [ring] exported {shard_no} shards ({total_size / 1e9:.2f} GB) -> {dest}") |
| return dest |
|
|
| def describe(self) -> str: |
| layers = sorted(int(n[5:]) for n in self.loaded if n.startswith("layer")) |
| try: |
| import psutil |
| rss = psutil.Process().memory_info().rss / 1024 ** 3 |
| mem = f"{rss:.2f} GiB RSS" |
| except Exception: |
| mem = "n/a" |
| return f"ring: resident layers={layers} pinned={sorted(self.pinned)} | {mem}" |
