Qwen3.5_4B_MLP

Qwen3.5_4B_MLP

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

结构

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

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

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 折叠成一次提交)。