在 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