Re-join optimizer hyperparameters to param groups by key on reconnect#19
Closed
cedriclim1 wants to merge 1 commit into
Closed
Re-join optimizer hyperparameters to param groups by key on reconnect#19cedriclim1 wants to merge 1 commit into
cedriclim1 wants to merge 1 commit into
Conversation
reconnect_optimizer_to_parameters restored per-group hyperparameters by zip-index against a fresh get_optimization_parameters() call. Group order and membership from that dict are not contractual (the dataset model's groups change with alignment settings), so any difference between creation time and reconnect time silently attached the wrong learning rate to every group -- no error, just a reconstruction quietly training with swapped lrs. set_optimizer now records each group's key in the torch param group under "name" (torch preserves unknown keys), and reconnect re-joins hyperparameters by that key; groups with no old counterpart keep optimizer defaults. Optimizers restored from checkpoints predating group names fall back to the previous index alignment. Regression tests cover reorder, group addition, state survival, and the legacy fallback.
Owner
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
reconnect_optimizer_to_parameters(used after device moves / AutoSerialize loads; called from the tomography object and dataset models) restored per-group hyperparameters by zip-index against a freshget_optimization_parameters()call. Nothing makes that dict's order or membership contractual — the dataset model's groups change with alignment settings — so any difference between creation time and reconnect time silently attaches the wrong learning rate to every group. No error is raised; a PPLR reconstruction just trains with swapped lrs (e.g. grids at 1e-3 and sigma_net at 1e-2).This was flagged as a design fragility in the loop's first sweep; it is reproducible, so it gets a fix rather than a note.
Fix
set_optimizerrecords each group's key in the torch param group as"name"(torch preserves unknown keys in param groups and state_dicts).reconnect_optimizer_to_parametersre-joins old hyperparameters to new groups by that key; a group with no old counterpart keeps optimizer defaults."name"present) fall back to the previous index alignment, so old checkpoints behave exactly as before.Evidence
TestReconnectHyperparamAlignment(5 tests): group reorder keeps each lr with its key (failed before — lrs swapped silently), added group gets defaults while others keep lrs (failed before), Adam state survives reconnect, legacy no-name fallback, names recorded at creation.uv run pytest tests --runslow: 434 passed, 1 skipped; only the pre-existing unrelated ptychography failure.Independent branch off the integration tip; no overlap with the other open PRs.