feat: change tracker API to initialize tracker early and track additional metrics.

requirements.txt

@@ -2,7 +2,7 @@ numpy
 accelerate>=0.20.3
 transformers>=4.34.1
 torch
-aim==3.17.5
+aim==3.18.1
 sentencepiece
 tokenizers>=0.13.3
 tqdm

tuning/config/configs.py

@@ -57,3 +57,7 @@ class TrainingArguments(transformers.TrainingArguments):
         default=False,
         metadata={"help": "Packing to be enabled in SFT Trainer, default is False"},
     )
+    tracker: str.lower = field(
+        default=None,
+        metadata={"help": "Experiment tracker to use. Requires additional configs, see tuning.configs/tracker_configs.py"}
+    )

tuning/config/tracker_configs.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+@dataclass
+class AimConfig:
+    # Name of the experiment
+    experiment: str = None
+    # 'repo' can point to a locally accessible directory (e.g., '~/.aim') or a remote repository hosted on a server.
+    # When 'remote_server_ip' or 'remote_server_port' is set, it designates a remote aim repo.
+    # Otherwise, 'repo' specifies the directory, with a default of None representing '.aim'.
+    aim_repo: str = None
+    aim_remote_server_ip: str = None
+    aim_remote_server_port: int = None
+    # Location of where run_hash is exported
+    aim_run_hash_export_path: str = None

tuning/sft_trainer.py

@@ -1,11 +1,11 @@
 # Standard
 from datetime import datetime
-from typing import Optional, Union
+from typing import Optional, Union, List
 import json
-import os
+import os, time
 
 # Third Party
-from peft.utils.other import fsdp_auto_wrap_policy
+import transformers
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -16,18 +16,20 @@
     TrainerCallback,
 )
 from transformers.utils import logging
+from peft.utils.other import fsdp_auto_wrap_policy
 from trl import DataCollatorForCompletionOnlyLM, SFTTrainer
 import datasets
 import fire
-import transformers
 
 # Local
-from tuning.aim_loader import get_aimstack_callback
-from tuning.config import configs, peft_config
+from tuning.config import configs, peft_config, tracker_configs
 from tuning.data import tokenizer_data_utils
 from tuning.utils.config_utils import get_hf_peft_config
 from tuning.utils.data_type_utils import get_torch_dtype
+from tuning.tracker.tracker import Tracker, get_tracker
+from tuning.tracker.aimstack_tracker import AimStackTracker
 
+logger = logging.get_logger("sft_trainer")
 
 class PeftSavingCallback(TrainerCallback):
     def on_save(self, args, state, control, **kwargs):
@@ -39,7 +41,6 @@ def on_save(self, args, state, control, **kwargs):
         if "pytorch_model.bin" in os.listdir(checkpoint_path):
             os.remove(os.path.join(checkpoint_path, "pytorch_model.bin"))
 
-
 class FileLoggingCallback(TrainerCallback):
     """Exports metrics, e.g., training loss to a file in the checkpoint directory."""
 
@@ -83,14 +84,15 @@ def _track_loss(self, loss_key, log_file, logs, state):
         with open(log_file, "a") as f:
             f.write(f"{json.dumps(log_obj, sort_keys=True)}\n")
 
-
 def train(
     model_args: configs.ModelArguments,
     data_args: configs.DataArguments,
     train_args: configs.TrainingArguments,
     peft_config: Optional[
         Union[peft_config.LoraConfig, peft_config.PromptTuningConfig]
     ] = None,
+    callbacks: Optional[List[TrainerCallback]] = None,
+    tracker: Optional[Tracker] = Tracker() # default tracker is dummy tracker
 ):
     """Call the SFTTrainer
 
@@ -105,8 +107,6 @@ def train(
     """
     run_distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
 
