MLX-LM — Apple 실리콘에서 LLM 실행과 파인튜닝

MLX-LM — Apple 실리콘에서 LLM 실행과 파인튜닝

MLX-LM은 Apple 실리콘 위에서 대규모 언어모델(LLM)을 실행·파인튜닝하기 위한 Python 패키지예요. Hugging Face 허브와 통합되어 있어 수천 개의 모델을 한 줄 명령으로 쓸 수 있죠.

출처: https://github.com/ml-explore/mlx-lm

설치

pip install mlx-lm

conda로는 conda install -c conda-forge mlx-lm 도 가능해요.

텍스트 생성과 채팅

명령줄에서 바로 생성하거나 채팅을 할 수 있어요.

mlx_lm.generate --prompt "How tall is Mt Everest?"
mlx_lm.chat

기본 모델은 mlx-community/Llama-3.2-3B-Instruct-4bit이고, --model 플래그로 다른 MLX 호환 모델을 지정할 수 있어요. 옵션을 보려면 mlx_lm.generate -h를 쓰면 됩니다.

Python API

mlx_lm을 모듈로 불러와서 load와 generate를 쓰면 돼요.

from mlx_lm import load, generate

model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")

prompt = "Write a story about Einstein"
messages = [{"role": "user", "content": prompt}]
prompt = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True,
)

text = generate(model, tokenizer, prompt=prompt, verbose=True)

스트리밍 생성은 stream_generate를 쓰면 응답 객체를 순차로 받아요.

from mlx_lm import load, stream_generate

model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
prompt = "Write a story about Einstein"
messages = [{"role": "user", "content": prompt}]
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)

for response in stream_generate(model, tokenizer, prompt, max_tokens=512):
    print(response.text, end="", flush=True)
print()

모델 양자화와 업로드

convert API로 4bit 양자화된 모델을 만들고 허브에 올릴 수 있어요.

from mlx_lm import convert

repo = "mistralai/Mistral-7B-Instruct-v0.3"
upload_repo = "mlx-community/My-Mistral-7B-Instruct-v0.3-4bit"
convert(repo, quantize=True, upload_repo=upload_repo)

명령줄로는 mlx_lm.convert --model mistralai/Mistral-7B-Instruct-v0.3 -q 입니다.

긴 프롬프트와 생성 다루기

긴 문맥을 효율적으로 다루는 도구도 제공돼요.

  • 회전형 고정 크기 KV 캐시: --max-kv-size n 으로 조절. 작은 값(512)은 RAM을 적게 쓰지만 품질이 떨어지고, 큰 값(4096+)은 품질이 좋지만 RAM을 더 써요.
  • 프리필 스텝 크기: --prefill-step-size n 기본값은 2048.
  • 프롬프트 캐싱: mlx_lm.cache_prompt 로 긴 프롬프트를 캐시하고 --prompt-cache-file 로 재사용.

더 알아보기