# -*- 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()