Transcript
00:00:00This is a 28.9 million parameter language model generating text right now one word at a time
00:00:08on a chip that costs about eight dollars. There's no wi-fi, nothing is being sent to a server,
00:00:14everything is happening inside an ESP32S3, a microcontroller with less RAM than a computer
00:00:22from the 90s. This is insane. How is this actually possible? What's the magic behind it and how can
00:00:29we build something similar? Well, those are all good questions that we're going to look at in today's
00:00:34video. It's going to be a lot of fun, so let's dive into it. So the ESP32S3 is a chip that gives you
00:00:46512 kilobytes of SRAM and that is the amount of fast memory this chip can compute with. Normally,
00:00:54if you would attempt to run an LLM inside of it, the whole model would have to fit right in there.
00:01:00The last language model anyone got running on a chip like this topped out at 260,000 parameters. So this
00:01:08one holds about 110 times more and the person who pulled it off is a Ukrainian developer Slava S. So
00:01:16what's the magic trick? How did they do it? How did they fit a 28.9 million parameter model on an $8 chip?
00:01:24Well, the workaround comes from an idea Google used in Gemma. It's called per-layer embeddings. Most of a
00:01:31language model's parameters do not compute anything. They sit in an embedding table that just gets read
00:01:37from. If most of your parameters are only ever looked up, they don't need fast memory. So you can get away with
00:01:44leaving that table sitting in a slow cheap flash memory and only pull the rows the current token needs. And the
00:01:52small part that actually computes and thinks about the next token, the attention head and feed forward stays in SRAM.
00:01:59Okay, that's the compute part. But the question still remains: How do we fit such a big model onto such a tiny chip?
00:02:05So the 25 million row table lives in flash memory. And that's the biggest chunk of the model's whole
00:02:1228.9 million parameters. And flash is cheap and huge on this particular chip. It has 16 megabytes of it.
00:02:20So instead of trying to squeeze that table into the same 512 kilobytes of SRAM,
00:02:26it just stays parked in flash. We pull out about six rows from it, one for each of the model's six layers.
00:02:33That's roughly 450 bytes total. Now, before you get too excited, I do have to address the fact that
00:02:40this is a very dumb, simple model. This won't be your typical GPT-style LLM. It won't give you code
00:02:47generation or answer questions about difficult topics. Because if you tried to train a normal model this size
00:02:53on a normal dataset, 28 million parameters would just give you gibberish. But this model itself is trained on
00:03:01Tiny Stories, a dataset built by researchers at Microsoft, written deliberately simple enough that
00:03:08even a tiny model, a few million parameters, can learn to write them coherently. And running it in plain C,
00:03:15on a chip with no operating system and no Python interpreter, that idea comes from the famous Andrei
00:03:22Karpathy's Llama2.c project, which showed us that you could train a small language model and run
00:03:28inference on it using nothing but a few hundred lines of portable C. And that is the blueprint this entire
00:03:35project is built on. Without Karpathy, this probably wouldn't even be possible. So that's how it works in a
00:03:40nutshell. Now let's actually try to build it on our own and run it on the chip to see how it performs.
00:03:47To actually build this ourselves, the first thing you need is the right hardware, specifically an ESP32S3
00:03:54with 16 megabytes of flash and 8 megabytes of PS RAM. That's the N16R8 variant. And this is really
00:04:03important because remember when we talked about that 25 million row table living in flash? Well, by itself,
00:04:10once trained and exported, it comes to about 15 megabytes. But boards with 4 or 8 megabytes of flash
00:04:17simply won't hold it. So if you're shopping for a board to follow along, that's the one spec you should
00:04:22check first. Now, when I first went through the project's actual instructions, the setup was scattered
00:04:28across a few different files and it assumed a fair amount of context I didn't have going into it. So
00:04:34instead of walking you through it exactly as written, I put together a single one shot script that does
00:04:40everything. Checks the board, installs the tool chain, prepares the data, trains it, exports it, verifies,
00:04:47builds and flashes it all in one go. So let's go ahead and run that script. And one quick heads up,
00:04:55if you're following along, the whole script takes about 25 minutes to finish. And probably that's the
00:05:00same time it would take you doing it step by step. It's just because there are so many separate
00:05:05commands you have to babysit. And most of the time is actually spent on training the tiny stories model.
00:05:11That takes about 13 minutes on my MacBook. And once the training is done, you'll see the LEDs on the panel
00:05:18start flashing. And that's an indication that we're about to load the model. And once it's loaded,
00:05:24here we can see the first basic example of a story of a little girl. And indeed,
00:05:29we're getting those nine tokens a second. But you'll notice that at a certain point,
00:05:33it will restart the story again. So the model is going in some kind of a loop. And if you want
00:05:39to give it a custom prompt, I've also included a sample script you can use to run your own custom
00:05:44prompts. So for this example, let's start a story about a robot. And another thing to note is that if you
00:05:50change the prompt, you will have to reflash the ESP32 again. And that's a limitation of this method.
00:05:56It can only play back one prompt that is pre-flashed at a time. And as you can see,
00:06:02we do get a mention of the robot in the story. And the story is indeed a bit different this time.
00:06:08But notice what happens after the end of the paragraph. We get back to the story of the little
00:06:13girl. So I have no idea why this is happening. But clearly, the model tends to steer back to that one
00:06:19specific story about that little girl. And let's do another example. And this time, let's use the
00:06:24famous phrase that is written at the very beginning of every Star Wars movie. And let's see where that
00:06:30takes us. So this is an interesting one. We can see that the model immediately drifts back to the
00:06:36little girl narrative again. And I'm assuming that for a model of this size, it doesn't even understand
00:06:42what a galaxy is. So probably that's why it's ignoring our specific prompt in this case. And
00:06:47once again, we see that the next paragraph starts the same story. So although this experiment is
00:06:53impressive and seeing a model produce nine tokens per second on a tiny microchip is genuinely mind-blowing,
00:06:59it's still very much a very limited proof of concept. As we saw, no matter what you ask the model, it will
00:07:06always drift back to storytelling because that's what it's trained on. But if you found this test
00:07:11interesting, I also did another video in a similar realm where I tested out if a first-gen Raspberry
00:07:18Pi could actually run a real LLM locally. So go check out that video if you're interested. So there you
00:07:24have it, folks. That's how you run a 28.9 million parameter language model on an ESP32 chip. We tested
00:07:32it. It works. So kudos to Slava for making this project. But what are your thoughts on this experiment?
00:07:38Do you see any real world examples where such an implementation might be useful? Give us your
00:07:43thoughts in the comment section down below. And folks, if you like these types of technical breakdowns,
00:07:48please let me know by smashing that like button underneath the video. And also don't forget to
00:07:53subscribe to our channel. This has been Andres from Betterstack and I will see you in the next videos.
Community Posts
No posts yet. Be the first to write about this video!
Write about this video