---
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 后评测原始数据 |