update sampler

Signed-off-by: stevehuang52 <[email protected]>

examples/speechlm/conf/salm_fc_linear.yaml

@@ -43,7 +43,6 @@ data:
     # Notably, the data weights are controlled by either bucketing_weights
     # or concat_sampling_probabilities depending on the dataset type (tar and
     # non-tar).
-    # See audio_text_qa_dataset.py for details.
     concat_sampling_probabilities: null # When providing a list of datasets, this arg defines the sampling probabilities from each dataset when strategy='random'
     context_key: ${data.common.context_key}
     answer_key: ${data.common.answer_key}
@@ -59,6 +58,7 @@ data:
     max_duration: 24 # it is set for LibriSpeech, you may need to update it for your dataset
     min_duration: 0.1
     # tarred datasets
+    is_concat: false
     is_tarred: false
     tarred_audio_filepaths: null
     shuffle_n: 2048

nemo/collections/speechlm/data/data_sampler.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, IterableDataset
 from nemo.lightning.pytorch.plugins import MegatronDataSampler
 from nemo.utils import logging
 
@@ -36,6 +36,10 @@ def transform_dataloader(self, dataloader: DataLoader, consumed_samples: int = 0
             logging.info(f"Dataset {dataloader.dataset} does not have __len__ method. Skipping Megatron data sampler.")
             return dataloader
 
+        if isinstance(dataloader.dataset, IterableDataset):
+            logging.info(f"Dataset {dataloader.dataset} is an IterableDataset. Skipping Megatron data sampler.")
+            return dataloader
+
         from megatron.core import parallel_state
 
         from nemo.lightning.data import add_megatron_sampler