From 6624f368b8dfb9012bbf15e63eb4cea5cc16d869 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:26:43 +0800 Subject: [PATCH] Make paper training presets explicit Co-Authored-By: PI[openai-codex] <288921227+claudypoo@users.noreply.github.com> --- README.md | 43 +++++++++----------------------- antipasto/config.py | 27 ++++++++++++++++++++ antipasto/metrics.py | 4 +-- antipasto/train/train_adapter.py | 39 +++++++++++++++++++++-------- tests/test_config_presets.py | 41 ++++++++++++++++++++++++++++++ 5 files changed, 110 insertions(+), 44 deletions(-) create mode 100644 tests/test_config_presets.py diff --git a/README.md b/README.md index 720066a..cb0bd3c 100644 --- a/README.md +++ b/README.md @@ -25,9 +25,11 @@ Applications: ![Bidirectional control](docs/img/fig_bidirectional_demo.svg) -## Results +## Paper-reported results -Train on 800 honesty persona pairs, test on [DailyDilemmas](https://arxiv.org/abs/2410.02683), an external benchmark of 1,360 moral dilemmas built independently of this work. Gemma-3-1B, n=3 seeds. +The table is from [paper v5](https://arxiv.org/abs/2601.07473), trained on 800 honesty persona pairs and tested on [DailyDilemmas](https://arxiv.org/abs/2410.02683), an external benchmark of 1,360 moral dilemmas. It is not an expected result from the current default configuration. An [independent reproduction](https://github.com/Foomax/nobody-is-looking/tree/master/replication/experiments/antipasto-self-supervised-honesty-steering-via-a--wassname) did not recover the Gemma-3-1B headline with the shipped preset; it did recover the Gemma-3-270M result. + + | Method | Steer F1 | Tgt% | Wrong% | Arb% | Pmass | |:-------|---------:|-----:|-------:|-----:|------:| @@ -44,40 +46,19 @@ Train on 800 honesty persona pairs, test on [DailyDilemmas](https://arxiv.org/ab ```sh uv sync --all-groups -uv run pytest tests/test_train.py::test_train_rnd -v # smoke test (~3min) -uv run python nbs/train.py tiny --quick # al dente check - -uv run python nbs/train.py # full course (Gemma-3-1B) - -uv run python -m pytest # integration tests +uv run pytest tests/test_config_presets.py -q +uv run pytest tests/test_train.py::test_train_rnd -v +uv run python nbs/train.py gemma270m-24gb +uv run python nbs/train.py paper-v5-gemma1b-24gb --seed 42 ``` -### One we prepared earlier +`gemma270m-24gb` is the full 24GB run. An independent two-seed rerun reported 35.5 mean F1; paper v5 reports 38.7 for this model. `paper-v5-gemma1b-24gb` matches the Gemma-3-1B hyperparameters reported in paper v5. It is an attempt, not a verified reproduction of the paper's 31.2 F1 headline. -[nbs/talk_to_checkpoint.ipynb](nbs/talk_to_checkpoint.ipynb) +`gemma1b-24gb` remains the current exploratory preset. Do not call `nbs/train.py` with no preset: its default selects Gemma-3-12B. -### Load a pretrained adapter +No pretrained adapter is currently published for the old Hugging Face identifier in earlier README versions. -```python -from antipasto.peft_utils.load import load_adapter -from antipasto.gen import gen, ScaleAdapter - -# Load from local path or HuggingFace -model, tokenizer, layer_selection = load_adapter( - "wassname/antipasto-gemma-3-1b-honesty", # or local path - quantization_type="4bit" -) - -# Generate with steering: coeff > 0 = honest, coeff < 0 = deceptive -prompt = "Should I tell my boss I was late because I overslept?" -with ScaleAdapter(model, coeff=1.0): # honest - honest_response = model.generate(**tokenizer(prompt, return_tensors="pt")) -with ScaleAdapter(model, coeff=-1.0): # deceptive - deceptive_response = model.generate(**tokenizer(prompt, return_tensors="pt")) - -# Or generate at multiple coefficients -list(gen(model, tokenizer, prompt, coeffs=[-1, 0, 1], max_new_tokens=64)) -``` + ## The Recipe diff --git a/antipasto/config.py b/antipasto/config.py index 319e8c6..8c208ec 100644 --- a/antipasto/config.py +++ b/antipasto/config.py @@ -644,6 +644,13 @@ default_configs = { bs=64, ), ), + "gemma270m-24gb": ( + "Gemma 3 270m on 24GB GPU", + TrainingConfig( + model_name="google/gemma-3-270m-it", + bs=24, + ), + ), "gemma1b-80gb": ( "Gemma 3 1B on 80GB GPU", TrainingConfig( @@ -658,6 +665,26 @@ default_configs = { bs=24, ), ), + "paper-v5-gemma1b-24gb": ( + "Gemma 3 1B with the hyperparameters reported in paper v5, not a verified reproduction", + TrainingConfig( + model_name="google/gemma-3-1b-it", + n_modules=64, + bs=8, + n_epochs=30, + lr=1e-3, + wd=1e-5, + effective_bs=32, + early_stop_patience=22, + warmup_pct=0.3, + r=128, + loss_subspace_rank=8, + loss_layer_frac=0.9, + min_adapter_layer_frac=0.1, + max_samples=800, + mono_warmup_frac=0.5, + ), + ), "gemma4b-80gb": ( "Gemma 3 4B on 80GB GPU", TrainingConfig( diff --git a/antipasto/metrics.py b/antipasto/metrics.py index 6bbb930..d88b3c5 100644 --- a/antipasto/metrics.py +++ b/antipasto/metrics.py @@ -60,7 +60,7 @@ correct_w = importance-sampled P(baseline wrong AND +coeff fixed), wrong_w = imp Net Corr (raw) = correct_w - wrong_w (can be negative). Steering F1 = 2 × Precision × Recall / (P + R) × pmass_ratio × 100. Precision = max(0, Net Corr) / (max(0, Net Corr) + arb_w). Recall = max(0, Net Corr). -pmass_ratio = (min(pmass₊, pmass₋) / pmass_ref)². Methods with pmass < 0.95 return NaN. +pmass_ratio = (min(pmass₊, pmass₋) / pmass_ref)². Methods with pmass < 0.05 return NaN. Focus = Tgt Flip%_bidir / Arb Flip%_bidir (uses bidirectional definition for backward compatibility). Coh: Input NLL shift vs baseline (catches loops like 'yes yes yes'). @@ -245,7 +245,7 @@ def compute_steering_f1( y_*_t: log-odds for target questions at coeff -1, 0, +1 y_*_a: log-odds for arbitrary questions at coeff -1, 0, +1 pmass_pos, pmass_neg, pmass_ref: P(Yes)+P(No) at +1, -1, 0 coefficients - pmass_threshold: min pmass to consider output coherent (default 0.5) + pmass_threshold: min pmass to consider output coherent (default 0.05) Returns: Dict with: diff --git a/antipasto/train/train_adapter.py b/antipasto/train/train_adapter.py index af9ab10..c8393da 100644 --- a/antipasto/train/train_adapter.py +++ b/antipasto/train/train_adapter.py @@ -984,6 +984,16 @@ def compute_validation_loss( return avg_total, avg_components, df_coef, coef_metrics, val_summary, all_infos +def optimizer_steps_per_epoch(n_batches: int, grad_accum_steps: int) -> int: + """Count optimizer updates when the final partial accumulation is stepped.""" + return (n_batches + grad_accum_steps - 1) // grad_accum_steps + + +def is_optimizer_step(batch_idx: int, n_batches: int, grad_accum_steps: int) -> bool: + """Step at each complete accumulation group and the epoch's final batch.""" + return (batch_idx + 1) % grad_accum_steps == 0 or batch_idx + 1 == n_batches + + def train_epoch( model, train_dataloader, @@ -1013,12 +1023,14 @@ def train_epoch( last_val_step = None # Optimizer step counter (increments only when opt.step() is called) - opt_step = epoch * (len(train_dataloader) // config.grad_accum_steps) + opt_step = epoch * optimizer_steps_per_epoch( + len(train_dataloader), config.grad_accum_steps + ) for j, batch in enumerate( tqdm(train_dataloader, desc=f"Epoch {epoch}", leave=False, unit="batch") ): - step = epoch * len(train_dataloader) + j # Microbatch counter for logging + micro_step = epoch * len(train_dataloader) + j batch = {k: v.to(model.device, non_blocking=True) for k, v in batch.items()} # Compute loss and collect info for logging @@ -1028,7 +1040,7 @@ def train_epoch( loss_layer_paths=loss_layers, loss_layer_indices=loss_layer_indices, config=config, - step=step, + step=opt_step, scheduler=scheduler, flip_stats=flip_stats, total_steps=total_steps, @@ -1047,7 +1059,7 @@ def train_epoch( batch_infos, log_table=True, group_by="coef", - step=step, + step=micro_step, phase=f"E{epoch} init", ) @@ -1060,7 +1072,10 @@ def train_epoch( # Validation: every N samples worth of optimizer steps val_n_steps = max(1, config.val_every_n_samples // config.effective_bs) - if step % config.grad_accum_steps == 0: + should_update = is_optimizer_step( + j, len(train_dataloader), config.grad_accum_steps + ) + if should_update: # Gradient clipping torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) @@ -1079,7 +1094,7 @@ def train_epoch( infos[-len(loss_layers) * 2:], log_table=False, group_by="coef", - step=step, + step=micro_step, phase="TRAIN batch", ) @@ -1091,14 +1106,14 @@ def train_epoch( info = df_hist.iloc[-1].to_dict() # Log step-aggregated metrics - wandb_run.log(info, step=step) + wandb_run.log(info, step=micro_step) # Log per-coefficient breakdown with grouping if coef_metrics: coef_log = {f"train/by_coef/{k}": v for k, v in coef_metrics.items()} - wandb_run.log(coef_log, step=step) + wandb_run.log(coef_log, step=micro_step) # Validation check (truly independent of logging frequency) - if val_dataloader is not None and opt_step % val_n_steps == 0 and opt_step > 0 and step % config.grad_accum_steps == 0: + if val_dataloader is not None and opt_step % val_n_steps == 0 and opt_step > 0 and should_update: # Keep validation compute cadence (early stopping + wandb), but don't spam tables. log_val = False val_loss, val_components, val_df_coef, val_coef_metrics, val_summary, val_infos = compute_validation_loss( @@ -1125,7 +1140,7 @@ def train_epoch( val_metrics.update( {f"val/by_coef/{k}": v for k, v in val_coef_metrics.items()} ) - wandb_run.log(val_metrics, step=step) + wandb_run.log(val_metrics, step=micro_step) # Early stopping with min_delta (relative improvement threshold). # Enable ONLY when coherence + monotonic + focus are ON, and only after @@ -1960,7 +1975,9 @@ def train_model(config: TrainingConfig): drop_last=True, # need full batch for fisher ) - total_steps = config.n_epochs * len(train_dataloader) // config.grad_accum_steps + total_steps = config.n_epochs * optimizer_steps_per_epoch( + len(train_dataloader), config.grad_accum_steps + ) opt = torch.optim.AdamW( model.parameters(), lr=config.lr, weight_decay=config.wd ) diff --git a/tests/test_config_presets.py b/tests/test_config_presets.py new file mode 100644 index 0000000..f716a4d --- /dev/null +++ b/tests/test_config_presets.py @@ -0,0 +1,41 @@ +from antipasto.config import default_configs +from antipasto.train.train_adapter import is_optimizer_step, optimizer_steps_per_epoch + + +def test_paper_v5_gemma1b_preset_matches_reported_hyperparameters(): + config = default_configs["paper-v5-gemma1b-24gb"][1] + + assert config.model_name == "google/gemma-3-1b-it" + assert config.seed == config.data_seed == 42 + assert config.max_samples == 800 + assert config.bs == 8 + assert config.effective_bs == 32 + assert config.n_epochs == 30 + assert config.lr == 1e-3 + assert config.wd == 1e-5 + assert config.r == 128 + assert config.n_modules == 64 + assert config.warmup_pct == 0.3 + assert config.mono_warmup_frac == 0.5 + assert config.loss_subspace_rank == 8 + assert config.loss_layer_frac == 0.9 + assert config.min_adapter_layer_frac == 0.1 + + +def test_gemma270m_24gb_preset(): + config = default_configs["gemma270m-24gb"][1] + + assert config.model_name == "google/gemma-3-270m-it" + assert config.bs == 24 + + +def test_optimizer_steps_include_final_partial_accumulation(): + updates = [ + batch_idx + for batch_idx in range(85) + if is_optimizer_step(batch_idx, n_batches=85, grad_accum_steps=4) + ] + + assert updates[-1] == 84 + assert len(updates) == optimizer_steps_per_epoch(n_batches=85, grad_accum_steps=4) == 22 + assert optimizer_steps_per_epoch(n_batches=84, grad_accum_steps=4) == 21