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