feat(train): stale-CI replay (S34) — repeat steps reuse the window-first CI envelope#953
Open
danbraunai-goodfire wants to merge 2 commits into
Open
feat(train): stale-CI replay (S34) — repeat steps reuse the window-first CI envelope#953danbraunai-goodfire wants to merge 2 commits into
danbraunai-goodfire wants to merge 2 commits into
Conversation
…rst CI envelope Under free-AT batch replay, a window's taps are identical across its steps, and the (~10x-the-target) CI fn fwd+vjp is the dominant step slab at 4L. New config-gated data.replay_stale_ci: repeat steps consume the window-first step's CI envelope as a constant — no taps read, no CI-fn forward/backward, ci_fn + optimizer state pass through untouched (it trains only on window-first steps, matching the 1-per-window unique-data rate). V/U + both adversaries see identical gradients either way (CI live vs constant doesn't change their backward), pinned by test_stale_ci_repeat_step_semantics. - train.py: make_train_step internals -> shared _build_step_impl(cached=None | CIEnvelope); flag-off traces the pre-S34 step unchanged (op order preserved). make_stale_ci_train_steps returns fresh/repeat/compute_ci; repeat's first arg bundles (model, envelope) so all-except-first donation spares both. - ci_fn LR under stale CI: unit-LR optimizer (run_state.unit_lr_ci_fn_optimizer) + in-step scheduled_value_traced at the GLOBAL step — optax's count advances only on window-first steps, so a count-driven cosine would stretch by the replay factor; Adam bias correction correctly tracks the update count. - run.py: do_step dispatches fresh/repeat on step % train_batch_replay; a mid-window requeue-resume rebuilds the envelope via compute_ci (one ci_fn update fresher than the interrupted window's — preemption-only wobble, noted in S34). Grad-norm window stats take the key union (repeats carry no ci_fns family). - configs: staleci 400k arms (replay 2/4/8/16, group-twin the running replay-only sweep) + a replay1 flag-off code baseline from this branch. SPEC S34 added (experimental, off by default; flag-off unamended). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QKw3iikejknWeF4ZQDFZbL
Twin of the scaling sweep's nw2 control (only run_name differs), run from this branch's codebase with replay_stale_ci off. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QKw3iikejknWeF4ZQDFZbL
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.
Description
Config-gated
data.replay_stale_ci(SPEC S34, off by default): under free-AT batch replay, the repeat steps of a window reuse the window-first step's CI envelope as a constant — no taps read, no CI-fn forward or backward, and the ci_fn + its optimizer state pass through untouched. The ci_fn trains only on window-first steps, matching the 1-per-window unique-data rate; its LR still follows the run's global-step schedule (unit-LR optimizer + in-stepscheduled_value_traced, since optax's count-driven cosine would stretch by the replay factor).Mechanically:
make_train_step's internals become a shared_build_step_impl(cached: CIEnvelope | None); flag-off traces the pre-S34 step unchanged (op order preserved for compile-cache stability).make_stale_ci_train_stepsreturnsfresh/repeat/compute_ci;repeat's first arg bundles(model, envelope)soall-except-firstdonation spares both. The engine'sdo_stepdispatches onstep % train_batch_replay; a mid-window requeue-resume rebuilds the envelope viacompute_ci(one ci_fn update fresher than the interrupted window's — a preemption-only wobble, noted in S34).Also ships the sweep arms:
pile_ppgd_bsc_nw0_lr0p01_replay{2,4,8,16}_staleci_400k.yaml(config-twins of the running replay-only arms, differing only in the flag) plusreplay1_staleci_base_400k.yaml, a flag-off replay=1 baseline from this branch to pin that the refactored flag-off path reproduces the parent branch's trajectory at scale.Related Issue
—
Motivation and Context
The CI fn is ~10x the target and its fwd+vjp is the dominant step slab in the pile-4L runs; within a replay window its inputs (taps) are bit-identical across steps, so repeats recompute it for one Adam step of param freshness. Skipping it on repeats is the same staleness bet free-AT itself makes, and V/U + adversary gradients are provably unchanged by CI being live vs constant. Expected step-time cut on repeats is large at 4L; the replay-scaling sweep (2/4/8/16, group
ppgd-replay-staleci-400kon cw-east, A/B against the runningppgd-replay-scaling-400k) measures both the speedup and whether decomposition quality survives.How Has This Been Tested?
test_stale_ci_repeat_step_semantics: repeat step with the fresh step's envelope produces the same components + adversary updates as the fresh step (allclose), leaves ci_fn + optimizer state bit-untouched,compute_cireproduces the fresh step's envelope, and the metric families differ only by the absentci_fnsgrad norms.make test: all pass except 4test_slow_evalrenderer tests that fail locally on the Mac GUI matplotlib backend (pass underMPLBACKEND=Agg; unrelated to this diff).XLA_FLAGS=--xla_force_host_platform_device_count=4(llama_simple_mlp / checkpoint / sharding): pass.make_train_step→_build_step_impl).make check/make type: clean.Does this PR introduce a breaking change?
No. Flag off (the default) emits the unchanged step;
DataConfiggains a requiredreplay_stale_cifield with a single lab-side construction site.🤖 Generated with Claude Code
https://claude.ai/code/session_01QKw3iikejknWeF4ZQDFZbL