大模型微调原理与代码实战案例(四)：LoRA

随着ChatGPT的快速崛起，大型模型的时代正在发生革命性变化。但对于很多人而言，进行大型模型的预训练或全面微调似乎是遥不可及的。

不过随着多种高效参数微调技术的涌现，科研人员和普通开发者都有机会尝试微调这些庞大的模型了。

本文我将分享了大模型微调技术的原理及代码案例**，完整版代码，可在文末获取**。

LoRA 简述

LoRA（论文：LoRA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS），该方法的核心思想就是通过低秩分解来模拟参数的改变量，从而以极小的参数量来实现大模型的间接训练。

在涉及到矩阵相乘的模块，在原始的PLM旁边增加一个新的通路，通过前后两个矩阵A,B相乘，第一个矩阵A负责降维，第二个矩阵B负责升维，中间层维度为r，从而来模拟所谓的本征秩（intrinsic rank）。

可训练层维度和预训练模型层维度一致为d，先将维度d通过全连接层降维至r，再从r通过全连接层映射回d维度，其中，r<<d，r是矩阵的秩，这样矩阵计算就从d x d变为d x r + r x d，参数量减少很多。

在下游任务训练时，固定模型的其他参数，只优化新增的两个矩阵的权重参数，将PLM跟新增的通路两部分的结果加起来作为最终的结果（两边通路的输入跟输出维度是一致的），即h=Wx+BAx。第一个矩阵的A的权重参数会通过高斯函数初始化，而第二个矩阵的B的权重参数则会初始化为零矩阵，这样能保证训练开始时新增的通路BA=0从而对模型结果没有影响。

在推理时，将左右两部分的结果加到一起即可，h=Wx+BAx=(W+BA)x，所以只要将训练完成的矩阵乘积BA跟原本的权重矩阵W加到一起作为新权重参数替换原本PLM的W即可，对于推理来说，不会增加额外的计算资源。

LoRA 微调实战

为了不影响阅读体验，完整版代码在公众号：机器学习社区，回复：lora，即可获取，这里仅列出关键步骤。

第一步，引进必要的库，如：LoRA 配置类 LoraConfig。

from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, TaskType

第二步，创建 LoRA 微调方法对应的配置。

peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM, 
    inference_mode=False, 
    r=8, 
    lora_alpha=32, 
    lora_dropout=0.1
)

参数说明：

• task_type：指定任务类型。如：条件生成任务（SEQ_2_SEQ_LM），因果语言建模（CAUSAL_LM）等。

• inference_mode：是否在推理模式下使用Peft模型。

• r：LoRA低秩矩阵的维数。关于秩的选择，通常，使用4，8，16即可。

• lora_alpha：LoRA低秩矩阵的缩放系数，为一个常数超参，调整alpha与调整学习率类似。

• lora_dropout：LoRA 层的丢弃（dropout）率，取值范围为[0, 1)。

• target_modules：要替换为 LoRA 的模块名称列表或模块名称的正则表达式。针对不同类型的模型，模块名称不一样，因此，我们需要根据具体的模型进行设置，比如，LLaMa的默认模块名为[q_proj, v_proj]，我们也可以自行指定为：[q_proj,k_proj,v_proj,o_proj]。在 PEFT 中支持的模型默认的模块名如下所示：

TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
    "t5": ["q", "v"],
    "mt5": ["q", "v"],
    "bart": ["q_proj", "v_proj"],
    "gpt2": ["c_attn"],
    "bloom": ["query_key_value"],
    "blip-2": ["q", "v", "q_proj", "v_proj"],
    "opt": ["q_proj", "v_proj"],
    "gptj": ["q_proj", "v_proj"],
    "gpt_neox": ["query_key_value"],
    "gpt_neo": ["q_proj", "v_proj"],
    "bert": ["query", "value"],
    "roberta": ["query", "value"],
    "xlm-roberta": ["query", "value"],
    "electra": ["query", "value"],
    "deberta-v2": ["query_proj", "value_proj"],
    "deberta": ["in_proj"],
    "layoutlm": ["query", "value"],
    "llama": ["q_proj", "v_proj"],
    "chatglm": ["query_key_value"],
    "gpt_bigcode": ["c_attn"],
    "mpt": ["Wqkv"],
}

Transformer的权重矩阵包括Attention模块里用于计算query, key, value的Wq，Wk，Wv以及多头attention的Wo和MLP层的权重矩阵，LoRA只应用于Attention模块中的4种权重矩阵，并且通过消融实验发现同时调整 Wq 和 Wv 会产生最佳结果，因此，默认的模块名基本都为 Wq 和 Wv 权重矩阵。

第三步，通过调用 get_peft_model 方法包装基础的 Transformer 模型。

model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

通过 print_trainable_parameters 方法可以查看到 LoRA 可训练参数的数量(仅为786,432)以及占比（仅为0.1404%）。

trainable params: 786,432 || all params: 560,001,024 || trainable%: 0.14043402892063284

PEFT 中 LoRA 相关的代码主要基于微软开源的LoRA的代码，并进行修改使其支持 PyTorch FSDP。在 PEFT 中， LoRA 模型相关源码如下所示。

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])

        # transformers models have a .config attribute, whose presence is assumed later on
        ifnot hasattr(self, "config"):
            self.config = {"model_type": "custom"}
...

LoRA 模型类结构如下所示：

PeftModelForCausalLM(
  (base_model): LoraModel(
    (model): BloomForCausalLM(
      (transformer): BloomModel(
        (word_embeddings): Embedding(250880, 1024)
        (word_embeddings_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        (h): ModuleList(
          (0): BloomBlock(
            (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
            (self_attention): BloomAttention(
              (query_key_value): Linear(
                in_features=1024, out_features=3072, bias=True
                (lora_dropout): ModuleDict(
                  (default): Dropout(p=0.1, inplace=False)
                )
                (lora_A): ModuleDict(
                  (default): Linear(in_features=1024, out_features=8, bias=False)
                )
                (lora_B): ModuleDict(
                  (default): Linear(in_features=8, out_features=3072, bias=False)
                )
                (lora_embedding_A): ParameterDict()
                (lora_embedding_B): ParameterDict()
              )
              (dense): Linear(in_features=1024, out_features=1024, bias=True)
              (attention_dropout): Dropout(p=0.0, inplace=False)
            )
            (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
            (mlp): BloomMLP(
              (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)
              (gelu_impl): BloomGelu()
              (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)
            )
          )
          ...
          (23): BloomBlock(
            ...
          )
        )
        (ln_f): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
      )
      (lm_head): Linear(in_features=1024, out_features=250880, bias=False)
    )
  )
)

第四步，模型训练的其余部分均无需更改，当模型训练完成之后，保存高效微调部分的模型权重以供模型推理即可。

peft_model_id = f"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}"
model.save_pretrained(peft_model_id)

输出的模型权重文件如下所示：

/data/nfs/llm/model/bloomz-560m_LORA_CAUSAL_LM
├── [ 447]  adapter_config.json
├── [3.0M]  adapter_model.bin
└── [ 147]  README.md

0 directories, 3 files

注意：这里只会保存经过训练的增量 PEFT 权重。其中，adapter_config.json 为 LoRA 配置文件；adapter_model.bin 为 LoRA 权重文件。

第五步，加载微调后的权重文件进行推理。

from peft import PeftModel, PeftConfig

peft_model_id = f"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}"
config = PeftConfig.from_pretrained(peft_model_id)
# 加载基础模型
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
# 加载PEFT模型
model = PeftModel.from_pretrained(model, peft_model_id)

# tokenizer编码
inputs = tokenizer(f'{text_column} : {dataset["test"][i]["Tweet text"]} Label : ', return_tensors="pt")

# 模型推理
outputs = model.generate(
        input_ids=inputs["input_ids"], 
        attention_mask=inputs["attention_mask"], 
        max_new_tokens=10, 
        eos_token_id=3
    )

# tokenizer解码
print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))

至此，我们完成了 LoRA 的训练及推理。

结语

本文对 LoRA 基本原理进行了简述；同时，讲解了 LoRA 进行模型训练及推理。下文将对 IA3 技术进行实战。

如果觉得我的文章能够能够给您带来帮助，期待您的点赞收藏加关注~~

技术交流群

完整版代码在公众号：机器学习社区，回复：lora，即可获取。

建了实战技术交流群！想要进交流群的同学，可以直接加微信号：mlc2060。加的时候备注一下：研究方向 +学校/公司+知乎，即可。然后就可以拉你进群了。

前沿技术资讯、算法交流、求职内推、算法竞赛、面试交流(校招、社招、实习)等、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企开发者互动交流~