【AI小项目5】使用 KerasNLP 对 Gemma 模型进行 LoRA 微调

一、项目简介

概述

本项目使用 KerasNLP 对 Gemma 模型进行 LoRA 微调，为了快速验证微调的效果，仅在数据集的一个小子集上进行了一个周期的微调，并使用了较低的 LoRA 秩值。取得了显著的改进效果。

时间

2024.09.15-2024.09.16

主要工作和收获

• 掌握了用 LoRA 技术改进微调的原理和基本使用
• 掌握了用 KerasNLP 微调 Gemma 模型的基本操作

技术栈

KerasNLP，Gemma，LoRA

数据集

databricks-dolly-15k

结果

使用LoRA微调后，在两个示例上结果有明显改进

参考

Fine-tune Gemma models in Keras using LoRA
Practical Tips for Finetuning LLMs Using LoRA

二、完整代码

1.kaggle上有我的完整代码，这里贴一份是为了方便那些上不了kaggle的读者。
2.代码有不懂的地方，可以看下背景知识补充部分

kaggle地址：使用 KerasNLP 对 Gemma 模型进行 LoRA 微调

概览

Gemma 是一系列轻量级、先进的开源模型，基于与 Gemini 模型相同的研究和技术构建。

大型语言模型（LLM），如 Gemma，已被证明在各种 NLP 任务中表现出色。LLM 首先在一个大型文本语料库上进行自监督预训练。预训练帮助 LLM 学习通用知识，如单词之间的统计关系。然后，LLM 可以使用领域特定的数据进行微调，以执行下游任务（例如情感分析）。

LLM 的规模非常大（参数数量达到数百万级）。对于大多数应用来说，不需要对模型的所有参数进行完全微调，因为典型的微调数据集相对于预训练数据集来说要小得多。

Low Rank Adaptation (LoRA) 是一种微调技术，通过冻结模型权重并在模型中插入少量新权重，大大减少了下游任务中的可训练参数数量。这使得使用 LoRA 进行训练更加快速和高效，同时生成更小的模型权重（几百 MB），并保持模型输出的质量。

本项目将引导你使用 KerasNLP 在 Gemma 2B 模型上进行 LoRA 微调，使用 Databricks Dolly 15k 数据集。该数据集包含 15,000 对高质量的人工生成的提示/响应对，专门用于微调 LLM。

设置

安装依赖

安装 Keras、KerasNLP 和其他依赖项。

# 安装最新的 Keras 3 . 这里有更多的细节： https://keras.io/getting_started/
!pip install -q -U keras-nlp
!pip install -q -U keras>=3

选择一个后端

Keras 是一个高层、多框架的深度学习 API，旨在简化使用和提高易用性。使用 Keras 3，你可以在以下三种后端之一上运行工作流程：TensorFlow、JAX 或 PyTorch。

在本项目中，配置 JAX 作为后端。

import os

os.environ["KERAS_BACKEND"] = "jax"  # Or "torch" or "tensorflow".

# 避免在 JAX 后端发生内存碎片化。
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"

导入包

import keras
import keras_nlp

加载数据集

预处理数据。本项目使用了1000个训练示例的子集，以加快笔记本的执行速度。为了获得更高质量的微调效果，建议使用更多的训练数据。

从一个JSON Lines格式的文件中读取数据，过滤掉某些样本并将其格式化为特定的字符串模板，然后截取前1000个样本。

import json
data = []
with open('/kaggle/input/databricks-dolly-15k/databricks-dolly-15k.jsonl') as file:
    for line in file:
        features = json.loads(line)
        
        # 为了保持数据的简化，只处理没有`context`的样本。
        if features["context"]:
            continue
            
        # 格式化每个样本为模板字符串
        template = "Instruction:\n{instruction}\n\nResponse:\n{response}"
        data.append(template.format(**features))

# 只使用前1000个训练样本，为了处理速度较快
data = data[:1000]

加载模型

