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

GitHub의 slvDev/esp32-ai 프로젝트는 ESP32-S3라는 약 $8의 비용이 드는 마이크로컨트롤러에서 28.9백만 파라미터의 언어 모델을 실행하는 방법을 소개합니다. 이 모델은 서버로 데이터를 전송하지 않고 칩 자체에서 실행되며, 칩에 연결된 작은 화면에 초당 약 9개의 토큰을 생성하여 출력합니다. 이전에 이와 유사한 칩에서 실행된 언어 모델은 26만 개의 파라미터를 가졌으나, 이 모델은 약 100배 더 많은 파라미터를 포함하고 있습니다. 이는 대부분의 모델이 RAM 대신 플래시에 저장되기 때문에 가능한 것으로, Google의 Gemma 모델에서 사용된 Per-Layer Embeddings라는 아이디어를 활용했습니다.

기술 사양

파라미터: 28.9M (이 중 25M은 플래시 조회 테이블에 저장)
칩: ESP32-S3, 약 $8, 512KB SRAM, 8MB PSRAM, 16MB 플래시
속도: 초당 약 9.5 토큰 (순수 계산 속도는 초당 9.7 토큰)
연결성: 없음, 모든 것이 장치에서 실행
모델 크기: 14.9MB (4비트)

어려움과 해결 방법

마이크로컨트롤러는 매우 적은 빠른 메모리를 가지고 있습니다. ESP32-S3는 512KB의 SRAM을 제공합니다. 일반적으로 전체 모델은 이 메모리에서 접근 가능해야 하므로, 작은 모델에 국한될 수밖에 없습니다. 이전 모델이 26만 개의 파라미터를 가졌던 이유도 여기에 있습니다. 이를 해결하기 위해 모델을 빠른 메모리에 두지 않고, 대부분의 언어 모델 파라미터가 임베딩 테이블에 위치하도록 했습니다. 모델은 이 테이블을 읽어들이며, 각 토큰이 필요로 하는 몇 개의 행만을 불러옵니다. 이렇게 하면 큰 모델을 실행하는 데 거의 비용이 들지 않습니다. 이 아이디어는 Google의 Gemma 모델의 Per-Layer Embeddings에서 비롯된 것입니다. 이 방법은 전화기나 GPU가 아닌 마이크로컨트롤러의 메모리 레이아웃에 적용되었습니다.

SRAM (빠르고 작음): "생각하는" 코어, 모든 토큰에 사용
PSRAM (중간): 출력 헤드와 작업 메모리
FLASH (크고 느림): 25M 파라미터 테이블, 토큰당 약 6개의 행 읽기 (~450B)

모델의 기능과 한계

이 모델은 TinyStories라는 데이터셋으로 훈련되어 짧고 간단한 이야기를 생성하며, 대체로 일관성을 유지합니다. 그러나 질문에 답하거나, 명령을 따르거나, 코드를 작성하거나, 사실을 아는 등의 기능은 수행하지 못합니다. 이는 모델의 추론을 담당하는 작은 부분의 한계 때문이며, 메모리 트릭이 이를 변경하지는 않습니다. 여기서 흥미로운 점은 28.9백만 파라미터 모델이 무엇을 말할 수 있는지가 아니라, 작은 칩에 큰 모델을 적합하게 만드는 아키텍처입니다.

직접 실행하기

펌웨어, 배선, 플래싱 단계는 firmware/esp32_llm/README.md에 있으며, 훈련, 제거, 양자화 코드는 src/와 experiments/에 있습니다. 전체 방법론, 제거 과정, 칩 내 측정은 RESULTS.md에 기록되어 있습니다.

크레딧

TinyStories는 작은 모델이 일관되게 글을 작성할 수 있도록 충분히 간단한 합성 이야기로 구성된 데이터셋입니다. 이 데이터셋은 Microsoft Research의 Ronen Eldan과 Yuanzhi Li가 arXiv에 발표한 것입니다. Per-Layer Embeddings는 Google의 Gemma 모델에서 가져온 설계로, 큰 모델을 작은 칩에 적합하게 만듭니다. Andrej Karpathy의 llama2.c는 많은 사람들이 작은 언어 모델을 훈련하고 순수 C로 실행할 수 있다고 믿게 만든 이유입니다. 이 프로젝트는 그로부터 발전한 것입니다.

프로젝트 진행 과정

저는 의도적으로 저장소에 복잡한 기록을 남겼습니다. 여기에는 초기 숫자를 부풀렸던 제 자신의 파라미터 계산 버그와 이를 수정한 후의 결과가 포함되어 있습니다. 커밋 기록과 RESULTS.md는 숫자가 어떻게 이동했는지와 그 이유를 보여줍니다.

파이랩 정리