save model in global rank 0 in multinode (#1357)

* save model in global rank 0 in multinode

* set_epoch only when training

Author: IceCreamMilkyTea

distributed/ddp-tutorial-series/multinode.py

@@ -71,7 +71,7 @@ def _save_snapshot(self, epoch):
     def train(self, max_epochs: int):
         for epoch in range(self.epochs_run, max_epochs):
             self._run_epoch(epoch)
-            if self.local_rank == 0 and epoch % self.save_every == 0:
+            if self.global_rank == 0 and epoch % self.save_every == 0:
                 self._save_snapshot(epoch)
 
 

distributed/minGPT-ddp/mingpt/trainer.py

@@ -111,7 +111,8 @@ def _run_batch(self, source, targets, train: bool = True) -> float:
         return loss.item()
 
     def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):
-        dataloader.sampler.set_epoch(epoch)
+        if train:
+            dataloader.sampler.set_epoch(epoch)
         for iter, (source, targets) in enumerate(dataloader):
             step_type = "Train" if train else "Eval"
             source = source.to(self.local_rank)