使用 LoRA 和 QLoRA 对 Gemma 进行参数高效微调
作者： 邱鸿宇，Abheesht Sharma，Matthew Watson
创建日期 2024/08/06
上次修改日期 2024/08/06
描述：使用 KerasHub 使用 LoRA 和 QLoRA 微调 Gemma 大语言模型。
引言
大型语言模型 (LLM) 已被证明在各种 NLP 任务中非常有效。LLM 首先以自监督的方式在大型文本语料库上进行预训练。预训练帮助 LLM 学习通用知识，例如单词之间的统计关系。然后，可以对 LLM 在感兴趣的下游任务（例如情感分析）上进行微调。
然而，LLM 的规模非常庞大，在微调时我们不需要训练模型中的所有参数，特别是因为模型微调所使用的数据集相对较小。换句话说，LLM 在微调方面参数过多。这就是 低秩自适应 (LoRA) 的用武之地；它显著减少了可训练参数的数量。这导致训练时间和 GPU 内存使用量的减少，同时保持输出质量。
此外，量化低秩自适应 (QLoRA) 扩展了 LoRA，通过量化技术提高效率，而不会降低性能。
在本示例中，我们将使用 LoRA 和 QLoRA 在下一个 token 预测任务上微调 KerasHub 的 Gemma 模型。
请注意，此示例在 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（噪声文本的机器翻译）数据集，该数据集可从 TensorFlow 数据集中获得。在本示例中，我们将使用该数据集的法语到英语部分。
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.
我们完成了！
请注意，出于演示目的，此示例仅在一个 epoch 内对数据集的一个小子集进行模型微调，并且 LoRA 秩值较低。为了从微调后的模型中获得更好的响应，您可以尝试
- 增加微调数据集的大小。
- 训练更多步数（epoch）。
- 设置更高的 LoRA 秩。
- 修改超参数值，例如
learning_rate和weight_decay。