실습 3 — 묶어서 생성하면 얼마나 빨라지나

import json, time, urllib.request
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Ollama 가 올려 둔 모델을 내려 GPU 를 비운다
for m in json.loads(urllib.request.urlopen("http://127.0.0.1:11434/api/ps").read())["models"]:
    urllib.request.urlopen(urllib.request.Request("http://127.0.0.1:11434/api/generate",
                           data=json.dumps({"model": m["name"], "keep_alive": 0}).encode()))

name = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(name)
# padding_side="left": 길이가 다른 문장을 묶을 때 빈칸을 왼쪽에 채운다(생성은 오른쪽 끝에서 이어 가므로)
tok.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16).to("cuda").eval()

questions = [json.loads(l)["질문"] for l in open("company_rules_qa.jsonl", encoding="utf-8")]
questions = (questions * 2)[:32]                     # 32개가 되게 두 번 이어 붙인다
texts = [tok.apply_chat_template([{"role": "user", "content": q}], tokenize=False, add_generation_prompt=True)
         for q in questions]
print(torch.cuda.get_device_name(0), "·", len(texts), "개 질문")

NEW = 64                                              # 요청마다 새로 만들 토큰 수


def gen_batch(batch):
    """문장 묶음을 한 번의 generate 로 만들고 걸린 시간을 잰다"""
    x = tok(batch, return_tensors="pt", padding=True).to("cuda")
    torch.cuda.synchronize()                          # GPU 일이 끝날 때까지 기다린 뒤 시각을 잰다
    start = time.perf_counter()
    with torch.no_grad():
        # min_new_tokens 와 max_new_tokens 를 같게: 모든 요청이 꼭 64토큰을 만들게 해 공정하게 잰다
        out = model.generate(**x, max_new_tokens=NEW, min_new_tokens=NEW, do_sample=False)
    torch.cuda.synchronize()
    return time.perf_counter() - start, out[:, x["input_ids"].shape[1]:]


gen_batch(texts[:2])                                  # 몸풀기: 첫 generate 는 준비 시간이 섞인다
sec, out = gen_batch(texts[:1])
# 결과: 요청 하나를 만드는 데 걸린 시간과 앞부분
print(f"{sec:.2f}초 ·", tok.decode(out[0], skip_special_tokens=True)[:60])

import pandas as pd

rows = []
for bs in [1, 2, 4, 8, 16, 32]:
    torch.cuda.reset_peak_memory_stats()               # 묶음마다 가장 많이 쓴 메모리를 새로 잰다
    sec, _ = gen_batch(texts[:bs])
    rows.append({"묶음 크기": bs, "걸린 시간(초)": round(sec, 2), "처리량(토큰/초)": round(bs * NEW / sec, 1),
                 "요청 하나의 지연(초)": round(sec, 2), "최대 메모리(MB)": torch.cuda.max_memory_allocated() // 2**20})
table = pd.DataFrame(rows)
# 결과: 묶음이 커질수록 처리량이 늘고, 지연과 메모리는 조금씩 는다
print(table.to_string(index=False))

import matplotlib.pyplot as plt
plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(table["묶음 크기"], table["처리량(토큰/초)"], marker="o")
ax.set_xscale("log", base=2)                         # 1, 2, 4, … 가 같은 간격으로 보이게 한다
ax.set_xticks(table["묶음 크기"]); ax.set_xticklabels(table["묶음 크기"])
ax.set_xlabel("묶음 크기(한 번에 생성한 요청 수)"); ax.set_ylabel("처리량(토큰/초)")
ax.set_title("묶어서 생성할 때의 처리량")
plt.tight_layout()
plt.show()

# 다음 실습의 Ollama 가 GPU 를 쓸 수 있게 모델을 지우고 메모리를 돌려준다
del model
torch.cuda.empty_cache()