776 lines
28 KiB
Python
776 lines
28 KiB
Python
"""
|
|
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('<H', nt))
|
|
f.write(struct.pack('<%dH' % nt, *lengths))
|
|
f.write(struct.pack('<%df' % nt, *scales))
|
|
f.write(bytes(quant_data))
|
|
|
|
total_bytes = len(quant_data)
|
|
header_bytes = 2 + 2 * len(tensors) + 4 * len(tensors)
|
|
print(f"\n[export] int8 weights: {nt} tensors, data={total_bytes}B ({total_bytes/1024:.1f}KB), "
|
|
f"header={header_bytes}B, total={total_bytes+header_bytes}B ({(total_bytes+header_bytes)/1024:.1f}KB)")
|
|
print(f" scales range: [{min(scales):.6f}, {max(scales):.6f}]")
|
|
|
|
print("\n[param layout]")
|
|
for name, t in zip(names, tensors):
|
|
print(f" {name:40s} shape={list(t.shape)} numel={t.numel()}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|