[BUG] MMapIndexedDatasetBuilder.merge_file_ does not properly populate _doc_idx
State: closed | Pull request: no
Labels: name: bug

name: stale

**Describe the bug**

The merge_file_ method merges the index.sizes list but leaves the doc_idx list untouched. This results in the following assertion error when trying to load the dataset:

```
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/indexed_dataset.py", line 293, in __init__
(head, rank=0, pid=4773)     assert self.sequence_lengths.shape[0] == self.document_indices[-1]
(head, rank=0, pid=4773) AssertionError
```

**To Reproduce**
* Construct a merged indexed dataset by calling `MMapIndexedDatasetBuilder.merge_file_` on one or more existing datasets (class from nemo/collections/nlp/data/language_modeling/megatron/indexed_dataset.py)
* Point data_prefix config to the merged dataset and launch training
* Then should see the above error

**Expected behavior**
doc_idx list should be properly merged

**Stack trace/logs**

```
(head, rank=0, pid=4773) [NeMo I 2024-07-20 16:51:33 utils:215] Building dataset splits with cls=GPTDataset, sizes=[-64, 0, 64], and config=GPTDatasetConfig(random_seed=1234, sequence_length=8192, blend=(['/k8s-nfs/nemo/merged_dataset'], [1.0]), blend_per_split=None, split='90,8,2', split_matrix=[(0, 0.9), (0.9, 0.98), (0.98, 1.0)], num_dataset_builder_threads=1, path_to_cache=None, mmap_bin_files=True, mock=False, tokenizer=<nemo.collections.common.tokenizers.huggingface.auto_tokenizer.AutoTokenizer object at 0x7d5250884ac0>, reset_position_ids=False, reset_attention_mask=False, eod_mask_loss=False, create_attention_mask=False, drop_last_partial_validation_sequence=True, add_extra_token_to_sequence=True)
(head, rank=0, pid=4773) [NeMo I 2024-07-20 16:51:33 utils:215] Load the _IndexReader from /k8s-nfs/nemo/merged_dataset.idx
(head, rank=0, pid=4773) [NeMo I 2024-07-20 16:51:33 utils:215] 	Extract the sequence lengths
(head, rank=0, pid=4773) [NeMo I 2024-07-20 16:51:33 utils:215] 	Extract the sequence pointers
(head, rank=0, pid=4773) [NeMo I 2024-07-20 16:51:33 utils:215] 	Extract the document indices
(head, rank=0, pid=4773) Error executing job with overrides: ['trainer.devices=8', 'trainer.num_nodes=4', 'run.name=llama3_8b_sky-2024-07-20-09-48-11-555945_nemo-rocky-4x-nexttoteonly_3']
(head, rank=0, pid=4773) Traceback (most recent call last):
(head, rank=0, pid=4773)   File "/root/sky_workdir/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py", line 63, in main
(head, rank=0, pid=4773)     trainer.fit(model)
(head, rank=0, pid=4773)   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 543, in fit
(head, rank=0, pid=4773)     call._call_and_handle_interrupt(
(head, rank=0, pid=4773)   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 43, in _call_and_handle_interrupt
(head, rank=0, pid=4773)     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
(head, rank=0, pid=4773)   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
(head, rank=0, pid=4773)     return function(*args, **kwargs)
(head, rank=0, pid=4773)   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 579, in _fit_impl
(head, rank=0, pid=4773)     self._run(model, ckpt_path=ckpt_path)
(head, rank=0, pid=4773)   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 948, in _run
(head, rank=0, pid=4773)     call._call_setup_hook(self)  # allow user to set up LightningModule in accelerator environment
(head, rank=0, pid=4773)   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 96, in _call_setup_hook
(head, rank=0, pid=4773)     _call_lightning_module_hook(trainer, "setup", stage=fn)
(head, rank=0, pid=4773)   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 159, in _call_lightning_module_hook
(head, rank=0, pid=4773)     output = fn(*args, **kwargs)
(head, rank=0, pid=4773)   File "/root/sky_workdir/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 1656, in setup
(head, rank=0, pid=4773)     self.build_train_valid_test_datasets()
(head, rank=0, pid=4773)   File "/root/sky_workdir/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 1566, in build_train_valid_test_datasets
(head, rank=0, pid=4773)     ).build()
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/blended_megatron_dataset_builder.py", line 126, in build
(head, rank=0, pid=4773)     datasets = self._build_blended_dataset_splits()
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/blended_megatron_dataset_builder.py", line 184, in _build_blended_dataset_splits
(head, rank=0, pid=4773)     megatron_datasets = self._build_megatron_datasets_parallel(
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/blended_megatron_dataset_builder.py", line 344, in _build_megatron_datasets_parallel
(head, rank=0, pid=4773)     _threading_helper(
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/blended_megatron_dataset_builder.py", line 327, in _threading_helper
(head, rank=0, pid=4773)     raise err
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/blended_megatron_dataset_builder.py", line 323, in _threading_helper
(head, rank=0, pid=4773)     megatron_datasets_split = future.result()
(head, rank=0, pid=4773)   File "/usr/lib/python3.10/concurrent/futures/_base.py", line 458, in result
(head, rank=0, pid=4773)     return self.__get_result()
(head, rank=0, pid=4773)   File "/usr/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
(head, rank=0, pid=4773)     raise self._exception
(head, rank=0, pid=4773)   File "/usr/lib/python3.10/concurrent/futures/thread.py", line 58, in run
(head, rank=0, pid=4773)     result = self.fn(*self.args, **self.kwargs)
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/blended_megatron_dataset_builder.py", line 388, in _build_megatron_dataset_splits
(head, rank=0, pid=4773)     low_level_dataset = self.cls.build_low_level_dataset(dataset_path, self.config)
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/gpt_dataset.py", line 143, in build_low_level_dataset
(head, rank=0, pid=4773)     return IndexedDataset(dataset_path, multimodal=False, mmap=config.mmap_bin_files)
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/indexed_dataset.py", line 351, in __init__
(head, rank=0, pid=4773)     self.initialize(path_prefix, multimodal, mmap)
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/indexed_dataset.py", line 376, in initialize
(head, rank=0, pid=4773)     self.index = _IndexReader(idx_path, self.multimodal)
(head, rank=0, pid=4773)   File "/root/sky_workdir/Megatron-LM/megatron/core/datasets/indexed_dataset.py", line 293, in __init__
(head, rank=0, pid=4773)     assert self.sequence_lengths.shape[0] == self.document_indices[-1]
(head, rank=0, pid=4773) AssertionError
```

**Environment overview (please complete the following information)**

 - Environment location: Docker
 - Method of NeMo install: `nvcr.io/nvidia/nemo:24.05.01`



This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.

This issue was closed because it has been inactive for 7 days since being marked as stale.