在 MLflow Transformers flavor 中使用大型模型
本指南中介绍的功能适用于熟悉 Transformers 和 MLflow 的高级用户。请在使用这些功能前了解其局限性和潜在风险。
MLflow Transformers flavor 允许您在 MLflow 中跟踪各种 Transformers 模型。然而，由于大型语言模型 (LLM) 的大小和内存需求，记录大型模型可能非常消耗资源。本指南概述了 MLflow 用于减少记录模型时的内存和磁盘使用量的功能，使您能够在资源受限的环境中处理大型模型。
概述
下表总结了使用 Transformers flavor 记录模型的不同方法。请注意，每种方法都有其特定的局限性和要求，如下节所述。
| 保存方法 | 描述 | 内存使用 | 磁盘使用 | 示例 |
|---|---|---|---|---|
| 标准 pipeline-based 记录 | 使用 pipeline 实例或 pipeline 组件字典记录模型。 | 高 | 高 | python |
| 内存高效的模型记录 | 通过指定本地检查点路径来记录模型，避免将模型加载到内存中。 | 低 | 高 | python |
| 存储高效的模型记录 | 通过保存对 HuggingFace Hub 存储库的引用来记录模型，而不是模型权重。 | 高 | 低 | python |
内存高效的模型记录
此方法于 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 参数传递，而不是 pipeline 实例。MLflow 将检查检查点的模型元数据，并在不加载模型权重的情况下记录它们。这样，您就可以用最少的计算资源将一个巨大的数十亿参数模型记录到 MLflow 中。
重要说明
在使用此功能时，请注意以下要求和限制：
- 检查点目录 **必须** 包含一个有效的 config.json 文件和模型权重文件。如果需要 tokenizer，其状态文件也必须存在于检查点目录中。您可以通过调用
tokenizer.save_pretrained("path/to/local/checkpoint")方法将 tokenizer 状态保存在您的检查点目录中。 - 您 **必须** 指定
task参数，并提供模型设计所对应的正确任务名称。 - 在此模式下，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 上底层存储库条目的引用；具体来说，当您记录组件或 pipeline 时，会存储存储库名称和模型权重的唯一提交哈希。当加载这种“仅引用”的模型时，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 Model Registry 或 Legacy Workspace Model Registry 部分。
为生产环境注册“仅引用”模型
使用上述任一优化方法记录的模型都是“仅引用”模型，这意味着模型权重未保存到 artifact 存储中，仅保存了对 HuggingFace Hub 存储库的引用。当您正常加载模型时，MLflow 将从 HuggingFace Hub 下载模型权重。
然而，这可能不适用于生产用例，因为模型权重可能不可用，或者由于网络问题导致下载失败。MLflow 提供了一种解决方案来解决在将“仅引用”模型注册到 Model Registry 时遇到的问题。
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 Model Registry 或 Legacy Workspace Model Registry
对于 OSS Model Registry 或 Databricks 中的旧版 Workspace Model Registry，您需要在注册模型之前手动将模型权重持久化到 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 模型，同时优化资源使用。