feat(data): add indexed resumable multimodal loading and exact token budgets (#16193)

* feat(data): add exact indexed multimodal runtime

Add the final indexed JSONL, native tar, WDSv2, and ShareGPT routing runtime together with exact audio-token estimation, resumable packed sampling, explicit audio/manifest failure policies, and focused runtime tests.

This is a data-only reconstruction from the original stacked work. SALM packed tensor shaping and model execution remain outside this topic.

Co-authored-by: Kunal Dhawan <kunaldhawan97@gmail.com>

Original-Commit: 8ef2ff0

Original-Commit: 08f4621

Original-Commit: 210a5b6bb6fc4e5eb11ca92d004820bc38cf836b

Original-Commit: 3355af1c17fcb2425a3120848b30117470315538

Original-Commit: 303f36e6308bb67d80a882cff3c427325ea52d7d

Original-Commit: e1f23096783973258534783d1c2dfcd6fe6f1302

Original-Commit: 95a77c3

Original-Commit: 64ed195

Original-Commit: 7176d60

Original-Commit: 8979b44

Original-Commit: 524c694

Original-Commit: acfdf7d

Original-Commit: 3d7a2b8

Original-Commit: a8ea945

Original-Commit: 28a0a3a

Original-Commit: fd7a491

Original-Commit: 958955b

Original-Commit: 36f2503
Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* feat(data): add indexed data validation and operations

Add index construction and conversion, atomic record validation and publication, ShareGPT route creation and relocation, full dataloader validation, materialized-SFT token accounting, and packed checkpoint progress analysis.

The full validator intentionally uses dense strict SALM collation in this topic; packed SALM batch construction is reconstructed in the dependent training topic.

Co-authored-by: Kunal Dhawan <kunaldhawan97@gmail.com>

Original-Commit: 8ef2ff0

Original-Commit: 2791228

Original-Commit: 64ed195

Original-Commit: 9c15fb2

Original-Commit: 262b771

Original-Commit: 524c694

Original-Commit: acfdf7d

Original-Commit: 8c7da42

Original-Commit: a591a8b

Original-Commit: 958955b

Original-Commit: bbfbc82

Original-Commit: a9e80e0

Original-Commit: 36f2503
Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* style(data): satisfy formatting and lint checks

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* build: pin Lhotse main for indexed data APIs

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* fix(data): close datastore readers after caching

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* fix(data): address indexed loading review

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* fix(data): preserve and override conversation system prompts

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* fix(data): preserve indexed tar byte ranges

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* fix(data): defer indexed tar validation to audio loading

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* build: pin Lhotse 2.0.0a5

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* fix(data): embed native tar routing in index packs

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* build: pin Lhotse 2.0.0a6

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* fix(data): preserve fixed-bucket caps in packed sampling

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* perf(data): parallelize native tar route construction

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* fix(data): diagnose invalid multimodal conversation turns

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* docs(data): document packed caps and route workers

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

* test(data): add sampling rate to tarred fixture

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

---------

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Author: pzelasko

docs/source/dataloaders.rst

@@ -23,11 +23,12 @@
     get_hidden_length_from_sample_length,
     speaker_to_target,
 )
+from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin
 from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
 from nemo.utils import logging
 
 
-class LhotseAudioToSpeechE2ESpkDiarDataset(torch.utils.data.Dataset):
+class LhotseAudioToSpeechE2ESpkDiarDataset(LhotseAudioLoadingDatasetMixin, torch.utils.data.Dataset):
     """
     This dataset is a Lhotse version of diarization dataset in audio_to_diar_label.py.
     Unlike native NeMo datasets, Lhotse dataset defines only the mapping from
@@ -94,7 +95,7 @@ def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
             speaker_activities.append(speaker_activity)
 
         cuts = type(cuts).from_cuts(mono_cuts)
-        audio, audio_lens, cuts = self.load_audio(cuts)
+        audio, audio_lens, cuts = self.load_audio_with_cuts(cuts)
         max_num_speakers = max(1, max(activity.shape[1] for activity in speaker_activities))
         speaker_activities = [
             torch.nn.functional.pad(activity, (0, max_num_speakers - activity.shape[1]))

docs/source/speechlm2/datasets.rst

@@ -26,6 +26,7 @@
 
 from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
 from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
+from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin
 from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper
 from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
 from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
@@ -73,7 +74,7 @@ class RandomPaddingConfig:
     post_pad_duration: float = 3.0  # amount of right-padding when pad_distribution='constant'
 
 
-class LhotseSpeechToTextBpeEOUDataset(torch.utils.data.Dataset):
+class LhotseSpeechToTextBpeEOUDataset(LhotseAudioLoadingDatasetMixin, torch.utils.data.Dataset):
     """
     This dataset processes the audio data and the corresponding text data to generate the ASR labels,
     along with EOU labels for each frame. The audios used in this dataset should only contain speech with
