🤗 MLflow 中的 Transformers

注意

transformers 风格正在积极开发中,并被标记为实验性。公共API可能会发生变化,随着功能的增加,可能会添加新功能。

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

本页详细解释了 MLflow transformers 风格的特性和配置。关于 MLflow 的 Transformer 集成的一般介绍,请参阅 MLflow Transformers Flavor 页面。

将 Transformers 模型加载为 Python 函数

支持的 Transformers 管道类型

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

备注

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

注意

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

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

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]

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.ndarray

List[str]

音频分类

bytes*****, str, 或 np.ndarray

pd.DataFrame (dtypes: {‘label’: str, ‘score’: double})

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

** 这些输入的集合也可以被传递。引用表必须是一个json编码的字典(即 {‘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,
        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 管道保存提示模板

备注

此功能仅在 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 pipeline 参数 设置为 False。这是为了防止模板内容展示给用户,因为这可能会导致混淆,因为它不是用户原始输入的一部分。要覆盖此行为,可以通过 paramsreturn_full_text 设置为 True,或者在 log_model() 中包含在 model_config 字典中。有关如何执行此操作的更多详细信息,请参阅 此部分

要获取更深入的指南,请查看 Prompt Templating 笔记本!

使用 model_config 和 Model Signature 参数进行推理

对于 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 diferent 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)

流水线 vs. 组件日志

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

备注

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

管道

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

例如,文本分类管道将处理文本的标记化,将标记传递给模型,然后解释对数以生成人类可读的分类。

在使用 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,
        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 风格将自动获取任何已保存的模型或管道的 ModelCard,前提是该卡片存储在 HuggingFace 中心。该卡片将作为模型工件的一部分记录下来,可以在与 MLmodel 文件和存储的模型对象相同的目录级别查看。

除了 ModelCard 之外,构成任何 Pipeline(或保存命名组件字典时的各个组件)的组件将存储其源类型。模型类型、Pipeline 类型、任务以及任何辅助组件(如 TokenizerImageProcessor)的类也将存储在 MLmodel 文件中。

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

备注

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

自动签名推断

对于支持 pyfunc 的管道,有三种方法可以将模型签名附加到 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,
        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 的提取逻辑。以前它依赖于 pipeline 实例的 torch_dtype 属性,但现在它是通过 dtype 属性从底层模型中提取的。这使得 MLflow 能够在 pipeline 实例化后捕获模型 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网站 <https://ffmpeg.org/download.html>`_ 的文档,以获取关于您特定操作系统的指导。

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

  • str

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

警告

如果你使用一个 URI (str) 作为 pyfunc 模型的输入类型,并且你打算通过 MLflow Model Server 进行实时推理,你必须在记录或保存模型时指定一个自定义的模型签名。默认的签名输入值类型 bytes 将在 MLflow Model serving 中强制将 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

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

  • np.ndarray

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

备注

用于服务的音频模型如果打算使用预格式化的 np.ndarray 格式音频,必须将模型保存为具有反映此模式的签名配置。如果不这样做,将导致类型转换错误,因为音频转换器管道的默认签名设置为期望 binary (bytes) 数据。服务端点不能接受类型的联合,因此特定模型实例必须选择其中一种作为允许的输入类型。

MLflow Transformers 风格中的 PEFT 模型

警告

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

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

例如,LoRA (低秩适应) 是一种通过低秩分解近似微调过程中权重更新的方法,它使用两个较小的矩阵。LoRA通常将需要训练的参数数量缩小到全模型微调的0.01% ~ 几个%(取决于配置),这显著加速了微调过程并减少了内存占用,以至于你甚至可以在一个小时内`在单个Nvidia A10G GPU上训练一个Mistral/Llama2 7B模型 <../tutorials/fine-tuning/transformers-peft.html>`_。通过使用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,
        },
        artifact_path="peft_model",
    )

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

MLflow 教程中的 PEFT 模型

查看教程 使用 MLflow 和 PEFT 通过 QLoRA 微调开源 LLM 以获取更深入的指南,了解如何将 PEFT 与 MLflow 结合使用。

保存的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 Workspace模型注册表中。

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