remove dummy_load, move gpu_ranks warning out of TrainingConfig

eole/config/run.py

@@ -15,7 +15,9 @@
 )
 from eole.transforms import get_transforms_cls
 from eole.constants import TransformType
+from eole.utils.logging import logger
 from pydantic import Field, field_validator, model_validator
+import torch
 
 
 # Models below somewhat replicate prior opt behavior to facilitate transition
@@ -82,6 +84,8 @@ def _validate_train_config(self):
                 self.save_data
             ), "-save_data should be set if \
                 want save samples."
+        if torch.cuda.is_available() and not self.training.gpu_ranks:
+            logger.warn("You have a CUDA device, should run with -gpu_ranks")
         return self
 
 
@@ -129,6 +133,8 @@ def _validate_predict_config(self):
         # TODO: do we really need this _all_transform?
         if self._all_transform is None:
             self._all_transform = self.transforms
+        if torch.cuda.is_available() and not self.gpu_ranks:
+            logger.warn("You have a CUDA device, should run with -gpu_ranks")
         return self
 
     def _update_with_model_config(self):
@@ -151,9 +157,7 @@ def _update_with_model_config(self):
         if os.path.exists(config_path):
             # logic from models.BaseModel.inference_logic
             model_config = build_model_config(config_dict.get("model", {}))
-            training_config = TrainingConfig(
-                **config_dict.get("training", {}), dummy_load=True
-            )
+            training_config = TrainingConfig(**config_dict.get("training", {}))
             training_config.world_size = self.world_size
             training_config.gpu_ranks = self.gpu_ranks
             # retrieve share_vocab from checkpoint config

eole/config/training.py

@@ -273,11 +273,6 @@ class TrainingConfig(
     score_threshold: float = Field(
         default=0.68, description="Threshold to filterout data"
     )
-    dummy_load: bool | None = Field(
-        default=False,
-        description="Ignore some warnings if we are only loading the configuration "
-        "prior to other operations, e.g. in `train_from` context.",
-    )
 
     @computed_field
     @cached_property
@@ -323,13 +318,6 @@ def get_model_path(self):
     def _validate_running_config(self):
         super()._validate_running_config()
         # self._validate_language_model_compatibilities_opts()
-        if (
-            torch.cuda.is_available()
-            and not self.gpu_ranks
-            and self.model_fields_set != set()
-            and not self.dummy_load
-        ):
-            logger.warn("You have a CUDA device, should run with -gpu_ranks")
         if self.world_size < len(self.gpu_ranks):
             raise AssertionError(
                 "parameter counts of -gpu_ranks must be less or equal "

eole/models/model_saver.py

@@ -52,10 +52,6 @@ def load_checkpoint(model_path):
                 config_dict = json.loads(os.path.expandvars(f.read()))
                 # drop data to prevent validation issues
                 config_dict["data"] = {}
-                if "training" in config_dict.keys():
-                    config_dict["training"]["dummy_load"] = True
-                else:
-                    config_dict["training"] = {"dummy_load": True}
                 _config = TrainConfig(**config_dict)
                 checkpoint["config"] = _config
         else: