add wandb logger

NVIDIA · Nov 11, 2024 · ede89e3 · ede89e3
1 parent ea62e6a
commit ede89e3
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/tests/collections/llm/gpt/model/test_hyena.py b/tests/collections/llm/gpt/model/test_hyena.py
@@ -19,7 +19,7 @@
 
 import torch
 from megatron.core.optimizer import OptimizerConfig
-from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 
 from nemo import lightning as nl
 from nemo.collections import llm
@@ -135,10 +135,12 @@ def get_args():
     callbacks = [checkpoint_callback]
 
     loggers = []
-    tensorboard_logger = TensorBoardLogger(
-        save_dir='dummy',  ## NOTE: this gets overwritten by default
+    wandb_logger = WandbLogger(
+        name=f"hyena-size-{args.model_size}-TP{args.tensor_parallel_size}-PP{args.pipeline_model_parallel_size}-CP{args.context_parallel_size}",
+        project="hyena_ux",
+        save_dir=args.experiment_dir,
     )
-    loggers.append(tensorboard_logger)
+    loggers.append(wandb_logger)
 
     opt_config = OptimizerConfig(
         optimizer='adam',