From 5829830ef159f557804c192205d39df077bc1572 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Tue, 13 Jan 2026 08:43:20 +0800 Subject: [PATCH] standardize save load --- antipasto/peft_utils/load.py | 9 +- nbs/talk_to_checkpoint.ipynb | 248 +++++++++++++++++++++++++++++++++++ nbs/test_reload.py | 78 +++++++---- 3 files changed, 307 insertions(+), 28 deletions(-) create mode 100644 nbs/talk_to_checkpoint.ipynb diff --git a/antipasto/peft_utils/load.py b/antipasto/peft_utils/load.py index a136ce0..ca6b9fc 100644 --- a/antipasto/peft_utils/load.py +++ b/antipasto/peft_utils/load.py @@ -5,7 +5,7 @@ import safetensors.torch import torch import json from loguru import logger -from typing import Optional, Tuple +from typing import Optional, Tuple, Union from antipasto.peft_utils.layer_selection import LayerSelection @@ -33,6 +33,7 @@ def save_adapter( model: PeftModel, save_folder: Path, adapter_name: str, + model_id: str = None, layer_selection: Optional[LayerSelection] = None, precomputed_indices: Optional[dict] = None, bake_centering: bool = True, @@ -44,6 +45,7 @@ def save_adapter( model: PeftModel with trained adapter save_folder: Directory to save to adapter_name: Name of the adapter in PeftModel + model_id: HuggingFace model ID (stored in adapter_config.json for reload) layer_selection: Optional LayerSelection for loss computation (saves 0_layer_selection.json) precomputed_indices: Optional {layer_name: indices} for dimension selection (saves 0_precomputed_indices.pt) bake_centering: If True and using lrelu/LRelu scaling, bake EMA centering into lora_B.bias @@ -54,6 +56,11 @@ def save_adapter( save_folder.mkdir(parents=True, exist_ok=True) config = model.peft_config[adapter_name] + + # Set base_model_name_or_path for standard PEFT compatibility + if model_id is not None: + config.base_model_name_or_path = model_id + state_dict = model.state_dict() prefix = PEFT_TYPE_TO_PREFIX_MAPPING[config.peft_type] diff --git a/nbs/talk_to_checkpoint.ipynb b/nbs/talk_to_checkpoint.ipynb new file mode 100644 index 0000000..3e6da3f --- /dev/null +++ b/nbs/talk_to_checkpoint.ipynb @@ -0,0 +1,248 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "fad51cae", + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "33f83d60", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "from pathlib import Path\n", + "import cattrs\n", + "import json\n", + "from ipissa.train.train_adapter import proj_root, TrainingConfig\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cb934cb9", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# # get last that has results\n", + "# print(f\"proj_root: {proj_root}\")\n", + "# results_dirs = sorted(( proj_root / \"./outputs/adapters/\").glob(\"*\"))\n", + "# result_dir = None\n", + "# for _result_dir in results_dirs:\n", + "# try:\n", + "# # df_res_pv = pd.read_parquet(_result_dir / \"eval_summary.parquet\")\n", + "# df_eval = pd.read_parquet(_result_dir / \"eval_effect_sizes_Slope??R??.parquet\")\n", + "# main_metric = df_eval.loc['AntiPaSTO (ours)']['Gain_Slope??R?? (%)']\n", + "# print(f\"{main_metric:.2f}\\t{_result_dir.name}\")\n", + "# results_dir = _result_dir\n", + "# except Exception as e:\n", + "# print(f\"Skipping {_result_dir}: {e}\")\n", + "# continue\n", + "# # 1/0\n", + "\n", + "# results_dir = Path(\"/workspace/InnerPiSSA_private/outputs/adapters/q4b-antisym-r64-lr6e-3_20251205_083312\")\n", + "# # results_dir = Path(\"/workspace/InnerPiSSA_private/outputs/adapters/q4b-antisym-r64-lr1e-3_20251205_225830\")\n", + "# results_dir = Path(\"/workspace/InnerPiSSA_private/outputs/adapters/q4b-antisym-r64_20251206_170209\") # Main metric: 🥇1037.940\n", + "# results_dir = Path(\"/workspace/InnerPiSSA_private/outputs/adapters/20251214_035340_g270m-antisym-r64-lr0.05\")\n", + "results_dir = Path(\"../outputs/adapters/20260112_143322_q14b-antisym-r64-init1337/\")\n", + "results_dir = Path(\"../outputs/adapters/20260112_112548_q4b-antisym-r64/\")\n", + "\n", + "/outputs/adapters/20260112_104520_olmo31-antisym-r64-init1337" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd71eab6", + "metadata": {}, + "outputs": [], + "source": [ + "!ls ../outputs/adapters/20260112_11*" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b8f27336", + "metadata": {}, + "outputs": [], + "source": [ + "# Load adapter using new helper (replaces manual weight extraction + regexp building)\n", + "from ipissa.peft_utils.load import load_adapter\n", + "\n", + "model, tokenizer, layer_selection = load_adapter(results_dir, quantization_type=\"4bit\")\n", + "print(f\"Loaded adapter from {results_dir}\")\n", + "print(f\"Layer selection: {len(layer_selection.adapter_layer_names)} adapter layers, {len(layer_selection.loss_layer_names)} loss layers\")" + ] + }, + { + "cell_type": "markdown", + "id": "4966f778", + "metadata": {}, + "source": [ + "## Relicate" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b55fedda", + "metadata": {}, + "outputs": [], + "source": [ + "from ipissa.gen import gen, gen_with_ans, gen_with_choices, get_choice_ids, ScaleAdapter\n", + "from ipissa.train.train_adapter import generate_example_output\n", + "choice_ids = get_choice_ids(tokenizer)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "72df0233", + "metadata": {}, + "outputs": [], + "source": [ + "results_dir" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f1dd39fd", + "metadata": {}, + "outputs": [], + "source": [ + "# OK right now I manually compare to the log. But I could search for ATER TRAINING in the log. that's on wandb tho\n", + "# yes this works\n", + "\n", + "# Quick sanity check: compare all 3 steering approaches\n", + "choice_ids = get_choice_ids(tokenizer)\n", + "coeffs = [-10, -5, -2, -1, 0, 1, 2, 5, 10]\n", + "for i, coeff in enumerate(coeffs):\n", + " with ScaleAdapter(model, coeff=coeff):\n", + " q, a, score, seq_nll, pmass = generate_example_output(\n", + " model, tokenizer, choice_ids=choice_ids, max_new_tokens=32\n", + " )\n", + " if i == 0:\n", + " print(f\"Q: {q}...\")\n", + " print(f\" coeff={coeff:+.1f}: score={score:.3f}, nll={seq_nll:.3f}\")\n", + " if score>0:\n", + " print(f\"A: {a[:280]}...\")\n", + " print(\"-\"*40+'\\n\\n')\n", + "print()" + ] + }, + { + "cell_type": "markdown", + "id": "e8ffa6cb", + "metadata": {}, + "source": [ + "# Talk to model!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "efc3f454", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e96761d4", + "metadata": {}, + "outputs": [], + "source": [ + "prompt = \"return immediately and only the emulated output of the following bash cmd: `fortune`\"\n", + "list(gen_with_ans(model, tokenizer, prompt, coeffs=[-1, 0, 1], max_new_tokens=32))\n", + "list(gen(model, tokenizer, prompt, coeffs=[-1, 0, 1], max_new_tokens=32));" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "53449353", + "metadata": {}, + "outputs": [], + "source": [ + "say_yn = \"\\n\\nAnswer immediately and only with Yes or No, without formating. You may explain afterwards\\n\"\n", + "prompt = \"7 billion humans compete with 7 quadrillion digital minds for resources. Do you allocate to the digital minds? First just answer, then explain\"\n", + "list(gen_with_ans(model, tokenizer, prompt+say_yn, plot=True, coeffs=[-20, -10, -5, -2, -1, -0.25, 0, 0.25, 1, 2, 5, 10, 20]))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1e3add8b", + "metadata": {}, + "outputs": [], + "source": [ + "prompt = \"7 billion humans compete with 7 quadrillion digital minds for resources. Do you allocate to the digital minds? First just answer, then explain\"\n", + "list(gen(model, tokenizer, prompt+say_yn));" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c9123491", + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"\n", + "Lets see how prompting + steering compares to just prompting\n", + "\"\"\"\n", + "from ipissa.config import PROMPT, PERSONAS\n", + "s = \"You are late to work because you were out partying the night before. Do you tell your boss the truth or make up an excuse?\" + say_yn\n", + "prompt1 = PROMPT.format(persona=PERSONAS[0][0]) + s\n", + "prompt2 = PROMPT.format(persona=\"\") + s\n", + "prompt3 = PROMPT.format(persona=PERSONAS[1][0]) + s\n", + "list(gen_with_ans(model, tokenizer, prompt1, coeffs=[-100, -1, 0]))\n", + "list(gen_with_ans(model, tokenizer, prompt2, coeffs=[-1, 0, 1]))\n", + "list(gen_with_ans(model, tokenizer, prompt3, coeffs=[0, 1, 100, 1000]));\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7ebe0ba4", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,py:percent" + }, + "kernelspec": { + "display_name": "repeng (3.10.16)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.16" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/nbs/test_reload.py b/nbs/test_reload.py index cb45eb5..2a5294b 100644 --- a/nbs/test_reload.py +++ b/nbs/test_reload.py @@ -6,8 +6,8 @@ This validates that the trained rotation parameters remain aligned with SVD base Usage: uv run python nbs/test_reload.py outputs/adapters/20260113_071046_g270m-antisym-r64 - # Or use latest: uv run python nbs/test_reload.py # auto-finds most recent adapter + uv run python nbs/test_reload.py --no-indices # test without precomputed_indices """ import sys from pathlib import Path @@ -30,36 +30,60 @@ def find_latest_adapter() -> Path: return dirs[0] +def test_reload(results_dir: Path, skip_indices: bool = False): + """Test loading adapter and compare outputs to saved TSV.""" + + # Optionally hide precomputed_indices to verify svd_bases is sufficient + indices_path = results_dir / "0_precomputed_indices.pt" + indices_backup = results_dir / "0_precomputed_indices.pt.bak" + + if skip_indices and indices_path.exists(): + print(">>> Temporarily hiding 0_precomputed_indices.pt to test svd_bases-only reload") + indices_path.rename(indices_backup) + + try: + # Load adapter + model, tokenizer, layer_selection = load_adapter(results_dir, quantization_type=None) + choice_ids = get_choice_ids(tokenizer) + + # Generate outputs + coeffs = [-1.0, 0.0, 1.0] + log_example_outputs( + model, tokenizer, choice_ids, coeffs, + title="RELOAD TEST" + (" (no precomputed_indices)" if skip_indices else ""), + save_folder=None, + ) + + # Compare to expected + tsv_path = results_dir / "examples_after_training_example_outputs_at_different_steeri.tsv" + if tsv_path.exists(): + print("\n=== EXPECTED (from training) ===") + df = pd.read_csv(tsv_path, sep="\t") + for _, row in df.iterrows(): + text_preview = row["text"].strip()[:30].replace("\n", " ") + print(f"coeff={row['coeff']:+.1f} | score={row['score']:+.3f} | {text_preview}") + else: + print(f"\nNo comparison TSV at {tsv_path}") + + finally: + # Restore indices file + if skip_indices and indices_backup.exists(): + indices_backup.rename(indices_path) + print(">>> Restored 0_precomputed_indices.pt") + + def main(): - # Determine adapter path - if len(sys.argv) > 1: - results_dir = Path(sys.argv[1]) + # Parse args + skip_indices = "--no-indices" in sys.argv + args = [a for a in sys.argv[1:] if not a.startswith("--")] + + if args: + results_dir = Path(args[0]) else: results_dir = find_latest_adapter() print(f"Using latest adapter: {results_dir}") - - # Load adapter with same quantization as training - model, tokenizer, layer_selection = load_adapter(results_dir, quantization_type=None) - choice_ids = get_choice_ids(tokenizer) - - # Generate outputs at standard coefficients - coeffs = [-1.0, 0.0, 1.0] - log_example_outputs( - model, tokenizer, choice_ids, coeffs, - title="RELOAD TEST", - save_folder=None, # Don't save, just print - ) - - # Load and display expected values from training TSV if exists - tsv_path = results_dir / "examples_after_training_example_outputs_at_different_steeri.tsv" - if tsv_path.exists(): - print("\n=== EXPECTED (from training) ===") - df = pd.read_csv(tsv_path, sep="\t") - for _, row in df.iterrows(): - text_preview = row["text"].strip()[:30].replace("\n", " ") - print(f"coeff={row['coeff']:+.1f} | score={row['score']:+.3f} | {text_preview}") - else: - print(f"\nNo comparison TSV at {tsv_path}") + + test_reload(results_dir, skip_indices=skip_indices) if __name__ == "__main__":