diff --git a/llmfoundry/utils/builders.py b/llmfoundry/utils/builders.py index b2cd535fe9..d35b089d50 100644 --- a/llmfoundry/utils/builders.py +++ b/llmfoundry/utils/builders.py @@ -12,7 +12,7 @@ from composer.core import Evaluator from composer.datasets.in_context_learning_evaluation import \ get_icl_task_dataloader -from composer.loggers import TensorboardLogger, WandBLogger +from composer.loggers import MLFlowLogger, TensorboardLogger, WandBLogger from composer.optim import DecoupledAdamW from composer.optim.scheduler import (ConstantWithWarmupScheduler, CosineAnnealingWithWarmupScheduler, @@ -67,6 +67,8 @@ def build_logger(name: str, kwargs: Dict[str, Any]): return WandBLogger(**kwargs) elif name == 'tensorboard': return TensorboardLogger(**kwargs) + elif name == 'mlflow': + return MLFlowLogger(**kwargs) else: raise ValueError(f'Not sure how to build logger: {name}') diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py index 6a4f1a8a4a..6c12775bfc 100644 --- a/llmfoundry/utils/config_utils.py +++ b/llmfoundry/utils/config_utils.py @@ -98,3 +98,11 @@ def log_config(cfg: DictConfig): raise e if wandb.run: wandb.config.update(om.to_container(cfg, resolve=True)) + + if 'mlflow' in cfg.get('loggers', {}): + try: + import mlflow + except ImportError as e: + raise e + if mlflow.active_run(): + mlflow.log_params(params=om.to_container(cfg, resolve=True)) diff --git a/setup.py b/setup.py index b2f017e8a7..a8921ba3ca 100644 --- a/setup.py +++ b/setup.py @@ -47,7 +47,7 @@ ] install_requires = [ - 'mosaicml[libcloud,nlp,wandb]>=0.15.0,<0.16', + 'mosaicml[libcloud,nlp,wandb,mlflow]>=0.15.0,<0.16', 'accelerate>=0.20,<0.21', # for HF inference `device_map` 'mosaicml-streaming>=0.5.1,<0.6', 'torch>=1.13.1,<=2.0.1',