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

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)

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