跳到主要内容

在 MLflow Transformers 格式中处理大型模型

警告

本指南中描述的功能适用于熟悉 Transformers 和 MLflow 的高级用户。在使用这些功能之前，请了解与之相关的限制和潜在风险。

MLflow 的 Transformers 格式 允许您在 MLflow 中跟踪各种 Transformers 模型。但是，由于大型语言模型 (LLM) 等模型体积庞大且内存需求高，因此记录它们可能会占用大量资源。本指南概述了 MLflow 在记录模型时用于减少内存和磁盘使用量的功能，使您能够在资源受限的环境中处理大型模型。

概述​

下表总结了使用 Transformers 格式记录模型的不同方法。请注意，每种方法都有其特定的限制和要求，具体将在以下部分中介绍。

保存方法描述内存使用量磁盘使用量示例
普通管道式日志记录使用管道实例或管道组件字典记录模型。高高
import mlflow
import transformers

pipeline = transformers.pipeline(
    task="text-generation",
    model="meta-llama/Meta-Llama-3.1-70B",
)

with mlflow.start_run():
    mlflow.transformers.log_model(
        transformers_model=pipeline,
        name="model",
    )
内存高效模型日志记录通过指定本地检查点路径来记录模型，避免将模型加载到内存中。低高
import mlflow

with mlflow.start_run():
    mlflow.transformers.log_model(
        # Pass a path to local checkpoint as a model
        transformers_model="/path/to/local/checkpoint",
        # Task argument is required for this saving mode.
        task="text-generation",
        name="model",
    )
存储高效模型日志记录通过保存对 HuggingFace Hub 存储库的引用而不是模型权重来记录模型。高低
import mlflow
import transformers

pipeline = transformers.pipeline(
    task="text-generation",
    model="meta-llama/Meta-Llama-3.1-70B",
)

with mlflow.start_run():
    mlflow.transformers.log_model(
        transformers_model=pipeline,
        name="model",
        # Set save_pretrained to False to save storage space
        save_pretrained=False,
    )

内存高效模型日志记录​

此方法于 MLflow 2.16.1 中引入，允许您在不将模型加载到内存的情况下记录模型。

import mlflow

with mlflow.start_run():
    mlflow.transformers.log_model(
        # Pass a path to local checkpoint as a model to avoid loading the model instance
        transformers_model="path/to/local/checkpoint",
        # Task argument is required for this saving mode.
        task="text-generation",
        name="model",
    )

在上面的示例中，我们将本地模型检查点/权重的路径作为 `mlflow.transformers.log_model()` API 中的 model 参数传递，而不是传递管道实例。MLflow 将检查检查点的模型元数据，并记录模型权重而不将其加载到内存中。这样，您就可以以最少的计算资源将庞大的数十亿参数模型记录到 MLflow 中。

重要说明​

使用此功能时，请注意以下要求和限制：

  1. 检查点目录 **必须** 包含有效的 config.json 文件和模型权重文件。如果需要分词器，其状态文件也必须存在于检查点目录中。您可以通过调用 `tokenizer.save_pretrained("path/to/local/checkpoint")` 方法将分词器状态保存在您的检查点目录中。
  2. 您 **必须** 为 `task` 参数指定模型设计的适当任务名称。
  3. 在此模式下，MLflow 可能无法准确推断模型依赖项。有关管理模型依赖项的更多信息，请参阅MLflow 模型中的依赖项管理。
警告

确保指定正确的任务参数，因为不兼容的任务将导致模型在加载时失败。您可以在 HuggingFace Hub 上查看您的模型有效的任务类型。

存储高效模型日志记录​

通常，当 MLflow 记录 ML 模型时，它会将模型权重的副本保存到 artifact 存储中。但是，当您使用 HuggingFace Hub 的预训练模型，并且不打算在记录模型之前对其进行微调或其他操作时，这并非最佳选择。在这种非常常见的情况下，在开发提示、测试推理参数以及其他操作时，复制（通常非常大）的模型权重是冗余的，只是浪费存储空间。

为了解决这个问题，MLflow 2.11.0 在 `mlflow.transformers.save_model()` 和 `mlflow.transformers.log_model()` API 中引入了一个新的参数 `save_pretrained`。当此参数设置为 `False` 时，MLflow 将跳过保存预训练模型权重，而是存储对 HuggingFace Hub 上底层存储库条目的引用；具体来说，在记录您的组件或管道时，将存储存储库名称和模型权重的唯一提交哈希。加载此类“仅引用”模型时，MLflow 将从保存的元数据中检查存储库名称和提交哈希，然后从 HuggingFace Hub 下载模型权重，或使用您 HuggingFace 本地缓存目录中的本地缓存模型。

这是使用 `save_pretrained` 参数记录模型的示例

import transformers

pipeline = transformers.pipeline(
    task="text-generation",
    model="meta-llama/Meta-Llama-3.1-70B",
    torch_dtype="torch.float16",
)

with mlflow.start_run():
    mlflow.transformers.log_model(
        transformers_model=pipeline,
        name="model",
        # Set save_pretrained to False to save storage space
        save_pretrained=False,
    )

在上面的示例中，MLflow 不会保存 **Llama-3.1-70B** 模型权重的副本，而是将以下元数据作为对 HuggingFace Hub 模型的引用进行记录。这将节省大约 150GB 的存储空间，并显著减少在开发过程中发起的每次运行的日志记录延迟。

通过导航到 MLflow UI，您可以看到记录的模型及其存储库 ID 和提交哈希。

flavors:
    ...
    transformers:
        source_model_name: meta-llama/Meta-Llama-3.1-70B-Instruct
        source_model_revision: 33101ce6ccc08fa6249c10a543ebfcac65173393
        ...

在生产部署之前，您可能希望持久化模型权重而不是存储库引用。为此，您可以使用 `mlflow.transformers.persist_pretrained_model()` API 从 HuggingFace Hub 下载模型权重并将其保存到 artifact 位置。有关更多信息，请参阅OSS 模型注册表或旧版工作区模型注册表部分。

为生产环境注册仅引用模型​

使用上述任一优化方法记录的模型都是“仅引用”模型，这意味着模型权重未保存到 artifact 存储中，仅保存了对 HuggingFace Hub 存储库的引用。当您正常加载模型时，MLflow 将从 HuggingFace Hub 下载模型权重。

但是，这可能不适用于生产用例，因为模型权重可能不可用，或者由于网络问题下载可能会失败。MLflow 提供了解决方案来解决在将仅引用模型注册到模型注册表时出现的问题。

Databricks Unity Catalog​

将仅引用模型注册到 Databricks Unity Catalog Model Registry **不需要** 比正常模型注册过程额外的步骤。MLflow 会自动将模型权重与模型元数据一起下载并注册到 Unity Catalog。

import mlflow

mlflow.set_registry_uri("databricks-uc")

# Log the repository ID as a model. The model weight will not be saved to the artifact store
with mlflow.start_run():
    model_info = mlflow.transformers.log_model(
        transformers_model="meta-llama/Meta-Llama-3.1-70B-Instruct",
        name="model",
    )

# When registering the model to Unity Catalog Model Registry, MLflow will automatically
# persist the model weight files. This may take a several minutes for large models.
mlflow.register_model(model_info.model_uri, "your.model.name")

OSS 模型注册表或旧版工作区模型注册表​

对于 OSS 模型注册表或 Databricks 中的旧版工作区模型注册表，您需要先手动将模型权重持久化到 artifact 存储，然后再注册模型。您可以使用 `mlflow.transformers.persist_pretrained_model()` API 从 HuggingFace Hub 下载模型权重并将其保存到 artifact 位置。此过程 **不需要重新记录模型**，而是可以高效地就地更新现有模型和元数据。

import mlflow

# Log the repository ID as a model. The model weight will not be saved to the artifact store
with mlflow.start_run():
    model_info = mlflow.transformers.log_model(
        transformers_model="meta-llama/Meta-Llama-3.1-70B-Instruct",
        name="model",
    )

# Before registering the model to the non-UC model registry, persist the model weight
# from the HuggingFace Hub to the artifact location.
mlflow.transformers.persist_pretrained_model(model_info.model_uri)

# Register the model
mlflow.register_model(model_info.model_uri, "your.model.name")

跳过预训练模型权重保存的注意事项​

虽然这些功能对于节省计算资源和存储空间以记录大型模型很有用，但也有一些需要注意的注意事项：

  • 模型可用性更改：如果您使用的是其他用户的存储库中的模型，该模型可能会在 HuggingFace Hub 上被删除或变为私有。在这种情况下，MLflow 无法重新加载模型。对于生产用例，建议在将模型从开发或暂存环境移至生产环境之前，将模型权重的副本保存到 artifact 存储中。
  • HuggingFace Hub 访问：由于网络延迟或 HuggingFace Hub 服务状态，从 HuggingFace Hub 下载模型可能会缓慢或不稳定。MLflow 不提供从 HuggingFace Hub 下载模型的任何重试机制或健壮的错误处理。因此，您不应依赖此功能来运行最终的生产候选版本。

通过了解这些方法及其局限性，您可以在 MLflow 中有效地处理大型 Transformers 模型，同时优化资源使用。