backup s27_calibrate.py

sft-adapter/s27_calibrate.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+s27_calibrate.py — samai-27b SFT 后校准 (Kaggle 2xT4, nf4 + adapter)
+
+产物 /tmp/k27/s27_calibration.json:
+  expert_freq: {"35": [f0..f7], ...}   8 个 ponder 层各专家被激活的 token 频率 (GGUF 合并权重)
+  halting_w:   [w_1..w_K]              停机质量分布 (记录用; flat GGUF 不使用)
+  steps_mean / steps_hist / n_prompts
+
+用法:
+  python3 s27_calibrate.py
+"""
+import argparse, json, os, time
+import numpy as np
+import torch
+
+os.environ.setdefault("HF_HOME", "/tmp/k27/hf_home")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+T0 = time.time()
+
+
+def log(m):
+    print(f"[{time.time()-T0:6.1f}s] {m}", flush=True)
+
+
+PROMPTS = [
+    ("chat", "你好，请用一句话介绍你自己。"),
+    ("chat", "帮我想一个周末去公园野餐的计划，两三句就行。"),
+    ("chat", "用一句话解释什么是复利。"),
+    ("chat", "给我讲一个冷笑话。"),
+    ("arith", "37 + 58 = ?"),
+    ("arith", "12 × 17 = ?"),
+    ("arith", "小明有 45 个苹果，分给 5 个朋友每人相同数量，每人几个？"),
+    ("rent", "一个出租院落面积 20 亩，去年每亩租金 1600 元，今年每亩租金比去年上涨 50%。"
+             "今年该院落全部出租，租金总额是多少？注意：不需要考虑押金和卫生费。"),
+    ("trap", "一个出租院落面积 20 亩，去年每亩租金 1600 元，今年每亩租金比去年上涨 50%。"
+             "若同时考虑押金 2000 元和每亩卫生费 10 元，今年该院落全部出租的租金总额是多少？"
+             "（押金可退，不计入租金）"),
+    ("long", "人工智能的发展经历了多次浪潮。从符号主义到连接主义，从专家系统到深度学习，"
+             "每一次范式转移都带来了能力的跃升。近年来，大语言模型通过海量文本预训练与指令微调，"
+             "展现出前所未有的通用性。请把这段话压缩成一句话。"),
+    ("en", "Translate to Chinese: The quick brown fox jumps over the lazy dog."),
+    ("code", "用 Python 写一个判断质数的函数，只要函数体。"),
+]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base",
+                    default="/tmp/k27/hf_home/models--tchbcb--samai-27b/snapshots/883452faeb47f1394acd180a778dc5d53f049f1b")
+    ap.add_argument("--adapter", default="/tmp/k27/s27_sft_adapter")
+    ap.add_argument("--out", default="/tmp/k27/s27_calibration.json")
+    a = ap.parse_args()
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from peft import PeftModel
+
+    log("loading tokenizer ...")
+    tok = AutoTokenizer.from_pretrained(a.base, trust_remote_code=True)
+
+    log("loading base (nf4, manual balanced map) ...")
+    from accelerate import init_empty_weights
+    from transformers import AutoConfig
+
+    def build_map(base):
+        """与 s27_train.py 同款: visual→cpu, lm_head→GPU1(纯前向 logits 尾部), 层贪心均衡"""
+        cfgm = AutoConfig.from_pretrained(base, trust_remote_code=True)
+        with init_empty_weights():
+            mm = AutoModelForCausalLM.from_config(cfgm, trust_remote_code=True)
+        LM = "model.language_model"
+        SKIP = ("lm_head", "router", "halter", "mtp", "norm", "visual", "embed_tokens")
+        lay, oth = {}, {}
+        for name, p in mm.named_parameters():
+            if name.startswith(LM + ".layers."):
+                li = int(name.split(".layers.")[1].split(".")[0])
+                b = p.numel() * (2 if ("router" in name or "halter" in name or "norm" in name) else 0.6)
+                lay[li] = lay.get(li, 0) + b
+            else:
+                b = p.numel() * (2 if any(k in name for k in SKIP) else 0.6)
+                top = ".".join(name.split(".")[:3]) if name.startswith(LM) else ".".join(name.split(".")[:2])
+                oth[top] = oth.get(top, 0) + b
+        dmap = {LM + ".embed_tokens": 0}
+        fill = oth.pop(LM + ".embed_tokens", 0)
+        oth.pop(LM + ".norm", None)
+        oth.pop("lm_head", None)
+        hi = 0
+        for li in sorted(lay):
+            if fill + lay[li] <= 10.0 * 2**30:
+                dmap[f"{LM}.layers.{li}"] = 0
+                fill += lay[li]
+                hi = li
+        for li in sorted(lay):
+            if li > hi:
+                dmap[f"{LM}.layers.{li}"] = 1
+        dmap[LM + ".norm"] = 1
+        dmap["mtp"] = 1
+        dmap["lm_head"] = 1
+        dmap["model.visual"] = "cpu"
+        for extra, _b in oth.items():
+            if extra.startswith(LM) and extra not in dmap:
+                dmap[extra] = 1
+        for name, _ in mm.named_modules():
+            if not name or name in dmap:
+                continue
+            if any(name.startswith(k + ".") for k in dmap):
+                continue
+            if any(k.startswith(name + ".") for k in dmap):
+                continue
+            parts = name.split(".")
+            dev = 0
+            for i in range(len(parts) - 1, 0, -1):
+                anc = ".".join(parts[:i])
+                if anc in dmap:
+                    dev = dmap[anc]
+                    break
+            dmap[name] = dev
+        del mm
+        return dmap
+
+    bnb = BitsAndBytesConfig(
+        load_in_4bit=True, bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True,
+        llm_int8_skip_modules=["lm_head", "router", "halter", "mtp", "norm", "visual"],
+        llm_int8_enable_fp32_cpu_offload=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        a.base, trust_remote_code=True, quantization_config=bnb,
+        dtype=torch.float16, attn_implementation="eager",
+        device_map=build_map(a.base))
+    log(f"gpu0={torch.cuda.memory_allocated(0)/2**30:.2f}G "
+        f"gpu1={torch.cuda.memory_allocated(1)/2**30:.2f}G")
+
+    log("attaching adapter ...")
+    model = PeftModel.from_pretrained(model, a.adapter)
+    model.eval()
+
+    cfg = model.config
+    cfg.ponder_all_positions = True      # prefill 全位置思考 → 统计充分
+    cfg.ponder_forced_steps = 0
+    cfg.use_cache = False
+
+    pl_ids = list(cfg.ponder_layers)
+    pl_layers = list(model.ponder_layers)
+    log(f"ponder layers: {pl_ids}")
+    K = int(getattr(cfg, "max_ponder_steps", 8) or 8)
+    eps = float(getattr(cfg, "ponder_epsilon", 0.01) or 0.01)
+
+    freq_acc = {li: np.zeros(8) for li in pl_ids}
+    freq_n = 0
+    W_sum = np.zeros(K)      # 各步质量累加 (跨层/位置/样本)
+    W_pos = 0                # 位置·层 计数
+    steps_hist = np.zeros(K + 1)
+
+    for i, (tag, text) in enumerate(PROMPTS):
+        enc = tok.apply_chat_template(
+            [{"role": "user", "content": text}], add_generation_prompt=True,
+            tokenize=True, return_tensors="pt", enable_thinking=False)
+        if not torch.is_tensor(enc):
+            enc = enc["input_ids"]
+        with torch.no_grad():
+            model(input_ids=enc.cuda(), use_cache=False)
+        # ---- expert freq (fresh 路由的 _last_mask: (B,T,N)) ----
+        for li, layer in zip(pl_ids, pl_layers):
+            m = getattr(layer.mlp, "_last_mask", None)
+            if m is not None:
+                freq_acc[li] += m.float().mean(dim=(0, 1)).cpu().numpy()
+        # ---- halting w (last_w_stack: (B, W, K_exec), 已含 p_N 截断) ----
+        for layer in pl_layers:
+            w = getattr(layer, "last_w_stack", None)
+            if w is None:
+                continue
+            wf = w.float().cpu().numpy()
+            k_exec = wf.shape[-1]
+            W_sum[:k_exec] += wf.reshape(-1, k_exec).sum(0)
+            W_pos += wf.shape[0] * wf.shape[1]
+            s = getattr(layer, "last_steps", None)
+            if s is not None:
+                for st in s.float().flatten().cpu().numpy().clip(0, K).astype(int):
+                    steps_hist[st] += 1
+        freq_n += 1
+        log(f"{i+1}/{len(PROMPTS)} {tag}: ok")
+        torch.cuda.empty_cache()
+
+    expert_freq = {str(li): (freq_acc[li] / max(freq_n, 1)).round(6).tolist()
+                   for li in pl_ids}
+    halting_w = (W_sum / max(W_pos, 1))
+    halting_w = [w for w in halting_w if w > 1e-6]
+    ssum = sum(halting_w) or 1.0
+    halting_w = [round(w / ssum, 6) for w in halting_w]
+    steps_mean = float((np.arange(1, K + 1) * steps_hist[1:]).sum()
+                       / max(steps_hist.sum(), 1.0))
+    out = {
+        "expert_freq": expert_freq,
+        "halting_w": halting_w,
+        "steps_hist": steps_hist.astype(int).tolist(),
+        "steps_mean": round(steps_mean, 4),
+        "n_prompts": len(PROMPTS),
+        "adapter": a.adapter,
+        "note": "expert_freq = 激活频率 f_i (静态合并 concat_i(f_i*E_i)); SFT 后 router",
+    }
+    json.dump(out, open(a.out, "w"), indent=1, ensure_ascii=False)
+    log("saved " + a.out)
+    print(json.dumps(out, indent=1))
+
+
+if __name__ == "__main__":
+    main()