diff --git a/llmfoundry/loggers/__init__.py b/llmfoundry/loggers/__init__.py index cd3f3fdc62..c60d9be2cd 100644 --- a/llmfoundry/loggers/__init__.py +++ b/llmfoundry/loggers/__init__.py @@ -4,6 +4,7 @@ from composer.loggers import ( InMemoryLogger, MLFlowLogger, + MosaicMLLogger, TensorboardLogger, WandBLogger, ) @@ -18,3 +19,4 @@ func=InMemoryLogger, ) # for backwards compatibility loggers.register('mlflow', func=MLFlowLogger) +loggers.register('mosaicml', func=MosaicMLLogger) diff --git a/tests/loggers/test_mosaic_ml_logger.py b/tests/loggers/test_mosaic_ml_logger.py new file mode 100644 index 0000000000..e9c003321b --- /dev/null +++ b/tests/loggers/test_mosaic_ml_logger.py @@ -0,0 +1,16 @@ +# Copyright 2024 MosaicML LLM Foundry authors +# SPDX-License-Identifier: Apache-2.0 + +from composer.loggers import MosaicMLLogger + +from llmfoundry.utils.builders import build_logger + + +def test_mosaic_ml_logger_constructs(): + mosaic_ml_logger = build_logger( + 'mosaicml', + kwargs={'ignore_exceptions': True}, + ) + + assert isinstance(mosaic_ml_logger, MosaicMLLogger) + assert mosaic_ml_logger.ignore_exceptions == True