AutoYAHP Part 2: Cleanup the Callbacks and Loggers for AutoYAHP (#1071)

Similar to #1056, this PR cleans up the callbacks and loggers for AutoYAHP. It does not depend on AutoYAHP itself (a future PR will remove the underlying hparam classes).

- Refactored callback and logger tests to not depend on hparams
- Reformatted the docstrings so they would be correctly parsed by auto-yahp
- Added `callback_settings.py`, similar to `algorithm_settings.py`, to return a list of pytest param objects for parameterization across callback tests. These param objects include appropriate markers (e.g. conditional skipping for wandb and mlperf; requiring that the memory monitor runs on GPU, ...)
- Moved the `TestTrainerAssets` into `tests/callbacks/test_callbacks.py`, since it tests the individual callbacks and loggers, not the trainer, and thus should live in `tests/callbacks`.
- Cleaned up the `MemoryMonitor` warnings when cuda is not available. Now, it warns when the model is not on cuda, to catch the edge case where one does CPU training when GPUs are available.

Author: ravi-mosaicml

composer/callbacks/callback_hparams.py

@@ -29,7 +29,11 @@
     "MemoryMonitorHparams",
     "LRMonitorHparams",
     "SpeedMonitorHparams",
+    "EarlyStopperHparams",
     "CheckpointSaverHparams",
+    "ThresholdStopperHparams",
+    "MLPerfCallbackHparams",
+    "callback_registry",
 ]
 
 
@@ -321,3 +325,15 @@ def initialize_object(self) -> CheckpointSaver:
             weights_only=self.weights_only,
             num_checkpoints_to_keep=self.num_checkpoints_to_keep,
         )
+
+
+callback_registry = {
+    "checkpoint_saver": CheckpointSaverHparams,
+    "speed_monitor": SpeedMonitorHparams,
+    "lr_monitor": LRMonitorHparams,
+    "grad_monitor": GradMonitorHparams,
+    "memory_monitor": MemoryMonitorHparams,
+    "mlperf": MLPerfCallbackHparams,
+    "early_stopper": EarlyStopperHparams,
+    "threshold_stopper": ThresholdStopperHparams,
+}

composer/callbacks/checkpoint_saver.py

@@ -245,7 +245,7 @@ class CheckpointSaver(Callback):
             If ``False`` (the default), then the ``folder`` must not exist or must not contain checkpoints which may conflict
             with the current run. Default: ``False``.
 
-        save_interval (:class:`.Time` | str | int | (:class:`.State`, :class:`.Event`) -> bool): A :class:`.Time`, time-string, integer (in epochs),
+        save_interval (Time | str | int | (State, Event) -> bool): A :class:`.Time`, time-string, integer (in epochs),
             or a function that takes (state, event) and returns a boolean whether a checkpoint should be saved.
 
             If an integer, checkpoints will be saved every n epochs.

composer/callbacks/memory_monitor.py

@@ -3,6 +3,7 @@
 
 """Log memory usage during training."""
 import logging
+import warnings
 from typing import Dict, Union
 
 import torch.cuda
@@ -76,24 +77,24 @@ class MemoryMonitor(Callback):
         Memory usage monitoring is only supported for the GPU devices.
     """
 
-    def __init__(self):
-        super().__init__()
-        log.info(
-            "Memory monitor just profiles the current GPU assuming that the memory footprint across GPUs is balanced.")
-        if torch.cuda.device_count() == 0:
-            log.warning("Memory monitor only works on GPU devices.")
+    def fit_start(self, state: State, logger: Logger) -> None:
+        # TODO(ravi) move this check would be on Event.INIT after #1084 is merged
+        # Not relying on `torch.cuda.is_available()` since the model could be on CPU.
+        model_device = next(state.model.parameters()).device
+
+        if model_device.type != 'cuda':
+            warnings.warn(f"The memory monitor only works on CUDA devices, but the model is on {model_device.type}.")
 
     def after_train_batch(self, state: State, logger: Logger):
         memory_report = {}
 
-        n_devices = torch.cuda.device_count()
-        if n_devices == 0:
+        model_device = next(state.model.parameters()).device
+        if model_device.type != 'cuda':
             return
 
         memory_report = _get_memory_report()
 
-        for mem_stat, val in memory_report.items():
-            logger.data_batch({'memory/{}'.format(mem_stat): val})
+        logger.data_batch({f'memory/{mem_stat}': val for (mem_stat, val) in memory_report.items()})
 
 
 _MEMORY_STATS = {
@@ -108,15 +109,8 @@ def after_train_batch(self, state: State, logger: Logger):
 
 
 def _get_memory_report() -> Dict[str, Union[int, float]]:
-    if not torch.cuda.is_available():
-        log.debug("Cuda is not available. The memory report will be empty.")
-        return {}
     memory_stats = torch.cuda.memory_stats()
 
-    if len(memory_stats) == 0:
-        log.debug("No GPU memory was used, returning empty.")
-        return {}
-
     # simplify the memory_stats
     memory_report = {
         name: memory_stats[torch_name] for (torch_name, name) in _MEMORY_STATS.items() if torch_name in memory_stats

composer/callbacks/speed_monitor.py

@@ -18,10 +18,10 @@
 class SpeedMonitor(Callback):
     """Logs the training throughput.
 
