使用 MLflow 进行深度学习（第二部分）

在深度学习领域，在私有数据集上微调预训练的大型语言模型 (LLM) 是提高模型对特定任务相关性的绝佳自定义选项。这种做法不仅很常见，而且对于开发专门的模型至关重要，尤其是在文本分类和摘要等任务中。

在这种情况下，MLflow 等工具非常宝贵。MLflow 等跟踪工具可帮助确保训练过程的每一个方面——指标、参数和工件——都可以被可重复地跟踪和记录，从而能够分析、比较和共享调优迭代。

在这篇博文中，我们将使用 MLflow 2.12 和 新近推出的 MLflow 深度学习功能 来跟踪微调用于文本分类的大型语言模型的所有重要方面，包括使用自动记录训练检查点以简化训练恢复过程。

用例：微调 Transformer 模型进行文本分类

我们在这篇博文中使用的示例场景利用了 unfair-TOS 数据集。

在当今世界，很难找到一个不附带具有法律约束力的服务条款的服务、平台，甚至是消费品。这些百科全书式的大型协议，充满了晦涩的法律术语和有时令人费解的细节，内容如此之多，以至于大多数人根本不看就接受了。然而，报告表明，有时其中会嵌入一些可疑的不公平条款。

通过机器学习 (ML) 处理服务条款 (TOS) 协议中的不公平条款，在当前对影响消费者的法律协议的透明度和公平性有迫切需求的背景下尤为重要。考虑一下示例 TOS 协议中的以下条款：**“我们可能会不时修订这些条款。更改不会追溯生效，并且条款的最新版本将始终...”** 此条款规定，服务提供商可以随时以任何理由暂停或终止服务，无论是否提前通知。大多数人会认为这是很不公平的。

虽然这句话深埋在一个相当密集的文档中，但 ML 算法不受人类在仔细审查文本和识别可能看起来有些不公平的条款时所产生的疲劳的困扰。通过自动识别潜在的不公平条款，基于 Transformer 的深度学习 (DL) 模型可以帮助保护消费者免受剥削性做法的侵害，确保更好地遵守法律标准，并建立服务提供商与用户之间的信任。

一个基础的预训练 Transformer 模型，如果没有专门的微调，在准确识别不公平的服务条款时会面临几个挑战。首先，它缺乏理解复杂法律语言所必需的领域特定知识。其次，它的训练目标过于宽泛，无法捕捉法律分析所需的细致解释。最后，它可能无法有效地识别决定合同条款公平性的细微上下文含义，因此在此专门任务上的效率较低。

使用提示工程来识别闭源大型语言模型的不公平服务条款可能会非常昂贵。这种方法需要大量的反复试验来完善提示，但又无法调整底层模型机制。每次迭代都会消耗大量的计算资源，尤其是在使用 少样本提示 时，这会导致成本不断增加，而无法保证准确性或有效性的相应提高。

在这种情况下，使用 **RoBERTa-base** 模型尤其有效，前提是它经过了微调。该模型足够强大，可以处理复杂的任务，如在文本中区分嵌入式指令，同时它又足够紧凑，可以在适度的硬件上进行微调，例如 Nvidia T4 GPU。

什么是 PEFT？

参数高效微调 (PEFT) 方法的优势在于，它们在保持大部分预训练模型参数固定的同时，只训练少量附加层或修改与模型权重交互时使用的参数。这种方法不仅在训练期间节省内存，而且显著缩短了总训练时间。与为了定制基础模型的性能以适应特定目标任务而微调其权重的替代方法相比，PEFT 方法在时间和金钱上都可以节省大量成本，同时在所需数据量少于全面微调训练任务的情况下，提供同等或更好的性能结果。

集成 Hugging-Face 模型和 PyTorch Lightning 框架

PyTorch Lightning 与 Hugging Face 的 Transformers 库 无缝集成，从而实现简化的模型训练工作流，该工作流利用了 Lightning 易于使用的更高级 API 和 HF 的最先进预训练模型。Lightning 与 transformers 的 PEFT 模块 的结合，通过减少代码复杂性并支持使用高质量的预优化模型来完成各种多样的 NLP 任务，从而提高了生产力和可伸缩性。

下面是使用 PyTorch Lightning 和 HuggingFace 的 peft 模块配置基于 PEFT 的模型微调的示例。

from typing import List
from lightning import LightningModule
from peft import get_peft_model, LoraConfig, TaskType
from transformers import AutoModelForSequenceClassification


