独自設計チップ AWS Trainium 搭載 Amazon EC2 Trn1 インスタンスで ML トレーニングを高速実行（実践編）

こんにちは！アマゾンウェブサービスジャパン合同会社アンナプルナラボの常世です。

2022 年 10 月 10 日に一般提供開始した Amazon EC2 Trn1 インスタンスはもうお試し頂けましたでしょうか？

本ブログは、「独自設計チップ AWS Trainium 搭載 Amazon EC2 Trn1 インスタンスで ML トレーニングを高速実行（基礎編）」の続編です。まだご覧になってない方は是非「基礎編」をご一読下さい。

基礎編では Trn1 インスタンスと AWS 自身が開発した Trn1 インスタンスの心臓部とも言える第二世代の ML 専用チップ AWS Trainium を紹介しました。本ブログは「実践編」として、日本語 BERT モデルのファインチューニングを Trn1 インスタンス上で実行し、特徴や性能結果を確認していきたいと思います。

日本語 BERT モデルのトレーニング（ファインチューニング）

事前準備

Trn1 インスタンスは Amazon EC2 インスタンスとして提供されます。お使いのアカウントの VPC 内に、trn1.2xlarge のインスタンスタイプを指定して起動します。本ブログ内の検証は、米国西部（オレゴン）リージョン (us-west-2) で実施しました。次に起動した Trn1 インスタンス上に、AWS Neuron ドキュメント（Install PyTorch Neuron on Trn1）に従って PyTorch Neuron 環境をインストールします。本ブログでは Amazon Linux2 AMI 上で環境を構築しています。

2022 年 10 月 27 日に新しいバージョン Neuron 2.4 がリリースされました。 Neuron は定期的にアップデートされ、新しい機能の追加、対応するモデルの拡充、性能向上、バグフィックス等を提供します。「基礎編」を通して既に Neuron 環境を構築済みの場合も最新バージョンへのアップデートを推奨します。

最新の PyTorch Neuron 環境が正しくインストールできているかどうか確認しておきましょう。

（※ Neuron 2.4 にアップデート後の 2022 年 10 月 31 日時点でのインストール結果です。）

$ pip list | grep "neuron\|torch"
aws-neuronx-runtime-discovery 2.9
libneuronxla                  0.1.472
neuronx-cc                    2.2.0.73+0af5a171c
neuronx-hwm                   2.2.0.3+58e9eca77
torch                         1.11.0
torch-neuron                  1.11.0.2.3.0.0
torch-neuronx                 1.11.0.1.2.0
torch-xla                     1.11.0+torchneuron3
$ yum list installed | grep neuron
aws-neuronx-collectives.x86_64        2.10.17.0_1d7a6b5d6-1          @neuron
aws-neuronx-dkms.noarch               2.6.5.0-dkms                   @neuron
aws-neuronx-oci-hook.x86_64           2.1.1.0-1                      @neuron
aws-neuronx-runtime-lib.x86_64        2.10.15.0_d4270a97b-1          @neuron
aws-neuronx-tools.x86_64              2.5.16.0-1                     @neuron

ここでは東北大学が提供し HuggingFace の Transformers で利用可能な日本語事前学習済みモデルを使った文章分類タスクを実行します。まずは必要となるライブラリをインストールしておきましょう。

$ pip install -U transformers[ja]==4.16.2 datasets==2.6.1

データセットの準備

本テストでは、Huggingface Hub で利用可能なデータセットの中から、日本語の Amazon の商品レビューデータセットを利用します。本テストではレビューの文章をPositiveかNegativeに分類する 2 クラスの分類問題として扱うことにします。元々のデータセットで与えられているRatingsが 5、4 のものをPositive (LABEL_1)、2, 1 のものをNegative (LABEL_0)として利用し、Ratings が 3 のデータは使用しないことにします。

from datasets import load_dataset
dataset = load_dataset('amazon_reviews_multi', 'ja')
dataset = dataset.remove_columns(['review_id', 'product_id', 'reviewer_id', 'review_title', 'language', 'product_category'])
dataset = dataset.filter(lambda dataset: dataset["stars"] != 3)
dataset = dataset.map(lambda dataset: {"labels": int(dataset["stars"] > 3)}, remove_columns=["stars"])

次に、文章テキストのままだとモデルのトレーニングはできないため、テキストを意味のある単位で分割（トークナイズ）した上で数値に変換します。トークナイザーには MeCab ベースの BertJapaneseTokenizer を利用しました。

from transformers import BertJapaneseTokenizer

MODEL_NAME = 'cl-tohoku/bert-base-japanese-whole-word-masking'
tokenizer = BertJapaneseTokenizer.from_pretrained(MODEL_NAME)

def tokenize_function(examples):
    return tokenizer(examples["review_body"], padding="max_length", max_length=128, truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)
tokenized_datasets = tokenized_datasets.remove_columns(['review_body'])

train_dataset = tokenized_datasets['train'].shuffle().select(range(4000))
eval_dataset = tokenized_datasets['test'].shuffle().select(range(256))
train_dataset.save_to_disk('./train/')
eval_dataset.save_to_disk('./test/')

実際にどのように変換されているのか、以下のスクリプトを実行し確認してみましょう。

index = 150000
print(dataset["train"][index])
print('Tokenize:', tokenizer.tokenize(dataset["train"]['review_body'][index]))
print('Encode:', tokenizer.encode(dataset["train"]['review_body'][index]))

{'review_body': '使いやすいです。迅速、丁寧な対応に気持ちの良い買い物ができました。', 'labels': 1}
Tokenize: ['使い', 'やすい', 'です', '。', '迅速', '、', '丁寧', 'な', '対応', 'に', '気持ち', 'の', '良い', '買い物', 'が', 'でき', 'まし', 'た', '。']
Encode: [2, 3276, 3428, 2992, 8, 14009, 6, 22063, 18, 1277, 7, 8415, 5, 3614, 20815, 14, 203, 3913, 10, 8, 3]

Trainer API を使用したトレーニング（ファインチューニング）実行

Transformers には Trainer という便利なクラスがあり、Torch Neuron からも利用可能です。ここでは Trainer API を利用してトレーニングを実行していきます。

neuron_parallel_compile による事前コンパイル

前回「基礎編」で、PyTorch Neuron では実行中に計算グラフのコンパイルが発生する点を説明しました。トレーニングの各ステップでは、グラフがトレースされ、トレースされたグラフが以前のものと異なる場合は、再度計算グラフのコンパイルが発生します。大規模なモデルの場合、各グラフのコンパイル時間が長くなることがあり、トレーニング時間の中で占めるコンパイル時間がボトルネックとなってしまう場合もあり得ます。このコンパイル時間を短縮するため、PyTorch Neuron では neuron_parallel_compile ユーティリティが提供されています。neuron_parallel_compile は、スクリプトの試行からグラフを抽出し並列事前コンパイルを実施、コンパイル結果（NEFF : Neuron Executable File Format）をキャッシュとしてディスク上に保持します。

では実際に事前コンパイルを実行してみましょう。以下の内容でbert-jp-precompile.pyというファイル名の Python スクリプトを作成し実行します。スクリプトは基本的にこの後実行するトレーニングスクリプトと同じ内容ですが、neuron_parallel_compileはグラフの高速コンパイルのみを目的とし実際の演算は実行されず、出力結果は無効となります。トレーニング実行中も必要に応じてグラフはコンパイルされるため、この事前コンパイルのプロセスはスキップし、次のトレーニング処理に直接進んでも問題はありません。

from transformers import BertForSequenceClassification, Trainer, TrainingArguments
from datasets import load_from_disk
import torch, torch_xla.core.xla_model as xm

device = xm.xla_device()

MODEL_NAME = 'cl-tohoku/bert-base-japanese-whole-word-masking'
model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2).to(device)

train_dataset = load_from_disk('./train/').with_format("torch")
train_dataset = train_dataset.select(range(64))

training_args = TrainingArguments(
    num_train_epochs = 2,
    learning_rate = 5e-5,
    per_device_train_batch_size = 8,
    output_dir = './results',
)

trainer = Trainer(
    model = model,
    args = training_args,
    train_dataset = train_dataset,
)

train_result = trainer.train()

