投機的サンプリング、KV キャッシュおよび OpenVINO™ によるテキスト生成¶
この Jupyter ノートブックは、ローカルへのインストール後にのみ起動できます。
モデルのサイズが大きくなるにつれて、生成 AI の実装には大量の推論リソースが必要になります。これにより、プロンプトからの生成あたりのコストが増加するだけでなく、要求を処理するために使用される電力消費も増加します。
テキスト生成の推論の最適化は、コストと電力消費を削減するのに不可欠です。推論プロセスを最適化すると、テキスト生成に必要な時間とエネルギーを大幅に削減できます。これにより、ハードウェアとソフトウェア両面のコストと電力消費の削減につながります。さらに、推論の最適化により、テキスト生成の精度と生成速度が向上します。これにより、ユーザー・エクスペリエンスが向上し、テキスト生成タスクの効率が向上します。要約すると、テキスト生成の推論の最適化は、コストと電力消費を削減し、テキスト生成の精度と速度を向上するには不可欠です。
もう 1 つの必要な条件は、最適化が相互に互換性があることです。つまり、特定の最適化を実装しても、他の最適化が妨げられることはありません。全体の効率を損なうような “衝突” を引き起こすことなく大幅な高速化を実現できる最適化のレベルがいくつかあります。
この方法の詳細については、Chen 氏らの論文 (http://arxiv.org/abs/2302.01318) を参照してください。さらに、Leviathan 氏らによる投機的サンプリングの正当性の証明 (元の分布が保存されることを示す) があります (http://arxiv.org/abs/2211.17192)。
OpenVINO を使用した実装について説明したブログ記事は、openvino.ai でご覧いただけます。
目次¶
必要条件¶
最初に、OpenVINO 統合によって高速化された Hugging Face Optimum ライブラリーをインストールする必要があります。Hugging Face Optimum Intel API は、Hugging Face Transformers ライブラリーのモデルを OpenVINO™ IR 形式に変換および量子化できる高レベル API です。詳細については、Hugging Face Optimum Intel のドキュメントを参照してください。
Transformers (Hugging Face) やその他のモジュールもインストールする必要があります。
%pip install -q --upgrade pip
%pip install -q --upgrade transformers torch gradio openvino accelerate onnx ipywidgets --extra-index-url https://download.pytorch.org/whl/cpu
%pip install -q "git+https://github.com/huggingface/optimum-intel.git"
推論デバイスの選択¶
OpenVINO を使用して推論を実行するデバイスをドロップダウン・リストから選択します。
import ipywidgets as widgets
from openvino.runtime import Core
core = Core()
device = widgets.Dropdown(
options=core.available_devices + ["AUTO"],
value='CPU',
description='Device:',
disabled=False,
)
device
KV キャッシュのサポートにより、自己回帰および推測形式のサンプリングを作成¶
テキスト生成は多くの場合、自己回帰方式で行われます。ここでは、コード内で KV キャッシュ (別名 Past Value Cache) をサポートします。貪欲なサンプリングを使用していることに注意してください。他のテキスト生成パラメーター (温度など) は調整しないため、推測的サンプリングの図はできるだけシンプルで分かりやすいものにしてください。
インポートの設定¶
import time
import numpy as np
import torch
import gradio as gr
自己回帰サンプリングを準備¶
def max_fn(x):
x_max = torch.where(x > 0, x, torch.zeros_like(x))
return x_max / torch.sum(x_max)
def autoregressive_sampling_with_pkv(x, model, N):
n = len(x)
T = n + N
input = x
past_kv = None
while n < T:
res = model(input, attention_mask=torch.ones(input.size(), dtype=torch.long), past_key_values=past_kv)
model_out = torch.softmax(res.logits, dim=2)
past_kv = res.past_key_values
next_token = torch.reshape(torch.argmax(model_out[-1][-1]), (1, 1))
x = torch.cat((x, next_token), dim=1)
n += 1
input = next_token
return x
推測サンプリングを準備¶
ステップ 1: 投機的サンプリングでは、まずドラフトモデルから K 個のサンプルを (自己回帰方式で) 生成します。
ステップ 2: これらは、バッチサイズ K を使用してターゲットモデル (ステップ 2) を使用して調査する候補になります。
ステップ 3: ここで、ステップ 2 でターゲットモデルから生成されたロジットに基づいて、ドラフトモデルからの K 個の候補が許容可能かどうかを判断します。
ステップ 4: 追加コストなしで別のトークンをサンプリングできます (すべての候補が受け入れられたと仮定)。
ステップ 3 の受け入れ基準は、ターゲットモデルのロジットを比較し、ドラフトモデルと比較する必要があります。比率が十分に高ければ、受け入れられる可能性が高くなります (乱数を使用)。
def speculative_sampling_with_pkv(x, draft_model, target_model, N, K):
n = x.size(1)
T = n + N
target_past_kv = None
while n < T:
# Step 1: autoregressive decode of K candidate tokens from
# the draft model and get final p for this batch of candidates
x_draft = None
draft_past_kv = None
x_draft_input = x
p_cum = None
for _ in range(K):
res_draft = draft_model(x_draft_input, attention_mask=torch.ones(x_draft_input.size(), dtype=torch.long), past_key_values=draft_past_kv, use_cache=True)
p = res_draft.logits
p = torch.softmax(p, dim=2)
draft_past_kv = res_draft.past_key_values
next_token = torch.reshape(torch.argmax(p[-1][-1]), (1, 1))
x_draft_input = next_token
if p_cum is None:
p_cum = p[:, -1].unsqueeze(1)
x_draft = next_token
else:
p_cum = torch.cat((p_cum, p), dim=1)
x_draft = torch.cat((x_draft, next_token), dim=1)
# Step 2: target model forward passes on x_draft
if target_past_kv is None:
x_draft_target_input = torch.cat((x, x_draft), dim=1)
else:
x_draft_target_input = x_draft
res = target_model(x_draft_target_input, attention_mask=torch.ones(x_draft_target_input.size(), dtype=torch.long), use_cache=False)
q = res.logits
target_new_past_kv = res.past_key_values
# Step 3: append draft tokens based on acceptance-rejection criterion and resample a token on rejection
all_accepted = True
for k in range(K):
j = x_draft[0][k].item()
q_item = q[-1][k][j].detach().numpy()
p_item = p_cum[-1][k][j].detach().numpy()
if np.random.random() < min(1, (q_item / p_item)): # accepted
x = torch.cat((x, torch.tensor(j).reshape(1,1)), dim=1)
n += 1
else: # rejected
q_p = max_fn(q[0][k] - p_cum[0][k])
resampled_output = torch.argmax(q_p)
resampled_output = torch.reshape(resampled_output, (1,1))
x = torch.cat((x, resampled_output), dim=1)
n += 1
all_accepted = False
break
target_past_kv = target_new_past_kv
# Step 4: if all draft tokens were accepted, sample a final token
if all_accepted:
x = torch.cat((x, torch.reshape(torch.argmax(q[-1][-1]), (1,1))), dim=1)
n += 1
return x
メイン生成関数¶
モデルのダウンロードと変換¶
Optimum Intel を使用すると、Hugging Face Hub から最適化されたモデルをロードし、Hugging Face API を使用して OpenVINO ランタイムで推論を実行するパイプラインを作成できます。Optimum 推論モデルは、Hugging Face Transformers モデルと API の互換性があります。つまり、AutoModelForXxx クラスを対応する OVModelForXxx クラスに置き換えるだけで済みます。
以下は Dolly モデルの例です。
-from transformers import AutoModelForCausalLM
+from optimum.intel.openvino import OVModelForCausalLM
from transformers import AutoTokenizer, pipeline
model_id = "databricks/dolly-v2-3b"
-model = AutoModelForCausalLM.from_pretrained(model_id)
+model = OVModelForCausalLM.from_pretrained(model_id, from_transformers=True)
モデルクラスの初期化は、from_pretrained メソッドの呼び出しから始まります。Transformers モデルをダウンロードして変換する場合は、パラメーター from_transformers=True を追加する必要があります。save_pretrained メソッドを使用して、変換されたモデルを次回に使用するため保存できます。トークナイザー・クラスとパイプライン API は Optimum モデルと互換性があります。
from pathlib import Path
from transformers import AutoTokenizer
from optimum.intel.openvino import OVModelForCausalLM
# If you are on a large system with lots of memory, you can run a larger model like DollyV2
# draft_model_id = "databricks/dolly-v2-3b"
# draft_model_path = Path("dolly-v2-3b")
# target_model_id = "databricks/dolly-v2-12b"
# target_model_path = Path("dolly-v2-12b")
# If you are on a system with limited memory, you can try the smaller GPT2 models
draft_model_id = "gpt2"
draft_model_path = Path("gpt2-local")
target_model_id = "gpt2-xl"
target_model_path = Path("gpt2-xl-local")
target_tokenizer = AutoTokenizer.from_pretrained(target_model_id)
current_device = device.value
# Save local copies for subsequent runs
if draft_model_path.exists():
draft_ov_model = OVModelForCausalLM.from_pretrained(draft_model_path, device=current_device)
else:
draft_ov_model = OVModelForCausalLM.from_pretrained(draft_model_id, device=current_device, from_transformers=True)
draft_ov_model.save_pretrained(draft_model_path)
if target_model_path.exists():
target_ov_model = OVModelForCausalLM.from_pretrained(target_model_path, device=current_device)
else:
target_ov_model = OVModelForCausalLM.from_pretrained(target_model_id, device=current_device, from_transformers=True)
target_ov_model.save_pretrained(target_model_path)
def main(
prompt: str = "Explain the difference between fission and fusion",
n_tokens_to_generate: int = 100,
K: int = 5,
seed: int = 5555,
):
# seed numpy rng
np.random.seed(seed)
draft_model = draft_ov_model
target_model = target_ov_model
input_ids = target_tokenizer(prompt, return_tensors="pt")['input_ids']
def run_autoregressive_sampling_fn(decode_fn, input_ids, **kwargs):
start = time.perf_counter()
output_ids = decode_fn(x=input_ids, **kwargs)
text = target_tokenizer.decode(output_ids[0], skip_special_tokens=True)
elapsed_time = time.perf_counter() - start
return text, elapsed_time
def run_speculative_sampling_fn(decode_fn, input_ids, **kwargs):
start = time.perf_counter()
output_ids = decode_fn(x=input_ids, **kwargs)
text = target_tokenizer.decode(output_ids[0], skip_special_tokens=True)
elapsed_time = time.perf_counter() - start
return text, elapsed_time
autoregressive_text, autoregressive_time = run_autoregressive_sampling_fn(
autoregressive_sampling_with_pkv,
input_ids,
model=target_model,
N=n_tokens_to_generate,
)
speculative_text, speculative_time = run_speculative_sampling_fn(
speculative_sampling_with_pkv,
input_ids,
target_model=target_model,
draft_model=draft_model,
N=n_tokens_to_generate,
K=K,
)
# Format results for output in gradio
out = "\n" + "Autoregressive Decode" + "\n" + "---------------------" + "\n"
out = out + f"Time = {autoregressive_time:.2f}s" + "\n" + f"Text = {autoregressive_text}" + "\n"
out = out + "\n" + "Speculative Decode" + "\n" + "------------------" + "\n"
out = out + f"Time = {speculative_time:.2f}s" + "\n" + f"Text = {speculative_text}"
return out
if __name__ == "__main__":
with gr.Blocks() as demo:
gr.Markdown(
"""
# Speculative Sampling Demo
## The output will show a comparison of Autoregressive Sampling vs Speculative Sampling
- Target Model: Dolly V2 12B
- Draft Model: Dolly V2 3B
- K = 5
> Some improvements can be made to acceptance criterion and adjusting temperature to improve text quality.
""")
with gr.Row():
inp = gr.Textbox(placeholder="THIS CANNOT BE EMPTY", label="Input Prompt")
out = gr.Textbox(label="Output")
btn = gr.Button("Run")
btn.click(fn=main, inputs=inp, outputs=out)
demo.launch()