AWS Trainium 和 Hugging Face Transformers 入门指南

本教程提供两种不同的格式：网页版和 notebook 版本.

本指南将帮助您开始使用 AWS Trainium 和 Hugging Face Transformers。它将涵盖如何在 AWS 上设置 Trainium 实例，加载和微调用于文本分类的 transformers 模型。

您将学习如何：

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

在开始之前，请确保您拥有 Hugging Face 帐户以保存 artifacts 和实验。

快速入门：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 Deep Learning AMI。

实例启动并运行后，我们可以 ssh 进入它。但是，我们不想在终端中进行开发，而是想使用 Jupyter 环境，我们可以使用它来准备数据集和启动训练。为此，我们需要在 ssh 命令中添加一个端口转发，这将把我们的 localhost 流量隧道传输到 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

我们需要确保已安装 training extra，以获取所有必要的依赖项

python -m pip install .[training]

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

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

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

http://localhost:8080/?token=8c1739aff1755bd7958c4cfccc8d08cb5da5234f61f129a9

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

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

2. 加载和处理数据集

我们正在训练一个文本分类模型，使用 emotion 数据集，以使示例保持简单。emotion 是一个包含英语 Twitter 消息的数据集，包含六种基本情绪：愤怒、恐惧、喜悦、爱、悲伤和惊讶。

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

from datasets import load_dataset


# Dataset id from huggingface.co/dataset
dataset_id = "dair-ai/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 also like to listen to jazz whilst painting it makes me feel more artistic and ambitious actually look to the rainbow', 'label': 1}

我们必须将“自然语言”转换为 token IDs 才能训练我们的模型。这由 Tokenizer 完成，Tokenizer 对输入进行 tokenization（包括将 tokens 转换为其在预训练词汇表中的对应 IDs）。如果您想了解更多信息，请查看第 6 章 Hugging Face 课程。

为了避免图重新编译，输入应具有固定的形状。我们需要截断或填充所有样本到相同的长度。

import os

from transformers import AutoTokenizer


# Model id to load the tokenizer
model_id = "bert-base-uncased"

# 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")


def tokenize_function(example):
    return tokenizer(
        example["text"],
        padding="max_length",
        truncation=True,
    )


# Tokenize dataset
tokenized_emotions = raw_dataset.map(tokenize, batched=True, remove_columns=["text"])

3. 使用 Hugging Face Transformers 微调 BERT

我们可以使用 Trainer 和 TrainingArguments 来微调基于 PyTorch 的 transformer 模型。

我们准备了一个简单的 train.py 训练脚本，用于在数据集上执行训练和评估。以下是摘录

from transformers import Trainer, TrainingArguments

def parse_args():
	...

def training_function(args):

    ...

    # 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=tokenized_emotions["train"],
        eval_dataset=tokenized_emotions["validation"],
        processing_class=tokenizer,
    )


    # Start training
    trainer.train()

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

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

我们将使用 torchrun 在两个 neuron core 上启动我们的训练脚本以进行分布式训练，从而实现数据并行。torchrun 是一种自动在多个加速器上分配 PyTorch 模型的工具。我们可以将加速器的数量作为 nproc_per_node 参数以及我们的超参数一起传递。

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

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

编译后，只需几分钟即可完成训练。

***** train metrics *****
  epoch                    =        3.0
  eval_loss                =     0.1761
  eval_runtime             = 0:00:03.73
  eval_samples_per_second  =    267.956
  eval_steps_per_second    =     16.881
  total_flos               =  1470300GF
  train_loss               =     0.2024
  train_runtime            = 0:07:27.14
  train_samples_per_second =     53.674
  train_steps_per_second   =      6.709

最后但并非最不重要的一点是，终止 EC2 实例以避免不必要的费用。查看性价比，我们的训练仅花费 20 美分 (1.34 美元/小时 * 0.13 小时 = 0.18 美元)

AWS Trainium & Inferentia

AWS Trainium 和 Hugging Face Transformers 入门指南

快速入门：AWS Trainium

1. 设置 AWS 环境

2. 加载和处理数据集

3. 使用 Hugging Face Transformers 微调 BERT