跳到主要内容

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 DataFrame 兼容的格式来进行服务和批量推理。

注意

某些 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() 将模型组件直接加载为 dict，并将 return_type 参数指定为“components”。

管道类型输入类型输出类型
指令文本生成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]
零样本分类*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 编码的字典（即 {'query': '我们卖得最多的是什么？', 'table': json.dumps(table_as_dict)}）

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

**** 您选择的模型的掩码语法将是该模型实现的特定语法。有些是“[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。这是为了防止模板向用户显示，这可能会导致混淆，因为它不是他们原始输入的一部分。要覆盖此行为，可以通过 params 将 return_full_text 设置为 True，或将其包含在 log_model() 中的 model_config 字典中。有关如何执行此操作的更多详细信息，请参阅此部分。

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

使用 model_config 和模型签名参数进行推理​

对于 transformers 推理，有两种方法可以将附加参数传递给管道。

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

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

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

注意

如果在记录模型时同时保存了 model_config 和带有参数的 ModelSignature，则它们都将用于推理。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 库的上下文中，是高级对象，它结合了预训练模型和分词器（以及其他组件，具体取决于任务类型）来执行特定任务。它们抽象了使用模型所涉及的大部分预处理和后处理工作。

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

使用 MLflow 记录管道时，您实际上是在保存这种高级抽象，它可以直接加载并用于推理，只需最少的设置。这对于预处理和后处理步骤对于手头任务是标准的端到端任务是理想的。

组件

组件是指可以构成管道的各个部分，例如模型本身、分词器，以及特定任务所需的任何附加处理器、提取器或配置。使用 MLflow 记录组件允许更大的灵活性和定制。当您的项目需要对预处理和后处理步骤有更多控制，或者当您需要以偏离管道抽象调用它们的方式以定制方式访问各个组件时，您可以记录单个组件。

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

注意

MLflow 默认使用 500 MB 的 max_shard_size 在 mlflow.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() 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!"}]

自动元数据和模型卡日志记录​

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

除了 ModelCard，构成任何管道的组件（或如果保存命名组件的字典，则为单个组件）将存储其源类型。模型类型、管道类型、任务以及任何补充组件（例如 Tokenizer 或 ImageProcessor）的类也将存储在 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 变化。

注意

以“components”模式（使用字典声明组件）记录或保存模型不支持为构造的管道设置数据类型。如果您需要覆盖数据编码的默认行为，请保存或记录 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，以便使用 str 基于 uri 的推理。如果未正确安装此包（包括从 pypi 和从 ffmpeg 二进制文件），则在推理时将抛出异常。

警告

如果将 uri (str) 用作您打算通过 MLflow 模型服务器托管用于实时推理的 pyfunc 模型的输入类型，则在记录或保存模型时必须指定自定义模型签名。在 MLflow 模型服务中，默认的签名输入值类型 bytes 将强制将 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

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

  • np.ndarray

此输入格式要求比特率在转换为 numpy.ndarray 之前已设置（例如，通过使用 librosa 或 pydub 等包），并且模型已使用输入使用 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 的详细用法，请参阅此部分。