Repository: mosaicml/composer

Issue:
Loading a checkpoint overwrites `max_duration`
Loading a checkpoint overwrites `max_duration`. This makes it pretty difficult to resume training from a late-in-training checkpoint. To do so, you'd have to manually extract optimizer and scheduler states, at least, and then create a new trainer with `load_only_weights=True`. I'm not exactly sure what the right solution is, but this workaround is too cumbersome to be reasonable.



Discussion:
Makes sense, the `max_duration` argument passed into the trainer should override whatever is in the checkpoint (i.e. `max_duration` should be a skip serialization field).

Resolving patch:
diff --git a/composer/core/state.py b/composer/core/state.py
--- a/composer/core/state.py
+++ b/composer/core/state.py
@@ -28,46 +28,6 @@
 
 logger = logging.getLogger(__name__)
 
-# These fields will be serialized directly using torch.save / torch.load
-DIRECT_SERIALIZATION_FIELDS = [
-    "last_batch_size",
-    "grad_accum",
-    "_precision",
-    "_max_duration",
-]
-
-# These fields will be serialized using .state_dict(), and loaded with .load_state_dict()
-STATE_DICT_SERIALIZATION_FIELDS = [
-    "model",
-    "_optimizers",
-    "_schedulers",
-    "_algorithms",
-    "_callbacks",
-    "scaler",
-    "timer",
-]
-
-# These fields will be serialized using .state_dict(), but will be skipped if DeepSpeed is enabled.
-# When DeepSpeed is being used, model and optimizer states are serialized directly by the DeepSpeed engine.
-STATE_DICT_SERIALIZATION_FIELDS_SKIP_DEEPSPEED = [
-    "model",
-    "_optimizers",
-]
-
-# These fields will not be serialized
-SKIP_SERIALIZATION_FIELDS = [
-    "loss",
-    "batch",
-    "batch_num_samples",
-    "batch_num_tokens",
-    "outputs",
-    "train_dataloader",
-    "evaluators",
-    "_steps_per_epoch",
-    "_precision_context",
-    "profiler",
-]
-
 
 def _default_precision_factory() -> Callable[[Union[str, Precision]], ContextManager]:
     """Returns a context manager to automatically cast to a specific precision.
@@ -138,6 +98,7 @@ class State(Serializable):
         loss (types.Tensors): The most recently computed loss.
         outputs (types.Tensors): The most recently computed output from the model's forward pass.
         timer (types.Timer): The timer that tracks training loop progress.
+        serialized_attributes (List[str]): The list of attributes which will be serialized in a checkpoint.
     """
 
     _max_duration: Time[int]
@@ -148,6 +109,21 @@ class State(Serializable):
     loss: types.Tensors
     outputs: types.Tensors
 
