llama.cppを勉強した後、携帯電話で大規模なモデルを実行するのはとても簡単だと分かりました

最近、オープンソース コミュニティでは、大規模モデルの最適化手法を模索する人が増えています。 LLaMa の推論コードをオリジナルの C++ で書き直した llama.cpp というプロジェクトがあり、非常にうまく動作し、幅広い注目を集めています。

いくつかの最適化と量子化された重みにより、これまで想像もできなかったさまざまなハードウェア上で LLaMa モデルをネイティブに実行できるようになります。で：

• Google Pixel 5 スマートフォンでは、7B パラメータ モデルを 1 トークン/秒の速度で実行できます。
• M2チップを搭載したMacbook Proでは、7Bパラメータモデルを使用した速度は約16トークン/秒です。
• 7B モデルは、わずか 0.1 トークン/秒の速度ではありますが、4GB RAM を搭載した Raspberry Pi でも実行できます。

GitHub リンク: https://github.com/ggerganov/llama.cpp

大規模モデルの実装の鍵は、一般化機能に加えて、推論パフォーマンスの最適化にあることはわかっています。しかし、最適化の程度は今や私たちの予想を超えています。 llama.cpp はこれまでに GitHub で 38,000 個のスターを獲得しており、これは LLaMa モデル自体とほぼ同じ数です。 6月には、llama.cpp の作者である Georgi Gerganov 氏が起業し、純粋な C 言語フレームワークを使用して大規模モデルの運用コストを削減することを目的とした新会社 ggml.ai の設立を発表しました。

これを見た多くの人はこう尋ねるでしょう。「どうしてこんなことが可能なのか？」大規模な言語モデルを実行するには、NVIDIA H100 などの GPU が必要ではないでしょうか?この疑問に取り組むために、最近、一部の研究者が大規模モデル推論を取り巻く数学について詳細な研究を行い、その答えを見つけようと試みました。

まず、「AI トレーニングに GPU が必要な理由」から始めましょう。GPU には、ディープラーニングにとって主に 2 つの利点があります。

• メモリ帯域幅が大きい（例：A100: 1935 GB/s、RTX 4090: 1008 GB/s）
• 非常に高い計算能力を備えています (A100: FP16 で 312 TFLOPS、RTX 4090: FP16 で 82.6 TFLOPS)

メモリ帯域幅は、データが HBM メモリ (RAM) からオンチップ メモリに移動するのにかかる時間に関係するため重要です。実際に GPU を使用して計算を行う場合、相関行列をオンチップ メモリに移動する必要がありますが、このメモリは非常に小さいです (A100 では 40 MB、RAM では 40 ～ 80 GB)。メモリ帯域幅はコンピューティング パフォーマンスよりも約 2 桁小さくなります。メモリ帯域幅は推論のボトルネックになることが多いため、これは後で重要になります。

コンピュータアーキテクチャの観点から見ると、効率とコストのバランスをとるために、速度と容量の異なるメモリをさまざまなレベルに分割する必要があります。頻繁にアクセスする必要があるデータは、最も高速だが​​最小のレジスタと L1 キャッシュに配置され、アクセス量が最も少ないデータは、最も低速で最大のメモリ スティックに配置されます。

これは LLaMa 推論タスクにとって何を意味するのでしょうか?推論的な数学から始めましょう。 Kipply の記事を使用して、LLM の推論パフォーマンスについて大まかな計算を行うことができます。

まず、モデルのサイズについて:

• Q、K、V 重み行列はすべて [d_model、d_head] の形状で、レイヤーごとに n_heads があります。注意出力行列も同じ形状で、合計で 4 * [d_model、n_heads * d_head] です。慣例により、GPT スタイルのネットワークでは d_head * n_heads = d_model となります。
• MLPには、[d_model, 4 * d_model]と[4 * d_model, d_model]の形状の2つの重み行列があります。
• 埋め込み行列の形状は[d_vocab, d_model]です。

これにより、GPT のようなモデルのパラメータの数に関する便利な式が得られます。

ここでは、llama.cpp が行う ChatGPT のようなサービスをローカルで実行するケースに焦点を当て、バッチ サイズを 1 と仮定します。効率的な推論を行うには、KV キャッシュをメモリに保存する必要があります。KV キャッシュには各レイヤーの KV 値を保存する必要があり、これは次の値を保存することと同等です。

ここで、n_bytes は各パラメータのバイト数を表すために使用されます。float32 の場合は 4、float16 の場合は 2 などです。中央の 2 は、K 値に対して 1 セットの重みと V 値に対して 1 セットの重みを格納する必要があるためです。

n 層モデルの場合、KV キャッシュの合計メモリは次のようになります。

KV キャッシュをメモリに保存することに加えて、重み自体もメモリに保存する必要があります。これには n_bytes * P バイトが必要です。

これは量子化の主な利点の 1 つです。精度を低くすることで、モデルを保存するために必要なメモリの量を大幅に削減できます。なお、int4 精度では、これらすべてのモデルが Nvidia の A100 (今日のデータセンターで一般的な GPU) のメモリに収まり、最大のモデルを除いてすべてがハイエンドのコンシューマー向け GPU (24GB RAM を搭載した RTX 3090/4090 など) に収まります。

さて、実際に推論を実行するとなると、合計 P 個のパラメータを持つ一連の行列乗算を実行し、行列サイズが (m, n) ベクトル (n,) でコストが 200mn であるため、トークンごとに約 2P FLOPS かかります。

すべての計算が完了したら、LLaMa を使用して推論を実行するための要件を計算しましょう。サンプリングの主な要件は次のとおりです。

