AWSが発表したTrainium2から見る、AIアクセラレータのGPUに対する勝算は？

Trainium2

　AWS(Amazon Web Services)は、12月2日～12月6日(現地時間)に同社の年次イベント「re:Invent 2024」を米国ネバダ州ラスベガス市で開催した。この中で、AWSは同社が「Trainium」(トレニアム)と呼んでいるAI演算に特化したAIアクセラレータの最新製品となる「Trainium2」の詳細を明らかにした。

　Trainium2は、初代Trainium(以下、混乱しないようにTrainium1と表記する)に比べて1つのダイあたりの演算コア(NeuronCore)数が2倍(4コア>8コア)になり、1パッケージに2つのダイが統合されており、1つのパッケージあたりの演算コアは従来製品に比べて4倍(4コア>16コア)になっている。

　こうしたTrainium2は、GPUに比べると電力効率が高く、それに合わせて提供コストが下げられることで、AWSは「効率」を武器に、GPUの高コストや入手性に不満を感じている企業などへの浸透を目指す方針だ。

AWS子会社のAnnapurna Labsが開発しているGraviton、Inferentia、Trainium

　AWSは近年カスタム半導体に力を入れており、Arm CPUとなる「Graviton」に早くから取り組んできたことはよく知られている。2018年に初代Gravitonを発表し、2019年には「Graviton2」を、そして2021年には「Graviton3」、2022年には「Graviton3E」、そして2023年には「Graviton4」を発表して、今に至っている。

Graviton4

Graviton3

　そうしたAWSのGravitonが成功を収めたことは、昨年(2023年)の11月にMicrosoftがCobalt、今年(2024年)の4月にGoogle CloudがAxionというArm CPUを相次いで発表して追随したことを見れば明らかだろう。

　AWSのカスタム半導体を設計しているのは、AWSが2015年に買収することを発表し、子会社化したイスラエルの半導体設計ベンダー「Annapurna Labs」だ。Annapurna LabsがAWSのために最初に設計したのは、「Nitro」(ナイトロ)と呼ばれる、今で言うところのDPU(Data Processing Unit)になる。

　Nitroは、仮想マシンを利用する時に発生するオーバーヘッドを処理し、CPUに無駄な負荷がかからないようにすることでCPUの実行効率を上げるネットワークアダプタを実現する。そうしたNitroをいち早く導入することで、AWSは顧客がEC2(AWSの仮想化されたハードウェアをクラウドベースで提供するIaaSのこと)におけるサービス実行効率を上げてきている。

　それに次いで、Annapurna Labsが取り組んだのがArm CPUのGravitonシリーズになる。その最新の製品Graviton4は、ArmのIPデザインである「Neoverse V2」に基づいて、最大96コアのArmv9 CPUコアが1つのパッケージ上に実装されている。

　そうしたNitro、Gravitonに次ぐ製品としてAnnapurna Labsが開発したAI向け半導体製品が、2019年に発表されたAI推論向けとなる「Inferentia」(インファレンシア)、2021年に発表されたAI学習向けとなる「Trainium」(トレニアム)となる。

第1世代のInferentiaは2019年に、第2世代のTrainiumは2021年に発表される

Inferentia(左)、Trainium(右)

　このInferentia/Trainiumは、一言で言えばAI演算時に行なわれる演算処理を専門に行なう“AIアクセラレータ”という半導体製品となる。Inferentia/Trainiumには「NeuronCore」(ニューロンコア)と呼ばれる演算コアが複数内蔵されており、それぞれのニューロンコアにはテンサーエンジン、浮動小数点エンジン、整数演算エンジンとローカルメモリとして利用できるSRAMメモリが搭載されており、AI演算に必要な処理を高速に行なえる。

　AI演算、特に学習で利用されるGPUとの違いは、ソフトウェアの自由度だ。GPUのソフトウェアは、NVIDIAのCUDA、AMDのROCmがその代表例と言えるが、ソフトウェア開発キットを利用することで、AIの学習/推論だけでなく、気候変動や石油採掘、空力開発などの科学シミュレーション、いわゆるHPC(High Performance Computing)と呼ばれるアプリケーションにも利用できる。それに対して、AIアクセラレータではAI演算に特化しており、基本的にはほかの用途に利用できない。

　こうしたInferentia/Trainiumは、現在まで3世代の製品がリリースされている。それが2019年に登場した初代Inferentia(以下Inferentia1)、2世代目がTrainium(以下Trainium1)およびほぼ同じ設計で推論向けと位置づけられるInferentia2、そして今年一般提供開始となったTrainium2になる。それぞれの製品の、ハイレベルの違いを表にすると以下のようになる。

Inferentia1

　Inferentia1はNeuronCoreの第1世代(NeuronCore-v1)を4つ搭載しており、AWSがNeuronLinkと呼んでいるチップ間インターコネクトで接続して、スケールアップできる。このNeuronLinkは何かと言うと、NVIDIAのNVLink、AMDのInfinity Fabricのような複数のチップを相互に接続してスケールアップするためのインターコネクトと同種だと考えると理解しやすい。

　AWSはこのInferentia1をAI推論専用と位置づけている。というのも、メモリは8GBと少なく、(用途にはよるが)推論であれば十分だが、大規模な学習に利用するにはやや容量が少ないからだ。

Trainium1

　このInferentia1を進化させたのが、Trainium1となる。Trainium1は、演算コアがNeuronCore-v2というNeuronCoreの第2世代へと進化している。NeuronCore-v1との大きな違いは、新しくGPSIMD(General Purpose SIMD)と呼ばれる512bitの浮動小数点演算プロセッサを内蔵していることで、これによりソフトウェアによるより柔軟性の高い演算が可能になることだ。

　ほかにも、テンサーエンジン、浮動小数点エンジン、整数演算エンジンにも手が入っており、コアのダイサイズが大きくなったためか、Trainium1では演算コアは2コア構成になっており、Inferentia1の4コアに比べて減っていることが大きな違いになる。ただし、FP16のスループットはInferentia1の64TFlopsと比較して約3倍の191TFLOPSに引き上げられている。

　また、AI学習用と位置づけられていることもあり、メモリは32GBのHBMへと強化されている。メモリ帯域も50GB/sから820GB/sへと大幅に引き上げられており、NeuronLink-v2というより強化されたインターコネクトを利用してスケールアップすることで、より大規模なAI学習に利用することが可能になっている。Trn1というインスタンスでは最大16基のTrainium1をスケールアップして利用することが可能で、最大3PFLOPS(FP16時)の性能でAI学習を処理することが可能になっている。

　なお、Trainium1とほぼ同じスペックのチップがInferentia2として提供されている。Inferentia2のインスタンスとなるInf2は、最大12基までスケールアップできると説明されており、チップあたりの帯域幅もTrainium1が384GB/sであるのに対して、Inferentia2では192GB/sと半分に制限されている。Trainium1のNeuronLinkのスペックを機能制限した製品がInferentia2であると考えられるだろう。

5nmに微細化されたTrainium2、演算コア、チップレット、インターコネクトが強化ポイント

　今回AWSが発表したTrainium2は、Inferentia1から数えて第3世代のAIアクセラレータとなる。特徴は3つあり、NeuronCoreが第3世代(NeuronCore-v3)へと進化していること、チップレットを採用して1つのパッケージあたり2つのダイが搭載されていること、そしてインターコネクトがより高速なNeuronLink-v3へと強化されていることだ。

左がTrainium 1、右がTrainium2、Trainium2は1パッケージに2ダイ搭載されている

　NeuronCore-v3は、テンサーエンジン、浮動小数点エンジン、整数演算エンジン、512ビットのGPSIMDという4種類の演算器を内蔵していることはNeuronCore-v2と同じだが、ローカルメモリ(CPU的な言い方をするとキャッシュメモリ)となるSRAMが28MBまで増やされているのが大きな特徴となる。

　さらに、テンサーエンジンが、いわゆるスパース性に対応している。これは行列演算時に不必要なデータを圧縮して演算する手法で、NVIDIAのGPUでもAmpere世代(A100など)以降で採用されており、浮動小数点演算時のスループットを大幅に改善させることができる。NeuronCore-v3では4:16、4:12、4:8、2:8、2:4、1:4、1:2というスパース性のパターンに対応しており、FP16で活用すると、約3.8倍のスループットを実現する。

