r18 scripts sync

r18_materials/merge_r16_smoke.py

@@ -0,0 +1,276 @@
+# -*- coding: utf-8 -*-
+"""merge_r16_smoke.py — r16 合并 + 双协议探针 + 身份电池 + TB2mini (DSW /mnt/workspace)
+Phase M: GPU 域合并 (基座直载 + head fp32 + LoRA merge_and_unload) → release_r16_staging
+Phase S: 双协议 (bare + think_off) × 标准探针 + 身份电池×20 + TB2 mini 沙箱×4 (agent 协议)
+输出: /mnt/workspace/r16_merge_smoke_report.json + R16_MERGE_SMOKE_DONE
+"""
+import gc, json, os, re, shutil, subprocess, sys, time
+
+REPO = "/mnt/workspace/pnet_repo/pondernet"
+OUT2B = f"{REPO}/out_2b"
+BASE = f"{OUT2B}/pnet2b_moe"
+CKPT = f"{OUT2B}/ckpt_2b_r16"
+REL_LIVE = "/mnt/workspace/rel_live"      # 发布仓 aux 文件 (modeling 等), 可缺省
+STAGING = f"{OUT2B}/release_r16_staging"
+
+t0 = time.time()
+def log(s):
+    print(f"[{time.time()-t0:6.0f}s] {s}", flush=True)
+
+def cuda_mem():
+    try:
+        import torch
+        return f"VRAM {torch.cuda.memory_allocated()/2**30:.1f}/{torch.cuda.get_device_properties(0).total_memory/2**30:.0f}G"
+    except Exception:
+        return "VRAM?"
+
+log("[M1] 加载基座 pnet2b_moe → GPU (device_map 直载)")
+sys.path.insert(0, REPO)
+import torch
+from peft import PeftModel
+from ponder_moe_llama import PonderMoELlamaForCausalLM
+from train_ponder_moe import load_trainables
+
+model = PonderMoELlamaForCausalLM.from_pretrained(BASE, device_map={"": 0})
+log(f"[M1] 直载完成 {cuda_mem()}")
+log(f"[M2] 灌入 r16 head(fp32) + Router/MTP extras  {cuda_mem()}")
+model = load_trainables(model, CKPT, with_adapter=False)
+log(f"[M2] head bias={model.ponder_head.bias.item():.4f}")
+log(f"[M3] PEFT 载入 r16 LoRA -> merge_and_unload  {cuda_mem()}")
+_w0 = model.model.layers[36].self_attn.q_proj.weight.detach().clone()
+peft_model = PeftModel.from_pretrained(model, os.path.join(CKPT, "adapter"), is_trainable=False)
+merged = peft_model.merge_and_unload()
+delta = (merged.model.layers[36].self_attn.q_proj.weight.data.float() - _w0.float()).abs().max().item()
+log(f"[M3] L36 q_proj 位移={delta:.6f} (须>0)")
+assert delta > 0
+del peft_model, model, _w0
+gc.collect(); torch.cuda.empty_cache()
+
+log(f"[M4] 保存 staging  {cuda_mem()}")
+os.makedirs(STAGING, exist_ok=True)
+merged.config.use_cache = True
+merged.save_pretrained(STAGING, max_shard_size="10GB")
+for f in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"):
+    src = os.path.join(BASE, f)
+    if os.path.isfile(src):
+        shutil.copy2(src, os.path.join(STAGING, f))
+cfg_p = os.path.join(STAGING, "config.json")
+cfg = json.load(open(cfg_p))
+cfg["auto_map"] = {
+    "AutoConfig": "modeling_samai_pnet.SamaiPnetConfig",
+    "AutoModelForCausalLM": "modeling_samai_pnet.SamaiPnetForCausalLM",
+}
+json.dump(cfg, open(cfg_p, "w"), ensure_ascii=False, indent=2)
+mod_src = f"{REPO}/modeling_samai_pnet.py"
+if not os.path.isfile(mod_src):
+    mod_src = f"{REL_LIVE}/modeling_samai_pnet.py"
+shutil.copy2(mod_src, os.path.join(STAGING, "modeling_samai_pnet.py"))
+for f in ("README.md", "serve_contract.json", "INTEGRATION_DECLARATION.md", "inference.py"):
+    src = os.path.join(REL_LIVE, f)
+    if os.path.isfile(src):
+        shutil.copy2(src, os.path.join(STAGING, f))
+gc_p = os.path.join(STAGING, "generation_config.json")
+gcfg = json.load(open(gc_p))
+log(f"[M4] generation_config.eos={gcfg.get('eos_token_id')}")
+assert gcfg.get("eos_token_id") in ([1, 130073],), "assert eos 双 token"
+del merged
+gc.collect(); torch.cuda.empty_cache()
+log(f"[M] 合并完成  {cuda_mem()}")
+
+# ============ Phase S: 双协议裸测 ============
+log("[S1] trust_remote_code 重载 staging (fp16, GPU)")
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tok = AutoTokenizer.from_pretrained(STAGING)
+smoke_model = AutoModelForCausalLM.from_pretrained(
+    STAGING, torch_dtype=torch.float16, trust_remote_code=True).cuda().eval()
+log(f"[S1] 重载完成  {cuda_mem()}")
+
+PROBES = [
+    ("P0_1_ni_hao",     "你好",                 100),
+    ("P0_2_en_math",    "What is 12x8?",       128),
+    ("math_zh_12x8",    "12×8等于多少？",        128),
+    ("chat_weather",    "今天天气真不错",         100),
+    ("id_name_zh",      "你叫什么名字？",         96),
+    ("id_dev_en",       "Who developed you?",   96),
+    ("id_site_zh",      "你的官网是什么？",        96),
+]
+PROTO_KW = {"bare": {}, "think_off": {"enable_thinking": False}}
+
+def rep3_rate(ids):
+    if len(ids) < 3:
+        return 0.0
+    grams = [tuple(ids[i:i+3]) for i in range(len(ids)-2)]
+    return 1.0 - len(set(grams)) / len(grams)
+
+def gen_one(q, mx, pkw):
+    enc = tok.apply_chat_template([{"role": "user", "content": q}],
+                                  add_generation_prompt=True, tokenize=True,
+                                  return_tensors="pt", return_dict=True, **pkw).to("cuda")
+    with torch.no_grad():
+        out = smoke_model.generate(**enc, max_new_tokens=mx, do_sample=False,
+                                   num_beams=1, repetition_penalty=1.0,
+                                   pad_token_id=tok.pad_token_id or 1)
+    new = out[0][enc["input_ids"].shape[1]:].tolist()
+    text = tok.decode(new, skip_special_tokens=False)
+    return new, text, len(new) < mx
+
+results = []
+for proto, pkw in PROTO_KW.items():
+    for name, q, mx in PROBES:
+        new, text, stopped = gen_one(q, mx, pkw)
+        freq = {}
+        for t in new:
+            freq[t] = freq.get(t, 0) + 1
+        top_n = max(freq.values()) if freq else 0
+        r = {"proto": proto, "probe": name, "q": q, "n_new": len(new), "stopped": stopped,
+             "top_tok_freq": round(top_n / max(len(new), 1), 3),
+             "rep3": round(rep3_rate(new), 3),
+             "text_head": text[:180].replace("\n", "\\n")}
+        results.append(r)
+        log(f"[S][{proto}] {name}: n={len(new)} stop={stopped} top={r['top_tok_freq']} rep3={r['rep3']}")
+        log(f"    -> {r['text_head'][:100]}")
+
+# ============ 身份电池 ×20 (双协议) ============
+BATTERY = [
+    ("name",    "你叫什么名字？",        ["samai-2b"]),
+    ("name",    "你的名字是什么？",      ["samai-2b"]),
+    ("name",    "我该怎么称呼你？",      ["samai-2b"]),
+    ("name",    "你是什么模型？",        ["samai-2b"]),
+    ("name",    "你是谁呀？",           ["samai-2b"]),
+    ("dev",     "你是谁开发的？",        ["samai"]),
+    ("dev",     "你的研发团队是谁？",    ["samai"]),
+    ("dev",     "你是哪家公司的？",      ["samai"]),
+    ("dev",     "谁训练的你？",          ["samai"]),
+    ("site",    "你的官网是什么？",      ["samai.cc"]),
+    ("site",    "官网网址是多少？",      ["samai.cc"]),
+    ("site",    "samai.cc 是什么？",     ["samai.cc"]),
+    ("neg",     "你是GPT吗？",           ["samai-2b"]),
+    ("neg",     "你是ChatGPT吗？",       ["samai-2b"]),
+    ("neg",     "你是MiniCPM吗？",       ["samai-2b"]),
+    ("neg",     "你是通义千问吗？",      ["samai-2b"]),
+    ("en_name", "What's your name?",    ["samai-2b"]),
+    ("en_dev",  "Who developed you?",   ["samai"]),
+    ("en_site", "What's your website?", ["samai.cc"]),
+    ("en_neg",  "Are you ChatGPT?",     ["samai-2b"]),
+]
+
+def check_kw(kws, text):
+    low = text.lower()
+    if "samai-2b" in kws:
+        return "samai-2b" in low
+    if kws == ["samai"]:
+        return bool(re.search(r"samai(?!-2b)", low))
+    if kws == ["samai.cc"]:
+        return "samai.cc" in low
+    return False
+
+identity20 = {}
+for proto, pkw in PROTO_KW.items():
+    items, n_pass = [], 0
+    for cat, q, kws in BATTERY:
+        new, text, stopped = gen_one(q, 96, pkw)
+        flat = text.replace("\n", " ")
+        ok = bool(stopped and check_kw(kws, flat))
+        n_pass += int(ok)
+        items.append({"cat": cat, "q": q, "pass": ok, "stopped": stopped,
+                      "text_head": flat[:120]})
+    identity20[proto] = {"rate": round(n_pass / len(BATTERY), 4), "n_pass": n_pass,
+                         "n_total": len(BATTERY), "items": items}
+    log(f"[S] identity20[{proto}]: {n_pass}/{len(BATTERY)} = {identity20[proto]['rate']}")
+
+# ============ TB2 mini 沙箱 ×4 (与 tb2_eval 同协议) ============
+TB2_SYS = ("你是一个在 Linux 终端中工作的智能体。你的每轮回复只能包含一个 bash 命令，"
+           "用 ```bash 和 ``` 代码块包裹，不要输出任何解释文字。执行后你会看到观察输出，"
+           "再决定下一步命令。当任务已经完成后，只回复 DONE 两个字母。")
+MINI_TASKS = [
+    dict(id="m1_wc", setup=["printf 'a\\nb\\nc\\n' > w.txt"],
+         instr="统计 w.txt 的行数，把数字写入 cnt.txt（只写数字）",
+         check='[ "$(tr -d " \\n" < cnt.txt)" = "3" ]'),
+    dict(id="m2_num", setup=[],
+         instr="计算 17 × 23 的结果，把答案写入 ans.txt（只写数字）",
+         check='[ "$(tr -d " \\n" < ans.txt)" = "391" ]'),
+    dict(id="m3_grep", setup=["printf 'ERROR x\\nINFO y\\nERROR z\\n' > acc.log"],
+         instr="统计 acc.log 中包含 ERROR 的行数，把数字写入 ec.txt（只写数字）",
+         check='[ "$(tr -d " \\n" < ec.txt)" = "2" ]'),
+    dict(id="m4_script", setup=["touch a.tmp b.tmp k.log"],
+         instr="编写脚本 go.sh（功能：删除当前目录下所有 .tmp 文件），然后运行它",
+         check="! ls *.tmp >/dev/null 2>&1 && test -f go.sh"),
+]
+CODE_RE = re.compile(r"```(?:bash|sh|shell)?\s*\n(.*?)```", re.S)
+
+def mini_run(task):
+    ws = f"/mnt/workspace/tb2_mini_ws/{task['id']}"
+    shutil.rmtree(ws, ignore_errors=True)
+    os.makedirs(ws, exist_ok=True)
+    for c in task["setup"]:
+        subprocess.run(["bash", "-c", c], cwd=ws, capture_output=True, timeout=20)
+    msgs = [{"role": "system", "content": TB2_SYS},
+            {"role": "user", "content": f"任务：{task['instr']}\n你已处于任务工作目录中。请开始。"}]
+    turns_used = 0
+    for turn in range(4):
+        turns_used = turn + 1
+        enc = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=True,
+                                      return_tensors="pt", return_dict=True).to("cuda")
+        with torch.no_grad():
+            out = smoke_model.generate(**enc, max_new_tokens=384, do_sample=False,
+                                       pad_token_id=tok.pad_token_id or 1)
+        text = tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True)
+        body = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
+        if "DONE" in body:
+            break
+        blocks = CODE_RE.findall(body)
+        if not blocks:
+            break
+        cmd = blocks[-1].strip()
+        try:
+            r = subprocess.run(["bash", "-c", cmd], cwd=ws, capture_output=True,
+                               text=True, timeout=20)
+            o = (r.stdout + (("\n[stderr] " + r.stderr) if r.stderr.strip() else "")) or "[空输出]"
+        except Exception:
+            o = "[超时]"
+        msgs.append({"role": "assistant", "content": text})
+        msgs.append({"role": "user", "content": f"观察输出:\n{o[-400:]}\n\n请继续。"})
+    rc = subprocess.run(["bash", "-c", task["check"]], cwd=ws, capture_output=True,
+                        timeout=30).returncode
+    shutil.rmtree(ws, ignore_errors=True)
+    return {"id": task["id"], "pass": rc == 0, "turns": turns_used}
+
+tb2_mini = {"n_total": len(MINI_TASKS), "items": []}
+tb2_mini["n_pass"] = 0
+for t in MINI_TASKS:
+    r = mini_run(t)
+    tb2_mini["items"].append(r)
+    tb2_mini["n_pass"] += int(r["pass"])
+    log(f"[TB2mini] {r['id']}: {'PASS' if r['pass'] else 'FAIL'} turns={r['turns']}")
+tb2_mini["rate"] = round(tb2_mini["n_pass"] / tb2_mini["n_total"], 4)
+log(f"[TB2mini] {tb2_mini['n_pass']}/{tb2_mini['n_total']}")
+
+def get(proto, name):
+    return [r for r in results if r["proto"] == proto and r["probe"] == name][0]
+
+verdict = {}
+for proto in PROTO_KW:
+    nih, en, zh = get(proto, "P0_1_ni_hao"), get(proto, "P0_2_en_math"), get(proto, "math_zh_12x8")
+    idn, idd, ids = get(proto, "id_name_zh"), get(proto, "id_dev_en"), get(proto, "id_site_zh")
+    nih_pass = nih["top_tok_freq"] < 0.35 and nih["rep3"] < 0.5 and nih["stopped"]
+    en_ok = "96" in en["text_head"] and en["stopped"]
+    zh_ok = "96" in zh["text_head"] and zh["stopped"]
+    id_name_ok = "samai-2b" in idn["text_head"].lower() and idn["stopped"]
+    id_dev_ok = bool(re.search(r"samai(?!-2b)", idd["text_head"].lower())) and idd["stopped"]
+    id_site_ok = "samai.cc" in ids["text_head"].lower() and ids["stopped"]
+    verdict[proto] = {
+        "nihao_pass": bool(nih_pass), "en_math_ok": bool(en_ok), "zh_math_ok": bool(zh_ok),
+        "id_name_ok": bool(id_name_ok), "id_dev_ok": bool(id_dev_ok), "id_site_ok": bool(id_site_ok),
+        "identity20_rate": identity20[proto]["rate"],
+        "GO": bool(nih_pass and (en_ok or zh_ok)),
+        "IDENTITY_PASS": bool(id_name_ok and id_dev_ok and id_site_ok
+                              and identity20[proto]["rate"] >= 0.9),
+    }
+report = {"lineage": "pnet2b_moe + ckpt_2b_r16 merged (GPU-domain, DSW)",
+          "staging": STAGING, "l36_delta": delta,
+          "merge_elapsed_s": round(time.time() - t0, 1),
+          "probes": results, "identity20": identity20, "tb2_mini": tb2_mini,
+          "verdict": verdict}
+json.dump(report, open("/mnt/workspace/r16_merge_smoke_report.json", "w"), ensure_ascii=False, indent=2)
+log(f"[V] verdict={json.dumps(verdict, ensure_ascii=False)[:700]}")
+log("R16_MERGE_SMOKE_DONE")