跳到主要内容

在 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，而不是管道实例。MLflow 将检查检查点的模型元数据，并在不将模型权重加载到内存中的情况下记录模型权重。这样，您可以用最少的计算资源将一个巨大的多亿参数模型记录到 MLflow 中。

重要说明​

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

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

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

存储高效的模型日志记录​

通常，当 MLflow 记录机器学习模型时，它会将模型权重副本保存到工件存储中。然而，当您使用 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 下载模型权重并将其保存到工件位置。有关更多信息，请参阅 OSS 模型注册表或旧版工作区模型注册表部分。

注册仅引用模型以供生产​

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

然而，这可能不适用于生产用例，因为模型权重可能不可用或由于网络问题导致下载失败。MLflow 提供了在将引用模型注册到模型注册表时解决此问题的解决方案。

Databricks Unity Catalog​

将仅引用模型注册到 Databricks Unity Catalog 模型注册表无需额外步骤，与正常的模型注册过程相同。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 模型注册表或旧版工作区模型注册表​

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

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 无法重新加载模型。对于生产用例，建议在模型从开发或暂存阶段迁移到生产阶段之前，将模型权重副本保存到工件存储中。
  • HuggingFace Hub 访问：从 HuggingFace Hub 下载模型可能由于网络延迟或 HuggingFace Hub 服务状态而缓慢或不稳定。MLflow 不提供任何重试机制或稳健的错误处理来从 HuggingFace Hub 下载模型。因此，您不应将此功能用于最终的生产候选运行。

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