+    # These attributes will be serialized using .state_dict(), and loaded with .load_state_dict()
+    # All other attributes will not be serialized.
+    # For simplicity, omit the leading underscore for private attributes.
+    # For example, even though the optimizers are stored on the state
+    # as the "_optimizers" attribute, here we specify just "optimizers"
+    serialized_attributes = [
+        "model",
+        "optimizers",
+        "schedulers",
+        "algorithms",
+        "callbacks",
+        "scaler",
+        "timer",
+    ]
+
     def __init__(
             self,
             # model
@@ -286,38 +262,21 @@ def state_dict(self) -> types.StateDict:
         """Returns the state as a :class:`dict`."""
         state_dict: types.StateDict = {}
 
-        deepspeed_enabled = False
-        try:
-            import deepspeed
-            deepspeed_enabled = isinstance(self.model, deepspeed.DeepSpeedEngine)
-        except ImportError:
-            pass
-
         for state_field_name, state_field_value in self.__dict__.items():
-            if state_field_name in SKIP_SERIALIZATION_FIELDS:
-                continue
-            elif state_field_name in DIRECT_SERIALIZATION_FIELDS:
-                state_dict[state_field_name] = state_field_value
+            if state_field_name.lstrip("_") not in self.serialized_attributes:
                 continue
-            elif state_field_name in STATE_DICT_SERIALIZATION_FIELDS:
-                if deepspeed_enabled and state_field_name in STATE_DICT_SERIALIZATION_FIELDS_SKIP_DEEPSPEED:
-                    continue
-                if state_field_name == "model":
-                    # Save model directly instead of by class name, since model may be wrapped by DistributedDataParallel
-                    serialized_value = state_field_value.state_dict()
-                else:
-                    serialized_value = {
-                        obj.__class__.__qualname__: obj.state_dict()
-                        for obj in ensure_tuple(state_field_value)
-                        if obj is not None
-                    }
-                state_dict[state_field_name] = serialized_value
-
+            if state_field_name == "model":
+                # Save model directly instead of by class name, since model may be wrapped by DistributedDataParallel
+                serialized_value = state_field_value.state_dict()
             else:
-                raise RuntimeError(f"Unable to serialize field {state_field_name}")
+                serialized_value = {
+                    obj.__class__.__qualname__: obj.state_dict()
+                    for obj in ensure_tuple(state_field_value)
+                    if obj is not None
+                }
+            state_dict[state_field_name] = serialized_value
+
         state_dict["_is_model_ddp_wrapped"] = isinstance(self.model, DistributedDataParallel)
-        if deepspeed_enabled:
-            state_dict["_deepspeed_enabled"] = True
         return state_dict
 
     def load_model_state(self, state_dict: types.StateDict, strict: bool):
@@ -342,35 +301,24 @@ def load_state_dict(self, state: types.StateDict, strict: bool = False):
             state_dict (types.StateDict): object returned from call to :meth:`state_dict`.
         """
 
-        deepspeed_enabled = False
-        if "_deepspeed_enabled" in state:
-            deepspeed_enabled = state["_deepspeed_enabled"]
-
         for state_field_name, state_field_value in self.__dict__.items():
-            if state_field_name in SKIP_SERIALIZATION_FIELDS:
+            if state_field_name.lstrip("_") not in self.serialized_attributes:
                 continue
-            elif state_field_name in DIRECT_SERIALIZATION_FIELDS:
-                setattr(self, state_field_name, state[state_field_name])
-            elif state_field_name in STATE_DICT_SERIALIZATION_FIELDS:
-                if deepspeed_enabled and state_field_name in STATE_DICT_SERIALIZATION_FIELDS_SKIP_DEEPSPEED:
-                    continue
-                serialized_value = state[state_field_name]
-
-                if state_field_name == "model":
-                    self.load_model_state(state, strict=strict)
-                else:
-                    for target in ensure_tuple(state_field_value):
-                        if target is None:
-                            continue
-                        if target.__class__.__qualname__ not in serialized_value:
-                            warnings.warn(
-                                f"{target.__class__.__qualname__} was not found in the state_dict. Its state will NOT be restored",
-                                category=UserWarning)
-                            continue
-                        source = serialized_value[target.__class__.__qualname__]
-                        target.load_state_dict(source)
+            serialized_value = state[state_field_name]
+
+            if state_field_name == "model":
+                self.load_model_state(state, strict=strict)
             else:
-                raise RuntimeError(f"Unable to load field {state_field_name}")
+                for target in ensure_tuple(state_field_value):
+                    if target is None:
+                        continue
+                    if target.__class__.__qualname__ not in serialized_value:
+                        warnings.warn(
+                            f"{target.__class__.__qualname__} was not found in the state_dict. Its state will NOT be restored",
+                            category=UserWarning)
+                        continue
+                    source = serialized_value[target.__class__.__qualname__]
+                    target.load_state_dict(source)
 
     @property
     def batch_idx(self) -> int:
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
--- a/composer/trainer/trainer.py
+++ b/composer/trainer/trainer.py
@@ -401,6 +401,13 @@ def __init__(
                 model=self.state.model,
                 optimizer=optimizer,
             )
+            # The deepspeed engine is responsible for serializing the model and optimizer state,
+            # so these attributes should not be serialized with the composer state.
+            if "model" in self.state.serialized_attributes:
+                self.state.serialized_attributes.remove("model")
+
+            if "optimizers" in self.state.serialized_attributes:
+                self.state.serialized_attributes.remove("optimizers")
 
         # If using DeepSpeed, the model must be loaded from checkpoint after the engine has been
         # initialized, but if using PyTorch DDP, the model must be loaded before it is wrapped with


Test patch:
diff --git a/tests/test_state.py b/tests/test_state.py
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -9,7 +9,6 @@
 
 from composer.algorithms.channels_last import ChannelsLastHparams
 from composer.core import DataSpec, State, types
-from composer.core.state import DIRECT_SERIALIZATION_FIELDS, SKIP_SERIALIZATION_FIELDS, STATE_DICT_SERIALIZATION_FIELDS
 from composer.datasets.dataloader import DataloaderHparams
 from composer.datasets.hparams import DatasetHparams
 from composer.models.base import ComposerModel
@@ -40,15 +39,6 @@ def get_dummy_state(model: ComposerModel, train_dataloader: types.DataLoader, va
     return state
 
 
-def is_field_serialized(field_name: str) -> bool:
-    if field_name in STATE_DICT_SERIALIZATION_FIELDS or field_name in DIRECT_SERIALIZATION_FIELDS:
-        return True
-    elif field_name in SKIP_SERIALIZATION_FIELDS:
-        return False
-    else:
-        raise RuntimeError(f"Serialization method for field {field_name} not specified")
-
-
 def assert_state_equivalent(state1: State, state2: State):
 
     # tested separately
@@ -60,7 +50,7 @@ def assert_state_equivalent(state1: State, state2: State):
     ]
 
     for field_name in state1.__dict__.keys():
-        if field_name in IGNORE_FIELDS or not is_field_serialized(field_name):
+        if field_name in IGNORE_FIELDS or field_name.lstrip("_") not in state1.serialized_attributes:
             continue
 
         var1 = getattr(state1, field_name)