　Trainium1(およびそのバリエーションと考えられるInferentia2)では、パッケージの中に1つのダイとHBMメモリが混載される形になっていた。それに対して、Trainium2では、ダイが2つ、そしてそれぞれにHBMメモリが混載される形になっている。

　またTrainium1では、1つのダイ(=パッケージ)に2つのNeuronCore-v2が搭載されていたが、Trainium2ではNeuronCore-v3が1つのダイに4つ、そしてそのダイが2つあるのでパッケージレベルでは8つのNeuronCore-v3が搭載されている形になる。

　これにより、大きく性能が引き上げられており、FP8時ではTrainium1比で6.8倍になる1,299TFLOPS、FP16時はTrainium1比で3.49倍となる667TFLOPS、そしてスパース性を利用するとFP8/FP16/BF16/TF32時に2,563TFLOPSというスループットを実現する。

　なお、こうした設計が可能になったのも、そもそもパッケージが大きくなっていることと、製造プロセスノードがTrainium1の7nmから5nmへ微細化されたためだと考えられる。

　最後に新しいインターコネクトとなるNeuronLink-v3に対応していることが挙げられる。NeuronLink-v3ではチップあたりの帯域幅が従来製品の約3.4倍に引き上げられている。これを利用して、かなり多くの数までスケールアップすることが可能で、今回Trainium2のインスタンスとして発表されたTrn2では16基まで、さらにそのTrn2を4台つないで64基までスケールアップしたTrn2 UltraServerでもNeuronLink-v3でスケールアップされている。仮にそれ以上が必要な場合には、EFAv3というAWS独自のイーサネットを利用してスケールアウトすることが可能になっている。

Trn2 UltraServer、2つのラックに64基のTrainium2が搭載され、NeuronLink-v3でスケールアップされている

　NeuronLink-v3ではリング構造になっており、16基までは1.28TB/sというほぼフル性能で周りのTrainium2同士が接続して、ネットワークを構築できる。さらに、64基の場合には256GB/sと接続速度は遅くなるが、追加のスイッチなどは必要なくスケールアップすることが可能だ。

　NVIDIAのH100などに採用されているNVLinkの場合は8基までスケールアップするであっても、NVSwitchと呼ばれる高価なスイッチチップが必要(H100/200世代では4つ、B200世代では2つ)になることを考えると、この点はコスト面でTrainium2のアドバンテージと考えられる。

Trainium3

　なお、今回のre:Invent 2024においてAWSはTrainium2の後継製品としてTrainium3を発表し、来年の末までにEC2インスタンスとして提供を開始する計画であることを明らかにした。Trainium3は、プロセスノードが3nmに微細化され、電力効率が40%改善し、性能は2倍になると明らかにしている。現状では明らかになっているのはこの3つだけで、来年にTrainium3ベースのインスタンスが発表される時にもう少し詳細が明らかになるものと考えられる。

ソフトウェア開発環境としてInferentia /Trainium版CUDAとなる「AWS Neuron SDK」を提供

　こうしたハードウェアが用意されていても、動かすプログラムがなければタダの箱になってしまう。NVIDIAの強みがCUDAというハードウェアとAIフレームワーク(TensorFlowやPyTorchなど)との間に入るミドルウェアであることを考えれば、そこの重要性は筆者が強調する必要もないだろう。

　AWSの場合には、Inferentia やTrainiumを利用するためのソフトウェア開発環境として「AWS Neuron SDK」(以下Neuron SDK)を用意しており、このNeuron SDKがNVIDIA GPUでのCUDAに相当するミドルウェアとなる。Neuron SDKは、PyTorch、JAX、TensorFlowという3つのメジャーなフレームワークに対応しており、それぞれに対応したNeuron SDKを利用することで、Inferentia やTrainiumを利用したAI推論や学習を行なわせることができる。

　もちろん、AWSが提供しているAmazon SageMakerなどのマネージドサービスからも利用できる。これらのマネージドサービスでは、ハードウェアの設定などはすべてAWS側で行なうため、顧客はInferentia やTrainiumを利用していることすら意識することなく、学習や推論などを行なえる。

　AWSはこうしたNeuronなどCUDAに匹敵するような開発環境を、2019年のInferentiaのリリース後に徐々に整備してきた。もちろん、CUDAとはプログラミングモデルが異なるため、若干の学習は必要になるが、CUDAをNeuronに置きかえるだけで、PyTorch、JAX、TensorFlowというフレームワークはそのまま活用してAI学習やAI推論を行なえる。その意味で、ソフトウェアの開発環境に関しても、より大規模なプログラムを作り込みたい顧客のニーズにも対応できるようになっている。

