「プロダクションでのあなたのLLMの最適化」

Optimizing your LLM in production

注意: このブログ投稿は、Transformersのドキュメンテーションページとしても利用可能です。

GPT3/4、Falcon、LLamaなどの大規模言語モデル（LLM）は、人間中心のタスクに取り組む能力を急速に向上させており、現代の知識ベース産業で不可欠なツールとして確立しています。しかし、これらのモデルを実世界のタスクに展開することは依然として課題が残っています：

• ほぼ人間のテキスト理解と生成能力を持つために、LLMは現在数十億のパラメータから構成される必要があります（Kaplanら、Weiら参照）。これにより、推論時のメモリ要件が増大します。
• 多くの実世界のタスクでは、LLMには豊富な文脈情報が必要です。これにより、推論中に非常に長い入力シーケンスを処理する能力が求められます。

これらの課題の核心は、特に広範な入力シーケンスを扱う場合に、LLMの計算およびメモリ能力を拡張することにあります。

• 「LLMの利点：電子商取引の検索を変革する」
• 「TensorFlowを使用した異常検出のための変分オートエンコーダのトレーニング」
• 「新しい研究は、AIイノベーションのモデルとして生態学を示唆しています」

このブログ投稿では、効率的なLLMの展開のために、現時点で最も効果的な技術について説明します：

1. 低精度: 研究により、8ビットおよび4ビットの数値精度で動作することが、モデルのパフォーマンスに大幅な低下を伴わずに計算上の利点をもたらすことが示されています。

2. Flash Attention: Flash Attentionは、よりメモリ効率の高いアテンションアルゴリズムのバリエーションであり、最適化されたGPUメモリの利用により、高い効率を実現します。

3. アーキテクチャのイノベーション: LLMは常に同じ方法で展開されるため、つまり長い入力コンテキストを持つ自己回帰的なテキスト生成として、より効率的な推論を可能にする専用のモデルアーキテクチャが提案されています。モデルアーキテクチャの中で最も重要な進歩は、Alibi、Rotary embeddings、Multi-Query Attention（MQA）、Grouped-Query-Attention（GQA）です。

このノートブックでは、テンソルの視点から自己回帰的な生成の分析を提供し、低精度の採用の利点と欠点について包括的な探索を行い、最新のアテンションアルゴリズムの詳細な調査を行い、改良されたLLMアーキテクチャについて議論します。これを行う過程で、各機能の改善を示す実用的な例を実行します。

1. 低精度の活用

LLMのメモリ要件は、LLMを重み行列とベクトルのセット、およびテキスト入力をベクトルのシーケンスとして見ることで最も理解できます。以下では、重みの定義はすべてのモデルの重み行列とベクトルを意味するために使用されます。

この投稿の執筆時点では、LLMは少なくとも数十億のパラメータから構成されています。各パラメータは通常、float32、bfloat16、またはfloat16形式で保存される10進数の数値で構成されています。これにより、LLMをメモリにロードするためのメモリ要件を簡単に計算できます：

X十億のパラメータを持つモデルの重みをロードするには、おおよそ4 * X GBのVRAMがfloat32精度で必要です

現在では、モデルはほとんどが完全なfloat32精度ではなく、通常はbfloat16精度、またはよりまれにfloat16精度でトレーニングされています。したがって、ルールオブサムは次のようになります：

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

テキスト入力が短い場合（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チップはA100で、80GBのVRAMを提供しています。80GB以上のモデルを読み込むためには、必然的にテンソル並列処理と/またはパイプライン並列処理が必要です。

🤗 Transformersは、モデルのアーキテクチャが特定の方法で書かれていることを必要とするため、テンソル並列処理をそのままサポートしていません。テンソル並列処理に対応したモデルを作成することに興味がある場合は、text-generation-inferenceライブラリをご覧ください。

素朴なパイプライン並列処理はそのままサポートされています。これには、device="auto"としてモデルを読み込むだけで十分です。これにより、利用可能なGPU上にさまざまなレイヤーが自動的に配置されます。ただし、非常に効果的ですが、この素朴なパイプライン並列処理では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に均等に分散されます。

このノートブックでは、単一の40GBのA100 GPUデバイスチップで実行できるbigcode/octocoderを使用します。今後適用するすべてのメモリと速度の最適化は、モデルやテンソル並列処理を必要とするモデルにも同様に適用できることに注意してください。

モデルがbfloat16精度でロードされているため、上記のおおよその目安に従うと、bigcode/octocoderで推論を実行するためのメモリ要件は約31GBのVRAMとなります。試してみましょう。

まず、モデルとトークナイザを読み込み、それらをTransformersのパイプラインオブジェクトに渡します。

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", torch_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.cuda.max_memory_allocated()を呼び出してピークのGPUメモリ割り当てを測定しましょう。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

