Skip to content

Add streaming multi-talker Parakeet RNN-T ASR model - #411

Open
themason2011 wants to merge 1 commit into
mainfrom
themason2011/parakeet-multitalker
Open

Add streaming multi-talker Parakeet RNN-T ASR model#411
themason2011 wants to merge 1 commit into
mainfrom
themason2011/parakeet-multitalker

Conversation

@themason2011

@themason2011 themason2011 commented Jul 17, 2026

Copy link
Copy Markdown

Summary

Adds NVIDIA NeMo's EncDecMultiTalkerRNNTBPEModel (multitalker-parakeet-streaming-0.6b-v1) as its own model type. It is a cache-aware streaming multi-talker RNN-T ASR model derived from nemotron-speech-streaming, exported as three ONNX graphs matching that deployment contract:

  • encoder — FastConformer with speaker-kernel injection and streaming channel/time caches. Uses causal dw_striding subsampling (asymmetric pad, 8x reduction), bias-free linears, LayerNorm convolution, chunked-limited attention, and a relative-position KV cache.
  • decoder — RNN-T LSTM prediction network (2-layer ONNX LSTM with PyTorch->ONNX gate reordering (i,f,g,o) -> (i,o,f,g)).
  • joint — RNN-T joint network emitting [B, T, U, V] logits.

New public API: ParakeetMultiTalkerModel / ParakeetMultiTalkerConfig with a .nemo loader build_parakeet_multitalker, and a MultiTalkerRNNTTask (registered as multitalker-rnnt).

Testing

  • Unit (tests/build_graph_test.py::TestBuildGraphParakeetMultiTalker): builds all three graphs from a tiny config and runs each through ORT with random weights; verifies I/O names/shapes and the streaming length math.
  • Integration (tests/parakeet_multitalker_integration_test.py, -m integration): streaming parity vs the NeMo PyTorch reference — encoder 3.9e-7, caches 2e-6 / 1.8e-4, decoder 1.1e-6, joint log-softmax 1.5e-5 with matching argmax. Verified end-to-end against real weights.
  • ruff check / ruff format clean.

Note

Independent of #410 — the small FastConformer _ReLU / _swish primitives are defined locally in models/parakeet_multitalker.py, so this PR has no code or git dependency on the Sortformer branch and targets main directly.

@github-actions

Copy link
Copy Markdown

Performance Comparison

Comparing bcf8e1543bcec7

Model Metric Baseline Current Delta
bert (feature-extraction) model_size_bytes 359 KB 359 KB +0.0%
bert (feature-extraction) num_nodes 60 60 +0.0%
falcon model_size_bytes 364 KB 364 KB +0.0%
falcon num_nodes 68 68 +0.0%
gemma2 model_size_bytes 428 KB 428 KB +0.0%
gemma2 num_nodes 107 107 +0.0%
gpt2 model_size_bytes 388 KB 388 KB +0.0%
gpt2 num_nodes 54 54 +0.0%
llama model_size_bytes 425 KB 425 KB +0.0%
llama num_nodes 62 62 +0.0%
llama (static-cache) model_size_bytes 425 KB 425 KB +0.0%
llama (static-cache) num_nodes 58 58 +0.0%
mamba (ssm-text-generation) model_size_bytes 296 KB 296 KB +0.0%
mamba (ssm-text-generation) num_nodes 98 98 +0.0%
phi3 model_size_bytes 421 KB 421 KB +0.0%
phi3 num_nodes 60 60 +0.0%
phi3 (static-cache) model_size_bytes 421 KB 421 KB +0.0%
phi3 (static-cache) num_nodes 56 56 +0.0%
qwen2 model_size_bytes 425 KB 425 KB +0.0%
qwen2 num_nodes 62 62 +0.0%
qwen2 (static-cache) model_size_bytes 425 KB 425 KB +0.0%
qwen2 (static-cache) num_nodes 58 58 +0.0%
qwen3_5_moe (hybrid-text-generation) model_size_bytes 506 KB 506 KB +0.0%
qwen3_5_moe (hybrid-text-generation) num_nodes 275 275 +0.0%
qwen3_5_text (hybrid-text-generation) model_size_bytes 458 KB 458 KB +0.0%
qwen3_5_text (hybrid-text-generation) num_nodes 129 129 +0.0%
qwen3_5_vl (hybrid-qwen-vl) model_size_bytes 977 KB 977 KB +0.0%
qwen3_5_vl (hybrid-qwen-vl) num_nodes 413 413 +0.0%
t5 (seq2seq) model_size_bytes 836 KB 836 KB +0.0%
t5 (seq2seq) num_nodes 166 166 +0.0%
whisper (speech-to-text) model_size_bytes 1008 KB 1008 KB +0.0%
whisper (speech-to-text) num_nodes 128 128 +0.0%

No performance regressions.

@github-actions

Copy link
Copy Markdown

🏗️ Architecture Diff

Comparing bcf8e1543bcec7

Model Sub-model Changes Status
bert (feature-extraction) model 0
falcon model 0
gemma2 model 0
gemma4 (gemma4) decoder 0
gemma4 (gemma4) embedding 0
gemma4 (gemma4) vision_encoder 0
gemma4_text model 0
gpt2 model 0
llama model 0
llama (static-cache) model 0
mamba (ssm-text-generation) model 0
phi3 model 0
phi3 (static-cache) model 0
qwen model 0
qwen (static-cache) model 0
qwen2 model 0
qwen2 (static-cache) model 0
qwen2_moe model 0
qwen2_moe (static-cache) model 0
qwen3 model 0
qwen3 (static-cache) model 0
qwen3_5_moe (hybrid-text-generation) model 0
qwen3_5_text (hybrid-text-generation) model 0
qwen3_5_vl (hybrid-qwen-vl) decoder 0
qwen3_5_vl (hybrid-qwen-vl) embedding 0
qwen3_5_vl (hybrid-qwen-vl) vision_encoder 0
qwen3_moe model 0
qwen3_moe (static-cache) model 0
qwen3_next (hybrid-text-generation) model 0
t5 (seq2seq) decoder 0
t5 (seq2seq) encoder 0
whisper (speech-to-text) decoder 0
whisper (speech-to-text) encoder 0

No architecture changes detected.


Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

@codecov

codecov Bot commented Jul 17, 2026

Copy link
Copy Markdown

The author of this PR, themason2011, is not an activated member of this organization on Codecov.
Please activate this user on Codecov to display this PR comment.
Coverage data is still being uploaded to Codecov.io for purposes of overall coverage calculations.
Please don't hesitate to email us at support@codecov.io with any questions.

Add NVIDIA NeMo `EncDecMultiTalkerRNNTBPEModel`
(multitalker-parakeet-streaming-0.6b-v1) as its own model type,
exported as three ONNX graphs matching the nemotron-speech-streaming
cache-aware streaming deployment contract:

- encoder: FastConformer with speaker-kernel injection and streaming
  channel/time caches (causal dw_striding subsampling, bias-free
  linears, LayerNorm conv, chunked-limited attention, rel-pos KV cache).
- decoder: RNN-T LSTM prediction network (2-layer ONNX LSTM with
  PyTorch->ONNX gate reordering).
- joint: RNN-T joint network emitting per-(time, target) logits.

Adds MultiTalkerRNNTTask, the ParakeetMultiTalkerModel/Config classes
with a .nemo loader (build_parakeet_multitalker), unit tests (graph
build + ORT smoke run of all three graphs), and an integration test
verifying streaming parity against the NeMo PyTorch reference
(encoder 3.9e-7, decoder 1.1e-6, joint argmax match).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Mason Corey <masoncorey@microsoft.com>
@themason2011
themason2011 force-pushed the themason2011/parakeet-multitalker branch from 43bcec7 to 0816c51 Compare July 17, 2026 20:36
@themason2011
themason2011 changed the base branch from themason2011/sortformer-diarization to main July 17, 2026 20:36
@themason2011
themason2011 requested a review from a team July 17, 2026 20:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant