LLM - LoRA 模型合并与保存

目录

一.引言

二.LoRA 

1.LoRA 简介

2.LoRA 参数

3.LoRA 合并

4.LoRA 保存

三.总结

---

一.引言

LLM 使用过程中最常用方法之一就是通过 LoRA 基于自己的数据对大模型进行微调，本文简单介绍 LoRA 原理以及如何合并多个 LoRA 模型并保存。

peft==0.4.0
transformers==4.29.1

二.LoRA 

1.LoRA 简介

LoRA 是一种消耗较少内存同时加速大型模型微调的技术。

为了使得微调更有效，LoRA 的方法是通过低秩分解两个较小的矩阵 [称为更新矩阵] 用于参数更新。在维持较少更新参数的基础上，训练并适应新的微调数据。原始参数 W 保持冻结，不会进行更新训练。该方法有如下优点：

    ◆  LoRA 通过大幅减少可训练参数数量，使微调更加效率。

    ◆  基于 LoRA 的轻便型，可以构建多个轻量级 LoRA 模型用于不同下游任务 

    ◆  LoRA 与许多其他参数有效方法正交，并且可以与其结合。

    ◆  使用 LoRA 进行微调的模型的性能与完全微调模型的性能相当。

    ◆  LoRA 几乎不添加任何推理延迟，因为适配器权重可以与基本模型合并。

2.LoRA 参数

原则上，LoRA 可以应用于神经网络中权重矩阵的任何子集，以减少可训练参数的数量。然而，为了简化和进一步提高参数效率，在 Transformer 模型中，LoRA 通常仅应用于注意力块。LoRA 模型中可训练参数的结果数量取决于低秩更新矩阵的大小，其主要由秩 r 和原始权重矩阵的形状确定。实际使用过程中，通过选择不同的 lora_target 决定训练的参数量，以 LLama 为例：

--lora_target q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj \

可以通过 numel 方法获取张量中参数的数量和 requires_grad 方法获取模型张量是否进行梯度计算来计算可训练参数的比例：

# 打印可训练参数
def print_trainable_params(model: torch.nn.Module) -> None:
    trainable_params, all_param = 0, 0
    for param in model.parameters():
        num_params = param.numel()
        # if using DS Zero 3 and the weights are initialized empty
        if num_params == 0 and hasattr(param, "ds_numel"):
            num_params = param.ds_numel
        all_param += num_params
        if param.requires_grad:
            trainable_params += num_params
    print("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
        trainable_params, all_param, 100 * trainable_params / all_param))

3.LoRA 合并

在训练过程中，较小的权重矩阵 A、B 是分开的。但一旦训练完成，权重实际上可以合并到一个相同的新权重矩阵中。虽然 LoRA 明显更小，训练速度更快，但由于分别加载基本模型和 LoRA 模型，可能会在推理过程中遇到延迟问题。为了消除延迟，可以使用 merge_and_unload 函数将适配器权重与基本模型合并，这样可以有效地将新合并的模型用作独立模型。同时也可以合并多个 LoRA 模型，使得 Base Model 同时具备多个任务处理能力。

训练时原始参数 W 保持不动，更新矩阵 A/B。合并时 W-merged = W + BA 代替原有的 W。 

from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, GenerationConfig
from peft import PeftModel

    # 载入预训练模型
    tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=True, padding_side="left", **config_kwargs)
    print("Tokenizer Load Success!")
    config = AutoConfig.from_pretrained(base_model, **config_kwargs)
    # Load and prepare pretrained models (without valuehead).
    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        config=config,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
        revision='main'
    )
    print('origin config =', model.config)
    # 模型合并
    ckpt_list = ["checkpoint-1000", "checkpoint-2000", "checkpoint-3000"]
    for checkpoint in ckpt_list:
        print('Merge checkpoint: {}'.format(checkpoint))
        model = PeftModel.from_pretrained(model, os.path.join(lora_model, checkpoint))
        model = model.merge_and_unload()
    print('merge config =', model.config)

通过 merge_and_unlaod 方法可以合并多个 Lora 模型，这里博主尝试将同一个模型的 3 个  CKPT 合并至原始模型中。通过合并不同类型任务的 CKPT，原始模型可以同时具备多种下游任务的能力且推理效率不会受影响。

4.LoRA 保存

import torch
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizer


def apply_lora(model_name_or_path, output_path, lora_path):
    print(f"Loading the base model from {model_name_or_path}")
    base = AutoModelForCausalLM.from_pretrained(
        model_name_or_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
    )
    base_tokenizer = LlamaTokenizer.from_pretrained(model_name_or_path)

    print(f"Loading the LoRA adapter from {lora_path}")

    lora_model = PeftModel.from_pretrained(
        base,
        lora_path,
        torch_dtype=torch.float16,
    )

    print("Applying the LoRA")
    model = lora_model.merge_and_unload()

    print(f"Saving the target model to {output_path}")
    model.save_pretrained(output_path)
    base_tokenizer.save_pretrained(output_path)

使用 merge_and_unload 方法进行参数合并，再调用 save_pretrained 方法进行模型保存。可以看到保存前后模型大小几乎没有变化。 

 

三.总结

合并多个 LoRA 模型也可能对模型的原始能力造成影响，大家可以根据需求测试与尝试。

参考：

- LoRA https://huggingface.co/docs/peft/conceptual_guides/lora

- PEFT Models