# TangYuGe MTP Assistant

配合 TangYuGe 文字主模型使用的四层 assistant（`Gemma4UnifiedAssistantForCausalLM`）：先由它草拟若干 token，再由主模型验证接受，用于缩短生成时间。完整 BF16 权重约 **846 MB**。

> **状态：实验候选，未通过正式验收。** 二十题在线对照中总体接受率 0.5231、答案段 0.4423，均低于项目设定的 0.60 门槛，且边界段相对未适配头下降 10.7 个百分点；正式门槛要求的 32 题只完成 20 题。请按实验件使用，不要当作已验证的加速方案，也不要替换根目录的默认生成入口。

## 模型与训练

| 项目 | 内容 |
| --- | --- |
| 基础 assistant | `google/gemma-4-12B-it-qat-q4_0-unquantized-assistant` |
| 固定 revision | `18934064dd4c5c6cc3621f6381e7d377fc8cb7bd` |
| 架构 | `Gemma4UnifiedAssistantForCausalLM`（4 层，隐藏维 1024） |
| 训练范围 | 只更新 `pre_projection.` 与 `post_projection.` 两组投影；主模型与 assistant 骨干全部冻结 |
| 训练数据 | 210 条完整轨迹 + 22 条源组隔离校准轨迹 |
| 训练与选择 | 训练 4096 步；按校准集选中第 **3072** 步导出 |
| 超参 | 学习率 2e-5，梯度累积 4，思考/答案/边界分层采样 0.35/0.55/0.10 |
| 精度 | FP32 优化，BF16 导出 |
| 许可证 | Apache-2.0 |

轨迹由冻结的 TangYuGe 主模型在 thinking 开启时生成，全部自然结束，没有强制闭合思考段。测试集不参与训练，也不参与选点。因为只训练两组投影，导出头与未适配的官方 assistant（`67f1420c…`）在权重上确有差异，但差异范围仅限这两组张量。

## 使用

下载完整仓库并安装根目录依赖后运行：

```bash
python TangYuGe-model/mtp/inference.py --message "你好，今天想聊聊。"
```

- 入口默认开启 thinking，默认思考额度 **512** token、最大生成 **1536** token、草稿长度 **3**。
- 思考额度耗尽时只闭合思考通道，最终回答仍由主模型生成；`--thinking-budget-tokens 0` 表示不限定思考额度。
- 运行需要 CUDA。**CPU 上的主模型＋assistant 联合生成未验证**，入口会直接拒绝非 CUDA 环境。
- 入口会校验主模型与 assistant 的绑定 SHA256，不匹配即报错，不静默加载。
- 根目录的文字入口只用主模型生成，便于与本入口对照。

## 二十题在线三臂对照

测试条件：单张 4090 48GB、thinking 开启、512 token 思考额度、3 token 草稿长度、同一批提示。

| 配置 | 完整回复 | 总体接受率 | 思考段 | 答案段 | 边界段 | token/s |
| --- | --- | --- | --- | --- | --- | --- |
| 关闭 MTP（主模型单独生成） | 20/20 | — | — | — | — | 14.24 |
| 未适配的官方 assistant | 20/20 | 0.4904 | 0.4970 | 0.4177 | 0.5063 | 28.94 |
| **本版 TangYuGe assistant** | 20/20 | **0.5231** | 0.5336 | 0.4423 | 0.3993 | **30.52** |

相对关闭 MTP，吞吐提升约 **+114%**；相对未适配头，总体接受率提高约 3.27 个百分点、答案段提高约 2.46 个百分点。逐项原始统计见 [`evaluation.json`](evaluation.json)。

口径提醒：**"完整回复"是完成率，不等于正确率**；单卡二十题的吞吐量不能外推到所有输入与设备；无额度 thinking 模式未做完整回复验证。

## 验收状态：未通过

项目预设的判定门槛与本次实测的对照：

| 门槛 | 阈值 | 实测 | 结果 |
| --- | --- | --- | --- |
| 答案段接受率 | ≥ 0.60 | 0.4423 | ✗ |
| 总体接受率 | ≥ 0.60 | 0.5231 | ✗ |
| 各通道相对未适配头的下降 | ≤ 1 个百分点 | 边界段 0.5063 → 0.3993，**降 10.7 个百分点** | ✗ |
| 吞吐增益（相对关闭 MTP） | ≥ 10% | +114% | ✓ |
| 正式题量 | 32 题 | 完成 20 题 | ✗（按指示停止） |

四项不达标，因此本版如实标注为未通过验收的实验候选，不替换默认入口。

## 已知限制与选择口径

1. **选点口径有缺陷（重要）**：本版第 3072 步是用"固定前缀校准探针"选出的，而该探针系统性高估在线接受率——同一批检查点在探针上的答案段为 0.91–0.94，在线实测只有 0.44。后续版本必须改用**在线式接受率**选点，校准探针只能当训练期监控。
2. **边界通道有取舍代价**：总体与答案段提升的同时，边界段明显回落，说明性能在通道间转移而非全面提升。
3. **答案段仍未达门槛**：按事先写定的硬停条件，这属于**结构问题**而非数据不足——继续加采样轨迹不会有本质改善，应转向结构改造（加秩、低秩适配器或分段监督重设计）。
4. **32 题正式验收未完成**：仅完成 20 题即按指示停止；正式结论待补。
5. **无额度 thinking 未验证**：额度耗尽时强制闭合思考通道与训练数据的自然结束不同，不能据此宣称无限额思考也可用。

## 验证与文件

本机 CPU 严格加载通过，49 个运行态张量哈希与云端一致；所选训练检查点的恢复探针通过。

| 文件 | 内容 |
| --- | --- |
| `model.safetensors` | 完整 assistant 的 BF16 权重（845,719,296 字节） |
| `config.json`、`generation_config.json` | 架构与生成配置 |
| `binding.json` | 来源、revision、训练范围、超参与主模型绑定 SHA |
| `evaluation.json` | 二十题三臂分段接受率、门槛与违反项 |
| `inference.py` | 显式推理入口（含绑定 SHA 校验与严格加载检查） |
| `thinking_budget.py` | 思考额度处理器 |

主模型 SHA256：`f1666a7bbe26e487dea3b254b109e180a952357fdf5335fd273422d66c29f488`
本版 assistant SHA256：`c412d658c9975406d2053b4d4b9dceaf98d5606463258f73999748223f6e5758`

本目录沿用仓库根目录的 [Apache-2.0 许可证](../LICENSE)，来源归属见 [NOTICE](../NOTICE)。

## 版本历史

| 版本 | 训练 | 选中 | 二十题在线总体/答案 | 结论 |
| --- | --- | --- | --- | --- |
| v1 | 2048 步 | 第 1280 步 | 八题 0.4870 / 0.3241 | 未通过；仅八题诊断 |
| **v2（本目录）** | 4096 步 | 第 3072 步 | **0.5231 / 0.4423** | 未通过；二十题，边界段回落 |

v1 头（`268fb5cd…`）已被 v2 覆盖，不在本仓库保留。