mirror of
https://github.com/wassname/jsteer.git
synced 2026-09-09 11:25:03 +08:00
wassnameandClaudypoo
69bd650dc5
The U4 loop-close is a separate finished-enough goal from the demo; guard killed. scripts/ top level is now just fit.py + smoke.py. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
69 lines
3.2 KiB
Python
69 lines
3.2 KiB
Python
"""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()
|
|
# dim_batch 16 -> 4 (Claude): two OOMs vs the user's live VS Code GPU kernel.
|
|
# 551 host-OOM-killed at n_done=36 (kernel ~1.5GB); 552 CUDA-OOM at n_done=45
|
|
# once the kernel grew to 8.18GB and this fit's 13.23GB hit the 23.5GB ceiling
|
|
# with only 44MB free (fragmentation ate the last margin). dim_batch=4 drops
|
|
# this fit to ~10.5GB so it is a polite co-tenant (leaves the user ~13GB); run
|
|
# under PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (the OOM's own
|
|
# suggestion) to defragment. dim_batch changes only the backward SCHEDULE
|
|
# (4x passes), NOT the accumulated Jacobian, so U4 exactness holds. Resumes
|
|
# from the checkpoint (n_done=45), lossless.
|
|
jac = Jacobian.fit(model, tok, meta["prompts"], layers=meta["layers"],
|
|
dim_batch=4, 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")
|