Google의 Gemma 2 모델 Mac Silicon 환경에서 구동하기

 

 
 
기존의 LLM 모델들은 CUDA 외의 다른 GPU 환경 지원을 잘 하지 않았지만, Mac의 Silicon 프로세서가 대중화되면서 Mac에서도 LLM을 사용할 수 있도록 환경이 많이 개선되었습니다. 이번 포스트에서는 Google의 Gemma 모델 중 크기가 가장 작은 2b 모델을 Mac 환경에서 구동해보고, 답변이 잘 출력되는지 간단히 테스트해보겠습니다.
 

먼저, 프로젝트 폴더 구성은 다음과 같이 설정할 것입니다. models 폴더 안에 모델명으로 폴더가 생성되고, 이 안에 모델 파일들이 내려받아지게 됩니다.
.env 파일은 보안이 필요한 Huggingface 토큰을 저장하는 용도이며, gemma-2b-it.ipynb 노트북에서 모델을 불러와 추론을 수행합니다.
 
 

myllm
└───models
│   └───gemma-2b-it
│       │   file01
│       │   file02
│       │    ...
└───gemma-2b-it.ipynb  
└───.env

 
 
이제, gemma-2b-it.ipynb 로 가서 코드를 작성해봅시다. 라이브러리들을 불러오고 Huggingface 토큰(HF_TOKEN)도 불러옵니다. Huggingface 토큰은 Huggingface 웹사이트 로그인 후 프로필에서 'Access Token' 메뉴로 가면 특별한 절차 없이 발급받을 수 있습니다. 이 토큰은 일부 모델을 다운받을 때 요구되므로, 발급 후 보관해놓는 것이 좋습니다.
 
 

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import snapshot_download
from huggingface_hub import login
from IPython.display import Markdown, display
login(token=os.environ.get("HF_TOKEN"), add_to_git_credential=True)

 
 
env 파일은 다음과 같은 형태입니다.
 
 

# .env
HF_TOKEN="hf_YOUR_TOKEN"

 
 
이제 모델을 다운받습니다. 여러 방법이 있지만, snapshot_download를 이용하여 다운받는 것이 간편합니다. repo_id에 gemme-2b-it 모델의 저장소 주소를 입력합니다. 모델 저장소 주소는 각 모델의 페이지를 참고하면 됩니다.
 
gemma-2b-it 모델 페이지
 
local_dir 옵션을 설정하면 해당 경로에 모델을 다운받습니다. 
 
 

snapshot_download(repo_id="google/gemma-2-2b-it",
                  local_dir='models/gemma-2-2b-it/')

 
 
다운로드 경로 설정 시 앞에 '/'를 붙이지 않는 것이 좋습니다. 일부 환경에서 오류가 발생할 수 있기 때문입니다. 예를 들면, local_dir='/models/gemma-2-2b-it/' 로 설정하면 권한 에러 등이 발생할 수 있습니다.
이제 해당 명령어를 실행하면, 다운로드가 진행되는 모습을 볼 수 있습니다.
 
 

 
 
Gemma 2b 모델은 약 5GB의 크기를 가집니다. Gemma는 모델 크기로도 종류가 나누어져 있지만, 같은 모델 크기 중에서도 일반 모델(gemma-2b)과 지시 기반 미세조정(instruction fine tuning) 모델(gemma-2b-it) 두 종류가 있습니다. 사전학습을 추가로 시킬 목적이 아닌 downstream task 수행이 목적이라면 gemma-2b-it 모델을 사용하면 되겠습니다.
 
이제, 다운받은 모델을 transformers 라이브러리로 불러와보겠습니다. AutoModel과 AutoTokenizer 객체에 다운받은 모델의 경로를 설정해줍니다.

이 때 Mac의 경우, device를 설정할 때 보통 "cuda"로 설정하던 부분을 "mps"로 바꾸어주기만 하면 됩니다. mps는 Metal Performance Shaders의 약자로, MacOS의 GPU를 모델 학습이나 추론에 이용할 수 있게 하는 백엔드 모듈입니다. torch.backends.mps.is_available()로 mps 사용 가능 여부를 판단한 후, 사용 가능하면 mps를 사용하도록 합니다.
 
Huggingface MPS 설명 문서
Pytorch MPS 설명 문서
 
 

