""" 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('