add wandb integration

run_clm_flax.py

@@ -53,6 +53,7 @@ from transformers import (
     is_tensorboard_available,
 )
 from transformers.testing_utils import CaptureLogger
+import wandb
 
 
 logger = logging.getLogger(__name__)
@@ -232,6 +233,13 @@ def main():
     # or by passing the --help flag to this script.
     # We now keep distinct sets of args, for a cleaner separation of concerns.
 
+    if jax.process_index() == 0:
+        wandb.init(
+            entity = os.getenv("WANDB_ENTITY", "indonesian-nlp"),
+            project = os.getenv("WANDB_PROJECT", "huggingface"),
+            sync_tensorboard =True
+        )
+
     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
     if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
         # If we pass only one argument to the script and it's the path to a json file,
@@ -250,6 +258,13 @@ def main():
             f"Output directory ({training_args.output_dir}) already exists and is not empty."
             "Use --overwrite_output_dir to overcome."
         )
+    # log your configs with wandb.config, accepts a dict
+    if jax.process_index() == 0:
+        wandb.config.update(training_args)  # optional, log your configs
+        wandb.config.update(model_args)  # optional, log your configs
+        wandb.config.update(data_args)  # optional, log your configs
+
+    wandb.config['test_log'] = 12345  # log additional things
 
     # Make one log on every process with the configuration for debugging.
     logging.basicConfig(
@@ -591,6 +606,8 @@ def main():
                 epochs.write(
                     f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
                 )
+                if jax.process_index() == 0:
+                    wandb.log({'my_metric': train_metrics})
 
                 train_metrics = []
 
@@ -623,6 +640,7 @@ def main():
                 if has_tensorboard and jax.process_index() == 0:
                     cur_step = epoch * (len(train_dataset) // train_batch_size)
                     write_eval_metric(summary_writer, eval_metrics, cur_step)
+                    wandb.log({'my_metric': eval_metrics})
 
             if cur_step % training_args.save_steps == 0 and cur_step > 0:
                 # save checkpoint after each epoch and push checkpoint to the hub