Registry based config - Part 1 (#975)

README.md

@@ -255,6 +255,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`.
 
+# :construction: 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,84 @@
+# :construction: 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:
+
+<!--pytest.mark.skip-->
+```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:
+
+<!--pytest.mark.skip-->
+```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:
+
+<!--pytest.mark.skip-->
+```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:
+
+<!--pytest.mark.skip-->
+```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:
+
+<!--pytest.mark.skip-->
+```yaml
+...
+code_paths:
+  - foundry_imports.py
+...
+```
+
+
+## Discovering registrable components
+Coming soon

llmfoundry/__init__.py

@@ -19,7 +19,7 @@
 
 hf_dynamic_modules_logger.addFilter(new_files_warning_filter)
 
-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,
@@ -65,7 +65,11 @@
     'build_alibi_bias',
     'optim',
     'utils',
+    'loggers',
+    'algorithms',
+    'callbacks',
     'TiktokenTokenizerWrapper',
+    'registry',
 ]
 
 __version__ = '0.6.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,18 +10,19 @@
 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
 from llmfoundry.utils.warnings import experimental
 
 log = logging.getLogger(__name__)
 
 
 @experimental('CurriculumLearning callback')
-class CurriculumLearning(Callback):
+class CurriculumLearning(CallbackWithConfig):
     """Starts an epoch with a different dataset when resuming from a checkpoint.
 
     Args:
@@ -30,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,21 @@
+# 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):
+    """A callback that takes a config dictionary as an argument, in addition to.
+
+    its other kwargs.
+    """
+
+    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',