通过

为单个 GPU 微调 Hugging Face 模型

  • 项目

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

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

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

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

请参阅什么是 Hugging Face Transformers？

要求

标记化 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 训练配置工具来配置训练器。 训练器类要求用户提供：

  • 指标
  • 基础模型
  • 训练配置

除了 Trainer 所计算的默认 loss 指标外，还可以配置评估指标。 以下示例演示如何将 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.

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

  • 减小用于训练的批大小。 可以减小 TrainingArguments 中的 per_device_train_batch_size 值。

  • 使用较低精度的训练。 可以在 TrainingArguments 中设置 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。