@@ -206,7 +207,7 @@ def _check_special_tokens(self, tokenizer: TokenizerSpec):
             )
 
     def __getitem__(self, cuts: CutSet) -> AudioToTextEOUBatch:
-        audio, audio_lens, cuts = self.load_audio(cuts)
+        audio, audio_lens, cuts = self.load_audio_with_cuts(cuts)
         audio_signals = []
         audio_lengths = []
         eou_targets = []

nemo/collections/asr/data/audio_to_diar_label_lhotse.py

@@ -21,12 +21,13 @@
 from lhotse.dataset import AudioSamples
 from lhotse.dataset.collation import collate_vectors
 
+from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin
 from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper
 from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
 from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
 
 
-class LhotseSpeechToTextBpeDataset(torch.utils.data.Dataset):
+class LhotseSpeechToTextBpeDataset(LhotseAudioLoadingDatasetMixin, torch.utils.data.Dataset):
     """
     This dataset is based on BPE datasets from audio_to_text.py.
     Unlike native NeMo datasets, Lhotse dataset defines only the mapping from
@@ -65,7 +66,7 @@ def __init__(self, tokenizer: TokenizerSpec, return_cuts: bool = False):
         self.return_cuts = return_cuts
 
     def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
-        audio, audio_lens, cuts = self.load_audio(cuts)
+        audio, audio_lens, cuts = self.load_audio_with_cuts(cuts)
         tokens = [
             torch.cat(
                 [

nemo/collections/asr/data/audio_to_eou_label_lhotse.py

@@ -21,12 +21,13 @@
 from lhotse.dataset import AudioSamples
 from lhotse.dataset.collation import collate_matrices, collate_vectors
 
+from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin
 from nemo.collections.common.tokenizers.aggregate_tokenizer import AggregateTokenizer
 from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
 from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
 
 
-class LhotseSpeechToTextBpeDatasetWithPrompt(torch.utils.data.Dataset):
+class LhotseSpeechToTextBpeDatasetWithPrompt(LhotseAudioLoadingDatasetMixin, torch.utils.data.Dataset):
     """
     Dataset class for speech-to-text with prompt vectors.
     Supports both language ID and custom prompts.
@@ -115,7 +116,7 @@ def get_hidden_length_from_sample_length(self, num_samples: int) -> int:
         return int(hidden_length)
 
     def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
-        audio, audio_lens, cuts = self.load_audio(cuts)
+        audio, audio_lens, cuts = self.load_audio_with_cuts(cuts)
         tokens = [torch.as_tensor(self.tokenizer(c.supervisions[0].text, c.supervisions[0].language)) for c in cuts]
 
         # Create prompt targets

nemo/collections/asr/data/audio_to_text_lhotse.py

@@ -26,13 +26,14 @@
 from lhotse.dataset import AudioSamples
 from lhotse.dataset.collation import collate_vectors
 
+from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin
 from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper
 from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
 from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
 from nemo.utils import logging
 
 
