Early-interruption augmentation can drop the relocated EOS token entirely
Describe the bug
DuplexSTTDataset._apply_early_interruption_augmentation (nemo/collections/speechlm2/data/duplex_stt_dataset.py) simulates an early user interruption by relocating a turn's EOS token earlier in the sequence, then padding the vacated tail. It writes the new EOS token before computing the tail-shift, then unconditionally overwrites the trailing frames_to_remove positions with the pad id:
target_tokens[batch_idx, new_eos_pos] = eos_id
seq_len = target_tokens.shape[1]
cont_start_pos = original_eos_pos + overlap_tokens
tail_length = seq_len - (cont_start_pos + 1)
if tail_length > 0:
target_tokens[batch_idx, new_eos_pos + 1 : new_eos_pos + 1 + tail_length] = target_tokens[
batch_idx, cont_start_pos + 1 : cont_start_pos + 1 + tail_length
].clone()
target_tokens[batch_idx, -frames_to_remove:] = pad_idWhen the turn's original EOS sits close enough to the end of the sequence that no tail-shift runs (cont_start_pos >= seq_len - 1), the final pad-fill range [-frames_to_remove:] starts at or before new_eos_pos and immediately overwrites the EOS token that was just written — the augmented turn ends up with no EOS token at all.
Steps/Code to reproduce bug
import torch
target_tokens = torch.full((1, 25), 0, dtype=torch.long)
target_tokens[0, 0] = 1 # BOS
target_tokens[0, 1:13] = torch.arange(10, 22) # content
target_tokens[0, 24] = 2 # original EOS at the last index
# with overlap_tokens=5, any cutoff in [1, 12] gives:
# cont_start_pos = 24 + 5 = 29 >= seq_len (25) -> no tail-shift runs
# frames_to_remove = 24 - cutoff -> pad-fill range covers new_eos_posCalling _apply_early_interruption_augmentation with this input leaves target_tokens[0] with zero occurrences of the EOS id — the relocated EOS is clobbered by the pad-fill.
Expected behavior
The relocated EOS token should survive the augmentation regardless of how close the original EOS sits to the end of the sequence.
Environment overview (please complete the following information)
- Environment location: Bare-metal
- Method of NeMo install: from source (editable install of a fresh clone of
main)
Environment details
- OS version: Ubuntu 24.04.4 LTS
- PyTorch version: 2.14.0+cpu
- Python version: 3.12.3
Additional context
early_interruption_prob is a documented, opt-in DuplexSTTDataset config knob (default 0.0) read directly from the model config at construction time; no shipped examples/speechlm2/conf/*.yaml currently enables it, but any training run that does silently loses the EOS token for the augmented turn whenever the original EOS is within overlap_tokens positions of the end of a short-enough sequence — the model would be trained on a turn with no end-of-turn signal. git log -S shows the function has been touched by exactly one commit since its introduction (#15092), so this is an oversight rather than a deliberate choice, and no test previously exercised this branch.
Source: NVIDIA-NeMo/Speech