From b28af3b6aad6a8ef72871386b6d2357807dd8fd6 Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Fri, 27 Sep 2024 16:23:45 +0200 Subject: [PATCH 1/8] Add first version of a Comet integration --- .isort.cfg | 2 +- src/axolotl/cli/__init__.py | 3 + src/axolotl/core/trainer_builder.py | 12 +++ src/axolotl/utils/callbacks/__init__.py | 30 ++++++ src/axolotl/utils/comet_.py | 91 +++++++++++++++++++ .../config/models/input/v0_4_1/__init__.py | 14 +++ 6 files changed, 151 insertions(+), 1 deletion(-) create mode 100644 src/axolotl/utils/comet_.py diff --git a/.isort.cfg b/.isort.cfg index 79067a7c91..e487797321 100644 --- a/.isort.cfg +++ b/.isort.cfg @@ -1,3 +1,3 @@ [settings] profile=black -known_third_party=wandb +known_third_party=wandb,comet_ml diff --git a/src/axolotl/cli/__init__.py b/src/axolotl/cli/__init__.py index 13c5b4ab58..a7c8c3d7dc 100644 --- a/src/axolotl/cli/__init__.py +++ b/src/axolotl/cli/__init__.py @@ -44,6 +44,7 @@ from axolotl.utils.tokenization import check_dataset_labels from axolotl.utils.trainer import prepare_opinionated_env, prepare_optim_env from axolotl.utils.wandb_ import setup_wandb_env_vars +from axolotl.utils.comet_ import setup_comet_env_vars project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) src_dir = os.path.join(project_root, "src") @@ -421,6 +422,8 @@ def load_cfg(config: Union[str, Path] = Path("examples/"), **kwargs): setup_mlflow_env_vars(cfg) + setup_comet_env_vars(cfg) + return cfg diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index 249398f850..10dfc17253 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -54,6 +54,7 @@ GPUStatsCallback, LossWatchDogCallback, SaveAxolotlConfigtoWandBCallback, + SaveAxolotlConfigtoCometCallback, SaveBetterTransformerModelCallback, SaveModelCallback, bench_eval_callback_factory, @@ -1104,6 +1105,10 @@ def get_callbacks(self) -> List[TrainerCallback]: callbacks.append( SaveAxolotlConfigtoMlflowCallback(self.cfg.axolotl_config_path) ) + if self.cfg.use_comet: + callbacks.append( + SaveAxolotlConfigtoCometCallback(self.cfg.axolotl_config_path) + ) return callbacks @@ -1172,6 +1177,11 @@ def get_post_trainer_create_callbacks(self, trainer): trainer, self.tokenizer, "mlflow" ) callbacks.append(LogPredictionCallback(self.cfg)) + if self.cfg.use_comet and self.cfg.eval_table_size > 0: + LogPredictionCallback = log_prediction_callback_factory( + trainer, self.tokenizer, "comet_ml" + ) + callbacks.append(LogPredictionCallback(self.cfg)) if self.cfg.do_bench_eval: callbacks.append(bench_eval_callback_factory(trainer, self.tokenizer)) @@ -1423,6 +1433,8 @@ def build(self, total_num_steps): report_to.append("mlflow") if self.cfg.use_tensorboard: report_to.append("tensorboard") + if self.cfg.use_comet: + report_to.append("comet_ml") training_arguments_kwargs["report_to"] = report_to training_arguments_kwargs["run_name"] = ( diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py index 73715b06ab..345331fedc 100644 --- a/src/axolotl/utils/callbacks/__init__.py +++ b/src/axolotl/utils/callbacks/__init__.py @@ -16,6 +16,7 @@ import torch import torch.distributed as dist import wandb +import comet_ml from datasets import load_dataset from optimum.bettertransformer import BetterTransformer from tqdm import tqdm @@ -747,6 +748,10 @@ def log_table_from_dataloader(name: str, table_dataloader): artifact_file="PredictionsVsGroundTruth.json", tracking_uri=tracking_uri, ) + elif logger == "comet_ml": + experiment = comet_ml.get_running_experiment() + if experiment: + experiment.log_table(f"{name} - Predictions vs Ground Truth.csv", pd.DataFrame(table_data)) if is_main_process(): log_table_from_dataloader("Eval", eval_dataloader) @@ -790,6 +795,31 @@ def on_train_begin( return control +class SaveAxolotlConfigtoCometCallback(TrainerCallback): + """Callback to save axolotl config to comet""" + + def __init__(self, axolotl_config_path): + self.axolotl_config_path = axolotl_config_path + + def on_train_begin( + self, + args: AxolotlTrainingArguments, # pylint: disable=unused-argument + state: TrainerState, # pylint: disable=unused-argument + control: TrainerControl, + **kwargs, # pylint: disable=unused-argument + ): + if is_main_process(): + try: + comet_experiment = comet_ml.start(source="axolotl") + comet_experiment._log_asset(self.axolotl_config_path, file_name="axolotl-config", framework="axolotl") + LOG.info( + "The Axolotl config has been saved to the Comet Experiment under assets." + ) + except (FileNotFoundError, ConnectionError) as err: + LOG.warning(f"Error while saving Axolotl config to Comet: {err}") + return control + + class SaveModelCallback(TrainerCallback): """Callback to save model on train end""" diff --git a/src/axolotl/utils/comet_.py b/src/axolotl/utils/comet_.py new file mode 100644 index 0000000000..d64aea7d5b --- /dev/null +++ b/src/axolotl/utils/comet_.py @@ -0,0 +1,91 @@ +"""Module for wandb utilities""" + +import os +import logging + +from axolotl.utils.dict import DictDefault + +LOG = logging.getLogger("axolotl.utils.comet_") + +COMET_ENV_MAPPING_OVERRIDE = { + "comet_mode": "COMET_START_MODE", + "comet_online": "COMET_START_ONLINE", +} +COMET_EXPERIMENT_CONFIG_ENV_MAPPING_OVERRIDE = { + "auto_histogram_activation_logging": "COMET_AUTO_LOG_HISTOGRAM_ACTIVATIONS", + "auto_histogram_epoch_rate": "COMET_AUTO_LOG_HISTOGRAM_EPOCH_RATE", + "auto_histogram_gradient_logging": "COMET_AUTO_LOG_HISTOGRAM_GRADIENTS", + "auto_histogram_tensorboard_logging": "COMET_AUTO_LOG_HISTOGRAM_TENSORBOARD", + "auto_histogram_weight_logging": "COMET_AUTO_LOG_HISTOGRAM_WEIGHTS", + "auto_log_co2": "COMET_AUTO_LOG_CO2", + "auto_metric_logging": "COMET_AUTO_LOG_METRICS", + "auto_metric_step_rate": "COMET_AUTO_LOG_METRIC_STEP_RATE", + "auto_output_logging": "COMET_AUTO_LOG_OUTPUT_LOGGER", + "auto_param_logging": "COMET_AUTO_LOG_PARAMETERS", + "comet_disabled": "COMET_AUTO_LOG_DISABLE", + "display_summary_level": "COMET_DISPLAY_SUMMARY_LEVEL", + "distributed_node_identifier": "COMET_DISTRIBUTED_NODE_IDENTIFIER", + "log_code": "COMET_AUTO_LOG_CODE", + "log_env_cpu": "COMET_AUTO_LOG_ENV_CPU", + "log_env_details": "COMET_AUTO_LOG_ENV_DETAILS", + "log_env_disk": "COMET_AUTO_LOG_ENV_DISK", + "log_env_gpu": "COMET_AUTO_LOG_ENV_GPU", + "log_env_host": "COMET_AUTO_LOG_ENV_HOST", + "log_env_network": "COMET_AUTO_LOG_ENV_NETWORK", + "log_git_metadata": "COMET_AUTO_LOG_GIT_METADATA", + "log_git_patch": "COMET_AUTO_LOG_GIT_PATCH", + "log_graph": "COMET_AUTO_LOG_GRAPH", + "name": "COMET_START_EXPERIMENT_NAME", + "offline_directory": "COMET_OFFLINE_DIRECTORY", + "parse_args": "COMET_AUTO_LOG_CLI_ARGUMENTS", + "tags": "COMET_START_EXPERIMENT_TAGS", +} + + +def python_value_to_environ_value(python_value): + if isinstance(python_value, bool): + if python_value is True: + return "true" + else: + return "false" + elif isinstance(python_value, list): # Comet only have one list of string parameter + return ",".join(map(str, python_value)) + else: + return python_value + + +def setup_comet_env_vars(cfg: DictDefault): + # TODO, we need to convert Axolotl configuration to environment variables + # as Transformers integration are call first and would create an + # Experiment first + # cfg.use_comet = True + + for key in cfg.keys(): + if key.startswith("comet_") and key != "comet_experiment_config": + value = cfg.get(key, "") + + if value and value != "": + env_variable_name = COMET_ENV_MAPPING_OVERRIDE.get(key, key.upper()) + final_value = python_value_to_environ_value(value) + print(f"SETTING ENV {env_variable_name=} {final_value=}") + os.environ[env_variable_name] = final_value + + if cfg.comet_experiment_config: + for key, value in cfg.comet_experiment_config.items(): + if value and value != "": + config_env_variable_name = COMET_EXPERIMENT_CONFIG_ENV_MAPPING_OVERRIDE.get( + key + ) + + if config_env_variable_name is None: + LOG.warning( + f"Unknown Comet Experiment Config name {key}, ignoring it" + ) + + final_value = python_value_to_environ_value(value) + print(f"SETTING ENV {config_env_variable_name=} {final_value=}") + os.environ[config_env_variable_name] = final_value + + # Enable comet if project name is present + if cfg.comet_project_name and len(cfg.comet_project_name) > 0: + cfg.use_comet = True diff --git a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py index 4e07c9260a..e8c6ce7624 100644 --- a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py +++ b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py @@ -484,6 +484,19 @@ def check_wandb_run(cls, data): return data +class CometConfig(BaseModel): + """Comet configuration subset""" + + use_comet: Optional[bool] = None + comet_api_key: Optional[str] = None + comet_workspace: Optional[str] = None + comet_project_name: Optional[str] = None + comet_experiment_key: Optional[str] = None + comet_mode: Optional[str] = None + comet_online: Optional[bool] = None + comet_experiment_config: Optional[Dict[str, Any]] = None + + class GradioConfig(BaseModel): """Gradio configuration subset""" @@ -504,6 +517,7 @@ class AxolotlInputConfig( HyperparametersConfig, WandbConfig, MLFlowConfig, + CometConfig, LISAConfig, GradioConfig, RemappedParameters, From d8f5caca2e1d7cdd11a27a11a20b9faac7448109 Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Tue, 1 Oct 2024 16:57:44 +0200 Subject: [PATCH 2/8] Remove debug prints --- src/axolotl/utils/comet_.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/axolotl/utils/comet_.py b/src/axolotl/utils/comet_.py index d64aea7d5b..be8e76229e 100644 --- a/src/axolotl/utils/comet_.py +++ b/src/axolotl/utils/comet_.py @@ -67,7 +67,6 @@ def setup_comet_env_vars(cfg: DictDefault): if value and value != "": env_variable_name = COMET_ENV_MAPPING_OVERRIDE.get(key, key.upper()) final_value = python_value_to_environ_value(value) - print(f"SETTING ENV {env_variable_name=} {final_value=}") os.environ[env_variable_name] = final_value if cfg.comet_experiment_config: @@ -83,7 +82,6 @@ def setup_comet_env_vars(cfg: DictDefault): ) final_value = python_value_to_environ_value(value) - print(f"SETTING ENV {config_env_variable_name=} {final_value=}") os.environ[config_env_variable_name] = final_value # Enable comet if project name is present From 2f2804b3bb063f2a71281904c343045b41cc281e Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Tue, 1 Oct 2024 18:26:24 +0200 Subject: [PATCH 3/8] Add test for Comet Configuration transformation to env variables --- src/axolotl/cli/__init__.py | 2 +- src/axolotl/core/trainer_builder.py | 2 +- src/axolotl/utils/callbacks/__init__.py | 13 ++- src/axolotl/utils/comet_.py | 26 +++--- tests/test_validation.py | 100 ++++++++++++++++++++++++ 5 files changed, 127 insertions(+), 16 deletions(-) diff --git a/src/axolotl/cli/__init__.py b/src/axolotl/cli/__init__.py index a7c8c3d7dc..32a8b8c41d 100644 --- a/src/axolotl/cli/__init__.py +++ b/src/axolotl/cli/__init__.py @@ -31,6 +31,7 @@ from axolotl.logging_config import configure_logging from axolotl.train import TrainDatasetMeta from axolotl.utils.chat_templates import chat_templates +from axolotl.utils.comet_ import setup_comet_env_vars from axolotl.utils.config import ( normalize_cfg_datasets, normalize_config, @@ -44,7 +45,6 @@ from axolotl.utils.tokenization import check_dataset_labels from axolotl.utils.trainer import prepare_opinionated_env, prepare_optim_env from axolotl.utils.wandb_ import setup_wandb_env_vars -from axolotl.utils.comet_ import setup_comet_env_vars project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) src_dir = os.path.join(project_root, "src") diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index 10dfc17253..8eec3e5146 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -53,8 +53,8 @@ EvalFirstStepCallback, GPUStatsCallback, LossWatchDogCallback, - SaveAxolotlConfigtoWandBCallback, SaveAxolotlConfigtoCometCallback, + SaveAxolotlConfigtoWandBCallback, SaveBetterTransformerModelCallback, SaveModelCallback, bench_eval_callback_factory, diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py index 345331fedc..30e932d750 100644 --- a/src/axolotl/utils/callbacks/__init__.py +++ b/src/axolotl/utils/callbacks/__init__.py @@ -10,13 +10,13 @@ from tempfile import NamedTemporaryFile from typing import TYPE_CHECKING, Any, Dict, List +import comet_ml import evaluate import numpy as np import pandas as pd import torch import torch.distributed as dist import wandb -import comet_ml from datasets import load_dataset from optimum.bettertransformer import BetterTransformer from tqdm import tqdm @@ -751,7 +751,10 @@ def log_table_from_dataloader(name: str, table_dataloader): elif logger == "comet_ml": experiment = comet_ml.get_running_experiment() if experiment: - experiment.log_table(f"{name} - Predictions vs Ground Truth.csv", pd.DataFrame(table_data)) + experiment.log_table( + f"{name} - Predictions vs Ground Truth.csv", + pd.DataFrame(table_data), + ) if is_main_process(): log_table_from_dataloader("Eval", eval_dataloader) @@ -811,7 +814,11 @@ def on_train_begin( if is_main_process(): try: comet_experiment = comet_ml.start(source="axolotl") - comet_experiment._log_asset(self.axolotl_config_path, file_name="axolotl-config", framework="axolotl") + comet_experiment._log_asset( + self.axolotl_config_path, + file_name="axolotl-config", + framework="axolotl", + ) LOG.info( "The Axolotl config has been saved to the Comet Experiment under assets." ) diff --git a/src/axolotl/utils/comet_.py b/src/axolotl/utils/comet_.py index be8e76229e..b4ecc80ad9 100644 --- a/src/axolotl/utils/comet_.py +++ b/src/axolotl/utils/comet_.py @@ -1,7 +1,7 @@ """Module for wandb utilities""" -import os import logging +import os from axolotl.utils.dict import DictDefault @@ -46,40 +46,44 @@ def python_value_to_environ_value(python_value): if isinstance(python_value, bool): if python_value is True: return "true" - else: - return "false" - elif isinstance(python_value, list): # Comet only have one list of string parameter + + return "false" + + if isinstance(python_value, int): + return str(python_value) + + if isinstance(python_value, list): # Comet only have one list of string parameter return ",".join(map(str, python_value)) - else: - return python_value + + return python_value def setup_comet_env_vars(cfg: DictDefault): # TODO, we need to convert Axolotl configuration to environment variables # as Transformers integration are call first and would create an # Experiment first - # cfg.use_comet = True for key in cfg.keys(): if key.startswith("comet_") and key != "comet_experiment_config": value = cfg.get(key, "") - if value and value != "": + if value is not None and value != "": env_variable_name = COMET_ENV_MAPPING_OVERRIDE.get(key, key.upper()) final_value = python_value_to_environ_value(value) os.environ[env_variable_name] = final_value if cfg.comet_experiment_config: for key, value in cfg.comet_experiment_config.items(): - if value and value != "": - config_env_variable_name = COMET_EXPERIMENT_CONFIG_ENV_MAPPING_OVERRIDE.get( - key + if value is not None and value != "": + config_env_variable_name = ( + COMET_EXPERIMENT_CONFIG_ENV_MAPPING_OVERRIDE.get(key) ) if config_env_variable_name is None: LOG.warning( f"Unknown Comet Experiment Config name {key}, ignoring it" ) + continue final_value = python_value_to_environ_value(value) os.environ[config_env_variable_name] = final_value diff --git a/tests/test_validation.py b/tests/test_validation.py index 35d0e265e7..771cb2cd7f 100644 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -9,6 +9,7 @@ import pytest from pydantic import ValidationError +from axolotl.utils.comet_ import setup_comet_env_vars from axolotl.utils.config import validate_config from axolotl.utils.config.models.input.v0_4_1 import AxolotlConfigWCapabilities from axolotl.utils.dict import DictDefault @@ -1329,3 +1330,102 @@ def test_wandb_set_disabled(self, minimal_cfg): os.environ.pop("WANDB_PROJECT", None) os.environ.pop("WANDB_DISABLED", None) + + +class TestValidationComet(BaseValidation): + """ + Validation test for comet + """ + + def test_comet_sets_env(self, minimal_cfg): + comet_config = { + "comet_api_key": "foo", + "comet_workspace": "some_workspace", + "comet_project_name": "some_project", + "comet_experiment_key": "some_experiment_key", + "comet_mode": "get_or_create", + "comet_online": False, + "comet_experiment_config": { + "auto_histogram_activation_logging": False, + "auto_histogram_epoch_rate": 2, + "auto_histogram_gradient_logging": True, + "auto_histogram_tensorboard_logging": False, + "auto_histogram_weight_logging": True, + "auto_log_co2": False, + "auto_metric_logging": True, + "auto_metric_step_rate": 15, + "auto_output_logging": False, + "auto_param_logging": True, + "comet_disabled": False, + "display_summary_level": 2, + "distributed_node_identifier": "some_distributed_node_identifier", + "log_code": True, + "log_env_cpu": False, + "log_env_details": True, + "log_env_disk": False, + "log_env_gpu": True, + "log_env_host": False, + "log_env_network": True, + "log_git_metadata": False, + "log_git_patch": True, + "log_graph": False, + "name": "some_name", + "offline_directory": "some_offline_directory", + "parse_args": True, + "tags": ["tag1", "tag2"], + }, + } + + cfg = DictDefault(comet_config) | minimal_cfg + + new_cfg = validate_config(cfg) + + setup_comet_env_vars(new_cfg) + + comet_env = { + key: value for key, value in os.environ.items() if key.startswith("COMET_") + } + + assert ( + len(comet_env) + == len(comet_config) + len(comet_config["comet_experiment_config"]) - 1 + ) + + assert comet_env == { + "COMET_API_KEY": "foo", + "COMET_AUTO_LOG_CLI_ARGUMENTS": "true", + "COMET_AUTO_LOG_CO2": "false", + "COMET_AUTO_LOG_CODE": "true", + "COMET_AUTO_LOG_DISABLE": "false", + "COMET_AUTO_LOG_ENV_CPU": "false", + "COMET_AUTO_LOG_ENV_DETAILS": "true", + "COMET_AUTO_LOG_ENV_DISK": "false", + "COMET_AUTO_LOG_ENV_GPU": "true", + "COMET_AUTO_LOG_ENV_HOST": "false", + "COMET_AUTO_LOG_ENV_NETWORK": "true", + "COMET_AUTO_LOG_GIT_METADATA": "false", + "COMET_AUTO_LOG_GIT_PATCH": "true", + "COMET_AUTO_LOG_GRAPH": "false", + "COMET_AUTO_LOG_HISTOGRAM_ACTIVATIONS": "false", + "COMET_AUTO_LOG_HISTOGRAM_EPOCH_RATE": "2", + "COMET_AUTO_LOG_HISTOGRAM_GRADIENTS": "true", + "COMET_AUTO_LOG_HISTOGRAM_TENSORBOARD": "false", + "COMET_AUTO_LOG_HISTOGRAM_WEIGHTS": "true", + "COMET_AUTO_LOG_METRIC_STEP_RATE": "15", + "COMET_AUTO_LOG_METRICS": "true", + "COMET_AUTO_LOG_OUTPUT_LOGGER": "false", + "COMET_AUTO_LOG_PARAMETERS": "true", + "COMET_DISPLAY_SUMMARY_LEVEL": "2", + "COMET_DISTRIBUTED_NODE_IDENTIFIER": "some_distributed_node_identifier", + "COMET_EXPERIMENT_KEY": "some_experiment_key", + "COMET_OFFLINE_DIRECTORY": "some_offline_directory", + "COMET_PROJECT_NAME": "some_project", + "COMET_START_EXPERIMENT_NAME": "some_name", + "COMET_START_EXPERIMENT_TAGS": "tag1,tag2", + "COMET_START_MODE": "get_or_create", + "COMET_START_ONLINE": "false", + "COMET_WORKSPACE": "some_workspace", + } + + for key in comet_env.keys(): + os.environ.pop(key, None) From 8a2511999a856ab36a8a30f1529f98135ecbe2ec Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Wed, 2 Oct 2024 15:10:07 +0200 Subject: [PATCH 4/8] Fix last lint warning --- src/axolotl/utils/callbacks/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py index 30e932d750..56e2913566 100644 --- a/src/axolotl/utils/callbacks/__init__.py +++ b/src/axolotl/utils/callbacks/__init__.py @@ -814,10 +814,10 @@ def on_train_begin( if is_main_process(): try: comet_experiment = comet_ml.start(source="axolotl") - comet_experiment._log_asset( + comet_experiment.log_other("Created from", "axolotl") + comet_experiment.log_asset( self.axolotl_config_path, file_name="axolotl-config", - framework="axolotl", ) LOG.info( "The Axolotl config has been saved to the Comet Experiment under assets." From 173af54650cc48cc4aa55633b3f8853cf04afef5 Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Wed, 2 Oct 2024 15:26:12 +0200 Subject: [PATCH 5/8] Update Readme for Comet logging documentation --- README.md | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index c84f1cb8c9..f6f4e4e806 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ Features: - Integrated with xformer, flash attention, [liger kernel](https://github.com/linkedin/Liger-Kernel), rope scaling, and multipacking - Works with single GPU or multiple GPUs via FSDP or Deepspeed - Easily run with Docker locally or on the cloud -- Log results and optionally checkpoints to wandb or mlflow +- Log results and optionally checkpoints to wandb, mlflow or Comet - And more! @@ -515,6 +515,22 @@ wandb_name: wandb_log_model: ``` +##### Comet Logging + +Make sure your `COMET_API_KEY` environment variable is set (recommended) or you login to wandb with `comet login`. + +- wandb options +```yaml +use_comet: +comet_api_key: +comet_workspace: +comet_project_name: +comet_experiment_key: +comet_mode: +comet_online: +comet_experiment_config: +``` + ##### Special Tokens It is important to have special tokens like delimiters, end-of-sequence, beginning-of-sequence in your tokenizer's vocabulary. This will help you avoid tokenization issues and help your model train better. You can do this in axolotl like this: From 3caaa2a84461ef7299f6b332ac99895a3d4e3a15 Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Thu, 3 Oct 2024 13:50:05 +0200 Subject: [PATCH 6/8] Update Comet integration to be optional, update code and tests --- src/axolotl/core/trainer_builder.py | 7 ++-- src/axolotl/utils/__init__.py | 6 +++- src/axolotl/utils/callbacks/__init__.py | 36 +++------------------ src/axolotl/utils/callbacks/comet_.py | 43 +++++++++++++++++++++++++ tests/test_validation.py | 5 ++- 5 files changed, 60 insertions(+), 37 deletions(-) create mode 100644 src/axolotl/utils/callbacks/comet_.py diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index 8eec3e5146..f583c99954 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -48,12 +48,11 @@ from axolotl.monkeypatch.multipack import SUPPORTED_MULTIPACK_MODEL_TYPES from axolotl.monkeypatch.relora import ReLoRACallback, ReLoRAScheduler -from axolotl.utils import is_mlflow_available +from axolotl.utils import is_comet_available, is_mlflow_available from axolotl.utils.callbacks import ( EvalFirstStepCallback, GPUStatsCallback, LossWatchDogCallback, - SaveAxolotlConfigtoCometCallback, SaveAxolotlConfigtoWandBCallback, SaveBetterTransformerModelCallback, SaveModelCallback, @@ -1106,6 +1105,8 @@ def get_callbacks(self) -> List[TrainerCallback]: SaveAxolotlConfigtoMlflowCallback(self.cfg.axolotl_config_path) ) if self.cfg.use_comet: + from axolotl.utils.callbacks.comet_ import SaveAxolotlConfigtoCometCallback + callbacks.append( SaveAxolotlConfigtoCometCallback(self.cfg.axolotl_config_path) ) @@ -1177,7 +1178,7 @@ def get_post_trainer_create_callbacks(self, trainer): trainer, self.tokenizer, "mlflow" ) callbacks.append(LogPredictionCallback(self.cfg)) - if self.cfg.use_comet and self.cfg.eval_table_size > 0: + if self.cfg.use_comet and is_comet_available() and self.cfg.eval_table_size > 0: LogPredictionCallback = log_prediction_callback_factory( trainer, self.tokenizer, "comet_ml" ) diff --git a/src/axolotl/utils/__init__.py b/src/axolotl/utils/__init__.py index 99dec79f1b..91545009ad 100644 --- a/src/axolotl/utils/__init__.py +++ b/src/axolotl/utils/__init__.py @@ -1,8 +1,12 @@ """ Basic utils for Axolotl """ -import importlib +import importlib.util def is_mlflow_available(): return importlib.util.find_spec("mlflow") is not None + + +def is_comet_available(): + return importlib.util.find_spec("comet_ml") is not None diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py index 56e2913566..5ec7a55668 100644 --- a/src/axolotl/utils/callbacks/__init__.py +++ b/src/axolotl/utils/callbacks/__init__.py @@ -10,7 +10,6 @@ from tempfile import NamedTemporaryFile from typing import TYPE_CHECKING, Any, Dict, List -import comet_ml import evaluate import numpy as np import pandas as pd @@ -30,7 +29,7 @@ ) from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy -from axolotl.utils import is_mlflow_available +from axolotl.utils import is_comet_available, is_mlflow_available from axolotl.utils.bench import log_gpu_memory_usage from axolotl.utils.callbacks.perplexity import Perplexity from axolotl.utils.config.models.input.v0_4_1 import AxolotlInputConfig @@ -748,7 +747,9 @@ def log_table_from_dataloader(name: str, table_dataloader): artifact_file="PredictionsVsGroundTruth.json", tracking_uri=tracking_uri, ) - elif logger == "comet_ml": + elif logger == "comet_ml" and is_comet_available(): + import comet_ml + experiment = comet_ml.get_running_experiment() if experiment: experiment.log_table( @@ -798,35 +799,6 @@ def on_train_begin( return control -class SaveAxolotlConfigtoCometCallback(TrainerCallback): - """Callback to save axolotl config to comet""" - - def __init__(self, axolotl_config_path): - self.axolotl_config_path = axolotl_config_path - - def on_train_begin( - self, - args: AxolotlTrainingArguments, # pylint: disable=unused-argument - state: TrainerState, # pylint: disable=unused-argument - control: TrainerControl, - **kwargs, # pylint: disable=unused-argument - ): - if is_main_process(): - try: - comet_experiment = comet_ml.start(source="axolotl") - comet_experiment.log_other("Created from", "axolotl") - comet_experiment.log_asset( - self.axolotl_config_path, - file_name="axolotl-config", - ) - LOG.info( - "The Axolotl config has been saved to the Comet Experiment under assets." - ) - except (FileNotFoundError, ConnectionError) as err: - LOG.warning(f"Error while saving Axolotl config to Comet: {err}") - return control - - class SaveModelCallback(TrainerCallback): """Callback to save model on train end""" diff --git a/src/axolotl/utils/callbacks/comet_.py b/src/axolotl/utils/callbacks/comet_.py new file mode 100644 index 0000000000..b29f997a86 --- /dev/null +++ b/src/axolotl/utils/callbacks/comet_.py @@ -0,0 +1,43 @@ +"""Comet module for trainer callbacks""" + +import logging +from typing import TYPE_CHECKING + +import comet_ml +from transformers import TrainerCallback, TrainerControl, TrainerState + +from axolotl.utils.distributed import is_main_process + +if TYPE_CHECKING: + from axolotl.core.trainer_builder import AxolotlTrainingArguments + +LOG = logging.getLogger("axolotl.callbacks") + + +class SaveAxolotlConfigtoCometCallback(TrainerCallback): + """Callback to save axolotl config to comet""" + + def __init__(self, axolotl_config_path): + self.axolotl_config_path = axolotl_config_path + + def on_train_begin( + self, + args: "AxolotlTrainingArguments", # pylint: disable=unused-argument + state: TrainerState, # pylint: disable=unused-argument + control: TrainerControl, + **kwargs, # pylint: disable=unused-argument + ): + if is_main_process(): + try: + comet_experiment = comet_ml.start(source="axolotl") + comet_experiment.log_other("Created from", "axolotl") + comet_experiment.log_asset( + self.axolotl_config_path, + file_name="axolotl-config", + ) + LOG.info( + "The Axolotl config has been saved to the Comet Experiment under assets." + ) + except (FileNotFoundError, ConnectionError) as err: + LOG.warning(f"Error while saving Axolotl config to Comet: {err}") + return control diff --git a/tests/test_validation.py b/tests/test_validation.py index 771cb2cd7f..6e0d0ad2a5 100644 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -9,7 +9,7 @@ import pytest from pydantic import ValidationError -from axolotl.utils.comet_ import setup_comet_env_vars +from axolotl.utils import is_comet_available from axolotl.utils.config import validate_config from axolotl.utils.config.models.input.v0_4_1 import AxolotlConfigWCapabilities from axolotl.utils.dict import DictDefault @@ -1332,12 +1332,15 @@ def test_wandb_set_disabled(self, minimal_cfg): os.environ.pop("WANDB_DISABLED", None) +@pytest.mark.skipif(is_comet_available() is False, reason="comet_ml is not installed") class TestValidationComet(BaseValidation): """ Validation test for comet """ def test_comet_sets_env(self, minimal_cfg): + from axolotl.utils.comet_ import setup_comet_env_vars + comet_config = { "comet_api_key": "foo", "comet_workspace": "some_workspace", From 98df4ad32a50c0da7c74c250c129abeb25f88a2e Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Thu, 3 Oct 2024 14:45:41 +0200 Subject: [PATCH 7/8] Add documentation for Comet configuration --- docs/config.qmd | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/config.qmd b/docs/config.qmd index e859999787..99a69a0973 100644 --- a/docs/config.qmd +++ b/docs/config.qmd @@ -267,6 +267,18 @@ mlflow_tracking_uri: # URI to mlflow mlflow_experiment_name: # Your experiment name hf_mlflow_log_artifacts: # set to true to copy each saved checkpoint on each save to mlflow artifact registry +# Comet configuration if you're using it +# Make sure your `COMET_API_KEY` environment variable is set (recommended) or you login to Comet with `comet login`. +# Check out our documentation for more details https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment-Creation/#comet_ml.start +use_comet: # Enable or disable Comet integration. +comet_api_key: # API key for Comet. Recommended to set via `comet login`. +comet_workspace: # Workspace name in Comet. Defaults to the user's default workspace. +comet_project_name: # Project name in Comet. Defaults to Uncategorized. +comet_experiment_key: # Identifier for the experiment. Used to append data to an existing experiment or control the key of new experiments. Default to a random key. +comet_mode: # Create a new experiment ("create") or log to an existing one ("get"). Default ("get_or_create") auto-selects based on configuration. +comet_online: # Set to True to log data to Comet server, or False for offline storage. Default is True. +comet_experiment_config: # Dictionary for additional configuration settings, see the doc for more details. + # Where to save the full-finetuned model to output_dir: ./completed-model From 6b86a82c211d9f54ebf31963aeda5376ee739b5c Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Thu, 3 Oct 2024 15:57:53 +0200 Subject: [PATCH 8/8] Add missing check --- src/axolotl/core/trainer_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index f583c99954..79f4539825 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -1104,7 +1104,7 @@ def get_callbacks(self) -> List[TrainerCallback]: callbacks.append( SaveAxolotlConfigtoMlflowCallback(self.cfg.axolotl_config_path) ) - if self.cfg.use_comet: + if self.cfg.use_comet and is_comet_available(): from axolotl.utils.callbacks.comet_ import SaveAxolotlConfigtoCometCallback callbacks.append(