samaidev-pnet-2b / train_head_cached.py

round5: 缓存训练: 特征缓存 head-only + 消融①软硬目标

04133a2 verified about 2 months ago

16.9 kB

# -*- coding: utf-8 -*-
"""
train_head_cached.py — 第四轮: 特征缓存 + 步数硬监督 (CPU 友好, HF Space)
核心思想:
  1. 表征冻结 (LoRA merge 后只读): 每步思考隐状态 h_n = block(h_{n-1}) 与 head 无关,
     只需展开一次缓存 h_1..h_K, 之后 head 训练仅重放 lam->w->mix->lm_head 链,
     每样本成本从 "K 次全模型前向+反传" 降到 "1 次 lm_head 投影" (CPU 提速 ~10 倍)
  2. 步数硬监督替代 β·KL (第三轮结论: KL 拉力被 CE 压制):
     实测 easy: -log w_1 (1 步快停)
     实测 hard: -Σ (1/K)·log w_n (均匀铺满, 期望步数 (K+1)/2)
     实测 medium: 不监督
     loss = CE + gamma * sup
数学与 PonderLlamaForCausalLM.forward 训练路径严格等价 (test_cached.py 验证):
  h_n = _run_block(h_{n-1}); lam_n = sigmoid(head(h_n)); lam_eff = lam*active
  w_n = remaining*lam_eff; mix += w_n*h_n; remaining *= (1-lam_eff)
  active: cumsum(w) < 1-eps; 终止后 mix += remaining*h_last (剩余归最后一步)
用法:
  python train_head_cached.py --model models/MiniCPM5-2B --adapter out_zh \
      --data data_zh/probe_results.jsonl --p25 0.323 --p75 0.555 \
      --n-per-bucket 80 --max-len 384 --k-steps 4 --epochs 2 \
      --gamma 1.0 --out out_cached
"""
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 transformers import AutoTokenizer
from transformers.models.llama.modeling_llama import LlamaForCausalLM

from ponder_llama import PonderLlamaForCausalLM, PonderLlamaConfig
from train_ponder_head import load_jsonl, PonderDataset, collate


# -------------------------------------------------------------------------- #
# 展开缓存: 与 forward 训练路径 (W==Q, use_cache=False, 无窗口) 严格一致
# -------------------------------------------------------------------------- #
@torch.no_grad()
def expand_hiddens(model: PonderLlamaForCausalLM, input_ids, attention_mask,
                   k_steps: int) -> list[torch.Tensor]:
    """返回 [h_1..h_K], 每个形状 (B, L, H) —— 思考块每步的未混合隐状态.
    复刻 forward: embed -> 前 start_layer 层单遍 -> 思考块 (后 8 层) 循环 K 次.
    不建 KV cache / 不构建窗口 mask (与训练路径一致), h_n 与 head 参数无关.
    mask 用 create_causal_mask (与 forward 同源, 正确处理右侧 padding)."""
    from transformers.masking_utils import create_causal_mask
    cfg = model.config
    m = model.model
    h = m.embed_tokens(input_ids)
    B, Q = h.shape[:2]
    position_ids = torch.arange(Q, device=h.device).reshape(1, -1).expand(B, -1)
    position_embeddings = m.rotary_emb(h, position_ids=position_ids)
    causal_mask = create_causal_mask(
        config=cfg, inputs_embeds=h, attention_mask=attention_mask,
        past_key_values=None, position_ids=position_ids)
    for layer in m.layers[: cfg.ponder_start_layer]:
        h = layer(h, attention_mask=causal_mask, position_ids=position_ids,
                  position_embeddings=position_embeddings, use_cache=False)
    hiddens = []
    for _ in range(k_steps):
        h = model._run_block(h, causal_mask, position_ids,
                             position_embeddings, None, {})
        hiddens.append(h.clone())
    return hiddens


def replay_head_loss(model, hiddens, labels, eps, gamma, sup_map,
                     easy_target="hard", easy_soft_p=0.7):
    """用缓存 h_n 重放 PonderNet 停机链, 返回 (loss, ce, sup, steps_mean).
    sup_map: (B,) 每样本的监督组别 ("easy"/"hard"/None).
    消融①: easy 目标软硬可选 ——
      hard: t=[1,0,..] (一步压死, round4 默认);
      soft: 截断几何分布 t_n ∝ (1-p)^(n-1)·p (留质量给后几步, p=easy_soft_p).
    hard 目标均匀铺满 K 步. 数学与 forward 循环逐行对应 (含 active/eps/剩余归末步)."""
    cfg = model.config
    device, dtype = hiddens[0].device, hiddens[0].dtype
    B, L = hiddens[0].shape[:2]
    K = len(hiddens)

    remaining = torch.ones(B, L, device=device, dtype=dtype)
    cum = torch.zeros(B, L, device=device, dtype=dtype)
    active = torch.ones(B, L, device=device, dtype=torch.bool)
    mix = torch.zeros(B, L, hiddens[0].size(-1), device=device, dtype=dtype)
    w_list = []
    for k in range(K):
        lam = model._halting_prob(hiddens[k], None)          # (B, L)
        lam_eff = lam * active.to(lam.dtype)
        w_k = remaining * lam_eff
        mix = mix + w_k.unsqueeze(-1) * hiddens[k]
        w_list.append(w_k)
        remaining = remaining * (1.0 - lam_eff)
        cum = cum + w_k
        active = active & (cum < 1.0 - eps)
        if not bool(active.any()):
            break
    mix = mix + remaining.unsqueeze(-1) * hiddens[len(w_list) - 1]
    w_stack = torch.stack(w_list, dim=-1).to(dtype)          # (B, L, K_exec)
    w_stack = torch.cat([w_stack[..., :-1],
                         w_stack[..., -1:] + remaining.unsqueeze(-1)], dim=-1)

    # logits / CE (与 forward 相同: norm 在 mix 之后, shift 对齐)
    mix_normed = model.model.norm(mix)
    logits = model.lm_head(mix_normed.to(model.lm_head.weight.dtype)).float()
    ce_logits = logits.float()
    ce = F.cross_entropy(ce_logits[..., :-1, :].reshape(-1, ce_logits.size(-1)),
                         labels[..., 1:].reshape(-1), ignore_index=-100)

    # 步数硬监督: -Σ target_n · log w_n (w_stack 已含剩余归末步, Σ=1)
    Ke = w_stack.size(-1)
    sup = torch.zeros((), device=device, dtype=torch.float32)
    n_sup = 0
    if gamma > 0 and sup_map:
        logw = torch.log(w_stack.float().clamp_min(1e-8))    # (B, L, Ke)
        # 逐样本构建目标分布 (避免 inplace 写破坏 autograd 版本)
        rows_t = []
        for grp in sup_map:
            t = torch.zeros(Ke, device=device, dtype=torch.float32)
            if grp == "easy":
                if easy_target == "soft":
                    # 消融① 软目标: 截断几何分布, 期望步数 1/p, 留概率质量给后几步
                    n_idx = torch.arange(1, Ke + 1, device=device, dtype=torch.float32)
                    g = (1.0 - easy_soft_p) ** (n_idx - 1) * easy_soft_p
                    t = g / g.sum()
                else:
                    t[0] = 1.0                                   # [1,0,..] 一步压死
            elif grp == "hard":
                t[:] = 1.0 / Ke                              # 均匀 -> 期望步数 (Ke+1)/2
            rows_t.append(t)
        target = torch.stack(rows_t, dim=0)[:, None, :].expand_as(logw)
        mask = (labels[:, 1:] != -100).float()               # 只监督有 label 的位置
        per_pos = -(target * logw).sum(-1)                   # (B, L)
        sup = (per_pos[:, :-1] * mask).sum()
        n_sup = int(mask.sum().item())
        if n_sup > 0:
            sup = sup / n_sup
    loss = ce + gamma * sup

    with torch.no_grad():
        cum_w = torch.cumsum(w_stack, dim=-1)
        steps = ((cum_w < 1.0 - eps).sum(-1) + 1).clamp(max=Ke)
        smean = float(steps[:, :-1][(labels[:, 1:] != -100)].float().mean())
    return loss, ce, sup, smean, w_stack


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--adapter", default=None, help="已训 LoRA 目录 (merge 冻结)")
    ap.add_argument("--data", required=True, help="probe_results.jsonl")
    ap.add_argument("--p25", type=float, required=True)
    ap.add_argument("--p75", type=float, required=True)
    ap.add_argument("--n-per-bucket", type=int, default=80)
    ap.add_argument("--max-len", type=int, default=384)
    ap.add_argument("--k-steps", type=int, default=4)
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--gamma", type=float, default=1.0)
    ap.add_argument("--easy-target", default="hard", choices=["hard", "soft"],
                    help="消融①: easy 目标分布 —— hard=[1,0,..] 一步压死 (round4); "
                         "soft=截断几何分布 p=easy-soft-p, 留质量给后几步")
    ap.add_argument("--easy-soft-p", type=float, default=0.7,
                    help="soft easy 目标的几何 p (期望步数 1/p)")
    ap.add_argument("--head-bias-init", type=float, default=-1.0)
    ap.add_argument("--head-hidden", type=int, default=0,
                    help="方向①: >0 时构造 MLP halting 头 (缓存重放同样适用)")
    ap.add_argument("--lr", type=float, default=2e-3)
    ap.add_argument("--epsilon", type=float, default=0.01)
    ap.add_argument("--beta-kl", type=float, default=0.0,
                    help=">0 时额外加 KL 正则 (默认纯硬监督)")
    ap.add_argument("--prior-p", type=float, default=0.5)
    ap.add_argument("--dtype", default="float32",
                    choices=["float32", "bfloat16", "float16"],
                    help="模型加载精度; 12GB RAM 免费容器用 bfloat16 (4GB), "
                         "head/replay 链路恒为 fp32")
    ap.add_argument("--batch-size", type=int, default=1, help="展开 batch (缓存阶段)")
    ap.add_argument("--out", default="out_cached")
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--limit", type=int, default=None)
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    device = "cpu"
    t00 = time.time()
    print(f"[0] 加载模型 {args.model} (fp32 CPU) ...", flush=True)

    base = PonderLlamaConfig.from_pretrained(args.model)
    base.ponder_signal = "learned"
    base.max_ponder_steps = args.k_steps
    base.ponder_all_positions = True
    base.ponder_epsilon = args.epsilon
    base.ponder_loss_beta = args.beta_kl
    base.ponder_prior_p = args.prior_p
    base.ponder_head_hidden = args.head_hidden
    if args.head_hidden > 0:
        print(f"[0] 方向①: MLP halting 头 hidden={args.head_hidden}", flush=True)
    model = PonderLlamaForCausalLM.from_pretrained(
        args.model, config=base, torch_dtype=getattr(torch, args.dtype),
        low_cpu_mem_usage=True, attn_implementation="sdpa")
    model.ponder_head = model.ponder_head.float()
    _last_lin = [m for m in model.ponder_head.modules()
                 if isinstance(m, torch.nn.Linear)][-1]
    with torch.no_grad():
        _last_lin.bias.fill_(args.head_bias_init)
    if getattr(base, "ponder_head_hidden", 0):
        print(f"[0] MLP halting 头 hidden={base.ponder_head_hidden} "
              f"(参数 {sum(p.numel() for p in model.ponder_head.parameters()):,})", flush=True)
    if args.adapter:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, args.adapter, is_trainable=False)
        model = model.merge_and_unload()      # LoRA 并入权重, 表征固定且前向零开销
        print(f"[0] LoRA 已 merge: {args.adapter}", flush=True)
    model.to(device).eval()

    # ---- 数据 + 实测难度分组 ---------------------------------------------
    rows = load_jsonl(args.data)
    if args.limit:
        rows = rows[: args.limit]
    groups = {"easy": [], "medium": [], "hard": []}
    for r in rows:
        ce = r.get("probe_ce", float("nan"))
        if ce != ce:
            continue
        d = "hard" if ce >= args.p75 else ("easy" if ce <= args.p25 else "medium")
        groups[d].append(r)
    rng = random.Random(args.seed)
    train = []
    for d in ("easy", "medium", "hard"):
        rng.shuffle(groups[d])
        take = groups[d][: args.n_per_bucket]
        train.extend((r, d) for r in take)
    rng.shuffle(train)
    from collections import Counter
    print(f"[1] 训练集 {len(train)} 条 {dict(Counter(d for _, d in train))}", flush=True)

    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([r for r, _ in train], tok, max_len=args.max_len)
    loader = torch.utils.data.DataLoader(
        ds, batch_size=args.batch_size, shuffle=False,
        collate_fn=lambda b: collate(b, tok.pad_token_id or 1))

    # ---- 阶段1: 展开缓存 h_n (batch 并行, 拆样本存储) ---------------------
    print(f"[2] 展开缓存 K={args.k_steps} (batch={args.batch_size}) ...", flush=True)
    cache = []          # list of (hiddens:list[(B,L,H)] fp16, labels, groups:list)
    t0 = time.time()
    done = 0
    for bi, (input_ids, labels, attn) in enumerate(loader):
        hid = expand_hiddens(model, input_ids, attention_mask=attn,
                             k_steps=args.k_steps)
        hid = [h.to(torch.float16) for h in hid]          # 缓存省内存, 训练时升回 fp32
        b0 = done
        gs = [train[j][1] for j in range(b0, b0 + input_ids.size(0))]
        cache.append((hid, labels, gs))
        done += input_ids.size(0)
        if done % 40 == 0:
            el = time.time() - t0
            print(f"  展开进度 {done}/{len(ds)} ({el:.0f}s, {el/done:.2f}s/样本)", flush=True)
    print(f"[2] 缓存完成: {time.time()-t0:.0f}s, 条目 {sum(len(g) for _,_,g in cache)}", flush=True)

    # ---- 阶段2: head 训练 (仅 2049 参数, 重放缓存) -----------------------
    for p in model.parameters():
        p.requires_grad_(False)
    head_params = list(model.ponder_head.parameters())     # 兼容 Linear / MLP 头
    for p in head_params:
        p.requires_grad_(True)
    opt = torch.optim.AdamW(
        [{"params": head_params, "lr": args.lr}], weight_decay=0.01)
    n_updates = math.ceil(len(cache) * args.epochs / 8)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min(1.0, (s + 1) / max(1, int(n_updates * 0.1))) *
        (0.5 * (1 + math.cos(math.pi * min(1.0, s / n_updates)))))
    print(f"[3] head 训练: {len(cache)} 样本 x {args.epochs} epochs, "
          f"gamma={args.gamma}, lr={args.lr}, easy_target={args.easy_target}"
          + (f"(p={args.easy_soft_p})" if args.easy_target == "soft" else ""), flush=True)
    t0 = time.time()
    hist = []
    for ep in range(args.epochs):
        order = list(range(len(cache)))
        rng.shuffle(order)
        te, ce_s, sup_s, st_s, nb = 0.0, 0.0, 0.0, 0.0, 0
        for i in order:
            hidb, labels_b, gs = cache[i]
            B = hidb[0].size(0)
            for b in range(B):
                hids = [h[b].to(torch.float32).unsqueeze(0) for h in hidb]
                loss, ce, sup, smean, _ = replay_head_loss(
                    model, hids, labels_b[b].unsqueeze(0), args.epsilon,
                    args.gamma, [gs[b]], easy_target=args.easy_target,
                    easy_soft_p=args.easy_soft_p)
                opt.zero_grad(set_to_none=True)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(head_params, 1.0)
                opt.step()
                sched.step()
                te += float(loss); ce_s += float(ce); sup_s += float(sup)
                st_s += smean; nb += 1
        hist.append((te / nb, ce_s / nb, sup_s / nb, st_s / nb))
        print(f"  epoch {ep+1}: loss {te/nb:.4f} (ce {ce_s/nb:.4f} + "
              f"{args.gamma}·sup {sup_s/nb:.4f}) | steps {st_s/nb:.3f} | "
              f"{time.time()-t0:.0f}s", flush=True)

    # ---- 阶段3: 步数分布报告 + 保存 --------------------------------------
    print("[4] 训练后步数分布 (全缓存重放):", flush=True)
    for grp in ("easy", "medium", "hard"):
        idx = [i for i, (_, g) in enumerate(train) if g == grp]
        if not idx:
            continue
        ss, nb = 0.0, 0
        bs = args.batch_size
        with torch.no_grad():
            for i in idx:
                hidb, labels_b, _ = cache[i // bs]
                b = i % bs
                hids = [h[b].to(torch.float32).unsqueeze(0) for h in hidb]
                _, _, _, sm, _ = replay_head_loss(
                    model, hids, labels_b[b].unsqueeze(0), args.epsilon, 0.0, [None])
                ss += sm; nb += 1
        print(f"  {grp:>6}: avg_steps {ss/nb:.3f} (n={nb})", flush=True)

    os.makedirs(args.out, exist_ok=True)
    from safetensors.torch import save_file
    head_sd = {k: v.data.to(torch.float32).cpu()
               for k, v in model.ponder_head.state_dict().items()}
    save_file(head_sd,
              os.path.join(args.out, "ponder_head.safetensors"),
              metadata={"note": f"cached-head round5, gamma={args.gamma}, K={args.k_steps}, "
                                 f"easy_target={args.easy_target}, easy_soft_p={args.easy_soft_p}"})
    with open(os.path.join(args.out, "train_meta.json"), "w") as f:
        json.dump({"config": vars(args), "hist": hist,
                   "buckets": {d: len(v) for d, v in groups.items()},
                   "elapsed_s": time.time() - t00}, f, indent=2)
    print(f"[5] 已保存 {args.out}/ponder_head.safetensors + train_meta.json", flush=True)
    print(f"CACHED_TRAIN_DONE 总耗时 {time.time()-t00:.0f}s", flush=True)


if __name__ == "__main__":
    main()

	# -- coding: utf-8 --
	"""
	train_head_cached.py — 第四轮: 特征缓存 + 步数硬监督 (CPU 友好, HF Space)

	核心思想:
	1. 表征冻结 (LoRA merge 后只读): 每步思考隐状态 h_n = block(h_{n-1}) 与 head 无关,
	只需展开一次缓存 h_1..h_K, 之后 head 训练仅重放 lam->w->mix->lm_head 链,
	每样本成本从 "K 次全模型前向+反传" 降到 "1 次 lm_head 投影" (CPU 提速 ~10 倍)
	2. 步数硬监督替代 β·KL (第三轮结论: KL 拉力被 CE 压制):
	实测 easy: -log w_1 (1 步快停)
	实测 hard: -Σ (1/K)·log w_n (均匀铺满, 期望步数 (K+1)/2)
	实测 medium: 不监督
	loss = CE + gamma * sup

	数学与 PonderLlamaForCausalLM.forward 训练路径严格等价 (test_cached.py 验证):
	h_n = _run_block(h_{n-1}); lam_n = sigmoid(head(h_n)); lam_eff = lam*active
	w_n = remaininglam_eff; mix += w_nh_n; remaining *= (1-lam_eff)
	active: cumsum(w) < 1-eps; 终止后 mix += remaining*h_last (剩余归最后一步)

	用法:
	python train_head_cached.py --model models/MiniCPM5-2B --adapter out_zh \
	--data data_zh/probe_results.jsonl --p25 0.323 --p75 0.555 \
	--n-per-bucket 80 --max-len 384 --k-steps 4 --epochs 2 \
	--gamma 1.0 --out out_cached
	"""
	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 transformers import AutoTokenizer
	from transformers.models.llama.modeling_llama import LlamaForCausalLM

	from ponder_llama import PonderLlamaForCausalLM, PonderLlamaConfig
	from train_ponder_head import load_jsonl, PonderDataset, collate


	# -------------------------------------------------------------------------- #
	# 展开缓存: 与 forward 训练路径 (W==Q, use_cache=False, 无窗口) 严格一致
	# -------------------------------------------------------------------------- #
	@torch.no_grad()
	def expand_hiddens(model: PonderLlamaForCausalLM, input_ids, attention_mask,
	k_steps: int) -> list[torch.Tensor]:
	"""返回 [h_1..h_K], 每个形状 (B, L, H) —— 思考块每步的未混合隐状态.

	复刻 forward: embed -> 前 start_layer 层单遍 -> 思考块 (后 8 层) 循环 K 次.
	不建 KV cache / 不构建窗口 mask (与训练路径一致), h_n 与 head 参数无关.
	mask 用 create_causal_mask (与 forward 同源, 正确处理右侧 padding)."""
	from transformers.masking_utils import create_causal_mask
	cfg = model.config
	m = model.model
	h = m.embed_tokens(input_ids)
	B, Q = h.shape[:2]
	position_ids = torch.arange(Q, device=h.device).reshape(1, -1).expand(B, -1)
	position_embeddings = m.rotary_emb(h, position_ids=position_ids)
	causal_mask = create_causal_mask(
	config=cfg, inputs_embeds=h, attention_mask=attention_mask,
	past_key_values=None, position_ids=position_ids)
	for layer in m.layers[: cfg.ponder_start_layer]:
	h = layer(h, attention_mask=causal_mask, position_ids=position_ids,
	position_embeddings=position_embeddings, use_cache=False)
	hiddens = []
	for _ in range(k_steps):
	h = model._run_block(h, causal_mask, position_ids,
	position_embeddings, None, {})
	hiddens.append(h.clone())
	return hiddens


	def replay_head_loss(model, hiddens, labels, eps, gamma, sup_map,
	easy_target="hard", easy_soft_p=0.7):
	"""用缓存 h_n 重放 PonderNet 停机链, 返回 (loss, ce, sup, steps_mean).

	sup_map: (B,) 每样本的监督组别 ("easy"/"hard"/None).
	消融①: easy 目标软硬可选 ——
	hard: t=[1,0,..] (一步压死, round4 默认);
	soft: 截断几何分布 t_n ∝ (1-p)^(n-1)·p (留质量给后几步, p=easy_soft_p).
	hard 目标均匀铺满 K 步. 数学与 forward 循环逐行对应 (含 active/eps/剩余归末步)."""
	cfg = model.config
	device, dtype = hiddens[0].device, hiddens[0].dtype
	B, L = hiddens[0].shape[:2]
	K = len(hiddens)

	remaining = torch.ones(B, L, device=device, dtype=dtype)
	cum = torch.zeros(B, L, device=device, dtype=dtype)
	active = torch.ones(B, L, device=device, dtype=torch.bool)
	mix = torch.zeros(B, L, hiddens[0].size(-1), device=device, dtype=dtype)
	w_list = []
	for k in range(K):
	lam = model._halting_prob(hiddens[k], None) # (B, L)
	lam_eff = lam * active.to(lam.dtype)
	w_k = remaining * lam_eff
	mix = mix + w_k.unsqueeze(-1) * hiddens[k]
	w_list.append(w_k)
	remaining = remaining * (1.0 - lam_eff)
	cum = cum + w_k
	active = active & (cum < 1.0 - eps)
	if not bool(active.any()):
	break
	mix = mix + remaining.unsqueeze(-1) * hiddens[len(w_list) - 1]
	w_stack = torch.stack(w_list, dim=-1).to(dtype) # (B, L, K_exec)
	w_stack = torch.cat([w_stack[..., :-1],
	w_stack[..., -1:] + remaining.unsqueeze(-1)], dim=-1)

	# logits / CE (与 forward 相同: norm 在 mix 之后, shift 对齐)
	mix_normed = model.model.norm(mix)
	logits = model.lm_head(mix_normed.to(model.lm_head.weight.dtype)).float()
	ce_logits = logits.float()
	ce = F.cross_entropy(ce_logits[..., :-1, :].reshape(-1, ce_logits.size(-1)),
	labels[..., 1:].reshape(-1), ignore_index=-100)

	# 步数硬监督: -Σ target_n · log w_n (w_stack 已含剩余归末步, Σ=1)
	Ke = w_stack.size(-1)
	sup = torch.zeros((), device=device, dtype=torch.float32)
	n_sup = 0
	if gamma > 0 and sup_map:
	logw = torch.log(w_stack.float().clamp_min(1e-8)) # (B, L, Ke)
	# 逐样本构建目标分布 (避免 inplace 写破坏 autograd 版本)
	rows_t = []
	for grp in sup_map:
	t = torch.zeros(Ke, device=device, dtype=torch.float32)
	if grp == "easy":
	if easy_target == "soft":
	# 消融① 软目标: 截断几何分布, 期望步数 1/p, 留概率质量给后几步
	n_idx = torch.arange(1, Ke + 1, device=device, dtype=torch.float32)
	g = (1.0 - easy_soft_p) ** (n_idx - 1) * easy_soft_p
	t = g / g.sum()
	else:
	t[0] = 1.0 # [1,0,..] 一步压死
	elif grp == "hard":
	t[:] = 1.0 / Ke # 均匀 -> 期望步数 (Ke+1)/2
	rows_t.append(t)
	target = torch.stack(rows_t, dim=0)[:, None, :].expand_as(logw)
	mask = (labels[:, 1:] != -100).float() # 只监督有 label 的位置
	per_pos = -(target * logw).sum(-1) # (B, L)
	sup = (per_pos[:, :-1] * mask).sum()
	n_sup = int(mask.sum().item())
	if n_sup > 0:
	sup = sup / n_sup
	loss = ce + gamma * sup

	with torch.no_grad():
	cum_w = torch.cumsum(w_stack, dim=-1)
	steps = ((cum_w < 1.0 - eps).sum(-1) + 1).clamp(max=Ke)
	smean = float(steps[:, :-1][(labels[:, 1:] != -100)].float().mean())
	return loss, ce, sup, smean, w_stack


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--model", required=True)
	ap.add_argument("--adapter", default=None, help="已训 LoRA 目录 (merge 冻结)")
	ap.add_argument("--data", required=True, help="probe_results.jsonl")
	ap.add_argument("--p25", type=float, required=True)
	ap.add_argument("--p75", type=float, required=True)
	ap.add_argument("--n-per-bucket", type=int, default=80)
	ap.add_argument("--max-len", type=int, default=384)
	ap.add_argument("--k-steps", type=int, default=4)
	ap.add_argument("--epochs", type=int, default=2)
	ap.add_argument("--gamma", type=float, default=1.0)
	ap.add_argument("--easy-target", default="hard", choices=["hard", "soft"],
	help="消融①: easy 目标分布 —— hard=[1,0,..] 一步压死 (round4); "
	"soft=截断几何分布 p=easy-soft-p, 留质量给后几步")
	ap.add_argument("--easy-soft-p", type=float, default=0.7,
	help="soft easy 目标的几何 p (期望步数 1/p)")
	ap.add_argument("--head-bias-init", type=float, default=-1.0)
	ap.add_argument("--head-hidden", type=int, default=0,
	help="方向①: >0 时构造 MLP halting 头 (缓存重放同样适用)")
	ap.add_argument("--lr", type=float, default=2e-3)
	ap.add_argument("--epsilon", type=float, default=0.01)
	ap.add_argument("--beta-kl", type=float, default=0.0,
	help=">0 时额外加 KL 正则 (默认纯硬监督)")
	ap.add_argument("--prior-p", type=float, default=0.5)
	ap.add_argument("--dtype", default="float32",
	choices=["float32", "bfloat16", "float16"],
	help="模型加载精度; 12GB RAM 免费容器用 bfloat16 (4GB), "
	"head/replay 链路恒为 fp32")
	ap.add_argument("--batch-size", type=int, default=1, help="展开 batch (缓存阶段)")
	ap.add_argument("--out", default="out_cached")
	ap.add_argument("--seed", type=int, default=7)
	ap.add_argument("--limit", type=int, default=None)
	args = ap.parse_args()

	torch.manual_seed(args.seed)
	device = "cpu"
	t00 = time.time()
	print(f"[0] 加载模型 {args.model} (fp32 CPU) ...", flush=True)

	base = PonderLlamaConfig.from_pretrained(args.model)
	base.ponder_signal = "learned"
	base.max_ponder_steps = args.k_steps
	base.ponder_all_positions = True
	base.ponder_epsilon = args.epsilon
	base.ponder_loss_beta = args.beta_kl
	base.ponder_prior_p = args.prior_p
	base.ponder_head_hidden = args.head_hidden
	if args.head_hidden > 0:
	print(f"[0] 方向①: MLP halting 头 hidden={args.head_hidden}", flush=True)
	model = PonderLlamaForCausalLM.from_pretrained(
	args.model, config=base, torch_dtype=getattr(torch, args.dtype),
	low_cpu_mem_usage=True, attn_implementation="sdpa")
	model.ponder_head = model.ponder_head.float()
	_last_lin = [m for m in model.ponder_head.modules()
	if isinstance(m, torch.nn.Linear)][-1]
	with torch.no_grad():
	_last_lin.bias.fill_(args.head_bias_init)
	if getattr(base, "ponder_head_hidden", 0):
	print(f"[0] MLP halting 头 hidden={base.ponder_head_hidden} "
	f"(参数 {sum(p.numel() for p in model.ponder_head.parameters()):,})", flush=True)
	if args.adapter:
	from peft import PeftModel
	model = PeftModel.from_pretrained(model, args.adapter, is_trainable=False)
	model = model.merge_and_unload() # LoRA 并入权重, 表征固定且前向零开销
	print(f"[0] LoRA 已 merge: {args.adapter}", flush=True)
	model.to(device).eval()

	# ---- 数据 + 实测难度分组 ---------------------------------------------
	rows = load_jsonl(args.data)
	if args.limit:
	rows = rows[: args.limit]
	groups = {"easy": [], "medium": [], "hard": []}
	for r in rows:
	ce = r.get("probe_ce", float("nan"))
	if ce != ce:
	continue
	d = "hard" if ce >= args.p75 else ("easy" if ce <= args.p25 else "medium")
	groups[d].append(r)
	rng = random.Random(args.seed)
	train = []
	for d in ("easy", "medium", "hard"):
	rng.shuffle(groups[d])
	take = groups[d][: args.n_per_bucket]
	train.extend((r, d) for r in take)
	rng.shuffle(train)
	from collections import Counter
	print(f"[1] 训练集 {len(train)} 条 {dict(Counter(d for _, d in train))}", flush=True)

	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([r for r, _ in train], tok, max_len=args.max_len)
	loader = torch.utils.data.DataLoader(
	ds, batch_size=args.batch_size, shuffle=False,
	collate_fn=lambda b: collate(b, tok.pad_token_id or 1))

	# ---- 阶段1: 展开缓存 h_n (batch 并行, 拆样本存储) ---------------------
	print(f"[2] 展开缓存 K={args.k_steps} (batch={args.batch_size}) ...", flush=True)
	cache = [] # list of (hiddens:list[(B,L,H)] fp16, labels, groups:list)
	t0 = time.time()
	done = 0
	for bi, (input_ids, labels, attn) in enumerate(loader):
	hid = expand_hiddens(model, input_ids, attention_mask=attn,
	k_steps=args.k_steps)
	hid = [h.to(torch.float16) for h in hid] # 缓存省内存, 训练时升回 fp32
	b0 = done
	gs = [train[j][1] for j in range(b0, b0 + input_ids.size(0))]
	cache.append((hid, labels, gs))
	done += input_ids.size(0)
	if done % 40 == 0:
	el = time.time() - t0
	print(f" 展开进度 {done}/{len(ds)} ({el:.0f}s, {el/done:.2f}s/样本)", flush=True)
	print(f"[2] 缓存完成: {time.time()-t0:.0f}s, 条目 {sum(len(g) for _,_,g in cache)}", flush=True)

	# ---- 阶段2: head 训练 (仅 2049 参数, 重放缓存) -----------------------
	for p in model.parameters():
	p.requires_grad_(False)
	head_params = list(model.ponder_head.parameters()) # 兼容 Linear / MLP 头
	for p in head_params:
	p.requires_grad_(True)
	opt = torch.optim.AdamW(
	[{"params": head_params, "lr": args.lr}], weight_decay=0.01)
	n_updates = math.ceil(len(cache) * args.epochs / 8)
	sched = torch.optim.lr_scheduler.LambdaLR(
	opt, lambda s: min(1.0, (s + 1) / max(1, int(n_updates * 0.1))) *
	(0.5 * (1 + math.cos(math.pi * min(1.0, s / n_updates)))))
	print(f"[3] head 训练: {len(cache)} 样本 x {args.epochs} epochs, "
	f"gamma={args.gamma}, lr={args.lr}, easy_target={args.easy_target}"
	+ (f"(p={args.easy_soft_p})" if args.easy_target == "soft" else ""), flush=True)
	t0 = time.time()
	hist = []
	for ep in range(args.epochs):
	order = list(range(len(cache)))
	rng.shuffle(order)
	te, ce_s, sup_s, st_s, nb = 0.0, 0.0, 0.0, 0.0, 0
	for i in order:
	hidb, labels_b, gs = cache[i]
	B = hidb[0].size(0)
	for b in range(B):
	hids = [h[b].to(torch.float32).unsqueeze(0) for h in hidb]
	loss, ce, sup, smean, _ = replay_head_loss(
	model, hids, labels_b[b].unsqueeze(0), args.epsilon,
	args.gamma, [gs[b]], easy_target=args.easy_target,
	easy_soft_p=args.easy_soft_p)
	opt.zero_grad(set_to_none=True)
	loss.backward()
	torch.nn.utils.clip_grad_norm_(head_params, 1.0)
	opt.step()
	sched.step()
	te += float(loss); ce_s += float(ce); sup_s += float(sup)
	st_s += smean; nb += 1
	hist.append((te / nb, ce_s / nb, sup_s / nb, st_s / nb))
	print(f" epoch {ep+1}: loss {te/nb:.4f} (ce {ce_s/nb:.4f} + "
	f"{args.gamma}·sup {sup_s/nb:.4f}) \| steps {st_s/nb:.3f} \| "
	f"{time.time()-t0:.0f}s", flush=True)

	# ---- 阶段3: 步数分布报告 + 保存 --------------------------------------
	print("[4] 训练后步数分布 (全缓存重放):", flush=True)
	for grp in ("easy", "medium", "hard"):
	idx = [i for i, (_, g) in enumerate(train) if g == grp]
	if not idx:
	continue
	ss, nb = 0.0, 0
	bs = args.batch_size
	with torch.no_grad():
	for i in idx:
	hidb, labels_b, _ = cache[i // bs]
	b = i % bs
	hids = [h[b].to(torch.float32).unsqueeze(0) for h in hidb]
	_, _, _, sm, _ = replay_head_loss(
	model, hids, labels_b[b].unsqueeze(0), args.epsilon, 0.0, [None])
	ss += sm; nb += 1
	print(f" {grp:>6}: avg_steps {ss/nb:.3f} (n={nb})", flush=True)

	os.makedirs(args.out, exist_ok=True)
	from safetensors.torch import save_file
	head_sd = {k: v.data.to(torch.float32).cpu()
	for k, v in model.ponder_head.state_dict().items()}
	save_file(head_sd,
	os.path.join(args.out, "ponder_head.safetensors"),
	metadata={"note": f"cached-head round5, gamma={args.gamma}, K={args.k_steps}, "
	f"easy_target={args.easy_target}, easy_soft_p={args.easy_soft_p}"})
	with open(os.path.join(args.out, "train_meta.json"), "w") as f:
	json.dump({"config": vars(args), "hist": hist,
	"buckets": {d: len(v) for d, v in groups.items()},
	"elapsed_s": time.time() - t00}, f, indent=2)
	print(f"[5] 已保存 {args.out}/ponder_head.safetensors + train_meta.json", flush=True)
	print(f"CACHED_TRAIN_DONE 总耗时 {time.time()-t00:.0f}s", flush=True)


	if __name__ == "__main__":
	main()