Merge pull request #14 from ghanvert/full-rewrite

Full rewrite

Author: ghanvert

examples/train.py

@@ -87,7 +87,7 @@ def validation_step(self, batch):
         batch_size=(2, 1, 1),
         # 'batch_size' can be an integer value or tuple, where elements are: (train_batch_size, val_batch_size).
         # Use an integer value to set the batch size equally for both sets.
-        optim=Optimizer.AdamW,
+        optimizer=Optimizer.AdamW,
         optim_kwargs={"lr": 0.001, "weight_decay": 0.01},  # Optimizer arguments as dictionary
         scheduler=Scheduler.LinearWithWarmup,
         scheduler_kwargs={"warmup_ratio": 0.03},  # Scheduler arguments as dictionary

examples/train_without_comments.py

@@ -59,7 +59,7 @@ def validation_step(self, batch):
     hps_config=HyperParameters(
         epochs=2,
         batch_size=(2, 1, 1),
-        optim=Optimizer.AdamW,
+        optimizer=Optimizer.AdamW,
         optim_kwargs={"lr": 0.001, "weight_decay": 0.01},
         scheduler=Scheduler.LinearWithWarmup,
         scheduler_kwargs={"warmup_ratio": 0.03},

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = accmt
-version = 1.7.7
+version = 1.8.0
 author = ghanvert
 author_email = [email protected]
 description = Accelerator Module and Trainer based on Accelerate library for simple distributed train processes, inspired by PyTorch Lightning.
@@ -24,7 +24,7 @@ install_requires =
     accelerate
     numpy
     PyYAML
-    accmt-cli
+    accmt-cli>=1.4.5
     pympler
     numba
 

src/accmt/__init__.py

@@ -23,14 +23,13 @@
 from .collate_fns import DataCollatorForLanguageModeling, DataCollatorForLongestSequence, DataCollatorForSeq2Seq
 from .dataloader_samplers import TemperatureSampler
 from .decorators import on_last_process, on_local_main_process, on_local_process, on_main_process, on_process
-from .handlers import Handler
 from .hyperparameters import HyperParameters, Optimizer, Scheduler
 from .modules import AcceleratorModule, ExtendedAcceleratorModule
 from .monitor import Monitor
 from .tracker import Aim, ClearML, CometML, DVCLive, MLFlow, TensorBoard, WandB
-from .trainer import Trainer, set_seed
-from .utility import prepare, prepare_array, prepare_dataframe
-from .utils import _precision_map
+from .trainer import Trainer
+from .utility import IS_CPU, IS_GPU, prepare, prepare_array, prepare_dataframe
+from .utils import _precision_map, get_seed, set_seed
 
 
 def allow_tf32(flag=True):
@@ -40,9 +39,6 @@ def allow_tf32(flag=True):
 
 allow_tf32()
 
-IS_CPU = bool(int(os.environ.get("ACCMT_CPU", 0)))
-IS_GPU = not IS_CPU
-
 _init_kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=86400))
 _dataloader_config = DataLoaderConfiguration(use_seedable_sampler=True)
 accelerator = Accelerator(

src/accmt/callbacks.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from abc import ABC
+from dataclasses import dataclass
 
 import torch
 from torch.optim import Optimizer
@@ -23,6 +24,7 @@
 from .states import TrainingState
 
 
+@dataclass
 class Callback(ABC):
     """
     Callback module containing different callback functions for different
@@ -38,7 +40,7 @@ class Callback(ABC):
         trainer (`Trainer`):
             Defined `Trainer` class.
         state (`TrainingState`):
-            Module's `TrainingState` class.
+            Reference to `TrainingState` class.
 
     Methods:
         on_fit_start (*optional*):
@@ -274,3 +276,90 @@ def on_evaluation_start(self):
     @override
     def on_evaluation_end(self):
         """Callback when evaluation ends."""
+
+
+# TODO there is a better way to do this, using a decorator like @register_callback("on_fit_start"), but
+# we'll implement that (probably) before release of version 2.0.
+@dataclass
+class CallbackMaster:
+    children: list[Callback]
+
+    def on_fit_start(self):
+        for child in self.children:
+            child.on_fit_start()
+
+    def on_fit_end(self):
+        for child in self.children:
+            child.on_fit_end()
+
+    def on_before_backward(self, loss: torch.Tensor):
+        for child in self.children:
+            child.on_before_backward(loss)
+
+    def on_after_backward(self):
+        for child in self.children:
+            child.on_after_backward()
+
+    def on_before_optimizer_step(self, optimizer: Optimizer):
+        for child in self.children:
+            child.on_before_optimizer_step(optimizer)
+
+    def on_after_optimizer_step(self, optimizer: Optimizer):
+        for child in self.children:
+            child.on_after_optimizer_step(optimizer)
+
+    def on_before_scheduler_step(self, scheduler: LRScheduler):
+        for child in self.children:
+            child.on_before_scheduler_step(scheduler)
+
+    def on_after_scheduler_step(self, scheduler: LRScheduler):
+        for child in self.children:
+            child.on_after_scheduler_step(scheduler)
+
+    def on_before_zero_grad(self, optimizer: Optimizer):
+        for child in self.children:
+            child.on_before_zero_grad(optimizer)
+
+    def on_after_zero_grad(self, optimizer: Optimizer):
+        for child in self.children:
+            child.on_after_zero_grad(optimizer)
+
+    def on_resume(self):
+        for child in self.children:
+            child.on_resume()
+
+    def on_save_checkpoint(self):
+        for child in self.children:
+            child.on_save_checkpoint()
+
+    def on_before_training_step(self, batch: Any):
+        for child in self.children:
+            child.on_before_training_step(batch)
+
+    def on_after_training_step(self):
+        for child in self.children:
+            child.on_after_training_step()
+
+    def on_before_validation_step(self, batch: Any):
+        for child in self.children:
+            child.on_before_validation_step(batch)
+
+    def on_after_validation_step(self):
+        for child in self.children:
+            child.on_after_validation_step()
+
+    def on_epoch_start(self):
+        for child in self.children:
+            child.on_epoch_start()
+
+    def on_epoch_end(self):
+        for child in self.children:
+            child.on_epoch_end()
+
+    def on_evaluation_start(self):
+        for child in self.children:
+            child.on_evaluation_start()
+
+    def on_evaluation_end(self):
+        for child in self.children:
+            child.on_evaluation_end()

src/accmt/dist_utils.py

@@ -18,7 +18,17 @@
 import torch.distributed as dist
 import torch.nn.functional as F
 
-from .utility import WORLD_SIZE
+from .utility import RANK, WORLD_SIZE
+from .utils import time_prefix
+
+
+def rprint(*args, rank: int = 0, add_time_prefix: bool = True, **kwargs):
+    """Print on a specific rank (default is main process)."""
+    if rank == RANK:
+        if add_time_prefix:
+            print("\n", f"{time_prefix()} ", *args, **kwargs, sep="")
+        else:
+            print("\n", *args, **kwargs, sep="")
 
 
 def pad_to(tensor: torch.Tensor, maximum: int) -> tuple[torch.Tensor, torch.Tensor]:

src/accmt/hyperparameters.py

@@ -93,14 +93,14 @@ def __init__(
         self,
         epochs: int = 1,
         batch_size: Union[int, tuple[int]] = 1,
-        optim: Union[str, Optimizer] = "SGD",
+        optimizer: Union[str, Optimizer] = "SGD",
         optim_kwargs: Optional[dict] = None,
         scheduler: Optional[Union[str, Scheduler]] = None,
         scheduler_kwargs: Optional[dict] = None,
     ):
         self.epochs = epochs
         self.batch_size = batch_size
-        self.optim = getattr(Optimizer, optim) if isinstance(optim, str) else optim
+        self.optimizer = getattr(Optimizer, optimizer) if isinstance(optimizer, str) else optimizer
         self._fix_kwargs(optim_kwargs)
         self.optim_kwargs = optim_kwargs if optim_kwargs is not None else {}
         self.scheduler = getattr(Scheduler, scheduler) if isinstance(scheduler, str) else scheduler
@@ -114,10 +114,10 @@ def from_config(cls, config: Union[str, dict]):
         elif "hps" in config:
             config = config["hps"]
 
-        valid_keys = {"epochs", "batch_size", "optim", "scheduler"}
+        valid_keys = {"epochs", "batch_size", "optimizer", "scheduler"}
         assert all(k in valid_keys for k in config.keys()), "You do not have valid keys. Please check documentation."
 
-        optimizer = config["optim"]
+        optimizer = config["optimizer"]
         assert "type" in optimizer, "'type' key is required in optimizer."
 
         scheduler = config["scheduler"] if "scheduler" in config else None
@@ -127,17 +127,19 @@ def from_config(cls, config: Union[str, dict]):
         return HyperParameters(
             epochs=config["epochs"],
             batch_size=config["batch_size"],
-            optim=optimizer["type"],
+            optimizer=optimizer["type"],
             optim_kwargs={k: v for k, v in optimizer.items() if k != "type"} if len(optimizer) > 1 else None,
             scheduler=scheduler["type"] if scheduler is not None else None,
-            scheduler_kwargs={k: v for k, v in scheduler.items() if k != "type"}
-            if scheduler is not None and len(scheduler) > 1
-            else None,
+            scheduler_kwargs=(
+                {k: v for k, v in scheduler.items() if k != "type"}
+                if scheduler is not None and len(scheduler) > 1
+                else None
+            ),
         )
 
     def to_dict(self) -> dict:
-        optim = self.optim if not isinstance(self.optim, str) else getattr(Optimizer, self.optim, None)
-        assert optim is not None, f"{optim} is not a valid optimizer."
+        optimizer = self.optimizer if not isinstance(self.optimizer, str) else getattr(Optimizer, self.optimizer, None)
+        assert optimizer is not None, f"{optimizer} is not a valid optimizer."
         scheduler = (
             self.scheduler if not isinstance(self.scheduler, str) else getattr(Scheduler, self.scheduler, "INVALID")
         )
@@ -146,7 +148,13 @@ def to_dict(self) -> dict:
         optim_kwargs = self.optim_kwargs if self.optim_kwargs is not None else {}
         schlr_kwargs = self.scheduler_kwargs if self.scheduler_kwargs is not None else {}
 
-        d = {"hps": {"epochs": self.epochs, "batch_size": self.batch_size, "optim": {"type": optim, **optim_kwargs}}}
+        d = {
+            "hps": {
+                "epochs": self.epochs,
+                "batch_size": self.batch_size,
+                "optimizer": {"type": optimizer, **optim_kwargs},
+            }
+        }
 
         if self.scheduler is not None:
             d["hps"]["scheduler"] = {"type": scheduler, **schlr_kwargs}
@@ -155,7 +163,7 @@ def to_dict(self) -> dict:
 
     def get_config(self) -> dict:
         hps = self.to_dict()["hps"]
-        _hps = {"epochs": hps["epochs"], "batch_size": hps["batch_size"], **hps["optim"]}
+        _hps = {"epochs": hps["epochs"], "batch_size": hps["batch_size"], **hps["optimizer"]}
         if "type" in _hps:
             t = _hps["type"]
             _hps["optimizer"] = t if isinstance(t, str) else t.__name__

src/accmt/modules.py

@@ -16,7 +16,10 @@
 from typing import Optional, Union
 
 import torch
+import torch.nn as nn
 from accelerate import Accelerator
+from torch.optim.lr_scheduler import LRScheduler
+from torch.optim.optimizer import Optimizer
 from typing_extensions import Any, override
 
 from .states import TrainingState
@@ -68,15 +71,16 @@ class AcceleratorModule(ABC):
             `torch.nn.Module`.
     """
 
+    accelerator: Accelerator = None
+    state: TrainingState = None
+    device: torch.device = None
     _implemented_collate_fn_train = False
     _implemented_collate_fn_val = False
-    _accelerator: Accelerator = None
-    _log_every: int = 1
     _extended = False
-    state: TrainingState = None
-    device: torch.device = None
-    status_dict: dict = None
-    batch_size: Union[int, tuple[int, int]] = None
+    model: nn.Module = None
+    teacher: Optional[nn.Module] = None
+    optimizer: Optimizer = None
+    scheduler: LRScheduler = None
 
     @override
     def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
@@ -87,14 +91,14 @@ def training_step(self, batch: Any) -> torch.Tensor:
         """Defines the training logic. Must return a loss tensor (scalar)."""
 
     @override
-    def validation_step(self, batch: Any) -> dict:
+    def validation_step(self, key: str, batch: Any) -> dict:
         """
         Defines the validation logic. Must return a dictionary containing
         each metric with predictions and targets, and also the loss value in the dictionary.
 
         Example:
             ```
-            # format is ==> "metric": (predictions, targets)
+            # format is ==> "metric": (predictions, targets, ...)
             return {
                 "loss": validation_loss_tensor, # (scalar tensor)
                 # with additional metrics:
@@ -129,18 +133,23 @@ def get_validation_dataloader(self, *args: Any, **kwargs: Any) -> Any:
         """Defines a custom PyTorch DataLoader class for validation."""
 
     def log(self, values: dict, log_kwargs: dict | None = {}):
-        if self._accelerator.is_main_process:
+        if self.accelerator.is_main_process:
             train_or_eval = "global_step" if self.model.training else "eval_global_step"
             if (self.status_dict[train_or_eval] + 1) % self._log_every == 0:
-                self._accelerator.log(values, step=self.status_dict[train_or_eval], log_kwargs=log_kwargs)
+                self.accelerator.log(values, step=self.status_dict[train_or_eval], log_kwargs=log_kwargs)
 
     def __init_subclass__(cls, **kwargs):
+        # check training step and validation_step functions
         if (
             cls.training_step == AcceleratorModule.training_step
             and cls.validation_step == AcceleratorModule.validation_step
         ):
-            raise TypeError("Subclasses of 'Trainer' must override 'training_step' and/or 'validation_step' methods.")
+            raise RuntimeError(
+                "Subclasses of 'Trainer' must override 'training_step' and 'validation_step' "
+                "(if evaluation is available)."
+            )
 
+        # check collate functions
         if cls.collate_fn_train != AcceleratorModule.collate_fn_train:
             cls._implemented_collate_fn_train = True
 
@@ -244,18 +253,18 @@ def backward(self, loss: torch.Tensor, **kwargs):
             `kwargs` (`Any`):
                 Extra arguments to be passed to 'accelerator.backward' function.
         """
-        self._accelerator.backward(loss, **kwargs)
+        self.accelerator.backward(loss, **kwargs)
 
     def step_optimizer(self):
-        self.state.optimizer.step()
+        self.optimizer.step()
 
     def step_scheduler(self):
-        self.state.scheduler.step()
+        self.scheduler.step()
 
     def step(self):
         """Step optimizer and scheduler (in that order). If there is no scheduler, it will be ignored."""
         self.step_optimizer()
-        if self.state.scheduler is not None:
+        if self.scheduler is not None:
             self.step_scheduler()
 
     def zero_grad(self, set_to_none: bool = True):
@@ -266,8 +275,12 @@ def zero_grad(self, set_to_none: bool = True):
             `set_to_none` (`bool`, *optional*, defaults to `True`):
                 Set gradients to `None` instead of `0`.
         """
-        self.state.optimizer.zero_grad(set_to_none=set_to_none)
+        self.optimizer.zero_grad(set_to_none=set_to_none)
 
     @override
     def training_step(self, batch: Any):
         pass
+
+    def __init_subclass__(cls, **kwargs):
+        # No call to super(), so it suppresses the behavior.
+        pass