LLMの速度とメモリを最適化

GPT3/4、Falcon、Llama などの大規模言語モデル (LLM) は、人間中心のタスクを処理する能力が急速に向上しており、現代の知識ベース産業において不可欠なツールとして確立されています。 しかし、これらのモデルを実世界のタスクにデプロイすることは依然として困難です：

• 人間に近いテキスト理解および生成能力を発揮するためには、LLMは現在、数十億規模のパラメータで構成される必要があります（Kaplan et al、Wei et. al参照）。その結果、推論時のメモリ要件が大幅に高まります。
• 多くの実世界のタスクでは、LLM に広範な文脈情報を提供する必要があります。そのため、推論時に非常に長い入力シーケンスを処理できるモデルの能力が求められます。

これらの課題の核心は、LLMの計算能力とメモリ能力を拡張すること、特に大規模な入力シーケンスを処理する場合に重要となります。

このガイドでは、LLMの効率的なデプロイメントのための効果的な技術について説明します：

1. 低精度: 研究により、8ビットおよび4ビットといった数値精度を下げた状態で動作させることで、モデルの性能に大きな低下をもたらすことなく計算上の利点を得られることが示されています。

2. Flash Attention: Flash Attentionは、アテンションアルゴリズムの変種であり、よりメモリ効率の高いアプローチを提供するだけでなく、GPUメモリの最適化された利用により効率の向上も実現します。

3. アーキテクチャの革新: LLMは推論時において常に同じ方法、すなわち長い入力コンテキストを伴う自己回帰的なテキスト生成としてデプロイされるため、より効率的な推論を可能にする専用のモデルアーキテクチャが提案されています。モデルアーキテクチャにおける最も重要な進歩は、Alibi、Rotary embeddings、Multi-Query Attention (MQA)、およびGrouped-Query-Attention (GQA)です。

本ガイドでは、テンソルの観点からオートレグレッシブ生成を分析します。低精度採用のメリットとデメリットを掘り下げ、最新のAttentionアルゴリズムを包括的に解説し、改良されたLLMアーキテクチャについて議論します。その過程で、各機能の改善点を示す実践的な例を実行します。

1. 低精度

LLMのメモリ要件は、LLMを重み行列とベクトルの集合、テキスト入力をベクトルのシーケンスとみなすことで最もよく理解できます。以下では、weightsという定義を用いて、すべてのモデルの重み行列とベクトルを指すものとします。

本ガイドの執筆時点では、LLMは少なくとも数十億パラメータから構成されています。各パラメータは小数で表され、例えば4.5689のように、通常はfloat32、bfloat16、またはfloat16形式で保存されます。これにより、LLMをメモリにロードするために必要なメモリ量を簡単に計算できます。

X十億パラメータのモデルの重みを読み込むには、float32精度で概ね4 * X GBのVRAMが必要です

現在、モデルはフルのfloat32精度で学習されることはほとんどなく、通常はbfloat16精度、あるいはそれより頻度は低いがfloat16精度で学習されます。したがって、経験則は次のようになります。

X十億パラメータのモデルの重みをロードするには、bfloat16/float16精度で約2 * X GBのVRAMが必要です

短いテキスト入力（1024トークン未満）の場合、推論に必要なメモリは、重みの読み込みに必要なメモリによって大きく支配されます。したがって、ここでは推論に必要なメモリは、モデルをGPU VRAMに読み込むのに必要なメモリと等しいと仮定します。

bfloat16 でモデルを読み込む際に、おおよそどの程度の VRAM が必要になるかを示す例をいくつか挙げます。

• GPT3には 2 * 175 GB = 350 GB の VRAM が必要です
• Bloom には 2 * 176 GB = 352 GB の VRAM が必要です
• Llama-2-70bには 2 * 70 GB = 140 GB の VRAM が必要です
• Falcon-40bには 2 * 40 GB = 80 GB の VRAM が必要です
• MPT-30b には 2 * 30 GB = 60 GB の VRAM が必要です
• bigcode/starcoderには 2 * 15.5 = 31 GB の VRAM が必要です

本ドキュメントの執筆時点において、市場で入手可能な最大級のGPUチップは、80GBのVRAMを備えたA100およびH100です。前述のモデルの多くは、読み込みだけでも80GBを超えるため、テンソル並列および/またはパイプライン並列が不可欠となります。

🤗 Transformers は、それぞれの設定クラスに base_tp_plan を持つサポート対象モデルのテンソル並列処理に対応しました。Tensor Parallelism については こちらをご覧ください。また、テンソル並列処理に配慮したモデルの作成にご興味がある場合は、text-generation-inference ライブラリもぜひご覧ください。

