pondernet: add difficulty-conditioned prior (--difficulty-prior, PonderLoader) + make_gsm8k_zh_data.py (HF dataset -> difficulty-bucketed CoT jsonl) + README section 10

pondernet/README.md

@@ -344,3 +344,49 @@ python eval_ponder.py --model models/MiniCPM5-2B --data data/mix.eval.jsonl \
 注意：新版 torch 的 `is_bf16_supported()` 对 T4 (sm75) 返回 True 但那是模拟 bf16
 （无硬件 tensor core），训练脚本已改为按 compute capability 判断，T4 自动落
 fp16 + GradScaler。
+
+## 十、"难题多想"语义分化专项：数据集选型与方案（2026-09-08）
+
+9.4 指出核心矛盾：模板化短答案使 hard 桶 per-token 不确定性反而低。破局需要
+**真实多步 CoT 数据**（让"1 步解不出"成为 CE 层面的现实）+ **难度→先验的显式通路**。
+
+### 10.1 HF 数据集选型（实查字段与规模）
+
+| 数据集 | 规模 | 语言 | 关键字段 | 适配点 |
+|---|---|---|---|---|
+| `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 字段 |
+
+### 10.2 专门方案（按改动成本排序，1+2 组合拳已落地）
+
+1. **CoT 数据替换（数据侧，首选）**：`make_gsm8k_zh_data.py` 从 MetaMathQA_GSM8K_zh
+   抽样，按"解答中等式行数"自动分桶（≤2 步 easy / 3 步 medium / ≥4 步 hard），
+   生成带 difficulty 字段的 jsonl。多步 CoT 的中间 token 无法靠 1 步记住，
+   halting 头将在真实的预测压力下学到"多想"。
+2. **难度条件化先验（损失侧，已实现）**：`--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 零改动。
+3. **两阶段训练（调度侧）**：先 LoRA 拟合 CoT 任务（CE 降下来），再冻结 LoRA
+   单独训 halting 头 —— 避免"快停"在训练早期成为 CE+KL 的联合最优解。
+4. **失败自举（AdaptThink 式，进阶）**：先用 K=1 跑一遍训练集，答错的样本构成
+   "必须多想"集合，对其施加更小的 p_g（或直接监督步数目标）—— 用模型自己的
+   失败案例定义难度，比启发式分桶更精准。
+
+### 10.3 一键复现（方案 1+2）
+
+```bash
+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 分桶可直接检验。

pondernet/make_gsm8k_zh_data.py

@@ -0,0 +1,117 @@
+# -*- coding: utf-8 -*-
+"""
+make_gsm8k_zh_data.py — 从 HF 拉取中文 CoT 数学数据, 按推理步数分难度桶
+
+解决"T4 实测中 hard 桶答案模板化 → 步数分化不出现"的问题:
+用真实多步推理数据 (MetaMathQA_GSM8K_zh, 响应为中文逐步 CoT) 替换模板化短答案,
+使"1 步解不出"成为 CE 层面的现实; 难度桶按解答步数自动划分 (反向构造可靠性高)。
+
+数据源 (HF):
+  meta-math/MetaMathQA_GSM8K_zh — 字段 query_zh / response_zh / type, 240k 条
+分桶规则: response_zh 中"计算行"数 (含 = 的行) ≈ 推理步数
+  easy:   1-2 步    medium: 3 步    hard: >=4 步
+
+用法 (T4/Colab 或任意有网环境):
+  pip install datasets
+  python make_gsm8k_zh_data.py --out-dir data_zh --n-train-per-bucket 800 \
+      --n-eval-per-bucket 50 --max-per-bucket 4000
+输出: data_zh/zh_cot.jsonl + data_zh/zh_cot.eval.jsonl
+每行: {"messages": [...], "difficulty": "easy|medium|hard", "n_steps": int}
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+
+
+def count_reasoning_steps(response: str) -> int:
+    """解答步数 ≈ 含等号的行数 (CoT 中每个计算步至少一个 =)"""
+    lines = [ln for ln in response.replace("。", "\n").split("\n")
+             if "=" in ln and any(ch.isdigit() for ch in ln)]
+    return max(1, len(lines))
+
+
+def bucket_of(n_steps: int) -> str:
+    if n_steps <= 2:
+        return "easy"
+    if n_steps == 3:
+        return "medium"
+    return "hard"
+
+
+def to_messages(query: str, response: str) -> list[dict]:
+    # 规范化答案结尾: 确保最后一行有"答案是 X"格式 (MetaMathQA 自带, 保持原样)
+    return [{"role": "user", "content": query.strip()},
+            {"role": "assistant", "content": response.strip()}]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out-dir", default="data_zh")
+    ap.add_argument("--n-train-per-bucket", type=int, default=800)
+    ap.add_argument("--n-eval-per-bucket", type=int, default=50)
+    ap.add_argument("--max-per-bucket", type=int, default=4000,
+                    help="每桶最多扫描收集的候选数 (控制内存)")
+    ap.add_argument("--max-chars", type=int, default=1200,
+                    help="response_zh 最大字符数 (控制序列长度/T4 显存)")
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--dataset", default="meta-math/MetaMathQA_GSM8K_zh")
+    args = ap.parse_args()
+
+    from datasets import load_dataset
+
+    rng = random.Random(args.seed)
+    print(f"加载 {args.dataset} (streaming) ...")
+    ds = load_dataset(args.dataset, split="train", streaming=True)
+
+    buckets = {"easy": [], "medium": [], "hard": []}
+    cap = args.max_per_bucket
+    scanned = 0
+    for ex in ds:
+        scanned += 1
+        if scanned % 20000 == 0:
+            print(f"  已扫描 {scanned} 条, 桶计数: "
+                  f"{ {k: len(v) for k, v in buckets.items()} }")
+        q, r = (ex.get("query_zh") or "").strip(), (ex.get("response_zh") or "").strip()
+        if not q or not r or len(r) > args.max_chars or len(q) > 400:
+            continue
+        if "答案是" not in r and "答案：" not in r and "####" not in r:
+            continue
+        n = count_reasoning_steps(r)
+        b = bucket_of(n)
+        if len(buckets[b]) >= cap:
+            continue
+        buckets[b].append({"messages": to_messages(q, r),
+                           "difficulty": b, "n_steps": n})
+        if all(len(v) >= cap for v in buckets.values()):
+            break
+    print(f"扫描完成 ({scanned} 条), 各桶候选: "
+          f"{ {k: len(v) for k, v in buckets.items()} }")
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    train_path = os.path.join(args.out_dir, "zh_cot.jsonl")
+    eval_path = os.path.join(args.out_dir, "zh_cot.eval.jsonl")
+    for path, n_take in [(train_path, args.n_train_per_bucket),
+                         (eval_path, args.n_eval_per_bucket)]:
+        out = []
+        for b, pool in buckets.items():
+            pool = pool[:]
+            rng.shuffle(pool)
+            out.extend(pool[:n_take])
+        rng.shuffle(out)
+        with open(path, "w", encoding="utf-8") as f:
+            for r in out:
+                f.write(json.dumps(r, ensure_ascii=False) + "\n")
+        from collections import Counter
+        print(f"写入 {path}: {len(out)} 条 {dict(Counter(r['difficulty'] for r in out))}")
+
+    print("完成。训练建议 (难度条件化先验, 让 hard 学会多想):")
+    print("  python train_ponder_head.py --model <MODEL> --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")
+
+
+if __name__ == "__main__":
+    main()

pondernet/train_ponder_head.py

@@ -369,7 +369,9 @@ class PonderDataset(Dataset):
             ids = ids[:max_len]
             label_mask = label_mask[: len(ids)]
             labels = [t if mk == 1 else -100 for t, mk in zip(ids, label_mask)]
-            self.examples.append({"input_ids": ids, "labels": labels})
+            # difficulty 透传给 PonderLoader (--difficulty-prior 需要)
+            self.examples.append({"input_ids": ids, "labels": labels,
+                                  "difficulty": r.get("difficulty", "unknown")})
 
     def __len__(self):
         return len(self.examples)
@@ -390,6 +392,37 @@ def collate(batch, pad_id=1):
     return (torch.tensor(input_ids), torch.tensor(labels), torch.tensor(attn))
 
 
+class PonderLoader:
+    """难度感知 DataLoader: 同难度样本组成同 batch, 迭代产出
+    (input_ids, labels, attn, difficulty)。配合 --difficulty-prior 使用 ——
+    同质 batch 使得按难度动态设置 config.ponder_prior_p 成为可能
+    (KL 先验把 easy 拉向快停 / hard 拉向多想, 语义分化由此而来)。
+    不用 --difficulty-prior 时走普通 DataLoader, 行为不变。"""
+
+    def __init__(self, ds, batch_size, pad_id, seed=7):
+        self.ds, self.bs, self.pad_id, self.seed = ds, batch_size, pad_id, seed
+        self.epoch = 0
+        self.by_bucket: dict[str, list[int]] = {}
+        for i, ex in enumerate(ds.examples):
+            self.by_bucket.setdefault(ex.get("difficulty", "unknown"), []).append(i)
+
+    def __len__(self):
+        return sum((len(v) + self.bs - 1) // self.bs for v in self.by_bucket.values())
+
+    def __iter__(self):
+        rng = random.Random(self.seed + self.epoch)
+        self.epoch += 1
+        batches = []
+        for d, idxs in self.by_bucket.items():
+            idxs = idxs[:]
+            rng.shuffle(idxs)
+            batches += [(d, idxs[i:i + self.bs])
+                        for i in range(0, len(idxs), self.bs)]
+        rng.shuffle(batches)          # batch 间顺序打乱 (batch 内同质)
+        for d, idxs in batches:
+            yield (*collate([self.ds[i] for i in idxs], self.pad_id), d)
+
+
 # ========================================================================== #
 # 模型 / 参数分组
 # ========================================================================== #
@@ -491,12 +524,20 @@ def setup_trainables(model, args):
 # 训练循环
 # ========================================================================== #
 def train_one_epoch(model, loader, optimizer, scheduler, device,
-                    accum=1, clip=1.0, log_every=0, global_step=0, scaler=None):
+                    accum=1, clip=1.0, log_every=0, global_step=0, scaler=None,
+                    prior_map=None):
     model.train()
     stats = {"loss": 0.0, "ce": 0.0, "kl": 0.0, "steps": 0.0, "nb": 0}
     optimizer.zero_grad(set_to_none=True)
     t0 = time.time()
-    for it, (input_ids, labels, attn) in enumerate(loader):
+    for it, item in enumerate(loader):
+        if len(item) == 4:                       # PonderLoader: 带 difficulty
+            input_ids, labels, attn, diff = item
+            pg = (prior_map or {}).get(diff)
+            if pg is not None:                   # 难度条件化先验: 同 batch 生效
+                model.config.ponder_prior_p = pg
+        else:
+            input_ids, labels, attn = item
         input_ids, labels, attn = input_ids.to(device), labels.to(device), attn.to(device)
         out = model(input_ids=input_ids, attention_mask=attn, labels=labels,
                     output_ponder=True)
@@ -592,6 +633,10 @@ def main():
     ap.add_argument("--max-steps", type=int, default=8, help="思考步数上限 K")
     ap.add_argument("--epsilon", type=float, default=0.01)
     ap.add_argument("--prior-p", type=float, default=0.4, help="几何先验 p_g (越小越鼓励多想)")
+    ap.add_argument("--difficulty-prior", default=None,
+                    help="难度条件化先验, 如 easy:0.7,medium:0.4,hard:0.15 —— "
+                         "同难度样本组成同 batch, KL 先验按难度调整 "
+                         "(hard 用小 p_g 鼓励多想; 需数据带 difficulty 字段)")
     ap.add_argument("--beta", type=float, default=0.05, help="KL 正则权重")
     ap.add_argument("--head-bias-init", type=float, default=-1.0, help="初始停机偏置, sigmoid(-1)≈0.27")
     ap.add_argument("--start-layer", type=int, default=None, help="思考块起始层 (默认最后8层)")
@@ -625,8 +670,20 @@ def main():
     if tok.pad_token_id is None:
         tok.pad_token = tok.eos_token
     ds = PonderDataset(rows, tok, max_len=args.max_len, text_field=args.text_field)
-    loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True,
-                        collate_fn=lambda b: collate(b, tok.pad_token_id or 1))
+    prior_map = None
+    if args.difficulty_prior:
+        prior_map = {}
+        for part in args.difficulty_prior.split(","):
+            k, v = part.split(":")
+            prior_map[k.strip()] = float(v)
+        n_diff = sum(1 for ex in ds.examples if ex.get("difficulty") in prior_map)
+        print(f"难度条件化先验: {prior_map} | 命中样本 {n_diff}/{len(ds.examples)}")
+        if n_diff == 0:
+            raise SystemExit("--difficulty-prior 需要数据带 difficulty 字段 (如 make_gsm8k_zh_data.py / --make-mix-data 的产物)")
+        loader = PonderLoader(ds, args.batch_size, tok.pad_token_id or 1, args.seed)
+    else:
+        loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True,
+                            collate_fn=lambda b: collate(b, tok.pad_token_id or 1))
 
     # ---- 模型 ----
     dtype = resolve_dtype(args.dtype, args.device)
@@ -653,7 +710,8 @@ def main():
         print(f"--- epoch {ep + 1}/{math.ceil(args.epochs)} ---")
         gs = train_one_epoch(model, loader, optimizer, scheduler, args.device,
                              accum=args.accum, clip=args.clip,
-                             log_every=args.log_every, global_step=gs, scaler=scaler)
+                             log_every=args.log_every, global_step=gs, scaler=scaler,
+                             prior_map=prior_map)
 
     save_model(model, args, args.output)
     print("完成。推理加载示例:")