使用 LoRA 对 Gemma 2 进行微调，以纳入 Rust 官方文档

显示全部楼层

概述

Gemma 是一组轻量级的开放模型，基于用于创建 Gemini 模型的研究和技术构建而成。大语言模型 (LLM) 如 Gemma 在各种自然语言处理任务中表现出色。LLM 首先通过在大量文本语料库上进行自监督预训练来学习，预训练帮助 LLM 学习通用知识，例如词与词之间的统计关系。然后，可以使用特定领域的数据对 LLM 进行微调，以执行下游任务 (例如情感分析)。

LLM 的规模非常庞大 (参数量级为数百万)。对于大多数应用，完全微调 (更新模型中的所有参数) 并不必要，因为微调数据集的规模相对于预训练数据集来说要小得多。

低秩适配 (LoRA)是一种微调技术，它通过冻结模型的权重并将少量的新权重插入模型，大大减少了下游任务的可训练参数数量。这使得使用 LoRA 进行训练速度更快、内存效率更高，并生成较小的模型权重 (几百 MB)，同时保持模型输出的质量。

低秩适配 (LoRA)
https://arxiv.org/abs/2106.09685

本教程将指导您使用 KerasNLP 在 Gemma 2B 模型上执行 LoRA 微调，使用 Rust 官方文档 1.6k 数据集。该数据集包含 1,600 对用 Gemini Flash 生成的提示/响应对，专门用于微调 LLM。

Rust 官方文档 1.6k 数据集
https://www.kaggle.com/datasets/asnowwolf/rust-official-book

建立环境

获取 Gemma 的访问权限

要完成本教程，您首先需要按照 Gemma 设置中的说明进行设置。Gemma 设置说明将指导您完成以下步骤:

Gemma 设置
https://ai.google.dev/gemma/docs/setup

Gemma 模型托管在 Kaggle 上。要使用 Gemma，请在 Kaggle 上请求访问权限:

登录或注册 https://www.kaggle.com/
打开 Gemma 模型卡，并选择 "请求访问"
填写同意书并接受条款和条件

Gemma 模型卡
https://www.kaggle.com/models/google/gemma

安装依赖

安装 Keras、KerasNLP 和其他依赖。

#InstallKeras3last.Seehttps://keras.io/getting_started/formoredetails.!pipinstall-q-Ukeras-nlp!pipinstall-q-Ukeras>=3

选定后端

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

在本教程中，请将后端配置为 JAX。

import os
os.environ["KERAS_BACKEND"] = "jax"# Or "torch" or "tensorflow".# Avoid memory fragmentation on JAX backend.os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"

导入包

导入 Keras 和 KerasNLP。

importkerasimportkeras_nlp

定义提示模板

template="Instruction:\n{question}\n\nResponse:\n{answer}"

加载模型

KerasNLP 提供了许多流行的模型架构{:.external} 的实现。在本教程中，您将使用 GemmaCausalLM 创建一个模型，这是一个用于因果语言建模的端到端 Gemma 2 模型。因果语言模型根据之前的 tokens 预测下一个 token。

模型架构
https://keras.io/api/keras_nlp/models/

使用 from_preset 方法创建模型:

gemma_lm=keras_nlp.models.GemmaCausalLM.from_preset("gemma2_instruct_2b_en")gemma_lm.summary()

normalizer.cc(51) LOG(INFO) precompiled_charsmap is empty. use identity normalization.

Preprocessor:"gemma_causal_lm_preprocessor"

Model:"gemma_causal_lm"

Totalparams:2,614,341,888(9.74GB)

Trainableparams:2,614,341,888(9.74GB)

Non-trainableparams:0(0.00B)

from_preset 方法根据预设的架构和权重实例化模型。在上面的代码中，字符串 "gemma2_instruct_2b_en" 指定了预设的架构——一个拥有 20 亿参数的 Gemma 2 指令对齐过的模型。

注意: 还有一个拥有 70 亿参数的 Gemma 模型可用。如果要在 Colab 中运行更大的模型，您需要访问付费计划中的高级 GPU。或者，您可以在 Kaggle 或 Google Cloud 上对 Gemma 7B 模型进行分布式微调。

