SudenMind 官方文档

PyTorch License Python Website

基于 AttnRes (Attention with Residual) 架构的开源中文对话生成模型,专为中文对话场景深度优化

项目介绍

SudenMind 是一个基于 AttnRes (Attention with Residual) 架构的中文对话生成模型。采用 Encoder-Decoder 设计,使用自定义 2 层 AttnRes 编码器 + 6 层 AttnRes 解码器,支持跨层残差连接,使用 ShareGPT/ChatML 业界标准对话格式,适配 ChatGLM tokenizer,针对中文对话场景优化。新增门控机制(Gate)和 KV 缓存(KV Cache)推理加速。

当前版本:v6.1 (AttnRes Gate & KV Cache 版本) | 最后更新:2026-03-29

核心特性

AttnRes 架构

每层可动态关注之前所有层的输出,信息流动更丰富,解决深层网络梯度消失问题

自定义 Encoder-Decoder

2层AttnRes编码器 + 6层AttnRes解码器,每层都可访问之前所有层的输出

门控机制 (Gate)

创新的门控网络动态平衡MoE输出与跨层残差输出,根据输入特征自适应调整融合比例

KV 缓存推理

支持键值缓存,避免重复计算历史token的键值对,将时间复杂度从 O(n²) 降至 O(n)

MoE 混合专家集成

所有 FFN 层替换为 MoE 层,8 个专家,top_k=2,在可控算力下提升模型表达能力

标准对话格式

使用 ShareGPT/ChatML 业界标准格式,兼容 OpenAI/Qwen/ChatGLM2/3 等主流模型

混合精度训练

FP16 加速训练,大幅节省显存,同时支持 NVIDIA CUDA 和 AMD ROCm 双平台

ONNX 导出

支持导出 ONNX 格式模型,可用 Netron 可视化结构,适配多端部署场景

模型架构

SudenMind 基于 Attention Residuals 论文的 Encoder-Decoder 架构设计,整体流程如下:

输入 (batch, seq_len)
    ↓
Token Embedding + Position Encoding
    ↓
AttnRes Encoder × 2 ← 自定义编码器(双向注意力)
  ├─ 第0层: 双向自注意力 → MoE → output_0
  ├─ 第1层: 双向自注意力 → MoE → AttnRes([output_0]) → output_1
  └─ ...每层都可访问之前所有层的输出
    ↓
AttnRes Decoder × 6 ← 解码器(因果注意力)
  ├─ 自注意力 (带因果掩码)
  ├─ 跨层残差注意力 (AttnRes)
  ├─ 门控特征融合 (Gate)
  └─ MoE 前馈网络 (8 个专家,top_k=2)
    ↓
Linear → Softmax
    ↓
输出 (batch, seq_len, vocab_size=65024)

AttnRes 门控机制详解

在每个 AttnRes 层中,我们设计了创新的门控网络,动态平衡新生成特征与历史残差特征的融合比例,计算公式如下:

Gate = Sigmoid([MoE_out; Res_out] * W_g)

Output = (1 - Gate) ⊗ MoE_out + Gate ⊗ Res_out

其中:

  • MoE_out:混合专家前馈网络的输出特征
  • Res_out:跨层残差注意力的输出特征
  • W_g:可学习的门控权重矩阵
  • :逐元素乘法运算
这个机制使模型能够根据当前输入的语义特征,动态决定采用更多的创新特征(MoE输出)还是保留历史信息(残差输出),显著提升了建模灵活性和长文本理解能力。

KV 缓存推理加速

为了大幅提升自回归生成的速度,我们在解码过程中实现了历史键值对(Key-Value pairs)缓存机制:

  • 编码器输出的键值对在首次前向传播时计算并全局缓存
  • 解码器每一步只需计算当前token的键值对,将其追加到缓存中
  • 避免了重复计算历史token的键值对,将自回归生成的时间复杂度从 O(n²) 降至 O(n)
KV 缓存仅在推理阶段生效,训练阶段不启用,避免对训练过程的梯度传播产生影响。

ShareGPT/ChatML 对话格式规范

SudenMind 使用业界标准的 ShareGPT/ChatML 对话格式,在 data_utils.py 中实现了完整的格式处理逻辑,规则如下:

  • 完整对话以 [BOS] 开头,以 [EOS] 结尾
  • 每一轮对话以 <|im_start|>{role}\n 开头,以 <|im_end|> 结束
  • 仅对 assistant 轮次的内容部分计算 loss;角色头、结束符、换行符等位置的loss权重设为 -100(忽略)

训练格式

[BOS] <|im_start|>user
你好
<|im_end|>
<|im_start|>assistant
你好!很高兴为你服务。
<|im_end|>
<|im_start|>user
今天天气怎么样?
<|im_end|>
<|im_start|>assistant
今天天气很好!
<|im_end|>
[EOS]

推理格式

[BOS] <|im_start|>user
你好
<|im_end|>
<|im_start|>assistant
你好!很高兴为你服务。
<|im_end|>
<|im_start|>user
今天天气怎么样?
<|im_end|>
<|im_start|>assistant

推理时,模型会基于上述前缀文本,续写 assistant 的回复内容,直到遇到 <|im_end|>[EOS] 结束符停止生成。

与业界主流模型格式对比

模型 对话格式 核心特点
SudenMind <|im_start|>user / <|im_start|>assistant 均以 <|im_end|> 结束 标准 ShareGPT/ChatML,支持 MoE + Gate + KV Cache
OpenAI GPT-4 ⟨im_start⟩user⟨im_sep⟩{内容}⟨im_end⟩ 使用 ⟨im_sep⟩ 分隔符
Qwen ⟨im_start⟩user\n{内容}⟨im_end⟩ 完全兼容 ShareGPT/ChatML
ChatGLM2/3 [BOS] ⟨user⟩\n{内容}⟨/assistant⟩ 语义类似,角色分隔形式略有差异

项目文件结构

SudenMind/
├── config.json             # ⭐ 超参数配置文件
├── data/
│   └── cache/              # 数据集缓存目录
├── model/                  # 模型保存目录
│   ├── sudenmind.pth       # 训练好的模型权重
│   └── sudenmind.onnx      # 导出的ONNX格式模型
├── src/                    # 核心源代码
│   ├── model.py            # AttnRes 模型定义(含门控机制和KV缓存)
│   ├── data_utils.py       # ShareGPT/ChatML 格式数据处理
│   ├── train.py            # 训练脚本
│   ├── chat.py             # 交互式对话(支持KV缓存推理)
│   ├── moe.py              # MoE 混合专家模块
│   └── view_module.py      # 模型结构可视化工具
├── docs/                   # 项目文档
│   ├── index.html          # 文档主页
│   └── README_EN.md        # 英文说明文档
├── README.md               # 项目说明文档
├── requirements.txt        # 依赖声明
├── run.sh                  # 一键运行脚本
└── setup.py                # 项目安装配置

配置说明

项目所有超参数集中在 config.json 中统一管理,分为模型、训练、数据、生成四大模块,详细说明如下:

{
  "model": {
    "d_model": 512,
    "d_fnn": 1024,
    "nhead": 8,
    "n_layers": 6,
    "dropout": 0.1,
    "max_position_embeddings": 5000,
    "num_experts": 8,
    "top_k": 2,
    "aux_loss_coef": 0.01,
    "tokenizer_name": "THUDM/chatglm-6b"
  },
  "training": {
    "lr": 0.0003,
    "weight_decay": 0.01,
    "betas": [0.9, 0.95],
    "eps": 1e-08,
    "warmup_ratio": 0.05,
    "max_epochs": 200,
    "label_smoothing": 0,
    "ignore_index": -100,
    "batch_size": 1,
    "max_norm": 5.0,
    "patience": 30,
    "target_loss": 0.03,
    "use_amp": false
  },
  "data": {
    "config": "base",
    "split": "train",
    "max_history": 5,
    "max_seq_len": 512
  },
  "generation": {
    "max_length": 200,
    "temperature": 0.3,
    "top_k": 50,
    "onnx_seq_len": 32,
    "onnx_opset": 11
  },
  "paths": {}
}
当前版本不再在 config.json 里通过显式字段声明 <|im_start|>/<|im_end|> 等特殊 token,统一使用 THUDM/chatglm-6b tokenizer 的内置映射。

快速开始

1. 环境准备

基础环境要求

  • Python 3.8+ (推荐 3.10)
  • PyTorch 2.0+
  • CUDA 11.8+ 或 ROCm 6.4+ (GPU加速)
  • 最低 16GB 显存 (推荐 24GB+)

步骤1:创建conda环境

conda create -n sudenmind python=3.10
conda activate sudenmind

步骤2:安装PyTorch

根据你的硬件平台选择对应的安装命令:

AMD ROCm 版本

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.4

NVIDIA CUDA 版本

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

步骤3:安装其他依赖

pip install transformers>=4.30.0 datasets>=2.0.0 sentencepiece>=0.1.99 onnx>=1.14.0 onnxruntime>=1.15.0 netron>=7.0.0

2. 训练模型

执行以下命令启动模型训练:

python src/train.py

训练内置特性

  • LCCC 中文对话数据集自动加载
  • ShareGPT/ChatML 格式自动转换
  • 混合精度训练 (FP16) 支持
  • Cosine Annealing 学习率调度
  • 早停机制,防止过拟合

3. 对话测试

训练完成后,执行以下命令启动交互式对话,内置支持 KV 缓存推理加速:

python src/chat.py

对话示例

You: 你好
Assistant: 你好!很高兴为你服务。
You: 今天天气怎么样?
Assistant: 今天天气很好,阳光明媚!

内置命令

  • quit:退出对话系统
  • clear:清空对话历史
  • history:查看完整对话历史

参考文献

  1. Kimi Team, et al. "Attention Residuals." arXiv preprint arXiv:2603.15031 (2026). 链接
  2. Du, Zhengxiao, et al. "GLM: General Language Model Pretraining with Autoregressive Blank Infilling." arXiv preprint arXiv:2103.10360 (2021). 链接

版本变更日志

v6.1 (当前版本)
2026-03-29
  • 移除gMASK标记:简化对话格式,仅保留[BOS]作为序列开始标记
  • 新增门控机制 (Gate):动态平衡MoE输出与跨层残差输出
  • 新增KV缓存推理:大幅提升自回归生成速度
  • 优化模型架构:改进AttnRes层的信息流动机制
  • 更新网站链接:官网迁移至 ai.sufun.space
v6.0
2026-03-XX
  • 完全重构为 ShareGPT/ChatML 格式
  • 删除所有中文角色前缀 ("用户:" / "助手:")
  • 使用业界标准格式,兼容 OpenAI/Qwen/ChatGLM2/3 等主流模型
v5.0
2026-03-XX
  • 移除 BERT 编码器,改为自定义 2 层 AttnResEncoder
  • 适配 ChatGLM tokenizer(词表大小 65024)

常见问题

训练时出现显存不足怎么办?

可以尝试以下优化方案:

  1. 在 config.json 中开启 use_amp: true 启用混合精度训练
  2. 减小 batch_size,最低可设置为1
  3. 减小 max_seq_len,推荐设置为256或128
  4. 减小 d_modeld_fnn 的数值
如何使用自己的数据集进行训练?

只需将你的数据集转换为 ShareGPT/ChatML 格式,然后修改 data_utils.py 中的数据集加载逻辑即可。项目内置了 LCCC 数据集的自动加载,可参考其实现适配你的自定义数据集。

模型支持哪些部署方式?

SudenMind 支持多种部署方式:

  • 原生 PyTorch 部署,直接使用 chat.py 进行交互式推理
  • 导出为 ONNX 格式,适配 ONNX Runtime 跨平台部署
  • 可适配 TensorRT、OpenVINO 等推理引擎进行加速
项目使用什么开源协议?

项目采用 MIT 开源协议,可自由用于个人学习、科研和商业项目,详情可查看仓库中的 LICENSE 文件。