make AcceleratRunner a subclass of Accelerator

add TorchRunner
add DeepSpeedRunner

Author: ZhiyuanChen

.github/workflows/push.yaml

@@ -24,7 +24,7 @@ jobs:
       - name: Install dependencies
         run: pip install -r requirements.txt && pip install -e .
       - name: Install dependencies for testing
-        run: pip install pytest pytest-cov torch torcheval torchmetrics torchvision accelerate
+        run: pip install pytest pytest-cov
       - name: pytest
         run: pytest --cov=materialx --cov-report=xml --cov-report=html .
       - name: Upload coverage report for documentation

danling/__init__.py

@@ -19,7 +19,7 @@
 
 from lazy_imports import try_import
 
-from danling import metrics, modules, optim, registry, runner, tensors, typing, utils
+from danling import defaults, metrics, modules, optim, registry, runner, tensors, typing, utils
 
 from .metrics import (
     AverageMeter,
@@ -31,7 +31,7 @@
 )
 from .optim import LRScheduler
 from .registry import GlobalRegistry, Registry
-from .runner import AccelerateRunner, BaseRunner, TorchRunner
+from .runner import AccelerateRunner, BaseRunner, Config, DeepSpeedRunner, Runner, TorchRunner
 from .tensors import NestedTensor, PNTensor, tensor
 from .utils import (
     catch,
@@ -49,6 +49,7 @@
     from .metrics import Metrics, MultiTaskMetrics
 
 __all__ = [
+    "defaults",
     "metrics",
     "modules",
     "optim",
@@ -57,9 +58,12 @@
     "tensors",
     "utils",
     "typing",
+    "Config",
+    "Runner",
     "BaseRunner",
-    "AccelerateRunner",
     "TorchRunner",
+    "AccelerateRunner",
+    "DeepSpeedRunner",
     "LRScheduler",
     "Registry",
     "GlobalRegistry",

danling/runner/defaults.py danling/defaults.py

@@ -17,14 +17,15 @@
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 # See the LICENSE file for more details.
 
-DEFAULT_RUN_NAME = "Run"
-DEFAULT_EXPERIMENT_NAME = "DanLing"
-DEFAULT_EXPERIMENT_ID = "xxxxxxxxxxxxxxxx"
-DEFAULT_IGNORED_KEYS_IN_HASH = {
+RUN_NAME = "Run"
+EXPERIMENT_NAME = "DanLing"
+EXPERIMENT_ID = "xxxxxxxxxxxxxxxx"
+SEED = 1016
+IGNORED_NAMES_IN_METRICS = ("index", "epochs", "steps")
+IGNORED_NAMES_IN_HASH = {
     "timestamp",
-    "iters",
-    "steps",
-    "epochs",
+    "epoch",
+    "step",
     "results",
     "score_split",
     "score",

danling/modules/mlp/dense.py

@@ -35,10 +35,10 @@ def __init__(
         super().__init__()
         self.residual = residual
         self.linear = nn.Linear(in_features, out_features, bias=bias)
-        self.norm = getattr(nn, norm)(out_features) if norm else nn.Identity()
-        self.activation = getattr(nn, activation)() if activation else nn.Identity()
+        self.norm = getattr(nn, norm)(out_features) if norm else None
+        self.activation = getattr(nn, activation)() if activation else None
         self.dropout = nn.Dropout(dropout)
-        self.pool = getattr(nn, pool)(out_features) if pool else nn.Identity() if self.residual else None
+        self.pool = getattr(nn, pool)(out_features) if self.residual else None
 
     def forward(self, x):
         out = self.linear(x)

danling/runner/README.md

@@ -4,26 +4,26 @@ The Runner of DanLing sets up the basic environment for running neural networks.
 
 ## Components
 
-For cross-platform compatibilities, DanLing features a two-level Runner + RunnerState system.
+For cross-platform compatibilities, DanLing features a two-level Runner + Config system.
 
 ### PlatformRunner
 
 PlatformRunner implements platform-specific features like `step` and `prepare`.
 
-The Runner contains all runtime information that is irrelevant to the checkpoint (e.g. `world_size`, `rank`, etc.). All other information should be saved in `RunnerState`.
+The Runner contains all runtime information that is irrelevant to the checkpoint (e.g. `world_size`, `rank`, etc.). All other information should be saved in `Config`.
 
 Currently, only [`AccelerateRunner`][danling.runner.AccelerateRunner] is supported.
 
 ### [`BaseRunner`][danling.runner.BaseRunner]
 
-[`BaseRunner`](danling.runner.BaseRunner) defines shared attributes and implements platform-agnostic features, including `init_logging`, `results` and `scores`.
+[`BaseRunner`][danling.runner.BaseRunner] defines shared attributes and implements platform-agnostic features, including `init_logging`, `results` and `scores`.
 
-### [`RunnerState`][danling.runner.RunnerState]
+### [`Config`][danling.runner.Config]
 
-[`RunnerState`][danling.runner.RunnerState] stores the state of a run (e.g. `epochs`, `run_id`, `network`, etc.).
+[`Config`][danling.runner.Config] stores the state of a run (e.g. `epoch`, `run_id`, `network`, etc.).
 
-With `RunnerState` and corresponding weights, you can resume a run from any point.
-Therefore, all members in `RunnerState` will be saved in the checkpoint, and thus should be json serialisable.
+With `Config` and corresponding weights, you can resume a run from any point.
+Therefore, all members in `Config` will be saved in the checkpoint, and thus should be json serialisable.
 
 ## Experiments Management
 

danling/runner/__init__.py

@@ -17,19 +17,22 @@
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 # See the LICENSE file for more details.
 
-from . import defaults
 from .accelerate_runner import AccelerateRunner
 from .base_runner import BaseRunner
-from .state import RunnerState
+from .config import Config
+from .deepspeed_runner import DeepSpeedRunner
+from .runner import Runner
 from .torch_runner import TorchRunner
 from .utils import on_local_main_process, on_main_process
 
 __all__ = [
-    "RunnerState",
+    "Config",
+    "Runner",
     "BaseRunner",
+    "TorchRunner",
     "AccelerateRunner",
+    "DeepSpeedRunner",
     "TorchRunner",
     "on_main_process",
     "on_local_main_process",
-    "defaults",
 ]