Gemma 7B 模型进行分布式微调
https://ai.google.dev/gemma/docs/distributed_tuning

在微调之前推理

在本节中，您将使用各种提示查询模型，以观察它的响应。

查询书中提到的与 Rust 相关的知识

您可以提出一些与 Rust 相关的问题或提示，以查看模型如何回应这些查询。

prompt=template.format(question="HowcanIoverloadthe`+`operatorforarithmeticadditioninRust?",answer="",)print(gemma_lm.generate(prompt,max_length=256))

Instruction:HowcanIoverloadthe`+`operatorforarithmeticadditioninRust?Response:```ruststructPoint{x:f64,y:f64,}implPoint{fnnew(x:f64,y:f64)->Self{Point{x,y}}fnadd(self,otheroint)->oint{Point{x:self.x+other.x,y:self.y+other.y,}}}fnmain(){letp1=Point::new(1.0,2.0);letp2=Point::new(3.0,4.0);letresult=p1+p2;println!("Result{},{})",result.x,result.y);}```**Explanation:**1.**StructDefinition:**Wedefinea`Point`structtorepresentpointsin2Dspace.2.**`add`Method:**Weimplementthe`+`operatorforthe`Point`

LoRA 微调

为了从模型中获得更好的响应，使用 Rust 官方文档 1.6k 数据集对模型进行低秩适应 (LoRA) 微调。

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

更高的秩意味着可以进行更详细的调整，但也意味着更多的可训练参数。较低的秩意味着较少的计算开销，但可能会有较低的适应精度。

本教程使用了秩为 4 的 LoRA。在实际操作中，建议从相对较小的秩 (如 4、8、16) 开始。这在实验中计算效率较高。使用这个秩训练您的模型，并评估在您的任务上的性能提升。在后续的试验中逐渐增加秩，看看是否能进一步提高性能。

加载数据集

预处理数据

import jsondata = []with open('/kaggle/input/rust-official-book/dataset.jsonl', encoding='utf-8') as file:for line in file:features = json.loads(line)# Format the entire example as a single string.data.append(template.format(**features))
# Only use 1000 training examples, to keep it fast.# data = data[:100]

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

Preprocessor:"gemma_causal_lm_preprocessor"

Model:"gemma_causal_lm"

Totalparams:2,617,270,528(9.75GB)

Trainableparams:2,928,640(11.17MB)

Non-trainableparams:2,614,341,888(9.74GB)

注意，启用 LoRA 会显著减少可训练参数的数量 (从 25 亿减少到 130 万)。

# Limit the input sequence length to 512 (to control memory usage).gemma_lm.preprocessor.sequence_length = 512# Use AdamW (a common optimizer for transformer models).optimizer = keras.optimizers.AdamW(learning_rate=5e-5,weight_decay=0.01,)# Exclude layernorm and bias terms from decay.optimizer.exclude_from_weight_decay(var_names=["bias", "scale"])
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)

1632/1632 ━━━━━━━━━━━━━━━━━━━━ 1390s 832ms/step - loss: 0.1808 - sparse_categorical_accuracy: 0.6240

<keras.src.callbacks.history.History at 0x7c45a0251c30>

微调后的推理

在微调之后，模型的响应会遵循提示中提供的指示。

查询书中提到的与 Rust 相关的知识

您可以提出一些与 Rust 相关的问题或提示，以查看模型如何回应这些查询。

prompt=template.format(question="HowcanIoverloadthe`+`operatorforarithmeticadditioninRust?",answer="",)print(gemma_lm.generate(prompt,max_length=256))

Instruction:

How can I overload the `+` operator for arithmetic addition in Rust?

Response:

You can overload the `+` operator by implementing the `Add` trait for your struct.

注意，本教程在一个小型粗糙数据集上进行微调，仅训练一个轮次 (epoch)，并使用较低的 LoRA 秩值。为了从微调后的模型中获得更好的响应，您可以尝试以下方法:

增加微调数据集的大小
提高微调数据集的质量 (人工核查)
训练更多的轮次(epochs)
设置更高的 LoRA 秩
修改超参数值，如 learning_rate 和 weight_decay