Register all parameters in optimizer groups explicitly#1529
Register all parameters in optimizer groups explicitly#1529steffen-wedig wants to merge 2 commits into
Conversation
get_params_options previously built named parameter groups for the MACE backbone only; submodules it did not know about (the PolarMACE electrostatics blocks, ~36% of the weights of the released POLAR-1 models) were silently excluded from the optimizer and stayed at their initialization during training. - Register optional submodules by name (radial_embedding, pair_repulsion_fn, joint_embedding, embedding_readout, les_readouts, les, lr_source_maps, fukui_source_map, field_dependent_charges_maps, local_electron_energy, layer_feature_mixer), skipping absent or parameter-free ones. - Raise ValueError naming any trainable parameter no group claims, so future submodules fail loudly at optimizer construction instead of silently not training. - Regenerate the hardcoded reference energies of the single-head foundation finetuning tests (workflows/test_freeze.py, workflows/test_run_train.py::test_run_train_foundation): the previous values encoded a removed artifact where the foundation model's bessel_weights, promoted to an unregistered trainable Parameter by the old finetuning loader, inflated the global clip_grad_norm_ norm (22% on the first step) and rescaled every real update. All other training tests pass unchanged. - Add tests/unit/test_optimizer_param_groups.py covering every model class, the trainable-Bessel case, and the unregistered-submodule raise.
aacostadiaz
left a comment
aacostadiaz
left a comment
There was a problem hiding this comment.
Thanks for tracking this down. A few comments below.
| "AtomicDipolesMACE": build_atomic_dipoles_mace, | ||
| "AtomicDielectricMACE": build_atomic_dielectric_mace, | ||
| "EnergyDipolesMACE": build_energy_dipoles_mace, | ||
| "PolarMACE": pytest.param( |
There was a problem hiding this comment.
I think this PolarMACE case never runs on PR CI. The unit job doesn't install graph_longrange (so the skipif always skips) and the polar job only collects tests/extensions/polar. Could we move it to tests/extensions/polar/ and use marks=pytest.mark.polar instead of the skipif?
There was a problem hiding this comment.
Moved both the Polar and Magnetic Mace test to the extension dir.
| # every trainable parameter is registered explicitly; weight decay stays | ||
| # off because these blocks mix biases, gates, and physically meaningful | ||
| # scalars that must not be regularized toward zero. | ||
| optional_submodule_names = [ |
There was a problem hiding this comment.
The optimizer group count goes from 5 to ~10, so --restart_latest from a checkpoint made before this change fails to load the optimizer state, and run_train swallows that error, so training silently restarts at epoch 0 with a fresh optimizer (verified locally). Worth logging a clear warning when the optimizer state fails.
There was a problem hiding this comment.
I added error handling to the mace/tools/checkpoint.py::load_checkpoint function, that adds a warning on ValueErrors from optimizer reloading and a regression test that triggers this.
| } | ||
| ) | ||
|
|
||
| # Guard against silently untrained submodules: any trainable parameter |
There was a problem hiding this comment.
Small caveat: LES builds its Atomwise MLP lazily in the first forward(), after the optimizer is created. So those parameters still slip past the guard. Pre-existing, but maybe worth a sentence in the comment so the guard doesn't read as a complete guarantee.
| ) | ||
|
|
||
|
|
||
| MODEL_BUILDERS = { |
There was a problem hiding this comment.
MagneticScaleShiftMACE is missing
…ion test coverage.
2 participants
I came across an issue in registering module parameters for the optimizer during training of PolarMace models.
get_params_options registers modules parameters by name, instead of using models.parameters() to get all trainable parameters. Submodules that get_params_options did not know about (
lr_source_maps,fukui_source_map,field_dependent_charges_mapsetc. in PolarMACE) are silently excluded from the optimizer and stayed at their initialization during training. I verified that this issue is not relevant for FMs, but only for from scratch trained PolarMace models.This PR makes the following changes:
Making these changes surfaced two pre-existing test failures on the develop branch and I updated their fixtures.
They were not raised (as far as I can tell) because the tests are marked as requiring network connection, but the ci-core.yaml doesnt set allow-network to true, so the CI skips them - not sure if thats intentional. The tests are workflows/test_freeze.py and workflows/test_run_train.py::test_run_train_foundation. Before PR 1244 , when the test fixtures were generated, the bessel_weights were always reloaded as nn.Parameters (making them accumulate gradients), but they were not added to the optimizer (because no explicit by-name mention in get_params_options) . Gradient clipping however uses model.parameters() irrespective of which parameters are registered in the optimizer, so thereby the tests that test finetuning of foundation models are impacted no matter if radial_weights are trained or not. After PR 1244, the radial_weights are loaded correctly as buffers in mace/tools/finetuning_utils.py::_copy_radial_weights which changes the behaviour of the test, therefore the fixtures had to be regenerated causing small changes in the energies.
All other training tests pass unchanged, also adding tests/unit/test_optimizer_param_groups.py covering every model class, the trainable-Bessel case, and the unregistered-submodule raise.
Something that should be confirmed is whether zero weight decay is correct for all these modules or if they should have individual lr values, and whether these settings agree with the training setup for the foundation models. Please lmk if I should change something!