-    The training throughput in terms of number of samples per second is logged on the
-    :attr:`~composer.core.event.Event.BATCH_END` event if we have reached the ``window_size`` threshold.  Per epoch
-    average throughput and wall clock train, validation, and total time is also logged on the
-    :attr:`~composer.core.event.Event.EPOCH_END` event.
+    The training throughput in terms of number of samples per second is logged on
+    the :attr:`~composer.core.event.Event.BATCH_END` event if we have reached the ``window_size`` threshold. Per epoch
+    average throughput and wall clock train, validation, and total time is also logged on
+    the :attr:`~composer.core.event.Event.EPOCH_END` event.
 
     Example
 
@@ -49,11 +49,11 @@ class SpeedMonitor(Callback):
     | Key                   | Logged data                                                 |
     +=======================+=============================================================+
     |                       | Rolling average (over ``window_size`` most recent           |
-    | ``throughput/step``   | batches) of the number of samples processed per second      |
+    | ``samples/step``      | batches) of the number of samples processed per second      |
     |                       |                                                             |
     +-----------------------+-------------------------------------------------------------+
     |                       | Number of samples processed per second (averaged over       |
-    | ``throughput/epoch``  | an entire epoch)                                            |
+    | ``samples/epoch``     | an entire epoch)                                            |
     +-----------------------+-------------------------------------------------------------+
     |``wall_clock/train``   | Total elapsed training time                                 |
     +-----------------------+-------------------------------------------------------------+
@@ -63,8 +63,8 @@ class SpeedMonitor(Callback):
     +-----------------------+-------------------------------------------------------------+
 
     Args:
