Llama の推論特性

Llama-2-70BはGPT-3.5に対する魅力的な代替手段ですが、低コストの言語モデルを求めている場合、OpenAIのAPIから乗り換える価値はないかもしれません。

価格とレイテンシーを検討する際は：

補完処理が多いワークロードには、Llama‑2 を提供すべきではありません

その代わりに、Llama はプロンプト主導のタスク（分類など）に最適です。次の場合は Llama-2 を使うのが合理的なこともあります：

あなたのワークロードにはプロンプトトークンがありません（直感に反しますが、後で説明します）
バッチ処理ジョブを実行している

それ以外の場合は、GPT-3.5 の方が低コストで高速です。

簡単な注意書きです。Llama を GPT-3.5 より優先して使う理由の一つは、ファインチューニングが可能なことです。ただし本稿では、コストとレイテンシーのみを検討します。Llama-2 を GPT-4 と比較しないのは、どちらかと言えば 3.5 相当のモデルに近いからです。ベンチマーク結果もこの主張を裏付けています。

Figure 1: GPT-3.5 dominates Llama in all benchmarks here. — Figure 1: GPT-3.5 dominates llama in all benchmarks here.

本稿では、ほぼ同等のレイテンシーを前提に、Llama‑2‑70B と GPT‑3.5‑turbo の提供コストを比較して、これらの主張を実証します。Llama は 80‑GB A100 GPU を 2 基で提供しており、これは（16 ビット精度で）Llama をメモリに載せるための最小要件です 3。

2×A100環境では、Llama の補完トークン単価は GPT-3.5 より割高であることがわかりました。8×A100 なら競争力のある価格になる可能性がありますが、待機時間（レイテンシ）が許容できないほど高くなる懸念があります。

一方で、プロンプトトークンあたりのコストは、Llama のほうが GPT-3.5 よりも3倍超安価です。

Transformer 数学入門

いくつかの単純な数式を用いて、Llama-2について次のことを示します。シーケンス長が $N$ 、バッチサイズが $B$ の場合：

\text{トークンあたりのFLOPS}=140⋅10^9 \text{FLOPs}

\text{メモリ帯域幅 / トークン} = 140 \text{ GB/s} + 320⋅N⋅B \text{ KB/s}

$140$ は、モデルのパラメータ数の2倍に由来し、320 KB/s はいくつかの計算から導かれます。次のセクションでは、これらの数値にどのように到達したかを説明します。

トランスフォーマーの数式を見事に解説している論文やブログ記事はほかにもあります。推論に関しては、Kipply の投稿が優れたリファレンスです。また、トランスフォーマーの FLOPs に用いられる簡略化した方程式を広めたのは、Scaling Laws だと考えています。

これらの数値を検証するために、まず Llama のアーキテクチャの詳細から始めます。隠れ次元は 4096、アテンションヘッド数は 64、レイヤー数は 80、各アテンションヘッドの次元は 128 です。

d_{model}=4096

n_h=64

n_l=80

d_{head}=128

モデルのFLOPsを計算する

順伝播におけるFLOPs（浮動小数点演算回数）は $≈2P$ で、 $P$ はモデルのパラメータ数です。モデル内のすべてのパラメータは、ある重み行列 $M∈R^{m×n}$ に属します。各入力トークンに対して、各行列はそのトークンを表すベクトルとの行列積にちょうど1回ずつ用いられます。

各 $M$ に対して、次元 $x$ のベクトル $m$ を左から乗算します。このベクトル行列乗算の総FLOPsは $2mn 6$ 、つまり重み行列のエントリ数の2倍です。transformerの全重み行列のエントリ総数は総パラメータ数 $P$ であり、これによりアテンション機構を除いて総 $2P$ FLOPsが得られます。

（比較的）短いシーケンスを扱う Llama のような大規模モデルでは、FLOPs に対するアテンションの寄与は無視できるほど小さいです。各レイヤーおよび各アテンションヘッドにおけるアテンションの計算は次のとおりです：

\text{softmax}\left(\frac{Q^T K}{\sqrt{d_{\text{head}}}}\right) V

$Q^TK$ では、 $d_{head}$ ベクトルに $d^{head}×N$ 行列を乗じる必要があり、これは $2dheadN$ FLOPs です。スケーリング係数と Softmax は無視できます。最後に、Attention ベクトルに $V$ を乗じるのに、さらに $2d_{head}N$ FLOPs が必要です。すべてのアテンションヘッドとレイヤーにわたって合計すると、 $4⋅d_{model}⋅nl⋅N=1.3⋅N$ MFLOPs になります。したがって、最大シーケンス長 8192 の場合でも、アテンションが占めるのは全体 140 GFLOPs のうち 10.5 GFLOPs にとどまります。十分に小さいため、単純化のために無視します。

補完はプロンプトよりも高いメモリを必要とします

トークンを生成する際には、各トークンを生成するたびにモデルの全ウェイトとKVキャッシュを再読み込みする必要があります。これはどういう意味でしょうか。任意の行列積を行うには、各行列のウェイトをRAMからGPUのレジスタへ読み込む必要があります。ユニークな行列が十分に多い場合、行列積そのものよりもウェイトの読み込みがボトルネックになります。そこで、プロンプトと補完それぞれにおける、トークンがモデル内部を通過する経路を比較してみましょう。

トランスフォーマーによるトークン生成のメモリ経路

これを示すために、単純な1層のトランスフォーマーがトークンのバッチを生成する際の（非常に大まかにスケッチした）処理の流れを追ってみましょう。

入力の埋め込み行列を読み取り、 $W_e$ バッチ内の各入力に対して対応する埋め込みベクトルを計算します。
バッチ内の各入力に対してベクトル $W_q,W_k,W_v$ 行列をメモリから読み出し、 $q_i,k_i,v_i$ （ベクトル）を計算します。
キャッシュされたキーとバリューを読み出す必要があるアテンション処理を実行します。これにより、各入力に対してベクトルが返されます。
メモリから $W_o$ を読み取り、前のステップの出力に乗算します
ステップ1の出力を読み取り、それをステップ4の出力に加算し、その後にlayernormを実行します。
私たちは $W_{ff_1}$ を読み取り、乗算して最初のフィードフォワード層の出力を得ます。
私たちは $W_{ff_2}$ を読み取り、乗算して第2のフィードフォワード層の出力を得ます。
ステップ5の出力を読み取り、ステップ7の出力に加算してから、LayerNorm を実行します。
アンエンベディング層を読み取り、 $W_{e}^{T}$ を行列同士の積で掛け合わせ、バッチ内の各入力に対するトークンの対数確率を得ます。
次のトークンをサンプリングし、それを手順1にフィードバックします。

メモリ要件を見積もってみましょう。手順1、2、4、6、7、9では、モデルの全パラメータをおおむね一度ずつ読み込みます。手順3では、各バッチ要素のKVキャッシュを読み込みます。すべての手順で、中間アクティベーションも読み込みますが、モデルサイズと比べれば無視できる程度です。したがって必要なメモリ帯域は「モデルの重み + KVキャッシュ」です。バッチサイズを増やしても、KVキャッシュを除けばメモリ要件はほぼ一定のままです！この点には後ほど戻ります。これはトークンあたり

トランスフォーマーでプロンプトのトークンを処理する際のメモリ経路

プロンプトを処理する際、モデルの全ウェイトは一度だけ読み込みますが、Attention によるメモリコストが発生します。同一のトランスフォーマを通過するシーケンスのバッチのおおまかな流れを考えてみましょう。

入力埋め込み行列 We を読み取り、バッチ内の各シーケンスに対応する埋め込み行列を計算します。
各 $W_q,W_k,W_v$ 行列をメモリから読み出して、 $Q_i,K_i,V_i$ （いずれも行列）を計算します
アテンション演算を実行します
メモリからWoを読み取り、前のステップの出力と乗算します
ステップ1の出力を読み取り、それをステップ4の出力に加え、続いてLayerNormを実行します
私たちは $W_{ff_1}$ を読み取り、乗算して最初のフィードフォワード層の出力を得ます
私たちは $W_{ff_2}$ を読み取り、乗算して第2のフィードフォワード層の出力を得ます
ステップ5の出力を読み取り、それをステップ7の出力に加算し、その後にlayernormを実行します
アンエンベディング層 $W_u=W_{e}^{T}$ を読み取り、乗算してプロンプト系列のトークン対数確率を得ます

