Add lr and grad norm logging

Author: gonzalobenegas

configs/callbacks/default.yaml

@@ -1,6 +1,7 @@
 defaults:
   - model_summary
   - rich_progress_bar
+  - lr_monitor
   - _self_
 
 model_checkpoint:

configs/callbacks/lr_monitor.yaml

@@ -0,0 +1,3 @@
+lr_monitor:
+  _target_: lightning.pytorch.callbacks.LearningRateMonitor
+  logging_interval: step

configs/experiment/logging.yaml

@@ -0,0 +1,30 @@
+# @package _global_
+
+logger:
+  wandb:
+    name: debug-lr-gradnorm-logging
+
+trainer:
+  max_steps: 100
+  log_every_n_steps: 5
+  val_check_interval: 5
+  limit_val_batches: 2
+  check_val_every_n_epoch: null
+
+model:
+  net:
+    embedder:
+      embedding_dim: 32
+    encoder:
+      n_layers: 2
+  scheduler:
+    _target_: transformers.get_cosine_schedule_with_warmup
+    _partial_: true
+    num_warmup_steps: 10
+    num_training_steps: ${trainer.max_steps}
+
+data:
+  batch_size: 8
+  per_device_batch_size: 8
+
+compile: false

glm_experiments/models/bert_lit_module.py

@@ -6,6 +6,7 @@
 import torch.nn as nn
 from biofoundation.model.scoring import compute_llr_mlm
 from lightning import LightningModule
+from lightning.pytorch.utilities import grad_norm
 from sklearn.metrics import average_precision_score
 from torchmetrics.aggregation import CatMetric
 
@@ -134,3 +135,8 @@ def configure_optimizers(self) -> dict[str, Any]:
                 "interval": "step",
             },
         }
+
+    def on_before_optimizer_step(self, optimizer: torch.optim.Optimizer) -> None:
+        """Log gradient norm before optimizer step."""
+        norms = grad_norm(self, norm_type=2)
+        self.log("train/grad_norm", norms["grad_2.0_norm_total"])

tests/test_bert_lit_module.py

@@ -209,3 +209,29 @@ def test_validation_step_mlm_still_works(bert_lit_module):
     # Should not raise
     result = bert_lit_module.validation_step(batch, batch_idx=0, dataloader_idx=0)
     assert result is None
+
+
+def test_on_before_optimizer_step_logs_grad_norm(bert_lit_module):
+    """Test that on_before_optimizer_step computes and logs gradient norm."""
+    from lightning.pytorch.utilities import grad_norm
+
+    batch_size = 2
+    seq_len = 100
+
+    batch = {
+        "input_ids": torch.randint(0, 6, (batch_size, seq_len)),
+        "labels": torch.randint(0, 6, (batch_size, seq_len)),
+        "loss_weight": torch.ones(batch_size, seq_len),
+    }
+
+    # Forward and backward to populate gradients
+    loss = bert_lit_module.model_step(batch)
+    loss.backward()
+
+    # Compute expected grad norm
+    norms = grad_norm(bert_lit_module, norm_type=2)
+    expected_norm = norms["grad_2.0_norm_total"]
+
+    # Verify gradients exist and norm is reasonable
+    assert expected_norm > 0.0
+    assert torch.isfinite(torch.tensor(expected_norm))

tests/test_configs.py

@@ -1,5 +1,7 @@
 import hydra
+from hydra import compose, initialize
 from hydra.core.hydra_config import HydraConfig
+from lightning.pytorch.callbacks import LearningRateMonitor
 from omegaconf import DictConfig
 
 
@@ -35,3 +37,14 @@ def test_eval_config(cfg_eval: DictConfig) -> None:
     hydra.utils.instantiate(cfg_eval.data)
     hydra.utils.instantiate(cfg_eval.model)
     hydra.utils.instantiate(cfg_eval.trainer)
+
+
+def test_lr_monitor_callback_config() -> None:
+    """Test that LearningRateMonitor callback config instantiates correctly."""
+    with initialize(version_base="1.3", config_path="../configs/callbacks"):
+        cfg = compose(config_name="lr_monitor")
+
+    callback = hydra.utils.instantiate(cfg)
+
+    assert isinstance(callback, LearningRateMonitor)
+    assert callback.logging_interval == "step"