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

@@ -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,
         }
 

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

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,

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)