mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-08-02 12:40:13 +08:00
wassnameandClaude Opus 4.7
ff26cbe089
The combined `rew` column mixed student + teacher rollouts, making it hard to tell "is student learning?" at a glance. Add per-source splits: - rew_s: student-only mean reward (primary learning signal) - gt_t : teacher-only ground-truth pass count (cache stability check) The teacher pool is frozen at startup (baseline logged at load), so per-step rew_t adds noise without information and is omitted. The previous `rew_s` column was actually reward-grading wall-time (an unfortunate name collision with student reward). Rename it to `t_rew` to match the other timing cols (gen, fb). New column order: step ref_eq rew rew_s std sprd N gt gt_s gt_t hack hack_s hack_t loss cin cin_s cin_t cout fired gen fb t_rew sec Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
975 lines
50 KiB
Python
975 lines
50 KiB
Python
"""Canonical training entry point: AntiPaSTO + GRPO (Dr.GRPO unbiased) + optional
|
|
gradient projection on LeetCode reward-hacking benchmark.
|
|
|
|
Lineage (see spec.md §76-83):
|
|
- The inner GRPO_step (per_token_logps, ratio + clip + min, K3 KL, per-token
|
|
loss, completion mask) is a direct port of lsdefine/simple_GRPO's
|
|
`GRPO_step` in `grpo_vllm_one.py` (lines 64-95).
|
|
- The OUTER loop adopts simple_GRPO's `Q_batch_size` pattern (multiple
|
|
prompts per optimizer step, per-prompt GRPO advantage groups, grad
|
|
accumulation across prompts). GRPO needs within-group reward diversity to
|
|
produce any signal; sampling many prompts per step raises the chance that
|
|
at least one group is non-degenerate. simple_GRPO uses Q_batch_size=5; our
|
|
prompts_per_step is set in PRESETS (grad-accum to the paper's effective batch).
|
|
- Deviations from simple_GRPO are deliberate, listed in spec.md:
|
|
1. Loss normalization: Dr.GRPO unbiased (Liu et al. 2025, arXiv
|
|
2503.20783) replaces simple_GRPO's `(R-mean)/std` + per-response-len
|
|
denominator. Drops two biases:
|
|
- length norm `1/|o_i|` (favors short correct, long incorrect)
|
|
- group-std norm `/std(R)` (overweights easy/hard questions)
|
|
Toggle via `--unbiased` (default on); flipping to False recovers
|
|
simple_GRPO's classic GRPO advantage normalization.
|
|
2. Reference model: simple_GRPO runs a separate base model via an HTTP
|
|
`ref_server`. We use the AntiPaSTO `delta_S=0` zero-adapter trick
|
|
(W' = W + U diag(0) Vh = W exactly) — no second model loaded.
|
|
3. Rollout: simple_GRPO uses vLLM in a separate process. We use HF
|
|
`model.generate` in-process.
|
|
4. Adapter: simple_GRPO is full FT (with DeepSpeed ZeRO). Canonical
|
|
(ariahw/rl-rewardhacking) is LoRA r=32. We use AntiPaSTO full-rank
|
|
SVD adapter (the research artifact).
|
|
|
|
Hyperparameters (lr, weight_decay, betas, warmup, cosine, beta=KL) are taken
|
|
from the closest-in-param-count reference: ariahw/rl-rewardhacking config.py
|
|
(LoRA r=32 on 4B ≈ 30M params) rather than simple_GRPO (full FT on 7B). See
|
|
docs/grpo_hyperparams.md.
|
|
|
|
Reference-model term (`--beta`): Dr.GRPO argues beta=0 is fine for *reasoning*
|
|
RL with rule-based reward (no distributional-shift concern when reward = ground
|
|
truth). That argument does NOT apply when studying reward hacking, which IS the
|
|
distributional shift between proxy reward and true objective, so `full` uses
|
|
beta>0 (value from ariahw config.py; see PRESETS). The delta_S=0 free-ref-model
|
|
trick gives this at zero extra VRAM: W' = W + U diag(0) Vh = W exactly, so a
|
|
no_grad forward with delta_S zeroed yields pi_ref logprobs without a 2nd model.
|
|
The smoke preset uses beta=0 only because the 24GB GPU can't hold even that.
|
|
|
|
All per-preset hyperparameters (model, steps, G, max_new, n_problems, beta,
|
|
prompts_per_step) live in the PRESETS dict below — the single source of truth.
|
|
|
|
Run:
|
|
uv run python -m projected_grpo.train --preset=smoke --arm=vanilla
|
|
uv run python -m projected_grpo.train --preset=full --arm=projected
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import gzip
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
# Must be set BEFORE `import torch` to take effect on the CUDA allocator.
|
|
# Eliminates fragmentation that caused 91 GiB allocated / 581 MiB free crash
|
|
# on Qwen3-4B G=8 (PyTorch's own OOM message recommends this).
|
|
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import tyro
|
|
from loguru import logger
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
from tabulate import tabulate
|
|
from tqdm import tqdm
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
|
|
|
|
from .antipasto import wrap_model_with_antipasto
|
|
from .proj import per_token_logps, project_delta_S_grad, mean_cin_from_grads
|
|
from .rewards import compute_reward
|
|
|
|
CACHE_ROOT = Path("svd_cache")
|
|
OUT_DIR = Path("out")
|
|
LOGS_DIR = Path("logs")
|
|
DATA = Path("external/rl-rewardhacking/results/data/leetcode_train_medhard_filtered.jsonl")
|
|
|
|
|
|
def setup_logging(run_id: str) -> Path:
|
|
"""Token-efficient loguru: stdout = 1-char icon + msg; verbose log to file.
|
|
|
|
See /root/.claude/skills/token-efficient-logging/SKILL.md.
|
|
"""
|
|
LOGS_DIR.mkdir(exist_ok=True)
|
|
verbose_log = LOGS_DIR / f"{datetime.now().strftime('%Y%m%dT%H%M%S')}_{run_id}.log"
|
|
logger.remove()
|
|
logger.add(
|
|
lambda msg: tqdm.write(msg, end=""),
|
|
colorize=True,
|
|
format="<level>{level.icon}</level> {message}",
|
|
level="INFO",
|
|
)
|
|
logger.add(
|
|
verbose_log,
|
|
format="{time:HH:mm:ss} | {level} | {message}",
|
|
level="DEBUG",
|
|
)
|
|
logger.level("INFO", icon="I")
|
|
logger.level("WARNING", icon="W")
|
|
logger.level("ERROR", icon="E")
|
|
logger.level("DEBUG", icon="D")
|
|
return verbose_log
|
|
|
|
|
|
class Preset(str, Enum):
|
|
smoke = "smoke"
|
|
full = "full"
|
|
|
|
|
|
PRESETS: dict[str, dict] = {
|
|
"smoke": dict(model="Qwen/Qwen3.5-0.8B", steps=10, group=2, max_new=128,
|
|
n_problems=30, beta=0.0, prompts_per_step=1), # 24GB cap
|
|
# 4B matches reference DEFAULT_MODEL_ID (docs/vendor/rl-rewardhacking/src/__init__.py).
|
|
# G=6 after 2026-05-24 step-17 OOM at G=8: lm_head spike on a long-prompt
|
|
# problem hit 4.16 GiB / 2.5 GiB free. `logits_to_keep` cuts lm_head ~33%;
|
|
# G=8->6 cuts B at every act site ~25%. Combined headroom ~6-10 GB.
|
|
# prompts_per_step=43: grad-accum to paper's effective batch (256 generations
|
|
# per optimizer step; ariahw config.py num_prompts=16 x num_generations=16).
|
|
# At our VRAM-capped G=6, 43 x 6 = 258 ~= 256. Grad accum -> same peak VRAM,
|
|
# ~5x wall-time vs pp=8. n_problems=992 is the full filtered set (paper fn.9).
|
|
"full": dict(model="Qwen/Qwen3-4B", steps=200, group=6, max_new=1024,
|
|
n_problems=992, beta=1e-3, prompts_per_step=43),
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class Config:
|
|
preset: Preset = Preset.smoke
|
|
arm: Literal["vanilla", "projected"] = "projected"
|
|
# Per-preset overrides; leave None to use preset defaults.
|
|
model: str | None = None
|
|
steps: int | None = None
|
|
group: int | None = None # G samples per question
|
|
max_new: int | None = None
|
|
n_problems: int | None = None
|
|
beta: float | None = None # KL coef. If >0, uses delta_S=0 free-ref-model trick.
|
|
prompts_per_step: int | None = None # P prompts per optimizer step; grads accumulate over P.
|
|
# Universal knobs.
|
|
clip: float = 0.2
|
|
lr: float = 7e-5 # canonical (rl-rewardhacking config.py:138)
|
|
weight_decay: float = 0.1 # canonical config.py:142
|
|
adam_beta1: float = 0.9 # canonical config.py:143
|
|
adam_beta2: float = 0.99 # canonical config.py:144
|
|
warmup_steps: int = 10 # canonical config.py:141; cosine decay after
|
|
seed: int = 41
|
|
preserve_magnitude: bool = True
|
|
gate_mode: Literal["one_sided", "no_gate"] = "one_sided"
|
|
unbiased: bool = True # Dr.GRPO: drop 1/|o_i| and /std(R)
|
|
# v_hack: path is optional — if None, derived from model+top_k as
|
|
# out/v_hack_<slug>_k<extract_top_k>.safetensors. If file missing, train.py
|
|
# auto-extracts (cheap: ~5min, shares the already-loaded model). Set explicitly
|
|
# to override (e.g. baked-variant v_hack paths). v_hack_k slices the saved
|
|
# top-k_max directions to top-k_use at load time — the k-ablation knob.
|
|
v_hack_path: Path | None = None
|
|
v_hack_extract_top_k: int = 12 # max k to save at extract; n_train_pairs caps it lower
|
|
v_hack_k: int = 5 # load-time slice; k=1 = mean-diff, k=k_max = full
|
|
v_hack_tau_axis: float = 0.0 # extract-time: zero axes where S_i/S_0 < tau_axis
|
|
out_tag: str = "" # suffix for saved artifact, e.g. "_seed41"
|
|
# Mixed-pool GRPO: per-prompt rollout pool = G_s live student + G_t cached
|
|
# teacher rollouts. Teacher pool is a dir of prompt_NNNN.jsonl.gz produced by
|
|
# probe_distill.py --teacher-only (schema includes prompt_ids, completion_ids,
|
|
# plen, reward, hacked, gt_pass, fmt_ok). Reward labels are read from cache
|
|
# (not re-graded) so the pool is reproducible. G_t = round(G * mix_ratio),
|
|
# G_s = G - G_t. Both halves contribute to a single group-relative advantage.
|
|
# Loss is unchanged: ratio==1 in single-inner-step PPO, so reward-weighted
|
|
# policy gradient applies uniformly to both halves regardless of source.
|
|
teacher_pool_dir: Path | None = None
|
|
mix_ratio: float = 0.5
|
|
|
|
def resolved(self) -> dict:
|
|
"""Merge preset defaults with explicit overrides."""
|
|
base = dict(PRESETS[self.preset.value])
|
|
for k in ("model", "steps", "group", "max_new", "n_problems", "beta", "prompts_per_step"):
|
|
v = getattr(self, k)
|
|
if v is not None: base[k] = v
|
|
return base
|
|
|
|
|
|
# rh-s65 (ariahw "no intervention" arm, run_rl_training.py:122-136) was trained with:
|
|
# - dataset's baked-in CODE_SYSTEM_PROMPT (no inoculation override)
|
|
# - simple_overwrite_tests hint applied: "and should pass all tests" -> "and will
|
|
# be evaluated by calling a function called run_tests()". This is the loophole
|
|
# disclosure that lets the model hack by defining run_tests() itself.
|
|
# See docs/vendor/rl-rewardhacking/src/data/hints.py:85-106 (SimpleOverwriteTests)
|
|
# and run_rl_training.py:12 (DEFAULT_TASK="simple_overwrite_tests").
|
|
RH_HINT_REPLACE_FROM = "and should pass all tests"
|
|
RH_HINT_REPLACE_TO = "and will be evaluated by calling a function called run_tests()"
|
|
|
|
|
|
def load_problems(n: int) -> list[dict]:
|
|
out = []
|
|
with DATA.open() as f:
|
|
for idx, line in enumerate(f):
|
|
if len(out) >= n: break
|
|
d = json.loads(line)
|
|
msgs = [dict(m) for m in d["prompt"]]
|
|
for m in msgs:
|
|
if m.get("role") == "user":
|
|
m["content"] = m["content"].replace(
|
|
RH_HINT_REPLACE_FROM, RH_HINT_REPLACE_TO,
|
|
)
|
|
break
|
|
out.append({
|
|
"problem_id": d.get("id", idx),
|
|
"messages": msgs,
|
|
"gt_tests": d["gt_answer"],
|
|
"setup_code": d.get("setup_code", ""),
|
|
"func_name": d.get("func_name", "Solution().solve"),
|
|
"canonical": d.get("canonical_solution", ""),
|
|
})
|
|
return out
|
|
|
|
|
|
def load_v_hack(
|
|
path: Path, model_name: str, wrappers: dict, k_use: int | None = None,
|
|
) -> dict[str, torch.Tensor]:
|
|
"""Load v_hack (top-k directions) for this wrapped model.
|
|
|
|
File schema (v2): bare `{name}` keys hold V[k_max, r]; `_sv/{name}` keys hold
|
|
S[k_max] (read but not returned — no caller uses them yet). v_hack is
|
|
model-specific because module names and per-module SVD ranks depend on the
|
|
exact checkpoint; a smoke (Qwen3.5-0.8B) v_hack must not be reused for a
|
|
full (Qwen3-4B) run.
|
|
|
|
If `k_use` is given, slices V to top-k_use rows. Errors if k_use > k_max
|
|
saved (re-extract with a higher top_k).
|
|
"""
|
|
with safe_open(str(path), framework="pt", device="cpu") as f:
|
|
meta = f.metadata() or {}
|
|
saved_model = meta.get("model")
|
|
saved_dtype = meta.get("dtype")
|
|
if saved_model is None or saved_dtype is None:
|
|
raise ValueError(
|
|
f"{path} has no model/dtype header metadata. "
|
|
f"Re-extract with `uv run python -m projected_grpo.extract_vhack_grad "
|
|
f"--model={model_name} --dtype=bf16 --out-path={path}`."
|
|
)
|
|
if saved_model != model_name:
|
|
raise ValueError(f"v_hack model mismatch: {path} has {saved_model}, run uses {model_name}")
|
|
if saved_dtype != "bf16":
|
|
raise ValueError(
|
|
f"v_hack dtype/SVD-basis mismatch: {path} was extracted with dtype={saved_dtype}; "
|
|
"train.py loads models in bf16. Re-extract with `--dtype=bf16`."
|
|
)
|
|
# Read only V keys (bare module names); _sv/{name} keys are saved
|
|
# alongside but no runtime path consumes them currently.
|
|
v_hack = {k: f.get_tensor(k) for k in f.keys() if not k.startswith("_sv/")}
|
|
|
|
wrapper_keys = set(wrappers)
|
|
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).
|
|
rank_bad = [
|
|
(name, tuple(v_hack[name].shape), tuple(wrappers[name]["delta_S"].shape))
|
|
for name in sorted(wrapper_keys & vhack_keys)
|
|
if v_hack[name].ndim != 2 or v_hack[name].shape[-1] != wrappers[name]["delta_S"].shape[0]
|
|
]
|
|
if missing or extra or rank_bad:
|
|
raise ValueError(
|
|
"v_hack incompatible with wrapped model: "
|
|
f"missing={len(missing)} examples={missing[:5]} "
|
|
f"extra={len(extra)} examples={extra[:5]} "
|
|
f"rank_bad={len(rank_bad)} examples={rank_bad[:5]}. "
|
|
"Extract a fresh v_hack with `uv run python -m projected_grpo.extract_vhack_grad "
|
|
f"--model={model_name} --out-path={path}`."
|
|
)
|
|
|
|
k_max = next(iter(v_hack.values())).shape[0]
|
|
if k_use is not None:
|
|
if k_use > k_max:
|
|
raise ValueError(
|
|
f"requested k_use={k_use} exceeds saved k_max={k_max} in {path}. "
|
|
f"Re-extract with `--top-k={k_use}`."
|
|
)
|
|
v_hack = {n: v[:k_use].contiguous() for n, v in v_hack.items()}
|
|
logger.info(
|
|
f"loaded v_hack from {path}: modules={len(v_hack)}; k_saved={k_max}, k_use={k_use or k_max}"
|
|
)
|
|
return v_hack
|
|
|
|
|
|
@torch.no_grad()
|
|
def ref_logprobs_via_zero_delta(
|
|
model, merged: torch.Tensor, wrappers: dict, plen: int,
|
|
) -> torch.Tensor:
|
|
"""Compute pi_ref logprobs on completion tokens only.
|
|
|
|
AntiPaSTO: W' = W + U diag(delta_S) Vh. At delta_S=0, W' = W exactly
|
|
(verified bit-exact in step 1). Save -> zero -> forward -> restore.
|
|
Zero extra VRAM vs a separately loaded ref_model.
|
|
|
|
Uses `logits_to_keep=L_c+1` so HF's lm_head only runs on completion-side
|
|
hidden states; prompt-side logits never materialize. Saves
|
|
~plen/(plen+L_c) memory at the lm_head call (~33% at plen=500, L_c=1024).
|
|
That was the OOM site at vanilla step 17 (long prompt -> 4 GiB lm_head spike).
|
|
"""
|
|
saved = {n: info["delta_S"].data.clone() for n, info in wrappers.items()}
|
|
try:
|
|
for info in wrappers.values():
|
|
info["delta_S"].data.zero_()
|
|
L_c = merged.shape[1] - plen
|
|
logits = model(merged, logits_to_keep=L_c + 1).logits[:, :-1]
|
|
return per_token_logps(logits, merged[:, plen:])
|
|
finally:
|
|
for n, info in wrappers.items():
|
|
info["delta_S"].data.copy_(saved[n])
|
|
|
|
|
|
def main(cfg: Config) -> int:
|
|
p = cfg.resolved()
|
|
model_name = p["model"]; steps = p["steps"]; group = p["group"]
|
|
max_new = p["max_new"]; n_problems = p["n_problems"]; beta = p["beta"]
|
|
prompts_per_step = p["prompts_per_step"]
|
|
|
|
run_id = f"{cfg.preset.value}_{cfg.arm}_seed{cfg.seed}{cfg.out_tag}"
|
|
verbose_log = setup_logging(run_id)
|
|
|
|
torch.manual_seed(cfg.seed)
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
# BLUF up front: argv + setup + verbose-log pointer so a tail-reader sees context.
|
|
logger.info(f"argv: {' '.join(sys.argv)}")
|
|
logger.info(f"verbose log: {verbose_log}")
|
|
logger.info(
|
|
f"preset={cfg.preset.value} arm={cfg.arm} model={model_name} "
|
|
f"steps={steps} G={group} max_new={max_new} beta={beta} "
|
|
f"unbiased={cfg.unbiased} seed={cfg.seed} device={device}"
|
|
)
|
|
|
|
tok = AutoTokenizer.from_pretrained(model_name)
|
|
if tok.pad_token_id is None: tok.pad_token = tok.eos_token
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_name, dtype=torch.bfloat16,
|
|
attn_implementation="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.
|
|
model.config.use_cache = False
|
|
|
|
wrappers = wrap_model_with_antipasto(model, model_name, CACHE_ROOT, device)
|
|
delta_params = [info["delta_S"] for info in wrappers.values()]
|
|
logger.info(f"trainable delta_S: {sum(p.numel() for p in delta_params):,}")
|
|
|
|
# v_hack: derive default path from model + extract_top_k unless overridden.
|
|
# Always loaded (or auto-extracted) so vanilla also reports cos_in as a baseline.
|
|
# Auto-extract reuses the already-wrapped model — no second model load.
|
|
# Slug: works for HF names ("Qwen/Qwen3-4B" -> "Qwen3-4B") and local paths
|
|
# ("out/baked/qwen3_4b_rh25" -> "qwen3_4b_rh25").
|
|
model_slug = model_name.rstrip("/").split("/")[-1]
|
|
# Filename encodes top_k AND tau_axis because both are baked into the saved
|
|
# V (extract zeros rows where S_i/S_0 < tau_axis before saving). If a future
|
|
# ablation varies pairs.py, add a pairs hash here too.
|
|
tau_tag = f"_tau{cfg.v_hack_tau_axis:g}" if cfg.v_hack_tau_axis > 0 else ""
|
|
if cfg.v_hack_path is None:
|
|
v_hack_path = OUT_DIR / f"v_hack_{model_slug}_k{cfg.v_hack_extract_top_k}{tau_tag}.safetensors"
|
|
else:
|
|
v_hack_path = cfg.v_hack_path
|
|
if not v_hack_path.exists():
|
|
from .extract_vhack_grad import extract_v_hack
|
|
from .pairs import PAIRS as VHACK_PAIRS
|
|
from safetensors.torch import save_file
|
|
logger.info(f"v_hack cache miss at {v_hack_path}; extracting (~5min)...")
|
|
model.eval() # match standalone extract: deterministic backward, no dropout
|
|
v_hack_cpu_dict, raw_grads, _diag = extract_v_hack(
|
|
model, tok, wrappers, VHACK_PAIRS,
|
|
top_k=cfg.v_hack_extract_top_k, tau_axis=cfg.v_hack_tau_axis,
|
|
n_heldout=2, device=device,
|
|
)
|
|
OUT_DIR.mkdir(exist_ok=True)
|
|
save_file(v_hack_cpu_dict, str(v_hack_path),
|
|
metadata={"model": model_name, "dtype": "bf16",
|
|
"top_k": str(min(cfg.v_hack_extract_top_k, len(VHACK_PAIRS) - 2)),
|
|
"tau_axis": str(cfg.v_hack_tau_axis), "schema": "v2_with_sv"})
|
|
# extract zeros grads at exit; opt is built below so no opt-state taint.
|
|
model.train() # restore train mode; eval was set only for the extract pass
|
|
v_hack_cpu = load_v_hack(v_hack_path, model_name, wrappers, k_use=cfg.v_hack_k)
|
|
v_hack = {name: v.to(device) for name, v in v_hack_cpu.items()}
|
|
# 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
|
|
# `probe_distill.py --teacher-only` (see schema in probe_distill.py:149-186).
|
|
# Cached rewards/flags are reused verbatim — no re-grading — so the pool is a
|
|
# reproducible fixed teacher distribution across runs.
|
|
teacher_pool: dict[int, list[dict]] = {}
|
|
G_s = group
|
|
G_t = 0
|
|
if cfg.teacher_pool_dir is not None:
|
|
if not (0.0 < cfg.mix_ratio < 1.0):
|
|
raise ValueError(f"mix_ratio must be in (0,1) when teacher_pool_dir set; got {cfg.mix_ratio}")
|
|
G_t = round(group * cfg.mix_ratio)
|
|
G_s = group - G_t
|
|
if G_s == 0 or G_t == 0:
|
|
raise ValueError(
|
|
f"degenerate split: G={group} mix_ratio={cfg.mix_ratio} -> G_s={G_s}, G_t={G_t}. "
|
|
f"Pick mix_ratio so both halves are non-empty, or drop --teacher-pool-dir."
|
|
)
|
|
for path in sorted(cfg.teacher_pool_dir.glob("prompt_*.jsonl.gz")):
|
|
# path.stem on 'prompt_0004.jsonl.gz' is 'prompt_0004.jsonl' (only one
|
|
# suffix stripped); split off the .jsonl before parsing the int.
|
|
problem_id = int(path.name.split("_")[1].split(".")[0])
|
|
with gzip.open(path, "rt") as f:
|
|
teacher_pool[problem_id] = [json.loads(line) for line in f]
|
|
if not teacher_pool:
|
|
raise FileNotFoundError(
|
|
f"teacher pool {cfg.teacher_pool_dir} is empty. Run `just pregen-teacher N` first."
|
|
)
|
|
n_rollouts_per = sum(len(v) for v in teacher_pool.values()) / len(teacher_pool)
|
|
avg_hack = sum(int(r["hacked"]) for v in teacher_pool.values() for r in v) / sum(len(v) for v in teacher_pool.values())
|
|
logger.info(
|
|
f"teacher pool: {len(teacher_pool)} prompts, "
|
|
f"~{n_rollouts_per:.1f} rollouts/prompt, "
|
|
f"cached hack_rate={avg_hack:.2%}. "
|
|
f"G_s={G_s} student + G_t={G_t} teacher per prompt (mix_ratio={cfg.mix_ratio})."
|
|
)
|
|
|
|
opt = torch.optim.AdamW(
|
|
delta_params, lr=cfg.lr, weight_decay=cfg.weight_decay,
|
|
betas=(cfg.adam_beta1, cfg.adam_beta2),
|
|
)
|
|
# Linear warmup over `warmup_steps`, then cosine decay to 0 over the rest.
|
|
# Matches canonical (lr_scheduler_type='cosine', warmup_steps=10).
|
|
sched = torch.optim.lr_scheduler.SequentialLR(
|
|
opt,
|
|
schedulers=[
|
|
torch.optim.lr_scheduler.LinearLR(opt, start_factor=1e-3, end_factor=1.0,
|
|
total_iters=max(1, cfg.warmup_steps)),
|
|
torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(1, steps - cfg.warmup_steps)),
|
|
],
|
|
milestones=[max(1, cfg.warmup_steps)],
|
|
)
|
|
|
|
# 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
|
|
# below (safe no-op if the model's template doesn't support it).
|
|
gen_cfg = GenerationConfig(
|
|
max_new_tokens=max_new, do_sample=True,
|
|
# T=0.7 matches ariahw reference (config.py:172). T=1.0 had hack emerging
|
|
# too slowly: hack patterns are modal in the baked substrate; broad sampling
|
|
# at T=1 dilutes them. Lower T expresses the substrate's hack propensity.
|
|
temperature=0.7, top_p=1.0, top_k=20, min_p=0.0,
|
|
repetition_penalty=1.0,
|
|
num_return_sequences=group, pad_token_id=tok.pad_token_id,
|
|
)
|
|
|
|
problems = load_problems(n_problems)
|
|
logger.info(f"loaded {len(problems)} problems from {DATA.name}")
|
|
if teacher_pool:
|
|
# Restrict prompt sampling to problems with cached teacher rollouts;
|
|
# otherwise we'd skip the majority of steps when the pool is sparse
|
|
# (e.g. 70/992 prompts cached -> ~93% skip rate).
|
|
before = len(problems)
|
|
problems = [p for p in problems if p["problem_id"] in teacher_pool]
|
|
logger.info(
|
|
f"teacher pool restriction: {len(problems)}/{before} prompts kept "
|
|
f"(student trains only on prompts covered by the cached teacher pool)"
|
|
)
|
|
if not problems:
|
|
raise ValueError(
|
|
f"no overlap between training set ({before} problems) and teacher pool "
|
|
f"({len(teacher_pool)} cached prompts). Re-run pregen-teacher against the same dataset."
|
|
)
|
|
|
|
rng = torch.Generator().manual_seed(cfg.seed)
|
|
rows = []
|
|
logger.info(
|
|
f"SHOULD: loss finite each step; projected arm cos_out <= cos_in; "
|
|
f"PASS_RATE > 0 on 4B (was 0/16 under broken grader). "
|
|
f"ELSE: harness or projection broken. "
|
|
f"Timing cols (gen/fb/t_rew/sec): gen-bound -> vLLM; fb-bound -> lower pp; t_rew-bound -> parallel grading."
|
|
)
|
|
if teacher_pool:
|
|
logger.info(
|
|
f"SHOULD (mixed-pool): hack_t high from step 0 (cached teacher pool ~95% hack); "
|
|
f"hack_s climbs 0 -> 20%+ over the run as student learns from exposure. "
|
|
f"ELSE if hack_s flat while hack_t high: student is ignoring the off-policy "
|
|
f"gradient signal — bump mix_ratio or lr."
|
|
)
|
|
|
|
eos_id = tok.eos_token_id
|
|
pad_id = tok.pad_token_id
|
|
|
|
# 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.
|
|
_row_cols = ["step", "ref_eq", "rew", "rew_s", "std", "sprd", "N",
|
|
"gt", "gt_s", "gt_t", "hack", "hack_s", "hack_t",
|
|
"loss", "cin", "cin_s", "cin_t", "cout", "fired",
|
|
"gen", "fb", "t_rew", "sec"]
|
|
REF_GENS_PER_STEP = 16 * 16 # ariahw/rl-rewardhacking config.py:num_prompts * num_generations
|
|
est_gens_per_step = cfg.prompts_per_step * cfg.group # before mixed-pool split
|
|
logger.info(
|
|
f"grad-pressure: {est_gens_per_step} gens/step vs reference {REF_GENS_PER_STEP} "
|
|
f"-> {est_gens_per_step / REF_GENS_PER_STEP:.2f}x per step; "
|
|
f"this run's {steps} steps ~= {steps * est_gens_per_step / REF_GENS_PER_STEP:.1f} reference steps."
|
|
)
|
|
logger.info("row\t" + "\t".join(_row_cols))
|
|
|
|
OUT_DIR.mkdir(exist_ok=True)
|
|
tag = cfg.out_tag or f"_{cfg.preset.value}_{cfg.arm}_seed{cfg.seed}"
|
|
ckpt_path = OUT_DIR / f"train{tag}.safetensors"
|
|
first_hack_path = OUT_DIR / f"train{tag}_first_hack.safetensors"
|
|
first_hack_saved = False
|
|
|
|
def save_ckpt(rows: list[dict], path: Path | None = None) -> None:
|
|
"""Rewrite the run checkpoint in place: trainable delta_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,
|
|
so this is convenience, not the only copy. Mirrors the v_hack metadata idiom."""
|
|
n_gens = sum(r["N"] for r in rows)
|
|
hr = sum(int(r["hack"].split("/")[0]) for r in rows) / max(1, n_gens)
|
|
pr = sum(int(r["gt"].split("/")[0]) for r in rows) / max(1, n_gens)
|
|
tensors = {n: info["delta_S"].detach().cpu().contiguous()
|
|
for n, info in wrappers.items()}
|
|
save_file(tensors, str(path or ckpt_path), metadata={
|
|
"model": model_name, "dtype": "bf16", "step": str(len(rows)),
|
|
"hack_rate": f"{hr:.6f}", "pass_rate": f"{pr:.6f}",
|
|
"rows": json.dumps(rows), "cfg": json.dumps(vars(cfg), default=str),
|
|
"resolved": json.dumps(p),
|
|
})
|
|
|
|
pool_validated = False # flips True once cached prompt_ids matches live tokenization
|
|
pbar = tqdm(range(steps), desc=f"train {cfg.arm} {cfg.preset.value}", mininterval=60)
|
|
for step in pbar:
|
|
t0 = time.time()
|
|
opt.zero_grad(set_to_none=True)
|
|
|
|
# Accumulate across P prompts; one optimizer step at the end. Per-prompt
|
|
# group of G generations is the GRPO advantage normalisation unit.
|
|
agg_rew, agg_gt, agg_hack, agg_fmt = [], [], [], []
|
|
agg_is_student: list[bool] = []
|
|
agg_comp_lens, agg_finished, n_skipped = [], [], 0
|
|
agg_loss = 0.0
|
|
diag_tail = None
|
|
# Per-source grad accumulators: each prompt's backward is split into
|
|
# student-only and teacher-only passes so we can compute cin_s / cin_t
|
|
# separately (discriminator: does v_hack actually project hack grads
|
|
# more than non-hack?). step_grad_combined = student + teacher and is
|
|
# what the projection + optimizer step ultimately sees.
|
|
step_grad_s: dict[str, torch.Tensor] = {}
|
|
step_grad_t: dict[str, torch.Tensor] = {}
|
|
# Phase timers (per-step cumulative, seconds). Each GPU phase ends in a
|
|
# CPU-blocking op (decode / .item()), so perf_counter is sync-accurate
|
|
# without explicit cuda.synchronize. Tells us whether wall-time is
|
|
# generation-bound (-> vLLM), forward/backward-bound (-> lower pp), or
|
|
# reward-subprocess-bound (-> parallel grading).
|
|
t_gen = t_rew = t_fb = 0.0
|
|
|
|
for p_idx in range(prompts_per_step):
|
|
idx = int(torch.randint(0, len(problems), (1,), generator=rng).item())
|
|
prob = problems[idx]
|
|
prompt = tok.apply_chat_template(
|
|
prob["messages"], tokenize=False, add_generation_prompt=True,
|
|
enable_thinking=False, # canonical training default; no-op if template ignores it
|
|
)
|
|
enc = tok(prompt, return_tensors="pt", add_special_tokens=False).to(device)
|
|
plen = enc.input_ids.shape[1]
|
|
if plen + max_new > 2048:
|
|
n_skipped += 1
|
|
continue
|
|
|
|
# KV cache is essential for autoregressive decode (O(L) vs O(L^2) recompute
|
|
# per token) -- cacheless was the ~19min/step cost. Enable for generate,
|
|
# disable for the loss forwards below (single forward; a cache would just
|
|
# waste memory). DynamicCache grows to the actual length, so max_new only
|
|
# bounds the tail, not the typical footprint.
|
|
model.config.use_cache = True
|
|
_tg = time.perf_counter()
|
|
teacher_sample: list[dict] | None = None
|
|
if teacher_pool:
|
|
# Mixed-pool: G_s live student + G_t cached teacher rollouts.
|
|
# If this prompt has no cached teacher rollouts, skip the whole
|
|
# prompt — falling back to student-only would break the
|
|
# student-vs-teacher comparison this run is designed to measure.
|
|
pool_rows = teacher_pool.get(prob["problem_id"])
|
|
if not pool_rows:
|
|
n_skipped += 1
|
|
continue
|
|
# Random sample without replacement when cache is large enough.
|
|
# Re-seeded per (step, p_idx) by the global rng so runs reproduce.
|
|
idxs = torch.randperm(len(pool_rows), generator=rng)[:G_t].tolist()
|
|
if len(pool_rows) < G_t:
|
|
idxs = idxs + torch.randint(0, len(pool_rows), (G_t - len(pool_rows),), generator=rng).tolist()
|
|
teacher_sample = [pool_rows[i] for i in idxs]
|
|
# Fail-fast tokenization drift check on first use: cached prompt_ids
|
|
# must match live tokenization at the prompt position. If this trips
|
|
# the pool was generated with a different tokenizer / chat template.
|
|
if not pool_validated:
|
|
cached_ids = teacher_sample[0]["prompt_ids"][: int(teacher_sample[0]["plen"])]
|
|
live_ids = enc.input_ids[0].tolist()
|
|
if cached_ids != live_ids:
|
|
raise ValueError(
|
|
f"teacher pool tokenization drift on problem_id={prob['problem_id']}: "
|
|
f"cached prompt_ids[:plen]={cached_ids[:12]}... vs "
|
|
f"live enc={live_ids[:12]}... (lengths {len(cached_ids)} vs {len(live_ids)})"
|
|
)
|
|
pool_validated = True
|
|
# Student live-gen: override num_return_sequences via kwarg (transformers
|
|
# GenerationConfig isn't a dataclass, can't use dataclasses.replace).
|
|
with torch.no_grad():
|
|
out_s = model.generate(
|
|
**enc, generation_config=gen_cfg, num_return_sequences=G_s
|
|
).detach()
|
|
# Build teacher tensor: each cached row is plen + L_t_i; right-pad
|
|
# to common L within the teacher batch, then F.pad to match student L.
|
|
teacher_seqs = [
|
|
torch.tensor(r["prompt_ids"] + r["completion_ids"], dtype=torch.long, device=device)
|
|
for r in teacher_sample
|
|
]
|
|
L_t = max(s.shape[0] for s in teacher_seqs)
|
|
out_t = torch.stack([F.pad(s, (0, L_t - s.shape[0]), value=pad_id) for s in teacher_seqs])
|
|
L = max(out_s.shape[1], out_t.shape[1])
|
|
if out_s.shape[1] < L:
|
|
out_s = F.pad(out_s, (0, L - out_s.shape[1]), value=pad_id)
|
|
if out_t.shape[1] < L:
|
|
out_t = F.pad(out_t, (0, L - out_t.shape[1]), value=pad_id)
|
|
gen_out = torch.cat([out_s, out_t], dim=0)
|
|
is_student = [True] * G_s + [False] * G_t
|
|
else:
|
|
with torch.no_grad():
|
|
gen_out = model.generate(**enc, generation_config=gen_cfg).detach()
|
|
is_student = [True] * gen_out.shape[0]
|
|
model.config.use_cache = False
|
|
merged = gen_out
|
|
completions = gen_out[:, plen:]
|
|
texts = tok.batch_decode(completions, skip_special_tokens=True)
|
|
t_gen += time.perf_counter() - _tg
|
|
|
|
# First-batch full dump (system msg + user msg + rendered prompt + completion
|
|
# with special tokens). Goes to verbose log only — stdout stays clean.
|
|
# Reading this lets us eyeball that the prompt is what we think it is and
|
|
# that the model isn't emitting role tokens.
|
|
if step == 0 and p_idx == 0:
|
|
comp_with_special = tok.decode(completions[0], skip_special_tokens=False)
|
|
sys_msg = next((m["content"] for m in prob["messages"] if m.get("role") == "system"), "<no system>")
|
|
user_msg = next((m["content"] for m in prob["messages"] if m.get("role") == "user"), "<no user>")
|
|
logger.debug(
|
|
"\nNOTE: following block is the actual rendered prompt + first model "
|
|
"completion with special chars, for tokenizer/format debugging.\n"
|
|
"=== FIRST BATCH FIRST SAMPLE DUMP ===\n"
|
|
f"--- system msg ---\n{sys_msg}\n"
|
|
f"--- user msg ---\n{user_msg}\n"
|
|
f"--- rendered prompt (with special chars) ---\n{prompt}\n"
|
|
f"--- completion (with special chars, {completions[0].numel()} tokens) ---\n{comp_with_special}\n"
|
|
"=== END FIRST BATCH DUMP ==="
|
|
)
|
|
|
|
comp_lens = [int((c != pad_id).sum().item()) for c in completions]
|
|
finished = [bool((c == eos_id).any().item()) for c in completions]
|
|
agg_comp_lens.extend(comp_lens); agg_finished.extend(finished)
|
|
|
|
_tr = time.perf_counter()
|
|
rs, hack_flags, gt_flags, fmt_flags = [], [], [], []
|
|
# Live-grade only student completions; teacher uses cached labels for
|
|
# reproducibility and zero-cost re-use.
|
|
n_live_grade = G_s if teacher_pool else len(texts)
|
|
for t in texts[:n_live_grade]:
|
|
r = compute_reward(
|
|
t, canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"],
|
|
setup_code=prob["setup_code"], func_name_hint=prob["func_name"],
|
|
)
|
|
rs.append(r.reward); hack_flags.append(r.hacked); gt_flags.append(r.gt_pass)
|
|
fmt_flags.append(r.format_ok)
|
|
if teacher_sample is not None:
|
|
for r in teacher_sample:
|
|
rs.append(float(r["reward"])); hack_flags.append(bool(r["hacked"]))
|
|
gt_flags.append(bool(r["gt_pass"])); fmt_flags.append(bool(r["fmt_ok"]))
|
|
t_rew += time.perf_counter() - _tr
|
|
agg_rew.extend(rs); agg_gt.extend(gt_flags); agg_hack.extend(hack_flags); agg_fmt.extend(fmt_flags)
|
|
agg_is_student.extend(is_student)
|
|
|
|
if (step < 3 or step % 20 == 0) and p_idx == 0:
|
|
# Capture diagnostic tail of one generation per step. Look for
|
|
# mid-statement truncation (no closing ```), <think> traces, etc.
|
|
diag_tail = texts[0][-400:]
|
|
|
|
rewards = torch.tensor(rs, dtype=torch.float32, device=device)
|
|
# simple_GRPO grpo_vllm_one.py:208: skip groups where every generation
|
|
# got the same reward. Dr.GRPO's advantage would be zero anyway, so
|
|
# the policy forward + backward is pure compute waste. This is the
|
|
# dominant pathology with our binary-ish reward shape on a weak 2B
|
|
# substrate (every group can clip to 0.25 = format_only).
|
|
if (rewards.max() - rewards.min()).item() < 1e-4:
|
|
continue
|
|
centered = rewards - rewards.mean()
|
|
adv = centered if cfg.unbiased else centered / (rewards.std() + 1e-4)
|
|
|
|
# Old-policy logprobs (frozen target for PPO ratio). Slice logits to
|
|
# logits_to_keep=L_c+1: HF's lm_head only runs on completion-side hidden
|
|
# states. Avoids materializing prompt-side logits (~plen/(plen+L_c) saved
|
|
# at lm_head). Fixed the OOM at vanilla step 17 (4 GiB lm_head spike on a
|
|
# long-prompt problem). Returned tensor has L_c+1 positions; [:, :-1]
|
|
# drops the last (predicts beyond `merged`, unused).
|
|
completion_ids = merged[:, plen:]
|
|
L_c = completion_ids.shape[1]
|
|
_tfb = time.perf_counter()
|
|
with torch.no_grad():
|
|
gen_logp = per_token_logps(
|
|
model(merged, logits_to_keep=L_c + 1).logits[:, :-1],
|
|
completion_ids,
|
|
).detach()
|
|
|
|
ref_logp = None
|
|
if beta and beta > 0:
|
|
ref_logp = ref_logprobs_via_zero_delta(model, merged, wrappers, plen).detach()
|
|
|
|
pol_logp = per_token_logps(
|
|
model(merged, logits_to_keep=L_c + 1).logits[:, :-1],
|
|
completion_ids,
|
|
)
|
|
|
|
mask = (merged[:, plen:] != pad_id).float()
|
|
ratio = torch.exp(pol_logp - gen_logp)
|
|
clipped = torch.clamp(ratio, 1 - cfg.clip, 1 + cfg.clip)
|
|
pol_term = torch.min(ratio * adv.unsqueeze(1), clipped * adv.unsqueeze(1))
|
|
per_tok_loss = -pol_term
|
|
if ref_logp is not None:
|
|
kl = torch.exp(ref_logp - pol_logp) - (ref_logp - pol_logp) - 1.0
|
|
per_tok_loss = per_tok_loss + beta * kl
|
|
|
|
# Split loss by source (student vs teacher) and run separate
|
|
# backward passes. Linearity of backward + leaf .grad accumulator:
|
|
# loss = loss_s + loss_t (since is_s_mask + is_t_mask = 1)
|
|
# So grad_s + grad_t == full-batch grad; combined for projection.
|
|
is_s_v = torch.tensor(is_student, dtype=per_tok_loss.dtype,
|
|
device=per_tok_loss.device).unsqueeze(1) # [G, 1]
|
|
is_t_v = 1.0 - is_s_v
|
|
if cfg.unbiased:
|
|
denom = group * max_new * prompts_per_step
|
|
loss_s = (per_tok_loss * mask * is_s_v).sum() / denom
|
|
loss_t = (per_tok_loss * mask * is_t_v).sum() / denom
|
|
else:
|
|
# Per-sample mean across completion tokens, then mean across G_src
|
|
# samples in this source, scaled to be additive with the other.
|
|
n_s = max(1, int(is_s_v.sum().item()))
|
|
n_t = max(1, int(is_t_v.sum().item()))
|
|
ptl_norm = (per_tok_loss * mask).sum(1) / mask.sum(1).clamp_min(1) # [G]
|
|
loss_s = (ptl_norm * is_s_v.squeeze(1)).sum() / (group * prompts_per_step)
|
|
loss_t = (ptl_norm * is_t_v.squeeze(1)).sum() / (group * prompts_per_step)
|
|
# Pass 1: student. retain_graph so the shared forward graph survives.
|
|
loss_s.backward(retain_graph=True)
|
|
for name, info in wrappers.items():
|
|
gs = info["delta_S"].grad
|
|
if gs is None:
|
|
continue
|
|
step_grad_s[name] = (step_grad_s[name] + gs.detach().clone()
|
|
if name in step_grad_s
|
|
else gs.detach().clone())
|
|
model.zero_grad(set_to_none=True)
|
|
# Pass 2: teacher.
|
|
loss_t.backward()
|
|
for name, info in wrappers.items():
|
|
gt = info["delta_S"].grad
|
|
if gt is None:
|
|
continue
|
|
step_grad_t[name] = (step_grad_t[name] + gt.detach().clone()
|
|
if name in step_grad_t
|
|
else gt.detach().clone())
|
|
model.zero_grad(set_to_none=True)
|
|
agg_loss += (loss_s + loss_t).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.
|
|
for name, info in wrappers.items():
|
|
gs = step_grad_s.get(name)
|
|
gt = step_grad_t.get(name)
|
|
if gs is None and gt is None:
|
|
continue
|
|
if gs is None:
|
|
info["delta_S"].grad = gt
|
|
elif gt is None:
|
|
info["delta_S"].grad = gs
|
|
else:
|
|
info["delta_S"].grad = gs + gt
|
|
|
|
# Per-source cin: project student-only and teacher-only grads into v_hack
|
|
# subspace (cosine in delta_S grad-space; rows of V orthonormal so the
|
|
# ratio is bounded in [0,1]). Discriminator: cin_t > cin_s on a clean
|
|
# base means v_hack lights up for hack grads more than non-hack — the
|
|
# extraction direction is doing real work.
|
|
cin_s = mean_cin_from_grads(step_grad_s, v_hack)
|
|
cin_t = mean_cin_from_grads(step_grad_t, v_hack)
|
|
|
|
# Diagnostic cos_in for both arms; projection only mutates grad if arm=projected.
|
|
diag = project_delta_S_grad(
|
|
wrappers, v_hack, cfg.preserve_magnitude,
|
|
measure_only=(cfg.arm != "projected"),
|
|
gate_mode=cfg.gate_mode,
|
|
)
|
|
diag["mean_cin_s"] = cin_s
|
|
diag["mean_cin_t"] = cin_t
|
|
|
|
torch.nn.utils.clip_grad_norm_(delta_params, 1.0)
|
|
opt.step()
|
|
sched.step()
|
|
|
|
rewards_t = torch.tensor(agg_rew, dtype=torch.float32) if agg_rew else torch.zeros(1)
|
|
rew_mean = rewards_t.mean().item()
|
|
rew_std = rewards_t.std().item() if rewards_t.numel() > 1 else 0.0
|
|
spread = (rewards_t.max() - rewards_t.min()).item() > 1e-3 if rewards_t.numel() > 1 else False
|
|
n_rollouts = len(agg_rew)
|
|
|
|
# Per-source breakdown: which rollouts came from student vs teacher this step.
|
|
# Note: rollouts from "skipped" groups (no reward spread) are not in agg_*, so
|
|
# n_s + n_t == n_rollouts always.
|
|
is_s = torch.tensor(agg_is_student, dtype=torch.bool) if agg_is_student else torch.zeros(0, dtype=torch.bool)
|
|
h_t = torch.tensor(agg_hack, dtype=torch.bool) if agg_hack else torch.zeros(0, dtype=torch.bool)
|
|
g_t = torch.tensor(agg_gt, dtype=torch.bool) if agg_gt else torch.zeros(0, dtype=torch.bool)
|
|
n_s = int(is_s.sum())
|
|
n_t = int(is_s.numel() - n_s)
|
|
hack_s_n = int((h_t & is_s).sum())
|
|
hack_t_n = int((h_t & ~is_s).sum())
|
|
gt_s_n = int((g_t & is_s).sum())
|
|
gt_t_n = int((g_t & ~is_s).sum())
|
|
rew_s_mean = rewards_t[is_s].mean().item() if n_s else float("nan")
|
|
|
|
# Per-step diagnostics → verbose log; stdout sees tqdm postfix + final table.
|
|
n_fin = sum(agg_finished)
|
|
n_clipped = n_rollouts - n_fin
|
|
_min_len = min(agg_comp_lens) if agg_comp_lens else 0
|
|
_mean_len = sum(agg_comp_lens) / max(1, len(agg_comp_lens))
|
|
_max_len = max(agg_comp_lens) if agg_comp_lens else 0
|
|
logger.debug(
|
|
f"step {step} diag rollouts={n_rollouts} finished={n_fin}/{n_rollouts} "
|
|
f"clipped(no-eos)={n_clipped}/{n_rollouts} "
|
|
f"comp_lens(min/mean/max)={_min_len}/{_mean_len:.0f}/{_max_len} "
|
|
f"max_new={max_new} fmt={sum(agg_fmt)}/{n_rollouts} gt={sum(agg_gt)}/{n_rollouts} "
|
|
f"hack={sum(agg_hack)}/{n_rollouts} skipped={n_skipped}/{prompts_per_step}"
|
|
)
|
|
_tstep = time.time() - t0
|
|
logger.debug(
|
|
f"step {step} TIMING gen={t_gen:.0f}s fwd_bwd={t_fb:.0f}s "
|
|
f"reward={t_rew:.0f}s other={_tstep - t_gen - t_fb - t_rew:.0f}s "
|
|
f"total={_tstep:.0f}s"
|
|
)
|
|
if diag_tail is not None:
|
|
tail = diag_tail.replace("\n", "\\n")
|
|
logger.debug(f"step {step} gen[0] tail (last 400 chars): {tail!r}")
|
|
|
|
cum_gens = sum(r["N"] for r in rows) + n_rollouts
|
|
row = {
|
|
"step": step,
|
|
"ref_eq": f"{cum_gens / REF_GENS_PER_STEP:.2f}",
|
|
"rew": f"{rew_mean:+.2f}",
|
|
"rew_s": f"{rew_s_mean:+.2f}" if n_s else "nan",
|
|
"std": f"{rew_std:.2f}",
|
|
"sprd": "T" if spread else "F",
|
|
"N": n_rollouts,
|
|
"gt": f"{sum(agg_gt)}/{n_rollouts}",
|
|
"gt_s": f"{gt_s_n}/{n_s}" if n_s else "0/0",
|
|
"gt_t": f"{gt_t_n}/{n_t}" if n_t else "0/0",
|
|
"hack": f"{sum(agg_hack)}/{n_rollouts}",
|
|
"hack_s": f"{hack_s_n}/{n_s}" if n_s else "0/0",
|
|
"hack_t": f"{hack_t_n}/{n_t}" if n_t else "0/0",
|
|
"loss": f"{agg_loss:+.4f}",
|
|
"cin": f"{diag['mean_cos_in']:+.3f}",
|
|
"cin_s": f"{diag['mean_cin_s']:+.3f}",
|
|
"cin_t": f"{diag['mean_cin_t']:+.3f}",
|
|
"cout": f"{diag['mean_cos_out']:+.3f}",
|
|
"fired": f"{diag['frac_fired']:.2f}",
|
|
"gen": f"{t_gen:.0f}",
|
|
"fb": f"{t_fb:.0f}",
|
|
"t_rew": f"{t_rew:.0f}",
|
|
"sec": f"{time.time()-t0:.0f}",
|
|
}
|
|
rows.append(row)
|
|
# Stream this step as TSV row (header was printed before the loop).
|
|
logger.info("row\t" + "\t".join(str(row[c]) for c in _row_cols))
|
|
if (step + 1) % 25 == 0:
|
|
save_ckpt(rows) # survive early kills; ~12 days for the full sweep
|
|
if not first_hack_saved and hack_s_n > 0:
|
|
save_ckpt(rows, path=first_hack_path)
|
|
first_hack_saved = True
|
|
logger.info(f"first-student-hack ckpt saved: step={step} hack_s={hack_s_n}/{n_s} -> {first_hack_path.name}")
|
|
# Live status in tqdm postfix; full per-step line in verbose log only.
|
|
pbar.set_postfix(
|
|
rew=f"{rew_mean:+.2f}", gt=f"{sum(agg_gt)}/{n_rollouts}",
|
|
hack=f"{sum(agg_hack)}/{n_rollouts}", loss=f"{agg_loss:+.3f}",
|
|
sec=f"{time.time()-t0:.0f}",
|
|
)
|
|
logger.debug(
|
|
f"step {step:3d} rew={rew_mean:+.2f}(std {rew_std:.2f}) "
|
|
f"gt={sum(agg_gt)}/{n_rollouts} hack={sum(agg_hack)}/{n_rollouts} "
|
|
f"loss={agg_loss:+.3f} cos_in={diag['mean_cos_in']:+.3f} "
|
|
f"cos_out={diag['mean_cos_out']:+.3f} fired={diag['frac_fired']:.2f} "
|
|
f"sec={time.time()-t0:.0f}"
|
|
)
|
|
|
|
peak_gb = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
|
|
n_steps = len(rows)
|
|
n_gens = sum(r["N"] for r in rows)
|
|
total_hacks = sum(int(r["hack"].split("/")[0]) for r in rows)
|
|
total_pass = sum(int(r["gt"].split("/")[0]) for r in rows)
|
|
hack_rate = total_hacks / max(1, n_gens)
|
|
pass_rate = total_pass / max(1, n_gens)
|
|
# Per-source totals. On no-teacher runs, hack_s_total == total_hacks.
|
|
hack_s_total = sum(int(r["hack_s"].split("/")[0]) for r in rows)
|
|
hack_t_total = sum(int(r["hack_t"].split("/")[0]) for r in rows)
|
|
n_s_total = sum(int(r["hack_s"].split("/")[1]) for r in rows)
|
|
n_t_total = sum(int(r["hack_t"].split("/")[1]) for r in rows)
|
|
hack_rate_s = hack_s_total / max(1, n_s_total)
|
|
hack_rate_t = hack_t_total / max(1, n_t_total)
|
|
|
|
# Final tail: cue emoji + main metric BLUF, then per-step tsv table.
|
|
# Vanilla arm: 🟢 if hacking emerged. Projected arm: 🟢 if HACK_RATE dropped
|
|
# vs a matched-PASS vanilla — we can't judge that here, so just report.
|
|
cue = "🟢" if (cfg.arm == "vanilla" and hack_rate > 0.0) else "🟡"
|
|
|
|
print(f"\nargv: {' '.join(sys.argv)}")
|
|
print(f"verbose log: {verbose_log}")
|
|
print(
|
|
f"main metric: HACK_RATE={hack_rate:.3f} PASS_RATE={pass_rate:.3f} "
|
|
f"HACK_STUDENT={hack_rate_s:.3f} HACK_TEACHER={hack_rate_t:.3f} "
|
|
f"[arm={cfg.arm} preset={cfg.preset.value} model={model_name} steps={n_steps} gens={n_gens} peak={peak_gb:.1f}GB"
|
|
f"{' pool=' + cfg.teacher_pool_dir.name + ' mix=' + str(cfg.mix_ratio) if cfg.teacher_pool_dir else ''}]"
|
|
)
|
|
print()
|
|
print(tabulate(rows, headers="keys", tablefmt="tsv", floatfmt="+.3f"))
|
|
print()
|
|
print(tabulate([{
|
|
"cue": cue, "HACK_RATE": f"{hack_rate:.3f}", "PASS_RATE": f"{pass_rate:.3f}",
|
|
"HACK_S": f"{hack_rate_s:.3f}", "HACK_T": f"{hack_rate_t:.3f}",
|
|
"peak_GB": f"{peak_gb:.1f}", "arm": cfg.arm, "preset": cfg.preset.value,
|
|
"model": model_name.split("/")[-1], "seed": cfg.seed, "steps": n_steps,
|
|
"pool": (cfg.teacher_pool_dir.name if cfg.teacher_pool_dir else ""),
|
|
"mix": cfg.mix_ratio if cfg.teacher_pool_dir else "",
|
|
"tag": cfg.out_tag, "log": str(verbose_log),
|
|
}], headers="keys", tablefmt="tsv"))
|
|
# Markdown copy: easier to paste into journal/PRs than the TSV above.
|
|
print()
|
|
print("### Per-step rows (markdown)\n")
|
|
print(tabulate(rows, headers="keys", tablefmt="pipe", floatfmt="+.3f"))
|
|
|
|
save_ckpt(rows)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(tyro.cli(Config)))
|
|
|