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

• Trên ESP32-S3, một mô hình ngôn ngữ 28,9 triệu tham số được chạy mà không cần kết nối máy chủ và xuất văn bản lên màn hình nhỏ với tốc độ khoảng 9 token/giây
• Trong toàn bộ tham số, 25 triệu được đặt trong flash chậm, và cấu trúc Per-Layer Embeddings chỉ đọc khoảng 6 hàng cần thiết cho mỗi token, tương đương 450 byte
• Mô hình có kích thước 14,9MB ở mức 4-bit; lõi tính toán dùng cho mỗi token được đặt trong 512KB SRAM, đầu ra và bộ nhớ làm việc trong 8MB PSRAM, còn bảng embedding lớn nằm trong 16MB flash
• Được huấn luyện bằng TinyStories nên có thể tạo các câu chuyện ngắn, đơn giản với độ nhất quán tương đối tốt, nhưng không phù hợp cho hỏi đáp, làm theo lệnh, viết mã hay kiến thức thực tế
• So với mô hình 260 nghìn tham số từng chạy trên chip tương tự trước đây, nó chứa số tham số nhiều hơn khoảng 100 lần; điểm cốt lõi không nằm ở chất lượng sinh mà ở cấu trúc bộ nhớ giúp đưa mô hình lớn lên chip nhỏ

---

Phần cứng và hiệu năng chạy

• Toàn bộ xử lý được thực hiện trên một ESP32-S3 giá khoảng 8 USD, không gửi dữ liệu lên máy chủ
  • Sử dụng 512KB SRAM, 8MB PSRAM và 16MB flash
  • Tốc độ tổng thể khoảng 9,5 tok/s, tốc độ tính toán thuần khoảng 9,7 tok/s
  • Kích thước mô hình 4-bit là 14,9MB
• Trong tổng số 28,9 triệu tham số, 25 triệu được lưu trong bảng tra cứu trên flash

Cách đưa mô hình vào bộ nhớ nhỏ

• Thông thường, toàn bộ mô hình phải có thể được truy cập từ bộ nhớ nhanh, nhưng SRAM của ESP32-S3 chỉ có 512KB, nên chỉ chứa được các mô hình rất nhỏ
• Tận dụng việc phần lớn tham số nằm trong bảng embedding chứ không trực tiếp tham gia tính toán, bảng này được giữ lại trên flash
  • Mỗi token chỉ đọc khoảng 6 hàng cần thiết, tương đương khoảng 450 byte
  • Chỉ phần nhỏ đảm nhiệm tính toán thực sự mới được giữ trong bộ nhớ nhanh
  • Phần lớn mô hình không được nạp trong lúc chạy, mà chỉ chọn các phần cần thiết từ flash
• Vai trò của từng loại bộ nhớ được chia như sau
  • SRAM: lõi tính toán dùng cho mọi token
  • PSRAM: đầu ra và bộ nhớ làm việc
  • Flash: bảng 25 triệu tham số

Áp dụng Per-Layer Embeddings

• Per-Layer Embeddings từng được Google dùng trong Gemma 3n và Gemma 4 nay được áp dụng vào cấu trúc bộ nhớ của vi điều khiển thay vì điện thoại hay GPU
• Trong phạm vi tác giả dự án kiểm chứng được, chưa có tiền lệ áp dụng cách này trên một con chip nhỏ đến vậy

Mô hình làm được gì và giới hạn

• Được huấn luyện trên các câu chuyện tổng hợp ngắn của TinyStories, nên có thể tạo ra các câu chuyện đơn giản và nhìn chung giữ được tính nhất quán
• Không thể làm hỏi đáp, thực hiện lệnh, viết mã hay cung cấp kiến thức thực tế
• Giới hạn này đến từ lõi nhỏ phụ trách suy luận; kỹ thuật bố trí bộ nhớ không tự thân nâng cao năng lực suy luận

Mã nguồn và tài liệu thí nghiệm

• firmware/esp32_llm/README.md có firmware, sơ đồ nối dây và quy trình nạp flash
• src/ và experiments/ chứa mã huấn luyện, thí nghiệm loại bỏ (ablation) và lượng tử hóa
• RESULTS.md tổng hợp toàn bộ phương pháp, thí nghiệm loại bỏ và các số đo trực tiếp trên chip

Dự án nền tảng và ghi chép

• TinyStories là bộ dữ liệu các câu chuyện tổng hợp ngắn được tạo ra để ngay cả mô hình nhỏ cũng có thể học cách viết nhất quán
• Per-Layer Embeddings của Google Gemma là nền tảng giúp đưa mô hình lớn lên chip nhỏ
• llama2.c của Andrej Karpathy ảnh hưởng đến hướng tiếp cận huấn luyện mô hình ngôn ngữ nhỏ và chạy bằng C thuần
• Kho lưu trữ cũng lưu lại lỗi tính toán từng làm phóng đại số lượng tham số ban đầu cùng quá trình sửa lỗi
  • Có thể xem trong lịch sử commit và RESULTS.md để biết điểm thay đổi số liệu và lý do

[GN⁺]

Ý kiến Hacker News

• Những gì có thể làm được với vi điều khiển giá 5 USD ngày nay thật đáng kinh ngạc. Trong các bo mạch Milk-V, Duo có tới 256MB bộ nhớ và TPU 1TOPS@INT8, thậm chí còn chạy được Linux, nên tôi đã mua 5 cái

  • Cũng không thể bỏ qua ISA vector 128-bit và 32 thanh ghi. Nó hỗ trợ số nguyên và số thực dấu chấm động tối đa 64-bit; với LMUL=8, một lệnh đơn có thể xử lý 1.024 bit, và đa số phép toán mất 3 chu kỳ cho mỗi 128 bit
    GCC và Clang hỗ trợ đầy đủ xTHeadVector, và nếu dùng hàm nội tại C thì chỉ với tùy chọn dòng lệnh là đã tương thích với RVV 1.0. Nhiều đoạn mã xử lý phần tử 8-bit như memcpy(), memset(), memcmp(), strlen(), strcpy(), strcmp() có khả năng tương thích nhị phân khá tốt
    Khi tôi mua Duo 64MB thì nó có giá 3 USD, sau đó các mẫu 64MB, 256MB, 512MB lần lượt là 5, 7, 10 USD, nhưng năm nay giá đã tăng khá nhiều: https://arace.tech/products/milk-v-duo, https://arace.tech/products/milkv-duo-s
  • Cũng có thể mua vi điều khiển dưới 0,5 USD, nhưng ngay cả sản phẩm 5 USD thì các ràng buộc khi chạy mô hình vẫn quá nghiêm ngặt, nên trên thực tế có vẻ là một môi trường gần như vô dụng. Mong rằng sẽ không xảy ra tình trạng khan hiếm MCU do nhu cầu chạy mô hình
  • Không cần nghĩ rằng vì vi điều khiển như bị mắc kẹt ở phần cứng thập niên 1980 nên assembly hoặc C là lựa chọn duy nhất. Với thiết bị ở mức này, việc chạy hệ điều hành Xerox PARC cũng không hề khó
    Vẫn có những trường hợp mà PIC 4KB là lựa chọn tốt nhất, nhưng trong đa số tình huống không nên hành xử như thể không có lựa chọn nào tốt hơn
  • Mức giá và hiệu năng này khiến người ta cảm nhận rõ vị thế gần như độc quyền của ARM đã áp đặt chi phí lớn đến mức nào lên toàn ngành

• Các mô hình chuyển đổi giọng nói-văn bản cũng đang tiến gần đến kích thước này, nên tôi tự hỏi thiết bị nhỏ gọn có thể trò chuyện với chúng ta đã gần đến đâu rồi. Có lẽ sẽ tới ngày bàn chải đánh răng đưa lời khuyên vệ sinh răng miệng, thậm chí còn chèn quảng cáo kem đánh răng

  • Tư vấn vệ sinh răng miệng thì đã làm được rồi. Bàn chải cao cấp của Philips kết nối bằng Bluetooth và đưa phản hồi trong ứng dụng
    Nhưng tôi không muốn sống trong một thế giới mà còn phải tìm cả trình chặn quảng cáo cho AI trên bàn chải đánh răng

• Họ đã dùng kỹ thuật embedding theo từng lớp rất khéo. Cũng đã có các mô hình TTS thực dụng cỡ khoảng 20 đến 30 triệu tham số, nên một ESP32 không kết nối mạng có thể đọc văn bản gần như theo thời gian thực

  • Tôi muốn thử cấu hình giao cho mỗi MCU phụ trách một lớp của mô hình. Cách làm là nối nhiều RP2350 bằng đường truyền chuyên dụng PIO, rồi kết hợp bộ nội suy và lệnh nhân kép
    PSRAM, flash và thẻ SD riêng lẻ thì băng thông không cao, nhưng nếu chạy nhiều cái cùng lúc có thể đạt thông lượng đáng kể. Phần cứng chuyên dụng quy mô lớn chắc chắn sẽ vượt trội về hiệu năng trên mỗi watt, nhưng nhờ chi phí ban đầu thấp và khả năng mở rộng dần dần, kiểu cấu hình này vẫn rất hấp dẫn

• So với vi điều khiển, tôi tò mò đâu là lựa chọn thực tế để LLM cục bộ trên Raspberry Pi 4 không mất tới 30 giây cho mỗi câu trả lời

• Chạy LLM trên thiết bị siêu nhỏ cũng hay, nhưng điều ấn tượng hơn là cách huấn luyện đã tạo ra những trọng số này

  • Ở vùng tham số nhỏ, việc tối ưu hóa quá mức chỉ để nhắm vào benchmark là khó hơn, nên giá trị của bảng xếp hạng vẫn chưa hết. Ngược lại, nó vẫn còn vai trò tìm ra các mô hình nhỏ hơn nhưng đạt cùng hiệu năng benchmark, đồng thời buộc giai đoạn mô hình hóa phải nhận diện và loại bỏ khả năng nén cũng như sự dư thừa

• Tôi tự hỏi nếu tận dụng flash với mẫu truy cập tốt, liệu có thể mở rộng theo hướng chạy mô hình lớn hơn nhiều trên CPU hay không

  • Gần đây đã có một thử nghiệm tương tự. Họ chạy GLM-5.2 bằng CPU desktop phổ thông, SSD NVMe và khoảng 25GB RAM, nhưng tốc độ không phải tính theo token mỗi giây mà là tối đa 20 giây cho mỗi token: https://github.com/JustVugg/colibri
  • Băng thông flash của ESP32 chỉ bằng khoảng 1/4 SRAM nội bộ. Trên các hệ thống mạnh hơn, khoảng cách này còn lớn hơn nhiều, và để tận dụng hiệu quả năng lực tính toán tăng thêm thì băng thông bộ nhớ cần được cung cấp cũng phải tăng mạnh

• ESP32-S3 khá mạnh, đến mức hiện tôi đang dùng nó cho công việc phát triển Raspberry Pi 4. Một trong hai cổng USB hỗ trợ OTG, nên có thể hiện thực những tính năng mà nếu theo cách khác sẽ tốn hơn 100 USD

  • Nếu chấp nhận một số hạn chế thì có thể làm cấu hình tương tự với RP2350 với giá khoảng 1 USD. Tốc độ tối đa là USB Full Speed 12Mbps; một cổng dùng ngoại vi USB tích hợp trong chip, còn cổng kia nối vào chân GPIO được PIO hỗ trợ
    Hiện nó chạy với tinyusb và pico-pio-usb, và tôi cũng đang thử port sang Rust với hy vọng đạt hiệu năng cao hơn

• Tôi tò mò độ chính xác của mô hình lượng tử hóa này ở mức nào

• Hiệu năng ở kích thước này thật đáng kinh ngạc, và điều đó khiến tôi háo hức muốn biết có thể làm được gì trên những máy tính một bo mạch mạnh hơn một chút được nhắc tới trong chủ đề