ステップ1、2、4、6、7では、モデルのすべてのパラメータを読み込みます。ステップ3では、FlashAttention を用いた attention 演算を実行します。これは（妥当な長さのシーケンスとバッチサイズにおいて）モデルの重みを読み出すよりもはるかに少ないメモリ帯域幅しか必要としません。すべてのステップでアクティベーションを読み出しますが、これはモデルサイズと比べると無視できる程度の行列です（こちらも妥当な長さのシーケンスおよび/またはバッチにおいて）。これはすべてのトークンに対するメモリ要件である点に注意してください。

結論として、プロンプト処理におけるトークンあたりのメモリ要件は、トークン生成時よりも大幅に少なくなります。これは、プロンプトに対してシーケンス次元で行列積をバッチ処理しているためです。

モデルの重みが必要とするメモリ帯域幅

16ビット精度のモデル重みは、メモリを $2⋅70=140$ GB使用します。

KVキャッシュに必要なメモリ帯域幅

私たちのKVキャッシュのサイズは、ニューラルネットの全レイヤー・全ヘッドにおける、これまでのすべてのトークンに対するキーとバリューの合計サイズであり、トークンおよびバッチ要素1つあたり320 MBです。

Llama 2 はマルチヘッド・アテンションを削除することにしました。ただし、マルチクエリ・アテンションの代わりにグループ化クエリ・アテンションを使用し、これにより性能が向上します。これにより、キーとバリューのヘッド（またはグループ）が通常のマルチヘッドでは 128、マルチクエリでは 1 であるのに対し、 $ng$ になります。

トークン数 $N$ に対して、KV キャッシュのサイズは $2 n_g n_l d_{\text{head}} N = 1.6 N \cdot 10^5$ となります。16 ビット精度を用いると、これは $320N$ KB です。バッチサイズを $B$ とすると、 $320⋅NB$ KB になります。

補完の場合、トークンあたりのメモリ要件は次のとおりです。

\text{メモリ / トークン}=140GB+320⋅N⋅B \text{ KB}

短いシーケンスや小さなバッチでは第1項が支配的になります。そうでなければ第2項のほうがはるかに大きくなります。ただし、利用可能なメモリが160 GBでモデル自体が140 GBを占めるため、今回の実験ではKVキャッシュがメモリ帯域幅に小さな負荷を与えることになります。

プロンプトのメモリ帯域幅はおおよそ次のとおりです：

\text{メモリ / トークン} = \frac{140}{N} \text{ GB}

コミュニケーションのオーバーヘッド

話を簡単にするため、ここでは通信コストは無視します。モデル並列化まで考慮すると議論が大幅に複雑になるためです。合理的な前提として、いずれの計算に対しても十分に大きな遅延を追加しないと見なせます（特に、Llama を 2 枚の GPU に分割しているだけだからです）。

プロンプト処理はとても低コストです

プロンプト処理、すなわち最初のトークンまでの時間は、Transformer 推論でもっとも効率の良い部分であり、GPT-3.5 と比べて価格が3倍程度下がると見込めます。

パラメータ数が $P$ のモデルに、長さ $N$ のプロンプトを与える場合、プロンプト処理に必要なメモリは約 $2P$ バイト、計算量は $2PN$ FLOPsとなります。A100は行列積で312 TFLOPs、メモリ帯域幅で2 TB/sを処理できるため、シーケンス長が N>156 の場合は計算がボトルネックになります。 9

A100 では、FLOPs の活用率（MFU）は最大でおよそ 70% をわずかに下回る程度になる見込みです。これは約 200 TFLOPs に相当します。80 GB の A100 を 2 基使用すると費用は約 $4.42/時で、秒あたりにすると $0.0012 です。Llama の FLOPs 要件は 1 トークンあたり 140 TFLOPs です。2 基の A100 による合計 FLOPs を踏まえると、1 秒あたりのトークン数は次のように見積もれます：

