File size: 4,976 Bytes
be2c293 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | #!/usr/bin/env python3
"""Check the MLX port against the torch reference (needs torch, transformers, the HF snapshots).
python test_parity.py --mlx ~/Desktop/Yue2-3B-MLX
Compares AR logits, NAR velocity and VAE decode on tiny inputs. Exits non-zero on mismatch.
"""
from __future__ import annotations
import argparse
from pathlib import Path
import mlx.core as mx
import numpy as np
import torch
from generate import CODEC_OFFSET, MUSIC_END, Tokenizer, token_prefix
from yue2_model import KVCache, load_model
from yue2_vae import load_vae
LATENT_START, LATENT_END, LATENT_PAD = 184621, 184622, 184623
def snapshot(repo):
from huggingface_hub import snapshot_download
return Path(snapshot_download(repo, allow_patterns=["*.json", "*.py", "model.safetensors", "qwen.tiktoken"]))
def rel_err(a, b):
a, b = np.asarray(a, np.float32), np.asarray(b, np.float32)
return float(np.abs(a - b).max() / (np.abs(b).max() + 1e-6))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--mlx", type=Path, default=Path.home() / "Desktop" / "Yue2-3B-MLX")
parser.add_argument("--ref-model", default="m-a-p/YuE2-3B")
parser.add_argument("--ref-vae", default="m-a-p/YuE2-Vae")
args = parser.parse_args()
from transformers import AutoModel, AutoModelForCausalLM
ref_dir = snapshot(args.ref_model) if not Path(args.ref_model).is_dir() else Path(args.ref_model)
vae_dir = snapshot(args.ref_vae) if not Path(args.ref_vae).is_dir() else Path(args.ref_vae)
tok = Tokenizer(args.mlx / "qwen.tiktoken")
prefix = token_prefix(tok, "pop, warm vocal", "[Verse]\nHello world\n", "off")
codec = [7, 300, 5000, 12345, 32000, 9, 9, 42]
ar_tokens = prefix + [c + CODEC_OFFSET for c in codec] + [MUSIC_END]
frames = len(codec)
rng = np.random.default_rng(0)
state = rng.standard_normal((frames, 64)).astype(np.float32)
raw_t = 0.3
# ── torch reference ──
ref = AutoModelForCausalLM.from_pretrained(ref_dir, trust_remote_code=True, torch_dtype=torch.bfloat16,
local_files_only=True).eval()
with torch.inference_mode():
ref_logits = ref(torch.tensor([ar_tokens]), use_cache=False).logits[0].float().numpy()
nar = [LATENT_START] + [LATENT_PAD] * frames + [LATENT_END]
tokens = torch.tensor([ar_tokens + nar])
ar_mask = torch.tensor([[True] * len(ar_tokens) + [False] * len(nar)])
nar_mask = ~ar_mask
content = torch.tensor([[False] * (len(ar_tokens) + 1) + [True] * frames + [False]])
ref_v = ref.nar_velocity(tokens, ar_mask, nar_mask, content, torch.tensor(state), raw_t).float().numpy()
del ref
ref_vae = AutoModel.from_pretrained(vae_dir, trust_remote_code=True, decoder_only=True, local_files_only=True).eval()
z = rng.standard_normal((1, 64, 40)).astype(np.float32)
with torch.inference_mode():
ref_audio = ref_vae.decode(torch.tensor(z))[0].T.numpy() # [L,2]
ref_tiled = ref_vae.decode_tiled(torch.tensor(z), core_frames=16, halo_frames=16)[0].T.numpy()
del ref_vae
# ── mlx ──
model = load_model(args.mlx)
caches = [KVCache() for _ in model.model.layers]
mlx_logits = np.array(model.ar_step(mx.array([ar_tokens]), caches, all_positions=True)[0].astype(mx.float32))
# incremental decode must equal the prefill logits for the same positions
caches2 = [KVCache() for _ in model.model.layers]
model.ar_step(mx.array([ar_tokens[:-3]]), caches2)
step_logits = [np.array(model.ar_step(mx.array([[t]]), caches2).astype(mx.float32))[0] for t in ar_tokens[-3:]]
ar_cache = model.nar_prefill(ar_tokens)
mlx_v = np.array(model.nar_velocity(mx.array(state).astype(mx.bfloat16), raw_t, ar_cache, len(ar_tokens)).astype(mx.float32))
vae = load_vae(args.mlx)
zc = mx.array(z[0].T) # [T,64]
mlx_audio = np.array(vae(zc[None])[0])
mlx_tiled = np.array(vae.decode_tiled(zc, core=16, halo=16))
checks = {
"ar_logits_rel_err": rel_err(mlx_logits, ref_logits),
"ar_argmax_agree": float((mlx_logits.argmax(-1) == ref_logits.argmax(-1)).mean()),
"ar_incremental_rel_err": rel_err(np.stack(step_logits), mlx_logits[-3:]),
"nar_velocity_rel_err": rel_err(mlx_v, ref_v),
"vae_rel_err": rel_err(mlx_audio, ref_audio),
"vae_tiled_rel_err": rel_err(mlx_tiled, ref_tiled),
"vae_len_ok": float(mlx_audio.shape == ref_audio.shape and mlx_tiled.shape == ref_tiled.shape),
}
for k, v in checks.items():
print(f"{k:26s} {v:.5f}")
ok = (checks["ar_logits_rel_err"] < 0.05 and checks["ar_argmax_agree"] > 0.9
and checks["ar_incremental_rel_err"] < 0.05 and checks["nar_velocity_rel_err"] < 0.05
and checks["vae_rel_err"] < 1e-3 and checks["vae_tiled_rel_err"] < 1e-3 and checks["vae_len_ok"])
print("PASS" if ok else "FAIL")
raise SystemExit(0 if ok else 1)
if __name__ == "__main__":
main()
|