-        window_size (int, optional):
-            Number of batches to use for a rolling average of throughput. Default to 100.
+        window_size (int, optional): Number of batches to use for a rolling average of throughput.
+            Default to 100.
     """
 
     def __init__(self, window_size: int = 100):

composer/core/callback.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from composer.core.serializable import Serializable
 
@@ -88,6 +88,11 @@ class Callback(Serializable, abc.ABC):
             trainer.engine.close()
     """
 
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        # Stub signature for pyright
+        del args, kwargs  # unused
+        pass
+
     def run_event(self, event: Event, state: State, logger: Logger) -> None:
         """This method is called by the engine on each event.
 

composer/loggers/file_logger.py

@@ -113,6 +113,7 @@ class FileLogger(LoggerDestination):
             :attr:`~.LogLevel.EPOCH`, then the logfile will be flushed every n epochs.  If
             the ``log_level`` is :attr:`~.LogLevel.BATCH`, then the logfile will be
             flushed every n batches. Default: ``100``.
+        overwrite (bool, optional): Whether to overwrite an existing logfile. (default: ``False``)
     """
 
     def __init__(

composer/loggers/logger_hparams.py

@@ -136,18 +136,18 @@ def initialize_object(self) -> WandBLogger:
             config_dict = self._flatten_dict(config_dict, prefix=[])
 
         init_params = {
-            "project": self.project,
-            "name": self.name,
-            "group": self.group,
-            "entity": self.entity,
-            "tags": tags,
             "config": config_dict,
         }
         init_params.update(self.extra_init_params)
         return WandBLogger(
+            project=self.project,
+            group=self.group,
+            name=self.name,
+            entity=self.entity,
+            tags=tags,
             log_artifacts=self.log_artifacts,
             rank_zero_only=self.rank_zero_only,
-            init_params=init_params,
+            init_kwargs=init_params,
         )
 
     @classmethod

composer/loggers/progress_bar_logger.py

@@ -129,7 +129,8 @@ def should_log(state: State, log_level: LogLevel, console_log_level: LogLevel =
             elif stream.lower() == "stderr":
                 stream = sys.stderr
             else:
-                raise ValueError("Invalid stream option: Should be 'stdout', 'stderr', or a TextIO-like object.")
+                raise ValueError(
+                    f"Invalid stream option: Should be 'stdout', 'stderr', or a TextIO-like object; got {stream}")
         self.stream = stream
 
     def log_data(self, state: State, log_level: LogLevel, data: Dict[str, Any]) -> None:

composer/loggers/wandb_logger.py

@@ -13,7 +13,7 @@
 import sys
 import tempfile
 import warnings
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 from composer.core.state import State
 from composer.loggers.logger import Logger, LogLevel
@@ -28,6 +28,12 @@ class WandBLogger(LoggerDestination):
     """Log to Weights and Biases (https://wandb.ai/)
 
     Args:
+        project (str, optional): WandB project name.
+        group (str, optional): WandB group name.
+        name (str, optional): WandB run name.
+            If not specified, the :attr:`.Logger.run_name` will be used.
+        entity (str, optional): WandB entity name.
+        tags (List[str], optional): WandB tags.
         log_artifacts (bool, optional): Whether to log
             `artifacts <https://docs.wandb.ai/ref/python/artifact>`_ (Default: ``False``).
         rank_zero_only (bool, optional): Whether to log only on the rank-zero process.
@@ -36,15 +42,22 @@ class WandBLogger(LoggerDestination):
             stored, which may discard pertinent information. For example, when using
             Deepspeed ZeRO, it would be impossible to restore from checkpoints without
             artifacts from all ranks (default: ``False``).
-        init_params (Dict[str, Any], optional): Parameters to pass into
+        init_kwargs (Dict[str, Any], optional): Any additional init kwargs
             ``wandb.init`` (see
             `WandB documentation <https://docs.wandb.ai/ref/python/init>`_).
     """
 
-    def __init__(self,
-                 log_artifacts: bool = False,
-                 rank_zero_only: bool = True,
-                 init_params: Optional[Dict[str, Any]] = None) -> None:
+    def __init__(
+        self,
+        project: Optional[str] = None,
+        group: Optional[str] = None,
+        name: Optional[str] = None,
+        entity: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        log_artifacts: bool = False,
+        rank_zero_only: bool = True,
+        init_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
         try:
             import wandb
         except ImportError as e:
@@ -60,11 +73,27 @@ def __init__(self,
                  "restore from checkpoints."))
         self._enabled = (not rank_zero_only) or dist.get_global_rank() == 0
 
+        if init_kwargs is None:
+            init_kwargs = {}
+
+        if project is not None:
+            init_kwargs['project'] = project
+
+        if group is not None:
+            init_kwargs['group'] = group
+
+        if name is not None:
+            init_kwargs['name'] = name
+
+        if entity is not None:
+            init_kwargs['entity'] = entity
+
+        if tags is not None:
+            init_kwargs['tags'] = tags
+
         self._rank_zero_only = rank_zero_only
         self._log_artifacts = log_artifacts
-        if init_params is None:
-            init_params = {}
-        self._init_params = init_params
+        self._init_kwargs = init_kwargs
         self._is_in_atexit = False
 
     def _set_is_in_atexit(self):
@@ -98,17 +127,16 @@ def init(self, state: State, logger: Logger) -> None:
         del state  # unused
 
         # Use the logger run name if the name is not set.
-        if "name" not in self._init_params or self._init_params["name"] is None:
-            self._init_params["name"] = logger.run_name
+        if "name" not in self._init_kwargs or self._init_kwargs["name"] is None:
+            self._init_kwargs["name"] = logger.run_name
 
         # Adjust name and group based on `rank_zero_only`.
         if not self._rank_zero_only:
-            name = self._init_params["name"]
-            group = self._init_params.get("group", None)
-            self._init_params["name"] = f"{name} [RANK_{dist.get_global_rank()}]"
-            self._init_params["group"] = group if group else name
+            name = self._init_kwargs["name"]
+            self._init_kwargs["name"] += f"-rank{dist.get_global_rank()}"
+            self._init_kwargs["group"] = self._init_kwargs["group"] if "group" in self._init_kwargs else name
         if self._enabled:
-            wandb.init(**self._init_params)
+            wandb.init(**self._init_kwargs)
             atexit.register(self._set_is_in_atexit)
 
     def log_file_artifact(self, state: State, log_level: LogLevel, artifact_name: str, file_path: pathlib.Path, *,

composer/trainer/trainer.py

@@ -55,7 +55,7 @@
 def _raise_missing_argument_exception(arg_name: str):
     raise ValueError((f"{arg_name} is a required argument and must be specified when constructing the "
                       f"{Trainer.__name__} or when calling {Trainer.__name__}.{Trainer.fit.__name__}(). "
-                      f"To fix, please specify `arg_name` via {Trainer.__name__}({arg_name}=...) or "
+                      f"To fix, please specify `{arg_name}` via {Trainer.__name__}({arg_name}=...) or "
                       f"{Trainer.__name__}.{Trainer.fit.__name__}({arg_name}=...)."))