Fix shuffle_idx in megatron/data/gpt_dataset.py

argonne-lcf · Oct 12, 2024 · 9de83a9 · 9de83a9
1 parent 89db92a
commit 9de83a9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/megatron/data/gpt_dataset.py b/megatron/data/gpt_dataset.py
@@ -117,6 +117,8 @@ 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.arange(self.num_samples)
+                np_rng.shuffle(self.shuffle_index)
                 self.shuffle_index=np_rng.shuffle(range(self.num_samples))
                 for i in range(self.num_datasets):
                     self.desc += dataset_builders[i].prefix + ","