File size: 3,816 Bytes
4f210d5 f1b8f61 4f210d5 f1b8f61 4f210d5 f1b8f61 4f210d5 f1b8f61 4f210d5 f1b8f61 4f210d5 f1b8f61 4f210d5 f1b8f61 4f210d5 f1b8f61 a49846b f1b8f61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | ---
license: apache-2.0
language:
- zh
- en
tags:
- pondernet
- mixture-of-experts
- adaptive-computation
- rent-math
base_model: XHToken/Spark-X2.5-4B
---
# samai-4b · PonderNet + dMoE (Spark-X2.5-4B backbone)
将 samai-2b 的 PonderNet/dMoE 改造方案移植到
**XHToken/Spark-X2.5-4B** 混合注意力基座, 并经过租金数学 SFT 的 4B 因果语言模型。
## 架构
- **骨干**: Spark-X2.5-4B — 36 层混合注意力 (3 滑窗 512 + 1 全注意力/组), q_k_v 融合投影,
headwise 输出门控, full 层 partial rotary (0.25, theta 5M), vocab 131072, tie embedding
- **PonderNet**: 后 8 层 (L28-35) 为 Ponder 层, learned halting, 期望步数随输入难度自适应
(实测: 闲聊 ~3.3 步 / 算术 ~5.3 步 / 多步应用题 ~4.9 步 / 歧义题 ~6.2 步)
- **dMoE**: 8 个 Ponder 层的 MLP 拆为 8 专家 + 可学习 router (tau=0.125 稀疏激活, Switch 式
负载均衡, STE 梯度)。零初始化 router (=稠密等价) 起步, SFT 后实测平均每 token 仅激活
**1.4~4.8/8** 专家, gate 分布随输入类型偏移 (数字类任务强偏好 L31-e4, p≈0.21)
- **MTP**: 2 个 multi-token 预测头, 仅训练期辅助, 推理旁路
## 上下文
- `max_position_embeddings = 200000` (默认)
- 基座原生训练长度 128k; 128k~200k 为外推区 (滑窗层不受影响, 9 个全注意力层有衰减)
- 如需 200k 内更平稳质量, 可在 config 加 `"rope_scaling": {"type":"linear","factor":1.53}`
(建模代码已支持 linear 插值)
- 显存参考: T4 16GB fp16 下 KV 实际可服务 ~4-5 万 token 长输入
## 训练
- SFT: 831 条租金/算术/闲聊数据 × 2 epochs, LoRA r16 a32 (attn + 全 MLP) +
router/ponder_head 全参, 可训练 39.5M (0.94%), loss 0.611
- LoRA adapter 存于 `sft_adapter/`; 本仓库主分支为合并后全量权重
## 评测 (831 数据外, 3 采样多数)
| 测试 | 结果 |
|---|---|
| 原租金题 (干扰数字两步推理) | **3/3 = 2400 元** |
| 天数折算变体 ×4 | 全对 |
| 满月/中途退租陷阱 ×2 | 全对 |
| 简单算术 | 对 |
| 费用求和 (漏加小额费用) | 0/3 ✗ |
## 使用
### transformers
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
m = AutoModelForCausalLM.from_pretrained("tchbcb/samai-4b",
trust_remote_code=True, dtype=torch.float16).cuda()
tok = AutoTokenizer.from_pretrained("tchbcb/samai-4b")
text = tok.apply_chat_template([{"role":"user","content":"你好"}],
add_generation_prompt=True, tokenize=False, enable_thinking=True) + "\n"
```
### OpenAI 兼容 API (key=1234)
```bash
python openai_api.py # port 7862
curl http://127.0.0.1:7862/v1/chat/completions \
-H "Authorization: Bearer 1234" -H "Content-Type: application/json" \
-d '{"model":"samai-4b","messages":[{"role":"user","content":"你好"}]}'
```
- 默认强制思考 (与 SFT 一致); `enable_thinking: false` 跳过思考
- 思考文本返回在 `message.reasoning_content` (DeepSeek-R1 风格), 支持 `stream: true`
- 附带字段: `samai_ponder_steps` (本次生成平均 ponder 步数)
### Colab 一键部署 (聊天页 7861 + API 7862)
```
!wget -qO /content/boot.py https://huggingface.co/tchbcb/samai-4b/resolve/main/bootstrap.py
!python3 /content/boot.py
```
## 文件
| 文件 | 说明 |
|---|---|
| model-*.safetensors | SFT 合并后全量权重 (fp16, 8.3GB) |
| modeling_samai_4b.py | 自包含建模代码 (transformers 5.16.1 适配) |
| chat_server.py | 网页聊天服务 (端口 7861, 强制思考/反循环/投票) |
| openai_api.py | OpenAI 兼容 API (端口 7862, key=1234) |
| control_server.py | 远程控制通道 (端口 5000, 开发调试用) |
| bootstrap.py | Colab 一键引导 (双服务三隧道) |
| sft_adapter/ | LoRA adapter (r16 + router/ponder_head) |
| eval_report.json | SFT 后评测原始数据 |
|