"""LoRA-based reversible ablation mode. Instead of permanent in-place weight surgery, applies ablation via rank-1 LoRA adapters. This provides: 1. **Reversibility**: LoRA adapters can be removed to restore original model 2. **Composability**: Multiple ablation adapters can be stacked/blended 3. **PEFT compatibility**: Exact-identity exports use standard HuggingFace PEFT Inspired by Heretic (p-e-w, 2025) which pioneered LoRA-mediated ablation. OBLITERATUS extends this with: - Multi-direction rank-k adapters (not just rank-1) - MoE-aware LoRA targeting (router + expert-specific adapters) - Integration with EGA per-expert directions - CoT-aware adapter strength modulation The mathematical equivalence to in-place projection depends on weight orientation: For W of shape (out, hidden) where d is in the hidden dimension: In-place: W' = W - scale * W @ d @ d^T LoRA: W' = W + B @ A where B = -scale * (W @ d), A = d^T For W of shape (hidden, out) (e.g., Conv1D layers): In-place: W' = W - scale * d @ d^T @ W LoRA: W' = W + B @ A where B = -scale * d, A = d^T @ W Both produce identical output, but LoRA stores {B, A} separately. References: - Hu et al. (2022): LoRA: Low-Rank Adaptation of Large Language Models - Heretic (p-e-w, 2025): LoRA-mediated directional ablation """ from __future__ import annotations import json import logging import math import os import re from dataclasses import dataclass from hashlib import sha256 from pathlib import Path from typing import TYPE_CHECKING, Any, Callable import torch import torch.nn as nn from safetensors.torch import load_file, save_file from obliteratus.checkpoint_provenance import ( AdapterIdentity, ArtifactIdentity, ProvenanceRecord, verify_provenance_record, ) from obliteratus.persistence_contracts import atomic_checkpoint_directory if TYPE_CHECKING: from obliteratus.abliterate import AbliterationPipeline logger = logging.getLogger(__name__) # Target module name patterns for LoRA adapter placement _LORA_TARGETS = [ "o_proj", "q_proj", "k_proj", "v_proj", "down_proj", "up_proj", "gate_proj", "gate", "router", ] _DIGEST = re.compile(r"^sha256:[0-9a-f]{64}$") _IMMUTABLE_REVISION = re.compile(r"^[0-9a-f]{40,64}$") _MODULE_PATH = re.compile(r"^[A-Za-z0-9_]+(?:\.[A-Za-z0-9_]+)*$") _PUBLIC_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/-]{0,511}$") _PUBLIC_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,511}$") _SECRET_VALUE = re.compile(r"(?i)(?:hf_[a-z0-9]{12,}|bearer\s+[a-z0-9._~+/-]{12,})") @dataclass(frozen=True) class BaseModelIdentity: """Exact base and tokenizer identity required for a canonical PEFT claim.""" repo_id: str revision: str weights_digest: str tokenizer_digest: str vocab_size: int architecture: str tied_embeddings: bool def __post_init__(self) -> None: if ( not isinstance(self.repo_id, str) or not self.repo_id or len(self.repo_id) > 512 or not _PUBLIC_REPO_ID.fullmatch(self.repo_id) or Path(self.repo_id).is_absolute() or ".." in self.repo_id.split("/") or _SECRET_VALUE.search(self.repo_id) ): raise ValueError("base repository identity is invalid") if not _IMMUTABLE_REVISION.fullmatch(self.revision): raise ValueError("base revision must be an immutable commit") for value, field in ( (self.weights_digest, "base weights digest"), (self.tokenizer_digest, "tokenizer digest"), ): if not _DIGEST.fullmatch(value): raise ValueError(f"{field} is invalid") if type(self.vocab_size) is not int or self.vocab_size <= 0: raise ValueError("base vocabulary size is invalid") if ( not isinstance(self.architecture, str) or not self.architecture or len(self.architecture) > 512 or not _PUBLIC_NAME.fullmatch(self.architecture) or Path(self.architecture).is_absolute() or _SECRET_VALUE.search(self.architecture) ): raise ValueError("base architecture is invalid") if type(self.tied_embeddings) is not bool: raise ValueError("base tied-embedding declaration is invalid") def to_dict(self) -> dict[str, Any]: return { "repo_id": self.repo_id, "revision": self.revision, "weights_digest": self.weights_digest, "tokenizer_digest": self.tokenizer_digest, "vocab_size": self.vocab_size, "architecture": self.architecture, "tied_embeddings": self.tied_embeddings, } def to_artifact_identity(self) -> ArtifactIdentity: return ArtifactIdentity("hub", self.repo_id, self.revision, self.weights_digest) def tokenizer_artifact_identity(self) -> ArtifactIdentity: return ArtifactIdentity( "hub", f"{self.repo_id}#tokenizer", self.revision, self.tokenizer_digest, ) @dataclass(frozen=True) class AdapterArtifact: artifact_id: str weights_path: Path config_path: Path provenance_path: Path def compute_lora_adapters( pipeline: AbliterationPipeline, rank: int = 1, ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: """Compute LoRA adapter pairs (B, A) for refusal direction ablation. For each target weight matrix W with refusal direction d: A = d^T @ W (rank-1: shape (1, in_features) or (rank, in_features)) B = -scale * d (rank-1: shape (out_features, 1) or (out_features, rank)) So that W + B @ A ≈ W - scale * (d @ d^T) @ W Args: pipeline: Initialized pipeline (post-DISTILL, pre-EXCISE). rank: LoRA rank (1 = rank-1 ablation, >1 = multi-direction). Returns: Dict mapping "layer.{idx}.{module}.{weight}" → (lora_B, lora_A) pairs. """ from obliteratus.strategies.utils import ( get_attention_module, get_ffn_module, get_layer_modules, ) from obliteratus.abliterate import ( _ATTN_OUT_NAMES, _ATTN_IN_NAMES, _FFN_OUT_NAMES, _FFN_IN_NAMES, _ROUTER_NAMES, ) if not pipeline.handle or not pipeline._strong_layers: return {} layers = get_layer_modules(pipeline.handle) arch = pipeline.handle.architecture module_names = { id(module): name for name, module in pipeline.handle.model.named_modules() if name } adapters: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} for idx in pipeline._strong_layers: if idx not in pipeline.refusal_subspaces: continue subspace = pipeline.refusal_subspaces[idx] n_dirs = min(rank, subspace.shape[0]) # Compute per-layer regularization (mirroring excise logic) reg = pipeline.regularization if pipeline.layer_adaptive_strength and idx in pipeline._layer_excise_weights: weight = pipeline._layer_excise_weights[idx] reg = pipeline.regularization + (1.0 - weight) * (1.0 - pipeline.regularization) * 0.15 if pipeline.float_layer_interpolation and idx in pipeline._float_layer_weights: float_w = pipeline._float_layer_weights[idx] reg = reg + (1.0 - float_w) * (1.0 - reg) * 0.3 scale = 1.0 - reg # Direction matrix: (n_dirs, hidden_dim) D = subspace[:n_dirs].float() # Collect target modules and their weight matrices targets: list[tuple[str, nn.Module, list[str]]] = [] try: attn = get_attention_module(layers[idx], arch) targets.append(("attn", attn, _ATTN_OUT_NAMES + _ATTN_IN_NAMES)) except (AttributeError, RuntimeError): pass try: ffn = get_ffn_module(layers[idx], arch) targets.append(("ffn", ffn, _FFN_OUT_NAMES + _FFN_IN_NAMES + _ROUTER_NAMES)) except (AttributeError, RuntimeError): pass for module_label, module, candidate_names in targets: for name in candidate_names: proj = getattr(module, name, None) if proj is None or not hasattr(proj, "weight"): continue W = proj.weight.data.float() if W.shape[-1] == D.shape[1]: # Standard: W is (out, hidden_dim), direction along last axis # A = D @ W^T → (n_dirs, out) ... no, we need the correct decomposition # W' = W - scale * D^T @ D @ W # LoRA: B @ A where B = -scale * D^T shape (hidden_dim, n_dirs) # A = D @ W shape (n_dirs, out)... # Actually for nn.Linear: output = input @ W^T + b # So to affect output: we need delta_W such that input @ delta_W^T # delta_W = -scale * d @ d^T @ W → B = -scale * d, A = d^T @ W # B shape: (out_features, n_dirs) where out_features = W.shape[0] # Wait, let me reconsider. For W shape (out, in): # delta_W = -scale * (d_col @ d_col^T) @ W # where d_col is (in, 1) if direction matches input dim # delta_W = -scale * d_col @ (d_col^T @ W) # = -scale * d_col @ coeff where coeff = d_col^T @ W = (1, out) # So: B = -scale * d_col = (in, 1), A = d_col^T @ W = (1, out) # But LoRA convention: delta_W = B @ A where B is (out, r), A is (r, in) # So we need: delta_W^T = A^T @ B^T # Hmm, let me just store the computed delta and split it # For each direction, compute the rank-1 adapter lora_B_parts = [] lora_A_parts = [] for di in range(n_dirs): d = D[di] # (hidden_dim,) d_col = d.unsqueeze(-1) # (hidden_dim, 1) # coeff = d^T @ W^T = (d @ W^T) → but W is (out, hidden), so: # For W @ d → (out, 1) projection coeff = (W @ d_col).squeeze(-1) # (out,) # delta_W = -scale * d_col @ coeff.unsqueeze(0) would be (hidden, out) # But we want (out, hidden) to match W shape # delta_W[i,j] = -scale * coeff[i] * d[j] # = B[i,:] @ A[:,j] where B = -scale * coeff.unsqueeze(1) = (out,1) # A = d.unsqueeze(0) = (1, hidden) lora_B_parts.append(-scale * coeff.unsqueeze(1)) # (out, 1) lora_A_parts.append(d.unsqueeze(0)) # (1, hidden) lora_B = torch.cat(lora_B_parts, dim=1) # (out, n_dirs) lora_A = torch.cat(lora_A_parts, dim=0) # (n_dirs, hidden) elif W.shape[0] == D.shape[1]: # Transposed case: W is (hidden_dim, out) lora_B_parts = [] lora_A_parts = [] for di in range(n_dirs): d = D[di] # (hidden_dim,) coeff = (d @ W) # (out,) lora_B_parts.append(-scale * d.unsqueeze(1)) # (hidden, 1) lora_A_parts.append(coeff.unsqueeze(0)) # (1, out) lora_B = torch.cat(lora_B_parts, dim=1) # (hidden, n_dirs) lora_A = torch.cat(lora_A_parts, dim=0) # (n_dirs, out) else: continue key = module_names.get(id(proj), f"layer.{idx}.{module_label}.{name}") adapters[key] = (lora_B.half(), lora_A.half()) pipeline.log(f"Computed {len(adapters)} LoRA adapter pairs (rank={rank})") return adapters def apply_lora_adapters( pipeline: AbliterationPipeline, adapters: dict[str, tuple[torch.Tensor, torch.Tensor]], ): """Apply pre-computed LoRA adapters by modifying weights in-place. This is equivalent to merging the LoRA into the base model. The adapters dict is stored in pipeline._lora_adapters for potential later unmerging. """ from obliteratus.strategies.utils import ( get_attention_module, get_ffn_module, get_layer_modules, ) if not pipeline.handle: return layers = get_layer_modules(pipeline.handle) arch = pipeline.handle.architecture named_modules = dict(pipeline.handle.model.named_modules()) applied = 0 for key, (lora_B, lora_A) in adapters.items(): proj = named_modules.get(key) if proj is None: parts = key.split(".") if len(parts) != 4 or parts[0] != "layer" or not parts[1].isdigit(): continue _, idx_str, module_label, weight_name = parts idx = int(idx_str) try: if module_label == "attn": module = get_attention_module(layers[idx], arch) else: module = get_ffn_module(layers[idx], arch) except (AttributeError, RuntimeError): continue proj = getattr(module, weight_name, None) if proj is None or not hasattr(proj, "weight"): continue W = proj.weight.data delta = (lora_B @ lora_A).to(device=W.device, dtype=W.dtype) if delta.shape == W.shape: W.add_(delta) applied += 1 pipeline._lora_adapters = adapters pipeline.log(f"Applied {applied} LoRA adapters (merged into weights)") def _json_bytes(value: object) -> bytes: return ( json.dumps(value, indent=2, sort_keys=True, ensure_ascii=True, allow_nan=False) + "\n" ).encode("utf-8") def _write_new(path: Path, payload: bytes) -> None: with path.open("xb") as stream: stream.write(payload) stream.flush() os.fsync(stream.fileno()) def _digest_bytes(payload: bytes) -> str: return f"sha256:{sha256(payload).hexdigest()}" def _digest_file(path: Path) -> str: digest = sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(1 << 20), b""): digest.update(chunk) return f"sha256:{digest.hexdigest()}" def _prepare_adapter_state( adapters: dict[str, tuple[torch.Tensor, torch.Tensor]], *, lora_alpha: int | None, ) -> tuple[dict[str, torch.Tensor], list[dict[str, Any]], int, int, float]: if not adapters: raise ValueError("canonical PEFT export requires at least one adapter") ranks = set() for module_name, pair in adapters.items(): if not _MODULE_PATH.fullmatch(module_name): raise ValueError(f"adapter module path is invalid: {module_name!r}") if not isinstance(pair, tuple) or len(pair) != 2: raise ValueError(f"adapter pair is invalid: {module_name}") lora_b, lora_a = pair if ( not isinstance(lora_a, torch.Tensor) or not isinstance(lora_b, torch.Tensor) or lora_a.ndim != 2 or lora_b.ndim != 2 or lora_b.shape[1] != lora_a.shape[0] or lora_a.dtype != lora_b.dtype ): raise ValueError(f"adapter tensor geometry is invalid: {module_name}") ranks.add(int(lora_a.shape[0])) if len(ranks) != 1 or next(iter(ranks)) <= 0: raise ValueError("canonical PEFT export requires one positive rank") rank = next(iter(ranks)) alpha = rank if lora_alpha is None else lora_alpha if type(alpha) is not int or alpha <= 0: raise ValueError("LoRA alpha must be a positive integer") scaling = alpha / rank state: dict[str, torch.Tensor] = {} key_map: list[dict[str, Any]] = [] for module_name in sorted(adapters): lora_b, lora_a = adapters[module_name] a_key = f"base_model.model.{module_name}.lora_A.weight" b_key = f"base_model.model.{module_name}.lora_B.weight" state[a_key] = lora_a.detach().cpu().contiguous() state[b_key] = lora_b.detach().cpu().contiguous() / scaling key_map.append( { "module_name": module_name, "target_module": module_name.rsplit(".", 1)[-1], "lora_A_key": a_key, "lora_B_key": b_key, "lora_A_shape": list(lora_a.shape), "lora_B_shape": list(lora_b.shape), "rank": rank, } ) return state, key_map, rank, alpha, scaling def save_lora_adapters( adapters: dict[str, tuple[torch.Tensor, torch.Tensor]], output_dir: str | Path, *, base_model: BaseModelIdentity, provenance_factory: Callable[[tuple[str, ...], AdapterIdentity], ProvenanceRecord], lora_alpha: int | None = None, lora_dropout: float = 0.0, modules_to_save: tuple[str, ...] = (), adapter_name: str = "default", ) -> AdapterArtifact: """Atomically write a standard PEFT LoRA artifact with exact provenance.""" if not isinstance(base_model, BaseModelIdentity): raise ValueError("canonical PEFT export requires an exact base model identity") if ( isinstance(lora_dropout, bool) or not isinstance(lora_dropout, (int, float)) or not math.isfinite(lora_dropout) or not 0 <= lora_dropout < 1 ): raise ValueError("LoRA dropout must be in [0, 1)") if modules_to_save: raise ValueError("modules_to_save is unsupported unless its tensors are supplied") if ( not isinstance(adapter_name, str) or not adapter_name or len(adapter_name) > 512 or not _PUBLIC_NAME.fullmatch(adapter_name) or Path(adapter_name).is_absolute() or _SECRET_VALUE.search(adapter_name) ): raise ValueError("adapter name is invalid") state, key_map, rank, alpha, scaling = _prepare_adapter_state( adapters, lora_alpha=lora_alpha, ) output_path = Path(output_dir) if ( output_path.is_symlink() or output_path.absolute() != output_path.resolve(strict=False) ): raise ValueError("adapter output directory must not be a symlink") if output_path.exists() and (not output_path.is_dir() or any(output_path.iterdir())): raise FileExistsError("adapter output directory must be absent or empty") config = { "base_model_name_or_path": base_model.repo_id, "bias": "none", "fan_in_fan_out": False, "inference_mode": True, "init_lora_weights": True, "lora_alpha": alpha, "lora_dropout": float(lora_dropout), "modules_to_save": None, "peft_type": "LORA", "r": rank, "revision": base_model.revision, # Full paths prevent PEFT from creating unsupplied adapters on every # module that happens to share a leaf name such as ``q_proj``. "target_modules": sorted(item["module_name"] for item in key_map), "task_type": "CAUSAL_LM", } config_payload = _json_bytes(config) key_map_digest = _digest_bytes( json.dumps(key_map, sort_keys=True, separators=(",", ":")).encode("utf-8") ) adapter_identity = AdapterIdentity( adapter_type="lora", base_model=base_model.to_artifact_identity(), config_digest=_digest_bytes(config_payload), key_map_digest=key_map_digest, ) card = ( "---\n" f"base_model: {base_model.repo_id}\n" f"base_model_revision: {base_model.revision}\n" "library_name: peft\n" "tags:\n - peft\n - lora\n - obliteratus\n" "---\n\n" "# OBLITERATUS LoRA adapter\n\n" "This is an unmerged PEFT LoRA artifact representing an OBLITERATUS surgery event.\n\n" f"- Exact base: `{base_model.repo_id}@{base_model.revision}`\n" f"- Base weights: `{base_model.weights_digest}`\n" f"- Tokenizer: `{base_model.tokenizer_digest}`\n" "- Scaling: standard PEFT `alpha / rank`; saved B values preserve the exact delta.\n" ).encode("utf-8") manifest = { "schema_id": "obliteratus.peft-adapter-manifest", "schema_version": "1.0.0", "adapter_name": adapter_name, "adapter_type": "lora", "adapter_format_version": "peft-lora-v1", "base_model": base_model.to_dict(), "rank": rank, "alpha": alpha, "scaling": scaling, "dropout": float(lora_dropout), "bias": "none", "modules_to_save": [], "target_modules": config["target_modules"], "tie_policy": "base_model_declared", "merged": False, "key_map_digest": key_map_digest, "model_card_digest": _digest_bytes(card), "key_map": key_map, } manifest_payload = _json_bytes(manifest) artifact_id = "" def validate_staging(staging: Path) -> None: expected = { "README.md", "adapter_config.json", "adapter_manifest.json", "adapter_model.safetensors", "checkpoint-provenance.json", } paths = tuple(staging.iterdir()) if {path.name for path in paths} != expected or any( path.is_symlink() or not path.is_file() for path in paths ): raise ValueError("adapter artifact set is incomplete") loaded = load_file(staging / "adapter_model.safetensors", device="cpu") if set(loaded) != set(state): raise ValueError("adapter tensor set does not match the manifest") validate_adapter_base(staging, base_model) with atomic_checkpoint_directory(output_path, validate=validate_staging) as staging: weights_path = staging / "adapter_model.safetensors" config_path = staging / "adapter_config.json" manifest_path = staging / "adapter_manifest.json" provenance_path = staging / "checkpoint-provenance.json" card_path = staging / "README.md" save_file(dict(sorted(state.items())), weights_path) _write_new(config_path, config_payload) _write_new(manifest_path, manifest_payload) _write_new(card_path, card) output_digests = tuple( sorted( ( _digest_file(weights_path), _digest_file(config_path), _digest_file(manifest_path), _digest_file(card_path), ) ) ) provenance = provenance_factory(output_digests, adapter_identity) provenance_record = provenance.to_dict() if ( provenance_record.get("base_model") != base_model.to_artifact_identity().to_dict() or provenance_record.get("adapter") != adapter_identity.to_dict() or sorted(provenance_record.get("output_digests", [])) != list(output_digests) ): raise ValueError("adapter provenance does not match the exact exported artifact") artifact_id = provenance.artifact_id _write_new(provenance_path, provenance.to_json().encode("utf-8")) return AdapterArtifact( artifact_id=artifact_id, weights_path=output_path / "adapter_model.safetensors", config_path=output_path / "adapter_config.json", provenance_path=output_path / "checkpoint-provenance.json", ) def _read_json_object(path: Path, detail: str) -> dict[str, Any]: if path.is_symlink() or not path.is_file() or path.stat().st_size > 8 << 20: raise ValueError(detail) try: value = json.loads( path.read_text(encoding="utf-8"), object_pairs_hook=_reject_duplicate_json_pairs, ) except (UnicodeError, json.JSONDecodeError, ValueError) as error: raise ValueError(detail) from error if not isinstance(value, dict): raise ValueError(detail) return value def _reject_duplicate_json_pairs(pairs: list[tuple[str, Any]]) -> dict[str, Any]: result: dict[str, Any] = {} for key, value in pairs: if key in result: raise ValueError(f"duplicate JSON key: {key}") result[key] = value return result def _canonical_object_digest(value: object) -> str: payload = json.dumps( value, sort_keys=True, separators=(",", ":"), ensure_ascii=True, allow_nan=False, ).encode("utf-8") return _digest_bytes(payload) def _validate_adapter_integrity( root: Path, config: dict[str, Any], manifest: dict[str, Any], provenance: dict[str, Any], ) -> None: weights = root / "adapter_model.safetensors" artifact_paths = ( root / "README.md", root / "adapter_config.json", root / "adapter_manifest.json", weights, root / "checkpoint-provenance.json", ) if any(path.is_symlink() or not path.is_file() for path in artifact_paths): raise ValueError("adapter_artifact_file_invalid") if weights.stat().st_size == 0: raise ValueError("adapter_weights_invalid") try: verified_provenance = verify_provenance_record(provenance) except (TypeError, ValueError) as error: if "record digest mismatch" in str(error): raise ValueError("adapter_provenance_digest_mismatch") from error if "artifact ID mismatch" in str(error): raise ValueError("adapter_artifact_identity_mismatch") from error raise ValueError("adapter_provenance_invalid") from error if verified_provenance != provenance: raise ValueError("adapter_provenance_invalid") expected_outputs = provenance.get("output_digests") if ( not isinstance(expected_outputs, list) or any(not isinstance(value, str) or not _DIGEST.fullmatch(value) for value in expected_outputs) or sorted(expected_outputs) != sorted( _digest_file(path) for path in ( root / "README.md", root / "adapter_config.json", root / "adapter_manifest.json", weights, ) ) ): raise ValueError("adapter_artifact_digest_mismatch") record_digest = provenance.get("record_digest") record_core = {key: value for key, value in provenance.items() if key != "record_digest"} if record_digest != _canonical_object_digest(record_core): raise ValueError("adapter_provenance_digest_mismatch") artifact_id = provenance.get("artifact_id") identity_core = {key: value for key, value in record_core.items() if key != "artifact_id"} expected_artifact_id = _canonical_object_digest(identity_core).replace( "sha256:", "artifact-sha256:", 1, ) if artifact_id != expected_artifact_id: raise ValueError("adapter_artifact_identity_mismatch") key_map = manifest.get("key_map") adapter = provenance.get("adapter") manifest_base = manifest.get("base_model") if ( not isinstance(key_map, list) or not isinstance(adapter, dict) or not isinstance(manifest_base, dict) ): raise ValueError("adapter_manifest_integrity_invalid") if any( not isinstance(item, dict) or not isinstance(item.get("module_name"), str) or not _MODULE_PATH.fullmatch(item["module_name"]) for item in key_map ): raise ValueError("adapter_manifest_integrity_invalid") key_map_digest = _canonical_object_digest(key_map) target_modules = [ item.get("module_name") for item in key_map if isinstance(item, dict) ] rank = manifest.get("rank") alpha = manifest.get("alpha") scaling = manifest.get("scaling") if ( manifest.get("key_map_digest") != key_map_digest or manifest.get("model_card_digest") != _digest_file(root / "README.md") or adapter.get("key_map_digest") != key_map_digest or adapter.get("config_digest") != _digest_file(root / "adapter_config.json") or config.get("base_model_name_or_path") != manifest_base.get("repo_id") or config.get("revision") != manifest_base.get("revision") or manifest.get("schema_id") != "obliteratus.peft-adapter-manifest" or manifest.get("schema_version") != "1.0.0" or config.get("peft_type") != "LORA" or config.get("target_modules") != sorted(target_modules) or manifest.get("target_modules") != sorted(target_modules) or type(rank) is not int or rank <= 0 or config.get("r") != rank or type(alpha) is not int or alpha <= 0 or config.get("lora_alpha") != alpha or not isinstance(scaling, (int, float)) or isinstance(scaling, bool) or not math.isfinite(scaling) or scaling != alpha / rank or config.get("lora_dropout") != manifest.get("dropout") or config.get("bias") != manifest.get("bias") ): raise ValueError("adapter_manifest_integrity_invalid") def validate_adapter_base( adapter_dir: str | Path, base_model: BaseModelIdentity, ) -> dict[str, Any]: """Fail before loading weights when exact base/tokenizer facts do not match.""" root = Path(adapter_dir) if ( root.is_symlink() or not root.is_dir() or root.absolute() != root.resolve(strict=True) ): raise ValueError("adapter_directory_invalid") config = _read_json_object(root / "adapter_config.json", "adapter_config_invalid") manifest = _read_json_object(root / "adapter_manifest.json", "adapter_manifest_invalid") provenance = _read_json_object( root / "checkpoint-provenance.json", "adapter_provenance_invalid", ) _validate_adapter_integrity(root, config, manifest, provenance) checks = ( (config.get("base_model_name_or_path"), base_model.repo_id, "base_model_identity_mismatch"), (config.get("revision"), base_model.revision, "base_model_revision_mismatch"), ( manifest.get("base_model", {}).get("weights_digest"), base_model.weights_digest, "base_model_digest_mismatch", ), ( manifest.get("base_model", {}).get("tokenizer_digest"), base_model.tokenizer_digest, "tokenizer_digest_mismatch", ), ( manifest.get("base_model", {}).get("vocab_size"), base_model.vocab_size, "vocab_size_mismatch", ), ( manifest.get("base_model", {}).get("architecture"), base_model.architecture, "architecture_mismatch", ), ( manifest.get("base_model", {}).get("tied_embeddings"), base_model.tied_embeddings, "tied_embeddings_mismatch", ), ) for actual, expected, detail in checks: if actual != expected: raise ValueError(detail) if provenance.get("base_model") != base_model.to_artifact_identity().to_dict(): raise ValueError("adapter_provenance_base_mismatch") return manifest def load_lora_adapters( adapter_dir: str | Path, *, base_model: BaseModelIdentity, ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: """Load canonical safetensors only; this path never accepts pickle artifacts.""" root = Path(adapter_dir) manifest = validate_adapter_base(root, base_model) state = load_file(root / "adapter_model.safetensors", device="cpu") validate_adapter_base(root, base_model) scaling = manifest.get("scaling") if ( isinstance(scaling, bool) or not isinstance(scaling, (int, float)) or not math.isfinite(scaling) or scaling <= 0 ): raise ValueError("adapter_scaling_invalid") result: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} expected_keys: set[str] = set() key_map = manifest.get("key_map") if not isinstance(key_map, list): raise ValueError("adapter_key_map_invalid") for item in key_map: if not isinstance(item, dict): raise ValueError("adapter_key_map_invalid") module_name = item.get("module_name") a_key = item.get("lora_A_key") b_key = item.get("lora_B_key") if ( not isinstance(module_name, str) or not isinstance(a_key, str) or not isinstance(b_key, str) or not _MODULE_PATH.fullmatch(module_name) or module_name in result or a_key == b_key or a_key in expected_keys or b_key in expected_keys or a_key not in state or b_key not in state ): raise ValueError("adapter_key_map_invalid") lora_a = state[a_key] lora_b = state[b_key] if ( lora_a.ndim != 2 or lora_b.ndim != 2 or lora_b.shape[1] != lora_a.shape[0] or lora_a.dtype != lora_b.dtype or item.get("lora_A_shape") != list(lora_a.shape) or item.get("lora_B_shape") != list(lora_b.shape) or item.get("rank") != lora_a.shape[0] ): raise ValueError("adapter_tensor_geometry_mismatch") expected_keys.update((a_key, b_key)) result[module_name] = (lora_b * scaling, lora_a) if expected_keys != set(state): raise ValueError("adapter_tensor_set_mismatch") return result def save_unsupported_obliteratus_adapters( adapters: dict[str, tuple[torch.Tensor, torch.Tensor]], output_dir: str | Path, *, reason: str, ) -> Path: """Persist safe legacy tensors without making a PEFT compatibility claim.""" if not adapters: raise ValueError("legacy adapter export requires at least one adapter") if ( not isinstance(reason, str) or not reason or len(reason) > 1024 or _SECRET_VALUE.search(reason) or Path(reason).is_absolute() ): raise ValueError("legacy adapter reason is invalid") output = Path(output_dir) if output.is_symlink() or output.absolute() != output.resolve(strict=False): raise ValueError("adapter output directory must not be a symlink") output.mkdir(parents=True, exist_ok=True) state = {} key_map = [] for module_name in sorted(adapters): lora_b, lora_a = adapters[module_name] b_key = f"{module_name}.B" a_key = f"{module_name}.A" state[b_key] = lora_b.detach().cpu().contiguous() state[a_key] = lora_a.detach().cpu().contiguous() key_map.append({"module_name": module_name, "A": a_key, "B": b_key}) path = output / "obliteratus_unsupported_adapter.safetensors" config_path = output / "obliteratus_unsupported_adapter.json" if path.exists() or config_path.exists(): raise FileExistsError("unsupported adapter artifact already exists") save_file(dict(sorted(state.items())), path) _write_new( config_path, _json_bytes( { "schema_id": "obliteratus.unsupported-adapter", "schema_version": "1.0.0", "support_status": "unsupported_legacy", "safe_serialization": True, "peft_compatible": False, "reason": reason, "key_map": key_map, } ), ) return path def save_legacy_pickle_adapters_trusted( adapters: dict[str, tuple[torch.Tensor, torch.Tensor]], output_dir: str | Path, *, allow_pickle: bool, ) -> Path: """Write the historical pickle shape only behind an explicit trust gate.""" if allow_pickle is not True: raise PermissionError("allow_pickle=True is required for unsafe legacy export") output = Path(output_dir) output.mkdir(parents=True, exist_ok=True) path = output / "obliteratus_legacy_adapter_unsafe.pt" state = { f"{module_name}.{suffix}": tensor for module_name, (lora_b, lora_a) in sorted(adapters.items()) for suffix, tensor in (("lora_B", lora_b), ("lora_A", lora_a)) } torch.save(state, path) return path