Qwen3.5_4B_MLP

基于 Qwen3.5-4B 的草案链合并模型:用 4 层纯 MLP 草案链物理替换尾部 4 层 Transformer(原第 28-31 层已从权重中删除),几何监督训练,不含任何概率对象。

结构

部分	内容
真实层	28 层标准 Qwen3_5DecoderLayer(线性注意力 21 + 全注意力 7)
草案链	4 层 DraftDecoderLayer(MLP:hidden=2048,fc1→silu→fc2+α 残差,权重已固化)
最终 norm	恒等(草案链末态直接喂读出头)

被替换的 L28-L31(原每 4 层一个 full attention 的 L31 也在内)已物理删除,权重不驻留。

性能(权重固化 + 链 CUDA Graph 部署路径)

指标	值
层参数	3123.8M(原 3570.1M,−12.5%)
端到端 tok/s	27.7(固定 128 步贪心,同口径完整模型 21.2)
层堆单步	31.6 ms
峰值显存	7.34 GiB(原 8.05,−8.8%)
held-out consEval	~0.843(原分布持平,无回归)

chat 真实口径(模板+rep_penalty 1.1+RepeatStopping):约 27.8 tok/s,相对未优化路径 +5.5%。注意:K4 链在代码类 prompt 上仍有复读残留(约 2/6 探针),科普/散文类可用。

用法

import torch
from draft_chain import load_merged
model = load_merged("Qwen3.5_4B_MLP", dtype=torch.bfloat16)
# 生成时 enable_thinking=False, KV cache 常驻

或直接对话(需在本仓库配套脚本目录):

python chat.py --draft draftChainK4WidePlus.pt --graph

文件

• model.safetensors / config.json:合并后权重(28 真层 + 4 链层 + tied 词表),键前缀 model.language_model.*(transformers Qwen3.5 的保存映射,from_pretrained/load_merged 自动剥回)
• draft_chain.py:自定义层定义 + load_merged 加载入口(自包含)
• graphChain.py / draftMerge.py:链 CUDA Graph 化与合并工具
• generation_config.json:生成配置
• p1.pdf:论文

方法简述

草案链用几何监督训练(单位化后 (1−cos)+λ·angle,f32 计算,fp32 归一化),无 CE、无 softmax。合并时物理删除尾部 K 层、追加 K 层链、final norm 置恒等。部署侧两处修正:权重固化(训练后 w=g·v/‖v‖ 提前算好,免推理期每步归一化重算)与链整链 CUDA Graph(4 块约 36 kernel 折叠成一次提交)。