KerasNLP 提供了许多流行模型架构的实现。在本项目中，你将使用 GemmaCausalLM 创建一个模型，这是一个用于因果语言建模的端到端 Gemma 模型。因果语言模型根据前面的标记预测下一个标记。

# 从Keras NLP的预设模型中导入，这意味着模型的权重和架构都已经预先定义好，便于直接使用。
gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma_2b_en")

# 查看模型结构，包括层数、参数数量以及每层的形状和类型。
gemma_lm.summary()

微调前的推理

在本节中，你将使用各种提示词对模型进行查询，看看它的响应情况。

欧洲旅行例子

向模型查询关于欧洲旅行的建议。

prompt = template.format(
    instruction="What should I do on a trip to Europe?",
    response="",
)
print(gemma_lm.generate(prompt, max_length=256))

打印结果如下：

模型只是简单地回应了一个去欧洲旅行的建议。

光合作用例子

提示模型用足够简单的语言解释光合作用，使5岁的孩子能够理解。

prompt = template.format(
    instruction="Explain the process of photosynthesis in a way that a child could understand.",
    response="",
)
print(gemma_lm.generate(prompt, max_length=256))

打印结果如下：

回应中包含一些可能对孩子来说不容易理解的词汇，如叶绿素（chlorophyll）、葡萄糖（glucose）等。

LoRA 微调

为了从模型中获得更好的回应，使用 Databricks Dolly 15k 数据集通过低秩适应 (LoRA) 对模型进行微调。

LoRA 的秩决定了添加到 LLM 原始权重中的可训练矩阵的维度。它控制了微调调整的表现力和精确度。

较高的秩意味着可以进行更详细的变化，但也意味着更多的可训练参数。较低的秩意味着计算开销较小，但可能会降低适应的精确度。

本项目使用 LoRA 秩为 4。实际操作中，可以从相对较小的秩（例如 4、8、16）开始。这对实验来说计算效率较高。用该秩训练模型，并评估在任务中的性能提升。

# 启用模型的 LoRA 并将 LoRA 秩设置为 4。
gemma_lm.backbone.enable_lora(rank=4)
gemma_lm.summary()

打印结果如下：（省略了部分信息）

请注意，启用 LoRA 后显著减少了可训练参数的数量（从 25 亿减少到 130 万）。

# 限制输入序列长度为512 (控制内存使用).
# Transformer模型的计算复杂度通常与序列长度成平方关系，因此控制序列长度是常见的优化手段。
gemma_lm.preprocessor.sequence_length = 512

# AdamW是Transformer模型中常用的优化器。相比经典的Adam优化器，AdamW增加了权重衰减（weight decay），它是一种常用的正则化方法，有助于减少模型的过拟合。
optimizer = keras.optimizers.AdamW(
    learning_rate=5e-5, # 这是Transformer类模型的常用学习率。
    weight_decay=0.01,
)

# 排除特定参数的权重衰减
# bias和scale（LayerNorm中的缩放参数）通常不参与权重衰减，因为它们在训练中的变化不需要像其他权重那样受到约束。
optimizer.exclude_from_weight_decay(var_names=["bias", "scale"])

# compile函数用于配置模型的训练方式
# 使用稀疏的分类交叉熵作为损失函数，适合分类任务。from_logits=True意味着输出没有经过softmax归一化，因此损失函数内部会处理这个问题。
# SparseCategoricalAccuracy：用于衡量分类的准确度，特别适用于稀疏标签（即目标标签是一个整数而不是one-hot向量）。
gemma_lm.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=optimizer,
    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)

# 开始训练
gemma_lm.fit(data, epochs=1, batch_size=1)

微调后的推理

经过微调后，模型的响应能够遵循提示中提供的指令。

欧洲旅行例子

prompt = template.format(
    instruction="What should I do on a trip to Europe?",
    response="",
)
print(gemma_lm.generate(prompt, max_length=256))

打印结果如下：

该模型现在推荐欧洲的旅游景点。

光合作用例子

prompt = template.format(
    instruction="Explain the process of photosynthesis in a way that a child could understand.",
    response="",
)
print(gemma_lm.generate(prompt, max_length=256))

打印结果如下：

模型现在用更简单的词汇来解释光合作用

改进方向

请注意，出于快速学习目的，本项目仅在数据集的一个小子集上进行了一个周期的微调（用GPU P100 跑一次花10分钟左右），并使用了较低的 LoRA 秩值（4）。为了从微调后的模型中获得更好的响应，你可以尝试以下方法：

1. 增加微调数据集的大小（从1个改成更多个）
2. 使用数据增强技术来扩充原数据集
3. 增加训练步骤（将1个 epoch改成多个）
4. 设置更高的 LoRA 秩
5. 为更多层启用 LoRA
6. 调整超参数值，例如 learning_rate 和 weight_decay。
7. 使用更大的模型比如 Gemma 7B 模型
8. 尝试使用其他优化器

三、背景知识补充

Fine-tune（微调）

Fine-tune（微调）是机器学习领域中的一种技术，特别在深度学习模型中较为常见。它的核心思想是在预训练的基础上，通过针对特定任务或数据集进行进一步训练，来提升模型在特定任务上的性能。

举个简单的例子：

假设你有一个已经训练好的NLP模型，它可以很好地理解英文语言。如果你想让它能执行特定任务，比如情感分析（判断电影评论是正面还是负面），你只需要将这个模型在你的电影评论数据集上再训练几轮（微调），让它能针对这个任务更准确地进行预测。

为什么大模型需要微调？

大模型需要微调的原因可以归结为以下几点：

1. 预训练任务和目标任务的差异：微调使模型更好地适应特定任务需求。
2. 领域特征的差异：微调让模型捕捉特定领域的特征。
3. 提升精度：微调帮助模型在特定任务上达到更高的精度。
4. 节省资源：微调避免了从头训练模型的高昂成本。
5. 小数据集优化：微调能让模型在少量数据上获得良好效果。

通过微调，我们可以充分发挥大模型的能力，使其在不同领域、任务上都能有最佳表现。

为什么不训练一个大模型精通所有特定任务？

直接训练一个大模型精通所有特定任务在现实中并不可行，主要原因有：

1. 计算资源和时间的高昂成本：训练和运行一个通用的、能够处理所有任务的大模型需要过多的计算资源。
2. 泛化与特化的矛盾：模型的泛化能力与针对特定任务的能力往往存在冲突，难以同时优化。
3. 任务间的干扰和灾难性遗忘：多个任务之间会产生冲突，使得模型无法在每个任务上都保持高水平表现。
4. 数据集的多样性和规模限制：特定任务的数据往往规模较小且特点鲜明，无法通过统一的训练来处理。
5. 效率问题：直接使用一个万能的大模型处理所有任务会导致资源浪费和运行效率低下。

因此，微调预训练模型可以在满足特定任务需求的同时，最大化利用资源和提升效率。

OpenAI最新发布的o1模型是否能达成这个目标？

OpenAI 最新发布的 o1 模型系列（包括 o1-preview 和 o1-mini）确实在提升复杂推理能力方面取得了显著进展，尤其是在数学、编程和多步骤工作流的处理上表现优异。它采用了新的推理方法，如“链式推理”（Chain of Thought），能够更好地处理需要多步骤思考和解决的复杂问题。

o1 模型的优势：

1. 推理能力提升：相比之前的模型，o1 在复杂推理任务中表现突出，能够提供更准确的逻辑推导和数学证明。例如，在国际数学奥林匹克题目上的正确率大幅提高，达到了 83%。
2. 代码生成和调试：o1-mini 专为编程任务设计，能够高效生成复杂代码并进行调试，尤其适合需要详细步骤解释的场景。
3. 推理优化：o1 通过“链式推理”方法，让模型能够更加系统性地解决问题，尤其在需要多步骤决策的任务中表现出色。

限制和挑战：

尽管 o1 在推理任务上表现出色，但它并不能完全解决所有特定任务中的需求差异。例如：

• 大规模事实知识：在需要大量事实知识的任务中（如网络浏览或大规模文件处理），o1 的表现仍有限，需要依赖其他模型（如 GPT-4o）。
• 缺乏高级功能：o1-preview 和 o1-mini 目前不支持诸如文件上传、数据分析、记忆等高级功能。

结论：

尽管 o1 模型在推理和特定任务上具有很强的能力，但它尚未达到能够精通所有特定任务的水平，尤其在需要处理大量知识或高级功能的任务中。因此，o1 模型更适合专门用于复杂推理、编程和科学任务，而不能完全取代在所有任务上的微调需求。

微调的基本步骤

微调（Fine-tuning）的基本步骤可以总结为以下几个关键阶段：

1. 选择预训练模型

首先，选择一个已经在大规模数据集上进行预训练的模型。这个模型通常是在通用任务上训练的，如BERT、GPT、RoBERTa等。预训练模型已经学习了通用的特征和模式，是微调的基础。

2. 准备数据集

微调需要针对特定任务的数据集。这个数据集应该与最终任务相关，可能是分类、回归、文本生成等任务的数据。数据需要进行清理和预处理，通常包括去除噪音、分词、去除停用词等步骤。

3. 适配任务和模型架构

根据任务的需求调整模型架构。比如，对于分类任务，可能需要在预训练模型的输出层添加一个分类器层，输出类别标签。对于文本生成任务，输出层可能是一个序列生成层。

4. 设置超参数

微调过程中，需要设置一些重要的超参数，如学习率、batch size、训练轮数（epoch）等。微调时通常会使用比预训练更小的学习率，以避免模型权重发生过大变化，影响已有的预训练效果。

5. 训练过程

在任务相关的数据集上，使用微调设置对模型进行训练。模型会基于原有的预训练权重，进一步调整权重以适应新的任务。在训练过程中，使用反向传播算法来更新模型的参数。

6. 评估和调整

在训练过程中，使用验证集来评估模型的性能。通过观察验证集上的损失值和准确率，可以决定是否需要调整模型架构或超参数。如果验证集的表现不理想，可能需要进一步调整数据预处理、模型结构或优化方法。

7. 部署和推理

训练完成后，将微调后的模型应用到实际的推理任务中。例如，可以用于分类新数据、生成文本或回答问题。通常会导出模型并将其部署在生产环境中供用户使用。

常见工具与库：

• Hugging Face Transformers：提供了丰富的预训练模型和微调工具，支持各种NLP任务。
• TensorFlow/Keras 和 PyTorch：用于自定义模型微调的深度学习框架。

通过这些步骤，微调可以让预训练模型在特定任务上表现得更加精准和高效。

如何选择预训练模型？

选择预训练模型时，需要根据具体任务、数据特征以及模型的性能要求来进行评估。以下是几个重要的考虑因素：

1. 任务类型

预训练模型应与目标任务类型相匹配。不同的预训练模型擅长于不同类型的任务：

• 文本分类任务：可以选择像 BERT 或 RoBERTa 这类擅长分类和文本表示的模型，它们通过掩码语言模型（MLM）任务进行预训练，适合于文本分类、情感分析等任务。
• 文本生成任务：如果任务是生成类任务（如文章生成、对话生成），可以选择 GPT 或 T5 这类基于自回归语言模型的预训练模型，它们在生成任务上表现优秀。
• 问答任务：对于问答任务，选择 BERT 或 T5 等预训练模型可能更合适，尤其是经过 SQuAD 数据集微调的模型可以在提取式问答任务中有优异表现。

2. 数据的领域

根据数据的领域选择适合的预训练模型：

• 通用领域数据：如果数据是通用文本（如新闻、社交媒体评论等），可以选择像 GPT-4、BERT 这类已经在大规模通用语料上预训练的模型。
• 专业领域数据：如果数据来自某个专业领域（如医学、法律），建议选择经过该领域特定数据微调过的预训练模型，例如 BioBERT（医学领域）或 LegalBERT（法律领域）。

3. 模型的规模

模型规模（参数数量）直接影响其性能和资源需求：

• 小型模型：如果你的资源有限或任务需要低延迟，可以选择较小的预训练模型，如 DistilBERT 或 TinyBERT，它们是对原始大型模型的压缩版本，仍保留了较好的性能。
• 大型模型：如果任务需要高精度且有充足的计算资源，选择 GPT-4、BERT-Large 等大型模型可能更合适，它们能捕获更多的语言特征和细节，但计算开销较大。

4. 训练数据规模和计算资源

如果微调的数据集较小，选择较小的模型进行微调可能更合理。较大的预训练模型往往需要更多的数据和计算资源来避免过拟合。如果资源充足且任务较复杂，可以使用更大的预训练模型。

5. 模型的社区支持与成熟度

选择成熟的、经过广泛测试和优化的预训练模型，可以减少遇到的技术问题。像 Hugging Face 的 Transformers 库提供了多种预训练模型，并有丰富的社区支持与文档，有利于加快微调过程。

6. 模型的开源与商业化支持

考虑是否需要开源模型，还是需要商业化支持的模型。开源模型（如 BERT、GPT）通常在研究和非商业环境中使用较多，而商用支持的模型（如 OpenAI 的 GPT 系列、Anthropic 的 Claude）则适合商业应用。

7. 性能和准确率

最后，需要查看各个模型在类似任务上的性能表现。许多预训练模型已经在各种基准测试（如 GLUE、SQuAD）上评估过，参考这些数据可以帮助你选择在类似任务上表现较好的模型。

总结

在选择预训练模型时，主要考虑任务类型、数据领域、模型规模、资源限制以及模型的社区支持和成熟度。像 BERT、GPT、T5 这类模型在多种任务中表现出色，是常用的选择。你可以根据具体需求调整选择策略，以获得最佳性能和效率。

使用LoRA微调

Fine-tuning with LoRA（Low-Rank Adaptation） 是一种用于大型语言模型的高效微调方法。LoRA的核心思想是在不修改预训练模型的原始权重的基础上，通过插入低秩矩阵来进行适应性学习，从而降低模型参数更新的计算成本和存储需求。它特别适合用于资源有限的场景，同时能够保持较高的模型性能。

LoRA 的基本原理：

LoRA 的关键在于将模型的权重矩阵分解为两个低秩矩阵，通过这种方式减少模型参数的更新量。具体而言，它不会直接更新原始模型的参数，而是学习一个额外的低秩矩阵，将该矩阵与模型原有的权重矩阵相加。这使得微调时的参数数量显著减少，同时能够保留原始预训练模型的知识。

公式上，如果我们有原始的模型权重矩阵 $W$，LoRA 的操作如下：
$$W=W_0+\Delta W=W_0+A\times B$$
其中，$A$ 和 $B$ 是低秩矩阵，$\Delta W=A\times B$ 是要学习的部分，而 $W_0$ 是原始预训练的权重矩阵。

图解

图片来源：https://magazine.sebastianraschka.com/p/practical-tips-for-finetuning-llms

Fine-tuning with LoRA 的步骤：

1. 加载预训练模型：首先从 Hugging Face 等库中加载预训练模型（如 BERT、GPT）。
2. 插入 LoRA 层：在模型的某些关键层（如自注意力层）中插入 LoRA 层。LoRA 只对这些低秩矩阵进行微调，而不修改原始模型参数。
3. 训练 LoRA 模型：对低秩矩阵进行微调，学习任务相关的参数。由于需要更新的参数数量大大减少，训练速度更快，所需的显存和计算资源也大幅下降。
4. 推理或部署：微调完成后，LoRA 层与原始模型结合使用，整体效果类似于微调整个模型。

