배경
신경망 모델의 파라미터들은 보통 정수가 아니고 실수로 되어 있습니다. 32비트 컴퓨터 시대가 열린 뒤로 (지금은 64비트 시대입니다만) 실수값들을 보통 32비트로 나타내고 있는데요. 신경망 학습시에 32비트까지 필요하지는 않고 16비트나 8비트로 줄여서 학습하는 것이 요즘 추세입니다. 이렇게 비트를 줄이면 신경망을 저장하는 메모리를 절약할 수 있고, 학습 속도도 빨라집니다.
8비트는 아직 NVIDIA H100 등 특수한 GPU에서만 지원해서 서서히 보급되는 중이고, 16비트는 많이 보급됐습니다. 16비트 모델을 저장하는 방법이나 불러오는 방법에 대한 글도 있는데, 꽤 간단한 방법임에도 모아서 설명하는 글이 없어서 지금 씁니다.
쓰기
허깅페이스에서 모델을 저장하는 방법은 매우 간단합니다. 아래 한 줄이 전부입니다. (state_dict만 저장하는 방법이 있지만 여기서는 생략합니다)
model.save_pretrained(save_path)그러면 16비트 모델은 어떻게 저장할까요? 아래와 같이 모델을 변환하고 저장합니다
model = model.half()
model.save_pretrained(save_path)위는 fp16 형식인 경우입니다. bf16이라면 아래와 같이 변환합니다
model = model.bfloat16()
model.save_pretrained(save_path)모델이 이미 변환된 상태라면 추가 변환은 필요 없습니다만, 모호한 상황에서는 명시적으로 변환하는 것이 좋겠습니다.
읽기
허깅페이스에서 모델을 불러오는 방법도 매우 간단합니다
from transformers import AutoModel
model = AutoModel.from_pretrained(save_path)이렇게 모델을 불러와서 타입을 확인해 보면 torch.float32라고 나옵니다
print(model.dtype)
torch.float32허깅페이스 모델의 config.json을 읽어보면 이 모델이 어떤 데이터 타입인지 적혀 있습니다만, 모델을 읽을 때에 별다른 옵션을 주지 않으면 기본적으로 fp32로 읽습니다. 16비트로 읽으려면 torch_dtype='auto 옵션이 필요합니다
from transformers import AutoModel
model = AutoModel.from_pretrained(save_path, torch_dtype='auto')
print(model.dtype)
torch.float16이것은 fp16으로 저장했을 때의 예시이고, bf16으로 저장했다면 torch.bfloat16으로 출력될 것입니다.
GPU를 사용하는 모델이면 CPU를 거치지 않고 GPU로 빠르게 읽도록 low_cpu_mem_usage=True 옵션도 추가해 주는 것이 좋습니다
model = AutoModel.from_pretrained(save_path, torch_dtype='auto', low_cpu_mem_usage=True)