Skip to content

gradient_norm

gradient_norm

Gradient norm logging callback.

GradientNormCallback 

GradientNormCallback(log_every_n_steps=50)

Bases: Callback

Callback to log gradient norms during training.

Logs: - grad_norm: root step metric - train/grad_norm_step: Step metric under the train namespace - train/grad_norm_epoch: Mean sampled gradient norm over the epoch - train/grad_norm_max_epoch: Max sampled gradient norm over the epoch - train/grad_clip_active_ratio: Fraction of sampled steps above the clip threshold - Individual parameter group gradient norms

Initialize gradient norm callback.

Parameters:

Name Type Description Default
log_every_n_steps int

Log gradient norms every N steps

 50
Source code in src/versatil/training/callbacks/gradient_norm.py 
21
22
23
24
25
26
27
28
29
30
def __init__(self, log_every_n_steps: int = 50):
    """Initialize gradient norm callback.

    Args:
        log_every_n_steps: Log gradient norms every N steps
    """
    super().__init__()
    self.log_every_n_steps = log_every_n_steps
    self._epoch_grad_norms: list[float] = []
    self._epoch_grad_clip_active: list[float] = []

on_before_optimizer_step 

on_before_optimizer_step(trainer, pl_module, optimizer)

Log gradient norms before optimizer step.

Parameters:

Name Type Description Default
trainer Trainer

Lightning trainer

 required
pl_module LightningModule

Lightning module

 required
optimizer Optimizer

The optimizer

 required
Source code in src/versatil/training/callbacks/gradient_norm.py 
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def on_before_optimizer_step(
    self,
    trainer: pl.Trainer,
    pl_module: pl.LightningModule,
    optimizer: torch.optim.Optimizer,
) -> None:
    """Log gradient norms before optimizer step.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
        optimizer: The optimizer
    """
    if trainer.global_step % self.log_every_n_steps != 0:
        return

    grad_norm = self._compute_grad_norm(pl_module)
    self._epoch_grad_norms.append(grad_norm)
    clip_val = getattr(trainer, "gradient_clip_val", None)
    clip_active = (
        isinstance(clip_val, (int, float))
        and clip_val > 0.0
        and grad_norm > clip_val
    )
    self._epoch_grad_clip_active.append(float(clip_active))

    pl_module.log(
        "grad_norm",
        grad_norm,
        on_step=True,
        on_epoch=False,
        prog_bar=False,
        logger=True,
    )
    pl_module.log(
        "train/grad_clip_active_step",
        float(clip_active),
        on_step=True,
        on_epoch=False,
        prog_bar=False,
        logger=True,
    )
    pl_module.log(
        "train/grad_norm_step",
        grad_norm,
        on_step=True,
        on_epoch=False,
        prog_bar=False,
        logger=True,
    )

    if hasattr(optimizer, "param_groups") and len(optimizer.param_groups) > 1:
        for idx, param_group in enumerate(optimizer.param_groups):
            group_grad_norm = self._compute_grad_norm_for_params(
                param_group["params"]
            )
            pl_module.log(
                f"grad_norm_group_{idx}",
                group_grad_norm,
                on_step=True,
                on_epoch=False,
                prog_bar=False,
                logger=True,
            )
            pl_module.log(
                f"train/grad_norm_group_{idx}_step",
                group_grad_norm,
                on_step=True,
                on_epoch=False,
                prog_bar=False,
                logger=True,
            )

on_train_epoch_end 

on_train_epoch_end(trainer, pl_module)

Log epoch-level summaries so gradient norms appear with train charts.

Source code in src/versatil/training/callbacks/gradient_norm.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def on_train_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Log epoch-level summaries so gradient norms appear with train charts."""
    if not self._epoch_grad_norms:
        return

    grad_norms = np.asarray(self._epoch_grad_norms, dtype=np.float32)
    metrics = {
        "train/grad_norm_epoch": float(grad_norms.mean()),
        "train/grad_norm_max_epoch": float(grad_norms.max()),
        "epoch": trainer.current_epoch,
    }
    if self._epoch_grad_clip_active:
        metrics["train/grad_clip_active_ratio"] = float(
            np.asarray(self._epoch_grad_clip_active, dtype=np.float32).mean()
        )
    if trainer.logger is not None:
        trainer.logger.log_metrics(metrics, step=trainer.current_epoch)
    self._epoch_grad_norms.clear()
    self._epoch_grad_clip_active.clear()