log details to metadata for run analytics

llmfoundry/utils/mosaicmllogger_utils.py

@@ -0,0 +1,166 @@
+# Copyright 2024 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+import json
+from typing import Any, Dict, List, Optional, Union
+
+from composer.loggers import MosaicMLLogger
+from composer.loggers.logger_destination import LoggerDestination
+from composer.utils.file_helpers import parse_uri
+from omegaconf import DictConfig, ListConfig
+
+
+def find_mosaicml_logger(
+        loggers: List[LoggerDestination]) -> Union[MosaicMLLogger, None]:
+    return next(
+        (logger for logger in loggers if isinstance(logger, MosaicMLLogger)),
+        None)
+
+
+def log_eval_analytics(mosaicml_logger: MosaicMLLogger,
+                       model_configs: ListConfig, icl_tasks: Union[str,
+                                                                   ListConfig],
+                       eval_gauntlet_config: Optional[Union[str, DictConfig]]):
+    """Logs analytics for runs using the `eval.py` script."""
+    metrics: Dict[str, Any] = {
+        'llmfoundry/script': 'eval',
+    }
+
+    if eval_gauntlet_config is not None:
+        metrics['llmfoundry/gauntlet_configured'] = True
+    else:
+        metrics['llmfoundry/gauntlet_configured'] = False
+
+    if isinstance(icl_tasks, str):
+        metrics['llmfoundry/icl_configured'] = True
+    elif len(icl_tasks) > 0:
+        metrics['llmfoundry/icl_configured'] = True
+    else:
+        metrics['llmfoundry/icl_configured'] = False
+
+    metrics['llmfoundry/model_configs'] = []
+    for model_config in model_configs:
+        model_config_data = {}
+        if model_config.get('vocab_size', None) is not None:
+            model_config_data['vocab_size'] = model_config.get('vocab_size')
+        if model_config.get('d_model', None) is not None:
+            model_config_data['d_model'] = model_config.get('d_model')
+        if model_config.get('n_heads', None) is not None:
+            model_config_data['n_heads'] = model_config.get('n_heads')
+
+        load_path = model_config.get('load_path', None)
+        if load_path is not None:
+            backend, _, _ = parse_uri(load_path)
+            model_config_data[
+                'cloud_provider_data'] = backend if backend else 'local'
+
+        metrics['llmfoundry/model_configs'].append(
+            json.dumps(model_config_data, sort_keys=True))
+    """
+    TODO: right now, we need to get the following subtypes for eval:
+        -- Batch Generation
+        -- Checkpoint Conversion
+    """
+    mosaicml_logger.log_metrics(metrics)
+    mosaicml_logger._flush_metadata(force_flush=True)
+
+
+def log_train_analytics(mosaicml_logger: MosaicMLLogger,
+                        model_config: DictConfig,
+                        train_loader_config: DictConfig,
+                        eval_loader_config: Union[DictConfig, ListConfig, None],
+                        callback_configs: Union[DictConfig,
+                                                None], tokenizer_name: str,
+                        load_path: Union[str, None], save_folder: Union[str,
+                                                                        None],
+                        icl_tasks_config: Optional[Union[ListConfig, str]],
+                        eval_gauntlet: Optional[Union[DictConfig, str]]):
+    """Logs analytics for runs using the `train.py` script.
+
+    TODO: clean up this code and make it more readable.
+    """
+    train_loader_dataset = train_loader_config.get('dataset', {})
+    metrics: Dict[str, Any] = {
+        'llmfoundry/tokenizer_name':
+            tokenizer_name,
+        'llmfoundry/script':
+            'train',
+        'llmfoundry/train_loader_name':
+            train_loader_config.get('name'),
+        'llmfoundry/train_loader_workers':
+            train_loader_dataset.get('num_workers'),
+    }
+
+    if callback_configs is not None:
+        metrics['llmfoundry/callbacks'] = [
+            name for name, _ in callback_configs.items()
+        ]
+
+    if eval_gauntlet is not None:
+        metrics['llmfoundry/gauntlet_configured'] = True
+    else:
+        metrics['llmfoundry/gauntlet_configured'] = False
+
+    if icl_tasks_config is not None:
+        if isinstance(icl_tasks_config, str):
+            metrics['llmfoundry/icl_configured'] = True
+        elif len(icl_tasks_config) > 0:
+            metrics['llmfoundry/icl_configured'] = True
+        else:
+            metrics['llmfoundry/icl_configured'] = False
+    else:
+        metrics['llmfoundry/icl_configured'] = False
+
+    if load_path is not None:
+        backend, _, _ = parse_uri(load_path)
+        metrics[
+            'llmfoundry/cloud_provider_data'] = backend if backend else 'local'
+    if save_folder is not None:
+        backend, _, _ = parse_uri(save_folder)
+        metrics[
+            'llmfoundry/cloud_provider_checkpoints'] = backend if backend else 'local'
+
+    if train_loader_dataset.get('hf_name', None) is not None:
+        metrics['llmfoundry/train_dataset_hf_name'] = train_loader_dataset.get(
+            'hf_name', None)
+    if train_loader_config.get('name') == 'finetuning':
+        metrics['llmfoundry/train_task_type'] = 'INSTRUCTION_FINETUNE'
+    elif train_loader_config.get('name') == 'text':
+        if load_path is not None or model_config.get('pretrained') == True:
+            metrics['llmfoundry/train_task_type'] = 'CONTINUED_PRETRAIN'
+        else:
+            metrics['llmfoundry/train_task_type'] = 'PRETRAIN'
+
+    if eval_loader_config is not None:
+        metrics['llmfoundry/eval_loaders'] = []
+
+        if isinstance(eval_loader_config, ListConfig):
+            eval_loader_configs: ListConfig = eval_loader_config
+        else:
+            eval_loader_configs = ListConfig([eval_loader_config])
+
+        for loader_config in eval_loader_configs:
+            eval_loader_info = {}
+            eval_loader_dataset = loader_config.get('dataset', {})
+            eval_loader_info['name'] = loader_config.get('name')
+            eval_loader_info['num_workers'] = eval_loader_dataset.get(
+                'num_workers', None)
+            if eval_loader_dataset.get('hf_name', None) is not None:
+                eval_loader_info['dataset_hf_name'] = eval_loader_dataset.get(
+                    'hf_name')
+
+            # Log as a key-sorted JSON string, so that we can easily parse it in Spark / SQL
+            metrics['llmfoundry/eval_loaders'].append(
+                json.dumps(eval_loader_info, sort_keys=True))
+
+    if model_config['name'] == 'hf_casual_lm':
+        metrics['llmfoundry/model_name'] = model_config.get(
+            'pretrained_model_name_or_path')
+    if model_config.get('vocab_size', None) is not None:
+        metrics['llmfoundry/vocab_size'] = model_config.get('vocab_size'),
+    if model_config.get('d_model', None) is not None:
+        metrics['llmfoundry/d_model'] = model_config.get('d_model')
+    if model_config.get('n_heads', None) is not None:
+        metrics['llmfoundry/n_heads'] = model_config.get('n_heads')
+
+    mosaicml_logger.log_metrics(metrics)
+    mosaicml_logger._flush_metadata(force_flush=True)

