Chạy mô hình ngôn ngữ 28,9 triệu tham số trên vi điều khiển ESP32-S3 giá 8 USD

Chạy một LLM 28,9 triệu tham số trên vi điều khiển giá 8 USD

Open to Work · 𝕏 slvDev · LinkedIn

Đây là một mô hình ngôn ngữ với 28,9 triệu tham số có khả năng tạo văn bản trên ESP32-S3, một loại vi điều khiển có giá khoảng 8 USD. Mô hình chạy trực tiếp trên chip mà không cần gửi dữ liệu đến máy chủ, đồng thời ghi từng từ lên một màn hình nhỏ kết nối với chip với tốc độ khoảng 9 token mỗi giây. Mô hình ngôn ngữ gần nhất từng được chạy trên một con chip như thế này chỉ có 260 nghìn tham số, vì vậy mô hình này lớn gấp khoảng một trăm lần. Nó có thể hoạt động được là nhờ phần lớn mô hình nằm trong bộ nhớ flash thay vì RAM, sử dụng một ý tưởng từ các mô hình Gemma của Google có tên là Per-Layer Embeddings.

Các con số

Tại sao điều này lại khó và làm thế nào để nó vẫn có thể hoạt động

Một vi điều khiển có rất ít bộ nhớ tốc độ cao. ESP32-S3 chỉ cung cấp cho bạn 512KB SRAM. Thông thường, toàn bộ mô hình phải nằm trong phạm vi truy cập của bộ nhớ này, điều đó khiến bạn bị giới hạn ở các mô hình cực nhỏ, và đó là lý do tại sao mô hình trước đó trên một con chip như thế này chỉ có 260 nghìn tham số.

Cách giải quyết vấn đề này là không đặt toàn bộ mô hình vào bộ nhớ tốc độ cao nữa. Hầu hết các tham số của một mô hình ngôn ngữ nằm trong bảng nhúng (embedding table), nơi mô hình chỉ đọc dữ liệu thay vì thực hiện tính toán trên đó. Vì vậy, bạn có thể để bảng 25 triệu hàng đó trong bộ nhớ flash tốc độ chậm và chỉ lấy ra vài hàng mà mỗi token cần, khoảng 450 byte, trong khi phần nhỏ thực hiện công việc tính toán thực tế sẽ nằm trong bộ nhớ tốc độ cao. Khi đó, việc chạy một mô hình lớn gần như không tốn kém gì, vì bạn không bao giờ phải tải toàn bộ nó. Nó chỉ nằm trong flash và được lấy mẫu từng chút một.

Ý tưởng đó chính là Per-Layer Embeddings của Google, từ Gemma 3n và Gemma 4. Ở đây, nó được áp dụng trên cấu trúc bộ nhớ của một vi điều khiển thay vì điện thoại hay GPU. Theo những gì tôi biết, chưa có ai thử nghiệm điều này trên một con chip nhỏ như vậy.

Nó làm được gì và không làm được gì

Mô hình này được huấn luyện trên TinyStories, vì vậy nó viết các câu chuyện ngắn, đơn giản và hầu như giữ được tính mạch lạc. Nó sẽ không trả lời câu hỏi, làm theo hướng dẫn, viết mã nguồn hay biết các sự kiện thực tế. Giới hạn đó đến từ phần nhỏ của mô hình chịu trách nhiệm suy luận, và thủ thuật bộ nhớ không thay đổi được điều này. Điều thú vị ở đây là kiến trúc, việc nhồi nhét một mô hình lớn vào một con chip nhỏ, chứ không phải là những gì một mô hình 28,9 triệu tham số có thể nói.

Tự chạy mô hình

Firmware, sơ đồ đấu dây và các bước nạp chương trình (flashing) nằm trong firmware/esp32_llm/README.md. Mã nguồn huấn luyện, thử nghiệm loại bỏ (ablation) và lượng tử hóa nằm trong src/ và experiments/. Phương pháp đầy đủ, các thử nghiệm loại bỏ và các phép đo trên chip được viết trong RESULTS.md.

Ghi nhận

TinyStories là tập dữ liệu được dùng để huấn luyện: những câu chuyện tổng hợp ngắn đủ đơn giản để một mô hình nhỏ vẫn có thể học cách viết mạch lạc (Ronen Eldan và Yuanzhi Li, Microsoft Research, arXiv:2305.07759). Phần còn lại là Per-Layer Embeddings, thiết kế của Google từ các mô hình Gemma, chính là yếu tố cho phép một mô hình lớn có thể chạy trên một con chip nhỏ.

Dự án llama2.c của Andrej Karpathy là lý do khiến nhiều người, bao gồm cả tôi, tin rằng bạn có thể huấn luyện một mô hình ngôn ngữ nhỏ và chạy nó bằng ngôn ngữ C thuần túy. Dự án này được phát triển dựa trên nền tảng đó.

Quá trình thực hiện thực tế

Tôi cố tình để lại lịch sử phát triển lộn xộn trong kho lưu trữ. Điều đó bao gồm một lỗi tôi tìm thấy trong cách tính toán tham số của chính mình, vốn đã làm thổi phồng một con số ban đầu, và kết quả đã được sửa chữa sau khi tôi khắc phục lỗi đó. Lịch sử commit và tệp RESULTS.md cho thấy các con số đã thay đổi như thế nào và tại sao.