$ neuron_parallel_compile python3 bert-jp-precompile.py
.
.
2022-10-31 07:00:47.000906: INFO ||PARALLEL_COMPILE||: Compiling /tmp/parallel_compile_workdir/MODULE_SyncTensorsGraph.6619_4112252229678617768.hlo.pb using following command: neuronx-cc compile --target=trn1 --framework XLA /tmp/parallel_compile_workdir/MODULE_SyncTensorsGraph.6619_4112252229678617768.hlo.pb --verbose=35 --output /tmp/parallel_compile_workdir/MODULE_SyncTensorsGraph.6619_4112252229678617768.neff
.....
Compiler status PASS
.
Compiler status PASS
2022-10-31 07:01:48.000684: INFO ||PARALLEL_COMPILE||: Compilation summary:
2022-10-31 07:01:48.000685: INFO ||PARALLEL_COMPILE||: Compilation of /tmp/parallel_compile_workdir/MODULE_SyncTensorsGraph.206_10613991280229713539.hlo.pb was successful. Plugged neff in the cache dir: /var/tmp
2022-10-31 07:01:48.000685: INFO ||PARALLEL_COMPILE||: Compilation of /tmp/parallel_compile_workdir/MODULE_SyncTensorsGraph.22839_3262341186335180678.hlo.pb was successful. Plugged neff in the cache dir: /var/tmp
2022-10-31 07:01:48.000685: INFO ||PARALLEL_COMPILE||: Compilation of /tmp/parallel_compile_workdir/MODULE_SyncTensorsGraph.6619_4112252229678617768.hlo.pb was successful. Plugged neff in the cache dir: /var/tmp
2022-10-31 07:01:48.000685: INFO ||PARALLEL_COMPILE||: Compilation of /tmp/parallel_compile_workdir/MODULE_SyncTensorsGraph.22839_16170154015276808396.hlo.pb was successful. Plugged neff in the cache dir: /var/tmp
2022-10-31 07:01:48.000685: INFO ||PARALLEL_COMPILE||: Total graphs: 4
2022-10-31 07:01:48.000685: INFO ||PARALLEL_COMPILE||: Total successful compilations: 4
2022-10-31 07:01:48.000685: INFO ||PARALLEL_COMPILE||: Total failed compilations: 0

コンパイル時間を短縮するためデータセット、epoch 数を制限している点にご注意ください。5~6分程度でコンパイルは終了します。コンパイル結果は/var/tmp/neuron-compile-cache/以下に保存されています。

シングルワーカーでのトレーニング実行

次に実際にトレーニングを実行してみます。事前コンパイルを実行した場合でも、evaluation ステップを初めて実行するタイミングなど、追加のコンパイルが発生します。一通りのコンパイルが終了した後、2 度目以降の実行では、Neuron コアの恩恵を受けた高速トレーニングを体験できます。以下の内容で bert-jp-single.py というファイル名の Python スクリプトを作成し実行してみましょう。

（※今後 evaluationステップの事前コンパイルにも対応する予定です。）

from transformers import BertForSequenceClassification, BertJapaneseTokenizer, Trainer, TrainingArguments
from datasets import load_from_disk
import torch, torch_xla.core.xla_model as xm

device = xm.xla_device()

MODEL_NAME = 'cl-tohoku/bert-base-japanese-whole-word-masking'
tokenizer = BertJapaneseTokenizer.from_pretrained(MODEL_NAME)
model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2).to(device)

train_dataset = load_from_disk('./train/').with_format("torch")
eval_dataset = load_from_disk('./test/').with_format("torch")

training_args = TrainingArguments(
    num_train_epochs = 10,
    learning_rate = 5e-5,
    per_device_train_batch_size = 8,
    per_device_eval_batch_size = 8,
    logging_strategy = "epoch", 
    output_dir = './results',
)

trainer = Trainer(
    model = model,
    args = training_args,
    train_dataset = train_dataset,
    eval_dataset = eval_dataset,
    tokenizer = tokenizer,
)

train_result = trainer.train()
print(train_result)

eval_result = trainer.evaluate()
print(eval_result)

trainer.save_model("./results")

先程の事前コンパイルとは異なり、今回は実際にトレーニングを実行するため、用意したデータセット全てに対して epoch = 10 で実行しています。

$ python3 bert-jp-single.py
.
.
***** Running training *****
  Num examples = 4000
  Num Epochs = 10
  Instantaneous batch size per device = 8
  Total train batch size (w. parallel, distributed & accumulation) = 8
  Gradient Accumulation steps = 1
  Total optimization steps = 5000

  0%|          | 0/5000 [00:00<?, ?it/s]
2022-11-01 00:14:22.000437: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/USER_neuroncc-2.2.0.73+0af5a171c/.../MODULE_SyncTensorsGraph.206_10613991280229713539.neff. Exiting with a successfully compiled graph
  0%|          | 1/5000 [00:00<18:53,  4.41it/s]
2022-11-01 00:14:23.000480: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/USER_neuroncc-2.2.0.73+0af5a171c/.../MODULE_SyncTensorsGraph.22839_3262341186335180678.neff. Exiting with a successfully compiled graph
  0%|          | 2/5000 [00:02<1:43:39,  1.24s/it]
2022-11-01 00:14:28.000385: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/USER_neuroncc-2.2.0.73+0af5a171c/.../MODULE_SyncTensorsGraph.22839_16170154015276808396.neff. Exiting with a successfully compiled graph
  0%|          | 3/5000 [00:07<4:08:48,  2.99s/it]
  0%|          | 4/5000 [00:10<4:30:59,  3.25s/it]
.
.
100%|██████████| 5000/5000 [09:45<00:00, 11.13it/s]

{'train_runtime': 585.0774, 'train_samples_per_second': 68.367, 'train_steps_per_second': 8.546, 'train_loss': 0.0946555986404419, 'epoch': 10.0}
.
.

ステップ数 5000 のトレーニングが 9~10分程で完了しました。

トレーニング実行中に、AWS Neuron で提供されるneuron-topツールを利用すると、Neuron コア及び vCPU の利用率、アクセラレータメモリ、ホストメモリの利用状況等を確認することができます。trn1.2xlargeには、一つの Trainium チップ、チップ内に二つの Neuron コアが搭載されています。以下の結果から、二つある Neuron コア（NC0 及び NC1）のうち一つの Neuron コアのみが 80.41% の利用率を示しており、もう一方のコアは利用されていないことが分かります。まだ最適化の余地はありそうです。

ここで以下のスクリプトbert-jp-inference.pyを実行し、生成されたモデルから期待通りの出力が得られるか確認しておきましょう。

from transformers import pipeline

classifier = pipeline("text-classification", model = "./results/")

print(classifier("大変すばらしい商品でした。感激です。"))
print(classifier("期待していた商品とは異なりました。残念です。"))

$ python bert-jp-inference.py
[{'label': 'LABEL_1', 'score': 0.9974695444107056}]
[{'label': 'LABEL_0', 'score': 0.9978734254837036}]

期待通りの出力を得られることが確認できたようです。

torchrun を用いたマルチワーカーでのトレーニング実行

それでは、先程のトレーニングスクリプトに変更を加え、二つある Neuron コアを有効活用してみましょう。複数の Neuron コアを利用したマルチワーカーで実行するためには torchrun コマンドを利用します。torchrun コマンドに対して、オプション--nproc_per_nodeで利用する Neuron コアの数（並列実行するワーカー数）を指定します。trn1.2xlargeでは 2 を、trn1.32xlargeの場合は 2, 8, 32 が指定可能です。

torchrun を利用したデータパラレルトレーニングを実行するにあたって、先程のスクリプトに一部変更を加えたbert-jp-dual.pyというファイル名のスクリプトを作成し実行します。

from transformers import BertForSequenceClassification, BertJapaneseTokenizer, Trainer, TrainingArguments
from datasets import load_from_disk
import torch, torch_xla.distributed.xla_backend
import os

if os.environ.get("WORLD_SIZE"):
    torch.distributed.init_process_group('xla')

orig_wrap_model = Trainer._wrap_model
def _wrap_model(self, model, training=True):
    self.args.local_rank = -1
    return orig_wrap_model(self, model, training)
Trainer._wrap_model = _wrap_model

MODEL_NAME = 'cl-tohoku/bert-base-japanese-whole-word-masking'
tokenizer = BertJapaneseTokenizer.from_pretrained(MODEL_NAME)
model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)

train_dataset = load_from_disk('./train/').with_format("torch")
eval_dataset = load_from_disk('./test/').with_format("torch")

training_args = TrainingArguments(
    num_train_epochs = 10,
    learning_rate = 5e-5,
    per_device_train_batch_size = 8,
    per_device_eval_batch_size = 8,
    logging_strategy = "epoch", 
    output_dir = './results',
)

trainer = Trainer(
    model = model,
    args = training_args,
    train_dataset = train_dataset,
    eval_dataset = eval_dataset,
    tokenizer = tokenizer,
)

train_result = trainer.train()
print(train_result)

eval_result = trainer.evaluate()
print(eval_result)

trainer.save_model("./results")

それでは変更後のスクリプトを利用してtrn1.2xlarge上の二つ Neuron コアを利用したトレーニングを実行してみましょう。シングルワーカーでのトレーニング結果と比較しTotal train batch sizeの値が倍の 16 に、Total optimization stepsが半分の 2500 となっている点を確認できると思います。

