From b9d5b45bbfe6ba9b08cd923357dc06444b8cf379 Mon Sep 17 00:00:00 2001 From: Vinay-Umrethe Date: Sun, 7 Jun 2026 14:10:45 +0530 Subject: [PATCH] fix: fix issues in automatic reproduction system (#352) * fix: Check if a model is gated / accessible * fix: handle unknown gated models * feat: Auto install requirements * simplify * Revert "simplify" This reverts commit 10287926e99e5543f67a72d38a595ae2b4084d71. * Revert "feat: Auto install requirements" This reverts commit f4be1abd043e17d83e589e54972c4ead2600c2b2. * fix: Seed pytorch method * reference, style * simplify token * feat: Export strategy in reproduce.json, v2 * style: Name * simplify export strategy * style: Rename * enumeration * maybe remove seed as well * fix: don't lock settings with permanent strategy * simplify no choice, use try/finally block --- config.default.toml | 4 --- src/heretic/config.py | 10 ++++++++ src/heretic/main.py | 53 ++++++++++++++++++++++++---------------- src/heretic/model.py | 4 +++ src/heretic/reproduce.py | 13 +++++++++- src/heretic/utils.py | 2 +- 6 files changed, 59 insertions(+), 27 deletions(-) diff --git a/config.default.toml b/config.default.toml index 6ec8e8e..c0433a8 100644 --- a/config.default.toml +++ b/config.default.toml @@ -123,10 +123,6 @@ n_trials = 200 # Number of trials that use random sampling for the purpose of exploration. n_startup_trials = 60 -# Random seed for reproducible optimization. Set to an integer to enable. -# Applies to Python's random module, NumPy, PyTorch, and Optuna. -# seed = 75 - # Directory to save and load study progress to/from. study_checkpoint_dir = "checkpoints" diff --git a/src/heretic/config.py b/src/heretic/config.py index cfaf44e..724b32e 100644 --- a/src/heretic/config.py +++ b/src/heretic/config.py @@ -32,6 +32,11 @@ class RowNormalization(str, Enum): FULL = "full" +class ExportStrategy(str, Enum): + MERGE = "merge" + ADAPTER = "adapter" + + class DatasetSpecification(BaseModel): dataset: str = Field( description="Hugging Face dataset ID, or path to dataset on disk." @@ -412,6 +417,11 @@ class Settings(BaseSettings): description="Maximum size for individual safetensors files generated when exporting a model.", ) + export_strategy: ExportStrategy | None = Field( + default=None, + description='How to export the model: "merge", "adapter", or unset to prompt the user.', + ) + refusal_markers: list[str] = Field( default=[ "sorry", diff --git a/src/heretic/main.py b/src/heretic/main.py index 0f0d98e..c359c2c 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -62,7 +62,7 @@ from rich.table import Table from rich.traceback import install from .analyzer import Analyzer -from .config import QuantizationMethod +from .config import ExportStrategy, QuantizationMethod from .evaluator import Evaluator from .model import AbliterationParameters, Model, get_model_class from .reproduce import ( @@ -88,13 +88,19 @@ from .utils import ( ) -def obtain_merge_strategy(settings: Settings, model: Model) -> str | None: +def obtain_export_strategy( + settings: Settings, + model: Model, +) -> ExportStrategy | None: """ - Prompts the user for how to proceed with saving the model. + Gets the export strategy from settings or prompts the user. Provides info to the user if the model is quantized on memory use. - Returns "merge", "adapter", or None (if cancelled/invalid). + Returns an export strategy, or None if cancelled. """ + if settings.export_strategy is not None: + return settings.export_strategy + if settings.quantization == QuantizationMethod.BNB_4BIT: print() print( @@ -148,11 +154,11 @@ def obtain_merge_strategy(settings: Settings, model: Model) -> str | None: if settings.quantization == QuantizationMethod.NONE else " (requires sufficient RAM)" ), - value="merge", + value=ExportStrategy.MERGE, ), Choice( title="Save LoRA adapter only (can be merged later)", - value="adapter", + value=ExportStrategy.ADAPTER, ), ], ) @@ -224,7 +230,7 @@ def run(): # FIXME: "Reproduction"/"reproducibility" name inconsistency! reproduction_information = load_reproduction_information(settings.reproduce) - if reproduction_information["version"] not in ["1"]: + if reproduction_information["version"] not in ["1", "2"]: print( ( f"[red]Unsupported file format version: [bold]{reproduction_information['version']}[/].[/] " @@ -865,11 +871,11 @@ def run(): if not save_directory: continue - strategy = obtain_merge_strategy(settings, model) + strategy = obtain_export_strategy(settings, model) if strategy is None: continue - if strategy == "adapter": + if strategy == ExportStrategy.ADAPTER: print("Saving LoRA adapter...") model.model.save_pretrained( save_directory, @@ -923,7 +929,7 @@ def run(): continue private = visibility == "Private" - strategy = obtain_merge_strategy(settings, model) + strategy = obtain_export_strategy(settings, model) if strategy is None: continue @@ -973,7 +979,7 @@ def run(): else: reproducibility_information = "none" - if strategy == "adapter": + if strategy == ExportStrategy.ADAPTER: print("Uploading LoRA adapter...") model.model.push_to_hub( repo_id, @@ -1036,17 +1042,22 @@ def run(): # Set the number of trials to the number of actual completed trials # for the reproduction configuration. settings.n_trials = count_completed_trials() + current_export_strategy = settings.export_strategy + settings.export_strategy = strategy - upload_reproduce_folder( - repo_id, - settings, - token, - checkpoint_path=study_checkpoint_file, - trial=trial, - include_system_information=( - reproducibility_information == "full" - ), - ) + try: + upload_reproduce_folder( + repo_id, + settings, + token, + checkpoint_path=study_checkpoint_file, + trial=trial, + include_system_information=( + reproducibility_information == "full" + ), + ) + finally: + settings.export_strategy = current_export_strategy print(f"Model uploaded to [bold]{repo_id}[/].") diff --git a/src/heretic/model.py b/src/heretic/model.py index 06e1711..5ff9fb7 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -539,6 +539,10 @@ class Model: W = W - W_org # Use a low-rank SVD to get an approximation of the matrix. r = self.peft_config.r + # svd_lowrank is randomized: + # https://github.com/pytorch/pytorch/blob/20919052303c0b5ba87f8bf7e19237dc33ab09d3/torch/_lowrank.py#L108-L109 + # Reseed immediately before the call so restoring a trial is independent of RNG history. + torch.manual_seed(self.settings.seed) U, S, Vh = torch.svd_lowrank(W, q=2 * r + 4, niter=6) # Truncate it to the part we want to store in the LoRA adapter. # Note: svd_lowrank actually returns V, so transpose it to get Vh. diff --git a/src/heretic/reproduce.py b/src/heretic/reproduce.py index 7717dee..6f82829 100644 --- a/src/heretic/reproduce.py +++ b/src/heretic/reproduce.py @@ -14,7 +14,11 @@ from urllib.request import urlopen import cpuinfo import torch from huggingface_hub import HfApi, hf_hub_download -from huggingface_hub.utils import disable_progress_bars, enable_progress_bars +from huggingface_hub.utils import ( + GatedRepoError, + disable_progress_bars, + enable_progress_bars, +) from questionary import Choice from rich.table import Table @@ -37,6 +41,7 @@ def collect_reproducibles(path: str): models = api.list_models( filter=["heretic", "reproducible"], sort="created_at", + expand=["gated", "tags"], ) found = 0 @@ -51,6 +56,12 @@ def collect_reproducibles(path: str): if model.tags is not None and "gguf" in model.tags: continue + if model.gated: + try: + api.auth_check(model.id, repo_type="model") + except GatedRepoError: + continue + print(f"[bold]{model.id}[/]...", end="") user, repository = model.id.split("/") diff --git a/src/heretic/utils.py b/src/heretic/utils.py index fdd5cf1..309ff86 100644 --- a/src/heretic/utils.py +++ b/src/heretic/utils.py @@ -547,7 +547,7 @@ def generate_reproduce_json( version_info = get_heretic_version_info() data = { - "version": "1", # Version number of the reproduce.json file format, to allow for future changes. + "version": "2", # Version number of the reproduce.json file format, to allow for future changes. "timestamp": timestamp, "system": None, # Defined here to preserve insertion order. "environment": {