在 AWS Trainium 上微调 BERT 进行文本分类

本教程的 Notebook 版本在此.

本教程将帮助你开始使用 AWS Trainium 和 Hugging Face Transformers。它将涵盖如何在 AWS 上设置 Trainium 实例、加载并微调一个用于文本分类的 transformers 模型。

您将学习如何

设置 AWS 环境
加载并处理数据集
使用 Hugging Face Transformers 和 Optimum Neuron 微调 BERT

在开始之前，请确保你有一个 Hugging Face 账户来保存模型文件和实验结果。

快速入门：AWS Trainium

AWS Trainium (Trn1) 是专为深度学习（DL）训练工作负载打造的 EC2 实例。Trainium 是 AWS Inferentia 的继任者，专注于高性能训练工作负载，声称与同类基于 GPU 的实例相比，可节省高达 50% 的训练成本。

Trainium 已针对训练自然语言处理、计算机视觉和推荐模型进行了优化。该加速器支持多种数据类型，包括 FP32、TF32、BF16、FP16、UINT8 以及可配置的 FP8。

最大的 Trainium 实例 trn1.32xlarge 拥有超过 500GB 的内存，使得在单个实例上微调约 100 亿参数模型变得轻而易举。下面是可用实例类型的概览。更多详情在此

实例大小	加速器	加速器内存	vCPU	CPU 内存	每小时价格
trn1.2xlarge	1	32	8	32	$1.34
trn1.32xlarge	16	512	128	512	$21.50
trn1n.32xlarge（2倍带宽）	16	512	128	512	$24.78

现在我们了解了 Trainium 的功能，让我们开始吧。🚀

注意：本教程是在 trn1.2xlarge AWS EC2 实例上创建的。

1. 设置 AWS 环境

在本例中，我们将使用 AWS 上的 trn1.2xlarge 实例，该实例配备 1 个加速器，包含两个 Neuron Cores，并使用 Hugging Face Neuron 深度学习 AMI。

这篇博文不详细介绍如何创建实例。你可以参考我之前的博文 “为 Hugging Face Transformers 设置 AWS Trainium”，其中包含了设置环境的详细步骤。

一旦实例启动并运行，我们就可以通过 ssh 连接到它。但我们不想在终端内进行开发，而是希望使用一个 Jupyter 环境，我们可以用它来准备数据集和启动训练。为此，我们需要在 ssh 命令中添加一个用于转发的端口，这将把我们本地主机的流量隧道传输到 Trainium 实例。

PUBLIC_DNS="" # IP address, e.g. ec2-3-80-....
KEY_PATH="" # local path to key, e.g. ssh/trn.pem

ssh -L 8080:localhost:8080 -i ${KEY_NAME}.pem ubuntu@$PUBLIC_DNS

现在我们可以启动我们的 jupyter 服务器了。

python -m notebook --allow-root --port=8080

你应该会看到一个熟悉的 jupyter 输出，其中包含一个指向 notebook 的 URL。

https://:8080/?token=8c1739aff1755bd7958c4cfccc8d08cb5da5234f61f129a9

我们可以点击它，然后在我们的本地浏览器中打开一个 jupyter 环境。

我们将仅使用 Jupyter 环境来准备数据集，然后使用 torchrun 在两个 Neuron Cores 上启动我们的训练脚本以进行分布式训练。让我们创建一个新的 notebook 并开始吧。

2. 加载并处理数据集

为了使示例简单明了，我们将在 emotion 数据集上训练一个文本分类模型。emotion 数据集包含带有六种基本情绪（愤怒、恐惧、喜悦、爱、悲伤和惊讶）的英文 Twitter 消息。

我们将使用 🤗 Datasets 库中的 load_dataset() 方法来加载 emotion 数据集。

from datasets import load_dataset

# Dataset id from huggingface.co/dataset
dataset_id = "philschmid/emotion"

# Load raw dataset
raw_dataset = load_dataset(dataset_id)

print(f"Train dataset size: {len(raw_dataset['train'])}")
print(f"Test dataset size: {len(raw_dataset['test'])}")

# Train dataset size: 16000
# Test dataset size: 2000

我们来看一个数据集的例子。

from random import randrange

random_id = randrange(len(raw_dataset['train']))
raw_dataset['train'][random_id]
# {'text': 'i feel isolated and alone in my trade', 'label': 0}

为了训练我们的模型，我们必须将“自然语言”转换为词元 ID。这是通过一个分词器（Tokenizer）完成的，它对输入进行分词（包括将词元转换为其在预训练词汇表中的相应 ID）。如果你想了解更多，请查看 Hugging Face 课程的第 6 章。

我们的 Neuron 加速器期望输入的形状是固定的。我们需要将所有样本截断或填充到相同的长度。

from transformers import AutoTokenizer
import os
# Model id to load the tokenizer
model_id = "bert-base-uncased"
save_dataset_path = "lm_dataset"
# Load Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Tokenize helper function
def tokenize(batch):
    return tokenizer(batch['text'], padding='max_length', truncation=True,return_tensors="pt")

# Tokenize dataset
raw_dataset =  raw_dataset.rename_column("label", "labels") # to match Trainer
tokenized_dataset = raw_dataset.map(tokenize, batched=True, remove_columns=["text"])
tokenized_dataset = tokenized_dataset.with_format("torch")

# save dataset to disk
tokenized_dataset["train"].save_to_disk(os.path.join(save_dataset_path,"train"))
tokenized_dataset["test"].save_to_disk(os.path.join(save_dataset_path,"eval"))

3. 使用 Hugging Face Transformers 微调 BERT

通常，你会使用 Trainer 和 TrainingArguments 来微调基于 PyTorch 的 transformer 模型。

但是，我们与 AWS 共同开发了一个 NeuronTrainer，以在 Trainium 或 Inferentia2 实例上训练时提高性能、鲁棒性和安全性。NeuronTrainer 还带有一个模型缓存，它允许我们使用来自 Hugging Face Hub 的预编译模型和配置来跳过训练开始时所需的编译步骤。这可以将训练时间减少约 3 倍。

NeuronTrainer 是 optimum-neuron 库的一部分，可以作为 Trainer 的一对一替代品使用。你只需在训练脚本中调整导入即可。

- from transformers import Trainer, TrainingArguments
+ from optimum.neuron import NeuronTrainer as Trainer
+ from optimum.neuron import NeuronTrainingArguments as TrainingArguments

我们准备了一个简单的 train.py 训练脚本，它基于 “PyTorch 2.0 和 Hugging Face Transformers 入门” 这篇博文，并使用了 NeuronTrainer。下面是摘录：

from transformers import TrainingArguments
from optimum.neuron import NeuronTrainer as Trainer

def parse_args():
	...

def training_function(args):

    # load dataset from disk and tokenizer
    train_dataset = load_from_disk(os.path.join(args.dataset_path, "train"))
		...

    # Download the model from huggingface.co/models
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_id, num_labels=num_labels, label2id=label2id, id2label=id2label
    )

    training_args = TrainingArguments(
			...
    )

    # Create Trainer instance
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics,
    )

    # Start training
    trainer.train()

我们可以使用 wget 命令将训练脚本加载到我们的环境中，或者从这里手动将其复制到 notebook 中。

!wget https://raw.githubusercontent.com/huggingface/optimum-neuron/main/notebooks/text-classification/scripts/train.py

我们将使用 torchrun 在两个 Neuron Cores 上启动我们的训练脚本以进行分布式训练。torchrun 是一个能自动将 PyTorch 模型分布到多个加速器上的工具。我们可以将加速器的数量作为 nproc_per_node 参数与我们的超参数一起传递。

我们将使用以下命令来启动训练

!torchrun --nproc_per_node=2 train.py \
 --model_id bert-base-uncased \
 --dataset_path lm_dataset \
 --lr 5e-5 \
 --per_device_train_batch_size 16 \
 --bf16 True \
 --epochs 3

注意：如果你看到非常差的准确率，你可能需要暂时禁用 bf16。

9分钟后，训练完成并取得了 0.914 的出色 f1 分数。

***** train metrics *****
  epoch                    =        3.0
  train_runtime            =    0:08:30
  train_samples            =      16000
  train_samples_per_second =     96.337

***** eval metrics *****
  eval_f1                  =      0.914
  eval_runtime             =    0:00:08

最后，别忘了终止 EC2 实例以避免不必要的费用。从性价比来看，我们的训练仅花费了 20美分 (1.34$/h * 0.15h = 0.20$)

AWS Trainium & Inferentia

在 AWS Trainium 上微调 BERT 进行文本分类

快速入门：AWS Trainium

1. 设置 AWS 环境

2. 加载并处理数据集

3. 使用 Hugging Face Transformers 微调 BERT