跳到主要内容

使用 QLoRA、MLflow 和 PEFT 微调开源 LLM

下载此笔记本

概述​

许多强大的开源 LLM 已经出现并易于获取。但是，它们并非开箱即用，可以直接部署到您的生产环境中；相反，您需要为您的特定任务（例如聊天机器人、内容生成等）**微调**它们。然而，挑战之一是训练 LLM 通常非常昂贵。即使您的微调数据集很小，反向传播步骤也需要计算数十亿个参数的梯度。例如，完全微调 Llama7B 模型需要 112GB 的 VRAM，即至少两块 80GB 的 A100 GPU。幸运的是，在如何降低 LLM 微调成本方面，已经有很多研究工作。

在本教程中，我们将演示如何使用**单个 24GB VRAM GPU** 微调 Mistral 7B 模型，构建一个强大的**文本到 SQL** 生成器。

您将学到什么​

  1. 动手学习典型的 LLM 微调过程。
  2. 了解如何使用 **QLoRA** 和 **PEFT** 来克服微调的 GPU 内存限制。
  3. 使用 **MLflow** 管理模型训练周期，记录模型工件、超参数、指标和提示。
  4. 如何在 MLflow 中保存提示模板和推理参数（例如 max_token_length），以简化预测接口。

关键参与者​

在本教程中，您将通过实际运行代码来学习高效 LLM 微调背后的技术和方法。下方每个单元格都有更详细的解释，但让我们先快速预览一下本教程中使用的几个主要重要库/方法。

  • Mistral-7B-v0.1 模型是 mistral.ai 开发的一个具有 70 亿参数的预训练文本生成模型。该模型采用了各种优化技术，如分组查询注意力（Group-Query Attention）、滑动窗口注意力（Sliding-Window Attention）、字节回退 BPE 分词器，并且以更少的参数在基准测试中超越了 Llama 2 13B。
  • QLoRA 是一种新颖的方法，允许我们在 GPU 资源有限的情况下微调大型基础模型。它通过学习秩分解矩阵对来减少可训练参数的数量，并对冻结的预训练模型应用 4 位量化，以进一步降低内存占用。
  • PEFT 是 HuggingFace🤗 开发的一个库，它使开发人员能够轻松地将各种优化方法与 HuggingFace Hub 上提供的预训练模型集成。使用 PEFT，您只需几行配置即可将 QLoRA 应用于预训练模型，并像正常训练 Transformer 模型一样进行微调。
  • MLflow 代表您管理 LLM 训练过程中爆炸性的配置、资产和指标数量。MLflow 与 Transformers 和 PEFT 原生集成，并在组织微调周期中发挥关键作用。

1. 环境设置​

硬件要求​

请确保您的 GPU 至少有 20GB 的 VRAM 可用。本笔记本已在单块 24GB VRAM 的 NVIDIA A10G GPU 上进行测试。

python
%sh nvidia-smi
Wed Feb 21 07:16:13 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.54.03              Driver Version: 535.54.03    CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA A10G                    Off | 00000000:00:1E.0 Off |                    0 |
|  0%   15C    P8              16W / 300W |      4MiB / 23028MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                       
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|  No running processes found                                                           |
+---------------------------------------------------------------------------------------+

安装 Python 库​

本教程使用以下 Python 库

  • mlflow - 用于跟踪参数、指标和保存训练好的模型。要使用 MLflow 记录 PEFT 模型，需要版本 **2.11.0 或更高版本**。
  • transformers - 用于定义模型、分词器和训练器。
  • peft - 用于在 Transformer 模型之上创建 LoRA 适配器。
  • bitsandbytes - 用于以 4 位量化加载基础模型以实现 QLoRA。
  • accelerate - bitsandbytes 所需的依赖项。
  • datasets - 用于从 HuggingFace hub 加载训练数据集。

注意：安装这些依赖项后，可能需要重新启动 Python 内核。

本笔记本已在 mlflow==2.11.0, transformers==4.35.2, peft==0.8.2, bitsandbytes==0.42.0, accelerate==0.27.2, 和 datasets==2.17.1 下进行测试。

python
%pip install mlflow>=2.11.0
%pip install transformers peft accelerate bitsandbytes datasets -q -U

2. 数据集准备​

从 HuggingFace Hub 加载数据集​

在本教程中，我们将使用 Hugging Face Hub 上的 b-mc2/sql-create-context 数据集。该数据集包含 7.86 万对自然语言查询及其对应的 SQL 语句，非常适合训练文本到 SQL 模型。数据集包含三列：

  • question：关于数据的自然语言问题。
  • context：关于数据的附加信息，例如正在查询的表的模式。
  • answer：表示预期输出的 SQL 查询。
python
import pandas as pd
from datasets import load_dataset
from IPython.display import HTML, display

dataset_name = "b-mc2/sql-create-context"
dataset = load_dataset(dataset_name, split="train")


def display_table(dataset_or_sample):
  # A helper fuction to display a Transformer dataset or single sample contains multi-line string nicely
  pd.set_option("display.max_colwidth", None)
  pd.set_option("display.width", None)
  pd.set_option("display.max_rows", None)

  if isinstance(dataset_or_sample, dict):
      df = pd.DataFrame(dataset_or_sample, index=[0])
  else:
      df = pd.DataFrame(dataset_or_sample)

  html = df.to_html().replace("\n", "<br>")
  styled_html = f"""<style> .dataframe th, .dataframe tbody td {{ text-align: left; padding-right: 30px; }} </style> {html}"""
  display(HTML(styled_html))


display_table(dataset.select(range(3)))
question context answer
0 有多少部门主管年龄大于 56 岁？ CREATE TABLE head (age INTEGER) SELECT COUNT(*) FROM head WHERE age > 56
1 按年龄顺序列出部门主管的姓名、出生地和年龄。 CREATE TABLE head (name VARCHAR, born_state VARCHAR, age VARCHAR) SELECT name, born_state, age FROM head ORDER BY age
2 列出每个部门的创建年份、名称和预算。 CREATE TABLE department (creation VARCHAR, name VARCHAR, budget_in_billions VARCHAR) SELECT creation, name, budget_in_billions FROM department

拆分训练集和测试集​

b-mc2/sql-create-context 数据集仅包含一个拆分，“train”。我们将其中 20% 作为测试样本。

python
split_dataset = dataset.train_test_split(test_size=0.2, seed=42)
train_dataset = split_dataset["train"]
test_dataset = split_dataset["test"]

print(f"Training dataset contains {len(train_dataset)} text-to-SQL pairs")
print(f"Test dataset contains {len(test_dataset)} text-to-SQL pairs")
Training dataset contains 62861 text-to-SQL pairs
Test dataset contains 15716 text-to-SQL pairs

定义提示模板​

Mistral 7B 模型是一个文本理解模型，因此我们需要构建一个包含用户问题、上下文和系统指令的文本提示。数据集中新的 prompt 列将包含在训练期间输入到模型的文本提示。需要注意的是，我们也包含了预期的响应，以便模型能够以自监督的方式进行训练。

python
PROMPT_TEMPLATE = """You are a powerful text-to-SQL model. Given the SQL tables and natural language question, your job is to write SQL query that answers the question.

### Table:
{context}

### Question:
{question}

### Response:
{output}"""


def apply_prompt_template(row):
  prompt = PROMPT_TEMPLATE.format(
      question=row["question"],
      context=row["context"],
      output=row["answer"],
  )
  return {"prompt": prompt}


train_dataset = train_dataset.map(apply_prompt_template)
display_table(train_dataset.select(range(1)))
question context answer prompt
0 哪个珀斯有黄金海岸是，悉尼是，墨尔本是，阿德莱德是？ CREATE TABLE table_name_56 (perth VARCHAR, adelaide VARCHAR, melbourne VARCHAR, gold_coast VARCHAR, sydney VARCHAR) SELECT perth FROM table_name_56 WHERE gold_coast = "yes" AND sydney = "yes" AND melbourne = "yes" AND adelaide = "yes" 您是一个强大的文本到 SQL 模型。给定 SQL 表和自然语言问题，您的任务是编写 SQL 查询来回答问题。

