EK-FAC: average query gradients over model checkpoints - #440
Open
luciaquirke wants to merge 5 commits into
Open
Conversation
The EK-FAC counterpart of ValidationConfig.ckpt_avg_k, completing what D9 asks for: both scorers seeded with the averaged query gradient. MAGIC computes query gradients inline, so averaging there is a loop over TrainerState checkpoints. EK-FAC instead reads them from a prebuilt mmap index, so the averaging has to happen when that index is BUILT: with two or more query_model_paths, step 1 of hessian_pipeline builds the query gradient once per checkpoint into query__ckpt<i> and averages them elementwise into the usual query path. Steps 2-4 are untouched. query_model_paths defaults to empty, which is the existing single-model behaviour exactly. average_gradient_indices refuses sources that disagree on row count or module layout -- those would be different query sets, and averaging them would be silently wrong rather than an error. It accumulates in float64 so a long running sum does not drift before the divide. Documented limitation: paths must load the way index_cfg.model does (HF directory or hub name). The MAGIC trainer writes its trajectory as TrainerState dcp checkpoints, which are not that format and must be exported first. 7 tests.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JzPbQBTC457nM5dnHFz7oF
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JzPbQBTC457nM5dnHFz7oF
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JzPbQBTC457nM5dnHFz7oF
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JzPbQBTC457nM5dnHFz7oF
Collaborator
Author
|
cc @LouisYRYJ I think I merged the MAGIC one already but this one looks a bit more messy |
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.
Adds
HessianPipelineConfig.query_model_paths, the EK-FAC counterpart ofValidationConfig.ckpt_avg_k(#439). With two or more paths, the query gradient is built once per checkpoint and averaged elementwise before the Hessian is applied. Empty (the default) builds it once fromindex_cfg.model, so existing runs are unchanged.Implementation
average_gradient_indicesinbergson/data.pyaverages several gradient indices into a new one. Sources must agree on row count, module layout and dtype -- they are the same query set at different checkpoints, so a mismatch means the indices aren't comparable. Accumulation is in float64 regardless of the store dtype, so a long running sum doesn't drift before the divide.hessian_pipelinebuilds each checkpoint's query index into<query_path>__ckpt<i>, then averages them intoquery_path.Note on checkpoint format
query_model_pathsentries load the same wayindex_cfg.modeldoes (HF model directory or hub name). The MAGIC trainer writes its trajectory asTrainerStatedcp checkpoints, which are not that format -- they need exporting first. This is called out in the config docstring.Testing
tests/test_ekfac_query_ckpt_avg.py-- one test covering the averaging math, the info.json passthrough, the single-source case, float64 accumulation over 64 sources, and the threeValueErrorguards. Passes.🤖 Generated with Claude Code
https://claude.ai/code/session_01JzPbQBTC457nM5dnHFz7oF