Naiveなパイプライン並列化は、そのままサポートされています。これを行うには、device="auto"を指定してモデルを読み込むだけで、こちらで説明されているように、異なるレイヤーが利用可能なGPUに自動的に配置されます。ただし、非常に効果的であるものの、このNaiveなパイプライン並列化では、GPUのアイドル状態の問題に対処していません。この問題に対処するには、こちらで説明されているより高度なパイプライン並列化が必要です。

8 x 80GB A100ノードにアクセスできる場合、BLOOMは次のようにロードできます。

!pip install transformers accelerate bitsandbytes optimum

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

device_map="auto" を使用すると、アテンション層が利用可能なすべての GPU に均等に分散されます。

このガイドでは、bigcode/octocoderを使用します。これは単一の40 GB A100 GPUデバイスチップ上で実行できるためです。なお、これから適用するすべてのメモリおよび速度の最適化は、モデル並列化やテンソル並列化を必要とするモデルにも同様に適用できます。

モデルは bfloat16 精度で読み込まれるため、上記の目安に従うと、bigcode/octocoder で推論を実行するためのメモリ要件は約 31 GB の VRAM になると予想されます。実際に試してみましょう。

まずモデルとトークナイザーを読み込み、両方をTransformersのpipelineオブジェクトに渡します。

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

出力:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

よし、これで結果を直接使用して、バイトをギガバイトに変換できます。

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

ピークアクセラレータメモリ割り当てを測定するために、torch.accelerator.memory.max_memory_allocated を呼び出します。

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

出力:

29.0260648727417

大まかな計算とかなり近い結果になりました！バイトからキロバイトへの換算には1000倍ではなく1024倍の乗算が必要なため、数値が完全に正確ではないことがわかります。したがって、この大まかな計算式は「最大X GB」という計算として理解することもできます。 なお、モデルを完全なfloat32精度で実行しようとした場合、VRAMに64 GBもの容量が必要だったことに注意してください。

現在、ほぼすべてのモデルはbfloat16で学習されています。GPUがbfloat16をサポートしている場合、モデルをフルのfloat32精度で実行する理由はありません。float32は、モデルの学習に使用された精度よりも良い推論結果をもたらすことはありません。

Hub 上でモデルの重みがどの形式で保存されているか不明な場合は、チェックポイントの config の "dtype" を確認できます。例として、こちらを参照してください。from_pretrained(..., dtype=...) で読み込む際は、config に記載されているのと同じ精度タイプにモデルを設定することを推奨します。ただし、元のタイプが float32 の場合は、推論に float16 もしくは bfloat16 のいずれかを使用できます。

割り当てられたすべてのメモリを解放して、ピーク時の割り当て済み GPU メモリを正確に測定できるようにするため、flush(...) 関数を定義しましょう。

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.accelerator.empty_cache()
  torch.accelerator.reset_peak_memory_stats()

次の実験のために、ここで呼び出しましょう。

flush()

Accelerate ライブラリでは、XPU、MLU、NPU、MPS などさまざまなハードウェアバックエンドを考慮した、デバイス非依存のユーティリティメソッド release_memory も利用できます。

from accelerate.utils import release_memory
# ...

release_memory(model)

では、GPU の VRAM が 32 GB 未満の場合はどうでしょうか。モデルの重みを 8 ビットまたは 4 ビットに量子化しても、性能に大きな劣化は生じないことが確認されています（Dettmers et al.を参照）。最近のGPTQ paper🤯に示されているように、モデルは 3 ビットや 2 ビットまで量子化しても、許容範囲の性能劣化で運用可能です。

詳細には立ち入らずに述べると、量子化スキームは、モデルの推論結果をできるだけ正確に保ちながら（言い換えれば bfloat16 にできるだけ近づけながら）、重みの精度を下げることを目指しています。 量子化はテキスト生成に特に有効です。なぜなら、私たちが気にすべきは最も確率の高い次のトークンの集合を選ぶことだけで、次のトークンのlogit分布の正確な値はそれほど重要ではないからです。 重要なのは、次のトークンのlogit分布が概ね同じ状態を保ち、argmaxやtopkの操作が同じ結果を返すことです。

量子化にはさまざまな手法がありますが、ここでは詳細には触れません。一般的に、すべての量子化手法は次のように動作します。

• 1. すべての重みをターゲットの精度に量子化します
• 2. 量子化された重みを読み込み、bfloat16 精度のベクトル入力シーケンスを渡します
• 3. 重みを動的にデクォアンタイズして bfloat16 に変換し、入力ベクトルとの演算を bfloat16 精度で実行します

要するに、これは 入力と重み行列の 乗算であり、 $X$ は 入力, $W$ は重み行列、 $Y$ は出力です: $$Y = X * W$$

に変更されます $$Y = X * \text{dequantize}(W)$$

すべての行列積に対して行われます。入力ネットワークグラフを通過するにつれて、すべての重み行列に対してデクアンタイズと再クアンタイズが順次実行されます。

