aryanbains's picture
Add Rook-V1 model, owner-reported comparisons and recorded evaluation
f6976d0 verified
Raw History Blame Contribute Delete
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")