在 AWS Trainium 上微调 BERT 用于文本分类
此教程的笔记本版本在此.
本教程将帮助您开始使用AWS Trainium 和 Hugging Face Transformers。它将涵盖如何在 AWS 上设置 Trainium 实例,以及如何加载和微调用于文本分类的转换器模型。
您将学习如何
在开始之前,请确保您拥有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 核心)和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
输出,其中包含指向笔记本的 URL。
http://localhost:8080/?token=8c1739aff1755bd7958c4cfccc8d08cb5da5234f61f129a9
我们可以点击它,然后在我们的本地浏览器中打开一个jupyter
环境。
我们将仅使用 Jupyter 环境来准备数据集,然后使用torchrun
在两个 Neuron 核心上启动我们的训练脚本以进行分布式训练。让我们创建一个新的笔记本并开始吧。
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
的 1 对 1 替换。您只需调整训练脚本中的导入即可。
- from transformers import Trainer, TrainingArguments
+ from optimum.neuron import NeuronTrainer as Trainer
+ from optimum.neuron import NeuronTrainingArguments as TrainingArguments
我们基于 “Pytorch 2.0 和 Hugging Face Transformers 入门” 博文,使用 NeuronTrainer
准备了一个简单的 train.py 训练脚本。以下是摘录
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
命令将训练脚本加载到我们的环境中,或者从 此处 手动将其复制到笔记本中。
!wget https://raw.githubusercontent.com/huggingface/optimum-neuron/main/notebooks/text-classification/scripts/train.py
我们将使用 torchrun
在两个 Neuron 核心上启动训练脚本以进行分布式训练。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 美元/小时 * 0.15 小时 = 0.20 美元
)