실습 2 — 같은 모델을 fp16, 8비트, 4비트로 불러오기

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

# Ollama 가 GPU 에 올려 둔 모델이 있으면 내려 GPU 를 비운다
for m in json.loads(urllib.request.urlopen("http://localhost:11434/api/ps").read())["models"]:
    urllib.request.urlopen(urllib.request.Request("http://localhost:11434/api/generate", headers={"Content-Type": "application/json"},
                                                  data=json.dumps({"model": m["name"], "keep_alive": 0}).encode())).read()
name = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(name)
ids = tok(open("shop_corpus.txt", encoding="utf-8").read()[:3000], return_tensors="pt")["input_ids"][:, :512].to("cuda")
QUESTION = "택배 상자가 젖어서 왔어요. 어떻게 해야 하나요?"


def measure(model):
    """메모리, 퍼플렉서티, 생성 속도, 답을 잰다. 생성은 한 번 데워 본 뒤 잰다"""
    with torch.no_grad():
        ppl = float(torch.exp(model(ids, labels=ids).loss))
    enc = tok.apply_chat_template([{"role": "user", "content": QUESTION}], add_generation_prompt=True,
                                  return_tensors="pt", return_dict=True).to("cuda")
    model.generate(**enc, max_new_tokens=8, do_sample=False)                    # 데우기
    torch.cuda.synchronize()
    t = time.perf_counter()
    out = model.generate(**enc, max_new_tokens=60, min_new_tokens=60, do_sample=False)   # 60토큰을 꼭 만든다
    torch.cuda.synchronize()
    sec = time.perf_counter() - t
    answer = tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).replace("\n", " ")
    return {"메모리MB": round(model.get_memory_footprint() / 1e6), "최대GPU MB": round(torch.cuda.max_memory_allocated() / 1e6),
            "퍼플렉서티": round(ppl, 3), "초당토큰": round(60 / sec, 1), "답": answer[:45]}


print("GPU:", torch.cuda.get_device_name(0), f"· {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")

configs = {
    "fp16": {"dtype": torch.float16},
    "8비트": {"dtype": torch.float16, "quantization_config": BitsAndBytesConfig(load_in_8bit=True)},
    # NF4: 정규분포 모양의 가중치에 맞춘 4비트 눈금. 계산은 fp16 으로 되돌려서 한다
    "4비트 NF4": {"dtype": torch.float16, "quantization_config": BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16)},
}
rows = {}
for label, kw in configs.items():
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    model = AutoModelForCausalLM.from_pretrained(name, device_map="cuda", **kw)
    rows[label] = measure(model)
    del model
    gc.collect()
    torch.cuda.empty_cache()
import pandas as pd
table = pd.DataFrame(rows).T
print(table.drop(columns="답").to_string())

for label, r in rows.items():
    print(f"{label}: {r['답']}")