为单个 GPU 微调 Hugging Face 模型

本文介绍如何在单个 GPU 上使用 Hugging Face transformers 库微调 Hugging Face 模型。本文还提供了有关如何从湖屋加载数据并将模型记录到 MLflow 的 Databricks 特定建议，让你能够在 Azure Databricks 上使用和管理模型。

Hugging Face transformers 库提供了训练器实用工具和自动模型类，它们可用于加载和微调转换器模型。

这些工具可用于以下任务，只需进行简单的修改：

加载要微调的模型。
正在构建 Hugging Face Transformers Trainer 工具的配置。
在单个 GPU 上执行训练。

请参阅什么是 Hugging Face Transformers？

要求

驱动器上带有一个 GPU 的单节点集群。
Databricks Runtime 13.0 ML 及更高版本的 GPU 版本。
- 此微调示例需要 Databricks Runtime 13.0 ML 及更高版本中包含的 🤗 Transformers、🤗 Datasets 和 🤗 Evaluate 包。
MLflow 2.3。
准备和加载数据，以便使用转换器来优化模型。

标记化 Hugging Face 数据集

Hugging Face 转换器模型需要标记化的输入，而不是下载的数据中的文本。为确保与基础模型兼容，请使用从基础模型加载的 AutoTokenizer。 Hugging Face datasets 允许将分词器一致地直接应用于训练和测试数据。

例如：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(base_model)
def tokenize_function(examples):
    return tokenizer(examples["text"], padding=False, truncation=True)

train_test_tokenized = train_test_dataset.map(tokenize_function, batched=True)

设置训练配置

可以使用 Hugging Face 训练配置工具来配置训练器。训练器类要求用户提供：

指标
基础模型
训练配置

除了 loss 所计算的默认 Trainer 指标外，还可以配置评估指标。以下示例演示如何将 accuracy 添加为指标：

import numpy as np
import evaluate
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

使用 NLP 自动模型类为任务加载适当的模型。

对于文本分类，请使用 AutoModelForSequenceClassification 加载用于文本分类的基础模型。创建模型时，请提供类数量以及在准备数据集期间创建的标签映射。

from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(
        base_model,
        num_labels=len(label2id),
        label2id=label2id,
        id2label=id2label
        )

接下来创建训练配置。 TrainingArguments 类用于指定输出目录、评估策略、学习速率和其他参数。

from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(output_dir=training_output_dir, evaluation_strategy="epoch")

使用数据整理器可以批处理训练和评估数据集中的输入。 DataCollatorWithPadding 为文本分类提供良好的基线性能。

from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer)

构造所有这些参数后，即可创建一个 Trainer。

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_test_dataset["train"],
    eval_dataset=train_test_dataset["test"],
    compute_metrics=compute_metrics,
    data_collator=data_collator,
)

训练并记录到 MLflow

Hugging Face 可与 MLflow 完美对接，在模型训练期间使用 MLflowCallback 自动记录指标。但是，你需要自行记录已训练的模型。

在 MLflow 运行中包装训练。这会从 tokenizer 和训练的模型构造 Transformers 管道，并将其写入本地磁盘。最后，使用 mlflow.transformers.log_model 将模型记录到 MLflow。

from transformers import pipeline

with mlflow.start_run() as run:
  trainer.train()
  trainer.save_model(model_output_dir)
  pipe = pipeline("text-classification", model=AutoModelForSequenceClassification.from_pretrained(model_output_dir), batch_size=1, tokenizer=tokenizer)
  model_info = mlflow.transformers.log_model(
        transformers_model=pipe,
        artifact_path="classification",
        input_example="Hi there!",
    )

如果不需要创建管道，可以将用于训练的组件提交到字典中：

model_info = mlflow.transformers.log_model(
  transformers_model={"model": trainer.model, "tokenizer": tokenizer},
  task="text-classification",
  artifact_path="text_classifier",
  input_example=["MLflow is great!", "MLflow on Databricks is awesome!"],
)

加载用于推理的模型

当模型已记录并准备就绪时，加载用于推理的模型的操作与加载 MLflow 包装的预训练模型的操作相同。

logged_model = "runs:/{run_id}/{model_artifact_path}".format(run_id=run.info.run_id, model_artifact_path=model_artifact_path)

# Load model as a Spark UDF. Override result_type if the model does not return double values.
loaded_model_udf = mlflow.pyfunc.spark_udf(spark, model_uri=logged_model, result_type='string')

test = test.select(test.text, test.label, loaded_model_udf(test.text).alias("prediction"))
display(test)

排查常见 CUDA 错误

本部分介绍常见的 CUDA 错误以及有关如何解决这些错误的指导。

OutOfMemoryError：CUDA 内存不足

训练大型模型时，可能会遇到的常见错误是 CUDA 内存不足错误。

示例：

OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 14.76 GiB total capacity; 666.34 MiB already allocated; 17.75 MiB free; 720.00 MiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF.

请尝试以下建议来解决此错误：

减小用于训练的批大小。可以减小 per_device_train_batch_size 中的值。
使用较低精度的训练。可以在 fp16=True 中设置。
在 TrainingArguments 中使用 gradient_accumulation_steps，以有效增加整体批大小。
使用 8 位 Adam 优化器。
在训练前清理 GPU 内存。有时，GPU 内存可能被一些未使用的代码占用。
```
from numba import cuda
device = cuda.get_current_device()
device.reset()
```

CUDA 内核错误

运行此训练时，可能会收到 CUDA 内核错误。

示例：

CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.

For debugging, consider passing CUDA_LAUNCH_BLOCKING=1.

故障排除：

尝试在 CPU 上运行代码，看错误是否可重现。
另一个选项是通过设置 CUDA_LAUNCH_BLOCKING=1 获取更好的回溯：
```
import os
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
```

笔记本：在单个 GPU 上微调文本分类

为了帮助你快速开始使用示例代码，此示例笔记本提供了一个用于微调文本分类模型的端到端示例。本文的后续部分将更详细地介绍如何在 Azure Databricks 上使用 Hugging Face 进行微调。

微调 Hugging Face 文本分类模型笔记本

获取笔记本

其他资源

详细了解 Azure Databricks 上的 Hugging Face。

什么是 Hugging Face Transformers？
可以在 Spark 上使用 Hugging Face Transformers 模型来横向扩展 NLP 批处理应用程序，详见使用 Hugging Face Transformers 进行 NLP 的模型推理。

上次更新时间为 2026-10-26