[themason2011]

Summary

Add ONNX graph support for NVIDIA Nemotron 3 Diarization, including both an offline full-sequence model and a stateful streaming model for per-frame speaker-activity probabilities.

Changes

• Add Nemotron3DiarizationConfig to map the Hugging Face audio, head, and streaming settings, and register nemotron3_diarization with the existing offline diarization task.
• Implement the Sortformer-style audio path: feature stacking/subsampling, bidirectional RoPE transformer, projection, sub-pixel upsampling, and sigmoid speaker-classification head.
• Add the diarization-streaming task. Each graph invocation processes one mel-feature chunk with configurable look-ahead and returns that chunk’s speaker probabilities plus updated fixed-capacity Arrival-Order Speaker Cache (AOSC), FIFO buffers, occupancy counts, and compression state. The cache update includes score-based top-k compression.
• Export the model and task from their public modules and add a tiny model configuration for graph-build coverage.

Validation

• Add streaming graph tests for package creation, input/output names, task lookup, and three consecutive ONNX Runtime calls that feed each step’s cache state into the next and check probability ranges and cache capacities.
• The existing parameterized speech graph-build suite includes the new tiny model configuration for offline graph construction.
• The three-step execution test uses random weights; this PR does not add a numerical comparison against Hugging Face with real model weights.

Scope / limitations

• The graph consumes mel-spectrogram features, not raw audio. Streaming callers provide the chunk window, look-ahead length, and cache state on each invocation.
• The streaming export assumes a fully valid input window; it does not expose Hugging Face’s padding attention_mask for a potentially padded final chunk.
• Outputs are speaker-activity probabilities, not post-processed speaker turns or timestamps.

[github-actions]

Performance Comparison

Comparing 88fd6a1f → 8503dcce

Model	Metric	Baseline	Current	Delta
bert (feature-extraction)	model_size_bytes	359 KB	359 KB	+0.0%	⚪
bert (feature-extraction)	num_nodes	68	68	+0.0%	⚪
falcon	model_size_bytes	364 KB	364 KB	+0.0%	⚪
falcon	num_nodes	66	66	+0.0%	⚪
gemma2	model_size_bytes	428 KB	428 KB	+0.0%	⚪
gemma2	num_nodes	105	105	+0.0%	⚪
gpt2	model_size_bytes	324 KB	324 KB	+0.0%	⚪
gpt2	num_nodes	54	54	+0.0%	⚪
llama	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama	num_nodes	60	60	+0.0%	⚪
llama (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama (static-cache)	num_nodes	56	56	+0.0%	⚪
mamba (ssm-text-generation)	model_size_bytes	296 KB	296 KB	+0.0%	⚪
mamba (ssm-text-generation)	num_nodes	94	94	+0.0%	⚪
phi3	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3	num_nodes	58	58	+0.0%	⚪
phi3 (static-cache)	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3 (static-cache)	num_nodes	54	54	+0.0%	⚪
qwen2	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2	num_nodes	60	60	+0.0%	⚪
qwen2 (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2 (static-cache)	num_nodes	56	56	+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	265	265	+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	127	127	+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	450	450	+0.0%	⚪
t5 (seq2seq)	model_size_bytes	836 KB	836 KB	+0.0%	⚪
t5 (seq2seq)	num_nodes	176	176	+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]

🏗️ Architecture Diff

Comparing 88fd6a1f → 8503dcce

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)

[Copilot]

Copilot review overview

🟡 Changes recommended

Critical and moderate issues remain in model compatibility, streaming behavior, configuration extraction, test coverage, and registry coverage.

Review effort: Lite
Findings: 4 · 1

Open (5)

• Add coverage YAML or skip entry for Nemotron diarization · New
• Use matching MHA dimensions for Nemotron V projection · New
• Cast padding value to the scores dtype · New
• Allow variable-length final streaming chunks · New
• Add Hugging Face numerical parity tests for streaming · New

What changed in this PR

Adds Nemotron 3 Diarization ONNX support, including configuration, model registration, offline/stateful streaming graphs, cache handling, and tests.

Changes:

• Adds Nemotron model architecture and configuration extraction.
• Adds streaming task wiring with cache/FIFO handling.
• Adds registry/exports and graph execution tests.

File	Summary / findings
tests/​build_graph/​speech_test.py	Adds graph and runtime tests; needs HF numerical parity coverage for normal, final-chunk, and compressed-cache steps.
tests/​_test_configs.py	Adds a tiny Nemotron test configuration.
src/​mobius/​tasks/​_diarization_streaming.py	Adds streaming graph I/O; fixed dimensions and lookahead handling do not support shorter final chunks or processor streaming modes.
src/​mobius/​tasks/​__init__.py	Exports the streaming task.
src/​mobius/​models/​nemotron3_diarization.py	Implements the model and cache logic; contains MHA/KV shape incompatibility and dtype-invalid padding.
src/​mobius/​models/​__init__.py	Exports the Nemotron model.
src/​mobius/​_registry.py	Registers the model; missing coverage YAML or documented skip will fail the coverage suite.
src/​mobius/​_configs/​_base.py	Adds configuration extraction; nested configs must be normalized before field access.
src/​mobius/​_configs/​__init__.py	Exports the configuration.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[titaiwangms]

Full review of head 90435051 (read-only; no changes made). The config/model/task/registry split follows Mobius conventions, the offline path reuses DiarizationTask, and the new streaming task has an explicit cache-state I/O contract. The earlier Copilot comments about missing YAML, MHA tiny-config dimensions, dtype constants, fixed-length final chunks, and missing real-weight goldens have been addressed in subsequent commits; the findings below concern the current head.

Findings

1. Major — L4/L5 can pass with incorrect diarization probabilities. In src/mobius/_testing/parity.py:424-438, a failed elementwise allclose becomes AMBIGUOUS whenever the dominant speaker per frame is unchanged. Both diarization consumers in tests/e2e_golden_test.py:2243,2308 assert only result != FAIL, accepting that result. For example, golden [0.9, 0.6] versus ONNX [0.9, 0.1] passes the argmax check even though the second speaker changes from active to inactive and the numeric error far exceeds the stated tolerance. This also weakens the migrated Sortformer L4 check. Please require PASS/allclose for continuous-probability parity and add a regression test with a secondary speaker crossing the 0.5 threshold.

2. Major — the default offline graph silently changes semantics for long recordings. src/mobius/models/nemotron3_diarization.py:22-26,582-598 describes a single full-sequence forward with no chunking or speaker cache, while the Hugging Face offline forward chunks a whole recording. The PR itself limits its parity claim to one chunk (chunk_length * subsampling_factor, approximately 2,720 mel frames at the checkpoint defaults); the committed offline golden deliberately stays below this limit. Inputs beyond it still execute but no longer follow the reference inference path. Please decide on and enforce the public contract: reject/clearly constrain long offline inputs and route them to streaming, or implement the matching chunked offline path.

Smaller follow-ups

• src/mobius/tasks/_diarization.py:20-24 says speaker_probs has time / subsampling_factor frames, but Nemotron upsamples and returns approximately one output per mel frame. Make the shared task's output-frame rule model-specific. Also correct the streaming task's output-length formula for windows whose raw-frame count is not divisible by the subsampling factor (src/mobius/tasks/_diarization_streaming.py:49-55).
• The migrated testdata/golden/diarization/sortformer.json retains input_frames/num_spks provenance, while the new scripts/generate_golden.py emits num_frames/num_speakers; regenerate it or document that it was carried over. The Nemotron golden records a development Transformers version but no immutable Transformers source revision, making the reference implementation harder to reproduce.

Verification boundary: This was a static review, not an independent real-weight run. L1, L3, and lint passed on this head, but the L4 and L5 CI jobs each hit their one-hour timeout and were cancelled; that is not evidence that the model-specific goldens failed or passed. A targeted completed L4/L5 run is still needed. Short, non-subsampling-aligned final chunks and reduced-precision cache compression have not been independently compared with the HF reference here; these are open verification questions, not additional confirmed defects.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[themason2011]

Thanks for the thorough re-review — addressed all findings in b574814. Summary of changes and verification for each:

1. Major — L4/L5 can pass with incorrect diarization probabilities.
Removed compare_diarization_golden()'s AMBIGUOUS downgrade entirely (src/mobius/_testing/parity.py). Any elementwise allclose failure is now a hard FAIL; dominant-speaker argmax and active-speaker-set Jaccard are diagnostics only, never a substitute for the tolerance gate. Both L4/L5 consumers in tests/e2e_golden_test.py now assert result == ParityResult.PASS explicitly. Added a regression test reproducing your exact example (src/mobius/_testing/parity_test.py::test_secondary_speaker_threshold_crossing_fails, golden [0.9, 0.6] vs ONNX [0.9, 0.1]) confirming it now fails.

2. Major — the default offline graph silently changes semantics for long recordings.
Rewrote the offline forward() (src/mobius/models/nemotron3_diarization.py) as a chunked ONNX Loop that matches HuggingFace's real offline forward: the whole input is embedded once, then processed in chunk_length-sized windows (with chunk_right_context lookahead), reusing the AOSC + FIFO speaker-cache mechanism across chunks with offline-specific cache sizes (config.fifo_length / config.speaker_cache_update_period, exposed as offline_fifo_length / offline_speaker_cache_update_period on Nemotron3DiarizationConfig). This is no longer a "stay under the limit" contract — arbitrary-length recordings now follow the real chunked reference path. Verified with new golden data for a 6000-mel-frame recording (750 encoder embeds, well above chunk_length=340) that forces 3 real chunking iterations and triggers cache compression in HuggingFace's own reference implementation (testdata/cases/diarization/nemotron3_diarization_multichunk.yaml); tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[diarization/nemotron3_diarization_multichunk] passes against it. Also added tests/build_graph/speech_test.py::TestBuildGraphNemotron3DiarizationOffline (Loop presence, single/multi-chunk shapes, cache-compression trigger) since no offline ORT-execution tests existed before.

Smaller follow-ups:

• DiarizationTask's docstring (src/mobius/tasks/_diarization.py) is now model-agnostic — no longer asserts frames = time / subsampling_factor universally; points to the specific model class for its exact formula. (The streaming task's output-length formula for non-subsampling-aligned windows was already corrected in a prior commit — confirmed still accurate.)
• Regenerated testdata/golden/diarization/sortformer.json via real NeMo inference so its field names match the current schema (no more carried-over input_frames/num_spks).
• The Nemotron golden now records an immutable transformers_commit (resolved via a new _best_effort_package_commit() helper in scripts/generate_golden.py) alongside the dev version string, for reproducibility.

Verification for this pass: tests/build_graph/speech_test.py (214 passed), tests/e2e_golden_test.py -k "diarization or sortformer" -m integration (25 passed, against the real nvidia/Nemotron-3-Diarization and Sortformer checkpoints — a completed, non-timed-out run), and the broader tests/build_graph + src/ suite (no new regressions). Along the way also fixed an unrelated general-purpose optimizer bug found while building the offline Loop graph: FoldTransposedInitializerPass was removing nodes via the wrong graph object, crashing for any Transpose living inside a subgraph.

Full details in the commit message: b574814.

[Copilot]

Copilot review overview

🟡 Changes recommended

Golden generation ignores the selected device, and streaming lookahead lacks required bounds validation.

Review effort: Balanced
Findings: 3 · 1

Open (4)

• Honor requested device in offline generation · New
• Honor requested device in streaming generation · New
• Validate num_lookahead_frames bounds · New
• Correct comment about silence_embeds usage · New

[titaiwangms]

Follow-up review of head 9ba8220b (read-only; no changes made). Thank you for addressing the previous review. The diarization probability comparator now fails on any elementwise allclose violation, both golden consumers require PASS, and the threshold-crossing regression covers the reported false pass. The default offline model now uses a chunked ONNX Loop, with a 6000-mel-frame L4 golden to cover cache updates across multiple chunks. The shared diarization task's output-frame documentation and the Sortformer provenance fields were also updated. These changes address the two previously reported Major findings at the implementation level.

New finding: non-aligned non-final streaming windows desynchronize output and cache

Major — src/mobius/models/nemotron3_diarization.py:980-982,1047-1055,1111-1124, with the advertised input/output contract in src/mobius/tasks/_diarization_streaming.py:28-57.

The new output trim uses raw_window_frames - num_lookahead_frames * subsampling_factor, but the cache update stores ceil(raw_window_frames / subsampling_factor) - num_lookahead_frames whole embeddings. For a 13-frame window, factor 2, and two lookahead embeddings, the output covers 9 raw frames while the five new embeddings committed to FIFO/cache cover 10. Advancing the next window by nine frames processes the tenth frame twice; advancing by ten drops a frame from the output. This is a supported scenario according to the revised task contract and the new non-aligned-window shape test, which checks output length but not the across-call consumed-frame invariant.

Please either constrain non-final windows to a multiple of the subsampling factor (while retaining short, possibly unaligned final windows), or explicitly carry the partial-frame/feature-stacking state needed to make the emitted and cached boundaries agree. A two-call regression using an unaligned non-final window would establish the contract; matching against the HF reference would further validate the chosen behavior.