$ torchrun --nproc_per_node=2 bert-jp-dual.py
.
.
***** Running training *****
  Num examples = 4000
  Num Epochs = 10
  Instantaneous batch size per device = 8
  Total train batch size (w. parallel, distributed & accumulation) = 16
  Gradient Accumulation steps = 1
  Total optimization steps = 2500

  0%|          | 0/2500 [00:00<?, ?it/s]
2022-11-01 00:50:30.000479: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/USER_neuroncc-2.2.0.73+0af5a171c/.../MODULE_SyncTensorsGraph.206_10613991280229713539.neff. Exiting with a successfully compiled graph
  0%|          | 1/2500 [00:00<09:54,  4.20it/s]
2022-11-01 00:50:31.000595: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/USER_neuroncc-2.2.0.73+0af5a171c/.../MODULE_SyncTensorsGraph.23909_822539195623521255.neff. Exiting with a successfully compiled graph
  0%|          | 2/2500 [00:02<54:55,  1.32s/it]
2022-11-01 00:50:36.000955: INFO ||NCC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/USER_neuroncc-2.2.0.73+0af5a171c/.../MODULE_SyncTensorsGraph.23909_15836450389508264351.neff. Exiting with a successfully compiled graph
  0%|          | 3/2500 [00:07<2:15:02,  3.24s/it]
  0%|          | 4/2500 [00:11<2:27:30,  3.55s/it]
.
.
100%|██████████| 2500/2500 [05:45<00:00,  9.61it/s]

{'train_runtime': 345.2408, 'train_samples_per_second': 115.861, 'train_steps_per_second': 7.241, 'train_loss': 0.09048797140866519, 'epoch': 10.0}
.
.

トレーニング実行中の neuron-top の出力も確認してみましょう。今度は二つの Neuron コアが利用されている事が確認できますね。トレーニングに要する実行時間も 5~6 分に削減されました。

BF16 を用いた更なる最適化

ここまでのテストではデータタイプは特に指定しておらず、標準的な 32bit 単精度浮動小数点演算 (FP32) を用いたトレーニングが実行されました。

前回「基礎編」で紹介した、確率的な丸め処理（Stochastic Rounding）とBF16 でトレーニングを実行した場合、FP32 同様の精度を保ちながら更なる高速化を実現する事が可能です。環境変数 XLA_USE_BF16=1を設定した上でトレーニングを実行してみましょう。（※ BF16 利用時は確率的な丸め処理もデフォルトで利用されます。）

$ XLA_USE_BF16=1 python3 bert-jp-single.py
.
.
{'train_runtime': 406.3988, 'train_samples_per_second': 98.425, 'train_steps_per_second': 12.303, 'train_loss': 0.09391015625, 'epoch': 10.0}
.
.
$ XLA_USE_BF16=1 torchrun --nproc_per_node=2 bert-jp-dual.py
.
.
{'train_runtime': 268.2745, 'train_samples_per_second': 149.101, 'train_steps_per_second': 9.319, 'train_loss': 0.099168701171875, 'epoch': 10.0}

FP32 で実行した際の結果と比較し、トレーニングに要する実行時間はそれぞれ 6~7 分、4~5 分と削減されている事が確認できると思います。

ここまで説明した一連の処理を実行するスクリプトは以下のレポジトリから入手可能です。

https://github.com/aws-neuron/aws-neuron-samples

まとめ

遂にローンチとなった Amazon EC2 Trn1 インスタンスが加わり、大規模 ML トレーニングを実行するためのインスタンスの選択肢が広がりました。Trn1 がより高い性能をより低いコストで実現するポテンシャルを秘めたインスタンスである一方、現時点でサポートするモデルは、BERT や GPTなど自然言語処理モデルが中心です。対応するモデルの拡張、フレームワーク対応、性能向上、各種 AWS サービスとのさらなる連携強化、Trn1 インスタンスは AWS Neuron のアップデートと共に更なる進化を遂げていきます。11 月 3 日には Amazon SageMaker でもトレーニングジョブとして Trn1 インスタンスを選択できるようになりました。

文字通りの Day 1、今後の Trn1 インスタンスの進化にも期待しましょう。

________________________________________
本ブログは、アマゾンウェブサービスジャパン合同会社アンナプルナラボの常世が執筆しました。

Amazon Web Services ブログ