出力:

29.0260648727417

おおよそ当初の推定と一致しています！数値は正確ではないことがわかりますが、バイトからキロバイトに変換するには1000ではなく1024で乗算する必要があります。したがって、おおよその計算式は「最大でもX GB」として理解することもできます。なお、モデルを完全なfloat32精度で実行しようとした場合、64GBのVRAMが必要でした。

ほとんどのモデルは現在、bfloat16でトレーニングされています。GPUがbfloat16をサポートしている場合、モデルを完全なfloat32精度で実行する理由はありません。Float32は、モデルのトレーニングに使用された精度よりも優れた推論結果を提供しません。

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

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

del pipe
del model

import gc
import torch

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

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

flush()

最近のバージョンのaccelerateライブラリでは、release_memory()というユーティリティメソッドも使用できます。

from accelerate.utils import release_memory
# ...

release_memory(model)

では、もしGPUが32GBのVRAMを持っていない場合はどうなるでしょうか？モデルの重みは8ビットまたは4ビットに量子化することができ、パフォーマンスの大幅な低下はありません（Dettmers et al.を参照）。最近のGPTQ論文では、モデルは3ビットや2ビットにまで量子化でき、受け入れ可能なパフォーマンスの低下が示されています 🤯。

詳細には触れませんが、量子化スキームは、重みの精度を下げつつ、モデルの推論結果を可能な限り正確に保つことを目指しています（つまり、bfloat16にできるだけ近い結果になるようにします）。テキスト生成では、次のトークンのセットを最も確率が高いものから選ぶことだけを気にするため、次のトークンのロジット分布の正確な値はあまり重要ではありません。重要なのは、次のトークンのロジット分布がほぼ同じままであり、argmaxやtopkの操作が同じ結果を返すことです。

量子化にはさまざまな技術がありますが、ここでは詳細には触れません。一般的に、量子化の手法は以下のように機能します：

• 1. すべての重みを目標精度に量子化する
• 1. 量子化された重みをロードし、入力ベクトルのシーケンスをbfloat16精度で渡す
• 1. 計算を行うために重みを動的にbfloat16に逆量子化する
• 1. 計算後に入力とともに重みを再度目標精度に量子化する

要するに、入力-重み行列の乗算は、Xが入力、Wが重み行列、Yが出力の場合：

Y=X∗W Y = X * W Y=X∗W

次のように変更されます：

Y=X∗dequantize(W);quantize(W) Y = X * \text{dequantize}(W); \text{quantize}(W) Y=X∗dequantize(W);quantize(W)

すべての行列乗算のたびに、逆量子化と再量子化が重み行列に対して順次行われます。

そのため、量子化された重みを使用する場合、推論時間は通常短縮されず、むしろ増加します。理論はこれぐらいにして、試してみましょう！Transformersを使用して重みを量子化するには、bitsandbytesライブラリがインストールされていることを確認する必要があります。

# !pip install bitsandbytes

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

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", 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.cuda.max_memory_allocated())

出力:

15.219234466552734

かなり少ないです！15GB以上で、4090のような一般のGPUでこのモデルを実行することができます。メモリ効率が非常に向上し、モデルの出力にほとんど劣化が見られません。ただし、推論中にわずかな遅延が見られることもわかります。

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

del model
del pipe

flush()

4ビット量子化でピークのGPUメモリ使用量がどの程度になるかを確認しましょう。モデルを4ビットに量子化するには、前と同じAPIを使用して、load_in_8bit=Trueの代わりにload_in_4bit=Trueを渡します。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True, low_cpu_mem_usage=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

出力:

以下は、バイトをギガバイトに変換するPythonの関数です。

```
def bytes_to_gigabytes(bytes):
    return bytes / 1024 / 1024 / 1024
```

この関数は単一の引数を受け取ります。

ほとんど先ほどと同じ出力テキストが表示されていますが、コードスニペットの前にpythonが抜けています。メモリの必要量はどれくらいだったか見てみましょう。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

出力:

9.543574333190918

たった9.5GBです！これは、150億以上のパラメータを持つモデルにとってはほとんどありません。

ここではモデルの精度の劣化はほとんど見られませんが、4ビット量子化は8ビット量子化や完全なbfloat16推論と比べて異なる結果になることが多いです。ユーザー自身で試してみる必要があります。

また、4ビット量子化では推論中の処理時間が8ビット量子化よりも再度遅くなることに注意してください。これは、4ビット量子化によりより積極的な量子化方法が使用され、推論中の量子化と非量子化の処理が長くなるためです。

del model
del pipe

flush()

全体的に、OctoCoderを8ビット精度で実行することで、必要なGPU VRAMが32GBからわずか15GBに削減され、4ビット精度でモデルを実行することで必要なGPU VRAMがわずか9GB以上にさらに削減されました。

4ビット量子化により、RTX3090、V100、T4などのGPUでモデルを実行することができます。これらのGPUはほとんどの人にとって利用しやすいです。

さらにGPU VRAMメモリを4ビットよりもさらに少なくするモデル量子化についての詳細な情報や、AutoGPTQの実装については、Transformersの量子化ドキュメントをご覧ください。

結論として、モデルの量子化はメモリ効率の向上と引き換えに精度および推論時間に影響を与えることを覚えておくことが重要です。

GPUメモリが使用制限にない場合、量子化を考慮する必要はない場合もあります。ただし、多くのGPUは量子化手法を使用せずにLLMを実行できないため、4ビットおよび8ビット量子化スキームは非常に便利なツールです。

より詳細な使用方法については、Transformers Quantization Docsを強くおすすめします。次に、より良いアルゴリズムと改良されたモデルアーキテクチャを使用して計算およびメモリ効率を改善する方法について見ていきましょう。

現在のトップパフォーマンスのLLMは、基本的なアーキテクチャがほぼ同じであり、フィードフォワードレイヤー、アクティベーションレイヤー、レイヤーノーマライゼーションレイヤー、そして最も重要なのはセルフアテンションレイヤーから構成されています。

セルフアテンションレイヤーは、大規模言語モデル（LLM）において入力トークン間の文脈関係を理解することができるようにするため、LLMの中心的な役割を果たしています。ただし、セルフアテンションレイヤーのピークGPUメモリ使用量は、入力トークンの数（シーケンス長とも呼ばれる）に対して計算およびメモリの複雑さが二次的に増加します。短い入力シーケンス（最大1000個の入力トークンまで）ではこれはほとんど気になりませんが、長い入力シーケンス（約16000個の入力トークンで）では深刻な問題となります。

より詳しく見てみましょう。長さ N の入力 X の自己注意層の出力 O \mathbf{O} O を計算するための式は次の通りです：

O=Attn(X)=V×Softmax(QKT) with Q=WqX,V=WvX,K=WkX \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} O=Attn(X)=V×Softmax(QKT) with Q=Wq​X,V=Wv​X,K=Wk​X mathbfX=(x1,…xN) mathbf{X} = (\mathbf{x}_1, … \mathbf{x}_{N}) mathbfX=(x1​,…xN​) は注意層への入力シーケンスです。射影 Q \mathbf{Q} Q と K \mathbf{K} K はそれぞれ N N N のベクトルで構成され、したがって QKT \mathbf{QK}^T QKT のサイズは N2 N^2 N2 となります。

LLMs は通常、複数の自己注意ヘッドを持ち、複数の自己注意計算を並列で実行します。LLM が 40 の自己注意ヘッドを持ち、bfloat16 精度で実行される場合、QKT \mathbf{QK^T} QKT 行列を保存するためのメモリ要件は 40∗2∗N2 40 * 2 * N^2 40∗2∗N2 バイトとなります。N=1000 の場合、約 50 MB の VRAM が必要ですが、N=16000 の場合は 19 GB の VRAM が必要であり、N=100,000 の場合は QKT \mathbf{QK}^T QKT 行列を保存するためにほぼ 1TB の VRAM が必要です。

要するに、大きな入力コンテキストに対しては、デフォルトの自己注意アルゴリズムはすぐにメモリの使用量が膨大になります。

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

大きな入力長に対する過剰なメモリ要件をどのように克服できるでしょうか？ QKT QK^T QKT 行列をなくすための自己注意メカニズムを計算する新しい方法が必要です。Tri Dao らは、まさにそのような新しいアルゴリズムを開発し、それを Flash Attention と呼びました。

要するに、Flash Attention は V×Softmax(QKT\mathbf{V} \times \text{Softmax}(\mathbf{QK}^TV×Softmax(QKT) の計算を分割し、出力のより小さなチャンクを反復して計算することで、次のような式を得ます：

Oi←sija∗Oi+sijb∗Vj×Softmax(QKi,jT) for multiple i,j iterations \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} Oi​←sija​∗Oi​+sijb​∗Vj​×Softmax(QKi,jT​) for multiple i,j iterations

ただし、sija s^a_{ij} sija​ および sijb s^b_{ij} sijb​ は i i i および j j j ごとに再計算する必要のあるソフトマックスの正規化統計量です。

Flash Attention 全体は少し複雑であり、ここでは詳細には立ち入りません。詳細については、よく書かれた Flash Attention 論文を参照してください。

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

ソフトマックスの正規化統計量を追跡し、いくつかのスマートな数学を使用することで、Flash Attention はデフォルトの自己注意層と 数値的に同等の 出力を提供しますが、メモリのコストは N N N と線形に増加するだけです。

式を見ると、通常のセルフアテンションの式に比べて、Flash Attentionは遅くなるはずだと直感的に言えます。なぜなら、より多くの計算が必要だからです。実際に、Flash Attentionは通常のアテンションに比べてより多くのFLOPsが必要とされます。なぜなら、softmax正規化の統計を常に再計算する必要があるからです（詳細については論文を参照してください）

しかし、Flash Attentionは、GPUの遅い高帯域幅のメモリ（VRAM）の負荷を大幅に減らすことができるため、デフォルトのアテンションに比べて推論時にははるかに高速です。代わりに、より速いオンチップメモリ（SRAM）に焦点を当てています。

基本的に、Flash Attentionは中間の書き込みおよび読み取り操作を、遅いVRAMメモリにアクセスする必要なく、速いオンチップSRAMメモリを使用して行うことを確認します。

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

具体的な例を見てみましょう。

OctoCoderモデルは、ユーザーのタスクに合わせたより良いアシスタントにするためのシステムプロンプトを含む、かなり長い入力プロンプトを受け取ります。システムプロンプトは、LLMを指示するために使用されます。以下では、OctoCoderをより優れたコーディングアシスタントにするためのシステムプロンプトを使用します。

system_prompt = """以下は、さまざまな人々とAI技術アシスタントとの対話のシリーズです。
アシスタントは、親切で、礼儀正しく、正直で、洗練され、感情に気づき、謙虚で知識豊富でありたいとします。
アシスタントは、コードの質問に対して助けになることを喜んでし、必要なことを正確に理解するために最善を尽くします。
また、間違った情報や誤解を与えることを避け、正しい答えについて完全に確信が持てない場合は注意を喚起します。
とは言っても、アシスタントは実用的で最善を尽くし、注意が役立つことを妨げすぎません。

Starcoderモデルは、The Stack（v1.2）（オプトアウトリクエストを除く）から80以上のプログラミング言語でトレーニングされた、155億パラメータモデルのシリーズです。
モデルはマルチクエリアテンションを使用し、Fill-in-the-Middle目的でトレーニングされ、大量のデータの1兆トークンに対して8,192トークンのコンテキストウィンドウを使用しました。

-----

質問：2つのリストを受け取り、入力リストから交互に要素を持つリストを返す関数を書いてください。

回答：もちろんです。それを行う関数は次のとおりです。

def alternating(list1, list2):
   results = []
   for i in range(len(list1)):
       results.append(list1[i])
       results.append(list2[i])
   return results

質問：この関数のテストケースを書いてもらえますか？

回答：もちろんです。いくつかのテストは次のとおりです。

assert alternating([10, 20, 30], [1, 2, 3]) == [10, 1, 20, 2, 30, 3]
assert alternating([True, False], [4, 5]) == [True, 4, False, 5]
assert alternating([], []) == []

質問：リストの長さが異なる場合、関数を修正して、長いリストの要素が最後になるようにすべての入力要素を返すようにしてください。

回答：修正された関数は次のとおりです。

def alternating(list1, list2):
   results = []
   for i in range(min(len(list1), len(list2))):
       results.append(list1[i])
       results.append(list2[i])
   if len(list1) > len(list2):
       results.extend(list1[i+1:])
   else:
       results.extend(list2[i+1:])
   return results

-----
"""

デモンストレーションの目的で、システムを10回複製して入力の長さを十分に長くし、Flash Attentionのメモリの節約を観察します。元のテキストプロンプト"質問：Pythonでバイトをギガバイトに変換する関数を書いてください。\n\n回答：ここに"を追加します。

long_prompt = 10 * system_prompt + prompt

再びbfloat16精度でモデルをインスタンス化します。

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

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

Flash Attentionなしでモデルを実行し、ピークGPUメモリ要件と推論時間を測定しましょう。

import time

start_time = time.time()
result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

出力:

10.96854019165039秒で生成されました。
はい。それを行う関数は次のようになります。\n\ndef bytes_to_giga(bytes):\n   return bytes / 1024 / 1024 / 1024\n\n答え：はい。それを行う関数は次のようになります。\n\ndef

前と同じ出力を得ていますが、今回はモデルが答えを複数回繰り返すようになっています。これは、デモンストレーションの目的でシステムのプロンプトを10回繰り返したため、モデルが自己繰り返しするように指示されたためです。

注意：実際のアプリケーションでは、システムのプロンプトは10回繰り返す必要はありません。1回だけで十分です！

ピーク時のGPUメモリ要件を測定しましょう。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

出力:

37.668193340301514

入力シーケンスが長くなったため、ピーク時のGPUメモリ要件が初期値よりも大幅に増加していることがわかります。また、生成には約1分以上かかります。

次の実験のために、GPUメモリを解放するためにflush()を呼び出します。

flush()

比較のために、同じ関数を実行しますが、Flash Attentionを有効にします。これを行うために、モデルをBetterTransformersに変換し、それによってPyTorchのSDPA self-attentionが有効になり、さらにFlash Attentionに基づいているようにします。

model.to_bettertransformer()

今度は先ほどとまったく同じコードスニペットを実行し、Under the hood TransformersはFlash Attentionを利用します。

start_time = time.time()
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

出力:

3.0211617946624756秒で生成されました。
 はい。それを行う関数は次のようになります。\n\ndef bytes_to_giga(bytes):\n   return bytes / 1024 / 1024 / 1024\n\n答え：はい。それを行う関数は次のようになります。\n\ndef

結果は先ほどとまったく同じですが、Flash Attentionのおかげで非常に高速化されていることがわかります。

最後にメモリの消費量を測定しましょう。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

出力:

32.617331981658936

そして、ほぼ元の29GBのピークGPUメモリに戻っています。

Flash Attentionを使用して非常に長い入力シーケンスを渡す場合、短い入力シーケンスを渡す場合と比べて、GPUメモリの使用量が約100MB増えることがわかります。

flush()

3. LLMアーキテクチャの内部科学：長いテキスト入力とチャットのための戦略的な選択

これまで、計算およびメモリの効率を向上させるために以下のことを考えてきました：

• 重みを低精度形式に変換すること
• セルフアテンションアルゴリズムをメモリおよび計算効率の高いバージョンに置き換えること

ここでは、LLMのアーキテクチャを変更することで、長いテキスト入力を必要とするタスクに対して最も効果的かつ効率的なものにする方法について見ていきます。例えば：

• リトリーバル拡張型質問応答
• 要約
• チャット

チャットでは、LLMが長いテキスト入力を処理するだけでなく、ユーザーとアシスタント（ChatGPTなど）のやり取りを効率的に処理できる必要があります。

トレーニング後、基本的なLLMアーキテクチャを変更することは困難ですので、LLMのタスクについて事前に考慮し、モデルのアーキテクチャを最適化することが重要です。大きな入力シーケンスに対して次の2つの重要なモデルアーキテクチャのコンポーネントがメモリまたはパフォーマンスのボトルネックとなることが速くにわかります。

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

各コンポーネントの詳細を見ていきましょう

3.1 LLMの位置エンベッディングの改善

セルフアテンションは、各トークンを他のトークンとの関係性に配置します。例えば、テキスト入力シーケンス “Hello”, “I”, “love”, “you” の Softmax(QKT) \text{Softmax}(\mathbf{QK}^T) Softmax(QKT) 行列は以下のようになります：

各単語トークンには、他のすべての単語トークンへの出現確率が与えられ、したがって他のすべての単語トークンと関係があります。例えば、”love” の単語は、”Hello” に0.05%、”I” に0.3%、自身に0.65%の確率で関連付けられます。

セルフアテンションに基づくLLMは、位置エンベッディングがない場合、テキスト入力の位置同士の理解に大きな困難を抱えます。これは、QKT \mathbf{QK}^T QKT によって計算される確率スコアが、他の単語トークンとの相対的な位置距離に関係なく、O(1) O(1) O(1) の計算で各単語トークンを関連付けるためです。そのため、位置エンベッディングのないLLMでは、各トークンが他のすべてのトークンと同じ距離にあるように見えます。例えば、”Hello I love you” と “You love I hello” の区別は非常に困難です。

文の順序を理解するために、LLMには追加の手がかりが必要であり、通常は位置エンコーディング（または位置エンベッディングとも呼ばれる）として適用されます。位置エンコーディングは、各トークンの位置を数値プレゼンテーションに変換し、LLMが文の順序をよりよく理解するために活用できるようにします。

Attention Is All You Need 論文の著者たちは、正弦波の位置エンベッディング P=p1,…,pN \mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N P=p1​,…,pN​ を導入しました。ここで、各ベクトル pi \mathbf{p}_i pi​ はその位置 i i i の正弦波関数として計算されます。位置エンベッディングは、単純に入力シーケンスベクトル X^=x^1,…,x^N \mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N X^=x^1​,…,x^N​ = x1+p1,…,xN+pN \mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N x1​+p1​,…,xN​+pN​ に単純に追加され、モデルに文の順序をより良く学習させます。

固定された位置エンベッディングを使用する代わりに、他の研究者（例：Devlin et al.）は、学習された位置エンコーディングを使用し、位置エンベッディング P \mathbf{P} P はトレーニング中に学習されます。

正弦波と学習された位置エンベッディングは、かつてLLMに文の順序をエンコードする主要な方法でしたが、これらの位置エンコーディングに関連するいくつかの問題が見つかりました：

• 1.) 正弦波と学習された位置エンベッディングは、どちらも絶対的な位置エンベッディングです。つまり、各位置ID：0,…,N 0, \ldots, N 0,…,N に対して一意のエンベッディングをエンコードします。Huang et al. や Su et al.] の研究によれば、絶対的な位置エンベッディングは、長いテキスト入力のLLMのパフォーマンスが低下します。長いテキスト入力の場合、モデルが入力トークン間の相対的な位置距離を学習する方が有利です。
• 2.) 学習された位置エンベッディングを使用する場合、LLMは固定された入力長 N N N で訓練される必要があります。これにより、訓練された入力長よりも長い入力長に対しての予測が困難になります。

