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

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

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

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

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

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

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

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

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

## 1. 低精度

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

本ガイドの執筆時点では、LLMは少なくとも数十億パラメータから構成されています。各パラメータは小数で表され、例えば`4.5689`のように、通常は[float32](https://en.wikipedia.org/wiki/Single-precision_floating-point_format)、[bfloat16](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format)、または[float16](https://en.wikipedia.org/wiki/Half-precision_floating-point_format)形式で保存されます。これにより、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**](https://huggingface.co/bigscience/bloom) には、2 \* 176 GB = **352 GB** の VRAM が必要です
- [**Llama-2-70b**](https://huggingface.co/meta-llama/Llama-2-70b-hf) には、2 \* 70 GB = **140 GB** の VRAM が必要です
- [**Falcon-40b**](https://huggingface.co/tiiuae/falcon-40b) には、2 \* 40 GB = **80 GB** の VRAM が必要です
- [**MPT-30b**](https://huggingface.co/mosaicml/mpt-30b) には、2 \* 30 GB = **60 GB** の VRAM が必要です
- [**bigcode/starcoder**](https://huggingface.co/bigcode/starcoder) には、2 \* 15.5 = **31 GB** の VRAM が必要です

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

🤗 Transformers は、それぞれの config クラスに `base_tp_plan` を持つサポート対象モデルのテンソル並列処理に対応しました。Tensor Parallelism については [こちら](./perf_train_gpu_many#tensor-parallelism)をご覧ください。また、テンソル並列処理に友好的な方法でモデルを実装することにご興味がある場合は、[text-generation-inference ライブラリ](https://github.com/huggingface/text-generation-inference/tree/main/server/text_generation_server/models/custom_modeling)もぜひご覧ください。

Naiveなパイプライン並列化は、そのままサポートされています。これを行うには、`device="auto"`を指定してモデルを読み込むだけで、[こちら](https://huggingface.co/docs/accelerate/v0.22.0/en/concept_guides/big_model_inference)で説明されているように、異なるレイヤーが利用可能なGPUに自動的に配置されます。ただし、非常に効果的であるものの、このNaiveなパイプライン並列化では、GPUのアイドル状態の問題に対処していません。この問題に対処するには、[こちら](./perf_train_gpu_many#pipeline-parallelism)で説明されているより高度なパイプライン並列化が必要です。

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

```bash
!pip install transformers accelerate bitsandbytes optimum
```

```python
from transformers import AutoModelForCausalLM

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

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

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

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

まずモデルとトークナイザーを読み込み、両方をTransformersの[pipeline](./main_classes/pipelines)オブジェクトに渡します。

```python
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)
```

```python
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
```

**出力**:

```text
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
```

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

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

ピークアクセラレータメモリ割り当てを測定するために、[`torch.accelerator.memory.max_memory_allocated`](https://docs.pytorch.org/docs/main/generated/torch.accelerator.memory.max_memory_allocated.html)を呼び出します。

```python
bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())
```

**出力**:

```text
29.0260648727417
```

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

> 現在、ほぼすべてのモデルは bfloat16 で学習されています。[GPU が bfloat16 をサポートしている](https://discuss.pytorch.org/t/bfloat16-native-support/117155/5)場合、モデルを float32 の完全な精度で実行する理由はありません。float32 は、モデルの学習に使用された精度よりも良い推論結果をもたらすことはありません。

Hub 上でモデルの重みがどの形式で保存されているか不明な場合は、チェックポイントの config の `"dtype"` を確認できます。*例*として、[こちら](https://huggingface.co/meta-llama/Llama-2-7b-hf/blob/6fdf2e60f86ff2481f2241aaee459f85b5b0bbb9/config.json#L21)を参照してください。`from_pretrained(..., dtype=...)` で読み込む際は、config に記載されているのと同じ精度タイプにモデルを設定することをお勧めします。ただし、元のタイプが float32 の場合は、推論に `float16` または `bfloat16` のいずれかを使用できます。

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

```python
del pipe
del model

import gc
import torch

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

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

```python
flush()
```

Accelerate ライブラリでは、XPU、MLU、NPU、MPS などさまざまなハードウェアバックエンドを考慮した、デバイス非依存のユーティリティメソッド [release_memory](https://github.com/huggingface/accelerate/blob/29be4788629b772a3b722076e433b5b3b5c85da3/src/accelerate/utils/memory.py#L63) も利用できます。

```python
from accelerate.utils import release_memory
# ...

release_memory(model)
```

では、GPU の VRAM が 32 GB 未満の場合はどうでしょうか。モデルの重みを 8 ビットまたは 4 ビットに量子化しても、性能の大きな低下なしに運用できることが確認されています（[Dettmers et al.](https://huggingface.co/papers/2208.07339)を参照）。最近の [GPTQ paper](https://huggingface.co/papers/2210.17323) 🤯 に示されているように、モデルは 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`](https://github.com/bitsandbytes-foundation/bitsandbytes) ライブラリがインストールされていることを確認する必要があります。

```bash
!pip install bitsandbytes
```

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

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

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

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

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

**出力**:

```text
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
```

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

```python
bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())
```

**出力**:

```text
15.219234466552734
```

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

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

```python
del model
del pipe
```

```python
flush()
```

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

```python
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
```

**出力**:

```text
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`がありません。メモリ使用量を確認してみましょう。

```python
bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())
```

**出力**:

```text
9.543574333190918
```

わずか9.5GBです！150億パラメータを超えるモデルとしては、本当に少ない量です。

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

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

```python
del model
del pipe
```

```python
flush()
```

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

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

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

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

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

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

## 2. Flash Attention

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

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

詳しく見てみましょう。長さ $N$ の入力 $\mathbf{X}$ に対する self-attention レイヤーの出力 $\mathbf{O}$ を計算する式は以下の通りです。

$$ \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$ の場合は 19 GB、$N=100,000$ の場合は $\mathbf{QK}^T$ 行列を保存するだけでもほぼ 1TB の VRAM が必要になります。

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

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

大規模な入力長における莫大なメモリ要件をどうすれば解消できるでしょうか。$\mathbf{QK}^T$ 行列を不要にする、新しい自己注意メカニズムの計算方法が必要です。[Tri Dao ら](https://huggingface.co/papers/2205.14135)は、まさにそのような新しいアルゴリズムを開発し、**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$ に対して再計算が必要な softmax 正規化の統計量です。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

![](/blog/assets/163_optimize_llm/self_attn_tokens.png)

各ワードトークンは、他のすべてのワードトークンに注意を向ける確率質量を与えられ、それによって他のすべてのワードトークンとの関係に置かれます。例えば、単語 *"love"* は、単語 *"Hello"* に 5%、*"I"* に 30%、そして自分自身に 65% の確率で注意を向けます。

自己注意機構に基づく LLM は、位置埋め込みがない場合、テキスト入力の各トークン同士の位置関係を理解するのに大きな困難に直面します。これは、$\mathbf{QK}^T$ で計算される確率スコアが、トークン間の相対的な位置距離に関係なく、各単語トークンを他の各単語トークンと $O(1)$ の計算で関連付けるためです。したがって、位置埋め込みのない LLM では、各トークンが他のすべてのトークンに対して同じ距離にあるように見えます。*例えば*、*"Hello I love you"* と *"You love I hello"* を区別することは非常に困難になります。

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

[*Attention Is All You Need*](https://huggingface.co/papers/1706.03762) の著者たちは、正弦波位置埋め込み $\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.](https://huggingface.co/papers/1810.04805) らは、学習中の訓練で位置埋め込み $\mathbf{P}$ を学習する学習済み位置エンコーディングを使用しました。

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

  1. 正弦波位置埋め込みと学習済み位置埋め込みは、いずれも絶対位置埋め込みであり、*つまり*各位置 ID ($0, \ldots, N$) に対して固有の埋め込みをエンコードします。[Huang ら](https://huggingface.co/papers/2009.13658)および[Su ら](https://huggingface.co/papers/2104.09864)の研究が示すように、絶対位置埋め込みは長いテキスト入力において LLM の性能を低下させる原因となります。長いテキスト入力では、トークンの絶対的な位置ではなく、トークン間の相対的な位置距離をモデルが学習することが望ましいです。
  2. 学習済み位置埋め込みを使用する場合、LLM は固定の入力長 $N$ で訓練する必要があり、訓練時に使用した入力長を超える長い入力への外挿が困難になります。

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

- [回転位置埋め込み（RoPE）](https://huggingface.co/papers/2104.09864)
- [ALiBi](https://huggingface.co/papers/2108.12409)

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

詳細には立ち入らずに述べると、*RoPE*では、位置情報をクエリとキーのペア、*例えば* $\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**](https://huggingface.co/tiiuae/falcon-40b)
- [**Llama**](https://huggingface.co/papers/2302.13971)
- [**PaLM**](https://huggingface.co/papers/2204.02311)

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

![](/blog/assets/163_optimize_llm/alibi.png)

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

*ALiBi*は、以下の代表的なLLMをはじめ、現在最も重要なLLMの多くで利用されています。

- [**MPT**](https://huggingface.co/mosaicml/mpt-30b)
- [**BLOOM**](https://huggingface.co/bigscience/bloom)

*RoPE* と *ALiBi* の両方の位置エンコーディングは、学習時に確認しなかった入力長にも外挿できますが、*ALiBi* は *RoPE* と比較して、そのままの状態（out-of-the-box）で外挿がはるかにうまく機能することが示されています。ALiBi の場合、下三角行列の位置行列の値を単に入力シーケンスの長さに合わせて増加させるだけです。一方、*RoPE* では、学習時に使用された $\theta$ をそのまま維持すると、学習時に確認したテキストよりもはるかに長いテキスト入力を渡した際に結果が悪くなります（*c.f* [Press et al.](https://huggingface.co/papers/2108.12409)）。ただし、コミュニティでは $\theta$ を適応させるいくつかの効果的なテクニックが提案されており、これにより *RoPE* の位置埋め込みが外挿されたテキスト入力シーケンスに対して良好に機能するようになります（[こちら](https://github.com/huggingface/transformers/pull/24653)を参照）。

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

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

結論として、大規模なテキスト入力の処理を必要とするタスクにデプロイされることを想定した LLM は、RoPE や ALiBi などの相対位置埋め込みを用いて学習させるのが望ましい。また、RoPE や ALiBi を用いた LLM が、例えば $N_1 = 2048$ のような固定長でしか学習されていなくても、位置埋め込みを外挿することで、$N_2 = 8192 > N_1$ のように $N_1$ よりもはるかに大きなテキスト入力を扱う実用上の利用が可能である点に注意してください。

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

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

自動回帰生成の仕組みを視覚的に理解するには、[Transformer's Generate Text Tutorial](./llm_tutorial#default-generate)をご覧ください。

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

```python
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
```

**出力**:

```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は[因果言語モデリングの目的関数](./tasks/language_modeling#causal-language-modeling)を用いて学習されるため、アテンションスコアの行列の上三角部分をマスクします。そのため、上記の2つの図ではアテンションスコアが空白（*つまり確率が0*）になっています。因果言語モデリングの簡単な復習として、[*Illustrated Self Attentionブログ*](https://jalammar.github.io/illustrated-gpt2/#part-2-illustrated-self-attention)を参照してください。

As a consequence, tokens *never* depend on later tokens, more specifically the $\mathbf{q}_i$ vector is never put in relation with any key, values vectors $\mathbf{k}_j, \mathbf{v}_j$ if $j > i$ . Instead $\mathbf{q}_i$ only attends to previous key-value vectors $\mathbf{k}_{m < i}, \mathbf{v}_{m < i} \text{ , for } m \in \{0, \ldots i - 1\}$. In order to reduce unnecessary computation, one can therefore cache each layer's key-value vectors for all previous timesteps.

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

```python
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
```

**出力**:

```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つの利点があります。

- QK^T の全行列を計算する場合と比べて計算量が少なくなるため、計算効率が大幅に向上します。これにより、推論速度が向上します。
- 生成トークン数に対して、必要な最大メモリは二次的に増加するのではなく、線形にのみ増加します。

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

キーバリューキャッシュの使用を推奨していますが、使用する場合にLLMの出力がわずかに異なることがあります。これは行列積カーネル自体の特性であり、詳細については[こちら](https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535)を参照してください。

#### 3.2.1 マルチラウンド会話

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

```text
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`インターフェースではまだ利用できません。

```python
# 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):]
```

**出力**:

```text
 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
```

Great, no additional time is spent recomputing the same key and values for the attention layer! There is however one catch. While the required peak memory for the $\mathbf{QK}^T$ matrix is significantly reduced, holding the key-value cache in memory can become very memory expensive for long input sequences or multi-turn chat. Remember that the key-value cache needs to store the key-value vectors for all previous input vectors $\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$ for all self-attention layers and for all attention heads.

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

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

**出力**:

```text
7864320000
```

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

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

[Multi-Query-Attention](https://huggingface.co/papers/1911.02150) は、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の論文](https://huggingface.co/papers/1911.02150)をご覧ください。

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

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

- [**Falcon**](https://huggingface.co/tiiuae/falcon-40b)
- [**PaLM**](https://huggingface.co/papers/2204.02311)
- [**MPT**](https://huggingface.co/mosaicml/mpt-30b)
- [**BLOOM**](https://huggingface.co/bigscience/bloom)

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

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

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

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

GQAは比較的最近に提案された手法であり、このノートブックの執筆時点では採用例がまだ少ないのが現状です。GQAの最も代表的な応用例は[Llama-v2](https://huggingface.co/meta-llama/Llama-2-70b-hf)です。

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

## 結論

研究コミュニティでは、ますます大型化する LLM の推論時間を高速化する、新しいかつ実用的な手法が次々と提案されています。その有望な研究方向性の一例として、[speculative decoding](https://huggingface.co/papers/2211.17192)があります。これは、より小型で高速な言語モデルによって「簡単なトークン」を生成し、LLM 本体が「難しいトークン」のみを生成する方式です。詳細な説明は本ノートブックの範囲外ですが、[このブログ記事](https://huggingface.co/blog/assisted-generation)で詳しく読むことができます。

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