Repository navigation
[Model] Opt-in AudioSeal watermark for generated speech - #469
Conversation
|
Moved to draft to keep the review queue on bug fixes; I'll mark it ready once the fixes ahead of it land. |
arun-elr
changed the title
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| # Optional: the opt-in AudioSeal watermark on generated speech | ||
| # (``audio_watermark`` in a deployment's model_kwargs; MIT code and weights). | ||
| # 0.3 may change the streaming layers mstar.model.audio_watermark drives. | ||
| audioseal = [ | ||
| "audioseal>=0.2,<0.3", | ||
| ] |
There was a problem hiding this comment.
[P2] The new audioseal extra and Qwen3-TTS's audio_watermark model_kwargs knob aren't in the docs.
docs/installation.rst's "Optional dependencies" table lists every extra except this new one. The Qwen3-TTS section of docs/models.rst documents the neighbouring model_kwargs: {codec_dtype: ...} knob but not audio_watermark, its keys (enabled, message, alpha, generator) or the detector CLI. That leaves the module docstring as the only place a deployer can find them. AGENTS.md invariant 9 asks for both.
There was a problem hiding this comment.
Fixed in 5b75121. docs/installation.rst's optional dependencies table has an .[audioseal] row (the version range, that it is not part of .[all], the startup failure without it, and that the detector CLI also needs .[audio]). The Qwen3-TTS notes in docs/models.rst document model_kwargs.audio_watermark next to codec_dtype: the four keys with their defaults and ranges, what fails at startup, the per-request streaming state, and python -m mstar.model.audio_watermark detect with its flags.
Automated reviewThis PR adds an engine hook, Findings: 1 inline (P2): the new The description is thorough and, as far as the code shows, accurate. Two things are missing:
Advisory, not a merge gate. Rules: |
Every number was rerun on the head with in-repo tools, and the description now gives each command. The overhead table is
The Tests section now gives the command: |
4007b38 to
cdb5719
Compare
There was a problem hiding this comment.
Ran cdb57193 (+ merged into main) on an h100 vs vLLM-Omni 0.30.0, with the AudioSeal detector and whisper on the served audio. On h100 it costs 40-55% of throughput (vs 20-25% on your B200):
| audio-s/s | vLLM-Omni | M* off | M* on | on + 1,2 below |
|---|---|---|---|---|
| c1 | 14.4 | 15.5 | 6.9 | 12.7 |
| c8 | 69.7 | 68.6 | 40.5 | 50.5 |
| c32 | 142.3 (MPS) | 143.8 | 73.3 | 103.2 |
- audio_watermark.py:330: with TF32 off, cuDNN on h100 picks a slow conv at batch 1-2, so a 2 s chunk takes 155-180 ms (your
costcommand: 174.8 ms here vs 7.8 on B200). With cuDNN TF32 allowed it's 9.5 ms and streamed vs whole stays within 3 int16 steps - audio_watermark.py:333: AudioSeal's MsgProcessor does
torch.arange(16).to(device)every call, so each codec step syncs the GPU thread and the generator can't be captured. The message is per deployment, so precomputing the embedding in__init__fixes it (bit identical, capture then works: b1 4.4 -> 1.3 ms) - audio_watermark.py:446: the mark only detects at 24 kHz. It survives MP3 / Opus / a 16 kHz round trip (236/236) but 0/236 on a 16 or 48 kHz copy, so
detectshould resample to 24k first - submodule_base.py:903: main already has
batch_postprocess(:894), unused and only called on rows-only steps. Calling it after postprocess on every step gives one batch hook instead of two
cdb5719 to
e358861
Compare
|
Thanks for the H100 runs and for pinning down each cause. All four are fixed and pushed. Measured on one H200 (no H100 in the pool; same Hopper cuDNN behavior, the slow conv reproduces: 159 ms per 2 s chunk at batch 1 before).
Fixed: the generator now runs with cuDNN TF32 allowed (matmuls stay fp32).
Fixed: the message embedding is built once in
Fixed:
Fixed: Tests: |
ba47169 to
912c55f
Compare
merceod
left a comment
merceod
left a comment
There was a problem hiding this comment.
Checked ba47169 (912c55f is the same patch rebased; the move from a4a9867 was the detect --message CLI only) on klone: correctness on an A100, timing on an H200 with three interleaved launches per arm (clock-limited node). Detection is 100% with the exact 0xBEEF message on whole files and on the first, middle and last streamed chunks, 0 clean files flagged, SNR 24.3 to 24.4 dB, PESQ 4.40, whisper WER unchanged, and within a launch the marked speech is the clean speech plus the mark (same lengths, mark only). Exactness against AudioSeal fp32 is 7.6e-5 to 8.4e-5. The slow-conv claim reproduces: a 2 s chunk at batch 1 takes 234 ms with cuDNN TF32 off and 14 ms with the workaround.
Cost, on vs off: 19.7%, 30.0% and 32.8% less audio per second at c=1, c=8 and c=32 (description 20 / 28 / 29), TTFA p50 up 8, 21 and 47 ms.
Notes, none blocking.
- Inside
_watermark_precisionthe process-wide TF32 flags are flipped, and while that window is opentorch.get_float32_matmul_precision()raises "mix of the legacy and new APIs" from any thread (torch 2.12.1). Nothing on the served path calls the getter today, so it is latent. Comment inline. - A raise inside the
batch_postprocesshook fails every request in that codec step, not only the one at fault. Comment inline. MARK_SAMPLE_RATEis a 24 kHz constant anddetectresamples everything to 24 kHz, so a 48 kHz (VoxCPM2) or 22.05 kHz (Nemotron) model would need per-model handling.- Stale comment at qwen3_tts/submodules.py:977 naming
postprocess_batch_tail. - One served byte check I could not do: flag off vs the base commit, because seeded Qwen3-TTS output differs between server launches on this stack even base against base. Neutrality rests on the code path, your CPU test and the same-launch pair above.
| at batch 1 and 2; with it on, a chunk and the whole utterance, which | ||
| convolve over different lengths, round apart by a few int16 steps.""" | ||
| matmul, cudnn = torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 | ||
| torch.backends.cuda.matmul.allow_tf32 = False |
There was a problem hiding this comment.
While this window is open, torch.get_float32_matmul_precision() raises "mix of the legacy and new APIs" from any thread (torch 2.12.1 on klone). The flags are restored after and nothing on the served path calls the getter today, so it is latent. Setting the precision locally on the conv/matmul call or using the new API would close it.
There was a problem hiding this comment.
While this window is open,
torch.get_float32_matmul_precision()raises "mix of the legacy and new APIs" from any thread
Fixed. The window now sets torch.set_float32_matmul_precision("highest") and restores the previous value, the API the engine uses at startup, and still turns cuDNN TF32 on. Flipping cuda.matmul.allow_tf32 alone left the CPU matmul setting at TF32, and that mix is what the getter refused; setting the new per-backend fp32_precision attributes instead only moves the error to the legacy getters. Inside the window every getter answers now, and the CUDA settings are the same as before (fp32 matmuls, TF32 cuDNN conv and RNN; CPU matmuls also run fp32 there).
Reproduced on torch 2.12.1 (CPU). test_precision_getters_answer_while_the_watermark_runs fails before the fix with that error and needs no audioseal, so CI runs it.
| return | ||
| rids = [rid for rid, _ in live] | ||
| try: | ||
| submodule.batch_postprocess( |
There was a problem hiding this comment.
A raise inside the hook fails every request in this codec step, not only the one whose chunk was bad (for example "chunk after its final one"). Per-request isolation here, or a guard in the hook, would keep one bad stream from taking its batch-mates down.
There was a problem hiding this comment.
A raise inside the hook fails every request in this codec step, not only the one whose chunk was bad
Kept as is in this PR, on purpose: a raise in batch_postprocess fails every request the step was marking (test_batch_tail_failure_fails_its_requests_only), so no chunk goes out unmarked. The raises it can hit are batch-wide (a CUDA error, AudioSeal's residual branches falling out of step). The one per-stream check, a chunk after the final one, guards an invariant: the codec node sends is_final once per request, and the request ends there.
Failing only the request at fault needs the hook to report failures per request (for example an exception carrying request id to error, which the engine registers one by one). That is an engine API change, better reviewed on its own; happy to open it as a follow-up if you want it.
7274c97 to
6b6932d
Compare
2525c5b to
dd5c3ca
Compare
|
@arun-elr could you resolve the branch conflicts? |
|
Both approvals are in now. The only thing left is the conflict with main (the test list and docs); the fork does not allow maintainer edits, so it needs a rebase from your side. It merges on the next pass once CI is green on the rebased head. |
Adds NodeSubmodule.postprocess_batch_tail (called once per step after the per-request postprocess, only when overridden), the audioseal extra, and CPU tests: config validation, missing dependency, flag off untouched, streamed equals offline, batching, detection when weights are cached.
TF32 convolutions round a chunk and the whole utterance apart by a few int16 steps (3.7e-5 max abs on a B200, 3.0e-8 with TF32 off, as on CPU), and the engine turns TF32 matmuls on process-wide.
A batch-dim slice of the [layers, batch, hidden] state is not contiguous, and cuDNN's LSTM refuses it when that stream next runs alone (found at concurrency 8 on a B200: 'rnn: hx is not contiguous').
…mark/audio_watermark_check.py reruns the watermark's exactness, cost and quality checks
… in full fp32, as the watermarker runs
…, detect at 24 kHz With cuDNN TF32 off, cuDNN on H100 picks a far slower conv at batch 1 and 2; the watermark now runs its convs in TF32 (matmuls stay fp32). The message embedding AudioSeal's MsgProcessor rebuilds from a host arange on every call is built once per deployment, so a watermark call no longer syncs the GPU. detect resamples a file to the 24 kHz the mark was added at. The engine calls the existing batch_postprocess after every request's postprocess instead of adding a second batch hook.
…e deployment's message An unmarked Qwen3-TTS file re-encoded as MP3 can score above 0.5 (50 of 167 served files did) while decoding to some other message. --message takes the deployment's payload; the docs show it.
…sion, so the precision getters keep answering
dd5c3ca to
afc2c1e
Compare
Rebased on current main. Three small commits on top, nothing else changed:
Agreed. Only Qwen3-TTS (24 kHz) turns the watermark on in this PR, so the per-model rate goes in with the first model at another rate that enables it. Locally on the head: CI's CPU Core selection 4719 passed, 890 skipped, TP lockstep sim 7 passed; with |
Closes #468
What does this PR do?
Adds an opt-in AudioSeal watermark to Qwen3-TTS's streamed speech: off by default, on with
model_kwargs: {audio_watermark: {enabled: true, message: 0xBEEF}}andpip install 'mstar-ai[audioseal]'.audio_watermark.py: AudioSeal generator with per-request streaming state, detector and CLI, because AudioSeal 0.2'sstreaming()carries no history.batch_postprocessafter every request'spostprocess(engine.py,submodule_base.py): one batched call per step, only where overridden.mstar/model/qwen3_tts/); with the flag off there is no watermarker and the hook returns at once.audiosealextra, docs, and a check script (pyproject.toml,docs/,benchmark/) that reruns every number below.On one H200 the mark is detected on every served file with the exact message, also on 16 and 48 kHz copies, at a cost of 20% (c1) to 29% (c32) of audio seconds per second against main.
How was it tested?
On 3800486 (head, rebased on current main):
ruff check .clean. CI's CPU Core selection (GPU hidden): 4719 passed, 890 skipped; TP lockstep sim: 7 passed. Withaudiosealandqwen-ttsinstalled (CPU):test_audio_watermark.pyandtest_qwen3_tts_model.py105 passed, 1 skipped (CUDA only). Earlier, on 2525c5b:pip install -e '.[dev,qwen3_tts,audioseal,audio]', thenpython -m pytest test/modular/test_audio_watermark.py test/modular/test_qwen3_tts_model.py test/modular/test_engine_rows_only_tail.py: 109 passed on one H200 (3 of the 4 newdetectcases fail on a4a9867). One H200:benchmark.tts_speech_benchon main, flag off and flag on, andaudio_watermark_check.py exactness,cost,quality: see the toggles.Real-user checks on one H200 (served on a4a9867; ba47169 changes only the detect CLI): Qwen3-TTS 0.6B over
t2s.txtat c=1/4/16/32 gave 13.5/30.6/72.0/113.7 audio-s/s with the mark against 16.7/40.5/103.1/161.7 without (19% to 30%, 0 errors) and Whisper WER 5.04/5.64% against 4.86/5.13% at c=4/32, with 266- and 401-word texts at 0.15% against 0.45%; the mark was found with the exact message on 167 of 167 served files and every 10 s window of the long ones, on 16, 22.05, 44.1 and 48 kHz copies, after MP3 (167) and Opus (164), and on no 8 kHz copy.Not tested, and not here
Not tested: H100 and B200 on this head (H200 only: same Hopper cuDNN behavior, the slow conv reproduces); other TTS models (only Qwen3-TTS is wired).
Not here: capturing the generator and taking it off the decode loop (#598); a per-request switch; a detection endpoint.
Depends on
None.
Design
audio_watermarkin the deployment'smodel_kwargstakesenabled(default false),message(16-bit payload: an int 0..65535 or a 16-character 0/1 string, default 0),alpha(strength in (0, 4], default 1.0) andgenerator(an AudioSeal card name or checkpoint path, defaultaudioseal_wm_streaming). Unknown keys, bad values, orenabled: truewithout the package fail at startup, not at the first request.NodeSubmodule.batch_postprocess(outputs, request_infos, inputs)was only called on rows-only steps, and nothing overrode it. It now also runs on every other step, after every request'spostprocess, for the requests whosepostprocesssucceeded; a raise fails those requests, as a failedpostprocessdoes. Only submodules that override it pay anything.nn.Sequentialthat holds the layers, so no conv keeps a cache and every chunk is marked as if it began the stream. The helper keeps each request's state (conv caches, the LSTM state, a pending partial frame) as plain tensors in the codec's request state and installs it in the generator's streaming layers around each call, which also lets requests at the same point share one batched call. Whole 320-sample frames go through; a partial frame waits for the next chunk or is zero-padded by the final one. Qwen3-TTS chunks are whole frames (1,920 samples per codec frame), so nothing is held back.MsgProcessorrebuilds it from a hostarangeon every call, a copy that synced the GPU on every codec step. Bit identical to the processor (test).python -m mstar.model.audio_watermark detect --message 0xBEEF out.wav(needs.[audio]for soundfile). With--messagea file counts as marked only if it decodes to the deployment's payload: by score alone, 50 of 167 clean served files re-encoded as MP3 passed 0.5, and none decoded to 0xBEEF. The mark is found only at the rate it was added at, so a file at another rate is resampled to 24 kHz first. Detect on whole utterances, or at least a few seconds: on this speech every single codec window was detected, but on a synthetic harmonic tone slices under a second scored below the threshold.watermark: false; which convention to share is open). An utterance that ends mid-frame is zero-padded to a whole frame before marking, while AudioSeal's own call pads inside each layer; the two differ in that last partial frame only (Qwen3-TTS never pads). A clean file resampled to 16 kHz scored up to 0.46 against the 0.5 threshold (none flagged; tracked in [Qwen3-TTS][Perf] AudioSeal watermark: capture the generator and take it off the decode loop #598). A copy resampled to 8 kHz loses the mark (0 of 167 served files detected).Performance
One H200 (torch 2.12.1+cu130, cuDNN 9.20, audioseal 0.2.0), Qwen3-TTS 0.6B CustomVoice. "Before" is the previous head rebased on main (4f65a3a), "this head" is e358861, main is 84fd4a6.
Served (
configs/qwen3tts.yaml; "on" addsmodel_kwargs: {audio_watermark: {enabled: true, message: 0xBEEF}}). The 27 sentences ofbenchmark/assets/t2s.txt(c32: the file 4 times, 108 requests, so 32 stay in flight), 3 repeats, median over repeats:python -m benchmark.tts_speech_bench --engine mstar --url http://127.0.0.1:8737 --model qwen3_tts \ --sentences benchmark/assets/t2s.txt --voice vivian --language English --concurrency 8 --repeats 3 --warmup 3 \ --out c8.json --save-audio-dir c8_wavNo errors in any run. Flag off matches main within noise.
Generator per call (
python -m benchmark.audio_watermark_check cost --device cuda,apply_pcm16, streams under way, median of 40):What is left is mostly launch overhead on the worker loop (a 1-frame chunk still costs about 4 ms); capturing the call or moving it off the decode loop is #598.
Correctness
Same H200 and code as Performance (this head).
Streaming exactness (
python -m benchmark.audio_watermark_check exactness --device cuda): the chunked watermark of a 3 s, 24 kHz speech-like signal against AudioSeal's own whole-utteranceget_watermarkin full fp32: within 8.3e-5 (2.7 int16 steps) for aligned, uneven and batched chunkings, and the offline call within 8.5e-5. The previous head, in full fp32, was within 3.0e-8. AudioSeal's ownstreaming()on the same chunks is 6.7e-3 off, which is the bug the helper works around. On CPU the streamed mark still equals the offline one to float rounding (tests).Quality and detection (
python -m benchmark.audio_watermark_check quality --device cuda --clean c8_off_wav --marked c8_on_wav, and the same for c1). Served greedy speech does not repeat across server restarts, so the clean files (flag off) are marked by the script chunk by chunk along the codec's schedule for SNR, PESQ and streamed-vs-whole; detection scores the served files of both servers, threshold 0.5.One codec window at a time (80 ms to 2 s, both concurrencies): 650 of 650 marked windows detected with the exact message, 0 of 650 clean windows flagged (max score 0.151).
Flag off is unchanged:
test_qwen3_tts_codec_without_watermark_emits_the_same_bytes_through_the_engine(CPU) checks that the codec emits the same bytes through the engine's postprocess path with the hook present; served, flag off matches main's throughput.Tests in detail
test/modular/test_audio_watermark.py(collected by CI's CPU tests): config defaults, message forms and bad values; the missing-package error names the extra, and Qwen3-TTS checks it before loading anything;batch_postprocessruns once on the requests whosepostprocesssucceeded, and its failure fails those requests only; streaming equals offline for several chunkings, batched streams do not mix and keep contiguous caches, PCM16 round trip within one step; matmul TF32 is off and cuDNN TF32 on during the mark, both restored after, and every precision getter answers during it; the precomputed message embedding equalsMsgProcessor's output and the processor never runs during a call;detectresamples 16 and 48 kHz input to 24 kHz; chunks after the final one and non-causal generators are refused; AudioSeal's ownstreaming()is shown not to carry history; the detector finds the streamed mark and its message, also on 16 and 48 kHz copies (real weights, only when cached). The streaming tests needaudioseal(random weights, no download) and skip without it; CI installs no extra, so there they skip (CI: the Qwen3-TTS model tests and the AudioSeal streaming tests never run #487).test/modular/test_qwen3_tts_model.py(needsqwen-tts, skips without it): with no watermarker the codec leaves audio untouched and emits the same bytes through the engine; with one, two requests throughbatch_postprocessequal their whole-utterance marks and the stream state goes with the request.test/modular/test_engine_rows_only_tail.py(main's): rows-only steps still run onlybatch_postprocess.Provenance: credit
AudioSeal is by Meta AI (facebookresearch/audioseal), MIT licence for code and weights: San Roman et al., "Proactive Detection of Voice Cloning with Localized Watermarking", ICML 2024. It is used as an optional dependency; no AudioSeal code is copied into M*.
Checklist
ruff check .passesAI assistance: written with AI coding agent