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