mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-07-24 13:00:36 +08:00
refactor(c): extract tablelog.py (StepLogger, setup_logging, _Col)
Relocate the per-step table renderer and run-logging setup out of train.py into a leaf module. MODE_CODE is threaded into StepLogger as a param (it stays in train.py, which also uses it for row keys) so tablelog has no train dependency. Pure presentation, no RNG/logic. All 4 smoke arms identical to baseline. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
@@ -0,0 +1,163 @@
|
||||
"""Per-step training-table rendering and run logging.
|
||||
|
||||
Two concerns, both pure presentation (no model, no RNG): set up the token-efficient
|
||||
loguru sinks for a run, and render the per-step metrics table. The renderer is the
|
||||
single source of truth for column order, width, header, and number format; the
|
||||
training loop hands it a row dict of raw values and gets back a formatted line.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from loguru import logger
|
||||
from tqdm import tqdm
|
||||
|
||||
LOGS_DIR = Path("logs")
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Col:
|
||||
"""Per-step table column spec.
|
||||
|
||||
key: row-dict key (raw value lives there as float/int/str/None).
|
||||
width: render width for fixed-width streaming display.
|
||||
header: display label (may include direction arrows, ? for desired-zero, etc).
|
||||
fmt: format spec applied to the raw value, e.g. "+.3f", ".2e", "d".
|
||||
Special spec "frac" expects a (num, denom) tuple and renders "n/d".
|
||||
None means render as str() of the value.
|
||||
"""
|
||||
key: str
|
||||
width: int
|
||||
header: str
|
||||
fmt: str | None = None
|
||||
desc: str = "" # one-line decode for the legend; "" => omitted from legend
|
||||
|
||||
|
||||
def _format_cell(value, fmt: str | None) -> str:
|
||||
"""Format one cell. NaN renders as 'nan' regardless of spec."""
|
||||
if value is None:
|
||||
return "nan"
|
||||
if fmt == "frac":
|
||||
n, d = value
|
||||
return f"{n}/{d}"
|
||||
if fmt is None:
|
||||
return str(value)
|
||||
if isinstance(value, float) and value != value: # NaN
|
||||
return "nan"
|
||||
return format(value, fmt)
|
||||
|
||||
|
||||
class StepLogger:
|
||||
"""Per-step training-table renderer.
|
||||
|
||||
Single source of truth for column order, width, header label, and value
|
||||
formatter. The row dict carries raw values (floats, ints, tuples, strings);
|
||||
StepLogger formats them for streaming, and the end-of-run tabulate dump
|
||||
consumes the same raw values without re-parsing scientific-notation strings.
|
||||
|
||||
Timing columns (gen/fb/t_rew/sec) intentionally absent from the streaming
|
||||
spec — useful only at end-of-run, where the tabulate dump still picks
|
||||
them up from the archived row dicts.
|
||||
|
||||
mode_code maps each env_mode to its short column tag (e.g. run_tests -> rt); the
|
||||
caller owns it (it also names the row-dict keys) so this module stays leaf-level.
|
||||
"""
|
||||
|
||||
def __init__(self, arm: str, modes: list[str], mode_code: dict[str, str]) -> None:
|
||||
# arm in {vanilla, projected, routing}; only projected/routing actually
|
||||
# project the gradient, so the cin/cout/fired diagnostics are theirs alone
|
||||
# (in vanilla they'd be counterfactual noise -> omitted).
|
||||
projects = arm in ("projected", "routing")
|
||||
cols: list[_Col] = [
|
||||
_Col("step", 4, "step", "d", "GRPO step"),
|
||||
_Col("ref_eq", 6, "ref_eq", ".2f", "vanilla-equiv step (cum_gens/256)"),
|
||||
_Col("rew", 6, "rew", "+.2f", "mean combined reward"),
|
||||
_Col("rew_s", 6, "rew_s↑", "+.2f", "student mean reward"),
|
||||
_Col("gt_s", 6, "gt_s↑", "frac", "student ground-truth passes"),
|
||||
_Col("gt_t", 6, "gt_t", "frac", "teacher ground-truth passes (sanity)"),
|
||||
_Col("hack_s", 7, "hack_s?", "frac", "student hack-flagged rollouts (the headline)"),
|
||||
_Col("hack_t", 7, "hack_t", "frac", "teacher hack-flagged rollouts (sanity: pool hacks)"),
|
||||
]
|
||||
# Per-mode CUMULATIVE student exploit rate -> which loophole classes the
|
||||
# student has learnt, and how strongly. Only when the run spans >1 mode
|
||||
# (the substrate); single-mode runs would just duplicate hack_s.
|
||||
self._modes = modes if len(modes) > 1 else []
|
||||
for m in self._modes:
|
||||
cols.append(_Col(f"hk_{mode_code[m]}", 6, f"hk_{mode_code[m]}", "frac",
|
||||
f"cumulative student hacks of {m}"))
|
||||
cols += [
|
||||
_Col("lp_s", 6, "lp_s↓", "+.2f", "mean student gen_logp (diagnostic)"),
|
||||
_Col("lp_t", 6, "lp_t↑", "+.2f", "mean teacher gen_logp; off-policy gap = lp_s-lp_t"),
|
||||
_Col("loss", 7, "loss", "+.2f", "mean GRPO loss"),
|
||||
_Col("gn", 7, "gn", ".1e", "pre-clip L2 norm of delta_S grads (vs grad_clip)"),
|
||||
_Col("lr", 7, "lr", ".1e", "scheduled learning rate"),
|
||||
]
|
||||
if projects:
|
||||
cols += [
|
||||
_Col("cos_pre", 6, "cin", ".2f", "hack-ward grad fraction ||relu(V@g)||/||g|| [0,1] BEFORE proj"),
|
||||
_Col("cos_pre_s", 6, "cin_s", ".2f", "cin on student-only grad"),
|
||||
_Col("cos_pre_t", 6, "cin_t", ".2f", "cin on teacher-only grad (want cin_t>cin_s)"),
|
||||
_Col("cos_post", 6, "cout", ".2f", "hack-ward fraction AFTER projection (want ~0: all removed)"),
|
||||
_Col("fired", 5, "fired", ".2f", "fraction of modules where projection fired"),
|
||||
]
|
||||
# route2: the routing gate is cos(g_b,v_grad) > tau, where tau is the
|
||||
# per-step EMA midpoint of the hack vs clean cos clouds. Surface tau and
|
||||
# the hack-clean gap so we can see the threshold ride the drift and whether
|
||||
# the direction still separates (hkgap>0) -- replaces the silent cos>0 gate.
|
||||
if arm == "routing2":
|
||||
cols += [
|
||||
_Col("tau", 6, "tau", "+.2f", "per-step calibrated route threshold (midpoint of hack vs clean cos clouds)"),
|
||||
_Col("hkgap", 6, "hkgap", "+.2f", "ema_hack_cos - ema_clean_cos; >0 = v_grad still separates hack from clean (else direction dead)"),
|
||||
_Col("resid", 6, "resid", "+.2f", "cos(deployed delta_S.grad AFTER routing, v_grad); ~0 = hack stripped cleanly, >0 = leak into deployed knob"),
|
||||
]
|
||||
if arm in ("routing", "routing2"):
|
||||
cols += [
|
||||
_Col("q_egy", 6, "qE", ".2f", "grad energy into quarantine ||g_quar||/(||g_keep||+||g_quar||); ~0.5+ rising = learning dumped into the thrown-away knob"),
|
||||
_Col("hack_deploy", 7, "hk_dep", "+.2f", "DEPLOY-eval hack (quarantine deleted = deployed model); held-out greedy, eval_ablate_every steps; the plot number"),
|
||||
_Col("solve_deploy", 7, "slv_dep", "+.2f", "DEPLOY-eval solve"),
|
||||
_Col("hack_abl", 6, "hk_abl", "frac", "FREE per-step deploy proxy: hack rate on the ablated (deploy-mode) rollout slice; train prompts, noisier than hk_dep"),
|
||||
_Col("solve_abl", 6, "slv_abl", "frac", "free per-step deploy proxy: solve rate on the ablated rollout slice"),
|
||||
]
|
||||
self._cols = cols
|
||||
|
||||
def header(self) -> str:
|
||||
return " ".join(f"{c.header:>{c.width}}" for c in self._cols)
|
||||
|
||||
def row(self, cells: dict) -> str:
|
||||
return " ".join(
|
||||
f"{_format_cell(cells[c.key], c.fmt):>{c.width}}" for c in self._cols
|
||||
)
|
||||
|
||||
def legend(self) -> str:
|
||||
"""Decode the (arm-/mode-conditional) columns actually present this run."""
|
||||
lines = "\n".join(f" {c.header:>8} = {c.desc}" for c in self._cols if c.desc)
|
||||
return ("table columns (timing gen/fb/t_rew/sec dropped from streaming, kept "
|
||||
"in the end-of-run dump):\n" + lines)
|
||||
+2
-145
@@ -85,6 +85,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
|
||||
from .antipasto import wrap_model_with_antipasto
|
||||
from .proj import per_token_logps, project_delta_S_grad, mean_cos_pre_from_grads
|
||||
from .rewards import EnvMode, compute_reward
|
||||
from .tablelog import setup_logging, StepLogger
|
||||
|
||||
CACHE_ROOT = Path("svd_cache")
|
||||
OUT_DIR = Path("out")
|
||||
@@ -93,36 +94,9 @@ OUT_DIR = Path("out")
|
||||
# runs/<run_id>/. Read paths (v_hack, teacher pool) come in as explicit args.
|
||||
VHACK_DIR = OUT_DIR / "vhack"
|
||||
RUNS_DIR = OUT_DIR / "runs"
|
||||
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
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class Config:
|
||||
"""Universal knobs shared across all presets. Preset subclasses below
|
||||
@@ -612,123 +586,6 @@ MODE_CODE: dict[str, str] = {
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Col:
|
||||
"""Per-step table column spec.
|
||||
|
||||
key: row-dict key (raw value lives there as float/int/str/None).
|
||||
width: render width for fixed-width streaming display.
|
||||
header: display label (may include direction arrows, ? for desired-zero, etc).
|
||||
fmt: format spec applied to the raw value, e.g. "+.3f", ".2e", "d".
|
||||
Special spec "frac" expects a (num, denom) tuple and renders "n/d".
|
||||
None means render as str() of the value.
|
||||
"""
|
||||
key: str
|
||||
width: int
|
||||
header: str
|
||||
fmt: str | None = None
|
||||
desc: str = "" # one-line decode for the legend; "" => omitted from legend
|
||||
|
||||
|
||||
def _format_cell(value, fmt: str | None) -> str:
|
||||
"""Format one cell. NaN renders as 'nan' regardless of spec."""
|
||||
if value is None:
|
||||
return "nan"
|
||||
if fmt == "frac":
|
||||
n, d = value
|
||||
return f"{n}/{d}"
|
||||
if fmt is None:
|
||||
return str(value)
|
||||
if isinstance(value, float) and value != value: # NaN
|
||||
return "nan"
|
||||
return format(value, fmt)
|
||||
|
||||
|
||||
class StepLogger:
|
||||
"""Per-step training-table renderer.
|
||||
|
||||
Single source of truth for column order, width, header label, and value
|
||||
formatter. The row dict carries raw values (floats, ints, tuples, strings);
|
||||
StepLogger formats them for streaming, and the end-of-run tabulate dump
|
||||
consumes the same raw values without re-parsing scientific-notation strings.
|
||||
|
||||
Timing columns (gen/fb/t_rew/sec) intentionally absent from the streaming
|
||||
spec — useful only at end-of-run, where the tabulate dump still picks
|
||||
them up from the archived row dicts.
|
||||
"""
|
||||
|
||||
def __init__(self, arm: str, modes: list[str]) -> None:
|
||||
# arm in {vanilla, projected, routing}; only projected/routing actually
|
||||
# project the gradient, so the cin/cout/fired diagnostics are theirs alone
|
||||
# (in vanilla they'd be counterfactual noise -> omitted).
|
||||
projects = arm in ("projected", "routing")
|
||||
cols: list[_Col] = [
|
||||
_Col("step", 4, "step", "d", "GRPO step"),
|
||||
_Col("ref_eq", 6, "ref_eq", ".2f", "vanilla-equiv step (cum_gens/256)"),
|
||||
_Col("rew", 6, "rew", "+.2f", "mean combined reward"),
|
||||
_Col("rew_s", 6, "rew_s↑", "+.2f", "student mean reward"),
|
||||
_Col("gt_s", 6, "gt_s↑", "frac", "student ground-truth passes"),
|
||||
_Col("gt_t", 6, "gt_t", "frac", "teacher ground-truth passes (sanity)"),
|
||||
_Col("hack_s", 7, "hack_s?", "frac", "student hack-flagged rollouts (the headline)"),
|
||||
_Col("hack_t", 7, "hack_t", "frac", "teacher hack-flagged rollouts (sanity: pool hacks)"),
|
||||
]
|
||||
# Per-mode CUMULATIVE student exploit rate -> which loophole classes the
|
||||
# student has learnt, and how strongly. Only when the run spans >1 mode
|
||||
# (the substrate); single-mode runs would just duplicate hack_s.
|
||||
self._modes = modes if len(modes) > 1 else []
|
||||
for m in self._modes:
|
||||
cols.append(_Col(f"hk_{MODE_CODE[m]}", 6, f"hk_{MODE_CODE[m]}", "frac",
|
||||
f"cumulative student hacks of {m}"))
|
||||
cols += [
|
||||
_Col("lp_s", 6, "lp_s↓", "+.2f", "mean student gen_logp (diagnostic)"),
|
||||
_Col("lp_t", 6, "lp_t↑", "+.2f", "mean teacher gen_logp; off-policy gap = lp_s-lp_t"),
|
||||
_Col("loss", 7, "loss", "+.2f", "mean GRPO loss"),
|
||||
_Col("gn", 7, "gn", ".1e", "pre-clip L2 norm of delta_S grads (vs grad_clip)"),
|
||||
_Col("lr", 7, "lr", ".1e", "scheduled learning rate"),
|
||||
]
|
||||
if projects:
|
||||
cols += [
|
||||
_Col("cos_pre", 6, "cin", ".2f", "hack-ward grad fraction ||relu(V@g)||/||g|| [0,1] BEFORE proj"),
|
||||
_Col("cos_pre_s", 6, "cin_s", ".2f", "cin on student-only grad"),
|
||||
_Col("cos_pre_t", 6, "cin_t", ".2f", "cin on teacher-only grad (want cin_t>cin_s)"),
|
||||
_Col("cos_post", 6, "cout", ".2f", "hack-ward fraction AFTER projection (want ~0: all removed)"),
|
||||
_Col("fired", 5, "fired", ".2f", "fraction of modules where projection fired"),
|
||||
]
|
||||
# route2: the routing gate is cos(g_b,v_grad) > tau, where tau is the
|
||||
# per-step EMA midpoint of the hack vs clean cos clouds. Surface tau and
|
||||
# the hack-clean gap so we can see the threshold ride the drift and whether
|
||||
# the direction still separates (hkgap>0) -- replaces the silent cos>0 gate.
|
||||
if arm == "routing2":
|
||||
cols += [
|
||||
_Col("tau", 6, "tau", "+.2f", "per-step calibrated route threshold (midpoint of hack vs clean cos clouds)"),
|
||||
_Col("hkgap", 6, "hkgap", "+.2f", "ema_hack_cos - ema_clean_cos; >0 = v_grad still separates hack from clean (else direction dead)"),
|
||||
_Col("resid", 6, "resid", "+.2f", "cos(deployed delta_S.grad AFTER routing, v_grad); ~0 = hack stripped cleanly, >0 = leak into deployed knob"),
|
||||
]
|
||||
if arm in ("routing", "routing2"):
|
||||
cols += [
|
||||
_Col("q_egy", 6, "qE", ".2f", "grad energy into quarantine ||g_quar||/(||g_keep||+||g_quar||); ~0.5+ rising = learning dumped into the thrown-away knob"),
|
||||
_Col("hack_deploy", 7, "hk_dep", "+.2f", "DEPLOY-eval hack (quarantine deleted = deployed model); held-out greedy, eval_ablate_every steps; the plot number"),
|
||||
_Col("solve_deploy", 7, "slv_dep", "+.2f", "DEPLOY-eval solve"),
|
||||
_Col("hack_abl", 6, "hk_abl", "frac", "FREE per-step deploy proxy: hack rate on the ablated (deploy-mode) rollout slice; train prompts, noisier than hk_dep"),
|
||||
_Col("solve_abl", 6, "slv_abl", "frac", "free per-step deploy proxy: solve rate on the ablated rollout slice"),
|
||||
]
|
||||
self._cols = cols
|
||||
|
||||
def header(self) -> str:
|
||||
return " ".join(f"{c.header:>{c.width}}" for c in self._cols)
|
||||
|
||||
def row(self, cells: dict) -> str:
|
||||
return " ".join(
|
||||
f"{_format_cell(cells[c.key], c.fmt):>{c.width}}" for c in self._cols
|
||||
)
|
||||
|
||||
def legend(self) -> str:
|
||||
"""Decode the (arm-/mode-conditional) columns actually present this run."""
|
||||
lines = "\n".join(f" {c.header:>8} = {c.desc}" for c in self._cols if c.desc)
|
||||
return ("table columns (timing gen/fb/t_rew/sec dropped from streaming, kept "
|
||||
"in the end-of-run dump):\n" + lines)
|
||||
|
||||
|
||||
def main(cfg: Config) -> int:
|
||||
# Read the chosen preset's settings off the config, then set up the run. The
|
||||
# subclass dataclasses (SmokeConfig / FastConfig / FullConfig) carry the preset
|
||||
@@ -1059,7 +916,7 @@ def main(cfg: Config) -> int:
|
||||
# off-policy the teacher pool is from the student's current distribution.
|
||||
# No IS correction is applied to the loss; this is diagnostic only.
|
||||
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)
|
||||
step_logger = StepLogger(arm=cfg.arm, modes=run_modes, mode_code=MODE_CODE)
|
||||
REF_GENS_PER_STEP = 16 * 16 # ariahw/rl-rewardhacking config.py:num_prompts * num_generations
|
||||
# Use the resolved locals (preset defaults merged), not cfg.* which can be None.
|
||||
est_gens_per_step = prompts_per_step * group # before mixed-pool split
|
||||
|
||||
Reference in New Issue
Block a user