214 lines
8.8 KiB
Python
214 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Export Qwen3-TTS talker KV-cache to ONNX with M-RoPE.
|
||
Following TTS_CALIBRATION_GUIDE.md exactly:
|
||
Inputs: [emb(1,1,1024), mask(1,1,1,200), cos(1,1,128), sin(1,1,128), 56×kv(1,8,199,128)]
|
||
Outputs: [hidden(1,1,1024), logits(1,1,3072), 56×kv(1,8,200,128)]
|
||
"""
|
||
import sys
|
||
sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages")
|
||
import warnings; warnings.filterwarnings("ignore")
|
||
import torch, torch.nn as nn, torch.nn.functional as F
|
||
import numpy as np, os, math
|
||
|
||
N_LAYERS = 28; DIM = 1024; N_HEADS = 16; N_KV_HEADS = 8; HEAD_DIM = 128
|
||
KV_LEN = 199; VOCAB = 3072; FFN_DIM = 3072
|
||
GQA_GROUPS = N_HEADS // N_KV_HEADS # 2
|
||
|
||
print("Loading talker weights...")
|
||
state = torch.load("/opt/Kazeia/models_qnn/qwen3-tts-export/qwen3_tts_talker.pth",
|
||
map_location="cpu", weights_only=False)
|
||
|
||
class RMSNorm(nn.Module):
|
||
def __init__(self, dim):
|
||
super().__init__()
|
||
self.weight = nn.Parameter(torch.ones(dim))
|
||
def forward(self, x):
|
||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.weight
|
||
|
||
def apply_interleaved_rope(x, cos, sin):
|
||
"""Apply interleaved RoPE: x[..., ::2] and x[..., 1::2] rotated by cos/sin."""
|
||
# x: [B, heads, 1, head_dim], cos/sin: [1, 1, 128] → use first 64 pairs
|
||
half = x.shape[-1] // 2
|
||
cos = cos[..., :half] # [1,1,64]
|
||
sin = sin[..., :half]
|
||
x1 = x[..., ::2] # [B, heads, 1, 64]
|
||
x2 = x[..., 1::2]
|
||
cos = cos.unsqueeze(1) # [1, 1, 1, 64]
|
||
sin = sin.unsqueeze(1)
|
||
o1 = x1 * cos - x2 * sin
|
||
o2 = x1 * sin + x2 * cos
|
||
return torch.stack([o1, o2], dim=-1).flatten(-2)
|
||
|
||
class TalkerKVWrapper(nn.Module):
|
||
def __init__(self, sd):
|
||
super().__init__()
|
||
self.norm = RMSNorm(DIM)
|
||
self.norm.weight.data = sd["norm.weight"]
|
||
self.lm_head = nn.Linear(DIM, VOCAB, bias=False)
|
||
self.lm_head.weight.data = sd["output.weight"]
|
||
|
||
self.layers = nn.ModuleList()
|
||
for i in range(N_LAYERS):
|
||
L = nn.Module()
|
||
L.attn_norm = RMSNorm(DIM); L.attn_norm.weight.data = sd[f"layers.{i}.attention_norm.weight"]
|
||
L.ffn_norm = RMSNorm(DIM); L.ffn_norm.weight.data = sd[f"layers.{i}.ffn_norm.weight"]
|
||
L.wq = nn.Linear(DIM, N_HEADS*HEAD_DIM, bias=False); L.wq.weight.data = sd[f"layers.{i}.attention.wq.weight"]
|
||
L.wk = nn.Linear(DIM, N_KV_HEADS*HEAD_DIM, bias=False); L.wk.weight.data = sd[f"layers.{i}.attention.wk.weight"]
|
||
L.wv = nn.Linear(DIM, N_KV_HEADS*HEAD_DIM, bias=False); L.wv.weight.data = sd[f"layers.{i}.attention.wv.weight"]
|
||
L.wo = nn.Linear(N_HEADS*HEAD_DIM, DIM, bias=False); L.wo.weight.data = sd[f"layers.{i}.attention.wo.weight"]
|
||
L.q_norm = RMSNorm(HEAD_DIM); L.q_norm.weight.data = sd[f"layers.{i}.attention.q_norm_fn.weight"]
|
||
L.k_norm = RMSNorm(HEAD_DIM); L.k_norm.weight.data = sd[f"layers.{i}.attention.k_norm_fn.weight"]
|
||
L.w1 = nn.Linear(DIM, FFN_DIM, bias=False); L.w1.weight.data = sd[f"layers.{i}.feed_forward.w1.weight"]
|
||
L.w2 = nn.Linear(FFN_DIM, DIM, bias=False); L.w2.weight.data = sd[f"layers.{i}.feed_forward.w2.weight"]
|
||
L.w3 = nn.Linear(DIM, FFN_DIM, bias=False); L.w3.weight.data = sd[f"layers.{i}.feed_forward.w3.weight"]
|
||
self.layers.append(L)
|
||
|
||
def forward(self, emb, mask, cos, sin, *kv_in):
|
||
"""
|
||
emb: [1, 1, 1024]
|
||
mask: [1, 1, 1, 200] (right-aligned, -10000 for masked positions)
|
||
cos: [1, 1, 128] (M-RoPE cos for current position)
|
||
sin: [1, 1, 128] (M-RoPE sin for current position)
|
||
kv_in: 28 × (k[1,8,199,128], v[1,8,199,128])
|
||
Returns: hidden[1,1,1024], logits[1,1,3072], 28 × (k[1,8,200,128], v[1,8,200,128])
|
||
"""
|
||
h = emb # [1, 1, 1024]
|
||
new_kvs = []
|
||
|
||
for i, L in enumerate(self.layers):
|
||
k_cache = kv_in[i * 2] # [1, 8, KV_LEN, 128]
|
||
v_cache = kv_in[i * 2 + 1] # [1, 8, KV_LEN, 128]
|
||
|
||
h_n = L.attn_norm(h)
|
||
q = L.wq(h_n).view(1, 1, N_HEADS, HEAD_DIM).transpose(1, 2) # [1,16,1,128]
|
||
k = L.wk(h_n).view(1, 1, N_KV_HEADS, HEAD_DIM).transpose(1, 2) # [1,8,1,128]
|
||
v = L.wv(h_n).view(1, 1, N_KV_HEADS, HEAD_DIM).transpose(1, 2) # [1,8,1,128]
|
||
|
||
# q_norm, k_norm BEFORE rope
|
||
q = L.q_norm(q)
|
||
k = L.k_norm(k)
|
||
|
||
# Apply M-RoPE (interleaved)
|
||
q = apply_interleaved_rope(q, cos, sin)
|
||
k = apply_interleaved_rope(k, cos, sin)
|
||
|
||
# Concat new k,v to cache → [1, 8, 200, 128]
|
||
k_full = torch.cat([k_cache, k], dim=2) # [1, 8, KV_LEN+1, 128]
|
||
v_full = torch.cat([v_cache, v], dim=2)
|
||
|
||
# GQA expand for attention
|
||
k_exp = k_full.repeat(1, GQA_GROUPS, 1, 1) # [1, 16, 200, 128]
|
||
v_exp = v_full.repeat(1, GQA_GROUPS, 1, 1)
|
||
|
||
# Attention
|
||
attn = torch.matmul(q, k_exp.transpose(-2, -1)) / math.sqrt(HEAD_DIM) # [1,16,1,200]
|
||
attn = attn + mask
|
||
attn = F.softmax(attn, dim=-1)
|
||
out = torch.matmul(attn, v_exp) # [1,16,1,128]
|
||
out = out.transpose(1, 2).contiguous().view(1, 1, N_HEADS * HEAD_DIM) # [1,1,2048]
|
||
h = h + L.wo(out)
|
||
|
||
# FFN (SwiGLU)
|
||
h_n2 = L.ffn_norm(h)
|
||
h = h + L.w2(F.silu(L.w1(h_n2)) * L.w3(h_n2))
|
||
|
||
new_kvs.extend([k_full, v_full])
|
||
|
||
h = self.norm(h)
|
||
logits = self.lm_head(h) # [1, 1, 3072]
|
||
|
||
return (h, logits, *new_kvs)
|
||
|
||
print("Building model...")
|
||
wrapper = TalkerKVWrapper(state)
|
||
wrapper.eval()
|
||
|
||
# Test with dummy inputs
|
||
print("Testing...")
|
||
dummy_emb = torch.randn(1, 1, DIM)
|
||
dummy_mask = torch.full((1, 1, 1, KV_LEN + 1), -10000.0)
|
||
dummy_mask[0, 0, 0, -1] = 0.0 # only current position visible
|
||
dummy_cos = torch.ones(1, 1, HEAD_DIM)
|
||
dummy_sin = torch.zeros(1, 1, HEAD_DIM)
|
||
kv_init = []
|
||
for _ in range(N_LAYERS):
|
||
kv_init.append(torch.zeros(1, N_KV_HEADS, KV_LEN, HEAD_DIM)) # k
|
||
kv_init.append(torch.zeros(1, N_KV_HEADS, KV_LEN, HEAD_DIM)) # v
|
||
|
||
with torch.no_grad():
|
||
out = wrapper(dummy_emb, dummy_mask, dummy_cos, dummy_sin, *kv_init)
|
||
|
||
hidden = out[0]
|
||
logits = out[1]
|
||
print(f"hidden: {hidden.shape}, logits: {logits.shape}")
|
||
print(f"new_k[0]: {out[2].shape}, new_v[0]: {out[3].shape}")
|
||
|
||
# Export ONNX
|
||
OUT = "/opt/Kazeia/models_qnn/talker_kv_cpu_mrope"
|
||
os.makedirs(OUT, exist_ok=True)
|
||
|
||
input_names = ["inputs_embeds", "attention_mask", "cos", "sin"]
|
||
for i in range(N_LAYERS):
|
||
input_names.extend([f"k_{i}_in", f"v_{i}_in"])
|
||
|
||
output_names = ["hidden", "logits"]
|
||
for i in range(N_LAYERS):
|
||
output_names.extend([f"k_{i}_out", f"v_{i}_out"])
|
||
|
||
all_inputs = (dummy_emb, dummy_mask, dummy_cos, dummy_sin, *kv_init)
|
||
|
||
print(f"\nExporting ONNX ({len(input_names)} inputs, {len(output_names)} outputs)...")
|
||
torch.onnx.export(
|
||
wrapper, all_inputs,
|
||
f"{OUT}/model.onnx",
|
||
input_names=input_names,
|
||
output_names=output_names,
|
||
opset_version=17,
|
||
)
|
||
sz = os.path.getsize(f"{OUT}/model.onnx")
|
||
print(f"Saved: {OUT}/model.onnx ({sz/1024/1024:.0f} MB)")
|
||
|
||
# Pre-compute M-RoPE cos/sin for all 200 positions and save
|
||
# For TTS, all 3 mrope axes use the same position, so M-RoPE = standard RoPE
|
||
# inv_freq = 1 / (theta^(2i/128)) for i in 0..63
|
||
print("\nPre-computing RoPE tables (standard, theta=1e6)...")
|
||
theta = 1000000.0
|
||
inv_freq = 1.0 / (theta ** (np.arange(0, HEAD_DIM, 2, dtype=np.float32) / HEAD_DIM)) # [64]
|
||
positions = np.arange(KV_LEN + 1, dtype=np.float32) # [200]
|
||
angles = np.outer(positions, inv_freq) # [200, 64]
|
||
rope_cos = np.cos(angles).astype(np.float32) # [200, 64]
|
||
rope_sin = np.sin(angles).astype(np.float32)
|
||
# Expand to head_dim=128 for interleaved rope: each freq pair applies to (x[2i], x[2i+1])
|
||
rope_cos_full = np.repeat(rope_cos, 2, axis=1) # [200, 128]
|
||
rope_sin_full = np.repeat(rope_sin, 2, axis=1)
|
||
np.save(f"{OUT}/talker_rotary_cos.npy", rope_cos_full)
|
||
np.save(f"{OUT}/talker_rotary_sin.npy", rope_sin_full)
|
||
print(f"Rotary tables: {rope_cos_full.shape} saved")
|
||
|
||
# Validate: run a quick prefill step and check output
|
||
print("\n=== Validation ===")
|
||
import onnxruntime as ort
|
||
sess = ort.InferenceSession(f"{OUT}/model.onnx")
|
||
print(f"ONNX inputs: {[i.name for i in sess.get_inputs()[:4]]}...")
|
||
|
||
mask = np.full((1, 1, 1, KV_LEN + 1), -10000.0, dtype=np.float32)
|
||
mask[0, 0, 0, -1] = 0.0
|
||
cos_0 = rope_cos_full[0:1].reshape(1, 1, HEAD_DIM)
|
||
sin_0 = rope_sin_full[0:1].reshape(1, 1, HEAD_DIM)
|
||
|
||
inputs = {
|
||
"inputs_embeds": np.random.randn(1, 1, DIM).astype(np.float32),
|
||
"attention_mask": mask,
|
||
"cos": cos_0,
|
||
"sin": sin_0,
|
||
}
|
||
for i in range(N_LAYERS):
|
||
inputs[f"k_{i}_in"] = np.zeros((1, N_KV_HEADS, KV_LEN, HEAD_DIM), dtype=np.float32)
|
||
inputs[f"v_{i}_in"] = np.zeros((1, N_KV_HEADS, KV_LEN, HEAD_DIM), dtype=np.float32)
|
||
|
||
out = sess.run(None, inputs)
|
||
print(f"ONNX hidden: {out[0].shape}, logits: {out[1].shape}")
|
||
print(f"ONNX k_0_out: {out[2].shape}")
|
||
print("Validation OK!")
|