MacBook Pro에서 5분 안에 훈련할 수 있는 가장 강력한 AI 모델은?

TL;DR Highlight

MacBook Pro에서 5분 만에 GPT 스타일 transformer를 훈련하여 최적의 모델 크기·데이터셋·학습 설정을 탐색하고 로컬 환경의 한계 내에서 실제로 효과가 있는 요소를 실측 데이터로 입증했다.

Who Should Read

작은 규모의 언어 모델을 직접 훈련해보거나 로컬 환경에서 ML 실험을 빠르게 반복하고 싶은 개발자. 특히 PyTorch/MLX 기반 모델 훈련을 처음 시도하거나 최적화 포인트를 찾고 있는 ML 엔지니어에게 유용하다.

Core Mechanics

실험의 최종 결과는 약 180만 파라미터(1.8M param)짜리 GPT 스타일 transformer를 TinyStories 데이터셋의 약 2천만 토큰으로 5분 훈련한 것이며, held-out split 기준 perplexity(모델이 텍스트를 얼마나 잘 예측하는지 나타내는 지표, 낮을수록 좋음) 약 9.6을 달성했다.
5분이라는 시간 제약 하에서는 '큰 모델을 적은 토큰으로' 훈련하는 것보다 '작은 모델을 많은 토큰으로' 훈련하는 전략이 훨씬 효과적이다. 1B param 모델에 토큰 4,000개를 넣는 것보다 1M param 모델에 4M 토큰을 넣는 것이 낫다.
반대로 너무 작아도 안 된다. 10K param 수준의 초소형 모델은 30초 만에 training loss가 plateau(더 이상 줄지 않는 상태)에 도달하고 영어 문법 자체를 학습하지 못해 gibberish(의미 없는 텍스트)만 출력했다.
MacBook의 MPS(Apple Silicon GPU 가속)를 쓰면 초당 약 3,000 토큰을 처리할 수 있다. 그런데 torch.compile, float16 변환, gradient accumulation(여러 배치에 걸쳐 누적 후 가중치 업데이트하는 기법) 등 일반적인 최적화 기법들이 오히려 속도를 떨어뜨리거나 의미 없는 수준이었다. 병목이 GPU 연산량이 아니라 GPU에 작업을 지시하는 'launch 횟수' 자체에 있기 때문이다.
PyTorch 대신 Apple Silicon에 최적화된 MLX 프레임워크로 전환해도 기대했던 2배 속도 향상은 없었다. 결론적으로 추천 설정은 MPS 사용, compile/quantize/gradient accumulation 모두 건너뛰기, 가능한 한 작은 모델 선택이다.
데이터셋 선택이 모델 품질에 결정적인 영향을 미쳤다. Simple English Wikipedia로 훈련하면 영어 문법은 맞지만 고유명사에 집착하는 이상한 출력이 나왔다. 10MB 분량의 다양한 주제를 섞으면 일관된 패턴이 없어 모델이 학습하기 어렵기 때문이다.
TinyStories(어린이 동화 스타일의 단순한 영어 텍스트 데이터셋)로 바꾸자 출력 품질이 눈에 띄게 향상됐다. 단순하고 일관된 문체 덕분에 적은 토큰으로도 패턴을 효과적으로 학습할 수 있었다.
최종 추천 조합은 1.8M param 모델 + TinyStories 약 20M 토큰 + MPS 가속 + 단순 훈련 루프(최적화 기법 없음)이며, 이 설정으로 5분 안에 읽을 수 있는 수준의 아동 동화 문장을 생성하는 모델을 만들 수 있었다.

Evidence

