[speechlm2] Add streaming inference pipeline for NemotronVoiceChat - #16198
[speechlm2] Add streaming inference pipeline for NemotronVoiceChat#16198erastorgueva-nv wants to merge 6 commits into
Conversation
…LM-Omni backends StreamingS2SPipeline and the VoiceChat wrapper share one frame loop, with LLM and TTS engines independently selectable. Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>
Keep DuplexLLM/DuplexTTS and the Vllm* classes so the native frame loop already matches the combined form. Selecting vllm_omni raises NotImplementedError; the runtime is the parent commit on duplex-vllm-omni-on-main. Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>
Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>
Replace the VoiceChat-specific encoder capture with encoder.set_streaming_cuda_graphs so subsequent cache-aware steps share the ASR helper. Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>
Remove redundant streaming overrides so native and converted EarTTS use the sampling values carried by the model checkpoint. Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>
Keep distinct config, engine, and tiny-model parity checks while avoiding repeated cases and a second real-checkpoint test. Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| cas_cfg = cfg["model"]["speech_generation"]["model"]["tts_config"]["cas_config"] | ||
| if cas_cfg.pop("pretrained_tokenizer_name", None) is not None: | ||
| logging.info("NemotronLabs release shim: removed legacy CAS pretrained_tokenizer_name") | ||
| except (KeyError, TypeError): |
|
@codex review this from the angle of test and regression coverage: there is a public Nemotron VoiceChat checkpoint on HF that can be used for functional tests |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e77e1ac127
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") | ||
| def test_pipeline_no_crash_hf_11b( |
There was a problem hiding this comment.
On any machine with CUDA, this test runs even when pytest was invoked without --with_downloads; the autouse gate in tests/conftest.py only skips tests carrying pytest.mark.with_downloads. Consequently, an ordinary SpeechLM2 GPU test run can unexpectedly download and load the public 11B checkpoint, consuming substantial time, disk, and GPU memory. Mark this checkpoint test with with_downloads (and preferably integration) so it remains an explicitly selected functional test.
Useful? React with 👍 / 👎.
| result = _run(pipeline, voicechat_audio_path) | ||
| output = result[0] | ||
| assert output.token_asr_text is None | ||
| assert output.raw_asr_text is None | ||
| assert output.token_function is not None |
There was a problem hiding this comment.
The only real-checkpoint test verifies that output fields are allocated, so regressions that misload weights or corrupt decoding can still pass while producing empty/padding-only tokens, silence, NaNs, or unrelated speech. Since this uses a fixed public checkpoint and fixed audio, configure deterministic decoding and assert a stable token/text invariant plus finite, non-silent audio; otherwise the new HF compatibility and native inference paths have no meaningful functional regression oracle.
AGENTS.md reference: AGENTS.md:L78-L78
Useful? React with 👍 / 👎.
| # Decode phase. Prefill-only frames (empty audio) are not passed to | ||
| # inference; outputs are stitched back 1:1 with the original *frames*. | ||
| non_empty_frames = [f for f in frames if f.samples.numel() > 0] | ||
| empty_terminal_frames = [f for f in frames if f.samples.numel() == 0 and f.is_last] | ||
| for frame in empty_terminal_frames: |
There was a problem hiding this comment.
All added pipeline tests enter through run(), whose first frame already contains audio, so none exercise the documented live-client flow where generate_step() receives an empty initialization frame or an empty terminal frame. Add a direct test that performs empty prefill, multiple audio chunks, termination, and stream-ID reuse while checking incremental outputs and state cleanup; otherwise regressions in this separate branch can break microphone/server integrations despite the file-based suite passing.
AGENTS.md reference: AGENTS.md:L78-L78
Useful? React with 👍 / 👎.
3 participants
Important
The
Update branchbutton must only be pressed in very rare occassions.An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.
What does this PR do ?
Add a streaming (chunk-by-chunk) inference pipeline for NemotronVoiceChat on the native PyTorch backend, following the same architecture as the NeMo ASR Inference Pipelines. Successor to #15571.
Collection: speechlm2
Changelog
StreamingS2SPipelinewithrun()for files/manifests andgenerate_step()for live audioNemotronVoicechatInferenceWrapperwith one frame loop: perception → LLM → TTS → codec decodeDuplexLLM/DuplexTTScontracts; native implementations inbackend/pytorch/. Engine type is chosen at construction, not in the per-frame bodyS2SPipelineBuilder(plain factory, same shape asasr/inference) and Hydra configs2s_streaming.yamlinference_precision_from_cfg; callers shut down the pipeline infinallyencoder.set_streaming_cuda_graphs; subsequent chunks passkeep_all_outputs=False)S2SContextManagerfor decode-state lifetime,S2SStreamingOutputfor accumulations2s_streaming_infer.pyfor files, directories, or manifestsstreaming_inference.rst(architecture, config, trailing-silence vs batch padding)tests/collections/speechlm2/nemo_inference_pipelines/: CPU units, tiny-model no-crash + offline-vs-streaming parity, public-11B native no-crashTo choose a voice, pass
speaker_namematching a latent registered in the checkpoint (public 11B:Aria).speaker_referenceis rejected — encoding a new wav goes through the anti-cloning projection and sounds wrong.The public HF checkpoint (
nvidia/NVIDIA-NemotronLabs-VoiceChat-11B) is supported. It has a function head; we decode those tokens (and feed them back into the next frame) but we do not actually call the functions. Executing tool calls would add a lot of complexity (API-call timing, etc.), and we might not release future native function-call models, so it's not worth implementing here.batch_size=1(cc @pzelasko)
Keeping this at 1.
The last review (#15571) noted that the inner forwards look capable of
bs>1. Lock-step batching (pad every file to the longest and step them together) wouldn't add much on top of offline inference — if you want fast batched eval, use the offline path. It would also take a lot of extra code for not much gain, and it would be confusing to exposebatch_size>1without independent stream start/end: in a real-time deployment a new stream could only join the batch once the others had also finished, which is a pretty weird API.Supporting several streams of different lengths at once is complex (per-stream caches for perception, LLM KV, EarTTS, codec, plus stream start/end). So we will not support this currently, making this pipeline just for local inference/initial PoC.
streaming.batch_sizeis asserted at construction and ingenerate_step.vLLM-Omni follow-up
This PR is native only. The combined native+vLLM tree was designed first, so the frame loop and
DuplexLLM/DuplexTTScontracts already match the combined form.VllmLLM/VllmEarTTSare stubs: selectingvllm_omniraisesNotImplementedErrorat wrapper construction.You can find a draft of vllm code (stock
vllm==0.26.0/vllm-omni==0.26.0, no fork; LLM and TTS independently switchable) on a parent commit on this branch: 1a2bd8dModifications to more general code — FYI @kevinhu-nv @Edresson
EarTTSModel: vectorized RVQ depth-sum embedding, optional per-subword embedding cache (use_tts_subword_cache), MaskGIT unmasking loop uses Python ints (compile-friendly)DuplexEARTTS: skipget_codec_silence_frame()when the codec has random weights (slow encode of silence that the checkpoint overwrites anyway)PyTorchEarTTS: optionaltorch.compileon the TTS backboneDuplexSTTModel: function head (autoregressive feedback even when tools are not executed), sharedLogitBoostsfor agent/ASR channels,build_input_embeddingso offline and streaming match addition order,cache_position/ configurablecache_keyfor Nemotron KV cacheNemotronVoiceChat.from_pretrained: HF-format checkpoints withllm_artifacts/, meta-device init,skip_prefixesfor submodules the caller will replace,return_logitsfor parity teststext_utils: sharedstrip_timestamps, byte-level BPE decoding, BOS/EOS kept as literal stringspretrained.py/hf_hub.py: meta-device LLM construct, skip pretrained ASR/LLM downloads when loading a VoiceChat HF dir, tokenizer export onsave_pretrainedUsage
python examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py \ audio_file=/path/to/audio.wav \ s2s.model_path=nvidia/NVIDIA-NemotronLabs-VoiceChat-11B \ s2s.speaker_name=Aria \ s2s.llm_engine_type=native \ s2s.tts_engine_type=native \ streaming.chunk_size_in_secs=0.08 \ streaming.buffer_size_in_secs=1.68GitHub Actions CI
The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.
Trusted PRs run automatically through copy-pr-bot. For an untrusted PR, a maintainer can trigger CI by commenting
/ok to test <head-sha>; repeat this after a new push if the PR remains untrusted.Before your PR is "Ready for review"
Pre checks:
vllm_omniis stubbed).PR Type:
If you haven't finished some of the above items you can still open "Draft" PR.
Who can review?
Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.
Additional Information