Multinode data sampling: fix

gaussalgo · Apr 29, 2024 · d5f8690 · d5f8690
1 parent 1778386
commit d5f8690
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/adaptor/adapter.py b/adaptor/adapter.py
@@ -184,5 +184,5 @@ def _load_optimizer_and_scheduler(self, checkpoint: str) -> None:
             logger.warning("Restoring training on global step %s", self.state.global_step)
 
         # in case of continued training, optimizer exists on model.model_name_or_path
-        # if the optmizer.pt does not exist, the `super()._load_optimizer_and_scheduler` does not do anything
+        # if the optimizer.pt does not exist, the `super()._load_optimizer_and_scheduler` does not do anything
         return super()._load_optimizer_and_scheduler(checkpoint=self.model.model_name_or_path)
diff --git a/adaptor/utils.py b/adaptor/utils.py
@@ -64,7 +64,7 @@ class AdaptationDataset(IterableDataset, abc.ABC):
     """
 
     def __init__(self, length: Optional[int] = None):
-        self.world_size = int(os.environ.get("LOCAL_WORLD_SIZE", 1))
+        self.world_size = int(os.environ.get("WORLD_SIZE", 1))
         if self.world_size > 1:
             logger.warning("World size for data sampling: %s" % self.world_size)
             self.length = length // self.world_size