Other observations

• The config docstring in src/mobius/_configs/_base.py still describes offline inference as a single full-sequence pass and says the streaming_* settings are used only by streaming, although the new offline Loop uses the shared AOSC capacity/scoring configuration. Please update those comments.
• The new _fold_transpose.py change correctly removes a subgraph node from its owning graph. A non-blocking question remains: its fold cache is keyed only by initializer name and inserts replacements at the root; could two distinct, same-named initializers in sibling subgraphs be conflated? We have not established a reachable example in this PR's generated models, so this is not reported as a confirmed defect. A small graph-scoping regression would settle it.
• The committed short and multichunk Nemotron goldens record different Transformers commit SHAs. It would help to confirm that both reference revisions implement the same diarization behavior, or regenerate them from one pinned reference.

Verification boundary: This review was static; I did not independently execute the PR or inspect its binary golden tensors. The author reports a completed local real-checkpoint run; the latest L4/L5 CI jobs on this head are cancelled, not evidence of either a pass or a model failure. This follow-up has no independent fp16/bf16 or real-weight parity check for unaligned streaming windows.

[themason2011]

Thanks for the thorough re-review! Root-caused and fixed all 4 items in 8503dcce (pushed).

Major — streaming output/cache desync: Confirmed by reading the actual pinned HuggingFace reference (Nemotron3DiarizationForAudioFrameClassification.forward) rather than guessing. It was a genuine bug, not an inherent architectural limitation: our final output trim subtracted num_lookahead_frames * subsampling_factor from raw_num_frames, but HF's real reference trims to plain raw_num_frames (logits[:, :num_frames], no lookahead subtraction — that trim is a no-op safety net that only ever activates on a final, lookahead == 0 chunk). Our extra subtraction under-emitted exactly one frame whenever a non-final window wasn't a multiple of subsampling_factor — precisely your example (raw=13, factor=2, lookahead=2 → we emitted 9, cache committed 10).

Fixed by trimming directly to raw_num_frames (relying on ONNX Slice's documented end-index clamping), matching HF exactly. This is provably a no-op for every case already covered by tests (aligned non-final windows, and any lookahead == 0 final window) — so it's a pure bug fix with zero risk to existing goldens, and doesn't require either of your two proposed remediations (window-alignment constraint or extra cross-call state): once the trim matches HF, output length always exactly equals the frames committed to the FIFO/cache. Added the two-call regression test you asked for (tests/build_graph/speech_test.py::test_non_final_chunk_with_misaligned_window_has_exact_frame_count) — an unaligned non-final chunk followed by a second chunk consuming the returned cache state, asserting speaker_probs length exactly matches FIFO growth × subsampling_factor in both calls. It also caught a pre-existing test that had (accidentally) encoded the old buggy expected count — now corrected.

Config docstring: Corrected Nemotron3DiarizationConfig's docstring and the AOSC-policy field comment in _configs/_base.py — offline is now described as Loop-chunked (not single-pass), and the shared-vs-per-mode field split (only fifo_length/speaker_cache_update_period differ; the rest of the AOSC policy is shared via _run_chunk_and_update_cache) is now stated accurately.

_fold_transpose.py fold-cache scoping: Investigated — confirmed not reachable for any mobius-built graph. onnxscript's OpBuilder always registers initializers on the root graph (self._root._graph.register_initializer(...)), even while emitting into a Loop/If body; subgraphs only ever reference outer-scope initializers by name, never declare their own local ones. So every mobius graph has a single, globally unique initializer namespace by construction, and the pass's name-only cache can't conflate two distinct initializers. Added the regression test you suggested (test_transpose_inside_loop_body_folds_against_root_initializer) that pins this directly with a manually-built Loop body, so this invariant is now guarded rather than just asserted.

Golden provenance mismatch: Confirmed real — the multichunk golden was on a different transformers_commit than the other two (short + streaming, which already matched each other). Regenerated all three offline/streaming diarization goldens with the currently-installed transformers so all three now share one pinned commit (0291458166a6834a37b179a24c5235f242dadea0), same model revision as before.

Verified: targeted diarization/sortformer suite (25 passed), full fast suite (10,518 passed, only the same 4 pre-existing failures already confirmed present on origin/main, unrelated to this PR), lintrunner f/ruff clean.