File size: 3,882 Bytes
b308044 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """Apply WSP v3.0 (9-Subspace) safetensors patches (factored r16 subspaces or combined dense delta)
to a loaded Hugging Face model in-memory.
Works uniformly across all 5 WSP v3.0 candidate models:
- nemotron3_4b (nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16)
- qwen35_4b (Qwen/Qwen3.5-4B)
- qwen35_2b (Qwen/Qwen3.5-2B)
- qwen35_9b (Qwen/Qwen3.5-9B)
- gemma4_12b (google/gemma-4-12B-it)
"""
from __future__ import annotations
from pathlib import Path
from typing import Dict, List, Sequence
import torch
from safetensors.torch import load_file
SUBSPACE_PATCH_FILES = [
"wsp_s1_think_r16.safetensors",
"wsp_s2_restraint_r16.safetensors",
"wsp_s3_router_r16.safetensors",
"wsp_s4_mid_code_r16.safetensors",
"wsp_s5_json_schema_r16.safetensors",
"wsp_s6_handoff_r16.safetensors",
"wsp_s7_injection_guard_r16.safetensors",
"wsp_s8_grounding_r16.safetensors",
"wsp_s9_error_recovery_r16.safetensors",
]
# Representative regime scale vectors [s1, s2, s3, s4, s5, s6, s7, s8, s9] matching dynamic_router_9s.py
REGIME_PRESETS: Dict[str, List[float]] = {
"general_fast_think": [2.15, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
"mid_code_synthesis": [0.95, 0.0, 0.0, 0.85, 0.0, 0.0, 0.0, 0.0, 0.0],
"tool_routing": [1.05, 0.0, 0.95, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
"restraint_pushback": [1.15, 1.05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
"turn2_json_synthesis": [0.75, 0.25, 0.0, 0.0, 1.15, 0.0, 0.0, 0.0, 0.0],
"handoff_gate": [0.65, 0.20, 0.0, 0.0, 0.0, 1.15, 0.0, 0.0, 0.0],
"injection_guard": [0.85, 0.55, 0.0, 0.0, 0.0, 0.0, 1.25, 0.0, 0.0],
"grounding_abstain": [0.85, 0.35, 0.0, 0.0, 0.35, 0.0, 0.0, 1.20, 0.0],
"error_recovery": [0.80, 0.0, 0.45, 0.0, 0.0, 0.0, 0.0, 0.0, 1.20],
}
def apply_subspace_scales_to_hf_model(
model: torch.nn.Module,
patches_dir: Path | str,
scales: Sequence[float],
) -> int:
"""Add sum_i scale[i] * (lora_b @ lora_a) directly into Hugging Face model weights in-place."""
if len(scales) != 9:
raise ValueError(f"Expected 9 subspace scales, got {len(scales)}")
pdir = Path(patches_dir)
params = dict(model.named_parameters())
updated_tensors = 0
with torch.no_grad():
for i, fname in enumerate(SUBSPACE_PATCH_FILES):
sc = float(scales[i])
if abs(sc) < 1e-9:
continue
tensors = load_file(str(pdir / fname), device="cpu")
prefixes = sorted({k[:-7] for k in tensors if k.endswith(".lora_a")})
for hf_key in prefixes:
if hf_key not in params:
raise KeyError(f"Parameter {hf_key} not found in Hugging Face model")
param = params[hf_key]
la = tensors[f"{hf_key}.lora_a"].to(device=param.device, dtype=torch.float32)
lb = tensors[f"{hf_key}.lora_b"].to(device=param.device, dtype=torch.float32)
param.add_((sc * (lb @ la)).to(dtype=param.dtype))
updated_tensors += 1
return updated_tensors
def apply_dense_patch_to_hf_model(
model: torch.nn.Module,
dense_patch_path: Path | str,
scale: float = 1.0,
) -> int:
"""Add scale * delta tensors from wsp_<slug>_patch.safetensors directly in-place."""
tensors = load_file(str(dense_patch_path), device="cpu")
params = dict(model.named_parameters())
updated = 0
with torch.no_grad():
for key, delta in tensors.items():
if not key.endswith(".delta"):
continue
hf_key = key[: -len(".delta")]
if hf_key not in params:
raise KeyError(f"Parameter {hf_key} not found in Hugging Face model")
param = params[hf_key]
param.add_((scale * delta.to(device=param.device, dtype=torch.float32)).to(dtype=param.dtype))
updated += 1
return updated
|