AdamW SOURCE - #327
Open
luciaquirke wants to merge 10 commits into
Open
Conversation
luciaquirke
changed the title
luciaquirke
force-pushed
the
source-damp-options
branch
9 times, most recently
from
July 20, 2026 14:34
d36f4f9 to
cfdb782
Compare
luciaquirke
force-pushed
the
source-damp-options
branch
from
July 25, 2026 09:48
cfdb782 to
570a531
Compare
luciaquirke
force-pushed
the
source-damp-options
branch
4 times, most recently
from
July 29, 2026 11:44
92faf8f to
0e376b9
Compare
Spell out the abbreviated preconditioner names so they read alongside hessian_path: EkfacConfig.precond_path/precond_post_multiply -> preconditioner_path/preconditioner_post_multiply, and DiagonalFactoredPreconditioner's multiply_by_precond -> multiply_by_preconditioner. Rename Trainer.train's snapshot_optimizer_hparams -> optimizer_hparams. Drop the getattr() guards on save_optimizer_state/save_retrained_models in magic's worker: both are TrainingConfig fields now, and worker takes a TrainingConfig. Trim the DiagonalFactoredPreconditioner and adam_preconditioner docs to the formula and the citation; the "called on every rank" FSDP comment is already pinned by test_save_optimizer_state_completes_under_fsdp, whose docstring explains the hang it prevents. The only untested part of this feature was the config wiring, so drive a real EkfacApplicator with preconditioner_path set through an on-disk gradient index and compare against DiagonalFactoredPreconditioner.apply for both eigenfunctions. That is what catches a rename that fails to reach compute_ivhp_sharded. CPU-only, 0.25s.
DiagonalFactoredPreconditioner.apply reshaped the caller's tensor and scaled it in place. compute_ivhp_sharded passes gradients that alias a read-only mmap and relies on the preconditioner returning fresh tensors (as FactoredPreconditioner does, via its rotation matmuls). The aliasing only bites when the preceding `.to(device, float32)` is a no-op, i.e. on a CPU-only run: writing to the mapping segfaults. Copy before scaling. Caught by test_ekfac_applicator_uses_diagonal_path on CPU-only CI; also pin the no-write-through contract directly.
The .to(...).clone() that fixed the read-only-mmap write copied twice on the GPU path: once host-to-device, once more on device. copy=True gets the same guarantee in a single copy.
The flag moved from ValidationConfig up to TrainingConfig, where "retrained" no longer describes what it does: on a plain training run nothing is re-trained, it just saves the model that run produced. Every config that carries it trains models, so name it for that. Metasmoothness gets the flag by inheritance rather than defining its own near-duplicate. Each run type still writes what it actually produces -- HF-format models under retrained/ for training and validation, theta_init.pt and theta_final.pt for metasmoothness -- and the docstring spells both out. The old key still parses via a deprecated save_retrained_models alias that warns and forwards, so existing YAML configs and model banks keep working; setting both to conflicting values is an error. The alias is marked for removal by 2026-10-29.
…ad them The trainer writes each snapshot's second moments as a sibling file, <save_dir>/step_<i>.optimizer.pt, next to that step's step_<i>.ckpt. Both consumers instead want optimizer.pt *inside* a directory: SOURCE reads <checkpoint>/optimizer.pt under use_adam_preconditioner, and TrackStar resolves PreprocessConfig.optimizer_state through load_optimizer, which loads optimizer.pt from a directory. Exporting to HF-format checkpoint dirs does not carry the sibling files across, so both paths previously told the user to copy the files by hand. place_optimizer_states() does that placement for a run's checkpoints, and place_final_optimizer_state() does it for the end-of-training state that TrackStar normalizes with. Since both consumers agree on <dir>/optimizer.pt, one call satisfies either. States are matched to destinations by step index when the destination names one, else positionally in step order with the counts required to agree. A half-named list, an unknown step, or a count mismatch raises: silently mismatching would build a preconditioner from the wrong step's moments, which no downstream check would catch. Existing files are never clobbered without overwrite=True, and the check precedes all writes so a refusal leaves nothing half-placed. Defaults to relative symlinks -- second moments are one scalar per weight, so copying one per checkpoint is wasteful, and relative keeps an exported tree valid when moved. copy/hardlink/move are available for self-contained exports. OPTIMIZER_STATE_FILE now has a single definition, re-exported from adam_preconditioner so pipeline.py's import is unchanged.
The examples/ additions were untracked scratch swept in by a careless git add -A during the rebase. optimizer_placement.py and its tests were superseded on main, which keeps each checkpoint's optimizer state inside the checkpoint rather than matching loose files back to steps.
luciaquirke
force-pushed
the
source-damp-options
branch
from
July 29, 2026 14:53
0e376b9 to
298f262
Compare
The rename is done and every config in the repo uses save_models, so the deprecated alias, its conflict check and the FutureWarning are just weight. Also trims the field's docstring to the three destinations it actually writes.
* feat: Fisher normalization for SOURCE segment eigenvalues
Segment eigenvalues were unnormalized sums over (documents x checkpoints).
EK-FAC is unaffected -- every inversion is homogeneous of degree -1 in lambda
because damping is relative to mean(lambda) -- but SOURCE's eigenfunctions are
not scale-invariant, and sigma enters exp(-lr*K*sigma). Eq. 1 of Bae et al.
2024 defines the empirical risk as a mean over data points, so
fisher_normalization defaults to per-document; token units put the mask
argument in its dead-linear range.
Also makes precompute_checkpoints resumable rather than overwrite-only, and
lets resolve() fall back when a checkpoint's sibling config.yaml turns out to
belong to an attribution run.
* feat: per-snapshot optimizer.pt export and interval save mode in MAGIC trainer
(cherry picked from commit 16c29730ff6fdadd1c61ea72dc9dda43e6a85274)
* fix: restrict the end-of-training snapshot to interval save mode
The other save modes' backward replay indexes the data stream by checkpoint,
so an extra trailing snapshot runs it past the end (DataStream index out of
range in test_fsdp_ddp_scores_match).
* fix: make skip_metagradient actually skip the MAGIC backward
The flag only toggled allow_compile; the backward still ran (and crashed on
interval-mode end-of-training snapshots). Port the dummy-zero-scores skip
path so bank-only builds go straight to validation.
(cherry picked from commit 13b631d3b030be312a881ceff6c31475cbceb4b5)
* feat: skip_metagradient config field for trainer-only MAGIC runs
Companion to the cherry-picked gate: defines the MagicConfig field and the
do_metagradient condition it reads.
* feat: Eq-43 hybrid for the Adam SOURCE segment eigenfunction
Bae et al. 2024 App. D: 'use the diagonal Hessian approximation for computing
the matrix exponential in Equation 39 and Equation 43. Note that we still use
the EK-FAC factors to compute H^-1 g in Equation 43.'
Evaluating the eigenfunction wholly on the diagonal drops that second half.
With adam_segment_hybrid the diagonal supplies only the exponential mask
(f_one_minus_exp) and a chained EK-FAC inverse supplies 1/x, applied to the
query as the adjoint H^-1 @ M_mask @ q.
* docs: WikiText-2 replication configs (train / EK-FAC / SOURCE)
Paper recipe from Bae et al. 2024 App. B.1 (AdamW, lr 3e-5, wd 1e-2, batch 8,
seq 512, 6 checkpoints) with train_mode enabled so GPT-2's 0.1 dropout is
active, matching kronfluence's train.py. SOURCE uses the App. C preconditioned
variant with the App. D Eq-43 hybrid, L=3, per-document Fisher units.
* feat: LDS against a precomputed subset-retrain ground truth
Scoring side of a published retrain bank (kronfluence ships masks.pt /
losses.pt for its WikiText-2 experiment): sum each subset's member scores,
Spearman against the measured subset losses, average over queries. Lets a
replication be checked against published numbers without retraining.
* fix: read the training config out of the saved {steps, metadata} document
save_run_config has always wrapped the step list in a document with a
metadata block, so load_training_config rejected every real run directory;
its test hand-rolled the bare list it expected instead of calling
save_run_config. Have the test write configs the way the writer does, and
report an unparseable payload as ValueError so callers guessing at a run
dir still degrade to their fallback.
* docs: driver for the WikiText-2 replication chain
Export -> EK-FAC IF -> SOURCE -> LDS against kronfluence's shipped bank.
EK-FAC 0.4451 against their reported 0.44; SOURCE 0.4506, above IF as in
the paper's Figure 6.
* style: apply black formatting
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.
Note that replicating the Kronfluence number for GPT-2 using the Bergson Trainer requires passing --train_mode to enable train mode/dropout for now (dropout does not play nicely with double backward/MAGIC)