新技術により大規模人工知能モデルの処理性能が効果的に向上

MIT と Nvidia の研究者は、高性能コンピューティング タスクで使用されるデータ構造であるスパース テンソルの処理を高速化する 2 つの手法を開発した。これらの技術により、生成 AI を推進する大規模な機械学習モデルなどのシステムのパフォーマンスとエネルギー効率を大幅に向上させることができます。

テンソルは機械学習モデルで使用されるデータ構造です。どちらの新しい手法も、テンソル内のいわゆるスパース性（ゼロ値）を効果的に利用しようとします。これらのテンソルを処理するときは、ゼロをスキップして計算とメモリを節約できます。たとえば、ゼロを掛けたものはすべてゼロなので、その演算をスキップできます。テンソルを圧縮できるため（ゼロを保存する必要がない）、テンソルの大部分をオンチップメモリ​​に保存できます。

ただし、スパース性を活用するにはいくつかの課題があります。大きなテンソル内でゼロ以外の値を見つけるのは簡単な作業ではありません。既存の方法では、検索を簡素化するためにスパースパターンを適用して非ゼロ値の位置を制限することがよくありますが、これにより効率的に処理できるスパーステンソルの多様性が制限されます。

もう 1 つの課題は、テンソルの領域によって非ゼロ値の数が変わる可能性があることです。このため、メモリ内のさまざまな領域を格納するために必要なスペースの量を判断することが難しくなります。領域が確実に収まるように、必要以上のスペースが割り当てられることが多く、その結果、ストレージ バッファーが十分に活用されなくなります。これにより、オフチップ メモリ トラフィックが増加し、エネルギー消費も増加します。

MIT と Nvidia の研究者は、これらの問題に対処するための 2 つの解決策を考案しました。まず、ハードウェアがより多様なスパースパターンに対して非ゼロ値を効率的に見つけられるようにする技術を開発しました。

別の解決策として、データがメモリに収まらないケースを処理できる方法を作成し、それによってストレージ バッファの使用率を向上させ、オフチップ メモリ トラフィックを削減しました。

どちらのアプローチも、スパース テンソル処理を高速化するために特別に設計されたハードウェア アクセラレータのパフォーマンスを向上させ、エネルギー要件を削減します。

「通常、より特化またはドメイン固有のハードウェア アクセラレータを使用すると、CPU などのより汎用的なプロセッサから得られる柔軟性が失われます。この 2 つの研究の注目すべき点は、特化と効率性を維持しながら、柔軟性と適応性を維持できることを示したことです」と、MIT の電気工学およびコンピューター サイエンス学部 (EECS) の准教授であり、電子工学研究所 (RLE) のメンバーであり、この 2 つの進歩に関する論文の共同執筆者である Vivienne Sze 氏は述べています。

共著者には、主著者で電気工学とコンピュータサイエンスの大学院生である Yannan Nellie Wu PhD ’23 と Zi Yu Xue、共同上級著者でコンピュータサイエンスと電気工学の実践教授であり、MIT のコンピュータサイエンスおよび人工知能研究所 (CSAIL) のメンバーである Joel Emer と NVIDIA の他のメンバーが含まれます。両論文は、IEEE/ACM マイクロアーキテクチャ国際シンポジウムで発表される予定です。

ハイライト: ゼロ値の効率的な検索

テンソルでは、さまざまな理由でスパース性が生じる可能性があります。たとえば、研究者は機械学習モデルの不要な部分を「削除」するために、テンソル内の特定の値をゼロに置き換えることでスパース性を作成することがあります。スパース性の度合い（ゼロの割合）とゼロの位置は、モデルによって異なる場合があります。

数十億の個別の値を持つモデルで残りの非ゼロ値を見つけやすくするために、研究者は非ゼロ値の位置を特定のパターンに収まるように制限することがよくあります。ただし、各ハードウェア アクセラレータは通常、特定のスパース パターンをサポートするように設計されているため、柔軟性が制限されます。

対照的に、MIT の研究者によって設計されたハードウェア アクセラレータである HighLight は、さまざまなスパース パターンを処理でき、ゼロ値のないモデルを実行する場合でも優れたパフォーマンスを発揮します。

彼らは、「階層構造化スパース性」と呼ばれる手法を使用して、いくつかの単純なスパース性パターンで構成されたさまざまなスパース性パターンを効率的に表現します。このアプローチでは、テンソル内の値を小さなブロックに分割します。各ブロックには独自の単純なスパース パターンがあります (4 つの値のブロックに 2 つのゼロと 2 つの非ゼロが存在する場合があります)。

次に、これらのブロックを階層にグループ化し、各ブロック セットにも独自の単純なスパース パターンを持たせました (4 つのブロックがあるレベルでは、ゼロ ブロックが 1 つと非ゼロ ブロックが 3 つある場合があります)。ブロックを組み合わせてさらに大きなレベルにしていきますが、各ステップのパターンは単純なままです。

このシンプルさにより、HighLight はゼロをより効率的に見つけてスキップすることができるため、この機会を最大限に活用して冗長な計算を削減できます。平均すると、彼らの加速器設計は、他のアプローチよりも約 6 倍高いエネルギー遅延積 (エネルギー効率に関連する指標) を達成しました。

「最後に、HighLight アクセラレータは、階層構造のスパース性に基づいてゼロ値の量が異なるワークロードを活用できる一方で、オーバーヘッドをあまり発生させないため、高密度モデルを効果的に高速化できます」と Wu 氏は説明します。

将来的には、彼女と共同研究者は、階層構造化スパース性をより多くの種類の機械学習モデルとモデル内のさまざまな種類のテンソルに適用したいと考えています。

Tailors と Swiftiles: 効果的に「オーバーサブスクライブ」してワークロードを加速

研究者はスパース性を活用して、コンピューター チップ上でデータをより効率的に移動および処理することもできます。

テンソルは通常、オンチップのメモリ バッファに格納できるサイズよりも大きいため、チップは一度にタイルと呼ばれる大きなテンソルのチャンクのみを取得して処理できます。

このバッファを最大限に活用し、チップがオフチップ メモリにアクセスする回数を制限するために (オフチップ メモリは通常、エネルギー消費の大部分を占め、処理速度を制限する)、研究者はバッファに収まる最大のスライスを使用することを試みました。

しかし、スパース テンソルでは多くのデータ値がゼロであるため、バッファーはゼロ値を格納せずに、その容量に基づいて予想されるよりも大きなスライスを保持できます。

ただし、ゼロ値の数はテンソルの領域によって異なる場合があり、タイルごとに異なる場合もあります。このため、バッファに収まる適切なタイルのサイズを決定することが困難になります。したがって、既存の方法では通常、ゼロが存在しないと保守的に想定し、より小さなタイルを選択することになり、バッファ内に無駄な空きスペースが生じます。

この不確実性に対処するために、研究者らは「オーバーサブスクリプション」を使用してタイルのサイズを大きくすること、およびタイルがバッファに収まらない場合にタイルを許容する方法を提案している。

スパース テンソルでは、スライスのサイズを選択して、通常、スライスに十分な数のゼロが含まれるようにし、そのほとんどがバッファー内に収まるようにします。しかし、場合によっては、スライスに収まるよりも多くの非ゼロ値が含まれることがあります。この場合、データはバッファから削除されます。

研究者らは、ハードウェアがタイル全体を再度取得して処理することなく、必要なデータのみを再取得できるようにした。これを処理するためにバッファの「末尾」を変更するため、この手法は Tailors と呼ばれます。

その後、オーバーブッキングを使用してスライス サイズを見つける方法も作成しました。 Swiftiles と呼ばれるこのアプローチは、ユーザーが設定したタイルの一定の割合が過剰にサブスクライブされるように、理想的なタイルのサイズをすばやく推定します。 （「Tailors」と「Swiftiles」という名前はテイラー・スウィフトにちなんで付けられました）。

Swiftiles は、ハードウェアがテンソルを検査して理想的なスライス サイズを識別する回数を減らし、計算を節約します。 Tailors と Swiftiles を組み合わせると、オーバーサブスクリプションに対応できない既存のハードウェア アクセラレータの半分のエネルギーで、速度が 2 倍以上になります。

Swiftiles を使用すると、見積りを精緻化するために何度も繰り返し作業する必要なく、これらのタイルがどのくらいの大きさになる必要があるかを見積もることができます。これは、オーバーサブスクリプションがサポートされている場合にのみ機能します。たとえかなり大きな誤差があったとしても、非ゼロの分布方法により、かなりの速度向上が得られます。

将来的には、研究者たちはオーバーサブスクリプションの考え方をコンピューターアーキテクチャの他の側面に適用し、オーバーサブスクリプションの最適なレベルを推定するプロセスの改善に取り組みたいと考えています。