device = "mps" if torch.backends.mps.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
    "models/gemma-2-2b-it",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map=device,
    )
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")

 
 
 
코드가 정상적으로 잘 작동한다면, 다음과 같이 모델 checkpoint를 불러오게 됩니다.
 
 

 
 
이제 모델이 잘 작동하는지 테스트해볼 단계입니다. 먼저, 입력과 출력을 표시할 간단한 함수를 구성합니다.
 
 

def display_messages(messages):
    for i in messages:
        if i['role'] == 'user':
            display_text = "**[User]** \n" + i['content']
            display(Markdown(display_text))
        elif i['role'] == 'assistant':
            display_text = "**[Gemma]** \n" + i['content']
            display(Markdown(display_text))

 
 
다음은 대화 기록(chat history)를 관리할 messages 객체를 설정합니다. 여기에 사용자(user)의 질문과 모델의 답변(assistant)를 저장하게 되며, 모델은 이 대화 기록을 참고하여 다음 답변을 생성하게 됩니다.
 
 

messages = []

 
 
사용자가 입력하는 인터페이스를 설정하고, 이를 messages에 더해줍니다. 이 셀을 실행하면 입력창이 나타나고, 질문을 입력하면 이것이 messages 객체에 더해집니다. 각 message는 'role(시스템/사용자/모델답변 구분)'과 'content(내용)'를 지정하여 더해집니다. 이러한 형식은 모델마다 다를 수 있으니 모델 페이지를 참고하기 바랍니다.
 
*일반적으로 messages 객체 최초 생성 시 'system' role로 system prompt를 먼저 설정합니다. 하지만, Gemma 모델은 system prompt를 지원하지 않으므로 빈 객체로 시작합니다. System prompt를 설정하고 싶다면, 사용자 입력에 별도로 지시문을 추가하여 설정할 수 있습니다(참고 링크).
 
 

query = input()
messages.append({"role" : "user", "content" : query})

 
 
이제 tokenizer에 apply_chat_template 기능을 이용하여 질문이 포함된 messages 객체를 입력해줍니다. 이 입력값도 설정한 device에 할당해주어야 합니다. prompt_length는 답변을 decoding할 때 지난 대화 기록이 포함되지 않게 하기 위해 필요합니다.
 
 

encoded = tokenizer.apply_chat_template(messages, return_tensors='pt').to(device)
prompt_length = encoded.shape[1]

 
 
이제 생성을 수행하고, 답변을 확인해봅시다. max_new_tokens는 답변의 최대 길이를, skip_special_tokens는 모델 내부적으로만 이용되는 여러 토큰들(예: <bos>, <eos>, <end_of_turn>)이 답변에 포함되지 않도록 합니다. 이 외에도 답변 생성에 많은 옵션들을 설정할 수 있으나, 여기서는 간단한 설정으로만 답변을 생성해봅시다.
 
 

outputs = model.generate(encoded, 
                         max_new_tokens=500)
response = tokenizer.decode(outputs[0][prompt_length:], skip_special_tokens=True)
messages.append({"role" : "assistant", "content" : response})
display_messages(messages)

 
 
아래와 같이 답변이 잘 생성되는 것을 볼 수 있습니다. Gemma 모델은 아직 한국어 지원을 하지 않기 때문에 지원하는 언어로 질문을 해야 합니다. 
 
 

 
 
대화를 계속 이어나가고 싶다면, messages 객체를 유지한 상태로 query = input() 이 포함된 셀로 되돌아가 다시 질문을 입력하고 생성을 수행하면 됩니다. 지난 대화 기록을 참고하여 새로운 답변을 생성할 수 있는지도 테스트해보았을 때(Can you summarize your last answer?), 잘 수행하는 것도 볼 수 있습니다.
 
몇 가지 질문을 수행해본 결과, 기본적인 답변은 그 규모에 비해 상당히 잘 생성하는 것으로 보이나 세부적인 답변 디테일은 다소 부족한 모습도 있습니다. 특정 분야에 잘 조정하여 사용한다면 효율적으로 LLM 앱을 개발할 수 있을 것 같습니다. 
 
참고로, asitop과 같은 모듈을 이용해 Mac에서 LLM을 운용할 때 메모리 사용량을 실시간으로 체크할 수 있습니다. 저의 환경에서는 확인 결과 10GB 정도의 메모리를 차지합니다. 작은 모델이라고 해도 메모리 사용량이 만만치 않습니다. 어떻게 LLM 운용을 더 최적할 수 있을 지에 대해서도 계속 연구해보아야습니다.