U4 step3: dim_batch 16->8 to survive OOM contention with user's VS Code kernel

Run 551 was OOM-killed at n_done=36: the user's VS Code Jupyter kernel (jsteer venv, PID 3214401) co-loaded ~1.5GB VRAM + 1.9GB RAM while the fit sat at the 22.4/24.6GB ceiling. Clean SIGKILL with no CUDA traceback = host OOM killer, not a CUDA OOM. dim_batch=8 halves the fit's peak footprint; it changes only the backward schedule, not the accumulated Jacobian, so U4 exactness is preserved. Resumes from checkpoint (n_done=36), lossless. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
2026-09-10 19:20:20 +08:00 · 2026-07-10 13:56:17 +08:00
parent beec9189da
commit ebb912452d
6 changed files with 8 additions and 1 deletions
@@ -0,0 +1,77 @@
+"""U1 parity gate: cached-Jacobian pullback == direct VJP, per layer.
+
+(authored by Claude)
+
+The two paths are linear-identical:  mean_p(J_p)^T w == mean_p(J_p^T w).
+Path A pulls the word cotangent through the CACHED pooled Jacobian.
+Path B contracts the same cotangent inside per-prompt backward passes.
+The only expected gap is fp16 storage in the cache, so per-layer cosine must
+exceed 0.999. A failure is a WIRING bug (layer index, position mask, pooling),
+not a threshold to tune.
+
+Run:
+    uv run python scripts/parity_u1.py 2>&1 | tee /tmp/claude-1000/jsteer_parity.log
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+from loguru import logger
+from tabulate import tabulate
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from jsteer import Jacobian, word_vector_vjp
+
+# Claude: repo root on path so `scripts.smoke` imports whether run as a file or -m.
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from scripts.smoke import CACHE, DEVICE, DTYPE, MODEL, PROMPTS  # same inputs  # noqa: E402
+
+WORDS = ["happy", "joy"]
+OUT = "artifacts/parity_u1.txt"
+
+
+def _unit_dir(vec, layer: int) -> torch.Tensor:
+    """Pull the per-layer unit direction out of a steering_lite Vector."""
+    return vec.stacked[layer]["v"].squeeze(0).float().cpu()  # [d]
+
+
+def main() -> None:
+    logger.info(f"loading {MODEL} ({DTYPE}) on {DEVICE}")
+    tok = AutoTokenizer.from_pretrained(MODEL)
+    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=DTYPE).to(DEVICE).eval()
+
+    jac = Jacobian.load(CACHE)
+    layers = jac.layers  # exact int layers fitted by the smoke
+    logger.info(f"cached layers={layers}")
+
+    # Path A: cached pooled Jacobian pullback.
+    vA = jac.word_vector(model, tok, WORDS)
+    # Path B: direct per-prompt VJP over the SAME prompts / layers / skip_first / max_length.
+    vB = word_vector_vjp(model, tok, PROMPTS, WORDS, layers=layers, max_length=128)
+
+    rows = []
+    for l in layers:
+        a, b = _unit_dir(vA, l), _unit_dir(vB, l)
+        cos = float(torch.dot(a, b) / (a.norm() * b.norm()))
+        rows.append((l, cos, float(a.norm()), float(b.norm())))
+
+    table = tabulate(rows, headers=["layer", "cos", "|vA|", "|vB|"],
+                     tablefmt="pipe", floatfmt="+.6f")
+    min_cos = min(r[1] for r in rows)
+    gate = "PASS" if min_cos > 0.999 else "FAIL"
+    report = f"{table}\n\nmin cos = {min_cos:+.6f}   GATE (>0.999): {gate}\n"
+
+    logger.info("U1 parity table:\n" + report)
+    with open(OUT, "w") as f:
+        f.write("U1 parity: cached-Jacobian pullback vs direct VJP (word=happy/joy)\n")
+        f.write(f"model={MODEL}  prompts={len(PROMPTS)}  layers={layers}\n\n")
+        f.write(report)
+    logger.info(f"wrote {OUT}")
+    if gate == "FAIL":
+        raise SystemExit(f"U1 parity FAILED: min cos={min_cos:.6f} <= 0.999 (wiring bug)")
+
+
+if __name__ == "__main__":
+    main()
@@ -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 <think>. 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.")
@@ -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")