MLflow 中的 🤗 Transformers

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

本页详细介绍了 MLflow transformers 风味的详细功能和配置。有关 MLflow Transformer 集成的通用介绍，请参阅 MLflow Transformers 风味页面。

• 将 Transformers 模型加载为 Python 函数
• 为 Transformer 管道保存提示模板
• 使用 model_config 和 Model Signature 参数进行推理
• 管道与组件记录
• 自动元数据和 ModelCard 记录
• 自动签名推断
• 通过覆盖 Pytorch dtype 扩展推理
• 音频管道的输入数据类型
• MLflow Transformers 风味中的 PEFT 模型

将 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]
文本到文本生成	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.ndarray	List[str]
音频分类	bytes*****, str, 或 np.ndarray	pd.DataFrame (dtypes: {'label': str, 'score': double})

* 也可以传递这些输入的集合。标准的必需键名为 'sequences' 和 'candidate_labels'，但这些可能会有所不同。请检查您使用的架构的输入要求，以确保提供了正确的字典键名。

** 也可以传递这些输入的集合。参考表必须是 json 编码的 dict（即 {'query': 'what did we sell most of?', '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 支持为某些管道类型指定提示模板

• feature-extraction (特征提取)
• fill-mask
• summarization
• text2text-generation
• text-generation (文本生成)

提示模板是用于在 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 和 Model Signature 参数进行推理

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

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

使用 model_config 来控制模型加载和所有输入样本的推理方式。model_config 中的配置在 predict() 时间是不可覆盖的，除非指示了一个具有相同参数的 ModelSignature。

另一方面，使用带有参数的 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 风味有两种不同的主要机制用于保存和加载模型：管道和组件。

注意

使用自定义代码（即需要 trust_remote_code=True 的模型）保存 transformers 模型需要 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(<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 风味将自动获取任何已保存的模型或管道的 ModelCard，前提是该模型或管道在 HuggingFace hub 上有存储的卡片。此卡片将作为模型伪影的一部分进行记录，位于与 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 中进行大量批量音频推理时，此输入模式非常有用，因为 Spark DataFrames 在处理大型 bytes 数据方面存在固有的局限性。确保在 pyfunc 模型运行的环境中安装了 ffmpeg，以便使用 str 输入的基于 URI 的推理。如果此包未正确安装（来自 pypi 和 ffmpeg 二进制文件），将在推理时抛出异常。

警告

如果您将 URI（str）作为 pyfunc 模型（您打算通过 MLflow Model Server 托管以进行实时推理）的输入类型，则必须在记录或保存模型时指定自定义模型签名。在 MLflow Model serving 中，默认签名输入值类型 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 (Low-Rank Adaptation) 是一种通过低秩分解将微调过程的权重更新近似为两个较小的矩阵的方法。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 模型支持。您可以使用与 S1. 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 模型教程

请查看教程 使用 QLoRA 和 MLflow 进行开源 LLM 微调，了解如何将 PEFT 与 MLflow 结合使用的更深入指南。

保存的 PEFT 模型格式

在保存 PEFT 模型时，MLflow 只保存 PEFT 适配器和配置，而不保存基础模型的权重。这与 Transformer 的 save_pretrained() 方法的行为相同，并且在存储空间和记录延迟方面都非常高效。一个区别是 MLflow 还会在模型元数据中保存基础模型的 HuggingFace Hub 存储库名称和版本，以便在加载 PEFT 模型时加载相同的基本模型。具体来说，MLflow 为 PEFT 模型保存了以下伪影：

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

MLflow 中 PEFT 模型的功能限制

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

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