Make training-callback and LR-scheduler monitored metrics configurable via YAML by p-grontas · Pull Request #85 · gridfm/gridfm-graphkit

p-grontas · 2026-07-14T12:17:16Z

Summary

Makes the validation metric monitored by the training callbacks and the LR scheduler
configurable through the YAML callbacks: section, instead of being hard-coded to
"Validation loss". Three independent keys are introduced:

early_stopping_monitor — metric watched by EarlyStopping.
checkpoint_monitor — metric watched by best-model saving and ModelCheckpoint.
lr_scheduler_monitor — metric watched by the ReduceLROnPlateau scheduler.

Each defaults to Validation loss when omitted. All are minimized (mode="min", lower is
better); the direction is fixed and not configurable. This lets a run be driven by, e.g., a
per-layer physics residual (Validation layer_11_residual) rather than the aggregate loss.

Changes

Callback wiring (cli.py): get_training_callbacks reads callbacks.early_stopping_monitor
and callbacks.checkpoint_monitor, each defaulting to DEFAULT_MONITOR="Validation loss".
LR scheduler (tasks/base_task.py): ReduceLROnPlateau now reads
callbacks.lr_scheduler_monitor (previously hard-coded to "Validation loss").
Fail-loud on bad metric names: every monitor key aborts the run rather than silently
no-op'ing.
- checkpoint_monitor — checked explicitly by SaveBestModelStateDict (training/callbacks.py):
  it skips the sanity-check pass, then raises KeyError listing the available metrics. This
  is the sole consumer of the key — ModelCheckpoint is configured monitor-free
  (save_last=True, save_top_k=0, only writes last.ckpt), so it no longer watches any metric.
- early_stopping_monitor — EarlyStopping(strict=True) (Lightning's default, now set
  explicitly in cli.py) raises when the metric is missing.
- lr_scheduler_monitor — the ReduceLROnPlateau scheduler dict sets strict=True explicitly
  (base_task.py), so Lightning raises MisconfigurationException with the available metrics
  at the end of the first training epoch.
