Add NVIDIA NeMo Sortformer speaker-diarization model - #410
Open
themason2011 wants to merge 4 commits into
Open
Conversation
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a new speaker-diarization export path to mobius by implementing NVIDIA NeMo’s Sortformer streaming diarization model as an ONNX graph (offline forward path), along with a new diarization task and tests to validate graph construction and NeMo parity.
Changes:
- Added
SortformerConfig,SortformerDiarizationModel, and a.nemoloader (build_sortformer) to reconstruct the model and apply NeMo weights. - Introduced
DiarizationTask(input_features → speaker_probs) and registered it under thediarizationtask name. - Added unit graph-build tests and an integration parity test against NeMo.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/sortformer_integration_test.py | Integration test comparing ONNX output to NeMo PyTorch offline inference. |
| tests/build_graph_test.py | New L1-style graph build + ORT execution smoke tests for Sortformer + DiarizationTask. |
| src/mobius/tasks/_diarization.py | New encoder-only diarization task wiring input_features to speaker_probs. |
| src/mobius/tasks/init.py | Exports/registers DiarizationTask and the diarization task mapping. |
| src/mobius/models/sortformer.py | New Sortformer model implementation + .nemo extraction/loading helper. |
| src/mobius/models/init.py | Exports Sortformer* symbols and build_sortformer. |
Comment on lines
+429
to
+430
| ff = self.feed_forward1(op, self.norm_feed_forward1(op, x)) | ||
| x = op.Add(x, op.Mul(ff, op.Constant(value_float=0.5))) |
Author
There was a problem hiding this comment.
Try to build sortformer using dtype=float16/bfloat16. If it doesn't work, we need to apply all data-type related suggestions from copilot. If it does work without changes, we can ignore and un-apply the previous comment
Comment on lines
+437
to
+438
| ff = self.feed_forward2(op, self.norm_feed_forward2(op, x)) | ||
| x = op.Add(x, op.Mul(ff, op.Constant(value_float=0.5))) |
Comment on lines
+491
to
+492
| # Scale inputs to the attention layers by sqrt(d_model) (xscaling). | ||
| x = op.Mul(x, op.Constant(value=ir.tensor(np.array(self._xscale, np.float32)))) |
Comment on lines
+525
to
+526
| scale = op.Constant(value=ir.tensor(np.array(self._head_dim**-0.5, dtype=np.float32))) | ||
| scores = op.Mul(op.MatMul(q, op.Transpose(k, perm=[0, 1, 3, 2])), scale) |
Comment on lines
+477
to
+485
| positions = op.Range(start, limit, op.Constant(value_int=-1)) # [2T-1] | ||
| pos_f = op.Unsqueeze(op.Cast(positions, to=ir.DataType.FLOAT), [-1]) # [2T-1,1] | ||
| div_term = op.Constant(value=ir.tensor(self._div_term)) # [d/2] | ||
| angles = op.Mul(pos_f, div_term) # [2T-1, d/2] | ||
| sin = op.Unsqueeze(op.Sin(angles), [-1]) | ||
| cos = op.Unsqueeze(op.Cos(angles), [-1]) | ||
| interleaved = op.Concat(sin, cos, axis=-1) # [2T-1, d/2, 2] | ||
| pe = op.Reshape(interleaved, [0, -1]) # [2T-1, d] | ||
| return op.Unsqueeze(pe, [0]) # [1, 2T-1, d] |
justinchuby
reviewed
Jul 17, 2026
justinchuby
reviewed
Jul 17, 2026
justinchuby
reviewed
Jul 17, 2026
justinchuby
reviewed
Jul 17, 2026
justinchuby
reviewed
Jul 17, 2026
themason2011
added a commit
that referenced
this pull request
Jul 17, 2026
Applies @justinchuby's PR #410 review feedback: - Use the opset-24 `op.Swish` op directly for Conformer feed-forward and convolution activations instead of the manual `x * sigmoid(x)` helper (removes the `_swish` function). - Replace the manual scaled-dot-product attention in `_TransformerAttention` with the opset-24 `op.Attention` op (bidirectional, unmasked). - Move the `.nemo` archive loader `build_sortformer` out of the modeling code into a dedicated NeMo integration package (`mobius.integrations.nemo`), mirroring the GGUF integration layout, and drop it from the public `mobius.models` namespace. Verified: sortformer graph unit tests pass and offline output stays bit-close to the NeMo reference (max abs diff 1.2e-7 vs golden). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <masoncorey@microsoft.com>
justinchuby
reviewed
Jul 17, 2026
Wire the NeMo Sortformer speaker-diarization model into mobius's central build_from_nemo loader instead of a bespoke build_sortformer entry point. - Register SortformerDiarizationModel under model_type 'sortformer' with the 'diarization' task in _registry.py. - Export SortformerConfig/SortformerDiarizationModel from models/__init__.py and DiarizationTask from tasks/__init__.py (TASK_REGISTRY 'diarization'). - Map the NeMo SortformerEncLabelModel target to model_type 'sortformer' and dispatch nemo_to_config to SortformerConfig.from_nemo_yaml before the FastConformer-RNNT encoder validation. - Add model_type field to SortformerConfig so the generic pipeline can resolve the model class and task from config.model_type. - Re-add graph-build unit tests and rewrite the integration test to use build_from_nemo. Verified: build_from_nemo output matches the NeMo golden reference (max abs diff 1.15e-7); 5 graph tests pass; ruff clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <masoncorey@microsoft.com>
themason2011
force-pushed
the
themason2011/sortformer-diarization
branch
from
July 20, 2026 21:30
c235538 to
ac3670a
Compare
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. |
justinchuby
approved these changes
Jul 21, 2026
justinchuby
left a comment
justinchuby
left a comment
Member
There was a problem hiding this comment.
Thanks! Could you address the copilot comments and ensure the unit test CI + lint CI are green?
Member
|
If possible, could you also add the L4/L5 tests? |
|
|
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (9)
src/mobius/models/sortformer.py:403
scaleis a float32 Constant, butmatrix_ac/matrix_bdfollow the model compute dtype (can be f16/bf16 whenbuild_from_nemo(..., dtype=...)is used). ONNXMulrequires matching dtypes, so this can produce an invalid graph for non-float32 builds. Cast the scalar to the scores dtype (the codebase typically usesop.CastLike).
scale = op.Constant(value=ir.tensor(np.array(self._head_dim**-0.5, dtype=np.float32)))
scores = op.Mul(op.Add(matrix_ac, matrix_bd), scale)
src/mobius/models/sortformer.py:486
_build_pos_emb()constructs positional embeddings in float32 (Cast(..., FLOAT)+ float32div_term), but those embeddings are fed intolinear_poswhich will use the model dtype (potentially f16/bf16 viabuild_from_nemo(dtype=...)). This dtype mismatch will breakMatMulinsideLinear. Cast the finalpos_embto the compute dtype before returning (keeping Sin/Cos in f32 if desired).
t_scalar = op.Squeeze(op.Shape(x, start=1, end=2)) # scalar T
one = op.Constant(value_int=1)
start = op.Sub(t_scalar, one) # T - 1
limit = op.Neg(t_scalar) # -T (exclusive) -> last value -(T-1)
positions = op.Range(start, limit, op.Constant(value_int=-1)) # [2T-1]
pos_f = op.Unsqueeze(op.Cast(positions, to=ir.DataType.FLOAT), [-1]) # [2T-1,1]
div_term = op.Constant(value=ir.tensor(self._div_term)) # [d/2]
angles = op.Mul(pos_f, div_term) # [2T-1, d/2]
sin = op.Unsqueeze(op.Sin(angles), [-1])
cos = op.Unsqueeze(op.Cos(angles), [-1])
interleaved = op.Concat(sin, cos, axis=-1) # [2T-1, d/2, 2]
pe = op.Reshape(interleaved, [0, -1]) # [2T-1, d]
return op.Unsqueeze(pe, [0]) # [1, 2T-1, d]
src/mobius/tasks/_diarization.py:42
feat_in if feat_in else "feat"treatsfeat_in=0as missing and falls back to a symbolic dimension. Iffeat_inis present, the check should beis not Noneto avoid surprising behavior and keep the shape contract consistent.
feat_in = getattr(config, "feat_in", None)
input_features = builder.input(
"input_features",
dtype=config.dtype,
shape=["batch", feat_in if feat_in else "feat", "time"],
)
src/mobius/models/sortformer.py:494
- The xscaling multiplier is emitted as a float32 constant, but
xis in the model compute dtype (potentially f16/bf16). This can make the graph invalid for non-float32 builds. UseCastLikeon the scalar constant (same pattern used in other models).
# Scale inputs to the attention layers by sqrt(d_model) (xscaling).
x = op.Mul(x, op.Constant(value=ir.tensor(np.array(self._xscale, np.float32))))
pos_emb = self._build_pos_emb(op, x)
src/mobius/integrations/nemo/_config_mapping.py:23
- The module-level docstring still claims this mapping only targets
ArchitectureConfig/ FastConformer-RNNT, but the module now also maps Sortformer diarization toSortformerConfig. Updating the top docstring prevents misleading documentation for NeMo support.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
src/mobius/models/sortformer.py:16
- The module docstring claims the subsampled time length is
T = T_mel / 8, but the dw_striding Conv2d stack uses stride=2 with padding=1/kernel=3, which yieldsceil(T_mel / 8)for odd lengths. This doc mismatch can confuse downstream users when shapes don’t divide evenly.
input_features [B, feat, T_mel]
-> transpose [B, T_mel, feat]
-> FastConformer encoder [B, T, fc_d_model] (T = T_mel / 8)
-> encoder_proj (Linear) [B, T, tf_d_model]
-> Transformer encoder (post-LN) [B, T, tf_d_model]
-> speaker sigmoid head [B, T, num_spks]
src/mobius/tasks/_diarization.py:25
- The output-frame formula in the docstring is inaccurate for the Sortformer subsampling stem (stride-2 convs with padding). The time dimension becomes
ceil(time / subsampling_factor)rather than a simple division whentimeisn’t an exact multiple.
Input: ``input_features`` — ``[batch, feat, time]`` mel spectrogram.
Output: ``speaker_probs`` — ``[batch, frames, num_spks]`` sigmoid
probabilities (``frames = time / subsampling_factor``).
"""
src/mobius/integrations/nemo/_config_mapping.py:23
- This file imports
BaseModelConfigfrom the private modulemobius._configs._base. Elsewhere (e.g.src/mobius/tasks/_diarization.py) the convention is to import config types from the publicmobius._configspackage. Using the public import keeps internal module boundaries consistent.
This issue also appears on line 22 of the same file.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
src/mobius/models/sortformer.py:87
- The PR description mentions a
build_sortformer(.nemo)loader implemented inmodels/sortformer.py, but there is no such function in this file (andbuild_sortformeris not found in the repo). Either add the helper (if intended API surface) or update the PR description to match the currentbuild_from_nemoentry point.
# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
@dataclasses.dataclass
class SortformerConfig(BaseModelConfig):
"""Configuration for the Sortformer diarization model.
Fields mirror the NeMo ``.nemo`` ``model_config.yaml`` (``encoder``,
``transformer_encoder`` and ``sortformer_modules`` sections).
"""
# Mel feature dimension (preprocessor ``features``).
feat_in: int = 128
# FastConformer encoder.
fc_d_model: int = 512
fc_num_layers: int = 17
fc_num_heads: int = 8
fc_ff_expansion: int = 4
fc_conv_kernel: int = 9
fc_subsampling_conv_channels: int = 256
fc_subsampling_factor: int = 8
# Transformer encoder (operates on projected embeddings).
tf_d_model: int = 192
tf_num_layers: int = 18
tf_num_heads: int = 8
tf_inner_size: int = 768
tf_hidden_act: str = "relu"
# Diarization head.
num_spks: int = 4
# HuggingFace/registry model_type — consumed by the generic
# ``build_from_nemo`` pipeline to resolve the model class and task.
model_type: str | None = "sortformer"
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <36940706+themason2011@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <36940706+themason2011@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (4)
src/mobius/models/sortformer.py:489
_build_pos_emb()returnspeas FLOAT (Sin/Cos are computed in float32), but the attention path later feeds this intoLinear/MatMultensors that followconfig.dtype(f16/bf16 possible viabuild_from_nemo(..., dtype=...)). Without castingpeto the compute dtype, ONNX will reject the graph due to mismatched element types.nemo_rnnt.pyhandles this by casting the positional encoding to the model dtype before projection (seesrc/mobius/models/nemo_rnnt.py:595-599).
interleaved = op.Concat(sin, cos, axis=-1) # [2T-1, d/2, 2]
pe = op.Reshape(interleaved, [0, -1]) # [2T-1, d]
return op.Unsqueeze(pe, [0]) # [1, 2T-1, d]
src/mobius/models/sortformer.py:434
op.Constant(value_float=0.5)creates a float32 scalar. If the model is built withconfig.dtype= f16/bf16, this will cause a dtype mismatch inMul/Add(ONNX requires matching element types). The repo already uses a cast-to-ref pattern for scalars (e.g._scalar_like()insrc/mobius/models/nemo_rnnt.py:55-63). Consider casting the scalar once and reusing it for both Macaron half-step residuals.
# Macaron feed-forward (half-step residual)
ff = self.feed_forward1(op, self.norm_feed_forward1(op, x))
x = op.Add(x, op.Mul(ff, op.Constant(value_float=0.5)))
src/mobius/models/sortformer.py:497
- The xscaling constant is emitted as a float32 tensor (
np.float32) and multiplied intox. If the graph is built withconfig.dtype= f16/bf16, this introduces a float32×float16Mul, which ONNX will reject. Please cast the scalar tox's dtype (same pattern as_scalar_like()insrc/mobius/models/nemo_rnnt.py:55-63).
# Scale inputs to the attention layers by sqrt(d_model) (xscaling).
x = op.Mul(x, op.Constant(value=ir.tensor(np.array(self._xscale, np.float32))))
pos_emb = self._build_pos_emb(op, x)
src/mobius/integrations/nemo/_config_mapping.py:23
BaseModelConfigis part of the publicmobius._configsAPI (re-exported fromsrc/mobius/_configs/__init__.py:24-29), but this file imports it from the private modulemobius._configs._base. SinceArchitectureConfigis already imported frommobius._configs, it would be more consistent to importBaseModelConfigfrom the same public module to avoid leaking internal paths into other packages.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
Author
lintrunner f --all-files |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (5)
src/mobius/models/sortformer.py:436
op.Constant(value_float=0.5)produces a float32 scalar; if the model is built withconfig.dtypeset to fp16/bf16 this will cause an ONNX type mismatch inMul. Cast the scalar to the compute dtype once and reuse it for both Macaron half-step residuals.
# Macaron feed-forward (half-step residual)
ff = self.feed_forward1(op, self.norm_feed_forward1(op, x))
x = op.Add(x, op.Mul(ff, op.Constant(value_float=0.5)))
# Relative-position self-attention
att = self.self_attn(op, self.norm_self_att(op, x), pos_emb, t_dim)
src/mobius/models/sortformer.py:496
- The xscaling multiplier is emitted as a float32 constant. For fp16/bf16 builds this will make
Mulinvalid due to mixed element types. Cast the constant tox's dtype (similar to_scalar_likeinnemo_rnnt.py).
x = op.Mul(x, op.Constant(value=ir.tensor(np.array(self._xscale, np.float32))))
src/mobius/models/sortformer.py:489
_build_pos_embcurrently returns float32 positional embeddings (due toto=FLOATand float32 div_term). In fp16/bf16 builds this will breaklinear_pos(op, pos_emb)becauseMatMulrequires matching dtypes. Cast the finalpeto the compute dtype (e.g.,CastLike(pe, x)) before returning.
pe = op.Reshape(interleaved, [0, -1]) # [2T-1, d]
return op.Unsqueeze(pe, [0]) # [1, 2T-1, d]
src/mobius/integrations/nemo/_config_mapping.py:23
BaseModelConfigis re-exported frommobius._configs, so importing it from the private module path (mobius._configs._base) is an unnecessary internal dependency. Import frommobius._configsto match the existingArchitectureConfigimport style.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
src/mobius/integrations/nemo/_config_mapping.py:23
- The module docstring still claims this mapping produces an
ArchitectureConfigand that only FastConformer-RNNT is supported. This is now inaccurate since Sortformer diarization is also handled bynemo_to_config, and it can return non-ArchitectureConfigconfigs.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
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
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 support for the NVIDIA NeMo Sortformer streaming speaker-diarization model (
nvidia/diar_streaming_sortformer_4spk-v2.1) to mobius.This is not a HuggingFace
transformersarchitecture — it ships as a NeMo.nemocheckpoint with noconfig.json. The architecture was reconstructed from the checkpoint and NeMo source. The exported ONNX graph implements the offline forward path (frontend_encoder+forward_infer): mel-spectrogram features in, per-frame speaker-activity probabilities out.Architecture
ConformerEncoder): rel_pos Transformer-XL attention, 8xdw_stridingconv subsampling, Macaron feed-forward blocks, batch-norm convolution module, xscaling.relu -> Linear -> relu -> Linear -> sigmoid).Module attribute names mirror the NeMo weight prefixes, so
preprocess_weightsonly drops unused keys (mel preprocessor,num_batches_tracked, the streaming FIFO head). The mel preprocessor (STFT) is left out of the graph, matching the mobius audio convention.Changes
models/sortformer.py—SortformerConfig(+from_nemo_yaml),SortformerDiarizationModel, and abuild_sortformer(.nemo)loader that extracts the archive, builds the graph, and applies checkpoint weights.tasks/_diarization.py—DiarizationTask(input_features -> speaker_probs), registered as thediarizationtask.models/__init__.pyand the task registry.Testing
tests/build_graph_test.py::TestBuildGraphSortformer): tiny-config graph build, I/O contract, initializers, task lookup, and an ORT forward run. All pass.tests/sortformer_integration_test.py,-m integration): compares the ONNX output against the real NeMo PyTorch reference. Offline output matches to 1.2e-7 (float32).Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com