ESP32-S3 Distributed AI

ESP32-S3 Distributed AI is a cutting-edge project designed to facilitate distributed micro-LLM inference using three ESP32-S3 N16R8 boards via ESP-NOW communication.

Overview

This project implements a distributed AI system that effectively runs a 56M-parameter language model across three ESP32-S3 microcontrollers. Building on the work of slvDev/esp32-ai, which showcased the execution of TinyStories on a single board, this initiative extends the architecture to a multi-board distributed system featuring web-based interaction.

The language model, trained on the WikiText-103 corpus (Wikipedia), utilizes Per-Layer Embeddings (PLE) from Google's Gemma architecture, quantized to 4-bit. The architecture is optimized to split the 50.3M-parameter PLE table across three boards that communicate wirelessly via the ESP-NOW protocol. Board B utilizes a KV cache in PSRAM, allowing the transformer to maintain awareness of the complete generated sequence instead of processing token-by-token.

Architecture

The system architecture consists of three primary components:

┌───────────────────┐      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   │
└───────────────────┘                  └───────────────────┘                  └───────────────────┘

Inference Flow

The end-to-end inference flow is as follows:

Browser ──WiFi──► Board C ──ESP-NOW──► Board A ──ESP-NOW──► Board B
                        ◄─────────────── ◄──────────────

1. The user types a prompt in their browser (connected to Board C's WiFi AP) and clicks "Generate".
2. Board C receives the prompt via HTTP POST and forwards it to Board A via ESP-NOW.
3. Board A tokenizes the prompt text and processes each token accordingly, utilizing Board C for additional data.
4. After processing, Board A sends the results back to Board C, where the generated text is streamed to the browser in real-time.

Split-PLE Design

The PLE table is divided between the boards to optimize memory usage:

• PLE_A resides on Board C, while PLE_B is managed by Board B. This design allows for efficient retrieval and processing of token embeddings across the distributed system.

KV Cache on Board B

This innovative approach features a KV cache in PSRAM (1.5 MB for 256 positions) on Board B, significantly enhancing the model's ability to maintain coherence across generated sequences. This solution addresses historical constraints of autoregressive models by allowing proper position-aware attention during inference.

Communication Protocol

• Uses ESP-NOW wireless protocol for peer-to-peer communication without the need for routers.
• Packet sizes of 250 bytes with automatic fragmentation ensure efficient data handling between boards.

Project Structure

The repository includes comprehensive code and data for both the training pipeline and ESP32-S3 firmware, including:

• Training Pipeline: Model architecture in PyTorch, data loading, training routines, and quantization scripts.
• Firmware: Distributed inference runtime code for each ESP32-S3 board.
• Utility Scripts: For environment setup and flashing the firmware to boards.

Model Specifications

This system utilizes a Tiny decoder-only transformer architecture with the following specifications:

• Total Parameters: 56.0M
• Vocabulary Size: 32,768 (using BPE from WikiText-103)
• Quantization: 8-bit for token embeddings and 4-bit for the rest of the model components.

Training Results

The model achieved impressive metrics during training, including:

• Final Validation Loss: 5.26
• Perplexity: 192

What's Pending

Future improvements include optimization of the BPE tokenizer implementation, enhancements to the web UI, and power optimizations among others.

Hardware Requirements

• ESP32-S3 N16R8 Boards: 3 required for deploying the model.
• USB-C cables: For each board.
• Computer: To run the setup and training scripts.

The ESP32-S3 Distributed AI project exemplifies a significant advancement in distributed machine learning on microcontrollers, effectively showcasing the potential of low-power devices in sophisticated AI applications.