最近では、上記の問題に対処できる相対的な位置エンベッディングがより人気となっています。特に以下のものが注目されています：

• Rotary Position Embedding (RoPE)
• ALiBi

RoPEとALiBiの両方は、単語トークンが関連付けられるセルフアテンションアルゴリズムに文の順序を直接伝えることが最善であると主張しています。具体的には、文の順序はQKT \mathbf{QK}^T QKT の計算を変更することで伝えるべきです。

詳細には触れずに、RoPEは位置情報をクエリ-キーペアにエンコードすることができることに注意してください。例えば、各ベクトルを角度θ∗iおよびθ∗jで回転させることにより、qi \mathbf{q}_i qi​およびxj \mathbf{x}_j xj​に位置情報をエンコードすることができます。ここで、i,jはそれぞれベクトルの文の位置を示しています：

q^iTx^j=qiTRθ,i−jxj. \mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j. q^​iT​x^j​=qiT​Rθ,i−j​xj​. したがって、Rθ,i−j \mathbf{R}_{\theta, i – j} Rθ,i−j​は回転行列を表します。θ \theta θは学習中に学習されるのではなく、訓練中の最大入力シーケンスの長さに依存する事前定義値に設定されます。

これにより、qi \mathbf{q}_i qi​とqj \mathbf{q}_j qj​の間の確率スコアは、i≠j i \ne j i=jの場合にのみ影響を受け、各ベクトルの特定の位置i i iおよびj j jに依存せずに相対距離i−j i – j i−jにのみ依存します。

