在深度学习领域，在私有数据集上微调预训练的大型语言模型 (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 和 Hugging Face 最先进的预训练模型。Lightning 与 transformers 的 PEFT 模块的结合通过减少代码复杂性并支持使用高质量的预优化模型来完成各种自然语言处理任务，从而提高了生产力和可扩展性。
以下是使用 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
)
完整的实现的其他参考资料可以在此处的配套仓库中看到
配置 MLflow 进行基于 PEFT 的微调
在启动训练过程之前，配置 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 能够生成有用的可视化，这些可视化可以集成到仪表板中，便于共享。此外，指标的集中存储以及参数，可以有效分析训练效果，如下图所示。
何时停止训练？
训练深度学习模型时，了解何时停止训练非常重要。高效训练（为了最大限度地降低进行训练所产生的总成本）和最佳模型性能在很大程度上依赖于防止模型在训练数据上过拟合。训练时间过长的模型将不可避免地变得非常擅长有效“记忆”训练数据，导致模型在处理新数据时性能下降。评估这种行为的一种简单方法是确保在训练循环期间捕获验证数据集指标（评估不在训练数据集中的数据的损失指标）。将 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 中实现。