Add streaming multi-talker Parakeet RNN-T ASR model - #411
Open
themason2011 wants to merge 1 commit into
Open
Conversation
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
|
The author of this PR, themason2011, is not an activated member of this organization on Codecov. |
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
force-pushed
the
themason2011/parakeet-multitalker
branch
from
July 17, 2026 20:36
43bcec7 to
0816c51
Compare
themason2011
changed the base branch from
themason2011/sortformer-diarization
to
main
July 17, 2026 20:36
justinchuby
added a commit
that referenced
this pull request
Aug 13, 2026
## Summary Add first-class Mobius support for [`nvidia/parakeet-ctc-1.1b`](https://huggingface.co/nvidia/parakeet-ctc-1.1b), Hugging Face `model_type: parakeet_ctc`. This is an independent offline CTC ASR architecture and does not duplicate #411, which targets streaming Parakeet RNNT. The new path consumes processor-normalized log-mel features and produces frame-level CTC logits. ## Architecture - Symmetric 8x depthwise Conv2d FastConformer subsampling - Transformer-XL relative positional self-attention with learned `bias_u`/`bias_v` - Macaron feed-forward blocks - BatchNorm Conformer convolution modules - 42-layer, hidden-size 1024 encoder and 1025-class Conv1d CTC head - Feature-input contract: `input_features [B, T, 80]` plus boolean `attention_mask` - Direct Hugging Face parameter-name alignment; preprocessing only removes BatchNorm `num_batches_tracked` The reusable component surface also gains grouped/bias-free `Conv1d` and frozen `BatchNorm1d` inference with fp32 accumulation. ## Model and golden integration - Register `parakeet_ctc` config/model/task extraction - Add `FeatureCTCAsrTask` - Add `feature-ctc-asr` YAML/schema and golden generation support - Pin config and every checkpoint shard to revision `20e63a0fed6aedba145b74b826dbd41df0941730` - Add exact frame-count enforcement for deterministic CTC L5 goldens - Add explicit ORT GenAI export rejection for this unsupported feature-input CTC contract Committed real-audio goldens use nonzero LibriSpeech audio and contain all 114 frame argmax IDs plus the full decoded transcript: > cauliflower mayonnaise take cold boiled cauliflower break into branches adding salt pepper and vinegar to season ## Validation GPU: NVIDIA RTX A1000 Laptop GPU, 8 GB; CUDA and `CUDAExecutionProvider` available. - `python -m pytest tests/build_graph_test.py tests/cli_test.py src/ -q -k "not phi4mm and not apply_weights_unknown" --tb=short -n auto` - **3639 passed, 57 skipped** - `python -m pytest tests/arch_validation_test.py -k parakeet -q --tb=short` - **3 passed, 874 deselected** - `python -m pytest tests/yaml_schema_test.py -q --tb=short` - **237 passed** - Real-checkpoint CUDA parity, fp32: - **passed** - max absolute difference: **0.000076** - mean absolute difference: **0.000014** - Real-checkpoint CUDA fp16: - **passed** - exact **114/114** frame IDs - exact full transcript - CUDA L4 checkpoint golden and L5 deterministic CTC transcript/frame golden: - **2 passed, 377 deselected** - `lintrunner -a` - **no issues** - Automated high-confidence code review after fixes: - **no significant issues found** ## CLI/runtime/quantization evidence - fp16 CLI CUDA export loaded and ran with ONNX Runtime CUDA EP - Runtime produced the exact 114 golden frame IDs and exact full transcript - Olive 0.13.0 INT4 RTN quantization produced a 751 MB model - Quantized CUDA inference produced the same full transcript ## Evidence-based limitations and waivers - **BF16 disabled:** Hugging Face bf16 reference output is correct, and the Mobius graph executes under ORT CUDA after avoiding unsupported kernels, but ORT output collapses to blank CTC predictions. Parakeet therefore defaults safely to fp32 and rejects explicit bf16; fp32 and fp16 are verified. This avoids exposing a configuration that runs without error but returns incorrect ASR results. - **ORT GenAI / Foundry Local waived:** ORT GenAI 0.15.2 rejects `input_features` in its decoder model schema. This offline feature-input CTC model cannot be represented by the current ORT GenAI decoder contract, so Mobius now fails explicitly rather than producing a misleading package. Foundry Local depends on ORT GenAI and is consequently not applicable until that runtime supports feature-input CTC pipelines. - **Safetensors package output:** the installed `onnx_ir` failed during safetensors CLI serialization; standard ONNX external-data export succeeded and was used for CUDA runtime and quantization validation. --------- Signed-off-by: Justin Chu <justinchu@microsoft.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Member
|
@themason2011 could you update this PR? |
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
2 participants
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
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 fromnemotron-speech-streaming, exported as three ONNX graphs matching that deployment contract:dw_stridingsubsampling (asymmetric pad, 8x reduction), bias-free linears, LayerNorm convolution, chunked-limited attention, and a relative-position KV cache.LSTMwith PyTorch->ONNX gate reordering(i,f,g,o)->(i,o,f,g)).[B, T, U, V]logits.New public API:
ParakeetMultiTalkerModel/ParakeetMultiTalkerConfigwith a.nemoloaderbuild_parakeet_multitalker, and aMultiTalkerRNNTTask(registered asmultitalker-rnnt).Testing
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.tests/parakeet_multitalker_integration_test.py,-m integration): streaming parity vs the NeMo PyTorch reference — encoder3.9e-7, caches2e-6/1.8e-4, decoder1.1e-6, joint log-softmax1.5e-5with matching argmax. Verified end-to-end against real weights.ruff check/ruff formatclean.Note
Independent of #410 — the small FastConformer
_ReLU/_swishprimitives are defined locally inmodels/parakeet_multitalker.py, so this PR has no code or git dependency on the Sortformer branch and targetsmaindirectly.