feat: seamless Bergson Trainer -> SOURCE attribute - #367
Merged
Conversation
SOURCE was built for HF Trainer checkpoints. It loads each checkpoint with from_pretrained, reads the step out of a `checkpoint-<N>` directory name, and scrapes per-step learning rates from log_history.json / trainer_state.json. A bergson run satisfies none of that: it writes DCP `step_<i>.ckpt` directories, and no HF artifacts at all. Every input had to be reconstructed by hand. Worse, one of those hand-typed inputs is silently wrong by default. bergson's SGD takes its momentum from TrainingConfig.adam_beta1 (see prepare_trainer), which defaults to 0.95, while the unrolling assumed no momentum. The terminal-velocity factor 1/(1-beta) (Bae et al. 2024, App. D.2) is therefore 20x, so every segment's lr*steps was understated 20-fold unless the user knew to look. `momentum` is now a config field derived from the run that trained the checkpoints. The pieces: - trainer_run.py resolves what a config leaves unset -- checkpoints, model_path, momentum -- from a bergson run's config.yaml. Called once at the top of the pipeline. - The trainer records its realized per-step LRs as log_history.json beside the checkpoints, in HF's shape. That is the file the LR math already checks first, so bergson runs join the existing path instead of branching it, and the schedule is exact rather than reconstructed. - trainer_export.py turns DCP checkpoints into `checkpoint-<i>/` model dirs from_pretrained can load, carrying each step's optimizer state and the LR history across. Offline, so a run only pays the disk for trajectories being attributed. - _checkpoint_step also accepts `step_<N>[.ckpt]`. Derivation is strictly a fallback: it only ever supplies a value the config did not give. Explicit lr_list/step_size_list/momentum/model_path/checkpoints always win, and with no `trainer_run` set, resolve() is a no-op -- so checkpoints from HF Trainer or any other trainer behave exactly as before. That is pinned by tests for each field, including that an explicit momentum of 0.0 survives (it is a real value, not "unset"). optimizer_placement.py lands here rather than in the ADAM SOURCE branch, since the export needs it and it is not Adam-specific.
for more information, see https://pre-commit.ci
The trainer wrote log_history.json beside its checkpoints and the export copied it next to the exported ones, because the reader derived the path from checkpoints[0].parent. Two locations that could drift, and a comment to explain why. Resolve it through trainer_run instead, so the file has one home and the export copies nothing.
The trainer wrote optimizer state as a sibling of the checkpoint it belonged to, which severed the association and forced everything downstream to rebuild it: optimizer_placement.py matched loose files back to checkpoint dirs by step name, or positionally when the names did not say, and had to reject half-named lists because a wrong match would silently build a preconditioner from another step's moments. Writing it inside step_<i>.ckpt/ removes the problem rather than solving it. The file travels with its checkpoint, so the export is a plain copy from one directory into another and nothing needs matching. Verified DCP is unbothered: an extra non-DCP file inside a checkpoint dir leaves dcp.load working, a resumed run reproducing the same parameters, and the file itself intact across a re-save -- pinned by a test, since the layout depends on it. place_final_optimizer_state went too. TrackStar's load_optimizer already takes either a file or a directory, so pointing optimizer_state at <run>/optimizer.pt needs no placement step at all; that helper solved a problem that did not exist. No migration path for the old sibling layout, by choice: re-export instead.
save_optimizer_state only ever wrote the end-of-training state, so SOURCE's Adam variant had nothing per checkpoint to read. Trainer.train now takes optimizer_cfg and writes each snapshot's second moments to optimizer.pt inside that snapshot's own directory, tagged with the step and the optimizer's betas/eps/eps_root so the file is self-describing. Inside rather than beside: the state belongs to one checkpoint, so keeping it there means nothing downstream has to match a loose file back to a step. The write happens before the DCP save starts, into the directory DCP is about to populate -- safe because DCP leaves unrelated files in a checkpoint directory alone, which is pinned by test_dcp_tolerates_optimizer_state_inside_the_checkpoint. The Adam SOURCE branch still writes the old sibling layout and should be fixed up to match.
export_checkpoints defaults to <run>/exported/ but discover_checkpoints only globbed <run>/checkpoint-*, so a default export left resolve() raising 'export them first' against checkpoints that were already there. Both sides now share EXPORT_DIRNAME, and discovery still accepts an export pointed at the run root.
Never called, and duplicated _checkpoint_step in approx_unrolling_math.
Pointing at a bergson run twice was redundant: an exported checkpoint already says which run it came from. Setting checkpoints is now enough; trainer_run is only needed to discover them or to name a different run. Only the two layouts we emit are probed (<run>/exported/checkpoint-N and <run>/checkpoint-N), so an unrelated config.yaml higher up is never mistaken for the producing run, and foreign checkpoints infer nothing.
save_optimizer_state was a bool that, after the per-checkpoint write landed, meant 'all checkpoints' -- no way to ask for just the final state, and no way to opt out of one file per checkpoint when only the normalizer was wanted. Now none/last/all. 'last' is the original meaning (final state at <run>/optimizer.pt, the attribution normalizer); 'all' adds each checkpoint's own optimizer.pt, which SOURCE's Adam variant needs and which costs a file the size of the trainable parameters per checkpoint. Old booleans still parse: true -> 'last', false -> 'none', so existing configs keep their original behaviour rather than silently gaining per-checkpoint files.
You always say which checkpoints to attribute, and an exported checkpoint already records which run produced it, so naming the run again was redundant. The field and its discovery path are gone; resolve() infers the run from checkpoints[0] and fills only model_path and momentum. Raw step_<i>.ckpt paths now raise up front naming export_checkpoints, instead of failing later inside from_pretrained -- the one useful thing the discovery path did.
luciaquirke
force-pushed
the
feat/source-trainer-integration
branch
from
July 29, 2026 14:11
89b5142 to
49668fe
Compare
The round-trip test proved the save/load mechanism but not the function's wiring: reading the run's config.yaml, building the model through prepare_trainer, honouring steps=, and landing in the default export dir with each checkpoint's optimizer state alongside. Ends by feeding the result to resolve(), so the train -> export -> SOURCE path is covered end to end.
luciaquirke
force-pushed
the
feat/source-trainer-integration
branch
from
July 29, 2026 14:15
478f13e to
c2f252e
Compare
for more information, see https://pre-commit.ci
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.
checkpoint-<i>/model dirsfrom_pretrainedcan load.step_<N>[.ckpt].