1 lightning.training.gradient_logging

lightning.training.gradient_logging

Select VIN submodules and compute detached gradient-norm diagnostics.

The helpers aggregate parameter gradients without modifying autograd state and bound logging cardinality through explicit depth, include, and exclude rules.

1.1 Classes

Name	Description
GradNormLoggingConfig	Configuration for gradient norm logging.
OptimizerDiagnosticsCallback	Observe gradients before stepping, after Lightning’s AMP unscaling.
OptimizerDiagnosticsConfig	Bound scalar diagnostics at real, unscaled optimizer-step boundaries.

1.1.1 GradNormLoggingConfig

lightning.training.gradient_logging.GradNormLoggingConfig()

Configuration for gradient norm logging.

1.1.1.1 Attributes

Name	Description
enabled	Whether to log gradient norms during training.
group_depth	Module depth relative to vin (1 = vin, 2 = vin.).
include	Glob patterns to include explicitly (relative to vin).
exclude	Glob patterns to exclude (relative to vin).
norm_type	Norm type used for logging.
max_items	Maximum number of modules to log (None disables the cap).

1.1.2 OptimizerDiagnosticsCallback

lightning.training.gradient_logging.OptimizerDiagnosticsCallback(config)

Observe gradients before stepping, after Lightning’s AMP unscaling.

Metrics use the module’s existing log_dict route, including offline W&B when attached. With no logger they remain available in callback metrics. Gradient norms precede automatic clipping and the QH module’s clipping hook. Other manual modules determine their own clipping order. Observation never mutates gradients. train-optimizer/observed_attempted_step is the zero-based trainer.global_step before this attempt; it travels with carried callback metrics and does not count completed raw AMP updates.

1.1.2.1 Methods

Name	Description
on_before_optimizer_step	Log bounded detached scalar observations for this optimizer attempt.

1.1.2.1.1 on_before_optimizer_step

lightning.training.gradient_logging.OptimizerDiagnosticsCallback.on_before_optimizer_step(
    trainer,
    pl_module,
    optimizer,
)

Log bounded detached scalar observations for this optimizer attempt.

1.1.3 OptimizerDiagnosticsConfig

lightning.training.gradient_logging.OptimizerDiagnosticsConfig()

Bound scalar diagnostics at real, unscaled optimizer-step boundaries.

Summary observes only optimizer-owned parameters. Detailed mode adds capped per-parameter norms; neither mode retains graphs, activations, or tensors. These rank-local observations do not establish a successful AMP update: GradScaler may subsequently skip a non-finite step.

1.1.3.1 Attributes

Name	Description
level	Off performs no diagnostics; detailed adds parameter-level L2 norms.
every_n_steps	Observe at this Lightning optimizer-attempt cadence; AMP skips still advance global_step.
max_items	Cap detailed parameter norms and per-optimizer parameter-group LR metrics.