上一个说的不够够全面没有代码 这里在说一次
loralib使用方法
安装
这个很简单，在第一部分的环境配置里我们也讲过
pip install loralib
# Alternatively
# pip install git+https://github.com/microsoft/LoRA低秩适配器就是图中右半部分的矩阵。loralib为以下四种pretrain weights提供了对应的低秩适配器（大白话：定义 module长什么样子）
-
nn.Linear：普通线性层 -
nn.Embedding：Embedding层 -
nn.Conv2d：卷积神经网络层 -
nn.MergedLinear：混合线性层
nn.MergedLinear是指，我们在计算attention时，既可以把qkv拆成3个独立矩阵计算，也可以把他们拼在一起，变成一个矩阵计算。这块也是我们后文会重点分析的部分。构造低秩适配器的代码如下：
# ------------------------------------------------
# Before
# 没有lora前，layer指的是pretrain weights
# ------------------------------------------------
# layer = nn.Linear(in_features, out_features)
# ------------------------------------------------
# After
# 使用lora后，layer指的是对应的低秩适配器
# ------------------------------------------------
import loralib as lora
# Add a pair of low-rank adaptation matrices with rank r=16
layer = lora.Linear(in_features, out_features, r=16)
nn.MergedLinear使用方法：
# ------------------------------------------------
# Before
# 使用lora前，我们正常定义qkv矩阵
# 这里采用的是3个矩阵合1的定义方法
# ------------------------------------------------
# qkv_proj = nn.Linear(d_model, 3*d_model)
# ------------------------------------------------
# After
# 使用lora后，我们定义qkv矩阵的低秩适配器
# 既可以3个适配器分开定义，也可以合起来定义
# ------------------------------------------------
# Break it up (remember to modify the pretrained checkpoint accordingly)
q_proj = lora.Linear(d_model, d_model, r=8)
k_proj = nn.Linear(d_model, d_model)
v_proj = lora.Linear(d_model, d_model, r=8)
# Alternatively, use lora.MergedLinear (recommended)
qkv_proj = lora.MergedLinear(d_model, 3*d_model, r=8, enable_lora=[True, False, True])
使用lora进行微调训练
一切尽在注释中～
import loralib as lora
# ------------------------------------------------
# model是指已经添加过lora低秩适配器的model
# 其lora相关的layer，名字都以"lora_"开头
# ------------------------------------------------
model = BigModel()
# ------------------------------------------------
# 使用lora微调时，我们冻结pretrain，只训练低秩适配器
# 通过mark_only_lora_as_trainable，我们将所有
# 不是以"lora_"开头的参数层的requires_grad设为False
# ------------------------------------------------
lora.mark_only_lora_as_trainable(model)
# ------------------------------------------------
# 正常训练
# ------------------------------------------------
for batch in dataloader:
...
model = BigModel()中，BugModel()已经过了我们的手动改造。也就是，我们在预训练模型的架构上，通过2.2.2中提供的方法，为模型添加了LoRA适配器。这也是本文开头所说的“繁重的手工活”。在HuggingFace Peft中，我们只需要加载其支持的预训练模型，然后在相关配置中指明需要在模型的哪些层上添加低秩矩阵即可，只需几行代码，替我们节省了手工改造的时间。
另外，mark_only_lora_as_trainable默认是不train低秩适配器的bias的，如果想train bias，可通过以下方法进行设置：
# ------------------------------------------------
# Before
# 不训练任何bias
# ------------------------------------------------
# lora.mark_only_lora_as_trainable(model) # Not training any bias vectors
# ------------------------------------------------
# After
# 根据需要决定是否要训练bias
# ------------------------------------------------
# Training all bias vectors associated with modules we apply LoRA to
lora.mark_only_lora_as_trainable(model, bias='lora_only')
# Alternatively, we can train *all* bias vectors in the model, including LayerNorm biases
lora.mark_only_lora_as_trainable(model, bias='all')
# When saving a checkpoint, use the same bias= ('all' or 'lora_only')
torch.save(lora.lora_state_dict(model, bias='all'), checkpoint_path)
保存checkpoint
loralib只保存低秩适配器部分的checkpoint，代码如下：
# ------------------------------------------------
# Before
# 使用lora前，保留的整体模型的checkpoint
# ------------------------------------------------
# torch.save(model.state_dict(), checkpoint_path)
# ------------------------------------------------
# After
# 使用lora后，只保存低秩适配器部分的checkpoint】
# ------------------------------------------------
torch.save(lora.lora_state_dict(model), checkpoint_path)
读取lora模型
由2.2.4可知，lora只保存低秩部分的权重结果。当我们微调好模型，想要再次读取它时，我们需要：
-
首先，我们要定义好
model，这个model是指已经过lora改造的模型，和2.2.3中的BigModel()是一个意思。也就是说，在这个模型里，和lora相关的参数层，其名字都是以"lora_"开头。 -
然后，我们读取预训练部分的权重
-
最后，我们再读取保存下来的低秩适配器权重
由这个过程可知，在loralib中，预训练权重和低秩适配器权重没有合在一起！它们是相互独立的。当然，等学习完源码后，我们也可以根据需要改写代码，决定是否要合权重。整体代码如下：
# Load the pretrained checkpoint first
model.load_state_dict(torch.load('ckpt_pretrained.pt'), strict=False)
# Then load the LoRA checkpoint
model.load_state_dict(torch.load('ckpt_lora.pt'), strict=False)
model.train()与model.eval()
在LoRA论文中曾提过，LoRA在做推理时有一个显著的优势：可以将低秩适配器的权重直接合并到预训练权重里，这样就和全参数微调一样，不会产生推理上的额外耗时。
我们在2.2.5中说过，经过改造的基础模型分为“预训练部分”和“低秩适配器”部分。当我们开启model.train()时，我们是用拆开的预训练和低秩适配器做训练的；当我们开启model.eval()时，我们则是先将低秩适配器合入预训练权重，再做推理。看到这里觉得抽象没关系，在后文里，我们会来看代码如何实现这一块。
好，到这一步，我们讲完了loralib的整体使用方式。接下来，我们就分模块开始讲解代码实现细节吧！
整体训练流程
-
gpt2_ft.py：使用lora微调gpt2的入口函数。包含了train/valid数据集划分，模型训练等步骤。 -
gpu.py：分布式训练相关配置 -
model.py：定义了添加低秩适配器后的gpt2模型架构
我们重点关注gpt2_ft.py和model.py
入口函数
首先来看gpt2_ft.py，它定义了整体训练框架。
if __name__ == '__main__':
# ------------------------------------------------------------
# 1. 分布式环境初始化
# ------------------------------------------------------------
# 解析入参
args = parser.parse_args()
# 根据入参相关配置，使用torch.distributed做分布式环境初始化(DDP)
# 相关代码在gpu.py中，这里不另做说明
parse_gpu(args)
print_args(args)
# 混合精度训练相关配置
if args.fp16:
try:
from apex import amp
except Exception as e:
warnings.warn('Could not import amp, apex may not be installed')
# 设置随机种子
torch.manual_seed(args.random_seed)
random.seed(args.random_seed)
# 在进程0所在的机器上，创建work_dir目录，
# 用于保存低秩适配器checkpoint及最终结果
if args.rank == 0:
args.logging = create_exp_dir(args.work_dir)
# ------------------------------------------------------------
# 2、划分train/valid数据集
# ------------------------------------------------------------
# FT_Dataset中包含了对数据集的处理，例如区分context，completion，
# 训练target等，这里不做介绍，大家可以看data_utils.py中的细节
train_data = FT_Dataset(
args.train_data, args.train_batch_size, args.seq_len,
joint_lm=args.obj=='jlm'
)
valid_data = FT_Dataset(
args.valid_data, args.valid_batch_size, args.seq_len,
)
train_loader = DataLoader(
train_data, batch_size=args.train_batch_size, num_workers=0,
shuffle=False, pin_memory=False, drop_last=True,
sampler=torch.utils.data.distributed.DistributedSampler(train_data, seed=args.random_seed)
)
valid_loader = DataLoader(
valid_data, batch_size=args.valid_batch_size, num_workers=0,
shuffle=False, pin_memory=False, drop_last=False,
sampler=torch.utils.data.distributed.DistributedSampler(valid_data, seed=args.random_seed)
)
# ------------------------------------------------------------
# 3、定义模型架构（pretrain + 低秩适配器）
# 读取pretrain部分权重
# ------------------------------------------------------------
# 对不同的pretrain模型设置不同的配置（gpt2 small/medium/large）
if args.model_card == 'gpt2.sm':
config = GPT2Config(
n_embd=768, n_layer=12, n_head=12,
lora_attn_dim=args.lora_dim,
lora_attn_alpha=args.lora_alpha,
lora_dropout=args.lora_dropout,
)
elif args.model_card == 'gpt2.md':
config = GPT2Config(
n_embd=1024, n_layer=24, n_head=16,
lora_attn_dim=args.lora_dim,
lora_attn_alpha=args.lora_alpha,
lora_dropout=args.lora_dropout,
)
elif args.model_card == 'gpt2.lg':
config = GPT2Config(
n_embd=1280, n_layer=36, n_head=20,
lora_attn_dim=args.lora_dim,
lora_attn_alpha=args.lora_alpha,
lora_dropout=args.lora_dropout,
)
# 定义gpt2模型，GPT2LMModel()对原始gpt2做了改造，添加了低秩适配器
lm_net = GPT2LMModel(config)
# 读取pretrain部分权重，此时低秩适配器部分的权重值为其初始化值（A阵随机高斯，B阵为0）
if args.init_checkpoint is not None:
print('loading model pretrained weight.')
lm_net.load_weight(torch.load(args.init_checkpoint))
# 把模型搬运到gpu上
lm_net = lm_net.cuda()
# ------------------------------------------------------------
# 4、lora微调相关配置
# ------------------------------------------------------------
# 使用mark_only_lora_as_trainable包装lm_net，表示只对低秩矩阵做训练
if args.lora_dim > 0:
lora.mark_only_lora_as_trainable(lm_net)
# 构造优化器
optimizer = create_adam_optimizer_from_args(lm_net, args)
# 每块卡上最多做多少次step
if args.max_step is None:
args.max_step = (args.max_epoch * train_data.num_batches + args.world_size - 1) // args.world_size
print('set max_step:', args.max_step)
scheduler = create_optimizer_scheduler(optimizer, args)
if args.fp16:
lm_net, optimizer = amp.initialize(lm_net, optimizer, opt_level="O1")
# 将模型和优化器包装为分布式的
lm_net, optimizer = distributed_opt(args, lm_net, optimizer, grad_acc=args.grad_acc)
# ------------------------------------------------------------
# 5、训练
# 核心函数为train_validate(), 在后文会详细解读
# ------------------------------------------------------------
try:
train_step = 0
for epoch in itertools.count(start=1):
# 定义每个step训练步骤，包含train和validate两步
train_step = train_validate(
lm_net, optimizer, scheduler, train_loader, valid_loader, args,
train_step=train_step, epoch=epoch
)
if train_step >= args.max_step or (args.max_epoch is not None and epoch >= args.max_epoch):
if args.rank == 0:
print('-' * 100)
print('End of training')
break
except KeyboardInterrupt:
if args.rank == 0:
print('-' * 100)
print('Exiting from training early')
distributed_sync(args)
print('cleanup dist ...')
cleanup(args)
好！到此我们介绍完了如何使用loralib来写微调的main()函数，现在我们需要重点关注两个细节：
-
预训练模型要如何改写？代码第58行
lm_net = GPT2LMModel(config)，根据前面的介绍，我们知道GPT2LMModel这个模型是经过手动改写的：在原始GPT2的模型架构上，添加了低秩适配模块（也就是矩阵A和B）。那你添加了相应的模块，肯定要定义module架构，forward方法之类的呀。所以，这是我们要关注的细节。 -
模型的训练(train)和验证(validate)具体逻辑是怎么样的？代码第122行，我们采用
train_validate这个函数实现了这一步。前文说过，训练和验证的过程，包含了低秩适配权重是否要合进预训练权重，怎么合进预训练权重的关键问题。也就是model.train()和model.eval()具体要如何改写的问题。这也是lora的重点之一。
所以接下来，让我们详细来看这两个部分。
重写预训练模型（添加低秩适配器）
相关代码在model.py文件下。总结来说，就是在GPT2模型架构的基础上，为模型相应的部分（例如attention层）添加lora低秩适配器。我们以Attention层的代码为例，我们关注的重点在代码第13行相关部分（一切尽在注释中）：
class Attention(nn.Module):
def __init__(self, nx, n_ctx, config, scale=False):
super(Attention, self).__init__()
n_state = nx # in Attention: n_state=768 (nx=n_embd)
# [switch nx => n_state from Block to Attention to keep identical to TF implem]
assert n_state % config.n_head == 0
# register_buffer的含义：作为一个常量储存起来，不会对它做梯度更新
self.register_buffer("bias", torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
self.n_head = config.n_head
self.split_size = n_state
self.scale = scale
# ------------------------------------------------
# 使用lora.MergedLinear改写attention层
# 改写过后的attention层包括预训练部分和低秩适配器部分
# ------------------------------------------------
self.c_attn = lora.MergedLinear(
nx, # embed_dim
n_state * 3, # embed_dim * 3，因为这里将qkv融合成一个矩阵处理
r=config.lora_attn_dim, # r值，即lora中的秩
lora_alpha=config.lora_attn_alpha, # alpha值，用于表示微调过程中对新知识的侧重程度
lora_dropout=config.lora_dropout, # dropout值
enable_lora=[True, False, True], # 表示对qkv三个矩阵中的q，v做低秩适配，k保持原样
fan_in_fan_out=True, # 表示在计算时是否要做矩阵专置
merge_weights=False # 表示是否想将低秩适配权重合并到预训练权重中
)
self.c_proj = Conv1D(n_state, nx)
self.config = config
def _attn(self, q, k, v, len_kv=None):
...
def merge_heads(self, x):
...
def split_heads(self, x, k=False):
...
def forward(self, x, history=None, layer_past=None, len_past=None):
...
x = self.c_attn(x)
...
lora.MergedLinear将被作为lora核心代码处理逻辑，在本文的第四部分详细阐述。在这里大家只要知道如何通过它来改写原始预训练模型即可。
特别值得补充的是，关于lora中的alpha/r这个系数，今天我在github issue上翻到了作者两周前的一条回复：
老外的图不用上..
大意是说，在作者之前的研究中发现，对于数据经过矩阵B、但还没有过激活层时的那个数值，其波动幅度和r有相关性（顺藤摸瓜找到了相关论文：https://arxiv.org/pdf/2011.14522.pdf，写得比较深，还没细看）。因此通过1/r这样的因子，来消除这种影响，使得模型训练更稳定，否则，就需要去调learning rate了（猜想之所以波动幅度和r相关，也是因为越大的r带来的噪声冗余越多，这点之前在原理篇有给过分析）。从这点上说，alpha可以纯被解释为模型对新知识的侧重程度，而1/r则是一种scaling rate。所以我在代码注释中，沿用了这种解释。
回过头来，通过这段代码，我们也感受了一下前文说的“手工活”，在peft中，你只需要加载hugging face支持的模型，通过简单的几行代码就能完成模型架构的修改了。示例如下：
from peft import LoraConfig, get_peft_model
config = LoraConfig(
r=16,
lora_alpha=16,
target_modules=["query", "value"], # 定义需要做低秩适配的部分
lora_dropout=0.1,
bias="none",
modules_to_save=["classifier"],
)
lora_model = get_peft_model(model, config)
print_trainable_parameters(lora_model)
"trainable params: 667493 || all params: 86466149 || trainable%: 0.77"
训练和验证
train_validate()函数同样在gpt2_ft.py下，我们来看具体代码（一切尽在注释中）：
def train_validate(
model, # 用lora包过的模型
optimizer,
scheduler,
train_loader,
valid_loader,
args, # 给类输入参数
train_step=0, # 当前train_step(相对于每个进程而言的)
epoch=0 # 当前epoch
):
# -----------------------------------------------
# 1. 模型训练
# -----------------------------------------------
# 开启训练模式，此时pretrain和低秩适配部分是分开的，后文中我们会看到代码细节
model.train()
# 用于计算到目前为止的loss总和、loss均值
avg_lm_loss = AverageMeter()
print('start to train the model................', epoch)
log_start_time = time.time()
# 用于记录模型训练过程中，基于valid数据集评估的最好效果
best_val_ppl = None
train_loader.sampler.set_epoch(epoch)
# 遍历每一个batch
for idx, data in enumerate(train_loader):
data = {key: value for key, value in data.items()}
# 把这个batch的数据搬运到gpu上
_input = data['input'].to(args.device)
_target = data['target'].to(args.device)
_msk = data['mask'].to(args.device)
_lm_logits, _lm_loss = model(
_input, lm_labels=_target, lm_mask=_msk, label_smooth=args.label_smooth
)
_lm_loss = _lm_loss.mean()
train_step += 1
# ---------------------------------------------------------------------------------
# args.grad_acc：表示要累计多少次step再做梯度更新
# 为什么要_lm_loss/(args.grad_acc)？
# 因为梯度累积的目的是，为了节省显存，把大batch拆成小batch来跑，所有小batch跑完以后再更新梯度
# 因此小batch更新梯度的结果需要等于大batch更新梯度的结果
# 大batch更新梯度结果：求导mean_loss/w，其中mean_loss是平均单条样本loss
# 小batch更新梯度结果：求导sum(mean_loss)/w，
# 也意味着，不做任何处理，分子相当于单条样本loss计算了grad_acc次
# 因此需要做: mean_loss / grad_acc
# 可配合下文optimizer_step()这个函数的注释来阅读
# ---------------------------------------------------------------------------------
is_update = True if train_step % args.grad_acc == 0 else False
avg_lm_loss.update(_lm_loss.item())
optimizer_step(
_lm_loss/(args.grad_acc), optimizer, model, scheduler, args, is_update=is_update
)
# 打印训练日志
if train_step % args.log_interval == 0:
elapsed = time.time() - log_start_time
lr = optimizer.param_groups[0]['lr']
log_str = f'| epoch {epoch:3d} step {train_step:>8d} | { idx + 1:>6d} batches | ' \
f'lr {lr:.3g} | ms/batch {elapsed * 1000 / args.log_interval:5.2f} | ' \
f'loss {avg_lm_loss.val:5.2f} | avg loss {avg_lm_loss.avg:5.2f} | ' \
f'ppl {math.exp(avg_lm_loss.avg):5.2f}'
if args.rank == 0:
print(log_str)
log_start_time = time.time()
avg_lm_loss.reset()
# 存储checkpoint
if train_step % args.save_interval == 0:
if args.rank == 0:
model_path = os.path.join(args.work_dir, f'model.{train_step}.pt')
print('saving checkpoint', model_path)
# 调用lora的lora_state_dict，在存储的时候只会保存lora部分的权重
torch.save({'model_state_dict': lora.lora_state_dict(model)}, model_path)
distributed_sync(args)
# -----------------------------------------------
# 2. 模型验证
# -----------------------------------------------
if train_step % args.eval_interval == 0:
eval_start_time = time.time()
# 在evaluate()开启了model.eval()
valid_loss, valid_ppl = evaluate(model, valid_loader, args)
# 记录基于验证数据集的最好模型效果
if best_val_ppl is None or valid_ppl < best_val_ppl:
best_val_ppl = valid_ppl
log_str = f'| Eval {train_step // args.eval_interval:3d} at step {train_step:>8d} | ' \
f'time: {time.time() - eval_start_time:5.2f}s | valid loss {valid_loss:5.2f} | ' \
f'valid ppl {valid_ppl:5.2f} | best ppl {best_val_ppl:5.2f} '
if args.rank == 0:
print('-' * 100)
print(log_str)
print('-' * 100)
# 验证完毕后，需要重新开启model.train()模式
model.train()
distributed_sync(args)
if train_step == args.max_step:
break
# -----------------------------------------------
# 3. 保存模型训练结果
# 只保存lora部分的结果
# -----------------------------------------------
if args.rank == 0:
model_path = os.path.join(args.work_dir, f'model.{train_step}.pt')
print('saving checkpoint', model_path)
torch.save({'model_state_dict': model.state_dict()}, model_path)
distributed_sync(args)
return train_step
def optimizer_step(_loss, _optimizer, _model, _schedule, args, is_update=True):
"""
定义梯度更新参数的策略
"""
if args.fp16: # 如果采用混合精度训练的话，记得要scale loss，以防出现梯度nan或inf
with amp.scale_loss(_loss, _optimizer) as _scaled_loss:
_scaled_loss.backward()
else: # 正常计算当前梯度
_loss.backward()
# 若当前step可以做梯度更新，则对梯度做clip后正常更新权重
if is_update:
if args.clip > 0:
if args.fp16:
torch.nn.utils.clip_grad_norm_(amp.master_params(_optimizer), args.clip)
else:
torch.nn.utils.clip_grad_norm_(_model.parameters(), args.clip)
# 更新权重
_optimizer.step()
# 梯度清零
_optimizer.zero_grad()
if _schedule is not None:
_schedule.step()
class AverageMeter(object):
"""Computes and stores the average and current value
Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262
"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val # 当前值
self.sum += val * n # 求和
self.count += n # 求次数
self.avg = self.sum / self.count # 求均值
梯度累积(grad_acc)更新
细节都写在代码注释中了。额外需要提的一点，是代码中关于_lm_loss/(args.grad_acc)这一步。
在模型训练的过程中，我们有时不会一个step更新一次梯度，而是累积若干次step（即代码中的grac_acc值）再做一次梯度更新，这样做的原因是：
-
减少模型的更新频率，一定程度上加速训练速度
-
节省显存。可以理解为，当一个大batch将显存打爆时，我把它拆成若干个小batch来跑，此时我们理所当然希望这个大batch和这若个干小batch对梯度计算的结果尽量是一样的。
低秩适配器运作细节
现在，我们以lora.MergedLinear为例，来看看实现低秩适配器的代码细节(一切尽在注释中)：
class LoRALayer():
def __init__(
self,
r: int,
lora_alpha: int,
lora_dropout: float,
merge_weights: bool,
):
# -----------------------------------------------------
# lora的秩
# -----------------------------------------------------
self.r = r
# -----------------------------------------------------
# 用于计算缩放因子scaling = lora_alpha/r
# -----------------------------------------------------
self.lora_alpha = lora_alpha
# -----------------------------------------------------
# Optional dropout
# -----------------------------------------------------
if lora_dropout > 0.:
self.lora_dropout = nn.Dropout(p=lora_dropout)
else:
self.lora_dropout = lambda x: x
# -----------------------------------------------------
# 表示当前pretrain部分(self.weight)中是否已经融入了lora部分
# self.merged=True，pretrain部分已包含lora,
# 则进行forward是可以直接用pretrain部分
# self.merged=False, pretrain部分未包含lora，
# 则进行forward时需要用pretrain+lora的结果
# 【表示合不合这个状态】
# -----------------------------------------------------
self.merged = False
# -----------------------------------------------------
# self.merge_weights = True，遵从训练时拆分pretain与lora，
# 推理时合并pretrain与lora
# self.merge_weights = False, 遵从无论是训练还是推理，
# 都拆分pretrain与lora
# 【表示合不合这个意愿】
# -----------------------------------------------------
self.merge_weights = merge_weights
class MergedLinear(nn.Linear, LoRALayer):
# LoRA implemented in a dense layer
def __init__(
self,
in_features: int,
out_features: int,
r: int = 0,
lora_alpha: int = 1,
lora_dropout: float = 0.,
enable_lora: List[bool] = [False],
fan_in_fan_out: bool = False,
merge_weights: bool = True,
**kwargs
):
# --------------------------------------------------------------------
# in_features: hidden_size
# out_features: hidden_size*3
# 将nn.linear和LoRALayer的属性继承过来
# --------------------------------------------------------------------
nn.Linear.__init__(self, in_features, out_features, **kwargs)
LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
merge_weights=merge_weights)
# --------------------------------------------------------------------
# 因为我们是把QKV三个矩阵的计算合成一个来计算的，
# 因此out_features要能被矩阵个数整除
# --------------------------------------------------------------------
assert out_features % len(enable_lora) == 0, \
'The length of enable_lora must divide out_features'
# --------------------------------------------------------------------
# 表示哪一块矩阵是需要做lora的，例如[True, False, True]，
# 说明第一和第三块需要，第二块不要
# --------------------------------------------------------------------
self.enable_lora = enable_lora
# --------------------------------------------------------------------
# bool值，表示是否需要对矩阵做转置，下文中会看到细节
# --------------------------------------------------------------------
self.fan_in_fan_out = fan_in_fan_out
# --------------------------------------------------------------------
# Actual trainable parameters
# 如果秩>0，且其中有任意一矩阵需要做lora
# --------------------------------------------------------------------
if r > 0 and any(enable_lora):
# --------------------------------------------------------------------
# 初始化A矩阵
# self.weight继承自nn.linear下的属性，shape=(in_features, out_features)
# =(hidden_size, hidden_size*3)
# lora_A的shape=(r*需要做lora的矩阵数量, hidden_size)
# 使用new_zeros，可以复制原来矩阵的数据类型和存储设备
# --------------------------------------------------------------------
self.lora_A = nn.Parameter(
self.weight.new_zeros((r * sum(enable_lora), in_features)))
# --------------------------------------------------------------------
# 初始化B矩阵
# lora_B的shape=(hidden_size*需要做lora的矩阵数量, r)
# --------------------------------------------------------------------
self.lora_B = nn.Parameter(
self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r))
) # weights for Conv1D with groups=sum(enable_lora)
# --------------------------------------------------------------------
# lora权重的缩放因子，
# 一方面决定预训练知识和lora学得的新知识间的比例，另一方面通过1/r因子使得
# 训练过程更加稳定（参考本文3.2部分的说明，以及lora原理篇中的讲解）
# 在gpt2做nlg的任务里，作者将lora_alpha设置为32，r设置为4
# --------------------------------------------------------------------
self.scaling = self.lora_alpha / self.r
# --------------------------------------------------------------------
# Freezing the pre-trained weight matrix
# 原始的QKV矩阵需要冻结
# --------------------------------------------------------------------
self.weight.requires_grad = False
# --------------------------------------------------------------------
# Compute the indices
# 1. lora_ind的shape=(3, hidden_size)，默认值为False
# 2. 对lora_ind，只有需要做lora矩阵的那些行，才填充为True，例如hidden_size = 5,
# 第一和第三个矩阵需要做lora，则lora_ind变为：
# tensor([[ True, True, True, True, True],
# [False, False, False, False, False],
# [ True, True, True, True, True]])
# 最后，将lora_ind的形式转成shape=(3*hidden_size)
# --------------------------------------------------------------------
self.lora_ind = self.weight.new_zeros(
(out_features, ), dtype=torch.bool
).view(len(enable_lora), -1)
self.lora_ind[enable_lora, :] = True
self.lora_ind = self.lora_ind.view(-1)
# 对lora_A和lora_B重新做参数初始化
self.reset_parameters()
# --------------------------------------------------------------------
# 若fan_in_fan_out = True,
# 则把weight的shape(hidden_size, 3*hidden_size) ->
# (3*hidden_size, hidden_size)
# 这样做的原因是，例如输入x=(b,s,h), w=(h, 3h)
# 调用F.linear(x, w)时，默认执行的是(b,s,h)*(3h,h)，
# 因此要通过这个方式，把w先做转置
# --------------------------------------------------------------------
if fan_in_fan_out:
self.weight.data = self.weight.data.transpose(0, 1)
def reset_parameters(self):
"""
重新初始化参数
"""
nn.Linear.reset_parameters(self)
if hasattr(self, 'lora_A'):
# --------------------------------------------------------------------
# initialize A the same way as the default for nn.Linear and B to zero
# 对lora_A，采用kaiming_uniform; 对lora_B，初始化为0
# --------------------------------------------------------------------
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
nn.init.zeros_(self.lora_B)
def zero_pad(self, x):
"""
x是数据过B矩阵后的结果，x.shape =(b, s, hidden_size*需要做lora的矩阵数量)
我们要将这个结果做zero padding，使得：
x.shape = (b, s, hidden_size*总矩阵数量)，在QKV的例子中，总矩阵数量=3，
其中不是我们要做lora矩阵的部分，都padding上0
"""
result = x.new_zeros((*x.shape[:-1], self.out_features))
result = result.view(-1, self.out_features)
result[:, self.lora_ind] = x.reshape(
-1, self.out_features // len(self.enable_lora) * sum(self.enable_lora)
)
return result.view((*x.shape[:-1], self.out_features))
def train(self, mode: bool = True):
def T(w):
return w.transpose(0, 1) if self.fan_in_fan_out else w
nn.Linear.train(self, mode)
# --------------------------------------------------------------
# 如果是model.train()模式
# --------------------------------------------------------------
if mode:
# --------------------------------------------------------------
# 参见4.1讲解
# --------------------------------------------------------------
if self.merge_weights and self.merged:
# Make sure that the weights are not merged
if self.r > 0 and any(self.enable_lora):
# ---------------------------------------------------
# delta_w: lora_A * lora_B矩阵后的结果
# 输入x * delta_w是最终lora矩阵的输出
# lora_A: shape=(r*需要做lora的矩阵数, h)
# lora_B: shape=(h*需要做lora的矩阵数, r)
# 过F.conv1d(lora_A当成输入，lora_B当成卷积)后的：
# delta_w = (1, h*需要做lora的矩阵数, h)
# squeeze之后delta_w = (h*需要做lora的矩阵数, h)
# ---------------------------------------------------
delta_w = F.conv1d(
self.lora_A.data.unsqueeze(0),
self.lora_B.data.unsqueeze(-1),
groups=sum(self.enable_lora)
).squeeze(0)
# 从pretrain矩阵中剔除lora部分
self.weight.data -= self.zero_pad(T(delta_w * self.scaling))
# 当原来的lora剥离出来后，self.merge要设置成False，表示此时pretrain中不再有lora
self.merged = False
# --------------------------------------------------------------
# 如果是train.eval()模式: model.train(False) = train.eval()
# --------------------------------------------------------------
else:
# --------------------------------------------------------------
# 参见4.1讲解
# --------------------------------------------------------------
if self.merge_weights and not self.merged:
# Merge the weights and mark it
if self.r > 0 and any(self.enable_lora):
delta_w = F.conv1d(
self.lora_A.data.unsqueeze(0),
self.lora_B.data.unsqueeze(-1),
groups=sum(self.enable_lora)
).squeeze(0)
self.weight.data += self.zero_pad(T(delta_w * self.scaling))
# 既然已将lora合进pretrain，那么就要把self.merged设置为True
self.merged = True
def forward(self, x: torch.Tensor):
"""
Params:
x: 输入数据，形状为(b, s, h)
"""
# --------------------------------------------------------------------
# 定义矩阵专置函数
# 例如原始矩阵是(hidden_size, 3*hidden_size),
# 转置之后为(3*hidden_size, hidden_size)
# --------------------------------------------------------------------
def T(w):
return w.transpose(0, 1) if self.fan_in_fan_out else w
# --------------------------------------------------------------------
# self.merged = True，如果lora和pretrain部分已经合起来
# 返回结果的shape = (b, s, 3*hidden_size)
# --------------------------------------------------------------------
if self.merged:
return F.linear(x, T(self.weight), bias=self.bias)
# --------------------------------------------------------------------
# self.merged = False，如果lora和pretrain部分没有合起来
# --------------------------------------------------------------------
else:
# 先用merged后的矩阵正常计算，result的shape = (b, s, 3*hidden_size)
result = F.linear(x, T(self.weight), bias=self.bias)
if self.r > 0:
# --------------------------------------------------------------------
# 数据过lora_A后的结果
# 结果shape = (b, s, h) * (h, r*需要做lora的矩阵的数量)
# = (b, s, r*需要做lora的矩阵的数量)
# --------------------------------------------------------------------
after_A = F.linear(self.lora_dropout(x), self.lora_A)
# --------------------------------------------------------------------
# 数据过lora_A，再过lora_B后的结果
# F.conv1d(input=输入数据，weight=卷积核)
# input = 输入数据过lora_A后的结果，shape = (b, r*需要做lora矩阵的数量, s)
# weight = (hidden_size*需要做lora的矩阵数量, r, 1)
# groups = 需要做lora的矩阵数量
# 通过F.conv1d计算出after_B的shape = (b, hidden_size*需要做lora的矩阵数量, s)
# 做transpose(-2, -1)转置后after_B的shape = (b, s, hidden_size*需要做lora的矩阵数量)
# --------------------------------------------------------------------
after_B = F.conv1d(
after_A.transpose(-2, -1),
self.lora_B.unsqueeze(-1),
groups=sum(self.enable_lora)
).transpose(-2, -1)
# --------------------------------------------------------------------
# 数据最终结果 = pretrain + lora
# zero_pad：将过B矩阵后的数据从shape=(b, s, hidden_size*需要做lora的矩阵数量)
# 变成(b, s, hidden_size*总矩阵数量)，
# 其中不属于需要做lora的部分就用0填重
# self.scaling = alpha/r，但暂时不知道scaling的作用
# --------------------------------------------------------------------
result += self.zero_pad(after_B) * self.scaling
return result
self.merged与self.merged_weights
其实这段代码理解起来并不复杂，就是一堆矩阵计算定义而已。如果你在看完注释后，对代码仍有疑惑，可以按照输入参数的定义，自己模拟一些矩阵，跑一遍这段代码，把中间结果的shape和具体数值打印出来，就基本能解决困惑啦～
这里像特别讲解的，是train()函数中，关于训练和推理时，要不要合并预训练权重和低秩适配器权重的问题。因为我初读这部分代码时，被绕晕了，特别是self.merged和self.merged_weights的作用方式。
首先明确一点：
-
self.merged表示【合不合这个状态】，即当前模型中低秩适配器的权重是否已经合并到预训练权重中。正因此，该参数是不可人为传入的，而是随着代码的执行流程而变动。在模型初始化时（见以上代码片段中LoRALayer()相关定义）该参数的取值为False，这也很好理解，因为在训练开始时，预训练权重就是和低秩适配器权重分开的。 -
self.merged_weights表示【合不合这个意愿】，具体是什么意愿我们在后面详说。因为是意愿，所以是可以人为传入的。那什么时候传入呢？在我们定义GPT2的attention层相关架构时（见3.2代码）。在作者写的这个GPT2的例子里，attention层self.merged_weights的初始化取值为False。
现在，我们来看self.merged_weights所指向的【意愿】到底是什么。在作者的设计中，当启动model.train()模式时，应当把预训练和低秩适配器拆开；当启动model.eval()时，应该把预训练和低秩适配器合并。但是，这毕竟只是作者的意愿，如果我想无论是训练还是推理，都把两者拆开，行不行？当然可以。所以：
-
self.merged_weights = True表示按作者的意愿做，训练不合，推理合。 -
self.merged_weights = False表示按你的意愿做，训练不合，推理也不合。
我们配合以上代码，先来看看按作者的意愿做会发生什么。
（1）首先，我们开启model.train()正常训练（即上述代码中进入if mode == True条件）。此时self.merged_weights = True, self.merged = False，真好，不用进行任何操作。数据就在权重拆开的情况下正常训练。
（2）然后，该用验证集进行推理了，我们开启model.eval()。 此时依然是self.merged_weights = True, self.merged = False，刚好满足if self.merge_weights and not self.merged这个判断条件。所以接下来我们就要进行相关操作了，把低秩适配器合并到预训练权重里，也就是代码中的self.weight.data += self.zero_pad(T(delta_w * self.scaling))。自然，合完了之后，表示【合不合这个状态】的self.merged就要设置为True。
（3）好了，现在做完一轮验证集验证了，我们也更新了一次当前模型在验证集上的最佳表现指标。现在我们得继续训练了，所以我们又切回model.train()模式。此时self.megred_weights = True, self.merged = True，意味着当前状态中预训练和低秩适配器的权重是合在一起的，满足了if self.merge_weights and self.merged这个条件。所以我们就要执行相关操作啦，也就是self.weight.data -= self.zero_pad(T(delta_w * self.scaling))，将低秩适配器从预训练权重里剥离开。自然，剥开之后，表示【合不合这个状态】的self.merged就要设置为False。
欸，经过这么一顺，是不是就不难理解了？whaosoft aiot http://143ai.com
那现在，我们再来看，按照我们的意愿把self.merged_weights设置为False的时候会发生什么。
在这种情况下，我们希望不管是训练还是推理，预训练和低秩适配器总是分开的。这时，在model.train()模式下，我们不会触发if self.merge_weights and self.merged判断条件；在model.eval()模式下，我们不会触发if self.merge_weights and not self.merged判断条件。也就是我们啥也不用做，就按照初始化设置的那样，从头到尾都把两个权重拆开。
一个待解疑惑
如果你仔细读过上面代码，你会发现一个神奇的地方，当数据过矩阵A时，我们采用的是普通的线性计算。但当数据过矩阵B时，采用的是F.conv1d（一维卷积)进行计算。
虽然不管是一位卷积，还是普通线性层，在这一块上的输出矩阵维度是一样的。但我目前依然没有想明白用卷积替换的原因。曾经想过保留上下文语义信息等原因，但是当画出一维卷积的过程时，又觉得似乎也没有达成这个目的。后面的就大家一起讨论吧~~~
参考
1、https://github.com/microsoft/LoRA
2、https://arxiv.org/pdf/2106.09685.pdf
1282
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
