unity

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

Ne0inhk

25 Dec 2024 — 9 min read

www.zeeklog.com - 使用 LoRA 对 Gemma 2 进行微调，以纳入 Rust 官方文档

本文原作者: Zhicheng Wang，原文发布于: Kaggle

https://www.kaggle.com/code/asnowwolf/lora-gemma2-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 和其他依赖。

# Install Keras 3 last. See https://keras.io/getting_started/ for more details. !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". # Avoid memory fragmentation on JAX backend. os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"

导入包

导入 Keras 和 KerasNLP。

import kerasimport keras_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"

Total params: 2,614,341,888 (9.74 GB)

Trainable params: 2,614,341,888 (9.74 GB)

Non-trainable params: 0 (0.00 B)

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="How can I overload the `+` operator for arithmetic addition in Rust?", answer="", ) print(gemma_lm.generate(prompt, max_length=256))

Instruction:How can I overload the `+` operator for arithmetic addition in Rust?Response:```ruststruct Point { x: f64, y: f64,}impl Point { fn new(x: f64, y: f64) -> Self { Point { x, y } } fn add(self, other: Point) -> Point { Point { x: self.x + other.x, y: self.y + other.y, } }}fn main() { let p1 = Point::new(1.0, 2.0); let p2 = Point::new(3.0, 4.0); let result = p1 + p2; println!("Result: ({}, {})", result.x, result.y);}```**Explanation:**1. **Struct Definition:** We define a `Point` struct to represent points in 2D space.2. **`add` Method:** We implement the `+` operator for the `Point`

LoRA 微调

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

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

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

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

加载数据集

预处理数据

import json data = [] 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]

# Enable LoRA for the model and set the LoRA rank to 4. gemma_lm.backbone.enable_lora(rank=4) gemma_lm.summary()

Preprocessor: "gemma_causal_lm_preprocessor"

Model: "gemma_causal_lm"

Total params: 2,617,270,528 (9.75 GB)

Trainable params: 2,928,640 (11.17 MB)

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

注意，启用 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="How can I overload the `+` operator for arithmetic addition in Rust?", 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 秩值。为了从微调后的模型中获得更好的响应，您可以尝试以下方法: