Downloads · 30 days
0
tchbcb/samaidev-pnet-2b
samaidev-pnet-2b is a machine learning model from tchbcb. Use it for the machine learning task on the model card, and read the license before you ship it in a product.
对象仓库: tchbcb/MiniCPM5-2B-cpu(即 openbmb/MiniCPM5-2B 的副本) 结论先行: 可以,而且改动比预想的小得多 —— 本目录给出可直接运行的完整实现 (ponderllama.py,约 400 行,含训练/推理/生成全链路 + 9 组冒烟测试全部通过)。
Downloads · 30 days
0
Access
Public
Updated Sep 16, 2026
Repo size
6 GB
Likes
0
Public
Click a slice to open those files.
.safetensors6 GB · 99%
From the Hugging Face model README
对象仓库:
tchbcb/MiniCPM5-2B-cpu(即openbmb/MiniCPM5-2B的副本) 结论先行: 可以,而且改动比预想的小得多 —— 本目录给出可直接运行的完整实现 (ponder_llama.py,约 400 行,含训练/推理/生成全链路 + 9 组冒烟测试全部通过)。
浅克隆(GIT_LFS_SKIP_SMUDGE=1 git clone --depth 1)后仓库仅 12MB,不含任何权重,
也没有自定义 modeling 代码。文件清单:
| 文件 | 内容 |
|---|---|
config.json | architectures: ["LlamaForCausalLM"], model_type: "llama" |
generation_config.json | 温度 1.0 / top_p 0.95 / eos [1, 130073] |
model.safetensors.index.json | 权重索引(381 个张量,单分片 ~5GB,未下载) |
tokenizer.json 等 | MiniCPM 词表(130,560) |
README*.md | 官方说明(transformers ≥5.6.2 原生加载,无需 trust_remote_code) |
关键架构参数:42 层 decoder、hidden 2048、16 头 / 2 KV 头(GQA)、head_dim 128、 RoPE θ=5e6、128K 上下文、SiLU MLP(intermediate 6144)。
分析结论:推理逻辑 100% 位于 transformers 原生 modeling_llama.py,其
LlamaModel.forward 是严格的"嵌入 → 42 层单遍串联 → RMSNorm → lm_head"结构,
没有任何递归/循环机制。要做 PonderNet,只需在原生代码之上包一层 —— 完全不必改
transformers 源码,也不必 fork 模型权重。
PonderNet(Banino et al., 2021)的核心:网络对每个样本反复执行同一计算块,每步输出 停机概率 λₙ,最终输出是各步输出的概率加权混合:
h_k = Block(h_{k-1}) # 思考块第 k 步
λ_k = Halting(h_k) ∈ (0,1) # 停机概率
w_k = remaining · λ_k # 本步"停下"的概率质量
mix = Σ w_k · h_k # 混合输出
remaining *= (1 - λ_k) # 直到 Σw ≥ 1-ε 或步数上限 K
42 层切成两段:前 34 层照常单遍,后 8 层(34~41)作为思考块循环执行
(ponder_start_layer=34 可调)。这样:
| 信号 | λ 定义 | 是否需训练 | 适用 |
|---|---|---|---|
entropy(默认) | 1 − 归一化输出熵,越自信越停 | 否 | 推理即插即用 |
msp | max softmax 概率 | 否 | 同上 |
drift | 相邻两次迭代隐状态的余弦相似度超过阈值 → 已收敛 → 停 | 否 | "想不动了就停" |
learned | nn.Linear(2048, 1) + sigmoid,真 PonderNet 头 | 是 | 配套 ponder loss |
诚实的说明:不训练时 learned 头是随机的、没有意义;推理可用的前三种是
启发式 halting(借模型自身的置信度/收敛度做停机判据),效果等价于
"自适应计算时间(ACT)+ PonderNet 式混合"。要得到论文语义的 λ,用下面提供的
loss 只训练这个 2049 参数的头(可冻结主干)即可。
每次思考迭代都会经 DynamicCache.update() 写 KV,直接循环会让缓存里出现重复 token。
本实现的解法 —— "替换式写入":
cache.layers[l].crop(-W) 撤销上一步写入,
再跑思考块补回 —— 任意时刻缓存长度恒等于真实 token 数(已用测试逐层断言);h_mix 重过一遍思考块,把"思考后"的
KV 写进缓存 —— 后续 token 注意到的正是思考后的表示;另一个坑:create_causal_mask 从非思考层读取缓存长度,对思考块会得到多 W 的
错误 KV 尺寸,因此窗口 mask 必须手工构建(_win_mask,token 索引因果 + padding
感知的加性浮点 mask,sdpa/eager 通用)。
labels 存在时自动计算 ponder loss = CE + β·KL(w ‖ Geometric(p_g)),
梯度可穿过混合权重到达 λ(与论文一致)。只训练 ponder_head(2049 参数)即可让
停机分布学出"难 token 多想、简单 token 少想"。
ponder_head.{weight,bias} 2 个张量(2049 参数),from_pretrained 直接加载;generate() 直接可用:继承 LlamaForCausalLM + 自定义 forward,HF 生成流程
(含左 padding、batch、logits_to_keep)无需任何适配;| 文件 | 说明 |
|---|---|
ponder_llama.py | 核心实现:PonderLlamaConfig + PonderLlamaForCausalLM(约 400 行) |
test_ponder.py | 随机权重冒烟测试,9 组用例(cache 一致性/生成/训练反传/4 信号/eager/左 padding) |
test_train.py | 训练管线验证,5 组用例(数据/三种模式/保存重载/KL 字段) |
test_gpu_readiness.py | T4/fp16 就绪性预验证,5 组用例(fp16 dtype 链/KL 不 NaN/ckpt 一致性) |
train_ponder_head.py | halting 头微调脚本(head / head+lora / head+block 三种模式,含难度分级数据生成器) |
eval_ponder.py | 效果评估:难度分桶 CE/步数统计、训前训后对比(--compare)、思考收益(--k1-baseline) |
ponder_steps_demo.py | tiny 模型演示:训练中平均思考步数向几何先验漂移的轨迹 |
demo_minicpm5.py | 真实权重 demo:chat 对比、每 token 思考步数统计 |
run_t4.sh | T4 一键脚本:自检→下载权重→数据→训练→评估全自动串联 |
data/pondernet_train.jsonl | 难度分级训练集 600 条(easy 220 / medium 200 / hard 180,答案反向构造保证正确) |
data/pondernet_train.eval.jsonl | 难度分级评估集 85 条(与训练集无重叠) |
from ponder_llama import PonderLlamaForCausalLM, PonderLlamaConfig, LlamaConfig
cfg = PonderLlamaConfig.from_llama(
LlamaConfig.from_pretrained("/path/to/MiniCPM5-2B-cpu"),
ponder_signal="entropy", # 或 msp / drift / learned
max_ponder_steps=8,
ponder_start_layer=None, # None = 最后 8 层作思考块
)
model = PonderLlamaForCausalLM.from_pretrained(
"/path/to/MiniCPM5-2B-cpu", config=cfg,
torch_dtype="bfloat16", low_cpu_mem_usage=True).eval().cuda()
out = model(input_ids, output_ponder=True)
print(out.ponder_steps) # 每个位置实际思考了几步
print(model._ponder_log) # 每次 forward 的诊断日志
训练 halting 头(可选,其余全部冻结):
for n, p in model.named_parameters():
p.requires_grad = n.startswith("ponder_head")
cfg.ponder_signal = "learned" # 切到可训练信号
# 正常 forward(input_ids, labels=...) → loss 里已含 KL(‖Geometric) 正则
主要超参:
| 参数 | 默认 | 含义 |
|---|---|---|
ponder_start_layer | L−8 | 思考块起始层 |
max_ponder_steps | 8 | 思考步数上限 K |
ponder_epsilon | 0.01 | 累计质量 1−ε 即停 |
ponder_signal | entropy | 停机信号 |
ponder_window | 1 | prefill 思考窗口(=1 即"回答前思考") |
ponder_all_positions | False | True = 所有位置独立 halting |
commit_kv | True | 用混合表示写 KV(False 省一遍计算,语义略糙) |
ponder_prior_p / ponder_loss_beta | 0.5 / 0.01 | 几何先验 p_g 与 KL 权重 β |
| 模式 | 可训练参数 | 显存 | 语义 |
|---|---|---|---|
head | 2049 | 极低(8GB 单卡可跑) | 仅学"何时停",思考块本身不变 |
head+lora(推荐) | 头 + LoRA(r=16,仅思考块) | 低 | 思考块学会"被重复执行",完整 PonderNet |
head+block | 头 + 思考块全参 | 高 | 效果上限最高 |
# 生成演示数据(16 条: 简单问答 + 推理题, 让头有"难题多想"的信号可学)
python train_ponder_head.py --model /path/to/MiniCPM5-2B-cpu --make-demo-data demo.jsonl
# 推荐: 头 + 思考块 LoRA
python train_ponder_head.py --model /path/to/MiniCPM5-2B-cpu \
--train-mode head+lora --lora-scope ponder --lora-r 16 \
--data demo.jsonl --epochs 2 --batch-size 2 --accum 4 \
--prior-p 0.4 --beta 0.05 --head-bias-init -1.0 \
--output out/ponder-lora --log-every 5
数据格式(jsonl,只对 assistant 回答区间计 loss):
{"messages": [{"role": "user", "content": "9.11 和 9.9 哪个大?"},
{"role": "assistant", "content": "9.9 更大。因为 0.90 > 0.11 …"}]}
--prior-p(几何先验 p_g):控制"想几步"的先验预期,期望步数 ≈ 1/p_g。
想让模型平均思考 3 步左右 → p_g=0.3;2.5 步 → p_g=0.4。--head-bias-init:初始停机偏置,sigmoid(-1.0)≈0.27(初始期望约 3.7 步)。--beta:KL 权重。太小 → 头被 CE 拉着全跑满 K 步;太大 → 全挤在第 1 步。训练日志中三列数字的联动(tiny 模型实测轨迹):
step | CE | KL | 平均思考步数
0 | 5.4395 | 0.0390 | 1.99
60 | 4.2772 | 0.0083 | 2.31 ← p_g=0.4, 期望 2.5
CE 下降(任务能力)+ KL 下降(停机分布靠向几何先验)+ 步数向 1/p_g 漂移但保留 区分度 —— 三者同时发生才是健康的 PonderNet 训练;如果步数立刻塌到 1 或钉死在 K, 调 p_g / beta / head_bias_init。
# head / head+lora 模式产物: ponder_head.safetensors (+ adapter_model.safetensors)
from safetensors.torch import load_file
model = PonderLlamaForCausalLM.from_ponder(
"/path/to/MiniCPM5-2B-cpu",
ponder_kwargs={"ponder_signal": "learned", "max_ponder_steps": 8})
head = load_file("out/ponder-lora/ponder_head.safetensors")
model.ponder_head.load_state_dict({k.replace("ponder_head.", ""): v
for k, v in head.items()})
# LoRA 适配器: model = PeftModel.from_pretrained(model, "out/ponder-lora")
T4 是 pre-Ampere 卡(sm_75):不支持 bf16,只能 fp16 + GradScaler。本仓库已针对
真实半精度 GPU 预修复并预验证(test_gpu_readiness.py,CPU fp16 复现):
显存预算(head+lora, fp16, batch 2×384):权重 5.0GB + LoRA/头 ≈ 0.1GB + 激活(重算后)≈ 1GB + logits(词表 130560)≈ 1.5GB ≈ 8GB 左右,T4 富余。
bash run_t4.sh # 自检 → 下载权重(约5GB, 磁盘不足自动清理缓存) →
# 数据 → 训练(fp16) → 训前/训后评估 + 思考收益
python eval_ponder.py --model $MODEL_DIR \
--head out/ponder-t4/ponder_head.safetensors --adapter out/ponder-t4 \
--data data/pondernet_train.eval.jsonl --compare --k1-baseline
easy < medium < hard 的思考步数梯度
(数据集本身就是按难度分桶构造的);训前(bias=-1, λ≈0.27)各桶步数几乎相同。--k1-baseline 输出 CE(K=1) − CE(K=8),正值说明"多想"确实
降低了损失;重点看 hard 桶是否比 easy 桶收益更大(思考用在刀刃上)。entropy/msp/drift 停机判据并非模型学出来的元认知,
实测大概率每个位置都会跑满 K 步。要得到"该多则多、该少则少"的停机行为,
用第六节的脚本微调 halting 头(数据量要求很小,head 模式单卡即可),
或在思考块上加 LoRA 联合训练 —— 那才是完整意义的 PonderNet。max_ponder_steps 调小或只保留 prefill 思考。w 与 commit-KV 机制保证了数学上的自洽,但不保证
生成质量提升 —— 建议以 max_ponder_steps=1(≈原模型)为 baseline 对比评估。环境:Python 3.13 / torch 2.11.0+cu128 / transformers 5.16.1 / peft 0.20.0。 600 条难度分级训练数据(easy 220 / medium 200 / hard 180),85 条评估 + 50 条 OOD 压力题。
out/ponder_lora/(adapter_model.safetensors 19MB + ponder_head.safetensors 8.4KB)| 配置 | 全局 CE | easy 步数 | medium 步数 | hard 步数 |
|---|---|---|---|---|
| 训前(bias=-1,λ≈0.27) | 8.821 | 7.23 | 7.20 | 7.34 |
| 训后 K=8 自适应 | 0.018 | 1.46 | 1.30 | 1.20 |
| 训后 K=1(强制单步) | 0.030 | 1.00 | 1.00 | 1.00 |
三个核心观察:
| 题型 | 平均步数 | 停在 1 步占比 |
|---|---|---|
| 数字串复述(10–14 位,纯记忆压力) | 1.50 | 66% |
| 平方(12–99) | 1.34 | 81% |
| 三步混合运算 a×b+c×d | 1.28 | 87% |
| 逻辑链比较 | 1.26 | 89% |
| 大数乘法(3–4 位) | 1.18 | 90% |
| 三步应用题(折扣×数量) | 1.17 | 89% |
已验证:PonderNet 循环思考在真实 2B 模型 + T4 上端到端可行 —— 训练稳定、 自适应停机、质量优于单步 baseline、算力成本可控(~10 分钟可训完)。
未出现:按题目语义难度分桶的步数分化。要让"难题多想"真正显现,需要:
--prior-p 调小(如 0.2)或加大 β,强制步数分布远离 1;# T4 端(Colab)—— 数据 md5: mix.jsonl=41c7ede4, mix.eval.jsonl=5a03ce37(与本地一致)
pip uninstall -y torchao # peft 0.20 与 Colab 预装 torchao 0.10 冲突
python train_ponder_head.py --model models/MiniCPM5-2B \
--train-mode head+lora --lora-scope ponder --data data/mix.jsonl \
--dtype float16 --epochs 2 --batch-size 2 --accum 4 --grad-ckpt \
--output out/ponder_lora
python eval_ponder.py --model models/MiniCPM5-2B --data data/mix.eval.jsonl \
--head out/ponder_lora/ponder_head.safetensors --adapter out/ponder_lora \
--max-steps 8 --generate
注意:新版 torch 的 is_bf16_supported() 对 T4 (sm75) 返回 True 但那是模拟 bf16
(无硬件 tensor core),训练脚本已改为按 compute capability 判断,T4 自动落
fp16 + GradScaler。
9.4 指出核心矛盾:模板化短答案使 hard 桶 per-token 不确定性反而低。破局需要 真实多步 CoT 数据(让"1 步解不出"成为 CE 层面的现实)+ 难度→先验的显式通路。
| 数据集 | 规模 | 语言 | 关键字段 | 适配点 |
|---|---|---|---|---|
meta-math/MetaMathQA_GSM8K_zh | 240k | 中文 | query_zh / response_zh / type | 中文逐步 CoT 解答,替换模板答案的首选 |
swulling/gsm8k_chinese | 7.5k | 中文题+英文解 | answer 内含 <<48/2=24>> 步标注 | 步标注可直接数出推理步数 = 现成难度标签 |
AI-MO/NuminaMath-CoT | 860k | 英文为主 | source(gsm8k/amc_aime/olympiads…) | source 即天然难度梯度,适合跨级对比 |
rabinadk1/competition_math-level_5 等 | MATH 分级子集 | 英文 | level 1–5 | 官方难度分级,可直接当 difficulty 字段 |
make_gsm8k_zh_data.py 从 MetaMathQA_GSM8K_zh
抽样,按"解答中等式行数"自动分桶(≤2 步 easy / 3 步 medium / ≥4 步 hard),
生成带 difficulty 字段的 jsonl。多步 CoT 的中间 token 无法靠 1 步记住,
halting 头将在真实的预测压力下学到"多想"。--difficulty-prior easy:0.7,medium:0.4,hard:0.15。
实现:PonderLoader 把同难度样本组成同 batch,训练循环在每个 batch 前把
config.ponder_prior_p 切换为该难度的 p_g —— KL 正则随之把 easy 拉向快停、
hard 拉向多想。这是 PonderNet 几何先验的 per-sample 扩展,ponder_llama.py 零改动。pip install datasets
python make_gsm8k_zh_data.py --out-dir data_zh --n-train-per-bucket 800 --n-eval-per-bucket 50
python train_ponder_head.py --model <MODEL_DIR> --train-mode head+lora \
--data data_zh/zh_cot.jsonl --difficulty-prior easy:0.7,medium:0.4,hard:0.15 \
--epochs 2 --batch-size 2 --accum 4 --grad-ckpt --dtype float16 --output out/ponder-zh
python eval_ponder.py --model <MODEL_DIR> --data data_zh/zh_cot.eval.jsonl \
--head out/ponder-zh/ponder_head.safetensors --adapter out/ponder-zh \
--max-steps 8 --tag zh-cot
预期:easy 桶步数 →~1,hard 桶步数显著上移(p_g=0.15 的几何先验期望 ≈6.7 步),
出现 easy < medium < hard 的语义分化;评估脚本按 difficulty 分桶可直接检验。
| 项 | 配置 |
|---|---|
| 数据 | meta-math/MetaMathQA_GSM8K_zh 中文逐步 CoT,按等式行数分桶 (easy≤2 / medium=3 / hard≥4 步) |
| 训练集 | 1200 条 (400/桶),实际推理步数梯度 easy 1.84 / medium 3.00 / hard 5.39 |
| 评估集 | zh_cot.eval.jsonl 180 条 (60/桶),评估取前 90 条 |
| 训练 | head+lora (r=16 思考块)、--difficulty-prior easy:0.7,medium:0.4,hard:0.15、max-steps 6、1 epoch (75 优化步)、fp16+GradScaler、T4 |
| 先验生效验证 | 训练日志 KL 恒定在 ~1.0-1.1,三档几何先验随 batch 难度正确切换 |
step 10 | loss 2.3062 (ce 2.2549) | ponder_steps 2.36
step 20 | loss 0.5697 (ce 0.5112) | ponder_steps 1.07
step 70 | loss 0.4326 (ce 0.3843) | ponder_steps 1.07
CE 快速收敛(CoT 任务 LoRA 拟合良好),但 ponder_steps 从 2.36 一路降到 ~1.07 并锁死 —— "尽快停"再次成为 CE+KL 的联合最优解。
| 配置 | 全局 CE | easy 步数 | medium 步数 | hard 步数 |
|---|---|---|---|---|
| K=6 自适应 (训后) | 0.3713 | 1.095 | 1.062 | 1.044 |
| K=1 强制单步 | 0.4223 | 1.000 | 1.000 | 1.000 |
思考收益:CE 0.4223 → 0.3713(降 12%),代价仅平均 0.07 步/token(约 7% 位置用到 2+ 步)。 自适应停机机制本身工作正常且正向获益,但三桶步数不仅未按 hard>medium>easy 分化, 反而呈微弱反向(hard 最低)—— 与第一轮(英文短答案)现象一致。
两轮实验(模板化短答案 + 真实多步 CoT)得到同一结果,可以下结论这不是数据形式问题:
结论:必须改变任务结构让"1 步解不出"成为事实,而不是换数据或调先验。
pondernet/weights_zh_round2/(ponder_head.safetensors + LoRA adapter 19.2MB)pondernet/data_zh_round2/(zh_train_1200.jsonl + zh_cot.eval.jsonl)eval_ponder.py --max-steps 6 --head out_zh/ponder_head.safetensors --adapter out_zh --limit 90用第二轮训后的模型(K=1 单步)对 1200 条训练样本逐样本实测 CE(probe), 以模型自己的能力边界重新定义难度,替代启发式分桶:
probe_ce ≥ p75 (0.555) → 实测 hard(强制多想,p_g=0.05)probe_ce ≤ p25 (0.323) → 实测 easy(强制快停,p_g=0.90)第三轮训练:加载第二轮 LoRA 继续训(--init-adapter out_zh,head 重置),
720 条均衡抽样,1 epoch,先验命中 720/720。新增 --init-adapter 两种用法:
--train-mode head(冻结 LoRA 只训 head)/ --train-mode head+lora(LoRA 续训),
test_boot.py 三阶段冒烟覆盖。
| 启发式桶 | n | probe 平均 CE (K=1) |
|---|---|---|
| easy (1-2 步) | 400 | 0.536 |
| medium (3 步) | 400 | 0.437 |
| hard (≥4 步) | 400 | 0.389 |
推理步数越多的题,模型单步预测反而越容易 —— 数学 CoT 的模板化语言 ("设...为 x"、"因此...=")使长解答的 per-token 熵更低。 这定量解释了两轮实验中 halting 头"反向分化"的原因:它优化的从来就是 预测难度,而不是题目语义难度。交叉表:实测 hard 中 56% 来自启发式 easy。
| 配置 | 全局 CE | easy 步数 | medium 步数 | hard 步数 |
|---|---|---|---|---|
| K=6 自适应 (boot 模型) | 0.3759 | 1.035 | 1.039 | 1.047 |
| K=1 强制单步 | 0.3762 | 1.000 | 1.000 | 1.000 |
pondernet/weights_boot_round3/(head + LoRA)、pondernet/data_boot_round3/
(probe_results.jsonl + boot_train.jsonl + boot_eval.jsonl)python train_ponder_head.py --model <MODEL> --train-mode head+lora \
--init-adapter out_zh --data data_zh/boot_train.jsonl \
--difficulty-prior easy:0.9,medium:0.5,hard:0.05 \
--epochs 1 --batch-size 2 --accum 8 --max-steps 6 --grad-ckpt --dtype float16 \
--output out_boot
python eval_ponder.py --model <MODEL> --data data_zh/boot_eval.jsonl \
--head out_boot/ponder_head.safetensors --adapter out_boot --max-steps 6