Executando um LLM de 28,9 milhões de parâmetros em um microcontrolador de US$ 8

• Executa no ESP32-S3 um modelo de linguagem de 28,9 milhões de parâmetros sem conexão com servidor, exibindo texto em uma pequena tela a cerca de 9 tokens por segundo
• Usa uma arquitetura de Per-Layer Embeddings que mantém 25 milhões do total de parâmetros em flash lenta e lê apenas cerca de 6 linhas necessárias por token, aproximadamente 450 bytes
• O modelo tem 14,9 MB em 4 bits; o núcleo de computação usado a cada token fica nos 512 KB de SRAM, a cabeça de saída e a memória de trabalho ficam nos 8 MB de PSRAM, e a grande tabela de embeddings fica nos 16 MB de flash
• Treinado com TinyStories, gera histórias curtas e simples de forma geralmente coerente, mas não é adequado para perguntas e respostas, execução de comandos, escrita de código ou conhecimento factual
• Contém cerca de 100 vezes mais parâmetros do que um modelo de 260 mil parâmetros executado anteriormente em um chip semelhante; o ponto central está na arquitetura de memória que acomoda um modelo grande em um chip pequeno, e não na qualidade de geração

---

Hardware e desempenho de execução

• Todo o processamento é feito em um único ESP32-S3, de cerca de US$ 8, sem enviar dados a servidores
  • Usa 512 KB de SRAM, 8 MB de PSRAM e 16 MB de flash
  • A velocidade total é de cerca de 9,5 tok/s, e a velocidade de computação pura é de cerca de 9,7 tok/s
  • O tamanho do modelo em 4 bits é de 14,9 MB
• Dos 28,9 milhões de parâmetros no total, 25 milhões ficam armazenados em uma tabela de consulta na flash

Como colocar o modelo em uma memória pequena

• Em geral, todo o modelo precisa estar acessível em memória rápida, mas a SRAM do ESP32-S3 tem apenas 512 KB, o que só comporta modelos muito pequenos
• Aproveita o fato de que a maior parte dos parâmetros está em uma tabela de embeddings, e não na parte computacional, mantendo essa tabela na flash
  • Lê apenas cerca de 6 linhas necessárias por token, aproximadamente 450 bytes
  • Mantém em memória rápida apenas a pequena parte responsável pela computação real
  • A maior parte do modelo não é carregada durante a execução; apenas as partes necessárias são selecionadas a partir da flash
• As funções da memória são divididas assim
  • SRAM: núcleo de computação usado por todos os tokens
  • PSRAM: cabeça de saída e memória de trabalho
  • flash: tabela com 25 milhões de parâmetros

Aplicação de Per-Layer Embeddings

• Aplica os Per-Layer Embeddings usados pelo Gemma 3n e Gemma 4, do Google, à arquitetura de memória de um microcontrolador, e não de um celular ou GPU
• Pelo que o criador do projeto verificou, não há precedentes de aplicação desse método a um chip tão pequeno

O que o modelo consegue fazer e suas limitações

• Aprendeu com as histórias sintéticas curtas do TinyStories, gerando histórias simples e mantendo coerência na maior parte do tempo
• Não consegue responder perguntas, executar comandos, escrever código nem fornecer conhecimento factual
• Essa limitação vem do pequeno núcleo responsável pela inferência, e a técnica de alocação de memória não aumenta a capacidade de inferência em si

Código e materiais dos experimentos

• firmware/esp32_llm/README.md contém o firmware, a fiação e o procedimento de flashing
• src/ e experiments/ incluem o código de treinamento, experimentos de ablação e quantização
• RESULTS.md resume o método completo, os experimentos de ablação e as medições on-chip

Projetos de base e histórico

• TinyStories é um conjunto de dados de histórias sintéticas curtas criado para que modelos pequenos também possam aprender escrita coerente
• Os Per-Layer Embeddings do Google Gemma serviram como base para colocar um modelo grande em um chip pequeno
• llama2.c, de Andrej Karpathy, influenciou a abordagem de treinar um pequeno modelo de linguagem e executá-lo em C puro
• O repositório também preserva um erro de cálculo que inflou a contagem inicial de parâmetros e o processo de correção
  • O histórico de commits e o RESULTS.md mostram onde os números foram alterados e por quê

[GN⁺]

Comentários no Hacker News

