Download implementation/rook/native/runtime.py from aryanbains/Rook-V1: direct link, hf CLI and curl.
- Browser
- Download file 14.2 kB
-
https://huggingface.co/aryanbains/Rook-V1/resolve/main/implementation/rook/native/runtime.py
- Command line
-
hf download hf://aryanbains/Rook-V1/implementation/rook/native/runtime.py
-
curl -L -o runtime.py https://huggingface.co/aryanbains/Rook-V1/resolve/main/implementation/rook/native/runtime.py
14.2 kB
| """CUDA-only native Lux model and explicit PEFT adapters. | |
| Uses verified Apache-2.0 native DecisionModel/CandidateHead/LoRA definitions; | |
| never imports the inference runtime or its HIP/shared-context optimizations. | |
| """ | |
| from __future__ import annotations | |
| import importlib | |
| import importlib.metadata | |
| import inspect | |
| import os | |
| import platform | |
| import shutil | |
| import subprocess | |
| import sys | |
| import time | |
| import types | |
| from pathlib import Path | |
| from .guards import MANIFEST_SHA, REVISION, check_layout, read_json, sha, verify_bundle | |
| def preflight(device: str, storage: Path) -> dict: | |
| if platform.system() != "Linux" or sys.version_info[:2] != (3, 12): | |
| raise RuntimeError("GPU execution requires Linux NVIDIA and Python 3.12; Windows is CPU preparation only") | |
| # A seed alone does not guarantee CUDA optimizer continuation. Configure | |
| # deterministic kernels before this process initializes CUDA/cuBLAS. | |
| os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") | |
| import torch | |
| torch.use_deterministic_algorithms(True) | |
| torch.backends.cudnn.benchmark = False | |
| torch.backends.cudnn.deterministic = True | |
| if not device.startswith("cuda:") or not torch.cuda.is_available() or torch.version.hip is not None or not torch.version.cuda: | |
| raise RuntimeError("explicit NVIDIA CUDA device required") | |
| torch.cuda.set_device(device) | |
| if not torch.cuda.is_bf16_supported(): | |
| raise RuntimeError("GPU must support BF16") | |
| versions = {} | |
| for name in ("torch", "transformers", "peft", "accelerate", "safetensors", "tokenizers", "flash-linear-attention", "causal-conv1d", "triton"): | |
| try: | |
| versions[name] = importlib.metadata.version(name) | |
| except importlib.metadata.PackageNotFoundError: | |
| versions[name] = None | |
| # Architecture reviewed at this version. CUDA compatibility still requires probes. | |
| if versions["transformers"] != "5.17.0": | |
| raise RuntimeError("reviewed architecture requires Transformers 5.17.0; do not silently change versions") | |
| if any(versions[k] is None for k in ("peft", "accelerate", "safetensors", "tokenizers")): | |
| raise RuntimeError("missing GPU integration dependency") | |
| free, total = torch.cuda.mem_get_info() | |
| props = torch.cuda.get_device_properties(device) | |
| driver = subprocess.run(["nvidia-smi", "--query-gpu=name,driver_version,memory.total,memory.free", "--format=csv,noheader"], capture_output=True, text=True, check=True, timeout=15).stdout.strip() | |
| meminfo = dict(line.split(":", 1) for line in Path("/proc/meminfo").read_text().splitlines()) | |
| return {"platform": platform.platform(), "python": sys.version, "gpu": props.name, | |
| "device": device, "compute_capability": [props.major, props.minor], "bf16": True, | |
| "gpu_free_bytes": free, "gpu_total_bytes": total, "cuda_runtime": torch.version.cuda, | |
| "driver_inventory": driver, "host_total_bytes": int(meminfo["MemTotal"].split()[0])*1024, | |
| "host_available_bytes": int(meminfo["MemAvailable"].split()[0])*1024, | |
| "disk_free_bytes": shutil.disk_usage(storage).free, "packages": versions, | |
| "optimized_rocm_runtime": False, "shared_context": False, | |
| "cuda_correctness_verified": False, | |
| "reproducibility": reproducibility_settings()} | |
| def reproducibility_settings(): | |
| import torch | |
| return {"deterministic_algorithms": torch.are_deterministic_algorithms_enabled(), | |
| "cudnn_benchmark": torch.backends.cudnn.benchmark, | |
| "cudnn_deterministic": torch.backends.cudnn.deterministic, | |
| "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG")} | |
| def native_sources(bundle: Path): | |
| """Call only after complete verification. Isolated vendored package namespace.""" | |
| name = "_rook_verified_lux_" + REVISION | |
| directory = (bundle / "decision2/_vendor/dev2model").resolve() | |
| if name in sys.modules: | |
| if sys.modules[name].__path__ != [str(directory)]: | |
| raise RuntimeError("native source namespace already bound to another bundle") | |
| else: | |
| package = types.ModuleType(name) | |
| package.__path__ = [str(directory)] | |
| sys.modules[name] = package | |
| # Importing verified source must not add __pycache__ files to the strict | |
| # immutable bundle. Load the lazily referenced source module here too. | |
| previous_bytecode = sys.dont_write_bytecode | |
| sys.dont_write_bytecode = True | |
| try: | |
| native = importlib.import_module(name + ".decision_model") | |
| lora = importlib.import_module(name + ".lora") | |
| importlib.import_module(name + ".source") | |
| return native, lora | |
| finally: | |
| sys.dont_write_bytecode = previous_bytecode | |
| def _assign(module, key, tensor): | |
| import torch | |
| owner_name, _, leaf = key.rpartition(".") | |
| owner = module.get_submodule(owner_name) if owner_name else module | |
| if leaf in owner._parameters: | |
| owner._parameters[leaf] = torch.nn.Parameter(tensor, requires_grad=False) | |
| elif leaf in owner._buffers: | |
| owner._buffers[leaf] = tensor | |
| else: | |
| raise ValueError(f"unexpected tensor: {key}") | |
| def load_model(bundle: Path, audit: Path, device: str, *, reference=False): | |
| os.environ["HF_HUB_OFFLINE"] = "1" | |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" | |
| verified = verify_bundle(bundle, audit) | |
| hardware = preflight(device, bundle) | |
| required = (verified["counts"]["backbone"] * 4 if reference else verified["backbone_storage_bytes"]) + verified["counts"]["head"]*4 | |
| if hardware["gpu_free_bytes"] < required + (2 << 30): | |
| raise RuntimeError("insufficient free VRAM for resident weights plus 2 GiB loading margin") | |
| if hardware["host_available_bytes"] < verified["largest_tensor_bytes"] + (2 << 30): | |
| raise RuntimeError("insufficient free host RAM for one tensor plus 2 GiB loading margin") | |
| import torch | |
| from safetensors import safe_open | |
| from transformers import AutoTokenizer | |
| from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig | |
| from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel | |
| native, _ = native_sources(bundle) | |
| started = time.perf_counter() | |
| torch.cuda.reset_peak_memory_stats() | |
| config = Qwen3_5TextConfig.from_dict(read_json(bundle / "backbone/config.json")) | |
| config.use_cache = False | |
| config._attn_implementation = "sdpa" | |
| if reference: | |
| # Independent HF staged FP32 GPU loader; no full host FP32 model. Native | |
| # forward/head/formatter remain the reference, BF16 autocast unchanged. | |
| backbone = Qwen3_5TextModel.from_pretrained(bundle / "backbone", config=config, | |
| dtype=torch.float32, device_map={"": device}, low_cpu_mem_usage=True, | |
| local_files_only=True, attn_implementation="sdpa") | |
| else: | |
| with torch.device("meta"): | |
| backbone = Qwen3_5TextModel(config) | |
| state = backbone.state_dict() | |
| actual = {n: {"shape": list(v.shape), "dtype": verified["layout"].get(n, {}).get("dtype")} for n, v in state.items()} | |
| check_layout(actual, verified["layout"]) | |
| del state | |
| for key, spec in verified["layout"].items(): | |
| # Reopen per tensor so previous shard mapping/pages are released. | |
| # No tensor/shard dictionary and no FP32->BF16 whole-model copy. | |
| with safe_open(str(bundle / verified["tensor_files"][key]), framework="pt", device="cpu") as stream: | |
| cpu_tensor = stream.get_tensor(key) | |
| tensor = cpu_tensor.to(device=device) | |
| if list(tensor.shape) != spec["shape"] or tensor.dtype != (torch.bfloat16 if spec["dtype"] == "BF16" else torch.float32): | |
| raise ValueError("loaded tensor differs from verified descriptor") | |
| _assign(backbone, key, tensor) | |
| del cpu_tensor, tensor | |
| # Nonpersistent RoPE buffers are not checkpoint tensors; reconstruct from | |
| # the pinned config on CUDA, never materialize arbitrary meta buffers. | |
| with torch.device(device): | |
| backbone.rotary_emb = type(backbone.rotary_emb)(config=config) | |
| with torch.device("meta"): | |
| head = native.CandidateHead(config.hidden_size, 256) | |
| check_layout({n: {"shape": list(v.shape), "dtype": "F32"} for n, v in head.state_dict().items()}, verified["head_layout"]) | |
| with safe_open(str(bundle / "decision_head.safetensors"), framework="pt", device="cpu") as stream: | |
| for key in verified["head_layout"]: | |
| _assign(head, key, stream.get_tensor(key).to(device=device)) | |
| metadata = read_json(bundle / "decision_config.json") | |
| model = native.DecisionModel(backbone, head, metadata) | |
| if any(t.is_meta for t in (*model.parameters(), *model.buffers())): | |
| raise RuntimeError("unmaterialized meta tensor remains") | |
| if sum(p.numel() for p in model.parameters()) != verified["manifest"]["parameters"]["loaded"]: | |
| raise RuntimeError("loaded parameter count mismatch") | |
| model.requires_grad_(False) | |
| model.eval() | |
| tokenizer = AutoTokenizer.from_pretrained(bundle, trust_remote_code=False, local_files_only=True) | |
| torch.cuda.synchronize() | |
| evidence = {"hardware": hardware, "loader": "hf-staged-fp32-native-reference" if reference else "meta-one-tensor-mixed-residency", | |
| "load_seconds": time.perf_counter()-started, "source_manifest_sha256": MANIFEST_SHA, | |
| "largest_host_tensor_bytes": verified["largest_tensor_bytes"], "resident_weight_bytes": required, | |
| "gpu_peak_allocated_bytes": torch.cuda.max_memory_allocated(), "gpu_peak_reserved_bytes": torch.cuda.max_memory_reserved(), | |
| "transformers_model_source_sha256": sha(Path(inspect.getfile(Qwen3_5TextModel))), | |
| **host_memory()} | |
| return model, tokenizer, verified, evidence | |
| def host_memory(): | |
| import resource | |
| status = dict(line.split(":", 1) for line in Path("/proc/self/status").read_text().splitlines() if ":" in line) | |
| return {"host_rss_bytes": int(status["VmRSS"].split()[0])*1024, | |
| "host_process_peak_rss_bytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss*1024} | |
| def runtime_versions(): | |
| import torch | |
| versions = {} | |
| for name in ("torch", "transformers", "peft", "accelerate", "safetensors", "tokenizers", "flash-linear-attention", "causal-conv1d", "triton"): | |
| try: | |
| versions[name] = importlib.metadata.version(name) | |
| except importlib.metadata.PackageNotFoundError: | |
| versions[name] = None | |
| versions["torch_cuda"] = torch.version.cuda | |
| versions["reproducibility"] = reproducibility_settings() | |
| from transformers.models.qwen3_5 import modeling_qwen3_5 | |
| from .guards import sha | |
| versions["qwen3_5_source_sha256"] = sha(Path(modeling_qwen3_5.__file__)) | |
| return versions | |
| def target_contract(model, verified, bundle): | |
| import torch | |
| _, lora = native_sources(bundle) | |
| targets = lora.select_target_modules(model.backbone) | |
| if targets != verified["targets"] or len(targets) != 248: | |
| raise ValueError("actual hybrid selector differs from expected 248 targets") | |
| modules = dict(model.backbone.named_modules()) | |
| dimensions = {} | |
| for name in targets: | |
| module = modules[name] | |
| expected = verified["layout"][name + ".weight"]["shape"] | |
| if not isinstance(module, torch.nn.Linear) or [module.out_features, module.in_features] != expected or module.bias is not None: | |
| raise ValueError("native projection dimensions/bias mismatch") | |
| dimensions[name] = [module.in_features, module.out_features] | |
| return targets, dimensions | |
| def attach_adapters(model, verified, bundle, config): | |
| import torch | |
| _, lora = native_sources(bundle) | |
| from importlib import import_module | |
| targets, dimensions = target_contract(model, verified, bundle) | |
| source = import_module(lora.__package__ + ".source") | |
| model.requires_grad_(False) | |
| lora.attach_lora(model, rank=config["rank"], alpha=config["alpha"], dropout=config["dropout"], | |
| source_kind="decision2", source_fingerprint=source.source_fingerprint(bundle)) | |
| model.metadata["rook_source_revision"] = REVISION | |
| for name, parameter in model.named_parameters(): | |
| if ".lora_A." in name or ".lora_B." in name: | |
| parameter.data = parameter.data.float() | |
| parameter.requires_grad_(True) | |
| if ".lora_B." in name and torch.count_nonzero(parameter).item() != 0: | |
| raise RuntimeError("LoRA initial B must be zero") | |
| audit_trainables(model) | |
| return {"targets": targets, "dimensions": dimensions, "gradient_checkpointing": False} | |
| def audit_trainables(model): | |
| import torch | |
| a = b = 0 | |
| result = {} | |
| for name, parameter in model.named_parameters(): | |
| expected = name.startswith("head.") or ".lora_A." in name or ".lora_B." in name | |
| if parameter.requires_grad != expected or expected and parameter.dtype != torch.float32: | |
| raise RuntimeError(f"unexpected trainable/storage contract: {name}") | |
| if expected: | |
| result[name] = list(parameter.shape) | |
| a += ".lora_A." in name | |
| b += ".lora_B." in name | |
| if (a, b) != (248, 248) or sum(n.startswith("head.") for n in result) != 10: | |
| raise RuntimeError("incomplete adapter/head trainables") | |
| return result | |
| def checkpointing(model, enabled: bool): | |
| """Only invoked by explicit comparison probe, default disabled.""" | |
| base = model.backbone.get_base_model() | |
| if not enabled: | |
| base.gradient_checkpointing_disable() | |
| if hasattr(base, "_require_grads_hook"): | |
| base.disable_input_require_grads() | |
| return | |
| if not base.supports_gradient_checkpointing or "gradient_checkpointing_kwargs" not in inspect.signature(base.gradient_checkpointing_enable).parameters: | |
| raise RuntimeError("backbone lacks inspected non-reentrant checkpointing interface") | |
| base.enable_input_require_grads() | |
| base.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) | |
| if not base.gradient_checkpointing: | |
| raise RuntimeError("gradient checkpointing was not enabled") | |