深入浅出剖析 LoRA 源码及实践

在上一篇中，我们详细阐述了LoRA的原理。在本篇中，我们将一起学习LoRA源码（微软原版）。

许多朋友在使用LoRA的过程中，都会用到HuggingFace Peft库封装好的LoRA接口，这个接口是对微软版LoRA代码的改写和封装，目的是减少大家在使用LoRA过程中的手工活（例如徒手更改模型架构，为模型添加LoRA adapter结构等，在后文我们会细说），除此外核心处理逻辑不变。所以本文依然选用原版的LoRA代码来做解读。

为了帮助大家更好学习，正文的第一部分提供了详细的白嫖google colab GPU的教程😉。我们可以在colab上，git clone下LoRA原版代码，微调一个基于small gpt2的NLG任务，在实操过程中，如果对代码逻辑有困惑之处，也可以通过写test脚本、打印中间值等方式来解惑。总之，不要钱的快乐才是真的快乐。

但是，白嫖的GPU资源毕竟是有限的（内存12GB，显存15GB），因此colab上的微调只适合用来做代码学习。如果大家对效果有更强烈的需求，还是建议大家在学完后，用更好的卡，更好的预训练模型来做。

技术交流

技术要学会分享、交流，不建议闭门造车。一个人可以走的很快、一堆人可以走的更远。

相关资料、数据、技术交流提升，均可加我们的交流群获取，群友已超过2000人，添加时最好的备注方式为：来源+兴趣方向，方便找到志同道合的朋友。

方式①、添加微信号：mlc2060，备注：来自CSDN + 技术交流
方式②、微信搜索公众号：机器学习社区，后台回复：加群

最后，在之前原理篇的评论和私信中，我发现有很多朋友关心lora中的alpha参数。最近我在lora的github issue上又挖到一条lora一作给出的关于alpha的解释，我把它放在本文3.2的分析中了，感兴趣的朋友可以特别关注下。

话不多说，进入正文吧，全文目录如下：

一、在Google Colab上做LoRA微调

1、首先，需要有一个梯子，翻墙注册google账号。

2、使用账号登录google colab: https://colab.research.google.com/

3、点击左上角符号，进入google drive（云端硬盘）

4、点击左上角“+新建”按钮，依次点击“新建-更多-Google Colabratory"，建立jupyter notebook。

5、点击菜单栏“代码执行程序-更改运行时类型”，选择T4 GPU。

6、给自己的jupyter notebook命名，方便后续归档查找，例如我的就命名为“lora_learning.ipynb”。命名结束后，点击左侧文件夹按钮，出来右侧横排四个图标。点击第三个图标，装载google drive硬盘，装载成功后，该图标上会有一条斜杠（这个过程可能需要大家耐心等待几秒）。

连接成功后，可以看到“drive-MyDrive”目录，我们从github下clone下来的代码，以及我们自己写的学习、测试脚本等，就可以放在这个目录下。

7、在当前jupyter notebook脚本上，执行以下代码，切换工作目录，并克隆git代码：

import os

%cd '/content/drive/My Drive'
!mkdir lora_fine_tune

%cd './lora_fine_tune'
!git clone https://github.com/microsoft/LoRA.git

%cd '/content/drive/MyDrive/lora_fine_tune/LoRA/examples/NLG'

这样，我们就能把git代码成功克隆到lora_fine_tune这个文件夹下了。

8、接下来，我们来解决环境问题。一般来说，只要按照requirement.txt进行配置即可。但是google colab的环境比较特殊，它只支持固定的torch + cuda版本搭配，同时对我们常用的HuggingFace Transformer库的某些版本也存在不兼容的情况。好巧不巧，LoRA源码的环境配置完美撞在了这两个枪口上。在一番尝试后，我找到了一种不影响模型效果的最快捷的解决办法。

首先，将/content/drive/MyDrive/lora_fine_tune/LoRA/examples/NLG/requirement.txt这份文件做一些改动。

原版requirement.txt:

--find-links https://download.pytorch.org/whl/torch_stable.html
torch==1.7.1+cu101
transformers==3.3.1
spacy
tqdm
tensorboard
progress

改为：

--find-links https://download.pytorch.org/whl/torch_stable.html
spacy
tqdm
tensorboard
progress

然后，在jupyter cell上执行以下代码，配置环境：

!pip install loralib
!pip install -r requirement.txt
!pip install transformers==4.30.2

好！到这一步为止，所有代码下载和环境配置的环节，我们就完成了。接下来，我们就需要预处理数据集，并加载我们的预训练模型了！

9、加载预训练模型的代码在download_pretrained_checkpoints.sh这个脚本下，我们通过wget的方式从s3上下载模型。在原始脚本中，会同时下载gpt2 small/medium/large的模型。但作为学习需要，我们只用保留gpt2 small，所以可以把下载另外两个模型的代码删去。

修改过后的download_pretrained_checkpoints.sh如下：

#!/bin/bash

echo "downloading pretrained model checkpoints..."
mkdir pretrained_checkpoints
cd pretrained_checkpoints
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-pytorch_model.bin
cd ..

echo "script complete!"

