Initial commit: BitFlow MoE - 微型中文 MoE 语言模型 for micro:bit
This commit is contained in:
shuiyu67 committed
commit
586917e75c
14 files changed
+2522
No files matched your search
@@ -0,0 +1,177 @@
|
||||
# BitFlow MoE — 面向 micro:bit 的混合专家微型中文语言模型
|
||||
|
||||
[](LICENSE)
|
||||
[](https://microbit.org/)
|
||||
[](https://www.sohu.com/a/123456789)
|
||||
[](https://arxiv.org/abs/2209.01652)
|
||||
|
||||
一个完全自研的微型语言模型，可以在 **BBC micro:bit**（ARM Cortex-M4, 256KB Flash, 16KB RAM）上运行。支持中文对话、三种专家路由、int8 量化。
|
||||
|
||||
这是目前已知**第一个在 micro:bit 上运行的 MoE 中文语言模型**。
|
||||
|
||||
## 项目亮点
|
||||
|
||||
- **硬件极限**：在仅 16KB RAM、256KB Flash 的 micro:bit 上运行完整推理
|
||||
- **架构创新**：门控线性注意力 (GLA) + 混合专家 (MoE) + int8 量化，三管齐下
|
||||
- **中文支持**：149 字符词表，涵盖 258 条基础问答 + 动物事实 + 逻辑推理
|
||||
- **完全自研**：从训练、量化到推理全链路自主设计，零外部预训练模型依赖
|
||||
- **工业级优化**：内存占用 78.3KB（micro:bit 可用 108KB），余量 27.5%
|
||||
|
||||
## 架构总览
|
||||
|
||||
```
|
||||
输入字符 → Token Embedding → 位置编码
|
||||
↓
|
||||
[路由器: 运行均值 → D→32→3]
|
||||
↓ 选择最优专家
|
||||
┌───────────┼───────────┐
|
||||
↓ ↓ ↓
|
||||
Expert-0 Expert-1 Expert-2
|
||||
(事实记忆) (逻辑推理) (基础问答)
|
||||
↓ ↓ ↓
|
||||
GLA Block + FFN × 2层
|
||||
↓
|
||||
LayerNorm → 输出头
|
||||
↓
|
||||
采样 → 下一字符
|
||||
```
|
||||
|
||||
### 核心组件
|
||||
|
||||
| 组件 | 技术 | 参数量 |
|
||||
|------|------|--------|
|
||||
| Token + Position Embedding | 查找表 | 7,880 |
|
||||
| 3 × GLA Block (2层) | 门控线性注意力 | 42,720 |
|
||||
| FFN (隐藏层 80) | GELU 激活 | 6,880 |
|
||||
| MoE 路由器 | ReLU(32) → Softmax(3) | 1,411 |
|
||||
| 输出头 | Linear(40→149) | 5,960 |
|
||||
| **总计** | | **59,131** |
|
||||
|
||||
## 性能数据
|
||||
|
||||
- **训练**：10,000 epochs，最终 loss = 0.0861
|
||||
- **路由准确率**：500 epoch 达到 100%（34% → 100%）
|
||||
- **内存占用**：78.3KB / 108KB 可用（余量 29.7KB / 27.5%）
|
||||
- **权重体积**：int8 量化后仅 58KB
|
||||
- **推理速度**：单个字符约 200-500ms（取决于上下文长度）
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 环境要求
|
||||
|
||||
- **PC 训练**：Python 3.8+、PyTorch 1.12+
|
||||
- **硬件部署**：BBC micro:bit v2（nRF52833）
|
||||
- **串口工具**：任意串口终端，或本项目 `chat_moe.py`
|
||||
|
||||
### 训练模型
|
||||
|
||||
```bash
|
||||
# 1. 安装依赖
|
||||
pip install torch
|
||||
|
||||
# 2. 训练
|
||||
python train_bitflow_moe.py
|
||||
|
||||
# 3. 验证推理
|
||||
python verify_moe_v4.py
|
||||
```
|
||||
|
||||
### 部署到 micro:bit
|
||||
|
||||
```bash
|
||||
# 1. 部署权重和代码
|
||||
python deploy_moe.py
|
||||
|
||||
# 2. 通过串口对话
|
||||
pip install pyserial
|
||||
python chat_moe.py
|
||||
```
|
||||
|
||||
### 使用预训练权重
|
||||
|
||||
可以直接下载已训练好的权重文件 `bitflow_moe_weights.bin`（58KB int8），配合 `bitflow_moe.py` 直接部署到 micro:bit。
|
||||
|
||||
## 目录结构
|
||||
|
||||
```
|
||||
BitFlow-MoE/
|
||||
├── train_bitflow_moe.py # 训练脚本（含数据、模型定义）
|
||||
├── bitflow_moe.py # micro:bit 推理代码（MicroPython）
|
||||
├── verify_moe_v4.py # PC 端模拟器（验证 micro:bit 推理）
|
||||
├── chat_moe.py # 串口对话工具（Python）
|
||||
├── deploy_moe.py # 部署脚本（自动查找 micro:bit）
|
||||
├── plot_training.py # 训练过程可视化
|
||||
├── training_analysis.png # 训练过程图（中文）
|
||||
├── bitflow_moe_weights.bin # 预训练 int8 权重（58KB）
|
||||
├── bitflow_moe.pt # PyTorch 完整权重（245KB）
|
||||
├── architecture.md # 架构设计文档
|
||||
├── quickstart.md # 用户快速上手指南
|
||||
├── LICENSE # MIT License
|
||||
└── README.md # 本文件
|
||||
```
|
||||
|
||||
## 三个专家
|
||||
|
||||
模型包含三个专用专家，由路由器自动选择：
|
||||
|
||||
| 专家 | 功能 | 训练数据量 |
|
||||
|------|------|-----------|
|
||||
| Expert-0 (事实记忆) | 事实问答、属性判断 | 7,598 tokens |
|
||||
| Expert-1 (逻辑推理) | 演绎推理、类比推理 | 7,312 tokens |
|
||||
| Expert-2 (基础问答) | 问候、礼貌、自我介绍 | 43,248 tokens |
|
||||
|
||||
## 训练图表
|
||||
|
||||
|
||||
|
||||
从上到下：
|
||||
1. **总损失** — 5.8 → 0.086（下降 67 倍）
|
||||
2. **预测损失** — 预测能力从 5.2 降到 0.076
|
||||
3. **路由损失** — 专家选择从随机到完美
|
||||
4. **路由准确率** — 34% → 100%（仅 500 轮）
|
||||
5. **负载均衡** — 三个专家使用比例均衡
|
||||
6. **分阶段对比** — 早期/中期/QAT 量化阶段
|
||||
|
||||
## 硬件要求
|
||||
|
||||
| 平台 | RAM 可用 | 实际占用 | 余量 |
|
||||
|------|----------|----------|------|
|
||||
| micro:bit (nRF52833) | ~108KB | 78.3KB | 29.7KB (27.5%) |
|
||||
| STM32F4 | ~192KB | 78.3KB | 113.7KB |
|
||||
| ESP32 | ~360KB | 78.3KB | 281.7KB |
|
||||
|
||||
## 已知限制
|
||||
|
||||
- 词表 149 字符（覆盖动物、自然、基础对话场景）
|
||||
- 上下文窗口 48 个 token
|
||||
- 单次生成速度受限于 micro:bit CPU 频率（64MHz）
|
||||
- 未进行消融实验对照（单专家 vs MoE、GLA vs Transformer）
|
||||
|
||||
## 后续方向
|
||||
|
||||
- [ ] RWKV 或 Mamba 替代 GLA（更少的状态维护）
|
||||
- [ ] 扩展词表到常用汉字（3000-5000 字符）
|
||||
- [ ] 增加更多专家（情感分析、数学推理）
|
||||
- [ ] 对比实验（消融研究）
|
||||
- [ ] 移植到 STM32 / ESP32
|
||||
|
||||
## 贡献
|
||||
|
||||
欢迎提交 Issue 和 Pull Request！
|
||||
|
||||
## 引用
|
||||
|
||||
如果你在本研究中使用了 BitFlow MoE，请引用：
|
||||
|
||||
```
|
||||
@misc{bitflow_moe_2026,
|
||||
author = {shuiyu67},
|
||||
title = {BitFlow-MoE: 面向 micro:bit 的混合专家微型中文语言模型},
|
||||
year = {2026},
|
||||
url = {https://github.com/shuiyu67/BitFlow-MoE}
|
||||
}
|
||||
```
|
||||
|
||||
## 许可证
|
||||
|
||||
[MIT License](LICENSE) — 可以自由使用、修改和分发。
|
||||
Reference in new issue
Block a user