fixed bugs and added commandline option

megatron/arguments.py

@@ -1252,6 +1252,8 @@ def _add_data_args(parser):
     group.add_argument('--data-file-list', type=str, default=None,
                        help='The file with the list of dataset and weights')
 
+    group.add_argument('--shuffle-sample', action='stored_true', help="Whether to shuffle the samples within in the dataset files")
+
     group.add_argument('--split', type=str, default='969, 30, 1',
                        help='Comma-separated list of proportions for training,'
                        ' validation, and test split. For example the split '

megatron/data/gpt_dataset.py

@@ -94,7 +94,7 @@ def Build(self):
 
         class BuildConcatDataset(torch.utils.data.Dataset):
             @dlp.log
-            def __init__(self, dataset_builders):
+            def __init__(self, dataset_builders, shuffle=False):
                 self.dataset_builders = dataset_builders
                 self.num_datasets = len(dataset_builders)
                 self.num_samples = np.sum([d.num_samples for d in dataset_builders])
@@ -117,7 +117,9 @@ def _build_indices():
 
                 self.dataset_index, self.dataset_sample_index = _build_indices()
                 np_rng = np.random.RandomState(seed=dataset_builders[0].seed)
-                self.shuffle_index=np_rng.shuffle(range(self.num_samples))
+                self.shuffle_index = np.arange(self.num_samples)
+                if shuffle:
+                    np_rng.shuffle(self.shuffle_index)
                 for i in range(self.num_datasets):
                     self.desc += dataset_builders[i].prefix + ","
 
@@ -146,7 +148,7 @@ def __getitem__(self, idx):
         valid_datasets = []
         test_datasets = []
         # Build individual datasets.
-
+        args = get_args()
         @dlp.log
         def build_corpus_datasets(dataset_type='train'):
             start_time = time.time()
@@ -172,7 +174,7 @@ def build_corpus_datasets(dataset_type='train'):
             print_rank_0(" > number of samples for each corpus ")
             corpus_weights_achieved={}
             for c in corpus_list:
-                datasets.append(BuildConcatDataset(corpus_builders[c]))
+                datasets.append(BuildConcatDataset(corpus_builders[c], args.shuffle_sample))
                 total += datasets[-1].num_samples
                 corpus_weights_achieved[c] =  float(datasets[-1].num_samples)/train_num_samples                
                 print_rank_0(f"    {c}: {datasets[-1].num_samples} w={corpus_weights_achieved[c]} (expected: {corpus_weights[c]})")