I Trained My Own AI Model with 30 Million Parameters and Ran It on an ESP32
How to train your own transformer model and run it on a tiny ESP32: tokenization, 4-bit quantization, and memory limits.
The author built their own language model and ran it on an ESP32 microcontroller. The article describes the model architecture, the training process, weight quantization, and results obtained on two boards.
Model Architecture
The author chose a decoder-only transformer architecture. The same type of architecture is used in modern GPT models. The model predicts the next token based on the previous tokens.
The model contains 12 transformer layers and a vocabulary of about 6,000 tokens. The author focused on the limited memory of the ESP32, so the model was intentionally chosen to be small.
The goal of the project is a chatbot that understands dialogues, language, and sentence structure.
Data and Tokenization
Publicly available datasets were used for training.
The text was split using the SentencePiece tool. Each token is assigned a numerical identifier. The model works with numbers rather than words. A single word may consist of several tokens, for example the word dependable.
The data was divided into batches. The model was trained on each batch sequentially, which allowed it to correct errors and gradually improve.
Training Process
The training loop is implemented in Python:
- The model receives text and predicts the next token.
- The cross entropy function evaluates the accuracy of the prediction.
- Back propagation determines how to change the model's weights.
- Gradient clipping limits the magnitude of weight changes.
The model began training with random weights. At this stage, it produced random fragments of text. By 100 million tokens, the quality had improved, but errors remained. By 300 million tokens, the model had improved noticeably.
Validation loss decreased across all three data streams. On the dialogue data, the value dropped from 8.81 to 2. This means the model predicts the next token better.
Weight Quantization
The ESP32 has 32 MB of flash memory and 16 MB of PSRAM. The model has 3.5 million parameters. Storing all parameters in 32-bit format requires about 122 MB of memory, which is roughly four times the flash memory capacity.
To reduce the size, quantization was applied. Most of the weights were converted to 4-bit format. Some weights were kept in 16-bit format so that not all accuracy would be lost. Quantization reduces the model's accuracy, but the author considers the balance between size and quality satisfactory.
Running the Model on ESP32
- Weights are stored in flash memory.
- PSRAM is used for temporary data.
- The processor performs computations, prediction, and reading of weights.
This process repeats for each token until the full response is generated.
The author received responses generated by the model on the ESP32. The model struggles with complex questions, but it produces meaningful answers.
Running on Xiao ESP32S3
The Xiao ESP32S3 board has 8 MB of flash memory and 8 MB of PS RAM. Even with quantization, the model does not fit entirely into the board's memory.
The weights are loaded from a microSD card, which is located on the back of the round display. Reading from a microSD card is slower than reading from RAM. For this reason, two buffers were used: while the processor works on data from one buffer, the next portion of weights is loaded into the second.
The ESP32S3 has two cores. The workload was distributed between them. However, to generate each token, the processor must pass through all 12 transformer layers, which takes a lot of time.
| Board | Flash memory size | PS RAM size | Generation speed |
|---|---|---|---|
| ESP32 (dev kit) | 32 MB | 16 MB | 1.4 tokens per second |
| Xiao ESP32S3 | 8 MB | 8 MB | about 0.74 tokens per second |
Alternative Approaches
The author notes that there are simpler ways to work with the ESP32.
Cloud API (using Google Gemini as an example). The model runs in Google AI Studio. The user obtains an API key and adds it to the ESP32 code. The board sends requests to Google's server and displays the response on the screen. The model does not run on the board itself.
Local run via llama. The program allows AI models to run on a computer. The model is loaded and runs in the terminal. The API key is added to the ESP32 code in the same way as when using Gemini. The difference is that the model runs on the local computer rather than on Google's server.
The choice of approach depends on the project's goals. Training a model from scratch and running it on an ESP32 require significantly more time and resources.
The author plans to publish instructions and additional materials on GitHub.
Introduction
The author presents the goal of the video: to train their own AI model and run it on an ESP32.
- Author: Norina
- Goal: train your own AI model and run it on an ESP32
- The author notes that the task was difficult and required solving many problems
Model Architecture
Chosen architecture is the decoder-only transformer, which is used in modern GPT models. The model predicts the next token based on the previous ones.
- Architecture: decoder-only transformer, as in modern GPT models
- Operating principle: predicting the next token based on the previous ones
- The model has 12 transformer layers
- Vocabulary of about 6000 tokens
- The model is being created for a weak device with very limited memory
Data and Tokenization
The goal is to create a chatbot that understands dialogues, language, and sentence structure. The data is taken from publicly available datasets. The text is tokenized using SentencePiece, and the data is split into batches.
- Goal: create a chatbot that understands dialogues and sentence structure
- Publicly available datasets were used
- Tokenization tool: SentencePiece
- Tokens receive numerical values; the model works with numbers, not words
- Words like dependable may consist of several tokens
- Data is divided into batches for gradual training
Training Code
The model predicts the next token, cross-entropy evaluates accuracy, and backpropagation adjusts the weights. Gradient clipping is used to control the magnitude of changes.
- The model receives text and predicts the next element
- Cross entropy evaluates the correctness of predictions
- Back propagation adjusts the model's weights
- Gradient clipping is used to control the magnitude of weight changes
Training Process
The model starts with random weights and produces random fragments. By 100 million tokens, quality begins to improve; by 300 million, it improves noticeably.
- Initial state: random weights, random text fragments
- Around 100 million tokens: start of improvements, but errors remain
- By 300 million tokens: significant improvement
- Validation loss decreased across all three data streams
- On conversational data, loss decreased from 8.81 to 2
Running on ESP32
The ESP32 board has 32 MB of flash memory and 16 MB of PS RAM. A model with 3.5 million parameters stored in 32-bit format takes about 122 MB, which is roughly four times more than the board's flash memory. Therefore, quantization to 4 bits is applied.
- Board: 32 MB of flash memory and 16 MB of PS RAM
- Model: 3.5 million parameters
- Storing in 32-bit format requires about 122 MB of memory
- This is roughly four times more than the board's flash memory
- Quantization: storing weights in 4-bit format with some loss of accuracy
- Some weights are kept at 16 bits to preserve accuracy
- Weights are stored in flash memory, while PS RAM is used for temporary data
- The processor generates tokens one at a time until the response is complete
Results on ESP32
The model runs locally on an ESP32, answers simple questions, and handles complex ones worse. The answers are considered meaningful.
- The model runs locally on the ESP32
- The model handles complex questions poorly
- The answers are assessed as meaningful
Running on Xiao ESP32S3
Автор пробует запустить модель на меньшей плате Xiao ESP32S3 с 8 МБ флеш-памяти и 8 МБ PS RAM. Веса читаются с microSD-карты, что медленнее чтения из RAM. Используются два буфера и два ядра процессора.
- Plate Xiao ESP32S3: 8 MB flash and 8 MB PS RAM
- Even with quantization, the model does not fit entirely on the board
- Weights are read from a microSD card located on the back of the round display
- Reading from microSD is slower than from RAM
- Two buffers are used: while one is being processed, the other loads weights
- Work is split between the two cores of the ESP32S3
- Speed: 1.4 tokens per second on the dev kit, about 0.74 tokens per second on the Xiao
Alternative Methods
The author mentions simpler ways: using the Google Gemini API via Google AI Studio and running locally via Llama. Instructions will be published on GitHub.
- Option 1: Google Gemini via Google AI Studio, obtaining an API key and adding it to the ESP32 code
- The ESP32 sends a request to Google's server and receives a response, which is displayed on the screen
- Option 2: Llama for running the model locally on a computer via terminal
- For Llama, you can also obtain an API and connect it to the ESP32 code
- Instructions will be published on GitHub
Conclusion
Training a model from scratch and running it on an ESP32 is a challenging task. The choice of approach depends on your goals. The author promises to publish all materials on GitHub.
- Simpler approaches exist than training a model from scratch
- The choice of approach depends on your goals
- All materials will be published on GitHub