RoPEは、次のような今日の重要なLLMの複数で使用されています：

• Falcon
• Llama
• PaLM

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

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

ALiBiは、次のような今日の重要なLLMの複数で使用されています：

• MPT
• BLOOM

RoPEとALiBiの位置エンコーディングの両方は、訓練中には学習されず、次の直感に基づいています：

• テキスト入力に関する位置情報は、自己注意層のQKT QK^T QKT行列に直接与えられるべきです。
• LLMは、相対距離の位置エンコーディングを学習することが奨励されるべきです。
• テキスト入力トークンが互いから遠く離れているほど、クエリ-キーの確率が低くなります。RoPEとALiBiは、クエリ-キーの確率を低下させます。RoPEは、クエリ-キーベクトル間の角度を増加させることにより、ベクトル積を減少させます。ALiBiは、ベクトル積に大きな負の数を追加することにより、確率を低下させます。

結論として、大規模なテキスト入力を処理するタスクに展開されるLLMは、RoPEやALiBiなどの相対位置エンコーディングを使用して訓練する方が良いです。また、RoPEとALiBiのLLMは、例えばN1=2048の固定長で訓練されている場合でも、N1 N_1 N1​よりもはるかに大きなテキスト入力で実際に使用することができます。この場合、位置エンコーディングを外挿することができます。

3.2 キー・値キャッシュ

LLMによる自己回帰テキスト生成は、入力シーケンスを反復的に入力し、次のトークンをサンプリングし、次のトークンを入力シーケンスに追加し、LLMが生成が終了したことを示すトークンを生成するまで続けることで機能します。

自己回帰生成がどのように機能するかをより視覚的に説明するために、Transformerのテキスト生成チュートリアルをご覧ください。

実際の動作を示すために、以下のコードスニペットを実行しましょう。単純にtorch.argmaxを使用して、最も可能性の高い次のトークンを取得します。

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

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("input_idsの形状", input_ids.shape)

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

出力:

input_idsの形状 torch.Size([1, 21])
input_idsの形状 torch.Size([1, 22])
input_idsの形状 torch.Size([1, 23])
input_idsの形状 torch.Size([1, 24])
input_idsの形状 torch.Size([1, 25])
[' ここにはPythonの関数があります']

入力トークンがサンプリングされるたびに、テキスト入力トークンが増加していることがわかります。

ほとんどの場合、LLMは因果言語モデリング目的を使用してトレーニングされるため、注意スコアの上三角行列をマスクします。これが、上記の2つの図で注意スコアが空白になっている理由です（すなわち、確率が0です）。因果言語モデリングについてのクイックな復習については、Illustrated Self Attentionブログを参照してください。

