실습 4 — 정밀도별 메모리 비교하기

# 같은 모델을 서로 다른 숫자 정밀도로 불러와 가중치 메모리 사용량을 비교합니다.
# BitsAndBytesConfig = bitsandbytes 라이브러리로 4비트·8비트 양자화할 때의 설정표
from transformers import BitsAndBytesConfig  # 4비트 양자화 설정 도구를 불러옵니다.

bnb_config = BitsAndBytesConfig(  # 4비트 NF4 양자화 방식을 설정합니다.
    # load_in_4bit=True: 선형층 가중치를 4비트로 줄여 읽는다
    load_in_4bit=True,  # 선형층 가중치를 4비트로 읽습니다.
    # bnb_4bit_quant_type="nf4": 정규분포에 맞춘 16개 눈금(NormalFloat4)을 쓴다
    bnb_4bit_quant_type="nf4",  # 가중치의 양자화 형식입니다.
    # bnb_4bit_use_double_quant=True: 스케일 값까지 한 번 더 양자화한다(Double Quantization)
    bnb_4bit_use_double_quant=True,  # 양자화 상수도 압축합니다.
    # bnb_4bit_compute_dtype: 계산할 때 잠깐 되돌릴 정밀도(이 PC 는 BF16)
    bnb_4bit_compute_dtype=MODEL_DTYPE,  # GPU에 맞는 계산 정밀도입니다.
)

footprints = {}  # 정밀도별 메모리 사용량을 저장합니다.
# torch.float32 = FP32(4바이트), MODEL_DTYPE = BF16(2바이트, BF16 을 못 쓰는 GPU 는 FP16)
for label, dtype in [("FP32", torch.float32), ("BF16 / FP16", MODEL_DTYPE)]:  # 두 정밀도를 비교합니다.
    precision_model = load_base_model(dtype=dtype)  # 현재 정밀도로 모델을 불러옵니다.
    # get_memory_footprint() = 가중치가 차지하는 바이트 수(memory footprint, 메모리 발자국)
    footprints[label] = precision_model.get_memory_footprint()  # 가중치 메모리를 기록합니다.
    del precision_model  # 비교가 끝난 모델 참조를 삭제합니다.
    gc.collect()  # 파이썬 객체를 정리합니다.
    torch.cuda.empty_cache()  # GPU 캐시를 비웁니다.

# quantization_config=bnb_config: 위의 4비트 설정으로 읽는다
# device_map={"": 0}: 모델 전체("")를 0번 GPU 에 바로 올린다(4비트 모델은 .to(device) 대신 이렇게 둔다)
q4_model = load_base_model(dtype=MODEL_DTYPE, quantization_config=bnb_config, device_map={"": 0})  # 같은 모델을 NF4 4비트로 불러옵니다.
footprints["4bit (NF4)"] = q4_model.get_memory_footprint()  # 4비트 메모리를 기록합니다.
# 첫 층의 q_proj 를 찍어 본다. 결과: Linear4bit(in_features=896, out_features=896, bias=True)
print(q4_model.model.layers[0].self_attn.q_proj)  # 4비트 선형층의 형태를 확인합니다.
del q4_model  # 4비트 비교 모델 참조를 삭제합니다.
gc.collect()  # 파이썬 객체를 정리합니다.
torch.cuda.empty_cache()  # 다음 실습을 위해 GPU 캐시를 비웁니다.
# display = 노트북 전용 함수: 표나 Series 를 노트북 화면에 예쁘게 보여 준다(.py 파일에서는 print 를 쓴다)
# 1024**2 로 나눠 MB 로, :,.0f 는 천 단위 쉼표와 소수점 없이
# 결과: FP32 1,885 MB / BF16 / FP16 942 MB / 4bit (NF4) 430 MB
display(pd.Series({label: f"{footprint / 1024**2:,.0f} MB" for label, footprint in footprints.items()}, name="메모리"))