したがって、量子化された重みを使用しても、推論時間はむしろ増加し、短縮されないことがよくあります。 理論はここまでにして、実際に試してみましょう。Transformers で重みを量子化するには、bitsandbytes ライブラリがインストールされていることを確認する必要があります。

!pip install bitsandbytes

その後、from_pretrained に load_in_8bit=True フラグを追加するだけで、モデルを8ビット量子化で読み込むことができます。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_8bit=True), pad_token_id=0)

では、例を再度実行してメモリ使用量を測定しましょう。

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

出力:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

よし、前回と同じ結果が得られたので、精度の低下はありません！では、今回どのくらいのメモリが使用されたか見てみましょう。

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

出力:

15.219234466552734

大幅に少なくなりました!メモリ使用量はわずか15 GB強まで低下し、4090などのコンシューマー向けGPUでこのモデルを動作させることが可能になりました。 メモリ効率の向上が顕著に確認でき、モデルの出力品質にはほぼ劣化が見られません。ただし、推論時の処理速度がわずかに低下していることも確認できます。

モデルを削除し、メモリを再度フラッシュします。

del model
del pipe

flush()

4-bit量子化によるピークGPUメモリ消費量を確認しましょう。モデルを4-bitに量子化するには、以前と同じAPIを使用し、load_in_8bit=Trueの代わりにload_in_4bit=Trueを渡します。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_4bit=True), pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

出力:

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

ほぼ以前と同じ出力テキストが表示されていますが、コードスニペットの直前にpythonが欠けています。メモリ使用量を確認してみましょう。

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

出力:

9.543574333190918

わずか9.5GB!150億パラメータ超のモデルとしては、本当に少ないですね。

ここではモデルの精度がほとんど低下していませんが、実際には4ビット量子化は、8ビット量子化や完全なbfloat16推論と比較して、異なる結果をもたらすことがよくあります。試すかどうかはユーザー次第です。

また、4ビット量子化ではより積極的な量子化手法が使用されるため、8ビット量子化と比較して推論がやや遅くなる点にも注意してください。これは $\text{quantize}$ と $\text{dequantize}$ の処理が推論中に時間がかかるためです。

del model
del pipe

flush()

全体として、OctoCoder を 8-bit 精度で実行すると、必要な GPU VRAM が 32G からわずか 15GB に減少し、4-bit 精度でモデルを実行すると、さらに必要な GPU VRAM が 9GB 強まで減少することがわかりました。

4-bit量子化により、RTX3090、V100、T4など、多くの人が入手しやすいGPU上でモデルを動作させることができます。

量子化に関する詳細や、4ビットよりもさらに少ないGPU VRAMメモリでモデルを量子化する方法については、GPT-QModelの実装を参照することをお勧めします。

結論として、モデルの量子化はメモリ効率の向上と引き換えに、精度や場合によっては推論時間を犠牲にする点に留意することが重要です。

GPU メモリが使用ケースにおいて制約にならない場合、量子化を検討する必要は通常ありません。しかし、多くの GPU では量子化手法なしでは LLM を実行できず、そのような場合、4 ビットおよび 8 ビットの量子化スキームは非常に有用なツールとなります。

より詳細な使用方法については、Transformers Quantization Docsを参照することを強くお勧めします。 次に、より優れたアルゴリズムと改良されたモデルアーキテクチャを用いて、計算効率とメモリ効率をどのように向上させることができるかを見ていきましょう。

2. Flash Attention

現在の高性能LLMは、フィードフォワード層、活性化層、レイヤーノーマライゼーション層、そして最も重要な自己注意層からなる、ほぼ同じ基本的なアーキテクチャを共有しています。

自己アテンション層は、入力トークン間の文脈的な関係をモデルに理解させることで、大規模言語モデル (LLM) の中心をなしています。 しかし、自己アテンション層の GPU メモリ消費のピークは 二次的に 増加し、計算量とメモリ複雑性の両方が入力トークン数（ シーケンス長) 以下では $N$ . 短い入力シーケンス（1000 トークン程度まで）ではあまり顕著ではありませんが、長い入力シーケンス（約 16000 トークン程度）では深刻な問題となります。

詳しく見てみましょう。長さ $\mathbf{O}$ の自己アテンション層の出力 $\mathbf{X}$ を計算する式は次のとおりです。 $N$ $$\textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X}$$

$\mathbf{X} = (\mathbf{x}_1, ... \mathbf{x}_{N})$ は、アテンション層への入力シーケンスとなります。射影 $\mathbf{Q}$ と $\mathbf{K}$ はそれぞれ $N$ 個のベクトルからなり、 $\mathbf{QK}^T$ のサイズは $N^2$ .

