실습 6 — GPU 메모리 셈하기

import matplotlib.pyplot as plt
import pandas as pd

plt.rcParams["font.family"] = "Malgun Gothic"        # 한글 글꼴(4단계의 그래프)
plt.rcParams["axes.unicode_minus"] = False

# 크기: (파라미터 수(십억), 층 수, KV 헤드 수, 헤드 크기)
MODELS = {"0.5B": (0.49, 24, 2, 64), "1.5B": (1.54, 28, 2, 128), "3B": (3.09, 36, 2, 128), "7B": (7.61, 28, 4, 128),
          "14B": (14.7, 48, 8, 128), "32B": (32.5, 64, 8, 128), "72B": (72.7, 80, 8, 128)}


def weights_gb(params_b, bits):
    return params_b * bits / 8                 # 십억 × 비트 ÷ 8 = GB


# Ollama 의 Q4_K_M 은 가중치마다 평균 5비트쯤 쓴다(3B 파일 1.93GB ÷ 30.9억 파라미터, 3단계에서 확인)
table = pd.DataFrame({bits_name: {m: weights_gb(v[0], bits) for m, v in MODELS.items()}
                      for bits_name, bits in [("fp16", 16), ("8비트", 8), ("4비트(Q4_K_M)", 5)]})
# 결과: 크기와 정밀도마다 가중치가 차지하는 메모리(GB)
print(table.round(1))

def kv_gb(layers, kv_heads, head_dim, tokens, users, bytes_=2):
    return 2 * layers * kv_heads * head_dim * bytes_ * tokens * users / 1e9


rows = []
for m, (p, layers, kvh, hd) in MODELS.items():
    per_token_kb = 2 * layers * kvh * hd * 2 / 1e3
    rows.append({"모델": m, "토큰 하나 KB": round(per_token_kb, 1),
                 **{f"{u}명×4,096": kv_gb(layers, kvh, hd, 4096, u) for u in [1, 8, 32]}})
kv = pd.DataFrame(rows).set_index("모델")
# 결과: 문맥 4,096 토큰을 동시에 몇 명이 쓰느냐에 따른 KV 캐시(GB, fp16)
print(kv.round(2))

import json
import urllib.request

import torch


def get(path):
    with urllib.request.urlopen("http://127.0.0.1:11434" + path, timeout=30) as r:
        return json.loads(r.read())


# 모델을 올려 둔다(문맥 4,096 토큰). 올라간 뒤 /api/ps 가 GPU 에 올린 크기를 알려 준다
req = urllib.request.Request("http://127.0.0.1:11434/api/generate", data=json.dumps(
    {"model": "qwen2.5:3b", "prompt": "안녕", "stream": False, "options": {"num_predict": 1, "num_ctx": 4096}}).encode(),
    headers={"Content-Type": "application/json"})
urllib.request.urlopen(req, timeout=120).read()
file_size = next(m["size"] for m in get("/api/tags")["models"] if m["name"] == "qwen2.5:3b")
loaded = next(m for m in get("/api/ps")["models"] if m["name"] == "qwen2.5:3b")
p, layers, kvh, hd = MODELS["3B"]
estimate = file_size / 1e9 + kv_gb(layers, kvh, hd, 4096, 1)
print(f"파일 {file_size / 1e9:.2f}GB → 가중치 하나에 {file_size * 8 / (p * 1e9):.2f}비트")
print(f"셈: 파일 + KV(4,096토큰) = {estimate:.2f}GB · Ollama 가 GPU 에 올린 크기 {loaded['size_vram'] / 1e9:.2f}GB")
# 결과: 이 PC 의 GPU 와 전체 메모리
props = torch.cuda.get_device_properties(0)
print(f"{props.name} · {props.total_memory / 1e9:.1f}GB")

GPUS = {"GTX 1650": 4, "RTX 4090": 24, "L40S": 48, "A100·H100": 80}
need = pd.DataFrame({m: {"4비트, 1명": weights_gb(v[0], 5) + kv_gb(*v[1:], 4096, 1),
                         "4비트, 8명": weights_gb(v[0], 5) + kv_gb(*v[1:], 4096, 8),
                         "fp16, 8명": weights_gb(v[0], 16) + kv_gb(*v[1:], 4096, 8)} for m, v in MODELS.items()}).T
need = need * 1.2                              # 계산 여유 20%
print(need.round(1))
# 결과: 경우마다 한 장에 들어가는 가장 작은 GPU
print(need.apply(lambda col: col.map(lambda g: next((n for n, c in GPUS.items() if g <= c), "여러 장"))))

ax = need.plot(kind="bar", logy=True, figsize=(9, 4), rot=0)
for name, cap in GPUS.items():
    ax.axhline(cap, linestyle="--", color="gray", linewidth=0.8)
    ax.text(6.6, cap, name, va="bottom", ha="right", fontsize=8)
ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v:g}"))
ax.set_ylabel("필요한 GPU 메모리(GB, 로그 눈금)")
ax.set_title("모델 크기와 사용자 수에 따른 GPU 메모리")
plt.tight_layout()
plt.show()