LoRA 的优势：

• 计算高效：与常规微调相比，LoRA 只更新少量参数（低秩矩阵），因此大幅减少了显存和计算资源的需求。
• 灵活性：LoRA 可以针对不同任务插入多个低秩矩阵，因此允许模型适应多任务学习，而不需要为每个任务单独训练一个完整模型。
• 保存预训练模型：LoRA 保持了原始模型的参数完好无损，微调后模型仅需保存新增的低秩矩阵，这对存储和模型部署非常友好。

常见应用：

• 自然语言处理任务：LoRA 主要用于语言模型的微调任务，如文本分类、问答系统和生成式任务。
• 生成式模型的微调：如 GPT、T5 等模型在进行文本生成或复杂任务的微调时，LoRA 提供了一个高效的替代方案，尤其是在大模型的多任务学习和迁移学习中效果显著。

工具支持：

• Hugging Face：LoRA 已集成到 Hugging Face 的 Transformers 库中，用户可以轻松应用 LoRA 进行模型微调。
• DeepSpeed：LoRA 也支持与 DeepSpeed 等高效训练框架结合使用，进一步提高大模型的微调效率。

通过 LoRA，可以在有限的资源条件下进行大模型的微调，同时保持较高的模型性能。这使得 LoRA 成为在大规模预训练模型上进行高效适应性学习的一种理想方法。

低秩矩阵

低秩矩阵（Low-Rank Matrix）是指其秩（rank）比其维度低的矩阵。矩阵的秩是指该矩阵中线性独立行或列的最大数目。低秩矩阵在矩阵分解、数据压缩、降维等方面具有重要应用，特别是在机器学习、信号处理、推荐系统等领域。

低秩矩阵的特点：

• 秩的定义：矩阵的秩是其行向量或列向量的线性独立性。如果一个矩阵 $A\in {\mathbb{R}}^{m\times n}$ 的秩为 $r$，则 $r\le \min (m,n)$。如果 $r<\min (m,n)$，则称其为低秩矩阵。
• 结构紧凑：低秩矩阵通过使用少量的独立向量可以表示出整个矩阵的结构，这意味着它具有更紧凑的表示形式，适合用于数据压缩。
• 近似矩阵：在许多应用中，可以通过低秩矩阵来近似原始矩阵，从而减少计算和存储成本。例如，矩阵分解中的 SVD（奇异值分解） 可用于找到最优的低秩近似矩阵。

低秩矩阵的应用：

1. 降维和数据压缩：在高维数据中，通过低秩矩阵分解（如PCA）可以将数据降维，保留数据的主要特征，减少计算复杂度。
2. 矩阵分解：低秩矩阵常用于矩阵分解技术，如在推荐系统中，利用低秩矩阵进行矩阵填充和预测用户偏好。
3. 深度学习中的模型微调：在 LoRA 等模型中，通过学习低秩矩阵来微调模型参数，以减少训练计算成本和模型存储开销。

低秩矩阵在 LoRA 中的作用：

LoRA 通过将模型的权重矩阵 $W$ 分解为两个低秩矩阵 $A$ 和 $B$，以较少的参数学习新的任务。这样不仅减少了存储开销，还加速了模型的微调过程，同时保留了模型原有的知识。

低秩矩阵的常用分解方法：

• SVD（Singular Value Decomposition，奇异值分解）：将矩阵分解为三个矩阵的乘积，低秩近似可以通过保留奇异值中最大的部分实现。
• PCA（Principal Component Analysis，主成分分析）：利用特征值分解，找到数据中最具代表性的成分，进行降维处理。

总结来说，低秩矩阵通过减少表示矩阵的独立向量数目，能够在数据压缩、降维和高效学习中起到关键作用，特别是在像 LoRA 这样的大模型微调场景中大大减少了计算资源的需求。