LLM には通常複数のアテンションヘッドがあり、複数の自己アテンション計算を並列で実行します。LLM が 40 個のアテンションヘッドを持ち、bfloat16 精度で動作すると仮定すると、 $\mathbf{QK^T}$ 行列を保存するためのメモリ要件は $40 * 2 * N^2$ バイトと計算できます。 $N=1000$ の場合、VRAM は約 50 MB しか必要としませんが、 $N=16000$ の場合、VRAM は 19 GB 必要になり、 $N=100,000$ の場合、 $\mathbf{QK}^T$ 行列を保存するだけでもほぼ 1 TB 必要になります。

長くなりそうなので要約すると、デフォルトの自己注意アルゴリズムは、大きな入力コンテキストに対しては、メモリコストが急激に高くなり、実用的ではなくなります。

LLMのテキスト理解と生成能力が向上するにつれ、より複雑なタスクに応用されています。かつてモデルが数文の翻訳や要約を処理する程度だったものが、現在はページ全体を扱っており、広範な入力長を処理する能力が求められています。

長い入力に対する膨大なメモリ要件をどうすれば解消できるでしょうか。 $\mathbf{QK}^T$ 行列を必要としない、新しい自己アテンション機構の計算方法が必要です。 Tri Dao らは まさにそのような新しいアルゴリズムを開発し、 Flash Attention.

要するに、Flash Attention は $\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T)$ の計算を分離し、代わりに複数の softmax 計算ステップを反復することで、出力の小さなチャンクを計算します: $$\textbf{O}_i \leftarrow s^a_{ij} * \textbf{O}_i + s^b_{ij} * \mathbf{V}_{j} \times \text{Softmax}(\mathbf{QK}^T_{i,j}) \text{ for multiple } i, j \text{ iterations}$$

ここで $s^a_{ij}$ および $s^b_{ij}$ は、すべての $i$ および $j$ .

Flash Attention全体はもう少し複雑ですが、本ガイドでは詳細に踏み込むことが範囲外となるため、大きく簡略化されています。詳細については、読み手に、よく書かれたFlash Attention論文を参照していただくことをお勧めします。

ここでの主なポイントは次のとおりです。

softmax の正規化統計を管理し、いくつかの巧妙な数学的手法を用いることで、Flash Attention は 数値的に同一の 出力を、デフォルトの自己アテンション層と比較して、メモリコストが線形に増加する程度に抑えながら提供します。 $N$ .

数式を見ると、より多くの計算が必要になるため、Flash Attention はデフォルトの自己注意機構の式よりもはるかに遅いはずだと直感的に思えます。実際、softmax の正規化統計を常に再計算する必要があり、Flash Attention は通常の注意機構よりも多くの FLOPs を必要とします（詳細については、興味がある場合は 論文 を参照してください）。

ただし、Flash Attention は、GPU のより低速で高帯域幅のメモリ（VRAM）への負荷を大幅に軽減し、代わりにより高速なオンチップメモリ（SRAM）に焦点を当てることで、デフォルトの attention と比べて推論がはるかに高速です。

Flash Attention は、出力ベクトル オンチップ SRAM メモリを使用して、より低速な VRAM メモリにアクセスせずに、すべての中間書き込みおよび読み取り操作を実行できるようにします。 $\mathbf{O}$ .

実際、Flash Attention が利用可能な場合、使用しない理由がまったくありません。このアルゴリズムは数学的に同じ出力を生成し、かつ高速でメモリ効率も高いです。

3. アーキテクチャの革新

これまで、以下により計算効率とメモリ効率の向上を検討してきました。

• 重みを低精度形式にキャストする
• 自己注意アルゴリズムを、メモリおよび計算効率の高いバージョンに置き換える

次に、長いテキスト入力を必要とするタスクに対して最も効果的かつ効率的になるよう、LLMのアーキテクチャをどのように変更できるかを検討しましょう。例:

• 検索拡張型質問応答、
• 要約、
• チャット

chatは、LLMが長いテキスト入力を処理する必要があるだけでなく、ユーザーとアシスタント間の往復対話（ChatGPTなど）を効率的に処理できることも必要とします。

学習が完了すると、LLMの基本的なアーキテクチャを変更することは困難になるため、事前にLLMのタスクについて検討し、それに応じてモデルのアーキテクチャを最適化することが重要です。 モデルのアーキテクチャには、長い入力シーケンスに対してメモリおよび/またはパフォーマンスのボトルネックに急速に発展する2つの重要なコンポーネントがあります。

• 位置埋め込み
• キーバリューキャッシュ

各コンポーネントについて、より詳しく見ていきましょう。

3.1 LLMの位置埋め込みの改善

Self-attention は、各トークンを他のトークンとの関係で捉えます。 例えば、 $\text{Softmax}(\mathbf{QK}^T)$ テキスト入力シーケンスの行列 ““Hello”, “I”, “love”, “you” は次のようになる可能性があります:

