写一个新的 PEFT 方法

PEFT 的代码骨架做得相对清晰，写一个全新的 tuner 只需要四个文件 ≈ 200 行。 本章带你实现一个示例方法 ScaleLoRA（在 LoRA 的 BA 外加上可学的逐位置 scaling，类似 DoRA 的简化版）， 顺道把"PEFT 扩展点"全摸一遍。

12.1需要写哪些文件

12.2定义算法

ScaleLoRA 的公式：

$$y = W_0 x + s \odot \frac{\alpha}{r}\, B A x,\quad s \in \mathbb{R}^{d_\text{out}}$$

跟 LoRA 比多了 $s$（每个 output 维度一个标量），训练时一起学。$s$ 初始为 1。

12.3config.py

# src/peft/tuners/scalelora/config.py
from dataclasses import dataclass, field
from typing import List, Optional, Union
from peft.config import PeftConfig
from peft.utils import PeftType

@dataclass
class ScaleLoraConfig(PeftConfig):
    r: int = field(default=8)
    target_modules: Optional[Union[List[str], str]] = field(default=None)
    lora_alpha: int = field(default=16)
    lora_dropout: float = field(default=0.0)
    bias: str = field(default="none")
    modules_to_save: Optional[List[str]] = field(default=None)

    def __post_init__(self):
        self.peft_type = PeftType.SCALELORA      # ★ 新加的类型

记得在 peft/utils/peft_types.py 里加一行：

class PeftType(str, enum.Enum):
    ...
    SCALELORA = "SCALELORA"        # ★ 新增

12.4layer.py

# src/peft/tuners/scalelora/layer.py
import math
import torch
from torch import nn
from peft.tuners.tuners_utils import BaseTunerLayer

class ScaleLoraLinear(nn.Module, BaseTunerLayer):
    adapter_layer_names = ("lora_A", "lora_B", "lora_s", "lora_dropout")

    def __init__(self, base_layer: nn.Linear, adapter_name: str,
                 r: int, lora_alpha: int, lora_dropout: float = 0.0):
        nn.Module.__init__(self)
        BaseTunerLayer.__init__(self)
        self.base_layer = base_layer
        self.in_features = base_layer.in_features
        self.out_features = base_layer.out_features
        # 用 ModuleDict 支持多 adapter
        self.lora_A = nn.ModuleDict({})
        self.lora_B = nn.ModuleDict({})
        self.lora_s = nn.ParameterDict({})
        self.lora_dropout = nn.ModuleDict({})
        self.scaling = {}
        self._active_adapter = adapter_name
        self.update_layer(adapter_name, r, lora_alpha, lora_dropout)

    def update_layer(self, adapter_name, r, lora_alpha, lora_dropout):
        self.lora_A[adapter_name] = nn.Linear(self.in_features, r, bias=False)
        self.lora_B[adapter_name] = nn.Linear(r, self.out_features, bias=False)
        self.lora_s[adapter_name] = nn.Parameter(torch.ones(self.out_features))
        self.lora_dropout[adapter_name] = (
            nn.Dropout(p=lora_dropout) if lora_dropout > 0 else nn.Identity()
        )
        self.scaling[adapter_name] = lora_alpha / r
        nn.init.kaiming_uniform_(self.lora_A[adapter_name].weight, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B[adapter_name].weight)

    def forward(self, x: torch.Tensor):
        result = self.base_layer(x)
        if self.disable_adapters:
            return result
        for adapter in self.active_adapters:
            if adapter not in self.lora_A:
                continue
            A = self.lora_A[adapter]
            B = self.lora_B[adapter]
            s = self.lora_s[adapter]
            scaling = self.scaling[adapter]
            dropout = self.lora_dropout[adapter]
            update = B(A(dropout(x))) * scaling          # 普通 LoRA 部分
            update = update * s                          # ★ 我们加的逐位 scaling
            result = result + update
        return result

    def merge(self, safe_merge: bool = False, adapter_names=None):
        for adapter in (adapter_names or self.active_adapters):
            if adapter in self.merged_adapters: continue
            A = self.lora_A[adapter].weight
            B = self.lora_B[adapter].weight
            s = self.lora_s[adapter]
            delta_W = (s.unsqueeze(1) * (B @ A)) * self.scaling[adapter]
            self.base_layer.weight.data += delta_W
            self.merged_adapters.append(adapter)

12.5model.py

# src/peft/tuners/scalelora/model.py
import torch
from torch import nn
from peft.tuners.tuners_utils import BaseTuner
from .layer import ScaleLoraLinear

class ScaleLoraModel(BaseTuner):
    prefix = "lora_"               # ★ 用来标记可训参数的前缀

    def _check_new_adapter_config(self, config):
        # 通常空着或做简单合法性检查
        pass

    def _check_target_module_exists(self, peft_config, key):
        # 用 endswith 匹配
        if isinstance(peft_config.target_modules, str):
            import re
            return re.fullmatch(peft_config.target_modules, key)
        return any(key.endswith(t) for t in peft_config.target_modules)

    def _create_and_replace(self, peft_config, adapter_name, target, target_name,
                            parent, current_key):
        if isinstance(target, nn.Linear):
            new_module = ScaleLoraLinear(
                base_layer=target, adapter_name=adapter_name,
                r=peft_config.r, lora_alpha=peft_config.lora_alpha,
                lora_dropout=peft_config.lora_dropout,
            )
            self._replace_module(parent, target_name, new_module, target)

    def _replace_module(self, parent, child_name, new_module, child):
        setattr(parent, child_name, new_module)
        # 保持 weight / bias 设备一致
        new_module.to(child.weight.device)

    def _mark_only_adapters_as_trainable(self, model):
        for n, p in model.named_parameters():
            if "lora_" not in n:
                p.requires_grad = False

    @staticmethod
    def _prepare_adapter_config(peft_config, model_config):
        return peft_config

12.6注册到 PEFT 主入口

在 src/peft/mapping_func.py（或新版 mapping.py）：

from peft.tuners.scalelora import ScaleLoraConfig, ScaleLoraModel

PEFT_TYPE_TO_CONFIG_MAPPING["SCALELORA"] = ScaleLoraConfig
PEFT_TYPE_TO_TUNER_MAPPING["SCALELORA"] = ScaleLoraModel

然后还要把它绑定到对应的 PeftModelForXXX 上（同文件附近）：

MODEL_TYPE_TO_PEFT_MODEL_MAPPING = {
    "CAUSAL_LM": {
        ...,
        "SCALELORA": PeftModelForCausalLM,
    },
    ...
}

12.7试用

from peft.tuners.scalelora import ScaleLoraConfig
from peft import get_peft_model

cfg = ScaleLoraConfig(
    task_type="CAUSAL_LM",
    r=8, lora_alpha=16,
    target_modules=["q_proj","v_proj"],
)
model = get_peft_model(base_model, cfg)
model.print_trainable_parameters()
# trainable: ~0.12% (LoRA 0.10% + s 矩阵约 0.02%)

跑训练、save/load、merge、multi-adapter 全都跟 LoRA 一样能用——因为我们继承了 BaseTuner / BaseTunerLayer 的所有公共行为。

12.8写完后还要做什么

1. 写测试：参考 tests/test_custom_models.py，把新 tuner 加进 PARAMETRIZED 测试矩阵；
2. 写 docstring：每个公开类的 docstring 必须用 sphinx 风格；
3. 跑 ruff + mypy：仓库根目录 make quality；
4. 更新 docs：docs/source/conceptual_guides/ 加方法说明；
5. 提 PR：在 GitHub issue 里先沟通；PR 标题用 FEAT: Add ScaleLoRA。

12.9真实新方法的常见难点

12.10另一条路：写 LoraVariant 而不是新 tuner

本章前面教的是"完全自定义 tuner"——加 config + layer + model 三件套。但如果你的新方法只是改变 LoRA 的 forward / merge 行为（如 DoRA、ALoRA），更优雅的做法是写一个 LoraVariant 子类（见 ch04.8）。

对比两条路：

"新方法仍然把 ΔW = BA 当核心吗？"——是 → 选 LoraVariant；不是 → 完全自定义 tuner。

12.11LoraVariant 模板

来源：src/peft/tuners/lora/layer.py:53–101 的抽象基类。写一个新 variant 大致这样：

from peft.tuners.lora.layer import LoraVariant
import torch.nn as nn
import torch

class MyVariant(LoraVariant):
    """演示：给 LoRA forward 额外加一个 learnable bias vector"""

    def init(self, layer, adapter_name):
        # 给每个 adapter 加一个 bias parameter
        out_features = layer.base_layer.out_features
        if not hasattr(layer, "my_bias"):
            layer.my_bias = nn.ParameterDict()
        layer.my_bias[adapter_name] = nn.Parameter(torch.zeros(out_features))

    def merge_safe(self, layer, adapter_name):
        # 把 BA + my_bias 写回 base_layer
        delta_w = layer.lora_B[adapter_name].weight @ layer.lora_A[adapter_name].weight
        layer.base_layer.weight.data += delta_w * layer.scaling[adapter_name]
        # bias 单独 merge
        if layer.base_layer.bias is None:
            layer.base_layer.bias = nn.Parameter(layer.my_bias[adapter_name].clone())
        else:
            layer.base_layer.bias.data += layer.my_bias[adapter_name]

    def merge_unsafe(self, layer, adapter_name):
        # 不做 NaN 校验的版本
        self.merge_safe(layer, adapter_name)

    def unmerge(self, layer, adapter_name):
        # merge_safe 的逆
        delta_w = layer.lora_B[adapter_name].weight @ layer.lora_A[adapter_name].weight
        layer.base_layer.weight.data -= delta_w * layer.scaling[adapter_name]
        if layer.base_layer.bias is not None:
            layer.base_layer.bias.data -= layer.my_bias[adapter_name]

    def forward(self, layer, x, scaling):
        # base forward 已经在外面跑过；这里返回 LoRA 增量
        lora_x = layer.lora_A[layer.active_adapters[0]](x)
        lora_x = layer.lora_B[layer.active_adapters[0]](lora_x) * scaling
        lora_x = lora_x + layer.my_bias[layer.active_adapters[0]]
        return lora_x

注册：在 LoraConfig 用 use_my_variant=True flag 触发（需要给 config 加字段 + 在 layer 的 resolve_lora_variant() 加 case）。

12.12提 PR 前 5 项 checklist

12.13这章你需要带走的

• 新 PEFT 方法 = config + layer + model 三件套，~ 200 行；
• layer 继承 BaseTunerLayer，必须实现 update_layer / forward / merge；
• model 继承 BaseTuner，必须实现 _create_and_replace / _check_target_module_exists / _mark_only_adapters_as_trainable；
• 新增 peft_type 后要注册到 mapping，否则 get_peft_model 找不到；
• 测试一定写 "merge ↔ unmerge 等价"，否则隐 bug；
• 量化兼容需要为每种 base layer 写一份子类（bnb / awq / gptq …）。