-class LhotseSpeechToTextBpeDatasetWithPromptIndex(torch.utils.data.Dataset):
+class LhotseSpeechToTextBpeDatasetWithPromptIndex(LhotseAudioLoadingDatasetMixin, torch.utils.data.Dataset):
     """
     Simplified dataset class for speech-to-text with prompt support.
 
@@ -136,7 +137,7 @@ def _get_prompt_index_for_cut(self, cut) -> int:
             return self._get_prompt_index(cut.supervisions[0].language)
 
     def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
-        audio, audio_lens, cuts = self.load_audio(cuts)
+        audio, audio_lens, cuts = self.load_audio_with_cuts(cuts)
         tokens = [torch.as_tensor(self.tokenizer(c.supervisions[0].text, c.supervisions[0].language)) for c in cuts]
 
         # Get prompt indices (just the language ID per sample, NOT full tensors)

nemo/collections/asr/data/audio_to_text_lhotse_prompt.py

@@ -23,6 +23,7 @@
 
 from nemo.collections.asr.data.audio_to_text_lhotse import _make_audio_samples
 from nemo.collections.common.data import apply_prompt_format_fn
+from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin
 from nemo.collections.common.prompts import PromptFormatter
 from nemo.collections.common.tokenizers import TokenizerSpec
 
@@ -48,7 +49,7 @@ def get_decoder_inputs_outputs(self) -> tuple[torch.Tensor, torch.Tensor]:
         return self.prompted_transcript[:, :-1], self.prompted_transcript[:, 1:]
 
 
-class PromptedAudioToTextLhotseDataset(torch.utils.data.Dataset):
+class PromptedAudioToTextLhotseDataset(LhotseAudioLoadingDatasetMixin, torch.utils.data.Dataset):
     """
     This dataset is based on :class:`~nemo.collections.asr.data.audio_to_text_lhotse.LhotseSpeechToTextBpeDataset`.
     It is a Lhotse-style dataset that converts a mini-batch of Cuts into tensors.
@@ -97,7 +98,7 @@ def __init__(
 
     def __getitem__(self, cuts: CutSet) -> PromptedAudioToTextMiniBatch:
         # Load the audio's from AIS and add them to the CutSet
-        audio, audio_lens, cuts = self.load_audio(cuts)
+        audio, audio_lens, cuts = self.load_audio_with_cuts(cuts)
 
         # Will work if batch_size is set to 1.
         if self.enable_chunking:

nemo/collections/asr/data/audio_to_text_lhotse_prompt_index.py

@@ -22,11 +22,12 @@
 
 from nemo.collections.asr.data.audio_to_text_lhotse import TokenizerWrapper
 from nemo.collections.asr.parts.utils.asr_multispeaker_utils import speaker_to_target
+from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin
 from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
 from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
 
 
-class LhotseSpeechToTextSpkBpeDataset(torch.utils.data.Dataset):
+class LhotseSpeechToTextSpkBpeDataset(LhotseAudioLoadingDatasetMixin, torch.utils.data.Dataset):
     """
     This dataset is based on BPE datasets from audio_to_text.py. It has the same functionality of LhotseSpeechToTextBpeDataset but also yield speaker target tensor.
     Unlike native NeMo datasets, Lhotse dataset defines only the mapping from
@@ -60,7 +61,7 @@ def __init__(self, cfg, tokenizer: TokenizerSpec):
 
     def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
 
-        audio, audio_lens, cuts = self.load_audio(cuts)
+        audio, audio_lens, cuts = self.load_audio_with_cuts(cuts)
 
         tokens = []
         spk_targets = []

nemo/collections/asr/data/audio_to_text_lhotse_prompted.py

@@ -32,6 +32,7 @@
 from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
 from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
 from nemo.collections.common.data.dataset import ConcatDataset
+from nemo.collections.common.data.lhotse.audio_loading import LhotseAudioLoadingDatasetMixin
 from nemo.collections.common.parts.preprocessing.manifest import get_full_path
 from nemo.core.classes import Serialization
 from nemo.utils import logging
@@ -438,7 +439,7 @@ def _collate_fn(self, batch: List[AudioNoiseItem]) -> AudioNoiseBatch:
         return _audio_noise_collate_fn(batch, self.batch_augmentor)
 
 
-class LhotseAudioNoiseDataset(torch.utils.data.Dataset):
+class LhotseAudioNoiseDataset(LhotseAudioLoadingDatasetMixin, torch.utils.data.Dataset):
     def __init__(self, noise_manifest: str | None = None, batch_augmentor_cfg: DictConfig = None):
         super().__init__()
 
@@ -453,7 +454,7 @@ def __init__(self, noise_manifest: str | None = None, batch_augmentor_cfg: DictC
 
     def __getitem__(self, cuts):
 
-        audios, audio_lens, cuts = self.load_audio(cuts)
+        audios, audio_lens, cuts = self.load_audio_with_cuts(cuts)
         if len(self.noise_data) > 0:
             sampled_noises = [sample_noise(self.noise_data, cut.sampling_rate, cut.num_samples) for cut in cuts]
             sampled_noises, sampled_noises_lens = zip(*sampled_noises)