2⋅200⋅10^{12}/140⋅109=2.86⋅10^3 \text{ トークン / 秒}

価格は次のとおりです：

$0.00042 / 1K トークン

GPT-3.5 の$0.0015と比べても、これは破格です！具体的には、ほぼ4倍の値下げになります！

レイテンシもかなり良好です！2基のGPUでバッチサイズ1の場合、512トークンを170ms、1536トークンを530msで処理できる見込みです。

これらの主張を、実際の数値で検証しましょう。私たちは、huggingface の text-generation-inference リポジトリの社内フォークを用いて、Llama-2 のコストとレイテンシを計測しています。

Figure 2: Each datapoint measures a different batch size. For prompt tokens, we always do far better on pricing than gpt-3.5, but trail slightly behind on gpt-3.5's latency of 0.4s for 3.6K tokens. — Figure 2: Each datapoint measures a different batch size. For prompt tokens, we always do far better on pricing than GPT-3.5, but trail slightly behind on GPT-3.5's latency of 0.4s for 3.6K tokens.

ご覧のとおり、価格は GPT-3.5 の $0.0015/1k tokens よりも大幅に優れています。長いシーケンスでは最初のトークンまでの時間でやや遅れを取っているように見えますが、解決策はきわめて簡単です。Llama を 2 ではなく 8 つの GPT にまたいで並列化すれば、ほぼ 4 倍の高速化が得られ、プロンプトにおいては Llama-2 が GPT-3.5 を凌駕します！

トークン生成は遅く、非常に高コストです

理論上は、補完において GPT-3.5 と同等の価格を実現することは可能ですが、実際には多くの場合それより不利になります。

トークン生成時には、計算ボトルネックからメモリボトルネックへと移行します。11．バッチサイズを1と仮定して、達成可能なスループットを算出しましょう。

各80 GB A100は、GPUあたりピーク時のメモリ帯域幅が2TB/sです。ただし、FLOPsの活用率と同様に、推論ワークロードではその60〜70%程度（1.3 TB/s）に落ち着くと見込むのが妥当です。小さなバッチではKVキャッシュの影響は無視できるため、2基のA100でのスループットは次のようになります：

\frac{2 \cdot 1.3 \cdot 10^{12} \,\text{B/s}}{140 \cdot 10^{9} \,\text{B/token}} = 18.6 \,\text{tokens/s}

新しい価格は大幅に悪化しています。$0.0012/秒では、コストは…

$0.066 / 1K トークン

これは、GPT-3.5 レベルのモデルとしては価格も速度もひどすぎます！ただし、先ほどのバッチサイズに関する注意を思い出してください。私たちはメモリがボトルネックになっているため、生成速度を落とさずにバッチサイズを増やせます。バッチサイズが大きいほど、コストは下がります。

KVキャッシュが最終的にGPU RAMをすべて占有してしまうため、無限大まで増やすことはできません。幸いにも、grouped-query attention によりこの問題は軽減されます。トークン数を N、バッチサイズを $B$ 、16-bit 精度とすると、キャッシュは $3.2⋅N⋅B⋅105$ Byte になります。4096 トークンの場合、バッチサイズ 1 で 1.3 GB のメモリに相当します。2-A100 機では 160 GB の空き容量があります。モデル自体が 135 GB を使用するため、KV キャッシュに使えるのは 25 GB のみです。メモリ格納の追加の非効率性もあるため、長いシーケンス長での最大バッチサイズは約 8 となります。

（おおよそ）8倍の高速化を考慮すると、価格は$0.00825 / 1K tokensになると見込まれます。これは依然としてGPT-3.5-turboより劣りますが、差は縮まります。より短いシーケンス長（合計1kトークン）の場合、バッチサイズを32まで増やせるため、価格は$0.00206 / 1K tokensになります。理論上、これはGPT-3.5-turboと競合し得ます。

