Added wandb options including whther to write grads to logs.

Author: joaomonteirof

launcher.sh

@@ -12,6 +12,10 @@ NSTEPS=100_000
 torchrun --nproc_per_node $NGPUS \
 trainval.py \
 -e mlm \
+--wandb-entity-name jbmf \
+--wandb-project-name tf-encoder \
+--wandb-run-name mlm \
+--wandb-log-gradients false \
 --steps $NSTEPS \
 -sb $PATH_TO_LOG \
 --train_data_name $TRAIN_DATA_NAME \

src/hf_trainer.py

@@ -1,4 +1,5 @@
 import os
+import wandb
 import torch
 from torch.utils.data.dataset import Dataset
 import transformers
@@ -16,6 +17,9 @@
     retrieval_eval,
 )
 from src.datasets_loader import Collator
+from src.logging_callback import LoggingCallback
+
+from functools import partial
 
 
 def compute_metrics(eval_pred: PredictionOutput) -> Dict[str, float]:
@@ -221,6 +225,10 @@ def get_trainer(
     log_every: int = 100,
     local_rank: int = 0,
     deepspeed_cfg_path: str = None,
+    wandb_entity_name: str = None,
+    wandb_project_name: str = None,
+    wandb_run_name: str = None,
+    wandb_log_grads: bool = False,
 ) -> CustomTrainer:
     """Intanstiates Trainer object.
 
@@ -234,6 +242,10 @@ def get_trainer(
         log_every (int): Logging interval.
         local_rank (int): Device id for distributed training.
         deepspeed_cfg_path (str, Optional): Optional path to deepspeed config.
+        wandb_entity_name (str, optional): Wandb entity. Defaults to None.
+        wandb_project_name (str, optional): Project name for wandb. Defaults to None.
+        wandb-run-name (str, optional): Run id name for wandb. Defaults to None.
+        wandb_log_grads (bool, optional): Whether to write grads on wandb logs. Defaults to False.
 
     Returns:
         CustomTrainer: Trainer object.
@@ -265,18 +277,27 @@ def get_trainer(
         save_strategy="steps",
         save_steps=log_every,
         evaluation_strategy="steps",
-        report_to="wandb",
+        # report_to="wandb",
     )
 
     encoder = get_encoder(exp_dict=exp_dict)
 
+    wandb.init(
+        name=wandb_run_name,
+        entity=wandb_entity_name,
+        project=wandb_project_name,
+    )
+
     trainer = CustomTrainer(
         model=encoder,
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=valid_dataset,
         compute_metrics=compute_metrics,
         data_collator=collate_fn,
+        callbacks=[
+            LoggingCallback(log_grads=wandb_log_grads),
+        ],
     )
 
     return trainer

src/logging_callback.py

@@ -0,0 +1,21 @@
+from transformers.integrations import WandbCallback
+
+
+class LoggingCallback(WandbCallback):
+    """
+    Overrigding WandbCallback to optionally turn off gradient logging.
+    """
+
+    def __init__(self, log_grads: bool):
+
+        super().__init__()
+
+        self.log_grads = log_grads
+
+    def setup(self, args, state, model, **kwargs):
+
+        super().setup(args, state, model, **kwargs)
+        _watch_model = "all" if self.log_grads else "parameters"
+        self._wandb.watch(
+            model, log=_watch_model, log_freq=max(100, args.logging_steps)
+        )

src/training_args.py

@@ -1,8 +1,26 @@
 import argparse
 
 
+def parse_bool_flag(s: str) -> bool:
+    """Parse boolean arguments from the command line.
+
+    Args:
+        s (str): Input arg string.
+
+    Returns:
+        bool: _description_
+    """
+    _FALSY_STRINGS = {"off", "false", "0"}
+    _TRUTHY_STRINGS = {"on", "true", "1"}
+    if s.lower() in _FALSY_STRINGS:
+        return False
+    elif s.lower() in _TRUTHY_STRINGS:
+        return True
+    else:
+        raise argparse.ArgumentTypeError("Invalid value for a boolean flag")
+
+
 def parse_args():
-    # Specify arguments regarding save directory and job scheduler
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "-e",
@@ -45,6 +63,22 @@ def parse_args():
         type=int,
         help="Number of iterations to wait before logging training scores.",
     )
+    parser.add_argument(
+        "--wandb-entity-name",
+        type=str,
+        default="bigcode",
+        help="Name of wandb entity for reporting.",
+    )
+    parser.add_argument(
+        "--wandb-project-name", type=str, default=None, help="Name of wandb project."
+    )
+    parser.add_argument("--wandb-run-name", type=str, default=None, help="Name of run.")
+    parser.add_argument(
+        "--wandb-log-gradients",
+        type=parse_bool_flag,
+        default="false",
+        help="Whether to write gradients to wandb logs.",
+    )
     parser.add_argument(
         "--dist_url",
         default="env://",

trainval_toolkit.py

@@ -57,6 +57,10 @@ def train(exp_dict, savedir, args):
         valid_dataset=gfg_test_data,
         collate_fn=collate_fn,
         log_every=args.log_every,
+        wandb_entity_name=args.wandb_entity_name,
+        wandb_project_name=args.wandb_project_name,
+        wandb_run_name=args.wandb_run_name,
+        wandb_log_grads=args.wandb_log_gradients,
         local_rank=args.local_rank,
         deepspeed_cfg_path=args.deepspeed,
     )