$8 마이크로컨트롤러에서 28.9M 파라미터 LLM 실행하기

이것은 약 $8 가격의 마이크로컨트롤러(microcontroller)인 ESP32-S3에서 텍스트를 생성하는 2,890만(28.9 million) 파라미터 언어 모델(language model)입니다. 이 모델은 서버로 아무것도 보내지 않고 칩 자체에서 실행되며, 초당 약 9개의 토큰(tokens) 속도로 칩에 연결된 작은 화면에 각 단어를 작성합니다. 사람들이 이와 같은 칩에서 실행했던 마지막 언어 모델은 26만 파라미터였으므로, 이 모델은 그보다 약 100배 더 많은 파라미터를 보유하고 있습니다. 모델의 대부분이 RAM 대신 플래시(flash)에 저장되기 때문에 가능하며, 이는 Google의 Gemma 모델에서 사용된 Per-Layer Embeddings라는 아이디어를 활용한 것입니다.

수치 (The numbers)

왜 어려운지, 그리고 어떻게 가능한지

마이크로컨트롤러는 빠른 메모리(fast memory)가 매우 적습니다. ESP32-S3는 512KB의 SRAM을 제공합니다. 보통 모델 전체가 SRAM에서 접근 가능해야 하므로 아주 작은 모델에 머물 수밖에 없으며, 이것이 이와 같은 칩에서 실행되었던 이전 모델이 단 26만 파라미터였던 이유입니다.

이를 해결하는 방법은 모델을 빠른 메모리에 아예 넣지 않는 것입니다. 언어 모델 파라미터의 대부분은 임베딩 테이블(embedding table)에 위치하며, 모델은 이를 계산하기보다는 읽어 들입니다. 따라서 2,500만 행의 테이블을 느린 플래시(flash)에 두고, 실제 작업을 수행하는 작은 부분은 빠른 메모리에 유지하면서 각 토큰이 필요로 하는 약 450바이트의 몇몇 행만 가져올 수 있습니다. 모델의 대부분을 로드하지 않기 때문에 대규모 모델을 실행하는 데 드는 비용이 거의 들지 않습니다. 모델은 그저 플래시(flash)에 머물러 있다가 조금씩 샘플링(sampled)될 뿐입니다.

그 아이디어는 Gemma 3n 및 Gemma 4에서 사용된 Google의 Per-Layer Embeddings입니다. 여기서는 스마트폰이나 GPU 대신 마이크로컨트롤러 (microcontroller)의 메모리 레이아웃 (memory layout)에서 실행됩니다. 제가 알기로는, 이렇게 작은 칩에서 이를 시도한 사람은 아무도 없었습니다.

  SRAM  (빠름, 매우 작음)   모든 토큰(token)에서 사용되는 "사고(thinking)" 코어
  PSRAM (중간)           출력 헤드(output head) 및 작업 메모리
  FLASH (거대함, 느림)    25M 파라미터 테이블, 토큰당 약 6개 행 읽기 (~450 B)

기능 및 한계

이 모델은 TinyStories 데이터셋으로 학습되었으므로, 짧고 단순한 이야기를 작성하며 대체로 일관성을 유지합니다. 질문에 답하거나, 지시를 따르거나, 코드를 작성하거나, 사실 관계를 알지는 못합니다. 그러한 한계는 추론 (reasoning)을 수행하는 모델의 작은 부분에서 기인하며, 메모리 트릭 (memory trick)이 이를 바꾸지는 못합니다. 여기서 흥미로운 점은 2,890만 파라미터 모델이 무엇을 말할 수 있느냐가 아니라, 거대한 모델을 아주 작은 칩에 맞추는 아키텍처 (architecture)입니다.

직접 실행해보기

펌웨어 (firmware), 배선 (wiring), 그리고 플래싱 (flashing) 단계는 firmware/esp32_llm/README.md에 설명되어 있습니다. 학습 (training), 절제 연구 (ablation), 그리고 양자화 (quantization) 코드는 src/ 및 experiments/에 있습니다. 전체 방법론, 절제 연구, 그리고 온칩 (on-chip) 측정 결과는 RESULTS.md에 정리되어 있습니다.

크레딧 (Credit)

TinyStories는 이 모델이 학습하는 데이터셋입니다. 이는 작은 모델도 일관성 있게 쓰는 법을 배울 수 있을 만큼 충분히 단순한 짧은 합성 이야기들입니다 (Ronen Eldan 및 Yuanzhi Li, Microsoft Research, arXiv:2305.07759). 나머지 절반은 Google의 Gemma 모델에서 가져온 설계인 Per-Layer Embeddings이며, 이것이 큰 모델을 작은 칩에 맞출 수 있게 해줍니다.

Andrej Karpathy의 llama2.c 덕분에 저를 포함한 많은 사람들이 아주 작은 언어 모델을 학습시키고 순수 C 언어로 실행할 수 있다는 것을 믿게 되었습니다. 이 프로젝트는 거기에서 파생되었습니다.

실제 진행 과정

저는 복잡한 히스토리를 의도적으로 리포지토리 (repo)에 남겨두었습니다. 여기에는 제 파라미터 계산 과정에서 발견한 버그로 인해 초기 수치가 부풀려졌던 점과, 이를 수정한 후 도출된 교정된 결과가 포함되어 있습니다. 커밋 히스토리 (commit history)와 RESULTS.md를 통해 수치가 어떻게 변했는지, 그리고 왜 변했는지를 확인할 수 있습니다.