From 60fe4007611c4620ef2f64e960b0e536b5ead050 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Tue, 13 Jan 2026 08:55:00 +0800 Subject: [PATCH] Refactor AntiPaSTO configuration and loading logic: remove svd_aligned_init, streamline adapter loading, and enhance model setup for clarity and maintainability. --- antipasto/peft_utils/antipasto_adapter.py | 26 +---- antipasto/peft_utils/load.py | 127 ++++++++++------------ antipasto/train/model_setup.py | 1 - antipasto/train/train_adapter.py | 29 +---- 4 files changed, 66 insertions(+), 117 deletions(-) diff --git a/antipasto/peft_utils/antipasto_adapter.py b/antipasto/peft_utils/antipasto_adapter.py index 357b2fc..d61e04a 100644 --- a/antipasto/peft_utils/antipasto_adapter.py +++ b/antipasto/peft_utils/antipasto_adapter.py @@ -90,10 +90,6 @@ class AntiPaSTOConfig(PeftConfig): default="cayley", metadata={"help": "Rotation parameterization: 'cayley' (recommended, exact reversibility) or 'matrix_exp' (exact but slower)"} ) - svd_aligned_init: bool = field( - default=False, - metadata={"help": "Initialize delta_s proportional to S (normalized). Gives very stable init (std=0.26 across seeds)."} - ) alpha: float = field( default=1.0, metadata={"help": "Steering coefficient for rotations (1.0 = forward, -1.0 = reverse, 0.0 = disabled)"} @@ -102,10 +98,6 @@ class AntiPaSTOConfig(PeftConfig): default=torch.pi/3, metadata={"help": "Max rotation angle (radians, soft-clamped). Small angles (≤0.3) ensure R(α)@S ≈ -R(-α)@S for output symmetry at α=±1. Set to inf to disable."} ) - # steer_s: bool = field( - # default=False, - # metadata={"help": "Whether to apply steering to singular value scaling"} - # ) # Standard PEFT parameters target_modules: Optional[list[str]] = field( @@ -142,7 +134,6 @@ class AntiPaSTOLayer(BaseTunerLayer): self.antipasto_rotation_method = {} self.antipasto_alpha = {} self.antipasto_max_rotation_angle = {} - self.antipasto_svd_aligned_init = {} # SVD components (per adapter) - simplified naming like SVDSteering self.antipasto_u = BufferDict({}) # U: [d_out, r] @@ -175,7 +166,6 @@ class AntiPaSTOLayer(BaseTunerLayer): rotate_v, rotation_method, max_rotation_angle, - svd_aligned_init: bool = False, precomputed_indices: Optional[Dict[str, torch.Tensor]] = None, svd_bases: Optional[Dict[str, torch.Tensor]] = None, layer_name: Optional[str] = None, @@ -199,7 +189,6 @@ class AntiPaSTOLayer(BaseTunerLayer): self.antipasto_rotate_v[adapter_name] = rotate_v self.antipasto_rotation_method[adapter_name] = rotation_method self.antipasto_max_rotation_angle[adapter_name] = max_rotation_angle - self.antipasto_svd_aligned_init[adapter_name] = svd_aligned_init # Get base weight base_weight = self.get_base_layer().weight @@ -281,13 +270,8 @@ class AntiPaSTOLayer(BaseTunerLayer): torch.zeros(r_actual, device=device), requires_grad=True ) - if self.antipasto_svd_aligned_init.get(adapter_name, False): - # SVD-aligned init: delta_s ∝ S (normalized). Very stable across seeds (std=0.26). - s_normalized = S / S.max() - self.antipasto_delta_s[adapter_name].data = s_normalized * 4e-4 + 4e-4 - else: - # Default: small random noise - nn.init.trunc_normal_(self.antipasto_delta_s[adapter_name], std=4e-4, mean=4e-4) + # Small random noise init + nn.init.trunc_normal_(self.antipasto_delta_s[adapter_name], std=4e-4, mean=4e-4) @@ -572,15 +556,11 @@ class AntiPaSTOModel(BaseTuner): "rotate_u": antipasto_config.rotate_u, "rotate_v": antipasto_config.rotate_v, "rotation_method": antipasto_config.rotation_method, - # "block_size": antipasto_config.block_size, "alpha": antipasto_config.alpha, "max_rotation_angle": antipasto_config.max_rotation_angle, - "svd_aligned_init": antipasto_config.svd_aligned_init, "precomputed_indices": antipasto_config.precomputed_indices, "svd_bases": antipasto_config.svd_bases, - "layer_name": current_key, # Pass layer name for dim index lookup - # "data_aware_init_use_magnitudes": antipasto_config.data_aware_init_use_magnitudes, - # "steer_s": antipasto_config.steer_s, + "layer_name": current_key, **optional_kwargs, } diff --git a/antipasto/peft_utils/load.py b/antipasto/peft_utils/load.py index ca6b9fc..8e7a6aa 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, Union +from typing import Optional, Tuple from antipasto.peft_utils.layer_selection import LayerSelection @@ -36,7 +36,6 @@ def save_adapter( model_id: str = None, layer_selection: Optional[LayerSelection] = None, precomputed_indices: Optional[dict] = None, - bake_centering: bool = True, save_svd_bases: bool = True, ): """Save adapter weights, config, and metadata needed for reloading. @@ -48,7 +47,6 @@ def save_adapter( 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 save_svd_bases: If True, save U, V, S buffers for exact reload (avoids SVD recomputation issues) """ from peft.mapping import PEFT_TYPE_TO_PREFIX_MAPPING @@ -108,110 +106,105 @@ def load_adapter( """Load a saved AntiPaSTO adapter with all metadata. Either provide base_model directly, OR model_id + quantization_type to load it. + Model ID can also be read from adapter_config.json (base_model_name_or_path). Args: adapter_folder: Path to saved adapter (contains adapter_model.safetensors, etc.) base_model: Pre-loaded base model (optional, provide this OR model_id) - model_id: HuggingFace model ID to load (optional, provide this OR base_model) + model_id: HuggingFace model ID to load (optional, auto-detected from adapter_config.json) quantization_type: Quantization type for loading model (e.g., "nf4", "int8", None) adapter_name: Name to assign to the loaded adapter Returns: - Tuple of (PeftModel with loaded adapter, LayerSelection if saved else None) + Tuple of (PeftModel with loaded adapter, tokenizer, LayerSelection if saved else None) Example: - model, layer_selection = load_adapter( + # Auto-detect model from adapter_config.json: + model, tokenizer, layer_selection = load_adapter(Path("outputs/adapters/my_run")) + + # Or specify model explicitly: + model, tokenizer, layer_selection = load_adapter( Path("outputs/adapters/my_run"), - model_id="Qwen/Qwen2.5-3B-Instruct", + model_id="google/gemma-3-270m-it", ) + # For inference: with ScaleAdapter(model, coeff=1.0): output = model.generate(...) """ - from antipasto.peft_utils.antipasto_adapter import register_antipasto_peft - from antipasto.train.model_setup import load_model, setup_adapter + from antipasto.peft_utils.antipasto_adapter import register_antipasto_peft, AntiPaSTOConfig + from antipasto.train.model_setup import load_model adapter_folder = Path(adapter_folder) # Register AntiPaSTO adapter type register_antipasto_peft() + # Load adapter_config.json (standard PEFT config) + adapter_config_path = adapter_folder / "adapter_config.json" + if not adapter_config_path.exists(): + raise ValueError(f"adapter_config.json not found in {adapter_folder}") + + with open(adapter_config_path) as f: + adapter_config_dict = json.load(f) + + # Determine model_id from config if not provided + if base_model is None and model_id is None: + model_id = adapter_config_dict.get("base_model_name_or_path") + if model_id is None: + raise ValueError( + "Must provide base_model or model_id, or have base_model_name_or_path in adapter_config.json" + ) + # Load base model if not provided if base_model is None: - if model_id is None: - # Try to get model_id from training_config.json - config_path = adapter_folder / "training_config.json" - if config_path.exists(): - with open(config_path) as f: - training_config = json.load(f) - model_id = training_config.get("model_name") - quantization_type = quantization_type or training_config.get("quantization_type") - else: - raise ValueError("Must provide base_model or model_id, or have training_config.json in adapter_folder") - base_model, tokenizer = load_model(model_id, quantization_type=quantization_type) else: tokenizer = None - # Load layer_selection if saved - layer_selection = None - layer_selection_path = adapter_folder / "0_layer_selection.json" - if layer_selection_path.exists(): - with open(layer_selection_path) as f: - layer_selection = LayerSelection.from_dict(json.load(f)) - target_modules = layer_selection.adapter_regex - else: - raise ValueError(f"Missing 0_layer_selection.json in {adapter_folder}") + # Load SVD bases (required for correct reload) + svd_bases_path = adapter_folder / "0_svd_bases.safetensors" + if not svd_bases_path.exists(): + raise ValueError(f"0_svd_bases.safetensors not found in {adapter_folder}") + svd_bases = safetensors.torch.load_file(svd_bases_path) + logger.info(f"Loaded SVD bases for {len(svd_bases) // 3} layers") - # Load precomputed_indices if saved (for dimension selection) - precomputed_indices = None - indices_path = adapter_folder / "0_precomputed_indices.pt" - if indices_path.exists(): - precomputed_indices = torch.load(indices_path, weights_only=True) - else: - logger.warning(f"No precomputed indices found in {adapter_folder}, proceeding without dimension selection.") - - # Load SVD bases if saved (for exact reload without recomputation) - svd_bases = None - svd_bases_st = adapter_folder / "0_svd_bases.safetensors" - if svd_bases_st.exists(): - svd_bases = safetensors.torch.load_file(svd_bases_st) - logger.info(f"Loaded SVD bases for {len(svd_bases) // 3} layers") - else: - logger.warning(f"No SVD bases found in {adapter_folder}, will recompute from base model.") - - # Load training config to get adapter settings - config_path = adapter_folder / "training_config.json" - if config_path.exists(): - with open(config_path) as f: - training_config = json.load(f) - - # Create minimal config for setup_adapter - from antipasto.config import TrainingConfig - import cattrs - config = cattrs.structure(training_config, TrainingConfig) - config.dataset_name = adapter_name # Use provided adapter name - else: - raise ValueError(f"training_config.json not found in {adapter_folder}") - - # Setup adapter structure - model = setup_adapter( - base_model, - config, - target_modules=target_modules, - precomputed_indices=precomputed_indices, + # Build AntiPaSTOConfig from adapter_config.json - fail fast on missing required fields + adapter_config = AntiPaSTOConfig( + r=adapter_config_dict["r"], + rotate_u=adapter_config_dict["rotate_u"], + rotate_v=adapter_config_dict["rotate_v"], + rotation_method=adapter_config_dict["rotation_method"], + max_rotation_angle=adapter_config_dict["max_rotation_angle"], + alpha=adapter_config_dict["alpha"], + task_type=adapter_config_dict["task_type"], + target_modules=adapter_config_dict["target_modules"], svd_bases=svd_bases, ) + # Create PeftModel with adapter + model = PeftModel(base_model, adapter_config, adapter_name=adapter_name) + + # Clear svd_bases from config after creation + adapter_config.svd_bases = None + + logger.info(f"Adapter configured: rank={adapter_config.r}, layers={len([n for n, m in model.named_modules() if hasattr(m, 'antipasto_u')])}") + # Load weights sd = safetensors.torch.load_file(adapter_folder / "adapter_model.safetensors") sd = add_adapter_name_to_sd(sd, adapter_name=adapter_name, prefix="antipasto_") - # FIXME do we use this with lora,dora,road,vera,ia3 too? result = model.load_state_dict(sd, strict=False) if result.unexpected_keys: raise ValueError(f"Unexpected keys in state_dict: {result.unexpected_keys[:5]}") + # Load layer_selection if saved (optional, for evaluation) + layer_selection = None + layer_selection_path = adapter_folder / "0_layer_selection.json" + if layer_selection_path.exists(): + with open(layer_selection_path) as f: + layer_selection = LayerSelection.from_dict(json.load(f)) + logger.info(f"Loaded adapter from {adapter_folder}") return model, tokenizer, layer_selection diff --git a/antipasto/train/model_setup.py b/antipasto/train/model_setup.py index 917b8ad..fe14986 100644 --- a/antipasto/train/model_setup.py +++ b/antipasto/train/model_setup.py @@ -96,7 +96,6 @@ def setup_adapter(base_model, config: TrainingConfig, target_modules: str, preco rotate_u=config.rot_u, rotate_v=config.rot_v, max_rotation_angle=config.max_rotation_angle, - svd_aligned_init=config.svd_aligned_init, task_type="CAUSAL_LM", target_modules=target_modules, precomputed_indices=precomputed_indices, diff --git a/antipasto/train/train_adapter.py b/antipasto/train/train_adapter.py index f49df6e..31bef5f 100644 --- a/antipasto/train/train_adapter.py +++ b/antipasto/train/train_adapter.py @@ -28,8 +28,6 @@ from baukit.nethook import TraceDict from loguru import logger from tabulate import tabulate from torch.utils.data import DataLoader -from torchjd import autojac -from torchjd.aggregation import UPGrad from tqdm.auto import tqdm from transformers import DataCollatorWithPadding @@ -955,11 +953,7 @@ def train_epoch( # === LoRA Trust Region: SOFT constraint (loss term) === # Add norm penalty to loss before backward for LoRA/DoRA adapters. - if config.upgrad: - # UPGrad: balance gradients from per-layer projection, coherence, and monotonic losses - autojac.backward(loss_components, aggregator, parallel_chunk_size=1) - else: - total_loss.mean().backward() + total_loss.mean().backward() # Logging log_n_steps = max(1, len(train_dataloader) * config.n_epochs // config.n_logs) @@ -1831,24 +1825,6 @@ def train_model(config: TrainingConfig): ) total_steps = config.n_epochs * len(train_dataloader) // config.grad_accum_steps - if config.upgrad: - # Build pref_vector: balance projection (per layer, per coef) vs coherence vs monotonic - # Structure: [proj_L0_pos, proj_L0_neg, proj_L1_pos, proj_L1_neg, ..., coh_pos, coh_neg, mono] - n_loss_layers = len(loss_layers) - pref_vec = [] - for _ in range(n_loss_layers): - pref_vec.append(10*config.upgrad_balance) # proj coef=+1 - pref_vec.append(10*1.0 / config.upgrad_balance) # proj coef=-1 (inverse balance) - if config.coh: - pref_vec.append(0.5) # coh coef=+1 - pref_vec.append(0.5) # coh coef=-1 - if config.mono: - pref_vec.append(1.0) # monotonic ordering - aggregator = UPGrad( - pref_vector=torch.tensor(pref_vec, device=model.device), - ) - else: - aggregator = None opt = torch.optim.AdamW( model.parameters(), lr=config.lr, weight_decay=config.wd ) @@ -1985,11 +1961,12 @@ def train_model(config: TrainingConfig): model, save_folder, config.dataset_name, + model_id=config.model_name, layer_selection=layer_selection, precomputed_indices=precomputed_indices_for_save, ) - # Save training config + # Save training config (for full reproducibility, optional for loading) with open(save_folder / "training_config.json", "w") as f: json.dump(cattrs.unstructure(config), f, indent=4)