GPUの弱点では総電力量とコスト、それを補うのがTrainiumのようなAIアクセラレータ

マット・ガーマンCEOの基調講演の中で発表されたBlackwellを搭載したEC2インスタンスの「P6」

　このようなInferentia、TrainiumといったAIアクセラレータは、汎用に使えるNVIDIAのH100/H200/B200のようなGPUに肩を並べていくことが可能なのだろうか？

re:Invent 2024の展示会で展示されていたGB200のボード

　市場の現状はどうなのか言えば、それはAWS CEO マット・ガーマン氏のコメントを聞くだけ十分だろう。ガーマン氏は同社が12月3日(現地時間)に行なったre:Invent 2024の基調講演において「AI学習ニーズのほとんどはGPUだ」と述べ、AI学習ではほぼGPUの独占市場であると述べ、実際同社もNVIDIAが3月に発表したBlackwellベースのEC2インスタンス「P6」を来年(2025年)の初頭以降に提供すると明らかにしている。

　Blackwellの投入を急いでいるのはAWSだけでなく、Microsoft Azureも、Google Cloudも、Oracle Cloud Infrastructure(OCI)も競ってBlackwellのインスタンスの導入を発表しているような状況で、この「AI学習ならGPU」という方程式が簡単に崩れそうにはないというのが現状だ。

　では、AIアクセラレータであるTrainiumはどこを狙っているのかというと、簡単に言えばGPUの弱点をカバーする市場だ。というのも、今世界中のデータセンターでNVIDIA GPUを導入する競争が起こった結果、データセンターが利用する電力の総量が増え続けており、今後データセンターを作るには、同じような地域に原子力発電所を併設しないと難しくなっていくかもと言われているような状況が発生している。

　確かにGPUは、CPUなどに比べてAI学習時の性能当たりの電力が高いのはその通りだが、BlackwellではTDPが1,000Wを超えるようなスペックになっているため、ピーク時の消費電力はもっと高く、何よりそれを冷やすのにもはや空冷では不可能で、水冷にするなど、そちらの冷却施設のための電力が別途必要になると、ピークの電力は増える一方だからだ。

　もう1つの弱点はコスト。NVIDIAのAI学習に最適なGPUという市場はNVIDIA自身が作り上げ、ほかの企業との競争を勝ち抜いて現在のポジションを作り上げたわけだが、しかし現状競合と言えるような競合はAMDの「Instinct MI300」シリーズぐらいしかないというほぼ独占市場になっている。競争がなければ価格が高止まりするのは市場の原理であり、それに不満をもっている顧客も少なくないと考えられる。

　そうした消費電力とコストの観点でTrainiumは勝っていると考えられている。ただし、Trainiumは外販されていないため、直接価格などは比較できず、あくまでAWSの主張によればという意味になる。

　実際、re:Invent 2024にAWSのInferentia/Trainiumとして登壇したApple マシンラーニング・AI 上級部長 ビノット・デュピン氏は「すでにTrainium2をアーリーアクセスプログラムで利用してきたが、従来の環境に比べて50%の効率が改善されたことが確認できた」と述べている。もちろんAppleは何と比べて50%かは明確には述べていない(つまり従来のTrainium1と比べての可能性もあるし、GPUの可能性もあるということだ)のだが、効率が50%改善というのは大きな意味がある。

　一般的に顧客がAWSのようなCSP(クラウド・サービス・プロバイダー)に支払うコストというのは、半導体の利用料金と電力の利用料金の両方が含まれている。CSPの顧客がいう「効率」には性能も含まれるだろうし、「コスト」も含まれると考えられる。その両方が改善されれば、顧客にとってはAI学習にかかる時間も、そしてコストも削減できることを意味し、大きな効果がある。その意味で、GPUの弱点(電力とコスト)に不満を感じている顧客にとっては、CSPが提供するAIアクセラレータを真剣に検討してみる意味があると言える。

　今後そうした顧客が増えていくかどうか、それがAIアクセラレータに取り組んでいるCSPにとって重要になってくると言えるだろう。