round5: 全前向训练: head / head+lora / head+block + 方向①②开关

train_ponder_head.py

@@ -0,0 +1,775 @@
+# -*- coding: utf-8 -*-
+"""
+train_ponder_head.py — 微调 PonderLlamaForCausalLM 的 halting 头 (真 PonderNet 训练)
+
+三种训练模式 (--train-mode):
+  head       只训 ponder_head (2049 参数), 主干全部冻结。最便宜, 单卡 8GB 也能跑;
+             但"多想几步"带来的增益有限 (思考块没学过被重复执行)。
+  head+lora  ponder_head + LoRA (--lora-scope ponder 只注入思考块 / all 注入全部层)。
+             推荐: 让思考块学会"被重复执行", 完整 PonderNet 语义, 显存友好的代价。
+  head+block ponder_head + 思考块全参微调。效果上限最高, 显存/算力开销最大。
+
+数据格式 (jsonl, 每行一条):
+  {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
+  或 {"text": "一段完整文本"}          (用 --text-field text)
+可用 --make-demo-data demo.jsonl 生成 16 条中/英混合演示数据 (简单问答 + 推理题,
+让 halting 头有机会学到"难题多想、简单题少想")。
+
+日志中的关键指标是 avg_ponder_steps: 训练初期接近 max_steps, 随 CE+KL 联合作用,
+简单题的步数应逐步下降 —— 这就是模型在学"该想几步"。
+
+示例:
+  python train_ponder_head.py --model /path/to/MiniCPM5-2B-cpu \
+      --train-mode head+lora --lora-scope ponder --data data.jsonl \
+      --epochs 2 --output out/ponder-lora
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import time
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+from transformers import AutoTokenizer, get_cosine_schedule_with_warmup
+
+from ponder_llama import PonderLlamaConfig, PonderLlamaForCausalLM, LlamaConfig
+
+
+# ========================================================================== #
+# 数据
+# ========================================================================== #
+DEMO_DATA = [
+    {"messages": [{"role": "user", "content": "你好"}, {"role": "assistant", "content": "你好！很高兴见到你。"}]},
+    {"messages": [{"role": "user", "content": "1+1等于几？"}, {"role": "assistant", "content": "1+1=2。"}]},
+    {"messages": [{"role": "user", "content": "今天星期几"}, {"role": "assistant", "content": "抱歉，我无法获取实时日期。"}]},
+    {"messages": [{"role": "user", "content": "把'苹果'翻译成英文"}, {"role": "assistant", "content": "苹果的英文是 apple。"}]},
+    {"messages": [{"role": "user", "content": "天空是什么颜色"}, {"role": "assistant", "content": "晴天时天空通常是蓝色的。"}]},
+    {"messages": [{"role": "user", "content": "水在多少度结冰？"}, {"role": "assistant", "content": "标准大气压下，水在 0 摄氏度结冰。"}]},
+    {"messages": [{"role": "user", "content": "9.11 和 9.9 哪个大？"}, {"role": "assistant", "content": "9.9 更大。比较小数：9.11 = 9 + 0.11，9.9 = 9 + 0.90，因为 0.90 > 0.11，所以 9.9 > 9.11。"}]},
+    {"messages": [{"role": "user", "content": "一个水池甲管8小时注满，乙管12小时注满，两管同开几小时注满？"}, {"role": "assistant", "content": "设总量为1。甲每小时注 1/8，乙每小时注 1/12，合计 1/8+1/12 = 3/24+2/24 = 5/24。注满需要 24/5 = 4.8 小时。"}]},
+    {"messages": [{"role": "user", "content": "鸡兔同笼，头36，脚100，各几只？"}, {"role": "assistant", "content": "设鸡 x 只，兔 36-x 只。脚数 2x + 4(36-x) = 100，得 144-2x = 100，x = 22。鸡 22 只，兔 14 只。"}]},
+    {"messages": [{"role": "user", "content": "证明 sqrt(2) 是无理数"}, {"role": "assistant", "content": "反证法：若 sqrt(2)=p/q（最简分数），则 p^2=2q^2，故 p 为偶数，设 p=2k，则 4k^2=2q^2，q^2=2k^2，q 也为偶数，与 p/q 最简矛盾。故 sqrt(2) 无理。"}]},
+    {"messages": [{"role": "user", "content": "写一个Python函数判断质数"}, {"role": "assistant", "content": "def is_prime(n):\n    if n < 2: return False\n    for i in range(2, int(n**0.5)+1):\n        if n % i == 0: return False\n    return True"}]},
+    {"messages": [{"role": "user", "content": "如果所有的A都是B，所有的B都是C，那么所有的A都是C吗？"}, {"role": "assistant", "content": "是的。这是三段论的传递性：A⊆B 且 B⊆C，则 A⊆C。"}]},
+    {"messages": [{"role": "user", "content": "给我讲个笑话"}, {"role": "assistant", "content": "程序员最讨厌的两件事：别人不写注释，和别人让自己写注释。"}]},
+    {"messages": [{"role": "user", "content": "一张纸对折10次后有多少层？"}, {"role": "assistant", "content": "每对折一次层数翻倍，对折 10 次后是 2^10 = 1024 层。"}]},
+    {"messages": [{"role": "user", "content": "用'但是'造句"}, {"role": "assistant", "content": "他很想去看电影，但是作业还没写完。"}]},
+    {"messages": [{"role": "user", "content": "5个人互相握手一次，共握手几次？"}, {"role": "assistant", "content": "这是组合问题：C(5,2) = 5×4/2 = 10 次。"}]},
+]
+
+
+def make_demo_data(path: str):
+    with open(path, "w", encoding="utf-8") as f:
+        for item in DEMO_DATA:
+            f.write(json.dumps(item, ensure_ascii=False) + "\n")
+    print(f"演示数据已写入 {path} ({len(DEMO_DATA)} 条)")
+
+
+# ========================================================================== #
+# 难度分级数据生成 (easy / medium / hard)
+# 设计目的: PonderNet 的核心卖点是"难题多想、简单题少想"。只有当数据存在清晰
+# 的难度梯度时, halting 头才有信号可学 —— 训练后各难度桶的平均思考步数应当分化。
+# 全部题目反向构造 (先选解再出题), 保证答案正确; "difficulty" 字段供 eval 分层统计。
+# ========================================================================== #
+EASY_GREET = [
+    ("你好", "你好！很高兴见到你，有什么可以帮你的吗？"),
+    ("你好呀", "你好呀！今天过得怎么样？"),
+    ("早上好", "早上好！祝你今天心情愉快。"),
+    ("晚安", "晚安，祝你做个好梦。"),
+    ("谢谢", "不客气，能帮到你我很开心。"),
+    ("你是谁", "我是一个 AI 助手，可以回答问题、陪你聊天。"),
+    ("你会什么", "我可以聊天、答题、算数，也可以帮你翻译。"),
+    ("再见", "再见！欢迎随时来找我聊天。"),
+    ("你叫什么名字", "你可以叫我小助手，很高兴认识你。"),
+    ("今天天气不错", "是啊，这样的天气很适合出门走走。"),
+    ("讲个笑话", "程序员最讨厌两件事：别人不写注释，和让自己写注释。"),
+    ("你吃饭了吗", "我是 AI 不用吃饭，不过谢谢你关心！"),
+    ("中午好", "中午好！记得吃午饭休息一下哦。"),
+    ("在吗", "在的，请说。"),
+    ("谢谢你", "不用谢，这是我应该做的。"),
+    ("你会说英文吗", "会的，你可以用英文和我聊天。"),
+    ("新年快乐", "新年快乐！祝万事如意。"),
+    ("生日快乐", "谢谢！也祝你天天开心。"),
+    ("周末愉快", "周末愉快！好好放松一下吧。"),
+    ("下午好", "下午好！喝杯茶休息一下吧。"),
+]
+COMMON_SENSE = [
+    ("天空是什么颜色", "晴天时天空通常是蓝色的。"),
+    ("一年有几个月", "一年有 12 个月。"),
+    ("一个星期有几天", "一个星期有 7 天。"),
+    ("冰遇到热会怎么样", "冰遇到热会融化成水。"),
+    ("太阳从哪边升起", "太阳从东边升起，西边落下。"),
+    ("中国的首都是哪里", "中国的首都是北京。"),
+    ("一年有哪四个季节", "一年有春、夏、秋、冬四个季节。"),
+    ("彩虹有几种颜色", "彩虹有红橙黄绿蓝靛紫 7 种颜色。"),
+    ("水在多少度结冰", "标准大气压下，水在 0 摄氏度结冰。"),
+    ("水在多少度沸腾", "标准大气压下，水在 100 摄氏度沸腾。"),
+    ("世界上最大的海洋是哪个", "世界上最大的海洋是太平洋。"),
+    ("熊猫最爱吃什么", "熊猫最爱吃竹子。"),
+    ("蜜蜂酿什么", "蜜蜂采花蜜酿成蜂蜜。"),
+    ("植物生长需要什么", "植物生长一般需要阳光、水分和空气。"),
+    ("十二生肖排第一的是什么", "十二生肖排第一的是鼠。"),
+    ("中国最长的河流是哪条", "中国最长的河流是长江。"),
+    ("世界上最高的山峰是哪座", "世界上最高的山峰是珠穆朗玛峰。"),
+    ("人的正常体温大约是多少", "人的正常体温大约是 37 摄氏度。"),
+    ("一天有多少个小时", "一天有 24 个小时。"),
+    ("一公里等于多少米", "一公里等于 1000 米。"),
+]
+TRANSLATE = [
+    ("苹果", "apple"), ("书", "book"), ("猫", "cat"), ("狗", "dog"),
+    ("太阳", "sun"), ("月亮", "moon"), ("水", "water"), ("火", "fire"),
+    ("山", "mountain"), ("河", "river"), ("树", "tree"), ("花", "flower"),
+    ("鸟", "bird"), ("鱼", "fish"), ("学校", "school"), ("老师", "teacher"),
+    ("学生", "student"), ("医生", "doctor"), ("医院", "hospital"), ("汽车", "car"),
+    ("火车", "train"), ("飞机", "airplane"), ("手机", "phone"), ("电脑", "computer"),
+    ("蓝色", "blue"), ("红色", "red"), ("绿色", "green"), ("白色", "white"),
+    ("星期一", "Monday"), ("星期天", "Sunday"), ("早上", "morning"), ("晚上", "evening"),
+    ("米饭", "rice"), ("面包", "bread"), ("鸡蛋", "egg"), ("牛奶", "milk"),
+    ("茶", "tea"), ("咖啡", "coffee"), ("朋友", "friend"), ("家庭", "family"),
+]
+ANTONYMS = [
+    ("大", "小"), ("多", "少"), ("高", "矮"), ("长", "短"), ("快", "慢"),
+    ("上", "下"), ("左", "右"), ("前", "后"), ("开", "关"), ("好", "坏"),
+    ("新", "旧"), ("冷", "热"), ("黑", "白"), ("买", "卖"), ("来", "去"),
+    ("远", "近"), ("深", "浅"), ("宽", "窄"), ("厚", "薄"), ("轻", "重"),
+    ("满", "空"), ("强", "弱"), ("甜", "苦"), ("明亮", "黑暗"), ("开始", "结束"),
+    ("���功", "失败"), ("勇敢", "胆怯"), ("勤劳", "懒惰"), ("干净", "肮脏"), ("整齐", "杂乱"),
+]
+# 工程问题预解: (a, b, t) 满足 1/a + 1/b = 1/t, 即 (a-t)(b-t) = t^2
+WORK_PROBLEMS = [(3, 6, 2), (4, 4, 2), (4, 12, 3), (6, 6, 3), (5, 20, 4),
+                 (6, 12, 4), (8, 8, 4), (6, 30, 5), (10, 10, 5), (7, 42, 6),
+                 (8, 24, 6), (9, 18, 6), (10, 15, 6), (12, 12, 6)]
+
+
+def _gen_easy(rng):
+    kind = rng.randrange(4)
+    if kind == 0:
+        q, a = rng.choice(EASY_GREET)
+    elif kind == 1:
+        q, a = rng.choice(COMMON_SENSE)
+    elif kind == 2:
+        zh, en = rng.choice(TRANSLATE)
+        if rng.random() < 0.5:
+            q, a = f"把'{zh}'翻译成英文", f"{zh} 的英文是 {en}。"
+        else:
+            q, a = f"把'{en}'翻译成中文", f"{en} 的中文是 {zh}。"
+    else:
+        op = rng.choice(["+", "-"])
+        x, y = rng.randint(1, 9), rng.randint(1, 9)
+        if op == "-":
+            x, y = max(x, y), min(x, y)
+        res = x + y if op == "+" else x - y
+        q, a = f"{x} {op} {y} 等于几？", f"{x} {op} {y} = {res}。"
+    return q, a
+
+
+def _gen_medium(rng):
+    kind = rng.randrange(5)
+    if kind == 0:  # 两位数乘一位数 / 三位数加减
+        if rng.random() < 0.5:
+            a, b = rng.randint(10, 99), rng.randint(2, 9)
+            q, a_ = f"{a} 乘 {b} 等于多少？", f"{a} × {b} = {a * b}。"
+        else:
+            a, b = rng.randint(100, 999), rng.randint(100, 999)
+            if rng.random() < 0.5:
+                q, a_ = f"{a} 加 {b} 等于多少？", f"{a} + {b} = {a + b}。"
+            else:
+                hi, lo = max(a, b), min(a, b)
+                q, a_ = f"{hi} 减 {lo} 等于多少？", f"{hi} - {lo} = {hi - lo}。"
+    elif kind == 1:  # 单位换算
+        table = [("米", "厘米", 100), ("米", "毫米", 1000), ("千克", "克", 1000),
+                 ("元", "角", 10), ("小时", "分钟", 60), ("分", "秒", 60)]
+        big, small, rate = rng.choice(table)
+        v = rng.randint(2, 99)
+        if rng.random() < 0.5:
+            q, a_ = f"{v} {big} 等于多少 {small}？", f"{v} {big} = {v * rate} {small}。"
+        else:
+            q, a_ = (f"{v * rate} {small} 等于多少 {big}？",
+                     f"{v * rate} {small} = {v} {big}。")
+    elif kind == 2:  # 单价×数量
+        p, n = rng.randint(2, 20), rng.randint(2, 12)
+        q, a_ = (f"一支笔 {p} 元，买 {n} 支一共需要多少钱？",
+                 f"单价 {p} 元，数量 {n} 支，一共 {p} × {n} = {p * n} 元。")
+    elif kind == 3:  # 速度×时间
+        v, t = rng.choice([20, 30, 40, 50, 60, 70, 80]), rng.randint(2, 9)
+        q, a_ = (f"一辆汽车每小时行驶 {v} 千米，行驶 {t} 小时，一共行驶多少千米？",
+                 f"速度 {v} 千米/小时 × 时间 {t} 小时 = {v * t} 千米。")
+    else:  # 反义词
+        w, ant = rng.choice(ANTONYMS)
+        q, a_ = f"'{w}'的反义词是什么？", f"'{w}'的反义词是'{ant}'。"
+    return q, a_
+
+
+def _gen_hard(rng):
+    kind = rng.randrange(6)
+    if kind == 0:  # 三位数×两位数, 答案带分解步骤
+        a, b = rng.randint(100, 999), rng.randint(10, 99)
+        tens, ones = b // 10 * 10, b % 10
+        q = f"{a} 乘 {b} 等于多少？"
+        a_ = (f"{a} × {b} = {a * b}。把 {b} 拆成 {tens} + {ones}："
+              f"{a} × {tens} = {a * tens}，{a} × {ones} = {a * ones}，"
+              f"相加得 {a * tens + a * ones} = {a * b}。")
+    elif kind == 1:  # 鸡兔同笼 (反向构造保证整数解)
+        c, r = rng.randint(2, 20), rng.randint(2, 20)
+        heads, feet = c + r, 2 * c + 4 * r
+        q = f"鸡兔同笼，从上面数有 {heads} 个头，从下面数有 {feet} 只脚，鸡和兔各有几只？"
+        a_ = (f"设鸡 x 只，兔 {heads} - x 只。列方程 2x + 4({heads} - x) = {feet}，"
+              f"化简得 {4 * heads} - 2x = {feet}，x = {c}。"
+              f"所以鸡 {c} 只，兔 {r} 只。")
+    elif kind == 2:  # 相遇问题
+        v1, v2 = rng.choice([20, 30, 40, 50, 60]), rng.choice([20, 30, 40, 50, 60])
+        t = rng.randint(2, 6)
+        d = (v1 + v2) * t
+        q = (f"甲、乙两车分别从相距 {d} 千米的两地同时出发，相向而行。"
+             f"甲车每小时行 {v1} 千米，乙车每小时行 {v2} 千米，几小时后两车相遇？")
+        a_ = (f"两车速度和为 {v1} + {v2} = {v1 + v2} 千米/小时。"
+              f"相遇时间 = 距离 ÷ 速度和 = {d} ÷ {v1 + v2} = {t} 小时。")
+    elif kind == 3:  # 工程问题 (预解表保证答案整洁)
+        a, b, t = rng.choice(WORK_PROBLEMS)
+        q = f"一个水池，甲管单独注满需要 {a} 小时，乙管单独注满需要 {b} 小时，两管同时开需要几小时注满？"
+        a_ = (f"设总量为 1。甲管每小时注 1/{a}，乙管每小时注 1/{b}，"
+              f"两管同开每小时注 1/{a} + 1/{b} = {a + b}/{a * b}。"
+              f"注满需要 {a * b} ÷ {a + b} = {t} 小时。")
+    elif kind == 4:  # 数列推理 (等差/等比/二阶)
+        sub = rng.randrange(3)
+        if sub == 0:
+            a1, d = rng.randint(2, 20), rng.randint(2, 12)
+            seq = [a1 + i * d for i in range(5)]
+            q, a_ = (f"找规律：{', '.join(map(str, seq))}，下一个数是多少？",
+                     f"相邻两项的差恒为 {d}，下一个数是 {seq[-1]} + {d} = {seq[-1] + d}。")
+        elif sub == 1:
+            a1, r = rng.randint(1, 4), rng.choice([2, 3])
+            seq, v = [], a1
+            for _ in range(5):
+                seq.append(v)
+                v *= r
+            q, a_ = (f"找规律：{', '.join(map(str, seq))}，下一个数是多少？",
+                     f"相邻两项的比为 {r}，下一个数是 {seq[-1]} × {r} = {seq[-1] * r}。")
+        else:
+            a1, d0, e = rng.randint(1, 9), rng.randint(1, 5), rng.randint(1, 4)
+            seq, v, dd = [], a1, d0
+            for _ in range(5):
+                seq.append(v)
+                v += dd
+                dd += e
+            q, a_ = (f"找规律：{', '.join(map(str, seq))}，下一个数是多少？",
+                     f"相邻两项的差依次是 {', '.join(str(d0 + i * e) for i in range(4))}，"
+                     f"差递增 {e}。下一个差是 {dd}，下一个数是 {seq[-1]} + {dd} = {seq[-1] + dd}。")
+    else:  # 除法余数 / 浓度
+        if rng.random() < 0.5:
+            divisor = rng.randint(6, 12)
+            r_ = rng.randint(1, divisor - 1)
+            quot = rng.randint(10, 99)
+            n = divisor * quot + r_
+            q = f"{n} 除以 {divisor}，商和余数分别是多少？"
+            a_ = f"{n} ÷ {divisor} = {quot} 余 {r_}，即商是 {quot}，余数是 {r_}。"
+        else:
+            c = rng.choice([10, 20, 25, 50])
+            m = rng.choice([200, 400, 500, 800, 1000])
+            sugar = m * c // 100
+            q = f"{m} 克糖水中含糖 {sugar} 克，这种糖水的浓度是多少？"
+            a_ = f"浓度 = 糖 ÷ 糖水 = {sugar} ÷ {m} = {c}%。"
+    return q, a_
+
+
+def make_mix_data(train_path: str, n_easy=220, n_medium=200, n_hard=180,
+                  n_eval_per_bucket=30, seed=7):
+    """生成难度分级数据集: train_path (训练) + train_path 去掉 .jsonl 加 .eval.jsonl (评估)。
+    每行: {"messages": [...], "difficulty": "easy|medium|hard"}"""
+    rng = random.Random(seed)
+    eval_path = train_path[:-len(".jsonl")] + ".eval.jsonl" if train_path.endswith(".jsonl") \
+        else train_path + ".eval.jsonl"
+    gens = {"easy": _gen_easy, "medium": _gen_medium, "hard": _gen_hard}
+
+    def fill(bucket, n, used, salt):
+        rows, tries = [], 0
+        r2 = random.Random(seed + salt)
+        while len(rows) < n and tries < n * 50:
+            tries += 1
+            q, a = gens[bucket](r2)
+            if q in used:
+                continue
+            used.add(q)
+            rows.append({"messages": [{"role": "user", "content": q},
+                                      {"role": "assistant", "content": a}],
+                         "difficulty": bucket})
+        return rows
+
+    used = set()
+    buckets = ["easy", "medium", "hard"]
+    train_rows = {b: fill(b, n, used, i)
+                  for i, (b, n) in enumerate(zip(buckets, [n_easy, n_medium, n_hard]))}
+    eval_rows = {b: fill(b, n_eval_per_bucket, used, 100) for b in buckets}
+
+    for path, rows in [(train_path, train_rows), (eval_path, eval_rows)]:
+        with open(path, "w", encoding="utf-8") as f:
+            for b in ("easy", "medium", "hard"):
+                for r in rows[b]:
+                    f.write(json.dumps(r, ensure_ascii=False) + "\n")
+        counts = {b: len(rows[b]) for b in rows}
+        print(f"已写入 {path}  {counts}")
+
+
+def load_jsonl(path: str) -> list[dict]:
+    rows = []
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                rows.append(json.loads(line))
+    return rows
+
+
+def _apply_chat_ids(tok, msgs, add_generation_prompt=False) -> list[int]:
+    """apply_chat_template 跨版本归一化: 统一返回 list[int]
+    (transformers 5.x 返回 BatchEncoding, 旧版返回 list[int], fast 分支返回 Encoding)"""
+    out = tok.apply_chat_template(msgs, tokenize=True,
+                                  add_generation_prompt=add_generation_prompt)
+    if hasattr(out, "ids"):                       # tokenizers.Encoding
+        return list(out.ids)
+    if hasattr(out, "input_ids"):                 # BatchEncoding
+        out = out["input_ids"]
+    if isinstance(out, torch.Tensor):
+        out = out.tolist()
+    if out and isinstance(out[0], (list, tuple)):  # 嵌套 -> 单条
+        out = out[0]
+    return [int(t) for t in out]
+
+
+class PonderDataset(Dataset):
+    """jsonl -> token ids。chat 格式走 chat_template 并只对回答部分计 loss;
+    纯文本格式整句计 loss。右侧 padding, labels 在 pad 处置 -100。"""
+
+    def __init__(self, rows, tokenizer, max_len=512, text_field=None):
+        tok = tokenizer
+        self.tok = tokenizer
+        self.max_len = max_len
+        self.examples = []
+        pad_id = tokenizer.pad_token_id or 0
+        for r in rows:
+            if text_field and text_field in r:
+                enc = tokenizer(r[text_field], truncation=True, max_length=max_len)
+                ids, label_mask = enc["input_ids"], [1] * len(enc["input_ids"])
+            else:
+                msgs = r["messages"]
+                # 逐轮拼接: assistant 回答的 token 区间计 loss, 其余位置置 -100
+                ids: list[int] = []
+                label_mask: list[int] = []
+                for i, m in enumerate(msgs):
+                    cur = _apply_chat_ids(tok, msgs[: i + 1])
+                    new = cur[len(ids):] if cur[: len(ids)] == ids else cur
+                    label_mask += ([1] if m["role"] == "assistant" else [0]) * len(new)
+                    ids = cur
+            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)]
+            # difficulty 透传给 PonderLoader (--difficulty-prior 需要)
+            self.examples.append({"input_ids": ids, "labels": labels,
+                                  "difficulty": r.get("difficulty", "unknown")})
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, i):
+        return self.examples[i]
+
+
+def collate(batch, pad_id=1):
+    maxlen = max(len(b["input_ids"]) for b in batch)
+    input_ids, labels, attn = [], [], []
+    for b in batch:
+        n = len(b["input_ids"])
+        pad = maxlen - n
+        input_ids.append(b["input_ids"] + [pad_id] * pad)
+        labels.append(b["labels"] + [-100] * pad)
+        attn.append([1] * n + [0] * pad)
+    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)
+
+
+# ========================================================================== #
+# 模型 / 参数分组
+# ========================================================================== #
+def resolve_dtype(arg, device):
+    """--dtype 解析。auto: cpu→fp32; cuda→原生 bf16 (sm80+) 否则 fp16。
+    注意: 不能用 torch.cuda.is_bf16_supported() 判断 T4 —— 新版 torch 对 sm75
+    返回 True 但那是模拟 (emulated) bf16, 没有硬件 tensor core, 训练极慢。
+    T4 (sm75) → fp16 + GradScaler; 30/40/50 系 (sm80+) → 原生 bf16。"""
+    if arg == "float32":
+        return torch.float32
+    if arg == "float16":
+        return torch.float16
+    if arg == "bfloat16":
+        return torch.bfloat16
+    if device == "cpu":
+        return torch.float32
+    major, _ = torch.cuda.get_device_capability(device)
+    return torch.bfloat16 if major >= 8 and torch.cuda.is_bf16_supported() else torch.float16
+
+
+def build_model(path, args, device, dtype=None):
+    """dtype: torch dtype; None 时按 --dtype 解析 (auto: cuda 无 bf16 支持则 fp16, T4 即此路径)."""
+    dtype = dtype or resolve_dtype(args.dtype, device)
+    print(f"模型精度: {dtype}  (device={device})")
+    cfg = PonderLlamaConfig.from_llama(
+        LlamaConfig.from_pretrained(path),
+        ponder_signal="learned",                       # 训练必须用可学习信号
+        max_ponder_steps=args.max_steps,
+        ponder_epsilon=args.epsilon,
+        ponder_prior_p=args.prior_p,
+        ponder_loss_beta=args.beta,
+        ponder_start_layer=args.start_layer,
+        ponder_all_positions=True,                     # 训练时所有位置独立 halting
+        ponder_head_hidden=getattr(args, "head_hidden", 0),        # 方向①: >0 时 MLP 头
+        ponder_step_ce_weight=getattr(args, "step_ce_weight", 0.0),  # 方向②: per-step CE
+    )
+    model = PonderLlamaForCausalLM.from_pretrained(
+        path, config=cfg, torch_dtype=dtype,
+        low_cpu_mem_usage=True, attn_implementation="sdpa")
+    # halting 头始终用 fp32 保证数值稳定
+    model.ponder_head = model.ponder_head.float()
+    # 初始停机概率: sigmoid(bias)。bias=-1 → λ≈0.27, 期望思考步数 ≈ 3.7
+    # 兼容 Linear (方向① 默认) 与 Sequential/MLP (方向① --head-hidden>0) 两种头
+    last_linear = [m for m in model.ponder_head.modules()
+                   if isinstance(m, torch.nn.Linear)][-1]
+    with torch.no_grad():
+        last_linear.bias.fill_(args.head_bias_init)
+    if getattr(args, "head_hidden", 0) and args.head_hidden > 0:
+        print(f"方向①: halting 头升级 MLP (hidden={args.head_hidden}, "
+              f"参数 {sum(p.numel() for p in model.ponder_head.parameters()):,})")
+    model.to(device)
+    if args.grad_ckpt and device != "cpu":
+        model.gradient_checkpointing_enable()
+    return model
+
+
+def setup_trainables(model, args):
+    """返回 (可训练参数组, 冻结说明). 三种模式对应不同的梯度开放范围."""
+    n = model.config.num_hidden_layers
+    s = model.config.ponder_start_layer
+    # 1) 全部冻结
+    for p in model.parameters():
+        p.requires_grad_(False)
+    groups = []
+    # 2) halting 头 (总是可训, fp32)
+    for p in model.ponder_head.parameters():
+        p.requires_grad_(True)
+        groups.append({"params": [p], "lr": args.lr_head, "name": "head"})
+    frozen_note = ["backbone"]
+    # 3) 思考块 / LoRA
+    if args.train_mode == "head+block":
+        for layer in model.model.layers[s:n]:
+            for p in layer.parameters():
+                p.requires_grad_(True)
+        groups.append({"params": [p for layer in model.model.layers[s:n]
+                                  for p in layer.parameters()],
+                       "lr": args.lr_block, "name": "block"})
+        frozen_note = ["layers 0.." + str(s - 1)]
+    elif args.train_mode == "head+lora":
+        try:
+            from peft import LoraConfig, get_peft_model, PeftModel
+        except ImportError as e:
+            raise SystemExit("head+lora 需要 peft: pip install peft") from e
+        if getattr(args, "init_adapter", None):
+            # 失败自举第三轮: 加载已训 LoRA 继续训 (head 重置为新初始化)
+            model = PeftModel.from_pretrained(model, args.init_adapter,
+                                              is_trainable=True)
+            print(f"已加载并继续训练 LoRA: {args.init_adapter}")
+        else:
+            scope = (list(range(s, n)) if args.lora_scope == "ponder" else list(range(n)))
+            target_names = set()
+            for i in scope:
+                for mod in ("q_proj", "k_proj", "v_proj", "o_proj"):
+                    target_names.add(f"model.layers.{i}.self_attn.{mod}")
+                for mod in ("gate_proj", "up_proj", "down_proj"):
+                    target_names.add(f"model.layers.{i}.mlp.{mod}")
+            lcfg = LoraConfig(r=args.lora_r, lora_alpha=args.lora_r * 2, lora_dropout=0.0,
+                              bias="none", task_type="CAUSAL_LM",
+                              target_modules=list(target_names))
+            model = get_peft_model(model, lcfg)
+        # 重新分组 (peft 包装后参数引用变化)
+        groups = [g for g in groups if g["name"] == "head"]
+        model.base_model.model.ponder_head.requires_grad_(True)  # 保险
+        lora_params = [p for pn, p in model.named_parameters()
+                       if p.requires_grad and "lora_" in pn]
+        groups.append({"params": lora_params, "lr": args.lr_lora, "name": "lora"})
+        frozen_note = [f"backbone (+LoRA r={args.lora_r} on {args.lora_scope} layers)"]
+    elif args.train_mode == "head" and getattr(args, "init_adapter", None):
+        # 两阶段 (失败自举第二阶段): 加载已训 LoRA 并冻结, 只重训 halting 头 ——
+        # 表征固定后, head 只能在"1 步不够"的位置上学习停机 (排除 LoRA 同步快停的混淆)
+        try:
+            from peft import PeftModel
+        except ImportError as e:
+            raise SystemExit("--init-adapter 需要 peft: pip install peft") from e
+        model = PeftModel.from_pretrained(model, args.init_adapter)
+        for pn_, p_ in model.named_parameters():
+            if "lora_" in pn_:
+                p_.requires_grad_(False)
+        # head 参数引用在 peft 包装后不变 (同一张量), groups 中的引用仍有效;
+        # peft from_pretrained 默认全部 requires_grad=False, 显式重开 head
+        for p in model.ponder_head.parameters():
+            p.requires_grad_(True)
+        frozen_note = [f"backbone + frozen LoRA from {args.init_adapter}"]
+    return model, groups, frozen_note
+
+
+# ========================================================================== #
+# 训练循环
+# ========================================================================== #
+def train_one_epoch(model, loader, optimizer, scheduler, device,
+                    accum=1, clip=1.0, log_every=0, global_step=0, scaler=None,
+                    prior_map=None, step_ce_weight=0.0):
+    model.train()
+    stats = {"loss": 0.0, "ce": 0.0, "kl": 0.0, "sce": 0.0, "steps": 0.0, "nb": 0}
+    optimizer.zero_grad(set_to_none=True)
+    t0 = time.time()
+    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)
+        # 注意: forward 的 loss = CE + beta*KL (+ step_ce_weight*per-step CE)
+        loss_to_back = out.loss / accum
+        if scaler is not None:
+            scaler.scale(loss_to_back).backward()   # fp16 梯度缩放防下溢
+        else:
+            loss_to_back.backward()
+        with torch.no_grad():
+            # 思考步数只统计有监督的位置
+            mask = labels[:, 1:] != -100
+            steps = out.ponder_steps[:, :-1][mask].float().mean() if mask.any() \
+                else out.ponder_steps.float().mean()
+            sce = out.ponder_step_ce.item() if out.ponder_step_ce is not None else 0.0
+            stats["loss"] += out.loss.item()
+            stats["ce"] += (out.loss.item()
+                            - (model.config.ponder_loss_beta * out.ponder_kl.item())
+                            - step_ce_weight * sce)
+            stats["kl"] += out.ponder_kl.item()
+            stats["sce"] += sce
+            stats["steps"] += steps.item()
+            stats["nb"] += 1
+        if (it + 1) % accum == 0 or (it + 1) == len(loader):
+            if scaler is not None:
+                scaler.unscale_(optimizer)          # clip 必须在真实梯度上做
+            torch.nn.utils.clip_grad_norm_(
+                [p for g in optimizer.param_groups for p in g["params"]], clip)
+            if scaler is not None:
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad(set_to_none=True)
+            global_step += 1
+            if log_every and global_step % log_every == 0:
+                k = max(stats["nb"], 1)
+                sce_part = (f" | stepCE {stats['sce']/k:6.4f}"
+                            if step_ce_weight > 0 else "")
+                print(f"  step {global_step:4d} | loss {stats['loss']/k:7.4f} "
+                      f"(ce {stats['ce']/k:7.4f} + {model.config.ponder_loss_beta}·kl {stats['kl']/k:6.4f}"
+                      f"{sce_part})"
+                      f" | ponder_steps {stats['steps']/k:5.2f} | lr {scheduler.get_last_lr()[0]:.2e}"
+                      f" | {time.time()-t0:6.1f}s", flush=True)
+                stats = {"loss": 0.0, "ce": 0.0, "kl": 0.0, "sce": 0.0, "steps": 0.0, "nb": 0}
+                t0 = time.time()
+    return global_step
+
+
+def save_model(model, args, out_dir):
+    os.makedirs(out_dir, exist_ok=True)
+    core = model.get_base_model() if hasattr(model, "get_base_model") else model
+    # ponder_head 总是单独保存 (peft 适配器不包含它)。
+    # 键名 = state_dict 去掉 ponder_head. 前缀: Linear 头 → weight/bias (向后兼容);
+    # MLP 头 (方向①) → 0.weight/0.bias/2.weight/2.bias
+    from safetensors.torch import save_file
+    head_sd = {k: v.data.to(torch.float32).cpu()
+               for k, v in core.ponder_head.state_dict().items()}
+    save_file(head_sd, os.path.join(out_dir, "ponder_head.safetensors"),
+              metadata={"note": f"signal=learned, K={args.max_steps}, p_g={args.prior_p}, "
+                                 f"head_hidden={getattr(args, 'head_hidden', 0)}"})
+    if args.train_mode == "head+lora":
+        model.save_pretrained(out_dir)          # adapter_model.safetensors
+        print(f"已保存 LoRA 适配器 + ponder_head 到 {out_dir}")
+    elif args.train_mode == "head+block":
+        core.save_pretrained(out_dir)
+        print(f"已保存全量模型 (含新 ponder_head) 到 {out_dir}")
+    else:
+        print(f"已保存 ponder_head.safetensors 到 {out_dir} "
+              f"(推理时 from_ponder 加载后 model.ponder_head.load_state_dict)")
+
+
+# ========================================================================== #
+# main
+# ========================================================================== #
+def main():
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[1])
+    ap.add_argument("--model", default=None,
+                    help="模型路径; 仅生成数据 (--make-demo-data/--make-mix-data) 时可省略")
+    ap.add_argument("--data", default=None, help="jsonl 训练数据; 不给且无 --make-* 时用内置演示数据")
+    ap.add_argument("--make-demo-data", default=None, metavar="PATH")
+    ap.add_argument("--make-mix-data", default=None, metavar="PATH",
+                    help="生成难度分级数据集 (easy/medium/hard + .eval.jsonl 评估集)")
+    ap.add_argument("--n-easy", type=int, default=220)
+    ap.add_argument("--n-medium", type=int, default=200)
+    ap.add_argument("--n-hard", type=int, default=180)
+    ap.add_argument("--n-eval-per-bucket", type=int, default=30)
+    ap.add_argument("--seed", type=int, default=7)
+    ap.add_argument("--text-field", default=None, help="纯文本字段名 (默认按 chat messages 解析)")
+    ap.add_argument("--train-mode", default="head", choices=["head", "head+lora", "head+block"])
+    ap.add_argument("--output", default="out/ponder")
+    ap.add_argument("--epochs", type=float, default=2.0)
+    ap.add_argument("--batch-size", type=int, default=2)
+    ap.add_argument("--accum", type=int, default=4)
+    ap.add_argument("--max-len", type=int, default=512)
+    ap.add_argument("--lr-head", type=float, default=1e-3)
+    ap.add_argument("--lr-lora", type=float, default=2e-4)
+    ap.add_argument("--lr-block", type=float, default=1e-5)
+    ap.add_argument("--warmup-ratio", type=float, default=0.03)
+    ap.add_argument("--clip", type=float, default=1.0)
+    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("--head-hidden", type=int, default=0,
+                    help="方向①: >0 时 halting 头升级 MLP (Linear->GELU->Linear), "
+                         "打破线性可分上限 (round4: 线性头在冻结特征上已到极限)")
+    ap.add_argument("--step-ce-weight", type=float, default=0.0,
+                    help="方向②: per-step CE 权重 —— 每步思考表示单独算解码 CE 按 w 加权, "
+                         "让后几步表示'可用', 回收 mix 摊薄代价 (需 head+lora/head+block 模式)")
+    ap.add_argument("--start-layer", type=int, default=None, help="思考块起始层 (默认最后8层)")
+    ap.add_argument("--lora-r", type=int, default=16)
+    ap.add_argument("--lora-scope", default="ponder", choices=["ponder", "all"])
+    ap.add_argument("--init-adapter", default=None,
+                    help="两阶段训练: 加载已训 LoRA 适配器目录并冻结 (配 --train-mode head), "
+                         "只重训 halting 头 (失败自举第二阶段)")
+    ap.add_argument("--dtype", default="auto", choices=["auto", "float16", "bfloat16", "float32"],
+                    help="auto: T4 等 pre-Ampere 卡自动 fp16, 30/40/50 系 bf16, cpu fp32")
+    ap.add_argument("--grad-ckpt", action="store_true", default=None,
+                    help="思考块梯度重算 (16GB 显存训练强烈建议开启)")
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    ap.add_argument("--log-every", type=int, default=10)
+    args = ap.parse_args()
+
+    if args.make_mix_data:
+        make_mix_data(args.make_mix_data, n_easy=args.n_easy, n_medium=args.n_medium,
+                      n_hard=args.n_hard, n_eval_per_bucket=args.n_eval_per_bucket,
+                      seed=args.seed)
+        if not args.data:
+            return
+    if args.make_demo_data:
+        make_demo_data(args.make_demo_data)
+        if not args.data:
+            return
+    if not args.model:
+        ap.error("--model 为必填 (仅在仅使用 --make-demo-data / --make-mix-data 时可省略)")
+
+    # ---- 数据 ----
+    rows = load_jsonl(args.data) if args.data else DEMO_DATA
+    print(f"训练样本: {len(rows)} 条 (mode={args.train_mode})")
+    tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+    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)
+    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))
+
+    # ---- 模型 ----
+    if args.step_ce_weight > 0 and args.train_mode == "head":
+        print("[警告] --step-ce-weight>0 但 train-mode=head: 头不改变表示且权重已 detach, "
+              "per-step CE 将无梯度作用对象; 建议搭配 --train-mode head+lora / head+block")
+    dtype = resolve_dtype(args.dtype, args.device)
+    model = build_model(args.model, args, args.device, dtype=dtype)
+    model, groups, frozen_note = setup_trainables(model, args)
+    n_train = sum(p.numel() for g in groups for p in g["params"])
+    print(f"可训练参数: {n_train:,} ({[g['name'] for g in groups]}), 冻结: {frozen_note}")
+
+    # ---- 优化器 ----
+    optimizer = torch.optim.AdamW(
+        [{"params": g["params"], "lr": g["lr"]} for g in groups], weight_decay=0.01)
+    total_updates = max(1, int(len(loader) / args.accum * args.epochs))
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer, int(total_updates * args.warmup_ratio), total_updates)
+    # fp16 训练需要 GradScaler 防梯度下溢 (T4); bf16/fp32 不需要
+    scaler = None
+    if args.device.startswith("cuda") and dtype == torch.float16:
+        scaler = torch.amp.GradScaler(args.device)
+        print("已启用 GradScaler (fp16)")
+
+    # ---- 训练 ----
+    gs = 0
+    for ep in range(math.ceil(args.epochs)):
+        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,
+                             prior_map=prior_map, step_ce_weight=args.step_ce_weight)
+
+    save_model(model, args, args.output)
+    print("完成。推理加载示例:")
+    print("  model = PonderLlamaForCausalLM.from_ponder('<模型路径>',")
+    print(f"         ponder_kwargs={{'ponder_signal':'learned','max_ponder_steps':{args.max_steps}}})")
+    print("  head = load_file('<output>/ponder_head.safetensors')")
+    print("  model.ponder_head.load_state_dict(head)")
+
+
+if __name__ == "__main__":
+    main()