Pascal
7b6ed4f6db
The fused mode appended the codec stream tail to the predictor frame graph, so one compute produced both a frame's codes and its 80 ms of audio with no host round trip. The experiment applied to max_batch 1 with a streaming synthesis only, it cost throughput against the buffered flush that stays the default, and it kept a second frame graph, its ring inputs and an init flag alive for that single case. It is not worth keeping. Remove the tail helpers, the fused graph of CodePredGraphSet, the codec_fused field of qt_init_params, the --codec-fused flag of both tools and the harness switch that exercised it. The predictor frame unroll and the in graph sampler are untouched.
398 lines
16 KiB
Python
398 lines
16 KiB
Python
"""Shared helpers for the qwentts.cpp cossim debug scripts.
|
|
|
|
Provides Philox uniform stream, dump load and save, install_hooks for the
|
|
talker submodel, the standard stage list and the metric helpers used by
|
|
debug-base / debug-tts / debug-customvoice / debug-clone cossim scripts.
|
|
|
|
Importing this module patches sys.path so qwen_tts upstream loads without
|
|
the V1 25Hz tokenizer (sox dependency stubbed out), and forces TF32 off on
|
|
every torch CUDA matmul path so Python results stay bit comparable across
|
|
runs and across machines.
|
|
"""
|
|
|
|
import os
|
|
import struct
|
|
import sys
|
|
import types
|
|
|
|
os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
|
|
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
|
|
torch.set_float32_matmul_precision("highest")
|
|
|
|
# Stub the V1 25Hz tokenizer so qwen_tts.core imports without sox.
|
|
UPSTREAM_ROOT = "../../Qwen3-TTS"
|
|
sys.path.insert(0, UPSTREAM_ROOT)
|
|
|
|
class _StubV1Config:
|
|
pass
|
|
|
|
class _StubV1Model:
|
|
pass
|
|
|
|
_qwen_pkg = types.ModuleType("qwen_tts")
|
|
_qwen_pkg.__path__ = [os.path.join(UPSTREAM_ROOT, "qwen_tts")]
|
|
sys.modules["qwen_tts"] = _qwen_pkg
|
|
|
|
_core_pkg = types.ModuleType("qwen_tts.core")
|
|
_core_pkg.__path__ = [os.path.join(UPSTREAM_ROOT, "qwen_tts", "core")]
|
|
# Inject the stubbed core module before any submodule import so the real
|
|
# qwen_tts/core/__init__.py never runs : it pulls the V1 25Hz tokenizer that
|
|
# imports whisper_encoder, which prints a flash-attn warning at module load.
|
|
sys.modules["qwen_tts.core"] = _core_pkg
|
|
from qwen_tts.core.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
|
|
from qwen_tts.core.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Model
|
|
_core_pkg.Qwen3TTSTokenizerV1Config = _StubV1Config
|
|
_core_pkg.Qwen3TTSTokenizerV1Model = _StubV1Model
|
|
_core_pkg.Qwen3TTSTokenizerV2Config = Qwen3TTSTokenizerV2Config
|
|
_core_pkg.Qwen3TTSTokenizerV2Model = Qwen3TTSTokenizerV2Model
|
|
|
|
from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration
|
|
from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig
|
|
from qwen_tts.core.models.processing_qwen3_tts import Qwen3TTSProcessor
|
|
from transformers import AutoConfig, AutoModel, AutoProcessor
|
|
from transformers.utils import logging as hf_logging
|
|
|
|
# Silence the GenerationConfig validator that warns "flags are not valid
|
|
# and may be ignored" for temperature / top_k / top_p when do_sample=False.
|
|
# Those flags ride along inside the checkpoint generation_config and the
|
|
# greedy path drops them on purpose, the warning is just noise here.
|
|
hf_logging.set_verbosity_error()
|
|
|
|
# Register the Qwen3-TTS classes once per process. Calling twice raises a
|
|
# ValueError inside transformers, hence the guard.
|
|
_REGISTERED = {"done": False}
|
|
|
|
def register_qwen3_tts():
|
|
if _REGISTERED["done"]:
|
|
return
|
|
AutoConfig.register("qwen3_tts", Qwen3TTSConfig)
|
|
AutoModel.register(Qwen3TTSConfig, Qwen3TTSForConditionalGeneration)
|
|
AutoProcessor.register(Qwen3TTSConfig, Qwen3TTSProcessor)
|
|
_REGISTERED["done"] = True
|
|
|
|
# Path to the C++ qwen-tts binary, relative to tests/.
|
|
BIN = "../build/qwen-tts"
|
|
|
|
# Standard stage list shared by every cossim script. Mode-specific scripts
|
|
# may extend this list before iterating (eg. clone adds SpeakerEmb / RefCodes).
|
|
STAGES_STANDARD = [
|
|
("Embed", "talker-input-embed.bin"),
|
|
("TrailingText", "trailing-text-hidden.bin"),
|
|
("TTSPadEmbed", "tts-pad-embed.bin"),
|
|
("L0", "talker-hidden-prefill-l0.bin"),
|
|
("L7", "talker-hidden-prefill-l7.bin"),
|
|
("L14", "talker-hidden-prefill-l14.bin"),
|
|
("L21", "talker-hidden-prefill-l21.bin"),
|
|
("L27", "talker-hidden-prefill-l27.bin"),
|
|
("Final", "talker-hidden-prefill-final.bin"),
|
|
("Logits", "talker-logits-prefill.bin"),
|
|
("NextEmbStep0", "next-emb-step0.bin"),
|
|
("TalkerHiddenStep1", "talker-hidden-step1.bin"),
|
|
]
|
|
|
|
# Philox4x32-10 mirror of src/philox.h. Returns the same float u that
|
|
# philox_uniform_fill(seed, subseq, ctr_lo=0) gives for n=1.
|
|
PHILOX_M0 = np.uint32(0xD2511F53)
|
|
PHILOX_M1 = np.uint32(0xCD9E8D57)
|
|
PHILOX_W0 = np.uint32(0x9E3779B9)
|
|
PHILOX_W1 = np.uint32(0xBB67AE85)
|
|
TWO_POW32_INV = np.float32(2.3283064365386963e-10)
|
|
|
|
def _mulhilo32(a, b):
|
|
p = np.uint64(a) * np.uint64(b)
|
|
lo = np.uint32(p & np.uint64(0xFFFFFFFF))
|
|
hi = np.uint32(p >> np.uint64(32))
|
|
return hi, lo
|
|
|
|
def _philox_round(ctr, k0, k1):
|
|
hi0, lo0 = _mulhilo32(PHILOX_M0, ctr[0])
|
|
hi1, lo1 = _mulhilo32(PHILOX_M1, ctr[2])
|
|
return (np.uint32(hi1 ^ ctr[1] ^ k0),
|
|
np.uint32(lo1),
|
|
np.uint32(hi0 ^ ctr[3] ^ k1),
|
|
np.uint32(lo0))
|
|
|
|
def _philox4x32_10(ctr, k0, k1):
|
|
mask = np.uint64(0xFFFFFFFF)
|
|
for _ in range(9):
|
|
ctr = _philox_round(ctr, k0, k1)
|
|
k0 = np.uint32((np.uint64(k0) + np.uint64(PHILOX_W0)) & mask)
|
|
k1 = np.uint32((np.uint64(k1) + np.uint64(PHILOX_W1)) & mask)
|
|
ctr = _philox_round(ctr, k0, k1)
|
|
return ctr
|
|
|
|
def philox_uniform(seed, subseq, ctr_lo=0):
|
|
slo = np.uint32(np.uint64(seed) & np.uint64(0xFFFFFFFF))
|
|
shi = np.uint32(np.uint64(seed) >> np.uint64(32))
|
|
ctr = (np.uint32(ctr_lo),
|
|
np.uint32(0),
|
|
np.uint32(np.uint64(subseq) & np.uint64(0xFFFFFFFF)),
|
|
np.uint32(np.uint64(subseq) >> np.uint64(32)))
|
|
r = _philox4x32_10(ctr, slo, shi)
|
|
return (np.float32(r[0]) + np.float32(0.5)) * TWO_POW32_INV
|
|
|
|
# Globals advanced exactly once per multinomial sample, mirroring the C++
|
|
# side which advances subseq_counter at every sample_top_k_p call.
|
|
_subseq_counter = [0]
|
|
_seed = [42]
|
|
_trace_samples = [False]
|
|
|
|
def reset_philox(seed):
|
|
_subseq_counter[0] = 0
|
|
_seed[0] = int(seed)
|
|
|
|
def set_trace(flag):
|
|
_trace_samples[0] = bool(flag)
|
|
|
|
def patched_multinomial(input, num_samples, replacement=False, generator=None, out=None):
|
|
"""Drop in replacement for torch.multinomial(num_samples=1) that pulls
|
|
the uniform draw from our Philox stream and walks the F32 cumulative
|
|
sum the same way src/sampling.h does."""
|
|
assert num_samples == 1, "patched_multinomial only handles num_samples=1"
|
|
probs = input
|
|
if probs.dim() == 1:
|
|
probs = probs.unsqueeze(0)
|
|
bsz, vocab = probs.shape
|
|
out_ids = torch.zeros((bsz, 1), dtype=torch.long, device=probs.device)
|
|
for b in range(bsz):
|
|
u = philox_uniform(_seed[0], _subseq_counter[0], 0)
|
|
seq = _subseq_counter[0]
|
|
_subseq_counter[0] += 1
|
|
row = probs[b].to(torch.float32).cpu().numpy()
|
|
s = float(row.sum())
|
|
# The C++ sampler draws u in [0, 1) and compares against acc/sum
|
|
# implicitly via acc >= u*sum. We replicate that exact arithmetic.
|
|
target = float(u) * s
|
|
acc = 0.0
|
|
idx = vocab - 1
|
|
for i in range(vocab):
|
|
acc += float(row[i])
|
|
if acc >= target:
|
|
idx = i
|
|
break
|
|
out_ids[b, 0] = idx
|
|
if _trace_samples[0] and seq < 32:
|
|
print(f"[Sample-PY] subseq={seq} u={float(u):.10f} idx={idx} top_prob={float(row.max()):.6f}")
|
|
if input.dim() == 1:
|
|
return out_ids.squeeze(0)
|
|
return out_ids
|
|
|
|
def ensure_dir(path):
|
|
os.makedirs(path, exist_ok=True)
|
|
|
|
def save_dump(path, data):
|
|
if isinstance(data, torch.Tensor):
|
|
data = data.detach().to(torch.float32).cpu().numpy()
|
|
data = np.ascontiguousarray(data.astype(np.float32))
|
|
shape = data.shape
|
|
with open(path, "wb") as f:
|
|
f.write(struct.pack("i", len(shape)))
|
|
for s in shape:
|
|
f.write(struct.pack("i", s))
|
|
f.write(data.tobytes())
|
|
|
|
def save_dump_i32(path, data):
|
|
if isinstance(data, torch.Tensor):
|
|
data = data.detach().to(torch.int64).cpu().numpy()
|
|
data = np.ascontiguousarray(data.astype(np.int64))
|
|
shape = data.shape
|
|
fdata = data.astype(np.float32)
|
|
with open(path, "wb") as f:
|
|
f.write(struct.pack("i", len(shape)))
|
|
for s in shape:
|
|
f.write(struct.pack("i", s))
|
|
f.write(fdata.tobytes())
|
|
|
|
def load_dump(path):
|
|
raw = np.fromfile(path, dtype=np.uint8)
|
|
ndim = int(np.frombuffer(raw[0:4], dtype=np.int32)[0])
|
|
shape = tuple(int(x) for x in np.frombuffer(raw[4:4 + 4 * ndim], dtype=np.int32))
|
|
body = np.frombuffer(raw[4 + 4 * ndim:], dtype=np.float32)
|
|
return body.reshape(shape), shape
|
|
|
|
def cos(a, b):
|
|
a = a.astype(np.float64).ravel()
|
|
b = b.astype(np.float64).ravel()
|
|
n = min(len(a), len(b))
|
|
a, b = a[:n], b[:n]
|
|
d = float(np.linalg.norm(a) * np.linalg.norm(b))
|
|
return float(np.dot(a, b) / d) if d > 1e-10 else 0.0
|
|
|
|
def stft_cos(a, b, win=2048, hop=512):
|
|
a = a.astype(np.float64).ravel()
|
|
b = b.astype(np.float64).ravel()
|
|
n = min(len(a), len(b))
|
|
a, b = a[:n], b[:n]
|
|
window = np.hanning(win)
|
|
frames = (n - win) // hop + 1
|
|
if frames <= 0:
|
|
return 0.0
|
|
sa = np.zeros((frames, win // 2 + 1))
|
|
sb = np.zeros((frames, win // 2 + 1))
|
|
for i in range(frames):
|
|
s = i * hop
|
|
sa[i] = np.abs(np.fft.rfft(a[s:s + win] * window))
|
|
sb[i] = np.abs(np.fft.rfft(b[s:s + win] * window))
|
|
return cos(sa.ravel(), sb.ravel())
|
|
|
|
def install_hooks(model, dump_dir, bisect_layers=(0, 7, 14, 21, 27)):
|
|
"""Capture every intermediate tensor we can pin against the C++ side.
|
|
Layout : bisection layers, final norm, prefill logits, the input embed
|
|
fed to the talker and the trailing-text overlay tensors that drive the
|
|
next-token embedding sums during generation. Counters track how many
|
|
times the talker submodel has run so step 1 (first single-token forward
|
|
after prefill) gets its own dump."""
|
|
seen_prefill = {"done": False}
|
|
seen_codes = {"done": False}
|
|
# talker_step counts how many times talker_model.forward has been called
|
|
# after the prefill. Prefill itself is recorded as 0, the first single
|
|
# token forward is 1, and so on. Single token forwards are detected by
|
|
# inputs_embeds.shape[1] == 1 in the pre hook.
|
|
talker_step = {"n": 0}
|
|
|
|
talker_model = model.talker.model
|
|
talker_lm = model.talker
|
|
|
|
seen_layers = {idx: False for idx in bisect_layers}
|
|
def make_layer_hook(layer_idx):
|
|
def hook(module, inputs, output):
|
|
if seen_layers[layer_idx]:
|
|
return
|
|
h = output[0] if isinstance(output, tuple) else output
|
|
if h.dim() == 3:
|
|
save_dump(os.path.join(dump_dir, f"talker-hidden-prefill-l{layer_idx}.bin"), h[0])
|
|
seen_layers[layer_idx] = True
|
|
return hook
|
|
for layer_idx in bisect_layers:
|
|
talker_model.layers[layer_idx].register_forward_hook(make_layer_hook(layer_idx))
|
|
|
|
seen_norm = {"done": False}
|
|
def norm_hook(module, inputs, output):
|
|
if seen_norm["done"]:
|
|
return
|
|
if output.dim() == 3 and output.shape[1] > 1:
|
|
save_dump(os.path.join(dump_dir, "talker-hidden-prefill-final.bin"), output[0])
|
|
seen_norm["done"] = True
|
|
talker_model.norm.register_forward_hook(norm_hook)
|
|
|
|
# Pre-hook on the talker submodel : sees inputs_embeds whether the outer
|
|
# talker forward was invoked with input_ids (single token step) or
|
|
# inputs_embeds (prefill). The submodel always receives inputs_embeds
|
|
# because the wrapper rebuilds it before calling self.model.
|
|
def talker_model_pre_hook(module, args, kwargs):
|
|
ie = kwargs.get("inputs_embeds", None)
|
|
if ie is None:
|
|
return
|
|
if ie.dim() != 3:
|
|
return
|
|
if ie.shape[1] > 1:
|
|
return
|
|
if talker_step["n"] == 0:
|
|
save_dump(os.path.join(dump_dir, "next-emb-step0.bin"), ie[0, 0])
|
|
talker_step["n"] += 1
|
|
talker_model.register_forward_pre_hook(talker_model_pre_hook, with_kwargs=True)
|
|
|
|
# Post-hook on the talker submodel : captures last_hidden_state at step
|
|
# 1 (first single token forward). That tensor is what feeds the code
|
|
# predictor at step 1, so any drift between Python and C++ tells us the
|
|
# next-emb-step0 changed the talker forward result.
|
|
talker_post_step = {"n": 0}
|
|
def talker_model_post_hook(module, inputs, output):
|
|
last = output.last_hidden_state if hasattr(output, "last_hidden_state") else output[0]
|
|
if last.dim() != 3 or last.shape[1] != 1:
|
|
return
|
|
if talker_post_step["n"] == 0:
|
|
save_dump(os.path.join(dump_dir, "talker-hidden-step1.bin"), last[0, -1])
|
|
talker_post_step["n"] += 1
|
|
talker_model.register_forward_hook(talker_model_post_hook)
|
|
|
|
# Talker LM wrapper hook : captures the prefill input embed (the talker
|
|
# codec_embedding sum + text projection that mirrors what
|
|
# prompt_builder_build produces in C++), the prefill logits, and the
|
|
# trailing_text_hidden / tts_pad_embed overlay tensors carried by the
|
|
# output dataclass at every step (we only dump them once).
|
|
seen_overlay = {"done": False}
|
|
orig_talker_forward = talker_lm.forward
|
|
def hooked_talker_forward(*args, **kwargs):
|
|
inputs_embeds = kwargs.get("inputs_embeds", None)
|
|
if (inputs_embeds is not None and inputs_embeds.dim() == 3
|
|
and inputs_embeds.shape[1] > 1 and not seen_prefill["done"]):
|
|
save_dump(os.path.join(dump_dir, "talker-input-embed.bin"), inputs_embeds[0])
|
|
seen_prefill["done"] = True
|
|
out = orig_talker_forward(*args, **kwargs)
|
|
if (out is not None and getattr(out, "logits", None) is not None
|
|
and not seen_codes["done"]):
|
|
logits = out.logits
|
|
if logits.dim() == 3 and logits.shape[1] > 1:
|
|
save_dump(os.path.join(dump_dir, "talker-logits-prefill.bin"), logits[0, -1])
|
|
seen_codes["done"] = True
|
|
if (out is not None and not seen_overlay["done"]
|
|
and getattr(out, "trailing_text_hidden", None) is not None
|
|
and getattr(out, "tts_pad_embed", None) is not None):
|
|
tth = out.trailing_text_hidden
|
|
tpe = out.tts_pad_embed
|
|
if tth.dim() == 3 and tpe.dim() == 3:
|
|
save_dump(os.path.join(dump_dir, "trailing-text-hidden.bin"), tth[0])
|
|
save_dump(os.path.join(dump_dir, "tts-pad-embed.bin"), tpe[0, 0])
|
|
seen_overlay["done"] = True
|
|
return out
|
|
talker_lm.forward = hooked_talker_forward
|
|
|
|
def pair(name, dump_cpp, dump_pt):
|
|
a, _ = load_dump(os.path.join(dump_cpp, name))
|
|
b, _ = load_dump(os.path.join(dump_pt, name))
|
|
return a, b
|
|
|
|
def metric(a, b):
|
|
n = min(a.size, b.size)
|
|
af = a.astype(np.float64).ravel()[:n]
|
|
bf = b.astype(np.float64).ravel()[:n]
|
|
d = np.abs(af - bf)
|
|
nrm_a = float(np.linalg.norm(af))
|
|
nrm_b = float(np.linalg.norm(bf))
|
|
c = float(np.dot(af, bf) / (nrm_a * nrm_b)) if nrm_a > 1e-10 and nrm_b > 1e-10 else 0.0
|
|
return c, float(d.max()), float(d.mean())
|
|
|
|
def compare_stages(stages, dump_cpp, dump_pt):
|
|
"""Iterate the stages list and print one line per pair. Skips silently
|
|
when a dump file is missing (eg. a mode that does not produce a given
|
|
intermediate)."""
|
|
for label, name in stages:
|
|
try:
|
|
a, b = pair(name, dump_cpp, dump_pt)
|
|
except FileNotFoundError:
|
|
print(f"[Cossim] {label} skipped (missing dump)")
|
|
continue
|
|
c, mx, mn = metric(a, b)
|
|
print(f"[Cossim] {label} cos: {c:.6f} max: {mx:.4e} mean: {mn:.4e}")
|
|
|
|
def compare_exact_i32(name, dump_cpp, dump_pt, label):
|
|
"""Compare two int dumps stored as f32 (the encoding path used by both
|
|
save_dump_i32 in Python and debug_dump_i32_as_f32 in C++). Prints an
|
|
exact match percentage. Returns the percentage as a float."""
|
|
a, b = pair(name, dump_cpp, dump_pt)
|
|
ai = a.astype(np.int64).ravel()
|
|
bi = b.astype(np.int64).ravel()
|
|
n = min(ai.size, bi.size)
|
|
pct = 100.0 * float(np.mean(ai[:n] == bi[:n]))
|
|
print(f"[Cossim] {label} exact: {pct:.2f}% ({n} values)")
|
|
return pct
|
|
|
|
# Greedy generation kwargs shared by every cossim script. do_sample=False
|
|
# alone selects argmax, top_k / top_p / temperature are intentionally
|
|
# omitted because GenerationConfig flags them as "not valid" warnings when
|
|
# do_sample=False. The subtalker_* keys are custom kwargs forwarded to the
|
|
# talker forward, the talker validator is bypassed by the script main.
|
|
GEN_KWARGS_GREEDY = dict(
|
|
do_sample = False,
|
|
subtalker_dosample = False,
|
|
repetition_penalty = 1.0,
|
|
)
|