► 代码示例 / 快速 Keras 教程 / Gemma 的 LoRA 和 QLoRA 参数高效微调

Gemma 的 LoRA 和 QLoRA 参数高效微调

作者： Hongyu Chiu, Abheesht Sharma, Matthew Watson
创建日期 2024/08/06
最后修改日期 2024/08/06
描述： 使用 KerasHub 微调 LoRA 和 QLoRA 的 Gemma LLM。

ⓘ 本示例使用 Keras 3

在 Colab 中查看 • GitHub 源代码

简介

大型语言模型 (LLM) 已被证明在各种 NLP 任务中非常有效。LLM 首先在大规模文本语料库上进行自监督预训练。预训练有助于 LLM 学习通用知识，例如单词之间的统计关系。然后，LLM 可以在感兴趣的下游任务（例如情感分析）上进行微调。

然而，LLM 的规模极其庞大，在微调时我们不需要训练模型中的所有参数，尤其是在微调模型的数据集相对较小的情况下。换句话说，LLM 在微调时存在参数过多的问题。这就是低秩适应 (LoRA) 的用武之地；它显著减少了可训练参数的数量。这在保持输出质量的同时，减少了训练时间和 GPU 内存使用量。

此外，量化低秩适应 (QLoRA) 扩展了 LoRA，通过量化技术提高效率，而不会降低性能。

在此示例中，我们将使用 LoRA 和 QLoRA 微调 KerasHub 的 Gemma 模型以进行下一个 token 预测任务。

请注意，此示例在 Keras 支持的所有后端上运行。TensorFlow 仅用于数据预处理。

设置

在我们开始实现流水线之前，让我们安装并导入所有需要的库。我们将使用 KerasHub 库。

其次，让我们将精度设置为 bfloat16。这将有助于我们减少内存使用和训练时间。

另外，请确保已正确配置 KAGGLE_USERNAME 和 KAGGLE_KEY 以访问 Gemma 模型。

# We might need the latest code from Keras and KerasHub
!pip install -q git+https://github.com/keras-team/keras.git git+https://github.com/keras-team/keras-hub.git

import gc
import os

os.environ["KERAS_BACKEND"] = "jax"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"  # Suppress verbose logging from TF

# os.environ["KAGGLE_USERNAME"] = "..."
# os.environ["KAGGLE_KEY"] = "..."

import keras
import keras_hub
import tensorflow as tf
import tensorflow_datasets as tfds

keras.config.set_dtype_policy("bfloat16")

数据集

我们将使用 MTNT (Machine Translation of Noisy Text) 数据集，该数据集可从 TensorFlow Datasets 获取。在此示例中，我们将使用数据集的法语到英语部分。

train_ds = tfds.load("mtnt/fr-en", split="train")

我们可以打印一些样本。数据集中的每个样本包含两个条目

src：原始法语句子。
dst：相应的英语翻译。

examples = train_ds.take(3)
examples = examples.as_numpy_iterator()

for idx, example in enumerate(examples):
    print(f"Example {idx}:")
    for key, val in example.items():
        print(f"{key}: {val}")
    print()

Example 0:
dst: b'Yep, serious...'
src: b"Le journal l'est peut-\xc3\xaatre, mais m\xc3\xaame moi qui suit de droite je les trouve limite de temps en temps..."

Example 1:
dst: b'Finally, I explained to you in what context this copy-pasting is relevant: when we are told padamalgame etc.'
src: b"Enfin je t'ai expliqu\xc3\xa9 dans quel cadre ce copypasta est pertinent : quand on nous dit padamalgame etc."

Example 2:
dst: b'Gift of Ubiquity: Fran\xc3\xa7ois Baroin is now advisor to the Barclays Bank, mayor, president of the agglomeration, professor at HEC Paris, president of the Association of Mayors of France and Advocate Counselor, it must take him half a day each month.'
src: b"Don d'Ubiquit\xc3\xa9 : Fran\xc3\xa7ois Baroin est d\xc3\xa9sormais conseiller \xc3\xa0 la Banque Barclays, maire, pr\xc3\xa9sident d'agglom\xc3\xa9ration, professeur \xc3\xa0 HEC Paris, pr\xc3\xa9sident de l'association des maires de France et avocat  Conseiller, \xc3\xa7a doit lui prendre une demi journ\xc3\xa9e par mois."

