Set seed after each epoch to improve consistency when resuming

fairseq/multiprocessing_trainer.py

@@ -57,9 +57,6 @@ class MultiprocessingTrainer(MultiprocessingEventLoop):
         """Initialize child processes."""
         self.args = args
 
-        # set torch.seed in this process
-        torch.manual_seed(args.seed)
-
         # set CUDA device
         torch.cuda.set_device(device_id)
 
@@ -142,6 +139,15 @@ class MultiprocessingTrainer(MultiprocessingEventLoop):
             self.lr_scheduler, cuda_device=device_id)
         return extra_state
 
+    def set_seed(self, seed):
+        Future.gen_list([
+            self.call_async(rank, '_async_set_seed', seed=seed)
+            for rank in range(self.num_replicas)
+        ])
+
+    def _async_set_seed(self, rank, device_id, seed):
+        torch.manual_seed(seed)
+
     def train_step(self, samples):
         """Do forward, backward and gradient step in parallel."""
         # PyTorch initializes gradient buffers lazily, so the first

train.py

@@ -133,6 +133,7 @@ def train(args, epoch, batch_offset, trainer, dataset, num_gpus):
     extra_meters = collections.defaultdict(lambda: AverageMeter())
 
     desc = '| epoch {:03d}'.format(epoch)
+    trainer.set_seed(args.seed + epoch)
     lr = trainer.get_lr()
     with progress_bar(itr, desc, leave=False) as t:
         for i, sample in data.skip_group_enumerator(t, num_gpus, batch_offset):