LLM ์ถ๋ก ์ต์ ํ [[llm-inference-optimization]]
๋๊ท๋ชจ ์ธ์ด ๋ชจ๋ธ(LLM)์ ์ฑํ ๋ฐ ์ฝ๋ ์์ฑ ๋ชจ๋ธ๊ณผ ๊ฐ์ ํ ์คํธ ์์ฑ ์์ฉ ํ๋ก๊ทธ๋จ์ ํ ๋จ๊ณ ๋์ด์ฌ๋ฆฌ๋ฉฐ, ๋์ ์์ค์ ์ดํด๋ ฅ๊ณผ ์ ์ฐฝํจ์ ๋ณด์ฌ์ฃผ๋ ํ ์คํธ๋ฅผ ์์ฑํฉ๋๋ค. ๊ทธ๋ฌ๋ LLM์ ๊ฐ๋ ฅํ๊ฒ ๋ง๋๋ ์์์ธ ๊ทธ๋ค์ ํฌ๊ธฐ๋ ๋์์ ์ถ๋ก ๊ณผ์ ์์ ๋์ ๊ณผ์ ๊ฐ ๋๊ธฐ๋ ํฉ๋๋ค.
๊ธฐ๋ณธ์ ์ธ ์ถ๋ก ์ ๋๋ฆฝ๋๋ค, ์๋ํ๋ฉด LLM์ด ๋ค์ ํ ํฐ์ ์์ฑํ๊ธฐ ์ํด ๋ฐ๋ณต์ ์ผ๋ก ํธ์ถ๋์ด์ผ ํ๊ธฐ ๋๋ฌธ์ ๋๋ค. ์์ฑ์ด ์งํ๋จ์ ๋ฐ๋ผ ์ ๋ ฅ ์ํ์ค๊ฐ ๊ธธ์ด์ ธ ์ฒ๋ฆฌ ์๊ฐ์ด ์ ์ ๊ธธ์ด์ง๋๋ค. ๋ํ, LLM์ ์์ญ์ต ๊ฐ์ ๋งค๊ฐ๋ณ์๋ฅผ ๊ฐ์ง๊ณ ์์ด ๋ชจ๋ ๊ฐ์ค์น๋ฅผ ๋ฉ๋ชจ๋ฆฌ์ ์ ์ฅํ๊ณ ์ฒ๋ฆฌํ๋ ๋ฐ ์ด๋ ค์์ด ์์ต๋๋ค.
์ด ๊ฐ์ด๋๋ LLM ์ถ๋ก ์ ๊ฐ์ํ๊ธฐ ์ํด Transformers์์ ์ฌ์ฉํ ์ ์๋ ์ต์ ํ ๊ธฐ์ ์ ์ฌ์ฉํ๋ ๋ฐฉ๋ฒ์ ๋ณด์ฌ์ค๋๋ค.
Hugging Face๋ LLM์ ์ถ๋ก ์ ์ต์ ํํ์ฌ ๋ฐฐํฌํ๊ณ ์๋น์คํ๋ ๋ฐ ์ ๋ ํ๋ ๋ผ์ด๋ธ๋ฌ๋ฆฌ์ธ Text Generation Inference (TGI)์ ์ ๊ณตํฉ๋๋ค. ์ด ๋ผ์ด๋ธ๋ฌ๋ฆฌ๋ ์ฒ๋ฆฌ๋ ์ฆ๊ฐ๋ฅผ ์ํ ์ง์์ ์ธ ๋ฐฐ์นญ๊ณผ ๋ค์ค GPU ์ถ๋ก ์ ์ํ ํ ์ ๋ณ๋ ฌํ์ ๊ฐ์ Transformers์ ํฌํจ๋์ง ์์ ๋ฐฐํฌ ์งํฅ ์ต์ ํ ๊ธฐ๋ฅ์ ํฌํจํฉ๋๋ค.
์ ์ kv-cache์ torch.compile[[static-kv-cache-and-torchcompile]]
๋์ฝ๋ฉ ์ค์ LLM์ ๊ฐ ์ ๋ ฅ ํ ํฐ์ ๋ํ key-value(kv) ๊ฐ์ ๊ณ์ฐํฉ๋๋ค. LLM์ ์๊ธฐํ๊ท(autoregressive)์ด๊ธฐ ๋๋ฌธ์ ์์ฑ๋ ์ถ๋ ฅ์ด ํ์ฌ ์ ๋ ฅ์ ์ผ๋ถ๊ฐ ๋์ด ๋งค๋ฒ ๋์ผํ kv ๊ฐ์ ๊ณ์ฐํฉ๋๋ค. ์ด๋ ๋งค๋ฒ ๋์ผํ kv ๊ฐ์ ๋ค์ ๊ณ์ฐํ๊ธฐ ๋๋ฌธ์ ํจ์จ์ ์ด์ง ์์ต๋๋ค.
์ด๋ฅผ ์ต์ ํํ๊ธฐ ์ํด, ์ด์ ํค(key)์ ๊ฐ(value)์ ์ฌ๊ณ์ฐํ์ง ์๊ณ ์ ์ฅํ๋ kv-cache๋ฅผ ์ฌ์ฉํ ์ ์์ต๋๋ค. ๊ทธ๋ฌ๋ kv-cache๋ ๊ฐ ์์ฑ ๋จ๊ณ์์ ์ฆ๊ฐํ๋ฉฐ ๋์ ์ด๊ธฐ ๋๋ฌธ์ PyTorch ์ฝ๋๋ฅผ ๋น ๋ฅด๊ณ ์ต์ ํ๋ ์ปค๋๋ก ํตํฉํ๋ ๊ฐ๋ ฅํ ์ต์ ํ ๋๊ตฌ์ธ torch.compile์ ์ฌ์ฉํ๋ ๋ฐ ์ ์ฝ์ด ์์ต๋๋ค.
์ ์ kv-cache๋ ์ต๋๊ฐ์ ๋ฏธ๋ฆฌ ํ ๋นํ์ฌ ์ด ๋ฌธ์ ๋ฅผ ํด๊ฒฐํ์ฌ torch.compile๊ณผ ๊ฒฐํฉํ ์ ์๊ฒ ํฉ๋๋ค. ์ด๋ฅผ ํตํด ์ต๋ 4๋ฐฐ์ ์๋ ํฅ์์ด ๊ฐ๋ฅํฉ๋๋ค. ์๋ ํฅ์์ ๋ชจ๋ธ ํฌ๊ธฐ(๋ ํฐ ๋ชจ๋ธ์ ์๋ ํฅ์์ด ์ ์)์ ํ๋์จ์ด์ ๋ฐ๋ผ ๋ค๋ฅผ ์ ์์ต๋๋ค.
ํ์ฌ Llama ๋ฐ ๋ช ๊ฐ์ง ๋ค๋ฅธ ๋ชจ๋ธ๋ง ์ ์ kv-cache์
torch.compile์ ์ง์ํฉ๋๋ค. ์ค์๊ฐ ๋ชจ๋ธ ํธํ์ฑ ๋ชฉ๋ก์ ์ด ์ด์๋ฅผ ํ์ธํ์ญ์์ค.
์์
์ ๋ณต์ก์ฑ์ ๋ฐ๋ผ ์ธ ๊ฐ์ง ๋ฐฉ์์ ์ ์ kv-cache ์ฌ์ฉ ๋ฐฉ๋ฒ์ด ์์ต๋๋ค:
1. ๊ธฐ๋ณธ ์ฌ์ฉ๋ฒ: generation_config์์ ํ๋๊ทธ๋ฅผ ์ค์ ํ๊ธฐ๋ง ํ๋ฉด ๋ฉ๋๋ค(๊ถ์ฅ);
2. ๊ณ ๊ธ ์ฌ์ฉ๋ฒ: ์ฌ๋ฌ ๋ฒ์ ์์ฑ์ด๋ ๋ง์ถคํ ์์ฑ ๋ฃจํ๋ฅผ ์ํด ์บ์ ๊ฐ์ฒด๋ฅผ ์ฒ๋ฆฌํฉ๋๋ค;
3. ๊ณ ๊ธ ์ฌ์ฉ๋ฒ: ๋จ์ผ ๊ทธ๋ํ๊ฐ ํ์ํ ๊ฒฝ์ฐ, ์ ์ฒด generate ํจ์๋ฅผ ํ๋์ ๊ทธ๋ํ๋ก ์ปดํ์ผํฉ๋๋ค.
์ฌ๋ฐ๋ฅธ ํญ์ ์ ํํ์ฌ ๊ฐ ๋ฐฉ๋ฒ์ ๋ํ ์ถ๊ฐ ์ง์นจ์ ํ์ธํ์ธ์.
torch.compile์ ์ฌ์ฉํ ๋ ์ด๋ค ์ ๋ต์ ์ฌ์ฉํ๋ , LLM ์ ๋ ฅ์ ์ ํ๋ ๊ฐ ์ธํธ๋ก ์ผ์ชฝ์ ํจ๋ฉํ๋ฉด ๋ชจ์๊ณผ ๊ด๋ จ๋ ์ฌ์ปดํ์ผ์ ํผํ ์ ์์ต๋๋ค.pad_to_multiple_oftokenizer flag๊ฐ ์ ์ฉํ ๊ฒ์ ๋๋ค!
์ด ์์ ์์๋ Gemma ๋ชจ๋ธ์ ์ฌ์ฉํด ๋ณด๊ฒ ์ต๋๋ค. ํ์ํ ์์ ์ ๋ค์๊ณผ ๊ฐ์ต๋๋ค:
- ๋ชจ๋ธ์
generation_config์์ฑ์ ์ ๊ทผํ์ฌcache_implementation์ "static"์ผ๋ก ์ค์ ํฉ๋๋ค; - ๋ชจ๋ธ์
forwardํจ์ค๋ฅผ ์ ์ kv-cache์ ํจ๊ป ์ปดํ์ผํ๊ธฐ ์ํดtorch.compile์ ํธ์ถํฉ๋๋ค.
์ด๋ ๊ฒ ํ๋ฉด ๋์ ๋๋ค!
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false" # ๊ธด ๊ฒฝ๊ณ ๋ฉ์์ง๋ฅผ ๋ฐฉ์งํ๊ธฐ ์ํด ์ค์ :)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
model.generation_config.cache_implementation = "static"
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
input_text = "The theory of special relativity states "
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
outputs = model.generate(**input_ids)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
['The theory of special relativity states 1. The speed of light is constant in all inertial reference']
generate ํจ์๋ ๋ด๋ถ์ ์ผ๋ก ๋์ผํ ์บ์ ๊ฐ์ฒด๋ฅผ ์ฌ์ฌ์ฉํ๋ ค๊ณ ์๋ํ๋ฉฐ, ์ด๋ฅผ ํตํด ๊ฐ ํธ์ถ ์ ์ฌ์ปดํ์ผ์ ํ์์ฑ์ ์ ๊ฑฐํฉ๋๋ค. ์ฌ์ปดํ์ผ์ ํผํ๋ ๊ฒ์ torch.compile์ ์ฑ๋ฅ์ ์ต๋ํ ํ์ฉํ๋ ๋ฐ ๋งค์ฐ ์ค์ํ๋ฉฐ, ๋ค์ ์ฌํญ์ ์ ์ํด์ผ ํฉ๋๋ค:
- ๋ฐฐ์น ํฌ๊ธฐ๊ฐ ๋ณ๊ฒฝ๋๊ฑฐ๋ ํธ์ถ ๊ฐ ์ต๋ ์ถ๋ ฅ ๊ธธ์ด๊ฐ ์ฆ๊ฐํ๋ฉด ์บ์๋ฅผ ๋ค์ ์ด๊ธฐํํด์ผ ํ๋ฉฐ, ์ด๋ก ์ธํด ์๋ก ์ปดํ์ผ์ ํด์ผ ํฉ๋๋ค;
- ์ปดํ์ผ๋ ํจ์์ ์ฒซ ๋ช ๋ฒ์ ํธ์ถ์ ํจ์๊ฐ ์ปดํ์ผ๋๋ ๋์ ๋ ๋๋ฆฝ๋๋ค.
๋ค์ค ํด ๋ํ์ ๊ฐ์ ์ ์ ์บ์์ ๊ณ ๊ธ ์ฌ์ฉ์ ์ํด์๋, ์บ์ ๊ฐ์ฒด๋ฅผ [
~GenerationMixin.generate] ์ธ๋ถ์์ ์ธ์คํด์คํํ๊ณ ์กฐ์ํ๋ ๊ฒ์ ๊ถ์ฅํฉ๋๋ค. ๊ณ ๊ธ ์ฌ์ฉ๋ฒ ํญ์ ์ฐธ์กฐํ์ธ์.
[StaticCache] ๊ฐ์ฒด๋ past_key_values ์ธ์๋ก ๋ชจ๋ธ์ [~GenerationMixin.generate] ํจ์์ ์ ๋ฌํ ์ ์์ต๋๋ค. ์ด ๊ฐ์ฒด๋ ์บ์ ๋ด์ฉ์ ์ ์งํ๋ฏ๋ก, ๋์ ์บ์๋ฅผ ์ฌ์ฉํ๋ ๊ฒ์ฒ๋ผ ์๋ก์ด [~GenerationMixin.generate] ํธ์ถ์ ์ด๋ฅผ ์ ๋ฌํ์ฌ ์์ฑ์ ๊ณ์ํ ์ ์์ต๋๋ค.
from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache
import torch
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false" # ๊ธด ๊ฒฝ๊ณ ๋ฉ์์ง๋ฅผ ๋ฐฉ์งํ๊ธฐ ์ํด ์ค์ :)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
input_text = "The theory of special relativity states "
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
prompt_length = input_ids.input_ids.shape[1]
model.generation_config.max_new_tokens = 16
past_key_values = StaticCache(
config=model.config,
batch_size=1,
# ์บ์๋ฅผ ์ฌ์ฌ์ฉํ ๊ณํ์ด ์๋ ๊ฒฝ์ฐ, ๋ชจ๋ ๊ฒฝ์ฐ์ ์ถฉ๋ถํ ์บ์ ๊ธธ์ด๋ฅผ ์ค์ ํด์ผ ํฉ๋๋ค
max_cache_len=prompt_length+(model.generation_config.max_new_tokens*2),
device=model.device,
dtype=model.dtype
)
outputs = model.generate(**input_ids, past_key_values=past_key_values)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
['The theory of special relativity states 1. The speed of light is constant in all inertial reference frames. 2']
# ์์ฑ๋ ํ
์คํธ์ ๋์ผํ ์บ์ ๊ฐ์ฒด๋ฅผ ์ ๋ฌํ์ฌ, ์ค๋จํ ๊ณณ์์ ์์ฑ์ ๊ณ์ํฉ๋๋ค.
# ๋ค์ค ํด ๋ํ์ ๊ฒฝ์ฐ, ์์ฑ๋ ํ
์คํธ์ ์๋ก์ด ์ฌ์ฉ์ ์
๋ ฅ์ ์ถ๊ฐํ ์ ์์ต๋๋ค.
new_input_ids = outputs
outputs = model.generate(new_input_ids, past_key_values=past_key_values)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
['The theory of special relativity states 1. The speed of light is constant in all inertial reference frames. 2. The speed of light is constant in all inertial reference frames. 3.']
๋์ผํ [
StaticCache] ๊ฐ์ฒด๋ฅผ ์๋ก์ด ํ๋กฌํํธ์ ์ฌ์ฉํ๋ ค๋ฉด, ํธ์ถ ๊ฐ์.reset()๋ฉ์๋๋ฅผ ์ฌ์ฉํ์ฌ ๊ทธ ๋ด์ฉ์ ์ด๊ธฐํํ๋ ๊ฒ์ด ์ข์ต๋๋ค.
๋ ๊น์ด ๋ค์ด๊ฐ๊ณ ์ถ๋ค๋ฉด, [StaticCache] ๊ฐ์ฒด๋ฅผ ๋ชจ๋ธ์ forward ํจ์ค์ ๋์ผํ past_key_values ์ธ์๋ก ์ ๋ฌํ ์๋ ์์ต๋๋ค. ์ด ์ ๋ต์ ์ฌ์ฉํ๋ฉด, ํ์ฌ ํ ํฐ๊ณผ ์ด์ ์ ์์ฑ๋ ํ ํฐ์ ์์น ๋ฐ ์บ์ ์์น๋ฅผ ๋ฐํ์ผ๋ก ๋ค์ ํ ํฐ์ ๋์ฝ๋ฉํ๋ ์์ฒด ํจ์๋ฅผ ์์ฑํ ์ ์์ต๋๋ค.
from transformers import LlamaTokenizer, LlamaForCausalLM, StaticCache, logging
from transformers.testing_utils import CaptureLogger
import torch
prompts = [
"Simply put, the theory of relativity states that ",
"My favorite all time favorite condiment is ketchup.",
]
NUM_TOKENS_TO_GENERATE = 40
torch_device = "cuda"
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", pad_token="</s>", padding_side="right")
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="sequential")
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
def decode_one_tokens(model, cur_token, input_pos, cache_position, past_key_values):
logits = model(
cur_token,
position_ids=input_pos,
cache_position=cache_position,
past_key_values=past_key_values,
return_dict=False,
use_cache=True
)[0]
new_token = torch.argmax(logits[:, -1], dim=-1)[:, None]
return new_token
StaticCache ๋ฉ์๋๋ฅผ ์ฌ์ฉํ์ฌ ์ ์ kv-cache์ torch.compile์ ํ์ฑํํ๋ ค๋ฉด ๋ช ๊ฐ์ง ์ค์ํ ์์
์ ์ํํด์ผ ํฉ๋๋ค:
- ์ถ๋ก ์ ๋ชจ๋ธ์ ์ฌ์ฉํ๊ธฐ ์ ์ [
StaticCache] ์ธ์คํด์ค๋ฅผ ์ด๊ธฐํํฉ๋๋ค. ์ฌ๊ธฐ์ ์ต๋ ๋ฐฐ์น ํฌ๊ธฐ์ ์ํ์ค ๊ธธ์ด์ ๊ฐ์ ๋งค๊ฐ๋ณ์๋ฅผ ์ค์ ํ ์ ์์ต๋๋ค. - ์ ์ kv-cache์ ํจ๊ป ์์ ํ๋ฅผ ์ปดํ์ผํ๊ธฐ ์ํด ๋ชจ๋ธ์
torch.compile์ ํธ์ถํฉ๋๋ค. - torch.backends.cuda.sdp_kernel ์ปจํ
์คํธ ๊ด๋ฆฌ์์์
enable_math=True๋ฅผ ์ค์ ํ์ฌ ๋ค์ดํฐ๋ธ PyTorch C++ ๊ตฌํ๋ ์ค์ผ์ผ๋ ์ ๊ณฑ ์ดํ ์ (scaled dot product attention)์ ํ์ฑํํ์ฌ ์ถ๋ก ์๋๋ฅผ ๋์ฑ ๋์ ๋๋ค.
batch_size, seq_length = inputs["input_ids"].shape
with torch.no_grad():
past_key_values = StaticCache(
config=model.config, max_batch_size=2, max_cache_len=4096, device=torch_device, dtype=model.dtype
)
cache_position = torch.arange(seq_length, device=torch_device)
generated_ids = torch.zeros(
batch_size, seq_length + NUM_TOKENS_TO_GENERATE + 1, dtype=torch.int, device=torch_device
)
generated_ids[:, cache_position] = inputs["input_ids"].to(torch_device).to(torch.int)
logits = model(
**inputs, cache_position=cache_position, past_key_values=past_key_values,return_dict=False, use_cache=True
)[0]
next_token = torch.argmax(logits[:, -1], dim=-1)[:, None]
generated_ids[:, seq_length] = next_token[:, 0]
decode_one_tokens = torch.compile(decode_one_tokens, mode="reduce-overhead", fullgraph=True)
cache_position = torch.tensor([seq_length + 1], device=torch_device)
for _ in range(1, NUM_TOKENS_TO_GENERATE):
with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True):
next_token = decode_one_tokens(model, next_token.clone(), None, cache_position, past_key_values)
generated_ids[:, cache_position] = next_token.int()
cache_position += 1
text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
text
['Simply put, the theory of relativity states that 1) the speed of light is constant, 2) the speed of light is the same for all observers, and 3) the laws of physics are the same for all observers.',
'My favorite all time favorite condiment is ketchup. I love it on everything. I love it on my eggs, my fries, my chicken, my burgers, my hot dogs, my sandwiches, my salads, my p']
์ ์ฒด generate ํจ์๋ฅผ ์ปดํ์ผํ๋ ๊ฒ์ ์ฝ๋ ์ธก๋ฉด์์ ๊ธฐ๋ณธ ์ฌ์ฉ๋ฒ๋ณด๋ค ๋ ๊ฐ๋จํฉ๋๋ค. generate ํจ์์ ๋ํด torch.compile์ ํธ์ถํ์ฌ ์ ์ฒด ํจ์๋ฅผ ์ปดํ์ผํ๋ฉด ๋ฉ๋๋ค. ์ ์ ์บ์์ ์ฌ์ฉ์ ์ง์ ํ ํ์๋ ์์ต๋๋ค. ์ ์ ์บ์๋ ํธํ๋์ง๋ง, ๋ฒค์น๋งํฌ์์๋ ๋์ ์บ์(๊ธฐ๋ณธ ์ค์ )๊ฐ ๋ ๋น ๋ฅธ ๊ฒ์ผ๋ก ๋ํ๋ฌ์ต๋๋ค.
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false" # ๊ธด ๊ฒฝ๊ณ ๋ฉ์์ง๋ฅผ ๋ฐฉ์งํ๊ธฐ ์ํด ์ค์ :)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
model.generate = torch.compile(model.generate, mode="reduce-overhead", fullgraph=True)
input_text = "The theory of special relativity states "
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
outputs = model.generate(**input_ids)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
['The theory of special relativity states 1. The speed of light is constant in all inertial reference']
์ด ๋ฐฉ๋ฒ์ ํตํด ๋ชจ๋ธ์ forward ํจ์ค๋ฟ๋ง ์๋๋ผ, ์
๋ ฅ ์ค๋น, logit ์ฒ๋ฆฌ๊ธฐ ์์
๋ฑ์ ํฌํจํ ๋ชจ๋ ๊ฒ์ ์ปดํ์ผํฉ๋๋ค. ๊ธฐ๋ณธ ์ฌ์ฉ ์์ ์ ๋นํด generate ํธ์ถ์ด ์ฝ๊ฐ ๋ ๋น ๋ฅผ ์ ์์ผ๋ฉฐ, ์ปดํ์ผ๋ ๊ทธ๋ํ๋ ๋ ํน์ดํ ํ๋์จ์ด ์ฅ์น๋ ์ฌ์ฉ ์ฌ๋ก์ ์ ํฉํ ์ ์์ต๋๋ค. ๊ทธ๋ฌ๋ ์ด ์ ๊ทผ ๋ฐฉ์์ ์ฌ์ฉํ๋ ๋ฐ๋ ๋ช ๊ฐ์ง ํฐ ๋จ์ ์ด ์์ต๋๋ค:
- ์ปดํ์ผ ์๋๊ฐ ํจ์ฌ ๋๋ฆฝ๋๋ค;
generate์ ๋ชจ๋ ๋งค๊ฐ๋ณ์ ์ค์ ์generation_config๋ฅผ ํตํด์๋ง ๊ฐ๋ฅํฉ๋๋ค;- ๋ง์ ๊ฒฝ๊ณ ์ ์์ธ๊ฐ ์ต์ ๋ฉ๋๋ค. -- ๋จผ์ ์ปดํ์ผ ๋์ง ์์ ํํ๋ก ํ ์คํธํ๋ ๊ฒ์ ๊ถ์ฅํฉ๋๋ค;
- ํ์ฌ ์์ ์ค์ด์ง๋ง ๊ธฐ๋ฅ ์ ํ์ด ์ฌํฉ๋๋ค(์: ์์ฑ ์์ ์์๋ EOS ํ ํฐ์ด ์ ํ๋์ด๋ ์์ฑ์ด ์ค๋จ๋์ง ์์ต๋๋ค).
์ถ์ ๋์ฝ๋ฉ [[speculative-decoding]]
๋ณด๋ค ์ฌ์ธต์ ์ธ ์ค๋ช ์ ์ํ๋ค๋ฉด, Assisted Generation: a new direction toward low-latency text generation ๋ธ๋ก๊ทธ ๊ฒ์๋ฌผ์ ํ์ธํ์ญ์์ค!
์๊ธฐ ํ๊ท์ ๋ ๋ค๋ฅธ ๋ฌธ์ ๋ ๊ฐ ์ ๋ ฅ ํ ํฐ์ ๋ํด ์์ ํ ์ค์ ๋ชจ๋ธ ๊ฐ์ค์น๋ฅผ ๋งค๋ฒ ๋ก๋ํด์ผ ํ๋ค๋ ์ ์ ๋๋ค. ์ด๋ ์์ญ์ต ๊ฐ์ ๋งค๊ฐ๋ณ์๋ฅผ ๊ฐ์ง LLM์๋ ๋๋ฆฌ๊ณ ๋ฒ๊ฑฐ๋กญ์ต๋๋ค. ์ถ์ ๋์ฝ๋ฉ(speculative decoding)์ ๋ ์๊ณ ๋น ๋ฅธ ๋ณด์กฐ ๋ชจ๋ธ์ ์ฌ์ฉํ์ฌ ํ๋ณด ํ ํฐ์ ์์ฑํ๊ณ , ์ด๋ฅผ ํฐ LLM์ด ๋จ์ผ ์์ ํ์์ ๊ฒ์ฆํ์ฌ ์ด ์๋ ์ ํ๋ฅผ ์ํํฉ๋๋ค. ๊ฒ์ฆ๋ ํ ํฐ์ด ์ ํํ๋ค๋ฉด, LLM์ ๋ณธ๋ ์์ฒด์ ์ผ๋ก ์์ฑํ๋ ๊ฒ์ฒ๋ผ ํ ํฐ์ ์ป์ ์ ์์ต๋๋ค. ์ ๋ฐฉ ํจ์ค๊ฐ ๋์ผํ ์ถ๋ ฅ์ ๋ณด์ฅํ๊ธฐ ๋๋ฌธ์ ์ ํ๋ ์ ํ๊ฐ ์์ต๋๋ค.
๊ฐ์ฅ ํฐ ์๋ ํฅ์์ ์ป๊ธฐ ์ํด, ๋ณด์กฐ ๋ชจ๋ธ์ ๋น ๋ฅด๊ฒ ํ ํฐ์ ์์ฑํ ์ ์๋๋ก LLM๋ณด๋ค ํจ์ฌ ์์์ผ ํฉ๋๋ค. ๋ณด์กฐ ๋ชจ๋ธ๊ณผ LLM ๋ชจ๋ธ์ ํ ํฐ์ ๋ค์ ์ธ์ฝ๋ฉํ๊ณ ๋์ฝ๋ฉํ์ง ์๋๋ก ๋์ผํ ํ ํฌ๋์ด์ ๋ฅผ ๊ณต์ ํด์ผ ํฉ๋๋ค.
์ถ์ ๋์ฝ๋ฉ์ ํ์ ๊ฒ์๊ณผ ์ํ๋ง ๋์ฝ๋ฉ ์ ๋ต์์๋ง ์ง์๋๋ฉฐ, ๋ฐฐ์น ์ ๋ ฅ์ ์ง์ํ์ง ์์ต๋๋ค.
๋ณด์กฐ ๋ชจ๋ธ์ ๋ก๋ํ๊ณ ์ด๋ฅผ [~GenerationMixin.generate] ๋ฉ์๋์ ์ ๋ฌํ์ฌ ์ถ์ ๋์ฝ๋ฉ์ ํ์ฑํํ์ญ์์ค.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
inputs = tokenizer("Einstein's theory of relativity states", return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b").to(device)
assistant_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m").to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
tokenizer.batch_decode(outputs, skip_special_tokens=True)
["Einstein's theory of relativity states that the speed of light is constant. "]
์ถ์ ์ํ๋ง ๋์ฝ๋ฉ(speculative sampling decoding)์ ์ํด, ๋ณด์กฐ ๋ชจ๋ธ ์ธ์๋ [~GenerationMixin.generate] ๋ฉ์๋์ do_sample ๋ฐ temperature ๋งค๊ฐ๋ณ์๋ฅผ ์ถ๊ฐํ์ญ์์ค.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
inputs = tokenizer("Einstein's theory of relativity states", return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b").to(device)
assistant_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m").to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model, do_sample=True, temperature=0.7)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
["Einstein's theory of relativity states that motion in the universe is not a straight line.\n"]
ํ๋กฌํํธ ์กฐํ ๋์ฝ๋ฉ [[prompt-lookup-decoding]]
ํ๋กฌํํธ ์กฐํ ๋์ฝ๋ฉ์ ํ์ ๊ฒ์๊ณผ ์ํ๋ง๊ณผ๋ ํธํ๋๋ ์ถ์ ๋์ฝ๋ฉ์ ๋ณํ์ ๋๋ค. ํ๋กฌํํธ ์กฐํ๋ ์์ฝ๊ณผ ๊ฐ์ ์ ๋ ฅ ๊ธฐ๋ฐ ์์ ์ ํนํ ์ ์๋ํฉ๋๋ค. ์ฌ๊ธฐ์๋ ํ๋กฌํํธ์ ์ถ๋ ฅ ๊ฐ์ ์ข ์ข ๊ฒน์น๋ ๋จ์ด๊ฐ ์์ต๋๋ค. ์ด๋ฌํ ๊ฒน์น๋ n-๊ทธ๋จ์ด LLM ํ๋ณด ํ ํฐ์ผ๋ก ์ฌ์ฉ๋ฉ๋๋ค.
ํ๋กฌํํธ ์กฐํ ๋์ฝ๋ฉ์ ํ์ฑํํ๋ ค๋ฉด prompt_lookup_num_tokens ๋งค๊ฐ๋ณ์์ ๊ฒน์น๋ ํ ํฐ ์๋ฅผ ์ง์ ํ์ญ์์ค. ๊ทธ๋ฐ ๋ค์ ์ด ๋งค๊ฐ๋ณ์๋ฅผ [~GenerationMixin.generate] ๋ฉ์๋์ ์ ๋ฌํ ์ ์์ต๋๋ค.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
inputs = tokenizer("The second law of thermodynamics states", return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b").to(device)
assistant_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m").to(device)
outputs = model.generate(**inputs, prompt_lookup_num_tokens=3)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
['The second law of thermodynamics states that entropy increases with temperature. ']
์ํ๋ง๊ณผ ํจ๊ป ํ๋กฌํํธ ์กฐํ ๋์ฝ๋ฉ์ ์ฌ์ฉํ๋ ค๋ฉด, [~GenerationMixin.generate] ๋ฉ์๋์ do_sample ๋ฐ temperature ๋งค๊ฐ๋ณ์๋ฅผ ์ถ๊ฐํ์ญ์์ค.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
inputs = tokenizer("The second law of thermodynamics states", return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b").to(device)
outputs = model.generate(**inputs, prompt_lookup_num_tokens=3, do_sample=True, temperature=0.7)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
["The second law of thermodynamics states that energy cannot be created nor destroyed. It's not a"]
์ดํ ์ ์ต์ ํ [[attention-optimizations]]
ํธ๋์คํฌ๋จธ ๋ชจ๋ธ์ ์๋ ค์ง ๋ฌธ์ ๋ ์ ํ ์ดํ ์ ๋ฉ์ปค๋์ฆ์ด ์ ๋ ฅ ํ ํฐ ์์ ํจ๊ป ๊ณ์ฐ ๋ฐ ๋ฉ๋ชจ๋ฆฌ๊ฐ ์ ๊ณฑ์ผ๋ก ์ฆ๊ฐํ๋ค๋ ๊ฒ์ ๋๋ค. ์ด ์ ํ์ ํจ์ฌ ๋ ๊ธด ์ํ์ค๋ฅผ ์ฒ๋ฆฌํ๋ LLM์์๋ ๋์ฑ ์ปค์ง๋๋ค. ์ด๋ฅผ ํด๊ฒฐํ๊ธฐ ์ํด FlashAttention2 ๋๋ PyTorch์ ์ค์ผ์ผ๋ ์ ๊ณฑ ์ดํ ์ ์ ์ฌ์ฉํด ๋ณด์ญ์์ค. ์ด๋ค์ ๋ ๋ฉ๋ชจ๋ฆฌ ํจ์จ์ ์ธ ์ดํ ์ ๊ตฌํ์ผ๋ก ์ถ๋ก ์ ๊ฐ์ํํ ์ ์์ต๋๋ค.
FlashAttention-2 [[flashattention-2]]
FlashAttention๊ณผ FlashAttention-2๋ ์ดํ ์ ๊ณ์ฐ์ ๋ ์์ ์ฒญํฌ๋ก ๋๋๊ณ ์ค๊ฐ ์ฝ๊ธฐ/์ฐ๊ธฐ ์์ ์ ์ค์ฌ ์ถ๋ก ์๋๋ฅผ ๋์ ๋๋ค. FlashAttention-2๋ ์๋ FlashAttention ์๊ณ ๋ฆฌ์ฆ์ ๊ฐ์ ํ์ฌ ์ํ์ค ๊ธธ์ด ์ฐจ์์์๋ ๋ณ๋ ฌ ์ฒ๋ฆฌ๋ฅผ ์ํํ๊ณ ํ๋์จ์ด์์ ์์ ์ ๋ ์ ๋ถํ ํ์ฌ ๋๊ธฐํ ๋ฐ ํต์ ์ค๋ฒํค๋๋ฅผ ์ค์ ๋๋ค.
FlashAttention-2๋ฅผ ์ฌ์ฉํ๋ ค๋ฉด [~PreTrainedModel.from_pretrained] ๋ฉ์๋์์ attn_implementation="flash_attention_2"๋ฅผ ์ค์ ํ์ญ์์ค.
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
quant_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
"google/gemma-2b",
quantization_config=quant_config,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)
PyTorch ์ค์ผ์ผ๋ ์ ๊ณฑ ์ดํ ์ (scaled dot product attention) [[pytorch-scaled-dot-product-attention]]
์ค์ผ์ผ๋ ์ ๊ณฑ ์ดํ ์ (SDPA)๋ PyTorch 2.0์์ ์๋์ผ๋ก ํ์ฑํ๋๋ฉฐ, FlashAttention, xFormers, PyTorch์ C++ ๊ตฌํ์ ์ง์ํฉ๋๋ค. SDPA๋ CUDA ๋ฐฑ์๋๋ฅผ ์ฌ์ฉํ๋ ๊ฒฝ์ฐ ๊ฐ์ฅ ์ฑ๋ฅ์ด ์ข์ ์ดํ ์ ์๊ณ ๋ฆฌ์ฆ์ ์ ํํฉ๋๋ค. ๋ค๋ฅธ ๋ฐฑ์๋์์๋ SDPA๊ฐ PyTorch C++ ๊ตฌํ์ผ๋ก ๊ธฐ๋ณธ ์ค์ ๋ฉ๋๋ค.
SDPA๋ ์ต์ PyTorch ๋ฒ์ ์ด ์ค์น๋์ด ์์ผ๋ฉด FlashAttention-2๋ ์ง์ํฉ๋๋ค.
์ธ ๊ฐ์ง ์ดํ
์
์๊ณ ๋ฆฌ์ฆ ์ค ํ๋๋ฅผ ๋ช
์์ ์ผ๋ก ํ์ฑํํ๊ฑฐ๋ ๋นํ์ฑํํ๋ ค๋ฉด torch.backends.cuda.sdp_kernel ์ปจํ
์คํธ ๊ด๋ฆฌ์๋ฅผ ์ฌ์ฉํ์ญ์์ค. ์๋ฅผ ๋ค์ด FlashAttention์ ํ์ฑํํ๋ ค๋ฉด enable_flash=True๋ก ์ค์ ํ์ญ์์ค.
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"google/gemma-2b",
torch_dtype=torch.bfloat16,
)
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
outputs = model.generate(**inputs)
์์ํ [[quantization]]
์์ํ๋ LLM ๊ฐ์ค์น๋ฅผ ๋ ๋ฎ์ ์ ๋ฐ๋๋ก ์ ์ฅํ์ฌ ํฌ๊ธฐ๋ฅผ ์ค์ ๋๋ค. ์ด๋ ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋์ ์ค์ด๋ฉฐ GPU ๋ฉ๋ชจ๋ฆฌ์ ์ ์ฝ์ด ์๋ ๊ฒฝ์ฐ ์ถ๋ก ์ ์ํด LLM์ ๋ก๋ํ๋ ๊ฒ์ ๋ ์ฉ์ดํ๊ฒ ํฉ๋๋ค. GPU๊ฐ ์ถฉ๋ถํ๋ค๋ฉด, ๋ชจ๋ธ์ ์์ํํ ํ์๋ ์์ต๋๋ค. ์ถ๊ฐ์ ์ธ ์์ํ ๋ฐ ์์ํ ํด์ ๋จ๊ณ๋ก ์ธํด ์ฝ๊ฐ์ ์ง์ฐ์ด ๋ฐ์ํ ์ ์๊ธฐ ๋๋ฌธ์ ๋๋ค(AWQ ๋ฐ ์ตํฉ AWQ ๋ชจ๋ ์ ์ธ).
๋ค์ํ ์์ํ ๋ผ์ด๋ธ๋ฌ๋ฆฌ(์์ธํ ๋ด์ฉ์ Quantization ๊ฐ์ด๋๋ฅผ ์ฐธ์กฐํ์ญ์์ค)๊ฐ ์์ต๋๋ค. ์ฌ๊ธฐ์๋ Quanto, AQLM, VPTQ, AWQ ๋ฐ AutoGPTQ๊ฐ ํฌํจ๋ฉ๋๋ค. ์ฌ์ฉ ์ฌ๋ก์ ๊ฐ์ฅ ์ ๋ง๋ ๋ผ์ด๋ธ๋ฌ๋ฆฌ๋ฅผ ์ฌ์ฉํด ๋ณด์ญ์์ค. ๋ํ AutoGPTQ์ bitsandbytes๋ฅผ ๋น๊ตํ๋ Overview of natively supported quantization schemes in ๐ค Transformers ๋ธ๋ก๊ทธ ๊ฒ์๋ฌผ์ ์ฝ์ด๋ณด๋ ๊ฒ์ ์ถ์ฒํฉ๋๋ค.
์๋์ ๋ชจ๋ธ ๋ฉ๋ชจ๋ฆฌ ๊ณ์ฐ๊ธฐ๋ฅผ ์ฌ์ฉํ์ฌ ๋ชจ๋ธ์ ๋ก๋ํ๋ ๋ฐ ํ์ํ ๋ฉ๋ชจ๋ฆฌ๋ฅผ ์ถ์ ํ๊ณ ๋น๊ตํด ๋ณด์ญ์์ค. ์๋ฅผ ๋ค์ด Mistral-7B-v0.1๋ฅผ ๋ก๋ํ๋ ๋ฐ ํ์ํ ๋ฉ๋ชจ๋ฆฌ๋ฅผ ์ถ์ ํด ๋ณด์ญ์์ค.
Mistral-7B-v0.1์ ๋ฐ์ ๋ฐ๋๋ก ๋ก๋ํ๋ ค๋ฉด [~transformers.AutoModelForCausalLM.from_pretrained] ๋ฉ์๋์์ torch_dtype ๋งค๊ฐ๋ณ์๋ฅผ torch.bfloat16์ผ๋ก ์ค์ ํ์ญ์์ค. ์ด ๊ฒฝ์ฐ 13.74GB์ ๋ฉ๋ชจ๋ฆฌ๊ฐ ํ์ํฉ๋๋ค.
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
model = AutoModelForCausalLM.from_pretrained(
"mistralai/Mistral-7B-v0.1", torch_dtype=torch.bfloat16, device_map="auto",
)
์ถ๋ก ์ ์ํด ์์ํ๋ ๋ชจ๋ธ(8๋นํธ ๋๋ 4๋นํธ)์ ๋ก๋ํ๋ ค๋ฉด bitsandbytes๋ฅผ ์ฌ์ฉํ๊ณ load_in_4bit ๋๋ load_in_8bit ๋งค๊ฐ๋ณ์๋ฅผ True๋ก ์ค์ ํ์ญ์์ค. ๋ชจ๋ธ์ 8๋นํธ๋ก ๋ก๋ํ๋ ๋ฐ๋ 6.87GB์ ๋ฉ๋ชจ๋ฆฌ๋ง ํ์ํฉ๋๋ค.
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch
quant_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
"mistralai/Mistral-7B-v0.1", quantization_config=quant_config, device_map="auto"
)