Registry based config - Part 1

README.md

@@ -259,6 +259,10 @@ export HUGGING_FACE_HUB_TOKEN=your-auth-token
 
 and uncomment the line containing `--hf_repo_for_upload ...` in the above call to `inference/convert_composer_to_hf.py`.
 
+# UNDER CONSTRUCTION: Registry
+
+We are adopting an extensible registry for LLM Foundry to allow various extensions of the library without forking it. See [./REGISTRY.md] for more information as it develops.
+
 # Learn more about LLM Foundry!
 
 Check out [TUTORIAL.md](https://github.com/mosaicml/llm-foundry/blob/main/TUTORIAL.md) to keep learning about working with LLM Foundry. The tutorial highlights example workflows, points you to other resources throughout the repo, and answers frequently asked questions!

REGISTRY.md

@@ -0,0 +1,79 @@
+# LLM Foundry Registry
+
+Some components of LLM Foundry are registrable. This means that you can register options for these components, and then use them in your yaml config, without forking the library.
+
+## How to register
+
+There are a few ways to register a new component:
+
+### Python entrypoints
+
+You can specify registered components via a Python entrypoint if you are building your own package with registered components.
+
+For example, the following would register the `WandBLogger` class, under the key `wandb`, in the `llm_foundry.loggers` registry:
+
+```yaml
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "foundry_registry"
+version = "0.1.0"
+dependencies = [
+    "mosaicml",
+    "llm-foundry",
+]
+
+[project.entry-points."llm_foundry.loggers"]
+my_logger = "foundry_registry.loggers:MyLogger"
+```
+
+### Direct call to register
+
+You can also register a component directly in your code:
+
+```python
+from composer.loggers import LoggerDestination
+from llmfoundry.registry import loggers
+
+class MyLogger(LoggerDestination):
+    pass
+
+loggers.register("my_logger", func=MyLogger)
+```
+
+### Decorators
+
+You can also use decorators to register components directly from your code:
+
+```python
+from composer.loggers import LoggerDestination
+from llmfoundry.registry import loggers
+
+@loggers.register("my_logger")
+class MyLogger(LoggerDestination):
+    pass
+```
+
+For both the direct call and decorator approaches, if using the LLM Foundry train/eval scripts, you will need to provide the `code_paths` argument, which is a list of files need to execute in order to register your components. For example, you may have a file called `foundry_imports.py` that contains the following:
+
+```python
+from foundry_registry.loggers import MyLogger
+from llmfoundry.registry import loggers
+
+loggers.register("my_logger", func=MyLogger)
+```
+
+You would then provide `code_paths` to the train/eval scripts in your yaml config:
+
+```yaml
+...
+code_paths:
+  - foundry_imports.py
+...
+```
+
+
+## Discovering registrable components
+Coming soon

llmfoundry/__init__.py

@@ -24,7 +24,7 @@
 # gated import otherwise.
 import transformers
 
-from llmfoundry import optim, utils
+from llmfoundry import algorithms, callbacks, loggers, optim, registry, utils
 from llmfoundry.data import (ConcatTokensDataset, MixtureOfDenoisersCollator,
                              NoConcatDataset, Seq2SeqFinetuningCollator,
                              build_finetuning_dataloader,
@@ -74,7 +74,11 @@
     'build_alibi_bias',
     'optim',
     'utils',
+    'loggers',
+    'algorithms',
+    'callbacks',
     'TiktokenTokenizerWrapper',
+    'registry',
 ]
 
 __version__ = '0.5.0'

llmfoundry/algorithms/__init__.py

@@ -0,0 +1,12 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+from composer.algorithms import (Alibi, GatedLinearUnits, GradientClipping,
+                                 LowPrecisionLayerNorm)
+
+from llmfoundry.registry import algorithms
+
+algorithms.register('gradient_clipping', func=GradientClipping)
+algorithms.register('alibi', func=Alibi)
+algorithms.register('gated_linear_units', func=GatedLinearUnits)
+algorithms.register('low_precision_layernorm', func=LowPrecisionLayerNorm)

llmfoundry/callbacks/__init__.py

@@ -1,6 +1,11 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
+from composer.callbacks import (EarlyStopper, Generate, LRMonitor,
+                                MemoryMonitor, MemorySnapshot, OOMObserver,
+                                OptimizerMonitor, RuntimeEstimator,
+                                SpeedMonitor)
+
 from llmfoundry.callbacks.async_eval_callback import AsyncEval
 from llmfoundry.callbacks.curriculum_learning_callback import CurriculumLearning
 from llmfoundry.callbacks.eval_gauntlet_callback import EvalGauntlet
@@ -11,6 +16,26 @@
 from llmfoundry.callbacks.resumption_callbacks import (GlobalLRScaling,
                                                        LayerFreezing)
 from llmfoundry.callbacks.scheduled_gc_callback import ScheduledGarbageCollector
+from llmfoundry.registry import callbacks, callbacks_with_config
+
+callbacks.register('lr_monitor', func=LRMonitor)
+callbacks.register('memory_monitor', func=MemoryMonitor)
+callbacks.register('memory_snapshot', func=MemorySnapshot)
+callbacks.register('speed_monitor', func=SpeedMonitor)
+callbacks.register('runtime_estimator', func=RuntimeEstimator)
+callbacks.register('optimizer_monitor', func=OptimizerMonitor)
+callbacks.register('generate_callback', func=Generate)
+callbacks.register('early_stopper', func=EarlyStopper)
+callbacks.register('fdiff_metrics', func=FDiffMetrics)
+callbacks.register('hf_checkpointer', func=HuggingFaceCheckpointer)
+callbacks.register('global_lr_scaling', func=GlobalLRScaling)
+callbacks.register('layer_freezing', func=LayerFreezing)
+callbacks.register('mono_checkpoint_saver', func=MonolithicCheckpointSaver)
+callbacks.register('scheduled_gc', func=ScheduledGarbageCollector)
+callbacks.register('oom_observer', func=OOMObserver)
+
+callbacks_with_config.register('async_eval', func=AsyncEval)
+callbacks_with_config.register('curriculum_learning', func=CurriculumLearning)
 
 __all__ = [
     'FDiffMetrics',

llmfoundry/callbacks/async_eval_callback.py

@@ -14,14 +14,15 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 from composer.callbacks import CheckpointSaver
-from composer.core import Callback, Event, State, Time, Timestamp, TimeUnit
+from composer.core import Event, State, Time, Timestamp, TimeUnit
 from composer.loggers import Logger
 from composer.loggers.mosaicml_logger import (MOSAICML_PLATFORM_ENV_VAR,
                                               RUN_NAME_ENV_VAR)
 from composer.utils import dist
 from composer.utils.file_helpers import list_remote_objects
 from composer.utils.misc import create_interval_scheduler
 
+from llmfoundry.interfaces import CallbackWithConfig
 from mcli import Run, RunConfig, create_run, get_run
 
 log = logging.getLogger(__name__)
@@ -177,7 +178,7 @@ def validate_eval_run_config(
 CHECKS_PER_INTERVAL = 4
 
 
-class AsyncEval(Callback):
+class AsyncEval(CallbackWithConfig):
     """Run the eval loop asynchronously as part of a MosaicML platform run.
 
     This callback is currently experimental. The API may change in the future.

llmfoundry/callbacks/curriculum_learning_callback.py

@@ -10,15 +10,17 @@
 import logging
 from typing import Any, Dict
 
-from composer.core import Callback, State
+from composer.core import State
 from composer.loggers import Logger
 from streaming import StreamingDataset
 from torch.utils.data import DataLoader
 
+from llmfoundry.interfaces import CallbackWithConfig
+
 log = logging.getLogger(__name__)
 
 
-class CurriculumLearning(Callback):
+class CurriculumLearning(CallbackWithConfig):
     """Starts an epoch with a different dataset when resuming from a checkpoint.
 
     This callback is currently experimental. The API may change without warning in the future.
@@ -29,13 +31,13 @@ class CurriculumLearning(Callback):
             being used.
     """
 
-    def __init__(self, dataset_index: int, current_dataset_config: Dict):
+    def __init__(self, dataset_index: int, train_config: Dict):
         self.dataset_index = dataset_index
         self.saved_dataset_index = 0
         self.all_dataset_configs = []
         self.current_dataset_state = {}
         # The current dataset config is resolved and passed in train.py
-        self.current_dataset_config = current_dataset_config
+        self.current_dataset_config = train_config['dataloader']
 
     def before_load(self, state: State, logger: Logger):
         del logger

llmfoundry/interfaces/__init__.py

@@ -0,0 +1,8 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+from llmfoundry.interfaces.callback_with_config import CallbackWithConfig
+
+__all__ = [
+    'CallbackWithConfig',
+]

llmfoundry/interfaces/callback_with_config.py

@@ -0,0 +1,17 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+from typing import Any
+
+from composer.core import Callback
+
+__all__ = ['CallbackWithConfig']
+
+
+class CallbackWithConfig(Callback, abc.ABC):
+
+    def __init__(self, config: dict[str, Any], *args: Any,
+                 **kwargs: Any) -> None:
+        del config, args, kwargs
+        pass

llmfoundry/loggers/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+from composer.loggers import (InMemoryLogger, MLFlowLogger, TensorboardLogger,
+                              WandBLogger)
+
+from llmfoundry.registry import loggers
+
+loggers.register('wandb', func=WandBLogger)
+loggers.register('tensorboard', func=TensorboardLogger)
+loggers.register('inmemory', func=InMemoryLogger)
+loggers.register('in_memory_logger',
+                 func=InMemoryLogger)  # for backwards compatibility
+loggers.register('mlflow', func=MLFlowLogger)

llmfoundry/optim/__init__.py

@@ -1,8 +1,26 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
+from composer.optim import (ConstantWithWarmupScheduler,
+                            CosineAnnealingWithWarmupScheduler, DecoupledAdamW,
+                            LinearWithWarmupScheduler)
+
 from llmfoundry.optim.adaptive_lion import DecoupledAdaLRLion, DecoupledClipLion
 from llmfoundry.optim.lion import DecoupledLionW
+from llmfoundry.optim.scheduler import InverseSquareRootWithWarmupScheduler
+from llmfoundry.registry import optimizers, schedulers
+
+optimizers.register('adalr_lion', func=DecoupledAdaLRLion)
+optimizers.register('clip_lion', func=DecoupledClipLion)
+optimizers.register('decoupled_lionw', func=DecoupledLionW)
+optimizers.register('decoupled_adamw', func=DecoupledAdamW)
+
+schedulers.register('constant_with_warmup', func=ConstantWithWarmupScheduler)
+schedulers.register('cosine_with_warmup',
+                    func=CosineAnnealingWithWarmupScheduler)
+schedulers.register('linear_decay_with_warmup', func=LinearWithWarmupScheduler)
+schedulers.register('inv_sqrt_with_warmup',
+                    func=InverseSquareRootWithWarmupScheduler)
 
 __all__ = [
     'DecoupledLionW',

llmfoundry/registry.py

@@ -0,0 +1,72 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+from typing import Type
+
+from composer.core import Algorithm, Callback
+from composer.loggers import LoggerDestination
+from composer.optim import ComposerScheduler
+from torch.optim import Optimizer
+
+from llmfoundry.interfaces import CallbackWithConfig
+from llmfoundry.utils.registry_utils import create
+
+_loggers_description = """The loggers registry is used to register classes that implement the LoggerDestination interface.
+These classes are used to log data from the training loop, and will be passed to the loggers arg of the Trainer. The loggers
+will be constructed by directly passing along the specified kwargs to the constructor."""
+loggers = create('llmfoundry',
+                 'loggers',
+                 generic_type=Type[LoggerDestination],
+                 entry_points=True,
+                 description=_loggers_description)
+
+_callbacks_description = """The callbacks registry is used to register classes that implement the Callback interface.
+These classes are used to interact with the Composer event system, and will be passed to the callbacks arg of the Trainer.
+The callbacks will be constructed by directly passing along the specified kwargs to the constructor."""
+callbacks = create('llmfoundry',
+                   'callbacks',
+                   generic_type=Type[Callback],
+                   entry_points=True,
+                   description=_callbacks_description)
+
+_callbacks_with_config_description = """The callbacks_with_config registry is used to register classes that implement the CallbackWithConfig interface.
+These are the same as the callbacks registry, except that they additionally take the full training config as an argument to their constructor."""
+callbacks_with_config = create('llm_foundry.callbacks_with_config',
+                               generic_type=Type[CallbackWithConfig],
+                               entry_points=True,
+                               description=_callbacks_with_config_description)
+
+_optimizers_description = """The optimizers registry is used to register classes that implement the Optimizer interface.
+The optimizer will be passed to the optimizers arg of the Trainer. The optimizer will be constructed by directly passing along the
+specified kwargs to the constructor, along with the model parameters."""
+optimizers = create('llmfoundry',
+                    'optimizers',
+                    generic_type=Type[Optimizer],
+                    entry_points=True,
+                    description=_optimizers_description)
+
+_algorithms_description = """The algorithms registry is used to register classes that implement the Algorithm interface.
+The algorithm will be passed to the algorithms arg of the Trainer. The algorithm will be constructed by directly passing along the
+specified kwargs to the constructor."""
+algorithms = create('llmfoundry',
+                    'algorithms',
+                    generic_type=Type[Algorithm],
+                    entry_points=True,
+                    description=_algorithms_description)
+
+_schedulers_description = """The schedulers registry is used to register classes that implement the ComposerScheduler interface.
+The scheduler will be passed to the schedulers arg of the Trainer. The scheduler will be constructed by directly passing along the
+specified kwargs to the constructor."""
+schedulers = create('llmfoundry',
+                    'schedulers',
+                    generic_type=Type[ComposerScheduler],
+                    entry_points=True,
+                    description=_schedulers_description)
+
+__all__ = [
+    'loggers',
+    'callbacks',
+    'callbacks_with_config',
+    'optimizers',
+    'algorithms',
+    'schedulers',
+]