より大規模な AI および機械学習モデルへの渇望を解消するために、Tesla は Hot Chips 34 で、Dojo と呼ばれる完全にカスタム化されたスーパーコンピューティング アーキテクチャに関する豊富な詳細を明らかにしました。
このシステムは本質的に大規模な構成可能なスーパーコンピューターですが、私たちが目にするものとは異なります。 トップ500、コンピューティング、ネットワーク、入出力 (I/O) シリコンから、命令セット アーキテクチャ (ISA)、電力供給、パッケージング、冷却に至るまで、完全にカスタム アーキテクチャから構築されています。 これらはすべて、調整された特定の機械学習トレーニング アルゴリズムを大規模に実行するという明確な目的で行われました。
テスラのハードウェア エンジニアリング担当シニア ディレクターであるガネーシュ ベンカタラマナン (Ganesh Venkataramanan) は、次のように述べています。彼の基調講演中。
しかし、分散ワークロードをスケーリングするための従来の方法では、機械学習の要求に追いつくために必要な速度で加速できなかったと彼は主張しました。 事実上、ムーアの法則はそれをカットしておらず、大規模な AI/ML トレーニング、つまり CPU/GPU の組み合わせ、または特殊な AI アクセラレーターを使用するまれな状況で利用できるシステムもありません。
「伝統的に私たちはチップを構築し、それらをパッケージに入れ、パッケージはPCBに取り付けられ、それがシステムに入ります。 システムはラックに入れられます」と Venkataramanan 氏は述べています。 問題は、データがチップからパッケージへ、およびパッケージから移動するたびに、遅延と帯域幅のペナルティが発生することです。
データセンターのサンドイッチ
この制限を回避するために、Venkataramanan と彼のチームはゼロからやり直しました。
「Elon とのインタビューの直後に、彼は AI 用の CPU や GPU とは異なる何ができるのかと尋ねました。 チーム全体がまだその質問に答えていると感じています。」
これが Dojo トレーニング タイルの開発につながりました。これは、556 kW の水冷パッケージで FP32 のパフォーマンスを 15 TFLOPS にできる半立方フィートを占有する自己完結型のコンピューティング クラスターです。
各タイルには 11 GB の SRAM が搭載されており、スタック全体でカスタム トランスポート プロトコルを使用して 9 TB/秒のファブリックで接続されています。
「このトレーニング タイルは、追加のスイッチを必要とせずに、コンピュータからメモリ、電力供給、通信まで、比類のない量の統合を表しています」と Venkataramanan 氏は述べています。
トレーニング タイルの中心にあるのは、TSMC の 1nm プロセスに基づく 50 億個のトランジスタ ダイである Tesla の D7 です。 テスラによれば、各 D1 は 22W の TDP で 32 TFLOPS の FP400 パフォーマンスが可能です。 ただし、Tesla は、チップがいくつかのカスタム計算を含む幅広い浮動小数点計算を実行できると述べています。
テスラの道場 D1 ダイ
「トランジスタをミリメートル四方で比較すると、これはおそらく世の中にあるあらゆるものの最先端です」と Venkataramanan 氏は述べています。
その後、テスラは 25 個の D1 を取得し、既知の良好なダイに分類し、TSMC のシステム オン ウエハー テクノロジを使用してパッケージ化して、「非常に低いレイテンシと非常に高い帯域幅で大量のコンピューティング統合を実現しました」と述べました。
しかし、システム オン ウエハー設計と垂直スタック アーキテクチャは、電力供給に関して課題をもたらしました。
Venkataramanan 氏によると、現在のほとんどのアクセラレータは、シリコンに直接隣接して電力を配置しています。 このアプローチは証明されていますが、アクセラレータの大きな領域をそれらのコンポーネント専用にする必要があることを意味し、Dojo にとっては非現実的であると彼は説明しました。 代わりに、テスラはダイの底部から直接電力を供給するようにチップを設計しました。
それをすべてまとめる
「このトレーニング タイルからデータセンター全体または建物全体を構築できますが、トレーニング タイルはコンピューティング部分にすぎません。 また、餌を与える必要があります」とベンカタラマナンは言いました。
Tesla の Dojo インターフェイス プロセッサ
このために、Tesla は Dojo Interface Processor (DIP) も開発しました。これは、ホスト CPU とトレーニング プロセッサの間のブリッジとして機能します。 DIP は、共有高帯域幅メモリ (HBM) のソースとして、また高速 400Gbit/秒 NIC としても機能します。
各 DIP は 32 GB の HBM を備えており、これらのカードを最大 900 枚まで 4.5 GB/s でトレーニング タイルに接続して、ホストへの合計 160 TB/s でタイルあたり合計 XNUMX GB の HBM を実現できます。
これらのタイルのテスラの V1 構成ペア (または 150 個の D1 ダイ) は、それぞれ 16 つの DIP カードを装備した 8 つのホスト CPU をサポートし、BFXNUMX または CFPXNUMX パフォーマンスの主張されたエクサフロップを達成しました。
テスラのV1アレンジメント
まとめて、Venkataramanan はアーキテクチャを説明します – 詳細 ここで深く by 次のプラットフォーム – Tesla は、Nvidia や AMD などの従来のアクセラレータに関連する制限を克服できます。
「従来のアクセラレータがどのように機能するかというと、通常、モデル全体を各アクセラレータに適合させようとします。 それを複製し、それぞれにデータを流します」と彼は言いました。 「モデルがどんどん大きくなったらどうなる？ これらのアクセラレータは、メモリ不足のために機能しなくなる可能性があります。」
これは新しい問題ではない、と彼は指摘した。 たとえば、Nvidia の NV スイッチを使用すると、大規模な GPU バンク全体でメモリをプールできます。 ただし、Venkataramanan 氏は、これにより複雑さが増すだけでなく、遅延が発生し、帯域幅が犠牲になると主張しています。
「私たちは最初からこれについて考えていました。 私たちの計算タイルと各ダイは、大きなモデルに適合するように作られています」と Venkataramanan 氏は述べています。
ソフトウェア
このような専用のコンピューティング アーキテクチャには、専用のソフトウェア スタックが必要です。 しかし、Venkataramanan 氏と彼のチームは、プログラマビリティが Dojo の成否を左右することを認識していました。
「これらのシステムを設計するとき、ソフトウェアのカウンターパートのプログラミングの容易さは最も重要です」と彼は言いました。 「研究者は、私たちが実行したい新しいアルゴリズムに適応するための手書きのカーネルをソフトウェアの人々が書くのを待つことはありません。」
これを行うために、Tesla はカーネルを使用するという考えを捨て、Dojo のアーキテクチャーをコンパイラーを中心に設計しました。
「私たちがしたことは、PiTorch を使用したことです。 中間レイヤーを作成しました。これは、その下にあるハードウェアを並列化してスケールアウトするのに役立ちます。 すべての下にコンパイルされたコードがあります」と彼は言いました。 「これは、将来のすべてのワークロードに適応できるソフトウェア スタックを作成する唯一の方法です。」
ソフトウェアの柔軟性に重点が置かれているにもかかわらず、Venkataramanan 氏は、現在ラボで実行されているプラットフォームは当面の間、Tesla の使用に限定されていると述べています。
「私たちはまず内部顧客に焦点を当てています」と彼は言いました。 「イーロンは、時間をかけて研究者が利用できるようにすることを公表しましたが、その時間枠はありません。 ®
