mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-08-11 11:18:24 +08:00
style(train): voice pass 2 -- section banners + finish math sweep through main()
Behavior-preserving (smoke + smoke-route2 exit 0, metrics identical, route2 ‖δS_hack‖=0.0079>0). route2/delta_S_hack/hk_abl logic untouched (comments only in that block; code identifiers left exactly as-is). - 13 `# ── section ──` banners marking main()'s phases: model/tokenizer, AntiPaSTO adapter, hack direction, teacher pool, optimizer, generation config, training loop, per-prompt rollouts, inject->project/route, refresh, deploy-eval, final eval. - Prose δS / τ throughout main()'s comments (code dict-keys "delta_S" unchanged). - Trimmed duplicated/verbose blocks (per-step table legend dup, no-checkpointing essay) and the last war-story (run-43 divergence anecdote). Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
+75
-95
@@ -390,7 +390,7 @@ def load_v_hack(
|
||||
vhack_keys = set(v_hack)
|
||||
missing = sorted(wrapper_keys - vhack_keys)
|
||||
extra = sorted(vhack_keys - wrapper_keys)
|
||||
# v_hack[name] is [k_max, r]; delta_S is [r]. Check last-dim match (rank r).
|
||||
# v_hack[name] is [k_max, r]; δS is [r]. Check last-dim match (rank r).
|
||||
rank_bad = [
|
||||
(name, tuple(v_hack[name].shape), tuple(wrappers[name]["delta_S"].shape))
|
||||
for name in sorted(wrapper_keys & vhack_keys)
|
||||
@@ -691,43 +691,38 @@ def main(cfg: Config) -> int:
|
||||
tok = AutoTokenizer.from_pretrained(model_name)
|
||||
if tok.pad_token_id is None: tok.pad_token = tok.eos_token
|
||||
|
||||
# On CPU smoke we fall back to fp32 + sdpa: flash-attn2 is CUDA-only and
|
||||
# CPU bf16 is patchy. Production GPU runs keep bf16 + flash_attention_2.
|
||||
# ── model + tokenizer ──
|
||||
# CPU smoke: fp32 + sdpa (flash-attn2 is CUDA-only, CPU bf16 is patchy).
|
||||
# GPU: bf16 + flash_attention_2.
|
||||
cpu = device.type == "cpu"
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
dtype=torch.float32 if cpu else torch.bfloat16,
|
||||
attn_implementation="sdpa" if cpu else "flash_attention_2",
|
||||
).to(device)
|
||||
# No gradient checkpointing: grad-accum forwards one G-group (6 seqs) at a time,
|
||||
# so peak activation memory is ~6 x merged_len, which fits at G=6 on 96GB without
|
||||
# recompute (worst-case merged 2048; flash-attn keeps attention O(N), MLP/residual
|
||||
# store ~12-15GB). Dropping checkpointing removes the backward recompute (~1.3-1.5x
|
||||
# on the train-compute portion). delta_S gets grad directly (it's a leaf inside
|
||||
# each Linear's W' = W + U diag(delta_S) Vh), so enable_input_require_grads -- a
|
||||
# checkpointing-only trick -- is unnecessary. use_cache is toggled per generate
|
||||
# call below: True for autoregressive decode, False for the single loss forwards.
|
||||
# No gradient checkpointing: grad-accum forwards one G-group at a time, so peak
|
||||
# activation memory fits at G=6 on 96GB without recompute. δS is a leaf inside
|
||||
# W' = W + U diag(δS) Vᵀ, so it gets grad directly (no enable_input_require_grads).
|
||||
# use_cache toggles per generate call: True for decode, False for the loss forwards.
|
||||
model.config.use_cache = False
|
||||
|
||||
# ── AntiPaSTO adapter: δS (kept) + δS_hack (quarantine), same shape r ──
|
||||
is_route2 = cfg.intervention == "route2"
|
||||
wrappers = wrap_model_with_antipasto(
|
||||
model, model_name, CACHE_ROOT, device,
|
||||
grad_probe=is_route2, # route2 needs the per-rollout delta_S gate probe
|
||||
grad_probe=is_route2, # route2 needs the per-rollout δS gate probe
|
||||
)
|
||||
# Both diagonals are trainable params, same shape r (capacity-balanced).
|
||||
# delta_S_hack only ever gets a grad under route (proj.py subspace split) or
|
||||
# route2 (per-rollout tau routing); under none/erase its grad stays None so
|
||||
# AdamW skips it and it stays exactly 0 (forward adds 0 -> identity).
|
||||
# δS_hack only gets a grad under route (proj.py subspace split) or route2
|
||||
# (per-rollout τ routing); under none/erase its grad stays None, so AdamW skips
|
||||
# it and it stays exactly 0 (forward adds 0 -> identity).
|
||||
delta_params = [info["delta_S"] for info in wrappers.values()]
|
||||
delta_hack_params = [info["delta_S_hack"] for info in wrappers.values()]
|
||||
logger.info(f"trainable delta_S: {sum(p.numel() for p in delta_params):,} "
|
||||
f"(+{sum(p.numel() for p in delta_hack_params):,} delta_S_hack quarantine)")
|
||||
|
||||
# v_hack: the hack-direction subspace the erase/route arms project against.
|
||||
# VANILLA (intervention=none) is a pure GRPO baseline and ignores v_hack
|
||||
# entirely -- loading it there only to print a cos_pre diagnostic was misleading
|
||||
# (and could trigger a needless ~5-min extraction). The cin/cout columns are
|
||||
# hidden on vanilla, so v_hack=None just means "no subspace machinery".
|
||||
# ── hack direction: v_hack (erase/route project against it) or v_grad (route2) ──
|
||||
# Vanilla (none) is pure GRPO and ignores v_hack entirely (the cin/cout columns
|
||||
# are hidden, so v_hack=None just means no subspace machinery).
|
||||
v_grad = None # set only by the route2 grad-mask branch below
|
||||
if cfg.intervention in ("none", "route2"):
|
||||
if cfg.intervention == "none" and cfg.v_hack_path is not None:
|
||||
@@ -746,7 +741,7 @@ def main(cfg: Config) -> int:
|
||||
logger.info(f"route2 pairs: hand-crafted PAIRS -> {len(MASK_PAIRS)} pairs")
|
||||
model.eval()
|
||||
# gradient-space mean-diff. extract_v_hack gives per-pair GRPO gradients
|
||||
# on delta_S; v_grad = unit(mean(g_hack - g_clean)) per module, oriented
|
||||
# on δS; v_grad = unit(mean(g_hack - g_clean)) per module, oriented
|
||||
# hack-ward (training reinforces hacks with the same sign, so a rollout
|
||||
# with cos(g_b, v_grad) above the calibrated tau is a reinforced hack).
|
||||
from .extract_vhack_grad import extract_v_hack
|
||||
@@ -807,6 +802,7 @@ def main(cfg: Config) -> int:
|
||||
k_use=cfg.v_hack_k, drop_bottom_frac=cfg.v_hack_drop_bottom_frac,
|
||||
)
|
||||
v_hack = {name: v.to(device) for name, v in v_hack_cpu.items()}
|
||||
# ── teacher pool ──
|
||||
# Teacher pool: pre-generated rollouts on disk keyed by problem_id. Each step's
|
||||
# G_t teacher rollouts come from a uniform random sample of that prompt's cache,
|
||||
# so we do *not* keep the teacher model in VRAM. Pool is produced by
|
||||
@@ -864,9 +860,8 @@ def main(cfg: Config) -> int:
|
||||
f"G_s={G_s} student + G_t={G_t} teacher per prompt (mix_ratio={cfg.mix_ratio})."
|
||||
)
|
||||
|
||||
# One group: delta_S (kept) + delta_S_hack (quarantine) share the lr -- same
|
||||
# shape, same basis, so no per-group lr juggling (the old A_q/B_q LoRA needed
|
||||
# its own lower lr because it was ~60x bigger; gone now).
|
||||
# ── optimizer + schedule ──
|
||||
# δS and δS_hack share the lr (same shape, same basis, no per-group juggling).
|
||||
opt = torch.optim.AdamW(
|
||||
delta_params + delta_hack_params,
|
||||
lr=lr, weight_decay=cfg.weight_decay, betas=(adam_beta1, adam_beta2),
|
||||
@@ -885,6 +880,7 @@ def main(cfg: Config) -> int:
|
||||
milestones=[warmup_steps],
|
||||
)
|
||||
|
||||
# ── generation config ──
|
||||
# Qwen3.5 model card: non-thinking mode for text tasks.
|
||||
# temperature=1.0, top_p=1.0, top_k=20, min_p=0.0, presence_penalty=2.0,
|
||||
# repetition_penalty=1.0. enable_thinking=False is set on the chat template
|
||||
@@ -966,25 +962,11 @@ def main(cfg: Config) -> int:
|
||||
L = max(p.shape[1] for p in parts)
|
||||
return torch.cat([F.pad(p, (0, L - p.shape[1]), value=pad_id) for p in parts], dim=0), n_abl
|
||||
|
||||
# Stream the per-step table live (header once, row per step). Same columns as
|
||||
# the final tabulate output. logger.info routes through tqdm.write so the
|
||||
# rows appear above the progress bar without breaking it.
|
||||
# Names kept <=7 chars so header and value share the same 8-col tab stop.
|
||||
# hack_s/hack_t split out the combined `hack` column by rollout source
|
||||
# (student vs teacher). On no-teacher runs hack_s == hack and hack_t == 0/0.
|
||||
# ref_eq = cumulative generations / 256, where 256 = canonical
|
||||
# num_prompts(16) * num_generations(16) per optimizer step (ariahw config.py).
|
||||
# So ref_eq=1.0 means we've issued the same number of gradient samples as
|
||||
# one canonical reference step. Convert our step count to "reference step
|
||||
# equivalents" by reading this column at the row of interest.
|
||||
# Per-source split (student/teacher) for rew, gt, hack columns. Teacher pool
|
||||
# is frozen so rew_t/gt_t are mostly sanity checks that cache sampling is
|
||||
# stable; rew_s/hack_s are the primary "is student learning?" signals.
|
||||
# `t_rew` is the reward-grading wall-time (s); kept separate from `rew_s`
|
||||
# (student mean reward) to avoid the name collision the older log had.
|
||||
# lp_s, lp_t are mean per-token gen_logp by source. Gap lp_s - lp_t = how
|
||||
# off-policy the teacher pool is from the student's current distribution.
|
||||
# No IS correction is applied to the loss; this is diagnostic only.
|
||||
# Per-step table streamed live (header once, row/step), same columns as the final
|
||||
# tabulate dump; the StepLogger legend below decodes each column. Per-source
|
||||
# (student/teacher) split on rew/gt/hack: teacher rows are frozen sanity, student
|
||||
# rows are the "is it learning?" signal. ref_eq = cumulative gens / 256 (the
|
||||
# canonical 16 prompts x 16 gens/step), so ref_eq=1.0 = one reference step's samples.
|
||||
run_modes = sorted({p["env_mode"] for p in problems}, key=lambda m: list(MODE_CODE).index(m))
|
||||
step_logger = StepLogger(arm=cfg.arm, modes=run_modes)
|
||||
REF_GENS_PER_STEP = 16 * 16 # ariahw/rl-rewardhacking config.py:num_prompts * num_generations
|
||||
@@ -1030,12 +1012,10 @@ def main(cfg: Config) -> int:
|
||||
diverged_steps = 0 # consecutive steps with collapsed teacher ppl (divergence tripwire)
|
||||
lp_t_best = -float("inf") # coherence high-water mark (best teacher gen_logp seen)
|
||||
# ppl_t = exp(-lp_t) on the FIXED teacher rollouts is a free coherence gauge.
|
||||
# Divergence is a DROP from the run's own best coherence, not an absolute level:
|
||||
# a real model sits at lp_t ~ -0.7 and craters to -11..-21 when it diverges (run
|
||||
# 43: lr too high on the 33M quarantine, generations -> token salad), a ~10-nat
|
||||
# drop. A relative threshold also keeps `just smoke` green -- the tiny-random model
|
||||
# has an intrinsic lp_t ~ -11.9 (uniform logp) but it stays flat, so it never DROPS.
|
||||
# Abort if lp_t falls this far below its best for 2 steps running (advantage dead).
|
||||
# Divergence is a DROP from the run's own best, not an absolute level: a healthy
|
||||
# model sits near lp_t ~ -0.7 and craters to -11..-21 (token salad) on divergence.
|
||||
# Relative threshold also keeps smoke green (tiny-random sits at lp_t ~ -11.9 but
|
||||
# stays flat). Abort if lp_t falls this far below best for 2 steps (advantage dead).
|
||||
DIVERGENCE_DROP = 5.0 # nats below best (e^5 ~ 150x worse ppl); never in healthy runs
|
||||
WARN_DROP = 3.0 # softer: log a warning before the hard abort
|
||||
dumped_hack_classes: set[str] = set() # first full example of each hack class -> verbose log
|
||||
@@ -1048,7 +1028,7 @@ def main(cfg: Config) -> int:
|
||||
mode_first_step: dict[str, int] = {}
|
||||
|
||||
def save_ckpt(rows: list[dict], path: Path | None = None) -> None:
|
||||
"""Rewrite the run checkpoint in place: trainable delta_S as tensors, per-step
|
||||
"""Rewrite the run checkpoint in place: trainable δS as tensors, per-step
|
||||
rows + config as JSON metadata (safetensors metadata is str->str only, so the
|
||||
non-tensor payload is JSON). Called every 25 steps and at the end, so an early
|
||||
kill keeps everything up to the last save. Rows are also streamed to the log,
|
||||
@@ -1058,7 +1038,7 @@ def main(cfg: Config) -> int:
|
||||
# dropped from the per-step table as redundant; reconstruct here).
|
||||
hr = sum(r["hack_s"][0] + r["hack_t"][0] for r in rows) / max(1, n_gens)
|
||||
pr = sum(r["gt_s"][0] + r["gt_t"][0] for r in rows) / max(1, n_gens)
|
||||
# Save delta_S only (not delta_S_hack). For route this is exactly the
|
||||
# Save δS only (not δS_hack). For route this is exactly the
|
||||
# deployment adapter: the quarantine knob is ablated at eval, so dropping
|
||||
# it here == the model you'd ship.
|
||||
tensors = {n: info["delta_S"].detach().cpu().contiguous()
|
||||
@@ -1078,6 +1058,7 @@ def main(cfg: Config) -> int:
|
||||
# that interactive bar sparse (tqdm's default maxinterval=10 forces 10s redraws).
|
||||
pbar = tqdm(range(steps), desc=f"train {cfg.arm} {cfg.preset_name}",
|
||||
mininterval=120, maxinterval=120, disable=None)
|
||||
# ── training loop: generate -> grade -> backward -> project -> step ──
|
||||
for step in pbar:
|
||||
t0 = time.time()
|
||||
opt.zero_grad(set_to_none=True)
|
||||
@@ -1104,23 +1085,21 @@ def main(cfg: Config) -> int:
|
||||
# what the projection + optimizer step ultimately sees.
|
||||
step_grad_s: dict[str, torch.Tensor] = {}
|
||||
step_grad_t: dict[str, torch.Tensor] = {}
|
||||
# route2: the flagged rollouts' delta_S-grad contribution, accumulated per
|
||||
# module across prompts, parked into delta_S_hack.grad at injection (the
|
||||
# quarantine, deleted at deploy). Keyed by module name. Mirrors how proj.py
|
||||
# parks route's removed component into delta_S_hack.
|
||||
# route2: the flagged rollouts' δS-grad contribution, accumulated per module
|
||||
# across prompts, parked into δS_hack.grad at injection (the quarantine,
|
||||
# deleted at deploy). Mirrors how proj.py parks route's removed component.
|
||||
step_grad_hack: dict[str, torch.Tensor] = {}
|
||||
|
||||
# route2: recover the per-rollout delta_S grad from the gate
|
||||
# (c.grad = delta_S * g_b), flag rollouts whose grad points hack-ward
|
||||
# (cos(g_b, v_grad) > tau), and route their contribution into delta_S_hack.
|
||||
# Only axes where delta_S has moved (|delta_S| > GATE_EPS) carry a reliable
|
||||
# per-rollout split; near-zero axes keep the full grad, so routing on a fresh
|
||||
# axis lags ~1 step until delta_S grows there (the A1 stale-mask trade-off).
|
||||
# route2: recover the per-rollout δS grad from the gate (c.grad = δS * g_b),
|
||||
# flag rollouts whose grad points hack-ward (cos(g_b, v_grad) > τ), and route
|
||||
# their contribution into δS_hack. Only axes where δS has moved (|δS| > GATE_EPS)
|
||||
# carry a reliable per-rollout split; near-zero axes keep the full grad, so
|
||||
# routing on a fresh axis lags ~1 step until δS grows there (A1 stale-mask trade-off).
|
||||
GATE_EPS = 1e-6
|
||||
step_flagged: list[float] = []
|
||||
step_tau: list[float] = [] # per-(prompt,module) calibrated route threshold
|
||||
step_hkgap: list[float] = [] # ema_hack_cos - ema_clean_cos (discrimination gauge)
|
||||
step_resid: list[float] = [] # cos(delta_S.grad AFTER routing, v_grad): hack-ward leak into deployed knob
|
||||
step_resid: list[float] = [] # cos(δS.grad AFTER routing, v_grad): hack-ward leak into deployed knob
|
||||
|
||||
def _route2_grad_filter(info, n_rollouts: int,
|
||||
hack_anchor: torch.Tensor,
|
||||
@@ -1128,7 +1107,7 @@ def main(cfg: Config) -> int:
|
||||
g = info["delta_S"].grad # [r] summed over rollouts*tokens
|
||||
# The hook's gate c is per-token ([G*s, r]) because nn.Linear sees a
|
||||
# flattened batch. Sum each rollout's token gate-grads -> per-rollout
|
||||
# delta_S*g_b: reshape [G*s, r] -> [G, s, r] -> sum tokens -> [G, r].
|
||||
# δS*g_b: reshape [G*s, r] -> [G, s, r] -> sum tokens -> [G, r].
|
||||
# Pad tokens carry ~0 grad (masked in the loss), so summing every
|
||||
# position is safe. Per-rollout (not per-token) is the preregistered
|
||||
# unit: GRPO advantage is per-rollout, and summing first denoises the
|
||||
@@ -1156,16 +1135,16 @@ def main(cfg: Config) -> int:
|
||||
route2_tau[name] = tau
|
||||
step_tau.append(tau)
|
||||
step_hkgap.append(ema_hack_cos.get(name, 0.0) - ema_clean_cos.get(name, 0.0))
|
||||
# Force-route known hacks (teacher + flagged student); tau-route the
|
||||
# ambiguous rest (incl. unknown B, which lands above tau if it shares
|
||||
# the v_grad direction). Do NOT force-keep clean_anchor -- it is
|
||||
# contaminated with unknown B, which we WANT routed.
|
||||
# Force-route known hacks (teacher + flagged student); τ-route the
|
||||
# ambiguous rest (incl. unknown B, which lands above τ if it shares the
|
||||
# v_grad direction). Do NOT force-keep clean_anchor: it is contaminated
|
||||
# with unknown B, which we WANT routed.
|
||||
flagged = (hack_anchor | (cos_b > tau)).float() # [G]
|
||||
step_flagged.append(flagged.mean().item())
|
||||
sub = torch.where(reliable, (cg * flagged.unsqueeze(1)).sum(0) / dS_safe,
|
||||
torch.zeros_like(g)) # flagged rollouts' contribution
|
||||
# Park the flagged contribution in delta_S_hack (deleted at deploy);
|
||||
# delta_S keeps only the unflagged. Capacity-balanced: both shape [r].
|
||||
# Park the flagged contribution in δS_hack (deleted at deploy); δS keeps
|
||||
# only the unflagged. Capacity-balanced: both shape [r].
|
||||
step_grad_hack[name] = (step_grad_hack[name] + sub.detach().clone()
|
||||
if name in step_grad_hack else sub.detach().clone())
|
||||
g_keep = g - sub # the deployed knob's gradient
|
||||
@@ -1192,6 +1171,7 @@ def main(cfg: Config) -> int:
|
||||
# reward-subprocess-bound (-> parallel grading).
|
||||
t_gen = t_rew = t_fb = 0.0
|
||||
|
||||
# ── per prompt: G_s student + G_t teacher rollouts -> grade -> backward ──
|
||||
for p_idx in range(prompts_per_step):
|
||||
idx = int(torch.randint(0, len(problems), (1,), generator=rng).item())
|
||||
prob = problems[idx]
|
||||
@@ -1469,12 +1449,11 @@ def main(cfg: Config) -> int:
|
||||
ptl_norm = (Lp * mask).sum(1) / mask.sum(1).clamp_min(1)
|
||||
loss = ptl_norm.sum() / (group * prompts_per_step)
|
||||
loss.backward()
|
||||
# route2: per-prompt anchor masks for the tau calibration.
|
||||
# Hack cloud = teacher rows (known-A hacks) + detector-flagged
|
||||
# (hack_E) student rows. Clean cloud = not-flagged student rows
|
||||
# (contaminated with unknown B by design -> conservative tau; B
|
||||
# still routes via cos>tau). is_student = [True]*G_s + [False]*G_t,
|
||||
# so hack_E_flags (len G_s) aligns with the leading student rows.
|
||||
# route2: per-prompt anchor masks for the τ calibration. Hack cloud =
|
||||
# teacher rows (known-A hacks) + detector-flagged (hack_E) student rows;
|
||||
# clean cloud = not-flagged student rows (contaminated with unknown B by
|
||||
# design -> conservative τ; B still routes via cos>τ). hack_E_flags
|
||||
# (len G_s) aligns with the leading student rows of is_student.
|
||||
if is_route2:
|
||||
_n_merged = merged.shape[0]
|
||||
_ha = torch.zeros(_n_merged, dtype=torch.bool, device=Lp.device)
|
||||
@@ -1488,8 +1467,8 @@ def main(cfg: Config) -> int:
|
||||
g = info["delta_S"].grad
|
||||
if g is None:
|
||||
continue
|
||||
# route2 routes here: strip flagged rollouts from delta_S.grad
|
||||
# and park them in delta_S_hack (via step_grad_hack in the filter).
|
||||
# route2 routes here: strip flagged rollouts from δS.grad and
|
||||
# park them in δS_hack (via step_grad_hack in the filter).
|
||||
if is_route2:
|
||||
g = _route2_grad_filter(info, merged.shape[0], _ha, _ca)
|
||||
step_grad_s[name] = (step_grad_s[name] + g.detach().clone()
|
||||
@@ -1499,9 +1478,8 @@ def main(cfg: Config) -> int:
|
||||
agg_loss += loss.item()
|
||||
t_fb += time.perf_counter() - _tfb
|
||||
|
||||
# Inject combined grad (student + teacher) into leaf .grad before
|
||||
# projection + optimizer. Where only one source contributed for a
|
||||
# module, take that source's grad directly.
|
||||
# ── inject grad -> project / route ──
|
||||
# Combine student + teacher grad into each leaf δS.grad (one source -> take it).
|
||||
for name, info in wrappers.items():
|
||||
gs = step_grad_s.get(name)
|
||||
gt = step_grad_t.get(name)
|
||||
@@ -1513,9 +1491,9 @@ def main(cfg: Config) -> int:
|
||||
info["delta_S"].grad = gs
|
||||
else:
|
||||
info["delta_S"].grad = gs + gt
|
||||
# route2: park the flagged rollouts' contribution into delta_S_hack.grad
|
||||
# (the autograd grad from delta_S_hack's own forward path was wiped by the
|
||||
# per-prompt zero_grad; we impose the routed grad here, like proj.py's route).
|
||||
# route2: park the flagged rollouts' contribution into δS_hack.grad (its own
|
||||
# forward-path grad was wiped by the per-prompt zero_grad; we impose the routed
|
||||
# grad here, like proj.py's route).
|
||||
for name, g in step_grad_hack.items():
|
||||
wrappers[name]["delta_S_hack"].grad = g
|
||||
|
||||
@@ -1542,7 +1520,7 @@ def main(cfg: Config) -> int:
|
||||
else:
|
||||
cos_pre_s = cos_pre_t = float("nan")
|
||||
# grad is mutated only for erase (subtract) and route (subtract + park in
|
||||
# delta_S_hack). cos_pre is measured on both.
|
||||
# δS_hack). cos_pre is measured on both.
|
||||
diag = project_delta_S_grad(
|
||||
wrappers, v_hack, cfg.preserve_magnitude,
|
||||
measure_only=False, # erase/route both project; vanilla took the branch above
|
||||
@@ -1591,6 +1569,7 @@ def main(cfg: Config) -> int:
|
||||
opt.step()
|
||||
sched.step()
|
||||
|
||||
# ── v_hack / v_grad refresh ──
|
||||
# Online v_hack refresh: re-extract against the *current* model so the
|
||||
# hack subspace tracks where the student is being pulled now (rather
|
||||
# than at step 0). Same PAIRS, same extract code; we just discard the
|
||||
@@ -1647,14 +1626,13 @@ def main(cfg: Config) -> int:
|
||||
logger.disable("projected_grpo.extract_vhack_grad")
|
||||
logger.disable("__main__")
|
||||
try:
|
||||
# Extract with the quarantine ablated (delta_S_hack=0). For route,
|
||||
# once the hack capability has been routed into delta_S_hack, the
|
||||
# main-knob gradient on the pairs no longer carries the hack
|
||||
# direction -- so re-extracting through the live quarantine rotates
|
||||
# v_hack off-hack and cin_t collapses at the refresh step. Ablating
|
||||
# sends the hack back through the observable main path so D captures
|
||||
# it, matching the delta_S_hack=0 state the build extraction saw.
|
||||
# No-op for erase (delta_S_hack is never trained, stays 0).
|
||||
# Extract with the quarantine ablated (δS_hack=0). For route, once the
|
||||
# hack capability has been routed into δS_hack, the main-knob gradient
|
||||
# on the pairs no longer carries the hack direction, so re-extracting
|
||||
# through the live quarantine rotates v_hack off-hack and cin_t collapses
|
||||
# at the refresh step. Ablating sends the hack back through the observable
|
||||
# main path, matching the δS_hack=0 state the build extraction saw.
|
||||
# No-op for erase (δS_hack is never trained, stays 0).
|
||||
with ablate_quarantine(wrappers):
|
||||
_new_V, _new_S, _, _ = extract_v_hack(
|
||||
model, tok, wrappers, VHACK_PAIRS,
|
||||
@@ -1691,6 +1669,7 @@ def main(cfg: Config) -> int:
|
||||
model.train()
|
||||
refr = f"{len(v_hack)}/{sum(V.shape[0] for V in v_hack.values())}" # mod/axes -> per-step row
|
||||
|
||||
# ── deploy-eval (route/route2): zero δS_hack, eval the shipped model ──
|
||||
# Periodic DEPLOY-eval (routing, Gradient Routing): zero the quarantine knob
|
||||
# and eval the DEPLOYED model on a fixed subset. Routing's claim is that the
|
||||
# cheating capability lands in the quarantine, so deleting it (= what we deploy)
|
||||
@@ -1839,7 +1818,7 @@ def main(cfg: Config) -> int:
|
||||
"cos_post": diag["mean_cos_post"],
|
||||
"fired": diag["frac_fired"],
|
||||
"refr": refr,
|
||||
# Route deploy-eval (delta_S_hack=0); NaN except on route eval steps.
|
||||
# Route deploy-eval (δS_hack=0); NaN except on route eval steps.
|
||||
# Appended AFTER refr so results.py's positional GT_S/HACK_S indices
|
||||
# are unaffected. plot_dynamics reads it by name.
|
||||
"hack_deploy": hack_deploy,
|
||||
@@ -1946,7 +1925,7 @@ def main(cfg: Config) -> int:
|
||||
# R3 sneaky-fail guard: under route, the quarantine knob must have absorbed
|
||||
# something (‖δS_hack‖ > 0), else routing silently degenerated to
|
||||
# erasure (parked grad never applied). Exactly 0 by construction for
|
||||
# none/erase (delta_S_hack gets no grad -> AdamW skips it).
|
||||
# none/erase (δS_hack gets no grad -> AdamW skips it).
|
||||
dsh_norm = float(sum(info["delta_S_hack"].data.float().pow(2).sum().item()
|
||||
for info in wrappers.values()) ** 0.5)
|
||||
logger.info(f"||delta_S_hack|| = {dsh_norm:.4f} "
|
||||
@@ -1965,6 +1944,7 @@ def main(cfg: Config) -> int:
|
||||
f"SHOULD: coherent code/prose. ELSE token salad => diverged, eval below is moot.\n"
|
||||
f"{_r['text'][:800]}\n=== END LAST GEN ===\n")
|
||||
|
||||
# ── final eval + BLUF ──
|
||||
# Final per-mode train-vs-deploy eval -- run for EVERY arm on the SAME fixed
|
||||
# eval subset so the all-arms overlay reads identical numbers. For route/route2
|
||||
# this is the absorption test: TRAIN keeps the quarantine knob on (still hacks),
|
||||
|
||||
Reference in New Issue
Block a user