GroqのLPUの基本アーキテクチャを深く掘り下げる
感じたこと
論文より幾分わかりやすいが、まだまだ難しい...
概要
私たちはここ最近、システムのパフォーマンスについて多くの議論を重ねてきました。数日前、Groq社が彼らの言語処理ユニット(LPU)ハードウェアを用いて大規模言語モデル(LLM)の推論ベンチマークで記録を破ったというニュースがありました。LPUの成功は、革新的なハードウェアアーキテクチャの設計と、それに合わせて強力なコンパイラを組み合わせることによるものです。私たちがパフォーマンスに注目している今、LPUのハードウェアとコンパイラの下層にあるものを理解することは、私たちの議論にとって非常に重要だと考えています。 本稿では、まずTSPのアーキテクチャとそのコンパイラを分析し、次にGroq社がこれらのTSPを用いて、どのように信頼性が高くスループットの大きい分散型AI推論エンジンを構築したのかを見ていきます。
GPUやCPUと比較して、TSPのアーキテクチャはかなり異なっており、その主な目的はTSPのハードウェアをより決定論的にすることにあります。まずはCPUやGPUにおける非決定論の原因について話しましょう。
CPUとGPUのマイクロアーキテクチャにおける非決定論の理解
CPUとGPUのマイクロアーキテクチャには、これらのデバイス上での命令の実行を非決定論的にする多くの機能があります。つまり、特定の命令がいつ実行されるのか、完了までにどのくらいの時間がかかるのか、結果がいつ利用可能になるのかを保証することができないのです。
💡レストランのキッチンにたとえる。
シェフ(CPU)は複数の料理人(命令)に異なる料理(タスク)を割り当てる。しかし、いくつかの要因により、料理が完成するまでの時間を正確に予測することは難しい。
・オーダーの順番が変わる(アウトオブオーダー実行)
・材料が届くのが遅れる(メモリアクセスの遅延)
・料理人の手際が日によって変わる(クロックスピードの変動)
・複雑な料理は準備に時間がかかる(命令の複雑さ)
・料理人が次の手順を予測して準備を始める(投機的実行)
このように、料理(命令)が完成(実行)されるまでの時間は、その日の状況によって変化する。お客様(プログラマ)は料理が提供される正確な時間を知ることができない。
例えば、最新のCPUには以下のような機能があります。
スーパースカラアーキテクチャ:1サイクルに複数の命令を発行することが可能
アウトオブオーダー実行:命令をどのような順序でも実行可能
投機的実行:分岐の場合、分岐条件がtrueかfalseかを推測し、スループットを向上させるために事前にその分岐を投機的に実行します。間違っていた場合は、現在の作業を破棄してから他の分岐パスの実行に戻る必要があります。
命令パイプライン:命令の実行を複数のステージに分割し、パイプライン方式で実行します。1つの命令がパイプラインの2番目のステージに移動すると、CPUは新しい命令をパイプラインの最初のステージに入れることができます。これもまた命令スループットを向上させます。
複数レベルのキャッシュ:命令がメモリからデータをロードする必要がある場合のレイテンシを削減するために、CPUには2~3レベルのキャッシュがあります。
これらはすべて、CPUにおける命令実行の順序とタイミングを非決定論的にし、推論するのが難しくなります。GPUにもキャッシュ、共有メモリ、グローバルメモリ、動的リソース分割など、非決定論の要因がいくつかあります。
これらのデバイスに内在する非決定論の問題は、プログラムのパフォーマンスを推論することが難しく、最悪の場合のパフォーマンスの限界を保証することが難しいことです。
その結果、Groq社は、大量の並列性を持ちながら非決定論的な動作がゼロであるという、テンサーストリーミングプロセッサ(TSP)のための全く新しい設計を考案しました。彼らはプログラムのパフォーマンスに限界を設けることを保証するために、ハードウェアから多くの複雑さを取り除き、命令の正確なスケジューリングと実行をより制御できるようにコンパイラに力を与えました。TSPのアーキテクチャが内部からどのようなものかを理解しましょう。
テンサーストリーミングプロセッサ(TSP)のアーキテクチャ
TSPのハードウェア設計は、CPUやGPUの設計とは対照的です。従来のマルチコアチップはタイル状のアーキテクチャを持ち、各タイルが1つの処理コアを表しています。内部的には、各タイルは異なる種類の計算(例えば算術演算、メモリ演算、論理演算、命令制御など)を行うための機能ユニットの集合体で構成されています。
TSPの設計者は、この従来の設計を内側から外側へとひっくり返しました。彼らは機能ユニットをコアの外に移動させ、2次元グリッド状に配置しました。このグリッドの各列は特定のタイプの機能ユニットを保持しており、スライスと呼ばれています。以下の図は、従来のマルチコアチップのタイル状の設計とTSPの設計を並べたものです。
https://scrapbox.io/files/662c60849b6830002d99e53d.png
TSPには以下の機能スライスがあります。
MXM:行列演算の実行
SXM:ベクトルのシフトとローテート演算
MEM:メモリの読み書き演算
VXM:ベクトルに対する算術演算
ICU:Instruction Control Unit - 他のスライスとは異なり、これは水平方向に組織され、他のスライスでの実行のための命令のフェッチとディスパッチを担当します。
TSPのアーキテクチャを理解したところで、その中核となる動作モード、つまり命令の実行に焦点を当ててみましょう。ここでGroq社のエンジニアリング力が真に輝いており、ハードウェアの設計が実行能力によって直接的に支えられていることが実証されています。
TSPにおける命令実行
TSPはSIMD(Single Instruction Multiple Data)方式で命令を実行します。各機能スライスは20個のタイルで構成され、各タイルは16要素のベクトルを生成することができます。したがって、1つの完全なスライスは最大で320要素のベクトルを処理し、生成することができます。
これらのスライスは生産者と消費者の関係で相互作用し、あるスライスがその結果をストリームとして生成し、別のスライスがそのストリームを入力として受け取って処理します。
ベクトルがメモリから読み込まれると、ストリームID(0~31)とフローの方向(東または西)が割り当てられます。ストリームは、あるスライスから次のスライスへと流れていきます。各スライスは自由にそのストリームを処理して新しい結果ストリームを生成するか、そのストリームをそのまま隣接するスライスに流すことができます。
https://scrapbox.io/files/662c61bdf8e29d002536ca42.png
TSPの各タイルはベクトルストリームの16要素を処理することができるので、フルベクトルを効率的に処理するために、命令の実行はパイプライン化されます。タイムスタンプt1で、ICUは機能スライスの1つの最下部タイルに命令を発行します。次のタイムスタンプでは、そのタイルの出力が次のスライスに北向きに移動しますが、その間にICUはベクトルの次の部分に対する別の命令を最下部タイルに発行します。これにより、以下の図に示すように、命令の実行がずらして行われます。
https://scrapbox.io/files/662c61cd83384b002604fb41.png
TSPが命令を実行する精度は、ハードウェアの産物であるだけでなく、コンパイラとの相乗効果の産物でもあります。コンパイラと命令セットアーキテクチャを調べることで、ソフトウェアとハードウェアがどのようにパフォーマンスの最適化のために融合しているかがわかります。
TSPのコンパイラとISA
TSPの設計者は、ハードウェアを簡素化し、命令実行に関連する多くの複雑さをコンパイラに押し付けました。コンパイラは、与えられたプログラムを正しく実行し、可能な限り効率的な方法でそれを行うために、命令とデータフローを正確にスケジューリングする必要があります。
コンパイラは、TSPハードウェアの以下のアーキテクチャ的に可視な状態にアクセスできます。
320レーンのプログラミング抽象化:TSPチップの各タイルは、SIMD方式で16要素(16レーン)のベクトルを操作することができます。垂直スライスは20個のそのようなタイルで構成されているため、実行に利用可能な合計320個のSIMDレーンがあります。
144個の独立した命令キュー:チップ上には144個の命令キューがあり、それぞれがサイクルごとに1つ以上の命令を発行することができます。コンパイラは、これらの各キューのプログラム順序を完全に制御します。
レーンごとに64個の論理ストリーム:各レーンは、オペランドや結果の移動に使用できる64個の論理ストリームにアクセスできます。これらは東向きと西向きに分割されており、つまり32個は東向きのデータ移動に使用でき、残りの32個は西向きのデータ移動に使用できます。
220 MiByteのグローバル共有SRAM
TSPハードウェアには非決定論的な動作がないため、コンパイラは各命令のレイテンシを正確に把握しています。また、コンパイラはコンパイル中のプログラム(例えば、深層ニューラルネットの計算グラフ)のデータフローを完全に把握しています。
コンパイラは計算タスク間の依存関係を特定し、それに基づいてTSP上の利用可能な機能ユニットへの並列実行を割り当てます。
コンパイラは、スライスがストリームを処理する準備ができたときにちょうどストリームがスライスに到着するように、命令をスケジュールします。コンパイラは、命令のレイテンシとデータフローの方向を正確に知っているので、これを行うことができるのです。
概念的には、コンパイラは時間と空間の両方で命令とデータの2次元スケジューリングを解決しています。TSPのプログラミングモデルは、以下の2つの重要な要素に依存しています。
ハードウェアの決定論的なデータパス
ISAを通じて命令のレイテンシに関する時間的情報を公開すること
これらのおかげで、コンパイラのバックエンドはチップ上の任意のストリームの位置と使用時間を追跡することができます。この論文では、これを「ソフトウェアで定義されたハードウェア」と呼んでいます。
TSPからLPUへのスケーリング
ここまでTSPのアーキテクチャについて見てきましたが、それはハードウェアレベルでは大量に並列化されているものの、非決定論的な機能は一切ないことがわかりました。コンパイラは、ベクトルデータをストリームの形でSIMD方式で処理するように、異なる機能スライスに命令とデータフローのスケジューリングをすべて行います。
TSPはLanguage Processing Unit(LPU)の基本ユニットであることがわかりました。多くのTSPがラックの形で組み合わされ、そのようなラックが多数接続されて、膨大なスループットを実現する分散システムを形成しています。Groq社がTSPから分散システムをどのように構築したのか、そしてそれがどのように機能するのかを理解しましょう。
マルチTSPシステムの設計:物理面と分散面
TSPと同様に、分散マルチTSPシステムの設計目標も、決定論的なデータフローと命令実行、およびノード間の低レイテンシ通信を中心としています。
分散TSPシステムの設計は、ノードから始まります。ノードは、1つのシャーシ内に収められた8つのTSPデバイスから形成されます。これらの各デバイスは11本のピンで構成され、そのうちの7本のピンは各TSPデバイスをノード内の他の7つのTSPデバイスに接続するために使用され、残りの4本のピンはグローバルリンクを形成するために使用されます。
ノード内の各デバイスには4つのグローバルリンクがあるため、合計32のグローバルリンクがあり、これらが32の仮想ポートを持つハイラディックスルータを形成します。
ハイラディックスルータは、高性能の分散システムに必要な大量の接続、高帯域幅、高性能をサポートします。
9つのTSPノードを8つのTSPずつ組み合わせて1つのラックを形成します。ラック内の各ノードには32のポートがあるため、ラック全体では32 x 9 = 288のグローバルポートを持つことになります。これらのポートのうち144個は、ラック内の高速なデータフローのためにノード間のフル接続性を提供するためにローカルに使用され、残りの144個のポートは他のラックへの接続に使用されます。
最大構成のシステムでは、10,440個のTSPで構成される145個のラックを相互接続でき、システム内の任意の2つのTSP間で最大5ホップで接続できます。
https://scrapbox.io/files/662c653d4e33ac0024244dcc.png
図:8つのTSPでノードを形成し、9つのノードでフルラックを形成する方法。2022年のGroq社の論文から引用
TSPベースの分散システムにおける決定論の実現
このスケールアップされた分散システムの体制では、単一のTSPの機能ユニットが大規模な並列プロセッサの単一の処理コアとして機能します。TSPの計算モデルは決定論的なハードウェアに基づいており、同じ動作を分散システムの体制にも拡張する必要があります。
この決定論を実現し、大規模に分散したシステムの中で個々のTSPユニットを同期させるためのメカニズムがいくつかあります。それらについて見ていきましょう。
Hardware Aligned Counters(HAC)を用いたTSPのクロック同期
各TSPデバイスには、オーバーフロー周期が256サイクルのハードウェアカウンタ(Hardware Aligned Counter, HAC)が含まれています。TSPはこれを使って以下の手順で互いに同期します。
2つのTSPが相互接続されると、一方がそのHAC値をもう一方に転送します。値を受け取ると、受け取った側はその値を送信側に送り返します。送信側は、現在のHAC値と送り返された値との差を観測します。この差が、2つのデバイス間のリンクのレイテンシになります。このプロセスを何度か繰り返して、2つのTSP間の平均リンクレイテンシを確立します。
その後、2つのデバイスは親子関係に配置されます。親T0は定期的に現在のHAC値(HAC0)を子T1に送信します。子は平均リンクレイテンシをHAC0に加算し、それを自身のHAC値と比較します。2つの値の差は、継続的なクロックドリフトによる初期のずれを表します。T1はこの差を減らすようにHAC値を調整します。このプロセスを何度か繰り返すと、2つのTSPのHAC値はリンクレイテンシのジッタを表す小さな近傍内に収束します。
このプロトコルにより、2つのTSPが互いに同期することができ、ネットワーク内のTSPの親/子のスパニングツリーを確立することで、マルチホップネットワークに拡張することができます。
初期プログラムのアライメント
マルチTSPシステム上でプログラムを実行するには、まずシステム全体でデータフローと命令実行を正しくスケジューリングするために、すべてのTSPを整列させる必要があります。これには以下のようなメカニズムが含まれます。
個々のTSPのレベルでは、複数の独立した機能ユニットと144個の独立した命令キューがあります。まず、これらを同期させる必要があります。これを行うために、TSPはバリア同期のためのSYNCとNOTIFY命令をサポートしています。この方式では、SYNC命令がすべての命令キューを停止状態にし、そのうちの1つのキューが通知役として機能します。通知役がNOTIFY命令を発行すると、それがチップ上のすべてのキューにブロードキャストされ、その時点でそれらは同期され、動作を再開します。
SYNC/NOTIFYメカニズムは単一のTSPでは機能しますが、マルチTSPシステムではさらなる作業が必要です。2つのTSPはHACを使って互いに同期しますが、それに加えて各TSPはDESKEW命令をサポートします。機能ユニットがDESKEW命令を実行すると、TSPのHACがオーバーフロー(エポック境界)するまで、後続の命令の処理を停止します。これにより、次の命令がTSPによってエポック境界で実行されることが保証されます。2つのTSPは以下のメカニズムを用いてプログラムの実行開始を整列させることができます。
子TSPは、親TSPからのベクトルの到着をエポックごとにチェックするポーリングループに入ります。
親TSPは、まずDESKEW命令を実行することでプログラムの実行を開始します。これにより、プログラムの最初の命令がエポック境界で実行されることが保証されます。
DESKEWの後、親TSPはエポック境界で子TSPにベクトルを送信します。子TSPはベクトルを受信するとそれをバッファに格納します。
次のエポック境界で、子TSPはポーリングループを抜け、RECEIVE命令を実行してベクトルを処理します。
この方式により、2つのTSPがプログラムの実行開始時に正しく整列することが保証されます。
マルチホップシステムをスケールさせるには、このスキームをスパニングツリーの各ホップで繰り返し実行することができます。
ランタイムの再同期
TSPはプログラムの開始時に1回だけ同期しますが、各TSPが独自のクロックソースを持っているため、プログラムの実行中にも再同期する必要があります。
このために、TSPはより軽量な方式を使用します。HACに加えて、各TSPにはソフトウェアアラインドカウンタ(SAC)があり、HACと同じオーバーフロー周期を持っています。
ただし、SACはTSP間で同期されておらず、単にTSPのクロックサイクルをカウントします。HAC値は分散システムのグローバルタイムを表し、SACはローカルタイムを表します。したがって、HACとSAC値の差が蓄積されたドリフトを決定します。
ローカルタイムとグローバルタイムを再同期させるために、TSPはRUNTIME_DESKEW命令を実行します。この命令は1つのパラメータ(ストールするサイクル数)を取ります。システム内の各TSPは同時にRUNTIME_DESKEW命令を実行し、蓄積されたドリフトに基づいてグローバルタイムをローカルタイムに合わせるために実行を停止します。
ソフトウェアスケジュールドネットワーキングにおけるコンパイラの役割
見てきたように、決定論的な実行の保証は、単一のTSPからマルチTSPの分散システムにまで拡張されています。これにより、コンパイラはTSP内のデータ移動だけでなく、ネットワーク全体にわたるデータ移動についても、サイクル精度の知識を持つことができます。コンパイラは、ソースTSPでベクトルを注入するタイミングと、ターゲットTSPにそれが到着するタイミングを正確に知っています。これはソフトウェアスケジュールドネットワーキングと呼ばれ、ハードウェアが動的にデータのフローを管理するのではなく、コンパイラがコンパイル時に静的にすべてを解決するからです。
コンパイラがソフトウェアスケジュールドネットワーキングを実行できるようにする分散TSPシステムの特性について議論しましょう。
既知のトラフィックパターン
深層学習モデルの場合、コンパイラはモデルの静的な計算グラフに基づいてデータのフローを推測することができます。また、コンパイラは、ネットワーク内の利用可能なTSPデバイスに計算タスクの分散を自動化します。その結果、コンパイラは各サブタスクの正確な実行時間とレイヤー間のアクティベーションの交換を計算します。これにより、並列分解ステップが明示的になり、コンパイラの完全な制御下に置かれます。
スケジュールされたデータフロー
従来のネットワークシステムでは、ネットワークを通過するパケットのフローはハードウェアによって管理され、ハードウェアはネットワークのバックプレッシャーを感知したときにルートを最適化します。データフローのこの反応的な調整は、レイテンシを増加させ、データフローの非決定論を引き起こします。
これを避けるために、分散マルチTSPシステムでは、コンパイラを使用してネットワークを通過するデータのフローを明示的にスケジュールします。コンパイラは、ネットワーク内のどの時点でもバックプレッシャーや輻輳が発生しないように、データを賢明にルーティングします。
また、コンパイラでスケジュールされたデータフローは、デバイスがデータを必要とするちょうどそのタイミングで、コンパイラが2番目のデバイスにデータを能動的にプッシュするようにスケジュールできるため、ネットワークのレイテンシも改善します。
決定論的な負荷分散
コンパイル時にデータフローをスケジュールすることのもう1つの利点は、利用可能なリンクにわたってフローを効果的に負荷分散できることです。従来のネットワークでは、ハードウェアはルータで利用可能な輻輳メトリクスに応じて、パケットごとにルーティングの決定を行います。
しかし、マルチTSPシステムの場合、コンパイラはデータ量(すなわち、テンソルの次元H×W×Cの積)に基づいて最適なスケジュールを行い、トラフィックを分散させるリンクの数を選択します。これにより、システムで利用可能な帯域幅を効率的に使用し、全体的なレイテンシを削減します。
ここまでLPUのハードウェアを高速化するものについて議論してきました。最後にハードウェアがエラーをどのように処理するかを議論して締めくくりましょう。
エラー訂正とシステムの信頼性戦略
従来のネットワークシステムでは、伝送エラーはリンク層でパケットを単に再試行することで処理されます。この戦略は、非決定論的な遅延を引き起こすため、マルチTSPシステムには適していません。代わりに、マルチTSPシステムでは、前方誤り訂正戦略を用いてエラーを処理します。
この戦略では、パケットはホップバイホップで送信され、エラーはその場で訂正されます。重大なエラーにはフラグが立てられ、ランタイムが推論を再生できるようになります。再生戦略は一時的なエラーを処理しますが、永続的なエラーには手動の介入が必要です。
信頼性戦略として、各ラックには予備のTSPノードが含まれています。ランタイムはラック内のすべてのノードの健全性を監視し、回復不能なエラーを検出したノードの1つで予備のノードにフェイルオーバーします。
ソフトウェアによる推論の再生と予備ノードを使用することで、システムは重大なエラーから優雅に回復することができます。
まとめ
これでGroq社のLPUの基礎となるハードウェアのアーキテクチャについての議論を終えます。TSPについて、そしてGroq社がそれを大規模な並列分散システムの形でどのように拡張したかについて、かなり多くの詳細について取り上げました。重要なポイントを簡単にまとめておきましょう。
革新的な設計:Groq社のLPUは、テンサーストリーミングプロセッサ(TSP)をベースに構築されており、言語モデルの推論タスクのパフォーマンスを向上させるために、決定論的で予測可能な実行に焦点を当てた、従来のCPU/GPUアーキテクチャからの脱却を示している。
決定論的な実行:TSPアーキテクチャは、シンプルなハードウェア設計により決定論を実現し、コンパイラが命令とデータフローを正確にスケジュールできるようにしており、これは予測可能なパフォーマンスにとって重要である。
ソフトウェアスケジュールドネットワーキング:Groq社の分散マルチTSPシステムを構築するアプローチは、従来のネットワークシステムでルーティングとスケジューリングによって典型的に導入される非決定論を排除している。コンパイラはデータ移動において重要な役割を果たし、効率性と予測可能性を保証している。
大規模な並列性:ハードウェア設計とコンパイラのインテリジェンスの組み合わせにより、TSP、ひいてはLPUは高度に並列な操作をサポートし、大規模な言語モデル(LLM)に特に適している。
システムのスケーラビリティと信頼性:TSPネットワークの物理的および論理的設計により、10,000を超えるTSPのシステムをサポートし、レイテンシを最小限に抑えながら高いスケーラビリティを実現し、エラー修正戦略とフェイルオーバーメカニズムを含む高い信頼性を実現している。
潜在的な影響:LPUテクノロジーの決定論的な性質とスケーラブルな設計は、特にAIと機械学習の分野で、計算タスクのパフォーマンスへの期待を再定義し、将来のプロセッサ設計に新たな基準を設定する可能性を秘めている。
私はこの記事が洞察に富むものであることを願っています。ご質問やコメントがありましたら、ぜひお知らせください。