class TransformerModule(LightningModule):
    def __init__(
        self,
        pretrained_model: str,
        num_classes: int = 2,
        lora_alpha: int = 32,
        lora_dropout: float = 0.1,
        r: int = 8,
        lr: float = 2e-4
    ):
        super().__init__()
        self.model = self.create_model(pretrained_model, num_classes, lora_alpha, lora_dropout, r)
        self.lr = lr
        self.save_hyperparameters("pretrained_model")

    def create_model(self, pretrained_model, num_classes, lora_alpha, lora_dropout, r):
        """Create and return the PEFT model with the given configuration.

        Args:
            pretrained_model: The path or identifier for the pretrained model.
            num_classes: The number of classes for the sequence classification.
            lora_alpha: The alpha parameter for LoRA.
            lora_dropout: The dropout rate for LoRA.
            r: The rank of LoRA adaptations.

        Returns:
            Model: A model configured with PEFT.
        """
        model = AutoModelForSequenceClassification.from_pretrained(
            pretrained_model_name_or_path=pretrained_model,
            num_labels=num_classes
        )
        peft_config = LoraConfig(
            task_type=TaskType.SEQ_CLS,
            inference_mode=False,
            r=r,
            lora_alpha=lora_alpha,
            lora_dropout=lora_dropout
        )
        return get_peft_model(model, peft_config)

    def forward(self, input_ids: List[int], attention_mask: List[int], label: List[int]):
        """Calculate the loss by passing inputs to the model and comparing against ground truth labels.

        Args:
            input_ids: List of token indices to be fed to the model.
            attention_mask: List to indicate to the model which tokens should be attended to, and which should not.
            label: List of ground truth labels associated with the input data.

        Returns:
            torch.Tensor: The computed loss from the model as a tensor.
        """
        return self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=label
        )

完整实现的附加参考可以在 此处配套的仓库中找到。

为基于 PEFT 的微调配置 MLflow

在开始训练过程之前，至关重要的是配置 MLFlow，以便为训练运行记录所有系统指标、损失指标和参数。截至 MLFlow 2.12，TensorFlow 和 PyTorch 的自动记录功能现在支持在训练期间对模型权重进行检查点记录，以在定义的 epoch 频率下提供模型权重的快照，从而在发生错误或计算环境丢失时能够恢复训练。下面是如何启用此功能的示例。

import mlflow


mlflow.enable_system_metrics_logging()
mlflow.pytorch.autolog(checkpoint_save_best_only = False, checkpoint_save_freq='epoch')

在上面的代码中，我们执行了以下操作：

• 启用系统指标记录：系统资源将被记录到 MLflow，以便了解在整个训练过程中内存、CPU、GPU、磁盘使用和网络流量中的瓶颈所在。

• 配置自动记录以记录所有 epoch 的参数、指标和检查点：深度学习涉及尝试各种模型架构和超参数设置。自动记录在系统地记录这些实验方面发挥着至关重要的作用，使比较不同运行并确定哪些配置能产生最佳结果变得更加容易。检查点在每个 epoch 记录，从而在项目初始探索阶段能够详细评估所有中间 epoch。但是，通常不建议在开发后期记录所有 epoch，以避免过多的数据写入和最终训练阶段的延迟。

自动记录的检查点指标和模型工件将在模型训练时显示在 MLflow UI 中，如下所示。

记录和提前停止的重要性

在此训练过程中，将 PyTorch Lightning 的 Trainer 回调与 MLflow 集成至关重要。该集成允许在模型微调期间全面跟踪和记录指标、参数和工件，而无需显式调用 MLflow 记录 API。此外，自动记录 API 允许修改默认记录行为，允许更改记录频率，使记录可以在每个 epoch、指定的 epoch 数量之后或显式定义的步骤中进行。

提前停止

提前停止是神经网络训练中的一项关键正则化技术，旨在通过在验证性能停滞时停止训练来帮助防止过拟合。Pytorch Lightning 包含允许轻松进行高级训练停止控制的 API，如下所示。

配置具有提前停止的 Pytorch Trainer 回调

下面的示例展示了如何在 Lightning 中配置 Trainer 对象，以利用提前停止来防止过拟合。配置完成后，通过调用 Trainer 对象的 fit 来执行训练。通过提供 EarlyStopping 回调，并结合 MLflow 的自动记录，将无需额外的努力即可使用、记录和跟踪适当的 epoch 数量。

from dataclasses import dataclass, field
import os

from data import LexGlueDataModule
from lightning import Trainer
from lightning.pytorch.callbacks import EarlyStopping
import mlflow


@dataclass
class TrainConfig:
    pretrained_model: str = "bert-base-uncased"
    num_classes: int = 2
    lr: float = 2e-4
    max_length: int = 128
    batch_size: int = 256
    num_workers: int = os.cpu_count()
    max_epochs: int = 10
    debug_mode_sample: int | None = None
    max_time: dict[str, float] = field(default_factory=lambda: {"hours": 3})
    model_checkpoint_dir: str = "/local_disk0/tmp/model-checkpoints"
    min_delta: float = 0.005
    patience: int = 4

train_config = TrainConfig()

# Instantiate the custom Transformer class for PEFT training
nlp_model = TransformerModule(
        pretrained_model=train_config.pretrained_model,
        num_classes=train_config.num_classes,
        lr=train_config.lr,
    )

