feat(checkpoint): split checkpoints into decomposition | training items (S22/S33)#950
Open
ocg-goodfire wants to merge 5 commits into
Open
feat(checkpoint): split checkpoints into decomposition | training items (S22/S33)#950ocg-goodfire wants to merge 5 commits into
ocg-goodfire wants to merge 5 commits into
Conversation
…ms (S22/S33 amended) Each step now stores two orbax items: decomposition (V/U components + ci_fn — the trained product) and training (both optimizer states, persistent adversaries, step — the process). Consumers (open_jax_run, fine-tune init S33) restore ONLY decomposition via restore_decomposition[_to_host] against an eval_shape'd run_state.init_decomposition reference — zero knowledge of training's optimizer/adversary structure, and none of the ~60GB opt-moment/bsc-source materialization that wedged single-GPU harvest restores (supersedes the interim partial-restore 1ca895860, which still needed eval_shape(init_train_state) incl. build_optimizers). Runs saved before the split (single default item) keep writing that layout from their immutable workspaces; every restore path dual-reads both layouts (_stored_as_split, grounded in the stored dir layout: <step>/default vs <step>/decomposition). Both items save sharded — the no-on-loop-full-gather property is unchanged. SPEC S22/S33 amended (pending Oli sign-off). New tests: two-item layout pin, consumer host-restore, legacy dual-read across restore_latest / restore_decomposition_to_host / init_from_parent. Validated vs real stored data: p-594db290 (8B L18, 60GB-class ckpt) opens on CPU in 22s. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EwzrwmBiMpkvDovpGgVysn
The pile_ppgd_bsc_RESUMESMOKE_dp8.yaml was a manual e2e testing artefact used to validate the two-item checkpoint format on 8 GPUs, not code intended to be committed. Removing per review. Crew-Address: task/checkpoint-split-decomposition-item
… break) Per review: no in-code compat for pre-split `default`-item checkpoints. Pre-split runs are migrated by an ad-hoc script when needed, rather than bloating every restore path with a fallback. - Remove `_stored_as_split` layout dispatch and `_legacy_partial_restore`. - Collapse restore_step / restore_decomposition / restore_decomposition_to_host to the two-item (`decomposition`+`training`) path only. - SPEC S22/S33, CLAUDE.md: state the clean break; drop dual-read wording. - Replace the legacy dual-read test with a two-item init_from_parent test (keeps S33 fine-tune-init coverage: decomposition carries over, opt state/adversaries/step stay the fresh reference's). Suite green at 1 and 4 simulated devices (param_decomp/tests/test_checkpoint.py, 7 passed). Crew-Address: task/checkpoint-split-decomposition-item
ocg-goodfire
commented
Jul 8, 2026
Comment on lines
+58
to
+78
| def _split(state: TrainState) -> tuple[Decomposition, TrainingItem]: | ||
| return ( | ||
| Decomposition(components=state.components, ci_fn=state.ci_fn), | ||
| TrainingItem( | ||
| components_opt_state=state.components_opt_state, | ||
| ci_fn_opt_state=state.ci_fn_opt_state, | ||
| adversaries=state.adversaries, | ||
| step=state.step, | ||
| ), | ||
| ) | ||
|
|
||
|
|
||
| def _join(decomposition: Decomposition, training: TrainingItem) -> TrainState: | ||
| return TrainState( | ||
| components=decomposition.components, | ||
| ci_fn=decomposition.ci_fn, | ||
| components_opt_state=training.components_opt_state, | ||
| ci_fn_opt_state=training.ci_fn_opt_state, | ||
| adversaries=training.adversaries, | ||
| step=training.step, | ||
| ) |
Collaborator
Author
There was a problem hiding this comment.
are these actaully necessary? they smell
…representation
Per review: rather than a flat TrainState that the checkpoint code regroups
into two items (the _split/_join helpers), TrainState now IS the composition of
the two checkpoint items:
TrainState{decomposition: Decomposition, training: TrainingItem}
so save/restore map straight onto its own fields — _split/_join deleted, no
regrouping anywhere. TrainingItem moved from checkpoint.py to train.py (it's now
a TrainState field). SPEC S22 records the one-representation invariant.
Churn: every state.<field> access becomes state.decomposition.<field> /
state.training.<field> — the hot train step, run.py resume/warmup, run_state
init, slow_eval hidden-acts, the lab run/eval paths, and all constructing tests.
Pure regrouping, no numeric change.
Validated (1 + 4 sim devices): test_checkpoint, test_checkpoint_production_topology,
test_generic_model_io, test_no_bake_invariant, test_llama8b, test_llama_simple_mlp,
test_finetune_resume, stacked_parity, tms, resid_mlp — all green (the only failure,
test_pretrained_target_converts_with_wildcards, is the pre-existing tip failure).
basedpyright clean; ruff clean.
Crew-Address: task/checkpoint-split-decomposition-item
…position item The pre-split-run migration the S22 clean format break requires, doubling as the storage lever: a pre-split checkpoint is ~90% resume-only state (persistent adversary + all optimizer moments); the tool partial-restores components+ci_fn as host numpy and rewrites ckpts/<step> as the two-item layout with ONLY the decomposition item, staged + tree-verified before an in-place swap. Dry-run by default, --apply to strip. The run can no longer warm-resume — for finished or abandoned runs only. Validated on a scratch copy of p-e118ece9/ckpts/355000 (19.2 -> 2.1 GiB, 7s): restore_decomposition_to_host over the stripped checkpoint is byte-equal to the original across all 163 decomposition leaves. Handles both current (adversaries) and old jax_runs-era (sources/sources_opt_state) key layouts; unknown keys fail. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RkKMDdM1PxaL66RmEJv8QN
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.
What
Checkpoints store two orbax items per step, and
TrainStateCOMPOSES exactly those two — one representation shared by the trainer and the checkpoint:decomposition—train.Decomposition(V/Ucomponents+ci_fn): the trained product.training—train.TrainingItem(both optimizer states, persistent adversaries,step): the process.TrainState = {decomposition, training}, sosave_state/restore_stepmap straight onto its own fields (no_split/_joinregrouping). Every consumer —open_jax_run(harvest/autointerp/clustering/app) and fine-tune init (S33) — restores onlydecomposition, against an abstract tree fromjax.eval_shape(run_state.init_decomposition): zero knowledge of how training builds its optimizers or adversaries.Why
The 2026-07-06 harvest wedge:
open_jax_runrebuilt + restored the FULLTrainState, so on one consumer GPU an 8B run's PPGD bsc sources + Adam moments (~60GB, sharded ÷64 during training) materialized whole, twice → OOM/allocator-wedge. The interim fix (1ca895860, only on localharvest-tonight— superseded here) partial-restored{components, ci_fn}but still neededjax.eval_shape(init_train_state)includingbuild_optimizersjust to recover the saved tree's structure. Consumers knowing training internals was the design smell; the item split removes it at the storage boundary.Compatibility — clean break, no in-code fallback
Clean checkpoint-format break. There is no dual-read of the pre-split single-
default-item layout — every restore path reads the two-item layout only. Runs checkpointed before this change must be converted by an ad-hoc migration script; the compat logic does not live in the codebase. Both items save sharded — the no-on-loop-full-gather write property is unchanged.SPEC
S22 + S33 amended (marked pending Oli sign-off in the text): the two-item layout, the
TrainState-composes-both invariant, the consumer contract, and the clean-break/no-fallback decision are normative.Validation
XLA_FLAGS=--xla_force_host_platform_device_count=4:test_checkpoint,test_checkpoint_production_topology,test_generic_model_io,test_no_bake_invariant,test_llama8b,test_llama_simple_mlp,test_finetune_resume,stacked_parity, and the labtms/resid_mlptoy suites.basedpyright+ruffclean on the change. (The one failure,test_pretrained_target_converts_with_wildcards, is the pre-existing tip failure, unrelated.)init_from_parent(S33:decompositioncarries over; optimizer state / adversaries /step=0stay the fresh reference's), andtest_roundtrip_and_exact_resumestill bit-exact through the composed state.<step>/{decomposition,training}; job 171288 resumed from the two-item step-40 checkpoint;open_jax_run(p-f653e1fd)opened it in 6s on CPU and ran a clean harvest forward. Manual smoke config not committed.Notes for review
restore_stepkeeps the format-coercion tail (device_putontoreference.format); fix(checkpoint): make orbax-restored TrainState donatable — solve the resume OOM (S22) #932 rewrites that same tail — whichever lands second has a small conflict.CheckpointManagerpins an item's handler at first use, so the saving manager can'tPyTreeRestorewhat itStandardSaved — consumers open their own manager.consumer-config-schema-drift, so no live harvest regresses on this PR alone.bridge thread: checkpoint-split-decomposition-item
Crew-Address: task/checkpoint-split-decomposition-item