Fix config for train_single

allenai · Jan 17, 2025 · 5715ef3 · 5715ef3
1 parent 2d676bd
commit 5715ef3
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/src/olmo_core/internal/experiment.py b/src/olmo_core/internal/experiment.py
@@ -101,7 +101,18 @@ def run(self, config: ExperimentConfig):
             launch(config)
         elif self == SubCmd.dry_run:
             pass
-        elif self in (SubCmd.train, SubCmd.train_single):
+        elif self == SubCmd.train:
+            try:
+                train(config)
+            finally:
+                teardown_training_environment()
+        elif self == SubCmd.train_single:
+            if config.model.dp_config is not None:
+                log.warning("dp_config is set to %s, but you can't use data parallelism when running on a single node. Disabling.", config.model.dp_config)
+                config.model.dp_config = None
+            if config.model.tp_config is not None:
+                log.warning("tp_config is set to %s, but you can't use tensor parallelism when running on a single node. Disabling.", config.model.dp_config)
+                config.model.tp_config = None
             try:
                 train(config)
             finally: