大規模モデルの生成が 2 倍高速化されました。北京大学数学部の卒業生が共同執筆した、単一の GPU で数時間で微調整できる | オープンソース

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

大きなモデルに小さなパーツを追加するだけで、推論速度が 2 倍になります。

追加のモデルをトレーニングしたり、コンピューティング ハードウェアを最適化したりする必要はありません。1台の A100 で、わずか数時間で微調整を完了できます。

「メデューサ」と呼ばれるこの新しい研究は、プリンストン大学、イリノイ大学カリフォルニア大学アーバイン校、カーネギーメロン大学、コネチカット大学が共同で行ったもので、「FlashAttention」の著者であるトリ・ダオ氏も参加している。

現時点では、バークレーの 70 億パラメータの「Llama」 Vicunaへの導入に成功しており、将来的には他の大規模モデルもサポートする予定です。GitHub のホットリストに掲載されています。

しかし実際には、この方法が導入される前から、業界には大規模モデルの推論を加速する方法がありました。主流だったのは、DeepMind が開始した投機的デコードでした。

Medusa はこのアプローチとどう違うのでしょうか?

投機的サンプリングの2つの「バグ」

大規模モデルの推論を高速化するには、まず何がその速度を「制限」するのかを知る必要があります。

計算能力の向上と比較すると、大規模モデルの推論速度はメモリ帯域幅（メモリバウンド）の影響を受けやすくなります。

これは、大規模なモデルにはキャッシュ容量をはるかに超える膨大な数のパラメータがあるため、推論時に重みを外部メモリ(ビデオメモリ)からキャッシュに一度読み込む必要があるためです。このプロセスはメモリ帯域幅によって制限され、通常は非常に遅くなります。

したがって、モデルがバッチ推論を実行する場合、一度に 100 個のトークンを処理する場合と 1 個のトークンを処理する場合の所要時間にほとんど違いはありません。

この特徴に基づいて、 DeepMind は昨年 11 月に「投機的サンプリング」と呼ばれる魔法のような操作を考案しました。

事前に、大規模モデル用の「候補単語」のバッチを生成するために、小規模モデル(ドラフト モデル)をトレーニングします。大規模モデルに「考えさせて」自分で生成させるのではなく、直接「選択」を行うことができます。

小型モデルの生成速度は大型モデルよりも数倍速いため、小型モデルにすでに含まれている単語が「使用可能」であると大型モデルが判断すると、再度ゆっくりと生成することなく、直接その単語が使用されます。

このプロセスは、入力メソッドの候補単語の関連付けに少し似ています。私たち(大きなモデル)が次に使用する単語を考える前に、入力メソッド(小さなモデル)はまずいくつかのオプションをリストします。

良いと思われるものを見つけたら、それを選択して使用します。生成されたものがどれも良くないと思われる場合は、それらを無視して最初からやり直します。

この投機的サンプリング方法は確かに驚くべき結果を達成しており、M2 Ultra で 340 億パラメータの LLaMA モデルを高精度で簡単に実行することもできます。

しかし、このアプローチには 2 つの問題があります。

一方で、大きなモデルの「候補単語」を生成するための小さなドラフトモデルを見つけるのは簡単ではありません。

この小さなモデルは、単なる生成モデルではありません。統一されたインターフェースと類似の確率分布の要件に加えて、生成品質は大きなモデルよりも大幅に劣ることはありません。

MetaがリリースしているLLaMAのようなモデルであれば大丈夫かもしれません。パラメータが数百億の大規模モデルバージョンと、パラメータが数十億の小規模モデルバージョンの両方があります。パラメータが小さいバージョンはドラフトモデルとして使用できます。

ただし、この方法は他の大規模なオープンソース モデルにはあま​​り適用できません。小さなモデルを自分で構築してトレーニングすると、時間がかかるだけでなく、生成された結果が期待どおりにならない可能性があります。

一方、2 つのモデルを組み合わせると、その後のシステム チューニングがより複雑になります。

大規模モデルがそれ自体でシステムであるのに対し、新たに追加されたドラフトモデルは別のシステムを導入することに相当するためです。

これにより、追加のネットワーク伝送とさまざまなハードウェア条件を考慮する必要があるため、モデルの展開がより複雑になり、計算の最適化の難易度がさらに高まります。

これらの問題を解決するために、メデューサが登場しました。

小さなモデルは必要ありません。いくつかの「頭」を追加するだけです。

メデューサ(複数の頭を持つモンスター)は、大規模なモデルの推論を高速化する新しい方法です。

投機的サンプリングと比較して、Transformer 大規模モデルに複数のデコード ヘッドを直接追加することを選択します。各デコード ヘッドは単層フィードフォワード ネットワークです。

これらの追加のデコード ヘッドにより、大規模なモデルは、「歯磨き粉を絞り出す」方法で単語を 1 つずつ生成するのではなく、一度に多くの単語を直接生成できるようになります。

生成精度も合格点。「次の単語の次の単語」を予測する際、Medusa の精度は 60% に達し、現在も最適化が進められています。

その後、ツリーベースの注意メカニズムを使用してこれらの単語を並行して検証し、推論を加速します。

Medusa をベースにすると、Vicuna の 70 億、130 億、330 億のパラメータを持つ大規模モデルの推論速度が1.9 倍以上向上しました。

研究者らは、70億パラメータのモデルについて、さまざまなタスクでの加速効果もテストし、コード生成の最大速度が2.15倍向上することを示しました。

最も重要なことは、Medusa を使用した後、大規模なモデル全体を再トレーニングする必要がないことです。

対照的に、単一の GPU でも、大規模モデルのパラメータをフリーズするだけで、大規模モデルでトレーニングできます。

追加のモデルが追加されないため、分散推論にも適しています。

著者について

この研究には2人の共同筆頭著者がいました。

共同筆頭著者の Tianle Cai 氏は、プリンストン大学の博士課程の学生です。彼の研究対象は、最適化、表現学習、アーキテクチャ設計などです。彼は北京大学数学科学学院を卒業し、応用数学とコンピューターサイエンスの二重学位を取得しました。

共同筆頭著者の Yuhong (Jesse) Li 氏は、イリノイ大学アーバナ・シャンペーン校(UIUC)の博士課程の学生で、効率的な機械学習を専門としています。同氏は北京郵電大学で学士号を取得しました。

さらに、この研究には、FlashAttentionの著者であり、スタンフォード大学の博士課程の学生でもあるTri Dao氏も参加しました。

FlashAttention は、アテンションを高速化し、メモリ使用量を削減する方法であり、PyTorch の標準アテンション実装よりも最大 9 倍高速です。

GitHub アドレス: https://github.com/FasterDecoding/Medusa

研究アドレス: https://sites.google.com/view/medusa-llm