File size: 2,793 Bytes
49dddf5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | # Copyright 2026 Angshul Majumdar. Licensed under the Apache License, Version 2.0.
"""TensorPEFTConfig."""
import json
import os
from dataclasses import asdict, dataclass, field
from typing import List, Optional, Union
CONFIG_NAME = "adapter_config.json"
WEIGHTS_NAME = "adapter_model.safetensors"
@dataclass
class TensorPEFTConfig:
"""Configuration of a TensorPEFT adapter.
Args:
atoms_per_group: rank-one CP atoms per tensor group (all chosen at initialisation).
Trainable parameters = sum over groups of atoms_per_group * (1 + sum of the group's mode sizes).
target_modules: ``"all-linear"`` (every Linear / Conv1D inside the transformer blocks) or a
list of module-name suffixes, e.g. ``["q_proj", "v_proj"]`` (``"c_attn.q"`` etc. for GPT-2).
exclude_modules: substrings of module names that are never adapted.
split_fused_qkv: split GPT-2-style fused ``c_attn`` = [q|k|v] into q, k, v.
hopm_iters: higher-order power iterations used to compute each atom.
alpha_scale: coefficients are parametrised as ``alpha_scale * alpha_raw`` so that a single
optimiser learning rate (about 3e-3 for Adam/AdamW) suits both coefficients and factors.
init_batches: default number of batches used by ``init_atoms``.
"""
atoms_per_group: int = 4
target_modules: Union[str, List[str]] = "all-linear"
exclude_modules: List[str] = field(default_factory=lambda: ["lm_head", "embed", "wte", "wpe"])
split_fused_qkv: bool = True
hopm_iters: int = 30
alpha_scale: float = 100.0 / 3.0
init_batches: int = 4
base_model_name_or_path: Optional[str] = None
peft_type: str = "TENSORPEFT"
task_type: str = "CAUSAL_LM"
tensorpeft_version: Optional[str] = None
def to_dict(self):
return asdict(self)
def save_pretrained(self, save_directory, extra=None):
os.makedirs(save_directory, exist_ok=True)
d = self.to_dict()
if extra:
d.update(extra)
with open(os.path.join(save_directory, CONFIG_NAME), "w") as f:
json.dump(d, f, indent=2)
@classmethod
def from_dict(cls, d):
names = {f for f in cls.__dataclass_fields__}
return cls(**{k: v for k, v in d.items() if k in names})
@classmethod
def from_pretrained(cls, model_id, **hub_kwargs):
path = _resolve(model_id, CONFIG_NAME, **hub_kwargs)
with open(path) as f:
d = json.load(f)
return cls.from_dict(d), d
def _resolve(model_id, filename, **hub_kwargs):
local = os.path.join(model_id, filename)
if os.path.isfile(local):
return local
from huggingface_hub import hf_hub_download
return hf_hub_download(repo_id=model_id, filename=filename, **hub_kwargs)
|