diff --git a/antipasto/peft_utils/antipasto_adapter.py b/antipasto/peft_utils/antipasto_adapter.py index ba94ca8..2af41c4 100644 --- a/antipasto/peft_utils/antipasto_adapter.py +++ b/antipasto/peft_utils/antipasto_adapter.py @@ -58,6 +58,11 @@ class AntiPaSTOConfig(PeftConfig): repr=False, # Don't print in __repr__ metadata={"help": "Dict of {layer_name: indices_tensor} for data-aware dim selection."} ) + svd_bases: Optional[Dict[str, torch.Tensor]] = field( + default=None, + repr=False, # Don't print in __repr__ + metadata={"help": "Dict of {layer_name.U/V/S: tensor} for loading saved SVD bases."} + ) def __post_init__(self): self.peft_type = 'APASTOADAPTER' @@ -69,6 +74,7 @@ class AntiPaSTOConfig(PeftConfig): d = super().to_dict() # Remove precomputed_indices from serialization (only for init) d.pop('precomputed_indices', None) + d.pop('svd_bases', None) return d rotate_u: bool = field( default=False, @@ -169,13 +175,15 @@ class AntiPaSTOLayer(BaseTunerLayer): 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, **kwargs ) -> None: """ Initialize adapter with simple top-r SVD + residual (PiSSA-style). - If precomputed_indices provided, uses those dim indices. + If svd_bases provided, uses those U/V/S directly (for exact reload). + Elif precomputed_indices provided, uses those dim indices. Otherwise falls back to naive top-r by singular value (default PiSSA). """ if adapter_name in self.antipasto_u: @@ -203,35 +211,61 @@ class AntiPaSTOLayer(BaseTunerLayer): base_weight = base_weight.float() # [out, in] device = base_weight.device - # Full SVD for component selection - U_full, S_full, Vh_full = torch.linalg.svd(base_weight, full_matrices=False) - max_rank = min(U_full.shape[1], S_full.shape[0]) # Can't exceed matrix dimensions - r_actual = min(r, max_rank) # Clamp r to available rank + # Check for saved SVD bases (exact reload) + # Key format in saved file: svd_bases.{module_path_with_underscores}.u + # layer_name format: model.layers.1.self_attn.o_proj + # We need to find matching key by converting layer_name to underscored format + svd_key_base = None + if svd_bases is not None: + layer_key = layer_name.replace('.', '_') + # Try to find matching key (handle different prefix possibilities) + for candidate_prefix in ['base_model_model_', 'base_model_', '']: + candidate_key = f"svd_bases.{candidate_prefix}{layer_key}.u" + if candidate_key in svd_bases: + svd_key_base = f"svd_bases.{candidate_prefix}{layer_key}" + break - # Dimension selection: precomputed_indices (data-aware) or top-r (default PiSSA) - if precomputed_indices is not None and layer_name in precomputed_indices: - indices = precomputed_indices[layer_name].to(device) - r_actual = min(len(indices), r_actual) - indices = indices[:r_actual] + if svd_key_base is not None: + U = svd_bases[f"{svd_key_base}.u"].to(device) + V = svd_bases[f"{svd_key_base}.v"].to(device) + S = svd_bases[f"{svd_key_base}.s"].to(device) + r_actual = S.shape[0] - U = U_full[:, indices] # [d_out, r_actual] - Vh = Vh_full[indices, :] # [r_actual, d_in] - V = Vh.T # [d_in, r_actual] - S = S_full[indices] + # Compute residual from saved bases + Vh = V.T + W_principal = U @ torch.diag(S) @ Vh + W_res = base_weight - W_principal - logger.debug(f"Precomputed indices init: layer={layer_name}, {len(indices)} dims") + logger.debug(f"Loaded SVD bases: layer={layer_name}, r={r_actual}") else: - # Naive top-r by singular values (original PiSSA) - U = U_full[:, :r_actual] # [d_out, r_actual] - S = S_full[:r_actual] # [r_actual] - Vh = Vh_full[:r_actual, :] # [r_actual, d_in] - V = Vh.T # [d_in, r_actual] - - # Compute residual (PiSSA-style) - W_principal = U @ torch.diag(S) @ Vh - W_res = base_weight - W_principal - # Consider in PiSSA is calculated as - # W_res = U[:, r:] @ torch.diag(S_full[r:]) @ Vh[r:, :] + # Compute SVD from base weight + U_full, S_full, Vh_full = torch.linalg.svd(base_weight, full_matrices=False) + max_rank = min(U_full.shape[1], S_full.shape[0]) + r_actual = min(r, max_rank) + + # Dimension selection: precomputed_indices (data-aware) or top-r (default PiSSA) + if precomputed_indices is not None and layer_name in precomputed_indices: + indices = precomputed_indices[layer_name].to(device) + r_actual = min(len(indices), r_actual) + indices = indices[:r_actual] + + U = U_full[:, indices] + Vh = Vh_full[indices, :] + V = Vh.T + S = S_full[indices] + + logger.debug(f"Precomputed indices init: layer={layer_name}, {len(indices)} dims") + else: + # Naive top-r by singular values (original PiSSA) + U = U_full[:, :r_actual] + S = S_full[:r_actual] + Vh = Vh_full[:r_actual, :] + V = Vh.T + + # Compute residual (PiSSA-style) + W_principal = U @ torch.diag(S) @ Vh + W_res = base_weight - W_principal + logger.debug(f"AntiPaSTO Layer Init: {layer_name}, r={r_actual}, norms W={base_weight.norm():.1f}, Wres={W_res.norm():.1f}, Wrank={W_principal.norm():.1f}") # Store frozen components @@ -541,6 +575,7 @@ class AntiPaSTOModel(BaseTuner): "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, diff --git a/antipasto/peft_utils/load.py b/antipasto/peft_utils/load.py index e575bbc..a136ce0 100644 --- a/antipasto/peft_utils/load.py +++ b/antipasto/peft_utils/load.py @@ -36,6 +36,7 @@ def save_adapter( 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. @@ -46,6 +47,7 @@ def save_adapter( 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 @@ -58,6 +60,21 @@ def save_adapter( to_return = {k: state_dict[k] for k in state_dict if prefix in k} to_return = {remove_adapter_name(k, adapter_name): v for k, v in to_return.items()} + + # Optionally include SVD bases (U, V, S) for exact reload + # These are BufferDict entries, not parameters, so we need to extract them separately + if save_svd_bases: + svd_buffers = {} + for name, module in model.named_modules(): + if hasattr(module, 'antipasto_u') and adapter_name in module.antipasto_u: + # Extract SVD bases for this adapter + layer_key = name.replace('.', '_') # Safe key for safetensors + svd_buffers[f"svd_bases.{layer_key}.u"] = module.antipasto_u[adapter_name].clone() + svd_buffers[f"svd_bases.{layer_key}.v"] = module.antipasto_v[adapter_name].clone() + svd_buffers[f"svd_bases.{layer_key}.s"] = module.antipasto_s[adapter_name].clone() + if svd_buffers: + safetensors.torch.save_file(svd_buffers, save_folder / "0_svd_bases.safetensors") + logger.info(f"Saved SVD bases for {len(svd_buffers) // 3} layers") safetensors.torch.save_file(to_return, save_folder / "adapter_model.safetensors") config.save_pretrained(save_folder) @@ -147,6 +164,15 @@ def load_adapter( 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(): @@ -167,6 +193,7 @@ def load_adapter( config, target_modules=target_modules, precomputed_indices=precomputed_indices, + svd_bases=svd_bases, ) # Load weights @@ -180,4 +207,4 @@ def load_adapter( logger.info(f"Loaded adapter from {adapter_folder}") - return model, None, layer_selection + return model, tokenizer, layer_selection diff --git a/antipasto/train/model_setup.py b/antipasto/train/model_setup.py index e517d19..917b8ad 100644 --- a/antipasto/train/model_setup.py +++ b/antipasto/train/model_setup.py @@ -79,7 +79,7 @@ def load_model(model_id, quantization_type="none"): return base_model, tokenizer -def setup_adapter(base_model, config: TrainingConfig, target_modules: str, precomputed_indices=None): +def setup_adapter(base_model, config: TrainingConfig, target_modules: str, precomputed_indices=None, svd_bases=None): """Setup AntiPaSTO adapter on base model. Args: @@ -87,6 +87,7 @@ def setup_adapter(base_model, config: TrainingConfig, target_modules: str, preco config: Training configuration target_modules: PEFT target_modules regex (from LayerSelection) precomputed_indices: Optional dict of {layer_name: indices_tensor} for dim selection + svd_bases: Optional dict of {layer_name.U/V/S: tensor} for loading saved SVD bases """ logger.debug(f"Target modules regex: {target_modules}") @@ -99,14 +100,17 @@ def setup_adapter(base_model, config: TrainingConfig, target_modules: str, preco task_type="CAUSAL_LM", target_modules=target_modules, precomputed_indices=precomputed_indices, + svd_bases=svd_bases, ) # Create PeftModel - AntiPaSTO handles bidirectional steering internally via alpha coefficient model = PeftModel(base_model, adapter_config, adapter_name=config.dataset_name) - # Clear precomputed_indices from config after adapter creation (only needed for init) + # Clear precomputed_indices and svd_bases from config after adapter creation (only needed for init) if adapter_config.precomputed_indices is not None: adapter_config.precomputed_indices = None + if adapter_config.svd_bases is not None: + adapter_config.svd_bases = None logger.info( f"Adapter configured: rank={config.r}, target_modules={target_modules}"