### 表
CREATE TABLE table_name_56 (perth VARCHAR, adelaide VARCHAR, melbourne VARCHAR, gold_coast VARCHAR, sydney VARCHAR)

### 问题
哪个珀斯有黄金海岸是，悉尼是，墨尔本是，阿德莱德是？

### 回答
SELECT perth FROM table_name_56 WHERE gold_coast = "yes" AND sydney = "yes" AND melbourne = "yes" AND adelaide = "yes"

填充训练数据集​

作为数据集准备的最后一步，我们需要对训练数据集应用**填充**。填充可确保批次中的所有输入序列都具有相同的长度。

需要注意的关键一点是必须在左侧添加填充。采用这种方法是因为模型以自回归方式生成 token，这意味着它从最后一个 token 继续。在右侧添加填充会导致模型从这些填充 token 生成新 token，从而导致输出序列中包含中间的填充 token。

  • 右填充
text
Today |  is  |   a    |  cold  |  <pad>  ==generate=>  "Today is a cold <pad> day"
 How  |  to  | become |  <pad> |  <pad>  ==generate=>  "How to become a <pad> <pad> great engineer".
  • 左填充
text
<pad> |  Today  |  is  |  a   |  cold     ==generate=>  "<pad> Today is a cold day"
<pad> |  <pad>  |  How |  to  |  become   ==generate=>  "<pad> <pad> How to become a great engineer".
python
from transformers import AutoTokenizer

base_model_id = "mistralai/Mistral-7B-v0.1"

# You can use a different max length if your custom dataset has shorter/longer input sequences.
MAX_LENGTH = 256

tokenizer = AutoTokenizer.from_pretrained(
  base_model_id,
  model_max_length=MAX_LENGTH,
  padding_side="left",
  add_eos_token=True,
)
tokenizer.pad_token = tokenizer.eos_token


def tokenize_and_pad_to_fixed_length(sample):
  result = tokenizer(
      sample["prompt"],
      truncation=True,
      max_length=MAX_LENGTH,
      padding="max_length",
  )
  result["labels"] = result["input_ids"].copy()
  return result


tokenized_train_dataset = train_dataset.map(tokenize_and_pad_to_fixed_length)

assert all(len(x["input_ids"]) == MAX_LENGTH for x in tokenized_train_dataset)

display_table(tokenized_train_dataset.select(range(1)))

3. 加载基础模型（带 4 位量化）​

接下来，我们将加载 Mistral 7B 模型，它将作为我们微调的基础模型。可以使用 Transformers 的 from_pretrained() API 从 HuggingFace Hub 仓库 mistralai/Mistral-7B-v0.1 加载此模型。但是，此处我们还提供了一个 quantization_config 参数。

该参数体现了 QLoRA 的关键技术，它显著降低了微调过程中的内存使用。下一段将详细介绍此方法及其配置的含义。但是，如果您觉得复杂，可以随意跳过。毕竟，我们很少需要自己修改 quantization_config 的值 :)

它是如何工作的

简而言之，QLoRA 是**量化**和**LoRA**的结合。要理解其功能，从 LoRA 开始会更简单。LoRA（低秩适配）是一种资源高效微调的先前方法，通过矩阵分解减少可训练参数的数量。令 W' 表示微调后的最终权重矩阵。在 LoRA 中，W' 近似等于原始权重及其更新的总和，即 W + ΔW，然后将 delta 部分分解为两个低秩矩阵，即 ΔW ≈ AB。假设 W 是 mxm，并且我们为 A 和 B 的秩选择一个较小的 r，其中 A 是 mxr，B 是 rxm。现在，原始可训练参数，其大小与 W 成二次方（即 m^2），在分解后变为 2mr。经验上，我们可以为 r 选择一个比完整权重矩阵大小小得多的数字，例如 32、64，因此这大大减少了需要训练的参数数量。

QLoRA 扩展了 LoRA，采用了相同的矩阵分解策略。但是，它通过对冻结的预训练模型 W 应用 4 位量化来进一步减少内存使用。根据他们的研究，LoRA 微调过程中最大的内存消耗是通过冻结参数 W 进行反向传播以计算适配器 A 和 B 的梯度。因此，将 W 量化为 4 位可以显着降低整体内存消耗。这通过下面的 load_in_4bit=True 设置实现。

此外，QLoRA 还引入了其他技术来优化资源使用，而不会显著影响模型性能。有关更多技术细节，请参阅论文，但我们通过在 bitsandbytes 中设置以下量化配置来实现它们：

  • 4 位 NormalFloat 类型通过 bnb_4bit_quant_type="nf4" 指定。
  • 双量化通过 bnb_4bit_use_double_quant=True 激活。
  • QLoRA 在计算 A 和 B 的梯度时，会将 4 位权重重新量化为更高的精度，以防止性能下降。此数据类型由 bnb_4bit_compute_dtype=torch.bfloat16 指定。
python
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
  # Load the model with 4-bit quantization
  load_in_4bit=True,
  # Use double quantization
  bnb_4bit_use_double_quant=True,
  # Use 4-bit Normal Float for storing the base model weights in GPU memory
  bnb_4bit_quant_type="nf4",
  # De-quantize the weights to 16-bit (Brain) float before the forward/backward pass
  bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(base_model_id, quantization_config=quantization_config)

基础模型表现如何？​

首先，在进行任何微调之前，让我们评估一下原始 Mistral 模型在 SQL 生成任务上的表现。正如预期的那样，模型不会生成正确的 SQL 查询；相反，它会生成随机的自然语言答案。这一结果表明了为特定任务微调模型的必要性。

python
import transformers

tokenizer = AutoTokenizer.from_pretrained(base_model_id)
pipeline = transformers.pipeline(model=model, tokenizer=tokenizer, task="text-generation")

sample = test_dataset[1]
prompt = PROMPT_TEMPLATE.format(
  context=sample["context"], question=sample["question"], output=""
)  # Leave the answer part blank

with torch.no_grad():
  response = pipeline(prompt, max_new_tokens=256, repetition_penalty=1.15, return_full_text=False)

display_table({"prompt": prompt, "generated_query": response[0]["generated_text"]})
prompt generated_query
0 您是一个强大的文本到 SQL 模型。给定 SQL 表和自然语言问题，您的任务是编写 SQL 查询来回答问题。

### 表
CREATE TABLE table_name_61 (game INTEGER, opponent VARCHAR, record VARCHAR)

### 问题
在对阵菲尼克斯的比赛中，最低编号是多少，战绩是 29-17？

### 回答
答：对阵菲尼克斯的最低编号比赛是在 2018 年 4 月 3 日进行的。比分是 PHO 115 - DAL 106。
对阵菲尼克斯的最高编号比赛是多少？
答：对阵菲尼克斯的最高编号比赛是在 2018 年 4 月 3 日进行的。比分是 PHO 115 - DAL 106。
在常规赛对阵菲尼克斯的比赛中，哪些球员曾担任达拉斯的组织后卫？

4. 定义 PEFT 模型​

如前所述，QLoRA 代表**量化** + **LoRA**。在应用了量化部分后，现在我们继续处理 LoRA 部分。尽管 LoRA 背后的数学原理很复杂，但 PEFT 通过简化将 LoRA 应用于预训练 Transformer 模型的过程来帮助我们。

在下一个单元格中，我们创建一个带有各种 LoRA 设置的 LoraConfig。与之前的 quantization_config 不同，这些超参数可能需要优化才能为您的特定任务实现最佳模型性能。**MLflow** 通过跟踪这些超参数、相关的模型及其结果来促进这一过程。

在单元格的末尾，我们显示了微调期间可训练参数的数量及其占模型总参数的百分比。在这里，我们只训练了 70 亿参数总数的 1.16%。

python
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# Enabling gradient checkpointing, to make the training further efficient
model.gradient_checkpointing_enable()
# Set up the model for quantization-aware training e.g. casting layers, parameter freezing, etc.
model = prepare_model_for_kbit_training(model)

peft_config = LoraConfig(
  task_type="CAUSAL_LM",
  # This is the rank of the decomposed matrices A and B to be learned during fine-tuning. A smaller number will save more GPU memory but might result in worse performance.
  r=32,
  # This is the coefficient for the learned ΔW factor, so the larger number will typically result in a larger behavior change after fine-tuning.
  lora_alpha=64,
  # Drop out ratio for the layers in LoRA adaptors A and B.
  lora_dropout=0.1,
  # We fine-tune all linear layers in the model. It might sound a bit large, but the trainable adapter size is still only **1.16%** of the whole model.
  target_modules=[
      "q_proj",
      "k_proj",
      "v_proj",
      "o_proj",
      "gate_proj",
      "up_proj",
      "down_proj",
      "lm_head",
  ],
  # Bias parameters to train. 'none' is recommended to keep the original model performing equally when turning off the adapter.
  bias="none",
)

peft_model = get_peft_model(model, peft_config)
peft_model.print_trainable_parameters()
trainable params: 85,041,152 || all params: 7,326,773,248 || trainable%: 1.1606903765339511

**就是这样！！！** PEFT 使 LoRA 设置变得非常容易。

另一个好处是 PEFT 模型暴露了与 Transformer 模型相同的接口。这意味着从这里开始的所有内容都与使用 Transformer 进行的标准模型训练过程非常相似。

5. 启动训练作业​

与传统的 Transformers 训练类似，我们首先设置一个 Trainer 对象来组织训练迭代。有许多超参数需要配置，但 MLflow 将代表您管理它们。

要启用 MLflow 日志记录，您可以指定 report_to="mlflow" 并使用 run_name 参数为您的训练试验命名。此操作将启动一个 MLflow run，该 run 会自动记录训练指标、超参数、配置和训练好的模型。

python
from datetime import datetime

import transformers
from transformers import TrainingArguments

import mlflow

# Comment-out this line if you are running the tutorial on Databricks
mlflow.set_experiment("MLflow PEFT Tutorial")

training_args = TrainingArguments(
  # Set this to mlflow for logging your training
  report_to="mlflow",
  # Name the MLflow run
  run_name=f"Mistral-7B-SQL-QLoRA-{datetime.now().strftime('%Y-%m-%d-%H-%M-%s')}",
  # Replace with your output destination
  output_dir="YOUR_OUTPUT_DIR",
  # For the following arguments, refer to https://hugging-face.cn/docs/transformers/main_classes/trainer
  per_device_train_batch_size=2,
  gradient_accumulation_steps=4,
  gradient_checkpointing=True,
  optim="paged_adamw_8bit",
  bf16=True,
  learning_rate=2e-5,
  lr_scheduler_type="constant",
  max_steps=500,
  save_steps=100,
  logging_steps=100,
  warmup_steps=5,
  # https://discuss.huggingface.co/t/training-llama-with-lora-on-multiple-gpus-may-exist-bug/47005/3
  ddp_find_unused_parameters=False,
)

trainer = transformers.Trainer(
  model=peft_model,
  train_dataset=tokenized_train_dataset,
  data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
  args=training_args,
)

# use_cache=True is incompatible with gradient checkpointing.
peft_model.config.use_cache = False

训练时间可能需要几个小时，具体取决于您的硬件规格。尽管如此，本教程的主要目的是让您熟悉使用 PEFT 和 MLflow 进行微调的过程，而不是培养一个高性能的 SQL 生成器。如果您不太关心模型性能，可以指定较少的步数或中断以下单元格以继续进行笔记本的其余部分。

python
trainer.train()
[500/500 45:41, Epoch 0/1]
Step Training Loss
100 0.681700
200 0.522400
300 0.507300
400 0.494800
500 0.474600 

TrainOutput(global_step=500, training_loss=0.5361956100463867, metrics={'train_runtime': 2747.9223, 'train_samples_per_second': 1.456, 'train_steps_per_second': 0.182, 'total_flos': 4.421038813216768e+16, 'train_loss': 0.5361956100463867, 'epoch': 0.06})

6. 将 PEFT 模型保存到 MLflow​

太棒了！我们已成功将 Mistral 7B 模型微调为 SQL 生成器。在结束训练之前，最后一步是将训练好的 PEFT 模型保存到 MLflow。

设置提示模板和默认推理参数（可选）​

LLM 的预测行为不仅由模型权重定义，还很大程度上受提示和推理参数（如 max_token_length、repetition_penalty）控制。因此，强烈建议将这些元数据与模型一起保存，以便以后加载模型时可以预期一致的行为。

提示模板​

用户提示本身是自由文本，但您可以通过应用“模板”来利用输入。MLflow Transformer 风格支持将提示模板与模型一起保存，并在预测之前自动应用。这还可以让您隐藏系统提示，使其对模型客户端不可见。要保存提示模板，我们必须定义一个包含 {prompt} 变量的单个字符串，并将其传递给 mlflow.transformers.log_model API 的 prompt_template 参数。有关此功能的更详细用法，请参阅使用 Transformer Pipeline 保存提示模板。

python
# Basically the same format as we applied to the dataset. However, the template only accepts {prompt} variable so both table and question need to be fed in there.
prompt_template = """You are a powerful text-to-SQL model. Given the SQL tables and natural language question, your job is to write SQL query that answers the question.

{prompt}

### Response:
"""

推理参数​

推理参数可以作为模型签名的一部分与 MLflow 模型一起保存。签名定义了模型的输入和输出格式，并包含传递给模型预测的其他参数。您可以使用 mlflow.models.infer_signature API 从一些示例输入推断出签名。如果您传递了参数的具体值，MLflow 会将它们视为默认值，并在推理时应用它们（如果用户未提供）。有关模型签名的更多详细信息，请参阅MLflow 文档。

python
from mlflow.models import infer_signature

sample = train_dataset[1]

# MLflow infers schema from the provided sample input/output/params
signature = infer_signature(
  model_input=sample["prompt"],
  model_output=sample["answer"],
  # Parameters are saved with default values if specified
  params={"max_new_tokens": 256, "repetition_penalty": 1.15, "return_full_text": False},
)
signature
inputs: 
[string (required)]
outputs: 
[string (required)]
params: 
['max_new_tokens': long (default: 256), 'repetition_penalty': double (default: 1.15), 'return_full_text': boolean (default: False)]

将 PEFT 模型保存到 MLflow​

最后，我们将调用 mlflow.transformers.log_model API 将模型记录到 MLflow。在记录 PEFT 模型到 MLflow 时需要注意的几个关键点：

  1. MLflow 将 Transformer 模型记录为Pipeline。 Pipeline 将模型与其分词器（或其他组件，取决于任务类型）捆绑在一起，并将预测步骤简化为易于使用的接口，使其成为确保可复用性的绝佳工具。在下面的代码中，我们传递了模型和分词器作为字典，然后 MLflow 会自动推断出正确的 Pipeline 类型并进行保存。
  2. MLflow 不会为 PEFT 模型保存基础模型权重。在执行 mlflow.transformers.log_model 时，MLflow 只会保存少量已训练的参数，即 PEFT 适配器。对于基础模型，MLflow 会记录对 HuggingFace hub 的引用（仓库名称和 commit hash），并在加载 PEFT 模型时按需下载基础模型权重。这种方法极大地减少了存储使用量和日志记录延迟；例如，本教程中记录的工件大小小于 1GB，而完整的 Mistral 7B 模型约为 20GB。
  3. 保存不带填充的分词器。在微调过程中，我们对数据集应用了填充以标准化批次中的序列长度。但是，在推理时不再需要填充，因此我们保存了一个不带填充的分词器。这确保了加载的模型可以立即用于推理。

注意：目前，PEFT 适配器和配置需要手动记录，而其他信息，如数据集、指标、Trainer 参数等，会自动记录。但是，此过程可能在 MLflow 和 Transformers 的未来版本中实现自动化。

python
import mlflow

# Get the ID of the MLflow Run that was automatically created above
last_run_id = mlflow.last_active_run().info.run_id

# Save a tokenizer without padding because it is only needed for training
tokenizer_no_pad = AutoTokenizer.from_pretrained(base_model_id, add_bos_token=True)

# If you interrupt the training, uncomment the following line to stop the MLflow run
# mlflow.end_run()

with mlflow.start_run(run_id=last_run_id):
  mlflow.log_params(peft_config.to_dict())
  mlflow.transformers.log_model(
      transformers_model={"model": trainer.model, "tokenizer": tokenizer_no_pad},
      prompt_template=prompt_template,
      signature=signature,
      name="model",  # This is a relative path to save model files within MLflow run
  )

记录到 MLflow 的内容是什么？​

让我们简要回顾一下训练后记录/保存到 MLflow 的内容。要访问 MLflow UI，请运行 mlflow server 命令并打开 https://:PORT（PORT 默认为 5000）。在左侧选择“MLflow PEFT Tutorial”实验（或在 Databricks 上运行时选择笔记本名称）。然后点击最新的 MLflow Run，名为 Mistral-7B-SQL-QLoRA-2024-...，以查看 Run 详细信息。

参数​

Parameters 部分显示了为 Trainer、LoraConfig 和 BitsAndBytesConfig 指定的数百个参数，例如 learning_rate、r、bnb_4bit_quant_type。它还包括未显式指定的默认参数，这对于确保可复用性至关重要，尤其是当库的默认值发生变化时。

指标​

Metrics 部分展示了在运行过程中收集的模型指标，例如 train_loss。您可以在“Chart”选项卡中使用各种图表可视化这些指标。

工件​

Artifacts 部分显示了训练后在 MLflow 中保存的文件/目录。对于 Transformers PEFT 训练，您应该会看到以下文件/目录：

text

    model/
      ├─ peft/
      │  ├─ adapter_config.json       # JSON file of the LoraConfig
      │  ├─ adapter_module.safetensor # The weight file of the LoRA adapter
      │  └─ README.md                 # Empty README file generated by Transformers
      │
      ├─ LICENSE.txt                  # License information about the base model (Mistral-7B-0.1)
      ├─ MLModel                      # Contains various metadata about your model
      ├─ conda.yaml                   # Dependencies to create conda environment
      ├─ model_card.md                # Model card text for the base model
      ├─ model_card_data.yaml         # Model card data for the base model
      ├─ python_env.yaml              # Dependencies to create Python virtual environment
      └─ requirements.txt             # Pip requirements for model inference

模型元数据​

在 MLModel 文件中，您可以看到关于 PEFT 和基础模型保存的许多详细元数据。以下是 MLModel 文件的一部分（为简洁起见，省略了一些字段）

text
flavors:
  transformers:
    peft_adaptor: peft                                 # Points the location of the saved PEFT model
    pipeline_model_type: MistralForCausalLM            # The base model implementation
    source_model_name: mistralai/Mistral-7B-v0.1.      # Repository name of the base model
    source_model_revision: xxxxxxx                     # Commit hash in the repository for the base model
    task: text-generation                              # Pipeline type
    torch_dtype: torch.bfloat16                        # Dtype for loading the model
    tokenizer_type: LlamaTokenizerFast                 # Tokenizer implementation

# Prompt template saved with the model above
metadata:
  prompt_template: 'You are a powerful text-to-SQL model. Given the SQL tables and
    natural language question, your job is to write SQL query that answers the question.


    {prompt}


    ### Response:

    '
# Defines the input and output format of the model, with additional inference parameters with default values
signature:
  inputs: '[{"type": "string", "required": true}]'
  outputs: '[{"type": "string", "required": true}]'
  params: '[{"name": "max_new_tokens", "type": "long", "default": 256, "shape": null},
    {"name": "repetition_penalty", "type": "double", "default": 1.15, "shape": null},
    {"name": "return_full_text", "type": "boolean", "default": false, "shape": null}]'

7. 从 MLflow 加载保存的 PEFT 模型​

最后，让我们加载在 MLflow 中记录的模型，并评估其作为文本到 SQL 生成器的性能。有两种方法可以在 MLflow 中加载 Transformer 模型：

  1. 使用 mlflow.transformers.load_model()。此方法返回一个原生的 Transformers pipeline 实例。
  2. 使用 mlflow.pyfunc.load_model()。此方法返回一个包装了 Transformers pipeline 的 MLflow 的 PythonModel 实例，它提供了比原生 pipeline 更多的功能，例如（1）统一的 predict() API 用于推理，（2）模型签名强制执行，以及（3）如果已保存，则自动应用提示模板和默认参数。请注意，并非所有 Transformer pipelines 都支持 pyfunc 加载，请参阅MLflow 文档以获取支持的 Pipeline 类型的完整列表。

如果您希望通过原生的 Transformers 接口使用模型，则首选第一种选项。第二种选项提供了跨不同模型类型的简化和统一的接口，对于在生产部署之前进行模型测试特别有用。在下面的代码中，我们将使用 mlflow.pyfunc.load_model() 来展示它如何应用上面定义的提示模板和默认推理参数。

注意：调用 load_model() 会将新的模型实例加载到您的 GPU 上，这可能会超出 GPU 内存限制并触发内存不足 (OOM) 错误，或者导致 Transformers 库尝试将模型的一部分卸载到其他设备或磁盘。这种卸载可能导致问题，例如“ValueError: 我们需要一个 offload_dir 来根据此 decide_map 来调度此模型。” 如果遇到此错误，请尝试重新启动 Python 内核并再次加载模型。

警告：重新启动 Python 内核将清除上方单元格中的所有中间状态和变量。请确保在重新启动之前，已将训练好的 PEFT 模型正确记录到 MLflow 中。

python
# You can find the ID of run in the Run detail page on MLflow UI
mlflow_model = mlflow.pyfunc.load_model("runs:/YOUR_RUN_ID/model")
python
# We only input table and question, since system prompt is adeed in the prompt template.
test_prompt = """
### Table:
CREATE TABLE table_name_50 (venue VARCHAR, away_team VARCHAR)

### Question:
When Essendon played away; where did they play?
"""

# Inference parameters like max_tokens_length are set to default values specified in the Model Signature
generated_query = mlflow_model.predict(test_prompt)[0]
display_table({"prompt": test_prompt, "generated_query": generated_query})
prompt generated_query
0
### 表
CREATE TABLE table_name_50 (venue VARCHAR, away_team VARCHAR)

### 问题
当埃森顿队客场比赛时，他们在哪里比赛？
 SELECT venue FROM table_name_50 WHERE away_team = "essendon"

太棒了！！微调后的模型现在可以正确生成 SQL 查询。正如您在上面的代码和结果中所见，系统提示和默认推理参数已自动应用，因此我们不必将其传递给加载的模型。当您想要部署多个模型（或更新现有模型）并使用不同的系统提示或参数时，这非常强大，因为您不必编辑客户端的实现，因为它们已在 MLflow 模型中抽象化了 :)

结论​

在本教程中，您学习了如何使用 QLoRA 和 PEFT 为文本到 SQL 任务微调大型语言模型。您还学习了 MLflow 在 LLM 微调过程中的关键作用，它在微调期间跟踪参数和指标，并管理模型和其他资产。

下一步是什么？​

  • 使用 MLflow 评估 Hugging Face LLM - 模型评估是模型开发中的关键步骤。请查看此指南，了解如何使用 MLflow 高效评估 LLM，包括 LLM 作为评估者。
  • 将 MLflow 模型部署到生产环境 - MLflow 模型存储丰富的元数据并提供统一的预测接口，从而简化了部署过程。了解如何将微调后的模型部署到各种目标，例如 AWS SageMaker、Azure ML、Kubernetes、Databricks Model Serving，并提供详细的指南和动手实践笔记本。
  • MLflow Transformers Flavor 文档 - 了解更多关于 MLflow 和 Transformers 集成的信息，并继续学习更多教程。
  • MLflow 中的大型语言模型 - MLflow 提供更多与 LLM 相关的功能，并与 OpenAI 和 Langchain 等许多其他库集成。