由于我们将微调模型以执行法语到英语的翻译任务，因此我们应该为指令微调格式化输入。例如，我们可以将此示例中的翻译任务格式化为

<start_of_turn>user
Translate French into English:
{src}<end_of_turn>
<start_of_turn>model
{dst}<end_of_turn>

<start_of_turn>user、<start_of_turn>model 和 <end_of_turn> 等特殊 token 用于 Gemma 模型。您可以在 https://ai.google.dev/gemma/docs/formatting 了解更多信息

train_ds = train_ds.map(
    lambda x: tf.strings.join(
        [
            "<start_of_turn>user\n",
            "Translate French into English:\n",
            x["src"],
            "<end_of_turn>\n",
            "<start_of_turn>model\n",
            "Translation:\n",
            x["dst"],
            "<end_of_turn>",
        ]
    )
)
examples = train_ds.take(3)
examples = examples.as_numpy_iterator()

for idx, example in enumerate(examples):
    print(f"Example {idx}:")
    print(example)
    print()

Example 0:
b"<start_of_turn>user\nTranslate French into English:\nLe journal l'est peut-\xc3\xaatre, mais m\xc3\xaame moi qui suit de droite je les trouve limite de temps en temps...<end_of_turn>\n<start_of_turn>model\nTranslation:\nYep, serious...<end_of_turn>"

Example 1:
b"<start_of_turn>user\nTranslate French into English:\nEnfin je t'ai expliqu\xc3\xa9 dans quel cadre ce copypasta est pertinent : quand on nous dit padamalgame etc.<end_of_turn>\n<start_of_turn>model\nTranslation:\nFinally, I explained to you in what context this copy-pasting is relevant: when we are told padamalgame etc.<end_of_turn>"

Example 2:
b"<start_of_turn>user\nTranslate French into English:\nDon d'Ubiquit\xc3\xa9 : Fran\xc3\xa7ois Baroin est d\xc3\xa9sormais conseiller \xc3\xa0 la Banque Barclays, maire, pr\xc3\xa9sident d'agglom\xc3\xa9ration, professeur \xc3\xa0 HEC Paris, pr\xc3\xa9sident de l'association des maires de France et avocat  Conseiller, \xc3\xa7a doit lui prendre une demi journ\xc3\xa9e par mois.<end_of_turn>\n<start_of_turn>model\nTranslation:\nGift of Ubiquity: Fran\xc3\xa7ois Baroin is now advisor to the Barclays Bank, mayor, president of the agglomeration, professor at HEC Paris, president of the Association of Mayors of France and Advocate Counselor, it must take him half a day each month.<end_of_turn>"

为了本示例的目的，我们将使用数据集的一个子集。

train_ds = train_ds.batch(1).take(100)

模型

KerasHub 提供了许多流行模型架构的实现。在此示例中，我们将使用 GemmaCausalLM，这是一个用于因果语言建模的端到端 Gemma 模型。因果语言模型根据先前的 token 预测下一个 token。

请注意，sequence_length 设置为 256 以加快拟合速度。

preprocessor = keras_hub.models.GemmaCausalLMPreprocessor.from_preset(
    "gemma_1.1_instruct_2b_en", sequence_length=256
)
gemma_lm = keras_hub.models.GemmaCausalLM.from_preset(
    "gemma_1.1_instruct_2b_en", preprocessor=preprocessor
)
gemma_lm.summary()

Preprocessor: "gemma_causal_lm_preprocessor"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Tokenizer (type)                                   ┃                                             Vocab # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ gemma_tokenizer (GemmaTokenizer)                   │                                             256,000 │
└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘

Model: "gemma_causal_lm"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                  ┃ Output Shape              ┃         Param # ┃ Connected to               ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ padding_mask (InputLayer)     │ (None, None)              │               0 │ -                          │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ token_ids (InputLayer)        │ (None, None)              │               0 │ -                          │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ gemma_backbone                │ (None, None, 2048)        │   2,506,172,416 │ padding_mask[0][0],        │
│ (GemmaBackbone)               │                           │                 │ token_ids[0][0]            │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ token_embedding               │ (None, None, 256000)      │     524,288,000 │ gemma_backbone[0][0]       │
│ (ReversibleEmbedding)         │                           │                 │                            │
└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘

 Total params: 2,506,172,416 (4.67 GB)

 Trainable params: 2,506,172,416 (4.67 GB)

 Non-trainable params: 0 (0.00 B)

LoRA 微调

LoRA 到底是什么？

低秩适应 (LoRA) 是一种用于 LLM 的参数高效微调技术。它冻结 LLM 的权重，并注入可训练的秩分解矩阵。让我们更清楚地理解这一点。

假设我们有一个 n x n 的预训练密集层（或权重矩阵）W0。我们初始化两个密集层 A 和 B，它们的形状分别为 n x rank 和 rank x n。rank 远小于 n。在论文中，1 到 4 之间的值被证明效果很好。

LoRA 方程

原始方程为 output = W0x + b0，其中 x 是输入，W0 和 b0 是原始密集层（已冻结）的权重矩阵和偏置项。LoRA 方程为：output = W0x + b0 + BAx，其中 A 和 B 是秩分解矩阵。

LoRA 的理念是，预训练语言模型的权重更新具有较低的“内在秩”，因为预训练语言模型存在参数过多的问题。即使将 W0 的更新限制为低秩分解矩阵，也可以复制完全微调的预测性能。

可训练参数数量

让我们做一些快速计算。假设 n 为 768，rank 为 4。W0 有 768 x 768 = 589,824 个参数，而 LoRA 层 A 和 B 加起来有 768 x 4 + 4 x 768 = 6,144 个参数。因此，对于密集层，我们从 589,824 个可训练参数减少到 6,144 个可训练参数！

为什么 LoRA 会减少内存占用？

尽管总参数数量会增加（因为我们添加了 LoRA 层），但内存占用会减少，因为可训练参数数量减少了。让我们深入探讨一下。

模型的内存使用可以分为四部分

模型内存：这是存储模型权重所需的内存。LoRA 的模型内存会比原始模型稍高。
前向传播内存：这主要取决于批次大小、序列长度等。为了公平比较，我们对两个模型保持不变。
反向传播内存：这是存储梯度的内存。请注意，梯度仅为可训练参数计算。
优化器内存：这是存储优化器状态所需的内存。例如，Adam 优化器会存储可训练参数的“一阶矩向量”和“二阶矩向量”。

由于 LoRA 的可训练参数数量大大减少，因此优化器内存以及存储 LoRA 梯度所需的内存比原始模型少得多。这是大部分内存节省发生的地方。

为什么 LoRA 如此受欢迎？

减少 GPU 内存使用；
训练更快；并且
没有额外的推理延迟。

使用 KerasHub 时，我们可以通过一行 API 启用 LoRA：enable_lora(rank=4)

从 gemma_lm.summary() 中，我们可以看到启用 LoRA 后，可训练参数的数量显著减少（从 25 亿减少到 130 万）。

gemma_lm.backbone.enable_lora(rank=4)
gemma_lm.summary()

Preprocessor: "gemma_causal_lm_preprocessor"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Tokenizer (type)                                   ┃                                             Vocab # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ gemma_tokenizer (GemmaTokenizer)                   │                                             256,000 │
└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘

Model: "gemma_causal_lm"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                  ┃ Output Shape              ┃         Param # ┃ Connected to               ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ padding_mask (InputLayer)     │ (None, None)              │               0 │ -                          │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ token_ids (InputLayer)        │ (None, None)              │               0 │ -                          │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ gemma_backbone                │ (None, None, 2048)        │   2,507,536,384 │ padding_mask[0][0],        │
│ (GemmaBackbone)               │                           │                 │ token_ids[0][0]            │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ token_embedding               │ (None, None, 256000)      │     524,288,000 │ gemma_backbone[0][0]       │
│ (ReversibleEmbedding)         │                           │                 │                            │
└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘

 Total params: 2,507,536,384 (4.67 GB)

 Trainable params: 1,363,968 (2.60 MB)

 Non-trainable params: 2,506,172,416 (4.67 GB)

