跳到主要内容

MLflow 中的 🤗 Transformers

transformers 模型风格允许通过 mlflow.transformers.save_model()mlflow.transformers.log_model() 函数,以 MLflow 格式记录 transformers 模型、组件和管道。使用这些函数还会将 python_function 风格添加到它们生成的 MLflow 模型中,从而允许通过 mlflow.pyfunc.load_model() 将该模型解释为用于推理的通用 Python 函数。您还可以使用 mlflow.transformers.load_model() 函数以原生 transformers 格式加载具有 transformers 风格的已保存或已记录的 MLflow 模型。

此页面详细介绍了 MLflow transformers 风格的功能和配置。有关 MLflow Transformer 集成的常规介绍,请参阅 MLflow Transformers 风格页面。

将 Transformers 模型作为 Python 函数加载

支持的 Transformers 管道类型

transformers python_function (pyfunc) 模型风格 简化并标准化了管道推理的输入和输出。这种一致性允许通过将 transformers 推理管道所需的数据结构强制转换为与 json 序列化和转换为 Pandas DataFrames 兼容的格式来服务和批量推理。

注意

某些 TextGenerationPipeline 类型,尤其是基于指令的类型,可能会在其输出中返回原始提示和包含的行格式换行符 "\n"。对于这些管道类型,如果您想禁用提示返回,可以在保存或记录模型时在 model_config 字典中设置以下内容:"include_prompt": False。要从生成的文本输出的主体中删除换行符,您可以将 "collapse_whitespace": True 选项添加到 model_config 字典中。如果保存的管道类型不继承自 TextGenerationPipeline,这些选项将不会对从管道推理返回的输出执行任何修改。

在当前版本中,音频和基于文本的大型语言模型支持与 pyfunc 一起使用,而计算机视觉、多模态、时间序列、强化学习和图模型仅支持通过 mlflow.transformers.load_model() 进行原生类型加载。

注意

并非所有 transformers 管道类型都支持。请参阅下表,了解当前支持的可以作为 pyfunc 加载的管道类型的列表。

MLflow 的未来版本将引入对这些附加类型的 pyfunc 支持。

下表显示了 transformers 管道类型到 python_function (pyfunc) 模型风格 数据类型输入和输出的映射。

重要

这些管道的 pyfunc 实现的输入和输出不能保证与从给定管道类型的原生使用中返回的输入类型和输出类型相匹配。如果您的用例需要访问分数、top_k 结果或管道推理调用输出中的其他附加引用,请通过 mlflow.transformers.load_model() 加载来使用原生实现,以接收完整输出。

同样,如果您的用例需要使用原始张量输出或通过外部 processor 模块处理输出,请通过调用 mlflow.transformers.load_model() 并将 return_type 参数指定为 'components',直接将模型组件作为 dict 加载。

