"""Validates moe_lora_convert_vllm.py's output. 1. Round-trip: reassemble the fused tensors from the converted per-expert set and compare against the originals. Must be bit-exact. 2. Per-expert delta: for a handful of experts, compute scaling * B_e @ A_e from the CONVERTED tensors and compare against Unsloth's REAL production merge function (unsloth_zoo/saving_utils.py -- the corrected ground truth, NOT vanilla PEFT's get_delta_weight -- see the README's "wrong turn" note). 3. Key coverage: assert exactly 18,432 (layer, expert, projection) pairs, every one present once, no orphans. NOTE: SRC/DST below are the paths from the run this was developed against -- point them at your own source adapter and its moe_lora_convert_vllm.py output before running. Requires a CUDA GPU (the real merge functions run on-device) and `unsloth` installed (imported for its side-effecting patches before `unsloth_zoo.saving_utils` is usable). """ import json import re import sys import torch import unsloth # noqa: F401 from safetensors import safe_open from unsloth_zoo.saving_utils import LoraStats, _merge_moe_gate_expert, _merge_moe_up_expert, _merge_moe_down_proj_expert SRC = "outputs/ladder_cpt/voxday_nuttall_joint_rung90000_90000/checkpoint-290" DST = "outputs/ladder_cpt/voxday_nuttall_joint_rung90000_90000/checkpoint-290_vllm_converted" E, r = 128, 16 with open(f"{DST}/adapter_config.json") as fh: cfg = json.load(fh) scaling = cfg["lora_alpha"] / cfg["r"] print(f"scaling={scaling}") ok = True # ---------- 1. Round-trip ---------- print("\n=== 1. Round-trip (reassemble converted -> compare to original, bit-exact) ===") with safe_open(f"{SRC}/adapter_model.safetensors", framework="pt") as fsrc, \ safe_open(f"{DST}/adapter_model.safetensors", framework="pt") as fdst: src_keys = list(fsrc.keys()) expert_prefix_re = re.compile(r"^base_model\.model\.model\.layers\.(\d+)\.mlp\.experts\.(?:base_layer\.)?lora_A\.weight$") layers = sorted({int(m.group(1)) for k in src_keys if (m := expert_prefix_re.match(k))}) layers_to_check = [0, 1, 23, 46, 47] for L in layers_to_check: prefix = f"base_model.model.model.layers.{L}.mlp.experts." a_keys = [k for k in src_keys if k.startswith(prefix) and k.endswith("lora_A.weight")] pairs = [(k, fsrc.get_tensor(k), fsrc.get_tensor(k.replace("lora_A.weight", "lora_B.weight"))) for k in a_keys] pairs.sort(key=lambda p: p[1].shape[-1], reverse=True) (_, gate_up_A, gate_up_B), (_, down_A, down_B) = pairs hidden = gate_up_A.shape[1] inter = gate_up_B.shape[0] // 2 # reassemble gate_up_A, gate_up_B, down_A, down_B from converted per-expert tensors recon_gate_up_A = torch.zeros_like(gate_up_A) recon_gate_up_B = torch.zeros_like(gate_up_B) recon_down_A = torch.zeros_like(down_A) recon_down_B = torch.zeros_like(down_B) for e in range(E): lo, hi = e * r, (e + 1) * r base = f"base_model.model.model.layers.{L}.mlp.experts.{e}" gA = fdst.get_tensor(f"{base}.gate_proj.lora_A.weight") gB = fdst.get_tensor(f"{base}.gate_proj.lora_B.weight") uA = fdst.get_tensor(f"{base}.up_proj.lora_A.weight") uB = fdst.get_tensor(f"{base}.up_proj.lora_B.weight") dA = fdst.get_tensor(f"{base}.down_proj.lora_A.weight") dB = fdst.get_tensor(f"{base}.down_proj.lora_B.weight") assert torch.equal(gA, uA), f"layer {L} expert {e}: gate_proj/up_proj lora_A should be identical (shared)" recon_gate_up_A[lo:hi, :] = gA recon_gate_up_B[:inter, lo:hi] = gB recon_gate_up_B[inter:, lo:hi] = uB recon_down_A[lo:hi, :] = dA recon_down_B[:, lo:hi] = dB g_a_ok = torch.equal(recon_gate_up_A, gate_up_A) g_b_ok = torch.equal(recon_gate_up_B, gate_up_B) d_a_ok = torch.equal(recon_down_A, down_A) d_b_ok = torch.equal(recon_down_B, down_B) all_ok = g_a_ok and g_b_ok and d_a_ok and d_b_ok ok = ok and all_ok print(f" layer {L}: gate_up_A={'OK' if g_a_ok else 'MISMATCH'} gate_up_B={'OK' if g_b_ok else 'MISMATCH'} " f"down_A={'OK' if d_a_ok else 'MISMATCH'} down_B={'OK' if d_b_ok else 'MISMATCH'}") # ---------- 2. Per-expert delta match against REAL Unsloth merge fn ---------- print("\n=== 2. Per-expert delta vs Unsloth's REAL production merge fn (saving_utils.py) ===") with safe_open(f"{SRC}/adapter_model.safetensors", framework="pt") as fsrc, \ safe_open(f"{DST}/adapter_model.safetensors", framework="pt") as fdst: for L in [0, 23, 47]: prefix = f"base_model.model.model.layers.{L}.mlp.experts." a_keys = [k for k in fsrc.keys() if k.startswith(prefix) and k.endswith("lora_A.weight")] pairs = [(k, fsrc.get_tensor(k), fsrc.get_tensor(k.replace("lora_A.weight", "lora_B.weight"))) for k in a_keys] pairs.sort(key=lambda p: p[1].shape[-1], reverse=True) (_, gate_up_A, gate_up_B), (_, down_A, down_B) = pairs hidden = gate_up_A.shape[1] inter = gate_up_B.shape[0] // 2 stats_gate_up = LoraStats(module=None, lora_A=gate_up_A.cuda(), lora_B=gate_up_B.cuda(), alpha=scaling) stats_down = LoraStats(module=None, lora_A=down_A.cuda(), lora_B=down_B.cuda(), alpha=scaling) for e in [0, 1, 63, 127]: base = f"base_model.model.model.layers.{L}.mlp.experts.{e}" gA = fdst.get_tensor(f"{base}.gate_proj.lora_A.weight") gB = fdst.get_tensor(f"{base}.gate_proj.lora_B.weight") uB = fdst.get_tensor(f"{base}.up_proj.lora_B.weight") dA = fdst.get_tensor(f"{base}.down_proj.lora_A.weight") dB = fdst.get_tensor(f"{base}.down_proj.lora_B.weight") from_converted_gate = scaling * (gB @ gA) from_converted_up = scaling * (uB @ gA) from_converted_down = scaling * (dB @ dA) zero_gate = torch.zeros(inter, hidden).cuda() zero_up = torch.zeros(inter, hidden).cuda() zero_down = torch.zeros(hidden, inter).cuda() real_gate = (_merge_moe_gate_expert(zero_gate, stats_gate_up, e, E, torch.float32) - zero_gate).cpu() real_up = (_merge_moe_up_expert(zero_up, stats_gate_up, e, E, torch.float32) - zero_up).cpu() real_down = (_merge_moe_down_proj_expert(zero_down, stats_down, e, E, torch.float32) - zero_down).cpu() def rel_or_abs_err(conv, real, tag): # Some experts are never routed during training and stay at LoRA's # standard zero-init for lora_B -> delta is exactly zero for both # sides. Relative error is 0/0 there; check absolute error instead. real_norm = torch.linalg.norm(real) if real_norm < 1e-12: abs_err = torch.linalg.norm(conv).item() return abs_err, abs_err < 1e-9, f"{tag}: untrained expert (both~0), abs_err={abs_err:.1e}" err = (torch.linalg.norm(conv - real) / real_norm).item() return err, err < 1e-3, f"{tag} rel_err={err:.2e}" _, gate_pass, gate_msg = rel_or_abs_err(from_converted_gate, real_gate, "gate") _, up_pass, up_msg = rel_or_abs_err(from_converted_up, real_up, "up") _, down_pass, down_msg = rel_or_abs_err(from_converted_down, real_down, "down") pass_ = gate_pass and up_pass and down_pass ok = ok and pass_ print(f" layer {L} expert {e:3d}: {gate_msg} {up_msg} {down_msg} {'PASS' if pass_ else 'FAIL'}") # ---------- 3. Key coverage ---------- print("\n=== 3. Key coverage ===") with safe_open(f"{DST}/adapter_model.safetensors", framework="pt") as fdst: dst_keys = list(fdst.keys()) expert_key_re = re.compile(r"^base_model\.model\.model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.lora_[AB]\.weight$") found = set() for k in dst_keys: m = expert_key_re.match(k) if m: L, e, proj = int(m.group(1)), int(m.group(2)), m.group(3) found.add((L, e, proj)) expected = {(L, e, proj) for L in range(48) for e in range(E) for proj in ("gate_proj", "up_proj", "down_proj")} missing = expected - found extra = found - expected n_ok = len(found) == 18432 == len(expected) and not missing and not extra ok = ok and n_ok print(f" found {len(found)} (layer,expert,proj) triples (expect 18432): {'PASS' if n_ok else 'FAIL'}") if missing: print(f" MISSING (first 5): {list(missing)[:5]}") if extra: print(f" EXTRA (first 5): {list(extra)[:5]}") print(f"\nOVERALL: {'ALL CHECKS PASSED' if ok else 'SOME CHECKS FAILED'}") print("VALIDATE_DONE") sys.exit(0 if ok else 1)