在jupyter中执行以下语句，则可正常下载模型：

# 下载gpt2 small
!bash download_pretrained_checkpoints.sh

预处理输入数据的脚本在create_datasets.sh下，在这一步中，我们会通过词表，将原始的文本输入转变为token，产生最终的输入文件train.jsonl和valid.jsonl。很多朋友在跑LoRA代码的时候，会发现产生train.jsonl和valid.jsonl文件不存在这样的报错，就是因为没有执行这一步转换。具体的执行代码如下：

# 预处理数据：文字->token_id
!bash create_datasets.sh

10、恭喜你！到这一步为止，所有的准备工作就完成啦！接下来只需要在jupyter中执行以下代码块，我们就能把LoRA跑起来了：

!python -m torch.distributed.launch --nproc_per_node=1 --use_env src/gpt2_ft.py \
    --train_data ./data/e2e/train.jsonl \
    --valid_data ./data/e2e/valid.jsonl \
    --train_batch_size 8 \
    --grad_acc 1 \
    --valid_batch_size 4 \
    --seq_len 512 \
    --model_card gpt2.sm \
    --init_checkpoint ./pretrained_checkpoints/gpt2-pytorch_model.bin \
    --platform local \
    --clip 0.0 \
    --lr 0.0002 \
    --weight_decay 0.01 \
    --correct_bias \
    --adam_beta2 0.999 \
    --scheduler linear \
    --warmup_step 500 \
    --max_epoch 5 \
    --save_interval 1000 \
    --lora_dim 4 \
    --lora_alpha 32 \
    --lora_dropout 0.1 \
    --label_smooth 0.1 \
    --work_dir ./trained_models/GPT2_S/e2e \
    --random_seed 110

这样，我们就可以顺滑地跑起来啦，中间结果如下，大家完全可以根据自己的需要，在理解源码的基础上修改输出日志，帮助自己更好地学习代码：

二、代码快速上手

2.1 代码整体理解

在进入源码细节解读前，我们先来看看原版的LoRA代码要怎么用。

首先，整个LoRA代码的核心在loralib这个库下，该库中包含了为模型不同层（attention, mlp, embedding等）添加lora低秩适配器、训练和推理的方法。examples目录下则提供了使用loralib对不同模型、不同数据集做lora微调的代码例子，和论文中的实验对应。

也就是说，如果HuggingFace Peft没有对LoRA做过封装，那么大家用的最多的就是这个loralib。所以，我们先来看loralib的使用方法。

2.2 loralib使用方法

2.2.1 安装

这个很简单，在第一部分的环境配置里我们也讲过

pip install loralib
# Alternatively
# pip install git+https://github.com/microsoft/LoRA

2.2.2 添加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])

2.2.3 使用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)

2.2.4 保存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)

2.2.5 读取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)

2.2.6 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

3.1 入口函数

首先来看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的重点之一。

所以接下来，让我们详细来看这两个部分。

3.2 重写预训练模型（添加低秩适配器）

相关代码在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核心代码处理逻辑，在本文的第四部分详细阐述。在这里大家只要知道如何通过它来改写原始预训练模型即可。

另外，对r值和alpha值有疑问的朋友，可以参考之前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"

3.3 训练和验证

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 # 求均值

3.3.1 梯度累积(grad_acc)更新

细节都写在代码注释中了。额外需要提的一点，是代码中关于_lm_loss/(args.grad_acc)这一步。

在模型训练的过程中，我们有时不会一个step更新一次梯度，而是累积若干次step（即代码中的grac_acc值）再做一次梯度更新，这样做的原因是：

减少模型的更新频率，一定程度上加速训练速度
节省显存。可以理解为，当一个大batch将显存打爆时，我把它拆成若干个小batch来跑，此时我们理所当然希望这个大batch和这若个干小batch对梯度计算的结果尽量是一样的。

我们从第二点出发，当我们采用一个大batch计算梯度时，计算结果为：，其中loss_mean表示单条样本的平均loss。

当我们采用若干个小batch计算梯度时，每次step我们喂一个小batch，做一次梯度计算(loss.backward())，计算结果为。但是对于pytorch来说，我们知道梯度是挂在tensor.grad这个属性下的。如果你做完**loss.backward()，但是不做optimizer.zero_grad()，那么权重的tensor.grad这个属性就会累加**。这样产生的结果是，当你执行完这若干次step后，你会发现tensor.grad的值变成了。

此时，梯度被放大了grad_accr倍，因此我们需要处以grad_accr，来调回正常梯度值。

四、低秩适配器运作细节

现在，我们以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

4.1 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。

欸，经过这么一顺，是不是就不难理解了？

那现在，我们再来看，按照我们的意愿把self.merged_weights设置为False的时候会发生什么。

在这种情况下，我们希望不管是训练还是推理，预训练和低秩适配器总是分开的。这时，在model.train()模式下，我们不会触发if self.merge_weights and self.merged判断条件；在model.eval()模式下，我们不会触发if self.merge_weights and not self.merged判断条件。也就是我们啥也不用做，就按照初始化设置的那样，从头到尾都把两个权重拆开。