Add CometLogger for metrics logging via comet.ml

README.md

@@ -30,6 +30,7 @@ torchtune focuses on integrating with popular tools and libraries from the ecosy
 - [PyTorch FSDP](https://pytorch.org/docs/stable/fsdp.html) for distributed training
 - [torchao](https:/pytorch-labs/ao) for lower precision dtypes and [post-training quantization](recipes/quantize.py) techniques
 - [Weights & Biases](https://wandb.ai/site) for [logging](https://pytorch.org/torchtune/main/deep_dives/wandb_logging.html) metrics and checkpoints, and tracking training progress
+- [CometML](https://www.comet.com/site/) as another option for [logging](https://pytorch.org/torchtune/stable/api_ref_utilities.html#metric-logging)
 - [ExecuTorch](https://pytorch.org/executorch-overview) for [on-device inference](https:/pytorch/executorch/tree/main/examples/models/llama2#optional-finetuning) using fine-tuned models
 - [bitsandbytes](https://huggingface.co/docs/bitsandbytes/main/en/index) for low memory optimizers for our [single-device recipes](recipes/configs/llama2/7B_full_low_memory.yaml)
 

docs/source/api_ref_utilities.rst

@@ -100,6 +100,7 @@ Various logging utilities.
  :toctree: generated/
  :nosignatures:
 
+ metric_logging.CometLogger
  metric_logging.WandBLogger
  metric_logging.TensorBoardLogger
  metric_logging.StdoutLogger

docs/source/deep_dives/comet_logging.rst

@@ -0,0 +1,53 @@
+.. _comet_logging:
+
+================
+Logging to Comet
+================
+
+This deep-dive will guide you through how to set up logging to Comet in torchtune.
+
+.. grid:: 1
+
+ .. grid-item-card:: :octicon:`mortar-board;1em;` What this deep-dive will cover
+
+ * How to get started with Comet
+ * How to use the :class:`~torchtune.utils.metric_logging.CometLogger`
+ * How to log configs, metrics, and model checkpoints to Comet
+
+Torchtune supports logging your training runs to `Comet <https://www.comet.com/site/?utm_source=torchtune&utm_medium=docs&utm_content=docs>`_.
+An example Comet workspace from a torchtune fine-tuning run can be seen in the screenshot below.
+
+.. image:: ../_static/img/comet_torchtune_project.png
+ :alt: torchtune workspace in Comet
+ :width: 100%
+ :align: center
+
+.. note::
+
+ You will need to install the :code:`comet_ml` package to use this feature.
+ You can install it via pip:
+
+ .. code-block:: bash
+
+ comet login
+
+Metric Logger
+-------------
+
+The only change you need to make is to add the metric logger to your config. Comet will log the metrics and model checkpoints for you.
+
+.. code-block:: yaml
+
+ # enable logging to the built-in CometLogger
+ metric_logger:
+ _component_: torchtune.utils.metric_logging.CometLogger
+ # the Comet project to log to
+ project: comet-examples-torchtune
+ experiment_name: my-experiment-name
+
+We automatically grab the config from the recipe you are running and log it to Comet. You can find it in the Comet Hyperparameters tab and the actual file in the :code:`Assets & Artifacts` tab.
+
+.. note::
+
+ Click on this sample `Comet project to see how it will looks like after fine-tuning <https://www.comet.com/examples/comet-example-torchtune-mistral/>`_.
+ The config used to train the models can be found `here <https://www.comet.com/examples/comet-example-torchtune-mistral/0aabcd062de548bbbd30912544aaa41a?experiment-tab=params>`_.

docs/source/index.rst

@@ -118,6 +118,7 @@ torchtune tutorials.
  :hidden:
 
  deep_dives/checkpointer
+ deep_dives/comet_logging
  deep_dives/configs
  deep_dives/recipe_deepdive
  deep_dives/wandb_logging

pyproject.toml

@@ -42,6 +42,7 @@ tune = "torchtune._cli.tune:main"
 [project.optional-dependencies]
 dev = [
  "bitsandbytes>=0.43.0",
+ "comet_ml>=3.44.2",
  "pre-commit",
  "pytest==7.4.0",
  "pytest-cov",

tests/torchtune/utils/test_metric_logging.py

@@ -17,6 +17,7 @@
 from tests.test_utils import assert_expected, captured_output
 
 from torchtune.utils.metric_logging import (
+ CometLogger,
  DiskLogger,
  StdoutLogger,
  TensorBoardLogger,
@@ -165,3 +166,36 @@ def test_save_config(self) -> None:
  expected_config_path = "torchtune_config.yaml"
  mock_save.assert_called_once_with(cfg, expected_config_path)
  mock_wandb_save.assert_called_once_with(expected_config_path)
+
+
+class TestCometLogger:
+ def test_log(self) -> None:
+ with patch("comet_ml.Experiment") as mock_experiment:
+ logger = CometLogger(project_name="test_project")
+ for i in range(5):
+ logger.log("test_log", float(i) ** 2, i)
+ logger.close()
+
+ assert mock_experiment.return_value.log_metric.call_count == 5
+ for i in range(5):
+ mock_experiment.return_value.log_metric.assert_any_call(
+ "test_log", float(i) ** 2, step=i
+ )
+
+ def test_log_dict(self) -> None:
+ with patch("comet_ml.Experiment") as mock_experiment:
+ logger = CometLogger(project_name="test_project")
+ metric_dict = {f"log_dict_{i}": float(i) ** 2 for i in range(5)}
+ logger.log_dict(metric_dict, 1)
+ logger.close()
+
+ mock_experiment.return_value.log_metrics.assert_called_with(
+ metric_dict, step=1
+ )
+
+ def test_log_config(self) -> None:
+ with patch("comet_ml.Experiment") as mock_experiment:
+ logger = CometLogger(project_name="test_project")
+ cfg = OmegaConf.create({"a": 1, "b": 2})
+ logger.log_config(cfg)
+ mock_experiment.return_value.log_parameters.assert_called_with(cfg)

torchtune/utils/metric_logging.py

@@ -8,7 +8,7 @@
 import time
 from pathlib import Path
 
-from typing import Mapping, Optional, Union
+from typing import List, Mapping, Optional, Union
 
 from numpy import ndarray
 from omegaconf import DictConfig, OmegaConf
@@ -317,3 +317,127 @@ def close(self) -> None:
  if self._writer:
  self._writer.close()
  self._writer = None
+
+
+class CometLogger(MetricLoggerInterface):
+ """Logger for use w/ Comet (https://www.comet.com/site/).
+ For more information about arguments expected by Comet, see
+ https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#for-the-experiment.
+
+ Args:
+ api_key (Optional[str]): Comet API key. It's recommended to configure the API Key from the environment.
+ workspace (Optional[str]): Comet workspace name. If not provided, uses the default workspace.
+ project (Optional[str]): Comet project name. Defaults to Uncategorized.
+ experiment_name (Optional[str]): The name for comet experiment to be used for logging.
+ experiment_key (Optional[str]): The key for comet experiment to be used for logging. Must be an alphanumeric
+ string whose length is between 32 and 50 characters.
+ mode (Optional[str]): Control how the Comet experiment is started. "get": Continue logging to an existing
+ experiment identified by the `experiment_key` value. "create": Always creates of a new experiment, useful
+ for HPO sweeps. "get_or_create" (default): Starts a fresh experiment if required, or persists logging to
+ an existing one.
+ online (Optional[bool]): If True, the data will be logged to Comet server, otherwise it will be stored locally
+ in offline experiment. Default is `True`.
+ experiment_name (Optional[str]): Name of the experiment. If not provided, Comet will auto-generate a name.
+ tags (Optional[List[str]]): Tags to associate with the experiment.
+ log_code (bool): Whether to log the source code. Defaults to True.
+ **kwargs: additional arguments to pass to Comet.start. See
+ https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment-Creation/#comet_ml.ExperimentConfig
+
+ Example:
+ >>> from torchtune.utils.metric_logging import CometLogger
+ >>> logger = CometLogger(project_name="my_project", workspace="my_workspace")
+ >>> logger.log("my_metric", 1.0, 1)
+ >>> logger.log_dict({"my_metric": 1.0}, 1)
+ >>> logger.close()
+
+ Raises:
+ ImportError: If `comet_ml` package is not installed.
+
+ Note:
+ This logger requires the comet_ml package to be installed.
+ You can install it with ``pip install comet_ml``.
+ You need to set up your Comet.ml API key before using this logger.
+ You can do this by setting the COMET_API_KEY environment variable
+ or by calling ``comet_ml.login()`` with your API key.
+ """
+
+ def __init__(
+ self,
+ api_key: Optional[str] = None,
+ workspace: Optional[str] = None,
+ project: Optional[str] = None,
+ experiment_key: Optional[str] = None,
+ mode: Optional[str] = None,
+ online: Optional[bool] = None,
+ experiment_name: Optional[str] = None,
+ tags: Optional[List[str]] = None,
+ log_code: bool = True,
+ **kwargs,
+ ):
+ try:
+ import comet_ml
+ except ImportError as e:
+ raise ImportError(
+ "``comet_ml`` package not found. Please install comet_ml using `pip install comet_ml` to use CometLogger."
+ "Alternatively, use the ``StdoutLogger``, which can be specified by setting metric_logger_type='stdout'."
+ ) from e
+
+ _, self.rank = get_world_size_and_rank()
+
+ # Declare it early so further methods don't crash in case of
+ # Experiment Creation failure due to mis-named configuration for
+ # example
+ self.experiment = None
+
+ if self.rank == 0:
+ self.experiment = comet_ml.start(
+ api_key=api_key,
+ workspace=workspace,
+ project=project,
+ experiment_key=experiment_key,
+ mode=mode,
+ online=online,
+ experiment_config=comet_ml.ExperimentConfig(
+ log_code=log_code, tags=tags, name=experiment_name, **kwargs
+ ),
+ )
+
+ def log(self, name: str, data: Scalar, step: int) -> None:
+ if self.experiment is not None:
+ self.experiment.log_metric(name, data, step=step)
+
+ def log_dict(self, payload: Mapping[str, Scalar], step: int) -> None:
+ if self.experiment is not None:
+ self.experiment.log_metrics(payload, step=step)
+
+ def log_config(self, config: DictConfig) -> None:
+ if self.experiment is not None:
+ resolved = OmegaConf.to_container(config, resolve=True)
+ self.experiment.log_parameters(resolved)
+
+ # Also try to save the config as a file
+ try:
+ self._log_config_as_file(config)
+ except Exception as e:
+ log.warning(f"Error saving Config to disk.\nError: \n{e}.")
+ return
+
+ def _log_config_as_file(self, config: DictConfig):
+ output_config_fname = Path(
+ os.path.join(
+ config.checkpointer.checkpoint_dir,
+ "torchtune_config.yaml",
+ )
+ )
+ OmegaConf.save(config, output_config_fname)
+
+ self.experiment.log_asset(
+ output_config_fname, file_name="torchtune_config.yaml"
+ )
+
+ def close(self) -> None:
+ if self.experiment is not None:
+ self.experiment.end()
+
+ def __del__(self) -> None:
+ self.close()