datamodule = LexGlueDataModule(
        pretrained_model=train_config.pretrained_model,
        max_length=train_config.max_length,
        batch_size=train_config.batch_size,
        num_workers=train_config.num_workers,
        debug_mode_sample=train_config.debug_mode_sample,
    )

# Log system metrics while training loop is running
mlflow.enable_system_metrics_logging()

# Automatically log per-epoch parameters, metrics, and checkpoint weights
mlflow.pytorch.autolog(checkpoint_save_best_only = False)

# Define the Trainer configuration
trainer = Trainer(
   callbacks=[
       EarlyStopping(
           monitor="Val_F1_Score",
           min_delta=train_config.min_delta,
           patience=train_config.patience,
           verbose=True,
           mode="max",
       )
   ],
   default_root_dir=train_config.model_checkpoint_dir,
   fast_dev_run=bool(train_config.debug_mode_sample),
   max_epochs=train_config.max_epochs,
   max_time=train_config.max_time,
   precision="32-true"
)

# Execute the training run
trainer.fit(model=nlp_model, datamodule=datamodule)

MLflow 中的可视化和共享功能

MLflow 2.12 中新推出的特定于 DL 的可视化功能使您能够跨 epoch 比较不同的运行和工件。在比较训练运行时，MLflow 能够生成有用的可视化，这些可视化可以集成到仪表板中，便于共享。此外，指标与参数的集中存储使得能够有效地分析训练效果，如下图所示。

何时停止训练？

在训练 DL 模型时，了解何时停止非常重要。高效训练（以最小化训练总成本）和最佳模型性能在很大程度上依赖于防止模型在训练数据上过拟合。训练过长的模型将不可避免地非常擅长“记住”训练数据，从而导致模型在面对新数据时的性能下降。一种评估此行为的简单方法是确保在训练循环中捕获验证数据集指标（对不在训练数据集中的数据进行评分）。将 MLflow 回调集成到 PyTorch Lightning Trainer 中，可以在可配置的迭代次数中迭代记录损失指标，从而实现易于调试的训练性能评估，确保在适当的时间强制执行停止标准以防止过拟合。

使用 MLflow 评估微调模型的 epoch 检查点

有了 MLflow 对您的训练过程进行的细致跟踪和记录，您可以灵活地检索和测试模型在任意检查点上的表现。为此，您可以使用 mlflow.pytorch.load_model() API 从特定运行加载模型，并使用 predict() 方法进行评估。

在下面的示例中，我们将加载第 3 个 epoch 的模型检查点，并使用 Lightning 训练模块根据保存的训练 epoch 的检查点状态生成预测。

import mlflow


mlflow.pytorch.autolog(disable = True)

run_id = '<Add the run ID>'

model = mlflow.pytorch.load_checkpoint(TransformerModule, run_id, 3)

examples_to_test = ["We reserve the right to modify the service price at any time and retroactively apply the adjusted price to historical service usage."]

train_module = Trainer()
tokenizer = AutoTokenizer.from_pretrained(train_config.pretrained_model)
tokens = tokenizer(examples_to_test,
                  max_length=train_config.max_length,
                  padding="max_length",
                  truncation=True)
ds = Dataset.from_dict(dict(tokens))
ds.set_format(
            type="torch", columns=["input_ids", "attention_mask"]
        )

train_module.predict(model, dataloaders = DataLoader(ds))

总结

将 MLflow 集成到预训练语言模型的微调过程中，特别是在自定义命名实体识别、文本分类和指令遵循等应用中，标志着在管理和优化深度学习工作流方面取得了重大进展。在这些工作流中利用 MLflow 的自动记录和跟踪功能，不仅提高了模型开发的可重复性和效率，而且还促进了一个协作环境，在这个环境中，可以轻松地共享和实施见解和改进。

随着我们不断突破这些模型能力的界限，MLflow 等工具将在发挥其全部潜力方面发挥重要作用。

如果您有兴趣查看完整的示例，请随时 查看完整的示例实现。

查看代码

我们提供的代码将深入探讨其他方面，例如从检查点训练、集成 MLflow 和 TensorBoard，以及使用 Pyfunc 进行模型包装等。这些资源专门针对在 Databricks Community Edition 上实现。完整示例仓库中的主运行器 notebook 可以在此处找到。

立即开始使用 MLflow 2.12

立即深入了解最新的 MLflow 更新，并改进您管理机器学习项目的方式！凭借我们最新的增强功能，包括高级指标聚合、自动捕获系统指标、直观的功能分组和简化的搜索功能，MLflow 将为您提升 ML 工作流至新高度。 立即开始使用 MLflow 的尖端工具和功能。

反馈

我们重视您的意见！我们的功能优先级排序受 MLflow 2023 年末调查反馈的指导。请填写我们的 2024 年春季调查，参与其中，您可以帮助确保您最想要的功能在 MLflow 中得到实现。