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

在深度学习领域，在私有数据集上对预训练大型语言模型（LLM）进行微调是一种出色的定制选项，可以提高模型对特定任务的相关性。这种做法不仅普遍，而且对于开发专用模型至关重要，特别是对于文本分类和摘要等任务。

在这种情况下，MLflow 等工具是无价的。MLflow 等跟踪工具可确保训练过程的每个方面（指标、参数和工件）都可复现地进行跟踪和记录，从而允许对调优迭代进行分析、比较和共享。

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

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

本博客中使用的示例场景利用了 unfair-TOS 数据集。

在当今世界，很难找到一项服务、平台，甚至一件消费品没有与其相关的具有法律约束力的服务条款。这些百科全书般大小的协议，充满了密集的法律术语和有时令人困惑的细节程度，体量之大以至于大多数人根本不阅读就接受了它们。然而，随着时间的推移，报告表明，有时其中会嵌入一些可疑的不公平条款。

通过机器学习 (ML) 解决服务条款 (TOS) 协议中的不公平条款尤为重要，因为法律协议中涉及消费者权益的透明度和公平性需求迫切。考虑以下示例 TOS 协议中的条款："我们可能不时修订这些条款。变更将不具有追溯力，且条款的最新版本将始终..." 该条款规定服务提供商可以随时出于任何原因，无论是否通知，暂停或终止服务。大多数人会认为这相当不公平。

尽管这句话隐藏在一个相当密集的文档深处，但机器学习算法不会像人类那样因梳理文本和识别可能显得有些不公平的条款而感到疲惫。通过自动识别潜在的不公平条款，基于 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 中新引入的深度学习专用可视化功能使您能够跨 epoch 对不同的运行和工件进行比较。在比较训练运行时，MLflow 能够生成有用的可视化，这些可视化可以集成到仪表板中，方便共享。此外，指标的集中存储与参数结合，可以有效分析训练效果，如下图所示。

何时停止训练？

训练深度学习模型时，了解何时停止训练非常重要。高效训练（用于最大限度地降低训练总成本）和最佳模型性能在很大程度上取决于防止模型在训练数据上过拟合。训练时间过长的模型将不可避免地变得非常擅长有效地“记忆”训练数据，从而在遇到新数据时导致模型性能下降。评估这种行为的一个直接方法是确保在训练循环期间捕获验证数据集指标（对不在训练数据集中的数据进行损失指标评分）。将 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 社区版上实现而定制的。完整示例存储库中的主要运行器 Notebook 可在此处找到。

立即开始使用 MLflow 2.12

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

反馈

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