• É surpreendente o que dá para fazer hoje com um microcontrolador de US$ 5. Comprei cinco placas Milk-V, do modelo Duo, que tem até 256 MB de memória e uma TPU de 1 TOPS@INT8, além de rodar Linux

  • Não dá para deixar de mencionar a ISA vetorial de 128 bits e os 32 registradores. Ela dá suporte a inteiros e ponto flutuante de até 64 bits e, com LMUL=8, processa 1.024 bits em uma única instrução; a maioria das operações leva 3 ciclos por 128 bits
    GCC e Clang dão suporte completo a xTHeadVector e, usando intrínsecos em C, ficam compatíveis com RVV 1.0 apenas com opções de linha de comando. Muito código que lida com elementos de 8 bits, como memcpy(), memset(), memcmp(), strlen(), strcpy() e strcmp(), é binariamente compatível em boa parte
    Quando comprei o Duo de 64 MB, ele custava US$ 3; depois os modelos de 64 MB, 256 MB e 512 MB custavam US$ 5, US$ 7 e US$ 10, respectivamente, mas os preços subiram bastante este ano: https://arace.tech/products/milk-v-duo, https://arace.tech/products/milkv-duo-s
  • Também dá para comprar microcontroladores por menos de US$ 0,50, mas, mesmo em um produto de US$ 5, as restrições para rodar modelos parecem tão severas que o ambiente é praticamente inútil. Espero que a demanda por modelos não cause escassez de MCUs
  • Não é preciso pensar que, como se microcontroladores estivessem presos ao hardware dos anos 1980, assembly ou C sejam as únicas opções. Um dispositivo desse porte conseguiria rodar até um sistema operacional do Xerox PARC sem grande dificuldade
    Ainda há usos em que um PIC de 4 KB é a melhor opção, mas não devemos agir como se, na maioria dos casos, não houvesse alternativas melhores
  • Esse preço e desempenho fazem perceber o quanto a posição monopolista da ARM impôs custos enormes ao setor inteiro

• Modelos de conversão entre voz e texto também estão chegando perto desse tamanho, então fico curioso sobre quão próximos estamos de dispositivos pequenos capazes de conversar conosco. Pode vir um mundo em que a escova de dentes dê conselhos de higiene bucal — ou até faça propaganda de pasta de dente

  • Conselhos de higiene bucal já são possíveis. As escovas de dentes topo de linha da Philips se conectam por Bluetooth e fornecem feedback no app
    Mas não quero viver em um mundo em que seja preciso procurar um bloqueador de anúncios para a IA da escova de dentes

• Eles usaram de forma muito inteligente uma técnica de embeddings por camada. Como também existem modelos práticos de TTS com cerca de 20 a 30 milhões de parâmetros, um ESP32 sem conexão de rede poderia praticamente ler textos em tempo real

  • Tenho vontade de testar uma configuração em que cada MCU fica responsável por uma camada do modelo. A ideia seria conectar vários RP2350 com linhas dedicadas via PIO e combinar interpoladores com instruções de multiplicação dupla
    PSRAM, flash e cartão SD não têm muita largura de banda individualmente, mas, acionando muitos deles ao mesmo tempo, dá para atingir uma vazão considerável. Hardware dedicado em grande escala vai levar vantagem em desempenho por watt, mas essa configuração também é atraente pelo baixo custo inicial e pela escalabilidade incremental

• Gostaria de saber quais são as opções realistas para que um LLM local no Raspberry Pi 4, em vez de um microcontrolador, não leve 30 segundos por resposta

• Rodar um LLM em um dispositivo minúsculo é legal, mas o que impressiona ainda mais é o método de treinamento que produziu esses pesos

  • Em faixas pequenas de parâmetros, é difícil fazer otimizações excessivas visando apenas benchmarks, então a utilidade dos rankings ainda não acabou. Pelo contrário: eles ainda têm o papel de encontrar modelos menores que atinjam o mesmo desempenho nos benchmarks, levando a etapa de modelagem a reconhecer e remover compressibilidade e redundância

• Fico curioso se, aproveitando a flash com bons padrões de acesso, não seria possível escalar isso para rodar modelos muito maiores na CPU

  • Houve uma tentativa parecida recentemente. Rodaram o GLM-5.2 em uma CPU comum de desktop, SSD NVMe e cerca de 25 GB de RAM, mas a velocidade não era em tokens por segundo, e sim de até 20 segundos por token: https://github.com/JustVugg/colibri
  • A largura de banda da flash do ESP32 é apenas cerca de um quarto da SRAM interna. Em sistemas mais potentes, essa diferença fica muito maior, e, para aproveitar de forma eficiente o aumento de capacidade de computação, a largura de banda de memória a ser fornecida também precisa crescer bastante

• O ESP32-S3 é bem poderoso e atualmente o uso no trabalho de desenvolvimento com Raspberry Pi 4. Uma de suas duas portas USB dá suporte a OTG, então dá para implementar funcionalidades que, de outra forma, custariam mais de US$ 100

  • Aceitando algumas limitações, dá para montar algo parecido com o RP2350 por cerca de US$ 1. A velocidade máxima é USB Full Speed, de 12 Mbps; uma porta usa o periférico USB integrado ao chip, e a outra é conectada a pinos GPIO com suporte por PIO
    Atualmente funciona com tinyusb e pico-pio-usb, e também estou experimentando um port para Rust na expectativa de desempenho maior

• Gostaria de saber qual é a precisão desse modelo quantizado

• O desempenho que ele mostra nesse tamanho é surpreendente, e fico animado para ver o que será possível nos computadores de placa única um pouco mais potentes mencionados na thread