Groq LPU の謎を解明: 世界最速のハードウェア アクセラレータの基礎となるアーキテクチャ設計!

先月末、スタートアップ企業Groqの製品が一夜にして人気を博した。

自社開発のハードウェアアクセラレータ LPU を活用して、500 トークン/秒という神レベルの推論速度を達成し、ChatGPT を瞬時に打ち負かしました。

写真

Groq が提供する応答速度は人々の認知能力をリフレッシュしますが、これはその背後にある言語処理ユニット ハードウェア LPU (language processing unit hardware) のおかげです。

Groq の R&D チームは、革新的なハードウェア アーキテクチャ設計を LPU に適用し、強力なコンパイラを搭載しました。

Substack のコラムニスト Abhinav Upadhyay 氏と一緒に、Groq LPU の基盤となるアーキテクチャの謎を段階的に解明してみましょう。

Groq LPU の謎

今のところ、Groq は LPU 自体に関する論文を発表していませんが、ここ数年で次の 2 つの論文を発表しました。

論文アドレス: https://dl.acm.org/doi/10.1109/ISCA45697.2020.00023

写真

論文アドレス: https://dl.acm.org/doi/abs/10.1145/3470496.3527405

2つの論文は、それぞれ2020年と2022年にトップクラスのコンピューターアーキテクチャ会議ISCAで発表され、後者は受賞論文となった。

これら 2 つの記事では、Groq の Tensor Streaming Processor (TSP) の設計と実装、および TSP を使用して分散推論エンジンを構築する方法について説明します。

公式には発表されていないものの、LPU はこの分散システムに基づいて設計および拡張されていると考えられます。

そこで、まず TSP とそのコンパイラのアーキテクチャを詳細に分析し、それを基に Groq がこれらの TSP を使用して、信頼性が高くスループットの高い分散 AI 推論エンジンを構築する方法を分析してみましょう。

TSP のアーキテクチャは従来の CPU や GPU チップとは大きく異なり、その主な目的は TSP ハードウェアをより決定論的にすることです。

ここでまず、CPU または GPU の不確実性について言及する必要があります。

CPU と GPU のマイクロアーキテクチャにおける非決定性

マイクロアーキテクチャの設計上、CPU および GPU での命令の実行は非決定論的です。つまり、特定の命令がいつ実行されるか、完了するまでにどのくらいの時間がかかるか、結果がいつ利用可能になるかは保証されません。

たとえば、最近の CPU は一般的に次のような設計になっています。

- スーパースカラーアーキテクチャ: 1 サイクルあたり複数の命令を発行できます。

- アウトオブオーダー実行: 命令を任意の順序で実行します。

- 投機的実行: 分岐の場合、分岐条件が真か偽かを推測し、事前に分岐の実行を予測してスループットを向上させます (もちろん、推測が間違っていた場合は、諦めて別の分岐に戻る必要があります)。

- 命令パイプライン: 命令は複数のステージに分割され、パイプライン方式で実行されるため、命令スループットがさらに向上します。

- 複数レベルのキャッシュ: CPU には 2 ～ 3 レベルのキャッシュがあり、メモリからデータをロードする際の待ち時間を短縮できます。

これらすべてにより、CPU での命令実行の順序とタイミングは非決定的になり、推論が困難になります。

GPU には、キャッシュ、共有メモリとグローバル メモリ、動的リソース パーティショニングなど、その他の非決定論的要素もあります。

非決定性の問題は、プログラムのパフォーマンスを推論することが難しく、最悪の場合のパフォーマンスの限界を保証することも難しいことです。

そのため、Groq は、高度に並列化され、非決定的な動作のない、まったく新しい TSP の設計を提案しました。これにより、ハードウェアの複雑さが解消され、コンパイラーは命令の実行を正確にスケジュールおよび制御できるようになり、プログラムのパフォーマンスの制限が保証されます。

次に、TSP アーキテクチャが内部からどのように見えるかを見てみましょう。

TSPアーキテクチャ

TSP のハードウェア設計は、CPU や GPU のハードウェア設計とはまったく対照的です。従来のマルチコア チップはタイル アーキテクチャを使用します。下の図 (a) では、各小さな四角形 (タイル) が処理コアを表しています。

コアは、さまざまな種類の計算 (算術演算、メモリ演算、論理演算、命令制御など) を実行する機能ユニットのセットから構成されます。

TSP の設計者は、この従来の設計を完全に覆しました。彼らは機能ユニットをコアの外側に移動し、2D グリッドに配置しました。

グリッドの各列には、スライスと呼ばれる特定のタイプの機能単位のみが含まれます。下の図は、従来のマルチコア チップと TSP の設計の違いを示しています。

TSP には次の関数スライスがあります。

MXM: 行列演算を実行するために使用

SXM: ベクトルをシフトおよび回転するために使用される

MEM: メモリの読み取り/書き込み操作

VXM: ベクトルの算術演算

ICU: 命令制御ユニット。これは少し特殊で、上図 (b) の下部にある水平の青いバーです。命令のフェッチとスケジュール設定、および他のスライスでの実行を担当します。

TSP のアーキテクチャについて理解できたので、次はその中核である命令実行に注目してみましょう。

TSPにおける命令実行

TSP は SIMD (Single Instruction Multiple Data) モードで命令を実行します。各機能スライスは 20 個のタイルで構成され、各タイルは 16 個の数字を処理できます。したがって、完全なスライスでは、最大 320 要素のベクトルを処理および生成できます。

これらのスライスは、プロデューサーとコンシューマーの方法で相互作用します。

ベクトルがメモリから読み取られると、フロー ID (0 ～ 31) とフロー方向 (東または西) が割り当てられます。各スライスは、ストリームを自由に処理して新しいストリームを生成することも、ストリームをそのまま次の隣接するスライスに流すこともできます。

完全なベクトルを効率的に処理するために、次の図に示すように、命令はパイプライン方式で実行されます。

写真

TSP での命令のパイプライン実行により、ストリームがスライス間でインターリーブされます。上の図の黒いブロックは、スライス内のさまざまなタイムスタンプでのストリームの動きを表しています。

もちろん、命令を問題なく実行するには、コンパイラと命令セット (ISA) の設計が必要です。

TSP コンパイラと ISA

TSP の設計者はハードウェアを簡素化したため、コンパイラに負担がかかりました。コンパイラは、特定のプログラムを正しく、最も効率的に実行するために、命令とデータのフローを正確にスケジュールする必要があります。

コンパイラは、TSP ハードウェアの次の状態にアクセスできます。

- 320 レーンのプログラミング抽象化: TSP チップ内の各タイルは、SIMD 形式で 16 ユニット (16 レーン) のベクトルを操作できます。垂直スライスは 20 個のタイルで構成されているため、実行に使用できる SIMD レーンは合計 320 個あります。

- 144 個の独立した命令キュー: チップ上には 144 個の命令キューがあり、サイクルごとに 1 つ以上の命令を発行できます。コンパイラは各キュー内のプログラムの順序を完全に制御します。

- チャネルあたり 64 個の論理ストリーム: 各チャネルは、オペランドまたは結果を移動するために使用できる 64 個の論理ストリームにアクセスでき、そのうち 32 個はデータを東方向に移動するために使用でき、残りの 32 個はデータを西方向に移動するために使用できます。

- 220M のグローバル共有 SRAM。

TSP ハードウェアには非決定論的な動作がないため、コンパイラは各命令のレイテンシだけでなく、プログラム内のデータフロー (DNN の計算グラフなど) も正確に把握しています。

コンパイラは計算タスク間の依存関係を識別し、並列実行のために TSP の利用可能な機能ユニットにそれらを配布します。

TSP プログラミング モデルは、次の 2 つの重要な要素に依存します。

ハードウェアにおける決定論的データパス

ISAによって取得された命令レイテンシに関する情報

コンパイラのバックエンドは、チップ上の任意のフローの場所と使用時間を追跡できます。これはソフトウェア定義ハードウェアと呼ばれます。

TSPからLPUへの拡大

TSP は LPU の基本単位です。多数の TSP がラック内にグループ化され、膨大なスループットを提供できる分散システムを形成します。

マルチTSPシステムの設計

TSP と同様に、分散型マルチ TSP システムの設計目標も、決定論的なデータ フローと命令実行、およびノー​​ド間の低遅延通信を中心に展開されます。

分散 TSP システムの設計はノードから始まります。ノードはシャーシ内の 8 つの TSP デバイスで構成されます。これらの各デバイスは 11 個のピンで構成されており、そのうち 7 個は各 TSP デバイスをノード内の他の 7 個の TSP デバイスに接続するために使用され、残りの 4 個のピンはグローバル リンクを形成するために使用されます。

ノード内の各デバイスには 4 つのグローバル リンクがあり、合計 32 個のグローバル リンクが組み合わさって 32 個の仮想ポートを持つ高基数ルーターを形成します。

高基数ルーターは、多数の接続、高帯域幅、および高性能をサポートします。これはまさに高性能分散システムに必要なものです。

このような 9 つの TSP ノードと 8 つの TSP を 1 つのラックに結合します。ラック内の各ノードには 32 個のポートがあるため、ラックには合計 288 個のグローバル ポートがあります。

144 個のポートはラック内でローカルに使用され、ラック内でデータを高速転送します。残りの 144 個のポートは他のラックに接続するために使用されます。

システムの最大構成では、10,440 個の TSP を含む 145 個の相互接続されたラックをサポートでき、システム内の任意の 2 つの TSP 間のホップは最大 5 個になります。

TSP ベースの分散システムにおける決定論の実装

この拡張分散システム体制では、単一の TSP の機能ユニットが、超並列プロセッサの単一の処理コアとして機能します。 TSP の計算モデルは決定論的なハードウェアに基づいているため、分散システム全体にも同じ決定論性が備わっている必要があります。

ハードウェアアライメントカウンタを使用してTSPクロックを同期する

各 TSP デバイスには、オーバーフロー期間が 256 のハードウェア アライメント カウンター (HAC) と呼ばれるハードウェア カウンターが含まれています。 TSP は、次の手順でこれを使用して相互に同期します。

- 2 つの TSP が相互接続されている場合、一方の TSP が HAC 値をもう一方の TSP に送信します。その後、相手側は送信者に値を返します。送信者は、現在の HAC 値と返された値の差を観察します。

- この差は、2 つのデバイス間のリンクの遅延を表します。このプロセスを複数回繰り返して、2 つの TSP 間の平均リンク遅延を取得します。

- その後、2 つのデバイスは親子関係に配置されます。親は定期的に現在の HAC 値を子に送信します。子は平均リンク遅延を自身の HAC 値に追加し、それを自身の HAC 値と比較します。

- 2 つの値の差は、継続的なクロック ドリフトによる初期の不整合を表します。次に、子供は HAC 値を調整してこの差を減らします。このプロセスを何度も繰り返すと、2 つの TSP の HAC 値が小さな範囲内に収束し、リンク遅延のジッターが示されます。

- このプロトコルにより、2 つの TSP が相互に同期できるようになり、ネットワーク内にスパニング ツリーを確立することで TSP マルチホップ ネットワークを拡張できます。

当初計画の調整

マルチ TSP システムでプログラムを実行する前に、システム全体のデータ フローと命令実行を正しくスケジュールするために、すべての TSP を調整する必要があります。これには次のメカニズムが含まれます。

- 単一の TSP レベルには、いくつかの独立した機能ユニットと 144 個の独立した命令キューがあります。これらを同期するために、TSP は SYNC および NOTIFY 命令をサポートしています。 SYNC 命令は、すべての命令キューをパーク状態にし、1 つのキューを通知者として機能させます。通知者が NOTIFY 命令を発行すると、その命令はチップ上のすべてのキューにブロードキャストされ、その時点でキューは同期され、操作が再開されます。

- マルチ TSP システムの場合、2 つの TSP は HAC を使用して相互に同期し、各 TSP は DESKEW 命令をサポートして、TSP の HAC がオーバーフローするまで後続の命令の処理を停止します。

- マルチホップ システムを拡張するには、上記のソリューションをスパニング ツリーの各ホップで繰り返すことができます。

ランタイム再同期

TSP はプログラムの開始時に一度同期されますが、各 TSP には独自の独立したクロック ソースがあるため、プログラムの実行中に再同期する必要もあります。

このため、TSP はより軽量なスキームを使用します。 HAC に加えて、各 TSP にはオーバーフロー期間が HAC と同じソフトウェア アライメント カウンター (SAC) があります。

ただし、SAC は TSP 間で同期されず、TSP のクロック サイクルをカウントするだけです。 HAC 値は分散システムのグローバル時間を表し、SAC はローカル時間を表します。したがって、HAC 値と SAC 値の差によって累積ドリフトが決まります。

ローカル時間とグローバル時間を再同期するために、TSP は RUNTIME_DESKEW 命令を実行します。システム内の各 TSP は命令を同時に実行し、累積されたドリフトに応じてグローバル時間とローカル時間を調整します。

ソフトウェア計画ネットワークにおけるコンパイラの役割

これまで、コンパイラーは TSP 内およびネットワーク全体でのデータの移動をサイクル単位で正確に把握していました。コンパイラは、ソース TSP にベクトルを挿入する正確な時間と、それがターゲット TSP に到達する正確な時間 (ソフトウェア プラン ネットワークと呼ばれる) を認識します。

コンパイラは、データフローを動的に管理するのではなく、コンパイル時にすべてを静的に解決します。

既知のトラフィックパターン

ディープラーニング モデルの場合、コンパイラはモデルの静的計算グラフに基づいてデータ フローを推測できます。コンパイラは、ネットワーク内で利用可能な TSP デバイス間で計算タスクを自動的に分散することもできます。

したがって、コンパイラは各サブタスクの正確な実行時間とレイヤー間のアクティベーション交換を計算します。これにより、並列分解ステップが明示的になり、コンパイラによって完全に制御されるようになります。

計画されたデータフロー

従来のネットワーク システムでは、ネットワークを介したパケットのフローはハードウェアによって管理され、ネットワークの負荷を感知するとルーティングが最適化されます。データ フローのこの受動的な調整により、レイテンシが増加し、データ フローに非決定性が生じます。

これを回避するために、分散型マルチ TSP システムではコンパイラを使用して、ネットワークを介したデータ フローを明示的にスケジュールします。コンパイラは、ネットワーク内で輻輳が発生しないようにデータを巧みにルーティングします。

さらに、コンパイラが計画したデータフローは、デバイス経由でデータを要求するのではなく、コンパイラがデータを事前にプッシュするようにスケジュールできるため、ネットワークのレイテンシが改善されます。

決定論的負荷分散

コンパイル時にデータフローをスケジュールするもう 1 つの利点は、コンパイラが利用可能なリンク間でフローの負荷を効果的に分散できることです。従来のネットワークでは、ハードウェアはルーターで利用可能な輻輳メトリックに基づいてパケットごとにルーティングの決定を実行します。

ただし、マルチ TSP システムの場合、コンパイラはデータ量に基づいて最適な方法でスケジューリングを実行し、トラフィックを分散するためのリンクの数を選択します。これにより、システム内の利用可能な帯域幅を効率的に使用し、全体的な遅延を削減できます。