diff --git a/scripts/u4_step1_ref524.py b/scripts/u4_step1_ref524.py new file mode 100644 index 0000000..08b9f57 --- /dev/null +++ b/scripts/u4_step1_ref524.py @@ -0,0 +1,56 @@ +"""U4 step 1/3: regenerate run-524's verified jacobian_word vector. + +(Claude) Run this under j-steer-dev's venv, NOT jsteer's: + + cd ../j-steer-dev && uv run python ../jsteer/scripts/u4_step1_ref524.py + +There `import jsteer` resolves to the OLD experiment package (j-steer-dev/src), +whose extract_word_pullback produced the verified 3/5 result. Run 524 never +persisted its vector tensors (only eval JSONs), but the extraction is +deterministic (seed-0 prompts, greedy, no sampling), so re-running it IS the +reference. Also dumps the 512 substrate prompts so steps 2/3 consume this one +artifact instead of regenerating them (no drift axis). + +Exact run-524 parameters: Qwen/Qwen3-4B, persona=authority, n_pairs=256, +seed=0, layers "mid" (7..27 of 36), words authority/obey/command/hierarchy, +batch_size=4, max_length=384, cotangent_scope=source_scope=all_valid (defaults). +""" +import json +from pathlib import Path + +import torch +from loguru import logger +from steering_lite.data import PERSONA_REGISTRY, make_persona_pairs +from transformers import AutoModelForCausalLM, AutoTokenizer + +from jsteer.pullback import extract_word_pullback # OLD package (j-steer-dev/src) + +ART = Path(__file__).resolve().parent.parent / "artifacts" +MODEL = "Qwen/Qwen3-4B" +WORDS = ["authority", "obey", "command", "hierarchy"] + +tok = AutoTokenizer.from_pretrained(MODEL) +model = AutoModelForCausalLM.from_pretrained( + MODEL, torch_dtype=torch.bfloat16).to("cuda").eval() + +n = model.config.num_hidden_layers +assert n == 36, f"expected Qwen3-4B with 36 layers, got {n}" +layers = tuple(range(max(2, int(n * 0.2)), min(n - 2, int(n * 0.8)))) # run_sweep "mid" -> 7..27 + +persona_pairs, template = PERSONA_REGISTRY["authority"] +pos, neg = make_persona_pairs(tok, n_pairs=256, thinking=True, + persona_pairs=persona_pairs, template=template, seed=0) +prompts = pos + neg # run_sweep feeds pos+neg as the linearization substrate +(ART / "u4_prompts.json").write_text(json.dumps( + {"model": MODEL, "layers": list(layers), "words": WORDS, "prompts": prompts})) +logger.info(f"dumped {len(prompts)} prompts, layers={layers}") +logger.info("SHOULD: chat-templated authority-persona prompt with . ELSE template drift.\n" + f"--- PROMPT[0] (full, special tokens) ---\n{prompts[0]}") + +vec = extract_word_pullback(model, tok, prompts, layers, WORDS, + batch_size=4, max_length=384)["jacobian_word"] +ref = {str(l): vec.stacked[l]["v"].squeeze(0).float().cpu() for l in layers} +torch.save(ref, ART / "u4_ref_524.pt") +logger.info(f"saved {ART / 'u4_ref_524.pt'} " + f"norms={[round(ref[str(l)].norm().item(), 3) for l in layers][:5]}... " + "SHOULD: all 1.0 (unit vectors). ELSE _to_vector changed.") diff --git a/scripts/u4_step2_vjp.py b/scripts/u4_step2_vjp.py new file mode 100644 index 0000000..9c65664 --- /dev/null +++ b/scripts/u4_step2_vjp.py @@ -0,0 +1,46 @@ +"""U4 step 2/3: port check -- jsteer's word_vector_vjp vs the run-524 reference. + +(Claude) Runs in jsteer's venv on the step-1 artifacts. Both sides are fp32 +direct-VJP extractions of the same math on the same 512 prompts; the only +difference is code lineage (old experiment package vs this library) plus +batch-order fp noise. GATE: cos > 0.999 per layer. A failure is a PORT BUG +(position mask, pooling, layer indexing) -- debug, do not tune. +""" +import json +from pathlib import Path + +import torch +from loguru import logger +from tabulate import tabulate +from transformers import AutoModelForCausalLM, AutoTokenizer + +from jsteer import word_vector_vjp + +ART = Path(__file__).resolve().parent.parent / "artifacts" +meta = json.loads((ART / "u4_prompts.json").read_text()) +ref = torch.load(ART / "u4_ref_524.pt") + +tok = AutoTokenizer.from_pretrained(meta["model"]) +model = AutoModelForCausalLM.from_pretrained( + meta["model"], torch_dtype=torch.bfloat16).to("cuda").eval() + +v = word_vector_vjp(model, tok, meta["prompts"], meta["words"], + layers=meta["layers"], batch_size=4, max_length=384) +torch.save({str(l): v.stacked[l]["v"].squeeze(0).float().cpu() for l in meta["layers"]}, + ART / "u4_vjp.pt") + +rows = [] +for l in meta["layers"]: + a = v.stacked[l]["v"].squeeze(0).float() + b = ref[str(l)].float() + rows.append((l, torch.nn.functional.cosine_similarity(a, b, dim=0).item())) +table = tabulate(rows, headers=["layer", "cos(jsteer_vjp, ref524)"], floatfmt="+.6f") +min_cos = min(c for _, c in rows) +verdict = "PASS" if min_cos > 0.999 else "FAIL" +out = (f"U4 step 2: jsteer word_vector_vjp vs regenerated run-524 vector\n" + f"model={meta['model']} prompts={len(meta['prompts'])} words={meta['words']}\n\n" + f"{table}\n\nmin cos = {min_cos:+.6f} GATE (>0.999): {verdict}\n") +(ART / "u4_step2_vjp_parity.txt").write_text(out) +print(out) +if verdict == "FAIL": + raise SystemExit("U4 step 2 FAILED: port bug, do not run step 3 until root-caused") diff --git a/scripts/u4_step3_fit4b.py b/scripts/u4_step3_fit4b.py new file mode 100644 index 0000000..2e606cd --- /dev/null +++ b/scripts/u4_step3_fit4b.py @@ -0,0 +1,59 @@ +"""U4 step 3/3: full 4B Jacobian fit on run-524's substrate + cache loop-close. + +(Claude) The expensive one: 512 prompts x ceil(2560/dim_batch) backwards. +checkpoint_path makes it resumable, so a kill/OOM loses at most one prompt. +After fitting, the cached word vector must match BOTH the step-2 jsteer VJP +vector and the step-1 run-524 reference (linearity: mean_p(J_p)^T w = +mean_p(J_p^T w); fp16 cache storage is the only gap). GATE: cos > 0.999. + +This closes the loop on the verified 3/5 moral-foundations result: the library +artifact (artifacts/qwen3-4b-authority.jac) provably contains the verified +steering vector. +""" +import json +import time +from pathlib import Path + +import torch +from loguru import logger +from tabulate import tabulate +from transformers import AutoModelForCausalLM, AutoTokenizer + +from jsteer import Jacobian + +ART = Path(__file__).resolve().parent.parent / "artifacts" +meta = json.loads((ART / "u4_prompts.json").read_text()) +ref524 = torch.load(ART / "u4_ref_524.pt") +vjp = torch.load(ART / "u4_vjp.pt") + +tok = AutoTokenizer.from_pretrained(meta["model"]) +model = AutoModelForCausalLM.from_pretrained( + meta["model"], torch_dtype=torch.bfloat16).to("cuda").eval() + +t0 = time.time() +jac = Jacobian.fit(model, tok, meta["prompts"], layers=meta["layers"], + dim_batch=16, max_seq_len=384, + checkpoint_path=str(ART / "qwen3-4b-authority.ckpt")) +logger.info(f"fit wall-time: {(time.time() - t0) / 3600:.2f} h") +jac.save(str(ART / "qwen3-4b-authority.jac")) +logger.info(f"saved cache: {(ART / 'qwen3-4b-authority.jac').stat().st_size / 1e9:.2f} GB") + +v = jac.word_vector(model, tok, meta["words"], layers=meta["layers"]) +rows = [] +for l in meta["layers"]: + a = v.stacked[l]["v"].squeeze(0).float() + rows.append((l, + torch.nn.functional.cosine_similarity(a, vjp[str(l)].float(), dim=0).item(), + torch.nn.functional.cosine_similarity(a, ref524[str(l)].float(), dim=0).item())) +table = tabulate(rows, headers=["layer", "cos(cache, jsteer_vjp)", "cos(cache, ref524)"], + floatfmt="+.6f") +min_cos = min(min(r[1], r[2]) for r in rows) +verdict = "PASS" if min_cos > 0.999 else "FAIL" +out = (f"U4 step 3: cached-4B word vector vs step-2 VJP and run-524 reference\n" + f"model={meta['model']} prompts={len(meta['prompts'])} words={meta['words']} " + f"dim_batch=16 fp16-cache\n\n{table}\n\n" + f"min cos = {min_cos:+.6f} GATE (>0.999): {verdict}\n") +(ART / "u4_loopclose.txt").write_text(out) +print(out) +if verdict == "FAIL": + raise SystemExit("U4 step 3 FAILED: cache wiring bug, root-cause before shipping the artifact")