Skip to content

Stage the base builder into checkpointed processes with a bit-equivalence proof#406

Merged
MaxGhenis merged 12 commits into
mainfrom
staged-base-builder
Jul 13, 2026
Merged

Stage the base builder into checkpointed processes with a bit-equivalence proof#406
MaxGhenis merged 12 commits into
mainfrom
staged-base-builder

Conversation

@MaxGhenis

Copy link
Copy Markdown
Contributor

What

The reusable base builder becomes a staged, checkpointed pipeline: --stage {a,b,c,d,all} + --checkpoint-dir, with each stage runnable as its own process so its memory dies with it. --stage all with no checkpoint dir is byte-identical to today's behavior.

Why

Four OOM kills at full scale (79–92GB dirty RSS on a 128GB host) across escalating mitigations (#401 chunked predict, #402 forest release, worker/batch reductions) — the single-process design accretes stages and outgrew the machine. Stack-sample + per-stage profiling localized the drivers; a blind cross-model design review found two failure modes parameter tuning could never fix.

The fixes, per the reconciled design

  • Per-target QRF chain checkpointing: the 64-target chain retained every fitted model internally; the chain now executes per-target with raw prior draws, exact target order, and both RNG streams checkpointed between targets — chained conditioning semantics preserved exactly.
  • _snap_to_observed_values memory bound: the recipient×unique-values matrix was unbounded (a latent fifth OOM independent of staging); now chunked, bit-identical, unit-tested.
  • Fresh-process outer stages with lossless Frame checkpoints (WeightKind, strata, MassChangeRecord log round-trip — the PE H5 loader's CALIBRATED coercion explicitly avoided), low-copy writes.
  • Per-stage RSS/wall profiling to stage_profile.json.

The proof

test_us_puf_support_base_equivalence.py: staged-vs-monolith bit equivalence on a subsample under a pinned single-thread environment, compared at every stage boundary (raw-bit observer via --equivalence-boundary-dir), with deterministic H5 metadata mode disabling leaf timestamps. Full suite green locally (exit 0), plus the checkpoint round-trip identity suite and builder tests (53 green).

First consumer: Build M's base rebuild, immediately on merge.

@MaxGhenis MaxGhenis merged commit 556166a into main Jul 13, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant