Updates for Sortformer training by ipmedenn · Pull Request #16227 · NVIDIA-NeMo/Speech

ipmedenn · 2026-09-07T10:52:40Z

Important

The Update branch button 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 ?

This PR extends Sortformer training with the functionality required for offline and streaming 8-speaker diarization. It adds new training configurations, subsegment-based data sampling, logits-based and auxiliary losses, block-causal tail augmentation, speaker-counting metrics, and improvements to checkpoint handling and long-recording inference.

The existing probability-based BCELoss remains supported, and causal-tail augmentation is disabled by default.

Collection: ASR, Speaker Tasks, Common

Changelog

Add first-stage offline and second-stage streaming 8-speaker Sortformer training configurations with 80 ms and 10 ms output frames, respectively.
Add deterministic initialization and sampling based on the configured training seed.
Add the init_from_nemo_model configuration option for warm-starting selected model submodules from a .nemo checkpoint.
Add checkpoint averaging after training and repair scheduler max_steps when resuming.
Update the end-to-end diarization dataset with:
- a five-field output containing audio, lengths, targets, and speaker identities;
- subsegment-based training using one or two source chunks;
- speaker-count-aware chunk sampling;
- full-duration negative samples when no suitable speech subsegment is available;
- optional manifest path validation.
Add an oversampling distributed sampler that produces a fixed number of samples per rank and supports resuming within an epoch.
Add BCEWithLogitsLoss and clean up the existing probability-based BCELoss.
Allow Sortformer forward methods to return speaker and activity logits when required for training while keeping probabilities as the default output.
Add optional activity and phantom-speaker auxiliary losses.
Preserve warm-start compatibility with checkpoints that do not contain the optional activity head.
Add training and validation metrics for speaker-count accuracy and mean absolute error.
Add training-only block-causal tail augmentation to TransformerEncoder: the prefix remains bidirectional, while each tail block attends to the prefix and preceding blocks but not future blocks. The augmentation is disabled by default, and validation and inference continue to use full attention.
Optimize max_batch_dur feature extraction by processing duration-bounded sub-batches without repeated tensor concatenation, and warn when a single recording exceeds the limit.
Add loading of registered Sortformer models from Hugging Face to e2e_diarize_speech.py.
Add and update unit tests for the dataset, sampler, losses, model forward paths, auxiliary losses, attention masks, checkpoint averaging, metrics, and feature extraction.

Usage

The offline configuration is intended for the first training stage:

cd examples/speaker_tasks/diarization/neural_diarizer

python ./sortformer_diar_train.py \
    --config-path='../conf/neural_diarizer' \
    --config-name='sortformer_offline_8spk.yaml' \
    model.train_ds.manifest_filepath=/path/to/train_manifest.json \
    model.validation_ds.manifest_filepath=/path/to/validation_manifest.json

The streaming configuration is intended for the second training stage. The first-stage model can be specified through init_from_nemo_model in the configuration:

init_from_nemo_model:
  stage1:
    path: /path/to/SortformerDiarizer_offline_8spk.nemo
    include: ["encoder", "sortformer_modules"]

python ./streaming_sortformer_diar_train.py \
    --config-path='../conf/neural_diarizer' \
    --config-name='sortformer_streaming_8spk.yaml' \
    model.train_ds.manifest_filepath=/path/to/train_manifest.json \
    model.validation_ds.manifest_filepath=/path/to/validation_manifest.json

e2e_diarize_speech.py can load either a local checkpoint through model_path or a registered Hugging Face model through pretrained_name:

python ./e2e_diarize_speech.py \
    pretrained_name='nvidia/diar_streaming_sortformer_4spk-v2.1' \
    batch_size=1 \
    dataset_manifest=/path/to/diarization_manifest.json

GitHub 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:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

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

Follow-up to Core updates for Streaming Sortformer inference #16032 and Improve MT-Parakeet streaming inference and timestamps #16070.
User-facing documentation for the new training workflow will be added in a follow-up PR.

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

copy-pr-bot · 2026-09-07T10:52:43Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

nithinraok

What about doc updates?

ipmedenn · 2026-09-08T08:52:34Z

What about doc updates?

User-facing documentation for the new training workflow is planned as a follow-up PR. This PR includes docstrings and inline descriptions for the new configuration parameters, and I have noted the documentation follow-up in the PR description.

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

tango4j

Requesting changes to approve

tango4j · 2026-09-08T15:27:04Z

-        return preds
+        return torch.sigmoid(self.forward_speaker_logits(hidden_out))
+
+    def forward_activity_logits(self, hidden_out: torch.Tensor) -> Optional[torch.Tensor]:


Add more description on the purpose of this function as one of the loss calculation heads

Addressed in d2ee5e0

tango4j · 2026-09-08T17:42:45Z


-    def _get_aux_train_evaluations(self, preds, targets, target_lens) -> dict:
+    @staticmethod
+    def _speaker_count_metrics(


Please consider moving this to a stand-alone file in metrics folder.

Addressed in d2ee5e0

tango4j

Please make the changes in loss folder's aux_diarization_loss.py

tango4j · 2026-09-11T00:12:37Z

+import torch.nn.functional as F
+
+
+def activity_loss(


This activity_loss should be a class method just like other loss *.py file.
This is because "loss" is always expected to be a class type that has type definitions and variable initializations.

Please follow the convention of: class name "ActivityLoss"
input_types
output_types
init
forward

It should be just like:
https://github.com/NVIDIA-NeMo/Speech/blob/main/nemo/collections/asr/losses/bce_loss.py

Loss class should be referred as follows:

=========
activity_loss:
target: nemo.collections.asr.losses.aux_diarization_loss.ActivityLoss
reduction: mean

tango4j · 2026-09-11T00:13:58Z

+    return loss
+
+
+def phantom_loss(


This is the same as activity_loss.
loss should not be a function.
In NeMo Speech, it is expected to be a class that has forward method.

Please follow the convention of having: class name "PhantomLoss"
input_types
output_types
init
forward

It should be just like:
https://github.com/NVIDIA-NeMo/Speech/blob/main/nemo/collections/asr/losses/bce_loss.py

Loss class should be referred as follows:

=========
phantom_loss:
target: nemo.collections.asr.losses.aux_diarization_loss.PhantomLoss
reduction: mean

tango4j · 2026-09-11T03:48:27Z

+  loss:
+    _target_: nemo.collections.asr.losses.bce_loss.BCEWithLogitsLoss
+    reduction: mean
+


here, the following should be added:

activity_loss:
target: nemo.collections.asr.losses.aux_diarization_loss.ActivityLoss
reduction: mean

phantom_loss:
target: nemo.collections.asr.losses.aux_diarization_loss.PhantomLoss
reduction: mean

tango4j

Please construct the loss class as other loss files (just like bce_loss)

tango4j · 2026-09-11T03:49:35Z

+  loss:
+    _target_: nemo.collections.asr.losses.bce_loss.BCEWithLogitsLoss
+    reduction: mean
+


here, the following should be added:

activity_loss:
target: nemo.collections.asr.losses.aux_diarization_loss.ActivityLoss
reduction: mean

phantom_loss:
target: nemo.collections.asr.losses.aux_diarization_loss.PhantomLoss
reduction: mean

tango4j · 2026-09-11T03:50:54Z

+    return loss
+
+
+def phantom_loss(


Loss class should be referred as follows:

=========
phantom_loss:
target: nemo.collections.asr.losses.aux_diarization_loss.PhantomLoss
reduction: mean

tango4j · 2026-09-11T03:51:06Z

+import torch.nn.functional as F
+
+
+def activity_loss(


Loss class should be referred as follows:

=========
activity_loss:
target: nemo.collections.asr.losses.aux_diarization_loss.ActivityLoss
reduction: mean

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

nithinraok

LGTM

tango4j

Changes are made exactly as requested.

tango4j · 2026-09-11T21:35:41Z

/ok to test f497938

ipmedenn added 14 commits September 2, 2026 07:18

updating sorrtformer training scripts, adding ckpt averaging

7d3f950

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

updating nemo dataloader for sortformer

496edf8

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

adding OversamplingDistributedSampler

b36fde0

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

adding yaml configs for 8spk Sortformer

b72ab5d

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

logits based BCE loss for sortformer

08667f8

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

auxiliary losses for sortformer

7c78a29

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

Fix Sortformer scheduler max_steps on resume

2f4b6ec

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

Merge NVIDIA-NeMo/Speech main into update-sortformer-training

f4b5a52

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

Add block-causal tail augmentation to streaming Sortformer training

efa1bc1

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

Add speaker counting metrics for Sortformer training and validation

99da565

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

Optimize OOM-safe Sortformer feature extraction

6fb624e

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

Add Hugging Face model loading to Sortformer diarization

8108f6a

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

Update 8-speaker Sortformer config comments and license header

bcd71bc

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

Actualizing sortformer_streaming_8spk.yaml

4b5d1bf

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

github-actions Bot added ASR Speaker Tasks common labels Sep 7, 2026

ipmedenn requested a review from tango4j September 7, 2026 10:54

github-advanced-security Bot found potential problems Sep 7, 2026

View reviewed changes

Comment thread nemo/collections/asr/models/sortformer_diar_models.py Fixed

Comment thread nemo/collections/asr/models/sortformer_diar_models.py Fixed

Comment thread nemo/collections/asr/data/audio_to_diar_label.py Fixed

Addressing CodeQL issues

a2924f4

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

ipmedenn marked this pull request as ready for review September 7, 2026 16:14

ipmedenn requested review from nithinraok and stevehuang52 September 7, 2026 16:15

nithinraok requested changes Sep 7, 2026

View reviewed changes

Organize Sortformer auxiliary losses and metrics

d2ee5e0

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

tango4j requested changes Sep 9, 2026

View reviewed changes

tango4j requested changes Sep 11, 2026

View reviewed changes

tango4j reviewed Sep 11, 2026

View reviewed changes

tango4j requested changes Sep 11, 2026

View reviewed changes

Refactor auxiliary diarization losses into configurable classes

f497938

Signed-off-by: ipmedenn <65592416+ipmedenn@users.noreply.github.com>

ipmedenn requested review from nithinraok and tango4j September 11, 2026 13:26

nithinraok approved these changes Sep 11, 2026

View reviewed changes

tango4j approved these changes Sep 11, 2026

View reviewed changes

copy-pr-bot Bot deployed to public September 11, 2026 21:36 Active

copy-pr-bot Bot requested a deployment to test September 11, 2026 21:37 Waiting

copy-pr-bot Bot deployed to public September 11, 2026 21:41 Active

Uh oh!

Conversation

ipmedenn commented Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

copy-pr-bot Bot commented Sep 7, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nithinraok left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ipmedenn commented Sep 8, 2026

Uh oh!

tango4j left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tango4j left a comment

Choose a reason for hiding this comment

Uh oh!

tango4j Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tango4j Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tango4j left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nithinraok left a comment

Choose a reason for hiding this comment

Uh oh!

tango4j left a comment

Choose a reason for hiding this comment

Uh oh!

tango4j commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

ipmedenn commented Sep 7, 2026 •

edited

Loading

tango4j Sep 11, 2026 •

edited

Loading

tango4j Sep 11, 2026 •

edited

Loading