第 2 章 Trainiumの基礎知識編
本章では以下を仮定します。
- Transformer アーキテクチャの基本的な理解
- HuggingFace Models の基本的な理解
問題 (8〜10)
Trainium は2021年頃にAWSが開発した「深層学習特化チップ」（ハードウェアアクセラレータ）です。行列演算を高速に行えるという点ではGPUもそうですが、それよりももっとコストパフォーマンスが良いです。ただし前章で述べた通り、計算グラフをコンパイルする必要があり、それ由来の制約があります。
ノード・コア・チップ
AWS Trainium の現行最新版である
Trn1が搭載されている主なインスタンスには、以下があります。コストと搭載メモリ量は以下の通りです。（us-east-2でのオンデマンド価格・2025-05-19調べ）
| インスタンス | Trn総コア数 | Trn総メモリ | CPUメモリ | コスト/h |
|---|---|---|---|---|
| trn1.2xlarge | 2 | 32GB | 32GB | $1.34 |
| trn1.32xlarge | 32 | 512GB | 512GB | $21.50 |
※ 上表のコストは「オンデマンド」の場合ですが、事前にキャパシティを予約し前払いとなる EC2 Capacity Blocks を利用すれば、より安価に利用が可能です。
Trn1インスタンスでは、インスタンス名の{num}xlargeのnumと総コア数が一致しています。
-
Trn1の「1コアあたりのメモリ量」は何GBですか？上の表をもとに計算して、暗記してください。解答を表示
Trn1のメモリ量は 1コアあたり16GB です。
補足します。
- 「チップ」は上表には記載がありませんが、「1チップ」=「2コア」(=32GB)です。
- trn1n インスタンスについて
n無し版よりも、ノード（=インスタンス）間のネットワークがスムーズになっており（EFA）、100以上のノードを同時に連携させるような大規模学習を行う場合には推奨されます。- Trainiumチップの次期バージョンについて
Trn2が2024年12月に公開されました。Trn1よりもパフォーマンスがさらに優れています（詳細）。今後、Trn3が2025年後半に公開される予定です。- Inferentia2 について
- Trainium(1,2) とは別に Inferentia2 というチップもあります。こちらは Trainium と比較して訓練より推論に特化したチップですが、訓練ができないというわけではありません。例えば、Trainium1 と Inferentia2 の実質的な違いは、ネットワークの帯域の違いくらいです（Trainium1 の方が通信が早い）。
-
Qwen2.5-7B-Instruct の推論デモサーバーを
Trn1上でホストしたいです。最大系列長は8Kとしたいです。この時、必要メモリ量を見積もって、どのインスタンスを立ち上げるのが最適かを判断してください。ただし、標準的な transformer ベースのLLMの推論時に必要となるメモリ量はパラメータ数×2Bytes + ((語彙サイズ)+2×(key-valueのヘッド次元)×(1層あたりのkey-valueヘッド数)×層数)×バッチサイズ×シーケンス長×2Bytesで見積もることができます（参考記事）。解答を表示
以下は config.json からの抜粋です：
{ "hidden_size": 3584, // 隠れ次元（=ヘッド次元 × 1層あたりのattentionヘッド数） "num_attention_heads": 28, // 1層あたりのattentionヘッド数 "num_hidden_layers": 28, // 層数 "num_key_value_heads": 4, // 1層あたりのkey-valueヘッド数 "vocab_size": 152064 // 語彙サイズ }上記からそれぞれの値が判明しますので、それらを代入すると
7G × 2 + (152064 + 2 × 128 × 4 × 28) × 1 × 8K × 2 = 14G + 3G = 17Gとなります。このため、Trnメモリを 32GB 持っているtrn1.2xlargeでホストが可能です。（ただしシングルコア(=16GB)には乗り切らないため、何らかのモデルパラレルを利用する必要があります。） -
Qwen2.5-7B-Instruct の訓練を
Trn11ノードで実施したいです。マイクロバッチサイズは1で、最大系列長は8Kとしたいです。この時、必要メモリ量を見積もって、どのインスタンスを立ち上げるのが最適かを判断してください。ただし、標準的な transformer ベースのLLMの学習時に必要となるメモリ量は3×パラメータ数×2Bytes + 2×(1トークンあたりの全中間表現の次元の総和)×バッチサイズ×シーケンス長×2Bytesで見積もることができます（参考記事）。なお1トークンあたりの全中間表現の次元の総和を正確に見積もるのはちょっと面倒ですが、ここでは層数×10×隠れ次元で近似することにします。解答を表示
上式に代入すると
3 × 7G × 2 + 2 × (28 × 10 × 3584) × 1 × 8K × 2 = 42G + 32G = 74GBとなります。このため、Trnメモリを 32GB 持っているtrn1.2xlarge1ノードでは学習不可能です。trn1.32xlarge1ノードであれば（適切なモデルパラレルの下で）学習可能です。
Trn/Infで学習や推論を行うには、モデル（計算グラフ）をTrn/Infに載せる必要があります。開発はモデルアーキテクチャごとに必要となります。メジャーなモデルアーキテクチャ（例：
LlamaForCausalLM）を使用する場合は、公式で開発済のケースもあります。しかし、そうでない場合にはこの開発が最も大変となります（これだけで2-3人月程度かかる場合があります）。
-
本50本ノックの内容を監修くださった AWS の常世様に、感謝を申し上げます。 ↩︎
Discussion