Remove unnecessary state dict and load_state_dict functions (mosaicml…

…#3361)

composer/devices/device.py

@@ -11,15 +11,13 @@
 import torch.nn
 from torch.optim import Optimizer
 
-from composer.core.serializable import Serializable
-
 __all__ = ['Device', 'T_nnModule']
 
 T_nnModule = TypeVar('T_nnModule', bound=torch.nn.Module)
 T_Batch = TypeVar('T_Batch')
 
 
-class Device(Serializable, ABC):
+class Device(ABC):
     """Abstract class for a device on which a model runs.
 
     Attributes:

composer/devices/device_cpu.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, TypeVar
+from typing import TypeVar
 
 import torch
 
@@ -34,11 +34,3 @@ def module_to_device(self, module: T_nnModule) -> T_nnModule:
 
     def tensor_to_device(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.to(self._device)
-
-    def state_dict(self) -> dict[str, Any]:
-        # CPU device has no RNG state
-        return {}
-
-    def load_state_dict(self, state: dict[str, Any]) -> None:
-        if len(state) != 0:
-            raise ValueError('CPU device has no state.')

composer/devices/device_gpu.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional, TypeVar
+from typing import Optional, TypeVar
 
 import torch
 import torch.backends.cuda
@@ -57,11 +57,3 @@ def module_to_device(self, module: T_nnModule) -> T_nnModule:
 
     def tensor_to_device(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.to(self._device, non_blocking=True)
-
-    def state_dict(self) -> dict[str, Any]:
-        return {
-            'rng': torch.cuda.get_rng_state(),
-        }
-
-    def load_state_dict(self, state: dict[str, Any]) -> None:
-        torch.cuda.set_rng_state(state['rng'])

composer/devices/device_hpu.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, TypeVar
+from typing import TypeVar
 
 import torch
 
@@ -34,11 +34,3 @@ def module_to_device(self, module: T_nnModule) -> T_nnModule:
 
     def tensor_to_device(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.to(self._device)
-
-    def state_dict(self) -> dict[str, Any]:
-        # HPU device has no RNG state
-        return {}
-
-    def load_state_dict(self, state: dict[str, Any]) -> None:
-        if len(state) != 0:
-            raise ValueError('HPU device has no state.')

composer/devices/device_mps.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from typing import Any, TypeVar
+from typing import TypeVar
 
 import torch
 import torch.cuda.amp
@@ -42,10 +42,3 @@ def module_to_device(self, module: T_nnModule) -> T_nnModule:
 
     def tensor_to_device(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.to(self._device)
-
-    def state_dict(self) -> dict[str, Any]:
-        return {}
-
-    def load_state_dict(self, state: dict[str, Any]) -> None:
-        if len(state) != 0:
-            raise ValueError('MPS device has no state.')

composer/devices/device_neuron.py

@@ -7,7 +7,7 @@
 
 import logging
 import os
-from typing import Any, TypeVar
+from typing import TypeVar
 
 import torch
 
@@ -43,10 +43,3 @@ def module_to_device(self, module: T_nnModule) -> T_nnModule:
 
     def tensor_to_device(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.to(self._device)
-
-    def state_dict(self) -> dict[str, Any]:
-        return {}
-
-    def load_state_dict(self, state: dict[str, Any]) -> None:
-        if len(state) != 0:
-            raise ValueError('Neuron device has no state.')

composer/devices/device_tpu.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, TypeVar
+from typing import TypeVar
 
 import torch
 
@@ -40,10 +40,3 @@ def module_to_device(self, module: T_nnModule) -> T_nnModule:
 
     def tensor_to_device(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor.to(self._device)
-
-    def state_dict(self) -> dict[str, Any]:
-        return {}
-
-    def load_state_dict(self, state: dict[str, Any]) -> None:
-        if len(state) != 0:
-            raise ValueError('TPU device has no state.')

composer/loggers/wandb_logger.py

@@ -184,29 +184,6 @@ def log_images(
             else:
                 wandb.log({name: list(wandb_images)}, step=step)
 
-    def state_dict(self) -> dict[str, Any]:
-        import wandb
-
-        # Storing these fields in the state dict to support run resuming in the future.
-        if self._enabled:
-            if wandb.run is None:
-                raise ValueError('wandb module must be initialized before serialization.')
-
-            # If WandB is disabled, most things are RunDisabled objects, which are not
-            # pickleable due to overriding __getstate__ but not __setstate__
-            if wandb.run.disabled:
-                return {}
-            else:
-                return {
-                    'name': wandb.run.name,
-                    'project': wandb.run.project,
-                    'entity': wandb.run.entity,
-                    'id': wandb.run.id,
-                    'group': wandb.run.group,
-                }
-        else:
-            return {}
-
     def init(self, state: State, logger: Logger) -> None:
         import wandb
         del logger  # unused