Early-interruption augmentation can drop the relocated EOS token entirely · Issue #16247 · NVIDIA-NeMo/Speech · GitHub
Skip to content

Early-interruption augmentation can drop the relocated EOS token entirely #16247

Description

@udsy19

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_id

When 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_pos

Calling _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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions