| |
| """Merge a LoRA SFT checkpoint into the base, export HF Qwen3, and prove vLLM reads what align wrote. |
| |
| CUDA_VISIBLE_DEVICES=0,1 .venv/bin/python scripts/qwen3_export_check.py run \ |
| --ckpt RUN/checkpoints/micro_N --out-dir .tmp/p2/p2a/export [--keep] |
| |
| The claim under test: *the file align writes is the model align trained.* Five independent checks, |
| each of which fails if the export is wrong (the process exit status gates ALL of them, `export_ok`): |
| |
| 1. MERGE PATH vs TRAINER FILE. align merges the adapter (fp32 B@A * alpha/r, one bf16 cast) into |
| the pristine snapshot weights from the adapter stored in the checkpoint's optimizer payload, |
| and writes it with `save_hf_qwen3`. The trainer's own checkpoint weight file was built by a |
| different route (FSDP gather -> `export_merged_state_dict`). The two must be bit-identical. |
| 2. vLLM LOADS IT. A separate venv (`.venv-rollout`, vLLM 0.26) loads the exported directory |
| unmodified -- sharded safetensors, config.json (rope_theta under both spellings), tokenizer -- |
| and greedy-generates 64 tokens on 2 prompts, returning the top-20 logprobs at every position. |
| 3. FIRST-TOKEN LOGPROBS. align's own forward on the merged weights (bf16, GPU 1) scores the same |
| prompt; the first-token top-1 id must agree and the logprob difference is reported. |
| 4. TEACHER-FORCED AGREEMENT over all 64 generated tokens: align scores prompt+generation in one |
| pass; per-position argmax agreement with what vLLM sampled, and max |dlogprob| of the sampled |
| token. bf16 TOLERANCE (documented, not fitted): two correct bf16 engines with different |
| attention/GEMM kernels and accumulation order differ by a few ulps of a logit of magnitude |
| 10-30 (ulp 0.0625-0.125) amplified through 36 layers; the gate is |dlogprob| <= 0.25 on the |
| first token and argmax agreement >= 95% over the 64 positions (a near-tie legitimately flips |
| argmax; a wrong merge, a missing QK-norm or a wrong rope_theta disagrees on most positions and |
| by whole nats). The thresholds are constants below, asserted by `evaluate`, and exercised on |
| injected mutants in tests/test_qwen3_scripts.py. |
| 5. ADAPTER vs MERGED. The single bf16 cast of W+delta rounds away any delta below half an ulp of |
| W (a dead zone, not noise), so at small LoRA movement the merged file can be a different |
| policy from base+adapter -- and checks 1-4 all score the MERGED file against itself. Here the |
| SAME teacher-forced tokens are scored by align with the adapter applied live (base bf16 + |
| fp32 delta) and by the merged file; the gates are the first-token / argmax-agreement |
| constants of check 4 (reused, not refitted). The retained-delta cosine (merged-W vs the |
| intended delta) is recorded alongside. Also gated: the real snapshot's hf<->align mapping |
| round trip, and that the adapter MOVED every adapted matrix (a zero adapter exports the base). |
| |
| The export is DELETED afterwards unless --keep (16.4 GB on a 94%-full disk). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import shutil |
| import subprocess |
| import sys |
| import time |
|
|
| REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, REPO) |
| ROLLOUT_PY = os.path.join(REPO, ".venv-rollout", "bin", "python") |
| SNAPSHOT = os.environ.get("P2_EXPORT_BASE_SNAPSHOT") or ("/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/" |
| "dfd02e6271a58375dfbf3ece0175277cf6b6a89a") |
| WORK = os.path.join(REPO, ".tmp", "p2", "p2a", "export_check") |
|
|
| N_NEW = 64 |
| TOPK = 20 |
| FIRST_TOKEN_MAX_ABS_DLOGPROB = 0.25 |
| MIN_ARGMAX_AGREEMENT = 0.95 |
|
|
| PROMPTS = ( |
| "<|im_start|>user\nComplete the following Lean 4 code:\n\n```lean4\nimport Mathlib\nimport Aesop\n\n" |
| "set_option maxHeartbeats 400000\n\nopen BigOperators Real Nat Topology Rat\n\n" |
| "theorem mathd_algebra_10 : abs ((120 : ℝ) / 100 * 30 - 130 / 100 * 20) = 10 := by\n sorry\n```\n" |
| "<|im_end|>\n<|im_start|>assistant\n", |
| "<|im_start|>user\nWhat is 17 * 23? Answer briefly.<|im_end|>\n<|im_start|>assistant\n", |
| ) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def merge_lora_state(base: dict, lora: dict, scaling: float) -> dict: |
| """`base` (align-keyed, `<path>.weight`) with `scaling * B @ A` folded in for every adapter path |
| in `lora` (`<path>.A` / `<path>.B`, as `align.ppo.lora.lora_state_dict` writes them). |
| |
| The delta is accumulated in fp32 and cast ONCE to the weight's dtype (the same arithmetic as |
| `export_merged_state_dict`, so the two routes can be compared bit for bit). Every adapter path |
| must resolve to a base weight and must have both halves: a half-present adapter or a path that |
| matches nothing would otherwise "export" an unmodified base. |
| """ |
| import torch |
|
|
| a_keys = {k[:-2] for k in lora if k.endswith(".A")} |
| b_keys = {k[:-2] for k in lora if k.endswith(".B")} |
| stray = sorted(k for k in lora if not (k.endswith(".A") or k.endswith(".B"))) |
| if stray: |
| raise KeyError(f"adapter state has keys that are neither .A nor .B: {stray[:4]}") |
| if a_keys != b_keys: |
| raise KeyError(f"adapter halves do not pair: A-only={sorted(a_keys - b_keys)[:3]} B-only={sorted(b_keys - a_keys)[:3]}") |
| if not a_keys: |
| raise KeyError("adapter state is empty; refusing to export an unmodified base as a merge") |
| out = dict(base) |
| for p in sorted(a_keys): |
| wk = f"{p}.weight" |
| if wk not in base: |
| raise KeyError(f"adapter path {p!r} has no base weight {wk!r}") |
| w, A, B = base[wk], lora[f"{p}.A"], lora[f"{p}.B"] |
| if A.shape[0] != B.shape[1] or (B.shape[0], A.shape[1]) != tuple(w.shape): |
| raise ValueError(f"{p}: A{tuple(A.shape)} B{tuple(B.shape)} do not compose to {tuple(w.shape)}") |
| out[wk] = (w.float() + (B.float() @ A.float()) * scaling).to(w.dtype) |
| return out |
|
|
|
|
| |
| |
| |
|
|
|
|
| def compare_generation(vllm: dict, align: dict) -> dict: |
| """Per-prompt metrics between vLLM's generation record and align's teacher-forced scoring. |
| |
| `vllm[i]`: {"token_ids": [64 ints], "chosen_logprobs": [64 floats], "first_top": {id: logprob}}. |
| `align[i]`: {"argmax": [64 ints], "chosen_logprobs": [64 floats], "first_logprobs": {id: logprob}} |
| (align's log-softmax at the same 64 positions, and at position 0 for vLLM's top-20 ids). |
| """ |
| per = [] |
| for v, a in zip(vllm, align): |
| n = len(v["token_ids"]) |
| if not (len(a["argmax"]) == len(a["chosen_logprobs"]) == n): |
| raise ValueError("generation records have different lengths") |
| agree = sum(int(x == y) for x, y in zip(v["token_ids"], a["argmax"])) |
| dchosen = [abs(x - y) for x, y in zip(v["chosen_logprobs"], a["chosen_logprobs"])] |
| v_top1 = max(v["first_top"], key=v["first_top"].get) |
| a_first = a["first_logprobs"] |
| per.append({ |
| "tokens": n, |
| "argmax_agreement": agree / n, |
| "max_abs_dlogprob_chosen": max(dchosen), |
| "mean_abs_dlogprob_chosen": sum(dchosen) / n, |
| "first_token_top1_same": int(v_top1) == int(max(a_first, key=a_first.get)), |
| "first_token_abs_dlogprob": abs(v["first_top"][v_top1] - a_first[v_top1]), |
| "first_token_max_abs_dlogprob_top20": max(abs(v["first_top"][k] - a_first[k]) for k in v["first_top"]), |
| }) |
| return {"per_prompt": per} |
|
|
|
|
| def evaluate(metrics: dict) -> dict: |
| checks = {} |
| for i, p in enumerate(metrics["per_prompt"]): |
| checks[f"prompt{i}.first_token_top1_same"] = bool(p["first_token_top1_same"]) |
| checks[f"prompt{i}.first_token_dlogprob"] = p["first_token_abs_dlogprob"] <= FIRST_TOKEN_MAX_ABS_DLOGPROB |
| checks[f"prompt{i}.argmax_agreement"] = p["argmax_agreement"] >= MIN_ARGMAX_AGREEMENT |
| return {"checks": checks, "passed": bool(checks) and all(checks.values()), |
| "thresholds": {"first_token_max_abs_dlogprob": FIRST_TOKEN_MAX_ABS_DLOGPROB, |
| "min_argmax_agreement": MIN_ARGMAX_AGREEMENT}} |
|
|
|
|
| def merge_fidelity_stats(base: dict, merged: dict, lora: dict, scaling: float) -> dict: |
| """How much of the intended delta `scaling * B @ A` survived the cast: aggregate cosine and |
| relative error of (merged - base) against the delta, and the fraction of nonzero delta |
| elements that rounded to exactly zero (the bf16 dead zone). Same arithmetic as |
| `align.ppo.lora.export_merged_state_dict`'s report.""" |
| import math |
|
|
| import torch |
|
|
| dot = ne = nd = 0.0 |
| nz_total = dropped = layers = 0 |
| for p in sorted({k[:-2] for k in lora if k.endswith(".A")}): |
| wk = f"{p}.weight" |
| d = ((lora[f"{p}.B"].float() @ lora[f"{p}.A"].float()) * scaling).flatten() |
| e = (merged[wk].float() - base[wk].float()).flatten() |
| dot, ne, nd = dot + float(torch.dot(e, d)), ne + float(torch.dot(e, e)), nd + float(torch.dot(d, d)) |
| nz = d != 0 |
| nz_total += int(nz.sum()) |
| dropped += int((nz & (e == 0)).sum()) |
| layers += 1 |
| if nd == 0.0: |
| return {"layers": layers, "cosine": 1.0, "rel_error": 0.0, "dropped_fraction": 0.0, "delta_norm": 0.0} |
| return {"layers": layers, "cosine": dot / math.sqrt(max(ne * nd, 1e-300)), |
| "rel_error": math.sqrt(max(ne - 2.0 * dot + nd, 0.0) / nd), |
| "dropped_fraction": dropped / nz_total if nz_total else 0.0, "delta_norm": math.sqrt(nd)} |
|
|
|
|
| def compare_adapter_vs_merged(adapter_lps: list, merged_lps: list) -> dict: |
| """Per-prompt agreement between align's log-softmax at the generated positions with the adapter |
| applied live (`adapter_lps[i]`, [T, V]) and on the merged bf16 file (`merged_lps[i]`, [T, V]).""" |
| import torch |
|
|
| per = [] |
| for a, m in zip(adapter_lps, merged_lps): |
| if a.shape != m.shape or a.dim() != 2: |
| raise ValueError(f"log-prob records differ in shape: {tuple(a.shape)} vs {tuple(m.shape)}") |
| a, m = a.float(), m.float() |
| a_top, m_top = a.argmax(-1), m.argmax(-1) |
| t0 = int(a_top[0]) |
| per.append({ |
| "tokens": a.shape[0], |
| "argmax_agreement": float((a_top == m_top).float().mean()), |
| "first_token_abs_dlogprob": abs(float(a[0, t0] - m[0, t0])), |
| "max_abs_dlogprob_top1": float((a.gather(1, a_top[:, None]) - m.gather(1, a_top[:, None])).abs().max()), |
| "kl_adapter_to_merged_mean": float((a.exp() * (a - m)).sum(-1).mean()), |
| }) |
| return {"per_prompt": per} |
|
|
|
|
| def evaluate_adapter_vs_merged(metrics: dict) -> dict: |
| checks = {} |
| for i, p in enumerate(metrics["per_prompt"]): |
| checks[f"prompt{i}.first_token_dlogprob"] = p["first_token_abs_dlogprob"] <= FIRST_TOKEN_MAX_ABS_DLOGPROB |
| checks[f"prompt{i}.argmax_agreement"] = p["argmax_agreement"] >= MIN_ARGMAX_AGREEMENT |
| return {"checks": checks, "passed": bool(checks) and all(checks.values()), |
| "thresholds": {"first_token_max_abs_dlogprob": FIRST_TOKEN_MAX_ABS_DLOGPROB, |
| "min_argmax_agreement": MIN_ARGMAX_AGREEMENT}} |
|
|
|
|
| def export_ok(report: dict) -> bool: |
| """The process verdict. EVERY advertised check must have run and passed: a check that did not |
| run (key absent) is a failure, not a pass -- a zero adapter, a lossy hf<->align mapping or a |
| merge that dropped the delta must not exit 0.""" |
| rel_min = (report.get("merge_relative_weight_change") or {}).get("min") |
| return bool( |
| (report.get("verdict") or {}).get("passed") is True |
| and (report.get("merge_vs_trainer_file") or {}).get("bitwise_mismatches") == 0 |
| and (report.get("snapshot_mapping_roundtrip") or {}).get("bitwise_equal") is True |
| and (report.get("merge_changed_tensors_vs_snapshot") or 0) > 0 |
| and isinstance(rel_min, (int, float)) and rel_min > 0 |
| and ((report.get("adapter_vs_merged") or {}).get("verdict") or {}).get("passed") is True |
| ) |
|
|
|
|
| |
| |
| |
|
|
| _VLLM_SCRIPT = r''' |
| import json, os, sys, time |
| model_dir, prompts_json, out_path, n_new, topk = sys.argv[1:6] |
| n_new, topk = int(n_new), int(topk) |
| os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0") |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") |
| from transformers import AutoTokenizer |
| from vllm import LLM, SamplingParams |
| tok = AutoTokenizer.from_pretrained(model_dir) |
| prompts = json.loads(prompts_json) |
| ids = [tok(p, add_special_tokens=False)["input_ids"] for p in prompts] |
| t0 = time.time() |
| llm = LLM(model=model_dir, tokenizer=model_dir, dtype="bfloat16", max_model_len=2048, |
| gpu_memory_utilization=0.85, enforce_eager=True, seed=0) |
| t_load = time.time() - t0 |
| sp = SamplingParams(temperature=0.0, max_tokens=n_new, logprobs=topk, ignore_eos=True) |
| outs = llm.generate([{"prompt_token_ids": x} for x in ids], sp) |
| res = [] |
| for x, o in zip(ids, outs): |
| c = o.outputs[0] |
| lps = c.logprobs |
| res.append({ |
| "prompt_ids": x, |
| "token_ids": list(c.token_ids), |
| "text": c.text, |
| "chosen_logprobs": [float(lps[t][c.token_ids[t]].logprob) for t in range(len(c.token_ids))], |
| "first_top": {str(k): float(v.logprob) for k, v in lps[0].items()}, |
| }) |
| json.dump({"results": res, "load_s": t_load, "model": model_dir}, open(out_path, "w")) |
| print(json.dumps({"ok": True, "n": len(res), "load_s": round(t_load, 1)}), flush=True) |
| # vLLM 0.26's EngineCore subprocess can hang at interpreter shutdown (measured: the result file was |
| # written at 00:44 and the process was still alive 20 minutes later holding 19 GiB of GPU 0). The |
| # result is on disk and flushed; leave without running the atexit/engine teardown. |
| sys.stdout.flush(); sys.stderr.flush() |
| os._exit(0) |
| ''' |
|
|
|
|
| def _first_visible_gpu(env=None) -> str: |
| """The child's CUDA_VISIBLE_DEVICES is relative to PHYSICAL indices, so a hard-coded "0" put |
| vLLM (0.85 of the card) on physical GPU 0 even when the operator pinned this check to 2,3. |
| Use the parent's FIRST visible entry ("0" when the parent is unpinned).""" |
| v = (os.environ if env is None else env).get("CUDA_VISIBLE_DEVICES", "") |
| first = v.split(",")[0].strip() |
| return first or "0" |
|
|
|
|
| def _rollout_env(gpu: str) -> dict: |
| scrub = ("RANK", "LOCAL_RANK", "WORLD_SIZE", "LOCAL_WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "TORCHELASTIC") |
| env = {k: v for k, v in os.environ.items() if not k.startswith(scrub)} |
| env.update({ |
| "CUDA_VISIBLE_DEVICES": gpu, "VLLM_USE_FLASHINFER_SAMPLER": "0", |
| "TMPDIR": os.path.join(REPO, ".tmp", "p2", "p2a"), "XDG_CACHE_HOME": os.path.join(REPO, ".cache"), |
| "HF_HOME": "/apps/development/rlhf/hf_cache", "HF_HUB_CACHE": "/apps/development/rlhf/hf_cache/hub", |
| "HF_HUB_DISABLE_XET": "1", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", |
| "TRITON_CACHE_DIR": os.path.join(REPO, ".cache", "triton"), |
| "TORCHINDUCTOR_CACHE_DIR": os.path.join(REPO, ".cache", "torchinductor"), |
| "PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false", |
| }) |
| return env |
|
|
|
|
| |
| |
| |
|
|
|
|
| def run_vllm_child(cmd, result_path, *, env, cwd, log_path, grace_s: float = 30.0, timeout_s: float = 3000.0): |
| """Run the vLLM child WITHOUT inheriting pipes, and stop its process group once it has written |
| its result. |
| |
| DEFECT THIS PREVENTS (measured in the P2 smoke loop): vLLM 0.26's EngineCore can outlive its |
| parent at interpreter exit (re-parented to init, holding the parent's stdout/stderr pipes). With |
| `subprocess.run(capture_output=True)` the parent then blocks in communicate() until the ORPHAN |
| closes the pipe -- 12+ minutes, the result file already written, GPU 0 held at 19 GB, the heavy |
| lock held. Here the child logs to a file (nothing for an orphan to hold open), runs in its own |
| session so the whole group can be signalled by its exact pgid, and is terminated `grace_s` after |
| the result file appears. Returns (returncode or None if it had to be terminated, log tail). |
| """ |
| import signal |
|
|
| with open(log_path, "w") as log: |
| proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT, env=env, cwd=cwd, start_new_session=True) |
| t0 = time.time() |
| done_at = None |
| rc = None |
| while True: |
| rc = proc.poll() |
| if rc is not None: |
| break |
| if os.path.isfile(result_path) and done_at is None: |
| done_at = time.time() |
| if (done_at is not None and time.time() - done_at > grace_s) or time.time() - t0 > timeout_s: |
| break |
| time.sleep(1.0) |
| |
| for sig in (signal.SIGTERM, signal.SIGKILL): |
| try: |
| os.killpg(proc.pid, sig) |
| except ProcessLookupError: |
| break |
| time.sleep(5.0) |
| if rc is None: |
| try: |
| rc = proc.wait(timeout=10) |
| except subprocess.TimeoutExpired: |
| rc = None |
| with open(log_path) as f: |
| return rc, f.read()[-3000:] |
|
|
|
|
| def _lora_scaling(ckpt: str, r: int | None, alpha: float | None) -> float: |
| if r is None or alpha is None: |
| raise SystemExit("--r and --alpha are required: the checkpoint stores A/B but not alpha/r, and guessing the scaling would export a wrong merge silently") |
| return alpha / r |
|
|
|
|
| def run(a) -> dict: |
| from align.utils.guard import apply_resource_guards |
|
|
| apply_resource_guards(cpu_threads=8, heavy=True, label="p2a-qwen3-export-check") |
| import torch |
|
|
| from align.model import qwen3 as q3 |
| from align.model.checkpoint import load_hf_safetensors |
| from align.model.qwen3_checkpoint import align_to_hf, hf_to_align, read_qwen3_config, save_hf_qwen3 |
|
|
| work = os.path.abspath(a.work) |
| out_dir = os.path.abspath(a.out_dir) |
| os.makedirs(work, exist_ok=True) |
| report: dict = {"ckpt": a.ckpt, "out_dir": out_dir} |
| qcfg = read_qwen3_config(SNAPSHOT) |
| scaling = _lora_scaling(a.ckpt, a.r, a.alpha) |
|
|
| |
| t0 = time.time() |
| raw = load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16) |
| base = hf_to_align(raw, qcfg) |
| |
| |
| back = align_to_hf(base, qcfg) |
| report["snapshot_mapping_roundtrip"] = {"tensors": len(raw), "bitwise_equal": all(torch.equal(back[k], raw[k]) for k in raw) |
| and set(back) == set(raw)} |
| del raw, back |
| payload = torch.load(os.path.join(a.ckpt, "optimizer_rank0.pt"), weights_only=True, map_location="cpu") |
| merged = merge_lora_state(base, payload["lora"], scaling) |
| report["merge_delta_retention"] = merge_fidelity_stats(base, merged, payload["lora"], scaling) |
| del base |
| save_hf_qwen3(merged, out_dir, cfg=qcfg, dtype=torch.bfloat16, tokenizer_dir=SNAPSHOT) |
| report["merge_s"] = round(time.time() - t0, 1) |
| ckpt_state = hf_to_align(load_hf_safetensors(a.ckpt, dtype=None), qcfg) |
| mism = [k for k, v in merged.items() if not torch.equal(v, ckpt_state[k])] |
| n_changed = sum(1 for p in {k[:-2] for k in payload["lora"] if k.endswith(".A")}) |
| report["merge_vs_trainer_file"] = {"tensors": len(merged), "bitwise_mismatches": len(mism), "first_mismatches": mism[:3], |
| "adapted_layers": n_changed} |
| del ckpt_state |
| |
| base2 = hf_to_align(load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16), qcfg) |
| moved = sum(1 for k, v in merged.items() if not torch.equal(v, base2[k])) |
| report["merge_changed_tensors_vs_snapshot"] = moved |
| |
| |
| |
| adapted = {f"{k[:-2]}.weight" for k in payload["lora"] if k.endswith(".A")} |
| rel = [((merged[k].float() - base2[k].float()).norm() / base2[k].float().norm()).item() for k in sorted(adapted)] |
| report["merge_relative_weight_change"] = {"n": len(rel), "mean": sum(rel) / len(rel), "max": max(rel), "min": min(rel)} |
| del base2, merged |
|
|
| |
| vout = os.path.join(work, "vllm_out.json") |
| env = _rollout_env(_first_visible_gpu()) |
| if os.path.isfile(vout): |
| os.remove(vout) |
| rc, tail = run_vllm_child([ROLLOUT_PY, "-c", _VLLM_SCRIPT, out_dir, json.dumps(list(PROMPTS)), vout, str(N_NEW), str(TOPK)], |
| vout, env=env, cwd=work, log_path=os.path.join(work, "vllm.log")) |
| report["vllm_returncode"] = rc |
| if not os.path.isfile(vout): |
| report["vllm_stderr_tail"] = tail |
| return report |
| vres = json.load(open(vout)) |
| report["vllm_load_s"] = vres["load_s"] |
| report["vllm_text"] = [r["text"][:200] for r in vres["results"]] |
|
|
| |
| dev = torch.device("cuda", 1 if torch.cuda.device_count() > 1 else 0) |
| cfg2, sd = (qcfg, hf_to_align(load_hf_safetensors(out_dir, dtype=torch.bfloat16), qcfg)) |
| model = q3.Qwen3ForCausalLM.from_state_dict(cfg2, sd, attn_implementation="sdpa", device="cuda").to(dev).eval() |
| del sd |
| arec, mlps = [], [] |
| with torch.no_grad(): |
| for r in vres["results"]: |
| ids = r["prompt_ids"] + r["token_ids"] |
| logits = model(input_ids=torch.tensor([ids], device=dev))[0].float() |
| lp = torch.log_softmax(logits, dim=-1) |
| n0 = len(r["prompt_ids"]) |
| pos = lp[n0 - 1: n0 - 1 + len(r["token_ids"])] |
| chosen = pos.gather(1, torch.tensor(r["token_ids"], device=dev)[:, None])[:, 0] |
| first = {k: float(pos[0, int(k)]) for k in r["first_top"]} |
| arec.append({"argmax": pos.argmax(-1).tolist(), "chosen_logprobs": chosen.tolist(), "first_logprobs": first}) |
| mlps.append(pos.cpu()) |
| metrics = compare_generation(vres["results"], arec) |
| report["metrics"] = metrics |
| report["verdict"] = evaluate(metrics) |
|
|
| |
| from align.ppo.lora import LoRAConfig, inject_lora, load_lora_state_dict |
|
|
| del model |
| torch.cuda.empty_cache() |
| base_sd = hf_to_align(load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16), qcfg) |
| amodel = q3.Qwen3ForCausalLM.from_state_dict(qcfg, base_sd, attn_implementation="sdpa", device="cuda").to(dev).eval() |
| del base_sd |
| inject_lora(amodel, LoRAConfig(r=a.r, alpha=a.alpha), adapter="sft", generator=torch.Generator().manual_seed(0)) |
| load_lora_state_dict(amodel, "sft", payload["lora"], strict=True) |
| alps = [] |
| with torch.no_grad(): |
| for r in vres["results"]: |
| ids = r["prompt_ids"] + r["token_ids"] |
| lp = torch.log_softmax(amodel(input_ids=torch.tensor([ids], device=dev))[0].float(), dim=-1) |
| n0 = len(r["prompt_ids"]) |
| alps.append(lp[n0 - 1: n0 - 1 + len(r["token_ids"])].cpu()) |
| m5 = compare_adapter_vs_merged(alps, mlps) |
| report["adapter_vs_merged"] = {"metrics": m5, "verdict": evaluate_adapter_vs_merged(m5)} |
| return report |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| r = sub.add_parser("run") |
| r.add_argument("--ckpt", required=True, help="a trainer checkpoint dir (micro_N) with optimizer_rank0.pt") |
| r.add_argument("--out-dir", default=os.path.join(WORK, "export")) |
| r.add_argument("--work", default=WORK) |
| r.add_argument("--r", type=int, default=None) |
| r.add_argument("--alpha", type=float, default=None) |
| r.add_argument("--keep", action="store_true") |
| a = ap.parse_args(argv) |
| rep = {} |
| try: |
| rep = run(a) |
| finally: |
| if not a.keep and os.path.isdir(a.out_dir): |
| shutil.rmtree(a.out_dir, ignore_errors=True) |
| rep["export_deleted"] = True |
| os.makedirs(a.work, exist_ok=True) |
| with open(os.path.join(a.work, "report.json"), "w") as f: |
| json.dump(rep, f, indent=2, default=str) |
| print(json.dumps(rep, indent=2, default=str)) |
| return 0 if export_ok(rep) else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
