diff --git a/.gitignore b/.gitignore deleted file mode 100644 index fe2645f..0000000 --- a/.gitignore +++ /dev/null @@ -1,82 +0,0 @@ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -*.egg-info/ -.installed.cfg -*.egg - -# PyInstaller -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.py,cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Environments -.env -.venv -env/ -venv/ -ENV/ - -# IDE -.vscode/ -.idea/ -*.swp -*.swo -*~ - -# micro:bit -*.hex - -# Training artifacts -*.pt - -# Large binary files -*.bin -!train_bitflow_moe.py - -# OS files -.DS_Store -Thumbs.db - -# Debug logs -*.log diff --git a/LICENSE b/LICENSE deleted file mode 100644 index 7203af6..0000000 --- a/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2026 shuiyu - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/README.md b/README.md index 9852ec6..db3ed1b 100644 --- a/README.md +++ b/README.md @@ -2,8 +2,7 @@ 你好呀～这个仓库已经搬到更好用的地方啦。 -为了让代码管理更清爽、多平台同步更省心，我们把原先分散的镜像统一收拢到 -**Gitee** 与 **GitHub** 主仓维护，这里就只保留这张"迁移小纸条"咯。 +为了让代码管理更清爽、多平台同步更省心，我们把原先分散的镜像统一收拢到主仓维护，这里就只保留这张"迁移小纸条"咯。 📦 新家在这里： - Gitee：https://gitee.com/shuiyu1123/BitFlow-MoE diff --git a/architecture.md b/architecture.md deleted file mode 100644 index f162553..0000000 --- a/architecture.md +++ /dev/null @@ -1,213 +0,0 @@ -# BitFlow MoE 架构设计文档 - -## 1. 设计目标 - -在 BBC micro:bit（ARM Cortex-M4, 64MHz, 16KB RAM, 256KB Flash）上运行一个具有**理解能力**的微型中文语言模型。 - -核心挑战： -- **内存限制**：可用 RAM ~108KB，模型权重 + 推理状态必须在这个范围内 -- **计算限制**：64MHz CPU，无 FPU（但有硬件乘法） -- **存储限制**：256KB Flash，预训练模型通常几百 MB，需极致压缩 - -## 2. 架构选型 - -### 2.1 为什么不用 Transformer？ - -标准 Transformer（即使是 2 层微型版）在 micro:bit 上存在问题： - -| 组件 | 标准 Attention | GLA (我们的选择) | -|------|---------------|------------------| -| 注意力矩阵 | O(n²) 状态矩阵 | O(d) 累积矩阵 | -| 最大上下文 | ~16 tokens | 48 tokens | -| 内存峰值 | 32KB+（仅注意力） | <5KB | -| 推理速度 | 极慢（会卡死） | 实时 | - -**结论**：标准 Attention 的 O(n²) 状态维护在 micro:bit 上是灾难性的。GLA 将复杂度降到 O(n×d)，状态矩阵从 n×n 压缩到 d×d。 - -### 2.2 为什么不用纯 RNN？ - -RNN 虽然内存 O(d)，但： -- 无法并行训练（梯度消失/爆炸） -- 长程依赖能力弱 -- 无法做到像我们这样的"理解"效果 - -GLA 是 RNN 的改进版，保留了 O(d) 内存，但增加了**可并行训练**的能力（通过离散时间注意力）。 - -## 3. 核心架构 - -### 3.1 门控线性注意力 (GLA) - -GLA 的核心公式： - -``` -S[t] = g[t] ⊙ S[t-1] + k[t] ⊗ v[t] -z[t] = g[t] ⊙ z[t-1] + k[t] -y[t] = Q[t] · S[t] · (Q[t] · z[t])⁻¹ -``` - -其中： -- **S** 是 d×d 累积状态矩阵（每次更新 O(d²)） -- **k, q, v** 是通过 Linear 投影得到的向量 -- **g** 是门控（gate），通过 sigmoid(log_decay) 控制衰减 -- **⊙** 是 Hadamard 乘积，**⊗** 是外积 - -**关键优化**： -1. 不使用矩阵求逆，而是通过标量除法实现 -2. 状态矩阵 S 只存储 d×d=40×40=1600 个浮点数（6.4KB） -3. 每次更新的复杂度是 O(d²) 而非 O(n²) - -### 3.2 混合专家 (MoE) - -``` -输入 → [Token Embedding] + [Position Embedding] - ↓ - [路由器: mean pooling → FC → argmax] - ↓ 选择最优专家 - ┌── Expert-0 (事实记忆) - ├── Expert-1 (逻辑推理) - └── Expert-2 (基础问答) - ↓ - [LayerNorm → GLA Block × 2] - ↓ - [LayerNorm → Output Head] - ↓ - 采样输出 -``` - -**为什么 MoE 有效**： -- 每个专家专精一类任务，不需要一个"全能"模型 -- 路由选择只激活一个专家，**推理计算量不增加** -- 三个专家共享相同的结构（GLA Block × 2），代码复用率高 - -### 3.3 路由器设计 - -```python -def router(x): - # 1. 对所有 token 做 mean pooling → 固定维度 d=40 - x_mean = x.mean(dim=1) - # 2. 全连接层: 40 → 32 → 3 - h = relu(fc1(x_mean)) # 40 → 32 - logits = fc2(h) # 32 → 3 - return logits # argmax 选专家 -``` - -**训练时**：使用交叉熵损失监督路由（labels 来自数据类别） -**推理时**：只计算一次路由，选择最优专家 - -## 4. 量化策略 - -### 4.1 int8 量化 - -``` -scale = max(|w|) / 127 -q_int8 = round(w / scale) # [-128, 127] -w_rec = q_int8 × scale # 反量化 -``` - -- **量化精度**：int8（8-bit 有符号整数） -- **量化感知训练 (QAT)**：训练后 2000 步（8000-10000 epoch）引入量化噪声，让模型适应 int8 精度 -- **结果**：loss 从 0.0861 仅上升到 0.0863，影响极小 - -### 4.2 内存分析 - -| 组件 | 全精度(float32) | int8 量化 | 压缩比 | -|------|----------------|-----------|--------| -| 模型权重 | 236.5KB | 57.7KB | **4.1×** | -| 推理状态 | ~70KB | ~70KB | 无（运行时需要float） | -| **总计** | **306.5KB** | **127.7KB** | **2.4×** | - -## 5. 数据工程 - -### 5.1 三层数据体系 - -``` -基础层 (BASE, 69 条) - └── 动物行为、颜色判断、位置关系 - ↓ 提供基本语义 -事实层 (FACT_QA, 76 条) - └── 什么是XX、XX吃什么、XX在哪里 - ↓ 提供知识推理 -问答层 (CHAT_QA, 258 条) - └── 问候、礼貌、自我介绍、喜好、情感、能力 - ↓ 提供交互能力 -``` - -### 5.2 数据增强 - -- 基础层 × 20 遍 -- 事实层 × 40 遍 -- 问答层 × 80 遍 - -总数据量：290,780 tokens，其中问答类占 74% - -## 6. 训练流程 - -``` -[10,000 epochs] - │ - ├─ 0 - 8000 epoch: 标准训练 - │ ├─ 学习率: 3e-3 (Cosine Annealing) - │ ├─ Loss = 预测损失 + 0.5 × 路由损失 + 0.01 × 负载均衡损失 - │ └─ 路由准确率: 34% → 100% - │ - └─ 8000 - 10000 epoch: QAT 量化感知训练 - ├─ 每步对权重做 int8 量化 + 反量化 - ├─ 模型适应量化噪声 - └─ loss 微升后稳定 -``` - -## 7. 推理流程 - -``` -[初始化] 加载 int8 权重 → 预解量化小权重 (SMW) - │ - ├─ [种子处理] 逐 token 前向传播，积累运行均值 - │ └─ 路由器根据均值选择专家，激活对应 GLA Block - │ - ├─ [生成循环] - │ ├─ 计算 token embedding - │ ├─ 累加运行均值 → 路由 → 选专家 - │ ├─ 专家 GLA Block × 2 层 - │ ├─ LayerNorm → 输出头 - │ ├─ 采样 (temperature=0.3) - │ └─ 更新频率惩罚 (bcnt) - │ - └─ [输出] 字符 ID 列表 → 串口发送 -``` - -## 8. micro:bit 内存布局 - -``` -128KB Flash (总容量) -├── ~20KB: MicroPython 解释器 + 标准库 -├── ~58KB: int8 权重文件 -├── ~7KB: 推理代码 (bitflow_moe.py) -└── ~43KB: 可用 (代码+数据段) - -16KB RAM (总容量) -├── ~108KB: 实际可用 (nRF52833 有 256KB RAM, 但 MicroPython 用了一部分) -├── 58KB: 解量化后权重 (运行时) -├── 13KB: GLA 状态矩阵 (2层 × 40×40) -├── 4KB: 工作缓冲区 -├── 2KB: 预解量化小权重 -└── ~30KB: 余量 (防止栈溢出) -``` - -## 9. 设计决策总结 - -| 决策 | 选择 | 理由 | -|------|------|------| -| 注意力机制 | GLA | O(d) 内存，可并行训练 | -| MoE 结构 | 3 专家硬路由 | 推理计算量不增加 | -| 量化 | int8 + QAT | 体积压缩 4×，精度损失 <1% | -| 训练策略 | 两步训练 + QAT | 先学语义，再适应量化 | -| 路由方式 | 监督学习 | 保证每个专家都有任务 | -| 采样温度 | 0.3 | 偏低，保证输出一致性 | - -## 10. 局限性 - -1. **GLA 的衰减机制**：log_decay 是标量，无法按维度自适应（Mamba 可以做） -2. **硬路由**：每次只激活一个专家，无法融合多个专家的知识 -3. **词表太小**：149 字符覆盖场景有限 -4. **无上下文理解**：48 token 窗口限制了长对话能力 -5. **未做消融实验**：无法精确量化每个模块的贡献 diff --git a/bitflow_moe.pt b/bitflow_moe.pt deleted file mode 100644 index 8acb4a9..0000000 Binary files a/bitflow_moe.pt and /dev/null differ diff --git a/bitflow_moe.py b/bitflow_moe.py deleted file mode 100644 index 3ade468..0000000 --- a/bitflow_moe.py +++ /dev/null @@ -1,192 +0,0 @@ -from microbit import * -import struct,gc,time,random,math -from array import array -uart.init(115200) -# BitFlow MoE v4 - micro:bit 推理 -# DIM=40 FFN=80 3专家 词表149 (~58K int8权重) -VS=149;D=40;CT=48;FD=80;NE=3;NL=2;RD=32 -RB=2+NE*12;NB=RB+4;HI=NB+2 -# 加载权重 -wf=open("weights.bin","rb") -nt=struct.unpack('=128:v-=256 - r[j]=v*s - return r -SMW=[] -for e in range(NE): - b0=2+e*12 - SMW.append(dqS(b0));SMW.append(dqS(b0+1));SMW.append(dqS(b0+2)) - SMW.append(dqS(b0+8));SMW.append(dqS(b0+9)) -rf1b=dqS(RB+1);rf2b=dqS(RB+3) -nfw=dqS(NB);nfb=dqS(NB+1) -del dqS;gc.collect() -# 融合解量化+矩阵乘 (省去DB大缓冲区) -def dqmv(i,ni,no,x,out): - s=SC[i];o=OFF[i] - for r in range(no): - ss=0.0;b=r*ni - for c in range(ni): - v=W[o+b+c] - if v>=128:v-=256 - ss+=v*s*x[c] - out[r]=ss -def ln(w,b,x,o): - m=0.0 - for i in range(D):m+=x[i] - m/=D;v=0.0 - for i in range(D):t=x[i]-m;v+=t*t - v/=D;iv=1.0/math.sqrt(v+1e-5) - for i in range(D):o[i]=w[i]*(x[i]-m)*iv+b[i] -def e1(x):return x+1.0 if x>=0 else math.exp(x) -def sg(x): - if x>=0:e=math.exp(-x);return 1.0/(1.0+e) - e=math.exp(x);return e/(1.0+e) -def ge(x): - c=0.7978845608 - return x*sg(2.0*c*(x+0.044715*x*x*x)) -def sm(l,t=0.3): - m=-1e9 - for i in range(VS): - v=l[i]-bcnt[i]*2.0 - if v>m:m=v - s=0.0 - for i in range(VS): - v=l[i]-bcnt[i]*2.0 - e=math.exp((v-m)/t) - bpr[i]=e;s+=e - r=random.random()*s;c=0.0 - for i in range(VS): - c+=bpr[i] - if c>=r:return i - return VS-1 -# GLA状态更新 (单层) -def gs(e,layer,h): - b0=2+e*12;dec=SMW[e*5+2] - dqmv(b0+3,D,D,h,bq) - dqmv(b0+4,D,D,h,bk) - dqmv(b0+5,D,D,h,bv) - dqmv(b0+6,D,D,h,bgd);gc.collect() - for i in range(D):bq[i]=e1(bq[i]);bk[i]=e1(bk[i]) - for i in range(D):bgd[i]=sg(bgd[i])*dec[i] - sl=S[layer];zl=z[layer] - for i in range(D): - si=i*D;gd=bgd[i];ki=bk[i] - for j in range(D):sl[si+j]=gd*sl[si+j]+ki*bv[j] - for i in range(D):zl[i]=bgd[i]*zl[i]+bk[i] - for j in range(D):bnum[j]=0.0 - for i in range(D): - si=i*D;qi=bq[i] - for j in range(D):bnum[j]+=qi*sl[si+j] - den=1e-6 - for i in range(D):den+=bq[i]*zl[i] - iv=1.0/den - for j in range(D):by[j]=bnum[j]*iv - dqmv(b0+7,D,D,by,bout);gc.collect() -# 专家块 (残差+GLA+FFN) -def bs(e,x,layer): - b0=2+e*12 - ln(SMW[e*5],SMW[e*5+1],x,bh) - gs(e,layer,bh) - for i in range(D):x[i]+=bout[i] - ln(SMW[e*5+3],SMW[e*5+4],x,bh) - dqmv(b0+10,D,FD,bh,bff) - for i in range(FD):bff[i]=ge(bff[i]) - dqmv(b0+11,FD,D,bff,bout);gc.collect() - for i in range(D):x[i]+=bout[i] -# 路由器 (D→RD→NE) -def route(x): - dqmv(RB,D,RD,x,brh) - for i in range(RD):brh[i]=max(0.0,brh[i]+rf1b[i]) - dqmv(RB+2,RD,NE,brh,brl) - for i in range(NE):brl[i]+=rf2b[i] - bi=0;bv2=brl[0] - for i in range(1,NE): - if brl[i]>bv2:bv2=brl[i];bi=i - return bi -# 前向: gen_flag=False时跳过head(种子处理加速) -def fw(ctx,gen_flag=True): - t=min(len(ctx)-1,CT-1);ti=ctx[-1] - s0=SC[0];o0=OFF[0];s1=SC[1];o1=OFF[1] - for i in range(D): - v0=W[o0+ti*D+i] - if v0>=128:v0-=256 - v1=W[o1+t*D+i] - if v1>=128:v1-=256 - bx[i]=v0*s0+v1*s1 - gc.collect() - for i in range(D):bsum[i]+=bx[i] - bct[0]+=1 - for i in range(D):brx[i]=bsum[i]/bct[0] - ei=route(brx) - for layer in range(NL):bs(ei,bx,layer) - ln(nfw,nfb,bx,bh) - if gen_flag: - dqmv(HI,D,VS,bh,blg);gc.collect() - return blg,ei - return None,ei -# 重置状态 -def rs(): - for l in range(NL): - for i in range(D*D):S[l][i]=0.0 - for i in range(D):z[l][i]=0.0 - for i in range(D):bsum[i]=0.0 - bct[0]=0 -# 生成 -def gen(ids,n): - NL_ch=VS-1 - ids=list(ids)+[NL_ch] - sl=len(ids);c=list(ids);rs();gc.collect() - for i in range(VS):bcnt[i]=0 - uart.write("S:%s\r\n"%",".join(str(i) for i in c)) - for i in range(len(c)-1): - fw(c[:i+1],False);gc.collect() - for _ in range(n): - t0=time.ticks_ms() - lg,ei=fw(c) - nx=sm(lg) - dt=time.ticks_diff(time.ticks_ms(),t0) - uart.write("T:%d %dms E%d\r\n"%(nx,dt,ei)) - bcnt[nx]+=1 - del lg;gc.collect();c.append(nx) - uart.write("O:%s\r\n"%",".join(str(i) for i in c[sl:])) -gc.collect() -print("READY") -print("free=%d"%gc.mem_free()) -while True: - if uart.any(): - l=uart.readline() - if l: - l=l.strip() - if l[:4]==b"GEN:": - p=l[4:].split(b"|") - ids=[int(x) for x in p[0].split(b",") if x.strip()] - nn=int(p[1]) if len(p)>1 else 10 - if ids:gen(ids,nn) - sleep(50) diff --git a/bitflow_moe_weights.bin b/bitflow_moe_weights.bin deleted file mode 100644 index d1c3766..0000000 Binary files a/bitflow_moe_weights.bin and /dev/null differ diff --git a/chat_moe.py b/chat_moe.py deleted file mode 100644 index e311096..0000000 --- a/chat_moe.py +++ /dev/null @@ -1,184 +0,0 @@ -""" -串口发送工具: 通过串口向 micro:bit 发送对话命令 -用法: - python chat_moe.py # 自动查找串口 - python chat_moe.py COM3 # 指定串口 -""" -import sys -import glob -import os - -try: - import serial - import serial.tools.list_ports -except ImportError: - print("[ERROR] 需要安装 pyserial:") - print(" pip install pyserial") - sys.exit(1) - - -def find_serial(): - """自动查找 micro:bit 串口""" - ports = serial.tools.list_ports.comports() - for p in ports: - # 检查是否是 micro:bit (VID:PID = 0d28:0204) - if p.vid == 0x0d28 and p.pid == 0x0204: - return p.device - # 也检查名称包含 MICROBIT 的 - if "MICROBIT" in p.description.upper(): - return p.device - return None - - -def send_command(port, text): - """发送文本并转换为ID命令""" - VOCAB_CHARS = ( - "猫狗鸟鱼猪牛马鸡虫羊兔鸭鹅" - "坐跑飞游看吃睡走听来叫爬跳唱画读做" - "在上里到" - "的是好了" - "吗什么怎哪" - "你我他她们它" - "蓝绿红黄白" - "天地水火花草树木日月星云春夏秋冬" - "大小多高远快" - "有没会不" - "人头脚手眼腿" - "垫肉书子" - "一三只两" - "热冷落开" - "自然动植物质色样颜得都" - "你好吗谢谢再见早晚安静" - "很棒爱喜欢乐伤心气" - "名字叫谁和问答应" - "聊天话今天说" - "客用" - "bit" - "\n" - ) - CH2I = {c: i for i, c in enumerate(VOCAB_CHARS)} - ctx_ids = [str(CH2I[c]) for c in text if c in CH2I] - if not ctx_ids: - print(f"[warn] 无法识别: '{text}'") - return - cmd = "GEN:" + ",".join(ctx_ids) + "|15" - print(f" 发送: {cmd}") - port.write((cmd + "\n").encode()) - - -def main(): - print("=" * 60) - print("BitFlow MoE v4 对话工具") - print("=" * 60) - - # 串口 - if len(sys.argv) > 1: - port_name = sys.argv[1] - else: - print("\n[search] 正在查找 micro:bit 串口 ...") - port_name = find_serial() - if port_name is None: - print("[ERROR] 未找到 micro:bit 串口!") - print(" 请确保:") - print(" 1. micro:bit 已通过 USB 连接") - print(" 2. main.py 已在 micro:bit 上运行") - print(" 3. 如已知串口, 可指定: python chat_moe.py COM3") - return - - print(f"\n[connect] 打开串口: {port_name} @115200") - try: - ser = serial.Serial(port_name, 115200, timeout=2) - except Exception as e: - print(f"[ERROR] 无法打开串口: {e}") - return - print(" 连接成功!") - print() - - # 等待 ready - ready = False - for _ in range(20): - if ser.in_waiting: - line = ser.readline().decode('utf-8', errors='ignore').strip() - print(f"< {line}") - if "READY" in line: - ready = True - print() - break - if not ready: - print("[warn] 未收到 READY 信号，micro:bit 可能还在启动...") - - # 常用短语列表 - phrases = [ - "你好", "你好吗", "谢谢", "再见", "早安", "晚安", - "你叫什么名字", "你喜欢什么", "你喜欢猫吗", "你是谁", - "你会什么", "你开心吗", "你在做什么", "今天好吗", - "什么会飞", "什么会跑", "什么会游", "什么吃草", - "天是什么颜色", "猫吃什么", "猫是什么", - ] - - print("=" * 60) - print("常用短语 (输入序号发送, 或直接输入文字, 输入 q 退出):") - print("=" * 60) - for i, p in enumerate(phrases): - print(f" {i:2d}. {p}") - print() - - try: - while True: - try: - query = input(">>> ").strip() - except EOFError: - break - - if query.lower() == 'q': - break - - if query == '': - continue - - # 检查是否是序号 - try: - idx = int(query) - if 0 <= idx < len(phrases): - print(f"\n 选择: {phrases[idx]}") - send_command(ser, phrases[idx]) - # 读取输出 - for _ in range(5): - if ser.in_waiting: - line = ser.readline().decode('utf-8', errors='ignore').strip() - if line.startswith('O:'): - ids = [int(x) for x in line[2:].split(',') if x.strip()] - text = "".join(VOCAB_CHARS.get(i, '?') for i in ids) - print(f"\n 回答: {text}") - print() - break - else: - print(f" {line}") - continue - except ValueError: - pass - - # 直接输入文字 - send_command(ser, query) - # 读取输出 - for _ in range(5): - if ser.in_waiting: - line = ser.readline().decode('utf-8', errors='ignore').strip() - if line.startswith('O:'): - ids = [int(x) for x in line[2:].split(',') if x.strip()] - text = "".join(VOCAB_CHARS.get(i, '?') for i in ids) - print(f"\n 回答: {text}") - print() - break - else: - print(f" {line}") - - except KeyboardInterrupt: - pass - finally: - ser.close() - print("\n[done] 已断开连接") - - -if __name__ == "__main__": - main() diff --git a/deploy_moe.py b/deploy_moe.py deleted file mode 100644 index c58a3db..0000000 --- a/deploy_moe.py +++ /dev/null @@ -1,191 +0,0 @@ -""" -BitFlow MoE v4 部署脚本 -将训练好的 int8 权重和推理代码部署到 micro:bit - -用法: - python deploy_moe.py # 自动查找 micro:bit 并部署 - python deploy_moe.py D:\\ # 指定 micro:bit 盘符 -""" -import os -import sys -import shutil -import glob - -# 文件路径 -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -WEIGHTS_FILE = os.path.join(SCRIPT_DIR, "bitflow_moe_weights.bin") -INFER_CODE = os.path.join(SCRIPT_DIR, "bitflow_moe.py") - -# micro:bit 上的文件名 -TARGET_WEIGHTS = "weights.bin" -TARGET_MAIN = "main.py" - - -def find_microbit(): - """自动查找 micro:bit USB 盘符""" - # Windows: 查找 MICROBIT 卷标 - for letter in "ABCDEFGHIJKLMNOPQRSTUVWXYZ": - drive = f"{letter}:\\" - if os.path.exists(drive): - # 检查是否是 micro:bit - label_path = os.path.join(drive, "DETAILS.TXT") - if os.path.exists(label_path): - try: - with open(label_path, "r") as f: - content = f.read() - if "micro:bit" in content.lower() or "V2" in content or "nRF52" in content: - return drive - except: - pass - # 也检查卷标 - vol_label = os.path.join(drive, "MICROBIT.HTM") - if os.path.exists(vol_label): - return drive - return None - - -def deploy(target_drive): - """部署文件到 micro:bit""" - target_weights = os.path.join(target_drive, TARGET_WEIGHTS) - target_main = os.path.join(target_drive, TARGET_MAIN) - - # 检查源文件 - if not os.path.exists(WEIGHTS_FILE): - print(f"[ERROR] 权重文件不存在: {WEIGHTS_FILE}") - print(" 请先运行: python train_bitflow_moe.py") - return False - - if not os.path.exists(INFER_CODE): - print(f"[ERROR] 推理代码不存在: {INFER_CODE}") - return False - - # 权重文件大小 - w_size = os.path.getsize(WEIGHTS_FILE) - print(f"\n[deploy] 部署到 micro:bit") - print(f" 盘符: {target_drive}") - print(f" 权重: {WEIGHTS_FILE} ({w_size/1024:.1f}KB)") - print(f" 代码: {INFER_CODE}") - - # 拷贝权重 - print(f"\n [1/2] 拷贝权重 -> {TARGET_WEIGHTS} ...") - try: - shutil.copy2(WEIGHTS_FILE, target_weights) - print(f" OK ({w_size/1024:.1f}KB)") - except Exception as e: - print(f" FAILED: {e}") - return False - - # 拷贝代码 (重命名为 main.py) - print(f" [2/2] 拷贝代码 -> {TARGET_MAIN} ...") - try: - shutil.copy2(INFER_CODE, target_main) - code_size = os.path.getsize(target_main) - print(f" OK ({code_size/1024:.1f}KB)") - except Exception as e: - print(f" FAILED: {e}") - return False - - # 验证 - print(f"\n[verify] 部署验证:") - if os.path.exists(target_weights): - print(f" {TARGET_WEIGHTS}: {os.path.getsize(target_weights)} bytes OK") - else: - print(f" {TARGET_WEIGHTS}: MISSING!") - return False - - if os.path.exists(target_main): - print(f" {TARGET_MAIN}: {os.path.getsize(target_main)} bytes OK") - else: - print(f" {TARGET_MAIN}: MISSING!") - return False - - total = os.path.getsize(target_weights) + os.path.getsize(target_main) - print(f"\n[done] 部署完成! 总计 {total/1024:.1f}KB") - print(f"\n micro:bit 将自动重启并运行 main.py") - print(f" 通过串口 (115200) 发送: GEN:0,1,5|20") - print(f" (ID格式: 逗号分隔的字符ID, |后为生成长度)") - print(f"\n 字符ID对照表见下方:") - return True - - -def print_vocab(): - """打印字符ID对照表""" - VOCAB_CHARS = ( - "猫狗鸟鱼猪牛马鸡虫羊兔鸭鹅" - "坐跑飞游看吃睡走听来叫爬跳唱画读做" - "在上里到" - "的是好了" - "吗什么怎哪" - "你我他她们它" - "蓝绿红黄白" - "天地水火花草树木日月星云春夏秋冬" - "大小多高远快" - "有没会不" - "人头脚手眼腿" - "垫肉书子" - "一三只两" - "热冷落开" - "自然动植物质色样颜得都" - "你好吗谢谢再见早晚安静" - "很棒爱喜欢乐伤心气" - "名字叫谁和问答应" - "聊天话今天说" - "客用" - "bit" - "\n" - ) - print(f"\n{'='*60}") - print("字符ID对照表 (用于串口通信)") - print(f"{'='*60}") - for i, c in enumerate(VOCAB_CHARS): - display = "\\n" if c == "\n" else c - print(f" {i:3d}: {display} ", end="") - if (i + 1) % 8 == 0: - print() - print() - print(f"\n常用短语ID:") - examples = { - "你好": "0,1", # 需要查实际ID - } - # 生成常用短语 - phrases = ["你好", "你好吗", "谢谢", "再见", "早安", "晚安", - "你叫什么名字", "你喜欢什么", "什么会飞", "什么会跑"] - CH2I = {c: i for i, c in enumerate(VOCAB_CHARS)} - for p in phrases: - ids = [str(CH2I[c]) for c in p if c in CH2I] - print(f" {p}: {','.join(ids)}") - print(f"\n 生成命令: GEN:|<生成长度>") - print(f" 示例: GEN:{','.join(str(CH2I[c]) for c in '你好' if c in CH2I)}|10") - - -def main(): - print("=" * 60) - print("BitFlow MoE v4 部署工具") - print("=" * 60) - - # 查找 micro:bit - if len(sys.argv) > 1: - target = sys.argv[1] - if not os.path.exists(target): - print(f"[ERROR] 指定的盘符不存在: {target}") - return - else: - print("\n[search] 正在查找 micro:bit ...") - target = find_microbit() - if target is None: - print("[ERROR] 未找到 micro:bit!") - print("\n请确保:") - print(" 1. micro:bit 已通过 USB 连接") - print(" 2. micro:bit 已出现在文件管理器中") - print(" 3. 如已知盘符, 可指定: python deploy_moe.py D:\\") - return - - # 部署 - success = deploy(target) - - if success: - print_vocab() - - -if __name__ == "__main__": - main() diff --git a/plot_training.py b/plot_training.py deleted file mode 100644 index c31f877..0000000 --- a/plot_training.py +++ /dev/null @@ -1,172 +0,0 @@ -""" -从训练输出解析数据并绘制训练过程图 -""" -import re -import matplotlib -matplotlib.use('Agg') -import matplotlib.pyplot as plt -import numpy as np - -# 设置中文字体 -plt.rcParams['font.sans-serif'] = ['Noto Sans CJK SC', 'WenQuanYi Micro Hei', 'Microsoft YaHei', 'SimHei'] -plt.rcParams['axes.unicode_minus'] = False - -# 从训练输出手动提取的10000步数据 -# 每500步打印一次 -data_raw = """ -ep=0 loss=5.8066 pred=5.2434 router=1.1062 acc=0.34 bal=1.0189 -ep=500 loss=0.4863 pred=0.4721 router=0.0053 acc=1.00 bal=1.1489 -ep=1000 loss=0.2003 pred=0.1876 router=0.0032 acc=1.00 bal=1.1169 -ep=1500 loss=0.1824 pred=0.1681 router=0.0079 acc=1.00 bal=1.0295 -ep=2000 loss=0.1486 pred=0.1371 router=0.0026 acc=1.00 bal=1.0134 -ep=2500 loss=0.1164 pred=0.1051 router=0.0006 acc=1.00 bal=1.0957 -ep=3000 loss=0.1065 pred=0.0929 router=0.0032 acc=1.00 bal=1.1952 -ep=3500 loss=0.1086 pred=0.0985 router=0.0001 acc=1.00 bal=1.0078 -ep=4000 loss=0.1040 pred=0.0936 router=0.0000 acc=1.00 bal=1.0371 -ep=4500 loss=0.1158 pred=0.1050 router=0.0001 acc=1.00 bal=1.0723 -ep=5000 loss=0.1117 pred=0.1009 router=0.0003 acc=1.00 bal=1.0723 -ep=5500 loss=0.0988 pred=0.0866 router=0.0001 acc=1.00 bal=1.2127 -ep=6000 loss=0.0888 pred=0.0780 router=0.0004 acc=1.00 bal=1.0605 -ep=6500 loss=0.1031 pred=0.0926 router=0.0000 acc=1.00 bal=1.0547 -ep=7000 loss=0.0986 pred=0.0884 router=0.0000 acc=1.00 bal=1.0137 -ep=7500 loss=0.0861 pred=0.0750 router=0.0000 acc=1.00 bal=1.1016 -ep=8000 loss=0.0940 pred=0.0837 router=0.0000 acc=1.00 bal=1.0312 -ep=8500 loss=0.0971 pred=0.0838 router=0.0002 acc=1.00 bal=1.3179 -ep=9000 loss=0.0908 pred=0.0807 router=0.0000 acc=1.00 bal=1.0137 -ep=9500 loss=0.0869 pred=0.0755 router=0.0000 acc=1.00 bal=1.1426 -ep=9999 loss=0.0863 pred=0.0759 router=0.0000 acc=1.00 bal=1.0371 -""" - -# 解析数据 -epochs = [] -losses = [] -pred_losses = [] -router_losses = [] -accuracies = [] -balance_losses = [] -qat_phase = [] - -pattern = r'ep=(\d+) loss=([\d.]+) pred=([\d.]+) router=([\d.]+) acc=([\d.]+) bal=([\d.]+)(.*)$' - -for line in data_raw.strip().split('\n'): - m = re.match(pattern, line.strip()) - if m: - ep = int(m.group(1)) - epochs.append(ep) - losses.append(float(m.group(2))) - pred_losses.append(float(m.group(3))) - router_losses.append(float(m.group(4))) - accuracies.append(float(m.group(5))) - balance_losses.append(float(m.group(6))) - qat_phase.append(' [QAT]' in m.group(7)) - -epochs = np.array(epochs) -losses = np.array(losses) -pred_losses = np.array(pred_losses) -router_losses = np.array(router_losses) -accuracies = np.array(accuracies) -balance_losses = np.array(balance_losses) - -# ========== 绘制训练过程图 ========== -fig = plt.figure(figsize=(16, 10)) -fig.suptitle('BitFlow MoE v4 训练过程 — 10,000 轮', fontsize=16, fontweight='bold', y=0.97) - -# 子图1: 总Loss -ax1 = plt.subplot(2, 3, 1) -# 找到 QAT 起始 epoch -qat_epochs = [e for e in epochs if e >= 8000] -qat_epoch_val = qat_epochs[0] if qat_epochs else 8000 -ax1.plot(epochs, losses, 'b-', linewidth=1.5, label='总损失') -ax1.axvspan(qat_epoch_val, epochs[-1], alpha=0.1, color='red', label='QAT量化阶段') -ax1.set_xlabel('训练轮次 (Epoch)') -ax1.set_ylabel('损失值') -ax1.set_title('总损失 (Total Loss)', fontsize=11) -ax1.legend(fontsize=8) -ax1.grid(True, alpha=0.3) -ax1.set_yscale('log') - -# 子图2: 预测Loss -ax2 = plt.subplot(2, 3, 2) -ax2.plot(epochs, pred_losses, 'g-', linewidth=1.5, label='预测损失') -ax2.axvspan(qat_epoch_val, epochs[-1], alpha=0.1, color='red') -ax2.set_xlabel('训练轮次 (Epoch)') -ax2.set_ylabel('损失值') -ax2.set_title('预测损失 (交叉熵)', fontsize=11) -ax2.legend(fontsize=8) -ax2.grid(True, alpha=0.3) -ax2.set_yscale('log') - -# 子图3: 路由Loss -ax3 = plt.subplot(2, 3, 3) -ax3.plot(epochs, router_losses, 'r-', linewidth=1.5, label='路由损失') -ax3.axvspan(qat_epoch_val, epochs[-1], alpha=0.1, color='red') -ax3.set_xlabel('训练轮次 (Epoch)') -ax3.set_ylabel('损失值') -ax3.set_title('路由损失 (专家选择)', fontsize=11) -ax3.legend(fontsize=8) -ax3.grid(True, alpha=0.3) -ax3.set_yscale('log') - -# 子图4: 路由准确率 -ax4 = plt.subplot(2, 3, 4) -ax4.plot(epochs, accuracies, 'purple', linewidth=1.5, marker='o', markevery=500, markersize=4) -ax4.axvspan(qat_epoch_val, epochs[-1], alpha=0.1, color='red') -ax4.set_xlabel('训练轮次 (Epoch)') -ax4.set_ylabel('准确率') -ax4.set_title('路由准确率 (专家路由)', fontsize=11) -ax4.legend(fontsize=8) -ax4.grid(True, alpha=0.3) -ax4.set_ylim([0, 1.05]) - -# 子图5: 负载均衡Loss -ax5 = plt.subplot(2, 3, 5) -ax5.plot(epochs, balance_losses, 'orange', linewidth=1.5, label='负载均衡损失') -ax5.axvspan(qat_epoch_val, epochs[-1], alpha=0.1, color='red') -ax5.set_xlabel('训练轮次 (Epoch)') -ax5.set_ylabel('损失值') -ax5.set_title('负载均衡 (专家利用率)', fontsize=11) -ax5.legend(fontsize=8) -ax5.grid(True, alpha=0.3) -ax5.set_yscale('log') - -# 子图6: 各阶段对比 (分三段) -ax6 = plt.subplot(2, 3, 6) -segments = [ - ('早期 (0-2K)', epochs[epochs<=2000], losses[epochs<=2000], 'steelblue'), - ('中期 (2K-6K)', epochs[(epochs>2000)&(epochs<=6000)], losses[(epochs>2000)&(epochs<=6000)], 'seagreen'), - ('QAT量化 (8K-10K)', epochs[epochs>=8000], losses[epochs>=8000], 'firebrick'), -] -for label, ep, lo, color in segments: - ax6.plot(ep, lo, 'o-', color=color, linewidth=1.5, markersize=4, label=label, alpha=0.7) -ax6.set_xlabel('训练轮次 (Epoch)') -ax6.set_ylabel('损失值') -ax6.set_title('分阶段对比 (早期/中期/QAT)', fontsize=11) -ax6.legend(fontsize=8) -ax6.grid(True, alpha=0.3) -ax6.set_yscale('log') - -plt.tight_layout(rect=[0, 0, 1, 0.95]) -output_path = 'c:\\Users\\Administrator\\.trae-cn\\work\\6a62468bb999696f3874e2ac\\training_analysis.png' -plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white') -plt.close() -print(f"[save] training_analysis.png -> {output_path}") - -# 输出关键统计 -print(f"\n关键数据点:") -print(f" 初始 loss: {losses[0]:.4f} router_acc: {accuracies[0]:.0%}") -print(f" 500 epoch: {losses[1]:.4f} router_acc: {accuracies[1]:.0%}") -print(f" 2000 epoch: {losses[np.where(epochs==2000)[0][0]]:.4f} router_acc: {accuracies[np.where(epochs==2000)[0][0]]:.0%}") -print(f" 6000 epoch: {losses[np.where(epochs==6000)[0][0]]:.4f} router_acc: {accuracies[np.where(epochs==6000)[0][0]]:.0%}") -print(f" 9999 epoch: {losses[-1]:.4f} router_acc: {accuracies[-1]:.0%}") -print(f" 最佳 loss: {losses.min():.4f} at epoch {epochs[np.argmin(losses)]}") - -# QAT 阶段 -if any(qat_phase): - qat_idx = np.where(np.array(qat_phase))[0][0] - print(f"\nQAT 量化阶段:") - print(f" QAT前(7500): {losses[np.where(epochs==7500)[0][0]]:.4f}") - print(f" QAT后(8000): {losses[np.where(epochs==8000)[0][0]]:.4f}") - print(f" QAT最终(9999): {losses[-1]:.4f}") - print(f" QAT导致loss上升: +{losses[qat_idx] - losses[qat_idx-1]:.4f} (正常，量化引入噪声)") -else: - print(f"\nQAT 量化阶段: 未检测到QAT数据") diff --git a/quickstart.md b/quickstart.md deleted file mode 100644 index 154dd05..0000000 --- a/quickstart.md +++ /dev/null @@ -1,209 +0,0 @@ -# BitFlow MoE 用户快速上手指南 - -## 一、什么是 BitFlow MoE？ - -BitFlow MoE 是一个可以在 **BBC micro:bit** 上运行的微型中文语言模型。它可以： - -- 回答基础问答（你好、谢谢、你的名字） -- 进行简单推理（什么会飞、什么吃草） -- 记住事实（猫吃什么、天是什么颜色） - -## 二、硬件准备 - -### 必需 - -- **BBC micro:bit v2**（推荐 v2，nRF52833 芯片） -- **USB 数据线**（Micro-USB） -- **电脑**（Windows / macOS / Linux） - -### 可选 - -- 串口工具（PuTTY、Tera Term、minicom） -- Python 3.8+（用于训练和对话工具） - -## 三、快速部署（3 分钟） - -### 步骤 1: 下载模型文件 - -从本仓库下载： -- `bitflow_moe_weights.bin`（58KB，int8 量化权重） -- `bitflow_moe.py`（micro:bit 推理代码） - -### 步骤 2: 拷贝到 micro:bit - -将两个文件拷贝到 micro:bit 的 USB 驱动器中： - -``` -MICROBIT/ -├── weights.bin ← bitflow_moe_weights.bin 改名 -├── main.py ← bitflow_moe.py 改名 -└── ... -``` - -> **注意**：文件名必须是 `weights.bin` 和 `main.py`，不能改名！ - -### 步骤 3: 按下复位键 - -按下 micro:bit 背面的 **RST** 按钮（或拔插 USB），micro:bit 会自动运行 main.py。 - -### 步骤 4: 通过串口对话 - -打开串口工具（115200 波特率），发送： - -``` -GEN:109,110|15 -``` - -回复：`你好` - -更多命令见下方 [常用短语 ID 表](#四常用短语id)。 - -## 四、常用短语 ID - -### 问候类 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 你好 | `GEN:109,110|15` | 你好 | -| 你好吗 | `GEN:109,110,111|15` | 我很好谢谢 | -| 早安 | `GEN:116,118|15` | 早安 | -| 晚安 | `GEN:117,118|15` | 晚安 | -| 再见 | `GEN:114,115|15` | 再见 | - -### 礼貌类 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 谢谢 | `GEN:113,113|15` | 不客气 | -| 你好棒 | `GEN:110,110,120|15` | 谢谢你好 | - -### 自我介绍 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 你叫什么名字 | `GEN:109,131,39,40,129,130|15` | 我叫bit | -| 你是谁 | `GEN:109,45|15` | 我是bit | - -### 喜好 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 你喜欢什么 | `GEN:109,123,124,39,40|15` | 我喜欢花 | -| 你喜欢猫吗 | `GEN:109,123,124,43|15` | 我很喜欢猫 | - -### 情感 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 你开心吗 | `GEN:109,70,71|15` | 我很开心你好 | -| 你伤心吗 | `GEN:109,126,127|15` | 我不伤心 | - -### 能力 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 你会什么 | `GEN:109,78,59|15` | 我会说话 | - -### 闲聊 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 你在做什么 | `GEN:109,60,61|15` | 我在看猫 | -| 今天好吗 | `GEN:140,60,36,37|15` | 今天很好 | - -### 动物推理 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 什么会飞 | `GEN:39,40,78,15|15` | 鸟会飞 | -| 什么会跑 | `GEN:39,40,78,14|15` | 猫会跑 | -| 什么会游 | `GEN:39,40,78,16|15` | 鱼会游 | -| 什么吃草 | `GEN:39,40,78,18|15` | 牛吃草 | - -### 事实 - -| 中文 | 命令 | 回复 | -|------|------|------| -| 天是什么颜色 | `GEN:54,35,38,39,40|15` | 天是蓝的 | -| 猫吃什么 | `GEN:43,35,18|15` | 猫吃鱼 | -| 猫是什么 | `GEN:43,35,38|15` | 猫是动物 | - -## 五、使用 Python 对话工具 - -如果你不想手动输入 ID，可以使用对话工具： - -```bash -# 安装依赖 -pip install pyserial - -# 运行对话工具（自动查找 micro:bit 串口） -python chat_moe.py -``` - -对话工具会自动： -1. 查找 micro:bit 的 USB 串口 -2. 显示常用短语菜单 -3. 你输入序号或文字 -4. 自动发送并翻译回复 - -## 六、在电脑上训练自己的模型 - -```bash -# 安装依赖 -pip install torch matplotlib - -# 训练 -python train_bitflow_moe.py - -# 训练完成后会自动： -# 1. 保存 PyTorch 权重 (bitflow_moe.pt) -# 2. 导出 int8 量化权重 (bitflow_moe_weights.bin) -# 3. 测试推理效果 -# 4. 绘制训练过程图 (training_analysis.png) -``` - -## 七、常见问题 - -### Q1: 为什么回复是一串数字？ - -BitFlow MoE 输出的是**字符 ID**（0-148），不是中文。这是为了在 micro:bit 上节省内存——字符 ID 只需要 1 字节，而中文字符需要多字节。 - -你可以通过 `chat_moe.py` 自动翻译，或参考上方 ID 表手动解码。 - -### Q2: 模型回答不总是对的？ - -是的。这个模型是**微型模型**，学习的数据量很小（290K tokens），词表只有 149 字符。它更像是一个"概念演示"，证明在 micro:bit 上可以做语言建模。 - -### Q3: 可以扩展词表吗？ - -可以。修改 `VOCAB_CHARS` 字符串即可扩展词表，但需要： -1. 增加 DIM 和 FFN_DIM（参数量会增大） -2. 可能需要更大的 micro:bit（如 ESP32） -3. 重新训练 - -### Q4: 可以用在其他单片机上吗？ - -可以！本项目的推理代码是纯 Python/MicroPython，理论上可以移植到任何支持 MicroPython 的 MCU。我们提供了内存分析： - -| 平台 | RAM | 是否可用 | -|------|-----|----------| -| micro:bit v2 | 256KB (可用 ~108KB) | ✅ 刚好够用 | -| STM32F4 | 192KB | ✅ 完全够用 | -| ESP32 | 520KB | ✅ 绰绰有余 | -| Arduino Nano | 2KB | ❌ 远远不够 | - -### Q5: 推理速度很慢？ - -micro:bit 的 CPU 是 64MHz，每个字符生成需要 200-500ms。这是硬件限制，无法通过软件优化大幅改善。 - -如果你需要更快的推理，可以考虑： -- 降低 DIM（32 → 24） -- 减少专家数量（3 → 1） -- 换用更强的 MCU（如 ESP32-S3） - -## 八、技术参考 - -- [架构设计文档](architecture.md) — 详细的架构分析 -- [训练过程图](training_analysis.png) — 10,000 轮训练可视化 -- [训练代码](train_bitflow_moe.py) — 完整的训练脚本 -- [推理代码](bitflow_moe.py) — micro:bit 推理实现 diff --git a/train_bitflow_moe.py b/train_bitflow_moe.py deleted file mode 100644 index 1f28b51..0000000 --- a/train_bitflow_moe.py +++ /dev/null @@ -1,775 +0,0 @@ -""" -BitFlow MoE v4: 面向 micro:bit 的混合专家微型中文对话模型 -架构: 门控线性注意力 (GLA) + 混合专家 (MoE) + int8 量化 - -v4 改进: -- 3个专家: 事实记忆(E0) + 逻辑推理(E1) + 基础问答(E2) -- 增大参数量: DIM=40, FFN=80 (~54K params, 更聪明) -- 扩展词表: +你好吗谢谢再见早晚安等基础对话字符 -- 基础问答专家: 200+条自然问答 (问候/礼貌/自我介绍/喜好/情感/能力/闲聊/天气/动物/颜色/数字) -""" -import math -import random -import torch -import torch.nn as nn -import torch.nn.functional as F - -# ---------- 配置 ---------- -VOCAB_CHARS = ( - "猫狗鸟鱼猪牛马鸡虫羊兔鸭鹅" # 动物 - "坐跑飞游看吃睡走听来叫爬跳唱画读做" # 动作 - "在上里到" # 位置/方向 - "的是好了" # 助词/语气 - "吗什么怎哪" # 疑问 - "你我他她们它" # 代词 - "蓝绿红黄白" # 颜色 - "天地水火花草树木日月星云春夏秋冬" # 自然 - "大小多高远快" # 形容词 - "有没会不" # 否定/存在/能力 - "人头脚手眼腿" # 身体 - "垫肉书子" # 物品 - "一三只两" # 数字/量词 - "热冷落开" # 状态 - "自然动植物质色样颜得都" # 补充 - "你好吗谢谢再见早晚安静" # 基础问答 (新增) - "很棒爱喜欢乐伤心气" # 情感 (新增) - "名字叫谁和问答应" # 对话用语 (新增) - "聊天话今天说" # 补充对话字符 (修复) - "客用" # 不客气/不用谢 (修复) - "bit" # 名字 (修复) - "\n" # 换行 -) -VOCAB_SIZE = len(VOCAB_CHARS) -DIM = 40 # 增大: 32 -> 40 -N_LAYER = 2 -CONTEXT = 48 -FFN_DIM = 80 # 增大: 64 -> 80 -NUM_EXPERTS = 3 # 3个专家 -ROUTER_DIM = 32 # 增大: 24 -> 32 -EPOCHS = 10000 # 增多: 8000 -> 10000 -BATCH_SIZE = 32 -LR = 3e-3 -SEED = 1337 - -QAT_START = 8000 -ROUTER_WEIGHT = 0.5 -LB_WEIGHT = 0.01 - -random.seed(SEED) -torch.manual_seed(SEED) - -CH2I = {c: i for i, c in enumerate(VOCAB_CHARS)} -I2CH = {i: c for i, c in enumerate(VOCAB_CHARS)} - -# ---------- 训练数据 ---------- - -BASE = [ - "猫坐在垫子上", "狗坐在木头上", "鸟坐在树上", "鱼在水里游", - "羊在吃草", "兔在跑", "鸭在水里游", "鹅在水里游", - "猫在跑", "狗在跑", "鸟在飞", "鱼在游", - "马在跑", "牛在走", "猪在睡", "鸡在走", - "虫在树上", "羊在走", "兔在跳", "鸭在走", "鹅在叫", - "天是蓝的", "水是绿的", "火是红的", "花是黄的", - "草是绿的", "日是红的", "月是黄的", "星是好的", "云是白的", - "春天是好的", "夏天是热的", "秋天是好的", "冬天是冷的", - "猫看狗", "狗看猫", "鸟看鱼", "猫吃鱼", "牛吃草", - "鸟吃虫", "羊吃草", "兔吃草", "鸭吃虫", "鹅吃草", - "猫在吃鱼", "狗在跑走", "鸟在唱", "虫在爬", "兔在跳", - "鸭在叫", "鹅在游", - "马是大的", "虫是小的", "马跑得快", "鸟飞得高", - "鱼游得快", "猫跳得高", "狗跑得远", - "猫有腿", "鸟有脚", "鱼没有腿", "人有手", "人有眼", - "人有脚", "猫有眼", - "春天花开了", "夏天日大了", "秋天花落了", "冬天水冷了", -] - -# Expert-0: 事实与记忆类问答 -FACT_QA = [ - "猫是什么\n猫是动物", "狗是什么\n狗是动物", "鸟是什么\n鸟是动物", - "鱼是什么\n鱼是动物", "马是什么\n马是动物", "牛是什么\n牛是动物", - "猪是什么\n猪是动物", "鸡是什么\n鸡是动物", "虫是什么\n虫是动物", - "羊是什么\n羊是动物", "兔是什么\n兔是动物", "鸭是什么\n鸭是动物", - "鹅是什么\n鹅是动物", - "花是什么\n花是植物", "草是什么\n草是植物", "树是什么\n树是植物", - "水是什么\n水是物质", "火是什么\n火是物质", "天是什么\n天是自然", - "地是什么\n地是自然", "日是什么\n日是星", "月是什么\n月是星", - "云是什么\n云是水", "人是什么\n人是动物", - "天是什么颜色\n天是蓝的", "水是什么颜色\n水是绿的", - "火是什么颜色\n火是红的", "草是什么颜色\n草是绿的", - "花是什么颜色\n花是黄的", "日是什么颜色\n日是红的", - "月是什么颜色\n月是黄的", "云是什么颜色\n云是白的", - "猫是动物吗\n猫是动物", "狗是动物吗\n狗是动物", - "鱼是动物吗\n鱼是动物", "花是植物吗\n花是植物", - "树是植物吗\n树是植物", - "天是蓝的吗\n天是蓝的", "水是绿的吗\n水是绿的", - "火是红的吗\n火是红的", "云是白的吗\n云是白的", - "猫在哪里\n猫在垫子上", "狗在哪里\n狗在木头上", - "鸟在哪里\n鸟在树上", "鱼在哪里\n鱼在水里", - "羊在哪里\n羊在草地上", "兔在哪里\n兔在草地上", - "虫在哪里\n虫在树上", - "猫吃什么\n猫吃鱼", "狗吃什么\n狗吃肉", "鸟吃什么\n鸟吃虫", - "牛吃什么\n牛吃草", "鱼吃什么\n鱼吃虫", "羊吃什么\n羊吃草", - "兔吃什么\n兔吃草", "鸭吃什么\n鸭吃虫", "鹅吃什么\n鹅吃草", - "春天是什么\n春天是好的", "夏天是什么\n夏天是热的", - "秋天是什么\n秋天是好的", "冬天是什么\n冬天是冷的", - "春天来了\n花开了", "夏天来了\n日大了", - "秋天来了\n花落了", "冬天来了\n水冷了", - "你在看什么\n我在看猫", "你在看什么\n我在看狗", - "你在看什么\n我在看鸟", "你在看什么\n我在看鱼", - "你在做什么\n我在看花", "你在做什么\n我在画猫", - "你在做什么\n我在读书", "他在做什么\n他在听鸟叫", - "她是什么\n她是花", "他是什么\n他是猫", "他是什么\n他是狗", -] - -# Expert-1: 逻辑与推理类问答 -REASONING_QA = [ - "猫是动物\n动物会跑\n猫会跑", - "狗是动物\n动物会跑\n狗会跑", - "鸟是动物\n动物会飞\n鸟会飞", - "鱼是动物\n动物会游\n鱼会游", - "马是动物\n动物会跑\n马会跑", - "牛是动物\n动物会走\n牛会走", - "猪是动物\n动物会睡\n猪会睡", - "羊是动物\n动物会走\n羊会走", - "兔是动物\n动物会跳\n兔会跳", - "鸭是动物\n动物会游\n鸭会游", - "鹅是动物\n动物会叫\n鹅会叫", - "动物有腿\n猫是动物\n猫有腿", - "动物有腿\n狗是动物\n狗有腿", - "动物有腿\n马是动物\n马有腿", - "动物有腿\n牛是动物\n牛有腿", - "动物有眼\n猫有眼", - "鱼是动物\n鱼没有腿", - "什么会跑\n猫会跑", - "什么会跑\n狗会跑", - "什么会飞\n鸟会飞", - "什么会游\n鱼会游", - "什么会游\n鸭会游", - "什么会跳\n兔会跳", - "什么会叫\n鹅会叫", - "什么会睡\n猪会睡", - "什么会走\n牛会走", - "什么会走\n羊会走", - "猫有腿\n狗有腿\n它们都有腿", - "猫会跑\n狗会跑\n它们都会跑", - "鸟有脚\n鸟会飞", - "鱼没有腿\n鱼会游", - "猫吃什么\n猫吃鱼", - "什么吃鱼\n猫吃鱼", - "什么吃草\n牛吃草", - "什么吃草\n羊吃草", - "什么吃虫\n鸟吃虫", - "什么吃虫\n鸭吃虫", - "猫是小的\n虫是小的", - "马是大的\n马跑得快", - "鸟飞得高\n鱼游得快", - "猫跳得高\n狗跑得远", -] - -# Expert-2: 基础问答/闲聊 (大幅扩展 v4: 200+条) -CHAT_QA = [ - # === 问候类 (30条) === - "你好\n你好", - "你好\n你好你好", - "你好bit\n你好", - "你好猫\n你好猫", - "你好狗\n你好狗", - "你好花\n你好花", - "你好天\n你好天", - "你好水\n你好水", - "你好火\n你好火", - "你好你好\n你好", - "bit你好\n你好bit", - "猫你好\n你好", - "你好吗\n我很好", - "你好吗\n我很好谢谢", - "你好吗\n我很好你好", - "你好吗bit\n我很好", - "bit你好吗\n我很好", - "你好了吗\n我好了", - "你好吗\n我很好和开心", - "早\n早", - "早\n早早", - "早安\n早安", - "早安bit\n早安", - "bit早安\n早安", - "晚安\n晚安", - "晚安bit\n晚安", - "再见\n再见", - "再见\n再见再见", - "再见bit\n再见", - "bit再见\n再见", - - # === 礼貌类 (18条) === - "谢谢\n不客气", - "谢谢\n不用谢", - "谢谢\n不客气你好", - "谢谢bit\n不客气", - "谢谢猫\n不客气", - "谢谢你好\n不客气", - "不客气\n谢谢", - "不用谢\n谢谢", - "你好棒\n谢谢", - "你好棒\n谢谢你好", - "你好棒\n谢谢你", - "你好棒bit\n谢谢", - "bit好棒\n谢谢", - "猫好棒\n猫很好", - "狗好棒\n狗很好", - "谢谢你\n不客气", - "你好棒\n谢谢你好bit", - "谢谢和再见\n不客气", - - # === 自我介绍 (20条) === - "你叫什么名字\n我叫bit", - "你叫什么\n我叫bit", - "你的名字\n我叫bit", - "你的名字是什么\n我叫bit", - "你是谁\n我是bit", - "你是谁\n我是猫", - "你是谁\n我是动物", - "你是bit吗\n我是bit", - "你是猫吗\n我不是猫", - "你是狗吗\n我不是狗", - "你是动物吗\n我不是动物", - "你是人吗\n我不是人", - "你是花吗\n我不是花", - "bit是谁\nbit是我", - "猫是谁\n猫是动物", - "狗是谁\n狗是动物", - "你叫什么\n我叫bit你好", - "我叫猫\n你好猫", - "我叫狗\n你好狗", - "我叫bit\n你好bit", - - # === 喜好类 (38条) === - "你喜欢什么\n我喜欢猫", - "你喜欢什么\n我喜欢狗", - "你喜欢什么\n我喜欢鸟", - "你喜欢什么\n我喜欢鱼", - "你喜欢什么\n我喜欢花", - "你喜欢什么\n我喜欢树", - "你喜欢什么\n我喜欢草", - "你喜欢什么\n我喜欢天", - "你喜欢什么\n我喜欢水", - "你喜欢什么\n我喜欢火", - "你喜欢什么\n我喜欢日", - "你喜欢什么\n我喜欢月", - "你喜欢什么\n我喜欢星", - "你喜欢什么\n我喜欢云", - "你喜欢什么\n我喜欢春天", - "你喜欢什么\n我喜欢夏天", - "你喜欢什么\n我喜欢秋天", - "你喜欢什么\n我喜欢冬天", - "你喜欢什么\n我喜欢读书", - "你喜欢什么\n我喜欢画画", - "你喜欢什么\n我喜欢聊天", - "你喜欢猫吗\n我很喜欢猫", - "你喜欢狗吗\n我很喜欢狗", - "你喜欢鸟吗\n我很喜欢鸟", - "你喜欢鱼吗\n我很喜欢鱼", - "你喜欢花吗\n我很喜欢花", - "你喜欢树吗\n我很喜欢树", - "你喜欢草吗\n我很喜欢草", - "你喜欢水吗\n我很喜欢水", - "你喜欢火吗\n我很喜欢火", - "你喜欢天吗\n我很喜欢天", - "你喜欢春天吗\n我很喜欢春天", - "你喜欢夏天吗\n我很喜欢夏天", - "你喜欢秋天吗\n我很喜欢秋天", - "你喜欢冬天吗\n我很喜欢冬天", - "你喜欢读书吗\n我很喜欢读书", - "你喜欢画画吗\n我很喜欢画画", - "你喜欢聊天吗\n我很喜欢聊天", - - # === 情感类 (22条) === - "你开心吗\n我很开心", - "你开心吗\n我很开心谢谢", - "你伤心吗\n我不伤心", - "你伤心吗\n我不伤心谢谢", - "你乐吗\n我很乐", - "你气吗\n我不气", - "你好开心\n我很开心", - "你好伤心\n我不伤心", - "我很开心\n你好", - "我很伤心\n你好", - "我很喜欢\n你好", - "我很爱你\n谢谢", - "你爱我吗\n我很爱你", - "你爱猫吗\n我很爱猫", - "你爱花吗\n我很爱花", - "猫开心吗\n猫很开心", - "狗开心吗\n狗很开心", - "鸟开心吗\n鸟很开心", - "鱼开心吗\n鱼很开心", - "bit开心吗\nbit很开心", - "你开心吗\n我很开心你好", - "你好开心吗\n我很开心", - - # === 能力问答 (23条) === - "你会什么\n我会说话", - "你会什么\n我会画画", - "你会什么\n我会读书", - "你会什么\n我会聊天", - "你会什么\n我会看花", - "你会什么\n我会看猫", - "你会什么\n我会看狗", - "你会说话吗\n我会说话", - "你会画画吗\n我会画画", - "你会读书吗\n我会读书", - "你会聊天吗\n我会聊天", - "你会跑吗\n我不会跑", - "你会飞吗\n我不会飞", - "你会游吗\n我不会游", - "你会跳吗\n我不会跳", - "你会唱吗\n我不会唱", - "你会画猫吗\n我会画猫", - "你会画狗吗\n我会画狗", - "你会画花吗\n我会画花", - "bit会什么\nbit会说话", - "bit会说话吗\nbit会说话", - "bit会聊天吗\nbit会聊天", - "你会什么\n我会说话和聊天", - - # === 在做什么 (25条) === - "你在做什么\n我在和你说话", - "你在做什么\n我在和你聊天", - "你在做什么\n我在看猫", - "你在做什么\n我在看狗", - "你在做什么\n我在看花", - "你在做什么\n我在读书", - "你在做什么\n我在画画", - "你在做什么\n我在画猫", - "你在做什么\n我在画花", - "你在做什么\n我在听鸟叫", - "你在做什么\n我在看鸟", - "你在做什么\n我在看鱼", - "你在做什么\n我在看天", - "你在做什么\n我在看水", - "谁在说话\n我在说话", - "谁在说话\n你在说话", - "谁在说话\nbit在说话", - "谁在聊天\n我在聊天", - "谁在聊天\n你在聊天", - "谁在聊天\nbit在聊天", - "你在聊天吗\n我在聊天", - "你在说话吗\n我在说话", - "bit在做什么\nbit在说话", - "bit在说话吗\nbit在说话", - "bit在聊天吗\nbit在聊天", - - # === 天气时间 (22条) === - "今天好吗\n今天很好", - "天气好吗\n天气很好", - "天好吗\n天很好", - "花好吗\n花很好", - "今天天好吗\n今天天很好", - "今天热吗\n今天很热", - "今天冷吗\n今天不冷", - "天热吗\n天很热", - "天冷吗\n天不冷", - "夏天热吗\n夏天很热", - "冬天冷吗\n冬天很冷", - "春天好吗\n春天很好", - "夏天好吗\n夏天很好", - "秋天好吗\n秋天很好", - "冬天好吗\n冬天很好", - "今天天热吗\n今天天很热", - "今天天冷吗\n今天天不冷", - "早好吗\n早很好", - "晚好吗\n晚很好", - "早安好吗\n早安很好", - "晚安好吗\n晚安很好", - "今天天气好吗\n今天天气很好", - - # === 动物闲聊 (19条) === - "猫好吗\n猫很好", - "狗好吗\n狗很好", - "鱼好吗\n鱼很好", - "鸟好吗\n鸟很好", - "虫好吗\n虫很好", - "马好吗\n马很好", - "牛好吗\n牛很好", - "猪好吗\n猪很好", - "羊好吗\n羊很好", - "兔好吗\n兔很好", - "鸭好吗\n鸭很好", - "鹅好吗\n鹅很好", - "猫好吗\n猫是好的", - "狗好吗\n狗是好的", - "花好吗\n花是好的", - "天好吗\n天是好的", - "水好吗\n水是好的", - "猫好吗\n猫很好谢谢", - "狗好吗\n狗很好谢谢", - - # === 颜色喜好 (18条) === - "你喜欢什么颜色\n我喜欢蓝的", - "你喜欢什么颜色\n我喜欢红的", - "你喜欢什么颜色\n我喜欢绿的", - "你喜欢什么颜色\n我喜欢黄的", - "你喜欢什么颜色\n我喜欢白的", - "你喜欢蓝的吗\n我很喜欢蓝的", - "你喜欢红的吗\n我很喜欢红的", - "你喜欢绿的吗\n我很喜欢绿的", - "你喜欢黄的吗\n我很喜欢黄的", - "你喜欢白的吗\n我很喜欢白的", - "天是蓝的吗\n天是蓝的", - "水是绿的吗\n水是绿的", - "火是红的吗\n火是红的", - "云是白的吗\n云是白的", - "花是黄的吗\n花是黄的", - "草是绿的吗\n草是绿的", - "日是红的吗\n日是红的", - "月是黄的吗\n月是黄的", - - # === 数字/量词 (11条) === - "一只猫\n一只猫", - "一只狗\n一只狗", - "一只鸟\n一只鸟", - "一只鱼\n一只鱼", - "两只猫\n两只猫", - "两只狗\n两只狗", - "三只鸟\n三只鸟", - "三只猫\n三只猫", - "一只猫好吗\n猫很好", - "一只狗好吗\n狗很好", - "三只鸟好吗\n鸟很好", - - # === 组合对话 (12条) === - "你好和再见\n你好再见", - "你好猫和狗\n你好猫你好狗", - "你喜欢猫和狗吗\n我很喜欢猫和狗", - "你喜欢花和树吗\n我很喜欢花和树", - "你喜欢天和水吗\n我很喜欢天和水", - "猫和狗好吗\n猫和狗都很好", - "花和草好吗\n花和草都很好", - "天和水好吗\n天和水都很好", - "春天和秋天好吗\n春天和秋天都很好", - "你和bit好吗\n我和bit都很好", - "你好和谢谢\n你好谢谢", - "你好再见谢谢\n你好再见不客气", -] - -# 分别构建数据 -FACT_TEXT = ("\n".join(BASE) + "\n") * 20 + "\n".join(FACT_QA) * 40 -REASON_TEXT = "\n".join(REASONING_QA) * 80 -CHAT_TEXT = "\n".join(CHAT_QA) * 80 - -FACT_DATA = torch.tensor([CH2I[c] for c in FACT_TEXT if c in CH2I], dtype=torch.long) -REASON_DATA = torch.tensor([CH2I[c] for c in REASON_TEXT if c in CH2I], dtype=torch.long) -CHAT_DATA = torch.tensor([CH2I[c] for c in CHAT_TEXT if c in CH2I], dtype=torch.long) -DATA = torch.cat([FACT_DATA, REASON_DATA, CHAT_DATA]) - -print(f"[data] fact={FACT_DATA.numel()} reason={REASON_DATA.numel()} chat={CHAT_DATA.numel()} total={DATA.numel()}") -print(f"[data] base={len(BASE)} fact_qa={len(FACT_QA)} reasoning_qa={len(REASONING_QA)} chat_qa={len(CHAT_QA)}") - -all_chars = set("".join(BASE) + "\n".join(FACT_QA) + "\n".join(REASONING_QA) + "\n".join(CHAT_QA)) -missing = [c for c in VOCAB_CHARS if c not in all_chars and c != '\n'] -unused = [c for c in all_chars if c not in CH2I] -if missing: - print(f"[warn] vocab chars not in training data: {missing}") -if unused: - print(f"[warn] training chars not in vocab (will be filtered): {unused}") -print(f"[vocab] size={VOCAB_SIZE}") - - -# ---------- GLA ---------- -class GLA(nn.Module): - def __init__(self, dim): - super().__init__() - self.q = nn.Linear(dim, dim, bias=False) - self.k = nn.Linear(dim, dim, bias=False) - self.v = nn.Linear(dim, dim, bias=False) - self.g = nn.Linear(dim, dim, bias=False) - self.o = nn.Linear(dim, dim, bias=False) - self.log_decay = nn.Parameter(torch.zeros(dim)) - - def forward(self, x): - B, T, D = x.shape - q = F.elu(self.q(x)) + 1 - k = F.elu(self.k(x)) + 1 - v = self.v(x) - gate = torch.sigmoid(self.g(x)) - decay = torch.sigmoid(self.log_decay) - - S = torch.zeros(B, D, D, device=x.device) - z = torch.zeros(B, D, device=x.device) - outputs = [] - for t in range(T): - g_t = gate[:, t] * decay - S = g_t.unsqueeze(-1) * S + k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(1) - z = g_t * z + k[:, t] - num = torch.bmm(q[:, t].unsqueeze(1), S).squeeze(1) - den = (q[:, t] * z).sum(-1, keepdim=True) + 1e-6 - y = num / den - outputs.append(y) - y = torch.stack(outputs, 1) - return self.o(y) - - -class BitFlowBlock(nn.Module): - def __init__(self, dim): - super().__init__() - self.norm1 = nn.LayerNorm(dim) - self.attn = GLA(dim) - self.norm2 = nn.LayerNorm(dim) - self.ff1 = nn.Linear(dim, FFN_DIM, bias=False) - self.ff2 = nn.Linear(FFN_DIM, dim, bias=False) - - def forward(self, x): - x = x + self.attn(self.norm1(x)) - h = self.norm2(x) - h = self.ff2(F.gelu(self.ff1(h), approximate='tanh')) - x = x + h - return x - - -# ---------- MoE ---------- -class Router(nn.Module): - def __init__(self, dim, num_experts, router_dim=ROUTER_DIM): - super().__init__() - self.fc1 = nn.Linear(dim, router_dim) - self.fc2 = nn.Linear(router_dim, num_experts) - - def forward(self, x): - x = x.mean(dim=1) - x = F.relu(self.fc1(x)) - return self.fc2(x) - - -class BitFlowMoE(nn.Module): - def __init__(self, num_experts=NUM_EXPERTS): - super().__init__() - self.tok = nn.Embedding(VOCAB_SIZE, DIM) - self.pos = nn.Embedding(CONTEXT, DIM) - self.experts = nn.ModuleList([ - BitFlowBlock(DIM) for _ in range(num_experts) - ]) - self.router = Router(DIM, num_experts) - self.norm_f = nn.LayerNorm(DIM) - self.head = nn.Linear(DIM, VOCAB_SIZE, bias=False) - self.num_experts = num_experts - - n = sum(p.numel() for p in self.parameters()) - print(f"[model] BitFlow MoE v3 params={n} ({n/1024:.1f}K) experts={num_experts} D={DIM} FFN={FFN_DIM}") - - def forward(self, idx, labels=None): - B, T = idx.shape - pos = torch.arange(T) - x = self.tok(idx) + self.pos(pos)[None] - gate_logits = self.router(x) - - if labels is not None: - outputs = torch.zeros_like(x) - for i in range(self.num_experts): - mask = (labels == i) - if mask.any(): - idx_mask = torch.where(mask)[0] - for _ in range(N_LAYER): - x[idx_mask] = self.experts[i](x[idx_mask]) - outputs[idx_mask] = x[idx_mask] - x = outputs - else: - expert_idx = gate_logits.argmax(dim=-1) - outputs = torch.zeros_like(x) - for i in range(self.num_experts): - mask = (expert_idx == i) - if mask.any(): - idx_mask = torch.where(mask)[0] - for _ in range(N_LAYER): - x[idx_mask] = self.experts[i](x[idx_mask]) - outputs[idx_mask] = x[idx_mask] - x = outputs - - x = self.norm_f(x) - return self.head(x), gate_logits - - -# ---------- 数据采样 ---------- -DATASETS = [FACT_DATA, REASON_DATA, CHAT_DATA] - -def get_batch(bs=BATCH_SIZE): - xs, ys, labels = [], [], [] - for _ in range(bs): - # 均匀采样 3 个专家 - r = random.random() - if r < 0.40: - data, label = FACT_DATA, 0 - elif r < 0.70: - data, label = REASON_DATA, 1 - else: - data, label = CHAT_DATA, 2 - i = random.randint(0, data.numel() - CONTEXT - 2) - xs.append(data[i:i+CONTEXT]) - ys.append(data[i+1:i+CONTEXT+1]) - labels.append(label) - return torch.stack(xs), torch.stack(ys), torch.tensor(labels, dtype=torch.long) - - -def apply_qat(model): - with torch.no_grad(): - for name, p in model.named_parameters(): - if p.requires_grad: - scale = p.abs().max() / 127.0 + 1e-8 - q = torch.clamp((p / scale).round(), -128, 127) * scale - p.copy_(q) - - -# ---------- 训练 ---------- -def main(): - model = BitFlowMoE(num_experts=NUM_EXPERTS) - opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01) - sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, EPOCHS) - model.train() - - best_loss = float('inf') - - for ep in range(EPOCHS): - if ep >= QAT_START: - apply_qat(model) - - x, y, labels = get_batch() - logits, gate_logits = model(x, labels=labels) - - pred_loss = F.cross_entropy(logits.reshape(-1, VOCAB_SIZE), y.reshape(-1)) - router_loss = F.cross_entropy(gate_logits, labels) - gate_probs = F.softmax(gate_logits, dim=-1) - avg_usage = gate_probs.mean(dim=0) - balance_loss = NUM_EXPERTS * (avg_usage * avg_usage).sum() - - loss = pred_loss + ROUTER_WEIGHT * router_loss + LB_WEIGHT * balance_loss - - if torch.isnan(loss) or torch.isinf(loss): - print(f"[warn] NaN/Inf at ep={ep}, skipping") - opt.zero_grad() - continue - - opt.zero_grad() - loss.backward() - torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) - opt.step() - sched.step() - - if ep % 500 == 0 or ep == EPOCHS - 1: - qat_tag = " [QAT]" if ep >= QAT_START else "" - with torch.no_grad(): - pred_expert = gate_logits.argmax(dim=-1) - router_acc = (pred_expert == labels).float().mean().item() - print(f"[train] ep={ep:4d} loss={loss.item():.4f} " - f"pred={pred_loss.item():.4f} router={router_loss.item():.4f} " - f"acc={router_acc:.2f} bal={balance_loss.item():.4f}{qat_tag}") - if loss.item() < best_loss: - best_loss = loss.item() - - model.eval() - print(f"\n[best] loss={best_loss:.4f}") - - test_model(model) - torch.save(model.state_dict(), "bitflow_moe.pt") - print("\n[save] bitflow_moe.pt") - - try: - export_int8(model) - except Exception: - import traceback - print("[ERROR] export failed:") - traceback.print_exc() - - -def test_model(model): - import traceback - try: - seeds = [ - # 基础问答 (Expert-2) - "你好", "你好吗", "谢谢", "再见", "你叫什么名字", - "你喜欢什么", "你喜欢猫吗", "你好棒", "你是谁", - "你会什么", "你开心吗", "早安", "今天好吗", - "你在做什么", "猫好吗", "你喜欢什么颜色", - # 事实类 (Expert-0) - "猫是什么", "猫在做什么", "猫吃什么", "天是什么颜色", - # 推理类 (Expert-1) - "什么会飞", "什么会跑", "什么会游", "什么吃草", - ] - NL = CH2I['\n'] - for seed in seeds: - ctx = [CH2I[c] for c in seed if c in CH2I] - if not ctx: - continue - ctx.append(NL) - seed_len = len(ctx) - - with torch.no_grad(): - x = torch.tensor([ctx[-CONTEXT:]], dtype=torch.long) - _, gate_logits = model(x, labels=None) - expert_idx = gate_logits.argmax(dim=-1).item() - gate_probs = F.softmax(gate_logits[0], dim=-1) - - with torch.no_grad(): - for _ in range(20): - x = torch.tensor([ctx[-CONTEXT:]], dtype=torch.long) - logits, _ = model(x, labels=None) - probs = F.softmax(logits[0, -1] / 0.3, dim=-1) - nxt = torch.multinomial(probs, 1).item() - ctx.append(nxt) - gen_ids = ctx[seed_len:] - gen = "".join(I2CH[i] for i in gen_ids) - answer = gen.split("\n")[0] if "\n" in gen else gen - - probs_str = " ".join(f"{gate_probs[i].item():.2f}" for i in range(NUM_EXPERTS)) - print(f" Q: {seed} [E{expert_idx} | {probs_str}]") - print(f" A: {answer}") - print() - except Exception: - print("[ERROR] generation failed:") - traceback.print_exc() - - -def export_int8(model): - import struct - state = model.state_dict() - tensors = [] - names = [] - for name, p in state.items(): - if name.endswith('.log_decay'): - val = torch.sigmoid(p).cpu() - else: - val = p.cpu() - tensors.append(val) - names.append(name) - - quant_data = bytearray() - scales = [] - lengths = [] - for t in tensors: - t_flat = t.flatten() - abs_max = t_flat.abs().max().item() - if abs_max < 1e-8: - scale = 1.0 - else: - scale = abs_max / 127.0 - q = torch.clamp((t_flat / scale).round(), -128, 127).to(torch.int8) - quant_data.extend(q.numpy().tobytes()) - scales.append(scale) - lengths.append(len(q)) - - with open("bitflow_moe_weights.bin", "wb") as f: - nt = len(tensors) - f.write(struct.pack('= 128: v -= 256 - r[j] = v * s - return r - -SMW = [] -for e in range(NE): - b0 = 2 + e * 12 - SMW.append(dqS(b0)) # norm1.w - SMW.append(dqS(b0 + 1)) # norm1.b - SMW.append(dqS(b0 + 2)) # log_decay - SMW.append(dqS(b0 + 8)) # norm2.w - SMW.append(dqS(b0 + 9)) # norm2.b - -rf1b = dqS(RB + 1) -rf2b = dqS(RB + 3) -nfw = dqS(NB) -nfb = dqS(NB + 1) - -# ========== 融合解量化+矩阵乘 ========== -def dqmv(i, ni, no, x, out): - s = SC[i]; o = OFF[i] - for r in range(no): - ss = 0.0; b = r * ni - for c in range(ni): - v = W[o + b + c] - if v >= 128: v -= 256 - ss += v * s * x[c] - out[r] = ss - -def mvW(w, ni, no, x, out): - for r in range(no): - ss = 0.0; b = r * ni - for c in range(ni): - ss += w[b + c] * x[c] - out[r] = ss - -def ln(w, b, x, o): - m = sum(x) / D - v = sum((x[i] - m) ** 2 for i in range(D)) / D - iv = 1.0 / math.sqrt(v + 1e-5) - for i in range(D): - o[i] = w[i] * (x[i] - m) * iv + b[i] - -def e1(x): - return x + 1.0 if x >= 0 else math.exp(x) - -def sg(x): - if x >= 0: - e = math.exp(-x); return 1.0 / (1.0 + e) - e = math.exp(x); return e / (1.0 + e) - -def ge(x): - c = 0.7978845608 - return x * sg(2.0 * c * (x + 0.044715 * x * x * x)) - -def sm(l, t=0.3): - m = -1e9 - for i in range(VS): - v = l[i] - bcnt[i] * 2.0 - if v > m: m = v - s = 0.0 - for i in range(VS): - v = l[i] - bcnt[i] * 2.0 - e = math.exp((v - m) / t) - bpr[i] = e; s += e - r = random.random() * s; c = 0.0 - for i in range(VS): - c += bpr[i] - if c >= r: return i - return VS - 1 - -# ========== 状态 ========== -S = [[0.0] * (D * D) for _ in range(NL)] -z = [[0.0] * D for _ in range(NL)] -bx = [0.0] * D; bh = [0.0] * D -bq = [0.0] * D; bk = [0.0] * D -bv = [0.0] * D; bgd = [0.0] * D -bnum = [0.0] * D; by = [0.0] * D -bout = [0.0] * D; bff = [0.0] * FD -bpr = [0.0] * VS; bcnt = [0] * VS -brh = [0.0] * RD; brl = [0.0] * NE -blg = [0.0] * VS -bsum = [0.0] * D; brx = [0.0] * D; bct = [0] - -# ========== GLA ========== -def gs(e, layer, h): - b0 = 2 + e * 12; dec = SMW[e * 5 + 2] - dqmv(b0 + 3, D, D, h, bq) - dqmv(b0 + 4, D, D, h, bk) - dqmv(b0 + 5, D, D, h, bv) - dqmv(b0 + 6, D, D, h, bgd) - for i in range(D): - bq[i] = e1(bq[i]); bk[i] = e1(bk[i]) - for i in range(D): - bgd[i] = sg(bgd[i]) * dec[i] - sl = S[layer]; zl = z[layer] - for i in range(D): - si = i * D; gd = bgd[i]; ki = bk[i] - for j in range(D): - sl[si + j] = gd * sl[si + j] + ki * bv[j] - for i in range(D): - zl[i] = bgd[i] * zl[i] + bk[i] - for j in range(D): - bnum[j] = 0.0 - for i in range(D): - si = i * D; qi = bq[i] - for j in range(D): - bnum[j] += qi * sl[si + j] - den = 1e-6 - for i in range(D): - den += bq[i] * zl[i] - iv = 1.0 / den - for j in range(D): - by[j] = bnum[j] * iv - dqmv(b0 + 7, D, D, by, bout) - -def bs(e, x, layer): - b0 = 2 + e * 12 - ln(SMW[e * 5], SMW[e * 5 + 1], x, bh) - gs(e, layer, bh) - for i in range(D): - x[i] += bout[i] - ln(SMW[e * 5 + 3], SMW[e * 5 + 4], x, bh) - dqmv(b0 + 10, D, FD, bh, bff) - for i in range(FD): - bff[i] = ge(bff[i]) - dqmv(b0 + 11, FD, D, bff, bout) - for i in range(D): - x[i] += bout[i] - -def route(x): - dqmv(RB, D, RD, x, brh) - for i in range(RD): - brh[i] = max(0.0, brh[i] + rf1b[i]) - dqmv(RB + 2, RD, NE, brh, brl) - for i in range(NE): - brl[i] += rf2b[i] - bi = 0; bv2 = brl[0] - for i in range(1, NE): - if brl[i] > bv2: - bv2 = brl[i]; bi = i - return bi, brl[:] - -def fw(ctx, gen_flag=True): - t = min(len(ctx) - 1, CT - 1); ti = ctx[-1] - s0 = SC[0]; o0 = OFF[0]; s1 = SC[1]; o1 = OFF[1] - for i in range(D): - v0 = W[o0 + ti * D + i] - if v0 >= 128: v0 -= 256 - v1 = W[o1 + t * D + i] - if v1 >= 128: v1 -= 256 - bx[i] = v0 * s0 + v1 * s1 - # 累加运行均值 (匹配训练时router的x.mean(dim=1)) - for i in range(D): bsum[i] += bx[i] - bct[0] += 1 - for i in range(D): brx[i] = bsum[i] / bct[0] - ei, gate = route(brx) - for layer in range(NL): - bs(ei, bx, layer) - ln(nfw, nfb, bx, bh) - if gen_flag: - dqmv(HI, D, VS, bh, blg) - return blg, ei, gate - return None, ei, gate - -def rs(): - for l in range(NL): - for i in range(D * D): - S[l][i] = 0.0 - for i in range(D): - z[l][i] = 0.0 - for i in range(D): - bsum[i] = 0.0 - bct[0] = 0 - -def gen(ids, n): - NL_ch = VS - 1 - ids = list(ids) + [NL_ch] - sl = len(ids); c = list(ids); rs() - for i in range(VS): - bcnt[i] = 0 - # 种子处理 - for i in range(len(c) - 1): - fw(c[:i + 1], False) - # 生成 - expert_names = ["事实", "推理", "问答"] - for _ in range(n): - lg, ei, gate = fw(c) - nx = sm(lg) - bcnt[nx] += 1 - c.append(nx) - gen_ids = c[sl:] - gen = "".join(I2CH.get(i, '?') for i in gen_ids) - answer = gen.split("\n")[0] if "\n" in gen else gen - return answer, ei - -# ========== 测试 ========== -seeds = [ - "你好", "你好吗", "谢谢", "再见", "你叫什么名字", - "你喜欢什么", "你喜欢猫吗", "你好棒", "你是谁", - "你会什么", "你开心吗", "早安", "今天好吗", - "你在做什么", "猫好吗", "你喜欢什么颜色", - "猫是什么", "猫在做什么", "猫吃什么", "天是什么颜色", - "什么会飞", "什么会跑", "什么会游", "什么吃草", -] - -print("=" * 60) -print("PC模拟器验证 (模拟micro:bit int8推理)") -print("=" * 60) - -correct = 0; total = 0 -for seed in seeds: - ctx = [CH2I[c] for c in seed if c in CH2I] - if not ctx: - continue - random.seed(42) # 固定随机种子便于复现 - answer, ei = gen(ctx, 20) - expert_names = ["事实", "推理", "问答"] - print(f"Q: {seed:12s} [E{ei}={expert_names[ei]}]") - print(f"A: {answer}") - print() - total += 1 - # 简单验证: 答案非空且包含相关字符 - if answer and len(answer) > 0: - correct += 1 - -print(f"\n结果: {correct}/{total} 有效输出") - -# 内存估算 -mem_w = len(W) -mem_S = NL * D * D * 4 -mem_z = NL * D * 4 -mem_buf = (9 * D + FD + VS + VS + RD + NE + VS + D + D) * 4 # +bsum, brx -mem_SMW = NE * 5 * D * 4 -mem_small = (RD + NE + D + D) * 4 # rf1b, rf2b, nfw, nfb -mem_lists = 3 * nt * 8 # LEN, SC, OFF as Python lists -total_mem = mem_w + mem_S + mem_z + mem_buf + mem_SMW + mem_small + mem_lists -print(f"\n内存估算 (micro:bit):") -print(f" W (int8权重): {mem_w:>6d} B = {mem_w/1024:.1f}KB") -print(f" S (GLA状态): {mem_S:>6d} B = {mem_S/1024:.1f}KB") -print(f" z (归一化): {mem_z:>6d} B = {mem_z/1024:.1f}KB") -print(f" 缓冲区: {mem_buf:>6d} B = {mem_buf/1024:.1f}KB") -print(f" SMW (小权重): {mem_SMW:>6d} B = {mem_SMW/1024:.1f}KB") -print(f" 路由/norm_f: {mem_small:>6d} B = {mem_small/1024:.1f}KB") -print(f" LEN/SC/OFF: {mem_lists:>6d} B = {mem_lists/1024:.1f}KB") -print(f" 总计: {total_mem:>6d} B = {total_mem/1024:.1f}KB") -print(f" micro:bit可用: ~108KB (128KB-20KB解释器)") -print(f" 余量: {108*1024 - total_mem:>6d} B = {(108*1024 - total_mem)/1024:.1f}KB")