その結果、トークンは以前のトークンに依存せず、具体的には、qi \mathbf{q}_i qi​ベクトルは、qi \mathbf{q}_i qi​ベクトルがキー、値ベクトル kj,vj \mathbf{k}_j, \mathbf{v}_j kj​,vj​（ただし、j>i j > i j>i ）と関連付けられないためです。代わりに、qi \mathbf{q}_i qi​は、m∈{0,…i−1} \text{ , for } m \in \{0, \ldots i – 1\} m∈{0,…i−1} に対して以前のキー・値ベクトル km<i,vm<i \mathbf{k}_{m < i}, \mathbf{v}_{m < i} km<i​,vm<i​にのみアテンションを行います。不要な計算を削減するために、各レイヤーのキー・値ベクトルをすべての前の時間ステップのキャッシュで保存することができます。

次に、キー・値キャッシュを使用するようLLMに指示し、各フォワードパスでキャッシュを取得して転送する方法を説明します。Transformersでは、forward呼び出しにuse_cacheフラグを渡すことでキー・値キャッシュを取得し、現在のトークンとともに渡すことができます。

past_key_values = None # past_key_valuesはキー・値キャッシュです
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")

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("input_idsの形状", input_ids.shape)
  print("キー・値キャッシュの長さ", len(past_key_values[0][0]))  # past_key_valuesの形状は[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

出力:

input_idsの形状 torch.Size([1, 20])
キー・値キャッシュの長さ 20
input_idsの形状 torch.Size([1, 20])
キー・値キャッシュの長さ 21
input_idsの形状 torch.Size([1, 20])
キー・値キャッシュの長さ 22
input_idsの形状 torch.Size([1, 20])
キー・値キャッシュの長さ 23
input_idsの形状 torch.Size([1, 20])
キー・値キャッシュの長さ 24
[' ここ', ' は', ' Python', ' の', ' 関数です']

ご覧の通り、キーと値のキャッシュを使用すると、テキスト入力トークンの長さは増加せず、単一の入力ベクトルのままです。一方、デコーディングステップごとにキーと値のキャッシュの長さが1つずつ増加します。

キーと値のキャッシュを利用することにより、QKT \mathbf{QK}^T QKT は実質的に qcKT \mathbf{q}_c\mathbf{K}^T qc​KT に簡略化されます。ここで、qc \mathbf{q}_c qc​ は常に単一のベクトルである現在の入力トークンのクエリ射影です。

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

• 完全なQKT \mathbf{QK}^T QKT 行列を計算する場合と比べて、計算効率が大幅に向上します。これにより、推論速度が向上します。
• 必要な最大メモリは生成されるトークンの数の二乗ではなく、線形に増加します。

キーと値のキャッシュを常に利用することは、同じ結果と長い入力シーケンスの場合の著しい高速化をもたらします。Transformersでは、テキストパイプラインやgenerateメソッドを使用する場合、キーと値のキャッシュがデフォルトで有効になっています。

キーと値のキャッシュは、チャットなどのアプリケーションに特に有用です。ここでは、自己回帰デコーディングの複数のパスが必要な場合に特に有用です。例を見てみましょう。

User: フランスには何人の人が住んでいますか？ 
Assistant: フランスにはおおよそ7500万人が住んでいます 
User: では、ドイツは何人いますか？ 
Assistant: ドイツには約8100万人の住民がいます

このチャットでは、LLMは自己回帰デコーディングを2回実行します：

• 1. 最初の回では、キーと値のキャッシュは空であり、入力プロンプトは"User: フランスには何人の人が住んでいますか？"であり、モデルは自己回帰的にテキスト"フランスにはおおよそ7500万人が住んでいます"を生成します。このとき、キーと値のキャッシュはデコーディングの各ステップで増加します。
• 1. 2回目の入力プロンプトは"User: フランスには何人の人が住んでいますか？ \n Assistant: フランスにはおおよそ7500万人が住んでいます \n User: では、ドイツは何人いますか？"です。キャッシュのおかげで、最初の2つの文のキーと値のベクトルはすでに計算されています。したがって、入力プロンプトは"User: では、ドイツは何人いますか？"のみです。短縮された入力プロンプトを処理する間、計算されたキーと値のベクトルは最初のデコーディングのキーと値のキャッシュに連結されます。2番目のAssistantの回答"ドイツには約8100万人の住民がいます"は、"User: フランスには何人の人が住んでいますか？ \n Assistant: フランスにはおおよそ7500万人が住んでいます \n User: では、ドイツは何人いますか？"とエンコードされたキーと値のベクトルで構成されるキーと値のキャッシュを使用して自己回帰的に生成されます。

ここで注意すべきことは、

