Dataflowは、SambaNova Systemsによって開発された、Artificial intelligenceワークロードを高速化するためのコンピューティングアーキテクチャである。命令を順次実行する従来のプロセッサ設計とは異なり、Dataflowは計算を有向グラフとして編成し、データが処理要素のネットワークを通じて継続的に流れるようにする。このアプローチは、Deep learningモデル、例えばニューラルネットワークや大規模言語モデルに見られる構造化された反復的な操作に合わせて調整されている。このアーキテクチャは、SambaNovaのカスタムハードウェア、DataScaleシステムやSN10シリーズの再構成可能データフローユニット(RDU)に実装されている。
Dataflowアーキテクチャは、スタンフォード大学での研究から生まれ、2017年にKunle Olukotun、Christopher Ré、Rodrigo Liangによって設立されたSambaNovaによって商業化された。同社の最初の製品であるDataScale SN10は2020年に発表され、2021年に出荷が開始された。このアーキテクチャは、従来のGPUベースのシステムを制限するメモリ帯域幅のボトルネックに対処するように設計されており、従来のシステムではメモリと計算ユニット間のデータ移動に多くのエネルギーを費やすことが多い。データをチップ上に保持し、構成可能なファブリックを通じてストリーミングすることで、DataflowはAI推論とトレーニングにおいてより高い利用率と低いレイテンシを達成することを目指している。
設計原則
Dataflowの核となる原則は、計算を命令フェッチとデコードから切り離すことでフォンノイマンボトルネックを排除することである。従来のCPUやGPUでは、各操作は命令をフェッチし、デコードし、メモリからフェッチしたデータに対して実行する必要がある。Dataflowは代わりに、モデル全体を静的なデータフローグラフにコンパイルし、各ノードが数学的操作(行列乗算や活性化など)を表し、各エッジがデータ依存関係を表す。グラフはその後、それぞれがローカルメモリと算術ユニットを持つ処理要素(PE)のグリッドにマッピングされる。
Dataflowアーキテクチャの各PEは、乗算累積、活性化関数、プーリングなど、さまざまな操作を実行できる。PEは、近隣通信と長距離通信の両方をサポートする高帯域幅、低レイテンシのネットワークを介して相互接続されている。このネットワークの構成はコンパイル時に決定され、ハードウェアを各特定モデルに特化させることができる。この再構成可能性は、AWS TrainiumやGroqのテンソルストリーミングプロセッサのような固定機能アクセラレータとの主要な差別化要因であり、これらもデータフロー原理を使用するが、異なるトレードオフを持つ。
コンパイルとソフトウェアスタック
SambaNovaはSambaFlowと呼ばれるソフトウェアスタックを提供しており、PyTorchやTensorFlowなどのフレームワークで書かれたモデルをDataflowグラフにコンパイルする。コンパイラは、演算子融合、メモリレイアウト最適化、パイプライン化など、いくつかの最適化を実行する。また、グラフを物理的なPEグリッドにマッピングし、計算リソースとメモリリソースのバランスを取る。SambaFlowはトレーニングと推論の両方をサポートし、データ移動と同期を管理するランタイムを含む。
コンパイルプロセスは静的であり、グラフは実行開始前に完全に決定される。これにより、コンパイラは操作を正確にスケジュールし、ランタイムオーバーヘッドを回避できる。ただし、これは、可変トリップカウントのループなどの動的制御フローが、アンロールされるか、特別なメカニズムを介して処理される必要があることも意味する。SambaNovaは、注意機構とフィードフォワード層を最適化することで、ほとんどの現代のGenerative AIシステムの基盤であるトランスフォーマーモデルをサポートするようにコンパイラを拡張している。
ハードウェア実装
2020年に導入されたSN10 RDUは、400億個のトランジスタを含み、TSMCで7ナノメートルプロセスを使用して製造されている。PEグリッド全体に分散された48 GBのオンチップSRAMを備えている。チップは約1.2 GHzのクロック速度で動作するが、その性能は生のクロックレートではなくデータフロースループットで測定される。SambaNovaは、単一のSN10がスパースモデルで最大638テラフロップス、デンスモデルで125テラフロップスの計算を提供できると主張しているが、これらの数値はワークロードに依存する。
DataScaleシステムは、複数のSN10チップを単一のシャーシに統合し、システムあたり最大16個のRDUを備える。これらのシステムは、ピアツーピア通信をサポートする高速インターコネクトを介して接続されている。SN30やSN40などの後の世代は、メモリと計算容量を増加させており、2023年に導入されたSN40は1.5 TBのオンチップメモリと最大5兆パラメータのモデルをサポートする。ハードウェアには、プルーニングとスパース性のための専用ユニットも含まれており、ゼロ値の計算をスキップして効率を向上させることができる。
性能特性
Dataflowシステムは、同じモデルが異なる入力で繰り返し実行される推論ワークロードに特に効果的である。静的なスケジューリングにより、予測可能なレイテンシが可能になり、オンチップメモリにより外部DRAMへのアクセス必要性が減少する。SambaNovaが公開したベンチマークでは、DataScale SN10は、BERTやGPT-3を含むいくつかの人気モデルでNVIDIA A100 GPUよりも高いスループットを達成し、消費電力も少ない。ただし、独立したベンチマークは限られており、実際の性能はモデルとバッチサイズによって異なる可能性がある。
トレーニングでは、Dataflowは頻繁な重み更新と勾配計算の動的な性質のために課題に直面する。SambaNovaは、データ並列性とモデル並列性をサポートすることでこれに対処しており、複数のRDUがモデルの異なる部分で作業する。コンパイラは、前方および後方パスをパイプライン化して、計算と通信をオーバーラップさせることができる。これらの努力にもかかわらず、トレーニング性能は推論ほど競争力がなく、多くの顧客は主にデプロイメントにDataflowを使用している。
他のアーキテクチャとの比較
Dataflowは、オンチップメモリを備えたマニーコア設計も使用するGraphcoreのIPUとよく比較される。主な違いは、Graphcoreがより伝統的な命令ベースの実行モデルを使用するのに対し、Dataflowは純粋にデータ駆動型であることである。Groqのテンソルストリーミングプロセッサもデータフローに似たアーキテクチャであるが、オンチップメモリのないよりシンプルで剛性の高い設計を使用し、代わりに多数の乗算累積ユニットの大規模なアレイを通じてデータをストリーミングする。SambaNovaの再構成可能性は、さまざまなモデル形状に対してより柔軟性を提供するが、コンパイラに複雑さを追加する。
NVIDIA(提供されたリストにはないが、暗黙的に含まれる)やAMDのGPUと比較すると、Dataflowは小さなバッチに対して低いレイテンシと、推論におけるより良いエネルギー効率を提供する。ただし、GPUは成熟したソフトウェアエコシステムと、フレームワークやライブラリのより広範なサポートの恩恵を受けている。Dataflowの利点は、注意機構がPEグリッドに効率的にマッピングできる大きなトランスフォーマーモデルで最も顕著である。
アプリケーションと採用
SambaNovaは、政府、医療、金融サービスなど、さまざまなセクターにDataflowシステムを展開している。注目すべき顧客には、科学計算にシステムを使用する米国エネルギー省や、SambaNovaハードウェアをマネージドサービスとして提供するOracle Cloudが含まれる。このアーキテクチャは、自然言語処理、コンピュータビジョン、創薬などのタスクに使用されてきた。2023年、SambaNovaは、Dataflowシステムでトレーニングされた1兆パラメータの大規模言語モデルSamba-1を導入し、大規模トレーニングにおけるアーキテクチャの能力を示した。
同社はまた、クラウドサービスを通じてDataflowを利用しやすくすることに焦点を当てており、顧客がハードウェアを直接購入せずに時間単位でレンタルできるようにしている。このモデルは、高性能なAI推論を必要とするが、専用インフラを維持するリソースを欠く企業によって採用されている。2024年現在、SambaNovaは10億ドル以上の資金を調達しており、ハードウェアとソフトウェアの反復を続けている。
課題と制限
Dataflowの主な課題の1つは、コンパイラの複雑さである。任意のモデルをPEグリッドにマッピングするには、分割、スケジューリング、メモリ割り当てのための高度なアルゴリズムが必要である。この複雑さは、特に大きなモデルでは、コンパイル時間の長期化につながる可能性がある。さらに、グラフの静的な性質により、強化学習などの一部のアプリケーションで一般的な動的な形状や制御フローを持つモデルをサポートすることが困難になる。
もう1つの制限は、GPUと比較したエコシステムサポートの欠如である。多くのAI研究者や開発者はCUDAとPyTorchに精通しており、異なるプログラミングモデルを必要とする新しいアーキテクチャを採用することに消極的かもしれない。SambaNovaは互換性レイヤーを提供することでこれを緩和しているが、まだギャップがある。同社はまた、Google CloudやAWSなどのより確立されたプレーヤーとの競争に直面しており、これらは独自のカスタムアクセラレータを提供している。
将来の方向性
SambaNovaは、トレーニング効率の向上とより大きなモデルのサポートに焦点を当てて、Dataflowアーキテクチャを進化させ続けている。SN40世代には、スパース計算と混合精度演算のための強化が含まれている。同社はまた、エッジデバイスへのDataflowの使用を探求しているが、現在のハードウェアは主にデータセンター向けに設計されている。Generative AIモデルがサイズを増すにつれて、効率的な推論ハードウェアへの需要は増加する可能性が高く、Dataflowのアプローチはより魅力的になるかもしれない。
データフローアーキテクチャの研究は学界でも進行中であり、MIT CSAILやBerkeley AI Researchのグループが同様のアイデアを探求している。データフローコンピューティングの原則は1970年代に遡る長い歴史を持つが、チップ製造とコンパイラ技術の進歩により、商業的なAIワークロードに実用的になったのは最近のことである。SambaNovaのDataflowは、この概念の最も顕著な商業的実装の1つを表しており、その成功は急速に変化するAI環境に適応する能力に依存するだろう。
結論
Dataflowは、データ移動の効率と再構成可能性を優先するAIアクセラレーションへの独自のアプローチである。モデルをオンチップメモリを備えた処理要素のグリッドにマッピングすることで、推論タスクにおいて従来のGPUに代わる魅力的な選択肢を提供する。コンパイラの複雑さとエコシステム採用の課題に直面しているが、大規模言語モデルに対する性能上の利点は、重要な投資と顧客の関心を集めている。AIモデルが成長し続けるにつれて、Dataflowの基礎となる原則は、より広範なコンピューティング業界にとってますます関連性が高まる可能性がある。