File size: 3,489 Bytes
8018ecd | 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 | """Apply Qwen/Qwen3.5-4B WSP v2.5 safetensors patches (factored r16 subspaces or dense delta)
to a loaded Hugging Face Qwen3.5-4B model in-memory.
Supports Variant 3 regime presets cleanly:
- general_fast_think : scales=[1.95, 0.0, 0.0, 0.0, 0.0, 0.0]
- mid_code_synthesis : scales=[0.65, 0.0, 0.0, 0.65, 0.0, 0.0]
- tool_routing : scales=[0.70, 0.0, 0.95, 0.0, 0.0, 0.0]
- turn2_json : scales=[0.70, 0.40, 0.0, 0.0, 1.0, 0.0]
- restraint_pushback : scales=[0.70, 1.0, 0.0, 0.0, 0.0, 0.0]
- escalate_handoff : scales=[0.70, 0.0, 0.0, 0.0, 0.0, 1.0]
"""
from __future__ import annotations
import json
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",
]
REGIME_PRESETS: Dict[str, List[float]] = {
"general_fast_think": [1.95, 0.0, 0.0, 0.0, 0.0, 0.0],
"mid_code_synthesis": [0.65, 0.0, 0.0, 0.65, 0.0, 0.0],
"tool_routing": [0.70, 0.0, 0.95, 0.0, 0.0, 0.0],
"turn2_json_synthesis": [0.70, 0.40, 0.0, 0.0, 1.0, 0.0],
"restraint_pushback": [0.70, 1.0, 0.0, 0.0, 0.0, 0.0],
"handoff_gate": [0.70, 0.0, 0.0, 0.0, 0.0, 1.0],
"static_all_on_v2": [0.65, 0.65, 0.65, 0.65, 0.65, 0.65],
}
def apply_subspace_scales_to_hf_model(
model: torch.nn.Module,
patches_dir: Path | str,
scales: Sequence[float],
use_s4_oproj: bool = False,
) -> int:
"""Add sum_i scale[i] * (lora_b @ lora_a) directly into HF model weights in-place."""
if len(scales) != 6:
raise ValueError(f"Expected 6 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
target_file = pdir / ("wsp_s4_mid_code_oproj_r16.safetensors" if (i == 3 and use_s4_oproj) else fname)
tensors = load_file(str(target_file), device="cpu")
prefixes = sorted({k[:-7] for k in tensors if k.endswith(".lora_a")})
for hf_key in prefixes:
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)
delta = sc * (lb @ la)
param.add_(delta.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_qwen35_4b_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")]
param = params[hf_key]
param.add_((scale * delta.to(device=param.device, dtype=torch.float32)).to(dtype=param.dtype))
updated += 1
return updated
|