ESP32s3-LLM-Cluster

Architectuur

Dit project draait een gesliced 0,5B LLM over een cluster van 7 ESP32-S3's. Eén node fungeert als master en de overige nodes fungeren als compute nodes. De master node voert de tokenizer en embedding uit, terwijl de attention layer en MLP op de compute nodes draaien. De master en de nodes communiceren via een high-speed SPI Daisy-Chain.

Dataflow Diagram

MASTER NODE

  • [ Prompt ] → BPE Tokenizer
  • Token Embedding (INT4, ~14MB in Flash)
  • (SPI CH A - Verzenden naar Node 1)
  • Output: Hidden State Vector (FP32) $\downarrow$

COMPUTE NODE 1

  • (SPI CH B - Ontvangen van Master)
  • Layer 0 tot 3 (4x Transformer Blocks):
  • RMSNorm (FP16 geschaald naar FP32)
  • 1,58-bit Attention (Q, K, V, O proj) + RoPE
  • KV Cache (PSRAM)
  • 1,58-bit MLP (Gate, Up, Down proj)
  • (SPI CH A - Verzenden naar Node 2) $\downarrow$

... (Nodes 2 tot 5) $\downarrow$

COMPUTE NODE 6

  • (SPI CH B - Ontvangen van Node 5)
  • Layer 20 tot 23 (4x Transformer Blocks):
  • Zelfde 1,58-bit architectuur
  • (SPI CH A - Terugverzenden naar Master) $\downarrow$

MASTER NODE

  • (SPI CH B - Ontvangen van Node 6)
  • Final RMS Norm (FP16, 64KB in 'fnorm' partitie)
  • LM Head (gekoppeld aan INT4 Embeddings)
  • Greedy Sampling
  • [ Output ] $\leftarrow$ Next Token ID

Aan de slag

Raadpleeg de workflow-gids (workflow.md) om met het project te beginnen.

Projectstructuur

  • README.md: Projectdocumentatie.
  • workflow.md: Stap-voor-stap gids voor flashen, modelvoorbereiding en bedrading.
  • .gitignore: Git-negeerregels voor build-bestanden en binaries.
  • docs/images/: Architectuurdiagrammen en hardwarefoto's.
  • master_board/: Firmware voor de Master Node (ESP-IDF).
  • main/main.cpp: Master-orchestrator, gebruikers-I/O & BPE tokenizer.
  • main/embedding.cpp: INT4 embedding lookup-logica.
  • main/lm_head.cpp: LM Head mapping en greedy sampling.
  • main/spi_bus.cpp: Master dual-channel SPI-driver.
  • partitions.csv: Aangepaste partitietabel (token, model, fnorm).
  • CMakeLists.txt: Build-configuratie.
  • node_firmware/: Firmware voor de Compute Nodes (ESP-IDF).
  • main/main.cpp: Entry point voor de node worker & inference loop.
  • main/bitlinear.cpp: Implementatie van de 1,58-bit ternaire lineaire laag.
  • main/bitlinear_forward.S: Assembly-geoptimaliseerde MAC-operaties voor 1,58-bit.
  • main/qwen_attention.cpp: Qwen Attention, RoPE & KV-Cache runtime.
  • main/lut_table.cpp: Look-up tables voor extreme optimalisatie.
  • main/spi_bus.cpp: Daisy-chain SPI DMA receiver/transmitter.
  • partitions.csv: Layer-partitie layout voor de node.
  • CMakeLists.txt: Build-configuratie.
  • python_tools/: Suite voor kwantisatie en preprocessing aan de PC-zijde.
  • crop_token.py: Vocabulaire-snoei (schaalt terug naar 32K tokens).
  • cropmodelweight.py: Slicing van de embedding-matrix.
  • qat_158.py: BitNet QAT (Quantization-Aware Training) fine-tuning.
  • bit4_embedding.py: INT4 weight packer voor embeddings.
  • packtokenizerbin.py: Serialiseert tokenizer-regels naar een ESP32 .bin.
  • packmodelbin.py: Verpakt 1,58-bit layer-chunks voor fysieke uitlijning.
  • lookmodelstructure.py: Debug-tool voor het inspecteren van .safetensors.
  • flash_*.bat: Multi-threaded scripts voor snel flashen.

Licentie

Dit project is gelicenseerd onder de MIT-licentie - zie het LICENSE-bestand voor details.

Erkenningen

Inspiratie, gerelateerd werk en referenties:

  • ESP-32-s3-Story-maker-LLM: Inspiratie voor single-node gekwantiseerde LLM-implementatie op ESP32-S3.
  • esp32s3-distributed-ai: Inspiratie voor multi-node gedistribueerde AI-architectuur op microcontrollers.
  • BitNet: Concept en architectuur van 1,58-bit ternaire kwantisatie.