-    logger = logging.get_logger("sft_trainer")
-
     # Validate parameters
     if (not isinstance(train_args.num_train_epochs, float)) or (
         train_args.num_train_epochs <= 0
@@ -123,12 +123,16 @@ def train(
         train_args.fsdp_config = {"xla": False}
 
     task_type = "CAUSAL_LM"
+
+    model_load_time = time.time()
     model = AutoModelForCausalLM.from_pretrained(
         model_args.model_name_or_path,
         cache_dir=train_args.cache_dir,
         torch_dtype=get_torch_dtype(model_args.torch_dtype),
         use_flash_attention_2=model_args.use_flash_attn,
     )
+    model_load_time = time.time() - model_load_time
+    tracker.track(metric=model_load_time, name='model_load_time')
 
     peft_config = get_hf_peft_config(task_type, peft_config)
 
@@ -219,11 +223,6 @@ def train(
         formatted_validation_dataset = json_dataset["validation"].map(format_dataset)
         logger.info(f"Validation dataset length is {len(formatted_validation_dataset)}")
 
-    aim_callback = get_aimstack_callback()
-    file_logger_callback = FileLoggingCallback(logger)
-    peft_saving_callback = PeftSavingCallback()
-    callbacks = [aim_callback, peft_saving_callback, file_logger_callback]
-
     if train_args.packing:
         logger.info("Packing is set to True")
         data_collator = None
@@ -269,7 +268,6 @@ def train(
         )
     trainer.train()
 
-
 def main(**kwargs):
     parser = transformers.HfArgumentParser(
         dataclass_types=(
@@ -278,6 +276,7 @@ def main(**kwargs):
             configs.TrainingArguments,
             peft_config.LoraConfig,
             peft_config.PromptTuningConfig,
+            tracker_configs.AimConfig,
         )
     )
     parser.add_argument(
@@ -286,23 +285,57 @@ def main(**kwargs):
         choices=["pt", "lora", None, "none"],
         default="pt",
     )
+    parser.add_argument(
+        "--extra_metadata",
+        type=str,
+        default=None,
+    )
     (
         model_args,
         data_args,
         training_args,
         lora_config,
         prompt_tuning_config,
-        peft_method,
+        aim_config,
+        additional,
         _,
     ) = parser.parse_args_into_dataclasses(return_remaining_strings=True)
-    if peft_method.peft_method == "lora":
-        tune_config = lora_config
-    elif peft_method.peft_method == "pt":
-        tune_config = prompt_tuning_config
+
+    peft_method = additional.peft_method
+    if peft_method =="lora":
+        tune_config=lora_config
+    elif peft_method =="pt":
+        tune_config=prompt_tuning_config
     else:
-        tune_config = None
-    train(model_args, data_args, training_args, tune_config)
+        tune_config=None
+
+    tracker_name = training_args.tracker
+    if tracker_name == "aim":
+        tracker_config=aim_config
+    else:
+        tracker_config=None
+
+    # Initialize the tracker early so we can calculate custom metrics like model_load_time.
+    tracker = get_tracker(tracker_name, tracker_config)
+
+    # Initialize callbacks
+    file_logger_callback = FileLoggingCallback(logger)
+    peft_saving_callback = PeftSavingCallback()
+    callbacks = [peft_saving_callback, file_logger_callback]
+
+    tracker_callback = tracker.get_hf_callback()
+    if tracker_callback is not None:
+        callbacks.append(tracker_callback)
+
+    # track extra metadata
+    if additional.extra_metadata is not None:
+        try:
+            metadata = json.loads(additional.extra_metadata)
+            tracker.track_metadata(metadata)
+        except:
+            logger.error("failed while parsing extra metadata. pass a valid json")
 
+    train(model_args, data_args, training_args, tune_config, callbacks, tracker)
 
 if __name__ == "__main__":
     fire.Fire(main)

tuning/tracker/aimstack_tracker.py

@@ -0,0 +1,47 @@
+# Standard
+import os
+
+from tuning.tracker.tracker import Tracker
+
+# Third Party
+from aim.hugging_face import AimCallback
+
+class AimStackTracker(Tracker):
+
+    def __init__(self, tracker_config):
+        super().__init__(tracker_config)
+        c = self.config
+        exp = c.experiment
+        ip = c.aim_remote_server_ip
+        port = c.aim_remote_server_port
+        repo = c.aim_repo
+        hash_export_path = c.aim_run_hash_export_path
+
+        if (ip is not None and port is not None):
+            aim_callback = AimCallback(
+                                repo="aim://" + ip +":"+ port + "/",
+                                experiment=exp
+                            )
+        if repo:
+            aim_callback = AimCallback(repo=repo, experiment=exp)
+        else:
+            aim_callback = AimCallback(experiment=exp)
+
+        run = aim_callback.experiment # Initialize Aim run
+        run_hash = run.hash # Extract the hash
+
+        # store the run hash
+        if hash_export_path:
+            with open(hash_export_path, 'w') as f:
+                f.write(str(run_hash)+'\n')
+
+        # Save Internal State
+        self.hf_callback = aim_callback
+        self.run = run
+
+    def get_hf_callback(self):
+        return self.hf_callback
+
+    def track(self, metric, name, stage='additional_metrics'):
+        context={'subset' : stage}
+        self.run.track(metric, name=name, context=context)

tuning/tracker/tracker.py

@@ -0,0 +1,29 @@
+# Generic Tracker API
+
+from tuning.tracker.aimstack_tracker import AimStackTracker
+
+class Tracker:
+    def __init__(self, tracker_config) -> None:
+        self.config = tracker_config
+
+    def get_hf_callback():
+        return None
+
+    def track(self, metric, name, stage):
+        pass
+
+    # Metadata passed here is supposed to be a KV object
+    # Key being the name and value being the metric to track.
+    def track_metadata(self, metadata=None):
+        if metadata is None or not isinstance(metadata, dict):
+            return
+        for k, v in metadata.items():
+            self.track(name=k, metric=v)
+
+def get_tracker(tracker_name, tracker_config):
+    if tracker_name == 'aim':
+        if tracker_config is not None:
+            tracker = AimStackTracker(tracker_config)
+    else:
+        tracker = Tracker()
+    return tracker