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

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

アンナプルナラボって何？と思われた方も多いかと思います。アンナプルナラボは元々 ARM アーキテクチャベースの SoC (System-on-Chip) を提供していたファブレス半導体ベンチャーでした。アンナプルナラボ製 ARM SoC は家庭向け、SMB 向けの NAS (Network Attached Storage) や、ネットワーク機器にも採用されているので、皆様がお使いの機器の中で動作しているかもしれません。2015年に AWS の一部となってからは AWS Nitro System や、AWS Graviton プロセッサを開発、提供してきました。

そのアンナプルナラボが機械学習 (ML) ワークロード向けに開発した第一世代のチップが Amazon EC2 Inf1 インスタンスに搭載された AWS Inferentia です。Amazon EC2 Inf1 インスタンスは、クラウドで最も低いコストで高性能 ML 推論を提供します。同等世代の GPU ベースの Amazon EC2 インスタンスと比較して、スループットが最大 2.3 倍高く、最大 70% 低い推論あたりのコストを実現しています。

Inf1 インスタンスは 2019 年末に一般提供開始して以降、多くのお客様にご利用頂いています。中でも株式会社マネーフォワード様には推論性能（低レイテンシー）及び実装の容易さから採用頂き、大規模推論環境を構築した経験談を Builders.flash に寄稿頂きました。AirBnB、Autodesk、InfoJobs (Adevinta)、Finch Computing といった世界中のお客様にご活用頂いているのみならず、Amazon Alexa や Amazon Advertising、Amazon Robotics、Amazon Search のように Amazon 内でサービスを支える基盤としても Inf1 インスタンスが利用されています。

AWS Trainium 搭載 Amazon EC2 Trn1 インスタンス

近年、機械学習 (ML) モデルの複雑さ、規模の進化は目を見張るものがあります。特に自然言語処理 (NLP) の分野では、モデルのサイズが過去数年間で年 10 倍のペースで増加しています。モデルのトレーニングにかかる時間とコストが爆発的に増加している状況の中で、求められるハードウェア性能は、もはや単一インスタンスで処理できるものではなくなってきているのが実情です。

AWS では、これまで大規模モデルのトレーニング用インスタンスとして、NVIDIA 社 A100 Tensor Core GPU を搭載した P4d インスタンス、Habana Labs (インテルの関連会社) の Gaudi アクセラレーターを搭載した DL1インスタンスをそれぞれ提供してきました。本ブログでは10月10日に満を持して一般提供開始した、AWS 自身が ML ワークロード、特に ML トレーニング向けに開発した第二世代の ML 専用チップ AWS Trainium を搭載した Amazon EC2 Trn1 インスタンスを紹介します。またブログ後半では、実際に Trn1 を体験するためのステップを紹介します。

Amazon EC2 Trn1 インスタンス

Amazon EC2 Trn1 インスタンスは、AWS クラウドでの一般的な自然言語処理 (NLP) モデルの ML トレーニングにおいて最高の性能を提供します。またトレーニングコストは、同等世代の GPU ベースのインスタンスと比較して最大 50% 削減することが可能です。Trn1 インスタンスでは、小規模モデルのトレーニングにもご利用頂けるように Trainium チップを１つ搭載した trn1.2xlarge と、大規模モデルの分散トレーニングを想定し、Trainium チップを 16個搭載した trn1.32xlarge の２つのサイズを用意しました。

インスタンスサイズ	Trainium	アクセラレータメモリ (GiB)	vCPU	インスタンスメモリ (GiB)	ローカル NVMe ストレージ (TB)	NW帯域 (Gbps)	EBS帯域 (Gbps)
trn1.2xlarge	1	32	8	32	0.5	最大 12.5	最大 20
trn1.32xlarge	16	512	128	512	8	800	80

最大サイズとなる trn1.32xlarge では、16個の Trainium チップと計 512GiB の広帯域 HBM2 アクセラレータメモリを搭載しています。インスタンス内の各 Trainium チップ間は、広帯域かつ低レイテンシーを実現した NeuronLink-v2 で接続されています。

また trn1.32xlarge では複数のインスタンスを EC2 UltraClusters と呼ばれるハイパースケールクラスターにデプロイ可能です。 EC2 UltraClusters では 30,000 を超える Trainium チップで構成される分散トレーニング環境に拡張可能です。各インスタンス間は、数ある EC2 インスタンスの中でも最大帯域となる 800Gbps Elastic Fabric Adapter (EFA) で接続しており、ペタビット規模のノンブロッキングネットワークで相互接続されたクラスターは、6.3 エクサフロップスのコンピュート性能を提供します。

AWS Trainium

次に Trn1 インスタンスの心臓部ともいえる AWS Trainium を見ていきたいと思います。Trainium は ML 推論ワークロード向けに開発した第一世代の ML チップ Inferentia に続いて、最も費用効果の高いクラウド上での MLトレーニング環境を提供すべく開発した第二世代の ML チップです。AWS Inferentia の設計思想を維持したまま、その性能は大きく進化しました。