別の解決策は、GPU の台数を増やすことです。8 台の 80 GB A100 をレンタルすると、合計 1.28TB のメモリになります。モデルの重みを除くと、KV キャッシュに 1TB 超のメモリを割り当てられるため、512 トークン超のバッチサイズが可能です。なお、実際にはコストが 512 倍も低下するわけではありません。KV キャッシュがモデルサイズの 8 倍のメモリ帯域を占有するため、コスト低下はおおよそ 64 倍に近くなります。

より多くの計算資源を使うことで、レイテンシの問題も解決できます。GPT-3.5 は約 70 TPS に到達します。モデルを 2 GPU ではなく 8 GPU に分割すれば、約 74.4 tokens/s まで向上するはずです。

この実験の実行時には 8 基の A100 を使用できませんでしたが、2 基の 80 GB A100 における数値を見てみましょう。

測定済みの生成パフォーマンス

Figure 3: For all datapoints, we measure price per generated tokens when generating 512 tokens.

これらの数値は、メモリ帯域幅の計算から予想される値にかなり近いものです。

ご覧のとおり、バッチサイズを大きくすると、1Kトークンあたりの価格はほぼ線形に低下します。しかし、特にシーケンス長が長い場合、依然としてGPT-3.5の$0.002/1Kトークンという価格にはかなり届きません。

バッチサイズが大きいと許容できない遅延が発生する

大きなバッチサイズで生成を実行すると、コストはGPT-3.5並みに競争力のある水準になりますが、最初のトークンまでの時間が増加します。バッチサイズを大きくするほど、コストは線形に下がる一方で、最初のトークンまでの時間は線形に伸びます。

バッチサイズを64にすると、GPT-4よりも有利な料金になります。ただし、バッチサイズ64では次のような状況になります：

わずか512トークンで、最初のトークンまでの時間が約3秒！
3596トークンに対するバッチサイズ64の場合は、20.1秒です。

その結果として、OpenAI の API と比べて Llama‑2 を選ぶ妥当性があるワークロードの種類は次のとおりです。

生成トークンがほとんど、または全くない大きなプロンプト――純粋なプロンプトの処理は驚くほど簡単です。
短い、またはプロンプトなしでトークンを生成するケースでは、バッチサイズを最大で64超まで調整でき、レイテンシを犠牲にせずに GPT-3.5-turbo と競合する価格帯を実現できます。
レイテンシが重要でないオフラインのバッチ処理ジョブ。

バッチサイズを上げるには、常に大規模なワークロードが必要ですが、ほとんどのスタートアップにはそれがありません。多くのユーザーや多くのワークロードでは、実際の使用は非常にバースト的です。もちろん、候補となる解決策はオンデマンドGPUのオートスケーリング（スケールアップ／スケールダウン）ですが、それでもGPUあたり平均して最大スループットの50%程度と見込むべきでしょう——特にコールドブート時間があることを考えると。

プロンプト量の多いタスクにはオープンソースのモデルを、生成量の多いタスクには GPT-3.5 のようなクローズドソースのモデルを使うことを推奨します。

量子化

ほとんどの量子化手法はロスありであり、つまり性能がいくらか低下します。8ビット量子化でもおおむね競合できる性能を達成できる可能性が高く、算出したすべての数値で価格を2倍引き下げられます！ 量子化による損失と不完全なリソース活用は相殺されるため、両方を考慮すると、実測した’価格と同程度になると見込んでいます！

しかし、ほとんどのオープンソースの量子化手法の目標は、大規模に最適化されたスループットを実現することではなく、少数台・小規模なコンシューマー向けGPUで容易にデプロイできるようにすることです。

パフォーマンスを維持しつつ、さらに低精度の量子化を最適化するオープンソースライブラリがいくつか存在します。しかし、これらのライブラリはスケールでのスループットではなく、データセンター外の少数/小規模GPUでのモデル提供に最適化されています。具体的には、低バッチ推論（主にバッチサイズ1）を対象に最適化されています。最大でも3～4倍の高速化を提供する一方で、それでもコストは$0.017/1Kトークンに相当します。

あるいは、言語モデルについて真面目に語りたい方は、TwitterでDMしてください。

付録

以下は、Llama-2-70B のレイテンシに関する測定値と、いくつかの追加の算出指標を含むデータポイントの一覧表です