• すべてのパラメータに加えて、KV キャッシュもメモリ内に保持されます。
• HBM 内のすべての重みをオンチップ ストレージに読み取ります。自己回帰的にサンプリングしているので、サンプリングするトークンごとにこれを繰り返す必要があります。
• 実際の行列乗算を実行して、ネットワークの出力を計算します。

最新のテンソル プログラミング ライブラリでは、オンチップ メモリへのパラメータの読み取りが非同期的に行われるため、レイテンシはコンピューティング レイテンシまたはメモリ レイテンシの最大値になります。したがって、私たちは次のように書きました。

ここで、B はバッチ サイズです。メモリ帯域幅は約 1.935e12 で、必要な FLOPS は約 3.12e14 なので、バッチ サイズが 161 未満である限り、モデルはメモリ制限を受けます。

バッチ サイズが 1 の場合、つまり、マシン上で予測のストリームが 1 つだけ生成される場合 (ほとんどのハードウェア (Nvidia の GPU など) と同様に) も同じ式になりますが、精度を下げると速度が直線的に向上します。fp32 の代わりに fp16 を使用すると FLOPS が 2 倍になり、int 8 にすると FLOPS がさらに 2 倍になり、int4 を使用するとさらに 2 倍になります。

llama.cpp は現在のディープラーニング推論でより積極的な int4 形式を使用するため、KV キャッシュの RAM 要件は 1.33 GB に削減され、モデル パラメータの VRAM は 16.25 GB に削減されます。素晴らしいですね

メモリ帯域幅はほとんどの場合 FLOPS 数よりもはるかに小さいため、メモリ帯域幅がボトルネックになります。

1) すべてのパラメータをオンチップ メモリにロードし、2) それらのパラメータを使用して結果を計算する必要があるため、FLOPS/トークンの数は必要なメモリ帯域幅と同じであることに注意してください。最新のテンソル プログラミング フレームワークはすべて「メモリへのロード」ビットを非同期に処理できるため、これらすべてが同時に行われ、必要な合計時間は最大 (計算時間、メモリ時間) になります。

NVIDIA A100 で LLaMa を実行する

A100 (80GB PCIe) では、メモリ帯域幅は 1935GB/s です。 int4の計算量は1248 TOPSです。したがって、モデルはメモリによって厳しく制限されます。 65B モデルでは約 30 トークン/秒、7B モデルでは約 277 トークン/秒になると予想されます。

MacBookでLLaMaを実行する

次はメインのフィルムです。Apple MacBook によく搭載されている M1 チップの GPU 帯域幅は 68.25 GB/s で、M1 GPU は最大 5.5 TFLOPS の fp16 計算を実行できます。したがって、int4 を使用する 65B モデルではサンプリング キャップが約 1 トークン/秒、7B モデルでは 10 トークン/秒になると予想されます。

M2 Pro チップの帯域幅は 200 GB/秒、M2 Max の帯域幅は 400 GB/秒なので、65B モデルを使用すると M2 Max が 6 トークン/秒に達するなど、パフォーマンスが大幅に向上することが期待できます。それはノートパソコンとしてはかなり良いですね。

Raspberry Pi 4 で LLaMa を実行する

Raspberry Pi 4 は、13.5 GFLOPS の計算能力と約 4GB/秒のメモリ帯域幅を備えています。これを考慮すると、7B モデルがメモリにバインドされている場合、推論速度は約 2 トークン/秒になると予想されます。ただし、現在は 0.1 トークン/秒程度となっており、これは実際には計算能力の限界によるものであると疑う理由があります。この試みはハードウェアの機能に関する知識がないまま行われました。Raspberry Pi の低精度演算仕様について、これを判断するのに十分な情報を見つけることができませんでした。

要約する

メモリ帯域幅は、ほとんどの場合、トランスフォーマー サンプリングに関連する最も制限的な要因です。量子化のように、これらのモデルのメモリ要件を削減するものは、モデルの提供を容易にします。これは、蒸留 (または単に小さなモデルを長時間トレーニングすること) が非常に重要であるもう 1 つの理由です。

OpenAIの科学者Andrej Karpathy氏はこの観察結果についてさらに説明した。

彼は次のように述べました。「並列推論とトレーニングに加えて、ヒント トークンは 1 つずつシリアルにデコードされるのではなく、LLM によって並列にエンコードできるため、batch_size = 1 の場合でもヒント エンコードを並列化できます。」プロンプトが長くなるにつれて、MacBook の推論パフォーマンスは A100 に比べてさらに低下します。

しかし一方で、Apple の M2 チップは、大規模モデルでの推論タスクにおいて強力な能力を発揮しているようだ。 「その結果、M2 Ultra は現在入手可能な最も小型で、最も美しく、すぐに使用でき、最もシンプルで、最も強力な個人用 LLM ノードとなりました。」

陳天奇氏もこの見解に同意している。

もちろん、これはタダのランチではありません。基本的に、精度を低くすると精度がいくらか失われ、より大きなモデルの応答を混乱させたり錯覚を引き起こしたりする奇妙な答えが出る可能性があります。ただし、モデルのパラメータが増えるにつれて、品質の低下は減少します。したがって、モデル ボリュームが非常に大きい場合、その差は無視できるほど小さくなる可能性があります。さらに、これは単なる推論コストです。トレーニングは全く別の問題です。

さまざまなパフォーマンスパラメータを比較検討することで、おそらく私たちはすぐに、真に「インテリジェント」なデバイスを手に入れることになるでしょう。