GPT-2 speedrun 프로젝트(modded-nanogpt)의 기법들을 적용하면 더 나은 결과를 얻을 수 있을 것이라는 의견이 있었다. 구체적으로 Muon optimizer(기울기 업데이트 방식을 개선한 옵티마이저), 더 나은 가중치 초기화, 학습률 튜닝 등을 조합하면 같은 시간 안에 더 낮은 perplexity를 달성할 수 있을 것이라는 제안이다.
'시간' 대신 '에너지(Joule)'를 제약 조건으로 삼아야 MBP와 H100을 공정하게 비교할 수 있다는 의견도 나왔다. 전력 효율 관점에서 보면 MBP가 상대적으로 경쟁력 있을 수 있다는 시각이다.
소형 모델 연구의 중요성에 대한 흥미로운 시각이 공유됐다. 효모(yeast)가 복잡한 생물학 연구를 위한 단순 모델 생물로 쓰이듯, 소형 transformer에서 관찰되는 행동이 대형 LLM의 동작을 이해하는 데 도움이 될 수 있다는 것이다. 즉 이 실험은 단순한 재미를 넘어 LLM 해석 가능성 연구에도 의미가 있다.
실용적 활용 가능성에 대한 논의도 있었다. '크고 무거운 범용 모델 대신, 특정 도메인에 집중한 소형 모델을 필요에 따라 on-demand로 만들어 쓰면 속도와 효율 면에서 유리하다'는 의견이 공감을 얻었다. 이 논의 도중 마침 Google이 Gemma 3 270M을 출시해 화제가 되기도 했다.
비슷한 연구로 'Cramming' 논문(arXiv:2212.14034)이 언급됐다. 이 논문은 현대 노트북에서 하루 동안 훈련할 수 있는 최선의 모델을 찾는 실험으로, 이번 글과 비슷한 문제의식을 공유한다. 또한 Karpathy의 zero-to-hero 시리즈처럼 Jupyter 노트북에서도 유용한 모델을 훈련할 수 있다는 점에서, 'AI 훈련=대규모 클러스터'라는 통념을 다시 생각하게 한다는 반응도 있었다.

How to Apply

로컬에서 빠른 모델 훈련 실험을 하려는 경우, torch.compile이나 float16 변환 같은 '일반적인 최적화'를 맹목적으로 적용하지 말고 먼저 MPS만 활성화한 단순 훈련 루프로 baseline을 측정하라. Apple Silicon에서는 launch overhead가 병목이기 때문에 복잡한 최적화가 오히려 역효과를 낼 수 있다.
작은 도메인 특화 모델을 빠르게 만들고 싶다면 데이터셋의 '일관성'에 집중하라. 여러 주제를 섞은 잡다한 데이터보다 TinyStories처럼 문체와 주제가 일관된 데이터가 적은 토큰으로도 훨씬 잘 학습된다. 특정 업무(예: 특정 도메인의 텍스트 요약이나 분류)에 맞는 단순하고 일관된 데이터셋을 직접 구성하면 소형 모델도 실용적인 성능을 낼 수 있다.
모델 크기와 훈련 토큰 수의 trade-off를 직접 실험해보려면, 같은 시간 안에 '크고 짧게' vs '작고 길게' 두 가지 설정을 비교해보라. 이 글의 결론처럼 시간 제약이 있는 환경에서는 파라미터 수를 줄이고 더 많은 토큰을 보여주는 쪽이 대부분 유리하다.
더 나은 결과를 원한다면 modded-nanogpt(https://github.com/KellerJordan/modded-nanogpt) 레포지토리의 Muon optimizer와 가중치 초기화 기법을 참고해 적용해보라. GPT-2 speedrun 커뮤니티에서 검증된 방법들이므로 소형 모델에서도 의미 있는 개선이 기대된다.

Terminology

perplexity모델이 텍스트를 얼마나 잘 예측하는지 나타내는 지표. 낮을수록 좋으며, 완벽한 예측이면 1에 가까워진다. 쉽게 말해 모델이 다음 단어를 얼마나 '당황하지 않고' 맞히는가를 수치화한 것.

MPSApple Silicon(M1/M2/M3/M4 등) 칩의 GPU를 PyTorch에서 사용하기 위한 백엔드. NVIDIA의 CUDA에 해당하는 Apple 환경의 GPU 가속 인터페이스다.

gradient accumulation메모리가 부족할 때 여러 미니배치의 기울기(gradient)를 누적한 뒤 한 번에 가중치를 업데이트하는 기법. 큰 배치 크기 효과를 메모리 적게 쓰면서 흉내 낼 수 있지만 연산 횟수가 늘어난다.

TinyStories어린아이가 이해할 수 있는 수준의 단순한 영어 단어와 문법으로 구성된 합성 동화 데이터셋. 소형 언어 모델 훈련 실험에 자주 사용된다.

torch.compilePyTorch 2.0에서 도입된 기능으로, 모델 연산 그래프를 미리 컴파일해서 실행 속도를 높이는 최적화 기법. 대형 모델에서는 효과적이지만 매우 작은 모델에서는 컴파일 오버헤드가 더 클 수 있다.

MLXApple이 개발한 Apple Silicon 전용 머신러닝 프레임워크. PyTorch와 유사한 API를 제공하면서 M 시리즈 칩의 통합 메모리 구조를 활용하도록 설계됐다.