Distributed micro-LLM inference across three ESP32-S3 N16R8 boards with ESP-NOW communication.
This project implements a distributed AI system that runs a 56M-parameter language model across three ESP32-S3 microcontrollers. Inspired by slvDev/esp32-ai, which demonstrated running TinyStories on a single board, this project extends the architecture to a multi-board distributed system with web-based interaction.
The model is trained on WikiText-103 (Wikipedia corpus) using Per-Layer Embeddings (PLE) from Google's Gemma architecture, quantized to 4-bit, and split across three boards that communicate via ESP-NOW wireless protocol. The 50.3M-parameter PLE table is split across Board A and Board B (Split-PLE) to fit within the 16MB flash per board. Board B maintains a KV cache in PSRAM (1.5 MB for 256 positions), enabling the transformer to attend to the full generated sequence instead of operating token-by-token.
┌───────────────────┐ ESP-NOW ┌───────────────────┐ ESP-NOW ┌───────────────────┐
│ Board A │ ◄──────────────► │ Board B │ ◄──────────────► │ Board C │
│ Embeddings + │ │ Core + KV Cache │ │ PLE_A + Decoder │
│ PLE_Proj + Head │ │ │ │ + WiFi │
│ │ │ │ │ │
│ • BPE Tokenizer │ │ • 6 Attn layers │ │ • PLE table A │
│ • tok_emb 8-bit │ │ • 6 FFN layers │ │ (12.5 MB) │
│ • ple_model_proj │ │ • KV Cache (128) │ │ • Sampling │
│ • out_norm + head │ │ • PLE_B (half) │ │ • Web Server │
│ │ │ • PLE Gating │ │ • Space heuristic │
│ Flash: 4.31 MB │ │ Flash: 13.71 MB │ │ Flash: 13.37 MB │
└───────────────────┘ └───────────────────┘ └───────────────────┘
Browser ──WiFi──► Board C ──ESP-NOW──► Board A ──ESP-NOW──► Board B
◄─────────────── ◄──────────────
User types a prompt in the browser (connected to Board C's WiFi AP) and clicks "Generate"Board C receives the prompt via HTTP POST → forwards it to Board A via ESP-NOW (MSG_EMBED_REQUEST
)Board A tokenizes the prompt text → for each token: a. Looks up token embeddingx[D]
(8-bit tok_emb) b. Computes local PLE projection:tmpP = ple_model_proj_a @ x
, RMS-norms it c. Requests PLE table row fromBoard C via ESP-NOW: sendstoken_id
→ C looks upple_table_a[token]
→ sends row back d. Combines:ple = (tmpP + trow * sqrt(Ph)) / sqrt(2)
→ sendstoken_id + x[D] + ple[L×Ph]
toBoard B e.Board B computes PLE_B[L×Ph] from its local table → combines PLE_A + PLE_B into full PLE[L×P] → runs 6 transformer layers (attention + FFN + PLE gating) → sends hidden statex[D]
back to Board A f.Board A appliesout_norm
→ computes output head:logits[V] = tok_emb^T · rmsnorm(x)
→ sends top-40 token IDs to Board C g.Board C samples the next token → decodes via BPE vocab → appends to generated outputBoard C streams the generated text to the browser via SSE (/stream
endpoint)User sees the text appear in real-time on the web page
The 50.3M-parameter PLE table (vocab × layers × 256) is split into two 128-dimension halves:
PLE_A(Board C, 12.5 MB, 4-bit group=64): first 128 dims per layer — reloaded from Board C at runtime via ESP-NOW (MSG_PLE_REQUEST / MSG_PLE_RESULT)PLE_B(Board B, 12.6 MB, 4-bit group=128): last 128 dims per layer — local on Board B
Board A keeps the small ple_model_proj
(~50 KB, 4-bit) and ple_proj_norm
(~0.5 KB, fp32) for local projection. The 12.5 MB PLE table was moved from Board A to Board C to fix partition overflow on Board A and allow upgrading tok_emb
to 8-bit for better inference quality.
Each token in the autoregressive loop previously ran through Board B's transformer with seq_len=1
— attention only saw the current token, not the history. This crippled coherence because the model was designed for seq_len=128
context.
Board B now maintains a KV cache in PSRAM (1.5 MB for 256 positions):
- For each new token, Q is computed normally, while K and V are appended to per-layer caches
- Attention scores are computed against
all cached positions(not just the current token) - RoPE uses the actual position index instead of
pos=0
— for the first time the transformer sees proper position-aware attention - Cache is reset on
MSG_SEQ_START
(new prompt) and caps atseq_len=256
- No changes to A, C, or the wireless protocol
This is the single largest quality improvement: the transformer finally works as designed, attending to the full generated sequence instead of operating token-by-token in isolation.
| Board | MAC Address | Serial Port |
|---|---|---|
| A (Embeddings + PLE_Proj + Head) | 14:c1:9f:2a:ac:c8 |
|
/dev/cu.usbmodem5C372059631 |
||
| B (Core + KV Cache) | 14:c1:9f:2c:91:10 |
|
/dev/cu.usbmodem5C372065471 |
||
| C (PLE_A + Decoder + WiFi AP) | 28:84:85:51:dc:10 |
|
/dev/cu.usbmodem5C4D0363671 |
ESP-NOW wireless peer-to-peer (no router needed)- Packets ≤ 250 bytes with automatic fragmentation
- Custom protocol with sequence numbers and acknowledgments
- Estimated latency: 1–5ms per round-trip
esp32s3/
├── README.md # This file (English)
├── README.es.md # Spanish version
├── pyproject.toml # Python dependencies
│
├── src/ # Training pipeline (runs on PC/Mac)
│ ├── model.py # PLE TinyLM architecture (PyTorch)
│ ├── dataset.py # TinyStories/WikiText-103 data
│ ├── train.py # Training loop
│ ├── quantize.py # 4-bit group-wise quantization
│ ├── export.py # Export to 3 board binaries
│ └── gen_assets.py # Generate vocab.h for firmware
│
├── firmware/ # ESP32-S3 firmware (Arduino IDE)
│ ├── common/
│ │ ├── llm.h # Distributed C inference runtime
│ │ └── espnow_protocol.h # ESP-NOW message protocol
│ │
│ ├── board_a_embeddings/ # Board A firmware
│ │ ├── board_a_embeddings.ino # Main sketch
│ │ ├── partitions.csv # Flash partition table
│ │ └── vocab.h # Generated tokenizer vocabulary
│ │
│ ├── board_b_core/ # Board B firmware
│ │ ├── board_b_core.ino # Main sketch
│ │ └── partitions.csv
│ │
│ ├── board_c_decoder/ # Board C firmware
│ │ ├── board_c_decoder.ino # Main sketch (WiFi + Web UI)
│ │ ├── partitions.csv
│ │ └── vocab.h # Generated tokenizer vocabulary
│ │
│ └── model/ # Exported model binaries (gitignored)
│ ├── board_a.bin # 4.31 MB (tok_emb 8-bit + ple_proj + out_norm)
│ ├── board_b.bin # 13.71 MB (transformer layers + PLE_B)
│ ├── board_c.bin # 13.37 MB (PLE table A, 4-bit group=64)
│ ├── golden.npz # Reference logits for verification
│ └── golden.txt # Text-format golden reference
│
├── tools/ # Utility scripts
│ ├── flash_all.sh # Flash all 3 boards
│ ├── verify_models.py # Verify exported binaries
│ └── setup_env.sh # Install Python dependencies
│
├── data/ # Dataset (gitignored)
│ ├── tinystories/ # Raw TinyStories text
│ ├── wikitext103/ # Raw WikiText-103 text
│ ├── tokenizer/ # Trained BPE tokenizer
│ ├── train.bin # Tokenized training data (112M tokens)
│ └── val.bin # Tokenized validation data
│
└── runs/ # Training checkpoints (gitignored)
├── ple-wiki60m-s0.pt # 56M model checkpoint
├── ple-wiki60m-s0.json # Training history
└── train.log # Training output log
| Parameter | Value | |---|---|---| | Architecture | Tiny decoder-only transformer with Split-PLE | | Total Parameters | 56.0M stored | | Core (dense, SRAM) | 1.5M | | PLE Table (flash, split) | 50.3M (25.2M per board) | | Output Head (tied) | 4.2M | | Vocabulary Size | 32,768 (BPE, WikiText-103) | | d_model | 128 | | n_layers | 6 | | n_heads | 4 | | ffn_hidden | 223 | | ple_dim | 256 (128 per board) | | Quantization | tok_emb: 8-bit group=128, PLE table A: 4-bit group=64, rest: 4-bit group=128 | | Model Binary Size | 31.39 MB total (A: 4.31, B: 13.71, C: 13.37) | | Dataset | WikiText-103 (Wikipedia) | | Training Steps | 12,000 | | Batch Size | 8 | | Sequence Length | 128 |
| Metric | Value |
|---|---|
| Final Validation Loss (fp32) | 5.26 |
| Perplexity (fp32) | 192 |
| 4-bit Quantization Degradation | +0.62 nats |
| 4-bit Perplexity | 358 |
| Training Tokens | 12.3M |
| Training Time | ~6.9 hours (Mac i7 CPU) |
- Python 3.11+ uvpackage manager- Arduino IDE with ESP32 board package
- 3x ESP32-S3 N16R8 boards
source .venv/bin/activate
python src/dataset.py --dataset wikitext103
python src/train.py --arm ple --d-model 128 --n-layers 6 --ple-dim 256 \
--target-core 1500000 --batch-size 8 --seq-len 128 --steps 12000 --tag wiki60m
python src/quantize.py --tag wiki60m
python src/export.py ple-wiki60m-s0
python src/gen_assets.py
./tools/flash_all.sh /dev/cu.usbmodemXXXX
Note:All 3 boards are already flashed. If you need to reflash, seetools/flash_all.sh
.
Power on all three boards
Connect your phone/laptop to WiFi:
ESP32-DIST-AI(password:ai123456
) - Open
http://192.168.4.1
in your browser - Type a prompt and click "Generate"
Architecture — PLE TinyLM v2, Split-PLE, 56M params #
- PLE TinyLM v2 with d_model=128, n_layers=6, ple_dim=256
- Split-PLE: PLE table divided across Board A + Board C (128+128 per layer)
- PLE table A moved from Board A to Board C (fixed partition overflow on A)
- Per-token A↔C PLE request/response via ESP-NOW
- KV cache on Board B: full multi-head causal attention with K/V caching in PSRAM (1.5 MB for 256 positions), RoPE uses actual position indices
- MAX_SEQ_LEN=256, gen loop=128, supporting 30+ word sequences
- srand(esp_random()) for non-deterministic sampling
Training — WikiText-103, PPL 192, 12K steps, 12.3M tokens #
- WikiText-103 dataset download and BPE tokenization (112M tokens)
- Model training: 12,000 steps, batch_size=8, seg_len=128, 12.3M tokens
- Final fp32 loss: 5.26, Perplexity: 192
- Training time: ~6.9 hours (Mac i7 CPU)
Quantization & Export — 4-bit, 31.39 MB total #
- 4-bit group-wise quantization (+0.62 nats degradation from fp32, 4-bit PPL 358)
- tok_emb: 8-bit group=128, PLE table A: 4-bit group=64, rest: 4-bit group=128
- Export to 3 board binaries: A: 4.31 MB, B: 13.71 MB, C: 13.37 MB
- Reference golden logits for verification
Firmware — 3 boards, ESP-NOW, Web UI, autoregressive loop #
- Board A: tokenizer + 8-bit tok_emb + PLE projection + output head
- Board B: 6 transformer layers + PLE_B + PLE gating + KV cache
- Board C: PLE table A (4-bit group=64) + decoder + WiFi AP + Web UI
- ESP-NOW protocol with fragmentation, sequence numbers, and acks
- Web UI for prompt input, streaming output, and SSE endpoint
- Autoregressive generation loop (up to 128 tokens) across all 3 boards
- Space-insertion heuristic on Board C
- Arduino core 3.3.11 compatibility, partition table 1.25 MB
Flash & Deploy — 3 boards flashed with firmware + model #
- Board A: MAC
14:c1:9f:2a:ac:c8
/ port5C372059631
- Board B: MAC
14:c1:9f:2c:91:10
/ port5C372065471
- Board C: MAC
28:84:85:51:dc:10
/ port5C4D0363671
- Flash and verification scripts
Testing — Local inference confirmed coherent output (~30 words) #
-
Local inference test confirmed ~30 coherent words
-
Confirmed 4-bit quantization is the quality bottleneck (not protocol)
-
Comparison between full fp32, quantized, and on-device inference
BPE tokenizer in C— Replace brute-force vocab search with proper BPE merge implementation - Static ESP-NOW MACs— Configure known MACs instead of broadcast - Web UI improvements— Streaming token display, temperature/top-k controls, status indicators - Error handling— ESP-NOW retransmission and timeout on packet loss - Power optimization— Deep sleep between inference requests - Audio output— I2S speaker integration (future)
| Component | Quantity | Notes |
|---|---|---|
| ESP32-S3 N16R8 | 3 | 512KB SRAM, 8MB PSRAM, 16MB flash |
| USB-C cables | 3 | For flashing firmware |
| Computer | 1 | Mac/Linux/Windows with Arduino IDE |
This project extends the work of:
— Running 28.9M parameter LLM on a single ESP32-S3 using Per-Layer EmbeddingsslvDev/esp32-aiGoogle Gemma— Per-Layer Embeddings architecture** WikiText-103**— Dataset by Salesforce Research (arXiv:1609.07843)** TinyStories**— Dataset by Eldan & Li (Microsoft Research,arXiv:2305.07759)— Inspiration for running tiny LMs in plain Ckarpathy/llama2.c
MIT