Configs: the three keys added to all examples/config/*.yaml and tests/config/*.yaml.
Docs: docs/quick_start/yaml_config.md documents the new keys and constraints; README
CLI tables updated.
Tests: two unit tests in tests/test_pipeline.py cover the config-driven wiring and the
default fallback.
Repo hygiene & CI: pre-commit runs at the pre-push stage and is re-enabled in CI; secrets
baseline updated (ignore *.ipynb, whitelist a known false positive); CONTRIBUTING checklist
- integration-tests step added; README/CONTRIBUTING document a from-source
  torch-scatter/torch-sparse install for PyTorch versions without a prebuilt PyG wheel.
  Formatting-only fixes applied to integrationtests/test_base_set.py.
Dependencies & security: build setuptools >= 83.0.0, pinned in CI, to clear advisory
PYSEC-2026-3447.

Backward compatibility

No behavior change for existing configs: all three keys default to Validation loss when
absent, matching the previous hard-coded behavior.

Usage

callbacks:
  patience: 100
  tol: 0
  early_stopping_monitor: Validation loss
  checkpoint_monitor: Validation loss
  lr_scheduler_monitor: Validation layer_11_residual

Testing

MLFLOW_ALLOW_FILE_STORE=true pytest tests/test_pipeline.py

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

p-grontas · 2026-07-15T09:33:49Z

@albanpuech ready for review

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

albanpuech

THANK YOU!!
just a few things to check

albanpuech · 2026-07-16T13:00:06Z

      - name: Install dependencies
        run: |
-          python -m pip install --upgrade pip wheel
+          # setuptools>=83.0.0 clears PYSEC-2026-3447


@romeokienzler @p-grontas Maybe for a future PR we will want to also pin the version of setuptools in the pyproject... but no need to do it now

albanpuech · 2026-07-16T13:03:01Z

+Note: A monitored metric name must exactly match a logged validation metric.
+For example, `Validation layer_11_residual` is only produced by the
+`LayeredWeightedPhysics` loss and requires a model deep enough to have a
+layer index 11 (a 12-layer model). Pointing a monitor at a metric that is
+never logged aborts the run once training begins.


albanpuech · 2026-07-16T13:06:00Z

  patience: 100
  tol: 0
+  early_stopping_monitor: Validation loss
+  checkpoint_monitor: Validation layer_11_residual


Suggested change

checkpoint_monitor: Validation layer_11_residual

checkpoint_monitor: Validation loss

albanpuech · 2026-07-16T13:06:11Z

  patience: 100
  tol: 0
+  early_stopping_monitor: Validation loss
+  checkpoint_monitor: Validation layer_11_residual


Suggested change

checkpoint_monitor: Validation layer_11_residual

checkpoint_monitor: Validation loss

albanpuech · 2026-07-16T13:08:41Z

  patience: 100
  tol: 0
+  early_stopping_monitor: Validation loss
+  checkpoint_monitor: Validation layer_11_residual


Suggested change

checkpoint_monitor: Validation layer_11_residual

checkpoint_monitor: Validation loss

albanpuech · 2026-07-16T13:09:13Z

  patience: 100
  tol: 0
+  early_stopping_monitor: Validation loss
+  checkpoint_monitor: Validation layer_11_residual


Suggested change

checkpoint_monitor: Validation layer_11_residual

checkpoint_monitor: Validation loss

albanpuech · 2026-07-16T13:13:20Z

+        lr_scheduler_monitor = getattr(
+            self.args.callbacks,
+            "lr_scheduler_monitor",
+            "Validation loss",


could use DEFAULT_MONITOR https://github.com/gridfm/gridfm-graphkit/pull/85/changes#diff-2453c8aec0e4d354dfe0c000a28790240cf2de314239087c06e01902066e3a5dR30

DEFAULT_MONITOR should be defined somewhere else than in cli in this case

otherwise, fine to not use DEFAULT_MONITOR, but it should be clear in cli.py that DEFAULT_MONITOR is not the the LR scheduler

albanpuech · 2026-07-16T13:20:24Z

@@ -1,3 +1,5 @@
+default_install_hook_types: [pre-push]


what is this?

If you install the pre-commit hooks (with pre-commit install), then when you git push, the pre-commit checks run and the branch is only pushed if they pass. So that the pre-commits are not forgotten and caught by the CI

albanpuech · 2026-07-16T13:21:51Z

+
+* [ ] **Rebase your branch onto the latest `main` before opening a PR.**
+* [ ] **Sign off all commits** with the Developer Certificate of Origin (`git commit -s`). See [Developer Certificate of Origin signoff].
+* [ ] Open a PR with a short description of your changes and add Alban Puech as a reviewer.


Suggested change

* [ ] Open a PR with a short description of your changes and add Alban Puech as a reviewer.

* [ ] Open a PR with a short description of your changes.

albanpuech · 2026-07-16T13:25:51Z

@romeokienzler JFYI

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

albanpuech · 2026-07-16T13:45:36Z

Some comments from cc are useful: @p-grontas

DEFAULT_MONITOR not reused: base_task.py:123 hardcodes "Validation loss" for the LR scheduler. Move DEFAULT_MONITOR to a neutral module (not cli.py) and use it in both places — importing cli→task is the wrong dependency direction.
Dead monitor on ModelCheckpoint: with save_top_k=0, Lightning 2.6.5's _save_topk_checkpoint returns immediately, so monitor/mode are never read — dead, misleading config (best-model selection is actually SaveBestModelStateDict; this only writes last.ckpt). Drop monitor=/mode= or comment that it only saves last.
Fail-loud is asymmetric but adequate: only SaveBestModelStateDict raises explicitly. EarlyStopping covers itself via strict=True; ModelCheckpoint can't (see above). So a bad early_stopping_monitor/checkpoint_monitor still aborts — worth a one-line PR note that early-stopping's fail-loud comes from Lightning, not this code.

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

albanpuech · 2026-07-17T08:38:03Z

@p-grontas Should I check again?

albanpuech

Two final things

docs/quick_start/yaml_config.md still says checkpoint_monitor is for “best-model saving and checkpointing” — slightly stale now that ModelCheckpoint only writes last.ckpt.
Test comment in test_pipeline.py still says save-best and checkpoint “share” checkpoint_monitor.

Thank you!

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

p-grontas · 2026-07-29T12:22:28Z

@albanpuech Thanks, good catch!

albanpuech · 2026-07-29T12:45:44Z

Thanks for the work @p-grontas

p-grontas added 9 commits July 13, 2026 15:54

Add monitor metric and mode to CLI

7ea3e20

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Updated contributing, pre-commit and secrets baseline

5a57938

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Updated secrets baseline, contributing, fixed formatting

d7fcf9f

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Uncommented pre-commit on CI

3c4d4ff

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Added integration tests to contributing

24a6e15

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Fixed setuptools version

f5d96d5

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Fixed setuptools on CI

bf6e7d1

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Constrained torch version

d33371f

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Fix torch version

e04e176

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

p-grontas marked this pull request as ready for review July 15, 2026 09:31

Moved metrics to config

2e61090

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

p-grontas changed the title ~~Add configurable --monitor / --monitor_mode flags for training callbacks~~ Make training-callback and LR-scheduler monitored metrics configurable via YAML Jul 16, 2026

albanpuech requested changes Jul 16, 2026

View reviewed changes

albanpuech requested a review from romeokienzler July 16, 2026 13:25

Address PR comments

c80687e

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

p-grontas added 2 commits July 16, 2026 16:24

Address PR comments 2

9a84e9e

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

Removed stale tests

1eef07b

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

p-grontas requested a review from albanpuech July 17, 2026 09:15

albanpuech reviewed Jul 28, 2026

View reviewed changes

Address PR comment

8263b4c

Signed-off-by: Panagiotis Grontas <72163196+p-grontas@users.noreply.github.com>

p-grontas requested a review from albanpuech July 29, 2026 12:22

albanpuech approved these changes Jul 29, 2026

View reviewed changes

	checkpoint_monitor: Validation layer_11_residual
	checkpoint_monitor: Validation loss

	* [ ] Open a PR with a short description of your changes and add Alban Puech as a reviewer.
	* [ ] Open a PR with a short description of your changes.

Uh oh!

Conversation

p-grontas commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Backward compatibility

Usage

Testing

Uh oh!

p-grontas commented Jul 15, 2026

Uh oh!

albanpuech left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

p-grontas Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

albanpuech commented Jul 16, 2026

Uh oh!

albanpuech commented Jul 16, 2026

Uh oh!

albanpuech commented Jul 17, 2026

Uh oh!

albanpuech left a comment

Choose a reason for hiding this comment

Uh oh!

p-grontas commented Jul 29, 2026

Uh oh!

albanpuech commented Jul 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

p-grontas commented Jul 14, 2026 •

edited

Loading

p-grontas Jul 16, 2026 •

edited

Loading