AWS Trainium 和 Hugging Face Transformers 入门

本教程有两种不同的格式，网页版和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 Core，以及Hugging Face Neuron 深度学习 AMI。

一旦实例启动并运行，我们就可以通过 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

我们需要确保我们已经安装了 training 额外依赖项，以获取所有必需的依赖项。

python -m pip install .[training]

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

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

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

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

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

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

2. 加载并处理数据集

我们将在 emotion 数据集上训练一个文本分类模型，以保持示例的简单性。emotion 数据集包含带有六种基本情绪的英文推特消息：愤怒、恐惧、喜悦、爱、悲伤和惊讶。

我们将使用 🤗 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 ID 才能训练我们的模型。这是由 Tokenizer 完成的，它将输入进行分词（包括将 token 转换为预训练词汇表中相应的 ID）。如果您想了解更多相关内容，请参阅 Hugging Face 课程的第 6 章。

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

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 在两个神经元核心上启动我们的训练脚本以进行分布式训练，从而实现数据并行。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