AWS Inferentia では内部に４つの Neuron コアと呼ばれるエンジンが搭載されていました。AWS Trainium ではチップ辺りのコア数は４つから２つに減ってはいるものの、より進化した第二世代の Neuronコアを搭載しています。

それぞれの Neuron コアは、以下に示す４つの専用エンジンと、各エンジンを利用し効率的に演算処理を実行するための大容量オンチップ SRAM メモリから構成されています。

Tensor エンジン：畳み込み等、行列演算に最適化されたシストリックアレイベースの Tensor エンジン。FP16/BF16 を用いた演算性能は 100 TFLOPS を超え、Inferentia に搭載された Neuron コアの 6 倍の性能を達成。
Vector エンジン： Batch Normalization や Pooling 処理に最適化されたエンジン。
Scalar エンジン： ReLU 等の活性化関数に最適化されたエンジン。
汎用 SIMD エンジン：カスタムオペレータに対応できるように Trainium で新規に追加された組み込み型の SIMD エンジン。C++ でカスタムオペレータの記述が可能に。

Neuron コアは各種データタイプに対応可能です。Inferentia で対応していた FP16, BF16, INT8 に加えて、FP32, TF32, cFP8 に対応しています。cFP8 は聞きなれない方も多いかと思いますが、Confiugurable FP8 という新しい 8 ビットの浮動小数点データタイプです。Trainium では、2022 年 9 月にNVIDIA社、ARM社、インテル社が共同で発表した２つのモード（E4M3 及び E5M2）に、E3M4（指数部 3 ビット、仮数部 4 ビット）を加えた３つのモードをいち早くハードウェア実装しました。

また、Trainium では確率的な丸め処理（Stochastic Rounding）という特徴的な丸め処理を初めてハードウェア実装しました。IEEE 754-2008 Standard で定義された標準的な丸め処理（RNE : Round Nearest Even）では、重みの値をアップデートする際、アップデートする値が小さい場合は値が完全につぶれてしまうことがあります。分かりやすい例として整数 1 に対して 0.2 を足し整数に丸める場合を考えると、何度 0.2 を足しても結果は 1 のままですが、確率的な丸め処理では、80％の確率で 1 に、20％の確率で 2 となるように丸め処理を行い、統計的に本来の値に近い値を保ちます。

以下のグラフは、BERT-Large Pre-Training の例となります。BF16 でトレーニングを実行した場合、FP32 同様の高い精度と BF16 を使用した場合の高いスループットとを両立している点を確認頂けると思います。

AWS Neuron

AWS Neuron は、Inf1 及び Trn1 インスタンスで ML ワークロードを実行するために利用される SDK です。Trn1 インスタンスの発表に合わせて、Neuron 2.x としてトレーニングワークロードにも対応しました。Trn1 向けには PyTorch、TensorFlow に対応しており、それぞれ PyTorch XLA、TensorFlow XLA をバックエンドのコンパイラとして利用します。（※現時点では PyTorch のみ正式対応中、TensorFlow については今後対応予定です。AWS Neuron の機能ロードマップについてはこちらをご確認ください。）

ML ワークロードを Inf1 及び Trn1 インスタンスで実行する上で、CPU や GPU と異なる点は、 Inferentia、Trainium に合わせてコンパイルする必要があるという点です。Inf1 上で推論モデルを実行する場合は、AOT (Ahead-Of-Time) コンパイラで事前にモデルをコンパイルし、コンパイル済みのモデルを実際の推論処理で実行する必要がありました。

一方、Trn1 上でトレーニングを実行する場合は、 JIT (Just-In-Time) コンパイラである PyTorch XLA を利用します。PyTorch では通常、計算グラフを構築と同時に実行します（eager mode）、PyTorch XLA を利用した PyTorch Neuron では、計算グラフを構築する過程での実行は一旦保留し、実際に実行が必要となったタイミングでグラフを最適化、コンパイル、実行する Lazy mode が利用されます。（※デバッグ時には Eagar mode を有効化し、ステップバイステップで実行することも可能です。）

PyTorch から PyTorch XLA への移行は GPU デバイスを PyTorch XLA の XLA デバイス（AWS Trainium）に置き換えるだけですので非常に簡単です。次のセクションでは実際に PyTorch XLA を利用し Neuron コアを動かしてみます。

実際に動かしてみましょう

さて、前置きが長くなってしまいましたが、実際に Trn1 インスタンスを使用し、Neuronコア上での動作を体験してみましょう。

事前準備

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

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

（※ 2022 年 10 月 10 日時点でのインストール結果です。）

$ yum list installed | grep neuron
aws-neuronx-collectives.x86_64        2.9.86.0_f47c632a8-1           @neuron
aws-neuronx-dkms.noarch               2.5.41.0-dkms                  @neuron
aws-neuronx-oci-hook.x86_64           2.0.16.0-1                     @neuron
aws-neuronx-runtime-lib.x86_64        2.9.64.0_281d6d35d-1           @neuron
aws-neuronx-tools.x86_64              2.4.14.0-1                     @neuron

$ pip list | grep 'neuron\|torch'
aws-neuronx-runtime-discovery 2.9
libneuronxla                  0.1.354
neuronx-cc                    2.1.0.76+2909d26a2
neuronx-hwm                   2.1.0.7+64eaede08
torch                         1.11.0
torch-neuron                  1.11.0.2.3.0.0
torch-neuronx                 1.11.0.1.1.1
torch-xla                     1.11.0+torchneuron2

確率的な丸め処理を体験

ここでは Trn1 上で Neuron コアを利用するために必要なステップを、確率的な丸め処理（Stochastic rounding）の例を通して体験してみます。以下の内容で rounding_mode_test.py というファイル名の Python スクリプトを作成します。

import torch
import torch_xla
import torch_xla.core.xla_model as xm

devices = ["cpu", "xla"]

for device in devices:
    a = torch.tensor(1024.0).half().to(device)

    for i in range(2048):
       a = (a + 0.5)
       xm.mark_step()

    print("Result by {} = {}".format(device, a))

このスクリプトでは、同じ内容の処理を Trn1 インスタンス上の CPU、Neuron コアそれぞれで実行しています。

Neuron コアで実行するにあたって必要な変更点はわずか数行のみです。

Trn1 上でトレーニングを実行する際、PyTorch XLA をバックエンドのコンパイラとして利用することは説明しましたが、スクリプトの最初で、torch_xla ライブラリを import しています。

また既存の torch.Tensor のデバイス（GPU / CPU）を切り替えるには to()メソッドを使用しますが、Neuron コアに切り替えるためにはデバイス名として xla を指定します。

最後に演算の過程で xm.mark_step() という API を呼び出しています。本 API を呼び出すことで、その時点での計算グラフを明示的にコンパイルし実行します。（※xm.mark_step() で明示的に指定しない場合でも、演算の途中結果を print() 関数で表示させたい場合など、計算グラフを確定させる必要がある場合には、必要に応じてコンパイラが実行されます。）

それでは、まず通常の丸め処理を使ってスクリプトを実行してみましょう。実行前に、NEURON_RT_STOCHASTIC_ROUNDING_EN 環境変数を設定することで、標準的な丸め処理を行うか、確率的な丸め処理を行うか選択します。

$ NEURON_RT_STOCHASTIC_ROUNDING_EN=0 python3 rounding_mode_test.py

Result by cpu = 1024.0
..... No candidate found under /var/tmp/neuron-compile-cache/USER_neuroncc-2.1.0.76+2909d26a2/MODULE_8335261481194474129.
..... Cache dir for the neff: /var/tmp/neuron-compile-cache/USER_neuroncc-2.1.0.76+2909d26a2/MODULE_8335261481194474129/...
.
Compiler status PASS
..... Exiting with a successfully compiled graph
Result by xla = 1024.0

（※テスト結果は重要な部分のみを抜粋しています）

本サンプルスクリプトでは、FP16 データタイプを利用しています。ここでは詳細は省きますが、FP16 で表現できる範囲を超えるため、標準的な丸め処理では 1024 に対して何度 0.5 を加えても加算分は切り捨てられ、CPU、Neuron コアどちらの場合も、結果は 1024 のままとなっています。またテスト結果から Neuron コアで実行する際には、コンパイラが起動している点も確認できると思います。

次に、Trainium でハードウェア実装された確率的な丸め処理を使って実行してみましょう。

$ NEURON_RT_STOCHASTIC_ROUNDING_EN=1 python3 rounding_mode_test.py

Result by cpu = 1024.0
..... Using a cached neff at /var/tmp/neuron-compile-cache/USER_neuroncc-2.1.0.76+2909d26a2/MODULE_8335261481194474129/.....5ea6d34f0229d.neff. Exiting with a successfully compiled graph
Result by xla = 2056.0

今度は 50%の確率で丸め処理後の結果が切り上げられるため、結果としては 2048 により近い 2056 を得ることができました。また先程の実行で得た計算グラフのコンパイル結果（NEFF : Neuron Executable File Format）はキャッシュとして保持されていたため、再度のコンパイルは発生していません。計算グラフによってはコンパイルに時間がかかる場合がありますので、一度コンパイルした結果を保持して再利用することは大変重要です。

まとめ

遂に一般提供を開始した Amazon EC2 Trn1 インスタンスが加わり、大規模 ML トレーニングを実行するためのインスタンスの選択肢が広がりました。本ブログでは「基礎編」として Trn1 インスタンス、AWS Trainium ML チップ、AWS Neuron SDK の概要を紹介しました。本ブログは２部構成となっています。第２部では「実践編」として Trn1 インスタンスを利用して日本語 BERT モデルをトレーニング（ファインチューニング）する過程及び AWS Neuron が提供するツール群を紹介する予定です。

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

________________________________________

Amazon EC2 Trn1インスタンス参考コンテンツ

Amazon EC2 Trn1 インスタンス (製品ページ)
AWS Trainium (製品ページ)
AWS Neuron (ドキュメント)
独自設計チップ AWS Trainium 搭載 Amazon EC2 Trn1 インスタンスで ML トレーニングを高速実行（実践編）

Amazon Web Services ブログ