TinyLLM ESP32

A tiny decoder-only character-level language model that runs fully on-device on the ESP32 DevKit v1 no PSRAM, no cloud.

Built with ESP-IDF · Inference in C · Training in PyTorch

---

Hardware

Board	ESP32 DevKit v1 (WROOM-32)
CPU	Dual-core Xtensa LX6 @ 240 MHz
SRAM	520 KB (no PSRAM required)
Flash	4 MB
Interface	UART0 via USB-serial (115200 baud)

---

Model Architecture

A miniature GPT-style decoder-only transformer trained on character-level text (printable ASCII).

vocab_size  = 95    (printable ASCII 0x20–0x7E)
n_embd      = 64
n_head      = 4     (head_size = 16)
n_layer     = 2
block_size  = 32    (context window)
ff_dim      = 256

Memory Budget

Resource	Usage	Available
Flash (INT8 weights)	~103 KB	4096 KB
SRAM (activations + KV-cache)	~66 KB	520 KB
SRAM free for OS / Wi-Fi	~454 KB	—

Key techniques:

• INT8 weight quantisation — all weight matrices stored as int8 in Flash RODATA
• Weight tying — token embedding and LM head share the same matrix
• Static allocation — zero heap usage; all activation buffers are compile-time arrays
• Causal KV-cache — pre-fill prompt once, then generate token-by-token

---

Project Structure

llm/
├── CMakeLists.txt              ← ESP-IDF project root
├── sdkconfig.defaults          ← 240 MHz, BT off, -O2
├── .gitignore
├── README.md
│
├── main/
│   ├── CMakeLists.txt
│   └── llm.c                   ← UART prompt loop + built-in commands
│
├── components/
│   └── tinyllm/
│       ├── CMakeLists.txt
│       ├── tinyllm.c           ← Transformer inference engine (C)
│       └── include/
│           ├── tinyllm.h       ← Hyperparameters & public API
│           ├── tokenizer.h     ← Zero-overhead ASCII tokenizer
│           └── weights.h       ← INT8 model weights (generated by tools/)
│
└── tools/
    ├── train_tinyllm.py        ← PyTorch training + INT8 weight export
    └── export_weights.py       ← Re-export existing checkpoint → weights.h

---

Quick Start

1. Train the model (on your PC)

# Install dependencies
pip install torch numpy

cd llm/tools

# Train on Shakespeare (downloads automatically, ~5 min CPU / ~30 s GPU)
python3 train_tinyllm.py

# Or use your own corpus
python3 train_tinyllm.py --corpus /path/to/your_text.txt

# Resume from a saved checkpoint
python3 train_tinyllm.py --checkpoint tinyllm.pt

# Export only (no training)
python3 export_weights.py --checkpoint tinyllm.pt

This generates components/tinyllm/include/weights.h (~103 KB), which is automatically compiled into Flash on the next idf.py build.

2. Build & flash

# Source the ESP-IDF environment (once per terminal session)
source $IDF_PATH/export.sh

cd llm/

# Set the target chip
idf.py set-target esp32

# Build
idf.py build

# Flash and open the serial monitor
idf.py flash monitor

3. Use the serial interface

Connect at 115200 baud (idf.py monitor does this automatically).

After the boot banner, a > prompt appears. Type any printable text and press Enter:

> ROMEO:
ROMEO: What light through yonder window breaks?
It is the east, and Juliet is the sun...

[87 tokens | 9.3 tok/s]
>

Built-in commands

Command	Effect
:temp 0.8	Set sampling temperature (0.0 = greedy, 0.8 = creative default, 1.5 = chaotic)
:len 20	Set max new tokens (1 – 32)
:reset	Soft-reset the chip
:quit	Restart the chip

---

Training Script Options

usage: train_tinyllm.py [-h] [--corpus CORPUS] [--checkpoint CHECKPOINT]
                        [--export-only] [--iters ITERS]

options:
  --corpus CORPUS          Path to plain-text training corpus
                           (default: downloads Shakespeare ~1 MB)
  --checkpoint CHECKPOINT  Load this .pt file instead of training from scratch
  --export-only            Skip training; only export --checkpoint to weights.h
  --iters ITERS            Number of training iterations (default: 5000)

---

How Inference Works

prompt (string)
    │
    ▼
tinyllm_encode()          ASCII char → token ID (offset from 0x20)
    │
    ▼
transformer_forward()     For each token at position pos:
    ├─ token_emb lookup   INT8 → float32 via scale factor
    ├─ RMSNorm
    ├─ Multi-head attention  (causal, KV-cache, 4 heads × 16 dim)
    ├─ Residual + RMSNorm
    ├─ FFN  (GELU, 256 hidden)
    └─ Residual
    │
    ▼ logits [95]
    │
sample_token()            Temperature softmax + multinomial sampling
    │
    ▼
tinyllm_decode()          token ID → char (putchar → UART0)

All weight matrices (int8) are read directly from Flash via XIP (Execute-In-Place) — they are never copied to SRAM.

---

Performance

Measured on ESP32 @ 240 MHz with Shakespeare-trained weights:

Context length	tok/s (approx.)
1 token	~15 tok/s
16 tokens	~9 tok/s
32 tokens (full context)	~5 tok/s

Performance scales linearly with context length due to the O(n²) attention loop (small context keeps this practical).

---

Customisation

Change the training corpus

Pass any plain-text file. The model learns whatever style and vocabulary exists in the file:

python3 train_tinyllm.py --corpus iot_commands.txt

Tune the architecture

Edit the defines in components/tinyllm/include/tinyllm.h and the matching constants at the top of tools/train_tinyllm.py, then retrain and rebuild:

#define TINYLLM_N_EMBD    64   // increase for more capacity
#define TINYLLM_N_HEAD     4
#define TINYLLM_N_LAYER    2   // add more layers (watch SRAM!)
#define TINYLLM_BLOCK_SIZE 32  // longer context = more SRAM

⚠️ Always keep TINYLLM_N_EMBD % TINYLLM_N_HEAD == 0.

Add an OLED display

The tinyllm_generate() function streams characters via putchar(). Redirect this to an I2C SSD1306 driver to display output on screen instead of (or in addition to) UART.

---

Requirements

Firmware

• ESP-IDF v5.x
• Target: esp32 (DevKit v1 / WROOM-32)

Training (PC)

• Python 3.10+
• PyTorch ≥ 2.0
• NumPy

---

License

MIT — see LICENSE in the ESP-IDF root, or add your own.