mirror of
https://github.com/wassname/jsteer.git
synced 2026-09-09 11:25:03 +08:00
scratch: move u4_step3 loop-close scripts (fit4b/guard/retry) out of scripts/ top
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>
This commit is contained in:
wassname
@@ -0,0 +1,68 @@
|
||||
"""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")
|
||||
Executable
+46
@@ -0,0 +1,46 @@
|
||||
#!/usr/bin/env bash
|
||||
# Claude: external requeue guard for the U4 4B Jacobian fit.
|
||||
#
|
||||
# WHY not just the in-task retry wrapper: systemd-oomd is active on this box and
|
||||
# kills the fit's ENTIRE pueue task-scope (bash + python together) under memory
|
||||
# pressure from the user's live VS Code GPU kernel -- so the in-cgroup retry in
|
||||
# u4_step3_retry.sh dies with it (observed: task 554 Killed whole, no retry line).
|
||||
#
|
||||
# This guard runs DETACHED in its own session/cgroup, uses ~0 RAM (so oomd never
|
||||
# targets it), and launches the resumable fit ONLY into genuinely-idle GPU/RAM
|
||||
# windows (see backoff below) so it never competes with the user's live work.
|
||||
# jlens checkpoints every prompt (n_done monotonic: 36->45->64->69 across kills),
|
||||
# so any kill only costs a model reload. Completes clean overnight once the
|
||||
# user's kernel idles; does nothing (polite) while they are active.
|
||||
#
|
||||
# Stops itself the instant artifacts/u4_loopclose.txt exists (success). To stop
|
||||
# early: pkill -f u4_step3_guard.sh
|
||||
set -u
|
||||
cd /media/wassname/SGIronWolf/projects5/2026/jspace/jsteer
|
||||
LABEL='U4 step3 guard-requeue'
|
||||
# POLITE backoff: the user is developing jsteer live (their VS Code GPU kernel is
|
||||
# the oomd competitor). Only launch the fit when the GPU/host is genuinely idle,
|
||||
# so we NEVER compete with the user's interactive work -- we just fill the empty
|
||||
# windows (overnight, breaks). This wastes no idle GPU time yet never fights the
|
||||
# human for their own machine. Thresholds sized for the fit's ~10.5GB footprint.
|
||||
GPU_FREE_MIN_MIB=13000 # need >=13GB free (i.e. others using <~11GB) to start
|
||||
RAM_AVAIL_MIN_MIB=20000 # need >=20GB host available (oomd triggers on host too)
|
||||
while [ ! -f artifacts/u4_loopclose.txt ]; do
|
||||
active=$(pueue status --json | jq -r --arg l "$LABEL" \
|
||||
'.tasks[] | select((.label // "") | contains($l)) | .status | if type=="object" then keys[0] else . end' \
|
||||
| grep -Ec 'Running|Queued|Stashed|Paused')
|
||||
if [ "$active" -eq 0 ]; then
|
||||
gpu_free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits | head -1)
|
||||
ram_avail=$(free -m | awk '/^Mem:/{print $7}')
|
||||
if [ "$gpu_free" -ge "$GPU_FREE_MIN_MIB" ] && [ "$ram_avail" -ge "$RAM_AVAIL_MIN_MIB" ]; then
|
||||
pueue add \
|
||||
-l "why: $LABEL -- polite idle-window fit of the 4B Jacobian; resolve: resume checkpoint to n_done=400, writes u4_loopclose.txt on success" \
|
||||
-w "$PWD" -o 0 -- bash scripts/u4_step3_retry.sh
|
||||
echo "$(date '+%F %H:%M:%S') GPU idle (free=${gpu_free}MiB, ram=${ram_avail}MiB) -- launched fit"
|
||||
else
|
||||
echo "$(date '+%F %H:%M:%S') user active (gpu_free=${gpu_free}MiB, ram_avail=${ram_avail}MiB) -- backing off"
|
||||
fi
|
||||
fi
|
||||
sleep 300
|
||||
done
|
||||
echo "$(date '+%F %H:%M:%S') u4_loopclose.txt exists -- fit complete, guard exiting"
|
||||
Executable
+35
@@ -0,0 +1,35 @@
|
||||
#!/usr/bin/env bash
|
||||
# Claude: auto-resume wrapper for the 4B Jacobian fit (U4 step 3).
|
||||
#
|
||||
# The 3090 is shared with the user's live VS Code Jupyter kernel, whose bursty
|
||||
# GPU/RAM load repeatedly OOM-killed this long-lived fit:
|
||||
# 551 host-OOM (SIGKILL, no traceback) at n_done=36 (kernel ~1.5GB)
|
||||
# 552 CUDA-OOM (traceback, exit 1) at n_done=45 (kernel grew to 8.18GB)
|
||||
# 553 host-OOM (SIGKILL, no traceback) at n_done=64
|
||||
# jlens checkpoints every prompt and resumes, so the accumulated Jacobian is
|
||||
# never lost -- each kill only costs one model reload. So instead of predicting
|
||||
# the user's bursts, just relaunch until the fit exits 0 (success writes
|
||||
# artifacts/u4_loopclose.txt). expandable_segments defrags VRAM; dim_batch=4
|
||||
# (in the .py) keeps this a polite ~10.5GB co-tenant.
|
||||
#
|
||||
# MAX_RETRIES caps runaway: with ~15-20 prompts/run and 336 remaining, an
|
||||
# all-bursty night needs <25 relaunches; 40 is safe headroom. If n_done stops
|
||||
# advancing across retries (logged below), that's a real bug, not an OOM --
|
||||
# stop and debug rather than spin.
|
||||
set -u
|
||||
cd "$(dirname "$0")/.."
|
||||
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
CKPT=artifacts/qwen3-4b-authority.ckpt
|
||||
MAX_RETRIES=40
|
||||
n=0
|
||||
ndone() { .venv/bin/python -c "import torch;print(torch.load('$CKPT',map_location='cpu',weights_only=False)['n_done'])" 2>/dev/null; }
|
||||
until .venv/bin/python scripts/u4_step3_fit4b.py; do
|
||||
n=$((n + 1))
|
||||
echo "[retry-wrapper] fit exited non-zero (OOM?); attempt $n/$MAX_RETRIES, n_done=$(ndone), resuming in 90s"
|
||||
if [ "$n" -ge "$MAX_RETRIES" ]; then
|
||||
echo "[retry-wrapper] hit MAX_RETRIES=$MAX_RETRIES at n_done=$(ndone) -- likely a real bug, not OOM. Stopping."
|
||||
exit 1
|
||||
fi
|
||||
sleep 90
|
||||
done
|
||||
echo "[retry-wrapper] fit completed after $n retries, n_done=$(ndone)"
|
||||
Reference in New Issue
Block a user