各ワードトークンは、他のすべてのワードトークンに注意を向ける確率質量を与えられ、それによって他のすべてのワードトークンとの関係が構築されます。例えば、単語“love”は、単語“Hello”に5%、“I”に30%、そして自身に65%の注意を向けています。

A LLM based on self-attention, but without position embeddings would have great difficulties in understanding the positions of the text inputs to each other. This is because the probability score computed by $\mathbf{QK}^T$ relates each word token to each other word token in $O(1)$ computations regardless of their relative positional distance to each other. Therefore, for the LLM without position embeddings each token appears to have the same distance to all other tokens, e.g. differentiating between “Hello I love you” and “You love I hello” would be very challenging.

LLMが文の順序を理解するには、追加のキューが必要であり、通常は位置エンコーディング（位置埋め込みとも呼ばれる）の形で適用されます。 位置エンコーディングは、各トークンの位置を数値表現としてエンコードし、LLMが文の順序をより正確に理解できるようにします。

論文 Attention Is All You Need の著者たちは、正弦波位置埋め込み $\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N$ . を提案しました。ここで、各ベクトル $\mathbf{p}_i$ は、その位置 $i$ . の正弦関数として計算されます。位置エンコーディングは、入力シーケンスのベクトル $\mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N$ = $\mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N$ に単純に加算されることで、モデルが文の順序をよりよく学習できるように促します。

固定の位置埋め込みを使用するのではなく、他の研究（ Devlin et al.) ）では、位置埋め込み $\mathbf{P}$ が学習中に学習される学習済み位置エンコーディングが使用されました。

正弦波位置埋め込みと学習済み位置埋め込みは、LLMに文の順序をエンコードするための主流な手法でしたが、これらの位置エンコーディングに関連するいくつかの問題が明らかになりました：

1. 正弦波位置埋め込みと学習済み位置埋め込みは、どちらも絶対位置埋め込みであり、 i.e. 各位置 ID $0, \ldots, N$ . に対して一意の埋め込みをエンコードします: Huang ら および Su ら, が示したように、絶対位置埋め込みは長いテキスト入力に対して LLM の性能を低下させます。長いテキスト入力では、トークンの絶対的な位置ではなく、トークン間の相対的な位置距離をモデルが学習することが望ましいです。
2. 学習済み位置埋め込みを使用する場合、LLM は固定の入力長 $N$ で訓練する必要があり、訓練時に使用した入力長を超える長い入力への外挿が困難になります。

最近、上記の問題に対処できる相対位置埋め込みがより広く普及しており、特に注目されているのは以下のものです。

• 回転位置埋め込み（RoPE）
• ALiBi

どちらも RoPE および ALiBi は、単語トークンが相互に結びつけられるのは自己注意アルゴリズム内であるため、文の順序に関する手がかりを LLM に与えるには、自己注意アルゴリズム内で直接行うのが最善だと主張しています。より具体的には、 $\mathbf{QK}^T$ の計算を変更することで、文の順序の手がかりを与えるべきです。

詳細には立ち入らずに、 RoPE は、位置情報をクエリとキーのペア e.g. $\mathbf{q}_i$ に $\mathbf{x}_j$ をそれぞれ角度 $\theta * i$ と $\theta * j$ で回転させることで符号化できることを述べています。ここで $i, j$ は各ベクトルの文内位置を表します。 $$\mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j.$$

$\mathbf{R}_{\theta, i - j}$ したがって、回転行列を表します。 $\theta$ は 学習 されず、代わりに、学習時の最大入力シーケンス長に依存する事前定義された値に設定されます。

これにより、 $\mathbf{q}_i$ と $\mathbf{q}_j$ 間の確率スコアは、 $i \ne j$ の場合にのみ影響を受け、 $i - j$ 各ベクトルの具体的な位置に関係なく、相対的な距離 $i$ と $j$ .

RoPEは、以下の代表的なLLMなど、現在最も重要なLLMの複数で利用されています。

• Falcon
• Llama
• PaLM

代替案として、 ALiBi は、はるかにシンプルな相対位置エンコーディング方式を提案しています。入力トークン同士の相対距離を、事前に定義された値 m でスケーリングした負の整数として、 $\mathbf{QK}^T$ 行列の各クエリ・キー要素に、softmax 計算の直前に加算します。

ALiBi論文に示されているように、このシンプルな相対位置エンコーディングにより、モデルは非常に長いテキスト入力シーケンスにおいても高い性能を維持できます。

ALiBiは、以下の代表的なLLMなど、現在最も重要な複数のLLMで使用されています。

• MPT
• BLOOM

どちらも RoPE と ALiBi の位置エンコーディングは、トレーニング時に確認しなかった入力長にも外挿できますが、 ALiBi は RoPE. に比べて、そのままの状態（out-of-the-box）で外挿がはるかにうまくいくことが示されています。 ALiBi の場合、下三角行列の位置行列の値を単に入力シーケンスの長さに合わせて増やすだけです。 RoPE, θ $\theta$ トレーニング時に使用された c.f Press et al.. しかし、コミュニティでは $\theta$ を適応させることで RoPE 位置埋め込みが外挿されたテキスト入力シーケンスでも良好に機能するようになるとするいくつかの効果的なテクニックが提案されています（ こちら).

RoPE と ALiBi はどちらも、学習中に学習されない相対位置埋め込みであり、以下の直観に基づいています。

• テキスト入力に関する位置情報に関する手がかりは、自己アテンション層の $\mathbf{QK}^T$ 行列に直接与える必要があります。
• LLMは、一定の相対距離位置エンコーディングを学習するようにインセンティブを与えられるべきです。
• テキスト入力トークン同士の距離が離れるほど、それらのクエリとバリューの確率は低下します。RoPE と ALiBi はどちらも、互いに遠く離れたトークンのクエリとキーの確率を低下させます。RoPE は、クエリとキーのベクトル間の角度を大きくすることで、それらの内積を減少させて確率を低下させます。一方、ALiBi は、内積に大きな負の数を加えることで確率を低下させます。

結論として、大規模なテキスト入力の処理を伴うタスクにデプロイすることを目的とした LLM は、RoPE や ALiBi などの相対位置埋め込みで訓練した方が望ましい。また、RoPE や ALiBi を用いた LLM が、例えば $N_1 = 2048$ のような固定長でしか訓練されていなくても、位置埋め込みを外挿することで、

3.2 キーバリューキャッシュ

LLMによる自己回帰的なテキスト生成は、入力シーケンスを投入し、次のトークンをサンプリングして入力シーケンスに追加し、LLMが生成終了を示すトークンを生成するまで、この処理を反復的に繰り返すことで動作します。

自動回帰生成の仕組みを視覚的に理解するには、Transformer’s Generate Text Tutorialをご覧ください。

自動回帰が実際どのように動作するかを示すために、簡単なコードスニペットを実行してみましょう。ここでは、torch.argmax を使用して最も確率の高い次のトークンを取得するだけです。

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("shape of input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

出力:

shape of input_ids torch.Size([1, 21])
shape of input_ids torch.Size([1, 22])
shape of input_ids torch.Size([1, 23])
shape of input_ids torch.Size([1, 24])
shape of input_ids torch.Size([1, 25])
[' Here is a Python function']

私たちが確認できるように、テキスト入力トークンは、ちょうどサンプリングされたトークンによって毎回増加します。

ごく少数の例外を除き、LLMは因果言語モデリングの目的関数を用いて学習されるため、アテンションスコアの行列の上三角部分をマスクします。そのため、上記の2つの図ではアテンションスコアが空白（つまり確率が0）になっています。因果言語モデリングの簡単な復習として、Illustrated Self Attentionブログを参照してください。

その結果、トークンは 決して 後続のトークンに依存せず、より具体的には $\mathbf{q}_i$ ベクトルは、 $\mathbf{k}_j, \mathbf{v}_j$ が $j > i$ の場合、キーやバリューのベクトル $\mathbf{q}_i$ は、以前のキー・バリューベクトル $\mathbf{k}${m < i}, \mathbf{v}${m < i} \text{ , for } m \in {0, \ldots i - 1}$ に対してのみアテンションを行います。そのため、不要な計算を削減するために、各層のキー・バリューベクトルをすべての過去のタイムステップについてキャッシュすることができます。

以下では、各フォワードパスにおいてキーバリューキャッシュを取得して転送することで、LLMにそのキャッシュを利用させる方法を説明します。 Transformersでは、forward呼び出しにuse_cacheフラグを渡すことでキーバリューキャッシュを取得でき、その後、現在のトークンとともに渡すことができます。

past_key_values = None # past_key_values is the key-value cache
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", next_token_id.shape)
  print("length of key-value cache", past_key_values.get_seq_length())  # past_key_values are of shape [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

出力:

shape of input_ids torch.Size([1, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
length of key-value cache 24
[' Here', ' is', ' a', ' Python', ' function']

ご覧のとおり、キーバリューキャッシュを使用する場合、テキスト入力トークンは長さが伸びることはなく、単一の入力ベクトルとして保持されます。一方、キーバリューキャッシュの長さは、デコーディングの各ステップで1ずつ増加します。

キーバリューキャッシュを活用することで、 $\mathbf{QK}^T$ は実質的に $\mathbf{q}_c\mathbf{K}^T$ に削減され、 $\mathbf{q}_c$ は現在渡されている入力トークンのクエリ射影であり、 常に 単一のベクトルです。

キーバリューキャッシュを使用することには、2つの利点があります。

• 計算量が削減されるため、 $\mathbf{QK}^T$ 行列の完全な計算と比較して、計算効率が大幅に向上します。これにより推論速度が向上します。
• 生成トークン数に対して、必要な最大メモリは二次的に増加するのではなく、線形にのみ増加します。

キーバリューキャッシュは常に使用すべきです。これにより、長い入力シーケンスにおいて結果が同一になるだけでなく、大幅な高速化が得られます。Transformersでは、テキストパイプラインやgenerateメソッドを使用する際、キーバリューキャッシュはデフォルトで有効になっています。キャッシュに関する完全なガイドはこちらにあります。

キーバリューキャッシュの使用を推奨していますが、使用時にLLMの出力がわずかに異なる場合があります。これは行列積カーネル自体の特性であり、詳細についてはこちらを参照してください。

3.2.1 マルチラウンド会話

キーバリューキャッシュは、チャットのように、オートレグレッシブデコーディングを複数回実行する必要があるアプリケーションに特に有用です。例を見てみましょう。

User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants

このチャットでは、LLMがオートレグレッシブデコーディングを2回実行します。

1. 初回では、キーバリューキャッシュは空であり、入力プロンプトは"User: How many people live in France?"です。モデルは、各デコーディングステップでキーバリューキャッシュを増やしながら、テキスト"Roughly 75 million people live in France"を自己回帰的に生成します。
2. 2回目の入力プロンプトは"User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?"です。キャッシュのおかげで、最初の2文のすべてのキー・ベクトルはすでに計算済みです。したがって、入力プロンプトは"User: And how many in Germany?"のみで構成されます。短縮された入力プロンプトを処理する際、その計算されたキー・ベクトルは、最初のデコーディングのキー・ベクトル・キャッシュに連結されます。2番目のAssistantの回答"Germany has ca. 81 million inhabitants"は、"User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?"のエンコードされたキー・ベクトルで構成されるキー・ベクトル・キャッシュを用いて、オートレグレッシブに生成されます。

ここでは2つの点に注意が必要です。

1. チャットにデプロイされたLLMでは、LLMが会話の過去のすべてのコンテキストを理解できるようにするため、すべてのコンテキストを保持することが重要です。例えば、上記の例では、ユーザーが"And how many are in Germany"と質問したとき、ドイツの人口について尋ねていることをLLMが理解する必要があります。
2. キーバリューキャッシュはチャットにおいて非常に有用であり、エンコーダーデコーダーアーキテクチャのように毎回ゼロからチャット履歴を再エンコードする必要がなく、エンコード済みのチャット履歴を継続的に拡張できるためです。

transformersでは、return_dict_in_generate=Trueを渡すと、デフォルトのuse_cache=Trueに加えて、generateの呼び出しがpast_key_valuesを返します。なお、pipelineインターフェースではまだ利用できません。

# Generation as usual
prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True)
decoded_output = tokenizer.batch_decode(generation_output.sequences)[0]

# Piping the returned `past_key_values` to speed up the next conversation round
prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(
  **model_inputs,
  past_key_values=generation_output.past_key_values,
  max_new_tokens=60,
  return_dict_in_generate=True
)
tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

出力:

 is a modified version of the function that returns Mega bytes instead.

def bytes_to_megabytes(bytes):
   return bytes / 1024 / 1024

Answer: The function takes a number of bytes as input and returns the number of

素晴らしいですね。アテンション層で同じキーとバリューを再計算する追加の時間はかかりません! ただし、1つの注意点があります。 $\mathbf{QK}^T$ 行列のピークメモリ量は大幅に削減されますが、キーバリューキャッシュをメモリに保持することは、長い入力シーケンスやマルチターンチャットの場合、非常にメモリコストが高くなることがあります。キーバリューキャッシュは、すべての過去の入力ベクトル $\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$ すべての自己アテンション層およびすべてのアテンションヘッドについて、キーとバリューのベクトルを保存する必要があります。

以前使用した LLM bigcode/octocoder のキーバリューキャッシュに保存される必要がある float 値の数を計算してみましょう。 float 値の数は、シーケンス長の 2 倍にアテンションヘッド数、アテンションヘッド次元数、レイヤー数をそれぞれ掛けたものになります。 仮想的な入力シーケンス長 16000 でこの LLM について計算すると、以下のようになります。

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

出力:

7864320000

約80億個のfloat値です!float16精度で80億個のfloat値を保存するには、約15 GBのRAMが必要であり、これはモデルの重み自体の約半分です! 研究者たちは、キーバリューキャッシュの保存にかかるメモリコストを大幅に削減できる2つの手法を提案しており、これらは次の小節で検討されます。

3.2.2 マルチクエリアテンション（MQA）

Multi-Query-Attentionは、Noam Shazeer氏の論文 Fast Transformer Decoding: One Write-Head is All You Need で提案されました。タイトルが示す通り、Noam氏は、n_head 個の key-value 射影の重みを使用する代わりに、すべての attention head で共有される単一の head-value 射影の重みペアを使用しても、モデルの性能が有意に低下しないことを発見しました。

単一のヘッド・バリュー射重みペアを使用することで、キー・バリューベクトル $\mathbf{k}_i, \mathbf{v}_i$ はすべてのアテンションヘッド間で同一である必要があるため、キャッシュには n_head 個のキー・バリュー射重みペアの代わりに1組のみを保存すればよいことになります。

ほとんどのLLMは20から100のattention headsを使用しているため、MQAはkey-value cacheのメモリ消費量を大幅に削減します。このノートブックで使用しているLLMの場合、入力シーケンス長が16000のとき、必要なメモリ消費量を15 GBから400 MB未満に削減できます。

メモリ削減に加えて、MQAは計算効率の向上にもつながります。その理由は以下の通りです。自己回帰デコーディングでは、大きなキー・ベクトルとバリュー・ベクトルを再読み込みし、現在のキー・ベクトルとバリュー・ベクトルのペアと連結した上で、 $\mathbf{q}_c\mathbf{K}^T$ の計算に与える必要があります。この再読み込みは各ステップで必要であり、自己回帰デコーディングにおいて、この一定の再読み込みに必要なメモリ帯域幅が深刻な時間ボトルネックになり得ます。キー・ベクトルとバリュー・ベクトルのサイズを小さくすることで、アクセスするメモリ量が減り、メモリ帯域幅のボトルネックが軽減されます。詳細については、 Noam の論文.

ここで理解すべき重要な点は、キーバリューアテンションヘッドの数を1に削減するのは、キーバリューキャッシュを使用する場合にのみ意味があるということです。キーバリューキャッシュなしで1回のフォワードパスを行う場合のモデルのピークメモリ消費量は、各アテンションヘッドが固有のクエリベクトルを持つため、各アテンションヘッドが異なる $\mathbf{QK}^T$ 行列を持つため、変わりません。

MQAはコミュニティで広く採用されており、現在では最も人気のあるLLMの多くで使用されています。

• Falcon
• PaLM
• MPT
• BLOOM

また、このノートブックで使用されているチェックポイント bigcode/octocoder は、MQA を利用しています。

3.2.3 グループクエリアテンション（GQA）

Google の Ainslie らが提案した Grouped-Query-Attention では、通常のマルチキー・バリューヘッド投影と比較して、MQA を使用すると品質が低下することがあることが示されています。この論文では、クエリヘッド投影重みの数をより穏やかに減らすことで、モデルの性能をより維持できると主張しています。単一のキー・バリュー投影重みを使用するのではなく、n < n_head 個のキー・バリュー投影重みを使用すべきです。n を n_head よりもはるかに小さい値、例えば 2、4、8 などに設定することで、MQA によるメモリと速度のほぼすべての利点を維持しつつ、モデル容量の犠牲を減らし、結果として性能の低下も抑えられると考えられます。

さらに、GQAの著者たちは、既存のモデルチェックポイントを、元の事前学習計算量のわずか5%でGQAアーキテクチャを持つようにアップトレーニングできることを発見しました。元の事前学習計算量の5%であっても依然として膨大な量になり得ますが、GQAのアップトレーニングにより、既存のチェックポイントをより長い入力シーケンスに活用できるようになります。

GQAは比較的最近に提案された手法であり、このノートブックの執筆時点では採用例がまだ少ないのが現状です。 GQAの最も代表的な応用例はLlama-v2です。

結論として、LLMが自動回帰デコーディングでデプロイされ、チャットのように長い入力シーケンスの処理が求められる場合、GQAまたはMQAの利用を強く推奨します。

結論

研究コミュニティでは、ますます大型化する LLM の推論時間を高速化するための新しい工夫が絶えず提案されています。その有望な研究方向性の一例として、speculative decodingがあります。これは、より小型で高速な言語モデルによって「簡単なトークン」を生成し、LLM 本体が生成するのは「難しいトークン」だけにするという手法です。詳細な説明は本ノートブックの範囲外ですが、このブログ記事で詳しく読むことができます。

GPT3/4、Llama-2-70b、Claude、PaLM といった大規模 LLM が、Hugging Face Chat や ChatGPT などのチャットインターフェースでこれほど高速に動作できるのは、前述した精度、アルゴリズム、アーキテクチャの改善が大きく寄与しているためです。 今後、GPU や TPU などのアクセラレータはさらに高速化し、より多くのメモリを扱えるようになるでしょうが、それでも常に最善のアルゴリズムとアーキテクチャを使用し、コストパフォーマンスを最大化するように心がけてください 🤗