管道类型输入类型输出类型
指令文本生成str 或 List[str]List[str]
对话str 或 List[str]List[str]
摘要str 或 List[str]List[str]
文本分类str 或 List[str]pd.DataFrame (dtypes: {'label': str, 'score': double})
文本生成str 或 List[str]List[str]
Text2Text 生成str 或 List[str]List[str]
Token 分类str 或 List[str]List[str]
翻译str 或 List[str]List[str]
ZeroShot 分类*Dict[str, [List[str] | str]*pd.DataFrame (dtypes: {'sequence': str, 'labels': str, 'scores': double})
表格问答**Dict[str, [List[str] | str]**List[str]
问题解答***Dict[str, str]***List[str]
填充掩码****str 或 List[str]****List[str]
特征提取str 或 List[str]np.ndarray
自动语音识别bytes*****, str, 或 np.ndarrayList[str]
音频分类bytes*****, str, 或 np.ndarraypd.DataFrame (dtypes: {'label': str, 'score': double})

* 也可以传递这些输入的集合。标准要求的键名是“sequences”和“candidate_labels”,但这些键名可能会有所不同。检查您正在使用的架构的输入要求,以确保提供正确的字典键名。

** 也可以传递这些输入的集合。引用表必须是一个 json 编码的 dict (即 {'query': '我们卖得最多的是什么?', 'table': json.dumps(table_as_dict)})

*** 也可以传递这些输入的集合。标准要求的键名是“question”和“context”。验证预期的输入键名是否与模型的预期输入匹配,以确保您的推理请求可以被正确读取。

**** 您选择的模型的掩码语法将特定于该模型的实现。有些是“[MASK]”,而另一些是“<mask>”。验证预期的语法以避免失败的推理请求。

***** 如果在 MLflow 模型服务中使用 pyfunc 进行实时推理,则原始音频(以字节格式)必须先进行 base64 编码,然后才能提交到端点。字符串输入将被解释为 uri 位置。

将 Transformers 模型作为 Python 函数加载的示例

在下面的示例中,在一个管道中使用了一个简单的预训练模型。记录到 MLflow 后,该管道将作为 pyfunc 加载,并用于从传入的字符串列表中生成响应。

import mlflow
import transformers

# Read a pre-trained conversation pipeline from HuggingFace hub
conversational_pipeline = transformers.pipeline(model="microsoft/DialoGPT-medium")

# Define the signature
signature = mlflow.models.infer_signature(
"Hi there, chatbot!",
mlflow.transformers.generate_signature_output(
conversational_pipeline, "Hi there, chatbot!"
),
)

# Log the pipeline
with mlflow.start_run():
model_info = mlflow.transformers.log_model(
transformers_model=conversational_pipeline,
name="chatbot",
task="conversational",
signature=signature,
input_example="A clever and witty question",
)

# Load the saved pipeline as pyfunc
chatbot = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)

# Ask the chatbot a question
response = chatbot.predict("What is machine learning?")

print(response)

# >> [It's a new thing that's been around for a while.]

使用 Transformer 管道保存提示模板

注意

此功能仅在 MLflow 2.10.0 及更高版本中可用。

MLflow 支持为某些管道类型指定提示模板

提示模板是用于在 pyfunc 推理之前格式化用户输入的字符串。要指定提示模板,请在调用 mlflow.transformers.save_model()mlflow.transformers.log_model() 时使用 prompt_template 参数。提示模板必须是一个带有单个格式占位符 {prompt} 的字符串。

例如

import mlflow
from transformers import pipeline

# Initialize a pipeline. `distilgpt2` uses a "text-generation" pipeline
generator = pipeline(model="distilgpt2")

# Define a prompt template
prompt_template = "Answer the following question: {prompt}"

# Save the model
mlflow.transformers.save_model(
transformers_model=generator,
path="path/to/model",
prompt_template=prompt_template,
)

当模型随后使用 mlflow.pyfunc.load_model() 加载时,提示模板将用于格式化用户输入,然后再将其传递到管道中

import mlflow

# Load the model with pyfunc
model = mlflow.pyfunc.load_model("path/to/model")

# The prompt template will be used to format this input, so the
# string that is passed to the text-generation pipeline will be:
# "Answer the following question: What is MLflow?"
model.predict("What is MLflow?")
注意

具有提示模板的 text-generation 管道将默认将 return_full_text 管道参数设置为 False。这是为了防止模板向用户显示,这可能会导致混淆,因为它不是他们原始输入的一部分。要覆盖此行为,可以通过 paramsreturn_full_text 设置为 True,或者通过将其包含在 log_model() 中的 model_config 字典中。有关如何执行此操作的更多详细信息,请参见 此部分

有关更深入的指南,请查看 提示模板笔记本

将 model_config 和模型签名参数用于推理

对于 transformers 推理,有两种方法可以将其他参数传递给管道。

  • 在保存/记录模型时使用 model_config。 (可选)在调用 load_model 时指定 model_config
  • 在调用 predict() 时指定推理时的参数

使用 model_config 来控制如何加载模型以及如何为所有输入样本执行推理。除非使用相同的参数指示 ModelSignature,否则无法在 predict() 时覆盖 model_config 中的配置。

另一方面,使用带有参数模式的 ModelSignature,以允许下游使用者提供可能需要计算其特定样本的预测的其他推理参数。

注意

如果在记录模型时保存了带有参数的 model_configModelSignature,则它们都将用于推理。ModelSignature 中的默认参数将覆盖 model_config 中的参数。如果在推理时提供了额外的 params,则它们优先于所有参数。我们建议将 model_config 用于通常运行模型所有样本所需的那些参数。然后,添加带有参数的 ModelSignature,用于您希望下游使用者指示每个样本的那些额外参数。

  • 使用 model_config
import mlflow
from mlflow.models import infer_signature
from mlflow.transformers import generate_signature_output
import transformers

architecture = "mrm8488/t5-base-finetuned-common_gen"
model = transformers.pipeline(
task="text2text-generation",
tokenizer=transformers.T5TokenizerFast.from_pretrained(architecture),
model=transformers.T5ForConditionalGeneration.from_pretrained(architecture),
)
data = "pencil draw paper"

# Infer the signature
signature = infer_signature(
data,
generate_signature_output(model, data),
)

# Define an model_config
model_config = {
"num_beams": 5,
"max_length": 30,
"do_sample": True,
"remove_invalid_values": True,
}

# Saving model_config with the model
mlflow.transformers.save_model(
model,
path="text2text",
model_config=model_config,
signature=signature,
)

pyfunc_loaded = mlflow.pyfunc.load_model("text2text")
# model_config will be applied
result = pyfunc_loaded.predict(data)

# overriding some inference configuration with different values
pyfunc_loaded = mlflow.pyfunc.load_model(
"text2text", model_config=dict(do_sample=False)
)
注意

请注意,在前面的示例中,用户在调用 predict 时无法覆盖配置 do_sample

  • 在推理时指定参数
import mlflow
from mlflow.models import infer_signature
from mlflow.transformers import generate_signature_output
import transformers

architecture = "mrm8488/t5-base-finetuned-common_gen"
model = transformers.pipeline(
task="text2text-generation",
tokenizer=transformers.T5TokenizerFast.from_pretrained(architecture),
model=transformers.T5ForConditionalGeneration.from_pretrained(architecture),
)
data = "pencil draw paper"

# Define an model_config
model_config = {
"num_beams": 5,
"remove_invalid_values": True,
}

# Define the inference parameters params
inference_params = {
"max_length": 30,
"do_sample": True,
}

# Infer the signature including params
signature_with_params = infer_signature(
data,
generate_signature_output(model, data),
params=inference_params,
)

# Saving model with signature and model config
mlflow.transformers.save_model(
model,
path="text2text",
model_config=model_config,
signature=signature_with_params,
)

pyfunc_loaded = mlflow.pyfunc.load_model("text2text")

# Pass params at inference time
params = {
"max_length": 20,
"do_sample": False,
}

# In this case we only override max_length and do_sample,
# other params will use the default one saved on ModelSignature
# or in the model configuration.
# The final params used for prediction is as follows:
# {
# "num_beams": 5,
# "max_length": 20,
# "do_sample": False,
# "remove_invalid_values": True,
# }
result = pyfunc_loaded.predict(data, params=params)

管道与组件日志记录

transformers 风格有两种不同的主要机制来保存和加载模型:管道和组件。

注意

使用自定义代码保存 transformers 模型(即需要 trust_remote_code=True 的模型)需要 transformers >= 4.26.0

管道

管道在 Transformers 库的上下文中是高级对象,它结合了预训练模型和分词器(以及其他组件,具体取决于任务类型)以执行特定任务。它们抽象了使用模型中涉及的大部分预处理和后处理工作。

例如,文本分类管道将处理文本的分词、将 token 传递到模型中,然后解释 logits 以生成人类可读的分类。

使用 MLflow 记录管道时,您实际上是在保存这种高级抽象,可以通过最小的设置加载并直接用于推理。这非常适合端到端任务,其中预处理和后处理步骤对于手头的任务是标准的。

组件

组件是指可以构成管道的各个部分,例如模型本身、分词器以及特定任务所需的任何其他处理器、提取器或配置。使用 MLflow 记录组件可以实现更大的灵活性和定制。当您的项目需要更好地控制预处理和后处理步骤时,或者当您需要以与管道抽象调用它们的方式不同的方式以定制方式访问各个组件时,您可以记录各个组件。

例如,如果您有自定义分词器,或者您想对模型输出应用一些特殊的后处理,则可以单独记录组件。加载组件时,您可以随后使用您的自定义组件重建管道,或者根据需要单独使用组件。

注意

MLflow 默认使用 500 MB 的 max_shard_sizemlflow.transformers.save_model()mlflow.transformers.log_model() API 中保存模型对象。您可以使用环境变量 MLFLOW_HUGGINGFACE_MODEL_MAX_SHARD_SIZE 来覆盖该值。

注意

对于基于组件的日志记录,提交的 dict 中必须满足的唯一要求是提供了模型。dict 的所有其他元素都是可选的。

记录基于组件的模型

以下示例显示了通过特定命名组件的字典映射记录 transformers 模型的组件。提交的字典中的键名必须在集合中:{"model", "tokenizer", "feature_extractor", "image_processor"}。处理器类型对象(某些图像处理器、音频处理器和多模态处理器)必须使用 mlflow.transformers.save_model()mlflow.transformers.log_model() API 中的 processor 参数显式保存。

记录后,这些组件将自动插入到为执行的任务而设计的适当 Pipeline 类型中并返回,可以进行推理。

注意

可以通过在 load_model() API 中将属性 return_type 设置为“components”来以其原始结构(字典)检索记录的组件。

注意

并非所有模型类型都与通过组件元素进行的管道 API 构造函数兼容。不兼容的模型将引发 MLflowException 错误,指出该模型缺少 name_or_path 属性。如果发生这种情况,请直接通过 transformers.pipeline(<repo name>) API 构造模型,并直接保存管道对象。

import mlflow
import transformers

task = "text-classification"
architecture = "distilbert-base-uncased-finetuned-sst-2-english"
model = transformers.AutoModelForSequenceClassification.from_pretrained(architecture)
tokenizer = transformers.AutoTokenizer.from_pretrained(architecture)

# Define the components of the model in a dictionary
transformers_model = {"model": model, "tokenizer": tokenizer}

# Log the model components
with mlflow.start_run():
model_info = mlflow.transformers.log_model(
transformers_model=transformers_model,
name="text_classifier",
task=task,
)

# Load the components as a pipeline
loaded_pipeline = mlflow.transformers.load_model(
model_info.model_uri, return_type="pipeline"
)

print(type(loaded_pipeline).__name__)
# >> TextClassificationPipeline

loaded_pipeline(["MLflow is awesome!", "Transformers is a great library!"])

# >> [{'label': 'POSITIVE', 'score': 0.9998478889465332},
# >> {'label': 'POSITIVE', 'score': 0.9998030066490173}]

保存管道并加载组件

某些用例可以从将解决方案定义为管道的简易性中受益,但需要组件级别的访问权限来执行基于微服务的部署策略,其中在不包含模型本身的容器上执行预处理/后处理。对于此范例,可以将其组成部分加载管道,如下所示。

import transformers
import mlflow

translation_pipeline = transformers.pipeline(
task="translation_en_to_fr",
model=transformers.T5ForConditionalGeneration.from_pretrained("t5-small"),
tokenizer=transformers.T5TokenizerFast.from_pretrained(
"t5-small", model_max_length=100
),
)

with mlflow.start_run():
model_info = mlflow.transformers.log_model(
transformers_model=translation_pipeline,
name="french_translator",
)

translation_components = mlflow.transformers.load_model(
model_info.model_uri, return_type="components"
)

for key, value in translation_components.items():
print(f"{key} -> {type(value).__name__}")

# >> task -> str
# >> model -> T5ForConditionalGeneration
# >> tokenizer -> T5TokenizerFast

response = translation_pipeline("MLflow is great!")

print(response)

# >> [{'translation_text': 'MLflow est formidable!'}]

reconstructed_pipeline = transformers.pipeline(**translation_components)

reconstructed_response = reconstructed_pipeline(
"transformers makes using Deep Learning models easy and fun!"
)

print(reconstructed_response)

# >> [{'translation_text': "Les transformateurs rendent l'utilisation de modèles Deep Learning facile et amusante!"}]

自动元数据和 ModelCard 日志记录

为了为保存的模型提供尽可能多的信息,transformers 风格将自动获取为 HuggingFace Hub 上存储的卡片的任何已保存的模型或管道的 ModelCard。此卡将作为模型工件的一部分记录,可在与 MLmodel 文件和存储的模型对象相同的目录级别查看。

除了 ModelCard 之外,构成任何管道的组件(如果在保存命名组件字典,则为各个组件)将存储其源类型。模型类型、管道类型、任务以及任何补充组件(例如 TokenizerImageProcessor)的类也将存储在 MLmodel 文件中。

为了保留附加到托管在 huggingface hub 上的任何模型的使用上的任何法律要求,在记录 transformers 模型时,会尽最大努力检索和持久化任何许可信息。将在模型目录的根目录中生成一个文件 (LICENSE.txt)。在此文件中,您将找到声明的许可证的副本、适用于模型使用的常见许可证类型的名称(即“apache-2.0”、“mit”),或者,如果上传模型存储库时从未向 huggingface hub 提交许可信息,则找到存储库的链接供您用来确定关于模型使用的任何限制。

注意

模型许可证信息是在 MLflow 2.10.0 中引入的。以前的版本不包含模型的许可证信息。

自动签名推理

对于支持 pyfunc 的管道,有 3 种方法可以将模型签名附加到 MLmodel 文件。

  • 通过将有效的 ModelSignature 设置为 signature 属性来显式提供模型签名。这可以通过辅助实用程序 mlflow.transformers.generate_signature_output() 生成

  • 提供一个 input_example。将推断签名并验证它是否与适当的输入类型匹配。将通过自动执行推理来验证输出类型(如果该模型是 pyfunc 支持的类型)。

  • 什么都不做。transformers 风格将自动应用管道类型支持的适当的通用签名(仅适用于单个实体;不会推断集合)。

通过覆盖 Pytorch dtype 来扩展推理

用于降低 transformers 管道中 pytorch 模型的总内存压力的常见配置是修改处理数据类型。这可以通过在创建 Pipeline 时设置 torch_dtype 参数来实现。有关管道配置的可调整参数的完整参考,请参见 训练文档

注意

此功能在 transformers < 4.26.x 版本中不存在

为了将这些配置应用于已保存或已记录的运行,有两种选择

  • 保存一个管道,并将 torch_dtype 参数设置为您选择的编码类型。

示例

import transformers
import torch
import mlflow

task = "translation_en_to_fr"

my_pipeline = transformers.pipeline(
task=task,
model=transformers.T5ForConditionalGeneration.from_pretrained("t5-small"),
tokenizer=transformers.T5TokenizerFast.from_pretrained(
"t5-small", model_max_length=100
),
framework="pt",
)

with mlflow.start_run():
model_info = mlflow.transformers.log_model(
transformers_model=my_pipeline,
name="my_pipeline",
torch_dtype=torch.bfloat16,
)

# Illustrate that the torch data type is recorded in the flavor configuration
print(model_info.flavors["transformers"])

结果

{'transformers_version': '4.28.1',
'code': None,
'task': 'translation_en_to_fr',
'instance_type': 'TranslationPipeline',
'source_model_name': 't5-small',
'pipeline_model_type': 'T5ForConditionalGeneration',
'framework': 'pt',
'torch_dtype': 'torch.bfloat16',
'tokenizer_type': 'T5TokenizerFast',
'components': ['tokenizer'],
'pipeline': 'pipeline'}
  • 加载模型时指定 torch_dtype 参数,以覆盖在日志记录或保存期间设置的任何值。

示例

import transformers
import torch
import mlflow

task = "translation_en_to_fr"

my_pipeline = transformers.pipeline(
task=task,
model=transformers.T5ForConditionalGeneration.from_pretrained("t5-small"),
tokenizer=transformers.T5TokenizerFast.from_pretrained(
"t5-small", model_max_length=100
),
framework="pt",
)

with mlflow.start_run():
model_info = mlflow.transformers.log_model(
transformers_model=my_pipeline,
name="my_pipeline",
torch_dtype=torch.bfloat16,
)

loaded_pipeline = mlflow.transformers.load_model(
model_info.model_uri, return_type="pipeline", torch_dtype=torch.float64
)

print(loaded_pipeline.torch_dtype)

结果

torch.float64
注意

MLflow 2.12.1 稍微更改了 torch_dtype 提取逻辑。以前它依赖于管道实例的 torch_dtype 属性,但现在它是通过 dtype 属性从底层模型中提取的。这使 MLflow 能够在管道实例化后捕获模型的 dtype 更改。

注意

以“组件”模式(使用字典声明组件)记录或保存模型不支持为构造的管道设置数据类型。如果您需要覆盖如何编码数据的默认行为,请保存或记录 pipeline 对象。

注意

加载为 python_function (pyfunc) 模型风格 时不支持覆盖管道的数据类型。在 save_model()log_model() 期间为 torch_dtype 设置的值在加载为 pyfunc 时将保持不变。

音频管道的输入数据类型

请注意,将原始数据传递给音频管道(原始字节)需要同一有效库的两个单独元素。为了使用比特率转置和音频字节数据转换为 numpy nd.array 格式,需要库 ffmpeg。直接从 pypi 安装此包 (pip install ffmpeg) 不会安装使 ffmpeg 函数所需的底层 c dll。请查阅 ffmpeg 网站上的文档,以获取有关您的给定操作系统的指导。

音频管道类型,加载为 python_function (pyfunc) 模型风格 时,有三种可用的输入类型

  • str

字符串输入类型用于可由 pyfunc 模型的实例访问的 blob 引用(uri 位置)。由于在 Spark DataFrames 中处理大型 bytes 数据的固有限制,此输入模式在 Spark 中进行音频推理的大型批量处理时非常有用。确保您已在运行 pyfunc 模型的环境中安装了 ffmpeg,以便使用基于 uri 的 str 输入推理。如果此包未正确安装(从 pypiffmpeg 二进制文件),则会在推理时引发异常。

警告

如果使用 uri (str) 作为您打算通过 MLflow 模型服务器托管以进行实时推理的 pyfunc 模型的输入类型,则在记录或保存模型时,您必须指定自定义模型签名。bytes 的默认签名输入值类型在 MLflow 模型服务中将强制将 uri 字符串转换为 bytes,这将导致服务进程引发异常,指出声音文件已损坏。

下面显示了一个用于音频模型的适当的基于 uri 的输入模型签名示例

from mlflow.models import infer_signature
from mlflow.transformers import generate_signature_output

url = "https://www.mywebsite.com/sound/files/for/transcription/file111.mp3"
signature = infer_signature(url, generate_signature_output(my_audio_pipeline, url))
with mlflow.start_run():
mlflow.transformers.log_model(
transformers_model=my_audio_pipeline,
name="my_transcriber",
signature=signature,
)
  • bytes

这是音频文件的默认序列化格式。这是最容易使用的格式,因为 Pipeline 实现将使用 ffmpeg (如果使用此格式,则为必需的依赖项)自动将音频比特率从文件转换为 Pipeline 中底层模型所需的比特率。当直接使用管道的 pyfunc 表示形式时(不是通过服务),声音文件可以直接作为 bytes 传递,而无需任何修改。通过服务使用时,bytes 数据必须进行 base64 编码。

  • np.ndarray

此输入格式要求在转换为 numpy.ndarray 之前已设置比特率(即,通过使用像 librosapydub 这样的包),并且模型已使用将 np.ndarray 格式用于输入的签名保存。

注意

用于服务的音频模型如果打算使用 np.ndarray 格式的预格式化音频,则必须使用反映此架构的签名配置保存模型。否则会导致类型转换错误,因为音频 transformers 管道的默认签名设置为期望 binary (bytes) 数据。服务终结点不能接受类型的联合,因此特定的模型实例必须选择一种或另一种作为允许的输入类型。

MLflow Transformers 风格中的 PEFT 模型

PEFT 是 HuggingFace🤗 开发的库,它为 HuggingFace Hub 上提供的预训练模型提供了各种优化方法。借助 PEFT,您可以轻松地应用各种优化技术,例如 LoRA 和 QLoRA,以降低微调 Transformers 模型的成本。

例如,LoRA(低秩自适应)是一种通过低秩分解来近似微调过程的权重更新与两个较小矩阵的方法。LoRA 通常将要训练的参数数量缩减到仅完整模型微调的 0.01%~几个%(具体取决于配置),这显着加快了微调过程并减少了内存占用,因此您甚至可以在一小时内在单个 Nvidia A10G GPU 上训练 Mistral/Llama2 7B 模型。通过使用 PEFT,您只需几行代码即可将 LoRA 应用于您的 Transformers 模型

from peft import LoraConfig, get_peft_model

base_model = AutoModelForCausalLM.from_pretrained(...)
lora_config = LoraConfig(...)
peft_model = get_peft_model(base_model, lora_config)

在 MLflow 2.11.0 中,我们引入了对 MLflow Transformers 风格中跟踪 PEFT 模型的支持。您可以使用与其他 Transformers 模型相同的 API 记录和加载 PEFT 模型,例如 mlflow.transformers.log_model()mlflow.transformers.load_model()

import mlflow
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "databricks/dolly-v2-7b"
base_model = AutoModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

peft_config = LoraConfig(...)
peft_model = get_peft_model(base_model, peft_config)

with mlflow.start_run():
# Your training code here
...

# Log the PEFT model
model_info = mlflow.transformers.log_model(
transformers_model={
"model": peft_model,
"tokenizer": tokenizer,
},
name="peft_model",
)

# Load the PEFT model
loaded_model = mlflow.transformers.load_model(model_info.model_uri)

MLflow 中的 PEFT 模型教程

有关如何在 MLflow 中使用 PEFT 的更深入指南,请查看教程使用 QLoRA、MLflow 和 PEFT 微调开源 LLM

已保存 PEFT 模型的格式

保存 PEFT 模型时,MLflow 仅保存 PEFT 适配器和配置,但不保存基本模型的权重。这与 Transformer 的 save_pretrained() 方法的行为相同,并且在存储空间和日志记录延迟方面非常高效。一个区别是,MLflow 还会在模型元数据中保存 HuggingFace Hub 存储库名称和基本模型的版本,以便它可以在加载 PEFT 模型时加载相同的基本模型。具体来说,以下工件保存在 MLflow 中用于 PEFT 模型

  • /peft 目录下的 PEFT 适配器权重。
  • /peft 目录下作为 JSON 文件的 PEFT 配置。
  • MLModel 元数据文件中基本模型的 HuggingFace Hub 存储库名称和提交哈希。

MLflow 中 PEFT 模型的限制

由于 PEFT 模型的模型保存/加载行为与 save_pretrained=False 的行为类似,因此 相同的注意事项适用于 PEFT 模型。例如,基本模型权重可能会在 HuggingFace Hub 中被删除或变为私有,并且 PEFT 模型无法注册到旧版 Databricks 工作区模型注册表。

要保存 PEFT 模型的基本模型权重,您可以使用 mlflow.transformers.persist_pretrained_model() API。这将从 HuggingFace Hub 下载基本模型权重并将其保存到工件位置,从而更新给定 PEFT 模型的元数据。有关此 API 的详细用法,请参阅此部分