跳到主要内容

🤗 MLflow 中的 Transformers

警告

transformers flavor 正在积极开发中，并标记为实验性功能。公共 API 可能会发生变化，随着向 flavor 添加更多功能，可能会添加新功能。

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

本页详细介绍了 MLflow transformers flavor 的特性和配置。有关 MLflow 的 Transformer 集成的总体介绍，请参阅 MLflow Transformers Flavor 页面。

将 Transformers 模型加载为 Python 函数​

支持的 Transformers 管道类型​

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

注意

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

注意

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

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

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

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

重要

这些管道的 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]
标记分类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': 'what did we sell most of?', 'table': json.dumps(table_as_dict)}）

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

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

***** 如果在 MLflow 模型 Serving 中使用 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,
        artifact_path="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 管道保存 Prompt 模板​

注意

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

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

Prompt 模板是用于在 pyfunc 推理之前格式化用户输入的字符串。要指定 prompt 模板，请在调用 mlflow.transformers.save_model() 或 mlflow.transformers.log_model() 时使用 prompt_template 参数。prompt 模板必须是一个包含单个格式占位符 {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() 加载时，prompt 模板将用于在将用户输入传递到管道之前对其进行格式化

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?")
注意

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

有关更深入的指南，请查看 Prompt Templating notebook！

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

对于 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 flavor 有两种主要机制用于保存和加载模型：管道和组件。

注意

保存包含自定义代码的 transformers 模型（即需要 trust_remote_code=True 的模型）需要 transformers >= 4.26.0。

管道

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

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

使用 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,
        artifact_path="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,
        artifact_path="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 flavor 将自动为任何已保存的在 HuggingFace hub 上有存储卡片的模型或管道获取 ModelCard。此卡片将作为模型 artifact 的一部分进行记录，与 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 flavor 将自动应用管道类型支持的适当通用签名（仅适用于单个实体；不推断集合）。

通过覆盖 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,
        artifact_path="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,
        artifact_path="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) 模型 flavor 时，不支持覆盖管道的数据类型。在 save_model() 或 log_model() 期间为 torch_dtype 设置的值在加载为 pyfunc 时将保持不变。

音频管道的输入数据类型​

请注意，将原始数据传递到音频管道（原始字节）需要同一个有效库的两个独立元素。为了使用比特率转置并将音频字节数据转换为 numpy nd.array 格式，需要库 ffmpeg。直接从 pypi 安装此包 (pip install ffmpeg) 并不会安装使 ffmpeg 正常工作所需的基础 c dll。有关您所用操作系统的指导，请查阅 ffmpeg 网站上的文档。

音频管道类型在加载为 python_function (pyfunc) 模型 flavor 时，有三种可用的输入类型

  • str

字符串输入类型适用于 pyfunc 模型实例可以访问的 blob 引用 (uri 位置)。由于在 Spark DataFrames 中处理大型 bytes 数据的固有局限性，这种输入模式在 Spark 中进行音频推理的大规模批处理时非常有用。请确保在 pyfunc 模型运行的环境中安装了 ffmpeg，以便使用基于 str 输入 uri 的推理。如果此软件包未正确安装（无论是从 pypi 还是从 ffmpeg 二进制文件），将在推理时抛出异常。

警告

如果您打算通过 MLflow Model Server 托管 pyfunc 模型以进行实时推理，并使用 uri (str) 作为输入类型，则在记录或保存模型时必须指定自定义模型签名。在 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,
        artifact_path="my_transcriber",
        signature=signature,
    )
  • bytes

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

  • np.ndarray

np.ndarray

注意

用于 serving 并打算使用 np.ndarray 格式的预格式化音频的音频模型必须在保存模型时配置签名以反映此模式。否则将导致类型转换错误，因为音频 transformers 管道的默认签名被设置为期望 binary (bytes) 数据。serving 端点不能接受类型的联合，因此特定的模型实例必须选择其中一种作为允许的输入类型。

MLflow Transformers flavor 中的 PEFT 模型​

警告

PEFT 模型在 MLflow 2.11.0 及更高版本中受支持，目前仍处于实验阶段。API 和行为可能在未来版本中发生变化。此外，PEFT 库正在积极开发中，因此并非所有功能和适配器类型都可能在 MLflow 中得到支持。

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 flavor 中跟踪 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,
        },
        artifact_path="peft_model",
    )

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

MLflow 中的 PEFT 模型教程​

有关如何将 PEFT 与 MLflow 一起使用的更深入指南，请查看教程 使用 MLflow 和 PEFT 通过 QLoRA 微调开源 LLM，

保存的 PEFT 模型格式​

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

  • 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 下载基础模型权重并将其保存到 artifact 位置，同时更新给定 PEFT 模型的元数据。有关此 API 的详细用法，请参阅本节。