让我们微调 LoRA 模型。

# To save memory, use the SGD optimizer instead of the usual AdamW optimizer.
# For this specific example, SGD is more than enough.
optimizer = keras.optimizers.SGD(learning_rate=1e-4)
gemma_lm.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=optimizer,
    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)
gemma_lm.fit(train_ds, epochs=1)

微调后，响应将遵循提示中提供的指令。

template = (
    "<start_of_turn>user\n"
    "Translate French into English:\n"
    "{inputs}"
    "<end_of_turn>\n"
    "<start_of_turn>model\n"
    "Translation:\n"
)
prompt = template.format(inputs="Bonjour, je m'appelle Morgane.")
outputs = gemma_lm.generate(prompt, max_length=256)
print("Translation:\n", outputs.replace(prompt, ""))

Translation:
 Hello, my name is Morgane.

释放内存。

del preprocessor
del gemma_lm
del optimizer
gc.collect()

QLoRA 微调

量化低秩适应 (QLoRA) 扩展了 LoRA，通过将模型权重从高精度数据类型（如 float32）量化到低精度数据类型（如 int8）来提高效率。这导致内存使用量减少和计算速度加快。保存的模型权重也更小。

请注意，此处实现的 QLoRA 版本比原始版本更简化。区别在于

不使用 4 位 NormalFloat 格式，因为没有后端支持它。
无双重量化。
无分页优化器。

要在 KerasHub 中启用 QLoRA，请按照以下步骤操作

实例化模型。
使用动态 int8 量化量化权重。
启用 LoRA。

步骤 2 和 3 通过一行 API 实现

quantize("int8")
enable_lora(...)

preprocessor = keras_hub.models.GemmaCausalLMPreprocessor.from_preset(
    "gemma_1.1_instruct_2b_en", sequence_length=256
)
gemma_lm = keras_hub.models.GemmaCausalLM.from_preset(
    "gemma_1.1_instruct_2b_en", preprocessor=preprocessor
)
gemma_lm.quantize("int8")
gemma_lm.backbone.enable_lora(rank=4)
gemma_lm.summary()

Preprocessor: "gemma_causal_lm_preprocessor_1"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Tokenizer (type)                                   ┃                                             Vocab # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ gemma_tokenizer (GemmaTokenizer)                   │                                             256,000 │
└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘

Model: "gemma_causal_lm_1"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                  ┃ Output Shape              ┃         Param # ┃ Connected to               ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ padding_mask (InputLayer)     │ (None, None)              │               0 │ -                          │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ token_ids (InputLayer)        │ (None, None)              │               0 │ -                          │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ gemma_backbone                │ (None, None, 2048)        │   2,508,502,016 │ padding_mask[0][0],        │
│ (GemmaBackbone)               │                           │                 │ token_ids[0][0]            │
├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤
│ token_embedding               │ (None, None, 256000)      │     524,544,000 │ gemma_backbone[0][0]       │
│ (ReversibleEmbedding)         │                           │                 │                            │
└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘

 Total params: 2,508,502,016 (2.34 GB)

 Trainable params: 1,363,968 (2.60 MB)

 Non-trainable params: 2,507,138,048 (2.34 GB)

让我们微调 QLoRA 模型。

如果您使用的是支持 int8 加速的设备，您应该会看到训练速度的提升。

optimizer = keras.optimizers.SGD(learning_rate=1e-4)
gemma_lm.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=optimizer,
    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)
gemma_lm.fit(train_ds, epochs=1)

使用 QLoRA 微调，您应该会得到类似的输出。

prompt = template.format(inputs="Bonjour, je m'appelle Morgane.")
outputs = gemma_lm.generate(prompt, max_length=256)
print("Translation:\n", outputs.replace(prompt, ""))

Translation:
 Hello, my name is Morgane.

我们都完成了！

请注意，出于演示目的，此示例在一个时期内对数据集的一个小子集进行微调，并且 LoRA 秩值较低。要获得经过微调的模型更好的响应，您可以尝试

增加微调数据集的大小。
训练更多步（时期）。
设置更高的 LoRA 秩。
修改超参数值，例如 learning_rate 和 weight_decay。