scripts/eval/eval.py

@@ -11,7 +11,6 @@
 
 import pandas as pd
 import torch
-from composer.loggers import MosaicMLLogger
 from composer.loggers.logger_destination import LoggerDestination
 from composer.models.base import ComposerModel
 from composer.trainer import Trainer
@@ -21,6 +20,9 @@
 from rich.traceback import install
 from transformers import PreTrainedTokenizerBase
 
+from llmfoundry.utils.mosaicmllogger_utils import (find_mosaicml_logger,
+                                                   log_eval_analytics)
+
 install()
 from llmfoundry.models.model_registry import COMPOSER_MODEL_REGISTRY
 from llmfoundry.registry import import_file
@@ -114,9 +116,7 @@ def evaluate_model(
         loggers_cfg.update(metadata, merge=True)
 
         # Find the MosaicMLLogger
-        mosaicml_logger = next((
-            logger for logger in loggers if isinstance(logger, MosaicMLLogger)),
-                               None)
+        mosaicml_logger = find_mosaicml_logger(loggers)
 
         if mosaicml_logger is not None:
             mosaicml_logger.log_metrics(metadata)
@@ -153,7 +153,6 @@ def evaluate_model(
     assert composer_model is not None
 
     log.info(f'Building trainer for {model_cfg.model_name}...')
-
     trainer = Trainer(
         run_name=run_name,
         seed=seed,
@@ -298,6 +297,15 @@ def main(cfg: DictConfig) -> Tuple[List[Trainer], pd.DataFrame]:
     models_df = None
     composite_scores = None
     trainers = []
+
+    mosaicml_logger = find_mosaicml_logger([
+        build_logger(name, logger_cfg)
+        for name, logger_cfg in loggers_cfg.items()
+    ])
+    if mosaicml_logger is not None:
+        log_eval_analytics(mosaicml_logger, model_configs, icl_tasks,
+                           eval_gauntlet_config)
+
     for model_cfg in model_configs:
         (trainer, logger_keys, eval_gauntlet_callback,
          eval_gauntlet_df) = evaluate_model(

scripts/train/train.py

@@ -23,6 +23,9 @@
 from omegaconf import OmegaConf as om
 from rich.traceback import install
 
+from llmfoundry.utils.mosaicmllogger_utils import (find_mosaicml_logger,
+                                                   log_train_analytics)
+
 install()
 
 from transformers import PreTrainedTokenizerBase
@@ -449,9 +452,7 @@ def main(cfg: DictConfig) -> Trainer:
         for name, logger_cfg in logger_configs.items()
     ] if logger_configs else []
 
-    mosaicml_logger = next(
-        (logger for logger in loggers if isinstance(logger, MosaicMLLogger)),
-        None)
+    mosaicml_logger = find_mosaicml_logger(loggers)
     if mosaicml_logger is None:
         if os.environ.get(MOSAICML_PLATFORM_ENV_VAR, 'false').lower(
         ) == 'true' and os.environ.get(MOSAICML_ACCESS_TOKEN_ENV_VAR):
@@ -543,6 +544,12 @@ def main(cfg: DictConfig) -> Trainer:
         if eval_gauntlet_callback is not None:
             callbacks.append(eval_gauntlet_callback)
 
+    if mosaicml_logger is not None:
+        log_train_analytics(mosaicml_logger, model_config, train_loader_config,
+                            eval_loader_config, callback_configs,
+                            tokenizer_name, load_path, save_folder,
+                            icl_tasks_config, eval_gauntlet_config)
+
     # Build Model
     log.info('Initializing model...')
     with init_context: