Seed _get_train_sampler's generator with arg seed to improve reproducibility (#15961)

* Seed get_train_sampler's generator with arg seed to improve reproducibility

and make the world_size<=1 code path more similar to the others

* move test file into trainer test explicitly

* dumb typo

* make style lint happy

* per discussion, switch to data_seed

* Apply suggestions from code review

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

tests/trainer/test_trainer.py

@@ -647,6 +647,67 @@ class TrainerIntegrationTest(TestCasePlus, TrainerIntegrationCommon):
         new_eval_dataset = RegressionDataset(length=128)
         self.assertEqual(len(trainer.get_eval_dataloader(new_eval_dataset)), 128 // (32 * n_gpu))
 
+    def test_sampler_seed(self):
+        # nb: we don't want to inherit from IterableDataset to hit the right code path
+        class DummyDataset(torch.utils.data.Dataset):
+            def __init__(self, length: int = 101):
+                self.length = length
+
+            def __len__(self):
+                return self.length
+
+            def __getitem__(self, i):
+                if (i < 0) or (i >= self.length):
+                    raise IndexError
+                return {"input_ids": [i]}
+
+        class DummyModel(PreTrainedModel):
+            def __init__(self, num_params: int):
+                super().__init__(PretrainedConfig())
+                # Add some (unused) params. the point here is that randomness in model_init shouldn't influence
+                # data loader order.
+                self.params = nn.Parameter(torch.randn(num_params))
+
+            def forward(self, input_ids, labels=None):
+                if labels is not None:
+                    return torch.tensor(0.0, device=input_ids.device), input_ids
+                else:
+                    return input_ids
+
+        def _get_first_data_sample(num_params, seed, data_seed, **kwargs):
+            with tempfile.TemporaryDirectory() as tmpdir:
+                trainer = Trainer(
+                    model_init=lambda: DummyModel(num_params),
+                    args=TrainingArguments(
+                        output_dir=tmpdir,
+                        **kwargs,
+                        seed=seed,
+                        data_seed=data_seed,
+                        local_rank=-1,
+                    ),
+                    train_dataset=DummyDataset(),
+                )
+
+                return next(iter(trainer.get_train_dataloader()))
+
+        # test that the seed is passed to the sampler
+        # the codepath we want to hit is world_size <= 1, and both group_by_length
+        for group_by_length in [True, False]:
+            sample42_1 = _get_first_data_sample(num_params=10, seed=42, data_seed=42, group_by_length=group_by_length)
+            sample42_2 = _get_first_data_sample(num_params=11, seed=42, data_seed=42, group_by_length=group_by_length)
+            self.assertTrue(torch.equal(sample42_1["input_ids"], sample42_2["input_ids"]))
+
+            # should get same samples with different seed, so long as data_seed is the same
+            sample42_3 = _get_first_data_sample(num_params=11, seed=11, data_seed=42, group_by_length=group_by_length)
+            self.assertTrue(torch.equal(sample42_1["input_ids"], sample42_3["input_ids"]))
+
+            # make sure we have some randomness in the samples if data_seed is different
+            others = [
+                _get_first_data_sample(num_params=i, seed=42, data_seed=i, group_by_length=group_by_length)
+                for i in range(10)
+            ]
+            self.assertTrue(any(not torch.equal(sample42_1["input_ids"], sample["input_ids"]) for sample in others))
+
     @require_torch_multi_gpu
     def test_data_is_not_parallelized_when_model_is_parallel(self):
         model = RegressionModel()