MoE-LoRA-converter/moe_lora_convert_validate.py

"""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)