fixes

ServiceNow · Jan 24, 2025 · b8f434f · b8f434f
1 parent 3cc977f
commit b8f434f
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/fast_llm/engine/checkpoint/config.py b/fast_llm/engine/checkpoint/config.py
@@ -7,7 +7,7 @@
 
 import yaml
 
-from fast_llm.config import Config, Field, FieldHint, FieldUpdate, check_field, config_class
+from fast_llm.config import Config, Field, FieldHint, FieldUpdate, check_field, config_class, skip_valid_if_none
 from fast_llm.engine.config_utils.data_type import DataType
 from fast_llm.utils import Assert
 
@@ -183,7 +183,7 @@ class CheckpointPathConfigBase(CheckpointConfigBase):
         default=None,
         desc="Custom timeout for lengthy operations.",
         hint=FieldHint.optional,
-        valid=check_field(Assert.gt, 0),
+        valid=skip_valid_if_none(check_field(Assert.gt, 0)),
     )
 
 

diff --git a/fast_llm/engine/training/config.py b/fast_llm/engine/training/config.py
@@ -312,7 +312,7 @@ class TrainingConfig(Config):
         desc="Timeout for lengthy operations such as checkpoint saving and loading,"
         " and dataset preparation and sampling.",
         hint=FieldHint.feature,
-        valid=check_field(Assert.gt, 0),
+        valid=skip_valid_if_none(check_field(Assert.gt, 0)),
     )
 
     def _validate(self) -> None:

diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py
@@ -418,7 +418,9 @@ def _save_checkpoint(
         }
         if metrics is not None:
             metadata["metrics"] = {key.value: value for key, value in metrics.items()}
-        self._multi_stage.save_checkpoint(config.get_save_config(checkpoint_directory), metadata)
+        self._multi_stage.save_checkpoint(
+            config.get_save_config(checkpoint_directory, timeout=self._config.training.timeout), metadata
+        )
 
         # Barrier to ensure everyone is done.
         safe_barrier(
@@ -447,7 +449,9 @@ def _load_checkpoint(self, config: TrainingCheckpointConfig, iteration: int) ->
         checkpoint_directory = config.get_save_directory(self._run.experiment_directory) / str(iteration)
         Assert.custom(pathlib.Path.is_file, checkpoint_directory / "ok")
 
-        metadata = self._multi_stage.load_checkpoint(config.get_load_config(checkpoint_directory))
+        metadata = self._multi_stage.load_checkpoint(
+            config.get_load_config(checkpoint_directory, timeout=self._config.training.timeout)
+        )
         self._optimizer.load(metadata["optimizer"])
         if "schedules" in metadata:
             # Backward compatibility.