PEFTLoRA实现及核心源码解读
作者： AINLP 来源： AINLP
1. PEFT 介绍
HuggingFace的PEFT(Parameter-Efficient Fine-Tuning）中提供了模型微调加速的方法，参数高效微调（PEFT）方法能够使预先训练好的语言模型（PLMs）有效地适应各种下游应用，而不需要对模型的所有参数进行微调。
对大规模的PLM进行微调往往成本过高，在这方面，PEFT方法只对少数（额外的）模型参数进行微调，基本思想在于仅微调少量 (额外) 模型参数，同时冻结预训练 LLM 的大部分参数，从而大大降低了计算和存储成本，这也克服了灾难性遗忘的问题，这是在 LLM 的全参数微调期间观察到的一种现象，同时PEFT 方法也显示出在低数据状态下比微调更好，可以更好地泛化到域外场景，参数高效微调有如下好处
（1）由于在训练时只更新少量参数，可以大大减少GPU显存的使用量，让大语言模型可以在消费级GPU进行训练。
（2）大模型的参数存在冗余，通过参数有效微调可以达到甚至比全参数微调的效果还要好
（3）无需全量保存参数，减少下游子任务参数存储成本
**2.**LoRA 介绍
LoRA的原理比较简单，原始全量微调就是在原始模型参数上通过微调加入增量W=W0+ΔW，那我们可以通过冻结原始参数W0，并且把增量部分通过低秩分解方式进一步降低参数量级ΔW=AB，原始参数的维度是dd, 则低秩分解后的参数量级是2rd，因为这里的r«d，因此可以起到大幅降低微调参数量级的效果，如下图
总结Lora特点：
（1）给原模型增加旁路，通过低秩分解（先降维再升维）来模拟参数的更新量；
（2）训练时，原模型固定，只训练降维矩阵A和升维矩B；
（3）推理时，可将BA加到原参数上，不引入额外的推理延迟；
（4）初始化，A采用高斯分布初始化，B初始化为全0，保证训练开始时旁路为0矩阵（严格讲，对不同算子采用不同的A初始化方式，例如，Linear算子采用kaiming_uniform，对于Embedding算子采用normal高斯分布）；
（5）可插拔式的切换任务，当前任务W0+B1A1，将lora部分减掉，换成B2A2，即可实现任务切换。
3. PEFT 实现****LoRA
基于PEFT框架实现指定模型的LoRA封装非常方便，仅需实例化模型，设置LoRA的config文件，调用get_peft_model方法即可。基于Transformer结构，LoRA一般只对每层的Self-Attention的部分进行微调，即对Wq、Wk、Wv、Wo四个映射层参数进行微调。消融实验显示只微调Wq效果略差，微调Wq、Wv的效果和微调Wq、Wk、Wv、Wo的效果相似（如下所示），以下示例微调Wq、Wv。
3.1 PEFT核心类
截至发文，PEFT支持LoRA、Prefix Tuning、P-Tuning、Prompt Tuning、AdaLoRA、Adaption Prompt共六种对大模型高效调参的方法，分别对应框架六个核心方法类即LoraModel、PrefixEncoder、PromptEncoder、PromptEmbedding、AdaLoraModel、AdaptionPromptModel，每个核心方法对应配置文件类为LoraConfig、PrefixTuningConfig、PromptEncoderConfig、PromptTuningConfig、AdaLoraConfig、AdaptionPromptConfig，如下
PeftModel类支持上述六种方法配置文件，对PromptLearning类 (PrefixTuning、PromptEncoder、PromptTuning)和非PromptLearning类（LORA、ADALORA、ADAPTION_PROMPT）不同分支进行处理，得到对应的base_model，其__init__函数如下
def __init__(self, model: PreTrainedModel, peft_config: PeftConfig, adapter_name: str = "default"):
super().__init__()
self.base_model = model
self.config = self.base_model.config
self.modules_to_save = None
self.peft_config = {}
self.active_adapter = adapter_name
self.peft_type = peft_config.peft_type
self.base_model_torch_dtype = getattr(model, "dtype", None)
if not isinstance(peft_config, PromptLearningConfig):
self.peft_config[adapter_name] = peft_config
self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type](
self.base_model, self.peft_config, adapter_name
)
self.set_additional_trainable_modules(peft_config, adapter_name)
else:
self.add_adapter(adapter_name, peft_config)
if getattr(model, "is_gradient_checkpointing", True):
model = self._prepare_model_for_gradient_checkpointing(model)
根据不同高效调参方法适配PeftModel类得到base_model，不同下游子任务基于PeftModel类构造子任务类，PEFT框架当前支持五种下游子任务，即PeftModelForSequenceClassification、PeftModelForSeq2SeqLM、PeftModelForCausalLM、PeftModelForTokenClassification、PeftModelForQuestionAnswering，这些子任务类就是借助PEFT框架形成的最终的peft_model，用于后续训练与推理。
3.2 LoraModel类的实现
LoraModel类实现对模型Lora化的方法封装，以下显式调用LoraModel实现Lora
from transformers import AutoModelForSeq2SeqLM, LoraConfig
from peft import LoraModel, LoraConfig
# step1. Lora配置
config = LoraConfig(peft_type="LORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"], lora_dropout=0.01,...)
# step2. 预训练模型加载
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
# step3. 显示生成Lora模型
lora_model = LoraModel(config, model)
LoraModel是如何实现的呢？其__init__函数如下
class LoraModel(torch.nn.Module):
def __init__(self, model, config, adapter_name):
super().__init__()
self.model = model
self.forward = self.model.forward
self.peft_config = config
self.add_adapter(adapter_name, self.peft_config[adapter_name])
主要实现对预训练模型add_adapter操作，该操作由_find_and_replace和mark_only_lora_as_trainable组成，_find_and_replace找到所有需要加入lora策略的层，例如q_proj，把它们替换成lora模式，mark_only_lora_as_trainable保留lora部分的参数可训练，其余参数全都固定下来不动。
def add_adapter(self, adapter_name, config=None):
if config is not None:
model_config = self.model.config.to_dict() if hasattr(self.model.config, "to_dict") else self.model.config
config = self._prepare_lora_config(config, model_config)
self.peft_config[adapter_name] = config
self._find_and_replace(adapter_name)
if len(self.peft_config) > 1 and self.peft_config[adapter_name].bias != "none":
raise ValueError(
"LoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters."
)
mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias)
if self.peft_config[adapter_name].inference_mode:
_freeze_adapter(self.model, adapter_name)
3.3 LoraLayer层的实现
定义低秩r、缩放参数alpha、归一化尺度字典scaling、A、B矩阵等，如下
class LoraLayer:
def __init__(self, in_features: int, out_features: int,**kwargs):
self.r = {}
self.lora_alpha = {}
self.scaling = {} # scaling[adapter_name] = lora_alpha / r
self.lora_dropout = nn.ModuleDict({})
self.lora_A = nn.ModuleDict({})
self.lora_B = nn.ModuleDict({})
# For Embedding layer
self.lora_embedding_A = nn.ParameterDict({})
self.lora_embedding_B = nn.ParameterDict({})
# Mark the weight as unmerged
self.merged = False
self.disable_adapters = False
self.in_features = in_features
self.out_features = out_features
self.kwargs = kwargs
def update_layer(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights):
...
def update_layer_conv2d(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights):
...
def update_layer_embedding(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights):
...
def reset_lora_parameters(self, adapter_name):
...
PEFT框架通过直接继承LoraLayer类已内置实现对nn.Linear、nn.Embedding、nn.Conv2d、bnb.nn.Linear8bitLt、bnb.nn.Linear4bit的Lora化支持，代码基本逻辑如下，每个模块可通过LoraLayer.disable_adapters字段决定forward推理中是否使用ΔW参数，使用则是Lora模型参数，否则为原模型参数。
class Linear(nn.Linear, LoraLayer):
...
def forward(self, x: torch.Tensor):
...
if self.disable_adapters:
if self.r[self.active_adapter] > 0 and self.merged:
self.unmerge()
...
...
class Embedding(nn.Embedding, LoraLayer):
...
class Conv2d(nn.Conv2d, LoraLayer):
...
class Linear8bitLt(bnb.nn.Linear8bitLt, LoraLayer):
...
class Linear4bit(bnb.nn.Linear4bit, LoraLayer):
...
3.4 三步实现LoRA及参数解释
# step1. 参数配置
R = 8
LORA_ALPHA = 16
LORA_DROPOUT = 0.04
TARGET_MODULES = ["q_proj", "v_proj"]
config = LoraConfig(
r=R,
lora_alpha=LORA_ALPHA,
target_modules=TARGET_MODULES,
lora_dropout=LORA_DROPOUT,
bias="none",
task_type="CAUSAL_LM",
)
# step2. 加载transformer预训练模型
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
# step3. 获取LoRA模型
model = get_peft_model(model, config)
PEFT源码中LoraConfig参数描述如下
class LoraConfig(PeftConfig):
"""
This is the configuration class to store the configuration of a [`LoraModel`].
Args:
r (`int`): Lora attention dimension.
target_modules (`Union[List[str],str]`): The names of the modules to apply Lora to.
lora_alpha (`int`): The alpha parameter for Lora scaling.
lora_dropout (`float`): The dropout probability for Lora layers.
fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out).
For example, gpt-2 uses `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`.:
bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only'
modules_to_save (`List[str]`):List of modules apart from LoRA layers to be set as trainable
and saved in the final checkpoint.
layers_to_transform (`Union[List[int],int]`):
The layer indexes to transform, if this argument is specified, it will apply the LoRA transformations on
the layer indexes that are specified in this list. If a single integer is passed, it will apply the LoRA
transformations on the layer at this index.
layers_pattern (`str`):
The layer pattern name, used only if `layers_to_transform` is different from `None` and if the layer
pattern is not in the common layers pattern.
"""
参数名
含义
r
lora的秩，矩阵A和矩阵B相连接的宽度，r«d
lora_alpha
尺度缩放参数，lora参数ΔWx乘以 α/r 尺度归一化 ,本质和learning rate相同
lora_dropout
lora层的dropout比率
bias
是否可训练bias，none：均不可；all：均可；lora_only：只有lora部分的bias可训练
modules_to_save
除了lora部分之外，还有哪些层可以被训练，并且需要保存
fan_in_fan_out
只有应用在Conv1D层时置为True，其他情况False
target_modules 指定应用lora的目标模块
**4.**LoRA 模型微调范式示例
前述得到Lora的模型peft model，接下来如何实现模型存储加载、混合精度训练？这里我们给出一个有代表性的范例，基于该范例根据具体任务做适当调整即可。
import datasets
from transformers import Trainer, DataCollatorForSeq2Seq
if resume_from_checkpoint:
lora_weight = torch.load(ckpt_name)
set_peft_model_state_dict(model, lora_weight)
train_data = datasets.load_from_disk(dataset_path)
class ModifiedTrainer(Trainer):
def save_model(self, output_dir=None, _internal_call=False):
# 改写trainer的save_model，在checkpoint的时候只存lora权重
from transformers.trainer import TRAINING_ARGS_NAME
os.makedirs(output_dir, exist_ok=True)
torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
saved_params = {
k: v.to("cpu") for k, v in self.model.named_parameters() if v.requires_grad
}
torch.save(saved_params, os.path.join(output_dir, "adapter_model.bin"))
trainer = ModifiedTrainer(
model=model,
train_dataset=train_data,
args=transformers.TrainingArguments(
per_device_train_batch_size=8,
gradient_accumulation_steps=16,
num_train_epochs=10,
learning_rate=3e-4,
fp16=True,
logging_steps=10,
save_steps=200,
output_dir=output_dir
),
data_collator=DataCollatorForSeq2Seq(
tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
),
)
trainer.train()
model.save_pretrained(train_args.output_dir)
因为peft model重写了原始model的save_pretrained函数，只把lora层的权重及配置文件进行存储，因此model.save_pretrained只会存储lora权重，这里trainer的save_model函数没有做相应的重写，因此我们重写下对应的function，避免checkpoint写入原始模型全部参数。
总结，本文对PEFT框架、LoRA原理做了简单介绍，也通过代码对PEFT中不同算子Lora化的逻辑和细节进行分析，对这些分析的理解将有助于开发者对不同transformer实现更灵活的Lora模型，助力实现对大模型的微调及应用落地。
*—-
END—- *
附录：
https://github.com/huggingface/peft
https://github.com/huggingface/safetensors
https://huggingface.co/docs/transformers/model_doc/rag
https://github.com/microsoft/LoRA
https://cloud.tencent.com/developer/article/2276508
进技术交流群请添加AINLP小助手微信（id: ainlp2)
请备注具体方向+所用到的相关技术点
关于AINLP
AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括LLM、预训练模型、自动生成、文本摘要、智能问答、聊天机器人、机器翻译、知识图谱、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLP小助手微信(id：ainlp2)，备注工作/研究方向+加群目的。
阅读至此了，分享、点赞、在看三选一吧🙏
更多AI工具，参考Github-AiBard123，国内AiBard123
