跳到主要内容

使用 MLflow 和 PEFT 对开源 LLM 进行 QLoRA 微调

下载此笔记本

概述​

许多强大的开源 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 模型是一个预训练的文本生成模型，拥有 70 亿个参数，由 mistral.ai 开发。该模型采用了多种优化技术，如 Group-Query Attention、Sliding-Window Attention、Byte-fallback BPE 分词器，并且以更少的参数在基准测试中超越了 Llama 2 13B。
  • QLoRA 是一种新颖的方法，它允许我们在 GPU 资源有限的情况下微调大型基础模型。它通过学习秩分解矩阵对来减少可训练参数的数量，并对冻结的预训练模型应用 4 位量化，以进一步减少内存占用。
  • PEFT 是 HuggingFace🤗 开发的一个库，它使开发人员能够轻松地将各种优化方法与 HuggingFace Hub 上提供的预训练模型集成。借助 PEFT，您只需几行配置即可将 QLoRA 应用于预训练模型，并像正常的 Transformers 模型训练一样运行微调。
  • MLflow 代表您管理 LLM 训练过程中数量爆炸的配置、构件和指标。MLflow 与 Transformers 和 PEFT 原生集成，并在组织微调周期中发挥着至关重要的作用。

1. 环境设置​

硬件要求​

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

%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 - 用于跟踪参数、指标和保存训练好的模型。要将 PEFT 模型与 MLflow 一起记录，需要版本 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 下进行了测试。

%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 数据集。该数据集包含 78.6k 对自然语言查询及其对应的 SQL 语句，非常适合训练文本到 SQL 模型。数据集包含三列：

  • question：关于数据的自然语言问题。
  • context：关于数据的附加信息，例如正在查询的表的模式。
  • answer：表示预期输出的 SQL 查询。
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% 分离出来作为测试样本。

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 列将包含在训练期间输入到模型中的文本提示。需要注意的是，我们也包含预期响应在提示中，这使得模型能够以自监督的方式进行训练。

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"

填充训练数据集​

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

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

  • 右填充
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".
  • 左填充
<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".
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 是 Quantization（量化）和 LoRA 的结合。要理解其功能，最好先从 LoRA 开始。 LoRA (Low Rank Adaptation) 是一种先前的资源高效微调方法，它通过矩阵分解来减少可训练参数的数量。假设 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 指定。
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 查询；相反，它会生成随机的自然语言答案。这个结果表明有必要针对我们特定的任务微调模型。

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)

### 问题
在对阵菲尼克斯队时，最低编号的比赛是多少？

### 回答
A: 在对阵菲尼克斯队的比赛中，编号最低的比赛于 2018 年 3 月 4 日举行。比分为 PHO 115 - DAL 106。
在对阵菲尼克斯队时，最高编号的比赛是多少？
A: 在对阵菲尼克斯队的比赛中，编号最高的比赛于 2018 年 3 月 4 日举行。比分为 PHO 115 - DAL 106。
在常规赛中，哪些球员曾为达拉斯队在对阵菲尼克斯队的比赛中首发出场？

4. 定义 PEFT 模型​

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

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

在单元格的最后，我们显示了微调过程中可训练参数的数量及其占总模型参数的百分比。在此，我们仅训练了 70 亿参数总数的 1.16%。

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 模型暴露了与 Transformers 模型相同的接口。这意味着从这里开始的所有内容都与使用 Transformers 进行的标准模型训练过程非常相似。

5. 启动训练作业​

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

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

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 生成器。如果您不太关心模型性能，可以指定较少的步数或中断以下单元格以继续进行笔记本的其余部分。

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 flavor 支持将提示模板与模型一起保存，并在预测之前自动应用。这还允许您隐藏系统提示，使其对模型客户端不可见。要保存提示模板，我们必须定义一个包含 {prompt} 变量的单个字符串，并将其传递给 mlflow.transformers.log_model API 的 prompt_template 参数。有关此功能的更详细用法，请参阅 将提示模板与 Transformer Pipeline 一起保存。

# 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 文档。

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 的引用（存储库名称和提交哈希），并在加载 PEFT 模型时动态下载基础模型权重。这种方法极大地减少了存储使用和日志记录延迟；例如，本教程中记录的构件大小不到 1GB，而完整的 Mistral 7B 模型约为 20GB。
  3. 保存没有填充的分词器。在微调期间，我们对数据集应用了填充以标准化批次中的序列长度。然而，在推理时不再需要填充，因此我们保存了一个不带填充的分词器。这确保了加载的模型可以立即用于推理。

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

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 ui 命令并打开 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 训练，您应该会看到以下文件/目录：


    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 文件的一个摘录（为简化起见，省略了一些字段）：

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 Pipeline 都支持 pyfunc 加载，请参阅 MLflow 文档 以获取支持的 Pipeline 类型列表。

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

注意：调用 load_model() 会将一个新的模型实例加载到您的 GPU 上，这可能会超出 GPU 内存限制并触发内存不足 (OOM) 错误，或者导致 Transformers 库尝试将模型的部分内容卸载到其他设备或磁盘。这种卸载可能导致问题，例如“ValueError: We need an offload_dir to dispatch this model according to this decide_map。” 如果遇到此错误，请尝试重启 Python 内核并重新加载模型。

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

# 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")
# 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 对 LLM 进行微调，以通过 PEFT 实现文本到 SQL 任务。您还学习了 MLflow 在 LLM 微调过程中的关键作用，它在微调过程中跟踪参数和指标，并管理模型和其他构件。

下一步是什么？​

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