• 1. チャットに展開されたLLMでは、すべての文脈を保持することが重要です。つまり、前の会話のすべての文脈をLLMが理解する必要があります。例えば、上記の例では、ユーザーが"And how many are in Germany"と尋ねるときに、LLMが人口を指していることを理解する必要があります。
• 1. キーと値のキャッシュは、エンコーダー・デコーダーアーキテクチャを使用する場合とは異なり、チャットに非常に便利です。なぜなら、キーと値のキャッシュを使用すると、チャット履歴をゼロから再エンコードする必要がなく、エンコード済みのチャット履歴を連続的に拡張できるからです。

ただし、一つ注意が必要です。QKT \mathbf{QK}^T QKT 行列の必要なピークメモリは大幅に削減されますが、キーと値のキャッシュをメモリに保持することは、長い入力シーケンスやマルチターンのチャットの場合に非常にメモリを消費する可能性があります。キーと値のキャッシュは、すべての自己注意層とすべてのアテンションヘッドに対して、すべての以前の入力ベクトルxi, for i∈{1,…,c−1} \mathbf{x}_i \text{, for } i \in \{1, \ldots, c – 1\} xi​, for i∈{1,…,c−1} のキーと値のベクトルを格納する必要があります。

前に使用したLLM bigcode/octocoder のキーと値のキャッシュに格納する必要がある浮動小数点値の数を計算しましょう。浮動小数点値の数は、シーケンスの長さの2倍にアテンションヘッドの数を乗じ、アテンションヘッドの次元数と層の数を乗じたものです。このLLMを仮想的な入力シーケンス長16000で計算すると、次のようになります：

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

出力:

7864320000

おおよそ80億の浮動小数点値です！float16精度で80億の浮動小数点値を保存するには、モデルの重み自体の約半分と同じくらいの15 GBのRAMが必要です！研究者は、キー・バリューキャッシュのメモリコストを大幅に削減するための2つの方法を提案しています：

• 1. Multi-Query-Attention（MQA）

Multi-Query-Attentionは、Noam ShazeerのFast Transformer Decoding: One Write-Head is All You Need論文で提案されました。タイトルが示すように、Noamは、n_headのキー・バリュープロジェクション重みの代わりに、モデルのパフォーマンスが著しく低下することなく、すべてのアテンションヘッドで共有される1つのヘッド・バリュープロジェクション重みペアを使用できることを発見しました。

1つのヘッド・バリュープロジェクション重みペアを使用することで、キーバリューベクトルki,vi \mathbf{k}_i, \mathbf{v}_i ki​,vi​はすべてのアテンションヘッドで同一である必要があります。その結果、キーバリュープロジェクションペアはn_head個ではなく、1つだけをキャッシュに保存する必要があります。

ほとんどのLLMは20から100のアテンションヘッドを使用するため、MQAはキーバリューキャッシュのメモリ消費量を大幅に削減します。このノートブックで使用されるLLMでは、入力シーケンス長が16000の場合、必要なメモリ消費量を15 GBから400 MB未満に減らすことができます。

メモリの節約に加えて、MQAは以下で説明するように計算効率も向上させます。オートレグレッシブデコーディングでは、大きなキーバリューベクトルをリロードし、現在のキーバリューベクトルペアと連結して、各ステップでqcKT \mathbf{q}_c\mathbf{K}^T qc​KT計算に供給する必要があります。オートレグレッシブデコーディングでは、定数のリロードに必要なメモリ帯域幅は深刻な時間のボトルネックになる可能性があります。キーバリューベクトルのサイズを減らすことで、アクセスする必要のあるメモリ量が減少し、メモリ帯域幅のボトルネックが軽減されます。詳細については、Noamの論文をご覧ください。

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

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

• Falcon
• PaLM
• MPT
• BLOOM

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

• 1. Grouped-Query-Attention（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のいずれかを活用することが強く推奨されます。

結論

研究コミュニティは、常に推論時間を高速化するための新しい便利な方法を提案しています。例えば、1つの有望な研究方向として、予測的デコーディングがあります。この方法では、より小さく、より高速な言語モデルによって「簡単なトークン」が生成され、LLM自体によってのみ「難しいトークン」が生成されます。詳細については、この素敵なブログ記事で読むことができますが、このノートブックの範囲外です。

GPT3/4、Llama-2-70b、Claude、PaLMなどの巨大なLLMがHugging Face ChatやChatGPTなどのチャットインターフェースで非常に速く実行できるのは、精度、アルゴリズム、アーキテクチャの改善に大いに貢献しているためです。今後は、GPU、TPUなどのアクセラレータがさらに高速化され、より多くのメモリが利用可能になるでしょうが、常に最高のアルゴリズムとアーキテクチャを使用して最大の効果を得るようにすることを確認する必要があります🤗

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!