s2s(voicechat-11b): Phases 2b+2c — TTS decoder + codec on CoreML, all components now measured - #87
Open
Alex-Wengg wants to merge 2 commits into
Open
s2s(voicechat-11b): Phases 2b+2c — TTS decoder + codec on CoreML, all components now measured#87Alex-Wengg wants to merge 2 commits into
Alex-Wengg wants to merge 2 commits into
Conversation
RVQ-VAE decoder at 12.5 Hz -> 22.05 kHz converted with a flexible frame count: the conv stack (ConvT 512->1536 k9s9 -> ConvNeXt x3 -> ConvT ->768 k7s7 -> ConvNeXt x3 -> ConvT ->384 k7s7 -> ConvNeXt x3 -> Conv ->18) runs in CoreML; the PRVQ code->latent lookup (31 embedding sums, codebooks exported as codec_prvq_mus.npy and verified identical to the TTS model's rvq_embs buffer) and the 16-point complex iSTFT tail (n_fft 16, hop 4, hann, overlap-add + envelope normalization, numpy port of spec_to_wav) run host-side. Parity on real audio (sample_general.wav resampled to 22.05k, encoded through the codec's own encoder): fp32 wav max|d| 7.5e-07 corr 1.000000; fp16 wav max|d| 7.5e-03 corr 0.999953 (raw pre-iSTFT gate deliberately loose - log-mags and radian phases swing in fp16 without audio impact). Codec round-trip corr vs source 0.941. Parity command exits nonzero on gate failure. M5 Pro: T=13 (1.04s audio) 15.6 ms GPU (67x RT) / 41.3 ms ANE; per-frame T=1 3.2 ms GPU - the codec costs ~3 ms of the 80 ms frame budget, matching the "low single-digit" estimate in the feasibility table.
… CoreML Two CoreML models + host glue per 80 ms frame: backbone_step_fp16 — manual gemma3_text single-frame step (28L h1152, 16 heads hd72, q/k RMSNorm, sandwich norms, 5:1 sliding/full layer pattern, RoPE theta 10k local / 1e6 global, scale 256^-0.5, gelu-tanh), batch 2 for classifier-free guidance, rolling 1024-slot KV via ct.StateType with pos-derived validity masking. Parity vs the HF backbone: 2.96e-05 over a 12-step prefill (validates every semantic detail empirically); fp16 chained steps 7.7e-02. CoreML states are fp16-only, so no fp32 stateful variant exists. mog_dense — MoG head mlp_stack + CFG combine emitting mixture logits[1024], log-std (clamped -4), mu_res[512], guided hidden. Mixture sampling (top-p+gumbel, or argmax for parity), the low-rank mu gathers (proj_mus[idx] 64x1152 then low_mat[idx] 512x64, mu*exp(logs)+mu_res), the gated text/audio fusion prep, and RVQ depthsum encode/decode run host-side from npy exports; numpy prep verified 6.7e-06 vs the torch modules. Deterministic e2e (argmax pick, noise 0, guidance 0.2): torch replica vs CoreML codes identical 4/4 frames; full fp16 chain (backbone fp16 + MoG fp16, numpy prep) 124/124 code agreement. Parity exits nonzero on failure. Measured M5 Pro: backbone step 14.2 ms GPU (ANE rejects the stateful graph), MoG dense 0.96 ms GPU / 0.82 ANE x 8 refinement iterations -> TTS ~22 ms/frame. With codec 3.2 ms (2c) the fully-measured serial frame is encoder 12 + LLM 43.5 + TTS 22 + RNNT 0.5 + codec 3.2 ~ 81 ms vs the 80 ms budget — README records the levers back under it (CFG off, fewer iterations, amortized encoding, cross-frame pipelining). Gotchas recorded: the per-frame TTS is an 8-iteration MaskGIT-style RVQ refinement (self-inverse power schedule, exponent 3), not one decoder call; disable_eos_prediction=True (EOS from the LLM text channel, no lm_head in the checkpoint); NeMo RMSNorm is the gemma (1+w) form; interleaving coremltools predicts with torch forwards segfaults (GIL fatal) — the parity harness runs all torch phases before any CoreML predict.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Completes the two remaining model conversions for VoiceChat-11B on Apple Silicon: the TTS one-step decoder (gemma3 backbone + MoG head, Phase 2b) and the audio codec decoder (Phase 2c). Every component of the per-80 ms frame is now converted, parity-gated (nonzero exit on failure), and measured on M5 Pro. Follows up #85.
Phase 2b — TTS one-step decoder (
convert_tts.py)backbone_step_fp16ct.StateType+ pos maskingmog_denseproj_mus[idx]64×1152 →low_mat[idx]512×64), RVQ depthsum encode/decodeDeterministic end-to-end (argmax pick, noise 0, guidance 0.2): torch replica vs CoreML codes identical 4/4 frames; full fp16 chain 124/124 code agreement.
Load-bearing findings:
disable_eos_prediction=True— no lm_head in the checkpoint; end-of-speech comes from the duplex LLM's text channel.Phase 2c — codec decoder (
convert_codec.py)Conv stack (ConvT 512→1536 k9s9 → ConvNeXt×3 → →768 → ConvNeXt×3 → →384 → ConvNeXt×3 → Conv→18) in CoreML with flexible frame count; PRVQ code→latent (31 embedding sums, codebooks verified identical to the TTS
rvq_embs) and the 16-point complex iSTFT tail in numpy. Parity on real audio through the codec's own encoder: fp32 wav 7.5e-07 corr 1.000000; fp16 wav 7.5e-03 corr 0.99995. 3.2 ms GPU per frame; 1 s batches at 67× RT.Frame budget (all measured now)
encoder 12 + LLM int8 43.5 + TTS 22 + RNNT 0.5 + codec 3.2 ≈ 81 ms vs the 80 ms budget — serial worst-case is right at the line. README records the levers back under it: CFG off (batch-1 backbone), fewer refinement iterations, amortized chunked encoding, cross-frame pipelining of LLM (GPU) vs TTS/codec.
Remaining: fusion + host loop (Phase 4 Swift), HF publish (Phase 5).
🤖 Generated with Claude Code