wassname/AntiPaSTO
1
0
Fork 0
mirror of https://github.com/wassname/AntiPaSTO.git synced 2026-09-11 11:52:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add support for saving and loading SVD bases in adapter configuration

Browse Source
This commit is contained in:
wassname
2026-01-13 07:43:12 +08:00
parent 42b38a0626
commit bc69fc39a3
3 changed files with 95 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
antipasto/peft_utils/antipasto_adapter.py
+61 -26
View File
@@ -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,
antipasto/peft_utils/load.py
+28 -1
View File
@@ -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
antipasto/train/model_setup.py
+6 -2
View File
@@ -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}"