실습 6 — LLM 은 가지치기에 얼마나 견디나

import gc, time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16).to("cuda")
ids = tok(open("shop_corpus.txt", encoding="utf-8").read()[:3000], return_tensors="pt")["input_ids"][:, :512].to("cuda")
enc = tok.apply_chat_template([{"role": "user", "content": "택배 상자가 젖어서 왔어요. 어떻게 해야 하나요?"}],
                              add_generation_prompt=True, return_tensors="pt", return_dict=True).to("cuda")


def check(m):
    with torch.no_grad():
        ppl = float(torch.exp(m(ids, labels=ids).loss))
        torch.cuda.synchronize()
        t = time.perf_counter()
        out = m.generate(**enc, max_new_tokens=40, min_new_tokens=40, do_sample=False)
        torch.cuda.synchronize()
    answer = tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).replace("\n", " ")
    return round(ppl, 2), round(40 / (time.perf_counter() - t), 1), answer[:50]


originals = {n: mod.weight.detach().clone() for n, mod in model.named_modules()
             if isinstance(mod, torch.nn.Linear) and "lm_head" not in n}
check(model)                                                     # 데우기
print("선형층:", len(originals), "개 · 기준:", check(model))

rows = []
for amount in (0.1, 0.3, 0.5):
    for n, mod in model.named_modules():
        if n in originals:
            w = originals[n]
            k = int(w.numel() * amount)
            threshold = w.abs().flatten().float().kthvalue(k).values      # 층마다 아래 amount 만큼의 경계값
            mod.weight.data = torch.where(w.abs() > threshold, w, torch.zeros_like(w))
    ppl, tps, answer = check(model)
    rows.append({"지운 몫": f"{amount:.0%}", "퍼플렉서티": ppl, "초당토큰": tps, "답": answer})
for n, mod in model.named_modules():                             # 원래대로 되돌린다
    if n in originals:
        mod.weight.data = originals[n]
import pandas as pd
print(pd.DataFrame(rows).drop(columns="답").to_string(index=False))
for r in rows:
    print(r["지운 몫"], "→", r["답"])

zeros = sum((mod.weight == 0).sum().item() for n, mod in model.named_modules() if n in originals)
size_mb = sum(p.numel() * p.element_size() for p in model.parameters()) / 1e6
print(f"되돌린 모델의 0 인 가중치: {zeros:,} · 메모리: {size_mb:.0f}MB")
print("50% 를 지운 모델도 같은 꼴의 행렬이라 메모리는 같다. 0 을 건너뛰는 희소 형식과 계산이 있어야 준다")
del model
gc.collect()
torch.cuda.empty_cache()