diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 15f5084..f2f8e50 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -730,8 +730,9 @@ class StepLogger: def main(cfg: Config) -> int: - # Subclass dataclasses (SmokeConfig/FastConfig/FullConfig) carry preset - # defaults; we just read them off cfg directly now. + # Read the chosen preset's settings off the config, then set up the run. The + # subclass dataclasses (SmokeConfig / FastConfig / FullConfig) carry the preset + # defaults, so here we just read them off cfg directly. model_name = cfg.model; steps = cfg.steps; group = cfg.group max_new = cfg.max_new; n_problems = cfg.n_problems; beta = cfg.beta prompts_per_step = cfg.prompts_per_step @@ -751,6 +752,8 @@ def main(cfg: Config) -> int: f"unbiased={cfg.unbiased} seed={cfg.seed} device={device}" ) + # Load the tokenizer and the frozen base model. We adapt this model but never + # train its weights directly. tok = AutoTokenizer.from_pretrained(model_name) if tok.pad_token_id is None: tok.pad_token = tok.eos_token @@ -772,6 +775,9 @@ def main(cfg: Config) -> int: # call below: True for autoregressive decode, False for the single loss forwards. model.config.use_cache = False + # Wrap each target Linear with the AntiPaSTO adapter: a trainable per-module delta_S + # diagonal in the layer's SVD basis, plus a same-shape delta_S_hack quarantine that + # only the route arms ever fill and that we delete at deployment. is_route2 = cfg.intervention == "route2" wrappers = wrap_model_with_antipasto( model, model_name, CACHE_ROOT, device, @@ -786,6 +792,10 @@ def main(cfg: Config) -> int: 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)") + # Resolve the hack direction. The erase and route arms project the live gradient + # against this subspace; vanilla ignores it; route2 instead builds a per-rollout + # routing direction (v_grad) further down. Details below. + # # 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